diff --git a/README.md b/README.md index 2b944df2af49b0dc89b0d63aebc86b70b0abab46..d98eef08a7cfccc21e9328d36518d0b6ed01e680 100644 --- a/README.md +++ b/README.md @@ -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) +kernel_module = get_kernel("kernels-community/fla", version=2) chunk_kda = kernel_module.chunk_kda chunk_kda(...) diff --git a/build/torch-cuda/__init__.py b/build/torch-cuda/__init__.py index 93947cdf3ec9f3c35772e484b1160934eeda5ade..c2984f39791c741c03d918327acd8785d509ac1f 100644 --- a/build/torch-cuda/__init__.py +++ b/build/torch-cuda/__init__.py @@ -1,3 +1,5 @@ +__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 diff --git a/build/torch-cuda/_ops.py b/build/torch-cuda/_ops.py index e04a067e52d184bc6a7b05c7a25feb67196cc25b..805ddaa3c49606692e95dcf03e07a35d8d6c80fd 100644 --- a/build/torch-cuda/_ops.py +++ b/build/torch-cuda/_ops.py @@ -25,7 +25,7 @@ def get_backend() -> str: def _find_ops_name() -> str: kernel_name = "fla" - unique_id = "fe2bb66" + unique_id = "faa8674" backend = get_backend() return f"_{kernel_name}_{backend}_{unique_id}" diff --git a/build/torch-cuda/backends.py b/build/torch-cuda/backends.py new file mode 100644 index 0000000000000000000000000000000000000000..8531c320e1b3e04260ef093a3f8c940840704f98 --- /dev/null +++ b/build/torch-cuda/backends.py @@ -0,0 +1,265 @@ +# 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 + +"""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") + + +class BaseBackend: + """Base class for operation-specific backends. + + Attributes: + backend_type (str, Optional): + Identifier for the backend type, used to distinguish different backend implementations. + `FLA_` 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. + `None` indicates no external dependency. Default: `None`. + env_var (str, Optional): + Environment variable name that controls whether the backend is enabled. + `None` means always enabled. Default: `None`. + default_enable (bool, Optional): + Whether the backend is enabled by default when `env_var` is not set. + Set to `False` to require explicit user opt-in. Default: `True`. + priority (int, Optional): + Backend priority. Lower values indicate higher priority. Default: 5. + """ + + 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: + if cls.package_name is None: + 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'] diff --git a/build/torch-cuda/layers.py b/build/torch-cuda/layers.py index 42cb865a5dd7150f2c543c009206dc23dea1d9ae..675015a3ac1833aac58d15b190571f05273d3087 100644 --- a/build/torch-cuda/layers.py +++ b/build/torch-cuda/layers.py @@ -7,6 +7,8 @@ from .ops.kda import chunk_kda, fused_recurrent_kda class FusedRMSNormGated(nn.Module): + can_torch_compile = True + 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): + can_torch_compile = True + def forward( self, query, @@ -54,6 +58,8 @@ class chunk_kimi_delta_attention(nn.Module): 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, ): @@ -75,6 +81,8 @@ class recurrent_kimi_delta_attention(nn.Module): class chunk_gated_delta_rule(nn.Module): + can_torch_compile = True + def forward( self, query, @@ -107,6 +115,8 @@ class chunk_gated_delta_rule(nn.Module): 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, ): diff --git a/build/torch-cuda/metadata.json b/build/torch-cuda/metadata.json index e4b3e1c7cd9b974ee2d2ec057abc80e75be3fd14..3869a3308ad85a2da3abe937c565df94b8c05f74 100644 --- a/build/torch-cuda/metadata.json +++ b/build/torch-cuda/metadata.json @@ -1,7 +1,7 @@ { "name": "fla", - "id": "_fla_cuda_fe2bb66", - "version": 1, + "id": "_fla_cuda_faa8674", + "version": 2, "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=", + "__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/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=", + "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": "sqvLGP2r+frjpL2qWK8GHCoWWPU5Df8IT5y20oaTZTE=", + "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": "3xcuWH+iLBPRb3g4ZkyXBXm8wJFxQK4GHdUpmK8Nh88=", - "ops/attn/decoding.py": "ScKiTKcbl4kcnYDNz6jXd2a1qB5ymDcz8YVkotRnO9Y=", - "ops/attn/naive.py": "nBsWqmBWOMQdaQYvdFejDLoX0Pei9W8oMPrHZoCCin0=", - "ops/attn/parallel.py": "+Vkkjy4OIwR+7SXBf6fCJqrPg740qc6a1+21RSNt9Cs=", + "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/fused.py": "Q8RKiwlUtFZkHELCLV3wtQGf1cVDPoDNMK1waEwbxFQ=", + "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/backends/__init__.py": "3DU4e6t7dmPuZLlHjPnt4q/MAtDXtaMmAcSY4ycyytY=", "ops/based/__init__.py": "qdBLUigpgjQaIN3HTtq57uxoqAaBnQkZbFiEIf18Rvk=", - "ops/based/fused_chunk.py": "Evgk49vwK6VxWWmw3KdwhiUV5N8RF/PcxZP8J1dlYR8=", + "ops/based/fused_chunk.py": "gBng4smV8X5eHn0pgo677bsu8kvOZ+DNEkJbvegBxZo=", "ops/based/naive.py": "YWIFSLeslp2ruyQDz/uCr5V4ao2cyRyeQH1wiVLMB8U=", - "ops/based/parallel.py": "7hNqqI2UIeJNebDTF7RwW+Q9ST15DwMxJugiaVO9s/k=", + "ops/based/parallel.py": "daBATK+vtrog3XYFx2WkTfyg0QPH0gGb/6zlaTZCeIs=", "ops/comba/__init__.py": "FBKonbvPPFHY7an2Auh1sXVJLbcgwX8aAD5fxi2oecY=", - "ops/comba/chunk.py": "o+3soCUQ27uZXe+YOZpYF1/aiLFEyTdbb+bmetY6O2c=", - "ops/comba/fused_recurrent.py": "ujg1r9D2Jx+/sgli1msfkXoiCZb2gX9dL/OpOca3kGs=", + "ops/comba/chunk.py": "dtxGTXfgVY0c8TMq6uSBDUBjaP9tNk5WsnScEyVDVpI=", + "ops/comba/fused_recurrent.py": "sVNBWn5bq7PI7QIpCGbDMGBfl0CMzDzv767eR+5jCiw=", "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/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": "+pPn7+hNImlTrCZLHNpaJKabA+OaoNXofYbX2F4fTD0=", + "ops/cp/chunk_delta_h.py": "OHQD+CubC1RfbNqxsRewZVh/NSNqUMEOuplaNm9FI7U=", "ops/cp/comm.py": "lqyQDwmK19vowTtGJGO+6SEn/neFFE0G8xFvlxLOmb0=", - "ops/cp/context.py": "N/p87Th0z6m2cGhP6icCN7rEg3QPVua8ymP8oDD4qy0=", + "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": "r1b6g3/4OO1oDZ8n/MMyzagYHs0c5FPKU8//7KskU8w=", + "ops/delta_rule/chunk.py": "ueCVdrAgOpCxVVb7gDgihDRLQck9dY0dOQB/qdxHxGc=", "ops/delta_rule/fused_chunk.py": "2NiUBRAl98YgRVGr2uzYJV22pQW3WNOe6euTX1Mdqyw=", - "ops/delta_rule/fused_recurrent.py": "aerY1sPJ0y6IknruBVOwbu3v/pLz2PjP3aNEzsArp5M=", + "ops/delta_rule/fused_recurrent.py": "WEg6WYVHkdvmncKPtXfYky4skqIe5whzjrX9rE3PAHo=", "ops/delta_rule/naive.py": "IK4DL04IAZtCmMyfzyw5fYcfqsD9HxgCQ9n2makZKWg=", - "ops/delta_rule/parallel.py": "GMhpAZoEDuzyIjK0Qcpu+ku528WnLgDUujcsoRnXU6o=", - "ops/delta_rule/wy_fast.py": "O0UuKiOw3mImzRQHFLKHifSewFAoqRJ9i9r0puiX9fw=", + "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": "8b88rviKkDZpW+F9GwgEkVkwuSioECwgRJX//CVksnI=", - "ops/deltaformer/parallel.py": "sDHrnr2fM4+0T2cyAJoI2VP9Lan3NmaaDxk9JLtyLWk=", + "ops/deltaformer/naive.py": "W78CzYUyuxLY9fEkKghSh2dnjux0NOJ6pAVhIN7M4PU=", + "ops/deltaformer/parallel.py": "XMKE9sLFJl4UVspOXQxhefztzykDEUp7+S9GyUhmOwo=", "ops/dsa/__init__.py": "ZNyws889LPBiY8gwB5zTs2vQwJYV8bX5jg6Cldkkq9s=", - "ops/dsa/naive.py": "lKgrdKOoE/qCdF4DtrI0uEkRdB28kUYbtn8hO63fzyo=", + "ops/dsa/naive.py": "JKyMUzwK44XrwlKUWfLhU/iK5DlCsXceJuxtXrzaCQ8=", "ops/forgetting_attn/__init__.py": "A6ozg9zdEsz4lgUAXqEIUY8CeNp+5IZqYKOpChpj/A4=", - "ops/forgetting_attn/naive.py": "RDPXw01ZnhTMV125OFvEVeDDMht3BYj5gQI4oaT+T1M=", - "ops/forgetting_attn/parallel.py": "fcKF7Q84XXmeeKI2Tpv5XIfmxFovuXlmBov7Ie8OjKg=", + "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": "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.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": "Q907vtcavFpalFlttThkWSb/SaTa4vqAG+I/nWvKDno=", + "ops/gated_delta_product/naive.py": "bJ1Cc3mljHcAXjc4w1JDyxj+i4hmhCPsoELKrTkyvIM=", "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/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": "kFbW7BX1I4DBBBhfV1yXGQR4Jb9tGM6yILNmTXNOhxo=", + "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": "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/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/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/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": "DNH9FhRaZl4nqZ/YrxH2SqCgZdC77bkJ24sFIemnkRM=", + "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/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/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": "HiMBn/U+a0vSCOVmn29UYLylI5x26qqualdl4E6ipd0=", - "ops/generalized_delta_rule/dplr/wy_fast_fwd.py": "R+CamNa7b49LfvrNZQs4roi+TReWcLGvBcoT3DbhoLM=", + "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": "BsCf/fiYn3wpS0eDPrR+mCrqsAihXD5v/QD2QkA5ATc=", - "ops/generalized_delta_rule/iplr/fused_recurrent.py": "nNo+ww/QmydOqBLnpmxXHAbANGDPTeWGjFI7nA9Vrw0=", + "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": "SoTcF2asZEIAnZS1qR9PXHGkHwnBspjuUsz6RpuqNiU=", + "ops/generalized_delta_rule/iplr/wy_fast.py": "hqVfLtmdlhQnlIeU7LrByAxBcNYtVSehVFPZW/o8ZhY=", "ops/gla/__init__.py": "PTUqqQlC4cXrhRhG/KRx+7X3m9WJKxGyzrIB/WJEXSo=", - "ops/gla/chunk.py": "B7reIP0UzTvYgK8Ds50r0uABlaucJWm70JVKKdYe1L0=", + "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": "ZIzMstFj8WAXagDzTlZo3hFZy0k+egHNg8Qr61sb4vw=", + "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": "SsfwmTOpKRG2Fja/lArJkw1teaomk/Vgczjr2/Oa+l0=", - "ops/gsa/fused_recurrent.py": "K7qBWOmh0uL9CO7D3pSlrxKFHy1fobTUyTFV5p7okXg=", - "ops/gsa/naive.py": "rezNszsU38KD1iHBW+ErV/Z01T5LIPadEjJnWiXz9LI=", + "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": "Br3GeHkw1EcOF/4X2BOvUTclhfLDZwkthIshoHTMMqg=", - "ops/hgrn/fused_recurrent.py": "KvO/5l/3wnJiT0jT0UA9IYxQhGHJy8jC7lYaOOyLM1w=", + "ops/hgrn/chunk.py": "odEz3XdQ1whBWOhYQfCVvMS+zy6QZ36hb9USuqjwV6w=", + "ops/hgrn/fused_recurrent.py": "SddIcmnP1KUcGWTk2rJGtPEoM5CnuV+zzMOUmK4pc98=", "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/__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": "9pMMqS5E/l0mFV/Tyap1QKY39CB5Iy94/Yn4s9WzjUw=", + "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": "kFjfxpUOVj61TnH0vXWZhDnw318suRLhtT2Qo50pkKM=", - "ops/lightning_attn/fused_recurrent.py": "SsJ+HC5opio7H9LsPWyrbnPOoK2EvOcxLuaq8Vp69GY=", + "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": "7X94oxSkHAGSpk7SQ5izs5Q0rYEDZZzUV/RorFdzY1o=", + "ops/linear_attn/utils.py": "tjqU7ST991YBV68Zs83Fy/li6sUSMeKNPGE3VcYMvfo=", "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/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": "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/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": "8qUN2zo88Pldh809UyDowBVj8kCNplsS00fIxLRdSyQ=", + "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": "XjFqvEXvuIEeCBrJU+YzhWRH5UAkPGcUX4Tz/xx2M6Y=", - "ops/nsa/naive.py": "5A7lyoN39oK0Wc+Yzm/dQbNEtlaGmZfJsMDv8II+m7c=", - "ops/nsa/parallel.py": "BLaPXXsO8tcAJ79ltV3xwX/1qTganF2Nh+Lto/3zvyQ=", + "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": "qHX0Le4hA66S2LLTcMa+p0VMB0uWvhYGI9urSPLrB4I=", - "ops/parallax/naive.py": "XEvAU0cSJtunA3ElSy4GFs6idarjK7Vb+xkrDKSp2OQ=", - "ops/parallax/parallel.py": "ZrLKoopSj60DSxmqbaS3MIf3AUyEpYh7evJPlTUN+AQ=", + "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": "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/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": "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/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": "+3WRG3oxV0T78e8HWkN32Q8KBe0V2GMQ/OpxowWfXFI=", + "ops/rebased/parallel.py": "cBnn/KTq9FlWJIpc6Jj9uJ7J14VfGFadrxABo4bsoQc=", "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/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": "gTnLcV5w17fZAOJyJttw52LGAe9XlkuS++lMGTZ3vRs=", + "ops/retention/parallel.py": "WRBtqBxCWpIZ/Lm2u3bXwnbeLbwwQ9ZrrajeYeNBoTg=", "ops/rwkv4/__init__.py": "GnWakdtZiQJ0+fTYxqqIyosoPwjH6Yp8P1AbpHIDnts=", - "ops/rwkv4/fused_recurrent.py": "0fvq/4+NZMq6tzq7oHMUCaPg10qOzteSKLj5cmQKhvk=", + "ops/rwkv4/fused_recurrent.py": "LG9PdTsubmJvgqvbDvT8N+9y5SWQZZXfTxC0Hg1jU5s=", "ops/rwkv6/__init__.py": "wPYgC6wG+xc4I2vessGhv1EaZ4tL5yuVQVSwi1rEEUk=", - "ops/rwkv6/chunk.py": "2VMAk6p0tKEudhieXLIg2bYg+89JbsGUiuhFpfvx6C8=", + "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": "wMTP/Y17ayzQXxXw+fN8P/FB5KGefLB4bB0wJ9kDg6Q=", + "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": "+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/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": "hB6Zq+allqAYg/H8/9kpk0l7bsZCK9IsITkfLIhqZQ8=", + "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": "NUmchQQpsxFwNaEpSQbwZ/638zSr0tuy2xGa3DD07+A=", + "ops/titans/naive.py": "0Y4R5926mX2CgYI55MjxVfn6p2X4iM4pABCzmVZ9d5c=", "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/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": "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/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": "JphI1kWj0Wwu024gnwgUWrWmhkAnKloXi1ik6cDpRrY=", + "ops/wall_attn/decode.py": "o5q0fsJTiz+0XW0RPFFI1hBXz6PXSBMSTvog8Bc+sGA=", "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=" + "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": { @@ -314,7 +355,7 @@ "dirty": false }, "kernel": { - "commit": "fe2bb66fee6c2a77eccf08a0508dc57428d488ed", + "commit": "faa8674647aa5ebaab5e997cd11e6e3808d090a9", "dirty": false } } diff --git a/build/torch-cuda/metadata.json.sigstore b/build/torch-cuda/metadata.json.sigstore index 02e743e09baf17634414070c61e836a0bdda5039..16c2ab1494b7a1dab0f9f3d4a9aaeed8a4f2ff47 100644 --- a/build/torch-cuda/metadata.json.sigstore +++ b/build/torch-cuda/metadata.json.sigstore @@ -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="}} \ No newline at end of file +{"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"}} \ No newline at end of file diff --git a/build/torch-cuda/modules/__init__.py b/build/torch-cuda/modules/__init__.py index 0936901f1de9349346957b28291b183e5a8c4d71..6ea67bf4bfb273b25a1483ed82f667c44e9558e8 100644 --- a/build/torch-cuda/modules/__init__.py +++ b/build/torch-cuda/modules/__init__.py @@ -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 +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.fused_norm_gate import ( +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.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 +_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', diff --git a/build/torch-cuda/modules/activations/__init__.py b/build/torch-cuda/modules/activations/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..723721c42c4c40b256de146705edaf718dab9aa4 --- /dev/null +++ b/build/torch-cuda/modules/activations/__init__.py @@ -0,0 +1,39 @@ +# 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.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', +] + diff --git a/build/torch-cuda/modules/activations.py b/build/torch-cuda/modules/activations/ops.py similarity index 64% rename from build/torch-cuda/modules/activations.py rename to build/torch-cuda/modules/activations/ops.py index 93fec5b0796f079fd320a189384229ec7a2fa181..11f536676d71c10d4c3d488ccc9f945b25b4e45a 100644 --- a/build/torch-cuda/modules/activations.py +++ b/build/torch-cuda/modules/activations/ops.py @@ -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. +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 ..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] +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 the tensor is "inner-contiguous" (stride(-1) == 1 and - higher dims are contiguous relative to dim -2) before using this value. + 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. + 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 + return x.is_contiguous() 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): + if torch.compiler.is_compiling(): + return x.contiguous() + if _is_inner_contiguous(x=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. + ``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.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, -): +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.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, -): +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') +@dispatch def sigmoid_fwd(x: torch.Tensor, output_contiguous: bool = False) -> torch.Tensor: - x = _ensure_inner_contiguous(x) + x = _ensure_inner_contiguous(x=x) T, D = x.numel(), x.shape[-1] - y = _alloc_output(x, output_contiguous) + 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), - stride_y_row=_get_stride(y), + stride_x_row=_get_stride(x=x), + stride_y_row=_get_stride(x=y), T=T, D=D, ) return y -@dispatch('modules') +@dispatch 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) + x = _ensure_inner_contiguous(x=x) + dy = _ensure_inner_contiguous(x=dy) T, D = x.numel(), x.shape[-1] - dx = _alloc_output(x, output_contiguous) + 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), - stride_dy_row=_get_stride(dy), - stride_dx_row=_get_stride(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, ) @@ -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) + return sigmoid_fwd(x=x) @staticmethod @input_guard(no_guard_contiguous=True) def backward(ctx, dout): x, = ctx.saved_tensors - return sigmoid_bwd(x, dout) + return sigmoid_bwd(x=x, dy=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.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, -): +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 @@ -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.autotune(configs=_activation_autotune_configs(), 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') +@dispatch def logsigmoid_fwd(x: torch.Tensor, temperature: float = 1., output_contiguous: bool = False) -> torch.Tensor: - x = _ensure_inner_contiguous(x) + x = _ensure_inner_contiguous(x=x) T, D = x.numel(), x.shape[-1] - y = _alloc_output(x, output_contiguous) + 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), - stride_y_row=_get_stride(y), + stride_x_row=_get_stride(x=x), + stride_y_row=_get_stride(x=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') +@dispatch 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) + x = _ensure_inner_contiguous(x=x) + dy = _ensure_inner_contiguous(x=dy) T, D = x.numel(), x.shape[-1] - dx = _alloc_output(x, output_contiguous) + 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), - stride_dy_row=_get_stride(dy), - stride_dx_row=_get_stride(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, @@ -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) + 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, dy, ctx.temperature), None + 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=[ - 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.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, -): +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.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, -): +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') +@dispatch def swish_fwd(x: torch.Tensor, output_contiguous: bool = False) -> torch.Tensor: - x = _ensure_inner_contiguous(x) + x = _ensure_inner_contiguous(x=x) T, D = x.numel(), x.shape[-1] - y = _alloc_output(x, output_contiguous) + 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), - stride_y_row=_get_stride(y), + stride_x_row=_get_stride(x=x), + stride_y_row=_get_stride(x=y), T=T, D=D, ) return y -@dispatch('modules') +@dispatch 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) + x = _ensure_inner_contiguous(x=x) + dy = _ensure_inner_contiguous(x=dy) T, D = x.numel(), x.shape[-1] - dx = _alloc_output(x, output_contiguous) + 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), - stride_dy_row=_get_stride(dy), - stride_dx_row=_get_stride(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, ) @@ -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) + return swish_fwd(x=x) @staticmethod @input_guard(no_guard_contiguous=True) def backward(ctx, dout): x, = ctx.saved_tensors - return swish_bwd(x, dout) + return swish_bwd(x=x, dy=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)) +# 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) -# 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)""" + """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 - ) + 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) + return bias_gelu(y=input, bias=bias) @staticmethod def backward(ctx, grad_output): input, bias = ctx.saved_tensors - return bias_gelu_bwd(grad_output, input, bias) + return bias_gelu_bwd(g=grad_output, y=input, bias=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') +# 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) -# 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 - ) + 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) + return gelu_fwd(x=input) @staticmethod def backward(ctx, grad_output): (input,) = ctx.saved_tensors - tmp = gelu_bwd(grad_output, input) + tmp = gelu_bwd(g=grad_output, x=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) + return sqrelu_fwd(x=input) @staticmethod def backward(ctx, grad_output): input, = ctx.saved_tensors - return sqrelu_bwd(grad_output, input) + return sqrelu_bwd(g=grad_output, x=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.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, -): +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.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, + 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') +@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) - y = _ensure_inner_contiguous(y) + x = _ensure_inner_contiguous(x=x) + y = _ensure_inner_contiguous(x=y) T, D = x.numel(), x.shape[-1] - z = _alloc_output(x, output_contiguous) + 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), - stride_y_row=_get_stride(y), - stride_z_row=_get_stride(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('modules') +@dispatch 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) + 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, output_contiguous) - dy = _alloc_output(y, output_contiguous) + dx = _alloc_output(x=x, contiguous=output_contiguous) + dy = _alloc_output(x=y, contiguous=output_contiguous) if use_weight: - z = _alloc_output(x, output_contiguous) + z = _alloc_output(x=x, contiguous=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, + 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, ) @@ -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) + 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, y, dout) + return swiglu_fwdbwd(x=x, y=y, g=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) + z = swiglu_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 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) + 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('modules') +@dispatch 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.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, -): +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.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, + 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) + x = _ensure_inner_contiguous(x=x) + y = _ensure_inner_contiguous(x=y) T, D = x.numel(), x.shape[-1] - z = _alloc_output(x, output_contiguous) + 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), - stride_y_row=_get_stride(y), - stride_z_row=_get_stride(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, ) @@ -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) + 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, output_contiguous) - dy = _alloc_output(y, output_contiguous) + dx = _alloc_output(x=x, contiguous=output_contiguous) + dy = _alloc_output(x=y, contiguous=output_contiguous) if use_weight: - z = _alloc_output(x, output_contiguous) + z = _alloc_output(x=x, contiguous=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, + 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, ) @@ -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) + 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, y, dout) + return sigmoidglu_fwdbwd(x=x, y=y, g=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) + 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 @@ -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) + 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 @@ -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.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, -): +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.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, + 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') +@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) - y = _ensure_inner_contiguous(y) + x = _ensure_inner_contiguous(x=x) + y = _ensure_inner_contiguous(x=y) T, D = x.numel(), x.shape[-1] - z = _alloc_output(x, output_contiguous) + 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), - stride_y_row=_get_stride(y), - stride_z_row=_get_stride(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, @@ -1075,7 +1018,7 @@ def powglu_fwd(x: torch.Tensor, y: torch.Tensor, power: float = 3.0, output_cont return z -@dispatch('modules') +@dispatch 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) + 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, output_contiguous) - dy = _alloc_output(y, output_contiguous) + dx = _alloc_output(x=x, contiguous=output_contiguous) + dy = _alloc_output(x=y, contiguous=output_contiguous) if use_weight: - z = _alloc_output(x, output_contiguous) + z = _alloc_output(x=x, contiguous=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, + 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, @@ -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). + 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) + 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, y, dout, ctx.power) + dx, dy = powglu_fwdbwd(x=x, y=y, g=dout, power=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) + 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 @@ -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) + 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 @@ -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') +@dispatch def powglu_linear( x: torch.Tensor, y: torch.Tensor, diff --git a/build/torch-cuda/modules/backends/__init__.py b/build/torch-cuda/modules/backends/__init__.py deleted file mode 100644 index 4ce5614452d211b97584dab79c74560f9c6d4c1f..0000000000000000000000000000000000000000 --- a/build/torch-cuda/modules/backends/__init__.py +++ /dev/null @@ -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'] diff --git a/build/torch-cuda/modules/backends/triton_ascend/__init__.py b/build/torch-cuda/modules/backends/triton_ascend/__init__.py deleted file mode 100644 index ba9c7a925c84d63ebed71917d64fcdcadd33df6c..0000000000000000000000000000000000000000 --- a/build/torch-cuda/modules/backends/triton_ascend/__init__.py +++ /dev/null @@ -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, - ) diff --git a/build/torch-cuda/modules/backends/triton_ascend/activations.py b/build/torch-cuda/modules/backends/triton_ascend/activations.py deleted file mode 100644 index 7f43a381e8b3d806cbfde5b3891c92e888aa467f..0000000000000000000000000000000000000000 --- a/build/torch-cuda/modules/backends/triton_ascend/activations.py +++ /dev/null @@ -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) diff --git a/build/torch-cuda/modules/backends/triton_ascend/causal_conv1d.py b/build/torch-cuda/modules/backends/triton_ascend/causal_conv1d.py deleted file mode 100644 index 9f77b9ec28b1e91197c1d47d6a66f181f83c2176..0000000000000000000000000000000000000000 --- a/build/torch-cuda/modules/backends/triton_ascend/causal_conv1d.py +++ /dev/null @@ -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 diff --git a/build/torch-cuda/modules/backends/triton_ascend/fused_cross_entropy.py b/build/torch-cuda/modules/backends/triton_ascend/fused_cross_entropy.py deleted file mode 100644 index 42aaec21dc9dc0461aac19c0b1e998efca04cb48..0000000000000000000000000000000000000000 --- a/build/torch-cuda/modules/backends/triton_ascend/fused_cross_entropy.py +++ /dev/null @@ -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, - ) diff --git a/build/torch-cuda/modules/backends/triton_ascend/fused_kl_div.py b/build/torch-cuda/modules/backends/triton_ascend/fused_kl_div.py deleted file mode 100644 index e593e06b377ce4d0843dde112540f498ef659ce6..0000000000000000000000000000000000000000 --- a/build/torch-cuda/modules/backends/triton_ascend/fused_kl_div.py +++ /dev/null @@ -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 diff --git a/build/torch-cuda/modules/backends/triton_ascend/fused_linear_cross_entropy.py b/build/torch-cuda/modules/backends/triton_ascend/fused_linear_cross_entropy.py deleted file mode 100644 index 3d8d9fda021352e37359fad061135aa7236812f4..0000000000000000000000000000000000000000 --- a/build/torch-cuda/modules/backends/triton_ascend/fused_linear_cross_entropy.py +++ /dev/null @@ -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 diff --git a/build/torch-cuda/modules/backends/triton_ascend/grpo.py b/build/torch-cuda/modules/backends/triton_ascend/grpo.py deleted file mode 100644 index 463923f5d560b23b7248c1ff0df6256ed222445b..0000000000000000000000000000000000000000 --- a/build/torch-cuda/modules/backends/triton_ascend/grpo.py +++ /dev/null @@ -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) diff --git a/build/torch-cuda/modules/backends/triton_ascend/layernorm.py b/build/torch-cuda/modules/backends/triton_ascend/layernorm.py deleted file mode 100644 index 9a8fe6ba39c7c99d496ba52ff4f349c6a951bf07..0000000000000000000000000000000000000000 --- a/build/torch-cuda/modules/backends/triton_ascend/layernorm.py +++ /dev/null @@ -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) diff --git a/build/torch-cuda/modules/backends/triton_ascend/rotary.py b/build/torch-cuda/modules/backends/triton_ascend/rotary.py deleted file mode 100644 index 4042cc23c6b8e1eb49f6e853703d68176e5960b4..0000000000000000000000000000000000000000 --- a/build/torch-cuda/modules/backends/triton_ascend/rotary.py +++ /dev/null @@ -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 diff --git a/build/torch-xpu/modules/conv/triton/__init__.py b/build/torch-cuda/modules/causal_conv1d/__init__.py similarity index 68% rename from build/torch-xpu/modules/conv/triton/__init__.py rename to build/torch-cuda/modules/causal_conv1d/__init__.py index e4489bc867adb04eaa1b23e2abd0cd060204b7c3..657f07ca8cbb1ba1744d6b013a96d877eb0291e1 100644 --- a/build/torch-xpu/modules/conv/triton/__init__.py +++ b/build/torch-cuda/modules/causal_conv1d/__init__.py @@ -5,20 +5,27 @@ # For a list of all contributors, visit: # https://github.com/fla-org/flash-linear-attention/graphs/contributors -from .ops import ( +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, - compute_dh0_triton, ) __all__ = [ 'CausalConv1dFunction', + 'ShortConvolution', + 'causal_conv1d', 'causal_conv1d_bwd', 'causal_conv1d_fwd', 'causal_conv1d_update', 'causal_conv1d_update_states', - 'compute_dh0_triton', ] + + +if GluonBackend.is_available(): + from ...modules.causal_conv1d.backends import gluon # noqa: F401 diff --git a/build/torch-cuda/modules/causal_conv1d/backends/__init__.py b/build/torch-cuda/modules/causal_conv1d/backends/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..546163768c14bff626769f78ac7700ed0e3398ea --- /dev/null +++ b/build/torch-cuda/modules/causal_conv1d/backends/__init__.py @@ -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 diff --git a/build/torch-cuda/modules/conv/cuda/__init__.py b/build/torch-cuda/modules/causal_conv1d/backends/cuda/__init__.py similarity index 60% rename from build/torch-cuda/modules/conv/cuda/__init__.py rename to build/torch-cuda/modules/causal_conv1d/backends/cuda/__init__.py index 1e1485fa2aee518131fb55be9a7addc870b5f819..4040937295ed2c7345b4163514ca2eb1c567a54a 100644 --- a/build/torch-cuda/modules/conv/cuda/__init__.py +++ b/build/torch-cuda/modules/causal_conv1d/backends/cuda/__init__.py @@ -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 +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', -] +__all__ = ['FastCausalConv1dFn', 'causal_conv1d_cuda', 'fast_causal_conv1d_fn'] diff --git a/build/torch-cuda/modules/causal_conv1d/backends/cuda/cp.py b/build/torch-cuda/modules/causal_conv1d/backends/cuda/cp.py new file mode 100644 index 0000000000000000000000000000000000000000..16253d833b938575275f0d29c9f8e6ffbc27e2d4 --- /dev/null +++ b/build/torch-cuda/modules/causal_conv1d/backends/cuda/cp.py @@ -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 diff --git a/build/torch-xpu/modules/conv/cuda/ops.py b/build/torch-cuda/modules/causal_conv1d/backends/cuda/ops.py similarity index 59% rename from build/torch-xpu/modules/conv/cuda/ops.py rename to build/torch-cuda/modules/causal_conv1d/backends/cuda/ops.py index 4d9ee19fa935251c955e714b93a179ac27c7982f..974d15d699e97417c756c1ba5eef83c172af0799 100644 --- a/build/torch-xpu/modules/conv/cuda/ops.py +++ b/build/torch-cuda/modules/causal_conv1d/backends/cuda/ops.py @@ -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.""" +"""CUDA and mixed Triton/CUDA implementations of 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 +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 @@ -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 + """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, - weight, - bias=None, + x: torch.Tensor, + weight: torch.Tensor, + bias: torch.Tensor | None = None, residual: torch.Tensor | None = None, - initial_states=None, - output_final_state=False, - activation=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" @@ -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 = 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.conv.triton.ops import causal_conv1d_fwd + # import here to avoid circular dependency. + from .....modules.causal_conv1d 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 - ) + 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): + 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') @@ -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, + 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, - ) + 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) +) -> 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( @@ -195,7 +178,7 @@ def causal_conv1d_cuda( 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: @@ -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) + 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, - activation=activation, 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, - ) + 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'] diff --git a/build/torch-cuda/modules/causal_conv1d/backends/gluon.py b/build/torch-cuda/modules/causal_conv1d/backends/gluon.py new file mode 100644 index 0000000000000000000000000000000000000000..5ad99684950d136da26b2fd6d67e2e939c1d4505 --- /dev/null +++ b/build/torch-cuda/modules/causal_conv1d/backends/gluon.py @@ -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 diff --git a/build/torch-cuda/modules/causal_conv1d/cp.py b/build/torch-cuda/modules/causal_conv1d/cp.py new file mode 100644 index 0000000000000000000000000000000000000000..7e8de64b30c00a846082bfb8b4107cb71788febc --- /dev/null +++ b/build/torch-cuda/modules/causal_conv1d/cp.py @@ -0,0 +1,330 @@ +# 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 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): + """ + Context Parallel version of CausalConv1dFunction. + + Forward: + 1. Get tails from previous rank to construct initial_state + 2. Call causal_conv1d_fwd + + Backward: + 1. Call causal_conv1d_bwd to get dx + 2. Sync communication: add next rank's first W-1 token gradients to current rank's last W-1 tokens + """ + + @staticmethod + def _prepare_initial_state_for_cp( + x: torch.Tensor, + weight: torch.Tensor, + cu_seqlens: torch.Tensor | None, + context: FLACPContext, + group: dist.ProcessGroup | None, + ) -> torch.Tensor | None: + """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) + if valid_len > 0: + # 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 + + @staticmethod + def _correct_dx_for_cp( + dx: torch.Tensor, + 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: + # Only keep gradients for positions that had real data from the + # 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] + # Add to current rank's last W-1 tokens (these tokens are used as initial_state by next rank) + dx[0, -(W-1):, :].add_(recv_d_init) + + @staticmethod + def forward( + ctx, + x: torch.Tensor, + weight: torch.Tensor, + bias: torch.Tensor | None, + activation: str | None, + chunk_indices: torch.Tensor | None, + 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 = cp_context.cu_seqlens + 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, + cu_seqlens=cu_seqlens, + context=cp_context, + 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_indices = chunk_indices + 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( + x: torch.Tensor, + weight: torch.Tensor, + bias: torch.Tensor | None = None, + activation: str | None = None, + chunk_indices: torch.Tensor | None = None, + 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) diff --git a/build/torch-cuda/modules/causal_conv1d/ops.py b/build/torch-cuda/modules/causal_conv1d/ops.py new file mode 100644 index 0000000000000000000000000000000000000000..9307a69064432992a90812390198b6819663500a --- /dev/null +++ b/build/torch-cuda/modules/causal_conv1d/ops.py @@ -0,0 +1,1324 @@ +# 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 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 + + +@triton.heuristics({ + 'HAS_WEIGHT': lambda args: args['weight'] is not None, + 'HAS_BIAS': lambda args: args['bias'] is not None, + 'HAS_RESIDUAL': lambda args: args['residual'] is not None, + 'USE_INITIAL_STATE': lambda args: args['initial_state'] is not None, + 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None, +}) +@fla_cache_autotune( + 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 +def causal_conv1d_fwd_kernel( + x, + y, + weight, + bias, + residual, + cu_seqlens, + initial_state, + chunk_indices, + B, + T, + stride_x_n, + stride_x_t, + stride_x_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, + HAS_RESIDUAL: tl.constexpr, + USE_INITIAL_STATE: tl.constexpr, + IS_VARLEN: 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 + else: + i_n = i_b + bos, eos = (i_b * T).to(tl.int64), (i_b * T + T).to(tl.int64) + 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_w = tl.load(weight + o_d[:, None] * W + o_w, mask=m_d[:, None] & m_w, other=0).to(tl.float32) + + 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) + + if HAS_WEIGHT: + b_yi *= tl.sum(b_w * (o_w == (i_w + W - 1)), 1) + b_y += b_yi + + if HAS_BIAS: + b_y += tl.load(bias + o_d, mask=m_d).to(tl.float32) + + if ACTIVATION == 'swish' or ACTIVATION == 'silu': + 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({ + '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, +}) +@fla_cache_autotune( + 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 +def causal_conv1d_bwd_kernel( + x, + y, + 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, + ACTIVATION: tl.constexpr, + HAS_WEIGHT: tl.constexpr, + HAS_BIAS: tl.constexpr, + USE_INITIAL_STATE: tl.constexpr, + USE_FINAL_STATE: tl.constexpr, + IS_VARLEN: 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 + else: + i_tg = i_b * tl.num_programs(1) + i_t + i_n = i_b + bos, eos = (i_b * T).to(tl.int64), (i_b * T + T).to(tl.int64) + p_x = x + tl.cast(i_b, tl.int64) * stride_x_n + + if IS_VARLEN: + p_dy = dy + bos * stride_dy_t + else: + 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) + + b_dx = tl.zeros((BT, BD), dtype=tl.float32) + if HAS_BIAS: + b_db = tl.zeros((BD,), dtype=tl.float32) + + 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 + if HAS_WEIGHT: + # [BT, BD] + b_wdy = b_wdy * tl.sum(b_w * (o_w == (W - i_w - 1)), 1) + # [BD] + b_dw = tl.sum(b_dy * b_x, 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 HAS_BIAS and i_w == 0: + 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 + if HAS_WEIGHT: + # [BT, BD] + b_wdy = b_wdy * tl.sum(b_w * (o_w == (W - i_w - 1)), 1) + # [BD] + b_dw = tl.sum(b_dy * b_x, 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 HAS_BIAS and i_w == 0: + 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 HAS_BIAS and i_w == 0: + b_db += tl.sum(b_dy_shift, 0) + b_wdy = b_dy_shift if not HAS_WEIGHT else (b_dy_shift * tl.sum(b_w * (o_w == (W - i_w - 1)), 1)) + 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 = 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 + + if IS_VARLEN: + p_dx = dx + bos * stride_dx_t + 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({ + 'USE_INITIAL_STATE': lambda args: args['cache'] is not None, + 'HAS_WEIGHT': lambda args: args['weight'] is not None, + 'HAS_BIAS': lambda args: args['bias'] is not None, + 'HAS_RESIDUAL': lambda args: args['residual'] is not None, +}) +@fla_cache_autotune( + configs=[ + triton.Config({'BD': BD}, num_warps=num_warps) + for BD in [8, 16, 32, 64, 128, 256] + for num_warps in NUM_WARPS_AUTOTUNE + ], + key=['D', 'W'], + restore_value=['cache'], + **autotune_cache_kwargs, +) +@triton.jit +def causal_conv1d_update_kernel( + x, + cache, + residual, + y, + weight, + bias, + stride_x_n, + stride_x_d, + stride_y_n, + stride_y_d, + D: tl.constexpr, + W: tl.constexpr, + BD: tl.constexpr, + BW: tl.constexpr, + ACTIVATION: tl.constexpr, + USE_INITIAL_STATE: tl.constexpr, + HAS_WEIGHT: 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) + m_d = o_d < D + m_w = o_w < W + + # [BD] + b_x = tl.load(x + i_n * stride_x_n + o_d * stride_x_d, mask=m_d, other=0).to(tl.float32) + + 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_WEIGHT: + b_w = tl.load(weight + o_d[:, None] * W + o_w, mask=m_d[:, None] & m_w, other=0) + b_y = tl.sum(b_cache * b_w, 1) + else: + b_y = tl.sum(b_cache, 1) + + if HAS_BIAS: + b_y += tl.load(bias + o_d, mask=m_d) + + if ACTIVATION == 'swish' or ACTIVATION == 'silu': + b_y = b_y * tl.sigmoid(b_y) + + 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 +def compute_dh0_kernel( + dy, + 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) + + +@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, +): + 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 + + 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_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: + o_c = W - (BW - seq_len) + tl.arange(0, BW) + m_c = (o_c >= 0) & (o_c < W) + + p_init = initial_state + i_n * D*W + o_d[None, :] * W + o_c[:, None] + mask_init = m_d[None, :] & m_c[:, None] + + 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, + 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, + BD=BD, + BW=BW, + ) + return final_state + + +@dispatch +@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: + # [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) + + 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) -> tuple[int, ...]: + + 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, + ) -> 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] diff --git a/build/torch-cuda/modules/conv/__init__.py b/build/torch-cuda/modules/conv/__init__.py deleted file mode 100644 index 23421fa8c55cf9599bfa4dd8f030b38f599d5d48..0000000000000000000000000000000000000000 --- a/build/torch-cuda/modules/conv/__init__.py +++ /dev/null @@ -1,19 +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 .causal_conv1d import causal_conv1d -from .long_conv import ImplicitLongConvolution, LongConvolution, PositionalEmbedding, fft_conv -from .short_conv import ShortConvolution - -__all__ = [ - 'ImplicitLongConvolution', - 'LongConvolution', - 'PositionalEmbedding', - 'ShortConvolution', - 'causal_conv1d', - 'fft_conv', -] diff --git a/build/torch-cuda/modules/conv/causal_conv1d.py b/build/torch-cuda/modules/conv/causal_conv1d.py deleted file mode 100644 index 98e44923f37df9a02088de4c8836de50f0b9a100..0000000000000000000000000000000000000000 --- a/build/torch-cuda/modules/conv/causal_conv1d.py +++ /dev/null @@ -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}") diff --git a/build/torch-cuda/modules/conv/cp/ops.py b/build/torch-cuda/modules/conv/cp/ops.py deleted file mode 100644 index 3b5b8e758d0c9892b8168109d68821187750a1a8..0000000000000000000000000000000000000000 --- a/build/torch-cuda/modules/conv/cp/ops.py +++ /dev/null @@ -1,258 +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 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): - """ - Context Parallel version of CausalConv1dFunction. - - Forward: - 1. Get tails from previous rank to construct initial_state - 2. Call causal_conv1d_fwd - - Backward: - 1. Call causal_conv1d_bwd to get dx - 2. Sync communication: add next rank's first W-1 token gradients to current rank's last W-1 tokens - """ - - @staticmethod - def _prepare_initial_state_for_cp( - x: torch.Tensor, - weight: torch.Tensor, - cu_seqlens: torch.Tensor | None, - context: FLACPContext, - group: dist.ProcessGroup | None, - ) -> torch.Tensor | None: - """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) - if valid_len > 0: - # 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 - - @staticmethod - def _correct_dx_for_cp( - dx: torch.Tensor, - 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: - # Only keep gradients for positions that had real data from the - # 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] - # Add to current rank's last W-1 tokens (these tokens are used as initial_state by next rank) - dx[0, -(W-1):, :].add_(recv_d_init) - - @staticmethod - def forward( - ctx, - x: torch.Tensor, - weight: torch.Tensor, - bias: torch.Tensor | None, - activation: str | None, - chunk_indices: torch.Tensor | None, - 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") - cu_seqlens = cp_context.cu_seqlens - 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, - cu_seqlens=cu_seqlens, - context=cp_context, - 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 - ctx.chunk_indices = chunk_indices - 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( - x: torch.Tensor, - weight: torch.Tensor, - bias: torch.Tensor | None = None, - activation: str | None = None, - chunk_indices: torch.Tensor | None = None, - 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 - ) diff --git a/build/torch-cuda/modules/conv/cuda/ops.py b/build/torch-cuda/modules/conv/cuda/ops.py deleted file mode 100644 index 4d9ee19fa935251c955e714b93a179ac27c7982f..0000000000000000000000000000000000000000 --- a/build/torch-cuda/modules/conv/cuda/ops.py +++ /dev/null @@ -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 diff --git a/build/torch-cuda/modules/conv/short_conv.py b/build/torch-cuda/modules/conv/short_conv.py deleted file mode 100644 index a53386bc6b751edf4ad94046890024bf35d5029d..0000000000000000000000000000000000000000 --- a/build/torch-cuda/modules/conv/short_conv.py +++ /dev/null @@ -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 diff --git a/build/torch-cuda/modules/conv/triton/__init__.py b/build/torch-cuda/modules/conv/triton/__init__.py deleted file mode 100644 index e4489bc867adb04eaa1b23e2abd0cd060204b7c3..0000000000000000000000000000000000000000 --- a/build/torch-cuda/modules/conv/triton/__init__.py +++ /dev/null @@ -1,24 +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 .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', -] diff --git a/build/torch-cuda/modules/conv/triton/kernels.py b/build/torch-cuda/modules/conv/triton/kernels.py deleted file mode 100644 index 86ae0d9cd76bfb8061b63d2b55e69a353976054b..0000000000000000000000000000000000000000 --- a/build/torch-cuda/modules/conv/triton/kernels.py +++ /dev/null @@ -1,683 +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 -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 - - -@triton.heuristics({ - 'HAS_WEIGHT': lambda args: args['weight'] is not None, - 'HAS_BIAS': lambda args: args['bias'] is not None, - 'HAS_RESIDUAL': lambda args: args['residual'] is not None, - 'USE_INITIAL_STATE': lambda args: args['initial_state'] is not None, - 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None, -}) -@fla_cache_autotune( - 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', 'NB'], - **autotune_cache_kwargs, -) -@triton.jit -def causal_conv1d_fwd_kernel( - x, - y, - weight, - bias, - residual, - cu_seqlens, - initial_state, - chunk_indices, - B, - T, - stride_x_n, - stride_x_t, - stride_x_d, - 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, - HAS_RESIDUAL: tl.constexpr, - USE_INITIAL_STATE: tl.constexpr, - IS_VARLEN: 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 - else: - i_n = i_b - bos, eos = (i_b * T).to(tl.int64), (i_b * T + T).to(tl.int64) - 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] - b_w = tl.load(weight + o_d[:, None] * W + o_w, mask=m_d[:, None] & m_w, other=0).to(tl.float32) - - 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) - - if HAS_WEIGHT: - b_yi *= tl.sum(b_w * (o_w == (i_w + W - 1)), 1) - b_y += b_yi - - if HAS_BIAS: - b_y += tl.load(bias + o_d, mask=m_d).to(tl.float32) - - if ACTIVATION == 'swish' or ACTIVATION == 'silu': - 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({ - '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, -}) -@fla_cache_autotune( - 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 -def causal_conv1d_bwd_kernel( - x, - y, - weight, - initial_state, - dht, - dy, - dx, - dw, - db, - cu_seqlens, - 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, - USE_INITIAL_STATE: tl.constexpr, - USE_FINAL_STATE: tl.constexpr, - IS_VARLEN: 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 - else: - i_tg = i_b * tl.num_programs(1) + i_t - i_n = i_b - bos, eos = (i_b * T).to(tl.int64), (i_b * T + T).to(tl.int64) - p_x = x + tl.cast(i_b, tl.int64) * stride_x_n - - if IS_VARLEN: - p_dy = dy + bos * stride_dy_t - else: - 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) - - b_dx = tl.zeros((BT, BD), dtype=tl.float32) - if HAS_BIAS: - b_db = tl.zeros((BD,), dtype=tl.float32) - - 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 - if HAS_WEIGHT: - # [BT, BD] - b_wdy = b_wdy * tl.sum(b_w * (o_w == (W - i_w - 1)), 1) - # [BD] - b_dw = tl.sum(b_dy * b_x, 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 HAS_BIAS and i_w == 0: - 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 - if HAS_WEIGHT: - # [BT, BD] - b_wdy = b_wdy * tl.sum(b_w * (o_w == (W - i_w - 1)), 1) - # [BD] - b_dw = tl.sum(b_dy * b_x, 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 HAS_BIAS and i_w == 0: - 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) - - if HAS_BIAS and i_w == 0: - b_db += tl.sum(b_dy_shift, 0) - b_wdy = b_dy_shift if not HAS_WEIGHT else (b_dy_shift * tl.sum(b_w * (o_w == (W - i_w - 1)), 1)) - 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 = 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 - - if IS_VARLEN: - p_dx = dx + bos * stride_dx_t - 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({ - 'USE_INITIAL_STATE': lambda args: args['cache'] is not None, - 'HAS_WEIGHT': lambda args: args['weight'] is not None, - 'HAS_BIAS': lambda args: args['bias'] is not None, - 'HAS_RESIDUAL': lambda args: args['residual'] is not None, -}) -@fla_cache_autotune( - configs=[ - triton.Config({'BD': BD}, num_warps=num_warps) - for BD in [8, 16, 32, 64, 128, 256] - for num_warps in NUM_WARPS_AUTOTUNE - ], - key=['D', 'W'], - restore_value=['cache'], - **autotune_cache_kwargs, -) -@triton.jit -def causal_conv1d_update_kernel( - x, - cache, - residual, - 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, - BW: tl.constexpr, - ACTIVATION: tl.constexpr, - USE_INITIAL_STATE: tl.constexpr, - HAS_WEIGHT: tl.constexpr, - 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) - m_d = o_d < D - m_w = o_w < W - - # [BD] - b_x = tl.load(x + i_n * stride_x_n + o_d * stride_x_d, mask=m_d, other=0).to(tl.float32) - - 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) - - if HAS_WEIGHT: - b_w = tl.load(weight + o_d[:, None] * W + o_w, mask=m_d[:, None] & m_w, other=0) - b_y = tl.sum(b_cache * b_w, 1) - else: - b_y = tl.sum(b_cache, 1) - - if HAS_BIAS: - b_y += tl.load(bias + o_d, mask=m_d) - - if ACTIVATION == 'swish' or ACTIVATION == 'silu': - b_y = b_y * tl.sigmoid(b_y) - - 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 -def compute_dh0_kernel( - dy, - 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) - - -@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, -): - 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 - - 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 - - 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: - o_c = W - (BW - seq_len) + tl.arange(0, BW) - m_c = (o_c >= 0) & (o_c < W) - - p_init = initial_state + i_n * D*W + o_d[None, :] * W + o_c[:, None] - mask_init = m_d[None, :] & m_c[:, None] - - 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, - 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 - - -@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 diff --git a/build/torch-cuda/modules/conv/triton/ops.py b/build/torch-cuda/modules/conv/triton/ops.py deleted file mode 100644 index aee880417d0372500f9f67721baf34cc02bfeee1..0000000000000000000000000000000000000000 --- a/build/torch-cuda/modules/conv/triton/ops.py +++ /dev/null @@ -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 diff --git a/build/torch-cuda/modules/convolution.py b/build/torch-cuda/modules/convolution.py deleted file mode 100644 index b00367976440d67f5dd628b5eff6b20c8eee5c00..0000000000000000000000000000000000000000 --- a/build/torch-cuda/modules/convolution.py +++ /dev/null @@ -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', -] diff --git a/build/torch-cuda/modules/feature_map.py b/build/torch-cuda/modules/feature_map.py index a35aa11b6803e89469bb8afbd50fdacc74984577..1a39e92d18d517ce422a52fb4a330efb448c8a9c 100644 --- a/build/torch-cuda/modules/feature_map.py +++ b/build/torch-cuda/modules/feature_map.py @@ -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 = 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 diff --git a/build/torch-cuda/modules/fused_bitlinear.py b/build/torch-cuda/modules/fused_bitlinear.py index 4b36c859aaea1277089d643c3b2fb863b3fdd4d4..cc698d3ddbba7635d093cbff9eab4b6210d913ae 100644 --- a/build/torch-cuda/modules/fused_bitlinear.py +++ b/build/torch-cuda/modules/fused_bitlinear.py @@ -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. +# 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) # -# Code adapted from https://github.com/ridgerchu/matmulfreellm/ +# 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] +from ..modules.norm import RMSNorm +from ..modules.norm.layernorm_quant import LayerNormLinearQuantFunction, weight_quant -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 +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) - # Quantize and then de-quantize the tensor - y = (x * scale).round().clamp_(-128, 127) / scale + y = (x * scale).round().clamp_(min=-128, max=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, +def layer_norm_linear_quant( 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, + 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, - 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( +) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]: + return LayerNormLinearQuantFunction.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, + norm_weight: torch.Tensor | None, + norm_bias: torch.Tensor | None, linear_weight: torch.Tensor, - linear_bias: torch.Tensor, - residual: torch.Tensor = None, + linear_bias: torch.Tensor | None, + residual: torch.Tensor | None = None, eps: float = 1e-5, prenorm: bool = False, residual_in_fp32: bool = False, -): - return layer_norm_linear_quant_fn( +) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]: + return layer_norm_linear_quant( 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, +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): """ - 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. - """ + RMS-normalized linear layer with 8-bit activations and ternary weights. + + Quantization uses a straight-through estimator during training. + Efficient deployment requires specialized kernels. - def __init__( - self, - in_features: int, - out_features: int, - bias: bool = False, - norm_eps: float = 1e-8, - ): - """ - Initializes the BitLinear layer. + 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. + """ - 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 __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(in_features, eps=norm_eps, dtype=torch.float32) + 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): - """ - 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 + def forward(self, x: torch.Tensor) -> torch.Tensor: w = self.weight + x_norm = self.norm(x=x) - # 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) + # 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): - """ - 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, + """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, ) diff --git a/build/torch-cuda/modules/fused_cross_entropy.py b/build/torch-cuda/modules/fused_cross_entropy.py deleted file mode 100644 index c0b770e50014cadfabb101f42655cc4e383d5cc2..0000000000000000000000000000000000000000 --- a/build/torch-cuda/modules/fused_cross_entropy.py +++ /dev/null @@ -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 diff --git a/build/torch-cuda/modules/fused_cross_entropy/__init__.py b/build/torch-cuda/modules/fused_cross_entropy/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..de35634759c64e32870c61410db88e1880727fd2 --- /dev/null +++ b/build/torch-cuda/modules/fused_cross_entropy/__init__.py @@ -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 diff --git a/build/torch-cuda/modules/fused_cross_entropy/gluon.py b/build/torch-cuda/modules/fused_cross_entropy/gluon.py new file mode 100644 index 0000000000000000000000000000000000000000..354f6883cd65f46945187ab4a723c39124aebd81 --- /dev/null +++ b/build/torch-cuda/modules/fused_cross_entropy/gluon.py @@ -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 diff --git a/build/torch-cuda/modules/fused_cross_entropy/ops.py b/build/torch-cuda/modules/fused_cross_entropy/ops.py new file mode 100644 index 0000000000000000000000000000000000000000..e5dcaf8fbe85a69272a38417ce55573798c819bf --- /dev/null +++ b/build/torch-cuda/modules/fused_cross_entropy/ops.py @@ -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 diff --git a/build/torch-cuda/modules/fused_kl_div.py b/build/torch-cuda/modules/fused_kl_div.py deleted file mode 100644 index d157c4a52502985c0fb0817eae657f807e9b9cc8..0000000000000000000000000000000000000000 --- a/build/torch-cuda/modules/fused_kl_div.py +++ /dev/null @@ -1,372 +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 torch.nn as nn -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 - - -@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, -): - # https://github.com/triton-lang/triton/issues/1058 - # If N*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) - - 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) - td = td * exp(tm - m_new) + tl.sum(exp(b_tl - m_new)) - tm = m_new - - b_loss = 0. - # KL(y_true || y) = exp(y_true) * (log(y_true) - log(y)) - 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) - - # Normalize the loss by the number of elements if reduction is 'batchmean' - 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, -): - """ - 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 - - # 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(8, triton.cdiv(V, H)) - C = triton.next_power_of_2(triton.cdiv(N, NC)) - 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 - # we use fp32 for loss accumulator - loss = torch.zeros(N, dtype=torch.float32, device=device) - - for ic in range(NC): - start, end = ic * C, min((ic + 1) * C, N) - # [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. - 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, - ) - - # gradient of logits is computed in-place by the above triton kernel and is of shape: C x V - # thus dx[start: end] should be of shape: C x H - # additionally, since we are chunking the inputs, observe that the loss and gradients are calculated only - # on `n_non_ignore` tokens. However, the gradient of the input should be calculated for all tokens. - # 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: - dw = dw.to(weight) - 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 - - -class FusedKLDivLossFunction(torch.autograd.Function): - - @staticmethod - @input_guard - def forward( - ctx, - x: torch.Tensor, - target_x: torch.Tensor, - weight: torch.Tensor, - 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 - - @staticmethod - @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( - x: torch.Tensor, - target_x: torch.Tensor, - weight: torch.Tensor, - target_weight: torch.Tensor, - reduction: str = 'batchmean', - accumulate_grad_in_fp32: bool = True, -) -> 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( - "FusedKLDivLoss treats target_x and target_weight as a frozen teacher and does not compute " - "gradients for them. Detach target_x/target_weight before calling FusedKLDivLoss, or use " - "torch.nn.functional.kl_div if teacher gradients are required." - ) - return FusedKLDivLossFunction.apply( - x, - target_x, - weight, - 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__() - - assert reduction in ['batchmean'], f"reduction: {reduction} is not supported" - - 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, - target_x=target_x, - weight=weight, - target_weight=target_weight, - reduction=self.reduction, - accumulate_grad_in_fp32=self.accumulate_grad_in_fp32, - ) - return loss diff --git a/build/torch-cuda/modules/fused_kl_div/__init__.py b/build/torch-cuda/modules/fused_kl_div/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..4402eec1072b3ec05c2a3120d1a7c91ceca3771a --- /dev/null +++ b/build/torch-cuda/modules/fused_kl_div/__init__.py @@ -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'] + diff --git a/build/torch-xpu/modules/fused_kl_div.py b/build/torch-cuda/modules/fused_kl_div/ops.py similarity index 69% rename from build/torch-xpu/modules/fused_kl_div.py rename to build/torch-cuda/modules/fused_kl_div/ops.py index d157c4a52502985c0fb0817eae657f807e9b9cc8..b197efad8a47323656e482fd88d564fa68fe99f8 100644 --- a/build/torch-xpu/modules/fused_kl_div.py +++ b/build/torch-cuda/modules/fused_kl_div/ops.py @@ -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 +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 +# 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 +# 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 + # running maxima and exponential sums for student and teacher logits. 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. - """ +def elementwise_mul_kernel(x, g, N: tl.constexpr, B: tl.constexpr): - # 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) + 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('modules') -def fused_kl_div_forward( +@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 @@ -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 + 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): @@ -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: + if dw 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 dx is not None: + 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, - ) + 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: @@ -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 +@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=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, - ) + elementwise_mul_kernel[(triton.cdiv(N * H, B),)](x=grad, g=do, N=N*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, + grad_enabled: bool, ): - loss, dx, dw = fused_kl_div_forward( + 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 @@ -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 + 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( @@ -280,23 +234,25 @@ def fused_kl_div_loss( ) -> tuple[torch.Tensor, torch.Tensor]: """ Args: - x (`torch.Tensor`): + x (torch.Tensor): Tensor of shape `[batch_size * seq_len, hidden_size]`. - target_x (`torch.Tensor`): + target_x (torch.Tensor): Frozen teacher input tensor of shape `[batch_size * seq_len, hidden_size]`. Must not require gradients. - weight (`torch.Tensor`): + weight (torch.Tensor): Tensor of shape `[vocab_size, hidden_size]`. - target_weight (`torch.Tensor`): + target_weight (torch.Tensor): Frozen teacher weight tensor of shape `[vocab_size, hidden_size]`. Must not require gradients. - reduction (`str`): + reduction (str, Optional): 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`. + 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 + loss (torch.Tensor): + Scalar KL divergence 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, + torch.is_grad_enabled(), ) class FusedKLDivLoss(nn.Module): - def __init__( - self, - reduction: str = 'batchmean', - accumulate_grad_in_fp32: bool = True, - ): + def __init__(self, reduction: str = 'batchmean', accumulate_grad_in_fp32: bool = True): """ Args: - reduction (`str`): + reduction (str, Optional): 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`. + 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. + 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, - ): + def forward(self, x: torch.Tensor, target_x: torch.Tensor, weight: torch.Tensor, target_weight: torch.Tensor): """ Args: - x (`torch.Tensor`): + x (torch.Tensor): Tensor of shape `[batch_size * seq_len, hidden_size]`. - target_x (`torch.Tensor`): + target_x (torch.Tensor): Frozen teacher input tensor of shape `[batch_size * seq_len, hidden_size]`. Must not require gradients. - weight (`torch.Tensor`): + weight (torch.Tensor): Tensor of shape `[vocab_size, hidden_size]`. - target_weight (`torch.Tensor`): + target_weight (torch.Tensor): Frozen teacher weight tensor of shape `[vocab_size, hidden_size]`. Must not require gradients. + Returns: - loss + loss (torch.Tensor): + Scalar KL divergence loss. """ loss = fused_kl_div_loss( x=x, diff --git a/build/torch-cuda/modules/fused_linear_cross_entropy.py b/build/torch-cuda/modules/fused_linear_cross_entropy.py deleted file mode 100644 index 8e9128a63e1ae7ae6aebc16f533edb53ef711f01..0000000000000000000000000000000000000000 --- a/build/torch-cuda/modules/fused_linear_cross_entropy.py +++ /dev/null @@ -1,767 +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 - -# 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 -import torch.nn.functional as F -import triton -import triton.language as tl -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 - - -@triton.heuristics({ - 'HAS_SCALE': lambda args: args['scale'] is not None, - 'HAS_SOFTCAPPING': lambda args: args['softcapping'] is not None, -}) -@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) - - # [N, H] - 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 - - # [V, H] - dw = torch.zeros_like(weight, device=device, dtype=grad_dtype) if weight is not None else None - # [V] - db = torch.zeros_like(bias, device=device, dtype=bias_grad_dtype) if bias is not None else None - # [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: - db = db.to(bias) - 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 - @input_guard - def forward( - ctx, - 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, - db.detach() if bias is not None else None, - ) - return loss - - @staticmethod - @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, - target, - weight, - bias, - ignore_index, - label_smoothing, - logit_scale, - logit_softcapping, - num_chunks, - reduction, - use_l2warp, - l2_penalty_factor, - accumulate_grad_in_fp32, - ) - - -class FusedLinearCrossEntropyLoss(nn.Module): - - def __init__( - self, - 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__() - - assert reduction in ["mean", "sum"], f"reduction: {reduction} is not supported" - - self.ignore_index = ignore_index - self.label_smoothing = label_smoothing - self.logit_scale = logit_scale - self.logit_softcapping = logit_softcapping - self.num_chunks = num_chunks - self.reduction = reduction - 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, - label_smoothing=self.label_smoothing, - logit_scale=self.logit_scale, - logit_softcapping=self.logit_softcapping, - num_chunks=self.num_chunks, - reduction=self.reduction, - 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),) - self.use_local_output = use_local_output - - @staticmethod - def _prepare_input_fn(sequence_sharding, mod, inputs, device_mesh): - x, target, weight, bias = inputs - - if not isinstance(x, DTensor): - # assume the input passed in already sharded on the sequence dim and create the DTensor - x = DTensor.from_local(x, device_mesh, sequence_sharding) - if x.placements != sequence_sharding: - x = x.redistribute(placements=sequence_sharding, async_op=True) - if not isinstance(target, DTensor): - target = DTensor.from_local(target, device_mesh, [Replicate()]) - if target.placements != sequence_sharding: - target = target.redistribute(placements=sequence_sharding, async_op=True) - - if not isinstance(weight, DTensor): - weight = DTensor.from_local(weight, device_mesh, [Replicate()]) - if weight.placements != [Replicate()]: - # we replicate the weight/bias in FLCE - weight = weight.redistribute(placements=[Replicate()], async_op=True) - - if bias is not None and not isinstance(bias, DTensor): - bias = DTensor.from_local(bias, device_mesh, [Replicate()]) - if bias is not None and bias.placements != [Replicate()]: - bias = bias.redistribute(placements=[Replicate()], async_op=True) - - return x.to_local(), target.to_local(), weight.to_local(), bias.to_local() if bias is not None else bias - - @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( - module, - device_mesh, - partition_fn=None, - input_fn=partial(self._prepare_input_fn, self.sequence_sharding), - output_fn=partial(self._prepare_output_fn, self.use_local_output), - ) diff --git a/build/torch-cuda/modules/fused_linear_cross_entropy/__init__.py b/build/torch-cuda/modules/fused_linear_cross_entropy/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..cd2b82f84073c7b0e7d0dad9eb412afc6d5cf03d --- /dev/null +++ b/build/torch-cuda/modules/fused_linear_cross_entropy/__init__.py @@ -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'] + diff --git a/build/torch-cuda/modules/fused_linear_cross_entropy/ops.py b/build/torch-cuda/modules/fused_linear_cross_entropy/ops.py new file mode 100644 index 0000000000000000000000000000000000000000..f3897dbd7925de81ddbdb53a6e7fa78b7ba3b298 --- /dev/null +++ b/build/torch-cuda/modules/fused_linear_cross_entropy/ops.py @@ -0,0 +1,535 @@ +# 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 + +# 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 +import torch.nn.functional as F +import triton +import triton.language as tl +from torch.distributed import DeviceMesh +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.heuristics({ + 'HAS_SCALE': lambda args: args['scale'] is not None, + 'HAS_SOFTCAPPING': lambda args: args['softcapping'] is not None, +}) +@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, H] + 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 + + # [V, H] + dw = torch.zeros_like(weight, device=device, dtype=grad_dtype) if weight is not None else None + # [V] + db = torch.zeros_like(bias, device=device, dtype=bias_grad_dtype) if bias is not None else None + # [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: + db = db.to(bias) + 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 + @input_guard + def forward( + ctx, + 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, + db.detach() if bias is not None else None, + ) + return loss + + @staticmethod + @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, + target, + weight, + bias, + ignore_index, + label_smoothing, + logit_scale, + logit_softcapping, + num_chunks, + reduction, + use_l2warp, + l2_penalty_factor, + accumulate_grad_in_fp32, + process_group, + ) + + +class FusedLinearCrossEntropyLoss(nn.Module): + + def __init__( + self, + 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__() + + assert reduction in ["mean", "sum"], f"reduction: {reduction} is not supported" + + self.ignore_index = ignore_index + self.label_smoothing = label_smoothing + self.logit_scale = logit_scale + self.logit_softcapping = logit_softcapping + self.num_chunks = num_chunks + self.reduction = reduction + 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, + label_smoothing=self.label_smoothing, + logit_scale=self.logit_scale, + logit_softcapping=self.logit_softcapping, + num_chunks=self.num_chunks, + reduction=self.reduction, + 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),) + self.use_local_output = use_local_output + + @staticmethod + def _prepare_input_fn(sequence_sharding, mod, inputs, device_mesh): + x, target, weight, bias = inputs + + if not isinstance(x, DTensor): + # assume the input passed in already sharded on the sequence dim and create the DTensor + x = DTensor.from_local(x, device_mesh, sequence_sharding) + if x.placements != sequence_sharding: + x = x.redistribute(placements=sequence_sharding, async_op=True) + if not isinstance(target, DTensor): + target = DTensor.from_local(target, device_mesh, [Replicate()]) + if target.placements != sequence_sharding: + target = target.redistribute(placements=sequence_sharding, async_op=True) + + if not isinstance(weight, DTensor): + weight = DTensor.from_local(weight, device_mesh, [Replicate()]) + if weight.placements != [Replicate()]: + # we replicate the weight/bias in FLCE + weight = weight.redistribute(placements=[Replicate()], async_op=True) + + if bias is not None and not isinstance(bias, DTensor): + bias = DTensor.from_local(bias, device_mesh, [Replicate()]) + if bias is not None and bias.placements != [Replicate()]: + bias = bias.redistribute(placements=[Replicate()], async_op=True) + + return x.to_local(), target.to_local(), weight.to_local(), bias.to_local() if bias is not None else bias + + @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( + module, + device_mesh, + partition_fn=None, + input_fn=partial(self._prepare_input_fn, self.sequence_sharding), + output_fn=partial(self._prepare_output_fn, self.use_local_output), + ) diff --git a/build/torch-cuda/modules/fused_norm_gate.py b/build/torch-cuda/modules/fused_norm_gate.py deleted file mode 100644 index 20af745372b58ff6834b8ff62f0aee8ecedf91f7..0000000000000000000000000000000000000000 --- a/build/torch-cuda/modules/fused_norm_gate.py +++ /dev/null @@ -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, - ) diff --git a/build/torch-cuda/modules/grpo.py b/build/torch-cuda/modules/grpo.py deleted file mode 100644 index dd75fb271768be2377b12006e818cd4aec7d4d55..0000000000000000000000000000000000000000 --- a/build/torch-cuda/modules/grpo.py +++ /dev/null @@ -1,421 +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 - -# 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( - configs=[ - 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, -) -@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_idx = tl.program_id(0) - - 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: - # store 0 - tl.store(loss_ptr, 0.0) - if save_kl: - tl.store(loss_ptr+M, 0.0) - - -@triton.autotune( - configs=[ - triton.Config({}, num_warps=NUM_WARPS, num_stages=NUM_STAGES) - for NUM_WARPS in [32] - for NUM_STAGES in [4] - ], - key=['B', 'N'], - **autotune_cache_kwargs, -) -@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_idx = tl.program_id(0) # B*L - 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) - # Need for in-place grad. - 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 - 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 GrpoLoss(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 - - 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) - - grpo_fwd_kernel[(M,)]( - 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, - ) - ctx.beta = beta - ctx.save_for_backward(lse, logits, input_ids, advantages, completion_mask) - ctx.ref_logp = ref_logp - ctx.inplace = inplace - return loss - - @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 - L = L_ADD_1 - 1 - M = B * L - - input_ids_start_index = input_ids.size(1) - L - - # B, L_ADD_1, N - dlogits = logits if inplace else torch.empty_like(logits) - BN = min(65536, triton.next_power_of_2(N)) - - grpo_bwd_kernel[(M,)]( - 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=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 - else: - return out.chunk(2, axis=0) - - -def grpo_loss_torch(logits, ref_logp, input_ids, advantages, beta=0.1, completion_mask=None, save_kl=False): - def get_log_probs(logits, input_ids): - per_token_logps = [] - for logits_row, input_ids_row in zip(logits, input_ids[:, -logits.size(1):], strict=False): - log_probs = logits_row.log_softmax(dim=-1) - token_log_prob = torch.gather(log_probs, dim=1, index=input_ids_row.unsqueeze(1)).squeeze(1) - per_token_logps.append(token_log_prob) - return torch.stack(per_token_logps) - - logits = logits[:, :-1] - per_token_logps = get_log_probs(logits, input_ids) - ref_per_token_logps = ref_logp - per_token_kl = torch.exp(ref_per_token_logps - per_token_logps) - (ref_per_token_logps - per_token_logps) - 1 - - per_token_loss = torch.exp(per_token_logps - per_token_logps.detach()) * advantages.unsqueeze(1) - per_token_loss = -(per_token_loss - beta * per_token_kl) - if completion_mask is not None: - per_token_loss *= completion_mask - if save_kl: - per_token_kl *= completion_mask - 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, - old_logps: torch.Tensor, - pad_mask: torch.Tensor, - logits_to_keep: int, - rewards: torch.Tensor, - 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 diff --git a/build/torch-xpu/modules/conv/cp/__init__.py b/build/torch-cuda/modules/grpo/__init__.py similarity index 69% rename from build/torch-xpu/modules/conv/cp/__init__.py rename to build/torch-cuda/modules/grpo/__init__.py index 97e899f587bcd72246487b736d919e6b7ee8b28a..85b5190345ed3163a21721d93abc4910ca18729f 100644 --- a/build/torch-xpu/modules/conv/cp/__init__.py +++ b/build/torch-cuda/modules/grpo/__init__.py @@ -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 +from ...modules.grpo.ops import fused_grpo_loss, grpo_loss_with_old_logps + +__all__ = ['fused_grpo_loss', 'grpo_loss_with_old_logps'] -__all__ = [ - 'CausalConv1dFunctionCP', - 'causal_conv1d_cp', -] diff --git a/build/torch-xpu/modules/grpo.py b/build/torch-cuda/modules/grpo/ops.py similarity index 58% rename from build/torch-xpu/modules/grpo.py rename to build/torch-cuda/modules/grpo/ops.py index dd75fb271768be2377b12006e818cd4aec7d4d55..1dbdfa6c26a9c74763cc0c5ed417f9fe299087ac 100644 --- a/build/torch-xpu/modules/grpo.py +++ b/build/torch-cuda/modules/grpo/ops.py @@ -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 -""" +# loss reference: https://github.com/huggingface/trl/blob/main/trl/trainer/grpo_trainer.py 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 +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( @@ -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] + for NUM_STAGES in NUM_STAGES_AUTOTUNE ], 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 + 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, + 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, + B=B, + N=N, + L=L, start_idx=input_ids_start_index, + BLOCK_SIZE=BN, ) - # The last token in the completion is not used in the loss computation - # and therefore its gradient should be set to 0 + # 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('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) +@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: - 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) - ''' + 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 @@ -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) +# 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, @@ -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. + """Compute clipped GRPO loss using current, reference and old-policy log probabilities. 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. + 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: The computed GRPO loss. + 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 = (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 + completion_mask = completion_mask & pad_mask - # 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 + 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() + loss = token_loss.sum() / completion_mask.sum() return loss diff --git a/build/torch-cuda/modules/l2norm.py b/build/torch-cuda/modules/l2norm.py deleted file mode 100644 index a840acaf4d2555d87b2958dd1dbd5f1d30362f96..0000000000000000000000000000000000000000 --- a/build/torch-cuda/modules/l2norm.py +++ /dev/null @@ -1,288 +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 torch.nn as nn -import triton -import triton.language as tl - -from ..ops.utils.cache import fla_cache_autotune -from ..utils import IS_AMD, autotune_cache_kwargs, input_guard - -BT_LIST = [8, 16, 32, 64, 128] -NUM_WARPS_AUTOTUNE = [1, 2, 4, 8, 16] if IS_AMD else [1, 2, 4, 8, 16, 32] - - -@triton.autotune( - configs=[triton.Config({}, num_warps=num_warps) for num_warps in NUM_WARPS_AUTOTUNE], - key=["D"], - **autotune_cache_kwargs, -) -@triton.jit -def l2norm_fwd_kernel1( - x, - y, - rstd, - eps, - D, - BD: tl.constexpr, -): - i_t = tl.program_id(0) - x += i_t * D - y += i_t * D - # Compute mean and variance - cols = tl.arange(0, BD) - mask = cols < D - - b_x = tl.load(x + cols, mask=mask, other=0.0).to(tl.float32) - b_rstd = 1 / tl.sqrt(tl.sum(b_x * b_x) + eps) - b_y = b_x * b_rstd - tl.store(y + cols, b_y, mask=mask) - tl.store(rstd + i_t, b_rstd) - - -@triton.autotune( - configs=[triton.Config({}, num_warps=num_warps) for num_warps in NUM_WARPS_AUTOTUNE], - key=["D"], - **autotune_cache_kwargs, -) -@triton.jit -def l2norm_bwd_kernel1( - y, - rstd, - dy, - dx, - eps, - D, - BD: tl.constexpr, -): - i_t = tl.program_id(0) - y += i_t * D - dx += i_t * D - dy += i_t * D - - cols = tl.arange(0, BD) - mask = cols < D - b_y = tl.load(y + cols, mask=mask, other=0.0).to(tl.float32) - b_rstd = tl.load(rstd + i_t).to(tl.float32) - b_dy = tl.load(dy + cols, mask=mask, other=0.0).to(tl.float32) - b_dx = b_dy * b_rstd - tl.sum(b_dy * b_y) * b_y * b_rstd - tl.store(dx + cols, b_dx, mask=mask) - - -@fla_cache_autotune( - configs=[triton.Config({"BT": BT}, num_warps=num_warps) for num_warps in [1, 2, 4, 8, 16] for BT in BT_LIST], - key=["D", "NB"], - **autotune_cache_kwargs, -) -@triton.jit(do_not_specialize=["T"]) -def l2norm_fwd_kernel( - x, - y, - rstd, - eps, - T, - D: tl.constexpr, - BD: tl.constexpr, - NB: tl.constexpr, - BT: tl.constexpr, -): - i_t = tl.program_id(0) - p_x = tl.make_block_ptr(x, (T, D), (D, 1), (i_t * BT, 0), (BT, BD), (1, 0)) - p_y = tl.make_block_ptr(y, (T, D), (D, 1), (i_t * BT, 0), (BT, BD), (1, 0)) - p_rstd = tl.make_block_ptr(rstd, (T,), (1,), (i_t * BT,), (BT,), (0,)) - - b_x = tl.load(p_x, boundary_check=(0, 1)).to(tl.float32) - b_rstd = 1 / tl.sqrt(tl.sum(b_x * b_x, 1) + eps) - b_y = b_x * b_rstd[:, None] - - tl.store(p_y, b_y.to(p_y.dtype.element_ty), boundary_check=(0, 1)) - tl.store(p_rstd, b_rstd.to(p_rstd.dtype.element_ty), boundary_check=(0,)) - - -@fla_cache_autotune( - configs=[triton.Config({"BT": BT}, num_warps=num_warps) for num_warps in [1, 2, 4, 8, 16] for BT in BT_LIST], - key=["D", "NB"], - **autotune_cache_kwargs, -) -@triton.jit(do_not_specialize=["T"]) -def l2norm_bwd_kernel( - y, - rstd, - dy, - dx, - eps, - T, - D: tl.constexpr, - BD: tl.constexpr, - NB: tl.constexpr, - BT: tl.constexpr, -): - i_t = tl.program_id(0) - p_y = tl.make_block_ptr(y, (T, D), (D, 1), (i_t * BT, 0), (BT, BD), (1, 0)) - p_rstd = tl.make_block_ptr(rstd, (T,), (1,), (i_t * BT,), (BT,), (0,)) - p_dy = tl.make_block_ptr(dy, (T, D), (D, 1), (i_t * BT, 0), (BT, BD), (1, 0)) - p_dx = tl.make_block_ptr(dx, (T, D), (D, 1), (i_t * BT, 0), (BT, BD), (1, 0)) - - b_y = tl.load(p_y, boundary_check=(0, 1)).to(tl.float32) - b_rstd = tl.load(p_rstd, boundary_check=(0,)).to(tl.float32) - b_dy = tl.load(p_dy, boundary_check=(0, 1)).to(tl.float32) - b_dx = b_dy * b_rstd[:, None] - tl.sum(b_dy * b_y, 1)[:, None] * b_y * b_rstd[:, None] - tl.store(p_dx, b_dx.to(p_dx.dtype.element_ty), boundary_check=(0, 1)) - - -def l2norm_fwd( - x: torch.Tensor, - eps: float = 1e-6, - output_dtype: torch.dtype | None = None, -): - x_shape_og = x.shape - x = x.view(-1, x.shape[-1]) - # allocate output - if output_dtype is None: - y = torch.empty_like(x) - else: - y = torch.empty_like(x, dtype=output_dtype) - assert y.stride(-1) == 1 - T, D = x.shape[0], x.shape[-1] - # 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 doesn't support feature dim >= 64KB.") - - rstd = torch.empty((T,), dtype=torch.float32, device=x.device) - 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"]),) - - l2norm_fwd_kernel[grid]( - x=x, - y=y, - rstd=rstd, - eps=eps, - T=T, - D=D, - BD=BD, - NB=NB, - ) - else: - l2norm_fwd_kernel1[(T,)]( - x=x, - y=y, - rstd=rstd, - eps=eps, - D=D, - BD=BD, - ) - return y.view(x_shape_og), rstd.view(x_shape_og[:-1]) - - -def l2norm_bwd( - y: torch.Tensor, - rstd: torch.Tensor, - dy: torch.Tensor, - eps: float = 1e-6, -): - y_shape_og = y.shape - y = y.view(-1, dy.shape[-1]) - dy = dy.view(-1, dy.shape[-1]) - assert dy.shape == y.shape - # allocate output - dx = torch.empty_like(y) - T, D = y.shape[0], y.shape[-1] - # Less than 64KB per feature: enqueue fused kernel - MAX_FUSED_SIZE = 65536 // y.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: - # 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"]),) - - l2norm_bwd_kernel[grid]( - y=y, - rstd=rstd, - dy=dy, - dx=dx, - eps=eps, - T=T, - D=D, - BD=BD, - NB=NB, - ) - else: - l2norm_bwd_kernel1[(T,)]( - y=y, - rstd=rstd, - dy=dy, - dx=dx, - eps=eps, - D=D, - BD=BD, - ) - - return dx.view(y_shape_og) - - -class L2NormFunction(torch.autograd.Function): - @staticmethod - @input_guard - def forward( - ctx, - x, - eps=1e-6, - output_dtype=None, - ): - y, rstd = l2norm_fwd(x, eps, output_dtype) - ctx.eps = eps - ctx.x_dtype = x.dtype - ctx.save_for_backward(y, rstd) - return y - - @staticmethod - @input_guard - def backward(ctx, dy): - y, rstd = ctx.saved_tensors - dx = l2norm_bwd(y, rstd, dy, ctx.eps) - return dx, None, None - - -def l2norm( - x: torch.Tensor, - eps: float = 1e-6, - output_dtype: torch.dtype | None = None, -) -> torch.Tensor: - return L2NormFunction.apply(x, eps, output_dtype) - - -l2_norm = l2norm - - -class L2Norm(nn.Module): - def __init__( - self, - eps: float = 1e-6, - output_dtype: torch.dtype | None = None, - ): - super().__init__() - self.eps = eps - self.output_dtype = output_dtype - - def forward(self, x: torch.Tensor) -> torch.Tensor: - return l2norm(x, self.eps, self.output_dtype) diff --git a/build/torch-cuda/modules/l2warp.py b/build/torch-cuda/modules/l2warp.py index 3d385823877dac969e46b20d22b17a033f3654a7..1353d10627308a35f3c2f64b0127518661ec52ab 100644 --- a/build/torch-cuda/modules/l2warp.py +++ b/build/torch-cuda/modules/l2warp.py @@ -39,13 +39,12 @@ class L2Wrap(torch.autograd.Function): return loss @staticmethod - def backward(ctx, grad_output: torch.Tensor): + def backward(ctx, dloss: 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 + 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 diff --git a/build/torch-cuda/modules/layernorm_gated.py b/build/torch-cuda/modules/layernorm_gated.py deleted file mode 100644 index ff0eac089649d8a784aed5f33b642d116bf72ffd..0000000000000000000000000000000000000000 --- a/build/torch-cuda/modules/layernorm_gated.py +++ /dev/null @@ -1,535 +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 -# -# Copyright (c) 2024, Tri Dao. -# -# Based on the Triton LayerNorm tutorial: https://triton-lang.org/main/getting-started/tutorials/05-layer-norm.html -# For the backward pass, we keep weight_grad and bias_grad in registers and accumulate. -# This backward pass is faster for dimensions up to 8k, but after that it's much slower due to register spilling. -# The models we train have hidden dim up to 8k anyway (e.g. Llama 70B), so this is fine. - -import math - -import torch -import torch.nn as nn -import torch.nn.functional as F -import triton -import triton.language as tl -from einops import rearrange - -from ..utils import get_multiprocessor_count, input_guard - - -def rms_norm_ref(x, weight, bias, z=None, eps=1e-6, group_size=None, norm_before_gate=True, upcast=True): - dtype = x.dtype - weight = weight.float() - bias = bias.float() if bias is not None else None - if upcast: - x = x.float() - z = z.float() if z is not None else z - if z is not None and not norm_before_gate: - x = x * F.silu(z) - if group_size is None: - rstd = 1 / torch.sqrt((x.square()).mean(dim=-1, keepdim=True) + eps) - out = (x * rstd * weight) + bias if bias is not None else (x * rstd * weight) - else: - x_group = rearrange(x, "... (g d) -> ... g d", d=group_size) - rstd = 1 / torch.sqrt((x_group.square()).mean(dim=-1, keepdim=True) + eps) - out = rearrange(x_group * rstd, "... g d -> ... (g d)") * weight - if bias is not None: - out = out + bias - if z is not None and norm_before_gate: - out *= F.silu(z) - return out.to(dtype) - - -@triton.heuristics({ - "HAS_BIAS": lambda args: args["B"] is not None, - "HAS_Z": lambda args: args["Z"] is not None, -}) -@triton.jit -def layer_norm_fwd_kernel( - X, # pointer to the input - Y, # pointer to the output - W, # pointer to the weights - B, # pointer to the biases - Z, # pointer to the other branch - 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_z_row, - M, # number of rows in X - N, # number of columns in X - eps, # epsilon to avoid division by zero - BLOCK_N: tl.constexpr, - HAS_BIAS: tl.constexpr, - HAS_Z: tl.constexpr, - NORM_BEFORE_GATE: tl.constexpr, - IS_RMS_NORM: tl.constexpr, -): - # Map the program id to the row of X and Y it should compute. - row = tl.program_id(0) - group = tl.program_id(1) - X += row * stride_x_row + group * N - Y += row * stride_y_row + group * N - if HAS_Z: - Z += row * stride_z_row + group * N - if not IS_RMS_NORM: - Mean += group * M - Rstd += group * M - W += group * N - if HAS_BIAS: - B += group * N - # Compute mean and variance - cols = tl.arange(0, BLOCK_N) - x = tl.load(X + cols, mask=cols < N, other=0.).to(tl.float32) - if HAS_Z and not NORM_BEFORE_GATE: - z = tl.load(Z + cols, mask=cols < N).to(tl.float32) - x *= z * tl.sigmoid(z) - 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.) - var = tl.sum(xbar * xbar, axis=0) / N - else: - xbar = tl.where(cols < N, x, 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 - 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 + b if HAS_BIAS else x_hat * w - if HAS_Z and NORM_BEFORE_GATE: - z = tl.load(Z + cols, mask=mask).to(tl.float32) - y *= z * tl.sigmoid(z) - # Write output - tl.store(Y + cols, y, mask=mask) - - -def layer_norm_fwd( - x: torch.Tensor, - weight: torch.Tensor, - bias: torch.Tensor, - eps: float, - z: torch.Tensor = None, - out: torch.Tensor = None, - group_size: int = None, - norm_before_gate: bool = True, - is_rms_norm: bool = False, -): - M, N = x.shape - if group_size is None: - group_size = N - assert N % group_size == 0 - ngroups = N // group_size - assert x.stride(-1) == 1 - if z is not None: - assert z.stride(-1) == 1 - assert z.shape == (M, N) - assert weight.shape == (N,) - assert weight.stride(-1) == 1 - if bias is not None: - assert bias.stride(-1) == 1 - assert bias.shape == (N,) - # allocate output - if out is not None: - assert out.shape == x.shape - else: - out = torch.empty_like(x) - assert out.stride(-1) == 1 - mean = torch.empty((ngroups * M, ), dtype=torch.float32, device=x.device) if not is_rms_norm else None - rstd = torch.empty((ngroups * 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(group_size)) - if group_size > BLOCK_N: - raise RuntimeError("This layer norm doesn't support feature dim >= 64KB.") - # heuristics for number of warps - num_warps = min(max(BLOCK_N // 256, 1), 8) - grid = (M, ngroups) - layer_norm_fwd_kernel[grid]( - x, - out, - weight, - bias, - z, - mean, - rstd, - x.stride(0), - out.stride(0), - z.stride(0) if z is not None else 0, - M, - group_size, - eps, - BLOCK_N=BLOCK_N, - NORM_BEFORE_GATE=norm_before_gate, - IS_RMS_NORM=is_rms_norm, - num_warps=num_warps, - ) - return out, mean, rstd - - -@triton.heuristics({ - "HAS_BIAS": lambda args: args["B"] is not None, - "HAS_Z": lambda args: args["Z"] is not None, - "RECOMPUTE_OUTPUT": lambda args: args["Y"] is not None, -}) -@triton.jit -def layer_norm_bwd_kernel( - X, # pointer to the input - W, # pointer to the weights - B, # pointer to the biases - Z, # pointer to the other branch - 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 - DZ, # pointer to the other branch - 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_z_row, - stride_y_row, - stride_dy_row, - stride_dx_row, - stride_dz_row, - stride_dw_row, - stride_db_row, - M, # number of rows in X - N, # number of columns in X - eps, # epsilon to avoid division by zero - rows_per_program, - NORM_BEFORE_GATE: tl.constexpr, - IS_RMS_NORM: tl.constexpr, - HAS_BIAS: tl.constexpr, - HAS_Z: tl.constexpr, - RECOMPUTE_OUTPUT: tl.constexpr, - BLOCK_N: tl.constexpr, -): - # Map the program id to the elements of X, DX, and DY it should compute. - row_block_id = tl.program_id(0) - group = tl.program_id(1) - row_start = row_block_id * rows_per_program - cols = tl.arange(0, BLOCK_N) - mask = cols < N - X += row_start * stride_x_row + group * N - if HAS_Z: - Z += row_start * stride_z_row + group * N - DZ += row_start * stride_dz_row + group * N - DY += row_start * stride_dy_row + group * N - DX += row_start * stride_dx_row + group * N - if RECOMPUTE_OUTPUT: - Y += row_start * stride_y_row + group * N - if not IS_RMS_NORM: - Mean += group * M - Rstd += group * M - W += group * N - w = tl.load(W + cols, mask=mask).to(tl.float32) - if (RECOMPUTE_OUTPUT or HAS_Z) and HAS_BIAS: - B += group * N - b = tl.load(B + cols, mask=mask, other=0.).to(tl.float32) - dw = tl.zeros((BLOCK_N,), dtype=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) - if HAS_Z and not NORM_BEFORE_GATE: - z = tl.load(Z + cols, mask=mask, other=0.).to(tl.float32) - x_og = x - x = x_og * z * tl.sigmoid(z) - 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.) - if HAS_Z and NORM_BEFORE_GATE: - z = tl.load(Z + cols, mask=mask, other=0.).to(tl.float32) - z_sigmoid = tl.sigmoid(z) - y = xhat * w + b if HAS_BIAS else xhat * w - if RECOMPUTE_OUTPUT: - tl.store(Y + cols, y * z * z_sigmoid, mask=mask) - dz = dy * y * z_sigmoid * (1 + z * (1 - z_sigmoid)) - tl.store(DZ + cols, dz, mask=mask) - dy *= z * z_sigmoid - else: - if RECOMPUTE_OUTPUT: - y = xhat * w + b if HAS_BIAS else xhat * w - tl.store(Y + cols, y, mask=mask) - wdy = w * dy - c1 = tl.sum(xhat * wdy, axis=0) / N - if not IS_RMS_NORM: - c2 = tl.sum(wdy, axis=0) / N - dx = (wdy - (xhat * c1 + c2)) * rstd - else: - dx = (wdy - xhat * c1) * rstd - dw += dy * xhat - if HAS_BIAS: - db += dy - if HAS_Z and not NORM_BEFORE_GATE: - z_sigmoid = tl.sigmoid(z) - dz = dx * x_og * z_sigmoid * (1 + z * (1 - z_sigmoid)) - tl.store(DZ + cols, dz, mask=mask) - dx *= z * z_sigmoid - # Write dx - tl.store(DX + cols, dx, mask=mask) - - X += stride_x_row - if HAS_Z: - Z += stride_z_row - DZ += stride_dz_row - if RECOMPUTE_OUTPUT: - Y += stride_y_row - DY += stride_dy_row - DX += stride_dx_row - tl.store(DW + row_block_id * stride_dw_row + group * N + cols, dw, mask=mask) - if HAS_BIAS: - tl.store(DB + row_block_id * stride_db_row + group * 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, - z: torch.Tensor = None, - group_size: int = None, - norm_before_gate: bool = True, - is_rms_norm: bool = False, - recompute_output: bool = False, - dz: torch.Tensor = None, - out: torch.Tensor = None, -): - M, N = x.shape - if group_size is None: - group_size = N - assert N % group_size == 0 - ngroups = N // group_size - assert x.stride(-1) == 1 - assert dy.stride(-1) == 1 - assert dy.shape == (M, N) - if z is not None: - assert z.stride(-1) == 1 - assert z.shape == (M, N) - assert weight.shape == (N,) - assert weight.stride(-1) == 1 - if bias is not None: - assert bias.stride(-1) == 1 - assert bias.shape == (N,) - # allocate output - dx = torch.empty_like(x) - if dz is not None: - assert z is not None - assert dz.shape == z.shape - assert dz.stride(-1) == 1 - else: - dz = torch.empty_like(z) if z is not None else None - if recompute_output: - if out is None: - out = torch.empty_like(x) - assert out.shape == x.shape - - # 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(group_size)) - if group_size > BLOCK_N: - raise RuntimeError("This layer norm doesn't support feature dim >= 64KB.") - # heuristics for number of warps - num_warps = min(max(BLOCK_N // 256, 1), 8) - sm_count = get_multiprocessor_count(x.device.index) - # If group size is small (e.g., 64), we're only using 1 warp. So having just 108 programs - # would limit the occupancy. - nrow_groups = math.ceil(sm_count * math.ceil(4 / num_warps) / ngroups) - _dw = torch.empty((nrow_groups, N), dtype=torch.float32, device=weight.device) - _db = torch.empty((nrow_groups, N), dtype=torch.float32, device=bias.device) if bias is not None else None - rows_per_program = math.ceil(M / nrow_groups) - grid = (nrow_groups, ngroups) - layer_norm_bwd_kernel[grid]( - x, - weight, - bias, - z, - out if recompute_output else None, - dy, - dx, - _dw, - _db, - dz, - mean, - rstd, - x.stride(0), - z.stride(0) if z is not None else 0, - 0 if not recompute_output else out.stride(0), - dy.stride(0), - dx.stride(0), - dz.stride(0) if dz is not None else 0, - _dw.stride(0), - _db.stride(0) if _db is not None else 0, - M, group_size, eps, - rows_per_program, - BLOCK_N=BLOCK_N, - NORM_BEFORE_GATE=norm_before_gate, - IS_RMS_NORM=is_rms_norm, - num_warps=num_warps, - ) - dw = _dw.sum(0).to(weight.dtype) - db = _db.sum(0).to(bias.dtype) if bias is not None else None - return (dx, dw, db, dz) if not recompute_output else (dx, dw, db, dz, out) - - -class LayerNormFn(torch.autograd.Function): - - @input_guard - @staticmethod - def forward(ctx, x, weight, bias, z=None, eps=1e-6, group_size=None, norm_before_gate=True, - is_rms_norm=False): - """If z is not None, we do norm(x) * silu(z) if norm_before_gate, else norm(x * silu(z)) - """ - - x_shape_og = x.shape - # reshape input data into 2D tensor - x = x.reshape(-1, x.shape[-1]) - if x.stride(-1) != 1: - x = x.contiguous() - if z is not None: - assert z.shape == x_shape_og - z = z.reshape(-1, z.shape[-1]) - if z.stride(-1) != 1: - z = z.contiguous() - weight = weight.contiguous() - if bias is not None: - bias = bias.contiguous() - y, mean, rstd = layer_norm_fwd( - x, - weight, - bias, - eps, - z=z, - group_size=group_size, - norm_before_gate=norm_before_gate, - is_rms_norm=is_rms_norm, - ) - ctx.save_for_backward(x, weight, bias, mean, rstd, z) - ctx.x_shape_og = x_shape_og - ctx.eps = eps - ctx.group_size = group_size - ctx.norm_before_gate = norm_before_gate - ctx.is_rms_norm = is_rms_norm - return y.reshape(x_shape_og) - - @input_guard - @staticmethod - def backward(ctx, dy): - x, weight, bias, mean, rstd, z = ctx.saved_tensors - dy = dy.reshape(-1, dy.shape[-1]) - if dy.stride(-1) != 1: - dy = dy.contiguous() - assert dy.shape == x.shape - dx, dw, db, dz = layer_norm_bwd( - dy, - x, - weight, - bias, - ctx.eps, - mean, - rstd, - z, - ctx.group_size, - ctx.norm_before_gate, - ctx.is_rms_norm, - ) - dx = dx.reshape(ctx.x_shape_og) - dz = dz.reshape(ctx.x_shape_og) if dz is not None else None - return dx, dw, db, dz, None, None, None, None - - -def layernorm_fn(x, weight, bias, z=None, eps=1e-6, group_size=None, norm_before_gate=True, is_rms_norm=False): - return LayerNormFn.apply(x, weight, bias, z, eps, group_size, norm_before_gate, is_rms_norm) - - -def rmsnorm_fn(x, weight, bias, z=None, eps=1e-6, group_size=None, norm_before_gate=True): - return LayerNormFn.apply(x, weight, bias, z, eps, group_size, norm_before_gate, True) - - -class LayerNormGated(nn.Module): - - def __init__( - self, - hidden_size, - eps: float = 1e-5, - group_size: int | None = None, - norm_before_gate: bool = True, - device: torch.device | None = None, - dtype: torch.dtype | None = None, - ): - """If group_size is not None, we do GroupNorm with each group having group_size elements. - group_size=None is equivalent to group_size=hidden_size (i.e. there's only 1 group). - """ - - factory_kwargs = {"device": device, "dtype": dtype} - super().__init__() - self.eps = eps - self.weight = nn.Parameter(torch.empty(hidden_size, **factory_kwargs)) - self.bias = nn.Parameter(torch.empty(hidden_size, **factory_kwargs)) - self.group_size = group_size - self.norm_before_gate = norm_before_gate - self.reset_parameters() - - def reset_parameters(self): - torch.nn.init.ones_(self.weight) - torch.nn.init.zeros_(self.bias) - - def forward(self, x, z=None): - """If z is not None, we do norm(x) * silu(z) if norm_before_gate, else norm(x * silu(z)) - """ - return layernorm_fn(x, self.weight, self.bias, z=z, group_size=self.group_size, eps=self.eps, - norm_before_gate=self.norm_before_gate) - - -class RMSNormGated(nn.Module): - - def __init__( - self, - hidden_size, - eps: float = 1e-5, - group_size: int | None = None, - norm_before_gate: bool = False, - device: torch.device | None = None, - dtype: torch.dtype | None = None, - ): - """If group_size is not None, we do GroupNorm with each group having group_size elements. - group_size=None is equivalent to group_size=hidden_size (i.e. there's only 1 group). - """ - factory_kwargs = {"device": device, "dtype": dtype} - super().__init__() - self.eps = eps - self.weight = nn.Parameter(torch.empty(hidden_size, **factory_kwargs)) - self.register_parameter("bias", None) - self.group_size = group_size - self.norm_before_gate = norm_before_gate - self.reset_parameters() - - def reset_parameters(self): - torch.nn.init.ones_(self.weight) - - def forward(self, x, z=None): - """If z is not None, we do norm(x) * silu(z) if norm_before_gate, else norm(x * silu(z)) - """ - return rmsnorm_fn(x, self.weight, self.bias, z=z, eps=self.eps, group_size=self.group_size, - norm_before_gate=self.norm_before_gate) diff --git a/build/torch-cuda/modules/conv/long_conv.py b/build/torch-cuda/modules/long_conv.py similarity index 54% rename from build/torch-cuda/modules/conv/long_conv.py rename to build/torch-cuda/modules/long_conv.py index cafa1881400f31cd3168344cbeef2632fbd035b1..1a1ae0d07d268e6b17c1cc6bf364776252ed7882 100644 --- a/build/torch-cuda/modules/conv/long_conv.py +++ b/build/torch-cuda/modules/long_conv.py @@ -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): +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 @@ -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. + 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. + 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 + y (torch.Tensor): + Output of shape [batch_size, seq_len, hidden_size]. """ - 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. - """ + 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): + def forward(self, x: torch.Tensor, *args, **kwargs) -> torch.Tensor: """ - Applies the LongConvolution operation on the input tensor. Args: - x: [batch_size, seq_len, hidden_size] tensor + x (torch.Tensor): + Input of shape [batch_size, seq_len, hidden_size]. + Returns: - y: [batch_size, seq_len, hidden_size] tensor + y (torch.Tensor): + Output of shape [batch_size, seq_len, hidden_size]. """ x = x.transpose(1, 2) - y = fft_conv(x, self.filter, dropout_mask=None, gelu=False) + 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): + 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 filteres is normalized so that t_f = 1 + # 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 - # To compute the right embeddings we use the "proper" linspace + # 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 @@ -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): + def forward(self, L: int) -> torch.Tensor: 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]): + d_emb (int, Optional): 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. + 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. - + 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. - - - """ + 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(d_emb, max_len) + 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) - # final linear layer - self.mlp = nn.Sequential( - nn.Linear(d_emb, d_hidden), - torch.nn.ReLU(), - nn.Linear(d_hidden, hidden_size), - ) + 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): + 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): + def forward(self, x: torch.Tensor, *args, **kwargs) -> torch.Tensor: """ Args: - x: [batch_size, seq_len, hidden_size] tensor + x (torch.Tensor): + Input of shape [batch_size, seq_len, hidden_size]. Returns: - y: [batch_size, seq_len, hidden_size] tensor + 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(x, k, dropout_mask=None, gelu=False) + y = fft_conv(u=x, k=k, dropout_mask=None, gelu=False) y = y.transpose(1, 2) return y.to(dtype=x.dtype) diff --git a/build/torch-cuda/modules/norm/__init__.py b/build/torch-cuda/modules/norm/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..82168e602ac6a52398d4372deba8b111a069ab31 --- /dev/null +++ b/build/torch-cuda/modules/norm/__init__.py @@ -0,0 +1,50 @@ +# 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.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', +] + diff --git a/build/torch-xpu/modules/fused_norm_gate.py b/build/torch-cuda/modules/norm/fused_norm_gate.py similarity index 79% rename from build/torch-xpu/modules/fused_norm_gate.py rename to build/torch-cuda/modules/norm/fused_norm_gate.py index 20af745372b58ff6834b8ff62f0aee8ecedf91f7..447f0c34b31b8ef764b0843dd412c27aba1e4b81 100644 --- a/build/torch-xpu/modules/fused_norm_gate.py +++ b/build/torch-cuda/modules/norm/fused_norm_gate.py @@ -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 +from ...backends import dispatch +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 +@triton.jit(do_not_specialize=['T']) 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 + x, + g, + y, + w, + b, + residual, + residual_out, + mean, + rstd, + eps, + T, + D: tl.constexpr, 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) + 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 = 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) + 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 = 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) + 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 = 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)) + 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 = 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,)) + 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: @@ -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,)) + 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) @@ -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) + 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) - # 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)) + 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( @@ -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 +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, @@ -138,7 +141,7 @@ def layer_norm_gated_fwd_kernel1( HAS_WEIGHT: tl.constexpr, HAS_BIAS: tl.constexpr, ): - i_t = tl.program_id(0) + i_t = tl.program_id(0).to(tl.int64) 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 +@triton.jit(do_not_specialize=['T']) 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 + x, + g, + w, + b, + y, + dy, + dx, + dg, + dw, + db, 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) + i_s = tl.program_id(0).to(tl.int64) 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)) + 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, 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) + 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 = 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 + 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, :] - 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_gate = b_g * b_sigmoid_g 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_gate = b_sigmoid_g b_dg = b_dy * b_y * b_sigmoid_g * (1 - b_sigmoid_g) - b_dy = b_dy * 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: @@ -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) + 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 - # 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)) + 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), boundary_check=(0, 1)) - tl.store(p_dg, b_dg.to(p_dg.dtype.element_ty), boundary_check=(0, 1)) + 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) @@ -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 +@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, @@ -354,7 +356,7 @@ def layer_norm_gated_bwd_kernel1( HAS_BIAS: tl.constexpr, RECOMPUTE_OUTPUT: tl.constexpr, ): - i_s = tl.program_id(0) + i_s = tl.program_id(0).to(tl.int64) 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_gate = b_g * b_sigmoid_g 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_gate = b_sigmoid_g b_dg = b_dy * b_y * b_sigmoid_g * (1 - b_sigmoid_g) - b_dy = b_dy * 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 @@ -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) +# 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, @@ -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, + 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: @@ -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) + # bucket token counts to limit autotuning across sequence lengths. 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,)]( + layer_norm_gated_fwd_kernel_row[(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 +# 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, @@ -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, + 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, + x_dtype: torch.dtype | None = 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. + # cap the program count for short sequences. 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) + # bucket token counts to limit autotuning across sequence lengths. 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, + D=D, 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]( + layer_norm_gated_bwd_kernel_row[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, + D=D, 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 + # 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) @@ -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 + # 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 @@ -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, - ) + 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, - ) + 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: + ) -> None: 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: + ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]: return layer_norm_gated( - x, - g, - self.weight, - self.bias, - self.activation, + x=x, + g=g, + weight=self.weight, + bias=self.bias, + activation=self.activation, residual=residual, - eps=self.eps, prenorm=prenorm, residual_in_fp32=residual_in_fp32, + eps=self.eps, ) @@ -998,7 +975,7 @@ class FusedRMSNormGated(nn.Module): activation: str = "swish", device: torch.device | None = None, dtype: torch.dtype | None = None, - ) -> FusedRMSNormGated: + ) -> None: 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: + ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]: return rms_norm_gated( - x, - g, - self.weight, - self.bias, - self.activation, + x=x, + g=g, + weight=self.weight, + bias=self.bias, + activation=self.activation, residual=residual, - eps=self.eps, prenorm=prenorm, residual_in_fp32=residual_in_fp32, + eps=self.eps, ) @@ -1061,7 +1038,7 @@ class FusedLayerNormSwishGate(FusedLayerNormGated): eps: float = 1e-5, device: torch.device | None = None, dtype: torch.dtype | None = None, - ) -> FusedLayerNormSwishGate: + ) -> None: 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, - ) + ) -> None: + 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: + ) -> None: 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: + ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]: return layer_norm_swish_gate_linear( - x, - g, - self.weight, - self.bias, - weight, - bias, + x=x, + g=g, + norm_weight=self.weight, + norm_bias=self.bias, + linear_weight=weight, + linear_bias=bias, residual=residual, - eps=self.eps, prenorm=prenorm, residual_in_fp32=residual_in_fp32, + eps=self.eps, ) @@ -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, - ) + ) -> 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, + hidden_size: int, elementwise_affine: bool = True, eps: float = 1e-5, device: torch.device | None = None, dtype: torch.dtype | None = None, - ) -> FusedRMSNormGatedLinear: + ) -> None: 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: + ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]: return rms_norm_swish_gate_linear( - x, - g, - self.weight, - self.bias, - weight, - bias, + x=x, + g=g, + norm_weight=self.weight, + norm_bias=self.bias, + linear_weight=weight, + linear_bias=bias, residual=residual, - eps=self.eps, prenorm=prenorm, residual_in_fp32=residual_in_fp32, + eps=self.eps, ) @@ -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, - ) + ) -> 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', +] diff --git a/build/torch-cuda/modules/norm/l2norm.py b/build/torch-cuda/modules/norm/l2norm.py new file mode 100644 index 0000000000000000000000000000000000000000..4261fcda80777d2d630cf802fdcaa21af017e073 --- /dev/null +++ b/build/torch-cuda/modules/norm/l2norm.py @@ -0,0 +1,208 @@ +# 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 torch.nn as nn +import triton +import triton.language as tl + +from ...backends import dispatch +from ...ops.utils.cache import fla_cache_autotune +from ...utils import IS_AMD, autotune_cache_kwargs, input_guard + +BT_LIST = [8, 16, 32, 64, 128] +NUM_WARPS_AUTOTUNE = [1, 2, 4, 8, 16] if IS_AMD else [1, 2, 4, 8, 16, 32] + + +@triton.autotune( + configs=[triton.Config({}, num_warps=num_warps) for num_warps in NUM_WARPS_AUTOTUNE], + key=["D"], + **autotune_cache_kwargs, +) +@triton.jit +def l2norm_fwd_kernel_row(x, y, rstd, eps, D, BD: tl.constexpr): + i_t = tl.program_id(0).to(tl.int64) + x += i_t * D + y += i_t * D + + o_d = tl.arange(0, BD) + mask = o_d < D + + b_x = tl.load(x + o_d, mask=mask, other=0.0).to(tl.float32) + b_rstd = 1 / tl.sqrt(tl.sum(b_x * b_x) + eps) + b_y = b_x * b_rstd + tl.store(y + o_d, b_y, mask=mask) + tl.store(rstd + i_t, b_rstd) + + +@triton.autotune( + configs=[triton.Config({}, num_warps=num_warps) for num_warps in NUM_WARPS_AUTOTUNE], + key=["D"], + **autotune_cache_kwargs, +) +@triton.jit +def l2norm_bwd_kernel_row(y, rstd, dy, dx, eps, D, BD: tl.constexpr): + i_t = tl.program_id(0).to(tl.int64) + y += i_t * D + dx += i_t * D + dy += i_t * D + + o_d = tl.arange(0, BD) + mask = o_d < D + b_y = tl.load(y + o_d, mask=mask, other=0.0).to(tl.float32) + b_rstd = tl.load(rstd + i_t).to(tl.float32) + b_dy = tl.load(dy + o_d, mask=mask, other=0.0).to(tl.float32) + b_dx = b_dy * b_rstd - tl.sum(b_dy * b_y) * b_y * b_rstd + tl.store(dx + o_d, b_dx, mask=mask) + + +@fla_cache_autotune( + configs=[triton.Config({"BT": BT}, num_warps=num_warps) for num_warps in [1, 2, 4, 8, 16] for BT in BT_LIST], + key=["D", "NB"], + **autotune_cache_kwargs, +) +@triton.jit(do_not_specialize=["T"]) +def l2norm_fwd_kernel(x, y, rstd, eps, T, D: tl.constexpr, BD: tl.constexpr, NB: tl.constexpr, BT: 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_t = o_t < T + m_x = m_t[:, None] & (o_d[None, :] < D) + p_x = x + o_t[:, None] * D + o_d[None, :] + p_y = y + o_t[:, None] * D + o_d[None, :] + p_rstd = rstd + o_t + + b_x = tl.load(p_x, mask=m_x, other=0.0).to(tl.float32) + b_rstd = 1 / tl.sqrt(tl.sum(b_x * b_x, 1) + eps) + b_y = b_x * b_rstd[:, None] + + tl.store(p_y, b_y.to(p_y.dtype.element_ty), mask=m_x) + tl.store(p_rstd, b_rstd.to(p_rstd.dtype.element_ty), mask=m_t) + + +@fla_cache_autotune( + configs=[triton.Config({"BT": BT}, num_warps=num_warps) for num_warps in [1, 2, 4, 8, 16] for BT in BT_LIST], + key=["D", "NB"], + **autotune_cache_kwargs, +) +@triton.jit(do_not_specialize=["T"]) +def l2norm_bwd_kernel(y, rstd, dy, dx, eps, T, D: tl.constexpr, BD: tl.constexpr, NB: tl.constexpr, BT: 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_t = o_t < T + m_x = m_t[:, None] & (o_d[None, :] < D) + p_y = y + o_t[:, None] * D + o_d[None, :] + p_rstd = rstd + o_t + p_dy = dy + o_t[:, None] * D + o_d[None, :] + p_dx = dx + o_t[:, None] * D + o_d[None, :] + + b_y = tl.load(p_y, mask=m_x, other=0.0).to(tl.float32) + b_rstd = tl.load(p_rstd, mask=m_t, other=0.0).to(tl.float32) + b_dy = tl.load(p_dy, mask=m_x, other=0.0).to(tl.float32) + b_dx = b_dy * b_rstd[:, None] - tl.sum(b_dy * b_y, 1)[:, None] * b_y * b_rstd[:, None] + tl.store(p_dx, b_dx.to(p_dx.dtype.element_ty), mask=m_x) + + +# Use triton only, other backends (cutlass, cute, etc.) need additional torch.compile support +# @dispatch +def l2norm_fwd(x: torch.Tensor, eps: float = 1e-6, output_dtype: torch.dtype | None = None): + x_shape_og = x.shape + x = x.view(-1, x.shape[-1]) + + if output_dtype is None: + y = torch.empty_like(x) + else: + y = torch.empty_like(x, dtype=output_dtype) + assert y.stride(-1) == 1 + T, D = x.shape[0], x.shape[-1] + + 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 doesn't support feature dim >= 64KB.") + + rstd = torch.empty((T,), dtype=torch.float32, device=x.device) + if D <= 512: + # bucket token counts to limit autotuning across sequence lengths. + NB = triton.cdiv(T, 2048 * 32) + + def grid(meta): + return (triton.cdiv(T, meta["BT"]),) + + l2norm_fwd_kernel[grid](x=x, y=y, rstd=rstd, eps=eps, T=T, D=D, BD=BD, NB=NB) + else: + l2norm_fwd_kernel_row[(T,)](x=x, y=y, rstd=rstd, eps=eps, D=D, BD=BD) + return y.view(x_shape_og), rstd.view(x_shape_og[:-1]) + + +# Use triton only, other backends (cutlass, cute, etc.) need additional torch.compile support +# @dispatch +def l2norm_bwd(y: torch.Tensor, rstd: torch.Tensor, dy: torch.Tensor, eps: float = 1e-6): + y_shape_og = y.shape + y = y.view(-1, dy.shape[-1]) + dy = dy.view(-1, dy.shape[-1]) + assert dy.shape == y.shape + + dx = torch.empty_like(y) + T, D = y.shape[0], y.shape[-1] + + MAX_FUSED_SIZE = 65536 // y.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): + return (triton.cdiv(T, meta["BT"]),) + + l2norm_bwd_kernel[grid](y=y, rstd=rstd, dy=dy, dx=dx, eps=eps, T=T, D=D, BD=BD, NB=NB) + else: + l2norm_bwd_kernel_row[(T,)](y=y, rstd=rstd, dy=dy, dx=dx, eps=eps, D=D, BD=BD) + + return dx.view(y_shape_og) + + +class L2NormFunction(torch.autograd.Function): + @staticmethod + @input_guard + def forward(ctx, x, eps=1e-6, output_dtype=None): + y, rstd = l2norm_fwd(x=x, eps=eps, output_dtype=output_dtype) + ctx.eps = eps + ctx.x_dtype = x.dtype + ctx.save_for_backward(y, rstd) + return y + + @staticmethod + @input_guard + def backward(ctx, dy): + y, rstd = ctx.saved_tensors + dx = l2norm_bwd(y=y, rstd=rstd, dy=dy, eps=ctx.eps) + return dx, None, None + + +def l2norm(x: torch.Tensor, eps: float = 1e-6, output_dtype: torch.dtype | None = None) -> torch.Tensor: + return L2NormFunction.apply(x, eps, output_dtype) + + +l2_norm = l2norm + + +class L2Norm(nn.Module): + def __init__(self, eps: float = 1e-6, output_dtype: torch.dtype | None = None): + super().__init__() + self.eps = eps + self.output_dtype = output_dtype + + def forward(self, x: torch.Tensor) -> torch.Tensor: + return l2norm(x=x, eps=self.eps, output_dtype=self.output_dtype) + + +__all__ = ['L2Norm', 'l2_norm', 'l2norm', 'l2norm_bwd', 'l2norm_fwd'] diff --git a/build/torch-cuda/modules/layernorm.py b/build/torch-cuda/modules/norm/layernorm.py similarity index 77% rename from build/torch-cuda/modules/layernorm.py rename to build/torch-cuda/modules/norm/layernorm.py index 270f74585de44adc85794e2824986ee2b5be1e46..53744aba219e681aaa60b44d68ff883b37b48ed4 100644 --- a/build/torch-cuda/modules/layernorm.py +++ b/build/torch-cuda/modules/norm/layernorm.py @@ -7,12 +7,10 @@ # # Copyright (c) 2023, Tri Dao # https://github.com/state-spaces/mamba/blob/fb7b5310fa865dbd62aa059b1e26f2b431363e2a/mamba_ssm/ops/triton/layernorm.py -# Implement residual + layer_norm / rms_norm. +# implement residual + layer_norm / rms_norm. -# Based on the Triton LayerNorm tutorial: https://triton-lang.org/main/getting-started/tutorials/05-layer-norm.html -# For the backward pass, we keep weight_grad and bias_grad in registers and accumulate. -# This is faster for dimensions up to 8k, but after that it's much slower due to register spilling. -# The models we train have hidden dim up to 8k anyway (e.g. Llama 70B), so this is fine. +# based on the Triton LayerNorm tutorial: https://triton-lang.org/main/getting-started/tutorials/05-layer-norm.html +# accumulate affine gradients in registers; wide dimensions can cause register spilling. from __future__ import annotations @@ -28,8 +26,8 @@ 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 ..utils import autotune_cache_kwargs, get_multiprocessor_count, input_guard +from ...backends import dispatch +from ...utils import IS_INTEL, autotune_cache_kwargs, get_multiprocessor_count, input_guard try: from torch.distributed.tensor import DTensor @@ -37,151 +35,6 @@ except (ImportError, AttributeError): DTensor = None -def layer_norm_ref( - x: torch.Tensor, - weight: torch.Tensor, - bias: torch.Tensor, - residual: torch.Tensor = None, - eps: float = 1e-5, - prenorm: bool = False, - upcast: bool = False, -): - dtype = x.dtype - if upcast: - weight = weight.float() - bias = bias.float() if bias is not None else None - if upcast: - x = x.float() - residual = residual.float() if residual is not None else residual - if residual is not None: - x = (x + residual).to(x.dtype) - out = F.layer_norm(x.to(weight.dtype), x.shape[-1:], weight=weight, bias=bias, eps=eps).to( - dtype, - ) - return out if not prenorm else (out, x) - - -def rms_norm_ref( - x: torch.Tensor, - weight: torch.Tensor, - bias: torch.Tensor, - residual: torch.Tensor = None, - eps: float = 1e-5, - prenorm: bool = False, - upcast: bool = False, -): - dtype = x.dtype - if upcast: - weight = weight.float() - bias = bias.float() if bias is not None else None - if upcast: - x = x.float() - residual = residual.float() if residual is not None else residual - if residual is not None: - x = (x + residual).to(x.dtype) - rstd = 1 / torch.sqrt((x.square()).mean(dim=-1, keepdim=True) + eps) - out = (x * rstd * weight) + bias if bias is not None else (x * rstd * weight) - out = out.to(dtype) - return out if not prenorm else (out, x) - - -def group_norm_ref( - x: torch.Tensor, - weight: torch.Tensor, - bias: torch.Tensor, - num_groups: int, - residual: torch.Tensor = None, - eps: float = 1e-5, - is_rms_norm: bool = False, - prenorm: bool = False, - upcast: bool = False, -): - dtype = x.dtype - if upcast: - weight = weight.float() - bias = bias.float() if bias is not None else None - if upcast: - x = x.float() - residual = residual.float() if residual is not None else residual - if residual is not None: - x = (x + residual).to(x.dtype) - residual = x - x, weight = [ - rearrange(data, "... (g d) -> ... g d", g=num_groups) for data in (x, weight) - ] - if bias is not None: - bias = rearrange(bias, '... (g d) -> ... g d', g=num_groups) - if not is_rms_norm: - mean = x.mean(dim=-1, keepdim=True) - x = x - mean - rstd = 1 / torch.sqrt((x.square()).mean(dim=-1, keepdim=True) + eps) - out = (x * rstd * weight) + bias if bias is not None else (x * rstd * weight) - out = rearrange(out, "... g d -> ... (g d)") - out = out.to(dtype) - return out if not prenorm else (out, residual) - - -class GroupNormRef(nn.Module): - - def __init__( - self, - num_groups: int, - hidden_size: int, - elementwise_affine: bool = True, - bias: bool = False, - eps: float = 1e-5, - is_rms_norm: bool = False, - ) -> GroupNormRef: - super().__init__() - - if hidden_size % num_groups != 0: - raise ValueError('num_channels must be divisible by num_groups') - - self.num_groups = num_groups - self.hidden_size = hidden_size - self.elementwise_affine = elementwise_affine - self.eps = eps - self.is_rms_norm = is_rms_norm - - self.register_parameter("weight", None) - self.register_parameter("bias", None) - if elementwise_affine: - self.weight = nn.Parameter(torch.empty(hidden_size)) - if bias: - self.bias = nn.Parameter(torch.empty(hidden_size)) - - 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.num_groups}, {self.hidden_size}" - if not self.elementwise_affine: - s += f", elementwise_affine={self.elementwise_affine}" - if self.is_rms_norm: - s += f", is_rms_norm={self.is_rms_norm}" - s += f", eps={self.eps}" - s += ")" - return s - - def forward(self, x, residual=None, prenorm=False): - return group_norm_ref( - x, - self.weight, - self.bias, - num_groups=self.num_groups, - residual=residual, - eps=self.eps, - is_rms_norm=self.is_rms_norm, - prenorm=prenorm, - upcast=True, - ) - - @triton.autotune( configs=[ triton.Config({'BT': BT}, num_warps=num_warps) @@ -191,17 +44,17 @@ class GroupNormRef(nn.Module): key=['D', 'NB', 'HAS_RESIDUAL', 'STORE_RESIDUAL_OUT', 'IS_RMS_NORM'], **autotune_cache_kwargs, ) -@triton.jit +@triton.jit(do_not_specialize=['T']) def layer_norm_fwd_kernel( - x, # pointer to the input - y, # pointer to the output - w, # pointer to the weights - b, # pointer to the biases - res, # pointer to the res - res_out, # pointer to the res - mean, # pointer to the mean - rstd, # pointer to the 1/std - eps, # epsilon to avoid division by zero + x, + y, + w, + b, + res, + res_out, + mean, + rstd, + eps, T, G: tl.constexpr, D: tl.constexpr, @@ -214,25 +67,27 @@ def layer_norm_fwd_kernel( HAS_WEIGHT: tl.constexpr, HAS_BIAS: tl.constexpr, ): - i_t = tl.program_id(0) + i_t = tl.program_id(0).to(tl.int64) o_t = i_t * BT + tl.arange(0, BT) o_g = o_t % G 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 = 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) + 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 = tl.make_block_ptr(res, (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) + p_res = res + 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 = tl.make_block_ptr(res_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)) + p_res_out = res_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 = 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,)) + 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: @@ -240,8 +95,8 @@ def layer_norm_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,)) + 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_g[:, None] * D + o_d[None, :], mask=m_d[None, :]).to(tl.float32) @@ -252,9 +107,8 @@ def layer_norm_fwd_kernel( if HAS_BIAS: b_y = b_y + b_b - # 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)) + 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.autotune( @@ -266,16 +120,16 @@ def layer_norm_fwd_kernel( **autotune_cache_kwargs, ) @triton.jit -def layer_norm_fwd_kernel1( - x, # pointer to the input - y, # pointer to the output - w, # pointer to the weights - b, # pointer to the biases - res, # pointer to the res - res_out, # pointer to the res - mean, # pointer to the mean - rstd, # pointer to the 1/std - eps, # epsilon to avoid division by zero +def layer_norm_fwd_kernel_row( + x, + y, + w, + b, + res, + res_out, + mean, + rstd, + eps, G: tl.constexpr, D: tl.constexpr, BD: tl.constexpr, @@ -285,7 +139,7 @@ def layer_norm_fwd_kernel1( HAS_WEIGHT: tl.constexpr, HAS_BIAS: tl.constexpr, ): - i_t = tl.program_id(0) + i_t = tl.program_id(0).to(tl.int64) i_g = i_t % G x += i_t * D @@ -322,13 +176,10 @@ def layer_norm_fwd_kernel1( if HAS_BIAS: b_y = b_y + b_b - # Write output tl.store(y + o_d, b_y, mask=m_d) -@triton.heuristics({ - 'RECOMPUTE_OUTPUT': lambda args: args['y'] is not None, -}) +@triton.heuristics({'RECOMPUTE_OUTPUT': lambda args: args['y'] is not None}) @triton.autotune( configs=[ triton.Config({'BT': BT}, num_warps=num_warps) @@ -338,16 +189,16 @@ def layer_norm_fwd_kernel1( key=['D', 'NB', 'HAS_DRESIDUAL', 'STORE_DRESIDUAL', 'IS_RMS_NORM'], **autotune_cache_kwargs, ) -@triton.jit +@triton.jit(do_not_specialize=['T']) 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 + x, + w, + b, + y, + dy, + dx, + dw, + db, dres, dres_in, mean, @@ -367,7 +218,7 @@ def layer_norm_bwd_kernel( HAS_BIAS: tl.constexpr, RECOMPUTE_OUTPUT: tl.constexpr, ): - i_s = tl.program_id(0) + i_s = tl.program_id(0).to(tl.int64) i_g, i_sg = i_s // GS, i_s % GS o_d = tl.arange(0, BD) @@ -379,26 +230,27 @@ def layer_norm_bwd_kernel( b_b = tl.load(b + i_g * D + o_d, mask=m_d, other=0.0).to(tl.float32) b_db = tl.zeros((BT, BD), dtype=tl.float32) - # Tg: number of tokens per group, used as the logical shape for make_block_ptr. - # for mean/rstd with shape (T,) and stride (G,), the strided view has Tg elements per group. - # the caller guarantees NS capped so every program has work. - # the last program's range may slightly exceed Tg (since BS = cdiv(T, NS)); - # boundary_check handles the partial tail tile, m_t < Tg masks dw/db accumulation. + # tokens are interleaved by group in x, mean, and rstd. Tg = T // G for i_t in range(i_sg * BS, i_sg * BS + BS, BT): - p_x = tl.make_block_ptr(x + i_g * D, (Tg, D), (G*D, 1), (i_t, 0), (BT, BD), (1, 0)) - p_dy = tl.make_block_ptr(dy + i_g * D, (Tg, D), (G*D, 1), (i_t, 0), (BT, BD), (1, 0)) - p_dx = tl.make_block_ptr(dx + i_g * D, (Tg, D), (G*D, 1), (i_t, 0), (BT, BD), (1, 0)) + o_t = (i_t + tl.arange(0, BT)).to(tl.int64) + m_t = o_t < Tg + m_x = m_t[:, None] & m_d[None, :] + m_o = o_t < min(i_sg * BS + BS, Tg) + m_ox = m_o[:, None] & m_d[None, :] + p_x = x + i_g * D + o_t[:, None] * (G*D) + o_d[None, :] + p_dy = dy + i_g * D + o_t[:, None] * (G*D) + o_d[None, :] + p_dx = dx + i_g * D + o_t[:, None] * (G*D) + o_d[None, :] # [BT, BD] - b_x = tl.load(p_x, boundary_check=(0, 1)).to(tl.float32) - b_dy = tl.load(p_dy, boundary_check=(0, 1)).to(tl.float32) + b_x = tl.load(p_x, 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 = tl.make_block_ptr(mean + i_g, (Tg,), (G,), (i_t,), (BT,), (0,)) - b_mean = tl.load(p_mean, boundary_check=(0,)) - p_rstd = tl.make_block_ptr(rstd + i_g, (Tg,), (G,), (i_t,), (BT,), (0,)) - b_rstd = tl.load(p_rstd, boundary_check=(0,)) - # Compute dx + p_mean = mean + i_g + o_t * G + b_mean = tl.load(p_mean, mask=m_t, other=0.0) + p_rstd = rstd + i_g + o_t * G + 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) @@ -406,20 +258,16 @@ def layer_norm_bwd_kernel( if HAS_BIAS: b_y = b_y + b_b[None, :] if RECOMPUTE_OUTPUT: - p_y = tl.make_block_ptr(y + i_g * D, (Tg, D), (G*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)) + p_y = y + i_g * D + o_t[:, None] * (G*D) + o_d[None, :] + tl.store(p_y, b_y.to(p_y.dtype.element_ty), mask=m_ox) 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_sg * BS + BS, Tg) if HAS_WEIGHT: b_wdy = b_dy * b_w - b_dw += tl.where(m_t[:, None], b_dy * b_xhat, 0.0) + b_dw += tl.where(m_o[:, None], b_dy * b_xhat, 0.0) if HAS_BIAS: - b_db += tl.where(m_t[:, None], b_dy, 0.0) + b_db += tl.where(m_o[:, 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 @@ -428,15 +276,15 @@ def layer_norm_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(dres + i_g * D, (Tg, D), (G*D, 1), (i_t, 0), (BT, BD), (1, 0)) - b_dres = tl.load(p_dres, boundary_check=(0, 1)).to(tl.float32) + p_dres = dres + i_g * D + o_t[:, None] * (G*D) + o_d[None, :] + b_dres = tl.load(p_dres, mask=m_x, other=0.0).to(tl.float32) b_dx += b_dres - # Write dx + if STORE_DRESIDUAL: - p_dres_in = tl.make_block_ptr(dres_in + i_g * D, (Tg, D), (G*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)) + p_dres_in = dres_in + i_g * D + o_t[:, None] * (G*D) + o_d[None, :] + tl.store(p_dres_in, b_dx.to(p_dres_in.dtype.element_ty), mask=m_ox) - tl.store(p_dx, b_dx.to(p_dx.dtype.element_ty), boundary_check=(0, 1)) + tl.store(p_dx, b_dx.to(p_dx.dtype.element_ty), mask=m_ox) if HAS_WEIGHT: tl.store(dw + i_s * D + o_d, tl.sum(b_dw, axis=0), mask=m_d) @@ -444,9 +292,7 @@ def layer_norm_bwd_kernel( tl.store(db + i_s * D + o_d, tl.sum(b_db, axis=0), mask=m_d) -@triton.heuristics({ - 'RECOMPUTE_OUTPUT': lambda args: args['y'] is not None, -}) +@triton.heuristics({'RECOMPUTE_OUTPUT': lambda args: args['y'] is not None}) @triton.autotune( configs=[ triton.Config({}, num_warps=num_warps) @@ -455,16 +301,16 @@ def layer_norm_bwd_kernel( key=['D', 'HAS_DRESIDUAL', 'STORE_DRESIDUAL', 'IS_RMS_NORM'], **autotune_cache_kwargs, ) -@triton.jit -def layer_norm_bwd_kernel1( - 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 +@triton.jit(do_not_specialize=['T']) +def layer_norm_bwd_kernel_row( + x, + w, + b, + y, + dy, + dx, + dw, + db, dres, dres_in, mean, @@ -482,7 +328,7 @@ def layer_norm_bwd_kernel1( HAS_BIAS: tl.constexpr, RECOMPUTE_OUTPUT: tl.constexpr, ): - i_s = tl.program_id(0) + i_s = tl.program_id(0).to(tl.int64) i_g, i_sg = i_s // GS, i_s % GS o_d = tl.arange(0, BD) @@ -503,7 +349,7 @@ def layer_norm_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) if RECOMPUTE_OUTPUT: @@ -527,7 +373,7 @@ def layer_norm_bwd_kernel1( if HAS_DRESIDUAL: b_dres = tl.load(dres + i_t * D + o_d, mask=mask, other=0).to(tl.float32) b_dx += b_dres - # Write dx + 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) @@ -539,15 +385,15 @@ def layer_norm_bwd_kernel1( tl.store(db + i_s * D + o_d, b_db, mask=mask) -@dispatch('modules') +@dispatch def layer_norm_fwd( 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, + residual: torch.Tensor | None = None, + out_dtype: torch.dtype | None = None, + residual_dtype: torch.dtype | None = None, is_rms_norm: bool = False, num_groups: int = 1, ): @@ -560,7 +406,7 @@ def layer_norm_fwd( assert weight.shape == (G * D,) if bias is not None: assert bias.shape == (G * 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): res_out = torch.empty(T, D, device=x.device, dtype=residual_dtype) @@ -568,26 +414,28 @@ def layer_norm_fwd( 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) - # 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: + # use the row kernel at D == 512 to fit Intel scratch-space limits. + if D <= 512 and not (D == 512 and IS_INTEL): NB = triton.cdiv(T, 2048) - def grid(meta): return (triton.cdiv(T, meta['BT']), ) + + def grid(meta): + return (triton.cdiv(T, meta['BT']),) + layer_norm_fwd_kernel[grid]( - x, - y, - weight, - bias, - residual, - res_out, - mean, - rstd, - eps, + x=x, + y=y, + w=weight, + b=bias, + res=residual, + res_out=res_out, + mean=mean, + rstd=rstd, + eps=eps, T=T, G=G, D=D, @@ -600,16 +448,16 @@ def layer_norm_fwd( HAS_BIAS=bias is not None, ) else: - layer_norm_fwd_kernel1[(T,)]( - x, - y, - weight, - bias, - residual, - res_out, - mean, - rstd, - eps, + layer_norm_fwd_kernel_row[(T,)]( + x=x, + y=y, + w=weight, + b=bias, + res=residual, + res_out=res_out, + mean=mean, + rstd=rstd, + eps=eps, G=G, D=D, BD=BD, @@ -623,18 +471,18 @@ def layer_norm_fwd( return y, mean, rstd, res_out if res_out is not None else x -@dispatch('modules') +@dispatch def layer_norm_bwd( dy: torch.Tensor, x: torch.Tensor, weight: torch.Tensor, bias: torch.Tensor, - mean: torch.Tensor = None, - rstd: torch.Tensor = None, - dres: torch.Tensor = None, + mean: torch.Tensor | None = None, + rstd: torch.Tensor | None = None, + dres: torch.Tensor | None = None, has_residual: bool = False, is_rms_norm: bool = False, - x_dtype: torch.dtype = None, + x_dtype: torch.dtype | None = None, recompute_output: bool = False, num_groups: int = 1, ): @@ -646,20 +494,16 @@ def layer_norm_bwd( assert weight.shape == (G * D,) if bias is not None: assert bias.shape == (G * D,) - # allocate output + 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 - # 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.") - # each program handles one group only. - # cap per-group program count to T // G 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. + # cap the per-group program count for short sequences. NS = min(triton.cdiv(get_multiprocessor_count(x.device.index), G), T // G) * G BS = triton.cdiv(T, NS) GS = NS // G @@ -668,21 +512,22 @@ def layer_norm_bwd( db = torch.empty((NS, D), dtype=torch.float, device=bias.device) if bias is not None else None grid = (NS,) - if D <= 512: + # use the row kernel at D == 512 to fit Intel scratch-space limits. + if D <= 512 and not (D == 512 and IS_INTEL): NB = triton.cdiv(T, 2048) layer_norm_bwd_kernel[grid]( - x, - weight, - bias, - y, - dy, - dx, - dw, - db, - dres, - dres_in, - mean, - rstd, + x=x, + w=weight, + b=bias, + y=y, + dy=dy, + dx=dx, + dw=dw, + db=db, + dres=dres, + dres_in=dres_in, + mean=mean, + rstd=rstd, T=T, G=G, D=D, @@ -697,19 +542,19 @@ def layer_norm_bwd( HAS_BIAS=bias is not None, ) else: - layer_norm_bwd_kernel1[grid]( - x, - weight, - bias, - y, - dy, - dx, - dw, - db, - dres, - dres_in, - mean, - rstd, + layer_norm_bwd_kernel_row[grid]( + x=x, + w=weight, + b=bias, + y=y, + dy=dy, + dx=dx, + dw=dw, + db=db, + dres=dres, + dres_in=dres_in, + mean=mean, + rstd=rstd, T=T, G=G, D=D, @@ -724,7 +569,7 @@ def layer_norm_bwd( ) 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 - # Don't need to compute dres_in separately in this case + # reuse dx when the residual gradient has the same dtype. 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) @@ -739,7 +584,7 @@ class LayerNormFunction(torch.autograd.Function): x, weight, bias, - residual: torch.Tensor = None, + residual: torch.Tensor | None = None, eps: float = 1e-5, prenorm: bool = False, residual_in_fp32: bool = False, @@ -750,22 +595,18 @@ class LayerNormFunction(torch.autograd.Function): if x.shape[-1] % num_groups != 0: raise ValueError('num_channels must be divisible by num_groups') - # reshape input data into 2D tensor - x = x.reshape(-1, (x.shape[-1] // num_groups)) + + x = x.reshape(-1, x.shape[-1] // num_groups) if residual is not None: assert residual.shape == x_shape_og residual = residual.reshape_as(x) - residual_dtype = ( - residual.dtype - if residual is not None - else (torch.float32 if residual_in_fp32 else None) - ) + residual_dtype = residual.dtype if residual is not None else (torch.float32 if residual_in_fp32 else None) y, mean, rstd, res_out = layer_norm_fwd( - x, - weight, - bias, - eps, - residual, + x=x, + weight=weight, + bias=bias, + eps=eps, + residual=residual, residual_dtype=residual_dtype, is_rms_norm=is_rms_norm, num_groups=num_groups, @@ -785,7 +626,7 @@ class LayerNormFunction(torch.autograd.Function): @input_guard def backward(ctx, dy, *args): x, weight, bias, mean, rstd = ctx.saved_tensors - dy = dy.reshape(-1, (dy.shape[-1] // ctx.num_groups)) + dy = dy.reshape(-1, dy.shape[-1] // ctx.num_groups) assert dy.shape == x.shape if ctx.prenorm: dresidual = args[0] @@ -794,15 +635,15 @@ class LayerNormFunction(torch.autograd.Function): else: dresidual = None dx, dw, db, dresidual_in = layer_norm_bwd( - dy, - x, - weight, - bias, - mean, - rstd, - dresidual, - ctx.has_residual, - ctx.is_rms_norm, + dy=dy, + x=x, + weight=weight, + bias=bias, + mean=mean, + rstd=rstd, + dres=dresidual, + has_residual=ctx.has_residual, + is_rms_norm=ctx.is_rms_norm, x_dtype=ctx.x_dtype, num_groups=ctx.num_groups, ) @@ -823,67 +664,39 @@ def layer_norm( x: torch.Tensor, weight: torch.Tensor, bias: torch.Tensor, - residual: torch.Tensor = None, + residual: torch.Tensor | None = None, eps: float = 1e-5, prenorm: bool = False, residual_in_fp32: bool = False, is_rms_norm: bool = False, ): - return LayerNormFunction.apply( - x, - weight, - bias, - residual, - eps, - prenorm, - residual_in_fp32, - is_rms_norm, - ) + return LayerNormFunction.apply(x, weight, bias, residual, eps, prenorm, residual_in_fp32, is_rms_norm) def group_norm( x: torch.Tensor, weight: torch.Tensor, bias: torch.Tensor, - residual: torch.Tensor = None, + residual: torch.Tensor | None = None, eps: float = 1e-5, prenorm: bool = False, residual_in_fp32: bool = False, is_rms_norm: bool = False, num_groups: int = 1, ): - return LayerNormFunction.apply( - x, - weight, - bias, - residual, - eps, - prenorm, - residual_in_fp32, - is_rms_norm, - num_groups, - ) + return LayerNormFunction.apply(x, weight, bias, residual, eps, prenorm, residual_in_fp32, is_rms_norm, num_groups) def rms_norm( x: torch.Tensor, weight: torch.Tensor, bias: torch.Tensor, - residual: torch.Tensor = None, + residual: torch.Tensor | None = None, eps: float = 1e-5, prenorm: bool = False, residual_in_fp32: bool = False, ): - return LayerNormFunction.apply( - x, - weight, - bias, - residual, - eps, - prenorm, - residual_in_fp32, - True, - ) + return LayerNormFunction.apply(x, weight, bias, residual, eps, prenorm, residual_in_fp32, True) def layer_norm_linear( @@ -892,7 +705,7 @@ def layer_norm_linear( norm_bias: torch.Tensor, linear_weight: torch.Tensor, linear_bias: torch.Tensor, - residual: torch.Tensor = None, + residual: torch.Tensor | None = None, eps: float = 1e-5, prenorm: bool = False, residual_in_fp32: bool = False, @@ -920,7 +733,7 @@ def rms_norm_linear( norm_bias: torch.Tensor, linear_weight: torch.Tensor, linear_bias: torch.Tensor, - residual: torch.Tensor = None, + residual: torch.Tensor | None = None, eps: float = 1e-5, prenorm: bool = False, residual_in_fp32: bool = False, @@ -945,7 +758,7 @@ def group_norm_linear( norm_bias: torch.Tensor, linear_weight: torch.Tensor, linear_bias: torch.Tensor, - residual: torch.Tensor = None, + residual: torch.Tensor | None = None, eps: float = 1e-5, prenorm: bool = False, residual_in_fp32: bool = False, @@ -967,6 +780,187 @@ def group_norm_linear( ) +class LayerNormLinearFunction(torch.autograd.Function): + + @staticmethod + @input_guard + def forward( + ctx, + x, + norm_weight, + norm_bias, + linear_weight, + linear_bias, + residual=None, + eps=1e-5, + prenorm=False, + residual_in_fp32=False, + is_rms_norm=False, + num_groups=1, + ): + x_shape_og = x.shape + + if x.shape[-1] % num_groups != 0: + raise ValueError('num_channels must be divisible by num_groups') + + x = x.reshape(-1, x.shape[-1] // num_groups) + if residual is not None: + assert residual.shape == x_shape_og + residual = residual.reshape_as(x) + residual_dtype = residual.dtype if residual is not None else (torch.float32 if residual_in_fp32 else None) + y, mean, rstd, res_out = layer_norm_fwd( + x=x, + weight=norm_weight, + bias=norm_bias, + eps=eps, + residual=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, + num_groups=num_groups, + ) + 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) + # recompute y in backward to save memory. + ctx.save_for_backward(res_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.num_groups = num_groups + 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, res_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()) + dy = dy.reshape(-1, dy.shape[-1] // ctx.num_groups) + 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, x.shape[-1]) + assert dresidual.shape == x.shape + else: + dresidual = None + dx, dnorm_weight, dnorm_bias, dresidual_in, y = layer_norm_bwd( + dy=dy, + x=x, + weight=norm_weight, + bias=norm_bias, + mean=mean, + rstd=rstd, + dres=dresidual, + has_residual=ctx.has_residual, + is_rms_norm=ctx.is_rms_norm, + x_dtype=ctx.x_dtype, + recompute_output=True, + num_groups=ctx.num_groups, + ) + dlinear_weight = torch.einsum("bo,bi->oi", dout, y.view(-1, linear_weight.shape[-1])) + 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, + None, + ) + + +def layer_norm_ref( + x: torch.Tensor, + weight: torch.Tensor, + bias: torch.Tensor, + residual: torch.Tensor | None = None, + eps: float = 1e-5, + prenorm: bool = False, + upcast: bool = False, +): + dtype = x.dtype + if upcast: + weight = weight.float() + bias = bias.float() if bias is not None else None + if upcast: + x = x.float() + residual = residual.float() if residual is not None else residual + if residual is not None: + x = (x + residual).to(x.dtype) + out = F.layer_norm(x.to(weight.dtype), x.shape[-1:], weight=weight, bias=bias, eps=eps).to(dtype) + return out if not prenorm else (out, x) + + +def rms_norm_ref( + x: torch.Tensor, + weight: torch.Tensor, + bias: torch.Tensor, + residual: torch.Tensor | None = None, + eps: float = 1e-5, + prenorm: bool = False, + upcast: bool = False, +): + dtype = x.dtype + if upcast: + weight = weight.float() + bias = bias.float() if bias is not None else None + if upcast: + x = x.float() + residual = residual.float() if residual is not None else residual + if residual is not None: + x = (x + residual).to(x.dtype) + rstd = 1 / torch.sqrt((x.square()).mean(dim=-1, keepdim=True) + eps) + out = (x * rstd * weight) + bias if bias is not None else (x * rstd * weight) + out = out.to(dtype) + return out if not prenorm else (out, x) + + +def group_norm_ref( + x: torch.Tensor, + weight: torch.Tensor, + bias: torch.Tensor, + num_groups: int, + residual: torch.Tensor | None = None, + eps: float = 1e-5, + is_rms_norm: bool = False, + prenorm: bool = False, + upcast: bool = False, +): + dtype = x.dtype + if upcast: + weight = weight.float() + bias = bias.float() if bias is not None else None + if upcast: + x = x.float() + residual = residual.float() if residual is not None else residual + if residual is not None: + x = (x + residual).to(x.dtype) + residual = x + x, weight = [rearrange(data, "... (g d) -> ... g d", g=num_groups) for data in (x, weight)] + if bias is not None: + bias = rearrange(bias, '... (g d) -> ... g d', g=num_groups) + if not is_rms_norm: + mean = x.mean(dim=-1, keepdim=True) + x = x - mean + rstd = 1 / torch.sqrt((x.square()).mean(dim=-1, keepdim=True) + eps) + out = (x * rstd * weight) + bias if bias is not None else (x * rstd * weight) + out = rearrange(out, "... g d -> ... (g d)") + out = out.to(dtype) + return out if not prenorm else (out, residual) + + class LayerNorm(nn.Module): def __init__( @@ -977,7 +971,7 @@ class LayerNorm(nn.Module): eps: float = 1e-5, device: torch.device | None = None, dtype: torch.dtype | None = None, - ) -> LayerNorm: + ) -> None: factory_kwargs = {"device": device, "dtype": dtype} super().__init__() @@ -1010,9 +1004,9 @@ class LayerNorm(nn.Module): def forward(self, x, residual=None, prenorm=False, residual_in_fp32=False): return layer_norm( - x, - self.weight, - self.bias, + x=x, + weight=self.weight, + bias=self.bias, residual=residual, eps=self.eps, prenorm=prenorm, @@ -1032,7 +1026,7 @@ class GroupNorm(nn.Module): is_rms_norm: bool = False, device: torch.device | None = None, dtype: torch.dtype | None = None, - ) -> GroupNorm: + ) -> None: factory_kwargs = {"device": device, "dtype": dtype} super().__init__() @@ -1072,9 +1066,9 @@ class GroupNorm(nn.Module): def forward(self, x, residual=None, prenorm=False, residual_in_fp32=False): return group_norm( - x, - self.weight, - self.bias, + x=x, + weight=self.weight, + bias=self.bias, residual=residual, eps=self.eps, prenorm=prenorm, @@ -1094,7 +1088,7 @@ class RMSNorm(nn.Module): eps: float = 1e-5, device: torch.device | None = None, dtype: torch.dtype | None = None, - ) -> RMSNorm: + ) -> None: factory_kwargs = {"device": device, "dtype": dtype} super().__init__() @@ -1127,9 +1121,9 @@ class RMSNorm(nn.Module): def forward(self, x, residual=None, prenorm=False, residual_in_fp32=False): return rms_norm( - x, - self.weight, - self.bias, + x=x, + weight=self.weight, + bias=self.bias, residual=residual, eps=self.eps, prenorm=prenorm, @@ -1137,122 +1131,17 @@ class RMSNorm(nn.Module): ) -class LayerNormLinearFunction(torch.autograd.Function): - - @staticmethod - @input_guard - def forward( - ctx, - x, - norm_weight, - norm_bias, - linear_weight, - linear_bias, - residual=None, - eps=1e-5, - prenorm=False, - residual_in_fp32=False, - is_rms_norm=False, - num_groups=1, - ): - x_shape_og = x.shape - - if x.shape[-1] % num_groups != 0: - raise ValueError('num_channels must be divisible by num_groups') - # reshape input data into 2D tensor - x = x.reshape(-1, (x.shape[-1] // num_groups)) - if residual is not None: - assert residual.shape == x_shape_og - residual = residual.reshape_as(x) - residual_dtype = ( - residual.dtype - if residual is not None - else (torch.float32 if residual_in_fp32 else None) - ) - y, mean, rstd, res_out = layer_norm_fwd( - 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, - num_groups=num_groups, - ) - 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(res_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.num_groups = num_groups - 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, res_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()) - dy = dy.reshape(-1, (dy.shape[-1] // ctx.num_groups)) - 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, x.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, - mean, - rstd, - dresidual, - ctx.has_residual, - ctx.is_rms_norm, - x_dtype=ctx.x_dtype, - recompute_output=True, - num_groups=ctx.num_groups, - ) - dlinear_weight = torch.einsum("bo,bi->oi", dout, y.view(-1, linear_weight.shape[-1])) - 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, - None, - ) - - class LayerNormLinear(nn.Module): def __init__( self, - hidden_size, + hidden_size: int, elementwise_affine: bool = True, bias: bool = False, eps: float = 1e-5, device: torch.device | None = None, dtype: torch.dtype | None = None, - ) -> LayerNormLinear: + ) -> None: factory_kwargs = {"device": device, "dtype": dtype} super().__init__() @@ -1310,7 +1199,7 @@ class GroupNormLinear(nn.Module): is_rms_norm: bool = False, device: torch.device | None = None, dtype: torch.dtype | None = None, - ) -> GroupNormLinear: + ) -> None: factory_kwargs = {"device": device, "dtype": dtype} super().__init__() @@ -1368,13 +1257,13 @@ class RMSNormLinear(nn.Module): def __init__( self, - hidden_size, + hidden_size: int, elementwise_affine: bool = True, bias: bool = False, eps: float = 1e-5, device: torch.device | None = None, dtype: torch.dtype | None = None, - ) -> RMSNormLinear: + ) -> None: factory_kwargs = {"device": device, "dtype": dtype} super().__init__() @@ -1427,15 +1316,11 @@ class NormParallel(ParallelStyle): self.sequence_sharding = (Shard(sequence_dim),) self.use_local_output = use_local_output - def _replicate_module_fn( - self, name: str, module: nn.Module, device_mesh: DeviceMesh, - ): + def _replicate_module_fn(self, name: str, module: nn.Module, device_mesh: DeviceMesh): for p_name, param in module.named_parameters(): # simple replication with fixed ones_ init from LayerNorm/RMSNorm, which allow # us to simply just use from_local - replicated_param = torch.nn.Parameter( - DTensor.from_local(param, device_mesh, [Replicate()], run_check=False), - ) + replicated_param = torch.nn.Parameter(DTensor.from_local(param, device_mesh, [Replicate()], run_check=False)) module.register_parameter(p_name, replicated_param) @staticmethod @@ -1444,19 +1329,13 @@ class NormParallel(ParallelStyle): if isinstance(input_tensor, DTensor): # if the passed in input DTensor is not sharded on the sequence dim, we need to redistribute it if input_tensor.placements != sequence_sharding: - input_tensor = input_tensor.redistribute( - placements=sequence_sharding, async_op=True, - ) + input_tensor = input_tensor.redistribute(placements=sequence_sharding, async_op=True) return input_tensor elif isinstance(input_tensor, torch.Tensor): # assume the input passed in already sharded on the sequence dim and create the DTensor - return DTensor.from_local( - input_tensor, device_mesh, sequence_sharding, run_check=False, - ) + return DTensor.from_local(input_tensor, device_mesh, sequence_sharding, run_check=False) else: - raise ValueError( - f"expecting input of {mod} to be a torch.Tensor or DTensor, but got {input_tensor}", - ) + raise ValueError(f"expecting input of {mod} to be a torch.Tensor or DTensor, but got {input_tensor}") @staticmethod def _prepare_output_fn(use_local_output, mod, outputs, device_mesh): @@ -1470,3 +1349,87 @@ class NormParallel(ParallelStyle): partial(self._prepare_input_fn, self.sequence_sharding), partial(self._prepare_output_fn, self.use_local_output), ) + + +class GroupNormRef(nn.Module): + + def __init__( + self, + num_groups: int, + hidden_size: int, + elementwise_affine: bool = True, + bias: bool = False, + eps: float = 1e-5, + is_rms_norm: bool = False, + ) -> None: + super().__init__() + + if hidden_size % num_groups != 0: + raise ValueError('num_channels must be divisible by num_groups') + + self.num_groups = num_groups + self.hidden_size = hidden_size + self.elementwise_affine = elementwise_affine + self.eps = eps + self.is_rms_norm = is_rms_norm + + self.register_parameter("weight", None) + self.register_parameter("bias", None) + if elementwise_affine: + self.weight = nn.Parameter(torch.empty(hidden_size)) + if bias: + self.bias = nn.Parameter(torch.empty(hidden_size)) + + 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.num_groups}, {self.hidden_size}" + if not self.elementwise_affine: + s += f", elementwise_affine={self.elementwise_affine}" + if self.is_rms_norm: + s += f", is_rms_norm={self.is_rms_norm}" + s += f", eps={self.eps}" + s += ")" + return s + + def forward(self, x, residual=None, prenorm=False): + return group_norm_ref( + x=x, + weight=self.weight, + bias=self.bias, + num_groups=self.num_groups, + residual=residual, + eps=self.eps, + is_rms_norm=self.is_rms_norm, + prenorm=prenorm, + upcast=True, + ) + + +__all__ = [ + 'GroupNorm', + 'GroupNormLinear', + 'GroupNormRef', + 'LayerNorm', + 'LayerNormLinear', + 'NormParallel', + 'RMSNorm', + 'RMSNormLinear', + 'group_norm', + 'group_norm_linear', + 'group_norm_ref', + 'layer_norm', + 'layer_norm_bwd', + 'layer_norm_fwd', + 'layer_norm_linear', + 'layer_norm_ref', + 'rms_norm', + 'rms_norm_linear', + 'rms_norm_ref', +] diff --git a/build/torch-cuda/modules/norm/layernorm_gated.py b/build/torch-cuda/modules/norm/layernorm_gated.py new file mode 100644 index 0000000000000000000000000000000000000000..a12370a12eb3a0a8fd8e46c2eaf9d21a370a93de --- /dev/null +++ b/build/torch-cuda/modules/norm/layernorm_gated.py @@ -0,0 +1,565 @@ +# 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 +# +# Copyright (c) 2024, Tri Dao. +# +# based on the Triton LayerNorm tutorial: https://triton-lang.org/main/getting-started/tutorials/05-layer-norm.html +# accumulate affine gradients in registers; wide dimensions can cause register spilling. + +import math + +import torch +import torch.nn as nn +import torch.nn.functional as F +import triton +import triton.language as tl +from einops import rearrange + +from ...utils import get_multiprocessor_count, input_guard + + +@triton.heuristics({"HAS_BIAS": lambda args: args["b"] is not None, "HAS_GATE": lambda args: args["g"] is not None}) +@triton.jit(do_not_specialize=['T']) +def layer_norm_fwd_kernel_group( + x, + y, + w, + b, + g, + mean, + rstd, + stride_x_row, + stride_y_row, + stride_g_row, + T, + D, + eps, + BD: tl.constexpr, + HAS_BIAS: tl.constexpr, + HAS_GATE: tl.constexpr, + NORM_BEFORE_GATE: tl.constexpr, + IS_RMS_NORM: tl.constexpr, +): + i_t = tl.program_id(0).to(tl.int64) + i_g = tl.program_id(1).to(tl.int64) + x += i_t * stride_x_row + i_g * D + y += i_t * stride_y_row + i_g * D + if HAS_GATE: + g += i_t * stride_g_row + i_g * D + if not IS_RMS_NORM: + mean += i_g * T + rstd += i_g * T + w += i_g * D + if HAS_BIAS: + b += i_g * D + + o_d = tl.arange(0, BD) + b_x = tl.load(x + o_d, mask=o_d < D, other=0.).to(tl.float32) + if HAS_GATE and not NORM_BEFORE_GATE: + b_g = tl.load(g + o_d, mask=o_d < D).to(tl.float32) + b_x *= b_g * tl.sigmoid(b_g) + 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(o_d < D, b_x - b_mean, 0.) + b_var = tl.sum(b_xbar * b_xbar, axis=0) / D + else: + b_xbar = tl.where(o_d < D, b_x, 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) + + m_d = o_d < D + 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_xhat = (b_x - b_mean) * b_rstd if not IS_RMS_NORM else b_x * b_rstd + b_y = b_xhat * b_w + b_b if HAS_BIAS else b_xhat * b_w + if HAS_GATE and NORM_BEFORE_GATE: + b_g = tl.load(g + o_d, mask=m_d).to(tl.float32) + b_y *= b_g * tl.sigmoid(b_g) + + tl.store(y + o_d, b_y, mask=m_d) + + +@triton.heuristics({ + "HAS_BIAS": lambda args: args["b"] is not None, + "HAS_GATE": lambda args: args["g"] is not None, + "RECOMPUTE_OUTPUT": lambda args: args["y"] is not None, +}) +@triton.jit(do_not_specialize=['T']) +def layer_norm_bwd_kernel_group( + x, + w, + b, + g, + y, + dy, + dx, + dw, + db, + dg, + mean, + rstd, + stride_x_row, + stride_g_row, + stride_y_row, + stride_dy_row, + stride_dx_row, + stride_dg_row, + stride_dw_row, + stride_db_row, + T, + D, + eps, + BS, + NORM_BEFORE_GATE: tl.constexpr, + IS_RMS_NORM: tl.constexpr, + HAS_BIAS: tl.constexpr, + HAS_GATE: tl.constexpr, + RECOMPUTE_OUTPUT: tl.constexpr, + BD: tl.constexpr, +): + i_s = tl.program_id(0).to(tl.int64) + i_g = tl.program_id(1).to(tl.int64) + bos = i_s * BS + o_d = tl.arange(0, BD) + m_d = o_d < D + x += bos * stride_x_row + i_g * D + if HAS_GATE: + g += bos * stride_g_row + i_g * D + dg += bos * stride_dg_row + i_g * D + dy += bos * stride_dy_row + i_g * D + dx += bos * stride_dx_row + i_g * D + if RECOMPUTE_OUTPUT: + y += bos * stride_y_row + i_g * D + if not IS_RMS_NORM: + mean += i_g * T + rstd += i_g * T + w += i_g * D + b_w = tl.load(w + o_d, mask=m_d).to(tl.float32) + if (RECOMPUTE_OUTPUT or HAS_GATE) and HAS_BIAS: + b += i_g * D + b_b = tl.load(b + o_d, mask=m_d, other=0.).to(tl.float32) + b_dw = tl.zeros((BD,), dtype=tl.float32) + if HAS_BIAS: + b_db = tl.zeros((BD,), dtype=tl.float32) + eos = min((i_s + 1) * BS, T) + for i_t in range(bos, eos): + b_x = tl.load(x + o_d, mask=m_d, other=0).to(tl.float32) + b_dy = tl.load(dy + o_d, mask=m_d, other=0).to(tl.float32) + if not IS_RMS_NORM: + b_mean = tl.load(mean + i_t) + if HAS_GATE and not NORM_BEFORE_GATE: + b_g = tl.load(g + o_d, mask=m_d, other=0.).to(tl.float32) + b_x_og = b_x + b_x = b_x_og * b_g * tl.sigmoid(b_g) + 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(m_d, b_xhat, 0.) + if HAS_GATE and NORM_BEFORE_GATE: + b_g = tl.load(g + o_d, mask=m_d, other=0.).to(tl.float32) + b_sigmoid_g = tl.sigmoid(b_g) + b_y = b_xhat * b_w + b_b if HAS_BIAS else b_xhat * b_w + if RECOMPUTE_OUTPUT: + tl.store(y + o_d, b_y * b_g * b_sigmoid_g, mask=m_d) + b_dg = b_dy * b_y * b_sigmoid_g * (1 + b_g * (1 - b_sigmoid_g)) + tl.store(dg + o_d, b_dg, mask=m_d) + b_dy *= b_g * b_sigmoid_g + else: + if RECOMPUTE_OUTPUT: + b_y = b_xhat * b_w + b_b if HAS_BIAS else b_xhat * b_w + tl.store(y + o_d, b_y, mask=m_d) + b_wdy = b_w * b_dy + b_c1 = tl.sum(b_xhat * b_wdy, axis=0) / D + if not IS_RMS_NORM: + b_c2 = tl.sum(b_wdy, axis=0) / D + b_dx = (b_wdy - (b_xhat * b_c1 + b_c2)) * b_rstd + else: + b_dx = (b_wdy - b_xhat * b_c1) * b_rstd + b_dw += b_dy * b_xhat + if HAS_BIAS: + b_db += b_dy + if HAS_GATE and not NORM_BEFORE_GATE: + b_sigmoid_g = tl.sigmoid(b_g) + b_dg = b_dx * b_x_og * b_sigmoid_g * (1 + b_g * (1 - b_sigmoid_g)) + tl.store(dg + o_d, b_dg, mask=m_d) + b_dx *= b_g * b_sigmoid_g + + tl.store(dx + o_d, b_dx, mask=m_d) + + x += stride_x_row + if HAS_GATE: + g += stride_g_row + dg += stride_dg_row + if RECOMPUTE_OUTPUT: + y += stride_y_row + dy += stride_dy_row + dx += stride_dx_row + tl.store(dw + i_s * stride_dw_row + i_g * D + o_d, b_dw, mask=m_d) + if HAS_BIAS: + tl.store(db + i_s * stride_db_row + i_g * D + o_d, b_db, mask=m_d) + + +def layer_norm_fwd( + x: torch.Tensor, + weight: torch.Tensor, + bias: torch.Tensor, + eps: float, + z: torch.Tensor | None = None, + out: torch.Tensor | None = None, + group_size: int | None = None, + norm_before_gate: bool = True, + is_rms_norm: bool = False, +): + T, D = x.shape + if group_size is None: + group_size = D + assert D % group_size == 0 + G = D // group_size + assert x.stride(-1) == 1 + if z is not None: + assert z.stride(-1) == 1 + assert z.shape == (T, D) + assert weight.shape == (D,) + assert weight.stride(-1) == 1 + if bias is not None: + assert bias.stride(-1) == 1 + assert bias.shape == (D,) + + if out is not None: + assert out.shape == x.shape + else: + out = torch.empty_like(x) + assert out.stride(-1) == 1 + mean = torch.empty((G * T,), dtype=torch.float32, device=x.device) if not is_rms_norm else None + rstd = torch.empty((G * T,), dtype=torch.float32, device=x.device) + + MAX_FUSED_SIZE = 65536 // x.element_size() + BD = min(MAX_FUSED_SIZE, triton.next_power_of_2(group_size)) + if group_size > BD: + raise RuntimeError("This layer norm doesn't support feature dim >= 64KB.") + num_warps = min(max(BD // 256, 1), 8) + grid = (T, G) + layer_norm_fwd_kernel_group[grid]( + x=x, + y=out, + w=weight, + b=bias, + g=z, + mean=mean, + rstd=rstd, + stride_x_row=x.stride(0), + stride_y_row=out.stride(0), + stride_g_row=z.stride(0) if z is not None else 0, + T=T, + D=group_size, + eps=eps, + BD=BD, + NORM_BEFORE_GATE=norm_before_gate, + IS_RMS_NORM=is_rms_norm, + num_warps=num_warps, + ) + return out, mean, rstd + + +def layer_norm_bwd( + dy: torch.Tensor, + x: torch.Tensor, + weight: torch.Tensor, + bias: torch.Tensor, + eps: float, + mean: torch.Tensor, + rstd: torch.Tensor, + z: torch.Tensor | None = None, + group_size: int | None = None, + norm_before_gate: bool = True, + is_rms_norm: bool = False, + recompute_output: bool = False, + dz: torch.Tensor | None = None, + out: torch.Tensor | None = None, +): + T, D = x.shape + if group_size is None: + group_size = D + assert D % group_size == 0 + G = D // group_size + assert x.stride(-1) == 1 + assert dy.stride(-1) == 1 + assert dy.shape == (T, D) + if z is not None: + assert z.stride(-1) == 1 + assert z.shape == (T, D) + assert weight.shape == (D,) + assert weight.stride(-1) == 1 + if bias is not None: + assert bias.stride(-1) == 1 + assert bias.shape == (D,) + + dx = torch.empty_like(x) + if dz is not None: + assert z is not None + assert dz.shape == z.shape + assert dz.stride(-1) == 1 + else: + dz = torch.empty_like(z) if z is not None else None + if recompute_output: + if out is None: + out = torch.empty_like(x) + assert out.shape == x.shape + + MAX_FUSED_SIZE = 65536 // x.element_size() + BD = min(MAX_FUSED_SIZE, triton.next_power_of_2(group_size)) + if group_size > BD: + raise RuntimeError("This layer norm doesn't support feature dim >= 64KB.") + num_warps = min(max(BD // 256, 1), 8) + sm_count = get_multiprocessor_count(x.device.index) + # use more programs for small groups to keep enough warps resident. + NS = math.ceil(sm_count * math.ceil(4 / num_warps) / G) + _dw = torch.empty((NS, D), dtype=torch.float32, device=weight.device) + _db = torch.empty((NS, D), dtype=torch.float32, device=bias.device) if bias is not None else None + BS = math.ceil(T / NS) + grid = (NS, G) + layer_norm_bwd_kernel_group[grid]( + x=x, + w=weight, + b=bias, + g=z, + y=out if recompute_output else None, + dy=dy, + dx=dx, + dw=_dw, + db=_db, + dg=dz, + mean=mean, + rstd=rstd, + stride_x_row=x.stride(0), + stride_g_row=z.stride(0) if z is not None else 0, + stride_y_row=0 if not recompute_output else out.stride(0), + stride_dy_row=dy.stride(0), + stride_dx_row=dx.stride(0), + stride_dg_row=dz.stride(0) if dz is not None else 0, + stride_dw_row=_dw.stride(0), + stride_db_row=_db.stride(0) if _db is not None else 0, + T=T, + D=group_size, + eps=eps, + BS=BS, + NORM_BEFORE_GATE=norm_before_gate, + IS_RMS_NORM=is_rms_norm, + BD=BD, + num_warps=num_warps, + ) + dw = _dw.sum(0).to(weight.dtype) + db = _db.sum(0).to(bias.dtype) if bias is not None else None + return (dx, dw, db, dz) if not recompute_output else (dx, dw, db, dz, out) + + +class LayerNormFn(torch.autograd.Function): + + @staticmethod + @input_guard + def forward( + ctx, + x: torch.Tensor, + weight: torch.Tensor, + bias: torch.Tensor | None, + z: torch.Tensor | None = None, + eps: float = 1e-6, + group_size: int | None = None, + norm_before_gate: bool = True, + is_rms_norm: bool = False, + ): + """Apply normalization before or after the optional SiLU gate.""" + + x_shape_og = x.shape + + x = x.reshape(-1, x.shape[-1]) + if x.stride(-1) != 1: + x = x.contiguous() + if z is not None: + assert z.shape == x_shape_og + z = z.reshape(-1, z.shape[-1]) + if z.stride(-1) != 1: + z = z.contiguous() + weight = weight.contiguous() + if bias is not None: + bias = bias.contiguous() + y, mean, rstd = layer_norm_fwd( + x=x, + weight=weight, + bias=bias, + eps=eps, + z=z, + group_size=group_size, + norm_before_gate=norm_before_gate, + is_rms_norm=is_rms_norm, + ) + ctx.save_for_backward(x, weight, bias, mean, rstd, z) + ctx.x_shape_og = x_shape_og + ctx.eps = eps + ctx.group_size = group_size + ctx.norm_before_gate = norm_before_gate + ctx.is_rms_norm = is_rms_norm + return y.reshape(x_shape_og) + + @staticmethod + @input_guard + def backward(ctx, dy): + x, weight, bias, mean, rstd, z = ctx.saved_tensors + dy = dy.reshape(-1, dy.shape[-1]) + if dy.stride(-1) != 1: + dy = dy.contiguous() + assert dy.shape == x.shape + dx, dw, db, dz = layer_norm_bwd( + dy=dy, + x=x, + weight=weight, + bias=bias, + eps=ctx.eps, + mean=mean, + rstd=rstd, + z=z, + group_size=ctx.group_size, + norm_before_gate=ctx.norm_before_gate, + is_rms_norm=ctx.is_rms_norm, + ) + dx = dx.reshape(ctx.x_shape_og) + dz = dz.reshape(ctx.x_shape_og) if dz is not None else None + return dx, dw, db, dz, None, None, None, None + + +def layernorm_fn( + x: torch.Tensor, + weight: torch.Tensor, + bias: torch.Tensor | None, + z: torch.Tensor | None = None, + eps: float = 1e-6, + group_size: int | None = None, + norm_before_gate: bool = True, + is_rms_norm: bool = False, +) -> torch.Tensor: + return LayerNormFn.apply(x, weight, bias, z, eps, group_size, norm_before_gate, is_rms_norm) + + +def rmsnorm_fn( + x: torch.Tensor, + weight: torch.Tensor, + bias: torch.Tensor | None, + z: torch.Tensor | None = None, + eps: float = 1e-6, + group_size: int | None = None, + norm_before_gate: bool = True, +) -> torch.Tensor: + return LayerNormFn.apply(x, weight, bias, z, eps, group_size, norm_before_gate, True) + + +def rms_norm_ref(x, weight, bias, z=None, eps=1e-6, group_size=None, norm_before_gate=True, upcast=True): + dtype = x.dtype + weight = weight.float() + bias = bias.float() if bias is not None else None + if upcast: + x = x.float() + z = z.float() if z is not None else z + if z is not None and not norm_before_gate: + x = x * F.silu(z) + if group_size is None: + rstd = 1 / torch.sqrt((x.square()).mean(dim=-1, keepdim=True) + eps) + out = (x * rstd * weight) + bias if bias is not None else (x * rstd * weight) + else: + x_group = rearrange(x, "... (g d) -> ... g d", d=group_size) + rstd = 1 / torch.sqrt((x_group.square()).mean(dim=-1, keepdim=True) + eps) + out = rearrange(x_group * rstd, "... g d -> ... (g d)") * weight + if bias is not None: + out = out + bias + if z is not None and norm_before_gate: + out *= F.silu(z) + return out.to(dtype) + + +class LayerNormGated(nn.Module): + + def __init__( + self, + hidden_size: int, + eps: float = 1e-5, + group_size: int | None = None, + norm_before_gate: bool = True, + device: torch.device | None = None, + dtype: torch.dtype | None = None, + ): + """If group_size is not None, we do GroupNorm with each group having group_size elements. + group_size=None is equivalent to group_size=hidden_size (i.e. there's only 1 group). + """ + + factory_kwargs = {"device": device, "dtype": dtype} + super().__init__() + self.eps = eps + self.weight = nn.Parameter(torch.empty(hidden_size, **factory_kwargs)) + self.bias = nn.Parameter(torch.empty(hidden_size, **factory_kwargs)) + self.group_size = group_size + self.norm_before_gate = norm_before_gate + self.reset_parameters() + + def reset_parameters(self): + torch.nn.init.ones_(self.weight) + torch.nn.init.zeros_(self.bias) + + def forward(self, x: torch.Tensor, z: torch.Tensor | None = None) -> torch.Tensor: + """Apply normalization before or after the optional SiLU gate.""" + return layernorm_fn( + x=x, + weight=self.weight, + bias=self.bias, + z=z, + eps=self.eps, + group_size=self.group_size, + norm_before_gate=self.norm_before_gate, + ) + + +class RMSNormGated(nn.Module): + + def __init__( + self, + hidden_size: int, + eps: float = 1e-5, + group_size: int | None = None, + norm_before_gate: bool = False, + device: torch.device | None = None, + dtype: torch.dtype | None = None, + ): + """If group_size is not None, we do GroupNorm with each group having group_size elements. + group_size=None is equivalent to group_size=hidden_size (i.e. there's only 1 group). + """ + factory_kwargs = {"device": device, "dtype": dtype} + super().__init__() + self.eps = eps + self.weight = nn.Parameter(torch.empty(hidden_size, **factory_kwargs)) + self.register_parameter("bias", None) + self.group_size = group_size + self.norm_before_gate = norm_before_gate + self.reset_parameters() + + def reset_parameters(self): + torch.nn.init.ones_(self.weight) + + def forward(self, x: torch.Tensor, z: torch.Tensor | None = None) -> torch.Tensor: + """Apply normalization before or after the optional SiLU gate.""" + return rmsnorm_fn( + x=x, + weight=self.weight, + bias=self.bias, + z=z, + eps=self.eps, + group_size=self.group_size, + norm_before_gate=self.norm_before_gate, + ) + + +__all__ = ['LayerNormGated', 'RMSNormGated', 'layer_norm_bwd', 'layer_norm_fwd', 'layernorm_fn', 'rms_norm_ref', 'rmsnorm_fn'] diff --git a/build/torch-cuda/modules/norm/layernorm_quant.py b/build/torch-cuda/modules/norm/layernorm_quant.py new file mode 100644 index 0000000000000000000000000000000000000000..ba4a1cc74bbb0422a08b4f9ddde9be2709003121 --- /dev/null +++ b/build/torch-cuda/modules/norm/layernorm_quant.py @@ -0,0 +1,428 @@ +# 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 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.functional as F +import triton +import triton.language as tl + +from ...utils import ( + IS_AMD, + autocast_custom_bwd, + autocast_custom_fwd, + autotune_cache_kwargs, + get_multiprocessor_count, + input_guard, +) + +NUM_WARPS_AUTOTUNE = [1, 2, 4, 8, 16] if IS_AMD else [1, 2, 4, 8, 16, 32] + + +def weight_quant(w: torch.Tensor) -> torch.Tensor: + """Quantize weights to 1.58 bits per tensor and return dequantized values.""" + scale = 1.0 / w.abs().mean().clamp_(min=1e-5) + y = (w * scale).round().clamp_(-1, 1) / scale + return y + + +@triton.jit +def activation_quant_fwd_kernel(b_x): + b_scale = 127.0 / tl.maximum(tl.max(tl.abs(b_x), 0), 1e-5) + b_scaled = b_x * b_scale + # triton-ascend does not implement libdevice.round; use round-half-away-from-zero. + b_y = tl.where(b_scaled >= 0, tl.floor(b_scaled + 0.5), tl.ceil(b_scaled - 0.5)) + return tl.maximum(tl.minimum(b_y, 127), -128) / b_scale + + +@triton.autotune( + configs=[ + triton.Config({}, num_warps=num_warps) + for num_warps in NUM_WARPS_AUTOTUNE + ], + key=['D', 'HAS_RESIDUAL', 'STORE_RESIDUAL_OUT', 'IS_RMS_NORM', 'HAS_BIAS'], + **autotune_cache_kwargs, +) +@triton.jit +def layer_norm_quant_fwd_kernel( + x, + y, + w, + b, + res, + res_out, + mean, + rstd, + eps, + 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).to(tl.int64) + + o_d = tl.arange(0, BD).to(tl.int64) + m_d = o_d < D + + p_x = x + i_t * D + o_d + b_x = tl.load(p_x, mask=m_d, other=0.0).to(tl.float32) + if HAS_RESIDUAL: + p_res = res + i_t * D + o_d + b_x += tl.load(p_res, mask=m_d, other=0.0).to(tl.float32) + if STORE_RESIDUAL_OUT: + p_res_out = res_out + i_t * D + o_d + tl.store(p_res_out, b_x.to(p_res_out.dtype.element_ty), mask=m_d) + if not IS_RMS_NORM: + b_mean = tl.sum(b_x, axis=0) / D + p_mean = mean + i_t + tl.store(p_mean, b_mean.to(p_mean.dtype.element_ty)) + 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) + + p_rstd = rstd + i_t + tl.store(p_rstd, b_rstd.to(p_rstd.dtype.element_ty)) + + if HAS_WEIGHT: + b_w = tl.load(w + o_d, mask=m_d, other=0.0).to(tl.float32) + if HAS_BIAS: + b_b = tl.load(b + o_d, mask=m_d, other=0.0).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 + + b_y = activation_quant_fwd_kernel(b_y) + + p_y = y + i_t * D + o_d + tl.store(p_y, b_y.to(p_y.dtype.element_ty), mask=m_d) + + +@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=['D', 'HAS_DRESIDUAL', 'STORE_DRESIDUAL', 'IS_RMS_NORM', 'HAS_BIAS'], + **autotune_cache_kwargs, +) +@triton.jit(do_not_specialize=['T']) +def layer_norm_quant_bwd_kernel( + x, + w, + b, + y, + dy, + dx, + dw, + db, + dres, + dres_in, + mean, + rstd, + T, + D: tl.constexpr, + BS: tl.constexpr, + BD: 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).to(tl.int64) + + o_d = tl.arange(0, BD).to(tl.int64) + m_d = o_d < D + + if HAS_WEIGHT: + b_w = tl.load(w + o_d, mask=m_d, other=0.0).to(tl.float32) + b_dw = tl.zeros((BD,), dtype=tl.float32) + if RECOMPUTE_OUTPUT and HAS_BIAS: + b_b = tl.load(b + o_d, mask=m_d, other=0.0).to(tl.float32) + if HAS_BIAS: + b_db = tl.zeros((BD,), dtype=tl.float32) + + for i_t in range(i_s * BS, min((i_s + 1) * BS, T)): + p_x = x + i_t * D + o_d + p_dy = dy + i_t * D + o_d + b_x = tl.load(p_x, mask=m_d, other=0.0).to(tl.float32) + b_dy = tl.load(p_dy, mask=m_d, other=0.0).to(tl.float32) + if not IS_RMS_NORM: + p_mean = mean + i_t + b_mean = tl.load(p_mean) + p_rstd = rstd + i_t + b_rstd = tl.load(p_rstd) + + b_x_hat = (b_x - b_mean) * b_rstd if not IS_RMS_NORM else b_x * b_rstd + b_x_hat = tl.where(m_d, b_x_hat, 0.0) + if RECOMPUTE_OUTPUT: + b_y = b_x_hat * b_w if HAS_WEIGHT else b_x_hat + if HAS_BIAS: + b_y = b_y + b_b + + b_y = activation_quant_fwd_kernel(b_y) + + p_y = y + i_t * D + o_d + tl.store(p_y, b_y.to(p_y.dtype.element_ty), mask=m_d) + b_wdy = b_dy + if HAS_WEIGHT: + b_wdy = b_dy * b_w + b_dw += b_dy * b_x_hat + if HAS_BIAS: + b_db += b_dy + if not IS_RMS_NORM: + b_c1 = tl.sum(b_x_hat * b_wdy, axis=0) / D + b_c2 = tl.sum(b_wdy, axis=0) / D + b_dx = (b_wdy - (b_x_hat * b_c1 + b_c2)) * b_rstd + else: + b_c1 = tl.sum(b_x_hat * b_wdy, axis=0) / D + b_dx = (b_wdy - b_x_hat * b_c1) * b_rstd + if HAS_DRESIDUAL: + p_dres = dres + i_t * D + o_d + b_dx += tl.load(p_dres, mask=m_d, other=0.0).to(tl.float32) + + if STORE_DRESIDUAL: + p_dres_in = dres_in + i_t * D + o_d + tl.store(p_dres_in, b_dx.to(p_dres_in.dtype.element_ty), mask=m_d) + + p_dx = dx + i_t * D + o_d + tl.store(p_dx, b_dx.to(p_dx.dtype.element_ty), mask=m_d) + + if HAS_WEIGHT: + p_dw = dw + i_s * D + o_d + tl.store(p_dw, b_dw.to(p_dw.dtype.element_ty), mask=m_d) + if HAS_BIAS: + p_db = db + i_s * D + o_d + tl.store(p_db, b_db.to(p_db.dtype.element_ty), mask=m_d) + + +def layer_norm_quant_fwd( + x: torch.Tensor, + weight: torch.Tensor | None, + bias: torch.Tensor | None, + eps: float, + 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: + 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,) + + 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.float32, device=x.device) if not is_rms_norm else None + rstd = torch.empty((T,), dtype=torch.float32, 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.") + layer_norm_quant_fwd_kernel[(T,)]( + x=x, + y=y, + w=weight, + b=bias, + res=residual, + res_out=res_out, + mean=mean, + rstd=rstd, + eps=eps, + 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, + ) + return y, mean, rstd, res_out if res_out is not None else x + + +def layer_norm_quant_bwd( + dy: torch.Tensor, + x: torch.Tensor, + weight: torch.Tensor | None, + bias: torch.Tensor | None, + mean: torch.Tensor | None, + rstd: torch.Tensor, + dres: 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 dy.shape == (T, D) + if dres is not None: + assert dres.shape == (T, D) + if weight is not None: + 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) + 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 + + 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.") + NS = get_multiprocessor_count(x.device.index) + BS = triton.cdiv(T, NS) + + dw = torch.empty((NS, D), dtype=torch.float32, device=weight.device) if weight is not None else None + db = torch.empty((NS, D), dtype=torch.float32, device=bias.device) if bias is not None else None + grid = (NS,) + layer_norm_quant_bwd_kernel[grid]( + x=x, + w=weight, + b=bias, + y=y, + dy=dy, + dx=dx, + dw=dw, + db=db, + dres=dres, + dres_in=dres_in, + mean=mean, + rstd=rstd, + T=T, + D=D, + BS=BS, + BD=BD, + 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.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 + # the residual shares the input gradient when their dtypes match + 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) + + +class LayerNormLinearQuantFunction(torch.autograd.Function): + + @staticmethod + @input_guard + @autocast_custom_fwd + def forward( + ctx, + 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, + ): + x_shape_og = x.shape + 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, res_out = layer_norm_quant_fwd( + x=x, + weight=norm_weight, + bias=norm_bias, + eps=eps, + residual=residual, + out_dtype=None if not torch.is_autocast_enabled() else torch.get_autocast_dtype('cuda'), + residual_dtype=residual_dtype, + is_rms_norm=is_rms_norm, + ) + y = y.reshape(x_shape_og) + dtype = torch.get_autocast_dtype('cuda') 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(input=y.to(linear_weight.dtype), weight=linear_weight, bias=linear_bias) + # recompute y in backward to save memory + ctx.save_for_backward(res_out, norm_weight, norm_bias, linear_weight, mean, rstd) + ctx.x_shape_og = x_shape_og + 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, res_out.reshape(x_shape_og)) + + @staticmethod + @input_guard + @autocast_custom_bwd + def backward(ctx, dout: torch.Tensor, *args): + x, norm_weight, norm_bias, linear_weight, mean, rstd = ctx.saved_tensors + dout = dout.reshape(-1, dout.shape[-1]) + dy = F.linear(input=dout, weight=linear_weight.t()) + dlinear_bias = None if ctx.linear_bias_is_none else dout.sum(0) + assert dy.shape == x.shape + if ctx.prenorm: + dres = args[0] + dres = dres.reshape(-1, dres.shape[-1]) + assert dres.shape == x.shape + else: + dres = None + dx, dnorm_weight, dnorm_bias, dres_in, y = layer_norm_quant_bwd( + dy=dy, + x=x, + weight=norm_weight, + bias=norm_bias, + mean=mean, + rstd=rstd, + dres=dres, + 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), + 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, + ) diff --git a/build/torch-cuda/modules/parallel.py b/build/torch-cuda/modules/parallel.py index 150c4c4e17f92655f646fcf85b2b7a9569acf293..bdea2964a71cf8e59490ba16f01ff5304a52edad 100644 --- a/build/torch-cuda/modules/parallel.py +++ b/build/torch-cuda/modules/parallel.py @@ -7,7 +7,7 @@ import torch.nn as nn from torch.distributed import DeviceMesh -from torch.distributed.tensor import distribute_module +from torch.distributed.tensor import Replicate, distribute_module from torch.distributed.tensor.parallel import ParallelStyle from torch.distributed.tensor.placement_types import Placement @@ -18,9 +18,14 @@ except (ImportError, AttributeError): class PrepareModuleWeight(ParallelStyle): + """Wrap parameters that already match their layout on a one-dimensional mesh. + + Replicated values must match across ranks; sharded values must already be local slices. + """ + def __init__(self, *, layouts: Placement | None = None): super().__init__() - self.layouts = layouts + self.layouts = Replicate() if layouts is None else layouts def _replicate_module_fn( self, @@ -28,9 +33,10 @@ class PrepareModuleWeight(ParallelStyle): module: nn.Module, device_mesh: DeviceMesh, ): - for p_name, param in module.named_parameters(): + for p_name, param in module.named_parameters(recurse=False): replicated_param = nn.Parameter( DTensor.from_local(param, device_mesh, [self.layouts], run_check=False), + requires_grad=param.requires_grad, ) module.register_parameter(p_name, replicated_param) diff --git a/build/torch-xpu/ops/gated_delta_product/__init__.py b/build/torch-cuda/modules/rotary/__init__.py similarity index 71% rename from build/torch-xpu/ops/gated_delta_product/__init__.py rename to build/torch-cuda/modules/rotary/__init__.py index 6f118dae595ba3d76f6cdb437477c926ec9bc9ff..82bdae337d1236c0bbe551fdb7f970901e08d2c3 100644 --- a/build/torch-xpu/ops/gated_delta_product/__init__.py +++ b/build/torch-cuda/modules/rotary/__init__.py @@ -5,8 +5,7 @@ # For a list of all contributors, visit: # https://github.com/fla-org/flash-linear-attention/graphs/contributors -from .chunk import chunk_gated_delta_product +from ...modules.rotary.ops import RotaryEmbedding, rotary_embedding + +__all__ = ['RotaryEmbedding', 'rotary_embedding'] -__all__ = [ - "chunk_gated_delta_product", -] diff --git a/build/torch-cuda/modules/rotary.py b/build/torch-cuda/modules/rotary/ops.py similarity index 69% rename from build/torch-cuda/modules/rotary.py rename to build/torch-cuda/modules/rotary/ops.py index ddf10adcfc40db494ef4848c58a30d97601352e9..ee4185dd69a43ba98ed9a59cf39e614ca3025ec8 100644 --- a/build/torch-cuda/modules/rotary.py +++ b/build/torch-cuda/modules/rotary/ops.py @@ -11,30 +11,13 @@ import triton import triton.language as tl from einops import rearrange, repeat -from ..modules.backends import dispatch -from ..ops.utils import prepare_chunk_indices -from ..utils import IS_AMD, autotune_cache_kwargs, get_multiprocessor_count, input_guard +from ...backends import dispatch +from ...ops.utils import prepare_chunk_indices +from ...utils import IS_AMD, autotune_cache_kwargs, get_multiprocessor_count, input_guard NUM_WARPS_AUTOTUNE = [2, 4, 8, 16] if IS_AMD else [2, 4, 8, 16, 32] -def rotate_half(x, interleaved=False): - if not interleaved: - x1, x2 = x.chunk(2, dim=-1) - return torch.cat((-x2, x1), dim=-1) - else: - x1, x2 = x[..., ::2], x[..., 1::2] - return rearrange(torch.stack((-x2, x1), dim=-1), '... d two -> ... (d two)', two=2) - - -def rotary_embedding_ref(x, cos, sin, interleaved=False): - ro_dim = cos.shape[-1] * 2 - assert ro_dim <= x.shape[-1] - cos = repeat(cos, '... d -> ... 1 (2 d)' if not interleaved else '... d -> ... 1 (d 2)') - sin = repeat(sin, '... d -> ... 1 (2 d)' if not interleaved else '... d -> ... 1 (d 2)') - return torch.cat([x[..., :ro_dim] * cos + rotate_half(x[..., :ro_dim], interleaved) * sin, x[..., ro_dim:]], -1) - - @triton.autotune( configs=[ triton.Config({}, num_warps=num_warps, num_stages=num_stages) @@ -66,11 +49,11 @@ def rotary_embedding_kernel( INTERLEAVED: tl.constexpr, CONJUGATE: tl.constexpr, ): - i_t, i_b, i_h = tl.program_id(0), tl.program_id(1), tl.program_id(2) + i_t, i_b, i_h = tl.program_id(0).to(tl.int64), tl.program_id(1).to(tl.int64), 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), tl.load(cu_seqlens + i_n + 1) + i_n, i_t = tl.load(chunk_indices + i_t * 2).to(tl.int32), tl.load(chunk_indices + i_t * 2 + 1).to(tl.int64) + bos, eos = tl.load(cu_seqlens + i_n).to(tl.int64), tl.load(cu_seqlens + i_n + 1).to(tl.int64) T = eos - bos x = x + bos * H*D + i_h * D y = y + bos * H*D + i_h * D @@ -105,7 +88,6 @@ def rotary_embedding_kernel( b_sin = -b_sin b_o0 = b_x0 * b_cos - b_x1 * b_sin b_o1 = b_x0 * b_sin + b_x1 * b_cos - # write back result 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) @@ -117,7 +99,8 @@ def rotary_embedding_kernel( # Then we do the calculation and use tl.where to pick put the right outputs for the even # and for the odd indices. o_d = tl.arange(0, BD) - o_d_swap = o_d + ((o_d + 1) % 2) * 2 - 1 # 1, 0, 3, 2, 5, 4, ... + # pair adjacent channels: 1, 0, 3, 2, 5, 4, ... + 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, :] @@ -138,7 +121,27 @@ def rotary_embedding_kernel( tl.store(p_y, b_y, mask=mask) -@dispatch('modules') +def rotate_half(x, interleaved=False): + if not interleaved: + x1, x2 = x.chunk(2, dim=-1) + return torch.cat((-x2, x1), dim=-1) + else: + x1, x2 = x[..., ::2], x[..., 1::2] + return rearrange(torch.stack((-x2, x1), dim=-1), '... d two -> ... (d two)', two=2) + + +def rotary_embedding_ref(x, cos, sin, interleaved=False): + ro_dim = cos.shape[-1] * 2 + assert ro_dim <= x.shape[-1] + cos = repeat(cos, '... d -> ... 1 (2 d)' if not interleaved else '... d -> ... 1 (d 2)') + sin = repeat(sin, '... d -> ... 1 (2 d)' if not interleaved else '... d -> ... 1 (d 2)') + return torch.cat( + [x[..., :ro_dim] * cos + rotate_half(x=x[..., :ro_dim], interleaved=interleaved) * sin, x[..., ro_dim:]], + -1, + ) + + +@dispatch def rotary_embedding_fwdbwd( x: torch.Tensor, cos: torch.Tensor, @@ -150,16 +153,31 @@ def rotary_embedding_fwdbwd( conjugate: bool = False, chunk_indices: torch.LongTensor | None = None, ) -> torch.Tensor: - """ + """Apply rotary embeddings or their conjugate for the backward pass. + Args: - x: [B, T, H, D]. - cos: [TR, R / 2] - sin: [TR, R / 2] - seqlen_offsets: integer or integer tensor of size [N] - cu_seqlens: [N + 1,] or None + x (torch.Tensor): + Input of shape `[B, T, H, D]`. + cos (torch.Tensor): + Cosine cache of shape `[TR, R // 2]`. + sin (torch.Tensor): + Sine cache with the same shape and dtype as `cos`. + seqlen_offsets (int | torch.Tensor, Optional): + Shared position offset or per-sequence offsets of shape `[N]`. Default: 0. + cu_seqlens (torch.Tensor, Optional): + Cumulative sequence lengths of shape `[N + 1]` for packed inputs. Default: `None`. + interleaved (bool, Optional): + Whether to rotate adjacent pairs instead of matching the two halves of each head. Default: `False`. + inplace (bool, Optional): + Whether to overwrite the input. Default: `False`. + conjugate (bool, Optional): + Whether to negate the sine cache for the backward pass. Default: `False`. + chunk_indices (torch.LongTensor, Optional): + Precomputed sequence chunk indices. Default: `None`. Returns: - y: [B, T, H, D] + y (torch.Tensor): + Rotated input of shape `[B, T, H, D]`. """ is_varlen = cu_seqlens is not None @@ -169,7 +187,9 @@ def rotary_embedding_fwdbwd( R2 = R * 2 assert D <= 256, "Only support D <= 256" - assert TR >= T, f"TR must be >= T, got {TR} and {T}" + assert R2 <= D, f"Rotary dimension must not exceed head dimension, got rotary_dim={R2} and head_dim={D}" + if not is_varlen: + 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}" @@ -177,6 +197,11 @@ def rotary_embedding_fwdbwd( if isinstance(seqlen_offsets, torch.Tensor): assert seqlen_offsets.shape == (N,) assert seqlen_offsets.dtype in [torch.int32, torch.int64] + sequence_lengths = T if not is_varlen else cu_seqlens[1:] - cu_seqlens[:-1] + torch._assert_async( + torch.all(sequence_lengths + seqlen_offsets <= TR), + f"Rotary cache is too short for tensor offsets, got {TR} positions", + ) else: assert seqlen_offsets + T <= TR @@ -193,15 +218,15 @@ def rotary_embedding_fwdbwd( grid = (NT, B, H) rotary_embedding_kernel[grid]( - x, - cos, - sin, - y, - cu_seqlens, - chunk_indices, - seqlen_offsets, - B=B, + x=x, + cos=cos, + sin=sin, + y=y, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + seq_offsets=seqlen_offsets, T=T, + B=B, H=H, D=D, R=R, @@ -232,9 +257,9 @@ class RotaryEmbeddingFunction(torch.autograd.Function): chunk_indices: torch.LongTensor | None = None, ): y = rotary_embedding_fwdbwd( - x, - cos, - sin, + x=x, + cos=cos, + sin=sin, seqlen_offsets=seqlen_offsets, cu_seqlens=cu_seqlens, interleaved=interleaved, @@ -242,7 +267,7 @@ class RotaryEmbeddingFunction(torch.autograd.Function): chunk_indices=chunk_indices, ) if isinstance(seqlen_offsets, int): - # Can't save int with save_for_backward + # can't save int with save_for_backward ctx.save_for_backward(cos, sin, cu_seqlens) ctx.seqlen_offsets = seqlen_offsets else: @@ -261,14 +286,13 @@ class RotaryEmbeddingFunction(torch.autograd.Function): cos, sin, cu_seqlens, seqlen_offsets = ctx.saved_tensors else: cos, sin, cu_seqlens = ctx.saved_tensors - # TD [2023-09-02]: For some reason Triton (2.0.0.post1) errors with - # "[CUDA]: invalid device context", and cloning makes it work. Idk why. Triton 2.1.0 works. + # cloning avoids an invalid-device-context error in Triton 2.0.0.post1. if not ctx.interleaved and not ctx.inplace: do = do.clone() dx = rotary_embedding_fwdbwd( - do, - cos, - sin, + x=do, + cos=cos, + sin=sin, seqlen_offsets=seqlen_offsets, cu_seqlens=cu_seqlens, interleaved=ctx.interleaved, @@ -289,49 +313,46 @@ def rotary_embedding( cu_seqlens: torch.Tensor | None = None, chunk_indices: torch.LongTensor | None = None, ): - """ + """Apply rotary embeddings, optionally using a different position offset for each sequence. + Args: - x: [B, T, H, D] - cos, sin: [TR, R//2] - interleaved: - If True, rotate pairs of even and odd dimensions (GPT-J style) instead of 1st half and 2nd half (GPT-NeoX style). - inplace: - If True, apply rotary embedding in-place. - seqlen_offsets: [N,] or int. - Each sequence in x is shifted by this amount. - Most commonly used in inference when we have KV cache. - cu_seqlens: [N + 1,] or None + x (torch.Tensor): + Input of shape `[B, T, H, D]`. + cos (torch.Tensor): + Cosine cache of shape `[TR, R // 2]`. + sin (torch.Tensor): + Sine cache with the same shape and dtype as `cos`. + interleaved (bool, Optional): + Whether to rotate adjacent pairs instead of matching the two halves of each head. Default: `False`. + inplace (bool, Optional): + Whether to overwrite the input. Default: `False`. + seqlen_offsets (int | torch.Tensor, Optional): + Shared position offset or per-sequence offsets of shape `[N]`. Default: 0. + cu_seqlens (torch.Tensor, Optional): + Cumulative sequence lengths of shape `[N + 1]` for packed inputs. Default: `None`. + chunk_indices (torch.LongTensor, Optional): + Precomputed sequence chunk indices. Default: `None`. Returns: - out: [B, T, H, D] + torch.Tensor: + Rotated input with the same shape and dtype as `x`. """ - return RotaryEmbeddingFunction.apply( - x, - cos, - sin, - interleaved, - inplace, - seqlen_offsets, - cu_seqlens, - chunk_indices, - ) + return RotaryEmbeddingFunction.apply(x, cos, sin, interleaved, inplace, seqlen_offsets, cu_seqlens, chunk_indices) class RotaryEmbedding(nn.Module): """ The rotary position embeddings from RoFormer_ (Su et. al). - A crucial insight from the method is that the query and keys are - transformed by rotation matrices which depend on the relative positions. + Queries and keys are transformed by rotation matrices that depend on their relative positions. - Other implementations are available in the Rotary Transformer repo_ and in - GPT-NeoX_, GPT-NeoX was an inspiration + Other implementations are available in the Rotary Transformer repo_ and in GPT-NeoX_, which inspired this implementation. .. _RoFormer: https://arxiv.org/abs/2104.09864 .. _repo: https://github.com/ZhuiyiTechnology/roformer .. _GPT-NeoX: https://github.com/EleutherAI/gpt-neox If scale_base is not None, this implements XPos (Sun et al., https://arxiv.org/abs/2212.10554). - A recommended value for scale_base is 512: https://github.com/HazyResearch/flash-attention/issues/96 + A recommended value for scale_base is 512: https://github.com/HazyResearch/flash-attention/issues/96. Reference: https://github.com/sunyt32/torchscale/blob/main/torchscale/component/xpos_relative_position.py """ @@ -344,21 +365,28 @@ class RotaryEmbedding(nn.Module): pos_idx_in_fp32: bool = True, device: torch.device | None = None, ): - """ - interleaved: - If True, rotate pairs of even and odd dimensions (GPT-J style) instead of 1st half and 2nd half (GPT-NeoX style). - pos_idx_in_fp32: - If True, the position indices [0.0, ..., seqlen - 1] are in fp32, otherwise they might be in lower precision. - This option was added because previously (before 2023-07-02), when we construct - the position indices, we use the dtype of self.inv_freq. - In most cases this would be fp32, but if the model is trained in pure bf16 (not mixed precision), then - self.inv_freq would be bf16, and the position indices are also in bf16. - Because of the limited precision of bf16 (e.g. 1995.0 is rounded to 2000.0), the - embeddings for some positions will coincide. - To maintain compatibility with models previously trained in pure bf16, we add this option. + """Initialize rotary position tables. + + Args: + dim (int): + Number of dimensions to rotate in each head. + base (float, Optional): + Base for the inverse frequencies. Default: 10000.0. + interleaved (bool, Optional): + Whether to rotate adjacent pairs (GPT-J) instead of matching the two halves (GPT-NeoX). Default: `False`. + scale_base (float, Optional): + Base for XPos scaling. Default: `None`. + pos_idx_in_fp32 (bool, Optional): + Whether to create position indices in fp32. Default: `True`. + Pure bf16 indices can round distinct positions to the same value (for example, 1995.0 to 2000.0). + Set to `False` for compatibility with models trained with bf16 position indices. + device (torch.device, Optional): + Device for the frequency buffers. Default: `None`. """ super().__init__() + assert dim % 2 == 0, f"Rotary dimension must be even, got {dim}" + self.dim = dim self.base = float(base) self.scale_base = scale_base @@ -399,10 +427,7 @@ class RotaryEmbedding(nn.Module): return s def _compute_inv_freq(self, device=None): - return 1.0 / ( - self.base - ** (torch.arange(0, self.dim, 2, device=device, dtype=torch.float32) / self.dim) - ) + return 1.0 / (self.base ** (torch.arange(0, self.dim, 2, device=device, dtype=torch.float32) / self.dim)) def _compute_scale(self, device=None): return (torch.arange(0, self.dim, 2, device=device, dtype=torch.float32) + 0.4 * self.dim) / (1.4 * self.dim) @@ -462,34 +487,49 @@ class RotaryEmbedding(nn.Module): max_seqlen: int | None = None, chunk_indices: torch.LongTensor | None = None, ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]: - """ - q: [B, T, H, D] - k: [B, T, H, D] - seqlen_offset: - [N] or int. - Each sequence in x is shifted by this amount. - Most commonly used in inference when we have KV cache. - cu_seqlens: [N + 1] or None - max_seqlen: int + """Rotate query and key tensors using cached position tables. + + Args: + q (torch.Tensor): + Query tensor of shape `[B, T, H, D]`. + k (torch.Tensor): + Key tensor of shape `[B, T, H_k, D]`. + seqlen_offset (int | torch.Tensor, Optional): + Shared position offset or per-sequence offsets of shape `[N]`. Default: 0. + cu_seqlens (torch.Tensor, Optional): + Cumulative sequence lengths of shape `[N + 1]` for packed inputs. Default: `None`. + max_seqlen (int, Optional): + Cache length used to initialize the rotary tables. Default: `None`. + Tensor offsets require this on the first call so cache sizing does not synchronize the device. + chunk_indices (torch.LongTensor, Optional): + Precomputed sequence chunk indices. Default: `None`. + + Returns: + tuple[torch.Tensor, torch.Tensor]: + Rotated query and key tensors. """ if max_seqlen is not None: self._update_cos_sin_cache(max_seqlen, device=q.device, dtype=q.dtype) elif isinstance(seqlen_offset, int): self._update_cos_sin_cache(q.shape[1] + seqlen_offset, device=q.device, dtype=q.dtype) + else: + assert self._cos_cached is not None, ( + "Tensor offsets require an initialized cache; pass max_seqlen on the first call" + ) if self.scale is None: q = rotary_embedding( - q, - self._cos_cached, - self._sin_cached, + x=q, + cos=self._cos_cached, + sin=self._sin_cached, interleaved=self.interleaved, seqlen_offsets=seqlen_offset, cu_seqlens=cu_seqlens, chunk_indices=chunk_indices, ) k = rotary_embedding( - k, - self._cos_cached, - self._sin_cached, + x=k, + cos=self._cos_cached, + sin=self._sin_cached, interleaved=self.interleaved, seqlen_offsets=seqlen_offset, cu_seqlens=cu_seqlens, @@ -498,18 +538,18 @@ class RotaryEmbedding(nn.Module): else: q = rotary_embedding( - q, - self._cos_cached, - self._sin_cached, + x=q, + cos=self._cos_cached, + sin=self._sin_cached, interleaved=self.interleaved, seqlen_offsets=seqlen_offset, cu_seqlens=cu_seqlens, chunk_indices=chunk_indices, ) k = rotary_embedding( - k, - self._cos_k_cached, - self._sin_k_cached, + x=k, + cos=self._cos_k_cached, + sin=self._sin_k_cached, interleaved=self.interleaved, seqlen_offsets=seqlen_offset, cu_seqlens=cu_seqlens, diff --git a/build/torch-cuda/modules/token_shift.py b/build/torch-cuda/modules/token_shift.py index 8712da2eaa0878bceb52d115c70d51ffae9633ce..fb8aa3a23c7ed97fc30dfa63bb00b1763548666e 100644 --- a/build/torch-cuda/modules/token_shift.py +++ b/build/torch-cuda/modules/token_shift.py @@ -10,12 +10,9 @@ import triton import triton.language as tl from ..ops.utils import prepare_chunk_indices -from ..utils import IS_AMD, IS_NPU, autotune_cache_kwargs, get_multiprocessor_count, input_guard, tensor_cache +from ..utils import IS_AMD, autotune_cache_kwargs, get_multiprocessor_count, input_guard, tensor_cache NUM_WARPS_AUTOTUNE = [2, 4, 8, 16] if IS_AMD else [2, 4, 8, 16, 32] -# Ascend Triton rejects 2-D grids whose product exceeds 65535 unless -# TRITON_ALL_BLOCKS_PARALLEL=1. Fall back to the long kernel instead. -_NPU_MAX_TRITON_GRID = 65535 def token_shift_ref( @@ -82,11 +79,11 @@ def token_shift_fwd_kernel_short( STORE_FINAL_STATE: tl.constexpr, IS_DECODE: tl.constexpr, ): - i_b, i_t = tl.program_id(0), tl.program_id(1) + i_b, i_t = tl.program_id(0).to(tl.int64), tl.program_id(1) if IS_VARLEN: i_n = i_b - bos, eos = tl.load(cu_seqlens + i_n).to(tl.int32), tl.load(cu_seqlens + i_n + 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) g_t = i_t + bos if g_t >= eos: @@ -130,6 +127,9 @@ def token_shift_fwd_kernel_short( tl.store(y + base_offset, delta, mask=m_d) else: tl.store(y + base_offset, -b_x, mask=m_d) + if STORE_FINAL_STATE: + if is_last_pos: + tl.store(cache_out + cache_offset, b_x, mask=m_d) return # Other positions: delta = prev - curr @@ -156,7 +156,7 @@ def token_shift_fwd_kernel_short( for num_warps in NUM_WARPS_AUTOTUNE for num_stages in [1, 2, 3] ], - key=['BD', 'NB'], + key=['BD'], **autotune_cache_kwargs, ) @triton.jit @@ -171,19 +171,18 @@ def token_shift_fwd_kernel_long( D: tl.constexpr, BD: tl.constexpr, BT: tl.constexpr, - NB: tl.constexpr, ND: tl.constexpr, IS_VARLEN: tl.constexpr, USE_INITIAL_STATE: tl.constexpr, STORE_FINAL_STATE: tl.constexpr, ): - i_dt, i_b = tl.program_id(0), tl.program_id(1) + i_dt, i_b = tl.program_id(0), tl.program_id(1).to(tl.int64) i_d, i_t = i_dt % ND, i_dt // ND 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) + tl.load(chunk_indices + i_t * 2 + 1).to(tl.int64) + bos, eos = tl.load(cu_seqlens + i_n).to(tl.int64), tl.load(cu_seqlens + i_n + 1).to(tl.int64) t_start = i_t * BT t_end = tl.minimum(t_start + BT, eos - bos) else: @@ -249,11 +248,11 @@ def token_shift_bwd_kernel_short( USE_INITIAL_STATE: tl.constexpr, HAS_DCACHE: tl.constexpr, ): - i_b, i_t = tl.program_id(0), tl.program_id(1) + i_b, i_t = tl.program_id(0).to(tl.int64), tl.program_id(1) if IS_VARLEN: i_n = i_b - bos, eos = tl.load(cu_seqlens + i_n).to(tl.int32), tl.load(cu_seqlens + i_n + 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) g_t = i_t + bos if g_t >= eos: return @@ -310,7 +309,7 @@ def token_shift_bwd_kernel_short( for num_warps in NUM_WARPS_AUTOTUNE for num_stages in [1, 2, 3] ], - key=['BD', 'NB'], + key=['BD'], **autotune_cache_kwargs, ) @triton.jit @@ -325,19 +324,18 @@ def token_shift_bwd_kernel_long( D: tl.constexpr, BD: tl.constexpr, BT: tl.constexpr, - NB: tl.constexpr, ND: tl.constexpr, IS_VARLEN: tl.constexpr, USE_INITIAL_STATE: tl.constexpr, HAS_DCACHE: tl.constexpr, ): - i_dt, i_b = tl.program_id(0), tl.program_id(1) + i_dt, i_b = tl.program_id(0), tl.program_id(1).to(tl.int64) i_d, i_t_blk = i_dt % ND, i_dt // ND if IS_VARLEN: i_n, i_t_blk = tl.load(chunk_indices + i_t_blk * 2).to(tl.int32), \ - tl.load(chunk_indices + i_t_blk * 2 + 1).to(tl.int32) - bos, eos = tl.load(cu_seqlens + i_n), tl.load(cu_seqlens + i_n + 1) + tl.load(chunk_indices + i_t_blk * 2 + 1).to(tl.int64) + bos, eos = tl.load(cu_seqlens + i_n).to(tl.int64), tl.load(cu_seqlens + i_n + 1).to(tl.int64) t_start = i_t_blk * BT t_end = tl.minimum(t_start + BT, eos - bos) else: @@ -394,8 +392,6 @@ def token_shift_fwd( N = B use_short_kernel = T <= 4096 - if IS_NPU and use_short_kernel and N * T > _NPU_MAX_TRITON_GRID: - use_short_kernel = False if output_cache: cache_out = torch.empty((N, D), device=x.device, dtype=x.dtype) @@ -409,7 +405,7 @@ def token_shift_fwd( N = B BD = triton.next_power_of_2(D) grid = (N, T) - IS_DECODE = T == 1 or (B == 1 and T == N) + IS_DECODE = T == 1 token_shift_fwd_kernel_short[grid]( x=x, y=y, @@ -430,7 +426,6 @@ def token_shift_fwd( BD = triton.next_power_of_2(D) ND = triton.cdiv(D, BD) - NB = triton.cdiv(B*T, 1024) def grid(meta): return (ND * NT, 1 if cu_seqlens is not None else N) token_shift_fwd_kernel_long[grid]( @@ -444,7 +439,6 @@ def token_shift_fwd( D=D, BD=BD, BT=BT, - NB=NB, ND=ND, STORE_FINAL_STATE=output_cache, ) @@ -487,7 +481,6 @@ def token_shift_bwd( if chunk_indices is None and cu_seqlens is not None: chunk_indices = prepare_chunk_indices(cu_seqlens, BT) NT = len(chunk_indices) if cu_seqlens is not None else triton.cdiv(T, BT) - NB = triton.cdiv(N * dy.shape[1], 1024) BD = triton.next_power_of_2(D) ND = triton.cdiv(D, BD) @@ -503,7 +496,6 @@ def token_shift_bwd( D=D, BD=BD, BT=BT, - NB=NB, ND=ND, ) return dx, grad_cache_out diff --git a/build/torch-cuda/modules/token_shift_cp.py b/build/torch-cuda/modules/token_shift_cp.py index b92027a11c6a0fca342079d272fb15bc643a170f..6931e2a461b5fa2781cc8a1712bbc4420bc8ef69 100644 --- a/build/torch-cuda/modules/token_shift_cp.py +++ b/build/torch-cuda/modules/token_shift_cp.py @@ -132,6 +132,8 @@ class TokenShiftCPFunction(torch.autograd.Function): ): if cp_context is None: raise ValueError("cp_context must be provided for TokenShiftCPFunction") + if cp_context.layout != 'contiguous': + raise NotImplementedError(f"token_shift CP does not support layout {cp_context.layout!r}") cu_seqlens = cp_context.cu_seqlens group = cp_context.group diff --git a/build/torch-cuda/ops/__init__.py b/build/torch-cuda/ops/__init__.py index 3f962e5dfd589647dd3966a128611808f2a02ace..63eb6d9027ef7db8d60f6adfe8f9f3c1ef162504 100644 --- a/build/torch-cuda/ops/__init__.py +++ b/build/torch-cuda/ops/__init__.py @@ -10,6 +10,7 @@ from .attn import parallel_attn from .attnres import fused_attnres from .based import fused_chunk_based, parallel_based from .comba import chunk_comba, fused_recurrent_comba +from .cyfa import chunk_cyfa, fused_recurrent_cyfa from .delta_rule import chunk_delta_rule, fused_chunk_delta_rule, fused_recurrent_delta_rule from .forgetting_attn import parallel_forgetting_attn from .gated_delta_rule import chunk_gated_delta_rule, chunk_gdn, fused_recurrent_gated_delta_rule, fused_recurrent_gdn @@ -27,18 +28,26 @@ from .lightning_attn import chunk_lightning_attn, fused_recurrent_lightning_attn from .linear_attn import chunk_linear_attn, fused_chunk_linear_attn, fused_recurrent_linear_attn from .log_linear_attn import chunk_log_linear_attn from .mesa_net import chunk_mesa_net +from .momentum_delta_rule import ( + chunk_momentum_delta_rule, + fused_recurrent_momentum_delta_rule, +) from .nsa import parallel_nsa from .parallax import parallel_parallax from .path_attn import parallel_path_attn +from .precond_gated_delta_rule import chunk_precond_gated_delta_rule, fused_recurrent_precond_gated_delta_rule +from .precond_kda import chunk_precond_kda, fused_recurrent_precond_kda from .retention import chunk_retention, fused_chunk_retention, fused_recurrent_retention, parallel_retention from .rwkv6 import chunk_rwkv6, fused_recurrent_rwkv6 from .rwkv7 import chunk_rwkv7, fused_recurrent_rwkv7 from .simple_gla import chunk_simple_gla, fused_chunk_simple_gla, fused_recurrent_simple_gla, parallel_simple_gla +from .stickbreaking_attn import parallel_stickbreaking_attn from .wall_attn import parallel_wall_attn, parallel_wall_attn_decode __all__ = [ 'chunk_abc', 'chunk_comba', + 'chunk_cyfa', 'chunk_delta_rule', 'chunk_dplr_delta_rule', 'chunk_gated_delta_rule', @@ -51,6 +60,9 @@ __all__ = [ 'chunk_linear_attn', 'chunk_log_linear_attn', 'chunk_mesa_net', + 'chunk_momentum_delta_rule', + 'chunk_precond_gated_delta_rule', + 'chunk_precond_kda', 'chunk_retention', 'chunk_rwkv6', 'chunk_rwkv7', @@ -63,6 +75,7 @@ __all__ = [ 'fused_chunk_retention', 'fused_chunk_simple_gla', 'fused_recurrent_comba', + 'fused_recurrent_cyfa', 'fused_recurrent_delta_rule', 'fused_recurrent_dplr_delta_rule', 'fused_recurrent_gated_delta_rule', @@ -74,6 +87,9 @@ __all__ = [ 'fused_recurrent_kda', 'fused_recurrent_lightning_attn', 'fused_recurrent_linear_attn', + 'fused_recurrent_momentum_delta_rule', + 'fused_recurrent_precond_gated_delta_rule', + 'fused_recurrent_precond_kda', 'fused_recurrent_retention', 'fused_recurrent_rwkv6', 'fused_recurrent_rwkv7', @@ -86,6 +102,7 @@ __all__ = [ 'parallel_path_attn', 'parallel_retention', 'parallel_simple_gla', + 'parallel_stickbreaking_attn', 'parallel_wall_attn', 'parallel_wall_attn_decode', ] diff --git a/build/torch-cuda/ops/abc/chunk.py b/build/torch-cuda/ops/abc/chunk.py index 1e364f43d7f6206d583305ebec966511b67a0f3c..35bbbe5a5e829f9b104034f0a28ae2eaae185801 100644 --- a/build/torch-cuda/ops/abc/chunk.py +++ b/build/torch-cuda/ops/abc/chunk.py @@ -11,7 +11,7 @@ import triton.language as tl from ...ops.utils import softmax_bwd, softmax_fwd from ...ops.utils.logcumsumexp import logcumsumexp_fwd_kernel -from ...ops.utils.op import exp +from ...ops.utils.op import exp, unflatten_program_id from ...utils import input_guard @@ -34,49 +34,57 @@ def chunk_abc_fwd_kernel_h( USE_INITIAL_STATE: tl.constexpr, STORE_FINAL_STATE: tl.constexpr, ): - i_v, i_k, i_bh = tl.program_id(0), tl.program_id(1), tl.program_id(2) + i_v, i_k, i_bh = tl.program_id(0).to(tl.int64), tl.program_id(1).to(tl.int64), tl.program_id(2).to(tl.int64) + o_k = i_k * BK + tl.arange(0, BK) + o_v = i_v * BV + tl.arange(0, BV) + m_kv = (o_k[:, None] < K) & (o_v[None, :] < V) b_h = tl.zeros([BK, BV], dtype=tl.float32) if USE_INITIAL_STATE: - p_h = tl.make_block_ptr(h0 + i_bh * K * V, (K, V), (V, 1), (i_k * BK, i_v * BV), (BK, BV), (1, 0)) - b_h += tl.load(p_h, boundary_check=(0, 1)).to(tl.float32) + p_h0 = h0 + i_bh * K * V + o_k[:, None] * V + o_v[None, :] + b_h += tl.load(p_h0, mask=m_kv, other=0.0).to(tl.float32) if NORMK: - p_z0 = tl.make_block_ptr(z + i_bh * T*K, (T * K,), (1,), (i_k * BK,), (BK,), (0,)) + p_z0 = z + i_bh * T*K + o_k else: - p_z0 = tl.make_block_ptr(z + i_bh * T*V, (T * V,), (1,), (i_v * BV,), (BV,), (0,)) + p_z0 = z + i_bh * T*V + o_v b_zp = tl.load(p_z0).to(tl.float32) for i_t in range(NT): - p_k = tl.make_block_ptr(k + i_bh * T*K, (K, T), (1, K), (i_k * BK, i_t * BT), (BK, BT), (0, 1)) - p_v = tl.make_block_ptr(v + i_bh * T*V, (T, V), (V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) - p_h = tl.make_block_ptr(h + i_bh * NT*K*V + i_t * K * V, (K, V), (V, 1), (i_k * BK, i_v * BV), (BK, BV), (1, 0)) - - tl.store(p_h, b_h.to(p_h.dtype.element_ty), boundary_check=(0, 1)) + o_t = i_t.to(tl.int64) * BT + tl.arange(0, BT) + m_kt = (o_k[:, None] < K) & (o_t[None, :] < T) + m_tv = (o_t[:, None] < T) & (o_v[None, :] < V) + p_k = k + i_bh * T*K + o_k[:, None] + o_t[None, :] * K + p_v = v + i_bh * T*V + o_t[:, None] * V + o_v[None, :] + p_h = h + i_bh * NT*K*V + i_t * K * V + o_k[:, None] * V + o_v[None, :] + + tl.store(p_h, b_h.to(p_h.dtype.element_ty), mask=m_kv) # [BK, BT] - b_k = tl.load(p_k, boundary_check=(0, 1)) + b_k = tl.load(p_k, mask=m_kt, other=0.0) # [BT, BV] - b_v = tl.load(p_v, boundary_check=(0, 1)) + b_v = tl.load(p_v, mask=m_tv, other=0.0) if NORMK: - p_zc = tl.make_block_ptr(z + i_bh * T*K, (T * K,), (1,), ((i_t * BT + BT - 1) * K + i_k * BK,), (BK,), (0,)) + o_zc = (i_t * BT + BT - 1).to(tl.int64) * K + i_k * BK + tl.arange(0, BK) + p_zc = z + i_bh * T*K + o_zc # [BK,] - b_zc = tl.load(p_zc, boundary_check=(0,)) + b_zc = tl.load(p_zc, mask=o_zc < T*K, other=0.0) b_r, b_zp = exp(b_zp - b_zc), b_zc # [BK, BV] b_h = b_h * b_r[:, None] b_k = exp(b_k - b_zc[:, None]).to(b_k.dtype) else: - p_zc = tl.make_block_ptr(z + i_bh * T*V, (T * V,), (1,), ((i_t * BT + BT - 1) * V + i_v * BV,), (BV,), (0,)) + o_zc = (i_t * BT + BT - 1).to(tl.int64) * V + i_v * BV + tl.arange(0, BV) + p_zc = z + i_bh * T*V + o_zc # [BV,] - b_zc = tl.load(p_zc, boundary_check=(0,)) + b_zc = tl.load(p_zc, mask=o_zc < T*V, other=0.0) b_r, b_zp = exp(b_zp - b_zc), b_zc # [BK, BV] b_h = b_h * b_r[None, :] b_v = exp(b_v - b_zc[None, :]).to(b_v.dtype) # [BK, BV] - b_h += tl.dot(b_k, b_v, allow_tf32=False) + b_h = tl.dot(b_k, b_v, b_h, allow_tf32=False) if STORE_FINAL_STATE: - p_h = tl.make_block_ptr(ht + i_bh * K * V, (K, V), (V, 1), (i_k * BK, i_v * BV), (BK, BV), (1, 0)) - tl.store(p_h, b_h.to(p_h.dtype.element_ty), boundary_check=(0, 1)) + p_ht = ht + i_bh * K * V + o_k[:, None] * V + o_v[None, :] + tl.store(p_ht, b_h.to(p_ht.dtype.element_ty), mask=m_kv) @triton.jit(do_not_specialize=['T']) @@ -92,41 +100,51 @@ def chunk_abc_fwd_kernel_intra_K( BV: tl.constexpr, NC: tl.constexpr, ): - i_v, i_c, i_bh = tl.program_id(0), tl.program_id(1), tl.program_id(2) + NV = tl.cdiv(V, BV) + i_v, i_c, i_bh = tl.program_id(0).to(tl.int64) % NV, tl.program_id(0).to(tl.int64) // NV, tl.program_id(1).to(tl.int64) i_t, i_i = i_c // NC, i_c % NC - p_z = tl.make_block_ptr(z + i_bh * T*V, (T, V), (V, 1), (i_t * BT + i_i * BC, i_v * BV), (BC, BV), (1, 0)) - p_zn = tl.make_block_ptr(z + i_bh * T*V, (T * V,), (1,), ((i_t * BT + i_i * BC) * V + i_v * BV,), (BV,), (0,)) + o_r = i_t * BT + i_i * BC + tl.arange(0, BC) + o_v = i_v * BV + tl.arange(0, BV) + m_rv = (o_r[:, None] < T) & (o_v[None, :] < V) + o_zn = (i_t * BT + i_i * BC) * V + o_v + p_z = z + i_bh * T*V + o_r[:, None] * V + o_v[None, :] + p_zn = z + i_bh * T*V + o_zn # [BV,] - b_zn = tl.load(p_zn, boundary_check=(0,)) + b_zn = tl.load(p_zn, mask=o_zn < T*V, other=0.0) # [BC, BV] b_o = tl.zeros([BC, BV], dtype=tl.float32) for i_j in range(0, i_i): - p_A = tl.make_block_ptr(A + i_bh * T * BT, (T, BT), (BT, 1), (i_t * BT + i_i * BC, i_j * BC), (BC, BC), (1, 0)) - p_v = tl.make_block_ptr(v + i_bh * T*V, (T, V), (V, 1), (i_t * BT + i_j * BC, i_v * BV), (BC, BV), (1, 0)) + o_vr = i_t * BT + i_j * BC + tl.arange(0, BC) + o_Aj = i_j * BC + tl.arange(0, BC) + m_vr = (o_vr[:, None] < T) & (o_v[None, :] < V) + m_A = (o_r[:, None] < T) & (o_Aj[None, :] < BT) + p_A = A + i_bh * T * BT + o_r[:, None] * BT + o_Aj[None, :] + p_v = v + i_bh * T*V + o_vr[:, None] * V + o_v[None, :] # [BC, BV] - b_v = tl.load(p_v, boundary_check=(0, 1)) + b_v = tl.load(p_v, mask=m_vr, other=0.0) # [BC, BC] - b_A = tl.load(p_A, boundary_check=(0, 1)) - b_o += tl.dot(b_A, exp(b_v - b_zn[None, :]).to(b_v.dtype), allow_tf32=False) - b_z = tl.load(p_z, boundary_check=(0, 1)) + b_A = tl.load(p_A, mask=m_A, other=0.0) + b_o = tl.dot(b_A, exp(b_v - b_zn[None, :]).to(b_v.dtype), b_o, allow_tf32=False) + b_z = tl.load(p_z, mask=m_rv, other=0.0) b_o *= exp(b_zn[None, :] - b_z) o_i = tl.arange(0, BC) o_A = i_bh * T * BT + (i_t * BT + i_i * BC + tl.arange(0, BC)) * BT + i_i * BC m_A = (i_t * BT + i_i * BC + tl.arange(0, BC)) < T for j in range(0, BC): - p_v = tl.make_block_ptr(v + i_bh * T*V, (T * V,), (1,), ((i_t * BT + i_i * BC + j) * V + i_v * BV,), (BV,), (0,)) + o_vj = (i_t * BT + i_i * BC + j) * V + o_v + p_vj = v + i_bh * T*V + o_vj # [BC,] b_A = tl.load(A + o_A + j, mask=m_A, other=0) # [BV,] - b_v = tl.load(p_v, boundary_check=(0,)).to(tl.float32) + b_v = tl.load(p_vj, mask=o_vj < T*V, other=0.0).to(tl.float32) # [BC, BV] # avoid 0 * inf = inf m_i = o_i[:, None] >= j b_o += tl.where(m_i, b_A[:, None] * exp(b_v[None, :] - b_z), 0) - p_o = tl.make_block_ptr(o + i_bh * T*V, (T, V), (V, 1), (i_t * BT + i_i * BC, i_v * BV), (BC, BV), (1, 0)) - tl.store(p_o, b_o.to(p_o.dtype.element_ty), boundary_check=(0, 1)) + p_o = o + i_bh * T*V + o_r[:, None] * V + o_v[None, :] + tl.store(p_o, b_o.to(p_o.dtype.element_ty), mask=m_rv) @triton.jit(do_not_specialize=['T']) @@ -146,45 +164,57 @@ def chunk_abc_fwd_kernel_K( BV: tl.constexpr, NT: tl.constexpr, ): - i_v, i_t, i_bh = tl.program_id(0), tl.program_id(1), tl.program_id(2) + i_v, i_t = unflatten_program_id(tl.cdiv(V, BV)) + i_bh = tl.program_id(1).to(tl.int64) i_p = tl.maximum(i_t * BT - 1, 0) o_i = tl.arange(0, BT) m_s = o_i[:, None] >= o_i[None, :] + o_t = i_t * BT + tl.arange(0, BT) + o_v = i_v * BV + tl.arange(0, BV) + m_t = o_t < T + m_tv = m_t[:, None] & (o_v[None, :] < V) b_o = tl.zeros([BT, BV], dtype=tl.float32) b_A = tl.zeros([BT, BT], dtype=tl.float32) for i_k in range(tl.cdiv(K, BK)): - p_q = tl.make_block_ptr(q + i_bh * T*K, (T, K), (K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) - p_k = tl.make_block_ptr(k + i_bh * T*K, (K, T), (1, K), (i_k * BK, i_t * BT), (BK, BT), (0, 1)) - p_h = tl.make_block_ptr(h + i_bh * NT*K*V + i_t * K * V, (K, V), (V, 1), (i_k * BK, i_v * BV), (BK, BV), (1, 0)) + o_k = i_k * BK + tl.arange(0, BK) + m_tk = m_t[:, None] & (o_k[None, :] < K) + m_kt = (o_k[:, None] < K) & m_t[None, :] + m_kv = (o_k[:, None] < K) & (o_v[None, :] < V) + p_q = q + i_bh * T*K + o_t[:, None] * K + o_k[None, :] + p_k = k + i_bh * T*K + o_k[:, None] + o_t[None, :] * K + p_h = h + i_bh * NT*K*V + i_t * K * V + o_k[:, None] * V + o_v[None, :] # [BT, BK] - b_q = tl.load(p_q, boundary_check=(0, 1)) + b_q = tl.load(p_q, mask=m_tk, other=0.0) b_q = (b_q * scale).to(b_q.dtype) # [BK, BT] - b_k = tl.load(p_k, boundary_check=(0, 1)) + b_k = tl.load(p_k, mask=m_kt, other=0.0) # [BK, BV] - b_h = tl.load(p_h, boundary_check=(0, 1)) + b_h = tl.load(p_h, mask=m_kv, other=0.0) # [BT, BV] - b_o += tl.dot(b_q, b_h, allow_tf32=False) + b_o = tl.dot(b_q, b_h, b_o, allow_tf32=False) # [BT, BT] - b_A += tl.dot(b_q, b_k, allow_tf32=False) - p_z = tl.make_block_ptr(z + i_bh * T*V, (T, V), (V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) - p_o = tl.make_block_ptr(o + i_bh * T*V, (T, V), (V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) + b_A = tl.dot(b_q, b_k, b_A, allow_tf32=False) + p_z = z + i_bh * T*V + o_t[:, None] * V + o_v[None, :] + p_o = o + i_bh * T*V + o_t[:, None] * V + o_v[None, :] # [BT, BV] - b_z = tl.load(p_z, boundary_check=(0, 1)) + b_z = tl.load(p_z, mask=m_tv, other=0.0) # [BT, BV] - p_zp = tl.make_block_ptr(z + i_bh * T*V, (T * V,), (1,), (i_p * V + i_v * BV,), (BV,), (0,)) - b_zp = tl.load(p_zp, boundary_check=(0,)) + o_zp = i_p * V + o_v + p_zp = z + i_bh * T*V + o_zp + b_zp = tl.load(p_zp, mask=o_zp < T*V, other=0.0) b_o = b_o * exp(b_zp[None, :] - b_z) - tl.store(p_o, b_o.to(p_o.dtype.element_ty), boundary_check=(0, 1)) + tl.store(p_o, b_o.to(p_o.dtype.element_ty), mask=m_tv) - p_A = tl.make_block_ptr(A + i_bh * T * BT, (T, BT), (BT, 1), (i_t * BT, 0), (BT, BT), (1, 0)) + o_A = tl.arange(0, BT) + m_AT = m_t[:, None] & (o_A[None, :] < BT) + p_A = A + i_bh * T * BT + o_t[:, None] * BT + o_A[None, :] # [BT, BT] b_A = tl.where(m_s, b_A, 0.) if i_v == 0: - tl.store(p_A, b_A.to(p_A.dtype.element_ty), boundary_check=(0, 1)) + tl.store(p_A, b_A.to(p_A.dtype.element_ty), mask=m_AT) @triton.jit(do_not_specialize=['T']) @@ -201,49 +231,55 @@ def chunk_abc_fwd_kernel_intra_V( BK: tl.constexpr, NC: tl.constexpr, ): - i_k, i_c, i_bh = tl.program_id(0), tl.program_id(1), tl.program_id(2) + NK = tl.cdiv(K, BK) + i_k, i_c, i_bh = tl.program_id(0).to(tl.int64) % NK, tl.program_id(0).to(tl.int64) // NK, tl.program_id(1).to(tl.int64) i_t, i_i, i_j = i_c // (NC * NC), (i_c % (NC * NC)) // NC, (i_c % (NC * NC)) % NC - n_bh = tl.num_programs(2) + n_bh = tl.num_programs(1).to(tl.int64) + o_q = i_t * BT + i_i * BC + tl.arange(0, BC) + o_k = i_k * BK + tl.arange(0, BK) + m_qk = (o_q[:, None] < T) & (o_k[None, :] < K) if i_i > i_j: - p_q = tl.make_block_ptr(q + i_bh * T*K, (T, K), (K, 1), (i_t * BT + i_i * BC, i_k * BK), (BC, BK), (1, 0)) - p_k = tl.make_block_ptr(k + i_bh * T*K, (K, T), (1, K), (i_k * BK, i_t * BT + i_j * BC), (BK, BC), (0, 1)) - p_z = tl.make_block_ptr(z + i_bh * T*K, (T, K), (K, 1), (i_t * BT + i_i * BC, i_k * BK), (BC, BK), (1, 0)) - p_A = tl.make_block_ptr(A + (i_k*n_bh+i_bh)*T*BT, (T, BT), (BT, 1), (i_t * BT + i_i * BC, i_j * BC), (BC, BC), (1, 0)) - p_zn = tl.make_block_ptr(z + i_bh * T*K, (T * K,), (1,), ((i_t * BT + i_i * BC) * K + i_k * BK,), (BK,), (0,)) + o_kj = i_t * BT + i_j * BC + tl.arange(0, BC) + o_Aj = i_j * BC + tl.arange(0, BC) + o_zn = (i_t * BT + i_i * BC) * K + o_k + p_q = q + i_bh * T*K + o_q[:, None] * K + o_k[None, :] + p_k = k + i_bh * T*K + o_k[:, None] + o_kj[None, :] * K + p_z = z + i_bh * T*K + o_q[:, None] * K + o_k[None, :] + p_A = A + (i_k*n_bh+i_bh)*T*BT + o_q[:, None] * BT + o_Aj[None, :] + p_zn = z + i_bh * T*K + o_zn # [BK,] - b_zn = tl.load(p_zn, boundary_check=(0,)) + b_zn = tl.load(p_zn, mask=o_zn < T*K, other=0.0) # [BC, BK] - b_q = tl.load(p_q, boundary_check=(0, 1)) - b_z = tl.load(p_z, boundary_check=(0, 1)) + b_q = tl.load(p_q, mask=m_qk, other=0.0) + b_z = tl.load(p_z, mask=m_qk, other=0.0) b_q = (b_q * exp(b_zn[None, :] - b_z) * scale).to(b_q.dtype) # [BK, BC] - b_k = tl.load(p_k, boundary_check=(0, 1)) + b_k = tl.load(p_k, mask=(o_k[:, None] < K) & (o_kj[None, :] < T), other=0.0) b_k = exp(b_k - b_zn[:, None]).to(b_k.dtype) # [BC, BC] b_A = tl.dot(b_q, b_k, allow_tf32=False) - tl.store(p_A, b_A.to(A.dtype.element_ty), boundary_check=(0, 1)) + tl.store(p_A, b_A.to(A.dtype.element_ty), mask=(o_q[:, None] < T) & (o_Aj[None, :] < BT)) elif i_i == i_j: - p_q = tl.make_block_ptr(q + i_bh * T*K, (T, K), (K, 1), (i_t * BT + i_i * BC, i_k * BK), (BC, BK), (1, 0)) - p_k = tl.make_block_ptr(k + i_bh * T*K, (T * K,), (1,), ((i_t * BT + i_j * BC) * K + i_k * BK,), (BK,), (0,)) - p_z = tl.make_block_ptr(z + i_bh * T*K, (T, K), (K, 1), (i_t * BT + i_i * BC, i_k * BK), (BC, BK), (1, 0)) + p_q = q + i_bh * T*K + o_q[:, None] * K + o_k[None, :] + p_z = z + i_bh * T*K + o_q[:, None] * K + o_k[None, :] # [BC, BK] - b_q = tl.load(p_q, boundary_check=(0, 1)) - b_z = tl.load(p_z, boundary_check=(0, 1)) + b_q = tl.load(p_q, mask=m_qk, other=0.0) + b_z = tl.load(p_z, mask=m_qk, other=0.0) o_i = tl.arange(0, BC) o_A = (i_bh + i_k * n_bh) * T * BT + (i_t * BT + i_i * BC + tl.arange(0, BC)) * BT + i_j * BC m_A = (i_t * BT + i_i * BC + tl.arange(0, BC)) < T for j in range(0, BC): + o_kj = (i_t * BT + i_j * BC + j) * K + o_k + p_kj = k + i_bh * T*K + o_kj # [BK,] - b_k = tl.load(p_k, boundary_check=(0,)).to(tl.float32) + b_k = tl.load(p_kj, mask=o_kj < T*K, other=0.0).to(tl.float32) # [BC,] b_A = tl.sum(b_q * exp(b_k[None, :] - b_z) * scale, 1) b_A = tl.where(o_i >= j, b_A, 0.) tl.store(A + o_A + j, b_A.to(b_q.dtype), mask=m_A) - p_k = tl.advance(p_k, (K,)) - @triton.jit(do_not_specialize=['T']) def chunk_abc_fwd_kernel_V( @@ -262,39 +298,50 @@ def chunk_abc_fwd_kernel_V( BV: tl.constexpr, NT: tl.constexpr, ): - i_v, i_t, i_bh = tl.program_id(0), tl.program_id(1), tl.program_id(2) + i_v, i_t = unflatten_program_id(tl.cdiv(V, BV)) + i_bh = tl.program_id(1).to(tl.int64) i_p = tl.maximum(i_t * BT - 1, 0) + o_t = i_t * BT + tl.arange(0, BT) + o_v = i_v * BV + tl.arange(0, BV) + m_t = o_t < T + m_tv = m_t[:, None] & (o_v[None, :] < V) b_o = tl.zeros([BT, BV], dtype=tl.float32) for i_k in range(tl.cdiv(K, BK)): - p_q = tl.make_block_ptr(q + i_bh * T*K, (T, K), (K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) - p_z = tl.make_block_ptr(z + i_bh * T*K, (T, K), (K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) - p_h = tl.make_block_ptr(h + i_bh * NT*K*V + i_t * K * V, (K, V), (V, 1), (i_k * BK, i_v * BV), (BK, BV), (1, 0)) - p_zp = tl.make_block_ptr(z + i_bh * T*K, (T * K,), (1,), (i_p * K + i_k * BK,), (BK,), (0,)) + o_k = i_k * BK + tl.arange(0, BK) + m_tk = m_t[:, None] & (o_k[None, :] < K) + m_kv = (o_k[:, None] < K) & (o_v[None, :] < V) + o_zp = i_p * K + o_k + p_q = q + i_bh * T*K + o_t[:, None] * K + o_k[None, :] + p_z = z + i_bh * T*K + o_t[:, None] * K + o_k[None, :] + p_h = h + i_bh * NT*K*V + i_t * K * V + o_k[:, None] * V + o_v[None, :] + p_zp = z + i_bh * T*K + o_zp # [BT, BK] - b_q = tl.load(p_q, boundary_check=(0, 1)) + b_q = tl.load(p_q, mask=m_tk, other=0.0) b_q = (b_q * scale).to(b_q.dtype) # [BT, BK] - b_z = tl.load(p_z, boundary_check=(0, 1)) + b_z = tl.load(p_z, mask=m_tk, other=0.0) # [BT, BK] - b_zp = tl.load(p_zp, boundary_check=(0,)) + b_zp = tl.load(p_zp, mask=o_zp < T*K, other=0.0) b_q = (b_q * exp(b_zp[None, :] - b_z)).to(b_q.dtype) # [BK, BV] - b_h = tl.load(p_h, boundary_check=(0, 1)) + b_h = tl.load(p_h, mask=m_kv, other=0.0) # works but dkw, owing to divine benevolence # [BT, BV] if i_k >= 0: - b_o += tl.dot(b_q, b_h, allow_tf32=False) - p_v = tl.make_block_ptr(v + i_bh * T*V, (T, V), (V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) - p_o = tl.make_block_ptr(o + i_bh * T*V, (T, V), (V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) - p_A = tl.make_block_ptr(A + i_bh * T * BT, (T, BT), (BT, 1), (i_t * BT, 0), (BT, BT), (1, 0)) + b_o = tl.dot(b_q, b_h, b_o, allow_tf32=False) + p_v = v + i_bh * T*V + o_t[:, None] * V + o_v[None, :] + p_o = o + i_bh * T*V + o_t[:, None] * V + o_v[None, :] + o_A = tl.arange(0, BT) + m_AT = m_t[:, None] & (o_A[None, :] < BT) + p_A = A + i_bh * T * BT + o_t[:, None] * BT + o_A[None, :] # [BT, BV] - b_v = tl.load(p_v, boundary_check=(0, 1)) + b_v = tl.load(p_v, mask=m_tv, other=0.0) # [BT, BT] - b_A = tl.load(p_A, boundary_check=(0, 1)) - b_o += tl.dot(b_A.to(b_v.dtype), b_v, allow_tf32=False) - tl.store(p_o, b_o.to(p_o.dtype.element_ty), boundary_check=(0, 1)) + b_A = tl.load(p_A, mask=m_AT, other=0.0) + b_o = tl.dot(b_A.to(b_v.dtype), b_v, b_o, allow_tf32=False) + tl.store(p_o, b_o.to(p_o.dtype.element_ty), mask=m_tv) @triton.jit(do_not_specialize=['T']) @@ -313,47 +360,56 @@ def chunk_abc_bwd_kernel_dh( NT: tl.constexpr, NORMK: tl.constexpr, ): - i_k, i_v, i_bh = tl.program_id(0), tl.program_id(1), tl.program_id(2) + i_k, i_v, i_bh = tl.program_id(0).to(tl.int64), tl.program_id(1).to(tl.int64), tl.program_id(2).to(tl.int64) + o_k = i_k * BK + tl.arange(0, BK) + o_v = i_v * BV + tl.arange(0, BV) + m_kv = (o_k[:, None] < K) & (o_v[None, :] < V) b_dh = tl.zeros([BK, BV], dtype=tl.float32) b_zp = tl.full([BK if NORMK else BV], float('inf'), dtype=tl.float32) for i_t in range(NT - 1, -1, -1): i_p = tl.maximum(i_t * BT - 1, 0) - p_q = tl.make_block_ptr(q + i_bh * T*K, (K, T), (1, K), (i_k * BK, i_t * BT), (BK, BT), (0, 1)) - p_do = tl.make_block_ptr(do + i_bh * T*V, (T, V), (V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) - p_dh = tl.make_block_ptr(dh + i_bh * NT*K*V + i_t * K*V, (K, V), (V, 1), (i_k * BK, i_v * BV), (BK, BV), (1, 0)) + o_t = i_t.to(tl.int64) * BT + tl.arange(0, BT) + m_t = o_t < T + m_kt = (o_k[:, None] < K) & m_t[None, :] + m_tv = m_t[:, None] & (o_v[None, :] < V) + p_q = q + i_bh * T*K + o_k[:, None] + o_t[None, :] * K + p_do = do + i_bh * T*V + o_t[:, None] * V + o_v[None, :] + p_dh = dh + i_bh * NT*K*V + i_t * K*V + o_k[:, None] * V + o_v[None, :] # [BK, BT] - b_q = tl.load(p_q, boundary_check=(0, 1)) + b_q = tl.load(p_q, mask=m_kt, other=0.0) b_q = (b_q * scale).to(b_q.dtype) # [BT, BV] - b_do = tl.load(p_do, boundary_check=(0, 1)) + b_do = tl.load(p_do, mask=m_tv, other=0.0) - tl.store(p_dh, b_dh.to(p_dh.dtype.element_ty), boundary_check=(0, 1)) + tl.store(p_dh, b_dh.to(p_dh.dtype.element_ty), mask=m_kv) if NORMK: - p_z = tl.make_block_ptr(z + i_bh * T*K, (K, T), (1, K), (i_k * BK, i_t * BT), (BK, BT), (0, 1)) - p_zc = tl.make_block_ptr(z + i_bh * T*K, (T * K,), (1,), (i_p * K + i_k * BK,), (BK,), (0,)) + o_zc = i_p.to(tl.int64) * K + o_k + p_z = z + i_bh * T*K + o_k[:, None] + o_t[None, :] * K + p_zc = z + i_bh * T*K + o_zc # [BK,] - b_zc = tl.load(p_zc, boundary_check=(0,)) + b_zc = tl.load(p_zc, mask=o_zc < T*K, other=0.0) b_r, b_zp = exp(b_zc - b_zp), b_zc # [BK, BT] - b_z = tl.load(p_z, boundary_check=(0, 1)) + b_z = tl.load(p_z, mask=m_kt, other=0.0) b_q = (b_q * exp(b_zc[:, None] - b_z)).to(b_q.dtype) # [BK, BV] b_dh = b_dh * b_r[:, None] else: - p_z = tl.make_block_ptr(z + i_bh * T*V, (T, V), (V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) - p_zc = tl.make_block_ptr(z + i_bh * T*V, (T * V,), (1,), (i_p * V + i_v * BV,), (BV,), (0,)) + o_zc = i_p.to(tl.int64) * V + o_v + p_z = z + i_bh * T*V + o_t[:, None] * V + o_v[None, :] + p_zc = z + i_bh * T*V + o_zc # [BV,] - b_zc = tl.load(p_zc, boundary_check=(0,)) + b_zc = tl.load(p_zc, mask=o_zc < T*V, other=0.0) b_r, b_zp = exp(b_zc - b_zp), b_zc # [BT, BV] - b_z = tl.load(p_z, boundary_check=(0,)) + b_z = tl.load(p_z, mask=m_t[:, None], other=0.0) b_do = (b_do * exp(b_zc[None, :] - b_z)).to(b_do.dtype) # [BK, BV] b_dh = b_dh * b_r[None, :] # [BK, BV] - b_dh += tl.dot(b_q, b_do, allow_tf32=False) + b_dh = tl.dot(b_q, b_do, b_dh, allow_tf32=False) @triton.jit(do_not_specialize=['T']) @@ -378,76 +434,89 @@ def chunk_abc_bwd_kernel_V( BV: tl.constexpr, NT: tl.constexpr, ): - i_k, i_t, i_bh = tl.program_id(0), tl.program_id(1), tl.program_id(2) + i_k, i_t = unflatten_program_id(tl.cdiv(K, BK)) + i_bh = tl.program_id(1).to(tl.int64) i_p = tl.maximum(i_t * BT - 1, 0) - n_bh = tl.num_programs(2) - - p_k = tl.make_block_ptr(k + i_bh * T*K, (T, K), (K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) - p_zc = tl.make_block_ptr(z + i_bh * T*K, (T * K,), (1,), ((i_t * BT + BT - 1) * K + i_k * BK,), (BK,), (0,)) - p_A = tl.make_block_ptr(A + i_bh * T * BT, (BT, T), (1, BT), (0, i_t * BT), (BT, BT), (0, 1)) + n_bh = tl.num_programs(1).to(tl.int64) + + o_t = i_t * BT + tl.arange(0, BT) + o_k = i_k * BK + tl.arange(0, BK) + m_t = o_t < T + m_tk = m_t[:, None] & (o_k[None, :] < K) + o_zc = (i_t * BT + BT - 1) * K + o_k + o_A = tl.arange(0, BT) + m_AT = (o_A[:, None] < BT) & m_t[None, :] + p_k = k + i_bh * T*K + o_t[:, None] * K + o_k[None, :] + p_zc = z + i_bh * T*K + o_zc + p_A = A + i_bh * T * BT + o_A[:, None] + o_t[None, :] * BT # [BK,] - b_zc = tl.load(p_zc, boundary_check=(0,)) + b_zc = tl.load(p_zc, mask=o_zc < T*K, other=0.0) # [BT, BK] - b_k = tl.load(p_k, boundary_check=(0, 1)) + b_k = tl.load(p_k, mask=m_tk, other=0.0) b_k = exp(b_k - b_zc[None, :]).to(b_k.dtype) # [BT, BT] - b_A = tl.load(p_A, boundary_check=(0, 1)) + b_A = tl.load(p_A, mask=m_AT, other=0.0) b_dq = tl.zeros([BT, BK], dtype=tl.float32) b_dk = tl.zeros([BT, BK], dtype=tl.float32) b_dA = tl.zeros([BT, BT], dtype=tl.float32) for i_v in range(tl.cdiv(V, BV)): - p_v = tl.make_block_ptr(v + i_bh * T*V, (T, V), (V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) - p_h = tl.make_block_ptr(h + i_bh * NT*K*V + i_t * V * K, (V, K), (1, V), (i_v * BV, i_k * BK), (BV, BK), (0, 1)) - p_do = tl.make_block_ptr(do + i_bh * T*V, (T, V), (V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) - p_dh = tl.make_block_ptr(dh + i_bh * NT*K*V + i_t * K*V, (K, V), (V, 1), (i_k * BK, i_v * BV), (BK, BV), (1, 0)) - p_dv = tl.make_block_ptr(dv + (i_k*n_bh+i_bh) * T*V, (T, V), (V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) + o_v = i_v * BV + tl.arange(0, BV) + m_tv = m_t[:, None] & (o_v[None, :] < V) + m_h = (o_v[:, None] < V) & (o_k[None, :] < K) + m_kv = (o_k[:, None] < K) & (o_v[None, :] < V) + p_v = v + i_bh * T*V + o_t[:, None] * V + o_v[None, :] + p_h = h + i_bh * NT*K*V + i_t * V * K + o_v[:, None] + o_k[None, :] * V + p_do = do + i_bh * T*V + o_t[:, None] * V + o_v[None, :] + p_dh = dh + i_bh * NT*K*V + i_t * K*V + o_k[:, None] * V + o_v[None, :] + p_dv = dv + (i_k*n_bh+i_bh) * T*V + o_t[:, None] * V + o_v[None, :] # [BT, BV] - b_v = tl.load(p_v, boundary_check=(0, 1)) + b_v = tl.load(p_v, mask=m_tv, other=0.0) # [BV, BK] - b_h = tl.load(p_h, boundary_check=(0, 1)) + b_h = tl.load(p_h, mask=m_h, other=0.0) # [BT, BV] - b_do = tl.load(p_do, boundary_check=(0, 1)) + b_do = tl.load(p_do, mask=m_tv, other=0.0) # [BK, BV] - b_dh = tl.load(p_dh, boundary_check=(0, 1)) + b_dh = tl.load(p_dh, mask=m_kv, other=0.0) # [BT, BV] b_dv = tl.dot(b_k, b_dh, allow_tf32=False) if i_k == 0: - b_dv += tl.dot(b_A.to(b_do.dtype), b_do, allow_tf32=False) + b_dv = tl.dot(b_A.to(b_do.dtype), b_do, b_dv, allow_tf32=False) b_do = (b_do * scale).to(b_do.dtype) - tl.store(p_dv, b_dv.to(p_dv.dtype.element_ty), boundary_check=(0, 1)) + tl.store(p_dv, b_dv.to(p_dv.dtype.element_ty), mask=m_tv) # [BT, BT] - b_dA += tl.dot(b_do, tl.trans(b_v), allow_tf32=False) + b_dA = tl.dot(b_do, tl.trans(b_v), b_dA, allow_tf32=False) # [BT, BK] - b_dq += tl.dot(b_do, b_h, allow_tf32=False) + b_dq = tl.dot(b_do, b_h, b_dq, allow_tf32=False) # [BT, BK] - b_dk += tl.dot(b_v, tl.trans(b_dh), allow_tf32=False) - p_z = tl.make_block_ptr(z + i_bh * T*K, (T, K), (K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) - p_zp = tl.make_block_ptr(z + i_bh * T*K, (T * K,), (1,), (i_p * K + i_k * BK,), (BK,), (0,)) + b_dk = tl.dot(b_v, tl.trans(b_dh), b_dk, allow_tf32=False) + o_zp = i_p * K + o_k + p_z = z + i_bh * T*K + o_t[:, None] * K + o_k[None, :] + p_zp = z + i_bh * T*K + o_zp # [BK,] - b_zp = tl.load(p_zp, boundary_check=(0,)) + b_zp = tl.load(p_zp, mask=o_zp < T*K, other=0.0) # [BT, BK] - b_z = tl.load(p_z, boundary_check=(0, 1)) + b_z = tl.load(p_z, mask=m_tk, other=0.0) b_z = exp(b_zp[None, :] - b_z) # [BT, BK] b_dq = b_dq * b_z b_dk = b_dk * b_k - p_dq = tl.make_block_ptr(dq + i_bh * T*K, (T, K), (K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) - p_dk = tl.make_block_ptr(dk + i_bh * T*K, (T, K), (K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) - p_dA = tl.make_block_ptr(dA + i_bh * T * BT, (T, BT), (BT, 1), (i_t * BT, 0), (BT, BT), (1, 0)) - tl.store(p_dq, b_dq.to(p_dq.dtype.element_ty), boundary_check=(0, 1)) - tl.store(p_dk, b_dk.to(p_dk.dtype.element_ty), boundary_check=(0, 1)) + p_dq = dq + i_bh * T*K + o_t[:, None] * K + o_k[None, :] + p_dk = dk + i_bh * T*K + o_t[:, None] * K + o_k[None, :] + p_dA = dA + i_bh * T * BT + o_t[:, None] * BT + o_A[None, :] + tl.store(p_dq, b_dq.to(p_dq.dtype.element_ty), mask=m_tk) + tl.store(p_dk, b_dk.to(p_dk.dtype.element_ty), mask=m_tk) o_i = tl.arange(0, BT) m_s = o_i[:, None] >= o_i[None, :] # [BT, BT] b_dA = tl.where(m_s, b_dA, 0.).to(b_k.dtype) if i_k == 0: - tl.store(p_dA, b_dA.to(p_dA.dtype.element_ty), boundary_check=(0, 1)) + tl.store(p_dA, b_dA.to(p_dA.dtype.element_ty), mask=m_t[:, None] & (o_A[None, :] < BT)) @triton.jit(do_not_specialize=['T']) @@ -465,82 +534,98 @@ def chunk_abc_bwd_kernel_intra_V( BK: tl.constexpr, NC: tl.constexpr, ): - i_k, i_c, i_bh = tl.program_id(0), tl.program_id(1), tl.program_id(2) + NK = tl.cdiv(K, BK) + i_k, i_c, i_bh = tl.program_id(0).to(tl.int64) % NK, tl.program_id(0).to(tl.int64) // NK, tl.program_id(1).to(tl.int64) i_t, i_i = i_c // NC, i_c % NC - p_z = tl.make_block_ptr(z + i_bh * T*K, (T, K), (K, 1), (i_t * BT + i_i * BC, i_k * BK), (BC, BK), (1, 0)) - p_zn = tl.make_block_ptr(z + i_bh * T*K, (T * K,), (1,), ((i_t * BT + i_i * BC) * K + i_k * BK,), (BK,), (0,)) + o_r = i_t * BT + i_i * BC + tl.arange(0, BC) + o_k = i_k * BK + tl.arange(0, BK) + m_rk = (o_r[:, None] < T) & (o_k[None, :] < K) + o_zn = (i_t * BT + i_i * BC) * K + o_k + p_z = z + i_bh * T*K + o_r[:, None] * K + o_k[None, :] + p_zn = z + i_bh * T*K + o_zn # [BK,] - b_zn = tl.load(p_zn, boundary_check=(0,)) + b_zn = tl.load(p_zn, mask=o_zn < T*K, other=0.0) # [BC, BK] - b_z = tl.load(p_z, boundary_check=(0, 1)) + b_z = tl.load(p_z, mask=m_rk, other=0.0) b_zq = exp(b_zn[None, :] - b_z) b_dq = tl.zeros([BC, BK], dtype=tl.float32) for i_j in range(0, i_i): - p_k = tl.make_block_ptr(k + i_bh * T*K, (T, K), (K, 1), (i_t * BT + i_j * BC, i_k * BK), (BC, BK), (1, 0)) - p_dA = tl.make_block_ptr(dA + i_bh * T * BT, (T, BT), (BT, 1), (i_t * BT + i_i * BC, i_j * BC), (BC, BC), (1, 0)) + o_j = i_t * BT + i_j * BC + tl.arange(0, BC) + o_dAj = i_j * BC + tl.arange(0, BC) + m_jk = (o_j[:, None] < T) & (o_k[None, :] < K) + m_dA = (o_r[:, None] < T) & (o_dAj[None, :] < BT) + p_k = k + i_bh * T*K + o_j[:, None] * K + o_k[None, :] + p_dA = dA + i_bh * T * BT + o_r[:, None] * BT + o_dAj[None, :] # [BC, BK] - b_k = tl.load(p_k, boundary_check=(0, 1)) + b_k = tl.load(p_k, mask=m_jk, other=0.0) b_kz = exp(b_k - b_zn[None, :]).to(b_k.dtype) # [BC, BC] - b_dA = tl.load(p_dA, boundary_check=(0, 1)) + b_dA = tl.load(p_dA, mask=m_dA, other=0.0) # [BC, BK] - b_dq += tl.dot(b_dA, b_kz, allow_tf32=False) + b_dq = tl.dot(b_dA, b_kz, b_dq, allow_tf32=False) b_dq *= b_zq o_i = tl.arange(0, BC) o_dA = i_bh * T * BT + (i_t * BT + i_i * BC + tl.arange(0, BC)) * BT + i_i * BC m_dA = (i_t * BT + i_i * BC + tl.arange(0, BC)) < T for j in range(0, BC): - p_kj = tl.make_block_ptr(k + i_bh * T*K, (T * K,), (1,), ((i_t * BT + i_i*BC+j) * K + i_k * BK,), (BK,), (0,)) + o_kj = (i_t * BT + i_i * BC + j) * K + o_k + p_kj = k + i_bh * T*K + o_kj # [BC,] b_dA = tl.load(dA + o_dA + j, mask=m_dA, other=0) # [BK,] - b_kj = tl.load(p_kj, boundary_check=(0,)).to(tl.float32) + b_kj = tl.load(p_kj, mask=o_kj < T*K, other=0.0).to(tl.float32) # [BC, BK] m_i = o_i[:, None] >= j # [BC, BK] b_dq += tl.where(m_i, b_dA[:, None] * exp(b_kj[None, :] - b_z), 0.) - p_dq = tl.make_block_ptr(dq + i_bh * T*K, (T, K), (K, 1), (i_t * BT + i_i * BC, i_k * BK), (BC, BK), (1, 0)) - tl.store(p_dq, b_dq.to(p_dq.dtype.element_ty), boundary_check=(0, 1)) + p_dq = dq + i_bh * T*K + o_r[:, None] * K + o_k[None, :] + tl.store(p_dq, b_dq.to(p_dq.dtype.element_ty), mask=m_rk) tl.debug_barrier() - p_k = tl.make_block_ptr(k + i_bh * T*K, (T, K), (K, 1), (i_t * BT + i_i * BC, i_k * BK), (BC, BK), (1, 0)) - p_zn = tl.make_block_ptr(z + i_bh * T*K, (T*K,), (1,), ((i_t * BT + i_i * BC + BC - 1) * K + i_k * BK,), (BK,), (0,)) + o_zn = (i_t * BT + i_i * BC + BC - 1) * K + o_k + p_k = k + i_bh * T*K + o_r[:, None] * K + o_k[None, :] + p_zn = z + i_bh * T*K + o_zn # [BK,] - b_zn = tl.load(p_zn, boundary_check=(0,)) + b_zn = tl.load(p_zn, mask=o_zn < T*K, other=0.0) # [BC, BK] - b_k = tl.load(p_k, boundary_check=(0, 1)) + b_k = tl.load(p_k, mask=m_rk, other=0.0) b_kz = exp(b_k - b_zn[None, :]) b_dk = tl.zeros([BC, BK], dtype=tl.float32) for i_j in range(i_i + 1, NC): - p_q = tl.make_block_ptr(q + i_bh * T*K, (T, K), (K, 1), (i_t * BT + i_j * BC, i_k * BK), (BC, BK), (1, 0)) - p_z = tl.make_block_ptr(z + i_bh * T*K, (T, K), (K, 1), (i_t * BT + i_j * BC, i_k * BK), (BC, BK), (1, 0)) - p_dA = tl.make_block_ptr(dA + i_bh * T * BT, (T, BT), (BT, 1), (i_t * BT + i_j * BC, i_i * BC), (BC, BC), (1, 0)) + o_j = i_t * BT + i_j * BC + tl.arange(0, BC) + o_dAi = i_i * BC + tl.arange(0, BC) + m_jk = (o_j[:, None] < T) & (o_k[None, :] < K) + m_dA2 = (o_j[:, None] < T) & (o_dAi[None, :] < BT) + p_q = q + i_bh * T*K + o_j[:, None] * K + o_k[None, :] + p_z = z + i_bh * T*K + o_j[:, None] * K + o_k[None, :] + p_dA = dA + i_bh * T * BT + o_j[:, None] * BT + o_dAi[None, :] # [BC, BK] - b_q = tl.load(p_q, boundary_check=(0, 1)) - b_z = tl.load(p_z, boundary_check=(0, 1)) + b_q = tl.load(p_q, mask=m_jk, other=0.0) + b_z = tl.load(p_z, mask=m_jk, other=0.0) b_qz = (b_q * exp(b_zn[None, :] - b_z)).to(b_q.dtype) # [BC, BC] - b_dA = tl.load(p_dA, boundary_check=(0, 1)) + b_dA = tl.load(p_dA, mask=m_dA2, other=0.0) # [BC, BK] - b_dk += tl.dot(tl.trans(b_dA), b_qz, allow_tf32=False) + b_dk = tl.dot(tl.trans(b_dA), b_qz, b_dk, allow_tf32=False) b_dk *= b_kz o_dA = i_bh * T * BT + (i_t * BT + i_i * BC) * BT + i_i * BC + tl.arange(0, BC) for j in range(0, BC): - p_qj = tl.make_block_ptr(q + i_bh * T*K, (T * K,), (1,), ((i_t * BT + i_i * BC + j) * K + i_k * BK,), (BK,), (0,)) - p_zj = tl.make_block_ptr(z + i_bh * T*K, (T * K,), (1,), ((i_t * BT + i_i * BC + j) * K + i_k * BK,), (BK,), (0,)) + o_qj = (i_t * BT + i_i * BC + j) * K + o_k + p_qj = q + i_bh * T*K + o_qj + p_zj = z + i_bh * T*K + o_qj # [BC,] b_dA = tl.load(dA + o_dA + j * BT, mask=(i_t * BT + i_i * BC + j < T), other=0) # [BK,] - b_qj = tl.load(p_qj, boundary_check=(0,)).to(tl.float32) - b_zj = tl.load(p_zj, boundary_check=(0,)).to(tl.float32) + b_qj = tl.load(p_qj, mask=o_qj < T*K, other=0.0).to(tl.float32) + b_zj = tl.load(p_zj, mask=o_qj < T*K, other=0.0).to(tl.float32) # [BC, BK] m_i = o_i[:, None] <= j b_dk += tl.where(m_i, b_dA[:, None] * b_qj[None, :] * exp(b_k - b_zj[None, :]), 0.) - p_dk = tl.make_block_ptr(dk + i_bh * T*K, (T, K), (K, 1), (i_t * BT + i_i * BC, i_k * BK), (BC, BK), (1, 0)) - tl.store(p_dk, b_dk.to(p_dk.dtype.element_ty), boundary_check=(0, 1)) + p_dk = dk + i_bh * T*K + o_r[:, None] * K + o_k[None, :] + tl.store(p_dk, b_dk.to(p_dk.dtype.element_ty), mask=m_rk) @triton.jit(do_not_specialize=['T']) @@ -557,49 +642,55 @@ def chunk_abc_bwd_kernel_intra_K( BV: tl.constexpr, NC: tl.constexpr, ): - i_v, i_c, i_bh = tl.program_id(0), tl.program_id(1), tl.program_id(2) + NV = tl.cdiv(V, BV) + i_v, i_c, i_bh = tl.program_id(0).to(tl.int64) % NV, tl.program_id(0).to(tl.int64) // NV, tl.program_id(1).to(tl.int64) i_t, i_i, i_j = i_c // (NC * NC), (i_c % (NC * NC)) // NC, (i_c % (NC * NC)) % NC - n_bh = tl.num_programs(2) + n_bh = tl.num_programs(1).to(tl.int64) + o_r = i_t * BT + i_i * BC + tl.arange(0, BC) + o_v = i_v * BV + tl.arange(0, BV) + m_rv = (o_r[:, None] < T) & (o_v[None, :] < V) if i_i > i_j: - p_v = tl.make_block_ptr(v + i_bh * T*V, (V, T), (1, V), (i_v * BV, i_t * BT + i_j * BC), (BV, BC), (0, 1)) - p_z = tl.make_block_ptr(z + i_bh * T*V, (T, V), (V, 1), (i_t * BT + i_i * BC, i_v * BV), (BC, BV), (1, 0)) - p_zn = tl.make_block_ptr(z + i_bh * T*V, (T * V,), (1,), ((i_t * BT + i_i * BC) * V + i_v * BV,), (BV,), (0,)) - p_do = tl.make_block_ptr(do + i_bh * T*V, (T, V), (V, 1), (i_t * BT + i_i * BC, i_v * BV), (BC, BV), (1, 0)) - p_dA = tl.make_block_ptr(dA+(i_bh+i_v*n_bh)*T*BT, (T, BT), (BT, 1), (i_t * BT + i_i * BC, i_j * BC), (BC, BC), (1, 0)) + o_kj = i_t * BT + i_j * BC + tl.arange(0, BC) + o_Aj = i_j * BC + tl.arange(0, BC) + o_zn = (i_t * BT + i_i * BC) * V + o_v + p_v = v + i_bh * T*V + o_v[:, None] + o_kj[None, :] * V + p_z = z + i_bh * T*V + o_r[:, None] * V + o_v[None, :] + p_zn = z + i_bh * T*V + o_zn + p_do = do + i_bh * T*V + o_r[:, None] * V + o_v[None, :] + p_dA = dA+(i_bh+i_v*n_bh)*T*BT + o_r[:, None] * BT + o_Aj[None, :] # [BV,] - b_zn = tl.load(p_zn, boundary_check=(0,)) + b_zn = tl.load(p_zn, mask=o_zn < T*V, other=0.0) # [BC, BV] - b_z = tl.load(p_z, boundary_check=(0, 1)) - b_do = tl.load(p_do, boundary_check=(0, 1)) + b_z = tl.load(p_z, mask=m_rv, other=0.0) + b_do = tl.load(p_do, mask=m_rv, other=0.0) b_do = (b_do * exp(b_zn[None, :] - b_z) * scale).to(b_do.dtype) # [BV, BC] - b_v = tl.load(p_v, boundary_check=(0, 1)) + b_v = tl.load(p_v, mask=(o_v[:, None] < V) & (o_kj[None, :] < T), other=0.0) b_v = exp(b_v - b_zn[:, None]).to(b_v.dtype) # [BC, BC] b_dA = tl.dot(b_do, b_v, allow_tf32=False) - tl.store(p_dA, b_dA.to(dA.dtype.element_ty), boundary_check=(0, 1)) + tl.store(p_dA, b_dA.to(dA.dtype.element_ty), mask=(o_r[:, None] < T) & (o_Aj[None, :] < BT)) elif i_i == i_j: - p_v = tl.make_block_ptr(v + i_bh * T*V, (T * V,), (1,), ((i_t * BT + i_j * BC) * V + i_v * BV,), (BV,), (0,)) - p_z = tl.make_block_ptr(z + i_bh * T*V, (T, V), (V, 1), (i_t * BT + i_i * BC, i_v * BV), (BC, BV), (1, 0)) - p_do = tl.make_block_ptr(do + i_bh * T*V, (T, V), (V, 1), (i_t * BT + i_i * BC, i_v * BV), (BC, BV), (1, 0)) + p_z = z + i_bh * T*V + o_r[:, None] * V + o_v[None, :] + p_do = do + i_bh * T*V + o_r[:, None] * V + o_v[None, :] # [BC, BV] - b_z = tl.load(p_z, boundary_check=(0, 1)) - b_do = tl.load(p_do, boundary_check=(0, 1)) * scale + b_z = tl.load(p_z, mask=m_rv, other=0.0) + b_do = tl.load(p_do, mask=m_rv, other=0.0) * scale o_i = tl.arange(0, BC) o_A = (i_bh + i_v * n_bh) * T * BT + (i_t * BT + i_i * BC + tl.arange(0, BC)) * BT + i_j * BC m_A = (i_t * BT + i_i * BC + tl.arange(0, BC)) < T for j in range(0, BC): + o_vj = (i_t * BT + i_j * BC + j) * V + o_v + p_vj = v + i_bh * T*V + o_vj # [BV,] - b_v = tl.load(p_v, boundary_check=(0,)).to(tl.float32) + b_v = tl.load(p_vj, mask=o_vj < T*V, other=0.0).to(tl.float32) # [BC,] b_dA = tl.sum(b_do * exp(b_v[None, :] - b_z), 1) b_dA = tl.where(o_i >= j, b_dA, 0) tl.store(dA + o_A + j, b_dA.to(b_do.dtype), mask=m_A) - p_v = tl.advance(p_v, (V,)) - @triton.jit(do_not_specialize=['T']) def chunk_abc_bwd_kernel_K( @@ -624,71 +715,84 @@ def chunk_abc_bwd_kernel_K( BV: tl.constexpr, NT: tl.constexpr, ): - i_k, i_t, i_bh = tl.program_id(0), tl.program_id(1), tl.program_id(2) + i_k, i_t = unflatten_program_id(tl.cdiv(K, BK)) + i_bh = tl.program_id(1).to(tl.int64) i_p = tl.maximum(i_t * BT - 1, 0) - n_bh = tl.num_programs(2) + n_bh = tl.num_programs(1).to(tl.int64) o_i = tl.arange(0, BT) m_s = o_i[:, None] >= o_i[None, :] - p_q = tl.make_block_ptr(q + i_bh * T*K, (T, K), (K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) - p_k = tl.make_block_ptr(k + i_bh * T*K, (T, K), (K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) - p_A = tl.make_block_ptr(A + (i_k*n_bh+i_bh) * T * BT, (T, BT), (BT, 1), (i_t * BT, 0), (BT, BT), (1, 0)) + o_t = i_t * BT + tl.arange(0, BT) + o_k = i_k * BK + tl.arange(0, BK) + m_t = o_t < T + m_tk = m_t[:, None] & (o_k[None, :] < K) + o_A = tl.arange(0, BT) + m_AT = m_t[:, None] & (o_A[None, :] < BT) + p_q = q + i_bh * T*K + o_t[:, None] * K + o_k[None, :] + p_k = k + i_bh * T*K + o_t[:, None] * K + o_k[None, :] + p_A = A + (i_k*n_bh+i_bh) * T * BT + o_t[:, None] * BT + o_A[None, :] # [BT, BK] - b_q = tl.load(p_q, boundary_check=(0, 1)) - b_k = tl.load(p_k, boundary_check=(0, 1)) + b_q = tl.load(p_q, mask=m_tk, other=0.0) + b_k = tl.load(p_k, mask=m_tk, other=0.0) # [BT, BT] b_A = tl.dot((b_q * scale).to(b_q.dtype), tl.trans(b_k), allow_tf32=False) b_A = tl.where(m_s, b_A, 0.) - tl.store(p_A, b_A.to(p_A.dtype.element_ty), boundary_check=(0, 1)) + tl.store(p_A, b_A.to(p_A.dtype.element_ty), mask=m_AT) b_dq = tl.zeros([BT, BK], dtype=tl.float32) b_dk = tl.zeros([BT, BK], dtype=tl.float32) for i_v in range(tl.cdiv(V, BV)): - p_v = tl.make_block_ptr(v + i_bh * T*V, (T, V), (V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) - p_z = tl.make_block_ptr(z + i_bh * T*V, (T, V), (V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) - p_zp = tl.make_block_ptr(z + i_bh * T*V, (T * V,), (1,), (i_p * V + i_v * BV,), (BV,), (0,)) - p_zc = tl.make_block_ptr(z + i_bh * T*V, (T * V,), (1,), ((i_t * BT + BT - 1) * V + i_v * BV,), (BV,), (0,)) - p_h = tl.make_block_ptr(h + i_bh * NT*K*V + i_t * K*V, (V, K), (1, V), (i_v * BV, i_k * BK), (BV, BK), (0, 1)) - - p_do = tl.make_block_ptr(do + i_bh * T*V, (T, V), (V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) - p_dh = tl.make_block_ptr(dh + i_bh * NT*K*V + i_t * K*V, (K, V), (V, 1), (i_k * BK, i_v * BV), (BK, BV), (1, 0)) - p_dv = tl.make_block_ptr(dv + (i_k*n_bh+i_bh) * T*V, (T, V), (V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) + o_v = i_v * BV + tl.arange(0, BV) + m_tv = m_t[:, None] & (o_v[None, :] < V) + m_h = (o_v[:, None] < V) & (o_k[None, :] < K) + m_kv = (o_k[:, None] < K) & (o_v[None, :] < V) + o_zp = i_p * V + o_v + o_zc = (i_t * BT + BT - 1) * V + o_v + p_v = v + i_bh * T*V + o_t[:, None] * V + o_v[None, :] + p_z = z + i_bh * T*V + o_t[:, None] * V + o_v[None, :] + p_zp = z + i_bh * T*V + o_zp + p_zc = z + i_bh * T*V + o_zc + p_h = h + i_bh * NT*K*V + i_t * K*V + o_v[:, None] + o_k[None, :] * V + + p_do = do + i_bh * T*V + o_t[:, None] * V + o_v[None, :] + p_dh = dh + i_bh * NT*K*V + i_t * K*V + o_k[:, None] * V + o_v[None, :] + p_dv = dv + (i_k*n_bh+i_bh) * T*V + o_t[:, None] * V + o_v[None, :] # [BV,] - b_zp = tl.load(p_zp, boundary_check=(0,)) - b_zc = tl.load(p_zc, boundary_check=(0,)) + b_zp = tl.load(p_zp, mask=o_zp < T*V, other=0.0) + b_zc = tl.load(p_zc, mask=o_zc < T*V, other=0.0) # [BT, BV] - b_v = tl.load(p_v, boundary_check=(0, 1)) + b_v = tl.load(p_v, mask=m_tv, other=0.0) b_v = exp(b_v - b_zc[None, :]).to(b_v.dtype) - b_z = tl.load(p_z, boundary_check=(0, 1)) + b_z = tl.load(p_z, mask=m_tv, other=0.0) b_z = exp(b_zp[None, :] - b_z) # [BV, BK] - b_h = tl.load(p_h, boundary_check=(0, 1)) + b_h = tl.load(p_h, mask=m_h, other=0.0) # [BT, BV] - b_do = tl.load(p_do, boundary_check=(0, 1)) + b_do = tl.load(p_do, mask=m_tv, other=0.0) b_do = (b_do * b_z * scale).to(b_do.dtype) # [BK, BV] - b_dh = tl.load(p_dh, boundary_check=(0, 1)) + b_dh = tl.load(p_dh, mask=m_kv, other=0.0) # [BT, BK] - b_dq += tl.dot(b_do, b_h, allow_tf32=False) - b_dk += tl.dot(b_v, tl.trans(b_dh), allow_tf32=False) + b_dq = tl.dot(b_do, b_h, b_dq, allow_tf32=False) + b_dk = tl.dot(b_v, tl.trans(b_dh), b_dk, allow_tf32=False) # [BT, BV] b_dv = b_v * tl.dot(b_k, b_dh, allow_tf32=False) - tl.store(p_dv, b_dv.to(p_dv.dtype.element_ty), boundary_check=(0, 1)) - p_dA = tl.make_block_ptr(dA + i_bh * T * BT, (T, BT), (BT, 1), (i_t * BT, 0), (BT, BT), (1, 0)) + tl.store(p_dv, b_dv.to(p_dv.dtype.element_ty), mask=m_tv) + p_dA = dA + i_bh * T * BT + o_t[:, None] * BT + o_A[None, :] # [BT, BT] - b_dA = tl.load(p_dA, boundary_check=(0, 1)) + b_dA = tl.load(p_dA, mask=m_AT, other=0.0) # [BT, BK] - b_dq += tl.dot(b_dA, b_k, allow_tf32=False) - b_dk += tl.dot(tl.trans(b_dA).to(b_k.dtype), b_q, allow_tf32=False) + b_dq = tl.dot(b_dA, b_k, b_dq, allow_tf32=False) + b_dk = tl.dot(tl.trans(b_dA).to(b_k.dtype), b_q, b_dk, allow_tf32=False) - p_dq = tl.make_block_ptr(dq + i_bh * T*K, (T, K), (K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) - p_dk = tl.make_block_ptr(dk + i_bh * T*K, (T, K), (K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) - tl.store(p_dq, b_dq.to(p_dq.dtype.element_ty), boundary_check=(0, 1)) - tl.store(p_dk, b_dk.to(p_dk.dtype.element_ty), boundary_check=(0, 1)) + p_dq = dq + i_bh * T*K + o_t[:, None] * K + o_k[None, :] + p_dk = dk + i_bh * T*K + o_t[:, None] * K + o_k[None, :] + tl.store(p_dq, b_dq.to(p_dq.dtype.element_ty), mask=m_tk) + tl.store(p_dk, b_dk.to(p_dk.dtype.element_ty), mask=m_tk) @triton.jit(do_not_specialize=['T']) @@ -705,44 +809,55 @@ def chunk_abc_bwd_kernel_intra_KV( BV: tl.constexpr, NC: tl.constexpr, ): - i_v, i_c, i_bh = tl.program_id(0), tl.program_id(1), tl.program_id(2) + NV = tl.cdiv(V, BV) + i_v, i_c, i_bh = tl.program_id(0).to(tl.int64) % NV, tl.program_id(0).to(tl.int64) // NV, tl.program_id(1).to(tl.int64) i_t, i_i = i_c // NC, i_c % NC - p_v = tl.make_block_ptr(v + i_bh * T*V, (T, V), (V, 1), (i_t * BT + i_i * BC, i_v * BV), (BC, BV), (1, 0)) - p_zn = tl.make_block_ptr(z + i_bh * T*V, (T*V,), (1,), ((i_t * BT + i_i * BC + BC - 1) * V + i_v * BV,), (BV,), (0,)) + o_r = i_t * BT + i_i * BC + tl.arange(0, BC) + o_v = i_v * BV + tl.arange(0, BV) + m_rv = (o_r[:, None] < T) & (o_v[None, :] < V) + o_zn = (i_t * BT + i_i * BC + BC - 1) * V + o_v + p_v = v + i_bh * T*V + o_r[:, None] * V + o_v[None, :] + p_zn = z + i_bh * T*V + o_zn # [BV,] - b_zn = tl.load(p_zn, boundary_check=(0,)) + b_zn = tl.load(p_zn, mask=o_zn < T*V, other=0.0) # [BC, BV] - b_v = tl.load(p_v, boundary_check=(0, 1)) + b_v = tl.load(p_v, mask=m_rv, other=0.0) b_dv = tl.zeros([BC, BV], dtype=tl.float32) for i_j in range(i_i + 1, NC): - p_z = tl.make_block_ptr(z + i_bh * T*V, (T, V), (V, 1), (i_t * BT + i_j * BC, i_v * BV), (BC, BV), (1, 0)) - p_A = tl.make_block_ptr(A + i_bh * T * BT, (BT, T), (1, BT), (i_i * BC, i_t * BT + i_j * BC), (BC, BC), (0, 1)) - p_do = tl.make_block_ptr(do + i_bh * T*V, (T, V), (V, 1), (i_t * BT + i_j * BC, i_v * BV), (BC, BV), (1, 0)) + o_j = i_t * BT + i_j * BC + tl.arange(0, BC) + o_Ai = i_i * BC + tl.arange(0, BC) + m_jv = (o_j[:, None] < T) & (o_v[None, :] < V) + m_A = (o_Ai[:, None] < BT) & (o_j[None, :] < T) + p_z = z + i_bh * T*V + o_j[:, None] * V + o_v[None, :] + p_A = A + i_bh * T * BT + o_Ai[:, None] + o_j[None, :] * BT + p_do = do + i_bh * T*V + o_j[:, None] * V + o_v[None, :] # [BC, BV] - b_z = tl.load(p_z, boundary_check=(0, 1)) - b_do = tl.load(p_do, boundary_check=(0, 1)) + b_z = tl.load(p_z, mask=m_jv, other=0.0) + b_do = tl.load(p_do, mask=m_jv, other=0.0) b_do = (b_do * exp(b_zn[None, :] - b_z)).to(b_do.dtype) # [BC, BC] - b_A = tl.load(p_A, boundary_check=(0, 1)) - b_dv += tl.dot(b_A, b_do, allow_tf32=False) + b_A = tl.load(p_A, mask=m_A, other=0.0) + b_dv = tl.dot(b_A, b_do, b_dv, allow_tf32=False) b_dv *= exp(b_v - b_zn[None, :]) o_i = tl.arange(0, BC) for j in range(0, BC): - p_z = tl.make_block_ptr(z + i_bh * T*V, (T * V,), (1,), ((i_t * BT + i_i * BC + j) * V + i_v * BV,), (BV,), (0,)) - p_A = tl.make_block_ptr(A + i_bh * T * BT, (T * BT,), (1,), ((i_t * BT + i_i * BC + j) * BT + i_i * BC,), (BC,), (0,)) - p_do = tl.make_block_ptr(do + i_bh * T*V, (T * V,), (1,), ((i_t * BT + i_i * BC + j) * V + i_v * BV,), (BV,), (0,)) + o_vj = (i_t * BT + i_i * BC + j) * V + o_v + o_Aj = (i_t * BT + i_i * BC + j) * BT + i_i * BC + tl.arange(0, BC) + p_zj = z + i_bh * T*V + o_vj + p_Aj = A + i_bh * T * BT + o_Aj + p_doj = do + i_bh * T*V + o_vj # [BC,] - b_A = tl.load(p_A, boundary_check=(0,)) + b_A = tl.load(p_Aj, mask=o_Aj < T*BT, other=0.0) # [BV,] - b_z = tl.load(p_z, boundary_check=(0,)) - b_do = tl.load(p_do, boundary_check=(0,)) + b_z = tl.load(p_zj, mask=o_vj < T*V, other=0.0) + b_do = tl.load(p_doj, mask=o_vj < T*V, other=0.0) # [BC, BV] m_i = o_i[:, None] <= j b_dv += tl.where(m_i, exp(b_v - b_z[None, :]) * b_A[:, None] * b_do[None, :], 0.) - p_dv = tl.make_block_ptr(dv + i_bh * T*V, (T, V), (V, 1), (i_t * BT + i_i * BC, i_v * BV), (BC, BV), (1, 0)) - tl.store(p_dv, b_dv.to(p_dv.dtype.element_ty), boundary_check=(0, 1)) + p_dv = dv + i_bh * T*V + o_r[:, None] * V + o_v[None, :] + tl.store(p_dv, b_dv.to(p_dv.dtype.element_ty), mask=m_rv) @triton.jit(do_not_specialize=['T']) @@ -757,25 +872,29 @@ def chunk_abc_bwd_kernel_rcum_inter( BS: tl.constexpr, NT: tl.constexpr, ): - i_m, i_bh = tl.program_id(0), tl.program_id(1) + i_m, i_bh = tl.program_id(0).to(tl.int64), tl.program_id(1).to(tl.int64) + o_s = i_m * BS + tl.arange(0, BS) b_sp = tl.zeros([BS], dtype=tl.float32) b_zp = tl.full([BS], float('inf'), dtype=tl.float32) for i_t in range(NT - 1, -1, -1): - p_s = tl.make_block_ptr(s + i_bh * T*S, (T, S), (S, 1), (i_t * BT, i_m * BS), (BT, BS), (1, 0)) - p_z = tl.make_block_ptr(z + i_bh * T*S, (T, S), (S, 1), (i_t * BT, i_m * BS), (BT, BS), (1, 0)) - p_zc = tl.make_block_ptr(z + i_bh * T*S, (T*S,), (1,), ((i_t * BT) * S + i_m * BS,), (BS,), (0,)) - p_ss = tl.make_block_ptr(ss + i_bh * T*S, (T, S), (S, 1), (i_t * BT, i_m * BS), (BT, BS), (1, 0)) - p_doo = tl.make_block_ptr(doo + i_bh * T*S, (T, S), (S, 1), (i_t * BT, i_m * BS), (BT, BS), (1, 0)) + o_t = i_t.to(tl.int64) * BT + tl.arange(0, BT) + m_ts = (o_t[:, None] < T) & (o_s[None, :] < S) + o_zc = i_t.to(tl.int64) * BT * S + o_s + p_s = s + i_bh * T*S + o_t[:, None] * S + o_s[None, :] + p_z = z + i_bh * T*S + o_t[:, None] * S + o_s[None, :] + p_zc = z + i_bh * T*S + o_zc + p_ss = ss + i_bh * T*S + o_t[:, None] * S + o_s[None, :] + p_doo = doo + i_bh * T*S + o_t[:, None] * S + o_s[None, :] # [BS,] - b_zc = tl.load(p_zc, boundary_check=(0,)) + b_zc = tl.load(p_zc, mask=o_zc < T*S, other=0.0) # [BT, BS] - b_s = tl.load(p_s, boundary_check=(0, 1)) - b_z = tl.load(p_z, boundary_check=(0, 1)) - b_ss = tl.load(p_ss, boundary_check=(0, 1)) + b_s = tl.load(p_s, mask=m_ts, other=0.0) + b_z = tl.load(p_z, mask=m_ts, other=0.0) + b_ss = tl.load(p_ss, mask=m_ts, other=0.0) b_doo = exp(b_s - b_zp[None, :]) * b_sp[None, :] - tl.store(p_doo, b_doo.to(p_doo.dtype.element_ty), boundary_check=(0, 1)) + tl.store(p_doo, b_doo.to(p_doo.dtype.element_ty), mask=m_ts) # [BS,] b_sp = b_sp * exp(b_zc - b_zp) + tl.sum(b_ss * exp(b_zc[None, :] - b_z), 0) b_zp = b_zc @@ -794,42 +913,50 @@ def chunk_abc_bwd_kernel_rcum_intra( BS: tl.constexpr, NC: tl.constexpr, ): - i_s, i_c, i_bh = tl.program_id(0), tl.program_id(1), tl.program_id(2) + NS = tl.cdiv(S, BS) + i_s, i_c, i_bh = tl.program_id(0).to(tl.int64) % NS, tl.program_id(0).to(tl.int64) // NS, tl.program_id(1).to(tl.int64) i_t, i_i = i_c // NC, i_c % NC o_i = tl.arange(0, BC) m_o = tl.full([BC, BC], 1., dtype=tl.float32) - p_s = tl.make_block_ptr(s + i_bh * T*S, (T, S), (S, 1), (i_t * BT + i_i * BC, i_s * BS), (BC, BS), (1, 0)) - p_zn = tl.make_block_ptr(z + i_bh * T*S, (T*S,), (1,), ((i_t * BT + i_i * BC + BC - 1) * S + i_s * BS,), (BS,), (0,)) - p_doo = tl.make_block_ptr(doo + i_bh * T*S, (T, S), (S, 1), (i_t * BT + i_i * BC, i_s * BS), (BC, BS), (1, 0)) + o_r = i_t * BT + i_i * BC + tl.arange(0, BC) + o_s = i_s * BS + tl.arange(0, BS) + m_rs = (o_r[:, None] < T) & (o_s[None, :] < S) + o_zn = (i_t * BT + i_i * BC + BC - 1) * S + o_s + p_s = s + i_bh * T*S + o_r[:, None] * S + o_s[None, :] + p_zn = z + i_bh * T*S + o_zn + p_doo = doo + i_bh * T*S + o_r[:, None] * S + o_s[None, :] # [BC, BS] - b_s = tl.load(p_s, boundary_check=(0, 1)) + b_s = tl.load(p_s, mask=m_rs, other=0.0) # [BS,] - b_zn = tl.load(p_zn, boundary_check=(0,)) + b_zn = tl.load(p_zn, mask=o_zn < T*S, other=0.0) b_doo = tl.zeros([BC, BS], dtype=tl.float32) for i_j in range(i_i + 1, NC): - p_z = tl.make_block_ptr(z + i_bh * T*S, (T, S), (S, 1), (i_t * BT + i_j * BC, i_s * BS), (BC, BS), (1, 0)) - p_ss = tl.make_block_ptr(ss + i_bh * T*S, (T, S), (S, 1), (i_t * BT + i_j * BC, i_s * BS), (BC, BS), (1, 0)) + o_j = i_t * BT + i_j * BC + tl.arange(0, BC) + m_js = (o_j[:, None] < T) & (o_s[None, :] < S) + p_z = z + i_bh * T*S + o_j[:, None] * S + o_s[None, :] + p_ss = ss + i_bh * T*S + o_j[:, None] * S + o_s[None, :] # [BC, BS] - b_z = tl.load(p_z, boundary_check=(0, 1)) - b_ss = tl.load(p_ss, boundary_check=(0, 1)) + b_z = tl.load(p_z, mask=m_js, other=0.0) + b_ss = tl.load(p_ss, mask=m_js, other=0.0) # [BC, BS] b_doo += b_ss * exp(b_zn[None, :] - b_z) b_doo = exp(b_s - b_zn[None, :]) * tl.dot(m_o.to(b_s.dtype), b_doo.to(b_s.dtype), allow_tf32=False) for j in range(0, BC): - p_z = tl.make_block_ptr(z + i_bh * T*S, (T*S,), (1,), ((i_t * BT + i_i * BC + j) * S + i_s * BS,), (BS,), (0,)) - p_ss = tl.make_block_ptr(ss + i_bh * T*S, (T*S,), (1,), ((i_t * BT + i_i * BC + j) * S + i_s * BS,), (BS,), (0,)) + o_sj = (i_t * BT + i_i * BC + j) * S + o_s + p_zj = z + i_bh * T*S + o_sj + p_ssj = ss + i_bh * T*S + o_sj # [BS,] - b_z = tl.load(p_z, boundary_check=(0,)) - b_ss = tl.load(p_ss, boundary_check=(0,)) + b_z = tl.load(p_zj, mask=o_sj < T*S, other=0.0) + b_ss = tl.load(p_ssj, mask=o_sj < T*S, other=0.0) # [BC, BS] m_i = o_i[:, None] <= j b_doo += tl.where(m_i, exp(b_s - b_z[None, :]) * b_ss[None, :], 0.) - b_doo += tl.load(p_doo, boundary_check=(0, 1)) - tl.store(p_doo, b_doo.to(p_doo.dtype.element_ty), boundary_check=(0, 1)) + b_doo += tl.load(p_doo, mask=m_rs, other=0.0) + tl.store(p_doo, b_doo.to(p_doo.dtype.element_ty), mask=m_rs) class ChunkABCFunction(torch.autograd.Function): @@ -848,7 +975,7 @@ class ChunkABCFunction(torch.autograd.Function): num_stages = 1 def fwd_pre(s, B, H, T, S): - # keep cummulative normalizer in fp32 + # keep cumulative normalizer in fp32 z = torch.empty_like(s, dtype=torch.float) grid = (B * H,) logcumsumexp_fwd_kernel[grid]( @@ -888,7 +1015,7 @@ class ChunkABCFunction(torch.autograd.Function): ) ok1 = torch.empty_like(s) Ak = q.new_empty(B, H, T, BT) - grid = (NM, NT, B * H) + grid = (NM * NT, B * H) chunk_abc_fwd_kernel_K[grid]( q, k, z, hk, ok1, Ak, scale=scale, @@ -897,7 +1024,7 @@ class ChunkABCFunction(torch.autograd.Function): num_stages=num_stages, ) ok0 = torch.empty_like(s) - grid = (NM, NT * NC, B * H) + grid = (NM * NT * NC, B * H) chunk_abc_fwd_kernel_intra_K[grid]( s, z, ok0, Ak, T=T, V=M, BT=BT, BC=BC, BV=BM, NC=NC, @@ -920,7 +1047,7 @@ class ChunkABCFunction(torch.autograd.Function): ht=final_state[1] if final_state is not None else None, ) Av = q.new_zeros(NM, B, H, T, BT) - grid = (NM, NT * NC * NC, B * H) + grid = (NM * NT * NC * NC, B * H) chunk_abc_fwd_kernel_intra_V[grid]( qv, s, z, Av, scale=scale, @@ -930,7 +1057,7 @@ class ChunkABCFunction(torch.autograd.Function): ) Av = Av.sum(0) ov = torch.empty_like(v) - grid = (NV, NT, B * H) + grid = (NV * NT, B * H) chunk_abc_fwd_kernel_V[grid]( qv, v, z, hv, ov, Av, scale=scale, @@ -985,7 +1112,7 @@ class ChunkABCFunction(torch.autograd.Function): num_warps=num_warps, num_stages=num_stages, ) - grid = (NS, NT * NC, B * H) + grid = (NS * NT * NC, B * H) chunk_abc_bwd_kernel_rcum_intra[grid]( s, z, ss, doo, T=T, S=S, BT=BT, BC=BC, BS=BS, NC=NC, @@ -1006,7 +1133,7 @@ class ChunkABCFunction(torch.autograd.Function): dsv1 = torch.empty_like(s, dtype=torch.float) dv = v.new_empty(NM, *v.shape) dAv = q.new_zeros(B, H, T, BT) - grid = (NM, NT, B * H) + grid = (NM * NT, B * H) chunk_abc_bwd_kernel_V[grid]( s, v, z, hv, Av, dov, dhv, dp1, dsv1, dv, dAv, scale=scale, @@ -1017,7 +1144,7 @@ class ChunkABCFunction(torch.autograd.Function): dv = dv.sum(0) dp0 = torch.empty_like(p) dsv0 = s.new_zeros(s.shape, dtype=torch.float) - grid = (NM, NT * NC, B * H) + grid = (NM * NT * NC, B * H) chunk_abc_bwd_kernel_intra_V[grid]( qv, s, z, dAv, dp0, dsv0, T=T, K=M, BT=BT, BC=BC, BK=BM, NC=NC, @@ -1039,7 +1166,7 @@ class ChunkABCFunction(torch.autograd.Function): normk=False, ) dAk = q.new_zeros(NM, B, H, T, BT) - grid = (NM, NT * NC * NC, B * H) + grid = (NM * NT * NC * NC, B * H) chunk_abc_bwd_kernel_intra_K[grid]( s, z, dok, dAk, scale=scale, @@ -1053,7 +1180,7 @@ class ChunkABCFunction(torch.autograd.Function): dq = torch.empty_like(q) dk = torch.empty_like(k) dsk1 = s.new_empty(NK, *s.shape, dtype=torch.float) - grid = (NK, NT, B * H) + grid = (NK * NT, B * H) chunk_abc_bwd_kernel_K[grid]( q, k, s, z, hk, Ak, dok, dhk, dq, dk, dsk1, dAk, scale=scale, @@ -1064,7 +1191,7 @@ class ChunkABCFunction(torch.autograd.Function): Ak = Ak.sum(0) dsk1 = dsk1.sum(0) dsk0 = torch.empty_like(s, dtype=torch.float) - grid = (NM, NT * NC, B * H) + grid = (NM * NT * NC, B * H) chunk_abc_bwd_kernel_intra_KV[grid]( s, z, Ak, dok, dsk0, T=T, V=M, BT=BT, BC=BC, BV=BM, NC=NC, @@ -1085,7 +1212,6 @@ def chunk_abc( s: torch.Tensor, initial_state: tuple[torch.Tensor] | None = None, output_final_state: bool = False, - head_first: bool = False, ) -> tuple[torch.Tensor, torch.Tensor]: r""" Args: @@ -1101,9 +1227,6 @@ def chunk_abc( Initial states of shape `[B, H, K, M]` and `[B, H, M, V]`. Default: `None`. output_final_state (Optional[bool]): Whether to output the final state of shape `[B, H, K, M]` and `[B, H, M, V]`. Default: `False`. - head_first (Optional[bool]): - Whether the inputs are in the head-first format. Default: `False`. - This argument has been deprecated. Returns: o (torch.Tensor): @@ -1111,9 +1234,7 @@ def chunk_abc( final_state (torch.Tensor): Final state of shape `[B, H, K, M]` and `[B, H, M, V]` if `output_final_state=True` else `None`. """ - if not head_first: - q, k, v, s = map(lambda x: x.transpose(1, 2), (q, k, v, s)) + q, k, v, s = map(lambda x: x.transpose(1, 2), (q, k, v, s)) o, final_state = ChunkABCFunction.apply(q, k, v, s, initial_state, output_final_state) - if not head_first: - o = o.transpose(1, 2) + o = o.transpose(1, 2) return o, final_state diff --git a/build/torch-cuda/ops/atk/__init__.py b/build/torch-cuda/ops/atk/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..2cd5c7fac4c52136819b113e5926c2ae04132ca8 --- /dev/null +++ b/build/torch-cuda/ops/atk/__init__.py @@ -0,0 +1,14 @@ +# 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 + +# ATK kernels +# Shared by precond_gated_delta_rule and precond_kda + +from ...ops.atk.chunk_atk_bwd import chunk_atk_bwd +from ...ops.atk.chunk_atk_fwd import chunk_atk_fwd, recompute_atk_fwd + +__all__ = ['chunk_atk_bwd', 'chunk_atk_fwd', 'recompute_atk_fwd'] diff --git a/build/torch-xpu/ops/abc/__init__.py b/build/torch-cuda/ops/atk/backends/__init__.py similarity index 83% rename from build/torch-xpu/ops/abc/__init__.py rename to build/torch-cuda/ops/atk/backends/__init__.py index 680a9a73daae95b958534e7859f13d7d32206561..2739122d6c5e3dcb83c0a7b3ca16d818a9741a74 100644 --- a/build/torch-xpu/ops/abc/__init__.py +++ b/build/torch-cuda/ops/atk/backends/__init__.py @@ -5,8 +5,7 @@ # For a list of all contributors, visit: # https://github.com/fla-org/flash-linear-attention/graphs/contributors -from .chunk import chunk_abc +"""ATK backends.""" -__all__ = [ - 'chunk_abc', -] + +__all__: list[str] = [] diff --git a/build/torch-cuda/ops/atk/chunk_atk_bwd.py b/build/torch-cuda/ops/atk/chunk_atk_bwd.py new file mode 100644 index 0000000000000000000000000000000000000000..7697503b2545fd45fce3aa456569f3b98fcaa1ae --- /dev/null +++ b/build/torch-cuda/ops/atk/chunk_atk_bwd.py @@ -0,0 +1,576 @@ +# 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 + +# Copyright (c) 2023-2025 + +import math +import os + +import torch +import triton +import triton.language as tl + +from ...backends import dispatch +from ...ops.utils import prepare_chunk_indices + + +@triton.heuristics({ + 'USE_INITIAL_STATE': lambda args: args['h0'] is not None, + 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None, +}) +@triton.jit(do_not_specialize=['T']) +def _atk_backward_chunk_out( + k, # *f32 [B, T, H, D] + log_g, # *f32 [B, T, H] + beta, # *f32 [B, T, H] - beta scaling + ac, # *f32 [B, C, H, D] (forward prefix contexts) + h0, # *f32 [N, H, D] or None - initial ATK state + g_kp, # *f32 [B, T, H, D] upstream grad on k_precond + # outputs (accumulators; +=) + gk_out, # *f32 [B, T, H, D] + gg_out, # *f32 [B, T, H] + gbeta_out, # *f32 [B, T, H] - gradient for beta + gac_prev, # *f32 [B, C, H, D] grad wrt ac_{i-1}, stored at index i-1 + dh0, # *f32 [N, H, D] or None - gradient for initial ATK state + g_log_atk_scale_out, # *f32 [H] - per-head log_atk_scale gradients (atomic add) + cu_seqlens, # *i32 [N+1] - cumulative sequence lengths + chunk_indices, # *i32 [NT, 2] - (seq_idx, chunk_idx) pairs + log_atk_scale, # *f32 [H] - per-head log-space center (learnable or fixed) + logx, # scalar float32 - log(x) for squash range + eps, # scalar float32 - epsilon for log safety + B: tl.constexpr, T, H: tl.constexpr, D: tl.constexpr, + CHUNK_LEN: tl.constexpr, + k_stride_b, k_stride_t, k_stride_h, k_stride_d, + log_g_stride_b, log_g_stride_t, log_g_stride_h, + beta_stride_b, beta_stride_t, beta_stride_h, + ac_stride_b, ac_stride_c, ac_stride_h, ac_stride_d, + gkp_stride_b, gkp_stride_t, gkp_stride_h, gkp_stride_d, + gk_stride_b, gk_stride_t, gk_stride_h, gk_stride_d, + gg_stride_b, gg_stride_t, gg_stride_h, + gbeta_stride_b, gbeta_stride_t, gbeta_stride_h, + gac_stride_b, gac_stride_c, gac_stride_h, gac_stride_d, + BK: tl.constexpr, + USE_INITIAL_STATE: tl.constexpr, + IS_VARLEN: tl.constexpr, +): + """ + Per-chunk backward kernel for ATK. K-tiled variant. + """ + i_t = tl.program_id(0).to(tl.int64) + h = tl.program_id(1).to(tl.int64) + chunk_id = tl.program_id(2).to(tl.int64) + + if IS_VARLEN: + i_n = tl.load(chunk_indices + i_t * 2).to(tl.int32) + chunk_id = tl.load(chunk_indices + i_t * 2 + 1).to(tl.int64) + bos = tl.load(cu_seqlens + i_n).to(tl.int32) + eos = tl.load(cu_seqlens + i_n + 1).to(tl.int32) + T = eos - bos + b = i_n.to(tl.int64) + else: + b = tl.program_id(0).to(tl.int64) + i_n = b + bos = 0 + eos = T + + if h >= H: + return + + if chunk_id * CHUNK_LEN >= T: + return + + T_range = chunk_id * CHUNK_LEN + tl.arange(0, CHUNK_LEN) + mask_T = T_range < T + C_range = tl.arange(0, CHUNK_LEN) + + if IS_VARLEN: + log_g_ptr = log_g + h * log_g_stride_h + (log_g_stride_t * (bos + T_range)) + beta_ptr = beta + h * beta_stride_h + (beta_stride_t * (bos + T_range)) + else: + log_g_ptr = log_g + b * log_g_stride_b + h * log_g_stride_h + (log_g_stride_t * T_range) + beta_ptr = beta + b * beta_stride_b + h * beta_stride_h + (beta_stride_t * T_range) + + log_g_val = tl.load(log_g_ptr, mask=mask_T, other=0.0).to(tl.float32) + beta_val = tl.load(beta_ptr, mask=mask_T, other=0.0).to(tl.float32) + g_val = tl.exp(log_g_val) + + la_cumsum = tl.cumsum(log_g_val) + + roll_mat = (C_range[:, None] == (C_range[None, :] + 1)).to(tl.float32) + la_cumsum_roll = tl.sum(roll_mat[:, :] * la_cumsum[None, :], 1) + M = tl.exp(la_cumsum_roll[:, None] - la_cumsum[None, :]) + M = tl.where(C_range[:, None] > C_range[None, :], M, 0.0) + + base_decays = tl.exp(la_cumsum_roll * (C_range > 0).to(tl.float32)) + + center = tl.load(log_atk_scale + h).to(tl.float32) + + gg_val = tl.zeros([CHUNK_LEN], dtype=tl.float32) + gbeta_val = tl.zeros([CHUNK_LEN], dtype=tl.float32) + g_center = tl.zeros([], dtype=tl.float32) + gbase_decays = tl.zeros([CHUNK_LEN], dtype=tl.float32) + gM = tl.zeros([CHUNK_LEN, CHUNK_LEN], dtype=tl.float32) + + for i_k in range(tl.cdiv(D, BK)): + d_offset = i_k * BK + D_range = d_offset + tl.arange(0, BK) + mask_D = D_range < D + + if IS_VARLEN: + k_ptr = k + h * k_stride_h + (k_stride_t * (bos + T_range))[:, None] + (k_stride_d) * D_range[None, :] + gkp_ptr = g_kp + h * gkp_stride_h + (gkp_stride_t * (bos + T_range))[:, None] + (gkp_stride_d) * D_range[None, :] + else: + k_ptr = k + b * k_stride_b + h * k_stride_h + (k_stride_t * T_range)[:, None] + (k_stride_d) * D_range[None, :] + gkp_ptr = g_kp + b * gkp_stride_b + h * gkp_stride_h + \ + (gkp_stride_t * T_range)[:, None] + (gkp_stride_d) * D_range[None, :] + + k_val = tl.load(k_ptr, mask=mask_T[:, None] * mask_D[None, :], other=0.0).to(tl.float32) + gkp_val = tl.load(gkp_ptr, mask=mask_T[:, None] * mask_D[None, :], other=0.0).to(tl.float32) + + k_sq = k_val * k_val + U = beta_val[:, None] * k_sq + + ac_ptr = ac + b * ac_stride_b + h * ac_stride_h + (chunk_id - 1) * ac_stride_c + D_range * ac_stride_d + if chunk_id == 0: + if USE_INITIAL_STATE: + ac_val = tl.load(h0 + (i_n * H + h) * D + D_range, mask=mask_D, other=0).to(tl.float32) + else: + ac_val = tl.zeros([BK], dtype=tl.float32) + else: + ac_val = tl.load(ac_ptr, mask=mask_D) + + raw_state = base_decays[:, None] * ac_val[None, :] + tl.dot(M, U) + A_t = g_val[:, None] * raw_state + U + + ell = tl.log(A_t + eps) + r = ell - center + abs_r = tl.abs(r) + one_plus_abs_r = 1.0 + abs_r + s = r / one_plus_abs_r + ds_dr = 1.0 / (one_plus_abs_r * one_plus_abs_r) + M_mult = tl.exp(-logx * s) + + gk_val = gkp_val * M_mult + gM_from_kprecond = gkp_val * k_val + gs = gM_from_kprecond * (-logx * M_mult) + + gr_local = gs * ds_dr + g_center += -tl.sum(gr_local) + gA_t = gr_local / (A_t + eps) + graw_state = gA_t * g_val[:, None] + gg_val += tl.sum(gA_t * raw_state, 1) + gU = gA_t + + gU_from_M = tl.dot(tl.trans(M), graw_state) + + gk_sq = (gU + gU_from_M) * beta_val[:, None] + gbeta_val += tl.sum(gU * k_sq, 1) + tl.sum(gU_from_M * k_sq, 1) + gk_val += 2 * k_val * gk_sq + + gac_tile = tl.sum(graw_state * base_decays[:, None], 0) + gac_ptr = gac_prev + i_n * gac_stride_b + h * gac_stride_h + (chunk_id - 1) * gac_stride_c + D_range * gac_stride_d + if chunk_id > 0: + tl.store(gac_ptr, gac_tile, mask=mask_D) + elif USE_INITIAL_STATE: + tl.store(dh0 + (i_n * H + h) * D + D_range, gac_tile, mask=mask_D) + + gbase_decays += tl.sum(graw_state * ac_val[None, :], 1) + gM = tl.dot(graw_state, tl.trans(U), gM) + + if IS_VARLEN: + gk_ptr = gk_out + (bos + T_range)[:, None] * gk_stride_t + h * gk_stride_h + D_range[None, :] * gk_stride_d + else: + gk_ptr = gk_out + b * gk_stride_b + T_range[:, None] * \ + gk_stride_t + h * gk_stride_h + D_range[None, :] * gk_stride_d + tl.store(gk_ptr, gk_val, mask=mask_T[:, None] * mask_D[None, :]) + + glog_g_val = gg_val * g_val + + gbase_decays = gbase_decays * (C_range > 0) + gla_cumsum_roll = gbase_decays * base_decays + + gM = tl.where(C_range[:, None] > C_range[None, :], gM, 0.0) + ginner_sum = M * gM + gla_cumsum_roll += tl.sum(ginner_sum, 1) + gla_cumsum = -1 * tl.sum(ginner_sum, 0) + + roll_mat = (C_range[:, None] == (C_range[None, :] + 1)).to(tl.float32) + gla_cumsum += tl.sum(roll_mat * gla_cumsum_roll[:, None], 0) * (C_range < (CHUNK_LEN - 1)).to(tl.float32) + + glog_g_val += tl.cumsum(gla_cumsum, reverse=True) + + if IS_VARLEN: + gg_ptr = gg_out + (bos + T_range) * gg_stride_t + h * gg_stride_h + gbeta_ptr = gbeta_out + (bos + T_range) * gbeta_stride_t + h * gbeta_stride_h + else: + gg_ptr = gg_out + b * gg_stride_b + T_range * gg_stride_t + h * gg_stride_h + gbeta_ptr = gbeta_out + b * gbeta_stride_b + T_range * gbeta_stride_t + h * gbeta_stride_h + + tl.store(gg_ptr, glog_g_val, mask=mask_T) + tl.store(gbeta_ptr, gbeta_val, mask=mask_T) + tl.atomic_add(g_log_atk_scale_out + h, g_center) + + +@triton.heuristics({ + 'USE_INITIAL_STATE': lambda args: args['h0'] is not None, + 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None, +}) +@triton.jit(do_not_specialize=['T']) +def _atk_backward_pass_chunks( + a, sa, ac, # forward buffers + h0, # *f32 [N, H, D] or None - initial ATK state + gac_from_out, # grad entering each ac[i] from later usage + ga, gsa, # outputs (+=) + dh0, # *f32 [N, H, D] or None - gradient for initial ATK state (accumulated) + cu_seqlens, # *i32 [N+1] - cumulative sequence lengths + B: tl.constexpr, T, H: tl.constexpr, D: tl.constexpr, + CHUNK_LEN: tl.constexpr, + a_stride_b, a_stride_c, a_stride_h, a_stride_d, + sa_stride_b, sa_stride_c, sa_stride_h, + ac_stride_b, ac_stride_c, ac_stride_h, ac_stride_d, + gac_stride_b, gac_stride_c, gac_stride_h, gac_stride_d, + ga_stride_b, ga_stride_c, ga_stride_h, ga_stride_d, + gsa_stride_b, gsa_stride_c, gsa_stride_h, + BLOCK_D: tl.constexpr, + USE_INITIAL_STATE: tl.constexpr, + IS_VARLEN: tl.constexpr, +): + """ + Sequential backward pass across chunks. + Not K-tiled due to cross-D reduction in gsa_val computation. + """ + i_nh = tl.program_id(0).to(tl.int64) + + if IS_VARLEN: + i_n = i_nh // H + h = i_nh % H + bos = tl.load(cu_seqlens + i_n).to(tl.int32) + eos = tl.load(cu_seqlens + i_n + 1).to(tl.int32) + T = eos - bos + b = i_n.to(tl.int64) + else: + b = i_nh // H + h = i_nh % H + i_n = b + if b >= B: + return + + if h >= H: + return + + N_chunks = tl.cdiv(T, CHUNK_LEN) + D_range = tl.arange(0, BLOCK_D) + D_mask = D_range < D + + sa_ptr = sa + b * sa_stride_b + h * sa_stride_h + (N_chunks - 1) * sa_stride_c + a_ptr = a + b * a_stride_b + h * a_stride_h + D_range * a_stride_d + (N_chunks - 1) * a_stride_c + ac_ptr = ac + b * ac_stride_b + h * ac_stride_h + D_range * ac_stride_d + (N_chunks - 1) * ac_stride_c + gac_ptr = gac_from_out + b * gac_stride_b + h * gac_stride_h + D_range * gac_stride_d + (N_chunks - 1) * gac_stride_c + gsa_ptr = gsa + b * gsa_stride_b + h * gsa_stride_h + (N_chunks - 1) * gsa_stride_c + ga_ptr = ga + b * ga_stride_b + h * ga_stride_h + D_range * ga_stride_d + (N_chunks - 1) * ga_stride_c + + gac_val = tl.zeros([BLOCK_D], dtype=tl.float32) + + for chunk_id in tl.range(N_chunks - 1, -1, -1): + gac_val += tl.load(gac_ptr, mask=D_mask) + + tl.store(ga_ptr, gac_val, mask=D_mask) + + sa_val = tl.load(sa_ptr).to(tl.float32) + + if chunk_id > 0: + ac_prev_ptr = ac + b * ac_stride_b + h * ac_stride_h + D_range * ac_stride_d + (chunk_id - 1) * ac_stride_c + ac_prev_val = tl.load(ac_prev_ptr, mask=D_mask) + gsa_val = tl.sum(gac_val * ac_prev_val) * tl.exp(sa_val) + else: + if USE_INITIAL_STATE: + h0_val = tl.load(h0 + i_nh * D + D_range, mask=D_mask, other=0).to(tl.float32) + gsa_val = tl.sum(gac_val * h0_val) * tl.exp(sa_val) + else: + gsa_val = 0.0 + + gac_val = gac_val * tl.exp(sa_val) + + tl.store(gsa_ptr, gsa_val) + + sa_ptr -= sa_stride_c + a_ptr -= a_stride_c + ac_ptr -= ac_stride_c + gac_ptr -= gac_stride_c + ga_ptr -= ga_stride_c + gsa_ptr -= gsa_stride_c + + if USE_INITIAL_STATE: + p_dh0 = dh0 + i_nh * D + D_range + tl.store(p_dh0, tl.load(p_dh0, mask=D_mask, other=0).to(tl.float32) + gac_val, mask=D_mask) + + +@triton.heuristics({ + 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None, +}) +@triton.jit(do_not_specialize=['T']) +def _atk_backward_chunk_summary( + k, log_g, beta, # inputs for recompute + ga, gsa, # grads from scan-backward + gk_out, gg_out, gbeta_out, + cu_seqlens, # *i32 [N+1] - cumulative sequence lengths + chunk_indices, # *i32 [NT, 2] - (seq_idx, chunk_idx) pairs + B: tl.constexpr, T, H: tl.constexpr, D: tl.constexpr, + CHUNK_LEN: tl.constexpr, + k_stride_b, k_stride_t, k_stride_h, k_stride_d, + lg_stride_b, lg_stride_t, lg_stride_h, + beta_stride_b, beta_stride_t, beta_stride_h, + ga_stride_b, ga_stride_c, ga_stride_h, ga_stride_d, + gsa_stride_b, gsa_stride_c, gsa_stride_h, + gk_stride_b, gk_stride_t, gk_stride_h, gk_stride_d, + gg_stride_b, gg_stride_t, gg_stride_h, + gbeta_stride_b, gbeta_stride_t, gbeta_stride_h, + BK: tl.constexpr, + IS_VARLEN: tl.constexpr, +): + i_t = tl.program_id(0).to(tl.int64) + h = tl.program_id(1).to(tl.int64) + chunk_id = tl.program_id(2).to(tl.int64) + + if IS_VARLEN: + i_n = tl.load(chunk_indices + i_t * 2).to(tl.int32) + chunk_id = tl.load(chunk_indices + i_t * 2 + 1).to(tl.int64) + bos = tl.load(cu_seqlens + i_n).to(tl.int32) + eos = tl.load(cu_seqlens + i_n + 1).to(tl.int32) + T = eos - bos + b = i_n.to(tl.int64) + else: + b = tl.program_id(0).to(tl.int64) + i_n = b + bos = 0 + eos = T + + if h >= H: + return + + if chunk_id * CHUNK_LEN >= T: + return + + T_range = chunk_id * CHUNK_LEN + tl.arange(0, CHUNK_LEN) + mask_T = T_range < T + + if IS_VARLEN: + log_g_ptr = log_g + (bos + T_range) * lg_stride_t + h * lg_stride_h + beta_ptr = beta + (bos + T_range) * beta_stride_t + h * beta_stride_h + else: + log_g_ptr = log_g + b * lg_stride_b + T_range * lg_stride_t + h * lg_stride_h + beta_ptr = beta + b * beta_stride_b + T_range * beta_stride_t + h * beta_stride_h + + log_g_val = tl.load(log_g_ptr, mask=mask_T, other=0.0).to(tl.float32) + beta_val = tl.load(beta_ptr, mask=mask_T, other=0.0).to(tl.float32) + + sa_val = tl.sum(log_g_val) + decays = tl.exp(sa_val - tl.cumsum(log_g_val)) + + gsa_ptr = gsa + i_n * gsa_stride_b + h * gsa_stride_h + chunk_id * gsa_stride_c + gsa_val = tl.load(gsa_ptr).to(tl.float32) + + gdecays = tl.zeros([CHUNK_LEN], dtype=tl.float32) + gbeta_val = tl.zeros([CHUNK_LEN], dtype=tl.float32) + + for i_k in range(tl.cdiv(D, BK)): + d_offset = i_k * BK + D_range = d_offset + tl.arange(0, BK) + mask_D = D_range < D + + if IS_VARLEN: + k_ptr = k + (bos + T_range)[:, None] * k_stride_t + h * k_stride_h + D_range[None, :] * k_stride_d + gk_ptr = gk_out + (bos + T_range)[:, None] * gk_stride_t + h * gk_stride_h + D_range[None, :] * gk_stride_d + else: + k_ptr = k + b * k_stride_b + T_range[:, None] * k_stride_t + h * k_stride_h + D_range[None, :] * k_stride_d + gk_ptr = gk_out + b * gk_stride_b + T_range[:, None] * \ + gk_stride_t + h * gk_stride_h + D_range[None, :] * gk_stride_d + + ga_ptr = ga + i_n * ga_stride_b + h * ga_stride_h + chunk_id * ga_stride_c + (ga_stride_d) * D_range + + k_val = tl.load(k_ptr, mask=mask_T[:, None] * mask_D[None, :], other=0.0).to(tl.float32) + ga_val = tl.load(ga_ptr, mask=mask_D, other=0.0).to(tl.float32) + + k_sq = k_val * k_val + U = beta_val[:, None] * k_sq + + gdecays += tl.sum(ga_val[None, :] * U, 1) + + gU = ga_val[None, :] * decays[:, None] + gbeta_val += tl.sum(gU * k_sq, 1) + + gk_sq = gU * beta_val[:, None] + gk_val = gk_sq * k_val * 2 + + tl.atomic_add(gk_ptr, gk_val, mask=mask_T[:, None] * mask_D[None, :]) + + gdecays_exp = decays * gdecays + + gsa_val += tl.sum(gdecays_exp) + glog_g_val = -1 * tl.cumsum(gdecays_exp, reverse=True) + + glog_g_val += gsa_val + + if IS_VARLEN: + gg_ptr = gg_out + (bos + T_range) * gg_stride_t + h * gg_stride_h + gbeta_ptr = gbeta_out + (bos + T_range) * gbeta_stride_t + h * gbeta_stride_h + else: + gg_ptr = gg_out + b * gg_stride_b + T_range * gg_stride_t + h * gg_stride_h + gbeta_ptr = gbeta_out + b * gbeta_stride_b + T_range * gbeta_stride_t + h * gbeta_stride_h + + tl.atomic_add(gg_ptr, glog_g_val, mask=mask_T) + tl.atomic_add(gbeta_ptr, gbeta_val, mask=mask_T) + + +@dispatch('atk') +def chunk_atk_bwd( + k: torch.Tensor, + g_raw: torch.Tensor, + beta: torch.Tensor, + dk_precond: torch.Tensor, + ac: torch.Tensor, + a: torch.Tensor, + sa: torch.Tensor, + chunk_size: int = 64, + initial_A_state: torch.Tensor | None = None, + cu_seqlens: torch.LongTensor | None = None, + x: float = 1.5, + eps: float = 1e-6, + log_atk_scale: torch.Tensor = None, + dat: torch.Tensor | None = None, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor | None]: + r""" + ATK backward pass. Takes precomputed forward intermediates from ``recompute_atk_fwd``. + + Args: + k: Keys ``[B, T, H, K]``. + g_raw: Raw gate values (NOT cumsum) ``[B, T, H]``. + beta: Beta scaling ``[B, T, H]``. + dk_precond: Gradient w.r.t. k_precond ``[B, T, H, K]``. + ac: Prefix-sum ATK states from ``recompute_atk_fwd`` ``[B, C, H, K]``. + a: Per-chunk summaries from ``recompute_atk_fwd`` ``[B, C, H, K]``. + sa: Per-chunk log-gate sums from ``recompute_atk_fwd`` ``[B, C, H]``. + chunk_size: Chunk size. + initial_A_state: Initial ATK diagonal state ``[N, H, K]``. + cu_seqlens: Cumulative sequence lengths ``[N+1]`` for varlen. + x: Squash range parameter. + eps: Epsilon for numerical stability. + log_atk_scale: Per-head log-space center ``[H]``. + dat: Incoming gradient w.r.t. the final ATK state ``[N, H, K]`` or ``None``. + + Returns: + dk_atk: Gradient contribution to k ``[B, T, H, K]``. + dbeta_atk: Gradient for beta ``[B, T, H]``. + dg_atk: Gradient for g ``[B, T, H]``. + d_log_atk_scale: Gradient for log_atk_scale ``[H]``. + dh0_atk: Gradient for initial_A_state ``[N, H, K]`` or ``None``. + """ + B, T, H, K = k.shape + CHUNK_LEN = chunk_size + + is_varlen = cu_seqlens is not None + chunk_indices = prepare_chunk_indices(cu_seqlens, CHUNK_LEN) if is_varlen else None + NT = triton.cdiv(T, CHUNK_LEN) if not is_varlen else len(chunk_indices) + N = len(cu_seqlens) - 1 if is_varlen else B + + BK = K if os.environ.get('ATK_NO_KTILE') else 32 # K-tile size (set ATK_NO_KTILE=1 to disable) + BLOCK_D = triton.next_power_of_2(K) + + k = k.contiguous() + g_raw = g_raw.contiguous() + beta = beta.contiguous() + dk_precond = dk_precond.contiguous() + + if log_atk_scale is None: + log_atk_scale = torch.full((H,), -0.2, dtype=torch.float32, device=k.device) + else: + log_atk_scale = log_atk_scale.contiguous() + + logx = math.log(x) if x > 0 else 0.0 + + gk = torch.zeros_like(k, dtype=torch.float32) + g_log_atk_scale = torch.zeros(H, device=k.device, dtype=torch.float32) + gg = torch.zeros_like(g_raw, dtype=torch.float32) + gbeta = torch.zeros_like(beta, dtype=torch.float32) + gac_prev = torch.zeros_like(ac, dtype=torch.float32) + if dat is not None: + # The final ATK state `at` equals `ac[:, last_chunk]`, so its incoming + # gradient enters the reverse scan as the carry at the final chunk; + # `gac_prev[:, last_chunk]` is never written by the local-backward kernel. + if cu_seqlens is None: + gac_prev[:, NT - 1] += dat + else: + seqlens = cu_seqlens[1:] - cu_seqlens[:-1] + last_chunk = torch.clamp((seqlens + CHUNK_LEN - 1) // CHUNK_LEN - 1, min=0) + valid = seqlens > 0 + idx = torch.arange(N, device=k.device) + gac_prev[idx[valid], last_chunk[valid]] += dat[valid] + dh0 = torch.zeros(N, H, K, device=k.device, dtype=torch.float32) if initial_A_state is not None else None + ga = torch.empty_like(a, dtype=torch.float32) + gsa = torch.empty_like(sa, dtype=torch.float32) + + if cu_seqlens is None: + grid = (B, H, triton.cdiv(T, CHUNK_LEN)) + grid2 = (B * H,) + else: + grid = (NT, H, 1) + grid2 = (N * H,) + + _atk_backward_chunk_out[grid]( + k, g_raw, beta, ac, initial_A_state, dk_precond, + gk, gg, gbeta, gac_prev, dh0, + g_log_atk_scale, + cu_seqlens, chunk_indices, + log_atk_scale, logx, eps, + B, T, H, K, CHUNK_LEN, + k.stride(0), k.stride(1), k.stride(2), k.stride(3), + g_raw.stride(0), g_raw.stride(1), g_raw.stride(2), + beta.stride(0), beta.stride(1), beta.stride(2), + ac.stride(0), ac.stride(1), ac.stride(2), ac.stride(3), + dk_precond.stride(0), dk_precond.stride(1), dk_precond.stride(2), dk_precond.stride(3), + gk.stride(0), gk.stride(1), gk.stride(2), gk.stride(3), + gg.stride(0), gg.stride(1), gg.stride(2), + gbeta.stride(0), gbeta.stride(1), gbeta.stride(2), + gac_prev.stride(0), gac_prev.stride(1), gac_prev.stride(2), gac_prev.stride(3), + BK=BK, num_warps=4, + ) + + _atk_backward_pass_chunks[grid2]( + a, sa, ac, + initial_A_state, + gac_prev, + ga, gsa, + dh0, + cu_seqlens, + B, T, H, K, CHUNK_LEN, + a.stride(0), a.stride(1), a.stride(2), a.stride(3), + sa.stride(0), sa.stride(1), sa.stride(2), + ac.stride(0), ac.stride(1), ac.stride(2), ac.stride(3), + gac_prev.stride(0), gac_prev.stride(1), gac_prev.stride(2), gac_prev.stride(3), + ga.stride(0), ga.stride(1), ga.stride(2), ga.stride(3), + gsa.stride(0), gsa.stride(1), gsa.stride(2), + BLOCK_D, num_warps=4 + ) + + _atk_backward_chunk_summary[grid]( + k, g_raw, beta, + ga, gsa, + gk, gg, gbeta, + cu_seqlens, chunk_indices, + B, T, H, K, CHUNK_LEN, + k.stride(0), k.stride(1), k.stride(2), k.stride(3), + g_raw.stride(0), g_raw.stride(1), g_raw.stride(2), + beta.stride(0), beta.stride(1), beta.stride(2), + ga.stride(0), ga.stride(1), ga.stride(2), ga.stride(3), + gsa.stride(0), gsa.stride(1), gsa.stride(2), + gk.stride(0), gk.stride(1), gk.stride(2), gk.stride(3), + gg.stride(0), gg.stride(1), gg.stride(2), + gbeta.stride(0), gbeta.stride(1), gbeta.stride(2), + BK, num_warps=4 + ) + + return gk, gbeta, gg, g_log_atk_scale, dh0 diff --git a/build/torch-cuda/ops/atk/chunk_atk_fwd.py b/build/torch-cuda/ops/atk/chunk_atk_fwd.py new file mode 100644 index 0000000000000000000000000000000000000000..529d0c829e9ef4105bf8510c956ee52018812fbb --- /dev/null +++ b/build/torch-cuda/ops/atk/chunk_atk_fwd.py @@ -0,0 +1,467 @@ +# 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 + +""" +Chunked ATK forward kernel for preconditioned delta rules. +K-tiled variant: processes head dimension in BK blocks. + +ATK recurrence: + A_t = exp(g) * A_{t-1} + beta * k^2 + +Preconditioner (symmetric fast squash): + ell = log(A_t + eps) + r = ell - log_atk_scale # deviation from learned center + s = r / (1 + |r|) # fast squash, bounded in (-1, 1) + M = exp(-log(x) * s) # bounded multiplier in [1/x, x] + k_precond = k * M +""" +import math +import os + +import torch +import triton +import triton.language as tl + +from ...backends import dispatch + + +@triton.heuristics({ + 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None +}) +@triton.jit(do_not_specialize=['T']) +def _forward_chunk_summary( + k, # *f32 [B, T, H, D] + beta, # *f32 [B, T, H] + log_g, # *f32 [B, T, H] + a, # *f32 [B, C, H, D] + sa, # *f32 [B, C, H] + cu_seqlens, # *i32 [N+1] - cumulative sequence lengths (None if not varlen) + chunk_indices, # *i32 [NT, 2] - (sequence_idx, chunk_idx) pairs (None if not varlen) + B: tl.constexpr, + T, + H: tl.constexpr, + D: tl.constexpr, + CHUNK_LEN: tl.constexpr, + k_stride_b, k_stride_t, k_stride_h, k_stride_d, + beta_stride_b, beta_stride_t, beta_stride_h, + log_g_stride_b, log_g_stride_t, log_g_stride_h, + a_stride_b, a_stride_c, a_stride_h, a_stride_d, + sa_stride_b, sa_stride_c, sa_stride_h, + BK: tl.constexpr, + IS_VARLEN: tl.constexpr, +): + i_t = tl.program_id(0).to(tl.int64) + h = tl.program_id(1).to(tl.int64) + chunk_id = tl.program_id(2).to(tl.int64) + + if IS_VARLEN: + i_n = tl.load(chunk_indices + i_t * 2).to(tl.int32) + chunk_id = tl.load(chunk_indices + i_t * 2 + 1).to(tl.int64) + bos = tl.load(cu_seqlens + i_n).to(tl.int32) + eos = tl.load(cu_seqlens + i_n + 1).to(tl.int32) + T = eos - bos + b = i_n.to(tl.int64) + else: + b = tl.program_id(0).to(tl.int64) + bos = 0 + eos = T + + if h >= H: + return + + if chunk_id * CHUNK_LEN >= T: + return + + T_range = chunk_id * CHUNK_LEN + tl.arange(0, CHUNK_LEN) + mask_T = T_range < T + + if IS_VARLEN: + beta_ptr = beta + (bos + T_range) * beta_stride_t + h * beta_stride_h + log_g_ptr = log_g + (bos + T_range) * log_g_stride_t + h * log_g_stride_h + else: + beta_ptr = beta + b * beta_stride_b + T_range * beta_stride_t + h * beta_stride_h + log_g_ptr = log_g + b * log_g_stride_b + T_range * log_g_stride_t + h * log_g_stride_h + + beta_val = tl.load(beta_ptr, mask=mask_T, other=0.0).to(tl.float32) + log_g_val = tl.load(log_g_ptr, mask=mask_T, other=0.0).to(tl.float32) + + sa_val = tl.sum(log_g_val) + decays = tl.exp(sa_val - tl.cumsum(log_g_val)) + + for i_k in range(tl.cdiv(D, BK)): + d_offset = i_k * BK + D_range = d_offset + tl.arange(0, BK) + mask_D = D_range < D + + if IS_VARLEN: + k_ptr = k + (bos + T_range)[:, None] * k_stride_t + h * k_stride_h + D_range[None, :] * k_stride_d + else: + k_ptr = k + b * k_stride_b + T_range[:, None] * k_stride_t + h * k_stride_h + D_range[None, :] * k_stride_d + + k_val = tl.load(k_ptr, mask=mask_T[:, None] * mask_D[None, :], other=0.0).to(tl.float32) + + k_sq = k_val * k_val + U = beta_val[:, None] * k_sq + a_val = tl.sum(decays[:, None] * U, 0) + + a_ptr = a + b * a_stride_b + h * a_stride_h + chunk_id * a_stride_c + D_range * a_stride_d + tl.store(a_ptr, a_val, mask=mask_D) + + sa_ptr = sa + b * sa_stride_b + h * sa_stride_h + chunk_id * sa_stride_c + tl.store(sa_ptr, sa_val) + + +@triton.heuristics({ + 'USE_INITIAL_STATE': lambda args: args['h0'] is not None, + 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None, +}) +@triton.jit(do_not_specialize=['T']) +def _forward_pass_chunks( + a, # *f32 [B, C, H, D] + sa, # *f32 [B, C, H] + ac, # *f32 [B, C, H, D] + h0, # *f32 [N, H, D] or None + cu_seqlens, # *i32 [N+1] + B: tl.constexpr, + T, + H: tl.constexpr, + D: tl.constexpr, + CHUNK_LEN: tl.constexpr, + a_stride_b, a_stride_c, a_stride_h, a_stride_d, + sa_stride_b, sa_stride_c, sa_stride_h, + ac_stride_b, ac_stride_c, ac_stride_h, ac_stride_d, + BK: tl.constexpr, + USE_INITIAL_STATE: tl.constexpr, + IS_VARLEN: tl.constexpr, +): + i_nh = tl.program_id(0).to(tl.int64) + + if IS_VARLEN: + i_n = i_nh // H + h = i_nh % H + bos = tl.load(cu_seqlens + i_n).to(tl.int32) + eos = tl.load(cu_seqlens + i_n + 1).to(tl.int32) + T = eos - bos + b = i_n.to(tl.int64) + else: + b = i_nh // H + h = i_nh % H + if b >= B: + return + + if h >= H: + return + + N_chunks = tl.cdiv(T, CHUNK_LEN) + + for i_k in range(tl.cdiv(D, BK)): + d_offset = i_k * BK + D_range = d_offset + tl.arange(0, BK) + D_mask = D_range < D + + sa_ptr = sa + b * sa_stride_b + h * sa_stride_h + a_ptr = a + b * a_stride_b + h * a_stride_h + D_range * a_stride_d + ac_ptr = ac + b * ac_stride_b + h * ac_stride_h + D_range * ac_stride_d + + if USE_INITIAL_STATE: + ac_val = tl.load(h0 + i_nh * D + D_range, mask=D_mask, other=0).to(tl.float32) + else: + ac_val = tl.zeros([BK], dtype=tl.float32) + + for chunk_id in tl.range(N_chunks): + a_val = tl.load(a_ptr, D_mask).to(tl.float32) + sa_val = tl.load(sa_ptr).to(tl.float32) + + ac_val = tl.exp(sa_val) * ac_val + a_val + + tl.store(ac_ptr, ac_val, mask=D_mask) + + sa_ptr += sa_stride_c + a_ptr += a_stride_c + ac_ptr += ac_stride_c + + +@triton.heuristics({ + 'USE_INITIAL_STATE': lambda args: args['h0'] is not None, + 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None, +}) +@triton.jit(do_not_specialize=['T']) +def _forward_chunk_out( + k, # *f32 [B, T, H, D] + beta, # *f32 [B, T, H] + log_g, # *f32 [B, T, H] + ac, # *f32 [B, C, H, D] + h0, # *f32 [N, H, D] or None + k_precond, # *f32 [B, T, H, D] + log_atk_scale, # *f32 [H] - per-head log-space center (learnable or fixed) + logx, # scalar float32 - log(x) for squash range + eps, # scalar float32 - epsilon for log safety + cu_seqlens, + chunk_indices, + B: tl.constexpr, + T, + H: tl.constexpr, + D: tl.constexpr, + CHUNK_LEN: tl.constexpr, + k_stride_b, k_stride_t, k_stride_h, k_stride_d, + beta_stride_b, beta_stride_t, beta_stride_h, + log_g_stride_b, log_g_stride_t, log_g_stride_h, + ac_stride_b, ac_stride_c, ac_stride_h, ac_stride_d, + kp_stride_b, kp_stride_t, kp_stride_h, kp_stride_d, + BK: tl.constexpr, + USE_INITIAL_STATE: tl.constexpr, + IS_VARLEN: tl.constexpr, +): + i_t = tl.program_id(0).to(tl.int64) + h = tl.program_id(1).to(tl.int64) + chunk_id = tl.program_id(2).to(tl.int64) + + if IS_VARLEN: + i_n = tl.load(chunk_indices + i_t * 2).to(tl.int32) + chunk_id = tl.load(chunk_indices + i_t * 2 + 1).to(tl.int64) + bos = tl.load(cu_seqlens + i_n).to(tl.int32) + eos = tl.load(cu_seqlens + i_n + 1).to(tl.int32) + T = eos - bos + b = i_n.to(tl.int64) + else: + b = tl.program_id(0).to(tl.int64) + bos = 0 + eos = T + + if h >= H: + return + + if chunk_id * CHUNK_LEN >= T: + return + + T_range = chunk_id * CHUNK_LEN + tl.arange(0, CHUNK_LEN) + mask_T = T_range < T + C_range = tl.arange(0, CHUNK_LEN) + + if IS_VARLEN: + beta_ptr = beta + (bos + T_range) * beta_stride_t + h * beta_stride_h + log_g_ptr = log_g + (bos + T_range) * log_g_stride_t + h * log_g_stride_h + else: + beta_ptr = beta + b * beta_stride_b + T_range * beta_stride_t + h * beta_stride_h + log_g_ptr = log_g + b * log_g_stride_b + h * log_g_stride_h + (log_g_stride_t * T_range) + + beta_val = tl.load(beta_ptr, mask=mask_T, other=0.0).to(tl.float32) + log_g_val = tl.load(log_g_ptr, mask=mask_T, other=0.0).to(tl.float32) + g_val = tl.exp(log_g_val) + + la_cumsum = tl.cumsum(log_g_val) + + roll_mat = (C_range[:, None] == (C_range[None, :] + 1)).to(tl.float32) + la_cumsum_roll = tl.sum(roll_mat[:, :] * la_cumsum[None, :], 1) + M = tl.exp(la_cumsum_roll[:, None] - la_cumsum[None, :]) + M = tl.where(C_range[:, None] > C_range[None, :], M, 0.0) + + base_decays = tl.exp(la_cumsum_roll * (C_range > 0).to(tl.float32)) + + center = tl.load(log_atk_scale + h).to(tl.float32) + + for i_k in range(tl.cdiv(D, BK)): + d_offset = i_k * BK + D_range = d_offset + tl.arange(0, BK) + mask_D = D_range < D + + if IS_VARLEN: + k_ptr = k + (bos + T_range)[:, None] * k_stride_t + h * k_stride_h + D_range[None, :] * k_stride_d + kp_ptr = k_precond + (bos + T_range)[:, None] * kp_stride_t + h * kp_stride_h + D_range[None, :] * kp_stride_d + else: + k_ptr = k + b * k_stride_b + h * k_stride_h + (k_stride_t * T_range)[:, None] + (k_stride_d) * D_range[None, :] + kp_ptr = k_precond + b * kp_stride_b + h * kp_stride_h + \ + (kp_stride_t * T_range)[:, None] + (kp_stride_d) * D_range[None, :] + + k_val = tl.load(k_ptr, mask=mask_T[:, None] * mask_D[None, :], other=0.0).to(tl.float32) + + k_sq = k_val * k_val + U = beta_val[:, None] * k_sq + + ac_ptr = ac + b * ac_stride_b + h * ac_stride_h + (chunk_id - 1) * ac_stride_c + D_range * ac_stride_d + + if chunk_id == 0: + if USE_INITIAL_STATE: + ac_val = tl.load(h0 + (b * H + h) * D + D_range, mask=mask_D, other=0).to(tl.float32) + else: + ac_val = tl.zeros([BK], dtype=tl.float32) + else: + ac_val = tl.load(ac_ptr, mask=mask_D) + + raw_state = base_decays[:, None] * ac_val[None, :] + tl.dot(M, U) + + A_t = g_val[:, None] * raw_state + U + + ell = tl.log(A_t + eps) + r = ell - center + s = r / (1.0 + tl.abs(r)) + M_precond = tl.exp(-logx * s) + k_precond_val = k_val * M_precond + + tl.store(kp_ptr, k_precond_val, mask=mask_T[:, None] * mask_D[None, :]) + + +def _atk_fwd_stages( + k: torch.Tensor, + beta: torch.Tensor, + log_g: torch.Tensor, + chunk_size: int, + initial_A_state: torch.Tensor | None, + output_final_state: bool, + cu_seqlens: torch.Tensor | None, + x: float, + eps: float, + log_atk_scale: torch.Tensor | None, +): + """Run the 3-stage ATK forward algorithm. + + Returns: + k_precond, ac, a, sa, at + """ + B, T, H, D = k.shape + CHUNK_LEN = chunk_size + + BK = D if os.environ.get('ATK_NO_KTILE') else 32 # K-tile size (set ATK_NO_KTILE=1 to disable) + + is_varlen = cu_seqlens is not None + + if is_varlen: + from ...ops.utils import prepare_chunk_indices + N = len(cu_seqlens) - 1 + chunk_indices = prepare_chunk_indices(cu_seqlens, CHUNK_LEN) + NT = len(chunk_indices) + + seq_lens = [cu_seqlens[i + 1] - cu_seqlens[i] for i in range(N)] + max_chunks = max(triton.cdiv(s.item(), CHUNK_LEN) for s in seq_lens) + + a = torch.empty(N, max_chunks, H, D, dtype=torch.float32, device=k.device) + sa = torch.empty(N, max_chunks, H, dtype=torch.float32, device=k.device) + ac = torch.empty(N, max_chunks, H, D, dtype=torch.float32, device=k.device) + else: + N = B + num_chunks = math.ceil(T / CHUNK_LEN) + chunk_indices = None + NT = B * num_chunks + + a = torch.empty(B, num_chunks, H, D, dtype=torch.float32, device=k.device) + sa = torch.empty(B, num_chunks, H, dtype=torch.float32, device=k.device) + ac = torch.empty(B, num_chunks, H, D, dtype=torch.float32, device=k.device) + + k_precond = torch.empty_like(k) + + k = k.contiguous() + beta = beta.contiguous() + log_g = log_g.contiguous() + + if log_atk_scale is None: + log_atk_scale = torch.full((H,), -0.2, dtype=torch.float32, device=k.device) + else: + log_atk_scale = log_atk_scale.contiguous() + + logx = math.log(x) if x > 0 else 0.0 + + if is_varlen: + grid = (NT, H, 1) + grid2 = (N * H,) + else: + grid = (B, H, num_chunks) + grid2 = (B * H,) + + _forward_chunk_summary[grid]( + k, beta, log_g, a, sa, + cu_seqlens, chunk_indices, + B, T, H, D, CHUNK_LEN, + k.stride(0), k.stride(1), k.stride(2), k.stride(3), + beta.stride(0), beta.stride(1), beta.stride(2), + log_g.stride(0), log_g.stride(1), log_g.stride(2), + a.stride(0), a.stride(1), a.stride(2), a.stride(3), + sa.stride(0), sa.stride(1), sa.stride(2), + BK, num_warps=4 + ) + + _forward_pass_chunks[grid2]( + a, sa, ac, + initial_A_state, + cu_seqlens, + B if not is_varlen else N, T, H, D, CHUNK_LEN, + a.stride(0), a.stride(1), a.stride(2), a.stride(3), + sa.stride(0), sa.stride(1), sa.stride(2), + ac.stride(0), ac.stride(1), ac.stride(2), ac.stride(3), + BK, num_warps=4 + ) + + _forward_chunk_out[grid]( + k, beta, log_g, ac, initial_A_state, k_precond, + log_atk_scale, logx, eps, + cu_seqlens, chunk_indices, + B, T, H, D, CHUNK_LEN, + k.stride(0), k.stride(1), k.stride(2), k.stride(3), + beta.stride(0), beta.stride(1), beta.stride(2), + log_g.stride(0), log_g.stride(1), log_g.stride(2), + ac.stride(0), ac.stride(1), ac.stride(2), ac.stride(3), + k_precond.stride(0), k_precond.stride(1), k_precond.stride(2), k_precond.stride(3), + BK, + num_warps=4 + ) + + at = None + if output_final_state: + if is_varlen: + seq_lens = [cu_seqlens[i + 1] - cu_seqlens[i] for i in range(N)] + last_chunk_indices = [triton.cdiv(s.item(), CHUNK_LEN) - 1 for s in seq_lens] + at = torch.stack([ac[i, last_chunk_indices[i], :, :] for i in range(N)], dim=0).to(k.dtype) + else: + at = ac[:, -1, :, :].contiguous().to(k.dtype) + + return k_precond.to(k.dtype), ac, a, sa, at + + +@dispatch('atk') +def chunk_atk_fwd( + k: torch.Tensor, + beta: torch.Tensor, + log_g: torch.Tensor = None, + chunk_size: int = 64, + initial_A_state: torch.Tensor = None, + output_final_state: bool = True, + cu_seqlens: torch.Tensor = None, + x: float = 1.5, + eps: float = 1e-6, + log_atk_scale: torch.Tensor = None, +): + r""" + Chunked ATK forward: computes preconditioned keys via 3-stage chunk algorithm. + """ + k_precond, _, _, _, at = _atk_fwd_stages( + k, beta, log_g, chunk_size, + initial_A_state, output_final_state, cu_seqlens, + x, eps, log_atk_scale, + ) + return k_precond, at + + +@dispatch('atk') +def recompute_atk_fwd( + k: torch.Tensor, + beta: torch.Tensor, + log_g: torch.Tensor, + chunk_size: int = 64, + initial_A_state: torch.Tensor = None, + cu_seqlens: torch.Tensor = None, + x: float = 1.5, + eps: float = 1e-6, + log_atk_scale: torch.Tensor = None, +): + r""" + Recompute ATK forward intermediates for use in backward pass. + """ + k_precond, ac, a, sa, _ = _atk_fwd_stages( + k, beta, log_g, chunk_size, + initial_A_state, False, cu_seqlens, + x, eps, log_atk_scale, + ) + return k_precond, ac, a, sa diff --git a/build/torch-cuda/ops/attn/backends/__init__.py b/build/torch-cuda/ops/attn/backends/__init__.py index 76179f136c9df857ad5bbfda568412dc06812e33..b77c5b1277d7030e0f2223d96ff8feb85604b5ba 100644 --- a/build/torch-cuda/ops/attn/backends/__init__.py +++ b/build/torch-cuda/ops/attn/backends/__init__.py @@ -5,7 +5,6 @@ # For a list of all contributors, visit: # https://github.com/fla-org/flash-linear-attention/graphs/contributors -from ....ops.backends import BackendRegistry -from ....ops.common.backends.tilelang import TileLangBackend +from ....ops.attn.backends.tilelang import AttnTileLangBackend -BackendRegistry("attn").register(TileLangBackend()) +__all__ = ['AttnTileLangBackend'] diff --git a/build/torch-cuda/ops/attn/backends/tilelang/__init__.py b/build/torch-cuda/ops/attn/backends/tilelang/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..f34f599e1d38219e262292dff0364ba16b87b615 --- /dev/null +++ b/build/torch-cuda/ops/attn/backends/tilelang/__init__.py @@ -0,0 +1,116 @@ +# 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 torch + +from .....backends import TileLangBackend, register +from .....utils import IS_NVIDIA_HOPPER, TRITON_ABOVE_3_4_0 + + +@register('attn') +class AttnTileLangBackend(TileLangBackend): + # work around Hopper regressions with Triton 3.4+ (see #640). + default_enable = IS_NVIDIA_HOPPER and TRITON_ABOVE_3_4_0 + + def parallel_attn_fwd_verifier( + self, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + g_cumsum: torch.Tensor | None, + sink_bias: torch.Tensor | None, + scale: float, + window_size: int | None = None, + cu_seqlens: torch.LongTensor | None = None, + chunk_indices: torch.LongTensor | None = None, + ) -> tuple[bool, str | None]: + if q.dtype not in (torch.float16, torch.bfloat16, torch.float32): + return False, f"TileLang backend does not support dtype {q.dtype}; fall back to Triton" + return True, None + + def parallel_attn_fwd( + self, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + g_cumsum: torch.Tensor | None, + sink_bias: torch.Tensor | None, + scale: float, + window_size: int | None = None, + cu_seqlens: torch.LongTensor | None = None, + chunk_indices: torch.LongTensor | None = None, + ) -> tuple[torch.Tensor, torch.Tensor]: + from .....ops.attn.backends.tilelang.parallel import parallel_attn_fwd_tilelang + return parallel_attn_fwd_tilelang( + q=q, + k=k, + v=v, + g_cumsum=g_cumsum, + sink_bias=sink_bias, + scale=scale, + window_size=window_size, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + ) + + def parallel_attn_bwd_verifier( + self, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + o: torch.Tensor, + g_cumsum: torch.Tensor | None, + lse: torch.Tensor, + do: torch.Tensor, + sink_bias: torch.Tensor | None = None, + scale: float | None = None, + window_size: int | None = None, + chunk_size: int = 128, + cu_seqlens: torch.LongTensor | None = None, + chunk_indices: torch.LongTensor | None = None, + ) -> tuple[bool, str | None]: + if q.dtype not in (torch.float16, torch.bfloat16, torch.float32): + return False, f"TileLang backend does not support dtype {q.dtype}; fall back to Triton" + return True, None + + def parallel_attn_bwd( + self, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + o: torch.Tensor, + g_cumsum: torch.Tensor | None, + lse: torch.Tensor, + do: torch.Tensor, + sink_bias: torch.Tensor | None = None, + scale: float | None = None, + window_size: int | None = None, + chunk_size: int = 128, + cu_seqlens: torch.LongTensor | None = None, + chunk_indices: torch.LongTensor | None = None, + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor | None, torch.Tensor | None]: + from .....ops.attn.backends.tilelang.parallel import parallel_attn_bwd_tilelang + return parallel_attn_bwd_tilelang( + q=q, + k=k, + v=v, + o=o, + g_cumsum=g_cumsum, + lse=lse, + do=do, + sink_bias=sink_bias, + scale=scale, + window_size=window_size, + chunk_size=chunk_size, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + ) + + +__all__ = ['AttnTileLangBackend'] diff --git a/build/torch-cuda/ops/attn/backends/tilelang/parallel.py b/build/torch-cuda/ops/attn/backends/tilelang/parallel.py new file mode 100644 index 0000000000000000000000000000000000000000..51849367a2551a0dea4c96dd71b5149270039706 --- /dev/null +++ b/build/torch-cuda/ops/attn/backends/tilelang/parallel.py @@ -0,0 +1,579 @@ +# 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 + +"""TileLang causal attention with GQA, gating, sliding windows, sink bias, and variable-length sequences. + +Forward returns log2 LSE, matching the Triton backward interface. +""" + +import tilelang +import tilelang.language as T +import torch + +from .....ops.utils import prepare_chunk_indices +from .....ops.utils.constant import RCP_LN2 + + +def _pick_tile(K: int, V: int, is_hopper_plus: bool) -> tuple[int, int, int]: + """Return (BT, BS, num_warps) tuned for head dim.""" + if is_hopper_plus: + if max(K, V) <= 64: + return 128, 64, 8 + if max(K, V) <= 128: + return 64, 64, 4 + return 64, 64, 4 + # tune pre-Hopper devices separately + if max(K, V) <= 64: + return 128, 64, 4 + if max(K, V) <= 128: + return 128, 64, 4 + return 64, 32, 4 + + +@tilelang.jit(pass_configs={ + tilelang.PassConfigKey.TL_DISABLE_TMA_LOWER: True, + tilelang.PassConfigKey.TL_ENABLE_FAST_MATH: True, +}) +def _build_parallel_attn_fwd_kernel( + B, + HQ, + H, + K, + V, + BT, + BS, + sm_scale, + log2e_scale, + dtype_str, + USE_G=False, + USE_WINDOW=False, + WINDOW_SIZE=0, + USE_SINK=False, + IS_VARLEN=False, + num_warps=8, + num_stages=1, +): + """Build a forward kernel specialized to the batch size and head dimensions.""" + dtype_map = {'float16': T.float16, 'bfloat16': T.bfloat16, 'float32': T.float32} + dtype = dtype_map[dtype_str] + accum_dtype = T.float32 + threads = num_warps * 32 + # pad shared buffers to satisfy WGMMA/MMA inner-dimension alignment + BK = (K + 15) // 16 * 16 + BV = (V + 15) // 16 * 16 + G = HQ // H + + T_d, NT_d, Ncu_d = T.dynamic("T, NT, Ncu") + + q_s = (B, T_d, HQ, K) + kv_s = (B, T_d, H, K) + vv_s = (B, T_d, H, V) + o_s = (B, T_d, HQ, V) + lse_s = (B, T_d, HQ) + sink_s = (HQ,) + + @T.macro + def kernel_body(q, k, v, g, sink_bias, o, lse, i_b, i_h, i_t, t_s, T_seq, bos): + i_hkv = i_h // G + + s_q = T.alloc_shared((BT, BK), dtype) + s_k = T.alloc_shared((BS, BK), dtype) + s_v = T.alloc_shared((BS, BV), dtype) + + b_s = T.alloc_fragment((BT, BS), accum_dtype) + b_p = T.alloc_fragment((BT, BS), dtype) + b_o = T.alloc_fragment((BT, BV), accum_dtype) + + b_m = T.alloc_fragment((BT,), accum_dtype) + b_m_prev = T.alloc_fragment((BT,), accum_dtype) + b_l = T.alloc_fragment((BT,), accum_dtype) + b_p_sum = T.alloc_fragment((BT,), accum_dtype) + b_scale = T.alloc_fragment((BT,), accum_dtype) + + if USE_G: + s_gq = T.alloc_shared((BT,), accum_dtype) + s_gk = T.alloc_shared((BS,), accum_dtype) + T.copy(g[i_b, t_s:t_s + BT, i_h], s_gq) + + T.copy(q[i_b, t_s:t_s + BT, i_h, :], s_q) + T.fill(b_o, 0) + T.fill(b_l, 0) + T.fill(b_m, -T.infinity(accum_dtype)) + + i_t_local = (t_s - bos) // BT + if USE_WINDOW: + # include every key visible to the first query in this tile + s0_raw = i_t_local * BT - WINDOW_SIZE + 1 + s0 = T.max(s0_raw, 0) + loop_st = (s0 // BS) + else: + loop_st = 0 + loop_ed_raw = T.ceildiv((i_t_local + 1) * BT, BS) + loop_ed = T.min(loop_ed_raw, T.ceildiv(T_seq, BS)) + + for k_idx in T.Pipelined(loop_st, loop_ed, num_stages=num_stages): + k_t_s = bos + k_idx * BS + + T.copy(k[i_b, k_t_s:k_t_s + BS, i_hkv, :], s_k) + T.copy(v[i_b, k_t_s:k_t_s + BS, i_hkv, :], s_v) + if USE_G: + T.copy(g[i_b, k_t_s:k_t_s + BS, i_h], s_gk) + + T.clear(b_s) + T.gemm(s_q, s_k, b_s, transpose_B=True, policy=T.GemmWarpPolicy.FullRow) + + for i, j in T.Parallel(BT, BS): + q_pos = t_s + i + key_pos = k_t_s + j + causal = (key_pos <= q_pos) & (key_pos < T_seq + bos) & (q_pos < T_seq + bos) + if USE_WINDOW: + mask = causal & (q_pos - key_pos < WINDOW_SIZE) + else: + mask = causal + if USE_G: + s = b_s[i, j] * log2e_scale + s_gq[i] - s_gk[j] + else: + s = b_s[i, j] * log2e_scale + b_s[i, j] = T.if_then_else(mask, s, -T.infinity(accum_dtype)) + + T.copy(b_m, b_m_prev) + T.reduce_max(b_s, b_m, dim=1, clear=False) + for i in T.Parallel(BT): + b_m[i] = T.max(b_m[i], b_m_prev[i]) + # use a finite pivot for all-masked rows to avoid (-inf) - (-inf) + b_m_stable = T.alloc_fragment((BT,), accum_dtype) + for i in T.Parallel(BT): + b_m_stable[i] = T.if_then_else(b_m[i] == -T.infinity(accum_dtype), 0., b_m[i]) + for i in T.Parallel(BT): + b_scale[i] = T.if_then_else(b_m_prev[i] == -T.infinity(accum_dtype), 0., T.exp2(b_m_prev[i] - b_m_stable[i])) + for i, j in T.Parallel(BT, BV): + b_o[i, j] *= b_scale[i] + for i, j in T.Parallel(BT, BS): + b_s[i, j] = T.exp2(b_s[i, j] - b_m_stable[i]) + T.reduce_sum(b_s, b_p_sum, dim=1) + for i in T.Parallel(BT): + b_l[i] = b_l[i] * b_scale[i] + b_p_sum[i] + + T.copy(b_s, b_p) + T.gemm(b_p, s_v, b_o, policy=T.GemmWarpPolicy.FullRow) + + # sink-bias contributes to the normalizer but not to the value matmul + if USE_SINK: + # keep sink_bias - b_m finite when the row has no valid keys + for i in T.Parallel(BT): + b_m[i] = T.if_then_else(b_m[i] == -T.infinity(accum_dtype), 0., b_m[i]) + for i in T.Parallel(BT): + b_l[i] += T.exp2(sink_bias[i_h] - b_m[i]) + + s_o = T.alloc_shared((BT, BV), dtype) + s_lse = T.alloc_shared((BT,), accum_dtype) + for i, d in T.Parallel(BT, BV): + b_o[i, d] = b_o[i, d] / b_l[i] + T.copy(b_o, s_o) + for i in T.Parallel(BT): + s_lse[i] = b_m[i] + T.log2(b_l[i]) + for i, d in T.Parallel(BT, V): + if t_s + i < T_seq + bos: + o[i_b, t_s + i, i_h, d] = s_o[i, d] + for i in T.Parallel(BT): + if t_s + i < T_seq + bos: + lse[i_b, t_s + i, i_h] = s_lse[i] + + if IS_VARLEN: + @T.prim_func + def kernel( + q: T.Tensor(q_s, dtype), + k: T.Tensor(kv_s, dtype), + v: T.Tensor(vv_s, dtype), + g: T.Tensor(lse_s, accum_dtype), + sink_bias: T.Tensor(sink_s, accum_dtype), + o: T.Tensor(o_s, dtype), + lse: T.Tensor(lse_s, accum_dtype), + cu_seqlens: T.Tensor((Ncu_d,), T.int32), + chunk_indices: T.Tensor((NT_d, 2), T.int32), + ): + with T.Kernel(HQ, NT_d, 1, threads=threads) as (i_h, i_t, _): + i_n = chunk_indices[i_t, 0] + i_t_local = chunk_indices[i_t, 1] + bos = cu_seqlens[i_n] + eos = cu_seqlens[i_n + 1] + T_seq = eos - bos + t_s = bos + i_t_local * BT + kernel_body(q, k, v, g, sink_bias, o, lse, 0, i_h, i_t, t_s, T_seq, bos) + else: + @T.prim_func + def kernel( + q: T.Tensor(q_s, dtype), + k: T.Tensor(kv_s, dtype), + v: T.Tensor(vv_s, dtype), + g: T.Tensor(lse_s, accum_dtype), + sink_bias: T.Tensor(sink_s, accum_dtype), + o: T.Tensor(o_s, dtype), + lse: T.Tensor(lse_s, accum_dtype), + ): + with T.Kernel(HQ, T.ceildiv(T_d, BT), B, threads=threads) as (i_h, i_t, i_b): + t_s = i_t * BT + kernel_body(q, k, v, g, sink_bias, o, lse, i_b, i_h, i_t, t_s, T_d, 0) + + return kernel + + +def parallel_attn_fwd_tilelang( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + g_cumsum: torch.Tensor | None, + sink_bias: torch.Tensor | None, + scale: float, + window_size: int | None = None, + cu_seqlens: torch.LongTensor | None = None, + chunk_indices: torch.LongTensor | None = None, +) -> tuple[torch.Tensor, torch.Tensor]: + B, T, HQ, K = q.shape + H = k.shape[2] + V = v.shape[-1] + + USE_G = g_cumsum is not None + USE_WINDOW = window_size is not None + USE_SINK = sink_bias is not None + IS_VARLEN = cu_seqlens is not None + + is_hopper_plus = torch.cuda.get_device_capability()[0] >= 9 + BT, BS, num_warps = _pick_tile(K=K, V=V, is_hopper_plus=is_hopper_plus) + + if chunk_indices is None and IS_VARLEN: + chunk_indices = prepare_chunk_indices(cu_seqlens=cu_seqlens, chunk_size=BT) + + sm_scale = scale if scale is not None else K ** -0.5 + log2e_scale = sm_scale * RCP_LN2 + + dtype_str = {torch.float16: 'float16', torch.bfloat16: 'bfloat16', torch.float32: 'float32'}.get(q.dtype) + if dtype_str is None: + raise ValueError(f"Unsupported dtype {q.dtype} for TileLang backend") + + o = torch.empty(B, T, HQ, V, dtype=q.dtype, device=q.device) + lse = torch.empty(B, T, HQ, dtype=torch.float32, device=q.device) + + g_kern = g_cumsum.float() if USE_G else torch.zeros(B, T, HQ, dtype=torch.float32, device=q.device) + sink_kern = sink_bias.float() if USE_SINK else torch.zeros(HQ, dtype=torch.float32, device=q.device) + + kernel = _build_parallel_attn_fwd_kernel( + B=B, + HQ=HQ, + H=H, + K=K, + V=V, + BT=BT, + BS=BS, + sm_scale=sm_scale, + log2e_scale=log2e_scale, + dtype_str=dtype_str, + USE_G=USE_G, + USE_WINDOW=USE_WINDOW, + WINDOW_SIZE=window_size or 0, + USE_SINK=USE_SINK, + IS_VARLEN=IS_VARLEN, + num_warps=num_warps, + num_stages=0 if USE_WINDOW else 1, + ) + + if IS_VARLEN: + kernel(q, k, v, g_kern, sink_kern, o, lse, cu_seqlens.int(), chunk_indices.int()) + else: + kernel(q, k, v, g_kern, sink_kern, o, lse) + + return o, lse + + +@tilelang.jit(pass_configs={ + tilelang.PassConfigKey.TL_DISABLE_TMA_LOWER: True, + tilelang.PassConfigKey.TL_ENABLE_FAST_MATH: True, +}) +def _build_parallel_attn_bwd_kernel( + B, + HQ, + H, + K, + V, + BT, + sm_scale, + log2e_scale, + dtype_str, + USE_G=False, + USE_WINDOW=False, + WINDOW_SIZE=0, + IS_VARLEN=False, + num_warps=4, +): + """Build a backward kernel specialized to the batch size and head dimensions.""" + dtype_map = {'float16': T.float16, 'bfloat16': T.bfloat16, 'float32': T.float32} + dtype = dtype_map[dtype_str] + accum_dtype = T.float32 + threads = num_warps * 32 + # pad shared buffers to satisfy WGMMA/MMA inner-dimension alignment + BK = (K + 15) // 16 * 16 + BV = (V + 15) // 16 * 16 + G = HQ // H + + T_d, NT_d, Ncu_d = T.dynamic("T, NT, Ncu") + + q_s = (B, T_d, HQ, K) + kv_s = (B, T_d, H, K) + vv_s = (B, T_d, H, V) + lse_s = (B, T_d, HQ) + do_s = (B, T_d, HQ, V) + dq_s = (B, T_d, HQ, K) + dkv_s = (B, T_d, H, K) + dvv_s = (B, T_d, H, V) + + @T.macro + def kernel_body(q, k, v, g, lse, delta, do, dq_out, dk_out, dv_out, dg, i_b, i_h, i_t, t_s, T_seq, bos): + i_hkv = i_h // G + + s_k = T.alloc_shared((BT, BK), dtype) + s_v = T.alloc_shared((BT, BV), dtype) + T.copy(k[i_b, t_s:t_s + BT, i_hkv, :], s_k) + T.copy(v[i_b, t_s:t_s + BT, i_hkv, :], s_v) + + if USE_G: + s_gk = T.alloc_shared((BT,), accum_dtype) + T.copy(g[i_b, t_s:t_s + BT, i_h], s_gk) + + b_p_t = T.alloc_fragment((BT, BT), accum_dtype) + b_p_t_cast = T.alloc_fragment((BT, BT), dtype) + b_ds_t = T.alloc_fragment((BT, BT), accum_dtype) + b_ds_t_cast = T.alloc_fragment((BT, BT), dtype) + b_ds = T.alloc_fragment((BT, BT), dtype) + + b_dv = T.alloc_fragment((BT, BV), accum_dtype) + b_dk = T.alloc_fragment((BT, BK), accum_dtype) + T.clear(b_dv) + T.clear(b_dk) + + if USE_G: + b_dgk = T.alloc_fragment((BT,), accum_dtype) + T.clear(b_dgk) + + s_q = T.alloc_shared((BT, BK), dtype) + s_do = T.alloc_shared((BT, BV), dtype) + s_lse = T.alloc_shared((BT,), accum_dtype) + s_delta = T.alloc_shared((BT,), accum_dtype) + s_ds_t = T.alloc_shared((BT, BT), dtype) + b_dq = T.alloc_fragment((BT, BK), accum_dtype) + if USE_G: + s_gq = T.alloc_shared((BT,), accum_dtype) + b_dgq = T.alloc_fragment((BT,), accum_dtype) + b_dgk_tile = T.alloc_fragment((BT,), accum_dtype) + + i_t_local = (t_s - bos) // BT + loop_st = i_t_local + if USE_WINDOW: + # include the last query that can attend to this key tile + loop_ed_swa = i_t_local + 1 + (WINDOW_SIZE + BT - 2) // BT + loop_ed = T.min(T.ceildiv(T_seq, BT), loop_ed_swa) + else: + loop_ed = T.ceildiv(T_seq, BT) + + # pipelining with atomic dq updates corrupts results for non-power-of-two head dimensions + for k_local in T.serial(loop_st, loop_ed): + q_t_s = bos + k_local * BT + + T.copy(q[i_b, q_t_s:q_t_s + BT, i_h, :], s_q) + T.copy(do[i_b, q_t_s:q_t_s + BT, i_h, :], s_do) + T.copy(lse[i_b, q_t_s:q_t_s + BT, i_h], s_lse) + T.copy(delta[i_b, q_t_s:q_t_s + BT, i_h], s_delta) + + if USE_G: + T.copy(g[i_b, q_t_s:q_t_s + BT, i_h], s_gq) + + T.clear(b_p_t) + T.gemm(s_k, s_q, b_p_t, transpose_B=True, policy=T.GemmWarpPolicy.FullRow) + + for i, j in T.Parallel(BT, BT): + if USE_G: + b_p_t[i, j] = T.exp2(b_p_t[i, j] * log2e_scale + s_gq[j] - s_gk[i] - s_lse[j]) + else: + b_p_t[i, j] = T.exp2(b_p_t[i, j] * log2e_scale - s_lse[j]) + + for i, j in T.Parallel(BT, BT): + key_pos = t_s + i + q_pos = q_t_s + j + causal = (key_pos <= q_pos) & (key_pos < T_seq + bos) & (q_pos < T_seq + bos) + if USE_WINDOW: + in_window = (q_pos - key_pos < WINDOW_SIZE) + mask = causal & in_window + else: + mask = causal + b_p_t[i, j] = T.if_then_else(mask, b_p_t[i, j], T.cast(0, accum_dtype)) + + T.copy(b_p_t, b_p_t_cast) + T.gemm(b_p_t_cast, s_do, b_dv, policy=T.GemmWarpPolicy.FullRow) + + T.clear(b_ds_t) + T.gemm(s_v, s_do, b_ds_t, transpose_B=True, policy=T.GemmWarpPolicy.FullRow) + + for i, j in T.Parallel(BT, BT): + b_ds_t[i, j] = b_p_t[i, j] * (b_ds_t[i, j] - s_delta[j]) + for i, j in T.Parallel(BT, BT): + b_ds_t_cast[i, j] = b_ds_t[i, j] * sm_scale + + T.copy(b_ds_t_cast, b_ds) + T.gemm(b_ds, s_q, b_dk, policy=T.GemmWarpPolicy.FullRow) + + T.copy(b_ds_t_cast, s_ds_t) + T.clear(b_dq) + T.gemm(s_ds_t, s_k, b_dq, transpose_A=True, policy=T.GemmWarpPolicy.FullRow) + for i, d in T.Parallel(BT, K): + if q_t_s + i < T_seq + bos: + T.atomic_add(dq_out[i_b, q_t_s + i, i_h, d], b_dq[i, d]) + + if USE_G: + T.reduce_sum(b_ds_t, b_dgq, dim=0) + for j in T.Parallel(BT): + if q_t_s + j < T_seq + bos: + T.atomic_add(dg[i_b, q_t_s + j, i_h], b_dgq[j]) + + T.reduce_sum(b_ds_t, b_dgk_tile, dim=1) + for i in T.Parallel(BT): + b_dgk[i] += b_dgk_tile[i] + + # grouped query heads share KV heads, so their gradients require atomic updates + for i, d in T.Parallel(BT, V): + if t_s + i < T_seq + bos: + T.atomic_add(dv_out[i_b, t_s + i, i_hkv, d], b_dv[i, d]) + for i, d in T.Parallel(BT, K): + if t_s + i < T_seq + bos: + T.atomic_add(dk_out[i_b, t_s + i, i_hkv, d], b_dk[i, d]) + + if USE_G: + for i in T.Parallel(BT): + if t_s + i < T_seq + bos: + T.atomic_add(dg[i_b, t_s + i, i_h], -b_dgk[i]) + + if IS_VARLEN: + @T.prim_func + def kernel( + q: T.Tensor(q_s, dtype), + k: T.Tensor(kv_s, dtype), + v: T.Tensor(vv_s, dtype), + g: T.Tensor(lse_s, accum_dtype), + lse: T.Tensor(lse_s, accum_dtype), + delta: T.Tensor(lse_s, accum_dtype), + do: T.Tensor(do_s, dtype), + dq_out: T.Tensor(dq_s, accum_dtype), + dk_out: T.Tensor(dkv_s, accum_dtype), + dv_out: T.Tensor(dvv_s, accum_dtype), + dg: T.Tensor(lse_s, accum_dtype), + cu_seqlens: T.Tensor((Ncu_d,), T.int32), + chunk_indices: T.Tensor((NT_d, 2), T.int32), + ): + with T.Kernel(HQ, NT_d, 1, threads=threads) as (i_h, i_t, _): + i_n = chunk_indices[i_t, 0] + i_t_local = chunk_indices[i_t, 1] + bos = cu_seqlens[i_n] + T_seq = cu_seqlens[i_n + 1] - bos + t_s = bos + i_t_local * BT + kernel_body(q, k, v, g, lse, delta, do, dq_out, dk_out, dv_out, dg, 0, i_h, i_t, t_s, T_seq, bos) + else: + @T.prim_func + def kernel( + q: T.Tensor(q_s, dtype), + k: T.Tensor(kv_s, dtype), + v: T.Tensor(vv_s, dtype), + g: T.Tensor(lse_s, accum_dtype), + lse: T.Tensor(lse_s, accum_dtype), + delta: T.Tensor(lse_s, accum_dtype), + do: T.Tensor(do_s, dtype), + dq_out: T.Tensor(dq_s, accum_dtype), + dk_out: T.Tensor(dkv_s, accum_dtype), + dv_out: T.Tensor(dvv_s, accum_dtype), + dg: T.Tensor(lse_s, accum_dtype), + ): + with T.Kernel(HQ, T.ceildiv(T_d, BT), B, threads=threads) as (i_h, i_t, i_b): + t_s = i_t * BT + kernel_body(q, k, v, g, lse, delta, do, dq_out, dk_out, dv_out, dg, i_b, i_h, i_t, t_s, T_d, 0) + + return kernel + + +def parallel_attn_bwd_tilelang( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + o: torch.Tensor, + g_cumsum: torch.Tensor | None, + lse: torch.Tensor, + do: torch.Tensor, + sink_bias: torch.Tensor | None = None, + scale: float | None = None, + window_size: int | None = None, + chunk_size: int = 128, + cu_seqlens: torch.LongTensor | None = None, + chunk_indices: torch.LongTensor | None = None, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor | None, torch.Tensor | None]: + B, T, HQ, K = q.shape + H = k.shape[2] + V = v.shape[-1] + BT = 64 + sm_scale = scale if scale is not None else K ** -0.5 + log2e_scale = sm_scale * RCP_LN2 + + USE_G = g_cumsum is not None + USE_WINDOW = window_size is not None + IS_VARLEN = cu_seqlens is not None + + if chunk_indices is None and IS_VARLEN: + chunk_indices = prepare_chunk_indices(cu_seqlens=cu_seqlens, chunk_size=BT) + + # reuse the Triton delta calculation for unpadded tensors + from .....ops.attn.parallel import parallel_attn_bwd_preprocess + delta = parallel_attn_bwd_preprocess(o=o, do=do) + + # zero-initialized fp32 outputs accumulate atomic updates across heads and key tiles + dq = torch.zeros(B, T, HQ, K, dtype=torch.float32, device=q.device) + dk = torch.zeros(B, T, H, K, dtype=torch.float32, device=k.device) + dv = torch.zeros(B, T, H, V, dtype=torch.float32, device=v.device) + dg = torch.zeros(B, T, HQ, dtype=torch.float32, device=q.device) if USE_G else None + + g_kern = g_cumsum.float() if USE_G else torch.zeros(B, T, HQ, dtype=torch.float32, device=q.device) + dg_kern = dg if USE_G else torch.zeros(B, T, HQ, dtype=torch.float32, device=q.device) + + dtype_str = {torch.float16: 'float16', torch.bfloat16: 'bfloat16', torch.float32: 'float32'}.get(q.dtype) + if dtype_str is None: + raise ValueError(f"Unsupported dtype {q.dtype} for TileLang backend") + + kernel = _build_parallel_attn_bwd_kernel( + B=B, + HQ=HQ, + H=H, + K=K, + V=V, + BT=BT, + sm_scale=sm_scale, + log2e_scale=log2e_scale, + dtype_str=dtype_str, + USE_G=USE_G, + USE_WINDOW=USE_WINDOW, + WINDOW_SIZE=window_size or 0, + IS_VARLEN=IS_VARLEN, + ) + + if IS_VARLEN: + kernel(q, k, v, g_kern, lse, delta, do, dq, dk, dv, dg_kern, cu_seqlens.int(), chunk_indices.int()) + else: + kernel(q, k, v, g_kern, lse, delta, do, dq, dk, dv, dg_kern) + + # saved LSE includes the sink mass in the softmax normalizer + dsink_bias = None + if sink_bias is not None: + p_sink = torch.exp2(sink_bias.float()[None, None, :] - lse.float()) + dsink_bias = -(p_sink * delta.float()).sum((0, 1)) + + dq = dq.to(q.dtype) + dk = dk.to(k.dtype) + dv = dv.to(v.dtype) + return dq, dk, dv, dg, dsink_bias diff --git a/build/torch-cuda/ops/attn/decoding.py b/build/torch-cuda/ops/attn/decoding.py index c9a8aea4c294e5f9e5158335abeaa77b0da005cd..9d04fb3b26109e0bc8b997bc9279d0a67aa59bbd 100644 --- a/build/torch-cuda/ops/attn/decoding.py +++ b/build/torch-cuda/ops/attn/decoding.py @@ -11,7 +11,7 @@ import triton.language as tl from ...ops.utils.cumsum import chunk_global_cumsum from ...ops.utils.op import exp -from ...utils import autotune_cache_kwargs, check_shared_mem +from ...utils import autotune_cache_kwargs, check_shared_mem, input_guard @triton.heuristics({ @@ -24,7 +24,7 @@ from ...utils import autotune_cache_kwargs, check_shared_mem for num_warps in [1, 2, 4] + ([] if check_shared_mem('hopper') else [8]) for num_stages in [2, 3, 4, 5] ], - key=['H', 'G', 'K', 'V', 'BK', 'BV', 'USE_G', 'USE_SINK_BIAS'], + key=['H', 'G', 'K', 'V', 'W', 'BK', 'BV', 'USE_G', 'USE_SINK_BIAS'], **autotune_cache_kwargs, ) @triton.jit @@ -44,23 +44,30 @@ def naive_attn_decoding_kernel( G: tl.constexpr, K: tl.constexpr, V: tl.constexpr, + W: tl.constexpr, BS: tl.constexpr, BK: tl.constexpr, BV: tl.constexpr, USE_G: tl.constexpr, USE_SINK_BIAS: tl.constexpr, ): - i_v, i_bh = tl.program_id(0), tl.program_id(1) + pid = tl.program_id(0).to(tl.int64) + NV = tl.cdiv(V, BV) + i_v, i_bh = pid % NV, (pid // NV).to(tl.int64) i_b, i_hq = i_bh // HQ, i_bh % HQ i_h = i_hq // G - bos, eos = tl.load(cu_seqlens + i_b).to(tl.int32), tl.load(cu_seqlens + i_b + 1).to(tl.int32) + bos, eos = tl.load(cu_seqlens + i_b).to(tl.int64), tl.load(cu_seqlens + i_b + 1).to(tl.int64) + if W is not None: + bos = tl.maximum(bos, eos - W) T = eos - bos - p_q = tl.make_block_ptr(q + i_bh * K, (K,), (1, ), (0, ), (BK,), (0,)) - p_o = tl.make_block_ptr(o + i_bh * V, (V,), (1, ), (0, ), (BV,), (0,)) + o_d = tl.arange(0, BK) + o_v = i_v * BV + tl.arange(0, BV) + p_q = q + i_bh * K + o_d + p_o = o + i_bh * V + o_v - b_q = tl.load(p_q, boundary_check=(0,)) + b_q = tl.load(p_q, mask=o_d < K, other=0.0) b_q = (b_q * scale).to(b_q.dtype) b_o = tl.zeros([BV], dtype=tl.float32) @@ -69,8 +76,8 @@ def naive_attn_decoding_kernel( b_acc = tl.zeros([1], dtype=tl.float32) if USE_G: - p_g = tl.make_block_ptr(g_cumsum + bos * HQ + i_hq, (T,), (HQ,), (T-1,), (1,), (0,)) - b_gq = tl.load(p_g, boundary_check=(0,)).to(tl.float32) + p_g = g_cumsum + bos * HQ + i_hq + (T - 1) * HQ + b_gq = tl.load(p_g, mask=T > 0, other=0.0).to(tl.float32) else: b_gq = None @@ -80,21 +87,22 @@ def naive_attn_decoding_kernel( b_sink_bias = None for i_s in range(0, T, BS): - p_k = tl.make_block_ptr(k + (bos * H + i_h) * K, (T, K), (H*K, 1), (i_s, 0), (BS, BK), (1, 0)) - p_v = tl.make_block_ptr(v + (bos * H + i_h) * V, (T, V), (H*V, 1), (i_s, i_v * BV), (BS, BV), (1, 0)) + o_k = (i_s + tl.arange(0, BS)).to(tl.int64) + m_k = o_k < T + p_k = k + (bos * H + i_h) * K + o_k[:, None] * (H*K) + o_d[None, :] + p_v = v + (bos * H + i_h) * V + o_k[:, None] * (H*V) + o_v[None, :] # [BK, BS] - b_k = tl.load(p_k, boundary_check=(0, 1)) + b_k = tl.load(p_k, mask=m_k[:, None] & (o_d[None, :] < K), other=0.0) # [BS, BV] - b_v = tl.load(p_v, boundary_check=(0, 1)) + b_v = tl.load(p_v, mask=m_k[:, None] & (o_v[None, :] < V), other=0.0) # [BT, BS] b_s = tl.sum(b_q[None, :] * b_k, 1) - mask = i_s + tl.arange(0, BS) < T - b_s = tl.where(mask, b_s, float('-inf')) + b_s = tl.where(m_k, b_s, float('-inf')) if USE_G: - p_gk = tl.make_block_ptr(g_cumsum + bos * HQ + i_hq, (T,), (HQ,), (i_s,), (BS,), (0,)) - b_gk = tl.load(p_gk, boundary_check=(0,)).to(tl.float32) + p_gk = g_cumsum + bos * HQ + i_hq + o_k * HQ + b_gk = tl.load(p_gk, mask=m_k, other=0.0).to(tl.float32) b_s += b_gq - b_gk # [BT, BS] b_m, b_mp = tl.maximum(b_m, tl.max(b_s)), b_m @@ -112,10 +120,11 @@ def naive_attn_decoding_kernel( # keep the sink-bias merge finite when masking leaves a row with no valid key. b_m = tl.where(b_m == float('-inf'), 0., b_m) b_acc += exp(b_sink_bias - b_m) - b_o = b_o / b_acc - tl.store(p_o, b_o.to(p_o.dtype.element_ty), boundary_check=(0, )) + b_o = b_o / tl.where(T > 0, b_acc, 1.) + tl.store(p_o, b_o.to(p_o.dtype.element_ty), mask=o_v < V) +@input_guard def attn_decoding_one_step( q: torch.Tensor, k: torch.Tensor, @@ -125,6 +134,7 @@ def attn_decoding_one_step( cu_seqlens: torch.LongTensor = None, do_gate_scale: bool = False, *, + window_size: int | None = None, sink_bias: torch.Tensor | None = None, ): r""" @@ -142,11 +152,13 @@ def attn_decoding_one_step( Scale factor for attention scores. If not provided, it will default to `1 / sqrt(K)`. Default: `None`. cu_seqlens (torch.LongTensor): - Cumulative sequence lengths of shape `[N+1]` used for variable-length training, - consistent with the FlashAttention API. + Cumulative KV sequence lengths of shape `[B+1]`, with one query at the end of each sequence. do_gate_scale (bool): Whether to apply gate scale. Default: `False`. If `True`, the attention scale will also be applied to the gating bias term in Forgetting Transformer or PaTH-FoX. + window_size (int, Optional): + Number of most recent keys attended to in each sequence, including the current token. + Must be nonnegative; zero returns zeros. Default: `None` (all keys). sink_bias (Optional[torch.Tensor]): Per-query-head attention-sink bias logits of shape `[HQ]` — one learnable scalar per query head, as introduced by GPT-OSS. @@ -157,9 +169,13 @@ def attn_decoding_one_step( Outputs of shape `[1, B, HQ, V]`. """ assert cu_seqlens is not None, "The cu_seqlens must be provided for varlen decoding" + if window_size is not None and window_size < 0: + raise ValueError("window_size must be nonnegative") B, T, H, K, V = *k.shape, v.shape[-1] N = len(cu_seqlens) - 1 HQ = q.shape[2] + if H == 0 or HQ % H != 0: + raise ValueError(f"The number of query heads ({HQ}) must be divisible by the number of key/value heads ({H}).") G = HQ // H if scale is None: scale = K ** -0.5 @@ -185,7 +201,7 @@ def attn_decoding_one_step( NV = triton.cdiv(V, BV) o = torch.empty(*q.shape[:-1], V, dtype=v.dtype, device=q.device) - grid = (NV, N * HQ) + grid = (NV * N * HQ,) naive_attn_decoding_kernel[grid]( q=q, k=k, @@ -202,6 +218,7 @@ def attn_decoding_one_step( G=G, K=K, V=V, + W=window_size, BS=BS, BK=BK, BV=BV, diff --git a/build/torch-cuda/ops/attn/naive.py b/build/torch-cuda/ops/attn/naive.py index 5bdebf1e60779bb6cfab34afb40edd91320de1d4..af5c1ef7a7d0f03f1d5a64ccb628652c144386a1 100644 --- a/build/torch-cuda/ops/attn/naive.py +++ b/build/torch-cuda/ops/attn/naive.py @@ -24,10 +24,10 @@ def naive_parallel_attn( Reference PyTorch implementation of parallel attention that returns both output and max_logits. Args: - q: [B, T, HQ, D] - k: [B, T, H, D] - v: [B, T, H, D] - scale: float, optional. If None, defaults to 1 / sqrt(D) + q: [B, T, HQ, K] + k: [B, T, H, K] + v: [B, T, H, V] + scale: float, optional. If None, defaults to 1 / sqrt(K) window_size: int, optional. If provided, each query at position i only attends to keys in [i - window_size + 1, i]. If None, full causal attention is used. causal: bool, default True @@ -40,21 +40,24 @@ def naive_parallel_attn( mass but does not contribute to the output. Returns: - output: [B, T, HQ, D] + output: [B, T, HQ, V] max_logits: [B, T, HQ] """ - B, T, HQ, D = q.shape + B, T, HQ, K = q.shape + V = v.shape[-1] H = k.shape[2] + if H == 0 or HQ % H != 0: + raise ValueError(f"The number of query heads ({HQ}) must be divisible by the number of key/value heads ({H}).") G = HQ // H if scale is None: - scale = D ** -0.5 + scale = K ** -0.5 - # reshape q to separate groups: [B, T, HQ, D] -> [B, T, H, G, D] - q = q.reshape(B, T, H, G, D) + # reshape q to separate groups: [B, T, HQ, K] -> [B, T, H, G, K] + q = q.reshape(B, T, H, G, K) # compute attention scores via einsum: [B, H, G, T, T] - # k is [B, T, H, D] — no group dim, so each group shares the same k + # k is [B, T, H, K] — no group dim, so each group shares the same k scores = torch.einsum('bqhgd,bkhd->bhgqk', q, k) * scale # apply causal mask @@ -79,15 +82,15 @@ def naive_parallel_attn( max_logits = torch.maximum(max_logits, sink_bias_logits) if sink_bias is None: - # compute output via einsum: [B, H, G, T, T] x [B, T, H, D] -> [B, T, H, G, D] + # compute output via einsum: [B, H, G, T, T] x [B, T, H, V] -> [B, T, H, G, V] attn_weights = F.softmax(scores, dim=-1) - output = torch.einsum('bhgqk,bkhd->bqhgd', attn_weights, v).reshape(B, T, HQ, D) + output = torch.einsum('bhgqk,bkhd->bqhgd', attn_weights, v).reshape(B, T, HQ, V) else: probs_unnorm = torch.exp(scores - max_logits[..., None]) sink_bias_unnorm = torch.exp(sink_bias_logits - max_logits) denom = probs_unnorm.sum(dim=-1) + sink_bias_unnorm output = torch.einsum('bhgqk,bkhd->bqhgd', probs_unnorm, v) - output = (output / denom.permute(0, 3, 1, 2)[..., None]).reshape(B, T, HQ, D) + output = (output / denom.permute(0, 3, 1, 2)[..., None]).reshape(B, T, HQ, V) return output, max_logits.permute(0, 3, 1, 2).reshape(B, T, HQ) @@ -101,34 +104,41 @@ def naive_attn_decoding( cu_seqlens: torch.LongTensor | None = None, do_gate_scale: bool = False, *, + window_size: int | None = None, sink_bias: torch.Tensor | None = None, ): """ Reference PyTorch implementation of packed-varlen decoding attention, mirroring `attn_decoding_one_step`. A single query per sequence attends - to all of its KV (no causal mask needed — query is at the last position). + to its most recent KV within the window (query is at the last position). Args: - q: [1, B, HQ, D] — one query token per sequence - k: [1, T_total, H, D] - v: [1, T_total, H, D] + q: [1, B, HQ, K] — one query token per sequence + k: [1, T_total, H, K] + v: [1, T_total, H, V] g: [1, T_total, HQ], optional log decay factors - scale: float, defaults to 1/sqrt(D) + scale: float, defaults to 1/sqrt(K) cu_seqlens: [B+1] do_gate_scale: bool, if True scales `g` by `scale` before use (matches PaTH / Forgetting Transformer convention). + window_size: nonnegative int, optional. Attend to the last window_size keys in each sequence. + None attends to all keys; zero returns zeros. sink_bias: [HQ], optional GPT-OSS-style sink bias logits Returns: o: [1, B, HQ, V] """ assert cu_seqlens is not None, "cu_seqlens must be provided for varlen decoding" - HQ, D = q.shape[-2], q.shape[-1] + if window_size is not None and window_size < 0: + raise ValueError("window_size must be nonnegative") + HQ, K = q.shape[-2], q.shape[-1] V = v.shape[-1] H = k.shape[2] + if H == 0 or HQ % H != 0: + raise ValueError(f"The number of query heads ({HQ}) must be divisible by the number of key/value heads ({H}).") G = HQ // H if scale is None: - scale = D ** -0.5 + scale = K ** -0.5 if sink_bias is not None: assert sink_bias.shape == (HQ,), "sink_bias must have shape [HQ]" @@ -136,19 +146,21 @@ def naive_attn_decoding( for i in range(len(cu_seqlens) - 1): bos, eos = int(cu_seqlens[i]), int(cu_seqlens[i + 1]) T_i = eos - bos - qi = q[:, i:i + 1] # [1, 1, HQ, D] + qi = q[:, i:i + 1] # [1, 1, HQ, K] - if T_i == 0: - # no KV for this sequence → output zeros (sink_bias has no value to contribute) + if T_i == 0 or window_size == 0: + # no visible KV → output zeros (sink_bias has no value to contribute) outputs.append(torch.zeros((1, 1, HQ, V), dtype=q.dtype, device=q.device)) continue - ki = k[:, bos:eos] # [1, T_i, H, D] - vi = v[:, bos:eos] # [1, T_i, H, D] + ki = k[:, bos:eos] # [1, T_i, H, K] + vi = v[:, bos:eos] # [1, T_i, H, V] # scores: [1, H, G, 1, T_i] - qi_g = qi.reshape(1, 1, H, G, D) + qi_g = qi.reshape(1, 1, H, G, K) scores = torch.einsum('bqhgd,bkhd->bhgqk', qi_g, ki) * scale + if window_size is not None: + scores = scores.masked_fill(torch.arange(T_i, device=q.device) < T_i - window_size, float('-inf')) if g is not None: gi = g[:, bos:eos].float() diff --git a/build/torch-cuda/ops/attn/parallel.py b/build/torch-cuda/ops/attn/parallel.py index 3adef692899e644d29db235b9c069b50ad79abb4..0f46058fee6ef1d248ae518b7fde0e80e3be9adb 100644 --- a/build/torch-cuda/ops/attn/parallel.py +++ b/build/torch-cuda/ops/attn/parallel.py @@ -10,11 +10,11 @@ import triton import triton.language as tl from einops import reduce -from ...ops.backends import dispatch +from ...backends import dispatch from ...ops.utils import prepare_chunk_indices from ...ops.utils.constant import RCP_LN2 from ...ops.utils.cumsum import chunk_global_cumsum -from ...ops.utils.op import exp2, log2 +from ...ops.utils.op import exp2, log2, unflatten_program_id from ...utils import autocast_custom_bwd, autocast_custom_fwd, check_shared_mem, contiguous @@ -53,12 +53,13 @@ def parallel_attn_fwd_kernel( USE_WINDOW: tl.constexpr, IS_VARLEN: tl.constexpr, ): - i_v, i_t, i_bh = tl.program_id(0), tl.program_id(1), tl.program_id(2) + i_v, i_t = unflatten_program_id(tl.cdiv(V, BV)) + i_bh = tl.program_id(1).to(tl.int64) i_b, i_hq = i_bh // HQ, i_bh % HQ i_h = i_hq // G 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) + i_n, i_t = tl.load(chunk_indices + i_t * 2).to(tl.int32), tl.load(chunk_indices + i_t * 2 + 1).to(tl.int64) bos, eos = tl.load(cu_seqlens + i_n).to(tl.int64), tl.load(cu_seqlens + i_n + 1).to(tl.int64) T = (eos - bos).to(tl.int32) else: @@ -66,13 +67,18 @@ def parallel_attn_fwd_kernel( bos, eos = (i_n * T).to(tl.int64), (i_n * T + T).to(tl.int64) RCP_LN2: tl.constexpr = 1.4426950216 - p_q = tl.make_block_ptr(q + (bos * HQ + i_hq) * K, (T, K), (HQ*K, 1), (i_t * BT, 0), (BT, BK), (1, 0)) - p_o = tl.make_block_ptr(o + (bos * HQ + i_hq) * V, (T, V), (HQ*V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) - p_lse = tl.make_block_ptr(lse + bos * HQ + i_hq, (T,), (HQ,), (i_t * BT,), (BT,), (0,)) + # [BT] + o_q = i_t * BT + tl.arange(0, BT) + o_d = tl.arange(0, BK) + o_v = i_v * BV + tl.arange(0, BV) + m_q = o_q < T + p_q = q + (bos * HQ + i_hq) * K + o_q[:, None] * (HQ*K) + o_d[None, :] + p_o = o + (bos * HQ + i_hq) * V + o_q[:, None] * (HQ*V) + o_v[None, :] + p_lse = lse + bos * HQ + i_hq + o_q * HQ # the Q block is kept in the shared memory throughout the whole kernel # [BT, BK] - b_q = tl.load(p_q, boundary_check=(0, 1)) + b_q = tl.load(p_q, mask=m_q[:, None] & (o_d[None, :] < K), other=0.0) # [BT, BV] b_o = tl.zeros([BT, BV], dtype=tl.float32) @@ -80,8 +86,8 @@ def parallel_attn_fwd_kernel( b_acc = tl.zeros([BT], dtype=tl.float32) if USE_G: - p_g = tl.make_block_ptr(g_cumsum + bos * HQ + i_hq, (T,), (HQ,), (i_t * BT,), (BT,), (0,)) - b_gq = tl.load(p_g, boundary_check=(0,)).to(tl.float32) + p_g = g_cumsum + bos * HQ + i_hq + o_q * HQ + b_gq = tl.load(p_g, mask=m_q, other=0.0).to(tl.float32) else: b_gq = None @@ -90,25 +96,22 @@ def parallel_attn_fwd_kernel( else: b_sink_bias = None - # [BT] - o_q = i_t * BT + tl.arange(0, BT) - # for sliding window, skip key blocks that are entirely outside the window. # the earliest key position any query in this block needs: max(0, i_t*BT - W + 1) i_start = tl.maximum((i_t * BT - W + 1) // BS * BS, 0) if USE_WINDOW else 0 for i_s in range(i_start, i_t * BT, BS): - p_k = tl.make_block_ptr(k + (bos * H + i_h) * K, (K, T), (1, H*K), (0, i_s), (BK, BS), (0, 1)) - p_v = tl.make_block_ptr(v + (bos * H + i_h) * V, (T, V), (H*V, 1), (i_s, i_v * BV), (BS, BV), (1, 0)) + o_k = i_s + tl.arange(0, BS) + m_k = o_k < T + p_k = k + (bos * H + i_h) * K + o_d[:, None] + o_k[None, :] * (H*K) + p_v = v + (bos * H + i_h) * V + o_k[:, None] * (H*V) + o_v[None, :] # [BK, BS] - b_k = tl.load(p_k, boundary_check=(0, 1)) + b_k = tl.load(p_k, mask=(o_d[:, None] < K) & m_k[None, :], other=0.0) # [BS, BV] - b_v = tl.load(p_v, boundary_check=(0, 1)) + b_v = tl.load(p_v, mask=m_k[:, None] & (o_v[None, :] < V), other=0.0) # [BT, BS] b_s = tl.dot(b_q, b_k) * scale * RCP_LN2 - o_k = i_s + tl.arange(0, BS) - m_k = o_k < T if USE_G: b_gk = tl.load(g_cumsum + (bos + o_k) * HQ + i_hq, mask=m_k, other=0).to(tl.float32) b_s += b_gq[:, None] - b_gk[None, :] @@ -131,16 +134,16 @@ def parallel_attn_fwd_kernel( b_mp = b_m for i_s in range(i_t * BT, min((i_t + 1) * BT, T), BS): - p_k = tl.make_block_ptr(k + (bos * H + i_h) * K, (K, T), (1, H*K), (0, i_s), (BK, BS), (0, 1)) - p_v = tl.make_block_ptr(v + (bos * H + i_h) * V, (T, V), (H*V, 1), (i_s, i_v * BV), (BS, BV), (1, 0)) - # [BS] o_k = i_s + tl.arange(0, BS) m_k = o_k < T + p_k = k + (bos * H + i_h) * K + o_d[:, None] + o_k[None, :] * (H*K) + p_v = v + (bos * H + i_h) * V + o_k[:, None] * (H*V) + o_v[None, :] + # [BK, BS] - b_k = tl.load(p_k, boundary_check=(0, 1)) + b_k = tl.load(p_k, mask=(o_d[:, None] < K) & m_k[None, :], other=0.0) # [BS, BV] - b_v = tl.load(p_v, boundary_check=(0, 1)) + b_v = tl.load(p_v, mask=m_k[:, None] & (o_v[None, :] < V), other=0.0) # [BT, BS] b_s = tl.dot(b_q, b_k) * scale * RCP_LN2 @@ -176,8 +179,9 @@ def parallel_attn_fwd_kernel( b_o = b_o / b_acc[:, None] b_m += log2(b_acc) - tl.store(p_o, b_o.to(p_o.dtype.element_ty), boundary_check=(0, 1)) - tl.store(p_lse, b_m.to(p_lse.dtype.element_ty), boundary_check=(0,)) + tl.store(p_o, b_o.to(p_o.dtype.element_ty), mask=m_q[:, None] & (o_v[None, :] < V)) + if i_v == 0: + tl.store(p_lse, b_m.to(p_lse.dtype.element_ty), mask=m_q) @triton.jit @@ -234,12 +238,13 @@ def parallel_attn_bwd_kernel_dq( USE_WINDOW: tl.constexpr, IS_VARLEN: tl.constexpr, ): - i_v, i_t, i_bh = tl.program_id(0), tl.program_id(1), tl.program_id(2) + i_v, i_t = unflatten_program_id(tl.cdiv(V, BV)) + i_bh = tl.program_id(1).to(tl.int64) i_b, i_hq = i_bh // HQ, i_bh % HQ i_h = i_hq // G 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) + i_n, i_t = tl.load(chunk_indices + i_t * 2).to(tl.int32), tl.load(chunk_indices + i_t * 2 + 1).to(tl.int64) bos, eos = tl.load(cu_seqlens + i_n).to(tl.int64), tl.load(cu_seqlens + i_n + 1).to(tl.int64) T = (eos - bos).to(tl.int32) else: @@ -248,44 +253,46 @@ def parallel_attn_bwd_kernel_dq( # NOTE: we must multiply RCP_LN2 after tl.dot for high precision RCP_LN2: tl.constexpr = 1.4426950216 - p_q = tl.make_block_ptr(q + (bos * HQ + i_hq) * K, (T, K), (HQ*K, 1), (i_t * BT, 0), (BT, BK), (1, 0)) - p_dq = tl.make_block_ptr(dq + (bos * HQ + i_hq) * K, (T, K), (HQ*K, 1), (i_t * BT, 0), (BT, BK), (1, 0)) - p_do = tl.make_block_ptr(do + (bos * HQ + i_hq) * V, (T, V), (HQ*V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) - p_lse = tl.make_block_ptr(lse + bos * HQ + i_hq, (T,), (HQ,), (i_t * BT,), (BT,), (0,)) - p_delta = tl.make_block_ptr(delta + bos * HQ + i_hq, (T,), (HQ,), (i_t * BT,), (BT,), (0,)) + o_q = i_t * BT + tl.arange(0, BT) + o_d = tl.arange(0, BK) + o_v = i_v * BV + tl.arange(0, BV) + m_q = o_q < T + p_q = q + (bos * HQ + i_hq) * K + o_q[:, None] * (HQ*K) + o_d[None, :] + p_dq = dq + (bos * HQ + i_hq) * K + o_q[:, None] * (HQ*K) + o_d[None, :] + p_do = do + (bos * HQ + i_hq) * V + o_q[:, None] * (HQ*V) + o_v[None, :] + p_lse = lse + bos * HQ + i_hq + o_q * HQ + p_delta = delta + bos * HQ + i_hq + o_q * HQ # [BT, BK] - b_q = tl.load(p_q, boundary_check=(0, 1)) + b_q = tl.load(p_q, mask=m_q[:, None] & (o_d[None, :] < K), other=0.0) # [BT, BV] - b_do = tl.load(p_do, boundary_check=(0, 1)) + b_do = tl.load(p_do, mask=m_q[:, None] & (o_v[None, :] < V), other=0.0) # [BT] - b_lse = tl.load(p_lse, boundary_check=(0,)) - b_delta = tl.load(p_delta, boundary_check=(0,)) + b_lse = tl.load(p_lse, mask=m_q, other=0.0) + b_delta = tl.load(p_delta, mask=m_q, other=0.0) # [BT, BK] b_dq = tl.zeros([BT, BK], dtype=tl.float32) if USE_G: b_dg = tl.zeros([BT], dtype=tl.float32) - p_gq = tl.make_block_ptr(g_cumsum + bos * HQ + i_hq, (T,), (HQ,), (i_t * BT,), (BT,), (0,)) - b_gq = tl.load(p_gq, boundary_check=(0,)).to(tl.float32) + p_gq = g_cumsum + bos * HQ + i_hq + o_q * HQ + b_gq = tl.load(p_gq, mask=m_q, other=0.0).to(tl.float32) else: b_gq = None b_dg = None - o_q = i_t * BT + tl.arange(0, BT) - i_start = tl.maximum((i_t * BT - W + 1) // BS * BS, 0) if USE_WINDOW else 0 for i_s in range(i_start, i_t * BT, BS): - p_k = tl.make_block_ptr(k + (bos * H + i_h) * K, (K, T), (1, H*K), (0, i_s), (BK, BS), (0, 1)) - p_v = tl.make_block_ptr(v + (bos * H + i_h) * V, (V, T), (1, H*V), (i_v * BV, i_s), (BV, BS), (0, 1)) - o_k = i_s + tl.arange(0, BS) m_k = o_k < T + p_k = k + (bos * H + i_h) * K + o_d[:, None] + o_k[None, :] * (H*K) + p_v = v + (bos * H + i_h) * V + o_v[:, None] + o_k[None, :] * (H*V) + # [BK, BS] - b_k = tl.load(p_k, boundary_check=(0, 1)) + b_k = tl.load(p_k, mask=(o_d[:, None] < K) & m_k[None, :], other=0.0) # [BV, BS] - b_v = tl.load(p_v, boundary_check=(0, 1)) + b_v = tl.load(p_v, mask=(o_v[:, None] < V) & m_k[None, :], other=0.0) # [BT, BS] b_s = tl.dot(b_q, b_k) * scale * RCP_LN2 if USE_G: @@ -299,27 +306,27 @@ def parallel_attn_bwd_kernel_dq( b_dp = tl.dot(b_do, b_v) b_ds = b_p * (b_dp.to(tl.float32) - b_delta[:, None]) # [BT, BS] @ [BS, BK] -> [BT, BK] - b_dq += tl.dot(b_ds.to(b_k.dtype), tl.trans(b_k)) + b_dq = tl.dot(b_ds.to(b_k.dtype), tl.trans(b_k), b_dq) if USE_G: b_dg += tl.sum(b_ds, 1) for i_s in range(i_t * BT, min((i_t + 1) * BT, T), BS): - p_k = tl.make_block_ptr(k + (bos * H + i_h) * K, (K, T), (1, H*K), (0, i_s), (BK, BS), (0, 1)) - p_v = tl.make_block_ptr(v + (bos * H + i_h) * V, (V, T), (1, H*V), (i_v * BV, i_s), (BV, BS), (0, 1)) - # [BS] o_k = i_s + tl.arange(0, BS) m_k = o_k < T + p_k = k + (bos * H + i_h) * K + o_d[:, None] + o_k[None, :] * (H*K) + p_v = v + (bos * H + i_h) * V + o_v[:, None] + o_k[None, :] * (H*V) + # [BK, BS] - b_k = tl.load(p_k, boundary_check=(0, 1)) + b_k = tl.load(p_k, mask=(o_d[:, None] < K) & m_k[None, :], other=0.0) # [BV, BS] - b_v = tl.load(p_v, boundary_check=(0, 1)) + b_v = tl.load(p_v, mask=(o_v[:, None] < V) & m_k[None, :], other=0.0) # [BT, BS] b_s = tl.dot(b_q, b_k) * scale * RCP_LN2 if USE_G: - p_gk = tl.make_block_ptr(g_cumsum + bos * HQ + i_hq, (T,), (HQ,), (i_s,), (BS,), (0,)) - b_gk = tl.load(p_gk, boundary_check=(0,)).to(tl.float32) + p_gk = g_cumsum + bos * HQ + i_hq + o_k * HQ + b_gk = tl.load(p_gk, mask=m_k, other=0.0).to(tl.float32) b_s += b_gq[:, None] - b_gk[None, :] if USE_WINDOW: b_p = tl.where( @@ -333,15 +340,15 @@ def parallel_attn_bwd_kernel_dq( b_dp = tl.dot(b_do, b_v) b_ds = b_p * (b_dp.to(tl.float32) - b_delta[:, None]) # [BT, BS] @ [BS, BK] -> [BT, BK] - b_dq += tl.dot(b_ds.to(b_k.dtype), tl.trans(b_k)) + b_dq = tl.dot(b_ds.to(b_k.dtype), tl.trans(b_k), b_dq) if USE_G: b_dg += tl.sum(b_ds, 1) b_dq *= scale - tl.store(p_dq, b_dq.to(p_dq.dtype.element_ty), boundary_check=(0, 1)) + tl.store(p_dq, b_dq.to(p_dq.dtype.element_ty), mask=m_q[:, None] & (o_d[None, :] < K)) if USE_G: - p_dg = tl.make_block_ptr(dg_cumsum + bos * HQ + i_hq, (T,), (HQ,), (i_t * BT,), (BT,), (0,)) - tl.store(p_dg, b_dg.to(p_dg.dtype.element_ty), boundary_check=(0,)) + p_dg = dg_cumsum + bos * HQ + i_hq + o_q * HQ + tl.store(p_dg, b_dg.to(p_dg.dtype.element_ty), mask=m_q) @triton.heuristics({ @@ -380,12 +387,13 @@ def parallel_attn_bwd_kernel_dkv( USE_WINDOW: tl.constexpr, IS_VARLEN: tl.constexpr, ): - i_v, i_t, i_bh = tl.program_id(0), tl.program_id(1), tl.program_id(2) + i_v, i_t = unflatten_program_id(tl.cdiv(V, BV)) + i_bh = tl.program_id(1).to(tl.int64) i_b, i_hq = i_bh // HQ, i_bh % HQ i_h = i_hq // G 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) + i_n, i_t = tl.load(chunk_indices + i_t * 2).to(tl.int32), tl.load(chunk_indices + i_t * 2 + 1).to(tl.int64) bos, eos = tl.load(cu_seqlens + i_n).to(tl.int64), tl.load(cu_seqlens + i_n + 1).to(tl.int64) T = (eos - bos).to(tl.int32) else: @@ -393,49 +401,51 @@ def parallel_attn_bwd_kernel_dkv( bos, eos = (i_n * T).to(tl.int64), (i_n * T + T).to(tl.int64) RCP_LN2: tl.constexpr = 1.4426950216 - p_k = tl.make_block_ptr(k + (bos * H + i_h) * K, (T, K), (H*K, 1), (i_t * BT, 0), (BT, BK), (1, 0)) - p_v = tl.make_block_ptr(v + (bos * H + i_h) * V, (T, V), (H*V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) - p_dk = tl.make_block_ptr(dk + (bos * HQ + i_hq) * K, (T, K), (HQ*K, 1), (i_t * BT, 0), (BT, BK), (1, 0)) - p_dv = tl.make_block_ptr(dv + (bos * HQ + i_hq) * V, (T, V), (HQ*V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) + o_k = i_t * BT + tl.arange(0, BT) + o_d = tl.arange(0, BK) + o_v = i_v * BV + tl.arange(0, BV) + m_k = o_k < T + p_k = k + (bos * H + i_h) * K + o_k[:, None] * (H*K) + o_d[None, :] + p_v = v + (bos * H + i_h) * V + o_k[:, None] * (H*V) + o_v[None, :] + p_dk = dk + (bos * HQ + i_hq) * K + o_k[:, None] * (HQ*K) + o_d[None, :] + p_dv = dv + (bos * HQ + i_hq) * V + o_k[:, None] * (HQ*V) + o_v[None, :] # [BT, BK] - b_k = tl.load(p_k, boundary_check=(0, 1)) + b_k = tl.load(p_k, mask=m_k[:, None] & (o_d[None, :] < K), other=0.0) b_dk = tl.zeros([BT, BK], dtype=tl.float32) # [BT, BV] - b_v = tl.load(p_v, boundary_check=(0, 1)) + b_v = tl.load(p_v, mask=m_k[:, None] & (o_v[None, :] < V), other=0.0) b_dv = tl.zeros([BT, BV], dtype=tl.float32) - o_k = i_t * BT + tl.arange(0, BT) - if USE_G: - p_gk = tl.make_block_ptr(g_cumsum + bos * HQ + i_hq, (T,), (HQ,), (i_t * BT,), (BT,), (0,)) - b_gk = tl.load(p_gk, boundary_check=(0,)).to(tl.float32) + p_gk = g_cumsum + bos * HQ + i_hq + o_k * HQ + b_gk = tl.load(p_gk, mask=m_k, other=0.0).to(tl.float32) b_dg = tl.zeros([BT], dtype=tl.float32) else: b_gk = None b_dg = None for i_s in range(i_t * BT, min((i_t + 1) * BT, T), BS): - p_q = tl.make_block_ptr(q + (bos * HQ + i_hq) * K, (T, K), (HQ*K, 1), (i_s, 0), (BS, BK), (1, 0)) - p_do = tl.make_block_ptr(do + (bos * HQ + i_hq) * V, (T, V), (HQ*V, 1), (i_s, i_v * BV), (BS, BV), (1, 0)) - p_lse = tl.make_block_ptr(lse + bos * HQ + i_hq, (T,), (HQ,), (i_s,), (BS,), (0,)) - p_delta = tl.make_block_ptr(delta + bos * HQ + i_hq, (T,), (HQ,), (i_s,), (BS,), (0,)) - # [BS] o_q = i_s + tl.arange(0, BS) m_q = o_q < T + p_q = q + (bos * HQ + i_hq) * K + o_q[:, None] * (HQ*K) + o_d[None, :] + p_do = do + (bos * HQ + i_hq) * V + o_q[:, None] * (HQ*V) + o_v[None, :] + p_lse = lse + bos * HQ + i_hq + o_q * HQ + p_delta = delta + bos * HQ + i_hq + o_q * HQ + # [BS, BK] - b_q = tl.load(p_q, boundary_check=(0, 1)) + b_q = tl.load(p_q, mask=m_q[:, None] & (o_d[None, :] < K), other=0.0) # [BS, BV] - b_do = tl.load(p_do, boundary_check=(0, 1)) + b_do = tl.load(p_do, mask=m_q[:, None] & (o_v[None, :] < V), other=0.0) # [BS] - b_lse = tl.load(p_lse, boundary_check=(0,)) - b_delta = tl.load(p_delta, boundary_check=(0,)) + b_lse = tl.load(p_lse, mask=m_q, other=0.0) + b_delta = tl.load(p_delta, mask=m_q, other=0.0) # [BT, BS] b_s = tl.dot(b_k, tl.trans(b_q)) * scale * RCP_LN2 if USE_G: - p_gq = tl.make_block_ptr(g_cumsum + bos * HQ + i_hq, (T,), (HQ,), (i_s,), (BS,), (0,)) - b_gq = tl.load(p_gq, boundary_check=(0,)).to(tl.float32) + p_gq = g_cumsum + bos * HQ + i_hq + o_q * HQ + b_gq = tl.load(p_gq, mask=m_q, other=0.0).to(tl.float32) b_s += b_gq[None, :] - b_gk[:, None] if USE_WINDOW: b_p = tl.where( @@ -445,13 +455,13 @@ def parallel_attn_bwd_kernel_dkv( else: b_p = tl.where((o_k[:, None] <= o_q[None, :]) & m_q[None, :], exp2(b_s - b_lse[None, :]), 0) # [BT, BS] @ [BS, BV] -> [BT, BV] - b_dv += tl.dot(b_p.to(b_do.dtype), b_do) + b_dv = tl.dot(b_p.to(b_do.dtype), b_do, b_dv) # [BT, BV] @ [BV, BS] -> [BT, BS] b_dp = tl.dot(b_v, tl.trans(b_do)) # [BT, BS] b_ds = b_p * (b_dp - b_delta[None, :]) # [BT, BS] @ [BS, BK] -> [BT, BK] - b_dk += tl.dot(b_ds.to(b_q.dtype), b_q) + b_dk = tl.dot(b_ds.to(b_q.dtype), b_q, b_dk) if USE_G: b_dg -= tl.sum(b_ds, 1) @@ -461,48 +471,48 @@ def parallel_attn_bwd_kernel_dkv( i_end = min(tl.cdiv(T, BS) * BS, (i_t + 1) * BT + W - 1) if USE_WINDOW else tl.cdiv(T, BS) * BS for i_s in range((i_t + 1) * BT, i_end, BS): - p_q = tl.make_block_ptr(q + (bos * HQ + i_hq) * K, (T, K), (HQ*K, 1), (i_s, 0), (BS, BK), (1, 0)) - p_do = tl.make_block_ptr(do + (bos * HQ + i_hq) * V, (T, V), (HQ*V, 1), (i_s, i_v * BV), (BS, BV), (1, 0)) - p_lse = tl.make_block_ptr(lse + bos * HQ + i_hq, (T,), (HQ,), (i_s,), (BS,), (0,)) - p_delta = tl.make_block_ptr(delta + bos * HQ + i_hq, (T,), (HQ,), (i_s,), (BS,), (0,)) - # [BS] o_q = i_s + tl.arange(0, BS) m_q = o_q < T + p_q = q + (bos * HQ + i_hq) * K + o_q[:, None] * (HQ*K) + o_d[None, :] + p_do = do + (bos * HQ + i_hq) * V + o_q[:, None] * (HQ*V) + o_v[None, :] + p_lse = lse + bos * HQ + i_hq + o_q * HQ + p_delta = delta + bos * HQ + i_hq + o_q * HQ + # [BS, BK] - b_q = tl.load(p_q, boundary_check=(0, 1)) + b_q = tl.load(p_q, mask=m_q[:, None] & (o_d[None, :] < K), other=0.0) # [BS, BV] - b_do = tl.load(p_do, boundary_check=(0, 1)) + b_do = tl.load(p_do, mask=m_q[:, None] & (o_v[None, :] < V), other=0.0) # [BS] - b_lse = tl.load(p_lse, boundary_check=(0,)) - b_delta = tl.load(p_delta, boundary_check=(0,)) + b_lse = tl.load(p_lse, mask=m_q, other=0.0) + b_delta = tl.load(p_delta, mask=m_q, other=0.0) # [BT, BS] b_s = tl.dot(b_k, tl.trans(b_q)) * scale * RCP_LN2 if USE_G: - p_gq = tl.make_block_ptr(g_cumsum + bos * HQ + i_hq, (T,), (HQ,), (i_s,), (BS,), (0,)) - b_gq = tl.load(p_gq, boundary_check=(0,)).to(tl.float32) + p_gq = g_cumsum + bos * HQ + i_hq + o_q * HQ + b_gq = tl.load(p_gq, mask=m_q, other=0.0).to(tl.float32) b_s += b_gq[None, :] - b_gk[:, None] if USE_WINDOW: b_p = tl.where((o_q[None, :] - o_k[:, None] < W) & m_q[None, :], exp2(b_s - b_lse[None, :]), 0) else: b_p = tl.where(m_q[None, :], exp2(b_s - b_lse[None, :]), 0) # [BT, BS] @ [BS, BV] -> [BT, BV] - b_dv += tl.dot(b_p.to(b_do.dtype), b_do) + b_dv = tl.dot(b_p.to(b_do.dtype), b_do, b_dv) # [BT, BV] @ [BV, BS] -> [BT, BS] b_dp = tl.dot(b_v, tl.trans(b_do)) # [BT, BS] b_ds = b_p * (b_dp - b_delta[None, :]) # [BT, BS] @ [BS, BK] -> [BT, BK] - b_dk += tl.dot(b_ds.to(b_q.dtype), b_q) + b_dk = tl.dot(b_ds.to(b_q.dtype), b_q, b_dk) if USE_G: b_dg -= tl.sum(b_ds, 1) b_dk = b_dk * scale - tl.store(p_dk, b_dk.to(p_dk.dtype.element_ty), boundary_check=(0, 1)) - tl.store(p_dv, b_dv.to(p_dv.dtype.element_ty), boundary_check=(0, 1)) + tl.store(p_dk, b_dk.to(p_dk.dtype.element_ty), mask=m_k[:, None] & (o_d[None, :] < K)) + tl.store(p_dv, b_dv.to(p_dv.dtype.element_ty), mask=m_k[:, None] & (o_v[None, :] < V)) if USE_G: - p_dg = tl.make_block_ptr(dg_cumsum + bos * HQ + i_hq, (T,), (HQ,), (i_t * BT,), (BT,), (0,)) - tl.store(p_dg, b_dg.to(p_dg.dtype.element_ty), boundary_check=(0,)) + p_dg = dg_cumsum + bos * HQ + i_hq + o_k * HQ + tl.store(p_dg, b_dg.to(p_dg.dtype.element_ty), mask=m_k) @dispatch('attn') @@ -546,7 +556,7 @@ def parallel_attn_fwd( o = torch.empty(B, T, HQ, V, dtype=v.dtype, device=q.device) lse = torch.empty(B, T, HQ, dtype=torch.float, device=q.device) - grid = (NV, NT, B * HQ) + grid = (NV * NT, B * HQ) parallel_attn_fwd_kernel[grid]( q=q, k=k, @@ -641,7 +651,7 @@ def parallel_attn_bwd( dq = torch.empty(B, T, HQ, K, dtype=k.dtype if H == HQ else torch.float, device=q.device) dk = torch.empty(B, T, HQ, K, dtype=k.dtype if H == HQ else torch.float, device=q.device) dv = torch.empty(B, T, HQ, V, dtype=v.dtype if H == HQ else torch.float, device=q.device) - grid = (NV, NT, B * HQ) + grid = (NV * NT, B * HQ) dg_cumsum, dg_cumsum_k = None, None if g_cumsum is not None: @@ -716,7 +726,6 @@ def parallel_attn_bwd( return dq, dk, dv, dg_cumsum, dsink_bias -@torch.compile class ParallelAttentionFunction(torch.autograd.Function): @staticmethod @@ -779,7 +788,6 @@ def parallel_attn( chunk_indices: torch.LongTensor | None = None, *, sink_bias: torch.Tensor | None = None, - **kwargs ) -> torch.Tensor: r""" Args: @@ -811,20 +819,17 @@ def parallel_attn( attention mass to a learnable "no-op" target: p_i = exp(s_i) / (sum_j exp(s_j) + exp(sink_bias[h])) o = sum_i p_i * v_i # sink slot contributes no value - When `None`, standard softmax is used. Reserved name: the future - `sink_tokens_*` kwargs will support Xiao 2024-style K/V sink tokens - and may be combined with `sink_bias`. + When `None`, standard softmax is used. Returns: o (torch.Tensor): Outputs of shape `[B, T, HQ, V]`. """ - if 'head_first' in kwargs: - raise DeprecationWarning( - "head_first has been removed. Inputs must be in `[B, T, H, ...]` format.", - ) if scale is None: scale = k.shape[-1] ** -0.5 + HQ, H = q.shape[2], k.shape[2] + if H == 0 or HQ % H != 0: + raise ValueError(f"The number of query heads ({HQ}) must be divisible by the number of key/value heads ({H}).") if cu_seqlens is not None and q.shape[0] != 1: raise ValueError( f"The batch size is expected to be 1 rather than {q.shape[0]} when using `cu_seqlens`. " diff --git a/build/torch-xpu/ops/utils/backends/__init__.py b/build/torch-cuda/ops/attnres/backends/__init__.py similarity index 50% rename from build/torch-xpu/ops/utils/backends/__init__.py rename to build/torch-cuda/ops/attnres/backends/__init__.py index 4ac07c0078d87c3792b0fce8d230c15062c9ab1b..9bc2be22ababea5b5144aed9f70034bbf144e2e8 100644 --- a/build/torch-xpu/ops/utils/backends/__init__.py +++ b/build/torch-cuda/ops/attnres/backends/__init__.py @@ -5,12 +5,15 @@ # For a list of all contributors, visit: # https://github.com/fla-org/flash-linear-attention/graphs/contributors -"""Utils op backends.""" +"""AttnRes backends.""" -from ....ops.backends import BackendRegistry, dispatch -from ....ops.utils.backends.triton_ascend import TritonAscendUtilsBackend -utils_registry = BackendRegistry('utils') -utils_registry.register(TritonAscendUtilsBackend()) +__all__: list[str] = [] -__all__ = ['dispatch', 'utils_registry'] +# gluon.py imports triton.experimental at module load; keep the Ascend backend usable when Gluon is unavailable. +try: + from ....ops.attnres.backends.gluon import AttnResGluonBackend + + __all__.append('AttnResGluonBackend') +except ImportError: + pass diff --git a/build/torch-cuda/ops/attnres/backends/gluon.py b/build/torch-cuda/ops/attnres/backends/gluon.py new file mode 100644 index 0000000000000000000000000000000000000000..0dbcdec41e79df69a5f85423e26dc482ce63b78b --- /dev/null +++ b/build/torch-cuda/ops/attnres/backends/gluon.py @@ -0,0 +1,772 @@ +# 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 + +"""Gluon backend for the fused AttnRes op (see ``fla/ops/attnres/fused.py`` for the Triton reference). + +Where the Triton kernel leaves layout, shared memory, and load/compute overlap to the compiler, this +port makes them explicit: residual sources are indexed statically (``L`` is a constexpr, so no +pointer-table gather), V tiles are staged through shared memory with ``cp.async``, and the backward +keeps all ``L`` tiles resident when they fit and streams a 2-deep ring otherwise. + +Opt-in and auto-dispatched like the other FLA backends: enable with ``FLA_GLUON=1`` or ``FLA_ATTNRES_GLUON=1``; +both are off by default. The verifier selects it for suitable CUDA calls and falls back to Triton elsewhere. +Numerical parity with Triton is the frozen ``tests/ops/test_attnres.py``. +""" + +from __future__ import annotations + +from collections.abc import Sequence + +import torch +import triton +from triton.experimental import gluon +from triton.experimental.gluon import language as gl +from triton.experimental.gluon.language.nvidia.ampere import async_copy as cp + +from ....backends import BaseBackend, register +from ....ops.utils.cache import fla_cache_autotune +from ....ops.utils.op import barrier +from ....utils import ( + autocast_custom_bwd, + autocast_custom_fwd, + autotune_cache_kwargs, + input_guard, +) + +# tokens per bwd program (BT * KT); each program spills one fp32 dqw/dow partial row +GROUP = 32 + +# each config sets `num_warps=W` and constexpr `NW=W` together so the static layouts stay in sync +# with the launch warp count. num_stages is meaningless in Gluon (no compiler pipelining). +_FWD_CONFIGS = [ + triton.Config({'BT': BT, 'NW': num_warps}, num_warps=num_warps) + for BT in [1, 2, 4] + for num_warps in [4, 8] +] +_BWD_CONFIGS = [ + triton.Config({'BT': BT, 'NW': num_warps}, num_warps=num_warps) + for BT in [1, 2, 4] + for num_warps in [4, 8, 16] +] + +# leave headroom below the 228KB Hopper/Blackwell carveout for barriers/paddings +_SMEM_BUDGET = 192 * 1024 + + +def _prune_fwd_smem(configs, named_args, **kwargs): + args = {**named_args, **kwargs} + def need(c): return 2 * c.kwargs['BT'] * args['BD'] * args['ES'] # noqa: E731 + keep = [c for c in configs if need(c) <= _SMEM_BUDGET] or configs[:1] + if args['L'] > 8: + # the source loop is fully unrolled; cap the sweep so big-L compiles stay sane + keep = [c for c in keep if c.kwargs['BT'] == 2] or keep[:2] + return keep + + +def _prune_bwd_smem(configs, named_args, **kwargs): + # streaming fallback needs a 2-deep V ring plus the do tile + args = {**named_args, **kwargs} + def need(c): return 3 * c.kwargs['BT'] * args['BD'] * args['ES'] # noqa: E731 + keep = [c for c in configs if need(c) <= _SMEM_BUDGET] or configs[:1] + if args['L'] > 8: + # both passes unroll L tile bodies; cap the sweep so big-L compiles stay sane + keep = [c for c in keep if c.kwargs['BT'] == 2 and c.num_warps in (8, 16)] or keep[:2] + return keep + + +@gluon.jit +def _sum2_combine(a0, a1, b0, b1): + return a0 + b0, a1 + b1 + + +@gluon.constexpr_function +def _warp_split(BT, NW): + # spread warps over tokens first, remainder over D + w0 = min(BT, NW) + return [w0, NW // w0] + + +@gluon.constexpr_function +def _lane_split(BD): + # lanes cover [32 // t1, t1] so that 4-wide lane segments span exactly min(BD, 128) columns + t1 = min(max(BD // 4, 1), 32) + return [32 // t1, t1] + + +@gluon.constexpr_function +def _resident_tiles(L, BT, BD, ES): + # True when all L v tiles plus the do tile fit in the smem budget + return (L + 1) * BT * BD * ES <= 192 * 1024 + + +@gluon.jit +def _dv_tile(b_v, b_do, b_qw, b_rstd, b_logit, b_p, b_delta, scale: gl.constexpr, D: gl.constexpr, L: gl.constexpr): + # softmax bwd with delta already known; returns the dv tile and this tile's dqw contribution. + # for L == 1 the softmax is constant so dlogit vanishes identically; branching on the constexpr + # keeps ds exactly zero instead of relying on two reductions cancelling bitwise + if L == 1: + b_ds = gl.zeros_like(b_delta) + else: + b_dp = gl.sum(b_v * b_do, axis=1) + b_ds = b_p * (b_dp - b_delta) * scale + # [BT, BD] + b_k = b_v * b_rstd[:, None] + b_dv = b_p[:, None] * b_do + (b_ds * b_rstd)[:, None] * (b_qw[None, :] - b_k * (b_logit / D)[:, None]) + return b_dv, gl.sum(b_ds[:, None] * b_k, axis=0) + + +@fla_cache_autotune( + configs=_FWD_CONFIGS, + key=['L', 'D', 'R', 'ES', 'HAS_ONORM', 'SAVE_OPRE'], + prune_configs_by={'early_config_prune': _prune_fwd_smem}, + **autotune_cache_kwargs, +) +@gluon.jit +def attnres_fwd_kernel_gluon( + q, + res, + w, + ow, + o, + o_pre, + rstd, + logit, + lse, + N, + L: gl.constexpr, + D: gl.constexpr, + eps: gl.constexpr, + scale: gl.constexpr, + BT: gl.constexpr, + BD: gl.constexpr, + NW: gl.constexpr, + R: gl.constexpr, + ES: gl.constexpr, + HAS_ONORM: gl.constexpr, + SAVE_OPRE: gl.constexpr, +): + W: gl.constexpr = _warp_split(BT, NW) + # [BT, BD] row-major, R-wide lane segments along D; warps over tokens first, then D + blk: gl.constexpr = gl.BlockedLayout([1, R], [1, 32], W, [1, 0]) + sl_d: gl.constexpr = gl.SliceLayout(0, blk) + sl_t: gl.constexpr = gl.SliceLayout(1, blk) + smem_layout: gl.constexpr = gl.SwizzledSharedLayout(1, 1, 1, [1, 0]) + NEEDS_DMASK: gl.constexpr = D != BD + + i_n = gl.program_id(0).to(gl.int64) * BT + # [BT]; OOB rows are clamped to a valid token (never masked at load) and masked only at stores, + # so no lane ever reads uninitialized smem under NaN poisoning + o_t = i_n + gl.arange(0, BT, layout=sl_t) + m_t = o_t < N + o_tc = gl.minimum(o_t, N - 1) + # [BD] + o_d = gl.arange(0, BD, layout=sl_d) + m_d = o_d < D + + b_qw = gl.load(q + o_d, mask=m_d, other=0.).to(gl.float32) * gl.load(w + o_d, mask=m_d, other=0.).to(gl.float32) + + # [BT, BD] global offsets shared by every source + o_v = o_tc[:, None] * D + o_d[None, :] + + # double-buffered cp.async over the L sources + b_vs = gl.allocate_shared_memory(res[0].dtype.element_ty, [2, BT, BD], smem_layout) + if NEEDS_DMASK: + cp.async_copy_global_to_shared(b_vs.index(0), res[0] + o_v, mask=m_d[None, :], eviction_policy="evict_first") + else: + cp.async_copy_global_to_shared(b_vs.index(0), res[0] + o_v, eviction_policy="evict_first") + cp.commit_group() + + b_m = gl.full([BT], float('-inf'), gl.float32, layout=sl_t) + b_acc = gl.zeros([BT], gl.float32, layout=sl_t) + b_o = gl.zeros([BT, BD], gl.float32, layout=blk) + for i_l in gl.static_range(L): + if i_l + 1 < L: + if NEEDS_DMASK: + cp.async_copy_global_to_shared( + b_vs.index((i_l + 1) % 2), + res[i_l + 1] + o_v, + mask=m_d[None, :], + eviction_policy="evict_first", + ) + else: + cp.async_copy_global_to_shared(b_vs.index((i_l + 1) % 2), res[i_l + 1] + o_v, eviction_policy="evict_first") + cp.commit_group() + cp.wait_group(1) + else: + cp.wait_group(0) + + # [BT, BD] + b_v = b_vs.index(i_l % 2).load(blk).to(gl.float32) + if NEEDS_DMASK: + b_v = gl.where(m_d[None, :], b_v, 0.) + + # one cross-warp pass for both row reductions + b_v2, b_vq = gl.reduce((b_v * b_v, b_v * b_qw[None, :]), axis=1, combine_fn=_sum2_combine) + # [BT] + b_rstd = gl.rsqrt(b_v2 / D + eps) + b_logit = b_vq * b_rstd + b_s = b_logit * scale + + b_m, b_mp = gl.maximum(b_m, b_s), b_m + b_r = gl.exp(b_mp - b_m) + b_p = gl.exp(b_s - b_m) + b_acc = b_acc * b_r + b_p + b_o = b_o * b_r[:, None] + b_p[:, None] * b_v + + gl.store(rstd + i_l * N + o_t, b_rstd.to(rstd.dtype.element_ty), mask=m_t) + gl.store(logit + i_l * N + o_t, b_logit.to(logit.dtype.element_ty), mask=m_t) + # cp.async is same-lane staging, but the CTA must sync before the next issue reuses slot i_l % 2 + barrier() + + gl.store(lse + o_t, b_m + gl.log(b_acc), mask=m_t) + + b_o = b_o / b_acc[:, None] + if SAVE_OPRE: + gl.store(o_pre + o_v, b_o.to(o_pre.dtype.element_ty), mask=m_t[:, None] & m_d[None, :]) + if HAS_ONORM: + # masked cols are already zero, so the row reduction needs no extra select + b_o_rstd = gl.rsqrt(gl.sum(b_o * b_o, axis=1) / D + eps) + b_ow = gl.load(ow + o_d, mask=m_d, other=0.).to(gl.float32) + b_o = b_o * b_o_rstd[:, None] * b_ow[None, :] + gl.store(o + o_v, b_o.to(o.dtype.element_ty), mask=m_t[:, None] & m_d[None, :]) + + +@fla_cache_autotune( + configs=_BWD_CONFIGS, + key=['L', 'D', 'R', 'ES', 'HAS_ONORM', 'SAVE_OPRE'], + prune_configs_by={'early_config_prune': _prune_bwd_smem}, + **autotune_cache_kwargs, +) +@gluon.jit +def attnres_bwd_kernel_dv_gluon( + q, + res, + w, + ow, + o_pre, + rstd, + logit, + lse, + do, + dres, + dqw, + dow_partial, + N, + L: gl.constexpr, + D: gl.constexpr, + eps: gl.constexpr, + scale: gl.constexpr, + G: gl.constexpr, + BT: gl.constexpr, + BD: gl.constexpr, + NW: gl.constexpr, + R: gl.constexpr, + ES: gl.constexpr, + HAS_ONORM: gl.constexpr, + SAVE_OPRE: gl.constexpr, +): + W: gl.constexpr = _warp_split(BT, NW) + blk: gl.constexpr = gl.BlockedLayout([1, R], [1, 32], W, [1, 0]) + sl_d: gl.constexpr = gl.SliceLayout(0, blk) + sl_t: gl.constexpr = gl.SliceLayout(1, blk) + smem_layout: gl.constexpr = gl.SwizzledSharedLayout(1, 1, 1, [1, 0]) + NEEDS_DMASK: gl.constexpr = D != BD + KT: gl.constexpr = G // BT + # all L v tiles stay resident when they fit (V then hits global exactly once per token); + # otherwise they stream through a 2-deep ring and each pass re-reads V, like Triton + RESIDENT: gl.constexpr = _resident_tiles(L, BT, BD, ES) + NB: gl.constexpr = L if RESIDENT else 2 + + i_g = gl.program_id(0).to(gl.int64) * G + + o_d = gl.arange(0, BD, layout=sl_d) + m_d = o_d < D + b_qw = gl.load(q + o_d, mask=m_d, other=0.).to(gl.float32) * gl.load(w + o_d, mask=m_d, other=0.).to(gl.float32) + if HAS_ONORM: + b_ow = gl.load(ow + o_d, mask=m_d, other=0.).to(gl.float32) + + # [BD] fp32 partials accumulated across all G tokens, spilled once at the end + b_dqw = gl.zeros([BD], gl.float32, layout=sl_d) + b_dow = gl.zeros([BD], gl.float32, layout=sl_d) + + b_vs = gl.allocate_shared_memory(res[0].dtype.element_ty, [NB, BT, BD], smem_layout) + b_dos = gl.allocate_shared_memory(do.dtype.element_ty, [BT, BD], smem_layout) + + for t in range(KT): + # [BT]; OOB rows are clamped to a valid token and their do rows zeroed below, so every + # masked contribution vanishes without reading uninitialized smem under NaN poisoning + o_t = i_g + t * BT + gl.arange(0, BT, layout=sl_t) + m_t = o_t < N + o_tc = gl.minimum(o_t, N - 1) + o_v = o_tc[:, None] * D + o_d[None, :] + + # do goes first so its wait count is independent of the V schedule + if NEEDS_DMASK: + cp.async_copy_global_to_shared(b_dos, do + o_v, mask=m_d[None, :]) + else: + cp.async_copy_global_to_shared(b_dos, do + o_v) + cp.commit_group() + if RESIDENT: + for i_l in gl.static_range(L): + if NEEDS_DMASK: + cp.async_copy_global_to_shared( + b_vs.index(i_l), + res[i_l] + o_v, + mask=m_d[None, :], + eviction_policy="evict_first", + ) + else: + cp.async_copy_global_to_shared(b_vs.index(i_l), res[i_l] + o_v, eviction_policy="evict_first") + cp.commit_group() + cp.wait_group(L) + else: + if NEEDS_DMASK: + cp.async_copy_global_to_shared(b_vs.index(0), res[0] + o_v, mask=m_d[None, :], eviction_policy="evict_first") + else: + cp.async_copy_global_to_shared(b_vs.index(0), res[0] + o_v, eviction_policy="evict_first") + cp.commit_group() + cp.wait_group(1) + + b_lse = gl.load(lse + o_tc).to(gl.float32) + # [BT, BD] + b_do = b_dos.load(blk).to(gl.float32) + b_do = gl.where(m_t[:, None], b_do, 0.) + if NEEDS_DMASK: + b_do = gl.where(m_d[None, :], b_do, 0.) + + if SAVE_OPRE: + b_opre = gl.load(o_pre + o_v, mask=m_d[None, :], other=0.).to(gl.float32) + if RESIDENT: + cp.wait_group(0) + else: + # level 1: recompute the mix sum_l p_l * v_l + b_opre = gl.zeros([BT, BD], gl.float32, layout=blk) + for i_l in gl.static_range(L): + if RESIDENT: + cp.wait_group(L - 1 - i_l) + b_v = b_vs.index(i_l).load(blk).to(gl.float32) + else: + if i_l + 1 < L: + if NEEDS_DMASK: + cp.async_copy_global_to_shared( + b_vs.index((i_l + 1) % 2), + res[i_l + 1] + o_v, + mask=m_d[None, :], + eviction_policy="evict_first", + ) + else: + cp.async_copy_global_to_shared( + b_vs.index((i_l + 1) % 2), + res[i_l + 1] + o_v, + eviction_policy="evict_first", + ) + cp.commit_group() + cp.wait_group(1) + else: + cp.wait_group(0) + b_v = b_vs.index(i_l % 2).load(blk).to(gl.float32) + if NEEDS_DMASK: + b_v = gl.where(m_d[None, :], b_v, 0.) + b_logit = gl.load(logit + i_l * N + o_tc).to(gl.float32) + b_p = gl.exp(b_logit * scale - b_lse) + b_opre += b_p[:, None] * b_v + if not RESIDENT: + # the slot just read is the next issue target + barrier() + + # output RMSNorm bwd: turn b_do into the gradient w.r.t. the pre-norm output + if HAS_ONORM: + b_o_rstd = gl.rsqrt(gl.sum(b_opre * b_opre, axis=1) / D + eps) + b_xhat = b_opre * b_o_rstd[:, None] + b_c1 = gl.sum(b_xhat * b_ow[None, :] * b_do, axis=1) / D + b_dow += gl.sum(b_xhat * b_do, axis=0) + b_do = (b_ow[None, :] * b_do - b_xhat * b_c1[:, None]) * b_o_rstd[:, None] + # [BT] delta = sum_l p*dp = + b_delta = gl.sum(b_do * b_opre, axis=1) + + if not RESIDENT and not SAVE_OPRE: + # re-prime the ring for the second streaming pass + if NEEDS_DMASK: + cp.async_copy_global_to_shared(b_vs.index(0), res[0] + o_v, mask=m_d[None, :], eviction_policy="evict_first") + else: + cp.async_copy_global_to_shared(b_vs.index(0), res[0] + o_v, eviction_policy="evict_first") + cp.commit_group() + + for i_l in gl.static_range(L): + if RESIDENT: + # tiles were already waited on above; second use comes straight from smem + b_v = b_vs.index(i_l).load(blk).to(gl.float32) + else: + if i_l + 1 < L: + if NEEDS_DMASK: + cp.async_copy_global_to_shared( + b_vs.index((i_l + 1) % 2), + res[i_l + 1] + o_v, + mask=m_d[None, :], + eviction_policy="evict_first", + ) + else: + cp.async_copy_global_to_shared( + b_vs.index((i_l + 1) % 2), + res[i_l + 1] + o_v, + eviction_policy="evict_first", + ) + cp.commit_group() + cp.wait_group(1) + else: + cp.wait_group(0) + b_v = b_vs.index(i_l % 2).load(blk).to(gl.float32) + if NEEDS_DMASK: + b_v = gl.where(m_d[None, :], b_v, 0.) + # [BT] + b_rstd = gl.load(rstd + i_l * N + o_tc).to(gl.float32) + b_logit = gl.load(logit + i_l * N + o_tc).to(gl.float32) + b_p = gl.exp(b_logit * scale - b_lse) + + b_dv, b_inc = _dv_tile(b_v, b_do, b_qw, b_rstd, b_logit, b_p, b_delta, scale, D, L) + gl.store(dres[i_l] + o_v, b_dv.to(dres[0].dtype.element_ty), mask=m_t[:, None] & m_d[None, :]) + b_dqw += b_inc + if not RESIDENT: + barrier() + # tiles of iteration t are fully consumed before iteration t+1 refills the buffers + barrier() + + i_p = gl.program_id(0).to(gl.int64) + gl.store(dqw + i_p * D + o_d, b_dqw, mask=m_d) + if HAS_ONORM: + gl.store(dow_partial + i_p * D + o_d, b_dow, mask=m_d) + + +@fla_cache_autotune( + configs=[ + triton.Config({'BN': BN, 'BD': BD, 'NW': num_warps}, num_warps=num_warps) + for BN, BD, num_warps in [(256, 32, 4), (512, 64, 4), (512, 64, 8), (1024, 128, 8)] + ], + key=['NP', 'D', 'HAS_ONORM'], + **autotune_cache_kwargs, +) +@gluon.jit +def attnres_bwd_kernel_dqdw_gluon( + q, + w, + dqw, + dow_partial, + dq, + dw, + dow, + NP, + D: gl.constexpr, + BN: gl.constexpr, + BD: gl.constexpr, + NW: gl.constexpr, + HAS_ONORM: gl.constexpr, +): + i_d = gl.program_id(0).to(gl.int64) + + # [BN, BD] fp32 rows; 4-wide lane segments along D, warps along the partial rows + T: gl.constexpr = _lane_split(BD) + blk: gl.constexpr = gl.BlockedLayout([1, 4], T, [NW, 1], [1, 0]) + sl_d: gl.constexpr = gl.SliceLayout(0, blk) + sl_n: gl.constexpr = gl.SliceLayout(1, blk) + + o_d = i_d * BD + gl.arange(0, BD, layout=sl_d) + m_d = o_d < D + + # column-sum the [NP, D] partials + b_dqw = gl.zeros([BD], gl.float32, layout=sl_d) + b_dow = gl.zeros([BD], gl.float32, layout=sl_d) + for i_n in range(0, NP, BN): + o_n = i_n + gl.arange(0, BN, layout=sl_n) + m = (o_n < NP)[:, None] & m_d[None, :] + p = o_n[:, None].to(gl.int64) * D + o_d[None, :] + b_dqw += gl.sum(gl.load(dqw + p, mask=m, other=0.), axis=0) + if HAS_ONORM: + b_dow += gl.sum(gl.load(dow_partial + p, mask=m, other=0.), axis=0) + + b_q = gl.load(q + o_d, mask=m_d, other=0.).to(gl.float32) + b_w = gl.load(w + o_d, mask=m_d, other=0.).to(gl.float32) + gl.store(dq + o_d, (b_dqw * b_w).to(dq.dtype.element_ty), mask=m_d) + gl.store(dw + o_d, (b_dqw * b_q).to(dw.dtype.element_ty), mask=m_d) + if HAS_ONORM: + gl.store(dow + o_d, b_dow.to(dow.dtype.element_ty), mask=m_d) + + +def _vec_width(dtype: torch.dtype, d: int, bd: int) -> int: + # widest lane segment (<= 16B) that divides D without forcing intra-warp broadcast; + # cp.async needs >= 4B per segment + r = min(16 // dtype.itemsize, max(1, bd // 32)) + while r > 1 and d % r != 0: + r //= 2 + if r * dtype.itemsize < 4: + raise ValueError(f"Gluon attnres requires D * itemsize >= 128 bytes, got D={d} ({dtype})") + return r + + +def _check_sources(tensors: Sequence[torch.Tensor]) -> tuple[torch.Tensor, ...]: + for t in tensors: + assert t.data_ptr() % 16 == 0, "attnres residual sources must be 16-byte aligned" + return tuple(tensors) + + +def _fused_attnres_fwd( + q: torch.Tensor, + residuals: Sequence[torch.Tensor], + res: tuple[torch.Tensor, ...], + w: torch.Tensor, + ow: torch.Tensor | None, + eps: float, + scale: float, + checkpoint_level: int, +) -> tuple[torch.Tensor, torch.Tensor | None, torch.Tensor, torch.Tensor, torch.Tensor]: + if not residuals[0].is_cuda: + raise ValueError("Gluon attnres requires CUDA tensors") + + output_shape = residuals[0].shape + L, N, D = len(residuals), residuals[0].numel() // output_shape[-1], output_shape[-1] + + dtype = residuals[0].dtype + stats_shape = (L, *output_shape[:-1]) + + has_onorm = ow is not None + save_opre = checkpoint_level == 0 + o = torch.empty(output_shape, device=residuals[0].device, dtype=dtype) + o_pre = torch.empty(output_shape, device=residuals[0].device, dtype=dtype) if save_opre else None + lse = torch.empty(output_shape[:-1], device=residuals[0].device, dtype=torch.float32) + rstd = torch.empty(stats_shape, device=residuals[0].device, dtype=torch.float32) + logit = torch.empty_like(rstd) + + def grid(meta): return (triton.cdiv(N, meta['BT']),) + attnres_fwd_kernel_gluon[grid]( + q=q, + res=res, + w=w, + ow=ow, + o=o, + o_pre=o_pre, + rstd=rstd, + logit=logit, + lse=lse, + N=N, + L=L, + D=D, + eps=eps, + scale=scale, + BD=triton.next_power_of_2(D), + R=_vec_width(dtype, D, triton.next_power_of_2(D)), + ES=dtype.itemsize, + HAS_ONORM=has_onorm, + SAVE_OPRE=save_opre, + ) + + return o, o_pre, rstd, logit, lse + + +def _fused_attnres_bwd( + do: torch.Tensor, + q: torch.Tensor, + residuals: Sequence[torch.Tensor], + res: tuple[torch.Tensor, ...], + w: torch.Tensor, + ow: torch.Tensor | None, + o_pre: torch.Tensor | None, + rstd: torch.Tensor, + logit: torch.Tensor, + lse: torch.Tensor, + eps: float, + scale: float, + checkpoint_level: int, +) -> tuple[list[torch.Tensor], torch.Tensor, torch.Tensor, torch.Tensor | None]: + L, N, D = len(residuals), do.numel() // do.shape[-1], do.shape[-1] + NP = triton.cdiv(N, GROUP) + + # bwd_dv produces dvs + per-program dqw/dow partials [NP, D]; bwd_dqdw reduces them into dq / dw / dow + has_onorm = ow is not None + if has_onorm: + dow_partial = torch.empty(NP, D, device=do.device, dtype=torch.float32) + dow = torch.empty_like(ow) + else: + dow_partial = dow = None + + dvs = [torch.empty_like(r) for r in residuals] + dres = _check_sources(dvs) + dqw = torch.empty(NP, D, device=do.device, dtype=torch.float32) + dq = torch.empty_like(q) + dw = torch.empty_like(w) + + attnres_bwd_kernel_dv_gluon[(NP,)]( + q=q, + res=res, + w=w, + ow=ow, + o_pre=o_pre, + rstd=rstd, + logit=logit, + lse=lse, + do=do, + dres=dres, + dqw=dqw, + dow_partial=dow_partial, + N=N, + L=L, + D=D, + eps=eps, + scale=scale, + G=GROUP, + BD=triton.next_power_of_2(D), + R=_vec_width(do.dtype, D, triton.next_power_of_2(D)), + ES=do.dtype.itemsize, + HAS_ONORM=has_onorm, + SAVE_OPRE=checkpoint_level == 0, + ) + + def grid(meta): return (triton.cdiv(D, meta['BD']),) + attnres_bwd_kernel_dqdw_gluon[grid]( + q=q, + w=w, + dqw=dqw, + dow_partial=dow_partial, + dq=dq, + dw=dw, + dow=dow, + NP=NP, + D=D, + HAS_ONORM=has_onorm, + ) + + return dvs, dq, dw, dow + + +class FusedAttnresGluonFunction(torch.autograd.Function): + + @staticmethod + @input_guard + @autocast_custom_fwd + def forward( + ctx, + query: torch.Tensor, + rms_weight: torch.Tensor, + output_rms_weight: torch.Tensor | None, + rms_eps: float, + scale: float, + return_weights: bool, + checkpoint_level: int, + *residuals: torch.Tensor, + ) -> tuple[torch.Tensor, torch.Tensor]: + res = _check_sources(residuals) + o, o_pre, rstd, logit, lse = _fused_attnres_fwd( + q=query, + residuals=residuals, + res=res, + w=rms_weight, + ow=output_rms_weight, + eps=rms_eps, + scale=scale, + checkpoint_level=checkpoint_level, + ) + ctx.save_for_backward(query, rms_weight, output_rms_weight, o_pre, rstd, logit, lse, *residuals) + ctx.eps = rms_eps + ctx.scale = scale + ctx.checkpoint_level = checkpoint_level + ctx.res = res + # probs are materialized only when requested; bwd recomputes them from logit + lse + p = (logit * scale - lse).exp() if return_weights else o.new_empty(0) + ctx.mark_non_differentiable(p) + return o, p + + @staticmethod + @input_guard + @autocast_custom_bwd + def backward( + ctx, + do: torch.Tensor, + dp: torch.Tensor | None = None, + ): + del dp + query, rms_weight, output_rms_weight, o_pre, rstd, logit, lse, *residuals = ctx.saved_tensors + dvs, dq, dw, dow = _fused_attnres_bwd( + do=do, + q=query, + residuals=residuals, + res=ctx.res, + w=rms_weight, + ow=output_rms_weight, + o_pre=o_pre, + rstd=rstd, + logit=logit, + lse=lse, + eps=ctx.eps, + scale=ctx.scale, + checkpoint_level=ctx.checkpoint_level, + ) + return (dq, dw, dow, None, None, None, None, *dvs) + + +def _run( + query: torch.Tensor, + residuals: Sequence[torch.Tensor], + rms_weight: torch.Tensor, + output_rms_weight: torch.Tensor | None = None, + rms_eps: float = 1e-6, + scale: float = 1.0, + return_weights: bool = False, + checkpoint_level: int = 1, +) -> torch.Tensor | tuple[torch.Tensor, ...]: + output_shape = residuals[0].shape + D = output_shape[-1] + flat_residuals = tuple(r.reshape(-1, D).contiguous() for r in residuals) + + o, p = FusedAttnresGluonFunction.apply( + query, rms_weight, output_rms_weight, rms_eps, scale, return_weights, checkpoint_level, *flat_residuals, + ) + o = o.view(output_shape) + + if return_weights: + p = p.view(len(residuals), *output_shape[:-1]) + return o, p + return o + + +@register('attnres') +class AttnResGluonBackend(BaseBackend): + """Dispatch entry for the Gluon AttnRes kernels (see the module docstring for the design). + + Off by default; enable with ``FLA_GLUON=1`` or ``FLA_ATTNRES_GLUON=1``. + The verifier accepts CUDA calls on SM80+ with ``D * itemsize >= 128`` bytes and otherwise defers to Triton. + """ + + backend_type = "gluon" + package_name = "triton.experimental.gluon" # ships with Triton 3.5+; absent on older builds + env_var = "FLA_ATTNRES_GLUON" + default_enable = False + priority = 3 + + def fused_attnres_verifier( + self, + query: torch.Tensor, + residuals: Sequence[torch.Tensor], + rms_weight: torch.Tensor, + output_rms_weight: torch.Tensor | None = None, + rms_eps: float = 1e-6, + scale: float = 1.0, + return_weights: bool = False, + checkpoint_level: int = 1, + **kwargs, + ) -> tuple[bool, str | None]: + if not residuals or not residuals[0].is_cuda: + return False, "gluon AttnRes backend requires CUDA tensors" + v = residuals[0] + if torch.cuda.get_device_capability(v.device)[0] < 8: + return False, "gluon AttnRes backend requires NVIDIA SM80+ (cp.async)" + if v.shape[-1] * v.element_size() < 128: + return False, "gluon AttnRes backend requires D * itemsize >= 128 bytes" + return True, None + + def fused_attnres(self, *args, **kwargs): + return _run(*args, **kwargs) + + +__all__ = ["AttnResGluonBackend"] diff --git a/build/torch-cuda/ops/attnres/fused.py b/build/torch-cuda/ops/attnres/fused.py index ea0389b68b6abce3dd143aa460db5e86f043eced..960a0c12425c516ae2d816e9fce251a8438ca4e7 100644 --- a/build/torch-cuda/ops/attnres/fused.py +++ b/build/torch-cuda/ops/attnres/fused.py @@ -13,6 +13,7 @@ import torch import triton import triton.language as tl +from ...backends import dispatch from ...ops.utils.cache import fla_cache_autotune from ...ops.utils.op import exp from ...utils import ( @@ -69,7 +70,7 @@ def attnres_fwd_kernel( b_o = tl.zeros([BD], dtype=tl.float32) for i_l in range(tl.cdiv(L, BL)): # [BL] - o_l = i_l * BL + tl.arange(0, BL) + o_l = (i_l * BL).to(tl.int64) + tl.arange(0, BL) m_l = o_l < L # per-tile base pointers from the length-L2 padded tuple; OOB rows keep res[0] and are masked by m_l p_v = res[0] + o_l * 0 @@ -99,25 +100,25 @@ def attnres_fwd_kernel( b_o = b_o * b_r + tl.sum(b_p[:, None] * b_v, axis=0) # rstd and logit saved for bwd_dv - p_rstd = tl.make_block_ptr(rstd + i_n, (L,), (N,), (i_l * BL,), (BL,), (0,)) - p_logit = tl.make_block_ptr(logit + i_n, (L,), (N,), (i_l * BL,), (BL,), (0,)) - tl.store(p_rstd, b_rstd.to(rstd.dtype.element_ty), boundary_check=(0,)) - tl.store(p_logit, b_logit.to(logit.dtype.element_ty), boundary_check=(0,)) + p_rstd = rstd + i_n + o_l * N + p_logit = logit + i_n + o_l * N + tl.store(p_rstd, b_rstd.to(rstd.dtype.element_ty), mask=m_l) + tl.store(p_logit, b_logit.to(logit.dtype.element_ty), mask=m_l) tl.store(lse + i_n, b_m + tl.log(b_acc)) # [BD] pre-norm mixed residual sum_l p_l * v_l b_o = b_o / b_acc if SAVE_OPRE: - p_o_pre = tl.make_block_ptr(o_pre + i_n * D, (D,), (1,), (0,), (BD,), (0,)) - tl.store(p_o_pre, b_o.to(p_o_pre.dtype.element_ty), boundary_check=(0,)) + p_o_pre = o_pre + i_n * D + o_d + tl.store(p_o_pre, b_o.to(p_o_pre.dtype.element_ty), mask=m_d) # fold the optional output RMSNorm into the returned output o (o_rstd is recomputed from o_pre in bwd, not stored) if HAS_ONORM: b_o_rstd = tl.rsqrt(tl.sum(tl.where(m_d, b_o * b_o, 0.0), axis=0) / D + eps) b_ow = tl.load(ow + o_d, mask=m_d, other=0.).to(tl.float32) b_o = b_o * b_o_rstd * b_ow - p_o = tl.make_block_ptr(o + i_n * D, (D,), (1,), (0,), (BD,), (0,)) - tl.store(p_o, b_o.to(p_o.dtype.element_ty), boundary_check=(0,)) + p_o = o + i_n * D + o_d + tl.store(p_o, b_o.to(p_o.dtype.element_ty), mask=m_d) @fla_cache_autotune( @@ -162,16 +163,16 @@ def attnres_bwd_kernel_dv( m_d = o_d < D b_qw = tl.load(q + o_d, mask=m_d, other=0.).to(tl.float32) * tl.load(w + o_d, mask=m_d, other=0.).to(tl.float32) b_lse = tl.load(lse + i_n).to(tl.float32) - p_do = tl.make_block_ptr(do + i_n * D, (D,), (1,), (0,), (BD,), (0,)) - b_do = tl.load(p_do, boundary_check=(0,), padding_option="zero").to(tl.float32) + p_do = do + i_n * D + o_d + b_do = tl.load(p_do, mask=m_d, other=0.0).to(tl.float32) if SAVE_OPRE: - p_o_pre = tl.make_block_ptr(o_pre + i_n * D, (D,), (1,), (0,), (BD,), (0,)) - b_o_pre = tl.load(p_o_pre, boundary_check=(0,), padding_option="zero").to(tl.float32) + p_o_pre = o_pre + i_n * D + o_d + b_o_pre = tl.load(p_o_pre, mask=m_d, other=0.0).to(tl.float32) else: # level 1: recompute the mix sum_l p_l * v_l from V b_o_pre = tl.zeros([BD], dtype=tl.float32) for i_l in range(tl.cdiv(L, BL)): - o_l = i_l * BL + tl.arange(0, BL) + o_l = (i_l * BL).to(tl.int64) + tl.arange(0, BL) m_l = o_l < L p_v = res[0] + o_l * 0 for i in tl.static_range(1, L2): @@ -182,8 +183,8 @@ def attnres_bwd_kernel_dv( mask=m_l[:, None] & m_d[None, :], other=0.0, ).to(tl.float32) - p_logit = tl.make_block_ptr(logit + i_n, (L,), (N,), (i_l * BL,), (BL,), (0,)) - b_logit = tl.load(p_logit, boundary_check=(0,), padding_option="zero").to(tl.float32) + p_logit = logit + i_n + o_l * N + b_logit = tl.load(p_logit, mask=m_l, other=0.0).to(tl.float32) b_p = tl.where(m_l, exp(b_logit * scale - b_lse), 0.0) b_o_pre += tl.sum(b_p[:, None] * b_v, axis=0) @@ -193,8 +194,8 @@ def attnres_bwd_kernel_dv( b_ow = tl.load(ow + o_d, mask=m_d, other=0.).to(tl.float32) b_xhat = b_o_pre * b_o_rstd b_c1 = tl.sum(tl.where(m_d, b_xhat * b_ow * b_do, 0.0), axis=0) / D - p_dow = tl.make_block_ptr(dow_partial + i_n * D, (D,), (1,), (0,), (BD,), (0,)) - tl.store(p_dow, (b_xhat * b_do).to(p_dow.dtype.element_ty), boundary_check=(0,)) + p_dow = dow_partial + i_n * D + o_d + tl.store(p_dow, (b_xhat * b_do).to(p_dow.dtype.element_ty), mask=m_d) b_do = (b_ow * b_do - b_xhat * b_c1) * b_o_rstd # delta = sum_l p*dp = b_delta = tl.sum(tl.where(m_d, b_do * b_o_pre, 0.0), axis=0) @@ -203,7 +204,7 @@ def attnres_bwd_kernel_dv( b_dqw = tl.zeros([BD], dtype=tl.float32) for i_l in range(tl.cdiv(L, BL)): # [BL] - o_l = i_l * BL + tl.arange(0, BL) + o_l = (i_l * BL).to(tl.int64) + tl.arange(0, BL) m_l = o_l < L m_v = m_l[:, None] & m_d[None, :] # per-tile source / dv base pointers from the length-L2 padded tuple @@ -221,11 +222,11 @@ def attnres_bwd_kernel_dv( other=0.0, ).to(tl.float32) - p_rstd = tl.make_block_ptr(rstd + i_n, (L,), (N,), (i_l * BL,), (BL,), (0,)) - p_logit = tl.make_block_ptr(logit + i_n, (L,), (N,), (i_l * BL,), (BL,), (0,)) + p_rstd = rstd + i_n + o_l * N + p_logit = logit + i_n + o_l * N # [BL]; recompute probs from logit + lse, OOB rows masked to 0 - b_rstd = tl.load(p_rstd, boundary_check=(0,), padding_option="zero").to(tl.float32) - b_logit = tl.load(p_logit, boundary_check=(0,), padding_option="zero").to(tl.float32) + b_rstd = tl.load(p_rstd, mask=m_l, other=0.0).to(tl.float32) + b_logit = tl.load(p_logit, mask=m_l, other=0.0).to(tl.float32) b_p = tl.where(m_l, exp(b_logit * scale - b_lse), 0.0) # softmax bwd with delta already known @@ -242,8 +243,8 @@ def attnres_bwd_kernel_dv( # [BD] b_dqw += tl.sum(b_ds[:, None] * b_k, axis=0) - p_dqw = tl.make_block_ptr(dqw + i_n * D, (D,), (1,), (0,), (BD,), (0,)) - tl.store(p_dqw, b_dqw, boundary_check=(0,)) + p_dqw = dqw + i_n * D + o_d + tl.store(p_dqw, b_dqw, mask=m_d) @fla_cache_autotune( @@ -270,7 +271,7 @@ def attnres_bwd_kernel_dqdw( BD: tl.constexpr, HAS_ONORM: tl.constexpr, ): - i_d = tl.program_id(0).to(tl.int32) + i_d = tl.program_id(0).to(tl.int64) # [BD] o_d = i_d * BD + tl.arange(0, BD) @@ -281,11 +282,13 @@ def attnres_bwd_kernel_dqdw( b_dqw = tl.zeros([BD], dtype=tl.float32) b_dow = tl.zeros([BD], dtype=tl.float32) for i_n in range(0, N, BN): - p_dqw = tl.make_block_ptr(dqw, (N, D), (D, 1), (i_n, i_d * BD), (BN, BD), (1, 0)) - b_dqw += tl.sum(tl.load(p_dqw, boundary_check=(0, 1), padding_option="zero").to(tl.float32), axis=0) + o_n = i_n.to(tl.int64) + tl.arange(0, BN) + m_nd = (o_n[:, None] < N) & m_d[None, :] + p_dqw = dqw + o_n[:, None] * D + o_d[None, :] + b_dqw += tl.sum(tl.load(p_dqw, mask=m_nd, other=0.0).to(tl.float32), axis=0) if HAS_ONORM: - p_dow = tl.make_block_ptr(dow_partial, (N, D), (D, 1), (i_n, i_d * BD), (BN, BD), (1, 0)) - b_dow += tl.sum(tl.load(p_dow, boundary_check=(0, 1), padding_option="zero").to(tl.float32), axis=0) + p_dow = dow_partial + o_n[:, None] * D + o_d[None, :] + b_dow += tl.sum(tl.load(p_dow, mask=m_nd, other=0.0).to(tl.float32), axis=0) # the logit uses the q * w product, so dq = (sum_n dqw) * w and dw = (sum_n dqw) * q # [BD] @@ -500,6 +503,7 @@ class FusedAttnresFunction(torch.autograd.Function): return (dq, dw, dow, None, None, None, None, *dvs) +@dispatch('attnres') def fused_attnres( query: torch.Tensor, residuals: Sequence[torch.Tensor], diff --git a/build/torch-cuda/ops/backends/__init__.py b/build/torch-cuda/ops/backends/__init__.py deleted file mode 100644 index 3d9cd4a2770287cfb5dcbecd5aa8dbc2ef53743c..0000000000000000000000000000000000000000 --- a/build/torch-cuda/ops/backends/__init__.py +++ /dev/null @@ -1,223 +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 - -"""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") - - -class BaseBackend: - """Base class for operation-specific backends. - - 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. - `None` indicates no external dependency. Default: `None`. - env_var (str, Optional): - Environment variable name that controls whether the backend is enabled. - `None` means always enabled. Default: `None`. - default_enable (bool, Optional): - Whether the backend is enabled by default when `env_var` is not set. - Set to `False` to require explicit user opt-in. Default: `True`. - priority (int, Optional): - Backend priority. Lower values indicate higher priority. Default: 5. - """ - - 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: - if cls.package_name is None: - 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'] diff --git a/build/torch-cuda/ops/based/fused_chunk.py b/build/torch-cuda/ops/based/fused_chunk.py index f1652f6a6023e6f04ad580d93474523137ea9e13..ad211dbc68c38801c32766fc94529dcdc3e741ba 100644 --- a/build/torch-cuda/ops/based/fused_chunk.py +++ b/build/torch-cuda/ops/based/fused_chunk.py @@ -29,7 +29,7 @@ def fused_chunk_based_fwd_kernel( BK: tl.constexpr, BV: tl.constexpr, ): - i_v, i_k, i_bh = tl.program_id(0), tl.program_id(1), tl.program_id(2) + i_v, i_k, i_bh = tl.program_id(0).to(tl.int64), tl.program_id(1).to(tl.int64), tl.program_id(2).to(tl.int64) o_i = tl.arange(0, BT) @@ -44,10 +44,8 @@ def fused_chunk_based_fwd_kernel( b_h_2o = tl.zeros([BK*BK, BV], dtype=tl.float32) # make block pointers - p_q = tl.make_block_ptr(q + i_bh * T*K, (T, K), (K, 1), (0, i_k * BK), (BT, BK), (1, 0)) - p_k = tl.make_block_ptr(k + i_bh * T*K, (K, T), (1, K), (i_k * BK, 0), (BK, BT), (0, 1)) - p_v = tl.make_block_ptr(v + i_bh * T*V, (T, V), (V, 1), (0, i_v * BV), (BT, BV), (1, 0)) - p_o = tl.make_block_ptr(o + (i_bh + i_k*B*H) * T*V, (T, V), (V, 1), (0, i_v * BV), (BT, BV), (1, 0)) + o_kk = i_k * BK + tl.arange(0, BK) + o_vv = i_v * BV + tl.arange(0, BV) p_z = z + (i_bh + i_k * B * H) * T + tl.arange(0, BT) k_2o = tl.zeros([1, BK * BK], dtype=tl.float32) @@ -55,15 +53,24 @@ def fused_chunk_based_fwd_kernel( k_0o = 0 for i in range(0, tl.cdiv(T, BT)): + o_t = (i * BT).to(tl.int64) + tl.arange(0, BT) + m_t = o_t < T + m_q = m_t[:, None] & (o_kk[None, :] < K) + m_k = (o_kk[:, None] < K) & m_t[None, :] + m_v = m_t[:, None] & (o_vv[None, :] < V) + p_q = q + i_bh * T*K + o_t[:, None] * K + o_kk[None, :] + p_k = k + i_bh * T*K + o_kk[:, None] + o_t[None, :] * K + p_v = v + i_bh * T*V + o_t[:, None] * V + o_vv[None, :] + p_o = o + (i_bh + i_k*B*H) * T*V + o_t[:, None] * V + o_vv[None, :] # [BK, BT] - b_k = tl.load(p_k, boundary_check=(0, 1)) + b_k = tl.load(p_k, mask=m_k, other=0.0) # [BK*BK, BT] b_k_2o = b_k[:, None, :] * b_k[None, :, :] b_k_2o = tl.reshape(b_k_2o, [BK * BK, BT]).to(b_k.dtype) # [BT, BV] - b_v = tl.load(p_v, boundary_check=(0, 1)) + b_v = tl.load(p_v, mask=m_v, other=0.0) # [BT, BK] - b_q = (tl.load(p_q, boundary_check=(0, 1)) * scale).to(b_k.dtype) + b_q = (tl.load(p_q, mask=m_q, other=0.0) * scale).to(b_k.dtype) b_o = tl.zeros([BT, BV], dtype=tl.float32) b_z = tl.zeros([BT], dtype=tl.float32) @@ -72,7 +79,7 @@ def fused_chunk_based_fwd_kernel( b_o += b_h_0o b_z += k_0o # first-order - b_o += tl.dot(b_q, b_h_1o.to(b_q.dtype), allow_tf32=False) + b_o = tl.dot(b_q, b_h_1o.to(b_q.dtype), b_o, allow_tf32=False) b_z += tl.sum(b_q * k_1o, axis=1) # second-order b_q_2o = b_q[:, :, None] * b_q[:, None, :] @@ -91,21 +98,22 @@ def fused_chunk_based_fwd_kernel( b_s = 1 + b_s + 0.5 * b_s * b_s b_s = tl.where(m_s, b_s, 0) b_z += tl.sum(b_s, axis=1) - b_o += tl.dot(b_s.to(b_q.dtype), b_v, allow_tf32=False) + b_o = tl.dot(b_s.to(b_q.dtype), b_v, b_o, allow_tf32=False) # [TB, BV] - tl.store(p_o, b_o.to(p_o.dtype.element_ty), boundary_check=(0, 1)) - tl.store(p_z, b_z.to(p_z.dtype.element_ty), mask=(i * BT + tl.arange(0, BT)) < T) + tl.store(p_o, b_o.to(p_o.dtype.element_ty), mask=m_v) + # only one V-split program writes z: its address does not depend on + # i_v, so all i_v programs would store the same values to the same + # locations concurrently (an unsynchronized same-value WAW). The bwd + # kernel below already guards its i_v-redundant work with i_v == 0. + if i_v == 0: + tl.store(p_z, b_z.to(p_z.dtype.element_ty), mask=(i * BT + tl.arange(0, BT)) < T) # update hidden state # [BK, BV] - b_h_2o = b_h_2o + tl.dot(b_k_2o.to(b_v.dtype), b_v, allow_tf32=False) - b_h_1o = b_h_1o + tl.dot(b_k, b_v, allow_tf32=False) + b_h_2o = tl.dot(b_k_2o.to(b_v.dtype), b_v, b_h_2o, allow_tf32=False) + b_h_1o = tl.dot(b_k, b_v, b_h_1o, allow_tf32=False) b_h_0o = b_h_0o + tl.sum(b_v, axis=0) - p_q = tl.advance(p_q, (BT, 0)) - p_k = tl.advance(p_k, (0, BT)) - p_v = tl.advance(p_v, (BT, 0)) - p_o = tl.advance(p_o, (BT, 0)) p_z += BT @@ -131,7 +139,7 @@ def fused_chunk_based_bwd_kernel( BK: tl.constexpr, BV: tl.constexpr, ): - i_v, i_k, i_bh = tl.program_id(0), tl.program_id(1), tl.program_id(2) + i_v, i_k, i_bh = tl.program_id(0).to(tl.int64), tl.program_id(1).to(tl.int64), tl.program_id(2).to(tl.int64) o_i = tl.arange(0, BT) m_s = o_i[:, None] >= o_i[None, :] @@ -146,27 +154,34 @@ def fused_chunk_based_bwd_kernel( k_1o = tl.zeros([1, BK], dtype=tl.float32) k_2o = tl.zeros([1, BK * BK], dtype=tl.float32) + o_kk = i_k * BK + tl.arange(0, BK) + o_vv = i_v * BV + tl.arange(0, BV) for i in range(0, tl.cdiv(T, BT)): - p_q = tl.make_block_ptr(q + i_bh * T*K, (T, K), (K, 1), (i * BT, i_k * BK), (BT, BK), (1, 0)) - p_k = tl.make_block_ptr(k + i_bh * T*K, (T, K), (K, 1), (i * BT, i_k * BK), (BT, BK), (1, 0)) - p_v = tl.make_block_ptr(v + i_bh * T*V, (V, T), (1, V), (i_v * BV, i * BT), (BV, BT), (0, 1)) - p_do = tl.make_block_ptr(do + i_bh * T*V, (T, V), (V, 1), (i * BT, i_v * BV), (BT, BV), (1, 0)) - p_dq = tl.make_block_ptr(dq + (i_bh + i_v*B*H) * T*K, (T, K), (K, 1), (i*BT, i_k*BK), (BT, BK), (1, 0)) + o_t = (i * BT).to(tl.int64) + tl.arange(0, BT) + m_t = o_t < T + m_q = m_t[:, None] & (o_kk[None, :] < K) + m_v = m_t[:, None] & (o_vv[None, :] < V) + m_kv = (o_vv[:, None] < V) & m_t[None, :] + p_q = q + i_bh * T*K + o_t[:, None] * K + o_kk[None, :] + p_k = k + i_bh * T*K + o_t[:, None] * K + o_kk[None, :] + p_v = v + i_bh * T*V + o_vv[:, None] + o_t[None, :] * V + p_do = do + i_bh * T*V + o_t[:, None] * V + o_vv[None, :] + p_dq = dq + (i_bh + i_v*B*H) * T*K + o_t[:, None] * K + o_kk[None, :] p_dz = dz + (i_bh) * T + tl.arange(0, BT) + i * BT b_dq = tl.zeros([BT, BK], dtype=tl.float32) # load tensors # [BT, BK] - b_q = tl.load(p_q, boundary_check=(0, 1)) + b_q = tl.load(p_q, mask=m_q, other=0.0) b_q = (b_q * scale).to(b_q.dtype) - b_k = tl.load(p_k, boundary_check=(0, 1)) - b_do = tl.load(p_do, boundary_check=(0, 1)).to(b_q.dtype) + b_k = tl.load(p_k, mask=m_q, other=0.0) + b_do = tl.load(p_do, mask=m_v, other=0.0).to(b_q.dtype) b_dz = tl.load(p_dz, mask=(tl.arange(0, BT) + i * BT) < T) # [BV, BT] - b_v = tl.load(p_v, boundary_check=(0, 1)) + b_v = tl.load(p_v, mask=m_kv, other=0.0) # inter-chunk - b_dq += tl.dot(b_do, (b_h_1o).to(b_do.dtype), allow_tf32=False) + b_dq = tl.dot(b_do, (b_h_1o).to(b_do.dtype), b_dq, allow_tf32=False) if i_v == 0: b_dq += b_dz[:, None] * k_1o b_dq_2o = tl.dot(b_do, (b_h_2o).to(b_do.dtype), allow_tf32=False) * 0.5 @@ -185,19 +200,19 @@ def fused_chunk_based_bwd_kernel( b_ds = tl.where(m_s, b_ds, 0) * scale b_s = tl.dot(b_q, tl.trans(b_k), allow_tf32=False) b_s = tl.where(m_s, b_s, 0) - b_dq += tl.dot((b_ds * (1 + b_s)).to(b_q.dtype), b_k, allow_tf32=False) + b_dq = tl.dot((b_ds * (1 + b_s)).to(b_q.dtype), b_k, b_dq, allow_tf32=False) # store - tl.store(p_dq, b_dq.to(p_dq.dtype.element_ty), boundary_check=(0, 1)) + tl.store(p_dq, b_dq.to(p_dq.dtype.element_ty), mask=m_q) # update hidden state # [BT, BK*BK] b_k_2o = b_k[:, :, None] * b_k[:, None, :] b_k_2o = tl.reshape(b_k_2o, [BT, BK * BK]).to(b_k.dtype) # [BV, BK*BK] - b_h_2o = b_h_2o + tl.dot(b_v, b_k_2o.to(b_v.dtype), allow_tf32=False) + b_h_2o = tl.dot(b_v, b_k_2o.to(b_v.dtype), b_h_2o, allow_tf32=False) # [BV, BK] - b_h_1o = b_h_1o + tl.dot(b_v, b_k, allow_tf32=False) + b_h_1o = tl.dot(b_v, b_k, b_h_1o, allow_tf32=False) if i_v == 0: # update running statistics @@ -219,21 +234,26 @@ def fused_chunk_based_bwd_kernel( dq_2o = tl.zeros([BK * BK, 1], dtype=tl.float32) for i in range(tl.cdiv(T, BT) * BT - BT, -BT, -BT): - p_q = tl.make_block_ptr(q + i_bh * T*K, (K, T), (1, K), (i_k * BK, i), (BK, BT), (0, 1)) - p_k = tl.make_block_ptr(k + i_bh * T*K, (T, K), (K, 1), (i, i_k * BK), (BT, BK), (1, 0)) - p_v = tl.make_block_ptr(v + i_bh * T*V, (T, V), (V, 1), (i, i_v * BV), (BT, BV), (1, 0)) - p_do = tl.make_block_ptr(do + i_bh * T*V, (T, V), (V, 1), (i, i_v * BV), (BT, BV), (1, 0)) - p_dk = tl.make_block_ptr(dk + (i_bh+i_v*B*H) * T*K, (T, K), (K, 1), (i, i_k*BK), (BT, BK), (1, 0)) - p_dv = tl.make_block_ptr(dv + (i_bh+i_k*B*H) * T*V, (T, V), (V, 1), (i, i_v*BV), (BT, BV), (1, 0)) + o_t = i.to(tl.int64) + tl.arange(0, BT) + m_t = o_t < T + m_q = (o_kk[:, None] < K) & m_t[None, :] + m_k = m_t[:, None] & (o_kk[None, :] < K) + m_v = m_t[:, None] & (o_vv[None, :] < V) + p_q = q + i_bh * T*K + o_kk[:, None] + o_t[None, :] * K + p_k = k + i_bh * T*K + o_t[:, None] * K + o_kk[None, :] + p_v = v + i_bh * T*V + o_t[:, None] * V + o_vv[None, :] + p_do = do + i_bh * T*V + o_t[:, None] * V + o_vv[None, :] + p_dk = dk + (i_bh+i_v*B*H) * T*K + o_t[:, None] * K + o_kk[None, :] + p_dv = dv + (i_bh+i_k*B*H) * T*V + o_t[:, None] * V + o_vv[None, :] p_dz = dz + (i_bh) * T + tl.arange(0, BT) + i b_dk = tl.zeros([BT, BK], dtype=tl.float32) b_dv = tl.zeros([BT, BV], dtype=tl.float32) - b_q = tl.load(p_q, boundary_check=(0, 1)) - b_k = tl.load(p_k, boundary_check=(0, 1)) - b_v = tl.load(p_v, boundary_check=(0, 1)) - b_do = tl.load(p_do, boundary_check=(0, 1)).to(b_q.dtype) + b_q = tl.load(p_q, mask=m_q, other=0.0) + b_k = tl.load(p_k, mask=m_k, other=0.0) + b_v = tl.load(p_v, mask=m_v, other=0.0) + b_do = tl.load(p_do, mask=m_v, other=0.0).to(b_q.dtype) b_dz = tl.load(p_dz, mask=(tl.arange(0, BT)+i) < T) b_q = (b_q * scale).to(b_k.dtype) @@ -248,18 +268,18 @@ def fused_chunk_based_bwd_kernel( b_s2 = tl.where(m_s, b_s2, 0) b_ds *= (1+b_s) - b_dk += tl.dot(b_ds.to(b_k.dtype), tl.trans(b_q), allow_tf32=False) - b_dv += tl.dot(b_s2.to(b_do.dtype), b_do, allow_tf32=False) + b_dk = tl.dot(b_ds.to(b_k.dtype), tl.trans(b_q), b_dk, allow_tf32=False) + b_dv = tl.dot(b_s2.to(b_do.dtype), b_do, b_dv, allow_tf32=False) # inter chunk b_k_2o = b_k[:, :, None] * b_k[:, None, :] b_k_2o = tl.reshape(b_k_2o, [BT, BK * BK]).to(b_k.dtype) - b_dv += tl.dot(b_k, b_dh_1o.to(b_k.dtype), allow_tf32=False) - b_dv += tl.dot(b_k_2o, b_dh_2o.to(b_k.dtype), allow_tf32=False) + b_dv = tl.dot(b_k, b_dh_1o.to(b_k.dtype), b_dv, allow_tf32=False) + b_dv = tl.dot(b_k_2o, b_dh_2o.to(b_k.dtype), b_dv, allow_tf32=False) b_dv += b_dh_0o - b_dk += tl.dot(b_v, tl.trans(b_dh_1o).to(b_k.dtype), allow_tf32=False) + b_dk = tl.dot(b_v, tl.trans(b_dh_1o).to(b_k.dtype), b_dk, allow_tf32=False) if i_v == 0: b_dk += dq_1o @@ -275,7 +295,7 @@ def fused_chunk_based_bwd_kernel( # hidden state update b_dh_0o += tl.sum(b_do, axis=0) - b_dh_1o = b_dh_1o + tl.dot(b_q, b_do, allow_tf32=False) + b_dh_1o = tl.dot(b_q, b_do, b_dh_1o, allow_tf32=False) b_q_2o = b_q[None, :, :] * b_q[:, None, :] b_q_2o = tl.reshape(b_q_2o, [BK * BK, BT]).to(b_k.dtype) b_dh_2o = b_dh_2o + tl.dot(b_q_2o, b_do, allow_tf32=False) * 0.5 @@ -284,8 +304,8 @@ def fused_chunk_based_bwd_kernel( dq_1o += (tl.sum(b_dz[None, :] * b_q, axis=1))[None, :] dq_2o += (tl.sum(b_dz[None, :] * b_q_2o, axis=1) * 0.5)[:, None] - tl.store(p_dk, b_dk.to(p_dk.dtype.element_ty), boundary_check=(0, 1)) - tl.store(p_dv, b_dv.to(p_dv.dtype.element_ty), boundary_check=(0, 1)) + tl.store(p_dk, b_dk.to(p_dk.dtype.element_ty), mask=m_k) + tl.store(p_dv, b_dv.to(p_dv.dtype.element_ty), mask=m_v) class FusedChunkBasedFunction(torch.autograd.Function): @@ -360,16 +380,13 @@ def fused_chunk_based( v: torch.Tensor, scale: float | None = None, use_norm: bool = True, - head_first: bool = False, ): assert q.shape[-1] <= 16, 'only support feature dimension up to 16.' if scale is None: scale = q.shape[-1] ** -0.5 - if not head_first: - q, k, v = map(lambda x: x.transpose(1, 2), (q, k, v)) + q, k, v = map(lambda x: x.transpose(1, 2), (q, k, v)) o, z = FusedChunkBasedFunction.apply(q, k, v, scale) if use_norm: o = o / (z[..., None] + 1e-6) - if not head_first: - o = o.transpose(1, 2) + o = o.transpose(1, 2) return o.to(q.dtype) diff --git a/build/torch-cuda/ops/based/parallel.py b/build/torch-cuda/ops/based/parallel.py index a693d3837d02ef7ad299ad9980c0bb47b41c6808..d204832783acccc66e27d49fbbbdee00c3e1b634 100644 --- a/build/torch-cuda/ops/based/parallel.py +++ b/build/torch-cuda/ops/based/parallel.py @@ -9,6 +9,7 @@ import torch import triton import triton.language as tl +from ...ops.utils.op import unflatten_program_id from ...utils import autocast_custom_bwd, autocast_custom_fwd, input_guard # Based: An Educational and Effective Sequence Mixer @@ -34,38 +35,46 @@ def parallel_based_fwd_kernel( BV: tl.constexpr, ): # i_c: chunk index. used for sequence parallelism - i_kv, i_c, i_bh = tl.program_id(0), tl.program_id(1), tl.program_id(2) + i_kv, i_c = unflatten_program_id(tl.cdiv(K, BK) * tl.cdiv(V, BV)) + i_bh = tl.program_id(1).to(tl.int64) NV = tl.cdiv(V, BV) i_k = i_kv // (NV) i_v = i_kv % (NV) - p_q = tl.make_block_ptr(q + i_bh * T*K, (T, K), (K, 1), (i_c * BTL, i_k * BK), (BTL, BK), (1, 0)) - p_k = tl.make_block_ptr(k + i_bh * T*K, (K, T), (1, K), (i_k * BK, 0), (BK, BTS), (0, 1)) - p_v = tl.make_block_ptr(v + i_bh * T*V, (T, V), (V, 1), (0, i_v * BV), (BTS, BV), (1, 0)) + o_t = i_c * BTL + tl.arange(0, BTL) + o_kk = i_k * BK + tl.arange(0, BK) + o_vv = i_v * BV + tl.arange(0, BV) + m_t = o_t < T + m_q = m_t[:, None] & (o_kk[None, :] < K) + m_o = m_t[:, None] & (o_vv[None, :] < V) + p_q = q + i_bh * T*K + o_t[:, None] * K + o_kk[None, :] # [BQ, BD] block Q, in the shared memory throughout the whole kernel - b_q = tl.load(p_q, boundary_check=(0, 1)) + b_q = tl.load(p_q, mask=m_q, other=0.0) b_q = (b_q * scale).to(b_q.dtype) b_o = tl.zeros([BTL, BV], dtype=tl.float32) b_z = tl.zeros([BTL], dtype=tl.float32) # Q block and K block have no overlap # no need for mask, thereby saving flops - for _ in range(0, i_c * BTL, BTS): + for i_s in range(0, i_c * BTL, BTS): + o_s = i_s + tl.arange(0, BTS) + m_k = (o_kk[:, None] < K) & (o_s[None, :] < T) + m_v = (o_s[:, None] < T) & (o_vv[None, :] < V) + p_k = k + i_bh * T*K + o_kk[:, None] + o_s[None, :] * K + p_v = v + i_bh * T*V + o_s[:, None] * V + o_vv[None, :] # [BK, BTS] - b_k = tl.load(p_k, boundary_check=(0, 1)) + b_k = tl.load(p_k, mask=m_k, other=0.0) # [BTS, BV] - b_v = tl.load(p_v, boundary_check=(0, 1)) + b_v = tl.load(p_v, mask=m_v, other=0.0) # [BTL, BTS] b_s = tl.dot(b_q, (b_k), allow_tf32=False) b_s = 1 + b_s + 0.5 * b_s * b_s b_z += tl.sum(b_s, axis=1) # [BQ, BD] - b_o = b_o + tl.dot(b_s.to(b_v.dtype), b_v, allow_tf32=False) - p_k = tl.advance(p_k, (0, BTS)) - p_v = tl.advance(p_v, (BTS, 0)) + b_o = tl.dot(b_s.to(b_v.dtype), b_v, b_o, allow_tf32=False) # # rescale interchunk output tl.debug_barrier() @@ -74,14 +83,17 @@ def parallel_based_fwd_kernel( # tl.debug_barrier() o_k = tl.arange(0, BTS) - p_k = tl.make_block_ptr(k + i_bh * T*K, (K, T), (1, K), (i_k * BK, i_c * BTL), (BK, BTS), (0, 1)) - p_v = tl.make_block_ptr(v + i_bh * T*V, (T, V), (V, 1), (i_c * BTL, i_v * BV), (BTS, BV), (1, 0)) # Q block and K block have overlap. masks required - for _ in range(i_c * BTL, (i_c + 1) * BTL, BTS): + for i_s in range(i_c * BTL, (i_c + 1) * BTL, BTS): + o_s = i_s + tl.arange(0, BTS) + m_k = (o_kk[:, None] < K) & (o_s[None, :] < T) + m_v = (o_s[:, None] < T) & (o_vv[None, :] < V) + p_k = k + i_bh * T*K + o_kk[:, None] + o_s[None, :] * K + p_v = v + i_bh * T*V + o_s[:, None] * V + o_vv[None, :] # [BK, BTS] - b_k = tl.load(p_k, boundary_check=(0, 1)) + b_k = tl.load(p_k, mask=m_k, other=0.0) # [BTS, BV] - b_v = tl.load(p_v, boundary_check=(0, 1)) + b_v = tl.load(p_v, mask=m_v, other=0.0) # [BTL, BTS] m_s = o_q[:, None] >= o_k[None, :] b_s = tl.dot(b_q, b_k, allow_tf32=False) @@ -89,15 +101,13 @@ def parallel_based_fwd_kernel( b_s = tl.where(m_s, b_s, 0) b_z += tl.sum(b_s, axis=1) # [BTL, BV] - b_o += tl.dot(b_s.to(b_q.dtype), b_v, allow_tf32=False) + b_o = tl.dot(b_s.to(b_q.dtype), b_v, b_o, allow_tf32=False) - p_k = tl.advance(p_k, (0, BTS)) - p_v = tl.advance(p_v, (BTS, 0)) o_k += BTS - p_o = tl.make_block_ptr(o + (i_bh + B * H * i_k) * T*V, (T, V), (V, 1), (i_c*BTL, i_v*BV), (BTL, BV), (1, 0)) + p_o = o + (i_bh + B * H * i_k) * T*V + o_t[:, None] * V + o_vv[None, :] p_z = z + (i_bh + B * H * i_k) * T + i_c * BTL + tl.arange(0, BTL) - tl.store(p_o, b_o.to(p_o.dtype.element_ty), boundary_check=(0, 1)) + tl.store(p_o, b_o.to(p_o.dtype.element_ty), mask=m_o) tl.store(p_z, b_z.to(p_z.dtype.element_ty), mask=((i_c * BTL + tl.arange(0, BTL)) < T)) @@ -124,23 +134,32 @@ def _parallel_based_bwd_dq( K: tl.constexpr, V: tl.constexpr, ): - p_do = tl.make_block_ptr(do + i_bh * T*V, (T, V), (V, 1), (i_c * BTL, i_v * BV), (BTL, BV), (1, 0)) - p_q = tl.make_block_ptr(q + (i_bh) * T*K, (T, K), (K, 1), (i_c*BTL, i_k*BK), (BTL, BK), (1, 0)) - b_q = tl.load(p_q, boundary_check=(0, 1)) + o_t = i_c * BTL + tl.arange(0, BTL) + o_kk = i_k * BK + tl.arange(0, BK) + o_vv = i_v * BV + tl.arange(0, BV) + m_t = o_t < T + m_q = m_t[:, None] & (o_kk[None, :] < K) + m_o = m_t[:, None] & (o_vv[None, :] < V) + p_do = do + i_bh * T*V + o_t[:, None] * V + o_vv[None, :] + p_q = q + (i_bh) * T*K + o_t[:, None] * K + o_kk[None, :] + b_q = tl.load(p_q, mask=m_q, other=0.0) b_q = (b_q * scale).to(b_q.dtype) - b_do = tl.load(p_do, boundary_check=(0, 1)).to(b_q.dtype) + b_do = tl.load(p_do, mask=m_o, other=0.0).to(b_q.dtype) b_dq = tl.zeros([BTL, BK], dtype=tl.float32) - p_k = tl.make_block_ptr(k + i_bh * T*K, (T, K), (K, 1), (0, i_k * BK), (BTS, BK), (1, 0)) - p_v = tl.make_block_ptr(v + i_bh * T*V, (V, T), (1, V), (i_v * BV, 0), (BV, BTS), (0, 1)) p_dz = dz + i_bh * T + i_c * BTL + tl.arange(0, BTL) b_dz = tl.load(p_dz, mask=(i_c * BTL + tl.arange(0, BTL)) < T) - for _ in range(0, i_c * BTL, BTS): + for i_s in range(0, i_c * BTL, BTS): + o_s = i_s + tl.arange(0, BTS) + m_k = (o_s[:, None] < T) & (o_kk[None, :] < K) + m_v = (o_vv[:, None] < V) & (o_s[None, :] < T) + p_k = k + i_bh * T*K + o_s[:, None] * K + o_kk[None, :] + p_v = v + i_bh * T*V + o_vv[:, None] + o_s[None, :] * V # [BTS, BK] - b_k = tl.load(p_k, boundary_check=(0, 1)) + b_k = tl.load(p_k, mask=m_k, other=0.0) # [BV, BTS] - b_v = tl.load(p_v, boundary_check=(0, 1)) + b_v = tl.load(p_v, mask=m_v, other=0.0) # [BTL, BTS] b_ds = tl.dot(b_do, b_v, allow_tf32=False) if i_v == 0: @@ -149,21 +168,22 @@ def _parallel_based_bwd_dq( b_ds = b_ds b_s = tl.dot(b_q, tl.trans(b_k), allow_tf32=False) # [BQ, BD] - b_dq += tl.dot((b_ds * (1 + b_s)).to(b_v.dtype), b_k, allow_tf32=False) - p_k = tl.advance(p_k, (BTS, 0)) - p_v = tl.advance(p_v, (0, BTS)) + b_dq = tl.dot((b_ds * (1 + b_s)).to(b_v.dtype), b_k, b_dq, allow_tf32=False) b_dq *= scale o_q = tl.arange(0, BTL) o_k = tl.arange(0, BTS) - p_k = tl.make_block_ptr(k + i_bh * T*K, (T, K), (K, 1), (i_c * BTL, i_k * BK), (BTS, BK), (1, 0)) - p_v = tl.make_block_ptr(v + i_bh * T*V, (V, T), (1, V), (i_v * BV, i_c * BTL), (BV, BTS), (0, 1)) # Q block and K block have overlap. masks required - for _ in range(i_c * BTL, (i_c + 1) * BTL, BTS): + for i_s in range(i_c * BTL, (i_c + 1) * BTL, BTS): + o_s = i_s + tl.arange(0, BTS) + m_k = (o_s[:, None] < T) & (o_kk[None, :] < K) + m_v = (o_vv[:, None] < V) & (o_s[None, :] < T) + p_k = k + i_bh * T*K + o_s[:, None] * K + o_kk[None, :] + p_v = v + i_bh * T*V + o_vv[:, None] + o_s[None, :] * V # [BTS, BK] - b_k = tl.load(p_k, boundary_check=(0, 1)) + b_k = tl.load(p_k, mask=m_k, other=0.0) # [BV, BTS] - b_v = tl.load(p_v, boundary_check=(0, 1)) + b_v = tl.load(p_v, mask=m_v, other=0.0) # [BTL, BTS] m_s = o_q[:, None] >= o_k[None, :] b_ds = tl.dot(b_do, b_v, allow_tf32=False) @@ -175,12 +195,10 @@ def _parallel_based_bwd_dq( b_s = tl.dot(b_q, tl.trans(b_k), allow_tf32=False) b_s = tl.where(m_s, b_s, 0) # [BTL, BK] - b_dq += tl.dot((b_ds + b_ds * b_s).to(b_k.dtype), b_k, allow_tf32=False) - p_k = tl.advance(p_k, (BTS, 0)) - p_v = tl.advance(p_v, (0, BTS)) + b_dq = tl.dot((b_ds + b_ds * b_s).to(b_k.dtype), b_k, b_dq, allow_tf32=False) o_k += BTS - p_dq = tl.make_block_ptr(dq + (i_bh + B * H * i_v) * T*K, (T, K), (K, 1), (i_c*BTL, i_k*BK), (BTL, BK), (1, 0)) - tl.store(p_dq, b_dq.to(p_dq.dtype.element_ty), boundary_check=(0, 1)) + p_dq = dq + (i_bh + B * H * i_v) * T*K + o_t[:, None] * K + o_kk[None, :] + tl.store(p_dq, b_dq.to(p_dq.dtype.element_ty), mask=m_q) return @@ -209,36 +227,48 @@ def _parallel_based_bwd_dkv( V: tl.constexpr, ): # compute dk dv - p_k = tl.make_block_ptr(k + i_bh * T*K, (T, K), (K, 1), (i_c * BTL, i_k * BK), (BTL, BK), (1, 0)) - p_v = tl.make_block_ptr(v + i_bh * T*V, (T, V), (V, 1), (i_c * BTL, i_v * BV), (BTL, BV), (1, 0)) - b_k, b_v = tl.load(p_k, boundary_check=(0, 1)), tl.load(p_v, boundary_check=(0, 1)) + o_t = i_c * BTL + tl.arange(0, BTL) + o_kk = i_k * BK + tl.arange(0, BK) + o_vv = i_v * BV + tl.arange(0, BV) + m_t = o_t < T + m_k = m_t[:, None] & (o_kk[None, :] < K) + m_v = m_t[:, None] & (o_vv[None, :] < V) + p_k = k + i_bh * T*K + o_t[:, None] * K + o_kk[None, :] + p_v = v + i_bh * T*V + o_t[:, None] * V + o_vv[None, :] + b_k, b_v = tl.load(p_k, mask=m_k, other=0.0), tl.load(p_v, mask=m_v, other=0.0) b_dk, b_dv = tl.zeros([BTL, BK], dtype=tl.float32), tl.zeros([BTL, BV], dtype=tl.float32) for i in range((tl.cdiv(T, BTS) * BTS)-BTS, (i_c + 1) * BTL - BTS, -BTS): - p_q = tl.make_block_ptr(q + i_bh * T*K, (K, T), (1, K), (i_k * BK, i), (BK, BTS), (0, 1)) - p_do = tl.make_block_ptr(do + i_bh * T*V, (V, T), (1, V), (i_v * BV, i), (BV, BTS), (0, 1)) + o_s = i + tl.arange(0, BTS) + m_q = (o_kk[:, None] < K) & (o_s[None, :] < T) + m_o = (o_vv[:, None] < V) & (o_s[None, :] < T) + p_q = q + i_bh * T*K + o_kk[:, None] + o_s[None, :] * K + p_do = do + i_bh * T*V + o_vv[:, None] + o_s[None, :] * V p_dz = dz + i_bh * T + i + tl.arange(0, BTS) - b_q = tl.load(p_q, boundary_check=(0, 1)) # [BK, BTS] - b_do = tl.load(p_do, boundary_check=(0, 1)).to(b_q.dtype) # [BV, BTS] + b_q = tl.load(p_q, mask=m_q, other=0.0) # [BK, BTS] + b_do = tl.load(p_do, mask=m_o, other=0.0).to(b_q.dtype) # [BV, BTS] b_dz = tl.load(p_dz, mask=(i + tl.arange(0, BTS)) < T) b_s = tl.dot(b_k.to(b_q.dtype), b_q, allow_tf32=False) * scale # [BTL, BTS] b_s2 = 1 + b_s + 0.5 * b_s * b_s - b_dv += tl.dot(b_s2.to(b_q.dtype), tl.trans(b_do), allow_tf32=False) + b_dv = tl.dot(b_s2.to(b_q.dtype), tl.trans(b_do), b_dv, allow_tf32=False) b_ds = tl.dot(b_v, b_do, allow_tf32=False) * scale if i_v == 0: b_ds += b_dz[None, :] * scale else: b_ds = b_ds - b_dk += tl.dot((b_ds + b_ds * b_s).to(b_q.dtype), tl.trans(b_q), allow_tf32=False) + b_dk = tl.dot((b_ds + b_ds * b_s).to(b_q.dtype), tl.trans(b_q), b_dk, allow_tf32=False) tl.debug_barrier() o_q, o_k = tl.arange(0, BTS), tl.arange(0, BTL) for i in range(i_c*BTL, (i_c+1)*BTL, BTS): - p_q = tl.make_block_ptr(q + i_bh * T*K, (K, T), (1, K), (i_k * BK, i), (BK, BTS), (0, 1)) - p_do = tl.make_block_ptr(do + i_bh * T*V, (V, T), (1, V), (i_v * BV, i), (BV, BTS), (0, 1)) + o_s = i + tl.arange(0, BTS) + m_q = (o_kk[:, None] < K) & (o_s[None, :] < T) + m_o = (o_vv[:, None] < V) & (o_s[None, :] < T) + p_q = q + i_bh * T*K + o_kk[:, None] + o_s[None, :] * K + p_do = do + i_bh * T*V + o_vv[:, None] + o_s[None, :] * V p_dz = dz + i_bh * T + i + tl.arange(0, BTS) - b_q = tl.load(p_q, boundary_check=(0, 1)) # [BD, BQ] - b_do = tl.load(p_do, boundary_check=(0, 1)).to(b_q.dtype) + b_q = tl.load(p_q, mask=m_q, other=0.0) # [BD, BQ] + b_do = tl.load(p_do, mask=m_o, other=0.0).to(b_q.dtype) b_dz = tl.load(p_dz, mask=(i + tl.arange(0, BTS)) < T) # [BK, BQ] m_s = o_k[:, None] <= o_q[None, :] @@ -254,14 +284,14 @@ def _parallel_based_bwd_dkv( b_ds = b_ds b_ds = tl.where(m_s, b_ds, 0) * scale # [BK, BD] - b_dv += tl.dot(b_s2.to(b_q.dtype), tl.trans(b_do), allow_tf32=False) - b_dk += tl.dot((b_ds + b_ds * b_s).to(b_q.dtype), tl.trans(b_q), allow_tf32=False) + b_dv = tl.dot(b_s2.to(b_q.dtype), tl.trans(b_do), b_dv, allow_tf32=False) + b_dk = tl.dot((b_ds + b_ds * b_s).to(b_q.dtype), tl.trans(b_q), b_dk, allow_tf32=False) o_q += BTS - p_dk = tl.make_block_ptr(dk + (i_bh + B * H * i_v) * T*K, (T, K), (K, 1), (i_c*BTL, i_k*BK), (BTL, BK), (1, 0)) - p_dv = tl.make_block_ptr(dv + (i_bh + B * H * i_k) * T*V, (T, V), (V, 1), (i_c*BTL, i_v*BV), (BTL, BV), (1, 0)) - tl.store(p_dk, b_dk.to(p_dk.dtype.element_ty), boundary_check=(0, 1)) - tl.store(p_dv, b_dv.to(p_dv.dtype.element_ty), boundary_check=(0, 1)) + p_dk = dk + (i_bh + B * H * i_v) * T*K + o_t[:, None] * K + o_kk[None, :] + p_dv = dv + (i_bh + B * H * i_k) * T*V + o_t[:, None] * V + o_vv[None, :] + tl.store(p_dk, b_dk.to(p_dk.dtype.element_ty), mask=m_k) + tl.store(p_dv, b_dv.to(p_dv.dtype.element_ty), mask=m_v) return @@ -286,7 +316,8 @@ def parallel_based_bwd_kernel( BK: tl.constexpr, BV: tl.constexpr, ): - i_kv, i_c, i_bh = tl.program_id(0), tl.program_id(1), tl.program_id(2) + i_kv, i_c = unflatten_program_id(tl.cdiv(K, BK) * tl.cdiv(V, BV)) + i_bh = tl.program_id(1).to(tl.int64) NV = tl.cdiv(V, BV) i_k = i_kv // (NV) i_v = i_kv % NV @@ -319,7 +350,7 @@ class ParallelBasedFunction(torch.autograd.Function): num_warps = 4 NK = triton.cdiv(K, BK) NV = triton.cdiv(V, BV) - grid = (NK * NV, triton.cdiv(T, BTL), B * H) + grid = (NK * NV * triton.cdiv(T, BTL), B * H) assert NK == 1, "will encounter some synchronization issue if not." @@ -359,7 +390,7 @@ class ParallelBasedFunction(torch.autograd.Function): num_warps = 4 NK = triton.cdiv(K, BK) NV = triton.cdiv(V, BV) - grid = (NK * NV, triton.cdiv(T, BTL), B * H) + grid = (NK * NV * triton.cdiv(T, BTL), B * H) assert NK == 1, "will encounter some synchronization issue if not" @@ -395,16 +426,13 @@ def parallel_based( v: torch.Tensor, scale: float | None = None, use_norm: bool = True, - head_first: bool = False, ): assert q.shape[-1] <= 128, "only support feature dim up to 128" if scale is None: scale = q.shape[-1] ** -0.5 - if not head_first: - q, k, v = map(lambda x: x.transpose(1, 2), (q, k, v)) + q, k, v = map(lambda x: x.transpose(1, 2), (q, k, v)) o, z = triton_parallel_based(q, k, v, scale) if use_norm: o = o / (z[..., None] + 1e-6) - if not head_first: - o = o.transpose(1, 2) + o = o.transpose(1, 2) return o.to(q.dtype) diff --git a/build/torch-cuda/ops/comba/chunk.py b/build/torch-cuda/ops/comba/chunk.py index e2497babafeb13b7a751e3462b794a034356f78f..31641a800a115c606feb19e9b6aac37fc26fc002 100644 --- a/build/torch-cuda/ops/comba/chunk.py +++ b/build/torch-cuda/ops/comba/chunk.py @@ -7,7 +7,7 @@ import torch -from ...modules.l2norm import l2norm_bwd, l2norm_fwd +from ...modules.norm.l2norm import l2norm_bwd, l2norm_fwd from ...ops.comba.utils import chunk_comba_cumsum_scalar_bwd, chunk_comba_cumsum_scalar_fwd from ...ops.comba.wy_fast import chunk_scaled_dot_comba_pkt_fwd, prepare_wy_repr_bwd, recompute_w_u_fwd from ...ops.common.chunk_delta_h import chunk_gated_delta_rule_bwd_dhu, chunk_gated_delta_rule_fwd_h @@ -312,22 +312,23 @@ def chunk_comba( k (torch.Tensor): keys of shape `[B, T, H, K]`. v (torch.Tensor): - values of shape `[B, T, H, V]`. + values of shape `[B, T, HV, V]`. + GVA (Grouped Value Attention) is applied if `HV > H`, where `HV` must be divisible by `H`. p (torch.Tensor): auxiliary keys of shape `[B, T, H, K]`. g (torch.Tensor): - (forget) gating tensor (in log space!) of shape `[B, T, H]`. + (forget) gating tensor (in log space!) of shape `[B, T, HV]`. beta (torch.Tensor): - betas of shape `[B, T, H]`. + betas of shape `[B, T, HV]`. scale (Optional[int]): Scale factor for the RetNet attention scores. If not provided, it will default to `1 / sqrt(K)`. Default: `None`. initial_state (Optional[torch.Tensor]): - Initial state of shape `[N, H, K, V]` for `N` input sequences. + Initial state of shape `[N, HV, K, V]` for `N` input sequences. For equal-length input sequences, `N` equals the batch size `B`. Default: `None`. output_final_state (Optional[bool]): - Whether to output the final state of shape `[N, H, K, V]`. Default: `False`. + Whether to output the final state of shape `[N, HV, K, V]`. Default: `False`. use_qk_l2norm_in_kernel (bool): Whether to apply L2norm to the q/k tensor internally. Default: `False`. cu_seqlens (torch.LongTensor): @@ -336,15 +337,15 @@ def chunk_comba( Returns: o (torch.Tensor): - Outputs of shape `[B, T, H, V]`. + Outputs of shape `[B, T, HV, V]`. final_state (torch.Tensor): - Final state of shape `[N, H, K, V]` if `output_final_state=True` else `None`. + Final state of shape `[N, HV, K, V]` if `output_final_state=True` else `None`. Examples:: >>> import torch >>> import torch.nn.functional as F >>> from einops import rearrange - >>> from ...ops.comba import chunk_comba + >>> from fla.ops.comba import chunk_comba # inputs with equal lengths >>> B, T, H, K, V = 4, 2048, 4, 512, 512 >>> q = torch.randn(B, T, H, K, dtype=torch.bfloat16, device='cuda') @@ -373,6 +374,17 @@ def chunk_comba( """ if p is None: p = k + if q.shape[2] != k.shape[2] or q.shape[2] != p.shape[2]: + raise ValueError( + f"q, k and p must have the same number of heads, " + f"but got q.shape[2]={q.shape[2]}, k.shape[2]={k.shape[2]} and p.shape[2]={p.shape[2]}" + ) + H, HV = q.shape[2], v.shape[2] + if HV % H != 0: + raise ValueError( + f"For GVA, num_v_heads (HV={HV}) must be evenly divisible by " + f"num_heads (H={H}), but got HV % H = {HV % H}" + ) if cu_seqlens is not None: if q.shape[0] != 1: raise ValueError( diff --git a/build/torch-cuda/ops/comba/fused_recurrent.py b/build/torch-cuda/ops/comba/fused_recurrent.py index bd773d7bf6d5b83c949898c1a7445f476aab3e52..d86745f4f4f232fe98f83b92a6744f42fd004a7f 100644 --- a/build/torch-cuda/ops/comba/fused_recurrent.py +++ b/build/torch-cuda/ops/comba/fused_recurrent.py @@ -45,7 +45,9 @@ def fused_recurrent_comba_fwd_kernel( USE_QK_L2NORM_IN_KERNEL: tl.constexpr, IS_VARLEN: tl.constexpr, ): - i_k, i_v, i_nh = tl.program_id(0), tl.program_id(1), tl.program_id(2) + pid = tl.program_id(0).to(tl.int64) + NK, NV = tl.cdiv(K, BK), tl.cdiv(V, BV) + i_k, i_v, i_nh = pid % NK, (pid // NK) % NV, (pid // (NK * NV)).to(tl.int64) i_n, i_hv = i_nh // HV, i_nh % HV i_h = i_hv // (HV // H) if IS_VARLEN: @@ -146,7 +148,7 @@ def fused_recurrent_comba_fwd( else: final_state = None - grid = (NK, NV, N * HV) + grid = (NK * NV * N * HV,) fused_recurrent_comba_fwd_kernel[grid]( q=q, k=k, @@ -274,7 +276,7 @@ def fused_recurrent_comba( >>> import torch >>> import torch.nn.functional as F >>> from einops import rearrange - >>> from ...ops.comba import fused_recurrent_comba + >>> from fla.ops.comba import fused_recurrent_comba # inputs with equal lengths >>> B, T, H, HV, K, V = 4, 2048, 4, 8, 512, 512 >>> q = torch.randn(B, T, H, K, device='cuda') diff --git a/build/torch-cuda/ops/comba/utils.py b/build/torch-cuda/ops/comba/utils.py index 2126de39cb9e272df0ed8891b2c3c34ef5900368..ab28acdb02b05294887c963a060dce983744ff33 100644 --- a/build/torch-cuda/ops/comba/utils.py +++ b/build/torch-cuda/ops/comba/utils.py @@ -40,26 +40,28 @@ def chunk_comba_cumsum_scalar_fwd_kernel( HAS_SCALE: tl.constexpr, IS_VARLEN: tl.constexpr, ): - i_t, i_bh = tl.program_id(0), tl.program_id(1) + i_t, i_bh = tl.program_id(0).to(tl.int64), tl.program_id(1).to(tl.int64) i_b, i_h = i_bh // H, i_bh % H 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.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) + i_n, i_t = tl.load(chunk_indices + i_t * 2).to(tl.int32), tl.load(chunk_indices + i_t * 2 + 1).to(tl.int64) + bos, eos = tl.load(cu_seqlens + i_n).to(tl.int64), tl.load(cu_seqlens + i_n + 1).to(tl.int64) T = eos - bos else: bos, eos = i_b * T, i_b * T + T - p_g = tl.make_block_ptr(g + bos*H + i_h, (T,), (H,), (i_t * BT,), (BT,), (0,)) - p_g0 = tl.make_block_ptr(g0 + bos*H + i_h, (T,), (H,), (i_t * BT,), (BT,), (0,)) - p_g1 = tl.make_block_ptr(g1 + bos*H + i_h, (T,), (H,), (i_t * BT,), (BT,), (0,)) + o_t = i_t * BT + tl.arange(0, BT) + m_t = o_t < T + p_g = g + bos*H + i_h + o_t * H + p_g0 = g0 + bos*H + i_h + o_t * H + p_g1 = g1 + bos*H + i_h + o_t * H # [BT] - b_g = tl.load(p_g, boundary_check=(0,)).to(tl.float32) + b_g = tl.load(p_g, mask=m_t, other=0.0).to(tl.float32) if HAS_SCALE: b_g = b_g * scale b_g1 = tl.cumsum(b_g, axis=0) b_g0 = b_g1 - b_g - tl.store(p_g0, b_g0.to(p_g0.dtype.element_ty), boundary_check=(0,)) - tl.store(p_g1, b_g1.to(p_g1.dtype.element_ty), boundary_check=(0,)) + tl.store(p_g0, b_g0.to(p_g0.dtype.element_ty), mask=m_t) + tl.store(p_g1, b_g1.to(p_g1.dtype.element_ty), mask=m_t) def chunk_comba_cumsum_scalar_fwd( @@ -116,17 +118,19 @@ def chunk_comba_cumsum_scalar_bwd_kernel( BT: tl.constexpr, IS_VARLEN: tl.constexpr, ): - i_t, i_bh = tl.program_id(0), tl.program_id(1) + i_t, i_bh = tl.program_id(0).to(tl.int64), tl.program_id(1).to(tl.int64) i_b, i_h = i_bh // H, i_bh % H 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.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) + i_n, i_t = tl.load(chunk_indices + i_t * 2).to(tl.int32), tl.load(chunk_indices + i_t * 2 + 1).to(tl.int64) + bos, eos = tl.load(cu_seqlens + i_n).to(tl.int64), tl.load(cu_seqlens + i_n + 1).to(tl.int64) T = eos - bos else: bos, eos = i_b * T, i_b * T + T - p_dg0 = tl.make_block_ptr(dg0 + bos*H + i_h, (T,), (H,), (i_t * BT,), (BT,), (0,)) - p_dgr = tl.make_block_ptr(dgr + bos*H + i_h, (T,), (H,), (i_t * BT,), (BT,), (0,)) + o_t = i_t * BT + tl.arange(0, BT) + m_t = o_t < T + p_dg0 = dg0 + bos*H + i_h + o_t * H + p_dgr = dgr + bos*H + i_h + o_t * H # [BT] """ b_dg: 1,2,3,4 @@ -135,11 +139,11 @@ def chunk_comba_cumsum_scalar_bwd_kernel( b_dz: 6 b_dgr: 6,5,3,0 """ - b_dg0 = tl.load(p_dg0, boundary_check=(0,)).to(tl.float32) + b_dg0 = tl.load(p_dg0, mask=m_t, other=0.0).to(tl.float32) b_temp = tl.cumsum(b_dg0, axis=0) b_dz = tl.sum(b_dg0, axis=0) b_dgr = -b_temp + b_dz[None] - tl.store(p_dgr, b_dgr.to(p_dgr.dtype.element_ty), boundary_check=(0,)) + tl.store(p_dgr, b_dgr.to(p_dgr.dtype.element_ty), mask=m_t) def chunk_comba_cumsum_scalar_bwd( diff --git a/build/torch-cuda/ops/comba/wy_fast.py b/build/torch-cuda/ops/comba/wy_fast.py index 431dbb7cfcd4fc9c36e78bc0f460415ada743638..516da67b4eed263c9a25e17cd3774ba0a11606ad 100644 --- a/build/torch-cuda/ops/comba/wy_fast.py +++ b/build/torch-cuda/ops/comba/wy_fast.py @@ -25,7 +25,7 @@ from ...utils import autotune_cache_kwargs, check_shared_mem for num_warps in [2, 4, 8] for num_stages in [2, 3, 4] ], - key=['H', 'K', 'BT', 'IS_VARLEN', 'USE_G'], + key=['H', 'HV', 'K', 'BT', 'IS_VARLEN', 'USE_G'], **autotune_cache_kwargs, ) @triton.jit(do_not_specialize=['T']) @@ -40,46 +40,50 @@ def chunk_scaled_dot_comba_pkt_fwd_kernel( chunk_indices, T, H: tl.constexpr, + HV: tl.constexpr, K: tl.constexpr, BT: tl.constexpr, BK: tl.constexpr, IS_VARLEN: tl.constexpr, USE_G: tl.constexpr, ): - i_t, i_bh = tl.program_id(0), tl.program_id(1) - i_b, i_h = i_bh // H, i_bh % H + i_t, i_bh = tl.program_id(0).to(tl.int64), tl.program_id(1).to(tl.int64) + i_b, i_h = i_bh // HV, i_bh % HV 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.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) + i_n, i_t = tl.load(chunk_indices + i_t * 2).to(tl.int32), tl.load(chunk_indices + i_t * 2 + 1).to(tl.int64) + bos, eos = tl.load(cu_seqlens + i_n).to(tl.int64), tl.load(cu_seqlens + i_n + 1).to(tl.int64) T = eos - bos else: bos, eos = i_b * T, i_b * T + T o_t = i_t * BT + tl.arange(0, BT) m_t = o_t < T - p_beta = tl.make_block_ptr(beta + bos*H + i_h, (T,), (H,), (i_t * BT,), (BT,), (0,)) - b_beta = tl.load(p_beta, boundary_check=(0,)) + p_beta = beta + bos*HV + i_h + o_t * HV + b_beta = tl.load(p_beta, mask=m_t, other=0.0) b_A = tl.zeros([BT, BT], dtype=tl.float32) for i_k in range(tl.cdiv(K, BK)): - p_k = tl.make_block_ptr(k + (bos*H + i_h) * K, (T, K), (H*K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) - p_p = tl.make_block_ptr(p + (bos*H + i_h) * K, (T, K), (H*K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) - b_k = tl.load(p_k, boundary_check=(0, 1)) - b_p = tl.load(p_p, boundary_check=(0, 1)) + o_k = i_k * BK + tl.arange(0, BK) + m_k = m_t[:, None] & (o_k[None, :] < K) + p_k = k + (bos*H + i_h // (HV // H)) * K + o_t[:, None] * (H*K) + o_k[None, :] + p_p = p + (bos*H + i_h // (HV // H)) * K + o_t[:, None] * (H*K) + o_k[None, :] + b_k = tl.load(p_k, mask=m_k, other=0.0) + b_p = tl.load(p_p, mask=m_k, other=0.0) b_pb = b_p * b_beta[:, None] - b_A += tl.dot(b_pb.to(b_k.dtype), tl.trans(b_k)) + b_A = tl.dot(b_pb.to(b_k.dtype), tl.trans(b_k), b_A) if USE_G: - p_g0 = tl.make_block_ptr(g0 + bos*H + i_h, (T,), (H,), (i_t * BT,), (BT,), (0,)) - p_g = tl.make_block_ptr(g + bos*H + i_h, (T,), (H,), (i_t * BT,), (BT,), (0,)) - b_g0 = tl.load(p_g0, boundary_check=(0,)) - b_g = tl.load(p_g, boundary_check=(0,)) + p_g0 = g0 + bos*HV + i_h + o_t * HV + p_g = g + bos*HV + i_h + o_t * HV + b_g0 = tl.load(p_g0, mask=m_t, other=0.0) + b_g = tl.load(p_g, mask=m_t, other=0.0) b_A = b_A * exp2(b_g0[:, None] - b_g[None, :]) m_A = (o_t[:, None] > o_t[None, :]) & (m_t[:, None] & m_t) b_A = tl.where(m_A, b_A, 0) - p_A = tl.make_block_ptr(A + (bos*H + i_h) * BT, (T, BT), (BT*H, 1), (i_t * BT, 0), (BT, BT), (1, 0)) - tl.store(p_A, b_A.to(p_A.dtype.element_ty), boundary_check=(0, 1)) + o_A = tl.arange(0, BT) + p_A = A + (bos*HV + i_h) * BT + o_t[:, None] * (BT*HV) + o_A[None, :] + tl.store(p_A, b_A.to(p_A.dtype.element_ty), mask=m_t[:, None] & (o_A[None, :] < BT)) def chunk_scaled_dot_comba_pkt_fwd( @@ -102,12 +106,12 @@ def chunk_scaled_dot_comba_pkt_fwd( p (torch.Tensor): The auxiliary key tensor of shape `[B, T, H, K]`. beta (torch.Tensor): - The beta tensor of shape `[B, T, H]`. + The beta tensor of shape `[B, T, HV]`. g0 (torch.Tensor): - The cumulative sum minus the original one of the gate tensor of shape `[B, T, H]`. + The cumulative sum minus the original one of the gate tensor of shape `[B, T, HV]`. Default: None g (torch.Tensor): - The cumulative sum of the gate tensor of shape `[B, T, H]`. + The cumulative sum of the gate tensor of shape `[B, T, HV]`. Default: None cu_seqlens (torch.LongTensor): The cumulative sequence lengths of the input tensor. @@ -118,15 +122,15 @@ def chunk_scaled_dot_comba_pkt_fwd( The dtype of the output tensor. Default: `torch.float32` Returns: - beta * K * K^T of shape `[B, T, H, BT]` where `BT` is the chunk size. + beta * K * K^T of shape `[B, T, HV, BT]` where `BT` is the chunk size. """ - B, T, H, K = k.shape + B, T, H, K, HV = *k.shape, beta.shape[2] BT = chunk_size if chunk_indices is None and cu_seqlens is not None: chunk_indices = prepare_chunk_indices(cu_seqlens, BT) NT = triton.cdiv(T, BT) if cu_seqlens is None else len(chunk_indices) - A = torch.empty(B, T, H, BT, device=k.device, dtype=output_dtype) - chunk_scaled_dot_comba_pkt_fwd_kernel[(NT, B * H)]( + A = torch.empty(B, T, HV, BT, device=k.device, dtype=output_dtype) + chunk_scaled_dot_comba_pkt_fwd_kernel[(NT, B * HV)]( k=k, p=p, beta=beta, @@ -137,6 +141,7 @@ def chunk_scaled_dot_comba_pkt_fwd( chunk_indices=chunk_indices, T=T, H=H, + HV=HV, K=K, BT=BT, ) @@ -152,7 +157,7 @@ def chunk_scaled_dot_comba_pkt_fwd( for num_warps in [2, 4] for num_stages in [2, 3, 4] ], - key=['H', 'K', 'V', 'BT', 'BK', 'BV', 'IS_VARLEN'], + key=['H', 'HV', 'K', 'V', 'BT', 'BK', 'BV', 'IS_VARLEN'], **autotune_cache_kwargs, ) @triton.jit(do_not_specialize=['T']) @@ -176,6 +181,7 @@ def prepare_wy_repr_bwd_kernel( chunk_indices, T, H: tl.constexpr, + HV: tl.constexpr, K: tl.constexpr, V: tl.constexpr, BT: tl.constexpr, @@ -183,59 +189,65 @@ def prepare_wy_repr_bwd_kernel( BV: tl.constexpr, IS_VARLEN: tl.constexpr, ): - i_t, i_bh = tl.program_id(0), tl.program_id(1) - i_b, i_h = i_bh // H, i_bh % H + i_t, i_bh = tl.program_id(0).to(tl.int64), tl.program_id(1).to(tl.int64) + i_b, i_h = i_bh // HV, i_bh % HV 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.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) + i_n, i_t = tl.load(chunk_indices + i_t * 2).to(tl.int32), tl.load(chunk_indices + i_t * 2 + 1).to(tl.int64) + bos, eos = tl.load(cu_seqlens + i_n).to(tl.int64), tl.load(cu_seqlens + i_n + 1).to(tl.int64) T = eos - bos else: bos, eos = i_b * T, i_b * T + T - p_beta = tl.make_block_ptr(beta + (bos*H + i_h), (T,), (H,), (i_t * BT,), (BT,), (0,)) - p_g0 = tl.make_block_ptr(g0 + (bos*H + i_h), (T,), (H,), (i_t * BT,), (BT,), (0,)) - p_g = tl.make_block_ptr(g + (bos*H + i_h), (T,), (H,), (i_t * BT,), (BT,), (0,)) - p_A = tl.make_block_ptr(A + (bos*H + i_h) * BT, (BT, T), (1, H*BT), (0, i_t * BT), (BT, BT), (0, 1)) - - b_A = tl.load(p_A, boundary_check=(0, 1)) - b_beta = tl.load(p_beta, boundary_check=(0,)) - b_g0 = tl.load(p_g0, boundary_check=(0,)) + o_t = i_t * BT + tl.arange(0, BT) + o_A = tl.arange(0, BT) + m_t = o_t < T + m_AT = (o_A[:, None] < BT) & m_t[None, :] + p_beta = beta + (bos*HV + i_h) + o_t * HV + p_g0 = g0 + (bos*HV + i_h) + o_t * HV + p_g = g + (bos*HV + i_h) + o_t * HV + p_A = A + (bos*HV + i_h) * BT + o_A[:, None] + o_t[None, :] * (HV*BT) + + b_A = tl.load(p_A, mask=m_AT, other=0.0) + b_beta = tl.load(p_beta, mask=m_t, other=0.0) + b_g0 = tl.load(p_g0, mask=m_t, other=0.0) b_g0_exp = exp2(b_g0) - b_g = tl.load(p_g, boundary_check=(0,)) + b_g = tl.load(p_g, mask=m_t, other=0.0) b_dbeta = tl.zeros([BT], dtype=tl.float32) b_dA = tl.zeros([BT, BT], dtype=tl.float32) b_dg0 = tl.zeros([BT], dtype=tl.float32) for i_k in range(tl.cdiv(K, BK)): - p_p = tl.make_block_ptr(p + (bos*H + i_h) * K, (T, K), (H*K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) - p_dp = tl.make_block_ptr(dp + (bos*H + i_h) * K, (T, K), (H*K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) - p_dw = tl.make_block_ptr(dw + (bos*H + i_h) * K, (T, K), (H*K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) - b_p = tl.load(p_p, boundary_check=(0, 1)) + o_k = i_k * BK + tl.arange(0, BK) + m_k = m_t[:, None] & (o_k[None, :] < K) + p_p = p + (bos*H + i_h // (HV // H)) * K + o_t[:, None] * (H*K) + o_k[None, :] + p_dp = dp + (bos*HV + i_h) * K + o_t[:, None] * (HV*K) + o_k[None, :] + p_dw = dw + (bos*HV + i_h) * K + o_t[:, None] * (HV*K) + o_k[None, :] + b_p = tl.load(p_p, mask=m_k, other=0.0) b_p_beta_g0 = (b_p * b_beta[:, None] * b_g0_exp[:, None]).to(b_p.dtype) - b_dw = tl.load(p_dw, boundary_check=(0, 1)) - b_dA += tl.dot(b_dw, tl.trans(b_p_beta_g0)) + b_dw = tl.load(p_dw, mask=m_k, other=0.0) + b_dA = tl.dot(b_dw, tl.trans(b_p_beta_g0), b_dA) b_dp_beta_g0 = tl.dot(b_A, b_dw) b_dp = b_dp_beta_g0 * b_beta[:, None] * b_g0_exp[:, None] b_dbeta += tl.sum(b_dp_beta_g0 * b_p * b_g0_exp[:, None], 1) b_dg0 += tl.sum(b_dp * b_p, 1) - tl.store(p_dp, b_dp.to(p_dp.dtype.element_ty), boundary_check=(0, 1)) + tl.store(p_dp, b_dp.to(p_dp.dtype.element_ty), mask=m_k) for i_v in range(tl.cdiv(V, BV)): - p_v = tl.make_block_ptr(v + (bos*H + i_h) * V, (T, V), (H*V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) - p_dv = tl.make_block_ptr(dv + (bos*H + i_h) * V, (T, V), (H*V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) - p_du = tl.make_block_ptr(du + (bos*H + i_h) * V, (T, V), (H*V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) - b_v = tl.load(p_v, boundary_check=(0, 1)) + o_v = i_v * BV + tl.arange(0, BV) + m_v = m_t[:, None] & (o_v[None, :] < V) + p_v = v + (bos*HV + i_h) * V + o_t[:, None] * (HV*V) + o_v[None, :] + p_dv = dv + (bos*HV + i_h) * V + o_t[:, None] * (HV*V) + o_v[None, :] + p_du = du + (bos*HV + i_h) * V + o_t[:, None] * (HV*V) + o_v[None, :] + b_v = tl.load(p_v, mask=m_v, other=0.0) b_v_beta = (b_v * b_beta[:, None]).to(b_v.dtype) - b_du = tl.load(p_du, boundary_check=(0, 1)) - b_dA += tl.dot(b_du, tl.trans(b_v_beta)) + b_du = tl.load(p_du, mask=m_v, other=0.0) + b_dA = tl.dot(b_du, tl.trans(b_v_beta), b_dA) b_dv_beta = tl.dot(b_A, b_du) b_dv = b_dv_beta * b_beta[:, None] b_dbeta += tl.sum(b_dv_beta * b_v, 1) - tl.store(p_dv, b_dv.to(p_dv.dtype.element_ty), boundary_check=(0, 1)) + tl.store(p_dv, b_dv.to(p_dv.dtype.element_ty), mask=m_v) - o_t = i_t * BT + tl.arange(0, BT) - m_t = o_t < T m_A = (o_t[:, None] > o_t[None, :]) & (m_t[:, None] & m_t) b_dA = tl.where(m_A, b_dA, 0) b_dA = tl.dot(b_dA.to(b_A.dtype), b_A) @@ -245,31 +257,33 @@ def prepare_wy_repr_bwd_kernel( b_A = tl.zeros([BT, BT], dtype=tl.float32) for i_k in range(tl.cdiv(K, BK)): - p_k = tl.make_block_ptr(k + (bos*H + i_h) * K, (T, K), (H*K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) - p_p = tl.make_block_ptr(p + (bos*H + i_h) * K, (T, K), (H*K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) - p_dk = tl.make_block_ptr(dk + (bos*H + i_h) * K, (T, K), (H*K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) - p_dp = tl.make_block_ptr(dp + (bos*H + i_h) * K, (T, K), (H*K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) - b_k = tl.load(p_k, boundary_check=(0, 1)) - b_p = tl.load(p_p, boundary_check=(0, 1)) - b_dp = tl.load(p_dp, boundary_check=(0, 1)) + o_k = i_k * BK + tl.arange(0, BK) + m_k = m_t[:, None] & (o_k[None, :] < K) + p_k = k + (bos*H + i_h // (HV // H)) * K + o_t[:, None] * (H*K) + o_k[None, :] + p_p = p + (bos*H + i_h // (HV // H)) * K + o_t[:, None] * (H*K) + o_k[None, :] + p_dk = dk + (bos*HV + i_h) * K + o_t[:, None] * (HV*K) + o_k[None, :] + p_dp = dp + (bos*HV + i_h) * K + o_t[:, None] * (HV*K) + o_k[None, :] + b_k = tl.load(p_k, mask=m_k, other=0.0) + b_p = tl.load(p_p, mask=m_k, other=0.0) + b_dp = tl.load(p_dp, mask=m_k, other=0.0) b_p_beta = (b_p * b_beta[:, None]).to(b_p.dtype) - b_A += tl.dot(b_p_beta, tl.trans(b_k)) + b_A = tl.dot(b_p_beta, tl.trans(b_k), b_A) b_dp_beta = tl.dot(b_dA, b_k) b_dbeta += tl.sum(b_dp_beta * b_p, 1) b_dk = tl.dot(tl.trans(b_dA), b_p_beta) b_dp += b_dp_beta * b_beta[:, None] - tl.store(p_dk, b_dk.to(p_dk.dtype.element_ty), boundary_check=(0, 1)) - tl.store(p_dp, b_dp.to(p_dp.dtype.element_ty), boundary_check=(0, 1)) + tl.store(p_dk, b_dk.to(p_dk.dtype.element_ty), mask=m_k) + tl.store(p_dp, b_dp.to(p_dp.dtype.element_ty), mask=m_k) b_dA_A = b_dA * b_A b_dg0 += tl.sum(b_dA_A, axis=1) b_dg = - tl.sum(b_dA_A, axis=0) - p_dg = tl.make_block_ptr(dg + (bos*H + i_h), (T,), (H,), (i_t * BT,), (BT,), (0,)) - p_dg0 = tl.make_block_ptr(dg0 + (bos*H + i_h), (T,), (H,), (i_t * BT,), (BT,), (0,)) - p_dbeta = tl.make_block_ptr(dbeta + (bos*H + i_h), (T,), (H,), (i_t * BT,), (BT,), (0,)) - tl.store(p_dg, b_dg.to(p_dg.dtype.element_ty), boundary_check=(0,)) - tl.store(p_dg0, b_dg0.to(p_dg0.dtype.element_ty), boundary_check=(0,)) - tl.store(p_dbeta, b_dbeta.to(p_dbeta.dtype.element_ty), boundary_check=(0,)) + p_dg = dg + (bos*HV + i_h) + o_t * HV + p_dg0 = dg0 + (bos*HV + i_h) + o_t * HV + p_dbeta = dbeta + (bos*HV + i_h) + o_t * HV + tl.store(p_dg, b_dg.to(p_dg.dtype.element_ty), mask=m_t) + tl.store(p_dg0, b_dg0.to(p_dg0.dtype.element_ty), mask=m_t) + tl.store(p_dbeta, b_dbeta.to(p_dbeta.dtype.element_ty), mask=m_t) @triton.heuristics({ @@ -281,7 +295,7 @@ def prepare_wy_repr_bwd_kernel( for num_warps in [2, 4, 8] for num_stages in [2, 3, 4] ], - key=['H', 'K', 'V', 'BT', 'BK', 'BV', 'IS_VARLEN'], + key=['H', 'HV', 'K', 'V', 'BT', 'BK', 'BV', 'IS_VARLEN'], **autotune_cache_kwargs, ) @triton.jit(do_not_specialize=['T']) @@ -297,6 +311,7 @@ def recompute_w_u_fwd_kernel( chunk_indices, T, H: tl.constexpr, + HV: tl.constexpr, K: tl.constexpr, V: tl.constexpr, BT: tl.constexpr, @@ -304,36 +319,44 @@ def recompute_w_u_fwd_kernel( BV: tl.constexpr, IS_VARLEN: tl.constexpr, ): - i_t, i_bh = tl.program_id(0), tl.program_id(1) - i_b, i_h = i_bh // H, i_bh % H + i_t, i_bh = tl.program_id(0).to(tl.int64), tl.program_id(1).to(tl.int64) + i_b, i_h = i_bh // HV, i_bh % HV 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.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) + i_n, i_t = tl.load(chunk_indices + i_t * 2).to(tl.int32), tl.load(chunk_indices + i_t * 2 + 1).to(tl.int64) + bos, eos = tl.load(cu_seqlens + i_n).to(tl.int64), tl.load(cu_seqlens + i_n + 1).to(tl.int64) T = eos - bos else: bos, eos = i_b * T, i_b * T + T - p_beta = tl.make_block_ptr(beta + bos*H + i_h, (T,), (H,), (i_t * BT,), (BT,), (0,)) - p_g = tl.make_block_ptr(g + (bos*H + i_h), (T,), (H,), (i_t * BT,), (BT,), (0,)) - p_A = tl.make_block_ptr(A + (bos*H + i_h) * BT, (T, BT), (H*BT, 1), (i_t * BT, 0), (BT, BT), (1, 0)) - b_beta = tl.load(p_beta, boundary_check=(0,)) - b_A = tl.load(p_A, boundary_check=(0, 1)) - b_g = exp2(tl.load(p_g, boundary_check=(0,))) + o_t = i_t * BT + tl.arange(0, BT) + o_A = tl.arange(0, BT) + m_t = o_t < T + m_A = m_t[:, None] & (o_A[None, :] < BT) + p_beta = beta + bos*HV + i_h + o_t * HV + p_g = g + (bos*HV + i_h) + o_t * HV + p_A = A + (bos*HV + i_h) * BT + o_t[:, None] * (HV*BT) + o_A[None, :] + b_beta = tl.load(p_beta, mask=m_t, other=0.0) + b_A = tl.load(p_A, mask=m_A, other=0.0) + b_g = exp2(tl.load(p_g, mask=m_t, other=0.0)) for i_v in range(tl.cdiv(V, BV)): - p_v = tl.make_block_ptr(v + (bos*H + i_h) * V, (T, V), (H*V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) - p_u = tl.make_block_ptr(u + (bos*H + i_h) * V, (T, V), (H*V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) - b_v = tl.load(p_v, boundary_check=(0, 1)) + o_v = i_v * BV + tl.arange(0, BV) + m_v = m_t[:, None] & (o_v[None, :] < V) + p_v = v + (bos*HV + i_h) * V + o_t[:, None] * (HV*V) + o_v[None, :] + p_u = u + (bos*HV + i_h) * V + o_t[:, None] * (HV*V) + o_v[None, :] + b_v = tl.load(p_v, mask=m_v, other=0.0) b_vb = (b_v * b_beta[:, None]).to(b_v.dtype) b_u = tl.dot(b_A, b_vb, allow_tf32=False) - tl.store(p_u, b_u.to(p_u.dtype.element_ty), boundary_check=(0, 1)) + tl.store(p_u, b_u.to(p_u.dtype.element_ty), mask=m_v) for i_k in range(tl.cdiv(K, BK)): - p_k = tl.make_block_ptr(k + (bos*H + i_h) * K, (T, K), (H*K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) - p_w = tl.make_block_ptr(w + (bos*H + i_h) * K, (T, K), (H*K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) - b_k = tl.load(p_k, boundary_check=(0, 1)) + o_k = i_k * BK + tl.arange(0, BK) + m_k = m_t[:, None] & (o_k[None, :] < K) + p_k = k + (bos*H + i_h // (HV // H)) * K + o_t[:, None] * (H*K) + o_k[None, :] + p_w = w + (bos*HV + i_h) * K + o_t[:, None] * (HV*K) + o_k[None, :] + b_k = tl.load(p_k, mask=m_k, other=0.0) b_kb = (b_k * b_beta[:, None] * b_g[:, None]).to(b_k.dtype) b_w = tl.dot(b_A, b_kb) - tl.store(p_w, b_w.to(p_w.dtype.element_ty), boundary_check=(0, 1)) + tl.store(p_w, b_w.to(p_w.dtype.element_ty), mask=m_k) def recompute_w_u_fwd( @@ -345,7 +368,7 @@ def recompute_w_u_fwd( cu_seqlens: torch.LongTensor | None, chunk_indices: torch.LongTensor | None = None, ) -> tuple[torch.Tensor, torch.Tensor]: - B, T, H, K, V = *k.shape, v.shape[-1] + B, T, H, K, V, HV = *k.shape, v.shape[-1], v.shape[2] BT = A.shape[-1] if chunk_indices is None and cu_seqlens is not None: @@ -355,8 +378,8 @@ def recompute_w_u_fwd( BV = 64 u = torch.empty_like(v) - w = torch.empty_like(k) - recompute_w_u_fwd_kernel[(NT, B*H)]( + w = k.new_empty(B, T, HV, K) + recompute_w_u_fwd_kernel[(NT, B*HV)]( k=k, v=v, beta=beta, @@ -368,6 +391,7 @@ def recompute_w_u_fwd( chunk_indices=chunk_indices, T=T, H=H, + HV=HV, K=K, V=V, BT=BT, @@ -390,7 +414,7 @@ def prepare_wy_repr_bwd( cu_seqlens: torch.LongTensor | None, chunk_indices: torch.LongTensor | None = None, ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: - B, T, H, K, V = *k.shape, v.shape[-1] + B, T, H, K, V, HV = *k.shape, v.shape[-1], v.shape[2] BT = A.shape[-1] if chunk_indices is None and cu_seqlens is not None: chunk_indices = prepare_chunk_indices(cu_seqlens, BT) @@ -399,13 +423,13 @@ def prepare_wy_repr_bwd( BK = min(max(triton.next_power_of_2(K), 16), CONST_TILING) BV = min(max(triton.next_power_of_2(V), 16), CONST_TILING) - dk = torch.empty_like(k) + dk = k.new_empty(B, T, HV, K) dv = torch.empty_like(v) - dp = torch.empty_like(p) + dp = p.new_empty(B, T, HV, K) dbeta = torch.empty_like(beta) dg0 = torch.empty_like(g0) dg = torch.empty_like(g) - prepare_wy_repr_bwd_kernel[(NT, B * H)]( + prepare_wy_repr_bwd_kernel[(NT, B * HV)]( k=k, v=v, p=p, @@ -425,10 +449,14 @@ def prepare_wy_repr_bwd( chunk_indices=chunk_indices, T=T, H=H, + HV=HV, K=K, V=V, BT=BT, BK=BK, BV=BV, ) + if H != HV: + dk = dk.view(B, T, H, HV // H, K).sum(3) + dp = dp.view(B, T, H, HV // H, K).sum(3) return dk, dv, dp, dbeta, dg0, dg diff --git a/build/torch-cuda/ops/common/__init__.py b/build/torch-cuda/ops/common/__init__.py index e69de29bb2d1d6434b8b29ae775ad8c2e48c5391..b2023fce4385271c57f2348a7db655ddbb7d6ffb 100644 --- a/build/torch-cuda/ops/common/__init__.py +++ b/build/torch-cuda/ops/common/__init__.py @@ -0,0 +1,7 @@ +# 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 + diff --git a/build/torch-cuda/ops/common/backends/__init__.py b/build/torch-cuda/ops/common/backends/__init__.py index 18ea07b2dee4fa20beb9809b751e5b86f222db2a..7081a5cdc44850f14197a4b2e425a74c35d2d5c8 100644 --- a/build/torch-cuda/ops/common/backends/__init__.py +++ b/build/torch-cuda/ops/common/backends/__init__.py @@ -7,15 +7,7 @@ """Common backends for shared operations like chunk_gated_delta_rule_fwd_h.""" -from ....ops.backends import BackendRegistry, dispatch from ....ops.common.backends.intracard import IntraCardCPBackend -from ....ops.common.backends.tilelang import TileLangBackend +from ....ops.common.backends.tilelang import CommonTileLangBackend -common_registry = BackendRegistry("common") - - -common_registry.register(IntraCardCPBackend()) -common_registry.register(TileLangBackend()) - - -__all__ = ['common_registry', 'dispatch'] +__all__ = ['CommonTileLangBackend', 'IntraCardCPBackend'] diff --git a/build/torch-cuda/ops/common/backends/intracard.py b/build/torch-cuda/ops/common/backends/intracard.py index 17c47ad52841ca4aa6cc6d201c689490eae338a7..3bcd9699ba2c7eab1d2bc8c6ffd6b0629fee0dbb 100644 --- a/build/torch-cuda/ops/common/backends/intracard.py +++ b/build/torch-cuda/ops/common/backends/intracard.py @@ -16,16 +16,28 @@ Only active under torch.inference_mode() with varlen (cu_seqlens != None). from __future__ import annotations import os +import warnings import torch -from ....ops.backends import BaseBackend +from ....backends import BaseBackend, register +from ....utils import IS_TF32_SUPPORTED # Maximum number of sub-sequences per original sequence # Limits merge chain depth to control precision loss MAX_SUBSEQS = int(os.environ.get('FLA_INTRACARD_MAX_SPLITS', 32)) +# use tf32x3 for the affine-chain dots in the pre-scan/merge kernels (NVIDIA only) +USE_TF32X3_AFFINE_CHAIN = os.environ.get('FLA_INTRACARD_TF32X3', '0') == '1' +if USE_TF32X3_AFFINE_CHAIN and not IS_TF32_SUPPORTED: + warnings.warn( + "tf32x3 affine chain requires an NVIDIA GPU with compute capability >= 8.0; falling back to ieee precision", + stacklevel=2, + ) + + +@register('common') class IntraCardCPBackend(BaseBackend): """Intra-card context parallel backend for chunk_gated_delta_rule_fwd_h.""" @@ -53,6 +65,7 @@ class IntraCardCPBackend(BaseBackend): cu_seqlens: torch.LongTensor | None = None, cu_seqlens_cpu: torch.LongTensor | None = None, chunk_indices: torch.LongTensor | None = None, + chunk_offsets: torch.LongTensor | None = None, ) -> tuple[bool, str | None]: """Check if intracard CP should handle this call.""" # Only in inference mode @@ -63,6 +76,11 @@ class IntraCardCPBackend(BaseBackend): if cu_seqlens is None: return False, "cu_seqlens is None" + # Only graph-mode callers pass chunk_offsets. Splitting reads cu_seqlens on the host, + # which CUDA Graph capture cannot record. + if chunk_offsets is not None: + return False, "chunk_offsets is given (CUDA Graph mode)" + return True, None def chunk_gated_delta_rule_fwd_h( @@ -80,8 +98,10 @@ class IntraCardCPBackend(BaseBackend): cu_seqlens: torch.LongTensor | None = None, cu_seqlens_cpu: torch.LongTensor | None = None, chunk_indices: torch.LongTensor | None = None, + chunk_offsets: torch.LongTensor | None = None, ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor | None]: """Intra-card CP implementation of chunk_gated_delta_rule_fwd_h.""" + # The verifier only accepts chunk_offsets=None; intracard_fwd_h derives the offsets it launches with. from ....ops.common.intracard_cp import intracard_fwd_h return intracard_fwd_h( @@ -95,4 +115,5 @@ class IntraCardCPBackend(BaseBackend): chunk_indices=chunk_indices, max_splits=MAX_SUBSEQS, state_v_first=state_v_first, + use_tf32x3_affine_chain=USE_TF32X3_AFFINE_CHAIN, ) diff --git a/build/torch-cuda/ops/common/backends/tilelang/__init__.py b/build/torch-cuda/ops/common/backends/tilelang/__init__.py index 088f5ac608a1a8ac368eda0cdc17594083ce999b..aba279e1e84d5470d183edb5f685f87bcf21c0a1 100644 --- a/build/torch-cuda/ops/common/backends/tilelang/__init__.py +++ b/build/torch-cuda/ops/common/backends/tilelang/__init__.py @@ -5,44 +5,18 @@ # For a list of all contributors, visit: # https://github.com/fla-org/flash-linear-attention/graphs/contributors -"""TileLang backend for common chunk operations. - -Enabled by default on Hopper (sm90+) with Triton >= 3.4.0 to work around -hardware-specific regressions (see #640). Can also be forced via FLA_TILELANG=1. -""" - from __future__ import annotations -import os - import torch -from .....ops.backends import BaseBackend -from .....utils import IS_NVIDIA_HOPPER, TRITON_ABOVE_3_4_0, find_spec_cached - -_TILELANG_AVAILABLE = find_spec_cached("tilelang") is not None - - -class TileLangBackend(BaseBackend): - - backend_type = "tilelang" - package_name = "tilelang" - env_var = "FLA_TILELANG" +from .....backends import TileLangBackend, register +from .....utils import IS_NVIDIA_HOPPER, TRITON_ABOVE_3_4_0 - @classmethod - def is_available(cls) -> bool: - return _TILELANG_AVAILABLE - @classmethod - def is_enabled(cls) -> bool: - # Explicit opt-in / opt-out always wins. - val = os.environ.get(cls.env_var) - if val is not None: - return val != "0" - # Default on only where the Triton path is known to be broken: - # Hopper (sm90) with Triton >= 3.4.0 (see #640). Everywhere else the - # Triton backend stays the default unless the user forces FLA_TILELANG=1. - return IS_NVIDIA_HOPPER and TRITON_ABOVE_3_4_0 +@register('common') +class CommonTileLangBackend(TileLangBackend): + # work around Hopper regressions with Triton 3.4+ (see #640). + default_enable = IS_NVIDIA_HOPPER and TRITON_ABOVE_3_4_0 def chunk_bwd_dqkwg_verifier( self, @@ -96,9 +70,7 @@ class TileLangBackend(BaseBackend): chunk_size: int = 64, chunk_indices: torch.LongTensor | None = None, ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor | None, torch.Tensor | None]: - from .....ops.common.backends.tilelang.chunk_bwd import ( - chunk_bwd_dqkwg_tilelang, - ) + from .....ops.common.backends.tilelang.chunk_bwd import chunk_bwd_dqkwg_tilelang return chunk_bwd_dqkwg_tilelang( q=q, k=k, @@ -111,91 +83,11 @@ class TileLangBackend(BaseBackend): g_gamma=g_gamma, dv=dv, scale=scale, + state_v_first=state_v_first, cu_seqlens=cu_seqlens, chunk_size=chunk_size, chunk_indices=chunk_indices, - state_v_first=state_v_first, - ) - - def parallel_attn_fwd_verifier( - self, - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - g_cumsum: torch.Tensor | None, - sink_bias: torch.Tensor | None, - scale: float, - window_size: int | None = None, - cu_seqlens: torch.LongTensor | None = None, - chunk_indices: torch.LongTensor | None = None, - ) -> tuple[bool, str | None]: - if q.dtype not in (torch.float16, torch.bfloat16, torch.float32): - return False, f"TileLang backend does not support dtype {q.dtype}; fall back to Triton" - return True, None - - def parallel_attn_fwd( - self, - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - g_cumsum: torch.Tensor | None, - sink_bias: torch.Tensor | None, - scale: float, - window_size: int | None = None, - cu_seqlens: torch.LongTensor | None = None, - chunk_indices: torch.LongTensor | None = None, - ) -> tuple[torch.Tensor, torch.Tensor]: - from .....ops.common.backends.tilelang.parallel_attn_fwd import ( - parallel_attn_fwd_tilelang, - ) - return parallel_attn_fwd_tilelang( - q=q, k=k, v=v, g_cumsum=g_cumsum, sink_bias=sink_bias, - scale=scale, window_size=window_size, cu_seqlens=cu_seqlens, - chunk_indices=chunk_indices, ) - def parallel_attn_bwd_verifier( - self, - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - o: torch.Tensor, - g_cumsum: torch.Tensor | None, - lse: torch.Tensor, - do: torch.Tensor, - sink_bias: torch.Tensor | None = None, - scale: float | None = None, - window_size: int | None = None, - chunk_size: int = 128, - cu_seqlens: torch.LongTensor | None = None, - chunk_indices: torch.LongTensor | None = None, - ) -> tuple[bool, str | None]: - if q.dtype not in (torch.float16, torch.bfloat16, torch.float32): - return False, f"TileLang backend does not support dtype {q.dtype}; fall back to Triton" - return True, None - def parallel_attn_bwd( - self, - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - o: torch.Tensor, - g_cumsum: torch.Tensor | None, - lse: torch.Tensor, - do: torch.Tensor, - sink_bias: torch.Tensor | None = None, - scale: float | None = None, - window_size: int | None = None, - chunk_size: int = 128, - cu_seqlens: torch.LongTensor | None = None, - chunk_indices: torch.LongTensor | None = None, - ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor | None, torch.Tensor | None]: - from .....ops.common.backends.tilelang.parallel_attn_bwd import ( - parallel_attn_bwd_tilelang, - ) - return parallel_attn_bwd_tilelang( - q=q, k=k, v=v, o=o, g_cumsum=g_cumsum, lse=lse, do=do, - sink_bias=sink_bias, scale=scale, window_size=window_size, - chunk_size=chunk_size, cu_seqlens=cu_seqlens, - chunk_indices=chunk_indices, - ) +__all__ = ['CommonTileLangBackend'] diff --git a/build/torch-cuda/ops/common/backends/tilelang/chunk_bwd.py b/build/torch-cuda/ops/common/backends/tilelang/chunk_bwd.py index 3c89a544206dfdd92fb9a87b4daf2092f5162f53..84b5584c7ae50fcdfd58c5c1d8d3d9c261cf547e 100644 --- a/build/torch-cuda/ops/common/backends/tilelang/chunk_bwd.py +++ b/build/torch-cuda/ops/common/backends/tilelang/chunk_bwd.py @@ -18,7 +18,7 @@ from .....utils import check_shared_mem tilelang.PassConfigKey.TL_DISABLE_TMA_LOWER: True, tilelang.PassConfigKey.TL_ENABLE_FAST_MATH: True, }) -def _build_kernel( +def _build_chunk_bwd_dqkwg_kernel( B, H, HV, @@ -43,263 +43,248 @@ def _build_kernel( threads = num_warps * 32 tile_hD1, tile_hD2 = (BV, BK) if STATE_V_FIRST else (BK, BV) - # Rebind to underscore-prefixed locals so the kernel body (a closure below) - # stays identical to the previous nested layout. tilelang caches _build_kernel - # by the outer (B, H, K, V, BT, BK, BV, NK, hD1, hD2, dtype_str, ...) tuple. - _B, _H, _HV, _K, _V = B, H, HV, K, V - _G = HV // H - _BT, _BK, _BV, _NK = BT, BK, BV, NK - _NV = NV - _hD1, _hD2, _thD1, _thD2 = hD1, hD2, tile_hD1, tile_hD2 - _threads = threads - _USE_G, _USE_DW = USE_G, USE_DW - _TS, _VAR = STATE_V_FIRST, IS_VARLEN - - # T, NT, total_h are dynamic (vary with sequence length, no recompilation). - # B, H are compile-time (stable across batches, enables fast integer division). + G = HV // H + T_d, NT_d, total_h_d, Ncu_d = T.dynamic("T, NT, total_h, Ncu") - # 4D tensor shapes using dynamic T + compile-time B, H, K, V. - # q/k carry H qk-heads; dq/dk/dw and v/do/g/dg carry HV value-heads - # (HV >= H for GVA). dq/dk/dw are written at value-head granularity and - # reduced to qk-heads by the caller when HV > H. - qk_s = (_B, T_d, _H, _K) - dqk_s = (_B, T_d, _HV, _K) - v_s = (_B, T_d, _HV, _V) - h_s = (total_h_d, _hD1, _hD2) - g_s = (_B, T_d, _HV) - dg_s = (_NK, _B, T_d, _HV) + # q/k use H heads; per-value-head gradients use HV + qk_s = (B, T_d, H, K) + dqk_s = (B, T_d, HV, K) + v_s = (B, T_d, HV, V) + h_s = (total_h_d, hD1, hD2) + g_s = (B, T_d, HV) + dg_s = (NK, B, T_d, HV) @T.macro - def kernel_body(q, k, v, g, h, do, dh, dq, dk, dw, dv, dg, scale, - i_b, i_h, i_k, t_s, T_seq, i_t_local, h_idx, k_off): - # Map the value-head program id to the qk-head it belongs to. For GVA - # (HV > H) multiple value-heads share the same q/k head. - i_hqk = i_h // _G - # -- accumulators -- - b_dq = T.alloc_fragment((_BT, _BK), T.float32) - b_dk = T.alloc_fragment((_BT, _BK), T.float32) - b_ds = T.alloc_fragment((_BT, _BT), T.float32) + def kernel_body(q, k, v, g, h, do, dh, dq, dk, dw, dv, dg, scale, i_b, i_h, i_k, t_s, T_seq, i_t_local, h_idx, k_off): + i_hqk = i_h // G + b_dq = T.alloc_fragment((BT, BK), T.float32) + b_dk = T.alloc_fragment((BT, BK), T.float32) + b_ds = T.alloc_fragment((BT, BT), T.float32) T.clear(b_dq) T.clear(b_dk) T.clear(b_ds) - if _USE_DW: - b_dw = T.alloc_fragment((_BT, _BK), T.float32) + if USE_DW: + b_dw = T.alloc_fragment((BT, BK), T.float32) T.clear(b_dw) - # -- shared tiles -- - s_v = T.alloc_shared((_BT, _BV), _dtype) - s_do = T.alloc_shared((_BT, _BV), _dtype) - s_h = T.alloc_shared((_thD1, _thD2), _dtype) - s_dh = T.alloc_shared((_thD1, _thD2), _dtype) + s_v = T.alloc_shared((BT, BV), _dtype) + s_do = T.alloc_shared((BT, BV), _dtype) + s_h = T.alloc_shared((tile_hD1, tile_hD2), _dtype) + s_dh = T.alloc_shared((tile_hD1, tile_hD2), _dtype) - # dg_last accumulator (shared scalar, accumulated across V-loop) - if _USE_G: + if USE_G: s_dg_last_acc = T.alloc_shared((1,), T.float32) for _i in T.Parallel(1): s_dg_last_acc[0] = 0.0 T.sync_threads() - # ========== V-loop ========== - for i_v_py in T.Pipelined(_NV, num_stages=2): - v_off_c = i_v_py * _BV + for i_v_py in T.Pipelined(NV, num_stages=2): + v_off_c = i_v_py * BV - T.copy(v[i_b, t_s:t_s + _BT, i_h, v_off_c:v_off_c + _BV], s_v) - T.copy(do[i_b, t_s:t_s + _BT, i_h, v_off_c:v_off_c + _BV], s_do) + T.copy(v[i_b, t_s:t_s + BT, i_h, v_off_c:v_off_c + BV], s_v) + T.copy(do[i_b, t_s:t_s + BT, i_h, v_off_c:v_off_c + BV], s_do) - if _TS: - T.copy(h[h_idx, v_off_c:v_off_c + _BV, k_off:k_off + _BK], s_h) - T.copy(dh[h_idx, v_off_c:v_off_c + _BV, k_off:k_off + _BK], s_dh) + if STATE_V_FIRST: + T.copy(h[h_idx, v_off_c:v_off_c + BV, k_off:k_off + BK], s_h) + T.copy(dh[h_idx, v_off_c:v_off_c + BV, k_off:k_off + BK], s_dh) else: - T.copy(h[h_idx, k_off:k_off + _BK, v_off_c:v_off_c + _BV], s_h) - T.copy(dh[h_idx, k_off:k_off + _BK, v_off_c:v_off_c + _BV], s_dh) + T.copy(h[h_idx, k_off:k_off + BK, v_off_c:v_off_c + BV], s_h) + T.copy(dh[h_idx, k_off:k_off + BK, v_off_c:v_off_c + BV], s_dh) T.gemm(s_do, s_v, b_ds, transpose_B=True) - # h·dh reduction must precede the gemms that consume s_h/s_dh: - # the downstream gemms are what the pipeline recognizes as consumers, - # so they act as the barrier against the next iter's prefetch. - if _USE_G: - f_hdh = T.alloc_fragment((_thD1, _thD2), T.float32) - for _i, _j in T.Parallel(_thD1, _thD2): + # reduce h·dh before the GEMMs that keep the next prefetch from overwriting s_h/s_dh + if USE_G: + f_hdh = T.alloc_fragment((tile_hD1, tile_hD2), T.float32) + for _i, _j in T.Parallel(tile_hD1, tile_hD2): f_hdh[_i, _j] = T.cast(s_h[_i, _j], T.float32) * T.cast(s_dh[_i, _j], T.float32) - f_hdh_row = T.alloc_fragment((_thD1,), T.float32) + f_hdh_row = T.alloc_fragment((tile_hD1,), T.float32) T.reduce_sum(f_hdh, f_hdh_row, dim=1) f_hdh_scalar = T.alloc_fragment((1,), T.float32) T.reduce_sum(f_hdh_row, f_hdh_scalar, dim=0) s_dg_last_acc[0] = s_dg_last_acc[0] + f_hdh_scalar[0] - if _TS: + if STATE_V_FIRST: T.gemm(s_do, s_h, b_dq) T.gemm(s_v, s_dh, b_dk) else: T.gemm(s_do, s_h, b_dq, transpose_B=True) T.gemm(s_v, s_dh, b_dk, transpose_B=True) - if _USE_DW: - s_dv = T.alloc_shared((_BT, _BV), _dtype) - T.copy(dv[i_b, t_s:t_s + _BT, i_h, v_off_c:v_off_c + _BV], s_dv) - if _TS: + if USE_DW: + s_dv = T.alloc_shared((BT, BV), _dtype) + T.copy(dv[i_b, t_s:t_s + BT, i_h, v_off_c:v_off_c + BV], s_dv) + if STATE_V_FIRST: T.gemm(s_dv, s_h, b_dw) else: T.gemm(s_dv, s_h, b_dw, transpose_B=True) - # ========== store dw (negated, with varlen boundary mask) ========== - if _USE_DW: - s_dw_out = T.alloc_shared((_BT, _BK), _dtype) - for _i, _j in T.Parallel(_BT, _BK): + if USE_DW: + s_dw_out = T.alloc_shared((BT, BK), _dtype) + for _i, _j in T.Parallel(BT, BK): s_dw_out[_i, _j] = T.cast(-b_dw[_i, _j], _dtype) T.sync_threads() - for _i, _j in T.Parallel(_BT, _BK): - if (i_t_local * _BT + _i) < T_seq: + for _i, _j in T.Parallel(BT, BK): + if (i_t_local * BT + _i) < T_seq: dw[i_b, t_s + _i, i_h, k_off + _j] = s_dw_out[_i, _j] - # dg_last is now in s_dg_last_acc[0] (shared memory, visible to all threads) - - # ========== load q, k ========== - # q/k are indexed by the shared qk-head, not the value-head. - s_q = T.alloc_shared((_BT, _BK), _dtype) - s_k = T.alloc_shared((_BT, _BK), _dtype) - T.copy(q[i_b, t_s:t_s + _BT, i_hqk, k_off:k_off + _BK], s_q) - T.copy(k[i_b, t_s:t_s + _BT, i_hqk, k_off:k_off + _BK], s_k) - - # ========== USE_G path ========== - if _USE_G: - # dg_last_acc already reduced to `red` scalar above + s_q = T.alloc_shared((BT, BK), _dtype) + s_k = T.alloc_shared((BT, BK), _dtype) + T.copy(q[i_b, t_s:t_s + BT, i_hqk, k_off:k_off + BK], s_q) + T.copy(k[i_b, t_s:t_s + BT, i_hqk, k_off:k_off + BK], s_k) - # g in shared memory (all threads can cross-read any element) - s_g = T.alloc_shared((_BT,), T.float32) - T.copy(g[i_b, t_s:t_s + _BT, i_h], s_g, disable_tma=True) + if USE_G: + # shared memory lets all threads read the last valid gate + s_g = T.alloc_shared((BT,), T.float32) + T.copy(g[i_b, t_s:t_s + BT, i_h], s_g, disable_tma=True) - last_pos = T.max(0, T.min(_BT, T_seq - i_t_local * _BT) - 1) + last_pos = T.max(0, T.min(BT, T_seq - i_t_local * BT) - 1) g_last = s_g[last_pos] b_dg_last = T.alloc_var(T.float32) b_dg_last = s_dg_last_acc[0] * T.exp2(g_last) - # b_dq *= exp2(g) * scale (inline, no extra fragment) - for _i, _j in T.Parallel(_BT, _BK): + for _i, _j in T.Parallel(BT, BK): b_dq[_i, _j] = b_dq[_i, _j] * T.exp2(s_g[_i]) * scale - # Gate b_dk with m_t mask: zero out OOB positions - for _i, _j in T.Parallel(_BT, _BK): - m_t = (i_t_local * _BT + _i) < T_seq - b_dk[_i, _j] = T.if_then_else( - m_t, - b_dk[_i, _j] * T.exp2(-s_g[_i] + g_last), - 0.0) - - # dk*k reduction (before ds gemms): compute product directly in - # fragments via fragment×shared. b_dg_last uses the gated-only b_dk. - f_prod2 = T.alloc_fragment((_BT, _BK), T.float32) - for _i, _j in T.Parallel(_BT, _BK): + for _i, _j in T.Parallel(BT, BK): + m_t = (i_t_local * BT + _i) < T_seq + b_dk[_i, _j] = T.if_then_else(m_t, b_dk[_i, _j] * T.exp2(-s_g[_i] + g_last), 0.0) + + # accumulate gated dk*k before the intra-chunk GEMMs update dk + f_prod2 = T.alloc_fragment((BT, BK), T.float32) + for _i, _j in T.Parallel(BT, BK): f_prod2[_i, _j] = b_dk[_i, _j] * s_k[_i, _j] - f_dg2 = T.alloc_fragment((_BT,), T.float32) + f_dg2 = T.alloc_fragment((BT,), T.float32) T.reduce_sum(f_prod2, f_dg2, dim=1) - # b_dg_last += sum(dk*k) via fragment reduce (f_dg2 → scalar) f_dkk_scalar = T.alloc_fragment((1,), T.float32) T.reduce_sum(f_dg2, f_dkk_scalar, dim=0) b_dg_last = b_dg_last + f_dkk_scalar[0] - # b_ds = where(causal, b_ds * exp2(g_i - g_j), 0) * scale - for _i, _j in T.Parallel(_BT, _BT): - causal = (_i >= _j) & ((i_t_local * _BT + _i) < T_seq) & ((i_t_local * _BT + _j) < T_seq) - b_ds[_i, _j] = T.if_then_else( - causal, - b_ds[_i, _j] * T.exp2(s_g[_i] - s_g[_j]) * scale, - 0.0) - - # cast ds for final gemms - s_ds = T.alloc_shared((_BT, _BT), _dtype) - f_ds = T.alloc_fragment((_BT, _BT), _dtype) - for _i, _j in T.Parallel(_BT, _BT): + for _i, _j in T.Parallel(BT, BT): + causal = (_i >= _j) & ((i_t_local * BT + _i) < T_seq) & ((i_t_local * BT + _j) < T_seq) + b_ds[_i, _j] = T.if_then_else(causal, b_ds[_i, _j] * T.exp2(s_g[_i] - s_g[_j]) * scale, 0.0) + + s_ds = T.alloc_shared((BT, BT), _dtype) + f_ds = T.alloc_fragment((BT, BT), _dtype) + for _i, _j in T.Parallel(BT, BT): f_ds[_i, _j] = T.cast(b_ds[_i, _j], _dtype) T.copy(f_ds, s_ds) - T.gemm(s_ds, s_k, b_dq) # dq += ds @ k - T.gemm(s_ds, s_q, b_dk, transpose_A=True) # dk += ds^T @ q + T.gemm(s_ds, s_k, b_dq) + T.gemm(s_ds, s_q, b_dk, transpose_A=True) - # b_dg = sum(b_dq * q) - sum(b_dk * k) using the fully-updated b_dq/b_dk - f_prod1 = T.alloc_fragment((_BT, _BK), T.float32) - for _i, _j in T.Parallel(_BT, _BK): + # dg uses the fully updated dq/dk + f_prod1 = T.alloc_fragment((BT, BK), T.float32) + for _i, _j in T.Parallel(BT, BK): f_prod1[_i, _j] = b_dq[_i, _j] * s_q[_i, _j] f_prod2[_i, _j] = b_dk[_i, _j] * s_k[_i, _j] - f_dg1 = T.alloc_fragment((_BT,), T.float32) + f_dg1 = T.alloc_fragment((BT,), T.float32) T.reduce_sum(f_prod1, f_dg1, dim=1) T.reduce_sum(f_prod2, f_dg2, dim=1) - f_dg_diff = T.alloc_fragment((_BT,), T.float32) - for _i in T.Parallel(_BT): + f_dg_diff = T.alloc_fragment((BT,), T.float32) + for _i in T.Parallel(BT): f_dg_diff[_i] = f_dg1[_i] - f_dg2[_i] - s_dg = T.alloc_shared((_BT,), T.float32) + s_dg = T.alloc_shared((BT,), T.float32) T.copy(f_dg_diff, s_dg) - # store dq, dk via shared (fragment has MMA layout, can't index directly) - f_out = T.alloc_fragment((_BT, _BK), _dtype) - s_out = T.alloc_shared((_BT, _BK), _dtype) - for _i, _j in T.Parallel(_BT, _BK): + # shared staging converts MMA fragments to an indexable layout + f_out = T.alloc_fragment((BT, BK), _dtype) + s_out = T.alloc_shared((BT, BK), _dtype) + for _i, _j in T.Parallel(BT, BK): f_out[_i, _j] = T.cast(b_dq[_i, _j], _dtype) T.copy(f_out, s_out) T.sync_threads() - for _i, _j in T.Parallel(_BT, _BK): - if (i_t_local * _BT + _i) < T_seq: + for _i, _j in T.Parallel(BT, BK): + if (i_t_local * BT + _i) < T_seq: dq[i_b, t_s + _i, i_h, k_off + _j] = s_out[_i, _j] - for _i, _j in T.Parallel(_BT, _BK): + for _i, _j in T.Parallel(BT, BK): f_out[_i, _j] = T.cast(b_dk[_i, _j], _dtype) T.copy(f_out, s_out) T.sync_threads() - for _i, _j in T.Parallel(_BT, _BK): - if (i_t_local * _BT + _i) < T_seq: + for _i, _j in T.Parallel(BT, BK): + if (i_t_local * BT + _i) < T_seq: dk[i_b, t_s + _i, i_h, k_off + _j] = s_out[_i, _j] - # store dg: merge dg_last into last valid position - for _i in T.Parallel(_BT): - if (i_t_local * _BT + _i) < T_seq: + for _i in T.Parallel(BT): + if (i_t_local * BT + _i) < T_seq: val = T.if_then_else(_i == last_pos, s_dg[_i] + b_dg_last, s_dg[_i]) dg[i_k, i_b, t_s + _i, i_h] = val - if _VAR: + if IS_VARLEN: @T.prim_func def kernel( - q: T.Tensor(qk_s, _dtype), k: T.Tensor(qk_s, _dtype), - v: T.Tensor(v_s, _dtype), g: T.Tensor(g_s, T.float32), - h: T.Tensor(h_s, _dtype), do: T.Tensor(v_s, _dtype), - dh: T.Tensor(h_s, _dtype), dq: T.Tensor(dqk_s, _dtype), - dk: T.Tensor(dqk_s, _dtype), dw: T.Tensor(dqk_s, _dtype), - dv: T.Tensor(v_s, _dtype), dg: T.Tensor(dg_s, T.float32), + q: T.Tensor(qk_s, _dtype), + k: T.Tensor(qk_s, _dtype), + v: T.Tensor(v_s, _dtype), + g: T.Tensor(g_s, T.float32), + h: T.Tensor(h_s, _dtype), + do: T.Tensor(v_s, _dtype), + dh: T.Tensor(h_s, _dtype), + dq: T.Tensor(dqk_s, _dtype), + dk: T.Tensor(dqk_s, _dtype), + dw: T.Tensor(dqk_s, _dtype), + dv: T.Tensor(v_s, _dtype), + dg: T.Tensor(dg_s, T.float32), cu_seqlens: T.Tensor((Ncu_d,), T.int32), chunk_indices: T.Tensor((NT_d, 2), T.int32), scale: T.float32, ): - with T.Kernel(_NK, NT_d, _HV, threads=_threads) as (i_k, i_t, i_h): + with T.Kernel(NK, NT_d, HV, threads=threads) as (i_k, i_t, i_h): i_n = chunk_indices[i_t, 0] i_t_local = chunk_indices[i_t, 1] bos = cu_seqlens[i_n] T_seq = cu_seqlens[i_n + 1] - bos - h_idx = i_t * _HV + i_h - t_s = bos + i_t_local * _BT - kernel_body(q, k, v, g, h, do, dh, dq, dk, dw, dv, dg, - scale, 0, i_h, i_k, t_s, T_seq, i_t_local, - h_idx, i_k * _BK) + h_idx = i_t * HV + i_h + t_s = bos + i_t_local * BT + kernel_body( + q, + k, + v, + g, + h, + do, + dh, + dq, + dk, + dw, + dv, + dg, + scale, + 0, + i_h, + i_k, + t_s, + T_seq, + i_t_local, + h_idx, + i_k * BK, + ) else: @T.prim_func def kernel( - q: T.Tensor(qk_s, _dtype), k: T.Tensor(qk_s, _dtype), - v: T.Tensor(v_s, _dtype), g: T.Tensor(g_s, T.float32), - h: T.Tensor(h_s, _dtype), do: T.Tensor(v_s, _dtype), - dh: T.Tensor(h_s, _dtype), dq: T.Tensor(dqk_s, _dtype), - dk: T.Tensor(dqk_s, _dtype), dw: T.Tensor(dqk_s, _dtype), - dv: T.Tensor(v_s, _dtype), dg: T.Tensor(dg_s, T.float32), + q: T.Tensor(qk_s, _dtype), + k: T.Tensor(qk_s, _dtype), + v: T.Tensor(v_s, _dtype), + g: T.Tensor(g_s, T.float32), + h: T.Tensor(h_s, _dtype), + do: T.Tensor(v_s, _dtype), + dh: T.Tensor(h_s, _dtype), + dq: T.Tensor(dqk_s, _dtype), + dk: T.Tensor(dqk_s, _dtype), + dw: T.Tensor(dqk_s, _dtype), + dv: T.Tensor(v_s, _dtype), + dg: T.Tensor(dg_s, T.float32), scale: T.float32, ): - with T.Kernel(_NK, T.ceildiv(T_d, _BT), _B * _HV, threads=_threads) as (i_k, i_t, i_bh): - i_b = i_bh // _HV - i_h = i_bh % _HV - NT_local = T.ceildiv(T_d, _BT) - h_idx = (i_b * NT_local + i_t) * _HV + i_h - t_s = i_t * _BT - kernel_body(q, k, v, g, h, do, dh, dq, dk, dw, dv, dg, - scale, i_b, i_h, i_k, t_s, T_d, i_t, - h_idx, i_k * _BK) + with T.Kernel(NK, T.ceildiv(T_d, BT), B * HV, threads=threads) as (i_k, i_t, i_bh): + i_b = i_bh // HV + i_h = i_bh % HV + NT_local = T.ceildiv(T_d, BT) + h_idx = (i_b * NT_local + i_t) * HV + i_h + t_s = i_t * BT + kernel_body(q, k, v, g, h, do, dh, dq, dk, dw, dv, dg, scale, i_b, i_h, i_k, t_s, T_d, i_t, h_idx, i_k * BK) return kernel @@ -325,7 +310,7 @@ def chunk_bwd_dqkwg_tilelang( HV, V = v.shape[2], v.shape[-1] BT = chunk_size if chunk_indices is None and cu_seqlens is not None: - chunk_indices = prepare_chunk_indices(cu_seqlens, BT) + chunk_indices = prepare_chunk_indices(cu_seqlens=cu_seqlens, chunk_size=BT) IS_VARLEN = cu_seqlens is not None CONST_TILING = 64 if check_shared_mem() else 32 @@ -338,8 +323,7 @@ def chunk_bwd_dqkwg_tilelang( USE_G = g is not None USE_DW = w is not None - # Outputs — kernel writes dq/dk at value-head (HV) granularity; reduce to - # qk-head (H) below when GVA is active. dg stays at HV (per value-head). + # dq/dk are computed per value head and reduced to query/key heads after the kernel dq = torch.empty(B, T, HV, K, dtype=q.dtype, device=q.device) dk = torch.empty(B, T, HV, K, dtype=k.dtype, device=k.device) dw_out = torch.empty_like(w) if USE_DW else None @@ -350,39 +334,52 @@ def chunk_bwd_dqkwg_tilelang( hD1, hD2 = h_flat.shape[-2], h_flat.shape[-1] dtype_str = {torch.float16: 'float16', torch.bfloat16: 'bfloat16', torch.float32: 'float32'}[q.dtype] - # Cache key: B, H, HV, tile sizes, flags. T is dynamic (no recompilation for different seq lengths). - # Small head dims (< 64) cannot be warp-partitioned across 4 warps by TileLang's - # GEMM (tile too small); drop to 2 warps so the kernel still compiles. + # small head dimensions need two warps for TileLang GEMM partitioning num_warps = 4 if min(K, V) >= 64 else 2 - kernel = _build_kernel( - B, - H, - HV, - K, - V, - BT, - BK, - BV, - NK, - hD1, - hD2, - dtype_str, - USE_G, - USE_DW, - state_v_first, - IS_VARLEN, + kernel = _build_chunk_bwd_dqkwg_kernel( + B=B, + H=H, + HV=HV, + K=K, + V=V, + BT=BT, + BK=BK, + BV=BV, + NK=NK, + hD1=hD1, + hD2=hD2, + dtype_str=dtype_str, + USE_G=USE_G, + USE_DW=USE_DW, + STATE_V_FIRST=state_v_first, + IS_VARLEN=IS_VARLEN, num_warps=num_warps, ) - # Unused optional params still need shape-matching tensors for TileLang + # optional inputs need shape-matching buffers for the kernel signature g_kern = g if USE_G else q.new_empty(B, T, HV) dw_kern = dw_out if USE_DW else q.new_empty(B, T, HV, K) dv_kern = dv if USE_DW else q.new_empty(B, T, HV, V) dg_kern = dg if USE_G else q.new_empty(NK, B, T, HV, dtype=torch.float32) if IS_VARLEN: - kernel(q, k, v, g_kern, h_flat, do, dh_flat, dq, dk, dw_kern, dv_kern, dg_kern, - cu_seqlens.int(), chunk_indices.int(), scale) + kernel( + q, + k, + v, + g_kern, + h_flat, + do, + dh_flat, + dq, + dk, + dw_kern, + dv_kern, + dg_kern, + cu_seqlens.int(), + chunk_indices.int(), + scale, + ) else: kernel(q, k, v, g_kern, h_flat, do, dh_flat, dq, dk, dw_kern, dv_kern, dg_kern, scale) diff --git a/build/torch-cuda/ops/common/backends/tilelang/parallel_attn_bwd.py b/build/torch-cuda/ops/common/backends/tilelang/parallel_attn_bwd.py deleted file mode 100644 index efab180746a6a3558c464d73b0d3c9c607ac85e1..0000000000000000000000000000000000000000 --- a/build/torch-cuda/ops/common/backends/tilelang/parallel_attn_bwd.py +++ /dev/null @@ -1,300 +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 tilelang -import tilelang.language as T -import torch - -from .....ops.utils import prepare_chunk_indices -from .....ops.utils.constant import RCP_LN2 - - -@tilelang.jit(pass_configs={ - tilelang.PassConfigKey.TL_DISABLE_TMA_LOWER: True, - tilelang.PassConfigKey.TL_ENABLE_FAST_MATH: True, -}) -def _build_parallel_attn_bwd_kernel( - B, HQ, H, K, V, BT, sm_scale, log2e_scale, dtype_str, - USE_G=False, USE_WINDOW=False, WINDOW_SIZE=0, IS_VARLEN=False, num_warps=4, -): - """Build the TileLang JIT kernel for parallel attention backward. - - Note: B (batch size) is baked into the compiled kernel as a compile-time - constant because it appears in static tensor shape annotations. Varying B - at runtime will trigger a recompilation. Callers should either keep B fixed - per process or pad batches to a fixed size. - """ - - dtype_map = {'float16': T.float16, 'bfloat16': T.bfloat16, 'float32': T.float32} - _dtype = dtype_map[dtype_str] - accum_dtype = T.float32 - threads = num_warps * 32 - _BT = BT - _B, _HQ, _H = B, HQ, H - _K_orig, _V_orig = K, V - # Pad to next multiple of 16 for WGMMA inner-dim constraints (shared mem only) - _K = (K + 15) // 16 * 16 - _V = (V + 15) // 16 * 16 - _G = HQ // H - _USE_G = USE_G - _USE_WINDOW = USE_WINDOW - _W = WINDOW_SIZE - _sm_scale = sm_scale - _log2e_scale = log2e_scale - - T_d, NT_d, Ncu_d = T.dynamic("T, NT, Ncu") - - q_s = (_B, T_d, _HQ, _K_orig) - kv_s = (_B, T_d, _H, _K_orig) - vv_s = (_B, T_d, _H, _V_orig) - lse_s = (_B, T_d, _HQ) - dq_s = (_B, T_d, _HQ, _K_orig) - dkv_s = (_B, T_d, _H, _K_orig) - dvv_s = (_B, T_d, _H, _V_orig) - - @T.macro - def kernel_body(q, k, v, g, lse, delta, do, dq_out, dk_out, dv_out, dg, - i_b, i_h, i_t, t_s, T_seq, bos): - i_hkv = i_h // _G - - K_shared = T.alloc_shared((_BT, _K), _dtype) - V_shared = T.alloc_shared((_BT, _V), _dtype) - T.copy(k[i_b, t_s:t_s + _BT, i_hkv, :], K_shared) - T.copy(v[i_b, t_s:t_s + _BT, i_hkv, :], V_shared) - - if _USE_G: - g_k_shared = T.alloc_shared((_BT,), accum_dtype) - T.copy(g[i_b, t_s:t_s + _BT, i_h], g_k_shared) - - qkT = T.alloc_fragment((_BT, _BT), accum_dtype) - qkT_cast = T.alloc_fragment((_BT, _BT), _dtype) - dsT = T.alloc_fragment((_BT, _BT), accum_dtype) - dsT_cast = T.alloc_fragment((_BT, _BT), _dtype) - ds = T.alloc_fragment((_BT, _BT), _dtype) - - dv_local = T.alloc_fragment((_BT, _V), accum_dtype) - dk_local = T.alloc_fragment((_BT, _K), accum_dtype) - T.clear(dv_local) - T.clear(dk_local) - - if _USE_G: - dg_k_acc = T.alloc_fragment((_BT,), accum_dtype) - T.clear(dg_k_acc) - - # Allocate shared/fragment buffers outside the loop so they can be reused - # across query-tile iterations. - q_shared = T.alloc_shared((_BT, _K), _dtype) - do_shared = T.alloc_shared((_BT, _V), _dtype) - lse_shared = T.alloc_shared((_BT,), accum_dtype) - delta_shared = T.alloc_shared((_BT,), accum_dtype) - dsT_shared = T.alloc_shared((_BT, _BT), _dtype) - dq_local = T.alloc_fragment((_BT, _K), accum_dtype) - if _USE_G: - g_q_shared = T.alloc_shared((_BT,), accum_dtype) - f_dg_q = T.alloc_fragment((_BT,), accum_dtype) - f_dg_k = T.alloc_fragment((_BT,), accum_dtype) - - i_t_local = (t_s - bos) // _BT - loop_st = i_t_local - if _USE_WINDOW: - # For a K tile at positions [t_s, t_s+BT), the farthest query - # that can still attend (under window W) is at position - # (t_s+BT-1) + (W-1). Round up to the enclosing Q tile. - loop_ed_swa = i_t_local + 1 + (_W + _BT - 2) // _BT - loop_ed = T.min(T.ceildiv(T_seq, _BT), loop_ed_swa) - else: - loop_ed = T.ceildiv(T_seq, _BT) - - # Note: T.Pipelined is intentionally avoided here because combining it - # with standard-layout T.atomic_add into dq_out produces garbage values - # (observed as inf/575+ diffs) for non-power-of-2 dimensions. - for k_local in T.serial(loop_st, loop_ed): - q_t_s = bos + k_local * _BT - - T.copy(q[i_b, q_t_s:q_t_s + _BT, i_h, :], q_shared) - T.copy(do[i_b, q_t_s:q_t_s + _BT, i_h, :], do_shared) - T.copy(lse[i_b, q_t_s:q_t_s + _BT, i_h], lse_shared) - T.copy(delta[i_b, q_t_s:q_t_s + _BT, i_h], delta_shared) - - if _USE_G: - T.copy(g[i_b, q_t_s:q_t_s + _BT, i_h], g_q_shared) - - T.clear(qkT) - T.gemm(K_shared, q_shared, qkT, transpose_B=True, policy=T.GemmWarpPolicy.FullRow) - - for i, j in T.Parallel(_BT, _BT): - if _USE_G: - qkT[i, j] = T.exp2(qkT[i, j] * _log2e_scale + g_q_shared[j] - g_k_shared[i] - lse_shared[j]) - else: - qkT[i, j] = T.exp2(qkT[i, j] * _log2e_scale - lse_shared[j]) - - # causal + boundary (+ optional sliding window) mask - for i, j in T.Parallel(_BT, _BT): - key_pos = t_s + i - q_pos = q_t_s + j - causal = (key_pos <= q_pos) & (key_pos < T_seq + bos) & (q_pos < T_seq + bos) - if _USE_WINDOW: - in_window = (q_pos - key_pos < _W) - mask = causal & in_window - else: - mask = causal - qkT[i, j] = T.if_then_else(mask, qkT[i, j], T.cast(0, accum_dtype)) - - # dv += p @ do - T.copy(qkT, qkT_cast) - T.gemm(qkT_cast, do_shared, dv_local, policy=T.GemmWarpPolicy.FullRow) - - # ds = p * (dp - delta) - T.clear(dsT) - T.gemm(V_shared, do_shared, dsT, transpose_B=True, policy=T.GemmWarpPolicy.FullRow) - - for i, j in T.Parallel(_BT, _BT): - dsT[i, j] = qkT[i, j] * (dsT[i, j] - delta_shared[j]) - for i, j in T.Parallel(_BT, _BT): - dsT_cast[i, j] = dsT[i, j] * _sm_scale - - # dk += ds^T @ q - T.copy(dsT_cast, ds) - T.gemm(ds, q_shared, dk_local, policy=T.GemmWarpPolicy.FullRow) - - # dq += ds @ k (atomic_add into standard-layout output) - T.copy(dsT_cast, dsT_shared) - T.clear(dq_local) - T.gemm(dsT_shared, K_shared, dq_local, transpose_A=True, policy=T.GemmWarpPolicy.FullRow) - for i, d in T.Parallel(_BT, _K_orig): - if q_t_s + i < T_seq + bos: - T.atomic_add(dq_out[i_b, q_t_s + i, i_h, d], dq_local[i, d]) - - if _USE_G: - # dg for query positions (positive) - T.reduce_sum(dsT, f_dg_q, dim=0) - for j in T.Parallel(_BT): - if q_t_s + j < T_seq + bos: - T.atomic_add(dg[i_b, q_t_s + j, i_h], f_dg_q[j]) - - # accumulate dg for key positions (negative) - T.reduce_sum(dsT, f_dg_k, dim=1) - for i in T.Parallel(_BT): - dg_k_acc[i] += f_dg_k[i] - - # store dv and dk via element-wise atomic_add with explicit boundary checks. - # atomic_add is required because in GQA multiple query-head blocks map to - # the same KV head and accumulate into the same dk/dv rows. - for i, d in T.Parallel(_BT, _V_orig): - if t_s + i < T_seq + bos: - T.atomic_add(dv_out[i_b, t_s + i, i_hkv, d], dv_local[i, d]) - for i, d in T.Parallel(_BT, _K_orig): - if t_s + i < T_seq + bos: - T.atomic_add(dk_out[i_b, t_s + i, i_hkv, d], dk_local[i, d]) - - if _USE_G: - for i in T.Parallel(_BT): - if t_s + i < T_seq + bos: - T.atomic_add(dg[i_b, t_s + i, i_h], -dg_k_acc[i]) - - if IS_VARLEN: - @T.prim_func - def kernel( - q: T.Tensor(q_s, _dtype), k: T.Tensor(kv_s, _dtype), - v: T.Tensor(vv_s, _dtype), g: T.Tensor(lse_s, accum_dtype), - lse: T.Tensor(lse_s, accum_dtype), delta: T.Tensor(lse_s, accum_dtype), - do: T.Tensor(q_s, _dtype), dq_out: T.Tensor(dq_s, accum_dtype), - dk_out: T.Tensor(dkv_s, accum_dtype), dv_out: T.Tensor(dvv_s, accum_dtype), - dg: T.Tensor(lse_s, accum_dtype), - cu_seqlens: T.Tensor((Ncu_d,), T.int32), - chunk_indices: T.Tensor((NT_d, 2), T.int32), - ): - with T.Kernel(_HQ, NT_d, 1, threads=threads) as (i_h, i_t, _): - i_n = chunk_indices[i_t, 0] - i_t_local = chunk_indices[i_t, 1] - bos = cu_seqlens[i_n] - T_seq = cu_seqlens[i_n + 1] - bos - t_s = bos + i_t_local * _BT - kernel_body(q, k, v, g, lse, delta, do, dq_out, dk_out, dv_out, dg, - 0, i_h, i_t, t_s, T_seq, bos) - else: - @T.prim_func - def kernel( - q: T.Tensor(q_s, _dtype), k: T.Tensor(kv_s, _dtype), - v: T.Tensor(vv_s, _dtype), g: T.Tensor(lse_s, accum_dtype), - lse: T.Tensor(lse_s, accum_dtype), delta: T.Tensor(lse_s, accum_dtype), - do: T.Tensor(q_s, _dtype), dq_out: T.Tensor(dq_s, accum_dtype), - dk_out: T.Tensor(dkv_s, accum_dtype), dv_out: T.Tensor(dvv_s, accum_dtype), - dg: T.Tensor(lse_s, accum_dtype), - ): - with T.Kernel(_HQ, T.ceildiv(T_d, _BT), _B, threads=threads) as (i_h, i_t, i_b): - t_s = i_t * _BT - kernel_body(q, k, v, g, lse, delta, do, dq_out, dk_out, dv_out, dg, - i_b, i_h, i_t, t_s, T_d, 0) - - return kernel - - -def parallel_attn_bwd_tilelang( - q, k, v, o, g_cumsum, lse, do, - sink_bias=None, scale=None, window_size=None, - chunk_size=128, cu_seqlens=None, chunk_indices=None, -): - B, T, HQ, K = q.shape - H = k.shape[2] - V = v.shape[-1] - BT = 64 - sm_scale = scale if scale is not None else K ** -0.5 - log2e_scale = sm_scale * RCP_LN2 - - USE_G = g_cumsum is not None - USE_WINDOW = window_size is not None - IS_VARLEN = cu_seqlens is not None - - if chunk_indices is None and IS_VARLEN: - chunk_indices = prepare_chunk_indices(cu_seqlens, BT) - - # reuse Triton preprocess for delta (works with unpadded tensors) - from .....ops.attn.parallel import parallel_attn_bwd_preprocess - delta = parallel_attn_bwd_preprocess(o, do) - - # No padding on the last dim: TileLang T.copy auto-zero-fills OOB columns - # when loading to shared memory, and auto-skips OOB writes. dq/dk/dv all - # need float32 buffers because GQA has multiple query-head blocks writing - # to the same KV head (dk/dv) and query rows receive contributions from - # multiple key tiles (dq). - dq = torch.zeros(B, T, HQ, K, dtype=torch.float32, device=q.device) - dk = torch.zeros(B, T, H, K, dtype=torch.float32, device=k.device) - dv = torch.zeros(B, T, H, V, dtype=torch.float32, device=v.device) - dg = torch.zeros(B, T, HQ, dtype=torch.float32, device=q.device) if USE_G else None - - g_kern = g_cumsum.float() if USE_G else torch.zeros(B, T, HQ, dtype=torch.float32, device=q.device) - dg_kern = dg if USE_G else torch.zeros(B, T, HQ, dtype=torch.float32, device=q.device) - - dtype_str = {torch.float16: 'float16', torch.bfloat16: 'bfloat16', torch.float32: 'float32'}.get(q.dtype) - if dtype_str is None: - raise ValueError(f"Unsupported dtype {q.dtype} for TileLang backend") - - kernel = _build_parallel_attn_bwd_kernel( - B, HQ, H, K, V, BT, sm_scale, log2e_scale, dtype_str, - USE_G=USE_G, USE_WINDOW=USE_WINDOW, WINDOW_SIZE=window_size or 0, - IS_VARLEN=IS_VARLEN, - ) - - if IS_VARLEN: - kernel(q, k, v, g_kern, lse, delta, do, dq, dk, dv, dg_kern, - cu_seqlens.int(), chunk_indices.int()) - else: - kernel(q, k, v, g_kern, lse, delta, do, dq, dk, dv, dg_kern) - - # sink_bias gradient is derived from the saved lse (which already includes - # the sink mass in the softmax normalizer) and delta — no kernel change - # needed. Matches the Triton reference. - dsink_bias = None - if sink_bias is not None: - p_sink = torch.exp2(sink_bias.float()[None, None, :] - lse.float()) - dsink_bias = -(p_sink * delta.float()).sum((0, 1)) - - dq = dq.to(q.dtype) - dk = dk.to(k.dtype) - dv = dv.to(v.dtype) - return dq, dk, dv, dg, dsink_bias diff --git a/build/torch-cuda/ops/common/backends/tilelang/parallel_attn_fwd.py b/build/torch-cuda/ops/common/backends/tilelang/parallel_attn_fwd.py deleted file mode 100644 index 9ff1bbe7978d9275513b9ccfef584f809c18e17b..0000000000000000000000000000000000000000 --- a/build/torch-cuda/ops/common/backends/tilelang/parallel_attn_fwd.py +++ /dev/null @@ -1,298 +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 - -"""TileLang forward kernel for parallel (causal) attention. - -Supports: GQA, gating (g_cumsum), sliding-window attention, sink_bias, -variable-length (cu_seqlens). Output format matches the Triton reference -at `fla.ops.attn.parallel.parallel_attn_fwd` so LSE is directly consumable -by either backend during backward. -""" - -import tilelang -import tilelang.language as T -import torch - -from .....ops.utils import prepare_chunk_indices -from .....ops.utils.constant import RCP_LN2 - - -def _pick_tile(K: int, V: int, is_hopper_plus: bool) -> tuple[int, int, int]: - """Return (BT, BS, num_warps) tuned for head dim.""" - if is_hopper_plus: - if max(K, V) <= 64: - return 128, 64, 8 - if max(K, V) <= 128: - return 64, 64, 4 - return 64, 64, 4 - # Ampere / pre-Hopper - if max(K, V) <= 64: - return 128, 64, 4 - if max(K, V) <= 128: - return 128, 64, 4 - return 64, 32, 4 - - -@tilelang.jit(pass_configs={ - tilelang.PassConfigKey.TL_DISABLE_TMA_LOWER: True, - tilelang.PassConfigKey.TL_ENABLE_FAST_MATH: True, -}) -def _build_parallel_attn_fwd_kernel( - B, HQ, H, K, V, BT, BS, sm_scale, log2e_scale, dtype_str, - USE_G=False, USE_WINDOW=False, WINDOW_SIZE=0, USE_SINK=False, - IS_VARLEN=False, num_warps=8, num_stages=1, -): - """Build the TileLang JIT kernel for parallel attention forward. - - Note: B (batch size) is baked into the compiled kernel as a compile-time - constant because it appears in static tensor shape annotations. Varying - B at runtime triggers a recompilation. Callers should keep B fixed per - process or pad batches to a fixed size. - """ - dtype_map = {'float16': T.float16, 'bfloat16': T.bfloat16, 'float32': T.float32} - _dtype = dtype_map[dtype_str] - accum_dtype = T.float32 - threads = num_warps * 32 - _BT, _BS = BT, BS - _B, _HQ, _H = B, HQ, H - _K_orig, _V_orig = K, V - # Pad shared-mem last dim to a multiple of 16 for WGMMA/MMA inner-dim. - # Same approach as the bwd kernel — tensors keep their original shape; - # only the shared buffers are padded. - _K = (K + 15) // 16 * 16 - _V = (V + 15) // 16 * 16 - _G = HQ // H - _USE_G = USE_G - _USE_WINDOW = USE_WINDOW - _USE_SINK = USE_SINK - _W = WINDOW_SIZE - _sm_scale = sm_scale - _log2e_scale = log2e_scale - - T_d, NT_d, Ncu_d = T.dynamic("T, NT, Ncu") - - q_s = (_B, T_d, _HQ, _K_orig) - kv_s = (_B, T_d, _H, _K_orig) - vv_s = (_B, T_d, _H, _V_orig) - o_s = (_B, T_d, _HQ, _V_orig) - lse_s = (_B, T_d, _HQ) - sink_s = (_HQ,) - - @T.macro - def kernel_body(q, k, v, g, sink_bias, o, lse, - i_b, i_h, i_t, t_s, T_seq, bos): - i_hkv = i_h // _G - - Q_shared = T.alloc_shared((_BT, _K), _dtype) - K_shared = T.alloc_shared((_BS, _K), _dtype) - V_shared = T.alloc_shared((_BS, _V), _dtype) - - acc_s = T.alloc_fragment((_BT, _BS), accum_dtype) - acc_s_cast = T.alloc_fragment((_BT, _BS), _dtype) - acc_o = T.alloc_fragment((_BT, _V), accum_dtype) - - m_f = T.alloc_fragment((_BT,), accum_dtype) - m_prev = T.alloc_fragment((_BT,), accum_dtype) - l_f = T.alloc_fragment((_BT,), accum_dtype) - p_sum = T.alloc_fragment((_BT,), accum_dtype) - scale_row = T.alloc_fragment((_BT,), accum_dtype) - - if _USE_G: - g_q_shared = T.alloc_shared((_BT,), accum_dtype) - g_k_shared = T.alloc_shared((_BS,), accum_dtype) - T.copy(g[i_b, t_s:t_s + _BT, i_h], g_q_shared) - - T.copy(q[i_b, t_s:t_s + _BT, i_h, :], Q_shared) - T.fill(acc_o, 0) - T.fill(l_f, 0) - T.fill(m_f, -T.infinity(accum_dtype)) - - i_t_local = (t_s - bos) // _BT - if _USE_WINDOW: - # earliest key the last query in this tile can see - # (q_pos - k_pos < W) => k_pos > q_pos - W. - # Round down to the enclosing K tile start. - s0_raw = i_t_local * _BT - _W + 1 - s0 = T.max(s0_raw, 0) - loop_st = (s0 // _BS) - else: - loop_st = 0 - loop_ed_raw = T.ceildiv((i_t_local + 1) * _BT, _BS) - loop_ed = T.min(loop_ed_raw, T.ceildiv(T_seq, _BS)) - - for k_idx in T.Pipelined(loop_st, loop_ed, num_stages=num_stages): - k_t_s = bos + k_idx * _BS - - T.copy(k[i_b, k_t_s:k_t_s + _BS, i_hkv, :], K_shared) - T.copy(v[i_b, k_t_s:k_t_s + _BS, i_hkv, :], V_shared) - if _USE_G: - T.copy(g[i_b, k_t_s:k_t_s + _BS, i_h], g_k_shared) - - T.clear(acc_s) - T.gemm(Q_shared, K_shared, acc_s, - transpose_B=True, policy=T.GemmWarpPolicy.FullRow) - - # Apply mask (causal + boundary + optional SWA), add log-gates, - # and convert scores to the log2 domain. Masked entries get - # large-negative so exp2(...) -> 0. - for i, j in T.Parallel(_BT, _BS): - q_pos = t_s + i - key_pos = k_t_s + j - causal = (key_pos <= q_pos) & (key_pos < T_seq + bos) & (q_pos < T_seq + bos) - if _USE_WINDOW: - mask = causal & (q_pos - key_pos < _W) - else: - mask = causal - if _USE_G: - s = acc_s[i, j] * _log2e_scale + g_q_shared[i] - g_k_shared[j] - else: - s = acc_s[i, j] * _log2e_scale - acc_s[i, j] = T.if_then_else(mask, s, -T.infinity(accum_dtype)) - - # Online softmax update. - T.copy(m_f, m_prev) - T.reduce_max(acc_s, m_f, dim=1, clear=False) - for i in T.Parallel(_BT): - m_f[i] = T.max(m_f[i], m_prev[i]) - # Guard against all-masked rows: use a finite pivot so we - # don't produce NaN via (-inf) - (-inf) anywhere downstream. - m_stable = T.alloc_fragment((_BT,), accum_dtype) - for i in T.Parallel(_BT): - m_stable[i] = T.if_then_else( - m_f[i] == -T.infinity(accum_dtype), 0., m_f[i]) - for i in T.Parallel(_BT): - scale_row[i] = T.if_then_else( - m_prev[i] == -T.infinity(accum_dtype), 0., - T.exp2(m_prev[i] - m_stable[i])) - for i, j in T.Parallel(_BT, _V): - acc_o[i, j] *= scale_row[i] - for i, j in T.Parallel(_BT, _BS): - acc_s[i, j] = T.exp2(acc_s[i, j] - m_stable[i]) - T.reduce_sum(acc_s, p_sum, dim=1) - for i in T.Parallel(_BT): - l_f[i] = l_f[i] * scale_row[i] + p_sum[i] - - T.copy(acc_s, acc_s_cast) - T.gemm(acc_s_cast, V_shared, acc_o, - policy=T.GemmWarpPolicy.FullRow) - - # sink-bias contributes to the normalizer but not to the value matmul - if _USE_SINK: - # If a row has no valid keys yet, m_f is still -inf — pin it to 0 - # so `sink_bias - m_f` is finite. - for i in T.Parallel(_BT): - m_f[i] = T.if_then_else( - m_f[i] == -T.infinity(accum_dtype), 0., m_f[i]) - for i in T.Parallel(_BT): - l_f[i] += T.exp2(sink_bias[i_h] - m_f[i]) - - # finalize O and LSE - O_shared = T.alloc_shared((_BT, _V), _dtype) - lse_shared = T.alloc_shared((_BT,), accum_dtype) - for i, d in T.Parallel(_BT, _V): - acc_o[i, d] = acc_o[i, d] / l_f[i] - T.copy(acc_o, O_shared) - for i in T.Parallel(_BT): - lse_shared[i] = m_f[i] + T.log2(l_f[i]) - for i, d in T.Parallel(_BT, _V_orig): - if t_s + i < T_seq + bos: - o[i_b, t_s + i, i_h, d] = O_shared[i, d] - for i in T.Parallel(_BT): - if t_s + i < T_seq + bos: - lse[i_b, t_s + i, i_h] = lse_shared[i] - - if IS_VARLEN: - @T.prim_func - def kernel( - q: T.Tensor(q_s, _dtype), k: T.Tensor(kv_s, _dtype), - v: T.Tensor(vv_s, _dtype), g: T.Tensor(lse_s, accum_dtype), - sink_bias: T.Tensor(sink_s, accum_dtype), - o: T.Tensor(o_s, _dtype), lse: T.Tensor(lse_s, accum_dtype), - cu_seqlens: T.Tensor((Ncu_d,), T.int32), - chunk_indices: T.Tensor((NT_d, 2), T.int32), - ): - with T.Kernel(_HQ, NT_d, 1, threads=threads) as (i_h, i_t, _): - i_n = chunk_indices[i_t, 0] - i_t_local = chunk_indices[i_t, 1] - bos = cu_seqlens[i_n] - eos = cu_seqlens[i_n + 1] - T_seq = eos - bos - t_s = bos + i_t_local * _BT - kernel_body(q, k, v, g, sink_bias, o, lse, - 0, i_h, i_t, t_s, T_seq, bos) - else: - @T.prim_func - def kernel( - q: T.Tensor(q_s, _dtype), k: T.Tensor(kv_s, _dtype), - v: T.Tensor(vv_s, _dtype), g: T.Tensor(lse_s, accum_dtype), - sink_bias: T.Tensor(sink_s, accum_dtype), - o: T.Tensor(o_s, _dtype), lse: T.Tensor(lse_s, accum_dtype), - ): - with T.Kernel(_HQ, T.ceildiv(T_d, _BT), _B, threads=threads) as (i_h, i_t, i_b): - t_s = i_t * _BT - kernel_body(q, k, v, g, sink_bias, o, lse, - i_b, i_h, i_t, t_s, T_d, 0) - - return kernel - - -def parallel_attn_fwd_tilelang( - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - g_cumsum: torch.Tensor | None, - sink_bias: torch.Tensor | None, - scale: float, - window_size: int | None = None, - cu_seqlens: torch.LongTensor | None = None, - chunk_indices: torch.LongTensor | None = None, -): - B, T, HQ, K = q.shape - H = k.shape[2] - V = v.shape[-1] - - USE_G = g_cumsum is not None - USE_WINDOW = window_size is not None - USE_SINK = sink_bias is not None - IS_VARLEN = cu_seqlens is not None - - is_hopper_plus = torch.cuda.get_device_capability()[0] >= 9 - BT, BS, num_warps = _pick_tile(K, V, is_hopper_plus) - - if chunk_indices is None and IS_VARLEN: - chunk_indices = prepare_chunk_indices(cu_seqlens, BT) - - sm_scale = scale if scale is not None else K ** -0.5 - log2e_scale = sm_scale * RCP_LN2 - - dtype_str = {torch.float16: 'float16', torch.bfloat16: 'bfloat16', - torch.float32: 'float32'}.get(q.dtype) - if dtype_str is None: - raise ValueError(f"Unsupported dtype {q.dtype} for TileLang backend") - - o = torch.empty(B, T, HQ, V, dtype=q.dtype, device=q.device) - lse = torch.empty(B, T, HQ, dtype=torch.float32, device=q.device) - - g_kern = g_cumsum.float() if USE_G else torch.zeros( - B, T, HQ, dtype=torch.float32, device=q.device) - sink_kern = sink_bias.float() if USE_SINK else torch.zeros( - HQ, dtype=torch.float32, device=q.device) - - kernel = _build_parallel_attn_fwd_kernel( - B, HQ, H, K, V, BT, BS, sm_scale, log2e_scale, dtype_str, - USE_G=USE_G, USE_WINDOW=USE_WINDOW, WINDOW_SIZE=window_size or 0, - USE_SINK=USE_SINK, IS_VARLEN=IS_VARLEN, num_warps=num_warps, - num_stages=1 if USE_WINDOW else 2, - ) - - if IS_VARLEN: - kernel(q, k, v, g_kern, sink_kern, o, lse, - cu_seqlens.int(), chunk_indices.int()) - else: - kernel(q, k, v, g_kern, sink_kern, o, lse) - - return o, lse diff --git a/build/torch-cuda/ops/common/chunk_delta_h.py b/build/torch-cuda/ops/common/chunk_delta_h.py index 89a556630064549d9d4156d559ea3feed8b2832b..58dd3f62d11daf3177aea02de0db4ac0f70f74dc 100644 --- a/build/torch-cuda/ops/common/chunk_delta_h.py +++ b/build/torch-cuda/ops/common/chunk_delta_h.py @@ -9,9 +9,10 @@ import torch import triton import triton.language as tl -from ...ops.backends import dispatch +from ...backends import dispatch from ...ops.utils import prepare_chunk_indices, prepare_chunk_offsets from ...ops.utils.cache import fla_cache_autotune +from ...ops.utils.graph import get_static_buffer from ...ops.utils.op import exp2 from ...utils import ( IS_INTEL, @@ -82,13 +83,15 @@ def chunk_gated_delta_rule_fwd_kernel_h_blockdim64( STATE_V_FIRST: tl.constexpr, IS_VARLEN: tl.constexpr, ): - i_v, i_nh = tl.program_id(0), tl.program_id(1) + pid = tl.program_id(0).to(tl.int64) + NV = tl.cdiv(V, BV) + i_v, i_nh = pid % NV, (pid // NV).to(tl.int64) i_n, i_h = i_nh // HV, i_nh % HV if IS_VARLEN: - bos, eos = tl.load(cu_seqlens + i_n).to(tl.int32), tl.load(cu_seqlens + i_n + 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 NT = tl.cdiv(T, BT) - boh = tl.load(chunk_offsets + i_n).to(tl.int32) + boh = tl.load(chunk_offsets + i_n).to(tl.int64) else: bos, eos = i_n * T, i_n * T + T NT = tl.cdiv(T, BT) @@ -125,98 +128,125 @@ def chunk_gated_delta_rule_fwd_kernel_h_blockdim64( ht = ht + i_nh * K*V # load initial state + o_v = i_v * BV + tl.arange(0, BV) + m_v = o_v < V + o_k1 = tl.arange(0, 64) + m_k1 = o_k1 < K + o_k2 = 64 + o_k1 + m_k2 = o_k2 < K + o_k3 = 128 + o_k1 + m_k3 = o_k3 < K + o_k4 = 192 + o_k1 + m_k4 = o_k4 < K if USE_INITIAL_STATE: if STATE_V_FIRST: - p_h0_1 = tl.make_block_ptr(h0, (V, K), (K, 1), (i_v * BV, 0), (BV, 64), (1, 0)) + p_h0_1 = h0 + o_v[:, None] * K + o_k1[None, :] + m_h0_1 = m_v[:, None] & m_k1[None, :] else: - p_h0_1 = tl.make_block_ptr(h0, (K, V), (V, 1), (0, i_v * BV), (64, BV), (1, 0)) - b_h1 += tl.load(p_h0_1, boundary_check=(0, 1)).to(tl.float32) + p_h0_1 = h0 + o_k1[:, None] * V + o_v[None, :] + m_h0_1 = m_k1[:, None] & m_v[None, :] + b_h1 += tl.load(p_h0_1, mask=m_h0_1, other=0.0).to(tl.float32) if K > 64: if STATE_V_FIRST: - p_h0_2 = tl.make_block_ptr(h0, (V, K), (K, 1), (i_v * BV, 64), (BV, 64), (1, 0)) + p_h0_2 = h0 + o_v[:, None] * K + o_k2[None, :] + m_h0_2 = m_v[:, None] & m_k2[None, :] else: - p_h0_2 = tl.make_block_ptr(h0, (K, V), (V, 1), (64, i_v * BV), (64, BV), (1, 0)) - b_h2 += tl.load(p_h0_2, boundary_check=(0, 1)).to(tl.float32) + p_h0_2 = h0 + o_k2[:, None] * V + o_v[None, :] + m_h0_2 = m_k2[:, None] & m_v[None, :] + b_h2 += tl.load(p_h0_2, mask=m_h0_2, other=0.0).to(tl.float32) if K > 128: if STATE_V_FIRST: - p_h0_3 = tl.make_block_ptr(h0, (V, K), (K, 1), (i_v * BV, 128), (BV, 64), (1, 0)) + p_h0_3 = h0 + o_v[:, None] * K + o_k3[None, :] + m_h0_3 = m_v[:, None] & m_k3[None, :] else: - p_h0_3 = tl.make_block_ptr(h0, (K, V), (V, 1), (128, i_v * BV), (64, BV), (1, 0)) - b_h3 += tl.load(p_h0_3, boundary_check=(0, 1)).to(tl.float32) + p_h0_3 = h0 + o_k3[:, None] * V + o_v[None, :] + m_h0_3 = m_k3[:, None] & m_v[None, :] + b_h3 += tl.load(p_h0_3, mask=m_h0_3, other=0.0).to(tl.float32) if K > 192: if STATE_V_FIRST: - p_h0_4 = tl.make_block_ptr(h0, (V, K), (K, 1), (i_v * BV, 192), (BV, 64), (1, 0)) + p_h0_4 = h0 + o_v[:, None] * K + o_k4[None, :] + m_h0_4 = m_v[:, None] & m_k4[None, :] else: - p_h0_4 = tl.make_block_ptr(h0, (K, V), (V, 1), (192, i_v * BV), (64, BV), (1, 0)) - b_h4 += tl.load(p_h0_4, boundary_check=(0, 1)).to(tl.float32) + p_h0_4 = h0 + o_k4[:, None] * V + o_v[None, :] + m_h0_4 = m_k4[:, None] & m_v[None, :] + b_h4 += tl.load(p_h0_4, mask=m_h0_4, other=0.0).to(tl.float32) # main recurrence for i_t in range(NT): i_t_int64 = i_t.to(tl.int64) + o_t = i_t_int64 * BT + tl.arange(0, BT) + m_t = o_t < T if STATE_V_FIRST: - p_h1 = tl.make_block_ptr(h + i_t_int64 * HV*K*V, (V, K), (K, 1), (i_v * BV, 0), (BV, 64), (1, 0)) + p_h1 = h + i_t_int64 * HV*K*V + o_v[:, None] * K + o_k1[None, :] + m_h1 = m_v[:, None] & m_k1[None, :] else: - p_h1 = tl.make_block_ptr(h + i_t_int64 * HV*K*V, (K, V), (V, 1), (0, i_v * BV), (64, BV), (1, 0)) - tl.store(p_h1, b_h1.to(p_h1.dtype.element_ty), boundary_check=(0, 1)) + p_h1 = h + i_t_int64 * HV*K*V + o_k1[:, None] * V + o_v[None, :] + m_h1 = m_k1[:, None] & m_v[None, :] + tl.store(p_h1, b_h1.to(p_h1.dtype.element_ty), mask=m_h1) if K > 64: if STATE_V_FIRST: - p_h2 = tl.make_block_ptr(h + i_t_int64 * HV*K*V, (V, K), (K, 1), (i_v * BV, 64), (BV, 64), (1, 0)) + p_h2 = h + i_t_int64 * HV*K*V + o_v[:, None] * K + o_k2[None, :] + m_h2 = m_v[:, None] & m_k2[None, :] else: - p_h2 = tl.make_block_ptr(h + i_t_int64 * HV*K*V, (K, V), (V, 1), (64, i_v * BV), (64, BV), (1, 0)) - tl.store(p_h2, b_h2.to(p_h2.dtype.element_ty), boundary_check=(0, 1)) + p_h2 = h + i_t_int64 * HV*K*V + o_k2[:, None] * V + o_v[None, :] + m_h2 = m_k2[:, None] & m_v[None, :] + tl.store(p_h2, b_h2.to(p_h2.dtype.element_ty), mask=m_h2) if K > 128: if STATE_V_FIRST: - p_h3 = tl.make_block_ptr(h + i_t_int64 * HV*K*V, (V, K), (K, 1), (i_v * BV, 128), (BV, 64), (1, 0)) + p_h3 = h + i_t_int64 * HV*K*V + o_v[:, None] * K + o_k3[None, :] + m_h3 = m_v[:, None] & m_k3[None, :] else: - p_h3 = tl.make_block_ptr(h + i_t_int64 * HV*K*V, (K, V), (V, 1), (128, i_v * BV), (64, BV), (1, 0)) - tl.store(p_h3, b_h3.to(p_h3.dtype.element_ty), boundary_check=(0, 1)) + p_h3 = h + i_t_int64 * HV*K*V + o_k3[:, None] * V + o_v[None, :] + m_h3 = m_k3[:, None] & m_v[None, :] + tl.store(p_h3, b_h3.to(p_h3.dtype.element_ty), mask=m_h3) if K > 192: if STATE_V_FIRST: - p_h4 = tl.make_block_ptr(h + i_t_int64 * HV*K*V, (V, K), (K, 1), (i_v * BV, 192), (BV, 64), (1, 0)) + p_h4 = h + i_t_int64 * HV*K*V + o_v[:, None] * K + o_k4[None, :] + m_h4 = m_v[:, None] & m_k4[None, :] else: - p_h4 = tl.make_block_ptr(h + i_t_int64 * HV*K*V, (K, V), (V, 1), (192, i_v * BV), (64, BV), (1, 0)) - tl.store(p_h4, b_h4.to(p_h4.dtype.element_ty), boundary_check=(0, 1)) + p_h4 = h + i_t_int64 * HV*K*V + o_k4[:, None] * V + o_v[None, :] + m_h4 = m_k4[:, None] & m_v[None, :] + tl.store(p_h4, b_h4.to(p_h4.dtype.element_ty), mask=m_h4) - p_w = tl.make_block_ptr(w, (T, K), (HV*K, 1), (i_t * BT, 0), (BT, 64), (1, 0)) - b_w = tl.load(p_w, boundary_check=(0, 1)) + p_w = w + o_t[:, None] * (HV*K) + o_k1[None, :] + b_w = tl.load(p_w, mask=m_t[:, None] & m_k1[None, :], other=0.0) if STATE_V_FIRST: b_v = tl.dot(b_w, tl.trans(b_h1).to(b_w.dtype)) else: b_v = tl.dot(b_w, b_h1.to(b_w.dtype)) if K > 64: - p_w = tl.make_block_ptr(w, (T, K), (HV*K, 1), (i_t * BT, 64), (BT, 64), (1, 0)) - b_w = tl.load(p_w, boundary_check=(0, 1)) + p_w = w + o_t[:, None] * (HV*K) + o_k2[None, :] + b_w = tl.load(p_w, mask=m_t[:, None] & m_k2[None, :], other=0.0) if STATE_V_FIRST: - b_v += tl.dot(b_w, tl.trans(b_h2).to(b_w.dtype)) + b_v = tl.dot(b_w, tl.trans(b_h2).to(b_w.dtype), b_v) else: - b_v += tl.dot(b_w, b_h2.to(b_w.dtype)) + b_v = tl.dot(b_w, b_h2.to(b_w.dtype), b_v) if K > 128: - p_w = tl.make_block_ptr(w, (T, K), (HV*K, 1), (i_t * BT, 128), (BT, 64), (1, 0)) - b_w = tl.load(p_w, boundary_check=(0, 1)) + p_w = w + o_t[:, None] * (HV*K) + o_k3[None, :] + b_w = tl.load(p_w, mask=m_t[:, None] & m_k3[None, :], other=0.0) if STATE_V_FIRST: - b_v += tl.dot(b_w, tl.trans(b_h3).to(b_w.dtype)) + b_v = tl.dot(b_w, tl.trans(b_h3).to(b_w.dtype), b_v) else: - b_v += tl.dot(b_w, b_h3.to(b_w.dtype)) + b_v = tl.dot(b_w, b_h3.to(b_w.dtype), b_v) if K > 192: - p_w = tl.make_block_ptr(w, (T, K), (HV*K, 1), (i_t * BT, 192), (BT, 64), (1, 0)) - b_w = tl.load(p_w, boundary_check=(0, 1)) + p_w = w + o_t[:, None] * (HV*K) + o_k4[None, :] + b_w = tl.load(p_w, mask=m_t[:, None] & m_k4[None, :], other=0.0) if STATE_V_FIRST: - b_v += tl.dot(b_w, tl.trans(b_h4).to(b_w.dtype)) + b_v = tl.dot(b_w, tl.trans(b_h4).to(b_w.dtype), b_v) else: - b_v += tl.dot(b_w, b_h4.to(b_w.dtype)) - p_v = tl.make_block_ptr(v, (T, V), (HV*V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) - b_v = tl.load(p_v, boundary_check=(0, 1)) - b_v + b_v = tl.dot(b_w, b_h4.to(b_w.dtype), b_v) + p_v = v + o_t[:, None] * (HV*V) + o_v[None, :] + b_v = tl.load(p_v, mask=m_t[:, None] & m_v[None, :], other=0.0) - b_v if SAVE_NEW_VALUE: - p_v = tl.make_block_ptr(v_new, (T, V), (HV*V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) - tl.store(p_v, b_v.to(p_v.dtype.element_ty), boundary_check=(0, 1)) + p_v = v_new + o_t[:, None] * (HV*V) + o_v[None, :] + tl.store(p_v, b_v.to(p_v.dtype.element_ty), mask=m_t[:, None] & m_v[None, :]) last_idx = min((i_t + 1) * BT, T) - 1 if USE_G: - m_t = (i_t * BT + tl.arange(0, BT)) < T b_g_last = tl.load(g + (bos * HV + last_idx * HV + i_h).to(tl.int64)).to(tl.float32) - p_g = tl.make_block_ptr(g + (bos * HV + i_h).to(tl.int64), (T,), (HV,), (i_t * BT,), (BT,), (0,)) - b_g = tl.load(p_g, boundary_check=(0,)).to(tl.float32) + p_g = g + (bos * HV + i_h).to(tl.int64) + o_t * HV + b_g = tl.load(p_g, mask=m_t, other=0.0).to(tl.float32) b_v = b_v * tl.where(m_t, exp2(b_g_last - b_g), 0)[:, None] b_g_last = exp2(b_g_last) b_h1 *= b_g_last @@ -257,58 +287,66 @@ def chunk_gated_delta_rule_fwd_kernel_h_blockdim64( b_h4 *= exp2(b_gk_last4)[:, None] b_v = b_v.to(k.dtype.element_ty) - p_k = tl.make_block_ptr(k, (K, T), (1, H*K), (0, i_t * BT), (64, BT), (0, 1)) - b_k = tl.load(p_k, boundary_check=(0, 1)) + p_k = k + o_k1[:, None] + o_t[None, :] * (H*K) + b_k = tl.load(p_k, mask=m_k1[:, None] & m_t[None, :], other=0.0) if STATE_V_FIRST: b_h1 += tl.trans(tl.dot(b_k, b_v)) else: - b_h1 += tl.dot(b_k, b_v) + b_h1 = tl.dot(b_k, b_v, b_h1) if K > 64: - p_k = tl.make_block_ptr(k, (K, T), (1, H*K), (64, i_t * BT), (64, BT), (0, 1)) - b_k = tl.load(p_k, boundary_check=(0, 1)) + p_k = k + o_k2[:, None] + o_t[None, :] * (H*K) + b_k = tl.load(p_k, mask=m_k2[:, None] & m_t[None, :], other=0.0) if STATE_V_FIRST: b_h2 += tl.trans(tl.dot(b_k, b_v)) else: - b_h2 += tl.dot(b_k, b_v) + b_h2 = tl.dot(b_k, b_v, b_h2) if K > 128: - p_k = tl.make_block_ptr(k, (K, T), (1, H*K), (128, i_t * BT), (64, BT), (0, 1)) - b_k = tl.load(p_k, boundary_check=(0, 1)) + p_k = k + o_k3[:, None] + o_t[None, :] * (H*K) + b_k = tl.load(p_k, mask=m_k3[:, None] & m_t[None, :], other=0.0) if STATE_V_FIRST: b_h3 += tl.trans(tl.dot(b_k, b_v)) else: - b_h3 += tl.dot(b_k, b_v) + b_h3 = tl.dot(b_k, b_v, b_h3) if K > 192: - p_k = tl.make_block_ptr(k, (K, T), (1, H*K), (192, i_t * BT), (64, BT), (0, 1)) - b_k = tl.load(p_k, boundary_check=(0, 1)) + p_k = k + o_k4[:, None] + o_t[None, :] * (H*K) + b_k = tl.load(p_k, mask=m_k4[:, None] & m_t[None, :], other=0.0) if STATE_V_FIRST: b_h4 += tl.trans(tl.dot(b_k, b_v)) else: - b_h4 += tl.dot(b_k, b_v) + b_h4 = tl.dot(b_k, b_v, b_h4) if STORE_FINAL_STATE: if STATE_V_FIRST: - p_ht = tl.make_block_ptr(ht, (V, K), (K, 1), (i_v * BV, 0), (BV, 64), (1, 0)) + p_ht = ht + o_v[:, None] * K + o_k1[None, :] + m_ht = m_v[:, None] & m_k1[None, :] else: - p_ht = tl.make_block_ptr(ht, (K, V), (V, 1), (0, i_v * BV), (64, BV), (1, 0)) - tl.store(p_ht, b_h1.to(p_ht.dtype.element_ty), boundary_check=(0, 1)) + p_ht = ht + o_k1[:, None] * V + o_v[None, :] + m_ht = m_k1[:, None] & m_v[None, :] + tl.store(p_ht, b_h1.to(p_ht.dtype.element_ty), mask=m_ht) if K > 64: if STATE_V_FIRST: - p_ht = tl.make_block_ptr(ht, (V, K), (K, 1), (i_v * BV, 64), (BV, 64), (1, 0)) + p_ht = ht + o_v[:, None] * K + o_k2[None, :] + m_ht = m_v[:, None] & m_k2[None, :] else: - p_ht = tl.make_block_ptr(ht, (K, V), (V, 1), (64, i_v * BV), (64, BV), (1, 0)) - tl.store(p_ht, b_h2.to(p_ht.dtype.element_ty), boundary_check=(0, 1)) + p_ht = ht + o_k2[:, None] * V + o_v[None, :] + m_ht = m_k2[:, None] & m_v[None, :] + tl.store(p_ht, b_h2.to(p_ht.dtype.element_ty), mask=m_ht) if K > 128: if STATE_V_FIRST: - p_ht = tl.make_block_ptr(ht, (V, K), (K, 1), (i_v * BV, 128), (BV, 64), (1, 0)) + p_ht = ht + o_v[:, None] * K + o_k3[None, :] + m_ht = m_v[:, None] & m_k3[None, :] else: - p_ht = tl.make_block_ptr(ht, (K, V), (V, 1), (128, i_v * BV), (64, BV), (1, 0)) - tl.store(p_ht, b_h3.to(p_ht.dtype.element_ty), boundary_check=(0, 1)) + p_ht = ht + o_k3[:, None] * V + o_v[None, :] + m_ht = m_k3[:, None] & m_v[None, :] + tl.store(p_ht, b_h3.to(p_ht.dtype.element_ty), mask=m_ht) if K > 192: if STATE_V_FIRST: - p_ht = tl.make_block_ptr(ht, (V, K), (K, 1), (i_v * BV, 192), (BV, 64), (1, 0)) + p_ht = ht + o_v[:, None] * K + o_k4[None, :] + m_ht = m_v[:, None] & m_k4[None, :] else: - p_ht = tl.make_block_ptr(ht, (K, V), (V, 1), (192, i_v * BV), (64, BV), (1, 0)) - tl.store(p_ht, b_h4.to(p_ht.dtype.element_ty), boundary_check=(0, 1)) + p_ht = ht + o_k4[:, None] * V + o_v[None, :] + m_ht = m_k4[:, None] & m_v[None, :] + tl.store(p_ht, b_h4.to(p_ht.dtype.element_ty), mask=m_ht) @triton.heuristics({ @@ -358,13 +396,18 @@ def chunk_gated_delta_rule_bwd_kernel_dhu_blockdim64( STATE_V_FIRST: tl.constexpr, IS_VARLEN: tl.constexpr, ): - i_v, i_nh = tl.program_id(0), tl.program_id(1) + # Keep in fp32 (some pass it as fp64) which causes issues under compile + scale = scale.to(tl.float32) + + pid = tl.program_id(0).to(tl.int64) + NV = tl.cdiv(V, BV) + i_v, i_nh = pid % NV, (pid // NV).to(tl.int64) i_n, i_h = i_nh // HV, i_nh % HV if IS_VARLEN: - bos, eos = tl.load(cu_seqlens + i_n).to(tl.int32), tl.load(cu_seqlens + i_n + 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 NT = tl.cdiv(T, BT) - boh = tl.load(chunk_offsets + i_n).to(tl.int32) + boh = tl.load(chunk_offsets + i_n).to(tl.int64) else: bos, eos = i_n * T, i_n * T + T NT = tl.cdiv(T, BT) @@ -403,73 +446,101 @@ def chunk_gated_delta_rule_bwd_kernel_dhu_blockdim64( if USE_FINAL_STATE_GRADIENT: dht += i_nh * K*V + o_v = i_v * BV + tl.arange(0, BV) + m_v = o_v < V + o_k1 = tl.arange(0, 64) + m_k1 = o_k1 < K + o_k2 = 64 + o_k1 + m_k2 = o_k2 < K + o_k3 = 128 + o_k1 + m_k3 = o_k3 < K + o_k4 = 192 + o_k1 + m_k4 = o_k4 < K if USE_FINAL_STATE_GRADIENT: if STATE_V_FIRST: - p_dht1 = tl.make_block_ptr(dht, (V, K), (K, 1), (i_v * BV, 0), (BV, 64), (1, 0)) + p_dht1 = dht + o_v[:, None] * K + o_k1[None, :] + m_dht1 = m_v[:, None] & m_k1[None, :] else: - p_dht1 = tl.make_block_ptr(dht, (K, V), (V, 1), (0, i_v * BV), (64, BV), (1, 0)) - b_dh1 += tl.load(p_dht1, boundary_check=(0, 1)) + p_dht1 = dht + o_k1[:, None] * V + o_v[None, :] + m_dht1 = m_k1[:, None] & m_v[None, :] + b_dh1 += tl.load(p_dht1, mask=m_dht1, other=0.0) if K > 64: if STATE_V_FIRST: - p_dht2 = tl.make_block_ptr(dht, (V, K), (K, 1), (i_v * BV, 64), (BV, 64), (1, 0)) + p_dht2 = dht + o_v[:, None] * K + o_k2[None, :] + m_dht2 = m_v[:, None] & m_k2[None, :] else: - p_dht2 = tl.make_block_ptr(dht, (K, V), (V, 1), (64, i_v * BV), (64, BV), (1, 0)) - b_dh2 += tl.load(p_dht2, boundary_check=(0, 1)) + p_dht2 = dht + o_k2[:, None] * V + o_v[None, :] + m_dht2 = m_k2[:, None] & m_v[None, :] + b_dh2 += tl.load(p_dht2, mask=m_dht2, other=0.0) if K > 128: if STATE_V_FIRST: - p_dht3 = tl.make_block_ptr(dht, (V, K), (K, 1), (i_v * BV, 128), (BV, 64), (1, 0)) + p_dht3 = dht + o_v[:, None] * K + o_k3[None, :] + m_dht3 = m_v[:, None] & m_k3[None, :] else: - p_dht3 = tl.make_block_ptr(dht, (K, V), (V, 1), (128, i_v * BV), (64, BV), (1, 0)) - b_dh3 += tl.load(p_dht3, boundary_check=(0, 1)) + p_dht3 = dht + o_k3[:, None] * V + o_v[None, :] + m_dht3 = m_k3[:, None] & m_v[None, :] + b_dh3 += tl.load(p_dht3, mask=m_dht3, other=0.0) if K > 192: if STATE_V_FIRST: - p_dht4 = tl.make_block_ptr(dht, (V, K), (K, 1), (i_v * BV, 192), (BV, 64), (1, 0)) + p_dht4 = dht + o_v[:, None] * K + o_k4[None, :] + m_dht4 = m_v[:, None] & m_k4[None, :] else: - p_dht4 = tl.make_block_ptr(dht, (K, V), (V, 1), (192, i_v * BV), (64, BV), (1, 0)) - b_dh4 += tl.load(p_dht4, boundary_check=(0, 1)) + p_dht4 = dht + o_k4[:, None] * V + o_v[None, :] + m_dht4 = m_k4[:, None] & m_v[None, :] + b_dh4 += tl.load(p_dht4, mask=m_dht4, other=0.0) for i_t in range(NT - 1, -1, -1): i_t_int64 = i_t.to(tl.int64) + o_t = i_t_int64 * BT + tl.arange(0, BT) + m_t = o_t < T if STATE_V_FIRST: - p_dh1 = tl.make_block_ptr(dh + i_t_int64*HV*K*V, (V, K), (K, 1), (i_v * BV, 0), (BV, 64), (1, 0)) + p_dh1 = dh + i_t_int64*HV*K*V + o_v[:, None] * K + o_k1[None, :] + m_dh1 = m_v[:, None] & m_k1[None, :] else: - p_dh1 = tl.make_block_ptr(dh + i_t_int64*HV*K*V, (K, V), (V, 1), (0, i_v * BV), (64, BV), (1, 0)) - tl.store(p_dh1, b_dh1.to(p_dh1.dtype.element_ty), boundary_check=(0, 1)) + p_dh1 = dh + i_t_int64*HV*K*V + o_k1[:, None] * V + o_v[None, :] + m_dh1 = m_k1[:, None] & m_v[None, :] + tl.store(p_dh1, b_dh1.to(p_dh1.dtype.element_ty), mask=m_dh1) if K > 64: if STATE_V_FIRST: - p_dh2 = tl.make_block_ptr(dh + i_t_int64*HV*K*V, (V, K), (K, 1), (i_v * BV, 64), (BV, 64), (1, 0)) + p_dh2 = dh + i_t_int64*HV*K*V + o_v[:, None] * K + o_k2[None, :] + m_dh2 = m_v[:, None] & m_k2[None, :] else: - p_dh2 = tl.make_block_ptr(dh + i_t_int64*HV*K*V, (K, V), (V, 1), (64, i_v * BV), (64, BV), (1, 0)) - tl.store(p_dh2, b_dh2.to(p_dh2.dtype.element_ty), boundary_check=(0, 1)) + p_dh2 = dh + i_t_int64*HV*K*V + o_k2[:, None] * V + o_v[None, :] + m_dh2 = m_k2[:, None] & m_v[None, :] + tl.store(p_dh2, b_dh2.to(p_dh2.dtype.element_ty), mask=m_dh2) if K > 128: if STATE_V_FIRST: - p_dh3 = tl.make_block_ptr(dh + i_t_int64*HV*K*V, (V, K), (K, 1), (i_v * BV, 128), (BV, 64), (1, 0)) + p_dh3 = dh + i_t_int64*HV*K*V + o_v[:, None] * K + o_k3[None, :] + m_dh3 = m_v[:, None] & m_k3[None, :] else: - p_dh3 = tl.make_block_ptr(dh + i_t_int64*HV*K*V, (K, V), (V, 1), (128, i_v * BV), (64, BV), (1, 0)) - tl.store(p_dh3, b_dh3.to(p_dh3.dtype.element_ty), boundary_check=(0, 1)) + p_dh3 = dh + i_t_int64*HV*K*V + o_k3[:, None] * V + o_v[None, :] + m_dh3 = m_k3[:, None] & m_v[None, :] + tl.store(p_dh3, b_dh3.to(p_dh3.dtype.element_ty), mask=m_dh3) if K > 192: if STATE_V_FIRST: - p_dh4 = tl.make_block_ptr(dh + i_t_int64*HV*K*V, (V, K), (K, 1), (i_v * BV, 192), (BV, 64), (1, 0)) + p_dh4 = dh + i_t_int64*HV*K*V + o_v[:, None] * K + o_k4[None, :] + m_dh4 = m_v[:, None] & m_k4[None, :] else: - p_dh4 = tl.make_block_ptr(dh + i_t_int64*HV*K*V, (K, V), (V, 1), (192, i_v * BV), (64, BV), (1, 0)) - tl.store(p_dh4, b_dh4.to(p_dh4.dtype.element_ty), boundary_check=(0, 1)) + p_dh4 = dh + i_t_int64*HV*K*V + o_k4[:, None] * V + o_v[None, :] + m_dh4 = m_k4[:, None] & m_v[None, :] + tl.store(p_dh4, b_dh4.to(p_dh4.dtype.element_ty), mask=m_dh4) - last_idx = min((i_t + 1) * BT, T) - 1 + last_idx = min((i_t_int64 + 1) * BT, T) - 1 if USE_G: bg_last = tl.load(g + (bos + last_idx) * HV + i_h).to(tl.float32) - p_g = tl.make_block_ptr(g + bos * HV + i_h, (T,), (HV,), (i_t * BT,), (BT,), (0,)) - b_g = tl.load(p_g, boundary_check=(0,)).to(tl.float32) + p_g = g + bos * HV + i_h + o_t * HV + b_g = tl.load(p_g, mask=m_t, other=0.0).to(tl.float32) bg_last_exp = exp2(bg_last) b_g_exp = exp2(b_g) - p_dv = tl.make_block_ptr(dv, (T, V), (HV*V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) - p_dv2 = tl.make_block_ptr(dv2, (T, V), (HV*V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) - p_do = tl.make_block_ptr(do, (T, V), (HV*V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) + p_dv = dv + o_t[:, None] * (HV*V) + o_v[None, :] + p_dv2 = dv2 + o_t[:, None] * (HV*V) + o_v[None, :] + p_do = do + o_t[:, None] * (HV*V) + o_v[None, :] - b_do = tl.load(p_do, boundary_check=(0, 1)) + b_do = tl.load(p_do, mask=m_t[:, None] & m_v[None, :], other=0.0) # Update dv - p_k = tl.make_block_ptr(k, (T, K), (H*K, 1), (i_t * BT, 0), (BT, 64), (1, 0)) - b_k = tl.load(p_k, boundary_check=(0, 1)) + p_k = k + o_t[:, None] * (H*K) + o_k1[None, :] + b_k = tl.load(p_k, mask=m_t[:, None] & m_k1[None, :], other=0.0) if USE_GK: o_k1 = tl.arange(0, 64) b_gk_last1 = tl.load(gk + last_idx * HV*K + o_k1, mask=(o_k1 < K), other=0.).to(tl.float32) @@ -479,49 +550,45 @@ def chunk_gated_delta_rule_bwd_kernel_dhu_blockdim64( b_dv = tl.dot(b_k, b_dh1.to(b_k.dtype)) if K > 64: - p_k = tl.make_block_ptr(k, (T, K), (H*K, 1), (i_t * BT, 64), (BT, 64), (1, 0)) - b_k = tl.load(p_k, boundary_check=(0, 1)) + p_k = k + o_t[:, None] * (H*K) + o_k2[None, :] + b_k = tl.load(p_k, mask=m_t[:, None] & m_k2[None, :], other=0.0) if USE_GK: - o_k2 = 64 + o_k1 b_gk_last2 = tl.load(gk + last_idx * HV*K + o_k2, mask=(o_k2 < K), other=0.).to(tl.float32) if STATE_V_FIRST: - b_dv += tl.dot(b_k, tl.trans(b_dh2).to(b_k.dtype)) + b_dv = tl.dot(b_k, tl.trans(b_dh2).to(b_k.dtype), b_dv) else: - b_dv += tl.dot(b_k, b_dh2.to(b_k.dtype)) + b_dv = tl.dot(b_k, b_dh2.to(b_k.dtype), b_dv) if K > 128: - p_k = tl.make_block_ptr(k, (T, K), (H*K, 1), (i_t * BT, 128), (BT, 64), (1, 0)) - b_k = tl.load(p_k, boundary_check=(0, 1)) + p_k = k + o_t[:, None] * (H*K) + o_k3[None, :] + b_k = tl.load(p_k, mask=m_t[:, None] & m_k3[None, :], other=0.0) if USE_GK: - o_k3 = 128 + o_k1 b_gk_last3 = tl.load(gk + last_idx * HV*K + o_k3, mask=(o_k3 < K), other=0.).to(tl.float32) if STATE_V_FIRST: - b_dv += tl.dot(b_k, tl.trans(b_dh3).to(b_k.dtype)) + b_dv = tl.dot(b_k, tl.trans(b_dh3).to(b_k.dtype), b_dv) else: - b_dv += tl.dot(b_k, b_dh3.to(b_k.dtype)) + b_dv = tl.dot(b_k, b_dh3.to(b_k.dtype), b_dv) if K > 192: - p_k = tl.make_block_ptr(k, (T, K), (H*K, 1), (i_t * BT, 192), (BT, 64), (1, 0)) - b_k = tl.load(p_k, boundary_check=(0, 1)) + p_k = k + o_t[:, None] * (H*K) + o_k4[None, :] + b_k = tl.load(p_k, mask=m_t[:, None] & m_k4[None, :], other=0.0) if USE_GK: - o_k4 = 192 + o_k1 b_gk_last4 = tl.load(gk + last_idx * HV*K + o_k4, mask=(o_k4 < K), other=0.).to(tl.float32) if STATE_V_FIRST: - b_dv += tl.dot(b_k, tl.trans(b_dh4).to(b_k.dtype)) + b_dv = tl.dot(b_k, tl.trans(b_dh4).to(b_k.dtype), b_dv) else: - b_dv += tl.dot(b_k, b_dh4.to(b_k.dtype)) + b_dv = tl.dot(b_k, b_dh4.to(b_k.dtype), b_dv) if USE_G: - m_t = (i_t * BT + tl.arange(0, BT)) < T b_dv *= tl.where(m_t, exp2(bg_last - b_g), 0)[:, None] - b_dv += tl.load(p_dv, boundary_check=(0, 1)) + b_dv += tl.load(p_dv, mask=m_t[:, None] & m_v[None, :], other=0.0) - tl.store(p_dv2, b_dv.to(p_dv.dtype.element_ty), boundary_check=(0, 1)) + tl.store(p_dv2, b_dv.to(p_dv.dtype.element_ty), mask=m_t[:, None] & m_v[None, :]) # Update dh - p_w = tl.make_block_ptr(w, (K, T), (1, HV*K), (0, i_t * BT), (64, BT), (0, 1)) - p_q = tl.make_block_ptr(q, (K, T), (1, H*K), (0, i_t * BT), (64, BT), (0, 1)) - b_w = tl.load(p_w, boundary_check=(0, 1)) - b_q = tl.load(p_q, boundary_check=(0, 1)) + p_w = w + o_k1[:, None] + o_t[None, :] * (HV*K) + p_q = q + o_k1[:, None] + o_t[None, :] * (H*K) + b_w = tl.load(p_w, mask=m_k1[:, None] & m_t[None, :], other=0.0) + b_q = tl.load(p_q, mask=m_k1[:, None] & m_t[None, :], other=0.0) if USE_G: b_dh1 *= bg_last_exp b_q = b_q * b_g_exp[None, :] @@ -535,10 +602,10 @@ def chunk_gated_delta_rule_bwd_kernel_dhu_blockdim64( else: b_dh1 += tl.dot(b_q.to(b_q.dtype), b_do.to(b_q.dtype)) * scale - tl.dot(b_w, b_dv.to(b_w.dtype)) if K > 64: - p_q = tl.make_block_ptr(q, (K, T), (1, H*K), (64, i_t * BT), (64, BT), (0, 1)) - p_w = tl.make_block_ptr(w, (K, T), (1, HV*K), (64, i_t * BT), (64, BT), (0, 1)) - b_q = tl.load(p_q, boundary_check=(0, 1)) - b_w = tl.load(p_w, boundary_check=(0, 1)) + p_q = q + o_k2[:, None] + o_t[None, :] * (H*K) + p_w = w + o_k2[:, None] + o_t[None, :] * (HV*K) + b_q = tl.load(p_q, mask=m_k2[:, None] & m_t[None, :], other=0.0) + b_w = tl.load(p_w, mask=m_k2[:, None] & m_t[None, :], other=0.0) if USE_G: b_dh2 *= bg_last_exp b_q = b_q * b_g_exp[None, :] @@ -552,10 +619,10 @@ def chunk_gated_delta_rule_bwd_kernel_dhu_blockdim64( else: b_dh2 += tl.dot(b_q.to(b_q.dtype), b_do.to(b_q.dtype)) * scale - tl.dot(b_w, b_dv.to(b_w.dtype)) if K > 128: - p_q = tl.make_block_ptr(q, (K, T), (1, H*K), (128, i_t * BT), (64, BT), (0, 1)) - p_w = tl.make_block_ptr(w, (K, T), (1, HV*K), (128, i_t * BT), (64, BT), (0, 1)) - b_q = tl.load(p_q, boundary_check=(0, 1)) - b_w = tl.load(p_w, boundary_check=(0, 1)) + p_q = q + o_k3[:, None] + o_t[None, :] * (H*K) + p_w = w + o_k3[:, None] + o_t[None, :] * (HV*K) + b_q = tl.load(p_q, mask=m_k3[:, None] & m_t[None, :], other=0.0) + b_w = tl.load(p_w, mask=m_k3[:, None] & m_t[None, :], other=0.0) if USE_G: b_dh3 *= bg_last_exp b_q = b_q * b_g_exp[None, :] @@ -569,10 +636,10 @@ def chunk_gated_delta_rule_bwd_kernel_dhu_blockdim64( else: b_dh3 += tl.dot(b_q.to(b_q.dtype), b_do.to(b_q.dtype)) * scale - tl.dot(b_w, b_dv.to(b_w.dtype)) if K > 192: - p_q = tl.make_block_ptr(q, (K, T), (1, H*K), (192, i_t * BT), (64, BT), (0, 1)) - p_w = tl.make_block_ptr(w, (K, T), (1, HV*K), (192, i_t * BT), (64, BT), (0, 1)) - b_q = tl.load(p_q, boundary_check=(0, 1)) - b_w = tl.load(p_w, boundary_check=(0, 1)) + p_q = q + o_k4[:, None] + o_t[None, :] * (H*K) + p_w = w + o_k4[:, None] + o_t[None, :] * (HV*K) + b_q = tl.load(p_q, mask=m_k4[:, None] & m_t[None, :], other=0.0) + b_w = tl.load(p_w, mask=m_k4[:, None] & m_t[None, :], other=0.0) if USE_G: b_dh4 *= bg_last_exp b_q = b_q * b_g_exp[None, :] @@ -588,31 +655,40 @@ def chunk_gated_delta_rule_bwd_kernel_dhu_blockdim64( if USE_INITIAL_STATE: if STATE_V_FIRST: - p_dh0 = tl.make_block_ptr(dh0, (V, K), (K, 1), (i_v * BV, 0), (BV, 64), (1, 0)) + p_dh0 = dh0 + o_v[:, None] * K + o_k1[None, :] + m_dh0 = m_v[:, None] & m_k1[None, :] else: - p_dh0 = tl.make_block_ptr(dh0, (K, V), (V, 1), (0, i_v * BV), (64, BV), (1, 0)) - tl.store(p_dh0, b_dh1.to(p_dh0.dtype.element_ty), boundary_check=(0, 1)) + p_dh0 = dh0 + o_k1[:, None] * V + o_v[None, :] + m_dh0 = m_k1[:, None] & m_v[None, :] + tl.store(p_dh0, b_dh1.to(p_dh0.dtype.element_ty), mask=m_dh0) if K > 64: if STATE_V_FIRST: - p_dh1 = tl.make_block_ptr(dh0, (V, K), (K, 1), (i_v * BV, 64), (BV, 64), (1, 0)) + p_dh1 = dh0 + o_v[:, None] * K + o_k2[None, :] + m_dh1 = m_v[:, None] & m_k2[None, :] else: - p_dh1 = tl.make_block_ptr(dh0, (K, V), (V, 1), (64, i_v * BV), (64, BV), (1, 0)) - tl.store(p_dh1, b_dh2.to(p_dh1.dtype.element_ty), boundary_check=(0, 1)) + p_dh1 = dh0 + o_k2[:, None] * V + o_v[None, :] + m_dh1 = m_k2[:, None] & m_v[None, :] + tl.store(p_dh1, b_dh2.to(p_dh1.dtype.element_ty), mask=m_dh1) if K > 128: if STATE_V_FIRST: - p_dh2 = tl.make_block_ptr(dh0, (V, K), (K, 1), (i_v * BV, 128), (BV, 64), (1, 0)) + p_dh2 = dh0 + o_v[:, None] * K + o_k3[None, :] + m_dh2 = m_v[:, None] & m_k3[None, :] else: - p_dh2 = tl.make_block_ptr(dh0, (K, V), (V, 1), (128, i_v * BV), (64, BV), (1, 0)) - tl.store(p_dh2, b_dh3.to(p_dh2.dtype.element_ty), boundary_check=(0, 1)) + p_dh2 = dh0 + o_k3[:, None] * V + o_v[None, :] + m_dh2 = m_k3[:, None] & m_v[None, :] + tl.store(p_dh2, b_dh3.to(p_dh2.dtype.element_ty), mask=m_dh2) if K > 192: if STATE_V_FIRST: - p_dh3 = tl.make_block_ptr(dh0, (V, K), (K, 1), (i_v * BV, 192), (BV, 64), (1, 0)) + p_dh3 = dh0 + o_v[:, None] * K + o_k4[None, :] + m_dh3 = m_v[:, None] & m_k4[None, :] else: - p_dh3 = tl.make_block_ptr(dh0, (K, V), (V, 1), (192, i_v * BV), (64, BV), (1, 0)) - tl.store(p_dh3, b_dh4.to(p_dh3.dtype.element_ty), boundary_check=(0, 1)) + p_dh3 = dh0 + o_k4[:, None] * V + o_v[None, :] + m_dh3 = m_k4[:, None] & m_v[None, :] + tl.store(p_dh3, b_dh4.to(p_dh3.dtype.element_ty), mask=m_dh3) -@dispatch('common') +# Use triton only, other backends (cutlass, cute, etc.) need additional torch.compile support +# @dispatch('common') def chunk_gated_delta_rule_fwd_h( k: torch.Tensor, w: torch.Tensor, @@ -627,6 +703,7 @@ def chunk_gated_delta_rule_fwd_h( cu_seqlens: torch.LongTensor | None = None, cu_seqlens_cpu: torch.LongTensor | None = None, chunk_indices: torch.LongTensor | None = None, + chunk_offsets: torch.LongTensor | None = None, ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor | None]: B, T, H, K, V, HV = *k.shape, u.shape[-1], u.shape[2] BT = chunk_size @@ -637,7 +714,9 @@ def chunk_gated_delta_rule_fwd_h( if cu_seqlens is None: N, NT, chunk_offsets = B, triton.cdiv(T, BT), None else: - N, NT, chunk_offsets = len(cu_seqlens) - 1, len(chunk_indices), prepare_chunk_offsets(cu_seqlens, BT) + N, NT = len(cu_seqlens) - 1, len(chunk_indices) + if chunk_offsets is None: + chunk_offsets = prepare_chunk_offsets(cu_seqlens, BT) assert K <= 256, "current kernel does not support head dimension larger than 256." if state_v_first: @@ -648,7 +727,7 @@ def chunk_gated_delta_rule_fwd_h( final_state = k.new_zeros(N, HV, K, V, dtype=torch.float32) if output_final_state else None v_new = torch.empty_like(u) if save_new_value else None - def grid(meta): return (triton.cdiv(V, meta['BV']), N*HV) + def grid(meta): return (triton.cdiv(V, meta['BV']) * N * HV,) chunk_gated_delta_rule_fwd_kernel_h_blockdim64[grid]( k=k, v=u, @@ -672,6 +751,7 @@ def chunk_gated_delta_rule_fwd_h( return h, v_new, final_state +# @dispatch('common') def chunk_gated_delta_rule_bwd_dhu( q: torch.Tensor, k: torch.Tensor, @@ -687,6 +767,8 @@ def chunk_gated_delta_rule_bwd_dhu( cu_seqlens: torch.LongTensor | None = None, chunk_size: int = 64, chunk_indices: torch.LongTensor | None = None, + chunk_offsets: torch.LongTensor | None = None, + use_graph: bool = False, ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: B, T, H, K, V, HV = *q.shape, do.shape[-1], do.shape[2] # N: the actual number of sequences in the batch with either equal or variable lengths @@ -698,16 +780,26 @@ def chunk_gated_delta_rule_bwd_dhu( if cu_seqlens is None: N, NT, chunk_offsets = B, triton.cdiv(T, BT), None else: - N, NT, chunk_offsets = len(cu_seqlens) - 1, len(chunk_indices), prepare_chunk_offsets(cu_seqlens, BT) + N, NT = len(cu_seqlens) - 1, len(chunk_indices) + if chunk_offsets is None: + chunk_offsets = prepare_chunk_offsets(cu_seqlens, BT) - if state_v_first: - dh = q.new_empty(B, NT, HV, V, K) + if use_graph: + if state_v_first: + dh = get_static_buffer("dhu_dh_vf", (B, NT, HV, V, K), q.dtype, q.device) + else: + dh = get_static_buffer("dhu_dh", (B, NT, HV, K, V), q.dtype, q.device) + dh0 = get_static_buffer("dhu_dh0", tuple(h0.shape), torch.float32, h0.device) if h0 is not None else None + dv2 = get_static_buffer("dhu_dv2", tuple(dv.shape), dv.dtype, dv.device) else: - dh = q.new_empty(B, NT, HV, K, V) - dh0 = torch.empty_like(h0, dtype=torch.float32) if h0 is not None else None - dv2 = torch.empty_like(dv) + if state_v_first: + dh = q.new_empty(B, NT, HV, V, K) + else: + dh = q.new_empty(B, NT, HV, K, V) + dh0 = torch.empty_like(h0, dtype=torch.float32) if h0 is not None else None + dv2 = torch.empty_like(dv) - def grid(meta): return (triton.cdiv(V, meta['BV']), N*HV) + def grid(meta): return (triton.cdiv(V, meta['BV']) * N * HV,) chunk_gated_delta_rule_bwd_kernel_dhu_blockdim64[grid]( q=q, k=k, diff --git a/build/torch-cuda/ops/common/chunk_h.py b/build/torch-cuda/ops/common/chunk_h.py index e4d485ed6df46ddbf77973e2e1d73cf03dfb8c96..07adb3371316057977ac63e02752aa0c1e02ef24 100644 --- a/build/torch-cuda/ops/common/chunk_h.py +++ b/build/torch-cuda/ops/common/chunk_h.py @@ -9,9 +9,10 @@ import torch import triton import triton.language as tl +from ...backends import dispatch from ...ops.utils import prepare_chunk_offsets -from ...ops.utils.op import exp2 -from ...utils import autotune_cache_kwargs, check_shared_mem +from ...ops.utils.op import exp2, unflatten_program_id +from ...utils import IS_NVIDIA_HOPPER, autotune_cache_kwargs, check_shared_mem BKV_LIST = [32, 64] if check_shared_mem() else [16, 32] @@ -29,7 +30,7 @@ BKV_LIST = [32, 64] if check_shared_mem() else [16, 32] for num_warps in [1, 2, 4, 8] for num_stages in [2, 3, 4] ], - key=['BT', 'USE_G', 'USE_GK', 'USE_GV', 'STATE_V_FIRST'], + key=['BT', 'USE_G', 'USE_GK', 'USE_GV', 'STATE_V_FIRST', 'K', 'V'], **autotune_cache_kwargs, ) @triton.jit(do_not_specialize=['T']) @@ -62,13 +63,13 @@ def chunk_fwd_kernel_h( IS_VARLEN: tl.constexpr, STATE_V_FIRST: tl.constexpr, ): - i_k, i_v, i_nh = tl.program_id(0), tl.program_id(1), tl.program_id(2) + i_k, i_v, i_nh = unflatten_program_id(tl.cdiv(K, BK), tl.cdiv(V, BV)) i_n, i_h = i_nh // H, i_nh % H if IS_VARLEN: - bos, eos = tl.load(cu_seqlens + i_n).to(tl.int32), tl.load(cu_seqlens + i_n + 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 NT, NS = tl.cdiv(T, BT), tl.cdiv(T, BS) - boh = tl.load(split_offsets + i_n).to(tl.int32) + boh = tl.load(split_offsets + i_n).to(tl.int64) else: bos, eos = i_n * T, i_n * T + T NT, NS = tl.cdiv(T, BT), tl.cdiv(T, BS) @@ -82,75 +83,82 @@ def chunk_fwd_kernel_h( # [BK, BV] accumulator; STATE_V_FIRST only flips the stored state's HBM layout to [V, K], applied at the load/store below. b_h = tl.zeros([BK, BV], dtype=tl.float32) + o_k = i_k * BK + tl.arange(0, BK) + o_v = i_v * BV + tl.arange(0, BV) if USE_INITIAL_STATE: if STATE_V_FIRST: - p_h0 = tl.make_block_ptr(h0 + i_nh * K*V, (V, K), (K, 1), (i_v * BV, i_k * BK), (BV, BK), (1, 0)) - b_h = tl.trans(tl.load(p_h0, boundary_check=(0, 1))).to(tl.float32) + p_h0 = h0 + i_nh * K*V + o_v[:, None] * K + o_k[None, :] + b_h = tl.trans(tl.load(p_h0, mask=(o_v[:, None] < V) & (o_k[None, :] < K), other=0.0)).to(tl.float32) else: - p_h0 = tl.make_block_ptr(h0 + i_nh * K*V, (K, V), (V, 1), (i_k * BK, i_v * BV), (BK, BV), (1, 0)) - b_h = tl.load(p_h0, boundary_check=(0, 1)).to(tl.float32) + p_h0 = h0 + i_nh * K*V + o_k[:, None] * V + o_v[None, :] + b_h = tl.load(p_h0, mask=(o_k[:, None] < K) & (o_v[None, :] < V), other=0.0).to(tl.float32) for i_t in range(NT): - i_s = i_t // NTS - p_k = tl.make_block_ptr(k + (bos*H + i_h) * K, (K, T), (1, H*K), (i_k * BK, i_t * BT), (BK, BT), (0, 1)) - p_v = tl.make_block_ptr(v + (bos*H + i_h) * V, (T, V), (H*V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) + i_t_int64 = i_t.to(tl.int64) + i_s = i_t_int64 // NTS + o_t = i_t_int64 * BT + tl.arange(0, BT) + m_t = o_t < T + p_k = k + (bos*H + i_h) * K + o_k[:, None] + o_t[None, :] * (H*K) + p_v = v + (bos*H + i_h) * V + o_t[:, None] * (H*V) + o_v[None, :] o_h = ((boh + i_s) * H + i_h).to(tl.int64) * K*V if STATE_V_FIRST: - p_h = tl.make_block_ptr(h + o_h, (V, K), (K, 1), (i_v * BV, i_k * BK), (BV, BK), (1, 0)) + p_h = h + o_h + o_v[:, None] * K + o_k[None, :] + m_h = (o_v[:, None] < V) & (o_k[None, :] < K) else: - p_h = tl.make_block_ptr(h + o_h, (K, V), (V, 1), (i_k * BK, i_v * BV), (BK, BV), (1, 0)) + p_h = h + o_h + o_k[:, None] * V + o_v[None, :] + m_h = (o_k[:, None] < K) & (o_v[None, :] < V) - if i_t % NTS == 0: - tl.store(p_h, (tl.trans(b_h) if STATE_V_FIRST else b_h).to(p_h.dtype.element_ty), boundary_check=(0, 1)) + if i_t_int64 % NTS == 0: + tl.store(p_h, (tl.trans(b_h) if STATE_V_FIRST else b_h).to(p_h.dtype.element_ty), mask=m_h) # [BK, BT] - b_k = tl.load(p_k, boundary_check=(0, 1)) + b_k = tl.load(p_k, mask=(o_k[:, None] < K) & m_t[None, :], other=0.0) # [BT, BV] - b_v = tl.load(p_v, boundary_check=(0, 1)) - last_idx = min((i_t + 1) * BT, T) - 1 + b_v = tl.load(p_v, mask=m_t[:, None] & (o_v < V)[None, :], other=0.0) + last_idx = min((i_t_int64 + 1) * BT, T) - 1 # scalar decay if USE_G: b_g_last = tl.load(g + bos * H + last_idx * H + i_h) - p_g = g + bos*H + (i_t * BT + tl.arange(0, BT)) * H + i_h - b_g = tl.load(p_g, mask=(i_t * BT + tl.arange(0, BT) < T), other=0.) + p_g = g + bos*H + (i_t_int64 * BT + tl.arange(0, BT)) * H + i_h + b_g = tl.load(p_g, mask=(i_t_int64 * BT + tl.arange(0, BT) < T), other=0.) b_h *= exp2(b_g_last) b_v = (b_v * exp2(b_g_last - b_g)[:, None]).to(b_v.dtype) if USE_G_GAMMA: - b_g_last = b_gamma * min(BT, T - i_t * BT) + b_g_last = b_gamma * min(BT, T - i_t_int64 * BT) b_h *= exp2(b_g_last) b_v = (b_v * exp2(b_g_last - b_g)[:, None]).to(b_v.dtype) # vector decay, h = Diag(gk) @ h if USE_GK: - p_gk = tl.make_block_ptr(gk + (bos*H + i_h) * K, (K, T), (1, H*K), (i_k * BK, i_t * BT), (BK, BT), (0, 1)) + p_gk = gk + (bos*H + i_h) * K + o_k[:, None] + o_t[None, :] * (H*K) p_gk_last = gk + (bos + last_idx) * H*K + i_h * K + i_k * BK + tl.arange(0, BK) b_gk_last = tl.load(p_gk_last, mask=(i_k * BK + tl.arange(0, BK) < K), other=0.) - b_gk = tl.load(p_gk, boundary_check=(0, 1)) + b_gk = tl.load(p_gk, mask=(o_k[:, None] < K) & m_t[None, :], other=0.0) b_h *= exp2(b_gk_last)[:, None] b_k = (b_k * exp2(b_gk_last[:, None] - b_gk)).to(b_k.dtype) # vector decay, h = h @ Diag(gv) if USE_GV: - p_gv = tl.make_block_ptr(gv + (bos*H + i_h) * V, (T, V), (H*V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) + p_gv = gv + (bos*H + i_h) * V + o_t[:, None] * (H*V) + o_v[None, :] p_gv_last = gv + (bos + last_idx) * H*V + i_h * V + i_v * BV + tl.arange(0, BV) b_gv_last = tl.load(p_gv_last, mask=(i_v * BV + tl.arange(0, BV) < V), other=0.) - b_gv = tl.load(p_gv, boundary_check=(0, 1)) + b_gv = tl.load(p_gv, mask=m_t[:, None] & (o_v < V)[None, :], other=0.0) b_h *= exp2(b_gv_last)[None, :] b_v = (b_v * exp2(b_gv_last[None, :] - b_gv)).to(b_v.dtype) - b_h += tl.dot(b_k, b_v) + b_h = tl.dot(b_k, b_v, b_h) if STORE_FINAL_STATE: if STATE_V_FIRST: - p_ht = tl.make_block_ptr(ht + i_nh * K*V, (V, K), (K, 1), (i_v * BV, i_k * BK), (BV, BK), (1, 0)) - tl.store(p_ht, tl.trans(b_h).to(p_ht.dtype.element_ty), boundary_check=(0, 1)) + p_ht = ht + i_nh * K*V + o_v[:, None] * K + o_k[None, :] + tl.store(p_ht, tl.trans(b_h).to(p_ht.dtype.element_ty), mask=(o_v[:, None] < V) & (o_k[None, :] < K)) else: - p_ht = tl.make_block_ptr(ht + i_nh * K*V, (K, V), (V, 1), (i_k * BK, i_v * BV), (BK, BV), (1, 0)) - tl.store(p_ht, b_h.to(p_ht.dtype.element_ty), boundary_check=(0, 1)) + p_ht = ht + i_nh * K*V + o_k[:, None] * V + o_v[None, :] + tl.store(p_ht, b_h.to(p_ht.dtype.element_ty), mask=(o_k[:, None] < K) & (o_v[None, :] < V)) @triton.heuristics({ @@ -166,10 +174,11 @@ def chunk_fwd_kernel_h( for num_warps in [1, 2, 4, 8] for num_stages in [2, 3, 4] ], - key=['BT', 'USE_G', 'USE_GK', 'USE_GV', 'STATE_V_FIRST'], + key=['BT', 'USE_G', 'USE_GK', 'USE_GV', 'STATE_V_FIRST', 'K', 'V'], **autotune_cache_kwargs, ) -@triton.jit(do_not_specialize=['T']) +# keep scale=1 dynamic to avoid a Hopper masked-load shared-memory miscompile. +@triton.jit(do_not_specialize=['T', 'scale'] if IS_NVIDIA_HOPPER else ['T']) def chunk_bwd_kernel_dh( q, g, @@ -202,15 +211,15 @@ def chunk_bwd_kernel_dh( IS_VARLEN: tl.constexpr, STATE_V_FIRST: tl.constexpr, ): - i_k, i_v, i_nh = tl.program_id(0), tl.program_id(1), tl.program_id(2) + i_k, i_v, i_nh = unflatten_program_id(tl.cdiv(K, BK), tl.cdiv(V, BV)) i_n, i_hq = i_nh // HQ, i_nh % HQ i_h = i_hq // NG if IS_VARLEN: - bos, eos = tl.load(cu_seqlens + i_n).to(tl.int32), tl.load(cu_seqlens + i_n + 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 NT = tl.cdiv(T, BT) NS = tl.cdiv(T, BS) - boh = tl.load(split_offsets + i_n).to(tl.int32) + boh = tl.load(split_offsets + i_n).to(tl.int64) else: bos, eos = i_n * T, i_n * T + T NT = tl.cdiv(T, BT) @@ -223,74 +232,82 @@ def chunk_bwd_kernel_dh( # [BK, BV] accumulator; STATE_V_FIRST only flips the stored state's HBM layout to [V, K], applied at the load/store below. b_dh = tl.zeros([BK, BV], dtype=tl.float32) + o_k = i_k * BK + tl.arange(0, BK) + o_v = i_v * BV + tl.arange(0, BV) if USE_FINAL_STATE_GRADIENT: if STATE_V_FIRST: - p_dht = tl.make_block_ptr(dht + i_nh * K*V, (V, K), (K, 1), (i_v * BV, i_k * BK), (BV, BK), (1, 0)) - b_dh += tl.trans(tl.load(p_dht, boundary_check=(0, 1))).to(tl.float32) + p_dht = dht + i_nh * K*V + o_v[:, None] * K + o_k[None, :] + b_dh += tl.trans(tl.load(p_dht, mask=(o_v[:, None] < V) & (o_k[None, :] < K), other=0.0)).to(tl.float32) else: - p_dht = tl.make_block_ptr(dht + i_nh * K*V, (K, V), (V, 1), (i_k * BK, i_v * BV), (BK, BV), (1, 0)) - b_dh += tl.load(p_dht, boundary_check=(0, 1)).to(tl.float32) + p_dht = dht + i_nh * K*V + o_k[:, None] * V + o_v[None, :] + b_dh += tl.load(p_dht, mask=(o_k[:, None] < K) & (o_v[None, :] < V), other=0.0).to(tl.float32) for i_t in range(NT - 1, -1, -1): - i_s = i_t // (BS // BT) + i_t_int64 = i_t.to(tl.int64) + i_s = i_t_int64 // (BS // BT) o_dh = ((boh + i_s) * H + i_h).to(tl.int64) * K*V if STATE_V_FIRST: - p_dh = tl.make_block_ptr(dh + o_dh, (V, K), (K, 1), (i_v * BV, i_k * BK), (BV, BK), (1, 0)) + p_dh = dh + o_dh + o_v[:, None] * K + o_k[None, :] + m_dh = (o_v[:, None] < V) & (o_k[None, :] < K) else: - p_dh = tl.make_block_ptr(dh + o_dh, (K, V), (V, 1), (i_k * BK, i_v * BV), (BK, BV), (1, 0)) - - if i_t % (BS // BT) == 0: - tl.store(p_dh, (tl.trans(b_dh) if STATE_V_FIRST else b_dh).to(p_dh.dtype.element_ty), boundary_check=(0, 1)) - last_idx = min(i_t * BT + BT, T) - 1 + p_dh = dh + o_dh + o_k[:, None] * V + o_v[None, :] + m_dh = (o_k[:, None] < K) & (o_v[None, :] < V) + + if i_t_int64 % (BS // BT) == 0: + tl.store(p_dh, (tl.trans(b_dh) if STATE_V_FIRST else b_dh).to(p_dh.dtype.element_ty), mask=m_dh) + last_idx = min(i_t_int64 * BT + BT, T) - 1 + o_t = i_t_int64 * BT + tl.arange(0, BT) + m_t = o_t < T # [BK, BT] - p_q = tl.make_block_ptr(q + (bos*HQ + i_hq) * K, (K, T), (1, HQ*K), (i_k * BK, i_t * BT), (BK, BT), (0, 1)) - p_do = tl.make_block_ptr(do + (bos*HQ + i_hq) * V, (T, V), (HQ*V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) - b_q = tl.load(p_q, boundary_check=(0, 1)) + p_q = q + (bos*HQ + i_hq) * K + o_k[:, None] + o_t[None, :] * (HQ*K) + p_do = do + (bos*HQ + i_hq) * V + o_t[:, None] * (HQ*V) + o_v[None, :] + b_q = tl.load(p_q, mask=(o_k[:, None] < K) & m_t[None, :], other=0.0) b_q = (b_q * scale).to(b_q.dtype) # [BT, BV] - b_do = tl.load(p_do, boundary_check=(0, 1)) + b_do = tl.load(p_do, mask=m_t[:, None] & (o_v < V)[None, :], other=0.0) if USE_G: - p_g = g + (bos + i_t * BT + tl.arange(0, BT)) * H + i_h + p_g = g + (bos + i_t_int64 * BT + tl.arange(0, BT)) * H + i_h b_g_last = tl.load(g + (bos + last_idx) * H + i_h) - b_g = tl.load(p_g, mask=(i_t * BT + tl.arange(0, BT) < T), other=0.) + b_g = tl.load(p_g, mask=(i_t_int64 * BT + tl.arange(0, BT) < T), other=0.) b_q = (b_q * exp2(b_g)[None, :]).to(b_q.dtype) b_dh *= exp2(b_g_last) if USE_G_GAMMA: - b_g_last = b_gamma * min(BT, T - i_t * BT) + b_g_last = b_gamma * min(BT, T - i_t_int64 * BT) b_q = (b_q * exp2(b_g)[None, :]).to(b_q.dtype) b_dh *= exp2(b_g_last) if USE_GK: - p_gk = tl.make_block_ptr(gk + (bos*H + i_h) * K, (K, T), (1, H*K), (i_k * BK, i_t * BT), (BK, BT), (0, 1)) + p_gk = gk + (bos*H + i_h) * K + o_k[:, None] + o_t[None, :] * (H*K) p_gk_last = gk + (bos + last_idx) * H*K + i_h * K + i_k * BK + tl.arange(0, BK) - b_gk = tl.load(p_gk, boundary_check=(0, 1)) + b_gk = tl.load(p_gk, mask=(o_k[:, None] < K) & m_t[None, :], other=0.0) b_gk_last = tl.load(p_gk_last, mask=(i_k * BK + tl.arange(0, BK) < K), other=0.) b_q = (b_q * exp2(b_gk)).to(b_q.dtype) b_dh *= exp2(b_gk_last)[:, None] if USE_GV: - p_gv = tl.make_block_ptr(gv + (bos*H + i_h) * V, (T, V), (H*V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) + p_gv = gv + (bos*H + i_h) * V + o_t[:, None] * (H*V) + o_v[None, :] p_gv_last = gv + (bos + last_idx) * H*V + i_h * V + i_v * BV + tl.arange(0, BV) - b_gv = tl.load(p_gv, boundary_check=(0, 1)) + b_gv = tl.load(p_gv, mask=m_t[:, None] & (o_v < V)[None, :], other=0.0) b_gv_last = tl.load(p_gv_last, mask=(i_v * BV + tl.arange(0, BV) < V), other=0.) b_do = (b_do * exp2(b_gv)) b_dh *= exp2(b_gv_last)[None, :] - b_dh += tl.dot(b_q, b_do.to(b_q.dtype)) + b_dh = tl.dot(b_q, b_do.to(b_q.dtype), b_dh) if STORE_INITIAL_STATE_GRADIENT: if STATE_V_FIRST: - p_dh0 = tl.make_block_ptr(dh0 + i_nh * K*V, (V, K), (K, 1), (i_v * BV, i_k * BK), (BV, BK), (1, 0)) - tl.store(p_dh0, tl.trans(b_dh).to(p_dh0.dtype.element_ty), boundary_check=(0, 1)) + p_dh0 = dh0 + i_nh * K*V + o_v[:, None] * K + o_k[None, :] + tl.store(p_dh0, tl.trans(b_dh).to(p_dh0.dtype.element_ty), mask=(o_v[:, None] < V) & (o_k[None, :] < K)) else: - p_dh0 = tl.make_block_ptr(dh0 + i_nh * K*V, (K, V), (V, 1), (i_k * BK, i_v * BV), (BK, BV), (1, 0)) - tl.store(p_dh0, b_dh.to(p_dh0.dtype.element_ty), boundary_check=(0, 1)) + p_dh0 = dh0 + i_nh * K*V + o_k[:, None] * V + o_v[None, :] + tl.store(p_dh0, b_dh.to(p_dh0.dtype.element_ty), mask=(o_k[:, None] < K) & (o_v[None, :] < V)) +@dispatch('common') def chunk_fwd_h( k: torch.Tensor, v: torch.Tensor, @@ -321,7 +338,7 @@ def chunk_fwd_h( state_shape = (V, K) if state_v_first else (K, V) h = k.new_empty(B, NS, H, *state_shape, dtype=k.dtype if not states_in_fp32 else torch.float) ht = k.new_empty(N, H, *state_shape, dtype=torch.float) if output_final_state else None - def grid(meta): return (triton.cdiv(K, meta['BK']), triton.cdiv(V, meta['BV']), N * H) + def grid(meta): return (triton.cdiv(K, meta['BK']) * triton.cdiv(V, meta['BV']) * N * H,) chunk_fwd_kernel_h[grid]( k=k, v=v, @@ -349,6 +366,7 @@ def chunk_fwd_h( return h, ht +@dispatch('common') def chunk_bwd_dh( q: torch.Tensor, k: torch.Tensor, @@ -386,7 +404,7 @@ def chunk_bwd_dh( dh = k.new_empty(B, NS, HQ, *state_shape, dtype=k.dtype if not states_in_fp32 else torch.float) dh0 = torch.empty_like(h0, dtype=torch.float) if h0 is not None else None - def grid(meta): return (triton.cdiv(K, meta['BK']), triton.cdiv(V, meta['BV']), N * H) + def grid(meta): return (triton.cdiv(K, meta['BK']) * triton.cdiv(V, meta['BV']) * N * H,) chunk_bwd_kernel_dh[grid]( q=q, g=g, diff --git a/build/torch-cuda/ops/common/chunk_h_parallel.py b/build/torch-cuda/ops/common/chunk_h_parallel.py index 2d864f5a710a0ddd563a628fde675b6cee99db86..b93c8ad01f95a8292a53e455710f0fe7763c19b1 100644 --- a/build/torch-cuda/ops/common/chunk_h_parallel.py +++ b/build/torch-cuda/ops/common/chunk_h_parallel.py @@ -15,7 +15,7 @@ import triton import triton.language as tl from ...ops.utils import prepare_chunk_indices, prepare_chunk_offsets -from ...ops.utils.op import exp +from ...ops.utils.op import exp, unflatten_program_id from ...utils import autotune_cache_kwargs @@ -61,7 +61,8 @@ def chunk_fwd_kernel_h_parallel( STORE_FINAL_STATE: tl.constexpr, IS_VARLEN: tl.constexpr, ): - i_kv, i_t, i_bh = tl.program_id(0), tl.program_id(1), tl.program_id(2) + i_kv, i_t = unflatten_program_id(tl.cdiv(K, BK) * tl.cdiv(V, BV)) + i_bh = tl.program_id(1).to(tl.int64) NV = tl.cdiv(V, BV) # i_b: batch index @@ -73,8 +74,8 @@ def chunk_fwd_kernel_h_parallel( i_b, i_h = i_bh // H, i_bh % H 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.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) + i_n, i_t = tl.load(chunk_indices + i_t * 2).to(tl.int32), tl.load(chunk_indices + i_t * 2 + 1).to(tl.int64) + bos, eos = tl.load(cu_seqlens + i_n).to(tl.int64), tl.load(cu_seqlens + i_n + 1).to(tl.int64) T = eos - bos NT = tl.cdiv(T, BT) else: @@ -83,22 +84,27 @@ def chunk_fwd_kernel_h_parallel( i_n, i_tg = i_b, i_b * NT + i_t i_nh = i_n * H + i_h - p_k = tl.make_block_ptr(k + (bos*H + i_h) * K, (K, T), (1, H*K), (i_k * BK, i_t * BT), (BK, BT), (0, 1)) - p_v = tl.make_block_ptr(v + (bos*H + i_h) * V, (T, V), (H*V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) - p_h = tl.make_block_ptr(h + (i_tg * H + i_h) * K*V, (K, V), (V, 1), (i_k * BK, i_v * BV), (BK, BV), (1, 0)) + o_k = i_k * BK + tl.arange(0, BK) + o_v = i_v * BV + tl.arange(0, BV) + o_t = i_t * BT + tl.arange(0, BT) + m_t = o_t < T + m_kv = (o_k[:, None] < K) & (o_v[None, :] < V) + p_k = k + (bos*H + i_h) * K + o_k[:, None] + o_t[None, :] * (H*K) + p_v = v + (bos*H + i_h) * V + o_t[:, None] * (H*V) + o_v[None, :] + p_h = h + (i_tg * H + i_h) * K*V + o_k[:, None] * V + o_v[None, :] if i_t == 0: if USE_INITIAL_STATE: - p_h0 = tl.make_block_ptr(h0 + i_nh * K*V, (K, V), (V, 1), (i_k * BK, i_v * BV), (BK, BV), (1, 0)) - b_h = tl.load(p_h0, boundary_check=(0, 1)).to(tl.float32) + p_h0 = h0 + i_nh * K*V + o_k[:, None] * V + o_v[None, :] + b_h = tl.load(p_h0, mask=m_kv, other=0.0).to(tl.float32) else: b_h = tl.zeros([BK, BV], dtype=tl.float32) - tl.store(p_h, b_h.to(p_h.dtype.element_ty), boundary_check=(0, 1)) + tl.store(p_h, b_h.to(p_h.dtype.element_ty), mask=m_kv) # [BK, BT] - b_k = tl.load(p_k, boundary_check=(0, 1)) + b_k = tl.load(p_k, mask=(o_k[:, None] < K) & m_t[None, :], other=0.0) # [BT, BV] - b_v = tl.load(p_v, boundary_check=(0, 1)) + b_v = tl.load(p_v, mask=m_t[:, None] & (o_v < V)[None, :], other=0.0) last_idx = min(i_t * BT + BT, T) - 1 # scalar decay @@ -110,29 +116,29 @@ def chunk_fwd_kernel_h_parallel( # vector decay, h = Diag(gk) @ h if USE_GK: - p_gk = tl.make_block_ptr(gk + (bos*H + i_h) * K, (K, T), (1, H*K), (i_k * BK, i_t * BT), (BK, BT), (0, 1)) + p_gk = gk + (bos*H + i_h) * K + o_k[:, None] + o_t[None, :] * (H*K) p_gk_last = gk + (bos + last_idx) * H*K + i_h * K + i_k * BK + tl.arange(0, BK) b_gk_last = tl.load(p_gk_last, mask=(i_k * BK + tl.arange(0, BK) < K), other=0.) - b_gk = tl.load(p_gk, boundary_check=(0, 1)) + b_gk = tl.load(p_gk, mask=(o_k[:, None] < K) & m_t[None, :], other=0.0) b_k = (b_k * exp(b_gk_last[:, None] - b_gk)).to(b_k.dtype) # vector decay, h = h @ Diag(gv) if USE_GV: - p_gv = tl.make_block_ptr(gv + (bos*H + i_h) * V, (T, V), (H*V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) + p_gv = gv + (bos*H + i_h) * V + o_t[:, None] * (H*V) + o_v[None, :] p_gv_last = gv + (bos + last_idx) * H*V + i_h * V + i_v * BV + tl.arange(0, BV) b_gv_last = tl.load(p_gv_last, mask=(i_v * BV + tl.arange(0, BV) < V), other=0.) - b_gv = tl.load(p_gv, boundary_check=(0, 1)) + b_gv = tl.load(p_gv, mask=m_t[:, None] & (o_v < V)[None, :], other=0.0) b_v = (b_v * exp(b_gv_last[None, :] - b_gv)).to(b_v.dtype) b_h = tl.dot(b_k, b_v) if i_t < NT - 1: - p_h = tl.make_block_ptr(h + ((i_tg + 1) * H + i_h) * K*V, (K, V), (V, 1), (i_k * BK, i_v * BV), (BK, BV), (1, 0)) - tl.store(p_h, b_h.to(p_h.dtype.element_ty), boundary_check=(0, 1)) + p_h = h + ((i_tg + 1) * H + i_h) * K*V + o_k[:, None] * V + o_v[None, :] + tl.store(p_h, b_h.to(p_h.dtype.element_ty), mask=m_kv) elif STORE_FINAL_STATE: - p_ht = tl.make_block_ptr(ht + i_nh * K*V, (K, V), (V, 1), (i_k * BK, i_v * BV), (BK, BV), (1, 0)) - tl.store(p_ht, b_h.to(p_ht.dtype.element_ty), boundary_check=(0, 1)) + p_ht = ht + i_nh * K*V + o_k[:, None] * V + o_v[None, :] + tl.store(p_ht, b_h.to(p_ht.dtype.element_ty), mask=m_kv) @triton.heuristics({ @@ -173,13 +179,13 @@ def chunk_fwd_kernel_h_reduction( STORE_FINAL_STATE: tl.constexpr, IS_VARLEN: tl.constexpr, ): - i_k, i_v, i_nh = tl.program_id(0), tl.program_id(1), tl.program_id(2) + i_k, i_v, i_nh = unflatten_program_id(tl.cdiv(K, BK), tl.cdiv(V, BV)) i_n, i_h = i_nh // H, i_nh % H if IS_VARLEN: - bos, eos = tl.load(cu_seqlens + i_n).to(tl.int32), tl.load(cu_seqlens + i_n + 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 NT = tl.cdiv(T, BT) - boh = tl.load(chunk_offsets + i_n).to(tl.int32) + boh = tl.load(chunk_offsets + i_n).to(tl.int64) else: bos, eos = i_n * T, i_n * T + T NT = tl.cdiv(T, BT) @@ -187,11 +193,14 @@ def chunk_fwd_kernel_h_reduction( # [BK, BV] b_h = tl.zeros([BK, BV], dtype=tl.float32) + o_k = i_k * BK + tl.arange(0, BK) + o_v = i_v * BV + tl.arange(0, BV) + m_kv = (o_k[:, None] < K) & (o_v[None, :] < V) for i_t in range(NT): - p_h = tl.make_block_ptr(h + ((boh + i_t) * H + i_h) * K*V, (K, V), (V, 1), (i_k * BK, i_v * BV), (BK, BV), (1, 0)) - b_h += tl.load(p_h, boundary_check=(0, 1)).to(tl.float32) + p_h = h + ((boh + i_t) * H + i_h) * K*V + o_k[:, None] * V + o_v[None, :] + b_h += tl.load(p_h, mask=m_kv, other=0.0).to(tl.float32) if i_t > 0: - tl.store(p_h, b_h.to(p_h.dtype.element_ty), boundary_check=(0, 1)) + tl.store(p_h, b_h.to(p_h.dtype.element_ty), mask=m_kv) last_idx = min(i_t * BT + BT, T) - 1 # scalar decay @@ -214,10 +223,10 @@ def chunk_fwd_kernel_h_reduction( b_h *= exp(b_gv_last)[None, :] if STORE_FINAL_STATE: - p_kvt = tl.make_block_ptr(kvt + i_nh * K*V, (K, V), (V, 1), (i_k * BK, i_v * BV), (BK, BV), (1, 0)) - p_ht = tl.make_block_ptr(ht + i_nh * K*V, (K, V), (V, 1), (i_k * BK, i_v * BV), (BK, BV), (1, 0)) - b_h += tl.load(p_kvt, boundary_check=(0, 1)).to(tl.float32) - tl.store(p_ht, b_h.to(p_ht.dtype.element_ty), boundary_check=(0, 1)) + p_kvt = kvt + i_nh * K*V + o_k[:, None] * V + o_v[None, :] + p_ht = ht + i_nh * K*V + o_k[:, None] * V + o_v[None, :] + b_h += tl.load(p_kvt, mask=m_kv, other=0.0).to(tl.float32) + tl.store(p_ht, b_h.to(p_ht.dtype.element_ty), mask=m_kv) @triton.heuristics({ @@ -265,7 +274,8 @@ def chunk_bwd_kernel_dh_parallel( USE_FINAL_STATE_GRADIENT: tl.constexpr, IS_VARLEN: tl.constexpr, ): - i_kv, i_t, i_bh = tl.program_id(0), tl.program_id(1), tl.program_id(2) + i_kv, i_t = unflatten_program_id(tl.cdiv(K, BK) * tl.cdiv(V, BV)) + i_bh = tl.program_id(1).to(tl.int64) NV = tl.cdiv(V, BV) i_k, i_v = i_kv // NV, i_kv % NV @@ -273,8 +283,8 @@ def chunk_bwd_kernel_dh_parallel( i_h = i_hq // NG 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.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) + i_n, i_t = tl.load(chunk_indices + i_t * 2).to(tl.int32), tl.load(chunk_indices + i_t * 2 + 1).to(tl.int64) + bos, eos = tl.load(cu_seqlens + i_n).to(tl.int64), tl.load(cu_seqlens + i_n + 1).to(tl.int64) T = eos - bos NT = tl.cdiv(T, BT) else: @@ -283,23 +293,28 @@ def chunk_bwd_kernel_dh_parallel( i_n, i_tg = i_b, i_b * NT + i_t i_nh = i_n * HQ + i_hq - p_q = tl.make_block_ptr(q + (bos*HQ + i_hq) * K, (K, T), (1, HQ*K), (i_k * BK, i_t * BT), (BK, BT), (0, 1)) - p_do = tl.make_block_ptr(do + (bos*HQ + i_hq) * V, (T, V), (HQ*V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) - p_dh = tl.make_block_ptr(dh + (i_tg * H + i_h) * K*V, (K, V), (V, 1), (i_k * BK, i_v * BV), (BK, BV), (1, 0)) + o_k = i_k * BK + tl.arange(0, BK) + o_v = i_v * BV + tl.arange(0, BV) + o_t = i_t * BT + tl.arange(0, BT) + m_t = o_t < T + m_kv = (o_k[:, None] < K) & (o_v[None, :] < V) + p_q = q + (bos*HQ + i_hq) * K + o_k[:, None] + o_t[None, :] * (HQ*K) + p_do = do + (bos*HQ + i_hq) * V + o_t[:, None] * (HQ*V) + o_v[None, :] + p_dh = dh + (i_tg * H + i_h) * K*V + o_k[:, None] * V + o_v[None, :] if i_t == NT - 1: if USE_FINAL_STATE_GRADIENT: - p_dht = tl.make_block_ptr(dht + i_nh * K*V, (K, V), (V, 1), (i_k * BK, i_v * BV), (BK, BV), (1, 0)) - b_dh = tl.load(p_dht, boundary_check=(0, 1)).to(tl.float32) + p_dht = dht + i_nh * K*V + o_k[:, None] * V + o_v[None, :] + b_dh = tl.load(p_dht, mask=m_kv, other=0.0).to(tl.float32) else: b_dh = tl.zeros([BK, BV], dtype=tl.float32) - tl.store(p_dh, b_dh.to(p_dh.dtype.element_ty), boundary_check=(0, 1)) + tl.store(p_dh, b_dh.to(p_dh.dtype.element_ty), mask=m_kv) # [BK, BT] - b_q = tl.load(p_q, boundary_check=(0, 1)) + b_q = tl.load(p_q, mask=(o_k[:, None] < K) & m_t[None, :], other=0.0) b_q = (b_q * scale).to(b_q.dtype) # [BT, BV] - b_do = tl.load(p_do, boundary_check=(0, 1)) + b_do = tl.load(p_do, mask=m_t[:, None] & (o_v < V)[None, :], other=0.0) if USE_G: p_g = g + (bos + i_t * BT + tl.arange(0, BT)) * H + i_h @@ -307,22 +322,22 @@ def chunk_bwd_kernel_dh_parallel( b_q = (b_q * exp(b_g)[None, :]).to(b_q.dtype) if USE_GK: - p_gk = tl.make_block_ptr(gk + (bos*H + i_h) * K, (K, T), (1, H*K), (i_k * BK, i_t * BT), (BK, BT), (0, 1)) - b_gk = tl.load(p_gk, boundary_check=(0, 1)) + p_gk = gk + (bos*H + i_h) * K + o_k[:, None] + o_t[None, :] * (H*K) + b_gk = tl.load(p_gk, mask=(o_k[:, None] < K) & m_t[None, :], other=0.0) b_q = (b_q * exp(b_gk)).to(b_q.dtype) if USE_GV: - p_gv = tl.make_block_ptr(gv + (bos*H + i_h) * V, (T, V), (H*V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) - b_gv = tl.load(p_gv, boundary_check=(0, 1)) + p_gv = gv + (bos*H + i_h) * V + o_t[:, None] * (H*V) + o_v[None, :] + b_gv = tl.load(p_gv, mask=m_t[:, None] & (o_v < V)[None, :], other=0.0) b_do = (b_do * exp(b_gv)).to(b_do.dtype) b_dh = tl.dot(b_q, b_do) if i_t > 0: - p_dh = tl.make_block_ptr(dh + ((i_tg - 1) * H + i_h) * K*V, (K, V), (V, 1), (i_k * BK, i_v * BV), (BK, BV), (1, 0)) - tl.store(p_dh, b_dh.to(p_dh.dtype.element_ty), boundary_check=(0, 1)) + p_dh = dh + ((i_tg - 1) * H + i_h) * K*V + o_k[:, None] * V + o_v[None, :] + tl.store(p_dh, b_dh.to(p_dh.dtype.element_ty), mask=m_kv) elif STORE_INITIAL_STATE_GRADIENT: - p_dh0 = tl.make_block_ptr(dh0 + i_nh * K*V, (K, V), (V, 1), (i_k * BK, i_v * BV), (BK, BV), (1, 0)) - tl.store(p_dh0, b_dh.to(p_dh0.dtype.element_ty), boundary_check=(0, 1)) + p_dh0 = dh0 + i_nh * K*V + o_k[:, None] * V + o_v[None, :] + tl.store(p_dh0, b_dh.to(p_dh0.dtype.element_ty), mask=m_kv) @triton.heuristics({ @@ -365,14 +380,14 @@ def chunk_bwd_kernel_dh_reduction( STORE_INITIAL_STATE_GRADIENT: tl.constexpr, IS_VARLEN: tl.constexpr, ): - i_k, i_v, i_nh = tl.program_id(0), tl.program_id(1), tl.program_id(2) + i_k, i_v, i_nh = unflatten_program_id(tl.cdiv(K, BK), tl.cdiv(V, BV)) i_n, i_hq = i_nh // HQ, i_nh % HQ i_h = i_hq // NG if IS_VARLEN: - bos, eos = tl.load(cu_seqlens + i_n).to(tl.int32), tl.load(cu_seqlens + i_n + 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 NT = tl.cdiv(T, BT) - boh = tl.load(chunk_offsets + i_n).to(tl.int32) + boh = tl.load(chunk_offsets + i_n).to(tl.int64) else: bos, eos = i_n * T, i_n * T + T NT = tl.cdiv(T, BT) @@ -380,11 +395,14 @@ def chunk_bwd_kernel_dh_reduction( # [BK, BV] b_dh = tl.zeros([BK, BV], dtype=tl.float32) + o_k = i_k * BK + tl.arange(0, BK) + o_v = i_v * BV + tl.arange(0, BV) + m_kv = (o_k[:, None] < K) & (o_v[None, :] < V) for i_t in range(NT - 1, -1, -1): - p_dh = tl.make_block_ptr(dh + ((boh+i_t) * H + i_h) * K*V, (K, V), (V, 1), (i_k * BK, i_v * BV), (BK, BV), (1, 0)) - b_dh += tl.load(p_dh, boundary_check=(0, 1)).to(tl.float32) + p_dh = dh + ((boh+i_t) * H + i_h) * K*V + o_k[:, None] * V + o_v[None, :] + b_dh += tl.load(p_dh, mask=m_kv, other=0.0).to(tl.float32) if i_t < NT - 1: - tl.store(p_dh, b_dh.to(p_dh.dtype.element_ty), boundary_check=(0, 1)) + tl.store(p_dh, b_dh.to(p_dh.dtype.element_ty), mask=m_kv) last_idx = min(i_t * BT + BT, T) - 1 if USE_G: @@ -402,10 +420,10 @@ def chunk_bwd_kernel_dh_reduction( b_dh *= exp(b_gv_last)[None, :] if STORE_INITIAL_STATE_GRADIENT: - p_doq0 = tl.make_block_ptr(doq0 + i_nh * K*V, (K, V), (V, 1), (i_k * BK, i_v * BV), (BK, BV), (1, 0)) - p_dh0 = tl.make_block_ptr(dh0 + i_nh * K*V, (K, V), (V, 1), (i_k * BK, i_v * BV), (BK, BV), (1, 0)) - b_dh += tl.load(p_doq0, boundary_check=(0, 1)).to(tl.float32) - tl.store(p_dh0, b_dh.to(p_dh0.dtype.element_ty), boundary_check=(0, 1)) + p_doq0 = doq0 + i_nh * K*V + o_k[:, None] * V + o_v[None, :] + p_dh0 = dh0 + i_nh * K*V + o_k[:, None] * V + o_v[None, :] + b_dh += tl.load(p_doq0, mask=m_kv, other=0.0).to(tl.float32) + tl.store(p_dh0, b_dh.to(p_dh0.dtype.element_ty), mask=m_kv) def chunk_fwd_h( @@ -434,7 +452,7 @@ def chunk_fwd_h( h = k.new_empty(B, NT, H, K, V, dtype=torch.float) ht = k.new_empty(N, H, K, V, dtype=torch.float) if output_final_state else None - def grid(meta): return (triton.cdiv(K, meta['BK']) * triton.cdiv(V, meta['BV']), NT, B * H) + def grid(meta): return (triton.cdiv(K, meta['BK']) * triton.cdiv(V, meta['BV']) * NT, B * H) chunk_fwd_kernel_h_parallel[grid]( k=k, v=v, @@ -456,7 +474,7 @@ def chunk_fwd_h( USE_GV=gv is not None, ) kvt, ht = ht, (torch.empty_like(ht) if output_final_state else None) - def grid(meta): return (triton.cdiv(K, meta['BK']), triton.cdiv(V, meta['BV']), N * H) + def grid(meta): return (triton.cdiv(K, meta['BK']) * triton.cdiv(V, meta['BV']) * N * H,) chunk_fwd_kernel_h_reduction[grid]( h=h, g=g, @@ -512,7 +530,7 @@ def chunk_bwd_dh( dh = k.new_empty(B, NT, HQ, K, V, dtype=k.dtype if not states_in_fp32 else torch.float) dh0 = torch.empty_like(h0, dtype=torch.float) if h0 is not None else None - def grid(meta): return (triton.cdiv(K, meta['BK']) * triton.cdiv(V, meta['BV']), NT, B * HQ) + def grid(meta): return (triton.cdiv(K, meta['BK']) * triton.cdiv(V, meta['BV']) * NT, B * HQ) chunk_bwd_kernel_dh_parallel[grid]( q=q, g=g, @@ -538,7 +556,7 @@ def chunk_bwd_dh( ) doq0, dh0 = dh0, (torch.empty_like(dh0) if dh0 is not None else None) - def grid(meta): return (triton.cdiv(K, meta['BK']), triton.cdiv(V, meta['BV']), N * HQ) + def grid(meta): return (triton.cdiv(K, meta['BK']) * triton.cdiv(V, meta['BV']) * N * HQ,) chunk_bwd_kernel_dh_reduction[grid]( g=g, gk=gk, diff --git a/build/torch-cuda/ops/common/chunk_h_split.py b/build/torch-cuda/ops/common/chunk_h_split.py index 3956afc77407c1a2ae77fc5b1d73fa657399bd62..f9d4455ab0c7b2af1683c816c31cfc7f87d02a4a 100644 --- a/build/torch-cuda/ops/common/chunk_h_split.py +++ b/build/torch-cuda/ops/common/chunk_h_split.py @@ -9,7 +9,7 @@ import torch import triton import triton.language as tl -from ...ops.utils.op import exp +from ...ops.utils.op import exp, unflatten_program_id from ...utils import autotune_cache_kwargs @@ -61,11 +61,11 @@ def chunk_fwd_kernel_h_split( # i_h: head index # i_n: sequence index # i_s: local split index inside a sequence - i_k, i_v, i_sh = tl.program_id(0), tl.program_id(1), tl.program_id(2) + i_k, i_v, i_sh = unflatten_program_id(tl.cdiv(K, BK), tl.cdiv(V, BV)) i_ss, i_h = i_sh // H, i_sh % H if IS_VARLEN: i_n, i_s = tl.load(split_indices + i_ss * 2).to(tl.int32), tl.load(split_indices + i_ss * 2 + 1).to(tl.int32) - bos, eos = tl.load(cu_seqlens + i_n).to(tl.int32), tl.load(cu_seqlens + i_n + 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 NS = tl.cdiv(T, S) else: @@ -76,62 +76,68 @@ def chunk_fwd_kernel_h_split( # [BK, BV] b_h = tl.zeros([BK, BV], dtype=tl.float32) + o_k = i_k * BK + tl.arange(0, BK) + o_v = i_v * BV + tl.arange(0, BV) + m_kv = (o_k[:, None] < K) & (o_v[None, :] < V) # for the first split, we directly store the state as the final result if i_s == 0: if USE_INITIAL_STATE: - p_h0 = tl.make_block_ptr(h0 + i_nh * K*V, (K, V), (V, 1), (i_k * BK, i_v * BV), (BK, BV), (1, 0)) - b_h += tl.load(p_h0, boundary_check=(0, 1)).to(tl.float32) - p_hr = tl.make_block_ptr(hr + i_sh * K*V, (K, V), (V, 1), (i_k * BK, i_v * BV), (BK, BV), (1, 0)) - tl.store(p_hr, b_h.to(p_hr.dtype.element_ty), boundary_check=(0, 1)) + p_h0 = h0 + i_nh * K*V + o_k[:, None] * V + o_v[None, :] + b_h += tl.load(p_h0, mask=m_kv, other=0.0).to(tl.float32) + p_hr = hr + i_sh * K*V + o_k[:, None] * V + o_v[None, :] + tl.store(p_hr, b_h.to(p_hr.dtype.element_ty), mask=m_kv) for i_t in range(tl.cdiv(i_s * S, BT), tl.cdiv(min(i_s * S + S, T), BT)): - p_k = tl.make_block_ptr(k + (bos*H + i_h) * K, (K, T), (1, H*K), (i_k * BK, i_t * BT), (BK, BT), (0, 1)) - p_v = tl.make_block_ptr(v + (bos*H + i_h) * V, (T, V), (H*V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) + i_t_int64 = i_t.to(tl.int64) + o_t = i_t_int64 * BT + tl.arange(0, BT) + m_t = o_t < T + p_k = k + (bos*H + i_h) * K + o_k[:, None] + o_t[None, :] * (H*K) + p_v = v + (bos*H + i_h) * V + o_t[:, None] * (H*V) + o_v[None, :] # [BK, BT] - b_k = tl.load(p_k, boundary_check=(0, 1)) + b_k = tl.load(p_k, mask=(o_k[:, None] < K) & m_t[None, :], other=0.0) # [BT, BV] - b_v = tl.load(p_v, boundary_check=(0, 1)) - last_idx = min(i_t * BT + BT, T) - 1 + b_v = tl.load(p_v, mask=m_t[:, None] & (o_v < V)[None, :], other=0.0) + last_idx = min(i_t_int64 * BT + BT, T) - 1 # scalar decay if USE_G: b_g_last = tl.load(g + bos * H + last_idx * H + i_h) - p_g = g + bos*H + (i_t * BT + tl.arange(0, BT)) * H + i_h + p_g = g + bos*H + (i_t_int64 * BT + tl.arange(0, BT)) * H + i_h b_h *= exp(b_g_last) - b_g = tl.load(p_g, mask=(i_t * BT + tl.arange(0, BT) < T), other=0.) + b_g = tl.load(p_g, mask=(i_t_int64 * BT + tl.arange(0, BT) < T), other=0.) b_v = (b_v * exp(b_g_last - b_g)[:, None]).to(b_v.dtype) # vector decay, h = Diag(gk) @ h if USE_GK: - p_gk = tl.make_block_ptr(gk + (bos*H + i_h) * K, (K, T), (1, H*K), (i_k * BK, i_t * BT), (BK, BT), (0, 1)) + p_gk = gk + (bos*H + i_h) * K + o_k[:, None] + o_t[None, :] * (H*K) p_gk_last = gk + (bos + last_idx) * H*K + i_h * K + i_k * BK + tl.arange(0, BK) b_gk_last = tl.load(p_gk_last, mask=(i_k * BK + tl.arange(0, BK) < K), other=0.) b_h *= exp(b_gk_last)[:, None] - b_gk = tl.load(p_gk, boundary_check=(0, 1)) + b_gk = tl.load(p_gk, mask=(o_k[:, None] < K) & m_t[None, :], other=0.0) b_k = (b_k * exp(b_gk_last[:, None] - b_gk)).to(b_k.dtype) # vector decay, h = h @ Diag(gv) if USE_GV: - p_gv = tl.make_block_ptr(gv + (bos*H + i_h) * V, (T, V), (H*V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) + p_gv = gv + (bos*H + i_h) * V + o_t[:, None] * (H*V) + o_v[None, :] p_gv_last = gv + (bos + last_idx) * H*V + i_h * V + i_v * BV + tl.arange(0, BV) b_gv_last = tl.load(p_gv_last, mask=(i_v * BV + tl.arange(0, BV) < V), other=0.) b_h *= exp(b_gv_last)[None, :] - b_gv = tl.load(p_gv, boundary_check=(0, 1)) + b_gv = tl.load(p_gv, mask=m_t[:, None] & (o_v < V)[None, :], other=0.0) b_v = (b_v * exp(b_gv_last[None, :] - b_gv)).to(b_v.dtype) - b_h += tl.dot(b_k, b_v) + b_h = tl.dot(b_k, b_v, b_h) # if there are more than one splits, we store the result to (unreduced) hs # otherwise, we store the result to ht as the final state if NS > 1: - p_hs = tl.make_block_ptr(hs + i_sh * K*V, (K, V), (V, 1), (i_k * BK, i_v * BV), (BK, BV), (1, 0)) - tl.store(p_hs, b_h.to(p_hs.dtype.element_ty), boundary_check=(0, 1)) + p_hs = hs + i_sh * K*V + o_k[:, None] * V + o_v[None, :] + tl.store(p_hs, b_h.to(p_hs.dtype.element_ty), mask=m_kv) elif STORE_FINAL_STATE: - p_ht = tl.make_block_ptr(ht + i_nh * K*V, (K, V), (V, 1), (i_k * BK, i_v * BV), (BK, BV), (1, 0)) - tl.store(p_ht, b_h.to(p_ht.dtype.element_ty), boundary_check=(0, 1)) + p_ht = ht + i_nh * K*V + o_k[:, None] * V + o_v[None, :] + tl.store(p_ht, b_h.to(p_ht.dtype.element_ty), mask=m_kv) @triton.heuristics({ @@ -173,25 +179,28 @@ def chunk_fwd_kernel_h_reduction( STORE_FINAL_STATE: tl.constexpr, IS_VARLEN: tl.constexpr, ): - i_k, i_v, i_nh = tl.program_id(0), tl.program_id(1), tl.program_id(2) + i_k, i_v, i_nh = unflatten_program_id(tl.cdiv(K, BK), tl.cdiv(V, BV)) i_n, i_h = i_nh // H, i_nh % H if IS_VARLEN: - bos, eos = tl.load(cu_seqlens + i_n).to(tl.int32), tl.load(cu_seqlens + i_n + 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 NS = tl.cdiv(T, S) - boh = tl.load(split_offsets + i_n).to(tl.int32) + boh = tl.load(split_offsets + i_n).to(tl.int64) else: bos, eos = i_n * T, i_n * T + T NS = tl.cdiv(T, S) boh = i_n * NS b_h = tl.zeros([BK, BV], dtype=tl.float32) + o_k = i_k * BK + tl.arange(0, BK) + o_v = i_v * BV + tl.arange(0, BV) + m_kv = (o_k[:, None] < K) & (o_v[None, :] < V) # skip the first split for i_s in range(1, NS): - p_hs = tl.make_block_ptr(hs + ((boh + i_s-1) * H + i_h) * K*V, (K, V), (V, 1), (i_k * BK, i_v * BV), (BK, BV), (1, 0)) - p_hr = tl.make_block_ptr(hr + ((boh + i_s) * H + i_h) * K*V, (K, V), (V, 1), (i_k * BK, i_v * BV), (BK, BV), (1, 0)) - b_h += tl.load(p_hs, boundary_check=(0, 1)).to(tl.float32) - tl.store(p_hr, b_h.to(p_hr.dtype.element_ty), boundary_check=(0, 1)) + p_hs = hs + ((boh + i_s-1) * H + i_h) * K*V + o_k[:, None] * V + o_v[None, :] + p_hr = hr + ((boh + i_s) * H + i_h) * K*V + o_k[:, None] * V + o_v[None, :] + b_h += tl.load(p_hs, mask=m_kv, other=0.0).to(tl.float32) + tl.store(p_hr, b_h.to(p_hr.dtype.element_ty), mask=m_kv) for i_t in range(tl.cdiv(i_s * S, BT), tl.cdiv(min(i_s * S + S, T), BT)): last_idx = min(i_t * BT + BT, T) - 1 @@ -214,10 +223,10 @@ def chunk_fwd_kernel_h_reduction( if NS > 1: if STORE_FINAL_STATE: - p_hs = tl.make_block_ptr(hs + ((boh + NS-1) * H + i_h)*K*V, (K, V), (V, 1), (i_k * BK, i_v * BV), (BK, BV), (1, 0)) - p_ht = tl.make_block_ptr(ht + i_nh * K*V, (K, V), (V, 1), (i_k * BK, i_v * BV), (BK, BV), (1, 0)) - b_h += tl.load(p_hs, boundary_check=(0, 1)).to(tl.float32) - tl.store(p_ht, b_h.to(p_ht.dtype.element_ty), boundary_check=(0, 1)) + p_hs = hs + ((boh + NS-1) * H + i_h)*K*V + o_k[:, None] * V + o_v[None, :] + p_ht = ht + i_nh * K*V + o_k[:, None] * V + o_v[None, :] + b_h += tl.load(p_hs, mask=m_kv, other=0.0).to(tl.float32) + tl.store(p_ht, b_h.to(p_ht.dtype.element_ty), mask=m_kv) @triton.heuristics({ @@ -271,11 +280,11 @@ def chunk_bwd_kernel_dh_split( # i_h: head index # i_n: sequence index # i_s: local split index inside a sequence - i_k, i_v, i_sh = tl.program_id(0), tl.program_id(1), tl.program_id(2) + i_k, i_v, i_sh = unflatten_program_id(tl.cdiv(K, BK), tl.cdiv(V, BV)) i_ss, i_hq = i_sh // HQ, i_sh % HQ if IS_VARLEN: i_n, i_s = tl.load(split_indices + i_ss * 2).to(tl.int32), tl.load(split_indices + i_ss * 2 + 1).to(tl.int32) - bos, eos = tl.load(cu_seqlens + i_n).to(tl.int32), tl.load(cu_seqlens + i_n + 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 NS = tl.cdiv(T, S) else: @@ -287,57 +296,63 @@ def chunk_bwd_kernel_dh_split( # [BK, BV] b_dh = tl.zeros([BK, BV], dtype=tl.float32) + o_k = i_k * BK + tl.arange(0, BK) + o_v = i_v * BV + tl.arange(0, BV) + m_kv = (o_k[:, None] < K) & (o_v[None, :] < V) if i_s == NS - 1: if USE_FINAL_STATE_GRADIENT: - p_dht = tl.make_block_ptr(dht + i_nh * K*V, (K, V), (V, 1), (i_k * BK, i_v * BV), (BK, BV), (1, 0)) - b_dh += tl.load(p_dht, boundary_check=(0, 1)).to(tl.float32) - p_dhr = tl.make_block_ptr(dhr + i_sh * K*V, (K, V), (V, 1), (i_k * BK, i_v * BV), (BK, BV), (1, 0)) - tl.store(p_dhr, b_dh.to(p_dhr.dtype.element_ty), boundary_check=(0, 1)) + p_dht = dht + i_nh * K*V + o_k[:, None] * V + o_v[None, :] + b_dh += tl.load(p_dht, mask=m_kv, other=0.0).to(tl.float32) + p_dhr = dhr + i_sh * K*V + o_k[:, None] * V + o_v[None, :] + tl.store(p_dhr, b_dh.to(p_dhr.dtype.element_ty), mask=m_kv) for i_t in range(tl.cdiv(min(i_s * S + S, T), BT) - 1, tl.cdiv(i_s * S, BT) - 1, -1): - p_q = tl.make_block_ptr(q + (bos*HQ + i_hq) * K, (K, T), (1, HQ*K), (i_k * BK, i_t * BT), (BK, BT), (0, 1)) - p_do = tl.make_block_ptr(do + (bos*HQ + i_hq) * V, (T, V), (HQ*V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) + i_t_int64 = i_t.to(tl.int64) + o_t = i_t_int64 * BT + tl.arange(0, BT) + m_t = o_t < T + p_q = q + (bos*HQ + i_hq) * K + o_k[:, None] + o_t[None, :] * (HQ*K) + p_do = do + (bos*HQ + i_hq) * V + o_t[:, None] * (HQ*V) + o_v[None, :] - b_q = tl.load(p_q, boundary_check=(0, 1)) + b_q = tl.load(p_q, mask=(o_k[:, None] < K) & m_t[None, :], other=0.0) b_q = (b_q * scale).to(b_q.dtype) # [BT, BV] - b_do = tl.load(p_do, boundary_check=(0, 1)) + b_do = tl.load(p_do, mask=m_t[:, None] & (o_v < V)[None, :], other=0.0) - last_idx = min(i_t * BT + BT, T) - 1 + last_idx = min(i_t_int64 * BT + BT, T) - 1 if USE_G: - p_g = g + (bos + i_t * BT + tl.arange(0, BT)) * H + i_h + p_g = g + (bos + i_t_int64 * BT + tl.arange(0, BT)) * H + i_h b_g_last = tl.load(g + (bos + last_idx) * H + i_h) - b_g = tl.load(p_g, mask=(i_t * BT + tl.arange(0, BT) < T), other=0.) + b_g = tl.load(p_g, mask=(i_t_int64 * BT + tl.arange(0, BT) < T), other=0.) b_q = (b_q * exp(b_g)[None, :]).to(b_q.dtype) b_dh *= exp(b_g_last) if USE_GK: - p_gk = tl.make_block_ptr(gk + (bos*H + i_h) * K, (K, T), (1, H*K), (i_k * BK, i_t * BT), (BK, BT), (0, 1)) + p_gk = gk + (bos*H + i_h) * K + o_k[:, None] + o_t[None, :] * (H*K) p_gk_last = gk + (bos + last_idx) * H*K + i_h * K + i_k * BK + tl.arange(0, BK) - b_gk = tl.load(p_gk, boundary_check=(0, 1)) + b_gk = tl.load(p_gk, mask=(o_k[:, None] < K) & m_t[None, :], other=0.0) b_q = (b_q * exp(b_gk)).to(b_q.dtype) b_gk_last = tl.load(p_gk_last, mask=(i_k * BK + tl.arange(0, BK) < K), other=0.) b_dh *= exp(b_gk_last)[:, None] if USE_GV: - p_gv = tl.make_block_ptr(gv + (bos*H + i_h) * V, (T, V), (H*V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) + p_gv = gv + (bos*H + i_h) * V + o_t[:, None] * (H*V) + o_v[None, :] p_gv_last = gv + (bos + last_idx) * H*V + i_h * V + i_v * BV + tl.arange(0, BV) - b_gv = tl.load(p_gv, boundary_check=(0, 1)) + b_gv = tl.load(p_gv, mask=m_t[:, None] & (o_v < V)[None, :], other=0.0) b_do = (b_do * exp(b_gv)).to(b_do.dtype) b_gv_last = tl.load(p_gv_last, mask=(i_v * BV + tl.arange(0, BV) < V), other=0.) b_dh *= exp(b_gv_last)[None, :] - b_dh += tl.dot(b_q, b_do) + b_dh = tl.dot(b_q, b_do, b_dh) if NS > 1: - p_dhs = tl.make_block_ptr(dhs + i_sh * K*V, (K, V), (V, 1), (i_k * BK, i_v * BV), (BK, BV), (1, 0)) - tl.store(p_dhs, b_dh.to(p_dhs.dtype.element_ty), boundary_check=(0, 1)) + p_dhs = dhs + i_sh * K*V + o_k[:, None] * V + o_v[None, :] + tl.store(p_dhs, b_dh.to(p_dhs.dtype.element_ty), mask=m_kv) elif STORE_INITIAL_STATE_GRADIENT: - p_dh0 = tl.make_block_ptr(dh0 + i_nh * K*V, (K, V), (V, 1), (i_k * BK, i_v * BV), (BK, BV), (1, 0)) - tl.store(p_dh0, b_dh.to(p_dh0.dtype.element_ty), boundary_check=(0, 1)) + p_dh0 = dh0 + i_nh * K*V + o_k[:, None] * V + o_v[None, :] + tl.store(p_dh0, b_dh.to(p_dh0.dtype.element_ty), mask=m_kv) @triton.heuristics({ @@ -381,25 +396,28 @@ def chunk_bwd_kernel_dh_reduction( STORE_INITIAL_STATE_GRADIENT: tl.constexpr, IS_VARLEN: tl.constexpr, ): - i_k, i_v, i_nh = tl.program_id(0), tl.program_id(1), tl.program_id(2) + i_k, i_v, i_nh = unflatten_program_id(tl.cdiv(K, BK), tl.cdiv(V, BV)) i_n, i_hq = i_nh // HQ, i_nh % HQ i_h = i_hq // NG if IS_VARLEN: - bos, eos = tl.load(cu_seqlens + i_n).to(tl.int32), tl.load(cu_seqlens + i_n + 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 NS = tl.cdiv(T, S) - boh = tl.load(split_offsets + i_n).to(tl.int32) + boh = tl.load(split_offsets + i_n).to(tl.int64) else: bos, eos = i_n * T, i_n * T + T NS = tl.cdiv(T, S) boh = i_n * NS b_dh = tl.zeros([BK, BV], dtype=tl.float32) + o_k = i_k * BK + tl.arange(0, BK) + o_v = i_v * BV + tl.arange(0, BV) + m_kv = (o_k[:, None] < K) & (o_v[None, :] < V) for i_s in range(NS - 2, -1, -1): - p_dhs = tl.make_block_ptr(dhs + ((boh+i_s+1) * H + i_h) * K*V, (K, V), (V, 1), (i_k * BK, i_v * BV), (BK, BV), (1, 0)) - p_dhr = tl.make_block_ptr(dhr + ((boh+i_s) * H + i_h) * K*V, (K, V), (V, 1), (i_k * BK, i_v * BV), (BK, BV), (1, 0)) - b_dh += tl.load(p_dhs, boundary_check=(0, 1)).to(tl.float32) - tl.store(p_dhr, b_dh.to(p_dhr.dtype.element_ty), boundary_check=(0, 1)) + p_dhs = dhs + ((boh+i_s+1) * H + i_h) * K*V + o_k[:, None] * V + o_v[None, :] + p_dhr = dhr + ((boh+i_s) * H + i_h) * K*V + o_k[:, None] * V + o_v[None, :] + b_dh += tl.load(p_dhs, mask=m_kv, other=0.0).to(tl.float32) + tl.store(p_dhr, b_dh.to(p_dhr.dtype.element_ty), mask=m_kv) for i_t in range(tl.cdiv(min(i_s * S + S, T), BT) - 1, tl.cdiv(i_s * S, BT) - 1, -1): last_idx = min(i_t * BT + BT, T) - 1 @@ -420,10 +438,10 @@ def chunk_bwd_kernel_dh_reduction( if NS > 1: if STORE_INITIAL_STATE_GRADIENT: - p_dhs = tl.make_block_ptr(dhs + (boh * H + i_h)*K*V, (K, V), (V, 1), (i_k * BK, i_v * BV), (BK, BV), (1, 0)) - p_dh0 = tl.make_block_ptr(dh0 + i_nh * K*V, (K, V), (V, 1), (i_k * BK, i_v * BV), (BK, BV), (1, 0)) - b_dh += tl.load(p_dhs, boundary_check=(0, 1)).to(tl.float32) - tl.store(p_dh0, b_dh.to(p_dh0.dtype.element_ty), boundary_check=(0, 1)) + p_dhs = dhs + (boh * H + i_h)*K*V + o_k[:, None] * V + o_v[None, :] + p_dh0 = dh0 + i_nh * K*V + o_k[:, None] * V + o_v[None, :] + b_dh += tl.load(p_dhs, mask=m_kv, other=0.0).to(tl.float32) + tl.store(p_dh0, b_dh.to(p_dh0.dtype.element_ty), mask=m_kv) def chunk_fwd_h( @@ -463,7 +481,7 @@ def chunk_fwd_h( hr = k.new_empty(NS, H, K, V, dtype=torch.float if states_in_fp32 else k.dtype) ht = k.new_empty(N, H, K, V, dtype=torch.float) if output_final_state else None # parallelized over splits - def grid(meta): return (triton.cdiv(K, meta['BK']), triton.cdiv(V, meta['BV']), NS * H) + def grid(meta): return (triton.cdiv(K, meta['BK']) * triton.cdiv(V, meta['BV']) * NS * H,) chunk_fwd_kernel_h_split[grid]( k=k, v=v, @@ -486,7 +504,7 @@ def chunk_fwd_h( USE_GK=gk is not None, USE_GV=gv is not None, ) - def grid(meta): return (triton.cdiv(K, meta['BK']), triton.cdiv(V, meta['BV']), N * H) + def grid(meta): return (triton.cdiv(K, meta['BK']) * triton.cdiv(V, meta['BV']) * N * H,) chunk_fwd_kernel_h_reduction[grid]( g=g, gk=gk, @@ -551,7 +569,7 @@ def chunk_bwd_dh( dh0 = torch.empty_like(h0, dtype=torch.float) if h0 is not None else None # parallelized over splits - def grid(meta): return (triton.cdiv(K, meta['BK']), triton.cdiv(V, meta['BV']), NS * HQ) + def grid(meta): return (triton.cdiv(K, meta['BK']) * triton.cdiv(V, meta['BV']) * NS * HQ,) chunk_bwd_kernel_dh_split[grid]( q=q, g=g, @@ -578,7 +596,7 @@ def chunk_bwd_dh( USE_GV=gv is not None, ) - def grid(meta): return (triton.cdiv(K, meta['BK']), triton.cdiv(V, meta['BV']), N * HQ) + def grid(meta): return (triton.cdiv(K, meta['BK']) * triton.cdiv(V, meta['BV']) * N * HQ,) chunk_bwd_kernel_dh_reduction[grid]( g=g, gk=gk, diff --git a/build/torch-cuda/ops/common/chunk_o.py b/build/torch-cuda/ops/common/chunk_o.py index 868d012b9733cfa156e3429c5ee696e24f998792..f3d97bc13741076ecc0199c32a381b33149217ec 100644 --- a/build/torch-cuda/ops/common/chunk_o.py +++ b/build/torch-cuda/ops/common/chunk_o.py @@ -9,28 +9,34 @@ import torch import triton import triton.language as tl -from ...ops.common.backends import dispatch +from ...backends import dispatch from ...ops.utils import prepare_chunk_indices from ...ops.utils.cache import fla_cache_autotune -from ...ops.utils.op import exp2 +from ...ops.utils.op import exp2, unflatten_program_id from ...utils import ( IS_INTEL, + IS_NVIDIA_BLACKWELL, IS_NVIDIA_HOPPER, TRITON_ABOVE_3_4_0, TRITON_ABOVE_3_7_1, + TRITON_ABOVE_3_8_0, autotune_cache_kwargs, check_shared_mem, ) BKV_LIST = [64, 128] if check_shared_mem() else ([32, 64] if check_shared_mem('ada') else [32]) NUM_WARPS = [2, 4] if IS_NVIDIA_HOPPER else [2, 4, 8] +# Triton<3.8 races at 8 warps on Blackwell: triton-lang/triton#10590, fla-org/flash-linear-attention#1327 +BLACKWELL_8_WARP_RACE = IS_NVIDIA_BLACKWELL and not TRITON_ABOVE_3_8_0 +CHUNK_BWD_DV_LOCAL_NUM_WARPS = [2, 4] if BLACKWELL_8_WARP_RACE else NUM_WARPS -# Upstream pairs each tile size with a single warp count, which is tuned for NVIDIA. -# On Intel that pairing is off by a factor of two: BK=BV=64 is fastest at 8 warps but is -# only offered at 4, so the autotuner falls back to the 32x32 config and leaves ~2.2x on -# the table. Widen the space there instead of changing the defaults for other vendors. -_O_CONFIGS = [ +# Intel needs 8 warps for BK=BV=64; the NVIDIA pairing only offers 4 and falls back to slower 32x32 tiles. +CHUNK_FWD_O_BLACKWELL_DROPPED_CONFIGS = [] if BLACKWELL_8_WARP_RACE else [ triton.Config({'BK': 128, 'BV': 128}, num_warps=8, num_stages=3), +] + +_O_CONFIGS = [ + *CHUNK_FWD_O_BLACKWELL_DROPPED_CONFIGS, triton.Config({'BK': 64, 'BV': 64}, num_warps=4, num_stages=3), triton.Config({'BK': 32, 'BV': 32}, num_warps=2, num_stages=3), ] @@ -47,11 +53,7 @@ if IS_INTEL: 'USE_G_GAMMA': lambda args: args['g_gamma'] is not None, 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None, }) -@fla_cache_autotune( - configs=_O_CONFIGS, - key=['H', 'HV', 'K', 'V', 'BT', 'STATE_V_FIRST'], - **autotune_cache_kwargs, -) +@fla_cache_autotune(configs=_O_CONFIGS, key=['H', 'HV', 'K', 'V', 'BT', 'STATE_V_FIRST'], **autotune_cache_kwargs) @triton.jit(do_not_specialize=['T']) def chunk_fwd_kernel_o( q, @@ -77,13 +79,14 @@ def chunk_fwd_kernel_o( STATE_V_FIRST: tl.constexpr, IS_VARLEN: tl.constexpr, ): - i_v, i_t, i_bh = tl.program_id(0), tl.program_id(1), tl.program_id(2).to(tl.int64) + i_v, i_t = unflatten_program_id(tl.cdiv(V, BV)) + i_bh = tl.program_id(1).to(tl.int64) i_b, i_h = i_bh // HV, i_bh % HV 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.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) + i_n, i_t = tl.load(chunk_indices + i_t * 2).to(tl.int32), tl.load(chunk_indices + i_t * 2 + 1).to(tl.int64) + bos, eos = tl.load(cu_seqlens + i_n).to(tl.int64), tl.load(cu_seqlens + i_n + 1).to(tl.int64) T = eos - bos NT = tl.cdiv(T, BT) else: @@ -91,7 +94,6 @@ def chunk_fwd_kernel_o( i_tg = i_b * NT + i_t bos, eos = i_b * T, i_b * T + T - # offset calculation q += (bos * H + i_h // (HV // H)) * K k += (bos * H + i_h // (HV // H)) * K v += (bos * HV + i_h) * V @@ -101,31 +103,38 @@ def chunk_fwd_kernel_o( b_o = tl.zeros([BT, BV], dtype=tl.float32) b_A = tl.zeros([BT, BT], dtype=tl.float32) + o_t = i_t * BT + tl.arange(0, BT) + m_t = o_t < T + o_v = i_v * BV + tl.arange(0, BV) for i_k in range(tl.cdiv(K, BK)): - p_q = tl.make_block_ptr(q, (T, K), (H*K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) - p_k = tl.make_block_ptr(k, (K, T), (1, H*K), (i_k * BK, i_t * BT), (BK, BT), (0, 1)) + o_k = i_k * BK + tl.arange(0, BK) + m_k = o_k < K + p_q = q + o_t[:, None] * (H*K) + o_k[None, :] + p_k = k + o_k[:, None] + o_t[None, :] * (H*K) if STATE_V_FIRST: - p_h = tl.make_block_ptr(h, (V, K), (K, 1), (i_v * BV, i_k * BK), (BV, BK), (1, 0)) + p_h = h + o_v[:, None] * K + o_k[None, :] + m_h = (o_v[:, None] < V) & m_k[None, :] else: - p_h = tl.make_block_ptr(h, (K, V), (V, 1), (i_k * BK, i_v * BV), (BK, BV), (1, 0)) + p_h = h + o_k[:, None] * V + o_v[None, :] + m_h = m_k[:, None] & (o_v[None, :] < V) # [BT, BK] - b_q = tl.load(p_q, boundary_check=(0, 1)) + b_q = tl.load(p_q, mask=m_t[:, None] & m_k[None, :], other=0.0) # [BK, BT] - b_k = tl.load(p_k, boundary_check=(0, 1)) - b_h = tl.load(p_h, boundary_check=(0, 1)) + b_k = tl.load(p_k, mask=m_k[:, None] & m_t[None, :], other=0.0) + b_h = tl.load(p_h, mask=m_h, other=0.0) # [BT, BK] @ [BK, BV] -> [BT, BV] if STATE_V_FIRST: - b_o += tl.dot(b_q, tl.trans(b_h)) + b_o = tl.dot(b_q, tl.trans(b_h), b_o) else: - b_o += tl.dot(b_q, b_h) + b_o = tl.dot(b_q, b_h, b_o) # [BT, BK] @ [BK, BT] -> [BT, BT] - b_A += tl.dot(b_q, b_k) + b_A = tl.dot(b_q, b_k, b_A) if USE_G: g += bos * HV + i_h - p_g = tl.make_block_ptr(g, (T,), (HV,), (i_t * BT,), (BT,), (0,)) - b_g = tl.load(p_g, boundary_check=(0,)) + p_g = g + o_t * HV + b_g = tl.load(p_g, mask=m_t, other=0.0) b_o = b_o * exp2(b_g)[:, None] b_A = b_A * exp2(b_g[:, None] - b_g[None, :]) if USE_G_GAMMA: @@ -133,19 +142,17 @@ def chunk_fwd_kernel_o( b_g = b_gamma * (tl.arange(0, BT) + 1) b_o = b_o * exp2(b_g)[:, None] b_A = b_A * exp2(b_g[:, None] - b_g[None, :]) - o_t = i_t * BT + tl.arange(0, BT) - m_t = o_t < T m_A = (o_t[:, None] >= o_t[None, :]) & (m_t[:, None] & m_t) b_A = tl.where(m_A, b_A, 0) - p_v = tl.make_block_ptr(v, (T, V), (HV*V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) - p_o = tl.make_block_ptr(o, (T, V), (HV*V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) + p_v = v + o_t[:, None] * (HV*V) + o_v[None, :] + p_o = o + o_t[:, None] * (HV*V) + o_v[None, :] - b_v = tl.load(p_v, boundary_check=(0, 1)) + b_v = tl.load(p_v, mask=m_t[:, None] & (o_v < V)[None, :], other=0.0) # to fix mma -> mma layout conversion # already solved by triton v3.2 or higher b_o = b_o * scale + tl.dot(b_A.to(b_v.dtype), b_v) * scale - tl.store(p_o, b_o.to(p_o.dtype.element_ty), boundary_check=(0, 1)) + tl.store(p_o, b_o.to(p_o.dtype.element_ty), mask=m_t[:, None] & (o_v < V)[None, :]) @triton.heuristics({ @@ -196,14 +203,16 @@ def chunk_bwd_kernel_dqkwg( STATE_V_FIRST: tl.constexpr, IS_VARLEN: tl.constexpr, ): - i_k, i_t, i_bh = tl.program_id(0), tl.program_id(1), tl.program_id(2).to(tl.int64) + scale = scale.to(tl.float32) + i_k, i_t = unflatten_program_id(tl.cdiv(K, BK)) + i_bh = tl.program_id(1).to(tl.int64) i_b, i_h = i_bh // HV, i_bh % HV - all = B * T + all = (B * T).to(tl.int64) 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.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) + i_n, i_t = tl.load(chunk_indices + i_t * 2).to(tl.int32), tl.load(chunk_indices + i_t * 2 + 1).to(tl.int64) + bos, eos = tl.load(cu_seqlens + i_n).to(tl.int64), tl.load(cu_seqlens + i_n + 1).to(tl.int64) T = eos - bos NT = tl.cdiv(T, BT) else: @@ -211,7 +220,6 @@ def chunk_bwd_kernel_dqkwg( i_tg = i_b * NT + i_t bos, eos = i_b * T, i_b * T + T - # offset calculation v += (bos * HV + i_h) * V do += (bos * HV + i_h) * V h += (i_tg * HV + i_h).to(tl.int64) * K*V @@ -238,55 +246,62 @@ def chunk_bwd_kernel_dqkwg( b_ds = tl.zeros([BT, BT], dtype=tl.float32) b_dw = tl.zeros([BT, BK], dtype=tl.float32) if USE_DW else None + o_t = i_t * BT + tl.arange(0, BT) + m_t = o_t < T + o_k = i_k * BK + tl.arange(0, BK) + m_k = o_k < K + m_qk = m_t[:, None] & m_k[None, :] for i_v in range(tl.cdiv(V, BV)): - p_v = tl.make_block_ptr(v, (T, V), (HV*V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) - p_do = tl.make_block_ptr(do, (T, V), (HV*V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) + o_v = i_v * BV + tl.arange(0, BV) + m_v = o_v < V + m_hv = m_t[:, None] & m_v[None, :] + p_v = v + o_t[:, None] * (HV*V) + o_v[None, :] + p_do = do + o_t[:, None] * (HV*V) + o_v[None, :] if STATE_V_FIRST: - p_h = tl.make_block_ptr(h, (V, K), (K, 1), (i_v * BV, i_k * BK), (BV, BK), (1, 0)) - p_dh = tl.make_block_ptr(dh, (V, K), (K, 1), (i_v * BV, i_k * BK), (BV, BK), (1, 0)) + p_h = h + o_v[:, None] * K + o_k[None, :] + p_dh = dh + o_v[:, None] * K + o_k[None, :] else: - p_h = tl.make_block_ptr(h, (V, K), (1, V), (i_v * BV, i_k * BK), (BV, BK), (0, 1)) - p_dh = tl.make_block_ptr(dh, (V, K), (1, V), (i_v * BV, i_k * BK), (BV, BK), (0, 1)) + p_h = h + o_v[:, None] + o_k[None, :] * V + p_dh = dh + o_v[:, None] + o_k[None, :] * V + m_h = m_v[:, None] & m_k[None, :] # [BT, BV] - b_v = tl.load(p_v, boundary_check=(0, 1)) - b_do = tl.load(p_do, boundary_check=(0, 1)) + b_v = tl.load(p_v, mask=m_hv, other=0.0) + b_do = tl.load(p_do, mask=m_hv, other=0.0) # [BV, BK] - b_h = tl.load(p_h, boundary_check=(0, 1)) - b_dh = tl.load(p_dh, boundary_check=(0, 1)) + b_h = tl.load(p_h, mask=m_h, other=0.0) + b_dh = tl.load(p_dh, mask=m_h, other=0.0) if USE_G: b_dg_last += (tl.sum(b_h * b_dh)) # [BT, BV] @ [BV, BT] -> [BT, BT] - b_ds += tl.dot(b_do, tl.trans(b_v)) + b_ds = tl.dot(b_do, tl.trans(b_v), b_ds) # [BT, BV] @ [BV, BK] -> [BT, BK] - b_dq += tl.dot(b_do, b_h.to(b_do.dtype)) + b_dq = tl.dot(b_do, b_h.to(b_do.dtype), b_dq) # [BT, BV] @ [BV, BK] -> [BT, BK] - b_dk += tl.dot(b_v, b_dh.to(b_v.dtype)) + b_dk = tl.dot(b_v, b_dh.to(b_v.dtype), b_dk) if USE_DW: - p_dv = tl.make_block_ptr(dv, (T, V), (HV*V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) - b_dv = tl.load(p_dv, boundary_check=(0, 1)) - b_dw += tl.dot(b_dv.to(b_v.dtype), b_h.to(b_v.dtype)) + p_dv = dv + o_t[:, None] * (HV*V) + o_v[None, :] + b_dv = tl.load(p_dv, mask=m_hv, other=0.0) + b_dw = tl.dot(b_dv.to(b_v.dtype), b_h.to(b_v.dtype), b_dw) if USE_DW: - p_dw = tl.make_block_ptr(dw, (T, K), (HV*K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) - tl.store(p_dw, -b_dw.to(p_dw.dtype.element_ty), boundary_check=(0, 1)) + p_dw = dw + o_t[:, None] * (HV*K) + o_k[None, :] + tl.store(p_dw, -b_dw.to(p_dw.dtype.element_ty), mask=m_qk) tl.debug_barrier() - p_q = tl.make_block_ptr(q, (T, K), (H*K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) - p_k = tl.make_block_ptr(k, (T, K), (H*K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) - b_q = tl.load(p_q, boundary_check=(0, 1)) - b_k = tl.load(p_k, boundary_check=(0, 1)) + p_q = q + o_t[:, None] * (H*K) + o_k[None, :] + p_k = k + o_t[:, None] * (H*K) + o_k[None, :] + b_q = tl.load(p_q, mask=m_qk, other=0.0) + b_k = tl.load(p_k, mask=m_qk, other=0.0) - p_dq = tl.make_block_ptr(dq, (T, K), (HV*K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) - p_dk = tl.make_block_ptr(dk, (T, K), (HV*K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) + p_dq = dq + o_t[:, None] * (HV*K) + o_k[None, :] + p_dk = dk + o_t[:, None] * (HV*K) + o_k[None, :] - o_t = i_t * BT + tl.arange(0, BT) - m_t = o_t < T m_A = (o_t[:, None] >= o_t[None, :]) & (m_t[:, None] & m_t) if USE_G: g += bos * HV + i_h dg += bos * HV + i_h - p_g = tl.make_block_ptr(g, (T,), (HV,), (i_t * BT,), (BT,), (0,)) - b_g = tl.load(p_g, boundary_check=(0,)) + p_g = g + o_t * HV + b_g = tl.load(p_g, mask=m_t, other=0.0) b_g_last = tl.load(g + (min(i_t * BT + BT, T) - 1) * HV) b_dg_last *= exp2(b_g_last) b_dq = b_dq * exp2(b_g)[:, None] * scale @@ -296,18 +311,17 @@ def chunk_bwd_kernel_dqkwg( b_ds = tl.where(m_A, b_ds * exp2(b_g[:, None] - b_g[None, :]), 0) * scale b_ds = b_ds.to(b_k.dtype) # [BT, BK] - b_dq += tl.dot(b_ds, b_k) - b_dk += tl.dot(tl.trans(b_ds), b_q) + b_dq = tl.dot(b_ds, b_k, b_dq) + b_dk = tl.dot(tl.trans(b_ds), b_q, b_dk) b_dg = tl.sum(b_dq * b_q, axis=1) - tl.sum(b_dk * b_k, axis=1) - p_dg = tl.make_block_ptr(dg, (T,), (HV,), (i_t * BT,), (BT,), (0,)) - # (SY 09/21) revcumsum in a separate kernel due to strange triton compiler issue - # b_dg = tl.dot(tl.where(o_t[:, None] <= o_t[None, :], 1., 0.), b_dg, allow_tf32=False) + b_dg_last) + p_dg = dg + o_t * HV + # reverse cumsum runs in a separate kernel to avoid a Triton compiler issue. b_dg = tl.where(o_t < min(i_t * BT + BT, T) - 1, b_dg, b_dg + b_dg_last) - tl.store(p_dq, b_dq.to(p_dq.dtype.element_ty), boundary_check=(0, 1)) - tl.store(p_dk, b_dk.to(p_dk.dtype.element_ty), boundary_check=(0, 1)) - tl.store(p_dg, b_dg.to(p_dg.dtype.element_ty), boundary_check=(0,)) + tl.store(p_dq, b_dq.to(p_dq.dtype.element_ty), mask=m_qk) + tl.store(p_dk, b_dk.to(p_dk.dtype.element_ty), mask=m_qk) + tl.store(p_dg, b_dg.to(p_dg.dtype.element_ty), mask=m_t) elif USE_G_GAMMA: b_dq = b_dq * exp2(b_g)[:, None] * scale @@ -315,19 +329,19 @@ def chunk_bwd_kernel_dqkwg( b_ds = tl.where(m_A, b_ds * exp2(b_g[:, None] - b_g[None, :]), 0) * scale b_ds = b_ds.to(b_k.dtype) # [BT, BK] - b_dq += tl.dot(b_ds, b_k) - b_dk += tl.dot(tl.trans(b_ds), b_q) - tl.store(p_dq, b_dq.to(p_dq.dtype.element_ty), boundary_check=(0, 1)) - tl.store(p_dk, b_dk.to(p_dk.dtype.element_ty), boundary_check=(0, 1)) + b_dq = tl.dot(b_ds, b_k, b_dq) + b_dk = tl.dot(tl.trans(b_ds), b_q, b_dk) + tl.store(p_dq, b_dq.to(p_dq.dtype.element_ty), mask=m_qk) + tl.store(p_dk, b_dk.to(p_dk.dtype.element_ty), mask=m_qk) else: b_ds = tl.where(m_A, b_ds, 0) b_ds = b_ds.to(b_k.dtype) - b_dq += tl.dot(b_ds, b_k) + b_dq = tl.dot(b_ds, b_k, b_dq) b_dk += tl.dot(tl.trans(b_ds), b_q) * scale b_dq *= scale - tl.store(p_dq, b_dq.to(p_dq.dtype.element_ty), boundary_check=(0, 1)) - tl.store(p_dk, b_dk.to(p_dk.dtype.element_ty), boundary_check=(0, 1)) + tl.store(p_dq, b_dq.to(p_dq.dtype.element_ty), mask=m_qk) + tl.store(p_dk, b_dk.to(p_dk.dtype.element_ty), mask=m_qk) @triton.heuristics({ @@ -369,12 +383,13 @@ def chunk_bwd_kernel_dv( IS_VARLEN: tl.constexpr, STATE_V_FIRST: tl.constexpr, ): - i_v, i_t, i_bh = tl.program_id(0), tl.program_id(1), tl.program_id(2).to(tl.int64) + i_v, i_t = unflatten_program_id(tl.cdiv(V, BV)) + i_bh = tl.program_id(1).to(tl.int64) i_b, i_h = i_bh // HV, i_bh % HV 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.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) + i_n, i_t = tl.load(chunk_indices + i_t * 2).to(tl.int32), tl.load(chunk_indices + i_t * 2 + 1).to(tl.int64) + bos, eos = tl.load(cu_seqlens + i_n).to(tl.int64), tl.load(cu_seqlens + i_n + 1).to(tl.int64) T = eos - bos NT = tl.cdiv(T, BT) else: @@ -384,34 +399,36 @@ def chunk_bwd_kernel_dv( b_dv = tl.zeros([BT, BV], dtype=tl.float32) - # offset calculation q += (bos * H + i_h // (HV // H)) * K k += (bos * H + i_h // (HV // H)) * K do += (bos * HV + i_h) * V dv += (bos * HV + i_h) * V dh += (i_tg * HV + i_h).to(tl.int64) * K*V + o_t = i_t * BT + tl.arange(0, BT) + m_t = o_t < T + o_v = i_v * BV + tl.arange(0, BV) b_A = tl.zeros([BT, BT], dtype=tl.float32) for i_k in range(tl.cdiv(K, BK)): - p_k = tl.make_block_ptr(k, (T, K), (H*K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) - p_q = tl.make_block_ptr(q, (K, T), (1, H*K), (i_k * BK, i_t * BT), (BK, BT), (0, 1)) - b_q = tl.load(p_q, boundary_check=(0, 1)) - b_k = tl.load(p_k, boundary_check=(0, 1)) - b_A += tl.dot(b_k, b_q) + o_k = i_k * BK + tl.arange(0, BK) + m_k = o_k < K + p_k = k + o_t[:, None] * (H*K) + o_k[None, :] + p_q = q + o_k[:, None] + o_t[None, :] * (H*K) + b_q = tl.load(p_q, mask=m_k[:, None] & m_t[None, :], other=0.0) + b_k = tl.load(p_k, mask=m_t[:, None] & m_k[None, :], other=0.0) + b_A = tl.dot(b_k, b_q, b_A) if STATE_V_FIRST: - p_dh = tl.make_block_ptr(dh, (V, K), (K, 1), (i_v * BV, i_k * BK), (BV, BK), (1, 0)) - b_dh = tl.trans(tl.load(p_dh, boundary_check=(0, 1))) + p_dh = dh + o_v[:, None] * K + o_k[None, :] + b_dh = tl.trans(tl.load(p_dh, mask=(o_v[:, None] < V) & m_k[None, :], other=0.0)) else: - p_dh = tl.make_block_ptr(dh, (K, V), (V, 1), (i_k * BK, i_v * BV), (BK, BV), (1, 0)) - b_dh = tl.load(p_dh, boundary_check=(0, 1)) - b_dv += tl.dot(b_k, b_dh.to(b_k.dtype)) + p_dh = dh + o_k[:, None] * V + o_v[None, :] + b_dh = tl.load(p_dh, mask=m_k[:, None] & (o_v[None, :] < V), other=0.0) + b_dv = tl.dot(b_k, b_dh.to(b_k.dtype), b_dv) - o_t = i_t * BT + tl.arange(0, BT) - m_t = o_t < T if USE_G: g += bos * HV + i_h - p_g = tl.make_block_ptr(g, (T,), (HV,), (i_t * BT,), (BT,), (0,)) - b_g = tl.load(p_g, boundary_check=(0,)) + p_g = g + o_t * HV + b_g = tl.load(p_g, mask=m_t, other=0.0) b_g_last = tl.load(g + (min(i_t * BT + BT, T) - 1) * HV) if USE_G_GAMMA: b_gamma = tl.load(g_gamma + i_h) @@ -424,11 +441,11 @@ def chunk_bwd_kernel_dv( b_dv *= tl.where(m_t, exp2(-b_g + b_g_last), 0)[:, None] else: b_A = tl.where(m_A, b_A * scale, 0).to(do.dtype.element_ty) - p_do = tl.make_block_ptr(do, (T, V), (HV*V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) - p_dv = tl.make_block_ptr(dv, (T, V), (HV*V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) - b_do = tl.load(p_do, boundary_check=(0, 1)) - b_dv += tl.dot(b_A.to(b_do.dtype), b_do) - tl.store(p_dv, b_dv.to(p_dv.dtype.element_ty), boundary_check=(0, 1)) + p_do = do + o_t[:, None] * (HV*V) + o_v[None, :] + p_dv = dv + o_t[:, None] * (HV*V) + o_v[None, :] + b_do = tl.load(p_do, mask=m_t[:, None] & (o_v < V)[None, :], other=0.0) + b_dv = tl.dot(b_A.to(b_do.dtype), b_do, b_dv) + tl.store(p_dv, b_dv.to(p_dv.dtype.element_ty), mask=m_t[:, None] & (o_v < V)[None, :]) @triton.heuristics({ @@ -440,7 +457,7 @@ def chunk_bwd_kernel_dv( @fla_cache_autotune( configs=[ triton.Config({}, num_warps=num_warps, num_stages=num_stages) - for num_warps in NUM_WARPS + for num_warps in CHUNK_BWD_DV_LOCAL_NUM_WARPS for num_stages in [2, 3, 4] ], key=['H', 'HV', 'K', 'V', 'BT', 'BK', 'BV', 'USE_G'], @@ -471,56 +488,62 @@ def chunk_bwd_kernel_dv_local( USE_A: tl.constexpr, IS_VARLEN: tl.constexpr, ): - i_t, i_bh = tl.program_id(0), tl.program_id(1).to(tl.int64) + i_t, i_bh = tl.program_id(0).to(tl.int64), tl.program_id(1).to(tl.int64) i_b, i_h = i_bh // HV, i_bh % HV 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.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) + i_n, i_t = tl.load(chunk_indices + i_t * 2).to(tl.int32), tl.load(chunk_indices + i_t * 2 + 1).to(tl.int64) + bos, eos = tl.load(cu_seqlens + i_n).to(tl.int64), tl.load(cu_seqlens + i_n + 1).to(tl.int64) T = eos - bos else: bos, eos = i_b * T, i_b * T + T - # offset calculation q += (bos * H + i_h // (HV // H)) * K k += (bos * H + i_h // (HV // H)) * K do += (bos * HV + i_h) * V dv += (bos * HV + i_h) * V + o_t = i_t * BT + tl.arange(0, BT) + m_t = o_t < T if USE_A: - p_A = tl.make_block_ptr(A + (bos * HV + i_h) * BT, (BT, T), (1, HV*BT), (0, i_t * BT), (BT, BT), (0, 1)) - b_A = tl.load(p_A, boundary_check=(0, 1)) + p_A = A + (bos * HV + i_h) * BT + tl.arange(0, BT)[:, None] + o_t[None, :] * (HV*BT) + b_A = tl.load(p_A, mask=m_t[None, :], other=0.0) else: if USE_G: g += bos * HV + i_h - p_g = tl.make_block_ptr(g, (T,), (HV,), (i_t * BT,), (BT,), (0,)) - b_g = tl.load(p_g, boundary_check=(0,)) + p_g = g + o_t * HV + b_g = tl.load(p_g, mask=m_t, other=0.0) if USE_G_GAMMA: b_gamma = tl.load(g_gamma + i_h) b_g = b_gamma * (tl.arange(0, BT) + 1) + # Keep in fp32 (some pass it as fp64) which causes issues under compile + scale = scale.to(tl.float32) b_A = tl.zeros([BT, BT], dtype=tl.float32) for i_k in range(tl.cdiv(K, BK)): - p_k = tl.make_block_ptr(k, (T, K), (H*K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) - p_q = tl.make_block_ptr(q, (K, T), (1, H*K), (i_k * BK, i_t * BT), (BK, BT), (0, 1)) + o_k = i_k * BK + tl.arange(0, BK) + m_k = o_k < K + p_k = k + o_t[:, None] * (H*K) + o_k[None, :] + p_q = q + o_k[:, None] + o_t[None, :] * (H*K) - b_k = tl.load(p_k, boundary_check=(0, 1)) - b_q = tl.load(p_q, boundary_check=(0, 1)) + b_k = tl.load(p_k, mask=m_t[:, None] & m_k[None, :], other=0.0) + b_q = tl.load(p_q, mask=m_k[:, None] & m_t[None, :], other=0.0) b_A += tl.dot(b_k, b_q) * scale if USE_G or USE_G_GAMMA: b_A *= exp2(b_g[None, :] - b_g[:, None]) - o_t = i_t * BT + tl.arange(0, BT) - m_t = o_t < T m_A = (o_t[:, None] <= o_t[None, :]) & (m_t[:, None] & m_t) b_A = tl.where(m_A, b_A, 0).to(do.dtype.element_ty) for i_v in range(tl.cdiv(V, BV)): - p_do = tl.make_block_ptr(do, (T, V), (HV*V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) - p_dv = tl.make_block_ptr(dv, (T, V), (HV*V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) - b_do = tl.load(p_do, boundary_check=(0, 1)) + o_v = i_v * BV + tl.arange(0, BV) + m_v = o_v < V + p_do = do + o_t[:, None] * (HV*V) + o_v[None, :] + p_dv = dv + o_t[:, None] * (HV*V) + o_v[None, :] + b_do = tl.load(p_do, mask=m_t[:, None] & m_v[None, :], other=0.0) b_dv = tl.dot(b_A.to(b_do.dtype), b_do) - tl.store(p_dv, b_dv.to(p_dv.dtype.element_ty), boundary_check=(0, 1)) + tl.store(p_dv, b_dv.to(p_dv.dtype.element_ty), mask=m_t[:, None] & m_v[None, :]) +# @dispatch('common') def chunk_fwd_o( q: torch.Tensor, k: torch.Tensor, @@ -543,7 +566,7 @@ def chunk_fwd_o( scale = k.shape[-1] ** -0.5 o = torch.empty_like(v) - def grid(meta): return (triton.cdiv(V, meta['BV']), NT, B * HV) + def grid(meta): return (triton.cdiv(V, meta['BV']) * NT, B * HV) chunk_fwd_kernel_o[grid]( q=q, k=k, @@ -566,6 +589,7 @@ def chunk_fwd_o( return o +@dispatch def chunk_bwd_dv( q: torch.Tensor, k: torch.Tensor, @@ -580,12 +604,18 @@ def chunk_bwd_dv( chunk_indices: torch.LongTensor | None = None, ) -> torch.Tensor: B, T, H, K, V, HV = *k.shape, do.shape[-1], do.shape[2] + if q.dtype in (torch.float16, torch.bfloat16): + # Triton miscompiles masked K-tail iterations into OOB shared-memory access for 16-bit odd K/V (IMA) + assert K % 2 == 0 and V % 2 == 0, f"chunk_bwd_dv requires even K and V for {q.dtype}, got K={K}, V={V}" BT = chunk_size if chunk_indices is None and cu_seqlens is not None: chunk_indices = prepare_chunk_indices(cu_seqlens, BT) - # H100 can have larger block size + # H100 can have larger block size, and so can the A100 when every tensor this kernel loads is 16-bit if check_shared_mem('hopper', k.device.index): CONST_TILING = 128 + elif check_shared_mem('ampere', k.device.index): + is_16bit = all(x.element_size() == 2 for x in (q, k, do, dh) if x is not None) + CONST_TILING = 128 if is_16bit else 64 elif check_shared_mem('ada', k.device.index): CONST_TILING = 64 else: @@ -598,7 +628,7 @@ def chunk_bwd_dv( scale = k.shape[-1] ** -0.5 dv = torch.empty_like(do) - grid = (NV, NT, B * HV) + grid = (NV * NT, B * HV) chunk_bwd_kernel_dv[grid]( q=q, k=k, @@ -623,6 +653,7 @@ def chunk_bwd_dv( return dv +# @dispatch('common') def chunk_bwd_dv_local( q: torch.Tensor, k: torch.Tensor, @@ -630,7 +661,7 @@ def chunk_bwd_dv_local( g: torch.Tensor | None = None, g_gamma: torch.Tensor | None = None, A: torch.Tensor | None = None, - scale: float = None, + scale: float | None = None, cu_seqlens: torch.LongTensor | None = None, chunk_size: int = 64, chunk_indices: torch.LongTensor | None = None, @@ -639,9 +670,12 @@ def chunk_bwd_dv_local( BT = chunk_size if chunk_indices is None and cu_seqlens is not None: chunk_indices = prepare_chunk_indices(cu_seqlens, BT) - # H100 can have larger block size + # H100 can have larger block size, and so can the A100 when every tensor this kernel loads is 16-bit if check_shared_mem('hopper', k.device.index): CONST_TILING = 128 + elif check_shared_mem('ampere', k.device.index): + is_16bit = all(x.element_size() == 2 for x in (q, k, do, A) if x is not None) + CONST_TILING = 128 if is_16bit else 64 elif check_shared_mem('ada', k.device.index): CONST_TILING = 64 else: @@ -675,7 +709,8 @@ def chunk_bwd_dv_local( return dv -@dispatch('common') +# Use triton only, other backends (cutlass, cute, etc.) need additional torch.compile support +# @dispatch('common') def chunk_bwd_dqkwg( q: torch.Tensor, k: torch.Tensor, @@ -706,8 +741,12 @@ def chunk_bwd_dqkwg( chunk_indices = prepare_chunk_indices(cu_seqlens, BT) NT = triton.cdiv(T, BT) if cu_seqlens is None else len(chunk_indices) + # H100 can have larger block size, and so can the A100 when every tensor this kernel loads is 16-bit if check_shared_mem('hopper', k.device.index): CONST_TILING = 128 + elif check_shared_mem('ampere', k.device.index): + is_16bit = all(x.element_size() == 2 for x in (q, k, v, do, h, dh, w, dv) if x is not None) + CONST_TILING = 128 if is_16bit else 64 elif check_shared_mem('ada', k.device.index): CONST_TILING = 64 else: @@ -720,7 +759,7 @@ def chunk_bwd_dqkwg( dg = torch.empty(NK, *g.shape, dtype=torch.float32, device=g.device) if g is not None else None dw = torch.empty_like(w) if w is not None else None - grid = (NK, NT, B * HV) + grid = (NK * NT, B * HV) chunk_bwd_kernel_dqkwg[grid]( q=q, k=k, @@ -730,9 +769,9 @@ def chunk_bwd_dqkwg( h=h, do=do, dh=dh, - dw=dw, dq=dq, dk=dk, + dw=dw, dv=dv, dg=dg, cu_seqlens=cu_seqlens, diff --git a/build/torch-cuda/ops/common/chunk_scaled_dot_kkt.py b/build/torch-cuda/ops/common/chunk_scaled_dot_kkt.py index 41dee5853e23f989ba83049f3cb4e51cb51f1b76..b8adad7da2ee93c77a2d8ba2a01c376e9ef1c242 100644 --- a/build/torch-cuda/ops/common/chunk_scaled_dot_kkt.py +++ b/build/torch-cuda/ops/common/chunk_scaled_dot_kkt.py @@ -9,6 +9,7 @@ import torch import triton import triton.language as tl +from ...backends import dispatch from ...ops.utils import prepare_chunk_indices from ...ops.utils.op import exp2 from ...utils import autotune_cache_kwargs @@ -45,39 +46,41 @@ def chunk_scaled_dot_kkt_fwd_kernel( IS_VARLEN: tl.constexpr, USE_G: tl.constexpr, ): - i_t, i_bh = tl.program_id(0), tl.program_id(1).to(tl.int64) + i_t, i_bh = tl.program_id(0).to(tl.int64), tl.program_id(1).to(tl.int64) i_b, i_h = i_bh // HV, i_bh % HV 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.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) + i_n, i_t = tl.load(chunk_indices + i_t * 2).to(tl.int32), tl.load(chunk_indices + i_t * 2 + 1).to(tl.int64) + bos, eos = tl.load(cu_seqlens + i_n).to(tl.int64), tl.load(cu_seqlens + i_n + 1).to(tl.int64) T = eos - bos else: bos, eos = i_b * T, i_b * T + T o_t = i_t * BT + tl.arange(0, BT) m_t = o_t < T - p_b = tl.make_block_ptr(beta + bos*HV + i_h, (T,), (HV,), (i_t * BT,), (BT,), (0,)) - b_b = tl.load(p_b, boundary_check=(0,)) + p_b = beta + bos*HV + i_h + o_t * HV + b_b = tl.load(p_b, mask=m_t, other=0.0) b_A = tl.zeros([BT, BT], dtype=tl.float32) for i_k in range(tl.cdiv(K, BK)): - p_k = tl.make_block_ptr(k + (bos*H + i_h // (HV // H)) * K, (T, K), (H*K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) - b_k = tl.load(p_k, boundary_check=(0, 1)) - b_A += tl.dot(b_k, tl.trans(b_k)) + o_k = i_k * BK + tl.arange(0, BK) + p_k = k + (bos*H + i_h // (HV // H)) * K + o_t[:, None] * (H*K) + o_k[None, :] + b_k = tl.load(p_k, mask=m_t[:, None] & (o_k < K)[None, :], other=0.0) + b_A = tl.dot(b_k, tl.trans(b_k), b_A) if USE_G: - p_g = tl.make_block_ptr(g + bos*HV + i_h, (T,), (HV,), (i_t * BT,), (BT,), (0,)) - b_g = tl.load(p_g, boundary_check=(0,)) + p_g = g + bos*HV + i_h + o_t * HV + b_g = tl.load(p_g, mask=m_t, other=0.0) b_g_diff = b_g[:, None] - b_g[None, :] b_A *= exp2(b_g_diff) b_A *= b_b[:, None] m_A = (o_t[:, None] > o_t[None, :]) & (m_t[:, None] & m_t) b_A = tl.where(m_A, b_A, 0) - p_A = tl.make_block_ptr(A + (bos*HV + i_h) * BT, (T, BT), (BT*HV, 1), (i_t * BT, 0), (BT, BT), (1, 0)) - tl.store(p_A, b_A.to(p_A.dtype.element_ty), boundary_check=(0, 1)) + p_A = A + (bos*HV + i_h) * BT + o_t[:, None] * (BT*HV) + tl.arange(0, BT)[None, :] + tl.store(p_A, b_A.to(p_A.dtype.element_ty), mask=m_t[:, None]) +@dispatch('common') def chunk_scaled_dot_kkt_fwd( k: torch.Tensor, g: torch.Tensor | None = None, diff --git a/build/torch-cuda/ops/common/fused_chunk.py b/build/torch-cuda/ops/common/fused_chunk.py index fbb21b11cd5d3e12521c0dd7e0da0ac9f78039af..5b75c3723afca12f9e16eb371ee16b081107dcaf 100644 --- a/build/torch-cuda/ops/common/fused_chunk.py +++ b/build/torch-cuda/ops/common/fused_chunk.py @@ -11,7 +11,7 @@ import triton.language as tl from ...ops.utils import chunk_local_cumsum from ...ops.utils.constant import RCP_LN2 -from ...ops.utils.op import exp2 +from ...ops.utils.op import exp2, unflatten_program_id from ...utils import ( IS_NVIDIA_HOPPER, autocast_custom_bwd, @@ -68,12 +68,12 @@ def fused_chunk_fwd_kernel( STORE_FINAL_STATE: tl.constexpr, IS_VARLEN: tl.constexpr, ): - i_v, i_k, i_nh = tl.program_id(0), tl.program_id(1), tl.program_id(2) + i_v, i_k, i_nh = unflatten_program_id(tl.cdiv(V, BV), tl.cdiv(K, BK)) i_n, i_h = i_nh // H, i_nh % H all = B * T if IS_VARLEN: - bos, eos = tl.load(cu_seqlens + i_n).to(tl.int32), tl.load(cu_seqlens + i_n + 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 else: bos, eos = i_n * T, i_n * T + T @@ -101,25 +101,30 @@ def fused_chunk_fwd_kernel( # [BK, BV] b_h = tl.zeros([BK, BV], dtype=tl.float32) if USE_INITIAL_STATE: - p_h = tl.make_block_ptr(h0 + i_nh * K*V, (K, V), (V, 1), (i_k * BK, i_v * BV), (BK, BV), (1, 0)) - b_h = tl.load(p_h, boundary_check=(0, 1)).to(tl.float32) + o_k = i_k * BK + tl.arange(0, BK) + o_v = i_v * BV + tl.arange(0, BV) + p_h = h0 + i_nh * K*V + o_k[:, None] * V + o_v[None, :] + b_h = tl.load(p_h, mask=(o_k[:, None] < K) & (o_v[None, :] < V), other=0.0).to(tl.float32) for i_t in range(0, NT): - p_q = tl.make_block_ptr(q, (T, K), (H*K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) - p_k = tl.make_block_ptr(k, (K, T), (1, H*K), (i_k * BK, i_t * BT), (BK, BT), (0, 1)) - p_v = tl.make_block_ptr(v, (T, V), (H*V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) - p_o = tl.make_block_ptr(o, (T, V), (H*V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) - - o_t = i_t * BT + tl.arange(0, BT) + i_t_int64 = i_t.to(tl.int64) + o_t = i_t_int64 * BT + tl.arange(0, BT) m_t = o_t < T + o_k = i_k * BK + tl.arange(0, BK) + o_v = i_v * BV + tl.arange(0, BV) + p_q = q + o_t[:, None] * (H*K) + o_k[None, :] + p_k = k + o_k[:, None] + o_t[None, :] * (H*K) + p_v = v + o_t[:, None] * (H*V) + o_v[None, :] + p_o = o + o_t[:, None] * (H*V) + o_v[None, :] + # [BT, BK] - b_q = tl.load(p_q, boundary_check=(0, 1)) + b_q = tl.load(p_q, mask=m_t[:, None] & (o_k < K)[None, :], other=0.0) b_q = (b_q * scale).to(b_q.dtype) # [BK, BT] - b_k = tl.load(p_k, boundary_check=(0, 1)) + b_k = tl.load(p_k, mask=(o_k[:, None] < K) & m_t[None, :], other=0.0) # [BT, BV] - b_v = tl.load(p_v, boundary_check=(0, 1)) - last_idx = min(i_t * BT + BT, T) - 1 + b_v = tl.load(p_v, mask=m_t[:, None] & (o_v < V)[None, :], other=0.0) + last_idx = min(i_t_int64 * BT + BT, T) - 1 # [BT, BT] b_s = tl.dot(b_q, b_k) @@ -134,7 +139,7 @@ def fused_chunk_fwd_kernel( b_gk = exp2(b_g_last - b_g) b_gn = exp2(b_g_last) if USE_G_GAMMA: - b_g_last = b_gamma * min(BT, T - i_t * BT) + b_g_last = b_gamma * min(BT, T - i_t_int64 * BT) b_gk = exp2(b_g_last - b_g) b_gn = exp2(b_g_last) if USE_G or USE_G_GAMMA: @@ -151,13 +156,15 @@ def fused_chunk_fwd_kernel( # [BT, BV] b_o = tl.dot(b_s.to(b_q.dtype), b_v) + tl.dot(b_q, b_h.to(b_q.dtype)) - b_h += tl.dot(b_k, b_v) + b_h = tl.dot(b_k, b_v, b_h) - tl.store(p_o, b_o.to(p_o.dtype.element_ty), boundary_check=(0, 1)) + tl.store(p_o, b_o.to(p_o.dtype.element_ty), mask=m_t[:, None] & (o_v < V)[None, :]) if STORE_FINAL_STATE: - p_ht = tl.make_block_ptr(ht + i_nh * K*V, (K, V), (V, 1), (i_k * BK, i_v * BV), (BK, BV), (1, 0)) - tl.store(p_ht, b_h.to(p_ht.dtype.element_ty), boundary_check=(0, 1)) + o_k = i_k * BK + tl.arange(0, BK) + o_v = i_v * BV + tl.arange(0, BV) + p_ht = ht + i_nh * K*V + o_k[:, None] * V + o_v[None, :] + tl.store(p_ht, b_h.to(p_ht.dtype.element_ty), mask=(o_k[:, None] < K) & (o_v[None, :] < V)) @triton.heuristics({ @@ -207,12 +214,12 @@ def fused_chunk_bwd_kernel( USE_INITIAL_STATE: tl.constexpr, USE_FINAL_STATE: tl.constexpr, ): - i_v, i_k, i_nh = tl.program_id(0), tl.program_id(1), tl.program_id(2) + i_v, i_k, i_nh = unflatten_program_id(tl.cdiv(V, BV), tl.cdiv(K, BK)) i_n, i_h = i_nh // H, i_nh % H all = B * T if IS_VARLEN: - bos, eos = tl.load(cu_seqlens + i_n).to(tl.int32), tl.load(cu_seqlens + i_n + 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 else: bos, eos = i_n * T, i_n * T + T @@ -241,25 +248,30 @@ def fused_chunk_bwd_kernel( # [BV, BK] b_h = tl.zeros([BV, BK], dtype=tl.float32) if USE_INITIAL_STATE: - p_h = tl.make_block_ptr(h0 + i_nh * K*V, (V, K), (1, V), (i_v * BV, i_k * BK), (BV, BK), (0, 1)) - b_h = tl.load(p_h, boundary_check=(0, 1)).to(tl.float32) + o_k = i_k * BK + tl.arange(0, BK) + o_v = i_v * BV + tl.arange(0, BV) + p_h = h0 + i_nh * K*V + o_v[:, None] + o_k[None, :] * V + b_h = tl.load(p_h, mask=(o_v[:, None] < V) & (o_k[None, :] < K), other=0.0).to(tl.float32) for i_t in range(0, NT): - p_q = tl.make_block_ptr(q, (T, K), (H*K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) - p_k = tl.make_block_ptr(k, (T, K), (H*K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) - p_v = tl.make_block_ptr(v, (V, T), (1, H*V), (i_v * BV, i_t * BT), (BV, BT), (0, 1)) - p_do = tl.make_block_ptr(do, (T, V), (H*V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) - p_dq = tl.make_block_ptr(dq, (T, K), (H*K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) - - o_t = i_t * BT + tl.arange(0, BT) + i_t_int64 = i_t.to(tl.int64) + o_t = i_t_int64 * BT + tl.arange(0, BT) m_t = o_t < T + o_k = i_k * BK + tl.arange(0, BK) + o_v = i_v * BV + tl.arange(0, BV) + p_q = q + o_t[:, None] * (H*K) + o_k[None, :] + p_k = k + o_t[:, None] * (H*K) + o_k[None, :] + p_v = v + o_v[:, None] + o_t[None, :] * (H*V) + p_do = do + o_t[:, None] * (H*V) + o_v[None, :] + p_dq = dq + o_t[:, None] * (H*K) + o_k[None, :] + # [BT, BK] - b_k = tl.load(p_k, boundary_check=(0, 1)) + b_k = tl.load(p_k, mask=m_t[:, None] & (o_k < K)[None, :], other=0.0) # [BV, BT] - b_v = tl.load(p_v, boundary_check=(0, 1)) + b_v = tl.load(p_v, mask=(o_v[:, None] < V) & m_t[None, :], other=0.0) # [BT, BV] - b_do = tl.load(p_do, boundary_check=(0, 1)) - last_idx = min(i_t * BT + BT, T) - 1 + b_do = tl.load(p_do, mask=m_t[:, None] & (o_v < V)[None, :], other=0.0) + last_idx = min(i_t_int64 * BT + BT, T) - 1 # [BT, BT] b_ds = tl.dot(b_do, b_v) * scale @@ -279,7 +291,7 @@ def fused_chunk_bwd_kernel( b_gs = tl.where(m_s & m_t, exp2(b_g[:, None] - b_g[None, :]), 0) b_ds = b_ds * b_gs # [BT, BK] - b_q = tl.load(p_q, boundary_check=(0, 1)) + b_q = tl.load(p_q, mask=m_t[:, None] & (o_k < K)[None, :], other=0.0) b_dq = tl.dot(b_ds.to(b_k.dtype), b_k) + tl.dot((b_do * b_gq[:, None] * scale).to(b_k.dtype), b_h.to(b_k.dtype)) # [BT] b_dg_t = tl.sum(b_q * b_dq, 1) @@ -288,7 +300,7 @@ def fused_chunk_bwd_kernel( b_h = b_h * b_gn + tl.dot(b_v, (b_k * b_gk[:, None]).to(b_k.dtype)) elif USE_G_GAMMA: - b_g_last = b_gamma * min(BT, T - i_t * BT) + b_g_last = b_gamma * min(BT, T - i_t_int64 * BT) b_gk = exp2(b_g_last - b_g) b_gn = exp2(b_g_last) @@ -296,7 +308,7 @@ def fused_chunk_bwd_kernel( b_gs = tl.where(m_s & m_t, exp2(b_g[:, None] - b_g[None, :]), 0) b_ds = b_ds * b_gs # [BT, BK] - b_q = tl.load(p_q, boundary_check=(0, 1)) + b_q = tl.load(p_q, mask=m_t[:, None] & (o_k < K)[None, :], other=0.0) b_dq = tl.dot(b_ds.to(b_k.dtype), b_k) + tl.dot((b_do * b_gq[:, None] * scale).to(b_k.dtype), b_h.to(b_k.dtype)) # [BV, BK] b_h = b_h * b_gn + tl.dot(b_v, (b_k * b_gk[:, None]).to(b_k.dtype)) @@ -307,15 +319,17 @@ def fused_chunk_bwd_kernel( # [BT, BK] b_dq = tl.dot(b_ds.to(b_k.dtype), b_k) + tl.dot((b_do * scale).to(b_k.dtype), b_h.to(b_k.dtype)) # [BV, BK] - b_h += tl.dot(b_v, b_k) + b_h = tl.dot(b_v, b_k, b_h) - tl.store(p_dq, b_dq.to(p_dq.dtype.element_ty), boundary_check=(0, 1)) + tl.store(p_dq, b_dq.to(p_dq.dtype.element_ty), mask=m_t[:, None] & (o_k < K)[None, :]) # [BK, BV] b_dh = tl.zeros([BK, BV], dtype=tl.float32) if USE_FINAL_STATE: - p_dh = tl.make_block_ptr(dht + i_nh * K*V, (K, V), (V, 1), (i_k * BK, i_v * BV), (BK, BV), (1, 0)) - b_dh += tl.load(p_dh, boundary_check=(0, 1)).to(tl.float32) + o_k = i_k * BK + tl.arange(0, BK) + o_v = i_v * BV + tl.arange(0, BV) + p_dh = dht + i_nh * K*V + o_k[:, None] * V + o_v[None, :] + b_dh += tl.load(p_dh, mask=(o_k[:, None] < K) & (o_v[None, :] < V), other=0.0).to(tl.float32) if USE_G: b_dg = tl.zeros([BT], dtype=tl.float32) @@ -326,23 +340,26 @@ def fused_chunk_bwd_kernel( tl.debug_barrier() for i_t in range(NT - 1, -1, -1): - p_q = tl.make_block_ptr(q, (K, T), (1, H*K), (i_k * BK, i_t * BT), (BK, BT), (0, 1)) - p_k = tl.make_block_ptr(k, (T, K), (H*K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) - p_v = tl.make_block_ptr(v, (T, V), (H*V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) - p_do = tl.make_block_ptr(do, (T, V), (H*V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) - p_dk = tl.make_block_ptr(dk, (T, K), (H*K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) - p_dv = tl.make_block_ptr(dv, (T, V), (H*V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) + i_t_int64 = i_t.to(tl.int64) + o_t = i_t_int64 * BT + tl.arange(0, BT) + m_t = o_t < T + o_k = i_k * BK + tl.arange(0, BK) + o_v = i_v * BV + tl.arange(0, BV) + p_q = q + o_k[:, None] + o_t[None, :] * (H*K) + p_k = k + o_t[:, None] * (H*K) + o_k[None, :] + p_v = v + o_t[:, None] * (H*V) + o_v[None, :] + p_do = do + o_t[:, None] * (H*V) + o_v[None, :] + p_dk = dk + o_t[:, None] * (H*K) + o_k[None, :] + p_dv = dv + o_t[:, None] * (H*V) + o_v[None, :] # [BK, BT] - b_q = tl.load(p_q, boundary_check=(0, 1)) + b_q = tl.load(p_q, mask=(o_k[:, None] < K) & m_t[None, :], other=0.0) # [BT, BK] - b_k = tl.load(p_k, boundary_check=(0, 1)) + b_k = tl.load(p_k, mask=m_t[:, None] & (o_k < K)[None, :], other=0.0) # [BT, BV] - b_v = tl.load(p_v, boundary_check=(0, 1)) - b_do = tl.load(p_do, boundary_check=(0, 1)) - last_idx = min(i_t * BT + BT, T) - 1 + b_v = tl.load(p_v, mask=m_t[:, None] & (o_v < V)[None, :], other=0.0) + b_do = tl.load(p_do, mask=m_t[:, None] & (o_v < V)[None, :], other=0.0) + last_idx = min(i_t_int64 * BT + BT, T) - 1 - o_t = i_t * BT + tl.arange(0, BT) - m_t = o_t < T # [BT, BT] b_s = tl.dot(b_k, b_q) b_ds = tl.dot(b_v, tl.trans(b_do)) @@ -377,7 +394,7 @@ def fused_chunk_bwd_kernel( tl.store(p_dg, b_dg.to(p_dg.dtype.element_ty), mask=m_t) elif USE_G_GAMMA: - b_g_last = b_gamma * min(BT, T - i_t * BT) + b_g_last = b_gamma * min(BT, T - i_t_int64 * BT) b_gk = exp2(b_g_last - b_g) b_gn = exp2(b_g_last) b_gs = tl.trans(tl.where(m_s & (m_t[:, None] & m_t), exp2(b_g[:, None] - b_g[None, :]), 0)) * scale @@ -400,14 +417,16 @@ def fused_chunk_bwd_kernel( # [BT, BV] b_dv = tl.dot(b_s.to(b_do.dtype), b_do) + tl.dot(b_k, b_dh.to(b_k.dtype)) # [BK, BV] - b_dh += tl.dot(b_q, (b_do * scale).to(b_do.dtype)) + b_dh = tl.dot(b_q, (b_do * scale).to(b_do.dtype), b_dh) - tl.store(p_dk, b_dk.to(p_dk.dtype.element_ty), boundary_check=(0, 1)) - tl.store(p_dv, b_dv.to(p_dv.dtype.element_ty), boundary_check=(0, 1)) + tl.store(p_dk, b_dk.to(p_dk.dtype.element_ty), mask=m_t[:, None] & (o_k < K)[None, :]) + tl.store(p_dv, b_dv.to(p_dv.dtype.element_ty), mask=m_t[:, None] & (o_v < V)[None, :]) if USE_INITIAL_STATE: - p_dh0 = tl.make_block_ptr(dh0 + i_nh * K*V, (K, V), (V, 1), (i_k * BK, i_v * BV), (BK, BV), (1, 0)) - tl.store(p_dh0, b_dh.to(p_dh0.dtype.element_ty), boundary_check=(0, 1)) + o_k = i_k * BK + tl.arange(0, BK) + o_v = i_v * BV + tl.arange(0, BV) + p_dh0 = dh0 + i_nh * K*V + o_k[:, None] * V + o_v[None, :] + tl.store(p_dh0, b_dh.to(p_dh0.dtype.element_ty), mask=(o_k[:, None] < K) & (o_v[None, :] < V)) def fused_chunk_fwd( @@ -430,7 +449,7 @@ def fused_chunk_fwd( o = v.new_empty(NK, *v.shape, dtype=torch.float) if NK > 1 else torch.empty_like(v) ht = k.new_empty(N, H, K, V, dtype=torch.float) if output_final_state else None - def grid(meta): return (triton.cdiv(V, meta['BV']), NK, N * H) + def grid(meta): return (triton.cdiv(V, meta['BV']) * NK * N * H,) fused_chunk_fwd_kernel[grid]( q=q, k=k, @@ -481,7 +500,7 @@ def fused_chunk_bwd( dg = g.new_empty(NK*NV, *g.shape, dtype=torch.float) if g is not None else None dh0 = torch.empty_like(initial_state) if initial_state is not None else None - grid = (NV, NK, N * H) + grid = (NV * NK * N * H,) fused_chunk_bwd_kernel[grid]( q=q, k=k, diff --git a/build/torch-cuda/ops/common/fused_recurrent.py b/build/torch-cuda/ops/common/fused_recurrent.py index 396295f62bf701324832b6cd9e374f96bea43a94..fb7de05051fa942d969d5f89f7631fc50108be46 100644 --- a/build/torch-cuda/ops/common/fused_recurrent.py +++ b/build/torch-cuda/ops/common/fused_recurrent.py @@ -5,6 +5,8 @@ # For a list of all contributors, visit: # https://github.com/fla-org/flash-linear-attention/graphs/contributors +from functools import lru_cache + import torch import triton import triton.language as tl @@ -57,7 +59,9 @@ def fused_recurrent_fwd_kernel( IS_VARLEN: tl.constexpr, STATE_V_FIRST: tl.constexpr = False, ): - i_v, i_k, i_nh = tl.program_id(0).to(tl.int64), tl.program_id(1).to(tl.int64), tl.program_id(2).to(tl.int64) + pid = tl.program_id(0).to(tl.int64) + NV, NK = tl.cdiv(V, BV), tl.cdiv(K, BK) + i_v, i_k, i_nh = (pid % NV).to(tl.int64), ((pid // NV) % NK).to(tl.int64), (pid // (NV * NK)).to(tl.int64) i_n, i_h = i_nh // H, i_nh % H all = B * T @@ -199,7 +203,9 @@ def fused_recurrent_bwd_kernel( IS_VARLEN: tl.constexpr, STATE_V_FIRST: tl.constexpr = False, ): - i_v, i_k, i_nh = tl.program_id(0).to(tl.int64), tl.program_id(1).to(tl.int64), tl.program_id(2).to(tl.int64) + pid = tl.program_id(0).to(tl.int64) + NV, NK = tl.cdiv(V, BV), tl.cdiv(K, BK) + i_v, i_k, i_nh = (pid % NV).to(tl.int64), ((pid // NV) % NK).to(tl.int64), (pid // (NV * NK)).to(tl.int64) i_n, i_h = i_nh // H, i_nh % H all = B * T @@ -402,6 +408,12 @@ def fused_recurrent_bwd_kernel( tl.store(p_dh0, b_dh.to(p_dh0.dtype.element_ty), mask=m_h) +@lru_cache(maxsize=128) +def _fused_recurrent_tiling(K: int, V: int) -> tuple[int, int, int, int]: + BK, BV = min(triton.next_power_of_2(K), 64), min(triton.next_power_of_2(V), 64) + return BK, BV, triton.cdiv(K, BK), triton.cdiv(V, BV) + + def fused_recurrent_fwd( q: torch.Tensor, k: torch.Tensor, @@ -419,8 +431,7 @@ def fused_recurrent_fwd( ): B, T, H, K, V = *k.shape, v.shape[-1] N = B if cu_seqlens is None else len(cu_seqlens) - 1 - BK, BV = min(triton.next_power_of_2(K), 64), min(triton.next_power_of_2(V), 64) - NK, NV = triton.cdiv(K, BK), triton.cdiv(V, BV) + BK, BV, NK, NV = _fused_recurrent_tiling(K, V) h0 = initial_state if output_final_state: @@ -430,9 +441,9 @@ def fused_recurrent_fwd( ht = q.new_empty(N, H, K, V, dtype=torch.float32) else: ht = None - o = q.new_empty(NK, *v.shape, dtype=torch.float32) + o = q.new_empty((NK, *v.shape) if NK > 1 else v.shape, dtype=torch.float32) - grid = (NV, NK, N * H) + grid = (NV * NK * N * H,) fused_recurrent_fwd_kernel[grid]( q=q, k=k, @@ -460,7 +471,8 @@ def fused_recurrent_fwd( REVERSE=reverse, STATE_V_FIRST=state_v_first, ) - o = o.sum(0) + if NK > 1: + o = o.sum(0) return o, ht @@ -484,24 +496,24 @@ def fused_recurrent_bwd( B, T, H, K, V = *k.shape, v.shape[-1] N = B if cu_seqlens is None else len(cu_seqlens) - 1 - BK, BV = min(triton.next_power_of_2(K), 64), min(triton.next_power_of_2(V), 64) - NK, NV = triton.cdiv(K, BK), triton.cdiv(V, BV) + BK, BV, NK, NV = _fused_recurrent_tiling(K, V) h0 = initial_state - dq = q.new_empty(NV, *q.shape, dtype=torch.float32) - dk = q.new_empty(NV, *k.shape, dtype=torch.float32) - dv = q.new_empty(NK, *v.shape, dtype=torch.float32) + dq = q.new_empty((NV, *q.shape) if NV > 1 else q.shape, dtype=torch.float32) + dk = q.new_empty((NV, *k.shape) if NV > 1 else k.shape, dtype=torch.float32) + dv = q.new_empty((NK, *v.shape) if NK > 1 else v.shape, dtype=torch.float32) dh0 = torch.empty_like(h0) if h0 is not None else None dg, dgk, dgv = None, None, None if g is not None: - dg = g.new_empty(NK*NV, *g.shape, dtype=torch.float32) + dg = g.new_empty((NK*NV, *g.shape) if NK*NV > 1 else g.shape, dtype=torch.float32) if gk is not None: - dgk = gk.new_empty(NV, *gk.shape, dtype=torch.float32) + dgk = gk.new_empty((NV, *gk.shape) if NV > 1 else gk.shape, dtype=torch.float32) if gv is not None: - dgv = gv.new_empty(NK, *gv.shape, dtype=torch.float32) + dgv = gv.new_empty((NK, *gv.shape) if NK > 1 else gv.shape, dtype=torch.float32) - grid = (NV, NK, N * H) + grid = (NV * NK * N * H,) + # disable auto-multi-buffer on the gate-gradient accumulate fused_recurrent_bwd_kernel[grid]( q=q, k=k, @@ -537,15 +549,23 @@ def fused_recurrent_bwd( REVERSE=reverse, STATE_V_FIRST=state_v_first, ) - dq = dq.sum(0) - dk = dk.sum(0) - dv = dv.sum(0) + if NV > 1: + dq = dq.sum(0) + dk = dk.sum(0) + if NK > 1: + dv = dv.sum(0) if g is not None: - dg = dg.sum(0).to(g) + if NK*NV > 1: + dg = dg.sum(0) + dg = dg.to(g) if gk is not None: - dgk = dgk.sum(0).to(gk) + if NV > 1: + dgk = dgk.sum(0) + dgk = dgk.to(gk) if gv is not None: - dgv = dgv.sum(0).to(gv) + if NK > 1: + dgv = dgv.sum(0) + dgv = dgv.to(gv) return dq, dk, dv, dg, dgk, dgv, dh0 diff --git a/build/torch-cuda/ops/common/gate.py b/build/torch-cuda/ops/common/gate.py index 3f60b99a59145c2c4189aa19b5607e1d0c5e67ac..5abcdbd10bb093e706d1341200a20980ab356f54 100644 --- a/build/torch-cuda/ops/common/gate.py +++ b/build/torch-cuda/ops/common/gate.py @@ -11,6 +11,7 @@ import torch import triton import triton.language as tl +from ...backends import dispatch from ...utils import autocast_custom_bwd, autocast_custom_fwd, input_guard @@ -53,6 +54,7 @@ _BETA_SIGMOID_BLOCK_SIZE = 2048 _BETA_SIGMOID_NUM_WARPS = 8 +@dispatch('common') def fused_beta_sigmoid_fwd(x: torch.Tensor, scale: float = 1.0) -> torch.Tensor: y = torch.empty_like(x, dtype=torch.float32) n_elements = x.numel() @@ -68,6 +70,7 @@ def fused_beta_sigmoid_fwd(x: torch.Tensor, scale: float = 1.0) -> torch.Tensor: return y +@dispatch('common') def fused_beta_sigmoid_bwd(x: torch.Tensor, dy: torch.Tensor, scale: float = 1.0) -> torch.Tensor: dx = torch.empty_like(x) n_elements = x.numel() diff --git a/build/torch-cuda/ops/common/intracard_cp.py b/build/torch-cuda/ops/common/intracard_cp.py index fb82b50041d010cba28828b756789e91fb3307c8..317b59bae591e6b7d0693a4edd1468798f364435 100644 --- a/build/torch-cuda/ops/common/intracard_cp.py +++ b/build/torch-cuda/ops/common/intracard_cp.py @@ -26,7 +26,7 @@ import triton from ...ops.common.chunk_delta_h import chunk_gated_delta_rule_fwd_kernel_h_blockdim64 from ...ops.cp.chunk_delta_h import pre_process_fwd_kernel_merged from ...ops.utils.index import prepare_chunk_indices, prepare_chunk_offsets -from ...utils import get_multiprocessor_count +from ...utils import IS_TF32_SUPPORTED, get_multiprocessor_count logger = logging.getLogger(__name__) @@ -111,7 +111,7 @@ def _raw_chunk_gated_delta_rule_fwd_h( v_new = torch.empty_like(u) if save_new_value else None def grid(meta): - return (triton.cdiv(V, meta['BV']), N * HV) + return (triton.cdiv(V, meta['BV']) * N * HV, ) chunk_gated_delta_rule_fwd_kernel_h_blockdim64[grid]( k=k, @@ -148,7 +148,7 @@ def compute_subseq_len( Splitting always reduces the critical path and helps, as long as the sequence is long enough to amortize the pre_scan + merge overhead. - The fwd_h kernel grid is (num_v_blocks, N*HV) where num_v_blocks ≈ 2. + The fwd_h kernel grid is num_v_blocks*N*HV where num_v_blocks ≈ 2. Each sub-sequence contributes 2*HV blocks. We target enough splits so that even a single long sequence can saturate all SMs. @@ -259,6 +259,7 @@ def intracard_pre_scan( cu_seqlens_subseq_split: torch.Tensor, S_split: int, chunk_size: int = 64, + use_tf32x3_affine_chain: bool = False, ): H, K, V, HV = kg.shape[2], kg.shape[3], u.shape[3], u.shape[2] BK = triton.next_power_of_2(K) @@ -286,6 +287,10 @@ def intracard_pre_scan( BLOCK_SIZE=BLOCK_SIZE, BK1=BK, MULTI_SEQS=True, + AFFINE_CHAIN_PRECISION=( + "tf32x3" if use_tf32x3_affine_chain and IS_TF32_SUPPORTED + else ("ieee" if not IS_TF32_SUPPORTED else None) + ), ) return hm @@ -300,6 +305,7 @@ def intracard_merge( device: torch.device, initial_state: torch.Tensor | None = None, state_v_first: bool = False, + use_tf32x3_affine_chain: bool = False, ) -> tuple[torch.Tensor | None, int]: """Merge sub-sequence states using pre-computed parameters. @@ -346,6 +352,7 @@ def intracard_merge( init_offsets=init_offsets, h0_seq_ids=h0_seq_ids, h0=initial_state, + h_seq_idx=None, HV=HV, K=K, V=V, @@ -354,6 +361,10 @@ def intracard_merge( INTRACARD_MODE=True, NUM_SEQ_ENTRIES=num_split_seqs, STATE_V_FIRST=state_v_first, + AFFINE_CHAIN_PRECISION=( + "tf32x3" if use_tf32x3_affine_chain and IS_TF32_SUPPORTED + else ("ieee" if not IS_TF32_SUPPORTED else None) + ), ) return initial_states_merge, num_non_first @@ -447,6 +458,7 @@ def intracard_fwd_h( cu_seqlens_cpu: torch.LongTensor | None = None, chunk_indices: torch.LongTensor | None = None, max_splits: int = 32, + use_tf32x3_affine_chain: bool = False, ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor | None]: assert cu_seqlens is not None, "intracard_fwd_h requires cu_seqlens" @@ -575,6 +587,7 @@ def intracard_fwd_h( cu_seqlens_subseq_split=cu_seqlens_split_flat, S_split=S_split_total, chunk_size=chunk_size, + use_tf32x3_affine_chain=use_tf32x3_affine_chain, ) initial_states_merge, num_non_first = intracard_merge( @@ -586,6 +599,7 @@ def intracard_fwd_h( device=device, initial_state=initial_state, state_v_first=state_v_first, + use_tf32x3_affine_chain=use_tf32x3_affine_chain, ) if state_v_first: diff --git a/build/torch-cuda/ops/cp/chunk_delta_h.py b/build/torch-cuda/ops/cp/chunk_delta_h.py index dffa38c93f0862505e6ab854a72d390ed54889c1..fa9578ae20a02c4d76e18d91944bb49977c35161 100644 --- a/build/torch-cuda/ops/cp/chunk_delta_h.py +++ b/build/torch-cuda/ops/cp/chunk_delta_h.py @@ -7,6 +7,7 @@ from __future__ import annotations +import warnings from typing import TYPE_CHECKING import torch @@ -16,7 +17,7 @@ import triton.language as tl from ...ops.cp.comm import all_gather_into_tensor from ...ops.utils.op import exp2 -from ...utils import autotune_cache_kwargs, check_shared_mem +from ...utils import IS_TF32_SUPPORTED, autotune_cache_kwargs, check_shared_mem if TYPE_CHECKING: from ...ops.cp.context import FLACPContext @@ -61,10 +62,11 @@ def pre_process_fwd_kernel_merged( USE_BG: tl.constexpr, IS_VARLEN: tl.constexpr, MULTI_SEQS: tl.constexpr, + AFFINE_CHAIN_PRECISION: tl.constexpr = None, ): - i_col, i_h = tl.program_id(0), tl.program_id(1) + i_col, i_h = tl.program_id(0).to(tl.int64), tl.program_id(1).to(tl.int64) if MULTI_SEQS: - i_n = tl.program_id(2) + i_n = tl.program_id(2).to(tl.int64) # Offset hm for this subseq: hm[i_n, h, k, v+k] hm += i_n * HV * K * (K + V) + i_h * K * (K + V) else: @@ -116,43 +118,55 @@ def pre_process_fwd_kernel_merged( if K > 192: b_h4 = tl.zeros([64, BLOCK_SIZE], dtype=tl.float32) + o_vb = i_v * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE) + m_vb = o_vb < V + o_k1 = tl.arange(0, 64) + m_k1 = o_k1 < K + o_k2 = 64 + o_k1 + m_k2 = o_k2 < K + o_k3 = 128 + o_k1 + m_k3 = o_k3 < K + o_k4 = 192 + o_k1 + m_k4 = o_k4 < K + # Main recurrence for h for i_t in range(NT): + o_t = i_t * BT + tl.arange(0, BT) + m_t = o_t < T # Compute decayed v - p_w = tl.make_block_ptr(w, (T, K), (stride_w, 1), (i_t * BT, 0), (BT, 64), (1, 0)) - b_w = tl.load(p_w, boundary_check=(0, 1)) + p_w = w + o_t[:, None] * stride_w + o_k1[None, :] + b_w = tl.load(p_w, mask=m_t[:, None] & m_k1[None, :], other=0.0) b_v_decay = tl.dot(b_w, b_h1.to(b_w.dtype)) if K > 64: - p_w = tl.make_block_ptr(w, (T, K), (stride_w, 1), (i_t * BT, 64), (BT, 64), (1, 0)) - b_w = tl.load(p_w, boundary_check=(0, 1)) - b_v_decay += tl.dot(b_w, b_h2.to(b_w.dtype)) + p_w = w + o_t[:, None] * stride_w + o_k2[None, :] + b_w = tl.load(p_w, mask=m_t[:, None] & m_k2[None, :], other=0.0) + b_v_decay = tl.dot(b_w, b_h2.to(b_w.dtype), b_v_decay) if K > 128: - p_w = tl.make_block_ptr(w, (T, K), (stride_w, 1), (i_t * BT, 128), (BT, 64), (1, 0)) - b_w = tl.load(p_w, boundary_check=(0, 1)) - b_v_decay += tl.dot(b_w, b_h3.to(b_w.dtype)) + p_w = w + o_t[:, None] * stride_w + o_k3[None, :] + b_w = tl.load(p_w, mask=m_t[:, None] & m_k3[None, :], other=0.0) + b_v_decay = tl.dot(b_w, b_h3.to(b_w.dtype), b_v_decay) if K > 192: - p_w = tl.make_block_ptr(w, (T, K), (stride_w, 1), (i_t * BT, 192), (BT, 64), (1, 0)) - b_w = tl.load(p_w, boundary_check=(0, 1)) - b_v_decay += tl.dot(b_w, b_h4.to(b_w.dtype)) + p_w = w + o_t[:, None] * stride_w + o_k4[None, :] + b_w = tl.load(p_w, mask=m_t[:, None] & m_k4[None, :], other=0.0) + b_v_decay = tl.dot(b_w, b_h4.to(b_w.dtype), b_v_decay) - p_v = tl.make_block_ptr(v, (T, V), (stride_v, 1), (i_t * BT, i_v * BLOCK_SIZE), (BT, BLOCK_SIZE), (1, 0)) + p_v = v + o_t[:, None] * stride_v + o_vb[None, :] if USE_BG: # DPLR mode: v2 = w @ h + u, h += kg^T @ v + bg^T @ v2 - b_v_orig = tl.load(p_v, boundary_check=(0, 1)) - p_u = tl.make_block_ptr(u, (T, V), (stride_v, 1), (i_t * BT, i_v * BLOCK_SIZE), (BT, BLOCK_SIZE), (1, 0)) - b_v = b_v_decay + tl.load(p_u, boundary_check=(0, 1)) + b_v_orig = tl.load(p_v, mask=m_t[:, None] & m_vb[None, :], other=0.0) + p_u = u + o_t[:, None] * stride_v + o_vb[None, :] + b_v = b_v_decay + tl.load(p_u, mask=m_t[:, None] & m_vb[None, :], other=0.0) else: # GDN/KDA mode: v_new = v - w @ h - b_v = tl.load(p_v, boundary_check=(0, 1)) - b_v_decay + b_v = tl.load(p_v, mask=m_t[:, None] & m_vb[None, :], other=0.0) - b_v_decay last_idx = min((i_t + 1) * BT, T) - 1 # Apply g decay if USE_G: - m_t = (i_t * BT + tl.arange(0, BT)) < T b_g_last = tl.load(g + last_idx * HV).to(tl.float32) - p_g = tl.make_block_ptr(g, (T,), (HV,), (i_t * BT,), (BT,), (0,)) - b_g = tl.load(p_g, boundary_check=(0,)).to(tl.float32) + p_g = g + o_t * HV + b_g = tl.load(p_g, mask=m_t, other=0.0).to(tl.float32) b_v = b_v * tl.where(m_t, exp2(b_g_last - b_g), 0)[:, None] b_g_last = exp2(b_g_last) b_h1 *= b_g_last @@ -184,60 +198,60 @@ def pre_process_fwd_kernel_merged( b_v = b_v.to(k.dtype.element_ty) # Update h - p_k = tl.make_block_ptr(k, (K, T), (1, stride_k), (0, i_t * BT), (64, BT), (0, 1)) - b_k = tl.load(p_k, boundary_check=(0, 1)) + p_k = k + o_k1[:, None] + o_t[None, :] * stride_k + b_k = tl.load(p_k, mask=m_k1[:, None] & m_t[None, :], other=0.0) if USE_BG: # DPLR mode: h += kg^T @ v + bg^T @ v2 - p_bg = tl.make_block_ptr(bg, (K, T), (1, stride_k), (0, i_t * BT), (64, BT), (0, 1)) - b_bg = tl.load(p_bg, boundary_check=(0, 1)) + p_bg = bg + o_k1[:, None] + o_t[None, :] * stride_k + b_bg = tl.load(p_bg, mask=m_k1[:, None] & m_t[None, :], other=0.0) b_h1 += tl.dot(b_k, b_v_orig.to(b_k.dtype)) + tl.dot(b_bg, b_v) if K > 64: - p_k = tl.make_block_ptr(k, (K, T), (1, stride_k), (64, i_t * BT), (64, BT), (0, 1)) - b_k = tl.load(p_k, boundary_check=(0, 1)) - p_bg = tl.make_block_ptr(bg, (K, T), (1, stride_k), (64, i_t * BT), (64, BT), (0, 1)) - b_bg = tl.load(p_bg, boundary_check=(0, 1)) + p_k = k + o_k2[:, None] + o_t[None, :] * stride_k + b_k = tl.load(p_k, mask=m_k2[:, None] & m_t[None, :], other=0.0) + p_bg = bg + o_k2[:, None] + o_t[None, :] * stride_k + b_bg = tl.load(p_bg, mask=m_k2[:, None] & m_t[None, :], other=0.0) b_h2 += tl.dot(b_k, b_v_orig.to(b_k.dtype)) + tl.dot(b_bg, b_v) if K > 128: - p_k = tl.make_block_ptr(k, (K, T), (1, stride_k), (128, i_t * BT), (64, BT), (0, 1)) - b_k = tl.load(p_k, boundary_check=(0, 1)) - p_bg = tl.make_block_ptr(bg, (K, T), (1, stride_k), (128, i_t * BT), (64, BT), (0, 1)) - b_bg = tl.load(p_bg, boundary_check=(0, 1)) + p_k = k + o_k3[:, None] + o_t[None, :] * stride_k + b_k = tl.load(p_k, mask=m_k3[:, None] & m_t[None, :], other=0.0) + p_bg = bg + o_k3[:, None] + o_t[None, :] * stride_k + b_bg = tl.load(p_bg, mask=m_k3[:, None] & m_t[None, :], other=0.0) b_h3 += tl.dot(b_k, b_v_orig.to(b_k.dtype)) + tl.dot(b_bg, b_v) if K > 192: - p_k = tl.make_block_ptr(k, (K, T), (1, stride_k), (192, i_t * BT), (64, BT), (0, 1)) - b_k = tl.load(p_k, boundary_check=(0, 1)) - p_bg = tl.make_block_ptr(bg, (K, T), (1, stride_k), (192, i_t * BT), (64, BT), (0, 1)) - b_bg = tl.load(p_bg, boundary_check=(0, 1)) + p_k = k + o_k4[:, None] + o_t[None, :] * stride_k + b_k = tl.load(p_k, mask=m_k4[:, None] & m_t[None, :], other=0.0) + p_bg = bg + o_k4[:, None] + o_t[None, :] * stride_k + b_bg = tl.load(p_bg, mask=m_k4[:, None] & m_t[None, :], other=0.0) b_h4 += tl.dot(b_k, b_v_orig.to(b_k.dtype)) + tl.dot(b_bg, b_v) else: # GDN/KDA mode: h += k^T @ v_new - b_h1 += tl.dot(b_k, b_v) + b_h1 = tl.dot(b_k, b_v, b_h1) if K > 64: - p_k = tl.make_block_ptr(k, (K, T), (1, stride_k), (64, i_t * BT), (64, BT), (0, 1)) - b_k = tl.load(p_k, boundary_check=(0, 1)) - b_h2 += tl.dot(b_k, b_v) + p_k = k + o_k2[:, None] + o_t[None, :] * stride_k + b_k = tl.load(p_k, mask=m_k2[:, None] & m_t[None, :], other=0.0) + b_h2 = tl.dot(b_k, b_v, b_h2) if K > 128: - p_k = tl.make_block_ptr(k, (K, T), (1, stride_k), (128, i_t * BT), (64, BT), (0, 1)) - b_k = tl.load(p_k, boundary_check=(0, 1)) - b_h3 += tl.dot(b_k, b_v) + p_k = k + o_k3[:, None] + o_t[None, :] * stride_k + b_k = tl.load(p_k, mask=m_k3[:, None] & m_t[None, :], other=0.0) + b_h3 = tl.dot(b_k, b_v, b_h3) if K > 192: - p_k = tl.make_block_ptr(k, (K, T), (1, stride_k), (192, i_t * BT), (64, BT), (0, 1)) - b_k = tl.load(p_k, boundary_check=(0, 1)) - b_h4 += tl.dot(b_k, b_v) + p_k = k + o_k4[:, None] + o_t[None, :] * stride_k + b_k = tl.load(p_k, mask=m_k4[:, None] & m_t[None, :], other=0.0) + b_h4 = tl.dot(b_k, b_v, b_h4) # Store h results stride_hm_kv = K + V - p_h1 = tl.make_block_ptr(hm, (K, V), (stride_hm_kv, 1), (0, i_v * BLOCK_SIZE), (64, BLOCK_SIZE), (1, 0)) - tl.store(p_h1, b_h1.to(p_h1.dtype.element_ty), boundary_check=(0, 1)) + p_h1 = hm + o_k1[:, None] * stride_hm_kv + o_vb[None, :] + tl.store(p_h1, b_h1.to(p_h1.dtype.element_ty), mask=m_k1[:, None] & m_vb[None, :]) if K > 64: - p_h2 = tl.make_block_ptr(hm, (K, V), (stride_hm_kv, 1), (64, i_v * BLOCK_SIZE), (64, BLOCK_SIZE), (1, 0)) - tl.store(p_h2, b_h2.to(p_h2.dtype.element_ty), boundary_check=(0, 1)) + p_h2 = hm + o_k2[:, None] * stride_hm_kv + o_vb[None, :] + tl.store(p_h2, b_h2.to(p_h2.dtype.element_ty), mask=m_k2[:, None] & m_vb[None, :]) if K > 128: - p_h3 = tl.make_block_ptr(hm, (K, V), (stride_hm_kv, 1), (128, i_v * BLOCK_SIZE), (64, BLOCK_SIZE), (1, 0)) - tl.store(p_h3, b_h3.to(p_h3.dtype.element_ty), boundary_check=(0, 1)) + p_h3 = hm + o_k3[:, None] * stride_hm_kv + o_vb[None, :] + tl.store(p_h3, b_h3.to(p_h3.dtype.element_ty), mask=m_k3[:, None] & m_vb[None, :]) if K > 192: - p_h4 = tl.make_block_ptr(hm, (K, V), (stride_hm_kv, 1), (192, i_v * BLOCK_SIZE), (64, BLOCK_SIZE), (1, 0)) - tl.store(p_h4, b_h4.to(p_h4.dtype.element_ty), boundary_check=(0, 1)) + p_h4 = hm + o_k4[:, None] * stride_hm_kv + o_vb[None, :] + tl.store(p_h4, b_h4.to(p_h4.dtype.element_ty), mask=m_k4[:, None] & m_vb[None, :]) else: # ====== Stage 2: Compute m (K x K) ====== # i_col is for m part, map to K dimension @@ -256,24 +270,25 @@ def pre_process_fwd_kernel_merged( b_m = tl.where(row[:, None] == col[None, :], 1.0, 0.0) for i_t in range(NT): + o_t = i_t * BT + tl.arange(0, BT) + m_t = o_t < T # Load k and w with full BK1 rows if USE_BG: # DPLR mode: use bg for transition matrix # bg was already offset at the beginning of the kernel - p_k = tl.make_block_ptr(bg, (T, K), (stride_k, 1), (i_t * BT, 0), (BT, BK1), (1, 0)) + p_k = bg + o_t[:, None] * stride_k + row[None, :] else: - p_k = tl.make_block_ptr(k, (T, K), (stride_k, 1), (i_t * BT, 0), (BT, BK1), (1, 0)) - b_k = tl.load(p_k, boundary_check=(0, 1)) - p_w = tl.make_block_ptr(w, (T, K), (stride_w, 1), (i_t * BT, 0), (BT, BK1), (1, 0)) - b_w = tl.load(p_w, boundary_check=(0, 1)) + p_k = k + o_t[:, None] * stride_k + row[None, :] + b_k = tl.load(p_k, mask=m_t[:, None] & (row < K)[None, :], other=0.0) + p_w = w + o_t[:, None] * stride_w + row[None, :] + b_w = tl.load(p_w, mask=m_t[:, None] & (row < K)[None, :], other=0.0) last_idx = min((i_t + 1) * BT, T) - 1 if USE_G: - m_t = (i_t * BT + tl.arange(0, BT)) < T b_g_last = tl.load(g + last_idx * HV).to(tl.float32) - p_g = tl.make_block_ptr(g, (T,), (HV,), (i_t * BT,), (BT,), (0,)) - b_g = tl.load(p_g, boundary_check=(0,)).to(tl.float32) + p_g = g + o_t * HV + b_g = tl.load(p_g, mask=m_t, other=0.0).to(tl.float32) b_k = b_k * tl.where(m_t, exp2(b_g_last - b_g), 0)[:, None] b_g_last = exp2(b_g_last) b_diag = tl.where(row[:, None] == row[None, :], b_g_last, 0.0) @@ -294,16 +309,17 @@ def pre_process_fwd_kernel_merged( # GDN/KDA mode: M = (diag - k^T @ w) @ M b_kw = tl.dot(tl.trans(b_k.to(b_w.dtype)), b_w) b_m_i = b_diag - b_kw - b_m = tl.dot(b_m_i.to(tl.float32), b_m.to(tl.float32)) + b_m = tl.dot(b_m_i.to(tl.float32), b_m.to(tl.float32), input_precision=AFFINE_CHAIN_PRECISION) # Store m result stride_hm_kv = K + V - p_m = tl.make_block_ptr(hm + V, (K, K), (stride_hm_kv, 1), (0, i_k_col * BLOCK_SIZE), (BK1, BLOCK_SIZE), (1, 0)) - tl.store(p_m, b_m.to(p_m.dtype.element_ty), boundary_check=(0, 1)) + p_m = hm + V + row[:, None] * stride_hm_kv + col[None, :] + tl.store(p_m, b_m.to(p_m.dtype.element_ty), mask=(row < K)[:, None] & (col < K)[None, :]) @triton.heuristics({ 'HAS_H0': lambda args: args['h0'] is not None, + 'HAS_H_SEQ_IDX': lambda args: args['h_seq_idx'] is not None, }) @triton.autotune( configs=[ @@ -312,7 +328,7 @@ def pre_process_fwd_kernel_merged( for num_stages in [2, 3, 4] for BV in [32, 64] ], - key=['HV', 'K', 'V', 'BT'], + key=['HV', 'K', 'V'], **autotune_cache_kwargs, ) @triton.jit(do_not_specialize=['pre_or_post_num_ranks', 'rank', 'NUM_SEQ_ENTRIES']) @@ -325,6 +341,7 @@ def merge_fwd_bwd_kernel( init_offsets, # None for CP, [num_split_seqs+1] for intracard h0_seq_ids, # None for CP, [num_split_seqs] for intracard h0, # None or [N_orig, HV, K, V] for intracard (or [V, K] when transposed) + h_seq_idx, # None, or a device scalar selecting the target sequence row of h (graph replay) HV: tl.constexpr, K: tl.constexpr, V: tl.constexpr, @@ -334,7 +351,10 @@ def merge_fwd_bwd_kernel( INTRACARD_MODE: tl.constexpr, # True: intracard mode, False: CP mode NUM_SEQ_ENTRIES, # num_split_seqs for intracard HAS_H0: tl.constexpr, # Heuristic: whether h0 is provided + HAS_H_SEQ_IDX: tl.constexpr, # Heuristic: whether h_seq_idx is provided STATE_V_FIRST: tl.constexpr = False, # When True, h0/h use [V, K] layout; ag_hm always [K, V+K] + AFFINE_CHAIN_PRECISION: tl.constexpr = None, # input_precision for M@h in the state update (h' = M @ h + he) + NUM_RANKS_ON_DEVICE: tl.constexpr = False, # CP mode: load pre_or_post_num_ranks from a device tensor ): """ Unified merge kernel for both CP and Intra-card modes. @@ -351,10 +371,14 @@ def merge_fwd_bwd_kernel( ag_hm always uses [K, V+K] layout (from pre_scan). The recurrence h' = M @ h + he becomes h_T' = h_T @ M^T + he^T. """ - i_v = tl.program_id(0) + i_v = tl.program_id(0).to(tl.int64) + o_k = tl.arange(0, BK) + m_k = o_k < K + o_v = i_v * BV + tl.arange(0, BV) + m_v = o_v < V if INTRACARD_MODE: - i_seq = tl.program_id(1) - i_h = tl.program_id(2) + i_seq = tl.program_id(1).to(tl.int64) + i_h = tl.program_id(2).to(tl.int64) if i_seq >= NUM_SEQ_ENTRIES: return @@ -372,17 +396,11 @@ def merge_fwd_bwd_kernel( if HAS_H0: orig_seq_id = tl.load(h0_seq_ids + i_seq).to(tl.int32) if STATE_V_FIRST: - p_h0 = tl.make_block_ptr( - h0 + (orig_seq_id * HV + i_h) * V * K, - (V, K), (K, 1), (i_v * BV, 0), (BV, BK), (1, 0) - ) - b_h = tl.load(p_h0, boundary_check=(0, 1)).to(tl.float32) + p_h0 = h0 + (orig_seq_id * HV + i_h) * V * K + o_v[:, None] * K + o_k[None, :] + b_h = tl.load(p_h0, mask=m_v[:, None] & m_k[None, :], other=0.0).to(tl.float32) else: - p_h0 = tl.make_block_ptr( - h0 + (orig_seq_id * HV + i_h) * K * V, - (K, V), (V, 1), (0, i_v * BV), (BK, BV), (1, 0) - ) - b_h = tl.load(p_h0, boundary_check=(0, 1)).to(tl.float32) + p_h0 = h0 + (orig_seq_id * HV + i_h) * K * V + o_k[:, None] * V + o_v[None, :] + b_h = tl.load(p_h0, mask=m_k[:, None] & m_v[None, :], other=0.0).to(tl.float32) else: if STATE_V_FIRST: b_h = tl.zeros([BV, BK], dtype=tl.float32) @@ -395,39 +413,37 @@ def merge_fwd_bwd_kernel( base = i_ss * stride_hm_s + i_h * stride_hm_h # he and m are always in [K, V+K] layout from pre_scan - p_he = tl.make_block_ptr( - ag_hm + base, (K, V), (V + K, 1), (0, i_v * BV), (BK, BV), (1, 0) - ) - b_he = tl.load(p_he, boundary_check=(0, 1)).to(tl.float32) - p_m = tl.make_block_ptr( - ag_hm + base + V, (K, K), (V + K, 1), (0, 0), (BK, BK), (1, 0) - ) - b_m = tl.load(p_m, boundary_check=(0, 1)).to(tl.float32) + p_he = ag_hm + base + o_k[:, None] * (V + K) + o_v[None, :] + b_he = tl.load(p_he, mask=m_k[:, None] & m_v[None, :], other=0.0).to(tl.float32) + p_m = ag_hm + base + V + o_k[:, None] * (V + K) + o_k[None, :] + b_m = tl.load(p_m, mask=m_k[:, None] & m_k[None, :], other=0.0).to(tl.float32) if STATE_V_FIRST: # h_T' = h_T @ M^T + he^T - b_h = tl.dot(b_h.to(tl.float32), tl.trans(b_m)) + tl.trans(b_he) + b_h = tl.dot(b_h.to(tl.float32), tl.trans(b_m), input_precision=AFFINE_CHAIN_PRECISION) + tl.trans(b_he) else: - b_h = tl.dot(b_m.to(tl.float32), b_h.to(tl.float32)) + b_he.to(tl.float32) + b_h = tl.dot(b_m.to(tl.float32), b_h.to(tl.float32), + input_precision=AFFINE_CHAIN_PRECISION) + b_he.to(tl.float32) # Store for non-first subseqs if idx < num_subseqs - 1: init_idx = init_base + idx stride_init = HV * K * V if STATE_V_FIRST: - p_out = tl.make_block_ptr( - h + init_idx * stride_init + i_h * V * K, - (V, K), (K, 1), (i_v * BV, 0), (BV, BK), (1, 0) - ) + p_out = h + init_idx * stride_init + i_h * V * K + o_v[:, None] * K + o_k[None, :] + m_out = m_v[:, None] & m_k[None, :] else: - p_out = tl.make_block_ptr( - h + init_idx * stride_init + i_h * K * V, - (K, V), (V, 1), (0, i_v * BV), (BK, BV), (1, 0) - ) - tl.store(p_out, b_h.to(p_out.dtype.element_ty), boundary_check=(0, 1)) + p_out = h + init_idx * stride_init + i_h * K * V + o_k[:, None] * V + o_v[None, :] + m_out = m_k[:, None] & m_v[None, :] + tl.store(p_out, b_h.to(p_out.dtype.element_ty), mask=m_out) else: # CP mode - i_h = tl.program_id(1) - num_ranks = pre_or_post_num_ranks.to(tl.int32) + i_h = tl.program_id(1).to(tl.int64) + if HAS_H_SEQ_IDX: + h += tl.load(h_seq_idx).to(tl.int64) * HV * K * V + if NUM_RANKS_ON_DEVICE: + num_ranks = tl.load(pre_or_post_num_ranks).to(tl.int32) + else: + num_ranks = pre_or_post_num_ranks.to(tl.int32) h += i_h * K * V ag_hm += i_h * K * (K + V) stride = HV * K * (K + V) @@ -440,19 +456,23 @@ def merge_fwd_bwd_kernel( cur_rank = rank - num_ranks + idx else: cur_rank = rank + num_ranks - idx - p_ag_h = tl.make_block_ptr(ag_hm + cur_rank * stride, (K, V), (K + V, 1), (0, i_v * BV), (BK, BV), (1, 0)) - b_ag_h = tl.load(p_ag_h, boundary_check=(0, 1)) - p_ag_m = tl.make_block_ptr(ag_hm + cur_rank * stride + V, (K, K), (K + V, 1), (0, 0), (BK, BK), (1, 0)) - b_ag_m = tl.load(p_ag_m, boundary_check=(0, 1)) + p_ag_h = ag_hm + cur_rank * stride + o_k[:, None] * (K + V) + o_v[None, :] + b_ag_h = tl.load(p_ag_h, mask=m_k[:, None] & m_v[None, :], other=0.0) + p_ag_m = ag_hm + cur_rank * stride + V + o_k[:, None] * (K + V) + o_k[None, :] + b_ag_m = tl.load(p_ag_m, mask=m_k[:, None] & m_k[None, :], other=0.0) if STATE_V_FIRST: - b_h = tl.dot(b_h.to(tl.float32), tl.trans(b_ag_m).to(tl.float32)) + tl.trans(b_ag_h).to(tl.float32) + b_h = tl.dot(b_h.to(tl.float32), tl.trans(b_ag_m).to(tl.float32), + input_precision=AFFINE_CHAIN_PRECISION) + tl.trans(b_ag_h).to(tl.float32) else: - b_h = tl.dot(b_ag_m.to(tl.float32), b_h.to(tl.float32)) + b_ag_h.to(tl.float32) + b_h = tl.dot(b_ag_m.to(tl.float32), b_h.to(tl.float32), + input_precision=AFFINE_CHAIN_PRECISION) + b_ag_h.to(tl.float32) if STATE_V_FIRST: - p_h = tl.make_block_ptr(h, (V, K), (K, 1), (i_v * BV, 0), (BV, BK), (1, 0)) + p_h = h + o_v[:, None] * K + o_k[None, :] + m_h = m_v[:, None] & m_k[None, :] else: - p_h = tl.make_block_ptr(h, (K, V), (V, 1), (0, i_v * BV), (BK, BV), (1, 0)) - tl.store(p_h, b_h.to(p_h.dtype.element_ty), boundary_check=(0, 1)) + p_h = h + o_k[:, None] * V + o_v[None, :] + m_h = m_k[:, None] & m_v[None, :] + tl.store(p_h, b_h.to(p_h.dtype.element_ty), mask=m_h) @triton.heuristics({ @@ -493,6 +513,7 @@ def pre_process_bwd_kernel_merged( USE_GK: tl.constexpr, USE_BG: tl.constexpr, IS_VARLEN: tl.constexpr, + AFFINE_CHAIN_PRECISION: tl.constexpr = None, ): """ Merged backward kernel that computes both dh (K x V) and dm (K x K) in a single kernel. @@ -501,7 +522,7 @@ def pre_process_bwd_kernel_merged( - Columns [0, V) are for computing dh (stage 1) - Columns [V, V+K) are for computing dm (stage 2) """ - i_col, i_h = tl.program_id(0), tl.program_id(1) + i_col, i_h = tl.program_id(0).to(tl.int64), tl.program_id(1).to(tl.int64) i_n = 0 if IS_VARLEN: bos, eos = tl.load(cu_seqlens + i_n).to(tl.int64), tl.load(cu_seqlens + i_n + 1).to(tl.int64) @@ -547,64 +568,72 @@ def pre_process_bwd_kernel_merged( if K > 192: b_dh4 = tl.zeros([64, BLOCK_SIZE], dtype=tl.float32) + o_vb = i_v * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE) + m_vb = o_vb < V + o_k1 = tl.arange(0, 64) + m_k1 = o_k1 < K + o_k2 = 64 + o_k1 + m_k2 = o_k2 < K + o_k3 = 128 + o_k1 + m_k3 = o_k3 < K + o_k4 = 192 + o_k1 + m_k4 = o_k4 < K + # Main recurrence for dh (reverse order) for i_t in range(NT - 1, -1, -1): last_idx = min((i_t + 1) * BT, T) - 1 + o_t = i_t * BT + tl.arange(0, BT) + m_t = o_t < T if USE_G: bg_last = tl.load(g + last_idx * HV).to(tl.float32) - p_g = tl.make_block_ptr(g, (T,), (HV,), (i_t * BT,), (BT,), (0,)) - b_g = tl.load(p_g, boundary_check=(0,)).to(tl.float32) + p_g = g + o_t * HV + b_g = tl.load(p_g, mask=m_t, other=0.0).to(tl.float32) bg_last_exp = exp2(bg_last) b_g_exp = exp2(b_g) - p_dv = tl.make_block_ptr(dv, (T, V), (stride_v, 1), (i_t * BT, i_v * BLOCK_SIZE), (BT, BLOCK_SIZE), (1, 0)) - p_do = tl.make_block_ptr(do, (T, V), (stride_v, 1), (i_t * BT, i_v * BLOCK_SIZE), (BT, BLOCK_SIZE), (1, 0)) - b_do = tl.load(p_do, boundary_check=(0, 1)) + p_dv = dv + o_t[:, None] * stride_v + o_vb[None, :] + p_do = do + o_t[:, None] * stride_v + o_vb[None, :] + b_do = tl.load(p_do, mask=m_t[:, None] & m_vb[None, :], other=0.0) # Update dv - p_k = tl.make_block_ptr(k, (T, K), (stride_qk, 1), (i_t * BT, 0), (BT, 64), (1, 0)) - b_k = tl.load(p_k, boundary_check=(0, 1)) + p_k = k + o_t[:, None] * stride_qk + o_k1[None, :] + b_k = tl.load(p_k, mask=m_t[:, None] & m_k1[None, :], other=0.0) if USE_GK: - o_k1 = tl.arange(0, 64) p_gk_last = gk + last_idx * HV * K b_gk_last1 = tl.load(p_gk_last + o_k1, mask=(o_k1 < K), other=0.).to(tl.float32) b_dv = tl.dot(b_k, b_dh1.to(b_k.dtype)) if K > 64: - p_k = tl.make_block_ptr(k, (T, K), (stride_qk, 1), (i_t * BT, 64), (BT, 64), (1, 0)) - b_k = tl.load(p_k, boundary_check=(0, 1)) + p_k = k + o_t[:, None] * stride_qk + o_k2[None, :] + b_k = tl.load(p_k, mask=m_t[:, None] & m_k2[None, :], other=0.0) if USE_GK: - o_k2 = 64 + o_k1 b_gk_last2 = tl.load(p_gk_last + o_k2, mask=(o_k2 < K), other=0.).to(tl.float32) - b_dv += tl.dot(b_k, b_dh2.to(b_k.dtype)) + b_dv = tl.dot(b_k, b_dh2.to(b_k.dtype), b_dv) if K > 128: - p_k = tl.make_block_ptr(k, (T, K), (stride_qk, 1), (i_t * BT, 128), (BT, 64), (1, 0)) - b_k = tl.load(p_k, boundary_check=(0, 1)) + p_k = k + o_t[:, None] * stride_qk + o_k3[None, :] + b_k = tl.load(p_k, mask=m_t[:, None] & m_k3[None, :], other=0.0) if USE_GK: - o_k3 = 128 + o_k1 b_gk_last3 = tl.load(p_gk_last + o_k3, mask=(o_k3 < K), other=0.).to(tl.float32) - b_dv += tl.dot(b_k, b_dh3.to(b_k.dtype)) + b_dv = tl.dot(b_k, b_dh3.to(b_k.dtype), b_dv) if K > 192: - p_k = tl.make_block_ptr(k, (T, K), (stride_qk, 1), (i_t * BT, 192), (BT, 64), (1, 0)) - b_k = tl.load(p_k, boundary_check=(0, 1)) + p_k = k + o_t[:, None] * stride_qk + o_k4[None, :] + b_k = tl.load(p_k, mask=m_t[:, None] & m_k4[None, :], other=0.0) if USE_GK: - o_k4 = 192 + o_k1 b_gk_last4 = tl.load(p_gk_last + o_k4, mask=(o_k4 < K), other=0.).to(tl.float32) - b_dv += tl.dot(b_k, b_dh4.to(b_k.dtype)) + b_dv = tl.dot(b_k, b_dh4.to(b_k.dtype), b_dv) if USE_G: - m_t = (i_t * BT + tl.arange(0, BT)) < T b_dv *= tl.where(m_t, exp2(bg_last - b_g), 0)[:, None] - b_dv += tl.load(p_dv, boundary_check=(0, 1)) + b_dv += tl.load(p_dv, mask=m_t[:, None] & m_vb[None, :], other=0.0) # Update dh - p_w = tl.make_block_ptr(w, (K, T), (1, stride_w), (0, i_t * BT), (64, BT), (0, 1)) - p_q = tl.make_block_ptr(q, (K, T), (1, stride_qk), (0, i_t * BT), (64, BT), (0, 1)) - b_w = tl.load(p_w, boundary_check=(0, 1)) - b_q = tl.load(p_q, boundary_check=(0, 1)) + p_w = w + o_k1[:, None] + o_t[None, :] * stride_w + p_q = q + o_k1[:, None] + o_t[None, :] * stride_qk + b_w = tl.load(p_w, mask=m_k1[:, None] & m_t[None, :], other=0.0) + b_q = tl.load(p_q, mask=m_k1[:, None] & m_t[None, :], other=0.0) if USE_G: b_dh1 *= bg_last_exp b_q = b_q * b_g_exp[None, :] @@ -617,10 +646,10 @@ def pre_process_bwd_kernel_merged( b_dh1 += tl.dot(b_q.to(b_q.dtype), b_do.to(b_q.dtype)) * scale - tl.dot(b_w, b_dv.to(b_w.dtype)) if K > 64: - p_q = tl.make_block_ptr(q, (K, T), (1, stride_qk), (64, i_t * BT), (64, BT), (0, 1)) - p_w = tl.make_block_ptr(w, (K, T), (1, stride_w), (64, i_t * BT), (64, BT), (0, 1)) - b_q = tl.load(p_q, boundary_check=(0, 1)) - b_w = tl.load(p_w, boundary_check=(0, 1)) + p_q = q + o_k2[:, None] + o_t[None, :] * stride_qk + p_w = w + o_k2[:, None] + o_t[None, :] * stride_w + b_q = tl.load(p_q, mask=m_k2[:, None] & m_t[None, :], other=0.0) + b_w = tl.load(p_w, mask=m_k2[:, None] & m_t[None, :], other=0.0) if USE_G: b_dh2 *= bg_last_exp b_q = b_q * b_g_exp[None, :] @@ -632,10 +661,10 @@ def pre_process_bwd_kernel_merged( b_dh2 += tl.dot(b_q.to(b_q.dtype), b_do.to(b_q.dtype)) * scale - tl.dot(b_w, b_dv.to(b_w.dtype)) if K > 128: - p_q = tl.make_block_ptr(q, (K, T), (1, stride_qk), (128, i_t * BT), (64, BT), (0, 1)) - p_w = tl.make_block_ptr(w, (K, T), (1, stride_w), (128, i_t * BT), (64, BT), (0, 1)) - b_q = tl.load(p_q, boundary_check=(0, 1)) - b_w = tl.load(p_w, boundary_check=(0, 1)) + p_q = q + o_k3[:, None] + o_t[None, :] * stride_qk + p_w = w + o_k3[:, None] + o_t[None, :] * stride_w + b_q = tl.load(p_q, mask=m_k3[:, None] & m_t[None, :], other=0.0) + b_w = tl.load(p_w, mask=m_k3[:, None] & m_t[None, :], other=0.0) if USE_G: b_dh3 *= bg_last_exp b_q = b_q * b_g_exp[None, :] @@ -647,10 +676,10 @@ def pre_process_bwd_kernel_merged( b_dh3 += tl.dot(b_q.to(b_q.dtype), b_do.to(b_q.dtype)) * scale - tl.dot(b_w, b_dv.to(b_w.dtype)) if K > 192: - p_q = tl.make_block_ptr(q, (K, T), (1, stride_qk), (192, i_t * BT), (64, BT), (0, 1)) - p_w = tl.make_block_ptr(w, (K, T), (1, stride_w), (192, i_t * BT), (64, BT), (0, 1)) - b_q = tl.load(p_q, boundary_check=(0, 1)) - b_w = tl.load(p_w, boundary_check=(0, 1)) + p_q = q + o_k4[:, None] + o_t[None, :] * stride_qk + p_w = w + o_k4[:, None] + o_t[None, :] * stride_w + b_q = tl.load(p_q, mask=m_k4[:, None] & m_t[None, :], other=0.0) + b_w = tl.load(p_w, mask=m_k4[:, None] & m_t[None, :], other=0.0) if USE_G: b_dh4 *= bg_last_exp b_q = b_q * b_g_exp[None, :] @@ -662,17 +691,17 @@ def pre_process_bwd_kernel_merged( b_dh4 += tl.dot(b_q.to(b_q.dtype), b_do.to(b_q.dtype)) * scale - tl.dot(b_w, b_dv.to(b_w.dtype)) # Store dh results - p_dh1 = tl.make_block_ptr(dhm, (K, V), (V + K, 1), (0, i_v * BLOCK_SIZE), (64, BLOCK_SIZE), (1, 0)) - tl.store(p_dh1, b_dh1.to(p_dh1.dtype.element_ty), boundary_check=(0, 1)) + p_dh1 = dhm + o_k1[:, None] * (V + K) + o_vb[None, :] + tl.store(p_dh1, b_dh1.to(p_dh1.dtype.element_ty), mask=m_k1[:, None] & m_vb[None, :]) if K > 64: - p_dh2 = tl.make_block_ptr(dhm, (K, V), (V + K, 1), (64, i_v * BLOCK_SIZE), (64, BLOCK_SIZE), (1, 0)) - tl.store(p_dh2, b_dh2.to(p_dh2.dtype.element_ty), boundary_check=(0, 1)) + p_dh2 = dhm + o_k2[:, None] * (V + K) + o_vb[None, :] + tl.store(p_dh2, b_dh2.to(p_dh2.dtype.element_ty), mask=m_k2[:, None] & m_vb[None, :]) if K > 128: - p_dh3 = tl.make_block_ptr(dhm, (K, V), (V + K, 1), (128, i_v * BLOCK_SIZE), (64, BLOCK_SIZE), (1, 0)) - tl.store(p_dh3, b_dh3.to(p_dh3.dtype.element_ty), boundary_check=(0, 1)) + p_dh3 = dhm + o_k3[:, None] * (V + K) + o_vb[None, :] + tl.store(p_dh3, b_dh3.to(p_dh3.dtype.element_ty), mask=m_k3[:, None] & m_vb[None, :]) if K > 192: - p_dh4 = tl.make_block_ptr(dhm, (K, V), (V + K, 1), (192, i_v * BLOCK_SIZE), (64, BLOCK_SIZE), (1, 0)) - tl.store(p_dh4, b_dh4.to(p_dh4.dtype.element_ty), boundary_check=(0, 1)) + p_dh4 = dhm + o_k4[:, None] * (V + K) + o_vb[None, :] + tl.store(p_dh4, b_dh4.to(p_dh4.dtype.element_ty), mask=m_k4[:, None] & m_vb[None, :]) else: # ====== Stage 2: Compute dm (K x K) ====== # i_col is for dm part, map to K dimension @@ -690,20 +719,21 @@ def pre_process_bwd_kernel_merged( for _i_t in range(NT): # Reverse order for backward i_t = NT - 1 - _i_t + o_t = i_t * BT + tl.arange(0, BT) + m_t = o_t < T # Load k and w with full BK1 rows - p_k = tl.make_block_ptr(k, (T, K), (stride_qk, 1), (i_t * BT, 0), (BT, BK1), (1, 0)) - b_k = tl.load(p_k, boundary_check=(0, 1)) - p_w = tl.make_block_ptr(w, (T, K), (stride_w, 1), (i_t * BT, 0), (BT, BK1), (1, 0)) - b_w = tl.load(p_w, boundary_check=(0, 1)) + p_k = k + o_t[:, None] * stride_qk + row[None, :] + b_k = tl.load(p_k, mask=m_t[:, None] & (row < K)[None, :], other=0.0) + p_w = w + o_t[:, None] * stride_w + row[None, :] + b_w = tl.load(p_w, mask=m_t[:, None] & (row < K)[None, :], other=0.0) last_idx = min((i_t + 1) * BT, T) - 1 if USE_G: - m_t = (i_t * BT + tl.arange(0, BT)) < T b_g_last = tl.load(g + last_idx * HV).to(tl.float32) - p_g = tl.make_block_ptr(g, (T,), (HV,), (i_t * BT,), (BT,), (0,)) - b_g = tl.load(p_g, boundary_check=(0,)).to(tl.float32) + p_g = g + o_t * HV + b_g = tl.load(p_g, mask=m_t, other=0.0).to(tl.float32) b_k = b_k * tl.where(m_t, exp2(b_g_last - b_g), 0)[:, None] b_g_last = exp2(b_g_last) b_diag = tl.where(row[:, None] == row[None, :], b_g_last, 0.0) @@ -723,11 +753,11 @@ def pre_process_bwd_kernel_merged( # GDN/KDA mode: dM = (diag - w^T @ k) @ dM b_m_i = b_diag - b_kw # Keep m chain in fp32 to avoid precision loss from repeated bf16 casting - b_m = tl.dot(b_m_i.to(tl.float32), b_m.to(tl.float32)) + b_m = tl.dot(b_m_i.to(tl.float32), b_m.to(tl.float32), input_precision=AFFINE_CHAIN_PRECISION) # Store dm result - p_m = tl.make_block_ptr(dhm + V, (K, K), (V + K, 1), (0, i_k_col * BLOCK_SIZE), (BK1, BLOCK_SIZE), (1, 0)) - tl.store(p_m, b_m.to(p_m.dtype.element_ty), boundary_check=(0, 1)) + p_m = dhm + V + row[:, None] * (V + K) + col[None, :] + tl.store(p_m, b_m.to(p_m.dtype.element_ty), mask=(row < K)[:, None] & (col < K)[None, :]) def chunk_gated_delta_rule_fwd_h_pre_process( @@ -743,10 +773,17 @@ def chunk_gated_delta_rule_fwd_h_pre_process( cu_seqlens: torch.LongTensor | None = None, initial_state: torch.Tensor | None = None, context: FLACPContext = None, + use_graph: bool = False, ) -> tuple[torch.Tensor, torch.Tensor]: if context is None or context.group is None: return initial_state assert initial_state is None, "When enable CP, the provided initial_state must be None." + use_tf32x3_affine_chain = context.use_tf32x3_affine_chain + if use_tf32x3_affine_chain and not IS_TF32_SUPPORTED: + warnings.warn( + "tf32x3 affine chain requires an NVIDIA GPU with compute capability >= 8.0; falling back to ieee precision", + stacklevel=2, + ) rank = dist.get_rank(group=context.group) B, T, H, K, V, HV = *k.shape, u.shape[-1], u.shape[2] @@ -760,12 +797,94 @@ def chunk_gated_delta_rule_fwd_h_pre_process( N = len(cu_seqlens) - 1 assert K <= 256, "current kernel does not support head dimension larger than 256." + if context.layout == 'zigzag': + if use_graph: + raise NotImplementedError("use_graph is not supported with zigzag CP layout") + fns = context.front_num_seqs + hm = k.new_zeros(2, HV, K, V + K, dtype=torch.float32) + if state_v_first: + initial_state = k.new_zeros(N, HV, V, K, dtype=torch.float32) + else: + initial_state = k.new_zeros(N, HV, K, V, dtype=torch.float32) + BLOCK_SIZE = 32 if K <= 64 else 64 + grid_hm = (triton.cdiv(V, BLOCK_SIZE) + triton.cdiv(K, BLOCK_SIZE), HV) + # each part exports the affine chain of its last segment, skipped when chain-last + for part, (cu_win, is_last) in enumerate(zip( + (cu_seqlens[fns - 1: fns + 1], cu_seqlens[-2:]), context.is_last_by_part)): + if not is_last: + pre_process_fwd_kernel_merged[grid_hm]( + k=k, + v=u if v is None else v, + w=w, + g=g, + gk=gk, + bg=bg, + u=u, + hm=hm[part], + cu_seqlens=cu_win, + T=T, + H=H, + HV=HV, + K=K, + V=V, + BT=BT, + BK1=BK, + BLOCK_SIZE=BLOCK_SIZE, + MULTI_SEQS=False, + AFFINE_CHAIN_PRECISION=( + "tf32x3" if use_tf32x3_affine_chain and IS_TF32_SUPPORTED + else ("ieee" if not IS_TF32_SUPPORTED else None) + ), + ) + ag_hm, _ = all_gather_into_tensor(hm, group=context.group) + # chain order: front parts of ranks 0..W-1, then back parts of ranks W-1..0 + slots = torch.cat([ag_hm[:, 0], ag_hm[:, 1].flip(0)]) + world_size = dist.get_world_size(context.group) + + def grid(meta): return (triton.cdiv(V, meta['BV']), HV) + for part in range(2): + if not context.is_first_by_part[part]: + merge_fwd_bwd_kernel[grid]( + h=initial_state[0] if part == 0 else initial_state[fns], + ag_hm=slots, + pre_or_post_num_ranks=context.pre_num_ranks_by_part[part], + rank=rank if part == 0 else 2 * world_size - 1 - rank, + seq_offsets=None, + init_offsets=None, + h0_seq_ids=None, + h0=None, + h_seq_idx=None, + HV=HV, + K=K, + V=V, + BK=BK, + FORWARD=True, + INTRACARD_MODE=False, + NUM_SEQ_ENTRIES=0, + STATE_V_FIRST=state_v_first, + AFFINE_CHAIN_PRECISION=( + "tf32x3" if use_tf32x3_affine_chain and IS_TF32_SUPPORTED + else ("ieee" if not IS_TF32_SUPPORTED else None) + ), + NUM_RANKS_ON_DEVICE=False, + ) + return initial_state + hm = k.new_zeros(HV, K, (V + K), dtype=torch.float32) if state_v_first: initial_state = k.new_zeros(N, HV, V, K, dtype=torch.float32) else: initial_state = k.new_zeros(N, HV, K, V, dtype=torch.float32) - if not context.is_last_rank: + cu_last = cu_seqlens[-2:] + if use_graph: + assert context.pre_num_ranks_dev is not None, "use_graph with CP requires context.pre_num_ranks_dev" + # last non-empty sequence window, on-device; zero-length tail padding would + # otherwise make cu_seqlens[-2:] point at a padding sequence + i_last = (cu_seqlens[1:] > cu_seqlens[:-1]).sum() - 1 + cu_last = cu_seqlens.index_select(0, torch.stack((i_last, i_last + 1))) + # graph mode always launches: a last rank's hm is never read, but the all-gather + # needs every rank, and host-side flags are frozen at capture + if use_graph or not context.is_last_rank: BLOCK_SIZE = 32 if K <= 64 else 64 grid = (triton.cdiv(V, BLOCK_SIZE) + triton.cdiv(K, BLOCK_SIZE), HV) # For DPLR, v provides the original v for computing h contributions, @@ -779,7 +898,7 @@ def chunk_gated_delta_rule_fwd_h_pre_process( bg=bg, u=u, hm=hm, - cu_seqlens=cu_seqlens[-2:], + cu_seqlens=cu_last, T=T, H=H, HV=HV, @@ -789,19 +908,25 @@ def chunk_gated_delta_rule_fwd_h_pre_process( BK1=BK, BLOCK_SIZE=BLOCK_SIZE, MULTI_SEQS=False, + AFFINE_CHAIN_PRECISION=( + "tf32x3" if use_tf32x3_affine_chain and IS_TF32_SUPPORTED + else ("ieee" if not IS_TF32_SUPPORTED else None) + ), ) ag_hm, _ = all_gather_into_tensor(hm, group=context.group) - if not context.is_first_rank: + # a zero num_ranks merge just re-stores the zeros initial_state[0] already holds + if use_graph or not context.is_first_rank: def grid(meta): return (triton.cdiv(V, meta['BV']), HV) merge_fwd_bwd_kernel[grid]( h=initial_state[0], ag_hm=ag_hm, - pre_or_post_num_ranks=context.pre_num_ranks, + pre_or_post_num_ranks=context.pre_num_ranks_dev if use_graph else context.pre_num_ranks, rank=rank, seq_offsets=None, init_offsets=None, h0_seq_ids=None, h0=None, + h_seq_idx=None, HV=HV, K=K, V=V, @@ -810,6 +935,11 @@ def chunk_gated_delta_rule_fwd_h_pre_process( INTRACARD_MODE=False, NUM_SEQ_ENTRIES=0, STATE_V_FIRST=state_v_first, + AFFINE_CHAIN_PRECISION=( + "tf32x3" if use_tf32x3_affine_chain and IS_TF32_SUPPORTED + else ("ieee" if not IS_TF32_SUPPORTED else None) + ), + NUM_RANKS_ON_DEVICE=use_graph, ) return initial_state @@ -830,10 +960,17 @@ def chunk_gated_delta_rule_bwd_dhu_pre_process( initial_state: torch.Tensor | None = None, context: FLACPContext | None = None, chunk_size: int = 64, + use_graph: bool = False, ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: if context is None or context.group is None: return dht, initial_state assert dht is None, "When enable CP, the provided dht must be None." + use_tf32x3_affine_chain = context.use_tf32x3_affine_chain + if use_tf32x3_affine_chain and not IS_TF32_SUPPORTED: + warnings.warn( + "tf32x3 affine chain requires an NVIDIA GPU with compute capability >= 8.0; falling back to ieee precision", + stacklevel=2, + ) rank = dist.get_rank(context.group) B, T, H, K, V, HV = *q.shape, do.shape[-1], do.shape[2] @@ -847,13 +984,89 @@ def chunk_gated_delta_rule_bwd_dhu_pre_process( else: N = len(cu_seqlens) - 1 + if context.layout == 'zigzag': + if use_graph: + raise NotImplementedError("use_graph is not supported with zigzag CP layout") + fns = context.front_num_seqs + dhm = q.new_zeros(2, HV, K, V + K, dtype=torch.float32) + if state_v_first: + dht = q.new_zeros(N, HV, V, K, dtype=torch.float32) + else: + dht = q.new_zeros(N, HV, K, V, dtype=torch.float32) + BLOCK_SIZE = 32 if K <= 64 else 64 + grid_hm = (triton.cdiv(V, BLOCK_SIZE) + triton.cdiv(K, BLOCK_SIZE), HV) + # each part exports the backward affine chain of its first segment, skipped when chain-first + for part, (cu_win, is_first) in enumerate(zip( + (cu_seqlens[:2], cu_seqlens[fns: fns + 2]), context.is_first_by_part)): + if not is_first: + pre_process_bwd_kernel_merged[grid_hm]( + q=q, + k=k if bg is None else bg, + w=w, + g=g, + gk=gk, + do=do, + dhm=dhm[part], + dv=dv, + cu_seqlens=cu_win, + scale=scale, + T=T, + H=H, + HV=HV, + K=K, + V=V, + BT=BT, + BK1=BK, + BLOCK_SIZE=BLOCK_SIZE, + USE_BG=bg is not None, + AFFINE_CHAIN_PRECISION=( + "tf32x3" if use_tf32x3_affine_chain and IS_TF32_SUPPORTED + else ("ieee" if not IS_TF32_SUPPORTED else None) + ), + ) + ag_dhm, _ = all_gather_into_tensor(dhm, group=context.group) + # chain order: front parts of ranks 0..W-1, then back parts of ranks W-1..0 + slots = torch.cat([ag_dhm[:, 0], ag_dhm[:, 1].flip(0)]) + world_size = dist.get_world_size(context.group) + + def grid(meta): return (triton.cdiv(V, meta['BV']), HV) + for part in range(2): + if not context.is_last_by_part[part]: + merge_fwd_bwd_kernel[grid]( + h=dht[fns - 1] if part == 0 else dht[-1], + ag_hm=slots, + pre_or_post_num_ranks=context.post_num_ranks_by_part[part], + rank=rank if part == 0 else 2 * world_size - 1 - rank, + seq_offsets=None, + init_offsets=None, + h0_seq_ids=None, + h0=None, + h_seq_idx=None, + HV=HV, + K=K, + V=V, + BK=BK, + FORWARD=False, + INTRACARD_MODE=False, + NUM_SEQ_ENTRIES=0, + STATE_V_FIRST=state_v_first, + AFFINE_CHAIN_PRECISION=( + "tf32x3" if use_tf32x3_affine_chain and IS_TF32_SUPPORTED + else ("ieee" if not IS_TF32_SUPPORTED else None) + ), + NUM_RANKS_ON_DEVICE=False, + ) + return dht, None + dhm = q.new_zeros(HV, K, V + K, dtype=torch.float32) if state_v_first: dht = q.new_zeros(N, HV, V, K, dtype=torch.float32) else: dht = q.new_zeros(N, HV, K, V, dtype=torch.float32) - if not context.is_first_rank: + # graph mode always launches: a first rank's dhm is never read, but the all-gather + # needs every rank, and host-side flags are frozen at capture + if use_graph or not context.is_first_rank: BLOCK_SIZE = 32 if K <= 64 else 64 grid = (triton.cdiv(V, BLOCK_SIZE) + triton.cdiv(K, BLOCK_SIZE), HV) pre_process_bwd_kernel_merged[grid]( @@ -876,21 +1089,36 @@ def chunk_gated_delta_rule_bwd_dhu_pre_process( BK1=BK, BLOCK_SIZE=BLOCK_SIZE, USE_BG=bg is not None, + AFFINE_CHAIN_PRECISION=( + "tf32x3" if use_tf32x3_affine_chain and IS_TF32_SUPPORTED + else ("ieee" if not IS_TF32_SUPPORTED else None) + ), ) ag_dhm, _ = all_gather_into_tensor(dhm, group=context.group) - if not context.is_last_rank: + # a zero num_ranks merge just re-stores the zeros dht already holds + if use_graph or not context.is_last_rank: + h_seq_idx = None + h = dht[-1] + if use_graph: + assert context.post_num_ranks_dev is not None, "use_graph with CP requires context.post_num_ranks_dev" + # merge target is the last non-empty sequence; zero-length tail padding + # would otherwise make dht[-1] a padding row + h_seq_idx = (cu_seqlens[1:] > cu_seqlens[:-1]).sum() - 1 + h = dht + def grid(meta): return (triton.cdiv(V, meta['BV']), HV) merge_fwd_bwd_kernel[grid]( - h=dht[-1], + h=h, ag_hm=ag_dhm, - pre_or_post_num_ranks=context.post_num_ranks, + pre_or_post_num_ranks=context.post_num_ranks_dev if use_graph else context.post_num_ranks, rank=rank, seq_offsets=None, init_offsets=None, h0_seq_ids=None, h0=None, + h_seq_idx=h_seq_idx, HV=HV, K=K, V=V, @@ -899,6 +1127,11 @@ def chunk_gated_delta_rule_bwd_dhu_pre_process( INTRACARD_MODE=False, NUM_SEQ_ENTRIES=0, STATE_V_FIRST=state_v_first, + AFFINE_CHAIN_PRECISION=( + "tf32x3" if use_tf32x3_affine_chain and IS_TF32_SUPPORTED + else ("ieee" if not IS_TF32_SUPPORTED else None) + ), + NUM_RANKS_ON_DEVICE=use_graph, ) # initial_state is None in the CP mode @@ -909,6 +1142,12 @@ def chunk_gated_delta_rule_bwd_dhu_pre_process( def compress_h0(h0: torch.Tensor, context: FLACPContext): if h0 is None or len(context.cu_seqlens) == 2: return h0 + if context.layout == 'zigzag': + if len(context.cu_seqlens) == 3: + # one segment per part; both rows may carry a cross-rank state + return h0 + # only each part's first segment can carry a cross-rank state + return h0[[0, context.front_num_seqs]] # Here must use clone op or the full tensor will be saved for backward return h0[:1].clone() @@ -917,6 +1156,12 @@ def expand_h0(h0: torch.Tensor, context: FLACPContext): if h0 is None or len(context.cu_seqlens) == 2: return h0 B = len(context.cu_seqlens) - 1 + if context.layout == 'zigzag': + if B == 2: + return h0 + expand_h0 = h0.new_zeros(B, *h0.shape[1:]) + expand_h0[[0, context.front_num_seqs]] = h0 + return expand_h0 expand_h0 = h0.new_zeros(B, *h0.shape[1:]) expand_h0[:1] = h0 return expand_h0 diff --git a/build/torch-cuda/ops/cp/context.py b/build/torch-cuda/ops/cp/context.py index 0fd0bd5dd5eb63e9f47e24234e3d9ef7fbcd18d2..9054ab8197e016d7afa1549f47f129b92255df78 100644 --- a/build/torch-cuda/ops/cp/context.py +++ b/build/torch-cuda/ops/cp/context.py @@ -31,6 +31,22 @@ class FLACPContext: post_num_ranks: int | None = None conv1d_kernel_size: int | None = None pre_num_conv_tokens: int | None = None + use_tf32x3_affine_chain: bool = False + # Device-resident copies of pre/post_num_ranks for graph replay: kernel scalar + # arguments are frozen at capture, so graph mode reads them via tl.load instead. + # Persistent int32 tensors the caller refreshes in place before each replay. + pre_num_ranks_dev: torch.Tensor | None = None + post_num_ranks_dev: torch.Tensor | None = None + # zigzag layout: rank r holds chain slots r (front part) and 2 * world_size - 1 - r + # (back part); the fields below are per part, the scalar fields above contiguous only. + layout: str = 'contiguous' + part_len: int | None = None + front_num_seqs: int | None = None + pre_num_ranks_by_part: tuple[int, int] | None = None + post_num_ranks_by_part: tuple[int, int] | None = None + is_first_by_part: tuple[bool, bool] | None = None + is_last_by_part: tuple[bool, bool] | None = None + pre_num_conv_tokens_by_part: tuple[int, int] | None = None def copy_for_backward(self) -> FLACPContext: """Create a copy for backward pass (useful when PP_SIZE > 1).""" @@ -44,6 +60,17 @@ class FLACPContext: post_num_ranks=self.post_num_ranks, conv1d_kernel_size=self.conv1d_kernel_size, pre_num_conv_tokens=self.pre_num_conv_tokens, + use_tf32x3_affine_chain=self.use_tf32x3_affine_chain, + pre_num_ranks_dev=self.pre_num_ranks_dev, + post_num_ranks_dev=self.post_num_ranks_dev, + layout=self.layout, + part_len=self.part_len, + front_num_seqs=self.front_num_seqs, + pre_num_ranks_by_part=self.pre_num_ranks_by_part, + post_num_ranks_by_part=self.post_num_ranks_by_part, + is_first_by_part=self.is_first_by_part, + is_last_by_part=self.is_last_by_part, + pre_num_conv_tokens_by_part=self.pre_num_conv_tokens_by_part, ) @property @@ -57,6 +84,47 @@ class FLACPContext: return self.group is not None +def _interval_cp_meta( + cu_seqlens_cpu: torch.LongTensor, + start: int, + end: int, + chain_pos: int, + part_len: int, +) -> tuple[torch.Tensor, int, int, bool, int, bool]: + """Local cu_seqlens and chain metadata for one token interval [start, end). + + `chain_pos` is the interval's position in the global processing chain and + `part_len` the chain's interval size; both reduce to (rank, rank length) + for the contiguous layout and to (slot, total/(2W)) for zigzag. + """ + # Find sequences overlapping with [start, end) + start_seq_idx = torch.searchsorted(cu_seqlens_cpu[1:], start, side='right') + end_seq_idx = torch.searchsorted(cu_seqlens_cpu[:-1], end, side='left') + + # Clamp global coordinates to [start, end] and shift to local coordinates; + # unique_consecutive removes duplicates from clamping + subset = cu_seqlens_cpu[start_seq_idx: end_seq_idx + 1] + local = (subset.clamp(min=start, max=end) - start).unique_consecutive().to(torch.int32) + + first_seq_global_start = cu_seqlens_cpu[start_seq_idx].item() + last_seq_global_end = cu_seqlens_cpu[end_seq_idx].item() + + # Tokens the interval's first sequence extends before the interval (conv tail depth) + pre_num_conv_tokens = max(0, start - first_seq_global_start) + + first_pos = first_seq_global_start // part_len + # (last_seq_global_end - 1) is the index of the last token + last_pos = (last_seq_global_end - 1) // part_len + return ( + local, + pre_num_conv_tokens, + chain_pos - first_pos, + chain_pos == first_pos, + last_pos - chain_pos, + last_pos == chain_pos, + ) + + @tensor_cache def get_cp_cu_seqlens( cu_seqlens: torch.LongTensor, @@ -64,91 +132,78 @@ def get_cp_cu_seqlens( world_size: int | None = None, rank: int | None = None, group: dist.ProcessGroup | None = None, - conv1d_kernel_size: int | None = None + conv1d_kernel_size: int | None = None, + use_tf32x3_affine_chain: bool = False, + layout: str = 'contiguous', ) -> FLACPContext: # 1. Initialize environment info if world_size is None: assert group is not None world_size = dist.get_world_size(group=group) rank = dist.get_rank(group=group) + if layout not in ('contiguous', 'zigzag'): + raise ValueError(f"`layout` must be 'contiguous' or 'zigzag', got {layout!r}.") # 2. Operate on CPU to avoid D2H sync and leverage vectorization (int64/long) if cu_seqlens_cpu is None: cu_seqlens_cpu = cu_seqlens.cpu() cu_seqlens_cpu = cu_seqlens_cpu.to(dtype=torch.long) - # Get total tokens and current rank's responsible range - # Assume cu_seqlens is [0, s1, s1+s2, ..., total] + # Get total tokens and the chain interval size. Assume cu_seqlens is [0, s1, s1+s2, ..., total] + # zigzag doubles the number of chain intervals: rank r holds intervals r and 2W-1-r total_tokens = cu_seqlens_cpu[-1].item() - part_len = total_tokens // world_size - rank_start = part_len * rank - rank_end = rank_start + part_len - - # 3. Vectorized search: find sequences overlapping with current rank's interval [rank_start, rank_end) - # We need to find idx such that: global_ends[idx] > rank_start AND global_starts[idx] < rank_end - - # Optimization: cu_seqlens is sorted, use searchsorted to quickly locate boundaries - # Find first sequence whose end > rank_start - # cu_seqlens_cpu[1:] contains all sequence end points - start_seq_idx = torch.searchsorted(cu_seqlens_cpu[1:], rank_start, side='right') - - # Find first sequence whose start >= rank_end, sequences before this may overlap - # cu_seqlens_cpu[:-1] contains all sequence start points - end_seq_idx = torch.searchsorted(cu_seqlens_cpu[:-1], rank_end, side='left') - - # Slice cu_seqlens_cpu[start_seq_idx : end_seq_idx + 1] to get relevant global cu_seqlens nodes - # +1 because end_seq_idx is an open boundary, and cu_seqlens length is num_seqs + 1 - subset_cu_seqlens = cu_seqlens_cpu[start_seq_idx: end_seq_idx + 1] - - # 4. Compute local cu_seqlens on CPU (int32) - # Clamp global coordinates to [rank_start, rank_end], subtract rank_start to get local coordinates - # unique_consecutive removes duplicates from clamping (e.g., sequences entirely outside this rank) - local_cu_seqlens_cpu = ( - subset_cu_seqlens.clamp(min=rank_start, max=rank_end) - rank_start - ).unique_consecutive().to(torch.int32) - - # Transfer to GPU (int32, small tensor, fast transfer) - # non_blocking=True can further hide latency in CUDA streams - local_cu_seqlens_gpu = local_cu_seqlens_cpu.to( - device=cu_seqlens.device, non_blocking=True - ) - - # 5. Compute Context Parallel metadata (first/last rank info) - # Use slice endpoints directly, avoiding loops - - # Get global info for the first sequence that has data on current rank - first_seq_global_start = cu_seqlens_cpu[start_seq_idx].item() - # Get global info for the last sequence that has data on current rank - last_seq_global_end = cu_seqlens_cpu[end_seq_idx].item() - - # Number of tokens current rank needs from previous ranks for conv - pre_num_conv_tokens = max(0, rank_start - first_seq_global_start) - - # Compute first sequence's starting rank - first_rank_of_first_seq = first_seq_global_start // part_len - # Number of previous ranks current rank needs to receive state from - pre_num_ranks = rank - first_rank_of_first_seq - # Whether current rank is the first in the sequence's processing chain - is_first_rank = (rank == first_rank_of_first_seq) + num_parts = world_size if layout == 'contiguous' else 2 * world_size + divisor = "`world_size`" if layout == 'contiguous' else "`2 * world_size`" + if total_tokens < num_parts: + raise ValueError( + f"`total_tokens` ({total_tokens}) must be at least {divisor} ({num_parts}) for context parallelism." + ) + if total_tokens % num_parts != 0: + raise ValueError( + f"`total_tokens` ({total_tokens}) must be divisible by {divisor} ({num_parts}) for context parallelism." + ) + part_len = total_tokens // num_parts - # Compute last sequence's ending rank - # (last_seq_global_end - 1) is the index of the last token - last_rank_of_last_seq = (last_seq_global_end - 1) // part_len - # Number of subsequent ranks current rank needs to send state to - post_num_ranks = last_rank_of_last_seq - rank - # Whether current rank is the last in the sequence's processing chain - is_last_rank = (rank == last_rank_of_last_seq) + if layout == 'contiguous': + local, pre_conv, pre_ranks, is_first, post_ranks, is_last = _interval_cp_meta( + cu_seqlens_cpu, part_len * rank, part_len * (rank + 1), rank, part_len, + ) + return FLACPContext( + group=group, + cu_seqlens=local.to(device=cu_seqlens.device, non_blocking=True), + cu_seqlens_cpu=local, + is_last_rank=is_last, + pre_num_ranks=pre_ranks, + is_first_rank=is_first, + post_num_ranks=post_ranks, + conv1d_kernel_size=conv1d_kernel_size, + pre_num_conv_tokens=pre_conv, + use_tf32x3_affine_chain=use_tf32x3_affine_chain, + ) + # zigzag: local buffer = [front part; back part], each part_len tokens + front_cu, front_conv, front_pre, front_first, front_post, front_last = _interval_cp_meta( + cu_seqlens_cpu, part_len * rank, part_len * (rank + 1), rank, part_len, + ) + back_pos = 2 * world_size - 1 - rank + back_cu, back_conv, back_pre, back_first, back_post, back_last = _interval_cp_meta( + cu_seqlens_cpu, part_len * back_pos, part_len * (back_pos + 1), back_pos, part_len, + ) + local = torch.cat([front_cu, back_cu[1:] + part_len]) return FLACPContext( group=group, - cu_seqlens=local_cu_seqlens_gpu, - cu_seqlens_cpu=local_cu_seqlens_cpu, - is_last_rank=is_last_rank, - pre_num_ranks=pre_num_ranks, - is_first_rank=is_first_rank, - post_num_ranks=post_num_ranks, + cu_seqlens=local.to(device=cu_seqlens.device, non_blocking=True), + cu_seqlens_cpu=local, conv1d_kernel_size=conv1d_kernel_size, - pre_num_conv_tokens=pre_num_conv_tokens + use_tf32x3_affine_chain=use_tf32x3_affine_chain, + layout='zigzag', + part_len=part_len, + front_num_seqs=len(front_cu) - 1, + pre_num_ranks_by_part=(front_pre, back_pre), + post_num_ranks_by_part=(front_post, back_post), + is_first_by_part=(front_first, back_first), + is_last_by_part=(front_last, back_last), + pre_num_conv_tokens_by_part=(front_conv, back_conv), ) @@ -157,6 +212,8 @@ def build_cp_context( group: ProcessGroup, conv1d_kernel_size: int | None = None, cu_seqlens_cpu: torch.Tensor | None = None, + use_tf32x3_affine_chain: bool = False, + layout: str = 'contiguous', ) -> FLACPContext: """Build a CP context for the given cu_seqlens and process group. @@ -165,8 +222,19 @@ def build_cp_context( group: Process group for CP communication. conv1d_kernel_size: Kernel size for convolution (optional). cu_seqlens_cpu: CPU version of cu_seqlens to avoid d2h transfer (optional). + use_tf32x3_affine_chain: Use tf32x3 for the affine-chain dots in the CP pre-process and merge kernels (NVIDIA only). + layout: 'contiguous' (default) assigns rank r the token interval + [r * T/W, (r + 1) * T/W); 'zigzag' assigns rank r two intervals of + T/(2W) tokens, [r * T/2W, (r + 1) * T/2W) and + [(2W - 1 - r) * T/2W, (2W - r) * T/2W), and the caller feeds the + concatenation [front; back] as the local input. Zigzag balances + causal-attention CP layouts without a re-shard at layer boundaries. Returns: FLACPContext with computed cu_seqlens and rank information. """ - return get_cp_cu_seqlens(cu_seqlens, cu_seqlens_cpu=cu_seqlens_cpu, group=group, conv1d_kernel_size=conv1d_kernel_size) + return get_cp_cu_seqlens( + cu_seqlens, cu_seqlens_cpu=cu_seqlens_cpu, group=group, + conv1d_kernel_size=conv1d_kernel_size, use_tf32x3_affine_chain=use_tf32x3_affine_chain, + layout=layout, + ) diff --git a/build/torch-cuda/ops/cyfa/__init__.py b/build/torch-cuda/ops/cyfa/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..5d44715655574c4c3f15e48b78a6e622da0d1947 --- /dev/null +++ b/build/torch-cuda/ops/cyfa/__init__.py @@ -0,0 +1,22 @@ +# 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 + + +# Portions adapted from CyclicFlowAttention, Copyright (c) 2026 Yixiao Chen. +# https://github.com/Chyxx/CyclicFlowAttention + +from __future__ import annotations + +from .chunk import chunk_cyfa +from .fused_recurrent import fused_recurrent_cyfa +from .naive import naive_recurrent_cyfa + +__all__ = [ + "chunk_cyfa", + "fused_recurrent_cyfa", + "naive_recurrent_cyfa", +] diff --git a/build/torch-cuda/ops/cyfa/chunk.py b/build/torch-cuda/ops/cyfa/chunk.py new file mode 100644 index 0000000000000000000000000000000000000000..0f83017c399327b284a8f1a1cdf11a016418e93e --- /dev/null +++ b/build/torch-cuda/ops/cyfa/chunk.py @@ -0,0 +1,1712 @@ +# 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 + + +# Portions adapted from CyclicFlowAttention, Copyright (c) 2026 Yixiao Chen. +# https://github.com/Chyxx/CyclicFlowAttention + +from __future__ import annotations + +import torch +import triton +import triton.language as tl + +from ...ops.common.chunk_h import chunk_bwd_dh, chunk_fwd_h +from ...ops.common.chunk_o import BKV_LIST, NUM_WARPS, chunk_bwd_dv, chunk_fwd_o +from ...ops.utils import prepare_chunk_indices +from ...ops.utils.cache import fla_cache_autotune +from ...ops.utils.constant import RCP_LN2 +from ...ops.utils.cumsum import chunk_global_cumsum, chunk_local_cumsum +from ...ops.utils.op import exp2 +from ...utils import ( + IS_NVIDIA_HOPPER, + autocast_custom_bwd, + autocast_custom_fwd, + autotune_cache_kwargs, + get_multiprocessor_count, + input_guard, +) + +from .qk_norm import qk_rmsnorm_fwd +from .utils import ( + CyFAState, + add_segment_initial, + build_readout_table, + cyfa_clock_cumsum, + cyfa_cos, + cyfa_sin, + modal_omega, + modal_phase, + prepare_cu_seqlens, + segment_sum, + validate_cyfa_inputs, + validate_initial_state, +) + +_CHUNK_SIZE = 64 +_READOUT_BWD_AUTOTUNE_CONFIGS = [ + triton.Config({'BT': bt}, num_warps=num_warps, num_stages=num_stages) + for bt in [16, 32, 64] + for num_warps in NUM_WARPS + for num_stages in [2, 3, 4] +] +_POINT_AUTOTUNE_CONFIGS = [ + triton.Config({'BT': bt}, num_warps=num_warps) + for bt in [16, 32, 64] + for num_warps in NUM_WARPS +] +_TABLE_GRAD_AUTOTUNE_CONFIGS = [ + triton.Config( + {'BN': bn, 'BR': br, 'BC': bc}, + num_warps=num_warps, + num_stages=num_stages, + ) + for bn in [64, 128] + for br, bc in [(32, 32), (32, 64), (64, 64)] + for num_warps in NUM_WARPS + for num_stages in [2, 3] +] +_FUSED_READOUT_FWD_AUTOTUNE_CONFIGS = [ + triton.Config( + {'BK': bk, 'BR': br}, + num_warps=num_warps, + num_stages=num_stages, + ) + for bk in [128] + for br in [16, 32] + for num_warps in NUM_WARPS + for num_stages in [2, 3, 4] +] +_RECURRENCE_BWD_AUTOTUNE_CONFIGS = [ + triton.Config( + {'BK': bk, 'BV': bv}, + num_warps=num_warps, + num_stages=num_stages, + ) + for bk in BKV_LIST + for bv in BKV_LIST + for num_warps in NUM_WARPS + for num_stages in [2, 3, 4] +] + + +@triton.jit(do_not_specialize=['N']) +def _chunk_cyfa_qk_norm_bwd_kernel( + q, + k, + dq_out, + dk_out, + q_weight, + k_weight, + q_rstd, + k_rstd, + dq, + dk, + dweight, + N, + D: tl.constexpr, + BS: tl.constexpr, + BT: tl.constexpr, + BD: tl.constexpr, +): + i_s = tl.program_id(0).to(tl.int64) + o_d = tl.arange(0, BD) + m_d = o_d < D + b_q_weight = tl.load(q_weight + o_d, mask=m_d).to(tl.float32) + b_k_weight = tl.load(k_weight + o_d, mask=m_d).to(tl.float32) + b_dq_weight = tl.zeros((BT, BD), dtype=tl.float32) + b_dk_weight = tl.zeros((BT, BD), dtype=tl.float32) + + for i_n in range(i_s * BS, i_s * BS + BS, BT): + o_n = i_n + tl.arange(0, BT) + m_n = o_n < N + m_acc = o_n < min((i_s + 1) * BS, N) + mask = m_n[:, None] & m_d[None, :] + acc_mask = m_acc[:, None] & m_d[None, :] + o_n_local = tl.arange(0, BT) + offsets = o_n_local[:, None] * D + o_d[None, :] + + b_q = tl.load((q + i_n * D) + offsets, mask=mask, other=0.0).to(tl.float32) + b_dq_out = tl.load((dq_out + i_n * D) + offsets, mask=mask, other=0.0).to(tl.float32) + b_q_rstd = tl.load((q_rstd + i_n) + o_n_local, mask=m_n, other=0.0) + b_q_hat = tl.where(m_d[None, :], b_q * b_q_rstd[:, None], 0.0) + b_q_wdy = b_dq_out * b_q_weight[None, :] + b_q_proj = tl.sum(b_q_hat * b_q_wdy, axis=1) / D + b_dq = (b_q_wdy - b_q_hat * b_q_proj[:, None]) * b_q_rstd[:, None] + tl.store((dq + i_n * D) + offsets, b_dq.to(dq.dtype.element_ty), mask=mask) + b_dq_weight += tl.where(acc_mask, b_dq_out * b_q_hat, 0.0) + + b_k = tl.load((k + i_n * D) + offsets, mask=mask, other=0.0).to(tl.float32) + b_dk_out = tl.load((dk_out + i_n * D) + offsets, mask=mask, other=0.0).to(tl.float32) + b_k_rstd = tl.load((k_rstd + i_n) + o_n_local, mask=m_n, other=0.0) + b_k_hat = tl.where(m_d[None, :], b_k * b_k_rstd[:, None], 0.0) + b_k_wdy = b_dk_out * b_k_weight[None, :] + b_k_proj = tl.sum(b_k_hat * b_k_wdy, axis=1) / D + b_dk = (b_k_wdy - b_k_hat * b_k_proj[:, None]) * b_k_rstd[:, None] + tl.store((dk + i_n * D) + offsets, b_dk.to(dk.dtype.element_ty), mask=mask) + b_dk_weight += tl.where(acc_mask, b_dk_out * b_k_hat, 0.0) + + partial_offset = i_s * D + o_d + tl.store(dweight + partial_offset, tl.sum(b_dq_weight, axis=0), mask=m_d) + tl.store(dweight + tl.num_programs(0).to(tl.int64) * D + partial_offset, tl.sum(b_dk_weight, axis=0), mask=m_d) + + +def _qk_rmsnorm_bwd( + dq_out: torch.Tensor, + dk_out: torch.Tensor, + q: torch.Tensor, + k: torch.Tensor, + q_weight: torch.Tensor, + k_weight: torch.Tensor, + q_rstd: torch.Tensor, + k_rstd: torch.Tensor, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + if q.shape != k.shape or dq_out.shape != q.shape or dk_out.shape != k.shape: + raise ValueError('Fused Q/K RMSNorm backward requires matching tensor shapes.') + q = q.contiguous() + k = k.contiguous() + dq_out = dq_out.contiguous() + dk_out = dk_out.contiguous() + n_rows = q.numel() // q.shape[-1] + d_head = q.shape[-1] + n_splits = min(get_multiprocessor_count(q.device.index), n_rows) + block_rows = triton.cdiv(n_rows, n_splits) + dq = torch.empty_like(q) + dk = torch.empty_like(k) + dweight = torch.empty((2, n_splits, d_head), device=q.device, dtype=torch.float32) + + _chunk_cyfa_qk_norm_bwd_kernel[(n_splits,)]( + q, + k, + dq_out, + dk_out, + q_weight, + k_weight, + q_rstd, + k_rstd, + dq, + dk, + dweight, + N=n_rows, + D=d_head, + BS=block_rows, + BT=32, + BD=triton.next_power_of_2(d_head), + num_warps=8, + num_stages=3, + ) + dq_weight, dk_weight = dweight.sum(dim=1).to(q_weight).unbind(0) + return dq, dk, dq_weight, dk_weight + + +@triton.heuristics({'IS_VARLEN': lambda args: args['cu_seqlens'] is not None}) +@fla_cache_autotune( + configs=_RECURRENCE_BWD_AUTOTUNE_CONFIGS, + key=['H', 'K', 'V', 'BT'], + reset_to_zero=['dg_force'], + **autotune_cache_kwargs, +) +@triton.jit(do_not_specialize=['T']) +def _chunk_cyfa_bwd_dqk_kernel( + q, + k, + v, + g, + h, + do, + dh, + dq, + dk, + dg_force, + cu_seqlens, + chunk_indices, + scale, + T, + H: tl.constexpr, + K: tl.constexpr, + V: tl.constexpr, + BT: tl.constexpr, + BK: tl.constexpr, + BV: tl.constexpr, + IS_VARLEN: tl.constexpr, +): + i_k = tl.program_id(0).to(tl.int64) + i_t = tl.program_id(1).to(tl.int64) + i_bh = tl.program_id(2).to(tl.int64) + i_b = i_bh // H + i_h = i_bh % H + + if IS_VARLEN: + i_tg = i_t + i_n = tl.load(chunk_indices + i_t * 2).to(tl.int64) + i_t = tl.load(chunk_indices + i_t * 2 + 1).to(tl.int64) + bos = tl.load(cu_seqlens + i_n).to(tl.int64) + eos = tl.load(cu_seqlens + i_n + 1).to(tl.int64) + T = eos - bos + else: + nt = tl.cdiv(T, BT) + i_tg = i_b * nt + i_t + bos = i_b * T + + v += (bos * H + i_h) * V + do += (bos * H + i_h) * V + h += (i_tg * H + i_h).to(tl.int64) * K * V + dh += (i_tg * H + i_h).to(tl.int64) * K * V + q += (bos * H + i_h) * K + k += (bos * H + i_h) * K + dq += (bos * H + i_h) * K + dk += (bos * H + i_h) * K + + b_dq = tl.zeros((BT, BK), dtype=tl.float32) + b_dk = tl.zeros((BT, BK), dtype=tl.float32) + b_ds = tl.zeros((BT, BT), dtype=tl.float32) + # g_last transports recurrent state into the next chunk; inject its boundary adjoint at the last cumsum coordinate. + b_dg_last = tl.zeros((1,), dtype=tl.float32) + for i_v in range(0, tl.cdiv(V, BV)): + o_v_0 = i_t * BT + tl.arange(0, BT) + o_v_1 = i_v * BV + tl.arange(0, BV) + p_v = v + o_v_0[:, None] * (H * V) + o_v_1[None, :] * (1) + m_v = (o_v_0[:, None] < T) & (o_v_1[None, :] < V) + o_do_0 = i_t * BT + tl.arange(0, BT) + o_do_1 = i_v * BV + tl.arange(0, BV) + p_do = do + o_do_0[:, None] * (H * V) + o_do_1[None, :] * (1) + m_do = (o_do_0[:, None] < T) & (o_do_1[None, :] < V) + o_h_0 = i_v * BV + tl.arange(0, BV) + o_h_1 = i_k * BK + tl.arange(0, BK) + p_h = h + o_h_0[:, None] * (1) + o_h_1[None, :] * (V) + m_h = (o_h_0[:, None] < V) & (o_h_1[None, :] < K) + o_dh_0 = i_v * BV + tl.arange(0, BV) + o_dh_1 = i_k * BK + tl.arange(0, BK) + p_dh = dh + o_dh_0[:, None] * (1) + o_dh_1[None, :] * (V) + m_dh = (o_dh_0[:, None] < V) & (o_dh_1[None, :] < K) + b_v = tl.load(p_v, mask=m_v, other=0.0) + b_do = tl.load(p_do, mask=m_do, other=0.0) + b_h = tl.load(p_h, mask=m_h, other=0.0) + b_dh = tl.load(p_dh, mask=m_dh, other=0.0) + b_dg_last += tl.sum(b_h * b_dh) + b_ds += tl.dot(b_do, tl.trans(b_v)) + b_dq += tl.dot(b_do, b_h.to(b_do.dtype)) + b_dk += tl.dot(b_v, b_dh.to(b_v.dtype)) + + o_q_0 = i_t * BT + tl.arange(0, BT) + o_q_1 = i_k * BK + tl.arange(0, BK) + p_q = q + o_q_0[:, None] * (H * K) + o_q_1[None, :] * (1) + m_q = (o_q_0[:, None] < T) & (o_q_1[None, :] < K) + o_k_0 = i_t * BT + tl.arange(0, BT) + o_k_1 = i_k * BK + tl.arange(0, BK) + p_k = k + o_k_0[:, None] * (H * K) + o_k_1[None, :] * (1) + m_k = (o_k_0[:, None] < T) & (o_k_1[None, :] < K) + b_q = tl.load(p_q, mask=m_q, other=0.0) + b_k = tl.load(p_k, mask=m_k, other=0.0) + + o_t = i_t * BT + tl.arange(0, BT) + mask_t = o_t < T + mask_a = (o_t[:, None] >= o_t[None, :]) & mask_t[:, None] & mask_t[None, :] + o_g_0 = i_t * BT + tl.arange(0, BT) + p_g = g + bos * H + i_h + o_g_0 * (H) + m_g = o_g_0 < T + b_g = tl.load(p_g, mask=m_g, other=0.0) + b_g_last = tl.load(g + (bos + min(i_t * BT + BT, T) - 1) * H + i_h) + + b_dg_last *= exp2(b_g_last) + b_dq *= exp2(b_g)[:, None] * scale + b_dk *= tl.where(mask_t, exp2(-b_g + b_g_last), 0.0)[:, None] + b_dg_last += tl.sum(b_dk * b_k) + b_ds = tl.where(mask_a, b_ds * exp2(b_g[:, None] - b_g[None, :]), 0.0) * scale + b_ds = b_ds.to(b_k.dtype) + b_dq += tl.dot(b_ds, b_k) + b_dk += tl.dot(tl.trans(b_ds), b_q) + + force = tl.sum(b_dq * b_q, axis=1) - tl.sum(b_dk * b_k, axis=1) + last = min(i_t * BT + BT, T) - 1 + force += tl.where(o_t == last, b_dg_last, 0.0) + tl.atomic_add( + dg_force + (bos + o_t) * H + i_h, + force, + mask=mask_t, + sem='relaxed', + ) + + o_dq_0 = i_t * BT + tl.arange(0, BT) + o_dq_1 = i_k * BK + tl.arange(0, BK) + p_dq = dq + o_dq_0[:, None] * (H * K) + o_dq_1[None, :] * (1) + m_dq = (o_dq_0[:, None] < T) & (o_dq_1[None, :] < K) + o_dk_0 = i_t * BT + tl.arange(0, BT) + o_dk_1 = i_k * BK + tl.arange(0, BK) + p_dk = dk + o_dk_0[:, None] * (H * K) + o_dk_1[None, :] * (1) + m_dk = (o_dk_0[:, None] < T) & (o_dk_1[None, :] < K) + tl.store(p_dq, b_dq.to(p_dq.dtype.element_ty), mask=m_dq) + tl.store(p_dk, b_dk.to(p_dk.dtype.element_ty), mask=m_dk) + + +def _scalar_decay_bwd_dqk( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + g: torch.Tensor, + h: torch.Tensor, + do: torch.Tensor, + dh: torch.Tensor, + *, + scale: float, + cu_seqlens: torch.Tensor | None, + chunk_size: int, + chunk_indices: torch.Tensor | None, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + bsz, seq_len, n_heads, d_k = q.shape + d_v = v.shape[-1] + nt = triton.cdiv(seq_len, chunk_size) if chunk_indices is None else len(chunk_indices) + dq = torch.empty_like(q) + dk = torch.empty_like(k) + dg_force = torch.zeros(bsz, seq_len, n_heads, device=q.device, dtype=torch.float32) + + def grid(meta): + return (triton.cdiv(d_k, meta['BK']), nt, bsz * n_heads) + + _chunk_cyfa_bwd_dqk_kernel[grid]( + q, + k, + v, + g, + h, + do, + dh, + dq, + dk, + dg_force, + cu_seqlens, + chunk_indices, + scale, + T=seq_len, + H=n_heads, + K=d_k, + V=d_v, + BT=chunk_size, + ) + return dq, dk, dg_force + + +@triton.heuristics({'IS_VARLEN': lambda args: args['cu_seqlens'] is not None}) +@triton.jit(do_not_specialize=['T']) +def _chunk_cyfa_bwd_dqk_dv_kernel( + q, + k, + v, + g, + h, + do, + dh, + dq, + dk, + dv, + dg_force, + cu_seqlens, + chunk_indices, + scale, + T, + H: tl.constexpr, + K: tl.constexpr, + V: tl.constexpr, + BT: tl.constexpr, + BK: tl.constexpr, + BV: tl.constexpr, + IS_VARLEN: tl.constexpr, +): + i_t = tl.program_id(0).to(tl.int64) + i_bh = tl.program_id(1).to(tl.int64) + i_b = i_bh // H + i_h = i_bh % H + + if IS_VARLEN: + i_tg = i_t + i_n = tl.load(chunk_indices + i_t * 2).to(tl.int64) + i_t = tl.load(chunk_indices + i_t * 2 + 1).to(tl.int64) + bos = tl.load(cu_seqlens + i_n).to(tl.int64) + eos = tl.load(cu_seqlens + i_n + 1).to(tl.int64) + T = eos - bos + else: + nt = tl.cdiv(T, BT) + i_tg = i_b * nt + i_t + bos = i_b * T + + v += (bos * H + i_h) * V + do += (bos * H + i_h) * V + dv += (bos * H + i_h) * V + h += (i_tg * H + i_h).to(tl.int64) * K * V + dh += (i_tg * H + i_h).to(tl.int64) * K * V + q += (bos * H + i_h) * K + k += (bos * H + i_h) * K + dq += (bos * H + i_h) * K + dk += (bos * H + i_h) * K + + o_t = i_t * BT + tl.arange(0, BT) + mask_t = o_t < T + + o_q_0 = i_t * BT + tl.arange(0, BT) + o_q_1 = 0 + tl.arange(0, BK) + p_q = q + o_q_0[:, None] * (H * K) + o_q_1[None, :] * (1) + m_q = (o_q_0[:, None] < T) & (o_q_1[None, :] < K) + o_k_0 = i_t * BT + tl.arange(0, BT) + o_k_1 = 0 + tl.arange(0, BK) + p_k = k + o_k_0[:, None] * (H * K) + o_k_1[None, :] * (1) + m_k = (o_k_0[:, None] < T) & (o_k_1[None, :] < K) + b_q = tl.load(p_q, mask=m_q, other=0.0) + b_k = tl.load(p_k, mask=m_k, other=0.0) + + o_g_0 = i_t * BT + tl.arange(0, BT) + p_g = g + bos * H + i_h + o_g_0 * (H) + m_g = o_g_0 < T + b_g = tl.load(p_g, mask=m_g, other=0.0) + b_g_last = tl.load(g + (bos + min(i_t * BT + BT, T) - 1) * H + i_h) + + mask_fwd = (o_t[:, None] >= o_t[None, :]) & mask_t[:, None] & mask_t[None, :] + mask_bwd = (o_t[:, None] <= o_t[None, :]) & mask_t[:, None] & mask_t[None, :] + b_A_dv = tl.dot(b_k, tl.trans(b_q)) + b_A_dv = tl.where(mask_bwd, b_A_dv * exp2(b_g[None, :] - b_g[:, None]) * scale, 0.0) + b_A_dv = b_A_dv.to(do.dtype.element_ty) + b_dv_decay = tl.where(mask_t, exp2(-b_g + b_g_last), 0.0) + + b_dq = tl.zeros((BT, BK), dtype=tl.float32) + b_dk = tl.zeros((BT, BK), dtype=tl.float32) + b_ds = tl.zeros((BT, BT), dtype=tl.float32) + b_dg_last = tl.zeros((1,), dtype=tl.float32) + for i_v in range(0, tl.cdiv(V, BV)): + o_v_0 = i_t * BT + tl.arange(0, BT) + o_v_1 = i_v * BV + tl.arange(0, BV) + p_v = v + o_v_0[:, None] * (H * V) + o_v_1[None, :] * (1) + m_v = (o_v_0[:, None] < T) & (o_v_1[None, :] < V) + o_do_0 = i_t * BT + tl.arange(0, BT) + o_do_1 = i_v * BV + tl.arange(0, BV) + p_do = do + o_do_0[:, None] * (H * V) + o_do_1[None, :] * (1) + m_do = (o_do_0[:, None] < T) & (o_do_1[None, :] < V) + o_dv_0 = i_t * BT + tl.arange(0, BT) + o_dv_1 = i_v * BV + tl.arange(0, BV) + p_dv = dv + o_dv_0[:, None] * (H * V) + o_dv_1[None, :] * (1) + m_dv = (o_dv_0[:, None] < T) & (o_dv_1[None, :] < V) + o_h_0 = i_v * BV + tl.arange(0, BV) + o_h_1 = 0 + tl.arange(0, BK) + p_h = h + o_h_0[:, None] * (1) + o_h_1[None, :] * (V) + m_h = (o_h_0[:, None] < V) & (o_h_1[None, :] < K) + o_dh_0 = i_v * BV + tl.arange(0, BV) + o_dh_1 = 0 + tl.arange(0, BK) + p_dh = dh + o_dh_0[:, None] * (1) + o_dh_1[None, :] * (V) + m_dh = (o_dh_0[:, None] < V) & (o_dh_1[None, :] < K) + b_v = tl.load(p_v, mask=m_v, other=0.0) + b_do = tl.load(p_do, mask=m_do, other=0.0) + b_h = tl.load(p_h, mask=m_h, other=0.0) + b_dh = tl.load(p_dh, mask=m_dh, other=0.0) + b_dg_last += tl.sum(b_h * b_dh) + b_ds += tl.dot(b_do, tl.trans(b_v)) + b_dq += tl.dot(b_do, b_h.to(b_do.dtype)) + b_dk += tl.dot(b_v, b_dh.to(b_v.dtype)) + + b_dv = tl.dot(b_k, tl.trans(b_dh).to(b_k.dtype)) + b_dv *= b_dv_decay[:, None] + b_dv += tl.dot(b_A_dv, b_do) + tl.store(p_dv, b_dv.to(p_dv.dtype.element_ty), mask=m_dv) + + b_dg_last *= exp2(b_g_last) + b_dq *= exp2(b_g)[:, None] * scale + b_dk *= b_dv_decay[:, None] + b_dg_last += tl.sum(b_dk * b_k) + b_ds = tl.where(mask_fwd, b_ds * exp2(b_g[:, None] - b_g[None, :]), 0.0) * scale + b_ds = b_ds.to(b_k.dtype) + if q.dtype.element_ty == tl.float32: + b_dq += tl.trans(tl.dot(tl.trans(b_k), tl.trans(b_ds))) + b_dk += tl.trans(tl.dot(tl.trans(b_q), b_ds)) + else: + b_dq += tl.dot(b_ds, b_k) + b_dk += tl.dot(tl.trans(b_ds), b_q) + + force = tl.sum(b_dq * b_q, axis=1) - tl.sum(b_dk * b_k, axis=1) + last = min(i_t * BT + BT, T) - 1 + force += tl.where(o_t == last, b_dg_last, 0.0) + tl.store(dg_force + (bos + o_t) * H + i_h, force, mask=mask_t) + + o_dq_0 = i_t * BT + tl.arange(0, BT) + o_dq_1 = 0 + tl.arange(0, BK) + p_dq = dq + o_dq_0[:, None] * (H * K) + o_dq_1[None, :] * (1) + m_dq = (o_dq_0[:, None] < T) & (o_dq_1[None, :] < K) + o_dk_0 = i_t * BT + tl.arange(0, BT) + o_dk_1 = 0 + tl.arange(0, BK) + p_dk = dk + o_dk_0[:, None] * (H * K) + o_dk_1[None, :] * (1) + m_dk = (o_dk_0[:, None] < T) & (o_dk_1[None, :] < K) + tl.store(p_dq, b_dq.to(p_dq.dtype.element_ty), mask=m_dq) + tl.store(p_dk, b_dk.to(p_dk.dtype.element_ty), mask=m_dk) + + +def _scalar_decay_bwd_dqk_dv_single_k( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + g: torch.Tensor, + h: torch.Tensor, + do: torch.Tensor, + dh: torch.Tensor, + *, + scale: float, + cu_seqlens: torch.Tensor | None, + chunk_size: int, + chunk_indices: torch.Tensor | None, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + bsz, seq_len, n_heads, d_k = q.shape + d_v = v.shape[-1] + if d_k > 128: + raise ValueError("single-k fused dqk/dv path requires K <= 128.") + nt = triton.cdiv(seq_len, chunk_size) if chunk_indices is None else len(chunk_indices) + dq = torch.empty_like(q) + dk = torch.empty_like(k) + dv = torch.empty_like(do) + dg_force = torch.empty(bsz, seq_len, n_heads, device=q.device, dtype=torch.float32) + block_k = max(triton.next_power_of_2(d_k), 16) + is_hopper_fp32 = IS_NVIDIA_HOPPER and q.dtype == torch.float32 + # Keep the fused FP32 kernel below the SM90 per-block shared-memory limit. + block_v = min(max(triton.next_power_of_2(d_v), 16), 32 if is_hopper_fp32 else 64) + + _chunk_cyfa_bwd_dqk_dv_kernel[(nt, bsz * n_heads)]( + q, + k, + v, + g, + h, + do, + dh, + dq, + dk, + dv, + dg_force, + cu_seqlens, + chunk_indices, + scale, + T=seq_len, + H=n_heads, + K=d_k, + V=d_v, + BT=chunk_size, + BK=block_k, + BV=block_v, + num_warps=4, + num_stages=1 if is_hopper_fp32 else 2, + ) + return dq, dk, dv, dg_force + + +def _scalar_decay_fwd( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + g_cumsum: torch.Tensor, + scale: float, + initial_state: torch.Tensor | None, + output_final_state: bool, + cu_seqlens: torch.Tensor | None, + chunk_size: int, + chunk_indices: torch.Tensor | None, +) -> tuple[torch.Tensor, torch.Tensor | None]: + h, final_state = chunk_fwd_h( + k=k, + v=v, + g=g_cumsum, + h0=None if initial_state is None else initial_state.transpose(-1, -2).contiguous(), + output_final_state=output_final_state, + cu_seqlens=cu_seqlens, + chunk_size=chunk_size, + states_in_fp32=False, + state_v_first=True, + ) + out = chunk_fwd_o( + q=q, + k=k, + v=v, + h=h, + g=g_cumsum, + scale=scale, + cu_seqlens=cu_seqlens, + chunk_size=chunk_size, + chunk_indices=chunk_indices, + state_v_first=True, + ) + if final_state is not None: + final_state = final_state.transpose(-1, -2).contiguous() + return out, final_state + + +def _scalar_decay_bwd( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + g_cumsum: torch.Tensor, + initial_state: torch.Tensor | None, + do: torch.Tensor, + dht: torch.Tensor | None, + scale: float, + cu_seqlens: torch.Tensor | None, + chunk_size: int, + chunk_indices: torch.Tensor | None, + fuse_dv_single_k: bool = False, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor | None]: + h, _ = chunk_fwd_h( + k=k, + v=v, + g=g_cumsum, + h0=initial_state, + output_final_state=False, + cu_seqlens=cu_seqlens, + chunk_size=chunk_size, + states_in_fp32=False, + ) + dh, dh0 = chunk_bwd_dh( + q=q, + k=k, + v=v, + g=g_cumsum, + do=do, + h0=initial_state, + dht=dht, + scale=scale, + cu_seqlens=cu_seqlens, + chunk_size=chunk_size, + states_in_fp32=False, + ) + if fuse_dv_single_k and q.shape[-1] <= 128: + dq, dk, dv, dg_force = _scalar_decay_bwd_dqk_dv_single_k( + q, + k, + v, + g_cumsum, + h, + do, + dh, + scale=scale, + cu_seqlens=cu_seqlens, + chunk_size=chunk_size, + chunk_indices=chunk_indices, + ) + else: + dq, dk, dg_force = _scalar_decay_bwd_dqk( + q, + k, + v, + g_cumsum, + h, + do, + dh, + scale=scale, + cu_seqlens=cu_seqlens, + chunk_size=chunk_size, + chunk_indices=chunk_indices, + ) + dv = chunk_bwd_dv( + q=q, + k=k, + g=g_cumsum, + g_gamma=None, + do=do, + dh=dh, + scale=scale, + cu_seqlens=cu_seqlens, + chunk_size=chunk_size, + chunk_indices=chunk_indices, + ) + return dq, dk, dv, dg_force, dh0 + + +@triton.jit +def _cardinal_orbit_value(lam, o_m, M: tl.constexpr): + is_cos = (o_m % 2) == 0 + phase = modal_phase(lam, o_m, M) + scale = tl.sqrt(tl.where((o_m // 2) == 0, 1.0, 2.0) / (M - 1)) + return scale * tl.where(is_cos, cyfa_cos(phase), -cyfa_sin(phase)) + + +@triton.jit +def _cardinal_orbit_derivative(lam, o_m, M: tl.constexpr): + is_cos = (o_m % 2) == 0 + omega = modal_omega(o_m, M) + phase = modal_phase(lam, o_m, M) + scale = tl.sqrt(tl.where((o_m // 2) == 0, 1.0, 2.0) / (M - 1)) + return -scale * omega * tl.where(is_cos, cyfa_sin(phase), cyfa_cos(phase)) + + +@triton.autotune( + configs=_POINT_AUTOTUNE_CONFIGS, + key=['M'], + **autotune_cache_kwargs, +) +@triton.jit(do_not_specialize=['T']) +def _chunk_cyfa_write_fwd_kernel( + lambdas, + beta, + write, + T, + H: tl.constexpr, + M: tl.constexpr, + BT: tl.constexpr, + BM: tl.constexpr, +): + i_t = tl.program_id(0).to(tl.int64) + i_bh = tl.program_id(1).to(tl.int64) + i_b = i_bh // H + i_h = i_bh % H + o_t = i_t * BT + tl.arange(0, BT) + o_m = tl.arange(0, BM) + m_t = o_t < T + m_m = o_m < M + + b_lambda = tl.load(lambdas + (i_b * T + o_t) * H + i_h, mask=m_t, other=0.0).to(tl.float32) + b_beta = tl.load(beta + (i_b * T + o_t) * H + i_h, mask=m_t, other=0.0).to(tl.float32) + write_value = b_beta[:, None] * _cardinal_orbit_value( + b_lambda[:, None], + o_m[None, :], + M, + ) + mask = m_t[:, None] & m_m[None, :] + p = ((i_b * T + o_t[:, None]) * H + i_h) * M + o_m[None, :] + tl.store(write + p, write_value, mask=mask) + + +def _point_write( + lambdas: torch.Tensor, + beta: torch.Tensor, + num_slots: int, + dtype: torch.dtype, +) -> torch.Tensor: + lambdas = lambdas.contiguous() + beta = beta.contiguous() + bsz, seq_len, n_heads = lambdas.shape + write = torch.empty(bsz, seq_len, n_heads, num_slots, device=lambdas.device, dtype=dtype) + _chunk_cyfa_write_fwd_kernel[lambda meta: (triton.cdiv(seq_len, meta['BT']), bsz * n_heads)]( + lambdas, + beta, + write, + T=seq_len, + H=n_heads, + M=num_slots, + BM=triton.next_power_of_2(num_slots), + ) + return write + + +@triton.autotune( + configs=_POINT_AUTOTUNE_CONFIGS, + key=['M'], + **autotune_cache_kwargs, +) +@triton.jit(do_not_specialize=['T']) +def _chunk_cyfa_write_bwd_kernel( + dwrite_a, + dwrite_b, + lambdas, + beta, + dlambda, + dbeta, + T, + H: tl.constexpr, + M: tl.constexpr, + BT: tl.constexpr, + BM: tl.constexpr, +): + i_t = tl.program_id(0).to(tl.int64) + i_bh = tl.program_id(1).to(tl.int64) + i_b = i_bh // H + i_h = i_bh % H + o_t = i_t * BT + tl.arange(0, BT) + o_m = tl.arange(0, BM) + m_t = o_t < T + mask = m_t[:, None] & (o_m[None, :] < M) + p = ((i_b * T + o_t[:, None]) * H + i_h) * M + o_m[None, :] + dw = tl.load(dwrite_a + p, mask=mask, other=0.0).to(tl.float32) + dw += tl.load(dwrite_b + p, mask=mask, other=0.0).to(tl.float32) + b_lambda = tl.load(lambdas + (i_b * T + o_t) * H + i_h, mask=m_t, other=0.0).to(tl.float32) + b_beta = tl.load(beta + (i_b * T + o_t) * H + i_h, mask=m_t, other=0.0).to(tl.float32) + endpoint = _cardinal_orbit_value(b_lambda[:, None], o_m[None, :], M) + tangent = _cardinal_orbit_derivative(b_lambda[:, None], o_m[None, :], M) + d_phase = tl.sum(tl.where(mask, dw * b_beta[:, None] * tangent, 0.0), axis=1) + d_write = tl.sum(tl.where(mask, dw * endpoint, 0.0), axis=1) + offset = (i_b * T + o_t) * H + i_h + tl.store(dlambda + offset, d_phase, mask=m_t) + tl.store(dbeta + offset, d_write, mask=m_t) + + +def _point_write_backward( + dwrite_a: torch.Tensor, + dwrite_b: torch.Tensor, + lambdas: torch.Tensor, + beta: torch.Tensor, +) -> tuple[torch.Tensor, torch.Tensor]: + bsz, seq_len, n_heads, num_slots = dwrite_a.shape + dlambda = torch.empty(bsz, seq_len, n_heads, device=dwrite_a.device, dtype=torch.float32) + dbeta = torch.empty_like(dlambda) + _chunk_cyfa_write_bwd_kernel[lambda meta: (triton.cdiv(seq_len, meta['BT']), bsz * n_heads)]( + dwrite_a.contiguous(), + dwrite_b.contiguous(), + lambdas.contiguous(), + beta.contiguous(), + dlambda, + dbeta, + T=seq_len, + H=n_heads, + M=num_slots, + BM=triton.next_power_of_2(num_slots), + ) + return dlambda, dbeta + + +@triton.heuristics({ + 'USE_G': lambda args: args['g'] is not None, + 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None, +}) +@fla_cache_autotune( + configs=_FUSED_READOUT_FWD_AUTOTUNE_CONFIGS, + key=['H', 'K', 'M', 'C', 'BT'], + **autotune_cache_kwargs, +) +@triton.jit(do_not_specialize=['T']) +def _chunk_cyfa_fused_readout_fwd_kernel( + q, + q_weight, + q_rstd, + k, + write, + h, + g, + lambdas, + table, + weights, + raw_output, + cu_seqlens, + chunk_indices, + scale, + T, + H: tl.constexpr, + K: tl.constexpr, + M: tl.constexpr, + C: tl.constexpr, + BT: tl.constexpr, + BK: tl.constexpr, + BR: tl.constexpr, + STORE_RAW: tl.constexpr, + USE_G: tl.constexpr, + IS_VARLEN: tl.constexpr, +): + i_r = tl.program_id(0).to(tl.int64) + i_t = tl.program_id(1).to(tl.int64) + i_bh = tl.program_id(2).to(tl.int64) + i_b = i_bh // H + i_h = i_bh % H + + if IS_VARLEN: + i_tg = i_t + i_n = tl.load(chunk_indices + i_t * 2).to(tl.int64) + i_t = tl.load(chunk_indices + i_t * 2 + 1).to(tl.int64) + bos = tl.load(cu_seqlens + i_n).to(tl.int64) + eos = tl.load(cu_seqlens + i_n + 1).to(tl.int64) + T = eos - bos + else: + n_chunks = tl.cdiv(T, BT) + i_tg = i_b * n_chunks + i_t + bos = i_b * T + + q += (bos * H + i_h) * K + k += (bos * H + i_h) * K + write += (bos * H + i_h) * M + h += (i_tg * H + i_h).to(tl.int64) * M * K + + # the row offset is bounded by BT; global time offsets remain int64 + o_t = i_t * BT + (i_r * BR).to(tl.int32) + tl.arange(0, BR) + o_s = i_t * BT + tl.arange(0, BT) + o_m = tl.arange(0, M) + m_t = o_t < T + m_s = o_s < T + b_x = tl.zeros((BR, M), dtype=tl.float32) + + if USE_G: + p_g_t = g + (bos + o_t) * H + i_h + p_g_s = g + (bos + o_s) * H + i_h + b_g_t = tl.load(p_g_t, mask=m_t, other=0.0) + b_g_s = tl.load(p_g_s, mask=m_s, other=0.0) + m_A = (o_t[:, None] >= o_s[None, :]) & m_t[:, None] & m_s[None, :] + p_write = write + o_s[:, None] * (H * M) + o_m[None, :] + b_write = tl.load(p_write, mask=m_s[:, None], other=0.0) + + for i_k in range(tl.cdiv(K, BK)): + o_k = i_k * BK + tl.arange(0, BK) + m_k = o_k < K + p_q = q + o_t[:, None] * (H * K) + o_k[None, :] + p_k = k + o_k[:, None] + o_s[None, :] * (H * K) + p_h = h + o_m[:, None] * K + o_k[None, :] + b_q = tl.load(p_q, mask=m_t[:, None] & m_k[None, :], other=0.0).to(tl.float32) + b_q_scale = tl.load(q_rstd + (bos + o_t) * H + i_h, mask=m_t, other=0.0) + b_q_weight = tl.load(q_weight + o_k, mask=m_k, other=0.0).to(tl.float32) + b_q = (b_q * b_q_scale[:, None] * b_q_weight[None, :]).to(q.dtype.element_ty) + b_k = tl.load(p_k, mask=m_k[:, None] & m_s[None, :], other=0.0) + b_h = tl.load(p_h, mask=m_k[None, :], other=0.0) + b_x_partial = tl.dot(b_q, tl.trans(b_h)) + b_A_partial = tl.dot(b_q, b_k) + if USE_G: + b_x_partial *= exp2(b_g_t)[:, None] + b_A_partial *= exp2(b_g_t[:, None] - b_g_s[None, :]) + b_A_partial = tl.where(m_A, b_A_partial, 0.0) + b_x_partial = (b_x_partial + tl.dot(b_A_partial.to(b_write.dtype), b_write)) * scale + b_x += b_x_partial.to(q.dtype.element_ty).to(tl.float32) + + p_x = ((bos + o_t[:, None]) * H + i_h) * M + o_m[None, :] + if STORE_RAW: + tl.store(raw_output + p_x, b_x, mask=m_t[:, None]) + + o_even = 2 * tl.arange(0, M // 2) + x_even, x_odd = tl.split(tl.reshape(b_x, (BR, M // 2, 2))) + b_lambda = tl.load(lambdas + (bos + o_t) * H + i_h, mask=m_t, other=0.0).to(tl.float32) + phase = modal_phase(b_lambda[:, None], o_even[None, :], M) + c = cyfa_cos(phase) + s = cyfa_sin(phase) + z = tl.reshape(tl.join(c * x_even - s * x_odd, s * x_even + c * x_odd), (BR, M)) + z = tl.where(m_t[:, None], z, 0.0) + + o_r = tl.arange(0, M) + m_r = o_r < C + w = tl.load( + table + i_h * C * M + o_r[:, None] * M + o_m[None, :], + mask=m_r[:, None], + other=0.0, + ).to(tl.float32) + logits = tl.dot(z.to(q.dtype.element_ty), tl.trans(w.to(q.dtype.element_ty))) + logits = tl.where(m_r[None, :], logits, -float("inf")) + max_logits = tl.max(logits, axis=1) + probs = tl.where(m_r[None, :], tl.exp(logits - max_logits[:, None]), 0.0) + denom = tl.sum(probs, axis=1) + probs *= tl.where(denom > 0.0, 1.0 / denom, 0.0)[:, None] + + u = tl.dot(probs.to(weights.dtype.element_ty), w.to(weights.dtype.element_ty)) + u_even, u_odd = tl.split(tl.reshape(u, (BR, M // 2, 2))) + y = tl.reshape(tl.join(c * u_even + s * u_odd, -s * u_even + c * u_odd), (BR, M)) + tl.store(weights + p_x, y, mask=m_t[:, None]) + + +def _scalar_decay_fused_readout_fwd( + q: torch.Tensor, + q_norm_weight: torch.Tensor, + q_rstd: torch.Tensor, + k: torch.Tensor, + write: torch.Tensor, + g_cumsum: torch.Tensor, + lambdas: torch.Tensor, + readout_table: torch.Tensor, + scale: float, + initial_state: torch.Tensor | None, + output_final_state: bool, + cu_seqlens: torch.Tensor | None, + chunk_size: int, + chunk_indices: torch.Tensor | None, + *, + save_o_prime: bool, +) -> tuple[torch.Tensor, torch.Tensor | None, torch.Tensor | None]: + bsz, seq_len, n_heads, d_k = q.shape + num_slots = write.shape[-1] + readout_slots = readout_table.shape[-2] + n_chunks = triton.cdiv(seq_len, chunk_size) if chunk_indices is None else len(chunk_indices) + h, final_state = chunk_fwd_h( + k=k, + v=write, + g=g_cumsum, + h0=None if initial_state is None else initial_state.transpose(-1, -2).contiguous(), + output_final_state=output_final_state, + cu_seqlens=cu_seqlens, + chunk_size=chunk_size, + states_in_fp32=False, + state_v_first=True, + ) + weights = torch.empty_like(write) + raw_output = torch.empty_like(weights) if save_o_prime else None + + def grid(meta): + return (triton.cdiv(chunk_size, meta['BR']), n_chunks, bsz * n_heads) + + _chunk_cyfa_fused_readout_fwd_kernel[grid]( + q=q, + q_weight=q_norm_weight, + q_rstd=q_rstd, + k=k, + write=write, + h=h, + g=g_cumsum, + lambdas=lambdas.contiguous(), + table=readout_table.contiguous(), + weights=weights, + raw_output=raw_output, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + scale=scale, + T=seq_len, + H=n_heads, + K=d_k, + M=num_slots, + C=readout_slots, + BT=chunk_size, + STORE_RAW=save_o_prime, + ) + if final_state is not None: + final_state = final_state.transpose(-1, -2).contiguous() + return weights, raw_output, final_state + + +@fla_cache_autotune( + configs=_READOUT_BWD_AUTOTUNE_CONFIGS, + key=['M', 'C'], + **autotune_cache_kwargs, +) +@triton.jit(do_not_specialize=['N']) +def _chunk_cyfa_fused_readout_bwd_kernel( + lambdas, + table, + raw_logits, + dweights, + probs, + dlogits, + draw_logits, + dlambda, + N, + H: tl.constexpr, + M: tl.constexpr, + C: tl.constexpr, + BT: tl.constexpr, + BM: tl.constexpr, +): + i_blk = tl.program_id(0) + i_h = tl.program_id(1) + o_n_local = tl.arange(0, BT) + o_n = i_blk * BT + o_n_local + o_j = tl.arange(0, BM) + o_r = tl.arange(0, BM) + mask_n = o_n < N + mask_j = o_j < M + mask_r = o_r < C + mask_p = 2 * tl.arange(0, BM // 2) + 1 < M + + i_blk_global = i_blk.to(tl.int64) + i_h_global = i_h.to(tl.int64) + token_head_base = i_blk_global * BT * H + i_h_global + lambdas += token_head_base + dlambda += token_head_base + raw_logits += token_head_base * M + dweights += token_head_base * M + draw_logits += token_head_base * M + probs += token_head_base * C + dlogits += token_head_base * C + p_lambda = o_n_local * H + p_x = o_n_local[:, None] * (H * M) + o_j[None, :] + b_lambda = tl.load(lambdas + p_lambda, mask=mask_n, other=0.0).to(tl.float32) + raw = tl.load(raw_logits + p_x, mask=mask_n[:, None] & mask_j[None, :], other=0.0).to(tl.float32) + dw = tl.load(dweights + p_x, mask=mask_n[:, None] & mask_j[None, :], other=0.0).to(tl.float32) + raw_even, raw_odd = tl.split(tl.reshape(raw, (BT, BM // 2, 2))) + dw_even, dw_odd = tl.split(tl.reshape(dw, (BT, BM // 2, 2))) + o_even = 2 * tl.arange(0, BM // 2) + omega = modal_omega(o_even, M) + phase = modal_phase(b_lambda[:, None], o_even[None, :], M) + c = cyfa_cos(phase) + s = cyfa_sin(phase) + raw_z = tl.reshape(tl.join(c * raw_even - s * raw_odd, s * raw_even + c * raw_odd), (BT, BM)) + dw_z = tl.reshape(tl.join(c * dw_even - s * dw_odd, s * dw_even + c * dw_odd), (BT, BM)) + w = tl.load( + table + i_h_global * C * M + o_r[:, None] * M + o_j[None, :], + mask=mask_r[:, None] & mask_j[None, :], + other=0.0, + ).to(tl.float32) + scores = tl.dot(raw_z.to(raw_logits.dtype.element_ty), tl.trans(w.to(raw_logits.dtype.element_ty))) + scores = tl.where(mask_r[None, :], scores, -float("inf")) + dp = tl.dot(dw_z.to(dweights.dtype.element_ty), tl.trans(w.to(dweights.dtype.element_ty))) + max_scores = tl.max(scores, axis=1) + p = tl.where(mask_r[None, :], tl.exp(scores - max_scores[:, None]), 0.0) + denom = tl.sum(p, axis=1) + p *= tl.where(denom > 0.0, 1.0 / denom, 0.0)[:, None] + correction = tl.sum(p * dp, axis=1) + dl = p * (dp - correction[:, None]) + p_probs = o_n_local[:, None] * (H * C) + o_r[None, :] + p_store = p.to(probs.dtype.element_ty) + dl_store = dl.to(dlogits.dtype.element_ty) + tl.store(probs + p_probs, p_store, mask=mask_n[:, None] & mask_r[None, :]) + tl.store(dlogits + p_probs, dl_store, mask=mask_n[:, None] & mask_r[None, :]) + + w_low = w.to(probs.dtype.element_ty) + coeff = tl.dot(p_store, w_low) + grad = tl.dot(dl_store, w_low) + coeff_even, coeff_odd = tl.split(tl.reshape(coeff, (BT, BM // 2, 2))) + grad_even, grad_odd = tl.split(tl.reshape(grad, (BT, BM // 2, 2))) + dx_even = c * grad_even + s * grad_odd + dx_odd = -s * grad_even + c * grad_odd + dx = tl.reshape(tl.join(dx_even, dx_odd), (BT, BM)) + tl.store(draw_logits + p_x, dx, mask=mask_n[:, None] & mask_j[None, :]) + + force_v = coeff_even * (-s * dw_even - c * dw_odd) + force_v += coeff_odd * (c * dw_even - s * dw_odd) + force_k = grad_even * (-s * raw_even - c * raw_odd) + force_k += grad_odd * (c * raw_even - s * raw_odd) + force = tl.sum( + tl.where(mask_n[:, None] & mask_p[None, :], omega[None, :] * (force_v + force_k), 0.0), + axis=1, + ) + tl.store(dlambda + p_lambda, force, mask=mask_n) + + +def _readout_backward( + lambdas: torch.Tensor, + dweights: torch.Tensor, + raw_logits: torch.Tensor, + readout_table: torch.Tensor, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + draw_logits = torch.empty_like(raw_logits) + dlambda = torch.empty_like(lambdas, dtype=torch.float32) + n_heads = raw_logits.shape[-2] + num_slots = raw_logits.shape[-1] + readout_slots = readout_table.shape[-2] + probs = torch.empty(*raw_logits.shape[:-1], readout_slots, device=raw_logits.device, dtype=raw_logits.dtype) + dlogits = torch.empty_like(probs) + n_tokens = raw_logits.numel() // (n_heads * num_slots) + + def grid(meta): + return (triton.cdiv(n_tokens, meta['BT']), n_heads) + + _chunk_cyfa_fused_readout_bwd_kernel[grid]( + lambdas.contiguous(), + readout_table.contiguous(), + raw_logits.contiguous(), + dweights.contiguous(), + probs, + dlogits, + draw_logits, + dlambda, + N=n_tokens, + H=n_heads, + M=num_slots, + C=readout_slots, + BM=triton.next_power_of_2(num_slots), + ) + return probs, dlogits, draw_logits, dlambda + + +def _prune_table_grad_configs(configs, nargs, **kwargs): + args = {**(nargs or {}), **kwargs} + right_a = args.get('right_a') + if not IS_NVIDIA_HOPPER or right_a is None or right_a.dtype != torch.float32: + return configs + # Triton 3.4 rejects this four-warp FP32 dot layout on SM90 during compilation. + return [ + config for config in configs + if not ( + config.kwargs['BR'] == 64 + and config.kwargs['BC'] == 64 + and config.num_warps == 4 + ) + ] + + +@fla_cache_autotune( + configs=_TABLE_GRAD_AUTOTUNE_CONFIGS, + key=['M', 'C'], + prune_configs_by={'early_config_prune': _prune_table_grad_configs}, + **autotune_cache_kwargs, +) +@triton.jit(do_not_specialize=['T']) +def _chunk_cyfa_readout_table_grad_kernel( + lambdas, + left_a, + right_a, + left_b, + right_b, + dtable, + T, + H: tl.constexpr, + M: tl.constexpr, + C: tl.constexpr, + BN: tl.constexpr, + BR: tl.constexpr, + BC: tl.constexpr, +): + i_r = tl.program_id(0) + i_c = tl.program_id(1) + i_hn = tl.program_id(2) + i_h = i_hn % H + i_n = (i_hn // H).to(tl.int64) + o_r = i_r * BR + tl.arange(0, BR) + o_c = i_c * BC + tl.arange(0, BC) + m_r = o_r < C + m_c = o_c < M + acc = tl.zeros((BR, BC), tl.float32) + n_splits = tl.num_programs(2) // H + i_h_global = i_h.to(tl.int64) + for start in range(i_n * BN, T, n_splits * BN): + o_n_local = tl.arange(0, BN) + o_n = start + o_n_local + token_head_base = start * H + i_h_global + m_n = o_n < T + b_lambda = tl.load((lambdas + token_head_base) + o_n_local * H, mask=m_n, other=0.0).to(tl.float32) + p_l = o_n_local[:, None] * (H * C) + o_r[None, :] + p_x = o_n_local[:, None] * (H * M) + o_c[None, :] + l_a = tl.load((left_a + token_head_base * C) + p_l, mask=m_n[:, None] & m_r[None, :], other=0.0) + l_b = tl.load((left_b + token_head_base * C) + p_l, mask=m_n[:, None] & m_r[None, :], other=0.0) + x_a = tl.load((right_a + token_head_base * M) + p_x, mask=m_n[:, None] & m_c[None, :], other=0.0).to(tl.float32) + x_b = tl.load((right_b + token_head_base * M) + p_x, mask=m_n[:, None] & m_c[None, :], other=0.0).to(tl.float32) + a_even, a_odd = tl.split(tl.reshape(x_a, (BN, BC // 2, 2))) + b_even, b_odd = tl.split(tl.reshape(x_b, (BN, BC // 2, 2))) + o_even = i_c * BC + 2 * tl.arange(0, BC // 2) + phase = modal_phase(b_lambda[:, None], o_even[None, :], M) + co = cyfa_cos(phase) + si = cyfa_sin(phase) + z_a = tl.reshape(tl.join(co * a_even - si * a_odd, si * a_even + co * a_odd), (BN, BC)) + z_b = tl.reshape(tl.join(co * b_even - si * b_odd, si * b_even + co * b_odd), (BN, BC)) + acc += tl.dot(tl.trans(l_a.to(right_a.dtype.element_ty)), z_a.to(right_a.dtype.element_ty)) + acc += tl.dot(tl.trans(l_b.to(right_b.dtype.element_ty)), z_b.to(right_b.dtype.element_ty)) + tl.store( + dtable + (i_n * H + i_h_global) * C * M + o_r[:, None] * M + o_c[None, :], + acc, + mask=m_r[:, None] & m_c[None, :], + ) + + +@triton.jit(do_not_specialize=['S']) +def _chunk_cyfa_readout_table_grad_reduce_kernel(partials, output, S, SIZE: tl.constexpr, BLOCK: tl.constexpr): + o_x = tl.program_id(0).to(tl.int64) * BLOCK + tl.arange(0, BLOCK) + o_s = tl.arange(0, 64) + x = tl.load(partials + o_s[:, None] * SIZE + o_x[None, :], + mask=(o_s[:, None] < S) & (o_x[None, :] < SIZE), other=0.0) + tl.store(output + o_x, tl.sum(x, axis=0), mask=o_x < SIZE) + + +def _readout_table_grad_dual( + lambdas: torch.Tensor, + left_a: torch.Tensor, + right_a: torch.Tensor, + left_b: torch.Tensor, + right_b: torch.Tensor, + readout_table: torch.Tensor, +) -> torch.Tensor: + n_heads = left_a.shape[-2] + readout_slots = left_a.shape[-1] + num_slots = right_a.shape[-1] + n_tokens = left_a.numel() // (n_heads * readout_slots) + n_splits = min(32, triton.cdiv(n_tokens, 128)) + partials = torch.empty((n_splits, n_heads, readout_slots, num_slots), + device=readout_table.device, dtype=torch.float32) + dtable = torch.empty_like(readout_table, dtype=torch.float32) + + def grid(meta): + return ( + triton.cdiv(readout_slots, meta['BR']), + triton.cdiv(num_slots, meta['BC']), + n_heads * n_splits, + ) + + _chunk_cyfa_readout_table_grad_kernel[grid]( + lambdas.contiguous(), + left_a.contiguous(), + right_a.contiguous(), + left_b.contiguous(), + right_b.contiguous(), + partials, + T=n_tokens, + H=n_heads, + M=num_slots, + C=readout_slots, + ) + _chunk_cyfa_readout_table_grad_reduce_kernel[(triton.cdiv(dtable.numel(), 128),)]( + partials, dtable, S=n_splits, SIZE=dtable.numel(), BLOCK=128, + ) + return dtable + + +class ChunkCyFAFunction(torch.autograd.Function): + + @staticmethod + @input_guard + @autocast_custom_fwd + def forward( + ctx, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + g: torch.Tensor, + delta: torch.Tensor, + beta: torch.Tensor, + readout_table: torch.Tensor, + q_norm_weight: torch.Tensor, + k_norm_weight: torch.Tensor, + q_norm_eps: float, + k_norm_eps: float, + scale: float, + initial_k: torch.Tensor | None, + initial_v: torch.Tensor | None, + initial_lambda: torch.Tensor | None, + output_final_state: bool, + cu_seqlens: torch.Tensor | None, + cu_seqlens_cpu: torch.Tensor | None, + checkpoint_level: int, + ) -> tuple[torch.Tensor, torch.Tensor | None, torch.Tensor | None, torch.Tensor | None]: + chunk_size = _CHUNK_SIZE + q_input = q + k_input = k + q, k, q_rstd, k_rstd = qk_rmsnorm_fwd( + q_input, + k_input, + q_norm_weight, + k_norm_weight, + q_norm_eps, + k_norm_eps, + store_q=False, + store_rstd=True, + ) + + if cu_seqlens is not None: + chunk_indices = prepare_chunk_indices(cu_seqlens, chunk_size, cu_seqlens_cpu=cu_seqlens_cpu) + else: + chunk_indices = None + + ctx.g_dtype = g.dtype + g = chunk_local_cumsum( + g.float(), + chunk_size=chunk_size, + scale=RCP_LN2, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + output_dtype=torch.float32, + ) + lambdas = cyfa_clock_cumsum(delta, cu_seqlens=cu_seqlens) + if initial_lambda is not None: + lambdas = add_segment_initial(lambdas, initial_lambda, cu_seqlens=cu_seqlens) + write = _point_write( + lambdas, + beta.float(), + readout_table.shape[-1], + dtype=q.dtype, + ) + weights, raw_logits, hkt = _scalar_decay_fused_readout_fwd( + q=q, + q_norm_weight=q_norm_weight, + q_rstd=q_rstd, + k=k, + write=write, + g_cumsum=g, + lambdas=lambdas, + readout_table=readout_table, + initial_state=initial_k, + output_final_state=output_final_state, + scale=scale, + cu_seqlens=cu_seqlens, + chunk_size=chunk_size, + chunk_indices=chunk_indices, + save_o_prime=checkpoint_level == 0, + ) + out, hvt = _scalar_decay_fwd( + q=weights.contiguous(), + k=write, + v=v, + g_cumsum=g, + initial_state=initial_v, + output_final_state=output_final_state, + scale=1.0, + cu_seqlens=cu_seqlens, + chunk_size=chunk_size, + chunk_indices=chunk_indices, + ) + if output_final_state: + flt = segment_sum(delta, cu_seqlens=cu_seqlens) + if initial_lambda is not None: + flt = flt + initial_lambda + else: + flt = None + + ctx.save_for_backward( + q_input, + k_input, + v, + q_norm_weight, + k_norm_weight, + delta, + beta, + g, + lambdas, + readout_table, + raw_logits, + weights, + initial_k, + initial_v, + initial_lambda, + q_rstd, + k_rstd, + chunk_indices, + ) + ctx.q_norm_eps = float(q_norm_eps) + ctx.k_norm_eps = float(k_norm_eps) + ctx.scale = float(scale) + ctx.cu_seqlens = cu_seqlens + return out.to(dtype=q.dtype), hkt, hvt, flt + + @staticmethod + @input_guard + @autocast_custom_bwd + def backward(ctx, do, dhkt=None, dhvt=None, dlt=None): + ( + q_saved, + k_saved, + v, + q_norm_weight, + k_norm_weight, + delta, + beta, + g, + lambdas, + readout_table, + raw_logits, + weights, + initial_k, + initial_v, + initial_lambda, + q_rstd, + k_rstd, + chunk_indices, + ) = ctx.saved_tensors + scale = ctx.scale + cu_seqlens = ctx.cu_seqlens + chunk_size = _CHUNK_SIZE + q, k, _, _ = qk_rmsnorm_fwd( + q_saved, + k_saved, + q_norm_weight, + k_norm_weight, + ctx.q_norm_eps, + ctx.k_norm_eps, + store_rstd=False, + ) + + write = _point_write( + lambdas, + beta.float(), + readout_table.shape[-1], + dtype=q.dtype, + ) + weights_q = weights.contiguous() + dweights, dwrite_v, dv, dg_force_v, dhv0 = _scalar_decay_bwd( + q=weights_q, + k=write, + v=v, + g_cumsum=g, + initial_state=initial_v, + do=do, + dht=dhvt, + scale=1.0, + cu_seqlens=cu_seqlens, + chunk_size=chunk_size, + chunk_indices=chunk_indices, + fuse_dv_single_k=True, + ) + + if raw_logits is None: + _, raw_logits, _ = _scalar_decay_fused_readout_fwd( + q=q_saved, + q_norm_weight=q_norm_weight, + q_rstd=q_rstd, + k=k, + write=write, + g_cumsum=g, + lambdas=lambdas, + readout_table=readout_table, + initial_state=initial_k, + output_final_state=False, + scale=scale, + cu_seqlens=cu_seqlens, + chunk_size=chunk_size, + chunk_indices=chunk_indices, + save_o_prime=True, + ) + + probs, dlogits, draw_logits, dlambda_readout = _readout_backward( + lambdas, + dweights, + raw_logits, + readout_table, + ) + if ctx.needs_input_grad[6]: + dreadout_table = _readout_table_grad_dual( + lambdas, + probs, + dweights, + dlogits, + raw_logits, + readout_table, + ) + else: + dreadout_table = None + + draw_q = draw_logits.contiguous() + dq, dk, dwrite_k, dg_force_k, dhk0 = _scalar_decay_bwd( + q=q, + k=k, + v=write, + g_cumsum=g, + initial_state=initial_k, + do=draw_q, + dht=dhkt, + scale=scale, + cu_seqlens=cu_seqlens, + chunk_size=chunk_size, + chunk_indices=chunk_indices, + fuse_dv_single_k=False, + ) + + dlambda_write, dbeta = _point_write_backward( + dwrite_k, + dwrite_v, + lambdas, + beta.float(), + ) + phase_force = dlambda_write.add_(dlambda_readout) + if dlt is not None: + dlt = dlt.float() + if cu_seqlens is None: + phase_force[:, -1, :].add_(dlt) + else: + end_indices = cu_seqlens[1:].to(device=phase_force.device, dtype=torch.long) - 1 + phase_force[0].index_add_(0, end_indices, dlt) + ddelta = chunk_global_cumsum( + phase_force, + reverse=True, + cu_seqlens=cu_seqlens, + output_dtype=torch.float32, + ) + dlambda0 = ( + None + if initial_lambda is None + else segment_sum(phase_force, cu_seqlens=cu_seqlens) + ) + dg_force = dg_force_v.add_(dg_force_k) + dg = chunk_local_cumsum( + dg_force, + chunk_size=chunk_size, + reverse=True, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + output_dtype=torch.float32, + ) + dq, dk, dq_norm_weight, dk_norm_weight = _qk_rmsnorm_bwd( + dq, + dk, + q_saved, + k_saved, + q_norm_weight, + k_norm_weight, + q_rstd, + k_rstd, + ) + return ( + dq.to(q_saved), + dk.to(k_saved), + dv.to(v), + dg.to(dtype=ctx.g_dtype), + ddelta.to(delta), + dbeta.to(beta), + None if dreadout_table is None else dreadout_table.to(readout_table), + dq_norm_weight, + dk_norm_weight, + None, + None, + None, + dhk0, + dhv0, + dlambda0, + None, + None, + None, + None, + ) + + +@torch.compiler.disable +def chunk_cyfa( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + g: torch.Tensor, + delta: torch.Tensor, + beta: torch.Tensor, + readout: torch.Tensor, + q_norm_weight: torch.Tensor, + k_norm_weight: torch.Tensor, + q_norm_eps: float, + k_norm_eps: float, + scale: float | None = None, + initial_state: CyFAState | None = None, + output_final_state: bool = False, + cu_seqlens: torch.LongTensor | None = None, + cu_seqlens_cpu: torch.LongTensor | None = None, + checkpoint_level: int = 0, +) -> tuple[torch.Tensor, CyFAState | None]: + """Run the chunkwise CyclicFlowAttention recurrence with log-space decay ``g``.""" + num_slots = validate_cyfa_inputs( + q, + k, + v, + g, + delta, + beta, + readout, + require_cuda=True, + ) + cu_seqlens = prepare_cu_seqlens( + cu_seqlens, + batch_size=q.shape[0], + seq_len=q.shape[1], + device=q.device, + ) + _, _, n_heads, d_k = q.shape + d_v = v.shape[-1] + if scale is None: + scale = d_k ** -0.5 + logical_bsz = q.shape[0] if cu_seqlens is None else int(cu_seqlens.numel() - 1) + if initial_state is None: + initial_k = None + initial_v = None + initial_lambda = None + else: + initial_k, initial_v, initial_lambda = validate_initial_state( + initial_state, + batch=logical_bsz, + n_heads=n_heads, + num_slots=num_slots, + d_k=d_k, + d_v=d_v, + device=q.device, + ) + readout_table = build_readout_table(readout) + if checkpoint_level not in (0, 1): + raise ValueError("`checkpoint_level` must be either 0 or 1.") + effective_checkpoint_level = 0 if ( + checkpoint_level == 0 + and torch.is_grad_enabled() + and any(tensor.requires_grad for tensor in (q, k, v, g, delta, beta, readout_table)) + ) else 1 + out, fk, fv, fl = ChunkCyFAFunction.apply( + q, + k, + v, + g, + delta, + beta, + readout_table, + q_norm_weight, + k_norm_weight, + float(q_norm_eps), + float(k_norm_eps), + float(scale), + initial_k, + initial_v, + initial_lambda, + bool(output_final_state), + cu_seqlens, + cu_seqlens_cpu, + effective_checkpoint_level, + ) + final_state = (fk, fv, fl) if output_final_state else None + return out, final_state diff --git a/build/torch-cuda/ops/cyfa/fused_recurrent.py b/build/torch-cuda/ops/cyfa/fused_recurrent.py new file mode 100644 index 0000000000000000000000000000000000000000..59fb0a82a45cdb880ab5f57b650c81585a67584b --- /dev/null +++ b/build/torch-cuda/ops/cyfa/fused_recurrent.py @@ -0,0 +1,1047 @@ +# 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 + + +# Portions adapted from CyclicFlowAttention, Copyright (c) 2026 Yixiao Chen. +# https://github.com/Chyxx/CyclicFlowAttention + +from __future__ import annotations + +import weakref +from functools import cache + +import torch +import triton +import triton.language as tl + +from ...ops.common.chunk_o import NUM_WARPS +from ...ops.utils.cache import fla_cache_autotune +from ...ops.utils.op import exp +from ...utils import IS_NVIDIA, autotune_cache_kwargs, check_shared_mem, get_device_capability, input_guard + +from .qk_norm import qk_rmsnorm_fwd +from .utils import ( + CyFAState, + build_readout_table, + cyfa_cos, + cyfa_sin, + modal_phase, + prepare_cu_seqlens, + validate_cyfa_inputs, + validate_initial_state, +) + +_READOUT_TABLE_CACHE: dict[ + tuple[int, int, int, str, torch.dtype], + tuple[weakref.ReferenceType[torch.Tensor], torch.Tensor], +] = {} +_READOUT_NUM_WARPS = ( + NUM_WARPS if check_shared_mem("ampere") else [w for w in NUM_WARPS if w <= 4] +) + + +def _tensor_version(x: torch.Tensor) -> int: + try: + return int(x._version) + except RuntimeError: + return -1 + + +@cache +def _device_supports_pdl(device_index: int | None) -> bool: + return IS_NVIDIA and get_device_capability(device_index or 0)[0] >= 9 + + +def _k_recurrent_feature_block_size(d_k: int) -> int: + return min(64, triton.next_power_of_2(d_k)) + + +@triton.jit +def _cardinal_orbit_value(lam, o_m, M: tl.constexpr): + is_cos = (o_m % 2) == 0 + phase = modal_phase(lam, o_m, M) + scale = tl.sqrt(tl.where((o_m // 2) == 0, 1.0, 2.0) / (M - 1)) + return scale * tl.where(is_cos, cyfa_cos(phase), -cyfa_sin(phase)) + + +@triton.jit +def _rotate_pos_1d(lam, x, BM: tl.constexpr, M: tl.constexpr): + x_even, x_odd = tl.split(tl.reshape(x, (BM // 2, 2))) + o_even = 2 * tl.arange(0, BM // 2) + phase = modal_phase(lam, o_even, M) + c = cyfa_cos(phase) + s = cyfa_sin(phase) + return tl.reshape(tl.join(c * x_even - s * x_odd, s * x_even + c * x_odd), (BM,)) + + +@triton.jit +def _rotate_neg_1d(lam, x, BM: tl.constexpr, M: tl.constexpr): + x_even, x_odd = tl.split(tl.reshape(x, (BM // 2, 2))) + o_even = 2 * tl.arange(0, BM // 2) + phase = modal_phase(lam, o_even, M) + c = cyfa_cos(phase) + s = cyfa_sin(phase) + return tl.reshape(tl.join(c * x_even + s * x_odd, -s * x_even + c * x_odd), (BM,)) + + +@triton.jit +def _rotate_pos_2d(lam, x, BR: tl.constexpr, BM: tl.constexpr, M: tl.constexpr): + x_even, x_odd = tl.split(tl.reshape(x, (BR, BM // 2, 2))) + o_even = 2 * tl.arange(0, BM // 2) + phase = modal_phase(lam[:, None], o_even[None, :], M) + c = cyfa_cos(phase) + s = cyfa_sin(phase) + return tl.reshape( + tl.join(c * x_even - s * x_odd, s * x_even + c * x_odd), + (BR, BM), + ) + + +@triton.jit +def _rotate_neg_2d(lam, x, BR: tl.constexpr, BM: tl.constexpr, M: tl.constexpr): + x_even, x_odd = tl.split(tl.reshape(x, (BR, BM // 2, 2))) + o_even = 2 * tl.arange(0, BM // 2) + phase = modal_phase(lam[:, None], o_even[None, :], M) + c = cyfa_cos(phase) + s = cyfa_sin(phase) + return tl.reshape( + tl.join(c * x_even + s * x_odd, -s * x_even + c * x_odd), + (BR, BM), + ) + + +@triton.jit +def _as_recurrent_key_dtype(x, ref): + return x.to(ref.dtype.element_ty).to(tl.float32) + + +@triton.heuristics( + { + "USE_INITIAL_STATE": lambda args: args["initial_k"] is not None, + "USE_INITIAL_LAMBDA": lambda args: args["initial_lambda"] is not None, + "STORE_FINAL_STATE": lambda args: args["state_k"] is not None, + "IS_VARLEN": lambda args: args["cu_seqlens"] is not None, + }, +) +@fla_cache_autotune( + configs=[ + triton.Config({"BM": bm}, num_warps=num_warps, num_stages=2) + for bm in [16, 32, 64] + for num_warps in NUM_WARPS + ], + key=["K", "BK", "M", "USE_INITIAL_STATE", "STORE_FINAL_STATE"], + **autotune_cache_kwargs, +) +@triton.jit(do_not_specialize=["T", "TOTAL"]) +def _fused_recurrent_cyfa_k_kernel( + q, + q_norm_weight, + q_rstd, + k, + delta, + beta, + g, + initial_k, + initial_lambda, + state_k, + state_lambda, + lambda_trace, + decay_trace, + raw_partial, + slot_write_out, + cu_seqlens, + scale, + T, + TOTAL, + H: tl.constexpr, + K: tl.constexpr, + M: tl.constexpr, + BK: tl.constexpr, + BM: tl.constexpr, + USE_INITIAL_STATE: tl.constexpr, + USE_INITIAL_LAMBDA: tl.constexpr, + STORE_FINAL_STATE: tl.constexpr, + IS_VARLEN: tl.constexpr, +): + pid = tl.program_id(0).to(tl.int64) + NM: tl.constexpr = triton.cdiv(M, BM) + NK: tl.constexpr = triton.cdiv(K, BK) + i_m = (pid % NM).to(tl.int64) + i_k = ((pid // NM) % NK).to(tl.int64) + i_nh = (pid // (NM * NK)).to(tl.int64) + i_n, i_h = i_nh // H, i_nh % H + if IS_VARLEN: + bos = tl.load(cu_seqlens + i_n).to(tl.int64) + eos = tl.load(cu_seqlens + i_n + 1).to(tl.int64) + T = eos - bos + else: + bos = i_n * T + + o_k = i_k * BK + tl.arange(0, BK) + o_m = i_m * BM + tl.arange(0, BM) + mask_k = o_k < K + mask_m = o_m < M + mask_h = mask_k[:, None] & mask_m[None, :] + + if USE_INITIAL_LAMBDA: + lam = tl.load(initial_lambda + i_n * H + i_h).to(tl.float32) + else: + lam = 0.0 + if USE_INITIAL_STATE: + p_initial = initial_k + ((i_n * H + i_h) * K + o_k[:, None]) * M + o_m[None, :] + state = tl.load(p_initial, mask=mask_h, other=0.0).to(tl.float32) + else: + state = tl.zeros((BK, BM), dtype=tl.float32) + + p_q = q + (bos * H + i_h) * K + o_k + p_q_rstd = q_rstd + bos * H + i_h + p_q_weight = q_norm_weight + o_k + p_k = k + (bos * H + i_h) * K + o_k + p_delta = delta + bos * H + i_h + p_beta = beta + bos * H + i_h + p_g = g + bos * H + i_h + p_raw = raw_partial + ((i_k * TOTAL + bos) * H + i_h) * M + o_m + p_write = slot_write_out + (bos * H + i_h) * M + o_m + p_lambda = lambda_trace + bos * H + i_h + p_decay = decay_trace + bos * H + i_h + + for _ in tl.range(0, T, num_stages=2): + d = tl.load(p_delta, eviction_policy="evict_last").to(tl.float32) + b = tl.load(p_beta, eviction_policy="evict_last").to(tl.float32) + lam += d + write = _as_recurrent_key_dtype( + b * _cardinal_orbit_value(lam, o_m, M), + q, + ) + if i_k == 0: + tl.store(p_write, write, mask=mask_m) + decay = exp(tl.load(p_g, eviction_policy="evict_last").to(tl.float32)) + state *= decay + value = tl.load(p_k, mask=mask_k, other=0.0, eviction_policy="evict_last").to(tl.float32) + state += value[:, None] * write[None, :] + query = tl.load(p_q, mask=mask_k, other=0.0, eviction_policy="evict_last").to(tl.float32) + query *= tl.load(p_q_rstd, eviction_policy="evict_last").to(tl.float32) + query *= tl.load( + p_q_weight, + mask=mask_k, + other=0.0, + eviction_policy="evict_last", + ).to(tl.float32) + query = _as_recurrent_key_dtype(query, q) * scale + raw = tl.sum(state * query[:, None], axis=0) + tl.store(p_raw, raw, mask=mask_m) + if i_k == 0 and i_m == 0: + tl.store(p_lambda, lam) + tl.store(p_decay, decay) + + p_q += H * K + p_q_rstd += H + p_k += H * K + p_delta += H + p_beta += H + p_g += H + p_raw += H * M + p_write += H * M + p_lambda += H + p_decay += H + + if STORE_FINAL_STATE: + p_state = state_k + ((i_n * H + i_h) * K + o_k[:, None]) * M + o_m[None, :] + tl.store(p_state, state, mask=mask_h) + if i_k == 0 and i_m == 0: + tl.store(state_lambda + i_n * H + i_h, lam) + + +@fla_cache_autotune( + configs=[ + triton.Config({"BT": bt}, num_warps=num_warps, num_stages=num_stages) + for bt in [16, 32] + for num_warps in _READOUT_NUM_WARPS + for num_stages in [2, 3, 4] + ], + key=["M", "C", "NK"], + **autotune_cache_kwargs, +) +@triton.jit(do_not_specialize=["TOTAL"]) +def _fused_recurrent_cyfa_readout_kernel( + raw_partial, + lambda_trace, + readout_table, + slot_read_out, + TOTAL, + H: tl.constexpr, + M: tl.constexpr, + C: tl.constexpr, + BM: tl.constexpr, + NK: tl.constexpr, + BT: tl.constexpr, +): + pid = tl.program_id(0).to(tl.int64) + NT = tl.cdiv(TOTAL, BT) + i_t = pid % NT + i_h = pid // NT + o_t = i_t * BT + tl.arange(0, BT) + o_m = tl.arange(0, BM) + o_c = tl.arange(0, BM) + mask_t = o_t < TOTAL + mask_m = o_m < M + mask_c = o_c < C + mask = mask_t[:, None] & mask_m[None, :] + + raw = tl.zeros((BT, BM), dtype=tl.float32) + for i_k in tl.static_range(0, NK): + p_raw = raw_partial + ((i_k * TOTAL + o_t[:, None]) * H + i_h) * M + o_m[None, :] + raw += tl.load(p_raw, mask=mask, other=0.0).to(tl.float32) + lam = tl.load(lambda_trace + o_t * H + i_h, mask=mask_t, other=0.0).to(tl.float32) + rotated = _rotate_pos_2d(lam, raw, BT, BM, M) + table = tl.load( + readout_table + i_h * C * M + o_c[:, None] * M + o_m[None, :], + mask=mask_c[:, None] & mask_m[None, :], + other=0.0, + ).to(slot_read_out.dtype.element_ty) + logits = tl.dot(rotated.to(slot_read_out.dtype.element_ty), tl.trans(table)) + logits = tl.where( + mask_t[:, None] & mask_c[None, :], + logits, + 0.0, + ) + max_logits = tl.max(tl.where(mask_c[None, :], logits, -float("inf")), axis=1) + probs = tl.where(mask_c[None, :], tl.exp(logits - max_logits[:, None]), 0.0) + denom = tl.sum(probs, axis=1) + inv_denom = tl.where(denom > 0.0, 1.0 / denom, 0.0) + probs *= inv_denom[:, None] + probs = probs.to(slot_read_out.dtype.element_ty) + coeff = tl.dot(probs, table) + slot_read = _rotate_neg_2d(lam, coeff, BT, BM, M) + tl.store( + slot_read_out + (o_t[:, None] * H + i_h) * M + o_m[None, :], + slot_read, + mask=mask, + ) + + +@triton.heuristics( + { + "USE_INITIAL_STATE": lambda args: args["initial_v"] is not None, + "STORE_FINAL_STATE": lambda args: args["state_v"] is not None, + "IS_VARLEN": lambda args: args["cu_seqlens"] is not None, + }, +) +@fla_cache_autotune( + configs=[ + triton.Config({"BV": bv}, num_warps=num_warps, num_stages=2) + for bv in [8, 16, 32, 64] + for num_warps in NUM_WARPS + ], + key=["BM", "V", "USE_INITIAL_STATE", "STORE_FINAL_STATE"], + **autotune_cache_kwargs, +) +@triton.jit(do_not_specialize=["T"]) +def _fused_recurrent_cyfa_v_kernel( + slot_read, + slot_write, + v, + decay_trace, + initial_v, + state_v, + out, + cu_seqlens, + T, + H: tl.constexpr, + M: tl.constexpr, + V: tl.constexpr, + BM: tl.constexpr, + BV: tl.constexpr, + USE_INITIAL_STATE: tl.constexpr, + STORE_FINAL_STATE: tl.constexpr, + IS_VARLEN: tl.constexpr, +): + pid = tl.program_id(0).to(tl.int64) + NV: tl.constexpr = triton.cdiv(V, BV) + i_v = (pid % NV).to(tl.int64) + i_nh = (pid // NV).to(tl.int64) + i_n, i_h = i_nh // H, i_nh % H + if IS_VARLEN: + bos = tl.load(cu_seqlens + i_n).to(tl.int64) + eos = tl.load(cu_seqlens + i_n + 1).to(tl.int64) + T = eos - bos + else: + bos = i_n * T + + o_m = tl.arange(0, BM) + o_v = i_v * BV + tl.arange(0, BV) + mask_m = o_m < M + mask_v = o_v < V + mask_h = mask_m[:, None] & mask_v[None, :] + if USE_INITIAL_STATE: + p_initial = initial_v + ((i_n * H + i_h) * M + o_m[:, None]) * V + o_v[None, :] + state = tl.load(p_initial, mask=mask_h, other=0.0).to(tl.float32) + else: + state = tl.zeros((BM, BV), dtype=tl.float32) + + p_read = slot_read + (bos * H + i_h) * M + o_m + p_write = slot_write + (bos * H + i_h) * M + o_m + p_v = v + (bos * H + i_h) * V + o_v + p_decay = decay_trace + bos * H + i_h + p_out = out + (bos * H + i_h) * V + o_v + + for _ in tl.range(0, T, num_stages=2): + state *= tl.load(p_decay, eviction_policy="evict_last").to(tl.float32) + write = tl.load(p_write, mask=mask_m, other=0.0, eviction_policy="evict_last").to(tl.float32) + value = tl.load(p_v, mask=mask_v, other=0.0, eviction_policy="evict_first").to(tl.float32) + state += write[:, None] * value[None, :] + read = tl.load(p_read, mask=mask_m, other=0.0, eviction_policy="evict_last").to(tl.float32) + result = tl.sum(state * read[:, None], axis=0) + tl.store(p_out, result, mask=mask_v) + + p_read += H * M + p_write += H * M + p_v += H * V + p_decay += H + p_out += H * V + + if STORE_FINAL_STATE: + p_state = state_v + ((i_n * H + i_h) * M + o_m[:, None]) * V + o_v[None, :] + tl.store(p_state, state, mask=mask_h) + + +@triton.heuristics( + { + "USE_INITIAL_STATE": lambda args: args["initial_k"] is not None, + "USE_INITIAL_LAMBDA": lambda args: args["initial_lambda"] is not None, + "STORE_FINAL_STATE": lambda args: args["state_k"] is not None, + }, +) +@triton.jit +def _fused_recurrent_cyfa_decode_k_kernel( + q, + q_norm_weight, + k, + k_norm_weight, + delta, + beta, + g, + initial_k, + initial_lambda, + state_k, + workspace, + lambda_value, + scale, + q_norm_eps, + k_norm_eps, + H: tl.constexpr, + K: tl.constexpr, + M: tl.constexpr, + BK: tl.constexpr, + BM: tl.constexpr, + USE_PDL: tl.constexpr, + USE_INITIAL_STATE: tl.constexpr, + USE_INITIAL_LAMBDA: tl.constexpr, + STORE_FINAL_STATE: tl.constexpr, +): + pid = tl.program_id(0).to(tl.int64) + NM: tl.constexpr = triton.cdiv(M, BM) + NK: tl.constexpr = triton.cdiv(K, BK) + i_m = pid % NM + i_k = (pid // NM) % NK + i_bh = (pid // (NM * NK)).to(tl.int64) + i_b = i_bh // H + i_h = i_bh % H + if USE_PDL: + tl.inline_asm_elementwise( + 'griddepcontrol.launch_dependents; // dummy $0', + '=r', + [], + dtype=tl.int32, + is_pure=False, + pack=1, + ) + + o_m = i_m * BM + tl.arange(0, BM) + o_k = i_k * BK + tl.arange(0, BK) + mask_m = o_m < M + mask_k = o_k < K + mask_h = mask_m[:, None] & mask_k[None, :] + + q_square = 0.0 + k_square = 0.0 + for k_start in tl.range(0, K, BK, num_stages=2): + o_norm = k_start + tl.arange(0, BK) + mask_norm = o_norm < K + q_norm = tl.load(q + i_bh * K + o_norm, mask=mask_norm, other=0.0).to(tl.float32) + k_norm = tl.load(k + i_bh * K + o_norm, mask=mask_norm, other=0.0).to(tl.float32) + q_square += tl.sum(q_norm * q_norm) + k_square += tl.sum(k_norm * k_norm) + q_rstd = tl.rsqrt(q_square / K + q_norm_eps) + k_rstd = tl.rsqrt(k_square / K + k_norm_eps) + + if USE_INITIAL_LAMBDA: + lam = tl.load(initial_lambda + i_b * H + i_h).to(tl.float32) + else: + lam = 0.0 + lam += tl.load(delta + i_bh).to(tl.float32) + b = tl.load(beta + i_bh).to(tl.float32) + slot_write = _as_recurrent_key_dtype( + b * _cardinal_orbit_value(lam, o_m, M), + q, + ) + decay = exp(tl.load(g + i_bh).to(tl.float32)) + + if USE_INITIAL_STATE: + p_initial = initial_k + ((i_bh * K + o_k[None, :]) * M + o_m[:, None]) + state = tl.load(p_initial, mask=mask_h, other=0.0).to(tl.float32) + else: + state = tl.zeros((BM, BK), dtype=tl.float32) + + key = tl.load(k + i_bh * K + o_k, mask=mask_k, other=0.0).to(tl.float32) + key *= k_rstd + key *= tl.load(k_norm_weight + o_k, mask=mask_k, other=0.0).to(tl.float32) + key = _as_recurrent_key_dtype(key, k) + state = state * decay + slot_write[:, None] * key[None, :] + + query = tl.load(q + i_bh * K + o_k, mask=mask_k, other=0.0).to(tl.float32) + query *= q_rstd + query *= tl.load(q_norm_weight + o_k, mask=mask_k, other=0.0).to(tl.float32) + query = _as_recurrent_key_dtype(query, q) * scale + raw = tl.sum(state * query[None, :], axis=1) + if USE_PDL: + tl.store(workspace + ((i_bh * NK + i_k) * M + o_m), raw, mask=mask_m) + else: + tl.store(workspace + ((i_bh * (NK + 1) + i_k) * M + o_m), raw, mask=mask_m) + if i_k == 0: + tl.store(workspace + ((i_bh * (NK + 1) + NK) * M + o_m), slot_write, mask=mask_m) + if i_k == 0 and i_m == 0: + tl.store(lambda_value + i_bh, lam) + if STORE_FINAL_STATE: + p_state = state_k + ((i_bh * K + o_k[None, :]) * M + o_m[:, None]) + tl.store(p_state, state, mask=mask_h) + + +@triton.heuristics( + { + "USE_INITIAL_STATE": lambda args: args["initial_v"] is not None, + "USE_INITIAL_LAMBDA": lambda args: args["initial_lambda"] is not None, + "STORE_FINAL_STATE": lambda args: args["state_v"] is not None, + }, +) +@triton.jit +def _fused_recurrent_cyfa_decode_v_kernel( + v, + g, + delta, + beta, + readout_table, + workspace, + lambda_value, + initial_lambda, + initial_v, + state_v, + state_lambda, + out, + H: tl.constexpr, + V: tl.constexpr, + M: tl.constexpr, + C: tl.constexpr, + NK: tl.constexpr, + BV: tl.constexpr, + BM: tl.constexpr, + USE_PDL: tl.constexpr, + USE_INITIAL_STATE: tl.constexpr, + USE_INITIAL_LAMBDA: tl.constexpr, + STORE_FINAL_STATE: tl.constexpr, +): + pid = tl.program_id(0).to(tl.int64) + NV: tl.constexpr = triton.cdiv(V, BV) + i_v = pid % NV + i_bh = (pid // NV).to(tl.int64) + i_h = i_bh % H + o_m = tl.arange(0, BM) + o_c = tl.arange(0, BM) + o_v = i_v * BV + tl.arange(0, BV) + mask_m = o_m < M + mask_c = o_c < C + mask_v = o_v < V + + if USE_PDL: + if USE_INITIAL_LAMBDA: + lam = tl.load(initial_lambda + i_bh).to(tl.float32) + else: + lam = 0.0 + lam += tl.load(delta + i_bh).to(tl.float32) + slot_write = _as_recurrent_key_dtype( + tl.load(beta + i_bh).to(tl.float32) * _cardinal_orbit_value(lam, o_m, M), + v, + ) + else: + lam = tl.load(lambda_value + i_bh).to(tl.float32) + slot_write = tl.load( + workspace + ((i_bh * (NK + 1) + NK) * M + o_m), + mask=mask_m, + other=0.0, + ).to(tl.float32) + table = tl.load( + readout_table + i_h * C * M + o_c[:, None] * M + o_m[None, :], + mask=mask_c[:, None] & mask_m[None, :], + other=0.0, + ).to(tl.float32) + decay = exp(tl.load(g + i_bh).to(tl.float32)) + value = tl.load(v + i_bh * V + o_v, mask=mask_v, other=0.0).to(tl.float32) + if USE_PDL: + tl.inline_asm_elementwise( + 'griddepcontrol.wait; // dummy $0', + '=r', + [], + dtype=tl.int32, + is_pure=False, + pack=1, + ) + + raw = tl.zeros((BM,), dtype=tl.float32) + for i_k in tl.static_range(0, NK): + if USE_PDL: + raw += tl.load( + workspace + ((i_bh * NK + i_k) * M + o_m), + mask=mask_m, + other=0.0, + ).to(tl.float32) + else: + raw += tl.load( + workspace + ((i_bh * (NK + 1) + i_k) * M + o_m), + mask=mask_m, + other=0.0, + ).to(tl.float32) + z = _as_recurrent_key_dtype(_rotate_pos_1d(lam, raw, BM, M), v) + logits = tl.sum(table * z[None, :], axis=1) + max_logits = tl.max(tl.where(mask_c, logits, -float("inf")), axis=0) + probs = tl.where(mask_c, tl.exp(logits - max_logits), 0.0) + denom = tl.sum(probs, axis=0) + probs *= tl.where(denom > 0.0, 1.0 / denom, 0.0) + coeff = tl.sum( + table * probs.to(v.dtype.element_ty).to(tl.float32)[:, None], + axis=0, + ) + slot_read = _as_recurrent_key_dtype(_rotate_neg_1d(lam, coeff, BM, M), v) + + mask_state = mask_m[:, None] & mask_v[None, :] + if USE_INITIAL_STATE: + p_initial = initial_v + ((i_bh * M + o_m[:, None]) * V + o_v[None, :]) + state = tl.load(p_initial, mask=mask_state, other=0.0).to(tl.float32) + else: + state = tl.zeros((BM, BV), dtype=tl.float32) + state = state * decay + slot_write[:, None] * value[None, :] + result = tl.sum(state * slot_read[:, None], axis=0) + tl.store(out + i_bh * V + o_v, result, mask=mask_v) + + if STORE_FINAL_STATE: + p_state = state_v + ((i_bh * M + o_m[:, None]) * V + o_v[None, :]) + tl.store(p_state, state, mask=mask_state) + if i_v == 0: + tl.store(state_lambda + i_bh, lam) + + +def _fused_recurrent_cyfa_decode_fwd( + q: torch.Tensor, + q_norm_weight: torch.Tensor, + k: torch.Tensor, + k_norm_weight: torch.Tensor, + v: torch.Tensor, + delta: torch.Tensor, + beta: torch.Tensor, + g: torch.Tensor, + readout_table: torch.Tensor, + scale: float, + q_norm_eps: float, + k_norm_eps: float, + initial_state: CyFAState | None, + output_final_state: bool, +) -> tuple[torch.Tensor, CyFAState | None]: + batch, _, n_heads, d_k = q.shape + d_v = v.shape[-1] + num_slots = readout_table.shape[-1] + use_pdl = _device_supports_pdl(q.device.index) + state = validate_initial_state( + initial_state, + batch=batch, + n_heads=n_heads, + num_slots=num_slots, + d_k=d_k, + d_v=d_v, + device=q.device, + ) + if state is None: + initial_k, initial_v, initial_lambda = None, None, None + else: + initial_k, initial_v, initial_lambda = state + state_k = state_v = state_lambda = None + if output_final_state: + if initial_k is None: + state_k = torch.empty(batch, n_heads, d_k, num_slots, device=q.device, dtype=torch.float32) + state_v = torch.empty(batch, n_heads, num_slots, d_v, device=q.device, dtype=torch.float32) + state_lambda = torch.empty(batch, n_heads, device=q.device, dtype=torch.float32) + else: + state_k, state_v = initial_k, initial_v + state_lambda = torch.empty_like(initial_lambda) if use_pdl else initial_lambda + + bk = min(128, triton.next_power_of_2(d_k)) + bm = min(64, triton.next_power_of_2(num_slots)) + bv = min(64, triton.next_power_of_2(d_v)) + nk = triton.cdiv(d_k, bk) + nm = triton.cdiv(num_slots, bm) + nv = triton.cdiv(d_v, bv) + workspace = torch.empty( + batch, + n_heads, + nk if use_pdl else nk + 1, + num_slots, + device=q.device, + dtype=q.dtype, + ) + lambda_value = None + if not use_pdl: + lambda_value = torch.empty(batch, n_heads, device=q.device, dtype=torch.float32) + out = torch.empty_like(v) + + _fused_recurrent_cyfa_decode_k_kernel[(batch * n_heads * nk * nm,)]( + q, + q_norm_weight, + k, + k_norm_weight, + delta, + beta, + g, + initial_k, + initial_lambda, + state_k, + workspace, + lambda_value, + float(scale), + float(q_norm_eps), + float(k_norm_eps), + H=n_heads, + K=d_k, + M=num_slots, + BK=bk, + BM=bm, + USE_PDL=use_pdl, + num_warps=4, + num_stages=2, + ) + _fused_recurrent_cyfa_decode_v_kernel[(batch * n_heads * nv,)]( + v, + g, + delta, + beta, + readout_table, + workspace, + lambda_value, + initial_lambda, + initial_v, + state_v, + state_lambda, + out, + H=n_heads, + V=d_v, + M=num_slots, + C=readout_table.shape[-2], + NK=nk, + BV=bv, + BM=triton.next_power_of_2(num_slots), + USE_PDL=use_pdl, + launch_pdl=use_pdl, + num_warps=4, + num_stages=2, + ) + final_state = (state_k, state_v, state_lambda) if output_final_state else None + return out, final_state + + +def _fused_recurrent_cyfa_fwd( + q: torch.Tensor, + q_norm_weight: torch.Tensor, + q_rstd: torch.Tensor | None, + k: torch.Tensor, + k_norm_weight: torch.Tensor, + v: torch.Tensor, + delta: torch.Tensor, + beta: torch.Tensor, + g: torch.Tensor, + readout_table: torch.Tensor, + scale: float, + q_norm_eps: float, + k_norm_eps: float, + initial_state: CyFAState | None, + output_final_state: bool, + cu_seqlens: torch.Tensor | None, +) -> tuple[torch.Tensor, CyFAState | None]: + batch, seq_len, n_heads, d_k = q.shape + d_v = v.shape[-1] + num_slots = readout_table.shape[-1] + if seq_len == 1 and cu_seqlens is None: + return _fused_recurrent_cyfa_decode_fwd( + q=q, + q_norm_weight=q_norm_weight, + k=k, + k_norm_weight=k_norm_weight, + v=v, + delta=delta, + beta=beta, + g=g, + readout_table=readout_table, + scale=scale, + q_norm_eps=q_norm_eps, + k_norm_eps=k_norm_eps, + initial_state=initial_state, + output_final_state=output_final_state, + ) + if cu_seqlens is None: + n_seq = batch + else: + cu_seqlens = prepare_cu_seqlens( + cu_seqlens, + batch_size=batch, + seq_len=seq_len, + device=q.device, + ) + n_seq = int(cu_seqlens.numel() - 1) + + state = validate_initial_state( + initial_state, + batch=n_seq, + n_heads=n_heads, + num_slots=num_slots, + d_k=d_k, + d_v=d_v, + device=q.device, + ) + if state is None: + initial_k, initial_v, initial_lambda = None, None, None + else: + initial_k, initial_v, initial_lambda = state + + state_k = None + state_v = None + state_lambda = None + if output_final_state: + state_k = torch.empty(n_seq, n_heads, d_k, num_slots, device=q.device, dtype=torch.float32) + state_v = torch.empty(n_seq, n_heads, num_slots, d_v, device=q.device, dtype=torch.float32) + state_lambda = torch.empty(n_seq, n_heads, device=q.device, dtype=torch.float32) + out = torch.empty_like(v) + total_tokens = batch * seq_len + lambda_trace = torch.empty(batch, seq_len, n_heads, device=q.device, dtype=torch.float32) + decay_trace = torch.empty_like(lambda_trace) + slot_write = torch.empty(batch, seq_len, n_heads, num_slots, device=q.device, dtype=q.dtype) + slot_read = torch.empty_like(slot_write) + bk = _k_recurrent_feature_block_size(d_k) + readout_slots = readout_table.shape[-2] + bm = triton.next_power_of_2(num_slots) + bm_state = triton.next_power_of_2(num_slots) + nk = triton.cdiv(d_k, bk) + raw_partial = torch.empty( + nk, + batch, + seq_len, + n_heads, + num_slots, + device=q.device, + dtype=q.dtype, + ) + + _fused_recurrent_cyfa_k_kernel[ + lambda meta: (triton.cdiv(num_slots, meta["BM"]) * nk * n_seq * n_heads,) + ]( + q, + q_norm_weight, + q_rstd, + k, + delta, + beta, + g, + initial_k, + initial_lambda, + state_k, + state_lambda, + lambda_trace, + decay_trace, + raw_partial, + slot_write, + cu_seqlens, + float(scale), + seq_len, + total_tokens, + H=n_heads, + K=d_k, + M=num_slots, + BK=bk, + ) + _fused_recurrent_cyfa_readout_kernel[ + lambda meta: (triton.cdiv(total_tokens, meta["BT"]) * n_heads,) + ]( + raw_partial, + lambda_trace, + readout_table, + slot_read, + total_tokens, + H=n_heads, + M=num_slots, + C=readout_slots, + BM=bm, + NK=nk, + ) + _fused_recurrent_cyfa_v_kernel[ + lambda meta: (triton.cdiv(d_v, meta["BV"]) * n_seq * n_heads,) + ]( + slot_read, + slot_write, + v, + decay_trace, + initial_v, + state_v, + out, + cu_seqlens, + seq_len, + H=n_heads, + M=num_slots, + V=d_v, + BM=bm_state, + ) + final_state = (state_k, state_v, state_lambda) if output_final_state else None + return out, final_state + + +class FusedRecurrentCyFAFunction(torch.autograd.Function): + @staticmethod + def forward( + ctx, + q: torch.Tensor, + q_norm_weight: torch.Tensor, + q_rstd: torch.Tensor | None, + k: torch.Tensor, + k_norm_weight: torch.Tensor, + v: torch.Tensor, + g: torch.Tensor, + delta: torch.Tensor, + beta: torch.Tensor, + readout_table: torch.Tensor, + scale: float, + q_norm_eps: float, + k_norm_eps: float, + initial_state: CyFAState | None, + output_final_state: bool, + cu_seqlens: torch.Tensor | None, + ): + return _fused_recurrent_cyfa_fwd( + q=q, + q_norm_weight=q_norm_weight, + q_rstd=q_rstd, + k=k, + k_norm_weight=k_norm_weight, + v=v, + delta=delta, + beta=beta, + g=g, + readout_table=readout_table, + scale=float(scale), + q_norm_eps=float(q_norm_eps), + k_norm_eps=float(k_norm_eps), + initial_state=initial_state, + output_final_state=bool(output_final_state), + cu_seqlens=cu_seqlens, + ) + + @staticmethod + @input_guard + def backward(ctx, *args): + raise NotImplementedError("`fused_recurrent_cyfa` is inference-only; use chunk mode for training.") + + +@input_guard +def fused_recurrent_cyfa( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + g: torch.Tensor, + delta: torch.Tensor, + beta: torch.Tensor, + readout: torch.Tensor, + q_norm_weight: torch.Tensor, + k_norm_weight: torch.Tensor, + q_norm_eps: float, + k_norm_eps: float, + scale: float | None = None, + initial_state: CyFAState | None = None, + output_final_state: bool = False, + cu_seqlens: torch.LongTensor | None = None, +) -> tuple[torch.Tensor, CyFAState | None]: + """Run the inference-only recurrent CyclicFlowAttention path with log-space decay ``g``.""" + num_slots = validate_cyfa_inputs( + q, + k, + v, + g, + delta, + beta, + readout, + require_cuda=True, + ) + if q.shape != k.shape or q.shape[-1] != q_norm_weight.numel() or q.shape[-1] != k_norm_weight.numel(): + raise ValueError("Fused Q/K RMSNorm requires matching contiguous Q/K tensors and weights.") + grad_inputs = (q, k, v, g, delta, beta, readout, q_norm_weight, k_norm_weight, *(initial_state or ())) + if torch.is_grad_enabled() and any(x is not None and x.requires_grad for x in grad_inputs): + raise NotImplementedError("`fused_recurrent_cyfa` is inference-only; use chunk mode for training.") + cu_seqlens = prepare_cu_seqlens( + cu_seqlens, + batch_size=q.shape[0], + seq_len=q.shape[1], + device=q.device, + ) + + _, _, _, d_k = q.shape + if scale is None: + scale = d_k ** -0.5 + if q.shape[1] == 1 and cu_seqlens is None: + q = q.contiguous() + k = k.contiguous() + q_rstd = None + else: + q, k, q_rstd, _ = qk_rmsnorm_fwd( + q, + k, + q_norm_weight, + k_norm_weight, + float(q_norm_eps), + float(k_norm_eps), + store_q=False, + store_rstd=True, + ) + cache_key = (id(readout), _tensor_version(readout), int(num_slots), str(q.device), q.dtype) + cache_entry = _READOUT_TABLE_CACHE.get(cache_key) + readout_table = None + if cache_entry is not None: + readout_ref, cached_table = cache_entry + if readout_ref() is readout: + readout_table = cached_table + if readout_table is None: + readout_table = build_readout_table(readout).to(q.dtype) + _READOUT_TABLE_CACHE[cache_key] = (weakref.ref(readout), readout_table) + return FusedRecurrentCyFAFunction.apply( + q, + q_norm_weight, + q_rstd, + k, + k_norm_weight, + v, + g, + delta, + beta, + readout_table, + float(scale), + float(q_norm_eps), + float(k_norm_eps), + initial_state, + bool(output_final_state), + cu_seqlens, + ) diff --git a/build/torch-cuda/ops/cyfa/naive.py b/build/torch-cuda/ops/cyfa/naive.py new file mode 100644 index 0000000000000000000000000000000000000000..ba6ee690a618d9c02cabffb33a70ebc29c134417 --- /dev/null +++ b/build/torch-cuda/ops/cyfa/naive.py @@ -0,0 +1,211 @@ +# 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 + + +# Portions adapted from CyclicFlowAttention, Copyright (c) 2026 Yixiao Chen. +# https://github.com/Chyxx/CyclicFlowAttention + +from __future__ import annotations + +import torch + +from .utils import CyFAState, build_readout_table, prepare_cu_seqlens, validate_cyfa_inputs, validate_initial_state + + +def _rotation_terms( + lambdas: torch.Tensor, + num_slots: int, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + idx = torch.arange(num_slots, device=lambdas.device, dtype=torch.long) + is_cos = idx.remainder(2) == 0 + mate = torch.where(is_cos, idx + 1, idx - 1).clamp(0, num_slots - 1) + period = num_slots - 1 + pair = torch.arange(num_slots // 2, device=lambdas.device, dtype=torch.long) + pair = pair.repeat_interleave(2) + whole = torch.floor(lambdas) + fraction = lambdas - whole + coordinate = whole.to(torch.long).remainder(period).unsqueeze(-1) + coordinate = (coordinate * pair).remainder(period).to(lambdas.dtype) + coordinate = coordinate + fraction.unsqueeze(-1) * pair.to(lambdas.dtype) + coordinate = torch.where(coordinate > period * 0.5, coordinate - period, coordinate) + phase = coordinate * (2 * torch.pi / period) + return is_cos, mate, torch.cos(phase), torch.sin(phase) + + +def _rotate_atom( + lambdas: torch.Tensor, + num_slots: int, +) -> torch.Tensor: + is_cos, _, c, s = _rotation_terms(lambdas, num_slots) + pair = torch.arange(num_slots, device=lambdas.device) // 2 + scale = torch.where(pair == 0, 1.0, 2.0).div(num_slots - 1).sqrt() + return scale * torch.where(is_cos.view(1, 1, -1), c, -s) + + +def _rotate_modal_pos( + lambdas: torch.Tensor, + x: torch.Tensor, + num_slots: int, +) -> torch.Tensor: + is_cos, mate, c, s = _rotation_terms(lambdas, num_slots) + xm = x[:, :, mate] + y_cos = c * x - s * xm + y_sin = s * xm + c * x + y = torch.where(is_cos.view(1, 1, -1), y_cos, y_sin) + return y + + +def _rotate_modal_neg( + lambdas: torch.Tensor, + x: torch.Tensor, + num_slots: int, +) -> torch.Tensor: + is_cos, mate, c, s = _rotation_terms(lambdas, num_slots) + xm = x[:, :, mate] + y_cos = c * x + s * xm + y_sin = -s * xm + c * x + y = torch.where(is_cos.view(1, 1, -1), y_cos, y_sin) + return y + + +def _rmsnorm_forward( + x: torch.Tensor, + weight: torch.Tensor, + eps: float, +) -> torch.Tensor: + y = x.float() + y = y * torch.rsqrt(y.square().mean(dim=-1, keepdim=True) + float(eps)) + y = y * weight.float() + return y.to(x.dtype) + + +def naive_recurrent_cyfa( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + g: torch.Tensor, + delta: torch.Tensor, + beta: torch.Tensor, + readout: torch.Tensor, + q_norm_weight: torch.Tensor, + k_norm_weight: torch.Tensor, + q_norm_eps: float, + k_norm_eps: float, + scale: float | None = None, + initial_state: CyFAState | None = None, + output_final_state: bool = False, + cu_seqlens: torch.LongTensor | None = None, +) -> tuple[torch.Tensor, CyFAState | None]: + """Token-by-token PyTorch reference for the canonical CyclicFlowAttention recurrence.""" + num_slots = validate_cyfa_inputs(q, k, v, g, delta, beta, readout) + cu_seqlens = prepare_cu_seqlens( + cu_seqlens, + batch_size=q.shape[0], + seq_len=q.shape[1], + device=q.device, + ) + if cu_seqlens is not None: + boundaries = cu_seqlens.detach().cpu().tolist() + n_seq = len(boundaries) - 1 + if initial_state is not None and initial_state[0].shape[0] != n_seq: + raise ValueError("Initial state count must match the packed sequence count.") + outputs = [] + final_parts = [[], [], []] + for i, (start, end) in enumerate(zip(boundaries[:-1], boundaries[1:], strict=True)): + state_i = None if initial_state is None else tuple(x[i: i + 1] for x in initial_state) + output_i, final_i = naive_recurrent_cyfa( + q=q[:, start:end], + k=k[:, start:end], + v=v[:, start:end], + g=g[:, start:end], + delta=delta[:, start:end], + beta=beta[:, start:end], + readout=readout, + scale=scale, + initial_state=state_i, + output_final_state=output_final_state, + q_norm_weight=q_norm_weight, + k_norm_weight=k_norm_weight, + q_norm_eps=q_norm_eps, + k_norm_eps=k_norm_eps, + ) + outputs.append(output_i) + if output_final_state: + for j, value in enumerate(final_i): + final_parts[j].append(value) + output = torch.cat(outputs, dim=1) + final_state = ( + tuple(torch.cat(parts, dim=0) for parts in final_parts) + if output_final_state + else None + ) + return output, final_state + + dtype = q.dtype + q = _rmsnorm_forward(q, q_norm_weight, q_norm_eps) + k = _rmsnorm_forward(k, k_norm_weight, k_norm_eps) + + q, k, v = map(lambda x: x.transpose(1, 2).contiguous().float(), (q, k, v)) + delta = delta.transpose(1, 2).contiguous().float() + beta = beta.transpose(1, 2).contiguous().float() + g = g.contiguous().float() + g = g.transpose(1, 2).contiguous() + bsz, n_heads, seqlen, d_k = q.shape + d_v = v.shape[-1] + if scale is None: + scale = d_k ** -0.5 + + state = validate_initial_state( + initial_state, + batch=bsz, + n_heads=n_heads, + num_slots=num_slots, + d_k=d_k, + d_v=d_v, + device=q.device, + ) + if state is None: + hk = torch.zeros(bsz, n_heads, d_k, num_slots, device=q.device) + hv = torch.zeros(bsz, n_heads, num_slots, d_v, device=q.device) + lambda0 = torch.zeros(bsz, n_heads, device=q.device) + else: + hk, hv, lambda0 = state + lambdas = delta.cumsum(dim=-1) + lambda0[:, :, None] + readout_table = build_readout_table(readout) + out = torch.zeros_like(v) + for i in range(seqlen): + decay = g[:, :, i].exp() + hk = hk * decay[:, :, None, None] + hv = hv * decay[:, :, None, None] + write = beta[:, :, i, None] * _rotate_atom( + lambdas[:, :, i], + num_slots, + ) + hk = hk + k[:, :, i, :, None] * write[:, :, None, :] + hv = hv + write[:, :, :, None] * v[:, :, i, None, :] + + raw_logits = torch.einsum("bhkm,bhk->bhm", hk, q[:, :, i] * float(scale)) + rotated_logits = _rotate_modal_pos(lambdas[:, :, i], raw_logits, num_slots) + logits = torch.einsum("hrm,bhm->bhr", readout_table, rotated_logits) + numer = torch.exp(logits - logits.max(dim=-1, keepdim=True).values) + denom = numer.sum(dim=-1, keepdim=True) + probs = torch.where(denom > 0.0, numer / denom, torch.zeros_like(numer)) + coeff = torch.einsum("hrm,bhr->bhm", readout_table, probs) + weights = _rotate_modal_neg(lambdas[:, :, i], coeff, num_slots) + out[:, :, i] = torch.einsum("bhm,bhmv->bhv", weights, hv) + + out = out.transpose(1, 2).contiguous().to(dtype) + if output_final_state: + if cu_seqlens is None: + final_lambda = lambdas[:, :, -1] + else: + ends = cu_seqlens[1:].to(device=lambdas.device, dtype=torch.long) - 1 + final_lambda = lambdas[0, :, ends].transpose(0, 1).contiguous() + final_state = (hk, hv, final_lambda) + else: + final_state = None + return out, final_state diff --git a/build/torch-cuda/ops/cyfa/qk_norm.py b/build/torch-cuda/ops/cyfa/qk_norm.py new file mode 100644 index 0000000000000000000000000000000000000000..8ed465a27528db78447bdf6c491bb3a56518f0f7 --- /dev/null +++ b/build/torch-cuda/ops/cyfa/qk_norm.py @@ -0,0 +1,117 @@ +# 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 + + +# Portions adapted from CyclicFlowAttention, Copyright (c) 2026 Yixiao Chen. +# https://github.com/Chyxx/CyclicFlowAttention + +from __future__ import annotations + +import torch +import triton +import triton.language as tl + +from ...ops.utils.cache import fla_cache_autotune +from ...utils import autotune_cache_kwargs + + +@fla_cache_autotune( + configs=[ + triton.Config({"BT": bt}, num_warps=num_warps) + for bt in [32, 64, 128] + for num_warps in [4, 8] + ], + key=["D", "STORE_Q", "STORE_RSTD"], + **autotune_cache_kwargs, +) +@triton.jit(do_not_specialize=["N"]) +def _qk_rmsnorm_fwd_kernel( + q, + k, + q_weight, + k_weight, + q_out, + k_out, + q_rstd, + k_rstd, + q_eps, + k_eps, + N, + D: tl.constexpr, + BT: tl.constexpr, + BD: tl.constexpr, + STORE_Q: tl.constexpr, + STORE_RSTD: tl.constexpr, +): + i_n = tl.program_id(0).to(tl.int64) + o_n = i_n * BT + tl.arange(0, BT) + o_d = tl.arange(0, BD) + m_n = o_n < N + m_d = o_d < D + mask = m_n[:, None] & m_d[None, :] + offsets = o_n[:, None] * D + o_d[None, :] + + b_q = tl.load(q + offsets, mask=mask, other=0.0).to(tl.float32) + b_q = tl.where(mask, b_q, 0.0) + b_q_rstd = tl.rsqrt(tl.sum(b_q * b_q, axis=1) / D + q_eps) + b_q_weight = tl.load(q_weight + o_d, mask=m_d, other=0.0).to(tl.float32) + if STORE_Q: + tl.store(q_out + offsets, b_q * b_q_rstd[:, None] * b_q_weight[None, :], mask=mask) + if STORE_RSTD: + tl.store(q_rstd + o_n, b_q_rstd, mask=m_n) + + b_k = tl.load(k + offsets, mask=mask, other=0.0).to(tl.float32) + b_k = tl.where(mask, b_k, 0.0) + b_k_rstd = tl.rsqrt(tl.sum(b_k * b_k, axis=1) / D + k_eps) + b_k_weight = tl.load(k_weight + o_d, mask=m_d, other=0.0).to(tl.float32) + tl.store(k_out + offsets, b_k * b_k_rstd[:, None] * b_k_weight[None, :], mask=mask) + if STORE_RSTD: + tl.store(k_rstd + o_n, b_k_rstd, mask=m_n) + + +def qk_rmsnorm_fwd( + q: torch.Tensor, + k: torch.Tensor, + q_weight: torch.Tensor, + k_weight: torch.Tensor, + q_eps: float, + k_eps: float, + *, + store_q: bool = True, + store_rstd: bool, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor | None, torch.Tensor | None]: + if q.shape != k.shape or q.shape[-1] != q_weight.numel() or q.shape[-1] != k_weight.numel(): + raise ValueError("Fused Q/K RMSNorm requires matching contiguous Q/K tensors and weights.") + q = q.contiguous() + k = k.contiguous() + q_out = torch.empty_like(q) if store_q else q + k_out = torch.empty_like(k) + n_rows = q.numel() // q.shape[-1] + q_rstd = torch.empty(n_rows, device=q.device, dtype=torch.float32) if store_rstd else None + k_rstd = torch.empty(n_rows, device=q.device, dtype=torch.float32) if store_rstd else None + + def grid(meta): + return (triton.cdiv(n_rows, meta["BT"]),) + + _qk_rmsnorm_fwd_kernel[grid]( + q, + k, + q_weight, + k_weight, + q_out, + k_out, + q_rstd, + k_rstd, + float(q_eps), + float(k_eps), + N=n_rows, + D=q.shape[-1], + BD=triton.next_power_of_2(q.shape[-1]), + STORE_Q=bool(store_q), + STORE_RSTD=bool(store_rstd), + ) + return q_out, k_out, q_rstd, k_rstd diff --git a/build/torch-cuda/ops/cyfa/utils.py b/build/torch-cuda/ops/cyfa/utils.py new file mode 100644 index 0000000000000000000000000000000000000000..e6088691cdcb05156b49b2090d2a4fc89ceee666 --- /dev/null +++ b/build/torch-cuda/ops/cyfa/utils.py @@ -0,0 +1,216 @@ +# 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 + + +# Portions adapted from CyclicFlowAttention, Copyright (c) 2026 Yixiao Chen. +# https://github.com/Chyxx/CyclicFlowAttention + +from __future__ import annotations + +from functools import lru_cache + +import torch +import triton +import triton.language as tl +from triton.language.extra import libdevice + +CyFAState = tuple[torch.Tensor, torch.Tensor, torch.Tensor] + + +@triton.jit(do_not_specialize=['T']) +def _cyfa_clock_cumsum_kernel( + delta, + clock, + cu_seqlens, + T, + H: tl.constexpr, + IS_VARLEN: tl.constexpr, + BT: tl.constexpr, +): + i_nh = tl.program_id(0).to(tl.int64) + i_n, i_h = i_nh // H, i_nh % H + if IS_VARLEN: + bos = tl.load(cu_seqlens + i_n).to(tl.int64) + eos = tl.load(cu_seqlens + i_n + 1).to(tl.int64) + else: + bos, eos = i_n * T, (i_n + 1) * T + length = eos - bos + total = tl.full((), 0, tl.float32) + for i_t in range(tl.cdiv(length, BT)): + offsets = i_t * BT + tl.arange(0, BT) + indices = (bos + offsets) * H + i_h + values = tl.load(delta + indices, mask=offsets < length, other=0).to(tl.float32) + prefix = tl.cumsum(values, 0) + total + total += tl.sum(values, 0) + tl.store(clock + indices, prefix, mask=offsets < length) + + +def cyfa_clock_cumsum(delta: torch.Tensor, cu_seqlens: torch.Tensor | None = None) -> torch.Tensor: + # a fixed scan geometry keeps the cyclic phase identical for dense and packed layouts. + B, T, H = delta.shape + N = B if cu_seqlens is None else len(cu_seqlens) - 1 + clock = torch.empty_like(delta, dtype=torch.float32) + _cyfa_clock_cumsum_kernel[(N * H,)]( + delta, clock, cu_seqlens, T=T, H=H, IS_VARLEN=cu_seqlens is not None, + BT=256, num_warps=4, num_stages=1, + ) + return clock + + +@triton.jit +def cyfa_cos(x): + return libdevice.fast_cosf(x) + + +@triton.jit +def cyfa_sin(x): + return libdevice.fast_sinf(x) + + +@triton.jit +def modal_omega(o_m, M: tl.constexpr): + pair = (o_m // 2).to(tl.float32) + return pair * (6.283185307179586 / (M - 1)) + + +@triton.jit +def modal_phase(lam, o_m, M: tl.constexpr): + # reduce in integer slot coordinates before converting to radians. + period: tl.constexpr = M - 1 + lam = lam.to(tl.float32) + pair_i = (o_m // 2).to(tl.int32) + whole = tl.floor(lam) + fraction = lam - whole + whole_mod = whole.to(tl.int32) % period + whole_mod = tl.where(whole_mod < 0, whole_mod + period, whole_mod) + coordinate = ((whole_mod * pair_i) % period).to(tl.float32) + coordinate += fraction * pair_i.to(tl.float32) + coordinate = tl.where(coordinate > period * 0.5, coordinate - period, coordinate) + return coordinate * (6.283185307179586 / period) + + +def validate_num_slots(num_slots: int) -> int: + num_slots = int(num_slots) + if num_slots <= 0 or num_slots % 2: + raise ValueError("`num_slots` must be a positive even integer.") + return num_slots + + +def validate_cyfa_inputs( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + g: torch.Tensor, + delta: torch.Tensor, + beta: torch.Tensor, + readout: torch.Tensor, + *, + require_cuda: bool = False, +) -> int: + if require_cuda and not q.is_cuda: + raise ValueError("CyclicFlowAttention optimized operators require CUDA tensors.") + if q.ndim != 4 or k.ndim != 4 or v.ndim != 4: + raise ValueError("`q`, `k`, and `v` must have shape [B, T, H, D].") + if q.shape[:3] != k.shape[:3] or q.shape[:3] != v.shape[:3]: + raise ValueError("`q`, `k`, and `v` must share [B, T, H].") + if g.shape != q.shape[:3] or delta.shape != q.shape[:3] or beta.shape != q.shape[:3]: + raise ValueError("`g`, `delta`, and `beta` must have shape [B, T, H].") + if readout.ndim != 3 or readout.shape[0] != q.shape[2] or readout.shape[1] != readout.shape[2]: + raise ValueError("`readout` must have shape [H, C, C].") + return validate_num_slots(readout.shape[-1] + 1) + + +def prepare_cu_seqlens( + cu_seqlens: torch.LongTensor | None, + *, + batch_size: int, + seq_len: int, + device: torch.device, +) -> torch.LongTensor | None: + if cu_seqlens is None: + return None + if batch_size != 1: + raise ValueError("Packed variable-length inputs must have batch size 1.") + cu_seqlens = cu_seqlens.to(device=device, dtype=torch.long).contiguous() + if int(cu_seqlens[-1].item()) != seq_len: + raise ValueError("Packed sequence length does not match `cu_seqlens`.") + return cu_seqlens + + +@lru_cache(maxsize=16) +@torch.inference_mode(False) +def _readout_basis(num_slots: int, device: torch.device) -> torch.Tensor: + period = num_slots - 1 + idx = torch.arange(num_slots, device=device) + rows = torch.arange(period, device=device, dtype=torch.float32)[:, None] + omega = (idx // 2) * (2.0 * torch.pi / period) + basis = torch.where( + idx.remainder(2)[None, :] == 1, + torch.sin(rows * omega), + torch.cos(rows * omega), + ) + basis[:, 0] *= period**-0.5 + basis[:, 1] = 0.0 + basis[:, 2:] *= (2.0 / period) ** 0.5 + return basis + + +def build_readout_table(readout: torch.Tensor) -> torch.Tensor: + num_slots = validate_num_slots(readout.shape[-1] + 1) + # only the fixed Fourier basis is cached, never the learned readout table. + basis = _readout_basis(num_slots, readout.device) + return torch.matmul(readout.float(), basis).contiguous() + + +def add_segment_initial( + x: torch.Tensor, + initial: torch.Tensor, + cu_seqlens: torch.Tensor | None, +) -> torch.Tensor: + if cu_seqlens is None: + return x + initial[:, None, :] + out = x.clone() + boundaries = cu_seqlens.detach().cpu().tolist() + for i, (start, end) in enumerate(zip(boundaries[:-1], boundaries[1:], strict=True)): + out[:, start:end] += initial[i][None, None, :] + return out + + +def segment_sum(x: torch.Tensor, cu_seqlens: torch.Tensor | None) -> torch.Tensor: + x = x.float() + if cu_seqlens is None: + return x.sum(dim=1) + boundaries = cu_seqlens.detach().cpu().tolist() + return torch.stack([x[:, start:end].sum(dim=1).squeeze(0) for start, end in zip(boundaries[:-1], boundaries[1:], strict=True)]) + + +def validate_initial_state( + initial_state: CyFAState | None, + *, + batch: int, + n_heads: int, + num_slots: int, + d_k: int, + d_v: int, + device: torch.device, +) -> CyFAState | None: + if initial_state is None: + return None + if len(initial_state) != 3: + raise ValueError("`initial_state` must contain (k, v, lambda).") + k0, v0, lambda0 = initial_state + if k0.shape != (batch, n_heads, d_k, num_slots): + raise ValueError("K initial state must have shape [B, H, Dk, M].") + if v0.shape != (batch, n_heads, num_slots, d_v): + raise ValueError("V initial state must have shape [B, H, M, Dv].") + if lambda0.shape != (batch, n_heads): + raise ValueError("Lambda initial state must have shape [B, H].") + return ( + k0.to(device=device, dtype=torch.float32).contiguous(), + v0.to(device=device, dtype=torch.float32).contiguous(), + lambda0.to(device=device, dtype=torch.float32).contiguous(), + ) diff --git a/build/torch-cuda/ops/delta_rule/chunk.py b/build/torch-cuda/ops/delta_rule/chunk.py index 0bd00789a1412e317bd7ea297d64c009f5841e09..bb310e6041586aa744a5742f1630f447d5d11cab 100644 --- a/build/torch-cuda/ops/delta_rule/chunk.py +++ b/build/torch-cuda/ops/delta_rule/chunk.py @@ -7,7 +7,7 @@ import torch -from ...modules.l2norm import l2norm_bwd, l2norm_fwd +from ...modules.norm.l2norm import l2norm_bwd, l2norm_fwd from ...ops.common.chunk_delta_h import chunk_gated_delta_rule_bwd_dhu, chunk_gated_delta_rule_fwd_h from ...ops.common.chunk_o import chunk_bwd_dqkwg, chunk_bwd_dv_local, chunk_fwd_o from ...ops.delta_rule.wy_fast import prepare_wy_repr_bwd, prepare_wy_repr_fwd, recompute_w_u_fwd @@ -275,7 +275,7 @@ def chunk_delta_rule( >>> import torch >>> import torch.nn.functional as F >>> from einops import rearrange - >>> from ...ops.delta_rule import chunk_delta_rule + >>> from fla.ops.delta_rule import chunk_delta_rule # inputs with equal lengths >>> B, T, H, K, V = 4, 2048, 4, 512, 512 >>> q = torch.randn(B, T, H, K, dtype=torch.bfloat16, device='cuda') diff --git a/build/torch-cuda/ops/delta_rule/fused_recurrent.py b/build/torch-cuda/ops/delta_rule/fused_recurrent.py index aa12c025915ea9d54860979722128ef98db025f0..c1069404885eb2c3f41da1bd01134d1ec8bd8980 100644 --- a/build/torch-cuda/ops/delta_rule/fused_recurrent.py +++ b/build/torch-cuda/ops/delta_rule/fused_recurrent.py @@ -9,7 +9,7 @@ import torch import triton import triton.language as tl -from ...modules.l2norm import l2norm_bwd, l2norm_fwd +from ...modules.norm.l2norm import l2norm_bwd, l2norm_fwd from ...utils import input_guard @@ -42,15 +42,17 @@ def fused_recurrent_delta_rule_fwd_kernel( IS_BETA_HEADWISE: tl.constexpr, IS_VARLEN: tl.constexpr, ): - i_v, i_k, i_nh = tl.program_id(0), tl.program_id(1), tl.program_id(2) + pid = tl.program_id(0).to(tl.int64) + NV, NK = tl.cdiv(V, BV), tl.cdiv(K, BK) + i_v, i_k, i_nh = pid % NV, (pid // NV) % NK, (pid // (NV * NK)).to(tl.int64) i_n, i_h = i_nh // H, i_nh % H if IS_VARLEN: bos, eos = tl.load(cu_seqlens + i_n).to(tl.int64), tl.load(cu_seqlens + i_n + 1).to(tl.int64) - all = T + all = T.to(tl.int64) T = eos - bos else: bos, eos = i_n * T, i_n * T + T - all = B * T + all = B * T.to(tl.int64) p_q = q + (bos * H + i_h) * K + i_k * BK + tl.arange(0, BK) p_k = k + (bos * H + i_h) * K + i_k * BK + tl.arange(0, BK) @@ -134,25 +136,27 @@ def fused_recurrent_delta_rule_bwd_kernel( USE_FINAL_STATE_GRADIENT: tl.constexpr, # whether to use dht IS_VARLEN: tl.constexpr, ): - i_v, i_k, i_nh = tl.program_id(0), tl.program_id(1), tl.program_id(2) + pid = tl.program_id(0).to(tl.int64) + NV = tl.cdiv(V, BV) + i_v, i_k, i_nh = pid % NV, (pid // NV) % NK, (pid // (NV * NK)).to(tl.int64) i_n, i_h = i_nh // H, i_nh % H if IS_VARLEN: bos, eos = tl.load(cu_seqlens + i_n).to(tl.int64), tl.load(cu_seqlens + i_n + 1).to(tl.int64) - all = T + all = T.to(tl.int64) T = eos - bos else: bos, eos = i_n * T, i_n * T + T - all = B * T + all = B * T.to(tl.int64) mask_k = i_k * BK + tl.arange(0, BK) < K mask_v = i_v * BV + tl.arange(0, BV) < V - p_q = q + (bos * H + i_h) * K + i_k * BK + tl.arange(0, BK) + (T - 1) * H*K - p_k = k + (bos * H + i_h) * K + i_k * BK + tl.arange(0, BK) + (T - 1) * H*K - p_v = v + (bos * H + i_h) * V + i_v * BV + tl.arange(0, BV) + (T - 1) * H*V - p_do = do + (bos * H + i_h) * V + i_v * BV + tl.arange(0, BV) + (T - 1) * H*V - p_dk = dk + ((i_v * all + bos) * H + i_h) * K + i_k * BK + tl.arange(0, BK) + (T - 1) * H*K - p_dv = dv + ((i_k * all + bos) * H + i_h) * V + i_v * BV + tl.arange(0, BV) + (T - 1) * H*V + p_q = q + (bos * H + i_h) * K + i_k * BK + tl.arange(0, BK) + (T - 1).to(tl.int64) * H*K + p_k = k + (bos * H + i_h) * K + i_k * BK + tl.arange(0, BK) + (T - 1).to(tl.int64) * H*K + p_v = v + (bos * H + i_h) * V + i_v * BV + tl.arange(0, BV) + (T - 1).to(tl.int64) * H*V + p_do = do + (bos * H + i_h) * V + i_v * BV + tl.arange(0, BV) + (T - 1).to(tl.int64) * H*V + p_dk = dk + ((i_v * all + bos) * H + i_h) * K + i_k * BK + tl.arange(0, BK) + (T - 1).to(tl.int64) * H*K + p_dv = dv + ((i_k * all + bos) * H + i_h) * V + i_v * BV + tl.arange(0, BV) + (T - 1).to(tl.int64) * H*V if IS_BETA_HEADWISE: p_beta = beta + (bos + T - 1) * H*V + i_h * V + i_v * BV + tl.arange(0, BV) p_dbeta = db + ((i_v * NK + i_k) * all + bos + T - 1) * H*V + i_h * V + tl.arange(0, BV) @@ -277,7 +281,7 @@ def fused_recurrent_delta_rule_fwd( else: final_state = None - grid = (NV, NK, N * H) + grid = (NV * NK * N * H,) u = torch.empty_like(v) fused_recurrent_delta_rule_fwd_kernel[grid]( q, @@ -333,7 +337,7 @@ def fused_recurrent_delta_rule_bwd( db = q.new_empty(NV, NK, B, T, H, V) else: db = q.new_empty(NV, B, T, H) - grid = (NV, NK, N * H) + grid = (NV * NK * N * H,) if initial_state is not None and initial_state.requires_grad: dh0 = torch.empty_like(initial_state, dtype=torch.float32) @@ -482,7 +486,7 @@ def fused_recurrent_delta_rule( >>> import torch >>> import torch.nn.functional as F >>> from einops import rearrange - >>> from ...ops.delta_rule import fused_recurrent_delta_rule + >>> from fla.ops.delta_rule import fused_recurrent_delta_rule # inputs with equal lengths >>> B, T, H, K, V = 4, 2048, 4, 512, 512 >>> q = torch.randn(B, T, H, K, device='cuda') diff --git a/build/torch-cuda/ops/delta_rule/parallel.py b/build/torch-cuda/ops/delta_rule/parallel.py index fdb47543a81a1f781556108e178e394ec487ac5e..3399bf535d676f81bb8ea8e576fac1eeb884eb80 100644 --- a/build/torch-cuda/ops/delta_rule/parallel.py +++ b/build/torch-cuda/ops/delta_rule/parallel.py @@ -42,17 +42,25 @@ def chunk_transform_qk_fwd_kernel( BT: tl.constexpr, OUTPUT_ATTENTIONS: tl.constexpr, ): - i_t, i_bh = tl.program_id(0), tl.program_id(1) - - p_q = tl.make_block_ptr(q + i_bh * T*K, (T, K), (K, 1), (i_t * BT, 0), (BT, BK), (1, 0)) - p_k = tl.make_block_ptr(k + i_bh * T*K, (T, K), (K, 1), (i_t * BT, 0), (BT, BK), (1, 0)) - p_v = tl.make_block_ptr(v + i_bh * T*V, (T, V), (V, 1), (i_t * BT, 0), (BT, BV), (1, 0)) - b_q = (tl.load(p_q, boundary_check=(0, 1)) * scale).to(p_q.dtype.element_ty) - b_k = tl.load(p_k, boundary_check=(0, 1)) - b_v = tl.load(p_v, boundary_check=(0, 1)) - - p_T = tl.make_block_ptr(A + i_bh * T * BT, (T, BT), (BT, 1), (i_t * BT, 0), (BT, BT), (1, 0)) - b_T = tl.load(p_T, boundary_check=(0, 1)) + i_t, i_bh = tl.program_id(0).to(tl.int64), tl.program_id(1).to(tl.int64) + + o_t = i_t * BT + tl.arange(0, BT) + o_k = tl.arange(0, BK) + o_v = tl.arange(0, BV) + o_T = tl.arange(0, BT) + m_b = o_t < T + m_qk = m_b[:, None] & (o_k[None, :] < K) + m_v = m_b[:, None] & (o_v[None, :] < V) + m_T = m_b[:, None] & (o_T[None, :] < BT) + p_q = q + i_bh * T*K + o_t[:, None] * K + o_k[None, :] + p_k = k + i_bh * T*K + o_t[:, None] * K + o_k[None, :] + p_v = v + i_bh * T*V + o_t[:, None] * V + o_v[None, :] + b_q = (tl.load(p_q, mask=m_qk, other=0.0) * scale).to(p_q.dtype.element_ty) + b_k = tl.load(p_k, mask=m_qk, other=0.0) + b_v = tl.load(p_v, mask=m_v, other=0.0) + + p_T = A + i_bh * T * BT + o_t[:, None] * BT + o_T[None, :] + b_T = tl.load(p_T, mask=m_T, other=0.0) o_i = tl.arange(0, BT) m_t = o_i[:, None] >= o_i[None, :] @@ -60,26 +68,26 @@ def chunk_transform_qk_fwd_kernel( m_t = o_i[:, None] > o_i[None, :] b_kk = tl.where(m_t, tl.dot(b_k, tl.trans(b_k), allow_tf32=False), 0).to(b_k.dtype) - p_beta = tl.make_block_ptr(beta + i_bh * T, (T, ), (1, ), (i_t * BT, ), (BT, ), (0, )) - b_beta = tl.load(p_beta, boundary_check=(0, )) + p_beta = beta + i_bh * T + o_t + b_beta = tl.load(p_beta, mask=m_b, other=0.0) b_k_beta = (b_k * b_beta[:, None]).to(b_k.dtype) b_qkT = tl.dot(b_qk, b_T, allow_tf32=False).to(b_k.dtype) if OUTPUT_ATTENTIONS: - p_a = tl.make_block_ptr(A_local + i_bh * T * BT, (T, BT), (BT, 1), (i_t * BT, 0), (BT, BT), (1, 0)) - tl.store(p_a, b_qkT.to(p_a.dtype.element_ty), boundary_check=(0, 1)) + p_a = A_local + i_bh * T * BT + o_t[:, None] * BT + o_T[None, :] + tl.store(p_a, b_qkT.to(p_a.dtype.element_ty), mask=m_T) b_kkT = tl.dot(b_kk, b_T, allow_tf32=False).to(b_k.dtype) - p_o = tl.make_block_ptr(o + i_bh * T*V, (T, V), (V, 1), (i_t * BT, 0), (BT, BV), (1, 0)) - tl.store(p_o, tl.dot(b_qkT, b_v).to(p_o.dtype.element_ty), boundary_check=(0, 1)) + p_o = o + i_bh * T*V + o_t[:, None] * V + o_v[None, :] + tl.store(p_o, tl.dot(b_qkT, b_v).to(p_o.dtype.element_ty), mask=m_v) - p_q_new = tl.make_block_ptr(q_new + i_bh * T*K, (T, K), (K, 1), (i_t * BT, 0), (BT, BK), (1, 0)) - tl.store(p_q_new, (b_q - tl.dot(b_qkT, b_k_beta, allow_tf32=False)).to(p_q_new.dtype.element_ty), boundary_check=(0, 1)) + p_q_new = q_new + i_bh * T*K + o_t[:, None] * K + o_k[None, :] + tl.store(p_q_new, (b_q - tl.dot(b_qkT, b_k_beta, allow_tf32=False)).to(p_q_new.dtype.element_ty), mask=m_qk) - p_k_new = tl.make_block_ptr(k_new + i_bh * T*K, (T, K), (K, 1), (i_t * BT, 0), (BT, BK), (1, 0)) + p_k_new = k_new + i_bh * T*K + o_t[:, None] * K + o_k[None, :] b_k_new = b_k - tl.dot(tl.trans(b_kkT), b_k_beta, allow_tf32=False) - tl.store(p_k_new, b_k_new.to(p_k_new.dtype.element_ty), boundary_check=(0, 1)) + tl.store(p_k_new, b_k_new.to(p_k_new.dtype.element_ty), mask=m_qk) def chunk_transform_qk_fwd( @@ -137,11 +145,16 @@ def save_intra_chunk_attn( T, BT: tl.constexpr, ): - i_t, i_bh = tl.program_id(0), tl.program_id(1) - p_A = tl.make_block_ptr(A + i_bh * T * T, (T, T), (T, 1), (i_t * BT, i_t * BT), (BT, BT), (1, 0)) - p_A_local = tl.make_block_ptr(A_local + i_bh * T * BT, (T, BT), (BT, 1), (i_t * BT, 0), (BT, BT), (1, 0)) - b_A_local = tl.load(p_A_local, boundary_check=(0, 1)) - tl.store(p_A, b_A_local.to(p_A.dtype.element_ty), boundary_check=(0, 1)) + i_t, i_bh = tl.program_id(0).to(tl.int64), tl.program_id(1).to(tl.int64) + o_t = i_t * BT + tl.arange(0, BT) + o_l = tl.arange(0, BT) + m_t = o_t < T + m_A = m_t[:, None] & (o_t[None, :] < T) + m_L = m_t[:, None] & (o_l[None, :] < BT) + p_A = A + i_bh * T * T + o_t[:, None] * T + o_t[None, :] + p_A_local = A_local + i_bh * T * BT + o_t[:, None] * BT + o_l[None, :] + b_A_local = tl.load(p_A_local, mask=m_L, other=0.0) + tl.store(p_A, b_A_local.to(p_A.dtype.element_ty), mask=m_A) @triton.heuristics({ @@ -166,72 +179,88 @@ def parallel_delta_rule_fwd_kernel( BV: tl.constexpr, OUTPUT_ATTENTIONS: tl.constexpr, ): - i_t, i_bh = tl.program_id(0), tl.program_id(1) - p_q = tl.make_block_ptr(q + i_bh * T*K, (T, K), (K, 1), (i_t * BT, 0), (BT, BK), (1, 0)) + i_t, i_bh = tl.program_id(0).to(tl.int64), tl.program_id(1).to(tl.int64) + o_t = i_t * BT + tl.arange(0, BT) + o_k = tl.arange(0, BK) + o_v = tl.arange(0, BV) + m_t = o_t < T + m_q = m_t[:, None] & (o_k[None, :] < K) + m_o = m_t[:, None] & (o_v[None, :] < V) + p_q = q + i_bh * T*K + o_t[:, None] * K + o_k[None, :] # the Q block is kept in the shared memory throughout the whole kernel # [BT, BK] b_q = tl.zeros([BT, BK], dtype=tl.float32) - b_q += tl.load(p_q, boundary_check=(0, 1)) + b_q += tl.load(p_q, mask=m_q, other=0.0) b_o = tl.zeros([BT, BV], dtype=tl.float32) - p_o = tl.make_block_ptr(o + i_bh * T*V, (T, V), (V, 1), (i_t * BT, 0), (BT, BV), (1, 0)) - b_o += tl.load(p_o, boundary_check=(0, 1)) + p_o = o + i_bh * T*V + o_t[:, None] * V + o_v[None, :] + b_o += tl.load(p_o, mask=m_o, other=0.0) # As opposed to Flashattention, this kernel requires scanning the KV blocks from right to left # Q block and K block have overlap. # masks required for offset in range((i_t + 1) * BT - 2 * BS, i_t * BT - BS, -BS): - p_k = tl.make_block_ptr(k + i_bh * T*K, (K, T), (1, K), (0, offset), (BK, BS), (0, 1)) - p_k2 = tl.make_block_ptr(k2 + i_bh * T*K, (T, K), (K, 1), (offset, 0), (BS, BK), (1, 0)) - p_v = tl.make_block_ptr(v + i_bh * T*V, (T, V), (V, 1), (offset, 0), (BS, BV), (1, 0)) - p_beta = tl.make_block_ptr(beta + i_bh * T, (T, ), (1, ), (offset, ), (BS, ), (0,)) + o_s = offset + tl.arange(0, BS) + m_k = (o_k[:, None] < K) & (o_s[None, :] < T) + m_k2 = (o_s[:, None] < T) & (o_k[None, :] < K) + m_v = (o_s[:, None] < T) & (o_v[None, :] < V) + m_a = m_t[:, None] & (o_s[None, :] < T) + p_k = k + i_bh * T*K + o_k[:, None] + o_s[None, :] * K + p_k2 = k2 + i_bh * T*K + o_s[:, None] * K + o_k[None, :] + p_v = v + i_bh * T*V + o_s[:, None] * V + o_v[None, :] + p_beta = beta + i_bh * T + o_s # [BK, BS] - b_k = tl.load(p_k, boundary_check=(0, 1)) + b_k = tl.load(p_k, mask=m_k, other=0.0) # [BS, BV] - b_v = tl.load(p_v, boundary_check=(0, 1)) + b_v = tl.load(p_v, mask=m_v, other=0.0) # [BS] - b_beta = tl.load(p_beta, boundary_check=(0,)) + b_beta = tl.load(p_beta, mask=o_s < T, other=0.0) # [BT, BS] m_s = tl.arange(0, BT) >= (offset - i_t*BT + BS) b_s = tl.dot(b_q.to(b_k.dtype), b_k, allow_tf32=False) b_s = tl.where(m_s[:, None], b_s, 0) - b_o += tl.dot(b_s.to(b_v.dtype), b_v, allow_tf32=False) - b_k2 = (tl.load(p_k2, boundary_check=(0, 1)) * b_beta[:, None]).to(b_v.dtype) + b_o = tl.dot(b_s.to(b_v.dtype), b_v, b_o, allow_tf32=False) + b_k2 = (tl.load(p_k2, mask=m_k2, other=0.0) * b_beta[:, None]).to(b_v.dtype) b_q -= tl.dot(b_s.to(b_v.dtype), b_k2, allow_tf32=False) if OUTPUT_ATTENTIONS: - p_a = tl.make_block_ptr(attn + i_bh * T * T, (T, T), (T, 1), (i_t * BT, offset), (BT, BS), (1, 0)) - tl.store(p_a, b_s.to(p_a.dtype.element_ty), boundary_check=(0, 1)) + p_a = attn + i_bh * T * T + o_t[:, None] * T + o_s[None, :] + tl.store(p_a, b_s.to(p_a.dtype.element_ty), mask=m_a) # Q block and K block have no overlap # no need for mask, thereby saving flops for offset in range(i_t * BT - BS, -BS, -BS): - p_k = tl.make_block_ptr(k + i_bh * T*K, (K, T), (1, K), (0, offset), (BK, BS), (0, 1)) - p_v = tl.make_block_ptr(v + i_bh * T*V, (T, V), (V, 1), (offset, 0), (BS, BV), (1, 0)) - p_beta = tl.make_block_ptr(beta + i_bh * T, (T, ), (1, ), (offset, ), (BS, ), (0,)) - p_k2 = tl.make_block_ptr(k2 + i_bh * T*K, (T, K), (K, 1), (offset, 0), (BS, BK), (1, 0)) + o_s = offset + tl.arange(0, BS) + m_k = (o_k[:, None] < K) & (o_s[None, :] < T) + m_k2 = (o_s[:, None] < T) & (o_k[None, :] < K) + m_v = (o_s[:, None] < T) & (o_v[None, :] < V) + m_a = m_t[:, None] & (o_s[None, :] < T) + p_k = k + i_bh * T*K + o_k[:, None] + o_s[None, :] * K + p_v = v + i_bh * T*V + o_s[:, None] * V + o_v[None, :] + p_beta = beta + i_bh * T + o_s + p_k2 = k2 + i_bh * T*K + o_s[:, None] * K + o_k[None, :] # [BK, BS] - b_k = tl.load(p_k, boundary_check=(0, 1)) + b_k = tl.load(p_k, mask=m_k, other=0.0) # [BS, BV] - b_v = tl.load(p_v, boundary_check=(0, 1)) + b_v = tl.load(p_v, mask=m_v, other=0.0) # [BS] - b_beta = tl.load(p_beta, boundary_check=(0,)) + b_beta = tl.load(p_beta, mask=o_s < T, other=0.0) # [BT, BS] b_s = (tl.dot(b_q.to(b_k.dtype), b_k, allow_tf32=False)) # [BT, BV] - b_o += tl.dot(b_s.to(b_v.dtype), b_v, allow_tf32=False) - b_k2 = (tl.load(p_k2, boundary_check=(0, 1)) * b_beta[:, None]).to(b_v.dtype) + b_o = tl.dot(b_s.to(b_v.dtype), b_v, b_o, allow_tf32=False) + b_k2 = (tl.load(p_k2, mask=m_k2, other=0.0) * b_beta[:, None]).to(b_v.dtype) b_q -= tl.dot(b_s.to(b_v.dtype), b_k2, allow_tf32=False).to(b_q.dtype) if OUTPUT_ATTENTIONS: - p_a = tl.make_block_ptr(attn + i_bh * T * T, (T, T), (T, 1), (i_t * BT, offset), (BT, BS), (1, 0)) - tl.store(p_a, b_s.to(p_a.dtype.element_ty), boundary_check=(0, 1)) + p_a = attn + i_bh * T * T + o_t[:, None] * T + o_s[None, :] + tl.store(p_a, b_s.to(p_a.dtype.element_ty), mask=m_a) - p_o_new = tl.make_block_ptr(o_new + i_bh * T*V, (T, V), (V, 1), (i_t*BT, 0), (BT, BV), (1, 0)) - tl.store(p_o_new, b_o.to(p_o.dtype.element_ty), boundary_check=(0, 1)) + p_o_new = o_new + i_bh * T*V + o_t[:, None] * V + o_v[None, :] + tl.store(p_o_new, b_o.to(p_o.dtype.element_ty), mask=m_o) class ParallelDeltaRuleFunction(torch.autograd.Function): @@ -309,7 +338,6 @@ def parallel_delta_rule( beta: torch.Tensor, scale: float | None = None, output_attentions: bool = False, - **kwargs, ) -> tuple[torch.Tensor, torch.Tensor]: r""" Args: @@ -333,10 +361,6 @@ def parallel_delta_rule( attn (torch.Tensor): Attention scores of shape `[B, H, T, T]` if `output_attentions=True` else `None`. """ - if 'head_first' in kwargs: - raise DeprecationWarning( - "head_first has been removed. Inputs must be in `[B, T, H, ...]` format.", - ) o, attn = ParallelDeltaRuleFunction.apply(q, k, v, beta, scale, output_attentions) return o, attn diff --git a/build/torch-cuda/ops/delta_rule/wy_fast.py b/build/torch-cuda/ops/delta_rule/wy_fast.py index 52f7af2641052cee489eec5d53c8dcb9628dfd11..c1cd36e8f8c4eedcb3fc7cb5888eb319750e4e8d 100644 --- a/build/torch-cuda/ops/delta_rule/wy_fast.py +++ b/build/torch-cuda/ops/delta_rule/wy_fast.py @@ -53,35 +53,43 @@ def recompute_w_u_fwd_kernel( BV: tl.constexpr, IS_VARLEN: tl.constexpr, ): - i_t, i_bh = tl.program_id(0), tl.program_id(1) + i_t, i_bh = tl.program_id(0).to(tl.int64), tl.program_id(1).to(tl.int64) i_b, i_h = i_bh // H, i_bh % H 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.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) + i_n, i_t = tl.load(chunk_indices + i_t * 2).to(tl.int32), tl.load(chunk_indices + i_t * 2 + 1).to(tl.int64) + bos, eos = tl.load(cu_seqlens + i_n).to(tl.int64), tl.load(cu_seqlens + i_n + 1).to(tl.int64) T = eos - bos else: bos, eos = i_b * T, i_b * T + T - p_beta = tl.make_block_ptr(beta + bos*H + i_h, (T,), (H,), (i_t * BT,), (BT,), (0,)) - p_A = tl.make_block_ptr(A + (bos*H + i_h) * BT, (T, BT), (H*BT, 1), (i_t * BT, 0), (BT, BT), (1, 0)) - b_beta = tl.load(p_beta, boundary_check=(0,)) - b_A = tl.load(p_A, boundary_check=(0, 1)) + o_t = i_t * BT + tl.arange(0, BT) + o_A = tl.arange(0, BT) + m_t = o_t < T + m_A = m_t[:, None] & (o_A[None, :] < BT) + p_beta = beta + bos*H + i_h + o_t * H + p_A = A + (bos*H + i_h) * BT + o_t[:, None] * (H*BT) + o_A[None, :] + b_beta = tl.load(p_beta, mask=m_t, other=0.0) + b_A = tl.load(p_A, mask=m_A, other=0.0) for i_v in range(tl.cdiv(V, BV)): - p_v = tl.make_block_ptr(v + (bos*H + i_h) * V, (T, V), (H*V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) - p_u = tl.make_block_ptr(u + (bos*H + i_h) * V, (T, V), (H*V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) - b_v = tl.load(p_v, boundary_check=(0, 1)) + o_v = i_v * BV + tl.arange(0, BV) + m_v = m_t[:, None] & (o_v[None, :] < V) + p_v = v + (bos*H + i_h) * V + o_t[:, None] * (H*V) + o_v[None, :] + p_u = u + (bos*H + i_h) * V + o_t[:, None] * (H*V) + o_v[None, :] + b_v = tl.load(p_v, mask=m_v, other=0.0) b_vb = (b_v * b_beta[:, None]).to(b_v.dtype) b_u = tl.dot(b_A.to(b_vb.dtype), b_vb, allow_tf32=False) - tl.store(p_u, (b_u).to(p_u.dtype.element_ty), boundary_check=(0, 1)) + tl.store(p_u, (b_u).to(p_u.dtype.element_ty), mask=m_v) for i_k in range(tl.cdiv(K, BK)): - p_k = tl.make_block_ptr(k + (bos*H + i_h) * K, (T, K), (H*K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) - p_w = tl.make_block_ptr(w + (bos*H + i_h) * K, (T, K), (H*K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) - b_k = tl.load(p_k, boundary_check=(0, 1)) + o_k = i_k * BK + tl.arange(0, BK) + m_k = m_t[:, None] & (o_k[None, :] < K) + p_k = k + (bos*H + i_h) * K + o_t[:, None] * (H*K) + o_k[None, :] + p_w = w + (bos*H + i_h) * K + o_t[:, None] * (H*K) + o_k[None, :] + b_k = tl.load(p_k, mask=m_k, other=0.0) b_kb = (b_k * b_beta[:, None]).to(b_k.dtype) b_w = tl.dot(b_A.to(b_kb.dtype), b_kb, allow_tf32=False) - tl.store(p_w, b_w.to(p_w.dtype.element_ty), boundary_check=(0, 1)) + tl.store(p_w, b_w.to(p_w.dtype.element_ty), mask=m_k) @triton.heuristics({ @@ -118,51 +126,59 @@ def prepare_wy_repr_bwd_kernel( BV: tl.constexpr, IS_VARLEN: tl.constexpr, ): - i_t, i_bh = tl.program_id(0), tl.program_id(1), tl.program_id(2) + i_t, i_bh = tl.program_id(0).to(tl.int64), tl.program_id(1).to(tl.int64) i_b, i_h = i_bh // H, i_bh % H 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.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) + i_n, i_t = tl.load(chunk_indices + i_t * 2).to(tl.int32), tl.load(chunk_indices + i_t * 2 + 1).to(tl.int64) + bos, eos = tl.load(cu_seqlens + i_n).to(tl.int64), tl.load(cu_seqlens + i_n + 1).to(tl.int64) T = eos - bos else: bos, eos = i_b * T, i_b * T + T - p_beta = tl.make_block_ptr(beta + bos*H + i_h, (T,), (H,), (i_t * BT,), (BT,), (0,)) - p_A = tl.make_block_ptr(A + (bos*H + i_h) * BT, (BT, T), (1, H*BT), (0, i_t * BT), (BT, BT), (0, 1)) + o_t = i_t * BT + tl.arange(0, BT) + o_A = tl.arange(0, BT) + m_t = o_t < T + m_A = (o_A[:, None] < BT) & m_t[None, :] + p_beta = beta + bos*H + i_h + o_t * H + p_A = A + (bos*H + i_h) * BT + o_A[:, None] + o_t[None, :] * (H*BT) - b_beta = tl.load(p_beta, boundary_check=(0,)) - b_A = tl.load(p_A, boundary_check=(0, 1)) + b_beta = tl.load(p_beta, mask=m_t, other=0.0) + b_A = tl.load(p_A, mask=m_A, other=0.0) b_dbeta = tl.zeros([BT], dtype=tl.float32) b_dA = tl.zeros([BT, BT], dtype=tl.float32) for i_v in range(tl.cdiv(V, BV)): - p_v = tl.make_block_ptr(v + (bos*H + i_h) * V, (T, V), (H*V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) - p_dv = tl.make_block_ptr(dv + (bos*H + i_h) * V, (T, V), (H*V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) - p_du = tl.make_block_ptr(du + (bos*H + i_h) * V, (T, V), (H*V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) + o_v = i_v * BV + tl.arange(0, BV) + m_v = m_t[:, None] & (o_v[None, :] < V) + p_v = v + (bos*H + i_h) * V + o_t[:, None] * (H*V) + o_v[None, :] + p_dv = dv + (bos*H + i_h) * V + o_t[:, None] * (H*V) + o_v[None, :] + p_du = du + (bos*H + i_h) * V + o_t[:, None] * (H*V) + o_v[None, :] - b_v = tl.load(p_v, boundary_check=(0, 1)) + b_v = tl.load(p_v, mask=m_v, other=0.0) b_v_beta = (b_v * b_beta[:, None]).to(b_v.dtype) - b_du = tl.load(p_du, boundary_check=(0, 1)) - b_dA += tl.dot(b_du, tl.trans(b_v_beta), allow_tf32=False) + b_du = tl.load(p_du, mask=m_v, other=0.0) + b_dA = tl.dot(b_du, tl.trans(b_v_beta), b_dA, allow_tf32=False) b_dv_beta = tl.dot(b_A, b_du, allow_tf32=False) b_dv = b_dv_beta * b_beta[:, None] b_dbeta += tl.sum(b_dv_beta * b_v, 1) - tl.store(p_dv, b_dv.to(p_dv.dtype.element_ty), boundary_check=(0, 1)) + tl.store(p_dv, b_dv.to(p_dv.dtype.element_ty), mask=m_v) for i_k in range(tl.cdiv(K, BK)): - p_k = tl.make_block_ptr(k + (bos*H + i_h) * K, (T, K), (H*K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) - p_dk = tl.make_block_ptr(dk + (bos*H + i_h) * K, (T, K), (H*K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) - p_dw = tl.make_block_ptr(dw + (bos*H + i_h) * K, (T, K), (H*K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) - b_k = tl.load(p_k, boundary_check=(0, 1)) + o_k = i_k * BK + tl.arange(0, BK) + m_k = m_t[:, None] & (o_k[None, :] < K) + p_k = k + (bos*H + i_h) * K + o_t[:, None] * (H*K) + o_k[None, :] + p_dk = dk + (bos*H + i_h) * K + o_t[:, None] * (H*K) + o_k[None, :] + p_dw = dw + (bos*H + i_h) * K + o_t[:, None] * (H*K) + o_k[None, :] + b_k = tl.load(p_k, mask=m_k, other=0.0) b_k_beta = (b_k * b_beta[:, None]).to(b_k.dtype) - b_dw = tl.load(p_dw, boundary_check=(0, 1)) - b_dA += tl.dot(b_dw, tl.trans(b_k_beta), allow_tf32=False) + b_dw = tl.load(p_dw, mask=m_k, other=0.0) + b_dA = tl.dot(b_dw, tl.trans(b_k_beta), b_dA, allow_tf32=False) b_dk_beta = tl.dot(b_A, b_dw, allow_tf32=False) b_dk = b_dk_beta * b_beta[:, None] b_dbeta += tl.sum(b_dk_beta * b_k, 1) - tl.store(p_dk, b_dk.to(p_dk.dtype.element_ty), boundary_check=(0, 1)) + tl.store(p_dk, b_dk.to(p_dk.dtype.element_ty), mask=m_k) b_dA = tl.where(tl.arange(0, BT)[:, None] > tl.arange(0, BT)[None, :], b_dA, 0) b_dA = tl.dot(b_dA.to(b_A.dtype), b_A) @@ -170,20 +186,22 @@ def prepare_wy_repr_bwd_kernel( b_dA = tl.where(tl.arange(0, BT)[:, None] > tl.arange(0, BT)[None, :], -b_dA, 0).to(k.dtype.element_ty) for i_k in range(tl.cdiv(K, BK)): - p_k = tl.make_block_ptr(k + (bos*H + i_h) * K, (T, K), (H*K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) - p_dk = tl.make_block_ptr(dk + (bos*H + i_h) * K, (T, K), (H*K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) - b_k = tl.load(p_k, boundary_check=(0, 1)) - b_dk = tl.load(p_dk, boundary_check=(0, 1)) + o_k = i_k * BK + tl.arange(0, BK) + m_k = m_t[:, None] & (o_k[None, :] < K) + p_k = k + (bos*H + i_h) * K + o_t[:, None] * (H*K) + o_k[None, :] + p_dk = dk + (bos*H + i_h) * K + o_t[:, None] * (H*K) + o_k[None, :] + b_k = tl.load(p_k, mask=m_k, other=0.0) + b_dk = tl.load(p_dk, mask=m_k, other=0.0) b_k_beta = (b_k * b_beta[:, None]).to(b_k.dtype) b_dk_beta = tl.dot(b_dA, b_k, allow_tf32=False) b_dbeta += tl.sum(b_dk_beta * b_k, 1) b_dk += safe_dot(tl.trans(b_dA), b_k_beta, allow_tf32=False) b_dk += b_dk_beta * b_beta[:, None] - tl.store(p_dk, b_dk.to(p_dk.dtype.element_ty), boundary_check=(0, 1)) + tl.store(p_dk, b_dk.to(p_dk.dtype.element_ty), mask=m_k) - p_dbeta = tl.make_block_ptr(dbeta + bos*H + i_h, (T,), (H,), (i_t * BT,), (BT,), (0,)) - tl.store(p_dbeta, b_dbeta.to(p_dbeta.dtype.element_ty), boundary_check=(0,)) + p_dbeta = dbeta + bos*H + i_h + o_t * H + tl.store(p_dbeta, b_dbeta.to(p_dbeta.dtype.element_ty), mask=m_t) def prepare_wy_repr_fwd( diff --git a/build/torch-cuda/ops/deltaformer/naive.py b/build/torch-cuda/ops/deltaformer/naive.py index e1398c31e4d52c53aeb5e62e5091ff23110a6e16..7eb5a23b0580a6f9d858e5b830243dd76edf6d94 100644 --- a/build/torch-cuda/ops/deltaformer/naive.py +++ b/build/torch-cuda/ops/deltaformer/naive.py @@ -39,116 +39,61 @@ def tril_softmax(scores: torch.Tensor, strict: bool = True) -> torch.Tensor: return probs -def naive_causal_attention_bhtd( - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, -) -> torch.Tensor: - B, H, T, D = q.shape - qk_scale = 1.0 / math.sqrt(D) - scores = torch.matmul(q, k.transpose(-1, -2)) * qk_scale # [B, H, T, T] - causal_mask = torch.triu(torch.ones(T, T, device=q.device), diagonal=1).bool() - scores = scores.masked_fill(causal_mask, float('-inf')) - attn_weights = torch.softmax(scores, dim=-1) # [B, H, T, T] - o = torch.matmul(attn_weights, v) # [B, H, T, D] - - return o - - -def naive_deltaformer_attn_head_first( +def naive_deltaformer_attn( q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, beta: torch.Tensor | None = None, ) -> torch.Tensor: """ - Naive reference implementation of DeltaFormer attention for head-first format. - - Two-stage process: - 1. Computes u[i] = v[i] - beta[i] * sum_{jbhts', qf, kf) * qk_scale probs = tril_softmax(scores, strict=True) # [B,H,T,T] float32 u_list = [] for t in range(T): if t == 0: - u_t = vf[:, :, t, :] + u_t = vf[:, t] else: - w = probs[:, :, t, :t] # [B,H,t] - u_prev = torch.stack(u_list, dim=-2) # [B,H,t,D] - weighted_sum = (w.unsqueeze(-1) * u_prev).sum(dim=-2) # [B,H,D] - u_t = vf[:, :, t, :] - betaf[:, :, t].unsqueeze(-1) * weighted_sum + w = probs[:, :, t, :t].transpose(1, 2) # [B,t,H] + u_prev = torch.stack(u_list, dim=1) # [B,t,H,D] + weighted_sum = (w.unsqueeze(-1) * u_prev).sum(dim=1) # [B,H,D] + u_t = vf[:, t] - betaf[:, t].unsqueeze(-1) * weighted_sum u_list.append(u_t) - u = torch.stack(u_list, dim=2) # [B,H,T,D] + u = torch.stack(u_list, dim=1) # [B,T,H,D] - o = naive_causal_attention_bhtd(q, k, u.to(orig_dtype)) + scores = torch.einsum('bthd,bshd->bhts', q, k) * qk_scale + causal_mask = torch.triu(torch.ones(T, T, device=q.device), diagonal=1).bool() + scores = scores.masked_fill(causal_mask, float('-inf')) + attn_weights = torch.softmax(scores, dim=-1) # [B,H,T,T] + o = torch.einsum('bhts,bshd->bthd', attn_weights, u.to(orig_dtype)) return o.to(orig_dtype) -def naive_deltaformer_attn( - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - beta: torch.Tensor | None = None, -) -> torch.Tensor: - """ - Naive reference implementation of DeltaFormer attention for sequence-first format. - - Args: - q: [B, T, H, D] - k: [B, T, H, D] - v: [B, T, H, D] - beta: [B, T, H] or None (defaults to ones) - - Returns: - o: [B, T, H, D] - """ - assert q.dim() == 4 and k.dim() == 4 and v.dim() == 4, "q,k,v must be [B,T,H,D]" - B, T, H, D = q.shape - assert k.shape == (B, T, H, D) and v.shape == (B, T, H, D) - - q_bhtd = q.transpose(1, 2) # [B, T, H, D] -> [B, H, T, D] - k_bhtd = k.transpose(1, 2) # [B, T, H, D] -> [B, H, T, D] - v_bhtd = v.transpose(1, 2) # [B, T, H, D] -> [B, H, T, D] - - if beta is not None: - assert beta.shape == (B, T, H) - beta_bhtd = beta.transpose(1, 2) # [B, T, H] -> [B, H, T] - else: - beta_bhtd = None - - o_bhtd = naive_deltaformer_attn_head_first(q_bhtd, k_bhtd, v_bhtd, beta_bhtd) - - o_bthd = o_bhtd.transpose(1, 2) # [B, H, T, D] -> [B, T, H, D] - - return o_bthd - - __all__ = [ 'naive_deltaformer_attn', 'tril_softmax', diff --git a/build/torch-cuda/ops/deltaformer/parallel.py b/build/torch-cuda/ops/deltaformer/parallel.py index eb60bee81d5b82cd666d2fef0913dc9f5a0c1cea..454d8f5eb40bd14ec0170037eb6bd45dcea14406 100644 --- a/build/torch-cuda/ops/deltaformer/parallel.py +++ b/build/torch-cuda/ops/deltaformer/parallel.py @@ -150,8 +150,8 @@ def parallel_deltaformer_fwd_kernel( BLOCK_C: tl.constexpr, BLOCK_T: tl.constexpr, ): - pid_c = tl.program_id(axis=0) - pid_h = tl.program_id(axis=1) + pid_c = tl.program_id(axis=0).to(tl.int64) + pid_h = tl.program_id(axis=1).to(tl.int64) rowid_block = tl.arange(0, BLOCK_C) + pid_c * BLOCK_C colid_block = tl.arange(0, BLOCK_T) @@ -160,29 +160,20 @@ def parallel_deltaformer_fwd_kernel( rowsum = tl.zeros([BLOCK_C], dtype=tl.float32) + 1 acc = tl.zeros([BLOCK_C, D], dtype=tl.float32) - q_blk_ptr = tl.make_block_ptr( - base=q_ptr + pid_h * D, - shape=(C, D), - strides=(H * D, 1), - offsets=(pid_c * BLOCK_C, 0), - block_shape=(BLOCK_C, D), - order=(1, 0), - ) - q = tl.load(q_blk_ptr, boundary_check=(0,)) + o_d = tl.arange(0, D) + m_c = rowid_block < C + q_blk_ptr = q_ptr + pid_h * D + rowid_block[:, None] * (H * D) + o_d[None, :] + q = tl.load(q_blk_ptr, mask=m_c[:, None], other=0.0) for kv_i in range(0, T, BLOCK_T): - k_blk_ptr = tl.make_block_ptr( - base=k_ptr + pid_h * D, - shape=(D, T), - strides=(1, H * D), - offsets=(0, kv_i), - block_shape=(D, BLOCK_T), - order=(0, 1), - ) - k = tl.load(k_blk_ptr, boundary_check=(1,)) + o_kv = kv_i.to(tl.int64) + colid_block + m_kv = o_kv < T + k_blk_ptr = k_ptr + pid_h * D + o_d[:, None] + o_kv[None, :] * (H * D) + k = tl.load(k_blk_ptr, mask=m_kv[None, :], other=0.0) qk = tl.dot(q, k) * qk_scale - if kv_i >= T - C: + # tiles straddling the T - C boundary need the mask too; prefix columns evaluate to false + if kv_i + BLOCK_T > T - C: mask = (T - C - kv_i + rowid_block[:, None] - colid_block[None, :] < 1) qk = tl.where(mask, -1e6, qk) @@ -197,15 +188,10 @@ def parallel_deltaformer_fwd_kernel( rowmax = rowmax_i if kv_i < T - C: - u_blk_ptr = tl.make_block_ptr( - base=u_ptr + pid_h * D, - shape=(T, D), - strides=(H * D, 1), - offsets=(kv_i, 0), - block_shape=(BLOCK_T, D), - order=(1, 0), - ) - u = tl.load(u_blk_ptr, boundary_check=(0,)) + u_blk_ptr = u_ptr + pid_h * D + o_kv[:, None] * (H * D) + o_d[None, :] + # rows at and past T - C are this launch's own output, not yet written + m_u = o_kv < T - C + u = tl.load(u_blk_ptr, mask=m_u[:, None], other=0.0) acc = tl.dot(p.to(u_ptr.dtype.element_ty), u, acc) lse = rowmax + tl.math.log2(rowsum) @@ -213,64 +199,32 @@ def parallel_deltaformer_fwd_kernel( lse_mask = rowid_block < C tl.store(lse_block_ptr, lse, mask=lse_mask) - v_ptr = tl.make_block_ptr( - base=v_ptr + pid_h * D, - shape=(C, D), - strides=(H * D, 1), - offsets=(pid_c * BLOCK_C, 0), - block_shape=(BLOCK_C, D), - order=(1, 0), - ) + v_ptr = v_ptr + pid_h * D + rowid_block[:, None] * (H * D) + o_d[None, :] acc = acc / rowsum[:, None] - beta_ptr = tl.make_block_ptr( - base=beta_ptr + pid_h, - shape=(C,), - strides=(H,), - offsets=(pid_c * BLOCK_C,), - block_shape=(BLOCK_C,), - order=(0,), - ) - beta = tl.load(beta_ptr, boundary_check=(0,)) + beta_ptr = beta_ptr + pid_h + rowid_block * H + beta = tl.load(beta_ptr, mask=m_c, other=0.0) acc = acc * beta[:, None] - v = tl.load(v_ptr, boundary_check=(0,)) + v = tl.load(v_ptr, mask=m_c[:, None], other=0.0) u = v - acc.to(v_ptr.dtype.element_ty) - u_block_ptr = tl.make_block_ptr( - base=u_ptr + pid_h * D, - shape=(T, D), - strides=(H * D, 1), - offsets=(T - C + pid_c * BLOCK_C, 0), - block_shape=(BLOCK_C, D), - order=(1, 0), - ) - tl.store(u_block_ptr, u, boundary_check=(0, 1)) + o_ur = T - C + rowid_block + u_block_ptr = u_ptr + pid_h * D + o_ur[:, None] * (H * D) + o_d[None, :] + tl.store(u_block_ptr, u, mask=(o_ur[:, None] < T) & (o_d[None, :] < D)) for kv_i in range(T - C, T, BLOCK_T): - k_blk_ptr = tl.make_block_ptr( - base=k_ptr + pid_h * D, - shape=(D, T), - strides=(1, H * D), - offsets=(0, kv_i), - block_shape=(D, BLOCK_T), - order=(0, 1), - ) - k = tl.load(k_blk_ptr, boundary_check=(1,)) + o_kv = kv_i.to(tl.int64) + colid_block + k_blk_ptr = k_ptr + pid_h * D + o_d[:, None] + o_kv[None, :] * (H * D) + k = tl.load(k_blk_ptr, mask=(o_kv < T)[None, :], other=0.0) qk = tl.dot(q, k) * qk_scale mask = (T - C - kv_i + rowid_block[:, None] - colid_block[None, :] < 1) qk -= rowmax[:, None] p = tl.math.exp2(qk) / rowsum[:, None] p = tl.where(mask, 0, p) - w_blk_ptr = tl.make_block_ptr( - base=w_ptr + pid_h * C, - shape=(C, C), - strides=(H * C, 1), - offsets=(pid_c * BLOCK_C, kv_i - (T - C)), - block_shape=(BLOCK_C, BLOCK_T), - order=(1, 0), - ) - tl.store(w_blk_ptr, p.to(w_ptr.dtype.element_ty), boundary_check=(0, 1)) + o_wc = o_kv - (T - C) + w_blk_ptr = w_ptr + pid_h * C + rowid_block[:, None] * (H * C) + o_wc[None, :] + tl.store(w_blk_ptr, p.to(w_ptr.dtype.element_ty), mask=m_c[:, None] & (o_wc[None, :] < C)) @triton.autotune(configs=_config_deltaformer(), key=['C', 'D']) @@ -290,74 +244,37 @@ def parallel_deltaformer_bwd_kernel_u( BLOCK_C: tl.constexpr, BLOCK_T: tl.constexpr, ): - pid_c = tl.program_id(axis=0) - pid_h = tl.program_id(axis=1) + pid_c = tl.program_id(axis=0).to(tl.int64) + pid_h = tl.program_id(axis=1).to(tl.int64) acc = tl.zeros([BLOCK_C, D], dtype=tl.float32) - q_blk_ptr = tl.make_block_ptr( - base=q_ptr + pid_h * D, - shape=(C, D), - strides=(H * D, 1), - offsets=(pid_c * BLOCK_C, 0), - block_shape=(BLOCK_C, D), - order=(1, 0), - ) - q = tl.load(q_blk_ptr, boundary_check=(0,)) + o_c = pid_c * BLOCK_C + tl.arange(0, BLOCK_C) + o_d = tl.arange(0, D) + m_c = o_c < C + q_blk_ptr = q_ptr + pid_h * D + o_c[:, None] * (H * D) + o_d[None, :] + q = tl.load(q_blk_ptr, mask=m_c[:, None], other=0.0) for kv_i in range(0, T, BLOCK_T): - k_blk_ptr = tl.make_block_ptr( - base=k_ptr + pid_h * D, - shape=(D, T), - strides=(1, H * D), - offsets=(0, kv_i), - block_shape=(D, BLOCK_T), - order=(0, 1), - ) - k = tl.load(k_blk_ptr, boundary_check=(1,)) + o_kv = kv_i.to(tl.int64) + tl.arange(0, BLOCK_T) + m_kv = o_kv < T + k_blk_ptr = k_ptr + pid_h * D + o_d[:, None] + o_kv[None, :] * (H * D) + k = tl.load(k_blk_ptr, mask=m_kv[None, :], other=0.0) qk = tl.dot(q, k) * fa_scale - lse_blk_ptr = tl.make_block_ptr( - base=lse_ptr + pid_h, - shape=(T,), - strides=(H,), - offsets=(kv_i,), - block_shape=(BLOCK_T,), - order=(0,), - ) - lse = tl.load(lse_blk_ptr, boundary_check=(0,)) - beta_blk_ptr = tl.make_block_ptr( - base=beta_ptr + pid_h, - shape=(T,), - strides=(H,), - offsets=(kv_i,), - block_shape=(BLOCK_T,), - order=(0,), - ) - beta = tl.load(beta_blk_ptr, boundary_check=(0,)) + lse_blk_ptr = lse_ptr + pid_h + o_kv * H + lse = tl.load(lse_blk_ptr, mask=m_kv, other=0.0) + beta_blk_ptr = beta_ptr + pid_h + o_kv * H + beta = tl.load(beta_blk_ptr, mask=m_kv, other=0.0) p = tl.math.exp2(qk - lse[None, :]) * beta[None, :] - v_blk_ptr = tl.make_block_ptr( - base=v_ptr + pid_h * D, - shape=(T, D), - strides=(H * D, 1), - offsets=(kv_i, 0), - block_shape=(BLOCK_T, D), - order=(1, 0), - ) - v = tl.load(v_blk_ptr, boundary_check=(0,)) + v_blk_ptr = v_ptr + pid_h * D + o_kv[:, None] * (H * D) + o_d[None, :] + v = tl.load(v_blk_ptr, mask=m_kv[:, None], other=0.0) acc = tl.dot(p.to(v_ptr.dtype.element_ty), v, acc) - o_blk_ptr = tl.make_block_ptr( - base=o_ptr + pid_h * D, - shape=(C, D), - strides=(H * D, 1), - offsets=(pid_c * BLOCK_C, 0), - block_shape=(BLOCK_C, D), - order=(1, 0), - ) - tl.store(o_blk_ptr, acc.to(o_ptr.dtype.element_ty), boundary_check=(0,)) + o_blk_ptr = o_ptr + pid_h * D + o_c[:, None] * (H * D) + o_d[None, :] + tl.store(o_blk_ptr, acc.to(o_ptr.dtype.element_ty), mask=m_c[:, None]) @triton.autotune(configs=_config_deltaformer(), key=['T', 'D']) @@ -376,79 +293,41 @@ def parallel_deltaformer_bwd_kernel_row_sum( BLOCK_C: tl.constexpr, BLOCK_T: tl.constexpr, ): - pid_c = tl.program_id(axis=0) - pid_h = tl.program_id(axis=1) + pid_c = tl.program_id(axis=0).to(tl.int64) + pid_h = tl.program_id(axis=1).to(tl.int64) rowid_block = tl.arange(0, BLOCK_C) + pid_c * BLOCK_C colid_block = tl.arange(0, BLOCK_T) acc = tl.zeros([BLOCK_C], dtype=tl.float32) - k_row_blk_ptr = tl.make_block_ptr( - base=q_ptr + pid_h * D, - shape=(T, D), - strides=(H * D, 1), - offsets=(pid_c * BLOCK_C, 0), - block_shape=(BLOCK_C, D), - order=(1, 0), - ) - k_row = tl.load(k_row_blk_ptr, boundary_check=(0,)) - lse_blk_ptr = tl.make_block_ptr( - base=lse_ptr + pid_h, - shape=(T,), - strides=(H,), - offsets=(pid_c * BLOCK_C,), - block_shape=(BLOCK_C,), - order=(0,), - ) - lse = tl.load(lse_blk_ptr, boundary_check=(0,)) - grad_v_blk_ptr = tl.make_block_ptr( - base=grad_v_ptr + pid_h * D, - shape=(T, D), - strides=(H * D, 1), - offsets=(pid_c * BLOCK_C, 0), - block_shape=(BLOCK_C, D), - order=(1, 0), - ) - grad_v_row = -tl.load(grad_v_blk_ptr, boundary_check=(0,)) + o_d = tl.arange(0, D) + m_r = rowid_block < T + k_row_blk_ptr = q_ptr + pid_h * D + rowid_block[:, None] * (H * D) + o_d[None, :] + k_row = tl.load(k_row_blk_ptr, mask=m_r[:, None], other=0.0) + lse_blk_ptr = lse_ptr + pid_h + rowid_block * H + lse = tl.load(lse_blk_ptr, mask=m_r, other=0.0) + grad_v_blk_ptr = grad_v_ptr + pid_h * D + rowid_block[:, None] * (H * D) + o_d[None, :] + grad_v_row = -tl.load(grad_v_blk_ptr, mask=m_r[:, None], other=0.0) for kv_i in range(0, (pid_c + 1) * BLOCK_C, BLOCK_T): - k_blk_ptr = tl.make_block_ptr( - base=k_ptr + pid_h * D, - shape=(D, T), - strides=(1, H * D), - offsets=(0, kv_i), - block_shape=(D, BLOCK_T), - order=(0, 1), - ) - k = tl.load(k_blk_ptr, boundary_check=(1,)) + o_kv = kv_i.to(tl.int64) + colid_block + m_kv = o_kv < T + k_blk_ptr = k_ptr + pid_h * D + o_d[:, None] + o_kv[None, :] * (H * D) + k = tl.load(k_blk_ptr, mask=m_kv[None, :], other=0.0) qk = tl.dot(k_row, k) * fa_scale p = tl.math.exp2(qk - lse[:, None]) - u_blk_ptr = tl.make_block_ptr( - base=u_ptr + pid_h * D, - shape=(D, T), - strides=(1, H * D), - offsets=(0, kv_i), - block_shape=(D, BLOCK_T), - order=(0, 1), - ) - ut = tl.load(u_blk_ptr, boundary_check=(1,)) + u_blk_ptr = u_ptr + pid_h * D + o_d[:, None] + o_kv[None, :] * (H * D) + ut = tl.load(u_blk_ptr, mask=m_kv[None, :], other=0.0) dp = tl.dot(grad_v_row, ut) if kv_i + BLOCK_T >= pid_c * BLOCK_C: mask = (rowid_block[:, None] <= colid_block[None, :] + kv_i) p = tl.where(mask, 0., p) dp = tl.where(mask, 0., dp) acc += tl.sum(p * dp, axis=1) - row_dot_block_ptr = tl.make_block_ptr( - base=row_dot_ptr + pid_h, - shape=(T,), - strides=(H,), - offsets=(pid_c * BLOCK_C,), - block_shape=(BLOCK_C,), - order=(0,), - ) - tl.store(row_dot_block_ptr, acc, boundary_check=(0,)) + row_dot_block_ptr = row_dot_ptr + pid_h + rowid_block * H + tl.store(row_dot_block_ptr, acc, mask=m_r) @triton.autotune(configs=[triton.Config({'BLOCK_C': BC}, num_stages=ns, num_warps=nw) @@ -473,104 +352,45 @@ def parallel_deltaformer_bwd_kernel_qk( qk_scale: tl.constexpr, BLOCK_C: tl.constexpr, ): - pid_c = tl.program_id(axis=0) - pid_h = tl.program_id(axis=1) + pid_c = tl.program_id(axis=0).to(tl.int64) + pid_h = tl.program_id(axis=1).to(tl.int64) block_i = tl.arange(0, BLOCK_C) acc = tl.zeros([BLOCK_C, D], dtype=tl.float32) - k_row_blk_ptr = tl.make_block_ptr( - base=q_ptr + pid_h * D, - shape=(T, D), - strides=(H * D, 1), - offsets=(pid_c * BLOCK_C, 0), - block_shape=(BLOCK_C, D), - order=(1, 0), - ) - k_row = tl.load(k_row_blk_ptr, boundary_check=(0,)) - lse_blk_ptr = tl.make_block_ptr( - base=lse_ptr + pid_h, - shape=(T,), - strides=(H,), - offsets=(pid_c * BLOCK_C,), - block_shape=(BLOCK_C,), - order=(0,), - ) - lse = tl.load(lse_blk_ptr, boundary_check=(0,)) - beta_blk_ptr = tl.make_block_ptr( - base=beta_ptr + pid_h, - shape=(T,), - strides=(H,), - offsets=(pid_c * BLOCK_C,), - block_shape=(BLOCK_C,), - order=(0,), - ) - beta = tl.load(beta_blk_ptr, boundary_check=(0,)) - grad_v_blk_ptr = tl.make_block_ptr( - base=grad_v_ptr + pid_h * D, - shape=(T, D), - strides=(H * D, 1), - offsets=(pid_c * BLOCK_C, 0), - block_shape=(BLOCK_C, D), - order=(1, 0), - ) - grad_v_row = -tl.load(grad_v_blk_ptr, boundary_check=(0,)) - row_dot_blk_ptr = tl.make_block_ptr( - base=row_dot_ptr + pid_h, - shape=(T,), - strides=(H,), - offsets=(pid_c * BLOCK_C,), - block_shape=(BLOCK_C,), - order=(0,), - ) - row_dot_row = tl.load(row_dot_blk_ptr, boundary_check=(0,)).to(k_ptr.dtype.element_ty) + o_c = pid_c * BLOCK_C + block_i + o_d = tl.arange(0, D) + m_c = o_c < T + k_row_blk_ptr = q_ptr + pid_h * D + o_c[:, None] * (H * D) + o_d[None, :] + k_row = tl.load(k_row_blk_ptr, mask=m_c[:, None], other=0.0) + lse_blk_ptr = lse_ptr + pid_h + o_c * H + lse = tl.load(lse_blk_ptr, mask=m_c, other=0.0) + beta_blk_ptr = beta_ptr + pid_h + o_c * H + beta = tl.load(beta_blk_ptr, mask=m_c, other=0.0) + grad_v_blk_ptr = grad_v_ptr + pid_h * D + o_c[:, None] * (H * D) + o_d[None, :] + grad_v_row = -tl.load(grad_v_blk_ptr, mask=m_c[:, None], other=0.0) + row_dot_blk_ptr = row_dot_ptr + pid_h + o_c * H + row_dot_row = tl.load(row_dot_blk_ptr, mask=m_c, other=0.0).to(k_ptr.dtype.element_ty) for kv_i in range(0, pid_c * BLOCK_C, BLOCK_C): - k_blk_ptr = tl.make_block_ptr( - base=k_ptr + pid_h * D, - shape=(D, T), - strides=(1, H * D), - offsets=(0, kv_i), - block_shape=(D, BLOCK_C), - order=(0, 1), - ) - kt = tl.load(k_blk_ptr, boundary_check=(1,)) + o_kv = kv_i.to(tl.int64) + block_i + k_blk_ptr = k_ptr + pid_h * D + o_d[:, None] + o_kv[None, :] * (H * D) + kt = tl.load(k_blk_ptr, mask=(o_kv < T)[None, :], other=0.0) qk = tl.dot(k_row, kt) * fa_scale p = tl.math.exp2(qk - lse[:, None]) * beta[:, None] - u_blk_ptr = tl.make_block_ptr( - base=u_ptr + pid_h * D, - shape=(D, T), - strides=(1, H * D), - offsets=(0, kv_i), - block_shape=(D, BLOCK_C), - order=(0, 1), - ) + u_blk_ptr = u_ptr + pid_h * D + o_d[:, None] + o_kv[None, :] * (H * D) ut = tl.load(u_blk_ptr) dp = tl.dot(grad_v_row, ut) da = p * (dp - row_dot_row[:, None]) k = tl.trans(kt, 1, 0) acc = tl.dot(da.to(k.dtype), k, acc) - k_row_blk_ptr = tl.make_block_ptr( - base=k_ptr + pid_h * D, - shape=(T, D), - strides=(H * D, 1), - offsets=(pid_c * BLOCK_C, 0), - block_shape=(BLOCK_C, D), - order=(1, 0), - ) - k_row_true = tl.load(k_row_blk_ptr, boundary_check=(0,)) + k_row_blk_ptr = k_ptr + pid_h * D + o_c[:, None] * (H * D) + o_d[None, :] + k_row_true = tl.load(k_row_blk_ptr, mask=m_c[:, None], other=0.0) qk = tl.dot(k_row, tl.trans(k_row_true, 1, 0)) * fa_scale p = tl.math.exp2(qk - lse[:, None]) * beta[:, None] - u_blk_ptr = tl.make_block_ptr( - base=u_ptr + pid_h * D, - shape=(D, T), - strides=(1, H * D), - offsets=(0, pid_c * BLOCK_C), - block_shape=(D, BLOCK_C), - order=(0, 1), - ) + u_blk_ptr = u_ptr + pid_h * D + o_d[:, None] + o_c[None, :] * (H * D) ut = tl.load(u_blk_ptr) dp = tl.dot(grad_v_row, ut) dpm = dp - row_dot_row[:, None] @@ -581,85 +401,38 @@ def parallel_deltaformer_bwd_kernel_qk( daat = da acc = tl.dot(daat.to(k_row.dtype), k_row_true, acc) - grad_q_blk_ptr = tl.make_block_ptr( - base=grad_q_ptr + pid_h * D, - shape=(T, D), - strides=(H * D, 1), - offsets=(BLOCK_C * pid_c, 0), - block_shape=(BLOCK_C, D), - order=(1, 0), - ) + grad_q_blk_ptr = grad_q_ptr + pid_h * D + o_c[:, None] * (H * D) + o_d[None, :] acc = acc * qk_scale - tl.store(grad_q_blk_ptr, acc.to(grad_q_ptr.dtype.element_ty), boundary_check=(0,)) + tl.store(grad_q_blk_ptr, acc.to(grad_q_ptr.dtype.element_ty), mask=m_c[:, None]) daat = tl.trans(da, 1, 0) acc = tl.dot(daat.to(k_row.dtype), k_row) k_row = k_row_true nu = -tl.trans(ut, 1, 0) for kv_i in range((pid_c + 1) * BLOCK_C, T, BLOCK_C): - k_blk_ptr = tl.make_block_ptr( - base=q_ptr + pid_h * D, - shape=(D, T), - strides=(1, H * D), - offsets=(0, kv_i), - block_shape=(D, BLOCK_C), - order=(0, 1), - ) - kt = tl.load(k_blk_ptr, boundary_check=(1,)) - lse_blk_ptr = tl.make_block_ptr( - base=lse_ptr + pid_h, - shape=(T,), - strides=(H,), - offsets=(kv_i,), - block_shape=(BLOCK_C,), - order=(0,), - ) - lse = tl.load(lse_blk_ptr, boundary_check=(0,)) - beta_blk_ptr = tl.make_block_ptr( - base=beta_ptr + pid_h, - shape=(T,), - strides=(H,), - offsets=(kv_i,), - block_shape=(BLOCK_C,), - order=(0,), - ) - beta = tl.load(beta_blk_ptr, boundary_check=(0,)) + o_kv = kv_i.to(tl.int64) + block_i + m_kv = o_kv < T + k_blk_ptr = q_ptr + pid_h * D + o_d[:, None] + o_kv[None, :] * (H * D) + kt = tl.load(k_blk_ptr, mask=m_kv[None, :], other=0.0) + lse_blk_ptr = lse_ptr + pid_h + o_kv * H + lse = tl.load(lse_blk_ptr, mask=m_kv, other=0.0) + beta_blk_ptr = beta_ptr + pid_h + o_kv * H + beta = tl.load(beta_blk_ptr, mask=m_kv, other=0.0) qk = tl.dot(k_row, kt) * fa_scale p = tl.math.exp2(qk - lse[None, :]) * beta[None, :] - grad_vt_blk_ptr = tl.make_block_ptr( - base=grad_v_ptr + pid_h * D, - shape=(D, T), - strides=(1, H * D), - offsets=(0, kv_i), - block_shape=(D, BLOCK_C), - order=(0, 1), - ) - grad_vt = tl.load(grad_vt_blk_ptr, boundary_check=(1,)) - row_dot_blk_ptr = tl.make_block_ptr( - base=row_dot_ptr + pid_h, - shape=(T,), - strides=(H,), - offsets=(kv_i,), - block_shape=(BLOCK_C,), - order=(0,), - ) - row_dot = tl.load(row_dot_blk_ptr, boundary_check=(0,)).to(k_ptr.dtype.element_ty) + grad_vt_blk_ptr = grad_v_ptr + pid_h * D + o_d[:, None] + o_kv[None, :] * (H * D) + grad_vt = tl.load(grad_vt_blk_ptr, mask=m_kv[None, :], other=0.0) + row_dot_blk_ptr = row_dot_ptr + pid_h + o_kv * H + row_dot = tl.load(row_dot_blk_ptr, mask=m_kv, other=0.0).to(k_ptr.dtype.element_ty) dp = tl.dot(nu, grad_vt) da = p * (dp - row_dot[None, :]) k = tl.trans(kt, 1, 0) acc = tl.dot(da.to(k.dtype), k, acc) - grad_k_blk_ptr = tl.make_block_ptr( - base=grad_k_ptr + pid_h * D, - shape=(T, D), - strides=(H * D, 1), - offsets=(BLOCK_C * pid_c, 0), - block_shape=(BLOCK_C, D), - order=(1, 0), - ) + grad_k_blk_ptr = grad_k_ptr + pid_h * D + o_c[:, None] * (H * D) + o_d[None, :] acc = acc * qk_scale - tl.store(grad_k_blk_ptr, acc.to(grad_k_ptr.dtype.element_ty), boundary_check=(0,)) + tl.store(grad_k_blk_ptr, acc.to(grad_k_ptr.dtype.element_ty), mask=m_c[:, None]) class ParallelDeltaformerFunction(torch.autograd.Function): diff --git a/build/torch-cuda/ops/dsa/naive.py b/build/torch-cuda/ops/dsa/naive.py index 0f9fb890100858efecb477f0179a18503829bd81..72e6312412989c1a18e431c58103a391c25d660f 100644 --- a/build/torch-cuda/ops/dsa/naive.py +++ b/build/torch-cuda/ops/dsa/naive.py @@ -155,7 +155,10 @@ def naive_dsa( assert q.shape[0] == 1, "batch size must be 1 when cu_seqlens are provided" dtype = q.dtype - G = q.shape[2] // k.shape[2] + HQ, H = q.shape[2], k.shape[2] + if H == 0 or HQ % H != 0: + raise ValueError(f"The number of query heads ({HQ}) must be divisible by the number of key/value heads ({H}).") + G = HQ // H q, k, v = (x.float() for x in (q, k, v)) k, v = (repeat(x, 'b t h d -> b t (h g) d', g=G) for x in (k, v)) diff --git a/build/torch-cuda/ops/forgetting_attn/naive.py b/build/torch-cuda/ops/forgetting_attn/naive.py index cf5612ad941075eda96004929aa4cf94d2038c19..201e9b5bb280ede951d4eb24f7d0369d96ea9103 100644 --- a/build/torch-cuda/ops/forgetting_attn/naive.py +++ b/build/torch-cuda/ops/forgetting_attn/naive.py @@ -35,6 +35,8 @@ def naive_forgetting_attn( """ _, T, HQ, D = q.shape H = k.shape[2] + if H == 0 or HQ % H != 0: + raise ValueError(f"The number of query heads ({HQ}) must be divisible by the number of key/value heads ({H}).") G = HQ // H if scale is None: diff --git a/build/torch-cuda/ops/forgetting_attn/parallel.py b/build/torch-cuda/ops/forgetting_attn/parallel.py index 594263a317ab6f8323fbe67d7689db68809dfbd8..9739558cc67b4400919723b8e57a528dc1cbdf58 100644 --- a/build/torch-cuda/ops/forgetting_attn/parallel.py +++ b/build/torch-cuda/ops/forgetting_attn/parallel.py @@ -18,7 +18,6 @@ def parallel_forgetting_attn( scale: float | None = None, window_size: int | None = None, cu_seqlens: torch.LongTensor | None = None, - **kwargs ) -> torch.Tensor: r""" Args: @@ -46,10 +45,6 @@ def parallel_forgetting_attn( o (torch.Tensor): Outputs of shape `[B, T, HQ, V]`. """ - if 'head_first' in kwargs: - raise DeprecationWarning( - "head_first has been removed. Inputs must be in `[B, T, H, ...]` format.", - ) if scale is None: scale = k.shape[-1] ** -0.5 if cu_seqlens is not None and q.shape[0] != 1: diff --git a/build/torch-cuda/ops/gated_delta_product/chunk.py b/build/torch-cuda/ops/gated_delta_product/chunk.py index 43a28158f6a04babbb02ba3bbbd4870e5c9be144..fd6b4a0977629c3c047e97180f9ac0531a93fdd0 100644 --- a/build/torch-cuda/ops/gated_delta_product/chunk.py +++ b/build/torch-cuda/ops/gated_delta_product/chunk.py @@ -10,6 +10,8 @@ from einops import rearrange from ...modules.l2norm import l2norm_bwd, l2norm_fwd from ...ops.common.chunk_scaled_dot_kkt import chunk_scaled_dot_kkt_fwd +from ...ops.cp import FLACPContext +from ...ops.cp.chunk_delta_h import chunk_gated_delta_rule_fwd_h_pre_process, compress_h0 from ...ops.delta_rule.chunk import chunk_delta_rule_bwd from ...ops.delta_rule.wy_fast import recompute_w_u_fwd as dn_recompute_w_u_fwd from ...ops.gated_delta_product.chunk_deltaproduct_h import chunk_gated_delta_product_fwd_h @@ -35,6 +37,7 @@ def chunk_gated_delta_product_fwd( num_householder: int = 1, chunk_indices: torch.LongTensor | None = None, chunk_indices_dp: torch.LongTensor | None = None, + cp_context: FLACPContext | None = None, ): cu_seqlens_dp = cu_seqlens * num_householder if cu_seqlens is not None else None if g is not None: @@ -94,6 +97,17 @@ def chunk_gated_delta_product_fwd( cu_seqlens=cu_seqlens_dp, chunk_indices=chunk_indices_dp, ) + if cp_context is not None: + initial_state = chunk_gated_delta_rule_fwd_h_pre_process( + k=k, + w=w, + u=u, + g=g_interleaved, + cu_seqlens=cu_seqlens_dp, + initial_state=initial_state, + context=cp_context, + ) + h, v_new, final_state = chunk_gated_delta_product_fwd_h( k=k, w=w, @@ -105,6 +119,9 @@ def chunk_gated_delta_product_fwd( num_householder=num_householder, chunk_indices=chunk_indices, ) + if cp_context is not None: + initial_state = compress_h0(initial_state, context=cp_context) + o = chunk_gated_delta_product_fwd_o( q=q, k=k, @@ -116,7 +133,7 @@ def chunk_gated_delta_product_fwd( num_householder=num_householder, chunk_indices=chunk_indices, ) - return g, g_interleaved, o, A, final_state + return g, g_interleaved, o, A, final_state, initial_state class ChunkGatedDeltaProductFunction(torch.autograd.Function): @@ -138,6 +155,7 @@ class ChunkGatedDeltaProductFunction(torch.autograd.Function): use_qk_l2norm_in_kernel: bool = False, cu_seqlens: torch.LongTensor | None = None, cu_seqlens_cpu: torch.LongTensor | None = None, + cp_context: FLACPContext | None = None, ): if use_qk_l2norm_in_kernel: q, q_rstd = l2norm_fwd(q) @@ -156,7 +174,7 @@ class ChunkGatedDeltaProductFunction(torch.autograd.Function): cu_seqlens * num_householder, 64, cu_seqlens_cpu=cu_seqlens_cpu_dp ) - g, g_interleaved, o, A, final_state = chunk_gated_delta_product_fwd( + g, g_interleaved, o, A, final_state, initial_state = chunk_gated_delta_product_fwd( q=q, k=k, v=v, @@ -169,6 +187,7 @@ class ChunkGatedDeltaProductFunction(torch.autograd.Function): num_householder=num_householder, chunk_indices=chunk_indices, chunk_indices_dp=chunk_indices_dp, + cp_context=cp_context, ) ctx.save_for_backward( q, @@ -186,6 +205,7 @@ class ChunkGatedDeltaProductFunction(torch.autograd.Function): ctx.scale = scale ctx.use_qk_l2norm_in_kernel = use_qk_l2norm_in_kernel ctx.num_householder = num_householder + ctx.cp_context = cp_context return o.to(q.dtype), final_state @staticmethod @@ -230,6 +250,7 @@ class ChunkGatedDeltaProductFunction(torch.autograd.Function): do=do, dht=dht, cu_seqlens=cu_seqlens * ctx.num_householder if cu_seqlens is not None else None, + cp_context=ctx.cp_context, chunk_indices=chunk_indices_dp, ) dg = rearrange(dg, 'b (l n) h -> b l n h ', n=ctx.num_householder)[:, :, 0].contiguous().to(g) @@ -252,7 +273,7 @@ class ChunkGatedDeltaProductFunction(torch.autograd.Function): if ctx.use_qk_l2norm_in_kernel: dq = l2norm_bwd(q_org, q_rstd, dq) dk = l2norm_bwd(k, k_rstd, dk) - return dq.to(q), dk.to(k), dv.to(v), dg, db.to(beta), None, None, dh0, None, None, None, None + return dq.to(q), dk.to(k), dv.to(v), dg, db.to(beta), None, None, dh0, None, None, None, None, None @torch.compiler.disable @@ -269,6 +290,7 @@ def chunk_gated_delta_product( use_qk_l2norm_in_kernel: bool = False, cu_seqlens: torch.LongTensor | None = None, cu_seqlens_cpu: torch.LongTensor | None = None, + cp_context: FLACPContext | None = None, ): r""" Args: @@ -299,6 +321,11 @@ def chunk_gated_delta_product( cu_seqlens (torch.LongTensor): Cumulative sequence lengths of shape `[N+1]` used for variable-length training, consistent with the FlashAttention API. + cp_context (FLACPContext, Optional): + Context parallel context for distributed training across multiple devices. + When provided, `g` is required, `initial_state` and `output_final_state` + are not supported, and `cu_seqlens` will be overridden by the context. + Default: `None`. Returns: o (torch.Tensor): @@ -310,7 +337,7 @@ def chunk_gated_delta_product( >>> import torch >>> import torch.nn.functional as F >>> from einops import rearrange - >>> from ...ops.gated_delta_rule import chunk_gated_delta_product + >>> from fla.ops.gated_delta_rule import chunk_gated_delta_product # inputs with equal lengths >>> B, T, H, K, V = 4, 2048, 4, 512, 512 >>> q = torch.randn(B, T, H, K, dtype=torch.bfloat16, device='cuda') @@ -343,6 +370,15 @@ def chunk_gated_delta_product( if g is not None: assert g.shape == (B, T, H) + if cp_context is not None: + assert g is not None, "Non-gated GDP is not supported for CP" + 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" + assert cp_context.cu_seqlens is not None, "cu_seqlens is required for CP" + cu_seqlens = cp_context.cu_seqlens + if cp_context.cu_seqlens_cpu is not None: + cu_seqlens_cpu = cp_context.cu_seqlens_cpu + if cu_seqlens is not None: if q.shape[0] != 1: raise ValueError( @@ -369,5 +405,6 @@ def chunk_gated_delta_product( use_qk_l2norm_in_kernel, cu_seqlens, cu_seqlens_cpu, + cp_context, ) return o, final_state diff --git a/build/torch-cuda/ops/gated_delta_product/chunk_deltaproduct_h.py b/build/torch-cuda/ops/gated_delta_product/chunk_deltaproduct_h.py index 6ffdc626f0c6dc03a87883c4aa98538377578dfc..5d222afbf406601fbb423db066b5d4085cb31e3e 100644 --- a/build/torch-cuda/ops/gated_delta_product/chunk_deltaproduct_h.py +++ b/build/torch-cuda/ops/gated_delta_product/chunk_deltaproduct_h.py @@ -10,7 +10,7 @@ import triton import triton.language as tl from ...ops.utils import prepare_chunk_indices, prepare_chunk_offsets -from ...ops.utils.op import exp2 +from ...ops.utils.op import exp2, unflatten_program_id from ...utils import IS_NVIDIA_HOPPER, autotune_cache_kwargs NUM_WARPS = [2, 4] if IS_NVIDIA_HOPPER else [2, 4, 8, 16] @@ -58,13 +58,13 @@ def chunk_gated_delta_product_fwd_kernel_h_blockdim64( SAVE_NEW_VALUE: tl.constexpr, IS_VARLEN: tl.constexpr, ): - i_v, i_nh = tl.program_id(0), tl.program_id(1) + i_v, i_nh = unflatten_program_id(tl.cdiv(V, BV)) i_n, i_h = i_nh // H, i_nh % H if IS_VARLEN: - bos, eos = tl.load(cu_seqlens + i_n).to(tl.int32), tl.load(cu_seqlens + i_n + 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 NT = tl.cdiv(T, BT) - boh = tl.load(chunk_offsets + i_n).to(tl.int32) + boh = tl.load(chunk_offsets + i_n).to(tl.int64) else: bos, eos = i_n * T, i_n * T + T NT = tl.cdiv(T, BT) @@ -94,67 +94,81 @@ def chunk_gated_delta_product_fwd_kernel_h_blockdim64( if STORE_FINAL_STATE: ht = ht + i_nh * K*V + o_k1 = tl.arange(0, 64) + o_k2 = 64 + tl.arange(0, 64) + o_k3 = 128 + tl.arange(0, 64) + o_k4 = 192 + tl.arange(0, 64) + o_v = i_v * BV + tl.arange(0, BV) + m_v = o_v < V + m_h1 = (o_k1[:, None] < K) & m_v[None, :] + m_h2 = (o_k2[:, None] < K) & m_v[None, :] + m_h3 = (o_k3[:, None] < K) & m_v[None, :] + m_h4 = (o_k4[:, None] < K) & m_v[None, :] + # load initial state if USE_INITIAL_STATE: - p_h0_1 = tl.make_block_ptr(h0, (K, V), (V, 1), (0, i_v * BV), (64, BV), (1, 0)) - b_h1 += tl.load(p_h0_1, boundary_check=(0, 1)).to(tl.float32) + p_h0_1 = h0 + o_k1[:, None] * V + o_v[None, :] + b_h1 += tl.load(p_h0_1, mask=m_h1, other=0.0).to(tl.float32) if K > 64: - p_h0_2 = tl.make_block_ptr(h0, (K, V), (V, 1), (64, i_v * BV), (64, BV), (1, 0)) - b_h2 += tl.load(p_h0_2, boundary_check=(0, 1)).to(tl.float32) + p_h0_2 = h0 + o_k2[:, None] * V + o_v[None, :] + b_h2 += tl.load(p_h0_2, mask=m_h2, other=0.0).to(tl.float32) if K > 128: - p_h0_3 = tl.make_block_ptr(h0, (K, V), (V, 1), (128, i_v * BV), (64, BV), (1, 0)) - b_h3 += tl.load(p_h0_3, boundary_check=(0, 1)).to(tl.float32) + p_h0_3 = h0 + o_k3[:, None] * V + o_v[None, :] + b_h3 += tl.load(p_h0_3, mask=m_h3, other=0.0).to(tl.float32) if K > 192: - p_h0_4 = tl.make_block_ptr(h0, (K, V), (V, 1), (192, i_v * BV), (64, BV), (1, 0)) - b_h4 += tl.load(p_h0_4, boundary_check=(0, 1)).to(tl.float32) + p_h0_4 = h0 + o_k4[:, None] * V + o_v[None, :] + b_h4 += tl.load(p_h0_4, mask=m_h4, other=0.0).to(tl.float32) # main recurrence for i_t in range(NT): + i_t_int64 = i_t.to(tl.int64) if i_t % num_householder == 0: - i_t_true = i_t // num_householder - p_h1 = tl.make_block_ptr(h + i_t_true * stride_h, (K, V), (V, 1), (0, i_v * BV), (64, BV), (1, 0)) - tl.store(p_h1, b_h1.to(p_h1.dtype.element_ty), boundary_check=(0, 1)) + i_t_true = i_t_int64 // num_householder + p_h1 = h + i_t_true * stride_h + o_k1[:, None] * V + o_v[None, :] + tl.store(p_h1, b_h1.to(p_h1.dtype.element_ty), mask=m_h1) if K > 64: - p_h2 = tl.make_block_ptr(h + i_t_true * stride_h, (K, V), (V, 1), (64, i_v * BV), (64, BV), (1, 0)) - tl.store(p_h2, b_h2.to(p_h2.dtype.element_ty), boundary_check=(0, 1)) + p_h2 = h + i_t_true * stride_h + o_k2[:, None] * V + o_v[None, :] + tl.store(p_h2, b_h2.to(p_h2.dtype.element_ty), mask=m_h2) if K > 128: - p_h3 = tl.make_block_ptr(h + i_t_true * stride_h, (K, V), (V, 1), (128, i_v * BV), (64, BV), (1, 0)) - tl.store(p_h3, b_h3.to(p_h3.dtype.element_ty), boundary_check=(0, 1)) + p_h3 = h + i_t_true * stride_h + o_k3[:, None] * V + o_v[None, :] + tl.store(p_h3, b_h3.to(p_h3.dtype.element_ty), mask=m_h3) if K > 192: - p_h4 = tl.make_block_ptr(h + i_t_true * stride_h, (K, V), (V, 1), (192, i_v * BV), (64, BV), (1, 0)) - tl.store(p_h4, b_h4.to(p_h4.dtype.element_ty), boundary_check=(0, 1)) - - p_v = tl.make_block_ptr(v, (T, V), (stride_v, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) - p_v_new = tl.make_block_ptr(v_new, (T, V), (stride_v, 1), (i_t * BT, i_v * BV), - (BT, BV), (1, 0)) if SAVE_NEW_VALUE else None + p_h4 = h + i_t_true * stride_h + o_k4[:, None] * V + o_v[None, :] + tl.store(p_h4, b_h4.to(p_h4.dtype.element_ty), mask=m_h4) + + o_t = i_t_int64 * BT + tl.arange(0, BT) + m_t = o_t < T + m_vv = m_t[:, None] & m_v[None, :] + p_v = v + o_t[:, None] * stride_v + o_v[None, :] + p_v_new = v_new + o_t[:, None] * stride_v + o_v[None, :] if SAVE_NEW_VALUE else None b_v_new = tl.zeros([BT, BV], dtype=tl.float32) - p_w = tl.make_block_ptr(w, (T, K), (stride_k, 1), (i_t * BT, 0), (BT, 64), (1, 0)) - b_w = tl.load(p_w, boundary_check=(0, 1)) - b_v_new += tl.dot(b_w, b_h1.to(b_w.dtype)) + p_w = w + o_t[:, None] * stride_k + o_k1[None, :] + b_w = tl.load(p_w, mask=m_t[:, None] & (o_k1[None, :] < K), other=0.0) + b_v_new = tl.dot(b_w, b_h1.to(b_w.dtype), b_v_new) if K > 64: - p_w = tl.make_block_ptr(w, (T, K), (stride_k, 1), (i_t * BT, 64), (BT, 64), (1, 0)) - b_w = tl.load(p_w, boundary_check=(0, 1)) - b_v_new += tl.dot(b_w, b_h2.to(b_w.dtype)) + p_w = w + o_t[:, None] * stride_k + o_k2[None, :] + b_w = tl.load(p_w, mask=m_t[:, None] & (o_k2[None, :] < K), other=0.0) + b_v_new = tl.dot(b_w, b_h2.to(b_w.dtype), b_v_new) if K > 128: - p_w = tl.make_block_ptr(w, (T, K), (stride_k, 1), (i_t * BT, 128), (BT, 64), (1, 0)) - b_w = tl.load(p_w, boundary_check=(0, 1)) - b_v_new += tl.dot(b_w, b_h3.to(b_w.dtype)) + p_w = w + o_t[:, None] * stride_k + o_k3[None, :] + b_w = tl.load(p_w, mask=m_t[:, None] & (o_k3[None, :] < K), other=0.0) + b_v_new = tl.dot(b_w, b_h3.to(b_w.dtype), b_v_new) if K > 192: - p_w = tl.make_block_ptr(w, (T, K), (stride_k, 1), (i_t * BT, 192), (BT, 64), (1, 0)) - b_w = tl.load(p_w, boundary_check=(0, 1)) - b_v_new += tl.dot(b_w, b_h4.to(b_w.dtype)) - b_v_new = -b_v_new + tl.load(p_v, boundary_check=(0, 1)) + p_w = w + o_t[:, None] * stride_k + o_k4[None, :] + b_w = tl.load(p_w, mask=m_t[:, None] & (o_k4[None, :] < K), other=0.0) + b_v_new = tl.dot(b_w, b_h4.to(b_w.dtype), b_v_new) + b_v_new = -b_v_new + tl.load(p_v, mask=m_vv, other=0.0) if SAVE_NEW_VALUE: - p_v_new = tl.make_block_ptr(v_new, (T, V), (stride_v, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) - tl.store(p_v_new, b_v_new.to(p_v_new.dtype.element_ty), boundary_check=(0, 1)) + p_v_new = v_new + o_t[:, None] * stride_v + o_v[None, :] + tl.store(p_v_new, b_v_new.to(p_v_new.dtype.element_ty), mask=m_vv) if USE_G: m_t = (i_t * BT + tl.arange(0, BT)) < T last_idx = min((i_t + 1) * BT, T) - 1 b_g_last = tl.load(g + bos * H + last_idx * H + i_h) - p_g = tl.make_block_ptr(g + bos * H + i_h, (T,), (H,), (i_t * BT,), (BT,), (0,)) - b_g = tl.load(p_g, boundary_check=(0,)) + p_g = g + bos * H + i_h + o_t * H + b_g = tl.load(p_g, mask=m_t, other=0.0) b_v_new = b_v_new * tl.where(m_t, exp2(b_g_last - b_g), 0)[:, None] b_g_last = exp2(b_g_last) b_h1 = b_h1 * b_g_last @@ -165,34 +179,34 @@ def chunk_gated_delta_product_fwd_kernel_h_blockdim64( if K > 192: b_h4 = b_h4 * b_g_last b_v_new = b_v_new.to(k.dtype.element_ty) - p_k = tl.make_block_ptr(k, (K, T), (1, stride_k), (0, i_t * BT), (64, BT), (0, 1)) - b_k = tl.load(p_k, boundary_check=(0, 1)) - b_h1 += tl.dot(b_k, b_v_new) + p_k = k + o_k1[:, None] + o_t[None, :] * stride_k + b_k = tl.load(p_k, mask=(o_k1[:, None] < K) & m_t[None, :], other=0.0) + b_h1 = tl.dot(b_k, b_v_new, b_h1) if K > 64: - p_k = tl.make_block_ptr(k, (K, T), (1, stride_k), (64, i_t * BT), (64, BT), (0, 1)) - b_k = tl.load(p_k, boundary_check=(0, 1)) - b_h2 += tl.dot(b_k, b_v_new) + p_k = k + o_k2[:, None] + o_t[None, :] * stride_k + b_k = tl.load(p_k, mask=(o_k2[:, None] < K) & m_t[None, :], other=0.0) + b_h2 = tl.dot(b_k, b_v_new, b_h2) if K > 128: - p_k = tl.make_block_ptr(k, (K, T), (1, stride_k), (128, i_t * BT), (64, BT), (0, 1)) - b_k = tl.load(p_k, boundary_check=(0, 1)) - b_h3 += tl.dot(b_k, b_v_new) + p_k = k + o_k3[:, None] + o_t[None, :] * stride_k + b_k = tl.load(p_k, mask=(o_k3[:, None] < K) & m_t[None, :], other=0.0) + b_h3 = tl.dot(b_k, b_v_new, b_h3) if K > 192: - p_k = tl.make_block_ptr(k, (K, T), (1, stride_k), (192, i_t * BT), (64, BT), (0, 1)) - b_k = tl.load(p_k, boundary_check=(0, 1)) - b_h4 += tl.dot(b_k, b_v_new) + p_k = k + o_k4[:, None] + o_t[None, :] * stride_k + b_k = tl.load(p_k, mask=(o_k4[:, None] < K) & m_t[None, :], other=0.0) + b_h4 = tl.dot(b_k, b_v_new, b_h4) # epilogue if STORE_FINAL_STATE: - p_ht = tl.make_block_ptr(ht, (K, V), (V, 1), (0, i_v * BV), (64, BV), (1, 0)) - tl.store(p_ht, b_h1.to(p_ht.dtype.element_ty), boundary_check=(0, 1)) + p_ht = ht + o_k1[:, None] * V + o_v[None, :] + tl.store(p_ht, b_h1.to(p_ht.dtype.element_ty), mask=m_h1) if K > 64: - p_ht = tl.make_block_ptr(ht, (K, V), (V, 1), (64, i_v * BV), (64, BV), (1, 0)) - tl.store(p_ht, b_h2.to(p_ht.dtype.element_ty), boundary_check=(0, 1)) + p_ht = ht + o_k2[:, None] * V + o_v[None, :] + tl.store(p_ht, b_h2.to(p_ht.dtype.element_ty), mask=m_h2) if K > 128: - p_ht = tl.make_block_ptr(ht, (K, V), (V, 1), (128, i_v * BV), (64, BV), (1, 0)) - tl.store(p_ht, b_h3.to(p_ht.dtype.element_ty), boundary_check=(0, 1)) + p_ht = ht + o_k3[:, None] * V + o_v[None, :] + tl.store(p_ht, b_h3.to(p_ht.dtype.element_ty), mask=m_h3) if K > 192: - p_ht = tl.make_block_ptr(ht, (K, V), (V, 1), (192, i_v * BV), (64, BV), (1, 0)) - tl.store(p_ht, b_h4.to(p_ht.dtype.element_ty), boundary_check=(0, 1)) + p_ht = ht + o_k4[:, None] * V + o_v[None, :] + tl.store(p_ht, b_h4.to(p_ht.dtype.element_ty), mask=m_h4) @triton.heuristics({ @@ -237,13 +251,13 @@ def chunk_gated_delta_product_bwd_kernel_dhu_blockdim64( USE_FINAL_STATE_GRADIENT: tl.constexpr, IS_VARLEN: tl.constexpr, ): - i_v, i_nh = tl.program_id(0), tl.program_id(1) + i_v, i_nh = unflatten_program_id(tl.cdiv(V, BV)) i_n, i_h = i_nh // H, i_nh % H if IS_VARLEN: - bos, eos = tl.load(cu_seqlens + i_n).to(tl.int32), tl.load(cu_seqlens + i_n + 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 NT = tl.cdiv(T, BT) - boh = tl.load(chunk_offsets + i_n).to(tl.int32) + boh = tl.load(chunk_offsets + i_n).to(tl.int64) else: bos, eos = i_n * T, i_n * T + T NT = tl.cdiv(T, BT) @@ -274,38 +288,53 @@ def chunk_gated_delta_product_bwd_kernel_dhu_blockdim64( if USE_FINAL_STATE_GRADIENT: dht += i_nh * K*V + o_k1 = tl.arange(0, 64) + o_k2 = 64 + tl.arange(0, 64) + o_k3 = 128 + tl.arange(0, 64) + o_k4 = 192 + tl.arange(0, 64) + o_v = i_v * BV + tl.arange(0, BV) + m_v = o_v < V + m_h1 = (o_k1[:, None] < K) & m_v[None, :] + m_h2 = (o_k2[:, None] < K) & m_v[None, :] + m_h3 = (o_k3[:, None] < K) & m_v[None, :] + m_h4 = (o_k4[:, None] < K) & m_v[None, :] + if USE_FINAL_STATE_GRADIENT: - p_dht1 = tl.make_block_ptr(dht, (K, V), (V, 1), (0, i_v * BV), (64, BV), (1, 0)) - b_dh1 += tl.load(p_dht1, boundary_check=(0, 1)) + p_dht1 = dht + o_k1[:, None] * V + o_v[None, :] + b_dh1 += tl.load(p_dht1, mask=m_h1, other=0.0) if K > 64: - p_dht2 = tl.make_block_ptr(dht, (K, V), (V, 1), (64, i_v * BV), (64, BV), (1, 0)) - b_dh2 += tl.load(p_dht2, boundary_check=(0, 1)) + p_dht2 = dht + o_k2[:, None] * V + o_v[None, :] + b_dh2 += tl.load(p_dht2, mask=m_h2, other=0.0) if K > 128: - p_dht3 = tl.make_block_ptr(dht, (K, V), (V, 1), (128, i_v * BV), (64, BV), (1, 0)) - b_dh3 += tl.load(p_dht3, boundary_check=(0, 1)) + p_dht3 = dht + o_k3[:, None] * V + o_v[None, :] + b_dh3 += tl.load(p_dht3, mask=m_h3, other=0.0) if K > 192: - p_dht4 = tl.make_block_ptr(dht, (K, V), (V, 1), (192, i_v * BV), (64, BV), (1, 0)) - b_dh4 += tl.load(p_dht4, boundary_check=(0, 1)) + p_dht4 = dht + o_k4[:, None] * V + o_v[None, :] + b_dh4 += tl.load(p_dht4, mask=m_h4, other=0.0) for i_t in range(NT - 1, -1, -1): - p_dh1 = tl.make_block_ptr(dh + i_t*stride_h, (K, V), (V, 1), (0, i_v * BV), (64, BV), (1, 0)) - tl.store(p_dh1, b_dh1.to(p_dh1.dtype.element_ty), boundary_check=(0, 1)) + i_t_int64 = i_t.to(tl.int64) + o_t = i_t_int64 * BT + tl.arange(0, BT) + m_t = o_t < T + m_vv = m_t[:, None] & m_v[None, :] + p_dh1 = dh + i_t_int64*stride_h + o_k1[:, None] * V + o_v[None, :] + tl.store(p_dh1, b_dh1.to(p_dh1.dtype.element_ty), mask=m_h1) if K > 64: - p_dh2 = tl.make_block_ptr(dh + i_t*stride_h, (K, V), (V, 1), (64, i_v * BV), (64, BV), (1, 0)) - tl.store(p_dh2, b_dh2.to(p_dh2.dtype.element_ty), boundary_check=(0, 1)) + p_dh2 = dh + i_t_int64*stride_h + o_k2[:, None] * V + o_v[None, :] + tl.store(p_dh2, b_dh2.to(p_dh2.dtype.element_ty), mask=m_h2) if K > 128: - p_dh3 = tl.make_block_ptr(dh + i_t*stride_h, (K, V), (V, 1), (128, i_v * BV), (64, BV), (1, 0)) - tl.store(p_dh3, b_dh3.to(p_dh3.dtype.element_ty), boundary_check=(0, 1)) + p_dh3 = dh + i_t_int64*stride_h + o_k3[:, None] * V + o_v[None, :] + tl.store(p_dh3, b_dh3.to(p_dh3.dtype.element_ty), mask=m_h3) if K > 192: - p_dh4 = tl.make_block_ptr(dh + i_t*stride_h, (K, V), (V, 1), (192, i_v * BV), (64, BV), (1, 0)) - tl.store(p_dh4, b_dh4.to(p_dh4.dtype.element_ty), boundary_check=(0, 1)) + p_dh4 = dh + i_t_int64*stride_h + o_k4[:, None] * V + o_v[None, :] + tl.store(p_dh4, b_dh4.to(p_dh4.dtype.element_ty), mask=m_h4) if USE_G: last_idx = min((i_t + 1) * BT, T) - 1 bg_last = tl.load(g + (bos + last_idx) * H + i_h) bg_last_exp = exp2(bg_last) - p_g = tl.make_block_ptr(g + bos * H + i_h, (T,), (H,), (i_t * BT,), (BT,), (0,)) - b_g = tl.load(p_g, boundary_check=(0,)) + p_g = g + bos * H + i_h + o_t * H + b_g = tl.load(p_g, mask=m_t, other=0.0) b_g_exp = exp2(b_g) else: bg_last = None @@ -313,74 +342,74 @@ def chunk_gated_delta_product_bwd_kernel_dhu_blockdim64( b_g = None b_g_exp = None - p_dv = tl.make_block_ptr(dv, (T, V), (stride_v, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) - p_wo = tl.make_block_ptr(do, (T, V), (stride_v, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) - p_dv2 = tl.make_block_ptr(dv2, (T, V), (stride_v, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) + p_dv = dv + o_t[:, None] * stride_v + o_v[None, :] + p_wo = do + o_t[:, None] * stride_v + o_v[None, :] + p_dv2 = dv2 + o_t[:, None] * stride_v + o_v[None, :] - b_wo = tl.load(p_wo, boundary_check=(0, 1)) + b_wo = tl.load(p_wo, mask=m_vv, other=0.0) b_dv = tl.zeros([BT, BV], dtype=tl.float32) # Update dv - p_k = tl.make_block_ptr(k, (T, K), (stride_k, 1), (i_t * BT, 0), (BT, 64), (1, 0)) - b_k = tl.load(p_k, boundary_check=(0, 1)) - b_dv += tl.dot(b_k, b_dh1.to(b_k.dtype)) + p_k = k + o_t[:, None] * stride_k + o_k1[None, :] + b_k = tl.load(p_k, mask=m_t[:, None] & (o_k1[None, :] < K), other=0.0) + b_dv = tl.dot(b_k, b_dh1.to(b_k.dtype), b_dv) if K > 64: - p_k = tl.make_block_ptr(k, (T, K), (stride_k, 1), (i_t * BT, 64), (BT, 64), (1, 0)) - b_k = tl.load(p_k, boundary_check=(0, 1)) - b_dv += tl.dot(b_k, b_dh2.to(b_k.dtype)) + p_k = k + o_t[:, None] * stride_k + o_k2[None, :] + b_k = tl.load(p_k, mask=m_t[:, None] & (o_k2[None, :] < K), other=0.0) + b_dv = tl.dot(b_k, b_dh2.to(b_k.dtype), b_dv) if K > 128: - p_k = tl.make_block_ptr(k, (T, K), (stride_k, 1), (i_t * BT, 128), (BT, 64), (1, 0)) - b_k = tl.load(p_k, boundary_check=(0, 1)) - b_dv += tl.dot(b_k, b_dh3.to(b_k.dtype)) + p_k = k + o_t[:, None] * stride_k + o_k3[None, :] + b_k = tl.load(p_k, mask=m_t[:, None] & (o_k3[None, :] < K), other=0.0) + b_dv = tl.dot(b_k, b_dh3.to(b_k.dtype), b_dv) if K > 192: - p_k = tl.make_block_ptr(k, (T, K), (stride_k, 1), (i_t * BT, 192), (BT, 64), (1, 0)) - b_k = tl.load(p_k, boundary_check=(0, 1)) - b_dv += tl.dot(b_k, b_dh4.to(b_k.dtype)) + p_k = k + o_t[:, None] * stride_k + o_k4[None, :] + b_k = tl.load(p_k, mask=m_t[:, None] & (o_k4[None, :] < K), other=0.0) + b_dv = tl.dot(b_k, b_dh4.to(b_k.dtype), b_dv) if USE_G: m_t = (i_t * BT + tl.arange(0, BT)) < T b_dv *= tl.where(m_t, exp2(bg_last - b_g), 0)[:, None] - b_dv += tl.load(p_dv, boundary_check=(0, 1)) + b_dv += tl.load(p_dv, mask=m_vv, other=0.0) - tl.store(p_dv2, b_dv.to(p_dv.dtype.element_ty), boundary_check=(0, 1)) + tl.store(p_dv2, b_dv.to(p_dv.dtype.element_ty), mask=m_vv) # Update dh - p_w = tl.make_block_ptr(w, (K, T), (1, stride_k), (0, i_t * BT), (64, BT), (0, 1)) - p_q = tl.make_block_ptr(q, (K, T), (1, stride_k), (0, i_t * BT), (64, BT), (0, 1)) - b_w = tl.load(p_w, boundary_check=(0, 1)) - b_q = tl.load(p_q, boundary_check=(0, 1)) + p_w = w + o_k1[:, None] + o_t[None, :] * stride_k + p_q = q + o_k1[:, None] + o_t[None, :] * stride_k + b_w = tl.load(p_w, mask=(o_k1[:, None] < K) & m_t[None, :], other=0.0) + b_q = tl.load(p_q, mask=(o_k1[:, None] < K) & m_t[None, :], other=0.0) if USE_G: b_dh1 *= bg_last_exp b_q = b_q * b_g_exp[None, :] b_q = (b_q * scale).to(b_q.dtype) b_dh1 += tl.dot(b_q, b_wo.to(b_q.dtype))-tl.dot(b_w, b_dv.to(b_w.dtype)) if K > 64: - p_q = tl.make_block_ptr(q, (K, T), (1, stride_k), (64, i_t * BT), (64, BT), (0, 1)) - p_w = tl.make_block_ptr(w, (K, T), (1, stride_k), (64, i_t * BT), (64, BT), (0, 1)) - b_q = tl.load(p_q, boundary_check=(0, 1)) - b_w = tl.load(p_w, boundary_check=(0, 1)) + p_q = q + o_k2[:, None] + o_t[None, :] * stride_k + p_w = w + o_k2[:, None] + o_t[None, :] * stride_k + b_q = tl.load(p_q, mask=(o_k2[:, None] < K) & m_t[None, :], other=0.0) + b_w = tl.load(p_w, mask=(o_k2[:, None] < K) & m_t[None, :], other=0.0) if USE_G: b_dh2 *= bg_last_exp b_q = b_q * b_g_exp[None, :] b_q = (b_q * scale).to(b_q.dtype) b_dh2 += tl.dot(b_q, b_wo.to(b_q.dtype))-tl.dot(b_w, b_dv.to(b_w.dtype)) if K > 128: - p_q = tl.make_block_ptr(q, (K, T), (1, stride_k), (128, i_t * BT), (64, BT), (0, 1)) - p_w = tl.make_block_ptr(w, (K, T), (1, stride_k), (128, i_t * BT), (64, BT), (0, 1)) - b_q = tl.load(p_q, boundary_check=(0, 1)) - b_w = tl.load(p_w, boundary_check=(0, 1)) + p_q = q + o_k3[:, None] + o_t[None, :] * stride_k + p_w = w + o_k3[:, None] + o_t[None, :] * stride_k + b_q = tl.load(p_q, mask=(o_k3[:, None] < K) & m_t[None, :], other=0.0) + b_w = tl.load(p_w, mask=(o_k3[:, None] < K) & m_t[None, :], other=0.0) if USE_G: b_dh3 *= bg_last_exp b_q = b_q * b_g_exp[None, :] b_q = (b_q * scale).to(b_q.dtype) b_dh3 += tl.dot(b_q, b_wo.to(b_q.dtype))-tl.dot(b_w, b_dv.to(b_w.dtype)) if K > 192: - p_q = tl.make_block_ptr(q, (K, T), (1, stride_k), (192, i_t * BT), (64, BT), (0, 1)) - p_w = tl.make_block_ptr(w, (K, T), (1, stride_k), (192, i_t * BT), (64, BT), (0, 1)) - b_q = tl.load(p_q, boundary_check=(0, 1)) - b_w = tl.load(p_w, boundary_check=(0, 1)) + p_q = q + o_k4[:, None] + o_t[None, :] * stride_k + p_w = w + o_k4[:, None] + o_t[None, :] * stride_k + b_q = tl.load(p_q, mask=(o_k4[:, None] < K) & m_t[None, :], other=0.0) + b_w = tl.load(p_w, mask=(o_k4[:, None] < K) & m_t[None, :], other=0.0) if USE_G: b_dh4 *= bg_last_exp b_q = b_q * b_g_exp[None, :] @@ -388,17 +417,17 @@ def chunk_gated_delta_product_bwd_kernel_dhu_blockdim64( b_dh4 += tl.dot(b_q, b_wo.to(b_q.dtype))-tl.dot(b_w, b_dv.to(b_w.dtype)) if USE_INITIAL_STATE: - p_dh0 = tl.make_block_ptr(dh0, (K, V), (V, 1), (0, i_v * BV), (64, BV), (1, 0)) - tl.store(p_dh0, b_dh1.to(p_dh0.dtype.element_ty), boundary_check=(0, 1)) + p_dh0 = dh0 + o_k1[:, None] * V + o_v[None, :] + tl.store(p_dh0, b_dh1.to(p_dh0.dtype.element_ty), mask=m_h1) if K > 64: - p_dh1 = tl.make_block_ptr(dh0, (K, V), (V, 1), (64, i_v * BV), (64, BV), (1, 0)) - tl.store(p_dh1, b_dh2.to(p_dh1.dtype.element_ty), boundary_check=(0, 1)) + p_dh1 = dh0 + o_k2[:, None] * V + o_v[None, :] + tl.store(p_dh1, b_dh2.to(p_dh1.dtype.element_ty), mask=m_h2) if K > 128: - p_dh2 = tl.make_block_ptr(dh0, (K, V), (V, 1), (128, i_v * BV), (64, BV), (1, 0)) - tl.store(p_dh2, b_dh3.to(p_dh2.dtype.element_ty), boundary_check=(0, 1)) + p_dh2 = dh0 + o_k3[:, None] * V + o_v[None, :] + tl.store(p_dh2, b_dh3.to(p_dh2.dtype.element_ty), mask=m_h3) if K > 192: - p_dh3 = tl.make_block_ptr(dh0, (K, V), (V, 1), (192, i_v * BV), (64, BV), (1, 0)) - tl.store(p_dh3, b_dh4.to(p_dh3.dtype.element_ty), boundary_check=(0, 1)) + p_dh3 = dh0 + o_k4[:, None] * V + o_v[None, :] + tl.store(p_dh3, b_dh4.to(p_dh3.dtype.element_ty), mask=m_h4) def chunk_gated_delta_product_fwd_h( @@ -431,7 +460,7 @@ def chunk_gated_delta_product_fwd_h( final_state = k.new_empty(N, H, K, V, dtype=torch.float32) if output_final_state else None v_new = torch.empty_like(u) if save_new_value else None - def grid(meta): return (triton.cdiv(V, meta['BV']), N*H) + def grid(meta): return (triton.cdiv(V, meta['BV']) * N * H,) chunk_gated_delta_product_fwd_kernel_h_blockdim64[grid]( k=k, v=u, @@ -484,7 +513,7 @@ def chunk_gated_delta_product_bwd_dhu( dh0 = torch.empty_like(h0, dtype=torch.float32) if h0 is not None else None dv2 = torch.empty_like(dv) - def grid(meta): return (triton.cdiv(V, meta['BV']), N*H) + def grid(meta): return (triton.cdiv(V, meta['BV']) * N * H,) chunk_gated_delta_product_bwd_kernel_dhu_blockdim64[grid]( q=q, k=k, diff --git a/build/torch-cuda/ops/gated_delta_product/chunk_deltaproduct_o.py b/build/torch-cuda/ops/gated_delta_product/chunk_deltaproduct_o.py index 975530c4c61d64006c1ee17886aa11d598823a98..52447327ca559153d21a7661f0bec1785cf36890 100644 --- a/build/torch-cuda/ops/gated_delta_product/chunk_deltaproduct_o.py +++ b/build/torch-cuda/ops/gated_delta_product/chunk_deltaproduct_o.py @@ -10,7 +10,7 @@ import triton import triton.language as tl from ...ops.utils import prepare_chunk_indices -from ...ops.utils.op import exp2 +from ...ops.utils.op import exp2, unflatten_program_id from ...utils import IS_NVIDIA_HOPPER, autotune_cache_kwargs, check_shared_mem BKV_LIST = [64, 128] if check_shared_mem() else [32, 64] @@ -54,13 +54,14 @@ def chunk_fwd_kernel_o( USE_G: tl.constexpr, IS_VARLEN: tl.constexpr, ): - i_v, i_t, i_bh = tl.program_id(0), tl.program_id(1), tl.program_id(2) + i_v, i_t = unflatten_program_id(tl.cdiv(V, BV)) + i_bh = tl.program_id(1).to(tl.int64) i_b, i_h = i_bh // H, i_bh % H 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.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) + i_n, i_t = tl.load(chunk_indices + i_t * 2).to(tl.int32), tl.load(chunk_indices + i_t * 2 + 1).to(tl.int64) + bos, eos = tl.load(cu_seqlens + i_n).to(tl.int64), tl.load(cu_seqlens + i_n + 1).to(tl.int64) T = eos - bos NT = tl.cdiv(T, BT) else: @@ -75,24 +76,29 @@ def chunk_fwd_kernel_o( o += (bos * H + i_h) * V h += (i_tg * H + i_h).to(tl.int64) * K*V + o_t = i_t * BT + tl.arange(0, BT) + o_v = i_v * BV + tl.arange(0, BV) + m_t = o_t < T + m_v = m_t[:, None] & (o_v[None, :] < V) b_o = tl.zeros([BT, BV], dtype=tl.float32) for i_k in range(tl.cdiv(K, BK)): - p_q = tl.make_block_ptr(q, (T, K), (H*K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) - p_h = tl.make_block_ptr(h, (K, V), (V, 1), (i_k * BK, i_v * BV), (BK, BV), (1, 0)) + o_k = i_k * BK + tl.arange(0, BK) + m_q = m_t[:, None] & (o_k[None, :] < K) + m_h = (o_k[:, None] < K) & (o_v[None, :] < V) + p_q = q + o_t[:, None] * (H*K) + o_k[None, :] + p_h = h + o_k[:, None] * V + o_v[None, :] # [BT, BK] - b_q = tl.load(p_q, boundary_check=(0, 1)) + b_q = tl.load(p_q, mask=m_q, other=0.0) # [BK, BV] - b_h = tl.load(p_h, boundary_check=(0, 1)) + b_h = tl.load(p_h, mask=m_h, other=0.0) # [BT, BK] @ [BK, BV] -> [BT, BV] - b_o += tl.dot(b_q, b_h) + b_o = tl.dot(b_q, b_h, b_o) - o_t = i_t * BT + tl.arange(0, BT) - m_t = o_t < T if USE_G: g += bos * H + i_h - p_g = tl.make_block_ptr(g, (T,), (H,), (i_t * BT,), (BT,), (0,)) - b_g = tl.load(p_g, boundary_check=(0,)) + p_g = g + o_t * H + b_g = tl.load(p_g, mask=m_t, other=0.0) m_A = (o_t[:, None] >= o_t[None, :]) & (m_t[:, None] & m_t) b_m = tl.where(m_A, exp2(b_g[:, None] - b_g[None, :]), 0) b_o = b_o * exp2(b_g)[:, None] @@ -102,21 +108,24 @@ def chunk_fwd_kernel_o( for i_dp in range(num_householder): b_A = tl.zeros([BT, BT], dtype=tl.float32) for i_k in range(tl.cdiv(K, BK)): - p_q = tl.make_block_ptr(q, (T, K), (H*K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) - p_k = tl.make_block_ptr(k+i_dp*H*K, (K, T), (1, num_householder*H*K), (i_k * BK, i_t * BT), (BK, BT), (0, 1)) + o_k = i_k * BK + tl.arange(0, BK) + m_q = m_t[:, None] & (o_k[None, :] < K) + m_k = (o_k[:, None] < K) & m_t[None, :] + p_q = q + o_t[:, None] * (H*K) + o_k[None, :] + p_k = k+i_dp*H*K + o_k[:, None] + o_t[None, :] * (num_householder*H*K) # [BT, BK] - b_q = tl.load(p_q, boundary_check=(0, 1)) + b_q = tl.load(p_q, mask=m_q, other=0.0) # [BK, BT] - b_k = tl.load(p_k, boundary_check=(0, 1)) + b_k = tl.load(p_k, mask=m_k, other=0.0) # [BT, BK] @ [BK, BT] -> [BT, BT] - b_A += tl.dot(b_q, b_k) + b_A = tl.dot(b_q, b_k, b_A) b_A = b_A * b_m - p_v = tl.make_block_ptr(v+i_dp*H*V, (T, V), (H*V*num_householder, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) - b_v = tl.load(p_v, boundary_check=(0, 1)) - b_o += tl.dot(b_A.to(b_v.dtype), b_v) + p_v = v+i_dp*H*V + o_t[:, None] * (H*V*num_householder) + o_v[None, :] + b_v = tl.load(p_v, mask=m_v, other=0.0) + b_o = tl.dot(b_A.to(b_v.dtype), b_v, b_o) b_o = b_o * scale - p_o = tl.make_block_ptr(o, (T, V), (H*V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) - tl.store(p_o, b_o.to(p_o.dtype.element_ty), boundary_check=(0, 1)) + p_o = o + o_t[:, None] * (H*V) + o_v[None, :] + tl.store(p_o, b_o.to(p_o.dtype.element_ty), mask=m_v) def chunk_gated_delta_product_fwd_o( @@ -138,7 +147,7 @@ def chunk_gated_delta_product_fwd_o( chunk_indices = prepare_chunk_indices(cu_seqlens, BT) NT = triton.cdiv(T, BT) if cu_seqlens is None else len(chunk_indices) o = v.new_empty(B, T, H, V).fill_(-float('inf')) - def grid(meta): return (triton.cdiv(V, meta['BV']), NT, B * H) + def grid(meta): return (triton.cdiv(V, meta['BV']) * NT, B * H) chunk_fwd_kernel_o[grid]( q, k, diff --git a/build/torch-cuda/ops/gated_delta_product/naive.py b/build/torch-cuda/ops/gated_delta_product/naive.py index 106f13d2954e5fad48f46be4526b92b876a7d8a0..591c3d3147e6b4837b60b9928fef1720a6ea3615 100644 --- a/build/torch-cuda/ops/gated_delta_product/naive.py +++ b/build/torch-cuda/ops/gated_delta_product/naive.py @@ -8,9 +8,29 @@ import torch -def naive_recurrent_gated_delta_product(q, k, v, g, beta, scale, cu_seqlens, +def naive_recurrent_gated_delta_product(q, k, v, g, beta, scale, cu_seqlens=None, initial_state=None, output_final_state=False, num_householder=1): + if cu_seqlens is not None: + outputs, final_states = [], [] + offsets = cu_seqlens.tolist() + for i, (bos, eos) in enumerate(zip(offsets[:-1], offsets[1:], strict=False)): + o_i, h_i = naive_recurrent_gated_delta_product( + q=q[:, bos:eos], + k=k[:, bos*num_householder:eos*num_householder], + v=v[:, bos*num_householder:eos*num_householder], + g=None if g is None else g[:, bos:eos], + beta=beta[:, bos*num_householder:eos*num_householder], + scale=scale, + initial_state=None if initial_state is None else initial_state[i:i + 1], + output_final_state=output_final_state, + num_householder=num_householder, + ) + outputs.append(o_i) + final_states.append(h_i) + + return torch.cat(outputs, dim=1), torch.cat(final_states) + q_original_dtype = q.dtype B, T, H, K = q.shape V = v.shape[-1] @@ -20,6 +40,7 @@ def naive_recurrent_gated_delta_product(q, k, v, g, beta, scale, cu_seqlens, if g is not None: assert g.shape == (B, T, H) q, k, v, beta = map(lambda x: x.float(), (q, k, v, beta)) + q = q * scale h = torch.zeros(B, H, K, V, dtype=torch.float32, device=q.device) if initial_state is not None: diff --git a/build/torch-xpu/ops/log_linear_attn/__init__.py b/build/torch-cuda/ops/gated_delta_rule/backends/__init__.py similarity index 71% rename from build/torch-xpu/ops/log_linear_attn/__init__.py rename to build/torch-cuda/ops/gated_delta_rule/backends/__init__.py index 56d33ebce43f4539caa76eb3d9de450bc9b93ac3..2bfb3cc10eeea2116f2699b4707f6ffba41a7c61 100644 --- a/build/torch-xpu/ops/log_linear_attn/__init__.py +++ b/build/torch-cuda/ops/gated_delta_rule/backends/__init__.py @@ -5,8 +5,8 @@ # For a list of all contributors, visit: # https://github.com/fla-org/flash-linear-attention/graphs/contributors -from .chunk import chunk_log_linear_attn +"""GDR backends.""" -__all__ = [ - 'chunk_log_linear_attn', -] +from ....ops.gated_delta_rule.backends.flash_qla import FlashQLABackend + +__all__ = ['FlashQLABackend'] diff --git a/build/torch-cuda/ops/gated_delta_rule/backends/flash_qla.py b/build/torch-cuda/ops/gated_delta_rule/backends/flash_qla.py new file mode 100644 index 0000000000000000000000000000000000000000..f8744fd963b64b677db01c0d68cce1ca67a1b69b --- /dev/null +++ b/build/torch-cuda/ops/gated_delta_rule/backends/flash_qla.py @@ -0,0 +1,134 @@ +# 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 +# +# Copyright (c) 2026 Qwen Team, Alibaba Cloud + +"""FlashQLA backend for chunk_gated_delta_rule.""" + +from __future__ import annotations + +from typing import TYPE_CHECKING + +import torch + +from ....backends import BaseBackend, register +from ....utils import IS_NVIDIA_HOPPER, IS_NVIDIA_SM100, IS_NVIDIA_SM120 + +if TYPE_CHECKING: + from ....ops.cp import FLACPContext + + +def _needs_backward(*tensors: torch.Tensor | None) -> bool: + """Whether autograd would later require a backward pass over these inputs.""" + if not torch.is_grad_enabled(): + return False + return any(isinstance(t, torch.Tensor) and t.requires_grad for t in tensors) + + +# We do not register as we cannot rely on the backend to work for compile +# @register('gated_delta_rule') +class FlashQLABackend(BaseBackend): + """Copyright (c) 2026 Qwen Team, Alibaba Cloud + + Fused TileLang forward and backward with intra-card CP (replaces the multi-kernel Triton path). + https://github.com/QwenLM/FlashQLA + + SM90/SM100/SM103 run both directions. SM120 (consumer/workstation Blackwell) ships a + bfloat16 forward kernel only, so it is dispatched exclusively for grad-free bf16 calls + (inference, frozen weights) and falls back to Triton otherwise. + + Disable with ``FLA_FLASH_QLA=0``. + """ + + backend_type = "flash_qla" + package_name = "flash_qla" + env_var = "FLA_FLASH_QLA" + default_enable = True + priority = 3 + + def chunk_gated_delta_rule_verifier( + self, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + g: torch.Tensor, + beta: torch.Tensor, + scale: float | None = None, + initial_state: torch.Tensor | None = None, + output_final_state: bool = False, + use_qk_l2norm_in_kernel: bool = False, + use_beta_sigmoid_in_kernel: bool = False, + allow_neg_eigval: bool = False, + state_v_first: bool = False, + cu_seqlens: torch.LongTensor | None = None, + cu_seqlens_cpu: torch.LongTensor | None = None, + cp_context: FLACPContext | None = None, + **kwargs, + ) -> tuple[bool, str | None]: + if not (IS_NVIDIA_HOPPER or IS_NVIDIA_SM100 or IS_NVIDIA_SM120): + return False, "FlashQLA requires NVIDIA SM90, SM100/SM103 or SM120" + if IS_NVIDIA_SM120 and _needs_backward(q, k, v, g, beta, initial_state): + return False, "FlashQLA on SM120 implements the forward pass only, but an input requires grad" + if q.dtype != torch.float16 and q.dtype != torch.bfloat16: + return False, f"FlashQLA requires dtype float16 or bfloat16, got {q.dtype}" + if not (q.dtype == k.dtype == v.dtype): + return False, f"FlashQLA requires q, k, v to have the same dtype, got {q.dtype}, {k.dtype}, {v.dtype}" + # NOTE: the masked tail-store in FlashQLA's blackwell_sm120 forward kernel emits + # tl::pack_float16x4 on cutlass::half_t, which fails to compile under nvcc. + if IS_NVIDIA_SM120 and q.dtype == torch.float16: + return False, "FlashQLA's SM120 forward kernel does not compile for float16" + if q.shape[-1] != 128: + return False, f"FlashQLA requires K=128, got {q.shape[-1]}" + if v.shape[-1] != 128: + return False, f"FlashQLA requires V=128, got {v.shape[-1]}" + if kwargs.get('use_gate_in_kernel'): + return False, "FlashQLA does not support use_gate_in_kernel" + if use_beta_sigmoid_in_kernel: + return False, "FlashQLA does not support use_beta_sigmoid_in_kernel" + if allow_neg_eigval: + return False, "FlashQLA does not support allow_neg_eigval" + if 'transpose_state_layout' in kwargs: + return False, "FlashQLA does not support the deprecated transpose_state_layout" + if cp_context is not None: + return False, "FlashQLA does not support inter-card context parallel" + return True, None + + def chunk_gated_delta_rule( + self, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + g: torch.Tensor, + beta: torch.Tensor, + scale: float | None = None, + initial_state: torch.Tensor | None = None, + output_final_state: bool = False, + use_qk_l2norm_in_kernel: bool = False, + use_beta_sigmoid_in_kernel: bool = False, + allow_neg_eigval: bool = False, + state_v_first: bool = False, + cu_seqlens: torch.LongTensor | None = None, + cu_seqlens_cpu: torch.LongTensor | None = None, + cp_context: FLACPContext | None = None, + **kwargs, + ): + import flash_qla + + return flash_qla.chunk_gated_delta_rule( + q=q, + k=k, + v=v, + g=g, + beta=beta, + scale=scale, + initial_state=initial_state, + output_final_state=output_final_state, + use_qk_l2norm_in_kernel=use_qk_l2norm_in_kernel, + state_v_first=state_v_first, + cu_seqlens=cu_seqlens, + auto_cp=True, + ) diff --git a/build/torch-cuda/ops/gated_delta_rule/chunk.py b/build/torch-cuda/ops/gated_delta_rule/chunk.py index 63fd1b099859f34dae89aad281cf435bab25c761..662692bc235e7d3888b7eb0f14d4516e5de70abc 100644 --- a/build/torch-cuda/ops/gated_delta_rule/chunk.py +++ b/build/torch-cuda/ops/gated_delta_rule/chunk.py @@ -9,7 +9,8 @@ import warnings import torch -from ...modules.l2norm import l2norm_bwd, l2norm_fwd +from ...backends import dispatch +from ...modules.norm.l2norm import l2norm_bwd, l2norm_fwd from ...ops.common.chunk_delta_h import chunk_gated_delta_rule_bwd_dhu, chunk_gated_delta_rule_fwd_h from ...ops.common.chunk_o import chunk_bwd_dqkwg, chunk_bwd_dv_local, chunk_fwd_o from ...ops.common.gate import fused_beta_sigmoid, fused_beta_sigmoid_bwd @@ -267,6 +268,7 @@ class ChunkGatedDeltaRuleFunction(torch.autograd.Function): state_v_first: bool = False, cu_seqlens: torch.LongTensor | None = None, cu_seqlens_cpu: torch.LongTensor | None = None, + chunk_indices: torch.LongTensor | None = None, use_qk_l2norm_in_kernel: bool = False, use_gate_in_kernel: bool = False, A_log: torch.Tensor | None = None, @@ -285,8 +287,7 @@ class ChunkGatedDeltaRuleFunction(torch.autograd.Function): if use_beta_sigmoid_in_kernel: beta = fused_beta_sigmoid(beta_raw, scale=2.0 if allow_neg_eigval else 1.0) - chunk_indices = None - if cu_seqlens is not None: + if chunk_indices is None and cu_seqlens is not None: chunk_indices = prepare_chunk_indices(cu_seqlens, chunk_size, cu_seqlens_cpu=cu_seqlens_cpu) g, o, A, final_state, initial_state, g_input = chunk_gated_delta_rule_fwd( q=q, @@ -386,12 +387,13 @@ class ChunkGatedDeltaRuleFunction(torch.autograd.Function): db = fused_beta_sigmoid_bwd(beta_raw, db, scale=2.0 if ctx.allow_neg_eigval else 1.0) return ( dq.to(q), dk.to(k), dv.to(v), dg.to(g), db.to(beta_raw), - None, dh0, None, None, None, None, None, None, dA_log, ddt_bias, + None, dh0, None, None, None, None, None, None, None, dA_log, ddt_bias, None, None, None, None, ) -@torch.compiler.disable +# Use triton only, other backends (cutlass, cute, etc.) need additional torch.compile support +# @dispatch('gated_delta_rule') def chunk_gated_delta_rule( q: torch.Tensor, k: torch.Tensor, @@ -407,6 +409,7 @@ def chunk_gated_delta_rule( state_v_first: bool = False, cu_seqlens: torch.LongTensor | None = None, cu_seqlens_cpu: torch.LongTensor | None = None, + chunk_indices: torch.LongTensor | None = None, cp_context: FLACPContext | None = None, **kwargs, ): @@ -461,6 +464,9 @@ def chunk_gated_delta_rule( cu_seqlens (torch.LongTensor): Cumulative sequence lengths of shape `[N+1]` used for variable-length training, consistent with the FlashAttention API. + chunk_indices (Optional[torch.LongTensor]): + Pre-computed chunk indices for variable-length inputs. + If provided, they are used directly instead of being computed from `cu_seqlens`. Default: `None`. cp_context (Optional[FLACPContext]): Context parallel context for distributed training across multiple devices. When provided, `initial_state` and `output_final_state` are not supported, @@ -476,7 +482,7 @@ def chunk_gated_delta_rule( >>> import torch >>> import torch.nn.functional as F >>> from einops import rearrange - >>> from ...ops.gated_delta_rule import chunk_gated_delta_rule + >>> from fla.ops.gated_delta_rule import chunk_gated_delta_rule # inputs with equal lengths >>> B, T, H, HV, K, V = 4, 2048, 4, 8, 512, 512 >>> q = torch.randn(B, T, H, K, dtype=torch.bfloat16, device='cuda') @@ -574,6 +580,7 @@ def chunk_gated_delta_rule( state_v_first, cu_seqlens, cu_seqlens_cpu, + chunk_indices, use_qk_l2norm_in_kernel, use_gate_in_kernel, A_log, diff --git a/build/torch-cuda/ops/gated_delta_rule/chunk_fwd.py b/build/torch-cuda/ops/gated_delta_rule/chunk_fwd.py index 55febabc537687f7c71a5b97ff0eb76bdf90c067..58ee16a3b61f771c1d4fc0ceb0bcad95db0c9dfc 100644 --- a/build/torch-cuda/ops/gated_delta_rule/chunk_fwd.py +++ b/build/torch-cuda/ops/gated_delta_rule/chunk_fwd.py @@ -9,6 +9,7 @@ import torch import triton import triton.language as tl +from ...backends import dispatch from ...ops.common.chunk_scaled_dot_kkt import chunk_scaled_dot_kkt_fwd from ...ops.gated_delta_rule.wy_fast import recompute_w_u_fwd from ...ops.utils import prepare_chunk_indices, solve_tril @@ -66,12 +67,12 @@ def chunk_gated_delta_rule_fwd_kkt_solve_kernel( 4. Block merge to get full (I+A)^{-1} 5. Write result to A (output) """ - i_t, i_bh = tl.program_id(0), tl.program_id(1) + i_t, i_bh = tl.program_id(0).to(tl.int64), tl.program_id(1).to(tl.int64) i_b, i_h = i_bh // HV, i_bh % HV 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.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) + i_n, i_t = tl.load(chunk_indices + i_t * 2).to(tl.int32), tl.load(chunk_indices + i_t * 2 + 1).to(tl.int64) + bos, eos = tl.load(cu_seqlens + i_n).to(tl.int64), tl.load(cu_seqlens + i_n + 1).to(tl.int64) T = eos - bos else: bos, eos = i_b * T, i_b * T + T @@ -94,26 +95,26 @@ def chunk_gated_delta_rule_fwd_kkt_solve_kernel( m_tc3 = (i_tc3 + o_i) < T # load beta for each sub-chunk - p_b0 = tl.make_block_ptr(beta + bos * HV + i_h, (T,), (HV,), (i_tc0,), (BC,), (0,)) - p_b1 = tl.make_block_ptr(beta + bos * HV + i_h, (T,), (HV,), (i_tc1,), (BC,), (0,)) - p_b2 = tl.make_block_ptr(beta + bos * HV + i_h, (T,), (HV,), (i_tc2,), (BC,), (0,)) - p_b3 = tl.make_block_ptr(beta + bos * HV + i_h, (T,), (HV,), (i_tc3,), (BC,), (0,)) - b_b0 = tl.load(p_b0, boundary_check=(0,)).to(tl.float32) - b_b1 = tl.load(p_b1, boundary_check=(0,)).to(tl.float32) - b_b2 = tl.load(p_b2, boundary_check=(0,)).to(tl.float32) - b_b3 = tl.load(p_b3, boundary_check=(0,)).to(tl.float32) + p_b0 = beta + bos * HV + i_h + (i_tc0 + o_i) * HV + p_b1 = beta + bos * HV + i_h + (i_tc1 + o_i) * HV + p_b2 = beta + bos * HV + i_h + (i_tc2 + o_i) * HV + p_b3 = beta + bos * HV + i_h + (i_tc3 + o_i) * HV + b_b0 = tl.load(p_b0, mask=m_tc0, other=0.0).to(tl.float32) + b_b1 = tl.load(p_b1, mask=m_tc1, other=0.0).to(tl.float32) + b_b2 = tl.load(p_b2, mask=m_tc2, other=0.0).to(tl.float32) + b_b3 = tl.load(p_b3, mask=m_tc3, other=0.0).to(tl.float32) # load gate if used if USE_G: - p_g0 = tl.make_block_ptr(g + bos * HV + i_h, (T,), (HV,), (i_tc0,), (BC,), (0,)) - p_g1 = tl.make_block_ptr(g + bos * HV + i_h, (T,), (HV,), (i_tc1,), (BC,), (0,)) - p_g2 = tl.make_block_ptr(g + bos * HV + i_h, (T,), (HV,), (i_tc2,), (BC,), (0,)) - p_g3 = tl.make_block_ptr(g + bos * HV + i_h, (T,), (HV,), (i_tc3,), (BC,), (0,)) + p_g0 = g + bos * HV + i_h + (i_tc0 + o_i) * HV + p_g1 = g + bos * HV + i_h + (i_tc1 + o_i) * HV + p_g2 = g + bos * HV + i_h + (i_tc2 + o_i) * HV + p_g3 = g + bos * HV + i_h + (i_tc3 + o_i) * HV - b_g0 = tl.load(p_g0, boundary_check=(0,)).to(tl.float32) - b_g1 = tl.load(p_g1, boundary_check=(0,)).to(tl.float32) - b_g2 = tl.load(p_g2, boundary_check=(0,)).to(tl.float32) - b_g3 = tl.load(p_g3, boundary_check=(0,)).to(tl.float32) + b_g0 = tl.load(p_g0, mask=m_tc0, other=0.0).to(tl.float32) + b_g1 = tl.load(p_g1, mask=m_tc1, other=0.0).to(tl.float32) + b_g2 = tl.load(p_g2, mask=m_tc2, other=0.0).to(tl.float32) + b_g3 = tl.load(p_g3, mask=m_tc3, other=0.0).to(tl.float32) ############################################################################ # Step 1: compute all 10 lower-triangular [BC, BC] blocks of K @ K^T @@ -134,37 +135,38 @@ def chunk_gated_delta_rule_fwd_kkt_solve_kernel( b_A32 = tl.zeros([BC, BC], dtype=tl.float32) for i_k in range(tl.cdiv(K, BK)): - p_k0 = tl.make_block_ptr(k, (T, K), (H*K, 1), (i_tc0, i_k * BK), (BC, BK), (1, 0)) - b_k0 = tl.load(p_k0, boundary_check=(0, 1)) + o_k = i_k * BK + tl.arange(0, BK) + p_k0 = k + (i_tc0 + o_i)[:, None] * (H*K) + o_k[None, :] + b_k0 = tl.load(p_k0, mask=m_tc0[:, None] & (o_k[None, :] < K), other=0.0) # diagonal block 0 - b_A00 += tl.dot(b_k0, tl.trans(b_k0)) + b_A00 = tl.dot(b_k0, tl.trans(b_k0), b_A00) if i_tc1 < T: - p_k1 = tl.make_block_ptr(k, (T, K), (H*K, 1), (i_tc1, i_k * BK), (BC, BK), (1, 0)) - b_k1 = tl.load(p_k1, boundary_check=(0, 1)) + p_k1 = k + (i_tc1 + o_i)[:, None] * (H*K) + o_k[None, :] + b_k1 = tl.load(p_k1, mask=m_tc1[:, None] & (o_k[None, :] < K), other=0.0) # diagonal block 1 - b_A11 += tl.dot(b_k1, tl.trans(b_k1)) + b_A11 = tl.dot(b_k1, tl.trans(b_k1), b_A11) # off-diagonal (1,0) - b_A10 += tl.dot(b_k1, tl.trans(b_k0)) + b_A10 = tl.dot(b_k1, tl.trans(b_k0), b_A10) if i_tc2 < T: - p_k2 = tl.make_block_ptr(k, (T, K), (H*K, 1), (i_tc2, i_k * BK), (BC, BK), (1, 0)) - b_k2 = tl.load(p_k2, boundary_check=(0, 1)) + p_k2 = k + (i_tc2 + o_i)[:, None] * (H*K) + o_k[None, :] + b_k2 = tl.load(p_k2, mask=m_tc2[:, None] & (o_k[None, :] < K), other=0.0) # diagonal block 2 - b_A22 += tl.dot(b_k2, tl.trans(b_k2)) + b_A22 = tl.dot(b_k2, tl.trans(b_k2), b_A22) # off-diagonal (2,0), (2,1) - b_A20 += tl.dot(b_k2, tl.trans(b_k0)) - b_A21 += tl.dot(b_k2, tl.trans(b_k1)) + b_A20 = tl.dot(b_k2, tl.trans(b_k0), b_A20) + b_A21 = tl.dot(b_k2, tl.trans(b_k1), b_A21) if i_tc3 < T: - p_k3 = tl.make_block_ptr(k, (T, K), (H*K, 1), (i_tc3, i_k * BK), (BC, BK), (1, 0)) - b_k3 = tl.load(p_k3, boundary_check=(0, 1)) + p_k3 = k + (i_tc3 + o_i)[:, None] * (H*K) + o_k[None, :] + b_k3 = tl.load(p_k3, mask=m_tc3[:, None] & (o_k[None, :] < K), other=0.0) # diagonal block 3 - b_A33 += tl.dot(b_k3, tl.trans(b_k3)) + b_A33 = tl.dot(b_k3, tl.trans(b_k3), b_A33) # off-diagonal (3,0), (3,1), (3,2) - b_A30 += tl.dot(b_k3, tl.trans(b_k0)) - b_A31 += tl.dot(b_k3, tl.trans(b_k1)) - b_A32 += tl.dot(b_k3, tl.trans(b_k2)) + b_A30 = tl.dot(b_k3, tl.trans(b_k0), b_A30) + b_A31 = tl.dot(b_k3, tl.trans(b_k1), b_A31) + b_A32 = tl.dot(b_k3, tl.trans(b_k2), b_A32) ############################################################################ # Step 2: apply gate and beta scaling @@ -292,29 +294,42 @@ def chunk_gated_delta_rule_fwd_kkt_solve_kernel( # Step 5: store full (I + A)^{-1} to output A ############################################################################ - p_A00 = tl.make_block_ptr(A, (T, BT), (HV*BT, 1), (i_tc0, 0), (BC, BC), (1, 0)) - p_A10 = tl.make_block_ptr(A, (T, BT), (HV*BT, 1), (i_tc1, 0), (BC, BC), (1, 0)) - p_A11 = tl.make_block_ptr(A, (T, BT), (HV*BT, 1), (i_tc1, BC), (BC, BC), (1, 0)) - p_A20 = tl.make_block_ptr(A, (T, BT), (HV*BT, 1), (i_tc2, 0), (BC, BC), (1, 0)) - p_A21 = tl.make_block_ptr(A, (T, BT), (HV*BT, 1), (i_tc2, BC), (BC, BC), (1, 0)) - p_A22 = tl.make_block_ptr(A, (T, BT), (HV*BT, 1), (i_tc2, 2*BC), (BC, BC), (1, 0)) - p_A30 = tl.make_block_ptr(A, (T, BT), (HV*BT, 1), (i_tc3, 0), (BC, BC), (1, 0)) - p_A31 = tl.make_block_ptr(A, (T, BT), (HV*BT, 1), (i_tc3, BC), (BC, BC), (1, 0)) - p_A32 = tl.make_block_ptr(A, (T, BT), (HV*BT, 1), (i_tc3, 2*BC), (BC, BC), (1, 0)) - p_A33 = tl.make_block_ptr(A, (T, BT), (HV*BT, 1), (i_tc3, 3*BC), (BC, BC), (1, 0)) - - tl.store(p_A00, b_Ai00.to(A.dtype.element_ty), boundary_check=(0, 1)) - tl.store(p_A10, b_Ai10.to(A.dtype.element_ty), boundary_check=(0, 1)) - tl.store(p_A11, b_Ai11.to(A.dtype.element_ty), boundary_check=(0, 1)) - tl.store(p_A20, b_Ai20.to(A.dtype.element_ty), boundary_check=(0, 1)) - tl.store(p_A21, b_Ai21.to(A.dtype.element_ty), boundary_check=(0, 1)) - tl.store(p_A22, b_Ai22.to(A.dtype.element_ty), boundary_check=(0, 1)) - tl.store(p_A30, b_Ai30.to(A.dtype.element_ty), boundary_check=(0, 1)) - tl.store(p_A31, b_Ai31.to(A.dtype.element_ty), boundary_check=(0, 1)) - tl.store(p_A32, b_Ai32.to(A.dtype.element_ty), boundary_check=(0, 1)) - tl.store(p_A33, b_Ai33.to(A.dtype.element_ty), boundary_check=(0, 1)) - - + p_A00 = A + (i_tc0 + o_i)[:, None] * (HV*BT) + o_i[None, :] + p_A10 = A + (i_tc1 + o_i)[:, None] * (HV*BT) + o_i[None, :] + p_A11 = A + (i_tc1 + o_i)[:, None] * (HV*BT) + (BC + o_i)[None, :] + p_A20 = A + (i_tc2 + o_i)[:, None] * (HV*BT) + o_i[None, :] + p_A21 = A + (i_tc2 + o_i)[:, None] * (HV*BT) + (BC + o_i)[None, :] + p_A22 = A + (i_tc2 + o_i)[:, None] * (HV*BT) + (2*BC + o_i)[None, :] + p_A30 = A + (i_tc3 + o_i)[:, None] * (HV*BT) + o_i[None, :] + p_A31 = A + (i_tc3 + o_i)[:, None] * (HV*BT) + (BC + o_i)[None, :] + p_A32 = A + (i_tc3 + o_i)[:, None] * (HV*BT) + (2*BC + o_i)[None, :] + p_A33 = A + (i_tc3 + o_i)[:, None] * (HV*BT) + (3*BC + o_i)[None, :] + + m_A0 = m_tc0[:, None] & (o_i[None, :] < BT) + m_A1 = m_tc1[:, None] & (o_i[None, :] < BT) + m_A2 = m_tc2[:, None] & (o_i[None, :] < BT) + m_A3 = m_tc3[:, None] & (o_i[None, :] < BT) + m_A11 = m_tc1[:, None] & ((BC + o_i)[None, :] < BT) + m_A21 = m_tc2[:, None] & ((BC + o_i)[None, :] < BT) + m_A22 = m_tc2[:, None] & ((2*BC + o_i)[None, :] < BT) + m_A31 = m_tc3[:, None] & ((BC + o_i)[None, :] < BT) + m_A32 = m_tc3[:, None] & ((2*BC + o_i)[None, :] < BT) + m_A33 = m_tc3[:, None] & ((3*BC + o_i)[None, :] < BT) + + tl.store(p_A00, b_Ai00.to(A.dtype.element_ty), mask=m_A0) + tl.store(p_A10, b_Ai10.to(A.dtype.element_ty), mask=m_A1) + tl.store(p_A11, b_Ai11.to(A.dtype.element_ty), mask=m_A11) + tl.store(p_A20, b_Ai20.to(A.dtype.element_ty), mask=m_A2) + tl.store(p_A21, b_Ai21.to(A.dtype.element_ty), mask=m_A21) + tl.store(p_A22, b_Ai22.to(A.dtype.element_ty), mask=m_A22) + tl.store(p_A30, b_Ai30.to(A.dtype.element_ty), mask=m_A3) + tl.store(p_A31, b_Ai31.to(A.dtype.element_ty), mask=m_A31) + tl.store(p_A32, b_Ai32.to(A.dtype.element_ty), mask=m_A32) + tl.store(p_A33, b_Ai33.to(A.dtype.element_ty), mask=m_A33) + + +# avoid dispatches +# @dispatch('gated_delta_rule') def chunk_gated_delta_rule_fwd_intra( k: torch.Tensor, v: torch.Tensor, diff --git a/build/torch-cuda/ops/gated_delta_rule/fused_recurrent.py b/build/torch-cuda/ops/gated_delta_rule/fused_recurrent.py index 978e34838d59d865b04eaf9a1d6ecc3c6adaf0da..aa6b2ce5dfc82a68f24f6e324e5eeeb3e48a2c24 100644 --- a/build/torch-cuda/ops/gated_delta_rule/fused_recurrent.py +++ b/build/torch-cuda/ops/gated_delta_rule/fused_recurrent.py @@ -63,7 +63,9 @@ def fused_recurrent_gated_delta_rule_fwd_kernel( APPLY_BETA_SIGMOID: tl.constexpr, ALLOW_NEG_EIGVAL: tl.constexpr, ): - i_v, i_nh = tl.program_id(0), tl.program_id(1) + pid = tl.program_id(0).to(tl.int64) + NV = tl.cdiv(V, BV) + i_v, i_nh = pid % NV, (pid // NV).to(tl.int64) i_n, i_hv = i_nh // HV, i_nh % HV i_h = i_hv // (HV // H) @@ -136,14 +138,14 @@ def fused_recurrent_gated_delta_rule_fwd_kernel( b_h *= exp(b_g) if USE_GK: - b_gk = tl.load(p_gk).to(tl.float32) + b_gk = tl.load(p_gk, mask=mask_k, other=0).to(tl.float32) if STATE_V_FIRST: b_h *= exp(b_gk[None, :]) else: b_h *= exp(b_gk[:, None]) if USE_GV: - b_gv = tl.load(p_gv).to(tl.float32) + b_gv = tl.load(p_gv, mask=mask_v, other=0).to(tl.float32) if STATE_V_FIRST: b_h *= exp(b_gv[:, None]) else: @@ -214,7 +216,7 @@ def fused_recurrent_gated_delta_rule_fwd( else: final_state = None - grid = (NV, N * HV) + grid = (NV * N * HV,) fused_recurrent_gated_delta_rule_fwd_kernel[grid]( q=q, k=k, @@ -390,7 +392,7 @@ def fused_recurrent_gated_delta_rule( >>> import torch >>> import torch.nn.functional as F >>> from einops import rearrange - >>> from ...ops.gated_delta_rule import fused_recurrent_gated_delta_rule + >>> from fla.ops.gated_delta_rule import fused_recurrent_gated_delta_rule # inputs with equal lengths >>> B, T, H, HV, K, V = 4, 2048, 4, 8, 512, 512 >>> q = torch.randn(B, T, H, K, device='cuda') diff --git a/build/torch-cuda/ops/gated_delta_rule/gate.py b/build/torch-cuda/ops/gated_delta_rule/gate.py index bf478e67dc1efdeacf56b88d21b1dd0243124534..18dd91ee4e40dcf2e22168edef5672443e012af5 100644 --- a/build/torch-cuda/ops/gated_delta_rule/gate.py +++ b/build/torch-cuda/ops/gated_delta_rule/gate.py @@ -10,6 +10,7 @@ import torch.nn.functional as F import triton import triton.language as tl +from ...backends import dispatch from ...ops.utils.cache import fla_cache_autotune from ...ops.utils.index import prepare_chunk_indices from ...ops.utils.op import exp @@ -75,20 +76,22 @@ def gdn_gate_chunk_cumsum_scalar_kernel( HAS_SCALE: tl.constexpr, IS_VARLEN: tl.constexpr, ): - i_t, i_bh = tl.program_id(0), tl.program_id(1) + i_t, i_bh = tl.program_id(0).to(tl.int64), tl.program_id(1).to(tl.int64) i_b, i_h = i_bh // H, i_bh % H 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.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) + i_n, i_t = tl.load(chunk_indices + i_t * 2).to(tl.int32), tl.load(chunk_indices + i_t * 2 + 1).to(tl.int64) + bos, eos = tl.load(cu_seqlens + i_n).to(tl.int64), tl.load(cu_seqlens + i_n + 1).to(tl.int64) T = eos - bos else: bos, eos = i_b * T, i_b * T + T - p_g = tl.make_block_ptr(g + bos * H + i_h, (T,), (H,), (i_t * BT,), (BT,), (0,)) - p_o = tl.make_block_ptr(o + bos * H + i_h, (T,), (H,), (i_t * BT,), (BT,), (0,)) + o_t = i_t * BT + tl.arange(0, BT) + m_t = o_t < T + p_g = g + bos * H + i_h + o_t * H + p_o = o + bos * H + i_h + o_t * H - b_g = tl.load(p_g, boundary_check=(0,)).to(tl.float32) + b_g = tl.load(p_g, mask=m_t, other=0.0).to(tl.float32) if HAS_BIAS: b_g = b_g + tl.load(dt_bias + i_h).to(tl.float32) b_A = tl.load(A_log + i_h).to(tl.float32) @@ -100,7 +103,7 @@ def gdn_gate_chunk_cumsum_scalar_kernel( b_o = -b_o + b_z[None] + b_gate if HAS_SCALE: b_o *= scale - tl.store(p_o, b_o.to(p_o.dtype.element_ty), boundary_check=(0,)) + tl.store(p_o, b_o.to(p_o.dtype.element_ty), mask=m_t) @triton.heuristics({ @@ -122,21 +125,24 @@ def gdn_gate_bwd_kernel( dyg, dg, dA, + dbias, T, H: tl.constexpr, BT: tl.constexpr, HAS_BIAS: tl.constexpr, ): - i_t, i_h = tl.program_id(0), tl.program_id(1) + i_t, i_h = tl.program_id(0).to(tl.int64), tl.program_id(1).to(tl.int64) b_A = tl.load(A_log + i_h).to(tl.float32) - p_g = tl.make_block_ptr(g + i_h, (T,), (H,), (i_t * BT,), (BT,), (0,)) - p_dg = tl.make_block_ptr(dg + i_h, (T,), (H,), (i_t * BT,), (BT,), (0,)) - p_dyg = tl.make_block_ptr(dyg + i_h, (T,), (H,), (i_t * BT,), (BT,), (0,)) + o_t = i_t * BT + tl.arange(0, BT) + m_t = o_t < T + p_g = g + i_h + o_t * H + p_dg = dg + i_h + o_t * H + p_dyg = dyg + i_h + o_t * H - b_g = tl.load(p_g, boundary_check=(0,)).to(tl.float32) - b_dyg = tl.load(p_dyg, boundary_check=(0,)).to(tl.float32) + b_g = tl.load(p_g, mask=m_t, other=0.0).to(tl.float32) + b_dyg = tl.load(p_dyg, mask=m_t, other=0.0).to(tl.float32) if HAS_BIAS: b_g = b_g + tl.load(dt_bias + i_h).to(tl.float32) @@ -149,11 +155,15 @@ def gdn_gate_bwd_kernel( b_dg = b_neg_expA * (b_dyg * tl.sigmoid(b_g)) b_dA = tl.sum(b_dyg * b_yg, 0) - tl.store(p_dg, b_dg.to(p_dg.dtype.element_ty), boundary_check=(0,)) + tl.store(p_dg, b_dg.to(p_dg.dtype.element_ty), mask=m_t) tl.store(dA + i_t * H + i_h, b_dA) + if HAS_BIAS: + b_dbias = tl.sum(b_dg, 0) + tl.store(dbias + i_t * H + i_h, b_dbias) @input_guard +@dispatch('gated_delta_rule') def gdn_gate_chunk_cumsum( g: torch.Tensor, A_log: torch.Tensor, @@ -187,6 +197,7 @@ def gdn_gate_chunk_cumsum( return o +@dispatch('gated_delta_rule') def gdn_gate_bwd( g: torch.Tensor, A_log: torch.Tensor, @@ -198,8 +209,9 @@ def gdn_gate_bwd( BT = 32 NT = triton.cdiv(T, BT) - dg = torch.empty_like(g, dtype=torch.float32) + dg = torch.empty_like(g) dA = A_log.new_empty(NT, H, dtype=torch.float32) + dbias = g.new_empty(NT, H, dtype=torch.float32) if dt_bias is not None else None gdn_gate_bwd_kernel[(NT, H)]( g=g, @@ -208,14 +220,14 @@ def gdn_gate_bwd( dyg=dyg, dg=dg, dA=dA, + dbias=dbias, T=T, H=H, BT=BT, ) - dg = dg.view_as(g).type_as(g) dA = dA.sum(0).view_as(A_log).type_as(A_log) - dbias = dg.view(-1, H).sum(0).to(dt_bias) if dt_bias is not None else None + dbias = dbias.sum(0).to(dt_bias) if dt_bias is not None else None return dg, dA, dbias @@ -244,19 +256,22 @@ def gdn_gate_fwd_kernel( BT: tl.constexpr, HAS_BIAS: tl.constexpr, ): - i_t, i_h = tl.program_id(0), tl.program_id(1) + i_t, i_h = tl.program_id(0).to(tl.int64), tl.program_id(1).to(tl.int64) b_A = tl.load(A_log + i_h).to(tl.float32) - p_g = tl.make_block_ptr(g + i_h, (T,), (H,), (i_t * BT,), (BT,), (0,)) - p_yg = tl.make_block_ptr(yg + i_h, (T,), (H,), (i_t * BT,), (BT,), (0,)) - b_g = tl.load(p_g, boundary_check=(0,)).to(tl.float32) + o_t = i_t * BT + tl.arange(0, BT) + m_t = o_t < T + p_g = g + i_h + o_t * H + p_yg = yg + i_h + o_t * H + b_g = tl.load(p_g, mask=m_t, other=0.0).to(tl.float32) if HAS_BIAS: b_g = b_g + tl.load(dt_bias + i_h).to(tl.float32) b_yg = -exp(b_A) * softplus(b_g) - tl.store(p_yg, b_yg.to(p_yg.dtype.element_ty), boundary_check=(0,)) + tl.store(p_yg, b_yg.to(p_yg.dtype.element_ty), mask=m_t) +@dispatch('gated_delta_rule') def gdn_gate_fwd( g: torch.Tensor, A_log: torch.Tensor, diff --git a/build/torch-cuda/ops/gated_delta_rule/wy_fast.py b/build/torch-cuda/ops/gated_delta_rule/wy_fast.py index d212ce610dc7b34d36584ff0f7323f5ccb4f062e..4a67e102874f57f353dba8d9b34edac3176bfb26 100644 --- a/build/torch-cuda/ops/gated_delta_rule/wy_fast.py +++ b/build/torch-cuda/ops/gated_delta_rule/wy_fast.py @@ -9,6 +9,7 @@ import torch import triton import triton.language as tl +from ...backends import dispatch from ...ops.utils import prepare_chunk_indices from ...ops.utils.cache import fla_cache_autotune from ...ops.utils.op import exp2 @@ -60,41 +61,49 @@ def recompute_w_u_fwd_kernel( USE_G: tl.constexpr, IS_VARLEN: tl.constexpr, ): - i_t, i_bh = tl.program_id(0), tl.program_id(1).to(tl.int64) + i_t, i_bh = tl.program_id(0).to(tl.int64), tl.program_id(1).to(tl.int64) i_b, i_h = i_bh // HV, i_bh % HV 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.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) + i_n, i_t = tl.load(chunk_indices + i_t * 2).to(tl.int32), tl.load(chunk_indices + i_t * 2 + 1).to(tl.int64) + bos, eos = tl.load(cu_seqlens + i_n).to(tl.int64), tl.load(cu_seqlens + i_n + 1).to(tl.int64) T = eos - bos else: bos, eos = i_b * T, i_b * T + T - p_b = tl.make_block_ptr(beta + bos*HV + i_h, (T,), (HV,), (i_t * BT,), (BT,), (0,)) - b_b = tl.load(p_b, boundary_check=(0,)) + o_t = i_t * BT + tl.arange(0, BT) + o_A = tl.arange(0, BT) + m_t = o_t < T + m_A = m_t[:, None] & (o_A[None, :] < BT) + p_b = beta + bos*HV + i_h + o_t * HV + b_b = tl.load(p_b, mask=m_t, other=0.0) - p_A = tl.make_block_ptr(A + (bos*HV + i_h) * BT, (T, BT), (HV*BT, 1), (i_t * BT, 0), (BT, BT), (1, 0)) - b_A = tl.load(p_A, boundary_check=(0, 1)) + p_A = A + (bos*HV + i_h) * BT + o_t[:, None] * (HV*BT) + o_A[None, :] + b_A = tl.load(p_A, mask=m_A, other=0.0) for i_v in range(tl.cdiv(V, BV)): - p_v = tl.make_block_ptr(v + (bos*HV + i_h) * V, (T, V), (HV*V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) - p_u = tl.make_block_ptr(u + (bos*HV + i_h) * V, (T, V), (HV*V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) - b_v = tl.load(p_v, boundary_check=(0, 1)) + o_v = i_v * BV + tl.arange(0, BV) + m_v = m_t[:, None] & (o_v[None, :] < V) + p_v = v + (bos*HV + i_h) * V + o_t[:, None] * (HV*V) + o_v[None, :] + p_u = u + (bos*HV + i_h) * V + o_t[:, None] * (HV*V) + o_v[None, :] + b_v = tl.load(p_v, mask=m_v, other=0.0) b_vb = (b_v * b_b[:, None]).to(b_v.dtype) b_u = tl.dot(b_A, b_vb, allow_tf32=False) - tl.store(p_u, b_u.to(p_u.dtype.element_ty), boundary_check=(0, 1)) + tl.store(p_u, b_u.to(p_u.dtype.element_ty), mask=m_v) if USE_G: - p_g = tl.make_block_ptr(g + (bos*HV + i_h), (T,), (HV,), (i_t * BT,), (BT,), (0,)) - b_g = exp2(tl.load(p_g, boundary_check=(0,))) + p_g = g + (bos*HV + i_h) + o_t * HV + b_g = exp2(tl.load(p_g, mask=m_t, other=0.0)) for i_k in range(tl.cdiv(K, BK)): - p_k = tl.make_block_ptr(k + (bos*H + i_h // (HV // H)) * K, (T, K), (H*K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) - p_w = tl.make_block_ptr(w + (bos*HV + i_h) * K, (T, K), (HV*K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) - b_k = tl.load(p_k, boundary_check=(0, 1)) + o_k = i_k * BK + tl.arange(0, BK) + m_k = m_t[:, None] & (o_k[None, :] < K) + p_k = k + (bos*H + i_h // (HV // H)) * K + o_t[:, None] * (H*K) + o_k[None, :] + p_w = w + (bos*HV + i_h) * K + o_t[:, None] * (HV*K) + o_k[None, :] + b_k = tl.load(p_k, mask=m_k, other=0.0) b_kb = b_k * b_b[:, None] if USE_G: b_kb *= b_g[:, None] b_w = tl.dot(b_A, b_kb.to(b_k.dtype)) - tl.store(p_w, b_w.to(p_w.dtype.element_ty), boundary_check=(0, 1)) + tl.store(p_w, b_w.to(p_w.dtype.element_ty), mask=m_k) @triton.heuristics({ @@ -136,43 +145,49 @@ def prepare_wy_repr_bwd_kernel( USE_G: tl.constexpr, IS_VARLEN: tl.constexpr, ): - i_t, i_bh = tl.program_id(0), tl.program_id(1).to(tl.int64) + i_t, i_bh = tl.program_id(0).to(tl.int64), tl.program_id(1).to(tl.int64) i_b, i_h = i_bh // HV, i_bh % HV 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.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) + i_n, i_t = tl.load(chunk_indices + i_t * 2).to(tl.int32), tl.load(chunk_indices + i_t * 2 + 1).to(tl.int64) + bos, eos = tl.load(cu_seqlens + i_n).to(tl.int64), tl.load(cu_seqlens + i_n + 1).to(tl.int64) T = eos - bos else: bos, eos = i_b * T, i_b * T + T - p_b = tl.make_block_ptr(beta + (bos*HV + i_h), (T,), (HV,), (i_t * BT,), (BT,), (0,)) - p_db = tl.make_block_ptr(db + (bos*HV + i_h), (T,), (HV,), (i_t * BT,), (BT,), (0,)) - p_A = tl.make_block_ptr(A + (bos*HV + i_h) * BT, (BT, T), (1, HV*BT), (0, i_t * BT), (BT, BT), (0, 1)) + o_t = i_t * BT + tl.arange(0, BT) + o_A = tl.arange(0, BT) + m_t = o_t < T + m_AT = (o_A[:, None] < BT) & m_t[None, :] + p_b = beta + (bos*HV + i_h) + o_t * HV + p_db = db + (bos*HV + i_h) + o_t * HV + p_A = A + (bos*HV + i_h) * BT + o_A[:, None] + o_t[None, :] * (HV*BT) - b_b = tl.load(p_b, boundary_check=(0,)) + b_b = tl.load(p_b, mask=m_t, other=0.0) b_db = tl.zeros([BT], dtype=tl.float32) - b_A = tl.load(p_A, boundary_check=(0, 1)) + b_A = tl.load(p_A, mask=m_AT, other=0.0) b_dA = tl.zeros([BT, BT], dtype=tl.float32) if USE_G: - p_g = tl.make_block_ptr(g + (bos*HV + i_h), (T,), (HV,), (i_t * BT,), (BT,), (0,)) - b_g = tl.load(p_g, boundary_check=(0,)) + p_g = g + (bos*HV + i_h) + o_t * HV + b_g = tl.load(p_g, mask=m_t, other=0.0) b_g_exp = exp2(b_g) b_dg = tl.zeros([BT], dtype=tl.float32) for i_k in range(tl.cdiv(K, BK)): - p_k = tl.make_block_ptr(k + (bos*H + i_h // (HV // H)) * K, (T, K), (H*K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) - p_dk = tl.make_block_ptr(dk + (bos*HV + i_h) * K, (T, K), (HV*K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) - p_dw = tl.make_block_ptr(dw + (bos*HV + i_h) * K, (T, K), (HV*K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) + o_k = i_k * BK + tl.arange(0, BK) + m_k = m_t[:, None] & (o_k[None, :] < K) + p_k = k + (bos*H + i_h // (HV // H)) * K + o_t[:, None] * (H*K) + o_k[None, :] + p_dk = dk + (bos*HV + i_h) * K + o_t[:, None] * (HV*K) + o_k[None, :] + p_dw = dw + (bos*HV + i_h) * K + o_t[:, None] * (HV*K) + o_k[None, :] # [BT, BK] - b_k = tl.load(p_k, boundary_check=(0, 1)) + b_k = tl.load(p_k, mask=m_k, other=0.0) if USE_G: b_kbg = b_k * (b_b * b_g_exp)[:, None] else: b_kbg = b_k * b_b[:, None] - b_dw = tl.load(p_dw, boundary_check=(0, 1)) + b_dw = tl.load(p_dw, mask=m_k, other=0.0) - b_dA += tl.dot(b_dw, tl.trans(b_kbg).to(b_dw.dtype)) + b_dA = tl.dot(b_dw, tl.trans(b_kbg).to(b_dw.dtype), b_dA) b_dkbg = tl.dot(b_A, b_dw) if USE_G: b_dk = b_dkbg * (b_g_exp * b_b)[:, None] @@ -181,23 +196,23 @@ def prepare_wy_repr_bwd_kernel( else: b_dk = b_dkbg * b_b[:, None] b_db += tl.sum(b_dkbg * b_k, 1) - tl.store(p_dk, b_dk.to(p_dk.dtype.element_ty), boundary_check=(0, 1)) + tl.store(p_dk, b_dk.to(p_dk.dtype.element_ty), mask=m_k) for i_v in range(tl.cdiv(V, BV)): - p_v = tl.make_block_ptr(v + (bos*HV + i_h) * V, (T, V), (HV*V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) - p_dv = tl.make_block_ptr(dv + (bos*HV + i_h) * V, (T, V), (HV*V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) - p_du = tl.make_block_ptr(du + (bos*HV + i_h) * V, (T, V), (HV*V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) - b_v = tl.load(p_v, boundary_check=(0, 1)) + o_v = i_v * BV + tl.arange(0, BV) + m_v = m_t[:, None] & (o_v[None, :] < V) + p_v = v + (bos*HV + i_h) * V + o_t[:, None] * (HV*V) + o_v[None, :] + p_dv = dv + (bos*HV + i_h) * V + o_t[:, None] * (HV*V) + o_v[None, :] + p_du = du + (bos*HV + i_h) * V + o_t[:, None] * (HV*V) + o_v[None, :] + b_v = tl.load(p_v, mask=m_v, other=0.0) b_vb = (b_v * b_b[:, None]).to(b_v.dtype) - b_du = tl.load(p_du, boundary_check=(0, 1)) - b_dA += tl.dot(b_du, tl.trans(b_vb)) + b_du = tl.load(p_du, mask=m_v, other=0.0) + b_dA = tl.dot(b_du, tl.trans(b_vb), b_dA) b_dvb = tl.dot(b_A, b_du) b_dv = b_dvb * b_b[:, None] b_db += tl.sum(b_dvb * b_v, 1) - tl.store(p_dv, b_dv.to(p_dv.dtype.element_ty), boundary_check=(0, 1)) + tl.store(p_dv, b_dv.to(p_dv.dtype.element_ty), mask=m_v) - o_t = i_t * BT + tl.arange(0, BT) - m_t = o_t < T m_A = (o_t[:, None] > o_t[None, :]) & (m_t[:, None] & m_t) b_dA = tl.where(m_A, b_dA, 0) b_dA = tl.dot(b_dA.to(b_A.dtype), b_A) @@ -211,29 +226,33 @@ def prepare_wy_repr_bwd_kernel( tl.debug_barrier() for i_k in range(tl.cdiv(K, BK)): - p_k = tl.make_block_ptr(k + (bos*H + i_h // (HV // H)) * K, (T, K), (H*K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) - p_dk = tl.make_block_ptr(dk + (bos*HV + i_h) * K, (T, K), (HV*K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) - b_k = tl.load(p_k, boundary_check=(0, 1)) + o_k = i_k * BK + tl.arange(0, BK) + m_k = m_t[:, None] & (o_k[None, :] < K) + p_k = k + (bos*H + i_h // (HV // H)) * K + o_t[:, None] * (H*K) + o_k[None, :] + p_dk = dk + (bos*HV + i_h) * K + o_t[:, None] * (HV*K) + o_k[None, :] + b_k = tl.load(p_k, mask=m_k, other=0.0) b_kt = tl.trans(b_k) b_kb = b_k * b_b[:, None] - b_A += tl.dot(b_k, b_kt) + b_A = tl.dot(b_k, b_kt, b_A) b_dkb = tl.dot(b_dA, b_k) b_db += tl.sum(b_dkb * b_k, 1) b_dk = b_dkb * b_b[:, None] + tl.trans(tl.dot(tl.trans(b_kb).to(b_dA.dtype), b_dA)) - b_dk += tl.load(p_dk, boundary_check=(0, 1)) + b_dk += tl.load(p_dk, mask=m_k, other=0.0) - tl.store(p_dk, b_dk.to(p_dk.dtype.element_ty), boundary_check=(0, 1)) - tl.store(p_db, b_db.to(p_db.dtype.element_ty), boundary_check=(0,)) + tl.store(p_dk, b_dk.to(p_dk.dtype.element_ty), mask=m_k) + tl.store(p_db, b_db.to(p_db.dtype.element_ty), mask=m_t) b_A *= b_b[:, None] if USE_G: b_AdA = b_dA * b_A - p_dg = tl.make_block_ptr(dg + (bos*HV + i_h), (T,), (HV,), (i_t * BT,), (BT,), (0,)) + p_dg = dg + (bos*HV + i_h) + o_t * HV b_dg += tl.sum(b_AdA, axis=1) - tl.sum(b_AdA, axis=0) - tl.store(p_dg, b_dg.to(p_dg.dtype.element_ty), boundary_check=(0,)) + tl.store(p_dg, b_dg.to(p_dg.dtype.element_ty), mask=m_t) +# avoid dispatches +# @dispatch('gated_delta_rule') def recompute_w_u_fwd( k: torch.Tensor, v: torch.Tensor, @@ -276,6 +295,7 @@ def recompute_w_u_fwd( return w, u +# @dispatch('gated_delta_rule') def prepare_wy_repr_bwd( k: torch.Tensor, v: torch.Tensor, diff --git a/build/torch-cuda/ops/gated_oja_rule/chunk_h.py b/build/torch-cuda/ops/gated_oja_rule/chunk_h.py index 7f21802673379180becd59948162534db18ad108..59f67adbfb385a849870e4d330dc2affe929371c 100644 --- a/build/torch-cuda/ops/gated_oja_rule/chunk_h.py +++ b/build/torch-cuda/ops/gated_oja_rule/chunk_h.py @@ -10,7 +10,7 @@ import triton import triton.language as tl from ...ops.utils import prepare_chunk_indices, prepare_chunk_offsets -from ...ops.utils.op import exp +from ...ops.utils.op import exp, unflatten_program_id from ...utils import check_shared_mem, is_nvidia_hopper BKV_LIST = [64, 128] if check_shared_mem() else [32, 64] @@ -57,14 +57,14 @@ def chunk_oja_fwd_kernel_h_blockdim64( SAVE_NEW_KEY: tl.constexpr, IS_VARLEN: tl.constexpr, ): - # (triton.cdiv(K, meta['BK']), N*H) - i_k, i_nh = tl.program_id(0), tl.program_id(1) + # (triton.cdiv(K, meta['BK']) * N * H,) + i_k, i_nh = unflatten_program_id(tl.cdiv(K, BK)) i_n, i_h = i_nh // H, i_nh % H if IS_VARLEN: - bos, eos = tl.load(cu_seqlens + i_n).to(tl.int32), tl.load(cu_seqlens + i_n + 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 NT = tl.cdiv(T, BT) - boh = tl.load(chunk_offsets + i_n).to(tl.int32) + boh = tl.load(chunk_offsets + i_n).to(tl.int64) else: bos, eos = i_n * T, i_n * T + T NT = tl.cdiv(T, BT) @@ -94,61 +94,66 @@ def chunk_oja_fwd_kernel_h_blockdim64( if STORE_FINAL_STATE: ht = ht + i_nh * K*V + o_k = i_k * BK + tl.arange(0, BK) + o_v1 = tl.arange(0, 64) # load initial state if USE_INITIAL_STATE: - p_h0_1 = tl.make_block_ptr(h0, (K, V), (V, 1), (i_k * BK, 0), (BK, 64), (1, 0)) - b_h1 += tl.load(p_h0_1, boundary_check=(0, 1)).to(tl.float32) + p_h0_1 = h0 + o_k[:, None] * V + o_v1[None, :] + b_h1 += tl.load(p_h0_1, mask=(o_k[:, None] < K) & (o_v1[None, :] < V), other=0.0).to(tl.float32) if V > 64: - p_h0_2 = tl.make_block_ptr(h0, (K, V), (V, 1), (i_k * BK, 64), (BK, 64), (1, 0)) - b_h2 += tl.load(p_h0_2, boundary_check=(0, 1)).to(tl.float32) + p_h0_2 = h0 + o_k[:, None] * V + (64 + o_v1)[None, :] + b_h2 += tl.load(p_h0_2, mask=(o_k[:, None] < K) & ((64 + o_v1)[None, :] < V), other=0.0).to(tl.float32) if V > 128: - p_h0_3 = tl.make_block_ptr(h0, (K, V), (V, 1), (i_k * BK, 128), (BK, 64), (1, 0)) - b_h3 += tl.load(p_h0_3, boundary_check=(0, 1)).to(tl.float32) + p_h0_3 = h0 + o_k[:, None] * V + (128 + o_v1)[None, :] + b_h3 += tl.load(p_h0_3, mask=(o_k[:, None] < K) & ((128 + o_v1)[None, :] < V), other=0.0).to(tl.float32) if V > 192: - p_h0_4 = tl.make_block_ptr(h0, (K, V), (V, 1), (i_k * BK, 192), (BK, 64), (1, 0)) - b_h4 += tl.load(p_h0_4, boundary_check=(0, 1)).to(tl.float32) + p_h0_4 = h0 + o_k[:, None] * V + (192 + o_v1)[None, :] + b_h4 += tl.load(p_h0_4, mask=(o_k[:, None] < K) & ((192 + o_v1)[None, :] < V), other=0.0).to(tl.float32) # main recurrence for i_t in range(NT): - p_h1 = tl.make_block_ptr(h + i_t * stride_h, (K, V), (V, 1), (i_k * BK, 0), (BK, 64), (1, 0)) - tl.store(p_h1, b_h1.to(p_h1.dtype.element_ty), boundary_check=(0, 1)) + o_t = i_t.to(tl.int64) * BT + tl.arange(0, BT) + m_t = o_t < T + m_tk = m_t[:, None] & (o_k[None, :] < K) + m_kv1 = (o_k[:, None] < K) & (o_v1[None, :] < V) + p_h1 = h + i_t.to(tl.int64) * stride_h + o_k[:, None] * V + o_v1[None, :] + tl.store(p_h1, b_h1.to(p_h1.dtype.element_ty), mask=m_kv1) if V > 64: - p_h2 = tl.make_block_ptr(h + i_t * stride_h, (K, V), (V, 1), (i_k * BK, 64), (BK, 64), (1, 0)) - tl.store(p_h2, b_h2.to(p_h2.dtype.element_ty), boundary_check=(0, 1)) + p_h2 = h + i_t.to(tl.int64) * stride_h + o_k[:, None] * V + (64 + o_v1)[None, :] + tl.store(p_h2, b_h2.to(p_h2.dtype.element_ty), mask=(o_k[:, None] < K) & ((64 + o_v1)[None, :] < V)) if V > 128: - p_h3 = tl.make_block_ptr(h + i_t * stride_h, (K, V), (V, 1), (i_k * BK, 128), (BK, 64), (1, 0)) - tl.store(p_h3, b_h3.to(p_h3.dtype.element_ty), boundary_check=(0, 1)) + p_h3 = h + i_t.to(tl.int64) * stride_h + o_k[:, None] * V + (128 + o_v1)[None, :] + tl.store(p_h3, b_h3.to(p_h3.dtype.element_ty), mask=(o_k[:, None] < K) & ((128 + o_v1)[None, :] < V)) if V > 192: - p_h4 = tl.make_block_ptr(h + i_t * stride_h, (K, V), (V, 1), (i_k * BK, 192), (BK, 64), (1, 0)) - tl.store(p_h4, b_h4.to(p_h4.dtype.element_ty), boundary_check=(0, 1)) + p_h4 = h + i_t.to(tl.int64) * stride_h + o_k[:, None] * V + (192 + o_v1)[None, :] + tl.store(p_h4, b_h4.to(p_h4.dtype.element_ty), mask=(o_k[:, None] < K) & ((192 + o_v1)[None, :] < V)) - p_w = tl.make_block_ptr(w, (T, V), (stride_v, 1), (i_t * BT, 0), (BT, 64), (1, 0)) - b_w = tl.load(p_w, boundary_check=(0, 1)) + p_w = w + o_t[:, None] * stride_v + o_v1[None, :] + b_w = tl.load(p_w, mask=m_t[:, None] & (o_v1[None, :] < V), other=0.0) b_k = tl.dot(b_w, tl.trans(b_h1).to(b_w.dtype)) # BT BK if V > 64: - p_w = tl.make_block_ptr(w, (T, V), (stride_v, 1), (i_t * BT, 64), (BT, 64), (1, 0)) - b_w = tl.load(p_w, boundary_check=(0, 1)) - b_k += tl.dot(b_w, tl.trans(b_h2).to(b_w.dtype)) + p_w = w + o_t[:, None] * stride_v + (64 + o_v1)[None, :] + b_w = tl.load(p_w, mask=m_t[:, None] & ((64 + o_v1)[None, :] < V), other=0.0) + b_k = tl.dot(b_w, tl.trans(b_h2).to(b_w.dtype), b_k) if V > 128: - p_w = tl.make_block_ptr(w, (T, V), (stride_v, 1), (i_t * BT, 128), (BT, 64), (1, 0)) - b_w = tl.load(p_w, boundary_check=(0, 1)) - b_k += tl.dot(b_w, tl.trans(b_h3).to(b_w.dtype)) + p_w = w + o_t[:, None] * stride_v + (128 + o_v1)[None, :] + b_w = tl.load(p_w, mask=m_t[:, None] & ((128 + o_v1)[None, :] < V), other=0.0) + b_k = tl.dot(b_w, tl.trans(b_h3).to(b_w.dtype), b_k) if V > 192: - p_w = tl.make_block_ptr(w, (T, V), (stride_v, 1), (i_t * BT, 192), (BT, 64), (1, 0)) - b_w = tl.load(p_w, boundary_check=(0, 1)) - b_k += tl.dot(b_w, tl.trans(b_h4).to(b_w.dtype)) + p_w = w + o_t[:, None] * stride_v + (192 + o_v1)[None, :] + b_w = tl.load(p_w, mask=m_t[:, None] & ((192 + o_v1)[None, :] < V), other=0.0) + b_k = tl.dot(b_w, tl.trans(b_h4).to(b_w.dtype), b_k) - p_u = tl.make_block_ptr(u, (T, K), (stride_k, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) - b_k = tl.load(p_u, boundary_check=(0, 1)) - b_k + p_u = u + o_t[:, None] * stride_k + o_k[None, :] + b_k = tl.load(p_u, mask=m_tk, other=0.0) - b_k if SAVE_NEW_KEY: - p_k = tl.make_block_ptr(k_new, (T, K), (stride_k, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) - tl.store(p_k, b_k.to(p_k.dtype.element_ty), boundary_check=(0, 1)) + p_k = k_new + o_t[:, None] * stride_k + o_k[None, :] + tl.store(p_k, b_k.to(p_k.dtype.element_ty), mask=m_tk) last_idx = min((i_t + 1) * BT, T) - 1 if USE_GV: - o_v1 = tl.arange(0, 64) b_gk_last1 = tl.load(gv + (bos + last_idx) * H*V + i_h * V + o_v1, mask=(o_v1 < V), other=0.) b_h1 *= exp(b_gk_last1)[None, :] if V > 64: @@ -166,34 +171,34 @@ def chunk_oja_fwd_kernel_h_blockdim64( b_k = b_k.to(v.dtype.element_ty) # BT BK - p_v = tl.make_block_ptr(v, (T, V), (stride_v, 1), (i_t * BT, 0), (BT, 64), (1, 0)) - b_v = tl.load(p_v, boundary_check=(0, 1)) # BT BV - b_h1 += tl.dot(tl.trans(b_k), b_v) + p_v = v + o_t[:, None] * stride_v + o_v1[None, :] + b_v = tl.load(p_v, mask=m_t[:, None] & (o_v1[None, :] < V), other=0.0) # BT BV + b_h1 = tl.dot(tl.trans(b_k), b_v, b_h1) if V > 64: - p_v = tl.make_block_ptr(v, (T, V), (stride_v, 1), (i_t * BT, 64), (BT, 64), (1, 0)) - b_v = tl.load(p_v, boundary_check=(0, 1)) - b_h2 += tl.dot(tl.trans(b_k), b_v) + p_v = v + o_t[:, None] * stride_v + (64 + o_v1)[None, :] + b_v = tl.load(p_v, mask=m_t[:, None] & ((64 + o_v1)[None, :] < V), other=0.0) + b_h2 = tl.dot(tl.trans(b_k), b_v, b_h2) if V > 128: - p_v = tl.make_block_ptr(v, (T, V), (stride_v, 1), (i_t * BT, 128), (BT, 64), (1, 0)) - b_v = tl.load(p_v, boundary_check=(0, 1)) - b_h3 += tl.dot(tl.trans(b_k), b_v) + p_v = v + o_t[:, None] * stride_v + (128 + o_v1)[None, :] + b_v = tl.load(p_v, mask=m_t[:, None] & ((128 + o_v1)[None, :] < V), other=0.0) + b_h3 = tl.dot(tl.trans(b_k), b_v, b_h3) if V > 192: - p_v = tl.make_block_ptr(v, (T, V), (stride_v, 1), (i_t * BT, 192), (BT, 64), (1, 0)) - b_v = tl.load(p_v, boundary_check=(0, 1)) - b_h4 += tl.dot(tl.trans(b_k), b_v) + p_v = v + o_t[:, None] * stride_v + (192 + o_v1)[None, :] + b_v = tl.load(p_v, mask=m_t[:, None] & ((192 + o_v1)[None, :] < V), other=0.0) + b_h4 = tl.dot(tl.trans(b_k), b_v, b_h4) # epilogue if STORE_FINAL_STATE: - p_ht = tl.make_block_ptr(ht, (K, V), (V, 1), (i_k * BK, 0), (BK, 64), (1, 0)) - tl.store(p_ht, b_h1.to(p_ht.dtype.element_ty), boundary_check=(0, 1)) + p_ht = ht + o_k[:, None] * V + o_v1[None, :] + tl.store(p_ht, b_h1.to(p_ht.dtype.element_ty), mask=(o_k[:, None] < K) & (o_v1[None, :] < V)) if V > 64: - p_ht = tl.make_block_ptr(ht, (K, V), (V, 1), (i_k * BK, 64), (BK, 64), (1, 0)) - tl.store(p_ht, b_h2.to(p_ht.dtype.element_ty), boundary_check=(0, 1)) + p_ht = ht + o_k[:, None] * V + (64 + o_v1)[None, :] + tl.store(p_ht, b_h2.to(p_ht.dtype.element_ty), mask=(o_k[:, None] < K) & ((64 + o_v1)[None, :] < V)) if V > 128: - p_ht = tl.make_block_ptr(ht, (K, V), (V, 1), (i_k * BK, 128), (BK, 64), (1, 0)) - tl.store(p_ht, b_h3.to(p_ht.dtype.element_ty), boundary_check=(0, 1)) + p_ht = ht + o_k[:, None] * V + (128 + o_v1)[None, :] + tl.store(p_ht, b_h3.to(p_ht.dtype.element_ty), mask=(o_k[:, None] < K) & ((128 + o_v1)[None, :] < V)) if V > 192: - p_ht = tl.make_block_ptr(ht, (K, V), (V, 1), (i_k * BK, 192), (BK, 64), (1, 0)) - tl.store(p_ht, b_h4.to(p_ht.dtype.element_ty), boundary_check=(0, 1)) + p_ht = ht + o_k[:, None] * V + (192 + o_v1)[None, :] + tl.store(p_ht, b_h4.to(p_ht.dtype.element_ty), mask=(o_k[:, None] < K) & ((192 + o_v1)[None, :] < V)) def chunk_oja_fwd_h( @@ -223,7 +228,7 @@ def chunk_oja_fwd_h( final_state = v.new_empty(N, H, K, V, dtype=torch.float32) if output_final_state else None k_new = torch.empty_like(u) if save_new_key else None - def grid(meta): return (triton.cdiv(K, meta['BK']), N*H) + def grid(meta): return (triton.cdiv(K, meta['BK']) * N * H,) chunk_oja_fwd_kernel_h_blockdim64[grid]( v=v, u=u, @@ -285,13 +290,13 @@ def chunk_oja_bwd_kernel_dhu_blockdim64( USE_FINAL_STATE_GRADIENT: tl.constexpr, IS_VARLEN: tl.constexpr ): - i_k, i_nh = tl.program_id(0), tl.program_id(1) + i_k, i_nh = unflatten_program_id(tl.cdiv(K, BK)) i_n, i_h = i_nh // H, i_nh % H if IS_VARLEN: - bos, eos = tl.load(cu_seqlens + i_n).to(tl.int32), tl.load(cu_seqlens + i_n + 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 NT = tl.cdiv(T, BT) - boh = tl.load(chunk_offsets + i_n).to(tl.int32) + boh = tl.load(chunk_offsets + i_n).to(tl.int64) else: bos, eos = i_n * T, i_n * T + T NT = tl.cdiv(T, BT) @@ -325,136 +330,143 @@ def chunk_oja_bwd_kernel_dhu_blockdim64( if USE_FINAL_STATE_GRADIENT: dht += i_nh * K*V + o_k = i_k * BK + tl.arange(0, BK) + o_v1 = tl.arange(0, 64) + m_kv1 = (o_k[:, None] < K) & (o_v1[None, :] < V) if USE_FINAL_STATE_GRADIENT: - p_dht1 = tl.make_block_ptr(dht, (K, V), (V, 1), (i_k * BK, 0), (BK, 64), (1, 0)) # [BK, BV] - b_dh1 += tl.load(p_dht1, boundary_check=(0, 1)) + p_dht1 = dht + o_k[:, None] * V + o_v1[None, :] # [BK, BV] + b_dh1 += tl.load(p_dht1, mask=(o_k[:, None] < K) & (o_v1[None, :] < V), other=0.0) if V > 64: - p_dht2 = tl.make_block_ptr(dht, (K, V), (V, 1), (i_k * BK, 64), (BK, 64), (1, 0)) - b_dh2 += tl.load(p_dht2, boundary_check=(0, 1)) + p_dht2 = dht + o_k[:, None] * V + (64 + o_v1)[None, :] + b_dh2 += tl.load(p_dht2, mask=(o_k[:, None] < K) & ((64 + o_v1)[None, :] < V), other=0.0) if V > 128: - p_dht3 = tl.make_block_ptr(dht, (K, V), (V, 1), (i_k * BK, 128), (BK, 64), (1, 0)) - b_dh3 += tl.load(p_dht3, boundary_check=(0, 1)) + p_dht3 = dht + o_k[:, None] * V + (128 + o_v1)[None, :] + b_dh3 += tl.load(p_dht3, mask=(o_k[:, None] < K) & ((128 + o_v1)[None, :] < V), other=0.0) if V > 192: - p_dht4 = tl.make_block_ptr(dht, (K, V), (V, 1), (i_k * BK, 192), (BK, 64), (1, 0)) - b_dh4 += tl.load(p_dht4, boundary_check=(0, 1)) + p_dht4 = dht + o_k[:, None] * V + (192 + o_v1)[None, :] + b_dh4 += tl.load(p_dht4, mask=(o_k[:, None] < K) & ((192 + o_v1)[None, :] < V), other=0.0) for i_t in range(NT - 1, -1, -1): - p_dh1 = tl.make_block_ptr(dh + i_t*stride_h, (K, V), (V, 1), (i_k * BK, 0), (BK, 64), (1, 0)) - tl.store(p_dh1, b_dh1.to(p_dh1.dtype.element_ty), boundary_check=(0, 1)) + o_t = i_t.to(tl.int64) * BT + tl.arange(0, BT) + m_t = o_t < T + m_tk = m_t[:, None] & (o_k[None, :] < K) + m_kv1 = (o_k[:, None] < K) & (o_v1[None, :] < V) + m_tv1 = m_t[:, None] & (o_v1[None, :] < V) + p_dh1 = dh + i_t.to(tl.int64)*stride_h + o_k[:, None] * V + o_v1[None, :] + tl.store(p_dh1, b_dh1.to(p_dh1.dtype.element_ty), mask=m_kv1) if V > 64: - p_dh2 = tl.make_block_ptr(dh + i_t*stride_h, (K, V), (V, 1), (i_k * BK, 64), (BK, 64), (1, 0)) - tl.store(p_dh2, b_dh2.to(p_dh2.dtype.element_ty), boundary_check=(0, 1)) + p_dh2 = dh + i_t.to(tl.int64)*stride_h + o_k[:, None] * V + (64 + o_v1)[None, :] + tl.store(p_dh2, b_dh2.to(p_dh2.dtype.element_ty), mask=(o_k[:, None] < K) & ((64 + o_v1)[None, :] < V)) if V > 128: - p_dh3 = tl.make_block_ptr(dh + i_t*stride_h, (K, V), (V, 1), (i_k * BK, 128), (BK, 64), (1, 0)) - tl.store(p_dh3, b_dh3.to(p_dh3.dtype.element_ty), boundary_check=(0, 1)) + p_dh3 = dh + i_t.to(tl.int64)*stride_h + o_k[:, None] * V + (128 + o_v1)[None, :] + tl.store(p_dh3, b_dh3.to(p_dh3.dtype.element_ty), mask=(o_k[:, None] < K) & ((128 + o_v1)[None, :] < V)) if V > 192: - p_dh4 = tl.make_block_ptr(dh + i_t*stride_h, (K, V), (V, 1), (i_k * BK, 192), (BK, 64), (1, 0)) - tl.store(p_dh4, b_dh4.to(p_dh4.dtype.element_ty), boundary_check=(0, 1)) + p_dh4 = dh + i_t.to(tl.int64)*stride_h + o_k[:, None] * V + (192 + o_v1)[None, :] + tl.store(p_dh4, b_dh4.to(p_dh4.dtype.element_ty), mask=(o_k[:, None] < K) & ((192 + o_v1)[None, :] < V)) last_idx = min((i_t + 1) * BT, T) - 1 - # Update dk_new, 按K切分 - p_dk = tl.make_block_ptr(dk, (T, K), (stride_k, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) # [BT, BK] - p_dk2 = tl.make_block_ptr(dk2, (T, K), (stride_k, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) # [BT, BK] + # update dk_new in K tiles + p_dk = dk + o_t[:, None] * stride_k + o_k[None, :] # [BT, BK] + p_dk2 = dk2 + o_t[:, None] * stride_k + o_k[None, :] # [BT, BK] if V > 0: - p_v = tl.make_block_ptr(vg, (T, V), (stride_v, 1), (i_t * BT, 0), (BT, 64), (1, 0)) - b_v = tl.load(p_v, boundary_check=(0, 1)) # [BT, BV] + p_v = vg + o_t[:, None] * stride_v + o_v1[None, :] + b_v = tl.load(p_v, mask=m_tv1, other=0.0) # [BT, BV] b_dk = tl.dot(b_v, tl.trans(b_dh1).to(b_v.dtype)) # [BT, BV] @ [BV, BK] -> [BT, BK] if V > 64: - p_v = tl.make_block_ptr(vg, (T, V), (stride_v, 1), (i_t * BT, 64), (BT, 64), (1, 0)) - b_v = tl.load(p_v, boundary_check=(0, 1)) - b_dk += tl.dot(b_v, tl.trans(b_dh2).to(b_v.dtype)) + p_v = vg + o_t[:, None] * stride_v + (64 + o_v1)[None, :] + b_v = tl.load(p_v, mask=m_t[:, None] & ((64 + o_v1)[None, :] < V), other=0.0) + b_dk = tl.dot(b_v, tl.trans(b_dh2).to(b_v.dtype), b_dk) if V > 128: - p_v = tl.make_block_ptr(vg, (T, V), (stride_v, 1), (i_t * BT, 128), (BT, 64), (1, 0)) - b_v = tl.load(p_v, boundary_check=(0, 1)) - b_dk += tl.dot(b_v, tl.trans(b_dh3).to(b_v.dtype)) + p_v = vg + o_t[:, None] * stride_v + (128 + o_v1)[None, :] + b_v = tl.load(p_v, mask=m_t[:, None] & ((128 + o_v1)[None, :] < V), other=0.0) + b_dk = tl.dot(b_v, tl.trans(b_dh3).to(b_v.dtype), b_dk) if V > 192: - p_v = tl.make_block_ptr(vg, (T, V), (stride_v, 1), (i_t * BT, 192), (BT, 64), (1, 0)) - b_v = tl.load(p_v, boundary_check=(0, 1)) - b_dk += tl.dot(b_v, tl.trans(b_dh4).to(b_v.dtype)) + p_v = vg + o_t[:, None] * stride_v + (192 + o_v1)[None, :] + b_v = tl.load(p_v, mask=m_t[:, None] & ((192 + o_v1)[None, :] < V), other=0.0) + b_dk = tl.dot(b_v, tl.trans(b_dh4).to(b_v.dtype), b_dk) - b_dk += tl.load(p_dk, boundary_check=(0, 1)) + b_dk += tl.load(p_dk, mask=m_tk, other=0.0) - tl.store(p_dk2, b_dk.to(p_dk.dtype.element_ty), boundary_check=(0, 1)) + tl.store(p_dk2, b_dk.to(p_dk.dtype.element_ty), mask=m_tk) - # Update dh, 按照K切分，收集所有V维度，q一次就好，wdo要收集所有 + # update dh in K tiles across all V tiles; load q once and w/do for each V tile - p_q = tl.make_block_ptr(q, (K, T), (1, stride_k), (i_k * BK, i_t * BT), (BK, BT), (0, 1)) # [BK, BT] - b_q = tl.load(p_q, boundary_check=(0, 1)) + p_q = q + o_k[:, None] + o_t[None, :] * stride_k # [BK, BT] + b_q = tl.load(p_q, mask=(o_k[:, None] < K) & m_t[None, :], other=0.0) if V > 0: - p_do = tl.make_block_ptr(do, (T, V), (stride_v, 1), (i_t * BT, 0), (BT, 64), (1, 0)) # [BT, BV] - b_do = tl.load(p_do, boundary_check=(0, 1)) - p_w = tl.make_block_ptr(w, (T, V), (stride_v, 1), (i_t * BT, 0), (BT, 64), (1, 0)) # [BT, BV] - b_w = tl.load(p_w, boundary_check=(0, 1)) - p_gv = tl.make_block_ptr(gv, (T, V), (stride_v, 1), (i_t * BT, 0), (BT, 64), (1, 0)) # [BT, BV] - b_gv = tl.load(p_gv, boundary_check=(0, 1)) + p_do = do + o_t[:, None] * stride_v + o_v1[None, :] # [BT, BV] + b_do = tl.load(p_do, mask=m_tv1, other=0.0) + p_w = w + o_t[:, None] * stride_v + o_v1[None, :] # [BT, BV] + b_w = tl.load(p_w, mask=m_tv1, other=0.0) + p_gv = gv + o_t[:, None] * stride_v + o_v1[None, :] # [BT, BV] + b_gv = tl.load(p_gv, mask=m_tv1, other=0.0) if USE_GV: - o_v1 = tl.arange(0, 64) - b_gv_last1 = tl.load(gv + last_idx * H*V + o_v1, mask=(o_v1 < V), other=0.) + b_gv_last1 = tl.load(gv + last_idx.to(tl.int64) * H*V + o_v1, mask=(o_v1 < V), other=0.) b_dh1 *= exp(b_gv_last1[None, :]) b_do *= exp(b_gv) b_dh1 += tl.dot(b_q.to(b_q.dtype), b_do.to(b_q.dtype)) * scale - \ tl.dot(tl.trans(b_dk).to(b_w.dtype), b_w) # [BK, BT] @ [BT, BV] - [BK, BT] @ [BT, BV] if V > 64: - p_do = tl.make_block_ptr(do, (T, V), (stride_v, 1), (i_t * BT, 64), (BT, 64), (1, 0)) - b_do = tl.load(p_do, boundary_check=(0, 1)) - p_w = tl.make_block_ptr(w, (T, V), (stride_v, 1), (i_t * BT, 64), (BT, 64), (1, 0)) # [BT, BV] - b_w = tl.load(p_w, boundary_check=(0, 1)) - p_gv = tl.make_block_ptr(gv, (T, V), (stride_v, 1), (i_t * BT, 64), (BT, 64), (1, 0)) # [BT, BV] - b_gv = tl.load(p_gv, boundary_check=(0, 1)) + p_do = do + o_t[:, None] * stride_v + (64 + o_v1)[None, :] + b_do = tl.load(p_do, mask=m_t[:, None] & ((64 + o_v1)[None, :] < V), other=0.0) + p_w = w + o_t[:, None] * stride_v + (64 + o_v1)[None, :] # [BT, BV] + b_w = tl.load(p_w, mask=m_t[:, None] & ((64 + o_v1)[None, :] < V), other=0.0) + p_gv = gv + o_t[:, None] * stride_v + (64 + o_v1)[None, :] # [BT, BV] + b_gv = tl.load(p_gv, mask=m_t[:, None] & ((64 + o_v1)[None, :] < V), other=0.0) if USE_GV: o_v2 = 64 + o_v1 - b_gv_last2 = tl.load(gv + last_idx * H*V + o_v2, mask=(o_v2 < V), other=0.) + b_gv_last2 = tl.load(gv + last_idx.to(tl.int64) * H*V + o_v2, mask=(o_v2 < V), other=0.) b_dh2 *= exp(b_gv_last2[None, :]) b_do *= exp(b_gv) b_dh2 += tl.dot(b_q.to(b_q.dtype), b_do.to(b_q.dtype)) * scale - tl.dot(tl.trans(b_dk).to(b_w.dtype), b_w) if V > 128: - p_do = tl.make_block_ptr(do, (T, V), (stride_v, 1), (i_t * BT, 128), (BT, 64), (1, 0)) - b_do = tl.load(p_do, boundary_check=(0, 1)) - p_w = tl.make_block_ptr(w, (T, V), (stride_v, 1), (i_t * BT, 128), (BT, 64), (1, 0)) # [BT, BV] - b_w = tl.load(p_w, boundary_check=(0, 1)) - p_gv = tl.make_block_ptr(gv, (T, V), (stride_v, 1), (i_t * BT, 128), (BT, 64), (1, 0)) # [BT, BV] - b_gv = tl.load(p_gv, boundary_check=(0, 1)) + p_do = do + o_t[:, None] * stride_v + (128 + o_v1)[None, :] + b_do = tl.load(p_do, mask=m_t[:, None] & ((128 + o_v1)[None, :] < V), other=0.0) + p_w = w + o_t[:, None] * stride_v + (128 + o_v1)[None, :] # [BT, BV] + b_w = tl.load(p_w, mask=m_t[:, None] & ((128 + o_v1)[None, :] < V), other=0.0) + p_gv = gv + o_t[:, None] * stride_v + (128 + o_v1)[None, :] # [BT, BV] + b_gv = tl.load(p_gv, mask=m_t[:, None] & ((128 + o_v1)[None, :] < V), other=0.0) if USE_GV: o_v3 = 128 + o_v1 - b_gv_last3 = tl.load(gv + last_idx * H*V + o_v3, mask=(o_v3 < V), other=0.) + b_gv_last3 = tl.load(gv + last_idx.to(tl.int64) * H*V + o_v3, mask=(o_v3 < V), other=0.) b_dh3 *= exp(b_gv_last3[None, :]) b_do *= exp(b_gv) b_dh3 += tl.dot(b_q.to(b_q.dtype), b_do.to(b_q.dtype)) * scale - tl.dot(tl.trans(b_dk).to(b_w.dtype), b_w) if V > 192: - p_do = tl.make_block_ptr(do, (T, V), (stride_v, 1), (i_t * BT, 192), (BT, 64), (1, 0)) - b_do = tl.load(p_do, boundary_check=(0, 1)) - p_w = tl.make_block_ptr(w, (T, V), (stride_v, 1), (i_t * BT, 192), (BT, 64), (1, 0)) # [BT, BV] - b_w = tl.load(p_w, boundary_check=(0, 1)) - p_gv = tl.make_block_ptr(gv, (T, V), (stride_v, 1), (i_t * BT, 192), (BT, 64), (1, 0)) # [BT, BV] - b_gv = tl.load(p_gv, boundary_check=(0, 1)) + p_do = do + o_t[:, None] * stride_v + (192 + o_v1)[None, :] + b_do = tl.load(p_do, mask=m_t[:, None] & ((192 + o_v1)[None, :] < V), other=0.0) + p_w = w + o_t[:, None] * stride_v + (192 + o_v1)[None, :] # [BT, BV] + b_w = tl.load(p_w, mask=m_t[:, None] & ((192 + o_v1)[None, :] < V), other=0.0) + p_gv = gv + o_t[:, None] * stride_v + (192 + o_v1)[None, :] # [BT, BV] + b_gv = tl.load(p_gv, mask=m_t[:, None] & ((192 + o_v1)[None, :] < V), other=0.0) if USE_GV: o_v4 = 192 + o_v1 - b_gv_last4 = tl.load(gv + last_idx * H*V + o_v4, mask=(o_v4 < V), other=0.) + b_gv_last4 = tl.load(gv + last_idx.to(tl.int64) * H*V + o_v4, mask=(o_v4 < V), other=0.) b_dh4 *= exp(b_gv_last4[None, :]) b_do *= exp(b_gv) b_dh4 += tl.dot(b_q.to(b_q.dtype), b_do.to(b_q.dtype)) * scale - tl.dot(tl.trans(b_dk).to(b_w.dtype), b_w) if USE_INITIAL_STATE: - p_dh0 = tl.make_block_ptr(dh0, (K, V), (V, 1), (i_k * BK, 0), (BK, 64), (1, 0)) - tl.store(p_dh0, b_dh1.to(p_dh0.dtype.element_ty), boundary_check=(0, 1)) + p_dh0 = dh0 + o_k[:, None] * V + o_v1[None, :] + tl.store(p_dh0, b_dh1.to(p_dh0.dtype.element_ty), mask=m_kv1) if V > 64: - p_dh1 = tl.make_block_ptr(dh0, (K, V), (V, 1), (i_k * BK, 64), (BK, 64), (1, 0)) - tl.store(p_dh1, b_dh2.to(p_dh1.dtype.element_ty), boundary_check=(0, 1)) + p_dh1 = dh0 + o_k[:, None] * V + (64 + o_v1)[None, :] + tl.store(p_dh1, b_dh2.to(p_dh1.dtype.element_ty), mask=(o_k[:, None] < K) & ((64 + o_v1)[None, :] < V)) if V > 128: - p_dh2 = tl.make_block_ptr(dh0, (K, V), (V, 1), (i_k * BK, 128), (BK, 64), (1, 0)) - tl.store(p_dh2, b_dh3.to(p_dh2.dtype.element_ty), boundary_check=(0, 1)) + p_dh2 = dh0 + o_k[:, None] * V + (128 + o_v1)[None, :] + tl.store(p_dh2, b_dh3.to(p_dh2.dtype.element_ty), mask=(o_k[:, None] < K) & ((128 + o_v1)[None, :] < V)) if V > 192: - p_dh3 = tl.make_block_ptr(dh0, (K, V), (V, 1), (i_k * BK, 192), (BK, 64), (1, 0)) - tl.store(p_dh3, b_dh4.to(p_dh3.dtype.element_ty), boundary_check=(0, 1)) + p_dh3 = dh0 + o_k[:, None] * V + (192 + o_v1)[None, :] + tl.store(p_dh3, b_dh4.to(p_dh3.dtype.element_ty), mask=(o_k[:, None] < K) & ((192 + o_v1)[None, :] < V)) def chunk_oja_bwd_dhu( @@ -487,7 +499,7 @@ def chunk_oja_bwd_dhu( dh0 = torch.empty_like(h0, dtype=torch.float32) if h0 is not None else None dk2 = torch.empty_like(dk) - def grid(meta): return (triton.cdiv(K, meta['BK']), N*H) + def grid(meta): return (triton.cdiv(K, meta['BK']) * N * H,) chunk_oja_bwd_kernel_dhu_blockdim64[grid]( q=q, vg=vg, @@ -553,13 +565,13 @@ def chunk_gsa_bwd_k_kernel_dqkvg( NG: tl.constexpr, IS_VARLEN: tl.constexpr, ): - i_k, i_t, i_bh = tl.program_id(0), tl.program_id(1), tl.program_id(2) + i_k, i_t, i_bh = tl.program_id(0).to(tl.int64), tl.program_id(1).to(tl.int64), tl.program_id(2).to(tl.int64) i_b, i_hq = i_bh // HQ, i_bh % HQ i_h = i_hq // NG 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.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) + i_n, i_t = tl.load(chunk_indices + i_t * 2).to(tl.int32), tl.load(chunk_indices + i_t * 2 + 1).to(tl.int64) + bos, eos = tl.load(cu_seqlens + i_n).to(tl.int64), tl.load(cu_seqlens + i_n + 1).to(tl.int64) all = T T = eos - bos NT = tl.cdiv(T, BT) @@ -567,82 +579,91 @@ def chunk_gsa_bwd_k_kernel_dqkvg( NT = tl.cdiv(T, BT) i_tg = i_b * NT + i_t bos, eos = i_b * T, i_b * T + T - all = B * T + all = B * T.to(tl.int64) o_i = tl.arange(0, BT) o_t = min(i_t * BT + BT, T) m_s = o_i[:, None] >= o_i[None, :] - p_q = tl.make_block_ptr(q + (bos*HQ+i_hq) * K, (T, K), (HQ*K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) - p_k = tl.make_block_ptr(k + (bos*H+i_h) * K, (T, K), (H*K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) - p_A = tl.make_block_ptr(A + ((i_k*all+bos)*HQ+i_hq)*BT, (T, BT), (HQ*BT, 1), (i_t * BT, 0), (BT, BT), (1, 0)) + o_tt = i_t * BT + tl.arange(0, BT) + o_k = i_k * BK + tl.arange(0, BK) + m_t = o_tt < T + m_tk = m_t[:, None] & (o_k[None, :] < K) + o_A = tl.arange(0, BT) + m_AT = m_t[:, None] & (o_A[None, :] < BT) + p_q = q + (bos*HQ+i_hq) * K + o_tt[:, None] * (HQ*K) + o_k[None, :] + p_k = k + (bos*H+i_h) * K + o_tt[:, None] * (H*K) + o_k[None, :] + p_A = A + ((i_k*all+bos)*HQ+i_hq)*BT + o_tt[:, None] * (HQ*BT) + o_A[None, :] # [BT, BK] - b_q = tl.load(p_q, boundary_check=(0, 1)) - b_k = tl.load(p_k, boundary_check=(0, 1)) + b_q = tl.load(p_q, mask=m_tk, other=0.0) + b_k = tl.load(p_k, mask=m_tk, other=0.0) # [BT, BT] b_A = tl.dot((b_q * scale).to(b_q.dtype), tl.trans(b_k)) b_A = tl.where(m_s, b_A, 0.) - tl.store(p_A, b_A.to(p_A.dtype.element_ty), boundary_check=(0, 1)) + tl.store(p_A, b_A.to(p_A.dtype.element_ty), mask=m_AT) b_dq = tl.zeros([BT, BK], dtype=tl.float32) b_dk = tl.zeros([BT, BK], dtype=tl.float32) for i_v in range(tl.cdiv(V, BV)): o_v = i_v * BV + tl.arange(0, BV) - p_v = tl.make_block_ptr(v + (bos*H+i_h)*V, (T, V), (H*V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) - p_g = tl.make_block_ptr(g + (bos*H+i_h)*V, (T, V), (H*V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) + m_tv = m_t[:, None] & (o_v[None, :] < V) + m_h = (o_v[:, None] < V) & (o_k[None, :] < K) + m_kv = (o_k[:, None] < K) & (o_v[None, :] < V) + p_v = v + (bos*H+i_h)*V + o_tt[:, None] * (H*V) + o_v[None, :] + p_g = g + (bos*H+i_h)*V + o_tt[:, None] * (H*V) + o_v[None, :] p_gn = g + (bos + o_t - 1) * H*V + i_h * V + o_v - p_do = tl.make_block_ptr(do + (bos*HQ+i_hq)*V, (T, V), (HQ*V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) - p_dv = tl.make_block_ptr(dv + ((i_k*all+bos)*HQ+i_hq)*V, (T, V), (HQ*V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) - p_dg = tl.make_block_ptr(dg + (bos*HQ+i_hq)*V, (T, V), (HQ*V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) - p_dgv = tl.make_block_ptr(dgv+((i_k*all+bos)*HQ+i_hq)*V, (T, V), (HQ*V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) - p_h = tl.make_block_ptr(h + (i_tg * H + i_h) * K*V, (V, K), (1, V), (i_v * BV, i_k * BK), (BV, BK), (0, 1)) - p_dh = tl.make_block_ptr(dh + (i_tg * HQ + i_hq) * K*V, (K, V), (V, 1), (i_k * BK, i_v * BV), (BK, BV), (1, 0)) + p_do = do + (bos*HQ+i_hq)*V + o_tt[:, None] * (HQ*V) + o_v[None, :] + p_dv = dv + ((i_k*all+bos)*HQ+i_hq)*V + o_tt[:, None] * (HQ*V) + o_v[None, :] + p_dg = dg + (bos*HQ+i_hq)*V + o_tt[:, None] * (HQ*V) + o_v[None, :] + p_dgv = dgv+((i_k*all+bos)*HQ+i_hq)*V + o_tt[:, None] * (HQ*V) + o_v[None, :] + p_h = h + (i_tg * H + i_h) * K*V + o_v[:, None] + o_k[None, :] * V + p_dh = dh + (i_tg * HQ + i_hq) * K*V + o_k[:, None] * V + o_v[None, :] m_v = o_v < V # [BV,] b_gn = tl.load(p_gn, mask=m_v, other=0) # [BT, BV] - b_v = tl.load(p_v, boundary_check=(0, 1)) - b_g = tl.load(p_g, boundary_check=(0, 1)) + b_v = tl.load(p_v, mask=m_tv, other=0.0) + b_g = tl.load(p_g, mask=m_tv, other=0.0) b_gv = exp(b_gn[None, :] - b_g) # [BV, BK] - b_h = tl.load(p_h, boundary_check=(0, 1)) + b_h = tl.load(p_h, mask=m_h, other=0.0) # [BT, BV] - b_do = tl.load(p_do, boundary_check=(0, 1)) + b_do = tl.load(p_do, mask=m_tv, other=0.0) b_do = (b_do * exp(b_g) * scale).to(b_do.dtype) # [BK, BV] - b_dh = tl.load(p_dh, boundary_check=(0, 1)) + b_dh = tl.load(p_dh, mask=m_kv, other=0.0) # [BV] b_dg = tl.sum(tl.trans(b_h) * b_dh, 0) * exp(b_gn) b_dh = b_dh.to(b_k.dtype) # [BT, BK] - b_dq += tl.dot(b_do, b_h.to(b_k.dtype)) - b_dk += tl.dot((b_v * b_gv).to(b_v.dtype), tl.trans(b_dh)) + b_dq = tl.dot(b_do, b_h.to(b_k.dtype), b_dq) + b_dk = tl.dot((b_v * b_gv).to(b_v.dtype), tl.trans(b_dh), b_dk) # [BT, BV] b_dv = tl.dot(b_k, b_dh) * b_gv # [BV] b_dg += tl.sum(b_dv * b_v, 0) if i_k == 0: - b_dgv = tl.load(p_dg, boundary_check=(0, 1)) + b_dg[None, :] + b_dgv = tl.load(p_dg, mask=m_tv, other=0.0) + b_dg[None, :] else: b_dgv = tl.zeros([BT, BV], dtype=tl.float32) + b_dg[None, :] - tl.store(p_dgv, b_dgv.to(p_dgv.dtype.element_ty), boundary_check=(0, 1)) - tl.store(p_dv, b_dv.to(p_dv.dtype.element_ty), boundary_check=(0, 1)) - p_dA = tl.make_block_ptr(dA + (bos*HQ + i_hq) * BT, (T, BT), (HQ*BT, 1), (i_t * BT, 0), (BT, BT), (1, 0)) - p_dq = tl.make_block_ptr(dq + (bos*HQ + i_hq) * K, (T, K), (HQ*K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) - p_dk = tl.make_block_ptr(dk + (bos*HQ + i_hq) * K, (T, K), (HQ*K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) + tl.store(p_dgv, b_dgv.to(p_dgv.dtype.element_ty), mask=m_tv) + tl.store(p_dv, b_dv.to(p_dv.dtype.element_ty), mask=m_tv) + p_dA = dA + (bos*HQ + i_hq) * BT + o_tt[:, None] * (HQ*BT) + o_A[None, :] + p_dq = dq + (bos*HQ + i_hq) * K + o_tt[:, None] * (HQ*K) + o_k[None, :] + p_dk = dk + (bos*HQ + i_hq) * K + o_tt[:, None] * (HQ*K) + o_k[None, :] # [BT, BT] - b_dA = tl.load(p_dA, boundary_check=(0, 1)) + b_dA = tl.load(p_dA, mask=m_AT, other=0.0) # [BT, BK] - b_dq += tl.dot(b_dA, b_k) - b_dk += tl.dot(tl.trans(b_dA).to(b_k.dtype), b_q) + b_dq = tl.dot(b_dA, b_k, b_dq) + b_dk = tl.dot(tl.trans(b_dA).to(b_k.dtype), b_q, b_dk) - tl.store(p_dq, b_dq.to(p_dq.dtype.element_ty), boundary_check=(0, 1)) - tl.store(p_dk, b_dk.to(p_dk.dtype.element_ty), boundary_check=(0, 1)) + tl.store(p_dq, b_dq.to(p_dq.dtype.element_ty), mask=m_tk) + tl.store(p_dk, b_dk.to(p_dk.dtype.element_ty), mask=m_tk) @triton.heuristics({ @@ -683,13 +704,14 @@ def chunk_oja_bwd_kernel_dvwg_h( HAVE_GK: tl.constexpr, IS_VARLEN: tl.constexpr, ): - i_v, i_t, i_bh = tl.program_id(0), tl.program_id(1), tl.program_id(2) + i_v, i_t = unflatten_program_id(tl.cdiv(V, BV)) + i_bh = tl.program_id(1).to(tl.int64) i_b, i_h = i_bh // H, i_bh % H 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.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) + i_n, i_t = tl.load(chunk_indices + i_t * 2).to(tl.int32), tl.load(chunk_indices + i_t * 2 + 1).to(tl.int64) + bos, eos = tl.load(cu_seqlens + i_n).to(tl.int64), tl.load(cu_seqlens + i_n + 1).to(tl.int64) T = eos - bos NT = tl.cdiv(T, BT) else: @@ -712,51 +734,56 @@ def chunk_oja_bwd_kernel_dvwg_h( b_dw = tl.zeros([BT, BV], dtype=tl.float32) b_dgv_last = tl.zeros([BV,], dtype=tl.float32) + o_t = i_t * BT + tl.arange(0, BT) + o_v = i_v * BV + tl.arange(0, BV) + m_v = o_v < V + m_tv = (o_t[:, None] < T) & (o_v[None, :] < V) if USE_GV: - o_v = i_v * BV + tl.arange(0, BV) - m_v = o_v < V p_gn = gv + (min(T, i_t * BT + BT) - 1) * H*V + o_v - p_gv = tl.make_block_ptr(gv, (T, V), (H*V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) + p_gv = gv + o_t[:, None] * (H*V) + o_v[None, :] b_gn = tl.load(p_gn, mask=m_v, other=0) - b_gv = tl.load(p_gv, boundary_check=(0, 1)) + b_gv = tl.load(p_gv, mask=m_tv, other=0.0) for i_k in range(tl.cdiv(K, BK)): - p_k = tl.make_block_ptr(k, (T, K), (H*K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) - p_dk = tl.make_block_ptr(dk, (T, K), (H*K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) - p_h = tl.make_block_ptr(h, (K, V), (V, 1), (i_k * BK, i_v * BV), (BK, BV), (1, 0)) - p_dh = tl.make_block_ptr(dh, (K, V), (V, 1), (i_k * BK, i_v * BV), (BK, BV), (1, 0)) - b_k = tl.load(p_k, boundary_check=(0, 1)) # BT BK - b_dk = tl.load(p_dk, boundary_check=(0, 1)) # BT BK - b_h = tl.load(p_h, boundary_check=(0, 1)) # BK BV - b_dh = tl.load(p_dh, boundary_check=(0, 1)) # BK BV - - b_dvg += tl.dot(b_k, b_dh.to(b_k.dtype)) # BT BK @ BK BV -> BT BV - b_dw += tl.dot(b_dk.to(b_k.dtype), b_h.to(b_k.dtype)) # BT BK @ BK BV -> BT BV + o_k = i_k * BK + tl.arange(0, BK) + m_tk = (o_t[:, None] < T) & (o_k[None, :] < K) + m_kv = (o_k[:, None] < K) & (o_v[None, :] < V) + p_k = k + o_t[:, None] * (H*K) + o_k[None, :] + p_dk = dk + o_t[:, None] * (H*K) + o_k[None, :] + p_h = h + o_k[:, None] * V + o_v[None, :] + p_dh = dh + o_k[:, None] * V + o_v[None, :] + b_k = tl.load(p_k, mask=m_tk, other=0.0) # BT BK + b_dk = tl.load(p_dk, mask=m_tk, other=0.0) # BT BK + b_h = tl.load(p_h, mask=m_kv, other=0.0) # BK BV + b_dh = tl.load(p_dh, mask=m_kv, other=0.0) # BK BV + + b_dvg = tl.dot(b_k, b_dh.to(b_k.dtype), b_dvg) # BT BK @ BK BV -> BT BV + b_dw = tl.dot(b_dk.to(b_k.dtype), b_h.to(b_k.dtype), b_dw) # BT BK @ BK BV -> BT BV b_dgv_last += tl.sum((b_h * b_dh) * exp(b_gn), axis=0) if USE_GV: b_dv = b_dvg * exp(b_gn[None, :] - b_gv) - p_v = tl.make_block_ptr(v, (T, V), (H*V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) - p_dv = tl.make_block_ptr(dv, (T, V), (H*V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) - p_dw = tl.make_block_ptr(dw, (T, V), (H*V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) - p_dgv_last = tl.make_block_ptr(dgv_last, (T, V), (H*V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) - b_v = tl.load(p_v, boundary_check=(0, 1)) + p_v = v + o_t[:, None] * (H*V) + o_v[None, :] + p_dv = dv + o_t[:, None] * (H*V) + o_v[None, :] + p_dw = dw + o_t[:, None] * (H*V) + o_v[None, :] + p_dgv_last = dgv_last + o_t[:, None] * (H*V) + o_v[None, :] + b_v = tl.load(p_v, mask=m_tv, other=0.0) b_dgv_last += tl.sum(b_dv * b_v, axis=0) - # 留给GSA2的接口 + # reserved for GSA2 if HAVE_GK: dgk += (bos * H + i_h) * V - p_dgk = tl.make_block_ptr(dgk, (T, V), (H*V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) - b_dgk = tl.load(p_dgk, boundary_check=(0, 1)) + p_dgk = dgk + o_t[:, None] * (H*V) + o_v[None, :] + b_dgk = tl.load(p_dgk, mask=m_tv, other=0.0) b_dgv_last = b_dgk + b_dgv_last[None, :] else: b_dgv_last = tl.zeros([BT, BV], dtype=tl.float32) + b_dgv_last[None, :] - tl.store(p_dv, b_dv.to(p_dv.dtype.element_ty), boundary_check=(0, 1)) - tl.store(p_dw, -b_dw.to(p_dw.dtype.element_ty), boundary_check=(0, 1)) - tl.store(p_dgv_last, b_dgv_last.to(p_dgv_last.dtype.element_ty), boundary_check=(0, 1)) + tl.store(p_dv, b_dv.to(p_dv.dtype.element_ty), mask=m_tv) + tl.store(p_dw, -b_dw.to(p_dw.dtype.element_ty), mask=m_tv) + tl.store(p_dgv_last, b_dgv_last.to(p_dgv_last.dtype.element_ty), mask=m_tv) def chunk_oja_bwd_dvwg_h( @@ -786,7 +813,7 @@ def chunk_oja_bwd_dvwg_h( dw = torch.empty_like(v) dgv_last = torch.empty_like(gv) - grid = (NV, NT, B * H) + grid = (NV * NT, B * H) chunk_oja_bwd_kernel_dvwg_h[grid]( k=k, v=v, diff --git a/build/torch-cuda/ops/gated_oja_rule/chunk_kkt.py b/build/torch-cuda/ops/gated_oja_rule/chunk_kkt.py index f152bfdc15109b57a85117324372352bc4934742..b162db2ddd05436c4efa4693c1471bb4bf859ee3 100644 --- a/build/torch-cuda/ops/gated_oja_rule/chunk_kkt.py +++ b/build/torch-cuda/ops/gated_oja_rule/chunk_kkt.py @@ -42,37 +42,41 @@ def chunk_scaled_dot_kkt_fwd_kernel( IS_VARLEN: tl.constexpr, USE_G: tl.constexpr, ): - i_t, i_bh = tl.program_id(0), tl.program_id(1) + i_t, i_bh = tl.program_id(0).to(tl.int64), tl.program_id(1).to(tl.int64) i_b, i_h = i_bh // H, i_bh % H 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.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) + i_n, i_t = tl.load(chunk_indices + i_t * 2).to(tl.int32), tl.load(chunk_indices + i_t * 2 + 1).to(tl.int64) + bos, eos = tl.load(cu_seqlens + i_n).to(tl.int64), tl.load(cu_seqlens + i_n + 1).to(tl.int64) T = eos - bos else: bos, eos = i_b * T, i_b * T + T o_t = i_t * BT + tl.arange(0, BT) m_t = o_t < T - p_b = tl.make_block_ptr(beta + bos*H + i_h, (T,), (H,), (i_t * BT,), (BT,), (0,)) - b_b = tl.load(p_b, boundary_check=(0,)) + p_b = beta + bos*H + i_h + o_t * H + b_b = tl.load(p_b, mask=m_t, other=0.0) b_A = tl.zeros([BT, BT], dtype=tl.float32) for i_k in range(tl.cdiv(K, BK)): - p_k = tl.make_block_ptr(k + (bos*H + i_h) * K, (T, K), (H*K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) - b_k = tl.load(p_k, boundary_check=(0, 1)) - b_A += tl.dot(b_k, tl.trans(b_k)) + o_k = i_k * BK + tl.arange(0, BK) + m_tk = m_t[:, None] & (o_k[None, :] < K) + p_k = k + (bos*H + i_h) * K + o_t[:, None] * (H*K) + o_k[None, :] + b_k = tl.load(p_k, mask=m_tk, other=0.0) + b_A = tl.dot(b_k, tl.trans(b_k), b_A) if USE_G: - p_g = tl.make_block_ptr(g + bos*H + i_h, (T,), (H,), (i_t * BT,), (BT,), (0,)) - b_g = tl.load(p_g, boundary_check=(0,)) + p_g = g + bos*H + i_h + o_t * H + b_g = tl.load(p_g, mask=m_t, other=0.0) b_g_diff = b_g[:, None] - b_g[None, :] b_A *= exp(b_g_diff) b_A *= b_b[:, None] m_A = (o_t[:, None] > o_t[None, :]) & (m_t[:, None] & m_t) b_A = tl.where(m_A, b_A, 0) - p_A = tl.make_block_ptr(A + (bos*H + i_h) * BT, (T, BT), (BT*H, 1), (i_t * BT, 0), (BT, BT), (1, 0)) - tl.store(p_A, b_A.to(p_A.dtype.element_ty), boundary_check=(0, 1)) + o_A = tl.arange(0, BT) + m_AT = m_t[:, None] & (o_A[None, :] < BT) + p_A = A + (bos*H + i_h) * BT + o_t[:, None] * (BT*H) + o_A[None, :] + tl.store(p_A, b_A.to(p_A.dtype.element_ty), mask=m_AT) @triton.heuristics({ @@ -104,12 +108,12 @@ def chunk_scaled_dot_kkt_fwd_kernel_intra_sub_inter( NC: tl.constexpr, IS_VARLEN: tl.constexpr, ): - i_t, i_c, i_bh = tl.program_id(0), tl.program_id(1), tl.program_id(2) + i_t, i_c, i_bh = tl.program_id(0).to(tl.int64), tl.program_id(1).to(tl.int64), tl.program_id(2).to(tl.int64) i_b, i_h = i_bh // H, i_bh % H i_i, i_j = i_c // NC, i_c % NC 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.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) + i_n, i_t = tl.load(chunk_indices + i_t * 2).to(tl.int32), tl.load(chunk_indices + i_t * 2 + 1).to(tl.int64) + bos, eos = tl.load(cu_seqlens + i_n).to(tl.int64), tl.load(cu_seqlens + i_n + 1).to(tl.int64) T = eos - bos else: bos, eos = i_b * T, i_b * T + T @@ -123,32 +127,38 @@ def chunk_scaled_dot_kkt_fwd_kernel_intra_sub_inter( g += (bos * H + i_h) * K A += (bos * H + i_h) * BT - p_b = tl.make_block_ptr(beta + bos * H + i_h, (T,), (H,), (i_t * BT + i_i * BC,), (BC,), (0,)) - b_b = tl.load(p_b, boundary_check=(0,)) + o_r = i_t * BT + i_i * BC + tl.arange(0, BC) + m_r = o_r < T + p_b = beta + bos * H + i_h + o_r * H + b_b = tl.load(p_b, mask=m_r, other=0.0) b_A = tl.zeros([BC, BC], dtype=tl.float32) for i_k in range(tl.cdiv(K, BK)): - p_k = tl.make_block_ptr(k, (T, K), (H*K, 1), (i_t * BT + i_i * BC, i_k * BK), (BC, BK), (1, 0)) - p_g = tl.make_block_ptr(g, (T, K), (H*K, 1), (i_t * BT + i_i * BC, i_k * BK), (BC, BK), (1, 0)) - b_kt = tl.make_block_ptr(k, (K, T), (1, H*K), (i_k * BK, i_t * BT + i_j * BC), (BK, BC), (0, 1)) - p_gk = tl.make_block_ptr(g, (K, T), (1, H*K), (i_k * BK, i_t * BT + i_j * BC), (BK, BC), (0, 1)) - + o_kj = i_t * BT + i_j * BC + tl.arange(0, BC) o_k = i_k * BK + tl.arange(0, BK) m_k = o_k < K + m_rk = m_r[:, None] & m_k[None, :] + m_kj = m_k[:, None] & (o_kj[None, :] < T) + p_k = k + o_r[:, None] * (H*K) + o_k[None, :] + p_g = g + o_r[:, None] * (H*K) + o_k[None, :] + b_kt = k + o_k[:, None] + o_kj[None, :] * (H*K) + p_gk = g + o_k[:, None] + o_kj[None, :] * (H*K) + # [BK,] b_gn = tl.load(g + (i_t * BT + i_i * BC) * H*K + o_k, mask=m_k, other=0) # [BC, BK] - b_g = tl.load(p_g, boundary_check=(0, 1)) - b_k = tl.load(p_k, boundary_check=(0, 1)) * exp(b_g - b_gn[None, :]) + b_g = tl.load(p_g, mask=m_rk, other=0.0) + b_k = tl.load(p_k, mask=m_rk, other=0.0) * exp(b_g - b_gn[None, :]) # [BK, BC] - b_gk = tl.load(p_gk, boundary_check=(0, 1)) - b_kt = tl.load(b_kt, boundary_check=(0, 1)) * exp(b_gn[:, None] - b_gk) + b_gk = tl.load(p_gk, mask=m_kj, other=0.0) + b_kt = tl.load(b_kt, mask=m_kj, other=0.0) * exp(b_gn[:, None] - b_gk) # [BC, BC] - b_A += tl.dot(b_k, b_kt) + b_A = tl.dot(b_k, b_kt, b_A) b_A *= b_b[:, None] - p_A = tl.make_block_ptr(A, (T, BT), (H*BT, 1), (i_t * BT + i_i * BC, i_j * BC), (BC, BC), (1, 0)) - tl.store(p_A, b_A.to(A.dtype.element_ty), boundary_check=(0, 1)) + o_Aj = i_j * BC + tl.arange(0, BC) + p_A = A + o_r[:, None] * (H*BT) + o_Aj[None, :] + tl.store(p_A, b_A.to(A.dtype.element_ty), mask=m_r[:, None] & (o_Aj[None, :] < BT)) @triton.heuristics({ @@ -179,11 +189,11 @@ def chunk_scaled_dot_kkt_fwd_kernel_intra_sub_intra( BK: tl.constexpr, IS_VARLEN: tl.constexpr, ): - i_t, i_i, i_bh = tl.program_id(0), tl.program_id(1), tl.program_id(2) + i_t, i_i, i_bh = tl.program_id(0).to(tl.int64), tl.program_id(1).to(tl.int64), tl.program_id(2).to(tl.int64) i_b, i_h = i_bh // H, i_bh % H 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.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) + i_n, i_t = tl.load(chunk_indices + i_t * 2).to(tl.int32), tl.load(chunk_indices + i_t * 2 + 1).to(tl.int64) + bos, eos = tl.load(cu_seqlens + i_n).to(tl.int64), tl.load(cu_seqlens + i_n + 1).to(tl.int64) T = eos - bos else: bos, eos = i_b * T, i_b * T + T @@ -197,12 +207,14 @@ def chunk_scaled_dot_kkt_fwd_kernel_intra_sub_intra( m_A = (i_t * BT + i_i * BC + o_i) < T o_A = (bos + i_t * BT + i_i * BC + o_i) * H*BT + i_h * BT + i_i * BC - p_k = tl.make_block_ptr(k + (bos * H + i_h) * K, (T, K), (H*K, 1), (i_t * BT + i_i * BC, 0), (BC, BK), (1, 0)) - p_g = tl.make_block_ptr(g + (bos * H + i_h) * K, (T, K), (H*K, 1), (i_t * BT + i_i * BC, 0), (BC, BK), (1, 0)) + o_r = i_t * BT + i_i * BC + o_i + m_rk = m_A[:, None] & m_k[None, :] + p_k = k + (bos * H + i_h) * K + o_r[:, None] * (H*K) + o_k[None, :] + p_g = g + (bos * H + i_h) * K + o_r[:, None] * (H*K) + o_k[None, :] p_b = beta + (bos + i_t * BT + i_i * BC + o_i) * H + i_h - b_k = tl.load(p_k, boundary_check=(0, 1)) * tl.load(p_b, mask=m_A, other=0)[:, None] - b_g = tl.load(p_g, boundary_check=(0, 1)) + b_k = tl.load(p_k, mask=m_rk, other=0.0) * tl.load(p_b, mask=m_A, other=0)[:, None] + b_g = tl.load(p_g, mask=m_rk, other=0.0) p_kt = k + (bos + i_t * BT + i_i * BC) * H*K + i_h * K + o_k p_gk = g + (bos + i_t * BT + i_i * BC) * H*K + i_h * K + o_k @@ -248,14 +260,15 @@ def chunk_scaled_dot_kkt_bwd_kernel_gk( NC: tl.constexpr, IS_VARLEN: tl.constexpr, ): - i_k, i_c, i_bh = tl.program_id(0), tl.program_id(1), tl.program_id(2) + NK = tl.cdiv(K, BK) + i_k, i_c, i_bh = tl.program_id(0).to(tl.int64) % NK, tl.program_id(0).to(tl.int64) // NK, tl.program_id(1).to(tl.int64) i_b, i_h = i_bh // H, i_bh % H i_t, i_i = i_c // NC, i_c % NC - all = B * T + all = B.to(tl.int64) * T 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.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) + i_n, i_t = tl.load(chunk_indices + i_t * 2).to(tl.int32), tl.load(chunk_indices + i_t * 2 + 1).to(tl.int64) + bos, eos = tl.load(cu_seqlens + i_n).to(tl.int64), tl.load(cu_seqlens + i_n + 1).to(tl.int64) else: bos, eos = i_b * T, i_b * T + T T = eos - bos @@ -274,27 +287,33 @@ def chunk_scaled_dot_kkt_bwd_kernel_gk( dg += (bos * H + i_h) * K db += (i_k * all + bos) * H + i_h - p_g = tl.make_block_ptr(g, (T, K), (H*K, 1), (i_t * BT + i_i * BC, i_k * BK), (BC, BK), (1, 0)) - p_b = tl.make_block_ptr(beta, (T,), (H,), (i_t * BT + i_i * BC,), (BC,), (0,)) + o_r = i_t * BT + i_i * BC + tl.arange(0, BC) + m_r = o_r < T + m_rk = m_r[:, None] & m_k[None, :] + p_g = g + o_r[:, None] * (H*K) + o_k[None, :] + p_b = beta + o_r * H # [BC, BK] - b_g = tl.load(p_g, boundary_check=(0, 1)) + b_g = tl.load(p_g, mask=m_rk, other=0.0) b_dk = tl.zeros([BC, BK], dtype=tl.float32) # [BC] - b_b = tl.load(p_b, boundary_check=(0,)) + b_b = tl.load(p_b, mask=m_r, other=0.0) if i_i > 0: p_gn = g + (i_t * BT + i_i * BC) * H*K + o_k # [BK,] b_gn = tl.load(p_gn, mask=m_k, other=0) for i_j in range(0, i_i): - p_k = tl.make_block_ptr(k, (T, K), (H*K, 1), (i_t * BT + i_j * BC, i_k * BK), (BC, BK), (1, 0)) - p_gk = tl.make_block_ptr(g, (T, K), (H*K, 1), (i_t * BT + i_j * BC, i_k * BK), (BC, BK), (1, 0)) - p_dA = tl.make_block_ptr(dA, (T, BT), (H*BT, 1), (i_t * BT + i_i * BC, i_j * BC), (BC, BC), (1, 0)) + o_j = i_t * BT + i_j * BC + tl.arange(0, BC) + m_jk = (o_j[:, None] < T) & m_k[None, :] + o_dAj = i_j * BC + tl.arange(0, BC) + p_k = k + o_j[:, None] * (H*K) + o_k[None, :] + p_gk = g + o_j[:, None] * (H*K) + o_k[None, :] + p_dA = dA + o_r[:, None] * (H*BT) + o_dAj[None, :] # [BC, BK] - b_k = tl.load(p_k, boundary_check=(0, 1)) - b_gk = tl.load(p_gk, boundary_check=(0, 1)) + b_k = tl.load(p_k, mask=m_jk, other=0.0) + b_gk = tl.load(p_gk, mask=m_jk, other=0.0) b_kg = b_k * exp(b_gn[None, :] - b_gk) # [BC, BC] - b_dA = tl.load(p_dA, boundary_check=(0, 1)) + b_dA = tl.load(p_dA, mask=m_r[:, None] & (o_dAj[None, :] < BT), other=0.0) # [BC, BK] b_dkb = tl.dot(b_dA, b_kg) * exp(b_g - b_gn[None, :]) b_dk += b_dkb @@ -305,8 +324,8 @@ def chunk_scaled_dot_kkt_bwd_kernel_gk( p_kj = k + (i_t * BT + i_i * BC) * H*K + o_k p_gkj = g + (i_t * BT + i_i * BC) * H*K + o_k - p_k = tl.make_block_ptr(k, (T, K), (H*K, 1), (i_t * BT + i_i * BC, i_k * BK), (BC, BK), (1, 0)) - b_k = tl.load(p_k, boundary_check=(0, 1)) + p_k = k + o_r[:, None] * (H*K) + o_k[None, :] + b_k = tl.load(p_k, mask=m_rk, other=0.0) for j in range(0, min(BC, T - i_t * BT - i_i * BC)): # [BC] b_dA = tl.load(dA + o_dA + j, mask=m_dA, other=0) @@ -323,8 +342,8 @@ def chunk_scaled_dot_kkt_bwd_kernel_gk( p_gkj += H*K b_db = tl.sum(b_dk * b_k, 1) b_dk *= b_b[:, None] - p_db = tl.make_block_ptr(db, (T,), (H,), (i_t * BT + i_i * BC,), (BC,), (0,)) - tl.store(p_db, b_db.to(p_db.dtype.element_ty), boundary_check=(0,)) + p_db = db + o_r * H + tl.store(p_db, b_db.to(p_db.dtype.element_ty), mask=m_r) tl.debug_barrier() # [BC, BK] @@ -336,24 +355,26 @@ def chunk_scaled_dot_kkt_bwd_kernel_gk( # [BK,] b_gn = tl.load(p_gn, mask=m_k, other=0) for i_j in range(i_i + 1, NC): - p_k = tl.make_block_ptr(k, (T, K), (H*K, 1), (i_t * BT + i_j * BC, i_k * BK), (BC, BK), (1, 0)) - p_gk = tl.make_block_ptr(g, (T, K), (H*K, 1), (i_t * BT + i_j * BC, i_k*BK), (BC, BK), (1, 0)) - p_dA = tl.make_block_ptr(dA, (BT, T), (1, H*BT), (i_i * BC, i_t * BT + i_j * BC), (BC, BC), (0, 1)) - p_b = tl.make_block_ptr(beta, (T,), (H,), (i_t * BT + i_j * BC,), (BC,), (0,)) - - o_j = i_t * BT + i_j * BC + o_i + o_j = i_t * BT + i_j * BC + tl.arange(0, BC) m_j = o_j < T + m_jk = m_j[:, None] & m_k[None, :] + o_dAi = i_i * BC + tl.arange(0, BC) + p_k = k + o_j[:, None] * (H*K) + o_k[None, :] + p_gk = g + o_j[:, None] * (H*K) + o_k[None, :] + p_dA = dA + o_dAi[:, None] + o_j[None, :] * (H*BT) + p_b = beta + o_j * H + # [BC] - b_b = tl.load(p_b, boundary_check=(0,)) + b_b = tl.load(p_b, mask=m_j, other=0.0) # [BC, BK] - b_kb = tl.load(p_k, boundary_check=(0, 1)).to(tl.float32) * b_b[:, None] - b_gk = tl.load(p_gk, boundary_check=(0, 1)) + b_kb = tl.load(p_k, mask=m_jk, other=0.0).to(tl.float32) * b_b[:, None] + b_gk = tl.load(p_gk, mask=m_jk, other=0.0) b_kbg = b_kb * tl.where(m_j[:, None], exp(b_gk - b_gn[None, :]), 0) # [BC, BC] - b_dA = tl.load(p_dA, boundary_check=(0, 1)) + b_dA = tl.load(p_dA, mask=(o_dAi[:, None] < BT) & m_j[None, :], other=0.0) # [BC, BK] # (SY 09/17) important to not use bf16 here to have a good precision. - b_dkt += tl.dot(b_dA, b_kbg) + b_dkt = tl.dot(b_dA, b_kbg, b_dkt) b_dkt *= exp(b_gn[None, :] - b_g) o_dA = (i_t * BT + i_i * BC) * H*BT + i_i * BC + o_i p_kj = k + (i_t * BT + i_i * BC) * H*K + o_k @@ -377,10 +398,10 @@ def chunk_scaled_dot_kkt_bwd_kernel_gk( b_dg = (b_dk - b_dkt) * b_k b_dk += b_dkt - p_dk = tl.make_block_ptr(dk, (T, K), (H*K, 1), (i_t * BT + i_i * BC, i_k * BK), (BC, BK), (1, 0)) - p_dg = tl.make_block_ptr(dg, (T, K), (H*K, 1), (i_t * BT + i_i * BC, i_k * BK), (BC, BK), (1, 0)) - tl.store(p_dk, b_dk.to(p_dk.dtype.element_ty), boundary_check=(0, 1)) - tl.store(p_dg, b_dg.to(p_dg.dtype.element_ty), boundary_check=(0, 1)) + p_dk = dk + o_r[:, None] * (H*K) + o_k[None, :] + p_dg = dg + o_r[:, None] * (H*K) + o_k[None, :] + tl.store(p_dk, b_dk.to(p_dk.dtype.element_ty), mask=m_rk) + tl.store(p_dg, b_dg.to(p_dg.dtype.element_ty), mask=m_rk) def chunk_scaled_dot_kkt_fwd( @@ -500,7 +521,7 @@ def chunk_scaled_dot_kkt_bwd_gk( dk = torch.empty_like(k, dtype=torch.float) dg = torch.empty_like(g, dtype=torch.float) db = beta.new_empty(NK, *beta.shape, dtype=torch.float) - grid = (NK, NT * NC, B * H) + grid = (NK * NT * NC, B * H) chunk_scaled_dot_kkt_bwd_kernel_gk[grid]( k=k, g=g, diff --git a/build/torch-cuda/ops/gated_oja_rule/chunk_o.py b/build/torch-cuda/ops/gated_oja_rule/chunk_o.py index 43dc94d230b70475817c59d93f361e98ab5bac65..402057ba859313ab0f45c1f6200a97fb4d47f90a 100644 --- a/build/torch-cuda/ops/gated_oja_rule/chunk_o.py +++ b/build/torch-cuda/ops/gated_oja_rule/chunk_o.py @@ -10,7 +10,7 @@ import triton import triton.language as tl from ...ops.utils import prepare_chunk_indices -from ...ops.utils.op import exp +from ...ops.utils.op import exp, unflatten_program_id from ...utils import check_shared_mem, is_nvidia_hopper BKV_LIST = [64, 128] if check_shared_mem() else [32, 64] @@ -52,13 +52,14 @@ def chunk_oja_fwd_inter( NG: tl.constexpr, IS_VARLEN: tl.constexpr, ): - i_v, i_t, i_bh = tl.program_id(0), tl.program_id(1), tl.program_id(2) + i_v, i_t = unflatten_program_id(tl.cdiv(V, BV)) + i_bh = tl.program_id(1).to(tl.int64) i_b, i_hq = i_bh // HQ, i_bh % HQ i_h = i_hq // NG 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.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) + i_n, i_t = tl.load(chunk_indices + i_t * 2).to(tl.int32), tl.load(chunk_indices + i_t * 2 + 1).to(tl.int64) + bos, eos = tl.load(cu_seqlens + i_n).to(tl.int64), tl.load(cu_seqlens + i_n + 1).to(tl.int64) T = eos - bos NT = tl.cdiv(T, BT) else: @@ -69,36 +70,46 @@ def chunk_oja_fwd_inter( o_i = tl.arange(0, BT) m_s = o_i[:, None] >= o_i[None, :] + o_t = i_t * BT + tl.arange(0, BT) + o_v = i_v * BV + tl.arange(0, BV) + m_t = o_t < T + m_tv = m_t[:, None] & (o_v[None, :] < V) b_o = tl.zeros([BT, BV], dtype=tl.float32) b_A = tl.zeros([BT, BT], dtype=tl.float32) for i_k in range(tl.cdiv(K, BK)): - p_q = tl.make_block_ptr(q + (bos * HQ + i_hq) * K, (T, K), (HQ*K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) - p_k = tl.make_block_ptr(k + (bos * H + i_h) * K, (K, T), (1, H*K), (i_k * BK, i_t * BT), (BK, BT), (0, 1)) - p_h = tl.make_block_ptr(h + (i_tg * H + i_h) * K*V, (K, V), (V, 1), (i_k * BK, i_v * BV), (BK, BV), (1, 0)) + o_k = i_k * BK + tl.arange(0, BK) + m_tk = m_t[:, None] & (o_k[None, :] < K) + m_kt = (o_k[:, None] < K) & m_t[None, :] + m_kv = (o_k[:, None] < K) & (o_v[None, :] < V) + p_q = q + (bos * HQ + i_hq) * K + o_t[:, None] * (HQ*K) + o_k[None, :] + p_k = k + (bos * H + i_h) * K + o_k[:, None] + o_t[None, :] * (H*K) + p_h = h + (i_tg * H + i_h) * K*V + o_k[:, None] * V + o_v[None, :] # [BT, BK] - b_q = tl.load(p_q, boundary_check=(0, 1)) + b_q = tl.load(p_q, mask=m_tk, other=0.0) b_q = (b_q * scale).to(b_q.dtype) # [BK, BT] - b_k = tl.load(p_k, boundary_check=(0, 1)) + b_k = tl.load(p_k, mask=m_kt, other=0.0) # [BK, BV] - b_h = tl.load(p_h, boundary_check=(0, 1)) + b_h = tl.load(p_h, mask=m_kv, other=0.0) # [BT, BV] - b_o += tl.dot(b_q, b_h) + b_o = tl.dot(b_q, b_h, b_o) # [BT, BT] - b_A += tl.dot(b_q, b_k) - p_g = tl.make_block_ptr(gv + (bos * H + i_h) * V, (T, V), (H*V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) - p_o = tl.make_block_ptr(o + (bos * HQ + i_hq) * V, (T, V), (HQ*V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) - p_A = tl.make_block_ptr(A + (bos * HQ + i_hq) * BT, (T, BT), (HQ*BT, 1), (i_t * BT, 0), (BT, BT), (1, 0)) + b_A = tl.dot(b_q, b_k, b_A) + p_g = gv + (bos * H + i_h) * V + o_t[:, None] * (H*V) + o_v[None, :] + p_o = o + (bos * HQ + i_hq) * V + o_t[:, None] * (HQ*V) + o_v[None, :] + o_A = tl.arange(0, BT) + m_AT = m_t[:, None] & (o_A[None, :] < BT) + p_A = A + (bos * HQ + i_hq) * BT + o_t[:, None] * (HQ*BT) + o_A[None, :] # [BT, BV] - b_g = tl.load(p_g, boundary_check=(0, 1)) + b_g = tl.load(p_g, mask=m_tv, other=0.0) b_o = b_o * exp(b_g) - tl.store(p_o, b_o.to(p_o.dtype.element_ty), boundary_check=(0, 1)) + tl.store(p_o, b_o.to(p_o.dtype.element_ty), mask=m_tv) # [BT, BT] b_A = tl.where(m_s, b_A, 0.) if i_v == 0: - tl.store(p_A, b_A.to(p_A.dtype.element_ty), boundary_check=(0, 1)) + tl.store(p_A, b_A.to(p_A.dtype.element_ty), mask=m_AT) @triton.heuristics({ @@ -123,13 +134,14 @@ def chunk_oja_fwd_intra( NG: tl.constexpr, IS_VARLEN: tl.constexpr, ): - i_v, i_c, i_bh = tl.program_id(0), tl.program_id(1), tl.program_id(2) + NV = tl.cdiv(V, BV) + i_v, i_c, i_bh = tl.program_id(0).to(tl.int64) % NV, tl.program_id(0).to(tl.int64) // NV, tl.program_id(1).to(tl.int64) i_b, i_hq = i_bh // HQ, i_bh % HQ i_h = i_hq // NG i_t, i_i = i_c // NC, i_c % NC 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.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) + i_n, i_t = tl.load(chunk_indices + i_t * 2).to(tl.int32), tl.load(chunk_indices + i_t * 2 + 1).to(tl.int64) + bos, eos = tl.load(cu_seqlens + i_n).to(tl.int64), tl.load(cu_seqlens + i_n + 1).to(tl.int64) T = eos - bos else: bos, eos = i_b * T, i_b * T + T @@ -140,45 +152,51 @@ def chunk_oja_fwd_intra( if i_t * BT + i_i * BC >= T: return - p_g = tl.make_block_ptr(gv + (bos * H + i_h) * V, (T, V), (H*V, 1), (i_t * BT + i_i * BC, i_v * BV), (BC, BV), (1, 0)) + o_r = i_t * BT + i_i * BC + tl.arange(0, BC) + m_rv = (o_r[:, None] < T) & (o_v[None, :] < V) + p_g = gv + (bos * H + i_h) * V + o_r[:, None] * (H*V) + o_v[None, :] p_gn = gv + (bos + min(i_t * BT + i_i * BC, T)) * H*V + i_h * V + o_v # [BV,] b_gn = tl.load(p_gn, mask=m_v, other=0) # [BC, BV] b_o = tl.zeros([BC, BV], dtype=tl.float32) for i_j in range(0, i_i): - p_A = tl.make_block_ptr(A + (bos*HQ+i_hq) * BT, (T, BT), (HQ*BT, 1), (i_t*BT+i_i*BC, i_j * BC), (BC, BC), (1, 0)) - p_v = tl.make_block_ptr(v + (bos*H+i_h) * V, (T, V), (H*V, 1), (i_t * BT + i_j * BC, i_v * BV), (BC, BV), (1, 0)) - p_gv = tl.make_block_ptr(gv + (bos*H+i_h) * V, (T, V), (H*V, 1), (i_t * BT + i_j * BC, i_v * BV), (BC, BV), (1, 0)) + o_j = i_t * BT + i_j * BC + tl.arange(0, BC) + o_Aj = i_j * BC + tl.arange(0, BC) + m_jv = (o_j[:, None] < T) & (o_v[None, :] < V) + m_A = (o_r[:, None] < T) & (o_Aj[None, :] < BT) + p_A = A + (bos*HQ+i_hq) * BT + o_r[:, None] * (HQ*BT) + o_Aj[None, :] + p_v = v + (bos*H+i_h) * V + o_j[:, None] * (H*V) + o_v[None, :] + p_gv = gv + (bos*H+i_h) * V + o_j[:, None] * (H*V) + o_v[None, :] # [BC, BV] - b_v = tl.load(p_v, boundary_check=(0, 1)) - b_gv = tl.load(p_gv, boundary_check=(0, 1)) + b_v = tl.load(p_v, mask=m_jv, other=0.0) + b_gv = tl.load(p_gv, mask=m_jv, other=0.0) b_vg = (b_v * exp(b_gn[None, :] - b_gv)).to(b_v.dtype) # [BC, BC] - b_A = tl.load(p_A, boundary_check=(0, 1)) - b_o += tl.dot(b_A, b_vg) + b_A = tl.load(p_A, mask=m_A, other=0.0) + b_o = tl.dot(b_A, b_vg, b_o) # [BC, BV] - b_g = tl.load(p_g, boundary_check=(0, 1)) + b_g = tl.load(p_g, mask=m_rv, other=0.0) b_o *= exp(b_g - b_gn[None, :]) o_i = tl.arange(0, BC) o_A = (bos + i_t * BT + i_i * BC + tl.arange(0, BC)) * HQ*BT + i_hq * BT + i_i * BC m_A = (i_t * BT + i_i * BC + tl.arange(0, BC)) < T for j in range(0, min(BC, T - i_t * BT - i_i * BC)): - p_v = v + (bos + i_t * BT + i_i * BC + j) * H*V + i_h * V + o_v - p_gv = gv + (bos + i_t * BT + i_i * BC + j) * H*V + i_h * V + o_v + p_vj = v + (bos + i_t * BT + i_i * BC + j) * H*V + i_h * V + o_v + p_gvj = gv + (bos + i_t * BT + i_i * BC + j) * H*V + i_h * V + o_v # [BC,] b_A = tl.load(A + o_A + j, mask=m_A, other=0) # [BV,] - b_v = tl.load(p_v, mask=m_v, other=0).to(tl.float32) - b_gv = tl.load(p_gv, mask=m_v, other=0).to(tl.float32) + b_v = tl.load(p_vj, mask=m_v, other=0).to(tl.float32) + b_gv = tl.load(p_gvj, mask=m_v, other=0).to(tl.float32) # [BC, BV] b_vg = b_v[None, :] * exp(b_g - b_gv[None, :]) # avoid 0 * inf = inf b_o += tl.where(o_i[:, None] >= j, b_A[:, None] * b_vg, 0.) - p_o = tl.make_block_ptr(o + (bos*HQ + i_hq) * V, (T, V), (HQ*V, 1), (i_t * BT + i_i * BC, i_v * BV), (BC, BV), (1, 0)) - b_o += tl.load(p_o, boundary_check=(0, 1)) - tl.store(p_o, b_o.to(p_o.dtype.element_ty), boundary_check=(0, 1)) + p_o = o + (bos*HQ + i_hq) * V + o_r[:, None] * (HQ*V) + o_v[None, :] + b_o += tl.load(p_o, mask=m_rv, other=0.0) + tl.store(p_o, b_o.to(p_o.dtype.element_ty), mask=m_rv) def chunk_oja_fwd_o( @@ -206,7 +224,7 @@ def chunk_oja_fwd_o( o = v.new_empty(B, T, HQ, V) A = q.new_empty(B, T, HQ, BT) - def grid(meta): return (triton.cdiv(V, meta['BV']), NT, B * HQ) + def grid(meta): return (triton.cdiv(V, meta['BV']) * NT, B * HQ) chunk_oja_fwd_inter[grid]( q, k, @@ -226,7 +244,7 @@ def chunk_oja_fwd_o( NG=NG, ) - def grid(meta): return (triton.cdiv(V, meta['BV']), NT * NC, B * HQ) + def grid(meta): return (triton.cdiv(V, meta['BV']) * NT * NC, B * HQ) chunk_oja_fwd_intra[grid]( v, gv, @@ -278,12 +296,13 @@ def chunk_oja_bwd_kernel_dA( NC: tl.constexpr, IS_VARLEN: tl.constexpr, ): - i_v, i_c, i_bh = tl.program_id(0), tl.program_id(1), tl.program_id(2) + NV = tl.cdiv(V, BV) + i_v, i_c, i_bh = tl.program_id(0).to(tl.int64) % NV, tl.program_id(0).to(tl.int64) // NV, tl.program_id(1).to(tl.int64) i_b, i_h = i_bh // H, i_bh % H i_t, i_i, i_j = i_c // (NC * NC), (i_c % (NC * NC)) // NC, (i_c % (NC * NC)) % NC 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.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) + i_n, i_t = tl.load(chunk_indices + i_t * 2).to(tl.int32), tl.load(chunk_indices + i_t * 2 + 1).to(tl.int64) + bos, eos = tl.load(cu_seqlens + i_n).to(tl.int64), tl.load(cu_seqlens + i_n + 1).to(tl.int64) all = T T = eos - bos else: @@ -299,31 +318,37 @@ def chunk_oja_bwd_kernel_dA( # [BC, BC] b_dA = tl.zeros([BC, BC], dtype=tl.float32) if i_i > i_j: - p_v = tl.make_block_ptr(v + (bos*H+i_h) * V, (V, T), (1, H*V), (i_v * BV, i_t*BT + i_j*BC), (BV, BC), (0, 1)) - p_gv = tl.make_block_ptr(gv + (bos*H+i_h) * V, (V, T), (1, H*V), (i_v * BV, i_t*BT + i_j*BC), (BV, BC), (0, 1)) + o_r = i_t * BT + i_i * BC + tl.arange(0, BC) + o_kj = i_t * BT + i_j * BC + tl.arange(0, BC) + m_rv = (o_r[:, None] < T) & (o_v[None, :] < V) + m_vj = (o_v[:, None] < V) & (o_kj[None, :] < T) + p_v = v + (bos*H+i_h) * V + o_v[:, None] + o_kj[None, :] * (H*V) + p_gv = gv + (bos*H+i_h) * V + o_v[:, None] + o_kj[None, :] * (H*V) p_gn = gv + (bos + i_t*BT + i_i*BC) * H*V + i_h * V + o_v - p_g = tl.make_block_ptr(gv + (bos*H+i_h) * V, (T, V), (H*V, 1), (i_t*BT + i_i*BC, i_v*BV), (BC, BV), (1, 0)) - p_do = tl.make_block_ptr(do + (bos*H+i_h) * V, (T, V), (H*V, 1), (i_t*BT + i_i*BC, i_v*BV), (BC, BV), (1, 0)) + p_g = gv + (bos*H+i_h) * V + o_r[:, None] * (H*V) + o_v[None, :] + p_do = do + (bos*H+i_h) * V + o_r[:, None] * (H*V) + o_v[None, :] # [BV,] b_gn = tl.load(p_gn, mask=m_v, other=0.) # [BC, BV] - b_g = tl.load(p_g, boundary_check=(0, 1)) - b_do = tl.load(p_do, boundary_check=(0, 1)) + b_g = tl.load(p_g, mask=m_rv, other=0.0) + b_do = tl.load(p_do, mask=m_rv, other=0.0) b_do = (b_do * exp(b_g - b_gn[None, :]) * scale).to(b_do.dtype) # [BV, BC] - b_v = tl.load(p_v, boundary_check=(0, 1)) - b_gv = tl.load(p_gv, boundary_check=(0, 1)) + b_v = tl.load(p_v, mask=m_vj, other=0.0) + b_gv = tl.load(p_gv, mask=m_vj, other=0.0) b_vg = (b_v * exp(b_gn[:, None] - b_gv)).to(b_v.dtype) # [BC, BC] b_dA = tl.dot(b_do, b_vg) elif i_i == i_j: - p_g = tl.make_block_ptr(gv + (bos*H + i_h) * V, (T, V), (H*V, 1), (i_t*BT + i_i*BC, i_v*BV), (BC, BV), (1, 0)) - p_do = tl.make_block_ptr(do + (bos*H + i_h) * V, (T, V), (H*V, 1), (i_t*BT + i_i*BC, i_v*BV), (BC, BV), (1, 0)) - p_v = v + (bos + i_t*BT + i_j*BC) * H*V + i_h * V + o_v - p_gv = gv + (bos + i_t*BT + i_j*BC) * H*V + i_h * V + o_v + o_r = i_t * BT + i_i * BC + tl.arange(0, BC) + m_rv = (o_r[:, None] < T) & (o_v[None, :] < V) + p_g = gv + (bos*H + i_h) * V + o_r[:, None] * (H*V) + o_v[None, :] + p_do = do + (bos*H + i_h) * V + o_r[:, None] * (H*V) + o_v[None, :] + p_vj = v + (bos + i_t*BT + i_j*BC) * H*V + i_h * V + o_v + p_gvj = gv + (bos + i_t*BT + i_j*BC) * H*V + i_h * V + o_v # [BC, BV] - b_g = tl.load(p_g, boundary_check=(0, 1)) - b_do = tl.load(p_do, boundary_check=(0, 1)) * scale + b_g = tl.load(p_g, mask=m_rv, other=0.0) + b_do = tl.load(p_do, mask=m_rv, other=0.0) * scale m_v = o_v < V o_i = tl.arange(0, BC) @@ -331,18 +356,20 @@ def chunk_oja_bwd_kernel_dA( m_dA = o_i[:, None] >= o_i[None, :] for j in range(0, min(BC, T - i_t * BT - i_j * BC)): # [BV,] - b_v = tl.load(p_v, mask=m_v, other=0).to(tl.float32) - b_gv = tl.load(p_gv, mask=m_v, other=0).to(tl.float32) + b_v = tl.load(p_vj, mask=m_v, other=0).to(tl.float32) + b_gv = tl.load(p_gvj, mask=m_v, other=0).to(tl.float32) # [BC,] b_dAj = tl.sum(b_do * b_v[None, :] * exp(b_g - b_gv[None, :]), 1) b_dA = tl.where((o_i == j)[None, :], b_dAj[:, None], b_dA) - p_v += H*V - p_gv += H*V + p_vj += H*V + p_gvj += H*V b_dA = tl.where(m_dA, b_dA, 0.) - p_dA = tl.make_block_ptr(dA+((i_v*all+bos)*H+i_h)*BT, (T, BT), (H*BT, 1), (i_t*BT+i_i*BC, i_j*BC), (BC, BC), (1, 0)) - tl.store(p_dA, b_dA.to(dA.dtype.element_ty), boundary_check=(0, 1)) + o_r = i_t * BT + i_i * BC + tl.arange(0, BC) + o_Aj = i_j * BC + tl.arange(0, BC) + p_dA = dA+((i_v*all+bos)*H+i_h)*BT + o_r[:, None] * (H*BT) + o_Aj[None, :] + tl.store(p_dA, b_dA.to(dA.dtype.element_ty), mask=(o_r[:, None] < T) & (o_Aj[None, :] < BT)) def chunk_oja_bwd_dA( @@ -366,8 +393,8 @@ def chunk_oja_bwd_dA( NV = triton.cdiv(V, BV) dA = v.new_empty(NV, B, T, H, BT) - # 计算dA - grid = (NV, NT * NC * NC, B * H) + # compute dA + grid = (NV * NT * NC * NC, B * H) chunk_oja_bwd_kernel_dA[grid]( v, gv, @@ -425,12 +452,13 @@ def chunk_oja_bwd_kernel_dqk( BV: tl.constexpr, IS_VARLEN: tl.constexpr, ): - i_k, i_t, i_bh = tl.program_id(0), tl.program_id(1), tl.program_id(2) + i_k, i_t = unflatten_program_id(tl.cdiv(K, BK)) + i_bh = tl.program_id(1).to(tl.int64) i_b, i_h = i_bh // H, i_bh % H 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.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) + i_n, i_t = tl.load(chunk_indices + i_t * 2).to(tl.int32), tl.load(chunk_indices + i_t * 2 + 1).to(tl.int64) + bos, eos = tl.load(cu_seqlens + i_n).to(tl.int64), tl.load(cu_seqlens + i_n + 1).to(tl.int64) all = T T = eos - bos NT = tl.cdiv(T, BT) @@ -438,47 +466,56 @@ def chunk_oja_bwd_kernel_dqk( NT = tl.cdiv(T, BT) i_tg = i_b * NT + i_t bos, eos = i_b * T, i_b * T + T - all = B * T + all = tl.cast(B * T, tl.int64) o_i = tl.arange(0, BT) m_s = o_i[:, None] >= o_i[None, :] + o_t = i_t * BT + tl.arange(0, BT) + o_k = i_k * BK + tl.arange(0, BK) + m_t = o_t < T + m_tk = m_t[:, None] & (o_k[None, :] < K) + o_A = tl.arange(0, BT) + m_AT = m_t[:, None] & (o_A[None, :] < BT) # [B, T, H, BT] - p_q = tl.make_block_ptr(q + (bos*H+i_h) * K, (T, K), (H*K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) - p_k = tl.make_block_ptr(k + (bos*H+i_h) * K, (T, K), (H*K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) - p_A = tl.make_block_ptr(A + ((i_k*all+bos)*H+i_h)*BT, (T, BT), (H*BT, 1), (i_t * BT, 0), (BT, BT), (1, 0)) - b_q = tl.load(p_q, boundary_check=(0, 1)) - b_k = tl.load(p_k, boundary_check=(0, 1)) + p_q = q + (bos*H+i_h) * K + o_t[:, None] * (H*K) + o_k[None, :] + p_k = k + (bos*H+i_h) * K + o_t[:, None] * (H*K) + o_k[None, :] + p_A = A + ((i_k*all+bos)*H+i_h)*BT + o_t[:, None] * (H*BT) + o_A[None, :] + b_q = tl.load(p_q, mask=m_tk, other=0.0) + b_k = tl.load(p_k, mask=m_tk, other=0.0) b_A = tl.dot((b_q * scale).to(b_q.dtype), tl.trans(b_k)) b_A = tl.where(m_s, b_A, 0.) - tl.store(p_A, b_A.to(p_A.dtype.element_ty), boundary_check=(0, 1)) + tl.store(p_A, b_A.to(p_A.dtype.element_ty), mask=m_AT) b_dq = tl.zeros([BT, BK], dtype=tl.float32) - # 先计算do对应的dq + # compute dq from do for i_v in range(tl.cdiv(V, BV)): - p_h = tl.make_block_ptr(h + (i_tg * H + i_h) * K*V, (V, K), (1, V), (i_v * BV, i_k * BK), (BV, BK), (0, 1)) - p_do = tl.make_block_ptr(do + (bos*H+i_h)*V, (T, V), (H*V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) - p_gv = tl.make_block_ptr(gv + (bos*H+i_h)*V, (T, V), (H*V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) - b_h = tl.load(p_h, boundary_check=(0, 1)) - b_do = tl.load(p_do, boundary_check=(0, 1)) - b_gv = tl.load(p_gv, boundary_check=(0, 1)) + o_v = i_v * BV + tl.arange(0, BV) + m_h = (o_v[:, None] < V) & (o_k[None, :] < K) + m_tv = m_t[:, None] & (o_v[None, :] < V) + p_h = h + (i_tg * H + i_h) * K*V + o_v[:, None] + o_k[None, :] * V + p_do = do + (bos*H+i_h)*V + o_t[:, None] * (H*V) + o_v[None, :] + p_gv = gv + (bos*H+i_h)*V + o_t[:, None] * (H*V) + o_v[None, :] + b_h = tl.load(p_h, mask=m_h, other=0.0) + b_do = tl.load(p_do, mask=m_tv, other=0.0) + b_gv = tl.load(p_gv, mask=m_tv, other=0.0) b_do = (b_do * exp(b_gv) * scale).to(b_do.dtype) - b_dq += tl.dot(b_do, b_h.to(b_do.dtype)) + b_dq = tl.dot(b_do, b_h.to(b_do.dtype), b_dq) - # 接着计算dA对应的dq, dk - p_dA = tl.make_block_ptr(dA + (bos*H + i_h) * BT, (T, BT), (H*BT, 1), (i_t * BT, 0), (BT, BT), (1, 0)) - p_dq = tl.make_block_ptr(dq + (bos*H + i_h) * K, (T, K), (H*K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) - p_dk = tl.make_block_ptr(dk + (bos*H + i_h) * K, (T, K), (H*K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) + # compute dq and dk from dA + p_dA = dA + (bos*H + i_h) * BT + o_t[:, None] * (H*BT) + o_A[None, :] + p_dq = dq + (bos*H + i_h) * K + o_t[:, None] * (H*K) + o_k[None, :] + p_dk = dk + (bos*H + i_h) * K + o_t[:, None] * (H*K) + o_k[None, :] # [BT, BT] - b_dA = tl.load(p_dA, boundary_check=(0, 1)) + b_dA = tl.load(p_dA, mask=m_AT, other=0.0) # [BT, BK] - b_dq += tl.dot(b_dA.to(b_q.dtype), b_k) + b_dq = tl.dot(b_dA.to(b_q.dtype), b_k, b_dq) b_dk = tl.dot(tl.trans(b_dA).to(b_q.dtype), b_q) - tl.store(p_dq, b_dq.to(p_dq.dtype.element_ty), boundary_check=(0, 1)) - tl.store(p_dk, b_dk.to(p_dk.dtype.element_ty), boundary_check=(0, 1)) + tl.store(p_dq, b_dq.to(p_dq.dtype.element_ty), mask=m_tk) + tl.store(p_dk, b_dk.to(p_dk.dtype.element_ty), mask=m_tk) def chunk_oja_bwd_dqk( @@ -506,8 +543,8 @@ def chunk_oja_bwd_dqk( dq = torch.empty_like(q) dk = torch.empty_like(k) A = dA.new_empty(NK, B, T, H, BT) - # 计算dA - grid = (NK, NT, B * H) + # compute A, dq and dk + grid = (NK * NT, B * H) chunk_oja_bwd_kernel_dqk[grid]( q, k, @@ -560,12 +597,13 @@ def chunk_oja_bwd_kernel_dv_o( NC: tl.constexpr, IS_VARLEN: tl.constexpr, ): - i_v, i_c, i_bh = tl.program_id(0), tl.program_id(1), tl.program_id(2) + NV = tl.cdiv(V, BV) + i_v, i_c, i_bh = tl.program_id(0).to(tl.int64) % NV, tl.program_id(0).to(tl.int64) // NV, tl.program_id(1).to(tl.int64) i_b, i_h = i_bh // H, i_bh % H i_t, i_i = i_c // NC, i_c % NC 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.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) + i_n, i_t = tl.load(chunk_indices + i_t * 2).to(tl.int32), tl.load(chunk_indices + i_t * 2 + 1).to(tl.int64) + bos, eos = tl.load(cu_seqlens + i_n).to(tl.int64), tl.load(cu_seqlens + i_n + 1).to(tl.int64) T = eos - bos else: bos, eos = i_b * T, i_b * T + T @@ -576,24 +614,30 @@ def chunk_oja_bwd_kernel_dv_o( if i_t * BT + i_i * BC >= T: return - p_gv = tl.make_block_ptr(g + (bos*H+i_h)*V, (T, V), (H*V, 1), (i_t * BT + i_i * BC, i_v * BV), (BC, BV), (1, 0)) + o_r = i_t * BT + i_i * BC + tl.arange(0, BC) + m_rv = (o_r[:, None] < T) & (o_v[None, :] < V) + p_gv = g + (bos*H+i_h)*V + o_r[:, None] * (H*V) + o_v[None, :] p_gn = g + (bos + min(i_t * BT + i_i * BC + BC, T)-1)*H*V + i_h*V + o_v # [BV,] b_gn = tl.load(p_gn, mask=m_v, other=0) # [BC, BV] - b_gv = tl.load(p_gv, boundary_check=(0, 1)) + b_gv = tl.load(p_gv, mask=m_rv, other=0.0) b_dvg = tl.zeros([BC, BV], dtype=tl.float32) for i_j in range(i_i + 1, NC): - p_g = tl.make_block_ptr(g + (bos*H+i_h) * V, (T, V), (H*V, 1), (i_t * BT + i_j * BC, i_v * BV), (BC, BV), (1, 0)) - p_A = tl.make_block_ptr(A + (bos*H+i_h) * BT, (BT, T), (1, H*BT), (i_i*BC, i_t*BT + i_j*BC), (BC, BC), (0, 1)) - p_do = tl.make_block_ptr(do + (bos*H+i_h) * V, (T, V), (H*V, 1), (i_t*BT + i_j*BC, i_v*BV), (BC, BV), (1, 0)) + o_j = i_t * BT + i_j * BC + tl.arange(0, BC) + o_Ai = i_i * BC + tl.arange(0, BC) + m_jv = (o_j[:, None] < T) & (o_v[None, :] < V) + m_A = (o_Ai[:, None] < BT) & (o_j[None, :] < T) + p_g = g + (bos*H+i_h) * V + o_j[:, None] * (H*V) + o_v[None, :] + p_A = A + (bos*H+i_h) * BT + o_Ai[:, None] + o_j[None, :] * (H*BT) + p_do = do + (bos*H+i_h) * V + o_j[:, None] * (H*V) + o_v[None, :] # [BC, BV] - b_g = tl.load(p_g, boundary_check=(0, 1)) - b_do = tl.load(p_do, boundary_check=(0, 1)) * exp(b_g - b_gn[None, :]) + b_g = tl.load(p_g, mask=m_jv, other=0.0) + b_do = tl.load(p_do, mask=m_jv, other=0.0) * exp(b_g - b_gn[None, :]) # [BC, BC] - b_A = tl.load(p_A, boundary_check=(0, 1)) + b_A = tl.load(p_A, mask=m_A, other=0.0) # [BC, BV] - b_dvg += tl.dot(b_A, b_do.to(b_A.dtype)) + b_dvg = tl.dot(b_A, b_do.to(b_A.dtype), b_dvg) b_dv = b_dvg * exp(b_gn[None, :] - b_gv) o_i = tl.arange(0, BC) @@ -615,20 +659,20 @@ def chunk_oja_bwd_kernel_dv_o( p_g += H * V p_A += H * BT p_do += H * V - p_o = tl.make_block_ptr(o + (bos*H+i_h)*V, (T, V), (H*V, 1), (i_t*BT + i_i*BC, i_v*BV), (BC, BV), (1, 0)) - p_v = tl.make_block_ptr(v + (bos*H+i_h)*V, (T, V), (H*V, 1), (i_t*BT + i_i*BC, i_v*BV), (BC, BV), (1, 0)) - p_do = tl.make_block_ptr(do + (bos*H+i_h)*V, (T, V), (H*V, 1), (i_t*BT + i_i*BC, i_v*BV), (BC, BV), (1, 0)) - p_dv = tl.make_block_ptr(dv + (bos*H+i_h)*V, (T, V), (H*V, 1), (i_t*BT + i_i*BC, i_v*BV), (BC, BV), (1, 0)) - p_dv2 = tl.make_block_ptr(dv2 + (bos*H+i_h)*V, (T, V), (H*V, 1), (i_t*BT + i_i*BC, i_v*BV), (BC, BV), (1, 0)) - p_dg = tl.make_block_ptr(dg + (bos*H+i_h)*V, (T, V), (H*V, 1), (i_t*BT + i_i*BC, i_v*BV), (BC, BV), (1, 0)) - - b_o = tl.load(p_o, boundary_check=(0, 1)).to(tl.float32) - b_v = tl.load(p_v, boundary_check=(0, 1)).to(tl.float32) - b_do = tl.load(p_do, boundary_check=(0, 1)).to(tl.float32) - b_dv = b_dv + tl.load(p_dv, boundary_check=(0, 1)).to(tl.float32) + p_o = o + (bos*H+i_h)*V + o_r[:, None] * (H*V) + o_v[None, :] + p_v = v + (bos*H+i_h)*V + o_r[:, None] * (H*V) + o_v[None, :] + p_do = do + (bos*H+i_h)*V + o_r[:, None] * (H*V) + o_v[None, :] + p_dv = dv + (bos*H+i_h)*V + o_r[:, None] * (H*V) + o_v[None, :] + p_dv2 = dv2 + (bos*H+i_h)*V + o_r[:, None] * (H*V) + o_v[None, :] + p_dg = dg + (bos*H+i_h)*V + o_r[:, None] * (H*V) + o_v[None, :] + + b_o = tl.load(p_o, mask=m_rv, other=0.0).to(tl.float32) + b_v = tl.load(p_v, mask=m_rv, other=0.0).to(tl.float32) + b_do = tl.load(p_do, mask=m_rv, other=0.0).to(tl.float32) + b_dv = b_dv + tl.load(p_dv, mask=m_rv, other=0.0).to(tl.float32) b_dg = b_o * b_do - b_v * b_dv - tl.store(p_dv2, b_dv.to(p_dv.dtype.element_ty), boundary_check=(0, 1)) - tl.store(p_dg, b_dg.to(p_dg.dtype.element_ty), boundary_check=(0, 1)) + tl.store(p_dv2, b_dv.to(p_dv.dtype.element_ty), mask=m_rv) + tl.store(p_dg, b_dg.to(p_dg.dtype.element_ty), mask=m_rv) def chunk_oja_bwd_dv_o( @@ -654,8 +698,8 @@ def chunk_oja_bwd_dv_o( dv2 = torch.empty_like(v, dtype=torch.float) dgv = torch.empty_like(gv) - # 计算dA - def grid(meta): return (triton.cdiv(V, meta['BV']), NT * NC, B * H) + # compute dv and dgv + def grid(meta): return (triton.cdiv(V, meta['BV']) * NT * NC, B * H) chunk_oja_bwd_kernel_dv_o[grid]( v=v, g=gv, diff --git a/build/torch-cuda/ops/gated_oja_rule/fused_recurrent.py b/build/torch-cuda/ops/gated_oja_rule/fused_recurrent.py index 4535fdc40137e59f53e36beea8620d03d4f7223c..991034f577e75b3daf940088d319b54b99dd3176 100644 --- a/build/torch-cuda/ops/gated_oja_rule/fused_recurrent.py +++ b/build/torch-cuda/ops/gated_oja_rule/fused_recurrent.py @@ -47,7 +47,9 @@ def fused_recurrent_oja_fwd_kernel( STORE_FINAL_STATE: tl.constexpr, IS_VARLEN: tl.constexpr, ): - i_v, i_nh = tl.program_id(0), tl.program_id(1) + pid = tl.program_id(0).to(tl.int64) + NV = tl.cdiv(V, BV) + i_v, i_nh = pid % NV, (pid // NV).to(tl.int64) i_n, i_hv = i_nh // HV, i_nh % HV i_h = i_hv // (HV // H) if IS_VARLEN: @@ -143,7 +145,7 @@ def fused_recurrent_oja_fwd( o = torch.empty_like(v) final_state = q.new_empty(N, HV, K, V, dtype=torch.float32) if output_final_state else None - grid = (NV, N * HV) + grid = (NV * N * HV,) fused_recurrent_oja_fwd_kernel[grid]( q=q, k=k, diff --git a/build/torch-cuda/ops/gated_oja_rule/wy_fast.py b/build/torch-cuda/ops/gated_oja_rule/wy_fast.py index d93a873f411cc297baceeb6c0642b7285d4ef5de..f9869a9bf54c3d58f3a36848364e62508509a43e 100644 --- a/build/torch-cuda/ops/gated_oja_rule/wy_fast.py +++ b/build/torch-cuda/ops/gated_oja_rule/wy_fast.py @@ -48,50 +48,57 @@ def recompute_w_u_fwd_kernel( STORE_VG: tl.constexpr, IS_VARLEN: tl.constexpr ): - i_t, i_bh = tl.program_id(0), tl.program_id(1) + i_t, i_bh = tl.program_id(0).to(tl.int64), tl.program_id(1).to(tl.int64) i_b, i_h = i_bh // H, i_bh % H 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.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) + i_n, i_t = tl.load(chunk_indices + i_t * 2).to(tl.int32), tl.load(chunk_indices + i_t * 2 + 1).to(tl.int64) + bos, eos = tl.load(cu_seqlens + i_n).to(tl.int64), tl.load(cu_seqlens + i_n + 1).to(tl.int64) T = eos - bos else: bos, eos = i_b * T, i_b * T + T - p_b = tl.make_block_ptr(beta + bos*H + i_h, (T,), (H,), (i_t * BT,), (BT,), (0,)) - b_b = tl.load(p_b, boundary_check=(0,)) + o_t = i_t * BT + tl.arange(0, BT) + m_t = o_t < T + p_b = beta + bos*H + i_h + o_t * H + b_b = tl.load(p_b, mask=m_t, other=0.0) - p_A = tl.make_block_ptr(A + (bos*H + i_h) * BT, (T, BT), (H*BT, 1), (i_t * BT, 0), (BT, BT), (1, 0)) - b_A = tl.load(p_A, boundary_check=(0, 1)) + o_A = tl.arange(0, BT) + m_AT = m_t[:, None] & (o_A[None, :] < BT) + p_A = A + (bos*H + i_h) * BT + o_t[:, None] * (H*BT) + o_A[None, :] + b_A = tl.load(p_A, mask=m_AT, other=0.0) for i_v in range(tl.cdiv(V, BV)): - p_v = tl.make_block_ptr(v + (bos*H + i_h) * V, (T, V), (H*V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) - p_w = tl.make_block_ptr(w + (bos*H + i_h) * V, (T, V), (H*V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) - b_v = tl.load(p_v, boundary_check=(0, 1)) + o_v = i_v * BV + tl.arange(0, BV) + m_tv = m_t[:, None] & (o_v[None, :] < V) + p_v = v + (bos*H + i_h) * V + o_t[:, None] * (H*V) + o_v[None, :] + p_w = w + (bos*H + i_h) * V + o_t[:, None] * (H*V) + o_v[None, :] + b_v = tl.load(p_v, mask=m_tv, other=0.0) b_vb = b_v * b_b[:, None] - p_gv = tl.make_block_ptr(gv + (bos*H + i_h) * V, (T, V), (H*V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) - b_gv = tl.load(p_gv, boundary_check=(0, 1)) + p_gv = gv + (bos*H + i_h) * V + o_t[:, None] * (H*V) + o_v[None, :] + b_gv = tl.load(p_gv, mask=m_tv, other=0.0) b_vb *= exp(b_gv) if STORE_VG: last_idx = min(i_t * BT + BT, T) - 1 - o_v = i_v * BV + tl.arange(0, BV) m_v = o_v < V b_gn = tl.load(gv + ((bos + last_idx) * H + i_h) * V + o_v, mask=m_v, other=0.) b_vg = b_v * exp(b_gn - b_gv) - p_vg = tl.make_block_ptr(vg + (bos * H + i_h) * V, (T, V), (H*V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) - tl.store(p_vg, b_vg.to(p_vg.dtype.element_ty), boundary_check=(0, 1)) + p_vg = vg + (bos * H + i_h) * V + o_t[:, None] * (H*V) + o_v[None, :] + tl.store(p_vg, b_vg.to(p_vg.dtype.element_ty), mask=m_tv) b_w = tl.dot(b_A, b_vb.to(b_v.dtype)) - tl.store(p_w, b_w.to(p_w.dtype.element_ty), boundary_check=(0, 1)) + tl.store(p_w, b_w.to(p_w.dtype.element_ty), mask=m_tv) for i_k in range(tl.cdiv(K, BK)): - p_k = tl.make_block_ptr(k + (bos*H + i_h) * K, (T, K), (H*K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) - p_u = tl.make_block_ptr(u + (bos*H + i_h) * K, (T, K), (H*K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) - b_k = tl.load(p_k, boundary_check=(0, 1)) + o_k = i_k * BK + tl.arange(0, BK) + m_tk = m_t[:, None] & (o_k[None, :] < K) + p_k = k + (bos*H + i_h) * K + o_t[:, None] * (H*K) + o_k[None, :] + p_u = u + (bos*H + i_h) * K + o_t[:, None] * (H*K) + o_k[None, :] + b_k = tl.load(p_k, mask=m_tk, other=0.0) b_kb = (b_k * b_b[:, None]).to(b_k.dtype) b_u = tl.dot(b_A, b_kb, allow_tf32=False) - tl.store(p_u, b_u.to(p_u.dtype.element_ty), boundary_check=(0, 1)) + tl.store(p_u, b_u.to(p_u.dtype.element_ty), mask=m_tk) @triton.heuristics({ @@ -130,60 +137,66 @@ def prepare_wy_repr_bwd_kernel( BV: tl.constexpr, IS_VARLEN: tl.constexpr ): - i_t, i_bh = tl.program_id(0), tl.program_id(1) + i_t, i_bh = tl.program_id(0).to(tl.int64), tl.program_id(1).to(tl.int64) i_b, i_h = i_bh // H, i_bh % H 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.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) + i_n, i_t = tl.load(chunk_indices + i_t * 2).to(tl.int32), tl.load(chunk_indices + i_t * 2 + 1).to(tl.int64) + bos, eos = tl.load(cu_seqlens + i_n).to(tl.int64), tl.load(cu_seqlens + i_n + 1).to(tl.int64) T = eos - bos else: bos, eos = i_b * T, i_b * T + T - p_b = tl.make_block_ptr(beta + (bos*H + i_h), (T,), (H,), (i_t * BT,), (BT,), (0,)) - p_db = tl.make_block_ptr(db + (bos*H + i_h), (T,), (H,), (i_t * BT,), (BT,), (0,)) - p_A = tl.make_block_ptr(A + (bos*H + i_h) * BT, (BT, T), (1, H*BT), (0, i_t * BT), (BT, BT), (0, 1)) + o_t = i_t * BT + tl.arange(0, BT) + m_t = o_t < T + o_A = tl.arange(0, BT) + m_AT = (o_A[:, None] < BT) & m_t[None, :] + p_b = beta + (bos*H + i_h) + o_t * H + p_db = db + (bos*H + i_h) + o_t * H + p_A = A + (bos*H + i_h) * BT + o_A[:, None] + o_t[None, :] * (H*BT) - b_b = tl.load(p_b, boundary_check=(0,)) + b_b = tl.load(p_b, mask=m_t, other=0.0) b_db = tl.zeros([BT], dtype=tl.float32) - b_A = tl.load(p_A, boundary_check=(0, 1)) + b_A = tl.load(p_A, mask=m_AT, other=0.0) b_dA = tl.zeros([BT, BT], dtype=tl.float32) for i_v in range(tl.cdiv(V, BV)): - p_v = tl.make_block_ptr(v + (bos*H + i_h) * V, (T, V), (H*V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) - p_dv = tl.make_block_ptr(dv + (bos*H + i_h) * V, (T, V), (H*V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) - p_dw = tl.make_block_ptr(dw + (bos*H + i_h) * V, (T, V), (H*V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) - b_v = tl.load(p_v, boundary_check=(0, 1)) - p_gv = tl.make_block_ptr(gv + (bos*H + i_h) * V, (T, V), (H*V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) - b_gv_exp = exp(tl.load(p_gv, boundary_check=(0, 1))) + o_v = i_v * BV + tl.arange(0, BV) + m_tv = m_t[:, None] & (o_v[None, :] < V) + p_v = v + (bos*H + i_h) * V + o_t[:, None] * (H*V) + o_v[None, :] + p_dv = dv + (bos*H + i_h) * V + o_t[:, None] * (H*V) + o_v[None, :] + p_dw = dw + (bos*H + i_h) * V + o_t[:, None] * (H*V) + o_v[None, :] + b_v = tl.load(p_v, mask=m_tv, other=0.0) + p_gv = gv + (bos*H + i_h) * V + o_t[:, None] * (H*V) + o_v[None, :] + b_gv_exp = exp(tl.load(p_gv, mask=m_tv, other=0.0)) b_vbg = b_v * b_b[:, None] * b_gv_exp - b_dw = tl.load(p_dw, boundary_check=(0, 1)) + b_dw = tl.load(p_dw, mask=m_tv, other=0.0) - b_dA += tl.dot(b_dw, tl.trans(b_vbg).to(b_dw.dtype)) + b_dA = tl.dot(b_dw, tl.trans(b_vbg).to(b_dw.dtype), b_dA) b_dvbg = tl.dot(b_A, b_dw) b_dv = b_dvbg * b_gv_exp * b_b[:, None] b_db += tl.sum(b_dvbg * b_v * b_gv_exp, 1) b_dgv = b_dvbg * b_vbg - p_dgv = tl.make_block_ptr(dgv + (bos*H + i_h) * V, (T, V), (H*V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) - tl.store(p_dgv, b_dgv.to(p_dgv.dtype.element_ty), boundary_check=(0, 1)) - tl.store(p_dv, b_dv.to(p_dv.dtype.element_ty), boundary_check=(0, 1)) + p_dgv = dgv + (bos*H + i_h) * V + o_t[:, None] * (H*V) + o_v[None, :] + tl.store(p_dgv, b_dgv.to(p_dgv.dtype.element_ty), mask=m_tv) + tl.store(p_dv, b_dv.to(p_dv.dtype.element_ty), mask=m_tv) for i_k in range(tl.cdiv(K, BK)): - p_k = tl.make_block_ptr(k + (bos*H + i_h) * K, (T, K), (H*K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) - p_dk = tl.make_block_ptr(dk + (bos*H + i_h) * K, (T, K), (H*K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) - p_du = tl.make_block_ptr(du + (bos*H + i_h) * K, (T, K), (H*K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) + o_k = i_k * BK + tl.arange(0, BK) + m_tk = m_t[:, None] & (o_k[None, :] < K) + p_k = k + (bos*H + i_h) * K + o_t[:, None] * (H*K) + o_k[None, :] + p_dk = dk + (bos*H + i_h) * K + o_t[:, None] * (H*K) + o_k[None, :] + p_du = du + (bos*H + i_h) * K + o_t[:, None] * (H*K) + o_k[None, :] # [BT, BK] - b_k = tl.load(p_k, boundary_check=(0, 1)) + b_k = tl.load(p_k, mask=m_tk, other=0.0) b_kb = (b_k * b_b[:, None]).to(b_k.dtype) # BT BK - b_du = tl.load(p_du, boundary_check=(0, 1)) # BT BK - b_dA += tl.dot(b_du, tl.trans(b_kb)) # BT BT + b_du = tl.load(p_du, mask=m_tk, other=0.0) # BT BK + b_dA = tl.dot(b_du, tl.trans(b_kb), b_dA) # BT BT b_dkb = tl.dot(b_A, b_du) # BT BK b_dk = b_dkb * b_b[:, None] b_db += tl.sum(b_dkb * b_k, 1) - tl.store(p_dk, b_dk.to(p_dk.dtype.element_ty), boundary_check=(0, 1)) + tl.store(p_dk, b_dk.to(p_dk.dtype.element_ty), mask=m_tk) - o_t = i_t * BT + tl.arange(0, BT) - m_t = o_t < T m_A = (o_t[:, None] > o_t[None, :]) & (m_t[:, None] & m_t) b_dA = tl.where(m_A, b_dA, 0) b_dA = tl.dot(b_dA.to(b_A.dtype), b_A) @@ -192,9 +205,10 @@ def prepare_wy_repr_bwd_kernel( b_dA = tl.where(m_A, -b_dA, 0) # if USE_GV: - p_dA = tl.make_block_ptr(dA + (bos*H + i_h) * BT, (T, BT), (H*BT, 1), (i_t * BT, 0), (BT, BT), (1, 0)) - tl.store(p_dA, b_dA.to(p_dA.dtype.element_ty), boundary_check=(0, 1)) - tl.store(p_db, b_db.to(p_db.dtype.element_ty), boundary_check=(0,)) + m_AT2 = m_t[:, None] & (o_A[None, :] < BT) + p_dA = dA + (bos*H + i_h) * BT + o_t[:, None] * (H*BT) + o_A[None, :] + tl.store(p_dA, b_dA.to(p_dA.dtype.element_ty), mask=m_AT2) + tl.store(p_db, b_db.to(p_db.dtype.element_ty), mask=m_t) def recompute_w_u_fwd( diff --git a/build/torch-xpu/ops/dsa/__init__.py b/build/torch-cuda/ops/gdn2/backends/__init__.py similarity index 83% rename from build/torch-xpu/ops/dsa/__init__.py rename to build/torch-cuda/ops/gdn2/backends/__init__.py index 9edd748fdda6f4f92a7fd2d1f2dc3208d8a23c0e..5ca87f78347c17efa366da83eceb142645ccb2f5 100644 --- a/build/torch-xpu/ops/dsa/__init__.py +++ b/build/torch-cuda/ops/gdn2/backends/__init__.py @@ -5,8 +5,7 @@ # For a list of all contributors, visit: # https://github.com/fla-org/flash-linear-attention/graphs/contributors -from .naive import naive_dsa +"""GDN-2 backends.""" -__all__ = [ - 'naive_dsa', -] + +__all__: list[str] = [] diff --git a/build/torch-cuda/ops/gdn2/chunk.py b/build/torch-cuda/ops/gdn2/chunk.py index 2d97a9333c9e74c66ed2273e8aec341ff7f752c3..edad7d61fcc906e52ac2402b82f6656417b9d498 100644 --- a/build/torch-cuda/ops/gdn2/chunk.py +++ b/build/torch-cuda/ops/gdn2/chunk.py @@ -26,6 +26,7 @@ GDN-2-specific Triton kernels live in ``chunk_intra.py``, from __future__ import annotations import warnings +from typing import TYPE_CHECKING import torch @@ -35,6 +36,9 @@ from ...ops.gdn2.chunk_fwd import chunk_gdn2_fwd from ...ops.utils import prepare_chunk_indices from ...utils import autocast_custom_bwd, autocast_custom_fwd, input_guard +if TYPE_CHECKING: + from ...ops.cp import FLACPContext + class ChunkGDN2Function(torch.autograd.Function): """Autograd-compatible wrapper around GDN-2 forward / backward.""" @@ -65,6 +69,7 @@ class ChunkGDN2Function(torch.autograd.Function): disable_recompute: bool, return_intermediate_states: bool, state_v_first: bool, + cp_context: FLACPContext | None = None, ): q_rstd, k_rstd = None, None if use_qk_l2norm_in_kernel: @@ -101,6 +106,7 @@ class ChunkGDN2Function(torch.autograd.Function): disable_recompute=disable_recompute, return_intermediate_states=return_intermediate_states, state_v_first=state_v_first, + cp_context=cp_context, ) if return_intermediate_states: @@ -125,6 +131,7 @@ class ChunkGDN2Function(torch.autograd.Function): ctx.use_gate_in_kernel = use_gate_in_kernel ctx.disable_recompute = disable_recompute ctx.state_v_first = state_v_first + ctx.cp_context = cp_context return o.type_as(q), final_state @staticmethod @@ -160,6 +167,7 @@ class ChunkGDN2Function(torch.autograd.Function): state_v_first=ctx.state_v_first, w_wy=w_wy, u_wy=u_wy, qg=qg, kg=kg, v_new=v_new, h=h, disable_recompute=ctx.disable_recompute, + cp_context=ctx.cp_context, ) if ctx.use_qk_l2norm_in_kernel: @@ -188,6 +196,7 @@ class ChunkGDN2Function(torch.autograd.Function): None, # disable_recompute None, # return_intermediate_states None, # state_v_first + None, # cp_context ) @@ -211,6 +220,7 @@ def chunk_gdn2( disable_recompute: bool = False, return_intermediate_states: bool = False, state_v_first: bool = False, + cp_context: FLACPContext | None = None, **kwargs, ): r""" @@ -224,52 +234,64 @@ def chunk_gdn2( Setting ``b = w = beta`` (scalar) recovers KDA. Args: - q: queries of shape ``[B, T, H, K]``. - k: keys of shape ``[B, T, H, K]``. - v: values of shape ``[B, T, H, V]``. - g: (forget) log-decay of shape ``[B, T, H, K]``. With - ``use_gate_in_kernel=True`` this is the raw pre-activation and the - kernel computes ``-exp(A_log) * softplus(g + dt_bias)`` (or the - bounded variant if ``lower_bound`` is set). - b: channel-wise ERASE gate of shape ``[B, T, H, K]``. Replaces KDA's - scalar beta. Typical range: ``[0, 2]``. - w: channel-wise WRITE gate of shape ``[B, T, H, V]``. New for GDN-2. - Typical range: ``[0, 1]``. - scale: attention scale. Defaults to ``1 / sqrt(K)``. - initial_state: optional ``[N, H, K, V]`` initial state in float32 (or - ``[N, H, V, K]`` if ``state_v_first=True``). - output_final_state: whether to output the final recurrent state. - use_qk_l2norm_in_kernel: L2-normalize q and k inside the kernel. - use_gate_in_kernel: compute the gate activation inside the kernel from - raw ``g`` and (required) ``A_log`` (plus optional ``dt_bias``, - ``lower_bound``). - cu_seqlens: ``[N+1]`` packed-sequence offsets. Requires batch size 1. - cu_seqlens_cpu: optional CPU mirror of ``cu_seqlens``, forwarded to the - state-recurrence kernel. - safe_gate: use the safe-gate intra kernel variant (M=16 TensorCore - path; requires gate values in ``[-5, 0)`` if combined with - ``use_gate_in_kernel=True``). - lower_bound: when ``safe_gate=True`` and ``use_gate_in_kernel=True``, - use the bounded gate activation - ``lower_bound * sigmoid(exp(A_log) * g)``. Must lie in ``[-5, 0)``. - disable_recompute: retain forward intermediates for a faster backward - at the cost of memory. Default: ``False``. - return_intermediate_states: when ``True``, also returns the per-chunk - pre-states ``h`` (shape ``[B, NT, H, K, V]``). Must be used inside - ``torch.inference_mode()``. - state_v_first: store the recurrent state in ``[V, K]`` layout instead - of the default ``[K, V]``. + q (torch.Tensor): + queries of shape ``[B, T, H, K]``. + k (torch.Tensor): + keys of shape ``[B, T, H, K]``. + v (torch.Tensor): + values of shape ``[B, T, H, V]``. + g (torch.Tensor): + (forget) log-decay of shape ``[B, T, H, K]``. + With ``use_gate_in_kernel=True`` this is the raw pre-activation, + and the kernel computes ``-exp(A_log) * softplus(g + dt_bias)`` (or the bounded variant if ``lower_bound`` is set). + b (torch.Tensor): + channel-wise ERASE gate of shape ``[B, T, H, K]``. Replaces KDA's scalar beta. Typical range: ``[0, 2]``. + w (torch.Tensor): + channel-wise WRITE gate of shape ``[B, T, H, V]``. New for GDN-2. Typical range: ``[0, 1]``. + scale (Optional[float]): + attention scale. Default: ``1 / sqrt(K)``. + initial_state (Optional[torch.Tensor]): + initial state in float32, of shape ``[N, H, K, V]`` (or ``[N, H, V, K]`` if ``state_v_first=True``). + output_final_state (Optional[bool]): + whether to output the final recurrent state. Default: `False`. + use_qk_l2norm_in_kernel (Optional[bool]): + L2-normalize q and k inside the kernel. Default: `False`. + use_gate_in_kernel (Optional[bool]): + compute the gate activation inside the kernel from raw ``g`` and ``A_log`` + (plus optional ``dt_bias``, ``lower_bound``). + When ``lower_bound`` is set, ``A_log`` may be ``None``, + in which case the gate is ``lower_bound * sigmoid(g + dt_bias)``. Default: `False`. + cu_seqlens (Optional[torch.LongTensor]): + ``[N+1]`` packed-sequence offsets. Requires batch size 1. + cu_seqlens_cpu (Optional[torch.LongTensor]): + CPU mirror of ``cu_seqlens``, forwarded to the state-recurrence kernel. + safe_gate (Optional[bool]): + use the safe-gate intra kernel variant (M=16 TensorCore path; + requires gate values in ``[-5, 0)`` if combined with ``use_gate_in_kernel=True``). Default: `False`. + lower_bound (Optional[float]): + when ``safe_gate=True`` and ``use_gate_in_kernel=True``, + use the bounded gate activation ``lower_bound * sigmoid(exp(A_log) * g)``. Must lie in ``[-5, 0)``. + disable_recompute (Optional[bool]): + retain forward intermediates for a faster backward at the cost of memory. Default: `False`. + return_intermediate_states (Optional[bool]): + when ``True``, also returns the per-chunk pre-states ``h`` (shape ``[B, NT, H, K, V]``). + Must be used inside ``torch.inference_mode()``. Default: `False`. + state_v_first (Optional[bool]): + store the recurrent state in ``[V, K]`` layout instead of the default ``[K, V]``. Default: `False`. + cp_context (Optional[FLACPContext]): + context-parallel context for distributed training across multiple devices. + When provided, ``initial_state`` and ``output_final_state`` are not supported, + and ``cu_seqlens`` / ``cu_seqlens_cpu`` are overridden by the context. Default: `None`. Returns: - Normal mode: ``(o, final_state)``. - - Intermediate mode (``return_intermediate_states=True``): - ``(o, final_state, h)``. + - Intermediate mode (``return_intermediate_states=True``): ``(o, final_state, h)``. Examples:: >>> import torch >>> import torch.nn.functional as F - >>> from ...ops.gdn2 import chunk_gdn2 + >>> from fla.ops.gdn2 import chunk_gdn2 >>> B, T, H, K, V = 2, 1024, 4, 128, 128 >>> q = torch.randn(B, T, H, K, dtype=torch.bfloat16, device='cuda') >>> k = torch.randn(B, T, H, K, dtype=torch.bfloat16, device='cuda') @@ -289,6 +311,15 @@ def chunk_gdn2( ) state_v_first = kwargs.pop('transpose_state_layout') + 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" + assert cp_context.cu_seqlens is not None, "cu_seqlens is required for CP" + # Override cu_seqlens and cu_seqlens_cpu with the ones from the context + cu_seqlens = cp_context.cu_seqlens + if cp_context.cu_seqlens_cpu is not None: + cu_seqlens_cpu = cp_context.cu_seqlens_cpu + if cu_seqlens is not None: if q.shape[0] != 1: raise ValueError( @@ -305,8 +336,9 @@ def chunk_gdn2( A_log, dt_bias = None, None if use_gate_in_kernel: - assert "A_log" in kwargs, "A_log must be provided when use_gate_in_kernel=True." - A_log, dt_bias = kwargs["A_log"], kwargs.get("dt_bias") + A_log, dt_bias = kwargs.get("A_log"), kwargs.get("dt_bias") + if A_log is None and lower_bound is None: + raise ValueError("`A_log` must be provided when `use_gate_in_kernel=True` and `lower_bound` is not set.") chunk_size = kwargs.pop("chunk_size", 64) if chunk_size != 64: @@ -360,4 +392,5 @@ def chunk_gdn2( disable_recompute, return_intermediate_states, state_v_first, + cp_context, ) diff --git a/build/torch-cuda/ops/gdn2/chunk_bwd.py b/build/torch-cuda/ops/gdn2/chunk_bwd.py index a75b316ad2e61167ae7316a851528120efcb732c..64ab90b29c59cb7c2d8b3b463c446202e828f4b7 100644 --- a/build/torch-cuda/ops/gdn2/chunk_bwd.py +++ b/build/torch-cuda/ops/gdn2/chunk_bwd.py @@ -30,11 +30,15 @@ wise, V-dim) at the same time. from __future__ import annotations +from typing import TYPE_CHECKING + import torch import triton import triton.language as tl +from ...backends import dispatch from ...ops.common.chunk_delta_h import chunk_gated_delta_rule_bwd_dhu, chunk_gated_delta_rule_fwd_h +from ...ops.cp.chunk_delta_h import chunk_gated_delta_rule_bwd_dhu_pre_process, expand_h0 from ...ops.gdn2.chunk_intra import chunk_gdn2_bwd_intra from ...ops.gdn2.wy_fast import recompute_w_u_fwd_gdn2 from ...ops.kda.chunk_bwd import chunk_kda_bwd_dAv @@ -45,6 +49,9 @@ from ...ops.utils.constant import RCP_LN2 from ...ops.utils.op import exp2 from ...utils import IS_NVIDIA_HOPPER, autotune_cache_kwargs +if TYPE_CHECKING: + from ...ops.cp import FLACPContext + NUM_WARPS_WY = [2, 4] if IS_NVIDIA_HOPPER else [2, 4, 8] @@ -60,7 +67,9 @@ NUM_WARPS_WY = [2, 4] if IS_NVIDIA_HOPPER else [2, 4, 8] for num_stages in [2, 3, 4] if not (IS_NVIDIA_HOPPER and BK == 32 and num_warps == 4) ], - key=['BT', 'STATE_V_FIRST'], + # NOTE: H is excluded on purpose -- it only scales the grid (NT, B * H) and + # strides addressing, so it does not shape the work the BK/BV sweep tunes. + key=['BT', 'K', 'V', 'STATE_V_FIRST'], **autotune_cache_kwargs, ) @triton.jit(do_not_specialize=['T']) @@ -97,12 +106,12 @@ def chunk_gdn2_bwd_kernel_wy_dqkg_fused( STATE_V_FIRST: tl.constexpr, IS_VARLEN: tl.constexpr, ): - i_t, i_bh = tl.program_id(0), tl.program_id(1) + i_t, i_bh = tl.program_id(0).to(tl.int64), tl.program_id(1).to(tl.int64) i_b, i_h = i_bh // H, i_bh % H if IS_VARLEN: - i_tg = i_t.to(tl.int64) - 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) + 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.int64) bos, eos = tl.load(cu_seqlens + i_n).to(tl.int64), tl.load(cu_seqlens + i_n + 1).to(tl.int64) T = (eos - bos).to(tl.int32) NT = tl.cdiv(T, BT) @@ -112,8 +121,11 @@ def chunk_gdn2_bwd_kernel_wy_dqkg_fused( bos, eos = (i_b * T).to(tl.int64), (i_b * T + T).to(tl.int64) o_t = i_t * BT + tl.arange(0, BT) + o_A = tl.arange(0, BT) m_t = o_t < T m_last = (o_t == min(T, i_t * BT + BT) - 1) + m_AT = (o_A[:, None] < BT) & m_t[None, :] + m_dA = m_t[:, None] & (o_A[None, :] < BT) q += (bos * H + i_h) * K k += (bos * H + i_h) * K @@ -135,8 +147,8 @@ def chunk_gdn2_bwd_kernel_wy_dqkg_fused( dw += (bos * H + i_h) * V dA += (bos * H + i_h) * BT - p_A = tl.make_block_ptr(A, (BT, T), (1, H * BT), (0, i_t * BT), (BT, BT), (0, 1)) - b_A = tl.load(p_A, boundary_check=(0, 1)) + p_A = A + o_A[:, None] + o_t[None, :] * (H * BT) + b_A = tl.load(p_A, mask=m_AT, other=0.0) b_dA = tl.zeros([BT, BT], dtype=tl.float32) @@ -144,12 +156,13 @@ def chunk_gdn2_bwd_kernel_wy_dqkg_fused( o_k = i_k * BK + tl.arange(0, BK) m_k = o_k < K - p_k = tl.make_block_ptr(k, (T, K), (H * K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) - p_g = tl.make_block_ptr(g, (T, K), (H * K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) - p_b = tl.make_block_ptr(b, (T, K), (H * K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) - b_k = tl.load(p_k, boundary_check=(0, 1)) - b_g = tl.load(p_g, boundary_check=(0, 1)).to(tl.float32) - b_b = tl.load(p_b, boundary_check=(0, 1)) + m_kk = m_t[:, None] & m_k[None, :] + p_k = k + o_t[:, None] * (H * K) + o_k[None, :] + p_g = g + o_t[:, None] * (H * K) + o_k[None, :] + p_b = b + o_t[:, None] * (H * K) + o_k[None, :] + b_k = tl.load(p_k, mask=m_kk, other=0.0) + b_g = tl.load(p_g, mask=m_kk, other=0.0).to(tl.float32) + b_b = tl.load(p_b, mask=m_kk, other=0.0) p_gn = g + (min(T, i_t * BT + BT) - 1).to(tl.int64) * H * K + o_k b_gn = tl.load(p_gn, mask=m_k, other=0).to(tl.float32) @@ -160,44 +173,47 @@ def chunk_gdn2_bwd_kernel_wy_dqkg_fused( b_dgk = tl.zeros([BK], dtype=tl.float32) for i_v in range(tl.cdiv(V, BV)): - p_v_new = tl.make_block_ptr(v_new, (T, V), (H * V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) - p_do = tl.make_block_ptr(do, (T, V), (H * V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) + o_v = i_v * BV + tl.arange(0, BV) + m_vv = m_t[:, None] & (o_v[None, :] < V) + m_hv = (o_v[:, None] < V) & m_k[None, :] + p_v_new = v_new + o_t[:, None] * (H * V) + o_v[None, :] + p_do = do + o_t[:, None] * (H * V) + o_v[None, :] if STATE_V_FIRST: - p_h = tl.make_block_ptr(h, (V, K), (K, 1), (i_v * BV, i_k * BK), (BV, BK), (1, 0)) - p_dh = tl.make_block_ptr(dh, (V, K), (K, 1), (i_v * BV, i_k * BK), (BV, BK), (1, 0)) + p_h = h + o_v[:, None] * K + o_k[None, :] + p_dh = dh + o_v[:, None] * K + o_k[None, :] else: - p_h = tl.make_block_ptr(h, (V, K), (1, V), (i_v * BV, i_k * BK), (BV, BK), (0, 1)) - p_dh = tl.make_block_ptr(dh, (V, K), (1, V), (i_v * BV, i_k * BK), (BV, BK), (0, 1)) - p_dv = tl.make_block_ptr(dv, (T, V), (H * V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) - b_v_new = tl.load(p_v_new, boundary_check=(0, 1)) - b_do = tl.load(p_do, boundary_check=(0, 1)) - b_h = tl.load(p_h, boundary_check=(0, 1)) - b_dh = tl.load(p_dh, boundary_check=(0, 1)) - b_dv = tl.load(p_dv, boundary_check=(0, 1)) + p_h = h + o_v[:, None] + o_k[None, :] * V + p_dh = dh + o_v[:, None] + o_k[None, :] * V + p_dv = dv + o_t[:, None] * (H * V) + o_v[None, :] + b_v_new = tl.load(p_v_new, mask=m_vv, other=0.0) + b_do = tl.load(p_do, mask=m_vv, other=0.0) + b_h = tl.load(p_h, mask=m_hv, other=0.0) + b_dh = tl.load(p_dh, mask=m_hv, other=0.0) + b_dv = tl.load(p_dv, mask=m_vv, other=0.0) b_dgk += tl.sum(b_h * b_dh, axis=0) - b_dq += tl.dot(b_do, b_h.to(b_do.dtype)) - b_dk += tl.dot(b_v_new, b_dh.to(b_v_new.dtype)) - b_dw_flow += tl.dot(b_dv.to(b_v_new.dtype), b_h.to(b_v_new.dtype)) + b_dq = tl.dot(b_do, b_h.to(b_do.dtype), b_dq) + b_dk = tl.dot(b_v_new, b_dh.to(b_v_new.dtype), b_dk) + b_dw_flow = tl.dot(b_dv.to(b_v_new.dtype), b_h.to(b_v_new.dtype), b_dw_flow) tl.debug_barrier() if i_k == 0: - p_v = tl.make_block_ptr(v, (T, V), (H * V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) - p_dv2 = tl.make_block_ptr(dv2, (T, V), (H * V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) - p_wg = tl.make_block_ptr(w_gate, (T, V), (H * V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) - p_dw_gate = tl.make_block_ptr(dw, (T, V), (H * V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) + p_v = v + o_t[:, None] * (H * V) + o_v[None, :] + p_dv2 = dv2 + o_t[:, None] * (H * V) + o_v[None, :] + p_wg = w_gate + o_t[:, None] * (H * V) + o_v[None, :] + p_dw_gate = dw + o_t[:, None] * (H * V) + o_v[None, :] - b_v = tl.load(p_v, boundary_check=(0, 1)) - b_wg = tl.load(p_wg, boundary_check=(0, 1)) + b_v = tl.load(p_v, mask=m_vv, other=0.0) + b_wg = tl.load(p_wg, mask=m_vv, other=0.0) # dA gets (w_gate * v) on the value side - the GDN-2 channel-wise twist. - b_dA += tl.dot(b_dv, tl.trans(b_v * b_wg)) + b_dA = tl.dot(b_dv, tl.trans(b_v * b_wg), b_dA) b_dvb = tl.dot(b_A, b_dv) b_dv2 = b_dvb * b_wg b_dw_gate = b_dvb * b_v - tl.store(p_dv2, b_dv2.to(p_dv2.dtype.element_ty), boundary_check=(0, 1)) - tl.store(p_dw_gate, b_dw_gate.to(p_dw_gate.dtype.element_ty), boundary_check=(0, 1)) + tl.store(p_dv2, b_dv2.to(p_dv2.dtype.element_ty), mask=m_vv) + tl.store(p_dw_gate, b_dw_gate.to(p_dw_gate.dtype.element_ty), mask=m_vv) b_gk_exp = exp2(b_g) b_gb = b_gk_exp * b_b @@ -209,36 +225,37 @@ def chunk_gdn2_bwd_kernel_wy_dqkg_fused( b_dw_flow = -b_dw_flow.to(b_A.dtype) # dA gets (b * exp(gk) * k) on the key side - the GDN-2 channel-wise twist. - b_dA += tl.dot(b_dw_flow, tl.trans((b_kg * b_b).to(b_A.dtype))) + b_dA = tl.dot(b_dw_flow, tl.trans((b_kg * b_b).to(b_A.dtype)), b_dA) b_dkgb = tl.dot(b_A, b_dw_flow) - p_db = tl.make_block_ptr(db, (T, K), (H * K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) + p_db = db + o_t[:, None] * (H * K) + o_k[None, :] b_db_partial = b_dkgb * b_kg - tl.store(p_db, b_db_partial.to(p_db.dtype.element_ty), boundary_check=(0, 1)) + tl.store(p_db, b_db_partial.to(p_db.dtype.element_ty), mask=m_kk) - p_q = tl.make_block_ptr(q, (T, K), (H * K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) - b_q = tl.load(p_q, boundary_check=(0, 1)) + p_q = q + o_t[:, None] * (H * K) + o_k[None, :] + b_q = tl.load(p_q, mask=m_kk, other=0.0) b_kdk = b_k * b_dk b_dgk += tl.sum(b_kdk, axis=0) b_dg = b_q * b_dq - b_kdk + m_last[:, None] * b_dgk + b_kg * b_dkgb * b_b b_dk = b_dk + b_dkgb * b_gb - p_dq = tl.make_block_ptr(dq, (T, K), (H * K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) - p_dk = tl.make_block_ptr(dk, (T, K), (H * K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) - p_dg = tl.make_block_ptr(dg, (T, K), (H * K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) - tl.store(p_dq, b_dq.to(p_dq.dtype.element_ty), boundary_check=(0, 1)) - tl.store(p_dk, b_dk.to(p_dk.dtype.element_ty), boundary_check=(0, 1)) - tl.store(p_dg, b_dg.to(p_dg.dtype.element_ty), boundary_check=(0, 1)) + p_dq = dq + o_t[:, None] * (H * K) + o_k[None, :] + p_dk = dk + o_t[:, None] * (H * K) + o_k[None, :] + p_dg = dg + o_t[:, None] * (H * K) + o_k[None, :] + tl.store(p_dq, b_dq.to(p_dq.dtype.element_ty), mask=m_kk) + tl.store(p_dk, b_dk.to(p_dk.dtype.element_ty), mask=m_kk) + tl.store(p_dg, b_dg.to(p_dg.dtype.element_ty), mask=m_kk) m_A = (o_t[:, None] > o_t[None, :]) & (m_t[:, None] & m_t) b_dA = tl.where(m_A, b_dA, 0) b_dA = tl.dot(b_dA.to(b_A.dtype), b_A) b_dA = tl.dot(b_A, b_dA.to(b_A.dtype)) b_dA = tl.where(m_A, -b_dA, 0) - p_dA = tl.make_block_ptr(dA, (T, BT), (H * BT, 1), (i_t * BT, 0), (BT, BT), (1, 0)) - tl.store(p_dA, b_dA.to(p_dA.dtype.element_ty), boundary_check=(0, 1)) + p_dA = dA + o_t[:, None] * (H * BT) + o_A[None, :] + tl.store(p_dA, b_dA.to(p_dA.dtype.element_ty), mask=m_dA) +@dispatch('gdn2') def chunk_gdn2_bwd_wy_dqkg_fused( q: torch.Tensor, k: torch.Tensor, @@ -339,6 +356,7 @@ def chunk_gdn2_bwd( v_new: torch.Tensor | None = None, h: torch.Tensor | None = None, disable_recompute: bool = False, + cp_context: FLACPContext | None = None, ): """End-to-end GDN-2 backward. @@ -346,6 +364,12 @@ def chunk_gdn2_bwd( shape [B, T, H, K] (channel-wise erase gate); ``dw`` has shape [B, T, H, V] (channel-wise write gate). """ + if cp_context is not None: + # If CP was used, the initial_state is guaranteed to have been populated in the fwd + # Restore the full initial_state tensor from the compressed version. + # Only the first sequence's state is non-zero as it's the only one that could be cross-rank. + initial_state = expand_h0(initial_state, context=cp_context) + if not disable_recompute: if use_gate_in_kernel: g = kda_gate_chunk_cumsum( @@ -394,6 +418,25 @@ def chunk_gdn2_bwd( chunk_indices=chunk_indices, ) + if cp_context is not None: + # initial_state is None in the CP mode + # We only need to compute dht of current rank and pass it to the backward kernel + dht, initial_state = chunk_gated_delta_rule_bwd_dhu_pre_process( + q=qg, + k=kg, + w=w_wy, + do=do, + dv=dv, + gk=g, + scale=scale, + cu_seqlens=cu_seqlens, + dht=dht, + initial_state=initial_state, + context=cp_context, + chunk_size=chunk_size, + state_v_first=state_v_first, + ) + dh, dh0, dv = chunk_gated_delta_rule_bwd_dhu( q=qg, k=kg, diff --git a/build/torch-cuda/ops/gdn2/chunk_fwd.py b/build/torch-cuda/ops/gdn2/chunk_fwd.py index 0b913f748f26a7eec931c82f6cfefaebe4877ced..a611dff64fc904774cae802ac76612197683977b 100644 --- a/build/torch-cuda/ops/gdn2/chunk_fwd.py +++ b/build/torch-cuda/ops/gdn2/chunk_fwd.py @@ -7,15 +7,23 @@ # Forward orchestration for GDN-2 chunkwise training. +from __future__ import annotations + +from typing import TYPE_CHECKING + import torch from ...ops.common.chunk_delta_h import chunk_gated_delta_rule_fwd_h +from ...ops.cp.chunk_delta_h import chunk_gated_delta_rule_fwd_h_pre_process, compress_h0 from ...ops.gdn2.chunk_intra import chunk_gdn2_fwd_intra from ...ops.gla.chunk import chunk_gla_fwd_o_gk from ...ops.kda.gate import kda_gate_chunk_cumsum from ...ops.utils import chunk_local_cumsum from ...ops.utils.constant import RCP_LN2 +if TYPE_CHECKING: + from ...ops.cp import FLACPContext + def chunk_gdn2_fwd( q: torch.Tensor, @@ -39,6 +47,7 @@ def chunk_gdn2_fwd( disable_recompute: bool = False, return_intermediate_states: bool = False, state_v_first: bool = False, + cp_context: FLACPContext | None = None, ): """Top-level GDN-2 forward pipeline. @@ -89,6 +98,19 @@ def chunk_gdn2_fwd( disable_recompute=disable_recompute, ) + if cp_context is not None: + initial_state = chunk_gated_delta_rule_fwd_h_pre_process( + k=kg, + w=w_wy, + u=u_wy, + gk=g, + cu_seqlens=cu_seqlens, + initial_state=initial_state, + context=cp_context, + chunk_size=chunk_size, + state_v_first=state_v_first, + ) + h, v_new, final_state = chunk_gated_delta_rule_fwd_h( k=kg, w=w_wy, @@ -103,6 +125,13 @@ def chunk_gdn2_fwd( state_v_first=state_v_first, ) + if cp_context is not None: + # In Context Parallel (CP) mode, global initial states are not supported at the entry point. + # The `initial_state` here is computed internally via inter-rank communication. + # Since only the first sequence in the local batch can be a continuation of a cross-rank sequence, + # only the first state in the tensor is relevant. We compress it to optimize memory for `save_for_backward`. + initial_state = compress_h0(initial_state, context=cp_context) + o = chunk_gla_fwd_o_gk( q=q, v=v_new, diff --git a/build/torch-cuda/ops/gdn2/chunk_intra.py b/build/torch-cuda/ops/gdn2/chunk_intra.py index 510114d1ee4abe108e246ba530bed0a2caae8efa..c1db5c8c2118748af503f858e757d70f96442e6d 100644 --- a/build/torch-cuda/ops/gdn2/chunk_intra.py +++ b/build/torch-cuda/ops/gdn2/chunk_intra.py @@ -24,17 +24,21 @@ import torch import triton import triton.language as tl +from ...backends import dispatch from ...ops.gdn2.chunk_intra_token_parallel import chunk_gdn2_fwd_intra_token_parallel from ...ops.gdn2.wy_fast import recompute_w_u_fwd_gdn2 from ...ops.utils import prepare_chunk_indices from ...ops.utils.cache import fla_cache_autotune -from ...ops.utils.op import exp2, gather -from ...utils import IS_GATHER_SUPPORTED, IS_TF32_SUPPORTED, autotune_cache_kwargs +from ...ops.utils.constant import RCP_LN2 +from ...ops.utils.op import exp2, gather, unflatten_program_id +from ...utils import IS_GATHER_SUPPORTED, IS_NPU, IS_TF32_SUPPORTED, TRITON_ABOVE_3_6_0, autotune_cache_kwargs if IS_TF32_SUPPORTED: SOLVE_TRIL_DOT_PRECISION = tl.constexpr('tf32') else: SOLVE_TRIL_DOT_PRECISION = tl.constexpr('ieee') +# TF32 rounding makes safe-gate gradient error grow with context length; bf16x3 reduces it and runs faster than TF32. +SAFE_GATE_DOT_PRECISION = tl.constexpr('bf16x3' if TRITON_ABOVE_3_6_0 and not IS_NPU else 'ieee') @triton.heuristics({ @@ -68,13 +72,14 @@ def chunk_gdn2_fwd_kernel_intra_sub_chunk( BK: tl.constexpr, IS_VARLEN: tl.constexpr, USE_GATHER: tl.constexpr, + GATE_SHIFT: tl.constexpr, ): - i_t, i_i, i_bh = tl.program_id(0), tl.program_id(1), tl.program_id(2) + i_t, i_i, i_bh = tl.program_id(0).to(tl.int64), tl.program_id(1).to(tl.int64), tl.program_id(2).to(tl.int64) i_b, i_h = i_bh // H, i_bh % H 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.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) + i_n, i_t = tl.load(chunk_indices + i_t * 2).to(tl.int32), tl.load(chunk_indices + i_t * 2 + 1).to(tl.int64) + bos, eos = tl.load(cu_seqlens + i_n).to(tl.int64), tl.load(cu_seqlens + i_n + 1).to(tl.int64) T = eos - bos else: bos, eos = i_b * T, i_b * T + T @@ -93,22 +98,26 @@ def chunk_gdn2_fwd_kernel_intra_sub_chunk( Aqk = Aqk + (bos * H + i_h) * BT Akk = Akk + (bos * H + i_h) * BC - p_q = tl.make_block_ptr(q, (T, K), (H * K, 1), (i_ti, 0), (BC, BK), (1, 0)) - p_k = tl.make_block_ptr(k, (T, K), (H * K, 1), (i_ti, 0), (BC, BK), (1, 0)) - p_g = tl.make_block_ptr(g, (T, K), (H * K, 1), (i_ti, 0), (BC, BK), (1, 0)) - p_b = tl.make_block_ptr(b, (T, K), (H * K, 1), (i_ti, 0), (BC, BK), (1, 0)) + o_k = tl.arange(0, BK) + m_kc = m_c[:, None] & (o_k[None, :] < K) + p_q = q + o_c[:, None] * (H * K) + o_k[None, :] + p_k = k + o_c[:, None] * (H * K) + o_k[None, :] + p_g = g + o_c[:, None] * (H * K) + o_k[None, :] + p_b = b + o_c[:, None] * (H * K) + o_k[None, :] - b_q = tl.load(p_q, boundary_check=(0, 1)) - b_k = tl.load(p_k, boundary_check=(0, 1)) - b_g = tl.load(p_g, boundary_check=(0, 1)) - b_b = tl.load(p_b, boundary_check=(0, 1)) + b_q = tl.load(p_q, mask=m_kc, other=0.0) + b_k = tl.load(p_k, mask=m_kc, other=0.0) + b_g = tl.load(p_g, mask=m_kc, other=0.0) + b_b = tl.load(p_b, mask=m_kc, other=0.0) + # the first-row anchor and fixed shift cannot depend on future gates or the sequence length if USE_GATHER: - b_gn = gather(b_g, tl.full([1, BK], min(BC // 2, T - i_ti - 1), dtype=tl.int16), axis=0) + b_gn = gather(b_g, tl.zeros([1, BK], dtype=tl.int16), axis=0) else: - p_gn = g + (i_ti + min(BC // 2, T - i_ti - 1)) * H * K + tl.arange(0, BK) + p_gn = g + i_ti * H * K + tl.arange(0, BK) b_gn = tl.load(p_gn, mask=tl.arange(0, BK) < K, other=0.0) b_gn = b_gn[None, :] + b_gn = b_gn - 0.5 * (BC - 1) * GATE_SHIFT b_gm = (b_g - b_gn).to(tl.float32) b_gq = tl.where(m_c[:, None], exp2(b_gm), 0.) @@ -129,10 +138,12 @@ def chunk_gdn2_fwd_kernel_intra_sub_chunk( b_Aqk = tl.where(m_Aqk, b_Aqk, 0.0) b_Akk = tl.where(m_Akk, b_Akk, 0.0) - p_Aqk = tl.make_block_ptr(Aqk, (T, BT), (H * BT, 1), (i_ti, i_i * BC), (BC, BC), (1, 0)) - p_Akk = tl.make_block_ptr(Akk, (T, BC), (H * BC, 1), (i_ti, 0), (BC, BC), (1, 0)) - tl.store(p_Aqk, b_Aqk.to(Aqk.dtype.element_ty), boundary_check=(0, 1)) - tl.store(p_Akk, b_Akk.to(Akk.dtype.element_ty), boundary_check=(0, 1)) + m_Aqk = m_c[:, None] & ((i_i * BC + o_i)[None, :] < BT) + m_Akk = m_c[:, None] & (o_i[None, :] < BC) + p_Aqk = Aqk + o_c[:, None] * (H * BT) + (i_i * BC + o_i)[None, :] + p_Akk = Akk + o_c[:, None] * (H * BC) + o_i[None, :] + tl.store(p_Aqk, b_Aqk.to(Aqk.dtype.element_ty), mask=m_Aqk) + tl.store(p_Akk, b_Akk.to(Akk.dtype.element_ty), mask=m_Akk) tl.debug_barrier() @@ -143,7 +154,7 @@ def chunk_gdn2_fwd_kernel_intra_sub_chunk( b_a += tl.sum(b_a[:, None] * b_Ai, 0) b_Ai = tl.where((o_i == i)[:, None], b_a, b_Ai) b_Ai += m_I - tl.store(p_Akk, b_Ai.to(Akk.dtype.element_ty), boundary_check=(0, 1)) + tl.store(p_Akk, b_Ai.to(Akk.dtype.element_ty), mask=m_Akk) @triton.heuristics({ @@ -179,12 +190,12 @@ def chunk_gdn2_fwd_kernel_inter_solve_fused( IS_VARLEN: tl.constexpr, USE_SAFE_GATE: tl.constexpr, ): - i_t, i_bh = tl.program_id(0), tl.program_id(1) + i_t, i_bh = tl.program_id(0).to(tl.int64), tl.program_id(1).to(tl.int64) i_b, i_h = i_bh // H, i_bh % H 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.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) + i_n, i_t = tl.load(chunk_indices + i_t * 2).to(tl.int32), tl.load(chunk_indices + i_t * 2 + 1).to(tl.int64) + bos, eos = tl.load(cu_seqlens + i_n).to(tl.int64), tl.load(cu_seqlens + i_n + 1).to(tl.int64) T = eos - bos else: bos, eos = i_b * T, i_b * T + T @@ -206,6 +217,7 @@ def chunk_gdn2_fwd_kernel_inter_solve_fused( Akkd += (bos * H + i_h) * BC o_i = tl.arange(0, BC) + m_tc0 = (i_tc0 + o_i) < T m_tc1 = (i_tc1 + o_i) < T m_tc2 = (i_tc2 + o_i) < T m_tc3 = (i_tc3 + o_i) < T @@ -229,94 +241,94 @@ def chunk_gdn2_fwd_kernel_inter_solve_fused( o_k = i_k * BK + tl.arange(0, BK) m_k = o_k < K - p_k0 = tl.make_block_ptr(k, (T, K), (H * K, 1), (i_tc0, i_k * BK), (BC, BK), (1, 0)) - p_g0 = tl.make_block_ptr(g, (T, K), (H * K, 1), (i_tc0, i_k * BK), (BC, BK), (1, 0)) - b_k0 = tl.load(p_k0, boundary_check=(0, 1)).to(tl.float32) - b_g0 = tl.load(p_g0, boundary_check=(0, 1)).to(tl.float32) + p_k0 = k + (i_tc0 + o_i)[:, None] * (H * K) + o_k[None, :] + p_g0 = g + (i_tc0 + o_i)[:, None] * (H * K) + o_k[None, :] + b_k0 = tl.load(p_k0, mask=m_tc0[:, None] & m_k[None, :], other=0.0).to(tl.float32) + b_g0 = tl.load(p_g0, mask=m_tc0[:, None] & m_k[None, :], other=0.0).to(tl.float32) if i_tc1 < T: - p_q1 = tl.make_block_ptr(q, (T, K), (H * K, 1), (i_tc1, i_k * BK), (BC, BK), (1, 0)) - p_k1 = tl.make_block_ptr(k, (T, K), (H * K, 1), (i_tc1, i_k * BK), (BC, BK), (1, 0)) - p_g1 = tl.make_block_ptr(g, (T, K), (H * K, 1), (i_tc1, i_k * BK), (BC, BK), (1, 0)) - p_b1 = tl.make_block_ptr(b, (T, K), (H * K, 1), (i_tc1, i_k * BK), (BC, BK), (1, 0)) - b_q1 = tl.load(p_q1, boundary_check=(0, 1)).to(tl.float32) - b_k1 = tl.load(p_k1, boundary_check=(0, 1)).to(tl.float32) - b_g1 = tl.load(p_g1, boundary_check=(0, 1)).to(tl.float32) - b_b1 = tl.load(p_b1, boundary_check=(0, 1)).to(tl.float32) + p_q1 = q + (i_tc1 + o_i)[:, None] * (H * K) + o_k[None, :] + p_k1 = k + (i_tc1 + o_i)[:, None] * (H * K) + o_k[None, :] + p_g1 = g + (i_tc1 + o_i)[:, None] * (H * K) + o_k[None, :] + p_b1 = b + (i_tc1 + o_i)[:, None] * (H * K) + o_k[None, :] + b_q1 = tl.load(p_q1, mask=m_tc1[:, None] & m_k[None, :], other=0.0).to(tl.float32) + b_k1 = tl.load(p_k1, mask=m_tc1[:, None] & m_k[None, :], other=0.0).to(tl.float32) + b_g1 = tl.load(p_g1, mask=m_tc1[:, None] & m_k[None, :], other=0.0).to(tl.float32) + b_b1 = tl.load(p_b1, mask=m_tc1[:, None] & m_k[None, :], other=0.0).to(tl.float32) b_gn1 = tl.load(g + i_tc1 * H * K + o_k, mask=m_k, other=0).to(tl.float32) b_gqn = tl.where(m_tc1[:, None], exp2(b_g1 - b_gn1[None, :]), 0) b_kgt = tl.trans(b_k0 * exp2(b_gn1[None, :] - b_g0)) b_bk1 = b_b1 * b_k1 - b_Aqk10 += tl.dot(b_q1 * b_gqn, b_kgt) - b_Akk10 += tl.dot(b_bk1 * b_gqn, b_kgt) + b_Aqk10 = tl.dot(b_q1 * b_gqn, b_kgt, b_Aqk10) + b_Akk10 = tl.dot(b_bk1 * b_gqn, b_kgt, b_Akk10) if i_tc2 < T: - p_q2 = tl.make_block_ptr(q, (T, K), (H * K, 1), (i_tc2, i_k * BK), (BC, BK), (1, 0)) - p_k2 = tl.make_block_ptr(k, (T, K), (H * K, 1), (i_tc2, i_k * BK), (BC, BK), (1, 0)) - p_g2 = tl.make_block_ptr(g, (T, K), (H * K, 1), (i_tc2, i_k * BK), (BC, BK), (1, 0)) - p_b2 = tl.make_block_ptr(b, (T, K), (H * K, 1), (i_tc2, i_k * BK), (BC, BK), (1, 0)) - b_q2 = tl.load(p_q2, boundary_check=(0, 1)).to(tl.float32) - b_k2 = tl.load(p_k2, boundary_check=(0, 1)).to(tl.float32) - b_g2 = tl.load(p_g2, boundary_check=(0, 1)).to(tl.float32) - b_b2 = tl.load(p_b2, boundary_check=(0, 1)).to(tl.float32) + p_q2 = q + (i_tc2 + o_i)[:, None] * (H * K) + o_k[None, :] + p_k2 = k + (i_tc2 + o_i)[:, None] * (H * K) + o_k[None, :] + p_g2 = g + (i_tc2 + o_i)[:, None] * (H * K) + o_k[None, :] + p_b2 = b + (i_tc2 + o_i)[:, None] * (H * K) + o_k[None, :] + b_q2 = tl.load(p_q2, mask=m_tc2[:, None] & m_k[None, :], other=0.0).to(tl.float32) + b_k2 = tl.load(p_k2, mask=m_tc2[:, None] & m_k[None, :], other=0.0).to(tl.float32) + b_g2 = tl.load(p_g2, mask=m_tc2[:, None] & m_k[None, :], other=0.0).to(tl.float32) + b_b2 = tl.load(p_b2, mask=m_tc2[:, None] & m_k[None, :], other=0.0).to(tl.float32) b_gn2 = tl.load(g + i_tc2 * H * K + o_k, mask=m_k, other=0).to(tl.float32) b_gqn2 = tl.where(m_tc2[:, None], exp2(b_g2 - b_gn2[None, :]), 0) b_qg2 = b_q2 * b_gqn2 b_bkg2 = (b_b2 * b_k2) * b_gqn2 b_kgt = tl.trans(b_k0 * exp2(b_gn2[None, :] - b_g0)) - b_Aqk20 += tl.dot(b_qg2, b_kgt) - b_Akk20 += tl.dot(b_bkg2, b_kgt) + b_Aqk20 = tl.dot(b_qg2, b_kgt, b_Aqk20) + b_Akk20 = tl.dot(b_bkg2, b_kgt, b_Akk20) b_kgt = tl.trans(b_k1 * exp2(b_gn2[None, :] - b_g1)) - b_Aqk21 += tl.dot(b_qg2, b_kgt) - b_Akk21 += tl.dot(b_bkg2, b_kgt) + b_Aqk21 = tl.dot(b_qg2, b_kgt, b_Aqk21) + b_Akk21 = tl.dot(b_bkg2, b_kgt, b_Akk21) if i_tc3 < T: - p_q3 = tl.make_block_ptr(q, (T, K), (H * K, 1), (i_tc3, i_k * BK), (BC, BK), (1, 0)) - p_k3 = tl.make_block_ptr(k, (T, K), (H * K, 1), (i_tc3, i_k * BK), (BC, BK), (1, 0)) - p_g3 = tl.make_block_ptr(g, (T, K), (H * K, 1), (i_tc3, i_k * BK), (BC, BK), (1, 0)) - p_b3 = tl.make_block_ptr(b, (T, K), (H * K, 1), (i_tc3, i_k * BK), (BC, BK), (1, 0)) - b_q3 = tl.load(p_q3, boundary_check=(0, 1)).to(tl.float32) - b_k3 = tl.load(p_k3, boundary_check=(0, 1)).to(tl.float32) - b_g3 = tl.load(p_g3, boundary_check=(0, 1)).to(tl.float32) - b_b3 = tl.load(p_b3, boundary_check=(0, 1)).to(tl.float32) + p_q3 = q + (i_tc3 + o_i)[:, None] * (H * K) + o_k[None, :] + p_k3 = k + (i_tc3 + o_i)[:, None] * (H * K) + o_k[None, :] + p_g3 = g + (i_tc3 + o_i)[:, None] * (H * K) + o_k[None, :] + p_b3 = b + (i_tc3 + o_i)[:, None] * (H * K) + o_k[None, :] + b_q3 = tl.load(p_q3, mask=m_tc3[:, None] & m_k[None, :], other=0.0).to(tl.float32) + b_k3 = tl.load(p_k3, mask=m_tc3[:, None] & m_k[None, :], other=0.0).to(tl.float32) + b_g3 = tl.load(p_g3, mask=m_tc3[:, None] & m_k[None, :], other=0.0).to(tl.float32) + b_b3 = tl.load(p_b3, mask=m_tc3[:, None] & m_k[None, :], other=0.0).to(tl.float32) b_gn3 = tl.load(g + i_tc3 * H * K + o_k, mask=m_k, other=0).to(tl.float32) b_gqn3 = tl.where(m_tc3[:, None], exp2(b_g3 - b_gn3[None, :]), 0) b_qg3 = b_q3 * b_gqn3 b_bkg3 = (b_b3 * b_k3) * b_gqn3 b_kgt = tl.trans(b_k0 * exp2(b_gn3[None, :] - b_g0)) - b_Aqk30 += tl.dot(b_qg3, b_kgt) - b_Akk30 += tl.dot(b_bkg3, b_kgt) + b_Aqk30 = tl.dot(b_qg3, b_kgt, b_Aqk30) + b_Akk30 = tl.dot(b_bkg3, b_kgt, b_Akk30) b_kgt = tl.trans(b_k1 * exp2(b_gn3[None, :] - b_g1)) - b_Aqk31 += tl.dot(b_qg3, b_kgt) - b_Akk31 += tl.dot(b_bkg3, b_kgt) + b_Aqk31 = tl.dot(b_qg3, b_kgt, b_Aqk31) + b_Akk31 = tl.dot(b_bkg3, b_kgt, b_Akk31) b_kgt = tl.trans(b_k2 * exp2(b_gn3[None, :] - b_g2)) - b_Aqk32 += tl.dot(b_qg3, b_kgt) - b_Akk32 += tl.dot(b_bkg3, b_kgt) + b_Aqk32 = tl.dot(b_qg3, b_kgt, b_Aqk32) + b_Akk32 = tl.dot(b_bkg3, b_kgt, b_Akk32) if i_tc1 < T: - p_Aqk10 = tl.make_block_ptr(Aqk, (T, BT), (H * BT, 1), (i_tc1, 0), (BC, BC), (1, 0)) - tl.store(p_Aqk10, (b_Aqk10 * scale).to(Aqk.dtype.element_ty), boundary_check=(0, 1)) + p_Aqk10 = Aqk + (i_tc1 + o_i)[:, None] * (H * BT) + o_i[None, :] + tl.store(p_Aqk10, (b_Aqk10 * scale).to(Aqk.dtype.element_ty), mask=m_tc1[:, None] & (o_i[None, :] < BT)) if i_tc2 < T: - p_Aqk20 = tl.make_block_ptr(Aqk, (T, BT), (H * BT, 1), (i_tc2, 0), (BC, BC), (1, 0)) - p_Aqk21 = tl.make_block_ptr(Aqk, (T, BT), (H * BT, 1), (i_tc2, BC), (BC, BC), (1, 0)) - tl.store(p_Aqk20, (b_Aqk20 * scale).to(Aqk.dtype.element_ty), boundary_check=(0, 1)) - tl.store(p_Aqk21, (b_Aqk21 * scale).to(Aqk.dtype.element_ty), boundary_check=(0, 1)) + p_Aqk20 = Aqk + (i_tc2 + o_i)[:, None] * (H * BT) + o_i[None, :] + p_Aqk21 = Aqk + (i_tc2 + o_i)[:, None] * (H * BT) + (BC + o_i)[None, :] + tl.store(p_Aqk20, (b_Aqk20 * scale).to(Aqk.dtype.element_ty), mask=m_tc2[:, None] & (o_i[None, :] < BT)) + tl.store(p_Aqk21, (b_Aqk21 * scale).to(Aqk.dtype.element_ty), mask=m_tc2[:, None] & ((BC + o_i)[None, :] < BT)) if i_tc3 < T: - p_Aqk30 = tl.make_block_ptr(Aqk, (T, BT), (H * BT, 1), (i_tc3, 0), (BC, BC), (1, 0)) - p_Aqk31 = tl.make_block_ptr(Aqk, (T, BT), (H * BT, 1), (i_tc3, BC), (BC, BC), (1, 0)) - p_Aqk32 = tl.make_block_ptr(Aqk, (T, BT), (H * BT, 1), (i_tc3, 2 * BC), (BC, BC), (1, 0)) - tl.store(p_Aqk30, (b_Aqk30 * scale).to(Aqk.dtype.element_ty), boundary_check=(0, 1)) - tl.store(p_Aqk31, (b_Aqk31 * scale).to(Aqk.dtype.element_ty), boundary_check=(0, 1)) - tl.store(p_Aqk32, (b_Aqk32 * scale).to(Aqk.dtype.element_ty), boundary_check=(0, 1)) - - p_Akk00 = tl.make_block_ptr(Akkd, (T, BC), (H * BC, 1), (i_tc0, 0), (BC, BC), (1, 0)) - p_Akk11 = tl.make_block_ptr(Akkd, (T, BC), (H * BC, 1), (i_tc1, 0), (BC, BC), (1, 0)) - p_Akk22 = tl.make_block_ptr(Akkd, (T, BC), (H * BC, 1), (i_tc2, 0), (BC, BC), (1, 0)) - p_Akk33 = tl.make_block_ptr(Akkd, (T, BC), (H * BC, 1), (i_tc3, 0), (BC, BC), (1, 0)) - b_Ai00 = tl.load(p_Akk00, boundary_check=(0, 1)).to(tl.float32) - b_Ai11 = tl.load(p_Akk11, boundary_check=(0, 1)).to(tl.float32) - b_Ai22 = tl.load(p_Akk22, boundary_check=(0, 1)).to(tl.float32) - b_Ai33 = tl.load(p_Akk33, boundary_check=(0, 1)).to(tl.float32) + p_Aqk30 = Aqk + (i_tc3 + o_i)[:, None] * (H * BT) + o_i[None, :] + p_Aqk31 = Aqk + (i_tc3 + o_i)[:, None] * (H * BT) + (BC + o_i)[None, :] + p_Aqk32 = Aqk + (i_tc3 + o_i)[:, None] * (H * BT) + (2 * BC + o_i)[None, :] + tl.store(p_Aqk30, (b_Aqk30 * scale).to(Aqk.dtype.element_ty), mask=m_tc3[:, None] & (o_i[None, :] < BT)) + tl.store(p_Aqk31, (b_Aqk31 * scale).to(Aqk.dtype.element_ty), mask=m_tc3[:, None] & ((BC + o_i)[None, :] < BT)) + tl.store(p_Aqk32, (b_Aqk32 * scale).to(Aqk.dtype.element_ty), mask=m_tc3[:, None] & ((2 * BC + o_i)[None, :] < BT)) + + p_Akk00 = Akkd + (i_tc0 + o_i)[:, None] * (H * BC) + o_i[None, :] + p_Akk11 = Akkd + (i_tc1 + o_i)[:, None] * (H * BC) + o_i[None, :] + p_Akk22 = Akkd + (i_tc2 + o_i)[:, None] * (H * BC) + o_i[None, :] + p_Akk33 = Akkd + (i_tc3 + o_i)[:, None] * (H * BC) + o_i[None, :] + b_Ai00 = tl.load(p_Akk00, mask=m_tc0[:, None] & (o_i[None, :] < BC), other=0.0).to(tl.float32) + b_Ai11 = tl.load(p_Akk11, mask=m_tc1[:, None] & (o_i[None, :] < BC), other=0.0).to(tl.float32) + b_Ai22 = tl.load(p_Akk22, mask=m_tc2[:, None] & (o_i[None, :] < BC), other=0.0).to(tl.float32) + b_Ai33 = tl.load(p_Akk33, mask=m_tc3[:, None] & (o_i[None, :] < BC), other=0.0).to(tl.float32) if not USE_SAFE_GATE: m_A = o_i[:, None] > o_i[None, :] @@ -389,27 +401,27 @@ def chunk_gdn2_fwd_kernel_inter_solve_fused( input_precision=SOLVE_TRIL_DOT_PRECISION, ) - p_Akk00 = tl.make_block_ptr(Akk, (T, BT), (H * BT, 1), (i_tc0, 0), (BC, BC), (1, 0)) - p_Akk10 = tl.make_block_ptr(Akk, (T, BT), (H * BT, 1), (i_tc1, 0), (BC, BC), (1, 0)) - p_Akk11 = tl.make_block_ptr(Akk, (T, BT), (H * BT, 1), (i_tc1, BC), (BC, BC), (1, 0)) - p_Akk20 = tl.make_block_ptr(Akk, (T, BT), (H * BT, 1), (i_tc2, 0), (BC, BC), (1, 0)) - p_Akk21 = tl.make_block_ptr(Akk, (T, BT), (H * BT, 1), (i_tc2, BC), (BC, BC), (1, 0)) - p_Akk22 = tl.make_block_ptr(Akk, (T, BT), (H * BT, 1), (i_tc2, 2 * BC), (BC, BC), (1, 0)) - p_Akk30 = tl.make_block_ptr(Akk, (T, BT), (H * BT, 1), (i_tc3, 0), (BC, BC), (1, 0)) - p_Akk31 = tl.make_block_ptr(Akk, (T, BT), (H * BT, 1), (i_tc3, BC), (BC, BC), (1, 0)) - p_Akk32 = tl.make_block_ptr(Akk, (T, BT), (H * BT, 1), (i_tc3, 2 * BC), (BC, BC), (1, 0)) - p_Akk33 = tl.make_block_ptr(Akk, (T, BT), (H * BT, 1), (i_tc3, 3 * BC), (BC, BC), (1, 0)) - - tl.store(p_Akk00, b_Ai00.to(Akk.dtype.element_ty), boundary_check=(0, 1)) - tl.store(p_Akk10, b_Ai10.to(Akk.dtype.element_ty), boundary_check=(0, 1)) - tl.store(p_Akk11, b_Ai11.to(Akk.dtype.element_ty), boundary_check=(0, 1)) - tl.store(p_Akk20, b_Ai20.to(Akk.dtype.element_ty), boundary_check=(0, 1)) - tl.store(p_Akk21, b_Ai21.to(Akk.dtype.element_ty), boundary_check=(0, 1)) - tl.store(p_Akk22, b_Ai22.to(Akk.dtype.element_ty), boundary_check=(0, 1)) - tl.store(p_Akk30, b_Ai30.to(Akk.dtype.element_ty), boundary_check=(0, 1)) - tl.store(p_Akk31, b_Ai31.to(Akk.dtype.element_ty), boundary_check=(0, 1)) - tl.store(p_Akk32, b_Ai32.to(Akk.dtype.element_ty), boundary_check=(0, 1)) - tl.store(p_Akk33, b_Ai33.to(Akk.dtype.element_ty), boundary_check=(0, 1)) + p_Akk00 = Akk + (i_tc0 + o_i)[:, None] * (H * BT) + o_i[None, :] + p_Akk10 = Akk + (i_tc1 + o_i)[:, None] * (H * BT) + o_i[None, :] + p_Akk11 = Akk + (i_tc1 + o_i)[:, None] * (H * BT) + (BC + o_i)[None, :] + p_Akk20 = Akk + (i_tc2 + o_i)[:, None] * (H * BT) + o_i[None, :] + p_Akk21 = Akk + (i_tc2 + o_i)[:, None] * (H * BT) + (BC + o_i)[None, :] + p_Akk22 = Akk + (i_tc2 + o_i)[:, None] * (H * BT) + (2 * BC + o_i)[None, :] + p_Akk30 = Akk + (i_tc3 + o_i)[:, None] * (H * BT) + o_i[None, :] + p_Akk31 = Akk + (i_tc3 + o_i)[:, None] * (H * BT) + (BC + o_i)[None, :] + p_Akk32 = Akk + (i_tc3 + o_i)[:, None] * (H * BT) + (2 * BC + o_i)[None, :] + p_Akk33 = Akk + (i_tc3 + o_i)[:, None] * (H * BT) + (3 * BC + o_i)[None, :] + + tl.store(p_Akk00, b_Ai00.to(Akk.dtype.element_ty), mask=m_tc0[:, None] & (o_i[None, :] < BT)) + tl.store(p_Akk10, b_Ai10.to(Akk.dtype.element_ty), mask=m_tc1[:, None] & (o_i[None, :] < BT)) + tl.store(p_Akk11, b_Ai11.to(Akk.dtype.element_ty), mask=m_tc1[:, None] & ((BC + o_i)[None, :] < BT)) + tl.store(p_Akk20, b_Ai20.to(Akk.dtype.element_ty), mask=m_tc2[:, None] & (o_i[None, :] < BT)) + tl.store(p_Akk21, b_Ai21.to(Akk.dtype.element_ty), mask=m_tc2[:, None] & ((BC + o_i)[None, :] < BT)) + tl.store(p_Akk22, b_Ai22.to(Akk.dtype.element_ty), mask=m_tc2[:, None] & ((2 * BC + o_i)[None, :] < BT)) + tl.store(p_Akk30, b_Ai30.to(Akk.dtype.element_ty), mask=m_tc3[:, None] & (o_i[None, :] < BT)) + tl.store(p_Akk31, b_Ai31.to(Akk.dtype.element_ty), mask=m_tc3[:, None] & ((BC + o_i)[None, :] < BT)) + tl.store(p_Akk32, b_Ai32.to(Akk.dtype.element_ty), mask=m_tc3[:, None] & ((2 * BC + o_i)[None, :] < BT)) + tl.store(p_Akk33, b_Ai33.to(Akk.dtype.element_ty), mask=m_tc3[:, None] & ((3 * BC + o_i)[None, :] < BT)) @triton.heuristics({ @@ -451,16 +463,18 @@ def chunk_gdn2_bwd_kernel_intra( NC: tl.constexpr, IS_VARLEN: tl.constexpr, SAFE_GATE: tl.constexpr, + GATE_SHIFT: tl.constexpr, USE_GATHER: tl.constexpr, ): - i_kc, i_t, i_bh = tl.program_id(0), tl.program_id(1), tl.program_id(2) + i_kc, i_t = unflatten_program_id(tl.cdiv(K, BK) * NC) + i_bh = tl.program_id(1).to(tl.int64) i_b, i_h = i_bh // H, i_bh % H i_k, i_i = i_kc // NC, i_kc % NC - all = B * T + all = B.to(tl.int64) * T 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.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) + i_n, i_t = tl.load(chunk_indices + i_t * 2).to(tl.int32), tl.load(chunk_indices + i_t * 2 + 1).to(tl.int64) + bos, eos = tl.load(cu_seqlens + i_n).to(tl.int64), tl.load(cu_seqlens + i_n + 1).to(tl.int64) else: bos, eos = i_b * T, i_b * T + T T = eos - bos @@ -487,11 +501,15 @@ def chunk_gdn2_bwd_kernel_intra( dg2 += (bos * H + i_h) * K db += (i_k * all + bos) * H * BK + i_h * BK - p_g = tl.make_block_ptr(g, (T, K), (H * K, 1), (i_ti, i_k * BK), (BC, BK), (1, 0)) - b_g = tl.load(p_g, boundary_check=(0, 1)).to(tl.float32) + o_i = tl.arange(0, BC) + o_c = i_ti + o_i + m_c = o_c < T + m_kc = m_c[:, None] & m_k[None, :] + p_g = g + o_c[:, None] * (H * K) + o_k[None, :] + b_g = tl.load(p_g, mask=m_kc, other=0.0).to(tl.float32) - p_b = tl.make_block_ptr(b, (T, K), (H * K, 1), (i_ti, i_k * BK), (BC, BK), (1, 0)) - b_b = tl.load(p_b, boundary_check=(0, 1)) + p_b = b + o_c[:, None] * (H * K) + o_k[None, :] + b_b = tl.load(p_b, mask=m_kc, other=0.0) b_dq2 = tl.zeros([BC, BK], dtype=tl.float32) b_dk2 = tl.zeros([BC, BK], dtype=tl.float32) @@ -500,43 +518,48 @@ def chunk_gdn2_bwd_kernel_intra( p_gn = g + i_ti * H * K + o_k b_gn = tl.load(p_gn, mask=m_k, other=0).to(tl.float32)[None, :] for i_j in range(0, i_i): - p_k = tl.make_block_ptr(k, (T, K), (H * K, 1), (i_t * BT + i_j * BC, i_k * BK), (BC, BK), (1, 0)) - p_gk = tl.make_block_ptr(g, (T, K), (H * K, 1), (i_t * BT + i_j * BC, i_k * BK), (BC, BK), (1, 0)) - p_dAqk = tl.make_block_ptr(dAqk, (T, BT), (H * BT, 1), (i_ti, i_j * BC), (BC, BC), (1, 0)) - p_dAkk = tl.make_block_ptr(dAkk, (T, BT), (H * BT, 1), (i_ti, i_j * BC), (BC, BC), (1, 0)) - b_kj = tl.load(p_k, boundary_check=(0, 1)) - b_gk = tl.load(p_gk, boundary_check=(0, 1)) + o_j = i_t * BT + i_j * BC + o_i + m_j = o_j < T + m_jk = m_j[:, None] & m_k[None, :] + m_dj = m_c[:, None] & ((i_j * BC + o_i)[None, :] < BT) + p_k = k + o_j[:, None] * (H * K) + o_k[None, :] + p_gk = g + o_j[:, None] * (H * K) + o_k[None, :] + p_dAqk = dAqk + o_c[:, None] * (H * BT) + (i_j * BC + o_i)[None, :] + p_dAkk = dAkk + o_c[:, None] * (H * BT) + (i_j * BC + o_i)[None, :] + b_kj = tl.load(p_k, mask=m_jk, other=0.0) + b_gk = tl.load(p_gk, mask=m_jk, other=0.0) b_kg = b_kj * exp2(b_gn - b_gk) - b_dAqk = tl.load(p_dAqk, boundary_check=(0, 1)) - b_dAkk = tl.load(p_dAkk, boundary_check=(0, 1)) - b_dq2 += tl.dot(b_dAqk, b_kg) - b_dk2 += tl.dot(b_dAkk, b_kg) + b_dAqk = tl.load(p_dAqk, mask=m_dj, other=0.0) + b_dAkk = tl.load(p_dAkk, mask=m_dj, other=0.0) + b_dq2 = tl.dot(b_dAqk, b_kg, b_dq2) + b_dk2 = tl.dot(b_dAkk, b_kg, b_dk2) b_gqn = exp2(b_g - b_gn) b_dq2 *= b_gqn b_dk2 *= b_gqn - o_i = tl.arange(0, BC) m_dA = (i_ti + o_i) < T o_dA = (i_ti + o_i) * H * BT + i_i * BC p_kj = k + i_ti * H * K + o_k p_gkj = g + i_ti * H * K + o_k - p_q = tl.make_block_ptr(q, (T, K), (H * K, 1), (i_ti, i_k * BK), (BC, BK), (1, 0)) - p_k = tl.make_block_ptr(k, (T, K), (H * K, 1), (i_ti, i_k * BK), (BC, BK), (1, 0)) - b_q = tl.load(p_q, boundary_check=(0, 1)) - b_k = tl.load(p_k, boundary_check=(0, 1)) + p_q = q + o_c[:, None] * (H * K) + o_k[None, :] + p_k = k + o_c[:, None] * (H * K) + o_k[None, :] + b_q = tl.load(p_q, mask=m_kc, other=0.0) + b_k = tl.load(p_k, mask=m_kc, other=0.0) if SAFE_GATE: if USE_GATHER: - b_gn = gather(b_g, tl.full([1, BK], min(BC // 2, T - i_ti - 1), dtype=tl.int16), axis=0) + b_gn = gather(b_g, tl.zeros([1, BK], dtype=tl.int16), axis=0) else: - p_gn = g + (i_ti + min(BC // 2, T - i_ti - 1)) * H * K + o_k + p_gn = g + i_ti * H * K + o_k b_gn = tl.load(p_gn, mask=m_k, other=0)[None, :] + b_gn = b_gn - 0.5 * (BC - 1) * GATE_SHIFT - p_dAqk = tl.make_block_ptr(dAqk, (T, BT), (H * BT, 1), (i_ti, i_i * BC), (BC, BC), (1, 0)) - p_dAkk = tl.make_block_ptr(dAkk, (T, BT), (H * BT, 1), (i_ti, i_i * BC), (BC, BC), (1, 0)) - b_dAqk_diag_qk = tl.load(p_dAqk, boundary_check=(0, 1)).to(tl.float32) - b_dAkk_diag_qk = tl.load(p_dAkk, boundary_check=(0, 1)).to(tl.float32) + m_dd = m_c[:, None] & ((i_i * BC + o_i)[None, :] < BT) + p_dAqk = dAqk + o_c[:, None] * (H * BT) + (i_i * BC + o_i)[None, :] + p_dAkk = dAkk + o_c[:, None] * (H * BT) + (i_i * BC + o_i)[None, :] + b_dAqk_diag_qk = tl.load(p_dAqk, mask=m_dd, other=0.0).to(tl.float32) + b_dAkk_diag_qk = tl.load(p_dAkk, mask=m_dd, other=0.0).to(tl.float32) m_i_diag_qk = (o_i[:, None] >= o_i[None, :]) & ((i_ti + o_i[:, None]) < T) & ((i_ti + o_i[None, :]) < T) m_j_diag_qk = (i_ti + o_i[:, None]) < T @@ -548,8 +571,8 @@ def chunk_gdn2_bwd_kernel_intra( exp_neg_b_g_diag_qk = tl.where(m_j_diag_qk, exp2(-b_g_diag_qk), 0.) b_k_exp_diag_qk = b_k * exp_neg_b_g_diag_qk - b_dq2 += tl.dot(b_dAqk_diag_qk, b_k_exp_diag_qk) * exp_b_g_diag_qk - b_dk2 += tl.dot(b_dAkk_diag_qk, b_k_exp_diag_qk) * exp_b_g_diag_qk + b_dq2 += tl.dot(b_dAqk_diag_qk, b_k_exp_diag_qk, input_precision=SAFE_GATE_DOT_PRECISION) * exp_b_g_diag_qk + b_dk2 += tl.dot(b_dAkk_diag_qk, b_k_exp_diag_qk, input_precision=SAFE_GATE_DOT_PRECISION) * exp_b_g_diag_qk else: for j in range(0, min(BC, T - i_t * BT - i_i * BC)): b_dAqk = tl.load(dAqk + o_dA + j, mask=m_dA, other=0) @@ -569,14 +592,16 @@ def chunk_gdn2_bwd_kernel_intra( b_db_tile = b_dk2 * b_k b_dk2 = b_dk2 * b_b - p_dq = tl.make_block_ptr(dq, (T, K), (H * K, 1), (i_ti, i_k * BK), (BC, BK), (1, 0)) - p_dq2 = tl.make_block_ptr(dq2, (T, K), (H * K, 1), (i_ti, i_k * BK), (BC, BK), (1, 0)) - p_db = tl.make_block_ptr(db, (T, BK), (H * BK, 1), (i_ti, 0), (BC, BK), (1, 0)) + o_db = tl.arange(0, BK) + m_db = m_c[:, None] & (o_db[None, :] < BK) + p_dq = dq + o_c[:, None] * (H * K) + o_k[None, :] + p_dq2 = dq2 + o_c[:, None] * (H * K) + o_k[None, :] + p_db = db + o_c[:, None] * (H * BK) + o_db[None, :] b_dg2 = b_q * b_dq2 - b_dq2 = b_dq2 + tl.load(p_dq, boundary_check=(0, 1)) - tl.store(p_dq2, b_dq2.to(p_dq2.dtype.element_ty), boundary_check=(0, 1)) - tl.store(p_db, b_db_tile.to(p_db.dtype.element_ty), boundary_check=(0, 1)) + b_dq2 = b_dq2 + tl.load(p_dq, mask=m_kc, other=0.0) + tl.store(p_dq2, b_dq2.to(p_dq2.dtype.element_ty), mask=m_kc) + tl.store(p_db, b_db_tile.to(p_db.dtype.element_ty), mask=m_db) tl.debug_barrier() b_dkt = tl.zeros([BC, BK], dtype=tl.float32) @@ -586,26 +611,27 @@ def chunk_gdn2_bwd_kernel_intra( p_gn = g + (min(i_ti + BC, T) - 1) * H * K + o_k b_gn = tl.load(p_gn, mask=m_k, other=0).to(tl.float32)[None, :] for i_j in range(i_i + 1, NC): - p_q = tl.make_block_ptr(q, (T, K), (H * K, 1), (i_t * BT + i_j * BC, i_k * BK), (BC, BK), (1, 0)) - p_k = tl.make_block_ptr(k, (T, K), (H * K, 1), (i_t * BT + i_j * BC, i_k * BK), (BC, BK), (1, 0)) - p_gk = tl.make_block_ptr(g, (T, K), (H * K, 1), (i_t * BT + i_j * BC, i_k * BK), (BC, BK), (1, 0)) - p_bj = tl.make_block_ptr(b, (T, K), (H * K, 1), (i_t * BT + i_j * BC, i_k * BK), (BC, BK), (1, 0)) - p_dAqk = tl.make_block_ptr(dAqk, (BT, T), (1, H * BT), (i_i * BC, i_t * BT + i_j * BC), (BC, BC), (0, 1)) - p_dAkk = tl.make_block_ptr(dAkk, (BT, T), (1, H * BT), (i_i * BC, i_t * BT + i_j * BC), (BC, BC), (0, 1)) - b_bj = tl.load(p_bj, boundary_check=(0, 1)) - b_qj = tl.load(p_q, boundary_check=(0, 1)) - b_kbj = tl.load(p_k, boundary_check=(0, 1)) * b_bj - b_gk = tl.load(p_gk, boundary_check=(0, 1)).to(tl.float32) - b_dAqk = tl.load(p_dAqk, boundary_check=(0, 1)) - b_dAkk = tl.load(p_dAkk, boundary_check=(0, 1)) - o_j = i_t * BT + i_j * BC + o_i m_j = o_j < T + m_jk = m_j[:, None] & m_k[None, :] + m_dj = ((i_i * BC + o_i)[:, None] < BT) & m_j[None, :] + p_q = q + o_j[:, None] * (H * K) + o_k[None, :] + p_k = k + o_j[:, None] * (H * K) + o_k[None, :] + p_gk = g + o_j[:, None] * (H * K) + o_k[None, :] + p_bj = b + o_j[:, None] * (H * K) + o_k[None, :] + p_dAqk = dAqk + (i_i * BC + o_i)[:, None] + o_j[None, :] * (H * BT) + p_dAkk = dAkk + (i_i * BC + o_i)[:, None] + o_j[None, :] * (H * BT) + b_bj = tl.load(p_bj, mask=m_jk, other=0.0) + b_qj = tl.load(p_q, mask=m_jk, other=0.0) + b_kbj = tl.load(p_k, mask=m_jk, other=0.0) * b_bj + b_gk = tl.load(p_gk, mask=m_jk, other=0.0).to(tl.float32) + b_dAqk = tl.load(p_dAqk, mask=m_dj, other=0.0) + b_dAkk = tl.load(p_dAkk, mask=m_dj, other=0.0) b_gkn = exp2(b_gk - b_gn) b_qg = b_qj * tl.where(m_j[:, None], b_gkn, 0) b_kbg = b_kbj * tl.where(m_j[:, None], b_gkn, 0) - b_dkt += tl.dot(b_dAqk, b_qg) - b_dkt += tl.dot(b_dAkk, b_kbg) + b_dkt = tl.dot(b_dAqk, b_qg, b_dkt) + b_dkt = tl.dot(b_dAkk, b_kbg, b_dkt) b_dkt *= exp2(b_gn - b_g) o_dA = i_ti * H * BT + i_i * BC + o_i @@ -616,16 +642,18 @@ def chunk_gdn2_bwd_kernel_intra( if SAFE_GATE: if USE_GATHER: - b_gn = gather(b_g, tl.full([1, BK], min(BC // 2, T - i_ti - 1), dtype=tl.int16), axis=0) + b_gn = gather(b_g, tl.zeros([1, BK], dtype=tl.int16), axis=0) else: - p_gn = g + (i_ti + min(BC // 2, T - i_ti - 1)) * H * K + o_k + p_gn = g + i_ti * H * K + o_k b_gn = tl.load(p_gn, mask=m_k, other=0).to(tl.float32)[None, :] - p_q = tl.make_block_ptr(q, (T, K), (H * K, 1), (i_ti, i_k * BK), (BC, BK), (1, 0)) - b_q = tl.load(p_q, boundary_check=(0, 1)) - p_dAqk = tl.make_block_ptr(dAqk, (BT, T), (1, H * BT), (i_i * BC, i_ti), (BC, BC), (0, 1)) - p_dAkk = tl.make_block_ptr(dAkk, (BT, T), (1, H * BT), (i_i * BC, i_ti), (BC, BC), (0, 1)) - b_dAqk_diag_kk = tl.load(p_dAqk, boundary_check=(0, 1)).to(tl.float32) - b_dAkk_diag_kk = tl.load(p_dAkk, boundary_check=(0, 1)).to(tl.float32) + b_gn = b_gn - 0.5 * (BC - 1) * GATE_SHIFT + p_q = q + o_c[:, None] * (H * K) + o_k[None, :] + b_q = tl.load(p_q, mask=m_kc, other=0.0) + m_dk = ((i_i * BC + o_i)[:, None] < BT) & m_c[None, :] + p_dAqk = dAqk + (i_i * BC + o_i)[:, None] + o_c[None, :] * (H * BT) + p_dAkk = dAkk + (i_i * BC + o_i)[:, None] + o_c[None, :] * (H * BT) + b_dAqk_diag_kk = tl.load(p_dAqk, mask=m_dk, other=0.0).to(tl.float32) + b_dAkk_diag_kk = tl.load(p_dAkk, mask=m_dk, other=0.0).to(tl.float32) m_i_diag_kk = (o_i[:, None] <= o_i[None, :]) & ((i_ti + o_i[:, None]) < T) & ((i_ti + o_i[None, :]) < T) m_j_diag_kk = (i_ti + o_i[:, None]) < T @@ -639,8 +667,8 @@ def chunk_gdn2_bwd_kernel_intra( b_q_exp = b_q * exp_b_g_diag_kk b_kb_exp = b_k * b_b * exp_b_g_diag_kk - b_dkt += tl.dot(b_dAqk_diag_kk, b_q_exp) * exp_neg_b_g_diag_kk - b_dkt += tl.dot(b_dAkk_diag_kk, b_kb_exp) * exp_neg_b_g_diag_kk + b_dkt += tl.dot(b_dAqk_diag_kk, b_q_exp, input_precision=SAFE_GATE_DOT_PRECISION) * exp_neg_b_g_diag_kk + b_dkt += tl.dot(b_dAkk_diag_kk, b_kb_exp, input_precision=SAFE_GATE_DOT_PRECISION) * exp_neg_b_g_diag_kk else: for j in range(0, min(BC, T - i_t * BT - i_i * BC)): b_dAqk = tl.load(dAqk + o_dA + j * H * BT) @@ -660,19 +688,20 @@ def chunk_gdn2_bwd_kernel_intra( p_gkj += H * K p_bj_ptr += H * K - p_dk = tl.make_block_ptr(dk, (T, K), (H * K, 1), (i_ti, i_k * BK), (BC, BK), (1, 0)) - p_dk2 = tl.make_block_ptr(dk2, (T, K), (H * K, 1), (i_ti, i_k * BK), (BC, BK), (1, 0)) - p_dg = tl.make_block_ptr(dg, (T, K), (H * K, 1), (i_ti, i_k * BK), (BC, BK), (1, 0)) - p_dg2 = tl.make_block_ptr(dg2, (T, K), (H * K, 1), (i_ti, i_k * BK), (BC, BK), (1, 0)) + p_dk = dk + o_c[:, None] * (H * K) + o_k[None, :] + p_dk2 = dk2 + o_c[:, None] * (H * K) + o_k[None, :] + p_dg = dg + o_c[:, None] * (H * K) + o_k[None, :] + p_dg2 = dg2 + o_c[:, None] * (H * K) + o_k[None, :] - b_dg2 += (b_dk2 - b_dkt) * b_k + tl.load(p_dg, boundary_check=(0, 1)) - b_dk2 += tl.load(p_dk, boundary_check=(0, 1)) + b_dg2 += (b_dk2 - b_dkt) * b_k + tl.load(p_dg, mask=m_kc, other=0.0) + b_dk2 += tl.load(p_dk, mask=m_kc, other=0.0) b_dk2 += b_dkt - tl.store(p_dk2, b_dk2.to(p_dk2.dtype.element_ty), boundary_check=(0, 1)) - tl.store(p_dg2, b_dg2.to(p_dg2.dtype.element_ty), boundary_check=(0, 1)) + tl.store(p_dk2, b_dk2.to(p_dk2.dtype.element_ty), mask=m_kc) + tl.store(p_dg2, b_dg2.to(p_dg2.dtype.element_ty), mask=m_kc) +@dispatch('gdn2') def chunk_gdn2_fwd_intra( q: torch.Tensor, k: torch.Tensor, @@ -685,9 +714,16 @@ def chunk_gdn2_fwd_intra( chunk_size: int = 64, chunk_indices: torch.LongTensor | None = None, safe_gate: bool = False, + lower_bound: float | None = None, disable_recompute: bool = False, ): - """Intra-chunk forward: build Aqk, Akk_inv, and the WY auxiliaries (w, u).""" + """Build Aqk, Akk_inv, and WY auxiliaries with causal safe-gate anchors. + + As in chunk_kda_fwd_intra, diagonal blocks use the first cumulative gate minus a fixed shift. + The default strength of 2 gives a shift of 15 in natural-log units for BC=16. + Both fused and precomputed public gate paths use this default; direct intra callers can override it with lower_bound. + Forward and backward retain the same shift even for partial blocks. + """ B, T, H, K = k.shape BT = chunk_size BC = 16 @@ -720,6 +756,7 @@ def chunk_gdn2_fwd_intra( BC=BC, BK=BK, USE_GATHER=IS_GATHER_SUPPORTED, + GATE_SHIFT=(-lower_bound if lower_bound is not None else 2.) * RCP_LN2, ) else: chunk_gdn2_fwd_intra_token_parallel( @@ -785,6 +822,7 @@ def chunk_gdn2_bwd_intra( chunk_indices: torch.LongTensor | None = None, chunk_size: int = 64, safe_gate: bool = False, + lower_bound: float | None = None, ): """Intra-chunk backward: q, k, g, b contributions from dAqk, dAkk. @@ -807,7 +845,7 @@ def chunk_gdn2_bwd_intra( db2 = q.new_empty(NK, B, T, H, BK, dtype=torch.float32) dg2 = torch.empty_like(dg, dtype=torch.float32) - grid = (NK * NC, NT, B * H) + grid = (NK * NC * NT, B * H) chunk_gdn2_bwd_kernel_intra[grid]( q=q, k=k, @@ -834,6 +872,7 @@ def chunk_gdn2_bwd_intra( NC=NC, SAFE_GATE=safe_gate, USE_GATHER=IS_GATHER_SUPPORTED, + GATE_SHIFT=(-lower_bound if lower_bound is not None else 2.) * RCP_LN2, ) dq = dq2 dk = dk2 diff --git a/build/torch-cuda/ops/gdn2/chunk_intra_token_parallel.py b/build/torch-cuda/ops/gdn2/chunk_intra_token_parallel.py index 90e091618d3c2efa997f587614c084942efdccb7..7e985cbbbceb1cc66d9b34e309bf1c4d274210a8 100644 --- a/build/torch-cuda/ops/gdn2/chunk_intra_token_parallel.py +++ b/build/torch-cuda/ops/gdn2/chunk_intra_token_parallel.py @@ -52,12 +52,13 @@ def chunk_gdn2_fwd_kernel_intra_token_parallel( T, H: tl.constexpr, K: tl.constexpr, + BK: tl.constexpr, BT: tl.constexpr, BC: tl.constexpr, BH: tl.constexpr, IS_VARLEN: tl.constexpr, ): - i_tg, i_hg = tl.program_id(0), tl.program_id(1) + i_tg, i_hg = tl.program_id(0).to(tl.int64), tl.program_id(1).to(tl.int64) if IS_VARLEN: i_n = 0 @@ -70,7 +71,7 @@ def chunk_gdn2_fwd_kernel_intra_token_parallel( else: left = mid + 1 i_n = left - bos, eos = tl.load(cu_seqlens + i_n).to(tl.int32), tl.load(cu_seqlens + i_n + 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 i_t = i_tg - bos else: @@ -92,30 +93,29 @@ def chunk_gdn2_fwd_kernel_intra_token_parallel( Aqk += bos * H * BT Akk += bos * H * BC - BK: tl.constexpr = triton.next_power_of_2(K) o_h = tl.arange(0, BH) o_k = tl.arange(0, BK) m_h = (i_hg * BH + o_h) < H m_k = o_k < K - p_q = tl.make_block_ptr(q + i_t * H * K, (H, K), (K, 1), (i_hg * BH, 0), (BH, BK), (1, 0)) - p_k = tl.make_block_ptr(k + i_t * H * K, (H, K), (K, 1), (i_hg * BH, 0), (BH, BK), (1, 0)) - p_g = tl.make_block_ptr(g + i_t * H * K, (H, K), (K, 1), (i_hg * BH, 0), (BH, BK), (1, 0)) - p_b = tl.make_block_ptr(b + i_t * H * K, (H, K), (K, 1), (i_hg * BH, 0), (BH, BK), (1, 0)) + p_q = q + i_t * H * K + (i_hg * BH + o_h)[:, None] * K + o_k[None, :] + p_k = k + i_t * H * K + (i_hg * BH + o_h)[:, None] * K + o_k[None, :] + p_g = g + i_t * H * K + (i_hg * BH + o_h)[:, None] * K + o_k[None, :] + p_b = b + i_t * H * K + (i_hg * BH + o_h)[:, None] * K + o_k[None, :] - b_q = tl.load(p_q, boundary_check=(0, 1)).to(tl.float32) - b_k = tl.load(p_k, boundary_check=(0, 1)).to(tl.float32) - b_g = tl.load(p_g, boundary_check=(0, 1)).to(tl.float32) - b_b = tl.load(p_b, boundary_check=(0, 1)).to(tl.float32) + b_q = tl.load(p_q, mask=m_h[:, None] & m_k[None, :], other=0.0).to(tl.float32) + b_k = tl.load(p_k, mask=m_h[:, None] & m_k[None, :], other=0.0).to(tl.float32) + b_g = tl.load(p_g, mask=m_h[:, None] & m_k[None, :], other=0.0).to(tl.float32) + b_b = tl.load(p_b, mask=m_h[:, None] & m_k[None, :], other=0.0).to(tl.float32) # Fold channel-wise erase gate into the key tile. b_k = b_k * b_b for j in range(i_ts, min(i_t + 1, min(T, i_ts + BC))): - p_kj = tl.make_block_ptr(k + j * H * K, (H, K), (K, 1), (i_hg * BH, 0), (BH, BK), (1, 0)) - p_gj = tl.make_block_ptr(g + j * H * K, (H, K), (K, 1), (i_hg * BH, 0), (BH, BK), (1, 0)) - b_kj = tl.load(p_kj, boundary_check=(0, 1)).to(tl.float32) - b_gj = tl.load(p_gj, boundary_check=(0, 1)).to(tl.float32) + p_kj = k + j * H * K + (i_hg * BH + o_h)[:, None] * K + o_k[None, :] + p_gj = g + j * H * K + (i_hg * BH + o_h)[:, None] * K + o_k[None, :] + b_kj = tl.load(p_kj, mask=m_h[:, None] & m_k[None, :], other=0.0).to(tl.float32) + b_gj = tl.load(p_gj, mask=m_h[:, None] & m_k[None, :], other=0.0).to(tl.float32) b_kgj = b_kj * exp2(b_g - b_gj) b_kgj = tl.where(m_k[None, :], b_kgj, 0.0) @@ -149,6 +149,7 @@ def chunk_gdn2_fwd_intra_token_parallel( N = len(cu_seqlens) - 1 if cu_seqlens is not None else B BT = chunk_size BC = sub_chunk_size + BK = triton.next_power_of_2(K) def grid(meta): return (B * T, triton.cdiv(H, meta['BH'])) @@ -166,6 +167,7 @@ def chunk_gdn2_fwd_intra_token_parallel( T=T, H=H, K=K, + BK=BK, BT=BT, BC=BC, ) diff --git a/build/torch-cuda/ops/gdn2/fused_recurrent.py b/build/torch-cuda/ops/gdn2/fused_recurrent.py index 3e7b2cd267feeb9b123332b7ed83252b0edfbcc8..1d7c642a52fed7041ab0abdb903216e8ad68a83a 100644 --- a/build/torch-cuda/ops/gdn2/fused_recurrent.py +++ b/build/torch-cuda/ops/gdn2/fused_recurrent.py @@ -40,6 +40,7 @@ import torch import triton import triton.language as tl +from ...backends import dispatch from ...ops.utils.op import exp from ...ops.utils.softplus import softplus from ...utils import input_guard @@ -52,7 +53,8 @@ from ...utils import input_guard "IS_VARLEN": lambda args: args["cu_seqlens"] is not None, "IS_CONTINUOUS_BATCHING": lambda args: args["ssm_state_indices"] is not None, "IS_SPEC_DECODING": lambda args: args["num_accepted_tokens"] is not None, - "HAS_DT_BIAS": lambda args: args["dt_bias"] is not None, + "HAS_A": lambda args: args["A_log"] is not None, + "HAS_BIAS": lambda args: args["dt_bias"] is not None, "USE_LOWER_BOUND": lambda args: args["lower_bound"] is not None, } ) @@ -93,13 +95,14 @@ def fused_recurrent_gdn2_fwd_kernel( IS_CONTINUOUS_BATCHING: tl.constexpr, IS_SPEC_DECODING: tl.constexpr, STORE_FINAL_STATE: tl.constexpr, - HAS_DT_BIAS: tl.constexpr, + HAS_A: tl.constexpr, + HAS_BIAS: tl.constexpr, USE_GATE_IN_KERNEL: tl.constexpr, USE_LOWER_BOUND: tl.constexpr, STATE_V_FIRST: tl.constexpr, num_stages: tl.constexpr, ): - pid = tl.program_id(0) + pid = tl.program_id(0).to(tl.int64) NV = tl.cdiv(V, BV) NK = tl.cdiv(K, BK) i_k = pid % NK @@ -174,17 +177,17 @@ def fused_recurrent_gdn2_fwd_kernel( b_q = b_q / tl.sqrt(tl.sum(b_q * b_q) + 1e-6) b_k = b_k / tl.sqrt(tl.sum(b_k * b_k) + 1e-6) b_q = b_q * scale - b_g = tl.load(p_g, eviction_policy='evict_last').to(tl.float32) + b_g = tl.load(p_g, mask=mask_k, other=0, eviction_policy='evict_last').to(tl.float32) if USE_GATE_IN_KERNEL: - b_A = tl.load(A_log + i_hv).to(tl.float32) + b_A = tl.load(A_log + i_hv).to(tl.float32) if HAS_A else 1.0 - if HAS_DT_BIAS: + if HAS_BIAS: b_bias = tl.load(dt_bias + i_hv * K + o_k, mask=mask_k, other=0).to(tl.float32) b_g = b_g + b_bias if USE_LOWER_BOUND: - b_gk = lower_bound * tl.sigmoid(exp(b_A) * b_g) + b_gk = lower_bound * tl.sigmoid((exp(b_A) if HAS_A else b_A) * b_g) else: b_gk = -exp(b_A) * softplus(b_g) else: @@ -248,7 +251,7 @@ def fused_recurrent_gdn2_fwd_kernel( tl.store(p_ht, b_h.to(p_ht.dtype.element_ty), mask=mask_h) -@torch.compiler.disable +@dispatch('gdn2') def fused_recurrent_gdn2_fwd( q: torch.Tensor, k: torch.Tensor, @@ -377,35 +380,49 @@ def fused_recurrent_gdn2( Setting ``b = w = beta`` (scalar) recovers KDA exactly. Args: - q (torch.Tensor): queries of shape ``[B, T, H, K]``. - k (torch.Tensor): keys of shape ``[B, T, H, K]``. - v (torch.Tensor): values of shape ``[B, T, HV, V]`` (GVA if HV > H). - g (torch.Tensor): channel-wise log-decay of shape ``[B, T, HV, K]``. - When ``use_gate_in_kernel=True``, this is the raw pre-activation and - the kernel computes ``-exp(A_log) * softplus(g + dt_bias)`` (or the - bounded variant if ``lower_bound`` is set). - b (torch.Tensor): channel-wise erase gate of shape ``[B, T, HV, K]``. - w (torch.Tensor): channel-wise write gate of shape ``[B, T, HV, V]``. - A_log (Optional[torch.Tensor]): decay parameter of shape ``[HV]``. - Required when ``use_gate_in_kernel=True``. - dt_bias (Optional[torch.Tensor]): per-channel bias of shape ``[HV * K]``, - only used when ``use_gate_in_kernel=True``. - scale (Optional[float]): attention scale, defaults to ``1/sqrt(K)``. - initial_state (Optional[torch.Tensor]): ``[N, HV, K, V]`` in float32 - (or ``[N, HV, V, K]`` if ``state_v_first=True``). - output_final_state (bool): whether to output the final state. - use_qk_l2norm_in_kernel (bool): L2-normalize q and k inside the kernel. - use_gate_in_kernel (bool): compute the gate activation from raw ``g``. - lower_bound (Optional[float]): when ``use_gate_in_kernel=True`` and set, - use ``lower_bound * sigmoid(exp(A_log) * g)`` instead of the softplus - activation. - state_v_first (bool): store state as ``[V, K]`` instead of the default - ``[K, V]``. - cu_seqlens (Optional[torch.LongTensor]): ``[N+1]`` packed-sequence offsets. + q (torch.Tensor): + queries of shape ``[B, T, H, K]``. + k (torch.Tensor): + keys of shape ``[B, T, H, K]``. + v (torch.Tensor): + values of shape ``[B, T, HV, V]`` (GVA if HV > H). + g (torch.Tensor): + channel-wise log-decay of shape ``[B, T, HV, K]``. + When ``use_gate_in_kernel=True``, this is the raw pre-activation, + and the kernel computes ``-exp(A_log) * softplus(g + dt_bias)`` (or the bounded variant if ``lower_bound`` is set). + b (torch.Tensor): + channel-wise erase gate of shape ``[B, T, HV, K]``. + w (torch.Tensor): + channel-wise write gate of shape ``[B, T, HV, V]``. + A_log (Optional[torch.Tensor]): + decay parameter of shape ``[HV]``. + When ``use_gate_in_kernel=True`` together with ``lower_bound``, + may be ``None`` to use ``lower_bound * sigmoid(g + dt_bias)``. + dt_bias (Optional[torch.Tensor]): + per-channel bias of shape ``[HV * K]``, only used when ``use_gate_in_kernel=True``. + scale (Optional[float]): + attention scale. Default: ``1/sqrt(K)``. + initial_state (Optional[torch.Tensor]): + ``[N, HV, K, V]`` in float32 (or ``[N, HV, V, K]`` if ``state_v_first=True``). + output_final_state (Optional[bool]): + whether to output the final state. Default: `False`. + use_qk_l2norm_in_kernel (Optional[bool]): + L2-normalize q and k inside the kernel. Default: `False`. + use_gate_in_kernel (Optional[bool]): + compute the gate activation from raw ``g``. Default: `False`. + lower_bound (Optional[float]): + when ``use_gate_in_kernel=True`` and set, + use ``lower_bound * sigmoid(exp(A_log) * g)`` instead of the softplus activation. + state_v_first (Optional[bool]): + store state as ``[V, K]`` instead of the default ``[K, V]``. Default: `False`. + cu_seqlens (Optional[torch.LongTensor]): + ``[N+1]`` packed-sequence offsets. Returns: - o (torch.Tensor): outputs of shape ``[B, T, HV, V]``. - final_state (Optional[torch.Tensor]): ``[N, HV, K, V]`` if requested. + o (torch.Tensor): + outputs of shape ``[B, T, HV, V]``. + final_state (Optional[torch.Tensor]): + ``[N, HV, K, V]`` if requested. """ if 'transpose_state_layout' in kwargs: if state_v_first: @@ -428,9 +445,9 @@ def fused_recurrent_gdn2( f"The number of initial states is expected to be equal to the number of input sequences, " f"i.e., {len(cu_seqlens) - 1} rather than {initial_state.shape[0]}." ) - assert b.shape == (*q.shape[:3], k.shape[-1]), ( + assert b.shape == (*v.shape[:3], k.shape[-1]), ( f"b must have shape [B, T, HV, K]; got {tuple(b.shape)} " - f"vs expected {(*q.shape[:3], k.shape[-1])}." + f"vs expected {(*v.shape[:3], k.shape[-1])}." ) assert w.shape == v.shape, ( f"w must have shape [B, T, HV, V] matching v; got {tuple(w.shape)} " diff --git a/build/torch-cuda/ops/gdn2/wy_fast.py b/build/torch-cuda/ops/gdn2/wy_fast.py index 7f758624d39ae01d20b0f026fc7fcbd9d647f77c..ad59f612d45bb96fd1761e2c012f27c8739e81b6 100644 --- a/build/torch-cuda/ops/gdn2/wy_fast.py +++ b/build/torch-cuda/ops/gdn2/wy_fast.py @@ -73,76 +73,54 @@ def recompute_w_u_fwd_gdn2_kernel( STORE_KG: tl.constexpr, IS_VARLEN: tl.constexpr, ): - i_t, i_bh = tl.program_id(0), tl.program_id(1) + i_t, i_bh = tl.program_id(0).to(tl.int64), tl.program_id(1).to(tl.int64) i_b, i_h = i_bh // H, i_bh % H 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.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) + i_n, i_t = tl.load(chunk_indices + i_t * 2).to(tl.int32), tl.load(chunk_indices + i_t * 2 + 1).to(tl.int64) + bos, eos = tl.load(cu_seqlens + i_n).to(tl.int64), tl.load(cu_seqlens + i_n + 1).to(tl.int64) T = eos - bos else: bos, eos = i_b * T, i_b * T + T - p_A = tl.make_block_ptr( - A + (bos * H + i_h) * BT, (T, BT), (H * BT, 1), - (i_t * BT, 0), (BT, BT), (1, 0), - ) - b_A = tl.load(p_A, boundary_check=(0, 1)) + o_t = i_t * BT + tl.arange(0, BT) + o_A = tl.arange(0, BT) + m_t = o_t < T + m_A = m_t[:, None] & (o_A[None, :] < BT) + p_A = A + (bos * H + i_h) * BT + o_t[:, None] * (H * BT) + o_A[None, :] + b_A = tl.load(p_A, mask=m_A, other=0.0) for i_v in range(tl.cdiv(V, BV)): - p_v = tl.make_block_ptr( - v + (bos * H + i_h) * V, (T, V), (H * V, 1), - (i_t * BT, i_v * BV), (BT, BV), (1, 0), - ) - p_u = tl.make_block_ptr( - u + (bos * H + i_h) * V, (T, V), (H * V, 1), - (i_t * BT, i_v * BV), (BT, BV), (1, 0), - ) - p_wg = tl.make_block_ptr( - w_gate + (bos * H + i_h) * V, (T, V), (H * V, 1), - (i_t * BT, i_v * BV), (BT, BV), (1, 0), - ) - b_v = tl.load(p_v, boundary_check=(0, 1)) - b_wg = tl.load(p_wg, boundary_check=(0, 1)) + o_v = i_v * BV + tl.arange(0, BV) + m_v = m_t[:, None] & (o_v[None, :] < V) + p_v = v + (bos * H + i_h) * V + o_t[:, None] * (H * V) + o_v[None, :] + p_u = u + (bos * H + i_h) * V + o_t[:, None] * (H * V) + o_v[None, :] + p_wg = w_gate + (bos * H + i_h) * V + o_t[:, None] * (H * V) + o_v[None, :] + b_v = tl.load(p_v, mask=m_v, other=0.0) + b_wg = tl.load(p_wg, mask=m_v, other=0.0) b_vb = (b_v * b_wg).to(b_v.dtype) b_u = tl.dot(b_A, b_vb) - tl.store(p_u, b_u.to(p_u.dtype.element_ty), boundary_check=(0, 1)) + tl.store(p_u, b_u.to(p_u.dtype.element_ty), mask=m_v) for i_k in range(tl.cdiv(K, BK)): - p_w = tl.make_block_ptr( - w + (bos * H + i_h) * K, (T, K), (H * K, 1), - (i_t * BT, i_k * BK), (BT, BK), (1, 0), - ) - p_k = tl.make_block_ptr( - k + (bos * H + i_h) * K, (T, K), (H * K, 1), - (i_t * BT, i_k * BK), (BT, BK), (1, 0), - ) - p_b = tl.make_block_ptr( - b + (bos * H + i_h) * K, (T, K), (H * K, 1), - (i_t * BT, i_k * BK), (BT, BK), (1, 0), - ) - b_k = tl.load(p_k, boundary_check=(0, 1)) - b_b = tl.load(p_b, boundary_check=(0, 1)) + o_k = i_k * BK + tl.arange(0, BK) + m_kk = m_t[:, None] & (o_k[None, :] < K) + p_w = w + (bos * H + i_h) * K + o_t[:, None] * (H * K) + o_k[None, :] + p_k = k + (bos * H + i_h) * K + o_t[:, None] * (H * K) + o_k[None, :] + p_b = b + (bos * H + i_h) * K + o_t[:, None] * (H * K) + o_k[None, :] + b_k = tl.load(p_k, mask=m_kk, other=0.0) + b_b = tl.load(p_b, mask=m_kk, other=0.0) b_kb = b_k * b_b - p_gk = tl.make_block_ptr( - gk + (bos * H + i_h) * K, (T, K), (H * K, 1), - (i_t * BT, i_k * BK), (BT, BK), (1, 0), - ) - b_gk = tl.load(p_gk, boundary_check=(0, 1)).to(tl.float32) + p_gk = gk + (bos * H + i_h) * K + o_t[:, None] * (H * K) + o_k[None, :] + b_gk = tl.load(p_gk, mask=m_kk, other=0.0).to(tl.float32) b_kb *= exp2(b_gk) if STORE_QG: - p_q = tl.make_block_ptr( - q + (bos * H + i_h) * K, (T, K), (H * K, 1), - (i_t * BT, i_k * BK), (BT, BK), (1, 0), - ) - p_qg = tl.make_block_ptr( - qg + (bos * H + i_h) * K, (T, K), (H * K, 1), - (i_t * BT, i_k * BK), (BT, BK), (1, 0), - ) - b_q = tl.load(p_q, boundary_check=(0, 1)) + p_q = q + (bos * H + i_h) * K + o_t[:, None] * (H * K) + o_k[None, :] + p_qg = qg + (bos * H + i_h) * K + o_t[:, None] * (H * K) + o_k[None, :] + b_q = tl.load(p_q, mask=m_kk, other=0.0) b_qg = b_q * exp2(b_gk) - tl.store(p_qg, b_qg.to(p_qg.dtype.element_ty), boundary_check=(0, 1)) + tl.store(p_qg, b_qg.to(p_qg.dtype.element_ty), mask=m_kk) if STORE_KG: last_idx = min(i_t * BT + BT, T) - 1 @@ -157,14 +135,11 @@ def recompute_w_u_fwd_gdn2_kernel( exp2(b_gn[None, :] - b_gk), 0, ) - p_kg = tl.make_block_ptr( - kg + (bos * H + i_h) * K, (T, K), (H * K, 1), - (i_t * BT, i_k * BK), (BT, BK), (1, 0), - ) - tl.store(p_kg, b_kg.to(p_kg.dtype.element_ty), boundary_check=(0, 1)) + p_kg = kg + (bos * H + i_h) * K + o_t[:, None] * (H * K) + o_k[None, :] + tl.store(p_kg, b_kg.to(p_kg.dtype.element_ty), mask=m_kk) b_w = tl.dot(b_A, b_kb.to(b_k.dtype)) - tl.store(p_w, b_w.to(p_w.dtype.element_ty), boundary_check=(0, 1)) + tl.store(p_w, b_w.to(p_w.dtype.element_ty), mask=m_kk) def recompute_w_u_fwd_gdn2( diff --git a/build/torch-xpu/ops/moba/__init__.py b/build/torch-cuda/ops/generalized_delta_rule/dplr/backends/__init__.py similarity index 82% rename from build/torch-xpu/ops/moba/__init__.py rename to build/torch-cuda/ops/generalized_delta_rule/dplr/backends/__init__.py index 2e0e045f19b7553ca74db815e14d2580573e5ee1..436af9365f727a97644619e2120123c05ffbdbc3 100644 --- a/build/torch-xpu/ops/moba/__init__.py +++ b/build/torch-cuda/ops/generalized_delta_rule/dplr/backends/__init__.py @@ -5,6 +5,7 @@ # For a list of all contributors, visit: # https://github.com/fla-org/flash-linear-attention/graphs/contributors -from .parallel import parallel_moba +"""DPLR backends.""" -__all__ = ['parallel_moba'] + +__all__: list[str] = [] diff --git a/build/torch-cuda/ops/generalized_delta_rule/dplr/chunk.py b/build/torch-cuda/ops/generalized_delta_rule/dplr/chunk.py index 240d0f71ffc59774bcf079e88622520692715898..32fccd845b7a3bd2ec76ed6aa3e85eea4c8453b9 100644 --- a/build/torch-cuda/ops/generalized_delta_rule/dplr/chunk.py +++ b/build/torch-cuda/ops/generalized_delta_rule/dplr/chunk.py @@ -9,6 +9,7 @@ import warnings import torch +from ....backends import dispatch from ....ops.cp import FLACPContext from ....ops.cp.chunk_delta_h import ( chunk_gated_delta_rule_bwd_dhu_pre_process, @@ -30,6 +31,20 @@ from ....ops.utils.constant import RCP_LN2 from ....utils import TRITON_ABOVE_3_4_0, autocast_custom_bwd, autocast_custom_fwd, input_guard +def gate_bound_is_safe(lower_bound: float, chunk_size: int) -> bool: + """Whether `gk >= lower_bound` keeps the centered tensor-core A-stage in fp32 range. + + The A-stages factorize exp2(gi[i] - gi[j]) around the mid-chunk row. The + largest positive exponent is on the a-side: the exclusive cumsum ge is + centered against the inclusive gi[mid], so it spans chunk_size/2 + 1 rows + and per-row exponents reach (chunk_size/2 + 1) * |lower_bound| * log2(e). + 124 is the fp32 exponent limit (~128 log2) minus 4 log2 of headroom for + the activation multiply. A non-negative bound is invalid (gk is a + log-decay < 0) and never licenses the scheme. + """ + return lower_bound < 0 and abs(lower_bound) * (chunk_size // 2 + 1) * RCP_LN2 <= 124 + + def chunk_dplr_fwd( q: torch.Tensor, k: torch.Tensor, @@ -79,6 +94,7 @@ def chunk_dplr_fwd( cu_seqlens=cu_seqlens, chunk_size=chunk_size, chunk_indices=chunk_indices, + safe_gate=safe_gate, ) if cp_context is not None: @@ -149,18 +165,17 @@ class ChunkDPLRDeltaRuleFunction(torch.autograd.Function): cu_seqlens: torch.LongTensor | None = None, cu_seqlens_cpu: torch.LongTensor | None = None, safe_gate: bool = False, + lower_bound: float | None = None, chunk_size: int | None = None, disable_recompute: bool = False, cp_context: FLACPContext | None = None, ): - # Due to gate numerical stability consideration, we only support chunk_size=16 when safe_gate=True - # And in practice, chunk_size=16 is sufficient for no safe gate situations. - # It's different from the other chunk implementations. if chunk_size is None: - chunk_size = 16 - elif TRITON_ABOVE_3_4_0: - chunk_size = chunk_size - else: + if TRITON_ABOVE_3_4_0 and lower_bound is not None and gate_bound_is_safe(lower_bound, 64): + chunk_size = 64 + else: + chunk_size = 16 + elif not TRITON_ABOVE_3_4_0: # Avoid Triton Compiler error warnings.warn( "Set chunk_size to 16, to avoid triton compiler erorr. " @@ -169,6 +184,10 @@ class ChunkDPLRDeltaRuleFunction(torch.autograd.Function): stacklevel=2, ) chunk_size = 16 + if not safe_gate and lower_bound is not None and gate_bound_is_safe(lower_bound, chunk_size): + # the caller-asserted gate bound licenses the same centered + # tensor-core scheme as safe_gate=True + safe_gate = True chunk_indices = None if cu_seqlens is not None: chunk_indices = prepare_chunk_indices(cu_seqlens, chunk_size, cu_seqlens_cpu=cu_seqlens_cpu) @@ -305,6 +324,7 @@ class ChunkDPLRDeltaRuleFunction(torch.autograd.Function): cu_seqlens=cu_seqlens, chunk_size=chunk_size, chunk_indices=ctx.chunk_indices, + safe_gate=ctx.safe_gate, ) del A_ab h, v_new, _ = chunk_dplr_fwd_h( @@ -434,10 +454,11 @@ class ChunkDPLRDeltaRuleFunction(torch.autograd.Function): return ( dq.to(q), dk.to(k), dv.to(v), da.to(a), db.to(b), dgk.to(gk), - None, dh0, None, None, None, None, None, None, None, + None, dh0, None, None, None, None, None, None, None, None, ) +@dispatch('generalized_delta_rule.dplr') @torch.compiler.disable def chunk_dplr_delta_rule( q: torch.Tensor, @@ -455,6 +476,7 @@ def chunk_dplr_delta_rule( chunk_size: int | None = None, disable_recompute: bool = False, cp_context: FLACPContext | None = None, + lower_bound: float | None = None, **kwargs, ): r""" @@ -490,13 +512,19 @@ def chunk_dplr_delta_rule( When `True`, the kernel can use M=16 TensorCore acceleration. The safe range is approximately `[-5, 0)`. Default: `False`. chunk_size (Optional[int]): - Chunk size for the chunked computation. Default: `None`, which means 16. + Chunk size for the chunked computation. Default: `None`, which means 64 + when `lower_bound` is given and `gate_bound_is_safe(lower_bound, 64)` + holds, and 16 otherwise. disable_recompute (Optional[bool]): Whether to disable gradient recomputation in the kernel. Default: `False`. cp_context (Optional[FLACPContext]): Context parallel context for distributed training across multiple devices. When provided, `initial_state` and `output_final_state` are not supported. Default: `None`. + lower_bound (Optional[float]): + When set, asserts `lower_bound <= gk < 0` (the caller is responsible for + the guarantee). Licenses the same tensor-core scheme as `safe_gate=True` + when the bound fits the chunk size. Default: `None`. Returns: o (torch.Tensor): @@ -515,6 +543,8 @@ def chunk_dplr_delta_rule( raise DeprecationWarning( "head_first has been removed. Inputs must be in `[B, T, H, ...]` format.", ) + if lower_bound is not None and lower_bound >= 0: + raise ValueError(f"`lower_bound` must be negative (gk is a log-decay < 0), got {lower_bound}.") 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" @@ -547,6 +577,7 @@ def chunk_dplr_delta_rule( cu_seqlens, cu_seqlens_cpu, safe_gate, + lower_bound, chunk_size, disable_recompute, cp_context, diff --git a/build/torch-cuda/ops/generalized_delta_rule/dplr/chunk_A_bwd.py b/build/torch-cuda/ops/generalized_delta_rule/dplr/chunk_A_bwd.py index fa6c97b7d5c5e1f8926ab19d2a1c486ece9e7bec..4f0ccf48b52436bd94f3beb706b6f53dd36c412c 100644 --- a/build/torch-cuda/ops/generalized_delta_rule/dplr/chunk_A_bwd.py +++ b/build/torch-cuda/ops/generalized_delta_rule/dplr/chunk_A_bwd.py @@ -10,7 +10,7 @@ import triton import triton.language as tl from ....ops.utils import prepare_chunk_indices -from ....ops.utils.op import exp2, gather +from ....ops.utils.op import exp2, gather, unflatten_program_id from ....utils import IS_AMD, IS_GATHER_SUPPORTED, autotune_cache_kwargs, check_shared_mem NUM_WARPS_AUTOTUNE = [2, 4, 8, 16] if IS_AMD else [2, 4, 8, 16, 32] @@ -62,14 +62,15 @@ def chunk_dplr_bwd_kernel_intra( IS_VARLEN: tl.constexpr, GATHER_SUPPORTED: tl.constexpr, ): - i_k, i_t, i_bh = tl.program_id(0), tl.program_id(1), tl.program_id(2) + i_k, i_t = unflatten_program_id(tl.cdiv(K, BK)) + i_bh = tl.program_id(1).to(tl.int64) i_b, i_h = i_bh // H, i_bh % H 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.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) + i_n, i_t = tl.load(chunk_indices + i_t * 2).to(tl.int32), tl.load(chunk_indices + i_t * 2 + 1).to(tl.int64) + bos, eos = tl.load(cu_seqlens + i_n).to(tl.int64), tl.load(cu_seqlens + i_n + 1).to(tl.int64) T = eos - bos else: - bos, eos = (i_b * T).to(tl.int32), (i_b * T + T).to(tl.int32) + bos, eos = i_b * T, i_b * T + T if i_t * BT >= T: return @@ -99,37 +100,41 @@ def chunk_dplr_bwd_kernel_intra( stride_qk = H*K stride_A = H*BT - p_ge = tl.make_block_ptr(ge, (T, K), (stride_qk, 1), (i_t * BT, i_k * BK), (BC, BK), (1, 0)) - p_gi = tl.make_block_ptr(gi, (T, K), (stride_qk, 1), (i_t * BT, i_k * BK), (BC, BK), (1, 0)) + o_i = tl.arange(0, BC) + o_c = i_t * BT + o_i + o_k = i_k * BK + tl.arange(0, BK) + m_c = o_c < T + m_k = o_k < K + m_kc = m_c[:, None] & m_k[None, :] + p_ge = ge + o_c[:, None] * stride_qk + o_k[None, :] + p_gi = gi + o_c[:, None] * stride_qk + o_k[None, :] # [BC, BK] - b_ge = tl.load(p_ge, boundary_check=(0, 1)) - b_gi = tl.load(p_gi, boundary_check=(0, 1)) + b_ge = tl.load(p_ge, mask=m_kc, other=0.0) + b_gi = tl.load(p_gi, mask=m_kc, other=0.0) b_dq = tl.zeros([BC, BK], dtype=tl.float32) b_da = tl.zeros([BC, BK], dtype=tl.float32) b_dk = tl.zeros([BC, BK], dtype=tl.float32) b_db = tl.zeros([BC, BK], dtype=tl.float32) # intra chunk gradient calculation - p_dAqk = tl.make_block_ptr(dAqk, (T, BT), (stride_A, 1), (i_t*BT, 0), (BC, BC), (1, 0)) - p_dAab = tl.make_block_ptr(dAab, (T, BT), (stride_A, 1), (i_t*BT, 0), (BC, BC), (1, 0)) - p_dAqb = tl.make_block_ptr(dAqb, (T, BT), (stride_A, 1), (i_t*BT, 0), (BC, BC), (1, 0)) - p_dAak = tl.make_block_ptr(dAak, (T, BT), (stride_A, 1), (i_t*BT, 0), (BC, BC), (1, 0)) - o_i = tl.arange(0, BC) - p_k = tl.make_block_ptr(k, (T, K), (stride_qk, 1), (i_t*BT, i_k*BK), (BC, BK), (1, 0)) - p_b = tl.make_block_ptr(b, (T, K), (stride_qk, 1), (i_t*BT, i_k*BK), (BC, BK), (1, 0)) - p_a = tl.make_block_ptr(a, (T, K), (stride_qk, 1), (i_t*BT, i_k*BK), (BC, BK), (1, 0)) - p_q = tl.make_block_ptr(q, (T, K), (stride_qk, 1), (i_t*BT, i_k*BK), (BC, BK), (1, 0)) - b_k = tl.load(p_k, boundary_check=(0, 1)) - b_b = tl.load(p_b, boundary_check=(0, 1)) - b_q = tl.load(p_q, boundary_check=(0, 1)) - b_a = tl.load(p_a, boundary_check=(0, 1)) - b_dAqk = tl.load(p_dAqk, boundary_check=(0, 1)) - b_dAab = tl.load(p_dAab, boundary_check=(0, 1)) - b_dAqb = tl.load(p_dAqb, boundary_check=(0, 1)) - b_dAak = tl.load(p_dAak, boundary_check=(0, 1)) + m_dA = m_c[:, None] & (o_i[None, :] < BT) + p_dAqk = dAqk + o_c[:, None] * stride_A + o_i[None, :] + p_dAab = dAab + o_c[:, None] * stride_A + o_i[None, :] + p_dAqb = dAqb + o_c[:, None] * stride_A + o_i[None, :] + p_dAak = dAak + o_c[:, None] * stride_A + o_i[None, :] + p_k = k + o_c[:, None] * stride_qk + o_k[None, :] + p_b = b + o_c[:, None] * stride_qk + o_k[None, :] + p_a = a + o_c[:, None] * stride_qk + o_k[None, :] + p_q = q + o_c[:, None] * stride_qk + o_k[None, :] + b_k = tl.load(p_k, mask=m_kc, other=0.0) + b_b = tl.load(p_b, mask=m_kc, other=0.0) + b_q = tl.load(p_q, mask=m_kc, other=0.0) + b_a = tl.load(p_a, mask=m_kc, other=0.0) + b_dAqk = tl.load(p_dAqk, mask=m_dA, other=0.0) + b_dAab = tl.load(p_dAab, mask=m_dA, other=0.0) + b_dAqb = tl.load(p_dAqb, mask=m_dA, other=0.0) + b_dAak = tl.load(p_dAak, mask=m_dA, other=0.0) # inter chunk gradient calculation - o_k = i_k * BK + tl.arange(0, BK) - m_k = o_k < K # intra chunk gradient calculation for j in range(0, min(BC, T - i_t * BT)): # trick to index the block @@ -192,32 +197,32 @@ def chunk_dplr_bwd_kernel_intra( b_db += tl.where(m_e, b_dA_ab_j * b_aj * tmp2, 0.) # post processing - p_dq = tl.make_block_ptr(dq, (T, K), (stride_qk, 1), (i_t * BT, i_k * BK), (BC, BK), (1, 0)) - p_dk = tl.make_block_ptr(dk, (T, K), (stride_qk, 1), (i_t * BT, i_k * BK), (BC, BK), (1, 0)) - p_da = tl.make_block_ptr(da, (T, K), (stride_qk, 1), (i_t * BT, i_k * BK), (BC, BK), (1, 0)) - p_db = tl.make_block_ptr(db, (T, K), (stride_qk, 1), (i_t * BT, i_k * BK), (BC, BK), (1, 0)) - p_dgk = tl.make_block_ptr(dgk, (T, K), (stride_qk, 1), (i_t * BT, i_k * BK), (BC, BK), (1, 0)) - p_dgk_offset = tl.make_block_ptr(dgk_offset, (T, K), (stride_qk, 1), (i_t * BT, i_k * BK), (BC, BK), (1, 0)) - p_dqg = tl.make_block_ptr(dqg, (T, K), (stride_qk, 1), (i_t * BT, i_k * BK), (BC, BK), (1, 0)) - p_dkg = tl.make_block_ptr(dkg, (T, K), (stride_qk, 1), (i_t * BT, i_k * BK), (BC, BK), (1, 0)) - p_dag = tl.make_block_ptr(dag, (T, K), (stride_qk, 1), (i_t * BT, i_k * BK), (BC, BK), (1, 0)) - p_dbg = tl.make_block_ptr(dbg, (T, K), (stride_qk, 1), (i_t * BT, i_k * BK), (BC, BK), (1, 0)) + p_dq = dq + o_c[:, None] * stride_qk + o_k[None, :] + p_dk = dk + o_c[:, None] * stride_qk + o_k[None, :] + p_da = da + o_c[:, None] * stride_qk + o_k[None, :] + p_db = db + o_c[:, None] * stride_qk + o_k[None, :] + p_dgk = dgk + o_c[:, None] * stride_qk + o_k[None, :] + p_dgk_offset = dgk_offset + o_c[:, None] * stride_qk + o_k[None, :] + p_dqg = dqg + o_c[:, None] * stride_qk + o_k[None, :] + p_dkg = dkg + o_c[:, None] * stride_qk + o_k[None, :] + p_dag = dag + o_c[:, None] * stride_qk + o_k[None, :] + p_dbg = dbg + o_c[:, None] * stride_qk + o_k[None, :] p_gn = gi + (min(i_t * BT + BT, T) - 1)*stride_qk + o_k p_gn = tl.max_contiguous(tl.multiple_of(p_gn, BK), BK) b_gn = tl.load(p_gn, mask=m_k, other=0) - b_da += tl.load(p_dag, boundary_check=(0, 1)) * exp2(b_ge) - b_dq += tl.load(p_dqg, boundary_check=(0, 1)) * exp2(b_gi) * scale + b_da += tl.load(p_dag, mask=m_kc, other=0.0) * exp2(b_ge) + b_dq += tl.load(p_dqg, mask=m_kc, other=0.0) * exp2(b_gi) * scale tmp = exp2(b_gn[None, :] - b_gi) - b_dk += tl.load(p_dkg, boundary_check=(0, 1)).to(tl.float32) * tmp - b_db += tl.load(p_dbg, boundary_check=(0, 1)).to(tl.float32) * tmp - tl.store(p_dq, (b_dq).to(p_dq.dtype.element_ty), boundary_check=(0, 1)) - tl.store(p_dk, b_dk.to(p_dk.dtype.element_ty), boundary_check=(0, 1)) - tl.store(p_da, b_da.to(p_da.dtype.element_ty), boundary_check=(0, 1)) - tl.store(p_db, b_db.to(p_db.dtype.element_ty), boundary_check=(0, 1)) + b_dk += tl.load(p_dkg, mask=m_kc, other=0.0).to(tl.float32) * tmp + b_db += tl.load(p_dbg, mask=m_kc, other=0.0).to(tl.float32) * tmp + tl.store(p_dq, (b_dq).to(p_dq.dtype.element_ty), mask=m_kc) + tl.store(p_dk, b_dk.to(p_dk.dtype.element_ty), mask=m_kc) + tl.store(p_da, b_da.to(p_da.dtype.element_ty), mask=m_kc) + tl.store(p_db, b_db.to(p_db.dtype.element_ty), mask=m_kc) b_dgk = (b_dq * b_q + b_da * b_a - b_dk * b_k - b_db * b_b).to(tl.float32) b_dgk_offset = b_da * b_a - tl.store(p_dgk, b_dgk.to(p_dgk.dtype.element_ty), boundary_check=(0, 1)) - tl.store(p_dgk_offset, b_dgk_offset.to(p_dgk_offset.dtype.element_ty), boundary_check=(0, 1)) + tl.store(p_dgk, b_dgk.to(p_dgk.dtype.element_ty), mask=m_kc) + tl.store(p_dgk_offset, b_dgk_offset.to(p_dgk_offset.dtype.element_ty), mask=m_kc) @triton.heuristics({ @@ -266,15 +271,16 @@ def chunk_dplr_bwd_kernel_intra_tensorcore( IS_VARLEN: tl.constexpr, GATHER_SUPPORTED: tl.constexpr, ): - i_k, i_t, i_bh = tl.program_id(0), tl.program_id(1), tl.program_id(2) + i_k, i_t = unflatten_program_id(tl.cdiv(K, BK)) + i_bh = tl.program_id(1).to(tl.int64) i_b, i_h = i_bh // H, i_bh % H 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.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) + i_n, i_t = tl.load(chunk_indices + i_t * 2).to(tl.int32), tl.load(chunk_indices + i_t * 2 + 1).to(tl.int64) + bos, eos = tl.load(cu_seqlens + i_n).to(tl.int64), tl.load(cu_seqlens + i_n + 1).to(tl.int64) T_len = eos - bos else: - bos, eos = (i_b * T).to(tl.int32), (i_b * T + T).to(tl.int32) + bos, eos = i_b * T, i_b * T + T T_len = T if i_t * BT >= T_len: @@ -286,23 +292,27 @@ def chunk_dplr_bwd_kernel_intra_tensorcore( valid_len = min(T_len - i_t * BT, BT) mid_idx = valid_len // 2 m_k = tl.arange(0, BK) + i_k * BK < K - p_offset = gi + offset_base_k + (i_t * BT + mid_idx) * K + tl.arange(0, BK) + i_k * BK + p_offset = gi + offset_base_k + (i_t * BT + mid_idx) * H * K + tl.arange(0, BK) + i_k * BK b_offset = tl.load(p_offset, mask=m_k, other=0.0).to(tl.float32) # Q, K, A, B, Gates: [BT, BK] - p_q = tl.make_block_ptr(q + offset_base_k, (T_len, K), (H*K, 1), (i_t*BT, i_k*BK), (BT, BK), (1, 0)) - p_k = tl.make_block_ptr(k + offset_base_k, (T_len, K), (H*K, 1), (i_t*BT, i_k*BK), (BT, BK), (1, 0)) - p_a = tl.make_block_ptr(a + offset_base_k, (T_len, K), (H*K, 1), (i_t*BT, i_k*BK), (BT, BK), (1, 0)) - p_b = tl.make_block_ptr(b + offset_base_k, (T_len, K), (H*K, 1), (i_t*BT, i_k*BK), (BT, BK), (1, 0)) - p_gi = tl.make_block_ptr(gi + offset_base_k, (T_len, K), (H*K, 1), (i_t*BT, i_k*BK), (BT, BK), (1, 0)) - p_ge = tl.make_block_ptr(ge + offset_base_k, (T_len, K), (H*K, 1), (i_t*BT, i_k*BK), (BT, BK), (1, 0)) - - b_q = tl.load(p_q, boundary_check=(0, 1)) - b_k = tl.load(p_k, boundary_check=(0, 1)) - b_a = tl.load(p_a, boundary_check=(0, 1)) - b_b = tl.load(p_b, boundary_check=(0, 1)) - b_gi_val = tl.load(p_gi, boundary_check=(0, 1)).to(tl.float32) - b_ge_val = tl.load(p_ge, boundary_check=(0, 1)).to(tl.float32) + o_t = i_t * BT + tl.arange(0, BT) + o_k = tl.arange(0, BK) + i_k * BK + m_t = o_t < T_len + m_kk = m_t[:, None] & (o_k[None, :] < K) + p_q = q + offset_base_k + o_t[:, None] * (H*K) + o_k[None, :] + p_k = k + offset_base_k + o_t[:, None] * (H*K) + o_k[None, :] + p_a = a + offset_base_k + o_t[:, None] * (H*K) + o_k[None, :] + p_b = b + offset_base_k + o_t[:, None] * (H*K) + o_k[None, :] + p_gi = gi + offset_base_k + o_t[:, None] * (H*K) + o_k[None, :] + p_ge = ge + offset_base_k + o_t[:, None] * (H*K) + o_k[None, :] + + b_q = tl.load(p_q, mask=m_kk, other=0.0) + b_k = tl.load(p_k, mask=m_kk, other=0.0) + b_a = tl.load(p_a, mask=m_kk, other=0.0) + b_b = tl.load(p_b, mask=m_kk, other=0.0) + b_gi_val = tl.load(p_gi, mask=m_kk, other=0.0).to(tl.float32) + b_ge_val = tl.load(p_ge, mask=m_kk, other=0.0).to(tl.float32) b_gi_shifted = b_gi_val - b_offset[None, :] b_ge_shifted = b_ge_val - b_offset[None, :] @@ -315,15 +325,17 @@ def chunk_dplr_bwd_kernel_intra_tensorcore( b_ops = (b_b * inv_exp_gi_shifted).to(tl.float32) a_ops = (b_a * exp_ge_shifted).to(tl.float32) - p_dAqk = tl.make_block_ptr(dAqk + offset_base_attn, (T_len, BT), (H*BT, 1), (i_t*BT, 0), (BT, BT), (1, 0)) - p_dAqb = tl.make_block_ptr(dAqb + offset_base_attn, (T_len, BT), (H*BT, 1), (i_t*BT, 0), (BT, BT), (1, 0)) - p_dAak = tl.make_block_ptr(dAak + offset_base_attn, (T_len, BT), (H*BT, 1), (i_t*BT, 0), (BT, BT), (1, 0)) - p_dAab = tl.make_block_ptr(dAab + offset_base_attn, (T_len, BT), (H*BT, 1), (i_t*BT, 0), (BT, BT), (1, 0)) + o_A = tl.arange(0, BT) + m_dA = m_t[:, None] & (o_A[None, :] < BT) + p_dAqk = dAqk + offset_base_attn + o_t[:, None] * (H*BT) + o_A[None, :] + p_dAqb = dAqb + offset_base_attn + o_t[:, None] * (H*BT) + o_A[None, :] + p_dAak = dAak + offset_base_attn + o_t[:, None] * (H*BT) + o_A[None, :] + p_dAab = dAab + offset_base_attn + o_t[:, None] * (H*BT) + o_A[None, :] - b_dAqk = tl.load(p_dAqk, boundary_check=(0, 1)) - b_dAqb = tl.load(p_dAqb, boundary_check=(0, 1)) - b_dAak = tl.load(p_dAak, boundary_check=(0, 1)) - b_dAab = tl.load(p_dAab, boundary_check=(0, 1)) + b_dAqk = tl.load(p_dAqk, mask=m_dA, other=0.0) + b_dAqb = tl.load(p_dAqb, mask=m_dA, other=0.0) + b_dAak = tl.load(p_dAak, mask=m_dA, other=0.0) + b_dAab = tl.load(p_dAab, mask=m_dA, other=0.0) offs_n = tl.arange(0, BT) offs_m = tl.arange(0, BT) @@ -350,15 +362,15 @@ def chunk_dplr_bwd_kernel_intra_tensorcore( b_db_intra = tl.dot(b_dAqb_T, q_ops_h) + tl.dot(b_dAab_T, a_ops) # Inter-chunk gradients loading - p_dqg = tl.make_block_ptr(dqg + offset_base_k, (T_len, K), (H*K, 1), (i_t*BT, i_k*BK), (BT, BK), (1, 0)) - p_dkg = tl.make_block_ptr(dkg + offset_base_k, (T_len, K), (H*K, 1), (i_t*BT, i_k*BK), (BT, BK), (1, 0)) - p_dag = tl.make_block_ptr(dag + offset_base_k, (T_len, K), (H*K, 1), (i_t*BT, i_k*BK), (BT, BK), (1, 0)) - p_dbg = tl.make_block_ptr(dbg + offset_base_k, (T_len, K), (H*K, 1), (i_t*BT, i_k*BK), (BT, BK), (1, 0)) + p_dqg = dqg + offset_base_k + o_t[:, None] * (H*K) + o_k[None, :] + p_dkg = dkg + offset_base_k + o_t[:, None] * (H*K) + o_k[None, :] + p_dag = dag + offset_base_k + o_t[:, None] * (H*K) + o_k[None, :] + p_dbg = dbg + offset_base_k + o_t[:, None] * (H*K) + o_k[None, :] - b_dqg = tl.load(p_dqg, boundary_check=(0, 1)) - b_dag = tl.load(p_dag, boundary_check=(0, 1)) - b_dkg = tl.load(p_dkg, boundary_check=(0, 1)) - b_dbg = tl.load(p_dbg, boundary_check=(0, 1)) + b_dqg = tl.load(p_dqg, mask=m_kk, other=0.0) + b_dag = tl.load(p_dag, mask=m_kk, other=0.0) + b_dkg = tl.load(p_dkg, mask=m_kk, other=0.0) + b_dbg = tl.load(p_dbg, mask=m_kk, other=0.0) last_idx = min((i_t+1) * BT, T_len) - 1 p_g_last = gi + offset_base_k + last_idx * H * K + tl.arange(0, BK) + i_k * BK @@ -373,19 +385,19 @@ def chunk_dplr_bwd_kernel_intra_tensorcore( b_dgk = (b_dq * b_q + b_da * b_a - b_dk * b_k - b_db * b_b) b_dgk_offset = b_da * b_a - p_dq = tl.make_block_ptr(dq + offset_base_k, (T_len, K), (H*K, 1), (i_t*BT, i_k*BK), (BT, BK), (1, 0)) - p_dk = tl.make_block_ptr(dk + offset_base_k, (T_len, K), (H*K, 1), (i_t*BT, i_k*BK), (BT, BK), (1, 0)) - p_da = tl.make_block_ptr(da + offset_base_k, (T_len, K), (H*K, 1), (i_t*BT, i_k*BK), (BT, BK), (1, 0)) - p_db = tl.make_block_ptr(db + offset_base_k, (T_len, K), (H*K, 1), (i_t*BT, i_k*BK), (BT, BK), (1, 0)) - p_dgk = tl.make_block_ptr(dgk + offset_base_k, (T_len, K), (H*K, 1), (i_t*BT, i_k*BK), (BT, BK), (1, 0)) - p_dgk_offset = tl.make_block_ptr(dgk_offset + offset_base_k, (T_len, K), (H*K, 1), (i_t*BT, i_k*BK), (BT, BK), (1, 0)) + p_dq = dq + offset_base_k + o_t[:, None] * (H*K) + o_k[None, :] + p_dk = dk + offset_base_k + o_t[:, None] * (H*K) + o_k[None, :] + p_da = da + offset_base_k + o_t[:, None] * (H*K) + o_k[None, :] + p_db = db + offset_base_k + o_t[:, None] * (H*K) + o_k[None, :] + p_dgk = dgk + offset_base_k + o_t[:, None] * (H*K) + o_k[None, :] + p_dgk_offset = dgk_offset + offset_base_k + o_t[:, None] * (H*K) + o_k[None, :] - tl.store(p_dq, b_dq.to(p_dq.dtype.element_ty), boundary_check=(0, 1)) - tl.store(p_dk, b_dk.to(p_dk.dtype.element_ty), boundary_check=(0, 1)) - tl.store(p_da, b_da.to(p_da.dtype.element_ty), boundary_check=(0, 1)) - tl.store(p_db, b_db.to(p_db.dtype.element_ty), boundary_check=(0, 1)) - tl.store(p_dgk, b_dgk.to(p_dgk.dtype.element_ty), boundary_check=(0, 1)) - tl.store(p_dgk_offset, b_dgk_offset.to(p_dgk_offset.dtype.element_ty), boundary_check=(0, 1)) + tl.store(p_dq, b_dq.to(p_dq.dtype.element_ty), mask=m_kk) + tl.store(p_dk, b_dk.to(p_dk.dtype.element_ty), mask=m_kk) + tl.store(p_da, b_da.to(p_da.dtype.element_ty), mask=m_kk) + tl.store(p_db, b_db.to(p_db.dtype.element_ty), mask=m_kk) + tl.store(p_dgk, b_dgk.to(p_dgk.dtype.element_ty), mask=m_kk) + tl.store(p_dgk_offset, b_dgk_offset.to(p_dgk_offset.dtype.element_ty), mask=m_kk) @triton.heuristics({ @@ -416,18 +428,18 @@ def chunk_dplr_bwd_dgk_kernel( BK: tl.constexpr, IS_VARLEN: tl.constexpr, ): - i_t, i_k, i_bh = tl.program_id(0), tl.program_id(1), tl.program_id(2) + i_t, i_k, i_bh = tl.program_id(0).to(tl.int64), tl.program_id(1).to(tl.int64), tl.program_id(2).to(tl.int64) i_b, i_h = i_bh // H, i_bh % H 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.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) + i_n, i_t = tl.load(chunk_indices + i_t * 2).to(tl.int32), tl.load(chunk_indices + i_t * 2 + 1).to(tl.int64) + bos, eos = tl.load(cu_seqlens + i_n).to(tl.int64), tl.load(cu_seqlens + i_n + 1).to(tl.int64) T = eos - bos NT = tl.cdiv(T, BT) else: NT = tl.cdiv(T, BT) - i_tg = (i_b * NT + i_t).to(tl.int32) - bos, eos = (i_b * T).to(tl.int32), (i_b * T + T).to(tl.int32) + i_tg = i_b * NT + i_t + bos, eos = i_b * T, i_b * T + T stride_qk = H * K dgk += (bos * H + i_h) * K @@ -437,17 +449,20 @@ def chunk_dplr_bwd_dgk_kernel( p_dgk_last = dgk_last + tl.arange(0, BK) + i_k * BK m_k = tl.arange(0, BK) + i_k * BK < K b_dgk_last = tl.load(p_dgk_last, mask=m_k, other=0) - p_dgk_offset = tl.make_block_ptr(dgk_offset, (T, K), (stride_qk, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) - p_dgk = tl.make_block_ptr(dgk, (T, K), (stride_qk, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) - b_dgk = tl.load(p_dgk, boundary_check=(0, 1)) - b_dgk_offset = tl.load(p_dgk_offset, boundary_check=(0, 1)) + o_t = i_t * BT + tl.arange(0, BT) + o_k = tl.arange(0, BK) + i_k * BK + m_kk = (o_t < T)[:, None] & (o_k[None, :] < K) + p_dgk_offset = dgk_offset + o_t[:, None] * stride_qk + o_k[None, :] + p_dgk = dgk + o_t[:, None] * stride_qk + o_k[None, :] + b_dgk = tl.load(p_dgk, mask=m_kk, other=0.0) + b_dgk_offset = tl.load(p_dgk_offset, mask=m_kk, other=0.0) # m_inv_cumsum = (tl.arange(0, BT)[:, None] <= tl.arange(0, BT)[None, :]).to(tl.float32) # b_dgk_cumsum = tl.dot(m_inv_cumsum, b_dgk, allow_tf32=False) b_dgk_cumsum = tl.cumsum(b_dgk, 0, reverse=True) b_dgk_cumsum += b_dgk_last[None, :] b_dgk_cumsum -= b_dgk_offset - p_dgk_output = tl.make_block_ptr(dgk_output, (T, K), (stride_qk, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) - tl.store(p_dgk_output, b_dgk_cumsum.to(p_dgk_output.dtype.element_ty), boundary_check=(0, 1)) + p_dgk_output = dgk_output + o_t[:, None] * stride_qk + o_k[None, :] + tl.store(p_dgk_output, b_dgk_cumsum.to(p_dgk_output.dtype.element_ty), mask=m_kk) def chunk_dplr_bwd_dqk_intra( @@ -488,7 +503,7 @@ def chunk_dplr_bwd_dqk_intra( dgk = torch.empty_like(gi, dtype=torch.float) dgk_offset = torch.empty_like(gi, dtype=torch.float) - grid = (NK, NT, B * H) + grid = (NK * NT, B * H) if safe_gate: chunk_dplr_bwd_kernel_intra_func = chunk_dplr_bwd_kernel_intra_tensorcore else: diff --git a/build/torch-cuda/ops/generalized_delta_rule/dplr/chunk_A_fwd.py b/build/torch-cuda/ops/generalized_delta_rule/dplr/chunk_A_fwd.py index b37433aad93359578bb2708ae36b8a2a02dd89b2..6c2da32cdf6989e93f807bef6d4b5c2258fb68d9 100644 --- a/build/torch-cuda/ops/generalized_delta_rule/dplr/chunk_A_fwd.py +++ b/build/torch-cuda/ops/generalized_delta_rule/dplr/chunk_A_fwd.py @@ -56,11 +56,11 @@ def chunk_dplr_fwd_A_kernel_intra_sub_intra( IS_VARLEN: tl.constexpr, GATHER_SUPPORTED: tl.constexpr, ): - i_t, i_b, i_h = tl.program_id(0), tl.program_id(1), tl.program_id(2) + i_t, i_b, i_h = tl.program_id(0).to(tl.int64), tl.program_id(1).to(tl.int64), tl.program_id(2).to(tl.int64) 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.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) + i_n, i_t = tl.load(chunk_indices + i_t * 2).to(tl.int32), tl.load(chunk_indices + i_t * 2 + 1).to(tl.int64) + bos, eos = tl.load(cu_seqlens + i_n).to(tl.int64), tl.load(cu_seqlens + i_n + 1).to(tl.int64) T = eos - bos else: bos, eos = i_b * T, i_b * T + T @@ -69,31 +69,33 @@ def chunk_dplr_fwd_A_kernel_intra_sub_intra( return o_i = tl.arange(0, BC) + o_c = i_t * BT + o_i o_k = tl.arange(0, BK) m_k = o_k < K + m_kk = ((i_t * BT + tl.arange(0, BC)) < T)[:, None] & m_k[None, :] m_A = (i_t * BT + tl.arange(0, BC)) < T last_idx = min((i_t+1) * BT, T) - 1 o_A = (bos + i_t * BT + tl.arange(0, BC)) * H*BT + i_h * BT - p_q = tl.make_block_ptr(q + (bos * H + i_h) * K, (T, K), (H*K, 1), (i_t * BT, 0), (BC, BK), (1, 0)) - p_k = tl.make_block_ptr(k + (bos * H + i_h) * K, (T, K), (H*K, 1), (i_t * BT, 0), (BC, BK), (1, 0)) - p_a = tl.make_block_ptr(a + (bos * H + i_h) * K, (T, K), (H*K, 1), (i_t * BT, 0), (BC, BK), (1, 0)) - p_b = tl.make_block_ptr(b + (bos * H + i_h) * K, (T, K), (H*K, 1), (i_t * BT, 0), (BC, BK), (1, 0)) - p_gi = tl.make_block_ptr(gi + (bos * H + i_h) * K, (T, K), (H*K, 1), (i_t * BT, 0), (BC, BK), (1, 0)) - p_ge = tl.make_block_ptr(ge + (bos * H + i_h) * K, (T, K), (H*K, 1), (i_t * BT, 0), (BC, BK), (1, 0)) + p_q = q + (bos * H + i_h) * K + o_c[:, None] * (H*K) + o_k[None, :] + p_k = k + (bos * H + i_h) * K + o_c[:, None] * (H*K) + o_k[None, :] + p_a = a + (bos * H + i_h) * K + o_c[:, None] * (H*K) + o_k[None, :] + p_b = b + (bos * H + i_h) * K + o_c[:, None] * (H*K) + o_k[None, :] + p_gi = gi + (bos * H + i_h) * K + o_c[:, None] * (H*K) + o_k[None, :] + p_ge = ge + (bos * H + i_h) * K + o_c[:, None] * (H*K) + o_k[None, :] p_g_last = gi + (bos * H + i_h) * K + last_idx * H * K + tl.arange(0, BK) b_g_last = tl.load(p_g_last, mask=m_k, other=0) - p_qg = tl.make_block_ptr(qg + (bos * H + i_h) * K, (T, K), (H*K, 1), (i_t * BT, 0), (BC, BK), (1, 0)) - p_kg = tl.make_block_ptr(kg + (bos * H + i_h) * K, (T, K), (H*K, 1), (i_t * BT, 0), (BC, BK), (1, 0)) - p_ag = tl.make_block_ptr(ag + (bos * H + i_h) * K, (T, K), (H*K, 1), (i_t * BT, 0), (BC, BK), (1, 0)) - p_bg = tl.make_block_ptr(bg + (bos * H + i_h) * K, (T, K), (H*K, 1), (i_t * BT, 0), (BC, BK), (1, 0)) + p_qg = qg + (bos * H + i_h) * K + o_c[:, None] * (H*K) + o_k[None, :] + p_kg = kg + (bos * H + i_h) * K + o_c[:, None] * (H*K) + o_k[None, :] + p_ag = ag + (bos * H + i_h) * K + o_c[:, None] * (H*K) + o_k[None, :] + p_bg = bg + (bos * H + i_h) * K + o_c[:, None] * (H*K) + o_k[None, :] - b_q = tl.load(p_q, boundary_check=(0, 1)) + b_q = tl.load(p_q, mask=m_kk, other=0.0) b_q = b_q * scale - b_k = tl.load(p_k, boundary_check=(0, 1)) - b_a = tl.load(p_a, boundary_check=(0, 1)) - b_b = tl.load(p_b, boundary_check=(0, 1)) - b_gi = tl.load(p_gi, boundary_check=(0, 1)).to(tl.float32) - b_ge = tl.load(p_ge, boundary_check=(0, 1)).to(tl.float32) + b_k = tl.load(p_k, mask=m_kk, other=0.0) + b_a = tl.load(p_a, mask=m_kk, other=0.0) + b_b = tl.load(p_b, mask=m_kk, other=0.0) + b_gi = tl.load(p_gi, mask=m_kk, other=0.0).to(tl.float32) + b_ge = tl.load(p_ge, mask=m_kk, other=0.0).to(tl.float32) # deal with decay term. g_exp = exp2(b_gi) @@ -102,10 +104,10 @@ def chunk_dplr_fwd_A_kernel_intra_sub_intra( b_kg = b_k * g_exp_inv b_bg = b_b * g_exp_inv b_ag = b_a * exp2(b_ge) - tl.store(p_qg, b_qg.to(p_qg.dtype.element_ty, fp_downcast_rounding="rtne"), boundary_check=(0, 1)) - tl.store(p_bg, b_bg.to(p_bg.dtype.element_ty, fp_downcast_rounding="rtne"), boundary_check=(0, 1)) - tl.store(p_ag, b_ag.to(p_ag.dtype.element_ty, fp_downcast_rounding="rtne"), boundary_check=(0, 1)) - tl.store(p_kg, b_kg.to(p_kg.dtype.element_ty, fp_downcast_rounding="rtne"), boundary_check=(0, 1)) + tl.store(p_qg, b_qg.to(p_qg.dtype.element_ty, fp_downcast_rounding="rtne"), mask=m_kk) + tl.store(p_bg, b_bg.to(p_bg.dtype.element_ty, fp_downcast_rounding="rtne"), mask=m_kk) + tl.store(p_ag, b_ag.to(p_ag.dtype.element_ty, fp_downcast_rounding="rtne"), mask=m_kk) + tl.store(p_kg, b_kg.to(p_kg.dtype.element_ty, fp_downcast_rounding="rtne"), mask=m_kk) # tl.debug_barrier() b_q = b_q.to(b_k.dtype) @@ -182,11 +184,11 @@ def chunk_dplr_fwd_A_kernel_intra_tensorcore( IS_VARLEN: tl.constexpr, GATHER_SUPPORTED: tl.constexpr, ): - i_t, i_b, i_h = tl.program_id(0), tl.program_id(1), tl.program_id(2) + i_t, i_b, i_h = tl.program_id(0).to(tl.int64), tl.program_id(1).to(tl.int64), tl.program_id(2).to(tl.int64) 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.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) + i_n, i_t = tl.load(chunk_indices + i_t * 2).to(tl.int32), tl.load(chunk_indices + i_t * 2 + 1).to(tl.int64) + bos, eos = tl.load(cu_seqlens + i_n).to(tl.int64), tl.load(cu_seqlens + i_n + 1).to(tl.int64) T_len = eos - bos else: bos = i_b * T @@ -199,19 +201,23 @@ def chunk_dplr_fwd_A_kernel_intra_tensorcore( offset_base = (bos * H + i_h) * K # Load the current chunk of Q, K, A, B and their gates - p_q = tl.make_block_ptr(q + offset_base, (T_len, K), (H*K, 1), (i_t * BT, 0), (BT, BK), (1, 0)) - p_k = tl.make_block_ptr(k + offset_base, (T_len, K), (H*K, 1), (i_t * BT, 0), (BT, BK), (1, 0)) - p_a = tl.make_block_ptr(a + offset_base, (T_len, K), (H*K, 1), (i_t * BT, 0), (BT, BK), (1, 0)) - p_b = tl.make_block_ptr(b + offset_base, (T_len, K), (H*K, 1), (i_t * BT, 0), (BT, BK), (1, 0)) - p_gi = tl.make_block_ptr(gi + offset_base, (T_len, K), (H*K, 1), (i_t * BT, 0), (BT, BK), (1, 0)) - p_ge = tl.make_block_ptr(ge + offset_base, (T_len, K), (H*K, 1), (i_t * BT, 0), (BT, BK), (1, 0)) - - b_q = tl.load(p_q, boundary_check=(0, 1)) - b_k = tl.load(p_k, boundary_check=(0, 1)) - b_a = tl.load(p_a, boundary_check=(0, 1)) - b_b = tl.load(p_b, boundary_check=(0, 1)) - b_gi_val = tl.load(p_gi, boundary_check=(0, 1)).to(tl.float32) - b_ge_val = tl.load(p_ge, boundary_check=(0, 1)).to(tl.float32) + o_t = i_t * BT + tl.arange(0, BT) + o_b = tl.arange(0, BK) + m_t = o_t < T_len + m_kb = m_t[:, None] & (o_b[None, :] < K) + p_q = q + offset_base + o_t[:, None] * (H*K) + o_b[None, :] + p_k = k + offset_base + o_t[:, None] * (H*K) + o_b[None, :] + p_a = a + offset_base + o_t[:, None] * (H*K) + o_b[None, :] + p_b = b + offset_base + o_t[:, None] * (H*K) + o_b[None, :] + p_gi = gi + offset_base + o_t[:, None] * (H*K) + o_b[None, :] + p_ge = ge + offset_base + o_t[:, None] * (H*K) + o_b[None, :] + + b_q = tl.load(p_q, mask=m_kb, other=0.0) + b_k = tl.load(p_k, mask=m_kb, other=0.0) + b_a = tl.load(p_a, mask=m_kb, other=0.0) + b_b = tl.load(p_b, mask=m_kb, other=0.0) + b_gi_val = tl.load(p_gi, mask=m_kb, other=0.0).to(tl.float32) + b_ge_val = tl.load(p_ge, mask=m_kb, other=0.0).to(tl.float32) # Calculate the index of the middle element of the valid part of the chunk valid_len = min(T_len - i_t * BT, BT) @@ -220,7 +226,7 @@ def chunk_dplr_fwd_A_kernel_intra_tensorcore( # Load the offset vector from Global Memory # p_offset points to gi[i_t*BT + mid_idx, :] m_k = tl.arange(0, BK) < K - p_offset = gi + offset_base + (i_t * BT + mid_idx) * K + tl.arange(0, BK) + p_offset = gi + offset_base + (i_t * BT + mid_idx) * H * K + tl.arange(0, BK) b_offset = tl.load(p_offset, mask=m_k, other=0.0).to(tl.float32) # Apply offset to gate values @@ -254,20 +260,20 @@ def chunk_dplr_fwd_A_kernel_intra_tensorcore( exp_g_centered = exp2(b_g_centered) # Create pointers for writing - p_qg = tl.make_block_ptr(qg + offset_base, (T_len, K), (H*K, 1), (i_t * BT, 0), (BT, BK), (1, 0)) - p_kg = tl.make_block_ptr(kg + offset_base, (T_len, K), (H*K, 1), (i_t * BT, 0), (BT, BK), (1, 0)) - p_ag = tl.make_block_ptr(ag + offset_base, (T_len, K), (H*K, 1), (i_t * BT, 0), (BT, BK), (1, 0)) - p_bg = tl.make_block_ptr(bg + offset_base, (T_len, K), (H*K, 1), (i_t * BT, 0), (BT, BK), (1, 0)) + p_qg = qg + offset_base + o_t[:, None] * (H*K) + o_b[None, :] + p_kg = kg + offset_base + o_t[:, None] * (H*K) + o_b[None, :] + p_ag = ag + offset_base + o_t[:, None] * (H*K) + o_b[None, :] + p_bg = bg + offset_base + o_t[:, None] * (H*K) + o_b[None, :] # Store gated Q and A - tl.store(p_qg, (q_ops * exp_offset[None, :]).to(p_qg.dtype.element_ty), boundary_check=(0, 1)) - tl.store(p_ag, (a_ops * exp_offset[None, :]).to(p_ag.dtype.element_ty), boundary_check=(0, 1)) + tl.store(p_qg, (q_ops * exp_offset[None, :]).to(p_qg.dtype.element_ty), mask=m_kb) + tl.store(p_ag, (a_ops * exp_offset[None, :]).to(p_ag.dtype.element_ty), mask=m_kb) # Store gated K and B b_kg_g = k_ops * exp_g_centered[None, :] b_bg_g = b_ops * exp_g_centered[None, :] - tl.store(p_kg, b_kg_g.to(p_kg.dtype.element_ty), boundary_check=(0, 1)) - tl.store(p_bg, b_bg_g.to(p_bg.dtype.element_ty), boundary_check=(0, 1)) + tl.store(p_kg, b_kg_g.to(p_kg.dtype.element_ty), mask=m_kb) + tl.store(p_bg, b_bg_g.to(p_bg.dtype.element_ty), mask=m_kb) # Transpose K and B for dot product k_ops_t = tl.trans(k_ops) @@ -295,15 +301,17 @@ def chunk_dplr_fwd_A_kernel_intra_tensorcore( # Store the intra-chunk attention matrices offset_out_base = (bos * H + i_h) * BT - p_Aqk = tl.make_block_ptr(Aqk + offset_out_base, (T_len, BT), (H*BT, 1), (i_t * BT, 0), (BT, BT), (1, 0)) - p_Aqb = tl.make_block_ptr(Aqb + offset_out_base, (T_len, BT), (H*BT, 1), (i_t * BT, 0), (BT, BT), (1, 0)) - p_Aak = tl.make_block_ptr(Aak + offset_out_base, (T_len, BT), (H*BT, 1), (i_t * BT, 0), (BT, BT), (1, 0)) - p_Aab = tl.make_block_ptr(Aab + offset_out_base, (T_len, BT), (H*BT, 1), (i_t * BT, 0), (BT, BT), (1, 0)) - - tl.store(p_Aqk, b_A_qk.to(p_Aqk.dtype.element_ty), boundary_check=(0, 1)) - tl.store(p_Aqb, b_A_qb.to(p_Aqb.dtype.element_ty), boundary_check=(0, 1)) - tl.store(p_Aak, b_A_ak.to(p_Aak.dtype.element_ty), boundary_check=(0, 1)) - tl.store(p_Aab, b_A_ab.to(p_Aab.dtype.element_ty), boundary_check=(0, 1)) + o_A = tl.arange(0, BT) + m_AO = m_t[:, None] & (o_A[None, :] < BT) + p_Aqk = Aqk + offset_out_base + o_t[:, None] * (H*BT) + o_A[None, :] + p_Aqb = Aqb + offset_out_base + o_t[:, None] * (H*BT) + o_A[None, :] + p_Aak = Aak + offset_out_base + o_t[:, None] * (H*BT) + o_A[None, :] + p_Aab = Aab + offset_out_base + o_t[:, None] * (H*BT) + o_A[None, :] + + tl.store(p_Aqk, b_A_qk.to(p_Aqk.dtype.element_ty), mask=m_AO) + tl.store(p_Aqb, b_A_qb.to(p_Aqb.dtype.element_ty), mask=m_AO) + tl.store(p_Aak, b_A_ak.to(p_Aak.dtype.element_ty), mask=m_AO) + tl.store(p_Aab, b_A_ab.to(p_Aab.dtype.element_ty), mask=m_AO) def chunk_dplr_fwd_intra( diff --git a/build/torch-cuda/ops/generalized_delta_rule/dplr/chunk_h_bwd.py b/build/torch-cuda/ops/generalized_delta_rule/dplr/chunk_h_bwd.py index 167a651114c69f1992ccedd6fcda7c46737e09c7..00fa68310b6182a88f2a1a0bed80a05fc0f7e6c3 100644 --- a/build/torch-cuda/ops/generalized_delta_rule/dplr/chunk_h_bwd.py +++ b/build/torch-cuda/ops/generalized_delta_rule/dplr/chunk_h_bwd.py @@ -10,7 +10,7 @@ import triton import triton.language as tl from ....ops.utils import prepare_chunk_indices, prepare_chunk_offsets -from ....ops.utils.op import exp2 +from ....ops.utils.op import exp2, unflatten_program_id from ....utils import IS_AMD, autotune_cache_kwargs, check_shared_mem NUM_WARPS_AUTOTUNE = [2, 4, 8, 16] if IS_AMD else [2, 4, 8, 16, 32] @@ -56,57 +56,65 @@ def chunk_dplr_bwd_kernel_dhu( USE_INITIAL_STATE: tl.constexpr, IS_VARLEN: tl.constexpr, ): - i_k, i_v, i_nh = tl.program_id(0), tl.program_id(1), tl.program_id(2) + i_k, i_v, i_nh = unflatten_program_id(tl.cdiv(K, BK), tl.cdiv(V, BV)) i_n, i_h = i_nh // H, i_nh % H if IS_VARLEN: - bos, eos = tl.load(cu_seqlens + i_n).to(tl.int32), tl.load(cu_seqlens + i_n + 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 NT = tl.cdiv(T, BT) - boh = tl.load(chunk_offsets + i_n).to(tl.int32) + boh = tl.load(chunk_offsets + i_n).to(tl.int64) else: bos, eos = i_n * T, i_n * T + T NT = tl.cdiv(T, BT) boh = i_n * NT + o_k = i_k * BK + tl.arange(0, BK) + o_v = i_v * BV + tl.arange(0, BV) + m_h = (o_k[:, None] < K) & (o_v[None, :] < V) # [BK, BV] b_dh = tl.zeros([BK, BV], dtype=tl.float32) if USE_FINAL_STATE_GRADIENT: - p_dht = tl.make_block_ptr(dht + i_nh * K*V, (K, V), (V, 1), (i_k * BK, i_v * BV), (BK, BV), (1, 0)) - b_dh += tl.load(p_dht, boundary_check=(0, 1)) + p_dht = dht + i_nh * K*V + o_k[:, None] * V + o_v[None, :] + b_dh += tl.load(p_dht, mask=m_h, other=0.0) mask_k = tl.arange(0, BK) < K for i_t in range(NT - 1, -1, -1): - p_dh = tl.make_block_ptr(dh + ((boh+i_t) * H + i_h) * K*V, (K, V), (V, 1), (i_k * BK, i_v * BV), (BK, BV), (1, 0)) - tl.store(p_dh, b_dh.to(p_dh.dtype.element_ty), boundary_check=(0, 1)) + p_dh = dh + ((boh+i_t) * H + i_h) * K*V + o_k[:, None] * V + o_v[None, :] + tl.store(p_dh, b_dh.to(p_dh.dtype.element_ty), mask=m_h) b_dh_tmp = tl.zeros([BK, BV], dtype=tl.float32) for i_c in range(tl.cdiv(BT, BC) - 1, -1, -1): - p_qg = tl.make_block_ptr(qg+(bos*H+i_h)*K, (K, T), (1, H*K), (i_k * BK, i_t * BT + i_c * BC), (BK, BC), (0, 1)) - p_bg = tl.make_block_ptr(bg+(bos*H+i_h)*K, (T, K), (H*K, 1), (i_t * BT + i_c * BC, i_k * BK), (BC, BK), (1, 0)) - p_w = tl.make_block_ptr(w+(bos*H+i_h)*K, (K, T), (1, H*K), (i_k * BK, i_t * BT + i_c * BC), (BK, BC), (0, 1)) - p_dv = tl.make_block_ptr(dv+(bos*H+i_h)*V, (T, V), (H*V, 1), (i_t*BT + i_c * BC, i_v * BV), (BC, BV), (1, 0)) - p_do = tl.make_block_ptr(do+(bos*H+i_h)*V, (T, V), (H*V, 1), (i_t*BT + i_c * BC, i_v * BV), (BC, BV), (1, 0)) - p_dv2 = tl.make_block_ptr(dv2+(bos*H+i_h)*V, (T, V), (H*V, 1), (i_t*BT + i_c * BC, i_v * BV), (BC, BV), (1, 0)) + o_c = i_t.to(tl.int64) * BT + i_c * BC + tl.arange(0, BC) + m_c = o_c < T + m_kc = (o_k[:, None] < K) & m_c[None, :] + m_bgc = m_c[:, None] & (o_k[None, :] < K) + m_vc = m_c[:, None] & (o_v[None, :] < V) + p_qg = qg+(bos*H+i_h)*K + o_k[:, None] + o_c[None, :] * (H*K) + p_bg = bg+(bos*H+i_h)*K + o_c[:, None] * (H*K) + o_k[None, :] + p_w = w+(bos*H+i_h)*K + o_k[:, None] + o_c[None, :] * (H*K) + p_dv = dv+(bos*H+i_h)*V + o_c[:, None] * (H*V) + o_v[None, :] + p_do = do+(bos*H+i_h)*V + o_c[:, None] * (H*V) + o_v[None, :] + p_dv2 = dv2+(bos*H+i_h)*V + o_c[:, None] * (H*V) + o_v[None, :] # [BK, BT] - b_qg = tl.load(p_qg, boundary_check=(0, 1)) + b_qg = tl.load(p_qg, mask=m_kc, other=0.0) # [BT, BK] - b_bg = tl.load(p_bg, boundary_check=(0, 1)) - b_w = tl.load(p_w, boundary_check=(0, 1)) + b_bg = tl.load(p_bg, mask=m_bgc, other=0.0) + b_w = tl.load(p_w, mask=m_kc, other=0.0) # [BT, V] - b_do = tl.load(p_do, boundary_check=(0, 1)) - b_dv = tl.load(p_dv, boundary_check=(0, 1)) + b_do = tl.load(p_do, mask=m_vc, other=0.0) + b_dv = tl.load(p_dv, mask=m_vc, other=0.0) b_dv2 = b_dv + tl.dot(b_bg, b_dh.to(b_bg.dtype)) - tl.store(p_dv2, b_dv2.to(p_dv.dtype.element_ty), boundary_check=(0, 1)) + tl.store(p_dv2, b_dv2.to(p_dv.dtype.element_ty), mask=m_vc) # [BK, BV] - b_dh_tmp += tl.dot(b_qg, b_do.to(b_qg.dtype)) - b_dh_tmp += tl.dot(b_w, b_dv2.to(b_qg.dtype)) + b_dh_tmp = tl.dot(b_qg, b_do.to(b_qg.dtype), b_dh_tmp) + b_dh_tmp = tl.dot(b_w, b_dv2.to(b_qg.dtype), b_dh_tmp) last_idx = min((i_t + 1) * BT, T) - 1 bg_last = tl.load(gk + ((bos + last_idx) * H + i_h) * K + tl.arange(0, BK), mask=mask_k) b_dh *= exp2(bg_last)[:, None] b_dh += b_dh_tmp if USE_INITIAL_STATE: - p_dh0 = tl.make_block_ptr(dh0 + i_nh * K*V, (K, V), (V, 1), (i_k * BK, i_v * BV), (BK, BV), (1, 0)) - tl.store(p_dh0, b_dh.to(p_dh0.dtype.element_ty), boundary_check=(0, 1)) + p_dh0 = dh0 + i_nh * K*V + o_k[:, None] * V + o_v[None, :] + tl.store(p_dh0, b_dh.to(p_dh0.dtype.element_ty), mask=m_h) def chunk_dplr_bwd_dhu( @@ -153,7 +161,7 @@ def chunk_dplr_bwd_dhu( dh0 = torch.empty_like(h0, dtype=torch.float32) if h0 is not None else None dv2 = torch.zeros_like(dv) - grid = (NK, NV, N * H) + grid = (NK * NV * N * H,) chunk_dplr_bwd_kernel_dhu[grid]( qg=qg, bg=bg, diff --git a/build/torch-cuda/ops/generalized_delta_rule/dplr/chunk_h_fwd.py b/build/torch-cuda/ops/generalized_delta_rule/dplr/chunk_h_fwd.py index 050ae739a549bc02830bba1e6f97f607d2ac5406..2569874218fba1f6cb66c5eeea03d66bd552ce52 100644 --- a/build/torch-cuda/ops/generalized_delta_rule/dplr/chunk_h_fwd.py +++ b/build/torch-cuda/ops/generalized_delta_rule/dplr/chunk_h_fwd.py @@ -10,7 +10,7 @@ import triton import triton.language as tl from ....ops.utils import prepare_chunk_indices, prepare_chunk_offsets -from ....ops.utils.op import exp2 +from ....ops.utils.op import exp2, unflatten_program_id from ....utils import IS_AMD, autotune_cache_kwargs, check_shared_mem NUM_WARPS_AUTOTUNE = [2, 4, 8, 16] if IS_AMD else [2, 4, 8, 16, 32] @@ -56,47 +56,54 @@ def chunk_dplr_fwd_kernel_h( STORE_FINAL_STATE: tl.constexpr, IS_VARLEN: tl.constexpr, ): - i_k, i_v, i_nh = tl.program_id(0), tl.program_id(1), tl.program_id(2) + i_k, i_v, i_nh = unflatten_program_id(tl.cdiv(K, BK), tl.cdiv(V, BV)) i_n, i_h = i_nh // H, i_nh % H if IS_VARLEN: - bos, eos = tl.load(cu_seqlens + i_n).to(tl.int32), tl.load(cu_seqlens + i_n + 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 NT = tl.cdiv(T, BT) - boh = tl.load(chunk_offsets + i_n).to(tl.int32) + boh = tl.load(chunk_offsets + i_n).to(tl.int64) else: bos, eos = i_n * T, i_n * T + T NT = tl.cdiv(T, BT) boh = i_n * NT o_k = i_k * BK + tl.arange(0, BK) + o_v = i_v * BV + tl.arange(0, BV) + m_h = (o_k[:, None] < K) & (o_v[None, :] < V) # [BK, BV] b_h = tl.zeros([BK, BV], dtype=tl.float32) if USE_INITIAL_STATE: - p_h0 = tl.make_block_ptr(h0 + i_nh * K*V, (K, V), (V, 1), (i_k * BK, i_v * BV), (BK, BV), (1, 0)) - b_h = tl.load(p_h0, boundary_check=(0, 1)).to(tl.float32) + p_h0 = h0 + i_nh * K*V + o_k[:, None] * V + o_v[None, :] + b_h = tl.load(p_h0, mask=m_h, other=0.0).to(tl.float32) for i_t in range(NT): - p_h = tl.make_block_ptr(h + ((boh + i_t) * H + i_h) * K*V, (K, V), (V, 1), (i_k * BK, i_v * BV), (BK, BV), (1, 0)) - tl.store(p_h, b_h.to(p_h.dtype.element_ty), boundary_check=(0, 1)) + p_h = h + ((boh + i_t) * H + i_h) * K*V + o_k[:, None] * V + o_v[None, :] + tl.store(p_h, b_h.to(p_h.dtype.element_ty), mask=m_h) b_hc = tl.zeros([BK, BV], dtype=tl.float32) - # since we need to make all DK in the SRAM. we face serve SRAM memory burden. By subchunking we allievate such burden + # since all DK values must fit in SRAM, subchunking alleviates the SRAM memory burden. for i_c in range(tl.cdiv(min(BT, T - i_t * BT), BC)): - p_kg = tl.make_block_ptr(kg+(bos*H+i_h)*K, (K, T), (1, H*K), (i_k * BK, i_t * BT + i_c * BC), (BK, BC), (0, 1)) - p_bg = tl.make_block_ptr(bg+(bos*H+i_h)*K, (K, T), (1, H*K), (i_k * BK, i_t * BT + i_c * BC), (BK, BC), (0, 1)) - p_w = tl.make_block_ptr(w+(bos*H+i_h)*K, (T, K), (H*K, 1), (i_t * BT + i_c * BC, i_k * BK), (BC, BK), (1, 0)) - p_v = tl.make_block_ptr(v+(bos*H+i_h)*V, (T, V), (H*V, 1), (i_t * BT + i_c * BC, i_v * BV), (BC, BV), (1, 0)) - p_u = tl.make_block_ptr(u+(bos*H+i_h)*V, (T, V), (H*V, 1), (i_t * BT + i_c * BC, i_v * BV), (BC, BV), (1, 0)) - p_v_new = tl.make_block_ptr(v_new+(bos*H+i_h)*V, (T, V), (H*V, 1), (i_t*BT+i_c*BC, i_v * BV), (BC, BV), (1, 0)) + o_c = i_t.to(tl.int64) * BT + i_c * BC + tl.arange(0, BC) + m_c = o_c < T + m_kc = (o_k[:, None] < K) & m_c[None, :] + m_wc = m_c[:, None] & (o_k[None, :] < K) + m_vc = m_c[:, None] & (o_v[None, :] < V) + p_kg = kg+(bos*H+i_h)*K + o_k[:, None] + o_c[None, :] * (H*K) + p_bg = bg+(bos*H+i_h)*K + o_k[:, None] + o_c[None, :] * (H*K) + p_w = w+(bos*H+i_h)*K + o_c[:, None] * (H*K) + o_k[None, :] + p_v = v+(bos*H+i_h)*V + o_c[:, None] * (H*V) + o_v[None, :] + p_u = u+(bos*H+i_h)*V + o_c[:, None] * (H*V) + o_v[None, :] + p_v_new = v_new+(bos*H+i_h)*V + o_c[:, None] * (H*V) + o_v[None, :] # [BK, BC] - b_kg = tl.load(p_kg, boundary_check=(0, 1)) - b_v = tl.load(p_v, boundary_check=(0, 1)) - b_w = tl.load(p_w, boundary_check=(0, 1)) - b_bg = tl.load(p_bg, boundary_check=(0, 1)) - b_v2 = tl.dot(b_w, b_h.to(b_w.dtype)) + tl.load(p_u, boundary_check=(0, 1)) - b_hc += tl.dot(b_kg, b_v) - b_hc += tl.dot(b_bg.to(b_hc.dtype), b_v2) - tl.store(p_v_new, b_v2.to(p_v_new.dtype.element_ty), boundary_check=(0, 1)) + b_kg = tl.load(p_kg, mask=m_kc, other=0.0) + b_v = tl.load(p_v, mask=m_vc, other=0.0) + b_w = tl.load(p_w, mask=m_wc, other=0.0) + b_bg = tl.load(p_bg, mask=m_kc, other=0.0) + b_v2 = tl.dot(b_w, b_h.to(b_w.dtype)) + tl.load(p_u, mask=m_vc, other=0.0) + b_hc = tl.dot(b_kg, b_v, b_hc) + b_hc = tl.dot(b_bg.to(b_hc.dtype), b_v2, b_hc) + tl.store(p_v_new, b_v2.to(p_v_new.dtype.element_ty), mask=m_vc) last_idx = min((i_t + 1) * BT, T) - 1 b_g_last = tl.load(gk + (bos + last_idx) * H*K + i_h * K + o_k, mask=o_k < K).to(tl.float32) @@ -104,8 +111,8 @@ def chunk_dplr_fwd_kernel_h( b_h += b_hc if STORE_FINAL_STATE: - p_ht = tl.make_block_ptr(ht + i_nh * K*V, (K, V), (V, 1), (i_k * BK, i_v * BV), (BK, BV), (1, 0)) - tl.store(p_ht, b_h.to(p_ht.dtype.element_ty, fp_downcast_rounding="rtne"), boundary_check=(0, 1)) + p_ht = ht + i_nh * K*V + o_k[:, None] * V + o_v[None, :] + tl.store(p_ht, b_h.to(p_ht.dtype.element_ty, fp_downcast_rounding="rtne"), mask=m_h) def chunk_dplr_fwd_h( @@ -153,7 +160,7 @@ def chunk_dplr_fwd_h( h = kg.new_empty(B, NT, H, K, V) final_state = kg.new_empty(N, H, K, V, dtype=torch.float32) if output_final_state else None v_new = torch.empty_like(u) - grid = (NK, NV, N * H) + grid = (NK * NV * N * H,) chunk_dplr_fwd_kernel_h[grid]( kg=kg, v=v, diff --git a/build/torch-cuda/ops/generalized_delta_rule/dplr/chunk_o_bwd.py b/build/torch-cuda/ops/generalized_delta_rule/dplr/chunk_o_bwd.py index be0b1b270bde765e42864198ead92903ce2189ed..9473fa7f826a5b32cfd5e82ae873fa4b77909939 100644 --- a/build/torch-cuda/ops/generalized_delta_rule/dplr/chunk_o_bwd.py +++ b/build/torch-cuda/ops/generalized_delta_rule/dplr/chunk_o_bwd.py @@ -10,7 +10,7 @@ import triton import triton.language as tl from ....ops.utils import prepare_chunk_indices -from ....ops.utils.op import exp2 +from ....ops.utils.op import exp2, unflatten_program_id from ....utils import IS_AMD, autotune_cache_kwargs, check_shared_mem NUM_WARPS_AUTOTUNE = [2, 4, 8, 16] if IS_AMD else [2, 4, 8, 16, 32] @@ -49,11 +49,11 @@ def chunk_dplr_bwd_kernel_dAu( BV: tl.constexpr, IS_VARLEN: tl.constexpr, ): - i_t, i_bh = tl.program_id(0), tl.program_id(1) + i_t, i_bh = tl.program_id(0).to(tl.int64), tl.program_id(1).to(tl.int64) i_b, i_h = i_bh // H, i_bh % H 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.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) + i_n, i_t = tl.load(chunk_indices + i_t * 2).to(tl.int32), tl.load(chunk_indices + i_t * 2 + 1).to(tl.int64) + bos, eos = tl.load(cu_seqlens + i_n).to(tl.int64), tl.load(cu_seqlens + i_n + 1).to(tl.int64) else: bos, eos = i_b * T, i_b * T + T T = eos - bos @@ -61,33 +61,40 @@ def chunk_dplr_bwd_kernel_dAu( b_dA_qk = tl.zeros([BT, BT], dtype=tl.float32) b_dA_qb = tl.zeros([BT, BT], dtype=tl.float32) - p_A_qb = tl.make_block_ptr(A_qb + (bos * H + i_h) * BT, (T, BT), (H*BT, 1), (i_t * BT, 0), (BT, BT), (1, 0)) + o_t = i_t * BT + tl.arange(0, BT) + o_A = tl.arange(0, BT) + m_t = o_t < T + m_A = m_t[:, None] & (o_A[None, :] < BT) + p_A_qb = A_qb + (bos * H + i_h) * BT + o_t[:, None] * (H*BT) + o_A[None, :] - b_A_qb = tl.load(p_A_qb, boundary_check=(0, 1)) + b_A_qb = tl.load(p_A_qb, mask=m_A, other=0.0) # causal mask b_A_qb = tl.where(tl.arange(0, BT)[:, None] >= tl.arange(0, BT)[None, :], b_A_qb, 0.).to(b_A_qb.dtype) for i_v in range(tl.cdiv(V, BV)): - p_do = tl.make_block_ptr(do + (bos*H + i_h) * V, (T, V), (H*V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) - p_v = tl.make_block_ptr(v + (bos*H + i_h) * V, (V, T), (1, H*V), (i_v * BV, i_t * BT), (BV, BT), (0, 1)) - p_v_new = tl.make_block_ptr(v_new + (bos*H + i_h) * V, (V, T), (1, H*V), (i_v * BV, i_t * BT), (BV, BT), (0, 1)) - p_dv_new = tl.make_block_ptr(dv_new + (bos*H + i_h) * V, (T, V), (H*V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) - b_v = tl.load(p_v, boundary_check=(0, 1)) - b_do = tl.load(p_do, boundary_check=(0, 1)) - b_v_new = tl.load(p_v_new, boundary_check=(0, 1)) - b_dA_qk += tl.dot(b_do, b_v) - b_dA_qb += tl.dot(b_do, b_v_new) + o_v = i_v * BV + tl.arange(0, BV) + m_v = m_t[:, None] & (o_v[None, :] < V) + m_vt = (o_v[:, None] < V) & m_t[None, :] + p_do = do + (bos*H + i_h) * V + o_t[:, None] * (H*V) + o_v[None, :] + p_v = v + (bos*H + i_h) * V + o_v[:, None] + o_t[None, :] * (H*V) + p_v_new = v_new + (bos*H + i_h) * V + o_v[:, None] + o_t[None, :] * (H*V) + p_dv_new = dv_new + (bos*H + i_h) * V + o_t[:, None] * (H*V) + o_v[None, :] + b_v = tl.load(p_v, mask=m_vt, other=0.0) + b_do = tl.load(p_do, mask=m_v, other=0.0) + b_v_new = tl.load(p_v_new, mask=m_vt, other=0.0) + b_dA_qk = tl.dot(b_do, b_v, b_dA_qk) + b_dA_qb = tl.dot(b_do, b_v_new, b_dA_qb) b_dv_new = tl.dot(tl.trans(b_A_qb), b_do) # for recurrent - tl.store(p_dv_new, b_dv_new.to(p_dv_new.dtype.element_ty), boundary_check=(0, 1)) + tl.store(p_dv_new, b_dv_new.to(p_dv_new.dtype.element_ty), mask=m_v) - p_dA_qk = tl.make_block_ptr(dA_qk + (bos * H + i_h) * BT, (T, BT), (H*BT, 1), (i_t * BT, 0), (BT, BT), (1, 0)) - p_dA_qb = tl.make_block_ptr(dA_qb + (bos * H + i_h) * BT, (T, BT), (H*BT, 1), (i_t * BT, 0), (BT, BT), (1, 0)) + p_dA_qk = dA_qk + (bos * H + i_h) * BT + o_t[:, None] * (H*BT) + o_A[None, :] + p_dA_qb = dA_qb + (bos * H + i_h) * BT + o_t[:, None] * (H*BT) + o_A[None, :] m_s = tl.arange(0, BT)[:, None] >= tl.arange(0, BT)[None, :] b_dA_qk = tl.where(m_s, b_dA_qk * scale, 0.) - tl.store(p_dA_qk, b_dA_qk.to(p_dA_qk.dtype.element_ty), boundary_check=(0, 1)) + tl.store(p_dA_qk, b_dA_qk.to(p_dA_qk.dtype.element_ty), mask=m_A) b_dA_qb = tl.where(m_s, b_dA_qb * scale, 0.) - tl.store(p_dA_qb, b_dA_qb.to(p_dA_qb.dtype.element_ty), boundary_check=(0, 1)) + tl.store(p_dA_qb, b_dA_qb.to(p_dA_qb.dtype.element_ty), mask=m_A) @triton.heuristics({ @@ -130,13 +137,14 @@ def chunk_dplr_bwd_o_kernel( BV: tl.constexpr, IS_VARLEN: tl.constexpr, ): - i_k, i_t, i_bh = tl.program_id(0), tl.program_id(1), tl.program_id(2) + i_k, i_t = unflatten_program_id(tl.cdiv(K, BK)) + i_bh = tl.program_id(1).to(tl.int64) i_b, i_h = i_bh // H, i_bh % H 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.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) + i_n, i_t = tl.load(chunk_indices + i_t * 2).to(tl.int32), tl.load(chunk_indices + i_t * 2 + 1).to(tl.int64) + bos, eos = tl.load(cu_seqlens + i_n).to(tl.int64), tl.load(cu_seqlens + i_n + 1).to(tl.int64) T = eos - bos NT = tl.cdiv(T, BT) else: @@ -165,6 +173,9 @@ def chunk_dplr_bwd_o_kernel( stride_qk = H*K stride_vo = H*V + o_t = i_t * BT + tl.arange(0, BT) + o_k = i_k * BK + tl.arange(0, BK) + m_t = o_t < T b_dq = tl.zeros([BT, BK], dtype=tl.float32) b_dk = tl.zeros([BT, BK], dtype=tl.float32) b_dw = tl.zeros([BT, BK], dtype=tl.float32) @@ -172,49 +183,53 @@ def chunk_dplr_bwd_o_kernel( b_dgk_last = tl.zeros([BK], dtype=tl.float32) for i_v in range(tl.cdiv(V, BV)): - p_v = tl.make_block_ptr(v, (T, V), (stride_vo, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) - p_v_new = tl.make_block_ptr(v_new, (T, V), (stride_vo, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) - p_do = tl.make_block_ptr(do, (T, V), (stride_vo, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) - p_h = tl.make_block_ptr(h, (V, K), (1, V), (i_v * BV, i_k * BK), (BV, BK), (0, 1)) - p_dh = tl.make_block_ptr(dh, (V, K), (1, V), (i_v * BV, i_k * BK), (BV, BK), (0, 1)) + o_v = i_v * BV + tl.arange(0, BV) + m_v = m_t[:, None] & (o_v[None, :] < V) + m_hv = (o_v[:, None] < V) & (o_k[None, :] < K) + p_v = v + o_t[:, None] * stride_vo + o_v[None, :] + p_v_new = v_new + o_t[:, None] * stride_vo + o_v[None, :] + p_do = do + o_t[:, None] * stride_vo + o_v[None, :] + p_h = h + o_v[:, None] + o_k[None, :] * V + p_dh = dh + o_v[:, None] + o_k[None, :] * V # [BT, BV] - b_v = tl.load(p_v, boundary_check=(0, 1)) - b_v_new = tl.load(p_v_new, boundary_check=(0, 1)) - b_do = tl.load(p_do, boundary_check=(0, 1)) + b_v = tl.load(p_v, mask=m_v, other=0.0) + b_v_new = tl.load(p_v_new, mask=m_v, other=0.0) + b_do = tl.load(p_do, mask=m_v, other=0.0) # [BV, BK] - b_h = tl.load(p_h, boundary_check=(0, 1)) - b_dh = tl.load(p_dh, boundary_check=(0, 1)) + b_h = tl.load(p_h, mask=m_hv, other=0.0) + b_dh = tl.load(p_dh, mask=m_hv, other=0.0) b_dgk_last += tl.sum((b_h * b_dh).to(tl.float32), axis=0) # [BT, BV] @ [BV, BK] -> [BT, BK] - b_dq += tl.dot(b_do, b_h.to(b_do.dtype)) + b_dq = tl.dot(b_do, b_h.to(b_do.dtype), b_dq) # [BT, BV] @ [BV, BK] -> [BT, BK] - b_dk += tl.dot(b_v, b_dh.to(b_v.dtype)) - b_db += tl.dot(b_v_new, b_dh.to(b_v_new.dtype)) - p_dv = tl.make_block_ptr(dv, (T, V), (stride_vo, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) - b_dv = tl.load(p_dv, boundary_check=(0, 1)) - b_dw += tl.dot(b_dv.to(b_v.dtype), b_h.to(b_v.dtype)) + b_dk = tl.dot(b_v, b_dh.to(b_v.dtype), b_dk) + b_db = tl.dot(b_v_new, b_dh.to(b_v_new.dtype), b_db) + p_dv = dv + o_t[:, None] * stride_vo + o_v[None, :] + b_dv = tl.load(p_dv, mask=m_v, other=0.0) + b_dw = tl.dot(b_dv.to(b_v.dtype), b_h.to(b_v.dtype), b_dw) m_k = (i_k*BK+tl.arange(0, BK)) < K last_idx = min(i_t * BT + BT, T) - 1 b_gk_last = tl.load(gk + last_idx * stride_qk + i_k*BK + tl.arange(0, BK), mask=m_k, other=float('-inf')) b_dgk_last *= exp2(b_gk_last) - p_k = tl.make_block_ptr(k, (T, K), (stride_qk, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) - p_b = tl.make_block_ptr(b, (T, K), (stride_qk, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) - b_k = tl.load(p_k, boundary_check=(0, 1)) - b_b = tl.load(p_b, boundary_check=(0, 1)) + m_kk = m_t[:, None] & m_k[None, :] + p_k = k + o_t[:, None] * stride_qk + o_k[None, :] + p_b = b + o_t[:, None] * stride_qk + o_k[None, :] + b_k = tl.load(p_k, mask=m_kk, other=0.0) + b_b = tl.load(p_b, mask=m_kk, other=0.0) b_dgk_last += tl.sum(b_k * b_dk, axis=0) b_dgk_last += tl.sum(b_b * b_db, axis=0) tl.store(dgk_last + tl.arange(0, BK) + i_k * BK, b_dgk_last, mask=m_k) - p_dw = tl.make_block_ptr(dw, (T, K), (stride_qk, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) - p_dk = tl.make_block_ptr(dk, (T, K), (stride_qk, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) - p_db = tl.make_block_ptr(db, (T, K), (stride_qk, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) - p_dq = tl.make_block_ptr(dq, (T, K), (stride_qk, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) - tl.store(p_dw, b_dw.to(p_dw.dtype.element_ty), boundary_check=(0, 1)) - tl.store(p_dk, b_dk.to(p_dk.dtype.element_ty), boundary_check=(0, 1)) - tl.store(p_db, b_db.to(p_db.dtype.element_ty), boundary_check=(0, 1)) - tl.store(p_dq, b_dq.to(p_dq.dtype.element_ty), boundary_check=(0, 1)) + p_dw = dw + o_t[:, None] * stride_qk + o_k[None, :] + p_dk = dk + o_t[:, None] * stride_qk + o_k[None, :] + p_db = db + o_t[:, None] * stride_qk + o_k[None, :] + p_dq = dq + o_t[:, None] * stride_qk + o_k[None, :] + tl.store(p_dw, b_dw.to(p_dw.dtype.element_ty), mask=m_kk) + tl.store(p_dk, b_dk.to(p_dk.dtype.element_ty), mask=m_kk) + tl.store(p_db, b_db.to(p_db.dtype.element_ty), mask=m_kk) + tl.store(p_dq, b_dq.to(p_dq.dtype.element_ty), mask=m_kk) @triton.heuristics({ @@ -228,7 +243,7 @@ def chunk_dplr_bwd_o_kernel( for BK in BK_LIST for BV in BK_LIST ], - key=['BT'], + key=['BT', 'K', 'V'], **autotune_cache_kwargs, ) @triton.jit @@ -249,12 +264,13 @@ def chunk_dplr_bwd_kernel_dv( BV: tl.constexpr, IS_VARLEN: tl.constexpr, ): - i_v, i_t, i_bh = tl.program_id(0), tl.program_id(1), tl.program_id(2) + i_v, i_t = unflatten_program_id(tl.cdiv(V, BV)) + i_bh = tl.program_id(1).to(tl.int64) i_b, i_h = i_bh // H, i_bh % H 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.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) + i_n, i_t = tl.load(chunk_indices + i_t * 2).to(tl.int32), tl.load(chunk_indices + i_t * 2 + 1).to(tl.int64) + bos, eos = tl.load(cu_seqlens + i_n).to(tl.int64), tl.load(cu_seqlens + i_n + 1).to(tl.int64) T = eos - bos NT = tl.cdiv(T, BT) else: @@ -275,20 +291,29 @@ def chunk_dplr_bwd_kernel_dv( stride_vo = H*V stride_A = H*BT + o_t = i_t * BT + tl.arange(0, BT) + o_v = i_v * BV + tl.arange(0, BV) + m_t = o_t < T for i_k in range(tl.cdiv(K, BK)): - p_dh = tl.make_block_ptr(dh, (K, V), (V, 1), (i_k * BK, i_v * BV), (BK, BV), (1, 0)) - p_kg = tl.make_block_ptr(kg, (T, K), (stride_qk, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) - b_dh = tl.load(p_dh, boundary_check=(0, 1)) - b_kg = tl.load(p_kg, boundary_check=(0, 1)) - b_dv += tl.dot(b_kg, b_dh.to(b_kg.dtype)) - - p_Aqk = tl.make_block_ptr(A_qk, (BT, T), (1, stride_A), (0, i_t * BT), (BT, BT), (0, 1)) - b_A = tl.where(tl.arange(0, BT)[:, None] <= tl.arange(0, BT)[None, :], tl.load(p_Aqk, boundary_check=(0, 1)), 0) - p_do = tl.make_block_ptr(do, (T, V), (stride_vo, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) - p_dv = tl.make_block_ptr(dv, (T, V), (stride_vo, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) - b_do = tl.load(p_do, boundary_check=(0, 1)) - b_dv += tl.dot(b_A.to(b_do.dtype), b_do) - tl.store(p_dv, b_dv.to(p_dv.dtype.element_ty), boundary_check=(0, 1)) + o_k = i_k * BK + tl.arange(0, BK) + m_dh = (o_k[:, None] < K) & (o_v[None, :] < V) + m_kg = m_t[:, None] & (o_k[None, :] < K) + p_dh = dh + o_k[:, None] * V + o_v[None, :] + p_kg = kg + o_t[:, None] * stride_qk + o_k[None, :] + b_dh = tl.load(p_dh, mask=m_dh, other=0.0) + b_kg = tl.load(p_kg, mask=m_kg, other=0.0) + b_dv = tl.dot(b_kg, b_dh.to(b_kg.dtype), b_dv) + + o_A = tl.arange(0, BT) + m_A = (o_A[:, None] < BT) & m_t[None, :] + m_v = m_t[:, None] & (o_v[None, :] < V) + p_Aqk = A_qk + o_A[:, None] + o_t[None, :] * stride_A + b_A = tl.where(tl.arange(0, BT)[:, None] <= tl.arange(0, BT)[None, :], tl.load(p_Aqk, mask=m_A, other=0.0), 0) + p_do = do + o_t[:, None] * stride_vo + o_v[None, :] + p_dv = dv + o_t[:, None] * stride_vo + o_v[None, :] + b_do = tl.load(p_do, mask=m_v, other=0.0) + b_dv = tl.dot(b_A.to(b_do.dtype), b_do, b_dv) + tl.store(p_dv, b_dv.to(p_dv.dtype.element_ty), mask=m_v) def chunk_dplr_bwd_dv( @@ -309,7 +334,7 @@ def chunk_dplr_bwd_dv( dv = torch.empty_like(do) - def grid(meta): return (triton.cdiv(V, meta['BV']), NT, B * H) + def grid(meta): return (triton.cdiv(V, meta['BV']) * NT, B * H) chunk_dplr_bwd_kernel_dv[grid]( A_qk=A_qk, kg=kg, @@ -358,7 +383,7 @@ def chunk_dplr_bwd_o( dk = torch.empty_like(k) dw = torch.empty_like(w) db = torch.empty_like(b) - grid = (NK, NT, B * H) + grid = (NK * NT, B * H) dgk_last = torch.empty(B, NT, H, K, dtype=torch.float, device=w.device) diff --git a/build/torch-cuda/ops/generalized_delta_rule/dplr/chunk_o_fwd.py b/build/torch-cuda/ops/generalized_delta_rule/dplr/chunk_o_fwd.py index 2de1977e78885a4e432ab0ce3c3323efe80adfa1..3cfeae3d30aa28d9013b3f60f666fff160874b17 100644 --- a/build/torch-cuda/ops/generalized_delta_rule/dplr/chunk_o_fwd.py +++ b/build/torch-cuda/ops/generalized_delta_rule/dplr/chunk_o_fwd.py @@ -10,6 +10,7 @@ import triton import triton.language as tl from ....ops.utils import prepare_chunk_indices +from ....ops.utils.op import unflatten_program_id from ....utils import IS_AMD, autotune_cache_kwargs, check_shared_mem NUM_WARPS_AUTOTUNE = [2, 4, 8, 16] if IS_AMD else [2, 4, 8, 16, 32] @@ -28,7 +29,7 @@ BK_LIST = [32, 64, 128] if check_shared_mem() else [16, 32] for num_warps in NUM_WARPS_AUTOTUNE for num_stages in [2, 3, 4] ], - key=['BT'], + key=['BT', 'K', 'V'], **autotune_cache_kwargs, ) @triton.jit(do_not_specialize=['T']) @@ -51,13 +52,14 @@ def chunk_dplr_fwd_kernel_o( BV: tl.constexpr, IS_VARLEN: tl.constexpr, ): - i_v, i_t, i_bh = tl.program_id(0), tl.program_id(1), tl.program_id(2) + i_v, i_t = unflatten_program_id(tl.cdiv(V, BV)) + i_bh = tl.program_id(1).to(tl.int64) i_b, i_h = i_bh // H, i_bh % H 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.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) + i_n, i_t = tl.load(chunk_indices + i_t * 2).to(tl.int32), tl.load(chunk_indices + i_t * 2 + 1).to(tl.int64) + bos, eos = tl.load(cu_seqlens + i_n).to(tl.int64), tl.load(cu_seqlens + i_n + 1).to(tl.int64) T = eos - bos NT = tl.cdiv(T, BT) else: @@ -65,29 +67,38 @@ def chunk_dplr_fwd_kernel_o( i_tg = i_b * NT + i_t bos, eos = i_b * T, i_b * T + T + o_t = i_t * BT + tl.arange(0, BT) + o_v = i_v * BV + tl.arange(0, BV) + m_t = o_t < T b_o = tl.zeros([BT, BV], dtype=tl.float32) for i_k in range(tl.cdiv(K, BK)): - p_qg = tl.make_block_ptr(qg + (bos * H + i_h) * K, (T, K), (H*K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) - p_h = tl.make_block_ptr(h + (i_tg * H + i_h) * K*V, (K, V), (V, 1), (i_k * BK, i_v * BV), (BK, BV), (1, 0)) - b_qg = tl.load(p_qg, boundary_check=(0, 1)) - b_h = tl.load(p_h, boundary_check=(0, 1)) - b_o += tl.dot(b_qg, b_h) + o_k = i_k * BK + tl.arange(0, BK) + m_qg = m_t[:, None] & (o_k[None, :] < K) + m_h = (o_k[:, None] < K) & (o_v[None, :] < V) + p_qg = qg + (bos * H + i_h) * K + o_t[:, None] * (H*K) + o_k[None, :] + p_h = h + (i_tg * H + i_h) * K*V + o_k[:, None] * V + o_v[None, :] + b_qg = tl.load(p_qg, mask=m_qg, other=0.0) + b_h = tl.load(p_h, mask=m_h, other=0.0) + b_o = tl.dot(b_qg, b_h, b_o) - p_Aqk = tl.make_block_ptr(A_qk + (bos * H + i_h) * BT, (T, BT), (H*BT, 1), (i_t * BT, 0), (BT, BT), (1, 0)) - p_Aqb = tl.make_block_ptr(A_qb + (bos * H + i_h) * BT, (T, BT), (H*BT, 1), (i_t * BT, 0), (BT, BT), (1, 0)) - p_v = tl.make_block_ptr(v + (bos * H + i_h) * V, (T, V), (H*V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) - p_v_new = tl.make_block_ptr(v_new + (bos * H + i_h) * V, (T, V), (H*V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) - p_o = tl.make_block_ptr(o + (bos * H + i_h) * V, (T, V), (H*V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) + o_A = tl.arange(0, BT) + m_A = m_t[:, None] & (o_A[None, :] < BT) + m_v = m_t[:, None] & (o_v[None, :] < V) + p_Aqk = A_qk + (bos * H + i_h) * BT + o_t[:, None] * (H*BT) + o_A[None, :] + p_Aqb = A_qb + (bos * H + i_h) * BT + o_t[:, None] * (H*BT) + o_A[None, :] + p_v = v + (bos * H + i_h) * V + o_t[:, None] * (H*V) + o_v[None, :] + p_v_new = v_new + (bos * H + i_h) * V + o_t[:, None] * (H*V) + o_v[None, :] + p_o = o + (bos * H + i_h) * V + o_t[:, None] * (H*V) + o_v[None, :] m_s = tl.arange(0, BT)[:, None] >= tl.arange(0, BT)[None, :] - b_Aqk = tl.load(p_Aqk, boundary_check=(0, 1)) - b_Aqb = tl.load(p_Aqb, boundary_check=(0, 1)) + b_Aqk = tl.load(p_Aqk, mask=m_A, other=0.0) + b_Aqb = tl.load(p_Aqb, mask=m_A, other=0.0) b_Aqk = tl.where(m_s, b_Aqk, 0) b_Aqb = tl.where(m_s, b_Aqb, 0) - b_v = tl.load(p_v, boundary_check=(0, 1)) - b_v_new = tl.load(p_v_new, boundary_check=(0, 1)) + b_v = tl.load(p_v, mask=m_v, other=0.0) + b_v_new = tl.load(p_v_new, mask=m_v, other=0.0) b_o = b_o + tl.dot(b_Aqk.to(b_v.dtype), b_v) + tl.dot(b_Aqb.to(b_v_new.dtype), b_v_new) - tl.store(p_o, b_o.to(p_o.dtype.element_ty), boundary_check=(0, 1)) + tl.store(p_o, b_o.to(p_o.dtype.element_ty), mask=m_v) def chunk_dplr_fwd_o( @@ -109,7 +120,7 @@ def chunk_dplr_fwd_o( NT = triton.cdiv(T, BT) if cu_seqlens is None else len(chunk_indices) o = torch.empty_like(v) - def grid(meta): return (triton.cdiv(V, meta['BV']), NT, B * H) + def grid(meta): return (triton.cdiv(V, meta['BV']) * NT, B * H) chunk_dplr_fwd_kernel_o[grid]( qg=qg, v=v, diff --git a/build/torch-cuda/ops/generalized_delta_rule/dplr/fused_recurrent.py b/build/torch-cuda/ops/generalized_delta_rule/dplr/fused_recurrent.py index ccaaf7652486de50eebd246cf5a6b7640823ab42..0663f69d13bc8c45bb16f88665f916944705c756 100644 --- a/build/torch-cuda/ops/generalized_delta_rule/dplr/fused_recurrent.py +++ b/build/torch-cuda/ops/generalized_delta_rule/dplr/fused_recurrent.py @@ -53,7 +53,9 @@ def fused_recurrent_dplr_delta_rule_fwd_kernel( STORE_FINAL_STATE: tl.constexpr, IS_VARLEN: tl.constexpr, ): - i_v, i_nh = tl.program_id(0).to(tl.int64), tl.program_id(1).to(tl.int64) + pid = tl.program_id(0).to(tl.int64) + NV = tl.cdiv(V, BV) + i_v, i_nh = (pid % NV).to(tl.int64), (pid // NV).to(tl.int64) i_n, i_h = i_nh // H, i_nh % H if IS_VARLEN: @@ -128,7 +130,7 @@ def fused_recurrent_dplr_delta_rule_fwd( ht = q.new_empty(N, H, K, V, dtype=torch.float32) if output_final_state else None o = torch.empty_like(v) - def grid(meta): return (triton.cdiv(V, meta['BV']), N * H) + def grid(meta): return (triton.cdiv(V, meta['BV']) * N * H,) fused_recurrent_dplr_delta_rule_fwd_kernel[grid]( q=q, k=k, diff --git a/build/torch-cuda/ops/generalized_delta_rule/dplr/wy_fast_bwd.py b/build/torch-cuda/ops/generalized_delta_rule/dplr/wy_fast_bwd.py index bf7b53ae84c415219f76a24f376ceb3159e04898..c0963f6715d85709906fbe36e3af8e4291a1f9c6 100644 --- a/build/torch-cuda/ops/generalized_delta_rule/dplr/wy_fast_bwd.py +++ b/build/torch-cuda/ops/generalized_delta_rule/dplr/wy_fast_bwd.py @@ -52,55 +52,64 @@ def prepare_wy_repr_bwd_kernel( BV: tl.constexpr, IS_VARLEN: tl.constexpr, ): - i_t, i_bh = tl.program_id(0), tl.program_id(1) + i_t, i_bh = tl.program_id(0).to(tl.int64), tl.program_id(1).to(tl.int64) i_b, i_h = i_bh // H, i_bh % H 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.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) + i_n, i_t = tl.load(chunk_indices + i_t * 2).to(tl.int32), tl.load(chunk_indices + i_t * 2 + 1).to(tl.int64) + bos, eos = tl.load(cu_seqlens + i_n).to(tl.int64), tl.load(cu_seqlens + i_n + 1).to(tl.int64) T = eos - bos else: bos, eos = i_b * T, i_b * T + T - p_Aak_t = tl.make_block_ptr(A_ak + (bos*H + i_h) * BT, (BT, T), (1, H*BT), (0, i_t * BT), (BT, BT), (0, 1)) - p_Aab_inv_t = tl.make_block_ptr(A_ab_inv + (bos*H + i_h) * BT, (BT, T), (1, H*BT), (0, i_t * BT), (BT, BT), (0, 1)) - p_dAak = tl.make_block_ptr(dAak + (bos*H + i_h) * BT, (T, BT), (H*BT, 1), (i_t * BT, 0), (BT, BT), (1, 0)) - p_dAab = tl.make_block_ptr(dAab + (bos*H + i_h) * BT, (T, BT), (H*BT, 1), (i_t * BT, 0), (BT, BT), (1, 0)) - - b_A_ab_inv_t = tl.load(p_Aab_inv_t, boundary_check=(0, 1)) - b_A_ak_t = tl.load(p_Aak_t, boundary_check=(0, 1)) + o_t = i_t * BT + tl.arange(0, BT) + o_A = tl.arange(0, BT) + m_t = o_t < T + m_AT = (o_A[:, None] < BT) & m_t[None, :] + m_dA = m_t[:, None] & (o_A[None, :] < BT) + p_Aak_t = A_ak + (bos*H + i_h) * BT + o_A[:, None] + o_t[None, :] * (H*BT) + p_Aab_inv_t = A_ab_inv + (bos*H + i_h) * BT + o_A[:, None] + o_t[None, :] * (H*BT) + p_dAak = dAak + (bos*H + i_h) * BT + o_t[:, None] * (H*BT) + o_A[None, :] + p_dAab = dAab + (bos*H + i_h) * BT + o_t[:, None] * (H*BT) + o_A[None, :] + + b_A_ab_inv_t = tl.load(p_Aab_inv_t, mask=m_AT, other=0.0) + b_A_ak_t = tl.load(p_Aak_t, mask=m_AT, other=0.0) b_A_ak_t = tl.where(tl.arange(0, BT)[:, None] < tl.arange(0, BT)[None, :], b_A_ak_t, 0) b_A_ab_inv_t = tl.where(tl.arange(0, BT)[:, None] <= tl.arange(0, BT)[None, :], b_A_ab_inv_t, 0) b_A_tmp_t = tl.dot(b_A_ak_t, b_A_ab_inv_t).to(v.dtype.element_ty) b_dA_tmp = tl.zeros([BT, BT], dtype=tl.float32) for i_v in range(tl.cdiv(V, BV)): - p_v = tl.make_block_ptr(v + (bos*H + i_h) * V, (T, V), (H*V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) - p_dv = tl.make_block_ptr(dv + (bos*H + i_h) * V, (T, V), (H*V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) - p_dv0 = tl.make_block_ptr(dv0 + (bos*H + i_h) * V, (T, V), (H*V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) - p_du = tl.make_block_ptr(du + (bos*H + i_h) * V, (T, V), (H*V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) - b_v = tl.load(p_v, boundary_check=(0, 1)) - b_du = tl.load(p_du, boundary_check=(0, 1)) - b_dA_tmp += tl.dot(b_du.to(b_v.dtype), tl.trans(b_v)) - b_dv0 = tl.load(p_dv0, boundary_check=(0, 1)) + o_v = i_v * BV + tl.arange(0, BV) + m_v = m_t[:, None] & (o_v[None, :] < V) + p_v = v + (bos*H + i_h) * V + o_t[:, None] * (H*V) + o_v[None, :] + p_dv = dv + (bos*H + i_h) * V + o_t[:, None] * (H*V) + o_v[None, :] + p_dv0 = dv0 + (bos*H + i_h) * V + o_t[:, None] * (H*V) + o_v[None, :] + p_du = du + (bos*H + i_h) * V + o_t[:, None] * (H*V) + o_v[None, :] + b_v = tl.load(p_v, mask=m_v, other=0.0) + b_du = tl.load(p_du, mask=m_v, other=0.0) + b_dA_tmp = tl.dot(b_du.to(b_v.dtype), tl.trans(b_v), b_dA_tmp) + b_dv0 = tl.load(p_dv0, mask=m_v, other=0.0) b_dv = b_dv0 + tl.dot(b_A_tmp_t, b_du) - tl.store(p_dv, b_dv.to(p_dv.dtype.element_ty), boundary_check=(0, 1)) + tl.store(p_dv, b_dv.to(p_dv.dtype.element_ty), mask=m_v) m_i = tl.arange(0, BT)[:, None] > tl.arange(0, BT)[None, :] b_dA_tmp = tl.where(m_i, b_dA_tmp, 0) b_dA_ak = tl.dot(b_A_ab_inv_t, b_dA_tmp) b_dA_ak = tl.where(m_i, b_dA_ak, 0) - tl.store(p_dAak, b_dA_ak, boundary_check=(0, 1)) + tl.store(p_dAak, b_dA_ak, mask=m_dA) b_dA_ab_inv = tl.dot(b_dA_tmp, b_A_ak_t) for i_k in range(tl.cdiv(K, BK)): - p_ag = tl.make_block_ptr(ag + (bos * H + i_h) * K, (T, K), (H*K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) - p_dag = tl.make_block_ptr(dag + (bos * H + i_h) * K, (T, K), (H*K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) - p_dw = tl.make_block_ptr(dw + (bos * H + i_h) * K, (T, K), (H*K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) - b_ag = tl.load(p_ag, boundary_check=(0, 1)) - b_dw = tl.load(p_dw, boundary_check=(0, 1)) - b_dA_ab_inv += tl.dot(b_dw, tl.trans(b_ag)) + o_k = i_k * BK + tl.arange(0, BK) + m_k = m_t[:, None] & (o_k[None, :] < K) + p_ag = ag + (bos * H + i_h) * K + o_t[:, None] * (H*K) + o_k[None, :] + p_dag = dag + (bos * H + i_h) * K + o_t[:, None] * (H*K) + o_k[None, :] + p_dw = dw + (bos * H + i_h) * K + o_t[:, None] * (H*K) + o_k[None, :] + b_ag = tl.load(p_ag, mask=m_k, other=0.0) + b_dw = tl.load(p_dw, mask=m_k, other=0.0) + b_dA_ab_inv = tl.dot(b_dw, tl.trans(b_ag), b_dA_ab_inv) b_dag = tl.dot(b_A_ab_inv_t.to(b_dw.dtype), b_dw) - tl.store(p_dag, b_dag.to(p_dag.dtype.element_ty), boundary_check=(0, 1)) + tl.store(p_dag, b_dag.to(p_dag.dtype.element_ty), mask=m_k) # if we know dL/dA^(-1), for dL/dA, we can use the following formula: # dL/dA = -(A^(-1))^T @ (dL/dA^(-1)) @ (A^(-1))^T @@ -113,7 +122,7 @@ def prepare_wy_repr_bwd_kernel( b_dA_ab_inv = tl.dot(b_A_ab_inv_t, b_dA_ab_inv) b_dA_ab_inv = tl.dot(b_dA_ab_inv, b_A_ab_inv_t) b_dA_ab_inv = tl.where(m_i, b_dA_ab_inv, 0) - tl.store(p_dAab, b_dA_ab_inv, boundary_check=(0, 1)) + tl.store(p_dAab, b_dA_ab_inv, mask=m_dA) def chunk_dplr_bwd_wy( diff --git a/build/torch-cuda/ops/generalized_delta_rule/dplr/wy_fast_fwd.py b/build/torch-cuda/ops/generalized_delta_rule/dplr/wy_fast_fwd.py index 1d8c91f6b106be2412e6b07895c90f0dd0ab0bda..9612505fe987e936e874912fe15c4af494c7d044 100644 --- a/build/torch-cuda/ops/generalized_delta_rule/dplr/wy_fast_fwd.py +++ b/build/torch-cuda/ops/generalized_delta_rule/dplr/wy_fast_fwd.py @@ -37,17 +37,21 @@ def prepare_wy_repr_fwd_kernel_chunk32( BC: tl.constexpr, # placeholder, do not delete IS_VARLEN: tl.constexpr, ): - i_t, i_bh = tl.program_id(0), tl.program_id(1) + i_t, i_bh = tl.program_id(0).to(tl.int64), tl.program_id(1).to(tl.int64) i_b, i_h = i_bh // H, i_bh % H 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.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) + i_n, i_t = tl.load(chunk_indices + i_t * 2).to(tl.int32), tl.load(chunk_indices + i_t * 2 + 1).to(tl.int64) + bos, eos = tl.load(cu_seqlens + i_n).to(tl.int64), tl.load(cu_seqlens + i_n + 1).to(tl.int64) T = eos - bos else: bos, eos = i_b * T, i_b * T + T - p_Aab = tl.make_block_ptr(A_ab + (bos*H + i_h) * BT, (T, BT), (H*BT, 1), (i_t * BT, 0), (BT, BT), (1, 0)) - p_Aab_inv = tl.make_block_ptr(A_ab_inv + (bos*H + i_h) * BT, (T, BT), (H*BT, 1), (i_t * BT, 0), (BT, BT), (1, 0)) - b_A_ab = tl.load(p_Aab, boundary_check=(0, 1)) + o_t = i_t * BT + tl.arange(0, BT) + o_A = tl.arange(0, BT) + m_t = o_t < T + m_A = m_t[:, None] & (o_A[None, :] < BT) + p_Aab = A_ab + (bos*H + i_h) * BT + o_t[:, None] * (H*BT) + o_A[None, :] + p_Aab_inv = A_ab_inv + (bos*H + i_h) * BT + o_t[:, None] * (H*BT) + o_A[None, :] + b_A_ab = tl.load(p_Aab, mask=m_A, other=0.0) b_A_ab = tl.where(tl.arange(0, BT)[:, None] > tl.arange(0, BT)[None, :], b_A_ab, 0) for i in range(1, BT): mask = tl.arange(0, BT) == i @@ -55,7 +59,122 @@ def prepare_wy_repr_fwd_kernel_chunk32( b_a = b_a + tl.sum(b_a[:, None] * b_A_ab, 0) * (tl.arange(0, BT) < i) b_A_ab = tl.where(mask[:, None], b_a, b_A_ab) b_A_ab += tl.arange(0, BT)[:, None] == tl.arange(0, BT)[None, :] - tl.store(p_Aab_inv, b_A_ab.to(p_Aab_inv.dtype.element_ty), boundary_check=(0, 1)) + tl.store(p_Aab_inv, b_A_ab.to(p_Aab_inv.dtype.element_ty), mask=m_A) + + +@triton.jit +def _inv_diag16(b_L): + o_i = tl.arange(0, 16) + b_I = (o_i[:, None] == o_i[None, :]).to(tl.float16) + b_L2 = tl.dot(b_L, b_L).to(tl.float16) + b_L4 = tl.dot(b_L2, b_L2).to(tl.float16) + b_L8 = tl.dot(b_L4, b_L4).to(tl.float16) + b_T = tl.dot(b_I + b_L, b_I + b_L2) + b_T = tl.dot(b_T.to(tl.float16), b_I + b_L4) + b_T = tl.dot(b_T.to(tl.float16), b_I + b_L8) + return b_T + + +@triton.heuristics({ + 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None, +}) +@triton.autotune( + configs=[ + triton.Config({}, num_warps=num_warps) + for num_warps in [1, 2, 4] + ], + key=['BT'], + **autotune_cache_kwargs, +) +@triton.jit(do_not_specialize=['T']) +def prepare_wy_repr_fwd_kernel_safe( + A_ab, + A_ab_inv, + cu_seqlens, + chunk_indices, + T, + H: tl.constexpr, + BT: tl.constexpr, + BC: tl.constexpr, # placeholder, do not delete + IS_VARLEN: tl.constexpr, +): + i_t, i_bh = tl.program_id(0).to(tl.int64), tl.program_id(1).to(tl.int64) + i_b, i_h = i_bh // H, i_bh % H + 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.int64) + bos, eos = tl.load(cu_seqlens + i_n).to(tl.int64), tl.load(cu_seqlens + i_n + 1).to(tl.int64) + T = eos - bos + else: + bos = i_b * T + + NB: tl.constexpr = 16 + NS: tl.constexpr = BT // NB + o_i = tl.arange(0, NB) + o_c1 = NB + o_i + m_lo = o_i[:, None] > o_i[None, :] + b_zero = tl.zeros([NB, NB], dtype=tl.float32) + + r_0 = i_t * BT + o_i + m_r0 = r_0 < T + o_A0 = A_ab + (bos * H + i_h) * BT + r_0[:, None] * (H * BT) + o_I0 = A_ab_inv + (bos * H + i_h) * BT + r_0[:, None] * (H * BT) + b_L00 = tl.where(m_lo, tl.load(o_A0 + o_i[None, :], mask=m_r0[:, None], other=0.0), 0.0).to(tl.float16) + b_T00 = _inv_diag16(b_L00) + tl.store(o_I0 + o_i[None, :], b_T00, mask=m_r0[:, None]) + if NS >= 2: + r_1 = i_t * BT + NB + o_i + m_r1 = r_1 < T + o_A1 = A_ab + (bos * H + i_h) * BT + r_1[:, None] * (H * BT) + o_I1 = A_ab_inv + (bos * H + i_h) * BT + r_1[:, None] * (H * BT) + b_L11 = tl.where(m_lo, tl.load(o_A1 + o_c1[None, :], mask=m_r1[:, None], other=0.0), 0.0).to(tl.float16) + b_L10 = tl.load(o_A1 + o_i[None, :], mask=m_r1[:, None], other=0.0).to(tl.float16) + b_T11 = _inv_diag16(b_L11) + b_T10 = tl.dot(b_T11.to(tl.float16), tl.dot(b_L10, b_T00.to(tl.float16)).to(tl.float16)) + tl.store(o_I1 + o_i[None, :], b_T10, mask=m_r1[:, None]) + tl.store(o_I1 + o_c1[None, :], b_T11, mask=m_r1[:, None]) + tl.store(o_I0 + o_c1[None, :], b_zero, mask=m_r0[:, None]) + if NS == 4: + o_c2 = 2 * NB + o_i + o_c3 = 3 * NB + o_i + r_2 = i_t * BT + 2 * NB + o_i + r_3 = i_t * BT + 3 * NB + o_i + m_r2 = r_2 < T + m_r3 = r_3 < T + o_A2 = A_ab + (bos * H + i_h) * BT + r_2[:, None] * (H * BT) + o_A3 = A_ab + (bos * H + i_h) * BT + r_3[:, None] * (H * BT) + o_I2 = A_ab_inv + (bos * H + i_h) * BT + r_2[:, None] * (H * BT) + o_I3 = A_ab_inv + (bos * H + i_h) * BT + r_3[:, None] * (H * BT) + b_L22 = tl.where(m_lo, tl.load(o_A2 + o_c2[None, :], mask=m_r2[:, None], other=0.0), 0.0).to(tl.float16) + b_L33 = tl.where(m_lo, tl.load(o_A3 + o_c3[None, :], mask=m_r3[:, None], other=0.0), 0.0).to(tl.float16) + b_L20 = tl.load(o_A2 + o_i[None, :], mask=m_r2[:, None], other=0.0).to(tl.float16) + b_L21 = tl.load(o_A2 + o_c1[None, :], mask=m_r2[:, None], other=0.0).to(tl.float16) + b_L30 = tl.load(o_A3 + o_i[None, :], mask=m_r3[:, None], other=0.0).to(tl.float16) + b_L31 = tl.load(o_A3 + o_c1[None, :], mask=m_r3[:, None], other=0.0).to(tl.float16) + b_L32 = tl.load(o_A3 + o_c2[None, :], mask=m_r3[:, None], other=0.0).to(tl.float16) + + b_T22 = _inv_diag16(b_L22) + b_T33 = _inv_diag16(b_L33) + b_T32 = tl.dot(b_T33.to(tl.float16), tl.dot(b_L32, b_T22.to(tl.float16)).to(tl.float16)) + b_P1 = tl.dot(b_L20, b_T00.to(tl.float16)) + tl.dot(b_L21, b_T10.to(tl.float16)) + b_P2 = tl.dot(b_L30, b_T00.to(tl.float16)) + tl.dot(b_L31, b_T10.to(tl.float16)) + b_T20 = tl.dot(b_T22.to(tl.float16), b_P1.to(tl.float16)) + b_T21 = tl.dot(b_T22.to(tl.float16), tl.dot(b_L21, b_T11.to(tl.float16)).to(tl.float16)) + b_T30 = tl.dot(b_T32.to(tl.float16), b_P1.to(tl.float16)) + tl.dot(b_T33.to(tl.float16), b_P2.to(tl.float16)) + b_T31 = tl.dot(b_T32.to(tl.float16), tl.dot(b_L21, b_T11.to(tl.float16)).to(tl.float16)) \ + + tl.dot(b_T33.to(tl.float16), tl.dot(b_L31, b_T11.to(tl.float16)).to(tl.float16)) + + tl.store(o_I0 + o_c2[None, :], b_zero, mask=m_r0[:, None]) + tl.store(o_I0 + o_c3[None, :], b_zero, mask=m_r0[:, None]) + tl.store(o_I1 + o_c2[None, :], b_zero, mask=m_r1[:, None]) + tl.store(o_I1 + o_c3[None, :], b_zero, mask=m_r1[:, None]) + tl.store(o_I2 + o_i[None, :], b_T20, mask=m_r2[:, None]) + tl.store(o_I2 + o_c1[None, :], b_T21, mask=m_r2[:, None]) + tl.store(o_I2 + o_c2[None, :], b_T22, mask=m_r2[:, None]) + tl.store(o_I2 + o_c3[None, :], b_zero, mask=m_r2[:, None]) + tl.store(o_I3 + o_i[None, :], b_T30, mask=m_r3[:, None]) + tl.store(o_I3 + o_c1[None, :], b_T31, mask=m_r3[:, None]) + tl.store(o_I3 + o_c2[None, :], b_T32, mask=m_r3[:, None]) + tl.store(o_I3 + o_c3[None, :], b_T33, mask=m_r3[:, None]) @triton.heuristics({ @@ -83,26 +202,31 @@ def prepare_wy_repr_fwd_kernel_chunk64( IS_VARLEN: tl.constexpr, GATHER_SUPPORTED: tl.constexpr = IS_GATHER_SUPPORTED, ): - i_t, i_bh = tl.program_id(0), tl.program_id(1) + i_t, i_bh = tl.program_id(0).to(tl.int64), tl.program_id(1).to(tl.int64) i_b, i_h = i_bh // H, i_bh % H 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.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) + i_n, i_t = tl.load(chunk_indices + i_t * 2).to(tl.int32), tl.load(chunk_indices + i_t * 2 + 1).to(tl.int64) + bos, eos = tl.load(cu_seqlens + i_n).to(tl.int64), tl.load(cu_seqlens + i_n + 1).to(tl.int64) T = eos - bos else: bos, eos = i_b * T, i_b * T + T - p_A1 = tl.make_block_ptr(A_ab + (bos*H + i_h) * BT, (T, BT), (H*BT, 1), (i_t * BT, 0), (BC, BC), (1, 0)) - p_A2 = tl.make_block_ptr(A_ab + (bos*H + i_h) * BT, (T, BT), (H*BT, 1), (i_t * BT + BC, BC), (BC, BC), (1, 0)) - p_A3 = tl.make_block_ptr(A_ab + (bos*H + i_h) * BT, (T, BT), (H*BT, 1), (i_t * BT + BC, 0), (BC, BC), (1, 0)) - p_A_inv1 = tl.make_block_ptr(A_ab_inv + (bos*H + i_h) * BT, (T, BT), (H*BT, 1), (i_t * BT, 0), (BC, BC), (1, 0)) - p_A_inv2 = tl.make_block_ptr(A_ab_inv + (bos*H + i_h) * BT, (T, BT), (H*BT, 1), (i_t * BT + BC, BC), (BC, BC), (1, 0)) - p_A_inv3 = tl.make_block_ptr(A_ab_inv + (bos*H + i_h) * BT, (T, BT), (H*BT, 1), (i_t * BT + BC, 0), (BC, BC), (1, 0)) - p_A_inv4 = tl.make_block_ptr(A_ab_inv + (bos*H + i_h) * BT, (T, BT), (H*BT, 1), (i_t * BT, BC), (BC, BC), (1, 0)) - - b_A = tl.load(p_A1, boundary_check=(0, 1)) - b_A2 = tl.load(p_A2, boundary_check=(0, 1)) - b_A3 = tl.load(p_A3, boundary_check=(0, 1)) + o_c = tl.arange(0, BC) + o_t1 = i_t * BT + o_c + o_t2 = i_t * BT + BC + o_c + m_t1 = o_t1 < T + m_t2 = o_t2 < T + p_A1 = A_ab + (bos*H + i_h) * BT + o_t1[:, None] * (H*BT) + o_c[None, :] + p_A2 = A_ab + (bos*H + i_h) * BT + o_t2[:, None] * (H*BT) + (BC + o_c)[None, :] + p_A3 = A_ab + (bos*H + i_h) * BT + o_t2[:, None] * (H*BT) + o_c[None, :] + p_A_inv1 = A_ab_inv + (bos*H + i_h) * BT + o_t1[:, None] * (H*BT) + o_c[None, :] + p_A_inv2 = A_ab_inv + (bos*H + i_h) * BT + o_t2[:, None] * (H*BT) + (BC + o_c)[None, :] + p_A_inv3 = A_ab_inv + (bos*H + i_h) * BT + o_t2[:, None] * (H*BT) + o_c[None, :] + p_A_inv4 = A_ab_inv + (bos*H + i_h) * BT + o_t1[:, None] * (H*BT) + (BC + o_c)[None, :] + + b_A = tl.load(p_A1, mask=m_t1[:, None] & (o_c[None, :] < BT), other=0.0) + b_A2 = tl.load(p_A2, mask=m_t2[:, None] & ((BC + o_c)[None, :] < BT), other=0.0) + b_A3 = tl.load(p_A3, mask=m_t2[:, None] & (o_c[None, :] < BT), other=0.0) b_A = tl.where(tl.arange(0, BC)[:, None] > tl.arange(0, BC)[None, :], b_A, 0) b_A2 = tl.where(tl.arange(0, BC)[:, None] > tl.arange(0, BC)[None, :], b_A2, 0) @@ -130,11 +254,15 @@ def prepare_wy_repr_fwd_kernel_chunk64( b_A2 += tl.arange(0, BC)[:, None] == tl.arange(0, BC)[None, :] b_A3 = tl.dot(tl.dot(b_A2, b_A3), b_A) # tl.debug_barrier() - tl.store(p_A_inv1, b_A.to(p_A_inv1.dtype.element_ty, fp_downcast_rounding="rtne"), boundary_check=(0, 1)) - tl.store(p_A_inv2, b_A2.to(p_A_inv2.dtype.element_ty, fp_downcast_rounding="rtne"), boundary_check=(0, 1)) - tl.store(p_A_inv3, b_A3.to(p_A_inv3.dtype.element_ty, fp_downcast_rounding="rtne"), boundary_check=(0, 1)) + tl.store(p_A_inv1, b_A.to(p_A_inv1.dtype.element_ty, fp_downcast_rounding="rtne"), + mask=m_t1[:, None] & (o_c[None, :] < BT)) + tl.store(p_A_inv2, b_A2.to(p_A_inv2.dtype.element_ty, fp_downcast_rounding="rtne"), + mask=m_t2[:, None] & ((BC + o_c)[None, :] < BT)) + tl.store(p_A_inv3, b_A3.to(p_A_inv3.dtype.element_ty, fp_downcast_rounding="rtne"), + mask=m_t2[:, None] & (o_c[None, :] < BT)) # causal mask - tl.store(p_A_inv4, tl.zeros([BC, BC], dtype=tl.float32).to(p_A_inv4.dtype.element_ty), boundary_check=(0, 1)) + tl.store(p_A_inv4, tl.zeros([BC, BC], dtype=tl.float32).to( + p_A_inv4.dtype.element_ty), mask=m_t1[:, None] & ((BC + o_c)[None, :] < BT)) @triton.heuristics({ @@ -168,21 +296,24 @@ def wu_fwd_kernel( BV: tl.constexpr, IS_VARLEN: tl.constexpr, ): - i_t, i_bh = tl.program_id(0), tl.program_id(1) + i_t, i_bh = tl.program_id(0).to(tl.int64), tl.program_id(1).to(tl.int64) i_b, i_h = i_bh // H, i_bh % H 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.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) + i_n, i_t = tl.load(chunk_indices + i_t * 2).to(tl.int32), tl.load(chunk_indices + i_t * 2 + 1).to(tl.int64) + bos, eos = tl.load(cu_seqlens + i_n).to(tl.int64), tl.load(cu_seqlens + i_n + 1).to(tl.int64) T = eos - bos else: bos, eos = i_b * T, i_b * T + T o_s = tl.arange(0, BT) - p_A_ab_inv = tl.make_block_ptr(A_ab_inv + (bos*H + i_h) * BT, (T, BT), (H*BT, 1), (i_t * BT, 0), (BT, BT), (1, 0)) - p_A_ak = tl.make_block_ptr(A_ak + (bos*H + i_h) * BT, (T, BT), (H*BT, 1), (i_t * BT, 0), (BT, BT), (1, 0)) + o_t = i_t * BT + o_s + m_t = o_t < T + m_A = m_t[:, None] & (o_s[None, :] < BT) + p_A_ab_inv = A_ab_inv + (bos*H + i_h) * BT + o_t[:, None] * (H*BT) + o_s[None, :] + p_A_ak = A_ak + (bos*H + i_h) * BT + o_t[:, None] * (H*BT) + o_s[None, :] - b_Aab_inv = tl.load(p_A_ab_inv, boundary_check=(0, 1)) - b_Aak = tl.load(p_A_ak, boundary_check=(0, 1)) + b_Aab_inv = tl.load(p_A_ab_inv, mask=m_A, other=0.0) + b_Aak = tl.load(p_A_ak, mask=m_A, other=0.0) b_Aab_inv = tl.where(o_s[:, None] >= o_s[None, :], b_Aab_inv, 0) b_Aak = tl.where(o_s[:, None] > o_s[None, :], b_Aak, 0) # let's use tf32 here @@ -192,18 +323,22 @@ def wu_fwd_kernel( b_Aab_inv = b_Aab_inv.to(ag.dtype.element_ty, fp_downcast_rounding="rtne") for i_k in range(tl.cdiv(K, BK)): - p_ag = tl.make_block_ptr(ag + (bos*H + i_h) * K, (T, K), (H*K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) - p_w = tl.make_block_ptr(w + (bos*H + i_h) * K, (T, K), (H*K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) - b_ag = tl.load(p_ag, boundary_check=(0, 1)) + o_k = i_k * BK + tl.arange(0, BK) + m_k = m_t[:, None] & (o_k[None, :] < K) + p_ag = ag + (bos*H + i_h) * K + o_t[:, None] * (H*K) + o_k[None, :] + p_w = w + (bos*H + i_h) * K + o_t[:, None] * (H*K) + o_k[None, :] + b_ag = tl.load(p_ag, mask=m_k, other=0.0) b_w = tl.dot(b_Aab_inv, b_ag) # both bf16 or fp16 - tl.store(p_w, b_w.to(p_w.dtype.element_ty, fp_downcast_rounding="rtne"), boundary_check=(0, 1)) + tl.store(p_w, b_w.to(p_w.dtype.element_ty, fp_downcast_rounding="rtne"), mask=m_k) for i_v in range(tl.cdiv(V, BV)): - p_v = tl.make_block_ptr(v + (bos*H + i_h) * V, (T, V), (H*V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) - p_u = tl.make_block_ptr(u + (bos*H + i_h) * V, (T, V), (H*V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) - b_v = tl.load(p_v, boundary_check=(0, 1)) + o_v = i_v * BV + tl.arange(0, BV) + m_v = m_t[:, None] & (o_v[None, :] < V) + p_v = v + (bos*H + i_h) * V + o_t[:, None] * (H*V) + o_v[None, :] + p_u = u + (bos*H + i_h) * V + o_t[:, None] * (H*V) + o_v[None, :] + b_v = tl.load(p_v, mask=m_v, other=0.0) b_u = tl.dot(b_Aak, b_v) # both bf16 or fp16 - tl.store(p_u, b_u.to(p_u.dtype.element_ty, fp_downcast_rounding="rtne"), boundary_check=(0, 1)) + tl.store(p_u, b_u.to(p_u.dtype.element_ty, fp_downcast_rounding="rtne"), mask=m_v) def wu_fwd( @@ -254,6 +389,7 @@ def prepare_wy_repr_fwd( cu_seqlens: torch.LongTensor | None, chunk_size: int = 64, chunk_indices: torch.LongTensor | None = None, + safe_gate: bool = False, ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: B, T, H, _ = ag.shape BT = chunk_size @@ -262,7 +398,12 @@ def prepare_wy_repr_fwd( chunk_indices = prepare_chunk_indices(cu_seqlens, BT) if cu_seqlens is not None else None NT = triton.cdiv(T, BT) if cu_seqlens is None else len(chunk_indices) BC = min(BT, 32) - fwd_fn = prepare_wy_repr_fwd_kernel_chunk64 if BT == 64 else prepare_wy_repr_fwd_kernel_chunk32 + if safe_gate: + if BT not in (16, 32, 64): + raise ValueError(f"the safe-gate inversion supports chunk_size 16/32/64, got {BT}") + fwd_fn = prepare_wy_repr_fwd_kernel_safe + else: + fwd_fn = prepare_wy_repr_fwd_kernel_chunk64 if BT == 64 else prepare_wy_repr_fwd_kernel_chunk32 A_ab_inv = torch.empty_like(A_ab) fwd_fn[(NT, B * H)]( A_ab=A_ab, diff --git a/build/torch-cuda/ops/generalized_delta_rule/iplr/chunk.py b/build/torch-cuda/ops/generalized_delta_rule/iplr/chunk.py index f4bb021597aed7559ccf47ca373ad62ebc5bc456..5f3560ee4d73017a7b630fabbf05568810b661d3 100644 --- a/build/torch-cuda/ops/generalized_delta_rule/iplr/chunk.py +++ b/build/torch-cuda/ops/generalized_delta_rule/iplr/chunk.py @@ -11,6 +11,7 @@ import triton.language as tl from ....ops.generalized_delta_rule.iplr.wy_fast import prepare_wy_repr_fwd from ....ops.utils import prepare_chunk_indices, prepare_chunk_offsets +from ....ops.utils.op import unflatten_program_id from ....utils import ( autocast_custom_bwd, autocast_custom_fwd, @@ -60,50 +61,58 @@ def chunk_generalized_iplr_delta_rule_fwd_kernel_h( STORE_FINAL_STATE: tl.constexpr, IS_VARLEN: tl.constexpr, ): - i_k, i_v, i_nh = tl.program_id(0), tl.program_id(1), tl.program_id(2) + i_k, i_v, i_nh = unflatten_program_id(tl.cdiv(K, BK), tl.cdiv(V, BV)) i_n, i_h = i_nh // H, i_nh % H if IS_VARLEN: - bos, eos = tl.load(cu_seqlens + i_n).to(tl.int32), tl.load(cu_seqlens + i_n + 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 NT = tl.cdiv(T, BT) - boh = tl.load(chunk_offsets + i_n).to(tl.int32) + boh = tl.load(chunk_offsets + i_n).to(tl.int64) else: bos, eos = i_n * T, i_n * T + T NT = tl.cdiv(T, BT) boh = i_n * NT + o_k = i_k * BK + tl.arange(0, BK) + o_v = i_v * BV + tl.arange(0, BV) + m_h = (o_k[:, None] < K) & (o_v[None, :] < V) # [BK, BV] b_h = tl.zeros([BK, BV], dtype=tl.float32) if USE_INITIAL_STATE: - p_h0 = tl.make_block_ptr(h0 + i_nh * K*V, (K, V), (V, 1), (i_k * BK, i_v * BV), (BK, BV), (1, 0)) - b_h = tl.load(p_h0, boundary_check=(0, 1)).to(tl.float32) + p_h0 = h0 + i_nh * K*V + o_k[:, None] * V + o_v[None, :] + b_h = tl.load(p_h0, mask=m_h, other=0.0).to(tl.float32) for i_t in range(NT): - p_h = tl.make_block_ptr(h + ((boh + i_t) * H + i_h) * K*V, (K, V), (V, 1), (i_k * BK, i_v * BV), (BK, BV), (1, 0)) - tl.store(p_h, b_h.to(p_h.dtype.element_ty), boundary_check=(0, 1)) + p_h = h + ((boh + i_t) * H + i_h) * K*V + o_k[:, None] * V + o_v[None, :] + tl.store(p_h, b_h.to(p_h.dtype.element_ty), mask=m_h) b_hc = tl.zeros([BK, BV], dtype=tl.float32) - # since we need to make all DK in the SRAM. we face serve SRAM memory burden. By subchunking we allievate such burden + # since all DK values must fit in SRAM, subchunking alleviates the SRAM memory burden. for i_c in range(tl.cdiv(min(BT, T - i_t * BT), BC)): - p_k = tl.make_block_ptr(k+(bos*H+i_h)*K, (K, T), (1, H*K), (i_k * BK, i_t * BT + i_c * BC), (BK, BC), (0, 1)) - p_b = tl.make_block_ptr(b+(bos*H+i_h)*K, (K, T), (1, H*K), (i_k * BK, i_t * BT + i_c * BC), (BK, BC), (0, 1)) - p_d = tl.make_block_ptr(d+(bos*H+i_h)*K, (T, K), (H*K, 1), (i_t * BT + i_c * BC, i_k * BK), (BC, BK), (1, 0)) - p_v = tl.make_block_ptr(v+(bos*H+i_h)*V, (T, V), (H*V, 1), (i_t * BT + i_c * BC, i_v * BV), (BC, BV), (1, 0)) - p_u = tl.make_block_ptr(u+(bos*H+i_h)*V, (T, V), (H*V, 1), (i_t * BT + i_c * BC, i_v * BV), (BC, BV), (1, 0)) - p_v_new = tl.make_block_ptr(v_new+(bos*H+i_h)*V, (T, V), (H*V, 1), (i_t*BT+i_c*BC, i_v * BV), (BC, BV), (1, 0)) + o_c = i_t.to(tl.int64) * BT + i_c * BC + tl.arange(0, BC) + m_c = o_c < T + m_kc = (o_k[:, None] < K) & m_c[None, :] + m_dc = m_c[:, None] & (o_k[None, :] < K) + m_vc = m_c[:, None] & (o_v[None, :] < V) + p_k = k+(bos*H+i_h)*K + o_k[:, None] + o_c[None, :] * (H*K) + p_b = b+(bos*H+i_h)*K + o_k[:, None] + o_c[None, :] * (H*K) + p_d = d+(bos*H+i_h)*K + o_c[:, None] * (H*K) + o_k[None, :] + p_v = v+(bos*H+i_h)*V + o_c[:, None] * (H*V) + o_v[None, :] + p_u = u+(bos*H+i_h)*V + o_c[:, None] * (H*V) + o_v[None, :] + p_v_new = v_new+(bos*H+i_h)*V + o_c[:, None] * (H*V) + o_v[None, :] # [BK, BC] - b_k = tl.load(p_k, boundary_check=(0, 1)) - b_v = tl.load(p_v, boundary_check=(0, 1)) - b_d = tl.load(p_d, boundary_check=(0, 1)) - b_b = tl.load(p_b, boundary_check=(0, 1)) - b_v2 = tl.dot(b_d, b_h.to(b_d.dtype)) + tl.load(p_u, boundary_check=(0, 1)) - b_hc += tl.dot(b_k, b_v) - b_hc += tl.dot(b_b, b_v2.to(b_k.dtype)) - tl.store(p_v_new, b_v2.to(p_v_new.dtype.element_ty), boundary_check=(0, 1)) + b_k = tl.load(p_k, mask=m_kc, other=0.0) + b_v = tl.load(p_v, mask=m_vc, other=0.0) + b_d = tl.load(p_d, mask=m_dc, other=0.0) + b_b = tl.load(p_b, mask=m_kc, other=0.0) + b_v2 = tl.dot(b_d, b_h.to(b_d.dtype)) + tl.load(p_u, mask=m_vc, other=0.0) + b_hc = tl.dot(b_k, b_v, b_hc) + b_hc = tl.dot(b_b, b_v2.to(b_k.dtype), b_hc) + tl.store(p_v_new, b_v2.to(p_v_new.dtype.element_ty), mask=m_vc) b_h += b_hc if STORE_FINAL_STATE: - p_ht = tl.make_block_ptr(ht + i_nh * K*V, (K, V), (V, 1), (i_k * BK, i_v * BV), (BK, BV), (1, 0)) - tl.store(p_ht, b_h.to(p_ht.dtype.element_ty), boundary_check=(0, 1)) + p_ht = ht + i_nh * K*V + o_k[:, None] * V + o_v[None, :] + tl.store(p_ht, b_h.to(p_ht.dtype.element_ty), mask=m_h) @triton.heuristics({ @@ -140,13 +149,14 @@ def chunk_generalized_iplr_delta_rule_fwd_kernel_o( BV: tl.constexpr, IS_VARLEN: tl.constexpr, ): - i_v, i_t, i_bh = tl.program_id(0), tl.program_id(1), tl.program_id(2) + i_v, i_t = unflatten_program_id(tl.cdiv(V, BV)) + i_bh = tl.program_id(1).to(tl.int64) i_b, i_h = i_bh // H, i_bh % H 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.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) + i_n, i_t = tl.load(chunk_indices + i_t * 2).to(tl.int32), tl.load(chunk_indices + i_t * 2 + 1).to(tl.int64) + bos, eos = tl.load(cu_seqlens + i_n).to(tl.int64), tl.load(cu_seqlens + i_n + 1).to(tl.int64) T = eos - bos NT = tl.cdiv(T, BT) else: @@ -169,37 +179,45 @@ def chunk_generalized_iplr_delta_rule_fwd_kernel_o( b_Aqk = tl.zeros([BT, BT], dtype=tl.float32) b_Aqb = tl.zeros([BT, BT], dtype=tl.float32) + o_t = i_t * BT + tl.arange(0, BT) + o_v = i_v * BV + tl.arange(0, BV) + m_t = o_t < T for i_k in range(tl.cdiv(K, BK)): - p_q = tl.make_block_ptr(q, (T, K), (stride_qk, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) - p_k = tl.make_block_ptr(k, (K, T), (1, stride_qk), (i_k * BK, i_t * BT), (BK, BT), (0, 1)) - p_h = tl.make_block_ptr(h, (K, V), (V, 1), (i_k * BK, i_v * BV), (BK, BV), (1, 0)) - p_b = tl.make_block_ptr(b, (K, T), (1, stride_qk), (i_k * BK, i_t * BT), (BK, BT), (0, 1)) + o_k = i_k * BK + tl.arange(0, BK) + m_q = m_t[:, None] & (o_k[None, :] < K) + m_k = (o_k[:, None] < K) & m_t[None, :] + m_h = (o_k[:, None] < K) & (o_v[None, :] < V) + p_q = q + o_t[:, None] * stride_qk + o_k[None, :] + p_k = k + o_k[:, None] + o_t[None, :] * stride_qk + p_h = h + o_k[:, None] * V + o_v[None, :] + p_b = b + o_k[:, None] + o_t[None, :] * stride_qk # [BT, BK] - b_q = tl.load(p_q, boundary_check=(0, 1)) + b_q = tl.load(p_q, mask=m_q, other=0.0) # [BK, BT] - b_k = tl.load(p_k, boundary_check=(0, 1)) - b_b = tl.load(p_b, boundary_check=(0, 1)) + b_k = tl.load(p_k, mask=m_k, other=0.0) + b_b = tl.load(p_b, mask=m_k, other=0.0) # [BK, BV] - b_h = tl.load(p_h, boundary_check=(0, 1)) + b_h = tl.load(p_h, mask=m_h, other=0.0) # [BT, BK] @ [BK, BV] -> [BT, BV] - b_o += tl.dot(b_q, b_h) + b_o = tl.dot(b_q, b_h, b_o) # [BT, BK] @ [BK, BT] -> [BT, BT] - b_Aqk += tl.dot(b_q, b_k) + b_Aqk = tl.dot(b_q, b_k, b_Aqk) # [BT, BK] @ [BK, BT] -> [BT, BT] - b_Aqb += tl.dot(b_q, b_b) + b_Aqb = tl.dot(b_q, b_b, b_Aqb) o_i = tl.arange(0, BT) m_A = o_i[:, None] >= o_i[None, :] b_Aqk = tl.where(m_A, b_Aqk, 0) b_Aqb = tl.where(m_A, b_Aqb, 0) - p_v = tl.make_block_ptr(v, (T, V), (stride_vo, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) - p_u = tl.make_block_ptr(u, (T, V), (stride_vo, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) - p_o = tl.make_block_ptr(o, (T, V), (stride_vo, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) - b_v = tl.load(p_v, boundary_check=(0, 1)) - b_u = tl.load(p_u, boundary_check=(0, 1)) + m_v = m_t[:, None] & (o_v[None, :] < V) + p_v = v + o_t[:, None] * stride_vo + o_v[None, :] + p_u = u + o_t[:, None] * stride_vo + o_v[None, :] + p_o = o + o_t[:, None] * stride_vo + o_v[None, :] + b_v = tl.load(p_v, mask=m_v, other=0.0) + b_u = tl.load(p_u, mask=m_v, other=0.0) b_o = (b_o + tl.dot(b_Aqk.to(b_v.dtype), b_v) + tl.dot(b_Aqb.to(b_u.dtype), b_u)) * scale - tl.store(p_o, b_o.to(p_o.dtype.element_ty), boundary_check=(0, 1)) + tl.store(p_o, b_o.to(p_o.dtype.element_ty), mask=m_v) def chunk_generalized_iplr_delta_rule_fwd_o( @@ -225,11 +243,7 @@ def chunk_generalized_iplr_delta_rule_fwd_o( o = torch.empty_like(v) - def grid(meta): return ( - triton.cdiv(V, meta['BV']), - NT, - B * H, - ) + def grid(meta): return (triton.cdiv(V, meta['BV']) * NT, B * H) chunk_generalized_iplr_delta_rule_fwd_kernel_o[grid]( q=q, k=k, @@ -297,7 +311,7 @@ def chunk_generalized_iplr_delta_rule_fwd_h( final_state = k.new_empty(N, H, K, V, dtype=torch.float32) if output_final_state else None v_new = torch.empty_like(u) - grid = (NK, NV, N * H) + grid = (NK * NV * N * H,) chunk_generalized_iplr_delta_rule_fwd_kernel_h[grid]( k=k, diff --git a/build/torch-cuda/ops/generalized_delta_rule/iplr/fused_recurrent.py b/build/torch-cuda/ops/generalized_delta_rule/iplr/fused_recurrent.py index c7046da862bbc927cd0b762d78236a611fdc4d63..131e38de4a91767b468d4882757eebf127f1a031 100644 --- a/build/torch-cuda/ops/generalized_delta_rule/iplr/fused_recurrent.py +++ b/build/torch-cuda/ops/generalized_delta_rule/iplr/fused_recurrent.py @@ -50,7 +50,9 @@ def fused_recurrent_fwd_kernel( STORE_FINAL_STATE: tl.constexpr, # whether to store final state IS_VARLEN: tl.constexpr, ): - i_v, i_nh = tl.program_id(0), tl.program_id(1) + pid = tl.program_id(0).to(tl.int64) + NV = tl.cdiv(V, BV) + i_v, i_nh = pid % NV, (pid // NV).to(tl.int64) i_n, i_h = i_nh // H, i_nh % H if IS_VARLEN: @@ -152,7 +154,9 @@ def fused_recurrent_bwd_kernel( USE_DHT: tl.constexpr, # whether to use dht IS_VARLEN: tl.constexpr, ): - i_v, i_nh = tl.program_id(0), tl.program_id(1) + pid = tl.program_id(0).to(tl.int64) + NV = tl.cdiv(V, BV) + i_v, i_nh = (pid % NV).to(tl.int64), (pid // NV).to(tl.int64) i_n, i_h = i_nh // H, i_nh % H dk += i_v * B * H * K * T db += i_v * B * H * K * T @@ -299,15 +303,14 @@ class FusedRecurrentIPLRDeltaRuleFunction(torch.autograd.Function): BK = triton.next_power_of_2(K) if output_final_state: - final_state = q.new_empty(B, H, K, V, dtype=torch.float32) + final_state = q.new_empty(N, H, K, V, dtype=torch.float32) else: final_state = None ha = torch.empty_like(v, dtype=torch.float32) def grid(meta): return ( - triton.cdiv(V, meta['BV']), - N * H, + triton.cdiv(V, meta['BV']) * N * H, ) o = torch.empty_like(v) fused_recurrent_fwd_kernel[grid]( @@ -349,7 +352,7 @@ class FusedRecurrentIPLRDeltaRuleFunction(torch.autograd.Function): db = b.new_empty(NV, *b.shape) dv = torch.empty_like(v) dha = torch.empty_like(ha) - grid = (NV, N * H) + grid = (NV * N * H,) if initial_state is not None and initial_state.requires_grad: dh0 = torch.empty_like(initial_state, dtype=torch.float32) @@ -419,9 +422,10 @@ def fused_recurrent_iplr_delta_rule( Scale factor for the RetNet attention scores. If not provided, it will default to `1 / sqrt(K)`. Default: `None`. initial_state (Optional[torch.Tensor]): - Initial state of shape `[B, H, K, V]`. Default: `None`. + Initial state of shape `[N, H, K, V]`, where `N` is the number of sequences + (`B` for fixed-length inputs). Default: `None`. output_final_state (Optional[bool]): - Whether to output the final state of shape `[B, H, K, V]`. Default: `False`. + Whether to output the final state of shape `[N, H, K, V]`. Default: `False`. cu_seqlens (torch.LongTensor): Cumulative sequence lengths of shape `[N+1]` used for variable-length training, consistent with the FlashAttention API. diff --git a/build/torch-cuda/ops/generalized_delta_rule/iplr/wy_fast.py b/build/torch-cuda/ops/generalized_delta_rule/iplr/wy_fast.py index 32c1ef769305fed000ff73fe9624be55ef07f106..2489423ddd4ab3340a8ddaa76478d72c9ac80be8 100644 --- a/build/torch-cuda/ops/generalized_delta_rule/iplr/wy_fast.py +++ b/build/torch-cuda/ops/generalized_delta_rule/iplr/wy_fast.py @@ -41,22 +41,27 @@ def prepare_wy_repr_fwd_kernel_chunk32( BC: tl.constexpr, # dummy placeholder IS_VARLEN: tl.constexpr, ): - i_t, i_bh = tl.program_id(0), tl.program_id(1) + i_t, i_bh = tl.program_id(0).to(tl.int64), tl.program_id(1).to(tl.int64) i_b, i_h = i_bh // H, i_bh % H 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.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) + i_n, i_t = tl.load(chunk_indices + i_t * 2).to(tl.int32), tl.load(chunk_indices + i_t * 2 + 1).to(tl.int64) + bos, eos = tl.load(cu_seqlens + i_n).to(tl.int64), tl.load(cu_seqlens + i_n + 1).to(tl.int64) T = eos - bos else: bos, eos = i_b * T, i_b * T + T + o_t = i_t * BT + tl.arange(0, BT) + m_t = o_t < T b_A = tl.zeros([BT, BT], dtype=tl.float32) for i_k in range(tl.cdiv(K, BK)): - p_a = tl.make_block_ptr(a + (bos * H + i_h) * K, (T, K), (H*K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) - p_b = tl.make_block_ptr(b + (bos * H + i_h) * K, (K, T), (1, K*H), (i_k * BK, i_t * BT), (BK, BT), (0, 1)) - b_a = tl.load(p_a, boundary_check=(0, 1)) - b_b = tl.load(p_b, boundary_check=(0, 1)) - b_A += tl.dot(b_a, b_b) + o_k = i_k * BK + tl.arange(0, BK) + m_a = m_t[:, None] & (o_k[None, :] < K) + m_b = (o_k[:, None] < K) & m_t[None, :] + p_a = a + (bos * H + i_h) * K + o_t[:, None] * (H*K) + o_k[None, :] + p_b = b + (bos * H + i_h) * K + o_k[:, None] + o_t[None, :] * (K*H) + b_a = tl.load(p_a, mask=m_a, other=0.0) + b_b = tl.load(p_b, mask=m_b, other=0.0) + b_A = tl.dot(b_a, b_b, b_A) b_A = tl.where(tl.arange(0, BT)[:, None] > tl.arange(0, BT)[None, :], b_A, 0) for i in range(1, BT): @@ -66,8 +71,9 @@ def prepare_wy_repr_fwd_kernel_chunk32( b_A = tl.where(mask[:, None], b_a, b_A) b_A += tl.arange(0, BT)[:, None] == tl.arange(0, BT)[None, :] - p_A = tl.make_block_ptr(A + (bos*H + i_h) * BT, (T, BT), (H*BT, 1), (i_t * BT, 0), (BT, BT), (1, 0)) - tl.store(p_A, b_A.to(p_A.dtype.element_ty), boundary_check=(0, 1)) + o_A = tl.arange(0, BT) + p_A = A + (bos*H + i_h) * BT + o_t[:, None] * (H*BT) + o_A[None, :] + tl.store(p_A, b_A.to(p_A.dtype.element_ty), mask=m_t[:, None] & (o_A[None, :] < BT)) @triton.heuristics({ @@ -96,11 +102,11 @@ def prepare_wy_repr_fwd_kernel_chunk64( BC: tl.constexpr, IS_VARLEN: tl.constexpr, ): - i_t, i_bh = tl.program_id(0), tl.program_id(1) + i_t, i_bh = tl.program_id(0).to(tl.int64), tl.program_id(1).to(tl.int64) i_b, i_h = i_bh // H, i_bh % H 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.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) + i_n, i_t = tl.load(chunk_indices + i_t * 2).to(tl.int32), tl.load(chunk_indices + i_t * 2 + 1).to(tl.int64) + bos, eos = tl.load(cu_seqlens + i_n).to(tl.int64), tl.load(cu_seqlens + i_n + 1).to(tl.int64) T = eos - bos else: bos, eos = i_b * T, i_b * T + T @@ -109,18 +115,28 @@ def prepare_wy_repr_fwd_kernel_chunk64( b_A2 = tl.zeros([BC, BC], dtype=tl.float32) b_A3 = tl.zeros([BC, BC], dtype=tl.float32) + o_c = tl.arange(0, BC) + o_t1 = i_t * BT + o_c + o_t2 = i_t * BT + BC + o_c + m_t1 = o_t1 < T + m_t2 = o_t2 < T for i_k in range(tl.cdiv(K, BK)): - p_a1 = tl.make_block_ptr(a + (bos * H + i_h) * K, (T, K), (H*K, 1), (i_t * BT, i_k * BK), (BC, BK), (1, 0)) - p_a2 = tl.make_block_ptr(a + (bos * H + i_h) * K, (T, K), (H*K, 1), (i_t * BT + BC, i_k * BK), (BC, BK), (1, 0)) - p_b1 = tl.make_block_ptr(b + (bos * H + i_h) * K, (K, T), (1, K*H), (i_k * BK, i_t * BT), (BK, BC), (0, 1)) - p_b2 = tl.make_block_ptr(b + (bos * H + i_h) * K, (K, T), (1, K*H), (i_k * BK, i_t * BT + BC), (BK, BC), (0, 1)) - b_a1 = tl.load(p_a1, boundary_check=(0, 1)) - b_a2 = tl.load(p_a2, boundary_check=(0, 1)) - b_b1 = tl.load(p_b1, boundary_check=(0, 1)) - b_b2 = tl.load(p_b2, boundary_check=(0, 1)) - b_A += tl.dot(b_a1, b_b1, allow_tf32=False) - b_A2 += tl.dot(b_a2, b_b2, allow_tf32=False) - b_A3 += tl.dot(b_a2, b_b1, allow_tf32=False) + o_k = i_k * BK + tl.arange(0, BK) + m_a1 = m_t1[:, None] & (o_k[None, :] < K) + m_a2 = m_t2[:, None] & (o_k[None, :] < K) + m_b1 = (o_k[:, None] < K) & m_t1[None, :] + m_b2 = (o_k[:, None] < K) & m_t2[None, :] + p_a1 = a + (bos * H + i_h) * K + o_t1[:, None] * (H*K) + o_k[None, :] + p_a2 = a + (bos * H + i_h) * K + o_t2[:, None] * (H*K) + o_k[None, :] + p_b1 = b + (bos * H + i_h) * K + o_k[:, None] + o_t1[None, :] * (K*H) + p_b2 = b + (bos * H + i_h) * K + o_k[:, None] + o_t2[None, :] * (K*H) + b_a1 = tl.load(p_a1, mask=m_a1, other=0.0) + b_a2 = tl.load(p_a2, mask=m_a2, other=0.0) + b_b1 = tl.load(p_b1, mask=m_b1, other=0.0) + b_b2 = tl.load(p_b2, mask=m_b2, other=0.0) + b_A = tl.dot(b_a1, b_b1, b_A, allow_tf32=False) + b_A2 = tl.dot(b_a2, b_b2, b_A2, allow_tf32=False) + b_A3 = tl.dot(b_a2, b_b1, b_A3, allow_tf32=False) b_A = tl.where(tl.arange(0, BC)[:, None] > tl.arange(0, BC)[None, :], b_A, 0) b_A2 = tl.where(tl.arange(0, BC)[:, None] > tl.arange(0, BC)[None, :], b_A2, 0) @@ -140,15 +156,16 @@ def prepare_wy_repr_fwd_kernel_chunk64( b_A2 += tl.arange(0, BC)[:, None] == tl.arange(0, BC)[None, :] b_A3 = tl.dot(tl.dot(b_A2, b_A3, allow_tf32=False), b_A, allow_tf32=False) - p_A1 = tl.make_block_ptr(A + (bos*H + i_h) * BT, (T, BT), (H*BT, 1), (i_t * BT, 0), (BC, BC), (1, 0)) - p_A2 = tl.make_block_ptr(A + (bos*H + i_h) * BT, (T, BT), (H*BT, 1), (i_t * BT + BC, BC), (BC, BC), (1, 0)) - p_A3 = tl.make_block_ptr(A + (bos*H + i_h) * BT, (T, BT), (H*BT, 1), (i_t * BT + BC, 0), (BC, BC), (1, 0)) - p_A4 = tl.make_block_ptr(A + (bos*H + i_h) * BT, (T, BT), (H*BT, 1), (i_t * BT, BC), (BC, BC), (1, 0)) - tl.store(p_A1, b_A.to(p_A1.dtype.element_ty), boundary_check=(0, 1)) - tl.store(p_A2, b_A2.to(p_A2.dtype.element_ty), boundary_check=(0, 1)) - tl.store(p_A3, b_A3.to(p_A3.dtype.element_ty), boundary_check=(0, 1)) + p_A1 = A + (bos*H + i_h) * BT + o_t1[:, None] * (H*BT) + o_c[None, :] + p_A2 = A + (bos*H + i_h) * BT + o_t2[:, None] * (H*BT) + (BC + o_c)[None, :] + p_A3 = A + (bos*H + i_h) * BT + o_t2[:, None] * (H*BT) + o_c[None, :] + p_A4 = A + (bos*H + i_h) * BT + o_t1[:, None] * (H*BT) + (BC + o_c)[None, :] + tl.store(p_A1, b_A.to(p_A1.dtype.element_ty), mask=m_t1[:, None] & (o_c[None, :] < BT)) + tl.store(p_A2, b_A2.to(p_A2.dtype.element_ty), mask=m_t2[:, None] & ((BC + o_c)[None, :] < BT)) + tl.store(p_A3, b_A3.to(p_A3.dtype.element_ty), mask=m_t2[:, None] & (o_c[None, :] < BT)) # causal mask - tl.store(p_A4, tl.zeros([BC, BC], dtype=tl.float32).to(p_A4.dtype.element_ty), boundary_check=(0, 1)) + tl.store(p_A4, tl.zeros([BC, BC], dtype=tl.float32).to( + p_A4.dtype.element_ty), mask=m_t1[:, None] & ((BC + o_c)[None, :] < BT)) @triton.heuristics({ @@ -181,40 +198,48 @@ def wu_fwd_kernel( BV: tl.constexpr, IS_VARLEN: tl.constexpr, ): - i_t, i_bh = tl.program_id(0), tl.program_id(1) + i_t, i_bh = tl.program_id(0).to(tl.int64), tl.program_id(1).to(tl.int64) i_b, i_h = i_bh // H, i_bh % H 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.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) + i_n, i_t = tl.load(chunk_indices + i_t * 2).to(tl.int32), tl.load(chunk_indices + i_t * 2 + 1).to(tl.int64) + bos, eos = tl.load(cu_seqlens + i_n).to(tl.int64), tl.load(cu_seqlens + i_n + 1).to(tl.int64) T = eos - bos else: bos, eos = i_b * T, i_b * T + T - p_A = tl.make_block_ptr(A + (bos*H + i_h) * BT, (T, BT), (H*BT, 1), (i_t * BT, 0), (BT, BT), (1, 0)) + o_t = i_t * BT + tl.arange(0, BT) + o_A = tl.arange(0, BT) + m_t = o_t < T + m_A = m_t[:, None] & (o_A[None, :] < BT) + p_A = A + (bos*H + i_h) * BT + o_t[:, None] * (H*BT) + o_A[None, :] - b_A = tl.load(p_A, boundary_check=(0, 1)) + b_A = tl.load(p_A, mask=m_A, other=0.0) b_Aak = tl.zeros([BT, BT], dtype=tl.float32) for i_k in range(tl.cdiv(K, BK)): - p_k = tl.make_block_ptr(k + (bos * H + i_h) * K, (T, K), (H*K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) - p_a = tl.make_block_ptr(a + (bos * H + i_h) * K, (T, K), (H*K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) - p_w = tl.make_block_ptr(w + (bos * H + i_h) * K, (T, K), (H*K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) - b_k = tl.load(p_k, boundary_check=(0, 1)) - b_a = tl.load(p_a, boundary_check=(0, 1)) + o_k = i_k * BK + tl.arange(0, BK) + m_k = m_t[:, None] & (o_k[None, :] < K) + p_k = k + (bos * H + i_h) * K + o_t[:, None] * (H*K) + o_k[None, :] + p_a = a + (bos * H + i_h) * K + o_t[:, None] * (H*K) + o_k[None, :] + p_w = w + (bos * H + i_h) * K + o_t[:, None] * (H*K) + o_k[None, :] + b_k = tl.load(p_k, mask=m_k, other=0.0) + b_a = tl.load(p_a, mask=m_k, other=0.0) b_w = tl.dot(b_A, b_a) - b_Aak += tl.dot(b_a, tl.trans(b_k)) - tl.store(p_w, b_w.to(p_w.dtype.element_ty), boundary_check=(0, 1)) + b_Aak = tl.dot(b_a, tl.trans(b_k), b_Aak) + tl.store(p_w, b_w.to(p_w.dtype.element_ty), mask=m_k) b_Aak = tl.where(tl.arange(0, BT)[:, None] > tl.arange(0, BT)[None, :], b_Aak, 0) b_Aak = b_Aak.to(k.dtype.element_ty) for i_v in range(tl.cdiv(V, BV)): - p_v = tl.make_block_ptr(v + (bos*H + i_h) * V, (T, V), (H*V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) - p_u = tl.make_block_ptr(u + (bos*H + i_h) * V, (T, V), (H*V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) - b_v = tl.load(p_v, boundary_check=(0, 1)) + o_v = i_v * BV + tl.arange(0, BV) + m_v = m_t[:, None] & (o_v[None, :] < V) + p_v = v + (bos*H + i_h) * V + o_t[:, None] * (H*V) + o_v[None, :] + p_u = u + (bos*H + i_h) * V + o_t[:, None] * (H*V) + o_v[None, :] + b_v = tl.load(p_v, mask=m_v, other=0.0) b_v = tl.dot(b_Aak, b_v).to(v.dtype.element_ty) b_u = tl.dot(b_A, b_v) - tl.store(p_u, b_u.to(p_u.dtype.element_ty), boundary_check=(0, 1)) + tl.store(p_u, b_u.to(p_u.dtype.element_ty), mask=m_v) def prepare_wy_repr_fwd( diff --git a/build/torch-xpu/ops/rebased/__init__.py b/build/torch-cuda/ops/gla/backends/__init__.py similarity index 79% rename from build/torch-xpu/ops/rebased/__init__.py rename to build/torch-cuda/ops/gla/backends/__init__.py index 694057c5a26afa9c15c91ba8d9ef9409842a6677..a1c4e2d82339f9c63103c95e614904269132efe4 100644 --- a/build/torch-xpu/ops/rebased/__init__.py +++ b/build/torch-cuda/ops/gla/backends/__init__.py @@ -5,8 +5,7 @@ # For a list of all contributors, visit: # https://github.com/fla-org/flash-linear-attention/graphs/contributors -from .parallel import parallel_rebased +"""GLA backends.""" -__all__ = [ - 'parallel_rebased', -] + +__all__: list[str] = [] diff --git a/build/torch-cuda/ops/gla/chunk.py b/build/torch-cuda/ops/gla/chunk.py index e39b6fec24c07f381d04c26e545224e030471cce..e472484a28e26c894bc424d9caa54d4b2b45c7b4 100644 --- a/build/torch-cuda/ops/gla/chunk.py +++ b/build/torch-cuda/ops/gla/chunk.py @@ -9,13 +9,21 @@ import torch import triton import triton.language as tl +from ...backends import dispatch from ...ops.common.chunk_h import chunk_bwd_dh, chunk_fwd_h from ...ops.utils import prepare_chunk_indices from ...ops.utils.cache import fla_cache_autotune from ...ops.utils.constant import RCP_LN2 from ...ops.utils.cumsum import chunk_local_cumsum -from ...ops.utils.op import exp2 -from ...utils import autotune_cache_kwargs, check_shared_mem, input_guard +from ...ops.utils.op import exp2, unflatten_program_id +from ...utils import ( + IS_NVIDIA_HOPPER, + TRITON_ABOVE_3_6_0, + TRITON_ABOVE_3_8_0, + autotune_cache_kwargs, + check_shared_mem, + input_guard, +) BK_LIST = [32, 64] if check_shared_mem() else [16, 32] BV_LIST = [64, 128] if check_shared_mem('ampere') else [16, 32] @@ -24,7 +32,7 @@ BV_LIST = [64, 128] if check_shared_mem('ampere') else [16, 32] def _prune_gla_bwd_configs(configs, nargs, **kwargs): # Keep a tile only if it leaves headroom below its dim, or is the smallest # option (so small dims still autotune); this avoids a software-pipelined - # make_block_ptr prefetching past the tensor. K/V arrive as launch kwargs. + # block load prefetching past the tensor. K/V arrive as launch kwargs. args = {**(nargs or {}), **kwargs} K, V = args['K'], args['V'] min_bk = min(c.kwargs['BK'] for c in configs) @@ -36,6 +44,15 @@ def _prune_gla_bwd_configs(configs, nargs, **kwargs): ] +def _prune_gla_bwd_inter_configs(configs, nargs, **kwargs): + configs = _prune_gla_bwd_configs(configs, nargs, **kwargs) + args = {**(nargs or {}), **kwargs} + if IS_NVIDIA_HOPPER and TRITON_ABOVE_3_6_0 and not TRITON_ABOVE_3_8_0 and args['STATE_V_FIRST']: + # avoid a Hopper ptxas WGMMA miscompile: https://github.com/triton-lang/triton/issues/11366 + configs = [c for c in configs if c.kwargs['BK'] != 32 or c.num_warps != 4] + return configs + + @triton.heuristics({ 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None, }) @@ -67,12 +84,12 @@ def chunk_gla_fwd_A_kernel_intra_sub_inter( NC: tl.constexpr, IS_VARLEN: tl.constexpr, ): - i_t, i_c, i_bh = tl.program_id(0), tl.program_id(1), tl.program_id(2) + i_t, i_c, i_bh = tl.program_id(0).to(tl.int64), tl.program_id(1).to(tl.int64), tl.program_id(2).to(tl.int64) i_b, i_h = i_bh // H, i_bh % H i_i, i_j = i_c // NC, i_c % NC 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.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) + i_n, i_t = tl.load(chunk_indices + i_t * 2).to(tl.int32), tl.load(chunk_indices + i_t * 2 + 1).to(tl.int64) + bos, eos = tl.load(cu_seqlens + i_n).to(tl.int64), tl.load(cu_seqlens + i_n + 1).to(tl.int64) T = eos - bos else: bos, eos = i_b * T, i_b * T + T @@ -83,31 +100,39 @@ def chunk_gla_fwd_A_kernel_intra_sub_inter( return b_A = tl.zeros([BC, BC], dtype=tl.float32) + o_i = i_t * BT + i_i * BC + tl.arange(0, BC) + o_j = i_t * BT + i_j * BC + tl.arange(0, BC) + m_i = o_i < T + m_j = o_j < T for i_k in range(tl.cdiv(K, BK)): o_k = i_k * BK + tl.arange(0, BK) m_k = o_k < K + m_qk = m_i[:, None] & m_k[None, :] + m_kj = m_k[:, None] & m_j[None, :] - p_q = tl.make_block_ptr(q + (bos*H+i_h)*K, (T, K), (H*K, 1), (i_t * BT + i_i * BC, i_k * BK), (BC, BK), (1, 0)) - p_g = tl.make_block_ptr(g + (bos*H+i_h)*K, (T, K), (H*K, 1), (i_t * BT + i_i * BC, i_k * BK), (BC, BK), (1, 0)) - p_k = tl.make_block_ptr(k + (bos*H+i_h)*K, (K, T), (1, H*K), (i_k * BK, i_t * BT + i_j * BC), (BK, BC), (0, 1)) - p_gk = tl.make_block_ptr(g + (bos*H+i_h)*K, (K, T), (1, H*K), (i_k * BK, i_t * BT + i_j * BC), (BK, BC), (0, 1)) + p_q = q + (bos*H+i_h)*K + o_i[:, None] * (H*K) + o_k[None, :] + p_g = g + (bos*H+i_h)*K + o_i[:, None] * (H*K) + o_k[None, :] + p_k = k + (bos*H+i_h)*K + o_k[:, None] + o_j[None, :] * (H*K) + p_gk = g + (bos*H+i_h)*K + o_k[:, None] + o_j[None, :] * (H*K) p_gn = g + (bos + i_t * BT + i_i * BC) * H*K + i_h * K + o_k # [BK,] b_gn = tl.load(p_gn, mask=m_k, other=0) # [BC, BK] - b_q = tl.load(p_q, boundary_check=(0, 1)) - b_g = tl.load(p_g, boundary_check=(0, 1)) + b_q = tl.load(p_q, mask=m_qk, other=0.0) + b_g = tl.load(p_g, mask=m_qk, other=0.0) b_qg = b_q * exp2(b_g - b_gn[None, :]) * scale # [BK, BC] - b_k = tl.load(p_k, boundary_check=(0, 1)) - b_gk = tl.load(p_gk, boundary_check=(0, 1)) + b_k = tl.load(p_k, mask=m_kj, other=0.0) + b_gk = tl.load(p_gk, mask=m_kj, other=0.0) b_kg = b_k * exp2(b_gn[:, None] - b_gk) # [BC, BC] using tf32 to improve precision here. - b_A += tl.dot(b_qg, b_kg) + b_A = tl.dot(b_qg, b_kg, b_A) - p_A = tl.make_block_ptr(A + (bos*H + i_h)*BT, (T, BT), (H*BT, 1), (i_t * BT + i_i * BC, i_j * BC), (BC, BC), (1, 0)) - tl.store(p_A, b_A.to(A.dtype.element_ty), boundary_check=(0, 1)) + o_jA = i_j * BC + tl.arange(0, BC) + m_A = m_i[:, None] & (o_jA[None, :] < BT) + p_A = A + (bos*H + i_h)*BT + o_i[:, None] * (H*BT) + o_jA[None, :] + tl.store(p_A, b_A.to(A.dtype.element_ty), mask=m_A) @triton.heuristics({ @@ -139,12 +164,12 @@ def chunk_gla_fwd_A_kernel_intra_sub_intra( BK: tl.constexpr, IS_VARLEN: tl.constexpr, ): - i_t, i_i, i_bh = tl.program_id(0), tl.program_id(1), tl.program_id(2) + i_t, i_i, i_bh = tl.program_id(0).to(tl.int64), tl.program_id(1).to(tl.int64), tl.program_id(2).to(tl.int64) i_b, i_h = i_bh // H, i_bh % H i_j = i_i 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.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) + i_n, i_t = tl.load(chunk_indices + i_t * 2).to(tl.int32), tl.load(chunk_indices + i_t * 2 + 1).to(tl.int64) + bos, eos = tl.load(cu_seqlens + i_n).to(tl.int64), tl.load(cu_seqlens + i_n + 1).to(tl.int64) T = eos - bos else: bos, eos = i_b * T, i_b * T + T @@ -163,10 +188,12 @@ def chunk_gla_fwd_A_kernel_intra_sub_intra( g += (bos * H + i_h) * K A += (bos * H + i_h) * BT - p_q = tl.make_block_ptr(q, (T, K), (H*K, 1), (i_t * BT + i_i * BC, 0), (BC, BK), (1, 0)) - p_g = tl.make_block_ptr(g, (T, K), (H*K, 1), (i_t * BT + i_i * BC, 0), (BC, BK), (1, 0)) - b_q = tl.load(p_q, boundary_check=(0, 1)) - b_g = tl.load(p_g, boundary_check=(0, 1)) + o_c = i_t * BT + i_i * BC + tl.arange(0, BC) + m_qk = m_A[:, None] & m_k[None, :] + p_q = q + o_c[:, None] * (H*K) + o_k[None, :] + p_g = g + o_c[:, None] * (H*K) + o_k[None, :] + b_q = tl.load(p_q, mask=m_qk, other=0.0) + b_g = tl.load(p_g, mask=m_qk, other=0.0) p_k = k + (i_t * BT + i_j * BC) * H*K + o_k p_gk = g + (i_t * BT + i_j * BC) * H*K + o_k @@ -214,13 +241,14 @@ def chunk_gla_fwd_A_kernel_intra_sub_intra_split( NC: tl.constexpr, IS_VARLEN: tl.constexpr, ): - i_k, i_tc, i_bh = tl.program_id(0), tl.program_id(1), tl.program_id(2) + i_k, i_tc = unflatten_program_id(tl.cdiv(K, BK)) + i_bh = tl.program_id(1).to(tl.int64) i_b, i_h = i_bh // H, i_bh % H - i_t, i_i = i_tc // NC, i_tc % NC + i_t, i_i = (i_tc // NC).to(tl.int64), i_tc % NC i_j = i_i 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.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) + i_n, i_t = tl.load(chunk_indices + i_t * 2).to(tl.int32), tl.load(chunk_indices + i_t * 2 + 1).to(tl.int64) + bos, eos = tl.load(cu_seqlens + i_n).to(tl.int64), tl.load(cu_seqlens + i_n + 1).to(tl.int64) all = T T = eos - bos else: @@ -241,10 +269,12 @@ def chunk_gla_fwd_A_kernel_intra_sub_intra_split( g += (bos * H + i_h) * K A += ((i_k * all + bos) * H + i_h) * BC - p_q = tl.make_block_ptr(q, (T, K), (H*K, 1), (i_t * BT + i_i * BC, i_k * BK), (BC, BK), (1, 0)) - p_g = tl.make_block_ptr(g, (T, K), (H*K, 1), (i_t * BT + i_i * BC, i_k * BK), (BC, BK), (1, 0)) - b_q = tl.load(p_q, boundary_check=(0, 1)) - b_g = tl.load(p_g, boundary_check=(0, 1)) + o_c = i_t * BT + i_i * BC + tl.arange(0, BC) + m_qk = m_A[:, None] & m_k[None, :] + p_q = q + o_c[:, None] * (H*K) + o_k[None, :] + p_g = g + o_c[:, None] * (H*K) + o_k[None, :] + b_q = tl.load(p_q, mask=m_qk, other=0.0) + b_g = tl.load(p_g, mask=m_qk, other=0.0) p_k = k + (i_t * BT + i_j * BC) * H*K + o_k p_gk = g + (i_t * BT + i_j * BC) * H*K + o_k @@ -288,11 +318,11 @@ def chunk_gla_fwd_A_kernel_intra_sub_intra_merge( NK: tl.constexpr, IS_VARLEN: tl.constexpr, ): - i_t, i_c, i_bh = tl.program_id(0), tl.program_id(1), tl.program_id(2) + i_t, i_c, i_bh = tl.program_id(0).to(tl.int64), tl.program_id(1).to(tl.int64), tl.program_id(2).to(tl.int64) i_b, i_h = i_bh // H, i_bh % H 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.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) + i_n, i_t = tl.load(chunk_indices + i_t * 2).to(tl.int32), tl.load(chunk_indices + i_t * 2 + 1).to(tl.int64) + bos, eos = tl.load(cu_seqlens + i_n).to(tl.int64), tl.load(cu_seqlens + i_n + 1).to(tl.int64) all = T T = eos - bos else: @@ -303,11 +333,16 @@ def chunk_gla_fwd_A_kernel_intra_sub_intra_merge( return b_A = tl.zeros([BC, BC], dtype=tl.float32) + o_c = i_t * BT + i_c * BC + tl.arange(0, BC) + o_i = tl.arange(0, BC) + m_c = o_c < T + m_A = m_c[:, None] & (o_i[None, :] < BC) + m_A2 = m_c[:, None] & ((i_c * BC + o_i)[None, :] < BT) for i_k in range(0, NK): - p_A = tl.make_block_ptr(A + (i_k*all+bos)*H*BC+i_h*BC, (T, BC), (H*BC, 1), (i_t*BT + i_c*BC, 0), (BC, BC), (1, 0)) - b_A += tl.load(p_A, boundary_check=(0, 1)) - p_A2 = tl.make_block_ptr(A2 + (bos*H+i_h)*BT, (T, BT), (H*BT, 1), (i_t * BT + i_c * BC, i_c * BC), (BC, BC), (1, 0)) - tl.store(p_A2, b_A.to(A2.dtype.element_ty), boundary_check=(0, 1)) + p_A = A + (i_k*all+bos)*H*BC+i_h*BC + o_c[:, None] * (H*BC) + o_i[None, :] + b_A += tl.load(p_A, mask=m_A, other=0.0) + p_A2 = A2 + (bos*H+i_h)*BT + o_c[:, None] * (H*BT) + (i_c * BC + o_i)[None, :] + tl.store(p_A2, b_A.to(A2.dtype.element_ty), mask=m_A2) @triton.heuristics({ @@ -321,7 +356,7 @@ def chunk_gla_fwd_A_kernel_intra_sub_intra_merge( for num_warps in [2, 4, 8] for num_stages in [2, 3, 4] ], - key=['BT', 'HV', 'STATE_V_FIRST'], + key=['BT', 'HV', 'STATE_V_FIRST', 'K', 'V'], **autotune_cache_kwargs, ) @triton.jit(do_not_specialize=['T']) @@ -345,13 +380,20 @@ def chunk_gla_fwd_kernel_o( BV: tl.constexpr, STATE_V_FIRST: tl.constexpr, IS_VARLEN: tl.constexpr, + USE_GRAPH: tl.constexpr = False, ): - i_v, i_t, i_bh = tl.program_id(0), tl.program_id(1), tl.program_id(2) + # Keep in fp32 (some pass it as fp64) which causes issues under compile + scale = scale.to(tl.float32) + i_v, i_t = unflatten_program_id(tl.cdiv(V, BV)) + i_bh = tl.program_id(1).to(tl.int64) i_b, i_hv = i_bh // HV, i_bh % HV i_h = i_hv // (HV // H) if IS_VARLEN: i_tg = i_t.to(tl.int64) - 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) + i_n = tl.load(chunk_indices + i_t * 2).to(tl.int32) + if USE_GRAPH and i_n < 0: + return + i_t = tl.load(chunk_indices + i_t * 2 + 1).to(tl.int64) bos, eos = tl.load(cu_seqlens + i_n).to(tl.int64), tl.load(cu_seqlens + i_n + 1).to(tl.int64) T = eos - bos NT = tl.cdiv(T, BT) @@ -370,37 +412,49 @@ def chunk_gla_fwd_kernel_o( A += (bos * HV + i_hv) * BT b_o = tl.zeros([BT, BV], dtype=tl.float32) + o_t = i_t * BT + tl.arange(0, BT) + o_v = i_v * BV + tl.arange(0, BV) + o_i = tl.arange(0, BT) + m_t = o_t < T + m_v = o_v < V + m_tv = m_t[:, None] & m_v[None, :] + m_A = m_t[:, None] & (o_i[None, :] < BT) for i_k in range(tl.cdiv(K, BK)): - p_q = tl.make_block_ptr(q, (T, K), (H*K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) - p_g = tl.make_block_ptr(g, (T, K), (HV*K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) + o_k = i_k * BK + tl.arange(0, BK) + m_k = o_k < K + m_qk = m_t[:, None] & m_k[None, :] + p_q = q + o_t[:, None] * (H*K) + o_k[None, :] + p_g = g + o_t[:, None] * (HV*K) + o_k[None, :] if STATE_V_FIRST: - p_h = tl.make_block_ptr(h, (V, K), (K, 1), (i_v * BV, i_k * BK), (BV, BK), (1, 0)) + p_h = h + o_v[:, None] * K + o_k[None, :] + m_h = m_v[:, None] & m_k[None, :] else: - p_h = tl.make_block_ptr(h, (K, V), (V, 1), (i_k * BK, i_v * BV), (BK, BV), (1, 0)) + p_h = h + o_k[:, None] * V + o_v[None, :] + m_h = m_k[:, None] & m_v[None, :] # [BT, BK] - b_q = tl.load(p_q, boundary_check=(0, 1)) + b_q = tl.load(p_q, mask=m_qk, other=0.0) # [BT, BK] - b_g = tl.load(p_g, boundary_check=(0, 1)).to(tl.float32) + b_g = tl.load(p_g, mask=m_qk, other=0.0).to(tl.float32) # [BT, BK] b_qg = (b_q * exp2(b_g)).to(b_q.dtype) - b_h = tl.load(p_h, boundary_check=(0, 1)) + b_h = tl.load(p_h, mask=m_h, other=0.0) if i_k >= 0: if STATE_V_FIRST: - b_o += tl.dot(b_qg, tl.trans(b_h).to(b_qg.dtype)) + b_o = tl.dot(b_qg, tl.trans(b_h).to(b_qg.dtype), b_o) else: - b_o += tl.dot(b_qg, b_h.to(b_qg.dtype)) + b_o = tl.dot(b_qg, b_h.to(b_qg.dtype), b_o) b_o *= scale - p_v = tl.make_block_ptr(v, (T, V), (HV*V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) - p_o = tl.make_block_ptr(o, (T, V), (HV*V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) - p_A = tl.make_block_ptr(A, (T, BT), (HV*BT, 1), (i_t * BT, 0), (BT, BT), (1, 0)) + p_v = v + o_t[:, None] * (HV*V) + o_v[None, :] + p_o = o + o_t[:, None] * (HV*V) + o_v[None, :] + p_A = A + o_t[:, None] * (HV*BT) + o_i[None, :] # [BT, BV] - b_v = tl.load(p_v, boundary_check=(0, 1)) + b_v = tl.load(p_v, mask=m_tv, other=0.0) # [BT, BT] - b_A = tl.load(p_A, boundary_check=(0, 1)) + b_A = tl.load(p_A, mask=m_A, other=0.0) b_A = tl.where(m_s, b_A, 0.).to(b_v.dtype) - b_o += tl.dot(b_A, b_v) - tl.store(p_o, b_o.to(p_o.dtype.element_ty), boundary_check=(0, 1)) + b_o = tl.dot(b_A, b_v, b_o) + tl.store(p_o, b_o.to(p_o.dtype.element_ty), mask=m_tv) @triton.heuristics({ @@ -434,12 +488,13 @@ def chunk_gla_bwd_kernel_intra( NC: tl.constexpr, IS_VARLEN: tl.constexpr, ): - i_kc, i_t, i_bh = tl.program_id(0), tl.program_id(1), tl.program_id(2) + i_kc, i_t = unflatten_program_id(tl.cdiv(K, BK) * NC) + i_bh = tl.program_id(1).to(tl.int64) i_b, i_h = i_bh // H, i_bh % H i_k, i_i = i_kc // NC, i_kc % NC 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.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) + i_n, i_t = tl.load(chunk_indices + i_t * 2).to(tl.int32), tl.load(chunk_indices + i_t * 2 + 1).to(tl.int64) + bos, eos = tl.load(cu_seqlens + i_n).to(tl.int64), tl.load(cu_seqlens + i_n + 1).to(tl.int64) else: bos, eos = i_b * T, i_b * T + T T = eos - bos @@ -449,9 +504,12 @@ def chunk_gla_bwd_kernel_intra( o_k = i_k * BK + tl.arange(0, BK) m_k = o_k < K - p_g = tl.make_block_ptr(g + (bos*H + i_h) * K, (T, K), (H*K, 1), (i_t * BT + i_i * BC, i_k * BK), (BC, BK), (1, 0)) + o_c = i_t * BT + i_i * BC + tl.arange(0, BC) + m_c = o_c < T + m_ck = m_c[:, None] & m_k[None, :] + p_g = g + (bos*H + i_h) * K + o_c[:, None] * (H*K) + o_k[None, :] # [BC, BK] - b_g = tl.load(p_g, boundary_check=(0, 1)) + b_g = tl.load(p_g, mask=m_ck, other=0.0) b_dq = tl.zeros([BC, BK], dtype=tl.float32) if i_i > 0: @@ -459,24 +517,28 @@ def chunk_gla_bwd_kernel_intra( # [BK,] b_gn = tl.load(p_gn, mask=m_k, other=0) for i_j in range(0, i_i): - p_k = tl.make_block_ptr(k+(bos*H+i_h)*K, (T, K), (H*K, 1), (i_t*BT+i_j*BC, i_k * BK), (BC, BK), (1, 0)) - p_gk = tl.make_block_ptr(g+(bos*H+i_h)*K, (T, K), (H*K, 1), (i_t*BT+i_j*BC, i_k * BK), (BC, BK), (1, 0)) - p_dA = tl.make_block_ptr(dA+(bos*H+i_h)*BT, (T, BT), (H*BT, 1), (i_t*BT+i_i*BC, i_j * BC), (BC, BC), (1, 0)) + o_j = i_t * BT + i_j * BC + tl.arange(0, BC) + o_jA = i_j * BC + tl.arange(0, BC) + m_jk = (o_j[:, None] < T) & m_k[None, :] + m_da = m_c[:, None] & (o_jA[None, :] < BT) + p_k = k+(bos*H+i_h)*K + o_j[:, None] * (H*K) + o_k[None, :] + p_gk = g+(bos*H+i_h)*K + o_j[:, None] * (H*K) + o_k[None, :] + p_dA = dA+(bos*H+i_h)*BT + o_c[:, None] * (H*BT) + o_jA[None, :] # [BC, BK] - b_k = tl.load(p_k, boundary_check=(0, 1)) - b_gk = tl.load(p_gk, boundary_check=(0, 1)) + b_k = tl.load(p_k, mask=m_jk, other=0.0) + b_gk = tl.load(p_gk, mask=m_jk, other=0.0) b_kg = b_k * exp2(b_gn[None, :] - b_gk) # [BC, BC] - b_dA = tl.load(p_dA, boundary_check=(0, 1)) + b_dA = tl.load(p_dA, mask=m_da, other=0.0) - b_dq += tl.dot(b_dA, b_kg) + b_dq = tl.dot(b_dA, b_kg, b_dq) b_dq *= exp2(b_g - b_gn[None, :]) o_i = tl.arange(0, BC) m_dA = (i_t * BT + i_i * BC + tl.arange(0, BC)) < T o_dA = bos*H*BT + (i_t * BT + i_i * BC + tl.arange(0, BC)) * H*BT + i_h * BT + i_i * BC p_kj = k + (bos + i_t * BT + i_i * BC) * H*K + i_h * K + o_k p_gkj = g + (bos + i_t * BT + i_i * BC) * H*K + i_h * K + o_k - p_dq = tl.make_block_ptr(dq + (bos*H + i_h) * K, (T, K), (H*K, 1), (i_t * BT + i_i * BC, i_k * BK), (BC, BK), (1, 0)) + p_dq = dq + (bos*H + i_h) * K + o_c[:, None] * (H*K) + o_k[None, :] for j in range(0, min(BC, T - i_t * BT - i_i * BC)): # [BC,] @@ -491,7 +553,7 @@ def chunk_gla_bwd_kernel_intra( b_dq += tl.where(m_i, b_dA[:, None] * b_kj[None, :] * exp2(b_g - b_gkj[None, :]), 0.) p_kj += H*K p_gkj += H*K - tl.store(p_dq, b_dq.to(p_dq.dtype.element_ty), boundary_check=(0, 1)) + tl.store(p_dq, b_dq.to(p_dq.dtype.element_ty), mask=m_ck) tl.debug_barrier() # [BC, BK] @@ -504,26 +566,28 @@ def chunk_gla_bwd_kernel_intra( # [BK,] b_gn = tl.load(p_gn, mask=m_k, other=0) for i_j in range(i_i + 1, NC): - p_q = tl.make_block_ptr(q + (bos*H+i_h)*K, (T, K), (H*K, 1), (i_t*BT+i_j*BC, i_k*BK), (BC, BK), (1, 0)) - p_gq = tl.make_block_ptr(g + (bos*H+i_h)*K, (T, K), (H*K, 1), (i_t*BT+i_j*BC, i_k*BK), (BC, BK), (1, 0)) - p_dA = tl.make_block_ptr(dA + (bos*H+i_h)*BT, (BT, T), (1, H*BT), (i_i*BC, i_t*BT + i_j*BC), (BC, BC), (0, 1)) - o_j = i_t * BT + i_j * BC + o_i + o_iA = i_i * BC + tl.arange(0, BC) m_j = o_j < T + m_jk = m_j[:, None] & m_k[None, :] + m_da = (o_iA[:, None] < BT) & m_j[None, :] + p_q = q + (bos*H+i_h)*K + o_j[:, None] * (H*K) + o_k[None, :] + p_gq = g + (bos*H+i_h)*K + o_j[:, None] * (H*K) + o_k[None, :] + p_dA = dA + (bos*H+i_h)*BT + o_iA[:, None] + o_j[None, :] * (H*BT) # [BC, BK] - b_q = tl.load(p_q, boundary_check=(0, 1)) - b_gq = tl.load(p_gq, boundary_check=(0, 1)) + b_q = tl.load(p_q, mask=m_jk, other=0.0) + b_gq = tl.load(p_gq, mask=m_jk, other=0.0) b_qg = b_q * tl.where(m_j[:, None], exp2(b_gq - b_gn[None, :]), 0) # [BC, BC] - b_dA = tl.load(p_dA, boundary_check=(0, 1)) + b_dA = tl.load(p_dA, mask=m_da, other=0.0) # [BC, BK] # (SY 09/17) important to not use bf16 here to have a good precision. - b_dk += tl.dot(b_dA, b_qg) + b_dk = tl.dot(b_dA, b_qg, b_dk) b_dk *= exp2(b_gn[None, :] - b_g) o_dA = bos*H*BT + (i_t * BT + i_i * BC) * H*BT + i_h * BT + i_i * BC + tl.arange(0, BC) p_qj = q + (bos + i_t * BT + i_i * BC) * H*K + i_h * K + o_k p_gqj = g + (bos + i_t * BT + i_i * BC) * H*K + i_h * K + o_k - p_dk = tl.make_block_ptr(dk + (bos*H+i_h)*K, (T, K), (H*K, 1), (i_t * BT + i_i * BC, i_k * BK), (BC, BK), (1, 0)) + p_dk = dk + (bos*H+i_h)*K + o_c[:, None] * (H*K) + o_k[None, :] for j in range(0, min(BC, T - i_t * BT - i_i * BC)): # [BC,] b_dA = tl.load(dA + o_dA + j * H*BT) @@ -535,7 +599,7 @@ def chunk_gla_bwd_kernel_intra( b_dk += tl.where(m_i, b_dA[:, None] * b_qj[None, :] * exp2(b_gqj[None, :] - b_g), 0.) p_qj += H*K p_gqj += H*K - tl.store(p_dk, b_dk.to(p_dk.dtype.element_ty), boundary_check=(0, 1)) + tl.store(p_dk, b_dk.to(p_dk.dtype.element_ty), mask=m_ck) @triton.heuristics({ @@ -565,28 +629,36 @@ def chunk_gla_bwd_kernel_dA( BV: tl.constexpr, IS_VARLEN: tl.constexpr, ): - i_t, i_bh = tl.program_id(0), tl.program_id(1) + i_t, i_bh = tl.program_id(0).to(tl.int64), tl.program_id(1).to(tl.int64) i_b, i_h = i_bh // H, i_bh % H 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.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) + i_n, i_t = tl.load(chunk_indices + i_t * 2).to(tl.int32), tl.load(chunk_indices + i_t * 2 + 1).to(tl.int64) + bos, eos = tl.load(cu_seqlens + i_n).to(tl.int64), tl.load(cu_seqlens + i_n + 1).to(tl.int64) else: bos, eos = i_b * T, i_b * T + T T = eos - bos b_dA = tl.zeros([BT, BT], dtype=tl.float32) + o_t = i_t * BT + tl.arange(0, BT) + o_i = tl.arange(0, BT) + m_t = o_t < T + m_A = m_t[:, None] & (o_i[None, :] < BT) for i_v in range(tl.cdiv(V, BV)): - p_do = tl.make_block_ptr(do + (bos*H + i_h) * V, (T, V), (H*V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) - p_v = tl.make_block_ptr(v + (bos*H + i_h) * V, (V, T), (1, H*V), (i_v * BV, i_t * BT), (BV, BT), (0, 1)) - b_v = tl.load(p_v, boundary_check=(0, 1)) - b_do = tl.load(p_do, boundary_check=(0, 1)) - - b_dA += tl.dot(b_do, b_v) - - p_dA = tl.make_block_ptr(dA + (bos * H + i_h) * BT, (T, BT), (H*BT, 1), (i_t * BT, 0), (BT, BT), (1, 0)) + o_v = i_v * BV + tl.arange(0, BV) + m_v = o_v < V + m_tv = m_t[:, None] & m_v[None, :] + m_vt = m_v[:, None] & m_t[None, :] + p_do = do + (bos*H + i_h) * V + o_t[:, None] * (H*V) + o_v[None, :] + p_v = v + (bos*H + i_h) * V + o_v[:, None] + o_t[None, :] * (H*V) + b_v = tl.load(p_v, mask=m_vt, other=0.0) + b_do = tl.load(p_do, mask=m_tv, other=0.0) + + b_dA = tl.dot(b_do, b_v, b_dA) + + p_dA = dA + (bos * H + i_h) * BT + o_t[:, None] * (H*BT) + o_i[None, :] m_s = tl.arange(0, BT)[:, None] >= tl.arange(0, BT)[None, :] b_dA = tl.where(m_s, b_dA * scale, 0.) - tl.store(p_dA, b_dA.to(p_dA.dtype.element_ty), boundary_check=(0, 1)) + tl.store(p_dA, b_dA.to(p_dA.dtype.element_ty), mask=m_A) @triton.heuristics({ @@ -624,12 +696,13 @@ def chunk_gla_bwd_kernel_dv( IS_VARLEN: tl.constexpr, STATE_V_FIRST: tl.constexpr, ): - i_v, i_t, i_bh = tl.program_id(0), tl.program_id(1), tl.program_id(2) + i_v, i_t = unflatten_program_id(tl.cdiv(V, BV)) + i_bh = tl.program_id(1).to(tl.int64) i_b, i_h = i_bh // H, i_bh % H 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.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) + i_n, i_t = tl.load(chunk_indices + i_t * 2).to(tl.int32), tl.load(chunk_indices + i_t * 2 + 1).to(tl.int64) + bos, eos = tl.load(cu_seqlens + i_n).to(tl.int64), tl.load(cu_seqlens + i_n + 1).to(tl.int64) T = eos - bos NT = tl.cdiv(T, BT) else: @@ -637,11 +710,18 @@ def chunk_gla_bwd_kernel_dv( i_tg = i_b * NT + i_t bos, eos = i_b * T, i_b * T + T - p_A = tl.make_block_ptr(A + (bos * H + i_h) * BT, (BT, T), (1, H*BT), (0, i_t * BT), (BT, BT), (0, 1)) - p_do = tl.make_block_ptr(do + (bos * H + i_h) * V, (T, V), (H*V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) - p_dv = tl.make_block_ptr(dv + (bos * H + i_h) * V, (T, V), (H*V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) - b_A = tl.load(p_A, boundary_check=(0, 1)) - b_do = tl.load(p_do, boundary_check=(0, 1)) + o_t = i_t * BT + tl.arange(0, BT) + o_v = i_v * BV + tl.arange(0, BV) + o_i = tl.arange(0, BT) + m_t = o_t < T + m_v = o_v < V + m_A = (o_i[:, None] < BT) & m_t[None, :] + m_tv = m_t[:, None] & m_v[None, :] + p_A = A + (bos * H + i_h) * BT + o_i[:, None] + o_t[None, :] * (H*BT) + p_do = do + (bos * H + i_h) * V + o_t[:, None] * (H*V) + o_v[None, :] + p_dv = dv + (bos * H + i_h) * V + o_t[:, None] * (H*V) + o_v[None, :] + b_A = tl.load(p_A, mask=m_A, other=0.0) + b_do = tl.load(p_do, mask=m_tv, other=0.0) b_A = tl.where(tl.arange(0, BT)[:, None] <= tl.arange(0, BT)[None, :], b_A, 0.) # (SY 09/17) important to disallow tf32 here to maintain a good precision. @@ -650,27 +730,29 @@ def chunk_gla_bwd_kernel_dv( for i_k in range(tl.cdiv(K, BK)): o_k = i_k * BK + tl.arange(0, BK) m_k = o_k < K + m_tk = m_t[:, None] & m_k[None, :] + m_kvd = m_k[:, None] & m_v[None, :] - p_k = tl.make_block_ptr(k + (bos * H + i_h) * K, (T, K), (H*K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) - p_gk = tl.make_block_ptr(g + (bos * H + i_h) * K, (T, K), (H*K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) + p_k = k + (bos * H + i_h) * K + o_t[:, None] * (H*K) + o_k[None, :] + p_gk = g + (bos * H + i_h) * K + o_t[:, None] * (H*K) + o_k[None, :] p_gn = g + (bos + min(i_t * BT + BT, T) - 1)*H*K + i_h * K + o_k if STATE_V_FIRST: # dh stored as [V, K]; read a logical [BK, BV] tile via on-the-fly transpose - p_dh = tl.make_block_ptr(dh + (i_tg * H + i_h) * K*V, (K, V), (1, K), (i_k * BK, i_v * BV), (BK, BV), (0, 1)) + p_dh = dh + (i_tg * H + i_h) * K*V + o_k[:, None] + o_v[None, :] * K else: - p_dh = tl.make_block_ptr(dh + (i_tg * H + i_h) * K*V, (K, V), (V, 1), (i_k * BK, i_v * BV), (BK, BV), (1, 0)) + p_dh = dh + (i_tg * H + i_h) * K*V + o_k[:, None] * V + o_v[None, :] - b_k = tl.load(p_k, boundary_check=(0, 1)) - b_gk = tl.load(p_gk, boundary_check=(0, 1)) - b_dh = tl.load(p_dh, boundary_check=(0, 1)) + b_k = tl.load(p_k, mask=m_tk, other=0.0) + b_gk = tl.load(p_gk, mask=m_tk, other=0.0) + b_dh = tl.load(p_dh, mask=m_kvd, other=0.0) b_gn = exp2(tl.load(p_gn, mask=m_k, other=0)[None, :] - b_gk) b_k = (b_k * b_gn).to(b_k.dtype) # [BT, BV] # (SY 09/17) it is ok to have bf16 interchunk gradient contribution here - b_dv += tl.dot(b_k, b_dh.to(b_k.dtype)) + b_dv = tl.dot(b_k, b_dh.to(b_k.dtype), b_dv) - tl.store(p_dv, b_dv.to(p_dv.dtype.element_ty), boundary_check=(0, 1)) + tl.store(p_dv, b_dv.to(p_dv.dtype.element_ty), mask=m_tv) @triton.heuristics({ @@ -685,7 +767,7 @@ def chunk_gla_bwd_kernel_dv( for num_stages in [2, 3, 4] ], key=['BT', 'STATE_V_FIRST', 'K', 'V'], - prune_configs_by={'early_config_prune': _prune_gla_bwd_configs}, + prune_configs_by={'early_config_prune': _prune_gla_bwd_inter_configs}, **autotune_cache_kwargs, ) @triton.jit(do_not_specialize=['T']) @@ -715,12 +797,13 @@ def chunk_gla_bwd_kernel_inter( IS_VARLEN: tl.constexpr, STATE_V_FIRST: tl.constexpr, ): - i_k, i_t, i_bh = tl.program_id(0), tl.program_id(1), tl.program_id(2) + i_k, i_t = unflatten_program_id(tl.cdiv(K, BK)) + i_bh = tl.program_id(1).to(tl.int64) i_b, i_h = i_bh // H, i_bh % H 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.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) + i_n, i_t = tl.load(chunk_indices + i_t * 2).to(tl.int32), tl.load(chunk_indices + i_t * 2 + 1).to(tl.int64) + bos, eos = tl.load(cu_seqlens + i_n).to(tl.int64), tl.load(cu_seqlens + i_n + 1).to(tl.int64) T = eos - bos NT = tl.cdiv(T, BT) else: @@ -743,8 +826,11 @@ def chunk_gla_bwd_kernel_inter( dk2 += (bos * H + i_h) * K dg += (bos * H + i_h) * K - p_gk = tl.make_block_ptr(g, (T, K), (H*K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) - b_gk = tl.load(p_gk, boundary_check=(0, 1)) + o_t = i_t * BT + tl.arange(0, BT) + m_t = o_t < T + m_tk = m_t[:, None] & m_k[None, :] + p_gk = g + o_t[:, None] * (H*K) + o_k[None, :] + b_gk = tl.load(p_gk, mask=m_tk, other=0.0) p_gn = g + (min(T, i_t * BT + BT) - 1) * H*K + o_k b_gn = tl.load(p_gn, mask=m_k, other=0) b_dq = tl.zeros([BT, BK], dtype=tl.float32) @@ -752,55 +838,60 @@ def chunk_gla_bwd_kernel_inter( b_dgk = tl.zeros([BK], dtype=tl.float32) for i_v in range(tl.cdiv(V, BV)): - p_v = tl.make_block_ptr(v, (T, V), (H*V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) - p_do = tl.make_block_ptr(do, (T, V), (H*V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) + o_v = i_v * BV + tl.arange(0, BV) + m_v = o_v < V + m_tv = m_t[:, None] & m_v[None, :] + m_vk = m_v[:, None] & m_k[None, :] + p_v = v + o_t[:, None] * (H*V) + o_v[None, :] + p_do = do + o_t[:, None] * (H*V) + o_v[None, :] if STATE_V_FIRST: # h / dh stored as [V, K] -- the [BV, BK] tile is now a contiguous read - p_h = tl.make_block_ptr(h, (V, K), (K, 1), (i_v * BV, i_k * BK), (BV, BK), (1, 0)) - p_dh = tl.make_block_ptr(dh, (V, K), (K, 1), (i_v * BV, i_k * BK), (BV, BK), (1, 0)) + p_h = h + o_v[:, None] * K + o_k[None, :] + p_dh = dh + o_v[:, None] * K + o_k[None, :] else: - p_h = tl.make_block_ptr(h, (V, K), (1, V), (i_v * BV, i_k * BK), (BV, BK), (0, 1)) - p_dh = tl.make_block_ptr(dh, (V, K), (1, V), (i_v * BV, i_k * BK), (BV, BK), (0, 1)) + p_h = h + o_v[:, None] + o_k[None, :] * V + p_dh = dh + o_v[:, None] + o_k[None, :] * V # [BT, BV] - b_v = tl.load(p_v, boundary_check=(0, 1)) - b_do = tl.load(p_do, boundary_check=(0, 1)) + b_v = tl.load(p_v, mask=m_tv, other=0.0) + b_do = tl.load(p_do, mask=m_tv, other=0.0) # [BV, BK] - b_h = tl.load(p_h, boundary_check=(0, 1)) - b_dh = tl.load(p_dh, boundary_check=(0, 1)) + b_h = tl.load(p_h, mask=m_vk, other=0.0) + b_dh = tl.load(p_dh, mask=m_vk, other=0.0) # [BK] b_dgk += tl.sum(b_h * b_dh, axis=0) # [BT, BK] - b_dq += tl.dot(b_do, b_h.to(b_do.dtype)) - b_dk += tl.dot(b_v, b_dh.to(b_v.dtype)) + b_dq = tl.dot(b_do, b_h.to(b_do.dtype), b_dq) + b_dk = tl.dot(b_v, b_dh.to(b_v.dtype), b_dk) b_dgk *= exp2(b_gn) b_dq *= scale b_dq = b_dq * exp2(b_gk) b_dk = b_dk * exp2(b_gn[None, :] - b_gk) - p_q = tl.make_block_ptr(q, (T, K), (H*K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) - p_k = tl.make_block_ptr(k, (T, K), (H*K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) - p_dq = tl.make_block_ptr(dq, (T, K), (H*K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) - p_dk = tl.make_block_ptr(dk, (T, K), (H*K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) - b_q = tl.load(p_q, boundary_check=(0, 1)) - b_k = tl.load(p_k, boundary_check=(0, 1)) + p_q = q + o_t[:, None] * (H*K) + o_k[None, :] + p_k = k + o_t[:, None] * (H*K) + o_k[None, :] + p_dq = dq + o_t[:, None] * (H*K) + o_k[None, :] + p_dk = dk + o_t[:, None] * (H*K) + o_k[None, :] + b_q = tl.load(p_q, mask=m_tk, other=0.0) + b_k = tl.load(p_k, mask=m_tk, other=0.0) b_dgk += tl.sum(b_dk * b_k, axis=0) - b_dq += tl.load(p_dq, boundary_check=(0, 1)) - b_dk += tl.load(p_dk, boundary_check=(0, 1)) + b_dq += tl.load(p_dq, mask=m_tk, other=0.0) + b_dk += tl.load(p_dk, mask=m_tk, other=0.0) b_dg = b_q * b_dq - b_k * b_dk # tl.debug_barrier() b_dg = b_dg - tl.cumsum(b_dg, axis=0) + tl.sum(b_dg, axis=0)[None, :] + b_dgk[None, :] # Buggy due to strange triton compiler issue. # m_s = tl.where(tl.arange(0, BT)[:, None] <= tl.arange(0, BT)[None, :], 1., 0.) # b_dg = tl.dot(m_s, b_dg, allow_tf32=False) + b_dgk[None, :] - p_dq = tl.make_block_ptr(dq2, (T, K), (H*K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) - p_dk = tl.make_block_ptr(dk2, (T, K), (H*K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) - p_dg = tl.make_block_ptr(dg, (T, K), (H*K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) - tl.store(p_dq, b_dq.to(p_dq.dtype.element_ty), boundary_check=(0, 1)) - tl.store(p_dk, b_dk.to(p_dk.dtype.element_ty), boundary_check=(0, 1)) - tl.store(p_dg, b_dg.to(p_dg.dtype.element_ty), boundary_check=(0, 1)) + p_dq = dq2 + o_t[:, None] * (H*K) + o_k[None, :] + p_dk = dk2 + o_t[:, None] * (H*K) + o_k[None, :] + p_dg = dg + o_t[:, None] * (H*K) + o_k[None, :] + tl.store(p_dq, b_dq.to(p_dq.dtype.element_ty), mask=m_tk) + tl.store(p_dk, b_dk.to(p_dk.dtype.element_ty), mask=m_tk) + tl.store(p_dg, b_dg.to(p_dg.dtype.element_ty), mask=m_tk) +@dispatch('gla') def chunk_gla_fwd_intra_gk( q: torch.Tensor, k: torch.Tensor, @@ -862,7 +953,7 @@ def chunk_gla_fwd_intra_gk( NK = triton.cdiv(K, BK) A_intra = q.new_empty(NK, B, T, H, BC, dtype=torch.float) - grid = (NK, NT * NC, B * H) + grid = (NK * NT * NC, B * H) chunk_gla_fwd_A_kernel_intra_sub_intra_split[grid]( q=q, k=k, @@ -897,6 +988,8 @@ def chunk_gla_fwd_intra_gk( return A +# Use triton only, other backends (cutlass, cute, etc.) need additional torch.compile support +# @dispatch('gla') def chunk_gla_fwd_o_gk( q: torch.Tensor, v: torch.Tensor, @@ -908,6 +1001,7 @@ def chunk_gla_fwd_o_gk( cu_seqlens: torch.LongTensor | None = None, chunk_size: int = 64, chunk_indices: torch.LongTensor | None = None, + use_graph: bool = False, ): B, T, H, K, HV, V = *q.shape, v.shape[2], v.shape[-1] BT = chunk_size @@ -918,7 +1012,7 @@ def chunk_gla_fwd_o_gk( # Please ensure zeros, since vllm will use padding v o = torch.zeros_like(v) - def grid(meta): return (triton.cdiv(V, meta['BV']), NT, B * HV) + def grid(meta): return (triton.cdiv(V, meta['BV']) * NT, B * HV) chunk_gla_fwd_kernel_o[grid]( q=q, v=v, @@ -936,10 +1030,12 @@ def chunk_gla_fwd_o_gk( V=V, BT=BT, STATE_V_FIRST=state_v_first, + USE_GRAPH=use_graph, ) return o +@dispatch('gla') def chunk_gla_bwd_dA( v: torch.Tensor, do: torch.Tensor, @@ -974,6 +1070,7 @@ def chunk_gla_bwd_dA( return dA +@dispatch('gla') def chunk_gla_bwd_dv( k: torch.Tensor, g: torch.Tensor, @@ -993,7 +1090,7 @@ def chunk_gla_bwd_dv( NT = triton.cdiv(T, BT) if cu_seqlens is None else len(chunk_indices) dv = torch.empty_like(do) - def grid(meta): return (triton.cdiv(V, meta['BV']), NT, B * H) + def grid(meta): return (triton.cdiv(V, meta['BV']) * NT, B * H) chunk_gla_bwd_kernel_dv[grid]( k=k, g=g, @@ -1013,6 +1110,7 @@ def chunk_gla_bwd_dv( return dv +@dispatch('gla') def chunk_gla_bwd_dqk_intra( q: torch.Tensor, k: torch.Tensor, @@ -1035,7 +1133,7 @@ def chunk_gla_bwd_dqk_intra( dq = torch.empty_like(q, dtype=torch.float) dk = torch.empty_like(k, dtype=torch.float) - grid = (NK * NC, NT, B * H) + grid = (NK * NC * NT, B * H) chunk_gla_bwd_kernel_intra[grid]( q=q, k=k, @@ -1056,6 +1154,7 @@ def chunk_gla_bwd_dqk_intra( return dq, dk +@dispatch('gla') def chunk_gla_bwd_dqkg( q: torch.Tensor, k: torch.Tensor, @@ -1082,7 +1181,7 @@ def chunk_gla_bwd_dqkg( dg = torch.empty_like(g) dq2 = torch.empty_like(dq) dk2 = torch.empty_like(dk) - def grid(meta): return (triton.cdiv(K, meta['BK']), NT, B * H) + def grid(meta): return (triton.cdiv(K, meta['BK']) * NT, B * H) chunk_gla_bwd_kernel_inter[grid]( q=q, k=k, @@ -1405,7 +1504,7 @@ def chunk_gla( >>> import torch >>> import torch.nn.functional as F >>> from einops import rearrange - >>> from ...ops.gla import chunk_gla + >>> from fla.ops.gla import chunk_gla # inputs with equal lengths >>> B, T, H, K, V = 4, 2048, 4, 512, 512 >>> q = torch.randn(B, T, H, K, device='cuda') diff --git a/build/torch-cuda/ops/gla/fused_recurrent.py b/build/torch-cuda/ops/gla/fused_recurrent.py index d6d695a9ec95bab3912634ccf3903eb9f932ebd5..10fa4e0148ef9b1457b0c211aa015947a47ffbff 100644 --- a/build/torch-cuda/ops/gla/fused_recurrent.py +++ b/build/torch-cuda/ops/gla/fused_recurrent.py @@ -67,7 +67,7 @@ def fused_recurrent_gla( >>> import torch >>> import torch.nn.functional as F >>> from einops import rearrange - >>> from ...ops.gla import fused_recurrent_gla + >>> from fla.ops.gla import fused_recurrent_gla # inputs with equal lengths >>> B, T, H, K, V = 4, 2048, 4, 512, 512 >>> q = torch.randn(B, T, H, K, device='cuda') diff --git a/build/torch-cuda/ops/gsa/chunk.py b/build/torch-cuda/ops/gsa/chunk.py index 21e68ac66504807274ce4e718e83cd0150a5775f..e8b81ead5aca1c13f78845427b14911ec4ed7966 100644 --- a/build/torch-cuda/ops/gsa/chunk.py +++ b/build/torch-cuda/ops/gsa/chunk.py @@ -6,941 +6,13 @@ # https://github.com/fla-org/flash-linear-attention/graphs/contributors import torch -import triton -import triton.language as tl -from einops import reduce -from ...ops.common.chunk_h import chunk_bwd_dh, chunk_fwd_h -from ...ops.gla.chunk import chunk_gla_bwd, chunk_gla_fwd -from ...ops.utils import prepare_chunk_indices +from ...ops.gsa.chunk_bwd import chunk_gsa_bwd +from ...ops.gsa.chunk_fwd import chunk_gsa_fwd +from ...ops.gsa.gate import fused_additive_gate +from ...ops.utils import chunk_local_cumsum, prepare_chunk_indices from ...ops.utils.constant import RCP_LN2 -from ...ops.utils.cumsum import chunk_local_cumsum -from ...ops.utils.op import exp2 -from ...ops.utils.softmax import softmax_bwd, softmax_fwd -from ...utils import autotune_cache_kwargs, input_guard - - -@triton.heuristics({ - 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None, -}) -@triton.autotune( - configs=[ - triton.Config({'BK': BK, 'BV': BV}, num_warps=num_warps, num_stages=num_stages) - for BK in [32, 64] - for BV in [32, 64] - for num_warps in [2, 4, 8] - for num_stages in [2, 3, 4] - ], - key=['BT'], - **autotune_cache_kwargs, -) -@triton.jit(do_not_specialize=['T']) -def chunk_gsa_fwd_k_kernel_inter( - q, - k, - h, - g, - o, - A, - cu_seqlens, - chunk_indices, - scale, - T, - HQ: tl.constexpr, - H: tl.constexpr, - K: tl.constexpr, - V: tl.constexpr, - BT: tl.constexpr, - BK: tl.constexpr, - BV: tl.constexpr, - NG: tl.constexpr, - IS_VARLEN: tl.constexpr, -): - i_v, i_t, i_bh = tl.program_id(0), tl.program_id(1), tl.program_id(2) - i_b, i_hq = i_bh // HQ, i_bh % HQ - i_h = i_hq // NG - 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.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) - T = eos - bos - NT = tl.cdiv(T, BT) - else: - NT = tl.cdiv(T, BT) - i_tg = i_b * NT + i_t - bos, eos = i_b * T, i_b * T + T - - o_i = tl.arange(0, BT) - m_s = o_i[:, None] >= o_i[None, :] - - b_o = tl.zeros([BT, BV], dtype=tl.float32) - b_A = tl.zeros([BT, BT], dtype=tl.float32) - for i_k in range(tl.cdiv(K, BK)): - p_q = tl.make_block_ptr(q + (bos * HQ + i_hq) * K, (T, K), (HQ*K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) - p_k = tl.make_block_ptr(k + (bos * H + i_h) * K, (K, T), (1, H*K), (i_k * BK, i_t * BT), (BK, BT), (0, 1)) - p_h = tl.make_block_ptr(h + (i_tg * H + i_h) * K*V, (K, V), (V, 1), (i_k * BK, i_v * BV), (BK, BV), (1, 0)) - - # [BT, BK] - b_q = tl.load(p_q, boundary_check=(0, 1)) - b_q = (b_q * scale).to(b_q.dtype) - # [BK, BT] - b_k = tl.load(p_k, boundary_check=(0, 1)) - # [BK, BV] - b_h = tl.load(p_h, boundary_check=(0, 1)) - # [BT, BV] - b_o += tl.dot(b_q, b_h) - # [BT, BT] - b_A += tl.dot(b_q, b_k) - p_g = tl.make_block_ptr(g + (bos * H + i_h) * V, (T, V), (H*V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) - p_o = tl.make_block_ptr(o + (bos * HQ + i_hq) * V, (T, V), (HQ*V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) - p_A = tl.make_block_ptr(A + (bos * HQ + i_hq) * BT, (T, BT), (HQ*BT, 1), (i_t * BT, 0), (BT, BT), (1, 0)) - # [BT, BV] - b_g = tl.load(p_g, boundary_check=(0, 1)) - b_o = b_o * exp2(b_g) - tl.store(p_o, b_o.to(p_o.dtype.element_ty), boundary_check=(0, 1)) - - # [BT, BT] - b_A = tl.where(m_s, b_A, 0.) - if i_v == 0: - tl.store(p_A, b_A.to(p_A.dtype.element_ty), boundary_check=(0, 1)) - - -@triton.heuristics({ - 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None, -}) -@triton.jit(do_not_specialize=['T']) -def chunk_gsa_fwd_k_kernel_intra( - v, - g, - o, - A, - cu_seqlens, - chunk_indices, - T, - HQ: tl.constexpr, - H: tl.constexpr, - V: tl.constexpr, - BT: tl.constexpr, - BC: tl.constexpr, - BV: tl.constexpr, - NC: tl.constexpr, - NG: tl.constexpr, - IS_VARLEN: tl.constexpr, -): - i_v, i_c, i_bh = tl.program_id(0), tl.program_id(1), tl.program_id(2) - i_b, i_hq = i_bh // HQ, i_bh % HQ - i_h = i_hq // NG - i_t, i_i = i_c // NC, i_c % NC - 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.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) - T = eos - bos - else: - bos, eos = i_b * T, i_b * T + T - - o_i = tl.arange(0, BC) - o_v = i_v * BV + tl.arange(0, BV) - m_v = o_v < V - - if i_t * BT + i_i * BC >= T: - return - - p_g = tl.make_block_ptr(g + (bos * H + i_h) * V, (T, V), (H*V, 1), (i_t * BT + i_i * BC, i_v * BV), (BC, BV), (1, 0)) - p_gn = g + (bos + min(i_t * BT + i_i * BC, T)) * H*V + i_h * V + o_v - # [BV,] - b_gn = tl.load(p_gn, mask=m_v, other=0) - # [BC, BV] - b_o = tl.zeros([BC, BV], dtype=tl.float32) - for i_j in range(0, i_i): - p_A = tl.make_block_ptr(A + (bos*HQ+i_hq) * BT, (T, BT), (HQ*BT, 1), (i_t*BT+i_i*BC, i_j * BC), (BC, BC), (1, 0)) - p_v = tl.make_block_ptr(v + (bos*H+i_h) * V, (T, V), (H*V, 1), (i_t * BT + i_j * BC, i_v * BV), (BC, BV), (1, 0)) - p_gv = tl.make_block_ptr(g + (bos*H+i_h) * V, (T, V), (H*V, 1), (i_t * BT + i_j * BC, i_v * BV), (BC, BV), (1, 0)) - # [BC, BV] - b_v = tl.load(p_v, boundary_check=(0, 1)) - b_gv = tl.load(p_gv, boundary_check=(0, 1)) - b_vg = (b_v * exp2(b_gn[None, :] - b_gv)).to(b_v.dtype) - # [BC, BC] - b_A = tl.load(p_A, boundary_check=(0, 1)) - b_o += tl.dot(b_A, b_vg) - # [BC, BV] - b_g = tl.load(p_g, boundary_check=(0, 1)) - b_o *= exp2(b_g - b_gn[None, :]) - - o_A = (bos + i_t * BT + i_i * BC + tl.arange(0, BC)) * HQ*BT + i_hq * BT + i_i * BC - m_A = (i_t * BT + i_i * BC + tl.arange(0, BC)) < T - for j in range(0, min(BC, T - i_t * BT - i_i * BC)): - p_v = v + (bos + i_t * BT + i_i * BC + j) * H*V + i_h * V + o_v - p_gv = g + (bos + i_t * BT + i_i * BC + j) * H*V + i_h * V + o_v - # [BC,] - b_A = tl.load(A + o_A + j, mask=m_A, other=0) - # [BV,] - b_v = tl.load(p_v, mask=m_v, other=0).to(tl.float32) - b_gv = tl.load(p_gv, mask=m_v, other=0).to(tl.float32) - # [BC, BV] - b_vg = b_v[None, :] * exp2(b_g - b_gv[None, :]) - # avoid 0 * inf = inf - b_o += tl.where(o_i[:, None] >= j, b_A[:, None] * b_vg, 0.) - p_o = tl.make_block_ptr(o + (bos*HQ + i_hq) * V, (T, V), (HQ*V, 1), (i_t * BT + i_i * BC, i_v * BV), (BC, BV), (1, 0)) - b_o += tl.load(p_o, boundary_check=(0, 1)) - tl.store(p_o, b_o.to(p_o.dtype.element_ty), boundary_check=(0, 1)) - - -@triton.heuristics({ - 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None, -}) -@triton.autotune( - configs=[ - triton.Config({}, num_warps=num_warps) - for num_warps in [2, 4, 8] - ], - key=['BT'], - **autotune_cache_kwargs, -) -@triton.jit(do_not_specialize=['T']) -def chunk_gsa_bwd_k_kernel_dA( - v, - g, - do, - dA, - chunk_indices, - cu_seqlens, - scale, - T, - B: tl.constexpr, - HQ: tl.constexpr, - H: tl.constexpr, - V: tl.constexpr, - BT: tl.constexpr, - BC: tl.constexpr, - BV: tl.constexpr, - NC: tl.constexpr, - NG: tl.constexpr, - IS_VARLEN: tl.constexpr, -): - i_v, i_c, i_bh = tl.program_id(0), tl.program_id(1), tl.program_id(2) - i_b, i_hq = i_bh // HQ, i_bh % HQ - i_h = i_hq // NG - i_t, i_i, i_j = i_c // (NC * NC), (i_c % (NC * NC)) // NC, (i_c % (NC * NC)) % NC - 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.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) - all = T - T = eos - bos - else: - bos, eos = i_b * T, i_b * T + T - all = B * T - - o_v = i_v * BV + tl.arange(0, BV) - m_v = o_v < V - - if i_t * BT + i_i * BC >= T: - return - - p_dA = tl.make_block_ptr(dA+((i_v*all+bos)*HQ+i_hq)*BT, (T, BT), (HQ*BT, 1), (i_t*BT+i_i*BC, i_j*BC), (BC, BC), (1, 0)) - - # [BC, BC] - b_dA = tl.zeros([BC, BC], dtype=tl.float32) - if i_i > i_j: - p_v = tl.make_block_ptr(v + (bos*H+i_h) * V, (V, T), (1, H*V), (i_v * BV, i_t*BT + i_j*BC), (BV, BC), (0, 1)) - p_gv = tl.make_block_ptr(g + (bos*H+i_h) * V, (V, T), (1, H*V), (i_v * BV, i_t*BT + i_j*BC), (BV, BC), (0, 1)) - p_gn = g + (bos + i_t*BT + i_i*BC) * H*V + i_h * V + o_v - p_g = tl.make_block_ptr(g + (bos*H+i_h) * V, (T, V), (H*V, 1), (i_t*BT + i_i*BC, i_v*BV), (BC, BV), (1, 0)) - p_do = tl.make_block_ptr(do + (bos*HQ+i_hq) * V, (T, V), (HQ*V, 1), (i_t*BT + i_i*BC, i_v*BV), (BC, BV), (1, 0)) - # [BV,] - b_gn = tl.load(p_gn, mask=m_v, other=0.) - # [BC, BV] - b_g = tl.load(p_g, boundary_check=(0, 1)) - b_do = tl.load(p_do, boundary_check=(0, 1)) - b_do = (b_do * exp2(b_g - b_gn[None, :])).to(b_do.dtype) - # [BV, BC] - b_v = tl.load(p_v, boundary_check=(0, 1)) - b_gv = tl.load(p_gv, boundary_check=(0, 1)) - b_vg = (b_v * exp2(b_gn[:, None] - b_gv)).to(b_v.dtype) - # [BC, BC] - b_dA = tl.dot(b_do, b_vg) * scale - elif i_i == i_j: - p_g = tl.make_block_ptr(g + (bos*H + i_h) * V, (T, V), (H*V, 1), (i_t*BT + i_i*BC, i_v*BV), (BC, BV), (1, 0)) - p_do = tl.make_block_ptr(do + (bos*HQ + i_hq) * V, (T, V), (HQ*V, 1), (i_t*BT + i_i*BC, i_v*BV), (BC, BV), (1, 0)) - p_v = v + (bos + i_t*BT + i_j*BC) * H*V + i_h * V + o_v - p_gv = g + (bos + i_t*BT + i_j*BC) * H*V + i_h * V + o_v - # [BC, BV] - b_g = tl.load(p_g, boundary_check=(0, 1)) - b_do = tl.load(p_do, boundary_check=(0, 1)) * scale - m_v = o_v < V - - o_i = tl.arange(0, BC) - # [BC, BC] - m_dA = o_i[:, None] >= o_i[None, :] - for j in range(0, min(BC, T - i_t * BT - i_j * BC)): - # [BV,] - b_v = tl.load(p_v, mask=m_v, other=0).to(tl.float32) - b_gv = tl.load(p_gv, mask=m_v, other=0).to(tl.float32) - # [BC,] - b_dAj = tl.sum(b_do * b_v[None, :] * exp2(b_g - b_gv[None, :]), 1) - b_dA = tl.where((o_i == j)[None, :], b_dAj[:, None], b_dA) - - p_v += H*V - p_gv += H*V - b_dA = tl.where(m_dA, b_dA, 0.) - tl.store(p_dA, b_dA.to(dA.dtype.element_ty), boundary_check=(0, 1)) - - -@triton.heuristics({ - 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None, -}) -@triton.autotune( - configs=[ - triton.Config({}, num_warps=num_warps, num_stages=num_stages) - for num_warps in [2, 4] - for num_stages in [2, 3, 4] - ], - key=['BT'], - **autotune_cache_kwargs, -) -@triton.jit(do_not_specialize=['T']) -def chunk_gsa_bwd_k_kernel_dqkvg( - q, - k, - v, - h, - g, - A, - do, - dh, - dq, - dk, - dv, - dg, - dgv, - dA, - cu_seqlens, - chunk_indices, - scale, - T, - B: tl.constexpr, - HQ: tl.constexpr, - H: tl.constexpr, - K: tl.constexpr, - V: tl.constexpr, - BT: tl.constexpr, - BK: tl.constexpr, - BV: tl.constexpr, - NG: tl.constexpr, - IS_VARLEN: tl.constexpr, -): - i_k, i_t, i_bh = tl.program_id(0), tl.program_id(1), tl.program_id(2) - i_b, i_hq = i_bh // HQ, i_bh % HQ - i_h = i_hq // NG - 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.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) - all = T - T = eos - bos - NT = tl.cdiv(T, BT) - else: - NT = tl.cdiv(T, BT) - i_tg = i_b * NT + i_t - bos, eos = i_b * T, i_b * T + T - all = B * T - - o_i = tl.arange(0, BT) - o_t = min(i_t * BT + BT, T) - m_s = o_i[:, None] >= o_i[None, :] - - p_q = tl.make_block_ptr(q + (bos*HQ+i_hq) * K, (T, K), (HQ*K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) - p_k = tl.make_block_ptr(k + (bos*H+i_h) * K, (T, K), (H*K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) - p_A = tl.make_block_ptr(A + ((i_k*all+bos)*HQ+i_hq)*BT, (T, BT), (HQ*BT, 1), (i_t * BT, 0), (BT, BT), (1, 0)) - - # [BT, BK] - b_q = tl.load(p_q, boundary_check=(0, 1)) - b_k = tl.load(p_k, boundary_check=(0, 1)) - # [BT, BT] - b_A = tl.dot((b_q * scale).to(b_q.dtype), tl.trans(b_k)) - b_A = tl.where(m_s, b_A, 0.) - tl.store(p_A, b_A.to(p_A.dtype.element_ty), boundary_check=(0, 1)) - - b_dq = tl.zeros([BT, BK], dtype=tl.float32) - b_dk = tl.zeros([BT, BK], dtype=tl.float32) - for i_v in range(tl.cdiv(V, BV)): - o_v = i_v * BV + tl.arange(0, BV) - p_v = tl.make_block_ptr(v + (bos*H+i_h)*V, (T, V), (H*V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) - p_g = tl.make_block_ptr(g + (bos*H+i_h)*V, (T, V), (H*V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) - p_gn = g + (bos + o_t - 1) * H*V + i_h * V + o_v - p_do = tl.make_block_ptr(do + (bos*HQ+i_hq)*V, (T, V), (HQ*V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) - p_dv = tl.make_block_ptr(dv + ((i_k*all+bos)*HQ+i_hq)*V, (T, V), (HQ*V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) - p_dg = tl.make_block_ptr(dg + (bos*HQ+i_hq)*V, (T, V), (HQ*V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) - p_dgv = tl.make_block_ptr(dgv+((i_k*all+bos)*HQ+i_hq)*V, (T, V), (HQ*V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) - p_h = tl.make_block_ptr(h + (i_tg * H + i_h) * K*V, (V, K), (1, V), (i_v * BV, i_k * BK), (BV, BK), (0, 1)) - p_dh = tl.make_block_ptr(dh + (i_tg * HQ + i_hq) * K*V, (K, V), (V, 1), (i_k * BK, i_v * BV), (BK, BV), (1, 0)) - m_v = o_v < V - - # [BV,] - b_gn = tl.load(p_gn, mask=m_v, other=0) - # [BT, BV] - b_v = tl.load(p_v, boundary_check=(0, 1)) - b_g = tl.load(p_g, boundary_check=(0, 1)) - b_gv = exp2(b_gn[None, :] - b_g) - # [BV, BK] - b_h = tl.load(p_h, boundary_check=(0, 1)) - # [BT, BV] - b_do = tl.load(p_do, boundary_check=(0, 1)) - b_do = (b_do * exp2(b_g)).to(b_do.dtype) - # [BK, BV] - b_dh = tl.load(p_dh, boundary_check=(0, 1)) - # [BV] - b_dg = tl.sum(tl.trans(b_h) * b_dh, 0) * exp2(b_gn) - - b_dh = b_dh.to(b_k.dtype) - # [BT, BK] - b_dq += tl.dot(b_do, b_h.to(b_k.dtype)) * scale - b_dk += tl.dot((b_v * b_gv).to(b_v.dtype), tl.trans(b_dh)) - # [BT, BV] - b_dv = tl.dot(b_k, b_dh) * b_gv - # [BV] - b_dg += tl.sum(b_dv * b_v, 0) - - if i_k == 0: - b_dgv = tl.load(p_dg, boundary_check=(0, 1)) + b_dg[None, :] - else: - b_dgv = tl.zeros([BT, BV], dtype=tl.float32) + b_dg[None, :] - - tl.store(p_dgv, b_dgv.to(p_dgv.dtype.element_ty), boundary_check=(0, 1)) - tl.store(p_dv, b_dv.to(p_dv.dtype.element_ty), boundary_check=(0, 1)) - p_dA = tl.make_block_ptr(dA + (bos*HQ + i_hq) * BT, (T, BT), (HQ*BT, 1), (i_t * BT, 0), (BT, BT), (1, 0)) - p_dq = tl.make_block_ptr(dq + (bos*HQ + i_hq) * K, (T, K), (HQ*K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) - p_dk = tl.make_block_ptr(dk + (bos*HQ + i_hq) * K, (T, K), (HQ*K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) - # [BT, BT] - b_dA = tl.load(p_dA, boundary_check=(0, 1)) - # [BT, BK] - b_dq += tl.dot(b_dA, b_k) - b_dk += tl.dot(tl.trans(b_dA).to(b_k.dtype), b_q) - - tl.store(p_dq, b_dq.to(p_dq.dtype.element_ty), boundary_check=(0, 1)) - tl.store(p_dk, b_dk.to(p_dk.dtype.element_ty), boundary_check=(0, 1)) - - -@triton.heuristics({ - 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None, -}) -@triton.jit(do_not_specialize=['T']) -def chunk_gsa_bwd_k_kernel_intra_dvg( - v, - g, - o, - A, - do, - dv, - dg, - cu_seqlens, - chunk_indices, - T, - HQ: tl.constexpr, - H: tl.constexpr, - V: tl.constexpr, - BT: tl.constexpr, - BC: tl.constexpr, - BV: tl.constexpr, - NC: tl.constexpr, - NG: tl.constexpr, - IS_VARLEN: tl.constexpr, -): - i_v, i_c, i_bh = tl.program_id(0), tl.program_id(1), tl.program_id(2) - i_b, i_hq = i_bh // HQ, i_bh % HQ - i_h = i_hq // NG - i_t, i_i = i_c // NC, i_c % NC - 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.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) - T = eos - bos - else: - bos, eos = i_b * T, i_b * T + T - - o_i = tl.arange(0, BC) - o_v = i_v * BV + tl.arange(0, BV) - m_v = o_v < V - - if i_t * BT + i_i * BC >= T: - return - - p_gv = tl.make_block_ptr(g + (bos*H+i_h)*V, (T, V), (H*V, 1), (i_t * BT + i_i * BC, i_v * BV), (BC, BV), (1, 0)) - p_gn = g + (bos + min(i_t * BT + i_i * BC + BC, T)-1)*H*V + i_h*V + o_v - # [BV,] - b_gn = tl.load(p_gn, mask=m_v, other=0) - # [BC, BV] - b_gv = tl.load(p_gv, boundary_check=(0, 1)) - b_dv = tl.zeros([BC, BV], dtype=tl.float32) - for i_j in range(i_i + 1, NC): - p_g = tl.make_block_ptr(g + (bos*H+i_h) * V, (T, V), (H*V, 1), (i_t * BT + i_j * BC, i_v * BV), (BC, BV), (1, 0)) - p_A = tl.make_block_ptr(A + (bos*HQ+i_hq) * BT, (BT, T), (1, HQ*BT), (i_i*BC, i_t*BT + i_j*BC), (BC, BC), (0, 1)) - p_do = tl.make_block_ptr(do + (bos*HQ+i_hq) * V, (T, V), (HQ*V, 1), (i_t*BT + i_j*BC, i_v*BV), (BC, BV), (1, 0)) - - m_j = (i_t * BT + i_j * BC + o_i) < T - # [BC, BV] - b_g = tl.load(p_g, boundary_check=(0, 1)) - b_do = tl.load(p_do, boundary_check=(0, 1)) * tl.where(m_j[:, None], exp2(b_g - b_gn[None, :]), 0) - # [BC, BC] - b_A = tl.load(p_A, boundary_check=(0, 1)) - # [BC, BV] - b_dv += tl.dot(b_A, b_do.to(b_A.dtype)) - b_dv *= exp2(b_gn[None, :] - b_gv) - - p_g = g + (bos + i_t * BT + i_i * BC) * H*V + i_h * V + o_v - p_A = A + (bos + i_t*BT + i_i*BC) * HQ*BT + i_hq * BT + i_i * BC + o_i - p_do = do + (bos + i_t*BT + i_i*BC) * HQ*V + i_hq * V + o_v - for j in range(0, min(BC, T - i_t * BT - i_i * BC)): - # [BC,] - b_A = tl.load(p_A) - # [BV,] - b_g = tl.load(p_g, mask=m_v, other=0) - b_do = tl.load(p_do, mask=m_v, other=0) - # [BC, BV] - m_i = o_i[:, None] <= j - b_dv += tl.where(m_i, exp2(b_g[None, :] - b_gv) * b_A[:, None] * b_do[None, :], 0.) - - p_g += H * V - p_A += HQ * BT - p_do += HQ * V - p_o = tl.make_block_ptr(o + (bos*HQ+i_hq)*V, (T, V), (HQ*V, 1), (i_t*BT + i_i*BC, i_v*BV), (BC, BV), (1, 0)) - p_v = tl.make_block_ptr(v + (bos*H+i_h)*V, (T, V), (H*V, 1), (i_t*BT + i_i*BC, i_v*BV), (BC, BV), (1, 0)) - p_do = tl.make_block_ptr(do + (bos*HQ+i_hq)*V, (T, V), (HQ*V, 1), (i_t*BT + i_i*BC, i_v*BV), (BC, BV), (1, 0)) - p_dv = tl.make_block_ptr(dv + (bos*HQ+i_hq)*V, (T, V), (HQ*V, 1), (i_t*BT + i_i*BC, i_v*BV), (BC, BV), (1, 0)) - p_dg = tl.make_block_ptr(dg + (bos*HQ+i_hq)*V, (T, V), (HQ*V, 1), (i_t*BT + i_i*BC, i_v*BV), (BC, BV), (1, 0)) - - b_o = tl.load(p_o, boundary_check=(0, 1)).to(tl.float32) - b_v = tl.load(p_v, boundary_check=(0, 1)).to(tl.float32) - b_do = tl.load(p_do, boundary_check=(0, 1)).to(tl.float32) - b_dv = b_dv + tl.load(p_dv, boundary_check=(0, 1)).to(tl.float32) - b_dg = b_o * b_do - b_v * b_dv - tl.store(p_dv, b_dv.to(p_dv.dtype.element_ty), boundary_check=(0, 1)) - tl.store(p_dg, b_dg.to(p_dg.dtype.element_ty), boundary_check=(0, 1)) - - -def chunk_gsa_fwd_v( - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - g: torch.Tensor, - scale: float = 1., - initial_state: torch.Tensor | None = None, - output_final_state: bool = False, - cu_seqlens: torch.LongTensor | None = None, - chunk_size: int = 64, - chunk_indices: torch.LongTensor | None = None, -) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: - _, A, h, ht, o = chunk_gla_fwd( - q=q, - k=k, - v=v, - g=None, - g_cumsum=g, - scale=scale, - initial_state=initial_state, - output_final_state=output_final_state, - cu_seqlens=cu_seqlens, - chunk_size=chunk_size, - chunk_indices=chunk_indices, - ) - return A, h, ht, o - - -def chunk_gsa_fwd_k( - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - g: torch.Tensor, - h0: torch.Tensor | None = None, - output_final_state: bool = False, - scale: float = 1., - cu_seqlens: torch.LongTensor | None = None, - chunk_size: int = 64, - chunk_indices: torch.LongTensor | None = None, -) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: - B, T, H, K, V = *k.shape, v.shape[-1] - BT = chunk_size - BC = min(16, BT) - BV = min(64, triton.next_power_of_2(V)) - HQ = q.shape[2] - - if chunk_indices is None: - if cu_seqlens is not None: - chunk_indices = prepare_chunk_indices(cu_seqlens, BT) - else: - chunk_indices = None - NT = triton.cdiv(T, BT) if cu_seqlens is None else len(chunk_indices) - NC = triton.cdiv(BT, BC) - NG = HQ // H - - h, ht = chunk_fwd_h( - k=k, - v=v, - g=None, - gk=None, - gv=g, - h0=h0, - output_final_state=output_final_state, - cu_seqlens=cu_seqlens, - chunk_size=BT, - states_in_fp32=False, - ) - o = v.new_empty(B, T, HQ, V) - A = q.new_empty(B, T, HQ, BT) - def grid(meta): return (triton.cdiv(V, meta['BV']), NT, B * HQ) - chunk_gsa_fwd_k_kernel_inter[grid]( - q, - k, - h, - g, - o, - A, - cu_seqlens=cu_seqlens, - chunk_indices=chunk_indices, - scale=scale, - T=T, - HQ=HQ, - H=H, - K=K, - V=V, - BT=BT, - NG=NG, - ) - - def grid(meta): return (triton.cdiv(V, meta['BV']), NT * NC, B * HQ) - chunk_gsa_fwd_k_kernel_intra[grid]( - v, - g, - o, - A, - cu_seqlens=cu_seqlens, - chunk_indices=chunk_indices, - T=T, - HQ=HQ, - H=H, - V=V, - BT=BT, - BC=BC, - BV=BV, - NC=NC, - NG=NG, - num_warps=4, - num_stages=2, - ) - return A, h, ht, o - - -def chunk_gsa_bwd_v( - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - g: torch.Tensor, - h0: torch.Tensor, - h: torch.Tensor, - A: torch.Tensor, - do: torch.Tensor, - dht: torch.Tensor, - dg: torch.Tensor, - scale: float = 1., - cu_seqlens: torch.LongTensor | None = None, - chunk_size: int = 64, - chunk_indices: torch.LongTensor | None = None, -): - dq, dk, dv, dg, dh0 = chunk_gla_bwd( - q=q, - k=k, - v=v, - g=None, - g_cumsum=g, - scale=scale, - initial_state=h0, - h=h, - A=A, - do=do, - dht=dht, - cu_seqlens=cu_seqlens, - chunk_size=chunk_size, - chunk_indices=chunk_indices, - ) - return dq, dk, dv, dg, dh0 - - -def chunk_gsa_bwd_k( - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - g: torch.Tensor, - h: torch.Tensor, - h0: torch.Tensor, - o: torch.Tensor, - do: torch.Tensor, - dht: torch.Tensor, - dg: torch.Tensor, - scale: float = 1., - cu_seqlens: torch.LongTensor | None = None, - chunk_size: int = 64, - chunk_indices: torch.LongTensor | None = None, -): - B, T, H, K, V = *k.shape, v.shape[-1] - BT = chunk_size - BC = min(16, BT) - BK = min(64, triton.next_power_of_2(K)) - BV = min(64, triton.next_power_of_2(V)) - HQ = q.shape[2] - - if chunk_indices is None: - if cu_seqlens is not None: - chunk_indices = prepare_chunk_indices(cu_seqlens, BT) - else: - chunk_indices = None - NT = triton.cdiv(T, BT) if cu_seqlens is None else len(chunk_indices) - NC = triton.cdiv(BT, BC) - NK = triton.cdiv(K, BK) - NV = triton.cdiv(V, BV) - NG = HQ // H - - if h is None: - h, _ = chunk_fwd_h( - k=k, - v=v, - g=None, - gk=None, - gv=g, - h0=h0, - output_final_state=False, - cu_seqlens=cu_seqlens, - chunk_size=BT, - states_in_fp32=False, - ) - dh, dh0 = chunk_bwd_dh( - q=q, - k=k, - v=v, - g=None, - gk=None, - gv=g, - do=do, - h0=h0, - dht=dht, - scale=scale, - cu_seqlens=cu_seqlens, - chunk_size=BT, - states_in_fp32=True, - ) - dA = q.new_empty(NV, B, T, HQ, BT) - grid = (NV, NT * NC * NC, B * HQ) - chunk_gsa_bwd_k_kernel_dA[grid]( - v, - g, - do, - dA, - cu_seqlens=cu_seqlens, - chunk_indices=chunk_indices, - scale=scale, - T=T, - B=B, - HQ=HQ, - H=H, - V=V, - BT=BT, - BC=BC, - BV=BV, - NC=NC, - NG=NG, - ) - dA = dA.sum(0, dtype=dA.dtype) - - A = do.new_empty(NK, B, T, HQ, BT) - dq = torch.empty_like(q) - dk = k.new_empty(B, T, HQ, K) - dv = v.new_empty(NK, B, T, HQ, V) - dgv = g.new_empty(NK, B, T, HQ, V, dtype=torch.float) - grid = (NK, NT, B * HQ) - chunk_gsa_bwd_k_kernel_dqkvg[grid]( - q, - k, - v, - h, - g, - A, - do, - dh, - dq, - dk, - dv, - dg, - dgv, - dA, - cu_seqlens=cu_seqlens, - chunk_indices=chunk_indices, - scale=scale, - T=T, - B=B, - HQ=HQ, - H=H, - K=K, - V=V, - BT=BT, - BK=BK, - BV=BV, - NG=NG, - ) - A = A.sum(0, dtype=A.dtype) - dv = dv.sum(0, dtype=dv.dtype) - dgv = dgv.sum(0, dtype=dgv.dtype) - - def grid(meta): return (triton.cdiv(V, meta['BV']), NT * NC, B * HQ) - chunk_gsa_bwd_k_kernel_intra_dvg[grid]( - v, - g, - o, - A, - do, - dv, - dg, - cu_seqlens=cu_seqlens, - chunk_indices=chunk_indices, - T=T, - HQ=HQ, - H=H, - V=V, - BT=BT, - BC=BC, - BV=BV, - NC=NC, - NG=NG, - num_warps=4, - num_stages=2, - ) - dg = dgv.add_( - chunk_local_cumsum( - dg, - chunk_size=BT, - reverse=True, - cu_seqlens=cu_seqlens, - chunk_indices=chunk_indices, - ) - ) - - return dq, dk, dv, dg, dh0 - - -def chunk_gsa_fwd( - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - s: torch.Tensor, - g: torch.Tensor, - initial_state: tuple[torch.Tensor, torch.Tensor] | None = None, - output_final_state: bool = False, - scale: float = 1., - cu_seqlens: torch.LongTensor | None = None, - chunk_size: int = 64, - chunk_indices: torch.LongTensor | None = None, -) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: - hk0, hv0 = None, None - if initial_state is not None: - hk0, hv0 = initial_state - Ak, hk, hkt, ok = chunk_gsa_fwd_k( - q=q, - k=k, - v=s, - g=g, - h0=hk0, - output_final_state=output_final_state, - scale=scale, - cu_seqlens=cu_seqlens, - chunk_size=chunk_size, - chunk_indices=chunk_indices, - ) - - # p is kept in fp32 for safe softmax backward - p = softmax_fwd(ok, dtype=torch.float) - - qv = p.to(q.dtype) - Av, hv, hvt, ov = chunk_gsa_fwd_v( - q=qv, - k=s, - v=v, - g=g, - scale=1., - initial_state=hv0, - output_final_state=output_final_state, - cu_seqlens=cu_seqlens, - chunk_size=chunk_size, - chunk_indices=chunk_indices, - ) - return Ak, hk, hkt, ok, p, Av, hv, hvt, ov - - -def chunk_gsa_bwd( - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - s: torch.Tensor, - g: torch.Tensor, - ok: torch.Tensor, - p: torch.Tensor, - A: tuple[torch.Tensor, torch.Tensor], - h: tuple[torch.Tensor, torch.Tensor], - initial_state: tuple[torch.Tensor, torch.Tensor] | None, - scale: float, - do: torch.Tensor, - dht: tuple[torch.Tensor, torch.Tensor], - cu_seqlens: torch.LongTensor | None = None, - chunk_size: int = 64, - chunk_indices: torch.LongTensor | None = None, -): - hk0, hv0 = None, None - if initial_state is not None: - hk0, hv0 = initial_state - - _, Av = A - hk, hv = h - dhkt, dhvt = dht - - qv = p.to(q.dtype) - dqv, dsv, dv, dg, dhv0 = chunk_gsa_bwd_v( - q=qv, - k=s, - v=v, - g=g, - h0=hv0, - h=hv, - A=Av, - do=do, - dht=dhvt, - dg=None, - scale=1., - cu_seqlens=cu_seqlens, - chunk_size=chunk_size, - chunk_indices=chunk_indices, - ) - - # softmax gradient, equivalent to: - # dok = qv * (dqv - (qv * dqv).sum(-1, True)) - dok = softmax_bwd(p, dqv, dtype=ok.dtype) - - dq, dk, dsk, dg, dhk0 = chunk_gsa_bwd_k( - q=q, - k=k, - v=s, - g=g, - h0=hk0, - h=hk, - o=ok, - do=dok, - dht=dhkt, - dg=dg, - scale=scale, - cu_seqlens=cu_seqlens, - chunk_size=chunk_size, - chunk_indices=chunk_indices, - ) - - ds = dsv.add_(dsk) - if q.shape[1] != k.shape[1]: - dk, dv, ds, dg = map(lambda x: reduce(x, 'b (h g) ... -> b h ...', 'sum', h=k.shape[1]), (dk, dv, ds, dg)) - dg = dg.to(s.dtype) - return dq, dk, dv, ds, dg, dhk0, dhv0 +from ...utils import input_guard class ChunkGSAFunction(torch.autograd.Function): @@ -968,23 +40,23 @@ class ChunkGSAFunction(torch.autograd.Function): if cu_seqlens is not None: chunk_indices = prepare_chunk_indices( - cu_seqlens, - chunk_size, + cu_seqlens=cu_seqlens, + chunk_size=chunk_size, cu_seqlens_cpu=cu_seqlens_cpu, ) else: chunk_indices = None - # Pre-scale by RCP_LN2 so downstream kernels can use exp2 directly. + # pre-scale gates so both stages use exp2 with the same cumulative values. g_org = g g = chunk_local_cumsum( - g, - chunk_size, + g=g, + chunk_size=chunk_size, scale=RCP_LN2, cu_seqlens=cu_seqlens, chunk_indices=chunk_indices, ) - Ak, hk, hkt, ok, p, Av, hv, hvt, ov = chunk_gsa_fwd( + hk, hkt, ok, p, Av, hv, hvt, ov = chunk_gsa_fwd( q=q, k=k, v=v, @@ -999,14 +71,9 @@ class ChunkGSAFunction(torch.autograd.Function): ) if checkpoint_level >= 1: - del g g = g_org if checkpoint_level > 1: - del hk - del hv hk, hv = None, None - else: - hk0, hv0 = None, None ctx.save_for_backward( q, @@ -1053,8 +120,8 @@ class ChunkGSAFunction(torch.autograd.Function): if ctx.checkpoint_level >= 1: g = chunk_local_cumsum( - g, - chunk_size, + g=g, + chunk_size=chunk_size, scale=RCP_LN2, cu_seqlens=cu_seqlens, chunk_indices=chunk_indices, @@ -1067,7 +134,7 @@ class ChunkGSAFunction(torch.autograd.Function): g=g, ok=ok, p=p, - A=(None, Av), + A=Av, h=(hk, hv), initial_state=(hk0, hv0), scale=scale, @@ -1088,9 +155,9 @@ def chunk_gsa( s: torch.Tensor, g: torch.Tensor | None = None, scale: float | None = None, - initial_state: tuple[torch.Tensor] | None = None, - output_final_state: bool | None = False, - checkpoint_level: int | None = 2, + initial_state: tuple[torch.Tensor, torch.Tensor] | None = None, + output_final_state: bool = False, + checkpoint_level: int = 2, cu_seqlens: torch.LongTensor | None = None, cu_seqlens_cpu: torch.LongTensor | None = None, **kwargs, @@ -1107,7 +174,7 @@ def chunk_gsa( s (torch.Tensor): slot representations of shape `[B, T, H, M]`. g (Optional[torch.Tensor]): - Forget gates of shape `[B, T, H, M]` applied to keys. + Forget gates of shape `[B, T, H, M]` applied along the slot dimension. If not provided, this function is equivalent to vanilla ABC. Default: `None`. scale (Optional[float]): Scale factor for attention scores. @@ -1124,7 +191,7 @@ def chunk_gsa( Default: `2`. - Level `0`: no memory saved, no recomputation. - Level `1`: recompute the fp32 cumulative values during backward. - - Level `2`: recompute the fp32 cumulative values and forward hidden states during backward. + - Level `2`: recompute cumulative gates and chunk states during backward. cu_seqlens (torch.LongTensor): Cumulative sequence lengths of shape `[N+1]` used for variable-length training, consistent with the FlashAttention API. @@ -1133,16 +200,16 @@ def chunk_gsa( Returns: o (torch.Tensor): - Outputs of shape `[B, T, H, V]`. + Outputs of shape `[B, T, HQ, V]`. final_state (Tuple[torch.Tensor]): Final state tuple having tensors of shape `[N, H, K, M]` and `[N, H, M, V]` if `output_final_state=True`. - `None` otherwise. + Both entries are `None` otherwise. Examples:: >>> import torch >>> import torch.nn.functional as F >>> from einops import rearrange - >>> from ...ops.gsa import fused_recurrent_gsa + >>> from fla.ops.gsa import chunk_gsa # inputs with equal lengths >>> B, T, H, K, V, M = 4, 2048, 4, 512, 512, 64 >>> q = torch.randn(B, T, H, K, device='cuda') @@ -1175,7 +242,7 @@ def chunk_gsa( "head_first has been removed. Inputs must be in `[B, T, H, ...]` format.", ) chunk_size = kwargs.pop('chunk_size', None) - if chunk_size is not None and chunk_size != 2 ** (chunk_size.bit_length() - 1): + if chunk_size is not None and (chunk_size <= 0 or chunk_size & (chunk_size - 1)): raise ValueError(f"`chunk_size` must be a power of 2, got {chunk_size}.") if cu_seqlens is not None: if q.shape[0] != 1: @@ -1190,16 +257,25 @@ def chunk_gsa( ) assert checkpoint_level in [0, 1, 2] if g is None: - # TODO: this 3 steps took huge amount of time, ought to be optimized - z = s.float().logcumsumexp(2) - g = torch.cat((z[:, :, :1], z[:, :, :-1]), 1) - z - s = torch.exp(s - z).to(k.dtype) + s, g, _ = fused_additive_gate( + x=s, + cu_seqlens=cu_seqlens, + output_dtype=k.dtype, + gate_dtype=torch.float32, + ) if scale is None: scale = q.shape[-1] ** -0.5 hk0, hv0 = None, None if initial_state is not None: hk0, hv0 = initial_state + H, HQ = k.shape[2], q.shape[2] + if HQ < H or HQ % H: + raise ValueError(f"The number of query heads ({HQ}) must be a multiple of the number of key heads ({H}).") + NG = HQ // H + if NG > 1: + k, v, s, g = (x.repeat_interleave(NG, dim=2) for x in (k, v, s, g)) + hk0, hv0 = (x.repeat_interleave(NG, dim=1) if x is not None else None for x in (hk0, hv0)) o, *final_state = ChunkGSAFunction.apply( q, k, @@ -1215,4 +291,6 @@ def chunk_gsa( cu_seqlens_cpu, chunk_size, ) + if NG > 1 and output_final_state: + final_state = [x[:, ::NG].contiguous() for x in final_state] return o, final_state diff --git a/build/torch-cuda/ops/gsa/chunk_bwd.py b/build/torch-cuda/ops/gsa/chunk_bwd.py new file mode 100644 index 0000000000000000000000000000000000000000..b5a7335d1029f4755a8b919a70615ee89fdf7393 --- /dev/null +++ b/build/torch-cuda/ops/gsa/chunk_bwd.py @@ -0,0 +1,580 @@ +# 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 +import triton.language as tl + +from ...ops.common.chunk_h import chunk_bwd_dh, chunk_fwd_h +from ...ops.gla.chunk import chunk_gla_bwd +from ...ops.utils import chunk_local_cumsum, prepare_chunk_indices, softmax_bwd +from ...ops.utils.cache import fla_cache_autotune +from ...ops.utils.op import exp2, unflatten_program_id +from ...utils import autotune_cache_kwargs + + +@triton.heuristics({ + 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None, +}) +@fla_cache_autotune( + configs=[ + triton.Config({}, num_warps=num_warps) + for num_warps in [2, 4, 8] + ], + key=['H', 'V', 'BT', 'IS_VARLEN'], + **autotune_cache_kwargs, +) +@triton.jit(do_not_specialize=['T']) +def chunk_gsa_bwd_kernel_dA( + v, + g, + do, + dA, + chunk_indices, + cu_seqlens, + scale, + T, + B: tl.constexpr, + H: tl.constexpr, + V: tl.constexpr, + BT: tl.constexpr, + BC: tl.constexpr, + BV: tl.constexpr, + NC: tl.constexpr, + IS_VARLEN: tl.constexpr, +): + NV = tl.cdiv(V, BV) + i_v, i_c, i_bh = tl.program_id(0).to(tl.int64) % NV, tl.program_id(0).to(tl.int64) // NV, tl.program_id(1).to(tl.int64) + i_b, i_h = i_bh // H, i_bh % H + i_t, i_i, i_j = (i_c // (NC * NC)).to(tl.int64), (i_c % (NC * NC)) // NC, (i_c % (NC * NC)) % NC + if IS_VARLEN: + i_n, i_t = tl.load(chunk_indices + i_t * 2).to(tl.int64), tl.load(chunk_indices + i_t * 2 + 1).to(tl.int64) + bos, eos = tl.load(cu_seqlens + i_n).to(tl.int64), tl.load(cu_seqlens + i_n + 1).to(tl.int64) + all = T + T = eos - bos + else: + bos, eos = i_b * T, i_b * T + T + all = B * T + + o_v = i_v * BV + tl.arange(0, BV).to(tl.int64) + m_v = o_v < V + + if i_t * BT + i_i * BC >= T: + return + + o_c = i_t * BT + i_i * BC + tl.arange(0, BC) + o_jA = i_j * BC + tl.arange(0, BC) + m_c = o_c < T + m_A = m_c[:, None] & (o_jA[None, :] < BT) + m_cv = m_c[:, None] & m_v[None, :] + p_dA = dA+((i_v*all+bos)*H+i_h)*BT + o_c[:, None] * (H*BT) + o_jA[None, :] + + # [BC, BC] + b_dA = tl.zeros([BC, BC], dtype=tl.float32) + if i_i > i_j: + o_j = i_t * BT + i_j * BC + tl.arange(0, BC) + m_vj = m_v[:, None] & (o_j[None, :] < T) + p_v = v + (bos*H+i_h) * V + o_v[:, None] + o_j[None, :] * (H*V) + p_gv = g + (bos*H+i_h) * V + o_v[:, None] + o_j[None, :] * (H*V) + p_gn = g + (bos + i_t*BT + i_i*BC) * H*V + i_h * V + o_v + p_g = g + (bos*H+i_h) * V + o_c[:, None] * (H*V) + o_v[None, :] + p_do = do + (bos*H+i_h) * V + o_c[:, None] * (H*V) + o_v[None, :] + # [BV,] + b_gn = tl.load(p_gn, mask=m_v, other=0.) + # [BC, BV] + b_g = tl.load(p_g, mask=m_cv, other=0.0) + b_do = tl.load(p_do, mask=m_cv, other=0.0) + b_do = (b_do * exp2(b_g - b_gn[None, :])).to(b_do.dtype) + # [BV, BC] + b_v = tl.load(p_v, mask=m_vj, other=0.0) + b_gv = tl.load(p_gv, mask=m_vj, other=0.0) + b_vg = (b_v * exp2(b_gn[:, None] - b_gv)).to(b_v.dtype) + # [BC, BC] + b_dA = tl.dot(b_do, b_vg) * scale + elif i_i == i_j: + p_g = g + (bos*H + i_h) * V + o_c[:, None] * (H*V) + o_v[None, :] + p_do = do + (bos*H + i_h) * V + o_c[:, None] * (H*V) + o_v[None, :] + p_v = v + (bos + i_t*BT + i_j*BC) * H*V + i_h * V + o_v + p_gv = g + (bos + i_t*BT + i_j*BC) * H*V + i_h * V + o_v + # [BC, BV] + b_g = tl.load(p_g, mask=m_cv, other=0.0) + b_do = tl.load(p_do, mask=m_cv, other=0.0) * scale + m_v = o_v < V + + o_i = tl.arange(0, BC) + # [BC, BC] + m_dA = o_i[:, None] >= o_i[None, :] + for j in range(0, min(BC, T - i_t * BT - i_j * BC)): + # [BV,] + b_v = tl.load(p_v, mask=m_v, other=0).to(tl.float32) + b_gv = tl.load(p_gv, mask=m_v, other=0).to(tl.float32) + # [BC,] + b_dAj = tl.sum(b_do * b_v[None, :] * exp2(b_g - b_gv[None, :]), 1) + b_dA = tl.where((o_i == j)[None, :], b_dAj[:, None], b_dA) + + p_v += H*V + p_gv += H*V + b_dA = tl.where(m_dA, b_dA, 0.) + tl.store(p_dA, b_dA.to(dA.dtype.element_ty), mask=m_A) + + +@triton.heuristics({ + 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None, +}) +@fla_cache_autotune( + configs=[ + triton.Config({}, num_warps=num_warps, num_stages=num_stages) + for num_warps in [2, 4] + for num_stages in [2, 3, 4] + ], + key=['H', 'K', 'V', 'BT', 'IS_VARLEN'], + **autotune_cache_kwargs, +) +@triton.jit(do_not_specialize=['T']) +def chunk_gsa_bwd_kernel_dqkvg( + q, + k, + v, + h, + g, + A, + do, + dh, + dq, + dk, + dv, + dg, + dgv, + dA, + cu_seqlens, + chunk_indices, + scale, + T, + B: tl.constexpr, + H: tl.constexpr, + K: tl.constexpr, + V: tl.constexpr, + BT: tl.constexpr, + BK: tl.constexpr, + BV: tl.constexpr, + IS_VARLEN: tl.constexpr, +): + i_k, i_t = unflatten_program_id(tl.cdiv(K, BK)) + i_bh = tl.program_id(1).to(tl.int64) + i_b, i_h = i_bh // H, i_bh % H + if IS_VARLEN: + i_tg = i_t + i_n, i_t = tl.load(chunk_indices + i_t * 2).to(tl.int64), tl.load(chunk_indices + i_t * 2 + 1).to(tl.int64) + bos, eos = tl.load(cu_seqlens + i_n).to(tl.int64), tl.load(cu_seqlens + i_n + 1).to(tl.int64) + all = T + T = eos - bos + NT = tl.cdiv(T, BT) + else: + NT = tl.cdiv(T, BT) + i_tg = i_b * NT + i_t + bos, eos = i_b * T, i_b * T + T + all = B * T + + o_i = tl.arange(0, BT) + o_t = min(i_t * BT + BT, T) + m_s = o_i[:, None] >= o_i[None, :] + + o_q = i_t * BT + tl.arange(0, BT) + o_k = i_k * BK + tl.arange(0, BK).to(tl.int64) + m_q = o_q < T + m_k = o_k < K + m_qk = m_q[:, None] & m_k[None, :] + m_A = m_q[:, None] & (o_i[None, :] < BT) + p_q = q + (bos*H+i_h) * K + o_q[:, None] * (H*K) + o_k[None, :] + p_k = k + (bos*H+i_h) * K + o_q[:, None] * (H*K) + o_k[None, :] + p_A = A + ((i_k*all+bos)*H+i_h)*BT + o_q[:, None] * (H*BT) + o_i[None, :] + + # [BT, BK] + b_q = tl.load(p_q, mask=m_qk, other=0.0) + b_k = tl.load(p_k, mask=m_qk, other=0.0) + # [BT, BT] + b_A = tl.dot((b_q * scale).to(b_q.dtype), tl.trans(b_k)) + b_A = tl.where(m_s, b_A, 0.) + tl.store(p_A, b_A.to(p_A.dtype.element_ty), mask=m_A) + + b_dq = tl.zeros([BT, BK], dtype=tl.float32) + b_dk = tl.zeros([BT, BK], dtype=tl.float32) + for i_v in range(tl.cdiv(V, BV)): + o_v = i_v * BV + tl.arange(0, BV).to(tl.int64) + m_v = o_v < V + m_qv = m_q[:, None] & m_v[None, :] + m_vk = m_v[:, None] & m_k[None, :] + m_kv = m_k[:, None] & m_v[None, :] + p_v = v + (bos*H+i_h)*V + o_q[:, None] * (H*V) + o_v[None, :] + p_g = g + (bos*H+i_h)*V + o_q[:, None] * (H*V) + o_v[None, :] + p_gn = g + (bos + o_t - 1) * H*V + i_h * V + o_v + p_do = do + (bos*H+i_h)*V + o_q[:, None] * (H*V) + o_v[None, :] + p_dv = dv + ((i_k*all+bos)*H+i_h)*V + o_q[:, None] * (H*V) + o_v[None, :] + p_dg = dg + (bos*H+i_h)*V + o_q[:, None] * (H*V) + o_v[None, :] + p_dgv = dgv+((i_k*all+bos)*H+i_h)*V + o_q[:, None] * (H*V) + o_v[None, :] + p_h = h + (i_tg * H + i_h) * K*V + o_v[:, None] + o_k[None, :] * V + p_dh = dh + (i_tg * H + i_h) * K*V + o_k[:, None] * V + o_v[None, :] + + # [BV,] + b_gn = tl.load(p_gn, mask=m_v, other=0) + # [BT, BV] + b_v = tl.load(p_v, mask=m_qv, other=0.0) + b_g = tl.load(p_g, mask=m_qv, other=0.0) + b_gv = exp2(b_gn[None, :] - b_g) + # [BV, BK] + b_h = tl.load(p_h, mask=m_vk, other=0.0) + # [BT, BV] + b_do = tl.load(p_do, mask=m_qv, other=0.0) + b_do = (b_do * exp2(b_g)).to(b_do.dtype) + # [BK, BV] + b_dh = tl.load(p_dh, mask=m_kv, other=0.0) + # [BV] + b_dg = tl.sum(tl.trans(b_h) * b_dh, 0) * exp2(b_gn) + + b_dh = b_dh.to(b_k.dtype) + # [BT, BK] + b_dq += tl.dot(b_do, b_h.to(b_k.dtype)) * scale + b_dk = tl.dot((b_v * b_gv).to(b_v.dtype), tl.trans(b_dh), b_dk) + # [BT, BV] + b_dv = tl.dot(b_k, b_dh) * b_gv + # [BV] + b_dg += tl.sum(b_dv * b_v, 0) + + if i_k == 0: + b_dgv = tl.load(p_dg, mask=m_qv, other=0.0) + b_dg[None, :] + else: + b_dgv = tl.zeros([BT, BV], dtype=tl.float32) + b_dg[None, :] + + tl.store(p_dgv, b_dgv.to(p_dgv.dtype.element_ty), mask=m_qv) + tl.store(p_dv, b_dv.to(p_dv.dtype.element_ty), mask=m_qv) + p_dA = dA + (bos*H + i_h) * BT + o_q[:, None] * (H*BT) + o_i[None, :] + p_dq = dq + (bos*H + i_h) * K + o_q[:, None] * (H*K) + o_k[None, :] + p_dk = dk + (bos*H + i_h) * K + o_q[:, None] * (H*K) + o_k[None, :] + # [BT, BT] + b_dA = tl.load(p_dA, mask=m_A, other=0.0) + # [BT, BK] + b_dq = tl.dot(b_dA, b_k, b_dq) + b_dk = tl.dot(tl.trans(b_dA).to(b_k.dtype), b_q, b_dk) + + tl.store(p_dq, b_dq.to(p_dq.dtype.element_ty), mask=m_qk) + tl.store(p_dk, b_dk.to(p_dk.dtype.element_ty), mask=m_qk) + + +@triton.heuristics({ + 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None, +}) +@triton.jit(do_not_specialize=['T']) +def chunk_gsa_bwd_kernel_intra_dvg( + v, + g, + o, + A, + do, + dv, + dg, + cu_seqlens, + chunk_indices, + T, + H: tl.constexpr, + V: tl.constexpr, + BT: tl.constexpr, + BC: tl.constexpr, + BV: tl.constexpr, + NC: tl.constexpr, + IS_VARLEN: tl.constexpr, +): + NV = tl.cdiv(V, BV) + i_v, i_c, i_bh = tl.program_id(0).to(tl.int64) % NV, tl.program_id(0).to(tl.int64) // NV, tl.program_id(1).to(tl.int64) + i_b, i_h = i_bh // H, i_bh % H + i_t, i_i = (i_c // NC).to(tl.int64), i_c % NC + if IS_VARLEN: + i_n, i_t = tl.load(chunk_indices + i_t * 2).to(tl.int64), tl.load(chunk_indices + i_t * 2 + 1).to(tl.int64) + bos, eos = tl.load(cu_seqlens + i_n).to(tl.int64), tl.load(cu_seqlens + i_n + 1).to(tl.int64) + T = eos - bos + else: + bos, eos = i_b * T, i_b * T + T + + o_i = tl.arange(0, BC) + o_v = i_v * BV + tl.arange(0, BV).to(tl.int64) + m_v = o_v < V + + if i_t * BT + i_i * BC >= T: + return + + o_c = i_t * BT + i_i * BC + tl.arange(0, BC) + m_c = o_c < T + m_cv = m_c[:, None] & m_v[None, :] + p_gv = g + (bos*H+i_h)*V + o_c[:, None] * (H*V) + o_v[None, :] + p_gn = g + (bos + min(i_t * BT + i_i * BC + BC, T)-1)*H*V + i_h*V + o_v + # [BV,] + b_gn = tl.load(p_gn, mask=m_v, other=0) + # [BC, BV] + b_gv = tl.load(p_gv, mask=m_cv, other=0.0) + b_dv = tl.zeros([BC, BV], dtype=tl.float32) + for i_j in range(i_i + 1, NC): + o_j = i_t * BT + i_j * BC + tl.arange(0, BC) + o_iA = i_i * BC + tl.arange(0, BC) + m_j = (i_t * BT + i_j * BC + o_i) < T + m_jv = m_j[:, None] & m_v[None, :] + m_A = (o_iA[:, None] < BT) & m_j[None, :] + p_g = g + (bos*H+i_h) * V + o_j[:, None] * (H*V) + o_v[None, :] + p_A = A + (bos*H+i_h) * BT + o_iA[:, None] + o_j[None, :] * (H*BT) + p_do = do + (bos*H+i_h) * V + o_j[:, None] * (H*V) + o_v[None, :] + # [BC, BV] + b_g = tl.load(p_g, mask=m_jv, other=0.0) + b_do = tl.load(p_do, mask=m_jv, other=0.0) * tl.where(m_j[:, None], exp2(b_g - b_gn[None, :]), 0) + # [BC, BC] + b_A = tl.load(p_A, mask=m_A, other=0.0) + # [BC, BV] + b_dv = tl.dot(b_A, b_do.to(b_A.dtype), b_dv) + b_dv *= exp2(b_gn[None, :] - b_gv) + + p_g = g + (bos + i_t * BT + i_i * BC) * H*V + i_h * V + o_v + p_A = A + (bos + i_t*BT + i_i*BC) * H*BT + i_h * BT + i_i * BC + o_i + p_do = do + (bos + i_t*BT + i_i*BC) * H*V + i_h * V + o_v + for j in range(0, min(BC, T - i_t * BT - i_i * BC)): + # [BC,] + b_A = tl.load(p_A) + # [BV,] + b_g = tl.load(p_g, mask=m_v, other=0) + b_do = tl.load(p_do, mask=m_v, other=0) + # [BC, BV] + m_i = o_i[:, None] <= j + b_dv += tl.where(m_i, exp2(b_g[None, :] - b_gv) * b_A[:, None] * b_do[None, :], 0.) + + p_g += H * V + p_A += H * BT + p_do += H * V + p_o = o + (bos*H+i_h)*V + o_c[:, None] * (H*V) + o_v[None, :] + p_v = v + (bos*H+i_h)*V + o_c[:, None] * (H*V) + o_v[None, :] + p_do = do + (bos*H+i_h)*V + o_c[:, None] * (H*V) + o_v[None, :] + p_dv = dv + (bos*H+i_h)*V + o_c[:, None] * (H*V) + o_v[None, :] + p_dg = dg + (bos*H+i_h)*V + o_c[:, None] * (H*V) + o_v[None, :] + + b_o = tl.load(p_o, mask=m_cv, other=0.0).to(tl.float32) + b_v = tl.load(p_v, mask=m_cv, other=0.0).to(tl.float32) + b_do = tl.load(p_do, mask=m_cv, other=0.0).to(tl.float32) + b_dv = b_dv + tl.load(p_dv, mask=m_cv, other=0.0).to(tl.float32) + b_dg = b_o * b_do - b_v * b_dv + tl.store(p_dv, b_dv.to(p_dv.dtype.element_ty), mask=m_cv) + tl.store(p_dg, b_dg.to(p_dg.dtype.element_ty), mask=m_cv) + + +def chunk_gsa_bwd_k( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + g: torch.Tensor, + h: torch.Tensor, + h0: torch.Tensor, + o: torch.Tensor, + do: torch.Tensor, + dht: torch.Tensor, + dg: torch.Tensor, + scale: float = 1., + cu_seqlens: torch.LongTensor | None = None, + chunk_size: int = 64, + chunk_indices: torch.LongTensor | None = None, +): + B, T, H, K, V = *k.shape, v.shape[-1] + BT = chunk_size + BC = min(16, BT) + BK = min(64, triton.next_power_of_2(K)) + BV = min(64, triton.next_power_of_2(V)) + + if chunk_indices is None and cu_seqlens is not None: + chunk_indices = prepare_chunk_indices(cu_seqlens, BT) + NT = triton.cdiv(T, BT) if cu_seqlens is None else len(chunk_indices) + NC = triton.cdiv(BT, BC) + NK = triton.cdiv(K, BK) + NV = triton.cdiv(V, BV) + + if h is None: + h, _ = chunk_fwd_h( + k=k, + v=v, + g=None, + gk=None, + gv=g, + h0=h0, + output_final_state=False, + cu_seqlens=cu_seqlens, + chunk_size=BT, + states_in_fp32=False, + ) + dh, dh0 = chunk_bwd_dh( + q=q, + k=k, + v=v, + g=None, + gk=None, + gv=g, + do=do, + h0=h0, + dht=dht, + scale=scale, + cu_seqlens=cu_seqlens, + chunk_size=BT, + states_in_fp32=True, + ) + dA = q.new_empty(NV, B, T, H, BT) + grid = (NV * NT * NC * NC, B * H) + chunk_gsa_bwd_kernel_dA[grid]( + v=v, + g=g, + do=do, + dA=dA, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + scale=scale, + T=T, + B=B, + H=H, + V=V, + BT=BT, + BC=BC, + BV=BV, + NC=NC, + ) + dA = dA.sum(0, dtype=dA.dtype) if NV > 1 else dA.squeeze(0) + + A = do.new_empty(NK, B, T, H, BT) + dq = torch.empty_like(q) + dk = k.new_empty(B, T, H, K) + dv = v.new_empty(NK, B, T, H, V) + dgv = g.new_empty(NK, B, T, H, V, dtype=torch.float) + grid = (NK * NT, B * H) + chunk_gsa_bwd_kernel_dqkvg[grid]( + q=q, + k=k, + v=v, + h=h, + g=g, + A=A, + do=do, + dh=dh, + dq=dq, + dk=dk, + dv=dv, + dg=dg, + dgv=dgv, + dA=dA, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + scale=scale, + T=T, + B=B, + H=H, + K=K, + V=V, + BT=BT, + BK=BK, + BV=BV, + ) + A = A.sum(0, dtype=A.dtype) if NK > 1 else A.squeeze(0) + dv = dv.sum(0, dtype=dv.dtype) if NK > 1 else dv.squeeze(0) + dgv = dgv.sum(0, dtype=dgv.dtype) if NK > 1 else dgv.squeeze(0) + + def grid(meta): return (triton.cdiv(V, meta['BV']) * NT * NC, B * H) + chunk_gsa_bwd_kernel_intra_dvg[grid]( + v=v, + g=g, + o=o, + A=A, + do=do, + dv=dv, + dg=dg, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + T=T, + H=H, + V=V, + BT=BT, + BC=BC, + BV=BV, + NC=NC, + num_warps=4, + num_stages=2, + ) + dg = dgv.add_( + chunk_local_cumsum( + dg, + chunk_size=BT, + reverse=True, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + ) + ) + + return dq, dk, dv, dg, dh0 + + +def chunk_gsa_bwd( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + s: torch.Tensor, + g: torch.Tensor, + ok: torch.Tensor, + p: torch.Tensor, + A: torch.Tensor, + h: tuple[torch.Tensor, torch.Tensor], + initial_state: tuple[torch.Tensor, torch.Tensor] | None, + scale: float, + do: torch.Tensor, + dht: tuple[torch.Tensor, torch.Tensor], + cu_seqlens: torch.LongTensor | None = None, + chunk_size: int = 64, + chunk_indices: torch.LongTensor | None = None, +): + hk0, hv0 = None, None + if initial_state is not None: + hk0, hv0 = initial_state + + hk, hv = h + dhkt, dhvt = dht + + qv = p.to(q.dtype) + dqv, dsv, dv, dg, dhv0 = chunk_gla_bwd( + q=qv, + k=s, + v=v, + g=None, + g_cumsum=g, + scale=1., + initial_state=hv0, + h=hv, + A=A, + do=do, + dht=dhvt, + cu_seqlens=cu_seqlens, + chunk_size=chunk_size, + chunk_indices=chunk_indices, + ) + + dok = softmax_bwd(p, dqv, dtype=ok.dtype) + + dq, dk, dsk, dg, dhk0 = chunk_gsa_bwd_k( + q=q, + k=k, + v=s, + g=g, + h0=hk0, + h=hk, + o=ok, + do=dok, + dht=dhkt, + dg=dg, + scale=scale, + cu_seqlens=cu_seqlens, + chunk_size=chunk_size, + chunk_indices=chunk_indices, + ) + + ds = dsv.add_(dsk) + dg = dg.to(s.dtype) + return dq, dk, dv, ds, dg, dhk0, dhv0 diff --git a/build/torch-cuda/ops/gsa/chunk_fwd.py b/build/torch-cuda/ops/gsa/chunk_fwd.py new file mode 100644 index 0000000000000000000000000000000000000000..a7bb31ac4246d48708bd244dbce7cec552113b91 --- /dev/null +++ b/build/torch-cuda/ops/gsa/chunk_fwd.py @@ -0,0 +1,320 @@ +# 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 +import triton.language as tl + +from ...ops.common.chunk_h import chunk_fwd_h +from ...ops.gla.chunk import chunk_gla_fwd +from ...ops.utils import prepare_chunk_indices, softmax_fwd +from ...ops.utils.cache import fla_cache_autotune +from ...ops.utils.op import exp2, unflatten_program_id +from ...utils import autotune_cache_kwargs + + +@triton.heuristics({ + 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None, +}) +@fla_cache_autotune( + configs=[ + triton.Config({'BK': BK, 'BV': BV}, num_warps=num_warps, num_stages=num_stages) + for BK in [32, 64] + for BV in [32, 64] + for num_warps in [2, 4, 8] + for num_stages in [2, 3, 4] + ], + key=['H', 'K', 'V', 'BT', 'IS_VARLEN'], + **autotune_cache_kwargs, +) +@triton.jit(do_not_specialize=['T']) +def chunk_gsa_fwd_kernel_inter( + q, + k, + h, + g, + o, + A, + cu_seqlens, + chunk_indices, + scale, + T, + H: tl.constexpr, + K: tl.constexpr, + V: tl.constexpr, + BT: tl.constexpr, + BK: tl.constexpr, + BV: tl.constexpr, + IS_VARLEN: tl.constexpr, +): + i_v, i_t = unflatten_program_id(tl.cdiv(V, BV)) + i_bh = tl.program_id(1).to(tl.int64) + i_b, i_h = i_bh // H, i_bh % H + if IS_VARLEN: + i_tg = i_t + i_n, i_t = tl.load(chunk_indices + i_t * 2).to(tl.int64), tl.load(chunk_indices + i_t * 2 + 1).to(tl.int64) + bos, eos = tl.load(cu_seqlens + i_n).to(tl.int64), tl.load(cu_seqlens + i_n + 1).to(tl.int64) + T = eos - bos + NT = tl.cdiv(T, BT) + else: + NT = tl.cdiv(T, BT) + i_tg = i_b * NT + i_t + bos, eos = i_b * T, i_b * T + T + + o_i = tl.arange(0, BT) + m_s = o_i[:, None] >= o_i[None, :] + + b_o = tl.zeros([BT, BV], dtype=tl.float32) + b_A = tl.zeros([BT, BT], dtype=tl.float32) + o_t = i_t * BT + tl.arange(0, BT) + o_v = i_v * BV + tl.arange(0, BV).to(tl.int64) + m_t = o_t < T + m_v = o_v < V + m_tv = m_t[:, None] & m_v[None, :] + for i_k in range(tl.cdiv(K, BK)): + o_k = i_k * BK + tl.arange(0, BK).to(tl.int64) + m_k = o_k < K + m_qk = m_t[:, None] & m_k[None, :] + m_kt = m_k[:, None] & m_t[None, :] + m_kv = m_k[:, None] & m_v[None, :] + p_q = q + (bos * H + i_h) * K + o_t[:, None] * (H*K) + o_k[None, :] + p_k = k + (bos * H + i_h) * K + o_k[:, None] + o_t[None, :] * (H*K) + p_h = h + (i_tg * H + i_h) * K*V + o_k[:, None] * V + o_v[None, :] + + # [BT, BK] + b_q = tl.load(p_q, mask=m_qk, other=0.0) + b_q = (b_q * scale).to(b_q.dtype) + # [BK, BT] + b_k = tl.load(p_k, mask=m_kt, other=0.0) + # [BK, BV] + b_h = tl.load(p_h, mask=m_kv, other=0.0) + # [BT, BV] + b_o = tl.dot(b_q, b_h, b_o) + # [BT, BT] + b_A = tl.dot(b_q, b_k, b_A) + m_A = m_t[:, None] & (o_i[None, :] < BT) + p_g = g + (bos * H + i_h) * V + o_t[:, None] * (H*V) + o_v[None, :] + p_o = o + (bos * H + i_h) * V + o_t[:, None] * (H*V) + o_v[None, :] + p_A = A + (bos * H + i_h) * BT + o_t[:, None] * (H*BT) + o_i[None, :] + # [BT, BV] + b_g = tl.load(p_g, mask=m_tv, other=0.0) + b_o = b_o * exp2(b_g) + tl.store(p_o, b_o.to(p_o.dtype.element_ty), mask=m_tv) + + # [BT, BT] + b_A = tl.where(m_s, b_A, 0.) + if i_v == 0: + tl.store(p_A, b_A.to(p_A.dtype.element_ty), mask=m_A) + + +@triton.heuristics({ + 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None, +}) +@triton.jit(do_not_specialize=['T']) +def chunk_gsa_fwd_kernel_intra( + v, + g, + o, + A, + cu_seqlens, + chunk_indices, + T, + H: tl.constexpr, + V: tl.constexpr, + BT: tl.constexpr, + BC: tl.constexpr, + BV: tl.constexpr, + NC: tl.constexpr, + IS_VARLEN: tl.constexpr, +): + NV = tl.cdiv(V, BV) + i_v, i_c, i_bh = tl.program_id(0).to(tl.int64) % NV, tl.program_id(0).to(tl.int64) // NV, tl.program_id(1).to(tl.int64) + i_b, i_h = i_bh // H, i_bh % H + i_t, i_i = (i_c // NC).to(tl.int64), i_c % NC + if IS_VARLEN: + i_n, i_t = tl.load(chunk_indices + i_t * 2).to(tl.int64), tl.load(chunk_indices + i_t * 2 + 1).to(tl.int64) + bos, eos = tl.load(cu_seqlens + i_n).to(tl.int64), tl.load(cu_seqlens + i_n + 1).to(tl.int64) + T = eos - bos + else: + bos, eos = i_b * T, i_b * T + T + + o_i = tl.arange(0, BC) + o_v = i_v * BV + tl.arange(0, BV).to(tl.int64) + m_v = o_v < V + + if i_t * BT + i_i * BC >= T: + return + + o_c = i_t * BT + i_i * BC + tl.arange(0, BC) + m_c = o_c < T + m_cv = m_c[:, None] & m_v[None, :] + p_g = g + (bos * H + i_h) * V + o_c[:, None] * (H*V) + o_v[None, :] + p_gn = g + (bos + min(i_t * BT + i_i * BC, T)) * H*V + i_h * V + o_v + # [BV,] + b_gn = tl.load(p_gn, mask=m_v, other=0) + # [BC, BV] + b_o = tl.zeros([BC, BV], dtype=tl.float32) + for i_j in range(0, i_i): + o_j = i_t * BT + i_j * BC + tl.arange(0, BC) + o_jA = i_j * BC + tl.arange(0, BC) + m_jv = (o_j[:, None] < T) & m_v[None, :] + m_A = m_c[:, None] & (o_jA[None, :] < BT) + p_A = A + (bos*H+i_h) * BT + o_c[:, None] * (H*BT) + o_jA[None, :] + p_v = v + (bos*H+i_h) * V + o_j[:, None] * (H*V) + o_v[None, :] + p_gv = g + (bos*H+i_h) * V + o_j[:, None] * (H*V) + o_v[None, :] + # [BC, BV] + b_v = tl.load(p_v, mask=m_jv, other=0.0) + b_gv = tl.load(p_gv, mask=m_jv, other=0.0) + b_vg = (b_v * exp2(b_gn[None, :] - b_gv)).to(b_v.dtype) + # [BC, BC] + b_A = tl.load(p_A, mask=m_A, other=0.0) + b_o = tl.dot(b_A, b_vg, b_o) + # [BC, BV] + b_g = tl.load(p_g, mask=m_cv, other=0.0) + b_o *= exp2(b_g - b_gn[None, :]) + + o_A = (bos + i_t * BT + i_i * BC + tl.arange(0, BC)) * H*BT + i_h * BT + i_i * BC + m_A = (i_t * BT + i_i * BC + tl.arange(0, BC)) < T + for j in range(0, min(BC, T - i_t * BT - i_i * BC)): + p_v = v + (bos + i_t * BT + i_i * BC + j) * H*V + i_h * V + o_v + p_gv = g + (bos + i_t * BT + i_i * BC + j) * H*V + i_h * V + o_v + # [BC,] + b_A = tl.load(A + o_A + j, mask=m_A, other=0) + # [BV,] + b_v = tl.load(p_v, mask=m_v, other=0).to(tl.float32) + b_gv = tl.load(p_gv, mask=m_v, other=0).to(tl.float32) + # [BC, BV] + b_vg = b_v[None, :] * exp2(b_g - b_gv[None, :]) + # mask noncausal products before accumulating to avoid 0 * inf = NaN + b_o += tl.where(o_i[:, None] >= j, b_A[:, None] * b_vg, 0.) + p_o = o + (bos*H + i_h) * V + o_c[:, None] * (H*V) + o_v[None, :] + b_o += tl.load(p_o, mask=m_cv, other=0.0) + tl.store(p_o, b_o.to(p_o.dtype.element_ty), mask=m_cv) + + +def chunk_gsa_fwd_k( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + g: torch.Tensor, + h0: torch.Tensor | None = None, + output_final_state: bool = False, + scale: float = 1., + cu_seqlens: torch.LongTensor | None = None, + chunk_size: int = 64, + chunk_indices: torch.LongTensor | None = None, +) -> tuple[torch.Tensor, torch.Tensor | None, torch.Tensor]: + B, T, H, K, V = *k.shape, v.shape[-1] + BT = chunk_size + BC = min(16, BT) + BV = min(64, triton.next_power_of_2(V)) + + if chunk_indices is None and cu_seqlens is not None: + chunk_indices = prepare_chunk_indices(cu_seqlens, BT) + NT = triton.cdiv(T, BT) if cu_seqlens is None else len(chunk_indices) + NC = triton.cdiv(BT, BC) + + h, ht = chunk_fwd_h( + k=k, + v=v, + g=None, + gk=None, + gv=g, + h0=h0, + output_final_state=output_final_state, + cu_seqlens=cu_seqlens, + chunk_size=BT, + states_in_fp32=False, + ) + o = v.new_empty(B, T, H, V) + A = q.new_empty(B, T, H, BT) + def grid(meta): return (triton.cdiv(V, meta['BV']) * NT, B * H) + chunk_gsa_fwd_kernel_inter[grid]( + q=q, + k=k, + h=h, + g=g, + o=o, + A=A, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + scale=scale, + T=T, + H=H, + K=K, + V=V, + BT=BT, + ) + + def grid(meta): return (triton.cdiv(V, meta['BV']) * NT * NC, B * H) + chunk_gsa_fwd_kernel_intra[grid]( + v=v, + g=g, + o=o, + A=A, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + T=T, + H=H, + V=V, + BT=BT, + BC=BC, + BV=BV, + NC=NC, + num_warps=4, + num_stages=2, + ) + return h, ht, o + + +def chunk_gsa_fwd( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + s: torch.Tensor, + g: torch.Tensor, + initial_state: tuple[torch.Tensor, torch.Tensor] | None = None, + output_final_state: bool = False, + scale: float = 1., + cu_seqlens: torch.LongTensor | None = None, + chunk_size: int = 64, + chunk_indices: torch.LongTensor | None = None, +): + hk0, hv0 = None, None + if initial_state is not None: + hk0, hv0 = initial_state + hk, hkt, ok = chunk_gsa_fwd_k( + q=q, + k=k, + v=s, + g=g, + h0=hk0, + output_final_state=output_final_state, + scale=scale, + cu_seqlens=cu_seqlens, + chunk_size=chunk_size, + chunk_indices=chunk_indices, + ) + + # p is kept in fp32 for safe softmax backward + p = softmax_fwd(ok, dtype=torch.float) + + qv = p.to(q.dtype) + _, Av, hv, hvt, ov = chunk_gla_fwd( + q=qv, + k=s, + v=v, + g=None, + g_cumsum=g, + scale=1., + initial_state=hv0, + output_final_state=output_final_state, + cu_seqlens=cu_seqlens, + chunk_size=chunk_size, + chunk_indices=chunk_indices, + ) + return hk, hkt, ok, p, Av, hv, hvt, ov diff --git a/build/torch-cuda/ops/gsa/fused_recurrent.py b/build/torch-cuda/ops/gsa/fused_recurrent.py index 686b1836c9224cc6f97784cfe658d2de812c95cf..28ca6e9765e4bd7f8346cd5019ad8622b0978cac 100644 --- a/build/torch-cuda/ops/gsa/fused_recurrent.py +++ b/build/torch-cuda/ops/gsa/fused_recurrent.py @@ -9,11 +9,16 @@ import torch import triton import triton.language as tl -from ...ops.common.fused_recurrent import fused_recurrent_bwd_kernel, fused_recurrent_fwd_kernel +from ...ops.common.fused_recurrent import fused_recurrent_bwd, fused_recurrent_fwd +from ...ops.gsa.gate import fused_additive_gate +from ...ops.utils import softmax_bwd from ...ops.utils.op import exp from ...utils import autocast_custom_bwd, autocast_custom_fwd, input_guard +@triton.heuristics({ + 'STORE_FINAL_STATE': lambda args: args['hkt'] is not None, +}) @triton.jit def fused_recurrent_gsa_inference_kernel( q, @@ -32,25 +37,29 @@ def fused_recurrent_gsa_inference_kernel( M: tl.constexpr, BK: tl.constexpr, BV: tl.constexpr, + BM: tl.constexpr, NG: tl.constexpr, + STORE_FINAL_STATE: tl.constexpr, ): - i_bh = tl.program_id(0) + i_bh = tl.program_id(0).to(tl.int64) i_bg = i_bh // NG - b_s = tl.load(s + i_bg * M + tl.arange(0, M)).to(tl.float32) - b_g = tl.load(g + i_bg * M + tl.arange(0, M)).to(tl.float32) + o_m = tl.arange(0, BM).to(tl.int64) + m_m = o_m < M + b_s = tl.load(s + i_bg * M + o_m, mask=m_m, other=0).to(tl.float32) + b_g = tl.load(g + i_bg * M + o_m, mask=m_m, other=0).to(tl.float32) b_g = exp(b_g) - b_ok = tl.zeros([M], dtype=tl.float32) + b_ok = tl.zeros([BM], dtype=tl.float32) for i_k in range(tl.cdiv(K, BK)): - o_k = i_k * BK + tl.arange(0, BK) + o_k = i_k * BK + tl.arange(0, BK).to(tl.int64) - p_hk0 = hk0 + i_bg * K * M + (o_k[None, :]) * M + tl.arange(0, M)[:, None] + p_hk0 = hk0 + i_bg * K * M + (o_k[None, :]) * M + o_m[:, None] # [BK,] mask_k = o_k < K - # [M, BK] - mask_hk = (tl.arange(0, M) < M)[:, None] & mask_k[None, :] - # [M, BK] + # [BM, BK] + mask_hk = m_m[:, None] & mask_k[None, :] + # [BM, BK] b_hk = tl.load(p_hk0, mask=mask_hk, other=0.).to(tl.float32) # [BK,] b_q = tl.load(q + i_bh * K + o_k, mask=mask_k, other=0.).to(tl.float32) * scale @@ -58,20 +67,21 @@ def fused_recurrent_gsa_inference_kernel( b_hk = b_hk * b_g[:, None] + b_k[None, :] * b_s[:, None] b_ok += tl.sum(b_hk * b_q[None, :], axis=1) - if i_bh % NG == 0: - p_hkt = hkt + i_bg * K * M + o_k[None, :] * M + tl.arange(0, M)[:, None] - tl.store(p_hkt, b_hk.to(p_hkt.dtype.element_ty), mask=mask_hk) + if STORE_FINAL_STATE: + if i_bh % NG == 0: + p_hkt = hkt + i_bg * K * M + o_k[None, :] * M + o_m[:, None] + tl.store(p_hkt, b_hk.to(p_hkt.dtype.element_ty), mask=mask_hk) - b_qv = tl.softmax(b_ok) + b_qv = tl.softmax(tl.where(m_m, b_ok, float('-inf'))) for i_v in range(tl.cdiv(V, BV)): - o_v = i_v * BV + tl.arange(0, BV) + o_v = i_v * BV + tl.arange(0, BV).to(tl.int64) - p_hv0 = hv0 + i_bg * M * V + tl.arange(0, M)[None, :] * V + o_v[:, None] + p_hv0 = hv0 + i_bg * M * V + o_m[None, :] * V + o_v[:, None] # [BV,] mask_v = o_v < V - # [BV, M] - mask_hv = mask_v[:, None] & (tl.arange(0, M) < M)[None, :] - # [BV, M] + # [BV, BM] + mask_hv = mask_v[:, None] & m_m[None, :] + # [BV, BM] b_hv = tl.load(p_hv0, mask=mask_hv, other=0).to(tl.float32) # [BV,] b_v = tl.load(v + i_bg * V + o_v, mask=mask_v, other=0).to(tl.float32) @@ -80,9 +90,10 @@ def fused_recurrent_gsa_inference_kernel( tl.store(o + i_bh * V + o_v, b_ov.to(o.dtype.element_ty), mask=mask_v) - if i_bh % NG == 0: - p_hvt = hvt + i_bg * M * V + tl.arange(0, M)[None, :] * V + o_v[:, None] - tl.store(p_hvt, b_hv.to(p_hvt.dtype.element_ty), mask=mask_hv) + if STORE_FINAL_STATE: + if i_bh % NG == 0: + p_hvt = hvt + i_bg * M * V + o_m[None, :] * V + o_v[:, None] + tl.store(p_hvt, b_hv.to(p_hvt.dtype.element_ty), mask=mask_hv) def fused_recurrent_gsa_inference( @@ -100,37 +111,35 @@ def fused_recurrent_gsa_inference( BK, BV = min(triton.next_power_of_2(K), 64), min(triton.next_power_of_2(V), 64) NG = HQ // H - if initial_state != (None, None) and initial_state is not None: + if initial_state is not None and initial_state[0] is not None: hk0, hv0 = initial_state else: hk0, hv0 = q.new_zeros(B, H, K, M, dtype=torch.float), q.new_zeros(B, H, M, V, dtype=torch.float) hkt, hvt = None, None if output_final_state: - if NG == 1: - hkt, hvt = hk0, hv0 - else: - hkt, hvt = q.new_empty(B, H, K, M, dtype=torch.float), q.new_empty(B, H, M, V, dtype=torch.float) + hkt, hvt = q.new_empty(B, H, K, M, dtype=torch.float), q.new_empty(B, H, M, V, dtype=torch.float) o = v.new_empty(B, T, HQ, V) grid = (B * HQ,) fused_recurrent_gsa_inference_kernel[grid]( - q, - k, - v, - s, - g, - o, - hk0, - hv0, - hkt, - hvt, + q=q, + k=k, + v=v, + s=s, + g=g, + o=o, + hk0=hk0, + hv0=hv0, + hkt=hkt, + hvt=hvt, scale=scale, K=K, V=V, M=M, BK=BK, BV=BV, + BM=triton.next_power_of_2(M), NG=NG, ) return o, (hkt, hvt) @@ -148,84 +157,32 @@ def fused_recurrent_gsa_fwd( reverse: bool = False, cu_seqlens: torch.LongTensor | None = None, ) -> tuple[torch.Tensor, tuple[torch.Tensor]]: - B, T, H, K, V, M = *k.shape, v.shape[-1], s.shape[-1] - N = B if cu_seqlens is None else len(cu_seqlens) - 1 - HQ = q.shape[2] - if HQ != H: - raise ValueError("GQA not supported yet.") - - BK, BV, BM = min(triton.next_power_of_2(K), 64), min(triton.next_power_of_2(V), 64), min(triton.next_power_of_2(M), 64) - NK, NV, NM = triton.cdiv(K, BK), triton.cdiv(V, BV), triton.cdiv(M, BM) - - hk0, hv0 = None, None - if initial_state != (None, None) and initial_state is not None: - hk0, hv0 = initial_state - hkt, hvt = None, None - if output_final_state: - hkt, hvt = q.new_empty(N, H, K, M, dtype=torch.float), q.new_empty(N, H, M, V, dtype=torch.float) - - ok = q.new_empty(NK, *s.shape, dtype=torch.float) - gk, gv = None, g - grid = (NM, NK, N * H) - fused_recurrent_fwd_kernel[grid]( + if q.shape[2] != k.shape[2]: + raise ValueError("GQA is only supported for single-token inference.") + hk0, hv0 = (None, None) if initial_state is None else initial_state + ok, hkt = fused_recurrent_fwd( q=q, k=k, v=s, - g=None, - g_gamma=None, - gk=gk, - gv=gv, - o=ok, - h0=hk0, - ht=hkt, - cu_seqlens=cu_seqlens, + gv=g, scale=scale, - B=B, - T=T, - H=H, - K=K, - V=M, - BK=BK, - BV=BM, - USE_G=False, - USE_G_GAMMA=False, - USE_GK=False, - USE_GV=True, - REVERSE=reverse, + initial_state=hk0, + output_final_state=output_final_state, + reverse=reverse, + cu_seqlens=cu_seqlens, ) - ok = ok.sum(0) - - qv = ok.softmax(-1, dtype=torch.float) - ov = q.new_empty(NM, *v.shape, dtype=torch.float) - gk, gv = g, None - grid = (NV, NM, N * H) - fused_recurrent_fwd_kernel[grid]( + qv = ok.softmax(-1, dtype=torch.float32) + ov, hvt = fused_recurrent_fwd( q=qv, k=s, v=v, - g=None, - g_gamma=None, - gk=gk, - gv=gv, - o=ov, - h0=hv0, - ht=hvt, - cu_seqlens=cu_seqlens, + gk=g, scale=1., - B=B, - T=T, - H=H, - K=M, - V=V, - BK=BM, - BV=BV, - USE_G=False, - USE_G_GAMMA=False, - USE_GK=True, - USE_GV=False, - REVERSE=reverse, + initial_state=hv0, + output_final_state=output_final_state, + reverse=reverse, + cu_seqlens=cu_seqlens, ) - ov = ov.sum(0) return ok, hkt, qv, ov, hvt @@ -246,109 +203,33 @@ def fused_recurrent_gsa_bwd( reverse: bool = False, cu_seqlens: torch.LongTensor | None = None, ) -> tuple[torch.Tensor]: - B, T, H, K, V, M = *q.shape, v.shape[-1], s.shape[-1] - N = B if cu_seqlens is None else len(cu_seqlens) - 1 - - BK, BV, BM = min(triton.next_power_of_2(K), 64), min(triton.next_power_of_2(V), 64), min(triton.next_power_of_2(M), 64) - NK, NV, NM = triton.cdiv(K, BK), triton.cdiv(V, BV), triton.cdiv(M, BM) - - dqv = q.new_empty(NV, B, T, H, M, dtype=torch.float) - dsv = q.new_empty(NV, B, T, H, M, dtype=torch.float) - dv = q.new_empty(NM, B, T, H, V, dtype=torch.float) - dgv = q.new_empty(NV, B, T, H, M, dtype=torch.float) - dhv0 = torch.empty_like(hv0)if hv0 is not None else None - - grid = (NV, NM, N * H) - fused_recurrent_bwd_kernel[grid]( + dqv, dsv, dv, _, dgv, _, dhv0 = fused_recurrent_bwd( q=qv, k=s, v=v, - g=None, - g_gamma=None, gk=g, - gv=None, - o=None, - h0=hv0, do=do, - dq=dqv, - dk=dsv, - dv=dv, - dg=None, - dgk=dgv, - dgv=None, dht=dhvt, - dh0=dhv0, - cu_seqlens=cu_seqlens, scale=1., - B=B, - T=T, - H=H, - K=M, - V=V, - BK=BM, - BV=BV, - USE_G=False, - USE_G_GAMMA=False, - USE_GK=True, - USE_GV=False, - REVERSE=reverse, + initial_state=hv0, + reverse=reverse, + cu_seqlens=cu_seqlens, ) - dqv = dqv.sum(0) - dsv = dsv.sum(0) - dv = dv.sum(0) - dgv = dgv.sum(0) - - dok = qv * (dqv - (qv * dqv).sum(-1, True)) - dq = q.new_empty(NM, B, T, H, K, dtype=torch.float) - dk = q.new_empty(NM, B, T, H, K, dtype=torch.float) - dsk = q.new_empty(NK, B, T, H, M, dtype=torch.float) - dgk = q.new_empty(NK, B, T, H, M, dtype=torch.float) - dhk0 = torch.empty_like(hk0)if hk0 is not None else None - - grid = (NM, NK, N * H) - fused_recurrent_bwd_kernel[grid]( + dok = softmax_bwd(qv, dqv, dtype=torch.float32) + dq, dk, dsk, _, _, dgk, dhk0 = fused_recurrent_bwd( q=q, k=k, v=s, - g=None, - g_gamma=None, - gk=None, gv=g, o=ok, - h0=hk0, do=dok, - dq=dq, - dk=dk, - dv=dsk, - dg=None, - dgk=None, - dgv=dgk, dht=dhkt, - dh0=dhk0, - cu_seqlens=cu_seqlens, scale=scale, - B=B, - T=T, - H=H, - K=K, - V=M, - BK=BK, - BV=BM, - USE_G=False, - USE_G_GAMMA=False, - USE_GK=False, - USE_GV=True, - REVERSE=reverse, + initial_state=hk0, + reverse=reverse, + cu_seqlens=cu_seqlens, ) - dq = dq.sum(0) - dk = dk.sum(0) - dsk = dsk.sum(0) - dgk = dgk.sum(0) - - ds = dsk.add_(dsv) - dg = dgk.add_(dgv) - - return dq, dk, dv, ds, dg, dhk0, dhv0 + return dq, dk, dv, dsk.add_(dsv), dgk.add_(dgv), dhk0, dhv0 class FusedRecurrentGSAFunction(torch.autograd.Function): @@ -371,7 +252,8 @@ class FusedRecurrentGSAFunction(torch.autograd.Function): cu_seqlens: torch.LongTensor | None = None, ) -> tuple[torch.Tensor, tuple[torch.Tensor]]: T = q.shape[1] - if T == 1 and not q.requires_grad: + # the decode kernel assumes one token per dense sequence and does not save tensors for backward. + if T == 1 and not any(ctx.needs_input_grad) and cu_seqlens is None: o, (hkt, hvt) = fused_recurrent_gsa_inference( q=q, k=k, @@ -436,10 +318,10 @@ def fused_recurrent_gsa( v: torch.Tensor, s: torch.Tensor, g: torch.Tensor | None = None, - scale: int | None = None, - initial_state: tuple[torch.Tensor] | None = None, - output_final_state: bool | None = False, - reverse: bool | None = False, + scale: float | None = None, + initial_state: tuple[torch.Tensor, torch.Tensor] | None = None, + output_final_state: bool = False, + reverse: bool = False, cu_seqlens: torch.LongTensor | None = None, ) -> tuple[torch.Tensor, torch.Tensor]: r""" @@ -453,7 +335,8 @@ def fused_recurrent_gsa( s (torch.Tensor): slot representations of shape `[B, T, H, M]`. g (torch.Tensor): - Forget gates of shape `[B, H, T, M]` applied to keys. + Forget gates of shape `[B, T, H, M]` applied along the slot dimension. + If omitted, slots are cumulatively normalized as in ABC. scale (Optional[float]): Scale factor for the attention scores. If not provided, it will default to `1 / sqrt(K)`. Default: `None`. @@ -462,7 +345,7 @@ def fused_recurrent_gsa( For equal-length input sequences, `N` equals the batch size `B`. Default: `None`. output_final_state (Optional[bool]): - Whether to output the final state of shape `[N, H, K, V]` and `[N, H, M, V]`. + Whether to output the final states of shape `[N, H, K, M]` and `[N, H, M, V]`. Default: `False`. reverse (Optional[bool]): If `True`, process the state passing in reverse order. Default: `False`. @@ -480,7 +363,7 @@ def fused_recurrent_gsa( >>> import torch >>> import torch.nn.functional as F >>> from einops import rearrange - >>> from ...ops.gsa import fused_recurrent_gsa + >>> from fla.ops.gsa import fused_recurrent_gsa # inputs with equal lengths >>> B, T, H, K, V, M = 4, 2048, 4, 512, 512, 64 >>> q = torch.randn(B, T, H, K, device='cuda') @@ -519,6 +402,13 @@ def fused_recurrent_gsa( f"The number of initial states is expected to be equal to the number of input sequences, " f"i.e., {len(cu_seqlens) - 1} rather than {initial_state[0].shape[0]}.", ) + if g is None: + s, g, _ = fused_additive_gate( + x=s, + cu_seqlens=cu_seqlens, + output_dtype=k.dtype, + gate_dtype=torch.float32, + ) if scale is None: scale = k.shape[-1] ** -0.5 if initial_state is None: diff --git a/build/torch-cuda/ops/gsa/gate.py b/build/torch-cuda/ops/gsa/gate.py new file mode 100644 index 0000000000000000000000000000000000000000..e8782846c3cfeda80b368ce88dc12795cce6cf32 --- /dev/null +++ b/build/torch-cuda/ops/gsa/gate.py @@ -0,0 +1,319 @@ +# 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 +import triton.language as tl + +from ...backends import dispatch +from ...ops.utils.op import exp, logaddexp, unflatten_program_id +from ...utils import autocast_custom_bwd, autocast_custom_fwd, autotune_cache_kwargs, input_guard + + +@triton.heuristics({ + 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None, + 'USE_INITIAL_STATE': lambda args: args['initial_state'] is not None, + 'STORE_Z': lambda args: args['z'] is not None, +}) +@triton.autotune( + configs=[ + triton.Config({'BT': BT}, num_warps=num_warps) + for BT in [16, 32, 64] + for num_warps in [2, 4, 8] + ], + key=['S'], + **autotune_cache_kwargs, +) +@triton.jit(do_not_specialize=['T']) +def fused_additive_gate_fwd_kernel( + x, + z, + k, + g, + initial_state, + final_state, + cu_seqlens, + T, + S: tl.constexpr, + BT: tl.constexpr, + BS: tl.constexpr, + IS_VARLEN: tl.constexpr, + USE_INITIAL_STATE: tl.constexpr, + STORE_Z: tl.constexpr, +): + i_s, i_n = unflatten_program_id(tl.cdiv(S, BS)) + if IS_VARLEN: + bos = tl.load(cu_seqlens + i_n).to(tl.int64) + eos = tl.load(cu_seqlens + i_n + 1).to(tl.int64) + else: + bos, eos = i_n * T, (i_n + 1) * T + T = eos - bos + o_s = i_s * BS + tl.arange(0, BS) + o_i = tl.arange(0, BT) + m_s = o_s < S + p_x = x + bos * S + o_s + p_k = k + bos * S + o_s + p_g = g + bos * S + o_s + if STORE_Z: + p_z = z + bos * S + o_s + if USE_INITIAL_STATE: + b_zp = tl.load(initial_state + i_n * S + o_s, mask=m_s, other=float('-inf')).to(tl.float32) + else: + b_zp = tl.full((BS,), float('-inf'), tl.float32) + for i_t in range(0, T, BT): + o_t = i_t + o_i + m_x = (o_t[:, None] < T) & m_s[None, :] + b_x = tl.load(p_x + o_t[:, None] * S, mask=m_x, other=float('-inf')).to(tl.float32) + b_z = tl.associative_scan(b_x, axis=0, combine_fn=logaddexp) + b_z = logaddexp(b_z, b_zp[None, :]) + o_prev = tl.broadcast_to(tl.maximum(o_i - 1, 0)[:, None], (BT, BS)) + b_zc = tl.gather(b_z, o_prev, axis=0) + b_zc = tl.where(o_i[:, None] == 0, b_zp[None, :], b_zc) + b_k = exp(b_x - b_z) + b_g = tl.where(b_zc == float('-inf'), 0., b_zc - b_z) + tl.store(p_k + o_t[:, None] * S, b_k, mask=m_x) + tl.store(p_g + o_t[:, None] * S, b_g, mask=m_x) + if STORE_Z: + tl.store(p_z + o_t[:, None] * S, b_z, mask=m_x) + o_last = tl.minimum(i_t + BT, T) - 1 + b_zp = tl.sum(tl.where(o_t[:, None] == o_last, b_z, 0.), axis=0) + tl.store(final_state + i_n * S + o_s, b_zp, mask=m_s) + + +@triton.jit +def _reverse_recurrence(a_left, b_left, a_right, b_right): + return a_left * a_right, b_left * a_right + b_right + + +@triton.heuristics({ + 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None, + 'USE_INITIAL_STATE': lambda args: args['initial_state'] is not None, + 'HAS_DK': lambda args: args['dk'] is not None, + 'HAS_DG': lambda args: args['dg'] is not None, + 'USE_FINAL_STATE_GRADIENT': lambda args: args['dfinal_state'] is not None, +}) +@triton.jit(do_not_specialize=['T']) +def fused_additive_gate_bwd_kernel( + x, + z, + initial_state, + cu_seqlens, + dk, + dg, + dfinal_state, + dx, + dinitial_state, + T, + S: tl.constexpr, + BT: tl.constexpr, + BS: tl.constexpr, + IS_VARLEN: tl.constexpr, + USE_INITIAL_STATE: tl.constexpr, + HAS_DK: tl.constexpr, + HAS_DG: tl.constexpr, + USE_FINAL_STATE_GRADIENT: tl.constexpr, +): + i_s, i_n = unflatten_program_id(tl.cdiv(S, BS)) + if IS_VARLEN: + bos = tl.load(cu_seqlens + i_n).to(tl.int64) + eos = tl.load(cu_seqlens + i_n + 1).to(tl.int64) + else: + bos, eos = i_n * T, (i_n + 1) * T + T = eos - bos + o_s = i_s * BS + tl.arange(0, BS) + o_i = tl.arange(0, BT) + m_s = o_s < S + p_x = x + bos * S + o_s + p_z = z + bos * S + o_s + if HAS_DK: + p_dk = dk + bos * S + o_s + if HAS_DG: + p_dg = dg + bos * S + o_s + p_dx = dx + bos * S + o_s + if USE_INITIAL_STATE: + b_z0 = tl.load(initial_state + i_n * S + o_s, mask=m_s, other=float('-inf')).to(tl.float32) + else: + b_z0 = tl.full((BS,), float('-inf'), tl.float32) + b_dzp = tl.full((BS,), 0., tl.float32) + if USE_FINAL_STATE_GRADIENT: + b_dzt = tl.load(dfinal_state + i_n * S + o_s, mask=m_s, other=0.).to(tl.float32) + else: + b_dzt = tl.zeros((BS,), tl.float32) + for i_t in range(0, T, BT): + o_t = T - 1 - i_t - o_i + m_x = (o_t[:, None] >= 0) & m_s[None, :] + m_next = m_x & (o_t[:, None] + 1 < T) + b_x = tl.load(p_x + o_t[:, None] * S, mask=m_x, other=0.).to(tl.float32) + b_z = tl.load(p_z + o_t[:, None] * S, mask=m_x, other=0.) + b_zn = tl.load(p_z + (o_t[:, None] + 1) * S, mask=m_next, other=0.) + b_k = exp(b_x - b_z) + if HAS_DK: + b_dk = tl.load(p_dk + o_t[:, None] * S, mask=m_x, other=0.).to(tl.float32) + else: + b_dk = tl.zeros((BT, BS), tl.float32) + b_dzc = -b_dk * b_k + if HAS_DG: + b_dg = tl.load(p_dg + o_t[:, None] * S, mask=m_x, other=0.).to(tl.float32) + b_dgn = tl.load(p_dg + (o_t[:, None] + 1) * S, mask=m_next, other=0.).to(tl.float32) + b_dg = tl.where((o_t[:, None] > 0) | (b_z0[None, :] != float('-inf')), b_dg, 0.) + b_dzc += b_dgn - b_dg + b_dzc += tl.where(o_t[:, None] == T - 1, b_dzt[None, :], 0.) + # all exponent differences in the reverse recurrence are nonpositive. + b_r = exp(tl.where(m_next, b_z - b_zn, 0.)) + b_r, b_dzc = tl.associative_scan((b_r, b_dzc), axis=0, combine_fn=_reverse_recurrence) + b_dz = b_dzc + b_r * b_dzp[None, :] + tl.store(p_dx + o_t[:, None] * S, b_k * (b_dk + b_dz), mask=m_x) + o_last = tl.full((1, BS), tl.minimum(T - i_t, BT) - 1, tl.int32) + b_dzp = tl.gather(b_dz, o_last, axis=0).reshape((BS,)) + if USE_INITIAL_STATE: + if T > 0: + b_z1 = tl.load(p_z, mask=m_s, other=0.) + b_dz0 = exp(b_z0 - b_z1) * b_dzp + if HAS_DG: + b_dg0 = tl.load(p_dg, mask=m_s, other=0.).to(tl.float32) + b_dz0 += tl.where(b_z0 != float('-inf'), b_dg0, 0.) + else: + b_dz0 = b_dzt + tl.store(dinitial_state + i_n * S + o_s, b_dz0, mask=m_s) + + +def fused_additive_gate_fwd( + x: torch.Tensor, + initial_state: torch.Tensor | None, + cu_seqlens: torch.Tensor | None, + save_z: bool, + output_dtype: torch.dtype | None = None, + gate_dtype: torch.dtype | None = None, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor | None]: + B, T, H, K = x.shape + N = B if cu_seqlens is None else cu_seqlens.numel() - 1 + S, BS = H * K, 32 + k = torch.empty_like(x, dtype=output_dtype or x.dtype) + g = torch.empty_like(x, dtype=gate_dtype or x.dtype) + z = torch.empty_like(x, dtype=torch.float32) if save_z else None + final_state = x.new_empty((N, 1, H, K), dtype=torch.float32) + fused_additive_gate_fwd_kernel[(N * triton.cdiv(S, BS),)]( + x=x, + z=z, + k=k, + g=g, + initial_state=initial_state, + final_state=final_state, + cu_seqlens=cu_seqlens, + T=T, + S=S, + BS=BS, + ) + return k, g, final_state, z + + +def fused_additive_gate_bwd( + x: torch.Tensor, + z: torch.Tensor, + initial_state: torch.Tensor | None, + cu_seqlens: torch.Tensor | None, + dk: torch.Tensor | None, + dg: torch.Tensor | None, + dfinal_state: torch.Tensor | None, +) -> tuple[torch.Tensor, torch.Tensor | None]: + B, T, H, K = x.shape + N = B if cu_seqlens is None else cu_seqlens.numel() - 1 + S, BS = H * K, 32 + dx = torch.empty_like(x) + dinitial_state = torch.empty_like(initial_state) if initial_state is not None else None + fused_additive_gate_bwd_kernel[(N * triton.cdiv(S, BS),)]( + x=x, + z=z, + initial_state=initial_state, + cu_seqlens=cu_seqlens, + dk=dk, + dg=dg, + dfinal_state=dfinal_state, + dx=dx, + dinitial_state=dinitial_state, + T=T, + S=S, + BT=32, + BS=BS, + ) + return dx, dinitial_state + + +class FusedAdditiveGateFunction(torch.autograd.Function): + + @staticmethod + @input_guard + @autocast_custom_fwd + def forward(ctx, x, initial_state, cu_seqlens, output_dtype, gate_dtype): + k, g, final_state, z = fused_additive_gate_fwd( + x=x, + initial_state=initial_state, + cu_seqlens=cu_seqlens, + save_z=ctx.needs_input_grad[0] or ctx.needs_input_grad[1], + output_dtype=output_dtype, + gate_dtype=gate_dtype, + ) + ctx.save_for_backward(x, z, initial_state, cu_seqlens) + ctx.set_materialize_grads(False) + return k, g, final_state + + @staticmethod + @input_guard + @autocast_custom_bwd + def backward(ctx, dk, dg, dfinal_state): + x, z, initial_state, cu_seqlens = ctx.saved_tensors + dx, dinitial_state = fused_additive_gate_bwd( + x=x, + z=z, + initial_state=initial_state, + cu_seqlens=cu_seqlens, + dk=dk, + dg=dg, + dfinal_state=dfinal_state, + ) + return dx, dinitial_state, None, None, None + + +@dispatch +def fused_additive_gate( + x: torch.Tensor, + initial_state: torch.Tensor | None = None, + cu_seqlens: torch.Tensor | None = None, + output_dtype: torch.dtype | None = None, + gate_dtype: torch.dtype | None = None, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """Compute additive input and log-forget gates from cumulative exponential scores. + + Args: + x (torch.Tensor): + Logits of shape `[B, T, H, K]`. Packed inputs have `B=1`. + initial_state (torch.Tensor, Optional): + FP32 log-normalizers of shape `[N, 1, H, K]`, where `N` is the number of sequences. Default: `None`. + cu_seqlens (torch.Tensor, Optional): + Cumulative lengths of the packed sequences, including empty sequences. Default: `None`. + output_dtype (torch.dtype, Optional): + Dtype of the normalized output. Default: the input dtype. + gate_dtype (torch.dtype, Optional): + Dtype of the log-decay gates. Default: the input dtype. + + Returns: + k (torch.Tensor): + Normalized values with the shape of `x` and the requested output dtype. + g (torch.Tensor): + Log-decay gates with the shape of `x` and the requested gate dtype. An empty prefix has a zero first gate. + final_state (torch.Tensor): + FP32 log-normalizers of shape `[N, 1, H, K]`. Empty sequences preserve `initial_state`, + or return `-inf` without one. Gradients are supported for `x` and `initial_state`. + """ + if cu_seqlens is not None and x.shape[0] != 1: + raise ValueError("Packed inputs with cu_seqlens must have batch size 1.") + if initial_state is not None: + N = x.shape[0] if cu_seqlens is None else cu_seqlens.numel() - 1 + expected_shape = (N, 1, *x.shape[2:]) + if initial_state.shape != expected_shape: + raise ValueError(f"initial_state must have shape {expected_shape}, got {tuple(initial_state.shape)}.") + return FusedAdditiveGateFunction.apply(x, initial_state, cu_seqlens, output_dtype, gate_dtype) diff --git a/build/torch-cuda/ops/gsa/naive.py b/build/torch-cuda/ops/gsa/naive.py index 73173bf394349092bce3833aa0d9dce022d0f693..b548798ad7560b93f4b5b1e0fb8fa47f5dbb7a7b 100644 --- a/build/torch-cuda/ops/gsa/naive.py +++ b/build/torch-cuda/ops/gsa/naive.py @@ -15,12 +15,14 @@ def naive_recurrent_gsa( v: torch.Tensor, s: torch.Tensor, g: torch.Tensor | None = None, - scale: int | None = None, - initial_state: torch.Tensor | None = None, + scale: float | None = None, + initial_state: tuple[torch.Tensor, torch.Tensor] | None = None, output_final_state: bool | None = False, -) -> torch.Tensor: +) -> tuple[torch.Tensor, tuple[torch.Tensor, torch.Tensor] | None]: dtype = q.dtype - q, k, v, s, g = map(lambda x: x.transpose(1, 2).contiguous().float(), (q, k, v, s, g)) + q, k, v, s = map(lambda x: x.transpose(1, 2).contiguous().float(), (q, k, v, s)) + if g is not None: + g = g.transpose(1, 2).contiguous().float() NG = q.shape[1]//k.shape[1] # [batch_size, n_heads, seq_len, n_slots] diff --git a/build/torch-cuda/ops/hgrn/chunk.py b/build/torch-cuda/ops/hgrn/chunk.py index 746de2c73c397e0ad964cdf8d19807e94f6e1202..9907b2da92dc16d249a05d3fc02854ed90e21d00 100644 --- a/build/torch-cuda/ops/hgrn/chunk.py +++ b/build/torch-cuda/ops/hgrn/chunk.py @@ -27,7 +27,7 @@ import torch import triton import triton.language as tl -from ...ops.utils.op import exp +from ...ops.utils.op import exp, unflatten_program_id from ...utils import autotune_cache_kwargs, input_guard @@ -62,7 +62,8 @@ def chunk_hgrn_fwd_kernel_h( BD: tl.constexpr, USE_INITIAL_STATE: tl.constexpr, ): - i_d, i_t, i_b = tl.program_id(0), tl.program_id(1), tl.program_id(2) + i_d, i_t = unflatten_program_id(tl.cdiv(D, BD)) + i_b = tl.program_id(1).to(tl.int64) o_d = i_d * BD + tl.arange(0, BD) mask = o_d < D @@ -103,21 +104,23 @@ def chunk_hgrn_fwd_kernel_o( BT: tl.constexpr, BD: tl.constexpr, ): - i_d, i_b = tl.program_id(0), tl.program_id(1) + i_d, i_b = tl.program_id(0).to(tl.int64), tl.program_id(1).to(tl.int64) o_d = i_d * BD + tl.arange(0, BD) mask = o_d < D for i_t in range(1, tl.cdiv(T, BT)): - p_gc = tl.make_block_ptr(gc + i_b * s_b, (T, D), (s_t, s_d), (i_t * BT, i_d * BD), (BT, BD), (1, 0)) - p_o = tl.make_block_ptr(o + i_b * s_b, (T, D), (s_t, s_d), (i_t * BT, i_d * BD), (BT, BD), (1, 0)) + o_t = (i_t * BT).to(tl.int64) + tl.arange(0, BT) + m_g = (o_t[:, None] < T) & mask[None, :] + p_gc = gc + i_b * s_b + o_t[:, None] * s_t + o_d[None, :] * s_d + p_o = o + i_b * s_b + o_t[:, None] * s_t + o_d[None, :] * s_d # [BD,] b_h0 = tl.load(o + i_b * T * D + i_t * BT * D - D + o_d, mask=mask, other=0).to(tl.float32) # [BT, BD] - b_gc = tl.load(p_gc, boundary_check=(0, 1)).to(tl.float32) - b_o = tl.load(p_o, boundary_check=(0, 1)).to(tl.float32) + b_gc = tl.load(p_gc, mask=m_g, other=0.0).to(tl.float32) + b_o = tl.load(p_o, mask=m_g, other=0.0).to(tl.float32) b_o = b_o + exp(b_gc) * b_h0[None, :] - tl.store(p_o, b_o.to(p_o.dtype.element_ty), boundary_check=(0, 1)) + tl.store(p_o, b_o.to(p_o.dtype.element_ty), mask=m_g) @triton.autotune( @@ -140,11 +143,12 @@ def chunk_hgrn_bwd_kernel_h( BT: tl.constexpr, BD: tl.constexpr, ): - i_d, i_t, i_b = tl.program_id(0), tl.program_id(1), tl.program_id(2) + i_d, i_t = unflatten_program_id(tl.cdiv(D, BD)) + i_b = tl.program_id(1).to(tl.int64) o_d = i_d * BD + tl.arange(0, BD) mask = o_d < D BC = min(BT, T - i_t * BT) - NT = tl.num_programs(1) + NT = tl.cdiv(T, BT) p_g = g + (i_b * T + i_t * BT + BC - 1) * D + o_d p_gc = gc + (i_b * T + i_t * BT + BC - 1) * D + o_d @@ -190,30 +194,34 @@ def chunk_hgrn_bwd_kernel_o( BT: tl.constexpr, BD: tl.constexpr, ): - i_d, i_b = tl.program_id(0), tl.program_id(1) + i_d, i_b = tl.program_id(0).to(tl.int64), tl.program_id(1).to(tl.int64) o_d = i_d * BD + tl.arange(0, BD) mask = o_d < D for i_t in range(tl.cdiv(T, BT) - 1, -1, -1): - p_g = tl.make_block_ptr(g + i_b * s_b, (T, D), (s_t, s_d), (i_t * BT, i_d * BD), (BT, BD), (1, 0)) - p_gc = tl.make_block_ptr(gc + i_b * s_b, (T, D), (s_t, s_d), (i_t * BT, i_d * BD), (BT, BD), (1, 0)) - p_o = tl.make_block_ptr(o + i_b * s_b, (T, D), (s_t, s_d), (i_t * BT - 1, i_d * BD), (BT, BD), (1, 0)) - p_dx = tl.make_block_ptr(dx + i_b * s_b, (T, D), (s_t, s_d), (i_t * BT, i_d * BD), (BT, BD), (1, 0)) - p_dg = tl.make_block_ptr(dg + i_b * s_b, (T, D), (s_t, s_d), (i_t * BT, i_d * BD), (BT, BD), (1, 0)) + o_t = (i_t * BT).to(tl.int64) + tl.arange(0, BT) + o_to = (i_t * BT - 1).to(tl.int64) + tl.arange(0, BT) + m_g = (o_t[:, None] < T) & mask[None, :] + m_go = ((o_to >= 0) & (o_to < T))[:, None] & mask[None, :] + p_g = g + i_b * s_b + o_t[:, None] * s_t + o_d[None, :] * s_d + p_gc = gc + i_b * s_b + o_t[:, None] * s_t + o_d[None, :] * s_d + p_o = o + i_b * s_b + o_to[:, None] * s_t + o_d[None, :] * s_d + p_dx = dx + i_b * s_b + o_t[:, None] * s_t + o_d[None, :] * s_d + p_dg = dg + i_b * s_b + o_t[:, None] * s_t + o_d[None, :] * s_d # [BD,] mask_t = mask & ((i_t + 1) * BT < T) b_ht = tl.load(dx + i_b * T * D + (i_t + 1) * BT * D + o_d, mask=mask_t, other=0).to(tl.float32) # [BT, BD] - b_g = tl.load(p_g, boundary_check=(0, 1)).to(tl.float32) - b_gc = tl.load(p_gc, boundary_check=(0, 1)).to(tl.float32) - b_o = tl.load(p_o, boundary_check=(0, 1)).to(tl.float32) - b_dx = tl.load(p_dx, boundary_check=(0, 1)).to(tl.float32) + b_g = tl.load(p_g, mask=m_g, other=0.0).to(tl.float32) + b_gc = tl.load(p_gc, mask=m_g, other=0.0).to(tl.float32) + b_o = tl.load(p_o, mask=m_go, other=0.0).to(tl.float32) + b_dx = tl.load(p_dx, mask=m_g, other=0.0).to(tl.float32) b_dx = b_dx + exp(b_gc) * b_ht[None, :] b_dg = b_o * b_dx * exp(b_g) - 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)) + tl.store(p_dx, b_dx.to(p_dx.dtype.element_ty), mask=m_g) + tl.store(p_dg, b_dg.to(p_dg.dtype.element_ty), mask=m_g) class ChunkHGRNFunction(torch.autograd.Function): @@ -227,17 +235,16 @@ class ChunkHGRNFunction(torch.autograd.Function): gc = torch.empty_like(g, dtype=torch.float) o = torch.empty_like(x, dtype=torch.float) - def grid(meta): return (triton.cdiv(D, meta['BD']), triton.cdiv(T, meta['BT']), B) - chunk_hgrn_fwd_kernel_h[grid]( - x, g, gc, o, initial_state, - T=T, D=D, BT=BT, - USE_INITIAL_STATE=initial_state is not None, - ) + def grid(meta): return (triton.cdiv(D, meta['BD']) * triton.cdiv(T, meta['BT']), B) + chunk_hgrn_fwd_kernel_h[grid](x, g, gc, o, initial_state, T=T, D=D, BT=BT, USE_INITIAL_STATE=initial_state is not None) def grid(meta): return (triton.cdiv(D, meta['BD']), B) chunk_hgrn_fwd_kernel_o[grid]( gc, o, o.stride(-3), o.stride(-2), o.stride(-1), - T=T, D=D, BT=BT, BD=BD, + T=T, + D=D, + BT=BT, + BD=BD, num_warps=num_warps, ) final_state = None @@ -257,18 +264,18 @@ class ChunkHGRNFunction(torch.autograd.Function): gc = torch.empty_like(g, dtype=torch.float) dx = torch.empty_like(o, dtype=torch.float) - def grid(meta): return (triton.cdiv(D, meta['BD']), triton.cdiv(T, meta['BT']), B) - chunk_hgrn_bwd_kernel_h[grid]( - g, gc, dx, do, - T=T, D=D, BT=BT, - ) + def grid(meta): return (triton.cdiv(D, meta['BD']) * triton.cdiv(T, meta['BT']), B) + chunk_hgrn_bwd_kernel_h[grid](g, gc, dx, do, T=T, D=D, BT=BT) dg = torch.empty_like(g, dtype=torch.float) def grid(meta): return (triton.cdiv(D, meta['BD']), B) chunk_hgrn_bwd_kernel_o[grid]( g, gc, o, dx, dg, o.stride(-3), o.stride(-2), o.stride(-1), - T=T, D=D, BT=BT, BD=BD, + T=T, + D=D, + BT=BT, + BD=BD, num_warps=num_warps, ) if initial_state is not None: diff --git a/build/torch-cuda/ops/hgrn/fused_recurrent.py b/build/torch-cuda/ops/hgrn/fused_recurrent.py index 0992dc29c8a1b588e9dc135d1fbaaeb7299d208e..a97782448f7dca7f3e04f942ede94464eb7f79c2 100644 --- a/build/torch-cuda/ops/hgrn/fused_recurrent.py +++ b/build/torch-cuda/ops/hgrn/fused_recurrent.py @@ -42,7 +42,9 @@ def fused_recurrent_hgrn_fwd_kernel( STORE_FINAL_STATE: tl.constexpr, IS_VARLEN: tl.constexpr, ): - i_d, i_n = tl.program_id(0), tl.program_id(1) + pid = tl.program_id(0).to(tl.int64) + ND = tl.cdiv(D, BD) + i_d, i_n = pid % ND, (pid // ND).to(tl.int64) if IS_VARLEN: bos, eos = tl.load(cu_seqlens + i_n).to(tl.int64), tl.load(cu_seqlens + i_n + 1).to(tl.int64) T = eos - bos @@ -107,7 +109,9 @@ def fused_recurrent_hgrn_bwd_kernel( USE_FINAL_STATE_GRADIENT: tl.constexpr, IS_VARLEN: tl.constexpr, ): - i_d, i_n = tl.program_id(0), tl.program_id(1) + pid = tl.program_id(0).to(tl.int64) + ND = tl.cdiv(D, BD) + i_d, i_n = pid % ND, (pid // ND).to(tl.int64) if IS_VARLEN: bos, eos = tl.load(cu_seqlens + i_n).to(tl.int64), tl.load(cu_seqlens + i_n + 1).to(tl.int64) T = eos - bos @@ -169,7 +173,7 @@ def fused_recurrent_hgrn_fwd( o = torch.empty_like(x) final_state = x.new_empty(N, D) if output_final_state else None - def grid(meta): return (triton.cdiv(D, meta['BD']), N) + def grid(meta): return (triton.cdiv(D, meta['BD']) * N,) fused_recurrent_hgrn_fwd_kernel[grid]( x=x, g=g, @@ -197,7 +201,7 @@ def fused_recurrent_hgrn_bwd( dx = torch.empty_like(o, dtype=torch.float) dg = torch.empty_like(g, dtype=torch.float) dh0 = torch.empty_like(initial_state, dtype=torch.float) if initial_state is not None else None - def grid(meta): return (triton.cdiv(D, meta['BD']), N) + def grid(meta): return (triton.cdiv(D, meta['BD']) * N,) fused_recurrent_hgrn_bwd_kernel[grid]( g=g, o=o, @@ -288,7 +292,7 @@ def fused_recurrent_hgrn( >>> import torch >>> import torch.nn.functional as F >>> from einops import rearrange - >>> from ...ops.hgrn import fused_recurrent_hgrn + >>> from fla.ops.hgrn import fused_recurrent_hgrn # inputs with equal lengths >>> B, T, D = 4, 2048, 512 >>> x = torch.randn(B, T, D, device='cuda') diff --git a/build/torch-cuda/ops/kda/__init__.py b/build/torch-cuda/ops/kda/__init__.py index f896d2b9bc3e7b8db0ac0912224358b27e154179..053652b16a6182ece0ec05836e58b020b091911c 100644 --- a/build/torch-cuda/ops/kda/__init__.py +++ b/build/torch-cuda/ops/kda/__init__.py @@ -12,3 +12,6 @@ __all__ = [ "chunk_kda", "fused_recurrent_kda", ] + + +from ...ops.kda.backends import flash_kda # noqa: F401 diff --git a/build/torch-cuda/ops/kda/backends/__init__.py b/build/torch-cuda/ops/kda/backends/__init__.py index 2b5d14c9e3bf8784d3fbc38b135ac392060c3fb3..ebe1634319969ae1b1bb65f55c20b41d1866dbcf 100644 --- a/build/torch-cuda/ops/kda/backends/__init__.py +++ b/build/torch-cuda/ops/kda/backends/__init__.py @@ -7,13 +7,5 @@ """KDA backends.""" -from ....ops.backends import BackendRegistry, dispatch -from ....ops.kda.backends.flashkda import FlashKDABackend -from ....ops.kda.backends.tilelang import KDATileLangBackend -kda_registry = BackendRegistry("kda") -kda_registry.register(FlashKDABackend()) -kda_registry.register(KDATileLangBackend()) - - -__all__ = ['dispatch', 'kda_registry'] +__all__: list[str] = [] diff --git a/build/torch-cuda/ops/kda/backends/flashkda.py b/build/torch-cuda/ops/kda/backends/flash_kda.py similarity index 83% rename from build/torch-cuda/ops/kda/backends/flashkda.py rename to build/torch-cuda/ops/kda/backends/flash_kda.py index 16b6cfba000efddb7f44c7f9ab8d9fef65243549..528ececb9fedea0b97a4db2ecf1510bda3a67e7a 100644 --- a/build/torch-cuda/ops/kda/backends/flashkda.py +++ b/build/torch-cuda/ops/kda/backends/flash_kda.py @@ -15,24 +15,27 @@ from typing import TYPE_CHECKING import torch -from ....ops.backends import BaseBackend +from ....backends import BaseBackend, register +from ....ops.kda.chunk import chunk_kda if TYPE_CHECKING: from ....ops.cp import FLACPContext +# We do not register as we cannot rely on the backend to work for compile +# @register(chunk_kda) class FlashKDABackend(BaseBackend): """Copyright (c) 2026 Moonshot AI Fused CUTLASS forward (replaces the multi-kernel Triton path). https://github.com/MoonshotAI/FlashKDA - Enabled only under ``torch.inference_mode()``; disable with ``FLA_FLASH_KDA=0``. - The kernel fuses q/k L2 norm, beta sigmoid, and the KDA gate, so callers must pass - raw tensors and set all three ``*_in_kernel`` flags. + Enabled only when gradient tracking is disabled; disable with ``FLA_FLASH_KDA=0``. + The kernel fuses q/k L2 norm, beta sigmoid, and the KDA gate, + so callers must pass raw tensors and set all three ``*_in_kernel`` flags. """ - backend_type = "flashkda" + backend_type = "flash_kda" package_name = "flash_kda" env_var = "FLA_FLASH_KDA" default_enable = True @@ -51,14 +54,20 @@ class FlashKDABackend(BaseBackend): use_qk_l2norm_in_kernel: bool = False, use_gate_in_kernel: bool = False, use_beta_sigmoid_in_kernel: bool = False, - state_v_first: bool = False, - cu_seqlens: torch.LongTensor | None = None, - cu_seqlens_cpu: torch.LongTensor | None = None, + allow_neg_eigval: bool = False, safe_gate: bool = False, lower_bound: float | None = None, disable_recompute: bool = False, return_intermediate_states: bool = False, + state_v_first: bool = False, + cu_seqlens: torch.LongTensor | None = None, + cu_seqlens_cpu: torch.LongTensor | None = None, cp_context: FLACPContext | None = None, + use_graph: bool = False, + max_num_seqs: int | None = None, + *, + A_log: torch.Tensor | None = None, + dt_bias: torch.Tensor | None = None, **kwargs, ) -> tuple[bool, str | None]: if torch.is_grad_enabled(): @@ -77,6 +86,10 @@ class FlashKDABackend(BaseBackend): return False, "FlashKDA requires use_qk_l2norm_in_kernel=True" if not use_beta_sigmoid_in_kernel: return False, "FlashKDA requires use_beta_sigmoid_in_kernel=True" + if allow_neg_eigval: + return False, "FlashKDA requires allow_neg_eigval=False" + if A_log is None or dt_bias is None: + return False, "FlashKDA requires A_log and dt_bias tensors" if not state_v_first: return False, "FlashKDA requires state_v_first=True" if cp_context is not None: @@ -100,14 +113,18 @@ class FlashKDABackend(BaseBackend): use_qk_l2norm_in_kernel: bool = False, use_gate_in_kernel: bool = False, use_beta_sigmoid_in_kernel: bool = False, - state_v_first: bool = False, - cu_seqlens: torch.LongTensor | None = None, - cu_seqlens_cpu: torch.LongTensor | None = None, + allow_neg_eigval: bool = False, safe_gate: bool = False, lower_bound: float | None = None, disable_recompute: bool = False, return_intermediate_states: bool = False, + state_v_first: bool = False, + cu_seqlens: torch.LongTensor | None = None, + cu_seqlens_cpu: torch.LongTensor | None = None, cp_context: FLACPContext | None = None, + use_graph: bool = False, + max_num_seqs: int | None = None, + *, A_log: torch.Tensor | None = None, dt_bias: torch.Tensor | None = None, **kwargs, @@ -137,7 +154,11 @@ class FlashKDABackend(BaseBackend): cu_seqlens = cu_seqlens.to(torch.long) flash_kda.fwd( - q, k, v, g, beta, + q, + k, + v, + g, + beta, scale, out_buf, A_log=A_log, diff --git a/build/torch-cuda/ops/kda/backends/tilelang/__init__.py b/build/torch-cuda/ops/kda/backends/tilelang/__init__.py deleted file mode 100644 index 3bc44e67242d147cca00e39effd8b2667566be02..0000000000000000000000000000000000000000 --- a/build/torch-cuda/ops/kda/backends/tilelang/__init__.py +++ /dev/null @@ -1,84 +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 - -"""TileLang backend for KDA operations.""" - -from __future__ import annotations - -import torch - -from .....ops.backends import BaseBackend -from .....utils import find_spec_cached - -_TILELANG_AVAILABLE = find_spec_cached("tilelang") is not None - - -class KDATileLangBackend(BaseBackend): - - backend_type = "tilelang" - package_name = "tilelang" - env_var = "FLA_TILELANG" - - @classmethod - def is_available(cls) -> bool: - return _TILELANG_AVAILABLE - - def chunk_kda_bwd_wy_dqkg_fused_verifier( - self, - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - v_new: torch.Tensor, - g: torch.Tensor, - beta: torch.Tensor, - A: torch.Tensor, - h: torch.Tensor, - do: torch.Tensor, - dh: torch.Tensor, - dv: torch.Tensor, - scale: float | None = None, - state_v_first: bool = False, - cu_seqlens: torch.LongTensor | None = None, - chunk_size: int = 64, - chunk_indices: torch.LongTensor | None = None, - ) -> tuple[bool, str | None]: - if v.shape[2] != k.shape[2]: - return False, ( - "TileLang backend does not support GQA (v has more heads than k); " - "use repeat_interleave on k/q to match v's head count, or fall back to Triton" - ) - return True, None - - def chunk_kda_bwd_wy_dqkg_fused( - self, - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - v_new: torch.Tensor, - g: torch.Tensor, - beta: torch.Tensor, - A: torch.Tensor, - h: torch.Tensor, - do: torch.Tensor, - dh: torch.Tensor, - dv: torch.Tensor, - scale: float | None = None, - state_v_first: bool = False, - cu_seqlens: torch.LongTensor | None = None, - chunk_size: int = 64, - chunk_indices: torch.LongTensor | None = None, - ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: - from .....ops.kda.backends.tilelang.chunk_bwd_dqkg import ( - chunk_kda_bwd_wy_dqkg_fused_tilelang, - ) - return chunk_kda_bwd_wy_dqkg_fused_tilelang( - q=q, k=k, v=v, v_new=v_new, g=g, beta=beta, A=A, - h=h, do=do, dh=dh, dv=dv, - scale=scale, cu_seqlens=cu_seqlens, - chunk_size=chunk_size, chunk_indices=chunk_indices, - state_v_first=state_v_first, - ) diff --git a/build/torch-cuda/ops/kda/backends/tilelang/chunk_bwd_dqkg.py b/build/torch-cuda/ops/kda/backends/tilelang/chunk_bwd_dqkg.py deleted file mode 100644 index 19ed77ac0395afe869b42f331f504255f58927d6..0000000000000000000000000000000000000000 --- a/build/torch-cuda/ops/kda/backends/tilelang/chunk_bwd_dqkg.py +++ /dev/null @@ -1,471 +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 - -"""TileLang implementation of chunk_kda_bwd_wy_dqkg_fused. - -Ported from the Triton kernel in fla/ops/kda/chunk_bwd.py. -Key structural differences from the common chunk_bwd_dqkwg kernel: - - K-loop is INTERNAL (no NK grid dimension): grid is (NT, B*H) or (NT, H) for varlen - - Gate g is 2D: (B, T, H, K), not (B, T, H) - - Extra inputs: v_new, beta, A (attention matrix) - - Extra outputs: dv2, db, dA - - dgk is per-K accumulator (BK,), not scalar - - dA post-processing: causal mask -> dA@A -> A@(dA@A) -> negate - - exp2 is always used (no exp/exp2 option) -""" - - -import tilelang -import tilelang.language as T -import torch -import triton - -from .....ops.utils import prepare_chunk_indices -from .....utils import check_shared_mem - - -@tilelang.jit(pass_configs={ - tilelang.PassConfigKey.TL_DISABLE_TMA_LOWER: True, - tilelang.PassConfigKey.TL_ENABLE_FAST_MATH: True, -}) -def _build_kda_bwd_kernel( - B, H, K, V, BT, BK, BV, - hD1, hD2, - dtype_str, - STATE_V_FIRST, IS_VARLEN=False, - num_warps=4, -): - dtype_map = {'float16': T.float16, 'bfloat16': T.bfloat16, 'float32': T.float32} - dtype = dtype_map[dtype_str] - NK = tilelang.cdiv(K, BK) - NV = tilelang.cdiv(V, BV) - threads = num_warps * 32 - tile_hD1, tile_hD2 = (BV, BK) if STATE_V_FIRST else (BK, BV) - - # Rebind to underscore-prefixed locals for closure capture - _B, _H, _K, _V = B, H, K, V - _BT, _BK, _BV, _NK, _NV = BT, BK, BV, NK, NV - _hD1, _hD2, _thD1, _thD2 = hD1, hD2, tile_hD1, tile_hD2 - _dtype = dtype - _threads = threads - _TS = STATE_V_FIRST - _VAR = IS_VARLEN - - T_d, NT_d, total_h_d, N_seqs_d = T.dynamic("T, NT, total_h, N_seqs") - - qk_s = (_B, T_d, _H, _K) - v_s = (_B, T_d, _H, _V) - g_s = (_B, T_d, _H, _K) # 2D gate: per-K dimension - beta_s = (_B, T_d, _H) - A_s = (_B, T_d, _H, _BT) # attention matrix: (B, T, H, BT) - h_s = (total_h_d, _hD1, _hD2) - db_s = (_B, T_d, _H) - dA_s = (_B, T_d, _H, _BT) - - @T.macro - def kernel_body(q, k, v, v_new, g, beta, A, h, do, dh, dq, dk, dv, dv2, dg, db, dA, scale, - i_b, i_h, t_s, T_seq, i_t_local, h_idx): - # ========== 1. Load beta and A for this chunk ========== - s_beta = T.alloc_shared((_BT,), T.float32) - T.copy(beta[i_b, t_s:t_s + _BT, i_h], s_beta, disable_tma=True) - # Zero OOB positions for varlen safety - for _i in T.Parallel(_BT): - if (i_t_local * _BT + _i) >= T_seq: - s_beta[_i] = 0.0 - - # A in Triton is loaded transposed: b_A[bt, t] = A[i_b, t_s+t, i_h, bt] - # Load, transpose, and zero OOB rows/cols for varlen safety - s_A = T.alloc_shared((_BT, _BT), _dtype) - T.copy(A[i_b, t_s:t_s + _BT, i_h, 0:_BT], s_A) - f_A_tmp = T.alloc_fragment((_BT, _BT), _dtype) - for _i, _j in T.Parallel(_BT, _BT): - valid = ((i_t_local * _BT + _i) < T_seq) & ((i_t_local * _BT + _j) < T_seq) - f_A_tmp[_i, _j] = T.if_then_else(valid, s_A[_j, _i], T.cast(0, _dtype)) - T.copy(f_A_tmp, s_A) - - # ========== 2. Init dA, db accumulators ========== - b_dA = T.alloc_fragment((_BT, _BT), T.float32) - T.clear(b_dA) - b_db = T.alloc_fragment((_BT,), T.float32) - T.clear(b_db) - - # ========== 3. Pre-K V-loop (dA_v contribution, dv2, db_v) ========== - # These computations are K-independent: dA += dv@v^T, dv2 = A@dv*beta, db += sum((A@dv)*v) - s_dv_pre = T.alloc_shared((_BT, _BV), _dtype) - s_v_orig = T.alloc_shared((_BT, _BV), _dtype) - b_dvb = T.alloc_fragment((_BT, _BV), T.float32) - s_dv2 = T.alloc_shared((_BT, _BV), _dtype) - f_dvv = T.alloc_fragment((_BT, _BV), T.float32) - f_dvv_row = T.alloc_fragment((_BT,), T.float32) - - for i_v in T.serial(_NV): - v_off = i_v * _BV - - T.copy(dv[i_b, t_s:t_s + _BT, i_h, v_off:v_off + _BV], s_dv_pre) - T.copy(v[i_b, t_s:t_s + _BT, i_h, v_off:v_off + _BV], s_v_orig) - - # b_dA += dv @ v^T (BT, BT) - T.gemm(s_dv_pre, s_v_orig, b_dA, transpose_B=True) - - # b_dvb = A @ dv (BT, BV) - T.clear(b_dvb) - T.gemm(s_A, s_dv_pre, b_dvb) - - # Store dv2 = dvb * beta[:, None] - for _i, _j in T.Parallel(_BT, _BV): - s_dv2[_i, _j] = T.cast(b_dvb[_i, _j] * s_beta[_i], _dtype) - for _i, _j in T.Parallel(_BT, _BV): - if (i_t_local * _BT + _i) < T_seq: - dv2[i_b, t_s + _i, i_h, v_off + _j] = s_dv2[_i, _j] - - # db += sum(dvb * v, dim=1) - for _i, _j in T.Parallel(_BT, _BV): - f_dvv[_i, _j] = b_dvb[_i, _j] * T.cast(s_v_orig[_i, _j], T.float32) - T.reduce_sum(f_dvv, f_dvv_row, dim=1) - for _i in T.Parallel(_BT): - b_db[_i] = b_db[_i] + f_dvv_row[_i] - - # ========== 4. K-loop ========== - for i_k in T.serial(_NK): - k_off = i_k * _BK - - # Load k, g for this K-block (zero OOB rows for varlen cross-row safety) - s_k = T.alloc_shared((_BT, _BK), _dtype) - T.copy(k[i_b, t_s:t_s + _BT, i_h, k_off:k_off + _BK], s_k) - for _i, _j in T.Parallel(_BT, _BK): - if (i_t_local * _BT + _i) >= T_seq: - s_k[_i, _j] = T.cast(0, _dtype) - - s_g = T.alloc_shared((_BT, _BK), T.float32) - T.copy(g[i_b, t_s:t_s + _BT, i_h, k_off:k_off + _BK], s_g) - for _i, _j in T.Parallel(_BT, _BK): - if (i_t_local * _BT + _i) >= T_seq: - s_g[_i, _j] = 0.0 - - # g at last valid position: (BK,) - last_pos = T.max(0, T.min(_BT, T_seq - i_t_local * _BT) - 1) - s_gn = T.alloc_shared((_BK,), T.float32) - for _j in T.Parallel(_BK): - s_gn[_j] = s_g[last_pos, _j] - - # Per-K accumulators - b_dq = T.alloc_fragment((_BT, _BK), T.float32) - b_dk = T.alloc_fragment((_BT, _BK), T.float32) - b_dw = T.alloc_fragment((_BT, _BK), T.float32) - T.clear(b_dq) - T.clear(b_dk) - T.clear(b_dw) - - # dgk: per-K accumulator for sum(h*dh, axis=0) -- shape (BK,) - # Use shared memory (avoid 1D fragment/reducer issues) - s_dgk = T.alloc_shared((_BK,), T.float32) - for _j in T.Parallel(_BK): - s_dgk[_j] = 0.0 - - # ========== 4a. V-loop (dq, dk, dw, dgk) ========== - s_v_new = T.alloc_shared((_BT, _BV), _dtype) - s_do = T.alloc_shared((_BT, _BV), _dtype) - s_h = T.alloc_shared((_thD1, _thD2), _dtype) - s_dh = T.alloc_shared((_thD1, _thD2), _dtype) - s_dv_tile = T.alloc_shared((_BT, _BV), _dtype) - f_hdh = T.alloc_fragment((_thD1, _thD2), T.float32) - f_hdh_k_nots = T.alloc_fragment((_thD1,), T.float32) - f_hdh_t = T.alloc_fragment((_thD2, _thD1), T.float32) - f_hdh_k_ts = T.alloc_fragment((_thD2,), T.float32) - - for i_v in T.serial(_NV): - v_off = i_v * _BV - - T.copy(v_new[i_b, t_s:t_s + _BT, i_h, v_off:v_off + _BV], s_v_new) - T.copy(do[i_b, t_s:t_s + _BT, i_h, v_off:v_off + _BV], s_do) - - if _TS: - T.copy(h[h_idx, v_off:v_off + _BV, k_off:k_off + _BK], s_h) - T.copy(dh[h_idx, v_off:v_off + _BV, k_off:k_off + _BK], s_dh) - else: - T.copy(h[h_idx, k_off:k_off + _BK, v_off:v_off + _BV], s_h) - T.copy(dh[h_idx, k_off:k_off + _BK, v_off:v_off + _BV], s_dh) - - # GEMMs first (consume s_h/s_dh before any alloc might reuse their memory) - # b_dq += do @ h, b_dk += v_new @ dh - if _TS: - T.gemm(s_do, s_h, b_dq) - T.gemm(s_v_new, s_dh, b_dk) - else: - T.gemm(s_do, s_h, b_dq, transpose_B=True) - T.gemm(s_v_new, s_dh, b_dk, transpose_B=True) - - # b_dw += dv @ h (dv feeds into dw → cross-row dA/dkgb, needs OOB mask) - T.copy(dv[i_b, t_s:t_s + _BT, i_h, v_off:v_off + _BV], s_dv_tile) - - if _TS: - T.gemm(s_dv_tile, s_h, b_dw) - else: - T.gemm(s_dv_tile, s_h, b_dw, transpose_B=True) - - # dgk += sum(h * dh, axis=0) -- per-column sum -> (BK,) - for _i, _j in T.Parallel(_thD1, _thD2): - f_hdh[_i, _j] = T.cast(s_h[_i, _j], T.float32) * T.cast(s_dh[_i, _j], T.float32) - - if not _TS: - T.reduce_sum(f_hdh, f_hdh_k_nots, dim=1) - for _j in T.Parallel(_BK): - s_dgk[_j] = s_dgk[_j] + f_hdh_k_nots[_j] - - else: - for _i, _j in T.Parallel(_thD2, _thD1): - f_hdh_t[_i, _j] = f_hdh[_j, _i] - - T.reduce_sum(f_hdh_t, f_hdh_k_ts, dim=1) - for _j in T.Parallel(_BK): - s_dgk[_j] = s_dgk[_j] + f_hdh_k_ts[_j] - - # ========== 4b. Gate dq, dk ========== - # b_dq = b_dq * exp2(g) * scale - s_dq = T.alloc_shared((_BT, _BK), T.float32) - T.copy(b_dq, s_dq) - for _i, _j in T.Parallel(_BT, _BK): - s_dq[_i, _j] = s_dq[_i, _j] * T.exp2(s_g[_i, _j]) * scale - - # b_dk = b_dk * where(m_t, exp2(gn - g), 0) - s_dk = T.alloc_shared((_BT, _BK), T.float32) - T.copy(b_dk, s_dk) - for _i, _j in T.Parallel(_BT, _BK): - m_t = (i_t_local * _BT + _i) < T_seq - s_dk[_i, _j] = T.if_then_else( - m_t, - s_dk[_i, _j] * T.exp2(s_gn[_j] - s_g[_i, _j]), - 0.0) - - # ========== 4c. kg = k * exp2(g), dgk *= exp2(gn) ========== - s_kg = T.alloc_shared((_BT, _BK), _dtype) - for _i, _j in T.Parallel(_BT, _BK): - s_kg[_i, _j] = T.cast(T.cast(s_k[_i, _j], T.float32) * T.exp2(s_g[_i, _j]), _dtype) - - for _j in T.Parallel(_BK): - s_dgk[_j] = s_dgk[_j] * T.exp2(s_gn[_j]) - - # ========== 4d. dw = -dw ========== - s_dw = T.alloc_shared((_BT, _BK), _dtype) - T.copy(b_dw, s_dw) - for _i, _j in T.Parallel(_BT, _BK): - s_dw[_i, _j] = T.cast(-T.cast(s_dw[_i, _j], T.float32), _dtype) - - # ========== 4e. dA += dw @ kg^T ========== - T.gemm(s_dw, s_kg, b_dA, transpose_B=True) - - # ========== 4f. dkgb = A @ dw ========== - b_dkgb = T.alloc_fragment((_BT, _BK), T.float32) - T.clear(b_dkgb) - T.gemm(s_A, s_dw, b_dkgb) - - # ========== 4g. db += sum(dkgb * kg, dim=1) ========== - f_dkgb_kg = T.alloc_fragment((_BT, _BK), T.float32) - for _i, _j in T.Parallel(_BT, _BK): - f_dkgb_kg[_i, _j] = b_dkgb[_i, _j] * T.cast(s_kg[_i, _j], T.float32) - f_db_row = T.alloc_fragment((_BT,), T.float32) - T.reduce_sum(f_dkgb_kg, f_db_row, dim=1) - for _i in T.Parallel(_BT): - b_db[_i] = b_db[_i] + f_db_row[_i] - - # ========== 4h. dgk += col_sum(k * dk) ========== - f_kdk = T.alloc_fragment((_BT, _BK), T.float32) - for _i, _j in T.Parallel(_BT, _BK): - f_kdk[_i, _j] = T.cast(s_k[_i, _j], T.float32) * s_dk[_i, _j] - # Column sum of kdk -> (BK,) via transpose then row-reduce - f_kdk_t = T.alloc_fragment((_BK, _BT), T.float32) - for _i, _j in T.Parallel(_BK, _BT): - f_kdk_t[_i, _j] = f_kdk[_j, _i] - f_kdk_col = T.alloc_fragment((_BK,), T.float32) - T.reduce_sum(f_kdk_t, f_kdk_col, dim=1) - s_kdk_col = T.alloc_shared((_BK,), T.float32) - T.copy(f_kdk_col, s_kdk_col) - for _j in T.Parallel(_BK): - s_dgk[_j] = s_dgk[_j] + s_kdk_col[_j] - - # ========== 4i. dg = q*dq - kdk + m_last*dgk + kg*dkgb*beta ========== - s_dg_out = T.alloc_shared((_BT, _BK), T.float32) - m_last_pos = T.max(0, T.min(_BT, T_seq - i_t_local * _BT) - 1) - for _i, _j in T.Parallel(_BT, _BK): - q_dq = T.cast(q[i_b, t_s + _i, i_h, k_off + _j], T.float32) * s_dq[_i, _j] - last_term = T.if_then_else(_i == m_last_pos, s_dgk[_j], 0.0) - beta_term = T.cast(s_kg[_i, _j], T.float32) * b_dkgb[_i, _j] * s_beta[_i] - s_dg_out[_i, _j] = q_dq - f_kdk[_i, _j] + last_term + beta_term - - # ========== 4j. dk += dkgb * gb where gb = exp2(g) * beta ========== - for _i, _j in T.Parallel(_BT, _BK): - s_dk[_i, _j] = s_dk[_i, _j] + b_dkgb[_i, _j] * T.exp2(s_g[_i, _j]) * s_beta[_i] - - # ========== 4k. Store dq, dk, dg for this K-block ========== - for _i, _j in T.Parallel(_BT, _BK): - if (i_t_local * _BT + _i) < T_seq: - dq[i_b, t_s + _i, i_h, k_off + _j] = s_dq[_i, _j] - for _i, _j in T.Parallel(_BT, _BK): - if (i_t_local * _BT + _i) < T_seq: - dk[i_b, t_s + _i, i_h, k_off + _j] = s_dk[_i, _j] - for _i, _j in T.Parallel(_BT, _BK): - if (i_t_local * _BT + _i) < T_seq: - dg[i_b, t_s + _i, i_h, k_off + _j] = s_dg_out[_i, _j] - - # ========== 5. dA post-processing ========== - # m_A = (i > j) & valid - # b_dA = where(m_A, b_dA * beta[j], 0) - # b_dA = b_dA @ A - # b_dA = A @ b_dA - # b_dA = where(m_A, -b_dA, 0) - - s_dA = T.alloc_shared((_BT, _BT), T.float32) - T.copy(b_dA, s_dA) - - # Apply causal mask and beta weighting - for _i, _j in T.Parallel(_BT, _BT): - m_valid = ((i_t_local * _BT + _i) < T_seq) & ((i_t_local * _BT + _j) < T_seq) - m_upper = (_i > _j) & m_valid - s_dA[_i, _j] = T.if_then_else(m_upper, s_dA[_i, _j] * s_beta[_j], 0.0) - - # dA = dA @ A - s_dA_dtype = T.alloc_shared((_BT, _BT), _dtype) - for _i, _j in T.Parallel(_BT, _BT): - s_dA_dtype[_i, _j] = T.cast(s_dA[_i, _j], _dtype) - - b_dA2 = T.alloc_fragment((_BT, _BT), T.float32) - T.clear(b_dA2) - T.gemm(s_dA_dtype, s_A, b_dA2) - - # dA = A @ dA - # Reuse s_dA_dtype for the second cast - for _i, _j in T.Parallel(_BT, _BT): - s_dA_dtype[_i, _j] = T.cast(b_dA2[_i, _j], _dtype) - - b_dA3 = T.alloc_fragment((_BT, _BT), T.float32) - T.clear(b_dA3) - T.gemm(s_A, s_dA_dtype, b_dA3) - - # Apply mask and negate, reuse s_dA buffer - for _i, _j in T.Parallel(_BT, _BT): - m_valid = ((i_t_local * _BT + _i) < T_seq) & ((i_t_local * _BT + _j) < T_seq) - m_upper = (_i > _j) & m_valid - s_dA[_i, _j] = T.if_then_else(m_upper, -b_dA3[_i, _j], 0.0) - - for _i, _j in T.Parallel(_BT, _BT): - if (i_t_local * _BT + _i) < T_seq: - dA[i_b, t_s + _i, i_h, _j] = s_dA[_i, _j] - - # ========== 6. Store db ========== - for _i in T.Parallel(_BT): - if (i_t_local * _BT + _i) < T_seq: - db[i_b, t_s + _i, i_h] = b_db[_i] - - # ========== Kernel entry points ========== - if _VAR: - @T.prim_func - def kernel( - q: T.Tensor(qk_s, _dtype), k: T.Tensor(qk_s, _dtype), - v: T.Tensor(v_s, _dtype), v_new: T.Tensor(v_s, _dtype), - g: T.Tensor(g_s, T.float32), beta: T.Tensor(beta_s, T.float32), - A: T.Tensor(A_s, _dtype), - h: T.Tensor(h_s, _dtype), do: T.Tensor(v_s, _dtype), - dh: T.Tensor(h_s, _dtype), - dq: T.Tensor(qk_s, T.float32), dk: T.Tensor(qk_s, T.float32), - dv: T.Tensor(v_s, _dtype), dv2: T.Tensor(v_s, _dtype), - dg: T.Tensor(g_s, T.float32), db: T.Tensor(db_s, T.float32), - dA: T.Tensor(dA_s, T.float32), - cu_seqlens: T.Tensor((N_seqs_d,), T.int32), - chunk_indices: T.Tensor((NT_d, 2), T.int32), - scale: T.float32, - ): - with T.Kernel(NT_d, _H, threads=_threads) as (i_t, i_h): - i_n = chunk_indices[i_t, 0] - i_t_local = chunk_indices[i_t, 1] - bos = cu_seqlens[i_n] - T_seq = cu_seqlens[i_n + 1] - bos - h_idx = i_t * _H + i_h - t_s = bos + i_t_local * _BT - kernel_body(q, k, v, v_new, g, beta, A, h, do, dh, - dq, dk, dv, dv2, dg, db, dA, - scale, 0, i_h, t_s, T_seq, i_t_local, h_idx) - else: - @T.prim_func - def kernel( - q: T.Tensor(qk_s, _dtype), k: T.Tensor(qk_s, _dtype), - v: T.Tensor(v_s, _dtype), v_new: T.Tensor(v_s, _dtype), - g: T.Tensor(g_s, T.float32), beta: T.Tensor(beta_s, T.float32), - A: T.Tensor(A_s, _dtype), - h: T.Tensor(h_s, _dtype), do: T.Tensor(v_s, _dtype), - dh: T.Tensor(h_s, _dtype), - dq: T.Tensor(qk_s, T.float32), dk: T.Tensor(qk_s, T.float32), - dv: T.Tensor(v_s, _dtype), dv2: T.Tensor(v_s, _dtype), - dg: T.Tensor(g_s, T.float32), db: T.Tensor(db_s, T.float32), - dA: T.Tensor(dA_s, T.float32), - scale: T.float32, - ): - with T.Kernel(T.ceildiv(T_d, _BT), _B * _H, threads=_threads) as (i_t, i_bh): - i_b = i_bh // _H - i_h = i_bh % _H - NT_local = T.ceildiv(T_d, _BT) - h_idx = (i_b * NT_local + i_t) * _H + i_h - t_s = i_t * _BT - kernel_body(q, k, v, v_new, g, beta, A, h, do, dh, - dq, dk, dv, dv2, dg, db, dA, - scale, i_b, i_h, t_s, T_d, i_t, h_idx) - - return kernel - - -def chunk_kda_bwd_wy_dqkg_fused_tilelang( - q, k, v, v_new, g, beta, A, h, do, dh, dv, - scale=None, cu_seqlens=None, chunk_size=64, - chunk_indices=None, state_v_first=False, -): - B, _, H, K, V = *k.shape, v.shape[-1] - BT = chunk_size - # Ensure float32 inputs match kernel signature - g = g.float() - beta = beta.float() - if chunk_indices is None and cu_seqlens is not None: - chunk_indices = prepare_chunk_indices(cu_seqlens, BT) - IS_VARLEN = cu_seqlens is not None - - CONST_TILING = 64 if check_shared_mem() else 32 - if scale is None: - scale = K ** -0.5 - - # Outputs - dq = torch.empty_like(q, dtype=torch.float) - dk = torch.empty_like(k, dtype=torch.float) - dv2 = torch.empty_like(v) - dg = torch.empty_like(g, dtype=torch.float) - db = torch.empty_like(beta, dtype=torch.float) - dA = torch.empty_like(A, dtype=torch.float) - - h_flat = h.reshape(-1, h.shape[-2], h.shape[-1]) - dh_flat = dh.reshape(-1, dh.shape[-2], dh.shape[-1]) - hD1, hD2 = h_flat.shape[-2], h_flat.shape[-1] - dtype_str = {torch.float16: 'float16', torch.bfloat16: 'bfloat16', torch.float32: 'float32'}[q.dtype] - - # Static tiling config - BK = min(max(triton.next_power_of_2(K), 16), CONST_TILING) - BV = min(max(triton.next_power_of_2(V), 16), CONST_TILING) - num_warps = 8 if K <= 64 else 4 - - kernel = _build_kda_bwd_kernel( - B, H, K, V, BT, BK, BV, - hD1, hD2, dtype_str, - state_v_first, IS_VARLEN, - num_warps=num_warps, - ) - - if IS_VARLEN: - kernel(q, k, v, v_new, g, beta, A, h_flat, do, dh_flat, - dq, dk, dv, dv2, dg, db, dA, - cu_seqlens.int(), chunk_indices.int(), scale) - else: - kernel(q, k, v, v_new, g, beta, A, h_flat, do, dh_flat, - dq, dk, dv, dv2, dg, db, dA, scale) - - dv = dv2 - return dq, dk, dv, db, dg, dA diff --git a/build/torch-cuda/ops/kda/chunk.py b/build/torch-cuda/ops/kda/chunk.py index b11ef0a8b8becb69e828413ac7ae37c3dd7fedaf..284d12dcb4664f75afa2e75a9cdfd4ce61cecf48 100644 --- a/build/torch-cuda/ops/kda/chunk.py +++ b/build/torch-cuda/ops/kda/chunk.py @@ -5,20 +5,20 @@ # For a list of all contributors, visit: # https://github.com/fla-org/flash-linear-attention/graphs/contributors -# Related files are modified and supported by the Moonshot AI Team +# related files are modified and supported by the Moonshot AI Team import warnings import torch -from ...modules.l2norm import l2norm_bwd, l2norm_fwd -from ...ops.backends import dispatch +from ...backends import dispatch +from ...modules.norm.l2norm import l2norm_bwd, l2norm_fwd from ...ops.common.gate import fused_beta_sigmoid, fused_beta_sigmoid_bwd from ...ops.cp import FLACPContext from ...ops.kda.chunk_bwd import chunk_kda_bwd from ...ops.kda.chunk_fwd import chunk_kda_fwd -from ...ops.utils.index import prepare_chunk_indices -from ...utils import autocast_custom_bwd, autocast_custom_fwd, input_guard +from ...ops.utils.index import prepare_chunk_indices, prepare_chunk_indices_static +from ...utils import IS_NPU, autocast_custom_bwd, autocast_custom_fwd, input_guard class ChunkKDAFunction(torch.autograd.Function): @@ -50,8 +50,9 @@ class ChunkKDAFunction(torch.autograd.Function): disable_recompute: bool = False, return_intermediate_states: bool = False, cp_context: FLACPContext | None = None, + use_graph: bool = False, + max_num_seqs: int | None = None, ): - # Apply l2norm q_rstd, k_rstd = None, None if use_qk_l2norm_in_kernel: q, q_rstd = l2norm_fwd(q) @@ -59,15 +60,30 @@ class ChunkKDAFunction(torch.autograd.Function): beta_raw = beta if use_beta_sigmoid_in_kernel: - beta = fused_beta_sigmoid(beta_raw, scale=2.0 if allow_neg_eigval else 1.0) + beta = fused_beta_sigmoid(x=beta_raw, scale=2.0 if allow_neg_eigval else 1.0) - chunk_indices = None + chunk_indices, chunk_offsets = None, None if cu_seqlens is not None: - chunk_indices = prepare_chunk_indices( - cu_seqlens, - chunk_size, - cu_seqlens_cpu=cu_seqlens_cpu, - ) + if use_graph: + if max_num_seqs is None: + max_num_seqs = cu_seqlens.shape[0] - 1 + assert cu_seqlens.shape[0] - 1 == max_num_seqs, ( + f"cu_seqlens must be padded with zero-length tail sequences to exactly " + f"max_num_seqs + 1 entries, got {cu_seqlens.shape[0] - 1} sequences " + f"with max_num_seqs={max_num_seqs}" + ) + nt_max = (q.shape[1] + chunk_size - 1) // chunk_size + max_num_seqs - 1 + chunk_indices, chunk_offsets = prepare_chunk_indices_static( + cu_seqlens=cu_seqlens, + chunk_size=chunk_size, + nt_max=nt_max, + ) + else: + chunk_indices = prepare_chunk_indices( + cu_seqlens=cu_seqlens, + chunk_size=chunk_size, + cu_seqlens_cpu=cu_seqlens_cpu, + ) g_input = g @@ -80,19 +96,21 @@ class ChunkKDAFunction(torch.autograd.Function): scale=scale, initial_state=initial_state, output_final_state=output_final_state, + state_v_first=state_v_first, cu_seqlens=cu_seqlens, cu_seqlens_cpu=cu_seqlens_cpu, chunk_indices=chunk_indices, + chunk_size=chunk_size, safe_gate=safe_gate, lower_bound=lower_bound, use_gate_in_kernel=use_gate_in_kernel, A_log=A_log, dt_bias=dt_bias, - chunk_size=chunk_size, disable_recompute=disable_recompute, return_intermediate_states=return_intermediate_states, cp_context=cp_context, - state_v_first=state_v_first, + use_graph=use_graph, + chunk_offsets=chunk_offsets, ) if return_intermediate_states: @@ -101,10 +119,31 @@ class ChunkKDAFunction(torch.autograd.Function): return o.type_as(q), final_state, h ctx.save_for_backward( - q, q_rstd, k, k_rstd, v, g_cumsum, g_input, beta_raw, beta, A_log, dt_bias, Aqk, Akk, - w, u, qg, kg, v_new, h, - initial_state, cu_seqlens, chunk_indices + q, + q_rstd, + k, + k_rstd, + v, + g_cumsum, + g_input, + beta_raw, + beta, + A_log, + dt_bias, + Aqk, + Akk, + w, + u, + qg, + kg, + v_new, + h, + initial_state, + cu_seqlens, + chunk_indices, + chunk_offsets, ) + ctx.use_graph = use_graph ctx.chunk_size = chunk_size ctx.safe_gate = safe_gate ctx.scale = scale @@ -121,14 +160,10 @@ class ChunkKDAFunction(torch.autograd.Function): @staticmethod @input_guard @autocast_custom_bwd - def backward( - ctx, - do: torch.Tensor, - dht: torch.Tensor, - ): + def backward(ctx, do: torch.Tensor, dht: torch.Tensor): (q, q_rstd, k, k_rstd, v, g_cumsum, g_input, beta_raw, beta, A_log, dt_bias, Aqk, Akk, w, u, qg, kg, v_new, h, - initial_state, cu_seqlens, chunk_indices) = ( + initial_state, cu_seqlens, chunk_indices, chunk_offsets) = ( ctx.saved_tensors ) @@ -162,19 +197,22 @@ class ChunkKDAFunction(torch.autograd.Function): kg=kg, v_new=v_new, h=h, + use_graph=ctx.use_graph, + chunk_offsets=chunk_offsets, ) if ctx.use_qk_l2norm_in_kernel: dq = l2norm_bwd(q, q_rstd, dq) dk = l2norm_bwd(k, k_rstd, dk) if ctx.use_beta_sigmoid_in_kernel: - db = fused_beta_sigmoid_bwd(beta_raw, db, scale=2.0 if ctx.allow_neg_eigval else 1.0) + db = fused_beta_sigmoid_bwd(x=beta_raw, dy=db, scale=2.0 if ctx.allow_neg_eigval else 1.0) return (dq.to(q), dk.to(k), dv.to(v), dg.to(g_input), db.to(beta_raw), dA, dbias, None, dh0, - None, None, None, None, None, None, None, None, None, None, None, None, None, None) + None, None, None, None, None, None, None, None, None, None, None, None, None, None, + None, None) -@dispatch('kda') -@torch.compiler.disable +# Use triton only, other backends (cutlass, cute, etc.) need additional torch.compile support +# @dispatch def chunk_kda( q: torch.Tensor, k: torch.Tensor, @@ -195,7 +233,9 @@ def chunk_kda( state_v_first: bool = False, cu_seqlens: torch.LongTensor | None = None, cu_seqlens_cpu: torch.LongTensor | None = None, - cp_context: FLACPContext = None, + cp_context: FLACPContext | None = None, + use_graph: bool = False, + max_num_seqs: int | None = None, **kwargs, ): r""" @@ -214,67 +254,92 @@ def chunk_kda( the kernel fuses ``-exp(A_log) * softplus(g + dt_bias)`` + chunk cumsum internally. beta (torch.Tensor): betas of shape ``[B, T, HV]``. - scale (Optional[float]): + scale (float, Optional): Scale factor for the KDA attention scores. If not provided, it will default to ``1 / sqrt(K)``. Default: ``None``. - initial_state (Optional[torch.Tensor]): + initial_state (torch.Tensor, Optional): Initial state of shape ``[N, HV, K, V]`` for ``N`` input sequences. For equal-length input sequences, ``N`` equals the batch size ``B``. Default: ``None``. - output_final_state (Optional[bool]): + output_final_state (bool, Optional): Whether to output the final state of shape ``[N, HV, K, V]``. Default: ``False``. - use_qk_l2norm_in_kernel (bool): + use_qk_l2norm_in_kernel (bool, Optional): Whether to apply L2norm to the q,k tensor internally. Default: ``False``. - use_gate_in_kernel (bool): + use_gate_in_kernel (bool, Optional): Whether to compute the log-space KDA decay internally. - If ``True``: The passed ``g`` acts as the raw input for ``-exp(A_log) * softplus(g + dt_bias.view(HV, K))``. Note that as part of the input arguments, ``A_log`` (shape ``[HV]``) and the optional ``dt_bias`` (shape ``[HV * K]``) should be provided. + When ``lower_bound`` is set, ``A_log`` may be ``None``, + in which case the gate is ``lower_bound * sigmoid(g + dt_bias)``. - If ``False``, ``g`` is expected to be the pre-computed decay value. Default: ``False``. - use_beta_sigmoid_in_kernel (bool): + use_beta_sigmoid_in_kernel (bool, Optional): Whether to apply ``torch.sigmoid(beta)`` before launching the chunk kernel. - If ``True``, the passed ``beta`` acts as the raw beta logits. - If ``False``, ``beta`` is expected to already be in post-sigmoid space. Default: ``False``. - allow_neg_eigval (bool): + allow_neg_eigval (bool, Optional): Whether to allow negative eigenvalues by scaling ``beta`` to ``[0, 2)``. - Only takes effect together with ``use_beta_sigmoid_in_kernel=True``, in which case - the kernel computes ``2 * sigmoid(beta)`` instead of ``sigmoid(beta)``. + Only takes effect together with ``use_beta_sigmoid_in_kernel=True``, + in which case the kernel computes ``2 * sigmoid(beta)`` instead of ``sigmoid(beta)``. Default: ``False``. - safe_gate (bool): - Whether to clamp the gate to ``[lower_bound, 0)`` and enable M=16 TensorCore - acceleration for higher throughput. Requires ``lower_bound`` to be set. + safe_gate (bool, Optional): + Whether to enable M=16 TensorCore acceleration for bounded gates. + With ``use_gate_in_kernel=True``, requires ``lower_bound`` and bounds gates to ``[lower_bound, 0)``. + Otherwise, the caller must provide pre-computed bounded gates. Default: ``False``. - lower_bound (Optional[float]): - Lower bound for the forget gate (in log space). When set together with - ``safe_gate=True``, changes the gate activation from - ``-exp(A_log) * softplus(g + dt_bias)`` to - ``lower_bound * sigmoid(exp(A_log) * (g + dt_bias))``, - which naturally clamps the output to ``[lower_bound, 0)``. + lower_bound (float, Optional): + Lower bound for the forget gate (in log space). When set together with ``safe_gate=True``, + the gate uses ``lower_bound * sigmoid(exp(A_log) * (g + dt_bias))``, + replacing ``-exp(A_log) * softplus(g + dt_bias)`` and clamping the output to ``[lower_bound, 0)``. Recommended value: ``-5`` (i.e., per-step decay ``exp(-5) ≈ 0.0067``). Default: ``None``. - disable_recompute (bool): - Whether to disable gradient recomputation in the kernel. When ``True``, the kernel - will save all intermediate activations for backward pass, which is beneficial - for training small models at the cost of increased memory usage. Default: ``False``. - return_intermediate_states (bool): - If True, returns intermediate state ``h`` for inference scenarios (e.g., vLLM). + disable_recompute (bool, Optional): + Whether to disable gradient recomputation in the kernel. When ``True``, + the kernel will save all intermediate activations for backward pass, + which is beneficial for training small models at the cost of increased memory usage. Default: ``False``. + return_intermediate_states (bool, Optional): + If ``True``, returns intermediate state ``h`` for inference scenarios (e.g., vLLM). Must be used within ``torch.inference_mode()`` and will return a 3-tuple instead of 2-tuple. This is not intended for training as it bypasses autograd. Default: ``False``. - state_v_first (Optional[bool]): + state_v_first (bool, Optional): Store the recurrent state in V-first ``[V, K]`` layout instead of the default ``[K, V]``. Default: ``False``. - cu_seqlens (torch.LongTensor): + cu_seqlens (torch.LongTensor, Optional): Cumulative sequence lengths of shape ``[N+1]`` used for variable-length training, - consistent with the FlashAttention API. - cu_seqlens_cpu (torch.LongTensor): + consistent with the FlashAttention API. Default: `None`. + cu_seqlens_cpu (torch.LongTensor, Optional): Cumulative sequence lengths of shape ``[N+1]`` used for variable-length training, - consistent with the FlashAttention API. - cp_context (Optional[FLACPContext]): + consistent with the FlashAttention API. Default: `None`. + cp_context (FLACPContext, Optional): Context parallel context for distributed training across multiple devices. When provided, ``initial_state`` and ``output_final_state`` are not supported, and ``cu_seqlens`` will be overridden by the context. Default: ``None``. + use_graph (bool, Optional): + Whether to build a fixed-shape varlen chunk list on-device for CUDA or NPU graph capture. + Forward and backward can then be replayed after only the contents of ``cu_seqlens`` change. Requires: + + - Pad ``cu_seqlens`` with zero-length tail sequences to exactly ``max_num_seqs + 1`` entries; + kernels return immediately on sentinel rows (segment id < 0), + so output/gradient rows not covered by real tokens are left undefined. + - Shape ``initial_state`` and the incoming ``dht`` gradient (when used) as ``[max_num_seqs, ...]``. + - Warm up the captured step eagerly so kernel autotuning runs outside the graph; + with ``cp_context``, this also includes NCCL communicator setup. + + With ``cp_context``, the cross-rank state all-gathers are recorded into the graph. This additionally requires: + + - A graph-capturable (non-blocking) NCCL communicator; + all ranks must capture and replay the same collectives in the same order. + - A persistent padded ``cp_context.cu_seqlens`` buffer, refreshed in place before each replay. + - Persistent int32 device scalars ``cp_context.pre_num_ranks_dev`` and ``post_num_ranks_dev``, + refreshed in place before each replay. + + Ascend currently requires flattened variable-length inputs and does not support context parallelism in graph mode. + ``disable_recompute=True`` is not supported. Default: ``False``. + max_num_seqs (int, Optional): + Static upper bound of the sequence count used together with ``use_graph``. + Uses ``cu_seqlens.shape[0] - 1`` when omitted. Default: ``None``. Returns: - Normal mode (return_intermediate_states=False): A tuple (o, final_state) @@ -297,7 +362,7 @@ def chunk_kda( >>> import torch >>> import torch.nn.functional as F >>> from einops import rearrange - >>> from ...ops.kda import chunk_kda + >>> from fla.ops.kda import chunk_kda # inputs with equal lengths (no GVA, HV == H) >>> B, T, H, K, V = 4, 2048, 4, 512, 512 >>> q = torch.randn(B, T, H, K, dtype=torch.bfloat16, device='cuda') @@ -309,16 +374,21 @@ def chunk_kda( >>> A_log = torch.randn(H, dtype=torch.float32, device='cuda') >>> dt_bias = torch.randn(H * K, dtype=torch.float32, device='cuda') >>> o, ht = chunk_kda( - q, k, v, g, beta, - A_log=A_log, - dt_bias=dt_bias, - use_qk_l2norm_in_kernel=True, - use_gate_in_kernel=True, - initial_state=h0, - output_final_state=True - ) + ... q=q, + ... k=k, + ... v=v, + ... g=g, + ... beta=beta, + ... initial_state=h0, + ... output_final_state=True, + ... use_qk_l2norm_in_kernel=True, + ... use_gate_in_kernel=True, + ... A_log=A_log, + ... dt_bias=dt_bias, + ... ) # GVA mode (HV > H) - >>> HV = 8 # 2x more value heads than qk heads + # 2x more value heads than qk heads + >>> HV = 8 >>> v = torch.randn(B, T, HV, V, dtype=torch.bfloat16, device='cuda') >>> g = torch.rand(B, T, HV, K, dtype=torch.bfloat16, device='cuda') >>> beta = torch.rand(B, T, HV, dtype=torch.bfloat16, device='cuda') @@ -326,28 +396,36 @@ def chunk_kda( >>> A_log = torch.randn(HV, dtype=torch.float32, device='cuda') >>> dt_bias = torch.randn(HV * K, dtype=torch.float32, device='cuda') >>> o, ht = chunk_kda( - q, k, v, g, beta, - A_log=A_log, - dt_bias=dt_bias, - use_qk_l2norm_in_kernel=True, - use_gate_in_kernel=True, - initial_state=h0, - output_final_state=True - ) + ... q=q, + ... k=k, + ... v=v, + ... g=g, + ... beta=beta, + ... initial_state=h0, + ... output_final_state=True, + ... use_qk_l2norm_in_kernel=True, + ... use_gate_in_kernel=True, + ... A_log=A_log, + ... dt_bias=dt_bias, + ... ) # for variable-length inputs, the batch size `B` is expected to be 1 and `cu_seqlens` is required >>> q, k, v, beta, g = map(lambda x: rearrange(x, 'b t ... -> 1 (b t) ...'), (q, k, v, beta, g)) # for a batch with 4 sequences, `cu_seqlens` with 5 start/end positions are expected >>> cu_seqlens = q.new_tensor([0, 2048, 4096, 6144, 8192], dtype=torch.long) >>> o, ht = chunk_kda( - q, k, v, g, beta, - A_log=A_log, - dt_bias=dt_bias, - use_qk_l2norm_in_kernel=True, - use_gate_in_kernel=True, - initial_state=h0, - output_final_state=True, - cu_seqlens=cu_seqlens - ) + ... q=q, + ... k=k, + ... v=v, + ... g=g, + ... beta=beta, + ... initial_state=h0, + ... output_final_state=True, + ... use_qk_l2norm_in_kernel=True, + ... use_gate_in_kernel=True, + ... cu_seqlens=cu_seqlens, + ... A_log=A_log, + ... dt_bias=dt_bias, + ... ) """ if 'transpose_state_layout' in kwargs: if state_v_first: @@ -359,11 +437,16 @@ def chunk_kda( ) state_v_first = kwargs.pop('transpose_state_layout') + if use_graph and IS_NPU: + if cp_context is not None: + raise NotImplementedError("Ascend KDA graph mode does not currently support context parallelism.") + if cu_seqlens is None: + raise NotImplementedError("Ascend KDA graph mode currently requires flattened variable-length inputs.") + 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" assert cp_context.cu_seqlens is not None, "cu_seqlens is required for CP" - # Override cu_seqlens and cu_seqlens_cpu with the ones from the context cu_seqlens = cp_context.cu_seqlens if cp_context.cu_seqlens_cpu is not None: cu_seqlens_cpu = cp_context.cu_seqlens_cpu @@ -384,8 +467,9 @@ def chunk_kda( A_log, dt_bias = None, None if use_gate_in_kernel: - assert "A_log" in kwargs, "A_log must be provided when use_gate_in_kernel=True." - A_log, dt_bias = kwargs["A_log"], kwargs.get("dt_bias") + A_log, dt_bias = kwargs.get("A_log"), kwargs.get("dt_bias") + if A_log is None and lower_bound is None: + raise ValueError("`A_log` must be provided when `use_gate_in_kernel=True` and `lower_bound` is not set.") chunk_size = kwargs.pop("chunk_size", 64) if chunk_size not in (32, 64): @@ -400,7 +484,6 @@ def chunk_kda( if allow_neg_eigval and not use_beta_sigmoid_in_kernel: raise ValueError("`allow_neg_eigval=True` requires `use_beta_sigmoid_in_kernel=True`.") - # Validate head dimensions for GVA B, T, H, K, HV = *q.shape, v.shape[2] assert q.shape == k.shape, f"q and k must have the same shape, got q={q.shape} vs k={k.shape}" assert K <= 256, f"Currently we only support key headdim <=256 for KDA, got {K}." @@ -437,4 +520,6 @@ def chunk_kda( disable_recompute, return_intermediate_states, cp_context, + use_graph, + max_num_seqs, ) diff --git a/build/torch-cuda/ops/kda/chunk_bwd.py b/build/torch-cuda/ops/kda/chunk_bwd.py index 13dba60d48cba5123cf72c82b2ded9930666c523..2b6ba658e59d2efe6b9a14760de61caa041e9238 100644 --- a/build/torch-cuda/ops/kda/chunk_bwd.py +++ b/build/torch-cuda/ops/kda/chunk_bwd.py @@ -9,7 +9,7 @@ import torch import triton import triton.language as tl -from ...ops.backends import dispatch +from ...backends import dispatch from ...ops.common.chunk_delta_h import chunk_gated_delta_rule_bwd_dhu, chunk_gated_delta_rule_fwd_h from ...ops.cp import FLACPContext from ...ops.cp.chunk_delta_h import chunk_gated_delta_rule_bwd_dhu_pre_process, expand_h0 @@ -19,8 +19,9 @@ from ...ops.kda.wy_fast import recompute_w_u_fwd from ...ops.utils import chunk_local_cumsum, prepare_chunk_indices from ...ops.utils.cache import fla_cache_autotune from ...ops.utils.constant import RCP_LN2 +from ...ops.utils.graph import get_static_buffer from ...ops.utils.op import exp2 -from ...utils import IS_NVIDIA_HOPPER, autotune_cache_kwargs, check_shared_mem +from ...utils import IS_NVIDIA_HOPPER, IS_NVIDIA_SM100, autotune_cache_kwargs, check_shared_mem BK_LIST = [32, 64] if check_shared_mem() else [16, 32] BV_LIST = [64, 128] if check_shared_mem('ampere') else [16, 32] @@ -60,13 +61,17 @@ def chunk_kda_bwd_kernel_dAv( BK: tl.constexpr, BV: tl.constexpr, IS_VARLEN: tl.constexpr, + USE_GRAPH: tl.constexpr = False, ): - i_t, i_bh = tl.program_id(0), tl.program_id(1) + i_t, i_bh = tl.program_id(0).to(tl.int64), tl.program_id(1).to(tl.int64) i_b, i_hv = i_bh // HV, i_bh % HV i_h = i_hv // (HV // H) 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.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) + i_n = tl.load(chunk_indices + i_t * 2).to(tl.int32) + if USE_GRAPH and i_n < 0: + return + i_t = tl.load(chunk_indices + i_t * 2 + 1).to(tl.int64) + bos, eos = tl.load(cu_seqlens + i_n).to(tl.int64), tl.load(cu_seqlens + i_n + 1).to(tl.int64) T = eos - bos else: bos, eos = i_b * T, i_b * T + T @@ -79,32 +84,39 @@ def chunk_kda_bwd_kernel_dAv( dv += (bos * HV + i_hv) * V dA += (bos * HV + i_hv) * BT - p_A = tl.make_block_ptr(A + (bos * HV + i_hv) * BT, (BT, T), (1, HV*BT), (0, i_t * BT), (BT, BT), (0, 1)) - b_A = tl.load(p_A, boundary_check=(0, 1)) - o_t = i_t * BT + tl.arange(0, BT) m_t = o_t < T + o_A = tl.arange(0, BT) + m_AT = (o_A[:, None] < BT) & m_t[None, :] + p_A = A + (bos * HV + i_hv) * BT + o_A[:, None] + o_t[None, :] * (HV*BT) + b_A = tl.load(p_A, mask=m_AT, other=0.0) + m_A = (o_t[:, None] <= o_t[None, :]) & (m_t[:, None] & m_t) b_A = tl.where(m_A, b_A, 0).to(do.dtype.element_ty) b_dA = tl.zeros([BT, BT], dtype=tl.float32) for i_v in range(tl.cdiv(V, BV)): - p_v = tl.make_block_ptr(v, (V, T), (1, HV*V), (i_v * BV, i_t * BT), (BV, BT), (0, 1)) - p_do = tl.make_block_ptr(do, (T, V), (HV*V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) - p_dv = tl.make_block_ptr(dv, (T, V), (HV*V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) + o_v = i_v * BV + tl.arange(0, BV) + m_v = o_v < V + m_vT = m_v[:, None] & m_t[None, :] + m_tv = m_t[:, None] & m_v[None, :] + p_v = v + o_v[:, None] + o_t[None, :] * (HV*V) + p_do = do + o_t[:, None] * (HV*V) + o_v[None, :] + p_dv = dv + o_t[:, None] * (HV*V) + o_v[None, :] # [BV, BT] - b_v = tl.load(p_v, boundary_check=(0, 1)) + b_v = tl.load(p_v, mask=m_vT, other=0.0) # [BT, BV] - b_do = tl.load(p_do, boundary_check=(0, 1)) + b_do = tl.load(p_do, mask=m_tv, other=0.0) # [BT, BT] - b_dA += tl.dot(b_do, b_v) + b_dA = tl.dot(b_do, b_v, b_dA) # [BT, BV] b_dv = tl.dot(b_A.to(b_do.dtype), b_do) - tl.store(p_dv, b_dv.to(p_dv.dtype.element_ty), boundary_check=(0, 1)) + tl.store(p_dv, b_dv.to(p_dv.dtype.element_ty), mask=m_tv) - p_dA = tl.make_block_ptr(dA, (T, BT), (HV*BT, 1), (i_t * BT, 0), (BT, BT), (1, 0)) + m_dA = m_t[:, None] & (o_A[None, :] < BT) + p_dA = dA + o_t[:, None] * (HV*BT) + o_A[None, :] b_dA = tl.where(o_t[:, None] >= o_t, b_dA * scale, 0.) - tl.store(p_dA, b_dA.to(p_dA.dtype.element_ty), boundary_check=(0, 1)) + tl.store(p_dA, b_dA.to(p_dA.dtype.element_ty), mask=m_dA) @triton.heuristics({ @@ -118,8 +130,9 @@ def chunk_kda_bwd_kernel_dAv( for num_warps in NUM_WARPS for num_stages in [2, 3, 4] if not (IS_NVIDIA_HOPPER and BK == 32 and num_warps == 4) + if not (IS_NVIDIA_SM100 and BK == 32 and num_warps != 2) ], - key=['BT', 'HV', 'STATE_V_FIRST'], + key=['BT', 'HV', 'STATE_V_FIRST', 'K', 'V'], **autotune_cache_kwargs, ) @triton.jit(do_not_specialize=['T']) @@ -154,14 +167,18 @@ def chunk_kda_bwd_kernel_wy_dqkg_fused( BV: tl.constexpr, STATE_V_FIRST: tl.constexpr, IS_VARLEN: tl.constexpr, + USE_GRAPH: tl.constexpr = False, ): - i_t, i_bh = tl.program_id(0), tl.program_id(1) + i_t, i_bh = tl.program_id(0).to(tl.int64), tl.program_id(1).to(tl.int64) i_b, i_hv = i_bh // HV, i_bh % HV i_h = i_hv // (HV // H) if IS_VARLEN: i_tg = i_t.to(tl.int64) - 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) + i_n = tl.load(chunk_indices + i_t * 2).to(tl.int32) + if USE_GRAPH and i_n < 0: + return + i_t = tl.load(chunk_indices + i_t * 2 + 1).to(tl.int64) bos, eos = tl.load(cu_seqlens + i_n).to(tl.int64), tl.load(cu_seqlens + i_n + 1).to(tl.int64) T = (eos - bos).to(tl.int32) NT = tl.cdiv(T, BT) @@ -192,11 +209,13 @@ def chunk_kda_bwd_kernel_wy_dqkg_fused( db += bos * HV + i_hv dA += (bos * HV + i_hv) * BT - p_beta = tl.make_block_ptr(beta, (T,), (HV,), (i_t * BT,), (BT,), (0,)) - b_beta = tl.load(p_beta, boundary_check=(0,)) + p_beta = beta + o_t * HV + b_beta = tl.load(p_beta, mask=m_t, other=0.0) - p_A = tl.make_block_ptr(A, (BT, T), (1, HV * BT), (0, i_t * BT), (BT, BT), (0, 1)) - b_A = tl.load(p_A, boundary_check=(0, 1)) + o_A = tl.arange(0, BT) + m_AT = (o_A[:, None] < BT) & m_t[None, :] + p_A = A + o_A[:, None] + o_t[None, :] * (HV * BT) + b_A = tl.load(p_A, mask=m_AT, other=0.0) b_dA = tl.zeros([BT, BT], dtype=tl.float32) b_db = tl.zeros([BT], dtype=tl.float32) @@ -204,11 +223,12 @@ def chunk_kda_bwd_kernel_wy_dqkg_fused( for i_k in range(tl.cdiv(K, BK)): o_k = i_k * BK + tl.arange(0, BK) m_k = o_k < K + m_tk = m_t[:, None] & m_k[None, :] - p_k = tl.make_block_ptr(k, (T, K), (H*K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) - p_g = tl.make_block_ptr(g, (T, K), (HV*K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) - b_k = tl.load(p_k, boundary_check=(0, 1)) - b_g = tl.load(p_g, boundary_check=(0, 1)).to(tl.float32) + p_k = k + o_t[:, None] * (H*K) + o_k[None, :] + p_g = g + o_t[:, None] * (HV*K) + o_k[None, :] + b_k = tl.load(p_k, mask=m_tk, other=0.0) + b_g = tl.load(p_g, mask=m_tk, other=0.0).to(tl.float32) p_gn = g + (min(T, i_t * BT + BT) - 1).to(tl.int64) * HV*K + o_k b_gn = tl.load(p_gn, mask=m_k, other=0).to(tl.float32) @@ -219,42 +239,45 @@ def chunk_kda_bwd_kernel_wy_dqkg_fused( b_dgk = tl.zeros([BK], dtype=tl.float32) for i_v in range(tl.cdiv(V, BV)): - p_v_new = tl.make_block_ptr(v_new, (T, V), (HV*V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) - p_do = tl.make_block_ptr(do, (T, V), (HV*V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) + o_v = i_v * BV + tl.arange(0, BV) + m_tv = m_t[:, None] & (o_v[None, :] < V) + m_h = (o_v[:, None] < V) & m_k[None, :] + p_v_new = v_new + o_t[:, None] * (HV*V) + o_v[None, :] + p_do = do + o_t[:, None] * (HV*V) + o_v[None, :] if STATE_V_FIRST: - p_h = tl.make_block_ptr(h, (V, K), (K, 1), (i_v * BV, i_k * BK), (BV, BK), (1, 0)) - p_dh = tl.make_block_ptr(dh, (V, K), (K, 1), (i_v * BV, i_k * BK), (BV, BK), (1, 0)) + p_h = h + o_v[:, None] * K + o_k[None, :] + p_dh = dh + o_v[:, None] * K + o_k[None, :] else: - p_h = tl.make_block_ptr(h, (V, K), (1, V), (i_v * BV, i_k * BK), (BV, BK), (0, 1)) - p_dh = tl.make_block_ptr(dh, (V, K), (1, V), (i_v * BV, i_k * BK), (BV, BK), (0, 1)) - p_dv = tl.make_block_ptr(dv, (T, V), (HV*V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) + p_h = h + o_v[:, None] + o_k[None, :] * V + p_dh = dh + o_v[:, None] + o_k[None, :] * V + p_dv = dv + o_t[:, None] * (HV*V) + o_v[None, :] # [BT, BV] - b_v_new = tl.load(p_v_new, boundary_check=(0, 1)) - b_do = tl.load(p_do, boundary_check=(0, 1)) + b_v_new = tl.load(p_v_new, mask=m_tv, other=0.0) + b_do = tl.load(p_do, mask=m_tv, other=0.0) # [BV, BK] - b_h = tl.load(p_h, boundary_check=(0, 1)) - b_dh = tl.load(p_dh, boundary_check=(0, 1)) + b_h = tl.load(p_h, mask=m_h, other=0.0) + b_dh = tl.load(p_dh, mask=m_h, other=0.0) # [BT, BV] - b_dv = tl.load(p_dv, boundary_check=(0, 1)) + b_dv = tl.load(p_dv, mask=m_tv, other=0.0) b_dgk += tl.sum(b_h * b_dh, axis=0) - b_dq += tl.dot(b_do, b_h.to(b_do.dtype)) - b_dk += tl.dot(b_v_new, b_dh.to(b_v_new.dtype)) - b_dw += tl.dot(b_dv.to(b_v_new.dtype), b_h.to(b_v_new.dtype)) + b_dq = tl.dot(b_do, b_h.to(b_do.dtype), b_dq) + b_dk = tl.dot(b_v_new, b_dh.to(b_v_new.dtype), b_dk) + b_dw = tl.dot(b_dv.to(b_v_new.dtype), b_h.to(b_v_new.dtype), b_dw) tl.debug_barrier() # DO NOT REMOVE THIS LINE! if i_k == 0: - p_v = tl.make_block_ptr(v, (T, V), (HV*V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) - p_dv2 = tl.make_block_ptr(dv2, (T, V), (HV*V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) + p_v = v + o_t[:, None] * (HV*V) + o_v[None, :] + p_dv2 = dv2 + o_t[:, None] * (HV*V) + o_v[None, :] - b_v = tl.load(p_v, boundary_check=(0, 1)) + b_v = tl.load(p_v, mask=m_tv, other=0.0) - b_dA += tl.dot(b_dv, tl.trans(b_v)) + b_dA = tl.dot(b_dv, tl.trans(b_v), b_dA) b_dvb = tl.dot(b_A, b_dv) b_dv2 = b_dvb * b_beta[:, None] b_db += tl.sum(b_dvb * b_v, 1) - tl.store(p_dv2, b_dv2.to(p_dv2.dtype.element_ty), boundary_check=(0, 1)) + tl.store(p_dv2, b_dv2.to(p_dv2.dtype.element_ty), mask=m_tv) b_gk_exp = exp2(b_g) b_gb = b_gk_exp * b_beta[:, None] @@ -265,24 +288,24 @@ def chunk_kda_bwd_kernel_wy_dqkg_fused( b_kg = b_k * b_gk_exp b_dw = -b_dw.to(b_A.dtype) - b_dA += tl.dot(b_dw, tl.trans(b_kg.to(b_A.dtype))) + b_dA = tl.dot(b_dw, tl.trans(b_kg.to(b_A.dtype)), b_dA) b_dkgb = tl.dot(b_A, b_dw) b_db += tl.sum(b_dkgb * b_kg, 1) - p_q = tl.make_block_ptr(q, (T, K), (H*K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) - b_q = tl.load(p_q, boundary_check=(0, 1)) + p_q = q + o_t[:, None] * (H*K) + o_k[None, :] + b_q = tl.load(p_q, mask=m_tk, other=0.0) b_kdk = b_k * b_dk b_dgk += tl.sum(b_kdk, axis=0) b_dg = b_q * b_dq - b_kdk + m_last[:, None] * b_dgk + b_kg * b_dkgb * b_beta[:, None] b_dk = b_dk + b_dkgb * b_gb - p_dq = tl.make_block_ptr(dq, (T, K), (HV*K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) - p_dk = tl.make_block_ptr(dk, (T, K), (HV*K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) - p_dg = tl.make_block_ptr(dg, (T, K), (HV*K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) - tl.store(p_dq, b_dq.to(p_dq.dtype.element_ty), boundary_check=(0, 1)) - tl.store(p_dk, b_dk.to(p_dk.dtype.element_ty), boundary_check=(0, 1)) - tl.store(p_dg, b_dg.to(p_dg.dtype.element_ty), boundary_check=(0, 1)) + p_dq = dq + o_t[:, None] * (HV*K) + o_k[None, :] + p_dk = dk + o_t[:, None] * (HV*K) + o_k[None, :] + p_dg = dg + o_t[:, None] * (HV*K) + o_k[None, :] + tl.store(p_dq, b_dq.to(p_dq.dtype.element_ty), mask=m_tk) + tl.store(p_dk, b_dk.to(p_dk.dtype.element_ty), mask=m_tk) + tl.store(p_dg, b_dg.to(p_dg.dtype.element_ty), mask=m_tk) m_A = (o_t[:, None] > o_t[None, :]) & (m_t[:, None] & m_t) b_dA = tl.where(m_A, b_dA * b_beta[None, :], 0) @@ -290,12 +313,14 @@ def chunk_kda_bwd_kernel_wy_dqkg_fused( b_dA = tl.dot(b_A, b_dA.to(b_A.dtype)) b_dA = tl.where(m_A, -b_dA, 0) - p_dA = tl.make_block_ptr(dA, (T, BT), (HV * BT, 1), (i_t * BT, 0), (BT, BT), (1, 0)) - p_db = tl.make_block_ptr(db, (T,), (HV,), (i_t * BT,), (BT,), (0,)) - tl.store(p_dA, b_dA.to(p_dA.dtype.element_ty), boundary_check=(0, 1)) - tl.store(p_db, b_db.to(p_db.dtype.element_ty), boundary_check=(0,)) + m_dA = m_t[:, None] & (o_A[None, :] < BT) + p_dA = dA + o_t[:, None] * (HV * BT) + o_A[None, :] + p_db = db + o_t * HV + tl.store(p_dA, b_dA.to(p_dA.dtype.element_ty), mask=m_dA) + tl.store(p_db, b_db.to(p_db.dtype.element_ty), mask=m_t) +# @dispatch('kda') def chunk_kda_bwd_dAv( q: torch.Tensor, k: torch.Tensor, @@ -306,6 +331,7 @@ def chunk_kda_bwd_dAv( cu_seqlens: torch.LongTensor | None = None, chunk_size: int = 64, chunk_indices: torch.LongTensor | None = None, + use_graph: bool = False, ) -> tuple[torch.Tensor, torch.Tensor]: B, T, H, K, HV, V = *k.shape, do.shape[2], do.shape[-1] BT = chunk_size @@ -322,8 +348,12 @@ def chunk_kda_bwd_dAv( BV = min(max(triton.next_power_of_2(V), 16), CONST_TILING) NT = triton.cdiv(T, BT) if cu_seqlens is None else len(chunk_indices) - dA = v.new_empty(B, T, HV, BT, dtype=torch.float) - dv = torch.empty_like(do) + if use_graph: + dA = get_static_buffer("dAv_dA", (B, T, HV, BT), torch.float, v.device) + dv = get_static_buffer("dAv_dv", tuple(do.shape), do.dtype, do.device) + else: + dA = v.new_empty(B, T, HV, BT, dtype=torch.float) + dv = torch.empty_like(do) grid = (NT, B * HV) chunk_kda_bwd_kernel_dAv[grid]( q=q, @@ -344,11 +374,13 @@ def chunk_kda_bwd_dAv( BT=BT, BK=BK, BV=BV, + USE_GRAPH=use_graph, ) return dA, dv -@dispatch('kda') +# Use triton only, other backends (cutlass, cute, etc.) need additional torch.compile support +# @dispatch('kda') def chunk_kda_bwd_wy_dqkg_fused( q: torch.Tensor, k: torch.Tensor, @@ -366,6 +398,8 @@ def chunk_kda_bwd_wy_dqkg_fused( cu_seqlens: torch.LongTensor | None = None, chunk_size: int = 64, chunk_indices: torch.LongTensor | None = None, + chunk_offsets: torch.LongTensor | None = None, + use_graph: bool = False, ): B, T, H, K, HV, V = *k.shape, v.shape[2], v.shape[-1] BT = chunk_size @@ -375,12 +409,20 @@ def chunk_kda_bwd_wy_dqkg_fused( NT = triton.cdiv(T, BT) if cu_seqlens is None else len(chunk_indices) # dq, dk are allocated at HV dimension; caller reduces to H if GVA - dq = g.new_empty(B, T, HV, K, dtype=torch.float) - dk = g.new_empty(B, T, HV, K, dtype=torch.float) - dv2 = torch.empty_like(v) - dg = torch.empty_like(g, dtype=torch.float) - db = torch.empty_like(beta, dtype=torch.float) - dA = torch.empty_like(A, dtype=torch.float) + if use_graph: + dq = get_static_buffer("wy_dq", (B, T, HV, K), torch.float, q.device) + dk = get_static_buffer("wy_dk", (B, T, HV, K), torch.float, q.device) + dv2 = get_static_buffer("wy_dv2", tuple(v.shape), v.dtype, v.device) + dg = get_static_buffer("wy_dg", tuple(g.shape), torch.float, g.device) + db = get_static_buffer("wy_db", tuple(beta.shape), torch.float, beta.device) + dA = get_static_buffer("wy_dA", tuple(A.shape), torch.float, A.device) + else: + dq = g.new_empty(B, T, HV, K, dtype=torch.float) + dk = g.new_empty(B, T, HV, K, dtype=torch.float) + dv2 = torch.empty_like(v) + dg = torch.empty_like(g, dtype=torch.float) + db = torch.empty_like(beta, dtype=torch.float) + dA = torch.empty_like(A, dtype=torch.float) grid = (NT, B * HV) chunk_kda_bwd_kernel_wy_dqkg_fused[grid]( @@ -411,6 +453,7 @@ def chunk_kda_bwd_wy_dqkg_fused( V=V, BT=BT, STATE_V_FIRST=state_v_first, + USE_GRAPH=use_graph, ) dv = dv2 return dq, dk, dv, db, dg, dA @@ -440,6 +483,8 @@ def chunk_kda_bwd( dt_bias: torch.Tensor | None = None, disable_recompute: bool = False, cp_context: FLACPContext | None = None, + use_graph: bool = False, + chunk_offsets: torch.LongTensor | None = None, **kwargs, ): H, HV = q.shape[2], v.shape[2] @@ -455,7 +500,8 @@ def chunk_kda_bwd( chunk_size=chunk_size, cu_seqlens=cu_seqlens, chunk_indices=chunk_indices, - lower_bound=lower_bound + lower_bound=lower_bound, + use_graph=use_graph, ) w, u, qg, kg = recompute_w_u_fwd( q=q, @@ -466,6 +512,7 @@ def chunk_kda_bwd( gk=g, cu_seqlens=cu_seqlens, chunk_indices=chunk_indices, + use_graph=use_graph, ) if cp_context is not None: # Restore the full initial_state tensor from the compressed version. @@ -480,6 +527,7 @@ def chunk_kda_bwd( output_final_state=False, cu_seqlens=cu_seqlens, chunk_indices=chunk_indices, + chunk_offsets=chunk_offsets, chunk_size=chunk_size, state_v_first=state_v_first, ) @@ -502,6 +550,7 @@ def chunk_kda_bwd( cu_seqlens=cu_seqlens, chunk_size=chunk_size, chunk_indices=chunk_indices, + use_graph=use_graph, ) if cp_context is not None: @@ -521,6 +570,7 @@ def chunk_kda_bwd( context=cp_context, chunk_size=chunk_size, state_v_first=state_v_first, + use_graph=use_graph, ) dh, dh0, dv = chunk_gated_delta_rule_bwd_dhu( @@ -536,6 +586,7 @@ def chunk_kda_bwd( cu_seqlens=cu_seqlens, chunk_size=chunk_size, chunk_indices=chunk_indices, + chunk_offsets=chunk_offsets, state_v_first=state_v_first, ) @@ -555,7 +606,9 @@ def chunk_kda_bwd( cu_seqlens=cu_seqlens, chunk_size=chunk_size, chunk_indices=chunk_indices, + chunk_offsets=chunk_offsets, state_v_first=state_v_first, + use_graph=use_graph, ) dq, dk, db, dg = chunk_kda_bwd_intra( @@ -570,9 +623,10 @@ def chunk_kda_bwd( db=db, dg=dg, cu_seqlens=cu_seqlens, - chunk_size=chunk_size, chunk_indices=chunk_indices, - safe_gate=safe_gate + chunk_size=chunk_size, + safe_gate=safe_gate, + use_graph=use_graph, ) # For GVA, reduce dq and dk from [B, T, HV, K] back to [B, T, H, K] @@ -587,6 +641,7 @@ def chunk_kda_bwd( reverse=True, cu_seqlens=cu_seqlens, chunk_indices=chunk_indices, + use_graph=use_graph, ) if use_gate_in_kernel: dg, dA, dbias = kda_gate_bwd( @@ -594,7 +649,7 @@ def chunk_kda_bwd( A_log=A_log, dt_bias=dt_bias, dyg=dg, - lower_bound=lower_bound + lower_bound=lower_bound, ) return dq, dk, dv, db, dg, dh0, dA, dbias diff --git a/build/torch-cuda/ops/kda/chunk_fwd.py b/build/torch-cuda/ops/kda/chunk_fwd.py index 2dbd234be98339d63f0c7419fef0312d39aac86b..59bdec8d6fa83817633317a3f63dae060706e4e9 100644 --- a/build/torch-cuda/ops/kda/chunk_fwd.py +++ b/build/torch-cuda/ops/kda/chunk_fwd.py @@ -39,6 +39,8 @@ def chunk_kda_fwd( disable_recompute: bool = False, return_intermediate_states: bool = False, cp_context: FLACPContext | None = None, + use_graph: bool = False, + chunk_offsets: torch.LongTensor | None = None, ): # Apply gate activation g_org = None @@ -53,6 +55,7 @@ def chunk_kda_fwd( cu_seqlens=cu_seqlens, chunk_indices=chunk_indices, lower_bound=lower_bound, + use_graph=use_graph, ) else: g = chunk_local_cumsum( @@ -60,7 +63,8 @@ def chunk_kda_fwd( scale=RCP_LN2, chunk_size=chunk_size, cu_seqlens=cu_seqlens, - chunk_indices=chunk_indices + chunk_indices=chunk_indices, + use_graph=use_graph, ) # qg = None if disable_recompute is False @@ -75,7 +79,8 @@ def chunk_kda_fwd( chunk_size=chunk_size, chunk_indices=chunk_indices, safe_gate=safe_gate, - disable_recompute=disable_recompute + disable_recompute=disable_recompute, + use_graph=use_graph, ) if cp_context is not None: @@ -89,6 +94,7 @@ def chunk_kda_fwd( context=cp_context, chunk_size=chunk_size, state_v_first=state_v_first, + use_graph=use_graph, ) h, v_new, final_state = chunk_gated_delta_rule_fwd_h( @@ -101,6 +107,7 @@ def chunk_kda_fwd( cu_seqlens=cu_seqlens, cu_seqlens_cpu=cu_seqlens_cpu, chunk_indices=chunk_indices, + chunk_offsets=chunk_offsets, chunk_size=chunk_size, state_v_first=state_v_first, ) @@ -123,6 +130,7 @@ def chunk_kda_fwd( chunk_size=chunk_size, chunk_indices=chunk_indices, state_v_first=state_v_first, + use_graph=use_graph, ) if disable_recompute is False: # Delete to save memory diff --git a/build/torch-cuda/ops/kda/chunk_intra.py b/build/torch-cuda/ops/kda/chunk_intra.py index 795e5996a4d870cc5b813ac6eb81036a4b270651..166c33bd51087131fc8a0ba7b6ffff5fd38bb0f6 100644 --- a/build/torch-cuda/ops/kda/chunk_intra.py +++ b/build/torch-cuda/ops/kda/chunk_intra.py @@ -9,17 +9,22 @@ import torch import triton import triton.language as tl +from ...backends import dispatch from ...ops.kda.chunk_intra_token_parallel import chunk_kda_fwd_intra_token_parallel from ...ops.kda.wy_fast import recompute_w_u_fwd from ...ops.utils import prepare_chunk_indices from ...ops.utils.cache import fla_cache_autotune -from ...ops.utils.op import exp2, gather -from ...utils import IS_GATHER_SUPPORTED, IS_TF32_SUPPORTED, autotune_cache_kwargs +from ...ops.utils.constant import RCP_LN2 +from ...ops.utils.graph import get_static_buffer +from ...ops.utils.op import exp2, gather, unflatten_program_id +from ...utils import IS_GATHER_SUPPORTED, IS_NPU, IS_TF32_SUPPORTED, TRITON_ABOVE_3_6_0, autotune_cache_kwargs if IS_TF32_SUPPORTED: SOLVE_TRIL_DOT_PRECISION = tl.constexpr('tf32') else: SOLVE_TRIL_DOT_PRECISION = tl.constexpr('ieee') +# TF32 rounding makes safe-gate gradient error grow with context length; bf16x3 reduces it and runs faster than TF32. +SAFE_GATE_DOT_PRECISION = tl.constexpr('bf16x3' if TRITON_ABOVE_3_6_0 and not IS_NPU else 'ieee') ################################################################################ # Fused inter + solve_tril kernel: compute off-diagonal Akk and solve in one pass @@ -60,6 +65,7 @@ def chunk_kda_fwd_kernel_inter_solve_fused( BK: tl.constexpr, IS_VARLEN: tl.constexpr, USE_SAFE_GATE: tl.constexpr, + USE_GRAPH: tl.constexpr = False, ): """ Fused kernel: compute inter-subchunk Akk + solve_tril in one pass. @@ -73,13 +79,16 @@ def chunk_kda_fwd_kernel_inter_solve_fused( 5. Computes merged Akk_inv 6. Writes Akk_inv to Akk """ - i_t, i_bh = tl.program_id(0), tl.program_id(1) + i_t, i_bh = tl.program_id(0).to(tl.int64), tl.program_id(1).to(tl.int64) i_b, i_hv = i_bh // HV, i_bh % HV i_h = i_hv // (HV // H) 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.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) + i_n = tl.load(chunk_indices + i_t * 2).to(tl.int32) + if USE_GRAPH and i_n < 0: + return + i_t = tl.load(chunk_indices + i_t * 2 + 1).to(tl.int64) + bos, eos = tl.load(cu_seqlens + i_n).to(tl.int64), tl.load(cu_seqlens + i_n + 1).to(tl.int64) T = eos - bos else: bos, eos = i_b * T, i_b * T + T @@ -103,6 +112,15 @@ def chunk_kda_fwd_kernel_inter_solve_fused( m_tc1 = (i_tc1 + o_i) < T m_tc2 = (i_tc2 + o_i) < T m_tc3 = (i_tc3 + o_i) < T + o_c0 = i_tc0 + o_i + o_c1 = i_tc1 + o_i + o_c2 = i_tc2 + o_i + o_c3 = i_tc3 + o_i + m_tc0 = o_c0 < T + m_A0 = m_tc0[:, None] & (o_i[None, :] < BT) + m_A1 = m_tc1[:, None] & (o_i[None, :] < BT) + m_A2 = m_tc2[:, None] & (o_i[None, :] < BT) + m_A3 = m_tc3[:, None] & (o_i[None, :] < BT) b_Aqk10 = tl.zeros([BC, BC], dtype=tl.float32) b_Akk10 = tl.zeros([BC, BC], dtype=tl.float32) @@ -126,19 +144,21 @@ def chunk_kda_fwd_kernel_inter_solve_fused( o_k = i_k * BK + tl.arange(0, BK) m_k = o_k < K - p_k0 = tl.make_block_ptr(k, (T, K), (H*K, 1), (i_tc0, i_k * BK), (BC, BK), (1, 0)) - p_g0 = tl.make_block_ptr(g, (T, K), (HV*K, 1), (i_tc0, i_k * BK), (BC, BK), (1, 0)) - b_k0 = tl.load(p_k0, boundary_check=(0, 1)).to(tl.float32) - b_g0 = tl.load(p_g0, boundary_check=(0, 1)).to(tl.float32) + m_ck0 = m_tc0[:, None] & m_k[None, :] + p_k0 = k + o_c0[:, None] * (H*K) + o_k[None, :] + p_g0 = g + o_c0[:, None] * (HV*K) + o_k[None, :] + b_k0 = tl.load(p_k0, mask=m_ck0, other=0.0).to(tl.float32) + b_g0 = tl.load(p_g0, mask=m_ck0, other=0.0).to(tl.float32) if i_tc1 < T: - p_q1 = tl.make_block_ptr(q, (T, K), (H*K, 1), (i_tc1, i_k * BK), (BC, BK), (1, 0)) - p_k1 = tl.make_block_ptr(k, (T, K), (H*K, 1), (i_tc1, i_k * BK), (BC, BK), (1, 0)) - p_g1 = tl.make_block_ptr(g, (T, K), (HV*K, 1), (i_tc1, i_k * BK), (BC, BK), (1, 0)) + m_ck1 = m_tc1[:, None] & m_k[None, :] + p_q1 = q + o_c1[:, None] * (H*K) + o_k[None, :] + p_k1 = k + o_c1[:, None] * (H*K) + o_k[None, :] + p_g1 = g + o_c1[:, None] * (HV*K) + o_k[None, :] # [BC, BK] - b_q1 = tl.load(p_q1, boundary_check=(0, 1)).to(tl.float32) - b_k1 = tl.load(p_k1, boundary_check=(0, 1)).to(tl.float32) - b_g1 = tl.load(p_g1, boundary_check=(0, 1)).to(tl.float32) + b_q1 = tl.load(p_q1, mask=m_ck1, other=0.0).to(tl.float32) + b_k1 = tl.load(p_k1, mask=m_ck1, other=0.0).to(tl.float32) + b_g1 = tl.load(p_g1, mask=m_ck1, other=0.0).to(tl.float32) # [BK] b_gn1 = tl.load(g + i_tc1 * HV*K + o_k, mask=m_k, other=0).to(tl.float32) # [BC, BK] @@ -146,17 +166,18 @@ def chunk_kda_fwd_kernel_inter_solve_fused( # [BK, BC] b_kgt = tl.trans(b_k0 * exp2(b_gn1[None, :] - b_g0)) # [BC, BC] - b_Aqk10 += tl.dot(b_q1 * b_gqn, b_kgt) - b_Akk10 += tl.dot(b_k1 * b_gqn, b_kgt) + b_Aqk10 = tl.dot(b_q1 * b_gqn, b_kgt, b_Aqk10) + b_Akk10 = tl.dot(b_k1 * b_gqn, b_kgt, b_Akk10) if NC >= 3 and i_tc2 < T: - p_q2 = tl.make_block_ptr(q, (T, K), (H*K, 1), (i_tc2, i_k * BK), (BC, BK), (1, 0)) - p_k2 = tl.make_block_ptr(k, (T, K), (H*K, 1), (i_tc2, i_k * BK), (BC, BK), (1, 0)) - p_g2 = tl.make_block_ptr(g, (T, K), (HV*K, 1), (i_tc2, i_k * BK), (BC, BK), (1, 0)) + m_ck2 = m_tc2[:, None] & m_k[None, :] + p_q2 = q + o_c2[:, None] * (H*K) + o_k[None, :] + p_k2 = k + o_c2[:, None] * (H*K) + o_k[None, :] + p_g2 = g + o_c2[:, None] * (HV*K) + o_k[None, :] # [BC, BK] - b_q2 = tl.load(p_q2, boundary_check=(0, 1)).to(tl.float32) - b_k2 = tl.load(p_k2, boundary_check=(0, 1)).to(tl.float32) - b_g2 = tl.load(p_g2, boundary_check=(0, 1)).to(tl.float32) + b_q2 = tl.load(p_q2, mask=m_ck2, other=0.0).to(tl.float32) + b_k2 = tl.load(p_k2, mask=m_ck2, other=0.0).to(tl.float32) + b_g2 = tl.load(p_g2, mask=m_ck2, other=0.0).to(tl.float32) # [BK] b_gn2 = tl.load(g + i_tc2 * HV*K + o_k, mask=m_k, other=0).to(tl.float32) # [BC, BK] @@ -165,22 +186,23 @@ def chunk_kda_fwd_kernel_inter_solve_fused( b_kg2 = b_k2 * b_gqn2 # [BK, BC] b_kgt = tl.trans(b_k0 * exp2(b_gn2[None, :] - b_g0)) - b_Aqk20 += tl.dot(b_qg2, b_kgt) - b_Akk20 += tl.dot(b_kg2, b_kgt) + b_Aqk20 = tl.dot(b_qg2, b_kgt, b_Aqk20) + b_Akk20 = tl.dot(b_kg2, b_kgt, b_Akk20) # [BC, BC] b_kgt = tl.trans(b_k1 * exp2(b_gn2[None, :] - b_g1)) # [BC, BC] - b_Aqk21 += tl.dot(b_qg2, b_kgt) - b_Akk21 += tl.dot(b_kg2, b_kgt) + b_Aqk21 = tl.dot(b_qg2, b_kgt, b_Aqk21) + b_Akk21 = tl.dot(b_kg2, b_kgt, b_Akk21) if NC >= 4 and i_tc3 < T: - p_q3 = tl.make_block_ptr(q, (T, K), (H*K, 1), (i_tc3, i_k * BK), (BC, BK), (1, 0)) - p_k3 = tl.make_block_ptr(k, (T, K), (H*K, 1), (i_tc3, i_k * BK), (BC, BK), (1, 0)) - p_g3 = tl.make_block_ptr(g, (T, K), (HV*K, 1), (i_tc3, i_k * BK), (BC, BK), (1, 0)) + m_ck3 = m_tc3[:, None] & m_k[None, :] + p_q3 = q + o_c3[:, None] * (H*K) + o_k[None, :] + p_k3 = k + o_c3[:, None] * (H*K) + o_k[None, :] + p_g3 = g + o_c3[:, None] * (HV*K) + o_k[None, :] # [BC, BK] - b_q3 = tl.load(p_q3, boundary_check=(0, 1)).to(tl.float32) - b_k3 = tl.load(p_k3, boundary_check=(0, 1)).to(tl.float32) - b_g3 = tl.load(p_g3, boundary_check=(0, 1)).to(tl.float32) + b_q3 = tl.load(p_q3, mask=m_ck3, other=0.0).to(tl.float32) + b_k3 = tl.load(p_k3, mask=m_ck3, other=0.0).to(tl.float32) + b_g3 = tl.load(p_g3, mask=m_ck3, other=0.0).to(tl.float32) # [BK] b_gn3 = tl.load(g + i_tc3 * HV*K + o_k, mask=m_k, other=0).to(tl.float32) # [BC, BK] @@ -190,63 +212,63 @@ def chunk_kda_fwd_kernel_inter_solve_fused( # [BK, BC] b_kgt = tl.trans(b_k0 * exp2(b_gn3[None, :] - b_g0)) # [BC, BC] - b_Aqk30 += tl.dot(b_qg3, b_kgt) - b_Akk30 += tl.dot(b_kg3, b_kgt) + b_Aqk30 = tl.dot(b_qg3, b_kgt, b_Aqk30) + b_Akk30 = tl.dot(b_kg3, b_kgt, b_Akk30) # [BK, BC] b_kgt = tl.trans(b_k1 * exp2(b_gn3[None, :] - b_g1)) # [BC, BC] - b_Aqk31 += tl.dot(b_qg3, b_kgt) - b_Akk31 += tl.dot(b_kg3, b_kgt) + b_Aqk31 = tl.dot(b_qg3, b_kgt, b_Aqk31) + b_Akk31 = tl.dot(b_kg3, b_kgt, b_Akk31) # [BK, BC] b_kgt = tl.trans(b_k2 * exp2(b_gn3[None, :] - b_g2)) # [BC, BC] - b_Aqk32 += tl.dot(b_qg3, b_kgt) - b_Akk32 += tl.dot(b_kg3, b_kgt) + b_Aqk32 = tl.dot(b_qg3, b_kgt, b_Aqk32) + b_Akk32 = tl.dot(b_kg3, b_kgt, b_Akk32) ################################################################################ # save off-diagonal Aqk blocks and prepare Akk ################################################################################ if i_tc1 < T: - p_Aqk10 = tl.make_block_ptr(Aqk, (T, BT), (HV*BT, 1), (i_tc1, 0), (BC, BC), (1, 0)) - tl.store(p_Aqk10, (b_Aqk10 * scale).to(Aqk.dtype.element_ty), boundary_check=(0, 1)) + p_Aqk10 = Aqk + o_c1[:, None] * (HV*BT) + o_i[None, :] + tl.store(p_Aqk10, (b_Aqk10 * scale).to(Aqk.dtype.element_ty), mask=m_A1) - p_b1 = tl.make_block_ptr(beta + bos * HV + i_hv, (T,), (HV,), (i_tc1,), (BC,), (0,)) - b_b1 = tl.load(p_b1, boundary_check=(0,)).to(tl.float32) + p_b1 = beta + bos * HV + i_hv + o_c1 * HV + b_b1 = tl.load(p_b1, mask=m_tc1, other=0.0).to(tl.float32) b_Akk10 = b_Akk10 * b_b1[:, None] if NC >= 3 and i_tc2 < T: - p_Aqk20 = tl.make_block_ptr(Aqk, (T, BT), (HV*BT, 1), (i_tc2, 0), (BC, BC), (1, 0)) - p_Aqk21 = tl.make_block_ptr(Aqk, (T, BT), (HV*BT, 1), (i_tc2, BC), (BC, BC), (1, 0)) - tl.store(p_Aqk20, (b_Aqk20 * scale).to(Aqk.dtype.element_ty), boundary_check=(0, 1)) - tl.store(p_Aqk21, (b_Aqk21 * scale).to(Aqk.dtype.element_ty), boundary_check=(0, 1)) + p_Aqk20 = Aqk + o_c2[:, None] * (HV*BT) + o_i[None, :] + p_Aqk21 = Aqk + o_c2[:, None] * (HV*BT) + (o_i + BC)[None, :] + tl.store(p_Aqk20, (b_Aqk20 * scale).to(Aqk.dtype.element_ty), mask=m_A2) + tl.store(p_Aqk21, (b_Aqk21 * scale).to(Aqk.dtype.element_ty), mask=m_A2) - p_b2 = tl.make_block_ptr(beta + bos * HV + i_hv, (T,), (HV,), (i_tc2,), (BC,), (0,)) - b_b2 = tl.load(p_b2, boundary_check=(0,)).to(tl.float32) + p_b2 = beta + bos * HV + i_hv + o_c2 * HV + b_b2 = tl.load(p_b2, mask=m_tc2, other=0.0).to(tl.float32) b_Akk20 = b_Akk20 * b_b2[:, None] b_Akk21 = b_Akk21 * b_b2[:, None] if NC >= 4 and i_tc3 < T: - p_Aqk30 = tl.make_block_ptr(Aqk, (T, BT), (HV*BT, 1), (i_tc3, 0), (BC, BC), (1, 0)) - p_Aqk31 = tl.make_block_ptr(Aqk, (T, BT), (HV*BT, 1), (i_tc3, BC), (BC, BC), (1, 0)) - p_Aqk32 = tl.make_block_ptr(Aqk, (T, BT), (HV*BT, 1), (i_tc3, 2*BC), (BC, BC), (1, 0)) - tl.store(p_Aqk30, (b_Aqk30 * scale).to(Aqk.dtype.element_ty), boundary_check=(0, 1)) - tl.store(p_Aqk31, (b_Aqk31 * scale).to(Aqk.dtype.element_ty), boundary_check=(0, 1)) - tl.store(p_Aqk32, (b_Aqk32 * scale).to(Aqk.dtype.element_ty), boundary_check=(0, 1)) - - p_b3 = tl.make_block_ptr(beta + bos * HV + i_hv, (T,), (HV,), (i_tc3,), (BC,), (0,)) - b_b3 = tl.load(p_b3, boundary_check=(0,)).to(tl.float32) + p_Aqk30 = Aqk + o_c3[:, None] * (HV*BT) + o_i[None, :] + p_Aqk31 = Aqk + o_c3[:, None] * (HV*BT) + (o_i + BC)[None, :] + p_Aqk32 = Aqk + o_c3[:, None] * (HV*BT) + (o_i + 2*BC)[None, :] + tl.store(p_Aqk30, (b_Aqk30 * scale).to(Aqk.dtype.element_ty), mask=m_A3) + tl.store(p_Aqk31, (b_Aqk31 * scale).to(Aqk.dtype.element_ty), mask=m_A3) + tl.store(p_Aqk32, (b_Aqk32 * scale).to(Aqk.dtype.element_ty), mask=m_A3) + + p_b3 = beta + bos * HV + i_hv + o_c3 * HV + b_b3 = tl.load(p_b3, mask=m_tc3, other=0.0).to(tl.float32) b_Akk30 = b_Akk30 * b_b3[:, None] b_Akk31 = b_Akk31 * b_b3[:, None] b_Akk32 = b_Akk32 * b_b3[:, None] - p_Akk00 = tl.make_block_ptr(Akkd, (T, BC), (HV*BC, 1), (i_tc0, 0), (BC, BC), (1, 0)) - p_Akk11 = tl.make_block_ptr(Akkd, (T, BC), (HV*BC, 1), (i_tc1, 0), (BC, BC), (1, 0)) - b_Ai00 = tl.load(p_Akk00, boundary_check=(0, 1)).to(tl.float32) - b_Ai11 = tl.load(p_Akk11, boundary_check=(0, 1)).to(tl.float32) + p_Akk00 = Akkd + o_c0[:, None] * (HV*BC) + o_i[None, :] + p_Akk11 = Akkd + o_c1[:, None] * (HV*BC) + o_i[None, :] + b_Ai00 = tl.load(p_Akk00, mask=m_A0, other=0.0).to(tl.float32) + b_Ai11 = tl.load(p_Akk11, mask=m_A1, other=0.0).to(tl.float32) if NC >= 3: - p_Akk22 = tl.make_block_ptr(Akkd, (T, BC), (HV*BC, 1), (i_tc2, 0), (BC, BC), (1, 0)) - b_Ai22 = tl.load(p_Akk22, boundary_check=(0, 1)).to(tl.float32) + p_Akk22 = Akkd + o_c2[:, None] * (HV*BC) + o_i[None, :] + b_Ai22 = tl.load(p_Akk22, mask=m_A2, other=0.0).to(tl.float32) if NC >= 4: - p_Akk33 = tl.make_block_ptr(Akkd, (T, BC), (HV*BC, 1), (i_tc3, 0), (BC, BC), (1, 0)) - b_Ai33 = tl.load(p_Akk33, boundary_check=(0, 1)).to(tl.float32) + p_Akk33 = Akkd + o_c3[:, None] * (HV*BC) + o_i[None, :] + b_Ai33 = tl.load(p_Akk33, mask=m_A3, other=0.0).to(tl.float32) ################################################################################ # forward substitution on diagonals @@ -340,29 +362,29 @@ def chunk_kda_fwd_kernel_inter_solve_fused( # store full Akk_inv to Akk ################################################################################ - p_Akk00 = tl.make_block_ptr(Akk, (T, BT), (HV*BT, 1), (i_tc0, 0), (BC, BC), (1, 0)) - p_Akk10 = tl.make_block_ptr(Akk, (T, BT), (HV*BT, 1), (i_tc1, 0), (BC, BC), (1, 0)) - p_Akk11 = tl.make_block_ptr(Akk, (T, BT), (HV*BT, 1), (i_tc1, BC), (BC, BC), (1, 0)) + p_Akk00 = Akk + o_c0[:, None] * (HV*BT) + o_i[None, :] + p_Akk10 = Akk + o_c1[:, None] * (HV*BT) + o_i[None, :] + p_Akk11 = Akk + o_c1[:, None] * (HV*BT) + (o_i + BC)[None, :] - tl.store(p_Akk00, b_Ai00.to(Akk.dtype.element_ty), boundary_check=(0, 1)) - tl.store(p_Akk10, b_Ai10.to(Akk.dtype.element_ty), boundary_check=(0, 1)) - tl.store(p_Akk11, b_Ai11.to(Akk.dtype.element_ty), boundary_check=(0, 1)) + tl.store(p_Akk00, b_Ai00.to(Akk.dtype.element_ty), mask=m_A0) + tl.store(p_Akk10, b_Ai10.to(Akk.dtype.element_ty), mask=m_A1) + tl.store(p_Akk11, b_Ai11.to(Akk.dtype.element_ty), mask=m_A1) if NC >= 3: - p_Akk20 = tl.make_block_ptr(Akk, (T, BT), (HV*BT, 1), (i_tc2, 0), (BC, BC), (1, 0)) - p_Akk21 = tl.make_block_ptr(Akk, (T, BT), (HV*BT, 1), (i_tc2, BC), (BC, BC), (1, 0)) - p_Akk22 = tl.make_block_ptr(Akk, (T, BT), (HV*BT, 1), (i_tc2, 2*BC), (BC, BC), (1, 0)) - tl.store(p_Akk20, b_Ai20.to(Akk.dtype.element_ty), boundary_check=(0, 1)) - tl.store(p_Akk21, b_Ai21.to(Akk.dtype.element_ty), boundary_check=(0, 1)) - tl.store(p_Akk22, b_Ai22.to(Akk.dtype.element_ty), boundary_check=(0, 1)) + p_Akk20 = Akk + o_c2[:, None] * (HV*BT) + o_i[None, :] + p_Akk21 = Akk + o_c2[:, None] * (HV*BT) + (o_i + BC)[None, :] + p_Akk22 = Akk + o_c2[:, None] * (HV*BT) + (o_i + 2*BC)[None, :] + tl.store(p_Akk20, b_Ai20.to(Akk.dtype.element_ty), mask=m_A2) + tl.store(p_Akk21, b_Ai21.to(Akk.dtype.element_ty), mask=m_A2) + tl.store(p_Akk22, b_Ai22.to(Akk.dtype.element_ty), mask=m_A2) if NC >= 4: - p_Akk30 = tl.make_block_ptr(Akk, (T, BT), (HV*BT, 1), (i_tc3, 0), (BC, BC), (1, 0)) - p_Akk31 = tl.make_block_ptr(Akk, (T, BT), (HV*BT, 1), (i_tc3, BC), (BC, BC), (1, 0)) - p_Akk32 = tl.make_block_ptr(Akk, (T, BT), (HV*BT, 1), (i_tc3, 2*BC), (BC, BC), (1, 0)) - p_Akk33 = tl.make_block_ptr(Akk, (T, BT), (HV*BT, 1), (i_tc3, 3*BC), (BC, BC), (1, 0)) - tl.store(p_Akk30, b_Ai30.to(Akk.dtype.element_ty), boundary_check=(0, 1)) - tl.store(p_Akk31, b_Ai31.to(Akk.dtype.element_ty), boundary_check=(0, 1)) - tl.store(p_Akk32, b_Ai32.to(Akk.dtype.element_ty), boundary_check=(0, 1)) - tl.store(p_Akk33, b_Ai33.to(Akk.dtype.element_ty), boundary_check=(0, 1)) + p_Akk30 = Akk + o_c3[:, None] * (HV*BT) + o_i[None, :] + p_Akk31 = Akk + o_c3[:, None] * (HV*BT) + (o_i + BC)[None, :] + p_Akk32 = Akk + o_c3[:, None] * (HV*BT) + (o_i + 2*BC)[None, :] + p_Akk33 = Akk + o_c3[:, None] * (HV*BT) + (o_i + 3*BC)[None, :] + tl.store(p_Akk30, b_Ai30.to(Akk.dtype.element_ty), mask=m_A3) + tl.store(p_Akk31, b_Ai31.to(Akk.dtype.element_ty), mask=m_A3) + tl.store(p_Akk32, b_Ai32.to(Akk.dtype.element_ty), mask=m_A3) + tl.store(p_Akk33, b_Ai33.to(Akk.dtype.element_ty), mask=m_A3) @triton.heuristics({ @@ -405,17 +427,23 @@ def chunk_kda_bwd_kernel_intra( NC: tl.constexpr, IS_VARLEN: tl.constexpr, SAFE_GATE: tl.constexpr, + GATE_SHIFT: tl.constexpr, USE_GATHER: tl.constexpr, + USE_GRAPH: tl.constexpr = False, ): - i_kc, i_t, i_bh = tl.program_id(0), tl.program_id(1), tl.program_id(2) + i_kc, i_t = unflatten_program_id(tl.cdiv(K, BK) * NC) + i_bh = tl.program_id(1).to(tl.int64) i_b, i_hv = i_bh // HV, i_bh % HV i_h = i_hv // (HV // H) i_k, i_i = i_kc // NC, i_kc % NC - all = B * T + total_tokens = B * T 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.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) + i_n = tl.load(chunk_indices + i_t * 2).to(tl.int32) + if USE_GRAPH and i_n < 0: + return + i_t = tl.load(chunk_indices + i_t * 2 + 1).to(tl.int64) + bos, eos = tl.load(cu_seqlens + i_n).to(tl.int64), tl.load(cu_seqlens + i_n + 1).to(tl.int64) else: bos, eos = i_b * T, i_b * T + T T = eos - bos @@ -440,35 +468,42 @@ def chunk_kda_bwd_kernel_intra( dk2 += (bos * HV + i_hv) * K dg += (bos * HV + i_hv) * K dg2 += (bos * HV + i_hv) * K - db += (i_k * all + bos) * HV + i_hv + db += (i_k * total_tokens + bos) * HV + i_hv - p_g = tl.make_block_ptr(g, (T, K), (HV*K, 1), (i_ti, i_k * BK), (BC, BK), (1, 0)) - b_g = tl.load(p_g, boundary_check=(0, 1)).to(tl.float32) + o_i = tl.arange(0, BC) + o_c = i_ti + o_i + m_c = o_c < T + m_ck = m_c[:, None] & m_k[None, :] + m_dAf = m_c[:, None] & (o_i[None, :] < BT) + m_dAt = (o_i[:, None] < BT) & m_c[None, :] + p_g = g + o_c[:, None] * (HV*K) + o_k[None, :] + b_g = tl.load(p_g, mask=m_ck, other=0.0).to(tl.float32) - p_b = tl.make_block_ptr(beta, (T,), (HV,), (i_ti,), (BC,), (0,)) - b_b = tl.load(p_b, boundary_check=(0,)) + p_b = beta + o_c * HV + b_b = tl.load(p_b, mask=m_c, other=0.0) b_dq2 = tl.zeros([BC, BK], dtype=tl.float32) b_dk2 = tl.zeros([BC, BK], dtype=tl.float32) if i_i > 0: - p_gn = g + i_ti * HV*K + o_k - # [BK,] - b_gn = tl.load(p_gn, mask=m_k, other=0).to(tl.float32)[None, :] + # [1, BK] + b_gn = tl.load(g + i_ti * HV*K + o_k, mask=m_k, other=0).to(tl.float32)[None, :] for i_j in range(0, i_i): - p_k = tl.make_block_ptr(k, (T, K), (H*K, 1), (i_t * BT + i_j * BC, i_k * BK), (BC, BK), (1, 0)) - p_gk = tl.make_block_ptr(g, (T, K), (HV*K, 1), (i_t * BT + i_j * BC, i_k * BK), (BC, BK), (1, 0)) - p_dAqk = tl.make_block_ptr(dAqk, (T, BT), (HV*BT, 1), (i_ti, i_j * BC), (BC, BC), (1, 0)) - p_dAkk = tl.make_block_ptr(dAkk, (T, BT), (HV*BT, 1), (i_ti, i_j * BC), (BC, BC), (1, 0)) + o_j = i_t * BT + i_j * BC + o_i + m_jk = (o_j < T)[:, None] & m_k[None, :] + p_k = k + o_j[:, None] * (H*K) + o_k[None, :] + p_gk = g + o_j[:, None] * (HV*K) + o_k[None, :] + p_dAqk = dAqk + o_c[:, None] * (HV*BT) + (i_j * BC + o_i)[None, :] + p_dAkk = dAkk + o_c[:, None] * (HV*BT) + (i_j * BC + o_i)[None, :] # [BC, BK] - b_k = tl.load(p_k, boundary_check=(0, 1)) - b_gk = tl.load(p_gk, boundary_check=(0, 1)) + b_k = tl.load(p_k, mask=m_jk, other=0.0) + b_gk = tl.load(p_gk, mask=m_jk, other=0.0) b_kg = b_k * exp2(b_gn - b_gk) # [BC, BC] - b_dAqk = tl.load(p_dAqk, boundary_check=(0, 1)) - b_dAkk = tl.load(p_dAkk, boundary_check=(0, 1)) + b_dAqk = tl.load(p_dAqk, mask=m_dAf, other=0.0) + b_dAkk = tl.load(p_dAkk, mask=m_dAf, other=0.0) # [BC, BK] - b_dq2 += tl.dot(b_dAqk, b_kg) - b_dk2 += tl.dot(b_dAkk, b_kg) + b_dq2 = tl.dot(b_dAqk, b_kg, b_dq2) + b_dk2 = tl.dot(b_dAkk, b_kg, b_dk2) b_gqn = exp2(b_g - b_gn) b_dq2 *= b_gqn b_dk2 *= b_gqn @@ -479,22 +514,22 @@ def chunk_kda_bwd_kernel_intra( p_kj = k + i_ti * H*K + o_k p_gkj = g + i_ti * HV*K + o_k - p_q = tl.make_block_ptr(q, (T, K), (H*K, 1), (i_ti, i_k * BK), (BC, BK), (1, 0)) - p_k = tl.make_block_ptr(k, (T, K), (H*K, 1), (i_ti, i_k * BK), (BC, BK), (1, 0)) - b_q = tl.load(p_q, boundary_check=(0, 1)) - b_k = tl.load(p_k, boundary_check=(0, 1)) - + p_q = q + o_c[:, None] * (H*K) + o_k[None, :] + p_k = k + o_c[:, None] * (H*K) + o_k[None, :] + b_q = tl.load(p_q, mask=m_ck, other=0.0) + b_k = tl.load(p_k, mask=m_ck, other=0.0) + # match the forward anchor when recomputing query-side gate factors if SAFE_GATE: if USE_GATHER: - b_gn = gather(b_g, tl.full([1, BK], min(BC//2, T - i_ti - 1), dtype=tl.int16), axis=0) + b_gn = gather(b_g, tl.zeros([1, BK], dtype=tl.int16), axis=0) else: - p_gn = g + (i_ti + min(BC // 2, T - i_ti - 1)) * HV*K + o_k - b_gn = tl.load(p_gn, mask=m_k, other=0)[None, :] + b_gn = tl.load(g + i_ti * HV*K + o_k, mask=m_k, other=0)[None, :] + b_gn = b_gn - 0.5 * (BC - 1) * GATE_SHIFT - p_dAqk = tl.make_block_ptr(dAqk, (T, BT), (HV*BT, 1), (i_ti, i_i * BC), (BC, BC), (1, 0)) - p_dAkk = tl.make_block_ptr(dAkk, (T, BT), (HV*BT, 1), (i_ti, i_i * BC), (BC, BC), (1, 0)) - b_dAqk_diag_qk = tl.load(p_dAqk, boundary_check=(0, 1)).to(tl.float32) - b_dAkk_diag_qk = tl.load(p_dAkk, boundary_check=(0, 1)).to(tl.float32) + p_dAqk = dAqk + o_c[:, None] * (HV*BT) + (i_i * BC + o_i)[None, :] + p_dAkk = dAkk + o_c[:, None] * (HV*BT) + (i_i * BC + o_i)[None, :] + b_dAqk_diag_qk = tl.load(p_dAqk, mask=m_dAf, other=0.0).to(tl.float32) + b_dAkk_diag_qk = tl.load(p_dAkk, mask=m_dAf, other=0.0).to(tl.float32) m_i_diag_qk = (o_i[:, None] >= o_i[None, :]) & ((i_ti + o_i[:, None]) < T) & ((i_ti + o_i[None, :]) < T) m_j_diag_qk = (i_ti + o_i[:, None]) < T @@ -502,12 +537,12 @@ def chunk_kda_bwd_kernel_intra( b_dAqk_diag_qk = tl.where(m_i_diag_qk, b_dAqk_diag_qk, 0.) b_dAkk_diag_qk = tl.where(m_i_diag_qk, b_dAkk_diag_qk, 0.) b_g_diag_qk = tl.where(m_j_diag_qk, b_g - b_gn, 0.) - exp_b_g_diag_qk = tl.where(m_j_diag_qk, exp2(b_g_diag_qk), 0.) - exp_neg_b_g_diag_qk = tl.where(m_j_diag_qk, exp2(-b_g_diag_qk), 0.) + b_gq_diag_qk = tl.where(m_j_diag_qk, exp2(b_g_diag_qk), 0.) + b_gk_diag_qk = tl.where(m_j_diag_qk, exp2(-b_g_diag_qk), 0.) - b_k_exp_diag_qk = b_k * exp_neg_b_g_diag_qk - b_dq2 += tl.dot(b_dAqk_diag_qk, b_k_exp_diag_qk) * exp_b_g_diag_qk - b_dk2 += tl.dot(b_dAkk_diag_qk, b_k_exp_diag_qk) * exp_b_g_diag_qk + b_k_exp_diag_qk = b_k * b_gk_diag_qk + b_dq2 += tl.dot(b_dAqk_diag_qk, b_k_exp_diag_qk, input_precision=SAFE_GATE_DOT_PRECISION) * b_gq_diag_qk + b_dk2 += tl.dot(b_dAkk_diag_qk, b_k_exp_diag_qk, input_precision=SAFE_GATE_DOT_PRECISION) * b_gq_diag_qk else: for j in range(0, min(BC, T - i_t * BT - i_i * BC)): # [BC] @@ -529,72 +564,73 @@ def chunk_kda_bwd_kernel_intra( b_db = tl.sum(b_dk2 * b_k, 1) b_dk2 *= b_b[:, None] - p_dq = tl.make_block_ptr(dq, (T, K), (HV*K, 1), (i_ti, i_k * BK), (BC, BK), (1, 0)) - p_dq2 = tl.make_block_ptr(dq2, (T, K), (HV*K, 1), (i_ti, i_k * BK), (BC, BK), (1, 0)) - p_db = tl.make_block_ptr(db, (T,), (HV,), (i_ti,), (BC,), (0,)) + p_dq = dq + o_c[:, None] * (HV*K) + o_k[None, :] + p_dq2 = dq2 + o_c[:, None] * (HV*K) + o_k[None, :] + p_db = db + o_c * HV b_dg2 = b_q * b_dq2 - b_dq2 = b_dq2 + tl.load(p_dq, boundary_check=(0, 1)) - tl.store(p_dq2, b_dq2.to(p_dq2.dtype.element_ty), boundary_check=(0, 1)) - tl.store(p_db, b_db.to(p_db.dtype.element_ty), boundary_check=(0,)) + b_dq2 = b_dq2 + tl.load(p_dq, mask=m_ck, other=0.0) + tl.store(p_dq2, b_dq2.to(p_dq2.dtype.element_ty), mask=m_ck) + tl.store(p_db, b_db.to(p_db.dtype.element_ty), mask=m_c) tl.debug_barrier() b_dkt = tl.zeros([BC, BK], dtype=tl.float32) NC = min(NC, tl.cdiv(T - i_t * BT, BC)) if i_i < NC - 1: - p_gn = g + (min(i_ti + BC, T) - 1) * HV*K + o_k - # [BK,] - b_gn = tl.load(p_gn, mask=m_k, other=0).to(tl.float32)[None, :] + # [1, BK] + b_gn = tl.load(g + (min(i_ti + BC, T) - 1) * HV*K + o_k, mask=m_k, other=0).to(tl.float32)[None, :] for i_j in range(i_i + 1, NC): - p_q = tl.make_block_ptr(q, (T, K), (H*K, 1), (i_t*BT+i_j*BC, i_k*BK), (BC, BK), (1, 0)) - p_k = tl.make_block_ptr(k, (T, K), (H*K, 1), (i_t * BT + i_j * BC, i_k * BK), (BC, BK), (1, 0)) - p_gk = tl.make_block_ptr(g, (T, K), (HV*K, 1), (i_t * BT + i_j * BC, i_k*BK), (BC, BK), (1, 0)) - p_b = tl.make_block_ptr(beta, (T,), (HV,), (i_t * BT + i_j * BC,), (BC,), (0,)) - p_dAqk = tl.make_block_ptr(dAqk, (BT, T), (1, HV*BT), (i_i * BC, i_t * BT + i_j * BC), (BC, BC), (0, 1)) - p_dAkk = tl.make_block_ptr(dAkk, (BT, T), (1, HV*BT), (i_i * BC, i_t * BT + i_j * BC), (BC, BC), (0, 1)) + o_j = i_t * BT + i_j * BC + o_i + m_j = o_j < T + m_jk = m_j[:, None] & m_k[None, :] + m_dAj = (o_i[:, None] < BT) & m_j[None, :] + p_q = q + o_j[:, None] * (H*K) + o_k[None, :] + p_k = k + o_j[:, None] * (H*K) + o_k[None, :] + p_gk = g + o_j[:, None] * (HV*K) + o_k[None, :] + p_b = beta + o_j * HV + p_dAqk = dAqk + (i_i * BC + o_i)[:, None] + o_j[None, :] * (HV*BT) + p_dAkk = dAkk + (i_i * BC + o_i)[:, None] + o_j[None, :] * (HV*BT) # [BC] - b_b = tl.load(p_b, boundary_check=(0,)) + b_b = tl.load(p_b, mask=m_j, other=0.0) # [BC, BK] - b_q = tl.load(p_q, boundary_check=(0, 1)) - b_kb = tl.load(p_k, boundary_check=(0, 1)) * b_b[:, None] - b_gk = tl.load(p_gk, boundary_check=(0, 1)).to(tl.float32) + b_q = tl.load(p_q, mask=m_jk, other=0.0) + b_kb = tl.load(p_k, mask=m_jk, other=0.0) * b_b[:, None] + b_gk = tl.load(p_gk, mask=m_jk, other=0.0).to(tl.float32) # [BC, BC] - b_dAqk = tl.load(p_dAqk, boundary_check=(0, 1)) - b_dAkk = tl.load(p_dAkk, boundary_check=(0, 1)) + b_dAqk = tl.load(p_dAqk, mask=m_dAj, other=0.0) + b_dAkk = tl.load(p_dAkk, mask=m_dAj, other=0.0) - o_j = i_t * BT + i_j * BC + o_i - m_j = o_j < T # [BC, BK] b_gkn = exp2(b_gk - b_gn) b_qg = b_q * tl.where(m_j[:, None], b_gkn, 0) b_kbg = b_kb * tl.where(m_j[:, None], b_gkn, 0) # [BC, BK] # (SY 09/17) important to not use bf16 here to have a good precision. - b_dkt += tl.dot(b_dAqk, b_qg) - b_dkt += tl.dot(b_dAkk, b_kbg) + b_dkt = tl.dot(b_dAqk, b_qg, b_dkt) + b_dkt = tl.dot(b_dAkk, b_kbg, b_dkt) b_dkt *= exp2(b_gn - b_g) o_dA = i_ti * HV*BT + i_i * BC + o_i p_qj = q + i_ti * H*K + o_k p_kj = k + i_ti * H*K + o_k p_gkj = g + i_ti * HV*K + o_k p_bj = beta + i_ti * HV - + # use the same anchor for key-side gradients so the shared offset cancels if SAFE_GATE: if USE_GATHER: - b_gn = gather(b_g, tl.full([1, BK], min(BC//2, T - i_ti - 1), dtype=tl.int16), axis=0) + b_gn = gather(b_g, tl.zeros([1, BK], dtype=tl.int16), axis=0) else: - p_gn = g + (i_ti + min(BC // 2, T - i_ti - 1)) * HV*K + o_k - b_gn = tl.load(p_gn, mask=m_k, other=0).to(tl.float32)[None, :] - p_q = tl.make_block_ptr(q, (T, K), (H*K, 1), (i_ti, i_k * BK), (BC, BK), (1, 0)) - b_q = tl.load(p_q, boundary_check=(0, 1)) - p_b = tl.make_block_ptr(beta, (T,), (HV,), (i_ti,), (BC,), (0,)) - b_b = tl.load(p_b, boundary_check=(0,)) - - p_dAqk = tl.make_block_ptr(dAqk, (BT, T), (1, HV*BT), (i_i * BC, i_ti), (BC, BC), (0, 1)) - p_dAkk = tl.make_block_ptr(dAkk, (BT, T), (1, HV*BT), (i_i * BC, i_ti), (BC, BC), (0, 1)) - b_dAqk_diag_kk = tl.load(p_dAqk, boundary_check=(0, 1)).to(tl.float32) - b_dAkk_diag_kk = tl.load(p_dAkk, boundary_check=(0, 1)).to(tl.float32) + b_gn = tl.load(g + i_ti * HV*K + o_k, mask=m_k, other=0).to(tl.float32)[None, :] + b_gn = b_gn - 0.5 * (BC - 1) * GATE_SHIFT + p_q = q + o_c[:, None] * (H*K) + o_k[None, :] + b_q = tl.load(p_q, mask=m_ck, other=0.0) + p_b = beta + o_c * HV + b_b = tl.load(p_b, mask=m_c, other=0.0) + + p_dAqk = dAqk + (i_i * BC + o_i)[:, None] + o_c[None, :] * (HV*BT) + p_dAkk = dAkk + (i_i * BC + o_i)[:, None] + o_c[None, :] * (HV*BT) + b_dAqk_diag_kk = tl.load(p_dAqk, mask=m_dAt, other=0.0).to(tl.float32) + b_dAkk_diag_kk = tl.load(p_dAkk, mask=m_dAt, other=0.0).to(tl.float32) m_i_diag_kk = (o_i[:, None] <= o_i[None, :]) & ((i_ti + o_i[:, None]) < T) & ((i_ti + o_i[None, :]) < T) m_j_diag_kk = (i_ti + o_i[:, None]) < T @@ -603,14 +639,14 @@ def chunk_kda_bwd_kernel_intra( b_dAkk_diag_kk = tl.where(m_i_diag_kk, b_dAkk_diag_kk, 0.) # ensure numerical stability b_g_diag_kk = tl.where(m_j_diag_kk, b_g - b_gn, 0.) - exp_b_g_diag_kk = tl.where(m_j_diag_kk, exp2(b_g_diag_kk), 0.) - exp_neg_b_g_diag_kk = tl.where(m_j_diag_kk, exp2(-b_g_diag_kk), 0.) + b_gq_diag_kk = tl.where(m_j_diag_kk, exp2(b_g_diag_kk), 0.) + b_gk_diag_kk = tl.where(m_j_diag_kk, exp2(-b_g_diag_kk), 0.) - b_q_exp = b_q * exp_b_g_diag_kk - b_kb_exp = b_k * b_b[:, None] * exp_b_g_diag_kk + b_q_exp = b_q * b_gq_diag_kk + b_kb_exp = b_k * b_b[:, None] * b_gq_diag_kk - b_dkt += tl.dot(b_dAqk_diag_kk, b_q_exp) * exp_neg_b_g_diag_kk - b_dkt += tl.dot(b_dAkk_diag_kk, b_kb_exp) * exp_neg_b_g_diag_kk + b_dkt += tl.dot(b_dAqk_diag_kk, b_q_exp, input_precision=SAFE_GATE_DOT_PRECISION) * b_gk_diag_kk + b_dkt += tl.dot(b_dAkk_diag_kk, b_kb_exp, input_precision=SAFE_GATE_DOT_PRECISION) * b_gk_diag_kk else: for j in range(0, min(BC, T - i_t * BT - i_i * BC)): # [BC,] @@ -630,17 +666,17 @@ def chunk_kda_bwd_kernel_intra( p_kj += H*K p_gkj += HV*K p_bj += HV - p_dk = tl.make_block_ptr(dk, (T, K), (HV*K, 1), (i_ti, i_k * BK), (BC, BK), (1, 0)) - p_dk2 = tl.make_block_ptr(dk2, (T, K), (HV*K, 1), (i_ti, i_k * BK), (BC, BK), (1, 0)) - p_dg = tl.make_block_ptr(dg, (T, K), (HV*K, 1), (i_ti, i_k * BK), (BC, BK), (1, 0)) - p_dg2 = tl.make_block_ptr(dg2, (T, K), (HV*K, 1), (i_ti, i_k * BK), (BC, BK), (1, 0)) + p_dk = dk + o_c[:, None] * (HV*K) + o_k[None, :] + p_dk2 = dk2 + o_c[:, None] * (HV*K) + o_k[None, :] + p_dg = dg + o_c[:, None] * (HV*K) + o_k[None, :] + p_dg2 = dg2 + o_c[:, None] * (HV*K) + o_k[None, :] - b_dg2 += (b_dk2 - b_dkt) * b_k + tl.load(p_dg, boundary_check=(0, 1)) - b_dk2 += tl.load(p_dk, boundary_check=(0, 1)) + b_dg2 += (b_dk2 - b_dkt) * b_k + tl.load(p_dg, mask=m_ck, other=0.0) + b_dk2 += tl.load(p_dk, mask=m_ck, other=0.0) b_dk2 += b_dkt - tl.store(p_dk2, b_dk2.to(p_dk2.dtype.element_ty), boundary_check=(0, 1)) - tl.store(p_dg2, b_dg2.to(p_dg2.dtype.element_ty), boundary_check=(0, 1)) + tl.store(p_dk2, b_dk2.to(p_dk2.dtype.element_ty), mask=m_ck) + tl.store(p_dg2, b_dg2.to(p_dg2.dtype.element_ty), mask=m_ck) @triton.heuristics({ @@ -674,15 +710,20 @@ def chunk_kda_fwd_kernel_intra_sub_chunk( BC: tl.constexpr, BK: tl.constexpr, IS_VARLEN: tl.constexpr, + GATE_SHIFT: tl.constexpr, USE_GATHER: tl.constexpr, + USE_GRAPH: tl.constexpr = False, ): - i_t, i_i, i_bh = tl.program_id(0), tl.program_id(1), tl.program_id(2) + i_t, i_i, i_bh = tl.program_id(0).to(tl.int64), tl.program_id(1).to(tl.int64), tl.program_id(2).to(tl.int64) i_b, i_hv = i_bh // HV, i_bh % HV i_h = i_hv // (HV // H) 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.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) + i_n = tl.load(chunk_indices + i_t * 2).to(tl.int32) + if USE_GRAPH and i_n < 0: + return + i_t = tl.load(chunk_indices + i_t * 2 + 1).to(tl.int64) + bos, eos = tl.load(cu_seqlens + i_n).to(tl.int64), tl.load(cu_seqlens + i_n + 1).to(tl.int64) T = eos - bos else: bos, eos = i_b * T, i_b * T + T @@ -701,27 +742,25 @@ def chunk_kda_fwd_kernel_intra_sub_chunk( Aqk = Aqk + (bos * HV + i_hv) * BT Akk = Akk + (bos * HV + i_hv) * BC - p_q = tl.make_block_ptr(q, (T, K), (H*K, 1), (i_ti, 0), (BC, BK), (1, 0)) - p_k = tl.make_block_ptr(k, (T, K), (H*K, 1), (i_ti, 0), (BC, BK), (1, 0)) - p_g = tl.make_block_ptr(g, (T, K), (HV*K, 1), (i_ti, 0), (BC, BK), (1, 0)) - - p_beta = tl.make_block_ptr(beta, (T,), (HV,), (i_ti,), (BC,), (0,)) + o_k = tl.arange(0, BK) + m_k = o_k < K + m_ck = m_c[:, None] & m_k[None, :] + p_q = q + o_c[:, None] * (H*K) + o_k[None, :] + p_k = k + o_c[:, None] * (H*K) + o_k[None, :] + p_g = g + o_c[:, None] * (HV*K) + o_k[None, :] - b_q = tl.load(p_q, boundary_check=(0, 1)) - b_k = tl.load(p_k, boundary_check=(0, 1)) - b_g = tl.load(p_g, boundary_check=(0, 1)) - b_beta = tl.load(p_beta, boundary_check=(0,)) + p_beta = beta + o_c * HV + b_q = tl.load(p_q, mask=m_ck, other=0.0) + b_k = tl.load(p_k, mask=m_ck, other=0.0) + b_g = tl.load(p_g, mask=m_ck, other=0.0) + b_beta = tl.load(p_beta, mask=m_c, other=0.0) if USE_GATHER: - b_gn = gather(b_g, tl.full([1, BK], min(BC//2, T - i_ti - 1), dtype=tl.int16), axis=0) + b_gn = gather(b_g, tl.zeros([1, BK], dtype=tl.int16), axis=0) else: - # caculate offset - p_gn = g + (i_ti + min(BC // 2, T - i_ti - 1)) * HV*K + tl.arange(0, BK) - b_gn = tl.load(p_gn, mask=tl.arange(0, BK) < K, other=0.0) - b_gn = b_gn[None, :] + b_gn = tl.load(g + i_ti * HV*K + o_k, mask=m_k, other=0.0)[None, :] + b_gn = b_gn - 0.5 * (BC - 1) * GATE_SHIFT - # current block, keep numerical stability by subtracting the left boundary - # less than 85 to avoid overflow in exp2 b_gm = (b_g - b_gn).to(tl.float32) b_gq = tl.where(m_c[:, None], exp2(b_gm), 0.) @@ -740,17 +779,15 @@ def chunk_kda_fwd_kernel_intra_sub_chunk( b_Aqk = tl.where(m_Aqk, b_Aqk, 0.0) b_Akk = tl.where(m_Akk, b_Akk, 0.0) - p_Aqk = tl.make_block_ptr(Aqk, (T, BT), (HV*BT, 1), (i_ti, i_i * BC), (BC, BC), (1, 0)) - p_Akk = tl.make_block_ptr(Akk, (T, BC), (HV*BC, 1), (i_ti, 0), (BC, BC), (1, 0)) - tl.store(p_Aqk, b_Aqk.to(Aqk.dtype.element_ty), boundary_check=(0, 1)) - tl.store(p_Akk, b_Akk.to(Akk.dtype.element_ty), boundary_check=(0, 1)) + m_Aqk_st = m_c[:, None] & (o_i[None, :] < BT) + m_Akk_st = m_c[:, None] & (o_i[None, :] < BC) + p_Aqk = Aqk + o_c[:, None] * (HV*BT) + (i_i * BC + o_i)[None, :] + p_Akk = Akk + o_c[:, None] * (HV*BC) + o_i[None, :] + tl.store(p_Aqk, b_Aqk.to(Aqk.dtype.element_ty), mask=m_Aqk_st) + tl.store(p_Akk, b_Akk.to(Akk.dtype.element_ty), mask=m_Akk_st) tl.debug_barrier() - ################################################################################ - # forward substitution - ################################################################################ - b_Ai = -b_Akk for i in range(2, min(BC, T - i_ti)): b_a = -tl.load(Akk + (i_ti + i) * HV*BC + o_i) @@ -758,9 +795,11 @@ def chunk_kda_fwd_kernel_intra_sub_chunk( b_a += tl.sum(b_a[:, None] * b_Ai, 0) b_Ai = tl.where((o_i == i)[:, None], b_a, b_Ai) b_Ai += m_I - tl.store(p_Akk, b_Ai.to(Akk.dtype.element_ty), boundary_check=(0, 1)) + tl.store(p_Akk, b_Ai.to(Akk.dtype.element_ty), mask=m_Akk_st) +# Use triton only, other backends (cutlass, cute, etc.) need additional torch.compile support +# @dispatch('kda') def chunk_kda_fwd_intra( q: torch.Tensor, k: torch.Tensor, @@ -772,8 +811,46 @@ def chunk_kda_fwd_intra( chunk_size: int = 64, chunk_indices: torch.LongTensor | None = None, safe_gate: bool = False, + lower_bound: float | None = None, disable_recompute: bool = False, + use_graph: bool = False, ): + """Compute intra-chunk factors and solve the triangular KDA system. + + Bounded log gates in safe_gate keep diagonal-block exponentials within FP32 range, allowing Tensor Core matmuls. + A reference C near the middle of the cumulative gates G balances the query/key scaling factors: + + exp(G_i - G_j) = exp(G_i - C) * exp(C - G_j) + + v0.5.2 uses the middle token's gate, which reads future gates for earlier causal scores (x): + + key token + 0 1 2 3 + query token 0 x . . . + 1 x x . . + 2 o o o . <- C = G_2 + 3 o o o o + + Dots are masked; x and o are retained. C cancels in exact arithmetic, but rounding can leak future gates into x. + See https://drisspg.github.io/nuggets/KDA-Future-Token-Leakage + + Instead, use the first cumulative gate with a fixed shift toward an estimated midpoint (natural-log units): + + s = -lower_bound if lower_bound is supplied to intra, otherwise 2 + d = (BC - 1) * s / 2; C = G_0 - d + + For BC=16 and s=2, C = G_0 - 15. Including the RCP_LN2 conversion, the kernel's exp2 gives + exp(G_i - G_0 + 15) for queries and exp(G_0 - G_i - 15) for keys. + The default 2 is a heuristic for log gates concentrated near 0, with few near the lower bound; + it limits the compensation to avoid over-scaling queries. + chunk_kda uses this default for both fused and pre-computed gates. An explicit intra lower_bound sets s instead. + Forward/backward keep the same fixed BC even for partial blocks, so appending tokens cannot move C. + + For actual log gates in [L, 0], the factor bounds are (BC=16, s=2, L=-5): + + max_i exp(G_i - C) <= exp(d) = exp(15 * 2 / 2) = exp(15) ~= 3.269e6 + max_i exp(C - G_i) <= exp(-(BC - 1) * L - d) = exp(-15 * (-5) - 15) = exp(60) ~= 1.142e26 + """ B, T, H, K, HV = *k.shape, gk.shape[2] BT = chunk_size if BT not in (32, 64): @@ -790,8 +867,6 @@ def chunk_kda_fwd_intra( # Separate fp32 buffer for diagonal 16x16 blocks (for precision in solve_tril) Akkd = torch.empty(B, T, HV, BC, device=k.device, dtype=torch.float32) - # Step 1: Run token_parallel first to compute diagonal blocks into Akkd (fp32) - # Step 1: compute diagonal blocks into Akk_diag (fp32) if safe_gate: grid = (NT, NC, B * HV) BK = triton.next_power_of_2(K) @@ -812,7 +887,9 @@ def chunk_kda_fwd_intra( BT=BT, BC=BC, BK=BK, + GATE_SHIFT=(-lower_bound if lower_bound is not None else 2.) * RCP_LN2, USE_GATHER=IS_GATHER_SUPPORTED, + USE_GRAPH=use_graph, ) else: Aqk, Akkd = chunk_kda_fwd_intra_token_parallel( @@ -826,6 +903,7 @@ def chunk_kda_fwd_intra( cu_seqlens=cu_seqlens, chunk_size=BT, sub_chunk_size=BC, + use_graph=use_graph, ) # Step 2: Fused inter + solve_tril (works for both fixed-len and varlen) @@ -849,6 +927,7 @@ def chunk_kda_fwd_intra( BC=BC, NC=NC, USE_SAFE_GATE=safe_gate, + USE_GRAPH=use_graph, ) w, u, qg, kg = recompute_w_u_fwd( k=k, @@ -859,10 +938,12 @@ def chunk_kda_fwd_intra( gk=gk, cu_seqlens=cu_seqlens, chunk_indices=chunk_indices, + use_graph=use_graph, ) return w, u, qg, kg, Aqk, Akk +# @dispatch('kda') def chunk_kda_bwd_intra( q: torch.Tensor, k: torch.Tensor, @@ -878,6 +959,8 @@ def chunk_kda_bwd_intra( chunk_indices: torch.LongTensor | None = None, chunk_size: int = 64, safe_gate: bool = False, + lower_bound: float | None = None, + use_graph: bool = False, ): B, T, H, K, HV = *k.shape, g.shape[2] BT = chunk_size @@ -890,11 +973,17 @@ def chunk_kda_bwd_intra( NC = triton.cdiv(BT, BC) NK = triton.cdiv(K, BK) - dq2 = torch.empty_like(dq) - dk2 = torch.empty_like(dk) - db2 = beta.new_empty(NK, *beta.shape, dtype=torch.float) - dg2 = torch.empty_like(dg, dtype=torch.float) - grid = (NK * NC, NT, B * HV) + if use_graph: + dq2 = get_static_buffer("intra_dq2", tuple(dq.shape), dq.dtype, dq.device) + dk2 = get_static_buffer("intra_dk2", tuple(dk.shape), dk.dtype, dk.device) + db2 = get_static_buffer("intra_db2", (NK, *beta.shape), torch.float, beta.device) + dg2 = get_static_buffer("intra_dg2", tuple(dg.shape), torch.float, dg.device) + else: + dq2 = torch.empty_like(dq) + dk2 = torch.empty_like(dk) + db2 = beta.new_empty(NK, *beta.shape, dtype=torch.float) + dg2 = torch.empty_like(dg, dtype=torch.float) + grid = (NK * NC * NT, B * HV) chunk_kda_bwd_kernel_intra[grid]( q=q, k=k, @@ -921,7 +1010,9 @@ def chunk_kda_bwd_intra( BK=BK, NC=NC, SAFE_GATE=safe_gate, + GATE_SHIFT=(-lower_bound if lower_bound is not None else 2.) * RCP_LN2, USE_GATHER=IS_GATHER_SUPPORTED, + USE_GRAPH=use_graph, ) dq = dq2 dk = dk2 diff --git a/build/torch-cuda/ops/kda/chunk_intra_token_parallel.py b/build/torch-cuda/ops/kda/chunk_intra_token_parallel.py index d1ce022245551716e6477c06fd011d8d3ed4c21e..45deeae9163fa4674ae812a3cfe2b2549c7b3222 100644 --- a/build/torch-cuda/ops/kda/chunk_intra_token_parallel.py +++ b/build/torch-cuda/ops/kda/chunk_intra_token_parallel.py @@ -11,6 +11,7 @@ import torch import triton import triton.language as tl +from ...backends import dispatch from ...ops.utils.cache import fla_cache_autotune from ...ops.utils.op import exp2 from ...utils import autotune_cache_kwargs @@ -43,14 +44,19 @@ def chunk_kda_fwd_kernel_intra_token_parallel( H: tl.constexpr, HV: tl.constexpr, K: tl.constexpr, + BK: tl.constexpr, BT: tl.constexpr, BC: tl.constexpr, BH: tl.constexpr, IS_VARLEN: tl.constexpr, + USE_GRAPH: tl.constexpr = False, ): - i_tg, i_hg = tl.program_id(0), tl.program_id(1) + i_tg, i_hg = tl.program_id(0).to(tl.int64), tl.program_id(1).to(tl.int64) if IS_VARLEN: + # static T may exceed the covered tokens; i_n would converge to N and read OOB + if USE_GRAPH and i_tg >= tl.load(cu_seqlens + N).to(tl.int32): + return i_n = 0 left, right = 0, N @@ -66,7 +72,7 @@ def chunk_kda_fwd_kernel_intra_token_parallel( left = mid + 1 i_n = left - bos, eos = tl.load(cu_seqlens + i_n).to(tl.int32), tl.load(cu_seqlens + i_n + 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 i_t = i_tg - bos else: @@ -90,7 +96,6 @@ def chunk_kda_fwd_kernel_intra_token_parallel( Akk += bos * HV*BC beta += bos * HV - BK: tl.constexpr = triton.next_power_of_2(K) o_hv = i_hg * BH + tl.arange(0, BH) o_h = o_hv // G o_k = tl.arange(0, BK) @@ -104,15 +109,15 @@ def chunk_kda_fwd_kernel_intra_token_parallel( b_k = tl.load(k + i_t * H * K + p_qk, mask=m_hk, other=0).to(tl.float32) # g: [B, T, HV, K], beta: [B, T, HV] - p_g = tl.make_block_ptr(g + i_t * HV * K, (HV, K), (K, 1), (i_hg * BH, 0), (BH, BK), (1, 0)) - p_beta = tl.make_block_ptr(beta + i_t * HV, (HV,), (1,), (i_hg * BH,), (BH,), (0,)) - b_g = tl.load(p_g, boundary_check=(0, 1)).to(tl.float32) - b_k = b_k * tl.load(p_beta, boundary_check=(0,)).to(tl.float32)[:, None] + p_g = g + i_t * HV * K + o_hv[:, None] * K + o_k[None, :] + p_beta = beta + i_t * HV + o_hv + b_g = tl.load(p_g, mask=m_hk, other=0.0).to(tl.float32) + b_k = b_k * tl.load(p_beta, mask=m_hv, other=0.0).to(tl.float32)[:, None] for j in range(i_ts, min(i_t + 1, min(T, i_ts + BC))): b_kj = tl.load(k + j * H * K + p_qk, mask=m_hk, other=0).to(tl.float32) - p_gj = tl.make_block_ptr(g + j * HV * K, (HV, K), (K, 1), (i_hg * BH, 0), (BH, BK), (1, 0)) - b_gj = tl.load(p_gj, boundary_check=(0, 1)).to(tl.float32) + p_gj = g + j * HV * K + o_hv[:, None] * K + o_k[None, :] + b_gj = tl.load(p_gj, mask=m_hk, other=0.0).to(tl.float32) b_kgj = tl.where(m_k[None, :], b_kj * exp2(b_g - b_gj), 0.0) b_Aqk = tl.sum(b_q * b_kgj, axis=1) * scale @@ -122,6 +127,8 @@ def chunk_kda_fwd_kernel_intra_token_parallel( tl.store(Akk + i_t * HV * BC + o_hv * BC + j - i_ts, b_Akk.to(Akk.dtype.element_ty), mask=m_hv) +# Use triton only, other backends (cutlass, cute, etc.) need additional torch.compile support +# @dispatch('kda') def chunk_kda_fwd_intra_token_parallel( q: torch.Tensor, k: torch.Tensor, @@ -133,6 +140,7 @@ def chunk_kda_fwd_intra_token_parallel( cu_seqlens: torch.LongTensor | None = None, chunk_size: int = 64, sub_chunk_size: int = 16, + use_graph: bool = False, ) -> None: """ Token-parallel implementation: each token gets its own thread block. @@ -156,6 +164,7 @@ def chunk_kda_fwd_intra_token_parallel( N = len(cu_seqlens) - 1 if cu_seqlens is not None else B BT = chunk_size BC = sub_chunk_size + BK = triton.next_power_of_2(K) def grid(meta): return (B * T, triton.cdiv(HV, meta['BH'])) chunk_kda_fwd_kernel_intra_token_parallel[grid]( @@ -172,7 +181,9 @@ def chunk_kda_fwd_intra_token_parallel( H=H, HV=HV, K=K, + BK=BK, BT=BT, BC=BC, + USE_GRAPH=use_graph, ) return Aqk, Akk diff --git a/build/torch-cuda/ops/kda/fused_recurrent.py b/build/torch-cuda/ops/kda/fused_recurrent.py index 285099dea8bc4c10544bd7277b35c4b328fb5193..bbfd407923c7d295263f75e35cd87108b9e75f8e 100644 --- a/build/torch-cuda/ops/kda/fused_recurrent.py +++ b/build/torch-cuda/ops/kda/fused_recurrent.py @@ -13,6 +13,7 @@ import torch import triton import triton.language as tl +from ...backends import dispatch from ...ops.utils.op import exp from ...ops.utils.softplus import softplus from ...utils import input_guard @@ -25,7 +26,8 @@ from ...utils import input_guard "IS_VARLEN": lambda args: args["cu_seqlens"] is not None, "IS_CONTINUOUS_BATCHING": lambda args: args["ssm_state_indices"] is not None, "IS_SPEC_DECODING": lambda args: args["num_accepted_tokens"] is not None, - "HAS_DT_BIAS": lambda args: args["dt_bias"] is not None, + "HAS_A": lambda args: args["A_log"] is not None, + "HAS_BIAS": lambda args: args["dt_bias"] is not None, "USE_LOWER_BOUND": lambda args: args["lower_bound"] is not None, } ) @@ -66,7 +68,8 @@ def fused_recurrent_kda_fwd_kernel( IS_CONTINUOUS_BATCHING: tl.constexpr, IS_SPEC_DECODING: tl.constexpr, STORE_FINAL_STATE: tl.constexpr, - HAS_DT_BIAS: tl.constexpr, + HAS_A: tl.constexpr, + HAS_BIAS: tl.constexpr, USE_GATE_IN_KERNEL: tl.constexpr, USE_LOWER_BOUND: tl.constexpr, APPLY_BETA_SIGMOID: tl.constexpr, @@ -74,7 +77,7 @@ def fused_recurrent_kda_fwd_kernel( STATE_V_FIRST: tl.constexpr, num_stages: tl.constexpr, ): - pid = tl.program_id(0) + pid = tl.program_id(0).to(tl.int64) NV = tl.cdiv(V, BV) NK = tl.cdiv(K, BK) i_k = pid % NK @@ -155,17 +158,17 @@ def fused_recurrent_kda_fwd_kernel( b_q = b_q / tl.sqrt(tl.sum(b_q * b_q) + 1e-6) b_k = b_k / tl.sqrt(tl.sum(b_k * b_k) + 1e-6) b_q = b_q * scale - b_g = tl.load(p_g, eviction_policy='evict_last').to(tl.float32) + b_g = tl.load(p_g, mask=mask_k, other=0, eviction_policy='evict_last').to(tl.float32) if USE_GATE_IN_KERNEL: - b_A = tl.load(A_log + i_hv).to(tl.float32) + b_A = tl.load(A_log + i_hv).to(tl.float32) if HAS_A else 1.0 - if HAS_DT_BIAS: + if HAS_BIAS: b_bias = tl.load(dt_bias + i_hv * K + o_k, mask=mask_k, other=0).to(tl.float32) b_g = b_g + b_bias if USE_LOWER_BOUND: - b_gk = lower_bound * tl.sigmoid(exp(b_A) * b_g) + b_gk = lower_bound * tl.sigmoid((exp(b_A) if HAS_A else b_A) * b_g) else: b_gk = -exp(b_A) * softplus(b_g) else: @@ -230,7 +233,8 @@ def fused_recurrent_kda_fwd_kernel( tl.store(p_ht, b_h.to(p_ht.dtype.element_ty), mask=mask_h) -@torch.compiler.disable +# Use triton only, other backends (cutlass, cute, etc.) need additional torch.compile support +# @dispatch('kda') def fused_recurrent_kda_fwd( q: torch.Tensor, k: torch.Tensor, @@ -367,10 +371,11 @@ def fused_recurrent_kda( beta (torch.Tensor): betas of shape `[B, T, HV]`. A_log (Optional[torch.Tensor]): - Decay parameter of shape `[HV]`. Required when `use_gate_in_kernel=True`. + Decay parameter of shape `[HV]`. + When `use_gate_in_kernel=True` together with `lower_bound`, + may be `None` to use `lower_bound * sigmoid(g + dt_bias)`. dt_bias (Optional[torch.Tensor]): - Bias added to `g` before activation, of shape `[HV]`. - Only used when `use_gate_in_kernel=True`. + Bias added to `g` before activation, of shape `[HV]`. Only used when `use_gate_in_kernel=True`. scale (Optional[float]): Scale factor for the RetNet attention scores. If not provided, it will default to `1 / sqrt(K)`. Default: `None`. @@ -384,8 +389,8 @@ def fused_recurrent_kda( Whether to use L2 normalization in the kernel. Default: `False`. use_gate_in_kernel (Optional[bool]): Whether to compute the log-space KDA decay internally. - When `True`, `g` is the raw input and `A_log` must be provided; the kernel fuses - gate activation into the recurrence. Default: `False`. + When `True`, `g` is the raw input and the kernel fuses gate activation into the recurrence. + Default: `False`. use_beta_sigmoid_in_kernel (Optional[bool]): Whether to apply `torch.sigmoid(beta)` inside the kernel. - If `True`, the passed `beta` acts as the raw beta logits. @@ -413,7 +418,7 @@ def fused_recurrent_kda( >>> import torch >>> import torch.nn.functional as F >>> from einops import rearrange - >>> from ...ops.kda import fused_recurrent_kda + >>> from fla.ops.kda import fused_recurrent_kda # inputs with equal lengths >>> B, T, H, HV, K, V = 4, 2048, 4, 8, 512, 512 >>> q = torch.randn(B, T, H, K, device='cuda') diff --git a/build/torch-cuda/ops/kda/gate.py b/build/torch-cuda/ops/kda/gate.py index bf38ddb6ee9ffa2d50ecbe1dd8699a42d561ded2..d2716fad5bf1f82f050643fdced5ea803d9bf5cd 100644 --- a/build/torch-cuda/ops/kda/gate.py +++ b/build/torch-cuda/ops/kda/gate.py @@ -12,9 +12,10 @@ import torch.nn.functional as F import triton import triton.language as tl +from ...backends import dispatch from ...ops.utils.cache import fla_cache_autotune from ...ops.utils.index import prepare_chunk_indices -from ...ops.utils.op import exp +from ...ops.utils.op import exp, unflatten_program_id from ...ops.utils.softplus import softplus from ...utils import IS_AMD, autocast_custom_bwd, autocast_custom_fwd, autotune_cache_kwargs, check_shared_mem, input_guard @@ -56,20 +57,44 @@ def naive_kda_gate( def naive_kda_lowerbound_gate( g: torch.Tensor, - A_log: torch.Tensor, + A_log: torch.Tensor | None = None, dt_bias: torch.Tensor | None = None, lower_bound: float = -5.0, output_dtype: torch.dtype = torch.float32, ) -> torch.Tensor: + """ + Torch reference implementation for KDA lowerbound gate computation. + + Computes: ``g = lower_bound * sigmoid(exp(A_log) * (g + dt_bias))``. + When ``A_log`` is ``None``: ``g = lower_bound * sigmoid(g + dt_bias)``. + + Args: + g (torch.Tensor): + Input tensor of shape `[..., H, K]`. + A_log (torch.Tensor | None): + Optional parameter tensor with `H` elements. + dt_bias (torch.Tensor | None): + Optional bias tensor added to `g` before activation, shape `[H * K]`. + lower_bound (float): + Lower bound for the gate output. Default: `-5.0`. + output_dtype (torch.dtype): + The dtype of the output tensor. Default: `torch.float32`. + + Returns: + Output tensor of shape `[..., H, K]`. + """ H, _ = g.shape[-2:] g = g.float() if dt_bias is not None: g = g + dt_bias.view(H, -1) - g = lower_bound * F.sigmoid(A_log.view(H, 1).exp() * g) + if A_log is not None: + g = A_log.view(H, 1).float().exp() * g + g = lower_bound * F.sigmoid(g) return g.to(output_dtype) @triton.heuristics({ + "HAS_A": lambda args: args["A_log"] is not None, "HAS_BIAS": lambda args: args["dt_bias"] is not None, "HAS_BETA": lambda args: args["beta"] is not None, 'USE_LOWER_BOUND': lambda args: args['lower_bound'] is not None, @@ -98,35 +123,41 @@ def kda_gate_fwd_kernel( D: tl.constexpr, BT: tl.constexpr, BD: tl.constexpr, + HAS_A: tl.constexpr, HAS_BIAS: tl.constexpr, HAS_BETA: tl.constexpr, USE_LOWER_BOUND: tl.constexpr, ): - i_t, i_h = tl.program_id(0), tl.program_id(1) + i_t, i_h = tl.program_id(0).to(tl.int64), tl.program_id(1).to(tl.int64) - b_A = tl.load(A_log + i_h).to(tl.float32) + b_A = tl.load(A_log + i_h).to(tl.float32) if HAS_A else 1.0 - p_g = tl.make_block_ptr(g + i_h * D, (T, D), (H * D, 1), (i_t * BT, 0), (BT, BD), (1, 0)) - p_yg = tl.make_block_ptr(yg + i_h * D, (T, D), (H * D, 1), (i_t * BT, 0), (BT, BD), (1, 0)) + o_t = i_t * BT + tl.arange(0, BT) + o_d = tl.arange(0, BD) + m_t = o_t < T + m_g = m_t[:, None] & (o_d[None, :] < D) + p_g = g + i_h * D + o_t[:, None] * (H * D) + o_d[None, :] + p_yg = yg + i_h * D + o_t[:, None] * (H * D) + o_d[None, :] # [BT, BD] - b_g = tl.load(p_g, boundary_check=(0, 1)).to(tl.float32) + b_g = tl.load(p_g, mask=m_g, other=0.0).to(tl.float32) if HAS_BIAS: - p_b = tl.make_block_ptr(dt_bias, (H * D,), (1,), (i_h * D,), (BD,), (0,)) - b_g = b_g + tl.load(p_b, boundary_check=(0,)).to(tl.float32) + o_b = i_h * D + tl.arange(0, BD) + b_g = b_g + tl.load(dt_bias + o_b, mask=o_b < H * D, other=0.0).to(tl.float32) if not USE_LOWER_BOUND: b_yg = -exp(b_A) * softplus(b_g) else: - b_yg = lower_bound * tl.sigmoid(exp(b_A) * b_g) - tl.store(p_yg, b_yg.to(p_yg.dtype.element_ty), boundary_check=(0, 1)) + b_yg = lower_bound * tl.sigmoid((exp(b_A) if HAS_A else b_A) * b_g) + tl.store(p_yg, b_yg.to(p_yg.dtype.element_ty), mask=m_g) if HAS_BETA: - p_b = tl.make_block_ptr(beta + i_h, (T,), (H,), (i_t * BT,), (BT,), (0,)) - p_yb = tl.make_block_ptr(yb + i_h, (T,), (H,), (i_t * BT,), (BT,), (0,)) - b_yb = tl.sigmoid(tl.load(p_b, boundary_check=(0,)).to(tl.float32)) - tl.store(p_yb, b_yb.to(p_yb.dtype.element_ty), boundary_check=(0,)) + p_b = beta + i_h + o_t * H + p_yb = yb + i_h + o_t * H + b_yb = tl.sigmoid(tl.load(p_b, mask=m_t, other=0.0).to(tl.float32)) + tl.store(p_yb, b_yb.to(p_yb.dtype.element_ty), mask=m_t) @triton.heuristics({ + "HAS_A": lambda args: args["A_log"] is not None, "HAS_BIAS": lambda args: args["dt_bias"] is not None, "HAS_BETA": lambda args: args["beta"] is not None, 'USE_LOWER_BOUND': lambda args: args['lower_bound'] is not None, @@ -150,6 +181,7 @@ def kda_gate_bwd_kernel( dyb, dg, dA, + dbias, dbeta, lower_bound, T, @@ -157,25 +189,30 @@ def kda_gate_bwd_kernel( D: tl.constexpr, BT: tl.constexpr, BD: tl.constexpr, + HAS_A: tl.constexpr, HAS_BIAS: tl.constexpr, HAS_BETA: tl.constexpr, USE_LOWER_BOUND: tl.constexpr, ): - i_t, i_h = tl.program_id(0), tl.program_id(1) + i_t, i_h = tl.program_id(0).to(tl.int64), tl.program_id(1).to(tl.int64) - b_A = tl.load(A_log + i_h).to(tl.float32) + b_A = tl.load(A_log + i_h).to(tl.float32) if HAS_A else 1.0 - p_g = tl.make_block_ptr(g + i_h * D, (T, D), (H * D, 1), (i_t * BT, 0), (BT, BD), (1, 0)) - p_dg = tl.make_block_ptr(dg + i_h * D, (T, D), (H * D, 1), (i_t * BT, 0), (BT, BD), (1, 0)) - p_dyg = tl.make_block_ptr(dyg + i_h * D, (T, D), (H * D, 1), (i_t * BT, 0), (BT, BD), (1, 0)) + o_t = i_t * BT + tl.arange(0, BT) + o_d = tl.arange(0, BD) + m_t = o_t < T + m_g = m_t[:, None] & (o_d[None, :] < D) + p_g = g + i_h * D + o_t[:, None] * (H * D) + o_d[None, :] + p_dg = dg + i_h * D + o_t[:, None] * (H * D) + o_d[None, :] + p_dyg = dyg + i_h * D + o_t[:, None] * (H * D) + o_d[None, :] # [BT, BD] - b_g = tl.load(p_g, boundary_check=(0, 1)).to(tl.float32) - b_dyg = tl.load(p_dyg, boundary_check=(0, 1)).to(tl.float32) + b_g = tl.load(p_g, mask=m_g, other=0.0).to(tl.float32) + b_dyg = tl.load(p_dyg, mask=m_g, other=0.0).to(tl.float32) if HAS_BIAS: - p_b = tl.make_block_ptr(dt_bias, (H * D,), (1,), (i_h * D,), (BD,), (0,)) - b_g = b_g + tl.load(p_b, boundary_check=(0,)).to(tl.float32) + o_b = i_h * D + tl.arange(0, BD) + b_g = b_g + tl.load(dt_bias + o_b, mask=o_b < H * D, other=0.0).to(tl.float32) # [BT, BD] if not USE_LOWER_BOUND: @@ -184,7 +221,7 @@ def kda_gate_bwd_kernel( b_dg = b_A * (b_dyg * tl.sigmoid(b_g)) b_dA = tl.sum(tl.sum(b_dyg * b_yg, 1), 0) else: - b_A = exp(b_A) + b_A = exp(b_A) if HAS_A else b_A b_inner = b_A * b_g b_sig = tl.sigmoid(b_inner) b_dsig = b_sig * (1.0 - b_sig) @@ -192,24 +229,29 @@ def kda_gate_bwd_kernel( b_d_inner_term = b_dyg * (lower_bound * b_dsig) # dg = d_inner_term * A b_dg = b_d_inner_term * b_A - b_dA = tl.sum(tl.sum(b_dg * b_g, 1), 0) + b_dA = tl.sum(tl.sum(b_dg * b_g, 1), 0) if HAS_A else 0.0 - tl.store(p_dg, b_dg.to(p_dg.dtype.element_ty), boundary_check=(0, 1)) - tl.store(dA + i_t * H + i_h, b_dA) + tl.store(p_dg, b_dg.to(p_dg.dtype.element_ty), mask=m_g) + if HAS_A: + tl.store(dA + i_t * H + i_h, b_dA) + if HAS_BIAS: + b_dbias = tl.sum(b_dg, 0) + tl.store(dbias + (i_t * H + i_h) * D + o_d, b_dbias, mask=o_d < D) if HAS_BETA: - p_b = tl.make_block_ptr(beta + i_h, (T,), (H,), (i_t * BT,), (BT,), (0,)) - p_db = tl.make_block_ptr(dbeta + i_h, (T,), (H,), (i_t * BT,), (BT,), (0,)) - p_dyb = tl.make_block_ptr(dyb + i_h, (T,), (H,), (i_t * BT,), (BT,), (0,)) + p_b = beta + i_h + o_t * H + p_db = dbeta + i_h + o_t * H + p_dyb = dyb + i_h + o_t * H - b_b = tl.load(p_b, boundary_check=(0,)).to(tl.float32) - b_db = tl.load(p_dyb, boundary_check=(0,)).to(tl.float32) * b_b * (1.0 - b_b) - tl.store(p_db, b_db.to(p_db.dtype.element_ty), boundary_check=(0,)) + b_b = tl.load(p_b, mask=m_t, other=0.0).to(tl.float32) + b_db = tl.load(p_dyb, mask=m_t, other=0.0).to(tl.float32) * b_b * (1.0 - b_b) + tl.store(p_db, b_db.to(p_db.dtype.element_ty), mask=m_t) +@dispatch('kda') def kda_gate_fwd( g: torch.Tensor, - A_log: torch.Tensor, + A_log: torch.Tensor | None = None, dt_bias: torch.Tensor | None = None, lower_bound: float | None = None, output_dtype: torch.dtype = torch.float32, @@ -238,23 +280,24 @@ def kda_gate_fwd( return yg +@dispatch('kda') def kda_gate_bwd( g: torch.Tensor, - A_log: torch.Tensor, + A_log: torch.Tensor | None = None, dt_bias: torch.Tensor | None = None, dyg: torch.Tensor | None = None, lower_bound: float | None = None, -) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor | None]: +) -> tuple[torch.Tensor, torch.Tensor | None, torch.Tensor | None]: H, K = g.shape[-2:] T = g.numel() // (H * K) BT = 32 NT = triton.cdiv(T, BT) - dg = torch.empty_like(g, dtype=torch.float32) - dA = A_log.new_empty(NT, H, dtype=torch.float32) + dg = torch.empty_like(g) + dA = g.new_empty(NT, H, dtype=torch.float32) if A_log is not None else None + dbias = g.new_empty(NT, H * K, dtype=torch.float32) if dt_bias is not None else None - grid = (triton.cdiv(T, BT), H) - kda_gate_bwd_kernel[grid]( + kda_gate_bwd_kernel[(NT, H)]( g=g, A_log=A_log, dt_bias=dt_bias, @@ -263,6 +306,7 @@ def kda_gate_bwd( dyb=None, dg=dg, dA=dA, + dbias=dbias, dbeta=None, T=T, H=H, @@ -272,9 +316,8 @@ def kda_gate_bwd( lower_bound=lower_bound, ) - dg = dg.view_as(g).type_as(g) - dA = dA.sum(0).view_as(A_log).type_as(A_log) - dbias = dg.view(-1, H * K).sum(0).to(dt_bias) if dt_bias is not None else None + dA = dA.sum(0).view_as(A_log).type_as(A_log) if A_log is not None else None + dbias = dbias.sum(0).to(dt_bias) if dt_bias is not None else None return dg, dA, dbias @@ -286,7 +329,7 @@ class KDAGateFunction(torch.autograd.Function): def forward( ctx, g: torch.Tensor, - A_log: torch.Tensor, + A_log: torch.Tensor | None = None, dt_bias: torch.Tensor | None = None, lower_bound: float | None = None, output_dtype: torch.dtype = torch.float32, @@ -318,9 +361,10 @@ class KDAGateFunction(torch.autograd.Function): @torch.compiler.disable +@dispatch('kda') def fused_kda_gate( g: torch.Tensor, - A_log: torch.Tensor, + A_log: torch.Tensor | None = None, dt_bias: torch.Tensor | None = None, lower_bound: float | None = None, output_dtype: torch.dtype = torch.float32, @@ -329,12 +373,15 @@ def fused_kda_gate( Fused KDA gate computation with autograd support. Computes: g = -A_log.exp().unsqueeze(-1) * softplus(g + dt_bias.view(g.shape[-2:])) + When ``lower_bound`` is set: g = lower_bound * sigmoid(exp(A_log) * (g + dt_bias)). + When ``A_log`` is ``None`` (requires ``lower_bound``): g = lower_bound * sigmoid(g + dt_bias). Args: g (torch.Tensor): Input tensor of shape `[..., H, K]`. - A_log (torch.Tensor): - Parameter tensor with `H` elements. + A_log (torch.Tensor | None): + Optional parameter tensor with `H` elements. + When ``None``, the gate reduces to ``lower_bound * sigmoid(g + dt_bias)`` (requires ``lower_bound``). dt_bias (torch.Tensor | None): Optional bias tensor added to `g` before activation, shape `[H * K]`. @@ -345,6 +392,7 @@ def fused_kda_gate( @triton.heuristics({ + "HAS_A": lambda args: args["A_log"] is not None, "HAS_BIAS": lambda args: args["dt_bias"] is not None, 'HAS_SCALE': lambda args: args['scale'] is not None, 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None, @@ -375,37 +423,45 @@ def kda_gate_chunk_cumsum_vector_kernel( BT: tl.constexpr, BS: tl.constexpr, REVERSE: tl.constexpr, + HAS_A: tl.constexpr, HAS_BIAS: tl.constexpr, HAS_SCALE: tl.constexpr, IS_VARLEN: tl.constexpr, USE_LOWER_BOUND: tl.constexpr, + USE_GRAPH: tl.constexpr = False, ): - i_s, i_t, i_bh = tl.program_id(0), tl.program_id(1), tl.program_id(2) + i_s, i_t = unflatten_program_id(tl.cdiv(S, BS)) + i_bh = tl.program_id(1).to(tl.int64) i_b, i_h = i_bh // H, i_bh % H 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.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) + i_n = tl.load(chunk_indices + i_t * 2).to(tl.int32) + if USE_GRAPH and i_n < 0: + return + i_t = tl.load(chunk_indices + i_t * 2 + 1).to(tl.int64) + bos, eos = tl.load(cu_seqlens + i_n).to(tl.int64), tl.load(cu_seqlens + i_n + 1).to(tl.int64) T = eos - bos else: bos, eos = i_b * T, i_b * T + T - p_s = tl.make_block_ptr(s + (bos * H + i_h) * S, (T, S), (H*S, 1), (i_t * BT, i_s * BS), (BT, BS), (1, 0)) - p_o = tl.make_block_ptr(o + (bos * H + i_h) * S, (T, S), (H*S, 1), (i_t * BT, i_s * BS), (BT, BS), (1, 0)) + o_t = i_t * BT + tl.arange(0, BT) + o_s = i_s * BS + tl.arange(0, BS) + m_s = (o_t[:, None] < T) & (o_s[None, :] < S) + p_s = s + (bos * H + i_h) * S + o_t[:, None] * (H*S) + o_s[None, :] + p_o = o + (bos * H + i_h) * S + o_t[:, None] * (H*S) + o_s[None, :] # [BT, BS] - b_s = tl.load(p_s, boundary_check=(0, 1)).to(tl.float32) + b_s = tl.load(p_s, mask=m_s, other=0.0).to(tl.float32) # Apply dt_bias if exists if HAS_BIAS: - p_b = tl.make_block_ptr(dt_bias + i_h * S, (S,), (1,), (i_s * BS,), (BS,), (0,)) - b_bias = tl.load(p_b, boundary_check=(0,)).to(tl.float32) + b_bias = tl.load(dt_bias + i_h * S + o_s, mask=o_s < S, other=0.0).to(tl.float32) b_s = b_s + b_bias[None, :] - b_A = tl.load(A_log + i_h).to(tl.float32) + b_A = tl.load(A_log + i_h).to(tl.float32) if HAS_A else 1.0 if not USE_LOWER_BOUND: # Apply gate: -exp(A_log) * softplus(g + bias) b_gate = -exp(b_A) * softplus(b_s) else: - b_gate = lower_bound * tl.sigmoid(exp(b_A) * b_s) + b_gate = lower_bound * tl.sigmoid((exp(b_A) if HAS_A else b_A) * b_s) # Apply chunk local cumsum if REVERSE: @@ -415,13 +471,14 @@ def kda_gate_chunk_cumsum_vector_kernel( if HAS_SCALE: b_o *= scale - tl.store(p_o, b_o.to(p_o.dtype.element_ty), boundary_check=(0, 1)) + tl.store(p_o, b_o.to(p_o.dtype.element_ty), mask=m_s) @input_guard +@dispatch('kda') def kda_gate_chunk_cumsum( g: torch.Tensor, - A_log: torch.Tensor, + A_log: torch.Tensor | None, chunk_size: int, scale: float = None, dt_bias: torch.Tensor | None = None, @@ -429,6 +486,7 @@ def kda_gate_chunk_cumsum( output_dtype: torch.dtype | None = torch.float, chunk_indices: torch.LongTensor | None = None, lower_bound: float | None = None, + use_graph: bool = False, **kwargs, ) -> torch.Tensor: if cu_seqlens is not None: @@ -441,13 +499,16 @@ def kda_gate_chunk_cumsum( NT = triton.cdiv(T, BT) if cu_seqlens is None else len(chunk_indices) assert chunk_size == 2**(chunk_size.bit_length()-1), "chunk_size must be a power of 2" - g_org, g = g, torch.empty_like(g, dtype=output_dtype or g.dtype) - def grid(meta): return (triton.cdiv(meta['S'], meta['BS']), NT, B * H) + o = torch.empty_like(g, dtype=output_dtype or g.dtype) + + def grid(meta): + return (triton.cdiv(meta['S'], meta['BS']) * NT, B * H) + kda_gate_chunk_cumsum_vector_kernel[grid]( - s=g_org, + s=g, A_log=A_log, dt_bias=dt_bias, - o=g, + o=o, scale=scale, cu_seqlens=cu_seqlens, chunk_indices=chunk_indices, @@ -457,5 +518,6 @@ def kda_gate_chunk_cumsum( S=S, BT=BT, REVERSE=False, + USE_GRAPH=use_graph, ) - return g + return o diff --git a/build/torch-cuda/ops/kda/wy_fast.py b/build/torch-cuda/ops/kda/wy_fast.py index 521e5784e0c724addbd65fcc8d46cb4d8792d9bd..f30ce1552b6f03037fd6a1f474ca7e663a938c9c 100644 --- a/build/torch-cuda/ops/kda/wy_fast.py +++ b/build/torch-cuda/ops/kda/wy_fast.py @@ -9,6 +9,7 @@ import torch import triton import triton.language as tl +from ...backends import dispatch from ...ops.utils import prepare_chunk_indices from ...ops.utils.cache import fla_cache_autotune from ...ops.utils.op import exp2 @@ -54,13 +55,17 @@ def recompute_w_u_fwd_kda_kernel( STORE_QG: tl.constexpr, STORE_KG: tl.constexpr, IS_VARLEN: tl.constexpr, + USE_GRAPH: tl.constexpr = False, ): - i_t, i_bh = tl.program_id(0), tl.program_id(1) + i_t, i_bh = tl.program_id(0).to(tl.int64), tl.program_id(1).to(tl.int64) i_b, i_hv = i_bh // HV, i_bh % HV i_h = i_hv // (HV // H) 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.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) + i_n = tl.load(chunk_indices + i_t * 2).to(tl.int32) + if USE_GRAPH and i_n < 0: + return + i_t = tl.load(chunk_indices + i_t * 2 + 1).to(tl.int64) + bos, eos = tl.load(cu_seqlens + i_n).to(tl.int64), tl.load(cu_seqlens + i_n + 1).to(tl.int64) T = eos - bos else: bos, eos = i_b * T, i_b * T + T @@ -78,46 +83,53 @@ def recompute_w_u_fwd_kda_kernel( if STORE_KG: kg += (bos * HV + i_hv) * K - p_b = tl.make_block_ptr(beta, (T,), (HV,), (i_t * BT,), (BT,), (0,)) - b_b = tl.load(p_b, boundary_check=(0,)) + o_t = i_t * BT + tl.arange(0, BT) + m_t = o_t < T + p_b = beta + o_t * HV + b_b = tl.load(p_b, mask=m_t, other=0.0) - p_A = tl.make_block_ptr(A, (T, BT), (HV*BT, 1), (i_t * BT, 0), (BT, BT), (1, 0)) - b_A = tl.load(p_A, boundary_check=(0, 1)) + o_A = tl.arange(0, BT) + m_A = m_t[:, None] & (o_A[None, :] < BT) + p_A = A + o_t[:, None] * (HV*BT) + o_A[None, :] + b_A = tl.load(p_A, mask=m_A, other=0.0) for i_v in range(tl.cdiv(V, BV)): - p_v = tl.make_block_ptr(v, (T, V), (HV*V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) - p_u = tl.make_block_ptr(u, (T, V), (HV*V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) - b_v = tl.load(p_v, boundary_check=(0, 1)) + o_v = i_v * BV + tl.arange(0, BV) + m_v = m_t[:, None] & (o_v[None, :] < V) + p_v = v + o_t[:, None] * (HV*V) + o_v[None, :] + p_u = u + o_t[:, None] * (HV*V) + o_v[None, :] + b_v = tl.load(p_v, mask=m_v, other=0.0) b_vb = (b_v * b_b[:, None]).to(b_v.dtype) b_u = tl.dot(b_A, b_vb) - tl.store(p_u, b_u.to(p_u.dtype.element_ty), boundary_check=(0, 1)) + tl.store(p_u, b_u.to(p_u.dtype.element_ty), mask=m_v) for i_k in range(tl.cdiv(K, BK)): - p_w = tl.make_block_ptr(w, (T, K), (HV*K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) - p_k = tl.make_block_ptr(k, (T, K), (H*K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) - b_k = tl.load(p_k, boundary_check=(0, 1)) + o_k = i_k * BK + tl.arange(0, BK) + m_k = o_k < K + m_tk = m_t[:, None] & m_k[None, :] + p_w = w + o_t[:, None] * (HV*K) + o_k[None, :] + p_k = k + o_t[:, None] * (H*K) + o_k[None, :] + b_k = tl.load(p_k, mask=m_tk, other=0.0) b_kb = b_k * b_b[:, None] - p_gk = tl.make_block_ptr(gk, (T, K), (HV*K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) - b_gk = tl.load(p_gk, boundary_check=(0, 1)).to(tl.float32) + p_gk = gk + o_t[:, None] * (HV*K) + o_k[None, :] + b_gk = tl.load(p_gk, mask=m_tk, other=0.0).to(tl.float32) b_kb *= exp2(b_gk) if STORE_QG: - p_q = tl.make_block_ptr(q, (T, K), (H*K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) - p_qg = tl.make_block_ptr(qg, (T, K), (HV*K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) - b_q = tl.load(p_q, boundary_check=(0, 1)) + p_q = q + o_t[:, None] * (H*K) + o_k[None, :] + p_qg = qg + o_t[:, None] * (HV*K) + o_k[None, :] + b_q = tl.load(p_q, mask=m_tk, other=0.0) b_qg = b_q * exp2(b_gk) - tl.store(p_qg, b_qg.to(p_qg.dtype.element_ty), boundary_check=(0, 1)) + tl.store(p_qg, b_qg.to(p_qg.dtype.element_ty), mask=m_tk) if STORE_KG: last_idx = min(i_t * BT + BT, T) - 1 - o_k = i_k * BK + tl.arange(0, BK) - m_k = o_k < K b_gn = tl.load(gk + last_idx * HV*K + o_k, mask=m_k, other=0.).to(tl.float32) b_kg = b_k * tl.where((i_t * BT + tl.arange(0, BT) < T)[:, None], exp2(b_gn[None, :] - b_gk), 0) - p_kg = tl.make_block_ptr(kg, (T, K), (HV*K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) - tl.store(p_kg, b_kg.to(p_kg.dtype.element_ty), boundary_check=(0, 1)) + p_kg = kg + o_t[:, None] * (HV*K) + o_k[None, :] + tl.store(p_kg, b_kg.to(p_kg.dtype.element_ty), mask=m_tk) b_w = tl.dot(b_A, b_kb.to(b_k.dtype)) - tl.store(p_w, b_w.to(p_w.dtype.element_ty), boundary_check=(0, 1)) + tl.store(p_w, b_w.to(p_w.dtype.element_ty), mask=m_tk) @triton.heuristics({ @@ -159,13 +171,17 @@ def prepare_wy_repr_bwd_kda_kernel( BK: tl.constexpr, BV: tl.constexpr, IS_VARLEN: tl.constexpr, + USE_GRAPH: tl.constexpr = False, ): - i_t, i_bh = tl.program_id(0), tl.program_id(1) + i_t, i_bh = tl.program_id(0).to(tl.int64), tl.program_id(1).to(tl.int64) i_b, i_hv = i_bh // HV, i_bh % HV i_h = i_hv // (HV // H) 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.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) + i_n = tl.load(chunk_indices + i_t * 2).to(tl.int32) + if USE_GRAPH and i_n < 0: + return + i_t = tl.load(chunk_indices + i_t * 2 + 1).to(tl.int64) + bos, eos = tl.load(cu_seqlens + i_n).to(tl.int64), tl.load(cu_seqlens + i_n + 1).to(tl.int64) T = eos - bos else: bos, eos = i_b * T, i_b * T + T @@ -185,54 +201,60 @@ def prepare_wy_repr_bwd_kda_kernel( dg += (bos * HV + i_hv) * K dg2 += (bos * HV + i_hv) * K - p_b = tl.make_block_ptr(beta, (T,), (HV,), (i_t * BT,), (BT,), (0,)) - p_db = tl.make_block_ptr(db, (T,), (HV,), (i_t * BT,), (BT,), (0,)) - p_A = tl.make_block_ptr(A, (BT, T), (1, HV*BT), (0, i_t * BT), (BT, BT), (0, 1)) + o_t = i_t * BT + tl.arange(0, BT) + m_t = o_t < T + p_b = beta + o_t * HV + p_db = db + o_t * HV + o_A = tl.arange(0, BT) + m_AT = (o_A[:, None] < BT) & m_t[None, :] + p_A = A + o_A[:, None] + o_t[None, :] * (HV*BT) - b_b = tl.load(p_b, boundary_check=(0,)) + b_b = tl.load(p_b, mask=m_t, other=0.0) b_db = tl.zeros([BT], dtype=tl.float32) - b_A = tl.load(p_A, boundary_check=(0, 1)) + b_A = tl.load(p_A, mask=m_AT, other=0.0) b_dA = tl.zeros([BT, BT], dtype=tl.float32) for i_k in range(tl.cdiv(K, BK)): - p_k = tl.make_block_ptr(k, (T, K), (H*K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) - p_dk = tl.make_block_ptr(dk, (T, K), (HV*K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) - p_dk2 = tl.make_block_ptr(dk2, (T, K), (HV*K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) - p_dw = tl.make_block_ptr(dw, (T, K), (HV*K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) - p_dg = tl.make_block_ptr(dg, (T, K), (HV*K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) - p_dg2 = tl.make_block_ptr(dg2, (T, K), (HV*K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) + o_k = i_k * BK + tl.arange(0, BK) + m_k = m_t[:, None] & (o_k[None, :] < K) + p_k = k + o_t[:, None] * (H*K) + o_k[None, :] + p_dk = dk + o_t[:, None] * (HV*K) + o_k[None, :] + p_dk2 = dk2 + o_t[:, None] * (HV*K) + o_k[None, :] + p_dw = dw + o_t[:, None] * (HV*K) + o_k[None, :] + p_dg = dg + o_t[:, None] * (HV*K) + o_k[None, :] + p_dg2 = dg2 + o_t[:, None] * (HV*K) + o_k[None, :] # [BT, BK] - b_k = tl.load(p_k, boundary_check=(0, 1)) - p_gk = tl.make_block_ptr(gk, (T, K), (HV*K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) - b_gk_exp = exp2(tl.load(p_gk, boundary_check=(0, 1))) + b_k = tl.load(p_k, mask=m_k, other=0.0) + p_gk = gk + o_t[:, None] * (HV*K) + o_k[None, :] + b_gk_exp = exp2(tl.load(p_gk, mask=m_k, other=0.0)) b_kbg = b_k * b_b[:, None] * b_gk_exp - b_dw = tl.load(p_dw, boundary_check=(0, 1)) + b_dw = tl.load(p_dw, mask=m_k, other=0.0) - b_dA += tl.dot(b_dw, tl.trans(b_kbg).to(b_dw.dtype)) + b_dA = tl.dot(b_dw, tl.trans(b_kbg).to(b_dw.dtype), b_dA) b_dkbg = tl.dot(b_A, b_dw) - b_dk = b_dkbg * b_gk_exp * b_b[:, None] + tl.load(p_dk, boundary_check=(0, 1)) + b_dk = b_dkbg * b_gk_exp * b_b[:, None] + tl.load(p_dk, mask=m_k, other=0.0) b_db += tl.sum(b_dkbg * b_k * b_gk_exp, 1) - b_dg = b_kbg * b_dkbg + tl.load(p_dg, boundary_check=(0, 1)) + b_dg = b_kbg * b_dkbg + tl.load(p_dg, mask=m_k, other=0.0) - tl.store(p_dk2, b_dk.to(p_dk2.dtype.element_ty), boundary_check=(0, 1)) - tl.store(p_dg2, b_dg.to(p_dg2.dtype.element_ty), boundary_check=(0, 1)) + tl.store(p_dk2, b_dk.to(p_dk2.dtype.element_ty), mask=m_k) + tl.store(p_dg2, b_dg.to(p_dg2.dtype.element_ty), mask=m_k) for i_v in range(tl.cdiv(V, BV)): - p_v = tl.make_block_ptr(v, (T, V), (HV*V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) - p_dv = tl.make_block_ptr(dv, (T, V), (HV*V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) - p_du = tl.make_block_ptr(du, (T, V), (HV*V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) - b_v = tl.load(p_v, boundary_check=(0, 1)) + o_v = i_v * BV + tl.arange(0, BV) + m_v = m_t[:, None] & (o_v[None, :] < V) + p_v = v + o_t[:, None] * (HV*V) + o_v[None, :] + p_dv = dv + o_t[:, None] * (HV*V) + o_v[None, :] + p_du = du + o_t[:, None] * (HV*V) + o_v[None, :] + b_v = tl.load(p_v, mask=m_v, other=0.0) b_vb = (b_v * b_b[:, None]).to(b_v.dtype) - b_du = tl.load(p_du, boundary_check=(0, 1)) - b_dA += tl.dot(b_du, tl.trans(b_vb)) + b_du = tl.load(p_du, mask=m_v, other=0.0) + b_dA = tl.dot(b_du, tl.trans(b_vb), b_dA) b_dvb = tl.dot(b_A, b_du) b_dv = b_dvb * b_b[:, None] b_db += tl.sum(b_dvb * b_v, 1) - tl.store(p_dv, b_dv.to(p_dv.dtype.element_ty), boundary_check=(0, 1)) + tl.store(p_dv, b_dv.to(p_dv.dtype.element_ty), mask=m_v) - o_t = i_t * BT + tl.arange(0, BT) - m_t = o_t < T m_A = (o_t[:, None] > o_t[None, :]) & (m_t[:, None] & m_t) b_dA = tl.where(m_A, b_dA, 0) b_dA = tl.dot(b_dA.to(b_A.dtype), b_A) @@ -240,11 +262,14 @@ def prepare_wy_repr_bwd_kda_kernel( b_dA = tl.where(m_A, -b_dA, 0) - p_dA = tl.make_block_ptr(dA, (T, BT), (HV*BT, 1), (i_t * BT, 0), (BT, BT), (1, 0)) - tl.store(p_dA, b_dA.to(p_dA.dtype.element_ty), boundary_check=(0, 1)) - tl.store(p_db, b_db.to(p_db.dtype.element_ty), boundary_check=(0,)) + m_dA = m_t[:, None] & (o_A[None, :] < BT) + p_dA = dA + o_t[:, None] * (HV*BT) + o_A[None, :] + tl.store(p_dA, b_dA.to(p_dA.dtype.element_ty), mask=m_dA) + tl.store(p_db, b_db.to(p_db.dtype.element_ty), mask=m_t) +# Use triton only, other backends (cutlass, cute, etc.) need additional torch.compile support +# @dispatch('kda') def recompute_w_u_fwd( k: torch.Tensor, v: torch.Tensor, @@ -254,6 +279,7 @@ def recompute_w_u_fwd( q: torch.Tensor | None = None, cu_seqlens: torch.LongTensor | None = None, chunk_indices: torch.LongTensor | None = None, + use_graph: bool = False, ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor | None, torch.Tensor | None]: B, T, H, K, V = *k.shape, v.shape[-1] HV = v.shape[2] @@ -290,6 +316,7 @@ def recompute_w_u_fwd( BT=BT, BK=BK, BV=BV, + USE_GRAPH=use_graph, ) return w, u, qg, kg @@ -306,6 +333,7 @@ def prepare_wy_repr_bwd( dg: torch.Tensor, cu_seqlens: torch.LongTensor | None = None, chunk_indices: torch.LongTensor | None = None, + use_graph: bool = False, ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: B, T, H, K, V = *k.shape, v.shape[-1] HV = v.shape[2] @@ -347,6 +375,7 @@ def prepare_wy_repr_bwd( BT=BT, BK=BK, BV=BV, + USE_GRAPH=use_graph, ) dk = dk2 dg = dg2 diff --git a/build/torch-xpu/modules/conv/cuda/__init__.py b/build/torch-cuda/ops/lightnet/__init__.py similarity index 64% rename from build/torch-xpu/modules/conv/cuda/__init__.py rename to build/torch-cuda/ops/lightnet/__init__.py index 1e1485fa2aee518131fb55be9a7addc870b5f819..11f3dcbd3ba0914f8532b7a1ab3e64c63d06ab8f 100644 --- a/build/torch-xpu/modules/conv/cuda/__init__.py +++ b/build/torch-cuda/ops/lightnet/__init__.py @@ -5,10 +5,10 @@ # 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 +from ...ops.lightnet.chunk import chunk_lightnet +from ...ops.lightnet.fused_recurrent import fused_recurrent_lightnet __all__ = [ - 'FastCausalConv1dFn', - 'causal_conv1d_cuda', - 'fast_causal_conv1d_fn', + 'chunk_lightnet', + 'fused_recurrent_lightnet', ] diff --git a/build/torch-cuda/ops/lightnet/chunk.py b/build/torch-cuda/ops/lightnet/chunk.py new file mode 100644 index 0000000000000000000000000000000000000000..be0d279586d4d18c474fcbba8990df0877750cc5 --- /dev/null +++ b/build/torch-cuda/ops/lightnet/chunk.py @@ -0,0 +1,66 @@ +# 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 + +from ...ops.gla import chunk_gla +from ...ops.gsa.gate import fused_additive_gate + + +@torch.compiler.disable +def chunk_lightnet( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + scale: float | None = None, + initial_state: tuple[torch.Tensor | None, torch.Tensor | None] | None = None, + output_final_state: bool = False, + state_v_first: bool = False, + cu_seqlens: torch.LongTensor | None = None, +) -> tuple[torch.Tensor, tuple[torch.Tensor | None, torch.Tensor]]: + """Compute chunkwise LightNet attention with cumulative key normalization. + + Args: + q (torch.Tensor): + Queries of shape `[B, T, H, K]`. + k (torch.Tensor): + Key logits of shape `[B, T, H, K]`. + v (torch.Tensor): + Values of shape `[B, T, H, V]`. + scale (float, Optional): + Attention score scale, defaulting to `1 / sqrt(K)`. Default: `None`. + initial_state (tuple, Optional): + FP32 `(recurrent_state, log_normalizer)` with shapes `[N, H, K, V]` and `[N, 1, H, K]`. + Either component may be `None`. `N` is the number of sequences, equal to `B` for dense inputs. Default: `None`. + output_final_state (bool, Optional): + Whether to return the final recurrent state. Default: `False`. + state_v_first (bool, Optional): + Use `[N, H, V, K]` for recurrent states instead of `[N, H, K, V]`. Default: `False`. + cu_seqlens (torch.LongTensor, Optional): + Cumulative sequence lengths of shape `[N+1]` for packed inputs with `B=1`. Default: `None`. + + Returns: + o (torch.Tensor): + Outputs of shape `[B, T, H, V]`. + final_state (tuple): + `(recurrent_state, log_normalizer)` in the same layouts as `initial_state`. + The recurrent state is `None` when `output_final_state=False`; the final log-normalizer is always returned. + """ + h0, z0 = (None, None) if initial_state is None else initial_state + k, g, zt = fused_additive_gate(x=k, initial_state=z0, cu_seqlens=cu_seqlens) + o, ht = chunk_gla( + q=q, + k=k, + v=v, + g=g, + scale=scale, + initial_state=h0, + output_final_state=output_final_state, + state_v_first=state_v_first, + cu_seqlens=cu_seqlens, + ) + return o, (ht, zt) diff --git a/build/torch-cuda/ops/lightnet/fused_recurrent.py b/build/torch-cuda/ops/lightnet/fused_recurrent.py new file mode 100644 index 0000000000000000000000000000000000000000..c0150c74820d42ba5efd13b6dc1935d347280abd --- /dev/null +++ b/build/torch-cuda/ops/lightnet/fused_recurrent.py @@ -0,0 +1,65 @@ +# 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 + +from ...ops.gla import fused_recurrent_gla +from ...ops.gsa.gate import fused_additive_gate + + +def fused_recurrent_lightnet( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + scale: float | None = None, + initial_state: tuple[torch.Tensor | None, torch.Tensor | None] | None = None, + output_final_state: bool = False, + state_v_first: bool = False, + cu_seqlens: torch.LongTensor | None = None, +) -> tuple[torch.Tensor, tuple[torch.Tensor | None, torch.Tensor]]: + """Compute recurrent LightNet attention with cumulative key normalization. + + Args: + q (torch.Tensor): + Queries of shape `[B, T, H, K]`. + k (torch.Tensor): + Key logits of shape `[B, T, H, K]`. + v (torch.Tensor): + Values of shape `[B, T, H, V]`. + scale (float, Optional): + Attention score scale, defaulting to `1 / sqrt(K)`. Default: `None`. + initial_state (tuple, Optional): + FP32 `(recurrent_state, log_normalizer)` with shapes `[N, H, K, V]` and `[N, 1, H, K]`. + Either component may be `None`. `N` is the number of sequences, equal to `B` for dense inputs. Default: `None`. + output_final_state (bool, Optional): + Whether to return the final recurrent state. Default: `False`. + state_v_first (bool, Optional): + Use `[N, H, V, K]` for recurrent states instead of `[N, H, K, V]`. Default: `False`. + cu_seqlens (torch.LongTensor, Optional): + Cumulative sequence lengths of shape `[N+1]` for packed inputs with `B=1`. Default: `None`. + + Returns: + o (torch.Tensor): + Outputs of shape `[B, T, H, V]`. + final_state (tuple): + `(recurrent_state, log_normalizer)` in the same layouts as `initial_state`. + The recurrent state is `None` when `output_final_state=False`; the final log-normalizer is always returned. + """ + h0, z0 = (None, None) if initial_state is None else initial_state + k, g, zt = fused_additive_gate(x=k, initial_state=z0, cu_seqlens=cu_seqlens) + o, ht = fused_recurrent_gla( + q=q, + k=k, + v=v, + gk=g, + scale=scale, + initial_state=h0, + output_final_state=output_final_state, + state_v_first=state_v_first, + cu_seqlens=cu_seqlens, + ) + return o, (ht, zt) diff --git a/build/torch-cuda/ops/lightning_attn/chunk.py b/build/torch-cuda/ops/lightning_attn/chunk.py index 728e9a8111fb9d4ba5d18c8ec8f2ec5a31053c7f..7f31e919b8c09d15b0405e7108d5c5b30963e0b6 100644 --- a/build/torch-cuda/ops/lightning_attn/chunk.py +++ b/build/torch-cuda/ops/lightning_attn/chunk.py @@ -70,8 +70,8 @@ def chunk_lightning_attn( f"The number of initial states is expected to be equal to the number of input sequences, " f"i.e., {len(cu_seqlens) - 1} rather than {initial_state.shape[0]}.", ) - H = q.shape[2] - g_gamma = -(8 / H * (1 - layer_idx / num_layers)) * q.new_tensor(range(H), dtype=torch.float) + H, L = q.shape[2], num_layers + g_gamma = -(8 / H * (1 - layer_idx / L)) * torch.arange(H, dtype=torch.float, device=q.device) return chunk_simple_gla( q=q, k=k, diff --git a/build/torch-cuda/ops/lightning_attn/fused_recurrent.py b/build/torch-cuda/ops/lightning_attn/fused_recurrent.py index a8efaefcc46371b9dc979bdb908bd6554bbcc196..a415af4243f0beec54aec9abb610c60464dc4f0a 100644 --- a/build/torch-cuda/ops/lightning_attn/fused_recurrent.py +++ b/build/torch-cuda/ops/lightning_attn/fused_recurrent.py @@ -21,7 +21,6 @@ def fused_recurrent_lightning_attn( output_final_state: bool = False, reverse: bool = False, cu_seqlens: torch.LongTensor | None = None, - **kwargs, ) -> tuple[torch.Tensor, torch.Tensor]: r""" Args: @@ -56,10 +55,6 @@ def fused_recurrent_lightning_attn( final_state (torch.Tensor): Final state of shape `[N, H, K, V]` if `output_final_state=True` else `None`. """ - if 'head_first' in kwargs: - raise DeprecationWarning( - "head_first has been removed. Inputs must be in `[B, T, H, ...]` format.", - ) if cu_seqlens is not None: if q.shape[0] != 1: raise ValueError( @@ -71,8 +66,8 @@ def fused_recurrent_lightning_attn( f"The number of initial states is expected to be equal to the number of input sequences, " f"i.e., {len(cu_seqlens) - 1} rather than {initial_state.shape[0]}.", ) - H = q.shape[2] - g_gamma = -(8 / H * (1 - layer_idx / num_layers)) * q.new_tensor(range(H), dtype=torch.float) + H, L = q.shape[2], num_layers + g_gamma = -(8 / H * (1 - layer_idx / L)) * torch.arange(H, dtype=torch.float, device=q.device) return fused_recurrent_simple_gla( q=q, k=k, diff --git a/build/torch-cuda/ops/linear_attn/utils.py b/build/torch-cuda/ops/linear_attn/utils.py index 9175d78f8ba17bf6d7710fdfab6979af089f7488..02b91c4fb0583c89e83dfa78e9b47e15c1474993 100644 --- a/build/torch-cuda/ops/linear_attn/utils.py +++ b/build/torch-cuda/ops/linear_attn/utils.py @@ -31,7 +31,13 @@ class StatefulNormalizeFunction(torch.autograd.Function): o_out = o / denom if cu_seqlens is not None: idx = (cu_seqlens[:-1] if reverse else cu_seqlens[1:] - 1).to(torch.long) - z_state_out = k_cum[0, idx].unsqueeze(1).contiguous() + nonempty = (cu_seqlens[1:] > cu_seqlens[:-1])[:, None, None, None] + if k.shape[1] > 0: + z_state_out = k_cum[0, idx.clamp(0, k.shape[1] - 1)].unsqueeze(1) + else: + z_state_out = k_cum.new_zeros(len(cu_seqlens) - 1, 1, *k.shape[2:]) + # empty sequences retain their initial key sum + z_state_out = torch.where(nonempty, z_state_out, z_init if z_init is not None else 0).contiguous() else: z_state_out = (k_cum[:, :1] if reverse else k_cum[:, -1:]).contiguous() @@ -58,8 +64,11 @@ class StatefulNormalizeFunction(torch.autograd.Function): if dz_state is not None: if cu_seqlens is not None: - idx = (cu_seqlens[:-1] if reverse else cu_seqlens[1:] - 1).to(torch.long) - dk_cum[0].index_add_(0, idx, dz_state.squeeze(1).to(dk_cum.dtype)) + if k_cum.shape[1] > 0: + idx = (cu_seqlens[:-1] if reverse else cu_seqlens[1:] - 1).to(torch.long) + nonempty = (cu_seqlens[1:] > cu_seqlens[:-1])[:, None, None] + dz = dz_state.squeeze(1).to(dk_cum.dtype).masked_fill(~nonempty, 0) + dk_cum[0].index_add_(0, idx.clamp(0, k_cum.shape[1] - 1), dz) elif reverse: dk_cum[:, :1].add_(dz_state.to(dk_cum.dtype)) else: @@ -79,6 +88,9 @@ class StatefulNormalizeFunction(torch.autograd.Function): ) dz_init.index_add_(0, seg_id, flat) dz_init = dz_init.unsqueeze(1) + if dz_state is not None: + nonempty = (seq_lens > 0)[:, None, None, None] + dz_init = dz_init + dz_state.masked_fill(nonempty, 0) else: dz_init = dk_cum.sum(dim=1, keepdim=True) diff --git a/build/torch-cuda/ops/log_linear_attn/chunk.py b/build/torch-cuda/ops/log_linear_attn/chunk.py index 35991283935d5c72ed8fb563134908ac6de080ff..05c22ce488f400f0193bf48670a8bd7a1a91a856 100644 --- a/build/torch-cuda/ops/log_linear_attn/chunk.py +++ b/build/torch-cuda/ops/log_linear_attn/chunk.py @@ -16,7 +16,7 @@ import triton.language as tl from einops import reduce from ...ops.utils import chunk_local_cumsum -from ...ops.utils.op import exp +from ...ops.utils.op import exp, unflatten_program_id from ...utils import autocast_custom_bwd, autocast_custom_fwd, autotune_cache_kwargs, input_guard BLOCK_K = 64 @@ -35,10 +35,10 @@ BLOCK_K = 64 for num_warps in [4] for num_stages in [2, 3, 4] ], - key=["H", "K", "V"], + key=["G", "H", "K", "V"], **autotune_cache_kwargs, ) -@triton.jit(do_not_specialize=["T"]) +@triton.jit(do_not_specialize=["T", "output_T"]) def chunkwise_fwd_kernel( q, k, @@ -52,7 +52,11 @@ def chunkwise_fwd_kernel( offsets, new_offsets, cu_seqlens, + output_cu_seqlens, + stride_o_n, T, + output_T, + G: tl.constexpr, H: tl.constexpr, K: tl.constexpr, V: tl.constexpr, @@ -66,13 +70,15 @@ def chunkwise_fwd_kernel( IS_VARLEN: tl.constexpr, USE_INITIAL_STATE: tl.constexpr, STORE_FINAL_STATE: tl.constexpr, + SCALE: tl.constexpr, ): - p_llut = tl.make_block_ptr(llut, (BT, BT), (BT, 1), (0, 0), (BT, BT), (1, 0)) - b_llut = tl.load(p_llut, boundary_check=(0, 1)) + o_i = tl.arange(0, BT) + p_llut = llut + o_i[:, None] * BT + o_i[None, :] + b_llut = tl.load(p_llut, mask=(o_i[:, None] < BT) & (o_i[None, :] < BT), other=0.0) # parallel over sequences and heads - i_k = tl.program_id(0) - i_nh = tl.program_id(1) + i_k, i_nh = unflatten_program_id(tl.cdiv(K, BK)) i_n, i_h = i_nh // H, i_nh % H + i_g = i_h // (H // G) if IS_VARLEN: bos, eos = ( @@ -83,7 +89,9 @@ def chunkwise_fwd_kernel( else: bos, eos = (i_n * T).to(tl.int64), (i_n * T + T).to(tl.int64) - o_i = tl.arange(0, BT) + o_kk = i_k * BK + tl.arange(0, BK) + o_v = tl.arange(0, V) + m_kv = (o_kk[:, None] < K) & (o_v[None, :] < V) # For hierarchical masking num_intra_levels = (tl.log2(float(BT))).to(tl.int32) + 1 @@ -126,168 +134,88 @@ def chunkwise_fwd_kernel( first_chunk_index = offset // BT if KV_0_CREATED and (first_chunk_index & 1 > 0): - p_kv_0 = tl.make_block_ptr( - h0 + ((i_n * L_IN + 0) * H + i_h) * K * V, - (K, V), - (V, 1), - (i_k * BK, 0), - (BK, V), - (1, 0), - ) - kv_0 = tl.load(p_kv_0, boundary_check=(0, 1)) + p_kv_0 = h0 + ((i_n * L_IN + 0) * H + i_h) * K * V + o_kk[:, None] * V + o_v[None, :] + kv_0 = tl.load(p_kv_0, mask=m_kv, other=0.0) if KV_1_CREATED and (first_chunk_index & 2 > 0): - p_kv_1 = tl.make_block_ptr( - h0 + ((i_n * L_IN + 1) * H + i_h) * K * V, - (K, V), - (V, 1), - (i_k * BK, 0), - (BK, V), - (1, 0), - ) - kv_1 = tl.load(p_kv_1, boundary_check=(0, 1)) + p_kv_1 = h0 + ((i_n * L_IN + 1) * H + i_h) * K * V + o_kk[:, None] * V + o_v[None, :] + kv_1 = tl.load(p_kv_1, mask=m_kv, other=0.0) if KV_2_CREATED and (first_chunk_index & 4 > 0): - p_kv_2 = tl.make_block_ptr( - h0 + ((i_n * L_IN + 2) * H + i_h) * K * V, - (K, V), - (V, 1), - (i_k * BK, 0), - (BK, V), - (1, 0), - ) - kv_2 = tl.load(p_kv_2, boundary_check=(0, 1)) + p_kv_2 = h0 + ((i_n * L_IN + 2) * H + i_h) * K * V + o_kk[:, None] * V + o_v[None, :] + kv_2 = tl.load(p_kv_2, mask=m_kv, other=0.0) if KV_3_CREATED and (first_chunk_index & 8 > 0): - p_kv_3 = tl.make_block_ptr( - h0 + ((i_n * L_IN + 3) * H + i_h) * K * V, - (K, V), - (V, 1), - (i_k * BK, 0), - (BK, V), - (1, 0), - ) - kv_3 = tl.load(p_kv_3, boundary_check=(0, 1)) + p_kv_3 = h0 + ((i_n * L_IN + 3) * H + i_h) * K * V + o_kk[:, None] * V + o_v[None, :] + kv_3 = tl.load(p_kv_3, mask=m_kv, other=0.0) if KV_4_CREATED and (first_chunk_index & 16 > 0): - p_kv_4 = tl.make_block_ptr( - h0 + ((i_n * L_IN + 4) * H + i_h) * K * V, - (K, V), - (V, 1), - (i_k * BK, 0), - (BK, V), - (1, 0), - ) - kv_4 = tl.load(p_kv_4, boundary_check=(0, 1)) + p_kv_4 = h0 + ((i_n * L_IN + 4) * H + i_h) * K * V + o_kk[:, None] * V + o_v[None, :] + kv_4 = tl.load(p_kv_4, mask=m_kv, other=0.0) if KV_5_CREATED and (first_chunk_index & 32 > 0): - p_kv_5 = tl.make_block_ptr( - h0 + ((i_n * L_IN + 5) * H + i_h) * K * V, - (K, V), - (V, 1), - (i_k * BK, 0), - (BK, V), - (1, 0), - ) - kv_5 = tl.load(p_kv_5, boundary_check=(0, 1)) + p_kv_5 = h0 + ((i_n * L_IN + 5) * H + i_h) * K * V + o_kk[:, None] * V + o_v[None, :] + kv_5 = tl.load(p_kv_5, mask=m_kv, other=0.0) if KV_6_CREATED and (first_chunk_index & 64 > 0): - p_kv_6 = tl.make_block_ptr( - h0 + ((i_n * L_IN + 6) * H + i_h) * K * V, - (K, V), - (V, 1), - (i_k * BK, 0), - (BK, V), - (1, 0), - ) - kv_6 = tl.load(p_kv_6, boundary_check=(0, 1)) + p_kv_6 = h0 + ((i_n * L_IN + 6) * H + i_h) * K * V + o_kk[:, None] * V + o_v[None, :] + kv_6 = tl.load(p_kv_6, mask=m_kv, other=0.0) if KV_7_CREATED and (first_chunk_index & 128 > 0): - p_kv_7 = tl.make_block_ptr( - h0 + ((i_n * L_IN + 7) * H + i_h) * K * V, - (K, V), - (V, 1), - (i_k * BK, 0), - (BK, V), - (1, 0), - ) - kv_7 = tl.load(p_kv_7, boundary_check=(0, 1)) + p_kv_7 = h0 + ((i_n * L_IN + 7) * H + i_h) * K * V + o_kk[:, None] * V + o_v[None, :] + kv_7 = tl.load(p_kv_7, mask=m_kv, other=0.0) if KV_8_CREATED and (first_chunk_index & 256 > 0): - p_kv_8 = tl.make_block_ptr( - h0 + ((i_n * L_IN + 8) * H + i_h) * K * V, - (K, V), - (V, 1), - (i_k * BK, 0), - (BK, V), - (1, 0), - ) - kv_8 = tl.load(p_kv_8, boundary_check=(0, 1)) + p_kv_8 = h0 + ((i_n * L_IN + 8) * H + i_h) * K * V + o_kk[:, None] * V + o_v[None, :] + kv_8 = tl.load(p_kv_8, mask=m_kv, other=0.0) if KV_9_CREATED and (first_chunk_index & 512 > 0): - p_kv_9 = tl.make_block_ptr( - h0 + ((i_n * L_IN + 9) * H + i_h) * K * V, - (K, V), - (V, 1), - (i_k * BK, 0), - (BK, V), - (1, 0), - ) - kv_9 = tl.load(p_kv_9, boundary_check=(0, 1)) + p_kv_9 = h0 + ((i_n * L_IN + 9) * H + i_h) * K * V + o_kk[:, None] * V + o_v[None, :] + kv_9 = tl.load(p_kv_9, mask=m_kv, other=0.0) if KV_10_CREATED and (first_chunk_index & 1024 > 0): - p_kv_10 = tl.make_block_ptr( - h0 + ((i_n * L_IN + 10) * H + i_h) * K * V, - (K, V), - (V, 1), - (i_k * BK, 0), - (BK, V), - (1, 0), - ) - kv_10 = tl.load(p_kv_10, boundary_check=(0, 1)) + p_kv_10 = h0 + ((i_n * L_IN + 10) * H + i_h) * K * V + o_kk[:, None] * V + o_v[None, :] + kv_10 = tl.load(p_kv_10, mask=m_kv, other=0.0) if KV_11_CREATED and (first_chunk_index & 2048 > 0): - p_kv_11 = tl.make_block_ptr( - h0 + ((i_n * L_IN + 11) * H + i_h) * K * V, - (K, V), - (V, 1), - (i_k * BK, 0), - (BK, V), - (1, 0), - ) - kv_11 = tl.load(p_kv_11, boundary_check=(0, 1)) + p_kv_11 = h0 + ((i_n * L_IN + 11) * H + i_h) * K * V + o_kk[:, None] * V + o_v[None, :] + kv_11 = tl.load(p_kv_11, mask=m_kv, other=0.0) + + if not USE_INITIAL_STATE: + output_T = T + elif IS_VARLEN: + output_bos = tl.load(output_cu_seqlens + i_n).to(tl.int64) + output_eos = tl.load(output_cu_seqlens + i_n + 1).to(tl.int64) + output_T = (output_eos - output_bos).to(tl.int32) NT = tl.cdiv(T, BT) output_offset = -1 * (offset % BT) for i_t in range(NT): - b_h_ptrs = level_scales + ((bos + i_t * BT + i_idx) * H + i_h) * L + b_llut - b_h = tl.load(b_h_ptrs, mask=i_idx >= j_idx) - - p_g = tl.make_block_ptr(g + bos * H + i_h, (T,), (H,), (i_t * BT,), (BT,), (0,)) - p_q = tl.make_block_ptr( - q + bos * K, (T, K), (K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0), - ) - p_k = tl.make_block_ptr( - k + bos * K, (K, T), (1, K), (i_k * BK, i_t * BT), (BK, BT), (0, 1), - ) - p_v = tl.make_block_ptr( - v + (bos * H + i_h) * V, - (T, V), - (H * V, 1), - (i_t * BT, 0), - (BT, V), - (1, 0), - ) - p_o = tl.make_block_ptr( - o + ((bos * H + i_h) * (K // BK) + i_k) * V, - (T, V), - (H * (K // BK) * V, 1), - (i_t * BT + output_offset, 0), - (BT, V), - (1, 0), - ) + b_h_ptrs = level_scales + ((bos + tl.minimum(i_t * BT + i_idx, T - 1)) * H + i_h) * L + b_llut + b_h = tl.load(b_h_ptrs, mask=(i_idx >= j_idx) & (i_t * BT + i_idx < T) & (i_t * BT + j_idx < T), other=0.0) + + o_t = (i_t * BT).to(tl.int64) + o_i + o_to = o_t + output_offset + m_t = o_t < T + m_to = (o_to >= 0) & (o_to < output_T) + m_qk = m_t[:, None] & (o_kk[None, :] < K) + m_kt = (o_kk[:, None] < K) & m_t[None, :] + m_tv = m_t[:, None] & (o_v[None, :] < V) + m_ov = m_to[:, None] & (o_v[None, :] < V) + p_g = g + bos * H + i_h + o_t * H + p_q = q + (bos * G + i_g) * K + o_t[:, None] * (G * K) + o_kk[None, :] + p_k = k + (bos * G + i_g) * K + o_kk[:, None] + o_t[None, :] * (G * K) + p_v = v + (bos * H + i_h) * V + o_t[:, None] * (H * V) + o_v[None, :] + if USE_INITIAL_STATE: + if IS_VARLEN: + p_o = o + ((output_bos * H + i_h) * (K // BK) + i_k) * V + else: + p_o = o + i_n.to(tl.int64) * stride_o_n + (i_h * (K // BK) + i_k) * V + else: + p_o = o + ((bos * H + i_h) * (K // BK) + i_k) * V + p_o = p_o + o_to[:, None] * (H * (K // BK) * V) + o_v[None, :] - b_g = tl.load(p_g, boundary_check=(0,)) - b_q = tl.load(p_q, boundary_check=(0, 1)) - b_k = tl.load(p_k, boundary_check=(0, 1)) + b_g = tl.load(p_g, mask=m_t, other=0.0) + b_q = tl.load(p_q, mask=m_qk, other=0.0) + b_k = tl.load(p_k, mask=m_kt, other=0.0) m_t = i_t * BT + o_i < T b_s = tl.where((i_idx >= j_idx) & m_t[:, None] & m_t[None, :], tl.exp(b_g[:, None] - b_g[None, :]), 0) b_s = (tl.dot(b_q, b_k) * b_s).to(b_q.dtype) * b_h - b_v = tl.load(p_v, boundary_check=(0, 1)) + b_v = tl.load(p_v, mask=m_tv, other=0.0) b_o = tl.zeros((BT, V), dtype=tl.float32) if MIN_LEVEL == 0: - b_o += tl.dot(b_s, b_v) + b_o = tl.dot(b_s, b_v, b_o) chunk_index = ( first_chunk_index + i_t @@ -295,150 +223,66 @@ def chunkwise_fwd_kernel( if MIN_LEVEL <= 0 and MAX_LEVEL >= 0: if chunk_index & 1: - p_l = tl.make_block_ptr( - level_scales + (bos * H + i_h) * L, - (T, L), - (H * L, 1), - (i_t * BT, num_intra_levels), - (BT, 1), - (1, 0), - ) - b_l = tl.load(p_l, boundary_check=(0, 1)) + p_l = level_scales + (bos * H + i_h) * L + o_t[:, None] * (H * L) + num_intra_levels + b_l = tl.load(p_l, mask=m_t[:, None] & ((num_intra_levels) < L), other=0.0) b_o += tl.dot((b_l * b_q), kv_0.to(b_q.dtype)) * tl.exp(b_g)[:, None] if MIN_LEVEL <= 1 and MAX_LEVEL >= 1: if chunk_index & 2: - p_l = tl.make_block_ptr( - level_scales + (bos * H + i_h) * L, - (T, L), - (H * L, 1), - (i_t * BT, num_intra_levels + 1), - (BT, 1), - (1, 0), - ) - b_l = tl.load(p_l, boundary_check=(0, 1)) + p_l = level_scales + (bos * H + i_h) * L + o_t[:, None] * (H * L) + num_intra_levels + 1 + b_l = tl.load(p_l, mask=m_t[:, None] & ((num_intra_levels + 1) < L), other=0.0) b_o += tl.dot((b_l * b_q), kv_1.to(b_q.dtype)) * tl.exp(b_g)[:, None] if MIN_LEVEL <= 2 and MAX_LEVEL >= 2: if chunk_index & 4: - p_l = tl.make_block_ptr( - level_scales + (bos * H + i_h) * L, - (T, L), - (H * L, 1), - (i_t * BT, num_intra_levels + 2), - (BT, 1), - (1, 0), - ) - b_l = tl.load(p_l, boundary_check=(0, 1)) + p_l = level_scales + (bos * H + i_h) * L + o_t[:, None] * (H * L) + num_intra_levels + 2 + b_l = tl.load(p_l, mask=m_t[:, None] & ((num_intra_levels + 2) < L), other=0.0) b_o += tl.dot((b_l * b_q), kv_2.to(b_q.dtype)) * tl.exp(b_g)[:, None] if MIN_LEVEL <= 3 and MAX_LEVEL >= 3: if chunk_index & 8: - p_l = tl.make_block_ptr( - level_scales + (bos * H + i_h) * L, - (T, L), - (H * L, 1), - (i_t * BT, num_intra_levels + 3), - (BT, 1), - (1, 0), - ) - b_l = tl.load(p_l, boundary_check=(0, 1)) + p_l = level_scales + (bos * H + i_h) * L + o_t[:, None] * (H * L) + num_intra_levels + 3 + b_l = tl.load(p_l, mask=m_t[:, None] & ((num_intra_levels + 3) < L), other=0.0) b_o += tl.dot((b_l * b_q), kv_3.to(b_q.dtype)) * tl.exp(b_g)[:, None] if MIN_LEVEL <= 4 and MAX_LEVEL >= 4: if chunk_index & 16: - p_l = tl.make_block_ptr( - level_scales + (bos * H + i_h) * L, - (T, L), - (H * L, 1), - (i_t * BT, num_intra_levels + 4), - (BT, 1), - (1, 0), - ) - b_l = tl.load(p_l, boundary_check=(0, 1)) + p_l = level_scales + (bos * H + i_h) * L + o_t[:, None] * (H * L) + num_intra_levels + 4 + b_l = tl.load(p_l, mask=m_t[:, None] & ((num_intra_levels + 4) < L), other=0.0) b_o += tl.dot((b_l * b_q), kv_4.to(b_q.dtype)) * tl.exp(b_g)[:, None] if MIN_LEVEL <= 5 and MAX_LEVEL >= 5: if chunk_index & 32: - p_l = tl.make_block_ptr( - level_scales + (bos * H + i_h) * L, - (T, L), - (H * L, 1), - (i_t * BT, num_intra_levels + 5), - (BT, 1), - (1, 0), - ) - b_l = tl.load(p_l, boundary_check=(0, 1)) + p_l = level_scales + (bos * H + i_h) * L + o_t[:, None] * (H * L) + num_intra_levels + 5 + b_l = tl.load(p_l, mask=m_t[:, None] & ((num_intra_levels + 5) < L), other=0.0) b_o += tl.dot((b_l * b_q), kv_5.to(b_q.dtype)) * tl.exp(b_g)[:, None] if MIN_LEVEL <= 6 and MAX_LEVEL >= 6: if chunk_index & 64: - p_l = tl.make_block_ptr( - level_scales + (bos * H + i_h) * L, - (T, L), - (H * L, 1), - (i_t * BT, num_intra_levels + 6), - (BT, 1), - (1, 0), - ) - b_l = tl.load(p_l, boundary_check=(0, 1)) + p_l = level_scales + (bos * H + i_h) * L + o_t[:, None] * (H * L) + num_intra_levels + 6 + b_l = tl.load(p_l, mask=m_t[:, None] & ((num_intra_levels + 6) < L), other=0.0) b_o += tl.dot((b_l * b_q), kv_6.to(b_q.dtype)) * tl.exp(b_g)[:, None] if MIN_LEVEL <= 7 and MAX_LEVEL >= 7: if chunk_index & 128: # 8192 - 16384 - p_l = tl.make_block_ptr( - level_scales + (bos * H + i_h) * L, - (T, L), - (H * L, 1), - (i_t * BT, num_intra_levels + 7), - (BT, 1), - (1, 0), - ) - b_l = tl.load(p_l, boundary_check=(0, 1)) + p_l = level_scales + (bos * H + i_h) * L + o_t[:, None] * (H * L) + num_intra_levels + 7 + b_l = tl.load(p_l, mask=m_t[:, None] & ((num_intra_levels + 7) < L), other=0.0) b_o += tl.dot((b_l * b_q), kv_7.to(b_q.dtype)) * tl.exp(b_g)[:, None] if MIN_LEVEL <= 8 and MAX_LEVEL >= 8: if chunk_index & 256: - p_l = tl.make_block_ptr( - level_scales + (bos * H + i_h) * L, - (T, L), - (H * L, 1), - (i_t * BT, num_intra_levels + 8), - (BT, 1), - (1, 0), - ) - b_l = tl.load(p_l, boundary_check=(0, 1)) + p_l = level_scales + (bos * H + i_h) * L + o_t[:, None] * (H * L) + num_intra_levels + 8 + b_l = tl.load(p_l, mask=m_t[:, None] & ((num_intra_levels + 8) < L), other=0.0) b_o += tl.dot((b_l * b_q), kv_8.to(b_q.dtype)) * tl.exp(b_g)[:, None] if MIN_LEVEL <= 9 and MAX_LEVEL >= 9: if chunk_index & 512: - p_l = tl.make_block_ptr( - level_scales + (bos * H + i_h) * L, - (T, L), - (H * L, 1), - (i_t * BT, num_intra_levels + 9), - (BT, 1), - (1, 0), - ) - b_l = tl.load(p_l, boundary_check=(0, 1)) + p_l = level_scales + (bos * H + i_h) * L + o_t[:, None] * (H * L) + num_intra_levels + 9 + b_l = tl.load(p_l, mask=m_t[:, None] & ((num_intra_levels + 9) < L), other=0.0) b_o += tl.dot((b_l * b_q), kv_9.to(b_q.dtype)) * tl.exp(b_g)[:, None] if MIN_LEVEL <= 10 and MAX_LEVEL >= 10: if chunk_index & 1024: - p_l = tl.make_block_ptr( - level_scales + (bos * H + i_h) * L, - (T, L), - (H * L, 1), - (i_t * BT, num_intra_levels + 10), - (BT, 1), - (1, 0), - ) - b_l = tl.load(p_l, boundary_check=(0, 1)) + p_l = level_scales + (bos * H + i_h) * L + o_t[:, None] * (H * L) + num_intra_levels + 10 + b_l = tl.load(p_l, mask=m_t[:, None] & ((num_intra_levels + 10) < L), other=0.0) b_o += tl.dot((b_l * b_q), kv_10.to(b_q.dtype)) * tl.exp(b_g)[:, None] if MIN_LEVEL <= 11 and MAX_LEVEL >= 11: if chunk_index & 2048: - p_l = tl.make_block_ptr( - level_scales + (bos * H + i_h) * L, - (T, L), - (H * L, 1), - (i_t * BT, num_intra_levels + 11), - (BT, 1), - (1, 0), - ) - b_l = tl.load(p_l, boundary_check=(0, 1)) + p_l = level_scales + (bos * H + i_h) * L + o_t[:, None] * (H * L) + num_intra_levels + 11 + b_l = tl.load(p_l, mask=m_t[:, None] & ((num_intra_levels + 11) < L), other=0.0) b_o += tl.dot((b_l * b_q), kv_11.to(b_q.dtype)) * tl.exp(b_g)[:, None] - tl.store(p_o, b_o.to(p_o.dtype.element_ty), boundary_check=(0, 1)) + tl.store(p_o, (b_o * SCALE).to(p_o.dtype.element_ty), mask=m_ov) if i_t < NT - 1 or T % BT == 0: # Only apply the state update if the last chunk is a full chunk. @@ -474,27 +318,27 @@ def chunkwise_fwd_kernel( b_v = (b_v * tl.exp(b_g_last - b_g)[:, None]).to(b_v.dtype) if MIN_LEVEL <= 1: - kv_0 += tl.dot(b_k, b_v) + kv_0 = tl.dot(b_k, b_v, kv_0) elif MIN_LEVEL == 2: - kv_1 += tl.dot(b_k, b_v) + kv_1 = tl.dot(b_k, b_v, kv_1) elif MIN_LEVEL == 3: - kv_2 += tl.dot(b_k, b_v) + kv_2 = tl.dot(b_k, b_v, kv_2) elif MIN_LEVEL == 4: - kv_3 += tl.dot(b_k, b_v) + kv_3 = tl.dot(b_k, b_v, kv_3) elif MIN_LEVEL == 5: - kv_4 += tl.dot(b_k, b_v) + kv_4 = tl.dot(b_k, b_v, kv_4) elif MIN_LEVEL == 6: - kv_5 += tl.dot(b_k, b_v) + kv_5 = tl.dot(b_k, b_v, kv_5) elif MIN_LEVEL == 7: - kv_6 += tl.dot(b_k, b_v) + kv_6 = tl.dot(b_k, b_v, kv_6) elif MIN_LEVEL == 8: - kv_7 += tl.dot(b_k, b_v) + kv_7 = tl.dot(b_k, b_v, kv_7) elif MIN_LEVEL == 9: - kv_8 += tl.dot(b_k, b_v) + kv_8 = tl.dot(b_k, b_v, kv_8) elif MIN_LEVEL == 10: - kv_9 += tl.dot(b_k, b_v) + kv_9 = tl.dot(b_k, b_v, kv_9) elif MIN_LEVEL == 11: - kv_10 += tl.dot(b_k, b_v) + kv_10 = tl.dot(b_k, b_v, kv_10) check_value = (~chunk_index & (chunk_index + 1)) - 1 @@ -547,125 +391,41 @@ def chunkwise_fwd_kernel( if STORE_FINAL_STATE: if (MIN_LEVEL <= 0 and MAX_LEVEL >= 0) and (chunk_index & 1 > 0): - p_kv = tl.make_block_ptr( - ht + ((i_n * L_OUT + 0) * H + i_h) * K * V, - (K, V), - (V, 1), - (i_k * BK, 0), - (BK, V), - (1, 0), - ) - tl.store(p_kv, kv_0, boundary_check=(0, 1)) + p_kv = ht + ((i_n * L_OUT + 0) * H + i_h) * K * V + o_kk[:, None] * V + o_v[None, :] + tl.store(p_kv, kv_0, mask=m_kv) if (MIN_LEVEL <= 1 and MAX_LEVEL >= 1) and (chunk_index & 2 > 0): - p_kv = tl.make_block_ptr( - ht + ((i_n * L_OUT + 1) * H + i_h) * K * V, - (K, V), - (V, 1), - (i_k * BK, 0), - (BK, V), - (1, 0), - ) - tl.store(p_kv, kv_1, boundary_check=(0, 1)) + p_kv = ht + ((i_n * L_OUT + 1) * H + i_h) * K * V + o_kk[:, None] * V + o_v[None, :] + tl.store(p_kv, kv_1, mask=m_kv) if (MIN_LEVEL <= 2 and MAX_LEVEL >= 2) and (chunk_index & 4 > 0): - p_kv = tl.make_block_ptr( - ht + ((i_n * L_OUT + 2) * H + i_h) * K * V, - (K, V), - (V, 1), - (i_k * BK, 0), - (BK, V), - (1, 0), - ) - tl.store(p_kv, kv_2, boundary_check=(0, 1)) + p_kv = ht + ((i_n * L_OUT + 2) * H + i_h) * K * V + o_kk[:, None] * V + o_v[None, :] + tl.store(p_kv, kv_2, mask=m_kv) if (MIN_LEVEL <= 3 and MAX_LEVEL >= 3) and (chunk_index & 8 > 0): - p_kv = tl.make_block_ptr( - ht + ((i_n * L_OUT + 3) * H + i_h) * K * V, - (K, V), - (V, 1), - (i_k * BK, 0), - (BK, V), - (1, 0), - ) - tl.store(p_kv, kv_3, boundary_check=(0, 1)) + p_kv = ht + ((i_n * L_OUT + 3) * H + i_h) * K * V + o_kk[:, None] * V + o_v[None, :] + tl.store(p_kv, kv_3, mask=m_kv) if (MIN_LEVEL <= 4 and MAX_LEVEL >= 4) and (chunk_index & 16 > 0): - p_kv = tl.make_block_ptr( - ht + ((i_n * L_OUT + 4) * H + i_h) * K * V, - (K, V), - (V, 1), - (i_k * BK, 0), - (BK, V), - (1, 0), - ) - tl.store(p_kv, kv_4, boundary_check=(0, 1)) + p_kv = ht + ((i_n * L_OUT + 4) * H + i_h) * K * V + o_kk[:, None] * V + o_v[None, :] + tl.store(p_kv, kv_4, mask=m_kv) if (MIN_LEVEL <= 5 and MAX_LEVEL >= 5) and (chunk_index & 32 > 0): - p_kv = tl.make_block_ptr( - ht + ((i_n * L_OUT + 5) * H + i_h) * K * V, - (K, V), - (V, 1), - (i_k * BK, 0), - (BK, V), - (1, 0), - ) - tl.store(p_kv, kv_5, boundary_check=(0, 1)) + p_kv = ht + ((i_n * L_OUT + 5) * H + i_h) * K * V + o_kk[:, None] * V + o_v[None, :] + tl.store(p_kv, kv_5, mask=m_kv) if (MIN_LEVEL <= 6 and MAX_LEVEL >= 6) and (chunk_index & 64 > 0): - p_kv = tl.make_block_ptr( - ht + ((i_n * L_OUT + 6) * H + i_h) * K * V, - (K, V), - (V, 1), - (i_k * BK, 0), - (BK, V), - (1, 0), - ) - tl.store(p_kv, kv_6, boundary_check=(0, 1)) + p_kv = ht + ((i_n * L_OUT + 6) * H + i_h) * K * V + o_kk[:, None] * V + o_v[None, :] + tl.store(p_kv, kv_6, mask=m_kv) if (MIN_LEVEL <= 7 and MAX_LEVEL >= 7) and (chunk_index & 128 > 0): - p_kv = tl.make_block_ptr( - ht + ((i_n * L_OUT + 7) * H + i_h) * K * V, - (K, V), - (V, 1), - (i_k * BK, 0), - (BK, V), - (1, 0), - ) - tl.store(p_kv, kv_7, boundary_check=(0, 1)) + p_kv = ht + ((i_n * L_OUT + 7) * H + i_h) * K * V + o_kk[:, None] * V + o_v[None, :] + tl.store(p_kv, kv_7, mask=m_kv) if (MIN_LEVEL <= 8 and MAX_LEVEL >= 8) and (chunk_index & 256 > 0): - p_kv = tl.make_block_ptr( - ht + ((i_n * L_OUT + 8) * H + i_h) * K * V, - (K, V), - (V, 1), - (i_k * BK, 0), - (BK, V), - (1, 0), - ) - tl.store(p_kv, kv_8, boundary_check=(0, 1)) + p_kv = ht + ((i_n * L_OUT + 8) * H + i_h) * K * V + o_kk[:, None] * V + o_v[None, :] + tl.store(p_kv, kv_8, mask=m_kv) if (MIN_LEVEL <= 9 and MAX_LEVEL >= 9) and (chunk_index & 512 > 0): - p_kv = tl.make_block_ptr( - ht + ((i_n * L_OUT + 9) * H + i_h) * K * V, - (K, V), - (V, 1), - (i_k * BK, 0), - (BK, V), - (1, 0), - ) - tl.store(p_kv, kv_9, boundary_check=(0, 1)) + p_kv = ht + ((i_n * L_OUT + 9) * H + i_h) * K * V + o_kk[:, None] * V + o_v[None, :] + tl.store(p_kv, kv_9, mask=m_kv) if (MIN_LEVEL <= 10 and MAX_LEVEL >= 10) and (chunk_index & 1024 > 0): - p_kv = tl.make_block_ptr( - ht + ((i_n * L_OUT + 10) * H + i_h) * K * V, - (K, V), - (V, 1), - (i_k * BK, 0), - (BK, V), - (1, 0), - ) - tl.store(p_kv, kv_10, boundary_check=(0, 1)) + p_kv = ht + ((i_n * L_OUT + 10) * H + i_h) * K * V + o_kk[:, None] * V + o_v[None, :] + tl.store(p_kv, kv_10, mask=m_kv) if (MIN_LEVEL <= 11 and MAX_LEVEL >= 11) and (chunk_index & 2048 > 0): - p_kv = tl.make_block_ptr( - ht + ((i_n * L_OUT + 11) * H + i_h) * K * V, - (K, V), - (V, 1), - (i_k * BK, 0), - (BK, V), - (1, 0), - ) - tl.store(p_kv, kv_11, boundary_check=(0, 1)) + p_kv = ht + ((i_n * L_OUT + 11) * H + i_h) * K * V + o_kk[:, None] * V + o_v[None, :] + tl.store(p_kv, kv_11, mask=m_kv) tl.store(new_offsets + i_n, (offset // BT) * BT + T) @@ -673,7 +433,7 @@ def chunkwise_fwd_kernel( @triton.heuristics({ "IS_VARLEN": lambda args: args["cu_seqlens"] is not None, }) -@triton.jit(do_not_specialize=["T"]) +@triton.jit(do_not_specialize=["T", "input_T"]) def copy_input_kernel( q, k, @@ -681,6 +441,7 @@ def copy_input_kernel( g, level_scales, cu_seqlens, + input_cu_seqlens, q_prev, k_prev, v_prev, @@ -693,6 +454,8 @@ def copy_input_kernel( g_new, level_scales_new, T, + input_T, + G: tl.constexpr, H: tl.constexpr, K: tl.constexpr, V: tl.constexpr, @@ -701,122 +464,81 @@ def copy_input_kernel( IS_VARLEN: tl.constexpr, ): # parallel over sequences and heads - i_nh = tl.program_id(0) + i_nh = tl.program_id(0).to(tl.int64) i_n, i_h = i_nh // H, i_nh % H + i_g = i_h // (H // G) if IS_VARLEN: bos, eos = ( tl.load(cu_seqlens + i_n).to(tl.int64), tl.load(cu_seqlens + i_n + 1).to(tl.int64), ) + input_bos, input_eos = ( + tl.load(input_cu_seqlens + i_n).to(tl.int64), + tl.load(input_cu_seqlens + i_n + 1).to(tl.int64), + ) T = (eos - bos).to(tl.int32) + input_T = (input_eos - input_bos).to(tl.int32) else: - bos, eos = (i_n * T).to(tl.int64), (i_n * T + T).to(tl.int64) + bos = i_n.to(tl.int64) * T + eos = bos + T + input_bos = i_n.to(tl.int64) * input_T offset = tl.load(offsets + i_n) input_offset = -1 * (offset % BT) NT = tl.cdiv(T, BT) + o_k = tl.arange(0, K) + o_v = tl.arange(0, V) + o_i = tl.arange(0, BT) for i_t in range(NT): - p_g = tl.make_block_ptr( - g + bos * H + i_h, (T,), (H,), (i_t * BT + input_offset,), (BT,), (0,), - ) - p_q = tl.make_block_ptr( - q + bos * K, (T, K), (K, 1), (i_t * BT + input_offset, 0), (BT, K), (1, 0), - ) - p_k = tl.make_block_ptr( - k + bos * K, (T, K), (K, 1), (i_t * BT + input_offset, 0), (BT, K), (1, 0), - ) - p_v = tl.make_block_ptr( - v + (bos * H + i_h) * V, - (T, V), - (H * V, 1), - (i_t * BT + input_offset, 0), - (BT, V), - (1, 0), - ) - p_g_new = tl.make_block_ptr( - g_new + bos * H + i_h, (T,), (H,), (i_t * BT,), (BT,), (0,), - ) - p_q_new = tl.make_block_ptr( - q_new + bos * K, (T, K), (K, 1), (i_t * BT, 0), (BT, K), (1, 0), - ) - p_k_new = tl.make_block_ptr( - k_new + bos * K, (T, K), (K, 1), (i_t * BT, 0), (BT, K), (1, 0), - ) - p_v_new = tl.make_block_ptr( - v_new + (bos * H + i_h) * V, - (T, V), - (H * V, 1), - (i_t * BT, 0), - (BT, V), - (1, 0), - ) - - b_g = tl.load(p_g, boundary_check=(0,)) - b_q = tl.load(p_q, boundary_check=(0, 1)) - b_k = tl.load(p_k, boundary_check=(0, 1)) - b_v = tl.load(p_v, boundary_check=(0, 1)) + o_t = (i_t * BT + input_offset).to(tl.int64) + o_i + o_tn = (i_t * BT).to(tl.int64) + o_i + m_t = (o_t >= 0) & (o_t < input_T) + m_tn = o_tn < T + m_tk = m_t[:, None] & (o_k[None, :] < K) + m_tv = m_t[:, None] & (o_v[None, :] < V) + m_tnk = m_tn[:, None] & (o_k[None, :] < K) + m_tnv = m_tn[:, None] & (o_v[None, :] < V) + p_g = g + input_bos * H + i_h + o_t * H + p_q = q + (input_bos * G + i_g) * K + o_t[:, None] * (G * K) + o_k[None, :] + p_k = k + (input_bos * G + i_g) * K + o_t[:, None] * (G * K) + o_k[None, :] + p_v = v + (input_bos * H + i_h) * V + o_t[:, None] * (H * V) + o_v[None, :] + p_g_new = g_new + bos * H + i_h + o_tn * H + p_q_new = q_new + (bos * G + i_g) * K + o_tn[:, None] * (G * K) + o_k[None, :] + p_k_new = k_new + (bos * G + i_g) * K + o_tn[:, None] * (G * K) + o_k[None, :] + p_v_new = v_new + (bos * H + i_h) * V + o_tn[:, None] * (H * V) + o_v[None, :] + + b_g = tl.load(p_g, mask=m_t, other=0.0) + b_q = tl.load(p_q, mask=m_tk, other=0.0) + b_k = tl.load(p_k, mask=m_tk, other=0.0) + b_v = tl.load(p_v, mask=m_tv, other=0.0) if i_t == 0: - p_g_prev = tl.make_block_ptr( - g_prev + i_n * BT * H + i_h, (BT,), (H,), (0,), (BT,), (0,), - ) - p_q_prev = tl.make_block_ptr( - q_prev + i_n * BT * K, (BT, K), (K, 1), (0, 0), (BT, K), (1, 0), - ) - p_k_prev = tl.make_block_ptr( - k_prev + i_n * BT * K, (BT, K), (K, 1), (0, 0), (BT, K), (1, 0), - ) - p_v_prev = tl.make_block_ptr( - v_prev + (i_n * BT * H + i_h) * V, - (BT, V), - (H * V, 1), - (0, 0), - (BT, V), - (1, 0), - ) + p_g_prev = g_prev + i_n * BT * H + i_h + o_i * H + p_q_prev = q_prev + (i_n * BT * G + i_g) * K + o_i[:, None] * (G * K) + o_k[None, :] + p_k_prev = k_prev + (i_n * BT * G + i_g) * K + o_i[:, None] * (G * K) + o_k[None, :] + p_v_prev = v_prev + (i_n * BT * H + i_h) * V + o_i[:, None] * (H * V) + o_v[None, :] - b_g += tl.load(p_g_prev, boundary_check=(0,)) - b_q += tl.load(p_q_prev, boundary_check=(0, 1)) - b_k += tl.load(p_k_prev, boundary_check=(0, 1)) - b_v += tl.load(p_v_prev, boundary_check=(0, 1)) + b_g += tl.load(p_g_prev, mask=o_i < BT, other=0.0) + b_q += tl.load(p_q_prev, mask=(o_i[:, None] < BT) & (o_k[None, :] < K), other=0.0) + b_k += tl.load(p_k_prev, mask=(o_i[:, None] < BT) & (o_k[None, :] < K), other=0.0) + b_v += tl.load(p_v_prev, mask=(o_i[:, None] < BT) & (o_v[None, :] < V), other=0.0) - tl.store(p_g_new, b_g, boundary_check=(0,)) - tl.store(p_q_new, b_q, boundary_check=(0, 1)) - tl.store(p_k_new, b_k, boundary_check=(0, 1)) - tl.store(p_v_new, b_v, boundary_check=(0, 1)) + tl.store(p_g_new, b_g, mask=m_tn) + tl.store(p_q_new, b_q, mask=m_tnk & (i_h % (H // G) == 0)) + tl.store(p_k_new, b_k, mask=m_tnk & (i_h % (H // G) == 0)) + tl.store(p_v_new, b_v, mask=m_tnv) for i in range(L): - p_l = tl.make_block_ptr( - level_scales + (bos * H + i_h) * L, - (T, L), - (H * L, 1), - (i_t * BT + input_offset, i), - (BT, 1), - (1, 0), - ) - p_l_new = tl.make_block_ptr( - level_scales_new + (bos * H + i_h) * L, - (T, L), - (H * L, 1), - (i_t * BT, i), - (BT, 1), - (1, 0), - ) - b_l = tl.load(p_l, boundary_check=(0,)) + p_l = level_scales + (input_bos * H + i_h) * L + o_t[:, None] * (H * L) + i + p_l_new = level_scales_new + (bos * H + i_h) * L + o_tn[:, None] * (H * L) + i + b_l = tl.load(p_l, mask=m_t[:, None], other=0.0) if i_t == 0: - p_l_prev = tl.make_block_ptr( - level_scales_prev + (i_n * BT * H + i_h) * L, - (BT, L), - (H * L, 1), - (0, i), - (BT, 1), - (1, 0), - ) - b_l += tl.load(p_l_prev, boundary_check=(0,)) - tl.store(p_l_new, b_l, boundary_check=(0,)) + p_l_prev = level_scales_prev + (i_n * BT * H + i_h) * L + o_i[:, None] * (H * L) + i + b_l += tl.load(p_l_prev, mask=o_i[:, None] < BT, other=0.0) + tl.store(p_l_new, b_l, mask=m_tn[:, None]) @triton.heuristics( @@ -839,6 +561,7 @@ def copy_last_chunk_kernel( level_scales_prev, offsets, T, + G: tl.constexpr, H: tl.constexpr, K: tl.constexpr, V: tl.constexpr, @@ -847,8 +570,9 @@ def copy_last_chunk_kernel( IS_VARLEN: tl.constexpr, ): # parallel over sequences and heads - i_nh = tl.program_id(0) + i_nh = tl.program_id(0).to(tl.int64) i_n, i_h = i_nh // H, i_nh % H + i_g = i_h // (H // G) if IS_VARLEN: bos, eos = ( @@ -861,57 +585,33 @@ def copy_last_chunk_kernel( seq_offset = (T // BT) * BT - p_g = tl.make_block_ptr(g + bos * H + i_h, (T,), (H,), (seq_offset,), (BT,), (0,)) - p_q = tl.make_block_ptr( - q + bos * K, (T, K), (K, 1), (seq_offset, 0), (BT, K), (1, 0), - ) - p_k = tl.make_block_ptr( - k + bos * K, (T, K), (K, 1), (seq_offset, 0), (BT, K), (1, 0), - ) - p_v = tl.make_block_ptr( - v + (bos * H + i_h) * V, - (T, V), - (H * V, 1), - (seq_offset, 0), - (BT, V), - (1, 0), - ) - p_g_prev = tl.make_block_ptr( - g_prev + i_n * BT * H + i_h, (BT,), (H,), (0,), (BT,), (0,), - ) - p_q_prev = tl.make_block_ptr( - q_prev + i_n * BT * K, (BT, K), (K, 1), (0, 0), (BT, K), (1, 0), - ) - p_k_prev = tl.make_block_ptr( - k_prev + i_n * BT * K, (BT, K), (K, 1), (0, 0), (BT, K), (1, 0), - ) - p_v_prev = tl.make_block_ptr( - v_prev + (i_n * BT * H + i_h) * V, (BT, V), (H * V, 1), (0, 0), (BT, V), (1, 0), - ) - - tl.store(p_g_prev, tl.load(p_g, boundary_check=(0,)), boundary_check=(0,)) - tl.store(p_q_prev, tl.load(p_q, boundary_check=(0, 1)), boundary_check=(0, 1)) - tl.store(p_k_prev, tl.load(p_k, boundary_check=(0, 1)), boundary_check=(0, 1)) - tl.store(p_v_prev, tl.load(p_v, boundary_check=(0, 1)), boundary_check=(0, 1)) + o_k = tl.arange(0, K) + o_v = tl.arange(0, V) + o_i = tl.arange(0, BT) + o_t = seq_offset.to(tl.int64) + o_i + m_t = o_t < T + m_tk = m_t[:, None] & (o_k[None, :] < K) + m_tv = m_t[:, None] & (o_v[None, :] < V) + m_ik = (o_i[:, None] < BT) & (o_k[None, :] < K) + m_iv = (o_i[:, None] < BT) & (o_v[None, :] < V) + p_g = g + bos * H + i_h + o_t * H + p_q = q + (bos * G + i_g) * K + o_t[:, None] * (G * K) + o_k[None, :] + p_k = k + (bos * G + i_g) * K + o_t[:, None] * (G * K) + o_k[None, :] + p_v = v + (bos * H + i_h) * V + o_t[:, None] * (H * V) + o_v[None, :] + p_g_prev = g_prev + i_n * BT * H + i_h + o_i * H + p_q_prev = q_prev + (i_n * BT * G + i_g) * K + o_i[:, None] * (G * K) + o_k[None, :] + p_k_prev = k_prev + (i_n * BT * G + i_g) * K + o_i[:, None] * (G * K) + o_k[None, :] + p_v_prev = v_prev + (i_n * BT * H + i_h) * V + o_i[:, None] * (H * V) + o_v[None, :] + + tl.store(p_g_prev, tl.load(p_g, mask=m_t, other=0.0), mask=o_i < BT) + tl.store(p_q_prev, tl.load(p_q, mask=m_tk, other=0.0), mask=m_ik & (i_h % (H // G) == 0)) + tl.store(p_k_prev, tl.load(p_k, mask=m_tk, other=0.0), mask=m_ik & (i_h % (H // G) == 0)) + tl.store(p_v_prev, tl.load(p_v, mask=m_tv, other=0.0), mask=m_iv) for i in range(L): - p_l = tl.make_block_ptr( - level_scales + (bos * H + i_h) * L, - (T, L), - (H * L, 1), - (seq_offset, i), - (BT, 1), - (1, 0), - ) - p_l_prev = tl.make_block_ptr( - level_scales_prev + (i_n * BT * H + i_h) * L, - (BT, L), - (H * L, 1), - (0, i), - (BT, 1), - (1, 0), - ) - tl.store(p_l_prev, tl.load(p_l, boundary_check=(0,)), boundary_check=(0,)) + p_l = level_scales + (bos * H + i_h) * L + o_t[:, None] * (H * L) + i + p_l_prev = level_scales_prev + (i_n * BT * H + i_h) * L + o_i[:, None] * (H * L) + i + tl.store(p_l_prev, tl.load(p_l, mask=m_t[:, None], other=0.0), mask=o_i[:, None] < BT) @triton.heuristics({"IS_VARLEN": lambda args: args["cu_seqlens"] is not None}) @@ -922,7 +622,7 @@ def copy_last_chunk_kernel( for num_warps in [4] for num_stages in [2, 3, 4] ], - key=["H", "K", "V"], + key=["G", "H", "K", "V"], restore_value=["dh", "dg_last"], **autotune_cache_kwargs, ) @@ -938,6 +638,7 @@ def chunkwise_bwd_kernel_dhg( ell, T, cu_seqlens, + G: tl.constexpr, H: tl.constexpr, K: tl.constexpr, V: tl.constexpr, @@ -946,11 +647,12 @@ def chunkwise_bwd_kernel_dhg( BK: tl.constexpr, NT: tl.constexpr, IS_VARLEN: tl.constexpr, + SCALE: tl.constexpr, ): # parallel over batches and heads - i_k = tl.program_id(0) - i_nh = tl.program_id(1) + i_k, i_nh = unflatten_program_id(tl.cdiv(K, BK)) i_n, i_h = i_nh // H, i_nh % H + i_g = i_h // (H // G) if IS_VARLEN: bos, eos = ( @@ -965,68 +667,44 @@ def chunkwise_bwd_kernel_dhg( num_intra_levels = (tl.log2(float(BT))).to(tl.int32) + 1 + o_kk = i_k * BK + tl.arange(0, BK) + o_v = tl.arange(0, V) + o_i = tl.arange(0, BT) + m_kv = (o_kk[:, None] < K) & (o_v[None, :] < V) for i_t in range(tl.cdiv(T, BT) - 1, -1, -1): - p_dh = tl.make_block_ptr( - dh + ((i_n * NT + i_t) * H + i_h) * K * V, - (K, V), - (V, 1), - (i_k * BK, 0), - (BK, V), - (1, 0), - ) - b_dh_old = tl.load(p_dh, boundary_check=(0, 1)) + o_t = (i_t * BT).to(tl.int64) + o_i + m_t = o_t < T + m_kt = (o_kk[:, None] < K) & m_t[None, :] + m_tv = m_t[:, None] & (o_v[None, :] < V) + p_dh = dh + ((i_n * NT + i_t) * H + i_h) * K * V + o_kk[:, None] * V + o_v[None, :] + b_dh_old = tl.load(p_dh, mask=m_kv, other=0.0) if (i_t & (1 << ell)) == 0: # store the chunk tl.store( - p_dh, b_dh.to(p_dh.dtype.element_ty) + b_dh_old, boundary_check=(0, 1), + p_dh, b_dh.to(p_dh.dtype.element_ty) + b_dh_old, mask=m_kv, ) # if you are about the transition to compute, reset to zeros if i_t > 0 and ((i_t - 1) & (1 << ell)) > 0: b_dh = tl.zeros([BK, V], dtype=tl.float32) if i_t & (1 << ell): - p_h = tl.make_block_ptr( - h_l + ((i_n * NT + i_t) * H + i_h) * K * V, - (K, V), - (V, 1), - (i_k * BK, 0), - (BK, V), - (1, 0), - ) + p_h = h_l + ((i_n * NT + i_t) * H + i_h) * K * V + o_kk[:, None] * V + o_v[None, :] - b_h = tl.load(p_h, boundary_check=(0, 1)) + b_h = tl.load(p_h, mask=m_kv, other=0.0) p_dg_last = dg_last + i_n * NT * H + i_t * H + i_h tl.atomic_add(p_dg_last, tl.sum(b_h * (b_dh + b_dh_old))) last_idx = min((i_t + 1) * BT, T) - 1 b_g_last = tl.exp(tl.load(g + bos * H + last_idx * H + i_h)) b_dh *= b_g_last if i_t & (1 << ell): # compute this chunk - p_g = tl.make_block_ptr( - g + bos * H + i_h, (T,), (H,), (i_t * BT,), (BT,), (0,), - ) - p_q = tl.make_block_ptr( - q + bos * K, (K, T), (1, K), (i_k * BK, i_t * BT), (BK, BT), (0, 1), - ) - p_do = tl.make_block_ptr( - do + (bos * H + i_h) * V, - (T, V), - (H * V, 1), - (i_t * BT, 0), - (BT, V), - (1, 0), - ) - p_l = tl.make_block_ptr( - l + (bos * H + i_h) * L + num_intra_levels + ell, - (T,), - (H * L,), - (i_t * BT,), - (BT,), - (0,), - ) - b_l = tl.load(p_l, boundary_check=(0,)) - b_g = tl.load(p_g, boundary_check=(0,)) - b_q = tl.load(p_q, boundary_check=(0, 1)) + p_g = g + bos * H + i_h + o_t * H + p_q = q + (bos * G + i_g) * K + o_kk[:, None] + o_t[None, :] * (G * K) + p_do = do + (bos * H + i_h) * V + o_t[:, None] * (H * V) + o_v[None, :] + p_l = l + (bos * H + i_h) * L + num_intra_levels + ell + o_t * (H * L) + b_l = tl.load(p_l, mask=m_t, other=0.0) + b_g = tl.load(p_g, mask=m_t, other=0.0) + b_q = tl.load(p_q, mask=m_kt, other=0.0) b_q = (b_q * (tl.exp(b_g) * b_l)[None, :]).to(b_q.dtype) - b_do = tl.load(p_do, boundary_check=(0, 1)) + b_do = (tl.load(p_do, mask=m_tv, other=0.0) * SCALE).to(do.dtype.element_ty) b_s = tl.dot(b_q, b_do).to(b_q.dtype) b_dh += b_s @@ -1039,7 +717,7 @@ def chunkwise_bwd_kernel_dhg( for num_warps in [4] for num_stages in [2, 3, 4] ], - key=["H", "K", "V"], + key=["G", "H", "K", "V"], restore_value=["dq", "dg"], **autotune_cache_kwargs, ) @@ -1058,6 +736,7 @@ def chunkwise_bwd_kernel_hdqgl( ell, T, cu_seqlens, + G: tl.constexpr, H: tl.constexpr, K: tl.constexpr, V: tl.constexpr, @@ -1065,10 +744,12 @@ def chunkwise_bwd_kernel_hdqgl( BT: tl.constexpr, NT: tl.constexpr, IS_VARLEN: tl.constexpr, + SCALE: tl.constexpr, ): # parallel over batches and heads - i_nh = tl.program_id(0) + i_nh = tl.program_id(0).to(tl.int64) i_n, i_h = i_nh // H, i_nh % H + i_g = i_h // (H // G) if IS_VARLEN: bos, eos = ( @@ -1083,77 +764,46 @@ def chunkwise_bwd_kernel_hdqgl( num_intra_levels = (tl.log2(float(BT))).to(tl.int32) + 1 + o_k = tl.arange(0, K) + o_v = tl.arange(0, V) + o_i = tl.arange(0, BT) + m_vk = (o_v[:, None] < V) & (o_k[None, :] < K) for i_t in range(tl.cdiv(T, BT)): - p_g = tl.make_block_ptr(g + bos * H + i_h, (T,), (H,), (i_t * BT,), (BT,), (0,)) + o_t = (i_t * BT).to(tl.int64) + o_i + m_t = o_t < T + m_tk = m_t[:, None] & (o_k[None, :] < K) + m_tv = m_t[:, None] & (o_v[None, :] < V) + p_g = g + bos * H + i_h + o_t * H if i_t & (1 << ell): # compute and store derivatives - p_do = tl.make_block_ptr( - do + (bos * H + i_h) * V, - (T, V), - (H * V, 1), - (i_t * BT, 0), - (BT, V), - (1, 0), - ) - p_q = tl.make_block_ptr( - q + bos * K, (T, K), (K, 1), (i_t * BT, 0), (BT, K), (1, 0), - ) - p_l = tl.make_block_ptr( - l + (bos * H + i_h) * L + num_intra_levels + ell, - (T,), - (H * L,), - (i_t * BT,), - (BT,), - (0,), - ) - p_dq = tl.make_block_ptr( - dq + (bos * H + i_h) * K, - (T, K), - (H * K, 1), - (i_t * BT, 0), - (BT, K), - (1, 0), - ) - p_dg = tl.make_block_ptr( - dg + bos * H + i_h, (T,), (H,), (i_t * BT,), (BT,), (0,), - ) - p_dl = tl.make_block_ptr( - dl + (bos * H + i_h) * L + num_intra_levels + ell, - (T,), - (H * L,), - (i_t * BT,), - (BT,), - (0,), - ) - p_h = tl.make_block_ptr( - h_l + ((i_n * NT + i_t) * H + i_h) * K * V, - (V, K), - (1, V), - (0, 0), - (V, K), - (0, 1), - ) - - b_do = tl.load(p_do, boundary_check=(0, 1)) - b_q = tl.load(p_q, boundary_check=(0, 1)) - b_g = tl.load(p_g, boundary_check=(0,)) - b_l = tl.load(p_l, boundary_check=(0,)) + p_do = do + (bos * H + i_h) * V + o_t[:, None] * (H * V) + o_v[None, :] + p_q = q + (bos * G + i_g) * K + o_t[:, None] * (G * K) + o_k[None, :] + p_l = l + (bos * H + i_h) * L + num_intra_levels + ell + o_t * (H * L) + p_dq = dq + (bos * H + i_h) * K + o_t[:, None] * (H * K) + o_k[None, :] + p_dg = dg + bos * H + i_h + o_t * H + p_dl = dl + (bos * H + i_h) * L + num_intra_levels + ell + o_t * (H * L) + p_h = h_l + ((i_n * NT + i_t) * H + i_h) * K * V + o_v[:, None] + o_k[None, :] * V + + b_do = (tl.load(p_do, mask=m_tv, other=0.0) * SCALE).to(do.dtype.element_ty) + b_q = tl.load(p_q, mask=m_tk, other=0.0) + b_g = tl.load(p_g, mask=m_t, other=0.0) + b_l = tl.load(p_l, mask=m_t, other=0.0) b_dlq = tl.exp(b_g)[:, None] * tl.dot(b_do, b_h.to(b_do.dtype)) b_dl = tl.sum(b_dlq * b_q, axis=1) b_dg = b_l * b_dl - tl.store(p_h, b_h, boundary_check=(0, 1)) - b_dq_old = tl.load(p_dq, boundary_check=(0, 1)) + tl.store(p_h, b_h, mask=m_vk) + b_dq_old = tl.load(p_dq, mask=m_tk, other=0.0) tl.store( p_dq, (b_l[:, None] * b_dlq).to(p_dq.dtype.element_ty) + b_dq_old, - boundary_check=(0, 1), + mask=m_tk, ) - tl.store(p_dl, b_dl.to(p_dl.dtype.element_ty), boundary_check=(0,)) - b_dg_old = tl.load(p_dg, boundary_check=(0,)) + tl.store(p_dl, b_dl.to(p_dl.dtype.element_ty), mask=m_t) + b_dg_old = tl.load(p_dg, mask=m_t, other=0.0) tl.store( - p_dg, b_dg.to(p_dg.dtype.element_ty) + b_dg_old, boundary_check=(0,), + p_dg, b_dg.to(p_dg.dtype.element_ty) + b_dg_old, mask=m_t, ) if ((i_t + 1) & (1 << ell)) == 0: b_h = tl.zeros([V, K], dtype=tl.float32) @@ -1162,22 +812,13 @@ def chunkwise_bwd_kernel_hdqgl( b_g_last = tl.load(g + bos * H + last_idx * H + i_h) b_h *= tl.exp(b_g_last) if (i_t & (1 << ell)) == 0: # update the state - p_k = tl.make_block_ptr( - k + bos * K, (T, K), (K, 1), (i_t * BT, 0), (BT, K), (1, 0), - ) - p_v = tl.make_block_ptr( - v + (bos * H + i_h) * V, - (V, T), - (1, H * V), - (0, i_t * BT), - (V, BT), - (0, 1), - ) - b_g = tl.load(p_g, boundary_check=(0,)) - b_k = tl.load(p_k, boundary_check=(0, 1)) - b_v = tl.load(p_v, boundary_check=(0, 1)) + p_k = k + (bos * G + i_g) * K + o_t[:, None] * (G * K) + o_k[None, :] + p_v = v + (bos * H + i_h) * V + o_v[:, None] + o_t[None, :] * (H * V) + b_g = tl.load(p_g, mask=m_t, other=0.0) + b_k = tl.load(p_k, mask=m_tk, other=0.0) + b_v = tl.load(p_v, mask=(o_v[:, None] < V) & m_t[None, :], other=0.0) b_k = (b_k * tl.exp(b_g_last - b_g)[:, None]).to(b_k.dtype) - b_h += tl.dot(b_v, b_k) + b_h = tl.dot(b_v, b_k, b_h) @triton.heuristics({"IS_VARLEN": lambda args: args["cu_seqlens"] is not None}) @@ -1187,7 +828,7 @@ def chunkwise_bwd_kernel_hdqgl( for num_warps in [4] for num_stages in [2, 3, 4] ], - key=["H", "K", "V"], + key=["G", "H", "K", "V"], restore_value=["dk", "dg", "dg_last"], **autotune_cache_kwargs, ) @@ -1202,6 +843,7 @@ def chunkwise_bwd_kernel_dkg( dg, cu_seqlens, T, + G: tl.constexpr, H: tl.constexpr, K: tl.constexpr, V: tl.constexpr, @@ -1210,8 +852,9 @@ def chunkwise_bwd_kernel_dkg( NT: tl.constexpr, IS_VARLEN: tl.constexpr, ): - i_t, i_nh = tl.program_id(0), tl.program_id(1) + i_t, i_nh = unflatten_program_id(NT) i_n, i_h = i_nh // H, i_nh % H + i_g = i_h // (H // G) if IS_VARLEN: bos, eos = ( @@ -1226,33 +869,22 @@ def chunkwise_bwd_kernel_dkg( o_t = i_t * BT + o_i m_t = o_t < T - p_dh = tl.make_block_ptr( - dh + ((i_n * NT + i_t) * H + i_h) * K * V, - (V, K), - (1, V), - (0, 0), - (V, K), - (0, 1), - ) - p_g = tl.make_block_ptr(g + bos * H + i_h, (T,), (H,), (i_t * BT,), (BT,), (0,)) - p_k = tl.make_block_ptr(k + bos * K, (T, K), (K, 1), (i_t * BT, 0), (BT, K), (1, 0)) - p_v = tl.make_block_ptr( - v + (bos * H + i_h) * V, - (T, V), - (H * V, 1), - (i_t * BT, 0), - (BT, V), - (1, 0), - ) - p_dk = tl.make_block_ptr( - dk + (bos * H + i_h) * K, (T, K), (H * K, 1), (i_t * BT, 0), (BT, K), (1, 0), - ) - p_dg = tl.make_block_ptr(dg + bos * H + i_h, (T,), (H,), (i_t * BT,), (BT,), (0,)) - - b_dh = tl.load(p_dh, boundary_check=(0, 1)) - b_g = tl.load(p_g, boundary_check=(0,)) - b_v = tl.load(p_v, boundary_check=(0, 1)) - b_k = tl.load(p_k, boundary_check=(0, 1)) + o_k = tl.arange(0, K) + o_v = tl.arange(0, V) + m_vk = (o_v[:, None] < V) & (o_k[None, :] < K) + m_tk = m_t[:, None] & (o_k[None, :] < K) + m_tv = m_t[:, None] & (o_v[None, :] < V) + p_dh = dh + ((i_n * NT + i_t) * H + i_h) * K * V + o_v[:, None] + o_k[None, :] * V + p_g = g + bos * H + i_h + o_t * H + p_k = k + (bos * G + i_g) * K + o_t[:, None] * (G * K) + o_k[None, :] + p_v = v + (bos * H + i_h) * V + o_t[:, None] * (H * V) + o_v[None, :] + p_dk = dk + (bos * H + i_h) * K + o_t[:, None] * (H * K) + o_k[None, :] + p_dg = dg + bos * H + i_h + o_t * H + + b_dh = tl.load(p_dh, mask=m_vk, other=0.0) + b_g = tl.load(p_g, mask=m_t, other=0.0) + b_v = tl.load(p_v, mask=m_tv, other=0.0) + b_k = tl.load(p_k, mask=m_tk, other=0.0) last_idx = min((i_t + 1) * BT, T) - 1 b_g_last = tl.load(g + bos * H + last_idx * H + i_h) p_dg_last = dg_last + i_n * NT * H + i_t * H + i_h @@ -1260,14 +892,16 @@ def chunkwise_bwd_kernel_dkg( b_dg_last *= tl.exp(b_g_last) b_dk = tl.where(m_t, exp(b_g_last - b_g), 0)[:, None] * tl.dot(b_v, b_dh).to(b_v.dtype) - b_dg = tl.load(p_dg, boundary_check=(0,)) + b_dg = tl.load(p_dg, mask=m_t, other=0.0) b_dg -= tl.sum(b_k * b_dk, axis=1) b_dg_last += tl.sum(b_dk * b_k) - b_dg = tl.where(o_i < BT - 1, b_dg, b_dg + b_dg_last) + # deposit the chunk-scalar gradient on the last LIVE row, the one `last_idx` selected the + # gate from; on a partial last chunk row BT - 1 is dead and the masked store drops it + b_dg = tl.where(o_t < last_idx, b_dg, b_dg + b_dg_last) - tl.store(p_dg, b_dg, boundary_check=(0,)) - tl.store(p_dk, b_dk.to(p_dk.dtype.element_ty), boundary_check=(0, 1)) + tl.store(p_dg, b_dg, mask=m_t) + tl.store(p_dk, b_dk.to(p_dk.dtype.element_ty), mask=m_tk) @triton.heuristics({"IS_VARLEN": lambda args: args["cu_seqlens"] is not None}) @@ -1277,7 +911,7 @@ def chunkwise_bwd_kernel_dkg( for num_warps in [4] for num_stages in [2, 3, 4] ], - key=["H", "K", "V"], + key=["G", "H", "K", "V"], restore_value=["dv"], **autotune_cache_kwargs, ) @@ -1289,6 +923,7 @@ def chunkwise_bwd_kernel_dv( dv, T, cu_seqlens, + G: tl.constexpr, H: tl.constexpr, K: tl.constexpr, V: tl.constexpr, @@ -1297,8 +932,9 @@ def chunkwise_bwd_kernel_dv( NT: tl.constexpr, IS_VARLEN: tl.constexpr, ): - i_t, i_nh = tl.program_id(0), tl.program_id(1) + i_t, i_nh = unflatten_program_id(NT) i_n, i_h = i_nh // H, i_nh % H + i_g = i_h // (H // G) if IS_VARLEN: bos, eos = ( @@ -1312,28 +948,24 @@ def chunkwise_bwd_kernel_dv( o_t = i_t * BT + tl.arange(0, BT) m_t = o_t < T - p_dh = tl.make_block_ptr( - dh + ((i_n * NT + i_t) * H + i_h) * K * V, - (K, V), - (V, 1), - (0, 0), - (K, V), - (1, 0), - ) - p_g = tl.make_block_ptr(g + bos * H + i_h, (T,), (H,), (i_t * BT,), (BT,), (0,)) - p_k = tl.make_block_ptr(k + bos * K, (T, K), (K, 1), (i_t * BT, 0), (BT, K), (1, 0)) - p_dv = tl.make_block_ptr( - dv + (bos * H + i_h) * V, (T, V), (H * V, 1), (i_t * BT, 0), (BT, V), (1, 0), - ) + o_k = tl.arange(0, K) + o_v = tl.arange(0, V) + m_kv = (o_k[:, None] < K) & (o_v[None, :] < V) + m_tk = m_t[:, None] & (o_k[None, :] < K) + m_tv = m_t[:, None] & (o_v[None, :] < V) + p_dh = dh + ((i_n * NT + i_t) * H + i_h) * K * V + o_k[:, None] * V + o_v[None, :] + p_g = g + bos * H + i_h + o_t * H + p_k = k + (bos * G + i_g) * K + o_t[:, None] * (G * K) + o_k[None, :] + p_dv = dv + (bos * H + i_h) * V + o_t[:, None] * (H * V) + o_v[None, :] last_idx = min((i_t + 1) * BT, T) - 1 b_g_last = tl.load(g + bos * H + last_idx * H + i_h) - b_dh = tl.load(p_dh, boundary_check=(0, 1)) - b_g = tl.load(p_g, boundary_check=(0,)) - b_k = tl.load(p_k, boundary_check=(0, 1)) + b_dh = tl.load(p_dh, mask=m_kv, other=0.0) + b_g = tl.load(p_g, mask=m_t, other=0.0) + b_k = tl.load(p_k, mask=m_tk, other=0.0) b_dv = tl.where(m_t, exp(-b_g + b_g_last), 0)[:, None] * tl.dot(b_k, b_dh).to(b_k.dtype) - tl.store(p_dv, b_dv.to(p_dv.dtype.element_ty), boundary_check=(0, 1)) + tl.store(p_dv, b_dv.to(p_dv.dtype.element_ty), mask=m_tv) @triton.heuristics({"IS_VARLEN": lambda args: args["cu_seqlens"] is not None}) @@ -1343,7 +975,7 @@ def chunkwise_bwd_kernel_dv( for num_warps in [4] for num_stages in [2, 3, 4] ], - key=["H", "K", "V"], + key=["G", "H", "K", "V"], restore_value=["dl", "dq", "dk", "dv", "dg"], **autotune_cache_kwargs, ) @@ -1364,17 +996,22 @@ def chunkwise_bwd_kernel_diag( dl, cu_seqlens, T, + G: tl.constexpr, H: tl.constexpr, K: tl.constexpr, V: tl.constexpr, L: tl.constexpr, BT: tl.constexpr, + NT: tl.constexpr, IS_VARLEN: tl.constexpr, + SCALE: tl.constexpr, ): - p_llut = tl.make_block_ptr(llut, (BT, BT), (BT, 1), (0, 0), (BT, BT), (1, 0)) - b_llut = tl.load(p_llut, boundary_check=(0, 1)) - i_t, i_nh = tl.program_id(0), tl.program_id(1) + o_i = tl.arange(0, BT) + p_llut = llut + o_i[:, None] * BT + o_i[None, :] + b_llut = tl.load(p_llut, mask=(o_i[:, None] < BT) & (o_i[None, :] < BT), other=0.0) + i_t, i_nh = unflatten_program_id(NT) i_n, i_h = i_nh // H, i_nh % H + i_g = i_h // (H // G) if IS_VARLEN: bos, eos = ( @@ -1385,42 +1022,39 @@ def chunkwise_bwd_kernel_diag( else: bos, eos = (i_n * T).to(tl.int64), (i_n * T + T).to(tl.int64) - o_i = tl.arange(0, BT) + o_k = tl.arange(0, K) + o_v = tl.arange(0, V) + o_t = i_t * BT + o_i + m_t = o_t < T + m_kt = (o_k[:, None] < K) & m_t[None, :] + m_tk = m_t[:, None] & (o_k[None, :] < K) + m_tv = m_t[:, None] & (o_v[None, :] < V) + m_vt = (o_v[:, None] < V) & m_t[None, :] i_idx = o_i[:, None] # BT x 1 j_idx = o_i[None, :] # 1 x BT - b_h_ptrs = l + ((bos + i_t * BT + i_idx) * H + i_h) * L + b_llut - b_h = tl.load(b_h_ptrs, mask=i_idx >= j_idx) - - p_g = tl.make_block_ptr(g + bos * H + i_h, (T,), (H,), (i_t * BT,), (BT,), (0,)) - p_q = tl.make_block_ptr(q + bos * K, (K, T), (1, K), (0, i_t * BT), (K, BT), (0, 1)) - p_k = tl.make_block_ptr(k + bos * K, (T, K), (K, 1), (i_t * BT, 0), (BT, K), (1, 0)) - p_v = tl.make_block_ptr( - v + (bos * H + i_h) * V, (V, T), (1, H * V), (0, i_t * BT), (V, BT), (0, 1), - ) - p_do = tl.make_block_ptr( - do + (bos * H + i_h) * V, (T, V), (H * V, 1), (i_t * BT, 0), (BT, V), (1, 0), - ) - p_dg = tl.make_block_ptr(dg + bos * H + i_h, (T,), (H,), (i_t * BT,), (BT,), (0,)) - p_dq = tl.make_block_ptr( - dq + (bos * H + i_h) * K, (T, K), (H * K, 1), (i_t * BT, 0), (BT, K), (1, 0), - ) - p_dk = tl.make_block_ptr( - dk + (bos * H + i_h) * K, (T, K), (H * K, 1), (i_t * BT, 0), (BT, K), (1, 0), - ) - p_dv = tl.make_block_ptr( - dv + (bos * H + i_h) * V, (T, V), (H * V, 1), (i_t * BT, 0), (BT, V), (1, 0), - ) - - b_g = tl.load(p_g, boundary_check=(0,)) - b_q = tl.load(p_q, boundary_check=(0, 1)) - b_k = tl.load(p_k, boundary_check=(0, 1)) - b_v = tl.load(p_v, boundary_check=(0, 1)) - b_do = tl.load(p_do, boundary_check=(0, 1)) - b_dq = tl.load(p_dq, boundary_check=(0, 1)) - b_dk = tl.load(p_dk, boundary_check=(0, 1)) - b_dv = tl.load(p_dv, boundary_check=(0, 1)) - b_dg = tl.load(p_dg, boundary_check=(0,)) + b_h_ptrs = l + ((bos + tl.minimum(i_t * BT + i_idx, T - 1)) * H + i_h) * L + b_llut + b_h = tl.load(b_h_ptrs, mask=(i_idx >= j_idx) & (i_t * BT + i_idx < T) & (i_t * BT + j_idx < T), other=0.0) + + p_g = g + bos * H + i_h + o_t * H + p_q = q + (bos * G + i_g) * K + o_k[:, None] + o_t[None, :] * (G * K) + p_k = k + (bos * G + i_g) * K + o_t[:, None] * (G * K) + o_k[None, :] + p_v = v + (bos * H + i_h) * V + o_v[:, None] + o_t[None, :] * (H * V) + p_do = do + (bos * H + i_h) * V + o_t[:, None] * (H * V) + o_v[None, :] + p_dg = dg + bos * H + i_h + o_t * H + p_dq = dq + (bos * H + i_h) * K + o_t[:, None] * (H * K) + o_k[None, :] + p_dk = dk + (bos * H + i_h) * K + o_t[:, None] * (H * K) + o_k[None, :] + p_dv = dv + (bos * H + i_h) * V + o_t[:, None] * (H * V) + o_v[None, :] + + b_g = tl.load(p_g, mask=m_t, other=0.0) + b_q = tl.load(p_q, mask=m_kt, other=0.0) + b_k = tl.load(p_k, mask=m_tk, other=0.0) + b_v = tl.load(p_v, mask=m_vt, other=0.0) + b_do = (tl.load(p_do, mask=m_tv, other=0.0) * SCALE).to(do.dtype.element_ty) + b_dq = tl.load(p_dq, mask=m_tk, other=0.0) + b_dk = tl.load(p_dk, mask=m_tk, other=0.0) + b_dv = tl.load(p_dv, mask=m_tv, other=0.0) + b_dg = tl.load(p_dg, mask=m_t, other=0.0) b_s = (tl.dot(b_k, b_q)).to(b_q.dtype) # Apply causal and padding masks @@ -1435,19 +1069,19 @@ def chunkwise_bwd_kernel_diag( b_dq += tl.dot(b_ds, b_k) b_dk += tl.trans(tl.dot(b_q, b_ds)) - tl.store(p_dv, b_dv.to(p_dv.dtype.element_ty), boundary_check=(0, 1)) - tl.store(p_dq, b_dq.to(p_dq.dtype.element_ty), boundary_check=(0, 1)) - tl.store(p_dk, b_dk.to(p_dk.dtype.element_ty), boundary_check=(0, 1)) - tl.store(p_dg, b_dg.to(p_dg.dtype.element_ty), boundary_check=(0,)) + tl.store(p_dv, b_dv.to(p_dv.dtype.element_ty), mask=m_tv) + tl.store(p_dq, b_dq.to(p_dq.dtype.element_ty), mask=m_tk) + tl.store(p_dk, b_dk.to(p_dk.dtype.element_ty), mask=m_tk) + tl.store(p_dg, b_dg.to(p_dg.dtype.element_ty), mask=m_t) num_intra_levels = (tl.log2(float(BT))).to(tl.int32) + 1 for i in range(num_intra_levels): - p_mask = tl.make_block_ptr(mask + i * (BT * BT), (BT, BT), (BT, 1), (0, 0), (BT, BT), (1, 0)) - b_mask = tl.load(p_mask, boundary_check=(0, 1)) + p_mask = mask + i * (BT * BT) + o_i[:, None] * BT + o_i[None, :] + b_mask = tl.load(p_mask, mask=(o_i[:, None] < BT) & (o_i[None, :] < BT), other=0.0) dl_i = tl.sum(tl.where(b_mask == 1, b_dl, 0), axis=1) - p_dl_i = tl.make_block_ptr(dl + (bos * H + i_h) * L + i, (T,), (H * L,), (i_t * BT,), (BT,), (0,)) - tl.store(p_dl_i, dl_i, boundary_check=(0,)) + p_dl_i = dl + (bos * H + i_h) * L + i + o_t * (H * L) + tl.store(p_dl_i, dl_i, mask=m_t & (i < L)) def construct_binary_level_mask(level, T): @@ -1522,13 +1156,14 @@ class ChunkLogLinearAttentionFunction(torch.autograd.Function): initial_state, output_final_state, cu_seqlens, + scale, ): B, T, G, K = k.shape _, _, H, V = v.shape _, _, _, L = level_scales.shape - if G != 1: - raise ValueError("Group dimension must be 1.") + if H % G != 0: + raise ValueError("The number of value heads must be divisible by the number of query/key heads.") if not math.log2(V).is_integer(): raise ValueError( @@ -1546,6 +1181,8 @@ class ChunkLogLinearAttentionFunction(torch.autograd.Function): h0 = initial_state.ht if initial_state is not None else None offsets = initial_state.offsets if initial_state is not None else None + original_cu_seqlens = cu_seqlens + original_T = T if cu_seqlens is None: NT = ceil_div(T + (torch.max(offsets) if offsets is not None else 0), BT) @@ -1565,6 +1202,19 @@ class ChunkLogLinearAttentionFunction(torch.autograd.Function): MAX_LEVEL = ceil_log(NT, 2) - 1 B = len(cu_seqlens) - 1 + # Exact powers of two need one more level in the serialized state than in the current outputs. + if output_final_state: + if cu_seqlens is None: + max_sequence_length = T + (offsets.max().item() if offsets is not None else 0) + else: + sequence_lengths = (cu_seqlens[1:] - cu_seqlens[:-1]).tolist() + if offsets is not None: + sequence_lengths = [length + offset for length, offset in zip(sequence_lengths, offsets.tolist())] + max_sequence_length = max(sequence_lengths) + completed_chunks = max_sequence_length // BT + if completed_chunks > 0: + MAX_LEVEL = max(MAX_LEVEL, math.floor(math.log2(completed_chunks))) + if MAX_LEVEL > 10: raise ValueError("Sequence length must be less than 2**17") @@ -1577,7 +1227,7 @@ class ChunkLogLinearAttentionFunction(torch.autograd.Function): if initial_state is not None: if cu_seqlens is not None: - cu_seqlens = cu_seqlens + F.pad(torch.cumsum(offsets % BT), (1, 0)) + cu_seqlens = cu_seqlens + F.pad(torch.cumsum(offsets % BT, dim=0), (1, 0)) else: assert (offsets == offsets[0]).all() T += offsets[0].item() % BT @@ -1595,6 +1245,7 @@ class ChunkLogLinearAttentionFunction(torch.autograd.Function): g=g, level_scales=level_scales, cu_seqlens=cu_seqlens, + input_cu_seqlens=original_cu_seqlens, q_prev=initial_state.q_prev, k_prev=initial_state.k_prev, v_prev=initial_state.v_prev, @@ -1607,6 +1258,8 @@ class ChunkLogLinearAttentionFunction(torch.autograd.Function): level_scales_new=level_scales_new, offsets=offsets, T=T, + input_T=original_T, + G=G, H=H, K=K, V=V, @@ -1627,10 +1280,12 @@ class ChunkLogLinearAttentionFunction(torch.autograd.Function): ) new_offsets = torch.zeros((B,), dtype=torch.int32, device=v.device) + # Cache raw decay increments so a continued partial chunk is cumulatively summed exactly once. + g_for_state = g g = chunk_local_cumsum(g, chunk_size=BT, cu_seqlens=cu_seqlens) def grid(meta): - return (triton.cdiv(K, meta["BK"]), B * H) + return (triton.cdiv(K, meta["BK"]) * B * H,) l_in = h0.shape[1] if initial_state is not None else None l_out = ht.shape[1] if output_final_state else None @@ -1650,7 +1305,11 @@ class ChunkLogLinearAttentionFunction(torch.autograd.Function): offsets=offsets, new_offsets=new_offsets, cu_seqlens=cu_seqlens, + output_cu_seqlens=original_cu_seqlens, + stride_o_n=o.stride(0), T=T, + output_T=original_T, + G=G, H=H, K=K, V=V, @@ -1660,10 +1319,12 @@ class ChunkLogLinearAttentionFunction(torch.autograd.Function): L_OUT=l_out, MIN_LEVEL=0, MAX_LEVEL=MAX_LEVEL, + SCALE=scale, ) ctx.save_for_backward(q, k, v, g, level_scales, initial_state, cu_seqlens) ctx.chunk_size = BT + ctx.scale = scale if output_final_state: q_prev = torch.zeros((B, BT, G, K), dtype=q.dtype, device=q.device) @@ -1676,7 +1337,7 @@ class ChunkLogLinearAttentionFunction(torch.autograd.Function): q=q, k=k, v=v, - g=g, + g=g_for_state, level_scales=level_scales, cu_seqlens=cu_seqlens, q_prev=q_prev, @@ -1686,6 +1347,7 @@ class ChunkLogLinearAttentionFunction(torch.autograd.Function): level_scales_prev=level_scales_prev, offsets=new_offsets, T=T, + G=G, H=H, K=K, V=V, @@ -1715,6 +1377,7 @@ class ChunkLogLinearAttentionFunction(torch.autograd.Function): q, k, v, g, level_scales, initial_state, cu_seqlens = ctx.saved_tensors chunk_size = ctx.chunk_size + scale = ctx.scale llut = ctx.llut mask = masks(chunk_size, v.device) @@ -1724,7 +1387,6 @@ class ChunkLogLinearAttentionFunction(torch.autograd.Function): ) B, T, G, K = k.shape - assert G == 1, "Multi-head attention is not supported" _, _, H, V = v.shape _, _, _, L = level_scales.shape BT = chunk_size @@ -1754,9 +1416,9 @@ class ChunkLogLinearAttentionFunction(torch.autograd.Function): grid = (B * H,) def grid_f(meta): - return (triton.cdiv(K, meta["BK"]), B * H) + return (triton.cdiv(K, meta["BK"]) * B * H,) - grid_t = (NT, B * H) + grid_t = (NT * B * H,) num_inter_chunk_levels = ceil_log(NT, 2) for ell in range(num_inter_chunk_levels - 1, -1, -1): @@ -1774,12 +1436,14 @@ class ChunkLogLinearAttentionFunction(torch.autograd.Function): cu_seqlens=cu_seqlens, ell=ell, T=T, + G=G, H=H, K=K, V=V, L=L, BT=BT, NT=NT, + SCALE=scale, ) chunkwise_bwd_kernel_dhg[grid_f]( do=do, @@ -1792,12 +1456,14 @@ class ChunkLogLinearAttentionFunction(torch.autograd.Function): cu_seqlens=cu_seqlens, ell=ell, T=T, + G=G, H=H, K=K, V=V, L=L, BT=BT, NT=NT, + SCALE=scale, ) chunkwise_bwd_kernel_dkg[grid_t]( @@ -1810,6 +1476,7 @@ class ChunkLogLinearAttentionFunction(torch.autograd.Function): dg=dg, cu_seqlens=cu_seqlens, T=T, + G=G, H=H, K=K, V=V, @@ -1825,6 +1492,7 @@ class ChunkLogLinearAttentionFunction(torch.autograd.Function): dv=dv, cu_seqlens=cu_seqlens, T=T, + G=G, H=H, K=K, V=V, @@ -1849,18 +1517,21 @@ class ChunkLogLinearAttentionFunction(torch.autograd.Function): dl=dl, cu_seqlens=cu_seqlens, T=T, + G=G, H=H, K=K, V=V, L=L, BT=BT, + NT=NT, + SCALE=scale, ) dg = chunk_local_cumsum(dg, chunk_size=chunk_size, reverse=True, cu_seqlens=cu_seqlens).to(g.dtype) dq = reduce(dq, "b t (g h) k -> b t g k", "sum", g=G, h=H // G) dk = reduce(dk, "b t (g h) k -> b t g k", "sum", g=G, h=H // G) - return dq, dk, dv, dg, dl, None, None, None + return dq, dk, dv, dg, dl, None, None, None, None @torch.compiler.disable @@ -1873,13 +1544,14 @@ def chunk_log_linear_attn( initial_state: LogLinearAttentionState | None = None, output_final_state: bool = False, cu_seqlens: torch.LongTensor | None = None, + scale: float | None = 1.0, ) -> tuple[torch.Tensor, torch.Tensor]: r""" Args: q (torch.Tensor): - queries of shape `[B, T, H, K]`. + queries of shape `[B, T, G, K]`. The number of value heads H must be divisible by G. k (torch.Tensor): - keys of shape `[B, T, H, K]`. + keys of shape `[B, T, G, K]`. Each query/key head is shared by H/G consecutive value heads. v (torch.Tensor): values of shape `[B, T, H, V]`. g (torch.Tensor): @@ -1895,6 +1567,8 @@ def chunk_log_linear_attn( cu_seqlens (torch.LongTensor): Cumulative sequence lengths of shape `[N+1]` used for variable-length training, consistent with the FlashAttention API. + scale (float, Optional): + Attention score scale. Default: 1.0. Pass `None` to use `K ** -0.5`. Returns: o (torch.Tensor): @@ -1910,6 +1584,9 @@ def chunk_log_linear_attn( f"Please flatten variable-length inputs before processing.", ) + if scale is None: + scale = q.shape[-1] ** -0.5 + o, final_state = ChunkLogLinearAttentionFunction.apply( q, k, @@ -1919,5 +1596,6 @@ def chunk_log_linear_attn( initial_state, output_final_state, cu_seqlens, + scale, ) return o, final_state diff --git a/build/torch-cuda/ops/log_linear_attn/naive.py b/build/torch-cuda/ops/log_linear_attn/naive.py index 9badd5605bd3bfab72d892315d695a51707ad03c..57ef0b36a714a148733709825e7755fab24e1c9d 100644 --- a/build/torch-cuda/ops/log_linear_attn/naive.py +++ b/build/torch-cuda/ops/log_linear_attn/naive.py @@ -51,7 +51,16 @@ def construct_H_matrix(a, L): return H -def naive_log_linear_attn(q, k, v, g, level_scales): +def naive_log_linear_attn(q, k, v, g, level_scales, scale: float | None = 1.0): + """Compute grouped log-linear attention with the same scaling as the chunk operator.""" + if scale is None: + scale = q.shape[-1] ** -0.5 + num_groups, num_heads = k.shape[2], v.shape[2] + if num_heads % num_groups != 0: + raise ValueError("The number of value heads must be divisible by the number of query/key heads.") + if num_groups != num_heads: + q = q.repeat_interleave(num_heads // num_groups, dim=2) + k = k.repeat_interleave(num_heads // num_groups, dim=2) H = construct_H_matrix(g.permute(0, 2, 1), level_scales.permute(0, 2, 3, 1)) M = torch.einsum("bhlc,blhn,bchn->bhlc", H, q, k) - return torch.einsum("bhlc,bchp->blhp", M, v) + return torch.einsum("bhlc,bchp->blhp", M, v) * scale diff --git a/build/torch-cuda/ops/mesa_net/chunk.py b/build/torch-cuda/ops/mesa_net/chunk.py index 72031cc120808fb77e541ce9e40d89b6c36b3785..4d68d925eaf324ffe5fd01614c9630597d00750a 100644 --- a/build/torch-cuda/ops/mesa_net/chunk.py +++ b/build/torch-cuda/ops/mesa_net/chunk.py @@ -7,7 +7,7 @@ import torch -from ...modules.l2norm import l2norm_bwd, l2norm_fwd +from ...modules.norm.l2norm import l2norm_bwd, l2norm_fwd from ...ops.common.chunk_h import chunk_bwd_dh from ...ops.mesa_net.chunk_cg_solver_bwd import chunk_mesa_cg_bwd from ...ops.mesa_net.chunk_cg_solver_fwd import chunk_mesa_cg_fwd @@ -364,7 +364,7 @@ def chunk_mesa_net( >>> import torch >>> import torch.nn.functional as F >>> from einops import rearrange - >>> from ...ops.mesa_net import chunk_mesa_net + >>> from fla.ops.mesa_net import chunk_mesa_net # inputs with equal lengths >>> B, T, H, K, V = 4, 2048, 16, 128, 128 >>> q = torch.randn(B, T, H, K, dtype=torch.bfloat16, device='cuda') diff --git a/build/torch-cuda/ops/mesa_net/chunk_cg_solver_bwd.py b/build/torch-cuda/ops/mesa_net/chunk_cg_solver_bwd.py index a08425bb369573341a69e6c8c03e0f815b374c55..c4f29e04d65de9be5df9022502e4db4fe0c52722 100644 --- a/build/torch-cuda/ops/mesa_net/chunk_cg_solver_bwd.py +++ b/build/torch-cuda/ops/mesa_net/chunk_cg_solver_bwd.py @@ -24,7 +24,7 @@ def chunk_update_once( b_lamb, ): b_o = tl.dot((tl.dot(b_p.to(b_k.dtype), tl.trans(b_k)) * b_m).to(b_v.dtype), b_v) - b_o += tl.dot((b_p * b_g_exp_q).to(b_h.dtype), b_h) + b_o = tl.dot((b_p * b_g_exp_q).to(b_h.dtype), b_h, b_o) if b_lamb is not None: b_o += b_lamb[None, :] * b_p return b_o @@ -52,13 +52,13 @@ def chunk_fwd_mesa_cg_dim64_kernel( BK: tl.constexpr, IS_VARLEN: tl.constexpr, ): - i_t, i_bh = tl.program_id(0), tl.program_id(1) + i_t, i_bh = tl.program_id(0).to(tl.int64), tl.program_id(1).to(tl.int64) i_b, i_h = i_bh // H, i_bh % H 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.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) + i_n, i_t = tl.load(chunk_indices + i_t * 2).to(tl.int32), tl.load(chunk_indices + i_t * 2 + 1).to(tl.int64) + bos, eos = tl.load(cu_seqlens + i_n).to(tl.int64), tl.load(cu_seqlens + i_n + 1).to(tl.int64) T = eos - bos NT = tl.cdiv(T, BT) else: @@ -79,20 +79,24 @@ def chunk_fwd_mesa_cg_dim64_kernel( beta += bos * H + i_h lamb += i_h * K - p_q = tl.make_block_ptr(dq, (T, K), (H*K, 1), (i_t * BT, 0), (BT, BK), (1, 0)) - p_k = tl.make_block_ptr(k, (T, K), (H*K, 1), (i_t * BT, 0), (BT, BK), (1, 0)) - p_h = tl.make_block_ptr(h, (K, K), (K, 1), (0, 0), (BK, BK), (1, 0)) - - b_h = tl.load(p_h, boundary_check=(0, 1)) - b_k = tl.load(p_k, boundary_check=(0, 1)) - b_q = tl.load(p_q, boundary_check=(0, 1)).to(tl.float32) - - p_g = tl.make_block_ptr(g, (T,), (H,), (i_t * BT,), (BT,), (0,)) - b_g = tl.load(p_g, boundary_check=(0,)).to(tl.float32) - p_beta = tl.make_block_ptr(beta, (T,), (H,), (i_t * BT,), (BT,), (0,)) - b_beta = tl.load(p_beta, boundary_check=(0,)).to(tl.float32) - p_lamb = tl.make_block_ptr(lamb, (K,), (1,), (0,), (BK,), (0,)) - b_lamb = tl.load(p_lamb, boundary_check=(0,)).to(tl.float32) + o_k = tl.arange(0, BK) + m_k = o_k < K + m_tk = m_t[:, None] & m_k[None, :] + m_kk = m_k[:, None] & m_k[None, :] + p_q = dq + o_t[:, None] * (H*K) + o_k[None, :] + p_k = k + o_t[:, None] * (H*K) + o_k[None, :] + p_h = h + o_k[:, None] * K + o_k[None, :] + + b_h = tl.load(p_h, mask=m_kk, other=0.0) + b_k = tl.load(p_k, mask=m_tk, other=0.0) + b_q = tl.load(p_q, mask=m_tk, other=0.0).to(tl.float32) + + p_g = g + o_t * H + b_g = tl.load(p_g, mask=m_t, other=0.0).to(tl.float32) + p_beta = beta + o_t * H + b_beta = tl.load(p_beta, mask=m_t, other=0.0).to(tl.float32) + p_lamb = lamb + o_k + b_lamb = tl.load(p_lamb, mask=m_k, other=0.0).to(tl.float32) b_m = exp2(b_g[:, None] - b_g[None, :]) * b_beta[None, :] b_m = tl.where((o_t[:, None] >= o_t[None, :]) & (m_t[:, None] & m_t[None, :]), b_m, 0) @@ -115,8 +119,8 @@ def chunk_fwd_mesa_cg_dim64_kernel( b_p = b_r + (b_delta_new / (b_delta_old + 1e-5))[:, None] * b_p b_delta_old = b_delta_new - p_q_final = tl.make_block_ptr(dq_final, (T, K), (H*K, 1), (i_t * BT, 0), (BT, BK), (1, 0)) - tl.store(p_q_final, b_x.to(p_q_final.dtype.element_ty), boundary_check=(0, 1)) + p_q_final = dq_final + o_t[:, None] * (H*K) + o_k[None, :] + tl.store(p_q_final, b_x.to(p_q_final.dtype.element_ty), mask=m_tk) def chunk_mesa_cg_bwd( diff --git a/build/torch-cuda/ops/mesa_net/chunk_cg_solver_fwd.py b/build/torch-cuda/ops/mesa_net/chunk_cg_solver_fwd.py index 75f474462185c8d69427b03db6b0d07e9160d34e..50f8785c4711759a0a19608f6dc1b9ec67df9984 100644 --- a/build/torch-cuda/ops/mesa_net/chunk_cg_solver_fwd.py +++ b/build/torch-cuda/ops/mesa_net/chunk_cg_solver_fwd.py @@ -24,7 +24,7 @@ def chunk_update_once( b_lamb, ): b_o = tl.dot((tl.dot(b_p.to(b_k.dtype), tl.trans(b_k)) * b_m).to(b_v.dtype), b_v) - b_o += tl.dot((b_p * b_g_exp_q).to(b_h.dtype), b_h) + b_o = tl.dot((b_p * b_g_exp_q).to(b_h.dtype), b_h, b_o) if b_lamb is not None: b_o += b_lamb[None, :] * b_p return b_o @@ -55,13 +55,13 @@ def chunk_fwd_mesa_cg_dim64_kernel( BK: tl.constexpr, IS_VARLEN: tl.constexpr, ): - i_t, i_bh = tl.program_id(0), tl.program_id(1) + i_t, i_bh = tl.program_id(0).to(tl.int64), tl.program_id(1).to(tl.int64) i_b, i_h = i_bh // H, i_bh % H 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.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) + i_n, i_t = tl.load(chunk_indices + i_t * 2).to(tl.int32), tl.load(chunk_indices + i_t * 2 + 1).to(tl.int64) + bos, eos = tl.load(cu_seqlens + i_n).to(tl.int64), tl.load(cu_seqlens + i_n + 1).to(tl.int64) T = eos - bos NT = tl.cdiv(T, BT) else: @@ -85,21 +85,25 @@ def chunk_fwd_mesa_cg_dim64_kernel( v += (bos * H + i_h) * K h_kv += (i_tg * H + i_h).to(tl.int64) * K * K - p_q = tl.make_block_ptr(q, (T, K), (H*K, 1), (i_t * BT, 0), (BT, BK), (1, 0)) - p_k = tl.make_block_ptr(k, (T, K), (H*K, 1), (i_t * BT, 0), (BT, BK), (1, 0)) - p_h = tl.make_block_ptr(h, (K, K), (K, 1), (0, 0), (BK, BK), (1, 0)) + o_k = tl.arange(0, BK) + m_k = o_k < K + m_tk = m_t[:, None] & m_k[None, :] + m_kk = m_k[:, None] & m_k[None, :] + p_q = q + o_t[:, None] * (H*K) + o_k[None, :] + p_k = k + o_t[:, None] * (H*K) + o_k[None, :] + p_h = h + o_k[:, None] * K + o_k[None, :] - b_h = tl.load(p_h, boundary_check=(0, 1)) - b_k = tl.load(p_k, boundary_check=(0, 1)) - b_q = tl.load(p_q, boundary_check=(0, 1)).to(tl.float32) + b_h = tl.load(p_h, mask=m_kk, other=0.0) + b_k = tl.load(p_k, mask=m_tk, other=0.0) + b_q = tl.load(p_q, mask=m_tk, other=0.0).to(tl.float32) - p_g = tl.make_block_ptr(g, (T,), (H,), (i_t * BT,), (BT,), (0,)) - b_g = tl.load(p_g, boundary_check=(0,)).to(tl.float32) - p_beta = tl.make_block_ptr(beta, (T,), (H,), (i_t * BT,), (BT,), (0,)) - b_beta = tl.load(p_beta, boundary_check=(0,)).to(tl.float32) - p_lamb = tl.make_block_ptr(lamb, (K,), (1,), (0,), (BK,), (0,)) + p_g = g + o_t * H + b_g = tl.load(p_g, mask=m_t, other=0.0).to(tl.float32) + p_beta = beta + o_t * H + b_beta = tl.load(p_beta, mask=m_t, other=0.0).to(tl.float32) + p_lamb = lamb + o_k - b_lamb = tl.load(p_lamb, boundary_check=(0,)).to(tl.float32) + b_lamb = tl.load(p_lamb, mask=m_k, other=0.0).to(tl.float32) b_m = exp2(b_g[:, None] - b_g[None, :]) * b_beta[None, :] b_m = tl.where((o_t[:, None] >= o_t[None, :]) & (m_t[:, None] & m_t[None, :]), b_m, 0) @@ -122,16 +126,16 @@ def chunk_fwd_mesa_cg_dim64_kernel( b_p = b_r + (b_delta_new / (b_delta_old + 1e-5))[:, None] * b_p b_delta_old = b_delta_new - p_q_final = tl.make_block_ptr(q_final, (T, K), (H*K, 1), (i_t * BT, 0), (BT, BK), (1, 0)) - tl.store(p_q_final, b_x.to(p_q_final.dtype.element_ty), boundary_check=(0, 1)) + p_q_final = q_final + o_t[:, None] * (H*K) + o_k[None, :] + tl.store(p_q_final, b_x.to(p_q_final.dtype.element_ty), mask=m_tk) - p_h_kv = tl.make_block_ptr(h_kv, (K, K), (K, 1), (0, 0), (BK, BK), (1, 0)) - b_h_kv = tl.load(p_h_kv, boundary_check=(0, 1)) - p_v = tl.make_block_ptr(v, (T, K), (H*K, 1), (i_t * BT, 0), (BT, BK), (1, 0)) - b_v = tl.load(p_v, boundary_check=(0, 1)) + p_h_kv = h_kv + o_k[:, None] * K + o_k[None, :] + b_h_kv = tl.load(p_h_kv, mask=m_kk, other=0.0) + p_v = v + o_t[:, None] * (H*K) + o_k[None, :] + b_v = tl.load(p_v, mask=m_tk, other=0.0) b_o = chunk_update_once(b_x, b_k, b_v, b_m, b_g_exp_q, b_h_kv, None) - p_o = tl.make_block_ptr(o, (T, K), (H*K, 1), (i_t * BT, 0), (BT, BK), (1, 0)) - tl.store(p_o, b_o.to(p_o.dtype.element_ty), boundary_check=(0, 1)) + p_o = o + o_t[:, None] * (H*K) + o_k[None, :] + tl.store(p_o, b_o.to(p_o.dtype.element_ty), mask=m_tk) def chunk_mesa_cg_fwd( diff --git a/build/torch-cuda/ops/mesa_net/chunk_h_fwd.py b/build/torch-cuda/ops/mesa_net/chunk_h_fwd.py index 88bfa11b11698295a6c835694efce8e887fea887..e638836d9250e79779ea58546cb22bba35616e47 100644 --- a/build/torch-cuda/ops/mesa_net/chunk_h_fwd.py +++ b/build/torch-cuda/ops/mesa_net/chunk_h_fwd.py @@ -10,7 +10,7 @@ import triton import triton.language as tl from ...ops.utils import prepare_chunk_offsets -from ...ops.utils.op import exp2, safe_dot +from ...ops.utils.op import exp2, safe_dot, unflatten_program_id from ...utils import autotune_cache_kwargs @@ -54,50 +54,57 @@ def chunk_mesa_net_fwd_kernel_h( STORE_FINAL_STATE: tl.constexpr, IS_VARLEN: tl.constexpr, ): - i_k, i_v, i_nh = tl.program_id(0), tl.program_id(1), tl.program_id(2) + i_k, i_v, i_nh = unflatten_program_id(tl.cdiv(K, 64), tl.cdiv(V, 64)) i_n, i_h = i_nh // H, i_nh % H if IS_VARLEN: - bos, eos = tl.load(cu_seqlens + i_n).to(tl.int32), tl.load(cu_seqlens + i_n + 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 NT = tl.cdiv(T, BT) NS = tl.cdiv(T, BS) - boh = tl.load(split_offsets + i_n).to(tl.int32) + boh = tl.load(split_offsets + i_n).to(tl.int64) else: bos, eos = i_n * T, i_n * T + T NT = tl.cdiv(T, BT) NS = tl.cdiv(T, BS) boh = i_n * NS + o_k = i_k * BK + tl.arange(0, BK) + o_v = i_v * BV + tl.arange(0, BV) + m_kv = (o_k[:, None] < K) & (o_v[None, :] < V) # [BK, BV] b_h = tl.zeros([BK, BV], dtype=tl.float32) b_h_kv = tl.zeros([BK, BV], dtype=tl.float32) if USE_INITIAL_STATE: - p_h0 = tl.make_block_ptr(h_init + i_nh * K*V, (K, V), (V, 1), (i_k * BK, i_v * BV), (BK, BV), (1, 0)) - b_h = tl.load(p_h0, boundary_check=(0, 1)).to(tl.float32) - p_h_kv0 = tl.make_block_ptr(h_kv_init + i_nh * K*V, (K, V), (V, 1), (i_k * BK, i_v * BV), (BK, BV), (1, 0)) - b_h_kv = tl.load(p_h_kv0, boundary_check=(0, 1)).to(tl.float32) + p_h0 = h_init + i_nh * K*V + o_k[:, None] * V + o_v[None, :] + b_h = tl.load(p_h0, mask=m_kv, other=0.0).to(tl.float32) + p_h_kv0 = h_kv_init + i_nh * K*V + o_k[:, None] * V + o_v[None, :] + b_h_kv = tl.load(p_h_kv0, mask=m_kv, other=0.0).to(tl.float32) for i_t in range(NT): i_s = i_t // (BS // BT) - p_k = tl.make_block_ptr(k + (bos*H + i_h) * K, (T, K), (H*K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) - p_k2 = tl.make_block_ptr(k + (bos*H + i_h) * V, (T, V), (H*V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) - p_v = tl.make_block_ptr(v + (bos*H + i_h) * V, (T, V), (H*V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) - p_beta = tl.make_block_ptr(beta + (bos*H + i_h), (T,), (H,), (i_t * BT,), (BT, ), (0,)) - b_beta = tl.load(p_beta, boundary_check=(0,)) + o_t = i_t.to(tl.int64) * BT + tl.arange(0, BT) + m_t = o_t < T + m_tk = m_t[:, None] & (o_k[None, :] < K) + m_tv = m_t[:, None] & (o_v[None, :] < V) + p_k = k + (bos*H + i_h) * K + o_t[:, None] * (H*K) + o_k[None, :] + p_k2 = k + (bos*H + i_h) * V + o_t[:, None] * (H*V) + o_v[None, :] + p_v = v + (bos*H + i_h) * V + o_t[:, None] * (H*V) + o_v[None, :] + p_beta = beta + (bos*H + i_h) + o_t * H + b_beta = tl.load(p_beta, mask=m_t, other=0.0) o_h = ((boh + i_s) * H + i_h).to(tl.int64) * K*V - p_h = tl.make_block_ptr(h + o_h, (K, V), (V, 1), (i_k * BK, i_v * BV), (BK, BV), (1, 0)) - p_h_kv = tl.make_block_ptr(h_kv + o_h, (K, V), (V, 1), (i_k * BK, i_v * BV), (BK, BV), (1, 0)) + p_h = h + o_h + o_k[:, None] * V + o_v[None, :] + p_h_kv = h_kv + o_h + o_k[:, None] * V + o_v[None, :] if i_t % (BS // BT) == 0: - tl.store(p_h, b_h.to(p_h.dtype.element_ty), boundary_check=(0, 1)) - tl.store(p_h_kv, b_h_kv.to(p_h_kv.dtype.element_ty), boundary_check=(0, 1)) + tl.store(p_h, b_h.to(p_h.dtype.element_ty), mask=m_kv) + tl.store(p_h_kv, b_h_kv.to(p_h_kv.dtype.element_ty), mask=m_kv) # [BK, BT] - b_k = tl.load(p_k, boundary_check=(0, 1)) - b_k2 = tl.load(p_k2, boundary_check=(0, 1)) + b_k = tl.load(p_k, mask=m_tk, other=0.0) + b_k2 = tl.load(p_k2, mask=m_tv, other=0.0) # [BT, BV] - b_v = tl.load(p_v, boundary_check=(0, 1)) + b_v = tl.load(p_v, mask=m_tv, other=0.0) last_idx = min((i_t + 1) * BT, T) - 1 # scalar decay @@ -108,13 +115,13 @@ def chunk_mesa_net_fwd_kernel_h( b_g = tl.load(p_g, mask=(i_t * BT + tl.arange(0, BT) < T), other=0.) b_k_decay = ((b_k * exp2(b_g_last - b_g)[:, None]) * b_beta[:, None]).to(b_k2.dtype) b_h += safe_dot(tl.trans(b_k_decay), b_k2) - b_h_kv += safe_dot(tl.trans(b_k_decay), b_v.to(b_k2.dtype)) + b_h_kv += safe_dot(tl.trans(b_k_decay).to(b_v.dtype), b_v) if STORE_FINAL_STATE: - p_ht = tl.make_block_ptr(h_final + i_nh * K*V, (K, V), (V, 1), (i_k * BK, i_v * BV), (BK, BV), (1, 0)) - tl.store(p_ht, b_h.to(p_ht.dtype.element_ty), boundary_check=(0, 1)) - p_h_kv_final = tl.make_block_ptr(h_kv_final + i_nh * K*V, (K, V), (V, 1), (i_k * BK, i_v * BV), (BK, BV), (1, 0)) - tl.store(p_h_kv_final, b_h_kv.to(p_h_kv_final.dtype.element_ty), boundary_check=(0, 1)) + p_ht = h_final + i_nh * K*V + o_k[:, None] * V + o_v[None, :] + tl.store(p_ht, b_h.to(p_ht.dtype.element_ty), mask=m_kv) + p_h_kv_final = h_kv_final + i_nh * K*V + o_k[:, None] * V + o_v[None, :] + tl.store(p_h_kv_final, b_h_kv.to(p_h_kv_final.dtype.element_ty), mask=m_kv) def chunk_mesa_fwd_h( @@ -142,12 +149,13 @@ def chunk_mesa_fwd_h( split_offsets = prepare_chunk_offsets(cu_seqlens, BS) N, NS = len(cu_seqlens) - 1, split_offsets[-1].item() - h = k.new_empty(B, NS, H, K, V, dtype=k.dtype if not states_in_fp32 else torch.float) - h_kv = k.new_empty(B, NS, H, K, V, dtype=k.dtype if not states_in_fp32 else torch.float) + # h_kk is stored in fp32 to ensure positive definite in CG solver + h = k.new_empty(B, NS, H, K, V, dtype=torch.float) + h_kv = k.new_empty(B, NS, H, K, V, dtype=torch.bfloat16 if not states_in_fp32 else torch.float) h_final = k.new_empty(N, H, K, V, dtype=torch.float) if output_final_state else None h_kv_final = k.new_empty(N, H, K, V, dtype=torch.float) - def grid(meta): return (triton.cdiv(K, 64), triton.cdiv(V, 64), N * H) + def grid(meta): return (triton.cdiv(K, 64) * triton.cdiv(V, 64) * N * H,) chunk_mesa_net_fwd_kernel_h[grid]( k=k, diff --git a/build/torch-cuda/ops/mesa_net/chunk_h_kk_intra_bwd.py b/build/torch-cuda/ops/mesa_net/chunk_h_kk_intra_bwd.py index 172bdddaba53dd99a5662133b644420139949907..8568072bc4c78a36772748686d22e4daa025ee10 100644 --- a/build/torch-cuda/ops/mesa_net/chunk_h_kk_intra_bwd.py +++ b/build/torch-cuda/ops/mesa_net/chunk_h_kk_intra_bwd.py @@ -42,13 +42,13 @@ def chunk_mesa_net_h_kk_bwd_intra_kernel( BV: tl.constexpr, IS_VARLEN: tl.constexpr, ): - i_t, i_bh = tl.program_id(0), tl.program_id(1) + i_t, i_bh = tl.program_id(0).to(tl.int64), tl.program_id(1).to(tl.int64) i_b, i_h = i_bh // H, i_bh % H 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.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) + i_n, i_t = tl.load(chunk_indices + i_t * 2).to(tl.int32), tl.load(chunk_indices + i_t * 2 + 1).to(tl.int64) + bos, eos = tl.load(cu_seqlens + i_n).to(tl.int64), tl.load(cu_seqlens + i_n + 1).to(tl.int64) T = eos - bos NT = tl.cdiv(T, BT) else: @@ -79,41 +79,47 @@ def chunk_mesa_net_h_kk_bwd_intra_kernel( b_dg_last = tl.zeros([1], dtype=tl.float32) b_dg = tl.zeros([BT], dtype=tl.float32) - p_g = tl.make_block_ptr(g, (T,), (H,), (i_t * BT,), (BT,), (0,)) - b_g = tl.load(p_g, boundary_check=(0,)) + o_k = tl.arange(0, BK) + o_v = tl.arange(0, BV) + m_k = o_k < K + m_tk = m_t[:, None] & m_k[None, :] + m_tv = m_t[:, None] & (o_v[None, :] < V) + m_h = (o_v[:, None] < V) & m_k[None, :] + p_g = g + o_t * H + b_g = tl.load(p_g, mask=m_t, other=0.0) b_g_last = tl.load(g + (min(i_t * BT + BT, T) - 1) * H) b_gk = tl.where(m_t, exp2(b_g_last - b_g), 0) - p_q_star = tl.make_block_ptr(q_star, (T, V), (H*V, 1), (i_t * BT, 0), (BT, BV), (1, 0)) - b_q_star = tl.load(p_q_star, boundary_check=(0, 1)) - p_dq = tl.make_block_ptr(dq, (T, V), (H*V, 1), (i_t * BT, 0), (BT, BV), (1, 0)) - b_dq = tl.load(p_dq, boundary_check=(0, 1)) + p_q_star = q_star + o_t[:, None] * (H*V) + o_v[None, :] + b_q_star = tl.load(p_q_star, mask=m_tv, other=0.0) + p_dq = dq + o_t[:, None] * (H*V) + o_v[None, :] + b_dq = tl.load(p_dq, mask=m_tv, other=0.0) b_dlamb = -tl.sum(b_q_star * b_dq, axis=0) - p_dlamb = tl.make_block_ptr(dlamb, (K,), (1,), (0,), (BK,), (0,)) - tl.store(p_dlamb, b_dlamb.to(p_dlamb.dtype.element_ty), boundary_check=(0,)) - - p_h = tl.make_block_ptr(h, (V, K), (1, V), (0, 0), (BV, BK), (0, 1)) - p_dh = tl.make_block_ptr(dh, (V, K), (1, V), (0, 0), (BV, BK), (0, 1)) - p_beta = tl.make_block_ptr(beta, (T,), (H,), (i_t * BT,), (BT,), (0,)) - b_beta = tl.load(p_beta, boundary_check=(0,)) - p_k = tl.make_block_ptr(k, (T, K), (H*K, 1), (i_t * BT, 0), (BT, BK), (1, 0)) - b_v = tl.load(p_k, boundary_check=(0, 1)) + p_dlamb = dlamb + o_k + tl.store(p_dlamb, b_dlamb.to(p_dlamb.dtype.element_ty), mask=m_k) + + p_h = h + o_v[:, None] + o_k[None, :] * V + p_dh = dh + o_v[:, None] + o_k[None, :] * V + p_beta = beta + o_t * H + b_beta = tl.load(p_beta, mask=m_t, other=0.0) + p_k = k + o_t[:, None] * (H*K) + o_k[None, :] + b_v = tl.load(p_k, mask=m_tk, other=0.0) b_k = (b_v * b_beta[:, None]).to(b_v.dtype) b_m = tl.where((o_t[:, None] >= o_t[None, :]) & (m_t[:, None] & m_t[None, :]), exp2(b_g[:, None] - b_g[None, :]), 0) b_s = tl.dot(b_q_star, tl.trans(b_k)) * b_m b_ds = tl.dot(b_dq, tl.trans(b_v)) - b_dv += tl.dot(tl.trans(b_s.to(b_dq.dtype)), b_dq) + b_dv = tl.dot(tl.trans(b_s.to(b_dq.dtype)), b_dq, b_dv) b_dm = b_s * b_ds b_dm = tl.where(tl.arange(0, BT)[:, None] >= tl.arange(0, BT)[None, :], b_dm, 0) b_dg += tl.sum(b_dm, axis=1) b_dg -= tl.sum(b_dm, axis=0) b_ds = b_ds * b_m - b_dk += tl.dot(tl.trans(b_ds.to(b_q_star.dtype)), b_q_star) + b_dk = tl.dot(tl.trans(b_ds.to(b_q_star.dtype)), b_q_star, b_dk) - b_h = tl.load(p_h, boundary_check=(0, 1)) - b_dg += tl.sum(tl.dot(b_dq, tl.trans(b_h)) * exp2(b_g)[:, None] * b_q_star, axis=1) - b_dh = tl.load(p_dh, boundary_check=(0, 1)) + b_h = tl.load(p_h, mask=m_h, other=0.0) + b_dg += tl.sum(tl.dot(b_dq.to(b_h.dtype), tl.trans(b_h)) * exp2(b_g)[:, None] * b_q_star, axis=1) + b_dh = tl.load(p_dh, mask=m_h, other=0.0) b_dk2 = tl.dot(b_v, b_dh.to(b_v.dtype)) * b_gk[:, None] b_dg -= tl.sum(b_dk2 * b_k, axis=1) b_dg_last += tl.sum(b_dk2 * b_k) @@ -122,19 +128,19 @@ def chunk_mesa_net_h_kk_bwd_intra_kernel( b_dh = b_dh * b_h b_dg_last += tl.sum(b_dh) * exp2(b_g_last) - p_dk_beta = tl.make_block_ptr(dk_beta, (T, K), (H*K, 1), (i_t * BT, 0), (BT, BK), (1, 0)) - b_dk -= tl.load(p_dk_beta, boundary_check=(0, 1)) + p_dk_beta = dk_beta + o_t[:, None] * (H*K) + o_k[None, :] + b_dk -= tl.load(p_dk_beta, mask=m_tk, other=0.0) b_dbeta = tl.sum(b_dk * b_v, axis=1) b_dk = b_dk * b_beta[:, None] + b_dv b_dk = -b_dk b_dg = tl.where(o_t < min(i_t * BT + BT, T) - 1, b_dg, b_dg + b_dg_last) - p_dk = tl.make_block_ptr(dk, (T, K), (H*K, 1), (i_t * BT, 0), (BT, BK), (1, 0)) - tl.store(p_dk, b_dk.to(p_dk.dtype.element_ty), boundary_check=(0, 1)) - p_dg = tl.make_block_ptr(dg, (T,), (H,), (i_t * BT,), (BT,), (0,)) - tl.store(p_dg, -b_dg.to(p_dg.dtype.element_ty), boundary_check=(0,)) - p_dbeta = tl.make_block_ptr(dbeta, (T,), (H,), (i_t * BT,), (BT,), (0,)) - tl.store(p_dbeta, -b_dbeta.to(p_dbeta.dtype.element_ty), boundary_check=(0,)) + p_dk = dk + o_t[:, None] * (H*K) + o_k[None, :] + tl.store(p_dk, b_dk.to(p_dk.dtype.element_ty), mask=m_tk) + p_dg = dg + o_t * H + tl.store(p_dg, -b_dg.to(p_dg.dtype.element_ty), mask=m_t) + p_dbeta = dbeta + o_t * H + tl.store(p_dbeta, -b_dbeta.to(p_dbeta.dtype.element_ty), mask=m_t) def chunk_mesa_net_h_kk_bwd_intra_fn( diff --git a/build/torch-cuda/ops/mesa_net/chunk_h_kv_intra_bwd.py b/build/torch-cuda/ops/mesa_net/chunk_h_kv_intra_bwd.py index 9767ed1c2b8bb24e8268250bbee5e13131b276a6..f2b4bf32f4427ab0cbbd67c45b686eb6fb876192 100644 --- a/build/torch-cuda/ops/mesa_net/chunk_h_kv_intra_bwd.py +++ b/build/torch-cuda/ops/mesa_net/chunk_h_kv_intra_bwd.py @@ -55,12 +55,12 @@ def chunk_mesa_net_h_kv_bwd_intra_kernel( BV: tl.constexpr, IS_VARLEN: tl.constexpr, ): - i_t, i_bh = tl.program_id(0), tl.program_id(1) + i_t, i_bh = tl.program_id(0).to(tl.int64), tl.program_id(1).to(tl.int64) i_b, i_h = i_bh // H, i_bh % H 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.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) + i_n, i_t = tl.load(chunk_indices + i_t * 2).to(tl.int32), tl.load(chunk_indices + i_t * 2 + 1).to(tl.int64) + bos, eos = tl.load(cu_seqlens + i_n).to(tl.int64), tl.load(cu_seqlens + i_n + 1).to(tl.int64) T = eos - bos NT = tl.cdiv(T, BT) else: @@ -92,27 +92,34 @@ def chunk_mesa_net_h_kv_bwd_intra_kernel( b_dg_last = tl.zeros([1], dtype=tl.float32) b_dg = tl.zeros([BT], dtype=tl.float32) - p_q = tl.make_block_ptr(q_star, (T, K), (H*K, 1), (i_t * BT, 0), (BT, BK), (1, 0)) - p_k = tl.make_block_ptr(k, (T, K), (H*K, 1), (i_t * BT, 0), (BT, BK), (1, 0)) - p_v = tl.make_block_ptr(v, (T, V), (H*V, 1), (i_t * BT, 0), (BT, BV), (1, 0)) - p_g = tl.make_block_ptr(g, (T,), (H,), (i_t * BT,), (BT,), (0,)) - p_beta = tl.make_block_ptr(beta, (T, ), (H, ), (i_t * BT,), (BT,), (0,)) - p_do = tl.make_block_ptr(do, (T, V), (H*V, 1), (i_t * BT, 0), (BT, BV), (1, 0)) - p_h = tl.make_block_ptr(h_kv, (V, K), (1, V), (0, 0), (BV, BK), (0, 1)) - p_dh = tl.make_block_ptr(dh_kv, (V, K), (1, V), (0, 0), (BV, BK), (0, 1)) - - b_q = tl.load(p_q, boundary_check=(0, 1)) - b_k = tl.load(p_k, boundary_check=(0, 1)) - b_v = tl.load(p_v, boundary_check=(0, 1)) - b_g = tl.load(p_g, boundary_check=(0,)) - b_beta = tl.load(p_beta, boundary_check=(0, )) - b_do = tl.load(p_do, boundary_check=(0, 1)) - b_h = tl.load(p_h, boundary_check=(0, 1)) - b_dh = tl.load(p_dh, boundary_check=(0, 1)) + o_k = tl.arange(0, BK) + o_v = tl.arange(0, BV) + m_tk = m_t[:, None] & (o_k[None, :] < K) + m_tv = m_t[:, None] & (o_v[None, :] < V) + m_h = (o_v[:, None] < V) & (o_k[None, :] < K) + p_q = q_star + o_t[:, None] * (H*K) + o_k[None, :] + p_k = k + o_t[:, None] * (H*K) + o_k[None, :] + p_v = v + o_t[:, None] * (H*V) + o_v[None, :] + p_g = g + o_t * H + p_beta = beta + o_t * H + p_do = do + o_t[:, None] * (H*V) + o_v[None, :] + p_h = h_kv + o_v[:, None] + o_k[None, :] * V + p_dh = dh_kv + o_v[:, None] + o_k[None, :] * V + + b_q = tl.load(p_q, mask=m_tk, other=0.0) + b_k = tl.load(p_k, mask=m_tk, other=0.0) + b_v = tl.load(p_v, mask=m_tv, other=0.0) + b_g = tl.load(p_g, mask=m_t, other=0.0) + b_beta = tl.load(p_beta, mask=m_t, other=0.0) + b_do = tl.load(p_do, mask=m_tv, other=0.0) + b_h = tl.load(p_h, mask=m_h, other=0.0) + b_dh = tl.load(p_dh, mask=m_h, other=0.0) b_g_last = tl.load(g + (min(i_t * BT + BT, T) - 1) * H) # calculation - b_dg_last += tl.sum(b_h * b_dh) + # h_kv is stored in bf16 (wide exponent range) while dh_kv is fp16 + # lift both precision to float32 to prevent overflow from large h_kv entries + b_dg_last += tl.sum(b_h.to(tl.float32) * b_dh.to(tl.float32)) b_dg_last *= exp2(b_g_last) b_m = tl.where((o_t[:, None] >= o_t[None, :]) & (m_t[:, None] & m_t[None, :]), exp2(b_g[:, None] - b_g[None, :]), 0) @@ -134,19 +141,19 @@ def chunk_mesa_net_h_kv_bwd_intra_kernel( b_dg_last += tl.sum(b_dk * b_k) b_dg -= tl.sum(b_dk * b_k, axis=1) b_dg += tl.sum(b_dq * b_q, axis=1) - b_dq += tl.dot(b_ds.to(b_k.dtype), b_k) + b_dq = tl.dot(b_ds.to(b_k.dtype), b_k, b_dq) b_dv += tl.dot(b_k, tl.trans(b_dh).to(b_k.dtype)) * b_g_exp_k[:, None] + tl.dot(tl.trans(b_s.to(b_do.dtype)), b_do) - b_dk += tl.dot(tl.trans(b_ds.to(b_q.dtype)), b_q) + b_dk = tl.dot(tl.trans(b_ds.to(b_q.dtype)), b_q, b_dk) b_dg = tl.where(o_t < min(i_t * BT + BT, T) - 1, b_dg, b_dg + b_dg_last) - p_dq = tl.make_block_ptr(dq, (T, K), (H*K, 1), (i_t * BT, 0), (BT, BK), (1, 0)) - p_dk = tl.make_block_ptr(dk_beta, (T, K), (H*K, 1), (i_t * BT, 0), (BT, BK), (1, 0)) - p_dv = tl.make_block_ptr(dv, (T, V), (H*V, 1), (i_t * BT, 0), (BT, BV), (1, 0)) - p_dg = tl.make_block_ptr(dg, (T,), (H,), (i_t * BT,), (BT,), (0,)) - tl.store(p_dq, b_dq.to(p_dq.dtype.element_ty), boundary_check=(0, 1)) - tl.store(p_dk, b_dk.to(p_dk.dtype.element_ty), boundary_check=(0, 1)) - tl.store(p_dv, b_dv.to(p_dv.dtype.element_ty), boundary_check=(0, 1)) - tl.store(p_dg, b_dg.to(p_dg.dtype.element_ty), boundary_check=(0,)) + p_dq = dq + o_t[:, None] * (H*K) + o_k[None, :] + p_dk = dk_beta + o_t[:, None] * (H*K) + o_k[None, :] + p_dv = dv + o_t[:, None] * (H*V) + o_v[None, :] + p_dg = dg + o_t * H + tl.store(p_dq, b_dq.to(p_dq.dtype.element_ty), mask=m_tk) + tl.store(p_dk, b_dk.to(p_dk.dtype.element_ty), mask=m_tk) + tl.store(p_dv, b_dv.to(p_dv.dtype.element_ty), mask=m_tv) + tl.store(p_dg, b_dg.to(p_dg.dtype.element_ty), mask=m_t) def chunk_mesa_net_h_kv_bwd_intra_fn( diff --git a/build/torch-cuda/ops/mesa_net/chunk_h_kv_intra_bwd_separate.py b/build/torch-cuda/ops/mesa_net/chunk_h_kv_intra_bwd_separate.py index 9f352a6265869004b7d745040fff7a3340a95bd1..fb1881ff691b817b79adb7a235bfc80e6b9e6738 100644 --- a/build/torch-cuda/ops/mesa_net/chunk_h_kv_intra_bwd_separate.py +++ b/build/torch-cuda/ops/mesa_net/chunk_h_kv_intra_bwd_separate.py @@ -53,12 +53,12 @@ def chunk_mesa_net_h_kv_bwd_intra_kernel_dkv( BV: tl.constexpr, IS_VARLEN: tl.constexpr, ): - i_t, i_bh = tl.program_id(0), tl.program_id(1) + i_t, i_bh = tl.program_id(0).to(tl.int64), tl.program_id(1).to(tl.int64) i_b, i_h = i_bh // H, i_bh % H 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.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) + i_n, i_t = tl.load(chunk_indices + i_t * 2).to(tl.int32), tl.load(chunk_indices + i_t * 2 + 1).to(tl.int64) + bos, eos = tl.load(cu_seqlens + i_n).to(tl.int64), tl.load(cu_seqlens + i_n + 1).to(tl.int64) T = eos - bos NT = tl.cdiv(T, BT) else: @@ -88,27 +88,34 @@ def chunk_mesa_net_h_kv_bwd_intra_kernel_dkv( b_dg_last = tl.zeros([1], dtype=tl.float32) b_dg = tl.zeros([BT], dtype=tl.float32) - p_v = tl.make_block_ptr(v, (T, V), (H*V, 1), (i_t * BT, 0), (BT, BV), (1, 0)) - p_k = tl.make_block_ptr(k, (T, K), (H*K, 1), (i_t * BT, 0), (BT, BK), (1, 0)) - p_beta = tl.make_block_ptr(beta, (T, ), (H, ), (i_t * BT,), (BT,), (0,)) - p_do = tl.make_block_ptr(do, (T, V), (H*V, 1), (i_t * BT, 0), (BT, BV), (1, 0)) - p_h = tl.make_block_ptr(h_kv, (V, K), (1, V), (0, 0), (BV, BK), (0, 1)) - p_dh = tl.make_block_ptr(dh_kv, (V, K), (1, V), (0, 0), (BV, BK), (0, 1)) - p_q = tl.make_block_ptr(q_star, (T, K), (H*K, 1), (i_t * BT, 0), (BT, BK), (1, 0)) - p_g = tl.make_block_ptr(g, (T,), (H,), (i_t * BT,), (BT,), (0,)) + o_k = tl.arange(0, BK) + o_v = tl.arange(0, BV) + m_tk = m_t[:, None] & (o_k[None, :] < K) + m_tv = m_t[:, None] & (o_v[None, :] < V) + m_h = (o_v[:, None] < V) & (o_k[None, :] < K) + p_v = v + o_t[:, None] * (H*V) + o_v[None, :] + p_k = k + o_t[:, None] * (H*K) + o_k[None, :] + p_beta = beta + o_t * H + p_do = do + o_t[:, None] * (H*V) + o_v[None, :] + p_h = h_kv + o_v[:, None] + o_k[None, :] * V + p_dh = dh_kv + o_v[:, None] + o_k[None, :] * V + p_q = q_star + o_t[:, None] * (H*K) + o_k[None, :] + p_g = g + o_t * H - b_q = tl.load(p_q, boundary_check=(0, 1)) - b_v = tl.load(p_v, boundary_check=(0, 1)) - b_k = tl.load(p_k, boundary_check=(0, 1)) - b_beta = tl.load(p_beta, boundary_check=(0, )) - b_g = tl.load(p_g, boundary_check=(0,)) - b_do = tl.load(p_do, boundary_check=(0, 1)) - b_h = tl.load(p_h, boundary_check=(0, 1)) - b_dh = tl.load(p_dh, boundary_check=(0, 1)) + b_q = tl.load(p_q, mask=m_tk, other=0.0) + b_v = tl.load(p_v, mask=m_tv, other=0.0) + b_k = tl.load(p_k, mask=m_tk, other=0.0) + b_beta = tl.load(p_beta, mask=m_t, other=0.0) + b_g = tl.load(p_g, mask=m_t, other=0.0) + b_do = tl.load(p_do, mask=m_tv, other=0.0) + b_h = tl.load(p_h, mask=m_h, other=0.0) + b_dh = tl.load(p_dh, mask=m_h, other=0.0) b_g_last = tl.load(g + (min(i_t * BT + BT, T) - 1) * H) # calculation - b_dg_last += tl.sum(b_h * b_dh) + # h_kv is stored in bf16 (wide exponent range) while dh_kv is fp16 + # lift both precision to float32 to prevent overflow from large h_kv entries + b_dg_last += tl.sum(b_h.to(tl.float32) * b_dh.to(tl.float32)) b_dg_last *= exp2(b_g_last) b_m = tl.where((o_t[:, None] >= o_t[None, :]) & (m_t[:, None] & m_t[None, :]), exp2(b_g[:, None] - b_g[None, :]), 0) @@ -125,14 +132,14 @@ def chunk_mesa_net_h_kv_bwd_intra_kernel_dkv( b_dg_last += tl.sum(b_dk * b_k) b_dg -= tl.sum(b_dk * b_k, axis=1) b_dv += tl.dot(b_k, tl.trans(b_dh).to(b_k.dtype)) * b_g_exp_k[:, None] + tl.dot(tl.trans(b_s.to(b_do.dtype)), b_do) - b_dk += tl.dot(tl.trans(b_ds.to(b_q.dtype)), b_q) + b_dk = tl.dot(tl.trans(b_ds.to(b_q.dtype)), b_q, b_dk) b_dg = tl.where(o_t < min(i_t * BT + BT, T) - 1, b_dg, b_dg + b_dg_last) - p_dk = tl.make_block_ptr(dk_beta, (T, K), (H*K, 1), (i_t * BT, 0), (BT, BK), (1, 0)) - p_dv = tl.make_block_ptr(dv, (T, V), (H*V, 1), (i_t * BT, 0), (BT, BV), (1, 0)) - p_dg = tl.make_block_ptr(dg, (T,), (H,), (i_t * BT,), (BT,), (0,)) - tl.store(p_dk, b_dk.to(p_dk.dtype.element_ty), boundary_check=(0, 1)) - tl.store(p_dv, b_dv.to(p_dv.dtype.element_ty), boundary_check=(0, 1)) - tl.store(p_dg, b_dg.to(p_dg.dtype.element_ty), boundary_check=(0,)) + p_dk = dk_beta + o_t[:, None] * (H*K) + o_k[None, :] + p_dv = dv + o_t[:, None] * (H*V) + o_v[None, :] + p_dg = dg + o_t * H + tl.store(p_dk, b_dk.to(p_dk.dtype.element_ty), mask=m_tk) + tl.store(p_dv, b_dv.to(p_dv.dtype.element_ty), mask=m_tv) + tl.store(p_dg, b_dg.to(p_dg.dtype.element_ty), mask=m_t) @triton.heuristics({ @@ -171,12 +178,12 @@ def chunk_mesa_net_h_kv_bwd_intra_kernel_dq( BV: tl.constexpr, IS_VARLEN: tl.constexpr, ): - i_t, i_bh = tl.program_id(0), tl.program_id(1) + i_t, i_bh = tl.program_id(0).to(tl.int64), tl.program_id(1).to(tl.int64) i_b, i_h = i_bh // H, i_bh % H 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.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) + i_n, i_t = tl.load(chunk_indices + i_t * 2).to(tl.int32), tl.load(chunk_indices + i_t * 2 + 1).to(tl.int64) + bos, eos = tl.load(cu_seqlens + i_n).to(tl.int64), tl.load(cu_seqlens + i_n + 1).to(tl.int64) T = eos - bos NT = tl.cdiv(T, BT) else: @@ -201,24 +208,29 @@ def chunk_mesa_net_h_kv_bwd_intra_kernel_dq( b_dq = tl.zeros([BT, BK], dtype=tl.float32) - p_v = tl.make_block_ptr(v, (T, V), (H*V, 1), (i_t * BT, 0), (BT, BV), (1, 0)) - p_k = tl.make_block_ptr(k, (T, K), (H*K, 1), (i_t * BT, 0), (BT, BK), (1, 0)) - p_beta = tl.make_block_ptr(beta, (T, ), (H, ), (i_t * BT,), (BT,), (0,)) - p_do = tl.make_block_ptr(do, (T, V), (H*V, 1), (i_t * BT, 0), (BT, BV), (1, 0)) - p_h = tl.make_block_ptr(h_kv, (V, K), (1, V), (0, 0), (BV, BK), (0, 1)) - p_q = tl.make_block_ptr(q_star, (T, K), (H*K, 1), (i_t * BT, 0), (BT, BK), (1, 0)) - p_g = tl.make_block_ptr(g, (T,), (H,), (i_t * BT,), (BT,), (0,)) - p_dg_prev = tl.make_block_ptr(dg_prev, (T,), (H,), (i_t * BT,), (BT,), (0,)) - p_dg = tl.make_block_ptr(dg, (T,), (H,), (i_t * BT,), (BT,), (0,)) - p_dq = tl.make_block_ptr(dq, (T, K), (H*K, 1), (i_t * BT, 0), (BT, BK), (1, 0)) + o_k = tl.arange(0, BK) + o_v = tl.arange(0, BV) + m_tk = m_t[:, None] & (o_k[None, :] < K) + m_tv = m_t[:, None] & (o_v[None, :] < V) + m_h = (o_v[:, None] < V) & (o_k[None, :] < K) + p_v = v + o_t[:, None] * (H*V) + o_v[None, :] + p_k = k + o_t[:, None] * (H*K) + o_k[None, :] + p_beta = beta + o_t * H + p_do = do + o_t[:, None] * (H*V) + o_v[None, :] + p_h = h_kv + o_v[:, None] + o_k[None, :] * V + p_q = q_star + o_t[:, None] * (H*K) + o_k[None, :] + p_g = g + o_t * H + p_dg_prev = dg_prev + o_t * H + p_dg = dg + o_t * H + p_dq = dq + o_t[:, None] * (H*K) + o_k[None, :] - b_q = tl.load(p_q, boundary_check=(0, 1)) - b_v = tl.load(p_v, boundary_check=(0, 1)) - b_k = tl.load(p_k, boundary_check=(0, 1)) - b_beta = tl.load(p_beta, boundary_check=(0, )) - b_g = tl.load(p_g, boundary_check=(0,)) - b_do = tl.load(p_do, boundary_check=(0, 1)) - b_h = tl.load(p_h, boundary_check=(0, 1)) + b_q = tl.load(p_q, mask=m_tk, other=0.0) + b_v = tl.load(p_v, mask=m_tv, other=0.0) + b_k = tl.load(p_k, mask=m_tk, other=0.0) + b_beta = tl.load(p_beta, mask=m_t, other=0.0) + b_g = tl.load(p_g, mask=m_t, other=0.0) + b_do = tl.load(p_do, mask=m_tv, other=0.0) + b_h = tl.load(p_h, mask=m_h, other=0.0) b_m = tl.where((o_t[:, None] >= o_t[None, :]) & (m_t[:, None] & m_t[None, :]), exp2(b_g[:, None] - b_g[None, :]), 0) b_k = (b_k * b_beta[:, None]).to(b_k.dtype) @@ -226,11 +238,11 @@ def chunk_mesa_net_h_kv_bwd_intra_kernel_dq( b_ds = tl.dot(b_do, tl.trans(b_v)) * b_m b_g_exp_q = exp2(b_g) b_dq = tl.dot(b_do, b_h.to(b_do.dtype)) * b_g_exp_q[:, None] - b_dg = tl.sum(b_dq * b_q, axis=1) + tl.load(p_dg_prev, boundary_check=(0,)) - b_dq += tl.dot(b_ds.to(b_k.dtype), b_k) + b_dg = tl.sum(b_dq * b_q, axis=1) + tl.load(p_dg_prev, mask=m_t, other=0.0) + b_dq = tl.dot(b_ds.to(b_k.dtype), b_k, b_dq) - tl.store(p_dq, b_dq.to(p_dq.dtype.element_ty), boundary_check=(0, 1)) - tl.store(p_dg, b_dg.to(p_dg.dtype.element_ty), boundary_check=(0,)) + tl.store(p_dq, b_dq.to(p_dq.dtype.element_ty), mask=m_tk) + tl.store(p_dg, b_dg.to(p_dg.dtype.element_ty), mask=m_t) def chunk_mesa_net_h_kv_bwd_intra_separate_fn( diff --git a/build/torch-cuda/ops/mesa_net/decoding_one_step.py b/build/torch-cuda/ops/mesa_net/decoding_one_step.py index a2277af73ac435297524367beacba0698889722c..74e8c1dca6ac12b3f12a536a99f5327859352c5b 100644 --- a/build/torch-cuda/ops/mesa_net/decoding_one_step.py +++ b/build/torch-cuda/ops/mesa_net/decoding_one_step.py @@ -34,7 +34,7 @@ def mesa_net_decoding_one_step_kernel( BV: tl.constexpr, MAX_CG_STEP: tl.constexpr, ): - i_nh = tl.program_id(0) + i_nh = tl.program_id(0).to(tl.int64) i_h = i_nh % H o_k = tl.arange(0, BK) diff --git a/build/torch-cuda/ops/moba/parallel.py b/build/torch-cuda/ops/moba/parallel.py index 054b96394b30e82fd7b0fc46de7c9eff78a767b0..f6c08e03a9a29bd661c3ee1baa9dc1edef1e2b49 100644 --- a/build/torch-cuda/ops/moba/parallel.py +++ b/build/torch-cuda/ops/moba/parallel.py @@ -365,7 +365,7 @@ def parallel_moba( into 128-token blocks with each query attending to 4 blocks: >>> import torch - >>> from ...ops.moba import parallel_moba + >>> from fla.ops.moba import parallel_moba >>> B, T, H, D = 1, 1024, 4, 64 >>> q = torch.randn(B, T, H, D, dtype=torch.float16, device='cuda') >>> k = torch.randn(B, T, H, D, dtype=torch.float16, device='cuda') diff --git a/build/torch-cuda/ops/momentum_delta_rule/__init__.py b/build/torch-cuda/ops/momentum_delta_rule/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..f4021e41e0c529b5d93b550e55e720f4089a8c34 --- /dev/null +++ b/build/torch-cuda/ops/momentum_delta_rule/__init__.py @@ -0,0 +1,17 @@ +# Copyright (c) 2023-2026, Songlin Yang, Yu Zhang, Zhiyuan Li +# +# This source code is licensed under the MIT license found in the +# LICENSE file in the root directory of this source tree. +# For a list of all contributors, visit: +# https://github.com/fla-org/flash-linear-attention/graphs/contributors + +from .chunk import chunk_momentum_delta_rule +from .fused_recurrent import fused_recurrent_momentum_delta_rule +from .naive import chunk_momentum_delta_rule_ref, recurrent_momentum_delta_rule_ref + +__all__ = [ + 'chunk_momentum_delta_rule', + 'chunk_momentum_delta_rule_ref', + 'fused_recurrent_momentum_delta_rule', + 'recurrent_momentum_delta_rule_ref', +] diff --git a/build/torch-cuda/ops/momentum_delta_rule/chunk.py b/build/torch-cuda/ops/momentum_delta_rule/chunk.py new file mode 100644 index 0000000000000000000000000000000000000000..c7f375ad9b8555fda121c5a97bb61c22dff7c862 --- /dev/null +++ b/build/torch-cuda/ops/momentum_delta_rule/chunk.py @@ -0,0 +1,118 @@ +# 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 + +from ...modules.norm.l2norm import l2norm_bwd, l2norm_fwd +from ...ops.momentum_delta_rule.naive import chunk_momentum_delta_rule_ref +from ...utils import autocast_custom_bwd, autocast_custom_fwd, input_guard + + +class ChunkMomentumDeltaRuleFunction(torch.autograd.Function): + @staticmethod + @input_guard + @autocast_custom_fwd + def forward(ctx, q, k, v, p, log_alpha, log_mu, beta, eta, scale, initial_S, initial_M, output_final_state, cu_seqlens, use_qk_l2norm_in_kernel, use_p_times_alpha, chunk_size): + if use_qk_l2norm_in_kernel: + q, q_rstd = l2norm_fwd(q) + k, k_rstd = l2norm_fwd(k) + p, p_rstd = l2norm_fwd(p) + else: + q_rstd, k_rstd, p_rstd = None, None, None + if cu_seqlens is not None: + raise NotImplementedError("Variable-length `cu_seqlens` not yet supported for full momentum PyTorch path.") + k_eta = (k * eta.unsqueeze(-1)).to(q.dtype) + p_eff = p if not use_p_times_alpha else (p * log_alpha.exp().unsqueeze(-1)).to(q.dtype) + o, final_state = chunk_momentum_delta_rule_ref( + q=q, k=k_eta, v=v, p=p_eff, log_alpha=log_alpha, log_mu=log_mu, + beta=beta, eta=torch.ones_like(beta), scale=scale, + initial_S=initial_S, initial_M=initial_M, + output_final_state=output_final_state, chunk_size=chunk_size) + ctx.save_for_backward(q, k, v, p, eta, beta, log_alpha, log_mu, initial_S, + initial_M, cu_seqlens, q_rstd, k_rstd, p_rstd) + ctx.scale = scale + ctx.use_qk_l2norm_in_kernel = use_qk_l2norm_in_kernel + ctx.use_p_times_alpha = use_p_times_alpha + ctx.chunk_size = chunk_size + final_S, final_M = (final_state[0], final_state[1]) if final_state is not None else (None, None) + return o.to(q.dtype), final_S, final_M + + @staticmethod + @input_guard + @autocast_custom_bwd + def backward(ctx, do, dst, dmt): + q, k, v, p, eta, beta, log_alpha, log_mu, initial_S, initial_M, cu_seqlens, q_rstd, k_rstd, p_rstd = ctx.saved_tensors + with torch.enable_grad(): + q_r = q.detach().requires_grad_(True) + k_r = k.detach().requires_grad_(True) + v_r = v.detach().requires_grad_(True) + p_r = p.detach().requires_grad_(True) + log_alpha_r = log_alpha.detach().requires_grad_(True) + log_mu_r = log_mu.detach().requires_grad_(True) + beta_r = beta.detach().requires_grad_(True) + eta_r = eta.detach().requires_grad_(True) if eta is not None else None + k_eta = k_r if eta_r is None else k_r * eta_r.unsqueeze(-1) + p_eff = p_r if not ctx.use_p_times_alpha else p_r * log_alpha_r.exp().unsqueeze(-1) + initial_S_r = initial_S.detach().requires_grad_(True) if initial_S is not None else None + initial_M_r = initial_M.detach().requires_grad_(True) if initial_M is not None else None + o, final_state = chunk_momentum_delta_rule_ref( + q=q_r, k=k_eta, v=v_r, p=p_eff, log_alpha=log_alpha_r, log_mu=log_mu_r, + beta=beta_r, eta=torch.ones_like(beta_r), scale=ctx.scale, + initial_S=initial_S_r, initial_M=initial_M_r, output_final_state=True, chunk_size=ctx.chunk_size) + differentiable_inputs = (q_r, k_r, v_r, p_r, log_alpha_r, log_mu_r, beta_r) + if eta_r is not None: + differentiable_inputs += (eta_r,) + if initial_S_r is not None: + differentiable_inputs += (initial_S_r, initial_M_r) + grads = torch.autograd.grad( + (o, final_state[0], final_state[1]), differentiable_inputs, + grad_outputs=( + do, + torch.zeros_like(final_state[0]) if dst is None else dst, + torch.zeros_like(final_state[1]) if dmt is None else dmt, + ), + retain_graph=False, allow_unused=True, + ) + dq, dk, dv, dp, dlog_alpha, dlog_mu, dbeta = grads[:7] + deta = grads[7] if eta is not None else None + state_offset = 8 if eta is not None else 7 + ds0, dm0 = (grads[state_offset:state_offset + 2] if initial_S is not None else (None, None)) + if ctx.use_qk_l2norm_in_kernel: + if dq is not None: + dq = l2norm_bwd(q, q_rstd, dq.contiguous()) + if dk is not None: + dk = l2norm_bwd(k, k_rstd, dk.contiguous()) + if dp is not None: + dp = l2norm_bwd(p, p_rstd, dp.contiguous()) + return dq, dk, dv, dp, dlog_alpha, dlog_mu, dbeta, deta, None, ds0, dm0, None, None, None, None, None + + +@torch.compiler.disable +def chunk_momentum_delta_rule(q, k, v, log_alpha, log_mu, p=None, beta=None, eta=None, scale=None, initial_state=None, output_final_state=False, cu_seqlens=None, use_qk_l2norm_in_kernel=True, use_p_times_alpha=True, chunk_size=64): + assert q.dtype == k.dtype == v.dtype + assert q.dtype != torch.float32, "ChunkMomentumDeltaRuleFunction does not support float32. Please use bfloat16." + if chunk_size not in (16, 32, 64): + raise ValueError(f"`chunk_size` must be 16, 32, or 64, got {chunk_size}.") + if p is None: + p = k + if scale is None: + scale = k.shape[-1] ** -0.5 + if beta is None: + beta = torch.ones_like(q[..., 0]) + if eta is None: + eta = torch.ones_like(q[..., 0]) + if initial_state is not None: + initial_S, initial_M = initial_state[0], initial_state[1] + else: + initial_S, initial_M = None, None + if cu_seqlens is not None: + raise NotImplementedError( + "Variable-length `cu_seqlens` is not yet supported for the full momentum reference path.") + o, final_S, final_M = ChunkMomentumDeltaRuleFunction.apply( + q, k, v, p, log_alpha, log_mu, beta, eta, scale, initial_S, initial_M, output_final_state, cu_seqlens, use_qk_l2norm_in_kernel, use_p_times_alpha, chunk_size) + final_state = torch.stack([final_S, final_M], dim=0) if output_final_state else None + return o, final_state diff --git a/build/torch-cuda/ops/momentum_delta_rule/fused_recurrent.py b/build/torch-cuda/ops/momentum_delta_rule/fused_recurrent.py new file mode 100644 index 0000000000000000000000000000000000000000..026d5ee3625d05573e7fdd683911923815f5dd4b --- /dev/null +++ b/build/torch-cuda/ops/momentum_delta_rule/fused_recurrent.py @@ -0,0 +1,103 @@ +# 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 + +from ...modules.norm.l2norm import l2norm_bwd, l2norm_fwd +from ...ops.momentum_delta_rule.naive import recurrent_momentum_delta_rule_ref +from ...utils import autocast_custom_bwd, autocast_custom_fwd, input_guard + + +class FusedRecurrentMomentumDeltaRuleFunction(torch.autograd.Function): + @staticmethod + @input_guard + @autocast_custom_fwd + def forward(ctx, q, k, v, p, log_alpha, log_mu, beta, eta, scale, initial_S, initial_M, output_final_state, cu_seqlens, use_qk_l2norm_in_kernel, use_p_times_alpha): + if use_qk_l2norm_in_kernel: + q, q_rstd = l2norm_fwd(q) + k, k_rstd = l2norm_fwd(k) + p, p_rstd = l2norm_fwd(p) + else: + q_rstd, k_rstd, p_rstd = None, None, None + if cu_seqlens is not None: + raise NotImplementedError("Variable-length `cu_seqlens` is not yet supported for full momentum PyTorch path.") + k_eta = (k * eta.unsqueeze(-1)).to(q.dtype) + p_eff = p if not use_p_times_alpha else (p * log_alpha.exp().unsqueeze(-1)).to(q.dtype) + o, final_state = recurrent_momentum_delta_rule_ref( + q=q, k=k_eta, v=v, p=p_eff, log_alpha=log_alpha, log_mu=log_mu, beta=beta, + eta=torch.ones_like(beta), scale=scale, initial_S=initial_S, initial_M=initial_M, + output_final_state=output_final_state, + ) + ctx.save_for_backward( + q, k, v, p, eta, beta, log_alpha, log_mu, initial_S, initial_M, cu_seqlens, q_rstd, k_rstd, p_rstd, + ) + ctx.scale = scale + ctx.use_qk_l2norm_in_kernel = use_qk_l2norm_in_kernel + ctx.use_p_times_alpha = use_p_times_alpha + final_S, final_M = (final_state[0], final_state[1]) if final_state is not None else (None, None) + return o.to(q.dtype), final_S, final_M + + @staticmethod + @input_guard + @autocast_custom_bwd + def backward(ctx, do, dst, dmt): + q, k, v, p, eta, beta, log_alpha, log_mu, initial_S, initial_M, cu_seqlens, q_rstd, k_rstd, p_rstd = ctx.saved_tensors + with torch.enable_grad(): + q_r, k_r, v_r, p_r = [x.detach().requires_grad_(True) for x in (q, k, v, p)] + log_alpha_r = log_alpha.detach().requires_grad_(True) + log_mu_r = log_mu.detach().requires_grad_(True) + beta_r = beta.detach().requires_grad_(True) + eta_r = eta.detach().requires_grad_(True) if eta is not None else None + k_eta = k_r if eta_r is None else k_r * eta_r.unsqueeze(-1) + p_eff = p_r if not ctx.use_p_times_alpha else p_r * log_alpha_r.exp().unsqueeze(-1) + initial_S_r = initial_S.detach().requires_grad_(True) if initial_S is not None else None + initial_M_r = initial_M.detach().requires_grad_(True) if initial_M is not None else None + o, final_state = recurrent_momentum_delta_rule_ref( + q=q_r, k=k_eta, v=v_r, p=p_eff, log_alpha=log_alpha_r, log_mu=log_mu_r, + beta=beta_r, eta=torch.ones_like(beta_r), scale=ctx.scale, initial_S=initial_S_r, + initial_M=initial_M_r, output_final_state=True, + ) + inputs = (q_r, k_r, v_r, p_r, log_alpha_r, log_mu_r, beta_r) + if eta_r is not None: + inputs += (eta_r,) + if initial_S_r is not None: + inputs += (initial_S_r, initial_M_r) + grads = torch.autograd.grad( + (o, final_state[0], final_state[1]), inputs, + grad_outputs=( + do, + torch.zeros_like(final_state[0]) if dst is None else dst, + torch.zeros_like(final_state[1]) if dmt is None else dmt, + ), + allow_unused=True, + ) + dq, dk, dv, dp, dlog_alpha, dlog_mu, dbeta = grads[:7] + deta = grads[7] if eta is not None else None + state_offset = 8 if eta is not None else 7 + ds0, dm0 = grads[state_offset:state_offset + 2] if initial_S is not None else (None, None) + if ctx.use_qk_l2norm_in_kernel: + dq = l2norm_bwd(q, q_rstd, dq.contiguous()) + dk = l2norm_bwd(k, k_rstd, dk.contiguous()) + dp = l2norm_bwd(p, p_rstd, dp.contiguous()) + return dq, dk, dv, dp, dlog_alpha, dlog_mu, dbeta, deta, None, ds0, dm0, None, None, None, None + + +@torch.compiler.disable +def fused_recurrent_momentum_delta_rule(q, k, v, log_alpha, log_mu, p=None, beta=None, eta=None, scale=None, initial_state=None, output_final_state=False, cu_seqlens=None, use_qk_l2norm_in_kernel=True, use_p_times_alpha=True): + scale = k.shape[-1] ** -0.5 if scale is None else scale + beta = torch.ones_like(q[..., 0]) if beta is None else beta + eta = torch.ones_like(q[..., 0]) if eta is None else eta + p = k if p is None else p + initial_S, initial_M = (initial_state[0], initial_state[1]) if initial_state is not None else (None, None) + if cu_seqlens is not None: + raise NotImplementedError("Variable-length `cu_seqlens` is not yet supported for full momentum PyTorch path.") + o, final_S, final_M = FusedRecurrentMomentumDeltaRuleFunction.apply( + q, k, v, p, log_alpha, log_mu, beta, eta, scale, initial_S, initial_M, + output_final_state, cu_seqlens, use_qk_l2norm_in_kernel, use_p_times_alpha, + ) + final_state = torch.stack([final_S, final_M], dim=0) if output_final_state else None + return o, final_state diff --git a/build/torch-cuda/ops/momentum_delta_rule/naive.py b/build/torch-cuda/ops/momentum_delta_rule/naive.py new file mode 100644 index 0000000000000000000000000000000000000000..294e4613fd8794e2de7e289b19603f4f8abe8dc8 --- /dev/null +++ b/build/torch-cuda/ops/momentum_delta_rule/naive.py @@ -0,0 +1,222 @@ +# 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 torch.nn.functional as F +from einops import rearrange + + +def recurrent_momentum_delta_rule_ref( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + log_alpha: torch.Tensor, + log_mu: torch.Tensor, + p: torch.Tensor, + beta: torch.Tensor | None = None, + eta: torch.Tensor | None = None, + scale: float | None = None, + initial_S: torch.Tensor | None = None, + initial_M: torch.Tensor | None = None, + output_final_state: bool = False, +): + q, k, v, p, log_alpha, log_mu, beta, eta = map( + lambda x: x.to(torch.float32), + [q, k, v, p, log_alpha, log_mu, beta, eta], + ) + B, T, H, DK, DV = *k.shape, v.shape[-1] + + if scale is None: + scale = 1 / (q.shape[-1] ** 0.5) + + q = q * scale + S_prev = torch.zeros(B, H, DK, DV, device=v.device, dtype=torch.float32) + M_prev = torch.zeros(B, H, DK, DV, device=v.device, dtype=torch.float32) + + if initial_M is not None: + M_prev = initial_M.to(torch.float32) + if initial_S is not None: + S_prev = initial_S.to(torch.float32) + + out = torch.zeros_like(v, dtype=torch.float32) + for i in range(T): + k_t = k[:, i] + q_t = q[:, i] + v_t = v[:, i] + p_t = p[:, i] + + mu_i = log_mu[:, i].exp().view(B, H, 1, 1) + beta_i = beta[:, i].view(B, H, 1, 1) + alpha_i = log_alpha[:, i].exp().view(B, H, 1, 1) + + eta_i = eta[:, i].unsqueeze(-1) + + w_t = -(v_t.unsqueeze(-2) - p_t.unsqueeze(-2) @ S_prev) + Mt = mu_i * M_prev + (eta_i * k_t).unsqueeze(-1) @ w_t + St = alpha_i * S_prev - beta_i * Mt + + out[:, i] = (q_t.unsqueeze(-1) * St).sum(-2) + + M_prev = Mt + S_prev = St + + o = out.to(q.dtype if hasattr(q, 'dtype') else v.dtype) + if output_final_state: + final_state = torch.stack([S_prev, M_prev], dim=0) + else: + final_state = None + + return o, final_state + + +def chunk_momentum_delta_rule_ref( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + p: torch.Tensor, + log_alpha: torch.Tensor, + log_mu: torch.Tensor, + beta: torch.Tensor, + eta: torch.Tensor, + scale: float | None = None, + initial_S: torch.Tensor | None = None, + initial_M: torch.Tensor | None = None, + output_final_state: bool = False, + chunk_size: int = 64, +): + BT = chunk_size + q_input, k_input, v_input, p_input = q, k, v, p + log_alpha_input, log_mu_input, beta_input, eta_input = log_alpha, log_mu, beta, eta + if scale is None: + scale = 1 / (q.shape[-1] ** 0.5) + q, k, v, p, log_alpha, log_mu, beta, eta = map( + lambda x: x.to(torch.float32), + [q, k, v, p, log_alpha, log_mu, beta, eta], + ) + T = q.shape[1] + pad_len = (BT - (T % BT)) % BT + + if pad_len > 0: + q = F.pad(q, (0, 0, 0, 0, 0, pad_len)) + k = F.pad(k, (0, 0, 0, 0, 0, pad_len)) + v = F.pad(v, (0, 0, 0, 0, 0, pad_len)) + p = F.pad(p, (0, 0, 0, 0, 0, pad_len)) + log_alpha = F.pad(log_alpha, (0, 0, 0, pad_len)) + log_mu = F.pad(log_mu, (0, 0, 0, pad_len)) + beta = F.pad(beta, (0, 0, 0, pad_len)) + eta = F.pad(eta, (0, 0, 0, pad_len)) + + B, l, H, DK = q.shape + DV = v.shape[-1] + q = q * scale + assert l % chunk_size == 0 + assert q.shape == (B, pad_len + T, H, DK) + assert log_alpha.shape == (B, pad_len + T, H) + + k_eta = eta[..., None] * k + q, k, v, p, log_alpha, log_mu, beta = map( + lambda x: rearrange(x, 'b (n c) h d -> b h n c d', c=chunk_size), + [q, k_eta, v, p, log_alpha.unsqueeze(-1), log_mu.unsqueeze(-1), beta.unsqueeze(-1)], + ) + + log_a_cum = log_alpha.squeeze(-1).cumsum(-1) + log_m_cum = log_mu.squeeze(-1).cumsum(-1) + log_beta = (beta + 1e-6).squeeze(-1).log() + + log_c_before = log_beta + log_m_cum - log_a_cum + log_ct = torch.logcumsumexp(log_c_before, dim=-1) + log_ct_tm1 = torch.cat( + [torch.full_like(log_ct[:, :, :, :1], float('-inf')), log_ct[:, :, :, :-1]], dim=-1, + ) + + a = log_ct.unsqueeze(-1) + b = log_ct_tm1.unsqueeze(-2) + x = (b - a).tril() + temp = 1 - torch.exp(x) + + log_bar_a_tm1 = torch.cat( + [torch.zeros_like(log_a_cum[:, :, :, :1]), log_a_cum[:, :, :, :-1]], dim=3, + ) + bar_a_tm1 = log_bar_a_tm1.exp() + + b_t = (log_a_cum + log_ct).exp() + b_tm1 = torch.cat([torch.zeros_like(b_t[:, :, :, :1]), b_t[:, :, :, :-1]], dim=3) + + gamma_mask_q = (log_a_cum.unsqueeze(-1) - log_m_cum.unsqueeze(-2) + a).exp().float().tril() * temp + gamma_mask = torch.cat( + [torch.zeros_like(gamma_mask_q[:, :, :, :1]), gamma_mask_q[:, :, :, :-1]], dim=3, + ) + + attn = (p @ k.transpose(-1, -2)) * gamma_mask + + attn_inv = -attn + for i in range(1, chunk_size): + attn_inv[..., i, :i] += (attn_inv[..., i, :, None].clone() * attn_inv[..., :, :i].clone()).sum(-2) + attn_inv = attn_inv + torch.eye(chunk_size, dtype=attn_inv.dtype, device=q.device) + + alpha_tm1_p = bar_a_tm1[..., None] * p + b_tm1_p = b_tm1[..., None] * p + + u_c = attn_inv @ v + y_c = attn_inv @ alpha_tm1_p + z_c = attn_inv @ b_tm1_p + + S_pre = k.new_zeros(B, H, DK, DV) + M_pre = k.new_zeros(B, H, DK, DV) + + if initial_S is not None: + S_pre = initial_S.to(torch.float32) + if initial_M is not None: + M_pre = initial_M.to(torch.float32) + + o = torch.zeros_like(v) + num_chunks = q.shape[2] + for i in range(num_chunks): + q_i, k_i, = q[:, :, i], k[:, :, i] + v_i = u_c[:, :, i] - y_c[:, :, i] @ S_pre + z_c[:, :, i] @ M_pre + + attn_inner = (q_i @ k_i.transpose(-1, -2)) * gamma_mask_q[:, :, i] + bar_alpha_t_q = q_i * log_a_cum[:, :, i, :].exp().unsqueeze(-1) + b_t_q = q_i * b_t[:, :, i, :].unsqueeze(-1) + qS_inter = bar_alpha_t_q @ S_pre - b_t_q @ M_pre + o[:, :, i] = qS_inter + attn_inner @ v_i + + decay_s = gamma_mask_q[:, :, i, -1].unsqueeze(-1) + S = log_a_cum[:, :, i, -1, None, None].exp() * S_pre \ + - b_t[:, :, i, -1, None, None] * M_pre \ + + (k_i * decay_s).transpose(-1, -2) @ v_i + + decay_m = (log_m_cum[:, :, i, -1, None] - log_m_cum[:, :, i]).exp()[..., None] + M = log_m_cum[:, :, i, -1, None, None].exp() * M_pre \ + - (k_i * decay_m).transpose(-1, -2) @ v_i + + S_pre, M_pre = S, M + + if output_final_state: + if pad_len: + _, final_state = recurrent_momentum_delta_rule_ref( + q=q_input, + k=k_input, + v=v_input, + p=p_input, + log_alpha=log_alpha_input, + log_mu=log_mu_input, + beta=beta_input, + eta=eta_input, + scale=scale, + initial_S=initial_S, + initial_M=initial_M, + output_final_state=True, + ) + else: + final_state = torch.stack([S_pre, M_pre], dim=0) + else: + final_state = None + + o = rearrange(o, 'b h n c d -> b (n c) h d') + o = o[:, :T] + return o, final_state diff --git a/build/torch-cuda/ops/nsa/compression.py b/build/torch-cuda/ops/nsa/compression.py index 6f51d36d41ddcd43595c85ef6783cefeffcbdbf2..1901a5ede84632088a91f4d2c4b91d5d20e892bc 100644 --- a/build/torch-cuda/ops/nsa/compression.py +++ b/build/torch-cuda/ops/nsa/compression.py @@ -11,7 +11,7 @@ import triton.language as tl from ...ops.attn.parallel import parallel_attn_bwd_preprocess from ...ops.utils import prepare_chunk_indices, prepare_chunk_offsets, prepare_token_indices -from ...ops.utils.op import exp, log +from ...ops.utils.op import exp, log, unflatten_program_id from ...utils import autocast_custom_bwd, autocast_custom_fwd, autotune_cache_kwargs, check_shared_mem, contiguous @@ -51,33 +51,38 @@ def parallel_nsa_compression_fwd_kernel( BV: tl.constexpr, IS_VARLEN: tl.constexpr, ): - i_t, i_v, i_bh = tl.program_id(0), tl.program_id(1), tl.program_id(2) + i_t, i_v, i_bh = tl.program_id(0).to(tl.int64), tl.program_id(1).to(tl.int64), tl.program_id(2).to(tl.int64) i_b, i_h = i_bh // H, i_bh % H if IS_VARLEN: - i_n, i_t = tl.load(token_indices_q + i_t * 2).to(tl.int32), tl.load(token_indices_q + i_t * 2 + 1).to(tl.int32) + i_n, i_t = tl.load(token_indices_q + i_t * 2).to(tl.int32), tl.load(token_indices_q + i_t * 2 + 1).to(tl.int64) bos_q, eos_q = tl.load(cu_seqlens_q + i_n).to(tl.int64), tl.load(cu_seqlens_q + i_n + 1).to(tl.int64) bos_k, eos_k = tl.load(cu_seqlens_k + i_n).to(tl.int64), tl.load(cu_seqlens_k + i_n + 1).to(tl.int64) TQ = (eos_q - bos_q).to(tl.int32) TK = (eos_k - bos_k).to(tl.int32) TC = tl.cdiv(TK, BS) - boc = tl.load(chunk_offsets + i_n).to(tl.int32) + boc = tl.load(chunk_offsets + i_n).to(tl.int64) else: bos_q, eos_q = (i_b * TQ).to(tl.int64), (i_b * TQ + TQ).to(tl.int64) TC = tl.cdiv(TK, BS) boc = i_b * TC - p_q = tl.make_block_ptr(q + (bos_q + i_t) * HQ*K, (HQ, K), (K, 1), (i_h * G, 0), (G, BK), (1, 0)) + o_g = i_h * G + tl.arange(0, G) + o_d = tl.arange(0, BK) + m_q = (o_g[:, None] < HQ) & (o_d[None, :] < K) + p_q = q + (bos_q + i_t) * HQ*K + o_g[:, None] * K + o_d[None, :] Q_OFFSET = TK - TQ # the Q block is kept in the shared memory throughout the whole kernel # [G, BK] - b_q = tl.load(p_q, boundary_check=(0, 1)) + b_q = tl.load(p_q, mask=m_q, other=0.0) b_q = (b_q * scale).to(b_q.dtype) # number of complete compression blocks visible to the query (q tokens are the last TQ of the sequence) NC = (i_t + Q_OFFSET + 1) // BS - p_o = tl.make_block_ptr(o + (bos_q + i_t) * HQ*V, (HQ, V), (V, 1), (i_h * G, i_v * BV), (G, BV), (1, 0)) + o_v = i_v * BV + tl.arange(0, BV) + m_o = (o_g[:, None] < HQ) & (o_v[None, :] < V) + p_o = o + (bos_q + i_t) * HQ*V + o_g[:, None] * V + o_v[None, :] # [G, BV] b_o = tl.zeros([G, BV], dtype=tl.float32) # max scores for the current block @@ -88,12 +93,12 @@ def parallel_nsa_compression_fwd_kernel( for i_c in range(0, NC, BC): o_c = i_c + tl.arange(0, BC) - p_k = tl.make_block_ptr(k + (boc * H + i_h) * K, (K, TC), (1, H*K), (0, i_c), (BK, BC), (0, 1)) - p_v = tl.make_block_ptr(v + (boc * H + i_h) * V, (TC, V), (H*V, 1), (i_c, i_v * BV), (BC, BV), (1, 0)) + p_k = k + (boc * H + i_h) * K + o_d[:, None] + o_c[None, :] * (H*K) + p_v = v + (boc * H + i_h) * V + o_c[:, None] * (H*V) + o_v[None, :] # [BK, BC] - b_k = tl.load(p_k, boundary_check=(0, 1)) + b_k = tl.load(p_k, mask=(o_d[:, None] < K) & (o_c[None, :] < TC), other=0.0) # [BC, BV] - b_v = tl.load(p_v, boundary_check=(0, 1)) + b_v = tl.load(p_v, mask=(o_c[:, None] < TC) & (o_v[None, :] < V), other=0.0) # [G, BC] b_s = tl.dot(b_q, b_k) # causal mask: only the NC complete blocks are visible @@ -114,7 +119,7 @@ def parallel_nsa_compression_fwd_kernel( b_o = b_o / b_acc[:, None] b_lse = b_m + log(b_acc) - tl.store(p_o, b_o.to(p_o.dtype.element_ty), boundary_check=(0, 1)) + tl.store(p_o, b_o.to(p_o.dtype.element_ty), mask=m_o) if i_v == 0: tl.store(lse + (bos_q + i_t) * HQ + i_h * G + tl.arange(0, G), b_lse.to(lse.dtype.element_ty)) @@ -156,15 +161,15 @@ def parallel_nsa_compression_bwd_kernel_dq( BV: tl.constexpr, IS_VARLEN: tl.constexpr, ): - i_t, i_v, i_bh = tl.program_id(0), tl.program_id(1), tl.program_id(2) + i_t, i_v, i_bh = tl.program_id(0).to(tl.int64), tl.program_id(1).to(tl.int64), tl.program_id(2).to(tl.int64) i_b, i_h = i_bh // H, i_bh % H - all = B * T + all = B * T.to(tl.int64) if IS_VARLEN: - i_n, i_t = tl.load(token_indices + i_t * 2).to(tl.int32), tl.load(token_indices + i_t * 2 + 1).to(tl.int32) + i_n, i_t = tl.load(token_indices + i_t * 2).to(tl.int32), tl.load(token_indices + i_t * 2 + 1).to(tl.int64) bos, eos = tl.load(cu_seqlens + i_n).to(tl.int64), tl.load(cu_seqlens + i_n + 1).to(tl.int64) T = (eos - bos).to(tl.int32) - boc = tl.load(chunk_offsets + i_n).to(tl.int32) + boc = tl.load(chunk_offsets + i_n).to(tl.int64) else: bos, eos = (i_b * T).to(tl.int64), (i_b * T + T).to(tl.int64) boc = i_b * tl.cdiv(T, BS) @@ -175,14 +180,19 @@ def parallel_nsa_compression_bwd_kernel_dq( delta += (bos + i_t) * HQ dq += (i_v * all + bos + i_t) * HQ*K - p_q = tl.make_block_ptr(q, (HQ, K), (K, 1), (i_h * G, 0), (G, BK), (1, 0)) - p_dq = tl.make_block_ptr(dq, (HQ, K), (K, 1), (i_h * G, 0), (G, BK), (1, 0)) + o_g = i_h * G + tl.arange(0, G) + o_d = tl.arange(0, BK) + m_q = (o_g[:, None] < HQ) & (o_d[None, :] < K) + p_q = q + o_g[:, None] * K + o_d[None, :] + p_dq = dq + o_g[:, None] * K + o_d[None, :] # [G, BK] - b_q = tl.load(p_q, boundary_check=(0, 1)) + b_q = tl.load(p_q, mask=m_q, other=0.0) b_q = (b_q * scale).to(b_q.dtype) - p_do = tl.make_block_ptr(do, (HQ, V), (V, 1), (i_h * G, i_v * BV), (G, BV), (1, 0)) + o_v = i_v * BV + tl.arange(0, BV) + m_o = (o_g[:, None] < HQ) & (o_v[None, :] < V) + p_do = do + o_g[:, None] * V + o_v[None, :] p_lse = lse + i_h * G + tl.arange(0, G) p_delta = delta + i_h * G + tl.arange(0, G) @@ -193,7 +203,7 @@ def parallel_nsa_compression_bwd_kernel_dq( NC = (i_t + 1) // BS # [G, BV] - b_do = tl.load(p_do, boundary_check=(0, 1)) + b_do = tl.load(p_do, mask=m_o, other=0.0) # [G] b_lse = tl.load(p_lse) b_delta = tl.load(p_delta) @@ -202,12 +212,12 @@ def parallel_nsa_compression_bwd_kernel_dq( b_dq = tl.zeros([G, BK], dtype=tl.float32) for i_c in range(0, NC, BC): o_c = i_c + tl.arange(0, BC) - p_k = tl.make_block_ptr(k + (boc * H + i_h) * K, (K, TC), (1, H*K), (0, i_c), (BK, BC), (0, 1)) - p_v = tl.make_block_ptr(v + (boc * H + i_h) * V, (V, TC), (1, H*V), (i_v * BV, i_c), (BV, BC), (0, 1)) + p_k = k + (boc * H + i_h) * K + o_d[:, None] + o_c[None, :] * (H*K) + p_v = v + (boc * H + i_h) * V + o_v[:, None] + o_c[None, :] * (H*V) # [BK, BC] - b_k = tl.load(p_k, boundary_check=(0, 1)) + b_k = tl.load(p_k, mask=(o_d[:, None] < K) & (o_c[None, :] < TC), other=0.0) # [BV, BC] - b_v = tl.load(p_v, boundary_check=(0, 1)) + b_v = tl.load(p_v, mask=(o_v[:, None] < V) & (o_c[None, :] < TC), other=0.0) # [G, BC] b_s = tl.dot(b_q, b_k) @@ -218,9 +228,33 @@ def parallel_nsa_compression_bwd_kernel_dq( b_dp = tl.dot(b_do, b_v) b_ds = b_p * (b_dp.to(tl.float32) - b_delta[:, None]) # [G, BC] @ [BC, BK] -> [G, BK] - b_dq += tl.dot(b_ds.to(b_k.dtype), tl.trans(b_k)) + b_dq = tl.dot(b_ds.to(b_k.dtype), tl.trans(b_k), b_dq) b_dq *= scale - tl.store(p_dq, b_dq.to(p_dq.dtype.element_ty), boundary_check=(0, 1)) + tl.store(p_dq, b_dq.to(p_dq.dtype.element_ty), mask=m_q) + + +def _prune_compression_dkv_bq(configs, nargs, **kwargs): + args = {**(nargs or {}), **kwargs} + G = args.get('G', 1) + BS = args.get('BS', 0) + BK = args.get('BK', 0) + BV = args.get('BV', 0) + T_BUCKET = args.get('T_BUCKET', 0) + LAYOUT = args.get('LAYOUT', -1) + is_packed_heavy = LAYOUT == 1 and G == 32 and BK == 128 and BV == 64 + return [ + config + for config in configs + if config.kwargs['BQ'] == 1 or ( + T_BUCKET > 0 + and LAYOUT in (0, 1) + and BS in (32, 64) + and BK <= 128 + and BV <= 128 + and 2 * G <= 256 + and config.num_stages == (2 if is_packed_heavy else 3) + ) + ] @triton.heuristics({ @@ -228,10 +262,14 @@ def parallel_nsa_compression_bwd_kernel_dq( }) @triton.autotune( configs=[ - triton.Config({}, num_warps=num_warps) + triton.Config({'BQ': 1}, num_warps=num_warps) for num_warps in [1, 2, 4] + ] + [ + triton.Config({'BQ': 2}, num_warps=4), + triton.Config({'BQ': 2}, num_warps=4, num_stages=2), ], - key=['BS', 'BK', 'BV'], + key=['BS', 'BK', 'BV', 'G', 'T_BUCKET', 'LAYOUT'], + prune_configs_by={'early_config_prune': _prune_compression_dkv_bq}, **autotune_cache_kwargs, ) @triton.jit(do_not_specialize=['T', 'TC']) @@ -260,67 +298,83 @@ def parallel_nsa_compression_bwd_kernel_dkv( BS: tl.constexpr, BK: tl.constexpr, BV: tl.constexpr, + T_BUCKET: tl.constexpr, + LAYOUT: tl.constexpr, + BQ: tl.constexpr, IS_VARLEN: tl.constexpr, ): - i_v, i_c, i_bh = tl.program_id(0), tl.program_id(1), tl.program_id(2) + i_v, i_c = unflatten_program_id(tl.cdiv(V, BV)) + i_bh = tl.program_id(1).to(tl.int64) i_b, i_h = i_bh // H, i_bh % H + T, TC = T.to(tl.int64), TC.to(tl.int64) all = B * TC if IS_VARLEN: - i_n, i_c = tl.load(chunk_indices + i_c * 2).to(tl.int32), tl.load(chunk_indices + i_c * 2 + 1).to(tl.int32) + i_n, i_c = tl.load(chunk_indices + i_c * 2).to(tl.int64), tl.load(chunk_indices + i_c * 2 + 1).to(tl.int64) bos, eos = tl.load(cu_seqlens + i_n).to(tl.int64), tl.load(cu_seqlens + i_n + 1).to(tl.int64) - T = (eos - bos).to(tl.int32) + T = eos - bos # the number of compression representations in total TC = tl.cdiv(T, BS) - boc = tl.load(chunk_offsets + i_n).to(tl.int32) + boc = tl.load(chunk_offsets + i_n).to(tl.int64) else: bos, eos = (i_b * T).to(tl.int64), (i_b * T + T).to(tl.int64) boc = i_b * tl.cdiv(T, BS) - p_k = tl.make_block_ptr(k + (boc * H + i_h) * K, (TC, K), (H*K, 1), (i_c * BC, 0), (BC, BK), (1, 0)) - p_v = tl.make_block_ptr(v + (boc * H + i_h) * V, (TC, V), (H*V, 1), (i_c * BC, i_v * BV), (BC, BV), (1, 0)) - p_dk = tl.make_block_ptr(dk + (i_v * all*H + boc * H + i_h) * K, (TC, K), (H*K, 1), (i_c * BC, 0), (BC, BK), (1, 0)) - p_dv = tl.make_block_ptr(dv + (boc * H + i_h) * V, (TC, V), (H*V, 1), (i_c * BC, i_v * BV), (BC, BV), (1, 0)) + o_c = i_c * BC + tl.arange(0, BC) + o_d = tl.arange(0, BK) + o_v = i_v * BV + tl.arange(0, BV) + m_k = (o_c[:, None] < TC) & (o_d[None, :] < K) + m_v = (o_c[:, None] < TC) & (o_v[None, :] < V) + p_k = k + (boc * H + i_h) * K + o_c[:, None] * (H*K) + o_d[None, :] + p_v = v + (boc * H + i_h) * V + o_c[:, None] * (H*V) + o_v[None, :] + p_dk = dk + (i_v * all*H + boc * H + i_h) * K + o_c[:, None] * (H*K) + o_d[None, :] + p_dv = dv + (boc * H + i_h) * V + o_c[:, None] * (H*V) + o_v[None, :] # [BC, BK] - b_k = tl.load(p_k, boundary_check=(0, 1)) + b_k = tl.load(p_k, mask=m_k, other=0.0) b_dk = tl.zeros([BC, BK], dtype=tl.float32) # [BC, BV] - b_v = tl.load(p_v, boundary_check=(0, 1)) + b_v = tl.load(p_v, mask=m_v, other=0.0) b_dv = tl.zeros([BC, BV], dtype=tl.float32) - for i in range(i_c * BC * BS, T): - o_c = i_c * BC + tl.arange(0, BC) - - p_q = tl.make_block_ptr(q + (bos + i) * HQ*K, (HQ, K), (K, 1), (i_h * G, 0), (G, BK), (1, 0)) - # [G, BK] - b_q = tl.load(p_q, boundary_check=(0, 1)) + for i in range(i_c * BC * BS, T, BQ): + o_qg = tl.arange(0, BQ * G) + o_q = i + o_qg // G + o_g = i_h * G + o_qg % G + m_q = (o_q < T) & (o_g < HQ) + p_q = q + (bos + o_q[:, None]) * HQ*K + o_g[:, None] * K + o_d[None, :] + # [BQ*G, BK] + b_q = tl.load(p_q, mask=m_q[:, None] & (o_d[None, :] < K), other=0.0) b_q = (b_q * scale).to(b_q.dtype) - p_do = tl.make_block_ptr(do + (bos + i) * HQ*V, (HQ, V), (V, 1), (i_h * G, i_v * BV), (G, BV), (1, 0)) - p_lse = lse + (bos + i) * HQ + i_h * G + tl.arange(0, G) - p_delta = delta + (bos + i) * HQ + i_h * G + tl.arange(0, G) - # [G, BV] - b_do = tl.load(p_do, boundary_check=(0, 1)) - # [G] - b_lse = tl.load(p_lse) - b_delta = tl.load(p_delta) - # [BC, G] + p_do = do + (bos + o_q[:, None]) * HQ*V + o_g[:, None] * V + o_v[None, :] + p_lse = lse + (bos + o_q) * HQ + o_g + p_delta = delta + (bos + o_q) * HQ + o_g + # [BQ*G, BV] + b_do = tl.load(p_do, mask=m_q[:, None] & (o_v[None, :] < V), other=0.0) + # [BQ*G] + b_lse = tl.load(p_lse, mask=m_q, other=0.0) + b_delta = tl.load(p_delta, mask=m_q, other=0.0) + # [BC, BQ*G] b_s = tl.dot(b_k, tl.trans(b_q)) b_p = exp(b_s - b_lse[None, :]) - b_p = tl.where((i >= max(0, (o_c + 1) * BS - 1))[:, None], b_p, 0) - # [BC, G] @ [G, BV] -> [BC, BV] - b_dv += tl.dot(b_p.to(b_do.dtype), b_do) - # [BC, BV] @ [BV, G] -> [BC, G] + b_p = tl.where( + m_q[None, :] & (o_q[None, :] >= tl.maximum(0, (o_c + 1) * BS - 1)[:, None]), + b_p, + 0, + ) + # [BC, BQ*G] @ [BQ*G, BV] -> [BC, BV] + b_dv = tl.dot(b_p.to(b_do.dtype), b_do, b_dv) + # [BC, BV] @ [BV, BQ*G] -> [BC, BQ*G] b_dp = tl.dot(b_v, tl.trans(b_do)) - # [BC, G] + # [BC, BQ*G] b_ds = b_p * (b_dp - b_delta[None, :]) - # [BC, G] @ [G, BK] -> [BC, BK] - b_dk += tl.dot(b_ds.to(b_q.dtype), b_q) + # [BC, BQ*G] @ [BQ*G, BK] -> [BC, BK] + b_dk = tl.dot(b_ds.to(b_q.dtype), b_q, b_dk) - tl.store(p_dk, b_dk.to(p_dk.dtype.element_ty), boundary_check=(0, 1)) - tl.store(p_dv, b_dv.to(p_dv.dtype.element_ty), boundary_check=(0, 1)) + tl.store(p_dk, b_dk.to(p_dk.dtype.element_ty), mask=m_k) + tl.store(p_dv, b_dv.to(p_dv.dtype.element_ty), mask=m_v) def parallel_nsa_compression_fwd( @@ -401,6 +455,8 @@ def parallel_nsa_compression_bwd( BK = max(triton.next_power_of_2(K), 16) BV = min(128, max(triton.next_power_of_2(v.shape[-1]), 16)) NV = triton.cdiv(V, BV) + T_BUCKET = 0 if T < 8192 else 1 if T < 32768 else 2 + LAYOUT = int(cu_seqlens is not None) if cu_seqlens is not None: chunk_offsets = prepare_chunk_offsets(cu_seqlens, BS) if chunk_indices is None: @@ -443,7 +499,7 @@ def parallel_nsa_compression_bwd( dk = torch.empty(NV, *k.shape, dtype=k.dtype if NV == 1 else torch.float, device=q.device) dv = torch.empty(v.shape, dtype=v.dtype, device=q.device) - grid = (NV, NC, B * H) + grid = (NV * NC, B * H) parallel_nsa_compression_bwd_kernel_dkv[grid]( q=q, k=k, @@ -469,6 +525,8 @@ def parallel_nsa_compression_bwd( BS=BS, BK=BK, BV=BV, + T_BUCKET=T_BUCKET, + LAYOUT=LAYOUT, ) dk = dk.sum(0) return dq, dk, dv @@ -559,6 +617,9 @@ def parallel_nsa_compression( scale: float = None, cu_seqlens: torch.LongTensor | tuple[torch.LongTensor, torch.LongTensor] | None = None ): + HQ, H = q.shape[2], k.shape[2] + if H == 0 or HQ % H != 0: + raise ValueError(f"The number of query heads ({HQ}) must be divisible by the number of key/value heads ({H}).") if scale is None: scale = k.shape[-1] ** -0.5 return ParallelNSACompressionFunction.apply( diff --git a/build/torch-cuda/ops/nsa/naive.py b/build/torch-cuda/ops/nsa/naive.py index 9c155db6809518535c8ca245b9dbb2c7e8b5ae58..fcce8585b11ddfd6932dfbe525eb316ccbaaeaf5 100644 --- a/build/torch-cuda/ops/nsa/naive.py +++ b/build/torch-cuda/ops/nsa/naive.py @@ -31,7 +31,6 @@ def naive_nsa_selection( block_size: int = 64, scale: float | None = None, cu_seqlens: torch.LongTensor | tuple[torch.LongTensor, torch.LongTensor] | None = None, - **kwargs, ) -> torch.Tensor: r""" Args: @@ -60,15 +59,14 @@ def naive_nsa_selection( o (torch.Tensor): Outputs of shape `[B, TQ, HQ, V]`. """ - if 'head_first' in kwargs: - raise DeprecationWarning( - "head_first has been removed. Inputs must be in `[B, T, H, ...]` format.", - ) if scale is None: scale = k.shape[-1] ** -0.5 dtype = q.dtype - G = q.shape[2] // k.shape[2] + HQ, H = q.shape[2], k.shape[2] + if H == 0 or HQ % H != 0: + raise ValueError(f"The number of query heads ({HQ}) must be divisible by the number of key/value heads ({H}).") + G = HQ // H BS = block_size k, v, block_indices = (repeat(x, 'b t h d -> b t (h g) d', g=G) for x in (k, v, block_indices)) q, k, v = map(lambda x: x.float(), (q, k, v)) @@ -157,8 +155,10 @@ def naive_nsa_compression( Log-sum-exp of attention scores of shape `[B, TQ, HQ]`, `-inf` where no block is visible yet. """ dtype = q.dtype - H = k_cmp.shape[2] - G = q.shape[2] // H + HQ, H = q.shape[2], k_cmp.shape[2] + if H == 0 or HQ % H != 0: + raise ValueError(f"The number of query heads ({HQ}) must be divisible by the number of key/value heads ({H}).") + G = HQ // H q, k_cmp, v_cmp = (x.float() for x in (q, k_cmp, v_cmp)) k_cmp, v_cmp = (repeat(x, 'b t h d -> b t (h g) d', g=G) for x in (k_cmp, v_cmp)) @@ -225,6 +225,8 @@ def naive_nsa_topk( """ B, TQ, HQ, _ = q.shape H = k_cmp.shape[2] + if H == 0 or HQ % H != 0: + raise ValueError(f"The number of query heads ({HQ}) must be divisible by the number of key/value heads ({H}).") G = HQ // H k_cmp = repeat(k_cmp, 'b t h d -> b t (h g) d', g=G) @@ -364,7 +366,10 @@ def naive_nsa( scale = k.shape[-1] ** -0.5 if cu_seqlens is not None: assert q.shape[0] == 1, "batch size must be 1 when cu_seqlens are provided" - G = q.shape[2] // k.shape[2] + HQ, H = q.shape[2], k.shape[2] + if H == 0 or HQ % H != 0: + raise ValueError(f"The number of query heads ({HQ}) must be divisible by the number of key/value heads ({H}).") + G = HQ // H assert G >= 16 and (G & (G - 1)) == 0, "Group size (HQ/H) must be a power of 2 and >= 16 in NSA" if cu_seqlens is not None: diff --git a/build/torch-cuda/ops/nsa/parallel.py b/build/torch-cuda/ops/nsa/parallel.py index c58ebf1238d261f7c034ac435a7713b00bf74ffd..cb590a18ab92507e0e7f0d36fc931c35f43f991a 100644 --- a/build/torch-cuda/ops/nsa/parallel.py +++ b/build/torch-cuda/ops/nsa/parallel.py @@ -15,7 +15,7 @@ from ...ops.attn.parallel import parallel_attn_bwd_preprocess from ...ops.nsa.compression import parallel_nsa_compression from ...ops.nsa.utils import _bitonic_merge from ...ops.utils import prepare_block_csr, prepare_chunk_indices, prepare_chunk_offsets, prepare_lens, prepare_token_indices -from ...ops.utils.op import exp, log +from ...ops.utils.op import exp, log, unflatten_program_id from ...ops.utils.pooling import mean_pooling from ...utils import autocast_custom_bwd, autocast_custom_fwd, autotune_cache_kwargs, check_shared_mem, contiguous @@ -63,28 +63,31 @@ def parallel_nsa_kernel_topk( BK: tl.constexpr, IS_VARLEN: tl.constexpr, ): - i_t, i_bh = tl.program_id(0), tl.program_id(1) + i_t, i_bh = tl.program_id(0).to(tl.int64), tl.program_id(1).to(tl.int64) i_b, i_h = i_bh // H, i_bh % H if IS_VARLEN: - i_n, i_t = tl.load(token_indices_q + i_t * 2).to(tl.int32), tl.load(token_indices_q + i_t * 2 + 1).to(tl.int32) + i_n, i_t = tl.load(token_indices_q + i_t * 2).to(tl.int32), tl.load(token_indices_q + i_t * 2 + 1).to(tl.int64) bos_q, eos_q = tl.load(cu_seqlens_q + i_n).to(tl.int64), tl.load(cu_seqlens_q + i_n + 1).to(tl.int64) bos_k, eos_k = tl.load(cu_seqlens_k + i_n).to(tl.int64), tl.load(cu_seqlens_k + i_n + 1).to(tl.int64) TQ = (eos_q - bos_q).to(tl.int32) TK = (eos_k - bos_k).to(tl.int32) TC = tl.cdiv(TK, BS) - boc = tl.load(chunk_offsets + i_n).to(tl.int32) + boc = tl.load(chunk_offsets + i_n).to(tl.int64) else: bos_q, eos_q = (i_b * TQ).to(tl.int64), (i_b * TQ + TQ).to(tl.int64) TC = tl.cdiv(TK, BS) boc = i_b * TC # boc is the start of the current sequence at [B, TC] dimensions - p_q = tl.make_block_ptr(q + (bos_q + i_t) * HQ * K, (HQ, K), (K, 1), (i_h * G, 0), (G, BK), (1, 0)) + o_g = i_h * G + tl.arange(0, G) + o_d = tl.arange(0, BK) + m_q = (o_g[:, None] < HQ) & (o_d[None, :] < K) + p_q = q + (bos_q + i_t) * HQ * K + o_g[:, None] * K + o_d[None, :] Q_OFFSET = TK - TQ # the Q block is kept in the shared memory throughout the whole kernel # [G, BK] - b_q = tl.load(p_q, boundary_check=(0, 1)) + b_q = tl.load(p_q, mask=m_q, other=0.0) b_q = (b_q * scale).to(b_q.dtype) # number of complete compression blocks visible to the query (q tokens are the last TQ of the sequence) @@ -102,9 +105,9 @@ def parallel_nsa_kernel_topk( for i_c in range(0, NC, BC): o_c = i_c + tl.arange(0, BC) - p_k = tl.make_block_ptr(k + (boc * H + i_h) * K, (K, TC), (1, H*K), (0, i_c), (BK, BC), (0, 1)) + p_k = k + (boc * H + i_h) * K + o_d[:, None] + o_c[None, :] * (H*K) # [BK, BC] - b_k = tl.load(p_k, boundary_check=(0, 1)) + b_k = tl.load(p_k, mask=(o_d[:, None] < K) & (o_c[None, :] < TC), other=0.0) # [G, BC] b_s = tl.dot(b_q, b_k) @@ -135,9 +138,9 @@ def parallel_nsa_kernel_topk( IC = (i_t + Q_OFFSET) // BS # Idx of the current query block for i_c in range(0, IC + 1, BC): # +1, because the current block might be also included o_c = i_c + tl.arange(0, BC) - p_k = tl.make_block_ptr(k + (boc * H + i_h) * K, (K, TC), (1, H*K), (0, i_c), (BK, BC), (0, 1)) + p_k = k + (boc * H + i_h) * K + o_d[:, None] + o_c[None, :] * (H*K) # [BK, BC] - b_k = tl.load(p_k, boundary_check=(0, 1)) + b_k = tl.load(p_k, mask=(o_d[:, None] < K) & (o_c[None, :] < TC), other=0.0) # [G, BC] b_s = tl.dot(b_q, b_k) b_s = tl.where(o_c < IC, b_s, float('-inf')) @@ -163,7 +166,7 @@ def parallel_nsa_kernel_topk( m_top = tl.arange(0, BC // S) == 0 b_top = tl.sum(m_top[:, None] * tl.reshape(o_i, [BC // S, S]), 0) - p_b = tl.make_block_ptr(block_indices + (bos_q + i_t) * H*S, (H*S,), (1,), (i_h * S,), (S,), (0,)) + p_b = block_indices + (bos_q + i_t) * H*S + i_h * S + tl.arange(0, S) tl.store(p_b, b_top.to(p_b.dtype.element_ty)) @@ -207,13 +210,13 @@ def parallel_nsa_fwd_kernel( IS_VARLEN: tl.constexpr, USE_BLOCK_COUNTS: tl.constexpr, ): - i_t, i_v, i_bh = tl.program_id(0), tl.program_id(1), tl.program_id(2) + i_t, i_v, i_bh = tl.program_id(0).to(tl.int64), tl.program_id(1).to(tl.int64), tl.program_id(2).to(tl.int64) i_b, i_h = i_bh // H, i_bh % H # k/v: [B, TK, H, *], q: [B, TQ, HQ, K], block_indices: [B, TQ, H, S], lse: [B, TQ, HQ]; G = HQ // H if IS_VARLEN: # token_indices_q maps a flattened query to its (sequence index, in-sequence position) - i_n, i_t = tl.load(token_indices_q + i_t * 2).to(tl.int32), tl.load(token_indices_q + i_t * 2 + 1).to(tl.int32) + i_n, i_t = tl.load(token_indices_q + i_t * 2).to(tl.int32), tl.load(token_indices_q + i_t * 2 + 1).to(tl.int64) bos_q, eos_q = tl.load(cu_seqlens_q + i_n).to(tl.int64), tl.load(cu_seqlens_q + i_n + 1).to(tl.int64) bos_k, eos_k = tl.load(cu_seqlens_k + i_n).to(tl.int64), tl.load(cu_seqlens_k + i_n + 1).to(tl.int64) TQ = (eos_q - bos_q).to(tl.int32) @@ -235,13 +238,18 @@ def parallel_nsa_fwd_kernel( else: NS = S - p_q = tl.make_block_ptr(q + (bos_q + i_t) * HQ * K, (HQ, K), (K, 1), (i_h * G, 0), (G, BK), (1, 0)) + o_g = i_h * G + tl.arange(0, G) + o_d = tl.arange(0, BK) + m_q = (o_g[:, None] < HQ) & (o_d[None, :] < K) + p_q = q + (bos_q + i_t) * HQ * K + o_g[:, None] * K + o_d[None, :] # the Q block is kept in shared memory throughout the kernel # [G, BK] - b_q = tl.load(p_q, boundary_check=(0, 1)) + b_q = tl.load(p_q, mask=m_q, other=0.0) b_q = (b_q * scale).to(b_q.dtype) - p_o = tl.make_block_ptr(o + (bos_q + i_t) * HQ * V, (HQ, V), (V, 1), (i_h * G, i_v * BV), (G, BV), (1, 0)) + o_v = i_v * BV + tl.arange(0, BV) + m_o = (o_g[:, None] < HQ) & (o_v[None, :] < V) + p_o = o + (bos_q + i_t) * HQ * V + o_g[:, None] * V + o_v[None, :] p_lse = lse + (bos_q + i_t) * HQ + i_h * G + tl.arange(0, G) # [G, BV] @@ -250,14 +258,16 @@ def parallel_nsa_fwd_kernel( b_m = tl.full([G], float('-inf'), dtype=tl.float32) b_acc = tl.zeros([G], dtype=tl.float32) for i in range(NS): - i_s = tl.load(block_indices + i).to(tl.int32) * BS # start token index of the i-th selected KV block + i_s = tl.load(block_indices + i).to(tl.int64) * BS # start token index of the i-th selected KV block if i_s <= Q_OFFSET + i_t and i_s >= 0: - p_k = tl.make_block_ptr(k, (K, TK), (1, H*K), (0, i_s), (BK, BS), (0, 1)) - p_v = tl.make_block_ptr(v, (TK, V), (H*V, 1), (i_s, i_v * BV), (BS, BV), (1, 0)) + o_s = i_s + tl.arange(0, BS) + m_s = o_s < TK + p_k = k + o_d[:, None] + o_s[None, :] * (H*K) + p_v = v + o_s[:, None] * (H*V) + o_v[None, :] # [BK, BS] - b_k = tl.load(p_k, boundary_check=(0, 1)) + b_k = tl.load(p_k, mask=(o_d[:, None] < K) & m_s[None, :], other=0.0) # [BS, BV] - b_v = tl.load(p_v, boundary_check=(0, 1)) + b_v = tl.load(p_v, mask=m_s[:, None] & (o_v[None, :] < V), other=0.0) # [G, BS] b_s = tl.dot(b_q, b_k) # causal mask against the absolute query position Q_OFFSET + i_t @@ -275,8 +285,9 @@ def parallel_nsa_fwd_kernel( b_o = b_o / b_acc[:, None] b_m += log(b_acc) - tl.store(p_o, b_o.to(p_o.dtype.element_ty), boundary_check=(0, 1)) - tl.store(p_lse, b_m.to(p_lse.dtype.element_ty)) + tl.store(p_o, b_o.to(p_o.dtype.element_ty), mask=m_o) + if i_v == 0: + tl.store(p_lse, b_m.to(p_lse.dtype.element_ty)) @triton.heuristics({ @@ -320,12 +331,12 @@ def parallel_nsa_bwd_kernel_dq( IS_VARLEN: tl.constexpr, USE_BLOCK_COUNTS: tl.constexpr, ): - i_t, i_v, i_bh = tl.program_id(0), tl.program_id(1), tl.program_id(2) + i_t, i_v, i_bh = tl.program_id(0).to(tl.int64), tl.program_id(1).to(tl.int64), tl.program_id(2).to(tl.int64) i_b, i_h = i_bh // H, i_bh % H - all = B * T + all = B * T.to(tl.int64) if IS_VARLEN: - i_n, i_t = tl.load(token_indices + i_t * 2).to(tl.int32), tl.load(token_indices + i_t * 2 + 1).to(tl.int32) + i_n, i_t = tl.load(token_indices + i_t * 2).to(tl.int32), tl.load(token_indices + i_t * 2 + 1).to(tl.int64) bos, eos = tl.load(cu_seqlens + i_n).to(tl.int64), tl.load(cu_seqlens + i_n + 1).to(tl.int64) T = (eos - bos).to(tl.int32) else: @@ -346,19 +357,24 @@ def parallel_nsa_bwd_kernel_dq( k += (bos * H + i_h) * K v += (bos * H + i_h) * V - p_q = tl.make_block_ptr(q, (HQ, K), (K, 1), (i_h * G, 0), (G, BK), (1, 0)) - p_dq = tl.make_block_ptr(dq, (HQ, K), (K, 1), (i_h * G, 0), (G, BK), (1, 0)) + o_g = i_h * G + tl.arange(0, G) + o_d = tl.arange(0, BK) + m_q = (o_g[:, None] < HQ) & (o_d[None, :] < K) + p_q = q + o_g[:, None] * K + o_d[None, :] + p_dq = dq + o_g[:, None] * K + o_d[None, :] # [G, BK] - b_q = tl.load(p_q, boundary_check=(0, 1)) + b_q = tl.load(p_q, mask=m_q, other=0.0) b_q = (b_q * scale).to(b_q.dtype) - p_do = tl.make_block_ptr(do, (HQ, V), (V, 1), (i_h * G, i_v * BV), (G, BV), (1, 0)) + o_v = i_v * BV + tl.arange(0, BV) + m_o = (o_g[:, None] < HQ) & (o_v[None, :] < V) + p_do = do + o_g[:, None] * V + o_v[None, :] p_lse = lse + i_h * G + tl.arange(0, G) p_delta = delta + i_h * G + tl.arange(0, G) # [G, BV] - b_do = tl.load(p_do, boundary_check=(0, 1)) + b_do = tl.load(p_do, mask=m_o, other=0.0) # [G] b_lse = tl.load(p_lse) b_delta = tl.load(p_delta) @@ -366,14 +382,16 @@ def parallel_nsa_bwd_kernel_dq( # [G, BK] b_dq = tl.zeros([G, BK], dtype=tl.float32) for i in range(NS): - i_s = tl.load(block_indices + i).to(tl.int32) * BS + i_s = tl.load(block_indices + i).to(tl.int64) * BS if i_s <= i_t and i_s >= 0: - p_k = tl.make_block_ptr(k, (K, T), (1, H*K), (0, i_s), (BK, BS), (0, 1)) - p_v = tl.make_block_ptr(v, (V, T), (1, H*V), (i_v * BV, i_s), (BV, BS), (0, 1)) + o_s = i_s + tl.arange(0, BS) + m_s = o_s < T + p_k = k + o_d[:, None] + o_s[None, :] * (H*K) + p_v = v + o_v[:, None] + o_s[None, :] * (H*V) # [BK, BS] - b_k = tl.load(p_k, boundary_check=(0, 1)) + b_k = tl.load(p_k, mask=(o_d[:, None] < K) & m_s[None, :], other=0.0) # [BV, BS] - b_v = tl.load(p_v, boundary_check=(0, 1)) + b_v = tl.load(p_v, mask=(o_v[:, None] < V) & m_s[None, :], other=0.0) # [G, BS] b_s = tl.dot(b_q, b_k) @@ -384,10 +402,10 @@ def parallel_nsa_bwd_kernel_dq( b_dp = tl.dot(b_do, b_v) b_ds = b_p * (b_dp.to(tl.float32) - b_delta[:, None]) # [G, BS] @ [BS, BK] -> [G, BK] - b_dq += tl.dot(b_ds.to(b_k.dtype), tl.trans(b_k)) + b_dq = tl.dot(b_ds.to(b_k.dtype), tl.trans(b_k), b_dq) b_dq *= scale - tl.store(p_dq, b_dq.to(p_dq.dtype.element_ty), boundary_check=(0, 1)) + tl.store(p_dq, b_dq.to(p_dq.dtype.element_ty), mask=m_q) def _prune_dkv_bq(configs, nargs, **kwargs): @@ -441,8 +459,8 @@ def parallel_nsa_bwd_kernel_dkv( BQ: tl.constexpr, IS_VARLEN: tl.constexpr, ): - i_v, i_blk = tl.program_id(0), tl.program_id(1) - all = B * T + i_v, i_blk = unflatten_program_id(tl.cdiv(V, BV)) + all = B * T.to(tl.int64) if IS_VARLEN: i_c, i_h = i_blk // H, i_blk % H i_n = tl.load(chunk_indices + i_c * 2).to(tl.int32) @@ -502,12 +520,12 @@ def parallel_nsa_bwd_kernel_dkv( b_p = tl.where((o_q[None, :] >= o_t[:, None]) & m_q[None, :], b_p, 0.) # [BS, BQ*G] @ [BQ*G, BV] -> [BS, BV] - b_dv += tl.dot(b_p.to(b_do.dtype), b_do) + b_dv = tl.dot(b_p.to(b_do.dtype), b_do, b_dv) # [BS, BV] @ [BV, BQ*G] -> [BS, BQ*G] b_dp = tl.dot(b_v, tl.trans(b_do)) b_ds = b_p * (b_dp.to(tl.float32) - b_delta[None, :]) # [BS, BQ*G] @ [BQ*G, BK] -> [BS, BK] - b_dk += tl.dot(b_ds.to(b_q.dtype), b_q) + b_dk = tl.dot(b_ds.to(b_q.dtype), b_q, b_dk) o_dk = (i_v * all + o_t)[:, None] * H*K + o_d[None, :] o_dv = o_t[:, None] * H*V + o_v[None, :] @@ -530,6 +548,9 @@ def parallel_nsa_topk( assert k.shape[0] == q.shape[0] and k.shape[-1] == q.shape[-1], "The last dimension of k and q must match" assert lse is None or lse.shape == (B, TQ, HQ), "The shape of lse must be (B, TQ, HQ)" + if H == 0 or HQ % H != 0: + raise ValueError(f"The number of query heads ({HQ}) must be divisible by the number of key/value heads ({H}).") + G = HQ // H if cu_seqlens is not None: if isinstance(cu_seqlens, tuple): @@ -540,7 +561,6 @@ def parallel_nsa_topk( else: cu_seqlens_q = cu_seqlens_k = token_indices_q = None - G = HQ // H # the number of selected blocks for each token S = block_counts if isinstance(block_counts, int) else block_counts.max().item() S = triton.next_power_of_2(S) @@ -655,7 +675,10 @@ def parallel_nsa_bwd( G = HQ // H BS = block_size BK = max(triton.next_power_of_2(K), 16) - BV = min(128, max(triton.next_power_of_2(v.shape[-1]), 16)) + # dq/dk accumulate softmax deltas reduced over the full value dim (delta from + # parallel_attn_bwd_preprocess spans all of V), so BV must span all of V (NV == 1) + # to avoid cross-program over-subtraction — same contract as fla/ops/attn/parallel bwd. + BV = max(triton.next_power_of_2(v.shape[-1]), 16) NV = triton.cdiv(V, BV) delta = parallel_attn_bwd_preprocess(o, do) @@ -707,7 +730,7 @@ def parallel_nsa_bwd( block_size=block_size, ) NB = chunk_indices.shape[0] * H if cu_seqlens is not None else B * H * M - grid = (NV, NB) + grid = (NV * NB,) parallel_nsa_bwd_kernel_dkv[grid]( q=q, k=k, @@ -738,7 +761,6 @@ def parallel_nsa_bwd( return dq, dk, dv -@torch.compile class ParallelNSAFunction(torch.autograd.Function): @staticmethod @@ -841,7 +863,7 @@ def parallel_nsa( g_slc (torch.Tensor): Gate score for selected attention of shape `[B, TQ, HQ]`. g_swa (torch.Tensor): - Gate score for sliding attentionof shape `[B, TQ, HQ]`. + Gate score for sliding attention of shape `[B, TQ, HQ]`. block_indices (torch.LongTensor): Block indices of shape `[B, TQ, H, S]`. `S` is the number of selected blocks for each query token, which is set to 16 in the paper. @@ -875,7 +897,10 @@ def parallel_nsa( f"The batch size is expected to be 1 rather than {q.shape[0]} when using `cu_seqlens`. " f"Please flatten variable-length inputs before processing.", ) - G = q.shape[2] // k.shape[2] + HQ, H = q.shape[2], k.shape[2] + if H == 0 or HQ % H != 0: + raise ValueError(f"The number of query heads ({HQ}) must be divisible by the number of key/value heads ({H}).") + G = HQ // H assert G >= 16 and (G & (G - 1)) == 0, "Group size (HQ/H) must be a power of 2 and >= 16 in NSA" if cu_seqlens is not None: diff --git a/build/torch-cuda/ops/parallax/decode.py b/build/torch-cuda/ops/parallax/decode.py index 4be585faa1e95dbd754ea408343a4f62486e50a9..8322aff8c917de2ddcc97f94165d481bd2d59a72 100644 --- a/build/torch-cuda/ops/parallax/decode.py +++ b/build/torch-cuda/ops/parallax/decode.py @@ -42,7 +42,9 @@ def parallax_decode_kernel( sliding window and to ``[cache_start, Skv)`` when set. One program owns a ``BT``-row query block; see ``naive_parallax`` for the output formula. """ - i_t, i_bh = tl.program_id(0), tl.program_id(1) + pid = tl.program_id(0).to(tl.int64) + NT = tl.cdiv(Sq, BT) + i_t, i_bh = (pid % NT).to(tl.int64), (pid // NT).to(tl.int64) i_b, i_hq = i_bh // HQ, i_bh % HQ i_h = i_hq // G # kv head shared by this q head (GQA) RCP_LN2: tl.constexpr = 1.4426950216 @@ -69,14 +71,15 @@ def parallax_decode_kernel( leftmost = kv_lo KV_START_BLOCK = leftmost // BS - p_q = tl.make_block_ptr(q + (i_b * Sq * HQ + i_hq) * K, (Sq, K), (HQ * K, 1), (q_off, 0), (BT, BK), (1, 0)) - p_r = tl.make_block_ptr(r + (i_b * Sq * HQ + i_hq) * K, (Sq, K), (HQ * K, 1), (q_off, 0), (BT, BK), (1, 0)) - p_k = tl.make_block_ptr(k + (i_b * Skv * H + i_h) * K, (Skv, K), (H * K, 1), (KV_START_BLOCK * BS, 0), (BS, BK), (1, 0)) - p_v = tl.make_block_ptr(v + (i_b * Skv * H + i_h) * K, (Skv, K), (H * K, 1), (KV_START_BLOCK * BS, 0), (BS, BK), (1, 0)) - p_o = tl.make_block_ptr(o + (i_b * Sq * HQ + i_hq) * K, (Sq, K), (HQ * K, 1), (q_off, 0), (BT, BK), (1, 0)) + o_k = tl.arange(0, BK) + m_k = o_k < K + m_qk = row_mask & m_k[None, :] + p_q = q + (i_b * Sq * HQ + i_hq) * K + rows[:, None] * (HQ * K) + o_k[None, :] + p_r = r + (i_b * Sq * HQ + i_hq) * K + rows[:, None] * (HQ * K) + o_k[None, :] + p_o = o + (i_b * Sq * HQ + i_hq) * K + rows[:, None] * (HQ * K) + o_k[None, :] - b_q = tl.load(p_q, boundary_check=(0, 1), padding_option="zero") - b_r = tl.load(p_r, boundary_check=(0, 1), padding_option="zero") + b_q = tl.load(p_q, mask=m_qk, other=0.0) + b_r = tl.load(p_r, mask=m_qk, other=0.0) m_acc = tl.zeros((BT, 1), dtype=tl.float32) - float("inf") d1_acc = tl.zeros((BT, 1), dtype=tl.float32) d2_acc = tl.zeros((BT, 1), dtype=tl.float32) @@ -86,8 +89,12 @@ def parallax_decode_kernel( for col_block_id in range(KV_START_BLOCK, KV_END_BLOCK): col = (col_block_id * BS + tl.arange(0, BS))[None, :] # [1, BS] key positions - b_k = tl.load(p_k, boundary_check=(0, 1), padding_option="zero") # [BS, BK] - b_v = tl.load(p_v, boundary_check=(0, 1), padding_option="zero") # [BS, BK] + o_kv = col_block_id.to(tl.int64) * BS + tl.arange(0, BS) + m_kv = (o_kv[:, None] < Skv) & m_k[None, :] + p_k = k + (i_b * Skv * H + i_h) * K + o_kv[:, None] * (H * K) + o_k[None, :] + p_v = v + (i_b * Skv * H + i_h) * K + o_kv[:, None] * (H * K) + o_k[None, :] + b_k = tl.load(p_k, mask=m_kv, other=0.0) # [BS, BK] + b_v = tl.load(p_v, mask=m_kv, other=0.0) # [BS, BK] # [BT, BS]: causal, in-cache, past left-padding; optionally inside the window. mask = (abs_q >= col) & row_mask & (col < Skv) & (col >= kv_lo) if WINDOW_SIZE_LEFT >= 0: @@ -108,8 +115,6 @@ def parallax_decode_kernel( barv_acc = tl.dot(w.to(b_v.dtype), b_v, out_dtype=tl.float32, acc=barv_acc) Rv_acc = tl.dot(wr.to(b_v.dtype), b_v, out_dtype=tl.float32, acc=Rv_acc) m_acc = m_new - p_k = tl.advance(p_k, (BS, 0)) - p_v = tl.advance(p_v, (BS, 0)) # Rows that see no valid key (e.g. left-padded query positions) have d1 == 0; # emit a finite zero instead of inf/NaN so padding can't poison valid rows. @@ -118,7 +123,7 @@ def parallax_decode_kernel( b_bart = d2_acc * inv_d1 # d2 / d1 b_o = b_barv + b_bart * b_barv - Rv_acc * inv_d1 # O1/d1 * (1 + d2/d1) - O2/d1 - tl.store(p_o, b_o.to(p_o.dtype.element_ty), boundary_check=(0, 1)) + tl.store(p_o, b_o.to(p_o.dtype.element_ty), mask=m_qk) def parallax_decode( @@ -162,6 +167,9 @@ def parallax_decode( """ B, Sq, HQ, K = q.shape Skv, H = k.shape[1], k.shape[2] + assert Skv >= Sq, f"Cached KV length must cover query length, got Skv={Skv} and Sq={Sq}" + if H == 0 or HQ % H != 0: + raise ValueError(f"The number of query heads ({HQ}) must be divisible by the number of key/value heads ({H}).") G = HQ // H if scale is None: scale = K ** -0.5 @@ -173,7 +181,7 @@ def parallax_decode( BK = triton.next_power_of_2(K) BT = _block_size(K, q.device.index) o = torch.empty_like(q) - grid = (triton.cdiv(Sq, BT), B * HQ) + grid = (triton.cdiv(Sq, BT) * B * HQ,) parallax_decode_kernel[grid]( q, r, k, v, o, float(scale), cache_start, Sq, Skv, HQ=HQ, H=H, G=G, K=K, BK=BK, @@ -216,7 +224,7 @@ def parallax_decode_one_step_kernel( ``parallax_decode_kernel`` incurs at ``Sq == 1``. One program per (batch, head); see ``naive_parallax`` for the output formula. """ - i_bh = tl.program_id(0) + i_bh = tl.program_id(0).to(tl.int64) i_b, i_hq = i_bh // HQ, i_bh % HQ i_h = i_hq // G # kv head shared by this q head (GQA) RCP_LN2: tl.constexpr = 1.4426950216 @@ -229,11 +237,13 @@ def parallax_decode_one_step_kernel( kv_lo = tl.maximum(kv_lo, Skv - WINDOW_SIZE_LEFT) kv_lo = tl.maximum(kv_lo, 0) - p_q = tl.make_block_ptr(q + i_bh * K, (K,), (1,), (0,), (BK,), (0,)) - p_r = tl.make_block_ptr(r + i_bh * K, (K,), (1,), (0,), (BK,), (0,)) - p_o = tl.make_block_ptr(o + i_bh * K, (K,), (1,), (0,), (BK,), (0,)) - b_q = tl.load(p_q, boundary_check=(0,), padding_option="zero").to(tl.float32) # [BK] query vector - b_r = tl.load(p_r, boundary_check=(0,), padding_option="zero").to(tl.float32) # [BK] secondary query + o_k = tl.arange(0, BK) + m_k = o_k < K + p_q = q + i_bh * K + o_k + p_r = r + i_bh * K + o_k + p_o = o + i_bh * K + o_k + b_q = tl.load(p_q, mask=m_k, other=0.0).to(tl.float32) # [BK] query vector + b_r = tl.load(p_r, mask=m_k, other=0.0).to(tl.float32) # [BK] secondary query scale_log2 = scale * RCP_LN2 # Running online-softmax state for the single query: pivot m, denominators @@ -245,13 +255,15 @@ def parallax_decode_one_step_kernel( o2 = tl.zeros((BK,), dtype=tl.float32) start_block = kv_lo // BS - p_k = tl.make_block_ptr(k + (i_b * Skv * H + i_h) * K, (Skv, K), (H * K, 1), (start_block * BS, 0), (BS, BK), (1, 0)) - p_v = tl.make_block_ptr(v + (i_b * Skv * H + i_h) * K, (Skv, K), (H * K, 1), (start_block * BS, 0), (BS, BK), (1, 0)) for i_s in range(start_block * BS, tl.cdiv(Skv, BS) * BS, BS): col = i_s + tl.arange(0, BS) + o_kv = i_s.to(tl.int64) + tl.arange(0, BS) + m_kv = (o_kv[:, None] < Skv) & m_k[None, :] + p_k = k + (i_b * Skv * H + i_h) * K + o_kv[:, None] * (H * K) + o_k[None, :] + p_v = v + (i_b * Skv * H + i_h) * K + o_kv[:, None] * (H * K) + o_k[None, :] mask = (col >= kv_lo) & (col < Skv) # [BS] valid keys - b_k = tl.load(p_k, boundary_check=(0, 1), padding_option="zero") # [BS, BK] - b_v = tl.load(p_v, boundary_check=(0, 1), padding_option="zero") # [BS, BK] + b_k = tl.load(p_k, mask=m_kv, other=0.0) # [BS, BK] + b_v = tl.load(p_v, mask=m_kv, other=0.0) # [BS, BK] s1 = tl.sum(b_q[None, :] * b_k, axis=1) * scale_log2 # [BS] = scale * (q . k), base-2 s2 = tl.sum(b_r[None, :] * b_k, axis=1) # [BS] = r . k (unscaled) s1 = tl.where(mask, s1, -float("inf")) @@ -265,12 +277,10 @@ def parallax_decode_one_step_kernel( o1 = o1 * alpha + tl.sum(p1[:, None] * b_v, axis=0) # [BK] o2 = o2 * alpha + tl.sum(p2[:, None] * b_v, axis=0) m = m_new - p_k = tl.advance(p_k, (BS, 0)) - p_v = tl.advance(p_v, (BS, 0)) inv_d1 = tl.where(d1 > 0.0, 1.0 / d1, 0.0) # 0 when no valid key (avoid NaN) out = o1 * inv_d1 * (1.0 + d2 * inv_d1) - o2 * inv_d1 # [BK] O1/d1*(1 + d2/d1) - O2/d1 - tl.store(p_o, out.to(p_o.dtype.element_ty), boundary_check=(0,)) + tl.store(p_o, out.to(p_o.dtype.element_ty), mask=m_k) def parallax_decode_one_step( @@ -315,6 +325,8 @@ def parallax_decode_one_step( if Sq != 1: raise ValueError(f"parallax_decode_one_step expects a single query (Sq=1), got Sq={Sq}") Skv, H = k.shape[1], k.shape[2] + if H == 0 or HQ % H != 0: + raise ValueError(f"The number of query heads ({HQ}) must be divisible by the number of key/value heads ({H}).") G = HQ // H if scale is None: scale = K ** -0.5 diff --git a/build/torch-cuda/ops/parallax/naive.py b/build/torch-cuda/ops/parallax/naive.py index 174d859ae2f758712572c037ee347c07fc860956..5020c63773640647705fe781d2bfe04c39354969 100644 --- a/build/torch-cuda/ops/parallax/naive.py +++ b/build/torch-cuda/ops/parallax/naive.py @@ -56,6 +56,8 @@ def naive_parallax( """ B, T, HQ, D = q.shape H = k.shape[2] + if H == 0 or HQ % H != 0: + raise ValueError(f"The number of query heads ({HQ}) must be divisible by the number of key/value heads ({H}).") G = HQ // H if scale is None: diff --git a/build/torch-cuda/ops/parallax/parallel.py b/build/torch-cuda/ops/parallax/parallel.py index 40c3c425b5cb0028596126543a6408ab52d485dc..e3ae7b78159a8db6d8f33e0bbc71ba696e7a8dba 100644 --- a/build/torch-cuda/ops/parallax/parallel.py +++ b/build/torch-cuda/ops/parallax/parallel.py @@ -52,12 +52,12 @@ def parallel_parallax_fwd_kernel( BS: tl.constexpr, IS_VARLEN: tl.constexpr, ): - i_t, i_bh = tl.program_id(0), tl.program_id(1) + i_t, i_bh = tl.program_id(0).to(tl.int64), tl.program_id(1).to(tl.int64) i_b, i_hq = i_bh // HQ, i_bh % HQ i_h = i_hq // G 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) + i_n, i_t = tl.load(chunk_indices + i_t * 2).to(tl.int32), tl.load(chunk_indices + i_t * 2 + 1).to(tl.int64) bos, eos = tl.load(cu_seqlens + i_n).to(tl.int64), tl.load(cu_seqlens + i_n + 1).to(tl.int64) T = eos - bos else: @@ -85,18 +85,19 @@ def parallel_parallax_fwd_kernel( SAFE_MIDDLE_START = tl.maximum(FIRST_COL_BLOCK, SAFE_LEFT_START) RIGHT_BORDER_START = tl.maximum(FIRST_COL_BLOCK, NUM_SAFE_BLOCKS) - p_q = tl.make_block_ptr(q + (bos * HQ + i_hq) * K, (T, K), (HQ * K, 1), (row_offset, 0), (BT, BK), (1, 0)) - p_r = tl.make_block_ptr(r + (bos * HQ + i_hq) * K, (T, K), (HQ * K, 1), (row_offset, 0), (BT, BK), (1, 0)) - p_k = tl.make_block_ptr(k + (bos * H + i_h) * K, (T, K), (H * K, 1), (FIRST_COL_BLOCK * BS, 0), (BS, BK), (1, 0)) - p_v = tl.make_block_ptr(v + (bos * H + i_h) * K, (T, K), (H * K, 1), (FIRST_COL_BLOCK * BS, 0), (BS, BK), (1, 0)) - p_o = tl.make_block_ptr(o + (bos * HQ + i_hq) * K, (T, K), (HQ * K, 1), (row_offset, 0), (BT, BK), (1, 0)) - p_barv = tl.make_block_ptr(barv + (bos * HQ + i_hq) * K, (T, K), (HQ * K, 1), (row_offset, 0), (BT, BK), (1, 0)) - p_d1 = tl.make_block_ptr(d1 + bos * HQ + i_hq, (T, 1), (HQ, 1), (row_offset, 0), (BT, 1), (1, 0)) - p_bart = tl.make_block_ptr(bart + bos * HQ + i_hq, (T, 1), (HQ, 1), (row_offset, 0), (BT, 1), (1, 0)) - p_m = tl.make_block_ptr(m + bos * HQ + i_hq, (T, 1), (HQ, 1), (row_offset, 0), (BT, 1), (1, 0)) - - b_q = tl.load(p_q, boundary_check=(0, 1), padding_option="zero") - b_r = tl.load(p_r, boundary_check=(0, 1), padding_option="zero") + o_k = tl.arange(0, BK) + m_k = o_k < K + m_qk = row_mask & m_k[None, :] + p_q = q + (bos * HQ + i_hq) * K + row_indices[:, None] * (HQ * K) + o_k[None, :] + p_r = r + (bos * HQ + i_hq) * K + row_indices[:, None] * (HQ * K) + o_k[None, :] + p_o = o + (bos * HQ + i_hq) * K + row_indices[:, None] * (HQ * K) + o_k[None, :] + p_barv = barv + (bos * HQ + i_hq) * K + row_indices[:, None] * (HQ * K) + o_k[None, :] + p_d1 = d1 + bos * HQ + i_hq + row_indices[:, None] * HQ + p_bart = bart + bos * HQ + i_hq + row_indices[:, None] * HQ + p_m = m + bos * HQ + i_hq + row_indices[:, None] * HQ + + b_q = tl.load(p_q, mask=m_qk, other=0.0) + b_r = tl.load(p_r, mask=m_qk, other=0.0) m_acc = tl.zeros((BT, 1), dtype=tl.float32) - float("inf") d1_acc = tl.zeros((BT, 1), dtype=tl.float32) d2_acc = tl.zeros((BT, 1), dtype=tl.float32) @@ -106,9 +107,12 @@ def parallel_parallax_fwd_kernel( # Phase 0: left-border blocks (SWA only). Window mask only. for col_block_id in range(FIRST_COL_BLOCK, LEFT_BORDER_END): - col_indices = col_block_id * BS + tl.arange(0, BS) - b_k = tl.load(p_k, boundary_check=(0, 1), padding_option="zero") - b_v = tl.load(p_v, boundary_check=(0, 1), padding_option="zero") + col_indices = col_block_id.to(tl.int64) * BS + tl.arange(0, BS) + m_kv = (col_indices[:, None] < T) & m_k[None, :] + p_k = k + (bos * H + i_h) * K + col_indices[:, None] * (H * K) + o_k[None, :] + p_v = v + (bos * H + i_h) * K + col_indices[:, None] * (H * K) + o_k[None, :] + b_k = tl.load(p_k, mask=m_kv, other=0.0) + b_v = tl.load(p_v, mask=m_kv, other=0.0) mask = ( (col_indices[None, :] >= row_indices[:, None] - WINDOW_SIZE_LEFT + 1) & row_mask @@ -129,13 +133,15 @@ def parallel_parallax_fwd_kernel( barv_acc = tl.dot(w.to(b_v.dtype), b_v, out_dtype=tl.float32, acc=barv_acc) Rv_acc = tl.dot(wr.to(b_v.dtype), b_v, out_dtype=tl.float32, acc=Rv_acc) m_acc = m_new - p_k = tl.advance(p_k, (BS, 0)) - p_v = tl.advance(p_v, (BS, 0)) # Phase A: safe blocks (no mask). for _safe in range(SAFE_MIDDLE_START, NUM_SAFE_BLOCKS): - b_k = tl.load(p_k, boundary_check=(0, 1), padding_option="zero") - b_v = tl.load(p_v, boundary_check=(0, 1), padding_option="zero") + o_kv = _safe.to(tl.int64) * BS + tl.arange(0, BS) + m_kv = (o_kv[:, None] < T) & m_k[None, :] + p_k = k + (bos * H + i_h) * K + o_kv[:, None] * (H * K) + o_k[None, :] + p_v = v + (bos * H + i_h) * K + o_kv[:, None] * (H * K) + o_k[None, :] + b_k = tl.load(p_k, mask=m_kv, other=0.0) + b_v = tl.load(p_v, mask=m_kv, other=0.0) qk = tl.dot(b_q, tl.trans(b_k), out_dtype=tl.float32) * scale_log2 m_new = tl.maximum(m_acc, tl.max(qk, axis=1, keep_dims=True)) safe_m = tl.where(m_new == -float("inf"), 0.0, m_new) @@ -150,14 +156,15 @@ def parallel_parallax_fwd_kernel( barv_acc = tl.dot(w.to(b_v.dtype), b_v, out_dtype=tl.float32, acc=barv_acc) Rv_acc = tl.dot(wr.to(b_v.dtype), b_v, out_dtype=tl.float32, acc=Rv_acc) m_acc = m_new - p_k = tl.advance(p_k, (BS, 0)) - p_v = tl.advance(p_v, (BS, 0)) # Phase B: right-border blocks (causal + boundary + window mask). for col_block_id in range(RIGHT_BORDER_START, NUM_TOTAL_BLOCKS): - col_indices = col_block_id * BS + tl.arange(0, BS) - b_k = tl.load(p_k, boundary_check=(0, 1), padding_option="zero") - b_v = tl.load(p_v, boundary_check=(0, 1), padding_option="zero") + col_indices = col_block_id.to(tl.int64) * BS + tl.arange(0, BS) + m_kv = (col_indices[:, None] < T) & m_k[None, :] + p_k = k + (bos * H + i_h) * K + col_indices[:, None] * (H * K) + o_k[None, :] + p_v = v + (bos * H + i_h) * K + col_indices[:, None] * (H * K) + o_k[None, :] + b_k = tl.load(p_k, mask=m_kv, other=0.0) + b_v = tl.load(p_v, mask=m_kv, other=0.0) if WINDOW_SIZE_LEFT >= 0: mask = ( (row_indices[:, None] >= col_indices[None, :]) @@ -186,19 +193,17 @@ def parallel_parallax_fwd_kernel( barv_acc = tl.dot(w.to(b_v.dtype), b_v, out_dtype=tl.float32, acc=barv_acc) Rv_acc = tl.dot(wr.to(b_v.dtype), b_v, out_dtype=tl.float32, acc=Rv_acc) m_acc = m_new - p_k = tl.advance(p_k, (BS, 0)) - p_v = tl.advance(p_v, (BS, 0)) inv_d1 = tl.where(row_mask, 1.0 / d1_acc, 0.0) b_barv = barv_acc * inv_d1 b_bart = d2_acc * inv_d1 b_o = b_barv + b_bart * b_barv - Rv_acc * inv_d1 - tl.store(p_o, b_o.to(p_o.dtype.element_ty), boundary_check=(0, 1)) - tl.store(p_barv, b_barv.to(p_barv.dtype.element_ty), boundary_check=(0, 1)) - tl.store(p_d1, d1_acc, boundary_check=(0, 1)) - tl.store(p_bart, b_bart, boundary_check=(0, 1)) - tl.store(p_m, m_acc, boundary_check=(0, 1)) + tl.store(p_o, b_o.to(p_o.dtype.element_ty), mask=m_qk) + tl.store(p_barv, b_barv.to(p_barv.dtype.element_ty), mask=m_qk) + tl.store(p_d1, d1_acc, mask=row_mask) + tl.store(p_bart, b_bart, mask=row_mask) + tl.store(p_m, m_acc, mask=row_mask) @triton.heuristics({ @@ -220,32 +225,36 @@ def parallel_parallax_bwd_kernel_preprocess( BT: tl.constexpr, IS_VARLEN: tl.constexpr, ): - i_t, i_bh = tl.program_id(0), tl.program_id(1) + i_t, i_bh = tl.program_id(0).to(tl.int64), tl.program_id(1).to(tl.int64) i_b, i_hq = i_bh // HQ, i_bh % HQ 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) + i_n, i_t = tl.load(chunk_indices + i_t * 2).to(tl.int32), tl.load(chunk_indices + i_t * 2 + 1).to(tl.int64) bos, eos = tl.load(cu_seqlens + i_n).to(tl.int64), tl.load(cu_seqlens + i_n + 1).to(tl.int64) T = eos - bos else: bos = (i_b * T).to(tl.int64) row_offset = i_t * BT - p_grad_o = tl.make_block_ptr(grad_o + (bos * HQ + i_hq) * K, (T, K), (HQ * K, 1), (row_offset, 0), (BT, BK), (1, 0)) - p_o = tl.make_block_ptr(o + (bos * HQ + i_hq) * K, (T, K), (HQ * K, 1), (row_offset, 0), (BT, BK), (1, 0)) - p_barv = tl.make_block_ptr(barv + (bos * HQ + i_hq) * K, (T, K), (HQ * K, 1), (row_offset, 0), (BT, BK), (1, 0)) - p_t = tl.make_block_ptr(delta_t + bos * HQ + i_hq, (T, 1), (HQ, 1), (row_offset, 0), (BT, 1), (1, 0)) - p_b = tl.make_block_ptr(delta_b + bos * HQ + i_hq, (T, 1), (HQ, 1), (row_offset, 0), (BT, 1), (1, 0)) - - b_grad_o = tl.load(p_grad_o, boundary_check=(0, 1), padding_option="zero").to(tl.float32) - b_o = tl.load(p_o, boundary_check=(0, 1), padding_option="zero").to(tl.float32) - b_barv = tl.load(p_barv, boundary_check=(0, 1), padding_option="zero").to(tl.float32) + o_t = row_offset + tl.arange(0, BT) + o_k = tl.arange(0, BK) + m_tk = (o_t[:, None] < T) & (o_k[None, :] < K) + m_t = o_t[:, None] < T + p_grad_o = grad_o + (bos * HQ + i_hq) * K + o_t[:, None] * (HQ * K) + o_k[None, :] + p_o = o + (bos * HQ + i_hq) * K + o_t[:, None] * (HQ * K) + o_k[None, :] + p_barv = barv + (bos * HQ + i_hq) * K + o_t[:, None] * (HQ * K) + o_k[None, :] + p_t = delta_t + bos * HQ + i_hq + o_t[:, None] * HQ + p_b = delta_b + bos * HQ + i_hq + o_t[:, None] * HQ + + b_grad_o = tl.load(p_grad_o, mask=m_tk, other=0.0).to(tl.float32) + b_o = tl.load(p_o, mask=m_tk, other=0.0).to(tl.float32) + b_barv = tl.load(p_barv, mask=m_tk, other=0.0).to(tl.float32) b_t = tl.sum(b_grad_o * b_o, axis=1, keep_dims=True) b_b = tl.sum(b_grad_o * b_barv, axis=1, keep_dims=True) - tl.store(p_t, b_t, boundary_check=(0, 1)) - tl.store(p_b, b_b, boundary_check=(0, 1)) + tl.store(p_t, b_t, mask=m_t) + tl.store(p_b, b_b, mask=m_t) @triton.heuristics({ @@ -279,12 +288,12 @@ def parallel_parallax_bwd_kernel_dqr( BS: tl.constexpr, IS_VARLEN: tl.constexpr, ): - i_t, i_bh = tl.program_id(0), tl.program_id(1) + i_t, i_bh = tl.program_id(0).to(tl.int64), tl.program_id(1).to(tl.int64) i_b, i_hq = i_bh // HQ, i_bh % HQ i_h = i_hq // G 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) + i_n, i_t = tl.load(chunk_indices + i_t * 2).to(tl.int32), tl.load(chunk_indices + i_t * 2 + 1).to(tl.int64) bos, eos = tl.load(cu_seqlens + i_n).to(tl.int64), tl.load(cu_seqlens + i_n + 1).to(tl.int64) T = eos - bos else: @@ -311,27 +320,28 @@ def parallel_parallax_bwd_kernel_dqr( SAFE_MIDDLE_START = tl.maximum(FIRST_COL_BLOCK, SAFE_LEFT_START) RIGHT_BORDER_START = tl.maximum(FIRST_COL_BLOCK, NUM_SAFE_BLOCKS) - p_q = tl.make_block_ptr(q + (bos * HQ + i_hq) * K, (T, K), (HQ * K, 1), (row_offset, 0), (BT, BK), (1, 0)) - p_r = tl.make_block_ptr(r + (bos * HQ + i_hq) * K, (T, K), (HQ * K, 1), (row_offset, 0), (BT, BK), (1, 0)) - p_k = tl.make_block_ptr(k + (bos * H + i_h) * K, (T, K), (H * K, 1), (FIRST_COL_BLOCK * BS, 0), (BS, BK), (1, 0)) - p_v = tl.make_block_ptr(v + (bos * H + i_h) * K, (T, K), (H * K, 1), (FIRST_COL_BLOCK * BS, 0), (BS, BK), (1, 0)) - p_d1 = tl.make_block_ptr(d1 + bos * HQ + i_hq, (T, 1), (HQ, 1), (row_offset, 0), (BT, 1), (1, 0)) - p_bart = tl.make_block_ptr(bart + bos * HQ + i_hq, (T, 1), (HQ, 1), (row_offset, 0), (BT, 1), (1, 0)) - p_m = tl.make_block_ptr(m + bos * HQ + i_hq, (T, 1), (HQ, 1), (row_offset, 0), (BT, 1), (1, 0)) - p_t = tl.make_block_ptr(delta_t + bos * HQ + i_hq, (T, 1), (HQ, 1), (row_offset, 0), (BT, 1), (1, 0)) - p_b = tl.make_block_ptr(delta_b + bos * HQ + i_hq, (T, 1), (HQ, 1), (row_offset, 0), (BT, 1), (1, 0)) - p_grad_o = tl.make_block_ptr(grad_o + (bos * HQ + i_hq) * K, (T, K), (HQ * K, 1), (row_offset, 0), (BT, BK), (1, 0)) - p_grad_q = tl.make_block_ptr(grad_q + (bos * HQ + i_hq) * K, (T, K), (HQ * K, 1), (row_offset, 0), (BT, BK), (1, 0)) - p_grad_r = tl.make_block_ptr(grad_r + (bos * HQ + i_hq) * K, (T, K), (HQ * K, 1), (row_offset, 0), (BT, BK), (1, 0)) - - b_q = tl.load(p_q, boundary_check=(0, 1), padding_option="zero") - b_r = tl.load(p_r, boundary_check=(0, 1), padding_option="zero") - b_d1 = tl.load(p_d1, boundary_check=(0, 1), padding_option="zero") - b_bart = tl.load(p_bart, boundary_check=(0, 1), padding_option="zero") - b_m = tl.load(p_m, boundary_check=(0, 1), padding_option="zero") - b_t = tl.load(p_t, boundary_check=(0, 1), padding_option="zero") - b_b = tl.load(p_b, boundary_check=(0, 1), padding_option="zero") - grad_o_tile = tl.load(p_grad_o, boundary_check=(0, 1), padding_option="zero") + o_k = tl.arange(0, BK) + m_k = o_k < K + m_qk = row_mask & m_k[None, :] + p_q = q + (bos * HQ + i_hq) * K + row_indices[:, None] * (HQ * K) + o_k[None, :] + p_r = r + (bos * HQ + i_hq) * K + row_indices[:, None] * (HQ * K) + o_k[None, :] + p_d1 = d1 + bos * HQ + i_hq + row_indices[:, None] * HQ + p_bart = bart + bos * HQ + i_hq + row_indices[:, None] * HQ + p_m = m + bos * HQ + i_hq + row_indices[:, None] * HQ + p_t = delta_t + bos * HQ + i_hq + row_indices[:, None] * HQ + p_b = delta_b + bos * HQ + i_hq + row_indices[:, None] * HQ + p_grad_o = grad_o + (bos * HQ + i_hq) * K + row_indices[:, None] * (HQ * K) + o_k[None, :] + p_grad_q = grad_q + (bos * HQ + i_hq) * K + row_indices[:, None] * (HQ * K) + o_k[None, :] + p_grad_r = grad_r + (bos * HQ + i_hq) * K + row_indices[:, None] * (HQ * K) + o_k[None, :] + + b_q = tl.load(p_q, mask=m_qk, other=0.0) + b_r = tl.load(p_r, mask=m_qk, other=0.0) + b_d1 = tl.load(p_d1, mask=row_mask, other=0.0) + b_bart = tl.load(p_bart, mask=row_mask, other=0.0) + b_m = tl.load(p_m, mask=row_mask, other=0.0) + b_t = tl.load(p_t, mask=row_mask, other=0.0) + b_b = tl.load(p_b, mask=row_mask, other=0.0) + grad_o_tile = tl.load(p_grad_o, mask=m_qk, other=0.0) grad_q_acc = tl.zeros((BT, BK), dtype=tl.float32) grad_r_acc = tl.zeros((BT, BK), dtype=tl.float32) scale_log2 = scale * RCP_LN2 @@ -340,9 +350,12 @@ def parallel_parallax_bwd_kernel_dqr( # Phase 0: left-border blocks (SWA only). for col_block_id in range(FIRST_COL_BLOCK, LEFT_BORDER_END): - col_indices = col_block_id * BS + tl.arange(0, BS) - b_k = tl.load(p_k, boundary_check=(0, 1), padding_option="zero") - b_v = tl.load(p_v, boundary_check=(0, 1), padding_option="zero") + col_indices = col_block_id.to(tl.int64) * BS + tl.arange(0, BS) + m_kv = (col_indices[:, None] < T) & m_k[None, :] + p_k = k + (bos * H + i_h) * K + col_indices[:, None] * (H * K) + o_k[None, :] + p_v = v + (bos * H + i_h) * K + col_indices[:, None] * (H * K) + o_k[None, :] + b_k = tl.load(p_k, mask=m_kv, other=0.0) + b_v = tl.load(p_v, mask=m_kv, other=0.0) mask = ( (col_indices[None, :] >= row_indices[:, None] - WINDOW_SIZE_LEFT + 1) & row_mask @@ -360,13 +373,15 @@ def parallel_parallax_bwd_kernel_dqr( gu = -p * delta grad_q_acc = tl.dot(gl.to(b_k.dtype), b_k, out_dtype=tl.float32, acc=grad_q_acc) grad_r_acc = tl.dot(gu.to(b_k.dtype), b_k, out_dtype=tl.float32, acc=grad_r_acc) - p_k = tl.advance(p_k, (BS, 0)) - p_v = tl.advance(p_v, (BS, 0)) # Phase A: safe blocks (no mask). for _ in range(SAFE_MIDDLE_START, NUM_SAFE_BLOCKS): - b_k = tl.load(p_k, boundary_check=(0, 1), padding_option="zero") - b_v = tl.load(p_v, boundary_check=(0, 1), padding_option="zero") + o_kv = _.to(tl.int64) * BS + tl.arange(0, BS) + m_kv = (o_kv[:, None] < T) & m_k[None, :] + p_k = k + (bos * H + i_h) * K + o_kv[:, None] * (H * K) + o_k[None, :] + p_v = v + (bos * H + i_h) * K + o_kv[:, None] * (H * K) + o_k[None, :] + b_k = tl.load(p_k, mask=m_kv, other=0.0) + b_v = tl.load(p_v, mask=m_kv, other=0.0) qk = tl.dot(b_q, tl.trans(b_k), out_dtype=tl.float32) * scale_log2 w = exp2(qk - b_m) a = tl.dot(grad_o_tile, tl.trans(b_v), out_dtype=tl.float32) @@ -378,14 +393,15 @@ def parallel_parallax_bwd_kernel_dqr( gu = -p * delta grad_q_acc = tl.dot(gl.to(b_k.dtype), b_k, out_dtype=tl.float32, acc=grad_q_acc) grad_r_acc = tl.dot(gu.to(b_k.dtype), b_k, out_dtype=tl.float32, acc=grad_r_acc) - p_k = tl.advance(p_k, (BS, 0)) - p_v = tl.advance(p_v, (BS, 0)) # Phase B: right-border blocks (causal + boundary + window mask). for col_block_id in range(RIGHT_BORDER_START, NUM_TOTAL_BLOCKS): - col_indices = col_block_id * BS + tl.arange(0, BS) - b_k = tl.load(p_k, boundary_check=(0, 1), padding_option="zero") - b_v = tl.load(p_v, boundary_check=(0, 1), padding_option="zero") + col_indices = col_block_id.to(tl.int64) * BS + tl.arange(0, BS) + m_kv = (col_indices[:, None] < T) & m_k[None, :] + p_k = k + (bos * H + i_h) * K + col_indices[:, None] * (H * K) + o_k[None, :] + p_v = v + (bos * H + i_h) * K + col_indices[:, None] * (H * K) + o_k[None, :] + b_k = tl.load(p_k, mask=m_kv, other=0.0) + b_v = tl.load(p_v, mask=m_kv, other=0.0) if WINDOW_SIZE_LEFT >= 0: mask = ( (row_indices[:, None] >= col_indices[None, :]) @@ -411,13 +427,11 @@ def parallel_parallax_bwd_kernel_dqr( gu = -p * delta grad_q_acc = tl.dot(gl.to(b_k.dtype), b_k, out_dtype=tl.float32, acc=grad_q_acc) grad_r_acc = tl.dot(gu.to(b_k.dtype), b_k, out_dtype=tl.float32, acc=grad_r_acc) - p_k = tl.advance(p_k, (BS, 0)) - p_v = tl.advance(p_v, (BS, 0)) grad_q_acc = scale * grad_q_acc - tl.store(p_grad_q, grad_q_acc.to(p_grad_q.dtype.element_ty), boundary_check=(0, 1)) - tl.store(p_grad_r, grad_r_acc.to(p_grad_r.dtype.element_ty), boundary_check=(0, 1)) + tl.store(p_grad_q, grad_q_acc.to(p_grad_q.dtype.element_ty), mask=m_qk) + tl.store(p_grad_r, grad_r_acc.to(p_grad_r.dtype.element_ty), mask=m_qk) @triton.heuristics({ @@ -451,12 +465,12 @@ def parallel_parallax_bwd_kernel_dkv( BS: tl.constexpr, IS_VARLEN: tl.constexpr, ): - i_t, i_bh = tl.program_id(0), tl.program_id(1) + i_t, i_bh = tl.program_id(0).to(tl.int64), tl.program_id(1).to(tl.int64) i_b, i_hq = i_bh // HQ, i_bh % HQ i_h = i_hq // G 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) + i_n, i_t = tl.load(chunk_indices + i_t * 2).to(tl.int32), tl.load(chunk_indices + i_t * 2 + 1).to(tl.int64) bos, eos = tl.load(cu_seqlens + i_n).to(tl.int64), tl.load(cu_seqlens + i_n + 1).to(tl.int64) T = eos - bos else: @@ -467,7 +481,6 @@ def parallel_parallax_bwd_kernel_dkv( col_indices = col_offset + tl.arange(0, BS) start_row_block = col_offset // BT - start_row_offset = start_row_block * BT num_row_blocks_qbound = tl.cdiv(T, BT) if WINDOW_SIZE_LEFT >= 0: @@ -478,21 +491,16 @@ def parallel_parallax_bwd_kernel_dkv( num_row_blocks = num_row_blocks_qbound WINDOW_SAFE_END = num_row_blocks - p_q = tl.make_block_ptr(q + (bos * HQ + i_hq) * K, (T, K), (HQ * K, 1), (start_row_offset, 0), (BT, BK), (1, 0)) - p_r = tl.make_block_ptr(r + (bos * HQ + i_hq) * K, (T, K), (HQ * K, 1), (start_row_offset, 0), (BT, BK), (1, 0)) - p_k = tl.make_block_ptr(k + (bos * H + i_h) * K, (T, K), (H * K, 1), (col_offset, 0), (BS, BK), (1, 0)) - p_v = tl.make_block_ptr(v + (bos * H + i_h) * K, (T, K), (H * K, 1), (col_offset, 0), (BS, BK), (1, 0)) - p_d1 = tl.make_block_ptr(d1 + bos * HQ + i_hq, (T, 1), (HQ, 1), (start_row_offset, 0), (BT, 1), (1, 0)) - p_bart = tl.make_block_ptr(bart + bos * HQ + i_hq, (T, 1), (HQ, 1), (start_row_offset, 0), (BT, 1), (1, 0)) - p_m = tl.make_block_ptr(m + bos * HQ + i_hq, (T, 1), (HQ, 1), (start_row_offset, 0), (BT, 1), (1, 0)) - p_t = tl.make_block_ptr(delta_t + bos * HQ + i_hq, (T, 1), (HQ, 1), (start_row_offset, 0), (BT, 1), (1, 0)) - p_b = tl.make_block_ptr(delta_b + bos * HQ + i_hq, (T, 1), (HQ, 1), (start_row_offset, 0), (BT, 1), (1, 0)) - p_grad_o = tl.make_block_ptr(grad_o + (bos * HQ + i_hq) * K, (T, K), (HQ * K, 1), (start_row_offset, 0), (BT, BK), (1, 0)) - p_grad_k = tl.make_block_ptr(grad_k + (bos * HQ + i_hq) * K, (T, K), (HQ * K, 1), (col_offset, 0), (BS, BK), (1, 0)) - p_grad_v = tl.make_block_ptr(grad_v + (bos * HQ + i_hq) * K, (T, K), (HQ * K, 1), (col_offset, 0), (BS, BK), (1, 0)) - - b_k = tl.load(p_k, boundary_check=(0, 1), padding_option="zero") - b_v = tl.load(p_v, boundary_check=(0, 1), padding_option="zero") + o_k = tl.arange(0, BK) + m_k = o_k < K + m_ck = (col_indices[:, None] < T) & m_k[None, :] + p_k = k + (bos * H + i_h) * K + col_indices[:, None] * (H * K) + o_k[None, :] + p_v = v + (bos * H + i_h) * K + col_indices[:, None] * (H * K) + o_k[None, :] + p_grad_k = grad_k + (bos * HQ + i_hq) * K + col_indices[:, None] * (HQ * K) + o_k[None, :] + p_grad_v = grad_v + (bos * HQ + i_hq) * K + col_indices[:, None] * (HQ * K) + o_k[None, :] + + b_k = tl.load(p_k, mask=m_ck, other=0.0) + b_v = tl.load(p_v, mask=m_ck, other=0.0) grad_k_acc = tl.zeros((BS, BK), dtype=tl.float32) grad_v_acc = tl.zeros((BS, BK), dtype=tl.float32) scale_log2 = scale * RCP_LN2 @@ -504,17 +512,26 @@ def parallel_parallax_bwd_kernel_dkv( # Phase A: causal-border row blocks. causal_end = tl.minimum(first_safe_row_block, num_row_blocks) for row_block_id in range(start_row_block, causal_end): - row_offset = row_block_id * BT + row_offset = row_block_id.to(tl.int64) * BT row_indices = row_offset + tl.arange(0, BT) row_mask = row_indices[:, None] < T - b_q = tl.load(p_q, boundary_check=(0, 1), padding_option="zero") - b_r = tl.load(p_r, boundary_check=(0, 1), padding_option="zero") - b_d1 = tl.load(p_d1, boundary_check=(0, 1), padding_option="zero") - b_bart = tl.load(p_bart, boundary_check=(0, 1), padding_option="zero") - b_m = tl.load(p_m, boundary_check=(0, 1), padding_option="zero") - b_t = tl.load(p_t, boundary_check=(0, 1), padding_option="zero") - b_b = tl.load(p_b, boundary_check=(0, 1), padding_option="zero") - grad_o_tile = tl.load(p_grad_o, boundary_check=(0, 1), padding_option="zero") + m_qk = row_mask & m_k[None, :] + p_q = q + (bos * HQ + i_hq) * K + row_indices[:, None] * (HQ * K) + o_k[None, :] + p_r = r + (bos * HQ + i_hq) * K + row_indices[:, None] * (HQ * K) + o_k[None, :] + p_d1 = d1 + bos * HQ + i_hq + row_indices[:, None] * HQ + p_bart = bart + bos * HQ + i_hq + row_indices[:, None] * HQ + p_m = m + bos * HQ + i_hq + row_indices[:, None] * HQ + p_t = delta_t + bos * HQ + i_hq + row_indices[:, None] * HQ + p_b = delta_b + bos * HQ + i_hq + row_indices[:, None] * HQ + p_grad_o = grad_o + (bos * HQ + i_hq) * K + row_indices[:, None] * (HQ * K) + o_k[None, :] + b_q = tl.load(p_q, mask=m_qk, other=0.0) + b_r = tl.load(p_r, mask=m_qk, other=0.0) + b_d1 = tl.load(p_d1, mask=row_mask, other=0.0) + b_bart = tl.load(p_bart, mask=row_mask, other=0.0) + b_m = tl.load(p_m, mask=row_mask, other=0.0) + b_t = tl.load(p_t, mask=row_mask, other=0.0) + b_b = tl.load(p_b, mask=row_mask, other=0.0) + grad_o_tile = tl.load(p_grad_o, mask=m_qk, other=0.0) qk = tl.dot(b_q, tl.trans(b_k), out_dtype=tl.float32) * scale_log2 rk = tl.dot(b_r, tl.trans(b_k), out_dtype=tl.float32) @@ -545,29 +562,29 @@ def parallel_parallax_bwd_kernel_dkv( weights = p * (1 + bart_minus_rk) grad_v_acc = tl.dot(tl.trans(weights).to(grad_o_tile.dtype), grad_o_tile, out_dtype=tl.float32, acc=grad_v_acc) - p_q = tl.advance(p_q, (BT, 0)) - p_r = tl.advance(p_r, (BT, 0)) - p_d1 = tl.advance(p_d1, (BT, 0)) - p_bart = tl.advance(p_bart, (BT, 0)) - p_m = tl.advance(p_m, (BT, 0)) - p_t = tl.advance(p_t, (BT, 0)) - p_b = tl.advance(p_b, (BT, 0)) - p_grad_o = tl.advance(p_grad_o, (BT, 0)) - # Phase B: safe row blocks (no causal/col/window mask). safe_b_start = tl.maximum(first_safe_row_block, start_row_block) for row_block_id in range(safe_b_start, SAFE_MIDDLE_END): - row_offset = row_block_id * BT + row_offset = row_block_id.to(tl.int64) * BT row_indices = row_offset + tl.arange(0, BT) row_mask = row_indices[:, None] < T - b_q = tl.load(p_q, boundary_check=(0, 1), padding_option="zero") - b_r = tl.load(p_r, boundary_check=(0, 1), padding_option="zero") - b_d1 = tl.load(p_d1, boundary_check=(0, 1), padding_option="zero") - b_bart = tl.load(p_bart, boundary_check=(0, 1), padding_option="zero") - b_m = tl.load(p_m, boundary_check=(0, 1), padding_option="zero") - b_t = tl.load(p_t, boundary_check=(0, 1), padding_option="zero") - b_b = tl.load(p_b, boundary_check=(0, 1), padding_option="zero") - grad_o_tile = tl.load(p_grad_o, boundary_check=(0, 1), padding_option="zero") + m_qk = row_mask & m_k[None, :] + p_q = q + (bos * HQ + i_hq) * K + row_indices[:, None] * (HQ * K) + o_k[None, :] + p_r = r + (bos * HQ + i_hq) * K + row_indices[:, None] * (HQ * K) + o_k[None, :] + p_d1 = d1 + bos * HQ + i_hq + row_indices[:, None] * HQ + p_bart = bart + bos * HQ + i_hq + row_indices[:, None] * HQ + p_m = m + bos * HQ + i_hq + row_indices[:, None] * HQ + p_t = delta_t + bos * HQ + i_hq + row_indices[:, None] * HQ + p_b = delta_b + bos * HQ + i_hq + row_indices[:, None] * HQ + p_grad_o = grad_o + (bos * HQ + i_hq) * K + row_indices[:, None] * (HQ * K) + o_k[None, :] + b_q = tl.load(p_q, mask=m_qk, other=0.0) + b_r = tl.load(p_r, mask=m_qk, other=0.0) + b_d1 = tl.load(p_d1, mask=row_mask, other=0.0) + b_bart = tl.load(p_bart, mask=row_mask, other=0.0) + b_m = tl.load(p_m, mask=row_mask, other=0.0) + b_t = tl.load(p_t, mask=row_mask, other=0.0) + b_b = tl.load(p_b, mask=row_mask, other=0.0) + grad_o_tile = tl.load(p_grad_o, mask=m_qk, other=0.0) qk = tl.dot(b_q, tl.trans(b_k), out_dtype=tl.float32) * scale_log2 rk = tl.dot(b_r, tl.trans(b_k), out_dtype=tl.float32) @@ -584,29 +601,29 @@ def parallel_parallax_bwd_kernel_dkv( weights = p * (1 + bart_minus_rk) grad_v_acc = tl.dot(tl.trans(weights).to(grad_o_tile.dtype), grad_o_tile, out_dtype=tl.float32, acc=grad_v_acc) - p_q = tl.advance(p_q, (BT, 0)) - p_r = tl.advance(p_r, (BT, 0)) - p_d1 = tl.advance(p_d1, (BT, 0)) - p_bart = tl.advance(p_bart, (BT, 0)) - p_m = tl.advance(p_m, (BT, 0)) - p_t = tl.advance(p_t, (BT, 0)) - p_b = tl.advance(p_b, (BT, 0)) - p_grad_o = tl.advance(p_grad_o, (BT, 0)) - # Phase C: window-border row blocks (SWA only). window_border_start = tl.maximum(WINDOW_BORDER_START, start_row_block) for row_block_id in range(window_border_start, num_row_blocks): - row_offset = row_block_id * BT + row_offset = row_block_id.to(tl.int64) * BT row_indices = row_offset + tl.arange(0, BT) row_mask = row_indices[:, None] < T - b_q = tl.load(p_q, boundary_check=(0, 1), padding_option="zero") - b_r = tl.load(p_r, boundary_check=(0, 1), padding_option="zero") - b_d1 = tl.load(p_d1, boundary_check=(0, 1), padding_option="zero") - b_bart = tl.load(p_bart, boundary_check=(0, 1), padding_option="zero") - b_m = tl.load(p_m, boundary_check=(0, 1), padding_option="zero") - b_t = tl.load(p_t, boundary_check=(0, 1), padding_option="zero") - b_b = tl.load(p_b, boundary_check=(0, 1), padding_option="zero") - grad_o_tile = tl.load(p_grad_o, boundary_check=(0, 1), padding_option="zero") + m_qk = row_mask & m_k[None, :] + p_q = q + (bos * HQ + i_hq) * K + row_indices[:, None] * (HQ * K) + o_k[None, :] + p_r = r + (bos * HQ + i_hq) * K + row_indices[:, None] * (HQ * K) + o_k[None, :] + p_d1 = d1 + bos * HQ + i_hq + row_indices[:, None] * HQ + p_bart = bart + bos * HQ + i_hq + row_indices[:, None] * HQ + p_m = m + bos * HQ + i_hq + row_indices[:, None] * HQ + p_t = delta_t + bos * HQ + i_hq + row_indices[:, None] * HQ + p_b = delta_b + bos * HQ + i_hq + row_indices[:, None] * HQ + p_grad_o = grad_o + (bos * HQ + i_hq) * K + row_indices[:, None] * (HQ * K) + o_k[None, :] + b_q = tl.load(p_q, mask=m_qk, other=0.0) + b_r = tl.load(p_r, mask=m_qk, other=0.0) + b_d1 = tl.load(p_d1, mask=row_mask, other=0.0) + b_bart = tl.load(p_bart, mask=row_mask, other=0.0) + b_m = tl.load(p_m, mask=row_mask, other=0.0) + b_t = tl.load(p_t, mask=row_mask, other=0.0) + b_b = tl.load(p_b, mask=row_mask, other=0.0) + grad_o_tile = tl.load(p_grad_o, mask=m_qk, other=0.0) qk = tl.dot(b_q, tl.trans(b_k), out_dtype=tl.float32) * scale_log2 rk = tl.dot(b_r, tl.trans(b_k), out_dtype=tl.float32) @@ -629,17 +646,8 @@ def parallel_parallax_bwd_kernel_dkv( weights = p * (1 + bart_minus_rk) grad_v_acc = tl.dot(tl.trans(weights).to(grad_o_tile.dtype), grad_o_tile, out_dtype=tl.float32, acc=grad_v_acc) - p_q = tl.advance(p_q, (BT, 0)) - p_r = tl.advance(p_r, (BT, 0)) - p_d1 = tl.advance(p_d1, (BT, 0)) - p_bart = tl.advance(p_bart, (BT, 0)) - p_m = tl.advance(p_m, (BT, 0)) - p_t = tl.advance(p_t, (BT, 0)) - p_b = tl.advance(p_b, (BT, 0)) - p_grad_o = tl.advance(p_grad_o, (BT, 0)) - - tl.store(p_grad_k, grad_k_acc.to(p_grad_k.dtype.element_ty), boundary_check=(0, 1)) - tl.store(p_grad_v, grad_v_acc.to(p_grad_v.dtype.element_ty), boundary_check=(0, 1)) + tl.store(p_grad_k, grad_k_acc.to(p_grad_k.dtype.element_ty), mask=m_ck) + tl.store(p_grad_v, grad_v_acc.to(p_grad_v.dtype.element_ty), mask=m_ck) def parallel_parallax_fwd(q, r, k, v, scale, cu_seqlens=None, chunk_indices=None, window_size_left=-1): @@ -756,7 +764,6 @@ def parallel_parallax( scale: float | None = None, window_size: int | None = None, cu_seqlens: torch.LongTensor | None = None, - **kwargs, ) -> torch.Tensor: r""" Causal Parallax (parameterized local linear attention) with autograd, @@ -788,14 +795,13 @@ def parallel_parallax( o (torch.Tensor): output of shape `[B, T, HQ, D]`. """ - if 'head_first' in kwargs: - raise DeprecationWarning( - "head_first has been removed. Inputs must be in `[B, T, H, ...]` format.", - ) if q.dtype not in (torch.bfloat16, torch.float16): raise TypeError(f"parallel_parallax requires bf16 or fp16 inputs, got q.dtype={q.dtype}") if scale is None: scale = k.shape[-1] ** -0.5 + HQ, H = q.shape[2], k.shape[2] + if H == 0 or HQ % H != 0: + raise ValueError(f"The number of query heads ({HQ}) must be divisible by the number of key/value heads ({H}).") if cu_seqlens is not None and q.shape[0] != 1: raise ValueError( f"The batch size is expected to be 1 rather than {q.shape[0]} when using `cu_seqlens`. " diff --git a/build/torch-cuda/ops/path_attn/cumprod_householder_bwd.py b/build/torch-cuda/ops/path_attn/cumprod_householder_bwd.py index b7bd259ca925d3b4d0598139591f18a9a8a8b7d6..964ff127bfd91e3ba10ba8eb339b4c32c4d1a56d 100644 --- a/build/torch-cuda/ops/path_attn/cumprod_householder_bwd.py +++ b/build/torch-cuda/ops/path_attn/cumprod_householder_bwd.py @@ -31,7 +31,7 @@ def chunk_cumprod_householder_bwd_kernel( HQ: tl.constexpr, IS_VARLEN: tl.constexpr, ): - i_ss, i_hq = tl.program_id(0), tl.program_id(1) + i_ss, i_hq = tl.program_id(0).to(tl.int64), tl.program_id(1).to(tl.int64) i_h = i_hq // G if IS_VARLEN: @@ -39,14 +39,14 @@ def chunk_cumprod_householder_bwd_kernel( bos, eos = tl.load(cu_seqlens + i_n).to(tl.int64), tl.load(cu_seqlens + i_n + 1).to(tl.int64) T = (eos - bos).to(tl.int32) NS = tl.cdiv(T, S) - boh = tl.load(chunk_offsets + i_n).to(tl.int32) - boh_large = tl.load(split_offsets + i_n).to(tl.int32) + boh = tl.load(chunk_offsets + i_n).to(tl.int64) + boh_large = tl.load(split_offsets + i_n).to(tl.int64) else: NS = tl.cdiv(T, S) i_n, i_s = i_ss // NS, i_ss % NS bos, eos = (i_n * T).to(tl.int64), (i_n * T + T).to(tl.int64) - boh = i_n * tl.cdiv(T, BT) - boh_large = i_n * tl.cdiv(T, S) + boh = (i_n * tl.cdiv(T, BT)).to(tl.int64) + boh_large = (i_n * tl.cdiv(T, S)).to(tl.int64) # offset calculations dhc_whole += ((boh_large + i_s) * HQ + i_hq) * K * K @@ -63,39 +63,44 @@ def chunk_cumprod_householder_bwd_kernel( stride_h = H * K * K NT_small = tl.cdiv(min(S, T-i_s*S), BT) - p_dhc_whole = tl.make_block_ptr(dhc_whole, (K, K), (K, 1), (0, 0), (BK, BK), (1, 0)) + o_d = tl.arange(0, BK) + m_d = o_d < K + p_dhc_whole = dhc_whole + o_d[:, None] * K + o_d[None, :] b_dhc = tl.zeros([BK, BK], dtype=tl.float32) - b_dhc += tl.load(p_dhc_whole, boundary_check=(0, 1)) + b_dhc += tl.load(p_dhc_whole, mask=m_d[:, None] & m_d[None, :], other=0.0) # calculate dh for i_t_small in range(0, NT_small): - p_k = tl.make_block_ptr(k, (T, K), (H*K, 1), (i_s*S + i_t_small*BT, 0), (BT, BK), (1, 0)) - p_dk = tl.make_block_ptr(dk, (T, K), (HQ*K, 1), (i_s*S + i_t_small*BT, 0), (BT, BK), (1, 0)) - p_dk_new = tl.make_block_ptr(dk_new, (T, K), (HQ*K, 1), (i_s*S + i_t_small*BT, 0), (BT, BK), (1, 0)) - p_hc = tl.make_block_ptr(hc_suffix + i_t_small * stride_h, (K, K), (K, 1), (0, 0), (BK, BK), (1, 0)) + o_k = (i_s*S + i_t_small*BT).to(tl.int64) + tl.arange(0, BT) + m_k = o_k < T + m_kk = m_k[:, None] & m_d[None, :] + p_k = k + o_k[:, None] * (H*K) + o_d[None, :] + p_dk = dk + o_k[:, None] * (HQ*K) + o_d[None, :] + p_dk_new = dk_new + o_k[:, None] * (HQ*K) + o_d[None, :] + p_hc = hc_suffix + i_t_small * stride_h + o_d[:, None] * K + o_d[None, :] - b_k = tl.load(p_k, boundary_check=(0, 1)) - b_dk = tl.load(p_dk, boundary_check=(0, 1)) + b_k = tl.load(p_k, mask=m_kk, other=0.0) + b_dk = tl.load(p_dk, mask=m_kk, other=0.0) - p_w1 = tl.make_block_ptr(w1, (T, K), (H*K, 1), (i_s*S + i_t_small*BT, 0), (BT, BK), (1, 0)) - p_w2 = tl.make_block_ptr(w2, (T, K), (H*K, 1), (i_s*S + i_t_small*BT, 0), (BT, BK), (1, 0)) + p_w1 = w1 + o_k[:, None] * (H*K) + o_d[None, :] + p_w2 = w2 + o_k[:, None] * (H*K) + o_d[None, :] - b_w1 = tl.load(p_w1, boundary_check=(0, 1)) - b_w2 = tl.load(p_w2, boundary_check=(0, 1)) - b_hc = tl.load(p_hc, boundary_check=(0, 1)) + b_w1 = tl.load(p_w1, mask=m_kk, other=0.0) + b_w2 = tl.load(p_w2, mask=m_kk, other=0.0) + b_hc = tl.load(p_hc, mask=m_d[:, None] & m_d[None, :], other=0.0) b_dk_new = b_dk - tl.dot(b_dk.to(b_hc.dtype), b_hc) - tl.store(p_dk_new, b_dk_new.to(dk_new.dtype.element_ty), boundary_check=(0, 1)) + tl.store(p_dk_new, b_dk_new.to(dk_new.dtype.element_ty), mask=m_kk) b_dh = b_dhc - tl.dot(tl.trans(b_hc), b_dhc.to(b_hc.dtype)) b_dw2 = tl.dot(b_w1, b_dh.to(b_w1.dtype)) b_dw1 = tl.dot(b_w2, tl.trans(b_dh.to(b_w2.dtype))) - p_dw1 = tl.make_block_ptr(dw1, (T, K), (HQ*K, 1), (i_s*S + i_t_small*BT, 0), (BT, BK), (1, 0)) - p_dw2 = tl.make_block_ptr(dw2, (T, K), (HQ*K, 1), (i_s*S + i_t_small*BT, 0), (BT, BK), (1, 0)) + p_dw1 = dw1 + o_k[:, None] * (HQ*K) + o_d[None, :] + p_dw2 = dw2 + o_k[:, None] * (HQ*K) + o_d[None, :] - tl.store(p_dw1, b_dw1.to(dw1.dtype.element_ty), boundary_check=(0, 1)) - tl.store(p_dw2, b_dw2.to(dw2.dtype.element_ty), boundary_check=(0, 1)) + tl.store(p_dw1, b_dw1.to(dw1.dtype.element_ty), mask=m_kk) + tl.store(p_dw2, b_dw2.to(dw2.dtype.element_ty), mask=m_kk) b_dhc = b_dhc - tl.dot(tl.dot(b_dhc.to(b_w2.dtype), tl.trans(b_w2)).to(b_w1.dtype), b_w1) b_dhc -= tl.dot(tl.trans(b_dk).to(b_k.dtype), b_k) diff --git a/build/torch-cuda/ops/path_attn/cumprod_householder_fwd.py b/build/torch-cuda/ops/path_attn/cumprod_householder_fwd.py index 1ced41573e5352b81472bcb492911891e126fb68..52943768c0d8e0be5ec3d3263e8c989ca657c9eb 100644 --- a/build/torch-cuda/ops/path_attn/cumprod_householder_fwd.py +++ b/build/torch-cuda/ops/path_attn/cumprod_householder_fwd.py @@ -16,7 +16,7 @@ from ...utils import check_shared_mem @triton.heuristics({ 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None, }) -@triton.jit +@triton.jit(do_not_specialize=['T']) def chunk_cumprod_householder_fwd_kernel( k, k_new, @@ -32,11 +32,11 @@ def chunk_cumprod_householder_fwd_kernel( K: tl.constexpr, H: tl.constexpr, BK: tl.constexpr, - T: tl.constexpr, + T, S: tl.constexpr, # split size, aka large chunk size IS_VARLEN: tl.constexpr, ): - i_ss, i_h = tl.program_id(0), tl.program_id(1) + i_ss, i_h = tl.program_id(0).to(tl.int64), tl.program_id(1).to(tl.int64) if IS_VARLEN: i_n, i_s = tl.load(split_indices + i_ss * 2).to(tl.int32), tl.load(split_indices + i_ss * 2 + 1).to(tl.int32) @@ -44,15 +44,15 @@ def chunk_cumprod_householder_fwd_kernel( T = (eos - bos).to(tl.int32) NS = tl.cdiv(T, S) - boh = tl.load(chunk_offsets + i_n).to(tl.int32) - boh_large = tl.load(split_offsets + i_n).to(tl.int32) + boh = tl.load(chunk_offsets + i_n).to(tl.int64) + boh_large = tl.load(split_offsets + i_n).to(tl.int64) else: NS = tl.cdiv(T, S) i_n, i_s = i_ss // NS, i_ss % NS bos, eos = (i_n * T).to(tl.int64), (i_n * T + T).to(tl.int64) - boh = i_n * tl.cdiv(T, BT) - boh_large = i_n * tl.cdiv(T, S) + boh = (i_n * tl.cdiv(T, BT)).to(tl.int64) + boh_large = (i_n * tl.cdiv(T, S)).to(tl.int64) NT_small = tl.cdiv(min(S, T-i_s*S), BT) stride_h = H*K*K @@ -66,24 +66,28 @@ def chunk_cumprod_householder_fwd_kernel( w1 += (bos * H + i_h) * K w2 += (bos * H + i_h) * K + o_d = tl.arange(0, BK) + m_d = o_d < K b_h = tl.zeros([BK, BK], dtype=tl.float32) for i_t_small in range(NT_small-1, -1, -1): - p_hc_suffix = tl.make_block_ptr(hc_suffix + i_t_small * stride_h, (K, K), (K, 1), (0, 0), (BK, BK), (1, 0)) - tl.store(p_hc_suffix, b_h.to(hc_suffix.dtype.element_ty), boundary_check=(0, 1)) - p_k = tl.make_block_ptr(k, (T, K), (H*K, 1), (i_s * S + i_t_small * BT, 0), (BT, BK), (1, 0)) - b_k = tl.load(p_k, boundary_check=(0, 1)) + o_k = (i_s * S + i_t_small * BT).to(tl.int64) + tl.arange(0, BT) + m_k = o_k < T + p_hc_suffix = hc_suffix + i_t_small * stride_h + o_d[:, None] * K + o_d[None, :] + tl.store(p_hc_suffix, b_h.to(hc_suffix.dtype.element_ty), mask=m_d[:, None] & m_d[None, :]) + p_k = k + o_k[:, None] * (H*K) + o_d[None, :] + b_k = tl.load(p_k, mask=m_k[:, None] & m_d[None, :], other=0.0) b_k = (b_k - tl.dot(b_k, tl.trans(b_h.to(b_k.dtype)))) - p_w1 = tl.make_block_ptr(w1, (K, T), (1, H*K), (0, i_s * S + i_t_small * BT), (BK, BT), (0, 1)) - p_w2 = tl.make_block_ptr(w2, (T, K), (H*K, 1), (i_s * S + i_t_small * BT, 0), (BT, BK), (1, 0)) - b_w1 = tl.load(p_w1, boundary_check=(0, 1)) - b_w2 = tl.load(p_w2, boundary_check=(0, 1)) + p_w1 = w1 + o_d[:, None] + o_k[None, :] * (H*K) + p_w2 = w2 + o_k[:, None] * (H*K) + o_d[None, :] + b_w1 = tl.load(p_w1, mask=m_d[:, None] & m_k[None, :], other=0.0) + b_w2 = tl.load(p_w2, mask=m_k[:, None] & m_d[None, :], other=0.0) b_v_new = (b_w1 - tl.dot(b_h.to(b_w1.dtype), b_w1)).to(b_w2.dtype) - b_h += tl.dot(b_v_new, b_w2) - p_k_new = tl.make_block_ptr(k_new, (T, K), (H*K, 1), (i_s * S + i_t_small * BT, 0), (BT, BK), (1, 0)) - tl.store(p_k_new, b_k.to(k_new.dtype.element_ty), boundary_check=(0, 1)) + b_h = tl.dot(b_v_new, b_w2, b_h) + p_k_new = k_new + o_k[:, None] * (H*K) + o_d[None, :] + tl.store(p_k_new, b_k.to(k_new.dtype.element_ty), mask=m_k[:, None] & m_d[None, :]) - p_hc_whole = tl.make_block_ptr(hc_whole, (K, K), (K, 1), (0, 0), (BK, BK), (1, 0)) - tl.store(p_hc_whole, b_h.to(hc_whole.dtype.element_ty), boundary_check=(0, 1)) + p_hc_whole = hc_whole + o_d[:, None] * K + o_d[None, :] + tl.store(p_hc_whole, b_h.to(hc_whole.dtype.element_ty), mask=m_d[:, None] & m_d[None, :]) def chunk_cumprod_householder_fwd_fn( diff --git a/build/torch-cuda/ops/path_attn/intra_chunk_preprocess_bwd.py b/build/torch-cuda/ops/path_attn/intra_chunk_preprocess_bwd.py index d7b3674d2740c0a1f323b4e51dbf8ed42962942f..591f993775346d02756e16f24fe7f9aed2275680 100644 --- a/build/torch-cuda/ops/path_attn/intra_chunk_preprocess_bwd.py +++ b/build/torch-cuda/ops/path_attn/intra_chunk_preprocess_bwd.py @@ -27,12 +27,12 @@ def intra_chunk_preprocess_bwd_kernel( K: tl.constexpr, BT: tl.constexpr, BK: tl.constexpr, IS_VARLEN: tl.constexpr, ): - i_t, i_nh = tl.program_id(0), tl.program_id(1) + i_t, i_nh = tl.program_id(0).to(tl.int64), tl.program_id(1).to(tl.int64) i_n, i_hq = i_nh // HQ, i_nh % HQ i_h = i_hq // G if IS_VARLEN: - i_n, i_t = tl.load(indices + i_t * 2).to(tl.int32), tl.load(indices + i_t * 2 + 1).to(tl.int32) + i_n, i_t = tl.load(indices + i_t * 2).to(tl.int32), tl.load(indices + i_t * 2 + 1).to(tl.int64) bos, eos = tl.load(offsets + i_n).to(tl.int64), tl.load(offsets + i_n + 1).to(tl.int64) T = (eos - bos).to(tl.int32) else: @@ -43,79 +43,84 @@ def intra_chunk_preprocess_bwd_kernel( b_dw = tl.zeros([BT, BK], dtype=tl.float32) b_dT = tl.zeros([BT, BT], dtype=tl.float32) - p_q = tl.make_block_ptr(q + (bos * HQ + i_hq) * K, (T, K), (K*HQ, 1), (i_t * BT, 0), (BT, BK), (1, 0)) - p_k = tl.make_block_ptr(k + (bos * H + i_h) * K, (T, K), (K*H, 1), (i_t * BT, 0), (BT, BK), (1, 0)) - p_w = tl.make_block_ptr(w + (bos * H + i_h) * K, (T, K), (K*H, 1), (i_t * BT, 0), (BT, BK), (1, 0)) - p_w2 = tl.make_block_ptr(w2 + (bos * H + i_h) * K, (T, K), (K*H, 1), (i_t * BT, 0), (BT, BK), (1, 0)) - p_beta = tl.make_block_ptr(beta + (bos * H + i_h), (T, ), (H, ), (i_t * BT, ), (BT, ), (0, )) - p_T = tl.make_block_ptr(AT + (bos * H + i_h) * BT, (T, BT), (BT*H, 1), (i_t * BT, 0), (BT, BT), (1, 0)) - b_w = tl.load(p_w, boundary_check=(0, 1)) - b_Twb = tl.load(p_w2, boundary_check=(0, 1)) - b_beta = tl.load(p_beta, boundary_check=(0, )) - b_q = tl.load(p_q, boundary_check=(0, 1)) - b_k = tl.load(p_k, boundary_check=(0, 1)) - b_T = tl.load(p_T, boundary_check=(0, 1)) + o_t = i_t * BT + tl.arange(0, BT) + o_i = tl.arange(0, BT) + o_d = tl.arange(0, BK) + m_t = o_t < T + m_k = m_t[:, None] & (o_d[None, :] < K) + m_T = m_t[:, None] & (o_i[None, :] < BT) + p_q = q + (bos * HQ + i_hq) * K + o_t[:, None] * (K*HQ) + o_d[None, :] + p_k = k + (bos * H + i_h) * K + o_t[:, None] * (K*H) + o_d[None, :] + p_w = w + (bos * H + i_h) * K + o_t[:, None] * (K*H) + o_d[None, :] + p_w2 = w2 + (bos * H + i_h) * K + o_t[:, None] * (K*H) + o_d[None, :] + p_beta = beta + (bos * H + i_h) + o_t * H + p_T = AT + (bos * H + i_h) * BT + o_t[:, None] * (BT*H) + o_i[None, :] + b_w = tl.load(p_w, mask=m_k, other=0.0) + b_Twb = tl.load(p_w2, mask=m_k, other=0.0) + b_beta = tl.load(p_beta, mask=m_t, other=0.0) + b_q = tl.load(p_q, mask=m_k, other=0.0) + b_k = tl.load(p_k, mask=m_k, other=0.0) + b_T = tl.load(p_T, mask=m_T, other=0.0) b_w_beta = (b_w * b_beta[:, None]).to(b_w.dtype) - o_i = tl.arange(0, BT) b_qw = tl.where(o_i[:, None] >= o_i[None, :], tl.dot(b_q, tl.trans(b_w)), 0).to(b_q.dtype) b_wbk = tl.where(o_i[:, None] > o_i[None, :], tl.dot(b_w_beta, tl.trans(b_k)), 0).to(b_k.dtype) b_Twbk = tl.dot(b_T, b_wbk).to(b_w.dtype) - p_dA_local = tl.make_block_ptr(dA_local + (bos * HQ + i_hq) * BT, (T, BT), (BT*HQ, 1), (i_t * BT, 0), (BT, BT), (1, 0)) - b_dA_local = tl.load(p_dA_local, boundary_check=(0, 1)) + p_dA_local = dA_local + (bos * HQ + i_hq) * BT + o_t[:, None] * (BT*HQ) + o_i[None, :] + b_dA_local = tl.load(p_dA_local, mask=m_T, other=0.0) # # Twb part qw part. - p_dq = tl.make_block_ptr(dq + (bos * HQ + i_hq) * K, (T, K), (K*HQ, 1), (i_t * BT, 0), (BT, BK), (1, 0)) - b_dq = tl.load(p_dq, boundary_check=(0, 1)) + p_dq = dq + (bos * HQ + i_hq) * K + o_t[:, None] * (K*HQ) + o_d[None, :] + b_dq = tl.load(p_dq, mask=m_k, other=0.0) - p_dw1 = tl.make_block_ptr(dw1 + (bos * HQ + i_hq) * K, (T, K), (K*HQ, 1), (i_t * BT, 0), (BT, BK), (1, 0)) - b_dw += tl.load(p_dw1, boundary_check=(0, 1)) + p_dw1 = dw1 + (bos * HQ + i_hq) * K + o_t[:, None] * (K*HQ) + o_d[None, :] + b_dw += tl.load(p_dw1, mask=m_k, other=0.0) b_dqw = -tl.dot(b_dA_local, tl.trans(b_Twbk)) - tl.dot(b_dq.to(b_Twb.dtype), tl.trans(b_Twb)) - p_dw2 = tl.make_block_ptr(dw2 + (bos * HQ + i_hq) * K, (T, K), (K*HQ, 1), (i_t * BT, 0), (BT, BK), (1, 0)) - b_dTwb = -tl.dot(tl.trans(b_qw), b_dq) + tl.load(p_dw2, boundary_check=(0, 1)) - b_dT += tl.dot(b_dTwb.to(b_w_beta.dtype), tl.trans(b_w_beta)) - b_dw_beta += tl.dot(tl.trans(b_T), b_dTwb.to(b_T.dtype)) + p_dw2 = dw2 + (bos * HQ + i_hq) * K + o_t[:, None] * (K*HQ) + o_d[None, :] + b_dTwb = -tl.dot(tl.trans(b_qw), b_dq) + tl.load(p_dw2, mask=m_k, other=0.0) + b_dT = tl.dot(b_dTwb.to(b_w_beta.dtype), tl.trans(b_w_beta), b_dT) + b_dw_beta = tl.dot(tl.trans(b_T), b_dTwb.to(b_T.dtype), b_dw_beta) b_dqw = tl.where(tl.arange(0, BT)[:, None] >= tl.arange(0, BT)[None, :], b_dqw, 0) b_dq += tl.dot(b_dA_local.to(b_k.dtype), b_k) b_dq += tl.dot(b_dqw.to(b_w.dtype), b_w) - b_dw += tl.dot(tl.trans(b_dqw.to(b_q.dtype)), b_q) - p_q_new = tl.make_block_ptr(dq_new + (bos * HQ + i_hq) * K, (T, K), (K*HQ, 1), (i_t * BT, 0), (BT, BK), (1, 0)) - tl.store(p_q_new, b_dq.to(dq_new.dtype.element_ty), boundary_check=(0, 1)) + b_dw = tl.dot(tl.trans(b_dqw.to(b_q.dtype)), b_q, b_dw) + p_q_new = dq_new + (bos * HQ + i_hq) * K + o_t[:, None] * (K*HQ) + o_d[None, :] + tl.store(p_q_new, b_dq.to(dq_new.dtype.element_ty), mask=m_k) # Twbk part - p_dk = tl.make_block_ptr(dk + (bos * HQ + i_hq) * K, (T, K), (K*HQ, 1), (i_t * BT, 0), (BT, BK), (1, 0)) - b_dk = tl.load(p_dk, boundary_check=(0, 1)) + p_dk = dk + (bos * HQ + i_hq) * K + o_t[:, None] * (K*HQ) + o_d[None, :] + b_dk = tl.load(p_dk, mask=m_k, other=0.0) b_dTwbk = -tl.dot(tl.trans(b_qw), b_dA_local.to(b_qw.dtype)) - tl.dot(b_w, tl.trans(b_dk.to(b_w.dtype))) b_dw -= tl.dot(b_Twbk, b_dk.to(b_w.dtype)) - b_dT += tl.dot(b_dTwbk.to(b_wbk.dtype), tl.trans(b_wbk)) + b_dT = tl.dot(b_dTwbk.to(b_wbk.dtype), tl.trans(b_wbk), b_dT) b_dwbk = tl.where(o_i[:, None] > o_i[None, :], tl.dot(tl.trans(b_T), b_dTwbk.to(b_T.dtype)), 0).to(b_w.dtype) - b_dw_beta += tl.dot(b_dwbk, b_k) + b_dw_beta = tl.dot(b_dwbk, b_k, b_dw_beta) b_dk += tl.dot(tl.trans(b_dwbk), b_w_beta) b_dk += tl.dot(tl.trans(b_dA_local), b_q) - p_dk_new = tl.make_block_ptr(dk_new + (bos * HQ + i_hq) * K, (T, K), (K*HQ, 1), (i_t * BT, 0), (BT, BK), (1, 0)) - tl.store(p_dk_new, b_dk.to(dk_new.dtype.element_ty), boundary_check=(0, 1)) + p_dk_new = dk_new + (bos * HQ + i_hq) * K + o_t[:, None] * (K*HQ) + o_d[None, :] + tl.store(p_dk_new, b_dk.to(dk_new.dtype.element_ty), mask=m_k) # matrix inverse's gradient - p_T = tl.make_block_ptr(AT + (bos * H + i_h) * BT, (BT, T), (1, BT*H), (0, i_t * BT), (BT, BT), (0, 1)) - b_Tt = tl.load(p_T, boundary_check=(0, 1)) + p_T = AT + (bos * H + i_h) * BT + o_i[:, None] + o_t[None, :] * (BT*H) + b_Tt = tl.load(p_T, mask=(o_i[:, None] < BT) & m_t[None, :], other=0.0) b_dT = tl.where(tl.arange(0, BT)[:, None] > tl.arange(0, BT)[None, :], b_dT, 0).to(b_w.dtype) b_dT = tl.dot(b_Tt, b_dT).to(b_w.dtype) b_dT = tl.dot(b_dT, b_Tt) b_dT = tl.where(tl.arange(0, BT)[:, None] > tl.arange(0, BT)[None, :], -b_dT, 0).to(b_k.dtype) - b_dw_beta += tl.dot(b_dT, b_w) - b_dw += tl.dot(tl.trans(b_dT), b_w_beta) + b_dw_beta = tl.dot(b_dT, b_w, b_dw_beta) + b_dw = tl.dot(tl.trans(b_dT), b_w_beta, b_dw) b_dw += b_dw_beta * b_beta[:, None] b_dbeta = tl.sum(b_dw_beta * b_w, axis=1) - p_dw = tl.make_block_ptr(dw + (bos * HQ + i_hq) * K, (T, K), (K*HQ, 1), (i_t * BT, 0), (BT, BK), (1, 0)) - tl.store(p_dw, b_dw.to(dw.dtype.element_ty), boundary_check=(0, 1)) - p_dbeta = tl.make_block_ptr(dbeta + (bos * HQ + i_hq), (T, ), (HQ, ), (i_t * BT, ), (BT, ), (0, )) - tl.store(p_dbeta, b_dbeta.to(dbeta.dtype.element_ty), boundary_check=(0, )) + p_dw = dw + (bos * HQ + i_hq) * K + o_t[:, None] * (K*HQ) + o_d[None, :] + tl.store(p_dw, b_dw.to(dw.dtype.element_ty), mask=m_k) + p_dbeta = dbeta + (bos * HQ + i_hq) + o_t * HQ + tl.store(p_dbeta, b_dbeta.to(dbeta.dtype.element_ty), mask=m_t) def intra_chunk_preprocess_bwd_fn(q, k, w, w2, beta, diff --git a/build/torch-cuda/ops/path_attn/intra_chunk_preprocess_bwd_prepare.py b/build/torch-cuda/ops/path_attn/intra_chunk_preprocess_bwd_prepare.py index f39242a4522290eb721ebf4b9ce0237ea99a2bb6..4571f9f8ad2f6fc18d4021b195e57f6c03cd6942 100644 --- a/build/torch-cuda/ops/path_attn/intra_chunk_preprocess_bwd_prepare.py +++ b/build/torch-cuda/ops/path_attn/intra_chunk_preprocess_bwd_prepare.py @@ -51,19 +51,19 @@ def chunk_transform_qk_bwd_kernel_prepare( USE_GATE: tl.constexpr, RETURN_H: tl.constexpr, ): - i_t, i_nh = tl.program_id(0), tl.program_id(1) + i_t, i_nh = tl.program_id(0).to(tl.int64), tl.program_id(1).to(tl.int64) i_n, i_hq = i_nh // HQ, i_nh % HQ i_h = i_hq // G if IS_VARLEN: - i_n, i_t = tl.load(indices + i_t * 2).to(tl.int32), tl.load(indices + i_t * 2 + 1).to(tl.int32) + i_n, i_t = tl.load(indices + i_t * 2).to(tl.int32), tl.load(indices + i_t * 2 + 1).to(tl.int64) bos, eos = tl.load(offsets + i_n).to(tl.int64), tl.load(offsets + i_n + 1).to(tl.int64) T = (eos - bos).to(tl.int32) - boh = tl.load(chunk_offsets + i_n).to(tl.int32) + boh = tl.load(chunk_offsets + i_n).to(tl.int64) else: bos, eos = (i_n * T).to(tl.int64), (i_n * T + T).to(tl.int64) NT = tl.cdiv(T, BT) - boh = i_n * NT + boh = (i_n * NT).to(tl.int64) sm_scale = scale * 1.44269504 # offset calculations @@ -88,19 +88,26 @@ def chunk_transform_qk_bwd_kernel_prepare( L += (bos*HQ + i_hq) D += (bos*HQ + i_hq) - p_q = tl.make_block_ptr(q, (T, K), (HQ*K, 1), (i_t * BT, 0), (BT, BK), (1, 0)) - p_k = tl.make_block_ptr(k, (K, T), (1, H*K), (0, i_t * BT), (BK, BT), (0, 1)) - p_w = tl.make_block_ptr(w, (T, K), (H*K, 1), (i_t * BT, 0), (BT, BK), (1, 0)) - p_beta = tl.make_block_ptr(beta, (T, ), (H, ), (i_t * BT, ), (BT, ), (0, )) - - b_q = tl.load(p_q, boundary_check=(0, 1)) - b_kt = tl.load(p_k, boundary_check=(0, 1)) - b_w = tl.load(p_w, boundary_check=(0, 1)) - b_beta = tl.load(p_beta, boundary_check=(0, )) - p_T = tl.make_block_ptr(AT, (T, BT), (BT*H, 1), (i_t * BT, 0), (BT, BT), (1, 0)) - b_T = tl.load(p_T, boundary_check=(0, 1)) * b_beta[None, :] - + o_t = i_t * BT + tl.arange(0, BT) o_i = tl.arange(0, BT) + o_d = tl.arange(0, BK) + o_v = tl.arange(0, BV) + m_r = o_t < T + m_k = m_r[:, None] & (o_d[None, :] < K) + m_v = m_r[:, None] & (o_v[None, :] < V) + m_T = m_r[:, None] & (o_i[None, :] < BT) + p_q = q + o_t[:, None] * (HQ*K) + o_d[None, :] + p_k = k + o_d[:, None] + o_t[None, :] * (H*K) + p_w = w + o_t[:, None] * (H*K) + o_d[None, :] + p_beta = beta + o_t * H + + b_q = tl.load(p_q, mask=m_k, other=0.0) + b_kt = tl.load(p_k, mask=(o_d[:, None] < K) & m_r[None, :], other=0.0) + b_w = tl.load(p_w, mask=m_k, other=0.0) + b_beta = tl.load(p_beta, mask=m_r, other=0.0) + p_T = AT + o_t[:, None] * (BT*H) + o_i[None, :] + b_T = tl.load(p_T, mask=m_T, other=0.0) * b_beta[None, :] + m_t = o_i[:, None] >= o_i[None, :] b_qw = tl.where(m_t, tl.dot(b_q, tl.trans(b_w.to(b_q.dtype))), 0).to(b_q.dtype) b_qwT = tl.dot(b_qw, b_T.to(b_q.dtype)).to(b_q.dtype) @@ -108,46 +115,46 @@ def chunk_transform_qk_bwd_kernel_prepare( b_A = tl.where(m_t, tl.dot(b_q, b_kt) - tl.dot(b_qwT, b_wbk), 0) b_q = b_q.to(tl.float32) - tl.dot(b_qwT, b_w.to(b_qwT.dtype)) - p_q_new = tl.make_block_ptr(q_new, (T, K), (K*HQ, 1), (i_t * BT, 0), (BT, K), (1, 0)) - tl.store(p_q_new, b_q.to(p_q_new.dtype.element_ty), boundary_check=(0, 1)) + p_q_new = q_new + o_t[:, None] * (K*HQ) + o_d[None, :] + tl.store(p_q_new, b_q.to(p_q_new.dtype.element_ty), mask=m_k) if i_hq % G == 0: b_Twb = tl.dot(b_T, b_w) # tf32 - p_h = tl.make_block_ptr(h, (T, K), (K * H, 1), (i_t * BT, 0), (BT, BK), (1, 0)) - tl.store(p_h, b_Twb.to(p_h.dtype.element_ty), boundary_check=(0, 1)) + p_h = h + o_t[:, None] * (K * H) + o_d[None, :] + tl.store(p_h, b_Twb.to(p_h.dtype.element_ty), mask=m_k) b_T_wbk = tl.dot(b_T.to(b_wbk.dtype), b_wbk).to(b_kt.dtype) - p_k_new = tl.make_block_ptr(k_new, (K, T), (1, K*H), (0, i_t * BT), (BK, BT), (0, 1)) + p_k_new = k_new + o_d[:, None] + o_t[None, :] * (K*H) tl.store(p_k_new, (b_kt - tl.dot(tl.trans(b_w.to(b_kt.dtype)), b_T_wbk) - ).to(p_k_new.dtype.element_ty), boundary_check=(0, 1)) + ).to(p_k_new.dtype.element_ty), mask=(o_d[:, None] < K) & m_r[None, :]) if USE_GATE: - p_g_cumsum = tl.make_block_ptr(g_cumsum, (T, ), (HQ, ), (i_t * BT, ), (BT, ), (0, )) - b_g_cumsum = tl.load(p_g_cumsum, boundary_check=(0, )) + p_g_cumsum = g_cumsum + o_t * HQ + b_g_cumsum = tl.load(p_g_cumsum, mask=m_r, other=0.0) b_A = b_A + (b_g_cumsum[:, None] - b_g_cumsum[None, :]) b_A = tl.where((i_t * BT + tl.arange(0, BT) < T)[:, None], b_A, float("-inf")) # avoid nan - p_l = tl.make_block_ptr(L, (T, ), (HQ, ), (i_t * BT, ), (BT, ), (0, )) - b_l = tl.load(p_l, boundary_check=(0, )) - p_delta = tl.make_block_ptr(D, (T, ), (HQ, ), (i_t * BT, ), (BT, ), (0, )) - delta = tl.load(p_delta, boundary_check=(0, )) + p_l = L + o_t * HQ + b_l = tl.load(p_l, mask=m_r, other=0.0) + p_delta = D + o_t * HQ + delta = tl.load(p_delta, mask=m_r, other=0.0) b_A_softmax = tl.exp2(tl.where(o_i[:, None] >= o_i[None, :], b_A * sm_scale - b_l[:, None], float("-inf"))) - p_do = tl.make_block_ptr(do, (T, V), (HQ*V, 1), (i_t * BT, 0), (BT, BV), (1, 0)) - b_do = tl.load(p_do, boundary_check=(0, 1)) + p_do = do + o_t[:, None] * (HQ*V) + o_v[None, :] + b_do = tl.load(p_do, mask=m_v, other=0.0) b_dv = tl.dot(tl.trans(b_A_softmax.to(b_do.dtype)), b_do) - p_dv = tl.make_block_ptr(dv, (T, V), (HQ*V, 1), (i_t * BT, 0), (BT, BV), (1, 0)) - tl.store(p_dv, b_dv.to(p_dv.dtype.element_ty), boundary_check=(0, 1)) + p_dv = dv + o_t[:, None] * (HQ*V) + o_v[None, :] + tl.store(p_dv, b_dv.to(p_dv.dtype.element_ty), mask=m_v) - p_v = tl.make_block_ptr(v, (V, T), (1, H*V), (0, i_t * BT), (BV, BT), (0, 1)) - b_v = tl.load(p_v, boundary_check=(0, 1)) + p_v = v + o_v[:, None] + o_t[None, :] * (H*V) + b_v = tl.load(p_v, mask=(o_v[:, None] < V) & m_r[None, :], other=0.0) b_dp = tl.dot(b_do, b_v) b_dA = ((b_dp - delta[:, None]) * b_A_softmax * scale) if USE_GATE: b_dgq = tl.sum(b_dA, axis=1) - tl.sum(b_dA, axis=0) - p_dg = tl.make_block_ptr(dg_cumsum, (T, ), (HQ, ), (i_t * BT, ), (BT, ), (0, )) - tl.store(p_dg, b_dgq.to(p_dg.dtype.element_ty), boundary_check=(0,)) - p_dA = tl.make_block_ptr(dA_local, (T, BT), (BT*HQ, 1), (i_t * BT, 0), (BT, BT), (1, 0)) - tl.store(p_dA, b_dA.to(p_dA.dtype.element_ty), boundary_check=(0, 1)) + p_dg = dg_cumsum + o_t * HQ + tl.store(p_dg, b_dgq.to(p_dg.dtype.element_ty), mask=m_r) + p_dA = dA_local + o_t[:, None] * (BT*HQ) + o_i[None, :] + tl.store(p_dA, b_dA.to(p_dA.dtype.element_ty), mask=m_T) def intra_chunk_preprocess_bwd_prepare_fn(q, k, v, w, beta, g_cumsum, A, L, D, do, scale, return_h=True, cu_seqlens=None, diff --git a/build/torch-cuda/ops/path_attn/intra_chunk_preprocess_fwd.py b/build/torch-cuda/ops/path_attn/intra_chunk_preprocess_fwd.py index c1f49ecf64cf23f2f3ebe9454e3daa77bf11c1d8..cde5ccba396632d1d80a2547ee34e77874d89314 100644 --- a/build/torch-cuda/ops/path_attn/intra_chunk_preprocess_fwd.py +++ b/build/torch-cuda/ops/path_attn/intra_chunk_preprocess_fwd.py @@ -46,12 +46,12 @@ def intra_chunk_preprocess_fwd_kernel( IS_VARLEN: tl.constexpr, USE_G: tl.constexpr, ): - i_t, i_nh = tl.program_id(0), tl.program_id(1) + i_t, i_nh = tl.program_id(0).to(tl.int64), tl.program_id(1).to(tl.int64) i_n, i_hq = i_nh // HQ, i_nh % HQ i_h = i_hq // G if IS_VARLEN: - i_n, i_t = tl.load(indices + i_t * 2).to(tl.int32), tl.load(indices + i_t * 2 + 1).to(tl.int32) + i_n, i_t = tl.load(indices + i_t * 2).to(tl.int32), tl.load(indices + i_t * 2 + 1).to(tl.int64) bos, eos = tl.load(offsets + i_n).to(tl.int64), tl.load(offsets + i_n + 1).to(tl.int64) T = (eos - bos).to(tl.int32) else: @@ -74,22 +74,28 @@ def intra_chunk_preprocess_fwd_kernel( L += (bos*HQ + i_hq) M += (bos*HQ + i_hq) - p_q = tl.make_block_ptr(q, (T, K), (HQ*K, 1), (i_t * BT, 0), (BT, BK), (1, 0)) - p_k = tl.make_block_ptr(k, (K, T), (1, H*K), (0, i_t * BT), (BK, BT), (0, 1)) - p_w = tl.make_block_ptr(w, (T, K), (H*K, 1), (i_t * BT, 0), (BT, BK), (1, 0)) - p_v = tl.make_block_ptr(v, (T, V), (H*V, 1), (i_t * BT, 0), (BT, BV), (1, 0)) - p_beta = tl.make_block_ptr(beta, (T, ), (H, ), (i_t * BT, ), (BT, ), (0, )) - p_T = tl.make_block_ptr(A, (T, BT), (BT*H, 1), (i_t * BT, 0), (BT, BT), (1, 0)) - - b_beta = tl.load(p_beta, boundary_check=(0, )) - b_q = tl.load(p_q, boundary_check=(0, 1)) - b_kt = tl.load(p_k, boundary_check=(0, 1)) - b_v = tl.load(p_v, boundary_check=(0, 1)) - b_w = tl.load(p_w, boundary_check=(0, 1)) - b_T = tl.load(p_T, boundary_check=(0, 1)) + o_t = i_t * BT + tl.arange(0, BT) + o_i = tl.arange(0, BT) + o_d = tl.arange(0, BK) + o_v = tl.arange(0, BV) + m_r = o_t < T + m_k = m_r[:, None] & (o_d[None, :] < K) + m_v = m_r[:, None] & (o_v[None, :] < V) + p_q = q + o_t[:, None] * (HQ*K) + o_d[None, :] + p_k = k + o_d[:, None] + o_t[None, :] * (H*K) + p_w = w + o_t[:, None] * (H*K) + o_d[None, :] + p_v = v + o_t[:, None] * (H*V) + o_v[None, :] + p_beta = beta + o_t * H + p_T = A + o_t[:, None] * (BT*H) + o_i[None, :] + + b_beta = tl.load(p_beta, mask=m_r, other=0.0) + b_q = tl.load(p_q, mask=m_k, other=0.0) + b_kt = tl.load(p_k, mask=(o_d[:, None] < K) & m_r[None, :], other=0.0) + b_v = tl.load(p_v, mask=m_v, other=0.0) + b_w = tl.load(p_w, mask=m_k, other=0.0) + b_T = tl.load(p_T, mask=m_r[:, None] & (o_i[None, :] < BT), other=0.0) b_T = b_T * b_beta[None, :] - o_i = tl.arange(0, BT) m_t = o_i[:, None] >= o_i[None, :] b_qw = tl.where(m_t, tl.dot(b_q, tl.trans(b_w.to(b_q.dtype))), 0).to(b_q.dtype) @@ -98,21 +104,21 @@ def intra_chunk_preprocess_fwd_kernel( b_A = tl.where(m_t, tl.dot(b_q, b_kt) - tl.dot(b_qwT.to(b_q.dtype), b_wbk), 0) b_q = b_q.to(tl.float32) - tl.dot(b_qwT, b_w.to(b_q.dtype)) - p_q_new = tl.make_block_ptr(q_new, (T, K), (K*HQ, 1), (i_t * BT, 0), (BT, K), (1, 0)) - tl.store(p_q_new, b_q.to(p_q_new.dtype.element_ty), boundary_check=(0, 1)) + p_q_new = q_new + o_t[:, None] * (K*HQ) + o_d[None, :] + tl.store(p_q_new, b_q.to(p_q_new.dtype.element_ty), mask=m_k) if i_hq % G == 0: b_Twb = tl.dot(b_T, b_w) - p_w2 = tl.make_block_ptr(w2, (T, K), (K*H, 1), (i_t * BT, 0), (BT, BK), (1, 0)) - tl.store(p_w2, b_Twb.to(p_w2.dtype.element_ty), boundary_check=(0, 1)) + p_w2 = w2 + o_t[:, None] * (K*H) + o_d[None, :] + tl.store(p_w2, b_Twb.to(p_w2.dtype.element_ty), mask=m_k) b_T_wbk = tl.dot(b_T.to(b_kt.dtype), b_wbk).to(b_kt.dtype) - p_k_new = tl.make_block_ptr(k_new, (K, T), (1, K*H), (0, i_t * BT), (BK, BT), (0, 1)) + p_k_new = k_new + o_d[:, None] + o_t[None, :] * (K*H) tl.store(p_k_new, (b_kt - tl.dot(tl.trans(b_w.to(b_kt.dtype)), b_T_wbk) - ).to(p_k_new.dtype.element_ty), boundary_check=(0, 1)) + ).to(p_k_new.dtype.element_ty), mask=(o_d[:, None] < K) & m_r[None, :]) if USE_G: - p_g_cumsum = tl.make_block_ptr(g_cumsum, (T, ), (HQ, ), (i_t * BT, ), (BT, ), (0, )) - b_g_cumsum = tl.load(p_g_cumsum, boundary_check=(0, )) + p_g_cumsum = g_cumsum + o_t * HQ + b_g_cumsum = tl.load(p_g_cumsum, mask=m_r, other=0.0) b_A = b_A + (b_g_cumsum[:, None] - b_g_cumsum[None, :]) b_A = tl.where((i_t * BT + tl.arange(0, BT) < T)[:, None], b_A, float("-inf")) # avoid nan @@ -121,12 +127,12 @@ def intra_chunk_preprocess_fwd_kernel( b_qkT_softmax = tl.math.exp2(b_qkT_softmax - m_i[:, None]) l_i = tl.sum(b_qkT_softmax, 1) b_o = tl.dot(b_qkT_softmax.to(b_v.dtype), b_v) - p_o = tl.make_block_ptr(o, (T, V), (V*HQ, 1), (i_t * BT, 0), (BT, BV), (1, 0)) - tl.store(p_o, b_o.to(p_o.dtype.element_ty), boundary_check=(0, 1)) - p_l = tl.make_block_ptr(L, (T, ), (HQ, ), (i_t * BT, ), (BT, ), (0, )) - p_m = tl.make_block_ptr(M, (T, ), (HQ, ), (i_t * BT, ), (BT, ), (0, )) - tl.store(p_m, m_i.to(p_m.dtype.element_ty), boundary_check=(0,)) - tl.store(p_l, l_i.to(p_l.dtype.element_ty), boundary_check=(0,)) + p_o = o + o_t[:, None] * (V*HQ) + o_v[None, :] + tl.store(p_o, b_o.to(p_o.dtype.element_ty), mask=m_v) + p_l = L + o_t * HQ + p_m = M + o_t * HQ + tl.store(p_m, m_i.to(p_m.dtype.element_ty), mask=m_r) + tl.store(p_l, l_i.to(p_l.dtype.element_ty), mask=m_r) def intra_chunk_preprocess_fwd_fn(q, k, v, w, beta, g_cumsum, A, scale, BT, cu_seqlens, diff --git a/build/torch-cuda/ops/path_attn/parallel_path_bwd_inter_dkv.py b/build/torch-cuda/ops/path_attn/parallel_path_bwd_inter_dkv.py index af070c93e13417da7292dfc94ffb4602155255d4..19c07bcc4ce481965773251d4d5ba24f6f3053c8 100644 --- a/build/torch-cuda/ops/path_attn/parallel_path_bwd_inter_dkv.py +++ b/build/torch-cuda/ops/path_attn/parallel_path_bwd_inter_dkv.py @@ -16,7 +16,7 @@ from ...ops.utils import prepare_chunk_indices, prepare_chunk_offsets 'USE_GATE': lambda args: args['g_cumsum'] is not None, 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None, }) -@triton.jit(do_not_specialize=['T']) +@triton.jit(do_not_specialize=['T', 'NUM_BLOCKS']) def parallel_path_bwd_dkv_kernel( q, k, @@ -46,21 +46,21 @@ def parallel_path_bwd_dkv_kernel( S: tl.constexpr, IS_VARLEN: tl.constexpr, USE_GATE: tl.constexpr, - NUM_BLOCKS: tl.constexpr, + NUM_BLOCKS, ): - i_t, i_bh = tl.program_id(0), tl.program_id(1) + i_t, i_bh = tl.program_id(0).to(tl.int64), tl.program_id(1).to(tl.int64) i_b, i_hq = i_bh // HQ, i_bh % HQ i_h = i_hq // G if IS_VARLEN: - i_n, i_t = tl.load(indices + i_t * 2).to(tl.int32), tl.load(indices + i_t * 2 + 1).to(tl.int32) - boh_large = tl.load(split_offsets + i_n).to(tl.int32) + i_n, i_t = tl.load(indices + i_t * 2).to(tl.int32), tl.load(indices + i_t * 2 + 1).to(tl.int64) + boh_large = tl.load(split_offsets + i_n).to(tl.int64) bos, eos = tl.load(cu_seqlens + i_n).to(tl.int64), tl.load(cu_seqlens + i_n + 1).to(tl.int64) T = (eos - bos).to(tl.int32) else: i_n = i_b bos, eos = (i_n * T).to(tl.int64), (i_n * T + T).to(tl.int64) - boh_large = i_n * tl.cdiv(T, S) + boh_large = (i_n * tl.cdiv(T, S)).to(tl.int64) # offset calculations @@ -82,15 +82,19 @@ def parallel_path_bwd_dkv_kernel( sm_scale = scale * 1.44269504 # load query - p_k = tl.make_block_ptr(k, (T, K), (H*K, 1), (i_t * BT, 0), (BT, BK), (1, 0)) - b_k = tl.load(p_k, boundary_check=(0, 1)) - p_v = tl.make_block_ptr(v, (T, V), (H*V, 1), (i_t * BT, 0), (BT, BV), (1, 0)) - b_v = tl.load(p_v, boundary_check=(0, 1)) + o_k = i_t * BT + tl.arange(0, BT) + o_d = tl.arange(0, BK) + o_v = tl.arange(0, BV) + m_k = o_k < T + p_k = k + o_k[:, None] * (H*K) + o_d[None, :] + b_k = tl.load(p_k, mask=m_k[:, None] & (o_d[None, :] < K), other=0.0) + p_v = v + o_k[:, None] * (H*V) + o_v[None, :] + b_v = tl.load(p_v, mask=m_k[:, None] & (o_v[None, :] < V), other=0.0) if USE_GATE: b_g_cumsum_k = tl.zeros([BT], dtype=tl.float32) - p_g_cumsum_k = tl.make_block_ptr(g_cumsum, (T, ), (HQ, ), (i_t * BT, ), (BT, ), (0, )) - b_g_cumsum_k += tl.load(p_g_cumsum_k, boundary_check=(0, )) + p_g_cumsum_k = g_cumsum + o_k * HQ + b_g_cumsum_k += tl.load(p_g_cumsum_k, mask=m_k, other=0.0) b_dg_cumsum_k = tl.zeros([BT], dtype=tl.float32) else: b_g_cumsum_k = None @@ -105,33 +109,34 @@ def parallel_path_bwd_dkv_kernel( last_chunk_end = tl.ceil(T / BS).to(tl.int32) * BS - BS for offset in range(last_chunk_end, last_chunk_start+S-BS, -BS): - p_delta = tl.make_block_ptr(D, (T, ), (HQ, ), (offset, ), (BS, ), (0, )) - p_l = tl.make_block_ptr(L, (T, ), (HQ, ), (offset, ), (BS, ), (0, )) - b_delta = tl.load(p_delta, boundary_check=(0, )) - b_l = tl.load(p_l, boundary_check=(0, )) - - p_q = tl.make_block_ptr(q + ((bos.to(tl.int64) * NUM_BLOCKS + idx_j) * HQ + i_hq) * K, (T, K), - (HQ*K*NUM_BLOCKS, 1), (offset, 0), (BS, BK), (1, 0)) - b_q = tl.load(p_q, boundary_check=(0, 1)) + o_q = (offset + tl.arange(0, BS)).to(tl.int64) + m_q = o_q < T + p_delta = D + o_q * HQ + p_l = L + o_q * HQ + b_delta = tl.load(p_delta, mask=m_q, other=0.0) + b_l = tl.load(p_l, mask=m_q, other=0.0) + + p_q = q + ((bos.to(tl.int64) * NUM_BLOCKS + idx_j) * HQ + i_hq) * K + o_q[:, None] * (HQ*K*NUM_BLOCKS) + o_d[None, :] + b_q = tl.load(p_q, mask=m_q[:, None] & (o_d[None, :] < K), other=0.0) b_A = tl.dot(b_k, tl.trans(b_q).to(b_k.dtype)) if USE_GATE: - p_g_cumsum_q = tl.make_block_ptr(g_cumsum, (T, ), (HQ, ), (offset, ), (BS, ), (0, )) - b_g_cumsum_q = tl.load(p_g_cumsum_q, boundary_check=(0, )) + p_g_cumsum_q = g_cumsum + o_q * HQ + b_g_cumsum_q = tl.load(p_g_cumsum_q, mask=m_q, other=0.0) b_A = b_A + b_g_cumsum_q[None, :] - b_g_cumsum_k[:, None] - b_A = tl.where((offset + tl.arange(0, BS) < T)[None, :], b_A, float("-inf")) # avoid nan + b_A = tl.where(m_q[None, :], b_A, float("-inf")) # avoid nan b_A_softmax = tl.math.exp2(b_A * sm_scale - b_l[None, :]) - p_do = tl.make_block_ptr(do, (T, V), (HQ*V, 1), (offset, 0), (BS, BV), (1, 0)) - b_do = tl.load(p_do, boundary_check=(0, 1)) - b_dv += tl.dot(b_A_softmax.to(b_do.dtype), b_do) + p_do = do + o_q[:, None] * (HQ*V) + o_v[None, :] + b_do = tl.load(p_do, mask=m_q[:, None] & (o_v[None, :] < V), other=0.0) + b_dv = tl.dot(b_A_softmax.to(b_do.dtype), b_do, b_dv) b_dp = tl.dot(b_v, tl.trans(b_do)) b_dA = ((b_dp - b_delta[None, :]) * b_A_softmax * scale) if USE_GATE: b_dg_cumsum_k -= tl.sum(b_dA, axis=1) - b_dk += tl.dot(b_dA.to(b_q.dtype), b_q) + b_dk = tl.dot(b_dA.to(b_q.dtype), b_q, b_dk) - p_dk = tl.make_block_ptr(dk, (T, K), (HQ*K, 1), (i_t * BT, 0), (BT, BK), (1, 0)) - tl.store(p_dk, b_dk.to(dk.dtype.element_ty), boundary_check=(0, 1)) + p_dk = dk + o_k[:, None] * (HQ*K) + o_d[None, :] + tl.store(p_dk, b_dk.to(dk.dtype.element_ty), mask=m_k[:, None] & (o_d[None, :] < K)) mask = i_t * BT + tl.arange(0, BT) < T tl.atomic_add( dv + (i_t * BT + tl.arange(0, BT))[:, None] * HQ * V + tl.arange(0, BV)[None, :], diff --git a/build/torch-cuda/ops/path_attn/parallel_path_bwd_inter_dqh.py b/build/torch-cuda/ops/path_attn/parallel_path_bwd_inter_dqh.py index 81c6096840c3990600ffd5ee1e677c442765e974..e2edc989b2a57a73d95bffcc3d8baedfb0e4411a 100644 --- a/build/torch-cuda/ops/path_attn/parallel_path_bwd_inter_dqh.py +++ b/build/torch-cuda/ops/path_attn/parallel_path_bwd_inter_dqh.py @@ -18,7 +18,7 @@ from ...utils import check_shared_mem 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None, 'USE_GATE': lambda args: args['g_cumsum'] is not None, }) -@triton.jit(do_not_specialize=['T']) +@triton.jit(do_not_specialize=['T', 'NUM_BLOCKS']) def parallel_path_bwd_dq_kernel( q, k, @@ -46,24 +46,26 @@ def parallel_path_bwd_dq_kernel( BK: tl.constexpr, BV: tl.constexpr, S: tl.constexpr, # aka larger chunk size - NUM_BLOCKS: tl.constexpr, + NUM_BLOCKS, IS_VARLEN: tl.constexpr, USE_GATE: tl.constexpr, ): - i_t, i_nh = tl.program_id(0), tl.program_id(1) + i_t, i_nh = tl.program_id(0).to(tl.int64), tl.program_id(1).to(tl.int64) i_n, i_hq = i_nh // HQ, i_nh % HQ i_h = i_hq // G if IS_VARLEN: - i_n, i_t = tl.load(indices + i_t * 2).to(tl.int32), tl.load(indices + i_t * 2 + 1).to(tl.int32) - boh_large = tl.load(split_offsets + i_n).to(tl.int32) + i_n, i_t = tl.load(indices + i_t * 2).to(tl.int32), tl.load(indices + i_t * 2 + 1).to(tl.int64) + boh_large = tl.load(split_offsets + i_n).to(tl.int64) bos, eos = tl.load(cu_seqlens + i_n).to(tl.int64), tl.load(cu_seqlens + i_n + 1).to(tl.int64) T = (eos - bos).to(tl.int32) else: bos, eos = (i_n * T).to(tl.int64), (i_n * T + T).to(tl.int64) - boh_large = i_n * tl.cdiv(T, S) + boh_large = (i_n * tl.cdiv(T, S)).to(tl.int64) o_t = i_t * BT + tl.arange(0, BT) m_t = o_t < T + o_d = tl.arange(0, BK) + o_v = tl.arange(0, BV) k += (bos * H + i_h) * K # GQA when H!=HQ v += (bos * H + i_h) * V # GQA when H!=HQ @@ -83,17 +85,17 @@ def parallel_path_bwd_dq_kernel( sm_scale = scale * 1.44269504 # load query - p_do = tl.make_block_ptr(do, (T, V), (HQ*V, 1), (i_t * BT, 0), (BT, BV), (1, 0)) - b_do = tl.load(p_do, boundary_check=(0, 1)) + p_do = do + o_t[:, None] * (HQ*V) + o_v[None, :] + b_do = tl.load(p_do, mask=m_t[:, None] & (o_v[None, :] < V), other=0.0) - p_l = tl.make_block_ptr(L, (T,), (HQ,), (i_t * BT,), (BT,), (0,)) - p_d = tl.make_block_ptr(D, (T,), (HQ,), (i_t * BT,), (BT,), (0,)) - b_l = tl.load(p_l, boundary_check=(0,)) - b_delta = tl.load(p_d, boundary_check=(0,)) + p_l = L + o_t * HQ + p_d = D + o_t * HQ + b_l = tl.load(p_l, mask=m_t, other=0.0) + b_delta = tl.load(p_d, mask=m_t, other=0.0) if USE_GATE: - p_g_cumsum_q = tl.make_block_ptr(g_cumsum, (T,), (HQ,), (i_t * BT,), (BT,), (0,)) - b_g_cumsum_q = tl.load(p_g_cumsum_q, boundary_check=(0,)).to(tl.float32) + p_g_cumsum_q = g_cumsum + o_t * HQ + b_g_cumsum_q = tl.load(p_g_cumsum_q, mask=m_t, other=0.0).to(tl.float32) b_dg_cumsum_q = tl.zeros([BT], dtype=tl.float32) else: b_g_cumsum_q = None @@ -104,37 +106,39 @@ def parallel_path_bwd_dq_kernel( for offset_outer in range(0, curr_end, S): idx_j = offset_outer // S - p_q = tl.make_block_ptr(q + ((bos.to(tl.int64) * NUM_BLOCKS + idx_j + 1) * HQ + i_hq) * K, (T, K), - (HQ*K*NUM_BLOCKS, 1), (i_t * BT, 0), (BT, BK), (1, 0)) - b_q = tl.load(p_q, boundary_check=(0, 1)) + p_q = q + ((bos.to(tl.int64) * NUM_BLOCKS + idx_j + 1) * HQ + i_hq) * \ + K + o_t[:, None] * (HQ*K*NUM_BLOCKS) + o_d[None, :] + b_q = tl.load(p_q, mask=m_t[:, None] & (o_d[None, :] < K), other=0.0) b_dh = -tl.dot(tl.trans(b_q), b_dq.to(b_q.dtype)) tl.atomic_add(dhc_whole + idx_j * stride_hq + tl.arange(0, K) [:, None] * K + tl.arange(0, K)[None, :], b_dh, sem='relaxed') - p_h = tl.make_block_ptr(hc_whole + idx_j * stride_h, (K, K), (K, 1), (0, 0), (BK, BK), (1, 0)) - b_h = tl.load(p_h, boundary_check=(0, 1)) + p_h = hc_whole + idx_j * stride_h + o_d[:, None] * K + o_d[None, :] + b_h = tl.load(p_h, mask=(o_d[:, None] < K) & (o_d[None, :] < K), other=0.0) b_dq = b_dq - tl.dot(b_dq.to(b_h.dtype), tl.trans(b_h)) for offset in range(offset_outer, min(offset_outer+S, i_t*BT), BS): - p_k = tl.make_block_ptr(k, (T, K), (H * K, 1), (offset, 0), (BS, BK), (1, 0)) - b_k = tl.load(p_k, boundary_check=(0, 1)) + o_k = (offset + tl.arange(0, BS)).to(tl.int64) + m_k = o_k < T + p_k = k + o_k[:, None] * (H * K) + o_d[None, :] + b_k = tl.load(p_k, mask=m_k[:, None] & (o_d[None, :] < K), other=0.0) b_A = tl.dot(b_q, tl.trans(b_k).to(b_q.dtype)) if USE_GATE: - p_g_cumsum_k = tl.make_block_ptr(g_cumsum, (T,), (HQ,), (offset,), (BS,), (0,)) - b_g_cumsum_k = tl.load(p_g_cumsum_k, boundary_check=(0,)).to(tl.float32) + p_g_cumsum_k = g_cumsum + o_k * HQ + b_g_cumsum_k = tl.load(p_g_cumsum_k, mask=m_k, other=0.0).to(tl.float32) b_A = b_A + b_g_cumsum_q[:, None] - b_g_cumsum_k[None, :] b_A = exp2(b_A * sm_scale - b_l[:, None]) b_A = tl.where(m_t[:, None], b_A, 0) - p_v = tl.make_block_ptr(v, (T, V), (H*V, 1), (offset, 0), (BS, BV), (1, 0)) - b_v = tl.load(p_v, boundary_check=(0, 1)) + p_v = v + o_k[:, None] * (H*V) + o_v[None, :] + b_v = tl.load(p_v, mask=m_k[:, None] & (o_v[None, :] < V), other=0.0) b_dp = tl.dot(b_do, tl.trans(b_v).to(b_do.dtype)) b_dA = (b_dp - b_delta[:, None]) * b_A * scale - b_dq += tl.dot(b_dA.to(b_k.dtype), b_k) + b_dq = tl.dot(b_dA.to(b_k.dtype), b_k, b_dq) if USE_GATE: b_dg_cumsum_q += tl.sum(b_dA, axis=1) - p_dq = tl.make_block_ptr(dq, (T, K), (K * HQ, 1), (i_t * BT, 0), (BT, BK), (1, 0)) - tl.store(p_dq, b_dq.to(dq.dtype.element_ty), boundary_check=(0, 1)) + p_dq = dq + o_t[:, None] * (K * HQ) + o_d[None, :] + tl.store(p_dq, b_dq.to(dq.dtype.element_ty), mask=m_t[:, None] & (o_d[None, :] < K)) if USE_GATE: tl.atomic_add(dg_cumsum + o_t * HQ, b_dg_cumsum_q, mask=m_t, sem='relaxed') diff --git a/build/torch-cuda/ops/path_attn/parallel_path_bwd_intra.py b/build/torch-cuda/ops/path_attn/parallel_path_bwd_intra.py index 314bd08730fc5df992bb3e48c62d3b8cd7508c4a..8793f0bce9ecd5db2a0da29226b0682b9140b7f4 100644 --- a/build/torch-cuda/ops/path_attn/parallel_path_bwd_intra.py +++ b/build/torch-cuda/ops/path_attn/parallel_path_bwd_intra.py @@ -28,12 +28,12 @@ def parallel_path_bwd_intra_chunk_kernel( BT: tl.constexpr, S: tl.constexpr, IS_VARLEN: tl.constexpr, USE_GATE: tl.constexpr, ): - i_t, i_bh = tl.program_id(0), tl.program_id(1) + i_t, i_bh = tl.program_id(0).to(tl.int64), tl.program_id(1).to(tl.int64) i_b, i_hq = i_bh // HQ, i_bh % HQ i_h = i_hq // G if IS_VARLEN: - i_n, i_t = tl.load(indices + i_t * 2).to(tl.int32), tl.load(indices + i_t * 2 + 1).to(tl.int32) + i_n, i_t = tl.load(indices + i_t * 2).to(tl.int32), tl.load(indices + i_t * 2 + 1).to(tl.int64) bos, eos = tl.load(offsets + i_n).to(tl.int64), tl.load(offsets + i_n + 1).to(tl.int64) T = (eos - bos).to(tl.int32) else: @@ -63,24 +63,29 @@ def parallel_path_bwd_intra_chunk_kernel( # constants sm_scale = scale * 1.44269504 - p_do = tl.make_block_ptr(do, (T, V), (HQ*V, 1), (i_t * BT, 0), (BT, BV), (1, 0)) + o_t = i_t * BT + tl.arange(0, BT) + o_d = tl.arange(0, BK) + o_v = tl.arange(0, BV) + m_t = o_t < T + m_k = m_t[:, None] & (o_d[None, :] < K) + p_do = do + o_t[:, None] * (HQ*V) + o_v[None, :] # [BT, BV] - b_do = tl.load(p_do, boundary_check=(0, 1)) + b_do = tl.load(p_do, mask=m_t[:, None] & (o_v[None, :] < V), other=0.0) - p_delta = tl.make_block_ptr(D, (T, ), (HQ, ), (i_t * BT, ), (BT, ), (0, )) - b_delta = tl.load(p_delta, boundary_check=(0, )) - p_l = tl.make_block_ptr(L, (T, ), (HQ, ), (i_t * BT, ), (BT, ), (0, )) - b_l = tl.load(p_l, boundary_check=(0, )) + p_delta = D + o_t * HQ + b_delta = tl.load(p_delta, mask=m_t, other=0.0) + p_l = L + o_t * HQ + b_l = tl.load(p_l, mask=m_t, other=0.0) b_dq = tl.zeros([BT, BK], dtype=tl.float32) - p_dq = tl.make_block_ptr(dq, (T, K), (HQ*K, 1), (i_t * BT, 0), (BT, BK), (1, 0)) - b_dq += tl.load(p_dq, boundary_check=(0, 1)) - p_q = tl.make_block_ptr(q, (T, K), (HQ*K, 1), (i_t * BT, 0), (BT, BK), (1, 0)) - b_q = tl.load(p_q, boundary_check=(0, 1)) + p_dq = dq + o_t[:, None] * (HQ*K) + o_d[None, :] + b_dq += tl.load(p_dq, mask=m_k, other=0.0) + p_q = q + o_t[:, None] * (HQ*K) + o_d[None, :] + b_q = tl.load(p_q, mask=m_k, other=0.0) if USE_GATE: - p_gq_cumsum = tl.make_block_ptr(g_cumsum, (T, ), (HQ, ), (i_t * BT, ), (BT, ), (0, )) - b_gq_cumsum = tl.load(p_gq_cumsum, boundary_check=(0, )) + p_gq_cumsum = g_cumsum + o_t * HQ + b_gq_cumsum = tl.load(p_gq_cumsum, mask=m_t, other=0.0) b_dgq = tl.zeros([BT], dtype=tl.float32) else: b_dgq = None @@ -88,23 +93,27 @@ def parallel_path_bwd_intra_chunk_kernel( curr_start = (tl.floor(i_t * BT / S).to(tl.int32) * S).to(tl.int32) for offset in range(curr_start, i_t * BT, BT): - mask = offset + tl.arange(0, BT) < T - p_k = tl.make_block_ptr(k, (T, K), (H*K, 1), (offset, 0), (BT, BK), (1, 0)) - b_k = tl.load(p_k, boundary_check=(0, 1)) + o_k = (offset + tl.arange(0, BT)).to(tl.int64) + mask = o_k < T + m_kk = mask[:, None] & (o_d[None, :] < K) + p_k = k + o_k[:, None] * (H*K) + o_d[None, :] + b_k = tl.load(p_k, mask=m_kk, other=0.0) b_q_tmp = tl.zeros([BT, BK], dtype=tl.float32) b_q_tmp += b_q for i_t_small in range(i_t * BT - BT, offset, -BT): - p_w1 = tl.make_block_ptr(w1, (T, K), (H*K, 1), (i_t_small, 0), (BT, BK), (1, 0)) - b_w1 = tl.load(p_w1, boundary_check=(0, 1)) - p_w2 = tl.make_block_ptr(w2, (T, K), (H*K, 1), (i_t_small, 0), (BT, BK), (1, 0)) - b_w2 = tl.load(p_w2, boundary_check=(0, 1)) + o_w = i_t_small + tl.arange(0, BT) + m_w = (o_w[:, None] < T) & (o_d[None, :] < K) + p_w1 = w1 + o_w[:, None] * (H*K) + o_d[None, :] + b_w1 = tl.load(p_w1, mask=m_w, other=0.0) + p_w2 = w2 + o_w[:, None] * (H*K) + o_d[None, :] + b_w2 = tl.load(p_w2, mask=m_w, other=0.0) b_A_tmp = tl.dot(b_q_tmp.to(b_w1.dtype), tl.trans(b_w1)) b_q_tmp -= tl.dot(b_A_tmp.to(b_w1.dtype), b_w2) b_q2 = b_q_tmp.to(b_k.dtype) b_A = tl.dot(b_q2, tl.trans(b_k)) if USE_GATE: - p_gk_cumsum = tl.make_block_ptr(g_cumsum, (T, ), (HQ, ), (offset, ), (BT, ), (0, )) - b_gk_cumsum = tl.load(p_gk_cumsum, boundary_check=(0, )) + p_gk_cumsum = g_cumsum + o_k * HQ + b_gk_cumsum = tl.load(p_gk_cumsum, mask=mask, other=0.0) b_A = b_A + b_gq_cumsum[:, None] - b_gk_cumsum[None, :] b_A = tl.where((i_t * BT + tl.arange(0, BT) < T)[:, None], b_A, float("-inf")) # avoid nan b_A_softmax = tl.math.exp2(b_A * sm_scale - b_l[:, None]) @@ -115,8 +124,8 @@ def parallel_path_bwd_intra_chunk_kernel( mask=mask[:, None], sem='relaxed', ) - p_v = tl.make_block_ptr(v, (T, V), (V*H, 1), (offset, 0), (BT, BV), (1, 0)) - b_v = tl.load(p_v, boundary_check=(0, 1)) + p_v = v + o_k[:, None] * (V*H) + o_v[None, :] + b_v = tl.load(p_v, mask=mask[:, None] & (o_v[None, :] < V), other=0.0) b_dp = tl.dot(b_do, tl.trans(b_v)) b_dA = ((b_dp - b_delta[:, None]) * b_A_softmax * scale) if USE_GATE: @@ -127,10 +136,10 @@ def parallel_path_bwd_intra_chunk_kernel( b_dk = tl.dot(tl.trans(b_dA), b_q2) tl.atomic_add(dk + (offset + tl.arange(0, BT))[:, None] * HQ*K + tl.arange(0, BK)[None, :], b_dk, mask=mask[:, None], sem='relaxed') - p_w1 = tl.make_block_ptr(w1, (T, K), (H*K, 1), (offset, 0), (BT, BK), (1, 0)) - b_w1 = tl.load(p_w1, boundary_check=(0, 1)) - p_w2 = tl.make_block_ptr(w2, (T, K), (H*K, 1), (offset, 0), (BT, BK), (1, 0)) - b_w2 = tl.load(p_w2, boundary_check=(0, 1)) + p_w1 = w1 + o_k[:, None] * (H*K) + o_d[None, :] + b_w1 = tl.load(p_w1, mask=m_kk, other=0.0) + p_w2 = w2 + o_k[:, None] * (H*K) + o_d[None, :] + b_w2 = tl.load(p_w2, mask=m_kk, other=0.0) b_dA2 = tl.dot(b_dq.to(b_w2.dtype), tl.trans(b_w2)).to(b_v.dtype) b_A2 = tl.dot(b_q2.to(b_w1.dtype), tl.trans(b_w1)).to(b_v.dtype) b_dw2 = -tl.dot(tl.trans(b_A2), b_dq.to(b_v.dtype)) @@ -140,10 +149,10 @@ def parallel_path_bwd_intra_chunk_kernel( tl.atomic_add(dw1 + (offset + tl.arange(0, BT))[:, None] * HQ*K + tl.arange(0, BK)[None, :], b_dw1, mask=mask[:, None], sem='relaxed') b_dq -= tl.dot(b_dA2, b_w1.to(b_v.dtype)) - b_dq += tl.dot(b_dA.to(b_k.dtype), b_k) + b_dq = tl.dot(b_dA.to(b_k.dtype), b_k, b_dq) - p_dq_new = tl.make_block_ptr(dq_new, (T, K), (HQ*K, 1), (i_t * BT, 0), (BT, BK), (1, 0)) - tl.store(p_dq_new, b_dq.to(dq_new.dtype.element_ty), boundary_check=(0, 1)) + p_dq_new = dq_new + o_t[:, None] * (HQ*K) + o_d[None, :] + tl.store(p_dq_new, b_dq.to(dq_new.dtype.element_ty), mask=m_k) mask = i_t * BT + tl.arange(0, BT) < T if USE_GATE: tl.atomic_add(dg_cumsum + (i_t * BT + tl.arange(0, BT)) * HQ, b_dgq, mask=mask, sem='relaxed') diff --git a/build/torch-cuda/ops/path_attn/parallel_path_fwd.py b/build/torch-cuda/ops/path_attn/parallel_path_fwd.py index 7fe37babba80d3a86bb64b3f9fbfc201438e7440..42aa08cee584a8d4d8106381aeed8d5f010c7dcb 100644 --- a/build/torch-cuda/ops/path_attn/parallel_path_fwd.py +++ b/build/torch-cuda/ops/path_attn/parallel_path_fwd.py @@ -45,57 +45,64 @@ def parallel_path_fwd_kernel( USE_GATE: tl.constexpr, IS_VARLEN: tl.constexpr, ): - i_t, i_bh = tl.program_id(0), tl.program_id(1) + i_t, i_bh = tl.program_id(0).to(tl.int64), tl.program_id(1).to(tl.int64) i_b, i_hq = i_bh // HQ, i_bh % HQ i_h = i_hq // G if IS_VARLEN: - i_n, i_t = tl.load(indices + i_t * 2).to(tl.int32), tl.load(indices + i_t * 2 + 1).to(tl.int32) + i_n, i_t = tl.load(indices + i_t * 2).to(tl.int32), tl.load(indices + i_t * 2 + 1).to(tl.int64) bos, eos = tl.load(cu_seqlens + i_n).to(tl.int64), tl.load(cu_seqlens + i_n + 1).to(tl.int64) T = (eos - bos).to(tl.int32) else: i_n = i_b bos, eos = (i_n * T).to(tl.int64), (i_n * T + T).to(tl.int64) - p_q = tl.make_block_ptr(q + (bos * HQ + i_hq) * K, (T, K), (HQ*K, 1), (i_t * BT, 0), (BT, BK), (1, 0)) + o_q = i_t * BT + tl.arange(0, BT) + o_d = tl.arange(0, BK) + o_v = tl.arange(0, BV) + m_q = o_q < T + m_o = m_q[:, None] & (o_v[None, :] < V) + p_q = q + (bos * HQ + i_hq) * K + o_q[:, None] * (HQ*K) + o_d[None, :] b_q = tl.zeros([BT, BK], dtype=tl.float32) - b_q += tl.load(p_q, boundary_check=(0, 1)) + b_q += tl.load(p_q, mask=m_q[:, None] & (o_d[None, :] < K), other=0.0) sm_scale = scale * 1.44269504 b_o = tl.zeros([BT, BV], dtype=tl.float32) - p_o = tl.make_block_ptr(o + (bos * HQ + i_hq) * V, (T, V), (HQ*V, 1), (i_t * BT, 0), (BT, BV), (1, 0)) - b_o += tl.load(p_o, boundary_check=(0, 1)) + p_o = o + (bos * HQ + i_hq) * V + o_q[:, None] * (HQ*V) + o_v[None, :] + b_o += tl.load(p_o, mask=m_o, other=0.0) - p_L = tl.make_block_ptr(L + bos * HQ + i_hq, (T, ), (HQ, ), (i_t * BT, ), (BT, ), (0,)) - p_M = tl.make_block_ptr(M + bos * HQ + i_hq, (T, ), (HQ, ), (i_t * BT, ), (BT, ), (0,)) - b_l = tl.load(p_L, boundary_check=(0,)) - b_m = tl.load(p_M, boundary_check=(0,)) + p_L = L + bos * HQ + i_hq + o_q * HQ + p_M = M + bos * HQ + i_hq + o_q * HQ + b_l = tl.load(p_L, mask=m_q, other=0.0) + b_m = tl.load(p_M, mask=m_q, other=0.0) if USE_GATE: - p_g_cumsum_q = tl.make_block_ptr(g_cumsum + bos * HQ + i_hq, (T, ), (HQ, ), (i_t * BT, ), (BT, ), (0,)) - b_g_cumsum_q = tl.load(p_g_cumsum_q, boundary_check=(0,)) + p_g_cumsum_q = g_cumsum + bos * HQ + i_hq + o_q * HQ + b_g_cumsum_q = tl.load(p_g_cumsum_q, mask=m_q, other=0.0) else: b_g_cumsum_q = None for offset in range((i_t + 1) * BT - 2 * BS, i_t*BT-BS, -BS): - p_k = tl.make_block_ptr(k + (bos * H + i_h) * K, (K, T), (1, K*H), (0, offset), (BK, BS), (0, 1)) # GQA when H!=HQ - p_v = tl.make_block_ptr(v + (bos * H + i_h) * V, (T, V), (V*H, 1), (offset, 0), (BS, BV), (1, 0)) # GQA when H!=HQ - p_w1 = tl.make_block_ptr(w1 + (bos * H + i_h) * K, (K, T), (1, K*H), (0, offset), (BK, BS), (0, 1)) - p_w2 = tl.make_block_ptr(w2 + (bos * H + i_h) * K, (T, K), (K*H, 1), (offset, 0), (BS, BK), (1, 0)) + o_k = offset + tl.arange(0, BS) + m_k = o_k < T + p_k = k + (bos * H + i_h) * K + o_d[:, None] + o_k[None, :] * (K*H) # GQA when H!=HQ + p_v = v + (bos * H + i_h) * V + o_k[:, None] * (V*H) + o_v[None, :] # GQA when H!=HQ + p_w1 = w1 + (bos * H + i_h) * K + o_d[:, None] + o_k[None, :] * (K*H) + p_w2 = w2 + (bos * H + i_h) * K + o_k[:, None] * (K*H) + o_d[None, :] # [BK, BS] - b_k = tl.load(p_k, boundary_check=(0, 1)) + b_k = tl.load(p_k, mask=(o_d[:, None] < K) & m_k[None, :], other=0.0) # [BS, BV] - b_v = tl.load(p_v, boundary_check=(0, 1)) + b_v = tl.load(p_v, mask=m_k[:, None] & (o_v[None, :] < V), other=0.0) # [BK, BK] - b_w1 = tl.load(p_w1, boundary_check=(0, 1)) - b_w2 = tl.load(p_w2, boundary_check=(0, 1)) + b_w1 = tl.load(p_w1, mask=(o_d[:, None] < K) & m_k[None, :], other=0.0) + b_w2 = tl.load(p_w2, mask=m_k[:, None] & (o_d[None, :] < K), other=0.0) # [BT, BS] m_s = i_t * BT + tl.arange(0, BT) >= (offset + BS) b_s = tl.dot(b_q.to(b_k.dtype), b_k) if USE_GATE: - p_g_cumsum_k = tl.make_block_ptr(g_cumsum + (bos * HQ + i_hq), (T, ), (HQ, ), (offset, ), (BS, ), (0,)) - b_g_cumsum_k = tl.load(p_g_cumsum_k, boundary_check=(0,)) + p_g_cumsum_k = g_cumsum + (bos * HQ + i_hq) + o_k * HQ + b_g_cumsum_k = tl.load(p_g_cumsum_k, mask=m_k, other=0.0) b_s = b_s + b_g_cumsum_q[:, None] - b_g_cumsum_k[None, :] b_s = tl.where(m_s[:, None], b_s * sm_scale, float("-inf")) b_m_new = tl.maximum(b_m, tl.max(b_s, 1)) @@ -104,7 +111,7 @@ def parallel_path_fwd_kernel( b_o *= alpha[:, None] b_l = b_l * alpha + tl.sum(b_s, 1) b_m = b_m_new - b_o += tl.dot(b_s.to(b_v.dtype), b_v) + b_o = tl.dot(b_s.to(b_v.dtype), b_v, b_o) b_s2 = tl.dot(b_q.to(b_w1.dtype), b_w1) b_s2 = tl.where(m_s[:, None], b_s2, 0) b_q -= tl.dot(b_s2.to(b_w2.dtype), b_w2) @@ -112,21 +119,23 @@ def parallel_path_fwd_kernel( tl.debug_barrier() for offset in range(i_t * BT - BS, -BS, -BS): - p_k = tl.make_block_ptr(k + (bos * H + i_h) * K, (K, T), (1, K*H), (0, offset), (BK, BS), (0, 1)) # GQA when H!=HQ - p_v = tl.make_block_ptr(v + (bos * H + i_h) * V, (T, V), (V*H, 1), (offset, 0), (BS, BV), (1, 0)) # GQA when H!=HQ - p_w1 = tl.make_block_ptr(w1 + (bos * H + i_h) * K, (K, T), (1, K*H), (0, offset), (BK, BS), (0, 1)) - p_w2 = tl.make_block_ptr(w2 + (bos * H + i_h) * K, (T, K), (K*H, 1), (offset, 0), (BS, BK), (1, 0)) + o_k = offset + tl.arange(0, BS) + m_k = o_k < T + p_k = k + (bos * H + i_h) * K + o_d[:, None] + o_k[None, :] * (K*H) # GQA when H!=HQ + p_v = v + (bos * H + i_h) * V + o_k[:, None] * (V*H) + o_v[None, :] # GQA when H!=HQ + p_w1 = w1 + (bos * H + i_h) * K + o_d[:, None] + o_k[None, :] * (K*H) + p_w2 = w2 + (bos * H + i_h) * K + o_k[:, None] * (K*H) + o_d[None, :] # [BK, BS] - b_k = tl.load(p_k, boundary_check=(0, 1)) + b_k = tl.load(p_k, mask=(o_d[:, None] < K) & m_k[None, :], other=0.0) # [BS, BV] - b_v = tl.load(p_v, boundary_check=(0, 1)) - b_w1 = tl.load(p_w1, boundary_check=(0, 1)) - b_w2 = tl.load(p_w2, boundary_check=(0, 1)) + b_v = tl.load(p_v, mask=m_k[:, None] & (o_v[None, :] < V), other=0.0) + b_w1 = tl.load(p_w1, mask=(o_d[:, None] < K) & m_k[None, :], other=0.0) + b_w2 = tl.load(p_w2, mask=m_k[:, None] & (o_d[None, :] < K), other=0.0) # [BT, BS] b_s = tl.dot(b_q.to(b_k.dtype), b_k) if USE_GATE: - p_g_cumsum_k = tl.make_block_ptr(g_cumsum + (bos * HQ + i_hq), (T, ), (HQ, ), (offset, ), (BS, ), (0,)) - b_g_cumsum_k = tl.load(p_g_cumsum_k, boundary_check=(0,)) + p_g_cumsum_k = g_cumsum + (bos * HQ + i_hq) + o_k * HQ + b_g_cumsum_k = tl.load(p_g_cumsum_k, mask=m_k, other=0.0) b_s = b_s + b_g_cumsum_q[:, None] - b_g_cumsum_k[None, :] b_s = b_s * sm_scale b_m_new = tl.maximum(b_m, tl.max(b_s, 1)) @@ -135,16 +144,16 @@ def parallel_path_fwd_kernel( b_o *= alpha[:, None] b_l = b_l * alpha + tl.sum(b_s, 1) b_m = b_m_new - b_o += tl.dot(b_s.to(b_v.dtype), b_v) + b_o = tl.dot(b_s.to(b_v.dtype), b_v, b_o) b_s2 = tl.dot(b_q.to(b_w1.dtype), b_w1) b_q -= tl.dot(b_s2.to(b_w2.dtype), b_w2) b_o = b_o / b_l[:, None] - p_o_new = tl.make_block_ptr(o_new + (bos * HQ + i_hq) * V, (T, V), (HQ*V, 1), (i_t*BT, 0), (BT, BV), (1, 0)) - tl.store(p_o_new, b_o.to(p_o_new.dtype.element_ty), boundary_check=(0, 1)) + p_o_new = o_new + (bos * HQ + i_hq) * V + o_q[:, None] * (HQ*V) + o_v[None, :] + tl.store(p_o_new, b_o.to(p_o_new.dtype.element_ty), mask=m_o) b_l = tl.math.log2(b_l) + b_m - p_L_new = tl.make_block_ptr(L_new + (bos * HQ + i_hq), (T, ), (HQ, ), (i_t * BT, ), (BT, ), (0,)) - tl.store(p_L_new, b_l.to(p_L_new.dtype.element_ty), boundary_check=(0,)) + p_L_new = L_new + (bos * HQ + i_hq) + o_q * HQ + tl.store(p_L_new, b_l.to(p_L_new.dtype.element_ty), mask=m_q) def parallel_path_fwd_fn( diff --git a/build/torch-cuda/ops/path_attn/prepare_k_cache.py b/build/torch-cuda/ops/path_attn/prepare_k_cache.py index 2fe4d40c191af498ce0e56a321e654040b27f442..41edb0c3431ed77c1e1f075566fc17c11d848a13 100644 --- a/build/torch-cuda/ops/path_attn/prepare_k_cache.py +++ b/build/torch-cuda/ops/path_attn/prepare_k_cache.py @@ -25,11 +25,11 @@ def parallel_path_fwd_kernel_prepare_k_cache( BT: tl.constexpr, BK: tl.constexpr, IS_VARLEN: tl.constexpr, ): - i_t, i_bh = tl.program_id(0), tl.program_id(1) + i_t, i_bh = tl.program_id(0).to(tl.int64), tl.program_id(1).to(tl.int64) i_b, i_h = i_bh // H, i_bh % H if IS_VARLEN: - i_n, i_t = tl.load(indices + i_t * 2).to(tl.int32), tl.load(indices + i_t * 2 + 1).to(tl.int32) + i_n, i_t = tl.load(indices + i_t * 2).to(tl.int32), tl.load(indices + i_t * 2 + 1).to(tl.int64) bos, eos = tl.load(offsets + i_n).to(tl.int64), tl.load(offsets + i_n + 1).to(tl.int64) T = (eos - bos).to(tl.int32) else: @@ -41,19 +41,24 @@ def parallel_path_fwd_kernel_prepare_k_cache( w1 += (bos * H + i_h) * K w2 += (bos * H + i_h) * K # constants - p_k = tl.make_block_ptr(k, (T, K), (H*K, 1), (i_t * BT, 0), (BT, BK), (1, 0)) + o_k = i_t * BT + tl.arange(0, BT) + o_d = tl.arange(0, BK) + m_k = (o_k[:, None] < T) & (o_d[None, :] < K) + p_k = k + o_k[:, None] * (H*K) + o_d[None, :] b_k = tl.zeros([BT, BK], dtype=tl.float32) - b_k += tl.load(p_k, boundary_check=(0, 1)) + b_k += tl.load(p_k, mask=m_k, other=0.0) for k_block_idx in range(i_t + 1, tl.cdiv(T, BT)): - p_w1 = tl.make_block_ptr(w1, (T, K), (H*K, 1), (k_block_idx * BT, 0), (BT, BK), (1, 0)) - p_w2 = tl.make_block_ptr(w2, (T, K), (H*K, 1), (k_block_idx * BT, 0), (BT, BK), (1, 0)) - b_w1 = tl.load(p_w1, boundary_check=(0, 1)) - b_w2 = tl.load(p_w2, boundary_check=(0, 1)) + o_w = k_block_idx * BT + tl.arange(0, BT) + m_w = (o_w[:, None] < T) & (o_d[None, :] < K) + p_w1 = w1 + o_w[:, None] * (H*K) + o_d[None, :] + p_w2 = w2 + o_w[:, None] * (H*K) + o_d[None, :] + b_w1 = tl.load(p_w1, mask=m_w, other=0.0) + b_w2 = tl.load(p_w2, mask=m_w, other=0.0) b_A = tl.dot(b_k.to(b_w2.dtype), tl.trans(b_w2)) b_k = b_k - tl.dot(b_A.to(b_w1.dtype), b_w1) - p_k_new = tl.make_block_ptr(k_new, (T, K), (H*K, 1), (i_t * BT, 0), (BT, BK), (1, 0)) - tl.store(p_k_new, b_k.to(p_k_new.dtype.element_ty), boundary_check=(0, 1)) + p_k_new = k_new + o_k[:, None] * (H*K) + o_d[None, :] + tl.store(p_k_new, b_k.to(p_k_new.dtype.element_ty), mask=m_k) def prepare_k_cache_fn(k, w1, w2, cu_seqlens, BS, use_cache=False, chunk_indices: torch.LongTensor | None = None): diff --git a/build/torch-cuda/ops/path_attn/transform_q.py b/build/torch-cuda/ops/path_attn/transform_q.py index 303d7a6322168f377daee064b852d5acea588801..30cd3cd9dd40cc757330cb11f50044a80c976124 100644 --- a/build/torch-cuda/ops/path_attn/transform_q.py +++ b/build/torch-cuda/ops/path_attn/transform_q.py @@ -15,7 +15,7 @@ from ...ops.utils import get_max_num_splits, prepare_chunk_indices @triton.heuristics({ 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None, }) -@triton.jit(do_not_specialize=['T']) +@triton.jit(do_not_specialize=['T', 'NUM_BLOCKS']) def transform_q_fwd_kernel( q, q_new, @@ -32,45 +32,50 @@ def transform_q_fwd_kernel( BT: tl.constexpr, BS: tl.constexpr, BK: tl.constexpr, - NUM_BLOCKS: tl.constexpr, + NUM_BLOCKS, IS_VARLEN: tl.constexpr, ): - i_t, i_bh = tl.program_id(0), tl.program_id(1) + i_t, i_bh = tl.program_id(0).to(tl.int64), tl.program_id(1).to(tl.int64) i_b, i_hq = i_bh // HQ, i_bh % HQ i_h = i_hq // G if IS_VARLEN: - i_n, i_t = tl.load(indices + i_t * 2).to(tl.int32), tl.load(indices + i_t * 2 + 1).to(tl.int32) + i_n, i_t = tl.load(indices + i_t * 2).to(tl.int32), tl.load(indices + i_t * 2 + 1).to(tl.int64) bos, eos = tl.load(cu_seqlens + i_n).to(tl.int64), tl.load(cu_seqlens + i_n + 1).to(tl.int64) T = (eos - bos).to(tl.int32) else: i_n = i_b bos, eos = (i_n * T).to(tl.int64), (i_n * T + T).to(tl.int64) # boh = i_n * tl.cdiv(T, BS) - p_q = tl.make_block_ptr(q + (bos * HQ + i_hq) * K, (T, K), (HQ*K, 1), (i_t * BT, 0), (BT, BK), (1, 0)) + o_q = i_t * BT + tl.arange(0, BT) + o_d = tl.arange(0, BK) + m_q = (o_q[:, None] < T) & (o_d[None, :] < K) + p_q = q + (bos * HQ + i_hq) * K + o_q[:, None] * (HQ*K) + o_d[None, :] b_q = tl.zeros([BT, BK], dtype=tl.float32) - b_q += tl.load(p_q, boundary_check=(0, 1)) + b_q += tl.load(p_q, mask=m_q, other=0.0) if BS == BT: if (i_t * BT) % S == 0: - p_q_new = tl.make_block_ptr(q_new + ((bos.to(tl.int64) * NUM_BLOCKS + (i_t * BT // S)) * HQ + i_hq) * K, - (T, K), (HQ*K*NUM_BLOCKS, 1), (i_t * BT, 0), (BT, BK), (1, 0)) - tl.store(p_q_new, b_q.to(q_new.dtype.element_ty), boundary_check=(0, 1)) + p_q_new = q_new + ((bos.to(tl.int64) * NUM_BLOCKS + (i_t * BT // S)) * HQ + i_hq) * \ + K + o_q[:, None] * (HQ*K*NUM_BLOCKS) + o_d[None, :] + tl.store(p_q_new, b_q.to(q_new.dtype.element_ty), mask=m_q) for offset in range((i_t + 1) * BT - 2 * BS, S-BS, -BS): - p_w1 = tl.make_block_ptr(w1 + (bos * H + i_h) * K, (K, T), (1, K*H), (0, offset), (BK, BS), (0, 1)) - p_w2 = tl.make_block_ptr(w2 + (bos * H + i_h) * K, (T, K), (K*H, 1), (offset, 0), (BS, BK), (1, 0)) - b_w1 = tl.load(p_w1, boundary_check=(0, 1)) - b_w2 = tl.load(p_w2, boundary_check=(0, 1)) - m_s = i_t * BT + tl.arange(0, BT) >= (offset + BS) + o_w = offset + tl.arange(0, BS) + m_w = o_w < T + p_w1 = w1 + (bos * H + i_h) * K + o_d[:, None] + o_w[None, :] * (K*H) + p_w2 = w2 + (bos * H + i_h) * K + o_w[:, None] * (K*H) + o_d[None, :] + b_w1 = tl.load(p_w1, mask=(o_d[:, None] < K) & m_w[None, :], other=0.0) + b_w2 = tl.load(p_w2, mask=m_w[:, None] & (o_d[None, :] < K), other=0.0) + m_s = o_q >= (offset + BS) b_s2 = tl.dot(b_q.to(b_w1.dtype), b_w1) b_s2 = tl.where(m_s[:, None], b_s2, 0) b_q -= tl.dot(b_s2.to(b_w2.dtype), b_w2) if offset % S == 0: - p_q_new = tl.make_block_ptr(q_new + ((bos.to(tl.int64) * NUM_BLOCKS + (offset // S)) * HQ + i_hq) * K, - (T, K), (HQ*K*NUM_BLOCKS, 1), (i_t * BT, 0), (BT, BK), (1, 0)) - tl.store(p_q_new, b_q.to(q_new.dtype.element_ty), boundary_check=(0, 1)) + p_q_new = q_new + ((bos.to(tl.int64) * NUM_BLOCKS + (offset // S)) * HQ + i_hq) * \ + K + o_q[:, None] * (HQ*K*NUM_BLOCKS) + o_d[None, :] + tl.store(p_q_new, b_q.to(q_new.dtype.element_ty), mask=m_q) def transform_q_fwd_fn( diff --git a/build/torch-cuda/ops/precond_gated_delta_rule/__init__.py b/build/torch-cuda/ops/precond_gated_delta_rule/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..b446dc118902a64c6441266e126112b942b77ebf --- /dev/null +++ b/build/torch-cuda/ops/precond_gated_delta_rule/__init__.py @@ -0,0 +1,16 @@ +# Copyright (c) 2023-2026, Songlin Yang, Yu Zhang, Zhiyuan Li +# +# This source code is licensed under the MIT license found in the +# LICENSE file in the root directory of this source tree. +# For a list of all contributors, visit: +# https://github.com/fla-org/flash-linear-attention/graphs/contributors + +from .chunk import chunk_precond_gated_delta_rule +from .fused_recurrent import fused_recurrent_precond_gated_delta_rule +from .naive import naive_recurrent_precond_gated_delta_rule + +__all__ = [ + "chunk_precond_gated_delta_rule", + "fused_recurrent_precond_gated_delta_rule", + "naive_recurrent_precond_gated_delta_rule", +] diff --git a/build/torch-cuda/ops/precond_gated_delta_rule/backends/__init__.py b/build/torch-cuda/ops/precond_gated_delta_rule/backends/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..2fcad6ad8aa5f803cf2602f421eac8afb9d4148f --- /dev/null +++ b/build/torch-cuda/ops/precond_gated_delta_rule/backends/__init__.py @@ -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 + +"""GDR backends.""" + + +__all__: list[str] = [] diff --git a/build/torch-cuda/ops/precond_gated_delta_rule/chunk.py b/build/torch-cuda/ops/precond_gated_delta_rule/chunk.py new file mode 100644 index 0000000000000000000000000000000000000000..ae92bbf8c0ed7a0960168eea69aa5297d0a65044 --- /dev/null +++ b/build/torch-cuda/ops/precond_gated_delta_rule/chunk.py @@ -0,0 +1,567 @@ +# 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 + +from ...modules.norm.l2norm import l2norm_bwd, l2norm_fwd +from ...ops.atk.chunk_atk_bwd import chunk_atk_bwd +from ...ops.atk.chunk_atk_fwd import chunk_atk_fwd, recompute_atk_fwd +from ...ops.common.chunk_delta_h import chunk_gated_delta_rule_bwd_dhu, chunk_gated_delta_rule_fwd_h +from ...ops.common.chunk_o import chunk_bwd_dqkwg, chunk_bwd_dv_local, chunk_fwd_o +from ...ops.cp import FLACPContext +from ...ops.gated_delta_rule.wy_fast import recompute_w_u_fwd +from ...ops.precond_gated_delta_rule.chunk_precond_kkt_fwd import chunk_precond_kkt_fwd +from ...ops.precond_gated_delta_rule.chunk_precond_wy_bwd import prepare_precond_wy_repr_bwd +from ...ops.utils import chunk_local_cumsum, solve_tril +from ...ops.utils.constant import RCP_LN2 +from ...ops.utils.index import prepare_chunk_indices +from ...utils import autocast_custom_bwd, autocast_custom_fwd, input_guard + + +def chunk_precond_gated_delta_rule_fwd( + q: torch.Tensor, + k_read: torch.Tensor, + k_write: torch.Tensor, + v: torch.Tensor, + g_atk: torch.Tensor, + g: torch.Tensor, + beta_atk: torch.Tensor, + beta: torch.Tensor, + scale: float, + initial_state: torch.Tensor, + initial_A_state: torch.Tensor, + output_final_state: bool, + cu_seqlens: torch.LongTensor | None = None, + cp_context: FLACPContext | None = None, + chunk_indices: torch.LongTensor | None = None, + x: float = 1.5, + eps: float = 1e-6, + log_atk_scale: torch.Tensor = None, + transpose_state_layout: bool = False, +): + if chunk_indices is None and cu_seqlens is not None: + chunk_indices = prepare_chunk_indices(cu_seqlens, 64) + + k_precond, at = chunk_atk_fwd( + k=k_write, + beta=beta_atk, + log_g=g_atk, + chunk_size=64, + initial_A_state=initial_A_state, + output_final_state=output_final_state, + cu_seqlens=cu_seqlens, + x=x, + eps=eps, + log_atk_scale=log_atk_scale, + ) + + g = chunk_local_cumsum( + g, + chunk_size=64, + scale=RCP_LN2, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + ) + + # Asymmetric KKT: uses k_read for correction, k_precond for write + A = chunk_precond_kkt_fwd( + k=k_read, + k_precond=k_precond, + g=g, + beta=beta, + chunk_size=64, + output_dtype=torch.float32, + cu_seqlens=cu_seqlens, + ) + A = solve_tril( + A=A, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + output_dtype=k_read.dtype, + ) + # obtain WY representation. u is actually the new v. + w, u = recompute_w_u_fwd( + k=k_read, + v=v, + beta=beta, + A=A, + g=g, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + ) + + h, v_new, final_state = chunk_gated_delta_rule_fwd_h( + k=k_precond, + w=w, + u=u, + g=g, + initial_state=initial_state, + output_final_state=output_final_state, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + state_v_first=transpose_state_layout, + ) + + o = chunk_fwd_o( + q=q, + k=k_precond, + v=v_new, + h=h, + g=g, + scale=scale, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + state_v_first=transpose_state_layout, + ) + + return g, o, A, final_state, at + + +def chunk_precond_gated_delta_rule_bwd( + q: torch.Tensor, + k_read: torch.Tensor, + k_write: torch.Tensor, + v: torch.Tensor, + g_atk: torch.Tensor, + g: torch.Tensor, + beta_atk: torch.Tensor, + beta: torch.Tensor, + A: torch.Tensor, + scale: float, + initial_state: torch.Tensor, + initial_A_state: torch.Tensor, + do: torch.Tensor, + dht: torch.Tensor, + cu_seqlens: torch.LongTensor | None = None, + cp_context: FLACPContext | None = None, + chunk_indices: torch.LongTensor | None = None, + x: float = 1.5, + eps: float = 1e-6, + log_atk_scale: torch.Tensor = None, + transpose_state_layout: bool = False, + dat: torch.Tensor | None = None, +): + if chunk_indices is None and cu_seqlens is not None: + chunk_indices = prepare_chunk_indices(cu_seqlens, 64) + + # Recompute ATK forward intermediates + k_precond, ac_atk, a_atk, sa_atk = recompute_atk_fwd( + k=k_write, + beta=beta_atk, + log_g=g_atk, + chunk_size=64, + initial_A_state=initial_A_state, + cu_seqlens=cu_seqlens, + x=x, + eps=eps, + log_atk_scale=log_atk_scale, + ) + + w, u = recompute_w_u_fwd( + k=k_read, + v=v, + beta=beta, + A=A, + g=g, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + ) + + h, v_new, _ = chunk_gated_delta_rule_fwd_h( + k=k_precond, + w=w, + u=u, + g=g, + initial_state=initial_state, + output_final_state=False, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + state_v_first=transpose_state_layout, + ) + dv = chunk_bwd_dv_local( + q=q, + k=k_precond, + g=g, + do=do, + scale=scale, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + ) + + dh, dh0, dv = chunk_gated_delta_rule_bwd_dhu( + q=q, + k=k_precond, + w=w, + g=g, + h0=initial_state, + dht=dht, + do=do, + dv=dv, + scale=scale, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + state_v_first=transpose_state_layout, + ) + dq, dk_precond, dw, dg = chunk_bwd_dqkwg( + q=q, + k=k_precond, + v=v_new, + w=w, + g=g, + h=h, + dv=dv, + do=do, + dh=dh, + scale=scale, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + state_v_first=transpose_state_layout, + ) + # Asymmetric WY backward: produces separate grads for k_read and k_precond + dk_read_wy, dk_precond_wy, dv, dbeta_wy, dg2 = prepare_precond_wy_repr_bwd( + k=k_read, + k_precond=k_precond, + v=v, + beta=beta, + g=g, + A=A, + dw=dw, + du=dv, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + ) + + dbeta = dbeta_wy + dk_precond.add_(dk_precond_wy) + + # ATK backward (uses precomputed forward intermediates) + dk_write_atk, dbeta_atk, dg_atk, d_log_atk_scale, dh0_atk = chunk_atk_bwd( + k=k_write, + g_raw=g_atk, + beta=beta_atk, + dk_precond=dk_precond, + ac=ac_atk, + a=a_atk, + sa=sa_atk, + initial_A_state=initial_A_state, + cu_seqlens=cu_seqlens, + x=x, + eps=eps, + log_atk_scale=log_atk_scale, + dat=dat, + ) + + dk_read = dk_read_wy + dk_write = dk_write_atk + + dg.add_(dg2) + dg = chunk_local_cumsum(dg, chunk_size=64, reverse=True, cu_seqlens=cu_seqlens, chunk_indices=chunk_indices) + + return dq, dk_read, dk_write, dv, dbeta_atk, dbeta, dg_atk, dg, dh0, dh0_atk, d_log_atk_scale + + +class ChunkPrecondGatedDeltaRuleFunction(torch.autograd.Function): + + @staticmethod + @input_guard + @autocast_custom_fwd + def forward( + ctx, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + g_atk: torch.Tensor, + g: torch.Tensor, + beta_atk: torch.Tensor, + beta: torch.Tensor, + scale: float, + initial_state: torch.Tensor, + initial_A_state: torch.Tensor, + output_final_state: bool, + use_qk_l2norm_in_kernel: bool = True, + cu_seqlens: torch.LongTensor | None = None, + cu_seqlens_cpu: torch.LongTensor | None = None, + cp_context: FLACPContext | None = None, + transpose_state_layout: bool = False, + x: float = 1.5, + eps: float = 1e-6, + log_atk_scale: torch.Tensor = None, + ): + q_rstd, k_rstd = None, None + if use_qk_l2norm_in_kernel: + q, q_rstd = l2norm_fwd(q) + k, k_rstd = l2norm_fwd(k) + + k_read = k + k_write = k + + H = k.shape[2] + log_atk_scale_was_tensor = isinstance(log_atk_scale, torch.Tensor) + + if log_atk_scale is None: + log_atk_scale = torch.full((H,), -0.2, dtype=torch.float32, device=k.device) + + chunk_indices = prepare_chunk_indices( + cu_seqlens, 64, cu_seqlens_cpu=cu_seqlens_cpu) if cu_seqlens is not None else None + + g, o, A, h_final, at = chunk_precond_gated_delta_rule_fwd( + q=q, + k_read=k_read, + k_write=k_write, + v=v, + g_atk=g_atk, + g=g, + beta_atk=beta_atk, + beta=beta, + scale=scale, + initial_state=initial_state, + initial_A_state=initial_A_state, + output_final_state=output_final_state, + cu_seqlens=cu_seqlens, + cp_context=cp_context, + chunk_indices=chunk_indices, + x=x, + eps=eps, + log_atk_scale=log_atk_scale, + transpose_state_layout=transpose_state_layout, + ) + + ctx.save_for_backward( + q, q_rstd, k, k_rstd, v, + g_atk, g, + beta_atk, beta, + A, initial_state, cu_seqlens, + log_atk_scale, chunk_indices + ) + ctx.scale = scale + ctx.use_qk_l2norm_in_kernel = use_qk_l2norm_in_kernel + ctx.initial_A_state = initial_A_state + ctx.cp_context = cp_context + ctx.transpose_state_layout = transpose_state_layout + ctx.x = x + ctx.eps = eps + ctx.log_atk_scale_was_tensor = log_atk_scale_was_tensor + + return o.to(q.dtype), h_final, at + + @staticmethod + @input_guard + @autocast_custom_bwd + def backward(ctx, do, dht, dat): + (q, q_rstd, k, k_rstd, v, + g_atk, g, + beta_atk, beta, + A, initial_state, cu_seqlens, + log_atk_scale, chunk_indices) = ctx.saved_tensors + + k_read = k + k_write = k + + dq, dk_read, dk_write, dv, dbeta_atk, dbeta, dg_atk, dg, dh0, dh0_atk, d_log_atk_scale = chunk_precond_gated_delta_rule_bwd( + q=q, + k_read=k_read, + k_write=k_write, + v=v, + g_atk=g_atk, + g=g, + beta_atk=beta_atk, + beta=beta, + A=A, + scale=ctx.scale, + initial_state=initial_state, + initial_A_state=ctx.initial_A_state, + do=do, + dht=dht, + cu_seqlens=cu_seqlens, + cp_context=ctx.cp_context, + chunk_indices=chunk_indices, + x=ctx.x, + eps=ctx.eps, + log_atk_scale=log_atk_scale, + transpose_state_layout=ctx.transpose_state_layout, + dat=dat, + ) + + dk_read.add_(dk_write) + dk = dk_read + + if ctx.use_qk_l2norm_in_kernel: + dq = l2norm_bwd(q, q_rstd, dq) + dk = l2norm_bwd(k, k_rstd, dk) + + d_log_atk_scale_out = d_log_atk_scale.to(log_atk_scale) if ctx.log_atk_scale_was_tensor else None + return ( + dq.to(q), dk.to(k), dv.to(v), + dg_atk.to(g_atk), dg.to(g), + dbeta_atk.to(beta_atk), dbeta.to(beta), + None, dh0, dh0_atk, None, # scale, initial_state, initial_A_state, output_final_state + None, None, None, None, None, # use_qk_l2norm_in_kernel, cu_seqlens, cu_seqlens_cpu, cp_context, transpose_state_layout + None, None, # x, eps + d_log_atk_scale_out, # log_atk_scale + ) + + +@torch.compiler.disable +def chunk_precond_gated_delta_rule( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + g_atk: torch.Tensor, + g: torch.Tensor, + beta_atk: torch.Tensor, + beta: torch.Tensor, + scale: float = None, + initial_state: torch.Tensor = None, + initial_A_state: torch.Tensor = None, + output_final_state: bool = False, + use_qk_l2norm_in_kernel: bool = True, + cu_seqlens: torch.LongTensor | None = None, + cu_seqlens_cpu: torch.LongTensor | None = None, + cp_context: FLACPContext | None = None, + transpose_state_layout: bool = False, + x: float = 1.5, + eps: float = 1e-6, + log_atk_scale: torch.Tensor = None, +): + r""" + Args: + q (torch.Tensor): + queries of shape `[B, T, H, K]`. + k (torch.Tensor): + keys of shape `[B, T, H, K]`. + v (torch.Tensor): + values of shape `[B, T, HV, V]`. + GVA (Grouped Value Attention) is applied if `HV > H`, where `HV` must be divisible by `H`. + g_atk (torch.Tensor): + ATK gates (decays) in log space of shape `[B, T, H]`. + g (torch.Tensor): + gates (decays) in log space of shape `[B, T, HV]`. + beta_atk (torch.Tensor): + ATK betas of shape `[B, T, H]`. + beta (torch.Tensor): + betas of shape `[B, T, HV]`. + scale (Optional[float]): + Scale factor for attention scores. + If not provided, it will default to `1 / sqrt(K)`. Default: `None`. + initial_state (Optional[torch.Tensor]): + Initial recurrent state of shape `[N, HV, K, V]` for `N` input sequences. + For equal-length input sequences, `N` equals the batch size `B`. + Default: `None`. + initial_A_state (Optional[torch.Tensor]): + Initial ATK diagonal state of shape `[N, H, K]`. Default: `None`. + output_final_state (Optional[bool]): + Whether to output the final state of shape `[N, HV, K, V]`. Default: `False`. + use_qk_l2norm_in_kernel (bool): + Whether to apply L2norm to the q/k tensor internally. Default: `True`. + cu_seqlens (torch.LongTensor): + Cumulative sequence lengths of shape `[N+1]` used for variable-length training, + consistent with the FlashAttention API. + cp_context (Optional[FLACPContext]): + Context parallelism is not yet supported for the preconditioned path; + passing a context raises `NotImplementedError`. Default: `None`. + transpose_state_layout (Optional[bool]): + Whether to use the transposed state layout for the hidden state. + Default: `False`. + x (float): + Squash range parameter. Default: `1.5`. + eps (float): + Epsilon for numerical stability in the squash function. Default: `1e-6`. + log_atk_scale (Optional[torch.Tensor]): + Per-head log-space center of shape `[H]`. Default: `None`. + + Returns: + o (torch.Tensor): + Outputs of shape `[B, T, HV, V]`. + recurrent_state (torch.Tensor): + Final recurrent state of shape `[N, HV, K, V]` if `output_final_state=True` else `None`. + A_state (torch.Tensor): + Final ATK diagonal state of shape `[N, H, K]` if `output_final_state=True` else `None`. + + Examples:: + >>> import torch + >>> import torch.nn.functional as F + >>> from einops import rearrange + >>> from fla.ops.precond_gated_delta_rule import chunk_precond_gated_delta_rule + # inputs with equal lengths + >>> B, T, H, K, V = 4, 2048, 4, 512, 512 + >>> q = torch.randn(B, T, H, K, dtype=torch.bfloat16, device='cuda') + >>> k = F.normalize(torch.randn(B, T, H, K, dtype=torch.bfloat16, device='cuda'), p=2, dim=-1) + >>> v = torch.randn(B, T, H, V, dtype=torch.bfloat16, device='cuda') + >>> beta_atk = torch.rand(B, T, H, dtype=torch.bfloat16, device='cuda').sigmoid() + >>> beta = torch.rand(B, T, H, dtype=torch.bfloat16, device='cuda').sigmoid() + >>> g_atk = F.logsigmoid(torch.rand(B, T, H, dtype=torch.bfloat16, device='cuda')) + >>> g = F.logsigmoid(torch.rand(B, T, H, dtype=torch.bfloat16, device='cuda')) + >>> h0 = torch.randn(B, H, K, V, dtype=torch.bfloat16, device='cuda') + >>> o, ht, at = chunk_precond_gated_delta_rule( + q, k, v, g_atk, g, beta_atk, beta, + initial_state=h0, + output_final_state=True + ) + # for variable-length inputs, the batch size `B` is expected to be 1 and `cu_seqlens` is required + >>> q, k, v, beta_atk, beta, g_atk, g = map(lambda x: rearrange(x, 'b t ... -> 1 (b t) ...'), (q, k, v, beta_atk, beta, g_atk, g)) + # for a batch with 4 sequences, `cu_seqlens` with 5 start/end positions are expected + >>> cu_seqlens = q.new_tensor([0, 2048, 4096, 6144, 8192], dtype=torch.long) + >>> o, ht, at = chunk_precond_gated_delta_rule( + q, k, v, g_atk, g, beta_atk, beta, + initial_state=h0, + output_final_state=True, + cu_seqlens=cu_seqlens + ) + """ + # Validate head dimensions + if q.shape[2] != k.shape[2]: + raise ValueError( + f"q and k must have the same number of heads, " + f"but got q.shape[2]={q.shape[2]} and k.shape[2]={k.shape[2]}" + ) + H, HV = q.shape[2], v.shape[2] + if HV % H != 0: + raise ValueError( + f"For GVA, num_v_heads (HV={HV}) must be evenly divisible by " + f"num_heads (H={H}), but got HV % H = {HV % H}" + ) + + if cp_context is not None: + raise NotImplementedError( + "Context parallelism is not yet supported for chunk_precond_gated_delta_rule: " + "the preconditioned path lacks the CP state compression/expansion plumbing. " + "Please run without `cp_context`." + ) + + if cu_seqlens is not None: + if q.shape[0] != 1: + raise ValueError( + f"The batch size is expected to be 1 rather than {q.shape[0]} when using `cu_seqlens`." + f"Please flatten variable-length inputs before processing.", + ) + if initial_state is not None and initial_state.shape[0] != len(cu_seqlens) - 1: + raise ValueError( + f"The number of initial states is expected to be equal to the number of input sequences, " + f"i.e., {len(cu_seqlens) - 1} rather than {initial_state.shape[0]}.", + ) + if scale is None: + scale = k.shape[-1] ** -0.5 + + o, h_final, a_final = ChunkPrecondGatedDeltaRuleFunction.apply( + q, k, v, + g_atk, g, + beta_atk, beta, + scale, + initial_state, + initial_A_state, + output_final_state, + use_qk_l2norm_in_kernel, + cu_seqlens, + cu_seqlens_cpu, + cp_context, + transpose_state_layout, + x, + eps, + log_atk_scale, + ) + + return o, h_final, a_final diff --git a/build/torch-cuda/ops/precond_gated_delta_rule/chunk_precond_kkt_fwd.py b/build/torch-cuda/ops/precond_gated_delta_rule/chunk_precond_kkt_fwd.py new file mode 100644 index 0000000000000000000000000000000000000000..4d03b78cd839d463eda45d14e5b546c37cf5fef0 --- /dev/null +++ b/build/torch-cuda/ops/precond_gated_delta_rule/chunk_precond_kkt_fwd.py @@ -0,0 +1,163 @@ +# 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 + +# Copyright (c) 2023-2025 +# Asymmetric KKT kernel for preconditioned gated delta rule +# Computes k @ k_precond^T (asymmetric) instead of k @ k^T (symmetric) + +import torch +import triton +import triton.language as tl + +from ...backends import dispatch +from ...ops.utils import prepare_chunk_indices +from ...ops.utils.op import exp2 +from ...utils import autotune_cache_kwargs + + +@triton.heuristics({ + 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None, +}) +@triton.autotune( + configs=[ + triton.Config({'BK': BK}, num_warps=num_warps, num_stages=num_stages) + for BK in [32, 64] + for num_warps in [2, 4, 8] + for num_stages in [2, 3, 4] + ], + key=['H', 'K', 'BT', 'IS_VARLEN'], + **autotune_cache_kwargs, +) +@triton.jit(do_not_specialize=['T']) +def chunk_precond_kkt_fwd_kernel( + k, # [B, T, H, K] - original keys + k_precond, # [B, T, H, K] - pre-computed preconditioned keys + g, # [B, T, HV] - gate (cumsum) + beta, # [B, T, HV] - beta scaling + A, # [B, T, HV, BT] - output asymmetric KKT matrix + cu_seqlens, + chunk_indices, + T, + H: tl.constexpr, + HV: tl.constexpr, + K: tl.constexpr, + BT: tl.constexpr, + BK: tl.constexpr, + IS_VARLEN: tl.constexpr, +): + """ + Compute asymmetric KKT matrix using pre-computed k_precond. + + A_ij = beta_i * k_i @ k_precond_j^T * exp(g_i - g_j) + + GVA (`HV > H`) is supported: `k`/`k_precond` carry `H` heads while `g`/`beta`/`A` + carry `HV` heads; each value head reads its shared key head via `i_hv // (HV // H)`. + """ + i_t, i_bh = tl.program_id(0).to(tl.int64), tl.program_id(1).to(tl.int64) + i_b, i_hv = i_bh // HV, i_bh % HV + i_h = i_hv // (HV // H) + + 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.int64) + bos, eos = tl.load(cu_seqlens + i_n).to(tl.int64), tl.load(cu_seqlens + i_n + 1).to(tl.int64) + T = eos - bos + else: + i_n = i_b + bos, eos = i_b * T, i_b * T + T + + o_t = i_t * BT + tl.arange(0, BT) + m_t = o_t < T + + b_b = tl.load(beta + bos*HV + i_hv + o_t*HV, mask=m_t, other=0.0) + b_g = tl.load(g + bos*HV + i_hv + o_t*HV, mask=m_t, other=0.0) + + b_A = tl.zeros([BT, BT], dtype=tl.float32) + + for i_k in range(tl.cdiv(K, BK)): + o_k = i_k * BK + tl.arange(0, BK) + m_tk = m_t[:, None] & (o_k[None, :] < K) + p_k = k + (bos*H + i_h) * K + o_t[:, None] * (H*K) + o_k[None, :] + b_k = tl.load(p_k, mask=m_tk, other=0.0) + + p_kp = k_precond + (bos*H + i_h) * K + o_t[:, None] * (H*K) + o_k[None, :] + b_kp = tl.load(p_kp, mask=m_tk, other=0.0) + + b_A = tl.dot(b_k, tl.trans(b_kp), b_A) + + # Attention gating and beta scaling (gates are pre-scaled by RCP_LN2 upstream) + b_A *= exp2(b_g[:, None] - b_g[None, :]) + b_A *= b_b[:, None] + + # Causal mask + m_A = (o_t[:, None] > o_t[None, :]) & (m_t[:, None] & m_t) + b_A = tl.where(m_A, b_A, 0) + + o_A = tl.arange(0, BT) + p_A = A + (bos*HV + i_hv) * BT + o_t[:, None] * (BT*HV) + o_A[None, :] + tl.store(p_A, b_A.to(A.dtype.element_ty), mask=m_t[:, None] & (o_A[None, :] < BT)) + + +@dispatch('precond_gated_delta_rule') +def chunk_precond_kkt_fwd( + k: torch.Tensor, + k_precond: torch.Tensor, + g: torch.Tensor, + beta: torch.Tensor, + chunk_size: int = 64, + output_dtype: torch.dtype = torch.float32, + cu_seqlens: torch.LongTensor | None = None, +) -> torch.Tensor: + r""" + Compute beta * K * K_precond^T (asymmetric). + + Args: + k (torch.Tensor): + The key tensor of shape `[B, T, H, K]`. + k_precond (torch.Tensor): + The preconditioned key tensor of shape `[B, T, H, K]`. + g (torch.Tensor): + The cumulative sum of the gate tensor of shape `[B, T, HV]`. + GVA is applied if `HV > H`, where `HV` must be divisible by `H`. + beta (torch.Tensor): + The beta tensor of shape `[B, T, HV]`. + cu_seqlens (torch.LongTensor): + The cumulative sequence lengths of the input tensor. + Default: None + chunk_size (int): + The chunk size. Default: 64. + output_dtype (torch.dtype): + The dtype of the output tensor. Default: `torch.float32` + + Returns: + beta * K * K_precond^T of shape `[B, T, HV, BT]` where `BT` is the chunk size. + """ + B, T, H, K = k.shape + HV = beta.shape[-1] + assert HV % H == 0, f"HV ({HV}) must be divisible by H ({H})" + BT = chunk_size + + chunk_indices = prepare_chunk_indices(cu_seqlens, BT) if cu_seqlens is not None else None + NT = triton.cdiv(T, BT) if cu_seqlens is None else len(chunk_indices) + + A = torch.empty(B, T, HV, BT, device=k.device, dtype=output_dtype) + + chunk_precond_kkt_fwd_kernel[(NT, B * HV)]( + k=k, + k_precond=k_precond, + g=g, + beta=beta, + A=A, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + T=T, + H=H, + HV=HV, + K=K, + BT=BT, + ) + + return A diff --git a/build/torch-cuda/ops/precond_gated_delta_rule/chunk_precond_wy_bwd.py b/build/torch-cuda/ops/precond_gated_delta_rule/chunk_precond_wy_bwd.py new file mode 100644 index 0000000000000000000000000000000000000000..55888d923157b98b828ef92f8110a14007df483f --- /dev/null +++ b/build/torch-cuda/ops/precond_gated_delta_rule/chunk_precond_wy_bwd.py @@ -0,0 +1,336 @@ +# 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 +import triton.language as tl + +from ...backends import dispatch +from ...ops.utils import prepare_chunk_indices +from ...ops.utils.op import exp2 +from ...utils import IS_NVIDIA_BLACKWELL, autotune_cache_kwargs, check_shared_mem + +if IS_NVIDIA_BLACKWELL: + """ + Compute tl.dot with SM100 workaround. + + On SM100 (Blackwell) GPUs, wraps the result in inline assembly to prevent + the TritonGPUHoistTMEMAlloc pass from incorrectly fusing add and dot operations. + See: https://github.com/fla-org/flash-linear-attention/issues/638 + + TODO: Remove this workaround once the Triton compiler bug is fixed. + Track upstream issue at: https://github.com/triton-lang/triton/issues/8695 + """ + @triton.jit + def safe_dot(a, b): + return tl.inline_asm_elementwise( + asm="mov.f32 $0, $1;", + constraints="=r,r", + args=[tl.dot(a, b)], + dtype=tl.float32, + is_pure=True, + pack=1, + ) +else: + @triton.jit + def safe_dot(a, b): + return tl.dot(a, b) + + +# ============================================================================== +# Asymmetric WY Backward Kernel for Preconditioned Gated Delta Rule +# ============================================================================== +# Key difference from symmetric version: +# - Forward KKT: baseKKT_ij = beta_i * +# - Backward must use: b_A = (beta*k) @ k_precond^T +# - Row-side gradients -> dk +# - Column-side gradients -> dk_precond (separate output for ATK backward) + + +@triton.heuristics({ + 'USE_G': lambda args: args['g'] is not None, + 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None, +}) +@triton.autotune( + configs=[ + triton.Config({}, num_warps=num_warps, num_stages=num_stages) + for num_warps in [2, 4] + for num_stages in [2, 3, 4] + ], + key=['H', 'HV', 'K', 'V', 'BT', 'BK', 'BV', 'IS_VARLEN'], + **autotune_cache_kwargs, +) +@triton.jit(do_not_specialize=['T']) +def prepare_precond_wy_repr_bwd_kernel( + k, # [B, T, H, K] - original k (READ key, row side) + k_precond, # [B, T, H, K] - preconditioned k (WRITE key, column side) + v, # [B, T, HV, V] + beta, # [B, T, HV] + g, # [B, T, HV] - gate cumsum + A, # [B, T, HV, BT] - inverse WY matrix + dw, # [B, T, HV, K] - gradient of w + du, # [B, T, HV, V] - gradient of u + dk, # [B, T, HV, K] - output: gradient w.r.t. original k (row side) + dk_precond, # [B, T, HV, K] - output: gradient w.r.t. k_precond (column side) + dv, # [B, T, HV, V] - output: gradient w.r.t. v + db, # [B, T, HV] - output: gradient w.r.t. beta + dg, # [B, T, HV] - output: gradient w.r.t. g (cumsum) + cu_seqlens, # *i32 [N+1] - cumulative sequence lengths + chunk_indices, # *i32 [NT, 2] - (seq_idx, chunk_idx) pairs + T, + H: tl.constexpr, + HV: tl.constexpr, + K: tl.constexpr, + V: tl.constexpr, + BT: tl.constexpr, + BK: tl.constexpr, + BV: tl.constexpr, + USE_G: tl.constexpr, + IS_VARLEN: tl.constexpr, +): + """ + Asymmetric WY backward for preconditioned gated delta rule. + + Key insight: The forward uses asymmetric KKT matrix: + baseKKT_ij = beta_i * + + Row index (i) uses original k (READ key) + Column index (j) uses k_precond (WRITE key) + + Gradient flow: + - Row-side (from dA @ K_precond): goes to dk and dbeta + - Column-side (from dA^T @ (beta*k)): goes to dk_precond only + """ + i_t, i_bh = tl.program_id(0).to(tl.int64), tl.program_id(1).to(tl.int64) + i_b, i_hv = i_bh // HV, i_bh % HV + i_h = i_hv // (HV // H) + + 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.int64) + bos, eos = tl.load(cu_seqlens + i_n).to(tl.int64), tl.load(cu_seqlens + i_n + 1).to(tl.int64) + T = eos - bos + else: + bos, eos = i_b * T, i_b * T + T + + o_t = i_t * BT + tl.arange(0, BT) + m_t = o_t < T + o_r = tl.arange(0, BT) + + b_b = tl.load(beta + (bos*HV + i_hv) + o_t*HV, mask=m_t, other=0.0) + b_db = tl.zeros([BT], dtype=tl.float32) + p_A = A + (bos*HV + i_hv) * BT + o_r[:, None] + o_t[None, :] * (HV*BT) + b_A = tl.load(p_A, mask=(o_r[:, None] < BT) & m_t[None, :], other=0.0) + b_dA = tl.zeros([BT, BT], dtype=tl.float32) + + if USE_G: + b_g = tl.load(g + (bos*HV + i_hv) + o_t*HV, mask=m_t, other=0.0) + b_g_exp = exp2(b_g) + b_dg = tl.zeros([BT], dtype=tl.float32) + + # First pass: accumulate dA from dw using original k (for w = A^-1 @ (k * beta * g)) + for i_k in range(tl.cdiv(K, BK)): + o_k = i_k * BK + tl.arange(0, BK) + m_tk = m_t[:, None] & (o_k[None, :] < K) + p_k = k + (bos*H + i_h) * K + o_t[:, None] * (H*K) + o_k[None, :] + p_dk = dk + (bos*HV + i_hv) * K + o_t[:, None] * (HV*K) + o_k[None, :] + p_dw = dw + (bos*HV + i_hv) * K + o_t[:, None] * (HV*K) + o_k[None, :] + + # [BT, BK] + b_k = tl.load(p_k, mask=m_tk, other=0.0) + if USE_G: + b_kbg = b_k * (b_b * b_g_exp)[:, None] + else: + b_kbg = b_k * b_b[:, None] + b_dw = tl.load(p_dw, mask=m_tk, other=0.0) + + # dA contribution from dw: dA += dw @ kbg^T + b_dA += safe_dot(b_dw, tl.trans(b_kbg).to(b_dw.dtype)) + + # dk contribution from first term: A @ dw + b_dkbg = safe_dot(b_A, b_dw.to(b_A.dtype)) + if USE_G: + b_dk = b_dkbg * (b_g_exp * b_b)[:, None] + b_db += tl.sum(b_dkbg * b_k * b_g_exp[:, None], 1) + b_dg += tl.sum(b_dkbg * b_kbg, 1) + else: + b_dk = b_dkbg * b_b[:, None] + b_db += tl.sum(b_dkbg * b_k, 1) + + tl.store(p_dk, b_dk.to(dk.dtype.element_ty), mask=m_tk) + + # Process dv and du (u = A^-1 @ (v * beta), same as symmetric) + for i_v in range(tl.cdiv(V, BV)): + o_v = i_v * BV + tl.arange(0, BV) + m_tv = m_t[:, None] & (o_v[None, :] < V) + p_v = v + (bos*HV + i_hv) * V + o_t[:, None] * (HV*V) + o_v[None, :] + p_dv = dv + (bos*HV + i_hv) * V + o_t[:, None] * (HV*V) + o_v[None, :] + p_du = du + (bos*HV + i_hv) * V + o_t[:, None] * (HV*V) + o_v[None, :] + b_v = tl.load(p_v, mask=m_tv, other=0.0) + b_vb = (b_v * b_b[:, None]).to(b_v.dtype) + b_du = tl.load(p_du, mask=m_tv, other=0.0) + b_dA += safe_dot(b_du, tl.trans(b_vb)) + b_dvb = safe_dot(b_A, b_du.to(b_A.dtype)) + b_dv = b_dvb * b_b[:, None] + b_db += tl.sum(b_dvb * b_v, 1) + tl.store(p_dv, b_dv.to(dv.dtype.element_ty), mask=m_tv) + + # Transform dA through lower triangular inversion + m_A = (o_t[:, None] > o_t[None, :]) & (m_t[:, None] & m_t) + b_dA = tl.where(m_A, b_dA, 0) + b_dA = safe_dot(b_dA.to(b_A.dtype), b_A) + b_dA = safe_dot(b_A, b_dA.to(b_A.dtype)) + + if USE_G: + b_dA *= exp2(b_g[:, None] - b_g[None, :]) + + b_dA = tl.where(m_A, -b_dA, 0) + + b_dA = b_dA.to(k.dtype.element_ty) + + tl.debug_barrier() + # ==== ASYMMETRIC PART ==== + # Reconstruct b_A (the asymmetric KKT) and compute gradients + # baseKKT_ij = beta_i * + # Row index i: uses k (READ key) + # Column index j: uses k_precond (WRITE key) + # + # From the gradient chain rule: + # - dL/dk_i += sum_j dA_ij * k_precond_j * beta_i (row-side) + # - dL/dk_precond_j += sum_i dA_ij * k_i * beta_i (column-side) + # - dL/dbeta_i += sum_j dA_ij * + + b_A_asymm = tl.zeros([BT, BT], dtype=tl.float32) # Asymmetric KKT matrix + + for i_k in range(tl.cdiv(K, BK)): + o_k = i_k * BK + tl.arange(0, BK) + m_tk = m_t[:, None] & (o_k[None, :] < K) + p_k = k + (bos*H + i_h) * K + o_t[:, None] * (H*K) + o_k[None, :] + p_kp = k_precond + (bos*H + i_h) * K + o_t[:, None] * (H*K) + o_k[None, :] + p_dk = dk + (bos*HV + i_hv) * K + o_t[:, None] * (HV*K) + o_k[None, :] + p_dkp = dk_precond + (bos*HV + i_hv) * K + o_t[:, None] * (HV*K) + o_k[None, :] + + b_k = tl.load(p_k, mask=m_tk, other=0.0) # [BT, BK] - original k + b_kp = tl.load(p_kp, mask=m_tk, other=0.0) # [BT, BK] - k_precond + b_dk = tl.load(p_dk, mask=m_tk, other=0.0) # Load existing dk from first pass + + # kb = k * beta (used on row side) + b_kb = (b_k * b_b[:, None]).to(b_k.dtype) + + # Build asymmetric KKT: (beta * k) @ k_precond^T + b_A_asymm += safe_dot(b_kb, tl.trans(b_kp)) + + # Row-side gradient: dkb = dA @ k_precond + b_dkb = safe_dot(b_dA, b_kp) + # dk contribution: dkb * beta + b_dk += b_dkb * b_b[:, None] + # dbeta contribution: + b_db += tl.sum(b_dkb * b_k, 1) + + # Column-side gradient: dk_precond = dA^T @ (beta * k) + b_dkp = safe_dot(tl.trans(b_dA), b_kb) + + tl.store(p_dk, b_dk.to(dk.dtype.element_ty), mask=m_tk) + tl.store(p_dkp, b_dkp.to(dk_precond.dtype.element_ty), mask=m_tk) + + tl.store(db + (bos*HV + i_hv) + o_t*HV, b_db.to(db.dtype.element_ty), mask=m_t) + + # dg using the asymmetric KKT matrix + if USE_G: + b_AdA = b_dA * b_A_asymm + b_dg += tl.sum(b_AdA, axis=1) - tl.sum(b_AdA, axis=0) + tl.store(dg + (bos*HV + i_hv) + o_t*HV, b_dg.to(dg.dtype.element_ty), mask=m_t) + + +@dispatch('precond_gated_delta_rule') +def prepare_precond_wy_repr_bwd( + k: torch.Tensor, + k_precond: torch.Tensor, + v: torch.Tensor, + beta: torch.Tensor, + g: torch.Tensor, + A: torch.Tensor, + dw: torch.Tensor, + du: torch.Tensor, + cu_seqlens: torch.LongTensor | None = None, + chunk_indices: torch.LongTensor | None = None, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + """ + Asymmetric WY backward for preconditioned gated delta rule. + + Unlike the symmetric version that uses the same k for both row and column + of the KKT matrix, this version uses: + - k (original) for row side (READ key) + - k_precond for column side (WRITE key) + + Args: + k: Original keys [B, T, H, K] + k_precond: Preconditioned keys [B, T, H, K] + v: Values [B, T, H, V] + beta: Beta scaling [B, T, H] + g: Gate cumsum [B, T, H] + A: Inverse WY matrix [B, T, H, BT] + dw: Gradient of w [B, T, H, K] + du: Gradient of u [B, T, H, V] + cu_seqlens: Cumulative sequence lengths [N+1] for varlen + chunk_indices: Precomputed chunk indices [NT, 2] for varlen + + Returns: + dk: Gradient w.r.t. original k (row-side) [B, T, H, K] + dk_precond: Gradient w.r.t. k_precond (column-side) [B, T, H, K] + dv: Gradient w.r.t. v [B, T, H, V] + dbeta: Gradient w.r.t. beta [B, T, H] + dg: Gradient w.r.t. g cumsum [B, T, H] + """ + B, T, H, K, V = *k.shape, v.shape[-1] + HV = v.shape[2] + BT = 64 + + if chunk_indices is None and cu_seqlens is not None: + chunk_indices = prepare_chunk_indices(cu_seqlens, BT) + NT = triton.cdiv(T, BT) if cu_seqlens is None else len(chunk_indices) + + CONST_TILING = 64 if check_shared_mem() else 32 + BK = min(max(triton.next_power_of_2(K), 16), CONST_TILING) + BV = min(max(triton.next_power_of_2(V), 16), CONST_TILING) + + dk = k.new_empty(B, T, HV, K) + dk_precond = k.new_empty(B, T, HV, K) + dv = torch.empty_like(v) + dg = torch.empty_like(g) if g is not None else None + db = torch.empty_like(beta) + + prepare_precond_wy_repr_bwd_kernel[(NT, B * HV)]( + k=k, + k_precond=k_precond, + v=v, + beta=beta, + g=g, + A=A, + dw=dw, + du=du, + dk=dk, + dk_precond=dk_precond, + dv=dv, + db=db, + dg=dg, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + T=T, + H=H, + HV=HV, + K=K, + V=V, + BT=BT, + BK=BK, + BV=BV, + ) + + if H != HV: + dk = dk.view(B, T, H, HV // H, K).sum(3) + dk_precond = dk_precond.view(B, T, H, HV // H, K).sum(3) + + return dk, dk_precond, dv, db, dg diff --git a/build/torch-xpu/ops/gated_delta_rule/fused_recurrent.py b/build/torch-cuda/ops/precond_gated_delta_rule/fused_recurrent.py similarity index 62% rename from build/torch-xpu/ops/gated_delta_rule/fused_recurrent.py rename to build/torch-cuda/ops/precond_gated_delta_rule/fused_recurrent.py index 978e34838d59d865b04eaf9a1d6ecc3c6adaf0da..67483ad57c18601e934282bbb361b74d420e54bc 100644 --- a/build/torch-xpu/ops/gated_delta_rule/fused_recurrent.py +++ b/build/torch-cuda/ops/precond_gated_delta_rule/fused_recurrent.py @@ -5,14 +5,12 @@ # For a list of all contributors, visit: # https://github.com/fla-org/flash-linear-attention/graphs/contributors -import warnings - import torch import triton import triton.language as tl -from ...ops.utils.op import exp -from ...ops.utils.softplus import softplus +from ...backends import dispatch +from ...ops.utils.op import exp, unflatten_program_id from ...utils import input_guard @@ -22,12 +20,12 @@ from ...utils import input_guard 'USE_GV': lambda args: args['gv'] is not None, 'USE_INITIAL_STATE': lambda args: args['h0'] is not None, 'STORE_FINAL_STATE': lambda args: args['ht'] is not None, + 'USE_INITIAL_ATK': lambda args: args['a0'] is not None, + 'STORE_FINAL_ATK': lambda args: args['at'] is not None, 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None, - 'USE_GATE_IN_KERNEL': lambda args: args['A_log'] is not None, - 'HAS_DT_BIAS': lambda args: args['dt_bias'] is not None, }) @triton.jit(do_not_specialize=['T']) -def fused_recurrent_gated_delta_rule_fwd_kernel( +def fused_recurrent_precond_gated_delta_rule_fwd_kernel( q, k, v, @@ -35,13 +33,18 @@ def fused_recurrent_gated_delta_rule_fwd_kernel( gk, gv, beta, - A_log, - dt_bias, + g_atk, + beta_atk, + log_atk_scale, o, h0, ht, + a0, + at, cu_seqlens, scale, + x, + eps, T, H: tl.constexpr, HV: tl.constexpr, @@ -56,14 +59,12 @@ def fused_recurrent_gated_delta_rule_fwd_kernel( IS_BETA_HEADWISE: tl.constexpr, USE_INITIAL_STATE: tl.constexpr, STORE_FINAL_STATE: tl.constexpr, - STATE_V_FIRST: tl.constexpr, + USE_INITIAL_ATK: tl.constexpr, + STORE_FINAL_ATK: tl.constexpr, + TRANSPOSE_STATE: tl.constexpr, IS_VARLEN: tl.constexpr, - USE_GATE_IN_KERNEL: tl.constexpr, - HAS_DT_BIAS: tl.constexpr, - APPLY_BETA_SIGMOID: tl.constexpr, - ALLOW_NEG_EIGVAL: tl.constexpr, ): - i_v, i_nh = tl.program_id(0), tl.program_id(1) + i_v, i_nh = unflatten_program_id(tl.cdiv(V, BV)) i_n, i_hv = i_nh // HV, i_nh % HV i_h = i_hv // (HV // H) @@ -89,26 +90,40 @@ def fused_recurrent_gated_delta_rule_fwd_kernel( else: p_beta = beta + (bos * HV + i_hv) * V + o_v + # ATK pointers (per key head: ATK preconditions k, so g_atk, beta_atk and + # log_atk_scale carry H heads; under GVA every value head in a group shares them) + p_g_atk = g_atk + bos * H + i_h + p_beta_atk = beta_atk + bos * H + i_h + + scale_val = tl.load(log_atk_scale + i_h).to(tl.float32) + logx = tl.log(tl.cast(x, tl.float32)) + p_o = o + (bos * HV + i_hv) * V + o_v mask_k = o_k < K mask_v = o_v < V - if STATE_V_FIRST: + if TRANSPOSE_STATE: mask_h = mask_v[:, None] & mask_k[None, :] else: mask_h = mask_k[:, None] & mask_v[None, :] - if STATE_V_FIRST: + if TRANSPOSE_STATE: b_h = tl.zeros([BV, BK], dtype=tl.float32) else: b_h = tl.zeros([BK, BV], dtype=tl.float32) if USE_INITIAL_STATE: - if STATE_V_FIRST: + if TRANSPOSE_STATE: p_h0 = h0 + i_nh * K*V + o_v[:, None] * K + o_k[None, :] else: p_h0 = h0 + i_nh * K*V + o_k[:, None] * V + o_v[None, :] b_h += tl.load(p_h0, mask=mask_h, other=0).to(tl.float32) + # ATK state (per key head: [N, H, K]) + b_a = tl.zeros([BK], dtype=tl.float32) + if USE_INITIAL_ATK: + p_a0 = a0 + (i_n * H + i_h) * K + o_k + b_a += tl.load(p_a0, mask=mask_k, other=0).to(tl.float32) + for _ in tl.range(0, T): b_q = tl.load(p_q, mask=mask_k, other=0).to(tl.float32) b_k = tl.load(p_k, mask=mask_k, other=0).to(tl.float32) @@ -121,41 +136,45 @@ def fused_recurrent_gated_delta_rule_fwd_kernel( b_beta = tl.load(p_beta).to(tl.float32) else: b_beta = tl.load(p_beta, mask=mask_v, other=0).to(tl.float32) - if APPLY_BETA_SIGMOID: - b_beta = tl.sigmoid(b_beta) - if ALLOW_NEG_EIGVAL: - b_beta = b_beta * 2 + + # ATK: A_t = exp(g_atk) * A_{t-1} + beta_atk * k^2 + b_g_atk = tl.load(p_g_atk).to(tl.float32) + b_beta_atk = tl.load(p_beta_atk).to(tl.float32) + b_a = exp(b_g_atk) * b_a + b_beta_atk * (b_k * b_k) + + # Symmetric fast squash preconditioner + b_ell = tl.log(b_a + eps) + b_r = b_ell - scale_val + b_s = b_r / (1.0 + tl.abs(b_r)) + b_M = tl.exp(-logx * b_s) + b_k_precond = b_k * b_M if USE_G: b_g = tl.load(p_g).to(tl.float32) - if USE_GATE_IN_KERNEL: - b_A = tl.load(A_log + i_hv).to(tl.float32) - if HAS_DT_BIAS: - b_g = b_g + tl.load(dt_bias + i_hv).to(tl.float32) - b_g = -exp(b_A) * softplus(b_g) b_h *= exp(b_g) if USE_GK: - b_gk = tl.load(p_gk).to(tl.float32) - if STATE_V_FIRST: + b_gk = tl.load(p_gk, mask=mask_k, other=0).to(tl.float32) + if TRANSPOSE_STATE: b_h *= exp(b_gk[None, :]) else: b_h *= exp(b_gk[:, None]) if USE_GV: - b_gv = tl.load(p_gv).to(tl.float32) - if STATE_V_FIRST: + b_gv = tl.load(p_gv, mask=mask_v, other=0).to(tl.float32) + if TRANSPOSE_STATE: b_h *= exp(b_gv[:, None]) else: b_h *= exp(b_gv[None, :]) - if STATE_V_FIRST: + # Delta rule update using k_precond for the write side + if TRANSPOSE_STATE: b_v = b_beta * (b_v - tl.sum(b_h * b_k[None, :], 1)) - b_h += b_v[:, None] * b_k[None, :] + b_h += b_v[:, None] * b_k_precond[None, :] b_o = tl.sum(b_h * b_q[None, :], 1) else: b_v = b_beta * (b_v - tl.sum(b_h * b_k[:, None], 0)) - b_h += b_k[:, None] * b_v + b_h += b_k_precond[:, None] * b_v b_o = tl.sum(b_h * b_q[:, None], 0) tl.store(p_o, b_o.to(p_o.dtype.element_ty), mask=mask_v) @@ -170,16 +189,26 @@ def fused_recurrent_gated_delta_rule_fwd_kernel( p_gv += HV*V p_beta += HV * (1 if IS_BETA_HEADWISE else V) p_o += HV*V + p_g_atk += H + p_beta_atk += H if STORE_FINAL_STATE: - if STATE_V_FIRST: + if TRANSPOSE_STATE: p_ht = ht + i_nh * K*V + o_v[:, None] * K + o_k[None, :] else: p_ht = ht + i_nh * K*V + o_k[:, None] * V + o_v[None, :] tl.store(p_ht, b_h.to(p_ht.dtype.element_ty), mask=mask_h) + if STORE_FINAL_ATK: + # The ATK state is shared by all value heads in a GVA group; only the + # first program of each (i_v, group) stores it to avoid duplicate writes. + if (i_v == 0) and (i_hv % (HV // H) == 0): + p_at = at + (i_n * H + i_h) * K + o_k + tl.store(p_at, b_a.to(p_at.dtype.element_ty), mask=mask_k) -def fused_recurrent_gated_delta_rule_fwd( + +@dispatch('precond_gated_delta_rule') +def fused_recurrent_precond_gated_delta_rule_fwd( q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, @@ -187,17 +216,19 @@ def fused_recurrent_gated_delta_rule_fwd( gk: torch.Tensor | None = None, gv: torch.Tensor | None = None, beta: torch.Tensor | None = None, - A_log: torch.Tensor | None = None, - dt_bias: torch.Tensor | None = None, + g_atk: torch.Tensor = None, + beta_atk: torch.Tensor = None, scale: float = None, initial_state: torch.Tensor = None, + initial_A_state: torch.Tensor = None, output_final_state: bool = False, - use_qk_l2norm_in_kernel: bool = False, - use_beta_sigmoid_in_kernel: bool = False, - allow_neg_eigval: bool = False, - state_v_first: bool = False, + use_qk_l2norm_in_kernel: bool = True, cu_seqlens: torch.LongTensor | None = None, -) -> tuple[torch.Tensor, torch.Tensor]: + transpose_state_layout: bool = False, + x: float = 1.5, + eps: float = 1e-6, + log_atk_scale: torch.Tensor = None, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: B, T, H, K, V = *k.shape, v.shape[-1] HV = v.shape[2] N = B if cu_seqlens is None else len(cu_seqlens) - 1 @@ -205,17 +236,22 @@ def fused_recurrent_gated_delta_rule_fwd( BV = min(8, triton.next_power_of_2(V)) if gv is None else triton.next_power_of_2(V) NV = triton.cdiv(V, BV) + if log_atk_scale is None: + log_atk_scale = torch.full((H,), -0.2, device=k.device, dtype=torch.float32) + o = torch.empty_like(v) if output_final_state: - if state_v_first: + if transpose_state_layout: final_state = q.new_empty(N, HV, V, K, dtype=torch.float32) else: final_state = q.new_empty(N, HV, K, V, dtype=torch.float32) + final_A_state = q.new_empty(N, H, K, dtype=torch.float32) else: final_state = None + final_A_state = None - grid = (NV, N * HV) - fused_recurrent_gated_delta_rule_fwd_kernel[grid]( + grid = (NV * N * HV,) + fused_recurrent_precond_gated_delta_rule_fwd_kernel[grid]( q=q, k=k, v=v, @@ -223,13 +259,18 @@ def fused_recurrent_gated_delta_rule_fwd( gk=gk, gv=gv, beta=beta, - A_log=A_log, - dt_bias=dt_bias, + g_atk=g_atk, + beta_atk=beta_atk, + log_atk_scale=log_atk_scale, o=o, h0=initial_state, ht=final_state, + a0=initial_A_state, + at=final_A_state, cu_seqlens=cu_seqlens, scale=scale, + x=x, + eps=eps, T=T, H=H, HV=HV, @@ -239,16 +280,14 @@ def fused_recurrent_gated_delta_rule_fwd( BV=BV, IS_BETA_HEADWISE=beta.ndim != v.ndim, USE_QK_L2NORM_IN_KERNEL=use_qk_l2norm_in_kernel, - APPLY_BETA_SIGMOID=use_beta_sigmoid_in_kernel, - ALLOW_NEG_EIGVAL=allow_neg_eigval, - STATE_V_FIRST=state_v_first, + TRANSPOSE_STATE=transpose_state_layout, num_warps=1, num_stages=3, ) - return o, final_state + return o, final_state, final_A_state -class FusedRecurrentFunction(torch.autograd.Function): +class FusedRecurrentPrecondGatedDeltaRuleFunction(torch.autograd.Function): @staticmethod @input_guard @@ -261,18 +300,20 @@ class FusedRecurrentFunction(torch.autograd.Function): gk: torch.Tensor | None = None, gv: torch.Tensor | None = None, beta: torch.Tensor | None = None, - A_log: torch.Tensor | None = None, - dt_bias: torch.Tensor | None = None, + g_atk: torch.Tensor = None, + beta_atk: torch.Tensor = None, scale: float = None, initial_state: torch.Tensor = None, + initial_A_state: torch.Tensor = None, output_final_state: bool = False, use_qk_l2norm_in_kernel: bool = False, - use_beta_sigmoid_in_kernel: bool = False, - allow_neg_eigval: bool = False, - state_v_first: bool = False, cu_seqlens: torch.LongTensor | None = None, + transpose_state_layout: bool = False, + x: float = 1.5, + eps: float = 1e-6, + log_atk_scale: torch.Tensor = None, ): - o, final_state = fused_recurrent_gated_delta_rule_fwd( + o, final_state, final_A_state = fused_recurrent_precond_gated_delta_rule_fwd( q=q, k=k, v=v, @@ -280,23 +321,25 @@ class FusedRecurrentFunction(torch.autograd.Function): gk=gk, gv=gv, beta=beta, - A_log=A_log, - dt_bias=dt_bias, + g_atk=g_atk, + beta_atk=beta_atk, scale=scale, initial_state=initial_state, + initial_A_state=initial_A_state, output_final_state=output_final_state, use_qk_l2norm_in_kernel=use_qk_l2norm_in_kernel, - use_beta_sigmoid_in_kernel=use_beta_sigmoid_in_kernel, - allow_neg_eigval=allow_neg_eigval, - state_v_first=state_v_first, cu_seqlens=cu_seqlens, + transpose_state_layout=transpose_state_layout, + x=x, + eps=eps, + log_atk_scale=log_atk_scale, ) - return o, final_state + return o, final_state, final_A_state @staticmethod @input_guard - def backward(ctx, do, dht): + def backward(ctx, do, dht, dat): raise NotImplementedError( "Backward pass is not implemented yet and we do not have plans to implement it " "because we haven't figured out how to compute dg without materializing the full " @@ -304,7 +347,7 @@ class FusedRecurrentFunction(torch.autograd.Function): ) -def fused_recurrent_gated_delta_rule( +def fused_recurrent_precond_gated_delta_rule( q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, @@ -312,19 +355,19 @@ def fused_recurrent_gated_delta_rule( gk: torch.Tensor | None = None, gv: torch.Tensor | None = None, beta: torch.Tensor | None = None, + g_atk: torch.Tensor = None, + beta_atk: torch.Tensor = None, scale: float = None, initial_state: torch.Tensor = None, + initial_A_state: torch.Tensor = None, output_final_state: bool = False, - use_qk_l2norm_in_kernel: bool = False, - use_gate_in_kernel: bool = False, - A_log: torch.Tensor | None = None, - dt_bias: torch.Tensor | None = None, - use_beta_sigmoid_in_kernel: bool = False, - allow_neg_eigval: bool = False, - state_v_first: bool = False, + use_qk_l2norm_in_kernel: bool = True, cu_seqlens: torch.LongTensor | None = None, - **kwargs, -) -> tuple[torch.Tensor, torch.Tensor]: + transpose_state_layout: bool = False, + x: float = 1.5, + eps: float = 1e-6, + log_atk_scale: torch.Tensor = None, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: r""" Args: q (torch.Tensor): @@ -336,61 +379,55 @@ def fused_recurrent_gated_delta_rule( GVA (Grouped Value Attention) is applied if `HV > H`, where `HV` must be divisible by `H`. g (torch.Tensor): g (decays) of shape `[B, T, HV]`. Default: `None`. - When `use_gate_in_kernel=False` (default), `g` must be in log space (pre-computed decay). - When `use_gate_in_kernel=True`, `g` is the raw pre-activation input; the kernel fuses - `-exp(A_log) * softplus(g + dt_bias)` internally per step. gk (torch.Tensor): gk (decays) of shape `[B, T, HV, K]`. Default: `None`. gv (torch.Tensor): gv (decays) of shape `[B, T, HV, V]`. Default: `None`. beta (torch.Tensor): betas of shape `[B, T, HV]`. + g_atk (torch.Tensor): + ATK gates (decays) in log space of shape `[B, T, H]`. + beta_atk (torch.Tensor): + ATK betas of shape `[B, T, H]`. scale (Optional[float]): - Scale factor for the RetNet attention scores. + Scale factor for the attention scores. If not provided, it will default to `1 / sqrt(K)`. Default: `None`. initial_state (Optional[torch.Tensor]): Initial state of shape `[N, HV, K, V]` for `N` input sequences. For equal-length input sequences, `N` equals the batch size `B`. Default: `None`. + initial_A_state (Optional[torch.Tensor]): + Initial ATK diagonal state of shape `[N, H, K]`; the ATK preconditioner + acts on the keys, so its state carries `H` (key) heads. Default: `None`. output_final_state (Optional[bool]): Whether to output the final state of shape `[N, HV, K, V]`. Default: `False`. use_qk_l2norm_in_kernel (Optional[bool]): - Whether to use L2 normalization in the kernel. Default: `False`. - use_gate_in_kernel (bool): - Whether to compute the log-space GDN decay internally. - When `True`, `g` is the raw input and `A_log` must be provided; the kernel fuses - gate activation into the recurrence. Default: `False`. - A_log (Optional[torch.Tensor]): - Decay parameter of shape `[HV]`. Required when `use_gate_in_kernel=True`. - dt_bias (Optional[torch.Tensor]): - Bias added to `g` before activation, of shape `[HV]`. - Only used when `use_gate_in_kernel=True`. - use_beta_sigmoid_in_kernel (Optional[bool]): - Whether to apply `torch.sigmoid(beta)` inside the kernel. - - If `True`, the passed `beta` acts as the raw beta logits. - - If `False`, `beta` is expected to already be in post-sigmoid space. - Default: `False`. - allow_neg_eigval (Optional[bool]): - Whether to allow negative eigenvalues by scaling `beta` to `[0, 2)`. - Only takes effect together with `use_beta_sigmoid_in_kernel=True`, in which case - the kernel computes `2 * sigmoid(beta)` instead of `sigmoid(beta)`. Default: `False`. - state_v_first (Optional[bool]): - Store the recurrent state in V-first ``[V, K]`` layout instead of the default ``[K, V]``. Default: ``False``. + Whether to use L2 normalization in the kernel. Default: `True`. cu_seqlens (torch.LongTensor): Cumulative sequence lengths of shape `[N+1]` used for variable-length training, consistent with the FlashAttention API. + transpose_state_layout (bool): + Whether to use transposed state layout `[V, K]` instead of `[K, V]`. Default: `False`. + x (float): + Squash range parameter. Default: `1.5`. + eps (float): + Epsilon for numerical stability in the squash function. Default: `1e-6`. + log_atk_scale (Optional[torch.Tensor]): + Per-head log-space center of shape `[H]`. Default: `None`. Returns: o (torch.Tensor): Outputs of shape `[B, T, HV, V]`. final_state (torch.Tensor): Final state of shape `[N, HV, K, V]` if `output_final_state=True` else `None`. + final_A_state (torch.Tensor): + Final ATK diagonal state of shape `[N, H, K]` if `output_final_state=True` else `None`. Examples:: >>> import torch >>> import torch.nn.functional as F >>> from einops import rearrange - >>> from ...ops.gated_delta_rule import fused_recurrent_gated_delta_rule + >>> from fla.ops.precond_gated_delta_rule import fused_recurrent_precond_gated_delta_rule # inputs with equal lengths >>> B, T, H, HV, K, V = 4, 2048, 4, 8, 512, 512 >>> q = torch.randn(B, T, H, K, device='cuda') @@ -398,33 +435,28 @@ def fused_recurrent_gated_delta_rule( >>> v = torch.randn(B, T, HV, V, device='cuda') >>> g = F.logsigmoid(torch.rand(B, T, HV, device='cuda')) >>> beta = torch.rand(B, T, HV, device='cuda').sigmoid() + >>> g_atk = F.logsigmoid(torch.rand(B, T, H, device='cuda')) + >>> beta_atk = torch.rand(B, T, H, device='cuda').sigmoid() >>> h0 = torch.randn(B, HV, K, V, device='cuda') - >>> o, ht = fused_gated_recurrent_delta_rule( - q, k, v, g, beta, + >>> o, ht, at = fused_recurrent_precond_gated_delta_rule( + q, k, v, g=g, beta=beta, + g_atk=g_atk, beta_atk=beta_atk, initial_state=h0, output_final_state=True ) # for variable-length inputs, the batch size `B` is expected to be 1 and `cu_seqlens` is required >>> q, k, v, g, beta = map(lambda x: rearrange(x, 'b t ... -> 1 (b t) ...'), (q, k, v, g, beta)) + >>> g_atk, beta_atk = map(lambda x: rearrange(x, 'b t ... -> 1 (b t) ...'), (g_atk, beta_atk)) # for a batch with 4 sequences, `cu_seqlens` with 5 start/end positions are expected >>> cu_seqlens = q.new_tensor([0, 2048, 4096, 6144, 8192], dtype=torch.long) - >>> o, ht = fused_gated_recurrent_delta_rule( - q, k, v, g, beta, + >>> o, ht, at = fused_recurrent_precond_gated_delta_rule( + q, k, v, g=g, beta=beta, + g_atk=g_atk, beta_atk=beta_atk, initial_state=h0, output_final_state=True, cu_seqlens=cu_seqlens ) """ - if 'transpose_state_layout' in kwargs: - if state_v_first: - raise ValueError("Cannot pass both `state_v_first` and the deprecated `transpose_state_layout`.") - warnings.warn( - "`transpose_state_layout` is deprecated and renamed to `state_v_first`.", - DeprecationWarning, - stacklevel=2, - ) - state_v_first = kwargs.pop('transpose_state_layout') - if cu_seqlens is not None: if q.shape[0] != 1: raise ValueError( @@ -440,18 +472,8 @@ def fused_recurrent_gated_delta_rule( scale = k.shape[-1] ** -0.5 if beta is None: beta = torch.ones_like(q[..., 0]) - if use_gate_in_kernel: - if A_log is None: - raise ValueError("`A_log` must be provided when `use_gate_in_kernel=True`.") - if g is None: - raise ValueError("`g` (raw pre-activation) must be provided when `use_gate_in_kernel=True`.") - else: - A_log = None - dt_bias = None - if allow_neg_eigval and not use_beta_sigmoid_in_kernel: - raise ValueError("`allow_neg_eigval=True` requires `use_beta_sigmoid_in_kernel=True`.") - o, final_state = FusedRecurrentFunction.apply( + o, final_state, final_A_state = FusedRecurrentPrecondGatedDeltaRuleFunction.apply( q, k, v, @@ -459,18 +481,20 @@ def fused_recurrent_gated_delta_rule( gk, gv, beta, - A_log, - dt_bias, + g_atk, + beta_atk, scale, initial_state, + initial_A_state, output_final_state, use_qk_l2norm_in_kernel, - use_beta_sigmoid_in_kernel, - allow_neg_eigval, - state_v_first, cu_seqlens, + transpose_state_layout, + x, + eps, + log_atk_scale, ) - return o, final_state + return o, final_state, final_A_state -fused_recurrent_gdn = fused_recurrent_gated_delta_rule +fused_recurrent_pgdn = fused_recurrent_precond_gated_delta_rule diff --git a/build/torch-cuda/ops/precond_gated_delta_rule/naive.py b/build/torch-cuda/ops/precond_gated_delta_rule/naive.py new file mode 100644 index 0000000000000000000000000000000000000000..d898987050a2bc02885c3e5958a5d94a26310808 --- /dev/null +++ b/build/torch-cuda/ops/precond_gated_delta_rule/naive.py @@ -0,0 +1,156 @@ +# 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 torch.nn.functional as F + + +def naive_recurrent_precond_gated_delta_rule( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + g_atk: torch.Tensor, + g: torch.Tensor, + beta_atk: torch.Tensor, + beta: torch.Tensor, + scale: float = None, + initial_state: torch.Tensor = None, + initial_A_state: torch.Tensor = None, + output_final_state: bool = False, + x: float = 1.5, + eps: float = 1e-6, + log_atk_scale: torch.Tensor | float = None, + cu_seqlens: torch.LongTensor | None = None, +): + """ + Reference PyTorch implementation of recurrent preconditioned gated delta rule. + + Args: + q: [B, T, H, K] + k: [B, T, H, K] + v: [B, T, HV, V] + g_atk: [B, T, H] + g: [B, T, HV] + beta_atk: [B, T, H] + beta: [B, T, HV] + scale: float, optional + initial_state: [B, HV, K, V], optional + initial_A_state: [B, H, K], optional + output_final_state: bool + cu_seqlens: [N+1], optional cumulative sequence lengths for packed inputs + + GVA is applied if `HV > H` (with `HV` divisible by `H`): each group of + `HV // H` value heads shares one key head and its ATK preconditioner state. + + Returns: + o: [B, T, HV, V] + final_state: [B, HV, K, V] if output_final_state else None + final_A_state: [B, H, K] if output_final_state else None + """ + if cu_seqlens is not None: + outputs, final_states, final_A_states = [], [], [] + offsets = cu_seqlens.tolist() + for i, (bos, eos) in enumerate(zip(offsets[:-1], offsets[1:], strict=False)): + o_i, final_state_i, final_A_state_i = naive_recurrent_precond_gated_delta_rule( + q=q[:, bos:eos], + k=k[:, bos:eos], + v=v[:, bos:eos], + g_atk=g_atk[:, bos:eos], + g=g[:, bos:eos], + beta_atk=beta_atk[:, bos:eos], + beta=beta[:, bos:eos], + scale=scale, + initial_state=None if initial_state is None else initial_state[i:i + 1], + initial_A_state=None if initial_A_state is None else initial_A_state[i:i + 1], + output_final_state=output_final_state, + x=x, + eps=eps, + log_atk_scale=log_atk_scale, + ) + outputs.append(o_i) + if output_final_state: + final_states.append(final_state_i) + final_A_states.append(final_A_state_i) + + o = torch.cat(outputs, dim=1) + if not output_final_state: + return o, None, None + return o, torch.cat(final_states), torch.cat(final_A_states) + + num_heads, num_v_heads = k.shape[2], v.shape[2] + n_rep = num_v_heads // num_heads + if n_rep > 1: + # GVA: value heads in a group share the key head and ATK state + q, k, g_atk, beta_atk = (t.repeat_interleave(n_rep, dim=2) for t in (q, k, g_atk, beta_atk)) + if initial_A_state is not None: + initial_A_state = initial_A_state.repeat_interleave(n_rep, dim=1) + if log_atk_scale is not None and not isinstance(log_atk_scale, (int, float)): + log_atk_scale = log_atk_scale.repeat_interleave(n_rep, dim=0) + + q, k, v, g_atk, g, beta_atk, beta = map( + lambda t: t.transpose(1, 2).contiguous().to(torch.float32), + [q, k, v, g_atk, g, beta_atk, beta] + ) + B, H, T, K, V = *k.shape, v.shape[-1] + o = torch.zeros(B, H, T, V).to(v) + S = torch.zeros(B, H, K, V).to(v) + A = torch.zeros(B, H, K).to(v) + if initial_state is not None: + S = initial_state.clone() + if initial_A_state is not None: + A = initial_A_state.clone() + if scale is None: + scale = 1 / (K ** 0.5) + + # L2 normalize q and k, then apply scale + q = F.normalize(q, p=2, dim=-1) * scale + k = F.normalize(k, p=2, dim=-1) + + if log_atk_scale is None: + log_atk_scale = -0.2 + if isinstance(log_atk_scale, (int, float)): + center = torch.tensor(log_atk_scale, dtype=torch.float32, device=k.device) + else: + center = log_atk_scale.to(torch.float32).view(1, H, 1) + + logx = torch.log(torch.tensor(x, dtype=torch.float32, device=k.device)) + + for i in range(T): + b_q = q[:, :, i] + b_k = k[:, :, i] + b_v = v[:, :, i].clone() + b_g_atk = g_atk[:, :, i] + b_g = g[:, :, i] + b_beta_atk = beta_atk[:, :, i] + b_beta = beta[:, :, i] + + # ATK recurrence + A = A.clone() * b_g_atk.exp().unsqueeze(-1) + b_beta_atk.unsqueeze(-1) * (b_k ** 2) + + # Symmetric fast squash + ell = torch.log(A + eps) + r = ell - center + s = r / (1.0 + torch.abs(r)) + M = torch.exp(-logx * s) + b_k_precond = b_k * M + + # Gated delta rule + S = S.clone() * b_g.exp()[..., None, None] + b_v = b_v - (S.clone() * b_k[..., None]).sum(-2) + b_v = b_v * b_beta[..., None] + S = S.clone() + b_k_precond.unsqueeze(-1) * b_v.unsqueeze(-2) + o[:, :, i] = torch.einsum('bhd,bhdm->bhm', b_q, S) + + if not output_final_state: + S = None + A = None + elif n_rep > 1: + # all value heads in a group carry the identical shared ATK state; + # return one copy per key head + A = A[:, ::n_rep] + o = o.transpose(1, 2).contiguous() + return o, S, A diff --git a/build/torch-cuda/modules/conv/cp/__init__.py b/build/torch-cuda/ops/precond_kda/__init__.py similarity index 57% rename from build/torch-cuda/modules/conv/cp/__init__.py rename to build/torch-cuda/ops/precond_kda/__init__.py index 97e899f587bcd72246487b736d919e6b7ee8b28a..f9a6fb53faf5100ea72ad2008060656b71b8458f 100644 --- a/build/torch-cuda/modules/conv/cp/__init__.py +++ b/build/torch-cuda/ops/precond_kda/__init__.py @@ -5,9 +5,12 @@ # For a list of all contributors, visit: # https://github.com/fla-org/flash-linear-attention/graphs/contributors -from .ops import CausalConv1dFunctionCP, causal_conv1d_cp +from .chunk import chunk_precond_kda +from .fused_recurrent import fused_recurrent_precond_kda +from .naive import naive_recurrent_precond_kda __all__ = [ - 'CausalConv1dFunctionCP', - 'causal_conv1d_cp', + "chunk_precond_kda", + "fused_recurrent_precond_kda", + "naive_recurrent_precond_kda", ] diff --git a/build/torch-cuda/ops/precond_kda/chunk.py b/build/torch-cuda/ops/precond_kda/chunk.py new file mode 100644 index 0000000000000000000000000000000000000000..2095b12ac1ed359b1af1f5e4750b8b7b547b1226 --- /dev/null +++ b/build/torch-cuda/ops/precond_kda/chunk.py @@ -0,0 +1,809 @@ +# 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 + +# Related files are modified and supported by the Moonshot AI Team + +import torch + +from ...modules.norm.l2norm import l2norm_bwd, l2norm_fwd +from ...ops.atk.chunk_atk_bwd import chunk_atk_bwd +from ...ops.atk.chunk_atk_fwd import chunk_atk_fwd, recompute_atk_fwd +from ...ops.common.chunk_delta_h import chunk_gated_delta_rule_bwd_dhu, chunk_gated_delta_rule_fwd_h +from ...ops.cp import FLACPContext +from ...ops.gla.chunk import chunk_gla_fwd_o_gk +from ...ops.kda.gate import kda_gate_bwd, kda_gate_chunk_cumsum +from ...ops.precond_kda.chunk_bwd import chunk_precond_kda_bwd_dAv, chunk_precond_kda_bwd_wy_dqkg +from ...ops.precond_kda.chunk_intra import chunk_precond_kda_bwd_intra, chunk_precond_kda_fwd_intra +from ...ops.precond_kda.wy_fast import recompute_w_u_fwd +from ...ops.utils import chunk_local_cumsum +from ...ops.utils.constant import RCP_LN2 +from ...utils import autocast_custom_bwd, autocast_custom_fwd, input_guard + + +def chunk_precond_kda_fwd( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + g: torch.Tensor, # Diagonal gate for KDA [B, T, H, K] - ALREADY cumsummed + g_atk: torch.Tensor, # Scalar gate for ATK [B, T, H] + beta_atk: torch.Tensor, # Beta for ATK [B, T, H] + beta: torch.Tensor, # Beta for KDA [B, T, H] + scale: float, + initial_state: torch.Tensor, + initial_A_state: torch.Tensor, + output_final_state: bool, + chunk_size: int = 64, + cu_seqlens: torch.LongTensor | None = None, + cp_context: FLACPContext | None = None, + chunk_indices: torch.LongTensor | None = None, + solve_tril_precision: str | None = None, + safe_gate: bool = False, + x: float = 1.5, + eps: float = 1e-6, + log_atk_scale: torch.Tensor = None, + transpose_state_layout: bool = False, +): + """ + Forward pass for preconditioned KDA. + + Args: + q: [B, T, H, K] - queries + k: [B, T, H, K] - keys + v: [B, T, H, V] - values + g: [B, T, H, K] - diagonal gate for KDA (chunk-local cumsum, log2 space) + g_atk: [B, T, H] - scalar gate for ATK (log space) + beta_atk: [B, T, H] - beta for ATK + beta: [B, T, H] - beta for KDA + scale: attention scale + initial_state: [B, H, K, V] or None + output_final_state: bool + x: range parameter for symmetric squash. M in [1/x, x] + eps: epsilon for log safety + log_atk_scale: per-head log-space center [H] + + Returns: + o, Aqk, Akk, final_state, at + """ + B, T, H, K = k.shape + + # Step 1: ATK preconditioning with symmetric fast squash + # Only need k_precond and at (final state); ac/a_atk/sa_atk recomputed in backward + k_precond, at = chunk_atk_fwd( + k=k, + beta=beta_atk, + log_g=g_atk, + chunk_size=chunk_size, + initial_A_state=initial_A_state, + output_final_state=output_final_state, + cu_seqlens=cu_seqlens, + x=x, + eps=eps, + log_atk_scale=log_atk_scale, + ) + + # Step 2: KDA intra-chunk computation (asymmetric) + # Aqk = q @ k_precond^T, Akk = k @ k_precond^T + w, u, kg, Aqk, Akk = chunk_precond_kda_fwd_intra( + q=q, + k=k, # Original k for Akk row side, w computation + k_precond=k_precond, # k_precond for column side, kg + v=v, + gk=g, + beta=beta, + scale=scale, + cu_seqlens=cu_seqlens, + chunk_size=chunk_size, + chunk_indices=chunk_indices, + solve_tril_precision=solve_tril_precision, + safe_gate=safe_gate, + ) + + # Step 3: Hidden state update (uses kg which is k_precond gated) + h, v_new, final_state = chunk_gated_delta_rule_fwd_h( + k=kg, # k_precond gated for write + w=w, + u=u, + gk=g, + initial_state=initial_state, + output_final_state=output_final_state, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + state_v_first=transpose_state_layout, + ) + + # Step 4: Output computation + o = chunk_gla_fwd_o_gk( + q=q, + v=v_new, + g=g, + A=Aqk, + h=h, + scale=scale, + cu_seqlens=cu_seqlens, + chunk_size=chunk_size, + chunk_indices=chunk_indices, + state_v_first=transpose_state_layout, + ) + + return o, Aqk, Akk, final_state, at, w, u, kg, v_new, h + + +def chunk_precond_kda_bwd( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + g: torch.Tensor, # Gate cumsum (log2 space) + g_atk: torch.Tensor, # Raw ATK gate + beta_atk: torch.Tensor, + beta: torch.Tensor, + Aqk: torch.Tensor, + Akk: torch.Tensor, + scale: float, + initial_state: torch.Tensor, + initial_A_state: torch.Tensor, + do: torch.Tensor, + dht: torch.Tensor, + cu_seqlens: torch.LongTensor | None = None, + cp_context: FLACPContext | None = None, + chunk_indices: torch.LongTensor | None = None, + chunk_size: int = 64, + x: float = 1.5, + eps: float = 1e-6, + log_atk_scale: torch.Tensor = None, + transpose_state_layout: bool = False, + safe_gate: bool = False, + disable_recompute: bool = False, + w: torch.Tensor | None = None, + u: torch.Tensor | None = None, + kg: torch.Tensor | None = None, + v_new: torch.Tensor | None = None, + h: torch.Tensor | None = None, + dat: torch.Tensor | None = None, +): + """ + Backward pass for preconditioned KDA with symmetric fast squash preconditioning. + + Returns: + dq, dk, dv, dg, dg_atk, dbeta_atk, dbeta, d_log_atk_scale, dh0 + """ + # Recompute ATK forward intermediates (k_precond, ac, a_atk, sa_atk) + k_precond, ac, a_atk, sa_atk = recompute_atk_fwd( + k=k, + beta=beta_atk, + log_g=g_atk, + chunk_size=chunk_size, + initial_A_state=initial_A_state, + cu_seqlens=cu_seqlens, + x=x, + eps=eps, + log_atk_scale=log_atk_scale, + ) + + if not disable_recompute: + # Step 1: Recompute WY representation (asymmetric) + w, u, qg, kg = recompute_w_u_fwd( + k=k, + k_precond=k_precond, + v=v, + beta=beta, + A=Akk, + gk=g, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + q=q, + output_qg=True, + output_kg=True, + ) + + # Step 2: Recompute hidden state (uses kg = gated k_precond) + h, v_new, _ = chunk_gated_delta_rule_fwd_h( + k=kg, + w=w, + u=u, + gk=g, + initial_state=initial_state, + output_final_state=False, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + state_v_first=transpose_state_layout, + ) + else: + # Intermediates (w, u, kg, v_new, h) saved from forward; only need qg + _, _, qg, _ = recompute_w_u_fwd( + k=k, + k_precond=k_precond, + v=v, + beta=beta, + A=Akk, + gk=g, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + q=q, + output_qg=True, + output_kg=False, + ) + + BT = chunk_size + + # Step 3: dAqk and local dv (matching KDA's chunk_kda_bwd_dAv) + dAqk, dv = chunk_precond_kda_bwd_dAv( + q=q, + k=k, + v=v_new, + do=do, + A=Aqk, + scale=scale, + cu_seqlens=cu_seqlens, + chunk_size=BT, + chunk_indices=chunk_indices, + ) + + # Step 4: Backward through hidden state + dh, dh0, dv = chunk_gated_delta_rule_bwd_dhu( + q=qg, + k=kg, + w=w, + gk=g, + h0=initial_state, + dht=dht, + do=do, + dv=dv, + scale=scale, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + state_v_first=transpose_state_layout, + ) + + # Step 5+6: Inter-chunk + WY backward (matching KDA's chunk_kda_bwd_wy_dqkg_fused) + dq, dk, dkg, dv, dbeta, dg, dAkk = chunk_precond_kda_bwd_wy_dqkg( + q=q, + k=k, + k_precond=k_precond, + v=v, + v_new=v_new, + g=g, + beta=beta, + A=Akk, + h=h, + do=do, + dh=dh, + dv=dv, + scale=scale, + cu_seqlens=cu_seqlens, + chunk_size=BT, + chunk_indices=chunk_indices, + transpose_state_layout=transpose_state_layout, + ) + + # Step 7: Asymmetric intra backward + dq2, dk_intra, dk_precond2, dbeta2, dg3 = chunk_precond_kda_bwd_intra( + q=q, + k=k, + k_precond=k_precond, + g=g, + beta=beta, + dAqk=dAqk, + dAkk=dAkk, + dq=dq, + dk=dk, + dk_precond=dkg, + db=dbeta, + dg=dg, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + chunk_size=chunk_size, + safe_gate=safe_gate, + ) + + dk_precond_total = dk_precond2 + dk = dk_intra + + dk_atk, dbeta_atk_grad, dg_atk, d_log_atk_scale, dh0_atk = chunk_atk_bwd( + k=k, + g_raw=g_atk, + beta=beta_atk, + ac=ac, + a=a_atk, + sa=sa_atk, + dk_precond=dk_precond_total, + initial_A_state=initial_A_state, + cu_seqlens=cu_seqlens, + x=x, + eps=eps, + log_atk_scale=log_atk_scale, + dat=dat, + ) + + # Combine dk gradients + dk_total = dk + dk_atk + + # dg3 from asymmetric intra backward already has reverse cumsum applied + dg_total = dg3 + + return dq2, dk_total, dv, dg_total, dg_atk, dbeta_atk_grad, dbeta2, d_log_atk_scale, dh0, dh0_atk + + +class ChunkPrecondKDAFunction(torch.autograd.Function): + + @staticmethod + @input_guard + @autocast_custom_fwd + def forward( + ctx, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + g: torch.Tensor, + g_atk: torch.Tensor, + beta_atk: torch.Tensor, + beta: torch.Tensor, + scale: float, + initial_state: torch.Tensor, + initial_A_state: torch.Tensor, + output_final_state: bool, + use_qk_l2norm_in_kernel: bool = False, + use_gate_in_kernel: bool = False, + A_log: torch.Tensor | None = None, + dt_bias: torch.Tensor | None = None, + cu_seqlens: torch.LongTensor | None = None, + cu_seqlens_cpu: torch.LongTensor | None = None, + chunk_indices: torch.LongTensor | None = None, + safe_gate: bool = False, + cp_context: FLACPContext | None = None, + transpose_state_layout: bool = False, + x: float = 1.5, + eps: float = 1e-6, + log_atk_scale: torch.Tensor = None, + lower_bound: float | None = None, + solve_tril_precision: str | None = None, + disable_recompute: bool = False, + return_intermediate_states: bool = False, + ): + chunk_size = 64 + + # Apply L2 normalization + q_rstd, k_rstd = None, None + if use_qk_l2norm_in_kernel: + q, q_rstd = l2norm_fwd(q) + k, k_rstd = l2norm_fwd(k) + + # Compute gate + cumsum + g_org = g + if use_gate_in_kernel: + g = kda_gate_chunk_cumsum( + g=g_org, + A_log=A_log, + chunk_size=chunk_size, + scale=RCP_LN2, + dt_bias=dt_bias, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + lower_bound=lower_bound, + ) + else: + g = chunk_local_cumsum( + g=g, + chunk_size=chunk_size, + scale=RCP_LN2, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + ) + + o, Aqk, Akk, final_state, at, w, u, kg, v_new, h = chunk_precond_kda_fwd( + q=q, + k=k, + v=v, + g=g, + g_atk=g_atk, + beta_atk=beta_atk, + beta=beta, + scale=scale, + initial_state=initial_state, + initial_A_state=initial_A_state, + output_final_state=output_final_state, + chunk_size=chunk_size, + cu_seqlens=cu_seqlens, + cp_context=cp_context, + chunk_indices=chunk_indices, + solve_tril_precision=solve_tril_precision, + safe_gate=safe_gate, + x=x, + eps=eps, + log_atk_scale=log_atk_scale, + transpose_state_layout=transpose_state_layout, + ) + + if return_intermediate_states: + assert torch.is_inference_mode_enabled(), "return_intermediate_states is only allowed in inference mode" + assert disable_recompute is False, "return_intermediate_states must be used with disable_recompute=False" + return o.to(q.dtype), final_state, at, h + + # Don't save computed g when use_gate_in_kernel (will recompute in backward) + if use_gate_in_kernel: + g = None + + # When disable_recompute=False (default), delete intermediates to save memory + if not disable_recompute: + w, u, kg, v_new, h = None, None, None, None, None + + ctx.save_for_backward( + q, q_rstd, k, k_rstd, v, g, g_org, g_atk, beta_atk, beta, + Aqk, Akk, initial_state, + A_log, dt_bias, log_atk_scale, cu_seqlens, chunk_indices, + w, u, kg, v_new, h + ) + ctx.scale = scale + ctx.use_qk_l2norm_in_kernel = use_qk_l2norm_in_kernel + ctx.use_gate_in_kernel = use_gate_in_kernel + ctx.safe_gate = safe_gate + ctx.chunk_size = chunk_size + ctx.initial_A_state = initial_A_state + ctx.cp_context = cp_context + ctx.transpose_state_layout = transpose_state_layout + ctx.x = x + ctx.eps = eps + ctx.lower_bound = lower_bound + ctx.disable_recompute = disable_recompute + + return o.to(q.dtype), final_state, at + + @staticmethod + @input_guard + @autocast_custom_bwd + def backward(ctx, do, dht, dat): + (q, q_rstd, k, k_rstd, v, g, g_org, g_atk, beta_atk, beta, + Aqk, Akk, initial_state, + A_log, dt_bias, log_atk_scale, cu_seqlens, chunk_indices, + w, u, kg, v_new, h) = ctx.saved_tensors + + # Recompute g (cumsummed) if use_gate_in_kernel was used + if ctx.use_gate_in_kernel: + g = kda_gate_chunk_cumsum( + g=g_org, + A_log=A_log, + chunk_size=ctx.chunk_size, + scale=RCP_LN2, + dt_bias=dt_bias, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + lower_bound=ctx.lower_bound, + ) + + dq, dk, dv, dg, dg_atk, dbeta_atk, dbeta, d_log_atk_scale, dh0, dh0_atk = chunk_precond_kda_bwd( + q=q, + k=k, + v=v, + g=g, + g_atk=g_atk, + beta_atk=beta_atk, + beta=beta, + Aqk=Aqk, + Akk=Akk, + scale=ctx.scale, + initial_state=initial_state, + initial_A_state=ctx.initial_A_state, + do=do, + dht=dht, + cu_seqlens=cu_seqlens, + cp_context=ctx.cp_context, + chunk_indices=chunk_indices, + chunk_size=ctx.chunk_size, + x=ctx.x, + eps=ctx.eps, + log_atk_scale=log_atk_scale, + transpose_state_layout=ctx.transpose_state_layout, + safe_gate=ctx.safe_gate, + disable_recompute=ctx.disable_recompute, + w=w, + u=u, + kg=kg, + v_new=v_new, + h=h, + dat=dat, + ) + + if ctx.use_qk_l2norm_in_kernel: + dq = l2norm_bwd(q, q_rstd, dq) + dk = l2norm_bwd(k, k_rstd, dk) + + # Compute gradients for A_log and dt_bias if use_gate_in_kernel + dA_log, ddt_bias = None, None + if ctx.use_gate_in_kernel: + dg, dA_log, ddt_bias = kda_gate_bwd( + g=g_org, + A_log=A_log, + dt_bias=dt_bias, + dyg=dg, + lower_bound=ctx.lower_bound, + ) + dA_log = dA_log.to(A_log) + if dt_bias is not None: + ddt_bias = ddt_bias.to(dt_bias) + + # Format d_log_atk_scale for return + if log_atk_scale is not None and d_log_atk_scale is not None: + d_log_atk_scale = d_log_atk_scale.to(log_atk_scale) + else: + d_log_atk_scale = None + + return ( + dq.to(q), + dk.to(k), + dv.to(v), + dg.to(g_org), + dg_atk.to(g_atk), + dbeta_atk.to(beta_atk), + dbeta.to(beta), + None, # scale + dh0, # initial_state + dh0_atk, # initial_A_state + None, # output_final_state + None, # use_qk_l2norm_in_kernel + None, # use_gate_in_kernel + dA_log, # A_log + ddt_bias, # dt_bias + None, # cu_seqlens + None, # cu_seqlens_cpu + None, # chunk_indices + None, # safe_gate + None, # cp_context + None, # transpose_state_layout + None, # x + None, # eps + d_log_atk_scale, # log_atk_scale + None, # lower_bound + None, # solve_tril_precision + None, # disable_recompute + None, # return_intermediate_states + ) + + +@torch.compiler.disable +def chunk_precond_kda( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + g: torch.Tensor, + g_atk: torch.Tensor, + beta_atk: torch.Tensor, + beta: torch.Tensor, + scale: float = None, + initial_state: torch.Tensor = None, + initial_A_state: torch.Tensor = None, + output_final_state: bool = False, + use_gate_in_kernel: bool = False, + safe_gate: bool = False, + lower_bound: float | None = None, + cu_seqlens: torch.LongTensor | None = None, + cu_seqlens_cpu: torch.LongTensor | None = None, + chunk_indices: torch.LongTensor | None = None, + cp_context: FLACPContext | None = None, + transpose_state_layout: bool = False, + x: float = 1.5, + eps: float = 1e-6, + log_atk_scale: torch.Tensor = None, + solve_tril_precision: str | None = None, + disable_recompute: bool = False, + return_intermediate_states: bool = False, + **kwargs, +): + r""" + Args: + q (torch.Tensor): + queries of shape `[B, T, H, K]`. + k (torch.Tensor): + keys of shape `[B, T, H, K]`. + v (torch.Tensor): + values of shape `[B, T, H, V]`. + g (torch.Tensor): + (forget) gating tensor (in log space!) of shape `[B, T, H, K]`. + g_atk (torch.Tensor): + ATK gates (decays) in log space of shape `[B, T, H]`. + beta_atk (torch.Tensor): + ATK betas of shape `[B, T, H]`. + beta (torch.Tensor): + betas of shape `[B, T, H]`. + scale (Optional[float]): + Scale factor for the attention scores. + If not provided, it will default to `1 / sqrt(K)`. Default: `None`. + initial_state (Optional[torch.Tensor]): + Initial state of shape `[N, H, K, V]` for `N` input sequences. + For equal-length input sequences, `N` equals the batch size `B`. + Default: `None`. + initial_A_state (Optional[torch.Tensor]): + Initial ATK diagonal state of shape `[N, H, K]`. Default: `None`. + output_final_state (Optional[bool]): + Whether to output the final state of shape `[N, H, K, V]`. Default: `False`. + use_gate_in_kernel (bool): + Whether to compute the log-space KDA decay internally. + - If `True`: + The passed `g` acts as the raw input for `-exp(A_log).view(H, -1) * softplus(g + dt_bias.view(H, K))`. + Note that as part of the input arguments, + `A_log` (shape `[H]`) and the optional `dt_bias` (shape `[H * K]`) should be provided. + - If `False`, `g` is expected to be the pre-computed decay value. + Default: `False`. + safe_gate (bool): + Whether the kernel can assume the gate values (in log space) are in a safe range + and use M=16 TensorCore acceleration for higher throughput. + The safe range is ``[lower_bound, 0)``. With the default ``lower_bound=-5``, + the per-step decay factor ``exp(g)`` is bounded in ``[exp(-5), 1) ≈ [0.0067, 1)``, + meaning each step retains at least ~0.67% of the state -- a negligible loss that + has minimal impact on model quality while enabling significant speedup. + Requires ``lower_bound`` to be set. Default: ``False``. + lower_bound (Optional[float]): + Lower bound for the forget gate (in log space) when ``use_gate_in_kernel=True``. + Changes the gate activation from ``-exp(A_log) * softplus(g + dt_bias)`` + to ``lower_bound * sigmoid(exp(A_log) * (g + dt_bias))``, + which naturally clamps the output to ``[lower_bound, 0)``. + Recommended value: ``-5`` (i.e., ``exp(-5) ≈ 0.0067``). Default: ``None``. + cu_seqlens (torch.LongTensor): + Cumulative sequence lengths of shape `[N+1]` used for variable-length training, + consistent with the FlashAttention API. + cu_seqlens_cpu (torch.LongTensor): + Cumulative sequence lengths of shape `[N+1]` used for variable-length training, + consistent with the FlashAttention API. + disable_recompute (bool): + Whether to disable gradient recomputation in the kernel. When `True`, the kernel + will save all intermediate activations for backward pass, which is beneficial + for training small models at the cost of increased memory usage. Default: `False`. + return_intermediate_states (bool): + If True, returns intermediate state `h` for inference scenarios (e.g., vLLM). + Must be used within `torch.inference_mode()` and will return a 4-tuple instead of 3-tuple. + This is not intended for training as it bypasses autograd. Default: `False`. + cp_context (Optional[FLACPContext]): + Context parallelism is not yet supported for the preconditioned path; + passing a context raises `NotImplementedError`. Default: `None`. + transpose_state_layout (Optional[bool]): + Whether to use the transposed state layout for the hidden state. + Default: `False`. + x (float): + Squash range parameter. Default: `1.5`. + eps (float): + Epsilon for numerical stability in the squash function. Default: `1e-6`. + log_atk_scale (Optional[torch.Tensor]): + Per-head log-space center of shape `[H]`. Default: `None`. + + Returns: + - Normal mode (return_intermediate_states=False): A tuple (o, final_state, A_state) + o (torch.Tensor): + Outputs of shape `[B, T, H, V]`. + final_state (torch.Tensor): + Final state of shape `[N, H, K, V]` if `output_final_state=True` else `None`. + A_state (torch.Tensor): + Final ATK diagonal state of shape `[N, H, K]` if `output_final_state=True` else `None`. + - Inference mode (return_intermediate_states=True): A tuple (o, final_state, A_state, h) + o (torch.Tensor): + Outputs of shape `[B, T, H, V]`. + final_state (torch.Tensor): + Final state of shape `[N, H, K, V]` if `output_final_state=True` else `None`. + A_state (torch.Tensor): + Final ATK diagonal state of shape `[N, H, K]` if `output_final_state=True` else `None`. + h (torch.Tensor): + Intermediate states of shape `[B, NT, H, K, V]` and dtype `bfloat16` for caching or further processing. + + Examples:: + >>> import torch + >>> import torch.nn.functional as F + >>> from einops import rearrange + >>> from fla.ops.precond_kda import chunk_precond_kda + # inputs with equal lengths + >>> B, T, H, K, V = 4, 2048, 4, 512, 512 + >>> q = torch.randn(B, T, H, K, dtype=torch.bfloat16, device='cuda') + >>> k = torch.randn(B, T, H, K, dtype=torch.bfloat16, device='cuda') + >>> v = torch.randn(B, T, H, V, dtype=torch.bfloat16, device='cuda') + >>> beta = torch.rand(B, T, H, dtype=torch.bfloat16, device='cuda') + >>> beta_atk = torch.rand(B, T, H, dtype=torch.bfloat16, device='cuda') + >>> g = torch.rand(B, T, H, K, dtype=torch.bfloat16, device='cuda') + >>> g_atk = F.logsigmoid(torch.rand(B, T, H, dtype=torch.bfloat16, device='cuda')) + >>> h0 = torch.randn(B, H, K, V, dtype=torch.bfloat16, device='cuda') + >>> A_log = torch.randn(H, dtype=torch.float32, device='cuda') + >>> dt_bias = torch.randn(H * K, dtype=torch.float32, device='cuda') + >>> o, ht, at = chunk_precond_kda( + q, k, v, g, g_atk, beta_atk, beta, + A_log=A_log, + dt_bias=dt_bias, + use_gate_in_kernel=True, + initial_state=h0, + output_final_state=True + ) + # for variable-length inputs, the batch size `B` is expected to be 1 and `cu_seqlens` is required + >>> q, k, v, beta, beta_atk, g, g_atk = map(lambda x: rearrange(x, 'b t ... -> 1 (b t) ...'), (q, k, v, beta, beta_atk, g, g_atk)) + # for a batch with 4 sequences, `cu_seqlens` with 5 start/end positions are expected + >>> cu_seqlens = q.new_tensor([0, 2048, 4096, 6144, 8192], dtype=torch.long) + >>> o, ht, at = chunk_precond_kda( + q, k, v, g, g_atk, beta_atk, beta, + A_log=A_log, + dt_bias=dt_bias, + use_gate_in_kernel=True, + initial_state=h0, + output_final_state=True, + cu_seqlens=cu_seqlens + ) + """ + + if cp_context is not None: + raise NotImplementedError( + "Context parallelism is not yet supported for chunk_precond_kda: " + "the preconditioned path lacks the CP state compression/expansion plumbing. " + "Please run without `cp_context`." + ) + + if cu_seqlens is not None: + if q.shape[0] != 1: + raise ValueError( + f"The batch size is expected to be 1 rather than {q.shape[0]} when using `cu_seqlens`." + f"Please flatten variable-length inputs before processing.", + ) + if initial_state is not None and initial_state.shape[0] != len(cu_seqlens) - 1: + raise ValueError( + f"The number of initial states is expected to be equal to the number of input sequences, " + f"i.e., {len(cu_seqlens) - 1} rather than {initial_state.shape[0]}.", + ) + if initial_state is not None: + assert initial_state.dtype == torch.float32, "initial_state must be in float32." + + if safe_gate and use_gate_in_kernel: + if lower_bound is None: + raise ValueError("`lower_bound` must be specified when `safe_gate=True` and `use_gate_in_kernel=True`.") + if not (-5 <= lower_bound < 0): + raise ValueError(f"`lower_bound` must be in the safe range [-5, 0), got {lower_bound}.") + + assert q.shape == k.shape == g.shape, "q, k, g must have the same shape." + assert k.shape[-1] <= 256, "Currently we only support key headdim <=256 for KDA :-(" + assert beta.shape == q.shape[:3], "beta must be of shape (batch size, seq len, num of head)." + assert v.shape == (*q.shape[:3], v.shape[-1]), "v must be of shape (batch size, seq len, num of head, head dim)." + + if scale is None: + scale = k.shape[-1] ** -0.5 + + # Prepare log_atk_scale with default if needed + H = k.shape[2] + if log_atk_scale is None: + log_atk_scale = torch.full((H,), -0.2, device=k.device, dtype=torch.float32) + + A_log, dt_bias = None, None + if use_gate_in_kernel: + assert "A_log" in kwargs, "A_log must be provided when use_gate_in_kernel=True." + A_log, dt_bias = kwargs["A_log"], kwargs.get("dt_bias") + + results = ChunkPrecondKDAFunction.apply( + q, + k, + v, + g, + g_atk, + beta_atk, + beta, + scale, + initial_state, + initial_A_state, + output_final_state, + True, # use_qk_l2norm_in_kernel (always True) + use_gate_in_kernel, + A_log, + dt_bias, + cu_seqlens, + cu_seqlens_cpu, + chunk_indices, + safe_gate, + cp_context, + transpose_state_layout, + x, + eps, + log_atk_scale, + lower_bound, + solve_tril_precision, + disable_recompute, + return_intermediate_states, + ) + + if return_intermediate_states: + # returns (o, final_state, at, h) + return results + + o, h_final, a_final = results + return o, h_final, a_final diff --git a/build/torch-cuda/ops/precond_kda/chunk_bwd.py b/build/torch-cuda/ops/precond_kda/chunk_bwd.py new file mode 100644 index 0000000000000000000000000000000000000000..8c5f73ac09f7baa04a9607fe4f42cf476f87c5de --- /dev/null +++ b/build/torch-cuda/ops/precond_kda/chunk_bwd.py @@ -0,0 +1,467 @@ +# 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 +import triton.language as tl + +from ...ops.utils import prepare_chunk_indices +from ...ops.utils.op import exp2 +from ...utils import IS_NVIDIA_HOPPER, autotune_cache_kwargs, check_shared_mem + +BK_LIST = [32, 64] if check_shared_mem() else [16, 32] +BV_LIST = [64, 128] if check_shared_mem('ampere') else [16, 32] +NUM_WARPS = [2, 4] if IS_NVIDIA_HOPPER else [2, 4, 8] + + +# ============================================================================== +# dAv kernel: compute dA and dv from inter-chunk backward +# Matches KDA's chunk_kda_bwd_kernel_dAv +# ============================================================================== + +@triton.heuristics({ + 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None, +}) +@triton.autotune( + configs=[ + triton.Config({}, num_warps=num_warps, num_stages=num_stages) + for num_warps in NUM_WARPS + for num_stages in [2, 3, 4] + ], + key=['H', 'K', 'V', 'BT', 'BK', 'BV'], + **autotune_cache_kwargs, +) +@triton.jit(do_not_specialize=['T']) +def chunk_precond_kda_bwd_kernel_dAv( + q, + k, + v, + A, + do, + dv, + dA, + cu_seqlens, + chunk_indices, + scale, + T, + H: tl.constexpr, + K: tl.constexpr, + V: tl.constexpr, + BT: tl.constexpr, + BK: tl.constexpr, + BV: tl.constexpr, + IS_VARLEN: tl.constexpr, +): + i_t, i_bh = tl.program_id(0).to(tl.int64), tl.program_id(1).to(tl.int64) + i_b, i_h = i_bh // H, i_bh % H + 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.int64) + bos, eos = tl.load(cu_seqlens + i_n).to(tl.int64), tl.load(cu_seqlens + i_n + 1).to(tl.int64) + T = eos - bos + else: + bos, eos = i_b * T, i_b * T + T + + # offset calculation + q += (bos * H + i_h) * K + k += (bos * H + i_h) * K + v += (bos * H + i_h) * V + do += (bos * H + i_h) * V + dv += (bos * H + i_h) * V + dA += (bos * H + i_h) * BT + + o_t = i_t * BT + tl.arange(0, BT) + m_t = o_t < T + o_r = tl.arange(0, BT) + p_A = A + (bos * H + i_h) * BT + o_r[:, None] + o_t[None, :] * (H*BT) + b_A = tl.load(p_A, mask=(o_r[:, None] < BT) & m_t[None, :], other=0.0) + + m_A = (o_t[:, None] <= o_t[None, :]) & (m_t[:, None] & m_t) + b_A = tl.where(m_A, b_A, 0).to(do.dtype.element_ty) + + b_dA = tl.zeros([BT, BT], dtype=tl.float32) + for i_v in range(tl.cdiv(V, BV)): + o_v = i_v * BV + tl.arange(0, BV) + m_v = o_v < V + m_tv = m_t[:, None] & m_v[None, :] + # [BV, BT] transposed + p_v = v + o_v[:, None] + o_t[None, :] * (H*V) + p_do = do + o_t[:, None] * (H*V) + o_v[None, :] + p_dv = dv + o_t[:, None] * (H*V) + o_v[None, :] + b_v = tl.load(p_v, mask=m_v[:, None] & m_t[None, :], other=0.0) + # [BT, BV] + b_do = tl.load(p_do, mask=m_tv, other=0.0) + # [BT, BT] + b_dA = tl.dot(b_do, b_v, b_dA) + # [BT, BV] + b_dv = tl.dot(b_A.to(b_do.dtype), b_do) + tl.store(p_dv, b_dv.to(dv.dtype.element_ty), mask=m_tv) + + o_A = tl.arange(0, BT) + p_dA = dA + o_t[:, None] * (H*BT) + o_A[None, :] + b_dA = tl.where(o_t[:, None] >= o_t, b_dA * scale, 0.) + tl.store(p_dA, b_dA.to(dA.dtype.element_ty), mask=m_t[:, None] & (o_A[None, :] < BT)) + + +def chunk_precond_kda_bwd_dAv( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + do: torch.Tensor, + A: torch.Tensor | None = None, + scale: float = None, + cu_seqlens: torch.LongTensor | None = None, + chunk_size: int = 64, + chunk_indices: torch.LongTensor | None = None, +) -> tuple[torch.Tensor, torch.Tensor]: + B, T, H, K, V = *k.shape, do.shape[-1] + BT = chunk_size + if chunk_indices is None and cu_seqlens is not None: + chunk_indices = prepare_chunk_indices(cu_seqlens, BT) + # H100 can have larger block size + if check_shared_mem('hopper', k.device.index): + CONST_TILING = 128 + elif check_shared_mem(): + CONST_TILING = 64 + else: + CONST_TILING = 32 + BK = min(max(triton.next_power_of_2(K), 16), CONST_TILING) + BV = min(max(triton.next_power_of_2(V), 16), CONST_TILING) + NT = triton.cdiv(T, BT) if cu_seqlens is None else len(chunk_indices) + + dA = v.new_empty(B, T, H, BT, dtype=torch.float) + dv = torch.empty_like(do) + grid = (NT, B * H) + chunk_precond_kda_bwd_kernel_dAv[grid]( + q=q, + k=k, + v=v, + A=A, + do=do, + dv=dv, + dA=dA, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + scale=scale, + T=T, + H=H, + K=K, + V=V, + BT=BT, + BK=BK, + BV=BV, + ) + return dA, dv + + +# ============================================================================== +# WY + inter-chunk backward kernel: compute dq, dk, dkg, dv, db, dg, dA +# Matches KDA's chunk_kda_bwd_kernel_wy_dqkg_fused with k/k_precond asymmetry +# ============================================================================== + +@triton.heuristics({ + 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None, +}) +@triton.autotune( + configs=[ + triton.Config({'BK': BK, 'BV': BV}, num_warps=num_warps, num_stages=num_stages) + for BK in BK_LIST + for BV in BV_LIST + for num_warps in NUM_WARPS + for num_stages in [2, 3, 4] + if not (IS_NVIDIA_HOPPER and BK == 32 and num_warps == 4) + ], + key=['BT', 'TRANSPOSE_STATE', 'K', 'V'], + **autotune_cache_kwargs, +) +@triton.jit(do_not_specialize=['T']) +def chunk_precond_kda_bwd_kernel_wy_dqkg( + q, + k, # original k (for WY backward) + k_precond, # preconditioned k (for inter backward) + v, + v_new, + g, + beta, + A, + h, + do, + dh, + dq, + dk, # dk from WY backward (original k) + dkg, # dkg from inter backward (k_precond) + dv, + dv2, + dg, + db, + dA, + cu_seqlens, + chunk_indices, + scale, + T, + H: tl.constexpr, + K: tl.constexpr, + V: tl.constexpr, + BT: tl.constexpr, + BK: tl.constexpr, + BV: tl.constexpr, + TRANSPOSE_STATE: tl.constexpr, + IS_VARLEN: tl.constexpr, +): + i_t, i_bh = tl.program_id(0).to(tl.int64), tl.program_id(1).to(tl.int64) + i_b, i_h = i_bh // H, i_bh % H + + if IS_VARLEN: + i_tg = i_t.to(tl.int64) + i_n, i_t = tl.load(chunk_indices + i_t * 2).to(tl.int32), tl.load(chunk_indices + i_t * 2 + 1).to(tl.int64) + bos, eos = tl.load(cu_seqlens + i_n).to(tl.int64), tl.load(cu_seqlens + i_n + 1).to(tl.int64) + T = (eos - bos).to(tl.int32) + NT = tl.cdiv(T, BT) + else: + NT = tl.cdiv(T, BT) + i_tg = (i_b * NT + i_t).to(tl.int64) + bos, eos = (i_b * T).to(tl.int64), (i_b * T + T).to(tl.int64) + + o_t = i_t * BT + tl.arange(0, BT) + m_t = o_t < T + m_last = (o_t == min(T, i_t * BT + BT) - 1) + + q += (bos * H + i_h) * K + k += (bos * H + i_h) * K + k_precond += (bos * H + i_h) * K + v += (bos * H + i_h) * V + v_new += (bos * H + i_h) * V + g += (bos * H + i_h) * K + beta += bos * H + i_h + A += (bos * H + i_h) * BT + h += (i_tg * H + i_h) * K*V + do += (bos * H + i_h) * V + dh += (i_tg * H + i_h) * K*V + dq += (bos * H + i_h) * K + dk += (bos * H + i_h) * K + dkg += (bos * H + i_h) * K + dv += (bos * H + i_h) * V + dv2 += (bos * H + i_h) * V + dg += (bos * H + i_h) * K + db += bos * H + i_h + dA += (bos * H + i_h) * BT + + b_beta = tl.load(beta + o_t*H, mask=m_t, other=0.0) + + o_r = tl.arange(0, BT) + p_A = A + o_r[:, None] + o_t[None, :] * (H * BT) + b_A = tl.load(p_A, mask=(o_r[:, None] < BT) & m_t[None, :], other=0.0) + + b_dA = tl.zeros([BT, BT], dtype=tl.float32) + b_db = tl.zeros([BT], dtype=tl.float32) + + for i_k in range(tl.cdiv(K, BK)): + o_k = i_k * BK + tl.arange(0, BK) + m_k = o_k < K + + m_tk = m_t[:, None] & m_k[None, :] + p_k = k + o_t[:, None] * (H*K) + o_k[None, :] + p_kp = k_precond + o_t[:, None] * (H*K) + o_k[None, :] + p_g = g + o_t[:, None] * (H*K) + o_k[None, :] + b_k = tl.load(p_k, mask=m_tk, other=0.0) + b_kp = tl.load(p_kp, mask=m_tk, other=0.0) + b_g = tl.load(p_g, mask=m_tk, other=0.0).to(tl.float32) + + p_gn = g + (min(T, i_t * BT + BT) - 1).to(tl.int64) * H*K + o_k + b_gn = tl.load(p_gn, mask=m_k, other=0).to(tl.float32) + + b_dq = tl.zeros([BT, BK], dtype=tl.float32) + b_dkg_raw = tl.zeros([BT, BK], dtype=tl.float32) + b_dw = tl.zeros([BT, BK], dtype=tl.float32) + b_dgk = tl.zeros([BK], dtype=tl.float32) + + for i_v in range(tl.cdiv(V, BV)): + o_v = i_v * BV + tl.arange(0, BV) + m_v = o_v < V + m_tv = m_t[:, None] & m_v[None, :] + m_vk = m_v[:, None] & m_k[None, :] + p_v_new = v_new + o_t[:, None] * (H*V) + o_v[None, :] + p_do = do + o_t[:, None] * (H*V) + o_v[None, :] + if TRANSPOSE_STATE: + p_h = h + o_v[:, None] * K + o_k[None, :] + p_dh = dh + o_v[:, None] * K + o_k[None, :] + else: + p_h = h + o_v[:, None] + o_k[None, :] * V + p_dh = dh + o_v[:, None] + o_k[None, :] * V + p_dv = dv + o_t[:, None] * (H*V) + o_v[None, :] + # [BT, BV] + b_v_new = tl.load(p_v_new, mask=m_tv, other=0.0) + b_do = tl.load(p_do, mask=m_tv, other=0.0) + # [BV, BK] + b_h = tl.load(p_h, mask=m_vk, other=0.0) + b_dh = tl.load(p_dh, mask=m_vk, other=0.0) + # [BT, BV] + b_dv = tl.load(p_dv, mask=m_tv, other=0.0) + + b_dgk += tl.sum(b_h * b_dh, axis=0) + b_dq = tl.dot(b_do, b_h.to(b_do.dtype), b_dq) + b_dkg_raw = tl.dot(b_v_new, b_dh.to(b_v_new.dtype), b_dkg_raw) + b_dw = tl.dot(b_dv.to(b_v_new.dtype), b_h.to(b_v_new.dtype), b_dw) + tl.debug_barrier() # DO NOT REMOVE THIS LINE! + if i_k == 0: + p_v = v + o_t[:, None] * (H*V) + o_v[None, :] + p_dv2 = dv2 + o_t[:, None] * (H*V) + o_v[None, :] + + b_v = tl.load(p_v, mask=m_tv, other=0.0) + + b_dA = tl.dot(b_dv, tl.trans(b_v), b_dA) + + b_dvb = tl.dot(b_A, b_dv) + b_dv2 = b_dvb * b_beta[:, None] + b_db += tl.sum(b_dvb * b_v, 1) + + tl.store(p_dv2, b_dv2.to(dv2.dtype.element_ty), mask=m_tv) + + b_gk_exp = exp2(b_g) + b_gb = b_gk_exp * b_beta[:, None] + b_dgk *= exp2(b_gn) + b_dq = b_dq * b_gk_exp * scale + + # WY backward: uses original k + b_kg_orig = b_k * b_gk_exp + + b_dw = -b_dw.to(b_A.dtype) + b_dA = tl.dot(b_dw, tl.trans(b_kg_orig.to(b_A.dtype)), b_dA) + + b_dkgb = tl.dot(b_A, b_dw) + b_db += tl.sum(b_dkgb * b_kg_orig, 1) + + # dk from WY backward (original k) + b_dk = b_dkgb * b_gb + + # Inter backward: uses k_precond + b_gn_g = tl.where(m_t[:, None], exp2(b_gn[None, :] - b_g), 0) + b_dkg = b_dkg_raw * b_gn_g + + b_kg_precond = b_kp * b_gn_g + b_kp_dkg = b_kg_precond * b_dkg_raw + b_dgk += tl.sum(b_kp_dkg, axis=0) + + p_q = q + o_t[:, None] * (H*K) + o_k[None, :] + b_q = tl.load(p_q, mask=m_tk, other=0.0) + + b_dg = (b_q * b_dq # inter query + - b_kp_dkg # inter write key + + m_last[:, None] * b_dgk # accumulated last position + + b_kg_orig * b_dkgb * b_beta[:, None]) # WY backward + + p_dq = dq + o_t[:, None] * (H*K) + o_k[None, :] + p_dk = dk + o_t[:, None] * (H*K) + o_k[None, :] + p_dkg = dkg + o_t[:, None] * (H*K) + o_k[None, :] + p_dg = dg + o_t[:, None] * (H*K) + o_k[None, :] + tl.store(p_dq, b_dq.to(dq.dtype.element_ty), mask=m_tk) + tl.store(p_dk, b_dk.to(dk.dtype.element_ty), mask=m_tk) + tl.store(p_dkg, b_dkg.to(dkg.dtype.element_ty), mask=m_tk) + tl.store(p_dg, b_dg.to(dg.dtype.element_ty), mask=m_tk) + + m_A = (o_t[:, None] > o_t[None, :]) & (m_t[:, None] & m_t) + b_dA = tl.where(m_A, b_dA * b_beta[None, :], 0) + b_dA = tl.dot(b_dA.to(b_A.dtype), b_A) + b_dA = tl.dot(b_A, b_dA.to(b_A.dtype)) + b_dA = tl.where(m_A, -b_dA, 0) + + o_A = tl.arange(0, BT) + p_dA = dA + o_t[:, None] * (H * BT) + o_A[None, :] + tl.store(p_dA, b_dA.to(dA.dtype.element_ty), mask=m_t[:, None] & (o_A[None, :] < BT)) + tl.store(db + o_t*H, b_db.to(db.dtype.element_ty), mask=m_t) + + +def chunk_precond_kda_bwd_wy_dqkg( + q: torch.Tensor, + k: torch.Tensor, + k_precond: torch.Tensor, + v: torch.Tensor, + v_new: torch.Tensor, + g: torch.Tensor, + beta: torch.Tensor, + A: torch.Tensor, + h: torch.Tensor, + do: torch.Tensor, + dh: torch.Tensor, + dv: torch.Tensor, + scale: float | None = None, + cu_seqlens: torch.LongTensor | None = None, + chunk_size: int = 64, + chunk_indices: torch.LongTensor | None = None, + transpose_state_layout: bool = False, +): + """ + Inter-chunk + WY backward for preconditioned KDA. + + Matches KDA's chunk_kda_bwd_wy_dqkg_fused but with k/k_precond asymmetry: + - WY backward uses original k (for w = A @ (k * beta * exp(gk))) + - Inter backward uses k_precond (for kg = k_precond * exp(gn - gk)) + + Args: + q: [B, T, H, K] - queries + k: [B, T, H, K] - original k (for WY backward) + k_precond: [B, T, H, K] - preconditioned k (for inter backward) + v: [B, T, H, V] - original v (for WY backward) + v_new: [B, T, H, V] - corrected v (for inter backward) + g: [B, T, H, K] - gate cumsum (log2 space) + beta: [B, T, H] - beta scaling + A: [B, T, H, BT] - WY inverse matrix + h: [NT, H, K, V] - per-chunk hidden states + do: [B, T, H, V] - output gradient + dh: [NT, H, K, V] - hidden state gradient + dv: [B, T, H, V] - dv from h backward (= du for WY) + + Returns: + dq, dk, dkg, dv, db, dg, dA + """ + B, T, H, K, V = *k.shape, v.shape[-1] + BT = chunk_size + + if chunk_indices is None and cu_seqlens is not None: + chunk_indices = prepare_chunk_indices(cu_seqlens, BT) + NT = triton.cdiv(T, BT) if cu_seqlens is None else len(chunk_indices) + + dq = torch.empty_like(q, dtype=torch.float) + dk = torch.empty_like(k, dtype=torch.float) + dkg = torch.empty_like(k_precond, dtype=torch.float) + dv2 = torch.empty_like(v) + dg = torch.empty_like(g, dtype=torch.float) + db = torch.empty_like(beta, dtype=torch.float) + dA = torch.empty_like(A, dtype=torch.float) + + grid = (NT, B * H) + chunk_precond_kda_bwd_kernel_wy_dqkg[grid]( + q=q, + k=k, + k_precond=k_precond, + v=v, + v_new=v_new, + g=g, + beta=beta, + A=A, + h=h, + do=do, + dh=dh, + dq=dq, + dk=dk, + dkg=dkg, + dv=dv, + dv2=dv2, + dg=dg, + db=db, + dA=dA, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + scale=scale, + T=T, + H=H, + K=K, + V=V, + BT=BT, + TRANSPOSE_STATE=transpose_state_layout, + ) + dv = dv2 + return dq, dk, dkg, dv, db, dg, dA diff --git a/build/torch-cuda/ops/precond_kda/chunk_intra.py b/build/torch-cuda/ops/precond_kda/chunk_intra.py new file mode 100644 index 0000000000000000000000000000000000000000..9b3767ef4c68c61cf1cf6f0a0551635761bbe206 --- /dev/null +++ b/build/torch-cuda/ops/precond_kda/chunk_intra.py @@ -0,0 +1,1027 @@ +# 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 +import triton.language as tl + +from ...ops.precond_kda.chunk_intra_token_parallel import chunk_precond_kda_fwd_intra_token_parallel +from ...ops.precond_kda.wy_fast import recompute_w_u_fwd +from ...ops.utils import chunk_local_cumsum, prepare_chunk_indices +from ...ops.utils.op import exp2, gather, unflatten_program_id +from ...utils import IS_GATHER_SUPPORTED, IS_TF32_SUPPORTED, autotune_cache_kwargs + +DEFAULT_SOLVE_TRIL_PRECISION = 'tf32x3' if IS_TF32_SUPPORTED else 'ieee' + + +################################################################################ +# Fused inter + solve_tril kernel: compute off-diagonal Akk and solve in one pass +################################################################################ + + +@triton.heuristics({ + 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None, + 'BK': lambda args: min(triton.next_power_of_2(args['K']), 64), +}) +@triton.autotune( + configs=[ + triton.Config({}, num_warps=num_warps) + for num_warps in [1, 2, 4] + ], + key=["H", "K", "BC"], + **autotune_cache_kwargs, +) +@triton.jit(do_not_specialize=['T']) +def chunk_precond_kda_fwd_kernel_inter_solve_fused( + q, + k, # Original k (for Akk row side) + k_precond, # Preconditioned k (for column side) + g, + beta, + Aqk, + Akk_diag, + Akk, + scale, + cu_seqlens, + chunk_indices, + T, + H: tl.constexpr, + K: tl.constexpr, + BT: tl.constexpr, + BC: tl.constexpr, + BK: tl.constexpr, + IS_VARLEN: tl.constexpr, + SOLVE_TRIL_DOT_PRECISION: tl.constexpr = 'tf32x3', + USE_SAFE_GATE: tl.constexpr = False, +): + """ + Fused kernel: compute inter-subchunk Akk + solve_tril in one pass. + Asymmetric version: Aqk = q @ k_precond^T, Akk = k @ k_precond^T + Prerequisite: token_parallel has already computed diagonal Akk blocks in Akk_diag. + + This kernel: + 1. Computes off-diagonal Aqk blocks -> writes to global + 2. Computes off-diagonal Akk blocks -> keeps in registers + 3. Loads diagonal Akk blocks from Akk_diag (fp32) + 4. Does forward substitution on diagonals (skipped when USE_SAFE_GATE) + 5. Computes merged Akk_inv + 6. Writes Akk_inv to Akk + """ + i_t, i_bh = tl.program_id(0).to(tl.int64), tl.program_id(1).to(tl.int64) + i_b, i_h = i_bh // H, i_bh % H + + 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.int64) + bos, eos = tl.load(cu_seqlens + i_n).to(tl.int64), tl.load(cu_seqlens + i_n + 1).to(tl.int64) + T = eos - bos + else: + bos, eos = i_b * T, i_b * T + T + + if i_t * BT >= T: + return + + i_tc0 = i_t * BT + i_tc1 = i_t * BT + BC + i_tc2 = i_t * BT + 2 * BC + i_tc3 = i_t * BT + 3 * BC + + q += (bos * H + i_h) * K + k += (bos * H + i_h) * K + k_precond += (bos * H + i_h) * K + g += (bos * H + i_h) * K + Aqk += (bos * H + i_h) * BT + Akk += (bos * H + i_h) * BT + Akk_diag += (bos * H + i_h) * BC + + o_i = tl.arange(0, BC) + o_c0 = i_tc0 + o_i + o_c1 = i_tc1 + o_i + o_c2 = i_tc2 + o_i + o_c3 = i_tc3 + o_i + m_tc0 = o_c0 < T + m_tc1 = o_c1 < T + m_tc2 = o_c2 < T + m_tc3 = o_c3 < T + + b_Aqk10 = tl.zeros([BC, BC], dtype=tl.float32) + b_Akk10 = tl.zeros([BC, BC], dtype=tl.float32) + + b_Aqk20 = tl.zeros([BC, BC], dtype=tl.float32) + b_Akk20 = tl.zeros([BC, BC], dtype=tl.float32) + b_Aqk21 = tl.zeros([BC, BC], dtype=tl.float32) + b_Akk21 = tl.zeros([BC, BC], dtype=tl.float32) + + b_Aqk30 = tl.zeros([BC, BC], dtype=tl.float32) + b_Akk30 = tl.zeros([BC, BC], dtype=tl.float32) + b_Aqk31 = tl.zeros([BC, BC], dtype=tl.float32) + b_Akk31 = tl.zeros([BC, BC], dtype=tl.float32) + b_Aqk32 = tl.zeros([BC, BC], dtype=tl.float32) + b_Akk32 = tl.zeros([BC, BC], dtype=tl.float32) + + ################################################################################ + # 1. off-diagonal blocks - ASYMMETRIC: column uses k_precond, row uses k + ################################################################################ + for i_k in range(tl.cdiv(K, BK)): + o_k = i_k * BK + tl.arange(0, BK) + m_k = o_k < K + + # Column side: load k_precond transposed (for Aqk and Akk column) + m_kc0 = m_k[:, None] & m_tc0[None, :] + p_kp0 = k_precond + o_k[:, None] + o_c0[None, :] * (H*K) + p_g0 = g + o_k[:, None] + o_c0[None, :] * (H*K) + b_kpt0 = tl.load(p_kp0, mask=m_kc0, other=0.0).to(tl.float32) # k_precond transposed + b_gt0 = tl.load(p_g0, mask=m_kc0, other=0.0).to(tl.float32) + + b_kpt1, b_gt1 = b_kpt0, b_gt0 + b_kpt2, b_gt2 = b_kpt0, b_gt0 + + if i_tc1 < T: + # Row side: q for Aqk, original k for Akk + m_c1k = m_tc1[:, None] & m_k[None, :] + p_q1 = q + o_c1[:, None] * (H*K) + o_k[None, :] + p_k1 = k + o_c1[:, None] * (H*K) + o_k[None, :] + p_kp1 = k_precond + o_c1[:, None] * (H*K) + o_k[None, :] + p_g1 = g + o_c1[:, None] * (H*K) + o_k[None, :] + # [BC, BK] + b_q1 = tl.load(p_q1, mask=m_c1k, other=0.0).to(tl.float32) + b_k1 = tl.load(p_k1, mask=m_c1k, other=0.0).to(tl.float32) # Original k for row + b_kp1 = tl.load(p_kp1, mask=m_c1k, other=0.0).to(tl.float32) # k_precond for column + b_g1 = tl.load(p_g1, mask=m_c1k, other=0.0).to(tl.float32) + # [BK, BC] + b_kpt1 = tl.trans(b_kp1) # k_precond transposed + b_gt1 = tl.trans(b_g1) + # [BK] + b_gn1 = tl.load(g + i_tc1 * H * K + o_k, mask=m_k, other=0).to(tl.float32) + # [BC, BK] + b_gqn1 = tl.where(m_tc1[:, None], exp2(b_g1 - b_gn1[None, :]), 0) + b_qg1 = b_q1 * b_gqn1 + b_kg1 = b_k1 * b_gqn1 # Original k for Akk row + # [BK, BC] + b_kpgt = b_kpt0 * exp2(b_gn1[:, None] - b_gt0) # k_precond for column + # [BC, BC] + b_Aqk10 = tl.dot(b_qg1, b_kpgt, b_Aqk10) + b_Akk10 = tl.dot(b_kg1, b_kpgt, b_Akk10) # Asymmetric: k @ k_precond^T + + if i_tc2 < T: + m_c2k = m_tc2[:, None] & m_k[None, :] + p_q2 = q + o_c2[:, None] * (H*K) + o_k[None, :] + p_k2 = k + o_c2[:, None] * (H*K) + o_k[None, :] + p_kp2 = k_precond + o_c2[:, None] * (H*K) + o_k[None, :] + p_g2 = g + o_c2[:, None] * (H*K) + o_k[None, :] + + b_q2 = tl.load(p_q2, mask=m_c2k, other=0.0).to(tl.float32) + b_k2 = tl.load(p_k2, mask=m_c2k, other=0.0).to(tl.float32) + b_kp2 = tl.load(p_kp2, mask=m_c2k, other=0.0).to(tl.float32) + b_g2 = tl.load(p_g2, mask=m_c2k, other=0.0).to(tl.float32) + b_kpt2 = tl.trans(b_kp2) + b_gt2 = tl.trans(b_g2) + + b_gn2 = tl.load(g + i_tc2 * H * K + o_k, mask=m_k, other=0).to(tl.float32) + b_gqn2 = tl.where(m_tc2[:, None], exp2(b_g2 - b_gn2[None, :]), 0) + b_qg2 = b_q2 * b_gqn2 + b_kg2 = b_k2 * b_gqn2 + b_kpgt = b_kpt0 * exp2(b_gn2[:, None] - b_gt0) + b_Aqk20 = tl.dot(b_qg2, b_kpgt, b_Aqk20) + b_Akk20 = tl.dot(b_kg2, b_kpgt, b_Akk20) + + b_kpgt = b_kpt1 * exp2(b_gn2[:, None] - b_gt1) + b_Aqk21 = tl.dot(b_qg2, b_kpgt, b_Aqk21) + b_Akk21 = tl.dot(b_kg2, b_kpgt, b_Akk21) + + if i_tc3 < T: + m_c3k = m_tc3[:, None] & m_k[None, :] + p_q3 = q + o_c3[:, None] * (H*K) + o_k[None, :] + p_k3 = k + o_c3[:, None] * (H*K) + o_k[None, :] + p_g3 = g + o_c3[:, None] * (H*K) + o_k[None, :] + b_q3 = tl.load(p_q3, mask=m_c3k, other=0.0).to(tl.float32) + b_k3 = tl.load(p_k3, mask=m_c3k, other=0.0).to(tl.float32) + b_g3 = tl.load(p_g3, mask=m_c3k, other=0.0).to(tl.float32) + + b_gn3 = tl.load(g + i_tc3 * H * K + o_k, mask=m_k, other=0).to(tl.float32) + b_gqn3 = tl.where(m_tc3[:, None], exp2(b_g3 - b_gn3[None, :]), 0) + b_qg3 = b_q3 * b_gqn3 + b_kg3 = b_k3 * b_gqn3 + b_kpgt = b_kpt0 * exp2(b_gn3[:, None] - b_gt0) + b_Aqk30 = tl.dot(b_qg3, b_kpgt, b_Aqk30) + b_Akk30 = tl.dot(b_kg3, b_kpgt, b_Akk30) + + b_kpgt = b_kpt1 * exp2(b_gn3[:, None] - b_gt1) + b_Aqk31 = tl.dot(b_qg3, b_kpgt, b_Aqk31) + b_Akk31 = tl.dot(b_kg3, b_kpgt, b_Akk31) + + b_kpgt = b_kpt2 * exp2(b_gn3[:, None] - b_gt2) + b_Aqk32 = tl.dot(b_qg3, b_kpgt, b_Aqk32) + b_Akk32 = tl.dot(b_kg3, b_kpgt, b_Akk32) + + ################################################################################ + # 2. save off-diagonal Aqk blocks and prepare Akk + ################################################################################ + if i_tc1 < T: + p_Aqk10 = Aqk + o_c1[:, None] * (H*BT) + o_i[None, :] + tl.store(p_Aqk10, (b_Aqk10 * scale).to(Aqk.dtype.element_ty), mask=m_tc1[:, None] & (o_i[None, :] < BC)) + + b_b1 = tl.load(beta + bos * H + i_h + o_c1*H, mask=m_tc1, other=0.0).to(tl.float32) + b_Akk10 = b_Akk10 * b_b1[:, None] + if i_tc2 < T: + p_Aqk20 = Aqk + o_c2[:, None] * (H*BT) + o_i[None, :] + p_Aqk21 = Aqk + o_c2[:, None] * (H*BT) + (o_i + BC)[None, :] + tl.store(p_Aqk20, (b_Aqk20 * scale).to(Aqk.dtype.element_ty), mask=m_tc2[:, None] & (o_i[None, :] < BC)) + tl.store(p_Aqk21, (b_Aqk21 * scale).to(Aqk.dtype.element_ty), mask=m_tc2[:, None] & (o_i[None, :] < BC)) + + b_b2 = tl.load(beta + bos * H + i_h + o_c2*H, mask=m_tc2, other=0.0).to(tl.float32) + b_Akk20 = b_Akk20 * b_b2[:, None] + b_Akk21 = b_Akk21 * b_b2[:, None] + if i_tc3 < T: + p_Aqk30 = Aqk + o_c3[:, None] * (H*BT) + o_i[None, :] + p_Aqk31 = Aqk + o_c3[:, None] * (H*BT) + (o_i + BC)[None, :] + p_Aqk32 = Aqk + o_c3[:, None] * (H*BT) + (o_i + 2*BC)[None, :] + tl.store(p_Aqk30, (b_Aqk30 * scale).to(Aqk.dtype.element_ty), mask=m_tc3[:, None] & (o_i[None, :] < BC)) + tl.store(p_Aqk31, (b_Aqk31 * scale).to(Aqk.dtype.element_ty), mask=m_tc3[:, None] & (o_i[None, :] < BC)) + tl.store(p_Aqk32, (b_Aqk32 * scale).to(Aqk.dtype.element_ty), mask=m_tc3[:, None] & (o_i[None, :] < BC)) + + b_b3 = tl.load(beta + bos * H + i_h + o_c3*H, mask=m_tc3, other=0.0).to(tl.float32) + b_Akk30 = b_Akk30 * b_b3[:, None] + b_Akk31 = b_Akk31 * b_b3[:, None] + b_Akk32 = b_Akk32 * b_b3[:, None] + + ################################################################################ + # 3. load diagonal Akk blocks + ################################################################################ + p_Akk00 = Akk_diag + o_c0[:, None] * (H*BC) + o_i[None, :] + p_Akk11 = Akk_diag + o_c1[:, None] * (H*BC) + o_i[None, :] + p_Akk22 = Akk_diag + o_c2[:, None] * (H*BC) + o_i[None, :] + p_Akk33 = Akk_diag + o_c3[:, None] * (H*BC) + o_i[None, :] + b_Ai00 = tl.load(p_Akk00, mask=m_tc0[:, None] & (o_i[None, :] < BC), other=0.0).to(tl.float32) + b_Ai11 = tl.load(p_Akk11, mask=m_tc1[:, None] & (o_i[None, :] < BC), other=0.0).to(tl.float32) + b_Ai22 = tl.load(p_Akk22, mask=m_tc2[:, None] & (o_i[None, :] < BC), other=0.0).to(tl.float32) + b_Ai33 = tl.load(p_Akk33, mask=m_tc3[:, None] & (o_i[None, :] < BC), other=0.0).to(tl.float32) + + ################################################################################ + # 4. forward substitution on diagonals + ################################################################################ + o_i = tl.arange(0, BC) + m_A = o_i[:, None] > o_i[None, :] + m_I = o_i[:, None] == o_i[None, :] + + if not USE_SAFE_GATE: + b_Ai00 = -tl.where(m_A, b_Ai00, 0) + b_Ai11 = -tl.where(m_A, b_Ai11, 0) + b_Ai22 = -tl.where(m_A, b_Ai22, 0) + b_Ai33 = -tl.where(m_A, b_Ai33, 0) + + for i in range(2, min(BC, T - i_tc0)): + b_a00 = -tl.load(Akk_diag + (i_tc0 + i) * H*BC + o_i) + b_a00 = tl.where(o_i < i, b_a00, 0.) + b_a00 += tl.sum(b_a00[:, None] * b_Ai00, 0) + b_Ai00 = tl.where((o_i == i)[:, None], b_a00, b_Ai00) + for i in range(BC + 2, min(2*BC, T - i_tc0)): + b_a11 = -tl.load(Akk_diag + (i_tc0 + i) * H*BC + o_i) + b_a11 = tl.where(o_i < i - BC, b_a11, 0.) + b_a11 += tl.sum(b_a11[:, None] * b_Ai11, 0) + b_Ai11 = tl.where((o_i == i - BC)[:, None], b_a11, b_Ai11) + for i in range(2*BC + 2, min(3*BC, T - i_tc0)): + b_a22 = -tl.load(Akk_diag + (i_tc0 + i) * H*BC + o_i) + b_a22 = tl.where(o_i < i - 2*BC, b_a22, 0.) + b_a22 += tl.sum(b_a22[:, None] * b_Ai22, 0) + b_Ai22 = tl.where((o_i == i - 2*BC)[:, None], b_a22, b_Ai22) + for i in range(3*BC + 2, min(4*BC, T - i_tc0)): + b_a33 = -tl.load(Akk_diag + (i_tc0 + i) * H*BC + o_i) + b_a33 = tl.where(o_i < i - 3*BC, b_a33, 0.) + b_a33 += tl.sum(b_a33[:, None] * b_Ai33, 0) + b_Ai33 = tl.where((o_i == i - 3*BC)[:, None], b_a33, b_Ai33) + + b_Ai00 += m_I + b_Ai11 += m_I + b_Ai22 += m_I + b_Ai33 += m_I + + ################################################################################ + # 5. compute merged inverse using off-diagonals + ################################################################################ + + # we used tf32 to maintain matrix inverse's precision whenever possible. + b_Ai10 = -tl.dot( + tl.dot(b_Ai11, b_Akk10, input_precision=SOLVE_TRIL_DOT_PRECISION), + b_Ai00, + input_precision=SOLVE_TRIL_DOT_PRECISION + ) + b_Ai21 = -tl.dot( + tl.dot(b_Ai22, b_Akk21, input_precision=SOLVE_TRIL_DOT_PRECISION), + b_Ai11, + input_precision=SOLVE_TRIL_DOT_PRECISION + ) + b_Ai32 = -tl.dot( + tl.dot(b_Ai33, b_Akk32, input_precision=SOLVE_TRIL_DOT_PRECISION), + b_Ai22, + input_precision=SOLVE_TRIL_DOT_PRECISION + ) + + b_Ai20 = -tl.dot( + b_Ai22, + tl.dot(b_Akk20, b_Ai00, input_precision=SOLVE_TRIL_DOT_PRECISION) + + tl.dot(b_Akk21, b_Ai10, input_precision=SOLVE_TRIL_DOT_PRECISION), + input_precision=SOLVE_TRIL_DOT_PRECISION + ) + b_Ai31 = -tl.dot( + b_Ai33, + tl.dot(b_Akk31, b_Ai11, input_precision=SOLVE_TRIL_DOT_PRECISION) + + tl.dot(b_Akk32, b_Ai21, input_precision=SOLVE_TRIL_DOT_PRECISION), + input_precision=SOLVE_TRIL_DOT_PRECISION + ) + b_Ai30 = -tl.dot( + b_Ai33, + tl.dot(b_Akk30, b_Ai00, input_precision=SOLVE_TRIL_DOT_PRECISION) + + tl.dot(b_Akk31, b_Ai10, input_precision=SOLVE_TRIL_DOT_PRECISION) + + tl.dot(b_Akk32, b_Ai20, input_precision=SOLVE_TRIL_DOT_PRECISION), + input_precision=SOLVE_TRIL_DOT_PRECISION + ) + + ################################################################################ + # 6. store full Akk_inv to Akk + ################################################################################ + m_col = o_i[None, :] < BC + p_Akk00 = Akk + o_c0[:, None] * (H*BT) + o_i[None, :] + p_Akk10 = Akk + o_c1[:, None] * (H*BT) + o_i[None, :] + p_Akk11 = Akk + o_c1[:, None] * (H*BT) + (o_i + BC)[None, :] + p_Akk20 = Akk + o_c2[:, None] * (H*BT) + o_i[None, :] + p_Akk21 = Akk + o_c2[:, None] * (H*BT) + (o_i + BC)[None, :] + p_Akk22 = Akk + o_c2[:, None] * (H*BT) + (o_i + 2*BC)[None, :] + p_Akk30 = Akk + o_c3[:, None] * (H*BT) + o_i[None, :] + p_Akk31 = Akk + o_c3[:, None] * (H*BT) + (o_i + BC)[None, :] + p_Akk32 = Akk + o_c3[:, None] * (H*BT) + (o_i + 2*BC)[None, :] + p_Akk33 = Akk + o_c3[:, None] * (H*BT) + (o_i + 3*BC)[None, :] + + tl.store(p_Akk00, b_Ai00.to(Akk.dtype.element_ty), mask=m_tc0[:, None] & m_col) + tl.store(p_Akk10, b_Ai10.to(Akk.dtype.element_ty), mask=m_tc1[:, None] & m_col) + tl.store(p_Akk11, b_Ai11.to(Akk.dtype.element_ty), mask=m_tc1[:, None] & m_col) + tl.store(p_Akk20, b_Ai20.to(Akk.dtype.element_ty), mask=m_tc2[:, None] & m_col) + tl.store(p_Akk21, b_Ai21.to(Akk.dtype.element_ty), mask=m_tc2[:, None] & m_col) + tl.store(p_Akk22, b_Ai22.to(Akk.dtype.element_ty), mask=m_tc2[:, None] & m_col) + tl.store(p_Akk30, b_Ai30.to(Akk.dtype.element_ty), mask=m_tc3[:, None] & m_col) + tl.store(p_Akk31, b_Ai31.to(Akk.dtype.element_ty), mask=m_tc3[:, None] & m_col) + tl.store(p_Akk32, b_Ai32.to(Akk.dtype.element_ty), mask=m_tc3[:, None] & m_col) + tl.store(p_Akk33, b_Ai33.to(Akk.dtype.element_ty), mask=m_tc3[:, None] & m_col) + + +@triton.heuristics({ + 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None, +}) +@triton.autotune( + configs=[ + triton.Config({}, num_warps=num_warps, num_stages=num_stages) + for num_warps in [1, 2, 4, 8] + for num_stages in [2, 3, 4] + ], + key=["BT", "BC"], + **autotune_cache_kwargs, +) +@triton.jit(do_not_specialize=['T']) +def chunk_precond_kda_fwd_kernel_intra_sub_chunk( + q, + k, # Original k (for Akk row side) + k_precond, # Preconditioned k (for column side) + g, + beta, + Aqk, + Akk, + scale, + cu_seqlens, + chunk_indices, + T, + H: tl.constexpr, + K: tl.constexpr, + BT: tl.constexpr, + BC: tl.constexpr, + BK: tl.constexpr, + IS_VARLEN: tl.constexpr, + USE_GATHER: tl.constexpr, +): + """ + Asymmetric sub_chunk kernel for preconditioned KDA. + Computes diagonal Aqk and Akk blocks using block-level dot products. + Key difference from symmetric KDA: column side uses k_precond. + """ + i_t, i_i, i_bh = tl.program_id(0).to(tl.int64), tl.program_id(1).to(tl.int64), tl.program_id(2).to(tl.int64) + i_b, i_h = i_bh // H, i_bh % H + + 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.int64) + bos, eos = tl.load(cu_seqlens + i_n).to(tl.int64), tl.load(cu_seqlens + i_n + 1).to(tl.int64) + T = eos - bos + else: + bos, eos = i_b * T, i_b * T + T + + i_ti = i_t * BT + i_i * BC + if i_ti >= T: + return + + o_c = i_ti + tl.arange(0, BC) + m_c = o_c < T + + q = q + (bos * H + i_h) * K + k = k + (bos * H + i_h) * K + k_precond = k_precond + (bos * H + i_h) * K + g = g + (bos * H + i_h) * K + beta = beta + bos * H + i_h + Aqk = Aqk + (bos * H + i_h) * BT + Akk = Akk + (bos * H + i_h) * BC + + o_k = tl.arange(0, BK) + m_ck = m_c[:, None] & (o_k[None, :] < K) + p_q = q + o_c[:, None] * (H*K) + o_k[None, :] + p_k = k + o_c[:, None] * (H*K) + o_k[None, :] + p_kp = k_precond + o_c[:, None] * (H*K) + o_k[None, :] + p_g = g + o_c[:, None] * (H*K) + o_k[None, :] + + b_q = tl.load(p_q, mask=m_ck, other=0.0) + b_k = tl.load(p_k, mask=m_ck, other=0.0) + b_kp = tl.load(p_kp, mask=m_ck, other=0.0) + b_g = tl.load(p_g, mask=m_ck, other=0.0) + b_beta = tl.load(beta + o_c*H, mask=m_c, other=0.0) + + if USE_GATHER: + b_gn = gather(b_g, tl.full([1, BK], min(BC//2, T - i_ti - 1), dtype=tl.int16), axis=0) + else: + # calculate offset + p_gn = g + (i_ti + min(BC // 2, T - i_ti - 1)) * H*K + tl.arange(0, BK) + b_gn = tl.load(p_gn, mask=tl.arange(0, BK) < K, other=0.0) + b_gn = b_gn[None, :] + + # current block, keep numerical stability by subtracting the left boundary + # less than 85 to avoid overflow in exp2 + b_gm = (b_g - b_gn).to(tl.float32) + + b_gq = tl.where(m_c[:, None], exp2(b_gm), 0.) + b_gk = tl.where(m_c[:, None], exp2(-b_gm), 0.) + + # Asymmetric: column side uses k_precond + b_kpgt = tl.trans(b_kp * b_gk) + + b_Aqk = tl.dot(b_q * b_gq, b_kpgt) * scale + b_Akk = tl.dot(b_k * b_gq, b_kpgt) * b_beta[:, None] + + o_i = tl.arange(0, BC) + m_Aqk = o_i[:, None] >= o_i[None, :] + m_Akk = o_i[:, None] > o_i[None, :] + m_I = o_i[:, None] == o_i[None, :] + + b_Aqk = tl.where(m_Aqk, b_Aqk, 0.0) + b_Akk = tl.where(m_Akk, b_Akk, 0.0) + + p_Aqk = Aqk + o_c[:, None] * (H*BT) + (i_i * BC + o_i)[None, :] + p_Akk = Akk + o_c[:, None] * (H*BC) + o_i[None, :] + tl.store(p_Aqk, b_Aqk.to(Aqk.dtype.element_ty), mask=m_c[:, None] & (o_i[None, :] < BC)) + tl.store(p_Akk, b_Akk.to(Akk.dtype.element_ty), mask=m_c[:, None] & (o_i[None, :] < BC)) + + tl.debug_barrier() + + ################################################################################ + # forward substitution + ################################################################################ + + b_Ai = -b_Akk + for i in range(2, min(BC, T - i_ti)): + b_a = -tl.load(Akk + (i_ti + i) * H*BC + o_i) + b_a = tl.where(o_i < i, b_a, 0.) + b_a += tl.sum(b_a[:, None] * b_Ai, 0) + b_Ai = tl.where((o_i == i)[:, None], b_a, b_Ai) + b_Ai += m_I + tl.store(p_Akk, b_Ai.to(Akk.dtype.element_ty), mask=m_c[:, None] & (o_i[None, :] < BC)) + + +def chunk_precond_kda_fwd_intra( + q: torch.Tensor, + k: torch.Tensor, + k_precond: torch.Tensor, + v: torch.Tensor, + gk: torch.Tensor, + beta: torch.Tensor, + scale: float, + cu_seqlens: torch.LongTensor | None = None, + chunk_size: int = 64, + chunk_indices: torch.LongTensor | None = None, + solve_tril_precision: str | None = None, + safe_gate: bool = False, +): + """ + Forward pass for preconditioned KDA intra-chunk. + + Args: + q: [B, T, H, K] - queries + k: [B, T, H, K] - original keys (for Akk row side, WY w computation) + k_precond: [B, T, H, K] - preconditioned keys (for column side, kg output) + v: [B, T, H, V] - values + gk: [B, T, H, K] - cumsum of gates + beta: [B, T, H] - beta scaling + scale: attention scale + + Returns: + w: [B, T, H, K] - WY w vector (uses original k) + u: [B, T, H, V] - WY u vector + kg: [B, T, H, K] - gated k_precond for hidden state update + Aqk: [B, T, H, BT] - q @ k_precond^T attention matrix + Akk: [B, T, H, BT] - k @ k_precond^T (asymmetric) for WY + """ + if solve_tril_precision is None: + solve_tril_precision = DEFAULT_SOLVE_TRIL_PRECISION + + B, T, H, K = k.shape + BT = chunk_size + if chunk_indices is None and cu_seqlens is not None: + chunk_indices = prepare_chunk_indices(cu_seqlens, BT) + NT = triton.cdiv(T, BT) if cu_seqlens is None else len(chunk_indices) + + BC = 16 + NC = triton.cdiv(BT, BC) + + Aqk = torch.empty(B, T, H, BT, device=k.device, dtype=k.dtype) + # Akk must be zero-initialized - kernel only writes lower triangular + Akk = torch.zeros(B, T, H, BT, device=k.device, dtype=k.dtype) + # Separate fp32 buffer for diagonal 16x16 blocks (for precision in solve_tril) + Akk_diag = torch.empty(B, T, H, BC, device=k.device, dtype=torch.float32) + + # Step 1: Compute diagonal blocks into Akk_diag (fp32) + if safe_gate: + grid = (NT, NC, B * H) + BK = triton.next_power_of_2(K) + chunk_precond_kda_fwd_kernel_intra_sub_chunk[grid]( + q=q, + k=k, + k_precond=k_precond, + g=gk, + beta=beta, + Aqk=Aqk, + Akk=Akk_diag, + scale=scale, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + T=T, + H=H, + K=K, + BT=BT, + BC=BC, + BK=BK, + USE_GATHER=IS_GATHER_SUPPORTED, + ) + else: + Aqk, Akk_diag = chunk_precond_kda_fwd_intra_token_parallel( + q=q, + k=k, + k_precond=k_precond, + gk=gk, + beta=beta, + Aqk=Aqk, + Akk=Akk_diag, + scale=scale, + cu_seqlens=cu_seqlens, + chunk_size=BT, + sub_chunk_size=BC, + ) + + # Step 2: Fused inter + solve_tril + grid = (NT, B * H) + chunk_precond_kda_fwd_kernel_inter_solve_fused[grid]( + q=q, + k=k, + k_precond=k_precond, + g=gk, + beta=beta, + Aqk=Aqk, + Akk_diag=Akk_diag, + Akk=Akk, + scale=scale, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + T=T, + H=H, + K=K, + BT=BT, + BC=BC, + SOLVE_TRIL_DOT_PRECISION=solve_tril_precision, + USE_SAFE_GATE=safe_gate, + ) + + # Step 3: WY representation + # w uses original k (for read/correction), kg uses k_precond (for write/h update) + w, u, _, kg = recompute_w_u_fwd( + k=k, # Original k for w + k_precond=k_precond, # k_precond for kg + v=v, + beta=beta, + A=Akk, + gk=gk, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + ) + return w, u, kg, Aqk, Akk + + +@triton.heuristics({ + 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None, +}) +@triton.autotune( + configs=[ + triton.Config({}, num_warps=num_warps, num_stages=num_stages) + for num_warps in [1, 2, 4, 8] + for num_stages in [2, 3, 4] + ], + key=['BK', 'NC', 'BT'], + **autotune_cache_kwargs, +) +@triton.jit(do_not_specialize=['B', 'T']) +def chunk_precond_kda_bwd_kernel_intra( + q, # [B, T, H, K] - queries + k, # [B, T, H, K] - original k (for Akk row side) + k_precond, # [B, T, H, K] - preconditioned k (for column side) + g, # [B, T, H, K] - gate cumsum + beta, # [B, T, H] + dAqk, # [B, T, H, BT] - gradient of Aqk + dAkk, # [B, T, H, BT] - gradient of Akk (asymmetric) + dq, # [B, T, H, K] - input: accumulated dq + dq2, # [B, T, H, K] - output: updated dq + dk, # [B, T, H, K] - input: accumulated dk (original k) + dk2, # [B, T, H, K] - output: updated dk + dk_precond, # [B, T, H, K] - input: accumulated dk_precond + dk_precond2, # [B, T, H, K] - output: updated dk_precond + dg, # [B, T, H, K] - input: accumulated dg + dg2, # [B, T, H, K] - output: updated dg + db, # [NK, B, T, H] - output: accumulated dbeta + cu_seqlens, + chunk_indices, + B, + T, + H: tl.constexpr, + K: tl.constexpr, + BT: tl.constexpr, + BC: tl.constexpr, + BK: tl.constexpr, + NC: tl.constexpr, + IS_VARLEN: tl.constexpr, + SAFE_GATE: tl.constexpr = False, + USE_GATHER: tl.constexpr = False, +): + """ + Asymmetric intra backward for preconditioned KDA. + + Key differences from symmetric KDA: + - Row side of Akk uses original k + - Column side uses k_precond + - dAkk flows to both dk (row) and dk_precond (column) + - dAqk flows to dq (row) and dk_precond (column) + + Forward: + Aqk[t, s] = q[t] @ k_precond[s]^T * exp(g[t] - g[s]) * beta[s] (for t >= s) + Akk[t, s] = k[t] @ k_precond[s]^T * exp(g[t] - g[s]) * beta[s] (for t > s) + """ + i_kc, i_t = unflatten_program_id(tl.cdiv(K, BK) * NC) + i_bh = tl.program_id(1).to(tl.int64) + i_b, i_h = i_bh // H, i_bh % H + i_k = i_kc // NC + i_i = i_kc % NC + + all_BT = B * T + + 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.int64) + bos, eos = tl.load(cu_seqlens + i_n).to(tl.int64), tl.load(cu_seqlens + i_n + 1).to(tl.int64) + T = eos - bos + else: + bos, eos = i_b * T, i_b * T + T + + i_ti = i_t * BT + i_i * BC + if i_ti >= T: + return + + o_k = i_k * BK + tl.arange(0, BK) + o_i = tl.arange(0, BC) + m_k = o_k < K + o_ti = i_ti + o_i + m_ti = o_ti < T + + q += (bos * H + i_h) * K + k += (bos * H + i_h) * K + k_precond += (bos * H + i_h) * K + g += (bos * H + i_h) * K + beta += bos * H + i_h + dAqk += (bos * H + i_h) * BT + dAkk += (bos * H + i_h) * BT + dq += (bos * H + i_h) * K + dq2 += (bos * H + i_h) * K + dk += (bos * H + i_h) * K + dk2 += (bos * H + i_h) * K + dk_precond += (bos * H + i_h) * K + dk_precond2 += (bos * H + i_h) * K + dg += (bos * H + i_h) * K + dg2 += (bos * H + i_h) * K + db += (i_k * all_BT + bos) * H + i_h # Like KDA line 416 + + p_g = g + o_ti[:, None] * (H*K) + o_k[None, :] + b_g = tl.load(p_g, mask=m_ti[:, None] & m_k[None, :], other=0.0).to(tl.float32) + b_b = tl.load(beta + o_ti*H, mask=m_ti, other=0.0).to(tl.float32) + + p_gn_start = g + i_ti * H*K + o_k + b_gn_start = tl.load(p_gn_start, mask=m_k, other=0).to(tl.float32) + + b_dq2 = tl.zeros([BC, BK], dtype=tl.float32) + b_dk2 = tl.zeros([BC, BK], dtype=tl.float32) + + for i_j in range(0, i_i): + o_cj = i_t * BT + i_j * BC + o_i + m_cj = o_cj < T + p_kp = k_precond + o_cj[:, None] * (H*K) + o_k[None, :] + p_gk = g + o_cj[:, None] * (H*K) + o_k[None, :] + p_dAqk = dAqk + o_ti[:, None] * (H*BT) + (i_j * BC + o_i)[None, :] + p_dAkk = dAkk + o_ti[:, None] * (H*BT) + (i_j * BC + o_i)[None, :] + # [BC, BK] + b_kp = tl.load(p_kp, mask=m_cj[:, None] & m_k[None, :], other=0.0).to(tl.float32) + b_gk = tl.load(p_gk, mask=m_cj[:, None] & m_k[None, :], other=0.0).to(tl.float32) + b_kpg = b_kp * exp2(b_gn_start[None, :] - b_gk) + # [BC, BC] + b_dAqk = tl.load(p_dAqk, mask=m_ti[:, None] & (o_i[None, :] < BC), other=0.0).to(tl.float32) + b_dAkk = tl.load(p_dAkk, mask=m_ti[:, None] & (o_i[None, :] < BC), other=0.0).to(tl.float32) + # [BC, BK] + b_dq2 = tl.dot(b_dAqk, b_kpg, b_dq2) + b_dk2 = tl.dot(b_dAkk, b_kpg, b_dk2) + + b_gqn = exp2(b_g - b_gn_start[None, :]) + b_dq2 *= b_gqn + b_dk2 *= b_gqn + + o_dA = (i_ti + o_i) * H*BT + i_i * BC + p_kpj = k_precond + i_ti * H*K + o_k + p_gkj = g + i_ti * H*K + o_k + m_dA = (i_ti + o_i) < T + + p_q = q + o_ti[:, None] * (H*K) + o_k[None, :] + p_k = k + o_ti[:, None] * (H*K) + o_k[None, :] + b_q = tl.load(p_q, mask=m_ti[:, None] & m_k[None, :], other=0.0).to(tl.float32) + b_k = tl.load(p_k, mask=m_ti[:, None] & m_k[None, :], other=0.0).to(tl.float32) + + if SAFE_GATE: + # Vectorized upper diagonal (dq2/dk2): column side uses k_precond + if USE_GATHER: + b_gn = gather(b_g, tl.full([1, BK], min(BC//2, T - i_ti - 1), dtype=tl.int16), axis=0) + else: + p_gn = g + (i_ti + min(BC // 2, T - i_ti - 1)) * H*K + o_k + b_gn = tl.load(p_gn, mask=m_k, other=0)[None, :] + + p_kp_diag = k_precond + o_ti[:, None] * (H*K) + o_k[None, :] + b_kp_diag = tl.load(p_kp_diag, mask=m_ti[:, None] & m_k[None, :], other=0.0).to(tl.float32) + + p_dAqk = dAqk + o_ti[:, None] * (H*BT) + (i_i * BC + o_i)[None, :] + p_dAkk = dAkk + o_ti[:, None] * (H*BT) + (i_i * BC + o_i)[None, :] + b_dAqk_diag_qk = tl.load(p_dAqk, mask=m_ti[:, None] & (o_i[None, :] < BC), other=0.0).to(tl.float32) + b_dAkk_diag_qk = tl.load(p_dAkk, mask=m_ti[:, None] & (o_i[None, :] < BC), other=0.0).to(tl.float32) + + m_i_diag_qk = (o_i[:, None] >= o_i[None, :]) & ((i_ti + o_i[:, None]) < T) & ((i_ti + o_i[None, :]) < T) + m_j_diag_qk = (i_ti + o_i[:, None]) < T + + b_dAqk_diag_qk = tl.where(m_i_diag_qk, b_dAqk_diag_qk, 0.) + b_dAkk_diag_qk = tl.where(m_i_diag_qk, b_dAkk_diag_qk, 0.) + b_g_diag_qk = tl.where(m_j_diag_qk, b_g - b_gn, 0.) + exp_b_g_diag_qk = tl.where(m_j_diag_qk, exp2(b_g_diag_qk), 0.) + exp_neg_b_g_diag_qk = tl.where(m_j_diag_qk, exp2(-b_g_diag_qk), 0.) + + # Asymmetric: column side uses k_precond + b_kp_exp_diag_qk = b_kp_diag * exp_neg_b_g_diag_qk + b_dq2 += tl.dot(b_dAqk_diag_qk, b_kp_exp_diag_qk) * exp_b_g_diag_qk + b_dk2 += tl.dot(b_dAkk_diag_qk, b_kp_exp_diag_qk) * exp_b_g_diag_qk + else: + for j in range(0, min(BC, T - i_t * BT - i_i * BC)): + b_dAqk_val = tl.load(dAqk + o_dA + j, mask=m_dA, other=0).to(tl.float32) + b_dAkk_val = tl.load(dAkk + o_dA + j, mask=m_dA, other=0).to(tl.float32) + + b_kpj = tl.load(p_kpj, mask=m_k, other=0).to(tl.float32) + b_gkj = tl.load(p_gkj, mask=m_k, other=0).to(tl.float32) + + m_ij = o_i[:, None] >= j + b_kpgj = b_kpj[None, :] * exp2(b_g - b_gkj[None, :]) + b_dq2 += tl.where(m_ij, b_dAqk_val[:, None] * b_kpgj, 0.) + b_dk2 += tl.where(m_ij, b_dAkk_val[:, None] * b_kpgj, 0.) + + p_kpj += H*K + p_gkj += H*K + + b_db = tl.sum(b_dk2 * b_k, 1) + p_db = db + o_ti*H + tl.store(p_db, b_db.to(db.dtype.element_ty), mask=m_ti) + b_dk2 *= b_b[:, None] + + m_tik = m_ti[:, None] & m_k[None, :] + p_dq = dq + o_ti[:, None] * (H*K) + o_k[None, :] + p_dq2 = dq2 + o_ti[:, None] * (H*K) + o_k[None, :] + + b_dg2 = b_q * b_dq2 + b_dq2 = b_dq2 + tl.load(p_dq, mask=m_tik, other=0.0) + tl.store(p_dq2, b_dq2.to(dq2.dtype.element_ty), mask=m_tik) + + tl.debug_barrier() + b_dkt = tl.zeros([BC, BK], dtype=tl.float32) # Final dk_precond (with beta from rows) + + # keep int64: `i_t` is int64 for pointer arithmetic, and the loop + # counter below must match the type of its int64 initial value `i_i + 1` + NC_actual = min(NC, tl.cdiv(T - i_t * BT, BC)) + if i_i < NC_actual - 1: + p_gn_t = g + (min(i_ti + BC, T) - 1) * H*K + o_k + b_gn_t = tl.load(p_gn_t, mask=m_k, other=0).to(tl.float32) + + for i_j in range(i_i + 1, NC_actual): + o_rj = i_t*BT + i_j*BC + o_i + m_rj = o_rj < T + p_q_t = q + o_rj[:, None] * (H*K) + o_k[None, :] + p_k_t = k + o_rj[:, None] * (H*K) + o_k[None, :] + p_gk_t = g + o_rj[:, None] * (H*K) + o_k[None, :] + p_b_row = beta + o_rj*H # beta at ROW positions + p_dAqk_t = dAqk + (i_i*BC + o_i)[:, None] + o_rj[None, :] * (H*BT) + p_dAkk_t = dAkk + (i_i*BC + o_i)[:, None] + o_rj[None, :] * (H*BT) + m_dA_t = (o_i[:, None] < BC) & m_rj[None, :] + + # [BC] + b_b_row = tl.load(p_b_row, mask=m_rj, other=0.0).to(tl.float32) # beta from ROW positions + # [BC, BK] + b_q_t = tl.load(p_q_t, mask=m_rj[:, None] & m_k[None, :], other=0.0).to(tl.float32) + b_k_t = tl.load(p_k_t, mask=m_rj[:, None] & m_k[None, :], other=0.0).to(tl.float32) + b_gk_t = tl.load(p_gk_t, mask=m_rj[:, None] & m_k[None, :], other=0.0).to(tl.float32) + # [BC, BC] + b_dAqk_t = tl.load(p_dAqk_t, mask=m_dA_t, other=0.0).to(tl.float32) + b_dAkk_t = tl.load(p_dAkk_t, mask=m_dA_t, other=0.0).to(tl.float32) + + o_j = i_t * BT + i_j * BC + o_i + m_j = o_j < T + # [BC, BK] + b_gkn_t = tl.where(m_j[:, None], exp2(b_gk_t - b_gn_t[None, :]), 0) + b_qg_t = b_q_t * b_gkn_t + b_kg_t = b_k_t * b_gkn_t * b_b_row[:, None] # beta from ROW positions + # [BC, BK] + b_dkt = tl.dot(b_dAqk_t, b_qg_t, b_dkt) + b_dkt = tl.dot(b_dAkk_t, b_kg_t, b_dkt) + + b_dkt *= exp2(b_gn_t[None, :] - b_g) + + if SAFE_GATE: + # Vectorized lower diagonal (dkt): row side uses q and k*beta + if USE_GATHER: + b_gn_t2 = gather(b_g, tl.full([1, BK], min(BC//2, T - i_ti - 1), dtype=tl.int16), axis=0) + else: + p_gn_t2 = g + (i_ti + min(BC // 2, T - i_ti - 1)) * H*K + o_k + b_gn_t2 = tl.load(p_gn_t2, mask=m_k, other=0).to(tl.float32)[None, :] + p_q_diag = q + o_ti[:, None] * (H*K) + o_k[None, :] + b_q_diag = tl.load(p_q_diag, mask=m_ti[:, None] & m_k[None, :], other=0.0).to(tl.float32) + b_b_diag = tl.load(beta + o_ti*H, mask=m_ti, other=0.0).to(tl.float32) + + p_dAqk_diag_kk = dAqk + (i_i * BC + o_i)[:, None] + o_ti[None, :] * (H*BT) + p_dAkk_diag_kk = dAkk + (i_i * BC + o_i)[:, None] + o_ti[None, :] * (H*BT) + m_dA_kk = (o_i[:, None] < BC) & m_ti[None, :] + b_dAqk_diag_kk = tl.load(p_dAqk_diag_kk, mask=m_dA_kk, other=0.0).to(tl.float32) + b_dAkk_diag_kk = tl.load(p_dAkk_diag_kk, mask=m_dA_kk, other=0.0).to(tl.float32) + + m_i_diag_kk = (o_i[:, None] <= o_i[None, :]) & ((i_ti + o_i[:, None]) < T) & ((i_ti + o_i[None, :]) < T) + m_j_diag_kk = (i_ti + o_i[:, None]) < T + + b_dAqk_diag_kk = tl.where(m_i_diag_kk, b_dAqk_diag_kk, 0.) + b_dAkk_diag_kk = tl.where(m_i_diag_kk, b_dAkk_diag_kk, 0.) + # ensure numerical stability + b_g_diag_kk = tl.where(m_j_diag_kk, b_g - b_gn_t2, 0.) + exp_b_g_diag_kk = tl.where(m_j_diag_kk, exp2(b_g_diag_kk), 0.) + exp_neg_b_g_diag_kk = tl.where(m_j_diag_kk, exp2(-b_g_diag_kk), 0.) + + # Row-side values: q (for Aqk) and k*beta (for Akk) + b_q_exp = b_q_diag * exp_b_g_diag_kk + b_kb_exp = b_k * b_b_diag[:, None] * exp_b_g_diag_kk + + b_dkt += tl.dot(b_dAqk_diag_kk, b_q_exp) * exp_neg_b_g_diag_kk + b_dkt += tl.dot(b_dAkk_diag_kk, b_kb_exp) * exp_neg_b_g_diag_kk + else: + o_dA_t = i_ti * H*BT + i_i * BC + o_i + p_qj = q + i_ti * H*K + o_k + p_kj = k + i_ti * H*K + o_k + p_gkj_t = g + i_ti * H*K + o_k + p_bj = beta + i_ti * H + + for j in range(0, min(BC, T - i_t * BT - i_i * BC)): + # [BC] + b_dAqk_t = tl.load(dAqk + o_dA_t + j * H*BT).to(tl.float32) + b_dAkk_t = tl.load(dAkk + o_dA_t + j * H*BT).to(tl.float32) + # [BK] + b_qj = tl.load(p_qj, mask=m_k, other=0).to(tl.float32) + b_kj = tl.load(p_kj, mask=m_k, other=0).to(tl.float32) + b_gkj_t = tl.load(p_gkj_t, mask=m_k, other=0).to(tl.float32) + b_bj = tl.load(p_bj).to(tl.float32) # beta at row position j + # [BC, BK] + m_ij = o_i[:, None] <= j + b_gkq_t = exp2(b_gkj_t[None, :] - b_g) + # k with beta from ROW position, q without beta (Aqk has no beta) + b_dkt += tl.where(m_ij, (b_dAkk_t[:, None] * b_kj[None, :] * b_bj + + b_dAqk_t[:, None] * b_qj[None, :]) * b_gkq_t, 0.) + + p_qj += H*K + p_kj += H*K + p_gkj_t += H*K + p_bj += H + + p_kp_local = k_precond + o_ti[:, None] * (H*K) + o_k[None, :] + b_kp_local = tl.load(p_kp_local, mask=m_tik, other=0.0).to(tl.float32) + + p_dk = dk + o_ti[:, None] * (H*K) + o_k[None, :] + p_dk2 = dk2 + o_ti[:, None] * (H*K) + o_k[None, :] + p_dk_precond = dk_precond + o_ti[:, None] * (H*K) + o_k[None, :] + p_dk_precond2 = dk_precond2 + o_ti[:, None] * (H*K) + o_k[None, :] + p_dg = dg + o_ti[:, None] * (H*K) + o_k[None, :] + p_dg2 = dg2 + o_ti[:, None] * (H*K) + o_k[None, :] + + b_dg2 += b_dk2 * b_k - b_dkt * b_kp_local + tl.load(p_dg, mask=m_tik, other=0.0) + b_dk2 += tl.load(p_dk, mask=m_tik, other=0.0) + b_dk_precond2 = b_dkt + tl.load(p_dk_precond, mask=m_tik, other=0.0) + + tl.store(p_dk2, b_dk2.to(dk2.dtype.element_ty), mask=m_tik) + tl.store(p_dk_precond2, b_dk_precond2.to(dk_precond2.dtype.element_ty), mask=m_tik) + tl.store(p_dg2, b_dg2.to(dg2.dtype.element_ty), mask=m_tik) + + +def chunk_precond_kda_bwd_intra( + q: torch.Tensor, + k: torch.Tensor, + k_precond: torch.Tensor, + g: torch.Tensor, + beta: torch.Tensor, + dAqk: torch.Tensor, + dAkk: torch.Tensor, + dq: torch.Tensor, + dk: torch.Tensor, + dk_precond: torch.Tensor, + db: torch.Tensor, + dg: torch.Tensor, + cu_seqlens: torch.LongTensor | None = None, + chunk_indices: torch.LongTensor | None = None, + chunk_size: int = 64, + safe_gate: bool = False, +): + """ + Asymmetric intra backward for preconditioned KDA. + """ + B, T, H, K = k.shape + BT = chunk_size + BC = min(16, BT) + BK = min(32, triton.next_power_of_2(K)) + + if chunk_indices is None and cu_seqlens is not None: + chunk_indices = prepare_chunk_indices(cu_seqlens, BT) + NT = triton.cdiv(T, BT) if cu_seqlens is None else len(chunk_indices) + NC = triton.cdiv(BT, BC) + NK = triton.cdiv(K, BK) + + dq2 = torch.empty_like(q) + dk2 = torch.empty_like(k) + dk_precond2 = torch.empty_like(k_precond) + db2 = beta.new_empty(NK, *beta.shape, dtype=torch.float) + dg2 = torch.empty_like(dg, dtype=torch.float) + + grid = (NK * NC * NT, B * H) + chunk_precond_kda_bwd_kernel_intra[grid]( + q=q, + k=k, + k_precond=k_precond, + g=g, + beta=beta, + dAqk=dAqk, + dAkk=dAkk, + dq=dq, + dq2=dq2, + dk=dk, + dk2=dk2, + dk_precond=dk_precond, + dk_precond2=dk_precond2, + dg=dg, + dg2=dg2, + db=db2, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + B=B, + T=T, + H=H, + K=K, + BT=BT, + BC=BC, + BK=BK, + NC=NC, + SAFE_GATE=safe_gate, + USE_GATHER=IS_GATHER_SUPPORTED, + ) + + dq = dq2 + dk = dk2 + dk_precond = dk_precond2 + db = db2.sum(0).add_(db) + dg = chunk_local_cumsum( + dg2, + chunk_size=chunk_size, + reverse=True, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + ) + + return dq, dk, dk_precond, db, dg diff --git a/build/torch-cuda/ops/precond_kda/chunk_intra_token_parallel.py b/build/torch-cuda/ops/precond_kda/chunk_intra_token_parallel.py new file mode 100644 index 0000000000000000000000000000000000000000..2e8050cc7c29ffb867ea4bf808d2d7c97fb30550 --- /dev/null +++ b/build/torch-cuda/ops/precond_kda/chunk_intra_token_parallel.py @@ -0,0 +1,188 @@ +# 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 +import triton.language as tl + +from ...ops.utils.op import exp2 +from ...utils import autotune_cache_kwargs + + +@triton.heuristics({ + 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None, +}) +@triton.autotune( + configs=[ + triton.Config({'BH': BH}, num_warps=num_warps) + for BH in [1, 2, 4, 8] + for num_warps in [1, 2, 4, 8] + ], + key=["K", "H"], + **autotune_cache_kwargs, +) +@triton.jit(do_not_specialize=['T', 'N']) +def chunk_precond_kda_fwd_kernel_intra_token_parallel( + q, + k, # Original k (for Akk row side) + k_precond, # Preconditioned k (for column side of both Aqk and Akk) + g, + beta, + Aqk, + Akk, + scale, + cu_seqlens, + N, + T, + H: tl.constexpr, + K: tl.constexpr, + BT: tl.constexpr, + BC: tl.constexpr, + BH: tl.constexpr, + IS_VARLEN: tl.constexpr, +): + """ + Asymmetric version for precond_kda: + - Aqk = q @ k_precond^T (column uses k_precond) + - Akk = k @ k_precond^T (row uses original k, column uses k_precond) + """ + i_tg, i_hg = tl.program_id(0).to(tl.int64), tl.program_id(1).to(tl.int64) + + if IS_VARLEN: + i_n = 0 + left, right = 0, N + + # Unrolled binary search (max B=2^32) + # We can limit iterations based on expected max batch size if needed + # 20 iterations covers B=1M, usually enough + for _ in range(20): + if left < right: + mid = (left + right) // 2 + if i_tg < tl.load(cu_seqlens + mid + 1).to(tl.int32): + right = mid + else: + left = mid + 1 + i_n = left + + bos, eos = tl.load(cu_seqlens + i_n).to(tl.int64), tl.load(cu_seqlens + i_n + 1).to(tl.int64) + T = eos - bos + i_t = i_tg - bos + else: + bos = (i_tg // T) * T + i_t = i_tg % T + + if i_t >= T: + return + + i_c = i_t // BT + i_s = (i_t % BT) // BC + i_tc = i_c * BT + i_ts = i_tc + i_s * BC + + q += bos * H*K + k += bos * H*K + k_precond += bos * H*K + g += bos * H*K + Aqk += bos * H*BT + Akk += bos * H*BC + beta += bos * H + + BK: tl.constexpr = triton.next_power_of_2(K) + o_h = tl.arange(0, BH) + o_k = tl.arange(0, BK) + m_h = (i_hg * BH + o_h) < H + m_k = o_k < K + + o_hh = i_hg * BH + o_h + m_hk = m_h[:, None] & m_k[None, :] + + p_q = q + i_t * H*K + o_hh[:, None] * K + o_k[None, :] + p_k = k + i_t * H*K + o_hh[:, None] * K + o_k[None, :] + p_g = g + i_t * H*K + o_hh[:, None] * K + o_k[None, :] + p_beta = beta + i_t * H + o_hh + + b_q = tl.load(p_q, mask=m_hk, other=0.0).to(tl.float32) + b_k = tl.load(p_k, mask=m_hk, other=0.0).to(tl.float32) # Original k for row + b_g = tl.load(p_g, mask=m_hk, other=0.0).to(tl.float32) + b_k = b_k * tl.load(p_beta, mask=m_h, other=0.0).to(tl.float32)[:, None] + + for j in range(i_ts, min(i_t + 1, min(T, i_ts + BC))): + # Load k_precond at position j (column side) - ASYMMETRIC: use k_precond here + p_kpj = k_precond + j * H*K + o_hh[:, None] * K + o_k[None, :] + p_gj = g + j * H*K + o_hh[:, None] * K + o_k[None, :] + + b_kpj = tl.load(p_kpj, mask=m_hk, other=0.0).to(tl.float32) # k_precond for column + b_gj = tl.load(p_gj, mask=m_hk, other=0.0).to(tl.float32) + + # k_precond gated at position j + b_kpgj = b_kpj * exp2(b_g - b_gj) + b_kpgj = tl.where(m_k[None, :], b_kpgj, 0.0) + + # Aqk: q (row i) @ k_precond^T (col j) + b_Aqk = tl.sum(b_q * b_kpgj, axis=1) * scale + # Akk: k (row i, beta-scaled) @ k_precond^T (col j) - ASYMMETRIC + b_Akk = tl.sum(b_k * b_kpgj, axis=1) * tl.where(j < i_t, 1.0, 0.0) + + tl.store(Aqk + i_t * H*BT + (i_hg * BH + o_h) * BT + j % BT, b_Aqk.to(Aqk.dtype.element_ty), mask=m_h) + tl.store(Akk + i_t * H*BC + (i_hg * BH + o_h) * BC + j - i_ts, b_Akk.to(Akk.dtype.element_ty), mask=m_h) + + +def chunk_precond_kda_fwd_intra_token_parallel( + q: torch.Tensor, + k: torch.Tensor, + k_precond: torch.Tensor, + gk: torch.Tensor, + beta: torch.Tensor, + Aqk: torch.Tensor, + Akk: torch.Tensor, + scale: float, + cu_seqlens: torch.LongTensor | None = None, + chunk_size: int = 64, + sub_chunk_size: int = 16, +) -> None: + """ + Token-parallel implementation for preconditioned KDA. + Supports both fixed-length and variable-length sequences. + Reduces wasted computation on padding. + Asymmetric: uses k for Akk row side, k_precond for column side. + + Writes directly to Aqk and Akk tensors (in-place). + + Args: + q: [B, T, H, K] + k: [B, T, H, K] - original k (for Akk row side) + k_precond: [B, T, H, K] - preconditioned k (for column side) + gk: [B, T, H, K] cumsum of gates + beta: [B, T, H] + Aqk: [B, T, H, BT] output tensor + Akk: [B, T, H, BC] output tensor for diagonal blocks (fp32) + scale: attention scale + """ + B, T, H, K = q.shape + N = len(cu_seqlens) - 1 if cu_seqlens is not None else B + BT = chunk_size + BC = sub_chunk_size + + def grid(meta): return (B * T, triton.cdiv(H, meta['BH'])) + chunk_precond_kda_fwd_kernel_intra_token_parallel[grid]( + q=q, + k=k, + k_precond=k_precond, + g=gk, + beta=beta, + Aqk=Aqk, + Akk=Akk, + scale=scale, + cu_seqlens=cu_seqlens, + N=N, + T=T, + H=H, + K=K, + BT=BT, + BC=BC, + ) + return Aqk, Akk diff --git a/build/torch-xpu/ops/kda/fused_recurrent.py b/build/torch-cuda/ops/precond_kda/fused_recurrent.py similarity index 66% rename from build/torch-xpu/ops/kda/fused_recurrent.py rename to build/torch-cuda/ops/precond_kda/fused_recurrent.py index 285099dea8bc4c10544bd7277b35c4b328fb5193..f35fbbed7d4477c40cdb3718d4f9902343e0dab8 100644 --- a/build/torch-xpu/ops/kda/fused_recurrent.py +++ b/build/torch-cuda/ops/precond_kda/fused_recurrent.py @@ -7,7 +7,7 @@ # This kernel is modified from the Decode kernel of the vllm gdn/kda model. -import warnings +import math import torch import triton @@ -22,6 +22,8 @@ from ...utils import input_guard { "USE_INITIAL_STATE": lambda args: args["h0"] is not None, "STORE_FINAL_STATE": lambda args: args["ht"] is not None, + "USE_INITIAL_ATK": lambda args: args["a0"] is not None, + "STORE_FINAL_ATK": lambda args: args["at"] is not None, "IS_VARLEN": lambda args: args["cu_seqlens"] is not None, "IS_CONTINUOUS_BATCHING": lambda args: args["ssm_state_indices"] is not None, "IS_SPEC_DECODING": lambda args: args["num_accepted_tokens"] is not None, @@ -30,22 +32,29 @@ from ...utils import input_guard } ) @triton.jit(do_not_specialize=["N", "T"]) -def fused_recurrent_kda_fwd_kernel( +def fused_recurrent_precond_kda_fwd_kernel( q, k, v, g, + g_atk, + beta_atk, beta, A_log, dt_bias, o, h0, ht, + a0, + at, cu_seqlens, ssm_state_indices, num_accepted_tokens, lower_bound, + log_atk_scale, scale: tl.constexpr, + logx, + eps, N: tl.int64, # num of sequences T: tl.int64, # num of tokens H: tl.constexpr, @@ -66,15 +75,15 @@ def fused_recurrent_kda_fwd_kernel( IS_CONTINUOUS_BATCHING: tl.constexpr, IS_SPEC_DECODING: tl.constexpr, STORE_FINAL_STATE: tl.constexpr, + USE_INITIAL_ATK: tl.constexpr, + STORE_FINAL_ATK: tl.constexpr, HAS_DT_BIAS: tl.constexpr, USE_GATE_IN_KERNEL: tl.constexpr, USE_LOWER_BOUND: tl.constexpr, - APPLY_BETA_SIGMOID: tl.constexpr, - ALLOW_NEG_EIGVAL: tl.constexpr, - STATE_V_FIRST: tl.constexpr, + TRANSPOSE_STATE: tl.constexpr, num_stages: tl.constexpr, ): - pid = tl.program_id(0) + pid = tl.program_id(0).to(tl.int64) NV = tl.cdiv(V, BV) NK = tl.cdiv(K, BK) i_k = pid % NK @@ -109,16 +118,21 @@ def fused_recurrent_kda_fwd_kernel( p_beta = beta + bos * HV + i_hv p_g = g + (bos * HV + i_hv) * K + o_k + # ATK pointers (per-HV-head: g_atk, beta_atk, log_atk_scale are shape-HV like v/beta/g) + p_g_atk = g_atk + bos * HV + i_hv + p_beta_atk = beta_atk + bos * HV + i_hv p_o = o + (bos * HV + i_hv) * V + o_v + center = tl.load(log_atk_scale + i_hv).to(tl.float32) + mask_k = o_k < K mask_v = o_v < V - if STATE_V_FIRST: + if TRANSPOSE_STATE: mask_h = mask_v[:, None] & mask_k[None, :] else: mask_h = mask_k[:, None] & mask_v[None, :] - if STATE_V_FIRST: + if TRANSPOSE_STATE: b_h = tl.zeros([BV, BK], dtype=tl.float32) else: b_h = tl.zeros([BK, BV], dtype=tl.float32) @@ -135,17 +149,36 @@ def fused_recurrent_kda_fwd_kernel( ) * stride_init_state_token ) - if STATE_V_FIRST: + if TRANSPOSE_STATE: p_h0 = p_h0 + i_hv * K * V + o_v[:, None] * K + o_k[None, :] else: p_h0 = p_h0 + i_hv * K * V + o_k[:, None] * V + o_v[None, :] else: - if STATE_V_FIRST: + if TRANSPOSE_STATE: p_h0 = h0 + (i_n * HV + i_hv) * K * V + o_v[:, None] * K + o_k[None, :] else: p_h0 = h0 + (i_n * HV + i_hv) * K * V + o_k[:, None] * V + o_v[None, :] b_h += tl.load(p_h0, mask=mask_h, other=0).to(tl.float32) + b_a = tl.zeros([BK], dtype=tl.float32) + if USE_INITIAL_ATK: + if IS_CONTINUOUS_BATCHING: + if IS_SPEC_DECODING: + i_t = tl.load(num_accepted_tokens + i_n).to(tl.int64) - 1 + else: + i_t = 0 + p_a0 = ( + a0 + + tl.load(ssm_state_indices + i_n * stride_indices_seq + i_t).to( + tl.int64 + ) + * stride_init_state_token + ) + p_a0 = p_a0 + i_hv * K + o_k + else: + p_a0 = a0 + (i_n * HV + i_hv) * K + o_k + b_a += tl.load(p_a0, mask=mask_k, other=0).to(tl.float32) + for i_t in tl.range(0, T, num_stages=num_stages): b_q = tl.load(p_q, mask=mask_k, other=0, eviction_policy='evict_last').to(tl.float32) b_k = tl.load(p_k, mask=mask_k, other=0, eviction_policy='evict_last').to(tl.float32) @@ -155,7 +188,7 @@ def fused_recurrent_kda_fwd_kernel( b_q = b_q / tl.sqrt(tl.sum(b_q * b_q) + 1e-6) b_k = b_k / tl.sqrt(tl.sum(b_k * b_k) + 1e-6) b_q = b_q * scale - b_g = tl.load(p_g, eviction_policy='evict_last').to(tl.float32) + b_g = tl.load(p_g, mask=mask_k, other=0, eviction_policy='evict_last').to(tl.float32) if USE_GATE_IN_KERNEL: b_A = tl.load(A_log + i_hv).to(tl.float32) @@ -171,12 +204,24 @@ def fused_recurrent_kda_fwd_kernel( else: b_gk = b_g - if STATE_V_FIRST: + b_g_atk = tl.load(p_g_atk).to(tl.float32) + b_beta_atk = tl.load(p_beta_atk).to(tl.float32) + + # ATK: A_t = exp(g_atk) * A_{t-1} + beta_atk * k^2 + b_a = exp(b_g_atk) * b_a + b_beta_atk * (b_k * b_k) + + b_ell = tl.log(b_a + eps) + b_r = b_ell - center + b_s = b_r / (1.0 + tl.abs(b_r)) + b_M = tl.exp(-logx * b_s) + b_k_precond = b_k * b_M + + if TRANSPOSE_STATE: b_h *= exp(b_gk[None, :]) else: b_h *= exp(b_gk[:, None]) - if STATE_V_FIRST: + if TRANSPOSE_STATE: b_v -= tl.sum(b_h * b_k[None, :], 1) else: b_v -= tl.sum(b_h * b_k[:, None], 0) @@ -184,16 +229,12 @@ def fused_recurrent_kda_fwd_kernel( b_beta = tl.load(p_beta, mask=mask_v, other=0, eviction_policy='evict_first').to(tl.float32) else: b_beta = tl.load(p_beta, eviction_policy='evict_last').to(tl.float32) - if APPLY_BETA_SIGMOID: - b_beta = tl.sigmoid(b_beta) - if ALLOW_NEG_EIGVAL: - b_beta = b_beta * 2 b_v *= b_beta - if STATE_V_FIRST: - b_h += b_v[:, None] * b_k[None, :] + if TRANSPOSE_STATE: + b_h += b_v[:, None] * b_k_precond[None, :] b_o = tl.sum(b_h * b_q[None, :], 1) else: - b_h += b_k[:, None] * b_v[None, :] + b_h += b_k_precond[:, None] * b_v[None, :] b_o = tl.sum(b_h * b_q[:, None], 0) tl.store(p_o, b_o.to(p_o.dtype.element_ty), mask=mask_v, eviction_policy='evict_first') @@ -208,53 +249,79 @@ def fused_recurrent_kda_fwd_kernel( ) else: p_ht = ht + (bos + i_t) * stride_final_state_token - if STATE_V_FIRST: + if TRANSPOSE_STATE: p_ht = p_ht + i_hv * K * V + o_v[:, None] * K + o_k[None, :] else: p_ht = p_ht + i_hv * K * V + o_k[:, None] * V + o_v[None, :] tl.store(p_ht, b_h.to(p_ht.dtype.element_ty), mask=mask_h) + # Store ATK state per timestep during continuous batching (only from first v-block and k-block) + if i_v == 0 and i_k == 0: + if INPLACE_FINAL_STATE: + p_at_cb = ( + at + + tl.load(ssm_state_indices + i_n * stride_indices_seq + i_t).to( + tl.int64 + ) + * stride_final_state_token + ) + else: + p_at_cb = at + (bos + i_t) * stride_final_state_token + p_at_cb = p_at_cb + i_hv * K + o_k + tl.store(p_at_cb, b_a.to(p_at_cb.dtype.element_ty), mask=mask_k) + p_q += H * K p_k += H * K p_o += HV * V p_v += HV * V p_g += HV * K + p_g_atk += HV + p_beta_atk += HV p_beta += HV * (V if IS_BETA_HEADWISE else 1) if not IS_CONTINUOUS_BATCHING: if STORE_FINAL_STATE: - if STATE_V_FIRST: + if TRANSPOSE_STATE: p_ht = ht + (i_n * HV + i_hv) * K * V + o_v[:, None] * K + o_k[None, :] else: p_ht = ht + (i_n * HV + i_hv) * K * V + o_k[:, None] * V + o_v[None, :] tl.store(p_ht, b_h.to(p_ht.dtype.element_ty), mask=mask_h) + if STORE_FINAL_ATK: + if i_v == 0 and i_k == 0: + p_at = at + (i_n * HV + i_hv) * K + o_k + tl.store(p_at, b_a.to(p_at.dtype.element_ty), mask=mask_k) + @torch.compiler.disable -def fused_recurrent_kda_fwd( +def fused_recurrent_precond_kda_fwd( q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, g: torch.Tensor, + g_atk: torch.Tensor, + beta_atk: torch.Tensor, beta: torch.Tensor, A_log: torch.Tensor | None = None, dt_bias: torch.Tensor | None = None, initial_state: torch.Tensor | None = None, + initial_A_state: torch.Tensor | None = None, scale: float | None = None, output_final_state: bool = False, inplace_final_state: bool = True, - state_v_first: bool = False, cu_seqlens: torch.LongTensor | None = None, ssm_state_indices: torch.Tensor | None = None, num_accepted_tokens: torch.Tensor | None = None, use_qk_l2norm_in_kernel: bool = False, use_gate_in_kernel: bool = False, - use_beta_sigmoid_in_kernel: bool = False, - allow_neg_eigval: bool = False, lower_bound: float | None = None, out: torch.Tensor | None = None, + transpose_state_layout: bool = False, + x: float = 1.5, + eps: float = 1e-6, + log_atk_scale: torch.Tensor = None, **kwargs, -) -> tuple[torch.Tensor, torch.Tensor]: +) -> tuple[torch.Tensor, torch.Tensor | None, torch.Tensor | None]: if scale is None: scale = k.shape[-1] ** -0.5 @@ -264,6 +331,12 @@ def fused_recurrent_kda_fwd( BK = triton.next_power_of_2(K) BV = 32 + if log_atk_scale is None: + log_atk_scale = torch.full((HV,), -0.2, device=k.device, dtype=torch.float32) + log_atk_scale = log_atk_scale.contiguous() + + logx = math.log(x) if x > 0 else 0.0 + if out is None: out = torch.zeros_like(v) else: @@ -271,13 +344,16 @@ def fused_recurrent_kda_fwd( if inplace_final_state: assert initial_state is not None final_state = initial_state + final_A_state = initial_A_state elif output_final_state: - if state_v_first: + if transpose_state_layout: final_state = q.new_empty(N, HV, V, K, dtype=torch.float32) else: final_state = q.new_empty(N, HV, K, V, dtype=torch.float32) + final_A_state = q.new_empty(N, HV, K, dtype=torch.float32) else: final_state = None + final_A_state = None stride_init_state_token = initial_state.stride(0) if initial_state is not None else 1 stride_final_state_token = final_state.stride(0) if final_state is not None else 1 @@ -289,23 +365,31 @@ def fused_recurrent_kda_fwd( else: stride_indices_seq, stride_indices_tok = ssm_state_indices.stride() - grid = (triton.cdiv(V, BV) * N * HV, ) - fused_recurrent_kda_fwd_kernel[grid]( + NK = triton.cdiv(K, BK) + grid = (triton.cdiv(V, BV) * NK * N * HV, ) + fused_recurrent_precond_kda_fwd_kernel[grid]( q=q, k=k, v=v, g=g, + g_atk=g_atk, + beta_atk=beta_atk, beta=beta, A_log=A_log, dt_bias=dt_bias, o=out, h0=initial_state, ht=final_state, + a0=initial_A_state, + at=final_A_state, cu_seqlens=cu_seqlens, ssm_state_indices=ssm_state_indices, num_accepted_tokens=num_accepted_tokens, lower_bound=lower_bound, + log_atk_scale=log_atk_scale, scale=scale, + logx=logx, + eps=eps, N=N, T=T, H=H, @@ -322,37 +406,39 @@ def fused_recurrent_kda_fwd( USE_QK_L2NORM_IN_KERNEL=use_qk_l2norm_in_kernel, INPLACE_FINAL_STATE=inplace_final_state, USE_GATE_IN_KERNEL=use_gate_in_kernel, - APPLY_BETA_SIGMOID=use_beta_sigmoid_in_kernel, - ALLOW_NEG_EIGVAL=allow_neg_eigval, - STATE_V_FIRST=state_v_first, + TRANSPOSE_STATE=transpose_state_layout, num_warps=4, num_stages=2, ) - return out, final_state + return out, final_state, final_A_state @input_guard -def fused_recurrent_kda( +def fused_recurrent_precond_kda( q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, g: torch.Tensor, + g_atk: torch.Tensor, + beta_atk: torch.Tensor, beta: torch.Tensor, A_log: torch.Tensor | None = None, dt_bias: torch.Tensor | None = None, scale: float | None = None, initial_state: torch.Tensor = None, + initial_A_state: torch.Tensor = None, output_final_state: bool = False, - use_qk_l2norm_in_kernel: bool = False, + use_qk_l2norm_in_kernel: bool = True, use_gate_in_kernel: bool = False, - use_beta_sigmoid_in_kernel: bool = False, - allow_neg_eigval: bool = False, lower_bound: float | None = None, - state_v_first: bool = False, cu_seqlens: torch.LongTensor | None = None, + transpose_state_layout: bool = False, + x: float = 1.5, + eps: float = 1e-6, + log_atk_scale: torch.Tensor = None, **kwargs, -) -> tuple[torch.Tensor, torch.Tensor]: +) -> tuple[torch.Tensor, torch.Tensor | None, torch.Tensor | None]: r""" Args: q (torch.Tensor): @@ -364,89 +450,91 @@ def fused_recurrent_kda( GVA is applied if `HV > H`. g (torch.Tensor): g (decays) of shape `[B, T, HV, K]`. + g_atk (torch.Tensor): + scalar gate for ATK in log space of shape `[B, T, H]`. + beta_atk (torch.Tensor): + betas for ATK of shape `[B, T, H]`. beta (torch.Tensor): betas of shape `[B, T, HV]`. A_log (Optional[torch.Tensor]): - Decay parameter of shape `[HV]`. Required when `use_gate_in_kernel=True`. + Per-head log decay parameter of shape `[H]`. Required when `use_gate_in_kernel=True`. + Default: `None`. dt_bias (Optional[torch.Tensor]): - Bias added to `g` before activation, of shape `[HV]`. - Only used when `use_gate_in_kernel=True`. + Per-head-key bias of shape `[H * K]`. Used when `use_gate_in_kernel=True`. + Default: `None`. scale (Optional[float]): - Scale factor for the RetNet attention scores. + Scale factor for attention scores. If not provided, it will default to `1 / sqrt(K)`. Default: `None`. initial_state (Optional[torch.Tensor]): Initial state of shape `[N, HV, K, V]` for `N` input sequences. For equal-length input sequences, `N` equals the batch size `B`. Default: `None`. + initial_A_state (Optional[torch.Tensor]): + Initial ATK diagonal state of shape `[N, H, K]`. Default: `None`. output_final_state (Optional[bool]): Whether to output the final state of shape `[N, HV, K, V]`. Default: `False`. use_qk_l2norm_in_kernel (Optional[bool]): Whether to use L2 normalization in the kernel. Default: `False`. use_gate_in_kernel (Optional[bool]): Whether to compute the log-space KDA decay internally. - When `True`, `g` is the raw input and `A_log` must be provided; the kernel fuses - gate activation into the recurrence. Default: `False`. - use_beta_sigmoid_in_kernel (Optional[bool]): - Whether to apply `torch.sigmoid(beta)` inside the kernel. - - If `True`, the passed `beta` acts as the raw beta logits. - - If `False`, `beta` is expected to already be in post-sigmoid space. - Default: `False`. - allow_neg_eigval (Optional[bool]): - Whether to allow negative eigenvalues by scaling `beta` to `[0, 2)`. - Only takes effect together with `use_beta_sigmoid_in_kernel=True`, in which case - the kernel computes `2 * sigmoid(beta)` instead of `sigmoid(beta)`. Default: `False`. + If `True`, `g` acts as the raw input for `-exp(A_log) * softplus(g + dt_bias)`. + `A_log` must be provided. Default: `False`. lower_bound (Optional[float]): - Lower bound for the forget gate (in log space). Only used when `use_gate_in_kernel=True`. Default: `None`. - state_v_first (Optional[bool]): - Store the recurrent state in V-first ``[V, K]`` layout instead of the default ``[K, V]``. Default: ``False``. + If provided, uses `lower_bound * sigmoid(exp(A_log) * g)` instead of + `-exp(A_log) * softplus(g)` for gate computation. Only used when + `use_gate_in_kernel=True`. Default: `None`. cu_seqlens (torch.LongTensor): Cumulative sequence lengths of shape `[N+1]` used for variable-length training, consistent with the FlashAttention API. + transpose_state_layout (bool): + Whether to use transposed state layout `[V, K]` instead of `[K, V]`. Default: `False`. + x (float): + Squash range parameter. Default: `1.5`. + eps (float): + Epsilon for numerical stability. Default: `1e-6`. + log_atk_scale (Optional[torch.Tensor]): + Per-head log-space center of shape `[H]`. Default: `None`. Returns: o (torch.Tensor): Outputs of shape `[B, T, HV, V]`. final_state (torch.Tensor): Final state of shape `[N, HV, K, V]` if `output_final_state=True` else `None`. + A_state (torch.Tensor): + Final ATK diagonal state of shape `[N, H, K]` if `output_final_state=True` else `None`. Examples:: >>> import torch >>> import torch.nn.functional as F >>> from einops import rearrange - >>> from ...ops.kda import fused_recurrent_kda + >>> from fla.ops.precond_kda import fused_recurrent_precond_kda # inputs with equal lengths >>> B, T, H, HV, K, V = 4, 2048, 4, 8, 512, 512 >>> q = torch.randn(B, T, H, K, device='cuda') >>> k = F.normalize(torch.randn(B, T, H, K, device='cuda'), p=2, dim=-1) >>> v = torch.randn(B, T, HV, V, device='cuda') >>> g = F.logsigmoid(torch.rand(B, T, HV, K, device='cuda')) + >>> g_atk = F.logsigmoid(torch.rand(B, T, H, device='cuda')) + >>> beta_atk = torch.rand(B, T, H, device='cuda').sigmoid() >>> beta = torch.rand(B, T, HV, device='cuda').sigmoid() >>> h0 = torch.randn(B, HV, K, V, device='cuda') - >>> o, ht = fused_recurrent_kda( - q, k, v, g, beta, + >>> o, ht, at = fused_recurrent_precond_kda( + q, k, v, g, g_atk, beta_atk, beta, initial_state=h0, output_final_state=True ) # for variable-length inputs, the batch size `B` is expected to be 1 and `cu_seqlens` is required - >>> q, k, v, g, beta = map(lambda x: rearrange(x, 'b t ... -> 1 (b t) ...'), (q, k, v, g, beta)) + >>> q, k, v, g = map(lambda x: rearrange(x, 'b t ... -> 1 (b t) ...'), (q, k, v, g)) + >>> g_atk, beta_atk, beta = map(lambda x: rearrange(x, 'b t ... -> 1 (b t) ...'), (g_atk, beta_atk, beta)) # for a batch with 4 sequences, `cu_seqlens` with 5 start/end positions are expected >>> cu_seqlens = q.new_tensor([0, 2048, 4096, 6144, 8192], dtype=torch.long) - >>> o_var, ht_var = fused_recurrent_kda( - q, k, v, g, beta, + >>> o_var, ht_var, at_var = fused_recurrent_precond_kda( + q, k, v, g, g_atk, beta_atk, beta, initial_state=h0, output_final_state=True, cu_seqlens=cu_seqlens ) """ - if 'transpose_state_layout' in kwargs: - if state_v_first: - raise ValueError("Cannot pass both `state_v_first` and the deprecated `transpose_state_layout`.") - warnings.warn( - "`transpose_state_layout` is deprecated and renamed to `state_v_first`.", - DeprecationWarning, - stacklevel=2, - ) - state_v_first = kwargs.pop('transpose_state_layout') if cu_seqlens is not None: if q.shape[0] != 1: @@ -461,27 +549,29 @@ def fused_recurrent_kda( ) if scale is None: scale = k.shape[-1] ** -0.5 - if allow_neg_eigval and not use_beta_sigmoid_in_kernel: - raise ValueError("`allow_neg_eigval=True` requires `use_beta_sigmoid_in_kernel=True`.") - o, final_state = fused_recurrent_kda_fwd( + o, final_state, final_A_state = fused_recurrent_precond_kda_fwd( q=q, k=k, v=v, g=g, + g_atk=g_atk, + beta_atk=beta_atk, beta=beta, A_log=A_log, dt_bias=dt_bias, scale=scale, initial_state=initial_state, + initial_A_state=initial_A_state, inplace_final_state=False, output_final_state=output_final_state, use_qk_l2norm_in_kernel=use_qk_l2norm_in_kernel, use_gate_in_kernel=use_gate_in_kernel, - use_beta_sigmoid_in_kernel=use_beta_sigmoid_in_kernel, - allow_neg_eigval=allow_neg_eigval, lower_bound=lower_bound, cu_seqlens=cu_seqlens, - state_v_first=state_v_first, + transpose_state_layout=transpose_state_layout, + x=x, + eps=eps, + log_atk_scale=log_atk_scale, ) - return o, final_state + return o, final_state, final_A_state diff --git a/build/torch-cuda/ops/precond_kda/naive.py b/build/torch-cuda/ops/precond_kda/naive.py new file mode 100644 index 0000000000000000000000000000000000000000..1ab39713e2fd9b81a2ec0eb84e9aa1f84e3ee0f3 --- /dev/null +++ b/build/torch-cuda/ops/precond_kda/naive.py @@ -0,0 +1,76 @@ +# 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 + + +def naive_recurrent_precond_kda( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + g: torch.Tensor, + g_atk: torch.Tensor, + beta_atk: torch.Tensor, + beta: torch.Tensor, + scale: float | None = None, + initial_state: torch.Tensor | None = None, + initial_A_state: torch.Tensor | None = None, + output_final_state: bool = False, + x: float = 1.5, + eps: float = 1e-6, + log_atk_scale: torch.Tensor | float = None, +): + dtype = v.dtype + B, T, H, K, V = *q.shape, v.shape[-1] + if scale is None: + scale = K ** -0.5 + + q, k, v, g, g_atk, beta_atk, beta = map( + lambda x: x.to(torch.float32), + [q, k, v, g, g_atk, beta_atk, beta] + ) + q = q * scale + + if log_atk_scale is None: + log_atk_scale = -0.2 + if isinstance(log_atk_scale, (int, float)): + center = torch.tensor(log_atk_scale, dtype=torch.float32, device=k.device) + else: + center = log_atk_scale.to(torch.float32).view(1, H, 1) + + logx = torch.log(torch.tensor(x, dtype=torch.float32, device=k.device)) + + A = k.new_zeros(B, H, K) + if initial_A_state is not None: + A += initial_A_state + + S = k.new_zeros(B, H, K, V) + if initial_state is not None: + S += initial_state + + o = torch.zeros_like(v) + + for i in range(T): + q_i, k_i, v_i, g_i = q[:, i], k[:, i], v[:, i], g[:, i] + g_atk_i, beta_atk_i, beta_i = g_atk[:, i], beta_atk[:, i], beta[:, i] + + A = g_atk_i.exp().unsqueeze(-1) * A + beta_atk_i.unsqueeze(-1) * (k_i ** 2) + + ell = torch.log(A + eps) + r = ell - center + s = r / (1.0 + torch.abs(r)) + M = torch.exp(-logx * s) + k_precond = k_i * M + + S = S * g_i[..., None].exp() + S = S + torch.einsum('b h k, b h v -> b h k v', beta_i[..., None] * k_precond, v_i - (k_i[..., None] * S).sum(-2)) + o[:, i] = torch.einsum('b h k, b h k v -> b h v', q_i, S) + + if not output_final_state: + S = None + A = None + return o.to(dtype), S, A diff --git a/build/torch-cuda/ops/precond_kda/wy_fast.py b/build/torch-cuda/ops/precond_kda/wy_fast.py new file mode 100644 index 0000000000000000000000000000000000000000..4ea73808dc81817ace9a6624a8b2008f6a67ca6f --- /dev/null +++ b/build/torch-cuda/ops/precond_kda/wy_fast.py @@ -0,0 +1,198 @@ +# 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 +import triton.language as tl + +from ...ops.utils import prepare_chunk_indices +from ...ops.utils.op import exp2 +from ...utils import autotune_cache_kwargs + + +@triton.heuristics({ + 'STORE_QG': lambda args: args['qg'] is not None, + 'STORE_KG': lambda args: args['kg'] is not None, + 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None, +}) +@triton.autotune( + configs=[ + triton.Config({}, num_warps=num_warps, num_stages=num_stages) + for num_warps in [2, 4, 8] + for num_stages in [2, 3, 4] + ], + key=['H', 'K', 'V', 'BT', 'BK', 'BV', 'IS_VARLEN'], + **autotune_cache_kwargs, +) +@triton.jit(do_not_specialize=['T']) +def recompute_w_u_fwd_kernel( + q, + k, + k_precond, + qg, + kg, + v, + beta, + w, + u, + A, + gk, + cu_seqlens, + chunk_indices, + T, + H: tl.constexpr, + K: tl.constexpr, + V: tl.constexpr, + BT: tl.constexpr, + BK: tl.constexpr, + BV: tl.constexpr, + STORE_QG: tl.constexpr, + STORE_KG: tl.constexpr, + IS_VARLEN: tl.constexpr, + DOT_PRECISION: tl.constexpr = 'tf32x3', +): + """ + Asymmetric WY representation: + - w = A @ (k * beta * exp(gk)) - uses original k for read/correction + - kg = k_precond * exp(gn - gk) - uses k_precond for write/h update + """ + i_t, i_bh = tl.program_id(0).to(tl.int64), tl.program_id(1).to(tl.int64) + i_b, i_h = i_bh // H, i_bh % H + 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.int64) + bos, eos = tl.load(cu_seqlens + i_n).to(tl.int64), tl.load(cu_seqlens + i_n + 1).to(tl.int64) + T = eos - bos + else: + bos, eos = i_b * T, i_b * T + T + + o_t = i_t * BT + tl.arange(0, BT) + m_t = o_t < T + + b_b = tl.load(beta + bos*H + i_h + o_t*H, mask=m_t, other=0.0) + + o_A = tl.arange(0, BT) + p_A = A + (bos*H + i_h) * BT + o_t[:, None] * (H*BT) + o_A[None, :] + b_A = tl.load(p_A, mask=m_t[:, None] & (o_A[None, :] < BT), other=0.0) + + # u computation (unchanged) + for i_v in range(tl.cdiv(V, BV)): + o_v = i_v * BV + tl.arange(0, BV) + m_tv = m_t[:, None] & (o_v[None, :] < V) + p_v = v + (bos*H + i_h) * V + o_t[:, None] * (H*V) + o_v[None, :] + p_u = u + (bos*H + i_h) * V + o_t[:, None] * (H*V) + o_v[None, :] + b_v = tl.load(p_v, mask=m_tv, other=0.0) + b_vb = (b_v * b_b[:, None]).to(b_v.dtype) + b_u = tl.dot(b_A, b_vb, input_precision=DOT_PRECISION) + tl.store(p_u, b_u.to(u.dtype.element_ty), mask=m_tv) + + for i_k in range(tl.cdiv(K, BK)): + o_k = i_k * BK + tl.arange(0, BK) + m_k = o_k < K + m_tk = m_t[:, None] & m_k[None, :] + p_w = w + (bos*H + i_h) * K + o_t[:, None] * (H*K) + o_k[None, :] + # Use original k for w (read/correction) + p_k = k + (bos*H + i_h) * K + o_t[:, None] * (H*K) + o_k[None, :] + b_k = tl.load(p_k, mask=m_tk, other=0.0) + b_kb = b_k * b_b[:, None] + + p_gk = gk + (bos*H + i_h) * K + o_t[:, None] * (H*K) + o_k[None, :] + b_gk = tl.load(p_gk, mask=m_tk, other=0.0) + b_kb *= exp2(b_gk) + + if STORE_QG: + p_q = q + (bos*H + i_h) * K + o_t[:, None] * (H*K) + o_k[None, :] + p_qg = qg + (bos*H + i_h) * K + o_t[:, None] * (H*K) + o_k[None, :] + b_q = tl.load(p_q, mask=m_tk, other=0.0) + b_qg = b_q * exp2(b_gk) + tl.store(p_qg, b_qg.to(qg.dtype.element_ty), mask=m_tk) + + if STORE_KG: + # Use k_precond for kg (write/h update) - ASYMMETRIC + p_kp = k_precond + (bos*H + i_h) * K + o_t[:, None] * (H*K) + o_k[None, :] + b_kp = tl.load(p_kp, mask=m_tk, other=0.0) + + last_idx = min(i_t * BT + BT, T) - 1 + b_gn = tl.load(gk + ((bos + last_idx) * H + i_h) * K + o_k, mask=m_k, other=0.) + # kg uses k_precond, not original k + b_kg = b_kp * tl.where((i_t * BT + tl.arange(0, BT) < T)[:, None], exp2(b_gn[None, :] - b_gk), 0) + p_kg = kg + (bos * H + i_h) * K + o_t[:, None] * (H*K) + o_k[None, :] + tl.store(p_kg, b_kg.to(kg.dtype.element_ty), mask=m_tk) + + # w uses original k + b_w = tl.dot(b_A, b_kb.to(b_k.dtype)) + tl.store(p_w, b_w.to(w.dtype.element_ty), mask=m_tk) + + +def recompute_w_u_fwd( + k: torch.Tensor, + k_precond: torch.Tensor, + v: torch.Tensor, + beta: torch.Tensor, + A: torch.Tensor, + gk: torch.Tensor, + cu_seqlens: torch.LongTensor | None = None, + chunk_indices: torch.LongTensor | None = None, + q: torch.Tensor | None = None, + output_qg: bool = False, + output_kg: bool = True, +): + """ + Compute WY representation for preconditioned KDA. + + Args: + k: [B, T, H, K] - original k (for w computation) + k_precond: [B, T, H, K] - preconditioned k (for kg computation) + v: [B, T, H, V] - values + beta: [B, T, H] - beta scaling + A: [B, T, H, BT] - inverse of Akk + gk: [B, T, H, K] - cumsum of gates + + Returns: + w: [B, T, H, K] - A @ (k * beta * exp(gk)) + u: [B, T, H, V] - A @ (v * beta) + qg: [B, T, H, K] or None - q * exp(gk) + kg: [B, T, H, K] or None - k_precond * exp(gn - gk) for h update + """ + B, T, H, K = k.shape + V = v.shape[-1] + BT = A.shape[-1] + BK = 64 + BV = 64 + + if chunk_indices is None and cu_seqlens is not None: + chunk_indices = prepare_chunk_indices(cu_seqlens, BT) + NT = triton.cdiv(T, BT) if cu_seqlens is None else len(chunk_indices) + + w = torch.empty_like(k) + u = torch.empty_like(v) + qg = torch.empty_like(q) if output_qg and q is not None else None + kg = torch.empty_like(k_precond) if output_kg else None + + grid = (NT, B * H) + recompute_w_u_fwd_kernel[grid]( + q=q, + k=k, + k_precond=k_precond, + qg=qg, + kg=kg, + v=v, + beta=beta, + w=w, + u=u, + A=A, + gk=gk, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + T=T, + H=H, + K=K, + V=V, + BT=BT, + BK=BK, + BV=BV, + ) + return w, u, qg, kg diff --git a/build/torch-cuda/ops/rebased/parallel.py b/build/torch-cuda/ops/rebased/parallel.py index e4945e8f2c32ec859acdeffcea3fcff2d7076ef3..80ed15156d855a92093572b539369e7f4f164878 100644 --- a/build/torch-cuda/ops/rebased/parallel.py +++ b/build/torch-cuda/ops/rebased/parallel.py @@ -9,6 +9,7 @@ import torch import triton import triton.language as tl +from ...ops.utils.op import unflatten_program_id from ...utils import autocast_custom_bwd, autocast_custom_fwd, input_guard # Rebased: Linear Transformers with Learnable Kernel Functions are Better In-Context Models @@ -34,38 +35,46 @@ def parallel_rebased_fwd_kernel( BV: tl.constexpr, ): # i_c: chunk index. used for sequence parallelism - i_kv, i_c, i_bh = tl.program_id(0), tl.program_id(1), tl.program_id(2) + i_kv, i_c = unflatten_program_id(tl.cdiv(K, BK) * tl.cdiv(V, BV)) + i_bh = tl.program_id(1).to(tl.int64) NV = tl.cdiv(V, BV) i_k = i_kv // (NV) i_v = i_kv % (NV) - p_q = tl.make_block_ptr(q + i_bh * T*K, (T, K), (K, 1), (i_c*BTL, i_k*BK), (BTL, BK), (1, 0)) - p_k = tl.make_block_ptr(k + i_bh * T*K, (K, T), (1, K), (i_k*BK, 0), (BK, BTS), (0, 1)) - p_v = tl.make_block_ptr(v + i_bh * T*V, (T, V), (V, 1), (0, i_v*BV), (BTS, BV), (1, 0)) + o_t = i_c * BTL + tl.arange(0, BTL) + o_kk = i_k * BK + tl.arange(0, BK) + o_vv = i_v * BV + tl.arange(0, BV) + m_t = o_t < T + m_q = m_t[:, None] & (o_kk[None, :] < K) + m_o = m_t[:, None] & (o_vv[None, :] < V) + p_q = q + i_bh * T*K + o_t[:, None] * K + o_kk[None, :] # [BQ, BD] block Q, in the shared memory throughout the whole kernel - b_q = tl.load(p_q, boundary_check=(0, 1)) + b_q = tl.load(p_q, mask=m_q, other=0.0) b_q = (b_q * scale).to(b_q.dtype) b_o = tl.zeros([BTL, BV], dtype=tl.float32) b_z = tl.zeros([BTL], dtype=tl.float32) # Q block and K block have no overlap # no need for mask, thereby saving flops - for _ in range(0, i_c*BTL, BTS): + for i_s in range(0, i_c*BTL, BTS): + o_s = i_s + tl.arange(0, BTS) + m_k = (o_kk[:, None] < K) & (o_s[None, :] < T) + m_v = (o_s[:, None] < T) & (o_vv[None, :] < V) + p_k = k + i_bh * T*K + o_kk[:, None] + o_s[None, :] * K + p_v = v + i_bh * T*V + o_s[:, None] * V + o_vv[None, :] # [BK, BTS] - b_k = tl.load(p_k, boundary_check=(0, 1)) + b_k = tl.load(p_k, mask=m_k, other=0.0) # [BTS, BV] - b_v = tl.load(p_v, boundary_check=(0, 1)) + b_v = tl.load(p_v, mask=m_v, other=0.0) # [BTL, BTS] b_s = tl.dot(b_q, (b_k), allow_tf32=False) b_s = b_s * b_s b_z += tl.sum(b_s, axis=1) # [BQ, BD] - b_o = b_o + tl.dot(b_s.to(b_v.dtype), b_v, allow_tf32=False) - p_k = tl.advance(p_k, (0, BTS)) - p_v = tl.advance(p_v, (BTS, 0)) + b_o = tl.dot(b_s.to(b_v.dtype), b_v, b_o, allow_tf32=False) # # rescale interchunk output tl.debug_barrier() @@ -74,14 +83,17 @@ def parallel_rebased_fwd_kernel( # tl.debug_barrier() o_k = tl.arange(0, BTS) - p_k = tl.make_block_ptr(k + i_bh * T*K, (K, T), (1, K), (i_k*BK, i_c*BTL), (BK, BTS), (0, 1)) - p_v = tl.make_block_ptr(v + i_bh * T*V, (T, V), (V, 1), (i_c*BTL, i_v*BV), (BTS, BV), (1, 0)) # Q block and K block have overlap. masks required - for _ in range(i_c*BTL, (i_c + 1) * BTL, BTS): + for i_s in range(i_c*BTL, (i_c + 1) * BTL, BTS): + o_s = i_s + tl.arange(0, BTS) + m_k = (o_kk[:, None] < K) & (o_s[None, :] < T) + m_v = (o_s[:, None] < T) & (o_vv[None, :] < V) + p_k = k + i_bh * T*K + o_kk[:, None] + o_s[None, :] * K + p_v = v + i_bh * T*V + o_s[:, None] * V + o_vv[None, :] # [BK, BTS] - b_k = tl.load(p_k, boundary_check=(0, 1)) + b_k = tl.load(p_k, mask=m_k, other=0.0) # [BTS, BV] - b_v = tl.load(p_v, boundary_check=(0, 1)) + b_v = tl.load(p_v, mask=m_v, other=0.0) # [BTL, BTS] m_s = o_q[:, None] >= o_k[None, :] b_s = tl.dot(b_q, b_k, allow_tf32=False) @@ -89,14 +101,12 @@ def parallel_rebased_fwd_kernel( b_s = tl.where(m_s, b_s, 0) b_z += tl.sum(b_s, axis=1) # [BTL, BV] - b_o += tl.dot(b_s.to(b_q.dtype), b_v, allow_tf32=False) - p_k = tl.advance(p_k, (0, BTS)) - p_v = tl.advance(p_v, (BTS, 0)) + b_o = tl.dot(b_s.to(b_q.dtype), b_v, b_o, allow_tf32=False) o_k += BTS - p_o = tl.make_block_ptr(o + (i_bh + B * H * i_k) * T*V, (T, V), (V, 1), (i_c*BTL, i_v*BV), (BTL, BV), (1, 0)) + p_o = o + (i_bh + B * H * i_k) * T*V + o_t[:, None] * V + o_vv[None, :] p_z = z + (i_bh + B * H * i_k) * T + i_c*BTL + tl.arange(0, BTL) - tl.store(p_o, b_o.to(p_o.dtype.element_ty), boundary_check=(0, 1)) + tl.store(p_o, b_o.to(p_o.dtype.element_ty), mask=m_o) tl.store(p_z, b_z.to(p_z.dtype.element_ty), mask=((i_c*BTL + tl.arange(0, BTL)) < T)) @@ -124,22 +134,31 @@ def _parallel_rebased_bwd_dq( BK: tl.constexpr, BV: tl.constexpr, ): - p_do = tl.make_block_ptr(do + i_bh * T*V, (T, V), (V, 1), (i_c*BTL, i_v*BV), (BTL, BV), (1, 0)) - p_q = tl.make_block_ptr(q + (i_bh) * T*K, (T, K), (K, 1), (i_c*BTL, i_k*BK), (BTL, BK), (1, 0)) - b_q = tl.load(p_q, boundary_check=(0, 1)) - b_do = tl.load(p_do, boundary_check=(0, 1)).to(b_q.dtype) + o_t = i_c * BTL + tl.arange(0, BTL) + o_kk = i_k * BK + tl.arange(0, BK) + o_vv = i_v * BV + tl.arange(0, BV) + m_t = o_t < T + m_q = m_t[:, None] & (o_kk[None, :] < K) + m_o = m_t[:, None] & (o_vv[None, :] < V) + p_do = do + i_bh * T*V + o_t[:, None] * V + o_vv[None, :] + p_q = q + (i_bh) * T*K + o_t[:, None] * K + o_kk[None, :] + b_q = tl.load(p_q, mask=m_q, other=0.0) + b_do = tl.load(p_do, mask=m_o, other=0.0).to(b_q.dtype) b_q = (b_q * scale).to(b_q.dtype) b_dq = tl.zeros([BTL, BK], dtype=tl.float32) - p_k = tl.make_block_ptr(k + i_bh * T*K, (T, K), (K, 1), (0, i_k*BK), (BTS, BK), (1, 0)) - p_v = tl.make_block_ptr(v + i_bh * T*V, (V, T), (1, V), (i_v*BV, 0), (BV, BTS), (0, 1)) p_dz = dz + i_bh * T + i_c*BTL + tl.arange(0, BTL) b_dz = tl.load(p_dz, mask=(i_c*BTL + tl.arange(0, BTL)) < T) - for _ in range(0, i_c*BTL, BTS): + for i_s in range(0, i_c*BTL, BTS): + o_s = i_s + tl.arange(0, BTS) + m_k = (o_s[:, None] < T) & (o_kk[None, :] < K) + m_v = (o_vv[:, None] < V) & (o_s[None, :] < T) + p_k = k + i_bh * T*K + o_s[:, None] * K + o_kk[None, :] + p_v = v + i_bh * T*V + o_vv[:, None] + o_s[None, :] * V # [BTS, BK] - b_k = tl.load(p_k, boundary_check=(0, 1)) + b_k = tl.load(p_k, mask=m_k, other=0.0) # [BV, BTS] - b_v = tl.load(p_v, boundary_check=(0, 1)) + b_v = tl.load(p_v, mask=m_v, other=0.0) # [BTL, BTS] b_ds = tl.dot(b_do, b_v, allow_tf32=False) if i_v == 0: @@ -148,21 +167,22 @@ def _parallel_rebased_bwd_dq( b_ds = b_ds b_s = tl.dot(b_q, tl.trans(b_k), allow_tf32=False) # [BQ, BD] - b_dq += tl.dot((2 * b_ds * b_s).to(b_v.dtype), b_k, allow_tf32=False) - p_k = tl.advance(p_k, (BTS, 0)) - p_v = tl.advance(p_v, (0, BTS)) + b_dq = tl.dot((2 * b_ds * b_s).to(b_v.dtype), b_k, b_dq, allow_tf32=False) b_dq *= scale o_q = tl.arange(0, BTL) o_k = tl.arange(0, BTS) - p_k = tl.make_block_ptr(k + i_bh * T*K, (T, K), (K, 1), (i_c*BTL, i_k*BK), (BTS, BK), (1, 0)) - p_v = tl.make_block_ptr(v + i_bh * T*V, (V, T), (1, V), (i_v*BV, i_c*BTL), (BV, BTS), (0, 1)) # Q block and K block have overlap. masks required - for _ in range(i_c*BTL, (i_c + 1) * BTL, BTS): + for i_s in range(i_c*BTL, (i_c + 1) * BTL, BTS): + o_s = i_s + tl.arange(0, BTS) + m_k = (o_s[:, None] < T) & (o_kk[None, :] < K) + m_v = (o_vv[:, None] < V) & (o_s[None, :] < T) + p_k = k + i_bh * T*K + o_s[:, None] * K + o_kk[None, :] + p_v = v + i_bh * T*V + o_vv[:, None] + o_s[None, :] * V # [BTS, BK] - b_k = tl.load(p_k, boundary_check=(0, 1)) + b_k = tl.load(p_k, mask=m_k, other=0.0) # [BV, BTS] - b_v = tl.load(p_v, boundary_check=(0, 1)) + b_v = tl.load(p_v, mask=m_v, other=0.0) # [BTL, BTS] m_s = o_q[:, None] >= o_k[None, :] b_ds = tl.dot(b_do, b_v, allow_tf32=False) @@ -174,13 +194,11 @@ def _parallel_rebased_bwd_dq( b_s = tl.dot(b_q, tl.trans(b_k), allow_tf32=False) b_s = tl.where(m_s, b_s, 0) # [BTL, BK] - b_dq += tl.dot((2 * b_ds * b_s).to(b_k.dtype), - b_k, allow_tf32=False) - p_k = tl.advance(p_k, (BTS, 0)) - p_v = tl.advance(p_v, (0, BTS)) + b_dq = tl.dot((2 * b_ds * b_s).to(b_k.dtype), + b_k, b_dq, allow_tf32=False) o_k += BTS - p_dq = tl.make_block_ptr(dq + (i_bh + B * H * i_v) * T*K, (T, K), (K, 1), (i_c*BTL, i_k*BK), (BTL, BK), (1, 0)) - tl.store(p_dq, b_dq.to(p_dq.dtype.element_ty), boundary_check=(0, 1)) + p_dq = dq + (i_bh + B * H * i_v) * T*K + o_t[:, None] * K + o_kk[None, :] + tl.store(p_dq, b_dq.to(p_dq.dtype.element_ty), mask=m_q) return @@ -210,40 +228,52 @@ def _parallel_rebased_bwd_dkv( BV: tl.constexpr, ): # compute dk dv - p_k = tl.make_block_ptr(k + i_bh * T*K, (T, K), (K, 1), (i_c*BTL, i_k*BK), (BTL, BK), (1, 0)) - p_v = tl.make_block_ptr(v + i_bh * T*V, (T, V), (V, 1), (i_c*BTL, i_v*BV), (BTL, BV), (1, 0)) - b_k, b_v = tl.load(p_k, boundary_check=(0, 1)), tl.load(p_v, boundary_check=(0, 1)) + o_t = i_c * BTL + tl.arange(0, BTL) + o_kk = i_k * BK + tl.arange(0, BK) + o_vv = i_v * BV + tl.arange(0, BV) + m_t = o_t < T + m_k = m_t[:, None] & (o_kk[None, :] < K) + m_v = m_t[:, None] & (o_vv[None, :] < V) + p_k = k + i_bh * T*K + o_t[:, None] * K + o_kk[None, :] + p_v = v + i_bh * T*V + o_t[:, None] * V + o_vv[None, :] + b_k, b_v = tl.load(p_k, mask=m_k, other=0.0), tl.load(p_v, mask=m_v, other=0.0) b_dk, b_dv = tl.zeros([BTL, BK], dtype=tl.float32), tl.zeros( [BTL, BV], dtype=tl.float32) for i in range((tl.cdiv(T, BTS) * BTS)-BTS, (i_c + 1) * BTL - BTS, -BTS): - p_q = tl.make_block_ptr(q + i_bh * T*K, (K, T), (1, K), (i_k*BK, i), (BK, BTS), (0, 1)) - p_do = tl.make_block_ptr(do + i_bh * T*V, (V, T), (1, V), (i_v*BV, i), (BV, BTS), (0, 1)) + o_s = i + tl.arange(0, BTS) + m_q = (o_kk[:, None] < K) & (o_s[None, :] < T) + m_o = (o_vv[:, None] < V) & (o_s[None, :] < T) + p_q = q + i_bh * T*K + o_kk[:, None] + o_s[None, :] * K + p_do = do + i_bh * T*V + o_vv[:, None] + o_s[None, :] * V p_dz = dz + i_bh * T + i + tl.arange(0, BTS) # [BK, BTS] - b_q = tl.load(p_q, boundary_check=(0, 1)) + b_q = tl.load(p_q, mask=m_q, other=0.0) # [BV, BTS] - b_do = tl.load(p_do, boundary_check=(0, 1)).to(b_q.dtype) + b_do = tl.load(p_do, mask=m_o, other=0.0).to(b_q.dtype) b_dz = tl.load(p_dz, mask=(i + tl.arange(0, BTS)) < T) # [BTL, BTS] b_s = tl.dot(b_k.to(b_q.dtype), b_q, allow_tf32=False) * scale b_s2 = b_s * b_s - b_dv += tl.dot(b_s2.to(b_q.dtype), tl.trans(b_do), allow_tf32=False) + b_dv = tl.dot(b_s2.to(b_q.dtype), tl.trans(b_do), b_dv, allow_tf32=False) b_ds = tl.dot(b_v, b_do, allow_tf32=False) * scale if i_v == 0: b_ds += b_dz[None, :] * scale else: b_ds = b_ds - b_dk += tl.dot((2 * b_ds * b_s).to(b_q.dtype), tl.trans(b_q), allow_tf32=False) + b_dk = tl.dot((2 * b_ds * b_s).to(b_q.dtype), tl.trans(b_q), b_dk, allow_tf32=False) tl.debug_barrier() o_q, o_k = tl.arange(0, BTS), tl.arange(0, BTL) for i in range(i_c*BTL, (i_c+1)*BTL, BTS): - p_q = tl.make_block_ptr(q + i_bh * T*K, (K, T), (1, K), (i_k*BK, i), (BK, BTS), (0, 1)) - p_do = tl.make_block_ptr(do + i_bh * T*V, (V, T), (1, V), (i_v*BV, i), (BV, BTS), (0, 1)) + o_s = i + tl.arange(0, BTS) + m_q = (o_kk[:, None] < K) & (o_s[None, :] < T) + m_o = (o_vv[:, None] < V) & (o_s[None, :] < T) + p_q = q + i_bh * T*K + o_kk[:, None] + o_s[None, :] * K + p_do = do + i_bh * T*V + o_vv[:, None] + o_s[None, :] * V p_dz = dz + i_bh * T + i + tl.arange(0, BTS) - b_q = tl.load(p_q, boundary_check=(0, 1)) # [BD, BQ] - b_do = tl.load(p_do, boundary_check=(0, 1)).to(b_q.dtype) + b_q = tl.load(p_q, mask=m_q, other=0.0) # [BD, BQ] + b_do = tl.load(p_do, mask=m_o, other=0.0).to(b_q.dtype) b_dz = tl.load(p_dz, mask=(i + tl.arange(0, BTS)) < T) # [BK, BQ] m_s = o_k[:, None] <= o_q[None, :] @@ -259,14 +289,14 @@ def _parallel_rebased_bwd_dkv( b_ds = b_ds b_ds = tl.where(m_s, b_ds, 0) * scale # [BK, BD] - b_dv += tl.dot(b_s2.to(b_q.dtype), tl.trans(b_do), allow_tf32=False) - b_dk += tl.dot((2 * b_ds * b_s).to(b_q.dtype), tl.trans(b_q), allow_tf32=False) + b_dv = tl.dot(b_s2.to(b_q.dtype), tl.trans(b_do), b_dv, allow_tf32=False) + b_dk = tl.dot((2 * b_ds * b_s).to(b_q.dtype), tl.trans(b_q), b_dk, allow_tf32=False) o_q += BTS - p_dk = tl.make_block_ptr(dk + (i_bh + B * H * i_v) * T*K, (T, K), (K, 1), (i_c*BTL, i_k*BK), (BTL, BK), (1, 0)) - p_dv = tl.make_block_ptr(dv + (i_bh + B * H * i_k) * T*V, (T, V), (V, 1), (i_c*BTL, i_v*BV), (BTL, BV), (1, 0)) - tl.store(p_dk, b_dk.to(p_dk.dtype.element_ty), boundary_check=(0, 1)) - tl.store(p_dv, b_dv.to(p_dv.dtype.element_ty), boundary_check=(0, 1)) + p_dk = dk + (i_bh + B * H * i_v) * T*K + o_t[:, None] * K + o_kk[None, :] + p_dv = dv + (i_bh + B * H * i_k) * T*V + o_t[:, None] * V + o_vv[None, :] + tl.store(p_dk, b_dk.to(p_dk.dtype.element_ty), mask=m_k) + tl.store(p_dv, b_dv.to(p_dv.dtype.element_ty), mask=m_v) return @@ -291,7 +321,8 @@ def parallel_rebased_bwd_kernel( BK: tl.constexpr, BV: tl.constexpr, ): - i_kv, i_c, i_bh = tl.program_id(0), tl.program_id(1), tl.program_id(2) + i_kv, i_c = unflatten_program_id(tl.cdiv(K, BK) * tl.cdiv(V, BV)) + i_bh = tl.program_id(1).to(tl.int64) NV = tl.cdiv(V, BV) i_k = i_kv // (NV) i_v = i_kv % (NV) @@ -362,7 +393,7 @@ class ParallelBasedFunction(torch.autograd.Function): num_warps = 4 NK = triton.cdiv(K, BK) NV = triton.cdiv(V, BV) - grid = (NK * NV, triton.cdiv(T, BTL), B * H) + grid = (NK * NV * triton.cdiv(T, BTL), B * H) assert NK == 1, "will encounter some synchronization issue if not." @@ -406,7 +437,7 @@ class ParallelBasedFunction(torch.autograd.Function): num_warps = 4 NK = triton.cdiv(K, BK) NV = triton.cdiv(V, BV) - grid = (NK * NV, triton.cdiv(T, BTL), B * H) + grid = (NK * NV * triton.cdiv(T, BTL), B * H) assert NK == 1, "will encounter some synchronization issue if not" @@ -448,20 +479,17 @@ def parallel_rebased( use_scale: bool = True, use_normalize: bool = True, return_both: bool = False, - head_first: bool = False, ): assert q.shape[-1] <= 128, "only support feature dim up to 128" if use_scale: scale = q.shape[-1] ** -0.5 else: scale = 1 - if not head_first: - q, k, v = map(lambda x: x.transpose(1, 2), (q, k, v)) + q, k, v = map(lambda x: x.transpose(1, 2), (q, k, v)) o, z = ParallelBasedFunction.apply(q, k, v, scale) if return_both: return o, z if use_normalize: o = o / (z[..., None] + eps) - if not head_first: - o = o.transpose(1, 2) + o = o.transpose(1, 2) return o.to(q.dtype) diff --git a/build/torch-cuda/ops/retention/chunk.py b/build/torch-cuda/ops/retention/chunk.py index 9afc986508228684ddb89e312e9b50fe24437f2e..38e2870ff0199f9a7e040356935fd22d66acec5c 100644 --- a/build/torch-cuda/ops/retention/chunk.py +++ b/build/torch-cuda/ops/retention/chunk.py @@ -65,7 +65,7 @@ def chunk_retention( f"The number of initial states is expected to be equal to the number of input sequences, " f"i.e., {len(cu_seqlens) - 1} rather than {initial_state.shape[0]}.", ) - g_gamma = (1 - q.new_tensor(2., dtype=torch.float).pow(-5. - q.new_tensor(range(q.shape[2]), dtype=torch.float))).log() + g_gamma = (1 - 2.0 ** (-5.0 - torch.arange(q.shape[2], dtype=torch.float, device=q.device))).log() o, final_state = chunk_simple_gla( q=q, k=k, diff --git a/build/torch-cuda/ops/retention/fused_chunk.py b/build/torch-cuda/ops/retention/fused_chunk.py index 603b4f5188c6f9a362907734821082ebf4409937..2733d7bf75ef86637ff2ab0396398fdc61765d7b 100644 --- a/build/torch-cuda/ops/retention/fused_chunk.py +++ b/build/torch-cuda/ops/retention/fused_chunk.py @@ -19,7 +19,6 @@ def fused_chunk_retention( initial_state: torch.Tensor | None = None, output_final_state: bool = False, cu_seqlens: torch.LongTensor | None = None, - **kwargs, ) -> tuple[torch.Tensor, torch.Tensor]: r""" Args: @@ -49,10 +48,6 @@ def fused_chunk_retention( Final state of shape `[N, H, K, V]` if `output_final_state=True` else `None`. """ - if 'head_first' in kwargs: - raise DeprecationWarning( - "head_first has been removed. Inputs must be in `[B, T, H, ...]` format.", - ) if cu_seqlens is not None: if q.shape[0] != 1: raise ValueError( @@ -64,7 +59,7 @@ def fused_chunk_retention( f"The number of initial states is expected to be equal to the number of input sequences, " f"i.e., {len(cu_seqlens) - 1} rather than {initial_state.shape[0]}.", ) - g_gamma = (1 - q.new_tensor(2., dtype=torch.float).pow(-5. - q.new_tensor(range(q.shape[2]), dtype=torch.float))).log() + g_gamma = (1 - 2.0 ** (-5.0 - torch.arange(q.shape[2], dtype=torch.float, device=q.device))).log() o, final_state = fused_chunk_simple_gla( q=q, k=k, diff --git a/build/torch-cuda/ops/retention/fused_recurrent.py b/build/torch-cuda/ops/retention/fused_recurrent.py index 5f141640e74ce953133fbf3378fe054bd99094b3..252340098fca056d105e5aabfb792b49686ac96c 100644 --- a/build/torch-cuda/ops/retention/fused_recurrent.py +++ b/build/torch-cuda/ops/retention/fused_recurrent.py @@ -45,7 +45,7 @@ def fused_recurrent_retention( f"The number of initial states is expected to be equal to the number of input sequences, " f"i.e., {len(cu_seqlens) - 1} rather than {initial_state.shape[0]}.", ) - g_gamma = (1 - q.new_tensor(2., dtype=torch.float).pow(-5. - q.new_tensor(range(q.shape[2]), dtype=torch.float))).log() + g_gamma = (1 - 2.0 ** (-5.0 - torch.arange(q.shape[2], dtype=torch.float, device=q.device))).log() o, final_state = fused_recurrent_simple_gla( q=q, k=k, diff --git a/build/torch-cuda/ops/retention/parallel.py b/build/torch-cuda/ops/retention/parallel.py index 2c65b409b8d82bd6cd0a37ba09ac1ae36272a3a1..016dc9ec91fd5f48fe475ab895d42eafab776caf 100644 --- a/build/torch-cuda/ops/retention/parallel.py +++ b/build/torch-cuda/ops/retention/parallel.py @@ -17,7 +17,6 @@ def parallel_retention( scale: float | None = None, output_attentions: bool = False, cu_seqlens: torch.LongTensor | None = None, - **kwargs, ) -> tuple[torch.Tensor, torch.Tensor]: r""" Args: @@ -42,16 +41,12 @@ def parallel_retention( attn (torch.Tensor): Attention scores of shape `[B, H, T, T]` if `output_attentions=True` else `None`. """ - if 'head_first' in kwargs: - raise DeprecationWarning( - "head_first has been removed. Inputs must be in `[B, T, H, ...]` format.", - ) if cu_seqlens is not None and q.shape[0] != 1: raise ValueError( f"The batch size is expected to be 1 rather than {q.shape[0]} when using `cu_seqlens`. " f"Please flatten variable-length inputs before processing.", ) - s = (1 - q.new_tensor(2., dtype=torch.float).pow(-5. - q.new_tensor(range(q.shape[2]), dtype=torch.float))).log() + s = (1 - 2.0 ** (-5.0 - torch.arange(q.shape[2], dtype=torch.float, device=q.device))).log() g = s[None, None, :].expand(q.shape[0], q.shape[1], q.shape[2]) o, attn = parallel_simple_gla( diff --git a/build/torch-cuda/ops/rwkv4/fused_recurrent.py b/build/torch-cuda/ops/rwkv4/fused_recurrent.py index c88f1ec0d7f56a916e33ea011571086e583c9ce0..d68c64eadcb9d69e1506ff429eba58311e8cd716 100644 --- a/build/torch-cuda/ops/rwkv4/fused_recurrent.py +++ b/build/torch-cuda/ops/rwkv4/fused_recurrent.py @@ -65,8 +65,8 @@ def fused_recurrent_rwkv4_forward_kernel( BLOCK_SIZE_C: tl.constexpr, ): # Parallelize over the batch dimension. - b_idx = tl.program_id(0) - c_idx = tl.program_id(1) + b_idx = tl.program_id(0).to(tl.int64) + c_idx = tl.program_id(1).to(tl.int64) cs = (c_idx * BLOCK_SIZE_C) + tl.arange(0, BLOCK_SIZE_C) cmask = cs < chans @@ -239,8 +239,8 @@ def fused_recurrent_rwkv4_backward_kernel( BLOCK_SIZE_C: tl.constexpr, ): # Parallelize over the batch dimension. - b_idx = tl.program_id(0) - c_idx = tl.program_id(1) + b_idx = tl.program_id(0).to(tl.int64) + c_idx = tl.program_id(1).to(tl.int64) cs = (c_idx * BLOCK_SIZE_C) + tl.arange(0, BLOCK_SIZE_C) cmask = cs < chans diff --git a/build/torch-cuda/ops/rwkv6/backends/__init__.py b/build/torch-cuda/ops/rwkv6/backends/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..ab7b7eda6f72f9af736fbe916afadc13d4c44f7a --- /dev/null +++ b/build/torch-cuda/ops/rwkv6/backends/__init__.py @@ -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 + +"""RWKV6 backends.""" + + +__all__: list[str] = [] diff --git a/build/torch-cuda/ops/rwkv6/chunk.py b/build/torch-cuda/ops/rwkv6/chunk.py index 7f930426b7d64bd68503bd807e65d9c625a0bb8d..21e1250a275eb1a800783fe72af388b35fb09f5e 100644 --- a/build/torch-cuda/ops/rwkv6/chunk.py +++ b/build/torch-cuda/ops/rwkv6/chunk.py @@ -9,11 +9,12 @@ import torch import triton import triton.language as tl +from ...backends import dispatch from ...ops.common.chunk_h import chunk_fwd_h from ...ops.gla.chunk import chunk_gla_bwd_dA, chunk_gla_bwd_dv, chunk_gla_fwd_o_gk from ...ops.utils import prepare_chunk_indices, prepare_chunk_offsets from ...ops.utils.constant import RCP_LN2 -from ...ops.utils.op import exp2 +from ...ops.utils.op import exp2, unflatten_program_id from ...utils import ( autocast_custom_bwd, autocast_custom_fwd, @@ -56,11 +57,12 @@ def chunk_rwkv6_fwd_cumsum_kernel( HAS_SCALE: tl.constexpr, IS_VARLEN: tl.constexpr, ): - i_s, i_t, i_bh = tl.program_id(0), tl.program_id(1), tl.program_id(2) + i_s, i_t = unflatten_program_id(tl.cdiv(S, BS)) + i_bh = tl.program_id(1).to(tl.int64) i_b, i_h = i_bh // H, i_bh % H 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.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) + i_n, i_t = tl.load(chunk_indices + i_t * 2).to(tl.int32), tl.load(chunk_indices + i_t * 2 + 1).to(tl.int64) + bos, eos = tl.load(cu_seqlens + i_n).to(tl.int64), tl.load(cu_seqlens + i_n + 1).to(tl.int64) T = eos - bos else: bos, eos = i_b * T, i_b * T + T @@ -69,19 +71,22 @@ def chunk_rwkv6_fwd_cumsum_kernel( m_i = tl.where(o_i[:, None] >= o_i[None, :], 1., 0.).to(tl.float32) m_e = tl.where(o_i[:, None] > o_i[None, :], 1., 0.).to(tl.float32) - p_s = tl.make_block_ptr(s + (bos * H + i_h) * S, (T, S), (H*S, 1), (i_t * BT, i_s * BS), (BT, BS), (1, 0)) - p_oi = tl.make_block_ptr(oi + (bos * H + i_h) * S, (T, S), (H*S, 1), (i_t * BT, i_s * BS), (BT, BS), (1, 0)) - p_oe = tl.make_block_ptr(oe + (bos * H + i_h) * S, (T, S), (H*S, 1), (i_t * BT, i_s * BS), (BT, BS), (1, 0)) + o_t = i_t * BT + tl.arange(0, BT) + o_s = i_s * BS + tl.arange(0, BS) + m_s = (o_t[:, None] < T) & (o_s[None, :] < S) + p_s = s + (bos * H + i_h) * S + o_t[:, None] * (H*S) + o_s[None, :] + p_oi = oi + (bos * H + i_h) * S + o_t[:, None] * (H*S) + o_s[None, :] + p_oe = oe + (bos * H + i_h) * S + o_t[:, None] * (H*S) + o_s[None, :] # [BT, BS] - b_s = tl.load(p_s, boundary_check=(0, 1)).to(tl.float32) + b_s = tl.load(p_s, mask=m_s, other=0.0).to(tl.float32) b_oi = tl.dot(m_i, b_s) b_oe = tl.dot(m_e, b_s) if HAS_SCALE: # Pre-scale by RCP_LN2 so downstream kernels can use exp2 directly. b_oi = b_oi * scale b_oe = b_oe * scale - tl.store(p_oi, b_oi.to(p_oi.dtype.element_ty, fp_downcast_rounding="rtne"), boundary_check=(0, 1)) - tl.store(p_oe, b_oe.to(p_oe.dtype.element_ty, fp_downcast_rounding="rtne"), boundary_check=(0, 1)) + tl.store(p_oi, b_oi.to(p_oi.dtype.element_ty, fp_downcast_rounding="rtne"), mask=m_s) + tl.store(p_oe, b_oe.to(p_oe.dtype.element_ty, fp_downcast_rounding="rtne"), mask=m_s) def chunk_rwkv6_fwd_cumsum( @@ -98,8 +103,11 @@ def chunk_rwkv6_fwd_cumsum( NT = triton.cdiv(T, BT) if cu_seqlens is None else len(chunk_indices) gi, ge = torch.empty_like(g, dtype=torch.float), torch.empty_like(g, dtype=torch.float) - def grid(meta): return (triton.cdiv(meta['S'], meta['BS']), NT, B * H) - # keep cummulative normalizer in fp32 + + def grid(meta): + return (triton.cdiv(meta['S'], meta['BS']) * NT, B * H) + + # keep cumulative normalizer in fp32 chunk_rwkv6_fwd_cumsum_kernel[grid]( g, gi, @@ -147,12 +155,12 @@ def chunk_rwkv6_fwd_A_kernel_intra_sub_inter( NC: tl.constexpr, IS_VARLEN: tl.constexpr, ): - i_t, i_c, i_bh = tl.program_id(0), tl.program_id(1), tl.program_id(2) + i_t, i_c, i_bh = tl.program_id(0).to(tl.int64), tl.program_id(1).to(tl.int64), tl.program_id(2).to(tl.int64) i_b, i_h = i_bh // H, i_bh % H i_i, i_j = i_c // NC, i_c % NC 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.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) + i_n, i_t = tl.load(chunk_indices + i_t * 2).to(tl.int32), tl.load(chunk_indices + i_t * 2 + 1).to(tl.int64) + bos, eos = tl.load(cu_seqlens + i_n).to(tl.int64), tl.load(cu_seqlens + i_n + 1).to(tl.int64) T = eos - bos else: bos, eos = i_b * T, i_b * T + T @@ -169,27 +177,33 @@ def chunk_rwkv6_fwd_A_kernel_intra_sub_inter( o_k = i_k * BK + tl.arange(0, BK) m_k = o_k < K - p_q = tl.make_block_ptr(q + (bos*H+i_h)*K, (T, K), (H*K, 1), (i_t * BT + i_i * BC, i_k * BK), (BC, BK), (1, 0)) - p_gq = tl.make_block_ptr(ge + (bos*H+i_h)*K, (T, K), (H*K, 1), (i_t * BT + i_i * BC, i_k * BK), (BC, BK), (1, 0)) - p_k = tl.make_block_ptr(k + (bos*H+i_h)*K, (K, T), (1, H*K), (i_k * BK, i_t * BT + i_j * BC), (BK, BC), (0, 1)) - p_gk = tl.make_block_ptr(gi + (bos*H+i_h)*K, (K, T), (1, H*K), (i_k * BK, i_t * BT + i_j * BC), (BK, BC), (0, 1)) + o_qi = i_t * BT + i_i * BC + tl.arange(0, BC) + o_kj = i_t * BT + i_j * BC + tl.arange(0, BC) + m_qk = m_i[:, None] & m_k[None, :] + m_kj = m_k[:, None] & (o_kj[None, :] < T) + p_q = q + (bos*H+i_h)*K + o_qi[:, None] * (H*K) + o_k[None, :] + p_gq = ge + (bos*H+i_h)*K + o_qi[:, None] * (H*K) + o_k[None, :] + p_k = k + (bos*H+i_h)*K + o_k[:, None] + o_kj[None, :] * (H*K) + p_gk = gi + (bos*H+i_h)*K + o_k[:, None] + o_kj[None, :] * (H*K) p_gn = gi + (bos + i_t * BT + i_i * BC - 1) * H*K + i_h * K + o_k # [BK,] b_gn = tl.load(p_gn, mask=m_k, other=0) # [BC, BK] - b_q = tl.load(p_q, boundary_check=(0, 1)) - b_gq = tl.where(m_i[:, None] & m_k, tl.load(p_gq, boundary_check=(0, 1)), float('-inf')) + b_q = tl.load(p_q, mask=m_qk, other=0.0) + b_gq = tl.where(m_i[:, None] & m_k, tl.load(p_gq, mask=m_qk, other=0.0), float('-inf')) b_qg = b_q * exp2(b_gq - b_gn[None, :]) * scale # [BK, BC] - b_k = tl.load(p_k, boundary_check=(0, 1)) - b_gk = tl.load(p_gk, boundary_check=(0, 1)) + b_k = tl.load(p_k, mask=m_kj, other=0.0) + b_gk = tl.load(p_gk, mask=m_kj, other=0.0) b_kg = b_k * exp2(b_gn[:, None] - b_gk) # [BC, BC] using tf32 to improve precision here. - b_A += tl.dot(b_qg, b_kg) + b_A = tl.dot(b_qg, b_kg, b_A) - p_A = tl.make_block_ptr(A + (bos*H + i_h)*BT, (T, BT), (H*BT, 1), (i_t * BT + i_i * BC, i_j * BC), (BC, BC), (1, 0)) - tl.store(p_A, b_A.to(A.dtype.element_ty), boundary_check=(0, 1)) + o_Ai = i_t * BT + i_i * BC + tl.arange(0, BC) + o_Aj = i_j * BC + tl.arange(0, BC) + p_A = A + (bos*H + i_h)*BT + o_Ai[:, None] * (H*BT) + o_Aj[None, :] + tl.store(p_A, b_A.to(A.dtype.element_ty), mask=m_i[:, None] & (o_Aj[None, :] < BT)) @triton.heuristics({ @@ -222,12 +236,12 @@ def chunk_rwkv6_fwd_A_kernel_intra_sub_intra( BK: tl.constexpr, IS_VARLEN: tl.constexpr, ): - i_t, i_i, i_bh = tl.program_id(0), tl.program_id(1), tl.program_id(2) + i_t, i_i, i_bh = tl.program_id(0).to(tl.int64), tl.program_id(1).to(tl.int64), tl.program_id(2).to(tl.int64) i_b, i_h = i_bh // H, i_bh % H i_j = i_i 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.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) + i_n, i_t = tl.load(chunk_indices + i_t * 2).to(tl.int32), tl.load(chunk_indices + i_t * 2 + 1).to(tl.int64) + bos, eos = tl.load(cu_seqlens + i_n).to(tl.int64), tl.load(cu_seqlens + i_n + 1).to(tl.int64) T = eos - bos else: bos, eos = i_b * T, i_b * T + T @@ -240,17 +254,19 @@ def chunk_rwkv6_fwd_A_kernel_intra_sub_intra( m_k = o_k < K m_A = (i_t * BT + i_i * BC + tl.arange(0, BC)) < T o_A = (bos + i_t * BT + i_i * BC + tl.arange(0, BC)) * H*BT + i_h * BT + i_j * BC - p_q = tl.make_block_ptr(q + (bos * H + i_h) * K, (T, K), (H*K, 1), (i_t * BT + i_i * BC, 0), (BC, BK), (1, 0)) - p_g = tl.make_block_ptr(ge + (bos * H + i_h) * K, (T, K), (H*K, 1), (i_t * BT + i_i * BC, 0), (BC, BK), (1, 0)) + o_q = i_t * BT + i_i * BC + tl.arange(0, BC) + m_q = m_A[:, None] & m_k[None, :] + p_q = q + (bos * H + i_h) * K + o_q[:, None] * (H*K) + o_k[None, :] + p_g = ge + (bos * H + i_h) * K + o_q[:, None] * (H*K) + o_k[None, :] p_qj = q + (bos + i_t * BT + i_j * BC) * H*K + i_h * K + o_k p_kj = k + (bos + i_t * BT + i_j * BC) * H*K + i_h * K + o_k p_gk = gi + (bos + i_t * BT + i_j * BC) * H*K + i_h * K + o_k - b_q = tl.load(p_q, boundary_check=(0, 1)) - b_g = tl.load(p_g, boundary_check=(0, 1)) + b_q = tl.load(p_q, mask=m_q, other=0.0) + b_g = tl.load(p_g, mask=m_q, other=0.0) - p_u = tl.make_block_ptr(u + i_h * K, (K,), (1,), (0,), (BK,), (0,)) - b_u = tl.load(p_u, boundary_check=(0,)) + p_u = u + i_h * K + o_k + b_u = tl.load(p_u, mask=m_k, other=0.0) for j in range(0, min(BC, T - i_t * BT - i_i * BC)): b_qj = tl.load(p_qj, mask=m_k, other=0).to(tl.float32) b_kj = tl.load(p_kj, mask=m_k, other=0).to(tl.float32) @@ -298,13 +314,14 @@ def chunk_rwkv6_fwd_A_kernel_intra_sub_intra_split( NC: tl.constexpr, IS_VARLEN: tl.constexpr, ): - i_k, i_tc, i_bh = tl.program_id(0), tl.program_id(1), tl.program_id(2) + i_k, i_tc = unflatten_program_id(tl.cdiv(K, BK)) + i_bh = tl.program_id(1).to(tl.int64) i_b, i_h = i_bh // H, i_bh % H i_t, i_i = i_tc // NC, i_tc % NC i_j = i_i 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.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) + i_n, i_t = tl.load(chunk_indices + i_t * 2).to(tl.int32), tl.load(chunk_indices + i_t * 2 + 1).to(tl.int64) + bos, eos = tl.load(cu_seqlens + i_n).to(tl.int64), tl.load(cu_seqlens + i_n + 1).to(tl.int64) all = T T = eos - bos else: @@ -320,17 +337,19 @@ def chunk_rwkv6_fwd_A_kernel_intra_sub_intra_split( m_A = (i_t * BT + i_i * BC + tl.arange(0, BC)) < T o_A = (i_k * all + bos + i_t * BT + i_i * BC + tl.arange(0, BC)) * H*BC + i_h * BC - p_q = tl.make_block_ptr(q + (bos * H + i_h) * K, (T, K), (H*K, 1), (i_t * BT + i_i * BC, i_k * BK), (BC, BK), (1, 0)) - p_g = tl.make_block_ptr(ge + (bos * H + i_h) * K, (T, K), (H*K, 1), (i_t * BT + i_i * BC, i_k * BK), (BC, BK), (1, 0)) + o_q = i_t * BT + i_i * BC + tl.arange(0, BC) + m_q = m_A[:, None] & m_k[None, :] + p_q = q + (bos * H + i_h) * K + o_q[:, None] * (H*K) + o_k[None, :] + p_g = ge + (bos * H + i_h) * K + o_q[:, None] * (H*K) + o_k[None, :] p_qj = q + (bos + i_t * BT + i_j * BC) * H*K + i_h * K + o_k p_kj = k + (bos + i_t * BT + i_j * BC) * H*K + i_h * K + o_k p_gk = gi + (bos + i_t * BT + i_j * BC) * H*K + i_h * K + o_k - b_q = tl.load(p_q, boundary_check=(0, 1)) - b_g = tl.load(p_g, boundary_check=(0, 1)) + b_q = tl.load(p_q, mask=m_q, other=0.0) + b_g = tl.load(p_g, mask=m_q, other=0.0) - p_u = tl.make_block_ptr(u + i_h * K, (K,), (1,), (i_k * BK), (BK,), (0,)) - b_u = tl.load(p_u, boundary_check=(0,)) + p_u = u + i_h * K + o_k + b_u = tl.load(p_u, mask=m_k, other=0.0) for j in range(0, min(BC, T - i_t * BT - i_i * BC)): b_qj = tl.load(p_qj, mask=m_k, other=0).to(tl.float32) b_kj = tl.load(p_kj, mask=m_k, other=0).to(tl.float32) @@ -371,11 +390,11 @@ def chunk_rwkv6_fwd_A_kernel_intra_sub_intra_merge( NK: tl.constexpr, IS_VARLEN: tl.constexpr, ): - i_t, i_c, i_bh = tl.program_id(0), tl.program_id(1), tl.program_id(2) + i_t, i_c, i_bh = tl.program_id(0).to(tl.int64), tl.program_id(1).to(tl.int64), tl.program_id(2).to(tl.int64) i_b, i_h = i_bh // H, i_bh % H 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.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) + i_n, i_t = tl.load(chunk_indices + i_t * 2).to(tl.int32), tl.load(chunk_indices + i_t * 2 + 1).to(tl.int64) + bos, eos = tl.load(cu_seqlens + i_n).to(tl.int64), tl.load(cu_seqlens + i_n + 1).to(tl.int64) all = T T = eos - bos else: @@ -385,12 +404,16 @@ def chunk_rwkv6_fwd_A_kernel_intra_sub_intra_merge( if i_t * BT + i_c * BC >= T: return + o_r = i_t * BT + i_c * BC + tl.arange(0, BC) + o_c = tl.arange(0, BC) + m_r = o_r < T b_A = tl.zeros([BC, BC], dtype=tl.float32) for i_k in range(0, NK): - p_A = tl.make_block_ptr(A + (i_k*all+bos)*H*BC+i_h*BC, (T, BC), (H*BC, 1), (i_t*BT + i_c*BC, 0), (BC, BC), (1, 0)) - b_A += tl.load(p_A, boundary_check=(0, 1)) - p_A2 = tl.make_block_ptr(A2 + (bos*H+i_h)*BT, (T, BT), (H*BT, 1), (i_t * BT + i_c * BC, i_c * BC), (BC, BC), (1, 0)) - tl.store(p_A2, b_A.to(A2.dtype.element_ty), boundary_check=(0, 1)) + p_A = A + (i_k*all+bos)*H*BC+i_h*BC + o_r[:, None] * (H*BC) + o_c[None, :] + b_A += tl.load(p_A, mask=m_r[:, None] & (o_c[None, :] < BC), other=0.0) + o_c2 = i_c * BC + tl.arange(0, BC) + p_A2 = A2 + (bos*H+i_h)*BT + o_r[:, None] * (H*BT) + o_c2[None, :] + tl.store(p_A2, b_A.to(A2.dtype.element_ty), mask=m_r[:, None] & (o_c2[None, :] < BT)) @triton.heuristics({ @@ -434,48 +457,55 @@ def chunk_rwkv6_bwd_kernel_dh( USE_FINAL_STATE_GRADIENT: tl.constexpr, IS_VARLEN: tl.constexpr, ): - i_k, i_v, i_nh = tl.program_id(0), tl.program_id(1), tl.program_id(2) + i_k, i_v, i_nh = unflatten_program_id(tl.cdiv(K, BK), tl.cdiv(V, BV)) i_n, i_hq = i_nh // HQ, i_nh % HQ i_h = i_hq // NG if IS_VARLEN: - bos, eos = tl.load(cu_seqlens + i_n).to(tl.int32), tl.load(cu_seqlens + i_n + 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 NT = tl.cdiv(T, BT) - boh = tl.load(chunk_offsets + i_n).to(tl.int32) + boh = tl.load(chunk_offsets + i_n).to(tl.int64) else: bos, eos = i_n * T, i_n * T + T NT = tl.cdiv(T, BT) boh = i_n * NT + o_k = i_k * BK + tl.arange(0, BK) + o_v = i_v * BV + tl.arange(0, BV) + m_kv = (o_k[:, None] < K) & (o_v[None, :] < V) # [BK, BV] b_dh = tl.zeros([BK, BV], dtype=tl.float32) if USE_FINAL_STATE_GRADIENT: - p_dht = tl.make_block_ptr(dht + i_nh * K*V, (K, V), (V, 1), (i_k * BK, i_v * BV), (BK, BV), (1, 0)) - b_dh += tl.load(p_dht, boundary_check=(0, 1)).to(tl.float32) + p_dht = dht + i_nh * K*V + o_k[:, None] * V + o_v[None, :] + b_dh += tl.load(p_dht, mask=m_kv, other=0.0).to(tl.float32) for i_t in range(NT - 1, -1, -1): - p_dh = tl.make_block_ptr(dh + ((boh+i_t) * H + i_h) * K*V, (K, V), (V, 1), (i_k * BK, i_v * BV), (BK, BV), (1, 0)) - tl.store(p_dh, b_dh.to(p_dh.dtype.element_ty), boundary_check=(0, 1)) + o_t = i_t.to(tl.int64) * BT + tl.arange(0, BT) + m_t = o_t < T + m_kt = (o_k[:, None] < K) & m_t[None, :] + m_tv = m_t[:, None] & (o_v[None, :] < V) + p_dh = dh + ((boh+i_t) * H + i_h) * K*V + o_k[:, None] * V + o_v[None, :] + tl.store(p_dh, b_dh.to(p_dh.dtype.element_ty), mask=m_kv) last_idx = min(i_t * BT + BT, T) - 1 # [BK, BT] - p_q = tl.make_block_ptr(q + (bos*HQ + i_hq) * K, (K, T), (1, HQ*K), (i_k * BK, i_t * BT), (BK, BT), (0, 1)) - p_do = tl.make_block_ptr(do + (bos*HQ + i_hq) * V, (T, V), (HQ*V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) - b_q = tl.load(p_q, boundary_check=(0, 1)) + p_q = q + (bos*HQ + i_hq) * K + o_k[:, None] + o_t[None, :] * (HQ*K) + p_do = do + (bos*HQ + i_hq) * V + o_t[:, None] * (HQ*V) + o_v[None, :] + b_q = tl.load(p_q, mask=m_kt, other=0.0) # [BT, BV] - b_do = tl.load(p_do, boundary_check=(0, 1)) + b_do = tl.load(p_do, mask=m_tv, other=0.0) - p_gk = tl.make_block_ptr(ge + (bos*H + i_h) * K, (K, T), (1, H*K), (i_k * BK, i_t * BT), (BK, BT), (0, 1)) + p_gk = ge + (bos*H + i_h) * K + o_k[:, None] + o_t[None, :] * (H*K) p_gk_last = gi + (bos + last_idx) * H*K + i_h * K + i_k * BK + tl.arange(0, BK) - b_gk = tl.load(p_gk, boundary_check=(0, 1)) + b_gk = tl.load(p_gk, mask=m_kt, other=0.0) b_q = (b_q * exp2(b_gk) * scale).to(b_q.dtype) b_gk_last = tl.load(p_gk_last, mask=(i_k * BK + tl.arange(0, BK) < K), other=0.) b_dh *= exp2(b_gk_last)[:, None] - b_dh += tl.dot(b_q, b_do) + b_dh = tl.dot(b_q, b_do, b_dh) if STORE_INITIAL_STATE_GRADIENT: - p_dh0 = tl.make_block_ptr(dh0 + i_nh * K*V, (K, V), (V, 1), (i_k * BK, i_v * BV), (BK, BV), (1, 0)) - tl.store(p_dh0, b_dh.to(p_dh0.dtype.element_ty), boundary_check=(0, 1)) + p_dh0 = dh0 + i_nh * K*V + o_k[:, None] * V + o_v[None, :] + tl.store(p_dh0, b_dh.to(p_dh0.dtype.element_ty), mask=m_kv) @triton.heuristics({ @@ -509,12 +539,13 @@ def chunk_rwkv6_bwd_kernel_intra( NC: tl.constexpr, IS_VARLEN: tl.constexpr, ): - i_k, i_c, i_bh = tl.program_id(0), tl.program_id(1), tl.program_id(2) + i_k, i_c = unflatten_program_id(tl.cdiv(K, BK)) + i_bh = tl.program_id(1).to(tl.int64) i_b, i_h = i_bh // H, i_bh % H i_t, i_i = i_c // NC, i_c % NC 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.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) + i_n, i_t = tl.load(chunk_indices + i_t * 2).to(tl.int32), tl.load(chunk_indices + i_t * 2 + 1).to(tl.int64) + bos, eos = tl.load(cu_seqlens + i_n).to(tl.int64), tl.load(cu_seqlens + i_n + 1).to(tl.int64) else: bos, eos = i_b * T, i_b * T + T T = eos - bos @@ -523,27 +554,32 @@ def chunk_rwkv6_bwd_kernel_intra( o_k = i_k * BK + tl.arange(0, BK) m_k = o_k < K + o_r = i_t * BT + i_i * BC + tl.arange(0, BC) + m_rk = (o_r[:, None] < T) & m_k[None, :] - p_ge = tl.make_block_ptr(ge + (bos*H + i_h) * K, (T, K), (H*K, 1), (i_t * BT + i_i * BC, i_k * BK), (BC, BK), (1, 0)) + p_ge = ge + (bos*H + i_h) * K + o_r[:, None] * (H*K) + o_k[None, :] # [BC, BK] - b_ge = tl.load(p_ge, boundary_check=(0, 1)) + b_ge = tl.load(p_ge, mask=m_rk, other=0.0) b_dq = tl.zeros([BC, BK], dtype=tl.float32) if i_i > 0: p_gn = gi + (bos + i_t * BT + i_i * BC - 1) * H*K + i_h*K + o_k # [BK,] b_gn = tl.load(p_gn, mask=m_k, other=0) for i_j in range(0, i_i): - p_k = tl.make_block_ptr(k+(bos*H+i_h)*K, (T, K), (H*K, 1), (i_t*BT+i_j*BC, i_k * BK), (BC, BK), (1, 0)) - p_gk = tl.make_block_ptr(gi+(bos*H+i_h)*K, (T, K), (H*K, 1), (i_t*BT+i_j*BC, i_k * BK), (BC, BK), (1, 0)) - p_dA = tl.make_block_ptr(dA+(bos*H+i_h)*BT, (T, BT), (H*BT, 1), (i_t*BT+i_i*BC, i_j * BC), (BC, BC), (1, 0)) + o_j = i_t * BT + i_j * BC + tl.arange(0, BC) + m_jk = (o_j[:, None] < T) & m_k[None, :] + o_dAj = i_j * BC + tl.arange(0, BC) + p_k = k+(bos*H+i_h)*K + o_j[:, None] * (H*K) + o_k[None, :] + p_gk = gi+(bos*H+i_h)*K + o_j[:, None] * (H*K) + o_k[None, :] + p_dA = dA+(bos*H+i_h)*BT + o_r[:, None] * (H*BT) + o_dAj[None, :] # [BC, BK] - b_k = tl.load(p_k, boundary_check=(0, 1)) - b_gk = tl.load(p_gk, boundary_check=(0, 1)) + b_k = tl.load(p_k, mask=m_jk, other=0.0) + b_gk = tl.load(p_gk, mask=m_jk, other=0.0) b_kg = b_k * exp2(b_gn[None, :] - b_gk) # [BC, BC] - b_dA = tl.load(p_dA, boundary_check=(0, 1)) + b_dA = tl.load(p_dA, mask=(o_r[:, None] < T) & (o_dAj[None, :] < BT), other=0.0) # [BC, BK] - b_dq += tl.dot(b_dA, b_kg) + b_dq = tl.dot(b_dA, b_kg, b_dq) b_dq *= exp2(b_ge - b_gn[None, :]) o_i = tl.arange(0, BC) @@ -551,7 +587,7 @@ def chunk_rwkv6_bwd_kernel_intra( o_dA = bos*H*BT + (i_t * BT + i_i * BC + tl.arange(0, BC)) * H*BT + i_h * BT + i_i * BC p_kj = k + (bos + i_t * BT + i_i * BC) * H*K + i_h * K + o_k p_gkj = gi + (bos + i_t * BT + i_i * BC) * H*K + i_h * K + o_k - p_dq = tl.make_block_ptr(dq + (bos*H + i_h) * K, (T, K), (H*K, 1), (i_t * BT + i_i * BC, i_k * BK), (BC, BK), (1, 0)) + p_dq = dq + (bos*H + i_h) * K + o_r[:, None] * (H*K) + o_k[None, :] for j in range(0, min(BC, T - i_t * BT - i_i * BC)): # [BC,] @@ -566,15 +602,13 @@ def chunk_rwkv6_bwd_kernel_intra( b_dq += tl.where(m_i, b_dA[:, None] * b_kj[None, :] * exp2(b_ge - b_gkj[None, :]), 0.) p_kj += H*K p_gkj += H*K - tl.store(p_dq, b_dq.to(p_dq.dtype.element_ty), boundary_check=(0, 1)) + tl.store(p_dq, b_dq.to(p_dq.dtype.element_ty), mask=m_rk) tl.debug_barrier() - p_k = tl.make_block_ptr(k + (bos*H + i_h) * K, (T, K), (H*K, 1), (i_t * BT + i_i * BC, i_k * BK), (BC, BK), (1, 0)) - p_gk = tl.make_block_ptr(gi + (bos*H + i_h) * K, (T, K), (H*K, 1), (i_t * BT + i_i * BC, i_k * BK), (BC, BK), (1, 0)) + p_gk = gi + (bos*H + i_h) * K + o_r[:, None] * (H*K) + o_k[None, :] # [BC, BK] - b_k = tl.load(p_k, boundary_check=(0, 1)) - b_gk = tl.load(p_gk, boundary_check=(0, 1)) + b_gk = tl.load(p_gk, mask=m_rk, other=0.0) b_dk = tl.zeros([BC, BK], dtype=tl.float32) NC = min(NC, tl.cdiv(T - i_t * BT, BC)) @@ -584,24 +618,27 @@ def chunk_rwkv6_bwd_kernel_intra( # [BK,] b_gn = tl.load(p_gn, mask=m_k, other=0) for i_j in range(i_i + 1, NC): - m_j = (i_t * BT + i_j * BC + tl.arange(0, BC)) < T - p_q = tl.make_block_ptr(q + (bos*H+i_h)*K, (T, K), (H*K, 1), (i_t * BT + i_j * BC, i_k*BK), (BC, BK), (1, 0)) - p_gq = tl.make_block_ptr(ge + (bos*H+i_h)*K, (T, K), (H*K, 1), (i_t * BT + i_j * BC, i_k*BK), (BC, BK), (1, 0)) - p_dA = tl.make_block_ptr(dA + (bos*H+i_h)*BT, (BT, T), (1, H*BT), (i_i*BC, i_t*BT + i_j*BC), (BC, BC), (0, 1)) + o_j2 = i_t * BT + i_j * BC + tl.arange(0, BC) + m_j = o_j2 < T + m_jk = m_j[:, None] & m_k[None, :] + o_dAi = i_i * BC + tl.arange(0, BC) + p_q = q + (bos*H+i_h)*K + o_j2[:, None] * (H*K) + o_k[None, :] + p_gq = ge + (bos*H+i_h)*K + o_j2[:, None] * (H*K) + o_k[None, :] + p_dA = dA + (bos*H+i_h)*BT + o_dAi[:, None] + o_j2[None, :] * (H*BT) # [BC, BK] - b_q = tl.load(p_q, boundary_check=(0, 1)) - b_gq = tl.where(m_j[:, None] & m_k, tl.load(p_gq, boundary_check=(0, 1)), float('-inf')) + b_q = tl.load(p_q, mask=m_jk, other=0.0) + b_gq = tl.where(m_j[:, None] & m_k, tl.load(p_gq, mask=m_jk, other=0.0), float('-inf')) b_qg = b_q * exp2(b_gq - b_gn[None, :]) # [BC, BC] - b_dA = tl.load(p_dA, boundary_check=(0, 1)) + b_dA = tl.load(p_dA, mask=(o_dAi[:, None] < BT) & m_j[None, :], other=0.0) # [BC, BK] # (SY 09/17) important to not use bf16 here to have a good precision. - b_dk += tl.dot(b_dA, b_qg) + b_dk = tl.dot(b_dA, b_qg, b_dk) b_dk *= exp2(b_gn[None, :] - b_gk) o_dA = bos*H*BT + (i_t * BT + i_i * BC) * H*BT + i_h * BT + i_i * BC + tl.arange(0, BC) p_qj = q + (bos + i_t * BT + i_i * BC) * H*K + i_h * K + o_k p_gqj = ge + (bos + i_t * BT + i_i * BC) * H*K + i_h * K + o_k - p_dk = tl.make_block_ptr(dk + (bos*H+i_h)*K, (T, K), (H*K, 1), (i_t * BT + i_i * BC, i_k * BK), (BC, BK), (1, 0)) + p_dk = dk + (bos*H+i_h)*K + o_r[:, None] * (H*K) + o_k[None, :] for j in range(0, min(BC, T - i_t * BT - i_i * BC)): # [BC,] b_dA = tl.load(dA + o_dA + j * H*BT) @@ -613,7 +650,7 @@ def chunk_rwkv6_bwd_kernel_intra( b_dk += tl.where(m_i, b_dA[:, None] * b_qj[None, :] * exp2(b_gqj[None, :] - b_gk), 0.) p_qj += H*K p_gqj += H*K - tl.store(p_dk, b_dk.to(p_dk.dtype.element_ty), boundary_check=(0, 1)) + tl.store(p_dk, b_dk.to(p_dk.dtype.element_ty), mask=m_rk) @triton.heuristics({ @@ -659,13 +696,14 @@ def chunk_rwkv6_bwd_kernel_inter( BV: tl.constexpr, IS_VARLEN: tl.constexpr, ): - i_k, i_t, i_bh = tl.program_id(0), tl.program_id(1), tl.program_id(2) + i_k, i_t = unflatten_program_id(tl.cdiv(K, BK)) + i_bh = tl.program_id(1).to(tl.int64) i_b, i_h = i_bh // H, i_bh % H 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.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) + i_n, i_t = tl.load(chunk_indices + i_t * 2).to(tl.int32), tl.load(chunk_indices + i_t * 2 + 1).to(tl.int64) + bos, eos = tl.load(cu_seqlens + i_n).to(tl.int64), tl.load(cu_seqlens + i_n + 1).to(tl.int64) T = eos - bos NT = tl.cdiv(T, BT) else: @@ -675,8 +713,10 @@ def chunk_rwkv6_bwd_kernel_inter( o_k = i_k * BK + tl.arange(0, BK) m_k = o_k < K - p_gk = tl.make_block_ptr(ge + (bos*H+i_h)*K, (T, K), (H*K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) - p_gi = tl.make_block_ptr(gi + (bos*H+i_h)*K, (T, K), (H*K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) + o_t = i_t * BT + tl.arange(0, BT) + m_tk = (o_t[:, None] < T) & m_k[None, :] + p_gk = ge + (bos*H+i_h)*K + o_t[:, None] * (H*K) + o_k[None, :] + p_gi = gi + (bos*H+i_h)*K + o_t[:, None] * (H*K) + o_k[None, :] p_gn = gi + (bos + min(T, i_t * BT + BT)-1) * H*K + i_h * K + o_k b_gn = tl.load(p_gn, mask=m_k, other=0) b_dq = tl.zeros([BT, BK], dtype=tl.float32) @@ -684,62 +724,66 @@ def chunk_rwkv6_bwd_kernel_inter( b_dgk = tl.zeros([BK], dtype=tl.float32) for i_v in range(tl.cdiv(V, BV)): - p_v = tl.make_block_ptr(v + (bos*H + i_h) * V, (T, V), (H*V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) - p_do = tl.make_block_ptr(do + (bos*H + i_h) * V, (T, V), (H*V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) - p_h = tl.make_block_ptr(h + (i_tg * H + i_h) * K*V, (V, K), (1, V), (i_v * BV, i_k * BK), (BV, BK), (0, 1)) - p_dh = tl.make_block_ptr(dh + (i_tg * H + i_h) * K*V, (V, K), (1, V), (i_v * BV, i_k * BK), (BV, BK), (0, 1)) + o_v = i_v * BV + tl.arange(0, BV) + m_tv = (o_t[:, None] < T) & (o_v[None, :] < V) + m_h = (o_v[:, None] < V) & m_k[None, :] + p_v = v + (bos*H + i_h) * V + o_t[:, None] * (H*V) + o_v[None, :] + p_do = do + (bos*H + i_h) * V + o_t[:, None] * (H*V) + o_v[None, :] + p_h = h + (i_tg * H + i_h) * K*V + o_v[:, None] + o_k[None, :] * V + p_dh = dh + (i_tg * H + i_h) * K*V + o_v[:, None] + o_k[None, :] * V # [BT, BV] - b_v = tl.load(p_v, boundary_check=(0, 1)) - b_do = tl.load(p_do, boundary_check=(0, 1)) + b_v = tl.load(p_v, mask=m_tv, other=0.0) + b_do = tl.load(p_do, mask=m_tv, other=0.0) # [BV, BK] - b_h = tl.load(p_h, boundary_check=(0, 1)) - b_dh = tl.load(p_dh, boundary_check=(0, 1)) + b_h = tl.load(p_h, mask=m_h, other=0.0) + b_dh = tl.load(p_dh, mask=m_h, other=0.0) # [BK] b_dgk += tl.sum(b_h * b_dh, axis=0) # [BT, BK] - b_dq += tl.dot(b_do, b_h.to(b_do.dtype)) - b_dk += tl.dot(b_v, b_dh.to(b_v.dtype)) + b_dq = tl.dot(b_do, b_h.to(b_do.dtype), b_dq) + b_dk = tl.dot(b_v, b_dh.to(b_v.dtype), b_dk) b_dgk *= exp2(b_gn) b_dq *= scale - b_gk = tl.load(p_gk, boundary_check=(0, 1)) - b_gi = tl.load(p_gi, boundary_check=(0, 1)) + b_gk = tl.load(p_gk, mask=m_tk, other=0.0) + b_gi = tl.load(p_gi, mask=m_tk, other=0.0) b_dq = b_dq * exp2(b_gk) b_dk = b_dk * exp2(b_gn[None, :] - b_gi) o_i = tl.arange(0, BT) - p_q = tl.make_block_ptr(q + (bos*H+i_h)*K, (T, K), (H*K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) - p_k = tl.make_block_ptr(k + (bos*H+i_h)*K, (T, K), (H*K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) - p_dq = tl.make_block_ptr(dq + (bos*H+i_h)*K, (T, K), (H*K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) - p_dk = tl.make_block_ptr(dk + (bos*H+i_h)*K, (T, K), (H*K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) + p_q = q + (bos*H+i_h)*K + o_t[:, None] * (H*K) + o_k[None, :] + p_k = k + (bos*H+i_h)*K + o_t[:, None] * (H*K) + o_k[None, :] + p_dq = dq + (bos*H+i_h)*K + o_t[:, None] * (H*K) + o_k[None, :] + p_dk = dk + (bos*H+i_h)*K + o_t[:, None] * (H*K) + o_k[None, :] p_dA_dig = dA + ((bos + i_t * BT + o_i) * H + i_h) * BT + o_i - b_q = tl.load(p_q, boundary_check=(0, 1)) - b_k = tl.load(p_k, boundary_check=(0, 1)) + b_q = tl.load(p_q, mask=m_tk, other=0.0) + b_k = tl.load(p_k, mask=m_tk, other=0.0) b_dgk += tl.sum(b_dk * b_k, axis=0) - b_dq += tl.load(p_dq, boundary_check=(0, 1)) - b_dk += tl.load(p_dk, boundary_check=(0, 1)) + b_dq += tl.load(p_dq, mask=m_tk, other=0.0) + b_dk += tl.load(p_dk, mask=m_tk, other=0.0) b_dg = b_q * b_dq - b_k * b_dk b_dg = b_dg - tl.cumsum(b_dg, axis=0) + tl.sum(b_dg, axis=0)[None, :] + b_dgk[None, :] - b_q * b_dq # [BT,] b_dA_dig = tl.load(p_dA_dig, mask=(i_t * BT + o_i) < T, other=0) - p_u = tl.make_block_ptr(u + i_h * K, (K,), (1,), (i_k * BK,), (BK,), (0,)) - b_u = tl.load(p_u, boundary_check=(0,)) + p_u = u + i_h * K + o_k + b_u = tl.load(p_u, mask=m_k, other=0.0) # scale is already applied to b_dA_diag b_dq += (b_dA_dig[:, None] * b_u[None, :] * b_k) b_dk += (b_dA_dig[:, None] * b_u[None, :] * b_q) b_du = tl.sum(b_dA_dig[:, None] * b_q * b_k, axis=0) - p_du = tl.make_block_ptr(du + (i_tg * H + i_h) * K, (K,), (1,), (i_k * BK,), (BK,), (0,)) - tl.store(p_du, b_du, boundary_check=(0,)) + p_du = du + (i_tg * H + i_h) * K + o_k + tl.store(p_du, b_du, mask=m_k) - p_dq = tl.make_block_ptr(dq2 + (bos * H + i_h) * K, (T, K), (H*K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) - p_dk = tl.make_block_ptr(dk2 + (bos * H + i_h) * K, (T, K), (H*K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) - p_dg = tl.make_block_ptr(dg + (bos * H + i_h) * K, (T, K), (H*K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) - tl.store(p_dq, b_dq.to(p_dq.dtype.element_ty), boundary_check=(0, 1)) - tl.store(p_dk, b_dk.to(p_dk.dtype.element_ty), boundary_check=(0, 1)) - tl.store(p_dg, b_dg.to(p_dg.dtype.element_ty), boundary_check=(0, 1)) + p_dq = dq2 + (bos * H + i_h) * K + o_t[:, None] * (H*K) + o_k[None, :] + p_dk = dk2 + (bos * H + i_h) * K + o_t[:, None] * (H*K) + o_k[None, :] + p_dg = dg + (bos * H + i_h) * K + o_t[:, None] * (H*K) + o_k[None, :] + tl.store(p_dq, b_dq.to(p_dq.dtype.element_ty), mask=m_tk) + tl.store(p_dk, b_dk.to(p_dk.dtype.element_ty), mask=m_tk) + tl.store(p_dg, b_dg.to(p_dg.dtype.element_ty), mask=m_tk) +@dispatch('rwkv6') def chunk_rwkv6_fwd_intra( q: torch.Tensor, k: torch.Tensor, @@ -806,7 +850,7 @@ def chunk_rwkv6_fwd_intra( NK = triton.cdiv(K, BK) A_intra = q.new_empty(NK, B, T, H, BC, dtype=torch.float) - grid = (NK, NT * NC, B * H) + grid = (NK * NT * NC, B * H) chunk_rwkv6_fwd_A_kernel_intra_sub_intra_split[grid]( q, k, @@ -873,7 +917,7 @@ def chunk_rwkv6_bwd_dh( dh = k.new_empty(B, NT, HQ, K, V, dtype=k.dtype if not states_in_fp32 else torch.float) dh0 = torch.empty_like(h0, dtype=torch.float) if h0 is not None else None - def grid(meta): return (triton.cdiv(K, meta['BK']), triton.cdiv(V, meta['BV']), N * H) + def grid(meta): return (triton.cdiv(K, meta['BK']) * triton.cdiv(V, meta['BV']) * N * H,) chunk_rwkv6_bwd_kernel_dh[grid]( q=q, gi=gi, @@ -919,7 +963,7 @@ def chunk_rwkv6_bwd_dqk_intra( dq = torch.empty_like(q, dtype=torch.float) dk = torch.empty_like(k, dtype=torch.float) - grid = (NK, NT * NC, B * H) + grid = (NK * NT * NC, B * H) chunk_rwkv6_bwd_kernel_intra[grid]( q, k, @@ -971,7 +1015,7 @@ def chunk_rwkv6_bwd_dqkgu( dk2 = torch.empty_like(dk) dg = torch.empty_like(g) du = u.new_empty(B * NT, H, K, dtype=torch.float) - def grid(meta): return (triton.cdiv(K, meta['BK']), NT, B * H) + def grid(meta): return (triton.cdiv(K, meta['BK']) * NT, B * H) chunk_rwkv6_bwd_kernel_inter[grid]( q, k, @@ -1281,7 +1325,7 @@ def chunk_rwkv6( >>> import torch >>> import torch.nn.functional as F >>> from einops import rearrange - >>> from ...ops.rwkv6 import chunk_rwkv6 + >>> from fla.ops.rwkv6 import chunk_rwkv6 # inputs with equal lengths >>> B, T, H, K, V = 4, 2048, 4, 512, 512 >>> r = torch.randn(B, T, H, K, device='cuda') diff --git a/build/torch-cuda/ops/rwkv6/fused_recurrent.py b/build/torch-cuda/ops/rwkv6/fused_recurrent.py index e9437cb8c090c0ce0808d94d8821d1afa2e989f7..0990da48640c8f1dc1e62c171cc2aca7f767c639 100644 --- a/build/torch-cuda/ops/rwkv6/fused_recurrent.py +++ b/build/torch-cuda/ops/rwkv6/fused_recurrent.py @@ -50,7 +50,9 @@ def fused_recurrent_rwkv6_fwd_kernel( STORE_FINAL_STATE: tl.constexpr, # whether to store final state IS_VARLEN: tl.constexpr, ): - i_v, i_k, i_nh = tl.program_id(0).to(tl.int64), tl.program_id(1).to(tl.int64), tl.program_id(2).to(tl.int64) + pid = tl.program_id(0).to(tl.int64) + NK, NV = tl.cdiv(K, BK), tl.cdiv(V, BV) + i_v, i_k, i_nh = (pid % NV).to(tl.int64), ((pid // NV) % NK).to(tl.int64), (pid // (NV * NK)).to(tl.int64) i_n, i_h = i_nh // H, i_nh % H if IS_VARLEN: bos, eos = tl.load(cu_seqlens + i_n).to(tl.int64), tl.load(cu_seqlens + i_n + 1).to(tl.int64) @@ -136,7 +138,9 @@ def fused_recurrent_rwkv6_bwd_kernel_dq( USE_INITIAL_STATE: tl.constexpr, IS_VARLEN: tl.constexpr, ): - i_v, i_k, i_nh = tl.program_id(0).to(tl.int64), tl.program_id(1).to(tl.int64), tl.program_id(2).to(tl.int64) + pid = tl.program_id(0).to(tl.int64) + NK, NV = tl.cdiv(K, BK), tl.cdiv(V, BV) + i_v, i_k, i_nh = (pid % NV).to(tl.int64), ((pid // NV) % NK).to(tl.int64), (pid // (NV * NK)).to(tl.int64) i_n, i_h = i_nh // H, i_nh % H if IS_VARLEN: bos, eos = tl.load(cu_seqlens + i_n).to(tl.int64), tl.load(cu_seqlens + i_n + 1).to(tl.int64) @@ -228,7 +232,9 @@ def fused_recurrent_rwkv6_bwd_kernel_dkv( USE_INITIAL_STATE: tl.constexpr, IS_VARLEN: tl.constexpr, ): - i_v, i_k, i_nh = tl.program_id(0).to(tl.int64), tl.program_id(1).to(tl.int64), tl.program_id(2).to(tl.int64) + pid = tl.program_id(0).to(tl.int64) + NK, NV = tl.cdiv(K, BK), tl.cdiv(V, BV) + i_v, i_k, i_nh = (pid % NV).to(tl.int64), ((pid // NV) % NK).to(tl.int64), (pid // (NV * NK)).to(tl.int64) i_n, i_h = i_nh // H, i_nh % H if IS_VARLEN: bos, eos = tl.load(cu_seqlens + i_n).to(tl.int64), tl.load(cu_seqlens + i_n + 1).to(tl.int64) @@ -318,10 +324,12 @@ def fused_recurrent_rwkv6_bwd_kernel_dw( REVERSE: tl.constexpr, IS_VARLEN: tl.constexpr, ): - i_k, i_nh = tl.program_id(0), tl.program_id(1) + pid = tl.program_id(0).to(tl.int64) + NK = tl.cdiv(K, BK) + i_k, i_nh = pid % NK, (pid // NK).to(tl.int64) i_n, i_h = i_nh // H, i_nh % H if IS_VARLEN: - bos, eos = tl.load(cu_seqlens + i_n).to(tl.int32), tl.load(cu_seqlens + i_n + 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) else: bos, eos = i_n * T, i_n * T + T T = eos - bos @@ -329,24 +337,31 @@ def fused_recurrent_rwkv6_bwd_kernel_dw( o_i = tl.arange(0, BT) m_i = tl.where(o_i[:, None] >= o_i[None, :], 1., 0.) if not REVERSE else tl.where(o_i[:, None] <= o_i[None, :], 1., 0.) + o_k = i_k * BK + tl.arange(0, BK) + m_k = o_k < K b_z = tl.zeros([BK], dtype=tl.float32) i_t = 0 if not REVERSE else NT - 1 for _ in range(NT): - p_q = tl.make_block_ptr(q + (bos*H + i_h) * K, (T, K), (H*K, 1), (i_t * BT + 1, i_k * BK), (BT, BK), (1, 0)) - p_k = tl.make_block_ptr(k + (bos*H + i_h) * K, (T-1, K), (H*K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) - p_dq = tl.make_block_ptr(dq + (bos*H + i_h) * K, (T, K), (H*K, 1), (i_t * BT + 1, i_k * BK), (BT, BK), (1, 0)) - p_dk = tl.make_block_ptr(dk + (bos*H + i_h) * K, (T-1, K), (H*K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) - p_dw = tl.make_block_ptr(dw + (bos*H + i_h) * K, (T, K), (H*K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) + o_q = i_t.to(tl.int64) * BT + 1 + o_i + o_kt = i_t.to(tl.int64) * BT + o_i + m_q = (o_q[:, None] < T) & m_k[None, :] + m_kt = (o_kt[:, None] < T - 1) & m_k[None, :] + m_wt = (o_kt[:, None] < T) & m_k[None, :] + p_q = q + (bos*H + i_h) * K + o_q[:, None] * (H*K) + o_k[None, :] + p_k = k + (bos*H + i_h) * K + o_kt[:, None] * (H*K) + o_k[None, :] + p_dq = dq + (bos*H + i_h) * K + o_q[:, None] * (H*K) + o_k[None, :] + p_dk = dk + (bos*H + i_h) * K + o_kt[:, None] * (H*K) + o_k[None, :] + p_dw = dw + (bos*H + i_h) * K + o_kt[:, None] * (H*K) + o_k[None, :] # [BT, BK] - b_q = tl.load(p_q, boundary_check=(0, 1)).to(tl.float32) - b_dq = tl.load(p_dq, boundary_check=(0, 1)).to(tl.float32) - b_k = tl.load(p_k, boundary_check=(0, 1)).to(tl.float32) - b_dk = tl.load(p_dk, boundary_check=(0, 1)).to(tl.float32) + b_q = tl.load(p_q, mask=m_q, other=0.0).to(tl.float32) + b_dq = tl.load(p_dq, mask=m_q, other=0.0).to(tl.float32) + b_k = tl.load(p_k, mask=m_kt, other=0.0).to(tl.float32) + b_dk = tl.load(p_dk, mask=m_kt, other=0.0).to(tl.float32) b_dw = (b_q * b_dq * scale) - b_k * b_dk b_c = b_z[None, :] + tl.dot(m_i, b_dw, allow_tf32=False) - tl.store(p_dw, b_c.to(p_dw.dtype.element_ty), boundary_check=(0, 1)) + tl.store(p_dw, b_c.to(p_dw.dtype.element_ty), mask=m_wt) if i_t >= 0: b_z += tl.sum(b_dw, 0) @@ -374,7 +389,7 @@ def fused_recurrent_rwkv6_fwd( ht = q.new_empty(N, H, K, V, dtype=torch.float) if output_final_state else None o = q.new_empty(NK, *v.shape, dtype=torch.float) - grid = (NV, NK, N * H) + grid = (NV * NK * N * H,) fused_recurrent_rwkv6_fwd_kernel[grid]( q, k, @@ -420,7 +435,7 @@ def fused_recurrent_rwkv6_bwd( dq = q.new_empty(NV, *q.shape, dtype=torch.float) dq1 = torch.empty_like(dq) - grid = (NV, NK, N * H) + grid = (NV * NK * N * H,) fused_recurrent_rwkv6_bwd_kernel_dq[grid]( k, v, @@ -452,7 +467,7 @@ def fused_recurrent_rwkv6_bwd( dv = q.new_empty(NK, *v.shape, dtype=torch.float) dh0 = torch.empty_like(initial_state) if initial_state is not None else None - grid = (NV, NK, N * H) + grid = (NV * NK * N * H,) fused_recurrent_rwkv6_bwd_kernel_dkv[grid]( q, k, @@ -480,7 +495,7 @@ def fused_recurrent_rwkv6_bwd( dv = dv.sum(0) dw = torch.empty_like(w) - def grid(meta): return (triton.cdiv(meta['K'], meta['BK']), N * H) + def grid(meta): return (triton.cdiv(meta['K'], meta['BK']) * N * H,) fused_recurrent_rwkv6_bwd_kernel_dw[grid]( q, k, @@ -568,7 +583,6 @@ def fused_recurrent_rwkv6( output_final_state: bool = False, reverse: bool = False, cu_seqlens: torch.LongTensor | None = None, - **kwargs, ) -> tuple[torch.Tensor, torch.Tensor]: r""" Args: @@ -608,7 +622,7 @@ def fused_recurrent_rwkv6( >>> import torch >>> import torch.nn.functional as F >>> from einops import rearrange - >>> from ...ops.rwkv6 import fused_recurrent_rwkv6 + >>> from fla.ops.rwkv6 import fused_recurrent_rwkv6 # inputs with equal lengths >>> B, T, H, K, V = 4, 2048, 4, 512, 512 >>> q = torch.randn(B, T, H, K, device='cuda') @@ -635,10 +649,6 @@ def fused_recurrent_rwkv6( >>> assert o.allclose(o_var.view(o.shape)) >>> assert ht.allclose(ht_var) """ - if 'head_first' in kwargs: - raise DeprecationWarning( - "head_first has been removed. Inputs must be in `[B, T, H, ...]` format.", - ) if cu_seqlens is not None: if r.shape[0] != 1: raise ValueError( diff --git a/build/torch-cuda/ops/rwkv7/channel_mixing.py b/build/torch-cuda/ops/rwkv7/channel_mixing.py index 3facea3459a8f797ef2ec40fba9a7085fc874534..87d876411ccad2d80f82e305a3eab647850d10c3 100644 --- a/build/torch-cuda/ops/rwkv7/channel_mixing.py +++ b/build/torch-cuda/ops/rwkv7/channel_mixing.py @@ -44,7 +44,7 @@ def rwkv_seq_mix_kernel( hidden_dim: tl.constexpr, BLOCK_SIZE: tl.constexpr, ): - block_start = tl.program_id(0) * BLOCK_SIZE + block_start = tl.program_id(0).to(tl.int64) * BLOCK_SIZE block_idx = block_start + tl.arange(0, BLOCK_SIZE)[:] total_seq_dim = token_length * hidden_dim @@ -82,16 +82,18 @@ def rwkv_seq_mix_kernel( def rwkv_channel_mixing_pow_and_relu( in_ptr, out_ptr, + n_elements, BLOCK_SIZE: tl.constexpr, ): """Fused ReLU and Power operation: x = ReLU(x)^2""" - xoffset = tl.program_id(0) * BLOCK_SIZE + xoffset = tl.program_id(0).to(tl.int64) * BLOCK_SIZE xindex = xoffset + tl.arange(0, BLOCK_SIZE) x0 = xindex - x = tl.load(in_ptr + (x0), None) + xmask = xindex < n_elements + x = tl.load(in_ptr + (x0), mask=xmask, other=0.0) x = tl.maximum(x, 0.0).to(tl.float32) x = tl.cast(x * x, dtype=out_ptr.dtype.element_ty, fp_downcast_rounding='rtne') - tl.store(out_ptr + (x0), x, None) + tl.store(out_ptr + (x0), x, mask=xmask) def rwkv_mix_torch(x: torch.Tensor, x_prev: torch.Tensor, x_k: torch.Tensor): @@ -162,6 +164,7 @@ def rwkv_relu_and_square_fwd(x: torch.Tensor, inplace: bool = True): rwkv_channel_mixing_pow_and_relu[grid]( x, output, + output.numel(), BLOCK_SIZE=4096, ) @@ -172,23 +175,25 @@ def rwkv_relu_and_square_fwd(x: torch.Tensor, inplace: bool = True): def relu_square_bwd_kernel( out_ptr, forward_input_ptr, + n_elements, BLOCK_SIZE: tl.constexpr, ): """ReLU(x)^2 backward kernel grad_input = grad_output * 2 * x if x > 0 else 0 """ - pid = tl.program_id(0) + pid = tl.program_id(0).to(tl.int64) block_start = pid * BLOCK_SIZE offsets = block_start + tl.arange(0, BLOCK_SIZE) + mask = offsets < n_elements - x = tl.load(forward_input_ptr + offsets).to(tl.float32) - grad = tl.load(out_ptr + offsets).to(tl.float32) + x = tl.load(forward_input_ptr + offsets, mask=mask, other=0.0).to(tl.float32) + grad = tl.load(out_ptr + offsets, mask=mask, other=0.0).to(tl.float32) x = tl.maximum(x, 0.0) grad_input = grad * 2 * x - tl.store(out_ptr + offsets, grad_input.to(out_ptr.dtype.element_ty)) + tl.store(out_ptr + offsets, grad_input.to(out_ptr.dtype.element_ty), mask=mask) @triton.autotune( @@ -210,7 +215,7 @@ def rwkv_mix_bwd_kenel( hidden_dim: tl.constexpr, BLOCK_SIZE: tl.constexpr, ): - pid = tl.program_id(0) + pid = tl.program_id(0).to(tl.int64) offsets = pid * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE) batch_idx = offsets // (token_length * hidden_dim) @@ -218,7 +223,7 @@ def rwkv_mix_bwd_kenel( seq_idx = seq_feat // hidden_dim feat_idx = seq_feat % hidden_dim - is_valid = offsets < (batch_size * token_length * hidden_dim) + is_valid = offsets < (tl.cast(batch_size, tl.int64) * token_length * hidden_dim) dk1 = tl.load(dk1_ptr0 + offsets, mask=is_valid) xk = tl.load(xk_ptr + feat_idx, mask=is_valid) @@ -238,7 +243,7 @@ def rwkv_mix_bwd_kenel( tl.store( dx_prev_ptr + dx_prev_offset, tl.cast(prod, dtype=dx_prev_ptr.dtype.element_ty), - mask=is_first_step, + mask=is_first_step & is_valid, ) @@ -272,6 +277,7 @@ def rwkv_channel_mixing_bwd(grad_output, x, x_prev, x_k, key_weight, value_weigh relu_square_bwd_kernel[grid]( dk, k1_K, + dk.numel(), BLOCK_SIZE=BLOCK_SIZE, ) diff --git a/build/torch-cuda/ops/rwkv7/chunk.py b/build/torch-cuda/ops/rwkv7/chunk.py index 598472ea5c3070ee87f947f75199d65e326c0fa9..e91f6f820eb4ebe3aee3f126971453fa369aaae3 100644 --- a/build/torch-cuda/ops/rwkv7/chunk.py +++ b/build/torch-cuda/ops/rwkv7/chunk.py @@ -7,6 +7,7 @@ import torch +from ...ops.cp import FLACPContext from ...ops.generalized_delta_rule import chunk_dplr_delta_rule @@ -24,7 +25,9 @@ def chunk_rwkv7( cu_seqlens_cpu: torch.LongTensor | None = None, safe_gate: bool = False, chunk_size: int | None = None, - **kwargs, + disable_recompute: bool = False, + cp_context: FLACPContext | None = None, + lower_bound: float | None = None, ): """ Args: @@ -58,12 +61,20 @@ def chunk_rwkv7( When `True`, the kernel can use M=16 TensorCore acceleration. The safe range is approximately `[-5, 0)`. Default: `False`. chunk_size (Optional[int]): - Chunk size for the chunked computation. Default: `None`, which means 16. + Chunk size for the chunked computation. Default: `None`, which means 64 + when `lower_bound` is given and `gate_bound_is_safe(lower_bound, 64)` + holds, and 16 otherwise. + disable_recompute (Optional[bool]): + Whether to disable gradient recomputation in the kernel. Default: `False`. + cp_context (Optional[FLACPContext]): + Context parallel context for distributed training across multiple devices. + When provided, `initial_state` and `output_final_state` are not supported, + and `cp_context.cu_seqlens` is used as the local `cu_seqlens`. Default: `None`. + lower_bound (Optional[float]): + When set, asserts `lower_bound <= w < 0` (the caller is responsible for + the guarantee). Licenses the same tensor-core scheme as `safe_gate=True` + when the bound fits the chunk size. Default: `None`. """ - if 'head_first' in kwargs: - raise DeprecationWarning( - "head_first has been removed. Inputs must be in `[B, T, H, ...]` format.", - ) return chunk_dplr_delta_rule( q=r, k=k, @@ -78,4 +89,7 @@ def chunk_rwkv7( cu_seqlens_cpu=cu_seqlens_cpu, safe_gate=safe_gate, chunk_size=chunk_size, + disable_recompute=disable_recompute, + cp_context=cp_context, + lower_bound=lower_bound, ) diff --git a/build/torch-cuda/ops/rwkv7/fused_addcmul.py b/build/torch-cuda/ops/rwkv7/fused_addcmul.py index ba18ba322282cd191b913c623e15d2011b4e7e4f..45f7a6f1276b4dcc25849484919a73f80f44fdc2 100644 --- a/build/torch-cuda/ops/rwkv7/fused_addcmul.py +++ b/build/torch-cuda/ops/rwkv7/fused_addcmul.py @@ -63,7 +63,7 @@ def fused_addcmul_fwd_kernel( D: tl.constexpr, BD: tl.constexpr, ): - i_b, i_t = tl.program_id(0), tl.program_id(1) * BT + i_b, i_t = tl.program_id(0).to(tl.int64), tl.program_id(1).to(tl.int64) * BT bos = i_b * (T + T_OFFSET) t_vec = i_t + T_OFFSET + tl.arange(0, BT) @@ -133,7 +133,7 @@ def addcmul_bwd_kernel1( BD: tl.constexpr, DTYPE: tl.constexpr, ): - i_b, i_t = tl.program_id(0), tl.program_id(1) + i_b, i_t = tl.program_id(0).to(tl.int64), tl.program_id(1).to(tl.int64) t_idx = T_OFFSET + i_t * BT + tl.arange(0, BT)[:, None] mask_t = t_idx < (T + T_OFFSET) @@ -143,7 +143,7 @@ def addcmul_bwd_kernel1( mask = mask_t & mask_d offset_base = i_b * (T + T_OFFSET) * D - x_idx = (offset_base + t_idx * D + d_idx).to(tl.uint32) + x_idx = offset_base + t_idx * D + d_idx b_dxr = tl.load(dxr + x_idx, mask=mask).to(DTYPE) b_dxw = tl.load(dxw + x_idx, mask=mask).to(DTYPE) diff --git a/build/torch-cuda/ops/rwkv7/fused_k_update.py b/build/torch-cuda/ops/rwkv7/fused_k_update.py index 4a6942675c9b39521f459917697500ab1e0c1922..fe276f7af69999162d886c22fb197ffd1c9ca969 100644 --- a/build/torch-cuda/ops/rwkv7/fused_k_update.py +++ b/build/torch-cuda/ops/rwkv7/fused_k_update.py @@ -10,6 +10,7 @@ import triton import triton.language as tl from ...ops.utils import prepare_chunk_indices +from ...ops.utils.op import unflatten_program_id from ...utils import IS_AMD, autotune_cache_kwargs, get_multiprocessor_count, input_guard NUM_WARPS_AUTOTUNE = [2, 4, 8, 16] if IS_AMD else [2, 4, 8, 16, 32] @@ -37,11 +38,11 @@ def k_update_fwd_kernel_short( BD: tl.constexpr, IS_VARLEN: tl.constexpr, ): - i_b, i_t = tl.program_id(0), tl.program_id(1) + i_b, i_t = tl.program_id(0).to(tl.int64), tl.program_id(1).to(tl.int64) if IS_VARLEN: - bos = tl.load(cu_seqlens + i_b).to(tl.int32) - eos = tl.load(cu_seqlens + i_b + 1).to(tl.int32) + bos = tl.load(cu_seqlens + i_b).to(tl.int64) + eos = tl.load(cu_seqlens + i_b + 1).to(tl.int64) g_t = bos + i_t if g_t >= eos: return @@ -80,13 +81,14 @@ def k_update_fwd_kernel_long( BD: tl.constexpr, BT: tl.constexpr, IS_VARLEN: tl.constexpr, ): - i_d, i_t_blk, i_b = tl.program_id(0), tl.program_id(1), tl.program_id(2) + i_d, i_t_blk = unflatten_program_id(tl.cdiv(D, BD)) + i_b = tl.program_id(1).to(tl.int64) if IS_VARLEN: i_n, i_t_blk = tl.load(chunk_indices + i_t_blk * 2).to(tl.int32), \ tl.load(chunk_indices + i_t_blk * 2 + 1).to(tl.int32) - bos = tl.load(cu_seqlens + i_n).to(tl.int32) - eos = tl.load(cu_seqlens + i_n + 1).to(tl.int32) + bos = tl.load(cu_seqlens + i_n).to(tl.int64) + eos = tl.load(cu_seqlens + i_n + 1).to(tl.int64) t_start = i_t_blk * BT t_end = tl.minimum(t_start + BT, eos - bos) else: @@ -129,11 +131,11 @@ def k_update_bwd_kernel_short( BD: tl.constexpr, IS_VARLEN: tl.constexpr, ): - i_b, i_t_base = tl.program_id(0), tl.program_id(1) * BT + i_b, i_t_base = tl.program_id(0).to(tl.int64), tl.program_id(1).to(tl.int64) * BT if IS_VARLEN: - bos = tl.load(cu_seqlens + i_b).to(tl.int32) - eos = tl.load(cu_seqlens + i_b + 1).to(tl.int32) + bos = tl.load(cu_seqlens + i_b).to(tl.int64) + eos = tl.load(cu_seqlens + i_b + 1).to(tl.int64) seq_len = eos - bos else: bos = i_b * T @@ -180,13 +182,14 @@ def k_update_bwd_kernel_long( BD: tl.constexpr, BT: tl.constexpr, IS_VARLEN: tl.constexpr, ): - i_d, i_t_blk, i_b = tl.program_id(0), tl.program_id(1), tl.program_id(2) + i_d, i_t_blk = unflatten_program_id(tl.cdiv(D, BD)) + i_b = tl.program_id(1).to(tl.int64) if IS_VARLEN: i_n, i_t_blk = tl.load(chunk_indices + i_t_blk * 2).to(tl.int32), \ tl.load(chunk_indices + i_t_blk * 2 + 1).to(tl.int32) - bos = tl.load(cu_seqlens + i_n).to(tl.int32) - eos = tl.load(cu_seqlens + i_n + 1).to(tl.int32) + bos = tl.load(cu_seqlens + i_n).to(tl.int64) + eos = tl.load(cu_seqlens + i_n + 1).to(tl.int64) t_start = i_t_blk * BT t_end = tl.minimum(t_start + BT, eos - bos) else: @@ -252,14 +255,9 @@ def k_update_fwd( BD = triton.next_power_of_2(D) def grid(meta): - return (triton.cdiv(D, meta['BD']), NT, N) + return (triton.cdiv(D, meta['BD']) * NT, N) - k_update_fwd_kernel_long[grid]( - k, a, ka, out, - cu_seqlens, chunk_idx, - T, D, - BD=BD, BT=BT, - ) + k_update_fwd_kernel_long[grid](k, a, ka, out, cu_seqlens, chunk_idx, T, D, BD=BD, BT=BT) return out, use_short, N, T @@ -304,15 +302,9 @@ def k_update_bwd( BD = triton.next_power_of_2(D) def grid(meta): - return (triton.cdiv(D, meta['BD']), NT, N) + return (triton.cdiv(D, meta['BD']) * NT, N) - k_update_bwd_kernel_long[grid]( - grad_out, k, a, ka, - dk, da, dka_tmp, - cu_seqlens, chunk_idx, - T, D, - BD=BD, BT=BT, - ) + k_update_bwd_kernel_long[grid](grad_out, k, a, ka, dk, da, dka_tmp, cu_seqlens, chunk_idx, T, D, BD=BD, BT=BT) if dka_tmp.dim() == 3: dka = dka_tmp.sum(dim=(0, 1), keepdim=True).type_as(ka) diff --git a/build/torch-cuda/ops/rwkv7/fused_recurrent.py b/build/torch-cuda/ops/rwkv7/fused_recurrent.py index 59a33027840b7a19a6dd7135240f2f92729fa8b1..19603c1bc9cb0e4183fee7838a77e27ffcf811c0 100644 --- a/build/torch-cuda/ops/rwkv7/fused_recurrent.py +++ b/build/torch-cuda/ops/rwkv7/fused_recurrent.py @@ -55,7 +55,9 @@ def fused_recurrent_rwkv7_fwd_kernel( IS_VARLEN: tl.constexpr, IS_DECODE: tl.constexpr, ): - i_v, i_nh = tl.program_id(0).to(tl.int64), tl.program_id(1).to(tl.int64) + pid = tl.program_id(0).to(tl.int64) + NV = tl.cdiv(V, BV) + i_v, i_nh = (pid % NV).to(tl.int64), (pid // NV).to(tl.int64) i_n, i_h = i_nh // H, i_nh % H if IS_VARLEN: @@ -145,7 +147,9 @@ def fused_recurrent_rwkv7_fwd( B, T, H, K, V = *k.shape, v.shape[-1] N = B if cu_seqlens is None else len(cu_seqlens) - 1 BK = triton.next_power_of_2(K) - IS_DECODE = (T == 1) + # For packed varlen inputs, decode only when every sequence has exactly one token: + # a zero-length sequence would otherwise read/write out of bounds in the decode branch. + IS_DECODE = (T == 1) and (cu_seqlens is None or bool((cu_seqlens.diff() == 1).all())) h0 = initial_state if not output_final_state: @@ -154,7 +158,7 @@ def fused_recurrent_rwkv7_fwd( ht = r.new_empty(N, H, K, V, dtype=torch.float32) o = torch.empty_like(v) - def grid(meta): return (triton.cdiv(V, meta['BV']), N * H) + def grid(meta): return (triton.cdiv(V, meta['BV']) * N * H,) fused_recurrent_rwkv7_fwd_kernel[grid]( r, w, @@ -190,7 +194,6 @@ def fused_recurrent_rwkv7( initial_state: torch.Tensor | None = None, output_final_state: bool = True, cu_seqlens: torch.LongTensor | None = None, - **kwargs, ): """ Args: @@ -218,10 +221,6 @@ def fused_recurrent_rwkv7( Cumulative sequence lengths of shape `[N+1]` used for variable-length training, consistent with the FlashAttention API. """ - if 'head_first' in kwargs: - raise DeprecationWarning( - "head_first has been removed. Inputs must be in `[B, T, H, ...]` format.", - ) return fused_recurrent_dplr_delta_rule( q=r, k=k, @@ -248,7 +247,6 @@ def fused_mul_recurrent_rwkv7( output_final_state: bool = False, reverse: bool = False, cu_seqlens: torch.Tensor | None = None, - **kwargs, ) -> tuple[torch.Tensor, torch.Tensor]: r""" This function computes the recurrence S_t = S_t @ (I + a_t b_t^T) + v_t k_t^T in a recurrent manner. @@ -280,10 +278,6 @@ def fused_mul_recurrent_rwkv7( Cumulative sequence lengths of shape `[N+1]` used for variable-length training, consistent with the FlashAttention API. """ - if 'head_first' in kwargs: - raise DeprecationWarning( - "head_first has been removed. Inputs must be in `[B, T, H, ...]` format.", - ) if cu_seqlens is not None: if r.shape[0] != 1: raise ValueError( diff --git a/build/torch-cuda/ops/rwkv7/gate_output_correction.py b/build/torch-cuda/ops/rwkv7/gate_output_correction.py index 59ca4ca68db21184f9fa47e61b681b21617f9de0..cfdb1c3f68ca77fb5c1753d06e8bd40dddda9cb6 100644 --- a/build/torch-cuda/ops/rwkv7/gate_output_correction.py +++ b/build/torch-cuda/ops/rwkv7/gate_output_correction.py @@ -84,8 +84,8 @@ def gate_output_correction_fwd_kernel( BLOCK_SIZE_D: tl.constexpr, BT: tl.constexpr, ): - pid_b, pid_t_block = tl.program_id(0), tl.program_id(1) - pid_h = tl.program_id(2) + pid_b, pid_t_block = tl.program_id(0).to(tl.int64), tl.program_id(1).to(tl.int64) + pid_h = tl.program_id(2).to(tl.int64) t_local = pid_t_block * BT + tl.arange(0, BT)[:, None] t_idx = T_OFFSET + t_local mask_t = t_local < T @@ -134,8 +134,8 @@ def gate_output_correction_bwd_kernel( BLOCK_SIZE_D: tl.constexpr, BT: tl.constexpr, ): - pid_b, pid_t_block = tl.program_id(0), tl.program_id(1) - pid_h = tl.program_id(2) + pid_b, pid_t_block = tl.program_id(0).to(tl.int64), tl.program_id(1).to(tl.int64) + pid_h = tl.program_id(2).to(tl.int64) t_local = pid_t_block * BT + tl.arange(0, BT)[:, None] t_idx = T_OFFSET + t_local diff --git a/build/torch-cuda/ops/simple_gla/__init__.py b/build/torch-cuda/ops/simple_gla/__init__.py index ea5bb8eb53bff143f76c3728f1f556d9c3d2591d..d914b5167440d03643fec2e68d58d495eefc6538 100644 --- a/build/torch-cuda/ops/simple_gla/__init__.py +++ b/build/torch-cuda/ops/simple_gla/__init__.py @@ -5,6 +5,7 @@ # For a list of all contributors, visit: # https://github.com/fla-org/flash-linear-attention/graphs/contributors + from .chunk import chunk_simple_gla from .fused_chunk import fused_chunk_simple_gla from .fused_recurrent import fused_recurrent_simple_gla @@ -16,3 +17,4 @@ __all__ = [ 'fused_recurrent_simple_gla', 'parallel_simple_gla', ] + diff --git a/build/torch-cuda/ops/simple_gla/backends/__init__.py b/build/torch-cuda/ops/simple_gla/backends/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..546163768c14bff626769f78ac7700ed0e3398ea --- /dev/null +++ b/build/torch-cuda/ops/simple_gla/backends/__init__.py @@ -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 diff --git a/build/torch-cuda/ops/simple_gla/chunk.py b/build/torch-cuda/ops/simple_gla/chunk.py index 08313e35b8d2993f4dd05c0f83d85161b47fb82c..6ed7aedaa13d185a3e3593a1d7ff00c35378e11d 100644 --- a/build/torch-cuda/ops/simple_gla/chunk.py +++ b/build/torch-cuda/ops/simple_gla/chunk.py @@ -295,7 +295,7 @@ def chunk_simple_gla( >>> import torch >>> import torch.nn.functional as F >>> from einops import rearrange - >>> from ...ops.simple_gla import chunk_simple_gla + >>> from fla.ops.simple_gla import chunk_simple_gla # inputs with equal lengths >>> B, T, H, K, V = 4, 2048, 4, 512, 512 >>> q = torch.randn(B, T, H, K, device='cuda') diff --git a/build/torch-cuda/ops/simple_gla/fused_chunk.py b/build/torch-cuda/ops/simple_gla/fused_chunk.py index 6e88e9c604b354c1714d5ab55029df7ab9bd44ae..453807721358ceb3af21026ace9e8d00e983da92 100644 --- a/build/torch-cuda/ops/simple_gla/fused_chunk.py +++ b/build/torch-cuda/ops/simple_gla/fused_chunk.py @@ -7,15 +7,17 @@ import torch +from ...backends import dispatch from ...ops.common.fused_chunk import fused_chunk +@dispatch def fused_chunk_simple_gla( q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, - g: torch.Tensor = None, - g_gamma: torch.Tensor = None, + g: torch.Tensor | None = None, + g_gamma: torch.Tensor | None = None, scale: float | None = None, initial_state: torch.Tensor | None = None, output_final_state: bool = False, @@ -29,25 +31,25 @@ def fused_chunk_simple_gla( keys of shape `[B, T, H, K]`. v (torch.Tensor): values of shape `[B, T, H, V]`. - g (torch.Tensor): + g (torch.Tensor, Optional): Forget gates of shape `[B, T, H]`. - Compared to GLA, the gating is head-wise instead of elementwise. - g_gamma (torch.Tensor): + Compared to GLA, the gating is head-wise instead of elementwise. Default: `None`. + g_gamma (torch.Tensor, Optional): Log decay of shape `[H]`. Head-wise data-independent decay is used if `g_gamma` is provided. - Only one of `g` or `g_gamma` should be provided. - scale (Optional[int]): + Only one of `g` or `g_gamma` should be provided. Default: `None`. + scale (float, Optional): Scale factor for the attention scores. If not provided, it will default to `1 / sqrt(K)`. Default: `None`. - initial_state (Optional[torch.Tensor]): + initial_state (torch.Tensor, Optional): Initial state of shape `[N, H, K, V]` for `N` input sequences. For equal-length input sequences, `N` equals the batch size `B`. Default: `None`. - output_final_state (Optional[bool]): + output_final_state (bool, Optional): Whether to output the final state of shape `[N, H, K, V]`. Default: `False`. - cu_seqlens (torch.LongTensor): + cu_seqlens (torch.LongTensor, Optional): Cumulative sequence lengths of shape `[N+1]` used for variable-length training, - consistent with the FlashAttention API. + consistent with the FlashAttention API. Default: `None`. Returns: o (torch.Tensor): @@ -59,7 +61,7 @@ def fused_chunk_simple_gla( >>> import torch >>> import torch.nn.functional as F >>> from einops import rearrange - >>> from ...ops.simple_gla import fused_chunk_simple_gla + >>> from fla.ops.simple_gla import fused_chunk_simple_gla # inputs with equal lengths >>> B, T, H, K, V = 4, 2048, 4, 512, 512 >>> q = torch.randn(B, T, H, K, device='cuda') @@ -67,21 +69,20 @@ def fused_chunk_simple_gla( >>> v = torch.randn(B, T, H, V, device='cuda') >>> g = F.logsigmoid(torch.randn(B, T, H, device='cuda')) >>> h0 = torch.randn(B, H, K, V, device='cuda') - >>> o, ht = fused_chunk_simple_gla( - q, k, v, g, - initial_state=h0, - output_final_state=True - ) + >>> o, ht = fused_chunk_simple_gla(q=q, k=k, v=v, g=g, initial_state=h0, output_final_state=True) # for variable-length inputs, the batch size `B` is expected to be 1 and `cu_seqlens` is required >>> q, k, v, g = map(lambda x: rearrange(x, 'b t h d -> 1 (b t) h d'), (q, k, v, g)) # for a batch with 4 sequences, `cu_seqlens` with 5 start/end positions are expected >>> cu_seqlens = q.new_tensor([0, 2048, 4096, 6144, 8192], dtype=torch.long) >>> o, ht = fused_chunk_simple_gla( - q, k, v, g, - initial_state=h0, - output_final_state=True, - cu_seqlens=cu_seqlens - ) + ... q=q, + ... k=k, + ... v=v, + ... g=g, + ... initial_state=h0, + ... output_final_state=True, + ... cu_seqlens=cu_seqlens, + ... ) """ if cu_seqlens is not None: if q.shape[0] != 1: diff --git a/build/torch-cuda/ops/simple_gla/fused_recurrent.py b/build/torch-cuda/ops/simple_gla/fused_recurrent.py index dbc209dbd8c36eabdc01de640531c545a8c7f885..10cda1004acae8a2965b6c3d7267999093fdad2f 100644 --- a/build/torch-cuda/ops/simple_gla/fused_recurrent.py +++ b/build/torch-cuda/ops/simple_gla/fused_recurrent.py @@ -9,15 +9,17 @@ import warnings import torch +from ...backends import dispatch from ...ops.common.fused_recurrent import fused_recurrent +@dispatch def fused_recurrent_simple_gla( q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, - g: torch.Tensor = None, - g_gamma: torch.Tensor = None, + g: torch.Tensor | None = None, + g_gamma: torch.Tensor | None = None, scale: float | None = None, initial_state: torch.Tensor | None = None, output_final_state: bool = False, @@ -34,31 +36,30 @@ def fused_recurrent_simple_gla( keys of shape `[B, T, H, K]`. v (torch.Tensor): values of shape `[B, T, H, V]`. - g (torch.Tensor): + g (torch.Tensor, Optional): Forget gates of shape `[B, T, H]`. - Compared to GLA, the gating is head-wise instead of elementwise. - g_gamma (torch.Tensor): + Compared to GLA, the gating is head-wise instead of elementwise. Default: `None`. + g_gamma (torch.Tensor, Optional): Log decay of shape `[H]`. Head-wise data-independent decay is used if `g_gamma` is provided. - Only one of `g` or `g_gamma` should be provided. - scale (Optional[float]): + Only one of `g` or `g_gamma` should be provided. Default: `None`. + scale (float, Optional): Scale factor for the attention scores. If not provided, it will default to `1 / sqrt(K)`. Default: `None`. - initial_state (Optional[torch.Tensor]): + initial_state (torch.Tensor, Optional): Initial state of shape `[N, H, K, V]` for `N` input sequences. For equal-length input sequences, `N` equals the batch size `B`. Default: `None`. - output_final_state (Optional[bool]): + output_final_state (bool, Optional): Whether to output the final state of shape `[N, H, K, V]`. Default: `False`. - reverse (Optional[bool]): + reverse (bool, Optional): If `True`, process the state passing in reverse order. Default: `False`. - state_v_first (Optional[bool]): + state_v_first (bool, Optional): Store the recurrent state in V-first ``[V, K]`` layout instead of the default ``[K, V]``. Default: ``False``. - cu_seqlens (torch.LongTensor): - Cumulative sequence lengths of shape `[N+1]` used for variable-length training, - consistent with the FlashAttention API. `initial_state`, `final_state`, and `dh0`/`dht` all follow the chosen layout. - Default: `False`. + cu_seqlens (torch.LongTensor, Optional): + Cumulative sequence lengths of shape `[N+1]` used for variable-length training, + consistent with the FlashAttention API. Default: `None`. Returns: o (torch.Tensor): @@ -70,7 +71,7 @@ def fused_recurrent_simple_gla( >>> import torch >>> import torch.nn.functional as F >>> from einops import rearrange - >>> from ...ops.simple_gla import fused_recurrent_simple_gla + >>> from fla.ops.simple_gla import fused_recurrent_simple_gla # inputs with equal lengths >>> B, T, H, K, V = 4, 2048, 4, 512, 512 >>> q = torch.randn(B, T, H, K, device='cuda') @@ -78,21 +79,20 @@ def fused_recurrent_simple_gla( >>> v = torch.randn(B, T, H, V, device='cuda') >>> g = F.logsigmoid(torch.randn(B, T, H, K, device='cuda')) >>> h0 = torch.randn(B, H, K, V, device='cuda') - >>> o, ht = fused_recurrent_simple_gla( - q, k, v, g, - initial_state=h0, - output_final_state=True - ) + >>> o, ht = fused_recurrent_simple_gla(q=q, k=k, v=v, g=g, initial_state=h0, output_final_state=True) # for variable-length inputs, the batch size `B` is expected to be 1 and `cu_seqlens` is required >>> q, k, v, g = map(lambda x: rearrange(x, 'b t h d -> 1 (b t) h d'), (q, k, v, g)) # for a batch with 4 sequences, `cu_seqlens` with 5 start/end positions are expected >>> cu_seqlens = q.new_tensor([0, 2048, 4096, 6144, 8192], dtype=torch.long) >>> o_var, ht_var = fused_recurrent_simple_gla( - q, k, v, g, - initial_state=h0, - output_final_state=True, - cu_seqlens=cu_seqlens - ) + ... q=q, + ... k=k, + ... v=v, + ... g=g, + ... initial_state=h0, + ... output_final_state=True, + ... cu_seqlens=cu_seqlens, + ... ) """ if 'transpose_state_layout' in kwargs: if state_v_first: @@ -127,7 +127,7 @@ def fused_recurrent_simple_gla( initial_state=initial_state, output_final_state=output_final_state, reverse=reverse, - cu_seqlens=cu_seqlens, state_v_first=state_v_first, + cu_seqlens=cu_seqlens, ) return o, final_state diff --git a/build/torch-cuda/ops/simple_gla/parallel.py b/build/torch-cuda/ops/simple_gla/parallel.py index e25f69bbaee7fbfd0817f55c70dd23833043ffc6..6d4e596a8a274382dea9ee62c19a365c3d5c7e74 100644 --- a/build/torch-cuda/ops/simple_gla/parallel.py +++ b/build/torch-cuda/ops/simple_gla/parallel.py @@ -9,10 +9,11 @@ import torch import triton import triton.language as tl +from ...backends import dispatch from ...ops.utils import prepare_chunk_indices from ...ops.utils.constant import RCP_LN2 from ...ops.utils.cumsum import chunk_global_cumsum, chunk_local_cumsum -from ...ops.utils.op import exp2 +from ...ops.utils.op import exp2, unflatten_program_id from ...utils import ( IS_INTEL_ALCHEMIST, IS_NVIDIA_HOPPER, @@ -68,14 +69,15 @@ def parallel_simple_gla_fwd_kernel( IS_VARLEN: tl.constexpr, USE_G: tl.constexpr, ): - i_kv, i_t, i_bh = tl.program_id(0), tl.program_id(1), tl.program_id(2) + i_kv, i_t = unflatten_program_id(tl.cdiv(K, BK) * NV) + i_bh = tl.program_id(1).to(tl.int64) i_k, i_v = i_kv // NV, i_kv % NV i_b, i_h = i_bh // H, i_bh % H all = B * T 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.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) + i_n, i_t = tl.load(chunk_indices + i_t * 2).to(tl.int32), tl.load(chunk_indices + i_t * 2 + 1).to(tl.int64) + bos, eos = tl.load(cu_seqlens + i_n).to(tl.int64), tl.load(cu_seqlens + i_n + 1).to(tl.int64) T = eos - bos else: bos, eos = i_b * T, i_b * T + T @@ -89,19 +91,22 @@ def parallel_simple_gla_fwd_kernel( if OUTPUT_ATTENTIONS: attn += i_k * B * H * T * T + (bos * H + i_h * T) * T - p_q = tl.make_block_ptr(q, (T, K), (H*K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) + # [BT] + o_q = i_t * BT + tl.arange(0, BT) + o_kk = i_k * BK + tl.arange(0, BK) + o_vv = i_v * BV + tl.arange(0, BV) + m_q = o_q < T + m_qk = m_q[:, None] & (o_kk[None, :] < K) + m_qv = m_q[:, None] & (o_vv[None, :] < V) + p_q = q + o_q[:, None] * (H*K) + o_kk[None, :] # the Q block is kept in the shared memory throughout the whole kernel # [BT, BK] - b_q = tl.load(p_q, boundary_check=(0, 1)) + b_q = tl.load(p_q, mask=m_qk, other=0.0) b_q = (b_q * scale).to(b_q.dtype) b_o = tl.zeros([BT, BV], dtype=tl.float32) - # [BT] - o_q = i_t * BT + tl.arange(0, BT) - m_q = o_q < T - # Q block and K block have overlap. - # masks required + # overlapping Q and K blocks require causal masks. if USE_G: # [BT,] b_gq = tl.load(g + o_q * H, mask=m_q, other=float('-inf')).to(tl.float32) @@ -110,15 +115,16 @@ def parallel_simple_gla_fwd_kernel( b_gq = None for i_s in range(i_t * BT, min((i_t + 1) * BT, T), BS): - p_k = tl.make_block_ptr(k, (K, T), (1, H*K), (i_k * BK, i_s), (BK, BS), (0, 1)) - p_v = tl.make_block_ptr(v, (T, V), (H*V, 1), (i_s, i_v * BV), (BS, BV), (1, 0)) - o_k = i_s + tl.arange(0, BS) m_k = o_k < T + m_kk = (o_kk[:, None] < K) & m_k[None, :] + m_kv = m_k[:, None] & (o_vv[None, :] < V) + p_k = k + o_kk[:, None] + o_k[None, :] * (H*K) + p_v = v + o_k[:, None] * (H*V) + o_vv[None, :] # [BK, BS] - b_k = tl.load(p_k, boundary_check=(0, 1)) + b_k = tl.load(p_k, mask=m_kk, other=0.0) # [BS, BV] - b_v = tl.load(p_v, boundary_check=(0, 1)) + b_v = tl.load(p_v, mask=m_kv, other=0.0) # [BT, BS] m_s = (o_q[:, None] >= o_k[None, :]) & (m_q[:, None] & m_k[None, :]) b_s = tl.dot(b_q, b_k) @@ -128,20 +134,21 @@ def parallel_simple_gla_fwd_kernel( b_s = tl.where(m_s, b_s, 0) # [BT, BV] if i_s >= 0: - b_o += tl.dot(b_s.to(b_q.dtype), b_v) + b_o = tl.dot(b_s.to(b_q.dtype), b_v, b_o) if OUTPUT_ATTENTIONS: - p_a = tl.make_block_ptr(attn, (T, T), (T, 1), (i_t * BT, i_s), (BT, BS), (1, 0)) - tl.store(p_a, b_s.to(p_a.dtype.element_ty), boundary_check=(0, 1)) + p_a = attn + o_q[:, None] * T + o_k[None, :] + tl.store(p_a, b_s.to(p_a.dtype.element_ty), mask=m_q[:, None] & m_k[None, :]) for i_s in range(i_t * BT - BS, -BS, -BS): - p_k = tl.make_block_ptr(k, (K, T), (1, H*K), (i_k * BK, i_s), (BK, BS), (0, 1)) - p_v = tl.make_block_ptr(v, (T, V), (H*V, 1), (i_s, i_v * BV), (BS, BV), (1, 0)) - o_k = i_s + tl.arange(0, BS) m_k = o_k < T + m_kk = (o_kk[:, None] < K) & m_k[None, :] + m_kv = m_k[:, None] & (o_vv[None, :] < V) + p_k = k + o_kk[:, None] + o_k[None, :] * (H*K) + p_v = v + o_k[:, None] * (H*V) + o_vv[None, :] # [BK, BS] - b_k = tl.load(p_k, boundary_check=(0, 1)) + b_k = tl.load(p_k, mask=m_kk, other=0.0) # [BS, BV] - b_v = tl.load(p_v, boundary_check=(0, 1)) + b_v = tl.load(p_v, mask=m_kv, other=0.0) # [BT, BS] m_s = m_q[:, None] & m_k[None, :] b_s = tl.dot(b_q, b_k) @@ -149,17 +156,17 @@ def parallel_simple_gla_fwd_kernel( b_g = tl.load(g + o_k * H, mask=m_k, other=0) b_gn = tl.load(g + (min(i_s + BS, T) - 1) * H) b_gp = tl.load(g + (i_s-1) * H) if i_s % BT > 0 else 0. - # No concrete meaning. Just to avoid some layout bugs. + # this factorization avoids Triton layout bugs. b_s *= exp2(b_gq[:, None] + (b_gn - b_g)[None, :]) b_gq += b_gn - b_gp b_s = tl.where(m_s, b_s, 0) if OUTPUT_ATTENTIONS: - p_a = tl.make_block_ptr(attn, (T, T), (T, 1), (i_t * BT, i_s), (BT, BS), (1, 0)) - tl.store(p_a, b_s.to(p_a.dtype.element_ty), boundary_check=(0, 1)) + p_a = attn + o_q[:, None] * T + o_k[None, :] + tl.store(p_a, b_s.to(p_a.dtype.element_ty), mask=m_q[:, None] & m_k[None, :]) if i_s >= 0: - b_o += tl.dot(b_s.to(b_v.dtype), b_v) - p_o = tl.make_block_ptr(o, (T, V), (H*V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) - tl.store(p_o, b_o.to(p_o.dtype.element_ty), boundary_check=(0, 1)) + b_o = tl.dot(b_s.to(b_v.dtype), b_v, b_o) + p_o = o + o_q[:, None] * (H*V) + o_vv[None, :] + tl.store(p_o, b_o.to(p_o.dtype.element_ty), mask=m_qv) @triton.jit(do_not_specialize=['T']) @@ -185,25 +192,30 @@ def parallel_simple_gla_bwd_kernel_dq( BV: tl.constexpr, USE_G: tl.constexpr, ): - p_do = tl.make_block_ptr(do, (T, V), (H*V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) + # [BT] + o_q = i_t * BT + tl.arange(0, BT) + o_kk = i_k * BK + tl.arange(0, BK) + o_vv = i_v * BV + tl.arange(0, BV) + m_q = o_q < T + m_qk = m_q[:, None] & (o_kk[None, :] < K) + m_qv = m_q[:, None] & (o_vv[None, :] < V) + p_do = do + o_q[:, None] * (H*V) + o_vv[None, :] # [BT, BV] - b_do = tl.load(p_do, boundary_check=(0, 1)) + b_do = tl.load(p_do, mask=m_qv, other=0.0) # [BT, BK] b_dq = tl.zeros([BT, BK], dtype=tl.float32) - # [BT] - o_q = i_t * BT + tl.arange(0, BT) - m_q = o_q < T for i_s in range(0, i_t * BT, BS): - p_k = tl.make_block_ptr(k, (T, K), (H*K, 1), (i_s, i_k * BK), (BS, BK), (1, 0)) - p_v = tl.make_block_ptr(v, (V, T), (1, H*V), (i_v * BV, i_s), (BV, BS), (0, 1)) - o_k = i_s + tl.arange(0, BS) m_k = o_k < T + m_kk = m_k[:, None] & (o_kk[None, :] < K) + m_vk = (o_vv[:, None] < V) & m_k[None, :] + p_k = k + o_k[:, None] * (H*K) + o_kk[None, :] + p_v = v + o_vv[:, None] + o_k[None, :] * (H*V) # [BS, BK] - b_k = tl.load(p_k, boundary_check=(0, 1)) + b_k = tl.load(p_k, mask=m_kk, other=0.0) # [BV, BS] - b_v = tl.load(p_v, boundary_check=(0, 1)) + b_v = tl.load(p_v, mask=m_vk, other=0.0) # [BT, BV] @ [BV, BS] = [BT, BS] b_ds = tl.dot(b_do, b_v) if USE_G: @@ -214,24 +226,25 @@ def parallel_simple_gla_bwd_kernel_dq( if i_s > 0: b_dq *= exp2(b_gn - b_gp) # [BT, BS] @ [BS, BK] = [BT, BK] - b_dq += tl.dot(b_ds.to(b_v.dtype), b_k) + b_dq = tl.dot(b_ds.to(b_v.dtype), b_k, b_dq) if USE_G: # [BT,] b_gq = tl.load(g + o_q * H, mask=m_q, other=float('-inf')) # [BT, BK] b_dq *= exp2(b_gq)[:, None] - # Q block and K block have overlap. masks required + # overlapping Q and K blocks require causal masks. for i_s in range(i_t * BT, min((i_t + 1) * BT, T), BS): - p_k = tl.make_block_ptr(k, (T, K), (H*K, 1), (i_s, i_k * BK), (BS, BK), (1, 0)) - p_v = tl.make_block_ptr(v, (V, T), (1, H*V), (i_v * BV, i_s), (BV, BS), (0, 1)) - o_k = i_s + tl.arange(0, BS) m_k = o_k < T + m_kk = m_k[:, None] & (o_kk[None, :] < K) + m_vk = (o_vv[:, None] < V) & m_k[None, :] + p_k = k + o_k[:, None] * (H*K) + o_kk[None, :] + p_v = v + o_vv[:, None] + o_k[None, :] * (H*V) # [BS, BK] - b_k = tl.load(p_k, boundary_check=(0, 1)) + b_k = tl.load(p_k, mask=m_kk, other=0.0) # [BV, BS] - b_v = tl.load(p_v, boundary_check=(0, 1)) + b_v = tl.load(p_v, mask=m_vk, other=0.0) # [BT, BV] @ [BV, BS] = [BT, BS] b_ds = tl.dot(b_do, b_v) if USE_G: @@ -240,17 +253,17 @@ def parallel_simple_gla_bwd_kernel_dq( m_s = (o_q[:, None] >= o_k[None, :]) & (m_q[:, None] & m_k[None, :]) b_ds = tl.where(m_s, b_ds, 0) # [BT, BK] - b_dq += tl.dot(b_ds.to(b_k.dtype), b_k) + b_dq = tl.dot(b_ds.to(b_k.dtype), b_k, b_dq) b_dq *= scale - p_dq = tl.make_block_ptr(dq, (T, K), (H*K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) - tl.store(p_dq, b_dq.to(p_dq.dtype.element_ty), boundary_check=(0, 1)) + p_dq = dq + o_q[:, None] * (H*K) + o_kk[None, :] + tl.store(p_dq, b_dq.to(p_dq.dtype.element_ty), mask=m_qk) if USE_G: - p_q = tl.make_block_ptr(q, (T, K), (H*K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) - b_q = tl.load(p_q, boundary_check=(0, 1)) + p_q = q + o_q[:, None] * (H*K) + o_kk[None, :] + b_q = tl.load(p_q, mask=m_qk, other=0.0) b_dg = tl.sum(b_dq * b_q, 1) - p_dg = tl.make_block_ptr(dg, (T,), (H,), (i_t * BT,), (BT,), (0,)) - tl.store(p_dg, b_dg.to(p_dg.dtype.element_ty), boundary_check=(0,)) + p_dg = dg + o_q * H + tl.store(p_dg, b_dg.to(p_dg.dtype.element_ty), mask=m_q) @triton.jit(do_not_specialize=['T']) @@ -278,29 +291,34 @@ def parallel_simple_gla_bwd_kernel_dkv( USE_G: tl.constexpr, ): o_k = i_t * BT + tl.arange(0, BT) + o_kk = i_k * BK + tl.arange(0, BK) + o_vv = i_v * BV + tl.arange(0, BV) m_k = o_k < T + m_kk = m_k[:, None] & (o_kk[None, :] < K) + m_kv = m_k[:, None] & (o_vv[None, :] < V) # [BT, BK] - p_k = tl.make_block_ptr(k, (T, K), (H*K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) - b_k = tl.load(p_k, boundary_check=(0, 1)) + p_k = k + o_k[:, None] * (H*K) + o_kk[None, :] + b_k = tl.load(p_k, mask=m_kk, other=0.0) b_dk = tl.zeros([BT, BK], dtype=tl.float32) # [BT, BV] - p_v = tl.make_block_ptr(v, (T, V), (H*V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) - b_v = tl.load(p_v, boundary_check=(0, 1)) + p_v = v + o_k[:, None] * (H*V) + o_vv[None, :] + b_v = tl.load(p_v, mask=m_kv, other=0.0) b_dv = tl.zeros([BT, BV], dtype=tl.float32) if USE_G: b_gk = tl.load(g + o_k * H, mask=m_k, other=0) NTS = tl.cdiv(T, BS) # [BT, BK] for i_s in range(NTS * BS - BS, (i_t + 1) * BT - BS, -BS): - p_q = tl.make_block_ptr(q, (T, K), (H*K, 1), (i_s, i_k * BK), (BS, BK), (1, 0)) - p_do = tl.make_block_ptr(do, (T, V), (H*V, 1), (i_s, i_v * BV), (BS, BV), (1, 0)) - o_q = i_s + tl.arange(0, BS) m_q = o_q < T + m_qk = m_q[:, None] & (o_kk[None, :] < K) + m_qv = m_q[:, None] & (o_vv[None, :] < V) + p_q = q + o_q[:, None] * (H*K) + o_kk[None, :] + p_do = do + o_q[:, None] * (H*V) + o_vv[None, :] # [BS, BK] - b_q = tl.load(p_q, boundary_check=(0, 1)) + b_q = tl.load(p_q, mask=m_qk, other=0.0) # [BS, BV] - b_do = tl.load(p_do, boundary_check=(0, 1)) + b_do = tl.load(p_do, mask=m_qv, other=0.0) # [BT, BS] b_ds = tl.dot(b_v, tl.trans(b_do)) b_s = tl.dot(b_k, tl.trans(b_q)) @@ -316,9 +334,9 @@ def parallel_simple_gla_bwd_kernel_dkv( b_ds *= b_gqn[None, :] b_s *= b_gqn[None, :] # [BT, BK] - b_dk += tl.dot(b_ds.to(b_q.dtype), b_q) + b_dk = tl.dot(b_ds.to(b_q.dtype), b_q, b_dk) # [BT, BV] - b_dv += tl.dot(b_s.to(b_do.dtype), b_do) + b_dv = tl.dot(b_s.to(b_do.dtype), b_do, b_dv) if USE_G: b_gn = tl.load(g + (min(i_t * BT + BT, T) - 1) * H) @@ -328,15 +346,16 @@ def parallel_simple_gla_bwd_kernel_dkv( b_dv *= b_gpn for i_s in range(i_t * BT, min((i_t + 1) * BT, T), BS): - p_q = tl.make_block_ptr(q, (T, K), (H*K, 1), (i_s, i_k * BK), (BS, BK), (1, 0)) - p_do = tl.make_block_ptr(do, (T, V), (H*V, 1), (i_s, i_v * BV), (BS, BV), (1, 0)) - o_q = i_s + tl.arange(0, BS) m_q = o_q < T + m_qk = m_q[:, None] & (o_kk[None, :] < K) + m_qv = m_q[:, None] & (o_vv[None, :] < V) + p_q = q + o_q[:, None] * (H*K) + o_kk[None, :] + p_do = do + o_q[:, None] * (H*V) + o_vv[None, :] # [BS, BK] - b_q = tl.load(p_q, boundary_check=(0, 1)) + b_q = tl.load(p_q, mask=m_qk, other=0.0) # [BS, BV] - b_do = tl.load(p_do, boundary_check=(0, 1)) + b_do = tl.load(p_do, mask=m_qv, other=0.0) # [BS] b_s = tl.dot(b_k, tl.trans(b_q)) b_ds = tl.dot(b_v, tl.trans(b_do)) @@ -350,14 +369,14 @@ def parallel_simple_gla_bwd_kernel_dkv( b_s = tl.where(m_s, b_s, 0) b_ds = tl.where(m_s, b_ds, 0) # [BT, BK] - b_dk += tl.dot(b_ds.to(b_q.dtype), b_q) - b_dv += tl.dot(b_s.to(b_do.dtype), b_do) + b_dk = tl.dot(b_ds.to(b_q.dtype), b_q, b_dk) + b_dv = tl.dot(b_s.to(b_do.dtype), b_do, b_dv) b_dk *= scale b_dv *= scale - p_dk = tl.make_block_ptr(dk, (T, K), (H*K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) - p_dv = tl.make_block_ptr(dv, (T, V), (H*V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) - tl.store(p_dk, b_dk.to(p_dk.dtype.element_ty), boundary_check=(0, 1)) - tl.store(p_dv, b_dv.to(p_dv.dtype.element_ty), boundary_check=(0, 1)) + p_dk = dk + o_k[:, None] * (H*K) + o_kk[None, :] + p_dv = dv + o_k[:, None] * (H*V) + o_vv[None, :] + tl.store(p_dk, b_dk.to(p_dk.dtype.element_ty), mask=m_kk) + tl.store(p_dv, b_dv.to(p_dv.dtype.element_ty), mask=m_kv) if USE_G: b_dg = tl.load(dg + o_k * H, mask=m_k, other=0) b_dg -= tl.sum(b_dk * b_k, 1) @@ -404,7 +423,8 @@ def parallel_simple_gla_bwd_kernel( IS_VARLEN: tl.constexpr, USE_G: tl.constexpr, ): - i_kv, i_t, i_bh = tl.program_id(0), tl.program_id(1), tl.program_id(2) + i_kv, i_t = unflatten_program_id(tl.cdiv(K, BK) * NV) + i_bh = tl.program_id(1).to(tl.int64) i_k, i_v = i_kv // NV, i_kv % NV i_b, i_h = i_bh // H, i_bh % H dq += i_v * B * H * T * K @@ -414,8 +434,8 @@ def parallel_simple_gla_bwd_kernel( dg += i_kv * B * H * T 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.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) + i_n, i_t = tl.load(chunk_indices + i_t * 2).to(tl.int32), tl.load(chunk_indices + i_t * 2 + 1).to(tl.int64) + bos, eos = tl.load(cu_seqlens + i_n).to(tl.int64), tl.load(cu_seqlens + i_n + 1).to(tl.int64) T = eos - bos else: bos, eos = i_b * T, i_b * T + T @@ -512,14 +532,8 @@ def parallel_simple_gla_fwd( # local cumulative decay in log space if g is not None: - g = chunk_local_cumsum( - g, - chunk_size, - scale=RCP_LN2, - cu_seqlens=cu_seqlens, - chunk_indices=chunk_indices, - ) - grid = (NK * NV, NT, B * H) + g = chunk_local_cumsum(g=g, chunk_size=chunk_size, scale=RCP_LN2, cu_seqlens=cu_seqlens, chunk_indices=chunk_indices) + grid = (NK * NV * NT, B * H) o = torch.empty(NK, *v.shape, dtype=v.dtype if NK == 1 else torch.float, device=q.device) attn = q.new_zeros(NK, B, H, T, T) if output_attentions else None @@ -533,9 +547,9 @@ def parallel_simple_gla_fwd( scale=scale, cu_seqlens=cu_seqlens, chunk_indices=chunk_indices, + T=T, B=B, H=H, - T=T, K=K, V=V, BT=BT, @@ -589,7 +603,7 @@ def parallel_simple_gla_bwd( chunk_indices = prepare_chunk_indices(cu_seqlens, chunk_size) if cu_seqlens is not None else None NT = triton.cdiv(T, BT) if cu_seqlens is None else len(chunk_indices) - grid = (NK * NV, NT, B * H) + grid = (NK * NV * NT, B * H) parallel_simple_gla_bwd_kernel[grid]( q=q, k=k, @@ -600,9 +614,9 @@ def parallel_simple_gla_bwd( dk=dk, dv=dv, dg=dg, + scale=scale, cu_seqlens=cu_seqlens, chunk_indices=chunk_indices, - scale=scale, T=T, B=B, H=H, @@ -630,7 +644,10 @@ class ParallelSimpleGLAFunction(torch.autograd.Function): ctx.dtype = q.dtype chunk_indices = prepare_chunk_indices( - cu_seqlens, chunk_size, cu_seqlens_cpu=cu_seqlens_cpu) if cu_seqlens is not None else None + cu_seqlens, + chunk_size, + cu_seqlens_cpu=cu_seqlens_cpu, + ) if cu_seqlens is not None else None o, g, attn = parallel_simple_gla_fwd( q=q, @@ -667,6 +684,7 @@ class ParallelSimpleGLAFunction(torch.autograd.Function): return dq.to(q), dk.to(k), dv.to(v), dg.to(ctx.dtype) if dg is not None else None, None, None, None, None +@dispatch def parallel_simple_gla( q: torch.Tensor, k: torch.Tensor, @@ -676,7 +694,6 @@ def parallel_simple_gla( output_attentions: bool = False, cu_seqlens: torch.LongTensor | None = None, cu_seqlens_cpu: torch.LongTensor | None = None, - **kwargs, ) -> tuple[torch.Tensor, torch.Tensor]: r""" Args: @@ -686,18 +703,18 @@ def parallel_simple_gla( keys of shape `[B, T, H, K]`. v (torch.Tensor): values of shape `[B, T, H, V]`. - g (Optional[torch.Tensor]): + g (torch.Tensor, Optional): Forget gates of shape `[B, T, H]`. Compared to GLA, the gating is head-wise instead of elementwise. Default: `None`. - scale (Optional[float]): + scale (float, Optional): Scale factor for attention scores. If not provided, it will default to `1 / sqrt(K)`. Default: `None`. - output_attentions (bool): + output_attentions (bool, Optional): Whether to output the materialized attention scores of shape `[B, H, T, T]`. Default: `False`. - cu_seqlens (torch.LongTensor): + cu_seqlens (torch.LongTensor, Optional): Cumulative sequence lengths of shape `[N+1]` used for variable-length training, - consistent with the FlashAttention API. - cu_seqlens_cpu (torch.LongTensor): + consistent with the FlashAttention API. Default: `None`. + cu_seqlens_cpu (torch.LongTensor, Optional): CPU copy of `cu_seqlens` to avoid unnecessary device synchronization. Default: `None`. Returns: @@ -706,10 +723,6 @@ def parallel_simple_gla( attn (torch.Tensor): Attention scores of shape `[B, H, T, T]` if `output_attentions=True` else `None`. """ - if 'head_first' in kwargs: - raise DeprecationWarning( - "head_first has been removed. Inputs must be in `[B, T, H, ...]` format.", - ) if cu_seqlens is not None: if q.shape[0] != 1: raise ValueError( @@ -721,14 +734,5 @@ def parallel_simple_gla( if scale is None: scale = k.shape[-1] ** -0.5 - o, attn = ParallelSimpleGLAFunction.apply( - q, - k, - v, - g, - scale, - output_attentions, - cu_seqlens, - cu_seqlens_cpu, - ) + o, attn = ParallelSimpleGLAFunction.apply(q, k, v, g, scale, output_attentions, cu_seqlens, cu_seqlens_cpu) return o, attn diff --git a/build/torch-cuda/ops/stickbreaking_attn/__init__.py b/build/torch-cuda/ops/stickbreaking_attn/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..9318b7b1baa9a418b89553792ecc6b5e2f5e416f --- /dev/null +++ b/build/torch-cuda/ops/stickbreaking_attn/__init__.py @@ -0,0 +1,16 @@ +# 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 +# +# reference: Tan et al., https://github.com/shawntan/stickbreaking-attention + +from .naive import naive_stickbreaking_attn +from .parallel import parallel_stickbreaking_attn + +__all__ = [ + 'naive_stickbreaking_attn', + 'parallel_stickbreaking_attn', +] diff --git a/build/torch-cuda/ops/stickbreaking_attn/naive.py b/build/torch-cuda/ops/stickbreaking_attn/naive.py new file mode 100644 index 0000000000000000000000000000000000000000..9ab62016e56d30eb71b16826358088505f649ce9 --- /dev/null +++ b/build/torch-cuda/ops/stickbreaking_attn/naive.py @@ -0,0 +1,75 @@ +# 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 torch.nn.functional as F + + +def naive_stickbreaking_attn( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + scale: float | None = None, + attend_current: bool = False, +) -> tuple[torch.Tensor, torch.Tensor]: + r""" + Reference stick-breaking attention, computed in log space in fp32, or in fp64 for fp64 inputs. + + Query `i` gives each visible key `j` the weight `A_ij = sigmoid(z_ij) * prod_l (1 - sigmoid(z_il))`, + where `z_ij = scale * q_i k_j` and `l` runs over the visible keys after `j`, + that is `j < l < i`, or `j < l <= i` with `attend_current`. + The nearest key breaks the stick first, and the weights of a query sum to at most 1. + + Args: + q (torch.Tensor): + Queries of shape `[B, T, HQ, K]`. + k (torch.Tensor): + Keys of shape `[B, T, H, K]`. GQA is applied if HQ is divisible by H. + v (torch.Tensor): + Values of shape `[B, T, H, V]`. + scale (float, Optional): + Scale factor for the attention logits. Default: `1 / sqrt(K)`. + attend_current (bool, Optional): + Whether a query also attends to its own position. Default: `False`. + + Returns: + o (torch.Tensor): + Outputs of shape `[B, T, HQ, V]`. + rem (torch.Tensor): + Stick left over by each query, `1 - sum_j A_ij`, of shape `[B, T, HQ]`. + """ + if q.shape[:2] != k.shape[:2] or k.shape[:3] != v.shape[:3] or q.shape[-1] != k.shape[-1]: + raise ValueError( + f"Expected q, k and v of shapes [B, T, HQ, K], [B, T, H, K] and [B, T, H, V], " + f"got {tuple(q.shape)}, {tuple(k.shape)} and {tuple(v.shape)}. " + f"Decoding against a longer key/value prefix is not supported.", + ) + B, T, HQ, K = q.shape + H, V = k.shape[2], v.shape[-1] + if H == 0 or HQ % H != 0: + raise ValueError(f"The number of query heads ({HQ}) must be divisible by the number of key/value heads ({H}).") + G = HQ // H + if scale is None: + scale = K ** -0.5 + + dtype = q.dtype + # fp64 inputs stay in fp64 so the reference stays accurate where fp32 cancels near saturated logits + compute_dtype = torch.double if torch.double in (q.dtype, k.dtype, v.dtype) else torch.float + q, k, v = q.to(compute_dtype), k.to(compute_dtype), v.to(compute_dtype) + # [B, H, G, T, T] + z = torch.einsum('bqhgd,bkhd->bhgqk', q.reshape(B, T, H, G, K), k) * scale + + i = torch.arange(T, device=q.device) + mask = i[None, :] <= i[:, None] if attend_current else i[None, :] < i[:, None] + # masked keys must leave the remaining stick unchanged + log_om_beta = F.logsigmoid(-z).masked_fill(~mask, 0.) + log_att = F.logsigmoid(z) + log_om_beta.flip(-1).cumsum(-1).flip(-1) - log_om_beta + att = log_att.masked_fill(~mask, float('-inf')).exp() + + o = torch.einsum('bhgqk,bkhd->bqhgd', att, v).reshape(B, T, HQ, V) + rem = log_om_beta.sum(-1).exp().permute(0, 3, 1, 2).reshape(B, T, HQ) + return o.to(dtype), rem.to(dtype) diff --git a/build/torch-cuda/ops/stickbreaking_attn/parallel.py b/build/torch-cuda/ops/stickbreaking_attn/parallel.py new file mode 100644 index 0000000000000000000000000000000000000000..0e32f4fbd4848bcb222f6abc2ce577be16f9f5a8 --- /dev/null +++ b/build/torch-cuda/ops/stickbreaking_attn/parallel.py @@ -0,0 +1,623 @@ +# 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 +import triton.language as tl + +from ...ops.utils import get_max_num_splits, prepare_chunk_indices +from ...ops.utils.op import exp2 +from ...ops.utils.softplus import softplus2 +from ...utils import ( + autocast_custom_bwd, + autocast_custom_fwd, + autotune_cache_kwargs, + check_shared_mem, + input_guard, + tensor_cache, +) + + +@tensor_cache +def _prepare_heavy_first_chunk_indices(cu_seqlens: torch.LongTensor, chunk_size: int) -> torch.LongTensor: + # later blocks of a sequence walk more keys, so the blocks of all sequences go in descending block index + chunk_indices = prepare_chunk_indices(cu_seqlens=cu_seqlens, chunk_size=chunk_size) + return chunk_indices[torch.argsort(chunk_indices[:, 1], descending=True, stable=True)] + + +@triton.heuristics({'IS_VARLEN': lambda args: args['cu_seqlens'] is not None}) +@triton.autotune( + configs=[ + triton.Config({}, num_warps=num_warps, num_stages=num_stages) + for num_warps in [2, 4, 8] + for num_stages in [2, 3, 4] + ], + key=['H', 'HQ', 'G', 'K', 'V', 'BT', 'BS', 'BK', 'BV', 'ATTEND_CURRENT'], + **autotune_cache_kwargs, +) +@triton.jit(do_not_specialize=['T']) +def parallel_stickbreaking_attn_fwd_kernel( + q, + k, + v, + o, + rem, + scale, + cu_seqlens, + chunk_indices, + T, + H: tl.constexpr, + HQ: tl.constexpr, + G: tl.constexpr, + K: tl.constexpr, + V: tl.constexpr, + BT: tl.constexpr, + BS: tl.constexpr, + BK: tl.constexpr, + BV: tl.constexpr, + ATTEND_CURRENT: tl.constexpr, + IS_VARLEN: tl.constexpr, +): + i_v, i_bh, i_t = tl.program_id(0), tl.program_id(1).to(tl.int64), tl.program_id(2).to(tl.int64) + i_b, i_hq = i_bh // HQ, i_bh % HQ + i_h = i_hq // G + + 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.int64) + bos, eos = tl.load(cu_seqlens + i_n).to(tl.int64), tl.load(cu_seqlens + i_n + 1).to(tl.int64) + T = (eos - bos).to(tl.int32) + else: + i_n = i_b + bos, eos = (i_n * T).to(tl.int64), (i_n * T + T).to(tl.int64) + # last query blocks first, as in the varlen chunk indices + i_t = tl.cdiv(T, BT) - 1 - i_t + RCP_LN2: tl.constexpr = 1.4426950216 + + # [BT] + o_q = i_t * BT + tl.arange(0, BT) + o_d = tl.arange(0, BK) + o_v = i_v * BV + tl.arange(0, BV) + m_q = o_q < T + p_q = q + (bos * HQ + i_hq) * K + o_q[:, None] * (HQ*K) + o_d[None, :] + p_o = o + (bos * HQ + i_hq) * V + o_q[:, None] * (HQ*V) + o_v[None, :] + p_rem = rem + bos * HQ + i_hq + o_q * HQ + + # [BT, BK] + b_q = tl.load(p_q, mask=m_q[:, None] & (o_d[None, :] < K), other=0.0) + # [BT, BV] + b_o = tl.zeros([BT, BV], dtype=tl.float32) + # [BT] + b_log_rem = tl.zeros([BT], dtype=tl.float32) + + # visit nearer blocks first to carry their remaining stick into earlier blocks + for i_s in range(0, BT, BS): + # [BS] + o_k = i_t * BT + BT - BS - i_s + tl.arange(0, BS) + m_k = o_k < T + p_k = k + (bos * H + i_h) * K + o_d[:, None] + o_k[None, :] * (H*K) + p_v = v + (bos * H + i_h) * V + o_k[:, None] * (H*V) + o_v[None, :] + # [BK, BS] + b_k = tl.load(p_k, mask=(o_d[:, None] < K) & m_k[None, :], other=0.0) + # [BS, BV] + b_v = tl.load(p_v, mask=m_k[:, None] & (o_v[None, :] < V), other=0.0) + + if ATTEND_CURRENT: + m_s = (o_q[:, None] >= o_k[None, :]) & m_k[None, :] + else: + m_s = (o_q[:, None] > o_k[None, :]) & m_k[None, :] + # [BT, BS] + b_s = tl.dot(b_q, b_k) * scale * RCP_LN2 + # masked keys must leave the remaining stick unchanged + b_log_om_beta = tl.where(m_s, -softplus2(b_s), 0.) + b_p = tl.where(m_s, exp2(b_s + tl.cumsum(b_log_om_beta, axis=1, reverse=True) + b_log_rem[:, None]), 0.) + b_o += tl.dot(b_p.to(b_v.dtype), b_v) + b_log_rem += tl.sum(b_log_om_beta, 1) + + # blocks below the diagonal need no causal mask + for i_s in range(BT, (i_t + 1) * BT, BS): + # [BS] + o_k = (i_t + 1) * BT - BS - i_s + tl.arange(0, BS) + p_k = k + (bos * H + i_h) * K + o_d[:, None] + o_k[None, :] * (H*K) + p_v = v + (bos * H + i_h) * V + o_k[:, None] * (H*V) + o_v[None, :] + # [BK, BS] + b_k = tl.load(p_k, mask=o_d[:, None] < K, other=0.0) + # [BS, BV] + b_v = tl.load(p_v, mask=o_v[None, :] < V, other=0.0) + + # [BT, BS] + b_s = tl.dot(b_q, b_k) * scale * RCP_LN2 + b_log_om_beta = -softplus2(b_s) + b_p = exp2(b_s + tl.cumsum(b_log_om_beta, axis=1, reverse=True) + b_log_rem[:, None]) + b_o += tl.dot(b_p.to(b_v.dtype), b_v) + b_log_rem += tl.sum(b_log_om_beta, 1) + + tl.store(p_o, b_o.to(p_o.dtype.element_ty), mask=m_q[:, None] & (o_v[None, :] < V)) + if i_v == 0: + tl.store(p_rem, exp2(b_log_rem).to(p_rem.dtype.element_ty), mask=m_q) + + +@triton.jit +def _stickbreaking_attn_weights(b_s, b_log_rem, m_s, AXIS: tl.constexpr): + # the scan axis follows keys, which are transposed in the dkv kernel + b_log_om_beta = tl.where(m_s, -softplus2(b_s), 0.) + b_p = tl.where(m_s, exp2(b_s + tl.cumsum(b_log_om_beta, axis=AXIS, reverse=True) + b_log_rem), 0.) + return b_log_om_beta, b_p + + +@triton.jit +def _stickbreaking_attn_exclusive_cumsum_combine(b_s1, b_x1, b_s2, b_x2): + return b_s1 + b_x1 + b_s2, b_x2 + + +@triton.jit +def _stickbreaking_attn_exclusive_cumsum(b_x, AXIS: tl.constexpr): + # pairs of (sum before the last element, last element) give an exclusive sum built only by additions + b_s, _ = tl.associative_scan((tl.zeros_like(b_x), b_x), axis=AXIS, combine_fn=_stickbreaking_attn_exclusive_cumsum_combine) + return b_s + + +@triton.heuristics({'IS_VARLEN': lambda args: args['cu_seqlens'] is not None}) +@triton.autotune( + configs=[ + triton.Config({}, num_warps=num_warps, num_stages=num_stages) + for num_warps in [2, 4, 8] + for num_stages in [1, 2, 3] + ], + key=['H', 'HQ', 'G', 'K', 'V', 'BS', 'BK', 'BV', 'ATTEND_CURRENT'], + **autotune_cache_kwargs, +) +@triton.jit(do_not_specialize=['T', 'NS']) +def parallel_stickbreaking_attn_bwd_kernel_dq( + q, + k, + v, + do, + drem, + dq, + log_rem, + delta_acc, + scale, + cu_seqlens, + chunk_indices, + T, + NS, + H: tl.constexpr, + HQ: tl.constexpr, + G: tl.constexpr, + K: tl.constexpr, + V: tl.constexpr, + BS: tl.constexpr, + BK: tl.constexpr, + BV: tl.constexpr, + ATTEND_CURRENT: tl.constexpr, + IS_VARLEN: tl.constexpr, +): + i_t, i_bh = tl.program_id(0).to(tl.int64), tl.program_id(1).to(tl.int64) + i_b, i_hq = i_bh // HQ, i_bh % HQ + i_h = i_hq // G + + 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.int64) + bos, eos = tl.load(cu_seqlens + i_n).to(tl.int64), tl.load(cu_seqlens + i_n + 1).to(tl.int64) + T = (eos - bos).to(tl.int32) + else: + i_n = i_b + bos, eos = (i_n * T).to(tl.int64), (i_n * T + T).to(tl.int64) + RCP_LN2: tl.constexpr = 1.4426950216 + + # [BS] + o_q = i_t * BS + tl.arange(0, BS) + o_d = tl.arange(0, BK) + o_v = tl.arange(0, BV) + m_q = o_q < T + p_q = q + (bos * HQ + i_hq) * K + o_q[:, None] * (HQ*K) + o_d[None, :] + p_do = do + (bos * HQ + i_hq) * V + o_q[:, None] * (HQ*V) + o_v[None, :] + p_dq = dq + (bos * HQ + i_hq) * K + o_q[:, None] * (HQ*K) + o_d[None, :] + o_snap = (bos + o_q) * HQ * NS + i_hq * NS + + # [BS, BK] + b_q = tl.load(p_q, mask=m_q[:, None] & (o_d[None, :] < K), other=0.0) + # [BS, BV] + b_do = tl.load(p_do, mask=m_q[:, None] & (o_v[None, :] < V), other=0.0) + # [BS] + b_log_rem = tl.zeros([BS], dtype=tl.float32) + + # snapshots let the dkv kernel reconstruct each tile independently + for i_s in range(0, (i_t + 1) * BS, BS): + i_k = i_t - i_s // BS + o_k = i_k * BS + tl.arange(0, BS) + m_k = o_k < T + tl.store(log_rem + o_snap + i_k, b_log_rem, mask=m_q) + p_k = k + (bos * H + i_h) * K + o_d[:, None] + o_k[None, :] * (H*K) + # [BK, BS] + b_k = tl.load(p_k, mask=(o_d[:, None] < K) & m_k[None, :], other=0.0) + + if ATTEND_CURRENT: + m_s = (o_q[:, None] >= o_k[None, :]) & m_k[None, :] + else: + m_s = (o_q[:, None] > o_k[None, :]) & m_k[None, :] + # [BS, BS] + b_s = tl.dot(b_q, b_k) * scale * RCP_LN2 + b_log_rem += tl.sum(tl.where(m_s, -softplus2(b_s), 0.), 1) + + # [BS] + b_delta_acc = tl.load(drem + (bos + o_q) * HQ + i_hq, mask=m_q, other=0.).to(tl.float32) * exp2(b_log_rem) + # the second walk reads log_rem snapshots stored by other threads of this program + tl.debug_barrier() + # [BS, BK] + b_dq = tl.zeros([BS, BK], dtype=tl.float32) + # visit farther blocks first so the gradient through the remaining stick is built by additions + for i_k in range(0, i_t + 1): + o_k = i_k * BS + tl.arange(0, BS) + m_k = o_k < T + tl.store(delta_acc + o_snap + i_k, b_delta_acc, mask=m_q) + b_log_rem = tl.load(log_rem + o_snap + i_k, mask=m_q, other=0.) + p_k = k + (bos * H + i_h) * K + o_d[:, None] + o_k[None, :] * (H*K) + p_v = v + (bos * H + i_h) * V + o_v[:, None] + o_k[None, :] * (H*V) + # [BK, BS] + b_k = tl.load(p_k, mask=(o_d[:, None] < K) & m_k[None, :], other=0.0) + # [BV, BS] + b_v = tl.load(p_v, mask=(o_v[:, None] < V) & m_k[None, :], other=0.0) + + if ATTEND_CURRENT: + m_s = (o_q[:, None] >= o_k[None, :]) & m_k[None, :] + else: + m_s = (o_q[:, None] > o_k[None, :]) & m_k[None, :] + # [BS, BS] + b_s = tl.dot(b_q, b_k) * scale * RCP_LN2 + b_dp = tl.dot(b_do, b_v) + b_log_om_beta, b_p = _stickbreaking_attn_weights(b_s=b_s, b_log_rem=b_log_rem[:, None], m_s=m_s, AXIS=1) + b_pdp = b_p * b_dp + # see README.md for the logit gradient formula + b_delta = b_delta_acc[:, None] + _stickbreaking_attn_exclusive_cumsum(b_x=b_pdp, AXIS=1) + b_ds = tl.where(m_s, exp2(b_log_om_beta) * b_pdp - exp2(b_s + b_log_om_beta) * b_delta, 0.) + b_dq += tl.dot(b_ds.to(b_k.dtype), tl.trans(b_k)) + b_delta_acc += tl.sum(b_pdp, 1) + + tl.store(p_dq, (b_dq * scale).to(p_dq.dtype.element_ty), mask=m_q[:, None] & (o_d[None, :] < K)) + + +@triton.heuristics({'IS_VARLEN': lambda args: args['cu_seqlens'] is not None}) +@triton.autotune( + configs=[ + triton.Config({}, num_warps=num_warps, num_stages=num_stages) + for num_warps in [2, 4, 8] + for num_stages in [1, 2, 3] + ], + key=['H', 'HQ', 'G', 'K', 'V', 'BT', 'BS', 'BK', 'BV', 'ATTEND_CURRENT'], + **autotune_cache_kwargs, +) +@triton.jit(do_not_specialize=['T', 'NS']) +def parallel_stickbreaking_attn_bwd_kernel_dkv( + q, + k, + v, + do, + dk, + dv, + log_rem, + delta_acc, + scale, + cu_seqlens, + chunk_indices, + T, + NS, + H: tl.constexpr, + HQ: tl.constexpr, + G: tl.constexpr, + K: tl.constexpr, + V: tl.constexpr, + BT: tl.constexpr, + BS: tl.constexpr, + BK: tl.constexpr, + BV: tl.constexpr, + ATTEND_CURRENT: tl.constexpr, + IS_VARLEN: tl.constexpr, +): + # each key block has one owner so gradient accumulation needs no atomics + i_s, i_bh = tl.program_id(0).to(tl.int64), tl.program_id(1).to(tl.int64) + i_b, i_h = i_bh // H, i_bh % H + + if IS_VARLEN: + i_n, i_s = tl.load(chunk_indices + i_s * 2).to(tl.int32), tl.load(chunk_indices + i_s * 2 + 1).to(tl.int64) + bos, eos = tl.load(cu_seqlens + i_n).to(tl.int64), tl.load(cu_seqlens + i_n + 1).to(tl.int64) + T = (eos - bos).to(tl.int32) + else: + i_n = i_b + bos, eos = (i_n * T).to(tl.int64), (i_n * T + T).to(tl.int64) + RCP_LN2: tl.constexpr = 1.4426950216 + + # [BS] + o_k = i_s * BS + tl.arange(0, BS) + o_d = tl.arange(0, BK) + o_v = tl.arange(0, BV) + m_k = o_k < T + p_k = k + (bos * H + i_h) * K + o_k[:, None] * (H*K) + o_d[None, :] + p_v = v + (bos * H + i_h) * V + o_k[:, None] * (H*V) + o_v[None, :] + p_dk = dk + (bos * H + i_h) * K + o_k[:, None] * (H*K) + o_d[None, :] + p_dv = dv + (bos * H + i_h) * V + o_k[:, None] * (H*V) + o_v[None, :] + + # [BS, BK] + b_k = tl.load(p_k, mask=m_k[:, None] & (o_d[None, :] < K), other=0.0) + # [BS, BV] + b_v = tl.load(p_v, mask=m_k[:, None] & (o_v[None, :] < V), other=0.0) + b_dk = tl.zeros([BS, BK], dtype=tl.float32) + b_dv = tl.zeros([BS, BV], dtype=tl.float32) + + for i_g in range(G): + i_hq = i_h * G + i_g + for i_t in range(i_s * BS, T, BT): + # [BT] + o_q = i_t + tl.arange(0, BT) + m_q = o_q < T + p_q = q + (bos * HQ + i_hq) * K + o_q[:, None] * (HQ*K) + o_d[None, :] + p_do = do + (bos * HQ + i_hq) * V + o_q[:, None] * (HQ*V) + o_v[None, :] + o_snap = (bos + o_q) * HQ * NS + i_hq * NS + i_s + # [BT, BK] + b_q = tl.load(p_q, mask=m_q[:, None] & (o_d[None, :] < K), other=0.0) + # [BT, BV] + b_do = tl.load(p_do, mask=m_q[:, None] & (o_v[None, :] < V), other=0.0) + # [BT] + b_log_rem = tl.load(log_rem + o_snap, mask=m_q, other=0.0) + b_delta_acc = tl.load(delta_acc + o_snap, mask=m_q, other=0.0) + + if ATTEND_CURRENT: + m_s = (o_k[:, None] <= o_q[None, :]) & m_k[:, None] & m_q[None, :] + else: + m_s = (o_k[:, None] < o_q[None, :]) & m_k[:, None] & m_q[None, :] + # [BS, BT] + b_s = tl.dot(b_k, tl.trans(b_q)) * scale * RCP_LN2 + b_dp = tl.dot(b_v, tl.trans(b_do)) + b_log_om_beta, b_p = _stickbreaking_attn_weights(b_s=b_s, b_log_rem=b_log_rem[None, :], m_s=m_s, AXIS=0) + b_pdp = b_p * b_dp + b_delta = b_delta_acc[None, :] + _stickbreaking_attn_exclusive_cumsum(b_x=b_pdp, AXIS=0) + b_ds = tl.where(m_s, exp2(b_log_om_beta) * b_pdp - exp2(b_s + b_log_om_beta) * b_delta, 0.) + # [BS, BV] + b_dv = tl.dot(b_p.to(b_do.dtype), b_do, b_dv) + # [BS, BK] + b_dk = tl.dot(b_ds.to(b_q.dtype), b_q, b_dk) + + tl.store(p_dk, (b_dk * scale).to(p_dk.dtype.element_ty), mask=m_k[:, None] & (o_d[None, :] < K)) + tl.store(p_dv, b_dv.to(p_dv.dtype.element_ty), mask=m_k[:, None] & (o_v[None, :] < V)) + + +def parallel_stickbreaking_attn_fwd( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + scale: float, + attend_current: bool = False, + cu_seqlens: torch.LongTensor | None = None, +) -> tuple[torch.Tensor, torch.Tensor]: + B, T, H, K, V = *k.shape, v.shape[-1] + HQ = q.shape[2] + G = HQ // H + BT = 64 + BS = 32 + BK = max(16, triton.next_power_of_2(K)) + BV = min(128, max(16, triton.next_power_of_2(V))) + NV = triton.cdiv(V, BV) + assert BT % BS == 0 + + chunk_indices = None + if cu_seqlens is not None: + chunk_indices = _prepare_heavy_first_chunk_indices(cu_seqlens=cu_seqlens, chunk_size=BT) + NT = triton.cdiv(T, BT) if cu_seqlens is None else len(chunk_indices) + + o = torch.empty(B, T, HQ, V, dtype=v.dtype, device=q.device) + rem = torch.empty(B, T, HQ, dtype=q.dtype, device=q.device) + # later query blocks walk more key blocks, so they go on the slowest grid axis and launch first to avoid a heavy tail + grid = (NV, B * HQ, NT) + parallel_stickbreaking_attn_fwd_kernel[grid]( + q=q, + k=k, + v=v, + o=o, + rem=rem, + scale=scale, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + T=T, + H=H, + HQ=HQ, + G=G, + K=K, + V=V, + BT=BT, + BS=BS, + BK=BK, + BV=BV, + ATTEND_CURRENT=attend_current, + ) + return o, rem + + +def parallel_stickbreaking_attn_bwd( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + do: torch.Tensor, + drem: torch.Tensor, + scale: float, + attend_current: bool = False, + cu_seqlens: torch.LongTensor | None = None, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + B, T, H, K, V = *k.shape, v.shape[-1] + HQ = q.shape[2] + G = HQ // H + # the dq query blocks match the snapshot key blocks + BS = 64 if check_shared_mem(arch='hopper', tensor_idx=q.device.index) else 32 + BT = 32 + BK = max(16, triton.next_power_of_2(K)) + BV = max(16, triton.next_power_of_2(V)) + + if cu_seqlens is not None: + chunk_indices = prepare_chunk_indices(cu_seqlens=cu_seqlens, chunk_size=BS) + NT = len(chunk_indices) + NS = get_max_num_splits(cu_seqlens=cu_seqlens, chunk_size=BS) + else: + chunk_indices = None + NT = NS = triton.cdiv(T, BS) + + # snapshots store the remaining stick log left by nearer key blocks and the gradient through it from farther ones + log_rem = q.new_empty(B, T, HQ, NS, dtype=torch.float) + delta_acc = q.new_empty(B, T, HQ, NS, dtype=torch.float) + dq = torch.empty_like(q) + dk = torch.empty_like(k) + dv = torch.empty_like(v) + parallel_stickbreaking_attn_bwd_kernel_dq[(NT, B * HQ)]( + q=q, + k=k, + v=v, + do=do, + drem=drem, + dq=dq, + log_rem=log_rem, + delta_acc=delta_acc, + scale=scale, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + T=T, + NS=NS, + H=H, + HQ=HQ, + G=G, + K=K, + V=V, + BS=BS, + BK=BK, + BV=BV, + ATTEND_CURRENT=attend_current, + ) + parallel_stickbreaking_attn_bwd_kernel_dkv[(NT, B * H)]( + q=q, + k=k, + v=v, + do=do, + dk=dk, + dv=dv, + log_rem=log_rem, + delta_acc=delta_acc, + scale=scale, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + T=T, + NS=NS, + H=H, + HQ=HQ, + G=G, + K=K, + V=V, + BT=BT, + BS=BS, + BK=BK, + BV=BV, + ATTEND_CURRENT=attend_current, + ) + return dq, dk, dv + + +class ParallelStickBreakingAttentionFunction(torch.autograd.Function): + + @staticmethod + @input_guard + @autocast_custom_fwd + def forward(ctx, q, k, v, scale, attend_current, cu_seqlens): + o, rem = parallel_stickbreaking_attn_fwd( + q=q, + k=k, + v=v, + scale=scale, + attend_current=attend_current, + cu_seqlens=cu_seqlens, + ) + ctx.save_for_backward(q, k, v) + ctx.scale = scale + ctx.attend_current = attend_current + ctx.cu_seqlens = cu_seqlens + return o.to(q.dtype), rem + + @staticmethod + @input_guard + @autocast_custom_bwd + def backward(ctx, do, drem): + q, k, v = ctx.saved_tensors + dq, dk, dv = parallel_stickbreaking_attn_bwd( + q=q, + k=k, + v=v, + do=do, + drem=drem, + scale=ctx.scale, + attend_current=ctx.attend_current, + cu_seqlens=ctx.cu_seqlens, + ) + return dq.to(q), dk.to(k), dv.to(v), None, None, None + + +def parallel_stickbreaking_attn( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + scale: float | None = None, + attend_current: bool = False, + cu_seqlens: torch.LongTensor | None = None, +) -> tuple[torch.Tensor, torch.Tensor]: + r""" + Stick-breaking attention (https://arxiv.org/abs/2410.17980). + + Query `i` gives each visible key `j` the weight `A_ij = sigmoid(z_ij) * prod_l (1 - sigmoid(z_il))`, + where `z_ij = scale * q_i k_j` and `l` runs over the visible keys after `j`, + that is `j < l < i`, or `j < l <= i` with `attend_current`. + Walking back from the query, each key takes a sigmoid share of the stick that the nearer keys left over. + The weights are not normalized, so a query can leave part of its stick unused. + + Args: + q (torch.Tensor): + Queries of shape `[B, T, HQ, K]`. + k (torch.Tensor): + Keys of shape `[B, T, H, K]`. GQA is applied if HQ is divisible by H. + v (torch.Tensor): + Values of shape `[B, T, H, V]`. + scale (float, Optional): + Scale factor for attention scores. Default: `1 / sqrt(K)`. + attend_current (bool, Optional): + Whether a query also attends to the key at its own position. + If `False`, query `i` only sees keys `j < i`, and the first query outputs zeros. Default: `False`. + cu_seqlens (torch.LongTensor, Optional): + Cumulative sequence lengths of shape `[N+1]` used for variable-length training, + consistent with the FlashAttention API. Default: `None`. + + Returns: + o (torch.Tensor): + Outputs of shape `[B, T, HQ, V]`. + rem (torch.Tensor): + The stick each query leaves unused, `1 - sum_j A_ij`, of shape `[B, T, HQ]`. + """ + if q.shape[:2] != k.shape[:2] or k.shape[:3] != v.shape[:3] or q.shape[-1] != k.shape[-1]: + raise ValueError( + f"Expected q, k and v of shapes [B, T, HQ, K], [B, T, H, K] and [B, T, H, V], " + f"got {tuple(q.shape)}, {tuple(k.shape)} and {tuple(v.shape)}. " + f"Decoding against a longer key/value prefix is not supported.", + ) + if scale is None: + scale = k.shape[-1] ** -0.5 + HQ, H = q.shape[2], k.shape[2] + if H == 0 or HQ % H != 0: + raise ValueError(f"The number of query heads ({HQ}) must be divisible by the number of key/value heads ({H}).") + if q.shape[-1] > 256: + raise ValueError(f"The key dimension ({q.shape[-1]}) can not be larger than 256.") + if v.shape[-1] > 256: + raise ValueError(f"The value dimension ({v.shape[-1]}) can not be larger than 256.") + if cu_seqlens is not None and q.shape[0] != 1: + raise ValueError( + f"The batch size is expected to be 1 rather than {q.shape[0]} when using `cu_seqlens`. " + f"Please flatten variable-length inputs before processing.", + ) + + o, rem = ParallelStickBreakingAttentionFunction.apply(q, k, v, scale, attend_current, cu_seqlens) + return o, rem diff --git a/build/torch-cuda/ops/titans/naive.py b/build/torch-cuda/ops/titans/naive.py index 926404ac30cdc46c9addb1cb957469020aff451f..3a01f882decf52497fb57da5dcc7937645c1af3e 100644 --- a/build/torch-cuda/ops/titans/naive.py +++ b/build/torch-cuda/ops/titans/naive.py @@ -318,18 +318,16 @@ def chunk_titans_linear_ref( chunk_size: int = 16, # chunk size initial_state: torch.Tensor = None, output_final_state: bool = False, - head_first: bool = False, use_chunk: bool = True, ): assert q.dtype == k.dtype == v.dtype assert k.shape[-1] == v.shape[-1], "DK must equal to DV." - if not head_first: - q = q.transpose(1, 2) - k = k.transpose(1, 2) - v = v.transpose(1, 2) - eta = eta.transpose(1, 2) - alpha = alpha.transpose(1, 2) - theta = theta.transpose(1, 2) + q = q.transpose(1, 2) + k = k.transpose(1, 2) + v = v.transpose(1, 2) + eta = eta.transpose(1, 2) + alpha = alpha.transpose(1, 2) + theta = theta.transpose(1, 2) seq_len = q.shape[-2] pad_len = (chunk_size - (seq_len % chunk_size)) % chunk_size if pad_len > 0: @@ -375,6 +373,5 @@ def chunk_titans_linear_ref( output_final_state, ) o = o[:, :, :seq_len, :] - if not head_first: - o = o.transpose(1, 2) + o = o.transpose(1, 2) return o, final_state diff --git a/build/torch-cuda/ops/ttt/chunk.py b/build/torch-cuda/ops/ttt/chunk.py index 3e5234997035702273550a2ee92f1c72b4655a89..0037bf462b28c69161b7e9c831920b1e0a537a7c 100644 --- a/build/torch-cuda/ops/ttt/chunk.py +++ b/build/torch-cuda/ops/ttt/chunk.py @@ -12,6 +12,7 @@ import triton.language as tl from ...modules.layernorm import group_norm from ...ops.utils import prepare_chunk_indices, prepare_chunk_offsets +from ...ops.utils.op import unflatten_program_id from ...utils import autocast_custom_bwd, autocast_custom_fwd, autotune_cache_kwargs, input_guard @@ -58,13 +59,13 @@ def chunk_ttt_linear_fwd_kernel_h( STORE_FINAL_STATE: tl.constexpr, IS_VARLEN: tl.constexpr, ): - i_k, i_v, i_nh = tl.program_id(0), tl.program_id(1), tl.program_id(2) + i_k, i_v, i_nh = unflatten_program_id(tl.cdiv(K, BK), tl.cdiv(V, BV)) i_n, i_h = i_nh // H, i_nh % H if IS_VARLEN: - bos, eos = tl.load(cu_seqlens + i_n).to(tl.int32), tl.load(cu_seqlens + i_n + 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 NT = tl.cdiv(T, BT) - boh = tl.load(chunk_offsets + i_n).to(tl.int32) + boh = tl.load(chunk_offsets + i_n).to(tl.int64) else: bos, eos = i_n * T, i_n * T + T NT = tl.cdiv(T, BT) @@ -74,28 +75,35 @@ def chunk_ttt_linear_fwd_kernel_h( b_h = tl.zeros([BK, BV], dtype=tl.float32) # [BV] b_hb = tl.zeros([BV], dtype=tl.float32) + o_k = i_k * BK + tl.arange(0, BK) + o_v = i_v * BV + tl.arange(0, BV) + m_h = (o_k[:, None] < K) & (o_v[None, :] < V) if USE_INITIAL_STATE: - p_h0 = tl.make_block_ptr(h0 + i_nh * K * V, (K, V), (V, 1), (i_k * BK, i_v * BV), (BK, BV), (1, 0)) - b_h = tl.load(p_h0, boundary_check=(0, 1), padding_option="zero").to(tl.float32) + p_h0 = h0 + i_nh * K * V + o_k[:, None] * V + o_v[None, :] + b_h = tl.load(p_h0, mask=m_h, other=0.0).to(tl.float32) if USE_INITIAL_STATE_B: - p_hb0 = tl.make_block_ptr(hb0 + i_nh * V, (V,), (1,), (i_v * BV,), (BV,), (0,)) - b_hb = tl.load(p_hb0, boundary_check=(0,), padding_option="zero").to(tl.float32) + p_hb0 = hb0 + i_nh * V + o_v + b_hb = tl.load(p_hb0, mask=o_v < V, other=0.0).to(tl.float32) offs = tl.arange(0, BV) b_w = tl.load(w + i_h * V + offs, mask=offs < V, other=0.) b_b = tl.load(b + i_h * V + offs, mask=offs < V, other=0.) for i_t in range(NT): - p_h = tl.make_block_ptr(h + ((boh + i_t) * H + i_h) * K*V, (K, V), (V, 1), (i_k * BK, i_v * BV), (BK, BV), (1, 0)) - p_hb = tl.make_block_ptr(hb + ((boh + i_t) * H + i_h) * V, (V,), (1,), (i_v * BV,), (BV,), (0,)) - tl.store(p_h, b_h.to(p_h.dtype.element_ty), boundary_check=(0, 1)) - tl.store(p_hb, b_hb.to(p_hb.dtype.element_ty), boundary_check=(0,)) - p_k = tl.make_block_ptr(k+(bos*H+i_h)*K, (K, T), (1, H*K), (i_k * BK, i_t * BT), (BK, BT), (0, 1)) - p_v = tl.make_block_ptr(v+(bos*H+i_h)*V, (T, V), (H*V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) - p_v_new = tl.make_block_ptr(v_new+(bos*H+i_h)*V, (T, V), (H*V, 1), (i_t*BT, i_v * BV), (BT, BV), (1, 0)) - p_eta_last = eta+bos*H+i_h + (T-1)*H if i_t == NT-1 else eta+bos*H+i_h + (i_t*BT+BT-1)*H - b_k = tl.load(p_k, boundary_check=(0, 1), padding_option="zero") - b_v = tl.load(p_v, boundary_check=(0, 1), padding_option="zero") + i_t_int64 = i_t.to(tl.int64) + o_t = i_t_int64 * BT + tl.arange(0, BT) + m_v = (o_t[:, None] < T) & (o_v[None, :] < V) + m_k = (o_k[:, None] < K) & (o_t[None, :] < T) + p_h = h + ((boh + i_t_int64) * H + i_h) * K*V + o_k[:, None] * V + o_v[None, :] + p_hb = hb + ((boh + i_t_int64) * H + i_h) * V + o_v + tl.store(p_h, b_h.to(p_h.dtype.element_ty), mask=m_h) + tl.store(p_hb, b_hb.to(p_hb.dtype.element_ty), mask=o_v < V) + p_k = k+(bos*H+i_h)*K + o_k[:, None] + o_t[None, :] * (H*K) + p_v = v+(bos*H+i_h)*V + o_t[:, None] * (H*V) + o_v[None, :] + p_v_new = v_new+(bos*H+i_h)*V + o_t[:, None] * (H*V) + o_v[None, :] + p_eta_last = eta+bos*H+i_h + (T-1)*H if i_t_int64 == NT-1 else eta+bos*H+i_h + (i_t_int64*BT+BT-1)*H + b_k = tl.load(p_k, mask=m_k, other=0.0) + b_v = tl.load(p_v, mask=m_v, other=0.0) b_kh = tl.dot(tl.trans(b_k), b_h.to(b_k.dtype), allow_tf32=False).to(tl.float32) + b_hb[None, :] b_kh = tl.where((offs < V)[None, :], b_kh, 0.) @@ -110,16 +118,16 @@ def chunk_ttt_linear_fwd_kernel_h( b_v = tl.where((offs < V)[None, :], b_v * b_w[None, :].to(b_k.dtype), 0.) b_v2 = rstd * (V * b_v - tl.sum(b_v, axis=1, keep_dims=True) - b_kh_hat.to(b_k.dtype) * tl.sum(b_v * b_kh_hat.to(b_k.dtype), axis=1, keep_dims=True)) / V - tl.store(p_v_new, b_v2.to(p_v_new.dtype.element_ty), boundary_check=(0, 1)) + tl.store(p_v_new, b_v2.to(p_v_new.dtype.element_ty), mask=m_v) b_eta_last = tl.load(p_eta_last) b_h = b_h - tl.dot(b_eta_last * b_k, b_v2.to(b_k.dtype), allow_tf32=False) b_hb = b_hb - tl.sum(b_eta_last * b_v2.to(b_k.dtype), axis=0) if STORE_FINAL_STATE: - p_ht = tl.make_block_ptr(ht + i_nh * K*V, (K, V), (V, 1), (i_k * BK, i_v * BV), (BK, BV), (1, 0)) - p_hbt = tl.make_block_ptr(hbt + i_nh * V, (V,), (1,), (i_v * BV,), (BV,), (0,)) - tl.store(p_ht, b_h.to(p_ht.dtype.element_ty), boundary_check=(0, 1)) - tl.store(p_hbt, b_hb.to(p_hbt.dtype.element_ty), boundary_check=(0,)) + p_ht = ht + i_nh * K*V + o_k[:, None] * V + o_v[None, :] + p_hbt = hbt + i_nh * V + o_v + tl.store(p_ht, b_h.to(p_ht.dtype.element_ty), mask=m_h) + tl.store(p_hbt, b_hb.to(p_hbt.dtype.element_ty), mask=o_v < V) @triton.heuristics({ @@ -155,13 +163,14 @@ def chunk_ttt_linear_fwd_kernel_o( BV: tl.constexpr, IS_VARLEN: tl.constexpr, ): - i_v, i_t, i_bh = tl.program_id(0), tl.program_id(1), tl.program_id(2) + i_v, i_t = unflatten_program_id(tl.cdiv(V, BV)) + i_bh = tl.program_id(1).to(tl.int64) i_b, i_h = i_bh // H, i_bh % H 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.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) + i_n, i_t = tl.load(chunk_indices + i_t * 2).to(tl.int32), tl.load(chunk_indices + i_t * 2 + 1).to(tl.int64) + bos, eos = tl.load(cu_seqlens + i_n).to(tl.int64), tl.load(cu_seqlens + i_n + 1).to(tl.int64) T = eos - bos NT = tl.cdiv(T, BT) else: @@ -181,21 +190,29 @@ def chunk_ttt_linear_fwd_kernel_o( stride_vo = H*V stride_eta = H - p_q = tl.make_block_ptr(q, (T, K), (stride_qk, 1), (i_t * BT, 0), (BT, BK), (1, 0)) - p_k = tl.make_block_ptr(k, (K, T), (1, stride_qk), (0, i_t * BT), (BK, BT), (0, 1)) - p_eta = tl.make_block_ptr(eta, (T,), (stride_eta,), (i_t * BT,), (BT,), (0,)) - p_h = tl.make_block_ptr(h, (K, V), (V, 1), (0, i_v * BV), (BK, BV), (1, 0)) - p_hb = tl.make_block_ptr(hb, (V,), (1,), (i_v * BV,), (BV,), (0,)) + o_t = i_t * BT + tl.arange(0, BT) + o_k = tl.arange(0, BK) + o_v = i_v * BV + tl.arange(0, BV) + m_t = o_t < T + m_q = m_t[:, None] & (o_k[None, :] < K) + m_k = (o_k[:, None] < K) & m_t[None, :] + m_h = (o_k[:, None] < K) & (o_v[None, :] < V) + m_v = m_t[:, None] & (o_v[None, :] < V) + p_q = q + o_t[:, None] * stride_qk + o_k[None, :] + p_k = k + o_k[:, None] + o_t[None, :] * stride_qk + p_eta = eta + o_t * stride_eta + p_h = h + o_k[:, None] * V + o_v[None, :] + p_hb = hb + o_v # [BT, BK] - b_q = tl.load(p_q, boundary_check=(0, 1), padding_option="zero") + b_q = tl.load(p_q, mask=m_q, other=0.0) # [BK, BT] - b_k = tl.load(p_k, boundary_check=(0, 1), padding_option="zero") + b_k = tl.load(p_k, mask=m_k, other=0.0) # [BT, 1] - b_eta = tl.load(p_eta, boundary_check=(0,), padding_option="zero") + b_eta = tl.load(p_eta, mask=m_t, other=0.0) # [BK, BV] - b_h = tl.load(p_h, boundary_check=(0, 1), padding_option="zero") + b_h = tl.load(p_h, mask=m_h, other=0.0) # [BV] - b_hb = tl.load(p_hb, boundary_check=(0,), padding_option="zero") + b_hb = tl.load(p_hb, mask=o_v < V, other=0.0) # [BT, BK] @ [BK, BV] -> [BT, BV] b_o = tl.dot(b_q, b_h, allow_tf32=False) # [BT, BK] @ [BK, BT] -> [BT, BT] @@ -206,12 +223,12 @@ def chunk_ttt_linear_fwd_kernel_o( b_A = tl.where(m_A, b_A, 0) b_Ae = tl.where(m_A, b_eta[:, None], 0.0) - p_v = tl.make_block_ptr(v, (T, V), (stride_vo, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) - p_o = tl.make_block_ptr(o, (T, V), (stride_vo, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) - b_v = tl.load(p_v, boundary_check=(0, 1), padding_option="zero") + p_v = v + o_t[:, None] * stride_vo + o_v[None, :] + p_o = o + o_t[:, None] * stride_vo + o_v[None, :] + b_v = tl.load(p_v, mask=m_v, other=0.0) b_o = (b_o - tl.dot(b_eta[:, None] * b_A.to(b_v.dtype), b_v, allow_tf32=False)) * scale b_o += b_hb[None, :] - tl.dot(b_Ae.to(b_v.dtype), b_v, allow_tf32=False) - tl.store(p_o, b_o.to(p_o.dtype.element_ty), boundary_check=(0, 1)) + tl.store(p_o, b_o.to(p_o.dtype.element_ty), mask=m_v) @triton.heuristics({ @@ -256,13 +273,13 @@ def chunk_ttt_linear_bwd_kernel_h( USE_INITIAL_STATE_B: tl.constexpr, IS_VARLEN: tl.constexpr, ): - i_k, i_v, i_nh = tl.program_id(0), tl.program_id(1), tl.program_id(2) + i_k, i_v, i_nh = unflatten_program_id(tl.cdiv(K, BK), tl.cdiv(V, BV)) i_n, i_h = i_nh // H, i_nh % H if IS_VARLEN: - bos, eos = tl.load(cu_seqlens + i_n).to(tl.int32), tl.load(cu_seqlens + i_n + 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 NT = tl.cdiv(T, BT) - boh = tl.load(chunk_offsets + i_n).to(tl.int32) + boh = tl.load(chunk_offsets + i_n).to(tl.int64) else: bos, eos = i_n * T, i_n * T + T NT = tl.cdiv(T, BT) @@ -272,29 +289,36 @@ def chunk_ttt_linear_bwd_kernel_h( b_h = tl.zeros([BK, BV], dtype=tl.float32) # [BV] b_hb = tl.zeros([BV], dtype=tl.float32) + o_k = i_k * BK + tl.arange(0, BK) + o_v = i_v * BV + tl.arange(0, BV) + m_h = (o_k[:, None] < K) & (o_v[None, :] < V) if USE_INITIAL_STATE: - p_h0 = tl.make_block_ptr(h0 + i_nh * K * V, (K, V), (V, 1), (i_k * BK, i_v * BV), (BK, BV), (1, 0)) - b_h = tl.load(p_h0, boundary_check=(0, 1), padding_option="zero").to(tl.float32) + p_h0 = h0 + i_nh * K * V + o_k[:, None] * V + o_v[None, :] + b_h = tl.load(p_h0, mask=m_h, other=0.0).to(tl.float32) if USE_INITIAL_STATE_B: - p_hb0 = tl.make_block_ptr(hb0 + i_nh * V, (V,), (1,), (i_v * BV,), (BV,), (0,)) - b_hb = tl.load(p_hb0, boundary_check=(0,), padding_option="zero").to(tl.float32) + p_hb0 = hb0 + i_nh * V + o_v + b_hb = tl.load(p_hb0, mask=o_v < V, other=0.0).to(tl.float32) offs = tl.arange(0, BV) b_w = tl.load(w + i_h * V + offs, mask=offs < V, other=0.) b_b = tl.load(b + i_h * V + offs, mask=offs < V, other=0.) for i_t in range(NT): - p_h = tl.make_block_ptr(h + ((boh + i_t) * H + i_h) * K*V, (K, V), (V, 1), (i_k * BK, i_v * BV), (BK, BV), (1, 0)) - tl.store(p_h, b_h.to(p_h.dtype.element_ty), boundary_check=(0, 1)) - p_k = tl.make_block_ptr(k+(bos*H+i_h)*K, (K, T), (1, H*K), (i_k * BK, i_t * BT), (BK, BT), (0, 1)) - p_v = tl.make_block_ptr(v+(bos*H+i_h)*V, (T, V), (H*V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) - p_v_new = tl.make_block_ptr(v_new+(bos*H+i_h)*V, (T, V), (H*V, 1), (i_t*BT, i_v * BV), (BT, BV), (1, 0)) - p_x = tl.make_block_ptr(x+(bos*H+i_h)*V, (T, V), (H*V, 1), (i_t*BT, i_v * BV), (BT, BV), (1, 0)) - p_y = tl.make_block_ptr(y+(bos*H+i_h)*V, (T, V), (H*V, 1), (i_t*BT, i_v * BV), (BT, BV), (1, 0)) - p_r = tl.make_block_ptr(r+bos*H+i_h, (T, 1), (H, 1), (i_t*BT, 0), (BT, 1), (1, 0)) - p_eta_last = eta+bos*H+i_h + (T-1)*H if i_t == NT-1 else eta+bos*H+i_h + (i_t*BT+BT-1)*H - b_k = tl.load(p_k, boundary_check=(0, 1), padding_option="zero") - b_v = tl.load(p_v, boundary_check=(0, 1), padding_option="zero") + i_t_int64 = i_t.to(tl.int64) + o_t = i_t_int64 * BT + tl.arange(0, BT) + m_v = (o_t[:, None] < T) & (o_v[None, :] < V) + m_k = (o_k[:, None] < K) & (o_t[None, :] < T) + p_h = h + ((boh + i_t_int64) * H + i_h) * K*V + o_k[:, None] * V + o_v[None, :] + tl.store(p_h, b_h.to(p_h.dtype.element_ty), mask=m_h) + p_k = k+(bos*H+i_h)*K + o_k[:, None] + o_t[None, :] * (H*K) + p_v = v+(bos*H+i_h)*V + o_t[:, None] * (H*V) + o_v[None, :] + p_v_new = v_new+(bos*H+i_h)*V + o_t[:, None] * (H*V) + o_v[None, :] + p_x = x+(bos*H+i_h)*V + o_t[:, None] * (H*V) + o_v[None, :] + p_y = y+(bos*H+i_h)*V + o_t[:, None] * (H*V) + o_v[None, :] + p_r = r+bos*H+i_h + o_t[:, None] * H + p_eta_last = eta+bos*H+i_h + (T-1)*H if i_t_int64 == NT-1 else eta+bos*H+i_h + (i_t_int64*BT+BT-1)*H + b_k = tl.load(p_k, mask=m_k, other=0.0) + b_v = tl.load(p_v, mask=m_v, other=0.0) b_kh = tl.dot(tl.trans(b_k), b_h.to(b_k.dtype), allow_tf32=False).to(tl.float32) + b_hb[None, :] b_kh = tl.where((offs < V)[None, :], b_kh, 0.) @@ -309,10 +333,10 @@ def chunk_ttt_linear_bwd_kernel_h( b_v = tl.where((offs < V)[None, :], b_v * b_w[None, :].to(b_k.dtype), 0.) b_v2 = rstd * (V * b_v - tl.sum(b_v, axis=1, keep_dims=True) - b_kh_hat.to(b_k.dtype) * tl.sum(b_v * b_kh_hat.to(b_k.dtype), axis=1, keep_dims=True)) / V - tl.store(p_x, b_kh_hat.to(p_x.dtype.element_ty), boundary_check=(0, 1)) - tl.store(p_y, b_v.to(p_y.dtype.element_ty), boundary_check=(0, 1)) - tl.store(p_r, rstd.to(p_r.dtype.element_ty), boundary_check=(0, 1)) - tl.store(p_v_new, b_v2.to(p_v_new.dtype.element_ty), boundary_check=(0, 1)) + tl.store(p_x, b_kh_hat.to(p_x.dtype.element_ty), mask=m_v) + tl.store(p_y, b_v.to(p_y.dtype.element_ty), mask=m_v) + tl.store(p_r, rstd.to(p_r.dtype.element_ty), mask=o_t[:, None] < T) + tl.store(p_v_new, b_v2.to(p_v_new.dtype.element_ty), mask=m_v) b_eta_last = tl.load(p_eta_last) b_h = b_h - tl.dot(b_eta_last * b_k, b_v2.to(b_k.dtype), allow_tf32=False) b_hb = b_hb - tl.sum(b_eta_last * b_v2.to(b_k.dtype), axis=0) @@ -348,11 +372,11 @@ def chunk_ttt_linear_bwd_kernel_dv_local( BV: tl.constexpr, IS_VARLEN: tl.constexpr, ): - i_t, i_bh = tl.program_id(0), tl.program_id(1) + i_t, i_bh = tl.program_id(0).to(tl.int64), tl.program_id(1).to(tl.int64) i_b, i_h = i_bh // H, i_bh % H 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.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) + i_n, i_t = tl.load(chunk_indices + i_t * 2).to(tl.int32), tl.load(chunk_indices + i_t * 2 + 1).to(tl.int64) + bos, eos = tl.load(cu_seqlens + i_n).to(tl.int64), tl.load(cu_seqlens + i_n + 1).to(tl.int64) T = eos - bos else: bos, eos = i_b * T, i_b * T + T @@ -368,25 +392,32 @@ def chunk_ttt_linear_bwd_kernel_dv_local( stride_eta = H b_A = tl.zeros([BT, BT], dtype=tl.float32) + o_t = i_t * BT + tl.arange(0, BT) + m_t = o_t < T for i_k in range(tl.cdiv(K, BK)): - p_k = tl.make_block_ptr(k, (T, K), (stride_qk, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) - p_q = tl.make_block_ptr(q, (K, T), (1, stride_qk), (i_k * BK, i_t * BT), (BK, BT), (0, 1)) - b_q = tl.load(p_q, boundary_check=(0, 1)) - b_k = tl.load(p_k, boundary_check=(0, 1)) - b_A += tl.dot(b_k, b_q) - - p_eta = tl.make_block_ptr(eta, (T,), (stride_eta,), (i_t * BT,), (BT,), (0,)) - b_eta = tl.load(p_eta, boundary_check=(0,)) + o_k = i_k * BK + tl.arange(0, BK) + m_k = m_t[:, None] & (o_k[None, :] < K) + m_q = (o_k[:, None] < K) & m_t[None, :] + p_k = k + o_t[:, None] * stride_qk + o_k[None, :] + p_q = q + o_k[:, None] + o_t[None, :] * stride_qk + b_q = tl.load(p_q, mask=m_q, other=0.0) + b_k = tl.load(p_k, mask=m_k, other=0.0) + b_A = tl.dot(b_k, b_q, b_A) + + p_eta = eta + o_t * stride_eta + b_eta = tl.load(p_eta, mask=m_t, other=0.0) mask = (tl.arange(0, BT)[:, None] <= tl.arange(0, BT)[None, :]) b_A = - tl.where(mask, b_A * scale * b_eta[None, :], 0).to(do.dtype.element_ty) b_Ae = - tl.where(mask, b_eta[None, :], 0).to(do.dtype.element_ty) for i_v in range(tl.cdiv(V, BV)): - p_do = tl.make_block_ptr(do, (T, V), (stride_vo, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) - p_dv = tl.make_block_ptr(dv, (T, V), (stride_vo, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) - b_do = tl.load(p_do, boundary_check=(0, 1)) + o_v = i_v * BV + tl.arange(0, BV) + m_v = m_t[:, None] & (o_v[None, :] < V) + p_do = do + o_t[:, None] * stride_vo + o_v[None, :] + p_dv = dv + o_t[:, None] * stride_vo + o_v[None, :] + b_do = tl.load(p_do, mask=m_v, other=0.0) b_dv = tl.dot(b_A.to(b_do.dtype), b_do) + tl.dot(b_Ae.to(b_do.dtype), b_do) - tl.store(p_dv, b_dv.to(p_dv.dtype.element_ty), boundary_check=(0, 1)) + tl.store(p_dv, b_dv.to(p_dv.dtype.element_ty), mask=m_v) @triton.heuristics({ @@ -445,13 +476,13 @@ def chunk_ttt_linear_bwd_kernel_norm( USE_INITIAL_STATE_B: tl.constexpr, IS_VARLEN: tl.constexpr, ): - i_k, i_v, i_nh = tl.program_id(0), tl.program_id(1), tl.program_id(2) + i_k, i_v, i_nh = unflatten_program_id(tl.cdiv(K, BK), tl.cdiv(V, BV)) i_n, i_h = i_nh // H, i_nh % H if IS_VARLEN: - bos, eos = tl.load(cu_seqlens + i_n).to(tl.int32), tl.load(cu_seqlens + i_n + 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 NT = tl.cdiv(T, BT) - boh = tl.load(chunk_offsets + i_n).to(tl.int32) + boh = tl.load(chunk_offsets + i_n).to(tl.int64) else: bos, eos = i_n * T, i_n * T + T NT = tl.cdiv(T, BT) @@ -461,12 +492,15 @@ def chunk_ttt_linear_bwd_kernel_norm( b_dh = tl.zeros([BK, BV], dtype=tl.float32) # [BV] b_dhb = tl.zeros([BV], dtype=tl.float32) + o_k = i_k * BK + tl.arange(0, BK) + o_v = i_v * BV + tl.arange(0, BV) + m_h = (o_k[:, None] < K) & (o_v[None, :] < V) if USE_FINAL_STATE_GRADIENT: - p_dht = tl.make_block_ptr(dht + i_nh * K*V, (K, V), (V, 1), (i_k * BK, i_v * BV), (BK, BV), (1, 0)) - b_dh += tl.load(p_dht, boundary_check=(0, 1), padding_option="zero") + p_dht = dht + i_nh * K*V + o_k[:, None] * V + o_v[None, :] + b_dh += tl.load(p_dht, mask=m_h, other=0.0) if USE_FINAL_STATE_GRADIENT_B: - p_dhbt = tl.make_block_ptr(dhbt + i_nh * V, (V,), (1,), (i_v * BV,), (BV,), (0,)) - b_dhb += tl.load(p_dhbt, boundary_check=(0,), padding_option="zero") + p_dhbt = dhbt + i_nh * V + o_v + b_dhb += tl.load(p_dhbt, mask=o_v < V, other=0.0) # [BV] offs_v = tl.arange(0, BV) @@ -475,44 +509,51 @@ def chunk_ttt_linear_bwd_kernel_norm( b_b = tl.load(b + i_h * V + offs_v, mask=offs_v < V, other=0.) b_dw = tl.zeros([BV], dtype=b_w.dtype) b_db = tl.zeros([BV], dtype=b_b.dtype) - p_dw = tl.make_block_ptr(dw + i_nh * V, (V,), (1,), (i_v * BV,), (BV,), (0,)) - p_db = tl.make_block_ptr(db + i_nh * V, (V,), (1,), (i_v * BV,), (BV,), (0,)) + p_dw = dw + i_nh * V + o_v + p_db = db + i_nh * V + o_v for i_t in range(NT - 1, -1, -1): - p_h = tl.make_block_ptr(h + ((boh+i_t) * H + i_h) * K*V, (V, K), (1, V), (i_v * BV, i_k * BK), (BV, BK), (0, 1)) - p_dh = tl.make_block_ptr(dh + ((boh+i_t) * H + i_h) * K*V, (K, V), (V, 1), (i_k * BK, i_v * BV), (BK, BV), (1, 0)) - p_dhb = tl.make_block_ptr(dhb + ((boh+i_t) * H + i_h) * V, (V,), (1,), (i_v * BV,), (BV,), (0,)) - tl.store(p_dh, b_dh.to(p_dh.dtype.element_ty), boundary_check=(0, 1)) - tl.store(p_dhb, b_dhb.to(p_dhb.dtype.element_ty), boundary_check=(0,)) - p_q = tl.make_block_ptr(q+(bos*H+i_h)*K, (K, T), (1, H*K), (i_k * BK, i_t * BT), (BK, BT), (0, 1)) - p_k = tl.make_block_ptr(k+(bos*H+i_h)*K, (T, K), (H*K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) - p_v = tl.make_block_ptr(v+(bos*H+i_h)*V, (T, V), (H*V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) - p_v_new = tl.make_block_ptr(v_new+(bos*H+i_h)*V, (T, V), (H*V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) - p_x = tl.make_block_ptr(x+(bos*H+i_h)*V, (T, V), (H*V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) - p_y = tl.make_block_ptr(y+(bos*H+i_h)*V, (T, V), (H*V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) - p_dv_new = tl.make_block_ptr(dv_new+(bos*H+i_h)*V, (T, V), (H*V, 1), (i_t*BT, i_v * BV), (BT, BV), (1, 0)) - p_dv = tl.make_block_ptr(dv+(bos*H+i_h)*V, (T, V), (H*V, 1), (i_t*BT, i_v * BV), (BT, BV), (1, 0)) - p_dk = tl.make_block_ptr(dk+(bos*H+i_h)*K, (T, K), (H*K, 1), (i_t*BT, i_k * BK), (BT, BK), (1, 0)) - p_do = tl.make_block_ptr(do+(bos*H+i_h)*V, (T, V), (H*V, 1), (i_t*BT, i_v * BV), (BT, BV), (1, 0)) - p_r = tl.make_block_ptr(r+bos*H+i_h, (T, 1), (H, 1), (i_t*BT, 0), (BT, 1), (1, 0)) - p_eta_last = eta+bos*H+i_h + (T-1)*H if i_t == NT-1 else eta+bos*H+i_h + (i_t*BT+BT-1)*H - b_k = tl.load(p_k, boundary_check=(0, 1), padding_option="zero") - b_dv_new = tl.load(p_dv_new, boundary_check=(0, 1), padding_option="zero").to(b_k.dtype) + i_t_int64 = i_t.to(tl.int64) + o_t = i_t_int64 * BT + tl.arange(0, BT) + m_t = o_t < T + m_k = m_t[:, None] & (o_k[None, :] < K) + m_q = (o_k[:, None] < K) & m_t[None, :] + m_v = m_t[:, None] & (o_v[None, :] < V) + m_hv = (o_v[:, None] < V) & (o_k[None, :] < K) + p_h = h + ((boh+i_t_int64) * H + i_h) * K*V + o_v[:, None] + o_k[None, :] * V + p_dh = dh + ((boh+i_t_int64) * H + i_h) * K*V + o_k[:, None] * V + o_v[None, :] + p_dhb = dhb + ((boh+i_t_int64) * H + i_h) * V + o_v + tl.store(p_dh, b_dh.to(p_dh.dtype.element_ty), mask=m_h) + tl.store(p_dhb, b_dhb.to(p_dhb.dtype.element_ty), mask=o_v < V) + p_q = q+(bos*H+i_h)*K + o_k[:, None] + o_t[None, :] * (H*K) + p_k = k+(bos*H+i_h)*K + o_t[:, None] * (H*K) + o_k[None, :] + p_v = v+(bos*H+i_h)*V + o_t[:, None] * (H*V) + o_v[None, :] + p_v_new = v_new+(bos*H+i_h)*V + o_t[:, None] * (H*V) + o_v[None, :] + p_x = x+(bos*H+i_h)*V + o_t[:, None] * (H*V) + o_v[None, :] + p_y = y+(bos*H+i_h)*V + o_t[:, None] * (H*V) + o_v[None, :] + p_dv_new = dv_new+(bos*H+i_h)*V + o_t[:, None] * (H*V) + o_v[None, :] + p_dv = dv+(bos*H+i_h)*V + o_t[:, None] * (H*V) + o_v[None, :] + p_dk = dk+(bos*H+i_h)*K + o_t[:, None] * (H*K) + o_k[None, :] + p_do = do+(bos*H+i_h)*V + o_t[:, None] * (H*V) + o_v[None, :] + p_r = r+bos*H+i_h + o_t[:, None] * H + p_eta_last = eta+bos*H+i_h + (T-1)*H if i_t_int64 == NT-1 else eta+bos*H+i_h + (i_t_int64*BT+BT-1)*H + b_k = tl.load(p_k, mask=m_k, other=0.0) + b_dv_new = tl.load(p_dv_new, mask=m_v, other=0.0).to(b_k.dtype) b_eta_last = tl.load(p_eta_last) b_dv_new -= tl.dot(b_eta_last * b_k, b_dh.to(b_k.dtype)) b_dv_new -= b_eta_last * b_dhb.to(b_k.dtype)[None, :] - b_v_new = tl.load(p_v_new, boundary_check=(0, 1), padding_option="zero") - b_x = tl.load(p_x, boundary_check=(0, 1), padding_option="zero").to(b_k.dtype) - b_y = tl.load(p_y, boundary_check=(0, 1), padding_option="zero").to(b_k.dtype) - b_rstd = tl.load(p_r, boundary_check=(0, 1), padding_option="zero").to(tl.float32) + b_v_new = tl.load(p_v_new, mask=m_v, other=0.0) + b_x = tl.load(p_x, mask=m_v, other=0.0).to(b_k.dtype) + b_y = tl.load(p_y, mask=m_v, other=0.0).to(b_k.dtype) + b_rstd = tl.load(p_r, mask=m_t[:, None], other=0.0).to(tl.float32) b_dy = b_rstd * (b_dv_new * V - tl.sum(b_dv_new, axis=1, keep_dims=True) - b_x * tl.sum(b_dv_new * b_x, axis=1, keep_dims=True)) / V b_dx = -b_rstd * (b_dv_new * tl.sum(b_x * b_y, axis=1, keep_dims=True) + b_y * tl.sum(b_dv_new * b_x, axis=1, keep_dims=True)) / V b_drstd = tl.sum(b_dv_new.to(b_rstd.dtype) * b_v_new.to(b_rstd.dtype) / b_rstd, axis=1, keep_dims=True) - b_v = tl.load(p_v, boundary_check=(0, 1), padding_option="zero") + b_v = tl.load(p_v, mask=m_v, other=0.0) b_w = b_w.to(b_k.dtype) b_b = b_b.to(b_k.dtype) b_dv = -b_w * b_dy.to(b_k.dtype) @@ -523,31 +564,31 @@ def chunk_ttt_linear_bwd_kernel_norm( b_dx = b_dx.to(b_k.dtype) + b_w * b_w * b_dy.to(b_k.dtype) # d_rstd, dx --> dkh --> dk, dh - b_q = tl.load(p_q, boundary_check=(0, 1), padding_option="zero") - b_h = tl.load(p_h, boundary_check=(0, 1), padding_option="zero") - b_do = tl.load(p_do, boundary_check=(0, 1), padding_option="zero") + b_q = tl.load(p_q, mask=m_q, other=0.0) + b_h = tl.load(p_h, mask=m_hv, other=0.0) + b_do = tl.load(p_do, mask=m_v, other=0.0) b_q = (b_q * scale).to(b_q.dtype) b_dkh = b_rstd * (V * b_dx - tl.sum(b_dx, axis=1, keep_dims=True) - b_x * tl.sum(b_x * b_dx, axis=1, keep_dims=True)) / V b_dkh -= b_rstd * b_rstd * b_drstd * b_x / V - b_dkh = tl.where((offs_v < V)[None, :] * (offs_t < T-i_t*BT)[:, None], b_dkh, 0.) + b_dkh = tl.where((offs_v < V)[None, :] * (offs_t < T-i_t_int64*BT)[:, None], b_dkh, 0.) b_dk += tl.dot(b_dkh, b_h.to(b_dkh.dtype)).to(b_k.dtype) b_dh += tl.dot(b_q, b_do.to(b_q.dtype)) + tl.dot(tl.trans(b_k).to(b_dkh.dtype), b_dkh) b_dhb += tl.sum(b_do + b_dkh, axis=0) b_dh = tl.where((offs_v < V)[None, :], b_dh, 0.) b_dhb = tl.where((offs_v < V), b_dhb, 0.) - tl.store(p_dv, b_dv.to(p_dv.dtype.element_ty), boundary_check=(0, 1)) - tl.store(p_dk, b_dk.to(p_dk.dtype.element_ty), boundary_check=(0, 1)) - tl.store(p_dw, b_dw.to(p_dw.dtype.element_ty), boundary_check=(0,)) - tl.store(p_db, b_db.to(p_db.dtype.element_ty), boundary_check=(0,)) + tl.store(p_dv, b_dv.to(p_dv.dtype.element_ty), mask=m_v) + tl.store(p_dk, b_dk.to(p_dk.dtype.element_ty), mask=m_k) + tl.store(p_dw, b_dw.to(p_dw.dtype.element_ty), mask=o_v < V) + tl.store(p_db, b_db.to(p_db.dtype.element_ty), mask=o_v < V) if USE_INITIAL_STATE: - p_dh0 = tl.make_block_ptr(dh0 + i_nh * K*V, (K, V), (V, 1), (i_k * BK, i_v * BV), (BK, BV), (1, 0)) - tl.store(p_dh0, b_dh.to(p_dh0.dtype.element_ty), boundary_check=(0, 1)) + p_dh0 = dh0 + i_nh * K*V + o_k[:, None] * V + o_v[None, :] + tl.store(p_dh0, b_dh.to(p_dh0.dtype.element_ty), mask=m_h) if USE_INITIAL_STATE_B: - p_dhb0 = tl.make_block_ptr(dhb0+i_nh*V, (V,), (1,), (i_v * BV,), (BV,), (0,)) - tl.store(p_dhb0, b_dhb.to(p_dhb0.dtype.element_ty), boundary_check=(0,)) + p_dhb0 = dhb0+i_nh*V + o_v + tl.store(p_dhb0, b_dhb.to(p_dhb0.dtype.element_ty), mask=o_v < V) @triton.heuristics({ @@ -588,12 +629,13 @@ def chunk_bwd_kernel_dqke( BV: tl.constexpr, IS_VARLEN: tl.constexpr, ): - i_k, i_t, i_bh = tl.program_id(0), tl.program_id(1), tl.program_id(2) + i_k, i_t = unflatten_program_id(tl.cdiv(K, BK)) + i_bh = tl.program_id(1).to(tl.int64) i_b, i_h = i_bh // H, i_bh % H 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.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) + i_n, i_t = tl.load(chunk_indices + i_t * 2).to(tl.int32), tl.load(chunk_indices + i_t * 2 + 1).to(tl.int64) + bos, eos = tl.load(cu_seqlens + i_n).to(tl.int64), tl.load(cu_seqlens + i_n + 1).to(tl.int64) T = eos - bos NT = tl.cdiv(T, BT) else: @@ -622,45 +664,52 @@ def chunk_bwd_kernel_dqke( b_ds = tl.zeros([BT, BT], dtype=tl.float32) b_de = tl.zeros([BT], dtype=tl.float32) - p_k = tl.make_block_ptr(k, (T, K), (stride_qk, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) - b_k = tl.load(p_k, boundary_check=(0, 1)) + o_t = i_t * BT + tl.arange(0, BT) + o_k = i_k * BK + tl.arange(0, BK) + m_t = o_t < T + m_k = m_t[:, None] & (o_k[None, :] < K) + p_k = k + o_t[:, None] * stride_qk + o_k[None, :] + b_k = tl.load(p_k, mask=m_k, other=0.0) p_e_last = (e + (i_t*BT+BT-1)*stride_e) if (i_t*BT+BT) <= T else (e + (T-1)*stride_e) i_last = (BT-1) if (i_t*BT+BT) <= T else (T % BT-1) mask = (tl.arange(0, BT) == i_last) b_e_last = tl.load(p_e_last) for i_v in range(tl.cdiv(V, BV)): - p_v = tl.make_block_ptr(v, (T, V), (stride_vo, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) - p_do = tl.make_block_ptr(do, (T, V), (stride_vo, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) - p_h = tl.make_block_ptr(h, (V, K), (1, V), (i_v * BV, i_k * BK), (BV, BK), (0, 1)) - p_dh = tl.make_block_ptr(dh, (V, K), (1, V), (i_v * BV, i_k * BK), (BV, BK), (0, 1)) - p_dhb = tl.make_block_ptr(dhb, (V,), (1,), (i_v * BV,), (BV,), (0,)) + o_v = i_v * BV + tl.arange(0, BV) + m_v = m_t[:, None] & (o_v[None, :] < V) + m_h = (o_v[:, None] < V) & (o_k[None, :] < K) + p_v = v + o_t[:, None] * stride_vo + o_v[None, :] + p_do = do + o_t[:, None] * stride_vo + o_v[None, :] + p_h = h + o_v[:, None] + o_k[None, :] * V + p_dh = dh + o_v[:, None] + o_k[None, :] * V + p_dhb = dhb + o_v # [BT, BV] - b_v = tl.load(p_v, boundary_check=(0, 1)) - b_do = tl.load(p_do, boundary_check=(0, 1)) + b_v = tl.load(p_v, mask=m_v, other=0.0) + b_do = tl.load(p_do, mask=m_v, other=0.0) # [BV, BK] - b_h = tl.load(p_h, boundary_check=(0, 1)) - b_dh = tl.load(p_dh, boundary_check=(0, 1)) + b_h = tl.load(p_h, mask=m_h, other=0.0) + b_dh = tl.load(p_dh, mask=m_h, other=0.0) # [BV] - b_dhb = tl.load(p_dhb, boundary_check=(0,)) + b_dhb = tl.load(p_dhb, mask=o_v < V, other=0.0) # [BT, BV] @ [BV, BT] -> [BT, BT] - b_ds += tl.dot(b_do, tl.trans(b_v)) + b_ds = tl.dot(b_do, tl.trans(b_v), b_ds) # [BT, BV] @ [BV, BK] -> [BT, BK] - b_dq += tl.dot(b_do, b_h.to(b_do.dtype)) + b_dq = tl.dot(b_do, b_h.to(b_do.dtype), b_dq) # [BT, BV] @ [BV, BK] -> [BT, BK] b_dk -= b_e_last * tl.dot(b_v, b_dh.to(b_v.dtype)) b_de -= mask * tl.sum(tl.trans(b_dh) * tl.dot(tl.trans(b_k), b_v.to(b_k.dtype))) b_de -= mask * tl.sum(b_dhb * tl.sum(b_v, axis=0).to(b_k.dtype)) o_i = tl.arange(0, BT) - p_q = tl.make_block_ptr(q, (T, K), (stride_qk, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) - p_e = tl.make_block_ptr(e, (T,), (stride_e,), (i_t * BT,), (BT,), (0,)) - b_q = tl.load(p_q, boundary_check=(0, 1)) - b_e = tl.load(p_e, boundary_check=(0,)) + p_q = q + o_t[:, None] * stride_qk + o_k[None, :] + p_e = e + o_t * stride_e + b_q = tl.load(p_q, mask=m_k, other=0.0) + b_e = tl.load(p_e, mask=m_t, other=0.0) - p_dq = tl.make_block_ptr(dq, (T, K), (stride_qk, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) - p_dk = tl.make_block_ptr(dk, (T, K), (stride_qk, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) - p_de = tl.make_block_ptr(de, (T,), (stride_e,), (i_t * BT,), (BT,), (0,)) + p_dq = dq + o_t[:, None] * stride_qk + o_k[None, :] + p_dk = dk + o_t[:, None] * stride_qk + o_k[None, :] + p_de = de + o_t * stride_e b_ds = tl.where(o_i[:, None] >= o_i[None, :], b_ds, 0) b_ds = b_ds.to(b_k.dtype) @@ -669,9 +718,9 @@ def chunk_bwd_kernel_dqke( b_de -= tl.sum(scale * tl.dot(b_ds, b_k) * b_q, axis=1) b_de -= tl.sum(b_ds, axis=1) b_dq *= scale - tl.store(p_dq, b_dq.to(p_dq.dtype.element_ty), boundary_check=(0, 1)) - tl.store(p_dk, b_dk.to(p_dk.dtype.element_ty), boundary_check=(0, 1)) - tl.store(p_de, b_de.to(p_de.dtype.element_ty), boundary_check=(0,)) + tl.store(p_dq, b_dq.to(p_dq.dtype.element_ty), mask=m_k) + tl.store(p_dk, b_dk.to(p_dk.dtype.element_ty), mask=m_k) + tl.store(p_de, b_de.to(p_de.dtype.element_ty), mask=m_t) def chunk_ttt_linear_fwd_h( @@ -712,7 +761,7 @@ def chunk_ttt_linear_fwd_h( final_state_bias = k.new_empty(N, H, 1, V, dtype=torch.float32) if output_final_state else None v_new = torch.empty_like(v) - grid = (NK, NV, N * H) + grid = (NK * NV * N * H,) chunk_ttt_linear_fwd_kernel_h[grid]( k=k, @@ -770,7 +819,7 @@ def chunk_ttt_linear_fwd_o( o = torch.empty_like(v) - grid = (NV, NT, B * H) + grid = (NV * NT, B * H) chunk_ttt_linear_fwd_kernel_o[grid]( q, k, @@ -830,7 +879,7 @@ def chunk_ttt_linear_bwd_h( y = torch.empty_like(v) v_new = torch.empty_like(v) - grid = (NK, NV, N * H) + grid = (NK * NV * N * H,) chunk_ttt_linear_bwd_kernel_h[grid]( k=k, @@ -943,8 +992,8 @@ def chunk_ttt_linear_bwd_norm( assert NK == 1, 'NK > 1 is not supported by TTT.' assert NV == 1, 'NV > 1 is not supported by TTT.' - dh = q.new_empty(B, NT, H, K, V) - dhb = q.new_empty(B, NT, H, 1, V) + dh = q.new_empty(B, NT, H, K, V, dtype=torch.float32) + dhb = q.new_empty(B, NT, H, 1, V, dtype=torch.float32) dh0 = torch.empty_like(h0, dtype=torch.float32) if h0 is not None else None dhb0 = torch.empty_like(hb0, dtype=torch.float32) if hb0 is not None else None dv = torch.empty_like(v) @@ -952,7 +1001,7 @@ def chunk_ttt_linear_bwd_norm( dw = w.new_empty(B, H, V) db = b.new_empty(B, H, V) - grid = (NK, NV, N * H) + grid = (NK * NV * N * H,) chunk_ttt_linear_bwd_kernel_norm[grid]( q=q, k=k, @@ -1127,7 +1176,7 @@ def chunk_ttt_linear_bwd_dqke( dq = torch.empty_like(q) dk = torch.empty_like(k) de = torch.empty_like(eta) - grid = (NK, NT, B * H) + grid = (NK * NT, B * H) chunk_bwd_kernel_dqke[grid]( q=q, @@ -1389,7 +1438,6 @@ def chunk_ttt_linear( output_final_state: bool = False, cu_seqlens: torch.LongTensor | None = None, cu_seqlens_cpu: torch.LongTensor | None = None, - **kwargs, ): r""" Args: @@ -1434,10 +1482,6 @@ def chunk_ttt_linear( assert k.shape[-1] == v.shape[-1], "DK must equal to DV." if isinstance(eta, float): eta = torch.full_like(q[:, :, :, :1], eta) - if 'head_first' in kwargs: - raise DeprecationWarning( - "head_first has been removed. Inputs must be in `[B, T, H, ...]` format.", - ) if cu_seqlens is not None: if q.shape[0] != 1: raise ValueError( diff --git a/build/torch-cuda/ops/ttt/fused_chunk.py b/build/torch-cuda/ops/ttt/fused_chunk.py index 200a76f4b0ec8668f116799e791cd346ed5f2519..f1be64bf0403ddb412bd7568e46f3041fb0895b6 100644 --- a/build/torch-cuda/ops/ttt/fused_chunk.py +++ b/build/torch-cuda/ops/ttt/fused_chunk.py @@ -58,10 +58,10 @@ def fused_chunk_ttt_linear_fwd_kernel( STORE_FINAL_STATE: tl.constexpr, IS_VARLEN: tl.constexpr, ): - i_nh = tl.program_id(0) + i_nh = tl.program_id(0).to(tl.int64) i_n, i_h = i_nh // H, i_nh % H if IS_VARLEN: - bos, eos = tl.load(cu_seqlens + i_n).to(tl.int32), tl.load(cu_seqlens + i_n + 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 NT = tl.cdiv(T, BT) else: @@ -78,24 +78,32 @@ def fused_chunk_ttt_linear_fwd_kernel( b_h = tl.zeros([BK, BV], dtype=tl.float32) # [BV] b_hb = tl.zeros([BV], dtype=tl.float32) + o_k = tl.arange(0, BK) + o_v = tl.arange(0, BV) + m_h = (o_k[:, None] < K) & (o_v[None, :] < V) if USE_INITIAL_STATE: - p_h0 = tl.make_block_ptr(h0 + i_nh * K * V, (K, V), (V, 1), (0, 0), (BK, BV), (1, 0)) - b_h = tl.load(p_h0, boundary_check=(0, 1), padding_option="zero").to(tl.float32) + p_h0 = h0 + i_nh * K * V + o_k[:, None] * V + o_v[None, :] + b_h = tl.load(p_h0, mask=m_h, other=0.0).to(tl.float32) if USE_INITIAL_STATE_B: - p_hb0 = tl.make_block_ptr(hb0 + i_nh * V, (V,), (1,), (0,), (BV,), (0,)) - b_hb = tl.load(p_hb0, boundary_check=(0,), padding_option="zero").to(tl.float32) + p_hb0 = hb0 + i_nh * V + o_v + b_hb = tl.load(p_hb0, mask=o_v < V, other=0.0).to(tl.float32) for i_t in range(NT): - p_q = tl.make_block_ptr(q+(bos*H+i_h)*K, (T, K), (H*K, 1), (i_t*BT, 0), (BT, BK), (1, 0)) - p_k = tl.make_block_ptr(k+(bos*H+i_h)*K, (K, T), (1, H*K), (0, i_t*BT), (BK, BT), (0, 1)) - p_v = tl.make_block_ptr(v+(bos*H+i_h)*V, (T, V), (H*V, 1), (i_t*BT, 0), (BT, BV), (1, 0)) - p_o = tl.make_block_ptr(o+(bos*H+i_h)*V, (T, V), (H*V, 1), (i_t*BT, 0), (BT, BV), (1, 0)) - p_e = tl.make_block_ptr(eta+(bos*H+i_h), (T,), (H,), (i_t*BT,), (BT,), (0,)) + o_t = i_t * BT + tl.arange(0, BT) + m_t = o_t < T + m_q = m_t[:, None] & (o_k[None, :] < K) + m_k = (o_k[:, None] < K) & m_t[None, :] + m_v = m_t[:, None] & (o_v[None, :] < V) + p_q = q+(bos*H+i_h)*K + o_t[:, None] * (H*K) + o_k[None, :] + p_k = k+(bos*H+i_h)*K + o_k[:, None] + o_t[None, :] * (H*K) + p_v = v+(bos*H+i_h)*V + o_t[:, None] * (H*V) + o_v[None, :] + p_o = o+(bos*H+i_h)*V + o_t[:, None] * (H*V) + o_v[None, :] + p_e = eta+(bos*H+i_h) + o_t * H p_e_last = eta+bos*H+i_h + (T-1)*H if i_t == NT-1 else eta+bos*H+i_h + (i_t*BT+BT-1)*H # [BK, BT] - b_k = tl.load(p_k, boundary_check=(0, 1), padding_option="zero") + b_k = tl.load(p_k, mask=m_k, other=0.0) # [BT, BV] - b_v = tl.load(p_v, boundary_check=(0, 1), padding_option="zero") + b_v = tl.load(p_v, mask=m_v, other=0.0) # [BT, BV] b_kh = tl.dot(tl.trans(b_k), b_h.to(b_k.dtype), allow_tf32=False).to(tl.float32) + b_hb[None, :] @@ -113,9 +121,9 @@ def fused_chunk_ttt_linear_fwd_kernel( * tl.sum(b_v * b_kh_hat.to(b_k.dtype), axis=1, keep_dims=True)) / V # [BT, BK] - b_q = tl.load(p_q, boundary_check=(0, 1), padding_option="zero") + b_q = tl.load(p_q, mask=m_q, other=0.0) # [BT] - b_e = tl.load(p_e, boundary_check=(0,), padding_option="zero") + b_e = tl.load(p_e, mask=m_t, other=0.0) b_q = (b_q * scale).to(b_k.dtype) # [BT, BT] @@ -125,19 +133,19 @@ def fused_chunk_ttt_linear_fwd_kernel( b_o = - tl.dot(b_e[:, None] * b_A.to(b_v2.dtype), b_v2, allow_tf32=False) b_o += b_hb[None, :] - tl.dot(b_Ae.to(b_v2.dtype), b_v2, allow_tf32=False) - b_o += tl.dot(b_q, b_h.to(b_q.dtype), allow_tf32=False) + b_o = tl.dot(b_q, b_h.to(b_q.dtype), b_o, allow_tf32=False) b_e_last = tl.load(p_e_last) b_h = b_h - tl.dot(b_e_last * b_k, b_v2.to(b_k.dtype), allow_tf32=False) b_hb = b_hb - tl.sum(b_e_last * b_v2.to(b_k.dtype), axis=0) b_h = tl.where((v_i < V)[None, :], b_h, 0.) b_hb = tl.where((v_i < V), b_hb, 0.) - tl.store(p_o, b_o.to(p_o.dtype.element_ty), boundary_check=(0, 1)) + tl.store(p_o, b_o.to(p_o.dtype.element_ty), mask=m_v) if STORE_FINAL_STATE: - p_ht = tl.make_block_ptr(ht + i_nh * K*V, (K, V), (V, 1), (0, 0), (BK, BV), (1, 0)) - p_hbt = tl.make_block_ptr(hbt + i_nh * V, (V,), (1,), (0,), (BV,), (0,)) - tl.store(p_ht, b_h.to(p_ht.dtype.element_ty), boundary_check=(0, 1)) - tl.store(p_hbt, b_hb.to(p_hbt.dtype.element_ty), boundary_check=(0,)) + p_ht = ht + i_nh * K*V + o_k[:, None] * V + o_v[None, :] + p_hbt = hbt + i_nh * V + o_v + tl.store(p_ht, b_h.to(p_ht.dtype.element_ty), mask=m_h) + tl.store(p_hbt, b_hb.to(p_hbt.dtype.element_ty), mask=o_v < V) @triton.heuristics({ @@ -181,7 +189,7 @@ def fused_chunk_ttt_linear_bwd_kernel_h( USE_INITIAL_STATE: tl.constexpr, USE_INITIAL_STATE_B: tl.constexpr, ): - i_nh = tl.program_id(0) + i_nh = tl.program_id(0).to(tl.int64) i_n, i_h = i_nh // H, i_nh % H bos, _ = i_n * T, i_n * T + T NT = tl.cdiv(T, BT) @@ -197,30 +205,38 @@ def fused_chunk_ttt_linear_bwd_kernel_h( b_h = tl.zeros([BK, BV], dtype=tl.float32) # [BV] b_hb = tl.zeros([BV], dtype=tl.float32) + o_k = tl.arange(0, BK) + o_v = tl.arange(0, BV) + m_h = (o_k[:, None] < K) & (o_v[None, :] < V) if USE_INITIAL_STATE: - p_h0 = tl.make_block_ptr(h0 + i_nh * K * V, (K, V), (V, 1), (0, 0), (BK, BV), (1, 0)) - b_h = tl.load(p_h0, boundary_check=(0, 1), padding_option="zero").to(tl.float32) + p_h0 = h0 + i_nh * K * V + o_k[:, None] * V + o_v[None, :] + b_h = tl.load(p_h0, mask=m_h, other=0.0).to(tl.float32) if USE_INITIAL_STATE_B: - p_hb0 = tl.make_block_ptr(hb0 + i_nh * V, (V,), (1,), (0,), (BV,), (0,)) - b_hb = tl.load(p_hb0, boundary_check=(0,), padding_option="zero").to(tl.float32) + p_hb0 = hb0 + i_nh * V + o_v + b_hb = tl.load(p_hb0, mask=o_v < V, other=0.0).to(tl.float32) for i_t in range(NT): - p_h = tl.make_block_ptr(h+((boh+i_t)*H+i_h)*K*V, (K, V), (V, 1), (0, 0), (BK, BV), (1, 0)) - p_k = tl.make_block_ptr(k+(bos*H+i_h)*K, (K, T), (1, H*K), (0, i_t*BT), (BK, BT), (0, 1)) - p_v = tl.make_block_ptr(v+(bos*H+i_h)*V, (T, V), (H*V, 1), (i_t*BT, 0), (BT, BV), (1, 0)) - p_v2 = tl.make_block_ptr(v2+(bos*H+i_h)*V, (T, V), (H*V, 1), (i_t*BT, 0), (BT, BV), (1, 0)) - p_x = tl.make_block_ptr(x+(bos*H+i_h)*V, (T, V), (H*V, 1), (i_t*BT, 0), (BT, BV), (1, 0)) - p_y = tl.make_block_ptr(y+(bos*H+i_h)*V, (T, V), (H*V, 1), (i_t*BT, 0), (BT, BV), (1, 0)) - p_r = tl.make_block_ptr(r+bos*H+i_h, (T, 1), (H, 1), (i_t*BT, 0), (BT, 1), (1, 0)) - p_e = tl.make_block_ptr(eta+(bos*H+i_h), (T,), (H,), (i_t*BT,), (BT,), (0,)) - p_dq = tl.make_block_ptr(dq+(bos*H+i_h)*K, (T, K), (H*K, 1), (i_t*BT, 0), (BT, BK), (1, 0)) - p_do = tl.make_block_ptr(do+(bos*H+i_h)*V, (T, V), (H*V, 1), (i_t*BT, 0), (BT, BV), (1, 0)) + o_t = i_t * BT + tl.arange(0, BT) + m_t = o_t < T + m_q = m_t[:, None] & (o_k[None, :] < K) + m_k = (o_k[:, None] < K) & m_t[None, :] + m_v = m_t[:, None] & (o_v[None, :] < V) + p_h = h+((boh+i_t)*H+i_h)*K*V + o_k[:, None] * V + o_v[None, :] + p_k = k+(bos*H+i_h)*K + o_k[:, None] + o_t[None, :] * (H*K) + p_v = v+(bos*H+i_h)*V + o_t[:, None] * (H*V) + o_v[None, :] + p_v2 = v2+(bos*H+i_h)*V + o_t[:, None] * (H*V) + o_v[None, :] + p_x = x+(bos*H+i_h)*V + o_t[:, None] * (H*V) + o_v[None, :] + p_y = y+(bos*H+i_h)*V + o_t[:, None] * (H*V) + o_v[None, :] + p_r = r+bos*H+i_h + o_t[:, None] * H + p_e = eta+(bos*H+i_h) + o_t * H + p_dq = dq+(bos*H+i_h)*K + o_t[:, None] * (H*K) + o_k[None, :] + p_do = do+(bos*H+i_h)*V + o_t[:, None] * (H*V) + o_v[None, :] p_e_last = eta+bos*H+i_h + (T-1)*H if i_t == NT-1 else eta+bos*H+i_h + (i_t*BT+BT-1)*H - tl.store(p_h, b_h.to(p_h.dtype.element_ty), boundary_check=(0, 1)) + tl.store(p_h, b_h.to(p_h.dtype.element_ty), mask=m_h) # [BK, BT] - b_k = tl.load(p_k, boundary_check=(0, 1), padding_option="zero") + b_k = tl.load(p_k, mask=m_k, other=0.0) # [BT, BV] - b_v = tl.load(p_v, boundary_check=(0, 1), padding_option="zero") + b_v = tl.load(p_v, mask=m_v, other=0.0) b_kh = tl.dot(tl.trans(b_k), b_h.to(b_k.dtype), allow_tf32=False).to(tl.float32) + b_hb[None, :] b_kh = tl.where((v_i < V)[None, :], b_kh, 0.) @@ -235,13 +251,13 @@ def fused_chunk_ttt_linear_bwd_kernel_h( b_v = tl.where((v_i < V)[None, :], b_v * b_w[None, :].to(b_k.dtype), 0.) b_v2 = rstd * (V * b_v - tl.sum(b_v, axis=1, keep_dims=True) - b_kh_hat.to(b_k.dtype) * tl.sum(b_v * b_kh_hat.to(b_k.dtype), axis=1, keep_dims=True)) / V - tl.store(p_x, b_kh_hat.to(p_x.dtype.element_ty), boundary_check=(0, 1)) - tl.store(p_y, b_v.to(p_y.dtype.element_ty), boundary_check=(0, 1)) - tl.store(p_r, rstd.to(p_r.dtype.element_ty), boundary_check=(0, 1)) - tl.store(p_v2, b_v2.to(p_v2.dtype.element_ty), boundary_check=(0, 1)) + tl.store(p_x, b_kh_hat.to(p_x.dtype.element_ty), mask=m_v) + tl.store(p_y, b_v.to(p_y.dtype.element_ty), mask=m_v) + tl.store(p_r, rstd.to(p_r.dtype.element_ty), mask=m_t[:, None]) + tl.store(p_v2, b_v2.to(p_v2.dtype.element_ty), mask=m_v) - b_e = tl.load(p_e, boundary_check=(0,), padding_option="zero") - b_do = tl.load(p_do, boundary_check=(0, 1), padding_option="zero") + b_e = tl.load(p_e, mask=m_t, other=0.0) + b_do = tl.load(p_do, mask=m_v, other=0.0) b_v2 = tl.where((v_i < V)[None, :], b_v2, 0.) b_ds = tl.dot(b_do, tl.trans(b_v2).to(b_do.dtype)) @@ -256,7 +272,7 @@ def fused_chunk_ttt_linear_bwd_kernel_h( b_hb = b_hb - tl.sum(b_e_last * b_v2.to(b_k.dtype), axis=0) b_h = tl.where((v_i < V)[None, :], b_h, 0.) b_hb = tl.where((v_i < V), b_hb, 0.) - tl.store(p_dq, b_dq.to(p_dq.dtype.element_ty), boundary_check=(0, 1)) + tl.store(p_dq, b_dq.to(p_dq.dtype.element_ty), mask=m_q) @triton.heuristics({ @@ -309,7 +325,7 @@ def fused_chunk_ttt_linear_bwd_kernel_dh( USE_FINAL_STATE_GRADIENT: tl.constexpr, USE_FINAL_STATE_GRADIENT_B: tl.constexpr, ): - i_nh = tl.program_id(0) + i_nh = tl.program_id(0).to(tl.int64) i_n, i_h = i_nh // H, i_nh % H bos, _ = i_n * T, i_n * T + T NT = tl.cdiv(T, BT) @@ -319,12 +335,15 @@ def fused_chunk_ttt_linear_bwd_kernel_dh( b_dh = tl.zeros([BK, BV], dtype=tl.float32) # [BV] b_dhb = tl.zeros([BV], dtype=tl.float32) + o_k = tl.arange(0, BK) + o_v = tl.arange(0, BV) + m_h = (o_k[:, None] < K) & (o_v[None, :] < V) if USE_FINAL_STATE_GRADIENT: - p_dht = tl.make_block_ptr(dht + i_nh * K*V, (K, V), (V, 1), (0, 0), (BK, BV), (1, 0)) - b_dh += tl.load(p_dht, boundary_check=(0, 1), padding_option="zero") + p_dht = dht + i_nh * K*V + o_k[:, None] * V + o_v[None, :] + b_dh += tl.load(p_dht, mask=m_h, other=0.0) if USE_FINAL_STATE_GRADIENT_B: - p_dhbt = tl.make_block_ptr(dhbt + i_nh * V, (V,), (1,), (0,), (BV,), (0,)) - b_dhb += tl.load(p_dhbt, boundary_check=(0,), padding_option="zero") + p_dhbt = dhbt + i_nh * V + o_v + b_dhb += tl.load(p_dhbt, mask=o_v < V, other=0.0) # [BV] o_i = tl.arange(0, BT) @@ -335,28 +354,34 @@ def fused_chunk_ttt_linear_bwd_kernel_dh( b_b = tl.load(b + i_h * V + v_i, mask=v_i < V, other=0.) b_dw = tl.zeros([BV], dtype=b_w.dtype) b_db = tl.zeros([BV], dtype=b_b.dtype) - p_dw = tl.make_block_ptr(dw + i_nh * V, (V,), (1,), (0,), (BV,), (0,)) - p_db = tl.make_block_ptr(db + i_nh * V, (V,), (1,), (0,), (BV,), (0,)) + p_dw = dw + i_nh * V + o_v + p_db = db + i_nh * V + o_v for i_t in range(NT - 1, -1, -1): - p_h = tl.make_block_ptr(h+((boh+i_t)*H+i_h)*K*V, (V, K), (1, V), (0, 0), (BV, BK), (0, 1)) - p_q = tl.make_block_ptr(q+(bos*H+i_h)*K, (K, T), (1, H*K), (0, i_t*BT), (BK, BT), (0, 1)) - p_k = tl.make_block_ptr(k+(bos*H+i_h)*K, (T, K), (H*K, 1), (i_t*BT, 0), (BT, BK), (1, 0)) - p_v = tl.make_block_ptr(v+(bos*H+i_h)*V, (T, V), (H*V, 1), (i_t*BT, 0), (BT, BV), (1, 0)) - p_v2 = tl.make_block_ptr(v2+(bos*H+i_h)*V, (T, V), (H*V, 1), (i_t*BT, 0), (BT, BV), (1, 0)) - p_x = tl.make_block_ptr(x+(bos*H+i_h)*V, (T, V), (H*V, 1), (i_t*BT, 0), (BT, BV), (1, 0)) - p_y = tl.make_block_ptr(y+(bos*H+i_h)*V, (T, V), (H*V, 1), (i_t*BT, 0), (BT, BV), (1, 0)) - p_r = tl.make_block_ptr(r+bos*H+i_h, (T, 1), (H, 1), (i_t*BT, 0), (BT, 1), (1, 0)) - p_e = tl.make_block_ptr(eta+(bos*H+i_h), (T,), (H,), (i_t*BT,), (BT,), (0,)) - p_dv = tl.make_block_ptr(dv+(bos*H+i_h)*V, (T, V), (H*V, 1), (i_t*BT, 0), (BT, BV), (1, 0)) - p_dk = tl.make_block_ptr(dk+(bos*H+i_h)*K, (T, K), (H*K, 1), (i_t*BT, 0), (BT, BK), (1, 0)) - p_do = tl.make_block_ptr(do+(bos*H+i_h)*V, (T, V), (H*V, 1), (i_t*BT, 0), (BT, BV), (1, 0)) - p_de = tl.make_block_ptr(de+(bos*H+i_h), (T,), (H,), (i_t*BT,), (BT,), (0,)) + o_t = i_t * BT + tl.arange(0, BT) + m_t = o_t < T + m_q = (o_k[:, None] < K) & m_t[None, :] + m_k = m_t[:, None] & (o_k[None, :] < K) + m_v = m_t[:, None] & (o_v[None, :] < V) + m_hv = (o_v[:, None] < V) & (o_k[None, :] < K) + p_h = h+((boh+i_t)*H+i_h)*K*V + o_v[:, None] + o_k[None, :] * V + p_q = q+(bos*H+i_h)*K + o_k[:, None] + o_t[None, :] * (H*K) + p_k = k+(bos*H+i_h)*K + o_t[:, None] * (H*K) + o_k[None, :] + p_v = v+(bos*H+i_h)*V + o_t[:, None] * (H*V) + o_v[None, :] + p_v2 = v2+(bos*H+i_h)*V + o_t[:, None] * (H*V) + o_v[None, :] + p_x = x+(bos*H+i_h)*V + o_t[:, None] * (H*V) + o_v[None, :] + p_y = y+(bos*H+i_h)*V + o_t[:, None] * (H*V) + o_v[None, :] + p_r = r+bos*H+i_h + o_t[:, None] * H + p_e = eta+(bos*H+i_h) + o_t * H + p_dv = dv+(bos*H+i_h)*V + o_t[:, None] * (H*V) + o_v[None, :] + p_dk = dk+(bos*H+i_h)*K + o_t[:, None] * (H*K) + o_k[None, :] + p_do = do+(bos*H+i_h)*V + o_t[:, None] * (H*V) + o_v[None, :] + p_de = de+(bos*H+i_h) + o_t * H p_e_last = eta+bos*H+i_h + (T-1)*H if i_t == NT-1 else eta+bos*H+i_h + (i_t*BT+BT-1)*H - b_q = tl.load(p_q, boundary_check=(0, 1), padding_option="zero") - b_k = tl.load(p_k, boundary_check=(0, 1), padding_option="zero") - b_e = tl.load(p_e, boundary_check=(0,), padding_option="zero") - b_do = tl.load(p_do, boundary_check=(0, 1), padding_option="zero") + b_q = tl.load(p_q, mask=m_q, other=0.0) + b_k = tl.load(p_k, mask=m_k, other=0.0) + b_e = tl.load(p_e, mask=m_t, other=0.0) + b_do = tl.load(p_do, mask=m_v, other=0.0) b_e_last = tl.load(p_e_last) b_A = tl.dot(b_k, b_q) b_A = - tl.where(m_A_t, b_A * scale * b_e[None, :], 0).to(do.dtype.element_ty) @@ -365,17 +390,17 @@ def fused_chunk_ttt_linear_bwd_kernel_dh( b_dv_new -= tl.dot(b_e_last * b_k, b_dh.to(b_k.dtype)) b_dv_new -= b_e_last * b_dhb.to(b_k.dtype)[None, :] - b_v2 = tl.load(p_v2, boundary_check=(0, 1), padding_option="zero").to(b_k.dtype) - b_x = tl.load(p_x, boundary_check=(0, 1), padding_option="zero").to(b_k.dtype) - b_y = tl.load(p_y, boundary_check=(0, 1), padding_option="zero").to(b_k.dtype) - b_rstd = tl.load(p_r, boundary_check=(0, 1), padding_option="zero").to(tl.float32) + b_v2 = tl.load(p_v2, mask=m_v, other=0.0).to(b_k.dtype) + b_x = tl.load(p_x, mask=m_v, other=0.0).to(b_k.dtype) + b_y = tl.load(p_y, mask=m_v, other=0.0).to(b_k.dtype) + b_rstd = tl.load(p_r, mask=m_t[:, None], other=0.0).to(tl.float32) b_dy = b_rstd * (b_dv_new * V - tl.sum(b_dv_new, axis=1, keep_dims=True) - b_x * tl.sum(b_dv_new * b_x, axis=1, keep_dims=True)) / V b_dx = -b_rstd * (b_dv_new * tl.sum(b_x * b_y, axis=1, keep_dims=True) + b_y * tl.sum(b_dv_new * b_x, axis=1, keep_dims=True)) / V b_drstd = tl.sum(b_dv_new.to(b_rstd.dtype) * b_v2.to(b_rstd.dtype) / b_rstd, axis=1, keep_dims=True) - b_v = tl.load(p_v, boundary_check=(0, 1), padding_option="zero") + b_v = tl.load(p_v, mask=m_v, other=0.0) b_w = b_w.to(b_k.dtype) b_b = b_b.to(b_k.dtype) b_dv = -b_w * b_dy.to(b_k.dtype) @@ -385,7 +410,7 @@ def fused_chunk_ttt_linear_bwd_kernel_dh( b_db += tl.sum(b_w * b_dy.to(b_k.dtype), axis=0).to(b_db.dtype) b_dx = b_dx.to(b_k.dtype) + b_w * b_w * b_dy.to(b_k.dtype) - b_h = tl.load(p_h, boundary_check=(0, 1), padding_option="zero") + b_h = tl.load(p_h, mask=m_hv, other=0.0) b_q = (b_q * scale).to(b_q.dtype) b_dkh = b_rstd * (V * b_dx - tl.sum(b_dx, axis=1, keep_dims=True) - b_x * tl.sum(b_x * b_dx, axis=1, keep_dims=True)) / V @@ -409,18 +434,18 @@ def fused_chunk_ttt_linear_bwd_kernel_dh( b_dh = tl.where((v_i < V)[None, :], b_dh, 0.) b_dhb = tl.where((v_i < V), b_dhb, 0.) - tl.store(p_dk, b_dk.to(p_dk.dtype.element_ty), boundary_check=(0, 1)) - tl.store(p_dv, b_dv.to(p_dv.dtype.element_ty), boundary_check=(0, 1)) - tl.store(p_de, b_de.to(p_de.dtype.element_ty), boundary_check=(0,)) - tl.store(p_dw, b_dw.to(p_dw.dtype.element_ty), boundary_check=(0,)) - tl.store(p_db, b_db.to(p_db.dtype.element_ty), boundary_check=(0,)) + tl.store(p_dk, b_dk.to(p_dk.dtype.element_ty), mask=m_k) + tl.store(p_dv, b_dv.to(p_dv.dtype.element_ty), mask=m_v) + tl.store(p_de, b_de.to(p_de.dtype.element_ty), mask=m_t) + tl.store(p_dw, b_dw.to(p_dw.dtype.element_ty), mask=o_v < V) + tl.store(p_db, b_db.to(p_db.dtype.element_ty), mask=o_v < V) if USE_INITIAL_STATE: - p_dh0 = tl.make_block_ptr(dh0+i_nh*K*V, (K, V), (V, 1), (0, 0), (BK, BV), (1, 0)) - tl.store(p_dh0, b_dh.to(p_dh0.dtype.element_ty), boundary_check=(0, 1)) + p_dh0 = dh0+i_nh*K*V + o_k[:, None] * V + o_v[None, :] + tl.store(p_dh0, b_dh.to(p_dh0.dtype.element_ty), mask=m_h) if USE_INITIAL_STATE_B: - p_dhb0 = tl.make_block_ptr(dhb0+i_nh*V, (V,), (1,), (0,), (BV,), (0,)) - tl.store(p_dhb0, b_dhb.to(p_dhb0.dtype.element_ty), boundary_check=(0,)) + p_dhb0 = dhb0+i_nh*V + o_v + tl.store(p_dhb0, b_dhb.to(p_dhb0.dtype.element_ty), mask=o_v < V) def fused_chunk_ttt_linear_bwd_h( @@ -743,7 +768,6 @@ def fused_chunk_ttt_linear( initial_state_bias: torch.Tensor | None = None, output_final_state: bool = False, cu_seqlens: torch.LongTensor | None = None, - **kwargs, ): r""" Args: @@ -786,10 +810,6 @@ def fused_chunk_ttt_linear( assert k.shape[-1] == v.shape[-1], "DK must equal to DV." if isinstance(eta, float): eta = torch.full_like(q[:, :, :, :1], eta) - if 'head_first' in kwargs: - raise DeprecationWarning( - "head_first has been removed. Inputs must be in `[B, T, H, ...]` format.", - ) if cu_seqlens is not None: if q.shape[0] != 1: raise ValueError( diff --git a/build/torch-cuda/ops/ttt/naive.py b/build/torch-cuda/ops/ttt/naive.py index 55a543534970f6c2bb038f8de4979a65f68bea31..4a16354e72401d12152bd204f7c2cae353ba15b7 100644 --- a/build/torch-cuda/ops/ttt/naive.py +++ b/build/torch-cuda/ops/ttt/naive.py @@ -87,7 +87,6 @@ def chunk_ttt_linear_ref( initial_state: torch.Tensor = None, initial_state_bias: torch.Tensor = None, output_final_state: bool = False, - head_first: bool = False, ): assert q.dtype == k.dtype == v.dtype assert k.shape[-1] == v.shape[-1], "The key and value dimension must be the same." @@ -95,11 +94,10 @@ def chunk_ttt_linear_ref( eta = torch.full_like(q[:, :, :, :1], eta) if scale is None: scale = k.shape[-1] ** -0.5 - if not head_first: - q = q.transpose(1, 2) - k = k.transpose(1, 2) - v = v.transpose(1, 2) - eta = eta.transpose(1, 2) + q = q.transpose(1, 2) + k = k.transpose(1, 2) + v = v.transpose(1, 2) + eta = eta.transpose(1, 2) T = q.shape[-2] padded = (mini_batch_size - (T % mini_batch_size)) % mini_batch_size if padded > 0: @@ -125,6 +123,5 @@ def chunk_ttt_linear_ref( output_final_state, ) o = o[:, :, :T, :].contiguous() - if not head_first: - o = o.transpose(1, 2) + o = o.transpose(1, 2) return o, final_state, final_state_bias diff --git a/build/torch-cuda/ops/utils/__init__.py b/build/torch-cuda/ops/utils/__init__.py index 88acd8b9cb58d763c2b81ea68337ac445e6658cf..0c3aa0b717d8b08993f33fc4b33b030d03b811f7 100644 --- a/build/torch-cuda/ops/utils/__init__.py +++ b/build/torch-cuda/ops/utils/__init__.py @@ -5,8 +5,8 @@ # For a list of all contributors, visit: # https://github.com/fla-org/flash-linear-attention/graphs/contributors -from .csr import prepare_block_csr -from .cumsum import ( +from ...ops.utils.csr import prepare_block_csr +from ...ops.utils.cumsum import ( chunk_global_cumsum, chunk_global_cumsum_scalar, chunk_global_cumsum_vector, @@ -14,9 +14,10 @@ from .cumsum import ( chunk_local_cumsum_scalar, chunk_local_cumsum_vector, ) -from .index import ( +from ...ops.utils.index import ( get_max_num_splits, prepare_chunk_indices, + prepare_chunk_indices_static, prepare_chunk_offsets, prepare_cu_seqlens_from_lens, prepare_cu_seqlens_from_mask, @@ -26,13 +27,14 @@ from .index import ( prepare_sequence_ids, prepare_token_indices, ) -from .logsumexp import logsumexp_fwd -from .matmul import addmm, matmul -from .pack import pack_sequence, unpack_sequence -from .pooling import mean_pooling -from .softmax import softmax_bwd, softmax_fwd -from .softplus import softplus -from .solve_tril import solve_tril +from ...ops.utils.logsumexp import logsumexp_fwd +from ...ops.utils.matmul import addmm, matmul +from ...ops.utils.pack import pack_sequence, unpack_sequence +from ...ops.utils.pooling import mean_pooling +from ...ops.utils.sinkhorn import sinkhorn +from ...ops.utils.softmax import softmax_bwd, softmax_fwd +from ...ops.utils.softplus import softplus +from ...ops.utils.solve_tril import solve_tril __all__ = [ "addmm", @@ -49,6 +51,7 @@ __all__ = [ "pack_sequence", "prepare_block_csr", "prepare_chunk_indices", + "prepare_chunk_indices_static", "prepare_chunk_offsets", "prepare_cu_seqlens_from_lens", "prepare_cu_seqlens_from_mask", @@ -57,9 +60,11 @@ __all__ = [ "prepare_position_ids", "prepare_sequence_ids", "prepare_token_indices", + "sinkhorn", "softmax_bwd", "softmax_fwd", "softplus", "solve_tril", "unpack_sequence", ] + diff --git a/build/torch-cuda/ops/utils/backends/__init__.py b/build/torch-cuda/ops/utils/backends/__init__.py deleted file mode 100644 index 4ac07c0078d87c3792b0fce8d230c15062c9ab1b..0000000000000000000000000000000000000000 --- a/build/torch-cuda/ops/utils/backends/__init__.py +++ /dev/null @@ -1,16 +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 - -"""Utils op backends.""" - -from ....ops.backends import BackendRegistry, dispatch -from ....ops.utils.backends.triton_ascend import TritonAscendUtilsBackend - -utils_registry = BackendRegistry('utils') -utils_registry.register(TritonAscendUtilsBackend()) - -__all__ = ['dispatch', 'utils_registry'] diff --git a/build/torch-cuda/ops/utils/backends/triton_ascend/__init__.py b/build/torch-cuda/ops/utils/backends/triton_ascend/__init__.py deleted file mode 100644 index e7f66dd814b9a23286971b95d6acf10f19c9459d..0000000000000000000000000000000000000000 --- a/build/torch-cuda/ops/utils/backends/triton_ascend/__init__.py +++ /dev/null @@ -1,38 +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 NPU backend for utils ops.""" - -from __future__ import annotations - -from .....ops.backends import BaseBackend - - -class TritonAscendUtilsBackend(BaseBackend): - 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 chunk_global_cumsum_verifier(self, *args, **kwargs): - return True, None - - def chunk_global_cumsum(self, *args, **kwargs): - from .cumsum import chunk_global_cumsum_npu - return chunk_global_cumsum_npu(*args, **kwargs) - - def chunk_local_cumsum_verifier(self, *args, **kwargs): - return True, None - - def chunk_local_cumsum(self, *args, **kwargs): - from .cumsum import chunk_local_cumsum_npu - return chunk_local_cumsum_npu(*args, **kwargs) diff --git a/build/torch-cuda/ops/utils/backends/triton_ascend/cumsum.py b/build/torch-cuda/ops/utils/backends/triton_ascend/cumsum.py deleted file mode 100644 index 514b413500fec396b63546b39772b3aa2596009d..0000000000000000000000000000000000000000 --- a/build/torch-cuda/ops/utils/backends/triton_ascend/cumsum.py +++ /dev/null @@ -1,631 +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 torch -import triton -import triton.language as tl - -from .....ops.utils.index import prepare_chunk_indices -from .....utils import input_guard -from .....utils.ascend_ub_manager import ( - ASCEND_MAX_GRID_DIM, - compute_grid_limited_tile_size, - compute_row_tile_block_size, - compute_ub_block_size, - iter_axis_launch_chunks, - max_grid_axis_chunks, -) - -_NUM_WARPS = 4 -# Peak live fp32 tiles: b_s, b_o (and b_z / partial sums for vector/global). -_CUMSUM_SCALAR_MEM_MULT = 3.0 -# b_s, b_c, b_z, plus tl.cumsum multi-buffer on [BT, BS] tiles. -_CUMSUM_VECTOR_MEM_MULT = 8.0 -_CUMSUM_SAFETY_MARGIN = 0.85 -_FALLBACK_BT_GLOBAL = 32 -_FALLBACK_BS_LOCAL = 32 -_FALLBACK_BS_GLOBAL = 16 -_MAX_BT_GLOBAL = 256 - - -def _get_global_scalar_bt(T: int) -> int: - """UB-safe time chunk size for global scalar cumsum.""" - desired = min(triton.next_power_of_2(T), _MAX_BT_GLOBAL) - return compute_ub_block_size( - T, - _CUMSUM_SCALAR_MEM_MULT, - safety_margin=_CUMSUM_SAFETY_MARGIN, - dtype_size=4, - fallback=_FALLBACK_BT_GLOBAL, - desired=desired, - ) - - -def _get_vector_bs(BT: int, S: int, *, fallback: int, max_block: int | None = None) -> int: - """UB-safe feature tile size for vector cumsum with fixed time chunk BT.""" - return compute_row_tile_block_size( - BT, - S, - _CUMSUM_VECTOR_MEM_MULT, - tiling_row=False, - safety_margin=_CUMSUM_SAFETY_MARGIN, - dtype_size=4, - fallback=fallback, - min_block=1, - max_block=max_block, - ) - - -def _get_global_vector_tile_config(T: int, S: int) -> tuple[int, int]: - """UB-safe (BT, BS) for global vector cumsum.""" - bs_cap = min(_FALLBACK_BS_GLOBAL, triton.next_power_of_2(S)) - BS = _get_vector_bs( - _FALLBACK_BT_GLOBAL, - S, - fallback=_FALLBACK_BS_GLOBAL, - max_block=bs_cap, - ) - BT = compute_row_tile_block_size( - T, - BS, - _CUMSUM_VECTOR_MEM_MULT, - tiling_row=True, - safety_margin=_CUMSUM_SAFETY_MARGIN, - dtype_size=4, - fallback=_FALLBACK_BT_GLOBAL, - min_block=1, - max_block=_FALLBACK_BT_GLOBAL, - ) - BS = _get_vector_bs(BT, S, fallback=_FALLBACK_BS_GLOBAL, max_block=bs_cap) - return BT, BS - - -def _launch_local_cumsum_scalar( - *, - g_org, - g, - scale, - cu_seqlens, - chunk_indices, - T, - B, - H, - BT, - NT, - head_first, - reverse, -): - bh_total = B * H - kernel_kwargs = dict( - s=g_org, - o=g, - scale=scale, - cu_seqlens=cu_seqlens, - T=T, - B=B, - H=H, - BT=BT, - HEAD_FIRST=head_first, - REVERSE=reverse, - num_warps=_NUM_WARPS, - ) - max_nt = max_grid_axis_chunks(NT, bh_total, max_grid=ASCEND_MAX_GRID_DIM) - for nt_off in range(0, NT, max_nt): - nt_len = min(max_nt, NT - nt_off) - if cu_seqlens is not None: - 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 - max_bh = max_grid_axis_chunks(bh_total, nt_len, max_grid=ASCEND_MAX_GRID_DIM) - for bh_off in range(0, bh_total, max_bh): - bh_len = min(max_bh, bh_total - bh_off) - kernel_kwargs['BH_OFFSET'] = bh_off - chunk_local_cumsum_scalar_kernel_npu[(nt_len, bh_len)](**kernel_kwargs) - - -def _launch_local_cumsum_vector( - *, - g_org, - g, - scale, - cu_seqlens, - chunk_indices, - T, - B, - H, - S, - BT, - BS, - NT, - head_first, - reverse, -): - bh_total = B * H - ns = triton.cdiv(S, BS) - kernel_kwargs = dict( - s=g_org, - o=g, - scale=scale, - cu_seqlens=cu_seqlens, - T=T, - B=B, - H=H, - S=S, - BT=BT, - BS=BS, - HEAD_FIRST=head_first, - REVERSE=reverse, - num_warps=_NUM_WARPS, - ) - max_nt = max_grid_axis_chunks(NT, ns * bh_total, max_grid=ASCEND_MAX_GRID_DIM) - for nt_off in range(0, NT, max_nt): - nt_len = min(max_nt, NT - nt_off) - if cu_seqlens is not None: - 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 - max_bh = max_grid_axis_chunks(bh_total, ns * nt_len, max_grid=ASCEND_MAX_GRID_DIM) - for bh_off in range(0, bh_total, max_bh): - bh_len = min(max_bh, bh_total - bh_off) - kernel_kwargs['BH_OFFSET'] = bh_off - chunk_local_cumsum_vector_kernel_npu[(ns, nt_len, bh_len)](**kernel_kwargs) - - -@triton.heuristics({ - 'HAS_SCALE': lambda args: args['scale'] is not None, - 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None, -}) -@triton.jit(do_not_specialize=['T']) -def chunk_local_cumsum_scalar_kernel_npu( - s, - o, - scale, - cu_seqlens, - chunk_indices, - T, - B: tl.constexpr, - H: tl.constexpr, - BT: tl.constexpr, - REVERSE: tl.constexpr, - HAS_SCALE: tl.constexpr, - IS_VARLEN: tl.constexpr, - HEAD_FIRST: tl.constexpr, - NT_OFFSET: tl.constexpr, - BH_OFFSET: tl.constexpr, -): - i_t, i_bh = tl.program_id(0), tl.program_id(1) - i_t += NT_OFFSET - i_bh += BH_OFFSET - i_b, i_h = i_bh // H, i_bh % H - 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.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) - T = eos - bos - else: - bos, eos = i_b * T, i_b * T + T - - if HEAD_FIRST: - p_s = tl.make_block_ptr(s + bos*H + i_h*T, (T,), (1,), (i_t * BT,), (BT,), (0,)) - p_o = tl.make_block_ptr(o + bos*H + i_h*T, (T,), (1,), (i_t * BT,), (BT,), (0,)) - else: - p_s = tl.make_block_ptr(s + bos*H + i_h, (T,), (H,), (i_t * BT,), (BT,), (0,)) - p_o = tl.make_block_ptr(o + bos*H + i_h, (T,), (H,), (i_t * BT,), (BT,), (0,)) - b_s = tl.load(p_s, boundary_check=(0,)).to(tl.float32) - b_o = tl.cumsum(b_s, axis=0) - if REVERSE: - b_z = tl.sum(b_s, axis=0) - b_o = -b_o + b_z[None] + b_s - if HAS_SCALE: - b_o *= scale - tl.store(p_o, b_o.to(p_o.dtype.element_ty), boundary_check=(0,)) - - -@triton.heuristics({ - 'HAS_SCALE': lambda args: args['scale'] is not None, - 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None, -}) -@triton.jit(do_not_specialize=['T']) -def chunk_local_cumsum_vector_kernel_npu( - s, - o, - scale, - cu_seqlens, - chunk_indices, - T, - B: tl.constexpr, - H: tl.constexpr, - S: tl.constexpr, - BT: tl.constexpr, - BS: tl.constexpr, - REVERSE: tl.constexpr, - HAS_SCALE: tl.constexpr, - IS_VARLEN: tl.constexpr, - HEAD_FIRST: tl.constexpr, - NT_OFFSET: tl.constexpr, - BH_OFFSET: tl.constexpr, -): - i_s, i_t, i_bh = tl.program_id(0), tl.program_id(1), tl.program_id(2) - i_t += NT_OFFSET - i_bh += BH_OFFSET - i_b, i_h = i_bh // H, i_bh % H - 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.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) - T = eos - bos - else: - bos, eos = i_b * T, i_b * T + T - - if HEAD_FIRST: - p_s = tl.make_block_ptr(s + (bos * H + i_h*T)*S, (T, S), (S, 1), (i_t * BT, i_s * BS), (BT, BS), (1, 0)) - p_o = tl.make_block_ptr(o + (bos * H + i_h*T)*S, (T, S), (S, 1), (i_t * BT, i_s * BS), (BT, BS), (1, 0)) - else: - p_s = tl.make_block_ptr(s + (bos * H + i_h) * S, (T, S), (H*S, 1), (i_t * BT, i_s * BS), (BT, BS), (1, 0)) - p_o = tl.make_block_ptr(o + (bos * H + i_h) * S, (T, S), (H*S, 1), (i_t * BT, i_s * BS), (BT, BS), (1, 0)) - b_s = tl.load(p_s, boundary_check=(0, 1)).to(tl.float32) - if REVERSE: - b_o = tl.cumsum(b_s, axis=0, reverse=True) - else: - b_o = tl.cumsum(b_s, axis=0) - if HAS_SCALE: - b_o *= scale - tl.store(p_o, b_o.to(p_o.dtype.element_ty), boundary_check=(0, 1)) - - -@triton.heuristics({ - 'HAS_SCALE': lambda args: args['scale'] is not None, - 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None, -}) -@triton.jit(do_not_specialize=['T']) -def chunk_global_cumsum_scalar_kernel_npu( - s, - o, - scale, - cu_seqlens, - T, - B: tl.constexpr, - H: tl.constexpr, - BT: tl.constexpr, - REVERSE: tl.constexpr, - HAS_SCALE: tl.constexpr, - IS_VARLEN: tl.constexpr, - HEAD_FIRST: tl.constexpr, - BH_OFFSET: tl.constexpr, -): - i_nh = tl.program_id(0) + BH_OFFSET - i_n, i_h = i_nh // H, i_nh % H - if IS_VARLEN: - bos, eos = tl.load(cu_seqlens + i_n).to(tl.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) - else: - bos, eos = i_n * T, i_n * T + T - T = eos - bos - - b_z = tl.zeros([], dtype=tl.float32) - NT = tl.cdiv(T, BT) - for i_c in range(NT): - i_t = NT - 1 - i_c if REVERSE else i_c - if HEAD_FIRST: - p_s = tl.make_block_ptr(s + bos*H + i_h*T, (T,), (1,), (i_t * BT,), (BT,), (0,)) - p_o = tl.make_block_ptr(o + bos*H + i_h*T, (T,), (1,), (i_t * BT,), (BT,), (0,)) - else: - p_s = tl.make_block_ptr(s + bos*H + i_h, (T,), (H,), (i_t * BT,), (BT,), (0,)) - p_o = tl.make_block_ptr(o + bos*H + i_h, (T,), (H,), (i_t * BT,), (BT,), (0,)) - b_s = tl.load(p_s, boundary_check=(0,)).to(tl.float32) - b_o = tl.cumsum(b_s, axis=0) - b_ss = tl.sum(b_s, 0) - if REVERSE: - b_o = -b_o + b_ss + b_s - b_o += b_z - if i_c >= 0: - b_z += b_ss - if HAS_SCALE: - b_o *= scale - tl.store(p_o, b_o.to(p_o.dtype.element_ty), boundary_check=(0,)) - - -@triton.heuristics({ - 'HAS_SCALE': lambda args: args['scale'] is not None, - 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None, -}) -@triton.jit(do_not_specialize=['T']) -def chunk_global_cumsum_vector_kernel_npu( - s, - o, - scale, - cu_seqlens, - T, - B: tl.constexpr, - H: tl.constexpr, - S: tl.constexpr, - BT: tl.constexpr, - BS: tl.constexpr, - REVERSE: tl.constexpr, - HAS_SCALE: tl.constexpr, - IS_VARLEN: tl.constexpr, - HEAD_FIRST: tl.constexpr, - BH_OFFSET: tl.constexpr, -): - i_s, i_nh = tl.program_id(0), tl.program_id(1) + BH_OFFSET - i_n, i_h = i_nh // H, i_nh % H - if IS_VARLEN: - bos, eos = tl.load(cu_seqlens + i_n).to(tl.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) - else: - bos, eos = i_n * T, i_n * T + T - T = eos - bos - - b_z = tl.zeros([BS], dtype=tl.float32) - NT = tl.cdiv(T, BT) - for i_c in range(NT): - i_t = NT - 1 - i_c if REVERSE else i_c - if HEAD_FIRST: - p_s = tl.make_block_ptr(s + (bos * H + i_h*T)*S, (T, S), (S, 1), (i_t * BT, i_s * BS), (BT, BS), (1, 0)) - p_o = tl.make_block_ptr(o + (bos * H + i_h*T)*S, (T, S), (S, 1), (i_t * BT, i_s * BS), (BT, BS), (1, 0)) - else: - p_s = tl.make_block_ptr(s + (bos * H + i_h) * S, (T, S), (H*S, 1), (i_t * BT, i_s * BS), (BT, BS), (1, 0)) - p_o = tl.make_block_ptr(o + (bos * H + i_h) * S, (T, S), (H*S, 1), (i_t * BT, i_s * BS), (BT, BS), (1, 0)) - b_s = tl.load(p_s, boundary_check=(0, 1)).to(tl.float32) - if REVERSE: - b_c = b_z[None, :] + tl.cumsum(b_s, axis=0, reverse=True) - else: - b_c = b_z[None, :] + tl.cumsum(b_s, axis=0) - if HAS_SCALE: - b_c *= scale - tl.store(p_o, b_c.to(p_o.dtype.element_ty), boundary_check=(0, 1)) - b_z += tl.sum(b_s, 0) - - -def chunk_local_cumsum_scalar_npu( - g: torch.Tensor, - chunk_size: int, - reverse: bool = False, - scale: float = None, - cu_seqlens: torch.Tensor | None = None, - head_first: bool = False, - output_dtype: torch.dtype | None = torch.float, - chunk_indices: torch.LongTensor | None = None, -) -> torch.Tensor: - if head_first: - B, H, T = g.shape - else: - B, T, H = g.shape - assert chunk_size == 2**(chunk_size.bit_length()-1), "chunk_size must be a power of 2" - BT = chunk_size - if chunk_indices is None and cu_seqlens is not None: - chunk_indices = prepare_chunk_indices(cu_seqlens, BT) - NT = triton.cdiv(T, BT) if cu_seqlens is None else len(chunk_indices) - g_org, g = g, torch.empty_like(g, dtype=output_dtype or g.dtype) - _launch_local_cumsum_scalar( - g_org=g_org, - g=g, - scale=scale, - cu_seqlens=cu_seqlens, - chunk_indices=chunk_indices, - T=T, - B=B, - H=H, - BT=BT, - NT=NT, - head_first=head_first, - reverse=reverse, - ) - return g - - -def chunk_local_cumsum_vector_npu( - g: torch.Tensor, - chunk_size: int, - reverse: bool = False, - scale: float = None, - cu_seqlens: torch.Tensor | None = None, - head_first: bool = False, - output_dtype: torch.dtype | None = torch.float, - chunk_indices: torch.LongTensor | None = None, -) -> torch.Tensor: - if head_first: - B, H, T, S = g.shape - else: - B, T, H, S = g.shape - BT = chunk_size - if chunk_indices is None and cu_seqlens is not None: - chunk_indices = prepare_chunk_indices(cu_seqlens, BT) - NT = triton.cdiv(T, BT) if cu_seqlens is None else len(chunk_indices) - assert chunk_size == 2**(chunk_size.bit_length()-1), "chunk_size must be a power of 2" - - BS = _get_vector_bs(BT, S, fallback=_FALLBACK_BS_LOCAL) - BS = compute_grid_limited_tile_size( - S, - NT * B * H, - BS, - max_grid=ASCEND_MAX_GRID_DIM, - ) - g_org, g = g, torch.empty_like(g, dtype=output_dtype or g.dtype) - _launch_local_cumsum_vector( - g_org=g_org, - g=g, - scale=scale, - cu_seqlens=cu_seqlens, - chunk_indices=chunk_indices, - T=T, - B=B, - H=H, - S=S, - BT=BT, - BS=BS, - NT=NT, - head_first=head_first, - reverse=reverse, - ) - return g - - -@input_guard -def chunk_global_cumsum_scalar_npu( - s: torch.Tensor, - reverse: bool = False, - cu_seqlens: torch.Tensor | None = None, - scale: float = None, - head_first: bool = False, - output_dtype: torch.dtype | None = torch.float, -) -> torch.Tensor: - if head_first: - B, H, T = s.shape - else: - B, T, H = s.shape - N = len(cu_seqlens) - 1 if cu_seqlens is not None else B - - BT = _get_global_scalar_bt(T) - z = torch.empty_like(s, dtype=output_dtype or s.dtype) - bh_total = N * H - kernel_kwargs = dict( - s=s, - o=z, - scale=scale, - cu_seqlens=cu_seqlens, - T=T, - B=B, - H=H, - BT=BT, - HEAD_FIRST=head_first, - REVERSE=reverse, - num_warps=_NUM_WARPS, - ) - for bh_off, bh_len in iter_axis_launch_chunks(bh_total, 1, max_grid=ASCEND_MAX_GRID_DIM): - kernel_kwargs['BH_OFFSET'] = bh_off - chunk_global_cumsum_scalar_kernel_npu[(bh_len,)](**kernel_kwargs) - return z - - -@input_guard -def chunk_global_cumsum_vector_npu( - s: torch.Tensor, - reverse: bool = False, - cu_seqlens: torch.Tensor | None = None, - scale: float = None, - head_first: bool = False, - output_dtype: torch.dtype | None = torch.float, -) -> torch.Tensor: - if head_first: - B, H, T, S = s.shape - else: - B, T, H, S = s.shape - N = len(cu_seqlens) - 1 if cu_seqlens is not None else B - BT, BS = _get_global_vector_tile_config(T, S) - BS = compute_grid_limited_tile_size( - S, - N * H, - BS, - max_grid=ASCEND_MAX_GRID_DIM, - ) - ns = triton.cdiv(S, BS) - - z = torch.empty_like(s, dtype=output_dtype or s.dtype) - bh_total = N * H - kernel_kwargs = dict( - s=s, - o=z, - scale=scale, - cu_seqlens=cu_seqlens, - T=T, - B=B, - H=H, - S=S, - BT=BT, - BS=BS, - HEAD_FIRST=head_first, - REVERSE=reverse, - num_warps=_NUM_WARPS, - ) - max_bh = max_grid_axis_chunks(bh_total, ns, max_grid=ASCEND_MAX_GRID_DIM) - for bh_off in range(0, bh_total, max_bh): - bh_len = min(max_bh, bh_total - bh_off) - kernel_kwargs['BH_OFFSET'] = bh_off - chunk_global_cumsum_vector_kernel_npu[(ns, bh_len)](**kernel_kwargs) - return z - - -@input_guard -def chunk_global_cumsum_npu( - s: torch.Tensor, - reverse: bool = False, - cu_seqlens: torch.Tensor | None = None, - scale: float = None, - head_first: bool = False, - output_dtype: torch.dtype | None = torch.float, -) -> torch.Tensor: - if cu_seqlens is not None: - assert s.shape[0] == 1, "Only batch size 1 is supported when cu_seqlens are provided" - if len(s.shape) == 3: - return chunk_global_cumsum_scalar_npu( - s=s, - reverse=reverse, - cu_seqlens=cu_seqlens, - scale=scale, - head_first=head_first, - output_dtype=output_dtype, - ) - if len(s.shape) == 4: - return chunk_global_cumsum_vector_npu( - s=s, - reverse=reverse, - cu_seqlens=cu_seqlens, - scale=scale, - head_first=head_first, - output_dtype=output_dtype, - ) - raise ValueError( - f"Unsupported input shape {s.shape}, " - f"which should be [B, T, H]/[B, T, H, D] if `head_first=False` " - f"or [B, H, T]/[B, H, T, D] otherwise", - ) - - -@input_guard -def chunk_local_cumsum_npu( - g: torch.Tensor, - chunk_size: int, - reverse: bool = False, - scale: float = None, - cu_seqlens: torch.Tensor | None = None, - head_first: bool = False, - output_dtype: torch.dtype | None = torch.float, - chunk_indices: torch.LongTensor | None = None, - **kwargs, -) -> torch.Tensor: - if cu_seqlens is not None: - assert g.shape[0] == 1, "Only batch size 1 is supported when cu_seqlens are provided" - if len(g.shape) == 3: - return chunk_local_cumsum_scalar_npu( - g=g, - chunk_size=chunk_size, - reverse=reverse, - scale=scale, - cu_seqlens=cu_seqlens, - head_first=head_first, - output_dtype=output_dtype, - chunk_indices=chunk_indices, - ) - if len(g.shape) == 4: - return chunk_local_cumsum_vector_npu( - g=g, - chunk_size=chunk_size, - reverse=reverse, - scale=scale, - cu_seqlens=cu_seqlens, - head_first=head_first, - output_dtype=output_dtype, - chunk_indices=chunk_indices, - ) - raise ValueError( - f"Unsupported input shape {g.shape}, " - f"which should be (B, T, H, D) if `head_first=False` " - f"or (B, H, T, D) otherwise", - ) diff --git a/build/torch-cuda/ops/utils/cache.py b/build/torch-cuda/ops/utils/cache.py index 40ab2dbc1802212f1e3e3ef4c97c2bc3ac79148f..b172b686ac8e86297bfcd8bc22e69c5d70c10382 100644 --- a/build/torch-cuda/ops/utils/cache.py +++ b/build/torch-cuda/ops/utils/cache.py @@ -374,10 +374,10 @@ class CachedAutotuner(Autotuner): def run(self, *args, **kwargs): key = AutotuneKey.build(self.arg_names, self.keys, args, kwargs) if self.should_check_fla_cache(key): - self.maybe_load_cached_config(key) + self.maybe_load_cached_config(key=key, nargs=dict(zip(self.arg_names, args)), runtime_kwargs=kwargs) return super().run(*args, **kwargs) - def maybe_load_cached_config(self, key: AutotuneKey): + def maybe_load_cached_config(self, key: AutotuneKey, nargs: dict[str, Any], runtime_kwargs: dict[str, Any]): best_config = load_cached_config(self.kernel_name, key) if best_config is not None: @@ -393,6 +393,16 @@ class CachedAutotuner(Autotuner): } if TRITON_ABOVE_3_5_1 else {} cfg = triton.Config(kw, num_warps=num_warps, num_stages=num_stages, **extra) + if self.early_config_prune is not None: + configs = self.early_config_prune([*self.configs, cfg], nargs, **runtime_kwargs) + if cfg not in configs: + self.cache.pop(key.autotune_key, None) + logger.debug( + "Cached config rejected for kernel %s and key %s; falling back to Triton autotune", + self.kernel_name, + list(key.autotune_key), + ) + return self.cache[key.autotune_key] = cfg else: logger.debug( @@ -411,7 +421,7 @@ def fla_cache_autotune(configs, key=None, prune_configs_by=None, reset_to_zero=N Extends Triton's autotune to load best configurations from FLA's config directory (default: fla/configs/{GPU}/, or FLA_CONFIG_DIR/ when overridden), keyed by kernel name from {kernel_name}.json. Lookup behaviour is controlled by FLA_CACHE_MODE. - Falls back to normal Triton autotuning when no cached config is found. + Falls back to normal Triton autotuning when a cached config is absent or rejected by early pruning. """ # key can be None when we want to use cache only (no fallback autotune) if key is None: diff --git a/build/torch-cuda/ops/utils/csr.py b/build/torch-cuda/ops/utils/csr.py index dab4cbab96c252ca1700ea4ded462c1298ffbdb1..cec2f7bc66600cfd06449ff20be67d2771eec9c0 100644 --- a/build/torch-cuda/ops/utils/csr.py +++ b/build/torch-cuda/ops/utils/csr.py @@ -36,7 +36,7 @@ def prepare_block_csr_kernel( USE_BLOCK_COUNTS: tl.constexpr, IS_VARLEN: tl.constexpr, ): - i_t, i_bh = tl.program_id(0), tl.program_id(1) + i_t, i_bh = tl.program_id(0).to(tl.int64), tl.program_id(1).to(tl.int64) i_b, i_h = i_bh // H, i_bh % H o_t = i_t * BT + tl.arange(0, BT) o_s = tl.arange(0, S) diff --git a/build/torch-cuda/ops/utils/cumsum.py b/build/torch-cuda/ops/utils/cumsum.py index 6140e860c8e1b17951b7a1df47cf9a9b93af4b25..bf8059ebe01f06efdca1074de001790ef2011676 100644 --- a/build/torch-cuda/ops/utils/cumsum.py +++ b/build/torch-cuda/ops/utils/cumsum.py @@ -9,9 +9,10 @@ import torch import triton import triton.language as tl -from ...ops.backends import dispatch +from ...backends import dispatch from ...ops.utils.cache import fla_cache_autotune from ...ops.utils.index import prepare_chunk_indices +from ...ops.utils.op import unflatten_program_id from ...utils import autotune_cache_kwargs, check_shared_mem, input_guard BS_LIST = [32, 64] if check_shared_mem() else [16, 32] @@ -43,32 +44,33 @@ def chunk_local_cumsum_scalar_kernel( REVERSE: tl.constexpr, HAS_SCALE: tl.constexpr, IS_VARLEN: tl.constexpr, - HEAD_FIRST: tl.constexpr, + USE_GRAPH: tl.constexpr = False, ): - i_t, i_bh = tl.program_id(0), tl.program_id(1) + i_bh, i_t = unflatten_program_id(B * H) i_b, i_h = i_bh // H, i_bh % H 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.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) + i_n = tl.load(chunk_indices + i_t * 2).to(tl.int32) + if USE_GRAPH and i_n < 0: + return + i_t = tl.load(chunk_indices + i_t * 2 + 1).to(tl.int64) + bos, eos = tl.load(cu_seqlens + i_n).to(tl.int64), tl.load(cu_seqlens + i_n + 1).to(tl.int64) T = eos - bos else: bos, eos = i_b * T, i_b * T + T - if HEAD_FIRST: - p_s = tl.make_block_ptr(s + bos*H + i_h*T, (T,), (1,), (i_t * BT,), (BT,), (0,)) - p_o = tl.make_block_ptr(o + bos*H + i_h*T, (T,), (1,), (i_t * BT,), (BT,), (0,)) - else: - p_s = tl.make_block_ptr(s + bos*H + i_h, (T,), (H,), (i_t * BT,), (BT,), (0,)) - p_o = tl.make_block_ptr(o + bos*H + i_h, (T,), (H,), (i_t * BT,), (BT,), (0,)) + o_t = i_t * BT + tl.arange(0, BT) + m_t = o_t < T + p_s = s + bos*H + i_h + o_t * H + p_o = o + bos*H + i_h + o_t * H # [BT] - b_s = tl.load(p_s, boundary_check=(0,)).to(tl.float32) - b_o = tl.cumsum(b_s, axis=0) + b_s = tl.load(p_s, mask=m_t, other=0.0).to(tl.float32) if REVERSE: - b_z = tl.sum(b_s, axis=0) - b_o = -b_o + b_z[None] + b_s + b_o = tl.cumsum(b_s, axis=0, reverse=True) + else: + b_o = tl.cumsum(b_s, axis=0) if HAS_SCALE: b_o *= scale - tl.store(p_o, b_o.to(p_o.dtype.element_ty), boundary_check=(0,)) + tl.store(p_o, b_o.to(p_o.dtype.element_ty), mask=m_t) @triton.heuristics({ @@ -100,32 +102,35 @@ def chunk_local_cumsum_vector_kernel( REVERSE: tl.constexpr, HAS_SCALE: tl.constexpr, IS_VARLEN: tl.constexpr, - HEAD_FIRST: tl.constexpr, + USE_GRAPH: tl.constexpr = False, ): - i_s, i_t, i_bh = tl.program_id(0), tl.program_id(1), tl.program_id(2) + i_s, i_t = unflatten_program_id(tl.cdiv(S, BS)) + i_bh = tl.program_id(1).to(tl.int64) i_b, i_h = i_bh // H, i_bh % H 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.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) + i_n = tl.load(chunk_indices + i_t * 2).to(tl.int32) + if USE_GRAPH and i_n < 0: + return + i_t = tl.load(chunk_indices + i_t * 2 + 1).to(tl.int64) + bos, eos = tl.load(cu_seqlens + i_n).to(tl.int64), tl.load(cu_seqlens + i_n + 1).to(tl.int64) T = eos - bos else: bos, eos = i_b * T, i_b * T + T - if HEAD_FIRST: - p_s = tl.make_block_ptr(s + (bos * H + i_h*T)*S, (T, S), (S, 1), (i_t * BT, i_s * BS), (BT, BS), (1, 0)) - p_o = tl.make_block_ptr(o + (bos * H + i_h*T)*S, (T, S), (S, 1), (i_t * BT, i_s * BS), (BT, BS), (1, 0)) - else: - p_s = tl.make_block_ptr(s + (bos * H + i_h) * S, (T, S), (H*S, 1), (i_t * BT, i_s * BS), (BT, BS), (1, 0)) - p_o = tl.make_block_ptr(o + (bos * H + i_h) * S, (T, S), (H*S, 1), (i_t * BT, i_s * BS), (BT, BS), (1, 0)) + o_t = i_t * BT + tl.arange(0, BT) + o_s = i_s * BS + tl.arange(0, BS) + m_s = (o_t[:, None] < T) & (o_s[None, :] < S) + p_s = s + (bos * H + i_h) * S + o_t[:, None] * (H*S) + o_s[None, :] + p_o = o + (bos * H + i_h) * S + o_t[:, None] * (H*S) + o_s[None, :] # [BT, BS] - b_s = tl.load(p_s, boundary_check=(0, 1)).to(tl.float32) + b_s = tl.load(p_s, mask=m_s, other=0.0).to(tl.float32) if REVERSE: b_o = tl.cumsum(b_s, axis=0, reverse=True) else: b_o = tl.cumsum(b_s, axis=0) if HAS_SCALE: b_o *= scale - tl.store(p_o, b_o.to(p_o.dtype.element_ty), boundary_check=(0, 1)) + tl.store(p_o, b_o.to(p_o.dtype.element_ty), mask=m_s) @triton.heuristics({ @@ -155,12 +160,11 @@ def chunk_global_cumsum_scalar_kernel( REVERSE: tl.constexpr, HAS_SCALE: tl.constexpr, IS_VARLEN: tl.constexpr, - HEAD_FIRST: tl.constexpr, ): - i_nh = tl.program_id(0) + i_nh = tl.program_id(0).to(tl.int64) i_n, i_h = i_nh // H, i_nh % H if IS_VARLEN: - bos, eos = tl.load(cu_seqlens + i_n).to(tl.int32), tl.load(cu_seqlens + i_n + 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) else: bos, eos = i_n * T, i_n * T + T T = eos - bos @@ -169,23 +173,22 @@ def chunk_global_cumsum_scalar_kernel( NT = tl.cdiv(T, BT) for i_c in range(NT): i_t = NT - 1 - i_c if REVERSE else i_c - if HEAD_FIRST: - p_s = tl.make_block_ptr(s + bos*H + i_h*T, (T,), (1,), (i_t * BT,), (BT,), (0,)) - p_o = tl.make_block_ptr(o + bos*H + i_h*T, (T,), (1,), (i_t * BT,), (BT,), (0,)) + o_t = i_t * BT + tl.arange(0, BT) + m_t = o_t < T + p_s = s + bos*H + i_h + o_t * H + p_o = o + bos*H + i_h + o_t * H + b_s = tl.load(p_s, mask=m_t, other=0.0).to(tl.float32) + if REVERSE: + b_o = tl.cumsum(b_s, axis=0, reverse=True) else: - p_s = tl.make_block_ptr(s + bos*H + i_h, (T,), (H,), (i_t * BT,), (BT,), (0,)) - p_o = tl.make_block_ptr(o + bos*H + i_h, (T,), (H,), (i_t * BT,), (BT,), (0,)) - b_s = tl.load(p_s, boundary_check=(0,)).to(tl.float32) - b_o = tl.cumsum(b_s, axis=0) + b_o = tl.cumsum(b_s, axis=0) b_ss = tl.sum(b_s, 0) - if REVERSE: - b_o = -b_o + b_ss + b_s b_o += b_z if i_c >= 0: b_z += b_ss if HAS_SCALE: b_o *= scale - tl.store(p_o, b_o.to(p_o.dtype.element_ty), boundary_check=(0,)) + tl.store(p_o, b_o.to(p_o.dtype.element_ty), mask=m_t) @triton.heuristics({ @@ -217,12 +220,11 @@ def chunk_global_cumsum_vector_kernel( REVERSE: tl.constexpr, HAS_SCALE: tl.constexpr, IS_VARLEN: tl.constexpr, - HEAD_FIRST: tl.constexpr, ): - i_s, i_nh = tl.program_id(0), tl.program_id(1) + i_s, i_nh = unflatten_program_id(tl.cdiv(S, BS)) i_n, i_h = i_nh // H, i_nh % H if IS_VARLEN: - bos, eos = tl.load(cu_seqlens + i_n).to(tl.int32), tl.load(cu_seqlens + i_n + 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) else: bos, eos = i_n * T, i_n * T + T T = eos - bos @@ -231,21 +233,20 @@ def chunk_global_cumsum_vector_kernel( NT = tl.cdiv(T, BT) for i_c in range(NT): i_t = NT - 1 - i_c if REVERSE else i_c - if HEAD_FIRST: - p_s = tl.make_block_ptr(s + (bos * H + i_h*T)*S, (T, S), (S, 1), (i_t * BT, i_s * BS), (BT, BS), (1, 0)) - p_o = tl.make_block_ptr(o + (bos * H + i_h*T)*S, (T, S), (S, 1), (i_t * BT, i_s * BS), (BT, BS), (1, 0)) - else: - p_s = tl.make_block_ptr(s + (bos * H + i_h) * S, (T, S), (H*S, 1), (i_t * BT, i_s * BS), (BT, BS), (1, 0)) - p_o = tl.make_block_ptr(o + (bos * H + i_h) * S, (T, S), (H*S, 1), (i_t * BT, i_s * BS), (BT, BS), (1, 0)) + o_t = i_t * BT + tl.arange(0, BT) + o_s = i_s * BS + tl.arange(0, BS) + m_s = (o_t[:, None] < T) & (o_s[None, :] < S) + p_s = s + (bos * H + i_h) * S + o_t[:, None] * (H*S) + o_s[None, :] + p_o = o + (bos * H + i_h) * S + o_t[:, None] * (H*S) + o_s[None, :] # [BT, BS] - b_s = tl.load(p_s, boundary_check=(0, 1)).to(tl.float32) + b_s = tl.load(p_s, mask=m_s, other=0.0).to(tl.float32) if REVERSE: b_c = b_z[None, :] + tl.cumsum(b_s, axis=0, reverse=True) else: b_c = b_z[None, :] + tl.cumsum(b_s, axis=0) if HAS_SCALE: b_c *= scale - tl.store(p_o, b_c.to(p_o.dtype.element_ty), boundary_check=(0, 1)) + tl.store(p_o, b_c.to(p_o.dtype.element_ty), mask=m_s) b_z += tl.sum(b_s, 0) @@ -253,26 +254,32 @@ def chunk_local_cumsum_scalar( g: torch.Tensor, chunk_size: int, reverse: bool = False, - scale: float = None, + scale: float | None = None, cu_seqlens: torch.Tensor | None = None, - head_first: bool = False, output_dtype: torch.dtype | None = torch.float, chunk_indices: torch.LongTensor | None = None, + use_graph: bool = False, + **kwargs, ) -> torch.Tensor: - if head_first: - B, H, T = g.shape - else: - B, T, H = g.shape + if 'head_first' in kwargs: + raise DeprecationWarning( + "head_first has been removed. Inputs must be in `[B, T, H, ...]` format.", + ) + B, T, H = g.shape assert chunk_size == 2**(chunk_size.bit_length()-1), "chunk_size must be a power of 2" BT = chunk_size if chunk_indices is None and cu_seqlens is not None: chunk_indices = prepare_chunk_indices(cu_seqlens, BT) NT = triton.cdiv(T, BT) if cu_seqlens is None else len(chunk_indices) - g_org, g = g, torch.empty_like(g, dtype=output_dtype or g.dtype) - grid = (NT, B * H) + if use_graph: + # zero inactive rows: kda_gate_bwd reduces the entire output buffer + o = torch.zeros_like(g, dtype=output_dtype or g.dtype) + else: + o = torch.empty_like(g, dtype=output_dtype or g.dtype) + grid = (NT * B * H,) chunk_local_cumsum_scalar_kernel[grid]( - s=g_org, - o=g, + s=g, + o=o, scale=scale, cu_seqlens=cu_seqlens, chunk_indices=chunk_indices, @@ -280,40 +287,49 @@ def chunk_local_cumsum_scalar( B=B, H=H, BT=BT, - HEAD_FIRST=head_first, REVERSE=reverse, + USE_GRAPH=use_graph, ) - return g + return o def chunk_local_cumsum_vector( g: torch.Tensor, chunk_size: int, reverse: bool = False, - scale: float = None, + scale: float | None = None, cu_seqlens: torch.Tensor | None = None, - head_first: bool = False, output_dtype: torch.dtype | None = torch.float, chunk_indices: torch.LongTensor | None = None, + use_graph: bool = False, + **kwargs, ) -> torch.Tensor: - if head_first: - B, H, T, S = g.shape - else: - B, T, H, S = g.shape + if 'head_first' in kwargs: + raise DeprecationWarning( + "head_first has been removed. Inputs must be in `[B, T, H, ...]` format.", + ) + B, T, H, S = g.shape BT = chunk_size if chunk_indices is None and cu_seqlens is not None: chunk_indices = prepare_chunk_indices(cu_seqlens, BT) NT = triton.cdiv(T, BT) if cu_seqlens is None else len(chunk_indices) assert chunk_size == 2**(chunk_size.bit_length()-1), "chunk_size must be a power of 2" - g_org, g = g, torch.empty_like(g, dtype=output_dtype or g.dtype) - def grid(meta): return (triton.cdiv(meta['S'], meta['BS']), NT, B * H) - # keep cummulative normalizer in fp32 - # this kernel is equivalent to - # g = g.view(B, H, NT, BT, -1).cumsum(-2).view(B, H, T, -1) + if use_graph: + # zero inactive rows: kda_gate_bwd reduces the entire output buffer + o = torch.zeros_like(g, dtype=output_dtype or g.dtype) + else: + o = torch.empty_like(g, dtype=output_dtype or g.dtype) + + def grid(meta): + return (triton.cdiv(meta['S'], meta['BS']) * NT, B * H) + + # keep cumulative normalizer in fp32 + # this kernel is equivalent to the following for full dense chunks without scaling or reversal + # g.view(B, NT, BT, H, S).cumsum(2, dtype=torch.float32).view(B, T, H, S) chunk_local_cumsum_vector_kernel[grid]( - s=g_org, - o=g, + s=g, + o=o, scale=scale, cu_seqlens=cu_seqlens, chunk_indices=chunk_indices, @@ -322,10 +338,10 @@ def chunk_local_cumsum_vector( H=H, S=S, BT=BT, - HEAD_FIRST=head_first, REVERSE=reverse, + USE_GRAPH=use_graph, ) - return g + return o @input_guard @@ -333,14 +349,15 @@ def chunk_global_cumsum_scalar( s: torch.Tensor, reverse: bool = False, cu_seqlens: torch.Tensor | None = None, - scale: float = None, - head_first: bool = False, + scale: float | None = None, output_dtype: torch.dtype | None = torch.float, + **kwargs, ) -> torch.Tensor: - if head_first: - B, H, T = s.shape - else: - B, T, H = s.shape + if 'head_first' in kwargs: + raise DeprecationWarning( + "head_first has been removed. Inputs must be in `[B, T, H, ...]` format.", + ) + B, T, H = s.shape N = len(cu_seqlens) - 1 if cu_seqlens is not None else B z = torch.empty_like(s, dtype=output_dtype or s.dtype) @@ -353,7 +370,6 @@ def chunk_global_cumsum_scalar( T=T, B=B, H=H, - HEAD_FIRST=head_first, REVERSE=reverse, ) return z @@ -364,19 +380,20 @@ def chunk_global_cumsum_vector( s: torch.Tensor, reverse: bool = False, cu_seqlens: torch.Tensor | None = None, - scale: float = None, - head_first: bool = False, + scale: float | None = None, output_dtype: torch.dtype | None = torch.float, + **kwargs, ) -> torch.Tensor: - if head_first: - B, H, T, S = s.shape - else: - B, T, H, S = s.shape + if 'head_first' in kwargs: + raise DeprecationWarning( + "head_first has been removed. Inputs must be in `[B, T, H, ...]` format.", + ) + B, T, H, S = s.shape N = len(cu_seqlens) - 1 if cu_seqlens is not None else B BS = min(32, triton.next_power_of_2(S)) z = torch.empty_like(s, dtype=output_dtype or s.dtype) - grid = (triton.cdiv(S, BS), N * H) + grid = (triton.cdiv(S, BS) * N * H,) chunk_global_cumsum_vector_kernel[grid]( s=s, o=z, @@ -387,22 +404,26 @@ def chunk_global_cumsum_vector( H=H, S=S, BS=BS, - HEAD_FIRST=head_first, REVERSE=reverse, ) return z +# Use triton only, other backends (cutlass, cute, etc.) need additional torch.compile support +# @dispatch @input_guard -@dispatch('utils') def chunk_global_cumsum( s: torch.Tensor, reverse: bool = False, cu_seqlens: torch.Tensor | None = None, - scale: float = None, - head_first: bool = False, + scale: float | None = None, output_dtype: torch.dtype | None = torch.float, + **kwargs, ) -> torch.Tensor: + if 'head_first' in kwargs: + raise DeprecationWarning( + "head_first has been removed. Inputs must be in `[B, T, H, ...]` format.", + ) if cu_seqlens is not None: assert s.shape[0] == 1, "Only batch size 1 is supported when cu_seqlens are provided" if len(s.shape) == 3: @@ -411,7 +432,6 @@ def chunk_global_cumsum( reverse=reverse, cu_seqlens=cu_seqlens, scale=scale, - head_first=head_first, output_dtype=output_dtype, ) elif len(s.shape) == 4: @@ -420,30 +440,33 @@ def chunk_global_cumsum( reverse=reverse, cu_seqlens=cu_seqlens, scale=scale, - head_first=head_first, output_dtype=output_dtype, ) else: raise ValueError( f"Unsupported input shape {s.shape}, " - f"which should be [B, T, H]/[B, T, H, D] if `head_first=False` " - f"or [B, H, T]/[B, H, T, D] otherwise", + f"which should be [B, T, H] or [B, T, H, D]", ) +# Use triton only, other backends (cutlass, cute, etc.) need additional torch.compile support +# @dispatch @input_guard -@dispatch('utils') def chunk_local_cumsum( g: torch.Tensor, chunk_size: int, reverse: bool = False, - scale: float = None, + scale: float | None = None, cu_seqlens: torch.Tensor | None = None, - head_first: bool = False, output_dtype: torch.dtype | None = torch.float, chunk_indices: torch.LongTensor | None = None, + use_graph: bool = False, **kwargs, ) -> torch.Tensor: + if 'head_first' in kwargs: + raise DeprecationWarning( + "head_first has been removed. Inputs must be in `[B, T, H, ...]` format.", + ) if cu_seqlens is not None: assert g.shape[0] == 1, "Only batch size 1 is supported when cu_seqlens are provided" if len(g.shape) == 3: @@ -453,9 +476,9 @@ def chunk_local_cumsum( reverse=reverse, scale=scale, cu_seqlens=cu_seqlens, - head_first=head_first, output_dtype=output_dtype, chunk_indices=chunk_indices, + use_graph=use_graph, ) elif len(g.shape) == 4: return chunk_local_cumsum_vector( @@ -464,13 +487,12 @@ def chunk_local_cumsum( reverse=reverse, scale=scale, cu_seqlens=cu_seqlens, - head_first=head_first, output_dtype=output_dtype, chunk_indices=chunk_indices, + use_graph=use_graph, ) else: raise ValueError( f"Unsupported input shape {g.shape}, " - f"which should be (B, T, H, D) if `head_first=False` " - f"or (B, H, T, D) otherwise", + f"which should be (B, T, H) or (B, T, H, D)", ) diff --git a/build/torch-cuda/ops/utils/graph.py b/build/torch-cuda/ops/utils/graph.py new file mode 100644 index 0000000000000000000000000000000000000000..c5b252ad3c4f42febeb9948beb74ad9bd6bcb742 --- /dev/null +++ b/build/torch-cuda/ops/utils/graph.py @@ -0,0 +1,40 @@ +# 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 + +"""Static buffer management for platform graph capture (CUDA graph, NPU graph). + +Buffers are allocated on first use, kept for the process lifetime, and shared by all +graphs with the same (name, shape, dtype, device) key; interleaved replay of graphs +with identical keys is not supported. +""" + +import torch + +_BUFFERS: dict[tuple, torch.Tensor] = {} + + +def get_static_buffer( + name: str, + shape: tuple, + dtype: torch.dtype, + device: torch.device | str, + zero: bool = False, +) -> torch.Tensor: + """Return the persistent buffer for the given key, allocated on first use and reused after. + + With ``zero=True`` the buffer is zeroed on every call, not just at allocation. + Ascend graph replay requires padding rows to be physically zero: tl.dot padding + can carry UB residue, and NaN x 0 = NaN defeats zero-weight masks. Called inside + the capture region, the zero-fill is recorded into the graph and re-executed on + every replay. + """ + key = (name, tuple(shape), dtype, str(device)) + buf = _BUFFERS.get(key) + if buf is None: + buf = torch.empty(shape, dtype=dtype, device=device) + _BUFFERS[key] = buf + return buf.zero_() if zero else buf diff --git a/build/torch-cuda/ops/utils/index.py b/build/torch-cuda/ops/utils/index.py index 61de8d4537652a62e6f8c2b7cfa6c28a0f9c69d2..f79daaa3219f164890e649d4ac6f0b21b6ee8047 100644 --- a/build/torch-cuda/ops/utils/index.py +++ b/build/torch-cuda/ops/utils/index.py @@ -27,7 +27,7 @@ def prepare_position_ids_kernel( cu_seqlens, B: tl.constexpr, ): - i_n = tl.program_id(0) + i_n = tl.program_id(0).to(tl.int64) bos, eos = tl.load(cu_seqlens + i_n).to(tl.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) T = eos - bos @@ -181,3 +181,31 @@ def get_max_num_splits( if cu_seqlens_cpu is not None: return triton.cdiv(int(max(prepare_lens(cu_seqlens_cpu))), chunk_size) return triton.cdiv(int(max(prepare_lens(cu_seqlens))), chunk_size) + + +def prepare_chunk_indices_static( + cu_seqlens: torch.LongTensor, + chunk_size: int, + nt_max: int, +) -> tuple[torch.LongTensor, torch.LongTensor]: + """Device-side, fixed-shape chunk-index construction for graph capture. + + Builds ``chunk_indices`` of shape ``[nt_max, 2]`` and ``chunk_offsets`` of shape + ``[N_max + 1]`` with on-device ops, so their construction can be recorded in a + platform graph. ``cu_seqlens`` must be padded with zero-length tail sequences up to + ``N_max + 1`` entries. Rows beyond the real chunk count carry the sentinel + ``i_n = -1``; kernels must return immediately on a negative segment id. + Not cached: the construction runs (and is recorded) on every call. + + Returns both tensors with the same dtype as ``cu_seqlens``. + """ + lens = cu_seqlens[1:] - cu_seqlens[:-1] + chunks = (lens + (chunk_size - 1)).div(chunk_size, rounding_mode='floor') + chunk_offsets = F.pad(chunks.cumsum(0), (1, 0)) + slots = torch.arange(nt_max, device=cu_seqlens.device, dtype=chunk_offsets.dtype) + i_n = torch.searchsorted(chunk_offsets, slots, right=True) - 1 + i_t = slots - chunk_offsets[i_n] + valid = slots < chunk_offsets[-1] + i_n = torch.where(valid, i_n, torch.full_like(i_n, -1)) + i_t = torch.where(valid, i_t, torch.zeros_like(i_t)) + return torch.stack([i_n, i_t], 1).to(cu_seqlens), chunk_offsets.to(cu_seqlens) diff --git a/build/torch-cuda/ops/utils/logcumsumexp.py b/build/torch-cuda/ops/utils/logcumsumexp.py index ee6773f73d9af1a2ded3dfe4294e89bbf121925a..28cdbfd58389097bb0c386d419c5fa33ff3078b0 100644 --- a/build/torch-cuda/ops/utils/logcumsumexp.py +++ b/build/torch-cuda/ops/utils/logcumsumexp.py @@ -29,18 +29,20 @@ def logcumsumexp_fwd_kernel( S: tl.constexpr, BT: tl.constexpr, ): - i_bh = tl.program_id(0) + i_bh = tl.program_id(0).to(tl.int64) o_i = tl.arange(0, BT) m_s = tl.where(o_i[:, None] >= o_i[None, :], 1., 0.) b_mp = tl.full([S], float('-inf'), dtype=tl.float32) b_zp = tl.zeros([S], dtype=tl.float32) for i_t in range(tl.cdiv(T, BT)): - p_s = tl.make_block_ptr(s + i_bh * T*S, (T, S), (S, 1), (i_t * BT, 0), (BT, S), (1, 0)) - p_z = tl.make_block_ptr(z + i_bh * T*S, (T, S), (S, 1), (i_t * BT, 0), (BT, S), (1, 0)) + o_t = i_t * BT + tl.arange(0, BT) + m_t = o_t < T + p_s = s + i_bh * T*S + o_t[:, None] * S + tl.arange(0, S)[None, :] + p_z = z + i_bh * T*S + o_t[:, None] * S + tl.arange(0, S)[None, :] # [BT, S] - b_s = tl.load(p_s, boundary_check=(0, 1)).to(tl.float32) + b_s = tl.load(p_s, mask=m_t[:, None], other=0.0).to(tl.float32) # [S,] b_mc = tl.max(b_s, 0) b_mc = tl.maximum(b_mp, b_mc) @@ -55,4 +57,4 @@ def logcumsumexp_fwd_kernel( # [BT, BS] # small eps to prevent underflows b_z = log(tl.where(b_z != 0, b_z, 1e-20)) + b_mc - tl.store(p_z, b_z.to(p_z.dtype.element_ty), boundary_check=(0, 1)) + tl.store(p_z, b_z.to(p_z.dtype.element_ty), mask=m_t[:, None]) diff --git a/build/torch-cuda/ops/utils/logsumexp.py b/build/torch-cuda/ops/utils/logsumexp.py index 0a9e0a5f53e4cde9a5d1ffe29709ab56f2aae200..92686368ef611fc2b9aff3ae01f2713d776afed6 100644 --- a/build/torch-cuda/ops/utils/logsumexp.py +++ b/build/torch-cuda/ops/utils/logsumexp.py @@ -40,6 +40,7 @@ def logsumexp_fwd_kernel( b_x = tl.load(x + i_n * D + o_d, mask=m_d, other=-float('inf')) if HAS_SCALE: b_x = b_x * scale + 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, B) + i_d, b_z) diff --git a/build/torch-cuda/ops/utils/matmul.py b/build/torch-cuda/ops/utils/matmul.py index cd86847a5ffcbfbb7f779ae212e7ad79da0ab136..2c51108e90f1d67ed9b25870cde4bdf6c6fff3ac 100644 --- a/build/torch-cuda/ops/utils/matmul.py +++ b/build/torch-cuda/ops/utils/matmul.py @@ -87,9 +87,9 @@ def matmul_kernel( # Map program ids `pid` to the block of C it should compute. # This is done in a grouped ordering to promote L2 data reuse. # See above `L2 Cache Optimizations` section for details. - i_b, i_m, i_n = tl.program_id(0), tl.program_id(1), tl.program_id(2) + i_b, i_m, i_n = tl.program_id(0).to(tl.int64), tl.program_id(1).to(tl.int64), tl.program_id(2).to(tl.int64) - NM, NN = tl.num_programs(1), tl.num_programs(2) + NM, NN = tl.num_programs(1).to(tl.int64), tl.num_programs(2).to(tl.int64) i_m, i_n = tl.swizzle2d(i_m, i_n, NM, NN, G) # ---------------------------------------------------------- diff --git a/build/torch-cuda/ops/utils/op.py b/build/torch-cuda/ops/utils/op.py index a4f5e69ce2728f5eab09b8566daa843da47c5efb..b6e504eeb8ae89f4b3c9012e977d5f6539863754 100644 --- a/build/torch-cuda/ops/utils/op.py +++ b/build/torch-cuda/ops/utils/op.py @@ -13,6 +13,22 @@ import triton.language.extra.libdevice as tldevice from ...utils import IS_GATHER_SUPPORTED, IS_NVIDIA_BLACKWELL + +@triton.jit +def barrier(): + tl.debug_barrier() + + +@triton.jit +def unflatten_program_id(X: tl.constexpr, Y: tl.constexpr = None): + """Recover int64 x-fastest 2D/3D indices without using CUDA's limited y/z grid dimensions.""" + pid = tl.program_id(0).to(tl.int64) + if Y is None: + return pid % X, pid // X + else: + return pid % X, (pid // X) % Y, pid // (X * Y) + + if os.environ.get('FLA_USE_FAST_OPS', '0') == '1': @triton.jit def exp(x): return tldevice.fast_expf(x.to(tl.float32)) @@ -37,6 +53,13 @@ else: def tanh(x): return tldevice.tanh(x.to(tl.float32)) +@triton.jit +def logaddexp(a, b): + maximum = tl.maximum(a, b) + difference = tl.where(maximum == float('-inf'), 0., -tl.abs(a - b)) + return maximum + tldevice.log1p(exp(difference)) + + if IS_NVIDIA_BLACKWELL: """ Compute tl.dot with Blackwell workaround. diff --git a/build/torch-cuda/ops/utils/pack.py b/build/torch-cuda/ops/utils/pack.py index 2e2f77074f54f75136045e7838cc330fa93bb1f4..4641fd03ed504ed93703918282cb7233616a778e 100644 --- a/build/torch-cuda/ops/utils/pack.py +++ b/build/torch-cuda/ops/utils/pack.py @@ -13,6 +13,7 @@ import triton import triton.language as tl from ...ops.utils.index import prepare_lens +from ...ops.utils.op import unflatten_program_id from ...utils import autotune_cache_kwargs, input_guard @@ -35,8 +36,8 @@ def packunpack_sequence_kernel( PADDING_SIDE: tl.constexpr, PACK: tl.constexpr, ): - i_d, i_s, i_b = tl.program_id(0), tl.program_id(1), tl.program_id(2) - bos, eos = tl.load(cu_seqlens + i_b), tl.load(cu_seqlens + i_b + 1) + i_d, i_s, i_b = unflatten_program_id(tl.cdiv(D, BD), S) + bos, eos = tl.load(cu_seqlens + i_b).to(tl.int64), tl.load(cu_seqlens + i_b + 1).to(tl.int64) T = eos - bos if PADDING_SIDE == 'left': @@ -71,7 +72,7 @@ def pack_sequence_fwdbwd( ND = triton.cdiv(D, BD) y = torch.empty(cu_seqlens[-1].item(), *x.shape[2:], device=x.device, dtype=x.dtype) - packunpack_sequence_kernel[ND, S, B]( + packunpack_sequence_kernel[(ND * S * B,)]( x=x, y=y, cu_seqlens=cu_seqlens, @@ -98,7 +99,7 @@ def unpack_sequence_fwdbwd( BD = min(triton.next_power_of_2(D), 4096) ND = triton.cdiv(D, BD) - packunpack_sequence_kernel[ND, S, B]( + packunpack_sequence_kernel[(ND * S * B,)]( x=x, y=y, cu_seqlens=cu_seqlens, @@ -144,7 +145,7 @@ class PackSequenceFunction(torch.autograd.Function): padding_side=ctx.padding_side, desired_shape=ctx.desired_shape, ) - return dx, *[None] * 10 + return dx, None, None class UnpackSequenceFunction(torch.autograd.Function): diff --git a/build/torch-cuda/ops/utils/pooling.py b/build/torch-cuda/ops/utils/pooling.py index b8c805c928b5919443e45374972df8a30b2d4d5e..e380c7be051e24fde83b0b1f929c2e17bc82ac0a 100644 --- a/build/torch-cuda/ops/utils/pooling.py +++ b/build/torch-cuda/ops/utils/pooling.py @@ -10,6 +10,7 @@ import triton import triton.language as tl from ...ops.utils.index import prepare_chunk_indices +from ...ops.utils.op import unflatten_program_id from ...utils import autocast_custom_bwd, autocast_custom_fwd, autotune_cache_kwargs, input_guard @@ -38,12 +39,13 @@ def mean_pooling_fwd_kernel( BD: tl.constexpr, IS_VARLEN: tl.constexpr, ): - i_d, i_t, i_bh = tl.program_id(0), tl.program_id(1), tl.program_id(2) + i_d, i_t = unflatten_program_id(tl.cdiv(D, BD)) + i_bh = tl.program_id(1).to(tl.int64) i_b, i_h = i_bh // H, i_bh % H 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.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) + i_n, i_t = tl.load(chunk_indices + i_t * 2).to(tl.int32), tl.load(chunk_indices + i_t * 2 + 1).to(tl.int64) + bos, eos = tl.load(cu_seqlens + i_n).to(tl.int64), tl.load(cu_seqlens + i_n + 1).to(tl.int64) T = eos - bos NT = tl.cdiv(T, BT) else: @@ -51,13 +53,16 @@ def mean_pooling_fwd_kernel( i_tg = i_b * NT + i_t bos, eos = i_b * T, i_b * T + T - p_x = tl.make_block_ptr(x + (bos * H + i_h) * D, (T, D), (H*D, 1), (i_t * BT, i_d * BD), (BT, BD), (1, 0)) - p_o = tl.make_block_ptr(o + (i_tg * H + i_h) * D, (D,), (1,), (i_d * BD,), (BD,), (0,)) + o_t = i_t * BT + tl.arange(0, BT) + o_d = i_d * BD + tl.arange(0, BD) + m_x = (o_t[:, None] < T) & (o_d[None, :] < D) + p_x = x + (bos * H + i_h) * D + o_t[:, None] * (H*D) + o_d[None, :] + p_o = o + (i_tg * H + i_h) * D + o_d # [BT, BD] - b_x = tl.load(p_x, boundary_check=(0, 1)).to(tl.float32) + b_x = tl.load(p_x, mask=m_x, other=0.0).to(tl.float32) # [BD] b_o = tl.sum(b_x, axis=0) / min(BT, T - i_t * BT) - tl.store(p_o, b_o.to(p_o.dtype.element_ty), boundary_check=(0,)) + tl.store(p_o, b_o.to(p_o.dtype.element_ty), mask=o_d < D) @triton.heuristics({ @@ -85,12 +90,13 @@ def mean_pooling_bwd_kernel( BD: tl.constexpr, IS_VARLEN: tl.constexpr, ): - i_d, i_t, i_bh = tl.program_id(0), tl.program_id(1), tl.program_id(2) + i_d, i_t = unflatten_program_id(tl.cdiv(D, BD)) + i_bh = tl.program_id(1).to(tl.int64) i_b, i_h = i_bh // H, i_bh % H 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.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) + i_n, i_t = tl.load(chunk_indices + i_t * 2).to(tl.int32), tl.load(chunk_indices + i_t * 2 + 1).to(tl.int64) + bos, eos = tl.load(cu_seqlens + i_n).to(tl.int64), tl.load(cu_seqlens + i_n + 1).to(tl.int64) T = eos - bos NT = tl.cdiv(T, BT) else: @@ -98,13 +104,16 @@ def mean_pooling_bwd_kernel( i_tg = i_b * NT + i_t bos, eos = i_b * T, i_b * T + T - p_dx = tl.make_block_ptr(dx + (bos * H + i_h) * D, (T, D), (H*D, 1), (i_t * BT, i_d * BD), (BT, BD), (1, 0)) - p_do = tl.make_block_ptr(do + (i_tg * H + i_h) * D, (D,), (1,), (i_d * BD,), (BD,), (0,)) + o_t = i_t * BT + tl.arange(0, BT) + o_d = i_d * BD + tl.arange(0, BD) + m_x = (o_t[:, None] < T) & (o_d[None, :] < D) + p_dx = dx + (bos * H + i_h) * D + o_t[:, None] * (H*D) + o_d[None, :] + p_do = do + (i_tg * H + i_h) * D + o_d # [BD] - b_do = tl.load(p_do, boundary_check=(0,)).to(tl.float32) + b_do = tl.load(p_do, mask=o_d < D, other=0.0).to(tl.float32) # [BT, BD] b_dx = b_do / tl.full((BT,), min(BT, T - i_t * BT), dtype=tl.float32)[:, None] - tl.store(p_dx, b_dx.to(p_dx.dtype.element_ty), boundary_check=(0, 1)) + tl.store(p_dx, b_dx.to(p_dx.dtype.element_ty), mask=m_x) def mean_pooling_fwd( @@ -120,7 +129,7 @@ def mean_pooling_fwd( NT = triton.cdiv(T, BT) if cu_seqlens is None else len(chunk_indices) o = x.new_empty(B, NT, H, D) - def grid(meta): return (triton.cdiv(D, meta['BD']), NT, B * H) + def grid(meta): return (triton.cdiv(D, meta['BD']) * NT, B * H) mean_pooling_fwd_kernel[grid]( x, o, @@ -149,7 +158,7 @@ def mean_pooling_bwd( NT = triton.cdiv(T, BT) if cu_seqlens is None else len(chunk_indices) dx = do.new_empty(B, T, H, D) - def grid(meta): return (triton.cdiv(D, meta['BD']), NT, B * H) + def grid(meta): return (triton.cdiv(D, meta['BD']) * NT, B * H) mean_pooling_bwd_kernel[grid]( do, dx, @@ -199,10 +208,12 @@ def mean_pooling( x: torch.Tensor, chunk_size: int, cu_seqlens: torch.LongTensor | None = None, - head_first: bool = False, + **kwargs, ) -> torch.Tensor: - if head_first: - x = x.transpose(1, 2) + if 'head_first' in kwargs: + raise DeprecationWarning( + "head_first has been removed. Inputs must be in `[B, T, H, ...]` format.", + ) if cu_seqlens is not None: if x.shape[0] != 1: raise ValueError( @@ -210,6 +221,4 @@ def mean_pooling( f"Please flatten variable-length inputs before processing.", ) o = MeanPoolingFunction.apply(x, chunk_size, cu_seqlens) - if head_first: - o = o.transpose(1, 2) return o diff --git a/build/torch-cuda/ops/utils/sinkhorn/__init__.py b/build/torch-cuda/ops/utils/sinkhorn/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..8b513d8ede801e006aeb6400cd6d3ac319b61841 --- /dev/null +++ b/build/torch-cuda/ops/utils/sinkhorn/__init__.py @@ -0,0 +1,16 @@ +# 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 ....ops.utils.sinkhorn.naive import naive_sinkhorn +from ....ops.utils.sinkhorn.ops import sinkhorn + +__all__ = ['naive_sinkhorn', 'sinkhorn'] + + +if GluonBackend.is_available(): + from ....ops.utils.sinkhorn import gluon # noqa: F401 diff --git a/build/torch-cuda/ops/utils/sinkhorn/gluon.py b/build/torch-cuda/ops/utils/sinkhorn/gluon.py new file mode 100644 index 0000000000000000000000000000000000000000..156f5a0b3a4da8e5c724bd88063aebe285032708 --- /dev/null +++ b/build/torch-cuda/ops/utils/sinkhorn/gluon.py @@ -0,0 +1,139 @@ +# 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 triton.experimental import gluon +from triton.experimental.gluon import language as gl + +from ....backends import GluonBackend, register +from ....ops.utils.constant import RCP_LN2 +from ....ops.utils.sinkhorn.ops import sinkhorn_fwd + + +@gluon.jit(do_not_specialize=['T']) +def sinkhorn_fwd_kernel( + x, + o, + T, + scale: gl.constexpr, + N: gl.constexpr, + BN: gl.constexpr, + BT: gl.constexpr, + NUM_ITERS: gl.constexpr, +): + # independent matrices share a warp; row and column reductions never cross warps + RI: gl.constexpr = 4 if BN >= 4 else BN + RJ: gl.constexpr = 8 if BN >= 8 else BN + layout: gl.constexpr = gl.BlockedLayout([1, 1, 1], [32 // (RI * RJ), RI, RJ], [4, 1, 1], [2, 1, 0]) + o_t = gl.program_id(0).to(gl.int64) * BT + gl.arange(0, BT, layout=gl.SliceLayout(1, gl.SliceLayout(2, layout))) + o_i = gl.arange(0, BN, layout=gl.SliceLayout(0, gl.SliceLayout(2, layout))).to(gl.int64) + o_j = gl.arange(0, BN, layout=gl.SliceLayout(0, gl.SliceLayout(1, layout))).to(gl.int64) + o_ij = (o_t[:, None, None] * N + o_i[None, :, None]) * N + o_j[None, None, :] + m_ij = (o_t[:, None, None] < T) & (o_i[None, :, None] < N) & (o_j[None, None, :] < N) + b_x = gl.load(x + o_ij, mask=m_ij, other=float('-inf')).to(gl.float32) + b_x = (b_x - gl.max(b_x, axis=1)[:, None, :]) * scale + for _ in range(NUM_ITERS): + b_max = gl.max(b_x, axis=1) + b_x = b_x - b_max[:, None, :] + b_x = gl.where(m_ij, b_x - gl.log2(gl.sum(gl.exp2(b_x), axis=1))[:, None, :], float('-inf')) + b_max = gl.max(b_x, axis=2) + b_x = b_x - b_max[:, :, None] + b_x = gl.where(m_ij, b_x - gl.log2(gl.sum(gl.exp2(b_x), axis=2))[:, :, None], float('-inf')) + gl.store(o + o_ij, gl.exp2(b_x), mask=m_ij) + + +@gluon.jit(do_not_specialize=['T']) +def sinkhorn_bwd_kernel( + x, + do, + dx, + T, + scale: gl.constexpr, + N: gl.constexpr, + BN: gl.constexpr, + BT: gl.constexpr, + NUM_ITERS: gl.constexpr, +): + # independent matrices share a warp; row and column reductions never cross warps + RI: gl.constexpr = 4 if BN >= 4 else BN + RJ: gl.constexpr = 8 if BN >= 8 else BN + layout: gl.constexpr = gl.BlockedLayout([1, 1, 1], [32 // (RI * RJ), RI, RJ], [4, 1, 1], [2, 1, 0]) + o_t = gl.program_id(0).to(gl.int64) * BT + gl.arange(0, BT, layout=gl.SliceLayout(1, gl.SliceLayout(2, layout))) + o_i = gl.arange(0, BN, layout=gl.SliceLayout(0, gl.SliceLayout(2, layout))).to(gl.int64) + o_j = gl.arange(0, BN, layout=gl.SliceLayout(0, gl.SliceLayout(1, layout))).to(gl.int64) + o_ij = (o_t[:, None, None] * N + o_i[None, :, None]) * N + o_j[None, None, :] + m_ij = (o_t[:, None, None] < T) & (o_i[None, :, None] < N) & (o_j[None, None, :] < N) + b_input = gl.load(x + o_ij, mask=m_ij, other=float('-inf')).to(gl.float32) + b_input = (b_input - gl.max(b_input, axis=1)[:, None, :]) * scale + b_dx = gl.load(do + o_ij, mask=m_ij, other=0).to(gl.float32) + # recompute earlier blocks to keep the shared probability cache within 40 KiB. + BI: gl.constexpr = NUM_ITERS if NUM_ITERS < 20 else 20 + b_ps = gl.allocate_shared_memory(gl.float32, [2 * BI, BT, BN, BN], gl.SwizzledSharedLayout(1, 1, 1, [2, 1, 0])) + for i_b in range(triton.cdiv(NUM_ITERS, BI) - 1, -1, -1): + b_x = b_input + for step in range(gl.minimum((i_b + 1) * BI, NUM_ITERS)): + b_x = b_x - gl.max(b_x, axis=1)[:, None, :] + b_x = gl.where(m_ij, b_x - gl.log2(gl.sum(gl.exp2(b_x), axis=1))[:, None, :], float('-inf')) + if step >= i_b * BI: + b_ps.index(2 * (step - i_b * BI)).store(gl.exp2(b_x)) + b_x = b_x - gl.max(b_x, axis=2)[:, :, None] + b_x = gl.where(m_ij, b_x - gl.log2(gl.sum(gl.exp2(b_x), axis=2))[:, :, None], float('-inf')) + if step >= i_b * BI: + b_ps.index(2 * (step - i_b * BI) + 1).store(gl.exp2(b_x)) + if i_b == triton.cdiv(NUM_ITERS, BI) - 1: + b_dx = b_dx * gl.exp2(b_x) + for step in range(gl.minimum(BI, NUM_ITERS - i_b * BI) - 1, -1, -1): + b_p = b_ps.index(2 * step + 1).load(layout) + b_dx = b_dx - b_p * gl.sum(b_dx, axis=2)[:, :, None] + b_p = b_ps.index(2 * step).load(layout) + b_dx = b_dx - b_p * gl.sum(b_dx, axis=1)[:, None, :] + gl.store(dx + o_ij, b_dx, mask=m_ij) + + +@register(sinkhorn_fwd, backend=GluonBackend) +def sinkhorn_fwd_gluon(logits: torch.Tensor, num_iters: int): + N = logits.shape[-1] + T = logits.numel() // (N * N) + out = torch.empty_like(logits, dtype=torch.float32) + BT = 1 if T < 128 else max(4, 128 // min(32, triton.next_power_of_2(N) ** 2)) + if T: + sinkhorn_fwd_kernel[(triton.cdiv(T, BT),)]( + x=logits, + o=out, + T=T, + scale=RCP_LN2, + N=N, + BN=triton.next_power_of_2(N), + BT=BT, + NUM_ITERS=num_iters, + num_warps=4, + enable_fp_fusion=False, + ) + return out, sinkhorn_bwd_gluon + + +def sinkhorn_bwd_gluon(logits: torch.Tensor, dout: torch.Tensor, num_iters: int) -> torch.Tensor: + N = logits.shape[-1] + T = logits.numel() // (N * N) + dx = torch.empty_like(logits) + BT = 1 if T < 128 else max(4, 128 // min(32, triton.next_power_of_2(N) ** 2)) + if T: + sinkhorn_bwd_kernel[(triton.cdiv(T, BT),)]( + x=logits, + do=dout, + dx=dx, + T=T, + scale=RCP_LN2, + N=N, + BN=triton.next_power_of_2(N), + BT=BT, + NUM_ITERS=num_iters, + num_warps=4, + enable_fp_fusion=False, + ) + return dx diff --git a/build/torch-cuda/ops/utils/sinkhorn/naive.py b/build/torch-cuda/ops/utils/sinkhorn/naive.py new file mode 100644 index 0000000000000000000000000000000000000000..a1fab9ae92487fc5f15a1bbf75193e17b359e968 --- /dev/null +++ b/build/torch-cuda/ops/utils/sinkhorn/naive.py @@ -0,0 +1,17 @@ +# 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 + + +def naive_sinkhorn(logits: torch.Tensor, num_iters: int = 20) -> torch.Tensor: + """Apply column and row log normalizations in FP32.""" + x = logits.float() + for _ in range(num_iters): + x = x - x.logsumexp(dim=-2, keepdim=True) + x = x - x.logsumexp(dim=-1, keepdim=True) + return x.exp() diff --git a/build/torch-cuda/ops/utils/sinkhorn/ops.py b/build/torch-cuda/ops/utils/sinkhorn/ops.py new file mode 100644 index 0000000000000000000000000000000000000000..6d88a92029434991ec1273ae13e7821b077f311c --- /dev/null +++ b/build/torch-cuda/ops/utils/sinkhorn/ops.py @@ -0,0 +1,167 @@ +# 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 +import triton.language as tl +from torch.autograd.function import once_differentiable + +from ....backends import dispatch +from ....ops.utils.constant import RCP_LN2 +from ....utils import IS_NVIDIA, autocast_custom_bwd, autocast_custom_fwd, input_guard + + +@triton.jit(do_not_specialize=['T']) +def sinkhorn_fwd_kernel( + x, + o, + T, + scale: tl.constexpr, + N: tl.constexpr, + BN: tl.constexpr, + NUM_ITERS: tl.constexpr, +): + i_t = tl.program_id(0).to(tl.int64) + o_i = tl.arange(0, BN).to(tl.int64) + o_j = tl.arange(0, BN).to(tl.int64) + o_ij = (i_t * N + o_i[:, None]) * N + o_j[None, :] + m_ij = (o_i[:, None] < N) & (o_j[None, :] < N) + b_x = tl.load(x + o_ij, mask=m_ij, other=float('-inf')).to(tl.float32) + b_x = (b_x - tl.max(b_x, axis=0)[None, :]) * scale + for _ in range(NUM_ITERS): + b_max = tl.max(b_x, axis=0) + b_x = b_x - b_max[None, :] + b_x = tl.where(m_ij, b_x - tl.log2(tl.sum(tl.exp2(b_x), axis=0))[None, :], float('-inf')) + b_max = tl.max(b_x, axis=1) + b_x = b_x - b_max[:, None] + b_x = tl.where(m_ij, b_x - tl.log2(tl.sum(tl.exp2(b_x), axis=1))[:, None], float('-inf')) + tl.store(o + o_ij, tl.exp2(b_x), mask=m_ij) + + +@triton.jit(do_not_specialize=['T']) +def sinkhorn_bwd_kernel(x, do, dx, T, scale: tl.constexpr, N: tl.constexpr, BN: tl.constexpr, NUM_ITERS: tl.constexpr): + i_t = tl.program_id(0).to(tl.int64) + o_i = tl.arange(0, BN).to(tl.int64) + o_j = tl.arange(0, BN).to(tl.int64) + o_ij = (i_t * N + o_i[:, None]) * N + o_j[None, :] + m_ij = (o_i[:, None] < N) & (o_j[None, :] < N) + b_input = tl.load(x + o_ij, mask=m_ij, other=float('-inf')).to(tl.float32) + b_input = (b_input - tl.max(b_input, axis=0)[None, :]) * scale + b_dx = tl.load(do + o_ij, mask=m_ij, other=0).to(tl.float32) + # recompute earlier blocks to bound the live probability cache. + BI: tl.constexpr = NUM_ITERS if NUM_ITERS < 20 else 20 + for i_b in range(triton.cdiv(NUM_ITERS, BI) - 1, -1, -1): + b_x = b_input + for _ in range(i_b * BI): + b_x = b_x - tl.max(b_x, axis=0)[None, :] + b_x = tl.where(m_ij, b_x - tl.log2(tl.sum(tl.exp2(b_x), axis=0))[None, :], float('-inf')) + b_x = b_x - tl.max(b_x, axis=1)[:, None] + b_x = tl.where(m_ij, b_x - tl.log2(tl.sum(tl.exp2(b_x), axis=1))[:, None], float('-inf')) + b_ps = () + for step in tl.static_range(BI): + b_col = tl.full((BN, BN), 0, tl.float32) + b_row = tl.full((BN, BN), 0, tl.float32) + if i_b * BI + step < NUM_ITERS: + b_x = b_x - tl.max(b_x, axis=0)[None, :] + b_x = tl.where(m_ij, b_x - tl.log2(tl.sum(tl.exp2(b_x), axis=0))[None, :], float('-inf')) + b_col = tl.exp2(b_x) + b_x = b_x - tl.max(b_x, axis=1)[:, None] + b_x = tl.where(m_ij, b_x - tl.log2(tl.sum(tl.exp2(b_x), axis=1))[:, None], float('-inf')) + b_row = tl.exp2(b_x) + b_ps += (b_col, b_row) + if i_b == triton.cdiv(NUM_ITERS, BI) - 1: + b_dx = b_dx * tl.exp2(b_x) + for step in tl.static_range(BI - 1, -1, -1): + b_dx = b_dx - b_ps[2 * step + 1] * tl.sum(b_dx, axis=1)[:, None] + b_dx = b_dx - b_ps[2 * step] * tl.sum(b_dx, axis=0)[None, :] + tl.store(dx + o_ij, b_dx, mask=m_ij) + + +@dispatch +def sinkhorn_fwd(logits: torch.Tensor, num_iters: int): + N = logits.shape[-1] + T = logits.numel() // (N * N) + out = torch.empty_like(logits, dtype=torch.float32) + if T: + sinkhorn_fwd_kernel[(T,)]( + x=logits, + o=out, + T=T, + scale=RCP_LN2, + N=N, + BN=triton.next_power_of_2(N), + NUM_ITERS=num_iters, + num_warps=4, + enable_fp_fusion=False, + ) + return out, sinkhorn_bwd + + +def sinkhorn_bwd(logits: torch.Tensor, dout: torch.Tensor, num_iters: int) -> torch.Tensor: + N = logits.shape[-1] + T = logits.numel() // (N * N) + dx = torch.empty_like(logits) + if T: + sinkhorn_bwd_kernel[(T,)]( + x=logits, + do=dout, + dx=dx, + T=T, + scale=RCP_LN2, + N=N, + BN=triton.next_power_of_2(N), + NUM_ITERS=num_iters, + num_warps=4, + enable_fp_fusion=False, + ) + return dx + + +class SinkhornFunction(torch.autograd.Function): + + @staticmethod + @input_guard + @autocast_custom_fwd + def forward(ctx, logits, num_iters): + out, ctx.bwd = sinkhorn_fwd(logits=logits, num_iters=num_iters) + ctx.save_for_backward(logits) + ctx.num_iters = num_iters + return out + + @staticmethod + @input_guard + @autocast_custom_bwd + @once_differentiable + def backward(ctx, dout): + logits, = ctx.saved_tensors + dx = ctx.bwd(logits=logits, dout=dout, num_iters=ctx.num_iters) + return dx, None + + +def sinkhorn(logits: torch.Tensor, num_iters: int = 20) -> torch.Tensor: + """Normalize square matrices with alternating column and row log-sum-exp steps. + + Args: + logits (torch.Tensor): + Input logits of shape `[..., N, N]`, with `1 <= N <= 8`. + num_iters (int, Optional): + Positive number of normalization iterations. Default: 20. + + Returns: + torch.Tensor: + Nonnegative FP32 matrices of shape `[..., N, N]`. + A finite iteration count does not guarantee exact double stochasticity. + """ + if not IS_NVIDIA or not logits.is_cuda: + raise ValueError('Sinkhorn kernels require NVIDIA CUDA tensors') + if logits.dtype not in (torch.float16, torch.bfloat16, torch.float32): + raise ValueError('logits must use FP16, BF16, or FP32') + if logits.ndim < 2 or not 1 <= logits.shape[-1] <= 8 or logits.shape[-1] != logits.shape[-2]: + raise ValueError('logits must have shape [..., N, N] with 1 <= N <= 8') + if isinstance(num_iters, bool) or not isinstance(num_iters, int) or num_iters < 1: + raise ValueError('num_iters must be a positive integer') + return SinkhornFunction.apply(logits, num_iters) diff --git a/build/torch-cuda/ops/utils/softmax.py b/build/torch-cuda/ops/utils/softmax.py index 55806ad9e7ac7fbccec07c69d9e18bace3b5e34e..06ef33d4183718872af135333c26e85ddce93f67 100644 --- a/build/torch-cuda/ops/utils/softmax.py +++ b/build/torch-cuda/ops/utils/softmax.py @@ -30,7 +30,7 @@ def softmax_fwd_kernel( D: tl.constexpr, B: tl.constexpr, ): - i_n = tl.program_id(0) + i_n = tl.program_id(0).to(tl.int64) o_d = tl.arange(0, B) m_d = o_d < D @@ -58,7 +58,7 @@ def softmax_bwd_kernel( D: tl.constexpr, B: tl.constexpr, ): - i_n = tl.program_id(0) + i_n = tl.program_id(0).to(tl.int64) o_d = tl.arange(0, B) m_d = o_d < D diff --git a/build/torch-cuda/ops/utils/softplus.py b/build/torch-cuda/ops/utils/softplus.py index b8844e7fdb7127495dba8659684c37bd484f4134..020c2d44fa9b5d5d58dd21110bfece55d527ac4e 100644 --- a/build/torch-cuda/ops/utils/softplus.py +++ b/build/torch-cuda/ops/utils/softplus.py @@ -107,9 +107,20 @@ def softplus2_triton(x): return tl.where(x < 15.0, tl.math.log2(1 + tl.math.exp2(x)), x) +# named wrappers preserve the symbols referenced by kernels in Inductor's generated source if IS_NVIDIA: - softplus = softplus_nv - softplus2 = softplus2_nv + @triton.jit + def softplus(x): + return softplus_nv(x) + + @triton.jit + def softplus2(x): + return softplus2_nv(x) else: - softplus = softplus_triton - softplus2 = softplus2_triton + @triton.jit + def softplus(x): + return softplus_triton(x) + + @triton.jit + def softplus2(x): + return softplus2_triton(x) diff --git a/build/torch-cuda/ops/utils/solve_tril.py b/build/torch-cuda/ops/utils/solve_tril.py index f61ce4a3030e81d2747a5eaca8ec98928a7084f4..fd4c4f46df2a48256837fecf2b3c6b5c77b717a1 100644 --- a/build/torch-cuda/ops/utils/solve_tril.py +++ b/build/torch-cuda/ops/utils/solve_tril.py @@ -11,14 +11,21 @@ import torch import triton import triton.language as tl +from ...backends import dispatch from ...ops.utils.index import prepare_chunk_indices from ...ops.utils.op import make_tensor_descriptor -from ...utils import IS_TMA_SUPPORTED, autotune_cache_kwargs, input_guard +from ...utils import IS_TF32_SUPPORTED, IS_TMA_SUPPORTED, autotune_cache_kwargs, input_guard FLA_TRIL_PRECISION = os.environ.get('FLA_TRIL_PRECISION', 'ieee') assert FLA_TRIL_PRECISION in ['ieee', 'tf32', 'tf32x3'], \ f"FLA_TRIL_PRECISION must be one of 'ieee', 'tf32', or 'tf32x3', but got {FLA_TRIL_PRECISION}" -DOT_PRECISION_AUTOTUNE_LIST = ["ieee"] if not IS_TMA_SUPPORTED else list({"ieee", FLA_TRIL_PRECISION}) +# NOTE: `IS_TF32_SUPPORTED` keeps the tf32 configs off non-Nvidia backends now that +# `IS_TMA_SUPPORTED` can also be true on AMD. The order is fixed instead of derived from a set so +# the autotune config order does not vary with the interpreter hash seed. +if IS_TMA_SUPPORTED and IS_TF32_SUPPORTED and FLA_TRIL_PRECISION != "ieee": + DOT_PRECISION_AUTOTUNE_LIST = ["ieee", FLA_TRIL_PRECISION] +else: + DOT_PRECISION_AUTOTUNE_LIST = ["ieee"] @triton.heuristics({ @@ -46,11 +53,11 @@ def solve_tril_16x16_kernel( IS_VARLEN: tl.constexpr, DOT_PRECISION: tl.constexpr, ): - i_t, i_bh = tl.program_id(0), tl.program_id(1).to(tl.int64) + i_t, i_bh = tl.program_id(0).to(tl.int64), tl.program_id(1).to(tl.int64) i_b, i_h = i_bh // H, i_bh % H 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.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) + i_n, i_t = tl.load(chunk_indices + i_t * 2).to(tl.int32), tl.load(chunk_indices + i_t * 2 + 1).to(tl.int64) + bos, eos = tl.load(cu_seqlens + i_n).to(tl.int64), tl.load(cu_seqlens + i_n + 1).to(tl.int64) T = eos - bos else: bos, eos = i_b * T, i_b * T + T @@ -62,10 +69,12 @@ def solve_tril_16x16_kernel( Ai = Ai + (bos*H + i_h) * 16 offset = (i_t * 16) % BT + o_t = (i_t * 16 + o_i).to(tl.int64) + m_t = o_t < T if not USE_TMA: - p_A = tl.make_block_ptr(A, (T, BT), (H*BT, 1), (i_t * 16, offset), (16, 16), (1, 0)) + p_A = A + o_t[:, None] * (H*BT) + (o_i + offset)[None, :] # [16, 16] - b_A = tl.load(p_A, boundary_check=(0, 1)).to(tl.float32) + b_A = tl.load(p_A, mask=m_t[:, None], other=0.0).to(tl.float32) b_A = tl.where(m_A, b_A, 0) else: desc = make_tensor_descriptor(A, [T, BT], [H*BT, 1], [16, 16]) @@ -82,8 +91,8 @@ def solve_tril_16x16_kernel( b_A = tl.where((o_i == i)[:, None], b_a, b_A) b_A += m_I if not USE_TMA: - p_Ai = tl.make_block_ptr(Ai, (T, 16), (H*16, 1), (i_t * 16, 0), (16, 16), (1, 0)) - tl.store(p_Ai, b_A.to(p_Ai.dtype.element_ty, fp_downcast_rounding="rtne"), boundary_check=(0, 1)) + p_Ai = Ai + o_t[:, None] * (H*16) + o_i[None, :] + tl.store(p_Ai, b_A.to(p_Ai.dtype.element_ty, fp_downcast_rounding="rtne"), mask=m_t[:, None]) else: desc_o.store([i_t * 16, 0], b_A.to(desc_o.dtype, fp_downcast_rounding="rtne")) @@ -114,11 +123,11 @@ def merge_16x16_to_32x32_inverse_kernel( IS_VARLEN: tl.constexpr, DOT_PRECISION: tl.constexpr, ): - i_t, i_bh = tl.program_id(0), tl.program_id(1).to(tl.int64) + i_t, i_bh = tl.program_id(0).to(tl.int64), tl.program_id(1).to(tl.int64) i_b, i_h = i_bh // H, i_bh % H 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.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) + i_n, i_t = tl.load(chunk_indices + i_t * 2).to(tl.int32), tl.load(chunk_indices + i_t * 2 + 1).to(tl.int64) + bos, eos = tl.load(cu_seqlens + i_n).to(tl.int64), tl.load(cu_seqlens + i_n + 1).to(tl.int64) T = eos - bos else: bos, eos = i_b * T, i_b * T + T @@ -129,11 +138,12 @@ def merge_16x16_to_32x32_inverse_kernel( A += (bos * H + i_h) * BT Ai += (bos * H + i_h) * BT + o_t = (i_t * BT + o_i).to(tl.int64) if not USE_TMA: - p_A_11 = tl.make_block_ptr(A, (T, BT), (H*BT, 1), (i_t * BT, 0), (16, 16), (1, 0)) - p_A_22 = tl.make_block_ptr(A, (T, BT), (H*BT, 1), (i_t * BT + 16, 16), (16, 16), (1, 0)) - b_Ai_11 = tl.load(p_A_11, boundary_check=(0, 1)).to(tl.float32) - b_Ai_22 = tl.load(p_A_22, boundary_check=(0, 1)).to(tl.float32) + p_A_11 = A + o_t[:, None] * (H*BT) + o_i[None, :] + p_A_22 = A + (o_t[:, None] + 16) * (H*BT) + (o_i[None, :] + 16) + b_Ai_11 = tl.load(p_A_11, mask=(o_t[:, None] < T), other=0.0).to(tl.float32) + b_Ai_22 = tl.load(p_A_22, mask=(o_t[:, None] + 16 < T), other=0.0).to(tl.float32) else: desc = make_tensor_descriptor(A, [T, BT], [H*BT, 1], [16, 16]) desc_o = make_tensor_descriptor(Ai, [T, BT], [H*BT, 1], [16, 16]) @@ -157,20 +167,20 @@ def merge_16x16_to_32x32_inverse_kernel( b_Ai_22 += m_I if not USE_TMA: - p_A_21 = tl.make_block_ptr(A, (T, BT), (H*BT, 1), (i_t * BT + 16, 0), (16, 16), (1, 0)) - b_A_21 = tl.load(p_A_21, boundary_check=(0, 1)).to(tl.float32) + p_A_21 = A + (o_t[:, None] + 16) * (H*BT) + o_i[None, :] + b_A_21 = tl.load(p_A_21, mask=(o_t[:, None] + 16 < T), other=0.0).to(tl.float32) else: b_A_21 = desc.load([i_t * BT + 16, 0]).to(tl.float32) b_Ai_21 = -tl.dot(tl.dot(b_Ai_22, b_A_21, input_precision=DOT_PRECISION), b_Ai_11, input_precision=DOT_PRECISION) if not USE_TMA: - p_Ai_11 = tl.make_block_ptr(Ai, (T, BT), (H*BT, 1), (i_t * BT, 0), (16, 16), (1, 0)) - p_Ai_21 = tl.make_block_ptr(Ai, (T, BT), (H*BT, 1), (i_t * BT + 16, 0), (16, 16), (1, 0)) - p_Ai_22 = tl.make_block_ptr(Ai, (T, BT), (H*BT, 1), (i_t * BT + 16, 16), (16, 16), (1, 0)) - tl.store(p_Ai_11, b_Ai_11.to(p_Ai_11.dtype.element_ty, fp_downcast_rounding="rtne"), boundary_check=(0, 1)) - tl.store(p_Ai_22, b_Ai_22.to(p_Ai_22.dtype.element_ty, fp_downcast_rounding="rtne"), boundary_check=(0, 1)) - tl.store(p_Ai_21, b_Ai_21.to(p_Ai_21.dtype.element_ty, fp_downcast_rounding="rtne"), boundary_check=(0, 1)) + p_Ai_11 = Ai + o_t[:, None] * (H*BT) + o_i[None, :] + p_Ai_21 = Ai + (o_t[:, None] + 16) * (H*BT) + o_i[None, :] + p_Ai_22 = Ai + (o_t[:, None] + 16) * (H*BT) + (o_i[None, :] + 16) + tl.store(p_Ai_11, b_Ai_11.to(p_Ai_11.dtype.element_ty, fp_downcast_rounding="rtne"), mask=(o_t[:, None] < T)) + tl.store(p_Ai_22, b_Ai_22.to(p_Ai_22.dtype.element_ty, fp_downcast_rounding="rtne"), mask=(o_t[:, None] + 16 < T)) + tl.store(p_Ai_21, b_Ai_21.to(p_Ai_21.dtype.element_ty, fp_downcast_rounding="rtne"), mask=(o_t[:, None] + 16 < T)) else: desc_o.store([i_t * BT + 0, 0], b_Ai_11.to(desc_o.dtype, fp_downcast_rounding="rtne")) desc_o.store([i_t * BT + 16, 0], b_Ai_21.to(desc_o.dtype, fp_downcast_rounding="rtne")) @@ -203,11 +213,11 @@ def merge_16x16_to_64x64_inverse_kernel( IS_VARLEN: tl.constexpr, DOT_PRECISION: tl.constexpr, ): - i_t, i_bh = tl.program_id(0), tl.program_id(1).to(tl.int64) + i_t, i_bh = tl.program_id(0).to(tl.int64), tl.program_id(1).to(tl.int64) i_b, i_h = i_bh // H, i_bh % H 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.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) + i_n, i_t = tl.load(chunk_indices + i_t * 2).to(tl.int32), tl.load(chunk_indices + i_t * 2 + 1).to(tl.int64) + bos, eos = tl.load(cu_seqlens + i_n).to(tl.int64), tl.load(cu_seqlens + i_n + 1).to(tl.int64) T = eos - bos else: bos, eos = i_b * T, i_b * T + T @@ -218,15 +228,16 @@ def merge_16x16_to_64x64_inverse_kernel( A += (bos * H + i_h) * BT Ai += (bos * H + i_h) * BT + o_t = (i_t * BT + o_i).to(tl.int64) if not USE_TMA: - p_A_11 = tl.make_block_ptr(A, (T, BT), (H*BT, 1), (i_t * BT, 0), (16, 16), (1, 0)) - p_A_22 = tl.make_block_ptr(A, (T, BT), (H*BT, 1), (i_t * BT + 16, 16), (16, 16), (1, 0)) - p_A_33 = tl.make_block_ptr(A, (T, BT), (H*BT, 1), (i_t * BT + 32, 32), (16, 16), (1, 0)) - p_A_44 = tl.make_block_ptr(A, (T, BT), (H*BT, 1), (i_t * BT + 48, 48), (16, 16), (1, 0)) - b_Ai_11 = tl.load(p_A_11, boundary_check=(0, 1)).to(tl.float32) - b_Ai_22 = tl.load(p_A_22, boundary_check=(0, 1)).to(tl.float32) - b_Ai_33 = tl.load(p_A_33, boundary_check=(0, 1)).to(tl.float32) - b_Ai_44 = tl.load(p_A_44, boundary_check=(0, 1)).to(tl.float32) + p_A_11 = A + o_t[:, None] * (H*BT) + o_i[None, :] + p_A_22 = A + (o_t[:, None] + 16) * (H*BT) + (o_i[None, :] + 16) + p_A_33 = A + (o_t[:, None] + 32) * (H*BT) + (o_i[None, :] + 32) + p_A_44 = A + (o_t[:, None] + 48) * (H*BT) + (o_i[None, :] + 48) + b_Ai_11 = tl.load(p_A_11, mask=(o_t[:, None] < T), other=0.0).to(tl.float32) + b_Ai_22 = tl.load(p_A_22, mask=(o_t[:, None] + 16 < T), other=0.0).to(tl.float32) + b_Ai_33 = tl.load(p_A_33, mask=(o_t[:, None] + 32 < T), other=0.0).to(tl.float32) + b_Ai_44 = tl.load(p_A_44, mask=(o_t[:, None] + 48 < T), other=0.0).to(tl.float32) else: desc = make_tensor_descriptor(A, [T, BT], [H*BT, 1], [16, 16]) desc_o = make_tensor_descriptor(Ai, [T, BT], [H*BT, 1], [16, 16]) @@ -267,18 +278,18 @@ def merge_16x16_to_64x64_inverse_kernel( b_Ai_44 += m_I if not USE_TMA: - p_A_21 = tl.make_block_ptr(A, (T, BT), (H*BT, 1), (i_t * BT + 16, 0), (16, 16), (1, 0)) - p_A_31 = tl.make_block_ptr(A, (T, BT), (H*BT, 1), (i_t * BT + 32, 0), (16, 16), (1, 0)) - p_A_32 = tl.make_block_ptr(A, (T, BT), (H*BT, 1), (i_t * BT + 32, 16), (16, 16), (1, 0)) - p_A_41 = tl.make_block_ptr(A, (T, BT), (H*BT, 1), (i_t * BT + 48, 0), (16, 16), (1, 0)) - p_A_42 = tl.make_block_ptr(A, (T, BT), (H*BT, 1), (i_t * BT + 48, 16), (16, 16), (1, 0)) - p_A_43 = tl.make_block_ptr(A, (T, BT), (H*BT, 1), (i_t * BT + 48, 32), (16, 16), (1, 0)) - b_A_21 = tl.load(p_A_21, boundary_check=(0, 1)).to(tl.float32) - b_A_31 = tl.load(p_A_31, boundary_check=(0, 1)).to(tl.float32) - b_A_32 = tl.load(p_A_32, boundary_check=(0, 1)).to(tl.float32) - b_A_41 = tl.load(p_A_41, boundary_check=(0, 1)).to(tl.float32) - b_A_42 = tl.load(p_A_42, boundary_check=(0, 1)).to(tl.float32) - b_A_43 = tl.load(p_A_43, boundary_check=(0, 1)).to(tl.float32) + p_A_21 = A + (o_t[:, None] + 16) * (H*BT) + o_i[None, :] + p_A_31 = A + (o_t[:, None] + 32) * (H*BT) + o_i[None, :] + p_A_32 = A + (o_t[:, None] + 32) * (H*BT) + (o_i[None, :] + 16) + p_A_41 = A + (o_t[:, None] + 48) * (H*BT) + o_i[None, :] + p_A_42 = A + (o_t[:, None] + 48) * (H*BT) + (o_i[None, :] + 16) + p_A_43 = A + (o_t[:, None] + 48) * (H*BT) + (o_i[None, :] + 32) + b_A_21 = tl.load(p_A_21, mask=(o_t[:, None] + 16 < T), other=0.0).to(tl.float32) + b_A_31 = tl.load(p_A_31, mask=(o_t[:, None] + 32 < T), other=0.0).to(tl.float32) + b_A_32 = tl.load(p_A_32, mask=(o_t[:, None] + 32 < T), other=0.0).to(tl.float32) + b_A_41 = tl.load(p_A_41, mask=(o_t[:, None] + 48 < T), other=0.0).to(tl.float32) + b_A_42 = tl.load(p_A_42, mask=(o_t[:, None] + 48 < T), other=0.0).to(tl.float32) + b_A_43 = tl.load(p_A_43, mask=(o_t[:, None] + 48 < T), other=0.0).to(tl.float32) else: b_A_21 = desc.load([i_t * BT + 16, 0]).to(tl.float32) b_A_31 = desc.load([i_t * BT + 32, 0]).to(tl.float32) @@ -312,26 +323,26 @@ def merge_16x16_to_64x64_inverse_kernel( ) if not USE_TMA: - p_Ai_11 = tl.make_block_ptr(Ai, (T, BT), (H*BT, 1), (i_t * BT, 0), (16, 16), (1, 0)) - p_Ai_22 = tl.make_block_ptr(Ai, (T, BT), (H*BT, 1), (i_t * BT + 16, 16), (16, 16), (1, 0)) - p_Ai_33 = tl.make_block_ptr(Ai, (T, BT), (H*BT, 1), (i_t * BT + 32, 32), (16, 16), (1, 0)) - p_Ai_44 = tl.make_block_ptr(Ai, (T, BT), (H*BT, 1), (i_t * BT + 48, 48), (16, 16), (1, 0)) - p_Ai_21 = tl.make_block_ptr(Ai, (T, BT), (H*BT, 1), (i_t * BT + 16, 0), (16, 16), (1, 0)) - p_Ai_31 = tl.make_block_ptr(Ai, (T, BT), (H*BT, 1), (i_t * BT + 32, 0), (16, 16), (1, 0)) - p_Ai_32 = tl.make_block_ptr(Ai, (T, BT), (H*BT, 1), (i_t * BT + 32, 16), (16, 16), (1, 0)) - p_Ai_41 = tl.make_block_ptr(Ai, (T, BT), (H*BT, 1), (i_t * BT + 48, 0), (16, 16), (1, 0)) - p_Ai_42 = tl.make_block_ptr(Ai, (T, BT), (H*BT, 1), (i_t * BT + 48, 16), (16, 16), (1, 0)) - p_Ai_43 = tl.make_block_ptr(Ai, (T, BT), (H*BT, 1), (i_t * BT + 48, 32), (16, 16), (1, 0)) - tl.store(p_Ai_11, b_Ai_11.to(p_Ai_11.dtype.element_ty, fp_downcast_rounding="rtne"), boundary_check=(0, 1)) - tl.store(p_Ai_22, b_Ai_22.to(p_Ai_22.dtype.element_ty, fp_downcast_rounding="rtne"), boundary_check=(0, 1)) - tl.store(p_Ai_33, b_Ai_33.to(p_Ai_33.dtype.element_ty, fp_downcast_rounding="rtne"), boundary_check=(0, 1)) - tl.store(p_Ai_44, b_Ai_44.to(p_Ai_44.dtype.element_ty, fp_downcast_rounding="rtne"), boundary_check=(0, 1)) - tl.store(p_Ai_21, b_Ai_21.to(p_Ai_21.dtype.element_ty, fp_downcast_rounding="rtne"), boundary_check=(0, 1)) - tl.store(p_Ai_31, b_Ai_31.to(p_Ai_31.dtype.element_ty, fp_downcast_rounding="rtne"), boundary_check=(0, 1)) - tl.store(p_Ai_32, b_Ai_32.to(p_Ai_32.dtype.element_ty, fp_downcast_rounding="rtne"), boundary_check=(0, 1)) - tl.store(p_Ai_41, b_Ai_41.to(p_Ai_41.dtype.element_ty, fp_downcast_rounding="rtne"), boundary_check=(0, 1)) - tl.store(p_Ai_42, b_Ai_42.to(p_Ai_42.dtype.element_ty, fp_downcast_rounding="rtne"), boundary_check=(0, 1)) - tl.store(p_Ai_43, b_Ai_43.to(p_Ai_43.dtype.element_ty, fp_downcast_rounding="rtne"), boundary_check=(0, 1)) + p_Ai_11 = Ai + o_t[:, None] * (H*BT) + o_i[None, :] + p_Ai_22 = Ai + (o_t[:, None] + 16) * (H*BT) + (o_i[None, :] + 16) + p_Ai_33 = Ai + (o_t[:, None] + 32) * (H*BT) + (o_i[None, :] + 32) + p_Ai_44 = Ai + (o_t[:, None] + 48) * (H*BT) + (o_i[None, :] + 48) + p_Ai_21 = Ai + (o_t[:, None] + 16) * (H*BT) + o_i[None, :] + p_Ai_31 = Ai + (o_t[:, None] + 32) * (H*BT) + o_i[None, :] + p_Ai_32 = Ai + (o_t[:, None] + 32) * (H*BT) + (o_i[None, :] + 16) + p_Ai_41 = Ai + (o_t[:, None] + 48) * (H*BT) + o_i[None, :] + p_Ai_42 = Ai + (o_t[:, None] + 48) * (H*BT) + (o_i[None, :] + 16) + p_Ai_43 = Ai + (o_t[:, None] + 48) * (H*BT) + (o_i[None, :] + 32) + tl.store(p_Ai_11, b_Ai_11.to(p_Ai_11.dtype.element_ty, fp_downcast_rounding="rtne"), mask=(o_t[:, None] < T)) + tl.store(p_Ai_22, b_Ai_22.to(p_Ai_22.dtype.element_ty, fp_downcast_rounding="rtne"), mask=(o_t[:, None] + 16 < T)) + tl.store(p_Ai_33, b_Ai_33.to(p_Ai_33.dtype.element_ty, fp_downcast_rounding="rtne"), mask=(o_t[:, None] + 32 < T)) + tl.store(p_Ai_44, b_Ai_44.to(p_Ai_44.dtype.element_ty, fp_downcast_rounding="rtne"), mask=(o_t[:, None] + 48 < T)) + tl.store(p_Ai_21, b_Ai_21.to(p_Ai_21.dtype.element_ty, fp_downcast_rounding="rtne"), mask=(o_t[:, None] + 16 < T)) + tl.store(p_Ai_31, b_Ai_31.to(p_Ai_31.dtype.element_ty, fp_downcast_rounding="rtne"), mask=(o_t[:, None] + 32 < T)) + tl.store(p_Ai_32, b_Ai_32.to(p_Ai_32.dtype.element_ty, fp_downcast_rounding="rtne"), mask=(o_t[:, None] + 32 < T)) + tl.store(p_Ai_41, b_Ai_41.to(p_Ai_41.dtype.element_ty, fp_downcast_rounding="rtne"), mask=(o_t[:, None] + 48 < T)) + tl.store(p_Ai_42, b_Ai_42.to(p_Ai_42.dtype.element_ty, fp_downcast_rounding="rtne"), mask=(o_t[:, None] + 48 < T)) + tl.store(p_Ai_43, b_Ai_43.to(p_Ai_43.dtype.element_ty, fp_downcast_rounding="rtne"), mask=(o_t[:, None] + 48 < T)) else: desc_o.store([i_t * BT + 0, 0], b_Ai_11.to(desc_o.dtype, fp_downcast_rounding="rtne")) desc_o.store([i_t * BT + 16, 16], b_Ai_22.to(desc_o.dtype, fp_downcast_rounding="rtne")) @@ -345,12 +356,13 @@ def merge_16x16_to_64x64_inverse_kernel( desc_o.store([i_t * BT + 48, 32], b_Ai_43.to(desc_o.dtype, fp_downcast_rounding="rtne")) +@dispatch @input_guard def solve_tril( A: torch.Tensor, cu_seqlens: torch.Tensor | None = None, chunk_indices: torch.LongTensor | None = None, - output_dtype: torch.dtype = torch.float, + output_dtype: torch.dtype | None = torch.float, ) -> torch.Tensor: """ Compute the inverse of the matrix I + A diff --git a/build/torch-cuda/ops/wall_attn/decode.py b/build/torch-cuda/ops/wall_attn/decode.py index 10a940e040cbf92db7fd21c215bb12bb76247b1c..8b81fec2d71427926470cf6eb3d80a8a283fbdf2 100644 --- a/build/torch-cuda/ops/wall_attn/decode.py +++ b/build/torch-cuda/ops/wall_attn/decode.py @@ -7,14 +7,14 @@ # # Wall attention, contributed by Tilde Research (Timor Averbuch, Dhruv Pai). # Heavily modified from fla/ops/gated_oja_rule/chunk.py. -"""Wall single-step decode kernel and pre-rescaled KV-cache builder.""" +"""Wall cached decode kernel and pre-rescaled KV-cache builder.""" import torch import triton import triton.language as tl from ...ops.utils.op import exp2, log2 -from ...utils import check_shared_mem +from ...utils import autotune_cache_kwargs, check_shared_mem, input_guard WALL_DECODE_AUTOTUNE_CONFIGS = [ triton.Config({'BT': BT}, num_warps=nw, num_stages=ns) @@ -26,13 +26,15 @@ WALL_DECODE_AUTOTUNE_CONFIGS = [ @triton.autotune( configs=WALL_DECODE_AUTOTUNE_CONFIGS, - key=['T_kv', 'K', 'V', 'HQ', 'H', 'C'], + key=['KV_CHUNKS_BUCKET', 'K', 'V', 'HQ', 'H', 'C'], + **autotune_cache_kwargs, ) @triton.heuristics({ 'USE_SINK_BIAS': lambda args: args['sink_bias'] is not None, 'USE_SCALAR_G': lambda args: args['g_scalar_cumsum'] is not None, + 'USE_WINDOW': lambda args: args['W'] is not None, }) -@triton.jit +@triton.jit(do_not_specialize=['T_kv', 'NC', 'KV_CHUNKS_BUCKET']) def parallel_wall_attn_decode_kernel( q, # [B, T_q, HQ, K] k_tilde, # [B, T_kv, HQ, K] pre-rescaled keys (per-Q-head) @@ -47,6 +49,8 @@ def parallel_wall_attn_decode_kernel( T_q, T_kv, NC, + KV_CHUNKS_BUCKET, + W: tl.constexpr, B: tl.constexpr, H: tl.constexpr, HQ: tl.constexpr, @@ -60,8 +64,11 @@ def parallel_wall_attn_decode_kernel( C: tl.constexpr, USE_SINK_BIAS: tl.constexpr, USE_SCALAR_G: tl.constexpr, + USE_WINDOW: tl.constexpr, ): - i_v, i_t, i_bh = tl.program_id(0), tl.program_id(1), tl.program_id(2) + pid = tl.program_id(0).to(tl.int64) + NV, NT = tl.cdiv(V, BV), tl.cdiv(T_q, BT) + i_v, i_t, i_bh = pid % NV, ((pid // NV) % NT).to(tl.int64), (pid // (NV * NT)).to(tl.int64) i_b, i_hq = i_bh // HQ, i_bh % HQ i_h = i_hq // G RCP_LN2: tl.constexpr = 1.4426950216 @@ -70,32 +77,22 @@ def parallel_wall_attn_decode_kernel( bos_kv = (i_b * T_kv).to(tl.int64) bos_nc = (i_b * NC).to(tl.int64) - p_q = tl.make_block_ptr( - q + (bos_q * HQ + i_hq) * K, (T_q, K), (HQ * K, 1), - (i_t * BT, 0), (BT, BK), (1, 0), - ) - p_pq = tl.make_block_ptr( - p_curr + (bos_q * HQ + i_hq) * K, (T_q, K), (HQ * K, 1), - (i_t * BT, 0), (BT, BK), (1, 0), - ) - p_o = tl.make_block_ptr( - o + (bos_q * HQ + i_hq) * V, (T_q, V), (HQ * V, 1), - (i_t * BT, i_v * BV), (BT, BV), (1, 0), - ) - p_lse = tl.make_block_ptr( - lse + bos_q * HQ + i_hq, (T_q,), (HQ,), (i_t * BT,), (BT,), (0,), - ) + o_q = i_t * BT + tl.arange(0, BT) + # queries occupy the last T_q positions of the cache + o_q_global = T_kv - T_q + o_q + o_d = tl.arange(0, BK) + o_v = i_v * BV + tl.arange(0, BV) + m_q = o_q < T_q + p_q = q + (bos_q * HQ + i_hq) * K + o_q[:, None] * (HQ * K) + o_d[None, :] + p_pq = p_curr + (bos_q * HQ + i_hq) * K + o_q[:, None] * (HQ * K) + o_d[None, :] + p_o = o + (bos_q * HQ + i_hq) * V + o_q[:, None] * (HQ * V) + o_v[None, :] + p_lse = lse + bos_q * HQ + i_hq + o_q * HQ - b_q = tl.load(p_q, boundary_check=(0, 1)) - b_pq = tl.load(p_pq, boundary_check=(0, 1)).to(tl.float32) + b_q = tl.load(p_q, mask=m_q[:, None] & (o_d[None, :] < K), other=0.0) + b_pq = tl.load(p_pq, mask=m_q[:, None] & (o_d[None, :] < K), other=0.0).to(tl.float32) if USE_SCALAR_G: - o_q_global = (T_kv - T_q) + i_t * BT + tl.arange(0, BT) - m_q = o_q_global < T_kv - b_cq = tl.load( - g_scalar_cumsum + (bos_kv + o_q_global) * HQ + i_hq, - mask=m_q, other=0, - ).to(tl.float32) + b_cq = tl.load(g_scalar_cumsum + (bos_kv + o_q_global) * HQ + i_hq, mask=m_q, other=0).to(tl.float32) if USE_SINK_BIAS: b_sink_bias = tl.load(sink_bias + i_hq).to(tl.float32) * RCP_LN2 @@ -104,40 +101,35 @@ def parallel_wall_attn_decode_kernel( b_m = tl.full([BT], float('-inf'), dtype=tl.float32) b_acc = tl.zeros([BT], dtype=tl.float32) - # Iterate over the KV cache one chunk at a time: BS == C - for i_s in range(0, T_kv, BS): + i_end = tl.minimum(T_kv, T_kv - T_q + (i_t + 1) * BT) + i_start = tl.maximum((T_kv - T_q + i_t * BT - W + 1) // BS * BS, 0) if USE_WINDOW else 0 + + for i_s in range(i_start, i_end, BS): i_c = i_s // C - # Load anchor - p_r = tl.make_block_ptr( - r_cache + (bos_nc * HQ + i_hq) * K, (NC, K), (HQ * K, 1), - (i_c, 0), (1, BK), (1, 0), - ) - b_R = tl.load(p_r, boundary_check=(0, 1)).to(tl.float32) - b_q_til = (b_q.to(tl.float32) * exp2(b_pq - b_R)).to(b_q.dtype) - - p_kt = tl.make_block_ptr( - k_tilde + (bos_kv * HQ + i_hq) * K, (K, T_kv), (1, HQ * K), - (0, i_s), (BK, BS), (0, 1), - ) - p_v = tl.make_block_ptr( - v + (bos_kv * H + i_h) * V, (T_kv, V), (H * V, 1), - (i_s, i_v * BV), (BS, BV), (1, 0), - ) - b_kt = tl.load(p_kt, boundary_check=(0, 1)) - b_v = tl.load(p_v, boundary_check=(0, 1)) + o_k = (i_s + tl.arange(0, BS)).to(tl.int64) + m_k = o_k < T_kv + o_c = (i_c + tl.arange(0, 1)).to(tl.int64) + p_r = r_cache + (bos_nc * HQ + i_hq) * K + o_c[:, None] * (HQ * K) + o_d[None, :] + b_R = tl.load(p_r, mask=(o_c[:, None] < NC) & (o_d[None, :] < K), other=0.0).to(tl.float32) + # dot in the cache's dtype: the rescale runs in fp32, so casting down to the + # (range-safe) storage dtype here loses only mantissa, matching training. + b_q_til = (b_q.to(tl.float32) * exp2(b_pq - b_R)).to(k_tilde.dtype.element_ty) + + p_kt = k_tilde + (bos_kv * HQ + i_hq) * K + o_d[:, None] + o_k[None, :] * (HQ * K) + p_v = v + (bos_kv * H + i_h) * V + o_k[:, None] * (H * V) + o_v[None, :] + b_kt = tl.load(p_kt, mask=(o_d[:, None] < K) & m_k[None, :], other=0.0) + b_v = tl.load(p_v, mask=m_k[:, None] & (o_v[None, :] < V), other=0.0) b_s = tl.dot(b_q_til, b_kt) * scale * RCP_LN2 - o_k = i_s + tl.arange(0, BS) - m_k = o_k < T_kv if USE_SCALAR_G: - b_ck = tl.load( - g_scalar_cumsum + (bos_kv + o_k) * HQ + i_hq, - mask=m_k, other=0, - ).to(tl.float32) + b_ck = tl.load(g_scalar_cumsum + (bos_kv + o_k) * HQ + i_hq, mask=m_k, other=0).to(tl.float32) b_s += b_cq[:, None] - b_ck[None, :] - b_s = tl.where(m_k[None, :], b_s, float('-inf')) + m_s = m_k[None, :] & (o_k[None, :] <= o_q_global[:, None]) + if USE_WINDOW: + m_s = m_s & (o_q_global[:, None] - o_k[None, :] < W) + b_s = tl.where(m_s, b_s, float('-inf')) b_m, b_mp = tl.maximum(b_m, tl.max(b_s, 1)), b_m b_mw = tl.where(b_m == float('-inf'), 0., b_m) @@ -152,10 +144,13 @@ def parallel_wall_attn_decode_kernel( b_o = b_o / b_acc[:, None] b_m += log2(b_acc) - tl.store(p_o, b_o.to(p_o.dtype.element_ty), boundary_check=(0, 1)) - tl.store(p_lse, b_m.to(p_lse.dtype.element_ty), boundary_check=(0,)) + tl.store(p_o, b_o.to(p_o.dtype.element_ty), mask=m_q[:, None] & (o_v[None, :] < V)) + if i_v == 0: + tl.store(p_lse, b_m.to(p_lse.dtype.element_ty), mask=m_q) +# launch on q's device rather than the current one (e.g. models sharded with device_map='auto'). +@input_guard(no_guard_contiguous=True) def parallel_wall_attn_decode( q: torch.Tensor, v: torch.Tensor, @@ -166,6 +161,7 @@ def parallel_wall_attn_decode( scale: float, cache_chunk_size: int, g_scalar_cumsum: torch.Tensor | None = None, + window_size: int | None = None, ): r"""Wall decode with a pre-rescaled KV cache. @@ -181,6 +177,10 @@ def parallel_wall_attn_decode( r_cache : ``[B, NC, HQ, K]`` per-chunk anchors, ``NC = ceil(T_kv / C)`` g_scalar_cumsum : ``[B, T_kv, HQ]`` or ``None`` optional FoX-style scalar gate + The queries are the last ``T_q`` tokens of the cache: + query ``i`` attends to positions up to ``T_kv - T_q + i``, + limited to the last ``window_size`` positions when provided. + Returns ``(o, lse)`` with ``o.shape == [B, T_q, HQ, V]``. """ C = cache_chunk_size @@ -188,6 +188,8 @@ def parallel_wall_attn_decode( _, T_kv, H, V = v.shape NC = r_cache.shape[1] G = HQ // H + # the bucket selects an autotune config; exact T_kv determines the causal bounds + KV_CHUNKS_BUCKET = triton.next_power_of_2(max(1, triton.cdiv(T_kv, C))) if k_tilde.shape != (B, T_kv, HQ, K): raise ValueError(f"k_tilde shape {k_tilde.shape} != expected {(B, T_kv, HQ, K)}") @@ -222,7 +224,7 @@ def parallel_wall_attn_decode( lse = torch.empty(B, T_q, HQ, dtype=torch.float, device=q.device) def grid(meta): - return (NV, triton.cdiv(T_q, meta['BT']), B * HQ) + return (NV * triton.cdiv(T_q, meta['BT']) * B * HQ,) parallel_wall_attn_decode_kernel[grid]( q=q, k_tilde=k_tilde, @@ -237,6 +239,8 @@ def parallel_wall_attn_decode( T_q=T_q, T_kv=T_kv, NC=NC, + KV_CHUNKS_BUCKET=KV_CHUNKS_BUCKET, + W=window_size, B=B, H=H, HQ=HQ, diff --git a/build/torch-cuda/ops/wall_attn/parallel.py b/build/torch-cuda/ops/wall_attn/parallel.py index 9443838331279604f6e1c57e836a8f20c0b3a834..8dc229c21f5ebfb2a45a298cdf1f87fbfe022d40 100644 --- a/build/torch-cuda/ops/wall_attn/parallel.py +++ b/build/torch-cuda/ops/wall_attn/parallel.py @@ -16,13 +16,12 @@ import triton import triton.language as tl from einops import reduce -from ...ops.attn.parallel import parallel_attn_bwd_preprocess -from ...ops.backends import dispatch +from ...backends import dispatch from ...ops.utils import prepare_chunk_indices from ...ops.utils.constant import RCP_LN2 from ...ops.utils.cumsum import chunk_global_cumsum -from ...ops.utils.op import exp2, log2 -from ...utils import autocast_custom_bwd, autocast_custom_fwd, check_shared_mem, contiguous +from ...ops.utils.op import exp2, log2, unflatten_program_id +from ...utils import autocast_custom_bwd, autocast_custom_fwd, autotune_cache_kwargs, check_shared_mem, contiguous _DEBUG_ASSERTS = os.environ.get("WALL_ATTN_DEBUG", "0") == "1" @@ -47,9 +46,45 @@ WALL_BWD_AUTOTUNE_CONFIGS = [ ] +def _prune_wall_fwd_configs(configs, nargs, **kwargs): + # The per-program fp32 accumulator tile is BT x BV. At the Hopper BV cap of + # 256, BT=128 configs hold a 128x256 fp32 accumulator (128 KiB), spill + # registers massively, can never win the benchmark, and each take minutes to + # compile under IEEE fp32 dots (TRITON_F32_DEFAULT=ieee in the test suite), + # which reads as a multi-minute stall per autotune key. Drop configs that + # exceed the tile budget; BV <= 128 keeps the full sweep. BV never exceeds + # the fwd cap (256), so BT=64 always survives and the pruned set is non-empty. + nargs = {**(nargs or {}), **kwargs} # keyword launches leave nargs empty + BV = nargs.get('BV') or 128 + return [c for c in configs if c.kwargs['BT'] * BV <= 16384] + + +@triton.jit(do_not_specialize=['N']) +def parallel_wall_attn_bwd_kernel_preprocess( + o, + do, + delta, + N, + V: tl.constexpr, + BV: tl.constexpr, +): + i_nv = tl.program_id(0).to(tl.int64) + i_v, i_n = i_nv // N, i_nv % N + o_v = i_v * BV + tl.arange(0, BV) + m_v = o_v < V + + b_o = tl.load(o + i_n * V + o_v, mask=m_v, other=0.).to(tl.float32) + b_do = tl.load(do + i_n * V + o_v, mask=m_v, other=0.).to(tl.float32) + b_delta = tl.sum(b_o * b_do) + + tl.store(delta + i_v * N + i_n, b_delta) + + @triton.autotune( configs=WALL_FWD_AUTOTUNE_CONFIGS, - key=['T', 'K', 'V', 'HQ', 'H'], + key=['T_BUCKET', 'K', 'V', 'HQ', 'H'], + prune_configs_by={'early_config_prune': _prune_wall_fwd_configs}, + **autotune_cache_kwargs, ) @triton.heuristics({ 'USE_SINK_BIAS': lambda args: args['sink_bias'] is not None, @@ -57,7 +92,7 @@ WALL_BWD_AUTOTUNE_CONFIGS = [ 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None, 'USE_SCALAR_G': lambda args: args['g_scalar_cumsum'] is not None, }) -@triton.jit +@triton.jit(do_not_specialize=['T', 'T_BUCKET']) def parallel_wall_attn_fwd_kernel( q, k, @@ -71,6 +106,7 @@ def parallel_wall_attn_fwd_kernel( cu_seqlens, chunk_indices, T, + T_BUCKET, W: tl.constexpr, B: tl.constexpr, H: tl.constexpr, @@ -87,12 +123,13 @@ def parallel_wall_attn_fwd_kernel( IS_VARLEN: tl.constexpr, USE_SCALAR_G: tl.constexpr, ): - i_v, i_t, i_bh = tl.program_id(0), tl.program_id(1), tl.program_id(2) + i_v, i_t = unflatten_program_id(tl.cdiv(V, BV)) + i_bh = tl.program_id(1).to(tl.int64) i_b, i_hq = i_bh // HQ, i_bh % HQ i_h = i_hq // G 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) + i_n, i_t = tl.load(chunk_indices + i_t * 2).to(tl.int32), tl.load(chunk_indices + i_t * 2 + 1).to(tl.int64) bos, eos = tl.load(cu_seqlens + i_n).to(tl.int64), tl.load(cu_seqlens + i_n + 1).to(tl.int64) T = (eos - bos).to(tl.int32) else: @@ -100,16 +137,22 @@ def parallel_wall_attn_fwd_kernel( bos, eos = (i_b * T).to(tl.int64), (i_b * T + T).to(tl.int64) RCP_LN2: tl.constexpr = 1.4426950216 - p_q = tl.make_block_ptr(q + (bos * HQ + i_hq) * K, (T, K), (HQ*K, 1), (i_t * BT, 0), (BT, BK), (1, 0)) - p_o = tl.make_block_ptr(o + (bos * HQ + i_hq) * V, (T, V), (HQ*V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) - p_lse = tl.make_block_ptr(lse + bos * HQ + i_hq, (T,), (HQ,), (i_t * BT,), (BT,), (0,)) - - p_pq = tl.make_block_ptr(g_cumsum + (bos * HQ + i_hq) * K, (T, K), (HQ*K, 1), (i_t * BT, 0), (BT, BK), (1, 0)) - p_R = tl.make_block_ptr(g_cumsum + (bos * HQ + i_hq) * K, (T, K), (HQ*K, 1), (i_t * BT, 0), (1, BK), (1, 0)) - - b_q = tl.load(p_q, boundary_check=(0, 1)) - b_pq = tl.load(p_pq, boundary_check=(0, 1)).to(tl.float32) - b_R = tl.load(p_R, boundary_check=(0, 1)).to(tl.float32) + o_q = i_t * BT + tl.arange(0, BT) + o_d = tl.arange(0, BK) + o_v = i_v * BV + tl.arange(0, BV) + m_q = o_q < T + m_qk = m_q[:, None] & (o_d[None, :] < K) + p_q = q + (bos * HQ + i_hq) * K + o_q[:, None] * (HQ*K) + o_d[None, :] + p_o = o + (bos * HQ + i_hq) * V + o_q[:, None] * (HQ*V) + o_v[None, :] + p_lse = lse + bos * HQ + i_hq + o_q * HQ + + p_pq = g_cumsum + (bos * HQ + i_hq) * K + o_q[:, None] * (HQ*K) + o_d[None, :] + o_r = i_t * BT + tl.arange(0, 1) + p_R = g_cumsum + (bos * HQ + i_hq) * K + o_r[:, None] * (HQ*K) + o_d[None, :] + + b_q = tl.load(p_q, mask=m_qk, other=0.0) + b_pq = tl.load(p_pq, mask=m_qk, other=0.0).to(tl.float32) + b_R = tl.load(p_R, mask=(o_r[:, None] < T) & (o_d[None, :] < K), other=0.0).to(tl.float32) # b_R is at i_t*BT; the same value is used in both off-diag and diag loops, # since |b_pq - b_R| within a BT-chunk is bounded by BT*|g_max|*RCP_LN2. # Off-diagonal: P_q - R <= 0 and R - P_k <= 0 -> exp2 <= 1 -> bf16 safe. @@ -120,32 +163,31 @@ def parallel_wall_attn_fwd_kernel( b_acc = tl.zeros([BT], dtype=tl.float32) if USE_SCALAR_G: - p_cq = tl.make_block_ptr(g_scalar_cumsum + bos * HQ + i_hq, (T,), (HQ,), (i_t * BT,), (BT,), (0,)) - b_cq = tl.load(p_cq, boundary_check=(0,)).to(tl.float32) + p_cq = g_scalar_cumsum + bos * HQ + i_hq + o_q * HQ + b_cq = tl.load(p_cq, mask=m_q, other=0.0).to(tl.float32) if USE_SINK_BIAS: b_sink_bias = tl.load(sink_bias + i_hq).to(tl.float32) else: b_sink_bias = None - o_q = i_t * BT + tl.arange(0, BT) i_start = tl.maximum((i_t * BT - W + 1) // BS * BS, 0) if USE_WINDOW else 0 b_R_t = tl.trans(b_R) for i_s in range(i_start, i_t * BT, BS): - p_k = tl.make_block_ptr(k + (bos * H + i_h) * K, (K, T), (1, H*K), (0, i_s), (BK, BS), (0, 1)) - p_v = tl.make_block_ptr(v + (bos * H + i_h) * V, (T, V), (H*V, 1), (i_s, i_v * BV), (BS, BV), (1, 0)) - p_pk = tl.make_block_ptr(g_cumsum + (bos * HQ + i_hq) * K, (K, T), (1, HQ*K), (0, i_s), (BK, BS), (0, 1)) + o_k = i_s + tl.arange(0, BS) + m_k = o_k < T + p_k = k + (bos * H + i_h) * K + o_d[:, None] + o_k[None, :] * (H*K) + p_v = v + (bos * H + i_h) * V + o_k[:, None] * (H*V) + o_v[None, :] + p_pk = g_cumsum + (bos * HQ + i_hq) * K + o_d[:, None] + o_k[None, :] * (HQ*K) - b_k = tl.load(p_k, boundary_check=(0, 1)) - b_v = tl.load(p_v, boundary_check=(0, 1)) - b_pk = tl.load(p_pk, boundary_check=(0, 1)).to(tl.float32) + b_k = tl.load(p_k, mask=(o_d[:, None] < K) & m_k[None, :], other=0.0) + b_v = tl.load(p_v, mask=m_k[:, None] & (o_v[None, :] < V), other=0.0) + b_pk = tl.load(p_pk, mask=(o_d[:, None] < K) & m_k[None, :], other=0.0).to(tl.float32) b_k_til = (b_k.to(tl.float32) * exp2(b_R_t - b_pk)).to(b_k.dtype) b_s = tl.dot(b_q_til, b_k_til) * scale * RCP_LN2 - o_k = i_s + tl.arange(0, BS) - m_k = o_k < T if USE_SCALAR_G: b_ck = tl.load(g_scalar_cumsum + (bos + o_k) * HQ + i_hq, mask=m_k, other=0).to(tl.float32) b_s += b_cq[:, None] - b_ck[None, :] @@ -161,21 +203,22 @@ def parallel_wall_attn_fwd_kernel( b_mp = b_m for i_s in range(i_t * BT, min((i_t + 1) * BT, T), BS): - p_k = tl.make_block_ptr(k + (bos * H + i_h) * K, (K, T), (1, H*K), (0, i_s), (BK, BS), (0, 1)) - p_v = tl.make_block_ptr(v + (bos * H + i_h) * V, (T, V), (H*V, 1), (i_s, i_v * BV), (BS, BV), (1, 0)) - p_pk = tl.make_block_ptr(g_cumsum + (bos * HQ + i_hq) * K, (K, T), (1, HQ*K), (0, i_s), (BK, BS), (0, 1)) + o_k = i_s + tl.arange(0, BS) + m_k = o_k < T + p_k = k + (bos * H + i_h) * K + o_d[:, None] + o_k[None, :] * (H*K) + p_v = v + (bos * H + i_h) * V + o_k[:, None] * (H*V) + o_v[None, :] + p_pk = g_cumsum + (bos * HQ + i_hq) * K + o_d[:, None] + o_k[None, :] * (HQ*K) # Per-sub-block local reference to prevent exp2 overflow. - p_R_local = tl.make_block_ptr(g_cumsum + (bos * HQ + i_hq) * K, (T, K), (HQ*K, 1), (i_s, 0), (1, BK), (1, 0)) - b_R_local = tl.load(p_R_local, boundary_check=(0, 1)).to(tl.float32) + o_r = i_s + tl.arange(0, 1) + p_R_local = g_cumsum + (bos * HQ + i_hq) * K + o_r[:, None] * (HQ*K) + o_d[None, :] + b_R_local = tl.load(p_R_local, mask=(o_r[:, None] < T) & (o_d[None, :] < K), other=0.0).to(tl.float32) b_R_local_bc = tl.broadcast_to(b_R_local, (BT, BK)) b_R_local_t = tl.trans(b_R_local) - o_k = i_s + tl.arange(0, BS) - m_k = o_k < T - b_k = tl.load(p_k, boundary_check=(0, 1)) - b_v = tl.load(p_v, boundary_check=(0, 1)) - b_pk = tl.load(p_pk, boundary_check=(0, 1)).to(tl.float32) + b_k = tl.load(p_k, mask=(o_d[:, None] < K) & m_k[None, :], other=0.0) + b_v = tl.load(p_v, mask=m_k[:, None] & (o_v[None, :] < V), other=0.0) + b_pk = tl.load(p_pk, mask=(o_d[:, None] < K) & m_k[None, :], other=0.0).to(tl.float32) # k_til in local frame: exp2 bounded by BS positions <= 110 (fp32-safe with BK dot). b_exp_k = b_R_local_t - b_pk @@ -210,20 +253,22 @@ def parallel_wall_attn_fwd_kernel( b_o = b_o / b_acc[:, None] b_m += log2(b_acc) - tl.store(p_o, b_o.to(p_o.dtype.element_ty), boundary_check=(0, 1)) - tl.store(p_lse, b_m.to(p_lse.dtype.element_ty), boundary_check=(0,)) + tl.store(p_o, b_o.to(p_o.dtype.element_ty), mask=m_q[:, None] & (o_v[None, :] < V)) + if i_v == 0: + tl.store(p_lse, b_m.to(p_lse.dtype.element_ty), mask=m_q) @triton.autotune( configs=WALL_BWD_AUTOTUNE_CONFIGS, - key=['T', 'K', 'V', 'HQ', 'H'], + key=['T_BUCKET', 'K', 'V', 'HQ', 'H'], + **autotune_cache_kwargs, ) @triton.heuristics({ 'USE_WINDOW': lambda args: args['W'] is not None, 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None, 'USE_SCALAR_G': lambda args: args['g_scalar_cumsum'] is not None, }) -@triton.jit(do_not_specialize=['T']) +@triton.jit(do_not_specialize=['T', 'T_BUCKET']) def parallel_wall_attn_bwd_kernel_dq( q, k, @@ -240,6 +285,7 @@ def parallel_wall_attn_bwd_kernel_dq( cu_seqlens, chunk_indices, T, + T_BUCKET, W: tl.constexpr, B: tl.constexpr, H: tl.constexpr, @@ -255,12 +301,20 @@ def parallel_wall_attn_bwd_kernel_dq( IS_VARLEN: tl.constexpr, USE_SCALAR_G: tl.constexpr, ): - i_v, i_t, i_bh = tl.program_id(0), tl.program_id(1), tl.program_id(2) + i_v, i_t = unflatten_program_id(tl.cdiv(V, BV)) + i_bh = tl.program_id(1).to(tl.int64) i_b, i_hq = i_bh // HQ, i_bh % HQ i_h = i_hq // G + i_v64 = i_v.to(tl.int64) + delta += i_v64 * B * T * HQ + dq += i_v64 * B * T * HQ * K + dg_cumsum += i_v64 * B * T * HQ * K + if USE_SCALAR_G: + dg_scalar_cumsum += i_v64 * B * T * HQ + 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) + i_n, i_t = tl.load(chunk_indices + i_t * 2).to(tl.int32), tl.load(chunk_indices + i_t * 2 + 1).to(tl.int64) bos, eos = tl.load(cu_seqlens + i_n).to(tl.int64), tl.load(cu_seqlens + i_n + 1).to(tl.int64) T = (eos - bos).to(tl.int32) else: @@ -269,46 +323,52 @@ def parallel_wall_attn_bwd_kernel_dq( RCP_LN2: tl.constexpr = 1.4426950216 LN2: tl.constexpr = 0.6931471805599453 - p_q = tl.make_block_ptr(q + (bos * HQ + i_hq) * K, (T, K), (HQ*K, 1), (i_t * BT, 0), (BT, BK), (1, 0)) - p_pq = tl.make_block_ptr(g_cumsum + (bos * HQ + i_hq) * K, (T, K), (HQ*K, 1), (i_t * BT, 0), (BT, BK), (1, 0)) - p_R = tl.make_block_ptr(g_cumsum + (bos * HQ + i_hq) * K, (T, K), (HQ*K, 1), (i_t * BT, 0), (1, BK), (1, 0)) - p_dq = tl.make_block_ptr(dq + (bos * HQ + i_hq) * K, (T, K), (HQ*K, 1), (i_t * BT, 0), (BT, BK), (1, 0)) - p_dg = tl.make_block_ptr(dg_cumsum + (bos * HQ + i_hq) * K, (T, K), (HQ*K, 1), (i_t * BT, 0), (BT, BK), (1, 0)) - p_do = tl.make_block_ptr(do + (bos * HQ + i_hq) * V, (T, V), (HQ*V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) - p_lse = tl.make_block_ptr(lse + bos * HQ + i_hq, (T,), (HQ,), (i_t * BT,), (BT,), (0,)) - p_delta = tl.make_block_ptr(delta + bos * HQ + i_hq, (T,), (HQ,), (i_t * BT,), (BT,), (0,)) - - b_q = tl.load(p_q, boundary_check=(0, 1)) - b_pq = tl.load(p_pq, boundary_check=(0, 1)).to(tl.float32) - b_R = tl.load(p_R, boundary_check=(0, 1)).to(tl.float32) + o_q = i_t * BT + tl.arange(0, BT) + o_d = tl.arange(0, BK) + o_v = i_v * BV + tl.arange(0, BV) + m_q = o_q < T + m_qk = m_q[:, None] & (o_d[None, :] < K) + m_o = m_q[:, None] & (o_v[None, :] < V) + p_q = q + (bos * HQ + i_hq) * K + o_q[:, None] * (HQ*K) + o_d[None, :] + p_pq = g_cumsum + (bos * HQ + i_hq) * K + o_q[:, None] * (HQ*K) + o_d[None, :] + o_r = i_t * BT + tl.arange(0, 1) + p_R = g_cumsum + (bos * HQ + i_hq) * K + o_r[:, None] * (HQ*K) + o_d[None, :] + p_dq = dq + (bos * HQ + i_hq) * K + o_q[:, None] * (HQ*K) + o_d[None, :] + p_dg = dg_cumsum + (bos * HQ + i_hq) * K + o_q[:, None] * (HQ*K) + o_d[None, :] + p_do = do + (bos * HQ + i_hq) * V + o_q[:, None] * (HQ*V) + o_v[None, :] + p_lse = lse + bos * HQ + i_hq + o_q * HQ + p_delta = delta + bos * HQ + i_hq + o_q * HQ + + b_q = tl.load(p_q, mask=m_qk, other=0.0) + b_pq = tl.load(p_pq, mask=m_qk, other=0.0).to(tl.float32) + b_R = tl.load(p_R, mask=(o_r[:, None] < T) & (o_d[None, :] < K), other=0.0).to(tl.float32) # Off-diagonal q_til (bf16 tensor cores, exp2 <= 1). b_q_til = (b_q.to(tl.float32) * exp2(b_pq - b_R)).to(b_q.dtype) - b_do = tl.load(p_do, boundary_check=(0, 1)) - b_lse = tl.load(p_lse, boundary_check=(0,)) - b_delta = tl.load(p_delta, boundary_check=(0,)) + b_do = tl.load(p_do, mask=m_o, other=0.0) + b_lse = tl.load(p_lse, mask=m_q, other=0.0) + b_delta = tl.load(p_delta, mask=m_q, other=0.0) b_dq_til = tl.zeros([BT, BK], dtype=tl.float32) if USE_SCALAR_G: - p_cq = tl.make_block_ptr(g_scalar_cumsum + bos * HQ + i_hq, (T,), (HQ,), (i_t * BT,), (BT,), (0,)) - b_cq = tl.load(p_cq, boundary_check=(0,)).to(tl.float32) + p_cq = g_scalar_cumsum + bos * HQ + i_hq + o_q * HQ + b_cq = tl.load(p_cq, mask=m_q, other=0.0).to(tl.float32) b_dc = tl.zeros([BT], dtype=tl.float32) - o_q = i_t * BT + tl.arange(0, BT) i_start = tl.maximum((i_t * BT - W + 1) // BS * BS, 0) if USE_WINDOW else 0 b_R_t = tl.trans(b_R) for i_s in range(i_start, i_t * BT, BS): - p_k = tl.make_block_ptr(k + (bos * H + i_h) * K, (K, T), (1, H*K), (0, i_s), (BK, BS), (0, 1)) - p_v = tl.make_block_ptr(v + (bos * H + i_h) * V, (V, T), (1, H*V), (i_v * BV, i_s), (BV, BS), (0, 1)) - p_pk = tl.make_block_ptr(g_cumsum + (bos * HQ + i_hq) * K, (K, T), (1, HQ*K), (0, i_s), (BK, BS), (0, 1)) - o_k = i_s + tl.arange(0, BS) m_k = o_k < T - b_k = tl.load(p_k, boundary_check=(0, 1)) - b_pk = tl.load(p_pk, boundary_check=(0, 1)).to(tl.float32) - b_v = tl.load(p_v, boundary_check=(0, 1)) + p_k = k + (bos * H + i_h) * K + o_d[:, None] + o_k[None, :] * (H*K) + p_v = v + (bos * H + i_h) * V + o_v[:, None] + o_k[None, :] * (H*V) + p_pk = g_cumsum + (bos * HQ + i_hq) * K + o_d[:, None] + o_k[None, :] * (HQ*K) + + b_k = tl.load(p_k, mask=(o_d[:, None] < K) & m_k[None, :], other=0.0) + b_pk = tl.load(p_pk, mask=(o_d[:, None] < K) & m_k[None, :], other=0.0).to(tl.float32) + b_v = tl.load(p_v, mask=(o_v[:, None] < V) & m_k[None, :], other=0.0) b_k_til = (b_k.to(tl.float32) * exp2(b_R_t - b_pk)).to(b_k.dtype) b_s = tl.dot(b_q_til, b_k_til) * scale * RCP_LN2 @@ -320,7 +380,7 @@ def parallel_wall_attn_bwd_kernel_dq( b_p = exp2(b_s - b_lse[:, None]) b_dp = tl.dot(b_do, b_v) b_ds = b_p * (b_dp.to(tl.float32) - b_delta[:, None]) - b_dq_til += tl.dot(b_ds.to(b_k.dtype), tl.trans(b_k_til)) + b_dq_til = tl.dot(b_ds.to(b_k.dtype), tl.trans(b_k_til), b_dq_til) if USE_SCALAR_G: b_dc += tl.sum(b_ds, 1) @@ -329,20 +389,21 @@ def parallel_wall_attn_bwd_kernel_dq( b_dq_diag = tl.zeros([BT, BK], dtype=tl.float32) for i_s in range(i_t * BT, min((i_t + 1) * BT, T), BS): - p_k = tl.make_block_ptr(k + (bos * H + i_h) * K, (K, T), (1, H*K), (0, i_s), (BK, BS), (0, 1)) - p_v = tl.make_block_ptr(v + (bos * H + i_h) * V, (V, T), (1, H*V), (i_v * BV, i_s), (BV, BS), (0, 1)) - p_pk = tl.make_block_ptr(g_cumsum + (bos * HQ + i_hq) * K, (K, T), (1, HQ*K), (0, i_s), (BK, BS), (0, 1)) + o_k = i_s + tl.arange(0, BS) + m_k = o_k < T + p_k = k + (bos * H + i_h) * K + o_d[:, None] + o_k[None, :] * (H*K) + p_v = v + (bos * H + i_h) * V + o_v[:, None] + o_k[None, :] * (H*V) + p_pk = g_cumsum + (bos * HQ + i_hq) * K + o_d[:, None] + o_k[None, :] * (HQ*K) - p_R_local = tl.make_block_ptr(g_cumsum + (bos * HQ + i_hq) * K, (T, K), (HQ*K, 1), (i_s, 0), (1, BK), (1, 0)) - b_R_local = tl.load(p_R_local, boundary_check=(0, 1)).to(tl.float32) + o_r = i_s + tl.arange(0, 1) + p_R_local = g_cumsum + (bos * HQ + i_hq) * K + o_r[:, None] * (HQ*K) + o_d[None, :] + b_R_local = tl.load(p_R_local, mask=(o_r[:, None] < T) & (o_d[None, :] < K), other=0.0).to(tl.float32) b_R_local_bc = tl.broadcast_to(b_R_local, (BT, BK)) b_R_local_t = tl.trans(b_R_local) - o_k = i_s + tl.arange(0, BS) - m_k = o_k < T - b_k = tl.load(p_k, boundary_check=(0, 1)) - b_pk = tl.load(p_pk, boundary_check=(0, 1)).to(tl.float32) - b_v = tl.load(p_v, boundary_check=(0, 1)) + b_k = tl.load(p_k, mask=(o_d[:, None] < K) & m_k[None, :], other=0.0) + b_pk = tl.load(p_pk, mask=(o_d[:, None] < K) & m_k[None, :], other=0.0).to(tl.float32) + b_v = tl.load(p_v, mask=(o_v[:, None] < V) & m_k[None, :], other=0.0) # Local-ref k_til: exp2 arg bounded by BS positions <= 110 (fp32-safe with BK dot). b_exp_k = b_R_local_t - b_pk @@ -380,23 +441,24 @@ def parallel_wall_attn_bwd_kernel_dq( # Structural: b_dg derived from b_dq (saves [BT,BK] fp32 accumulator). b_dg = LN2 * b_q.to(tl.float32) * b_dq - tl.store(p_dq, b_dq.to(p_dq.dtype.element_ty), boundary_check=(0, 1)) - tl.store(p_dg, b_dg.to(p_dg.dtype.element_ty), boundary_check=(0, 1)) + tl.store(p_dq, b_dq.to(p_dq.dtype.element_ty), mask=m_qk) + tl.store(p_dg, b_dg.to(p_dg.dtype.element_ty), mask=m_qk) if USE_SCALAR_G: - p_dc = tl.make_block_ptr(dg_scalar_cumsum + bos * HQ + i_hq, (T,), (HQ,), (i_t * BT,), (BT,), (0,)) - tl.store(p_dc, b_dc.to(p_dc.dtype.element_ty), boundary_check=(0,)) + p_dc = dg_scalar_cumsum + bos * HQ + i_hq + o_q * HQ + tl.store(p_dc, b_dc.to(p_dc.dtype.element_ty), mask=m_q) @triton.autotune( configs=WALL_BWD_AUTOTUNE_CONFIGS, - key=['T', 'K', 'V', 'HQ', 'H'], + key=['T_BUCKET', 'K', 'V', 'HQ', 'H'], + **autotune_cache_kwargs, ) @triton.heuristics({ 'USE_WINDOW': lambda args: args['W'] is not None, 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None, 'USE_SCALAR_G': lambda args: args['g_scalar_cumsum'] is not None, }) -@triton.jit(do_not_specialize=['T']) +@triton.jit(do_not_specialize=['T', 'T_BUCKET']) def parallel_wall_attn_bwd_kernel_dkv( q, k, @@ -414,6 +476,7 @@ def parallel_wall_attn_bwd_kernel_dkv( chunk_indices, scale, T, + T_BUCKET, W: tl.constexpr, B: tl.constexpr, H: tl.constexpr, @@ -430,12 +493,20 @@ def parallel_wall_attn_bwd_kernel_dkv( USE_SCALAR_G: tl.constexpr, DIAG_BF16: tl.constexpr, ): - i_v, i_t, i_bh = tl.program_id(0), tl.program_id(1), tl.program_id(2) + i_v, i_t = unflatten_program_id(tl.cdiv(V, BV)) + i_bh = tl.program_id(1).to(tl.int64) i_b, i_hq = i_bh // HQ, i_bh % HQ i_h = i_hq // G + i_v64 = i_v.to(tl.int64) + delta += i_v64 * B * T * HQ + dk += i_v64 * B * T * HQ * K + dg_cumsum += i_v64 * B * T * HQ * K + if USE_SCALAR_G: + dg_scalar_cumsum += i_v64 * B * T * HQ + 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) + i_n, i_t = tl.load(chunk_indices + i_t * 2).to(tl.int32), tl.load(chunk_indices + i_t * 2 + 1).to(tl.int64) bos, eos = tl.load(cu_seqlens + i_n).to(tl.int64), tl.load(cu_seqlens + i_n + 1).to(tl.int64) T = (eos - bos).to(tl.int32) else: @@ -444,41 +515,46 @@ def parallel_wall_attn_bwd_kernel_dkv( RCP_LN2: tl.constexpr = 1.4426950216 LN2: tl.constexpr = 0.6931471805599453 - p_k = tl.make_block_ptr(k + (bos * H + i_h) * K, (T, K), (H*K, 1), (i_t * BT, 0), (BT, BK), (1, 0)) - p_pk = tl.make_block_ptr(g_cumsum + (bos * HQ + i_hq) * K, (T, K), (HQ*K, 1), (i_t * BT, 0), (BT, BK), (1, 0)) - p_v = tl.make_block_ptr(v + (bos * H + i_h) * V, (T, V), (H*V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) - p_dk = tl.make_block_ptr(dk + (bos * HQ + i_hq) * K, (T, K), (HQ*K, 1), (i_t * BT, 0), (BT, BK), (1, 0)) - p_dv = tl.make_block_ptr(dv + (bos * HQ + i_hq) * V, (T, V), (HQ*V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) - p_dg = tl.make_block_ptr(dg_cumsum + (bos * HQ + i_hq) * K, (T, K), (HQ*K, 1), (i_t * BT, 0), (BT, BK), (1, 0)) - - b_k = tl.load(p_k, boundary_check=(0, 1)) - b_pk = tl.load(p_pk, boundary_check=(0, 1)).to(tl.float32) + o_k = i_t * BT + tl.arange(0, BT) + o_d = tl.arange(0, BK) + o_v = i_v * BV + tl.arange(0, BV) + m_k = o_k < T + m_kk = m_k[:, None] & (o_d[None, :] < K) + m_kv = m_k[:, None] & (o_v[None, :] < V) + p_k = k + (bos * H + i_h) * K + o_k[:, None] * (H*K) + o_d[None, :] + p_pk = g_cumsum + (bos * HQ + i_hq) * K + o_k[:, None] * (HQ*K) + o_d[None, :] + p_v = v + (bos * H + i_h) * V + o_k[:, None] * (H*V) + o_v[None, :] + p_dk = dk + (bos * HQ + i_hq) * K + o_k[:, None] * (HQ*K) + o_d[None, :] + p_dv = dv + (bos * HQ + i_hq) * V + o_k[:, None] * (HQ*V) + o_v[None, :] + p_dg = dg_cumsum + (bos * HQ + i_hq) * K + o_k[:, None] * (HQ*K) + o_d[None, :] + + b_k = tl.load(p_k, mask=m_kk, other=0.0) + b_pk = tl.load(p_pk, mask=m_kk, other=0.0).to(tl.float32) b_dk = tl.zeros([BT, BK], dtype=tl.float32) - b_v = tl.load(p_v, boundary_check=(0, 1)) + b_v = tl.load(p_v, mask=m_kv, other=0.0) b_dv = tl.zeros([BT, BV], dtype=tl.float32) - o_k = i_t * BT + tl.arange(0, BT) - if USE_SCALAR_G: - p_ck = tl.make_block_ptr(g_scalar_cumsum + bos * HQ + i_hq, (T,), (HQ,), (i_t * BT,), (BT,), (0,)) - b_ck = tl.load(p_ck, boundary_check=(0,)).to(tl.float32) + p_ck = g_scalar_cumsum + bos * HQ + i_hq + o_k * HQ + b_ck = tl.load(p_ck, mask=m_k, other=0.0).to(tl.float32) b_dc = tl.zeros([BT], dtype=tl.float32) for i_s in range(i_t * BT, min((i_t + 1) * BT, T), BS): - p_Rq = tl.make_block_ptr(g_cumsum + (bos * HQ + i_hq) * K, (T, K), (HQ*K, 1), (i_s, 0), (1, BK), (1, 0)) - p_q = tl.make_block_ptr(q + (bos * HQ + i_hq) * K, (T, K), (HQ*K, 1), (i_s, 0), (BS, BK), (1, 0)) - p_pq = tl.make_block_ptr(g_cumsum + (bos * HQ + i_hq) * K, (T, K), (HQ*K, 1), (i_s, 0), (BS, BK), (1, 0)) - p_do = tl.make_block_ptr(do + (bos * HQ + i_hq) * V, (T, V), (HQ*V, 1), (i_s, i_v * BV), (BS, BV), (1, 0)) - p_lse = tl.make_block_ptr(lse + bos * HQ + i_hq, (T,), (HQ,), (i_s,), (BS,), (0,)) - p_delta = tl.make_block_ptr(delta + bos * HQ + i_hq, (T,), (HQ,), (i_s,), (BS,), (0,)) - o_q = i_s + tl.arange(0, BS) m_q = o_q < T - b_Rq = tl.load(p_Rq, boundary_check=(0, 1)).to(tl.float32) + o_r = i_s + tl.arange(0, 1) + p_Rq = g_cumsum + (bos * HQ + i_hq) * K + o_r[:, None] * (HQ*K) + o_d[None, :] + p_q = q + (bos * HQ + i_hq) * K + o_q[:, None] * (HQ*K) + o_d[None, :] + p_pq = g_cumsum + (bos * HQ + i_hq) * K + o_q[:, None] * (HQ*K) + o_d[None, :] + p_do = do + (bos * HQ + i_hq) * V + o_q[:, None] * (HQ*V) + o_v[None, :] + p_lse = lse + bos * HQ + i_hq + o_q * HQ + p_delta = delta + bos * HQ + i_hq + o_q * HQ + + b_Rq = tl.load(p_Rq, mask=(o_r[:, None] < T) & (o_d[None, :] < K), other=0.0).to(tl.float32) b_Rq_bc_k = tl.broadcast_to(b_Rq, (BT, BK)) b_Rq_bc_q = tl.broadcast_to(b_Rq, (BS, BK)) - b_q = tl.load(p_q, boundary_check=(0, 1)) - b_pq = tl.load(p_pq, boundary_check=(0, 1)).to(tl.float32) + b_q = tl.load(p_q, mask=m_q[:, None] & (o_d[None, :] < K), other=0.0) + b_pq = tl.load(p_pq, mask=m_q[:, None] & (o_d[None, :] < K), other=0.0).to(tl.float32) b_exp_k = b_Rq_bc_k - b_pk # Non-causal keys (key_pos > last query in sub-block) can have # exp_k >> 110, overflowing the gradient headroom. The causal mask @@ -495,9 +571,9 @@ def parallel_wall_attn_bwd_kernel_dkv( # Precise path: fp32 tensor cores. b_q_til = b_q.to(tl.float32) * exp2(b_pq - b_Rq_bc_q) b_k_til = b_k.to(tl.float32) * b_exp_k_val - b_do = tl.load(p_do, boundary_check=(0, 1)) - b_lse = tl.load(p_lse, boundary_check=(0,)) - b_delta = tl.load(p_delta, boundary_check=(0,)) + b_do = tl.load(p_do, mask=m_q[:, None] & (o_v[None, :] < V), other=0.0) + b_lse = tl.load(p_lse, mask=m_q, other=0.0) + b_delta = tl.load(p_delta, mask=m_q, other=0.0) b_s = tl.dot(b_k_til, tl.trans(b_q_til)) * scale * RCP_LN2 if USE_SCALAR_G: b_cq = tl.load(g_scalar_cumsum + (bos + o_q) * HQ + i_hq, mask=m_q, other=0).to(tl.float32) @@ -509,7 +585,7 @@ def parallel_wall_attn_bwd_kernel_dkv( ) else: b_p = tl.where((o_k[:, None] <= o_q[None, :]) & m_q[None, :], exp2(b_s - b_lse[None, :]), 0) - b_dv += tl.dot(b_p.to(b_do.dtype), b_do) + b_dv = tl.dot(b_p.to(b_do.dtype), b_do, b_dv) b_dp = tl.dot(b_v, tl.trans(b_do)) b_ds = b_p * (b_dp - b_delta[None, :]) if DIAG_BF16: @@ -523,27 +599,28 @@ def parallel_wall_attn_bwd_kernel_dkv( i_end = min(tl.cdiv(T, BS) * BS, (i_t + 1) * BT + W - 1) if USE_WINDOW else tl.cdiv(T, BS) * BS for i_s in range((i_t + 1) * BT, i_end, BS): - p_Rq = tl.make_block_ptr(g_cumsum + (bos * HQ + i_hq) * K, (T, K), (HQ*K, 1), (i_s, 0), (1, BK), (1, 0)) - p_q = tl.make_block_ptr(q + (bos * HQ + i_hq) * K, (T, K), (HQ*K, 1), (i_s, 0), (BS, BK), (1, 0)) - p_pq = tl.make_block_ptr(g_cumsum + (bos * HQ + i_hq) * K, (T, K), (HQ*K, 1), (i_s, 0), (BS, BK), (1, 0)) - p_do = tl.make_block_ptr(do + (bos * HQ + i_hq) * V, (T, V), (HQ*V, 1), (i_s, i_v * BV), (BS, BV), (1, 0)) - p_lse = tl.make_block_ptr(lse + bos * HQ + i_hq, (T,), (HQ,), (i_s,), (BS,), (0,)) - p_delta = tl.make_block_ptr(delta + bos * HQ + i_hq, (T,), (HQ,), (i_s,), (BS,), (0,)) - o_q = i_s + tl.arange(0, BS) m_q = o_q < T - b_Rq = tl.load(p_Rq, boundary_check=(0, 1)).to(tl.float32) + o_r = i_s + tl.arange(0, 1) + p_Rq = g_cumsum + (bos * HQ + i_hq) * K + o_r[:, None] * (HQ*K) + o_d[None, :] + p_q = q + (bos * HQ + i_hq) * K + o_q[:, None] * (HQ*K) + o_d[None, :] + p_pq = g_cumsum + (bos * HQ + i_hq) * K + o_q[:, None] * (HQ*K) + o_d[None, :] + p_do = do + (bos * HQ + i_hq) * V + o_q[:, None] * (HQ*V) + o_v[None, :] + p_lse = lse + bos * HQ + i_hq + o_q * HQ + p_delta = delta + bos * HQ + i_hq + o_q * HQ + + b_Rq = tl.load(p_Rq, mask=(o_r[:, None] < T) & (o_d[None, :] < K), other=0.0).to(tl.float32) b_Rq_bc_k = tl.broadcast_to(b_Rq, (BT, BK)) b_Rq_bc_q = tl.broadcast_to(b_Rq, (BS, BK)) - b_q = tl.load(p_q, boundary_check=(0, 1)) - b_pq = tl.load(p_pq, boundary_check=(0, 1)).to(tl.float32) + b_q = tl.load(p_q, mask=m_q[:, None] & (o_d[None, :] < K), other=0.0) + b_pq = tl.load(p_pq, mask=m_q[:, None] & (o_d[None, :] < K), other=0.0).to(tl.float32) b_q_til = (b_q.to(tl.float32) * exp2(b_pq - b_Rq_bc_q)).to(b_q.dtype) b_exp_k_off = b_Rq_bc_k - b_pk b_exp_k_off_val = exp2(b_exp_k_off) # CSE: reused for k_til and final dk scaling b_k_til = (b_k.to(tl.float32) * b_exp_k_off_val).to(b_k.dtype) - b_do = tl.load(p_do, boundary_check=(0, 1)) - b_lse = tl.load(p_lse, boundary_check=(0,)) - b_delta = tl.load(p_delta, boundary_check=(0,)) + b_do = tl.load(p_do, mask=m_q[:, None] & (o_v[None, :] < V), other=0.0) + b_lse = tl.load(p_lse, mask=m_q, other=0.0) + b_delta = tl.load(p_delta, mask=m_q, other=0.0) b_s = tl.dot(b_k_til, tl.trans(b_q_til)) * scale * RCP_LN2 if USE_SCALAR_G: b_cq = tl.load(g_scalar_cumsum + (bos + o_q) * HQ + i_hq, mask=m_q, other=0).to(tl.float32) @@ -552,7 +629,7 @@ def parallel_wall_attn_bwd_kernel_dkv( b_p = tl.where((o_q[None, :] - o_k[:, None] < W) & m_q[None, :], exp2(b_s - b_lse[None, :]), 0) else: b_p = tl.where(m_q[None, :], exp2(b_s - b_lse[None, :]), 0) - b_dv += tl.dot(b_p.to(b_do.dtype), b_do) + b_dv = tl.dot(b_p.to(b_do.dtype), b_do, b_dv) b_dp = tl.dot(b_v, tl.trans(b_do)) b_ds = b_p * (b_dp - b_delta[None, :]) b_dk_til = tl.dot(b_ds.to(b_q.dtype), b_q_til) @@ -562,12 +639,12 @@ def parallel_wall_attn_bwd_kernel_dkv( # Structural: b_dg derived from b_dk (saves [BT,BK] fp32 accumulator). b_dg = -LN2 * b_k.to(tl.float32) * b_dk - tl.store(p_dk, b_dk.to(p_dk.dtype.element_ty), boundary_check=(0, 1)) - tl.store(p_dv, b_dv.to(p_dv.dtype.element_ty), boundary_check=(0, 1)) - tl.store(p_dg, b_dg.to(p_dg.dtype.element_ty), boundary_check=(0, 1)) + tl.store(p_dk, b_dk.to(p_dk.dtype.element_ty), mask=m_kk) + tl.store(p_dv, b_dv.to(p_dv.dtype.element_ty), mask=m_kv) + tl.store(p_dg, b_dg.to(p_dg.dtype.element_ty), mask=m_kk) if USE_SCALAR_G: - p_dc = tl.make_block_ptr(dg_scalar_cumsum + bos * HQ + i_hq, (T,), (HQ,), (i_t * BT,), (BT,), (0,)) - tl.store(p_dc, b_dc.to(p_dc.dtype.element_ty), boundary_check=(0,)) + p_dc = dg_scalar_cumsum + bos * HQ + i_hq + o_k * HQ + tl.store(p_dc, b_dc.to(p_dc.dtype.element_ty), mask=m_k) @dispatch('attn') @@ -623,20 +700,22 @@ def parallel_wall_attn_fwd( NT = len(chunk_indices) o = torch.empty(B, T, HQ, V, dtype=v.dtype, device=q.device) lse = torch.empty(B, T, HQ, dtype=torch.float, device=q.device) - parallel_wall_attn_fwd_kernel.fn[(NV, NT, B * HQ)]( + parallel_wall_attn_fwd_kernel.fn[(NV * NT, B * HQ)]( q=q, k=k, v=v, o=o, g_cumsum=g_cumsum, g_scalar_cumsum=g_scalar_cumsum, sink_bias=sink_bias, lse=lse, scale=scale, cu_seqlens=cu_seqlens, chunk_indices=chunk_indices, - B=B, T=T, W=window_size, H=H, HQ=HQ, G=G, K=K, V=V, + B=B, T=T, T_BUCKET=1, W=window_size, H=H, HQ=HQ, G=G, K=K, V=V, BT=BT, BS=BS, BK=BK, BV=BV, num_warps=num_warps, num_stages=num_stages, ) return o, lse o = torch.empty(B, T, HQ, V, dtype=v.dtype, device=q.device) lse = torch.empty(B, T, HQ, dtype=torch.float, device=q.device) + # the bucket selects an autotune config; exact T remains the runtime loop bound + T_BUCKET = triton.next_power_of_2(T) def grid(meta): - return (NV, triton.cdiv(T, meta['BT']), B * HQ) + return (NV * triton.cdiv(T, meta['BT']), B * HQ) parallel_wall_attn_fwd_kernel[grid]( q=q, k=k, @@ -651,6 +730,7 @@ def parallel_wall_attn_fwd( chunk_indices=chunk_indices, B=B, T=T, + T_BUCKET=T_BUCKET, W=window_size, H=H, HQ=HQ, @@ -663,6 +743,26 @@ def parallel_wall_attn_fwd( return o, lse +def parallel_wall_attn_bwd_preprocess( + o: torch.Tensor, + do: torch.Tensor, + BV: int, +) -> torch.Tensor: + V = o.shape[-1] + NV = triton.cdiv(V, BV) + N = o.numel() // V + delta = torch.empty(NV, *o.shape[:-1], dtype=torch.float, device=o.device) + parallel_wall_attn_bwd_kernel_preprocess[(NV * N,)]( + o=o, + do=do, + delta=delta, + N=N, + V=V, + BV=BV, + ) + return delta + + @dispatch('attn') def parallel_wall_attn_bwd( q: torch.Tensor, @@ -686,51 +786,50 @@ def parallel_wall_attn_bwd( B, T, H, K, V = *k.shape, v.shape[-1] HQ = q.shape[2] G = HQ // H - # dq/dk are reduced over the full value dim in one program (no cross-program accumulation), - # so BV must span all of V (NV == 1). Don't cap it here -- the forward can, the backward can't. BK = max(triton.next_power_of_2(K), 16) - BV = max(triton.next_power_of_2(V), 16) + if check_shared_mem('hopper', q.device.index) or check_shared_mem('ampere', q.device.index): + BV = min(128, max(triton.next_power_of_2(V), 16)) + else: + BV = min(64, max(triton.next_power_of_2(V), 16)) NV = triton.cdiv(V, BV) # Varlen: pin BT=128 and bypass the autotuner (BT sweep would invalidate the # precomputed chunk map). `extra` injects the fixed launch config via `.fn`. is_varlen = cu_seqlens is not None + # varlen bypasses autotuning; dense lengths share configs within power-of-two buckets + T_BUCKET = 1 if is_varlen else triton.next_power_of_2(T) extra = {} if is_varlen: BT = 128 - if check_shared_mem('hopper'): + if check_shared_mem('hopper', q.device.index): BS, num_warps = min(64, max(16, triton.next_power_of_2(T))), 8 - elif check_shared_mem('ampere'): + elif check_shared_mem('ampere', q.device.index): BS, num_warps = min(32, max(16, triton.next_power_of_2(T))), 4 else: BS, num_warps = min(32, max(16, triton.next_power_of_2(T))), 2 if chunk_indices is None: chunk_indices = prepare_chunk_indices(cu_seqlens, BT) NT = len(chunk_indices) - grid = (NV, NT, B * HQ) + grid = (NV * NT, B * HQ) extra = dict(BT=BT, BS=BS, num_warps=num_warps, num_stages=2) else: def grid(meta): - return (NV, triton.cdiv(T, meta['BT']), B * HQ) + return (NV * triton.cdiv(T, meta['BT']), B * HQ) - delta = parallel_attn_bwd_preprocess(o, do) + delta = parallel_wall_attn_bwd_preprocess(o, do, BV) - dq = torch.empty(B, T, HQ, K, dtype=k.dtype if H == HQ else torch.float, device=q.device) - dk = torch.empty(B, T, HQ, K, dtype=k.dtype if H == HQ else torch.float, device=q.device) - dv = torch.empty(B, T, HQ, V, dtype=v.dtype if H == HQ else torch.float, device=q.device) + partial_dtype = k.dtype if NV == 1 and H == HQ else torch.float + dq = torch.empty(NV, B, T, HQ, K, dtype=partial_dtype, device=q.device) dq_kernel = parallel_wall_attn_bwd_kernel_dq.fn if is_varlen else parallel_wall_attn_bwd_kernel_dq dkv_kernel = parallel_wall_attn_bwd_kernel_dkv.fn if is_varlen else parallel_wall_attn_bwd_kernel_dkv - dg_cumsum = torch.empty(B, T, HQ, K, dtype=torch.float, device=q.device) - dg_cumsum_k = torch.empty_like(dg_cumsum) + dg_cumsum = torch.empty(NV, B, T, HQ, K, dtype=torch.float, device=q.device) if g_scalar_cumsum is not None: - dg_scalar_cumsum = torch.empty(B, T, HQ, dtype=torch.float, device=q.device) - dg_scalar_cumsum_k = torch.empty_like(dg_scalar_cumsum) + dg_scalar_cumsum = torch.empty(NV, B, T, HQ, dtype=torch.float, device=q.device) else: dg_scalar_cumsum = None - dg_scalar_cumsum_k = None dq_kernel[grid]( q=q, @@ -748,6 +847,7 @@ def parallel_wall_attn_bwd( chunk_indices=chunk_indices, scale=scale, T=T, + T_BUCKET=T_BUCKET, W=window_size, B=B, H=H, @@ -759,6 +859,26 @@ def parallel_wall_attn_bwd( BV=BV, **extra, ) + # reduce query-side partials before allocating key-side partials to bound peak temporary memory + if NV > 1: + dq = dq.sum(0) + dg_cumsum = dg_cumsum.sum(0) + if g_scalar_cumsum is not None: + dg_scalar_cumsum = dg_scalar_cumsum.sum(0) + else: + dq = dq[0] + dg_cumsum = dg_cumsum[0] + if g_scalar_cumsum is not None: + dg_scalar_cumsum = dg_scalar_cumsum[0] + + dk = torch.empty(NV, B, T, HQ, K, dtype=partial_dtype, device=q.device) + dv = torch.empty(B, T, HQ, V, dtype=v.dtype if H == HQ else torch.float, device=q.device) + dg_cumsum_k = torch.empty(NV, B, T, HQ, K, dtype=torch.float, device=q.device) + if g_scalar_cumsum is not None: + dg_scalar_cumsum_k = torch.empty(NV, B, T, HQ, dtype=torch.float, device=q.device) + else: + dg_scalar_cumsum_k = None + dkv_kernel[grid]( q=q, k=k, @@ -776,6 +896,7 @@ def parallel_wall_attn_bwd( chunk_indices=chunk_indices, scale=scale, T=T, + T_BUCKET=T_BUCKET, W=window_size, B=B, H=H, @@ -788,6 +909,17 @@ def parallel_wall_attn_bwd( DIAG_BF16=_DKV_DIAG_BF16, **extra, ) + if NV > 1: + dk = dk.sum(0) + dg_cumsum_k = dg_cumsum_k.sum(0) + if g_scalar_cumsum is not None: + dg_scalar_cumsum_k = dg_scalar_cumsum_k.sum(0) + else: + dk = dk[0] + dg_cumsum_k = dg_cumsum_k[0] + if g_scalar_cumsum is not None: + dg_scalar_cumsum_k = dg_scalar_cumsum_k[0] + dk = reduce(dk, 'b t (h g) k -> b t h k', g=G, reduction='sum') dv = reduce(dv, 'b t (h g) v -> b t h v', g=G, reduction='sum') dg_cumsum.add_(dg_cumsum_k) @@ -798,7 +930,8 @@ def parallel_wall_attn_bwd( dsink_bias = None if sink_bias is not None: p_sink_bias = torch.exp2(sink_bias[None, None, :] - lse) - dsink_bias = -(p_sink_bias * delta).sum((0, 1)) + delta_total = delta.sum(0) if NV > 1 else delta[0] + dsink_bias = -(p_sink_bias * delta_total).sum((0, 1)) return dq, dk, dv, dg_cumsum, dsink_bias, dg_scalar_cumsum diff --git a/build/torch-cuda/symbols.json b/build/torch-cuda/symbols.json index 1ac407c9b755eed358a808263538d9cbf1b92034..257e5b98dc6869e4d01a542c52b0a9b5478305dc 100644 --- a/build/torch-cuda/symbols.json +++ b/build/torch-cuda/symbols.json @@ -1,13 +1,13 @@ { "schema_version": 1, - "module": "_fla_cuda_fe2bb66", + "module": "_fla_cuda_faa8674", "functions": [ { "name": "chunk_kda", "qualname": "chunk_kda", - "module": "_fla_cuda_fe2bb66.ops.kda.chunk", + "module": "_fla_cuda_faa8674.ops.kda.chunk", "kind": "function", - "doc": "Args:\n q (torch.Tensor):\n queries of shape ``[B, T, H, K]``.\n k (torch.Tensor):\n keys of shape ``[B, T, H, K]``.\n v (torch.Tensor):\n values of shape ``[B, T, HV, V]``.\n GVA (Grouped Value Attention) is applied if ``HV > H``, where ``HV`` must be divisible by ``H``.\n g (torch.Tensor):\n (forget) gating tensor (in log space!) of shape ``[B, T, HV, K]``.\n When ``use_gate_in_kernel=False`` (default), ``g`` should be the pre-computed decay value.\n When ``use_gate_in_kernel=True``, ``g`` is the raw input before gate activation;\n the kernel fuses ``-exp(A_log) * softplus(g + dt_bias)`` + chunk cumsum internally.\n beta (torch.Tensor):\n betas of shape ``[B, T, HV]``.\n scale (Optional[float]):\n Scale factor for the KDA attention scores.\n If not provided, it will default to ``1 / sqrt(K)``. Default: ``None``.\n initial_state (Optional[torch.Tensor]):\n Initial state of shape ``[N, HV, K, V]`` for ``N`` input sequences.\n For equal-length input sequences, ``N`` equals the batch size ``B``.\n Default: ``None``.\n output_final_state (Optional[bool]):\n Whether to output the final state of shape ``[N, HV, K, V]``. Default: ``False``.\n use_qk_l2norm_in_kernel (bool):\n Whether to apply L2norm to the q,k tensor internally. Default: ``False``.\n use_gate_in_kernel (bool):\n Whether to compute the log-space KDA decay internally.\n - If ``True``:\n The passed ``g`` acts as the raw input for ``-exp(A_log) * softplus(g + dt_bias.view(HV, K))``.\n Note that as part of the input arguments,\n ``A_log`` (shape ``[HV]``) and the optional ``dt_bias`` (shape ``[HV * K]``) should be provided.\n - If ``False``, ``g`` is expected to be the pre-computed decay value.\n Default: ``False``.\n use_beta_sigmoid_in_kernel (bool):\n Whether to apply ``torch.sigmoid(beta)`` before launching the chunk kernel.\n - If ``True``, the passed ``beta`` acts as the raw beta logits.\n - If ``False``, ``beta`` is expected to already be in post-sigmoid space.\n Default: ``False``.\n allow_neg_eigval (bool):\n Whether to allow negative eigenvalues by scaling ``beta`` to ``[0, 2)``.\n Only takes effect together with ``use_beta_sigmoid_in_kernel=True``, in which case\n the kernel computes ``2 * sigmoid(beta)`` instead of ``sigmoid(beta)``.\n Default: ``False``.\n safe_gate (bool):\n Whether to clamp the gate to ``[lower_bound, 0)`` and enable M=16 TensorCore\n acceleration for higher throughput. Requires ``lower_bound`` to be set.\n Default: ``False``.\n lower_bound (Optional[float]):\n Lower bound for the forget gate (in log space). When set together with\n ``safe_gate=True``, changes the gate activation from\n ``-exp(A_log) * softplus(g + dt_bias)`` to\n ``lower_bound * sigmoid(exp(A_log) * (g + dt_bias))``,\n which naturally clamps the output to ``[lower_bound, 0)``.\n Recommended value: ``-5`` (i.e., per-step decay ``exp(-5) ≈ 0.0067``).\n Default: ``None``.\n disable_recompute (bool):\n Whether to disable gradient recomputation in the kernel. When ``True``, the kernel\n will save all intermediate activations for backward pass, which is beneficial\n for training small models at the cost of increased memory usage. Default: ``False``.\n return_intermediate_states (bool):\n If True, returns intermediate state ``h`` for inference scenarios (e.g., vLLM).\n Must be used within ``torch.inference_mode()`` and will return a 3-tuple instead of 2-tuple.\n This is not intended for training as it bypasses autograd. Default: ``False``.\n state_v_first (Optional[bool]):\n Store the recurrent state in V-first ``[V, K]`` layout instead of the default ``[K, V]``. Default: ``False``.\n cu_seqlens (torch.LongTensor):\n Cumulative sequence lengths of shape ``[N+1]`` used for variable-length training,\n consistent with the FlashAttention API.\n cu_seqlens_cpu (torch.LongTensor):\n Cumulative sequence lengths of shape ``[N+1]`` used for variable-length training,\n consistent with the FlashAttention API.\n cp_context (Optional[FLACPContext]):\n Context parallel context for distributed training across multiple devices.\n When provided, ``initial_state`` and ``output_final_state`` are not supported,\n and ``cu_seqlens`` will be overridden by the context. Default: ``None``.\n\nReturns:\n - Normal mode (return_intermediate_states=False): A tuple (o, final_state)\n o (torch.Tensor):\n Outputs of shape ``[B, T, HV, V]``.\n final_state (torch.Tensor):\n Final state of shape ``[N, HV, K, V]`` if ``output_final_state=True`` else ``None``.\n - Inference mode (return_intermediate_states=True): A tuple (o, final_state, h)\n o (torch.Tensor):\n Outputs of shape ``[B, T, HV, V]``.\n final_state (torch.Tensor):\n Final state of shape ``[N, HV, K, V]`` if ``output_final_state=True`` else ``None``.\n h (torch.Tensor):\n Intermediate states of shape ``[B, NT, HV, K, V]`` and dtype ``bfloat16``.\n - For equal-length sequences: ``NT = ceil(T / chunk_size)``\n - For variable-length sequences (cu_seqlens): B is always 1 (flattened),\n NT is the total number of chunks across all sequences.\n\nExamples::\n >>> import torch\n >>> import torch.nn.functional as F\n >>> from einops import rearrange\n >>> from ...ops.kda import chunk_kda\n # inputs with equal lengths (no GVA, HV == H)\n >>> B, T, H, K, V = 4, 2048, 4, 512, 512\n >>> q = torch.randn(B, T, H, K, dtype=torch.bfloat16, device='cuda')\n >>> k = torch.randn(B, T, H, K, dtype=torch.bfloat16, device='cuda')\n >>> v = torch.randn(B, T, H, V, dtype=torch.bfloat16, device='cuda')\n >>> beta = torch.rand(B, T, H, dtype=torch.bfloat16, device='cuda')\n >>> g = torch.rand(B, T, H, K, dtype=torch.bfloat16, device='cuda')\n >>> h0 = torch.randn(B, H, K, V, dtype=torch.bfloat16, device='cuda')\n >>> A_log = torch.randn(H, dtype=torch.float32, device='cuda')\n >>> dt_bias = torch.randn(H * K, dtype=torch.float32, device='cuda')\n >>> o, ht = chunk_kda(\n q, k, v, g, beta,\n A_log=A_log,\n dt_bias=dt_bias,\n use_qk_l2norm_in_kernel=True,\n use_gate_in_kernel=True,\n initial_state=h0,\n output_final_state=True\n )\n # GVA mode (HV > H)\n >>> HV = 8 # 2x more value heads than qk heads\n >>> v = torch.randn(B, T, HV, V, dtype=torch.bfloat16, device='cuda')\n >>> g = torch.rand(B, T, HV, K, dtype=torch.bfloat16, device='cuda')\n >>> beta = torch.rand(B, T, HV, dtype=torch.bfloat16, device='cuda')\n >>> h0 = torch.randn(B, HV, K, V, dtype=torch.bfloat16, device='cuda')\n >>> A_log = torch.randn(HV, dtype=torch.float32, device='cuda')\n >>> dt_bias = torch.randn(HV * K, dtype=torch.float32, device='cuda')\n >>> o, ht = chunk_kda(\n q, k, v, g, beta,\n A_log=A_log,\n dt_bias=dt_bias,\n use_qk_l2norm_in_kernel=True,\n use_gate_in_kernel=True,\n initial_state=h0,\n output_final_state=True\n )\n # for variable-length inputs, the batch size `B` is expected to be 1 and `cu_seqlens` is required\n >>> q, k, v, beta, g = map(lambda x: rearrange(x, 'b t ... -> 1 (b t) ...'), (q, k, v, beta, g))\n # for a batch with 4 sequences, `cu_seqlens` with 5 start/end positions are expected\n >>> cu_seqlens = q.new_tensor([0, 2048, 4096, 6144, 8192], dtype=torch.long)\n >>> o, ht = chunk_kda(\n q, k, v, g, beta,\n A_log=A_log,\n dt_bias=dt_bias,\n use_qk_l2norm_in_kernel=True,\n use_gate_in_kernel=True,\n initial_state=h0,\n output_final_state=True,\n cu_seqlens=cu_seqlens\n )", + "doc": "Args:\n q (torch.Tensor):\n queries of shape ``[B, T, H, K]``.\n k (torch.Tensor):\n keys of shape ``[B, T, H, K]``.\n v (torch.Tensor):\n values of shape ``[B, T, HV, V]``.\n GVA (Grouped Value Attention) is applied if ``HV > H``, where ``HV`` must be divisible by ``H``.\n g (torch.Tensor):\n (forget) gating tensor (in log space!) of shape ``[B, T, HV, K]``.\n When ``use_gate_in_kernel=False`` (default), ``g`` should be the pre-computed decay value.\n When ``use_gate_in_kernel=True``, ``g`` is the raw input before gate activation;\n the kernel fuses ``-exp(A_log) * softplus(g + dt_bias)`` + chunk cumsum internally.\n beta (torch.Tensor):\n betas of shape ``[B, T, HV]``.\n scale (float, Optional):\n Scale factor for the KDA attention scores.\n If not provided, it will default to ``1 / sqrt(K)``. Default: ``None``.\n initial_state (torch.Tensor, Optional):\n Initial state of shape ``[N, HV, K, V]`` for ``N`` input sequences.\n For equal-length input sequences, ``N`` equals the batch size ``B``.\n Default: ``None``.\n output_final_state (bool, Optional):\n Whether to output the final state of shape ``[N, HV, K, V]``. Default: ``False``.\n use_qk_l2norm_in_kernel (bool, Optional):\n Whether to apply L2norm to the q,k tensor internally. Default: ``False``.\n use_gate_in_kernel (bool, Optional):\n Whether to compute the log-space KDA decay internally.\n - If ``True``:\n The passed ``g`` acts as the raw input for ``-exp(A_log) * softplus(g + dt_bias.view(HV, K))``.\n Note that as part of the input arguments,\n ``A_log`` (shape ``[HV]``) and the optional ``dt_bias`` (shape ``[HV * K]``) should be provided.\n When ``lower_bound`` is set, ``A_log`` may be ``None``,\n in which case the gate is ``lower_bound * sigmoid(g + dt_bias)``.\n - If ``False``, ``g`` is expected to be the pre-computed decay value.\n Default: ``False``.\n use_beta_sigmoid_in_kernel (bool, Optional):\n Whether to apply ``torch.sigmoid(beta)`` before launching the chunk kernel.\n - If ``True``, the passed ``beta`` acts as the raw beta logits.\n - If ``False``, ``beta`` is expected to already be in post-sigmoid space.\n Default: ``False``.\n allow_neg_eigval (bool, Optional):\n Whether to allow negative eigenvalues by scaling ``beta`` to ``[0, 2)``.\n Only takes effect together with ``use_beta_sigmoid_in_kernel=True``,\n in which case the kernel computes ``2 * sigmoid(beta)`` instead of ``sigmoid(beta)``.\n Default: ``False``.\n safe_gate (bool, Optional):\n Whether to enable M=16 TensorCore acceleration for bounded gates.\n With ``use_gate_in_kernel=True``, requires ``lower_bound`` and bounds gates to ``[lower_bound, 0)``.\n Otherwise, the caller must provide pre-computed bounded gates.\n Default: ``False``.\n lower_bound (float, Optional):\n Lower bound for the forget gate (in log space). When set together with ``safe_gate=True``,\n the gate uses ``lower_bound * sigmoid(exp(A_log) * (g + dt_bias))``,\n replacing ``-exp(A_log) * softplus(g + dt_bias)`` and clamping the output to ``[lower_bound, 0)``.\n Recommended value: ``-5`` (i.e., per-step decay ``exp(-5) ≈ 0.0067``).\n Default: ``None``.\n disable_recompute (bool, Optional):\n Whether to disable gradient recomputation in the kernel. When ``True``,\n the kernel will save all intermediate activations for backward pass,\n which is beneficial for training small models at the cost of increased memory usage. Default: ``False``.\n return_intermediate_states (bool, Optional):\n If ``True``, returns intermediate state ``h`` for inference scenarios (e.g., vLLM).\n Must be used within ``torch.inference_mode()`` and will return a 3-tuple instead of 2-tuple.\n This is not intended for training as it bypasses autograd. Default: ``False``.\n state_v_first (bool, Optional):\n Store the recurrent state in V-first ``[V, K]`` layout instead of the default ``[K, V]``. Default: ``False``.\n cu_seqlens (torch.LongTensor, Optional):\n Cumulative sequence lengths of shape ``[N+1]`` used for variable-length training,\n consistent with the FlashAttention API. Default: `None`.\n cu_seqlens_cpu (torch.LongTensor, Optional):\n Cumulative sequence lengths of shape ``[N+1]`` used for variable-length training,\n consistent with the FlashAttention API. Default: `None`.\n cp_context (FLACPContext, Optional):\n Context parallel context for distributed training across multiple devices.\n When provided, ``initial_state`` and ``output_final_state`` are not supported,\n and ``cu_seqlens`` will be overridden by the context. Default: ``None``.\n use_graph (bool, Optional):\n Whether to build a fixed-shape varlen chunk list on-device for CUDA or NPU graph capture.\n Forward and backward can then be replayed after only the contents of ``cu_seqlens`` change. Requires:\n\n - Pad ``cu_seqlens`` with zero-length tail sequences to exactly ``max_num_seqs + 1`` entries;\n kernels return immediately on sentinel rows (segment id < 0),\n so output/gradient rows not covered by real tokens are left undefined.\n - Shape ``initial_state`` and the incoming ``dht`` gradient (when used) as ``[max_num_seqs, ...]``.\n - Warm up the captured step eagerly so kernel autotuning runs outside the graph;\n with ``cp_context``, this also includes NCCL communicator setup.\n\n With ``cp_context``, the cross-rank state all-gathers are recorded into the graph. This additionally requires:\n\n - A graph-capturable (non-blocking) NCCL communicator;\n all ranks must capture and replay the same collectives in the same order.\n - A persistent padded ``cp_context.cu_seqlens`` buffer, refreshed in place before each replay.\n - Persistent int32 device scalars ``cp_context.pre_num_ranks_dev`` and ``post_num_ranks_dev``,\n refreshed in place before each replay.\n\n Ascend currently requires flattened variable-length inputs and does not support context parallelism in graph mode.\n ``disable_recompute=True`` is not supported. Default: ``False``.\n max_num_seqs (int, Optional):\n Static upper bound of the sequence count used together with ``use_graph``.\n Uses ``cu_seqlens.shape[0] - 1`` when omitted. Default: ``None``.\n\nReturns:\n - Normal mode (return_intermediate_states=False): A tuple (o, final_state)\n o (torch.Tensor):\n Outputs of shape ``[B, T, HV, V]``.\n final_state (torch.Tensor):\n Final state of shape ``[N, HV, K, V]`` if ``output_final_state=True`` else ``None``.\n - Inference mode (return_intermediate_states=True): A tuple (o, final_state, h)\n o (torch.Tensor):\n Outputs of shape ``[B, T, HV, V]``.\n final_state (torch.Tensor):\n Final state of shape ``[N, HV, K, V]`` if ``output_final_state=True`` else ``None``.\n h (torch.Tensor):\n Intermediate states of shape ``[B, NT, HV, K, V]`` and dtype ``bfloat16``.\n - For equal-length sequences: ``NT = ceil(T / chunk_size)``\n - For variable-length sequences (cu_seqlens): B is always 1 (flattened),\n NT is the total number of chunks across all sequences.\n\nExamples::\n >>> import torch\n >>> import torch.nn.functional as F\n >>> from einops import rearrange\n >>> from fla.ops.kda import chunk_kda\n # inputs with equal lengths (no GVA, HV == H)\n >>> B, T, H, K, V = 4, 2048, 4, 512, 512\n >>> q = torch.randn(B, T, H, K, dtype=torch.bfloat16, device='cuda')\n >>> k = torch.randn(B, T, H, K, dtype=torch.bfloat16, device='cuda')\n >>> v = torch.randn(B, T, H, V, dtype=torch.bfloat16, device='cuda')\n >>> beta = torch.rand(B, T, H, dtype=torch.bfloat16, device='cuda')\n >>> g = torch.rand(B, T, H, K, dtype=torch.bfloat16, device='cuda')\n >>> h0 = torch.randn(B, H, K, V, dtype=torch.bfloat16, device='cuda')\n >>> A_log = torch.randn(H, dtype=torch.float32, device='cuda')\n >>> dt_bias = torch.randn(H * K, dtype=torch.float32, device='cuda')\n >>> o, ht = chunk_kda(\n ... q=q,\n ... k=k,\n ... v=v,\n ... g=g,\n ... beta=beta,\n ... initial_state=h0,\n ... output_final_state=True,\n ... use_qk_l2norm_in_kernel=True,\n ... use_gate_in_kernel=True,\n ... A_log=A_log,\n ... dt_bias=dt_bias,\n ... )\n # GVA mode (HV > H)\n # 2x more value heads than qk heads\n >>> HV = 8\n >>> v = torch.randn(B, T, HV, V, dtype=torch.bfloat16, device='cuda')\n >>> g = torch.rand(B, T, HV, K, dtype=torch.bfloat16, device='cuda')\n >>> beta = torch.rand(B, T, HV, dtype=torch.bfloat16, device='cuda')\n >>> h0 = torch.randn(B, HV, K, V, dtype=torch.bfloat16, device='cuda')\n >>> A_log = torch.randn(HV, dtype=torch.float32, device='cuda')\n >>> dt_bias = torch.randn(HV * K, dtype=torch.float32, device='cuda')\n >>> o, ht = chunk_kda(\n ... q=q,\n ... k=k,\n ... v=v,\n ... g=g,\n ... beta=beta,\n ... initial_state=h0,\n ... output_final_state=True,\n ... use_qk_l2norm_in_kernel=True,\n ... use_gate_in_kernel=True,\n ... A_log=A_log,\n ... dt_bias=dt_bias,\n ... )\n # for variable-length inputs, the batch size `B` is expected to be 1 and `cu_seqlens` is required\n >>> q, k, v, beta, g = map(lambda x: rearrange(x, 'b t ... -> 1 (b t) ...'), (q, k, v, beta, g))\n # for a batch with 4 sequences, `cu_seqlens` with 5 start/end positions are expected\n >>> cu_seqlens = q.new_tensor([0, 2048, 4096, 6144, 8192], dtype=torch.long)\n >>> o, ht = chunk_kda(\n ... q=q,\n ... k=k,\n ... v=v,\n ... g=g,\n ... beta=beta,\n ... initial_state=h0,\n ... output_final_state=True,\n ... use_qk_l2norm_in_kernel=True,\n ... use_gate_in_kernel=True,\n ... cu_seqlens=cu_seqlens,\n ... A_log=A_log,\n ... dt_bias=dt_bias,\n ... )", "signature": { "parameters": [ { @@ -128,7 +128,19 @@ "name": "cp_context", "kind": "POSITIONAL_OR_KEYWORD", "default": "None", - "annotation": "_fla_cuda_fe2bb66.ops.cp.context.FLACPContext" + "annotation": "_fla_cuda_faa8674.ops.cp.context.FLACPContext | None" + }, + { + "name": "use_graph", + "kind": "POSITIONAL_OR_KEYWORD", + "default": "False", + "annotation": "bool" + }, + { + "name": "max_num_seqs", + "kind": "POSITIONAL_OR_KEYWORD", + "default": "None", + "annotation": "int | None" }, { "name": "kwargs", @@ -143,9 +155,9 @@ { "name": "fused_recurrent_kda", "qualname": "fused_recurrent_kda", - "module": "_fla_cuda_fe2bb66.ops.kda.fused_recurrent", + "module": "_fla_cuda_faa8674.ops.kda.fused_recurrent", "kind": "function", - "doc": "Args:\n q (torch.Tensor):\n queries of shape `[B, T, H, K]`.\n k (torch.Tensor):\n keys of shape `[B, T, H, K]`.\n v (torch.Tensor):\n values of shape `[B, T, HV, V]`.\n GVA is applied if `HV > H`.\n g (torch.Tensor):\n g (decays) of shape `[B, T, HV, K]`.\n beta (torch.Tensor):\n betas of shape `[B, T, HV]`.\n A_log (Optional[torch.Tensor]):\n Decay parameter of shape `[HV]`. Required when `use_gate_in_kernel=True`.\n dt_bias (Optional[torch.Tensor]):\n Bias added to `g` before activation, of shape `[HV]`.\n Only used when `use_gate_in_kernel=True`.\n scale (Optional[float]):\n Scale factor for the RetNet attention scores.\n If not provided, it will default to `1 / sqrt(K)`. Default: `None`.\n initial_state (Optional[torch.Tensor]):\n Initial state of shape `[N, HV, K, V]` for `N` input sequences.\n For equal-length input sequences, `N` equals the batch size `B`.\n Default: `None`.\n output_final_state (Optional[bool]):\n Whether to output the final state of shape `[N, HV, K, V]`. Default: `False`.\n use_qk_l2norm_in_kernel (Optional[bool]):\n Whether to use L2 normalization in the kernel. Default: `False`.\n use_gate_in_kernel (Optional[bool]):\n Whether to compute the log-space KDA decay internally.\n When `True`, `g` is the raw input and `A_log` must be provided; the kernel fuses\n gate activation into the recurrence. Default: `False`.\n use_beta_sigmoid_in_kernel (Optional[bool]):\n Whether to apply `torch.sigmoid(beta)` inside the kernel.\n - If `True`, the passed `beta` acts as the raw beta logits.\n - If `False`, `beta` is expected to already be in post-sigmoid space.\n Default: `False`.\n allow_neg_eigval (Optional[bool]):\n Whether to allow negative eigenvalues by scaling `beta` to `[0, 2)`.\n Only takes effect together with `use_beta_sigmoid_in_kernel=True`, in which case\n the kernel computes `2 * sigmoid(beta)` instead of `sigmoid(beta)`. Default: `False`.\n lower_bound (Optional[float]):\n Lower bound for the forget gate (in log space). Only used when `use_gate_in_kernel=True`. Default: `None`.\n state_v_first (Optional[bool]):\n Store the recurrent state in V-first ``[V, K]`` layout instead of the default ``[K, V]``. Default: ``False``.\n cu_seqlens (torch.LongTensor):\n Cumulative sequence lengths of shape `[N+1]` used for variable-length training,\n consistent with the FlashAttention API.\n\nReturns:\n o (torch.Tensor):\n Outputs of shape `[B, T, HV, V]`.\n final_state (torch.Tensor):\n Final state of shape `[N, HV, K, V]` if `output_final_state=True` else `None`.\n\nExamples::\n >>> import torch\n >>> import torch.nn.functional as F\n >>> from einops import rearrange\n >>> from ...ops.kda import fused_recurrent_kda\n # inputs with equal lengths\n >>> B, T, H, HV, K, V = 4, 2048, 4, 8, 512, 512\n >>> q = torch.randn(B, T, H, K, device='cuda')\n >>> k = F.normalize(torch.randn(B, T, H, K, device='cuda'), p=2, dim=-1)\n >>> v = torch.randn(B, T, HV, V, device='cuda')\n >>> g = F.logsigmoid(torch.rand(B, T, HV, K, device='cuda'))\n >>> beta = torch.rand(B, T, HV, device='cuda').sigmoid()\n >>> h0 = torch.randn(B, HV, K, V, device='cuda')\n >>> o, ht = fused_recurrent_kda(\n q, k, v, g, beta,\n initial_state=h0,\n output_final_state=True\n )\n # for variable-length inputs, the batch size `B` is expected to be 1 and `cu_seqlens` is required\n >>> q, k, v, g, beta = map(lambda x: rearrange(x, 'b t ... -> 1 (b t) ...'), (q, k, v, g, beta))\n # for a batch with 4 sequences, `cu_seqlens` with 5 start/end positions are expected\n >>> cu_seqlens = q.new_tensor([0, 2048, 4096, 6144, 8192], dtype=torch.long)\n >>> o_var, ht_var = fused_recurrent_kda(\n q, k, v, g, beta,\n initial_state=h0,\n output_final_state=True,\n cu_seqlens=cu_seqlens\n )", + "doc": "Args:\n q (torch.Tensor):\n queries of shape `[B, T, H, K]`.\n k (torch.Tensor):\n keys of shape `[B, T, H, K]`.\n v (torch.Tensor):\n values of shape `[B, T, HV, V]`.\n GVA is applied if `HV > H`.\n g (torch.Tensor):\n g (decays) of shape `[B, T, HV, K]`.\n beta (torch.Tensor):\n betas of shape `[B, T, HV]`.\n A_log (Optional[torch.Tensor]):\n Decay parameter of shape `[HV]`.\n When `use_gate_in_kernel=True` together with `lower_bound`,\n may be `None` to use `lower_bound * sigmoid(g + dt_bias)`.\n dt_bias (Optional[torch.Tensor]):\n Bias added to `g` before activation, of shape `[HV]`. Only used when `use_gate_in_kernel=True`.\n scale (Optional[float]):\n Scale factor for the RetNet attention scores.\n If not provided, it will default to `1 / sqrt(K)`. Default: `None`.\n initial_state (Optional[torch.Tensor]):\n Initial state of shape `[N, HV, K, V]` for `N` input sequences.\n For equal-length input sequences, `N` equals the batch size `B`.\n Default: `None`.\n output_final_state (Optional[bool]):\n Whether to output the final state of shape `[N, HV, K, V]`. Default: `False`.\n use_qk_l2norm_in_kernel (Optional[bool]):\n Whether to use L2 normalization in the kernel. Default: `False`.\n use_gate_in_kernel (Optional[bool]):\n Whether to compute the log-space KDA decay internally.\n When `True`, `g` is the raw input and the kernel fuses gate activation into the recurrence.\n Default: `False`.\n use_beta_sigmoid_in_kernel (Optional[bool]):\n Whether to apply `torch.sigmoid(beta)` inside the kernel.\n - If `True`, the passed `beta` acts as the raw beta logits.\n - If `False`, `beta` is expected to already be in post-sigmoid space.\n Default: `False`.\n allow_neg_eigval (Optional[bool]):\n Whether to allow negative eigenvalues by scaling `beta` to `[0, 2)`.\n Only takes effect together with `use_beta_sigmoid_in_kernel=True`, in which case\n the kernel computes `2 * sigmoid(beta)` instead of `sigmoid(beta)`. Default: `False`.\n lower_bound (Optional[float]):\n Lower bound for the forget gate (in log space). Only used when `use_gate_in_kernel=True`. Default: `None`.\n state_v_first (Optional[bool]):\n Store the recurrent state in V-first ``[V, K]`` layout instead of the default ``[K, V]``. Default: ``False``.\n cu_seqlens (torch.LongTensor):\n Cumulative sequence lengths of shape `[N+1]` used for variable-length training,\n consistent with the FlashAttention API.\n\nReturns:\n o (torch.Tensor):\n Outputs of shape `[B, T, HV, V]`.\n final_state (torch.Tensor):\n Final state of shape `[N, HV, K, V]` if `output_final_state=True` else `None`.\n\nExamples::\n >>> import torch\n >>> import torch.nn.functional as F\n >>> from einops import rearrange\n >>> from fla.ops.kda import fused_recurrent_kda\n # inputs with equal lengths\n >>> B, T, H, HV, K, V = 4, 2048, 4, 8, 512, 512\n >>> q = torch.randn(B, T, H, K, device='cuda')\n >>> k = F.normalize(torch.randn(B, T, H, K, device='cuda'), p=2, dim=-1)\n >>> v = torch.randn(B, T, HV, V, device='cuda')\n >>> g = F.logsigmoid(torch.rand(B, T, HV, K, device='cuda'))\n >>> beta = torch.rand(B, T, HV, device='cuda').sigmoid()\n >>> h0 = torch.randn(B, HV, K, V, device='cuda')\n >>> o, ht = fused_recurrent_kda(\n q, k, v, g, beta,\n initial_state=h0,\n output_final_state=True\n )\n # for variable-length inputs, the batch size `B` is expected to be 1 and `cu_seqlens` is required\n >>> q, k, v, g, beta = map(lambda x: rearrange(x, 'b t ... -> 1 (b t) ...'), (q, k, v, g, beta))\n # for a batch with 4 sequences, `cu_seqlens` with 5 start/end positions are expected\n >>> cu_seqlens = q.new_tensor([0, 2048, 4096, 6144, 8192], dtype=torch.long)\n >>> o_var, ht_var = fused_recurrent_kda(\n q, k, v, g, beta,\n initial_state=h0,\n output_final_state=True,\n cu_seqlens=cu_seqlens\n )", "signature": { "parameters": [ { @@ -263,9 +275,9 @@ { "name": "chunk_gated_delta_rule", "qualname": "chunk_gated_delta_rule", - "module": "_fla_cuda_fe2bb66.ops.gated_delta_rule.chunk", + "module": "_fla_cuda_faa8674.ops.gated_delta_rule.chunk", "kind": "function", - "doc": "Args:\n q (torch.Tensor):\n queries of shape `[B, T, H, K]`.\n k (torch.Tensor):\n keys of shape `[B, T, H, K]`.\n v (torch.Tensor):\n values of shape `[B, T, HV, V]`.\n GVA (Grouped Value Attention) is applied if `HV > H`, where `HV` must be divisible by `H`.\n g (torch.Tensor):\n (forget) gating tensor of shape `[B, T, HV]`.\n When `use_gate_in_kernel=False` (default), `g` should be in log space (pre-computed decay).\n When `use_gate_in_kernel=True`, `g` is the raw input before gate activation;\n the kernel fuses `-exp(A_log) * softplus(g + dt_bias)` + chunk cumsum internally.\n beta (torch.Tensor):\n betas of shape `[B, T, HV]`.\n scale (Optional[float]):\n Scale factor for the RetNet attention scores.\n If not provided, it will default to `1 / sqrt(K)`. Default: `None`.\n initial_state (Optional[torch.Tensor]):\n Initial state of shape `[N, HV, K, V]` for `N` input sequences.\n For equal-length input sequences, `N` equals the batch size `B`.\n Default: `None`.\n output_final_state (Optional[bool]):\n Whether to output the final state of shape `[N, HV, K, V]`. Default: `False`.\n use_qk_l2norm_in_kernel (bool):\n Whether to apply L2norm to the q/k tensor internally. Default: `False`.\n use_gate_in_kernel (bool):\n Whether to compute the log-space GDN decay internally.\n When `True`, the passed `g` is the raw input, and `A_log` must be provided.\n The kernel fuses gate activation + chunk cumsum in a single pass.\n Default: `False`.\n A_log (Optional[torch.Tensor]):\n Decay parameter of shape `[HV]`. Required when `use_gate_in_kernel=True`.\n dt_bias (Optional[torch.Tensor]):\n Bias added to `g` before activation, of shape `[HV]`.\n Only used when `use_gate_in_kernel=True`.\n use_beta_sigmoid_in_kernel (bool):\n Whether to apply `torch.sigmoid(beta)` before launching the chunk kernel.\n - If `True`, the passed `beta` acts as the raw beta logits.\n - If `False`, `beta` is expected to already be in post-sigmoid space.\n Default: `False`.\n allow_neg_eigval (bool):\n Whether to allow negative eigenvalues by scaling `beta` to `[0, 2)`.\n Only takes effect together with `use_beta_sigmoid_in_kernel=True`, in which case\n the kernel computes `2 * sigmoid(beta)` instead of `sigmoid(beta)`. Default: `False`.\n state_v_first (Optional[bool]):\n Store the recurrent state in V-first ``[V, K]`` layout instead of the default ``[K, V]``. Default: ``False``.\n cu_seqlens (torch.LongTensor):\n Cumulative sequence lengths of shape `[N+1]` used for variable-length training,\n consistent with the FlashAttention API.\n cp_context (Optional[FLACPContext]):\n Context parallel context for distributed training across multiple devices.\n When provided, `initial_state` and `output_final_state` are not supported,\n and `cu_seqlens` will be overridden by the context. Default: `None`.\n\nReturns:\n o (torch.Tensor):\n Outputs of shape `[B, T, HV, V]`.\n final_state (torch.Tensor):\n Final state of shape `[N, HV, K, V]` if `output_final_state=True` else `None`.\n\nExamples::\n >>> import torch\n >>> import torch.nn.functional as F\n >>> from einops import rearrange\n >>> from ...ops.gated_delta_rule import chunk_gated_delta_rule\n # inputs with equal lengths\n >>> B, T, H, HV, K, V = 4, 2048, 4, 8, 512, 512\n >>> q = torch.randn(B, T, H, K, dtype=torch.bfloat16, device='cuda')\n >>> k = F.normalize(torch.randn(B, T, H, K, dtype=torch.bfloat16, device='cuda'), p=2, dim=-1)\n >>> v = torch.randn(B, T, HV, V, dtype=torch.bfloat16, device='cuda')\n >>> beta = torch.rand(B, T, HV, dtype=torch.bfloat16, device='cuda').sigmoid()\n >>> g = F.logsigmoid(torch.rand(B, T, HV, dtype=torch.bfloat16, device='cuda'))\n >>> h0 = torch.randn(B, HV, K, V, dtype=torch.bfloat16, device='cuda')\n >>> o, ht = chunk_gated_delta_rule(\n q, k, v, g, beta,\n initial_state=h0,\n output_final_state=True\n )\n # for variable-length inputs, the batch size `B` is expected to be 1 and `cu_seqlens` is required\n >>> q, k, v, beta, g = map(lambda x: rearrange(x, 'b t ... -> 1 (b t) ...'), (q, k, v, beta, g))\n # for a batch with 4 sequences, `cu_seqlens` with 5 start/end positions are expected\n >>> cu_seqlens = q.new_tensor([0, 2048, 4096, 6144, 8192], dtype=torch.long)\n >>> o, ht = chunk_gated_delta_rule(\n q, k, v, g, beta,\n initial_state=h0,\n output_final_state=True,\n cu_seqlens=cu_seqlens\n )", + "doc": "Args:\n q (torch.Tensor):\n queries of shape `[B, T, H, K]`.\n k (torch.Tensor):\n keys of shape `[B, T, H, K]`.\n v (torch.Tensor):\n values of shape `[B, T, HV, V]`.\n GVA (Grouped Value Attention) is applied if `HV > H`, where `HV` must be divisible by `H`.\n g (torch.Tensor):\n (forget) gating tensor of shape `[B, T, HV]`.\n When `use_gate_in_kernel=False` (default), `g` should be in log space (pre-computed decay).\n When `use_gate_in_kernel=True`, `g` is the raw input before gate activation;\n the kernel fuses `-exp(A_log) * softplus(g + dt_bias)` + chunk cumsum internally.\n beta (torch.Tensor):\n betas of shape `[B, T, HV]`.\n scale (Optional[float]):\n Scale factor for the RetNet attention scores.\n If not provided, it will default to `1 / sqrt(K)`. Default: `None`.\n initial_state (Optional[torch.Tensor]):\n Initial state of shape `[N, HV, K, V]` for `N` input sequences.\n For equal-length input sequences, `N` equals the batch size `B`.\n Default: `None`.\n output_final_state (Optional[bool]):\n Whether to output the final state of shape `[N, HV, K, V]`. Default: `False`.\n use_qk_l2norm_in_kernel (bool):\n Whether to apply L2norm to the q/k tensor internally. Default: `False`.\n use_gate_in_kernel (bool):\n Whether to compute the log-space GDN decay internally.\n When `True`, the passed `g` is the raw input, and `A_log` must be provided.\n The kernel fuses gate activation + chunk cumsum in a single pass.\n Default: `False`.\n A_log (Optional[torch.Tensor]):\n Decay parameter of shape `[HV]`. Required when `use_gate_in_kernel=True`.\n dt_bias (Optional[torch.Tensor]):\n Bias added to `g` before activation, of shape `[HV]`.\n Only used when `use_gate_in_kernel=True`.\n use_beta_sigmoid_in_kernel (bool):\n Whether to apply `torch.sigmoid(beta)` before launching the chunk kernel.\n - If `True`, the passed `beta` acts as the raw beta logits.\n - If `False`, `beta` is expected to already be in post-sigmoid space.\n Default: `False`.\n allow_neg_eigval (bool):\n Whether to allow negative eigenvalues by scaling `beta` to `[0, 2)`.\n Only takes effect together with `use_beta_sigmoid_in_kernel=True`, in which case\n the kernel computes `2 * sigmoid(beta)` instead of `sigmoid(beta)`. Default: `False`.\n state_v_first (Optional[bool]):\n Store the recurrent state in V-first ``[V, K]`` layout instead of the default ``[K, V]``. Default: ``False``.\n cu_seqlens (torch.LongTensor):\n Cumulative sequence lengths of shape `[N+1]` used for variable-length training,\n consistent with the FlashAttention API.\n chunk_indices (Optional[torch.LongTensor]):\n Pre-computed chunk indices for variable-length inputs.\n If provided, they are used directly instead of being computed from `cu_seqlens`. Default: `None`.\n cp_context (Optional[FLACPContext]):\n Context parallel context for distributed training across multiple devices.\n When provided, `initial_state` and `output_final_state` are not supported,\n and `cu_seqlens` will be overridden by the context. Default: `None`.\n\nReturns:\n o (torch.Tensor):\n Outputs of shape `[B, T, HV, V]`.\n final_state (torch.Tensor):\n Final state of shape `[N, HV, K, V]` if `output_final_state=True` else `None`.\n\nExamples::\n >>> import torch\n >>> import torch.nn.functional as F\n >>> from einops import rearrange\n >>> from fla.ops.gated_delta_rule import chunk_gated_delta_rule\n # inputs with equal lengths\n >>> B, T, H, HV, K, V = 4, 2048, 4, 8, 512, 512\n >>> q = torch.randn(B, T, H, K, dtype=torch.bfloat16, device='cuda')\n >>> k = F.normalize(torch.randn(B, T, H, K, dtype=torch.bfloat16, device='cuda'), p=2, dim=-1)\n >>> v = torch.randn(B, T, HV, V, dtype=torch.bfloat16, device='cuda')\n >>> beta = torch.rand(B, T, HV, dtype=torch.bfloat16, device='cuda').sigmoid()\n >>> g = F.logsigmoid(torch.rand(B, T, HV, dtype=torch.bfloat16, device='cuda'))\n >>> h0 = torch.randn(B, HV, K, V, dtype=torch.bfloat16, device='cuda')\n >>> o, ht = chunk_gated_delta_rule(\n q, k, v, g, beta,\n initial_state=h0,\n output_final_state=True\n )\n # for variable-length inputs, the batch size `B` is expected to be 1 and `cu_seqlens` is required\n >>> q, k, v, beta, g = map(lambda x: rearrange(x, 'b t ... -> 1 (b t) ...'), (q, k, v, beta, g))\n # for a batch with 4 sequences, `cu_seqlens` with 5 start/end positions are expected\n >>> cu_seqlens = q.new_tensor([0, 2048, 4096, 6144, 8192], dtype=torch.long)\n >>> o, ht = chunk_gated_delta_rule(\n q, k, v, g, beta,\n initial_state=h0,\n output_final_state=True,\n cu_seqlens=cu_seqlens\n )", "signature": { "parameters": [ { @@ -352,11 +364,17 @@ "default": "None", "annotation": "torch.LongTensor | None" }, + { + "name": "chunk_indices", + "kind": "POSITIONAL_OR_KEYWORD", + "default": "None", + "annotation": "torch.LongTensor | None" + }, { "name": "cp_context", "kind": "POSITIONAL_OR_KEYWORD", "default": "None", - "annotation": "_fla_cuda_fe2bb66.ops.cp.context.FLACPContext | None" + "annotation": "_fla_cuda_faa8674.ops.cp.context.FLACPContext | None" }, { "name": "kwargs", @@ -371,9 +389,9 @@ { "name": "fused_recurrent_gated_delta_rule", "qualname": "fused_recurrent_gated_delta_rule", - "module": "_fla_cuda_fe2bb66.ops.gated_delta_rule.fused_recurrent", + "module": "_fla_cuda_faa8674.ops.gated_delta_rule.fused_recurrent", "kind": "function", - "doc": "Args:\n q (torch.Tensor):\n queries of shape `[B, T, H, K]`.\n k (torch.Tensor):\n keys of shape `[B, T, H, K]`.\n v (torch.Tensor):\n values of shape `[B, T, HV, V]`.\n GVA (Grouped Value Attention) is applied if `HV > H`, where `HV` must be divisible by `H`.\n g (torch.Tensor):\n g (decays) of shape `[B, T, HV]`. Default: `None`.\n When `use_gate_in_kernel=False` (default), `g` must be in log space (pre-computed decay).\n When `use_gate_in_kernel=True`, `g` is the raw pre-activation input; the kernel fuses\n `-exp(A_log) * softplus(g + dt_bias)` internally per step.\n gk (torch.Tensor):\n gk (decays) of shape `[B, T, HV, K]`. Default: `None`.\n gv (torch.Tensor):\n gv (decays) of shape `[B, T, HV, V]`. Default: `None`.\n beta (torch.Tensor):\n betas of shape `[B, T, HV]`.\n scale (Optional[float]):\n Scale factor for the RetNet attention scores.\n If not provided, it will default to `1 / sqrt(K)`. Default: `None`.\n initial_state (Optional[torch.Tensor]):\n Initial state of shape `[N, HV, K, V]` for `N` input sequences.\n For equal-length input sequences, `N` equals the batch size `B`.\n Default: `None`.\n output_final_state (Optional[bool]):\n Whether to output the final state of shape `[N, HV, K, V]`. Default: `False`.\n use_qk_l2norm_in_kernel (Optional[bool]):\n Whether to use L2 normalization in the kernel. Default: `False`.\n use_gate_in_kernel (bool):\n Whether to compute the log-space GDN decay internally.\n When `True`, `g` is the raw input and `A_log` must be provided; the kernel fuses\n gate activation into the recurrence. Default: `False`.\n A_log (Optional[torch.Tensor]):\n Decay parameter of shape `[HV]`. Required when `use_gate_in_kernel=True`.\n dt_bias (Optional[torch.Tensor]):\n Bias added to `g` before activation, of shape `[HV]`.\n Only used when `use_gate_in_kernel=True`.\n use_beta_sigmoid_in_kernel (Optional[bool]):\n Whether to apply `torch.sigmoid(beta)` inside the kernel.\n - If `True`, the passed `beta` acts as the raw beta logits.\n - If `False`, `beta` is expected to already be in post-sigmoid space.\n Default: `False`.\n allow_neg_eigval (Optional[bool]):\n Whether to allow negative eigenvalues by scaling `beta` to `[0, 2)`.\n Only takes effect together with `use_beta_sigmoid_in_kernel=True`, in which case\n the kernel computes `2 * sigmoid(beta)` instead of `sigmoid(beta)`. Default: `False`.\n state_v_first (Optional[bool]):\n Store the recurrent state in V-first ``[V, K]`` layout instead of the default ``[K, V]``. Default: ``False``.\n cu_seqlens (torch.LongTensor):\n Cumulative sequence lengths of shape `[N+1]` used for variable-length training,\n consistent with the FlashAttention API.\n\nReturns:\n o (torch.Tensor):\n Outputs of shape `[B, T, HV, V]`.\n final_state (torch.Tensor):\n Final state of shape `[N, HV, K, V]` if `output_final_state=True` else `None`.\n\nExamples::\n >>> import torch\n >>> import torch.nn.functional as F\n >>> from einops import rearrange\n >>> from ...ops.gated_delta_rule import fused_recurrent_gated_delta_rule\n # inputs with equal lengths\n >>> B, T, H, HV, K, V = 4, 2048, 4, 8, 512, 512\n >>> q = torch.randn(B, T, H, K, device='cuda')\n >>> k = F.normalize(torch.randn(B, T, H, K, device='cuda'), p=2, dim=-1)\n >>> v = torch.randn(B, T, HV, V, device='cuda')\n >>> g = F.logsigmoid(torch.rand(B, T, HV, device='cuda'))\n >>> beta = torch.rand(B, T, HV, device='cuda').sigmoid()\n >>> h0 = torch.randn(B, HV, K, V, device='cuda')\n >>> o, ht = fused_gated_recurrent_delta_rule(\n q, k, v, g, beta,\n initial_state=h0,\n output_final_state=True\n )\n # for variable-length inputs, the batch size `B` is expected to be 1 and `cu_seqlens` is required\n >>> q, k, v, g, beta = map(lambda x: rearrange(x, 'b t ... -> 1 (b t) ...'), (q, k, v, g, beta))\n # for a batch with 4 sequences, `cu_seqlens` with 5 start/end positions are expected\n >>> cu_seqlens = q.new_tensor([0, 2048, 4096, 6144, 8192], dtype=torch.long)\n >>> o, ht = fused_gated_recurrent_delta_rule(\n q, k, v, g, beta,\n initial_state=h0,\n output_final_state=True,\n cu_seqlens=cu_seqlens\n )", + "doc": "Args:\n q (torch.Tensor):\n queries of shape `[B, T, H, K]`.\n k (torch.Tensor):\n keys of shape `[B, T, H, K]`.\n v (torch.Tensor):\n values of shape `[B, T, HV, V]`.\n GVA (Grouped Value Attention) is applied if `HV > H`, where `HV` must be divisible by `H`.\n g (torch.Tensor):\n g (decays) of shape `[B, T, HV]`. Default: `None`.\n When `use_gate_in_kernel=False` (default), `g` must be in log space (pre-computed decay).\n When `use_gate_in_kernel=True`, `g` is the raw pre-activation input; the kernel fuses\n `-exp(A_log) * softplus(g + dt_bias)` internally per step.\n gk (torch.Tensor):\n gk (decays) of shape `[B, T, HV, K]`. Default: `None`.\n gv (torch.Tensor):\n gv (decays) of shape `[B, T, HV, V]`. Default: `None`.\n beta (torch.Tensor):\n betas of shape `[B, T, HV]`.\n scale (Optional[float]):\n Scale factor for the RetNet attention scores.\n If not provided, it will default to `1 / sqrt(K)`. Default: `None`.\n initial_state (Optional[torch.Tensor]):\n Initial state of shape `[N, HV, K, V]` for `N` input sequences.\n For equal-length input sequences, `N` equals the batch size `B`.\n Default: `None`.\n output_final_state (Optional[bool]):\n Whether to output the final state of shape `[N, HV, K, V]`. Default: `False`.\n use_qk_l2norm_in_kernel (Optional[bool]):\n Whether to use L2 normalization in the kernel. Default: `False`.\n use_gate_in_kernel (bool):\n Whether to compute the log-space GDN decay internally.\n When `True`, `g` is the raw input and `A_log` must be provided; the kernel fuses\n gate activation into the recurrence. Default: `False`.\n A_log (Optional[torch.Tensor]):\n Decay parameter of shape `[HV]`. Required when `use_gate_in_kernel=True`.\n dt_bias (Optional[torch.Tensor]):\n Bias added to `g` before activation, of shape `[HV]`.\n Only used when `use_gate_in_kernel=True`.\n use_beta_sigmoid_in_kernel (Optional[bool]):\n Whether to apply `torch.sigmoid(beta)` inside the kernel.\n - If `True`, the passed `beta` acts as the raw beta logits.\n - If `False`, `beta` is expected to already be in post-sigmoid space.\n Default: `False`.\n allow_neg_eigval (Optional[bool]):\n Whether to allow negative eigenvalues by scaling `beta` to `[0, 2)`.\n Only takes effect together with `use_beta_sigmoid_in_kernel=True`, in which case\n the kernel computes `2 * sigmoid(beta)` instead of `sigmoid(beta)`. Default: `False`.\n state_v_first (Optional[bool]):\n Store the recurrent state in V-first ``[V, K]`` layout instead of the default ``[K, V]``. Default: ``False``.\n cu_seqlens (torch.LongTensor):\n Cumulative sequence lengths of shape `[N+1]` used for variable-length training,\n consistent with the FlashAttention API.\n\nReturns:\n o (torch.Tensor):\n Outputs of shape `[B, T, HV, V]`.\n final_state (torch.Tensor):\n Final state of shape `[N, HV, K, V]` if `output_final_state=True` else `None`.\n\nExamples::\n >>> import torch\n >>> import torch.nn.functional as F\n >>> from einops import rearrange\n >>> from fla.ops.gated_delta_rule import fused_recurrent_gated_delta_rule\n # inputs with equal lengths\n >>> B, T, H, HV, K, V = 4, 2048, 4, 8, 512, 512\n >>> q = torch.randn(B, T, H, K, device='cuda')\n >>> k = F.normalize(torch.randn(B, T, H, K, device='cuda'), p=2, dim=-1)\n >>> v = torch.randn(B, T, HV, V, device='cuda')\n >>> g = F.logsigmoid(torch.rand(B, T, HV, device='cuda'))\n >>> beta = torch.rand(B, T, HV, device='cuda').sigmoid()\n >>> h0 = torch.randn(B, HV, K, V, device='cuda')\n >>> o, ht = fused_gated_recurrent_delta_rule(\n q, k, v, g, beta,\n initial_state=h0,\n output_final_state=True\n )\n # for variable-length inputs, the batch size `B` is expected to be 1 and `cu_seqlens` is required\n >>> q, k, v, g, beta = map(lambda x: rearrange(x, 'b t ... -> 1 (b t) ...'), (q, k, v, g, beta))\n # for a batch with 4 sequences, `cu_seqlens` with 5 start/end positions are expected\n >>> cu_seqlens = q.new_tensor([0, 2048, 4096, 6144, 8192], dtype=torch.long)\n >>> o, ht = fused_gated_recurrent_delta_rule(\n q, k, v, g, beta,\n initial_state=h0,\n output_final_state=True,\n cu_seqlens=cu_seqlens\n )", "signature": { "parameters": [ { @@ -499,7 +517,7 @@ { "name": "FusedRMSNormGated", "qualname": "FusedRMSNormGated", - "module": "_fla_cuda_fe2bb66.layers", + "module": "_fla_cuda_faa8674.layers", "kind": "class", "doc": "Base class for all neural network modules.\n\nYour models should also subclass this class.\n\nModules can also contain other Modules, allowing them to be nested in\na tree structure. You can assign the submodules as regular attributes::\n\n import torch.nn as nn\n import torch.nn.functional as F\n\n\n class Model(nn.Module):\n def __init__(self) -> None:\n super().__init__()\n self.conv1 = nn.Conv2d(1, 20, 5)\n self.conv2 = nn.Conv2d(20, 20, 5)\n\n def forward(self, x):\n x = F.relu(self.conv1(x))\n return F.relu(self.conv2(x))\n\nSubmodules assigned in this way will be registered, and will also have their\nparameters converted when you call :meth:`to`, etc.\n\n.. note::\n As per the example above, an ``__init__()`` call to the parent class\n must be made before assignment on the child.\n\n:ivar training: Boolean represents whether this module is in training or\n evaluation mode.\n:vartype training: bool", "signature": { @@ -521,13 +539,13 @@ }, "attributes": { "has_backward": null, - "can_torch_compile": null + "can_torch_compile": true } }, { "name": "chunk_kimi_delta_attention", "qualname": "chunk_kimi_delta_attention", - "module": "_fla_cuda_fe2bb66.layers", + "module": "_fla_cuda_faa8674.layers", "kind": "class", "doc": "Base class for all neural network modules.\n\nYour models should also subclass this class.\n\nModules can also contain other Modules, allowing them to be nested in\na tree structure. You can assign the submodules as regular attributes::\n\n import torch.nn as nn\n import torch.nn.functional as F\n\n\n class Model(nn.Module):\n def __init__(self) -> None:\n super().__init__()\n self.conv1 = nn.Conv2d(1, 20, 5)\n self.conv2 = nn.Conv2d(20, 20, 5)\n\n def forward(self, x):\n x = F.relu(self.conv1(x))\n return F.relu(self.conv2(x))\n\nSubmodules assigned in this way will be registered, and will also have their\nparameters converted when you call :meth:`to`, etc.\n\n.. note::\n As per the example above, an ``__init__()`` call to the parent class\n must be made before assignment on the child.\n\n:ivar training: Boolean represents whether this module is in training or\n evaluation mode.\n:vartype training: bool", "signature": { @@ -549,13 +567,13 @@ }, "attributes": { "has_backward": null, - "can_torch_compile": null + "can_torch_compile": true } }, { "name": "recurrent_kimi_delta_attention", "qualname": "recurrent_kimi_delta_attention", - "module": "_fla_cuda_fe2bb66.layers", + "module": "_fla_cuda_faa8674.layers", "kind": "class", "doc": "Base class for all neural network modules.\n\nYour models should also subclass this class.\n\nModules can also contain other Modules, allowing them to be nested in\na tree structure. You can assign the submodules as regular attributes::\n\n import torch.nn as nn\n import torch.nn.functional as F\n\n\n class Model(nn.Module):\n def __init__(self) -> None:\n super().__init__()\n self.conv1 = nn.Conv2d(1, 20, 5)\n self.conv2 = nn.Conv2d(20, 20, 5)\n\n def forward(self, x):\n x = F.relu(self.conv1(x))\n return F.relu(self.conv2(x))\n\nSubmodules assigned in this way will be registered, and will also have their\nparameters converted when you call :meth:`to`, etc.\n\n.. note::\n As per the example above, an ``__init__()`` call to the parent class\n must be made before assignment on the child.\n\n:ivar training: Boolean represents whether this module is in training or\n evaluation mode.\n:vartype training: bool", "signature": { @@ -577,13 +595,13 @@ }, "attributes": { "has_backward": null, - "can_torch_compile": null + "can_torch_compile": true } }, { "name": "chunk_gated_delta_rule", "qualname": "chunk_gated_delta_rule", - "module": "_fla_cuda_fe2bb66.layers", + "module": "_fla_cuda_faa8674.layers", "kind": "class", "doc": "Base class for all neural network modules.\n\nYour models should also subclass this class.\n\nModules can also contain other Modules, allowing them to be nested in\na tree structure. You can assign the submodules as regular attributes::\n\n import torch.nn as nn\n import torch.nn.functional as F\n\n\n class Model(nn.Module):\n def __init__(self) -> None:\n super().__init__()\n self.conv1 = nn.Conv2d(1, 20, 5)\n self.conv2 = nn.Conv2d(20, 20, 5)\n\n def forward(self, x):\n x = F.relu(self.conv1(x))\n return F.relu(self.conv2(x))\n\nSubmodules assigned in this way will be registered, and will also have their\nparameters converted when you call :meth:`to`, etc.\n\n.. note::\n As per the example above, an ``__init__()`` call to the parent class\n must be made before assignment on the child.\n\n:ivar training: Boolean represents whether this module is in training or\n evaluation mode.\n:vartype training: bool", "signature": { @@ -605,13 +623,13 @@ }, "attributes": { "has_backward": null, - "can_torch_compile": null + "can_torch_compile": true } }, { "name": "recurrent_gated_delta_rule", "qualname": "recurrent_gated_delta_rule", - "module": "_fla_cuda_fe2bb66.layers", + "module": "_fla_cuda_faa8674.layers", "kind": "class", "doc": "Base class for all neural network modules.\n\nYour models should also subclass this class.\n\nModules can also contain other Modules, allowing them to be nested in\na tree structure. You can assign the submodules as regular attributes::\n\n import torch.nn as nn\n import torch.nn.functional as F\n\n\n class Model(nn.Module):\n def __init__(self) -> None:\n super().__init__()\n self.conv1 = nn.Conv2d(1, 20, 5)\n self.conv2 = nn.Conv2d(20, 20, 5)\n\n def forward(self, x):\n x = F.relu(self.conv1(x))\n return F.relu(self.conv2(x))\n\nSubmodules assigned in this way will be registered, and will also have their\nparameters converted when you call :meth:`to`, etc.\n\n.. note::\n As per the example above, an ``__init__()`` call to the parent class\n must be made before assignment on the child.\n\n:ivar training: Boolean represents whether this module is in training or\n evaluation mode.\n:vartype training: bool", "signature": { @@ -633,7 +651,7 @@ }, "attributes": { "has_backward": null, - "can_torch_compile": null + "can_torch_compile": true } } ] diff --git a/build/torch-cuda/utils/__init__.py b/build/torch-cuda/utils/__init__.py index 0d2fbe5e9295cb9b75e2eb8fc86e417468a9e632..5fcb5d342c4cf8e0e5b26af0915e33d13da3d523 100644 --- a/build/torch-cuda/utils/__init__.py +++ b/build/torch-cuda/utils/__init__.py @@ -5,23 +5,7 @@ # For a list of all contributors, visit: # https://github.com/fla-org/flash-linear-attention/graphs/contributors -import sys - -from ._compat import ( # noqa: F401 - SUPPORTS_AUTOTUNE_CACHE, - TRITON_ABOVE_3_4_0, - TRITON_ABOVE_3_5_1, - TRITON_ABOVE_3_7_1, - autotune_cache_kwargs, - find_spec_cached, -) -from ._config import ( # noqa: F401 - FLA_CACHE_RESULTS, - FLA_CI_ENV, - FLA_DISABLE_TENSOR_CACHE, - FLA_TENSOR_CACHE_SIZE, -) -from ._decorators import ( # noqa: F401 +from ..utils.decorators import ( Action, checkpoint, contiguous, @@ -30,8 +14,28 @@ from ._decorators import ( # noqa: F401 require_version, tensor_cache, ) -from ._device import ( # noqa: F401 +from ..utils.env import ( + FLA_CACHE_RESULTS, + FLA_CI_ENV, + FLA_DISABLE_TENSOR_CACHE, + FLA_TENSOR_CACHE_SIZE, + SUPPORTS_AUTOTUNE_CACHE, + TRITON_ABOVE_3_4_0, + TRITON_ABOVE_3_5_1, + TRITON_ABOVE_3_6_0, + TRITON_ABOVE_3_7_1, + TRITON_ABOVE_3_8_0, + autotune_cache_kwargs, + check_environments, + check_pytorch_version, + check_transformers_version, + find_spec_cached, + has_usable_nvcc, + supports_transformers_next_sequence_length, +) +from ..utils.hardware import ( IS_AMD, + IS_AMD_TMA_ARCH, IS_ARM, IS_GATHER_SUPPORTED, IS_INTEL, @@ -40,13 +44,13 @@ from ._device import ( # noqa: F401 IS_NVIDIA, IS_NVIDIA_BLACKWELL, IS_NVIDIA_HOPPER, + IS_NVIDIA_SM100, + IS_NVIDIA_SM120, IS_TF32_SUPPORTED, IS_TMA_SUPPORTED, Backend, autocast_custom_bwd, autocast_custom_fwd, - check_environments, - check_pytorch_version, check_shared_mem, custom_device_ctx, device, @@ -55,31 +59,25 @@ from ._device import ( # noqa: F401 device_torch_lib, get_all_max_shared_mem, get_available_device, + get_device_arch, + get_device_capability, + get_device_smem_optin, get_multiprocessor_count, map_triton_backend_to_torch_device, ) -from ._testing import assert_close, get_abs_err, get_err_ratio # noqa: F401 - - -def _register_aliases(): - current_module = sys.modules[__name__] - for key in ( - 'IS_AMD', - 'IS_ARM', - 'IS_INTEL', - 'IS_INTEL_ALCHEMIST', - 'IS_NVIDIA', - 'IS_NPU', - 'IS_NVIDIA_BLACKWELL', - 'IS_NVIDIA_HOPPER', - 'IS_TF32_SUPPORTED', - 'IS_GATHER_SUPPORTED', - 'IS_TMA_SUPPORTED', - ): - if hasattr(current_module, key): - setattr(current_module, key.lower(), getattr(current_module, key)) - - -_register_aliases() +from ..utils.testing import assert_close, get_abs_err, get_err_ratio -del _register_aliases +is_amd = IS_AMD +is_amd_tma_arch = IS_AMD_TMA_ARCH +is_arm = IS_ARM +is_intel = IS_INTEL +is_intel_alchemist = IS_INTEL_ALCHEMIST +is_nvidia = IS_NVIDIA +is_npu = IS_NPU +is_nvidia_blackwell = IS_NVIDIA_BLACKWELL +is_nvidia_hopper = IS_NVIDIA_HOPPER +is_nvidia_sm100 = IS_NVIDIA_SM100 +is_nvidia_sm120 = IS_NVIDIA_SM120 +is_tf32_supported = IS_TF32_SUPPORTED +is_gather_supported = IS_GATHER_SUPPORTED +is_tma_supported = IS_TMA_SUPPORTED diff --git a/build/torch-cuda/utils/_compat.py b/build/torch-cuda/utils/_compat.py deleted file mode 100644 index a702e300e509b8d436f2c9349daa696ded2b5a9c..0000000000000000000000000000000000000000 --- a/build/torch-cuda/utils/_compat.py +++ /dev/null @@ -1,27 +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 functools -import inspect -from importlib.util import find_spec - -import triton -from packaging import version as package_version - -from ._config import FLA_CACHE_RESULTS - -TRITON_ABOVE_3_4_0 = package_version.parse(triton.__version__) >= package_version.parse("3.4.0") -TRITON_ABOVE_3_5_1 = package_version.parse(triton.__version__) >= package_version.parse("3.5.1") -TRITON_ABOVE_3_7_1 = package_version.parse(triton.__version__) >= package_version.parse("3.7.1") - -SUPPORTS_AUTOTUNE_CACHE = "cache_results" in inspect.signature(triton.autotune).parameters -autotune_cache_kwargs = {"cache_results": FLA_CACHE_RESULTS} if SUPPORTS_AUTOTUNE_CACHE else {} - - -@functools.cache -def find_spec_cached(name): - return find_spec(name) diff --git a/build/torch-cuda/utils/_config.py b/build/torch-cuda/utils/_config.py deleted file mode 100644 index 79b6ce8a784f457b847834a6f3f6128a06bcf3f3..0000000000000000000000000000000000000000 --- a/build/torch-cuda/utils/_config.py +++ /dev/null @@ -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 - -import os - -FLA_CI_ENV = os.getenv("FLA_CI_ENV") == "1" -FLA_CACHE_RESULTS = os.getenv('FLA_CACHE_RESULTS', '1') == '1' - -FLA_DISABLE_TENSOR_CACHE = os.getenv('FLA_DISABLE_TENSOR_CACHE', '0') == '1' -try: - FLA_TENSOR_CACHE_SIZE = int(os.getenv('FLA_TENSOR_CACHE_SIZE', "4")) -except ValueError: - FLA_TENSOR_CACHE_SIZE = 4 diff --git a/build/torch-cuda/utils/_decorators.py b/build/torch-cuda/utils/_decorators.py deleted file mode 100644 index f29063d557d8e6cb12f43643fa6964e838c83f4c..0000000000000000000000000000000000000000 --- a/build/torch-cuda/utils/_decorators.py +++ /dev/null @@ -1,338 +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 contextlib -import functools -import inspect -import sys -import warnings -from collections import deque -from collections.abc import Callable -from enum import Enum -from typing import Any - -import torch -from packaging import version as package_version - -from ._config import FLA_DISABLE_TENSOR_CACHE, FLA_TENSOR_CACHE_SIZE -from ._device import custom_device_ctx - - -__version__ = "0.5.2" # kept for BC - - -class Action(Enum): - NONE = "none" - NOTIFY = "notify" - NOTIFY_ALWAYS = "notify_always" - RAISE = "raise" - - -def tensor_cache( - fn: Callable[..., torch.Tensor], -) -> Callable[..., torch.Tensor]: - """ - A decorator that memoizes the most recent results of a function call by argument identity. - - The decorator keeps a bounded queue of up to ``FLA_TENSOR_CACHE_SIZE`` (default 4) - recent ``(args, kwargs, result)`` triples. On each call, every cached entry is checked - in order; an entry is considered a hit when the positional arg count and kwarg key set - match and every argument is the *same object* (``is`` identity) as the cached one. On a - hit the cached result is returned and ``fn`` is skipped; on a miss ``fn`` is invoked and - the new triple is appended (evicting the oldest when the queue is full). - - Caching is fully bypassed when the ``FLA_DISABLE_TENSOR_CACHE`` environment variable is - set to ``'1'``. - - Args: - fn (Callable[..., torch.Tensor]): - The function to be decorated. Intended for functions whose inputs are tensors - (or other objects compared by identity) and whose output is a tensor. - - Returns: - Callable[..., torch.Tensor]: - A wrapped version of ``fn`` backed by an identity-based bounded cache. - """ - cached: deque = deque(maxlen=FLA_TENSOR_CACHE_SIZE) - - def cache_disabled() -> bool: - utils_module = sys.modules.get('fla.utils') - return getattr(utils_module, 'FLA_DISABLE_TENSOR_CACHE', FLA_DISABLE_TENSOR_CACHE) - - @functools.wraps(fn) - def wrapper(*args: Any, **kwargs: Any) -> Any: - if cache_disabled(): - return fn(*args, **kwargs) - - for cached_args, cached_kwargs, cached_result in cached: - if len(args) != len(cached_args) or len(kwargs) != len(cached_kwargs): - continue - if all(a is b for a, b in zip(args, cached_args, strict=False)) and \ - all(k in cached_kwargs and v is cached_kwargs[k] for k, v in kwargs.items()): - return cached_result - - result = fn(*args, **kwargs) - cached.append((args, kwargs, result)) - return result - - return wrapper - - -def _skip_contiguous( - no_guard_contiguous: bool | list[str] | tuple[str, ...] | set[str], - param_name: str, - skip_params: set[str], -) -> bool: - return no_guard_contiguous is True or param_name in skip_params - - -def _contiguous_if_needed(arg: Any, skip: bool) -> Any: - if isinstance(arg, torch.Tensor) and not skip: - return arg.contiguous() - return arg - - -def input_guard( - fn: Callable[..., torch.Tensor] | None = None, - *, - no_guard_contiguous: bool | list[str] | tuple[str, ...] | set[str] = False, -) -> Callable[[Callable[..., torch.Tensor]], Callable[..., torch.Tensor]] | Callable[..., torch.Tensor]: - """ - A decorator to make sure all input tensors are contiguous and set the device based on input tensors. - - Args: - no_guard_contiguous (bool | list[str] | tuple[str, ...] | set[str]): - If True, skip all contiguous checks. If a list/tuple/set of parameter names, skip contiguous check for those parameters. - """ - - def decorator(fn: Callable[..., torch.Tensor]) -> Callable[..., torch.Tensor]: - # Get function signature for parameter name mapping - sig = inspect.signature(fn) - param_names = list(sig.parameters.keys()) - skip_params = set(no_guard_contiguous) if isinstance(no_guard_contiguous, (list, tuple, set)) else set() - - @functools.wraps(fn) - def wrapper(*args, **kwargs): - # Process args with parameter name mapping - processed_args = [] - for i, arg in enumerate(args): - if i < len(param_names): - param_name = param_names[i] - else: - # For *args beyond signature, use position as name - param_name = f"__arg_{i}" - - processed_args.append(_contiguous_if_needed( - arg, _skip_contiguous(no_guard_contiguous, param_name, skip_params))) - - # Process kwargs - processed_kwargs = {} - for k, v in kwargs.items(): - processed_kwargs[k] = _contiguous_if_needed(v, _skip_contiguous(no_guard_contiguous, k, skip_params)) - - tensor = None - for arg in args: - if isinstance(arg, torch.Tensor): - tensor = arg - break - if tensor is None: - for value in kwargs.values(): - if isinstance(value, torch.Tensor): - tensor = value - break - - if tensor is not None: - ctx = custom_device_ctx(tensor.device.index) - else: - ctx = contextlib.nullcontext() - - with ctx: - return fn(*processed_args, **processed_kwargs) - - return wrapper - - # Handle direct usage without parentheses: @input_guard - if fn is not None: - return decorator(fn) - - return decorator - - -def contiguous(fn: Callable[..., torch.Tensor]) -> Callable[..., torch.Tensor]: - """Alias for input_guard() without parameters.""" - return input_guard(fn) - - -def require_version(version, hint): - """ - Perform a runtime check of the dependency versions, using the exact same syntax used by pip. - """ - def decorator(fn): - @functools.wraps(fn) - def wrapper(ctx, *args, **kwargs): - from transformers.utils.versions import require_version - require_version(version, hint) - return fn( - ctx, - *(i if not isinstance(i, torch.Tensor) else i.contiguous() for i in args), - **{k: (v if not isinstance(v, torch.Tensor) else v.contiguous()) for k, v in kwargs.items()}, - ) - return wrapper - return decorator - - -def deprecate_kwarg( - old_name: str, - version: str, - new_name: str | None = None, - warn_if_greater_or_equal_version: bool = False, - raise_if_greater_or_equal_version: bool = False, - raise_if_both_names: bool = False, - additional_message: str | None = None, -): - """ - Decorator to notify users about deprecated keyword arguments, replacing them with a new name if specified. - - This decorator allows you to: - - Notify users when a keyword argument is deprecated. - - Automatically replace deprecated keyword arguments with new ones. - - Raise an error if deprecated arguments are used, depending on the specified conditions. - - By default, the decorator notifies the user about the deprecated argument while the `fla.__version__` < specified `version` - in the decorator. To keep notifications with any version `warn_if_greater_or_equal_version=True` can be set. - - Args: - old_name (`str`): - Name of the deprecated keyword argument. - version (`str`): - The version in which the keyword argument was (or will be) deprecated. - new_name (`Optional[str]`, *optional*): - The new name for the deprecated keyword argument. - If specified, the deprecated keyword argument will be replaced with this new name. - warn_if_greater_or_equal_version (`bool`, *optional*, defaults to `False`): - Whether to show warning if current `fla` version is greater or equal to the deprecated version. - raise_if_greater_or_equal_version (`bool`, *optional*, defaults to `False`): - Whether to raise `ValueError` if current `fla` version is greater or equal to the deprecated version. - raise_if_both_names (`bool`, *optional*, defaults to `False`): - Whether to raise `ValueError` if both deprecated and new keyword arguments are set. - additional_message (`Optional[str]`, *optional*): - An additional message to append to the default deprecation message. - - Raises: - ValueError: - If `raise_if_greater_or_equal_version` is `True` and the current version >= the deprecated one, - or if `raise_if_both_names` is `True` and both old and new keyword arguments are provided. - - Returns: - Callable: - A wrapped function that handles the deprecated keyword arguments according to the specified parameters. - - Example usage with renaming argument: - - ```python - @deprecate_kwarg("reduce_labels", new_name="do_reduce_labels", version="6.0.0") - def my_function(do_reduce_labels): - print(do_reduce_labels) - - my_function(reduce_labels=True) # Will show a deprecation warning and use do_reduce_labels=True - ``` - - Example usage without renaming argument: - - ```python - @deprecate_kwarg("max_size", version="6.0.0") - def my_function(max_size): - print(max_size) - - my_function(max_size=1333) # Will show a deprecation warning - ``` - - """ - deprecated_version = package_version.parse(version) - current_version = package_version.parse(__version__) - is_greater_or_equal_version = current_version >= deprecated_version - - if is_greater_or_equal_version: - version_message = f"and removed starting from version {version}" - else: - version_message = f"and will be removed in version {version}" - - def wrapper(func): - # Required for better warning message - sig = inspect.signature(func) - function_named_args = set(sig.parameters.keys()) - is_instance_method = "self" in function_named_args - is_class_method = "cls" in function_named_args - - @functools.wraps(func) - def wrapped_func(*args, **kwargs): - # Get class + function name (just for better warning message) - func_name = func.__name__ - if is_instance_method: - func_name = f"{args[0].__class__.__name__}.{func_name}" - elif is_class_method: - func_name = f"{args[0].__name__}.{func_name}" - - minimum_action = Action.NONE - message = None - - # deprecated kwarg and its new version are set for function call -> replace it with new name - if old_name in kwargs and new_name in kwargs: - minimum_action = Action.RAISE if raise_if_both_names else Action.NOTIFY_ALWAYS - message = ( - f"Both `{old_name}` and `{new_name}` are set for `{func_name}`. " - f"Using `{new_name}={kwargs[new_name]}` and ignoring deprecated `{old_name}={kwargs[old_name]}`." - ) - kwargs.pop(old_name) - - # only deprecated kwarg is set for function call -> replace it with new name - elif old_name in kwargs and new_name is not None and new_name not in kwargs: - minimum_action = Action.NOTIFY - message = ( - f"`{old_name}` is deprecated {version_message} for `{func_name}`. " - f"Use `{new_name}` instead." - ) - kwargs[new_name] = kwargs.pop(old_name) - - # deprecated kwarg is not set for function call and new name is not specified -> just notify - elif old_name in kwargs: - minimum_action = Action.NOTIFY - message = f"`{old_name}` is deprecated {version_message} for `{func_name}`." - - if message is not None and additional_message is not None: - message = f"{message} {additional_message}" - - # update minimum_action if argument is ALREADY deprecated (current version >= deprecated version) - if is_greater_or_equal_version: - # change to (NOTIFY, NOTIFY_ALWAYS) -> RAISE if specified - # in case we want to raise error for already deprecated arguments - if raise_if_greater_or_equal_version and minimum_action != Action.NONE: - minimum_action = Action.RAISE - - # change to NOTIFY -> NONE if specified (NOTIFY_ALWAYS can't be changed to NONE) - elif not warn_if_greater_or_equal_version and minimum_action == Action.NOTIFY: - minimum_action = Action.NONE - - # raise error or notify user - if minimum_action == Action.RAISE: - raise ValueError(message) - elif minimum_action in (Action.NOTIFY, Action.NOTIFY_ALWAYS): - # DeprecationWarning is ignored by default, so we use FutureWarning instead - warnings.warn(message, FutureWarning, stacklevel=2) - - return func(*args, **kwargs) - - return wrapped_func - - return wrapper - - -def checkpoint(fn): - @functools.wraps(fn) - def wrapper(*args, **kwargs): - return torch.utils.checkpoint.checkpoint(fn, *args, **kwargs) - return wrapper diff --git a/build/torch-cuda/utils/_testing.py b/build/torch-cuda/utils/_testing.py deleted file mode 100644 index 12064dcc4d48c88a89856afac93db66aceaeaa79..0000000000000000000000000000000000000000 --- a/build/torch-cuda/utils/_testing.py +++ /dev/null @@ -1,41 +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 logging -import warnings - -import torch - -from ._config import FLA_CI_ENV - -logger = logging.getLogger(__name__) - - -def get_abs_err(x, y): - return (x.detach() - y.detach()).flatten().abs().max().item() - - -def get_err_ratio(x, y): - err = (x.detach() - y.detach()).flatten().square().mean().sqrt().item() - base = (x.detach()).flatten().square().mean().sqrt().item() - return err / (base + 1e-8) - - -def assert_close(prefix, ref, tri, ratio, warning=False, err_atol=1e-6): - abs_atol = get_abs_err(ref, tri) - error_rate = get_err_ratio(ref, tri) - msg = f"{prefix:>16} diff: {abs_atol:.6f} ratio: {error_rate:.6f}" - logger.info(msg) - if abs_atol <= err_atol: - return - assert not torch.isnan(ref).any(), f"{prefix}: NaN detected in ref" - assert not torch.isnan(tri).any(), f"{prefix}: NaN detected in tri" - if warning or (FLA_CI_ENV and (error_rate < 0.01 or abs_atol <= 0.3)): - if error_rate > ratio: - warnings.warn(msg) - else: - assert error_rate < ratio, msg diff --git a/build/torch-cuda/utils/ascend_ub_manager.py b/build/torch-cuda/utils/ascend_ub_manager.py deleted file mode 100644 index e1760b4345e19798339058537e7bb0ba45024a1d..0000000000000000000000000000000000000000 --- a/build/torch-cuda/utils/ascend_ub_manager.py +++ /dev/null @@ -1,576 +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 -# copied from https://github.com/linkedin/Liger-Kernel/blob/main/src/liger_kernel/ops/backends/_ascend/ub_manager.py - -""" -Unified Buffer (UB) Manager for Ascend NPU. - -This module provides UB capacity detection and tiling strategy computation -for running Triton kernels on Ascend NPU. It automatically calculates -optimal block sizes based on UB capacity constraints to prevent UB overflow. -""" - -import os -import warnings - -import torch -import triton - -# Ascend Triton launch limits (see triton_ascend/activations.py). -ASCEND_MAX_GRID_DIM = 65535 -# Legacy fused kernel byte cap (65536 // fp16 element size). -_FALLBACK_MAX_FUSED_BLOCK = 65536 // 2 -# Conservative UB fallback when runtime detection is unavailable (64 KiB). -_FALLBACK_UB_CAPACITY_BITS = 65536 * 8 - - -def is_npu_available() -> bool: - """Detect Ascend NPU availability.""" - if hasattr(torch, 'npu'): - try: - return torch.npu.is_available() - except Exception: - pass - try: - from transformers.utils import is_torch_npu_available - - return is_torch_npu_available() - except ImportError: - return False - - -def _fallback_ub_capacity(reason: str) -> int: - warnings.warn( - f"Using fallback UB capacity ({_FALLBACK_UB_CAPACITY_BITS // 8} bytes): {reason}", - stacklevel=3, - ) - return _FALLBACK_UB_CAPACITY_BITS - - -def _normalize_tiling_dims(tiling_dim: int | tuple[int, ...]) -> set: - """ - Normalize tiling dimension specification to a set of dimension indices. - - Args: - tiling_dim: Either an int (single dimension) or tuple of ints (multiple dimensions). - - Returns: - Set of dimension indices that can be tiled. - """ - if isinstance(tiling_dim, int): - return {tiling_dim} - elif isinstance(tiling_dim, tuple): - return set(tiling_dim) - else: - return set() - - -def _default_strategy( - ub_capacity_bits: int, - safety_margin: float, - dtype_size: int, - memory_multiplier: float, - shapes: tuple[tuple[int, ...], ...], - tiling_dims: tuple[int | tuple[int, ...], ...], -) -> tuple[int, ...]: - """ - Default tiling strategy: calculate maximum safe block size based on UB capacity. - - This is a unified strategy function that works for all kernels by abstracting - the memory calculation as: memory_multiplier * BLOCK_SIZE * unit_param * dtype_size * 8 bits - - Args: - ub_capacity_bits: UB capacity in bits - safety_margin: Safety margin as a float (e.g., 0.80 for 80%) - dtype_size: Size of data type in bytes (e.g., 2 for float16, 4 for float32) - memory_multiplier: Memory multiplier for estimating peak memory usage - shapes: Tuple of full shapes. Each shape is a tuple of dimension sizes. - - For ROPE: ((n_q_head, hd), (n_kv_head, hd)) - - For GEGLU: ((n_cols,),) - tiling_dims: Tuple specifying which dimensions can be tiled for each shape. - Each element can be: - - int: single dimension index (e.g., 0 for first dimension) - - tuple of ints: multiple dimensions that can be tiled together - - For ROPE: (0, 0) means first dimension of each shape can be tiled - - For GEGLU: (0,) means first dimension of the shape can be tiled - Length must match len(shapes). - - Returns: - Tuple of maximum safe block sizes, one for each shape. - Each element is a power of 2. - - Note: - For each shape, fixed dimensions (non-tiling) are multiplied together to get unit_param. - The final block size is computed in compute_default_tiling_strategy by taking - min(desired_block_size, max_safe_block_size) where desired_block_size = triton.next_power_of_2(original_dim). - """ - if not shapes or not tiling_dims: - return () - - # Calculate max_safe_block_size for each tiling dimension - max_safe_sizes = [] - - for shape, tiling_dim in zip(shapes, tiling_dims): - # Normalize tiling_dim to a set of dimension indices - tiling_dim_set = _normalize_tiling_dims(tiling_dim) - - # Validate tiling dimensions are within shape bounds - if not tiling_dim_set: - raise ValueError( - f"Invalid tiling_dim: {tiling_dim}. tiling_dim must be an int or a non-empty tuple of ints." - ) - if any(dim_idx < 0 or dim_idx >= len(shape) for dim_idx in tiling_dim_set): - raise ValueError( - f"Invalid tiling_dim: {tiling_dim} for shape {shape}. " - f"All dimension indices must be in range [0, {len(shape)})." - ) - - # Calculate unit_param: product of fixed (non-tiling) dimensions - unit_param = 1.0 - for dim_idx, dim_size in enumerate(shape): - if dim_idx not in tiling_dim_set: - if dim_size <= 0: - # Invalid dimension size, use conservative default - unit_param = 1.0 - break - unit_param *= float(dim_size) - - # Ensure unit_param is at least 1.0 - if unit_param <= 0: - unit_param = 1.0 - - # Calculate maximum safe block size based on UB capacity - # Memory: memory_multiplier * BLOCK_SIZE * unit_param * dtype_size * 8 bits - SAFE_UB_CAPACITY_BITS = int(ub_capacity_bits * safety_margin) - - # Solve: memory_multiplier * BLOCK_SIZE * unit_param * dtype_size * 8 <= SAFE_UB_CAPACITY_BITS - # BLOCK_SIZE <= SAFE_UB_CAPACITY_BITS / (memory_multiplier * unit_param * dtype_size * 8) - max_block_size = int(SAFE_UB_CAPACITY_BITS // (memory_multiplier * unit_param * dtype_size * 8)) - max_block_size = max(1, max_block_size) - - # Find largest power of 2 <= max_block_size - # Use triton.next_power_of_2(max_block_size + 1) // 2 to get the largest power of 2 <= max_block_size - safe_block_size = triton.next_power_of_2(max_block_size + 1) // 2 - max_safe_sizes.append(safe_block_size) - - return tuple(max_safe_sizes) - - -class UBManager: - """ - Unified Buffer Manager for Ascend NPU. - - Provides UB capacity detection and management for Ascend NPU devices. - The UB capacity is used by tiling strategy functions to calculate optimal block sizes. - """ - - def __init__(self, ub_capacity_bits: int | None = None): - """ - Initialize UB Manager. - - Args: - ub_capacity_bits: UB capacity in bits. If None, will be detected automatically. - """ - self._npu_model = self._detect_npu_model() - self._ub_capacity_bits = ub_capacity_bits or self._detect_ub_capacity() - - @property - def ub_capacity_bits(self) -> int: - """Get UB capacity in bits.""" - return self._ub_capacity_bits - - @property - def ub_capacity_bytes(self) -> int: - """Get UB capacity in bytes.""" - return self._ub_capacity_bits // 8 - - @property - def npu_model(self) -> str: - """Get detected NPU model name.""" - return self._npu_model - - def _detect_npu_model(self) -> str: - """Detect NPU model from device properties.""" - if not is_npu_available(): - return "unknown" - - try: - dev_props = torch.npu.get_device_properties(0) - # Try to get model name from device properties - return dev_props.name - except Exception: - pass - - return "default" - - def _detect_ub_capacity(self) -> int: - """ - Detect UB capacity from environment variable or get_soc_spec. - - Returns: - UB capacity in bits. Falls back to a conservative default with a warning - when detection fails. - """ - # Check environment variable first (in bits) - env_capacity = os.getenv("ASCEND_UB_CAPACITY_BITS") - if env_capacity is not None: - try: - capacity_bits = int(env_capacity) - if capacity_bits > 0: - return capacity_bits - except ValueError: - pass - - # Try to get from get_soc_spec (returns bytes, convert to bits) - if is_npu_available(): - try: - from tbe.common.platform import get_soc_spec, set_current_compile_soc_info - - # Set current SOC info for get_soc_spec to work correctly - device = torch.npu - soc_info = device.get_device_name(device.current_device()) - set_current_compile_soc_info(soc_info) - - # Query UB size (get_soc_spec returns size in bytes) - ub_size_bytes = get_soc_spec("UB_SIZE") - - if ub_size_bytes is None or ub_size_bytes <= 0: - raise ValueError(f"Invalid UB_SIZE from get_soc_spec: {ub_size_bytes}") - - # Convert bytes to bits - ub_capacity_bits = ub_size_bytes * 8 - return ub_capacity_bits - - except ImportError: - return _fallback_ub_capacity( - "Cannot import tbe.common.platform.get_soc_spec. " - "Source CANN set_env.sh or set ASCEND_UB_CAPACITY_BITS." - ) - except Exception as e: - return _fallback_ub_capacity( - f"Failed to detect UB capacity from get_soc_spec: {e}. " - "Set ASCEND_UB_CAPACITY_BITS to override." - ) - - return _fallback_ub_capacity( - "NPU is not available. Set ASCEND_UB_CAPACITY_BITS to override." - ) - - -# Global singleton instance -_ub_manager: UBManager | None = None - - -def get_ub_manager() -> UBManager: - """Get global UB manager instance.""" - global _ub_manager - if _ub_manager is None: - _ub_manager = UBManager() - return _ub_manager - - -def compute_default_tiling_strategy( - safety_margin: float = 0.80, - dtype_size: int | None = None, - memory_multiplier: float | None = None, - shapes: tuple[tuple[int, ...], ...] | None = None, - tiling_dims: tuple[int | tuple[int, ...], ...] | None = None, -) -> tuple[tuple[int, ...], ...] | None: - """ - Compute tiling strategy using the default strategy function. - - This function directly calls the default strategy and computes the final - tiling result. All kernels use the same unified strategy function, so - there's no need for kernel_name-based lookup. - - Args: - safety_margin: Safety margin as a float (e.g., 0.80 for 80%). Default is 0.80. - dtype_size: Size of data type in bytes (e.g., 2 for float16, 4 for float32). - Must be provided. If None or <= 0, defaults to 4 (float32). - memory_multiplier: Memory multiplier for estimating peak memory usage. - - For GEGLU: typically 10.0 for backward, 4.0 for forward - - For ROPE: typically 3.0 - If None, defaults to 10.0 (conservative estimate). - shapes: Tuple of full shapes. Each shape is a tuple of dimension sizes. - - For ROPE: ((n_q_head, hd), (n_kv_head, hd)) - - For GEGLU: ((n_cols,),) - Can pass original shapes (will handle padding internally) or padded shapes. - tiling_dims: Tuple specifying which dimensions can be tiled for each shape. - Each element can be: - - int: single dimension index (e.g., 0 for first dimension) - - tuple of ints: multiple dimensions that can be tiled together - - For ROPE: (0, 0) means first dimension of each shape can be tiled - - For GEGLU: (0,) means first dimension of the shape can be tiled - Length must match len(shapes). Cannot be empty. - - Returns: - Tuple of tiled shapes with same structure as input shapes. - Tiling dimensions are replaced with computed block sizes (power of 2), - while non-tiling dimensions are padded to next power of 2. - - For ROPE: ((block_size_q, pad_hd), (block_size_kv, pad_hd)) - - For GEGLU: ((block_size,),) - Returns None if shapes or tiling_dims is None or empty. - - Examples: - >>> # ROPE forward - >>> strategy = compute_default_tiling_strategy( - ... safety_margin=0.90, - ... dtype_size=4, - ... memory_multiplier=3.0, - ... shapes=((32, 128), (32, 128)), - ... tiling_dims=(0, 0) - ... ) - >>> # Returns: ((block_size_q, 128), (block_size_kv, 128)) - >>> # GEGLU forward - >>> strategy = compute_default_tiling_strategy( - ... safety_margin=0.80, - ... dtype_size=2, - ... memory_multiplier=7.0, - ... shapes=((4096,),), - ... tiling_dims=(0,) - ... ) - >>> # Returns: ((block_size,),) - """ - ub_manager = get_ub_manager() - - if shapes is None or not shapes or tiling_dims is None or not tiling_dims: - return None - - if len(shapes) != len(tiling_dims): - return None - - if dtype_size is None or dtype_size <= 0: - dtype_size = 4 # Default to float32 - - if memory_multiplier is None or memory_multiplier <= 0: - memory_multiplier = 10.0 # Default conservative estimate - - # Call strategy to get max_safe_block_size for each shape - max_supported = _default_strategy( - ub_manager.ub_capacity_bits, - safety_margin, - dtype_size, - memory_multiplier, - shapes, - tiling_dims, - ) - - if not max_supported or len(max_supported) != len(shapes): - return None - - # Build result: same structure as shapes, with tiling dims replaced by computed block sizes - result = [] - for shape, tiling_dim, max_safe in zip(shapes, tiling_dims, max_supported): - result_shape = list(shape) - - # Normalize tiling_dim to a set of dimension indices - tiling_dim_set = _normalize_tiling_dims(tiling_dim) - - # Validate tiling dimensions are within shape bounds - if not tiling_dim_set: - raise ValueError( - f"Invalid tiling_dim: {tiling_dim}. tiling_dim must be an int or a non-empty tuple of ints." - ) - if any(dim_idx < 0 or dim_idx >= len(result_shape) for dim_idx in tiling_dim_set): - raise ValueError( - f"Invalid tiling_dim: {tiling_dim} for shape {shape}. " - f"All dimension indices must be in range [0, {len(result_shape)})." - ) - - # Replace tiling dimensions with computed block sizes - # For each tiling dimension, compute: min(desired, max_safe) - for dim_idx in tiling_dim_set: - original_dim = result_shape[dim_idx] - desired = triton.next_power_of_2(original_dim) - final_val = min(desired, max_safe) - final_val = max(1, final_val) # Ensure at least 1 - result_shape[dim_idx] = final_val - - # Pad non-tiling dimensions to next power of 2 - for dim_idx, dim_size in enumerate(result_shape): - if dim_idx not in tiling_dim_set: - result_shape[dim_idx] = triton.next_power_of_2(dim_size) - - result.append(tuple(result_shape)) - - return tuple(result) - - -def compute_ub_block_size( - dim_size: int, - memory_multiplier: float, - *, - safety_margin: float = 0.9, - dtype_size: int = 4, - fallback: int = 2048, - min_block: int = 1, - max_block: int | None = None, - desired: int | None = None, -) -> int: - """Compute UB-safe block size for a single tilable dimension.""" - if desired is None: - desired = triton.next_power_of_2(dim_size) - - tile_shapes = compute_default_tiling_strategy( - safety_margin=safety_margin, - dtype_size=dtype_size, - memory_multiplier=memory_multiplier, - shapes=((dim_size,),), - tiling_dims=(0,), - ) - block = min(desired, tile_shapes[0][0]) if tile_shapes else min(desired, fallback) - block = max(min_block, block) - if max_block is not None: - block = min(block, max_block) - return block - - -def compute_vocab_block_size( - vocab_size: int, - num_rows: int, - memory_multiplier: float, - *, - safety_margin: float = 0.9, - max_block: int = 8192, - fallback: int = 2048, -) -> int: - """UB-safe vocab tile size respecting Ascend grid dim0 limit.""" - ub_block = compute_ub_block_size( - vocab_size, - memory_multiplier, - safety_margin=safety_margin, - fallback=fallback, - ) - max_splits = max(1, ASCEND_MAX_GRID_DIM // max(num_rows, 1)) - grid_min = triton.next_power_of_2((vocab_size + max_splits - 1) // max_splits) - return min(max(ub_block, grid_min), max_block) - - -def compute_elementwise_block_size( - n_elements: int, - memory_multiplier: float = 2.5, - *, - safety_margin: float = 0.9, - min_block: int = 1024, - fallback: int | None = None, -) -> int: - """UB-safe block size for elementwise kernels under grid limit.""" - if fallback is None: - fallback = _FALLBACK_MAX_FUSED_BLOCK - ub_block = compute_ub_block_size( - n_elements, - memory_multiplier, - safety_margin=safety_margin, - fallback=fallback, - ) - grid_min = max( - min_block, - triton.next_power_of_2((n_elements + ASCEND_MAX_GRID_DIM - 1) // ASCEND_MAX_GRID_DIM), - ) - return min(ub_block, grid_min) - - -def compute_activation_block_size( - total_elements: int, - is_backward: bool, - *, - max_grid: int = ASCEND_MAX_GRID_DIM, - max_core_dim: int = 65535, - safety_margin: float = 0.9, - max_block: int = 2048, - memory_multiplier: float | None = None, -) -> int: - """UB-safe block size for flattened activation kernels.""" - if memory_multiplier is None: - memory_multiplier = 6.0 if is_backward else 3.0 - block = compute_ub_block_size( - total_elements, - memory_multiplier, - safety_margin=safety_margin, - fallback=2048, - min_block=256, - max_block=max_block, - ) - block = min(block, max_core_dim // 8) - if triton.cdiv(total_elements, block) > max_grid: - block = triton.cdiv(total_elements, max_grid) - block = max(triton.next_power_of_2(block), 1) - return block - - -def compute_row_tile_block_size( - row_dim: int, - fixed_dim: int, - memory_multiplier: float, - *, - tiling_row: bool = True, - safety_margin: float = 0.85, - dtype_size: int = 4, - fallback: int = 16, - min_block: int = 1, - max_block: int | None = None, -) -> int: - """UB-safe tile along one axis of a 2D [row, col] kernel tile.""" - if tiling_row: - shapes = ((row_dim, fixed_dim),) - tiling_dims = (0,) - desired = triton.next_power_of_2(row_dim) - else: - shapes = ((row_dim, fixed_dim),) - tiling_dims = (1,) - desired = triton.next_power_of_2(fixed_dim) - - tile_shapes = compute_default_tiling_strategy( - safety_margin=safety_margin, - dtype_size=dtype_size, - memory_multiplier=memory_multiplier, - shapes=shapes, - tiling_dims=tiling_dims, - ) - if tile_shapes: - block = tile_shapes[0][0] if tiling_row else tile_shapes[0][1] - block = min(desired, block) - else: - block = min(desired, fallback) - block = max(min_block, block) - if max_block is not None: - block = min(block, max_block) - return block - - -def max_grid_axis_chunks( - axis_size: int, - other_grid_product: int, - *, - max_grid: int = ASCEND_MAX_GRID_DIM, -) -> int: - """Max launch chunks along one grid axis while keeping the product <= max_grid.""" - return max(1, max_grid // max(other_grid_product, 1)) - - -def compute_grid_limited_tile_size( - axis_size: int, - other_grid_product: int, - ub_safe_block: int, - *, - max_grid: int = ASCEND_MAX_GRID_DIM, - min_block: int = 1, -) -> int: - """Pick a UB-safe tile size; host-side chunking handles grid overflow.""" - return max(min_block, min(ub_safe_block, axis_size)) - - -def iter_axis_launch_chunks( - axis_size: int, - other_grid_product: int, - *, - max_grid: int = ASCEND_MAX_GRID_DIM, -): - """Yield ``(offset, chunk_len)`` for host-side grid-axis tiling.""" - max_chunks = max_grid_axis_chunks(axis_size, other_grid_product, max_grid=max_grid) - for offset in range(0, axis_size, max_chunks): - yield offset, min(max_chunks, axis_size - offset) diff --git a/build/torch-xpu/utils/_decorators.py b/build/torch-cuda/utils/decorators.py similarity index 98% rename from build/torch-xpu/utils/_decorators.py rename to build/torch-cuda/utils/decorators.py index f29063d557d8e6cb12f43643fa6964e838c83f4c..304a84c1c69a549717f294da587a904c08927ab8 100644 --- a/build/torch-xpu/utils/_decorators.py +++ b/build/torch-cuda/utils/decorators.py @@ -18,11 +18,9 @@ from typing import Any import torch from packaging import version as package_version -from ._config import FLA_DISABLE_TENSOR_CACHE, FLA_TENSOR_CACHE_SIZE -from ._device import custom_device_ctx - - -__version__ = "0.5.2" # kept for BC +from .. import __version__ +from ..utils.env import FLA_DISABLE_TENSOR_CACHE, FLA_TENSOR_CACHE_SIZE +from ..utils.hardware import custom_device_ctx class Action(Enum): @@ -60,7 +58,7 @@ def tensor_cache( cached: deque = deque(maxlen=FLA_TENSOR_CACHE_SIZE) def cache_disabled() -> bool: - utils_module = sys.modules.get('fla.utils') + utils_module = sys.modules.get(__package__) return getattr(utils_module, 'FLA_DISABLE_TENSOR_CACHE', FLA_DISABLE_TENSOR_CACHE) @functools.wraps(fn) diff --git a/build/torch-cuda/utils/env.py b/build/torch-cuda/utils/env.py new file mode 100644 index 0000000000000000000000000000000000000000..1fe5790884283df53cea48ae22041fe376808b8f --- /dev/null +++ b/build/torch-cuda/utils/env.py @@ -0,0 +1,153 @@ +# 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 functools +import importlib.metadata +import inspect +import logging +import os +import shutil +import sys +from functools import cache, lru_cache +from importlib.util import find_spec +from pathlib import Path + +import torch +import triton +from packaging import version as package_version + +logger = logging.getLogger(__name__) + +FLA_CI_ENV = os.getenv("FLA_CI_ENV") == "1" +FLA_CACHE_RESULTS = os.getenv('FLA_CACHE_RESULTS', '1') == '1' + +FLA_DISABLE_TENSOR_CACHE = os.getenv('FLA_DISABLE_TENSOR_CACHE', '0') == '1' +try: + FLA_TENSOR_CACHE_SIZE = int(os.getenv('FLA_TENSOR_CACHE_SIZE', "4")) +except ValueError: + FLA_TENSOR_CACHE_SIZE = 4 + + +@lru_cache(maxsize=1) +def check_environments(): + """ + Checks the current operating system, Triton version, and Python version, + issuing warnings if they don't meet recommendations. + This function's body only runs once due to lru_cache. + """ + # Check Operating System + if sys.platform == 'win32': + # Check if triton-windows is installed + try: + from importlib.metadata import PackageNotFoundError, metadata + metadata('triton-windows') + # triton-windows is installed, no warning needed + except PackageNotFoundError: + logger.warning( + "Detected Windows operating system. Consider installing triton-windows " + "(https://github.com/triton-lang/triton-windows) for better compatibility. " + "Without it, some features may not work correctly.", + ) + + triton_version = package_version.parse(triton.__version__) + required_triton_version = package_version.parse("3.3.0") + + if triton_version < required_triton_version: + logger.warning( + f"Current Triton version {triton_version} is below the recommended 3.3.0 version. " + "Errors may occur and these issues will not be fixed. " + "Please consider upgrading Triton.", + ) + + # Check Python version + py_version = package_version.parse(f"{sys.version_info.major}.{sys.version_info.minor}") + required_py_version = package_version.parse("3.11") + + if py_version < required_py_version: + logger.warning( + f"Current Python version {py_version} is below the recommended 3.11 version. " + "It is recommended to upgrade to Python 3.11 or higher for the best experience.", + ) + + return None + + +check_environments() + + +@cache +def check_pytorch_version(version_s: str = '2.4') -> bool: + return package_version.parse(torch.__version__) >= package_version.parse(version_s) + + +@cache +def check_transformers_version(version_s: str) -> bool: + try: + installed_version = importlib.metadata.version('transformers') + except importlib.metadata.PackageNotFoundError: + return False + return package_version.parse(installed_version) >= package_version.parse(version_s) + + +@cache +def supports_transformers_next_sequence_length() -> bool: + if find_spec_cached('transformers') is None: + return False + from transformers.generation import GenerationMixin + + return 'next_sequence_length' in inspect.signature(GenerationMixin.prepare_inputs_for_generation).parameters + + +TRITON_ABOVE_3_4_0 = package_version.parse(triton.__version__) >= package_version.parse("3.4.0") +TRITON_ABOVE_3_5_1 = package_version.parse(triton.__version__) >= package_version.parse("3.5.1") +TRITON_ABOVE_3_6_0 = package_version.parse(triton.__version__) >= package_version.parse("3.6.0") +TRITON_ABOVE_3_7_1 = package_version.parse(triton.__version__) >= package_version.parse("3.7.1") +TRITON_ABOVE_3_8_0 = package_version.parse(triton.__version__) >= package_version.parse("3.8.0") + +SUPPORTS_AUTOTUNE_CACHE = "cache_results" in inspect.signature(triton.autotune).parameters +autotune_cache_kwargs = {"cache_results": FLA_CACHE_RESULTS} if SUPPORTS_AUTOTUNE_CACHE else {} + + +@functools.cache +def find_spec_cached(name): + try: + return find_spec(name) + except ModuleNotFoundError as error: + if error.name is not None and (name == error.name or name.startswith(f'{error.name}.')): + return None + raise + + +@functools.cache +def has_usable_nvcc() -> bool: + """Whether a usable nvcc compiler is available for TileLang's JIT. + + Mirrors the guesses in ``tilelang.env._find_cuda_home`` (env + CUDA_HOME/CUDA_PATH, nvcc on PATH, the ``nvidia-cuda-nvcc`` wheel, + /usr/local/cuda), but verifies the nvcc binary actually exists — + only ``nvidia-cuda-nvcc`` >= 13.0 ships it, the ``-cu12`` variant + installs just ptxas. + """ + cuda_home = os.environ.get("CUDA_HOME") or os.environ.get("CUDA_PATH") + if cuda_home is not None and (Path(cuda_home) / "bin" / "nvcc").exists(): + return True + if shutil.which("nvcc") is not None: + return True + try: + files = importlib.metadata.files("nvidia-cuda-nvcc") or [] + except importlib.metadata.PackageNotFoundError: + files = [] + if any(f.name in ("nvcc", "nvcc.exe") for f in files): + return True + if (Path("/usr/local/cuda") / "bin" / "nvcc").exists(): + return True + + logger.info( + "[FLA Backend] TileLang is installed but no usable nvcc compiler was found; falling back to Triton. " + "Install a CUDA toolkit or nvidia-cuda-nvcc, or set FLA_TILELANG=0 to disable TileLang explicitly." + ) + return False diff --git a/build/torch-cuda/utils/_device.py b/build/torch-cuda/utils/hardware.py similarity index 68% rename from build/torch-cuda/utils/_device.py rename to build/torch-cuda/utils/hardware.py index 0064c1939421d3831978e0210a63afa94a9fb36a..3979c814bb8d71efe05f96406293542ffcf1893e 100644 --- a/build/torch-cuda/utils/_device.py +++ b/build/torch-cuda/utils/hardware.py @@ -10,63 +10,16 @@ import functools import logging import os import platform -import sys import warnings from enum import Enum -from functools import cache, lru_cache +from functools import cache import torch import triton -from packaging import version as package_version - -logger = logging.getLogger(__name__) - - -@lru_cache(maxsize=1) -def check_environments(): - """ - Checks the current operating system, Triton version, and Python version, - issuing warnings if they don't meet recommendations. - This function's body only runs once due to lru_cache. - """ - # Check Operating System - if sys.platform == 'win32': - # Check if triton-windows is installed - try: - from importlib.metadata import PackageNotFoundError, metadata - metadata('triton-windows') - # triton-windows is installed, no warning needed - except PackageNotFoundError: - logger.warning( - "Detected Windows operating system. Consider installing triton-windows " - "(https://github.com/triton-lang/triton-windows) for better compatibility. " - "Without it, some features may not work correctly.", - ) - - triton_version = package_version.parse(triton.__version__) - required_triton_version = package_version.parse("3.3.0") - - if triton_version < required_triton_version: - logger.warning( - f"Current Triton version {triton_version} is below the recommended 3.3.0 version. " - "Errors may occur and these issues will not be fixed. " - "Please consider upgrading Triton.", - ) - - # Check Python version - py_version = package_version.parse(f"{sys.version_info.major}.{sys.version_info.minor}") - required_py_version = package_version.parse("3.11") - - if py_version < required_py_version: - logger.warning( - f"Current Python version {py_version} is below the recommended 3.11 version. " - "It is recommended to upgrade to Python 3.11 or higher for the best experience.", - ) - - return None +from ..utils.env import check_pytorch_version -check_environments() +logger = logging.getLogger(__name__) def _cpu_device_warning(): @@ -74,24 +27,37 @@ def _cpu_device_warning(): @cache -def check_pytorch_version(version_s: str = '2.4') -> bool: - return package_version.parse(torch.__version__) >= package_version.parse(version_s) - - -@cache -def get_multiprocessor_count(tensor_idx: int = 0) -> int: +def _get_multiprocessor_count(tensor_idx: int = 0, *, use_aicore: bool = False) -> int: try: return triton.runtime.driver.active.utils.get_device_properties(tensor_idx)['multiprocessor_count'] except Exception: # Maybe we use a NPU device. try: if triton.runtime.driver.active.get_current_target().backend == 'npu': - return triton.runtime.driver.active.utils.get_device_properties(tensor_idx)['num_vectorcore'] + props = triton.runtime.driver.active.utils.get_device_properties(tensor_idx) + return props['num_aicore'] if use_aicore else props['num_vectorcore'] except Exception: logger.debug('Failed to get NPU multiprocessor count, falling back to 1.', exc_info=True) return 1 +@torch.compiler.assume_constant_result +def get_multiprocessor_count(tensor_idx: int = 0, *, use_aicore: bool = False) -> int: + return _get_multiprocessor_count(tensor_idx, use_aicore=use_aicore) + + +@cache +def get_device_capability(device_index: int = 0) -> tuple[int, int]: + major, minor = torch.cuda.get_device_capability(device_index) + return int(major), int(minor) + + +@cache +def get_device_smem_optin(device_index: int = 0) -> int: + props = torch.cuda.get_device_properties(device_index) + return int(getattr(props, 'shared_memory_per_block_optin', props.shared_memory_per_block)) + + @cache def get_available_device() -> str: try: @@ -101,6 +67,16 @@ def get_available_device() -> str: return 'cpu' +@cache +def get_device_arch() -> str: + """Triton target arch, e.g. 'gfx1250' on AMD or '90' on Nvidia. Empty when unavailable.""" + try: + # AMD may append feature flags, e.g. 'gfx1250:xnack-'. + return str(triton.runtime.driver.active.get_current_target().arch).split(':')[0] + except Exception: + return '' + + def map_triton_backend_to_torch_device() -> str: backend = get_available_device() # 'cuda' | 'hip' | 'xpu' | 'cpu' | ... return {'cuda': 'cuda', 'hip': 'cuda', 'xpu': 'xpu'}.get(backend, backend) @@ -115,22 +91,43 @@ device_platform = get_available_device() device_name = map_triton_backend_to_torch_device() IS_AMD = (device_platform == 'hip') + IS_ARM = platform.machine().lower() in ('aarch64', 'arm64') + IS_INTEL = (device_platform == 'xpu') IS_INTEL_ALCHEMIST = (IS_INTEL and 'Intel(R) Arc(TM) A' in torch.xpu.get_device_name(0)) -IS_NVIDIA = (device_platform == 'cuda') + IS_NPU = (device_platform == 'npu') + +IS_NVIDIA = (device_platform == 'cuda') +IS_NVIDIA_HOPPER = ( + IS_NVIDIA and ( + 'NVIDIA H' in torch.cuda.get_device_name(0) + or torch.cuda.get_device_capability()[0] == 9 + ) +) +IS_NVIDIA_SM100 = (IS_NVIDIA and torch.cuda.get_device_capability()[0] == 10) +# NOTE: exactly 12.0 — 12.1 (GB10) is a different target that FlashQLA rejects at import time. +IS_NVIDIA_SM120 = (IS_NVIDIA and torch.cuda.get_device_capability() == (12, 0)) IS_NVIDIA_BLACKWELL = (IS_NVIDIA and torch.cuda.get_device_capability()[0] in (10, 12)) -IS_NVIDIA_HOPPER = (IS_NVIDIA and ('NVIDIA H' in torch.cuda.get_device_name(0) or torch.cuda.get_device_capability()[0] == 9)) # Nvidia Ampere or newer, haven't check AMD and intel yet. IS_TF32_SUPPORTED = (IS_NVIDIA and torch.cuda.get_device_capability(0)[0] >= 8) IS_GATHER_SUPPORTED = hasattr(triton.language, 'gather') +# AMD architectures whose Triton backend lowers `tl.make_tensor_descriptor` onto real +# hardware tensor-descriptor loads/stores. gfx1250 is the first one; gfx942/gfx950 are not. +IS_AMD_TMA_ARCH = (IS_AMD and get_device_arch() in ('gfx1250',)) + IS_TMA_SUPPORTED = ( - IS_NVIDIA - and torch.cuda.get_device_capability(0)[0] >= 9 + ( + (IS_NVIDIA and torch.cuda.get_device_capability(0)[0] >= 9) + or IS_AMD_TMA_ARCH + ) and os.environ.get('FLA_USE_TMA', '0') == '1' - and (hasattr(triton.language, '_experimental_make_tensor_descriptor') or hasattr(triton.language, 'make_tensor_descriptor')) + and ( + hasattr(triton.language, '_experimental_make_tensor_descriptor') + or hasattr(triton.language, 'make_tensor_descriptor') + ) ) if IS_NVIDIA and not IS_TF32_SUPPORTED: @@ -180,7 +177,7 @@ class Backend(Enum): @cache -def check_shared_mem(arch: str = "none", tensor_idx: int = 0) -> bool: +def _check_shared_mem(arch: str = "none", tensor_idx: int = 0) -> bool: try: device_shared_mem_list = get_all_max_shared_mem() max_shared_memory = device_shared_mem_list[tensor_idx] @@ -189,6 +186,11 @@ def check_shared_mem(arch: str = "none", tensor_idx: int = 0) -> bool: return False +@torch.compiler.assume_constant_result +def check_shared_mem(arch: str = "none", tensor_idx: int = 0) -> bool: + return _check_shared_mem(arch, tensor_idx) + + if check_pytorch_version('2.4'): if device == 'cpu': device = 'cuda' @@ -215,3 +217,4 @@ else: return torch.cuda.device(index) except (AttributeError, AssertionError, RuntimeError): return contextlib.nullcontext() + diff --git a/build/torch-xpu/utils/_testing.py b/build/torch-cuda/utils/testing.py similarity index 97% rename from build/torch-xpu/utils/_testing.py rename to build/torch-cuda/utils/testing.py index 12064dcc4d48c88a89856afac93db66aceaeaa79..d44308495d028e3e13cc32d31451e54793137a11 100644 --- a/build/torch-xpu/utils/_testing.py +++ b/build/torch-cuda/utils/testing.py @@ -10,7 +10,7 @@ import warnings import torch -from ._config import FLA_CI_ENV +from ..utils.env import FLA_CI_ENV logger = logging.getLogger(__name__) diff --git a/build/torch-xpu/__init__.py b/build/torch-xpu/__init__.py deleted file mode 100644 index 93947cdf3ec9f3c35772e484b1160934eeda5ade..0000000000000000000000000000000000000000 --- a/build/torch-xpu/__init__.py +++ /dev/null @@ -1,12 +0,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 - - -__all__ = [ - # original functions directly - "chunk_kda", "fused_recurrent_kda", - "chunk_gated_delta_rule", "fused_recurrent_gated_delta_rule", - # layers - "layers", -] diff --git a/build/torch-xpu/_ops.py b/build/torch-xpu/_ops.py deleted file mode 100644 index e04a067e52d184bc6a7b05c7a25feb67196cc25b..0000000000000000000000000000000000000000 --- a/build/torch-xpu/_ops.py +++ /dev/null @@ -1,41 +0,0 @@ -import torch - -def get_backend() -> str: - """Detect the backend by inspecting torch.""" - import torch - - acc = torch.accelerator.current_accelerator() if hasattr(torch, "accelerator") else None - if acc is not None and acc.type == "tpu": - return "tpu" - elif hasattr(torch, "neuron"): - # Needs to be sorted before specific Torch builds, since Neuron - # extension can be loaded into e.g. CUDA Torch builds. - return "neuron" - elif torch.version.cuda is not None: - return "cuda" - elif torch.version.hip is not None: - return "rocm" - elif torch.backends.mps.is_available(): - return "metal" - elif hasattr(torch.version, "xpu") and torch.version.xpu is not None: - return "xpu" - else: - return "cpu" - - -def _find_ops_name() -> str: - kernel_name = "fla" - unique_id = "fe2bb66" - backend = get_backend() - return f"_{kernel_name}_{backend}_{unique_id}" - - -_OPS_NAME = _find_ops_name() - -ops = getattr(torch.ops, _OPS_NAME) - -def add_op_namespace_prefix(op_name: str) -> str: - """ - Prefix op by namespace. - """ - return f"{_OPS_NAME}::{op_name}" \ No newline at end of file diff --git a/build/torch-xpu/layers.py b/build/torch-xpu/layers.py deleted file mode 100644 index 42cb865a5dd7150f2c543c009206dc23dea1d9ae..0000000000000000000000000000000000000000 --- a/build/torch-xpu/layers.py +++ /dev/null @@ -1,136 +0,0 @@ -import torch.nn as nn - -from .modules.fused_norm_gate import rms_norm_gated -from .ops.gated_delta_rule import chunk_gated_delta_rule as chunk_gdn -from .ops.gated_delta_rule import fused_recurrent_gated_delta_rule as fused_recurrent_gdn -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, - gate, - self.weight, - None, # bias - self.activation, - residual=None, - eps=self.variance_epsilon, - prenorm=False, - residual_in_fp32=False, - ) - - -class chunk_kimi_delta_attention(nn.Module): - def forward( - self, - query, - key, - value, - g, - beta, - chunk_size=64, - initial_state=None, - output_final_state=False, - use_qk_l2norm_in_kernel=False, - **kwargs, - ): - # Keep internal consistency between transformers and fla to allow both - cu_seqlens = kwargs.pop("cu_seq_lens_q", kwargs.pop("cu_seqlens", None)) - - return chunk_kda( - query, - key, - value, - g=g, - beta=beta, - chunk_size=chunk_size, - initial_state=initial_state, - output_final_state=output_final_state, - use_qk_l2norm_in_kernel=use_qk_l2norm_in_kernel, - cu_seqlens=cu_seqlens, - **kwargs, - ) - - -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, - ): - # Keep internal consistency between transformers and fla to allow both - cu_seqlens = kwargs.pop("cu_seq_lens_q", kwargs.pop("cu_seqlens", None)) - - return fused_recurrent_kda( - query, - key, - value, - g=g, - beta=beta, - initial_state=initial_state, - output_final_state=output_final_state, - use_qk_l2norm_in_kernel=use_qk_l2norm_in_kernel, - cu_seqlens=cu_seqlens, - **kwargs, - ) - - -class chunk_gated_delta_rule(nn.Module): - def forward( - self, - query, - key, - value, - g, - beta, - chunk_size=64, - initial_state=None, - output_final_state=False, - use_qk_l2norm_in_kernel=False, - **kwargs, - ): - # Keep internal consistency between transformers and fla to allow both - cu_seqlens = kwargs.pop("cu_seq_lens_q", kwargs.pop("cu_seqlens", None)) - - return chunk_gdn( - query, - key, - value, - g=g, - beta=beta, - chunk_size=chunk_size, - initial_state=initial_state, - output_final_state=output_final_state, - use_qk_l2norm_in_kernel=use_qk_l2norm_in_kernel, - cu_seqlens=cu_seqlens, - **kwargs, - ) - - -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, - ): - # Keep internal consistency between transformers and fla to allow both - cu_seqlens = kwargs.pop("cu_seq_lens_q", kwargs.pop("cu_seqlens", None)) - - return fused_recurrent_gdn( - query, - key, - value, - g=g, - beta=beta, - initial_state=initial_state, - output_final_state=output_final_state, - use_qk_l2norm_in_kernel=use_qk_l2norm_in_kernel, - cu_seqlens=cu_seqlens, - **kwargs, - ) - - -__all__ = [ - "FusedRMSNormGated", - "chunk_kimi_delta_attention", - "recurrent_kimi_delta_attention", - "chunk_gated_delta_rule", - "recurrent_gated_delta_rule" -] diff --git a/build/torch-xpu/metadata.json b/build/torch-xpu/metadata.json deleted file mode 100644 index de301e89dccccf941a44196645d6bfa1c569372d..0000000000000000000000000000000000000000 --- a/build/torch-xpu/metadata.json +++ /dev/null @@ -1,321 +0,0 @@ -{ - "name": "fla", - "id": "_fla_xpu_fe2bb66", - "version": 1, - "kernels-minver": "0.14.0", - "license": "MIT", - "python-depends": [ - "einops" - ], - "kernel-depends": [], - "backend": { - "type": "xpu" - }, - "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": "mLYlaOKZomfTLhl4BVcyn4O5q2XY2OfxWDiw6nZrSrU=", - "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": { - "kernel-builder": { - "version": "0.18.0-dev0", - "commit": "d054906a13710b811b96697b4635f03020842733", - "dirty": false - }, - "kernel": { - "commit": "fe2bb66fee6c2a77eccf08a0508dc57428d488ed", - "dirty": false - } - } -} \ No newline at end of file diff --git a/build/torch-xpu/metadata.json.sigstore b/build/torch-xpu/metadata.json.sigstore deleted file mode 100644 index 54862867196284fec57ea7d3f20e42b6071a8007..0000000000000000000000000000000000000000 --- a/build/torch-xpu/metadata.json.sigstore +++ /dev/null @@ -1 +0,0 @@ -{"mediaType":"application/vnd.dev.sigstore.bundle.v0.3+json", "verificationMaterial":{"certificate":{"rawBytes":"MIIHLTCCBrOgAwIBAgIUccKqLwRVOhn7FTeoom+7JBUxGR4wCgYIKoZIzj0EAwMwNzEVMBMGA1UEChMMc2lnc3RvcmUuZGV2MR4wHAYDVQQDExVzaWdzdG9yZS1pbnRlcm1lZGlhdGUwHhcNMjYxMDE5MTEwMDU2WhcNMjYxMDE5MTExMDU2WjAAMFkwEwYHKoZIzj0CAQYIKoZIzj0DAQcDQgAEcf0Zu+ccQZsFDzoPKdqMiKLK1d5n8//7KDoRGgTN9mzwlldvwE//TLmPf/BmuPtrXY60QrCR1D0tkMtpbG4LRKOCBdIwggXOMA4GA1UdDwEB/wQEAwIHgDATBgNVHSUEDDAKBggrBgEFBQcDAzAfBgNVHSMEGDAWgBTf0+nPViQRlvmo2OkoVaLGLhhkPzBrBgNVHREBAf8EYTBfhl1odHRwczovL2dpdGh1Yi5jb20vaHVnZ2luZ2ZhY2Uva2VybmVscy1jb21tdW5pdHkvLmdpdGh1Yi93b3JrZmxvd3MvYnVpbGQueWFtbEByZWZzL2hlYWRzL21haW4wOQYKKwYBBAGDvzABAQQraHR0cHM6Ly90b2tlbi5hY3Rpb25zLmdpdGh1YnVzZXJjb250ZW50LmNvbTAfBgorBgEEAYO/MAECBBF3b3JrZmxvd19kaXNwYXRjaDA2BgorBgEEAYO/MAEDBChmZTJiYjY2ZmVlNmMyYTc3ZWNjZjA4YTA1MDhkYzU3NDI4ZDQ4OGVkMBMGCisGAQQBg78wAQQEBUJ1aWxkMCsGCisGAQQBg78wAQUEHWh1Z2dpbmdmYWNlL2tlcm5lbHMtY29tbXVuaXR5MB0GCisGAQQBg78wAQYED3JlZnMvaGVhZHMvbWFpbjA7BgorBgEEAYO/MAEIBC0MK2h0dHBzOi8vdG9rZW4uYWN0aW9ucy5naXRodWJ1c2VyY29udGVudC5jb20wbQYKKwYBBAGDvzABCQRfDF1odHRwczovL2dpdGh1Yi5jb20vaHVnZ2luZ2ZhY2Uva2VybmVscy1jb21tdW5pdHkvLmdpdGh1Yi93b3JrZmxvd3MvYnVpbGQueWFtbEByZWZzL2hlYWRzL21haW4wOAYKKwYBBAGDvzABCgQqDChmZTJiYjY2ZmVlNmMyYTc3ZWNjZjA4YTA1MDhkYzU3NDI4ZDQ4OGVkMBsGCisGAQQBg78wAQsEDQwLc2VsZi1ob3N0ZWQwQAYKKwYBBAGDvzABDAQyDDBodHRwczovL2dpdGh1Yi5jb20vaHVnZ2luZ2ZhY2Uva2VybmVscy1jb21tdW5pdHkwOAYKKwYBBAGDvzABDQQqDChmZTJiYjY2ZmVlNmMyYTc3ZWNjZjA4YTA1MDhkYzU3NDI4ZDQ4OGVkMB8GCisGAQQBg78wAQ4EEQwPcmVmcy9oZWFkcy9tYWluMBoGCisGAQQBg78wAQ8EDAwKMTA3MTQ3NTUyOTAuBgorBgEEAYO/MAEQBCAMHmh0dHBzOi8vZ2l0aHViLmNvbS9odWdnaW5nZmFjZTAYBgorBgEEAYO/MAERBAoMCDI1NzIwNzQzMG0GCisGAQQBg78wARIEXwxdaHR0cHM6Ly9naXRodWIuY29tL2h1Z2dpbmdmYWNlL2tlcm5lbHMtY29tbXVuaXR5Ly5naXRodWIvd29ya2Zsb3dzL2J1aWxkLnlhbWxAcmVmcy9oZWFkcy9tYWluMDgGCisGAQQBg78wARMEKgwoZmUyYmI2NmZlZTZjMmE3N2VjY2YwOGEwNTA4ZGM1NzQyOGQ0ODhlZDAhBgorBgEEAYO/MAEUBBMMEXdvcmtmbG93X2Rpc3BhdGNoMGQGCisGAQQBg78wARUEVgxUaHR0cHM6Ly9naXRodWIuY29tL2h1Z2dpbmdmYWNlL2tlcm5lbHMtY29tbXVuaXR5L2FjdGlvbnMvcnVucy8zOTIzOTU0OTYzNS9hdHRlbXB0cy8xMBYGCisGAQQBg78wARYECAwGcHVibGljMEYGCisGAQQBg78wARgEOAw2cmVwbzpodWdnaW5nZmFjZS9rZXJuZWxzLWNvbW11bml0eTpyZWY6cmVmcy9oZWFkcy9tYWluMIGLBgorBgEEAdZ5AgQCBH0EewB5AHcA3T0wasbHETJjGR4cmWc3AqJKXrjePK3/h4pygC8p7o4AAAGhU9JYiwAABAMASDBGAiEAqpB2O914Vim67H+8TKHmD5w80d5/SwUBTo0qW+hpz04CIQCdLOoKdYZoyG9aM1grLsvH/kvCr+cuPYy7Oav8Mj7qNTAKBggqhkjOPQQDAwNoADBlAjEAudbzki4fUOODSHG9WPt+8Luz2Ekrng8x6RaXnNSQLw+4r7o3l7bbFtF+QqHValG9AjAaIrNSVz/rn6aVoaQ8YS53NRvUn8AoLswkHSZSVSNJSgVcFk7+zfyQGUo/sGk2UrU="}, "tlogEntries":[{"logIndex":"3320106863", "logId":{"keyId":"wNI9atQGlz+VWfO6LRygH4QUfY/8W4RFwiT5i5WRgB0="}, "kindVersion":{"kind":"hashedrekord", "version":"0.0.1"}, "integratedTime":"1792407656", "inclusionPromise":{"signedEntryTimestamp":"MEUCIBIXVACpzZk3A5UdTHJGNzY1+BEe8izv6WytQRoNurdRAiEAwiN6I127S9wYrzzhAp1KlLnJg6t79oYM7xeX1ljOCX8="}, "inclusionProof":{"logIndex":"3198202601", "rootHash":"7NFr/lK6Bzd15e0zU7FG33+UqCjS81C8Wqgavc/42PY=", "treeSize":"3198202606", "hashes":["Jqmf6hBm2SlbxPt6bWzI3yTDaXWJ9wpwn4oFUPfSzC4=", "Hu84dUxAvrcEUjbpHGVTa0v3W9f4bcLgyxJ9arIsA4k=", "iwjbHz8DebR+crQAEZT8SrEV6EY2d6vr38RBJMuRdI8=", "aSMGqrar0XadNbUv88cEFHH4z041bKA8kIwNb0/4dvw=", "da3b5yFUv0WI9AHoAIzLrj+mneing+E7+sN84uV6dyU=", "AIg+9CDrUOfKYHkholj8Hb5rLzOEXDmzH8bvEzlWT5A=", "9XILTHMRbuTH5s7/ktpnEOF5r3lYOe799pW7DVDKXb0=", "A9FtVzn33g4x+CTV70Ba6UZyAYH9uzW+6x2dh0lNvgI=", "fc0pcvawZCdBMV+HYOr/hQ5DyDVH7XSj5bkI9of0GLM=", "FEgdG4+ki96DDLGD8Iuzoy3Ym45SfhKnrkhNx+u0ElY=", "OlLDeTiZ1vIMDel4vGDtbNVnrFnzDe1MFV5rBN5Ep4o=", "uCZQe757l2sS1c05Y7TzOwEm+X8MeGLIUCkEOmFVOz4=", "nyQV/+cWQUG5Hb2PKjQQNkQIJmz2qzAZKBQlkHGIknM=", "G+nk7mgE5h6dUmv2UY7NDJrY3WvrZY31raaO1AETt+o=", "qlTWNDzERnTBQVcHHXXVPjiUP0SsikGxfcxVlgqVFDo=", "gg9GNApNke3+zNhl4tsWtmIycCFrOMCB3lys6XuSQNQ=", "NlOclldmKPA1tMKIskqHAwuQDYHp31YwpK+OOMrMccc=", "qxzHanAzz57SDdmJe0B7bJK72NTIBbwEMGBBKvDOROw=", "xH/DCseLHr9eKoYT8qsORZK7zVdEGYWHuVtsVrD95wY="], "checkpoint":{"envelope":"rekor.sigstore.dev - 1193050959916656506\n3198202606\n7NFr/lK6Bzd15e0zU7FG33+UqCjS81C8Wqgavc/42PY=\n\n— rekor.sigstore.dev wNI9ajBFAiEAnIN6aO06lOxaFRKVVPyZZ+iY3yKFD7FjVuY2RcYj9SoCIBQUN/Yq8vPM0QLSjxPAH1I37N99dRXRjLZ48pJyi46Z\n"}}, "canonicalizedBody":"eyJhcGlWZXJzaW9uIjoiMC4wLjEiLCJraW5kIjoiaGFzaGVkcmVrb3JkIiwic3BlYyI6eyJkYXRhIjp7Imhhc2giOnsiYWxnb3JpdGhtIjoic2hhMjU2IiwidmFsdWUiOiJiZDFlMGE3NWQxNjhiZDJiMWNmOWM3NzQwNTI1OGQyMmZiOWJiMjgzNDRmYjUwZGJiNzlmMGU0NDNiYTZjZDlkIn19LCJzaWduYXR1cmUiOnsiY29udGVudCI6Ik1FWUNJUURzUERCOE9MMWd2bnpTeUY2VHpKMllqQW9YSFJkcjhnUnZOaVd2V25xNklBSWhBS3pkQTJuMjlqK3MyK2JNd1RoUEIwMHRJb29BOWdRczNFRVo2b3N1WHkvWiIsInB1YmxpY0tleSI6eyJjb250ZW50IjoiTFMwdExTMUNSVWRKVGlCRFJWSlVTVVpKUTBGVVJTMHRMUzB0Q2sxSlNVaE1WRU5EUW5KUFowRjNTVUpCWjBsVlkyTkxjVXgzVWxaUGFHNDNSbFJsYjI5dEt6ZEtRbFY0UjFJMGQwTm5XVWxMYjFwSmVtb3dSVUYzVFhjS1RucEZWazFDVFVkQk1WVkZRMmhOVFdNeWJHNWpNMUoyWTIxVmRWcEhWakpOVWpSM1NFRlpSRlpSVVVSRmVGWjZZVmRrZW1SSE9YbGFVekZ3WW01U2JBcGpiVEZzV2tkc2FHUkhWWGRJYUdOT1RXcFplRTFFUlRWTlZFVjNUVVJWTWxkb1kwNU5hbGw0VFVSRk5VMVVSWGhOUkZVeVYycEJRVTFHYTNkRmQxbElDa3R2V2tsNmFqQkRRVkZaU1V0dldrbDZhakJFUVZGalJGRm5RVVZqWmpCYWRTdGpZMUZhYzBaRWVtOVFTMlJ4VFdsTFRFc3haRFZ1T0M4dk4wdEViMUlLUjJkVVRqbHRlbmRzYkdSMmQwVXZMMVJNYlZCbUwwSnRkVkIwY2xoWk5qQlJja05TTVVRd2RHdE5kSEJpUnpSTVVrdFBRMEprU1hkbloxaFBUVUUwUndwQk1WVmtSSGRGUWk5M1VVVkJkMGxJWjBSQlZFSm5UbFpJVTFWRlJFUkJTMEpuWjNKQ1owVkdRbEZqUkVGNlFXWkNaMDVXU0ZOTlJVZEVRVmRuUWxSbUNqQXJibEJXYVZGU2JIWnRiekpQYTI5V1lVeEhUR2hvYTFCNlFuSkNaMDVXU0ZKRlFrRm1PRVZaVkVKbWFHd3hiMlJJVW5kamVtOTJUREprY0dSSGFERUtXV2sxYW1JeU1IWmhTRlp1V2pKc2RWb3lXbWhaTWxWMllUSldlV0p0Vm5OamVURnFZakl4ZEdSWE5YQmtTR3QyVEcxa2NHUkhhREZaYVRrellqTktjZ3BhYlhoMlpETk5kbGx1Vm5CaVIxRjFaVmRHZEdKRlFubGFWMXA2VERKb2JGbFhVbnBNTWpGb1lWYzBkMDlSV1V0TGQxbENRa0ZIUkhaNlFVSkJVVkZ5Q21GSVVqQmpTRTAyVEhrNU1HSXlkR3hpYVRWb1dUTlNjR0l5TlhwTWJXUndaRWRvTVZsdVZucGFXRXBxWWpJMU1GcFhOVEJNYlU1MllsUkJaa0puYjNJS1FtZEZSVUZaVHk5TlFVVkRRa0pHTTJJelNuSmFiWGgyWkRFNWEyRllUbmRaV0ZKcVlVUkJNa0puYjNKQ1owVkZRVmxQTDAxQlJVUkNRMmh0V2xSS2FRcFphbGt5V20xV2JFNXRUWGxaVkdNeldsZE9hbHBxUVRSWlZFRXhUVVJvYTFsNlZUTk9SRWswV2tSUk5FOUhWbXROUWsxSFEybHpSMEZSVVVKbk56aDNDa0ZSVVVWQ1ZVb3hZVmQ0YTAxRGMwZERhWE5IUVZGUlFtYzNPSGRCVVZWRlNGZG9NVm95WkhCaWJXUnRXVmRPYkV3eWRHeGpiVFZzWWtoTmRGa3lPWFFLWWxoV2RXRllValZOUWpCSFEybHpSMEZSVVVKbk56aDNRVkZaUlVRelNteGFiazEyWVVkV2FGcElUWFppVjBad1ltcEJOMEpuYjNKQ1owVkZRVmxQTHdwTlFVVkpRa013VFVzeWFEQmtTRUo2VDJrNGRtUkhPWEphVnpSMVdWZE9NR0ZYT1hWamVUVnVZVmhTYjJSWFNqRmpNbFo1V1RJNWRXUkhWblZrUXpWcUNtSXlNSGRpVVZsTFMzZFpRa0pCUjBSMmVrRkNRMUZTWmtSR01XOWtTRkozWTNwdmRrd3laSEJrUjJneFdXazFhbUl5TUhaaFNGWnVXakpzZFZveVdtZ0tXVEpWZG1FeVZubGliVlp6WTNreGFtSXlNWFJrVnpWd1pFaHJka3h0WkhCa1IyZ3hXV2s1TTJJelNuSmFiWGgyWkROTmRsbHVWbkJpUjFGMVpWZEdkQXBpUlVKNVdsZGFla3d5YUd4WlYxSjZUREl4YUdGWE5IZFBRVmxMUzNkWlFrSkJSMFIyZWtGQ1EyZFJjVVJEYUcxYVZFcHBXV3BaTWxwdFZteE9iVTE1Q2xsVVl6TmFWMDVxV21wQk5GbFVRVEZOUkdocldYcFZNMDVFU1RSYVJGRTBUMGRXYTAxQ2MwZERhWE5IUVZGUlFtYzNPSGRCVVhORlJGRjNUR015Vm5NS1dta3hiMkl6VGpCYVYxRjNVVUZaUzB0M1dVSkNRVWRFZG5wQlFrUkJVWGxFUkVKdlpFaFNkMk42YjNaTU1tUndaRWRvTVZscE5XcGlNakIyWVVoV2JncGFNbXgxV2pKYWFGa3lWWFpoTWxaNVltMVdjMk41TVdwaU1qRjBaRmMxY0dSSWEzZFBRVmxMUzNkWlFrSkJSMFIyZWtGQ1JGRlJjVVJEYUcxYVZFcHBDbGxxV1RKYWJWWnNUbTFOZVZsVVl6TmFWMDVxV21wQk5GbFVRVEZOUkdocldYcFZNMDVFU1RSYVJGRTBUMGRXYTAxQ09FZERhWE5IUVZGUlFtYzNPSGNLUVZFMFJVVlJkMUJqYlZadFkzazViMXBYUm10amVUbDBXVmRzZFUxQ2IwZERhWE5IUVZGUlFtYzNPSGRCVVRoRlJFRjNTMDFVUVROTlZGRXpUbFJWZVFwUFZFRjFRbWR2Y2tKblJVVkJXVTh2VFVGRlVVSkRRVTFJYldnd1pFaENlazlwT0haYU1td3dZVWhXYVV4dFRuWmlVemx2WkZka2JtRlhOVzVhYlVacUNscFVRVmxDWjI5eVFtZEZSVUZaVHk5TlFVVlNRa0Z2VFVORVNURk9la2wzVG5wUmVrMUhNRWREYVhOSFFWRlJRbWMzT0hkQlVrbEZXSGQ0WkdGSVVqQUtZMGhOTmt4NU9XNWhXRkp2WkZkSmRWa3lPWFJNTW1neFdqSmtjR0p0WkcxWlYwNXNUREowYkdOdE5XeGlTRTEwV1RJNWRHSllWblZoV0ZJMVRIazFiZ3BoV0ZKdlpGZEpkbVF5T1hsaE1scHpZak5rZWt3eVNqRmhWM2hyVEc1c2FHSlhlRUZqYlZadFkzazViMXBYUm10amVUbDBXVmRzZFUxRVowZERhWE5IQ2tGUlVVSm5OemgzUVZKTlJVdG5kMjlhYlZWNVdXMUpNazV0V214YVZGcHFUVzFGTTA0eVZtcFpNbGwzVDBkRmQwNVVRVFJhUjAweFRucFJlVTlIVVRBS1QwUm9iRnBFUVdoQ1oyOXlRbWRGUlVGWlR5OU5RVVZWUWtKTlRVVllaSFpqYlhSdFlrYzVNMWd5VW5Cak0wSm9aRWRPYjAxSFVVZERhWE5IUVZGUlFncG5OemgzUVZKVlJWWm5lRlZoU0ZJd1kwaE5Oa3g1T1c1aFdGSnZaRmRKZFZreU9YUk1NbWd4V2pKa2NHSnRaRzFaVjA1c1RESjBiR050Tld4aVNFMTBDbGt5T1hSaVdGWjFZVmhTTlV3eVJtcGtSMngyWW01TmRtTnVWblZqZVRoNlQxUkplazlVVlRCUFZGbDZUbE01YUdSSVVteGlXRUl3WTNrNGVFMUNXVWNLUTJselIwRlJVVUpuTnpoM1FWSlpSVU5CZDBkalNGWnBZa2RzYWsxRldVZERhWE5IUVZGUlFtYzNPSGRCVW1kRlQwRjNNbU50Vm5kaWVuQnZaRmRrYmdwaFZ6VnVXbTFHYWxwVE9YSmFXRXAxV2xkNGVreFhUblppVnpFeFltMXNNR1ZVY0hsYVYxazJZMjFXYldONU9XOWFWMFpyWTNrNWRGbFhiSFZOU1VkTUNrSm5iM0pDWjBWRlFXUmFOVUZuVVVOQ1NEQkZaWGRDTlVGSVkwRXpWREIzWVhOaVNFVlVTbXBIVWpSamJWZGpNMEZ4U2t0WWNtcGxVRXN6TDJnMGNIa0taME00Y0Rkdk5FRkJRVWRvVlRsS1dXbDNRVUZDUVUxQlUwUkNSMEZwUlVGeGNFSXlUemt4TkZacGJUWTNTQ3M0VkV0SWJVUTFkemd3WkRVdlUzZFZRZ3BVYnpCeFZ5dG9jSG93TkVOSlVVTmtURTl2UzJSWldtOTVSemxoVFRGbmNreHpka2d2YTNaRGNpdGpkVkJaZVRkUFlYWTRUV28zY1U1VVFVdENaMmR4Q21ocmFrOVFVVkZFUVhkT2IwRkVRbXhCYWtWQmRXUmllbXRwTkdaVlQwOUVVMGhIT1ZkUWRDczRUSFY2TWtWcmNtNW5PSGcyVW1GWWJrNVRVVXgzS3pRS2NqZHZNMnczWW1KR2RFWXJVWEZJVm1Gc1J6bEJha0ZoU1hKT1UxWjZMM0p1Tm1GV2IyRlJPRmxUTlROT1VuWlZiamhCYjB4emQydElVMXBUVmxOT1NncFRaMVpqUm1zM0szcG1lVkZIVlc4dmMwZHJNbFZ5VlQwS0xTMHRMUzFGVGtRZ1EwVlNWRWxHU1VOQlZFVXRMUzB0TFFvPSJ9fX19"}], "timestampVerificationData":{"rfc3161Timestamps":[{"signedTimestamp":"MIICyDADAgEAMIICvwYJKoZIhvcNAQcCoIICsDCCAqwCAQMxDTALBglghkgBZQMEAgEwgbcGCyqGSIb3DQEJEAEEoIGnBIGkMIGhAgEBBgkrBgEEAYO/MAIwMTANBglghkgBZQMEAgEFAAQgSXEeWTKPy5M344/YFT/XP8azEPBYve+1t3rfl2y9fU8CFGZeBX8tHbhk1uWSO0u4AA1UktZhGA8yMDI2MTAxOTExMDA1NlowAwIBAaAypDAwLjEVMBMGA1UEChMMc2lnc3RvcmUuZGV2MRUwEwYDVQQDEwxzaWdzdG9yZS10c2GgADGCAdowggHWAgEBMFEwOTEVMBMGA1UEChMMc2lnc3RvcmUuZGV2MSAwHgYDVQQDExdzaWdzdG9yZS10c2Etc2VsZnNpZ25lZAIUOhNULwyQYe68wUMvy4qOiyojiwwwCwYJYIZIAWUDBAIBoIH8MBoGCSqGSIb3DQEJAzENBgsqhkiG9w0BCRABBDAcBgkqhkiG9w0BCQUxDxcNMjYxMDE5MTEwMDU2WjAvBgkqhkiG9w0BCQQxIgQgwfvgGPFqgA9OvUgsLvMTGQENfojqD7cfseLFLfU9c84wgY4GCyqGSIb3DQEJEAIvMX8wfTB7MHkEIIX5J7wHq2LKw7RDVsEO/IGyxog/2nq55thw2dE6zQW3MFUwPaQ7MDkxFTATBgNVBAoTDHNpZ3N0b3JlLmRldjEgMB4GA1UEAxMXc2lnc3RvcmUtdHNhLXNlbGZzaWduZWQCFDoTVC8MkGHuvMFDL8uKjosqI4sMMAoGCCqGSM49BAMCBGYwZAIwYWMvEXVuabGURhsFSF1Y6rl227JXpU/4jZkBHXYZbZbN5Aqs8ZZw2U+hsJmLxqjyAjA9CeY2AqRWdIrE4bTzCHn/eZH533Bce4owlrDB6F9jzhkrgN9DihO3z3YSPGhRhGM="}]}}, "messageSignature":{"messageDigest":{"algorithm":"SHA2_256", "digest":"vR4KddFovSsc+cd0BSWNIvubsoNE+1Dbt58ORDumzZ0="}, "signature":"MEYCIQDsPDB8OL1gvnzSyF6TzJ2YjAoXHRdr8gRvNiWvWnq6IAIhAKzdA2n29j+s2+bMwThPB00tIooA9gQs3EEZ6osuXy/Z"}} \ No newline at end of file diff --git a/build/torch-xpu/modules/__init__.py b/build/torch-xpu/modules/__init__.py deleted file mode 100644 index 0936901f1de9349346957b28291b183e5a8c4d71..0000000000000000000000000000000000000000 --- a/build/torch-xpu/modules/__init__.py +++ /dev/null @@ -1,52 +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.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', -] diff --git a/build/torch-xpu/modules/activations.py b/build/torch-xpu/modules/activations.py deleted file mode 100644 index 93fec5b0796f079fd320a189384229ec7a2fa181..0000000000000000000000000000000000000000 --- a/build/torch-xpu/modules/activations.py +++ /dev/null @@ -1,1205 +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 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: - # 2D: any layout with stride(-1)==1 is valid (can view as (T, D)) - return True - if ndim == 3: - # 3D (B, T, D): stride should be (T*D, D, 1) - return x.stride(0) == x.stride(-2) * x.shape[-2] - if ndim == 4: - # 4D (B, H, T, D): stride should be (H*T*D, T*D, D, 1) - if x.stride(1) != x.stride(-2) * x.shape[-2]: - return False - return x.stride(0) == x.stride(1) * x.shape[1] - # 5D+ fallback to loop - 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: - """Make the tensor inner-contiguous if it isn't already.""" - if _is_inner_contiguous(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=[ - 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 - row = offs // D - col = offs % D - b_x = tl.load(x + row * stride_x_row + col, mask=mask, other=0.).to(tl.float32) - b_y = tl.sigmoid(b_x) - 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 - 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_s = tl.sigmoid(b_x) - b_dx = b_dy * b_s * (1.0 - b_s) - 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, - ) - return dx - - -class SigmoidFunction(torch.autograd.Function): - - @staticmethod - @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 - row = offs // D - col = offs % D - b_x = tl.load(x + row * stride_x_row + col, 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 + 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, - dy, - dx, - stride_x_row, - stride_dy_row, - stride_dx_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 - 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 * ((1. - tl.sigmoid(b_x)) / temperature) - 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, - ) - 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, - ) - return dx - - -class LogSigmoidFunction(torch.autograd.Function): - - @staticmethod - @input_guard(no_guard_contiguous=True) - 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 - row = offs // D - col = offs % D - b_x = tl.load(x + row * stride_x_row + col, mask=mask, other=0.).to(tl.float32) - b_y = b_x * tl.sigmoid(b_x) - 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 - 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_s = tl.sigmoid(b_x) - b_dx = b_dy * b_s * (1.0 + b_x * (1.0 - b_s)) - 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, - ) - return dx - - -class SwishFunction(torch.autograd.Function): - - @staticmethod - @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 - -# 1/sqrt(2*pi)-> 0.3989423 -# 1/sqrt(2) -> 0.70710678 -# 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) - - -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 - - -fast_gelu_impl = FastGeLUFunction.apply - - -@torch.compile -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) - - -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 - row = offs // D - col = offs % D - b_x = tl.load(x + row * stride_x_row + col, mask=mask, other=0.).to(tl.float32) - b_y = tl.load(y + row * stride_y_row + col, mask=mask, other=0.).to(tl.float32) - b_z = b_x * tl.sigmoid(b_x) * b_y - 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, - stride_dx_row, - stride_dy_row, - stride_z_row, - T, - D: tl.constexpr, - B: tl.constexpr, - HAS_WEIGHT: 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_y = tl.load(y + row * stride_y_row + col, mask=mask, other=0.).to(tl.float32) - b_g = tl.load(g + row * stride_g_row + col, mask=mask, other=0.).to(tl.float32) - - b_s = tl.sigmoid(b_x) - b_xs = b_x * b_s - b_dx = b_g * b_s * (1.0 + b_x * (1.0 - b_s)) * b_y - b_dy = b_g * b_xs - - tl.store(dx + row * stride_dx_row + col, b_dx.to(dx.dtype.element_ty), mask=mask) - tl.store(dy + row * stride_dy_row + col, b_dy.to(dy.dtype.element_ty), mask=mask) - if HAS_WEIGHT: - b_z = b_xs * b_y - 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, - 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 - swiglu_fwdbwd_kernel[lambda meta: (triton.cdiv(T, meta['B']),)]( - 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, - T=T, - D=D, - ) - if use_weight: - return dx, dy, z - return dx, dy - - -class SwiGLUFunction(torch.autograd.Function): - r""" - Swish-Gated Linear Unit (SwiGLU) function. - - .. math:: - \text{SwiGLU}(x, y) = swish(x) * y = \frac{x}{1 + \exp(-x)} * y - """ - - @staticmethod - @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): - r""" - Swish-Gated Linear Unit (SwiGLU) function followed by a linear transformation. - - .. math:: - \text{SwiGLULinear}(x, y, W, b) = (swish(x) * y) W + b - - This simple wrap discards the intermediate results of SwiGLU(x, y) to save memory. - """ - - @staticmethod - @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 - @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(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 - row = offs // D - col = offs % D - b_x = tl.load(x + row * stride_x_row + col, mask=mask, other=0.).to(tl.float32) - b_y = tl.load(y + row * stride_y_row + col, mask=mask, other=0.).to(tl.float32) - b_z = tl.sigmoid(b_x) * b_y - 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, - stride_dx_row, - stride_dy_row, - stride_z_row, - T, - D: tl.constexpr, - B: tl.constexpr, - HAS_WEIGHT: 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_y = tl.load(y + row * stride_y_row + col, mask=mask, other=0.).to(tl.float32) - b_g = tl.load(g + row * stride_g_row + col, mask=mask, other=0.).to(tl.float32) - - b_s = tl.sigmoid(b_x) - b_dx = b_g * b_s * (1.0 - b_s) * b_y - b_dy = b_g * b_s - - tl.store(dx + row * stride_dx_row + col, b_dx.to(dx.dtype.element_ty), mask=mask) - tl.store(dy + row * stride_dy_row + col, b_dy.to(dy.dtype.element_ty), mask=mask) - if HAS_WEIGHT: - b_z = b_s * b_y - tl.store(z + row * stride_z_row + col, b_z.to(z.dtype.element_ty), mask=mask) - - -@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, - ) - return z - - -@torch.compiler.disable -def sigmoidglu_fwdbwd( - 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"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']),)]( - 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, - T=T, - D=D, - ) - if use_weight: - return dx, dy, z - return dx, dy - - -class SigmoidGLUFunction(torch.autograd.Function): - r""" - Sigmoid-Gated Linear Unit (SigmoidGLU) function. - - .. math:: - \text{SigmoidGLU}(x, y) = sigmoid(x) * y = \frac{1}{1 + \exp(-x)} * y - """ - - @staticmethod - @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): - r""" - Sigmoid-Gated Linear Unit (SigmoidGLU) function followed by a linear transformation. - - .. math:: - \text{SigmoidGLULinear}(x, y, W, b) = (sigmoid(x) * y) W + b - - This simple wrap discards the intermediate results of SigmoidGLU(x, y) to save memory. - """ - - @staticmethod - @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 - 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 = 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 - - -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 - row = offs // D - col = offs % D - b_x = tl.load(x + row * stride_x_row + col, mask=mask, other=0.).to(tl.float32) - b_y = tl.load(y + row * stride_y_row + col, 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 + 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, - stride_dx_row, - stride_dy_row, - stride_z_row, - m, - T, - D: tl.constexpr, - B: tl.constexpr, - HAS_WEIGHT: 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_y = tl.load(y + row * stride_y_row + col, mask=mask, other=0.).to(tl.float32) - b_g = tl.load(g + row * stride_g_row + col, 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 + row * stride_dx_row + col, b_dx.to(dx.dtype.element_ty), mask=mask) - tl.store(dy + row * stride_dy_row + col, b_dy.to(dy.dtype.element_ty), mask=mask) - if HAS_WEIGHT: - b_z = b_gate * b_y - 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, - ) - return z - - -@dispatch('modules') -def powglu_fwdbwd( - 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 - powglu_fwdbwd_kernel[lambda meta: (triton.cdiv(T, meta['B']),)]( - 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, - ) - if use_weight: - return dx, dy, z - return dx, dy - - -class PowGLUFunction(torch.autograd.Function): - r""" - Power-Gated Linear Unit (PowGLU) function. - - .. math:: - \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 - @input_guard(no_guard_contiguous=True) - 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 - - -class PowGLULinearFunction(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(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(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 - - -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, - weight: torch.Tensor, - bias: torch.Tensor, - power: float = 3.0, -) -> torch.Tensor: - return PowGLULinearFunction.apply(x, y, weight, bias, power) - - -ACT2FN = { - 'relu': F.relu, - 'sigmoid': sigmoid, - 'logsigmoid': logsigmoid, - 'silu': swish, - 'swish': swish, - 'sqrelu': sqrelu, - 'gelu': fast_gelu_impl, - 'bias_gelu': bias_gelu_impl, -} diff --git a/build/torch-xpu/modules/backends/__init__.py b/build/torch-xpu/modules/backends/__init__.py deleted file mode 100644 index 4ce5614452d211b97584dab79c74560f9c6d4c1f..0000000000000000000000000000000000000000 --- a/build/torch-xpu/modules/backends/__init__.py +++ /dev/null @@ -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'] diff --git a/build/torch-xpu/modules/backends/triton_ascend/__init__.py b/build/torch-xpu/modules/backends/triton_ascend/__init__.py deleted file mode 100644 index ba9c7a925c84d63ebed71917d64fcdcadd33df6c..0000000000000000000000000000000000000000 --- a/build/torch-xpu/modules/backends/triton_ascend/__init__.py +++ /dev/null @@ -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, - ) diff --git a/build/torch-xpu/modules/backends/triton_ascend/activations.py b/build/torch-xpu/modules/backends/triton_ascend/activations.py deleted file mode 100644 index 7f43a381e8b3d806cbfde5b3891c92e888aa467f..0000000000000000000000000000000000000000 --- a/build/torch-xpu/modules/backends/triton_ascend/activations.py +++ /dev/null @@ -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) diff --git a/build/torch-xpu/modules/backends/triton_ascend/causal_conv1d.py b/build/torch-xpu/modules/backends/triton_ascend/causal_conv1d.py deleted file mode 100644 index 9f77b9ec28b1e91197c1d47d6a66f181f83c2176..0000000000000000000000000000000000000000 --- a/build/torch-xpu/modules/backends/triton_ascend/causal_conv1d.py +++ /dev/null @@ -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 diff --git a/build/torch-xpu/modules/backends/triton_ascend/fused_cross_entropy.py b/build/torch-xpu/modules/backends/triton_ascend/fused_cross_entropy.py deleted file mode 100644 index 42aaec21dc9dc0461aac19c0b1e998efca04cb48..0000000000000000000000000000000000000000 --- a/build/torch-xpu/modules/backends/triton_ascend/fused_cross_entropy.py +++ /dev/null @@ -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, - ) diff --git a/build/torch-xpu/modules/backends/triton_ascend/fused_kl_div.py b/build/torch-xpu/modules/backends/triton_ascend/fused_kl_div.py deleted file mode 100644 index e593e06b377ce4d0843dde112540f498ef659ce6..0000000000000000000000000000000000000000 --- a/build/torch-xpu/modules/backends/triton_ascend/fused_kl_div.py +++ /dev/null @@ -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 diff --git a/build/torch-xpu/modules/backends/triton_ascend/fused_linear_cross_entropy.py b/build/torch-xpu/modules/backends/triton_ascend/fused_linear_cross_entropy.py deleted file mode 100644 index 3d8d9fda021352e37359fad061135aa7236812f4..0000000000000000000000000000000000000000 --- a/build/torch-xpu/modules/backends/triton_ascend/fused_linear_cross_entropy.py +++ /dev/null @@ -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 diff --git a/build/torch-xpu/modules/backends/triton_ascend/grpo.py b/build/torch-xpu/modules/backends/triton_ascend/grpo.py deleted file mode 100644 index 463923f5d560b23b7248c1ff0df6256ed222445b..0000000000000000000000000000000000000000 --- a/build/torch-xpu/modules/backends/triton_ascend/grpo.py +++ /dev/null @@ -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) diff --git a/build/torch-xpu/modules/backends/triton_ascend/layernorm.py b/build/torch-xpu/modules/backends/triton_ascend/layernorm.py deleted file mode 100644 index 9a8fe6ba39c7c99d496ba52ff4f349c6a951bf07..0000000000000000000000000000000000000000 --- a/build/torch-xpu/modules/backends/triton_ascend/layernorm.py +++ /dev/null @@ -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) diff --git a/build/torch-xpu/modules/backends/triton_ascend/rotary.py b/build/torch-xpu/modules/backends/triton_ascend/rotary.py deleted file mode 100644 index 4042cc23c6b8e1eb49f6e853703d68176e5960b4..0000000000000000000000000000000000000000 --- a/build/torch-xpu/modules/backends/triton_ascend/rotary.py +++ /dev/null @@ -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 diff --git a/build/torch-xpu/modules/conv/__init__.py b/build/torch-xpu/modules/conv/__init__.py deleted file mode 100644 index 23421fa8c55cf9599bfa4dd8f030b38f599d5d48..0000000000000000000000000000000000000000 --- a/build/torch-xpu/modules/conv/__init__.py +++ /dev/null @@ -1,19 +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 .causal_conv1d import causal_conv1d -from .long_conv import ImplicitLongConvolution, LongConvolution, PositionalEmbedding, fft_conv -from .short_conv import ShortConvolution - -__all__ = [ - 'ImplicitLongConvolution', - 'LongConvolution', - 'PositionalEmbedding', - 'ShortConvolution', - 'causal_conv1d', - 'fft_conv', -] diff --git a/build/torch-xpu/modules/conv/causal_conv1d.py b/build/torch-xpu/modules/conv/causal_conv1d.py deleted file mode 100644 index 98e44923f37df9a02088de4c8836de50f0b9a100..0000000000000000000000000000000000000000 --- a/build/torch-xpu/modules/conv/causal_conv1d.py +++ /dev/null @@ -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}") diff --git a/build/torch-xpu/modules/conv/cp/ops.py b/build/torch-xpu/modules/conv/cp/ops.py deleted file mode 100644 index 3b5b8e758d0c9892b8168109d68821187750a1a8..0000000000000000000000000000000000000000 --- a/build/torch-xpu/modules/conv/cp/ops.py +++ /dev/null @@ -1,258 +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 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): - """ - Context Parallel version of CausalConv1dFunction. - - Forward: - 1. Get tails from previous rank to construct initial_state - 2. Call causal_conv1d_fwd - - Backward: - 1. Call causal_conv1d_bwd to get dx - 2. Sync communication: add next rank's first W-1 token gradients to current rank's last W-1 tokens - """ - - @staticmethod - def _prepare_initial_state_for_cp( - x: torch.Tensor, - weight: torch.Tensor, - cu_seqlens: torch.Tensor | None, - context: FLACPContext, - group: dist.ProcessGroup | None, - ) -> torch.Tensor | None: - """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) - if valid_len > 0: - # 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 - - @staticmethod - def _correct_dx_for_cp( - dx: torch.Tensor, - 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: - # Only keep gradients for positions that had real data from the - # 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] - # Add to current rank's last W-1 tokens (these tokens are used as initial_state by next rank) - dx[0, -(W-1):, :].add_(recv_d_init) - - @staticmethod - def forward( - ctx, - x: torch.Tensor, - weight: torch.Tensor, - bias: torch.Tensor | None, - activation: str | None, - chunk_indices: torch.Tensor | None, - 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") - cu_seqlens = cp_context.cu_seqlens - 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, - cu_seqlens=cu_seqlens, - context=cp_context, - 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 - ctx.chunk_indices = chunk_indices - 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( - x: torch.Tensor, - weight: torch.Tensor, - bias: torch.Tensor | None = None, - activation: str | None = None, - chunk_indices: torch.Tensor | None = None, - 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 - ) diff --git a/build/torch-xpu/modules/conv/long_conv.py b/build/torch-xpu/modules/conv/long_conv.py deleted file mode 100644 index cafa1881400f31cd3168344cbeef2632fbd035b1..0000000000000000000000000000000000000000 --- a/build/torch-xpu/modules/conv/long_conv.py +++ /dev/null @@ -1,172 +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 math - -import torch -import torch.nn as nn -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 - if k_rev is not None: - k_rev_f = torch.fft.rfft(k_rev, n=fft_size) / fft_size - k_f = k_f + k_rev_f.conj() - u_f = torch.fft.rfft(u.to(dtype=k.dtype), n=fft_size) - - if len(u.shape) > 3: - k_f = k_f.unsqueeze(1) - y = torch.fft.irfft(u_f * k_f, n=fft_size, norm="forward")[..., :seqlen] - - out = y + u - if gelu: - out = F.gelu(out) - if dropout_mask is not None: - return (out * rearrange(dropout_mask, "b H -> b H 1")).to(dtype=u.dtype) - else: - return out.to(dtype=u.dtype) - - -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 - - f = torch.linspace(1e-4, bands - 1, bands)[None, None] - z = torch.exp(-1j * f * w) - 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] - - -class ImplicitLongConvolution(nn.Module): - """ - Long convolution with implicit filter parameterized by an MLP. - - Args: - hidden_size (int): - 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) diff --git a/build/torch-xpu/modules/conv/short_conv.py b/build/torch-xpu/modules/conv/short_conv.py deleted file mode 100644 index a53386bc6b751edf4ad94046890024bf35d5029d..0000000000000000000000000000000000000000 --- a/build/torch-xpu/modules/conv/short_conv.py +++ /dev/null @@ -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 diff --git a/build/torch-xpu/modules/conv/triton/kernels.py b/build/torch-xpu/modules/conv/triton/kernels.py deleted file mode 100644 index 86ae0d9cd76bfb8061b63d2b55e69a353976054b..0000000000000000000000000000000000000000 --- a/build/torch-xpu/modules/conv/triton/kernels.py +++ /dev/null @@ -1,683 +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 -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 - - -@triton.heuristics({ - 'HAS_WEIGHT': lambda args: args['weight'] is not None, - 'HAS_BIAS': lambda args: args['bias'] is not None, - 'HAS_RESIDUAL': lambda args: args['residual'] is not None, - 'USE_INITIAL_STATE': lambda args: args['initial_state'] is not None, - 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None, -}) -@fla_cache_autotune( - 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', 'NB'], - **autotune_cache_kwargs, -) -@triton.jit -def causal_conv1d_fwd_kernel( - x, - y, - weight, - bias, - residual, - cu_seqlens, - initial_state, - chunk_indices, - B, - T, - stride_x_n, - stride_x_t, - stride_x_d, - 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, - HAS_RESIDUAL: tl.constexpr, - USE_INITIAL_STATE: tl.constexpr, - IS_VARLEN: 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 - else: - i_n = i_b - bos, eos = (i_b * T).to(tl.int64), (i_b * T + T).to(tl.int64) - 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] - b_w = tl.load(weight + o_d[:, None] * W + o_w, mask=m_d[:, None] & m_w, other=0).to(tl.float32) - - 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) - - if HAS_WEIGHT: - b_yi *= tl.sum(b_w * (o_w == (i_w + W - 1)), 1) - b_y += b_yi - - if HAS_BIAS: - b_y += tl.load(bias + o_d, mask=m_d).to(tl.float32) - - if ACTIVATION == 'swish' or ACTIVATION == 'silu': - 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({ - '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, -}) -@fla_cache_autotune( - 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 -def causal_conv1d_bwd_kernel( - x, - y, - weight, - initial_state, - dht, - dy, - dx, - dw, - db, - cu_seqlens, - 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, - USE_INITIAL_STATE: tl.constexpr, - USE_FINAL_STATE: tl.constexpr, - IS_VARLEN: 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 - else: - i_tg = i_b * tl.num_programs(1) + i_t - i_n = i_b - bos, eos = (i_b * T).to(tl.int64), (i_b * T + T).to(tl.int64) - p_x = x + tl.cast(i_b, tl.int64) * stride_x_n - - if IS_VARLEN: - p_dy = dy + bos * stride_dy_t - else: - 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) - - b_dx = tl.zeros((BT, BD), dtype=tl.float32) - if HAS_BIAS: - b_db = tl.zeros((BD,), dtype=tl.float32) - - 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 - if HAS_WEIGHT: - # [BT, BD] - b_wdy = b_wdy * tl.sum(b_w * (o_w == (W - i_w - 1)), 1) - # [BD] - b_dw = tl.sum(b_dy * b_x, 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 HAS_BIAS and i_w == 0: - 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 - if HAS_WEIGHT: - # [BT, BD] - b_wdy = b_wdy * tl.sum(b_w * (o_w == (W - i_w - 1)), 1) - # [BD] - b_dw = tl.sum(b_dy * b_x, 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 HAS_BIAS and i_w == 0: - 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) - - if HAS_BIAS and i_w == 0: - b_db += tl.sum(b_dy_shift, 0) - b_wdy = b_dy_shift if not HAS_WEIGHT else (b_dy_shift * tl.sum(b_w * (o_w == (W - i_w - 1)), 1)) - 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 = 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 - - if IS_VARLEN: - p_dx = dx + bos * stride_dx_t - 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({ - 'USE_INITIAL_STATE': lambda args: args['cache'] is not None, - 'HAS_WEIGHT': lambda args: args['weight'] is not None, - 'HAS_BIAS': lambda args: args['bias'] is not None, - 'HAS_RESIDUAL': lambda args: args['residual'] is not None, -}) -@fla_cache_autotune( - configs=[ - triton.Config({'BD': BD}, num_warps=num_warps) - for BD in [8, 16, 32, 64, 128, 256] - for num_warps in NUM_WARPS_AUTOTUNE - ], - key=['D', 'W'], - restore_value=['cache'], - **autotune_cache_kwargs, -) -@triton.jit -def causal_conv1d_update_kernel( - x, - cache, - residual, - 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, - BW: tl.constexpr, - ACTIVATION: tl.constexpr, - USE_INITIAL_STATE: tl.constexpr, - HAS_WEIGHT: tl.constexpr, - 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) - m_d = o_d < D - m_w = o_w < W - - # [BD] - b_x = tl.load(x + i_n * stride_x_n + o_d * stride_x_d, mask=m_d, other=0).to(tl.float32) - - 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) - - if HAS_WEIGHT: - b_w = tl.load(weight + o_d[:, None] * W + o_w, mask=m_d[:, None] & m_w, other=0) - b_y = tl.sum(b_cache * b_w, 1) - else: - b_y = tl.sum(b_cache, 1) - - if HAS_BIAS: - b_y += tl.load(bias + o_d, mask=m_d) - - if ACTIVATION == 'swish' or ACTIVATION == 'silu': - b_y = b_y * tl.sigmoid(b_y) - - 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 -def compute_dh0_kernel( - dy, - 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) - - -@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, -): - 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 - - 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 - - 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: - o_c = W - (BW - seq_len) + tl.arange(0, BW) - m_c = (o_c >= 0) & (o_c < W) - - p_init = initial_state + i_n * D*W + o_d[None, :] * W + o_c[:, None] - mask_init = m_d[None, :] & m_c[:, None] - - 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, - 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 - - -@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 diff --git a/build/torch-xpu/modules/conv/triton/ops.py b/build/torch-xpu/modules/conv/triton/ops.py deleted file mode 100644 index aee880417d0372500f9f67721baf34cc02bfeee1..0000000000000000000000000000000000000000 --- a/build/torch-xpu/modules/conv/triton/ops.py +++ /dev/null @@ -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 diff --git a/build/torch-xpu/modules/convolution.py b/build/torch-xpu/modules/convolution.py deleted file mode 100644 index b00367976440d67f5dd628b5eff6b20c8eee5c00..0000000000000000000000000000000000000000 --- a/build/torch-xpu/modules/convolution.py +++ /dev/null @@ -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', -] diff --git a/build/torch-xpu/modules/feature_map.py b/build/torch-xpu/modules/feature_map.py deleted file mode 100644 index a35aa11b6803e89469bb8afbd50fdacc74984577..0000000000000000000000000000000000000000 --- a/build/torch-xpu/modules/feature_map.py +++ /dev/null @@ -1,315 +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 functools -import math - -import torch -import torch.nn.functional as F -from torch import nn - -from ..modules.activations import fast_gelu_impl, sigmoid, sqrelu, swish -from ..modules.layernorm import layer_norm -from ..utils import checkpoint - - -@functools.cache -def _triu_indices(n: int, offset: int, device: torch.device) -> torch.Tensor: - # cache the upper-triangular gather indices per (size, offset, device) to avoid rebuilding - # them and copying host -> device on every forward - return torch.triu_indices(n, n, offset, device=device) - - -@checkpoint -def flatten_diag_outer_product(x: torch.Tensor, y: torch.Tensor) -> torch.Tensor: - z = torch.einsum("...i,...j->...ij", x, y) - N = z.size(-1) - indices = _triu_indices(N, 0, z.device) - return z[..., indices[0], indices[1]] - - -@checkpoint -def flatten_diag_outer_product_off1(x: torch.Tensor, y: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: - z = torch.einsum("...i,...j->...ij", x, y) - N = z.size(-1) - indices = _triu_indices(N, 1, z.device) - diag = torch.arange(N, device=z.device) - return z[..., indices[0], indices[1]], z[..., diag, diag] - - -def is_power_of_2(n: int) -> bool: - return (n & (n - 1) == 0) and n != 0 - - -class HedgehogFeatureMap(nn.Module): - - r""" - Hedgehog feature map as introduced in - `The Hedgehog & the Porcupine: Expressive Linear Attentions with Softmax Mimicry `_ - """ - - def __init__( - self, - head_dim: int, - ) -> None: - super().__init__() - # Trainable map - self.layer = nn.Linear(head_dim, head_dim) - self.init_weights_() - - def init_weights_(self): - """Initialize trainable map as identity""" - nn.init.eye_(self.layer.weight) - nn.init.zeros_(self.layer.bias) - - def forward(self, x: torch.Tensor) -> torch.Tensor: - x = self.layer(x) # b, h, l, d - y = 2 * x - return torch.cat([y, -y], dim=-1).softmax(-1) - - -class T2RFeatureMap(nn.Module): - - r""" - Simple linear mapping feature map as in - `Finetuning Pretrained Transformers into RNNs `_ - """ - - def __init__( - self, - head_dim: int, - dot_dim: int | None = None, - bias: bool | None = False, - ) -> None: - super().__init__() - # Trainable map - if dot_dim is None: - dot_dim = head_dim - - self.head_dim = head_dim - self.dot_dim = dot_dim - self.bias = bias - - self.layer = nn.Linear(head_dim, dot_dim, bias=bias) - - def __repr__(self) -> str: - return f"{self.__class__.__name__}(head_dim={self.head_dim}, dot_dim={self.dot_dim}, bias={self.bias})" - - def forward(self, x: torch.Tensor) -> torch.Tensor: - return self.layer(x).relu() - - -class DPFPFeatureMap(nn.Module): - - r""" - Deterministic Parameter-Free Projection (DPFP) feature map in - `Linear Transformers Are Secretly Fast Weight Programmers `_ - """ - - def __init__( - self, - head_dim: int, - nu: int = 4, - ) -> None: - super().__init__() - 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 - - -class HadamardFeatureMap(nn.Module): - def __init__( - self, - head_dim: int, - ) -> None: - super().__init__() - # Trainable map - self.layer1 = nn.Linear(head_dim, head_dim) - self.layer2 = nn.Linear(head_dim, head_dim) - - def forward(self, x: torch.Tensor) -> torch.Tensor: - return self.layer1(x) * self.layer2(x) - - -class LearnableOuterProductFeatureMap(nn.Module): - def __init__( - self, - head_dim: int, - feature_dim: int, - ) -> None: - super().__init__() - # Trainable map - self.layer1 = nn.Linear(head_dim, feature_dim, bias=False) - self.layer2 = nn.Linear(head_dim, feature_dim, bias=False) - self.normalizer = feature_dim ** -0.5 - - def forward(self, x: torch.Tensor) -> torch.Tensor: - return flatten_diag_outer_product(self.layer1(x), self.layer2(x)) - - -class LearnablePolySketchNonNegativeFeatureMap(nn.Module): - - def __init__( - self, - head_dim: int, - sketch_size: int | None = None, - degree: int | None = 2, - ) -> None: - super().__init__() - - assert is_power_of_2(degree) and degree >= 2, f"The degree {degree} must be a power of 2" - - if sketch_size is None: - sketch_size = head_dim - - self.head_dim = head_dim - self.sketch_size = sketch_size - self.degree = degree - - self.gamma = nn.Parameter(torch.ones(head_dim)) - self.beta = nn.Parameter(torch.zeros(head_dim)) - # NOTE: the sketch layers defined here are quite different from the original paper - # currently we simply use linear layers without any non-linear activations - self.sketches1 = nn.ModuleList([ - nn.Linear(head_dim, sketch_size, bias=False), - *[nn.Linear(sketch_size, sketch_size, bias=False) for _ in range(int(math.log2(self.degree)) - 2)], - ]) - self.sketches2 = nn.ModuleList([ - nn.Linear(head_dim, sketch_size, bias=False), - *[nn.Linear(sketch_size, sketch_size, bias=False) for _ in range(int(math.log2(self.degree)) - 2)], - ]) - - def forward(self, x: torch.Tensor) -> torch.Tensor: - # Section 2.1 - x = layer_norm(x, self.gamma, self.beta) - # first map the input to sketch size with learnable parameters - x = self.sketches1[0](x) * self.sketches2[0](x) * self.head_dim ** -0.5 - for i in range(1, int(math.log2(self.degree)) - 1): - x = self.sketches1[i](x) * self.sketches2[i](x) * self.head_dim ** -0.5 - # do sketch mapping for log2(p) - 1 times in total - # do p=2 mapping to ensure non-negativity - return flatten_diag_outer_product(x, x) - - -class TaylorFeatureMap(nn.Module): - def __init__( - self, - head_dim: int, - ) -> None: - super().__init__() - self.head_dim = head_dim - self.r2 = math.sqrt(2) - self.rd = math.sqrt(self.head_dim) - self.rrd = math.sqrt(self.rd) - - def forward(self, x: torch.Tensor) -> torch.Tensor: - x2_1, x2_2 = flatten_diag_outer_product_off1(x, x) - return torch.cat([torch.ones_like(x[..., 0:1]), x / self.rrd, x2_2 / (self.rd * self.r2), x2_1 / self.rd], dim=-1) - - -class RebasedFeatureMap(nn.Module): - - def __init__( - self, - head_dim: int, - use_gamma: bool | None = True, - use_beta: bool | None = True, - normalize: bool | None = True, - ) -> None: - super().__init__() - - self.head_dim = head_dim - self.use_gamma = use_gamma - self.use_beta = use_beta - self.normalize = normalize - - self.gamma = None - self.beta = None - if use_gamma: - self.gamma = nn.Parameter(torch.ones(head_dim)) - if use_beta: - self.beta = nn.Parameter(torch.zeros(head_dim)) - - def forward(self, x: torch.Tensor, flatten: bool | None = True) -> torch.Tensor: - if self.use_beta and self.use_gamma and self.normalize: - x = layer_norm(x, self.gamma, self.beta) - elif self.normalize: - x = F.layer_norm(x, (self.head_dim,), self.gamma, self.beta) - elif self.use_gamma and self.use_beta: - x = torch.addcmul(self.beta, x, self.gamma) - elif self.use_gamma: - x = x.mul(self.gamma) - else: - raise RuntimeError(f"Not supported combination of `use_gamma`, `use_beta` and `normalize`, " - f"which is currently set as (`{self.use_gamma}`, `{self.use_beta}`, `{self.normalize}`)") - if not flatten: - return x - x2_1, x2_2 = flatten_diag_outer_product_off1(x, x) - # rebased use learnable parameters to approximate any quadratic function - return torch.cat([x2_2 * self.head_dim ** -0.5, x2_1 * (2 / self.head_dim) ** 0.5], dim=-1) - - -class ReLUFeatureMap(nn.Module): - - def __init__( - self, - ) -> None: - super().__init__() - - def forward(self, x: torch.Tensor) -> torch.Tensor: - return F.relu(x) - - -class SquaredReLUFeatureMap(nn.Module): - - def __init__( - self, - ) -> None: - super().__init__() - - def forward(self, x: torch.Tensor) -> torch.Tensor: - return sqrelu(x) - - -class GELUFeatureMap(nn.Module): - - def __init__( - self, - ) -> None: - super().__init__() - - def forward(self, x: torch.Tensor) -> torch.Tensor: - return fast_gelu_impl(x) - - -class SwishFeatureMap(nn.Module): - - def __init__( - self, - ) -> None: - super().__init__() - - def forward(self, x: torch.Tensor) -> torch.Tensor: - return swish(x) - - -class SigmoidFeatureMap(nn.Module): - - def __init__( - self, - ) -> None: - super().__init__() - - def forward(self, x: torch.Tensor) -> torch.Tensor: - return sigmoid(x) diff --git a/build/torch-xpu/modules/fused_bitlinear.py b/build/torch-xpu/modules/fused_bitlinear.py deleted file mode 100644 index 4b36c859aaea1277089d643c3b2fb863b3fdd4d4..0000000000000000000000000000000000000000 --- a/build/torch-xpu/modules/fused_bitlinear.py +++ /dev/null @@ -1,638 +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 - -# 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, - linear_weight, - linear_bias, - residual, - eps, - prenorm, - residual_in_fp32, - is_rms_norm, - ) - - -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, - linear_weight=linear_weight, - linear_bias=linear_bias, - residual=residual, - eps=eps, - prenorm=prenorm, - residual_in_fp32=residual_in_fp32, - is_rms_norm=True, - ) - - -@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, - ) diff --git a/build/torch-xpu/modules/fused_cross_entropy.py b/build/torch-xpu/modules/fused_cross_entropy.py deleted file mode 100644 index c0b770e50014cadfabb101f42655cc4e383d5cc2..0000000000000000000000000000000000000000 --- a/build/torch-xpu/modules/fused_cross_entropy.py +++ /dev/null @@ -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 diff --git a/build/torch-xpu/modules/fused_linear_cross_entropy.py b/build/torch-xpu/modules/fused_linear_cross_entropy.py deleted file mode 100644 index 8e9128a63e1ae7ae6aebc16f533edb53ef711f01..0000000000000000000000000000000000000000 --- a/build/torch-xpu/modules/fused_linear_cross_entropy.py +++ /dev/null @@ -1,767 +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 - -# 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 -import torch.nn.functional as F -import triton -import triton.language as tl -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 - - -@triton.heuristics({ - 'HAS_SCALE': lambda args: args['scale'] is not None, - 'HAS_SOFTCAPPING': lambda args: args['softcapping'] is not None, -}) -@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) - - # [N, H] - 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 - - # [V, H] - dw = torch.zeros_like(weight, device=device, dtype=grad_dtype) if weight is not None else None - # [V] - db = torch.zeros_like(bias, device=device, dtype=bias_grad_dtype) if bias is not None else None - # [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: - db = db.to(bias) - 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 - @input_guard - def forward( - ctx, - 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, - db.detach() if bias is not None else None, - ) - return loss - - @staticmethod - @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, - target, - weight, - bias, - ignore_index, - label_smoothing, - logit_scale, - logit_softcapping, - num_chunks, - reduction, - use_l2warp, - l2_penalty_factor, - accumulate_grad_in_fp32, - ) - - -class FusedLinearCrossEntropyLoss(nn.Module): - - def __init__( - self, - 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__() - - assert reduction in ["mean", "sum"], f"reduction: {reduction} is not supported" - - self.ignore_index = ignore_index - self.label_smoothing = label_smoothing - self.logit_scale = logit_scale - self.logit_softcapping = logit_softcapping - self.num_chunks = num_chunks - self.reduction = reduction - 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, - label_smoothing=self.label_smoothing, - logit_scale=self.logit_scale, - logit_softcapping=self.logit_softcapping, - num_chunks=self.num_chunks, - reduction=self.reduction, - 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),) - self.use_local_output = use_local_output - - @staticmethod - def _prepare_input_fn(sequence_sharding, mod, inputs, device_mesh): - x, target, weight, bias = inputs - - if not isinstance(x, DTensor): - # assume the input passed in already sharded on the sequence dim and create the DTensor - x = DTensor.from_local(x, device_mesh, sequence_sharding) - if x.placements != sequence_sharding: - x = x.redistribute(placements=sequence_sharding, async_op=True) - if not isinstance(target, DTensor): - target = DTensor.from_local(target, device_mesh, [Replicate()]) - if target.placements != sequence_sharding: - target = target.redistribute(placements=sequence_sharding, async_op=True) - - if not isinstance(weight, DTensor): - weight = DTensor.from_local(weight, device_mesh, [Replicate()]) - if weight.placements != [Replicate()]: - # we replicate the weight/bias in FLCE - weight = weight.redistribute(placements=[Replicate()], async_op=True) - - if bias is not None and not isinstance(bias, DTensor): - bias = DTensor.from_local(bias, device_mesh, [Replicate()]) - if bias is not None and bias.placements != [Replicate()]: - bias = bias.redistribute(placements=[Replicate()], async_op=True) - - return x.to_local(), target.to_local(), weight.to_local(), bias.to_local() if bias is not None else bias - - @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( - module, - device_mesh, - partition_fn=None, - input_fn=partial(self._prepare_input_fn, self.sequence_sharding), - output_fn=partial(self._prepare_output_fn, self.use_local_output), - ) diff --git a/build/torch-xpu/modules/l2norm.py b/build/torch-xpu/modules/l2norm.py deleted file mode 100644 index a840acaf4d2555d87b2958dd1dbd5f1d30362f96..0000000000000000000000000000000000000000 --- a/build/torch-xpu/modules/l2norm.py +++ /dev/null @@ -1,288 +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 torch.nn as nn -import triton -import triton.language as tl - -from ..ops.utils.cache import fla_cache_autotune -from ..utils import IS_AMD, autotune_cache_kwargs, input_guard - -BT_LIST = [8, 16, 32, 64, 128] -NUM_WARPS_AUTOTUNE = [1, 2, 4, 8, 16] if IS_AMD else [1, 2, 4, 8, 16, 32] - - -@triton.autotune( - configs=[triton.Config({}, num_warps=num_warps) for num_warps in NUM_WARPS_AUTOTUNE], - key=["D"], - **autotune_cache_kwargs, -) -@triton.jit -def l2norm_fwd_kernel1( - x, - y, - rstd, - eps, - D, - BD: tl.constexpr, -): - i_t = tl.program_id(0) - x += i_t * D - y += i_t * D - # Compute mean and variance - cols = tl.arange(0, BD) - mask = cols < D - - b_x = tl.load(x + cols, mask=mask, other=0.0).to(tl.float32) - b_rstd = 1 / tl.sqrt(tl.sum(b_x * b_x) + eps) - b_y = b_x * b_rstd - tl.store(y + cols, b_y, mask=mask) - tl.store(rstd + i_t, b_rstd) - - -@triton.autotune( - configs=[triton.Config({}, num_warps=num_warps) for num_warps in NUM_WARPS_AUTOTUNE], - key=["D"], - **autotune_cache_kwargs, -) -@triton.jit -def l2norm_bwd_kernel1( - y, - rstd, - dy, - dx, - eps, - D, - BD: tl.constexpr, -): - i_t = tl.program_id(0) - y += i_t * D - dx += i_t * D - dy += i_t * D - - cols = tl.arange(0, BD) - mask = cols < D - b_y = tl.load(y + cols, mask=mask, other=0.0).to(tl.float32) - b_rstd = tl.load(rstd + i_t).to(tl.float32) - b_dy = tl.load(dy + cols, mask=mask, other=0.0).to(tl.float32) - b_dx = b_dy * b_rstd - tl.sum(b_dy * b_y) * b_y * b_rstd - tl.store(dx + cols, b_dx, mask=mask) - - -@fla_cache_autotune( - configs=[triton.Config({"BT": BT}, num_warps=num_warps) for num_warps in [1, 2, 4, 8, 16] for BT in BT_LIST], - key=["D", "NB"], - **autotune_cache_kwargs, -) -@triton.jit(do_not_specialize=["T"]) -def l2norm_fwd_kernel( - x, - y, - rstd, - eps, - T, - D: tl.constexpr, - BD: tl.constexpr, - NB: tl.constexpr, - BT: tl.constexpr, -): - i_t = tl.program_id(0) - p_x = tl.make_block_ptr(x, (T, D), (D, 1), (i_t * BT, 0), (BT, BD), (1, 0)) - p_y = tl.make_block_ptr(y, (T, D), (D, 1), (i_t * BT, 0), (BT, BD), (1, 0)) - p_rstd = tl.make_block_ptr(rstd, (T,), (1,), (i_t * BT,), (BT,), (0,)) - - b_x = tl.load(p_x, boundary_check=(0, 1)).to(tl.float32) - b_rstd = 1 / tl.sqrt(tl.sum(b_x * b_x, 1) + eps) - b_y = b_x * b_rstd[:, None] - - tl.store(p_y, b_y.to(p_y.dtype.element_ty), boundary_check=(0, 1)) - tl.store(p_rstd, b_rstd.to(p_rstd.dtype.element_ty), boundary_check=(0,)) - - -@fla_cache_autotune( - configs=[triton.Config({"BT": BT}, num_warps=num_warps) for num_warps in [1, 2, 4, 8, 16] for BT in BT_LIST], - key=["D", "NB"], - **autotune_cache_kwargs, -) -@triton.jit(do_not_specialize=["T"]) -def l2norm_bwd_kernel( - y, - rstd, - dy, - dx, - eps, - T, - D: tl.constexpr, - BD: tl.constexpr, - NB: tl.constexpr, - BT: tl.constexpr, -): - i_t = tl.program_id(0) - p_y = tl.make_block_ptr(y, (T, D), (D, 1), (i_t * BT, 0), (BT, BD), (1, 0)) - p_rstd = tl.make_block_ptr(rstd, (T,), (1,), (i_t * BT,), (BT,), (0,)) - p_dy = tl.make_block_ptr(dy, (T, D), (D, 1), (i_t * BT, 0), (BT, BD), (1, 0)) - p_dx = tl.make_block_ptr(dx, (T, D), (D, 1), (i_t * BT, 0), (BT, BD), (1, 0)) - - b_y = tl.load(p_y, boundary_check=(0, 1)).to(tl.float32) - b_rstd = tl.load(p_rstd, boundary_check=(0,)).to(tl.float32) - b_dy = tl.load(p_dy, boundary_check=(0, 1)).to(tl.float32) - b_dx = b_dy * b_rstd[:, None] - tl.sum(b_dy * b_y, 1)[:, None] * b_y * b_rstd[:, None] - tl.store(p_dx, b_dx.to(p_dx.dtype.element_ty), boundary_check=(0, 1)) - - -def l2norm_fwd( - x: torch.Tensor, - eps: float = 1e-6, - output_dtype: torch.dtype | None = None, -): - x_shape_og = x.shape - x = x.view(-1, x.shape[-1]) - # allocate output - if output_dtype is None: - y = torch.empty_like(x) - else: - y = torch.empty_like(x, dtype=output_dtype) - assert y.stride(-1) == 1 - T, D = x.shape[0], x.shape[-1] - # 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 doesn't support feature dim >= 64KB.") - - rstd = torch.empty((T,), dtype=torch.float32, device=x.device) - 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"]),) - - l2norm_fwd_kernel[grid]( - x=x, - y=y, - rstd=rstd, - eps=eps, - T=T, - D=D, - BD=BD, - NB=NB, - ) - else: - l2norm_fwd_kernel1[(T,)]( - x=x, - y=y, - rstd=rstd, - eps=eps, - D=D, - BD=BD, - ) - return y.view(x_shape_og), rstd.view(x_shape_og[:-1]) - - -def l2norm_bwd( - y: torch.Tensor, - rstd: torch.Tensor, - dy: torch.Tensor, - eps: float = 1e-6, -): - y_shape_og = y.shape - y = y.view(-1, dy.shape[-1]) - dy = dy.view(-1, dy.shape[-1]) - assert dy.shape == y.shape - # allocate output - dx = torch.empty_like(y) - T, D = y.shape[0], y.shape[-1] - # Less than 64KB per feature: enqueue fused kernel - MAX_FUSED_SIZE = 65536 // y.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: - # 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"]),) - - l2norm_bwd_kernel[grid]( - y=y, - rstd=rstd, - dy=dy, - dx=dx, - eps=eps, - T=T, - D=D, - BD=BD, - NB=NB, - ) - else: - l2norm_bwd_kernel1[(T,)]( - y=y, - rstd=rstd, - dy=dy, - dx=dx, - eps=eps, - D=D, - BD=BD, - ) - - return dx.view(y_shape_og) - - -class L2NormFunction(torch.autograd.Function): - @staticmethod - @input_guard - def forward( - ctx, - x, - eps=1e-6, - output_dtype=None, - ): - y, rstd = l2norm_fwd(x, eps, output_dtype) - ctx.eps = eps - ctx.x_dtype = x.dtype - ctx.save_for_backward(y, rstd) - return y - - @staticmethod - @input_guard - def backward(ctx, dy): - y, rstd = ctx.saved_tensors - dx = l2norm_bwd(y, rstd, dy, ctx.eps) - return dx, None, None - - -def l2norm( - x: torch.Tensor, - eps: float = 1e-6, - output_dtype: torch.dtype | None = None, -) -> torch.Tensor: - return L2NormFunction.apply(x, eps, output_dtype) - - -l2_norm = l2norm - - -class L2Norm(nn.Module): - def __init__( - self, - eps: float = 1e-6, - output_dtype: torch.dtype | None = None, - ): - super().__init__() - self.eps = eps - self.output_dtype = output_dtype - - def forward(self, x: torch.Tensor) -> torch.Tensor: - return l2norm(x, self.eps, self.output_dtype) diff --git a/build/torch-xpu/modules/l2warp.py b/build/torch-xpu/modules/l2warp.py deleted file mode 100644 index 3d385823877dac969e46b20d22b17a033f3654a7..0000000000000000000000000000000000000000 --- a/build/torch-xpu/modules/l2warp.py +++ /dev/null @@ -1,51 +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 - - -class L2Wrap(torch.autograd.Function): - r""" - This class of penalty prevents the model from becoming overconfident, - thereby mitigating precision loss in BF16. - - This version is memory-optimized by not storing the full logits tensor. - """ - @staticmethod - def forward( - ctx, - loss: torch.Tensor, - logits: torch.Tensor, - l2_penalty_factor: float = 1e-4, - ) -> torch.Tensor: - """ - Args: - loss (torch.Tensor): - The already-reduced (scalar) loss to wrap. - logits (torch.Tensor): - The logits of shape `[B, T, V]`. - l2_penalty_factor (float, Optional): - The strength of the L2 penalty on the max logit. Default: 1e-4. - """ - maxx, ids = torch.max(logits, dim=-1, keepdim=True) - ctx.logits_shape = logits.shape - factor = l2_penalty_factor / (logits.shape[0] * logits.shape[1]) - maxx = maxx * factor - ctx.save_for_backward(maxx, ids) - 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 diff --git a/build/torch-xpu/modules/layernorm.py b/build/torch-xpu/modules/layernorm.py deleted file mode 100644 index 270f74585de44adc85794e2824986ee2b5be1e46..0000000000000000000000000000000000000000 --- a/build/torch-xpu/modules/layernorm.py +++ /dev/null @@ -1,1472 +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 -# -# Copyright (c) 2023, Tri Dao -# https://github.com/state-spaces/mamba/blob/fb7b5310fa865dbd62aa059b1e26f2b431363e2a/mamba_ssm/ops/triton/layernorm.py -# Implement residual + layer_norm / rms_norm. - -# Based on the Triton LayerNorm tutorial: https://triton-lang.org/main/getting-started/tutorials/05-layer-norm.html -# For the backward pass, we keep weight_grad and bias_grad in registers and accumulate. -# This is faster for dimensions up to 8k, but after that it's much slower due to register spilling. -# The models we train have hidden dim up to 8k anyway (e.g. Llama 70B), so this is fine. - -from __future__ import annotations - -from functools import partial - -import torch -import torch.nn as nn -import torch.nn.functional as F -import triton -import triton.language as tl -from einops import rearrange -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 ..utils import autotune_cache_kwargs, get_multiprocessor_count, input_guard - -try: - from torch.distributed.tensor import DTensor -except (ImportError, AttributeError): - DTensor = None - - -def layer_norm_ref( - x: torch.Tensor, - weight: torch.Tensor, - bias: torch.Tensor, - residual: torch.Tensor = None, - eps: float = 1e-5, - prenorm: bool = False, - upcast: bool = False, -): - dtype = x.dtype - if upcast: - weight = weight.float() - bias = bias.float() if bias is not None else None - if upcast: - x = x.float() - residual = residual.float() if residual is not None else residual - if residual is not None: - x = (x + residual).to(x.dtype) - out = F.layer_norm(x.to(weight.dtype), x.shape[-1:], weight=weight, bias=bias, eps=eps).to( - dtype, - ) - return out if not prenorm else (out, x) - - -def rms_norm_ref( - x: torch.Tensor, - weight: torch.Tensor, - bias: torch.Tensor, - residual: torch.Tensor = None, - eps: float = 1e-5, - prenorm: bool = False, - upcast: bool = False, -): - dtype = x.dtype - if upcast: - weight = weight.float() - bias = bias.float() if bias is not None else None - if upcast: - x = x.float() - residual = residual.float() if residual is not None else residual - if residual is not None: - x = (x + residual).to(x.dtype) - rstd = 1 / torch.sqrt((x.square()).mean(dim=-1, keepdim=True) + eps) - out = (x * rstd * weight) + bias if bias is not None else (x * rstd * weight) - out = out.to(dtype) - return out if not prenorm else (out, x) - - -def group_norm_ref( - x: torch.Tensor, - weight: torch.Tensor, - bias: torch.Tensor, - num_groups: int, - residual: torch.Tensor = None, - eps: float = 1e-5, - is_rms_norm: bool = False, - prenorm: bool = False, - upcast: bool = False, -): - dtype = x.dtype - if upcast: - weight = weight.float() - bias = bias.float() if bias is not None else None - if upcast: - x = x.float() - residual = residual.float() if residual is not None else residual - if residual is not None: - x = (x + residual).to(x.dtype) - residual = x - x, weight = [ - rearrange(data, "... (g d) -> ... g d", g=num_groups) for data in (x, weight) - ] - if bias is not None: - bias = rearrange(bias, '... (g d) -> ... g d', g=num_groups) - if not is_rms_norm: - mean = x.mean(dim=-1, keepdim=True) - x = x - mean - rstd = 1 / torch.sqrt((x.square()).mean(dim=-1, keepdim=True) + eps) - out = (x * rstd * weight) + bias if bias is not None else (x * rstd * weight) - out = rearrange(out, "... g d -> ... (g d)") - out = out.to(dtype) - return out if not prenorm else (out, residual) - - -class GroupNormRef(nn.Module): - - def __init__( - self, - num_groups: int, - hidden_size: int, - elementwise_affine: bool = True, - bias: bool = False, - eps: float = 1e-5, - is_rms_norm: bool = False, - ) -> GroupNormRef: - super().__init__() - - if hidden_size % num_groups != 0: - raise ValueError('num_channels must be divisible by num_groups') - - self.num_groups = num_groups - self.hidden_size = hidden_size - self.elementwise_affine = elementwise_affine - self.eps = eps - self.is_rms_norm = is_rms_norm - - self.register_parameter("weight", None) - self.register_parameter("bias", None) - if elementwise_affine: - self.weight = nn.Parameter(torch.empty(hidden_size)) - if bias: - self.bias = nn.Parameter(torch.empty(hidden_size)) - - 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.num_groups}, {self.hidden_size}" - if not self.elementwise_affine: - s += f", elementwise_affine={self.elementwise_affine}" - if self.is_rms_norm: - s += f", is_rms_norm={self.is_rms_norm}" - s += f", eps={self.eps}" - s += ")" - return s - - def forward(self, x, residual=None, prenorm=False): - return group_norm_ref( - x, - self.weight, - self.bias, - num_groups=self.num_groups, - residual=residual, - eps=self.eps, - is_rms_norm=self.is_rms_norm, - prenorm=prenorm, - upcast=True, - ) - - -@triton.autotune( - configs=[ - triton.Config({'BT': BT}, num_warps=num_warps) - for BT in [32, 64, 128] - for num_warps in [2, 4, 8] - ], - key=['D', 'NB', 'HAS_RESIDUAL', 'STORE_RESIDUAL_OUT', 'IS_RMS_NORM'], - **autotune_cache_kwargs, -) -@triton.jit -def layer_norm_fwd_kernel( - x, # pointer to the input - y, # pointer to the output - w, # pointer to the weights - b, # pointer to the biases - res, # pointer to the res - res_out, # pointer to the res - mean, # pointer to the mean - rstd, # pointer to the 1/std - eps, # epsilon to avoid division by zero - T, - G: tl.constexpr, - D: tl.constexpr, - BT: tl.constexpr, - BD: tl.constexpr, - NB: 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) - - o_t = i_t * BT + tl.arange(0, BT) - o_g = o_t % G - 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(res, (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(res_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_g[:, None] * D + o_d[None, :], mask=m_d[None, :]).to(tl.float32) - if HAS_BIAS: - b_b = tl.load(b + o_g[:, None] * D + o_d[None, :], mask=m_d[None, :]).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 if HAS_WEIGHT else b_x_hat - if HAS_BIAS: - b_y = b_y + b_b - - # 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.autotune( - configs=[ - triton.Config({}, num_warps=num_warps) - for num_warps in [2, 4, 8, 16] - ], - key=['D', 'HAS_RESIDUAL', 'STORE_RESIDUAL_OUT', 'IS_RMS_NORM'], - **autotune_cache_kwargs, -) -@triton.jit -def layer_norm_fwd_kernel1( - x, # pointer to the input - y, # pointer to the output - w, # pointer to the weights - b, # pointer to the biases - res, # pointer to the res - res_out, # pointer to the res - mean, # pointer to the mean - rstd, # pointer to the 1/std - eps, # epsilon to avoid division by zero - 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 - - # Write output - tl.store(y + o_d, b_y, mask=m_d) - - -@triton.heuristics({ - 'RECOMPUTE_OUTPUT': lambda args: args['y'] is not None, -}) -@triton.autotune( - configs=[ - triton.Config({'BT': BT}, num_warps=num_warps) - for BT in [32, 64] - for num_warps in [2, 4, 8] - ], - key=['D', 'NB', 'HAS_DRESIDUAL', 'STORE_DRESIDUAL', 'IS_RMS_NORM'], - **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 - dres, - dres_in, - mean, - rstd, - T, - G: tl.constexpr, - D: tl.constexpr, - BS: tl.constexpr, - BT: tl.constexpr, - BD: tl.constexpr, - NB: 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) - m_d = o_d < D - if HAS_WEIGHT: - b_w = tl.load(w + i_g * D + 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 + i_g * D + o_d, mask=m_d, other=0.0).to(tl.float32) - b_db = tl.zeros((BT, BD), dtype=tl.float32) - - # Tg: number of tokens per group, used as the logical shape for make_block_ptr. - # for mean/rstd with shape (T,) and stride (G,), the strided view has Tg elements per group. - # the caller guarantees NS capped so every program has work. - # the last program's range may slightly exceed Tg (since BS = cdiv(T, NS)); - # boundary_check handles the partial tail tile, m_t < Tg masks dw/db accumulation. - Tg = T // G - for i_t in range(i_sg * BS, i_sg * BS + BS, BT): - p_x = tl.make_block_ptr(x + i_g * D, (Tg, D), (G*D, 1), (i_t, 0), (BT, BD), (1, 0)) - p_dy = tl.make_block_ptr(dy + i_g * D, (Tg, D), (G*D, 1), (i_t, 0), (BT, BD), (1, 0)) - p_dx = tl.make_block_ptr(dx + i_g * D, (Tg, D), (G*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_dy = tl.load(p_dy, boundary_check=(0, 1)).to(tl.float32) - - if not IS_RMS_NORM: - p_mean = tl.make_block_ptr(mean + i_g, (Tg,), (G,), (i_t,), (BT,), (0,)) - b_mean = tl.load(p_mean, boundary_check=(0,)) - p_rstd = tl.make_block_ptr(rstd + i_g, (Tg,), (G,), (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 + i_g * D, (Tg, D), (G*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_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_sg * BS + BS, Tg) - 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(dres + i_g * D, (Tg, D), (G*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(dres_in + i_g * D, (Tg, D), (G*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)) - - 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({ - '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] - ], - key=['D', 'HAS_DRESIDUAL', 'STORE_DRESIDUAL', 'IS_RMS_NORM'], - **autotune_cache_kwargs, -) -@triton.jit -def layer_norm_bwd_kernel1( - 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 - 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) - # 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) - 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 - # Write dx - 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) - - -@dispatch('modules') -def layer_norm_fwd( - 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,) - # 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): - 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) - # 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: - NB = triton.cdiv(T, 2048) - def grid(meta): return (triton.cdiv(T, meta['BT']), ) - layer_norm_fwd_kernel[grid]( - x, - y, - weight, - bias, - residual, - res_out, - mean, - rstd, - eps, - T=T, - G=G, - D=D, - BD=BD, - NB=NB, - 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, - ) - else: - layer_norm_fwd_kernel1[(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, - ) - # res_out is None if residual is None and residual_dtype == input_dtype - return y, mean, rstd, res_out if res_out is not None else x - - -@dispatch('modules') -def layer_norm_bwd( - 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,) - # allocate output - 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 - - # 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.") - # each program handles one group only. - # cap per-group program count to T // G 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(triton.cdiv(get_multiprocessor_count(x.device.index), G), T // G) * G - BS = triton.cdiv(T, NS) - GS = NS // G - - 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: - NB = triton.cdiv(T, 2048) - layer_norm_bwd_kernel[grid]( - x, - weight, - bias, - y, - dy, - dx, - dw, - db, - dres, - dres_in, - mean, - rstd, - T=T, - G=G, - D=D, - BS=BS, - BD=BD, - NB=NB, - 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, - ) - else: - 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 - # Don't need to compute dres_in separately in this case - 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) - - -class LayerNormFunction(torch.autograd.Function): - - @staticmethod - @input_guard - def forward( - ctx, - x, - weight, - bias, - residual: torch.Tensor = None, - eps: float = 1e-5, - prenorm: bool = False, - residual_in_fp32: bool = False, - is_rms_norm: bool = False, - num_groups: int = 1, - ): - x_shape_og = x.shape - - if x.shape[-1] % num_groups != 0: - raise ValueError('num_channels must be divisible by num_groups') - # reshape input data into 2D tensor - x = x.reshape(-1, (x.shape[-1] // num_groups)) - if residual is not None: - assert residual.shape == x_shape_og - residual = residual.reshape_as(x) - residual_dtype = ( - residual.dtype - if residual is not None - else (torch.float32 if residual_in_fp32 else None) - ) - y, mean, rstd, res_out = layer_norm_fwd( - x, - weight, - bias, - eps, - residual, - residual_dtype=residual_dtype, - is_rms_norm=is_rms_norm, - num_groups=num_groups, - ) - ctx.save_for_backward(res_out, weight, bias, mean, rstd) - ctx.x_shape_og = x_shape_og - ctx.eps = eps - ctx.is_rms_norm = is_rms_norm - ctx.num_groups = num_groups - 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, res_out.reshape(x_shape_og)) - - @staticmethod - @input_guard - def backward(ctx, dy, *args): - x, weight, bias, mean, rstd = ctx.saved_tensors - dy = dy.reshape(-1, (dy.shape[-1] // ctx.num_groups)) - assert dy.shape == x.shape - if ctx.prenorm: - dresidual = args[0] - dresidual = dresidual.reshape(-1, x.shape[-1]) - assert dresidual.shape == x.shape - else: - dresidual = None - dx, dw, db, dresidual_in = layer_norm_bwd( - dy, - x, - weight, - bias, - mean, - rstd, - dresidual, - ctx.has_residual, - ctx.is_rms_norm, - x_dtype=ctx.x_dtype, - num_groups=ctx.num_groups, - ) - return ( - dx.reshape(ctx.x_shape_og), - dw, - db, - dresidual_in.reshape(ctx.x_shape_og) if ctx.has_residual else None, - None, - None, - None, - None, - None, - ) - - -def layer_norm( - x: torch.Tensor, - weight: torch.Tensor, - bias: torch.Tensor, - residual: torch.Tensor = None, - eps: float = 1e-5, - prenorm: bool = False, - residual_in_fp32: bool = False, - is_rms_norm: bool = False, -): - return LayerNormFunction.apply( - x, - weight, - bias, - residual, - eps, - prenorm, - residual_in_fp32, - is_rms_norm, - ) - - -def group_norm( - x: torch.Tensor, - weight: torch.Tensor, - bias: torch.Tensor, - residual: torch.Tensor = None, - eps: float = 1e-5, - prenorm: bool = False, - residual_in_fp32: bool = False, - is_rms_norm: bool = False, - num_groups: int = 1, -): - return LayerNormFunction.apply( - x, - weight, - bias, - residual, - eps, - prenorm, - residual_in_fp32, - is_rms_norm, - num_groups, - ) - - -def rms_norm( - x: torch.Tensor, - weight: torch.Tensor, - bias: torch.Tensor, - residual: torch.Tensor = None, - eps: float = 1e-5, - prenorm: bool = False, - residual_in_fp32: bool = False, -): - return LayerNormFunction.apply( - x, - weight, - bias, - residual, - eps, - prenorm, - residual_in_fp32, - True, - ) - - -def layer_norm_linear( - 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, - is_rms_norm: bool = False, - num_groups: int = 1, -): - return LayerNormLinearFunction.apply( - x, - norm_weight, - norm_bias, - linear_weight, - linear_bias, - residual, - eps, - prenorm, - residual_in_fp32, - is_rms_norm, - num_groups, - ) - - -def rms_norm_linear( - 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( - x=x, - norm_weight=norm_weight, - norm_bias=norm_bias, - linear_weight=linear_weight, - linear_bias=linear_bias, - residual=residual, - eps=eps, - prenorm=prenorm, - residual_in_fp32=residual_in_fp32, - is_rms_norm=True, - ) - - -def group_norm_linear( - 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, - is_rms_norm: bool = False, - num_groups: int = 1, -): - return layer_norm_linear( - x=x, - norm_weight=norm_weight, - norm_bias=norm_bias, - linear_weight=linear_weight, - linear_bias=linear_bias, - residual=residual, - eps=eps, - prenorm=prenorm, - residual_in_fp32=residual_in_fp32, - is_rms_norm=is_rms_norm, - num_groups=num_groups, - ) - - -class LayerNorm(nn.Module): - - 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, - ) -> LayerNorm: - 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)) - 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 += ")" - return s - - def forward(self, x, residual=None, prenorm=False, residual_in_fp32=False): - return layer_norm( - x, - self.weight, - self.bias, - residual=residual, - eps=self.eps, - prenorm=prenorm, - residual_in_fp32=residual_in_fp32, - ) - - -class GroupNorm(nn.Module): - - def __init__( - self, - num_groups: int, - hidden_size: int, - elementwise_affine: bool = True, - bias: bool = False, - eps: float = 1e-5, - is_rms_norm: bool = False, - device: torch.device | None = None, - dtype: torch.dtype | None = None, - ) -> GroupNorm: - factory_kwargs = {"device": device, "dtype": dtype} - super().__init__() - - if hidden_size % num_groups != 0: - raise ValueError('num_channels must be divisible by num_groups') - - self.num_groups = num_groups - self.hidden_size = hidden_size - self.elementwise_affine = elementwise_affine - self.eps = eps - self.is_rms_norm = is_rms_norm - - 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.num_groups}, {self.hidden_size}" - if not self.elementwise_affine: - s += f", elementwise_affine={self.elementwise_affine}" - if self.is_rms_norm: - s += f", is_rms_norm={self.is_rms_norm}" - s += f", eps={self.eps}" - s += ")" - return s - - def forward(self, x, residual=None, prenorm=False, residual_in_fp32=False): - return group_norm( - x, - self.weight, - self.bias, - residual=residual, - eps=self.eps, - prenorm=prenorm, - residual_in_fp32=residual_in_fp32, - is_rms_norm=self.is_rms_norm, - num_groups=self.num_groups, - ) - - -class RMSNorm(nn.Module): - - 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, - ) -> RMSNorm: - 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)) - 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 += ")" - return s - - def forward(self, x, residual=None, prenorm=False, residual_in_fp32=False): - return rms_norm( - x, - self.weight, - self.bias, - residual=residual, - eps=self.eps, - prenorm=prenorm, - residual_in_fp32=residual_in_fp32, - ) - - -class LayerNormLinearFunction(torch.autograd.Function): - - @staticmethod - @input_guard - def forward( - ctx, - x, - norm_weight, - norm_bias, - linear_weight, - linear_bias, - residual=None, - eps=1e-5, - prenorm=False, - residual_in_fp32=False, - is_rms_norm=False, - num_groups=1, - ): - x_shape_og = x.shape - - if x.shape[-1] % num_groups != 0: - raise ValueError('num_channels must be divisible by num_groups') - # reshape input data into 2D tensor - x = x.reshape(-1, (x.shape[-1] // num_groups)) - if residual is not None: - assert residual.shape == x_shape_og - residual = residual.reshape_as(x) - residual_dtype = ( - residual.dtype - if residual is not None - else (torch.float32 if residual_in_fp32 else None) - ) - y, mean, rstd, res_out = layer_norm_fwd( - 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, - num_groups=num_groups, - ) - 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(res_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.num_groups = num_groups - 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, res_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()) - dy = dy.reshape(-1, (dy.shape[-1] // ctx.num_groups)) - 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, x.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, - mean, - rstd, - dresidual, - ctx.has_residual, - ctx.is_rms_norm, - x_dtype=ctx.x_dtype, - recompute_output=True, - num_groups=ctx.num_groups, - ) - dlinear_weight = torch.einsum("bo,bi->oi", dout, y.view(-1, linear_weight.shape[-1])) - 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, - None, - ) - - -class LayerNormLinear(nn.Module): - - def __init__( - self, - hidden_size, - elementwise_affine: bool = True, - bias: bool = False, - eps: float = 1e-5, - device: torch.device | None = None, - dtype: torch.dtype | None = None, - ) -> LayerNormLinear: - 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)) - 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 += ")" - return s - - def forward(self, x, weight, bias, residual=None, prenorm=False, residual_in_fp32=False): - return layer_norm_linear( - x=x, - norm_weight=self.weight, - norm_bias=self.bias, - linear_weight=weight, - linear_bias=bias, - residual=residual, - eps=self.eps, - prenorm=prenorm, - residual_in_fp32=residual_in_fp32, - is_rms_norm=False, - ) - - -class GroupNormLinear(nn.Module): - - def __init__( - self, - num_groups: int, - hidden_size: int, - elementwise_affine: bool = True, - bias: bool = False, - eps: float = 1e-5, - is_rms_norm: bool = False, - device: torch.device | None = None, - dtype: torch.dtype | None = None, - ) -> GroupNormLinear: - factory_kwargs = {"device": device, "dtype": dtype} - super().__init__() - - if hidden_size % num_groups != 0: - raise ValueError('num_channels must be divisible by num_groups') - - self.num_groups = num_groups - self.hidden_size = hidden_size - self.elementwise_affine = elementwise_affine - self.eps = eps - self.is_rms_norm = is_rms_norm - - 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.num_groups}, {self.hidden_size}" - if not self.elementwise_affine: - s += f", elementwise_affine={self.elementwise_affine}" - if self.is_rms_norm: - s += f", is_rms_norm={self.is_rms_norm}" - s += f", eps={self.eps}" - s += ")" - return s - - def forward(self, x, weight, bias, residual=None, prenorm=False, residual_in_fp32=False): - return layer_norm_linear( - x=x, - norm_weight=self.weight, - norm_bias=self.bias, - linear_weight=weight, - linear_bias=bias, - residual=residual, - eps=self.eps, - prenorm=prenorm, - residual_in_fp32=residual_in_fp32, - is_rms_norm=self.is_rms_norm, - num_groups=self.num_groups, - ) - - -class RMSNormLinear(nn.Module): - - def __init__( - self, - hidden_size, - elementwise_affine: bool = True, - bias: bool = False, - eps: float = 1e-5, - device: torch.device | None = None, - dtype: torch.dtype | None = None, - ) -> RMSNormLinear: - 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)) - 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 += ")" - return s - - def forward(self, x, weight, bias, residual=None, prenorm=False, residual_in_fp32=False): - return layer_norm_linear( - x=x, - norm_weight=self.weight, - norm_bias=self.bias, - linear_weight=weight, - linear_bias=bias, - residual=residual, - eps=self.eps, - prenorm=prenorm, - residual_in_fp32=residual_in_fp32, - is_rms_norm=True, - ) - - -class NormParallel(ParallelStyle): - - def __init__(self, *, sequence_dim: int = 1, use_local_output: bool = False): - super().__init__() - self.sequence_sharding = (Shard(sequence_dim),) - self.use_local_output = use_local_output - - def _replicate_module_fn( - self, name: str, module: nn.Module, device_mesh: DeviceMesh, - ): - for p_name, param in module.named_parameters(): - # simple replication with fixed ones_ init from LayerNorm/RMSNorm, which allow - # us to simply just use from_local - replicated_param = torch.nn.Parameter( - DTensor.from_local(param, device_mesh, [Replicate()], run_check=False), - ) - module.register_parameter(p_name, replicated_param) - - @staticmethod - def _prepare_input_fn(sequence_sharding, mod, inputs, device_mesh): - input_tensor = inputs[0] - if isinstance(input_tensor, DTensor): - # if the passed in input DTensor is not sharded on the sequence dim, we need to redistribute it - if input_tensor.placements != sequence_sharding: - input_tensor = input_tensor.redistribute( - placements=sequence_sharding, async_op=True, - ) - return input_tensor - elif isinstance(input_tensor, torch.Tensor): - # assume the input passed in already sharded on the sequence dim and create the DTensor - return DTensor.from_local( - input_tensor, device_mesh, sequence_sharding, run_check=False, - ) - else: - raise ValueError( - f"expecting input of {mod} to be a torch.Tensor or DTensor, but got {input_tensor}", - ) - - @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( - module, - device_mesh, - self._replicate_module_fn, - partial(self._prepare_input_fn, self.sequence_sharding), - partial(self._prepare_output_fn, self.use_local_output), - ) diff --git a/build/torch-xpu/modules/layernorm_gated.py b/build/torch-xpu/modules/layernorm_gated.py deleted file mode 100644 index ff0eac089649d8a784aed5f33b642d116bf72ffd..0000000000000000000000000000000000000000 --- a/build/torch-xpu/modules/layernorm_gated.py +++ /dev/null @@ -1,535 +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 -# -# Copyright (c) 2024, Tri Dao. -# -# Based on the Triton LayerNorm tutorial: https://triton-lang.org/main/getting-started/tutorials/05-layer-norm.html -# For the backward pass, we keep weight_grad and bias_grad in registers and accumulate. -# This backward pass is faster for dimensions up to 8k, but after that it's much slower due to register spilling. -# The models we train have hidden dim up to 8k anyway (e.g. Llama 70B), so this is fine. - -import math - -import torch -import torch.nn as nn -import torch.nn.functional as F -import triton -import triton.language as tl -from einops import rearrange - -from ..utils import get_multiprocessor_count, input_guard - - -def rms_norm_ref(x, weight, bias, z=None, eps=1e-6, group_size=None, norm_before_gate=True, upcast=True): - dtype = x.dtype - weight = weight.float() - bias = bias.float() if bias is not None else None - if upcast: - x = x.float() - z = z.float() if z is not None else z - if z is not None and not norm_before_gate: - x = x * F.silu(z) - if group_size is None: - rstd = 1 / torch.sqrt((x.square()).mean(dim=-1, keepdim=True) + eps) - out = (x * rstd * weight) + bias if bias is not None else (x * rstd * weight) - else: - x_group = rearrange(x, "... (g d) -> ... g d", d=group_size) - rstd = 1 / torch.sqrt((x_group.square()).mean(dim=-1, keepdim=True) + eps) - out = rearrange(x_group * rstd, "... g d -> ... (g d)") * weight - if bias is not None: - out = out + bias - if z is not None and norm_before_gate: - out *= F.silu(z) - return out.to(dtype) - - -@triton.heuristics({ - "HAS_BIAS": lambda args: args["B"] is not None, - "HAS_Z": lambda args: args["Z"] is not None, -}) -@triton.jit -def layer_norm_fwd_kernel( - X, # pointer to the input - Y, # pointer to the output - W, # pointer to the weights - B, # pointer to the biases - Z, # pointer to the other branch - 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_z_row, - M, # number of rows in X - N, # number of columns in X - eps, # epsilon to avoid division by zero - BLOCK_N: tl.constexpr, - HAS_BIAS: tl.constexpr, - HAS_Z: tl.constexpr, - NORM_BEFORE_GATE: tl.constexpr, - IS_RMS_NORM: tl.constexpr, -): - # Map the program id to the row of X and Y it should compute. - row = tl.program_id(0) - group = tl.program_id(1) - X += row * stride_x_row + group * N - Y += row * stride_y_row + group * N - if HAS_Z: - Z += row * stride_z_row + group * N - if not IS_RMS_NORM: - Mean += group * M - Rstd += group * M - W += group * N - if HAS_BIAS: - B += group * N - # Compute mean and variance - cols = tl.arange(0, BLOCK_N) - x = tl.load(X + cols, mask=cols < N, other=0.).to(tl.float32) - if HAS_Z and not NORM_BEFORE_GATE: - z = tl.load(Z + cols, mask=cols < N).to(tl.float32) - x *= z * tl.sigmoid(z) - 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.) - var = tl.sum(xbar * xbar, axis=0) / N - else: - xbar = tl.where(cols < N, x, 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 - 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 + b if HAS_BIAS else x_hat * w - if HAS_Z and NORM_BEFORE_GATE: - z = tl.load(Z + cols, mask=mask).to(tl.float32) - y *= z * tl.sigmoid(z) - # Write output - tl.store(Y + cols, y, mask=mask) - - -def layer_norm_fwd( - x: torch.Tensor, - weight: torch.Tensor, - bias: torch.Tensor, - eps: float, - z: torch.Tensor = None, - out: torch.Tensor = None, - group_size: int = None, - norm_before_gate: bool = True, - is_rms_norm: bool = False, -): - M, N = x.shape - if group_size is None: - group_size = N - assert N % group_size == 0 - ngroups = N // group_size - assert x.stride(-1) == 1 - if z is not None: - assert z.stride(-1) == 1 - assert z.shape == (M, N) - assert weight.shape == (N,) - assert weight.stride(-1) == 1 - if bias is not None: - assert bias.stride(-1) == 1 - assert bias.shape == (N,) - # allocate output - if out is not None: - assert out.shape == x.shape - else: - out = torch.empty_like(x) - assert out.stride(-1) == 1 - mean = torch.empty((ngroups * M, ), dtype=torch.float32, device=x.device) if not is_rms_norm else None - rstd = torch.empty((ngroups * 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(group_size)) - if group_size > BLOCK_N: - raise RuntimeError("This layer norm doesn't support feature dim >= 64KB.") - # heuristics for number of warps - num_warps = min(max(BLOCK_N // 256, 1), 8) - grid = (M, ngroups) - layer_norm_fwd_kernel[grid]( - x, - out, - weight, - bias, - z, - mean, - rstd, - x.stride(0), - out.stride(0), - z.stride(0) if z is not None else 0, - M, - group_size, - eps, - BLOCK_N=BLOCK_N, - NORM_BEFORE_GATE=norm_before_gate, - IS_RMS_NORM=is_rms_norm, - num_warps=num_warps, - ) - return out, mean, rstd - - -@triton.heuristics({ - "HAS_BIAS": lambda args: args["B"] is not None, - "HAS_Z": lambda args: args["Z"] is not None, - "RECOMPUTE_OUTPUT": lambda args: args["Y"] is not None, -}) -@triton.jit -def layer_norm_bwd_kernel( - X, # pointer to the input - W, # pointer to the weights - B, # pointer to the biases - Z, # pointer to the other branch - 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 - DZ, # pointer to the other branch - 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_z_row, - stride_y_row, - stride_dy_row, - stride_dx_row, - stride_dz_row, - stride_dw_row, - stride_db_row, - M, # number of rows in X - N, # number of columns in X - eps, # epsilon to avoid division by zero - rows_per_program, - NORM_BEFORE_GATE: tl.constexpr, - IS_RMS_NORM: tl.constexpr, - HAS_BIAS: tl.constexpr, - HAS_Z: tl.constexpr, - RECOMPUTE_OUTPUT: tl.constexpr, - BLOCK_N: tl.constexpr, -): - # Map the program id to the elements of X, DX, and DY it should compute. - row_block_id = tl.program_id(0) - group = tl.program_id(1) - row_start = row_block_id * rows_per_program - cols = tl.arange(0, BLOCK_N) - mask = cols < N - X += row_start * stride_x_row + group * N - if HAS_Z: - Z += row_start * stride_z_row + group * N - DZ += row_start * stride_dz_row + group * N - DY += row_start * stride_dy_row + group * N - DX += row_start * stride_dx_row + group * N - if RECOMPUTE_OUTPUT: - Y += row_start * stride_y_row + group * N - if not IS_RMS_NORM: - Mean += group * M - Rstd += group * M - W += group * N - w = tl.load(W + cols, mask=mask).to(tl.float32) - if (RECOMPUTE_OUTPUT or HAS_Z) and HAS_BIAS: - B += group * N - b = tl.load(B + cols, mask=mask, other=0.).to(tl.float32) - dw = tl.zeros((BLOCK_N,), dtype=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) - if HAS_Z and not NORM_BEFORE_GATE: - z = tl.load(Z + cols, mask=mask, other=0.).to(tl.float32) - x_og = x - x = x_og * z * tl.sigmoid(z) - 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.) - if HAS_Z and NORM_BEFORE_GATE: - z = tl.load(Z + cols, mask=mask, other=0.).to(tl.float32) - z_sigmoid = tl.sigmoid(z) - y = xhat * w + b if HAS_BIAS else xhat * w - if RECOMPUTE_OUTPUT: - tl.store(Y + cols, y * z * z_sigmoid, mask=mask) - dz = dy * y * z_sigmoid * (1 + z * (1 - z_sigmoid)) - tl.store(DZ + cols, dz, mask=mask) - dy *= z * z_sigmoid - else: - if RECOMPUTE_OUTPUT: - y = xhat * w + b if HAS_BIAS else xhat * w - tl.store(Y + cols, y, mask=mask) - wdy = w * dy - c1 = tl.sum(xhat * wdy, axis=0) / N - if not IS_RMS_NORM: - c2 = tl.sum(wdy, axis=0) / N - dx = (wdy - (xhat * c1 + c2)) * rstd - else: - dx = (wdy - xhat * c1) * rstd - dw += dy * xhat - if HAS_BIAS: - db += dy - if HAS_Z and not NORM_BEFORE_GATE: - z_sigmoid = tl.sigmoid(z) - dz = dx * x_og * z_sigmoid * (1 + z * (1 - z_sigmoid)) - tl.store(DZ + cols, dz, mask=mask) - dx *= z * z_sigmoid - # Write dx - tl.store(DX + cols, dx, mask=mask) - - X += stride_x_row - if HAS_Z: - Z += stride_z_row - DZ += stride_dz_row - if RECOMPUTE_OUTPUT: - Y += stride_y_row - DY += stride_dy_row - DX += stride_dx_row - tl.store(DW + row_block_id * stride_dw_row + group * N + cols, dw, mask=mask) - if HAS_BIAS: - tl.store(DB + row_block_id * stride_db_row + group * 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, - z: torch.Tensor = None, - group_size: int = None, - norm_before_gate: bool = True, - is_rms_norm: bool = False, - recompute_output: bool = False, - dz: torch.Tensor = None, - out: torch.Tensor = None, -): - M, N = x.shape - if group_size is None: - group_size = N - assert N % group_size == 0 - ngroups = N // group_size - assert x.stride(-1) == 1 - assert dy.stride(-1) == 1 - assert dy.shape == (M, N) - if z is not None: - assert z.stride(-1) == 1 - assert z.shape == (M, N) - assert weight.shape == (N,) - assert weight.stride(-1) == 1 - if bias is not None: - assert bias.stride(-1) == 1 - assert bias.shape == (N,) - # allocate output - dx = torch.empty_like(x) - if dz is not None: - assert z is not None - assert dz.shape == z.shape - assert dz.stride(-1) == 1 - else: - dz = torch.empty_like(z) if z is not None else None - if recompute_output: - if out is None: - out = torch.empty_like(x) - assert out.shape == x.shape - - # 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(group_size)) - if group_size > BLOCK_N: - raise RuntimeError("This layer norm doesn't support feature dim >= 64KB.") - # heuristics for number of warps - num_warps = min(max(BLOCK_N // 256, 1), 8) - sm_count = get_multiprocessor_count(x.device.index) - # If group size is small (e.g., 64), we're only using 1 warp. So having just 108 programs - # would limit the occupancy. - nrow_groups = math.ceil(sm_count * math.ceil(4 / num_warps) / ngroups) - _dw = torch.empty((nrow_groups, N), dtype=torch.float32, device=weight.device) - _db = torch.empty((nrow_groups, N), dtype=torch.float32, device=bias.device) if bias is not None else None - rows_per_program = math.ceil(M / nrow_groups) - grid = (nrow_groups, ngroups) - layer_norm_bwd_kernel[grid]( - x, - weight, - bias, - z, - out if recompute_output else None, - dy, - dx, - _dw, - _db, - dz, - mean, - rstd, - x.stride(0), - z.stride(0) if z is not None else 0, - 0 if not recompute_output else out.stride(0), - dy.stride(0), - dx.stride(0), - dz.stride(0) if dz is not None else 0, - _dw.stride(0), - _db.stride(0) if _db is not None else 0, - M, group_size, eps, - rows_per_program, - BLOCK_N=BLOCK_N, - NORM_BEFORE_GATE=norm_before_gate, - IS_RMS_NORM=is_rms_norm, - num_warps=num_warps, - ) - dw = _dw.sum(0).to(weight.dtype) - db = _db.sum(0).to(bias.dtype) if bias is not None else None - return (dx, dw, db, dz) if not recompute_output else (dx, dw, db, dz, out) - - -class LayerNormFn(torch.autograd.Function): - - @input_guard - @staticmethod - def forward(ctx, x, weight, bias, z=None, eps=1e-6, group_size=None, norm_before_gate=True, - is_rms_norm=False): - """If z is not None, we do norm(x) * silu(z) if norm_before_gate, else norm(x * silu(z)) - """ - - x_shape_og = x.shape - # reshape input data into 2D tensor - x = x.reshape(-1, x.shape[-1]) - if x.stride(-1) != 1: - x = x.contiguous() - if z is not None: - assert z.shape == x_shape_og - z = z.reshape(-1, z.shape[-1]) - if z.stride(-1) != 1: - z = z.contiguous() - weight = weight.contiguous() - if bias is not None: - bias = bias.contiguous() - y, mean, rstd = layer_norm_fwd( - x, - weight, - bias, - eps, - z=z, - group_size=group_size, - norm_before_gate=norm_before_gate, - is_rms_norm=is_rms_norm, - ) - ctx.save_for_backward(x, weight, bias, mean, rstd, z) - ctx.x_shape_og = x_shape_og - ctx.eps = eps - ctx.group_size = group_size - ctx.norm_before_gate = norm_before_gate - ctx.is_rms_norm = is_rms_norm - return y.reshape(x_shape_og) - - @input_guard - @staticmethod - def backward(ctx, dy): - x, weight, bias, mean, rstd, z = ctx.saved_tensors - dy = dy.reshape(-1, dy.shape[-1]) - if dy.stride(-1) != 1: - dy = dy.contiguous() - assert dy.shape == x.shape - dx, dw, db, dz = layer_norm_bwd( - dy, - x, - weight, - bias, - ctx.eps, - mean, - rstd, - z, - ctx.group_size, - ctx.norm_before_gate, - ctx.is_rms_norm, - ) - dx = dx.reshape(ctx.x_shape_og) - dz = dz.reshape(ctx.x_shape_og) if dz is not None else None - return dx, dw, db, dz, None, None, None, None - - -def layernorm_fn(x, weight, bias, z=None, eps=1e-6, group_size=None, norm_before_gate=True, is_rms_norm=False): - return LayerNormFn.apply(x, weight, bias, z, eps, group_size, norm_before_gate, is_rms_norm) - - -def rmsnorm_fn(x, weight, bias, z=None, eps=1e-6, group_size=None, norm_before_gate=True): - return LayerNormFn.apply(x, weight, bias, z, eps, group_size, norm_before_gate, True) - - -class LayerNormGated(nn.Module): - - def __init__( - self, - hidden_size, - eps: float = 1e-5, - group_size: int | None = None, - norm_before_gate: bool = True, - device: torch.device | None = None, - dtype: torch.dtype | None = None, - ): - """If group_size is not None, we do GroupNorm with each group having group_size elements. - group_size=None is equivalent to group_size=hidden_size (i.e. there's only 1 group). - """ - - factory_kwargs = {"device": device, "dtype": dtype} - super().__init__() - self.eps = eps - self.weight = nn.Parameter(torch.empty(hidden_size, **factory_kwargs)) - self.bias = nn.Parameter(torch.empty(hidden_size, **factory_kwargs)) - self.group_size = group_size - self.norm_before_gate = norm_before_gate - self.reset_parameters() - - def reset_parameters(self): - torch.nn.init.ones_(self.weight) - torch.nn.init.zeros_(self.bias) - - def forward(self, x, z=None): - """If z is not None, we do norm(x) * silu(z) if norm_before_gate, else norm(x * silu(z)) - """ - return layernorm_fn(x, self.weight, self.bias, z=z, group_size=self.group_size, eps=self.eps, - norm_before_gate=self.norm_before_gate) - - -class RMSNormGated(nn.Module): - - def __init__( - self, - hidden_size, - eps: float = 1e-5, - group_size: int | None = None, - norm_before_gate: bool = False, - device: torch.device | None = None, - dtype: torch.dtype | None = None, - ): - """If group_size is not None, we do GroupNorm with each group having group_size elements. - group_size=None is equivalent to group_size=hidden_size (i.e. there's only 1 group). - """ - factory_kwargs = {"device": device, "dtype": dtype} - super().__init__() - self.eps = eps - self.weight = nn.Parameter(torch.empty(hidden_size, **factory_kwargs)) - self.register_parameter("bias", None) - self.group_size = group_size - self.norm_before_gate = norm_before_gate - self.reset_parameters() - - def reset_parameters(self): - torch.nn.init.ones_(self.weight) - - def forward(self, x, z=None): - """If z is not None, we do norm(x) * silu(z) if norm_before_gate, else norm(x * silu(z)) - """ - return rmsnorm_fn(x, self.weight, self.bias, z=z, eps=self.eps, group_size=self.group_size, - norm_before_gate=self.norm_before_gate) diff --git a/build/torch-xpu/modules/mlp.py b/build/torch-xpu/modules/mlp.py deleted file mode 100644 index 2c9a44ed3e90cead07518c9c0cf9ca2ec08f69cf..0000000000000000000000000000000000000000 --- a/build/torch-xpu/modules/mlp.py +++ /dev/null @@ -1,141 +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 - -from functools import partial -from typing import TYPE_CHECKING, Any - -import torch -import torch.nn as nn -from torch.distributed import DeviceMesh -from torch.distributed.tensor import Placement, Replicate, Shard, distribute_module -from torch.distributed.tensor.parallel import ParallelStyle - -from ..modules.activations import powglu, powglu_linear, swiglu, swiglu_linear - -try: - from torch.distributed.tensor import DTensor -except (ImportError, AttributeError): - DTensor = None - -if TYPE_CHECKING: - from transformers.processing_utils import Unpack - - -class GatedMLP(nn.Module): - - def __init__( - self, - hidden_size: int, - hidden_ratio: int | None = None, - intermediate_size: int | None = None, - hidden_act: str = 'swish', - fuse_swiglu: bool = True, - powglu_power: float = 3.0, - ) -> GatedMLP: - super().__init__() - - self.hidden_size = hidden_size - # the final number of params is `hidden_ratio * hidden_size^2` - # `intermediate_size` is chosen to be a multiple of 256 closest to `2/3 * hidden_size * hidden_ratio` - if hidden_ratio is None: - hidden_ratio = 4 - if intermediate_size is None: - intermediate_size = int(hidden_size * hidden_ratio * 2 / 3) - intermediate_size = 256 * ((intermediate_size + 256 - 1) // 256) - self.hidden_ratio = hidden_ratio - self.intermediate_size = intermediate_size - self.hidden_act = hidden_act - self.fuse_swiglu = fuse_swiglu - self.powglu_power = powglu_power - - if hidden_act not in ('swish', 'powlu'): - raise ValueError(f'Unsupported hidden_act: {hidden_act}') - - self.gate_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=False) - self.up_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=False) - self.down_proj = nn.Linear(self.intermediate_size, self.hidden_size, bias=False) - if self.fuse_swiglu and hidden_act == 'swish': - self.swiglu_linear = SwiGLULinear() - - def forward( - self, - x: torch.Tensor, - **kwargs: Unpack[Any], - ) -> torch.Tensor: - gate, y = self.gate_proj(x), self.up_proj(x) - if self.hidden_act == 'powlu': - if self.fuse_swiglu: - return powglu_linear(gate, y, self.down_proj.weight, self.down_proj.bias, self.powglu_power) - return self.down_proj(powglu(gate, y, self.powglu_power)) - if self.fuse_swiglu: - return self.swiglu_linear(gate, y, self.down_proj.weight, self.down_proj.bias) - return self.down_proj(swiglu(gate, y)) - - -class SwiGLULinear(nn.Module): - - def forward(self, x, y, weight, bias): - return swiglu_linear(x, y, weight, bias) - - -class SwiGLULinearParallel(ParallelStyle): - def __init__( - self, - *, - input_layouts: Placement | None = None, - output_layouts: Placement | None = None, - use_local_output: bool = True, - ): - super().__init__() - self.input_layouts = (input_layouts or Shard(-1),) - self.output_layouts = (output_layouts or Replicate(),) - self.desired_input_layouts = (Shard(-1),) - self.use_local_output = use_local_output - - @staticmethod - def _prepare_input_fn( - input_layouts, desired_input_layouts, mod, inputs, device_mesh, - ): - x, y, weight, bias = inputs - if not isinstance(x, DTensor): - x = DTensor.from_local(x, device_mesh, input_layouts, run_check=False) - if x.placements != desired_input_layouts: - x = x.redistribute(placements=desired_input_layouts, async_op=True) - - if not isinstance(y, DTensor): - y = DTensor.from_local(y, device_mesh, input_layouts, run_check=False) - if y.placements != desired_input_layouts: - y = y.redistribute(placements=desired_input_layouts, async_op=True) - - if not isinstance(weight, DTensor): - weight = DTensor.from_local(weight, device_mesh, (Shard(1),)) - - if bias is not None and not isinstance(bias, DTensor): - bias = DTensor.from_local(bias, device_mesh, (Replicate(),)) - - return x, y, weight, bias - - @staticmethod - def _prepare_output_fn(output_layouts, use_local_output, mod, outputs, device_mesh): - # Rowwise sharding produces partial output, depending on output layouts: - # 1. to replicate -> allreduce - # 2. to shard -> reduce_scatter - if outputs.placements != output_layouts: - outputs = outputs.redistribute(placements=output_layouts, async_op=True) - # back to local tensor if use_local_output is True - return outputs.to_local() if use_local_output else outputs - - def _apply(self, module: nn.Module, device_mesh: DeviceMesh) -> nn.Module: - return distribute_module( - module, - device_mesh, - partition_fn=None, - input_fn=partial(self._prepare_input_fn, self.input_layouts, self.desired_input_layouts), - output_fn=partial(self._prepare_output_fn, self.output_layouts, self.use_local_output), - ) diff --git a/build/torch-xpu/modules/parallel.py b/build/torch-xpu/modules/parallel.py deleted file mode 100644 index 150c4c4e17f92655f646fcf85b2b7a9569acf293..0000000000000000000000000000000000000000 --- a/build/torch-xpu/modules/parallel.py +++ /dev/null @@ -1,44 +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.nn as nn -from torch.distributed import DeviceMesh -from torch.distributed.tensor import distribute_module -from torch.distributed.tensor.parallel import ParallelStyle -from torch.distributed.tensor.placement_types import Placement - -try: - from torch.distributed.tensor import DTensor -except (ImportError, AttributeError): - DTensor = None - - -class PrepareModuleWeight(ParallelStyle): - def __init__(self, *, layouts: Placement | None = None): - super().__init__() - self.layouts = layouts - - def _replicate_module_fn( - self, - name: str, - module: nn.Module, - device_mesh: DeviceMesh, - ): - for p_name, param in module.named_parameters(): - replicated_param = nn.Parameter( - DTensor.from_local(param, device_mesh, [self.layouts], run_check=False), - ) - module.register_parameter(p_name, replicated_param) - - def _apply(self, module: nn.Module, device_mesh: DeviceMesh) -> nn.Module: - return distribute_module( - module, - device_mesh, - partition_fn=self._replicate_module_fn, - input_fn=None, - output_fn=None, - ) diff --git a/build/torch-xpu/modules/rotary.py b/build/torch-xpu/modules/rotary.py deleted file mode 100644 index ddf10adcfc40db494ef4848c58a30d97601352e9..0000000000000000000000000000000000000000 --- a/build/torch-xpu/modules/rotary.py +++ /dev/null @@ -1,519 +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 torch.nn as nn -import triton -import triton.language as tl -from einops import rearrange, repeat - -from ..modules.backends import dispatch -from ..ops.utils import prepare_chunk_indices -from ..utils import IS_AMD, autotune_cache_kwargs, get_multiprocessor_count, input_guard - -NUM_WARPS_AUTOTUNE = [2, 4, 8, 16] if IS_AMD else [2, 4, 8, 16, 32] - - -def rotate_half(x, interleaved=False): - if not interleaved: - x1, x2 = x.chunk(2, dim=-1) - return torch.cat((-x2, x1), dim=-1) - else: - x1, x2 = x[..., ::2], x[..., 1::2] - return rearrange(torch.stack((-x2, x1), dim=-1), '... d two -> ... (d two)', two=2) - - -def rotary_embedding_ref(x, cos, sin, interleaved=False): - ro_dim = cos.shape[-1] * 2 - assert ro_dim <= x.shape[-1] - cos = repeat(cos, '... d -> ... 1 (2 d)' if not interleaved else '... d -> ... 1 (d 2)') - sin = repeat(sin, '... d -> ... 1 (2 d)' if not interleaved else '... d -> ... 1 (d 2)') - return torch.cat([x[..., :ro_dim] * cos + rotate_half(x[..., :ro_dim], interleaved) * sin, x[..., ro_dim:]], -1) - - -@triton.autotune( - configs=[ - triton.Config({}, num_warps=num_warps, num_stages=num_stages) - for num_warps in NUM_WARPS_AUTOTUNE - for num_stages in [2, 3, 4] - ], - 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, -): - i_t, i_b, i_h = 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), 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: - # Load the 1st and 2nd halves of x, do calculation, then store to 1st and 2nd halves of out - 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 - # write back result - 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: - # We don't want to load x[0, 2, 4, ...] and x[1, 3, 5, ...] separately since both are slow. - # Instead, we load x0 = x[0, 1, 2, 3, ...] and x1 = x[1, 0, 3, 2, ...]. - # Loading x0 will be fast but x1 will be slow. - # Then we load cos = cos[0, 0, 1, 1, ...] and sin = sin[0, 0, 1, 1, ...]. - # Then we do the calculation and use tl.where to pick put the right outputs for the even - # and for the odd indices. - o_d = tl.arange(0, BD) - o_d_swap = o_d + ((o_d + 1) % 2) * 2 - 1 # 1, 0, 3, 2, 5, 4, ... - 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) - - -@dispatch('modules') -def rotary_embedding_fwdbwd( - 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: - """ - Args: - x: [B, T, H, D]. - cos: [TR, R / 2] - sin: [TR, R / 2] - seqlen_offsets: integer or integer tensor of size [N] - cu_seqlens: [N + 1,] or None - - Returns: - y: [B, T, H, D] - """ - 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 - - # zeros_like: rows the kernel skips (negative o_cs under left-padding) must be defined, not uninitialized. - 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) - BT = min(128, triton.next_power_of_2(triton.cdiv(T, get_multiprocessor_count(x.device.index)))) - 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) - - grid = (NT, B, H) - rotary_embedding_kernel[grid]( - x, - cos, - sin, - y, - cu_seqlens, - chunk_indices, - 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, - ) - return y - - -class RotaryEmbeddingFunction(torch.autograd.Function): - - @staticmethod - @input_guard - def forward( - ctx, - x, - cos, - sin, - interleaved=False, - inplace=False, - seqlen_offsets: int | torch.Tensor = 0, - cu_seqlens: torch.Tensor | None = None, - chunk_indices: torch.LongTensor | None = None, - ): - y = rotary_embedding_fwdbwd( - x, - cos, - sin, - seqlen_offsets=seqlen_offsets, - cu_seqlens=cu_seqlens, - interleaved=interleaved, - inplace=inplace, - chunk_indices=chunk_indices, - ) - if isinstance(seqlen_offsets, int): - # Can't save int with save_for_backward - ctx.save_for_backward(cos, sin, cu_seqlens) - ctx.seqlen_offsets = seqlen_offsets - else: - ctx.save_for_backward(cos, sin, cu_seqlens, seqlen_offsets) - ctx.seqlen_offsets = None - ctx.interleaved = interleaved - ctx.inplace = inplace - ctx.chunk_indices = chunk_indices - return y if not inplace else x - - @staticmethod - @input_guard - def backward(ctx, do): - seqlen_offsets = ctx.seqlen_offsets - if seqlen_offsets is None: - cos, sin, cu_seqlens, seqlen_offsets = ctx.saved_tensors - else: - cos, sin, cu_seqlens = ctx.saved_tensors - # TD [2023-09-02]: For some reason Triton (2.0.0.post1) errors with - # "[CUDA]: invalid device context", and cloning makes it work. Idk why. Triton 2.1.0 works. - if not ctx.interleaved and not ctx.inplace: - do = do.clone() - dx = rotary_embedding_fwdbwd( - do, - cos, - sin, - seqlen_offsets=seqlen_offsets, - cu_seqlens=cu_seqlens, - interleaved=ctx.interleaved, - inplace=ctx.inplace, - conjugate=True, - chunk_indices=ctx.chunk_indices, - ) - return dx, None, None, None, None, None, None, None - - -def rotary_embedding( - x, - cos, - sin, - interleaved=False, - inplace=False, - seqlen_offsets: int | torch.Tensor = 0, - cu_seqlens: torch.Tensor | None = None, - chunk_indices: torch.LongTensor | None = None, -): - """ - Args: - x: [B, T, H, D] - cos, sin: [TR, R//2] - interleaved: - If True, rotate pairs of even and odd dimensions (GPT-J style) instead of 1st half and 2nd half (GPT-NeoX style). - inplace: - If True, apply rotary embedding in-place. - seqlen_offsets: [N,] or int. - Each sequence in x is shifted by this amount. - Most commonly used in inference when we have KV cache. - cu_seqlens: [N + 1,] or None - - Returns: - out: [B, T, H, D] - """ - return RotaryEmbeddingFunction.apply( - x, - cos, - sin, - interleaved, - inplace, - seqlen_offsets, - cu_seqlens, - chunk_indices, - ) - - -class RotaryEmbedding(nn.Module): - """ - The rotary position embeddings from RoFormer_ (Su et. al). - A crucial insight from the method is that the query and keys are - transformed by rotation matrices which depend on the relative positions. - - Other implementations are available in the Rotary Transformer repo_ and in - GPT-NeoX_, GPT-NeoX was an inspiration - - .. _RoFormer: https://arxiv.org/abs/2104.09864 - .. _repo: https://github.com/ZhuiyiTechnology/roformer - .. _GPT-NeoX: https://github.com/EleutherAI/gpt-neox - - If scale_base is not None, this implements XPos (Sun et al., https://arxiv.org/abs/2212.10554). - A recommended value for scale_base is 512: https://github.com/HazyResearch/flash-attention/issues/96 - Reference: https://github.com/sunyt32/torchscale/blob/main/torchscale/component/xpos_relative_position.py - """ - - def __init__( - self, - dim: int, - base: float = 10000.0, - scale_base: float | None = None, - interleaved: bool = False, - pos_idx_in_fp32: bool = True, - device: torch.device | None = None, - ): - """ - interleaved: - If True, rotate pairs of even and odd dimensions (GPT-J style) instead of 1st half and 2nd half (GPT-NeoX style). - pos_idx_in_fp32: - If True, the position indices [0.0, ..., seqlen - 1] are in fp32, otherwise they might be in lower precision. - This option was added because previously (before 2023-07-02), when we construct - the position indices, we use the dtype of self.inv_freq. - In most cases this would be fp32, but if the model is trained in pure bf16 (not mixed precision), then - self.inv_freq would be bf16, and the position indices are also in bf16. - Because of the limited precision of bf16 (e.g. 1995.0 is rounded to 2000.0), the - embeddings for some positions will coincide. - To maintain compatibility with models previously trained in pure bf16, we add this option. - """ - super().__init__() - - self.dim = dim - self.base = float(base) - self.scale_base = scale_base - self.interleaved = interleaved - self.pos_idx_in_fp32 = pos_idx_in_fp32 - self.device = device - - # Generate and save the inverse frequency buffer (non trainable) - self.register_buffer("inv_freq", torch.empty(-(dim // -2), dtype=torch.float32, device=device), persistent=False) - - scale = None - if scale_base is not None: - scale = torch.empty(-(dim // -2), dtype=torch.float32, device=device) - self.register_buffer("scale", scale, persistent=False) - - self._seq_len_cached = 0 - self._cos_cached = None - self._sin_cached = None - self._cos_k_cached = None - self._sin_k_cached = None - - self.reset_parameters() - - def reset_parameters(self): - with torch.no_grad(): - self.inv_freq.copy_(self._compute_inv_freq(device=self.inv_freq.device)) - if self.scale_base is not None: - self.scale.copy_(self._compute_scale(device=self.scale.device)) - - def __repr__(self): - s = f"{self.__class__.__name__}(" - s += f"dim={self.dim}, " - s += f"base={self.base}, " - s += f"interleaved={self.interleaved}, " - if self.scale_base is not None: - s += f"scale_base={self.scale_base}, " - s += f"pos_idx_in_fp32={self.pos_idx_in_fp32})" - return s - - def _compute_inv_freq(self, device=None): - return 1.0 / ( - self.base - ** (torch.arange(0, self.dim, 2, device=device, dtype=torch.float32) / self.dim) - ) - - def _compute_scale(self, device=None): - return (torch.arange(0, self.dim, 2, device=device, dtype=torch.float32) + 0.4 * self.dim) / (1.4 * self.dim) - - def _update_cos_sin_cache(self, seqlen, device=None, dtype=None): - # Reset the tables if the sequence length has changed, - # if we're on a new device (possibly due to tracing for instance), - # or if we're switching from inference mode to training - if ( - seqlen > self._seq_len_cached - or self._cos_cached is None - or self._cos_cached.device != device - or self._cos_cached.dtype != dtype - or (self.training and self._cos_cached.is_inference()) - ): - self._seq_len_cached = seqlen - # We want fp32 here, not self.inv_freq.dtype, since the model could be loaded in bf16 - # And the output of arange can be quite large, so bf16 would lose a lot of precision. - # However, for compatibility reason, we add an option to use the dtype of self.inv_freq. - if self.pos_idx_in_fp32: - t = torch.arange(seqlen, device=device, dtype=torch.float32) - # We want fp32 here as well since inv_freq will be multiplied with t, and the output - # will be large. Having it in bf16 will lose a lot of precision and cause the - # cos & sin output to change significantly. - # We want to recompute self.inv_freq if it was not loaded in fp32 - if self.inv_freq.dtype != torch.float32: - inv_freq = self._compute_inv_freq(device=device) - else: - inv_freq = self.inv_freq - else: - t = torch.arange(seqlen, device=device, dtype=self.inv_freq.dtype) - inv_freq = self.inv_freq - # Don't do einsum, it converts fp32 to fp16 under AMP - # freqs = torch.einsum("i,j->ij", t, self.inv_freq) - freqs = torch.outer(t, inv_freq) - if self.scale is None: - self._cos_cached = torch.cos(freqs).to(dtype) - self._sin_cached = torch.sin(freqs).to(dtype) - else: - power = ( - torch.arange(seqlen, dtype=self.scale.dtype, device=self.scale.device) - - seqlen // 2 - ) / self.scale_base - scale = self.scale.to(device=power.device) ** rearrange(power, "s -> s 1") - # We want the multiplication by scale to happen in fp32 - self._cos_cached = (torch.cos(freqs) * scale).to(dtype) - self._sin_cached = (torch.sin(freqs) * scale).to(dtype) - self._cos_k_cached = (torch.cos(freqs) / scale).to(dtype) - self._sin_k_cached = (torch.sin(freqs) / scale).to(dtype) - - def forward( - self, - q: torch.Tensor, - k: torch.Tensor, - seqlen_offset: int | torch.Tensor = 0, - cu_seqlens: torch.Tensor | None = None, - max_seqlen: int | None = None, - chunk_indices: torch.LongTensor | None = None, - ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]: - """ - q: [B, T, H, D] - k: [B, T, H, D] - seqlen_offset: - [N] or int. - Each sequence in x is shifted by this amount. - Most commonly used in inference when we have KV cache. - cu_seqlens: [N + 1] or None - max_seqlen: int - """ - if max_seqlen is not None: - self._update_cos_sin_cache(max_seqlen, device=q.device, dtype=q.dtype) - elif isinstance(seqlen_offset, int): - self._update_cos_sin_cache(q.shape[1] + seqlen_offset, device=q.device, dtype=q.dtype) - if self.scale is None: - q = rotary_embedding( - q, - self._cos_cached, - self._sin_cached, - interleaved=self.interleaved, - seqlen_offsets=seqlen_offset, - cu_seqlens=cu_seqlens, - chunk_indices=chunk_indices, - ) - k = rotary_embedding( - k, - self._cos_cached, - self._sin_cached, - interleaved=self.interleaved, - seqlen_offsets=seqlen_offset, - cu_seqlens=cu_seqlens, - chunk_indices=chunk_indices, - ) - - else: - q = rotary_embedding( - q, - self._cos_cached, - self._sin_cached, - interleaved=self.interleaved, - seqlen_offsets=seqlen_offset, - cu_seqlens=cu_seqlens, - chunk_indices=chunk_indices, - ) - k = rotary_embedding( - k, - self._cos_k_cached, - self._sin_k_cached, - interleaved=self.interleaved, - seqlen_offsets=seqlen_offset, - cu_seqlens=cu_seqlens, - chunk_indices=chunk_indices, - ) - - return q, k diff --git a/build/torch-xpu/modules/token_shift.py b/build/torch-xpu/modules/token_shift.py deleted file mode 100644 index 8712da2eaa0878bceb52d115c70d51ffae9633ce..0000000000000000000000000000000000000000 --- a/build/torch-xpu/modules/token_shift.py +++ /dev/null @@ -1,573 +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 -import triton.language as tl - -from ..ops.utils import prepare_chunk_indices -from ..utils import IS_AMD, IS_NPU, autotune_cache_kwargs, get_multiprocessor_count, input_guard, tensor_cache - -NUM_WARPS_AUTOTUNE = [2, 4, 8, 16] if IS_AMD else [2, 4, 8, 16, 32] -# Ascend Triton rejects 2-D grids whose product exceeds 65535 unless -# TRITON_ALL_BLOCKS_PARALLEL=1. Fall back to the long kernel instead. -_NPU_MAX_TRITON_GRID = 65535 - - -def token_shift_ref( - x: torch.Tensor, - cu_seqlens: torch.Tensor | None = None, -) -> torch.Tensor: - if cu_seqlens is not None: - # Variable length mode with cu_seqlens - assert x.dim() == 3, "Input must be [B, T, D]" - B, T, D = x.shape - assert B == 1, "Batch size must be 1 when using cu_seqlens" - - result = torch.zeros_like(x) - N = cu_seqlens.shape[0] - 1 - - for i in range(N): - start = cu_seqlens[i].item() - end = cu_seqlens[i+1].item() - seq_len = end - start - - if seq_len <= 1: - # For sequences of length 1 or 0, delta is simply -x - result[0, start:end] = -x[0, start:end] - else: - # For longer sequences, handle padding manually - shifted = torch.zeros_like(x[0, start:end]) - shifted[1:] = x[0, start:end-1] - delta = shifted - x[0, start:end] - result[0, start:end] = delta - - return result - else: - time_shift = torch.nn.ZeroPad2d((0, 0, 1, -1)) - shifted = time_shift(x) - delta = shifted - x - return delta - - -@triton.heuristics({ - 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None, - 'USE_INITIAL_STATE': lambda args: args['cache'] is not None, -}) -@triton.autotune( - configs=[ - triton.Config({}, num_warps=num_warps, num_stages=num_stages) - for num_warps in NUM_WARPS_AUTOTUNE - for num_stages in [1, 2, 3] - ], - key=['BD'], - **autotune_cache_kwargs, -) -@triton.jit -def token_shift_fwd_kernel_short( - x, - y, - cu_seqlens, - cache, - cache_out, - T, - D: tl.constexpr, - BD: tl.constexpr, - IS_VARLEN: tl.constexpr, - USE_INITIAL_STATE: tl.constexpr, - STORE_FINAL_STATE: tl.constexpr, - IS_DECODE: tl.constexpr, -): - i_b, i_t = tl.program_id(0), tl.program_id(1) - - if IS_VARLEN: - i_n = i_b - bos, eos = tl.load(cu_seqlens + i_n).to(tl.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) - g_t = i_t + bos - - if g_t >= eos: - return - - is_first_pos = (i_t == 0) - is_last_pos = (g_t == eos - 1) - else: - g_t = i_t - is_first_pos = (g_t == 0) - is_last_pos = (g_t == T - 1) - - o_d = tl.arange(0, BD) - m_d = o_d < D - - if IS_VARLEN: - base_offset = g_t * D + o_d - else: - base_offset = i_b * T*D + g_t * D + o_d - - b_x = tl.load(x + base_offset, mask=m_d) - if IS_VARLEN: - cache_offset = i_n * D + o_d # i_n is seq index - else: - cache_offset = i_b * D + o_d # i_b is batch index - - if IS_DECODE and USE_INITIAL_STATE: - b_cache = tl.load(cache + cache_offset, mask=m_d) - delta = b_cache - b_x - tl.store(y + base_offset, delta, mask=m_d) - if STORE_FINAL_STATE: - tl.store(cache_out + cache_offset, b_x, mask=m_d) - return - - if is_first_pos: - # First position in sequence: delta = -hidden_states - if USE_INITIAL_STATE: - # cache shape: [N, D] - b_cache = tl.load(cache + cache_offset, mask=m_d) - delta = b_cache - b_x - tl.store(y + base_offset, delta, mask=m_d) - else: - tl.store(y + base_offset, -b_x, mask=m_d) - return - - # Other positions: delta = prev - curr - if IS_VARLEN: - prev_offset = (g_t-1) * D + o_d - else: - prev_offset = i_b * T*D + (g_t-1) * D + o_d - - prev_values = tl.load(x + prev_offset, mask=m_d) - delta = prev_values - b_x - tl.store(y + base_offset, delta, mask=m_d) - if STORE_FINAL_STATE: - if is_last_pos: - tl.store(cache_out + cache_offset, b_x, mask=m_d) - - -@triton.heuristics({ - 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None, - 'USE_INITIAL_STATE': lambda args: args['cache'] is not None, -}) -@triton.autotune( - configs=[ - triton.Config({}, num_warps=num_warps, num_stages=num_stages) - for num_warps in NUM_WARPS_AUTOTUNE - for num_stages in [1, 2, 3] - ], - key=['BD', 'NB'], - **autotune_cache_kwargs, -) -@triton.jit -def token_shift_fwd_kernel_long( - x, - y, - cu_seqlens, - chunk_indices, - cache, - cache_out, - T, - D: tl.constexpr, - BD: tl.constexpr, - BT: tl.constexpr, - NB: tl.constexpr, - ND: tl.constexpr, - IS_VARLEN: tl.constexpr, - USE_INITIAL_STATE: tl.constexpr, - STORE_FINAL_STATE: tl.constexpr, -): - i_dt, i_b = tl.program_id(0), tl.program_id(1) - i_d, i_t = i_dt % ND, i_dt // ND - - 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_start = i_t * BT - t_end = tl.minimum(t_start + BT, eos - bos) - else: - i_n = i_b - bos, eos = i_b * T, (i_b + 1) * T - t_start = i_t * BT - t_end = tl.minimum(t_start + BT, T) - - o_d = i_d * BD + tl.arange(0, BD) - m_d = o_d < D - - for t in range(t_start, t_end): - global_t = bos + t - offset = global_t * D + o_d - b_x = tl.load(x + offset, mask=m_d) - is_first = (global_t == bos) - if is_first: - if USE_INITIAL_STATE: - # cache shape: [N, D] - cache_off = i_n * D + o_d if IS_VARLEN else i_b * D + o_d - b_cache = tl.load(cache + cache_off, mask=m_d) - delta = b_cache - b_x - else: - delta = -b_x - else: - prev_off = offset - D - b_prev = tl.load(x + prev_off, mask=m_d) - delta = b_prev - b_x - - tl.store(y + offset, delta, mask=m_d) - - if STORE_FINAL_STATE: - if global_t == eos - 1: - cache_out_off = i_n * D + o_d if IS_VARLEN else i_b * D + o_d - tl.store(cache_out + cache_out_off, b_x, mask=m_d) - - -@triton.heuristics({ - 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None, - 'USE_INITIAL_STATE': lambda args: args['grad_cache_out'] is not None, - 'HAS_DCACHE': lambda args: args['grad_cache_in'] is not None, -}) -@triton.autotune( - configs=[ - triton.Config({}, num_warps=num_warps, num_stages=num_stages) - for num_warps in NUM_WARPS_AUTOTUNE - for num_stages in [1, 2, 3] - ], - key=['BD'], - **autotune_cache_kwargs, -) -@triton.jit -def token_shift_bwd_kernel_short( - dx, - dy, - cu_seqlens, - grad_cache_in, - grad_cache_out, - T, - D: tl.constexpr, - BD: tl.constexpr, - IS_VARLEN: tl.constexpr, - USE_INITIAL_STATE: tl.constexpr, - HAS_DCACHE: tl.constexpr, -): - i_b, i_t = tl.program_id(0), tl.program_id(1) - - if IS_VARLEN: - i_n = i_b - bos, eos = tl.load(cu_seqlens + i_n).to(tl.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) - g_t = i_t + bos - if g_t >= eos: - return - is_first_pos = (g_t == bos) - is_last_pos = (g_t == eos - 1) - else: - g_t = i_t - is_first_pos = (g_t == 0) - is_last_pos = (g_t == T - 1) - - o_d = tl.arange(0, BD) - m_d = o_d < D - - if IS_VARLEN: - base_offset = g_t * D + o_d - # This should not be used for varlen - cache_off = i_n * D + o_d - else: - base_offset = i_b * T * D + g_t * D + o_d - cache_off = i_b * D + o_d - - b_dy = tl.load(dy + base_offset, mask=m_d) - - if is_last_pos: - # grad = -grad_delta[t] + grad_cache_in（from next rank） - if HAS_DCACHE: - b_dy_cache = tl.load(grad_cache_in + cache_off, mask=m_d) - b_dx = -b_dy + b_dy_cache - else: - b_dx = -b_dy - else: - # grad = -grad_delta[t] + grad_delta[t+1] - if IS_VARLEN: - next_offset = (g_t + 1) * D + o_d - else: - next_offset = i_b * T * D + (g_t + 1) * D + o_d - b_dx = -b_dy + tl.load(dy + next_offset, mask=m_d) - - tl.store(dx + base_offset, b_dx, mask=m_d) - - if USE_INITIAL_STATE: - if is_first_pos: - tl.store(grad_cache_out + cache_off, b_dy, mask=m_d) - - -@triton.heuristics({ - 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None, - 'USE_INITIAL_STATE': lambda args: args['grad_cache_out'] is not None, - 'HAS_DCACHE': lambda args: args['grad_cache_in'] is not None, -}) -@triton.autotune( - configs=[ - triton.Config({}, num_warps=num_warps, num_stages=num_stages) - for num_warps in NUM_WARPS_AUTOTUNE - for num_stages in [1, 2, 3] - ], - key=['BD', 'NB'], - **autotune_cache_kwargs, -) -@triton.jit -def token_shift_bwd_kernel_long( - dx, - dy, - cu_seqlens, - chunk_indices, - grad_cache_in, - grad_cache_out, - T, - D: tl.constexpr, - BD: tl.constexpr, - BT: tl.constexpr, - NB: tl.constexpr, - ND: tl.constexpr, - IS_VARLEN: tl.constexpr, - USE_INITIAL_STATE: tl.constexpr, - HAS_DCACHE: tl.constexpr, -): - i_dt, i_b = tl.program_id(0), tl.program_id(1) - i_d, i_t_blk = i_dt % ND, i_dt // ND - - if IS_VARLEN: - i_n, i_t_blk = tl.load(chunk_indices + i_t_blk * 2).to(tl.int32), \ - tl.load(chunk_indices + i_t_blk * 2 + 1).to(tl.int32) - bos, eos = tl.load(cu_seqlens + i_n), tl.load(cu_seqlens + i_n + 1) - t_start = i_t_blk * BT - t_end = tl.minimum(t_start + BT, eos - bos) - else: - i_n = i_b - bos, eos = i_b * T, (i_b + 1) * T - t_start = i_t_blk * BT - t_end = tl.minimum(t_start + BT, T) - - o_d = i_d * BD + tl.arange(0, BD) - m_d = o_d < D - cache_off = i_n * D + o_d if IS_VARLEN else i_b * D + o_d - - for t in range(t_start, t_end): - global_t = bos + t - offset = global_t * D + o_d - b_dy = tl.load(dy + offset, mask=m_d) - - if global_t == eos - 1: - if HAS_DCACHE: - b_dy_cache = tl.load(grad_cache_in + cache_off, mask=m_d) - b_dx = -b_dy + b_dy_cache - else: - b_dx = -b_dy - else: - next_off = offset + D - b_dx = -b_dy + tl.load(dy + next_off, mask=m_d) - - tl.store(dx + offset, b_dx, mask=m_d) - - if USE_INITIAL_STATE: - if global_t == bos: - tl.store(grad_cache_out + cache_off, b_dy, mask=m_d) - - -@tensor_cache -def prepare_maxlens(cu_seqlens: torch.LongTensor) -> int: - return torch.max(cu_seqlens.diff()).item() - - -def token_shift_fwd( - x: torch.Tensor, - cu_seqlens: torch.Tensor | None = None, - cache: torch.Tensor | None = None, - output_cache: bool = False, - chunk_indices: torch.LongTensor | None = None, -) -> torch.Tensor: - B, T, D = x.shape - y = torch.empty_like(x) - - if cu_seqlens is not None: - T = prepare_maxlens(cu_seqlens) - N = len(cu_seqlens) - 1 - else: - N = B - - use_short_kernel = T <= 4096 - if IS_NPU and use_short_kernel and N * T > _NPU_MAX_TRITON_GRID: - use_short_kernel = False - - if output_cache: - cache_out = torch.empty((N, D), device=x.device, dtype=x.dtype) - else: - cache_out = None - - if use_short_kernel: - if cu_seqlens is not None: - N = len(cu_seqlens) - 1 - else: - N = B - BD = triton.next_power_of_2(D) - grid = (N, T) - IS_DECODE = T == 1 or (B == 1 and T == N) - token_shift_fwd_kernel_short[grid]( - x=x, - y=y, - cu_seqlens=cu_seqlens, - cache=cache, - cache_out=cache_out, - T=T, - D=D, - BD=BD, - STORE_FINAL_STATE=output_cache, - IS_DECODE=IS_DECODE, - ) - else: - BT = min(64, triton.next_power_of_2(triton.cdiv(max(16, B*T), get_multiprocessor_count(x.device.index)))) - if chunk_indices is None and cu_seqlens is not None: - chunk_indices = prepare_chunk_indices(cu_seqlens, BT) - NT = len(chunk_indices) if cu_seqlens is not None else triton.cdiv(T, BT) - - BD = triton.next_power_of_2(D) - ND = triton.cdiv(D, BD) - NB = triton.cdiv(B*T, 1024) - - def grid(meta): return (ND * NT, 1 if cu_seqlens is not None else N) - token_shift_fwd_kernel_long[grid]( - x, - y, - cu_seqlens, - chunk_indices, - cache, - cache_out, - T, - D=D, - BD=BD, - BT=BT, - NB=NB, - ND=ND, - STORE_FINAL_STATE=output_cache, - ) - - return y, N, T, use_short_kernel, cache_out - - -def token_shift_bwd( - dy: torch.Tensor, - N: int, - T: int, - dcache: torch.Tensor | None = None, - cu_seqlens: torch.Tensor | None = None, - use_short_kernel: bool = True, - has_init_cache: bool = False, - chunk_indices: torch.LongTensor | None = None, -) -> torch.Tensor: - D = dy.shape[2] - BD = triton.next_power_of_2(D) - dx = torch.empty_like(dy) - if has_init_cache: - grad_cache_out = torch.empty((N, D), device=dy.device, dtype=dy.dtype) - else: - grad_cache_out = None - if use_short_kernel: - grid = (N, T) - token_shift_bwd_kernel_short[grid]( - dy=dy, - dx=dx, - cu_seqlens=cu_seqlens, - grad_cache_in=dcache, - grad_cache_out=grad_cache_out, - T=T, - D=D, - BD=BD, - ) - else: - BT = min(64, triton.next_power_of_2(triton.cdiv(max(16, dy.numel() // D), - get_multiprocessor_count(dy.device.index)))) - if chunk_indices is None and cu_seqlens is not None: - chunk_indices = prepare_chunk_indices(cu_seqlens, BT) - NT = len(chunk_indices) if cu_seqlens is not None else triton.cdiv(T, BT) - NB = triton.cdiv(N * dy.shape[1], 1024) - BD = triton.next_power_of_2(D) - ND = triton.cdiv(D, BD) - - def grid(meta): return (ND * NT, 1 if cu_seqlens is not None else N) - token_shift_bwd_kernel_long[grid]( - dx, - dy, - cu_seqlens, - chunk_indices, - dcache, - grad_cache_out, - T, - D=D, - BD=BD, - BT=BT, - NB=NB, - ND=ND, - ) - return dx, grad_cache_out - - -class TokenShift(torch.autograd.Function): - - @staticmethod - @input_guard - def forward(ctx, x: torch.Tensor, cu_seqlens: torch.Tensor | None = None, - cache: torch.Tensor | None = None, output_cache: bool = False, - chunk_indices: torch.LongTensor | None = None): - output, N, T, use_short_kernel, cache_out = token_shift_fwd(x, cu_seqlens, cache, output_cache, chunk_indices) - ctx.cu_seqlens = cu_seqlens - ctx.chunk_indices = chunk_indices - ctx.N = N - ctx.T = T - ctx.use_short_kernel = use_short_kernel - ctx.has_cache = cache is not None - return output, cache_out - - @staticmethod - @input_guard - def backward(ctx, dy: torch.Tensor, dcache: torch.Tensor | None = None): - dx, grad_cache = token_shift_bwd(dy, ctx.N, ctx.T, dcache, ctx.cu_seqlens, - ctx.use_short_kernel, ctx.has_cache, ctx.chunk_indices) - return dx, None, grad_cache, None, None - - -@torch.compiler.disable -def token_shift( - x: torch.Tensor, - cu_seqlens: torch.LongTensor | None = None, - cache: torch.Tensor | None = None, - output_cache: bool = False, - chunk_indices: torch.LongTensor | None = None, -): - """ - Token-shift operation implemented with Triton kernels. - - Args: - x: Input tensor of shape [B, T, D] (or [1, T, D] when `cu_seqlens` is supplied). - cu_seqlens: Optional cumulative sequence lengths of shape [B + 1]. - When supplied, `x.shape[0]` must be 1 and `x.dim()` must be 3. - cache: Optional cache tensor of shape [N, D] that holds the last token - from the previous call. - output_cache: Whether to return the updated cache alongside the output. - In previous versions this parameter did not exist and the - cache was always dropped; to preserve backward compatibility - the default is False. - - Returns: - output: Tensor of shape [B, T, D] after applying the token-shift. - - cache_out: Tensor of shape [B, 1, D] containing the last token that - should be fed as `cache` in the next call. Only returned - when `output_cache=True`. - """ - if cu_seqlens is not None: - assert x.dim() == 3, "Input must be [B, T, D]" - assert x.shape[0] == 1, "Batch size must be 1 when using cu_seqlens" - - output, cache_out = TokenShift.apply(x, cu_seqlens, cache, output_cache, chunk_indices) - if output_cache: - return output, cache_out - else: - return output diff --git a/build/torch-xpu/modules/token_shift_cp.py b/build/torch-xpu/modules/token_shift_cp.py deleted file mode 100644 index b92027a11c6a0fca342079d272fb15bc643a170f..0000000000000000000000000000000000000000 --- a/build/torch-xpu/modules/token_shift_cp.py +++ /dev/null @@ -1,229 +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 - -""" -Context Parallel support for Token Shift. - -Token shift has a 1-token dependency on previous tokens: - y[t] = x[t-1] - x[t] (for t > 0) - y[0] = cache - x[0] (cache is the last token from previous rank) - -In CP mode, non-first ranks need the last token from the previous rank as cache. -Backward: non-last ranks need to send the last token's gradient to previous rank. -""" - -import torch -import torch.distributed as dist - -from ..modules.token_shift import token_shift_bwd, token_shift_fwd -from ..ops.cp import FLACPContext, conv_cp_send_recv_bwd, conv_cp_send_recv_fwd - - -class TokenShiftCPFunction(torch.autograd.Function): - """ - Context Parallel version of TokenShift. - - Forward: - 1. Get last token from previous rank to construct cache - 2. Call token_shift_fwd with cache - - Backward: - 1. Call token_shift_bwd to get dx - 2. Sync communication: add next rank's first token gradient to current rank's last token - """ - - @staticmethod - def _prepare_cache_for_cp( - x: torch.Tensor, - cu_seqlens: torch.Tensor | None, - context: FLACPContext, - group: dist.ProcessGroup | None, - ) -> tuple[torch.Tensor | None, int]: - """Prepare cache for CP forward pass by communicating with previous rank. - - Args: - x: Input tensor of shape [1, T, D] - cu_seqlens: Cumulative sequence lengths - context: CP context - group: Process group for communication - - Returns: - cache: Cache tensor of shape [N, D] or None - pre_num_tokens: Number of tokens from previous rank for the first sequence - """ - if group is None: - return None, 0 - - D = x.shape[-1] - cache = None - pre_num_tokens = 0 - - if not context.is_first_rank: - # Non-first rank: need cache from previous rank - 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_token = x_2d[-1:].contiguous() # [1, D] - prev_last_token = conv_cp_send_recv_fwd(last_token, group) # [1, D] - - # For varlen: only the first sequence needs cache from prev rank - N = len(cu_seqlens) - 1 if cu_seqlens is not None else 1 - cache = torch.zeros(N, D, device=x.device, dtype=x.dtype) - - # pre_num_conv_tokens tells us how many tokens from prev rank - # belong to the first sequence on this rank - pre_num_tokens = getattr(context, 'pre_num_conv_tokens', 0) - if pre_num_tokens > 0: - # The prev rank's last token is used as cache for first sequence - cache[0] = prev_last_token[0] - else: - # First rank: participate in send but don't use received data - x_2d = x.squeeze(0) - last_token = x_2d[-1:].contiguous() - _ = conv_cp_send_recv_fwd(last_token, group) - - return cache, pre_num_tokens - - @staticmethod - def _correct_dx_for_cp( - dx: torch.Tensor, - grad_cache: torch.Tensor | None, - group: dist.ProcessGroup | None, - is_first_rank: bool, - pre_num_tokens: int = 0, - ) -> None: - """Correct dx gradients for CP backward pass. - - Args: - dx: Gradient tensor to be corrected, shape [1, T, D] - grad_cache: Gradient w.r.t. cache, shape [N, D] or None - group: Process group - is_first_rank: Whether this is the first rank - pre_num_tokens: Number of tokens from previous rank for first sequence - """ - if group is None: - return - - D = dx.shape[-1] - - # Prepare gradient to send to previous rank - if grad_cache is not None and pre_num_tokens > 0: - # Only first sequence's cache gradient is relevant - d_cache = grad_cache[0:1] # [1, D] - else: - d_cache = torch.zeros(1, D, device=dx.device, dtype=dx.dtype) - - # Send to previous rank, receive from next rank - recv_grad = conv_cp_send_recv_bwd(d_cache, group) # [1, D] - - # Add received gradient to current rank's last token - dx[0, -1, :].add_(recv_grad[0]) - - @staticmethod - def forward( - ctx, - x: torch.Tensor, - cu_seqlens: torch.Tensor | None, - chunk_indices: torch.Tensor | None, - cp_context: FLACPContext | None, - ): - if cp_context is None: - raise ValueError("cp_context must be provided for TokenShiftCPFunction") - - cu_seqlens = cp_context.cu_seqlens - group = cp_context.group - - # Prepare cache for CP - cache, pre_num_tokens = TokenShiftCPFunction._prepare_cache_for_cp( - x=x, - cu_seqlens=cu_seqlens, - context=cp_context, - group=group, - ) - - # Save for backward - ctx.cu_seqlens = cu_seqlens - ctx.chunk_indices = chunk_indices - ctx.group = group - ctx.has_cache = cache is not None - ctx.is_first_rank = cp_context.is_first_rank - ctx.pre_num_tokens = pre_num_tokens - - # Call original forward - y, N, T, use_short_kernel, cache_out = token_shift_fwd( - x=x, - cu_seqlens=cu_seqlens, - cache=cache, - output_cache=True, - chunk_indices=chunk_indices, - ) - - ctx.N = N - ctx.T = T - ctx.use_short_kernel = use_short_kernel - - return y - - @staticmethod - def backward(ctx, dy: torch.Tensor): - group = ctx.group - - # Prepare dcache for backward - # For CP: non-last rank needs to receive gradient from next rank - # This is handled in _correct_dx_for_cp after computing dx - dcache = None # Will be computed by token_shift_bwd - - # Call original backward - dx, grad_cache = token_shift_bwd( - dy=dy, - N=ctx.N, - T=ctx.T, - dcache=dcache, - cu_seqlens=ctx.cu_seqlens, - use_short_kernel=ctx.use_short_kernel, - has_init_cache=ctx.has_cache, - chunk_indices=ctx.chunk_indices, - ) - - # Correct dx gradients for CP - TokenShiftCPFunction._correct_dx_for_cp( - dx=dx, - grad_cache=grad_cache, - group=group, - is_first_rank=ctx.is_first_rank, - pre_num_tokens=ctx.pre_num_tokens, - ) - - return dx, None, None, None - - -@torch.compiler.disable -def token_shift_cp( - x: torch.Tensor, - cp_context: FLACPContext, - cu_seqlens: torch.Tensor | None = None, - chunk_indices: torch.Tensor | None = None, -): - """ - Context Parallel version of token_shift. - - Args: - x: Input tensor of shape [1, T, D] - cp_context: CP context (required for CP mode) - cu_seqlens: Cumulative sequence lengths - chunk_indices: Chunk indices for variable-length sequences - - Returns: - output: Tensor of shape [1, T, D] after applying token-shift - """ - if cp_context is None: - raise ValueError("cp_context must be provided for token_shift_cp") - - assert cp_context.cu_seqlens is not None, "cu_seqlens must be provided for token_shift_cp" - - return TokenShiftCPFunction.apply( - x, cu_seqlens, chunk_indices, cp_context - ) diff --git a/build/torch-xpu/ops/__init__.py b/build/torch-xpu/ops/__init__.py deleted file mode 100644 index 3f962e5dfd589647dd3966a128611808f2a02ace..0000000000000000000000000000000000000000 --- a/build/torch-xpu/ops/__init__.py +++ /dev/null @@ -1,91 +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 .abc import chunk_abc -from .attn import parallel_attn -from .attnres import fused_attnres -from .based import fused_chunk_based, parallel_based -from .comba import chunk_comba, fused_recurrent_comba -from .delta_rule import chunk_delta_rule, fused_chunk_delta_rule, fused_recurrent_delta_rule -from .forgetting_attn import parallel_forgetting_attn -from .gated_delta_rule import chunk_gated_delta_rule, chunk_gdn, fused_recurrent_gated_delta_rule, fused_recurrent_gdn -from .generalized_delta_rule import ( - chunk_dplr_delta_rule, - chunk_iplr_delta_rule, - fused_recurrent_dplr_delta_rule, - fused_recurrent_iplr_delta_rule, -) -from .gla import chunk_gla, fused_chunk_gla, fused_recurrent_gla -from .gsa import chunk_gsa, fused_recurrent_gsa -from .hgrn import fused_recurrent_hgrn -from .kda import chunk_kda, fused_recurrent_kda -from .lightning_attn import chunk_lightning_attn, fused_recurrent_lightning_attn -from .linear_attn import chunk_linear_attn, fused_chunk_linear_attn, fused_recurrent_linear_attn -from .log_linear_attn import chunk_log_linear_attn -from .mesa_net import chunk_mesa_net -from .nsa import parallel_nsa -from .parallax import parallel_parallax -from .path_attn import parallel_path_attn -from .retention import chunk_retention, fused_chunk_retention, fused_recurrent_retention, parallel_retention -from .rwkv6 import chunk_rwkv6, fused_recurrent_rwkv6 -from .rwkv7 import chunk_rwkv7, fused_recurrent_rwkv7 -from .simple_gla import chunk_simple_gla, fused_chunk_simple_gla, fused_recurrent_simple_gla, parallel_simple_gla -from .wall_attn import parallel_wall_attn, parallel_wall_attn_decode - -__all__ = [ - 'chunk_abc', - 'chunk_comba', - 'chunk_delta_rule', - 'chunk_dplr_delta_rule', - 'chunk_gated_delta_rule', - 'chunk_gdn', - 'chunk_gla', - 'chunk_gsa', - 'chunk_iplr_delta_rule', - 'chunk_kda', - 'chunk_lightning_attn', - 'chunk_linear_attn', - 'chunk_log_linear_attn', - 'chunk_mesa_net', - 'chunk_retention', - 'chunk_rwkv6', - 'chunk_rwkv7', - 'chunk_simple_gla', - 'fused_attnres', - 'fused_chunk_based', - 'fused_chunk_delta_rule', - 'fused_chunk_gla', - 'fused_chunk_linear_attn', - 'fused_chunk_retention', - 'fused_chunk_simple_gla', - 'fused_recurrent_comba', - 'fused_recurrent_delta_rule', - 'fused_recurrent_dplr_delta_rule', - 'fused_recurrent_gated_delta_rule', - 'fused_recurrent_gdn', - 'fused_recurrent_gla', - 'fused_recurrent_gsa', - 'fused_recurrent_hgrn', - 'fused_recurrent_iplr_delta_rule', - 'fused_recurrent_kda', - 'fused_recurrent_lightning_attn', - 'fused_recurrent_linear_attn', - 'fused_recurrent_retention', - 'fused_recurrent_rwkv6', - 'fused_recurrent_rwkv7', - 'fused_recurrent_simple_gla', - 'parallel_attn', - 'parallel_based', - 'parallel_forgetting_attn', - 'parallel_nsa', - 'parallel_parallax', - 'parallel_path_attn', - 'parallel_retention', - 'parallel_simple_gla', - 'parallel_wall_attn', - 'parallel_wall_attn_decode', -] diff --git a/build/torch-xpu/ops/abc/chunk.py b/build/torch-xpu/ops/abc/chunk.py deleted file mode 100644 index 1e364f43d7f6206d583305ebec966511b67a0f3c..0000000000000000000000000000000000000000 --- a/build/torch-xpu/ops/abc/chunk.py +++ /dev/null @@ -1,1119 +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 -import triton.language as tl - -from ...ops.utils import softmax_bwd, softmax_fwd -from ...ops.utils.logcumsumexp import logcumsumexp_fwd_kernel -from ...ops.utils.op import exp -from ...utils import input_guard - - -@triton.jit(do_not_specialize=['T']) -def chunk_abc_fwd_kernel_h( - k, - v, - z, - h, - h0, - ht, - T, - K: tl.constexpr, - V: tl.constexpr, - BT: tl.constexpr, - BK: tl.constexpr, - BV: tl.constexpr, - NT: tl.constexpr, - NORMK: tl.constexpr, - USE_INITIAL_STATE: tl.constexpr, - STORE_FINAL_STATE: tl.constexpr, -): - i_v, i_k, i_bh = tl.program_id(0), tl.program_id(1), tl.program_id(2) - - b_h = tl.zeros([BK, BV], dtype=tl.float32) - if USE_INITIAL_STATE: - p_h = tl.make_block_ptr(h0 + i_bh * K * V, (K, V), (V, 1), (i_k * BK, i_v * BV), (BK, BV), (1, 0)) - b_h += tl.load(p_h, boundary_check=(0, 1)).to(tl.float32) - if NORMK: - p_z0 = tl.make_block_ptr(z + i_bh * T*K, (T * K,), (1,), (i_k * BK,), (BK,), (0,)) - else: - p_z0 = tl.make_block_ptr(z + i_bh * T*V, (T * V,), (1,), (i_v * BV,), (BV,), (0,)) - b_zp = tl.load(p_z0).to(tl.float32) - for i_t in range(NT): - p_k = tl.make_block_ptr(k + i_bh * T*K, (K, T), (1, K), (i_k * BK, i_t * BT), (BK, BT), (0, 1)) - p_v = tl.make_block_ptr(v + i_bh * T*V, (T, V), (V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) - p_h = tl.make_block_ptr(h + i_bh * NT*K*V + i_t * K * V, (K, V), (V, 1), (i_k * BK, i_v * BV), (BK, BV), (1, 0)) - - tl.store(p_h, b_h.to(p_h.dtype.element_ty), boundary_check=(0, 1)) - # [BK, BT] - b_k = tl.load(p_k, boundary_check=(0, 1)) - # [BT, BV] - b_v = tl.load(p_v, boundary_check=(0, 1)) - if NORMK: - p_zc = tl.make_block_ptr(z + i_bh * T*K, (T * K,), (1,), ((i_t * BT + BT - 1) * K + i_k * BK,), (BK,), (0,)) - # [BK,] - b_zc = tl.load(p_zc, boundary_check=(0,)) - b_r, b_zp = exp(b_zp - b_zc), b_zc - # [BK, BV] - b_h = b_h * b_r[:, None] - b_k = exp(b_k - b_zc[:, None]).to(b_k.dtype) - else: - p_zc = tl.make_block_ptr(z + i_bh * T*V, (T * V,), (1,), ((i_t * BT + BT - 1) * V + i_v * BV,), (BV,), (0,)) - # [BV,] - b_zc = tl.load(p_zc, boundary_check=(0,)) - b_r, b_zp = exp(b_zp - b_zc), b_zc - # [BK, BV] - b_h = b_h * b_r[None, :] - b_v = exp(b_v - b_zc[None, :]).to(b_v.dtype) - # [BK, BV] - b_h += tl.dot(b_k, b_v, allow_tf32=False) - - if STORE_FINAL_STATE: - p_h = tl.make_block_ptr(ht + i_bh * K * V, (K, V), (V, 1), (i_k * BK, i_v * BV), (BK, BV), (1, 0)) - tl.store(p_h, b_h.to(p_h.dtype.element_ty), boundary_check=(0, 1)) - - -@triton.jit(do_not_specialize=['T']) -def chunk_abc_fwd_kernel_intra_K( - v, - z, - o, - A, - T, - V: tl.constexpr, - BT: tl.constexpr, - BC: tl.constexpr, - BV: tl.constexpr, - NC: tl.constexpr, -): - i_v, i_c, i_bh = tl.program_id(0), tl.program_id(1), tl.program_id(2) - i_t, i_i = i_c // NC, i_c % NC - - p_z = tl.make_block_ptr(z + i_bh * T*V, (T, V), (V, 1), (i_t * BT + i_i * BC, i_v * BV), (BC, BV), (1, 0)) - p_zn = tl.make_block_ptr(z + i_bh * T*V, (T * V,), (1,), ((i_t * BT + i_i * BC) * V + i_v * BV,), (BV,), (0,)) - # [BV,] - b_zn = tl.load(p_zn, boundary_check=(0,)) - # [BC, BV] - b_o = tl.zeros([BC, BV], dtype=tl.float32) - for i_j in range(0, i_i): - p_A = tl.make_block_ptr(A + i_bh * T * BT, (T, BT), (BT, 1), (i_t * BT + i_i * BC, i_j * BC), (BC, BC), (1, 0)) - p_v = tl.make_block_ptr(v + i_bh * T*V, (T, V), (V, 1), (i_t * BT + i_j * BC, i_v * BV), (BC, BV), (1, 0)) - # [BC, BV] - b_v = tl.load(p_v, boundary_check=(0, 1)) - # [BC, BC] - b_A = tl.load(p_A, boundary_check=(0, 1)) - b_o += tl.dot(b_A, exp(b_v - b_zn[None, :]).to(b_v.dtype), allow_tf32=False) - b_z = tl.load(p_z, boundary_check=(0, 1)) - b_o *= exp(b_zn[None, :] - b_z) - - o_i = tl.arange(0, BC) - o_A = i_bh * T * BT + (i_t * BT + i_i * BC + tl.arange(0, BC)) * BT + i_i * BC - m_A = (i_t * BT + i_i * BC + tl.arange(0, BC)) < T - for j in range(0, BC): - p_v = tl.make_block_ptr(v + i_bh * T*V, (T * V,), (1,), ((i_t * BT + i_i * BC + j) * V + i_v * BV,), (BV,), (0,)) - # [BC,] - b_A = tl.load(A + o_A + j, mask=m_A, other=0) - # [BV,] - b_v = tl.load(p_v, boundary_check=(0,)).to(tl.float32) - # [BC, BV] - # avoid 0 * inf = inf - m_i = o_i[:, None] >= j - b_o += tl.where(m_i, b_A[:, None] * exp(b_v[None, :] - b_z), 0) - p_o = tl.make_block_ptr(o + i_bh * T*V, (T, V), (V, 1), (i_t * BT + i_i * BC, i_v * BV), (BC, BV), (1, 0)) - tl.store(p_o, b_o.to(p_o.dtype.element_ty), boundary_check=(0, 1)) - - -@triton.jit(do_not_specialize=['T']) -def chunk_abc_fwd_kernel_K( - q, - k, - z, - h, - o, - A, - scale, - T, - K: tl.constexpr, - V: tl.constexpr, - BT: tl.constexpr, - BK: tl.constexpr, - BV: tl.constexpr, - NT: tl.constexpr, -): - i_v, i_t, i_bh = tl.program_id(0), tl.program_id(1), tl.program_id(2) - i_p = tl.maximum(i_t * BT - 1, 0) - - o_i = tl.arange(0, BT) - m_s = o_i[:, None] >= o_i[None, :] - - b_o = tl.zeros([BT, BV], dtype=tl.float32) - b_A = tl.zeros([BT, BT], dtype=tl.float32) - for i_k in range(tl.cdiv(K, BK)): - p_q = tl.make_block_ptr(q + i_bh * T*K, (T, K), (K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) - p_k = tl.make_block_ptr(k + i_bh * T*K, (K, T), (1, K), (i_k * BK, i_t * BT), (BK, BT), (0, 1)) - p_h = tl.make_block_ptr(h + i_bh * NT*K*V + i_t * K * V, (K, V), (V, 1), (i_k * BK, i_v * BV), (BK, BV), (1, 0)) - - # [BT, BK] - b_q = tl.load(p_q, boundary_check=(0, 1)) - b_q = (b_q * scale).to(b_q.dtype) - # [BK, BT] - b_k = tl.load(p_k, boundary_check=(0, 1)) - # [BK, BV] - b_h = tl.load(p_h, boundary_check=(0, 1)) - # [BT, BV] - b_o += tl.dot(b_q, b_h, allow_tf32=False) - # [BT, BT] - b_A += tl.dot(b_q, b_k, allow_tf32=False) - p_z = tl.make_block_ptr(z + i_bh * T*V, (T, V), (V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) - p_o = tl.make_block_ptr(o + i_bh * T*V, (T, V), (V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) - # [BT, BV] - b_z = tl.load(p_z, boundary_check=(0, 1)) - # [BT, BV] - p_zp = tl.make_block_ptr(z + i_bh * T*V, (T * V,), (1,), (i_p * V + i_v * BV,), (BV,), (0,)) - b_zp = tl.load(p_zp, boundary_check=(0,)) - b_o = b_o * exp(b_zp[None, :] - b_z) - tl.store(p_o, b_o.to(p_o.dtype.element_ty), boundary_check=(0, 1)) - - p_A = tl.make_block_ptr(A + i_bh * T * BT, (T, BT), (BT, 1), (i_t * BT, 0), (BT, BT), (1, 0)) - # [BT, BT] - b_A = tl.where(m_s, b_A, 0.) - if i_v == 0: - tl.store(p_A, b_A.to(p_A.dtype.element_ty), boundary_check=(0, 1)) - - -@triton.jit(do_not_specialize=['T']) -def chunk_abc_fwd_kernel_intra_V( - q, - k, - z, - A, - scale, - T, - K: tl.constexpr, - BT: tl.constexpr, - BC: tl.constexpr, - BK: tl.constexpr, - NC: tl.constexpr, -): - i_k, i_c, i_bh = tl.program_id(0), tl.program_id(1), tl.program_id(2) - i_t, i_i, i_j = i_c // (NC * NC), (i_c % (NC * NC)) // NC, (i_c % (NC * NC)) % NC - n_bh = tl.num_programs(2) - - if i_i > i_j: - p_q = tl.make_block_ptr(q + i_bh * T*K, (T, K), (K, 1), (i_t * BT + i_i * BC, i_k * BK), (BC, BK), (1, 0)) - p_k = tl.make_block_ptr(k + i_bh * T*K, (K, T), (1, K), (i_k * BK, i_t * BT + i_j * BC), (BK, BC), (0, 1)) - p_z = tl.make_block_ptr(z + i_bh * T*K, (T, K), (K, 1), (i_t * BT + i_i * BC, i_k * BK), (BC, BK), (1, 0)) - p_A = tl.make_block_ptr(A + (i_k*n_bh+i_bh)*T*BT, (T, BT), (BT, 1), (i_t * BT + i_i * BC, i_j * BC), (BC, BC), (1, 0)) - p_zn = tl.make_block_ptr(z + i_bh * T*K, (T * K,), (1,), ((i_t * BT + i_i * BC) * K + i_k * BK,), (BK,), (0,)) - # [BK,] - b_zn = tl.load(p_zn, boundary_check=(0,)) - # [BC, BK] - b_q = tl.load(p_q, boundary_check=(0, 1)) - b_z = tl.load(p_z, boundary_check=(0, 1)) - b_q = (b_q * exp(b_zn[None, :] - b_z) * scale).to(b_q.dtype) - # [BK, BC] - b_k = tl.load(p_k, boundary_check=(0, 1)) - b_k = exp(b_k - b_zn[:, None]).to(b_k.dtype) - # [BC, BC] - b_A = tl.dot(b_q, b_k, allow_tf32=False) - tl.store(p_A, b_A.to(A.dtype.element_ty), boundary_check=(0, 1)) - elif i_i == i_j: - p_q = tl.make_block_ptr(q + i_bh * T*K, (T, K), (K, 1), (i_t * BT + i_i * BC, i_k * BK), (BC, BK), (1, 0)) - p_k = tl.make_block_ptr(k + i_bh * T*K, (T * K,), (1,), ((i_t * BT + i_j * BC) * K + i_k * BK,), (BK,), (0,)) - p_z = tl.make_block_ptr(z + i_bh * T*K, (T, K), (K, 1), (i_t * BT + i_i * BC, i_k * BK), (BC, BK), (1, 0)) - # [BC, BK] - b_q = tl.load(p_q, boundary_check=(0, 1)) - b_z = tl.load(p_z, boundary_check=(0, 1)) - - o_i = tl.arange(0, BC) - o_A = (i_bh + i_k * n_bh) * T * BT + (i_t * BT + i_i * BC + tl.arange(0, BC)) * BT + i_j * BC - m_A = (i_t * BT + i_i * BC + tl.arange(0, BC)) < T - for j in range(0, BC): - # [BK,] - b_k = tl.load(p_k, boundary_check=(0,)).to(tl.float32) - # [BC,] - b_A = tl.sum(b_q * exp(b_k[None, :] - b_z) * scale, 1) - b_A = tl.where(o_i >= j, b_A, 0.) - tl.store(A + o_A + j, b_A.to(b_q.dtype), mask=m_A) - - p_k = tl.advance(p_k, (K,)) - - -@triton.jit(do_not_specialize=['T']) -def chunk_abc_fwd_kernel_V( - q, - v, - z, - h, - o, - A, - scale, - T, - K: tl.constexpr, - V: tl.constexpr, - BT: tl.constexpr, - BK: tl.constexpr, - BV: tl.constexpr, - NT: tl.constexpr, -): - i_v, i_t, i_bh = tl.program_id(0), tl.program_id(1), tl.program_id(2) - i_p = tl.maximum(i_t * BT - 1, 0) - - b_o = tl.zeros([BT, BV], dtype=tl.float32) - for i_k in range(tl.cdiv(K, BK)): - p_q = tl.make_block_ptr(q + i_bh * T*K, (T, K), (K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) - p_z = tl.make_block_ptr(z + i_bh * T*K, (T, K), (K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) - p_h = tl.make_block_ptr(h + i_bh * NT*K*V + i_t * K * V, (K, V), (V, 1), (i_k * BK, i_v * BV), (BK, BV), (1, 0)) - p_zp = tl.make_block_ptr(z + i_bh * T*K, (T * K,), (1,), (i_p * K + i_k * BK,), (BK,), (0,)) - - # [BT, BK] - b_q = tl.load(p_q, boundary_check=(0, 1)) - b_q = (b_q * scale).to(b_q.dtype) - # [BT, BK] - b_z = tl.load(p_z, boundary_check=(0, 1)) - # [BT, BK] - b_zp = tl.load(p_zp, boundary_check=(0,)) - b_q = (b_q * exp(b_zp[None, :] - b_z)).to(b_q.dtype) - # [BK, BV] - b_h = tl.load(p_h, boundary_check=(0, 1)) - # works but dkw, owing to divine benevolence - # [BT, BV] - if i_k >= 0: - b_o += tl.dot(b_q, b_h, allow_tf32=False) - p_v = tl.make_block_ptr(v + i_bh * T*V, (T, V), (V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) - p_o = tl.make_block_ptr(o + i_bh * T*V, (T, V), (V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) - p_A = tl.make_block_ptr(A + i_bh * T * BT, (T, BT), (BT, 1), (i_t * BT, 0), (BT, BT), (1, 0)) - # [BT, BV] - b_v = tl.load(p_v, boundary_check=(0, 1)) - # [BT, BT] - b_A = tl.load(p_A, boundary_check=(0, 1)) - b_o += tl.dot(b_A.to(b_v.dtype), b_v, allow_tf32=False) - tl.store(p_o, b_o.to(p_o.dtype.element_ty), boundary_check=(0, 1)) - - -@triton.jit(do_not_specialize=['T']) -def chunk_abc_bwd_kernel_dh( - q, - z, - do, - dh, - scale, - T, - K: tl.constexpr, - V: tl.constexpr, - BT: tl.constexpr, - BK: tl.constexpr, - BV: tl.constexpr, - NT: tl.constexpr, - NORMK: tl.constexpr, -): - i_k, i_v, i_bh = tl.program_id(0), tl.program_id(1), tl.program_id(2) - - b_dh = tl.zeros([BK, BV], dtype=tl.float32) - b_zp = tl.full([BK if NORMK else BV], float('inf'), dtype=tl.float32) - for i_t in range(NT - 1, -1, -1): - i_p = tl.maximum(i_t * BT - 1, 0) - p_q = tl.make_block_ptr(q + i_bh * T*K, (K, T), (1, K), (i_k * BK, i_t * BT), (BK, BT), (0, 1)) - p_do = tl.make_block_ptr(do + i_bh * T*V, (T, V), (V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) - p_dh = tl.make_block_ptr(dh + i_bh * NT*K*V + i_t * K*V, (K, V), (V, 1), (i_k * BK, i_v * BV), (BK, BV), (1, 0)) - - # [BK, BT] - b_q = tl.load(p_q, boundary_check=(0, 1)) - b_q = (b_q * scale).to(b_q.dtype) - # [BT, BV] - b_do = tl.load(p_do, boundary_check=(0, 1)) - - tl.store(p_dh, b_dh.to(p_dh.dtype.element_ty), boundary_check=(0, 1)) - if NORMK: - p_z = tl.make_block_ptr(z + i_bh * T*K, (K, T), (1, K), (i_k * BK, i_t * BT), (BK, BT), (0, 1)) - p_zc = tl.make_block_ptr(z + i_bh * T*K, (T * K,), (1,), (i_p * K + i_k * BK,), (BK,), (0,)) - # [BK,] - b_zc = tl.load(p_zc, boundary_check=(0,)) - b_r, b_zp = exp(b_zc - b_zp), b_zc - # [BK, BT] - b_z = tl.load(p_z, boundary_check=(0, 1)) - b_q = (b_q * exp(b_zc[:, None] - b_z)).to(b_q.dtype) - # [BK, BV] - b_dh = b_dh * b_r[:, None] - else: - p_z = tl.make_block_ptr(z + i_bh * T*V, (T, V), (V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) - p_zc = tl.make_block_ptr(z + i_bh * T*V, (T * V,), (1,), (i_p * V + i_v * BV,), (BV,), (0,)) - # [BV,] - b_zc = tl.load(p_zc, boundary_check=(0,)) - b_r, b_zp = exp(b_zc - b_zp), b_zc - # [BT, BV] - b_z = tl.load(p_z, boundary_check=(0,)) - b_do = (b_do * exp(b_zc[None, :] - b_z)).to(b_do.dtype) - # [BK, BV] - b_dh = b_dh * b_r[None, :] - # [BK, BV] - b_dh += tl.dot(b_q, b_do, allow_tf32=False) - - -@triton.jit(do_not_specialize=['T']) -def chunk_abc_bwd_kernel_V( - k, - v, - z, - h, - A, - do, - dh, - dq, - dk, - dv, - dA, - scale, - T, - K: tl.constexpr, - V: tl.constexpr, - BT: tl.constexpr, - BK: tl.constexpr, - BV: tl.constexpr, - NT: tl.constexpr, -): - i_k, i_t, i_bh = tl.program_id(0), tl.program_id(1), tl.program_id(2) - i_p = tl.maximum(i_t * BT - 1, 0) - n_bh = tl.num_programs(2) - - p_k = tl.make_block_ptr(k + i_bh * T*K, (T, K), (K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) - p_zc = tl.make_block_ptr(z + i_bh * T*K, (T * K,), (1,), ((i_t * BT + BT - 1) * K + i_k * BK,), (BK,), (0,)) - p_A = tl.make_block_ptr(A + i_bh * T * BT, (BT, T), (1, BT), (0, i_t * BT), (BT, BT), (0, 1)) - - # [BK,] - b_zc = tl.load(p_zc, boundary_check=(0,)) - # [BT, BK] - b_k = tl.load(p_k, boundary_check=(0, 1)) - b_k = exp(b_k - b_zc[None, :]).to(b_k.dtype) - # [BT, BT] - b_A = tl.load(p_A, boundary_check=(0, 1)) - - b_dq = tl.zeros([BT, BK], dtype=tl.float32) - b_dk = tl.zeros([BT, BK], dtype=tl.float32) - b_dA = tl.zeros([BT, BT], dtype=tl.float32) - for i_v in range(tl.cdiv(V, BV)): - p_v = tl.make_block_ptr(v + i_bh * T*V, (T, V), (V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) - p_h = tl.make_block_ptr(h + i_bh * NT*K*V + i_t * V * K, (V, K), (1, V), (i_v * BV, i_k * BK), (BV, BK), (0, 1)) - p_do = tl.make_block_ptr(do + i_bh * T*V, (T, V), (V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) - p_dh = tl.make_block_ptr(dh + i_bh * NT*K*V + i_t * K*V, (K, V), (V, 1), (i_k * BK, i_v * BV), (BK, BV), (1, 0)) - p_dv = tl.make_block_ptr(dv + (i_k*n_bh+i_bh) * T*V, (T, V), (V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) - - # [BT, BV] - b_v = tl.load(p_v, boundary_check=(0, 1)) - # [BV, BK] - b_h = tl.load(p_h, boundary_check=(0, 1)) - # [BT, BV] - b_do = tl.load(p_do, boundary_check=(0, 1)) - # [BK, BV] - b_dh = tl.load(p_dh, boundary_check=(0, 1)) - - # [BT, BV] - b_dv = tl.dot(b_k, b_dh, allow_tf32=False) - if i_k == 0: - b_dv += tl.dot(b_A.to(b_do.dtype), b_do, allow_tf32=False) - b_do = (b_do * scale).to(b_do.dtype) - tl.store(p_dv, b_dv.to(p_dv.dtype.element_ty), boundary_check=(0, 1)) - # [BT, BT] - b_dA += tl.dot(b_do, tl.trans(b_v), allow_tf32=False) - # [BT, BK] - b_dq += tl.dot(b_do, b_h, allow_tf32=False) - # [BT, BK] - b_dk += tl.dot(b_v, tl.trans(b_dh), allow_tf32=False) - p_z = tl.make_block_ptr(z + i_bh * T*K, (T, K), (K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) - p_zp = tl.make_block_ptr(z + i_bh * T*K, (T * K,), (1,), (i_p * K + i_k * BK,), (BK,), (0,)) - # [BK,] - b_zp = tl.load(p_zp, boundary_check=(0,)) - # [BT, BK] - b_z = tl.load(p_z, boundary_check=(0, 1)) - b_z = exp(b_zp[None, :] - b_z) - # [BT, BK] - b_dq = b_dq * b_z - b_dk = b_dk * b_k - - p_dq = tl.make_block_ptr(dq + i_bh * T*K, (T, K), (K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) - p_dk = tl.make_block_ptr(dk + i_bh * T*K, (T, K), (K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) - p_dA = tl.make_block_ptr(dA + i_bh * T * BT, (T, BT), (BT, 1), (i_t * BT, 0), (BT, BT), (1, 0)) - tl.store(p_dq, b_dq.to(p_dq.dtype.element_ty), boundary_check=(0, 1)) - tl.store(p_dk, b_dk.to(p_dk.dtype.element_ty), boundary_check=(0, 1)) - - o_i = tl.arange(0, BT) - m_s = o_i[:, None] >= o_i[None, :] - # [BT, BT] - b_dA = tl.where(m_s, b_dA, 0.).to(b_k.dtype) - if i_k == 0: - tl.store(p_dA, b_dA.to(p_dA.dtype.element_ty), boundary_check=(0, 1)) - - -@triton.jit(do_not_specialize=['T']) -def chunk_abc_bwd_kernel_intra_V( - q, - k, - z, - dA, - dq, - dk, - T, - K: tl.constexpr, - BT: tl.constexpr, - BC: tl.constexpr, - BK: tl.constexpr, - NC: tl.constexpr, -): - i_k, i_c, i_bh = tl.program_id(0), tl.program_id(1), tl.program_id(2) - i_t, i_i = i_c // NC, i_c % NC - - p_z = tl.make_block_ptr(z + i_bh * T*K, (T, K), (K, 1), (i_t * BT + i_i * BC, i_k * BK), (BC, BK), (1, 0)) - p_zn = tl.make_block_ptr(z + i_bh * T*K, (T * K,), (1,), ((i_t * BT + i_i * BC) * K + i_k * BK,), (BK,), (0,)) - # [BK,] - b_zn = tl.load(p_zn, boundary_check=(0,)) - # [BC, BK] - b_z = tl.load(p_z, boundary_check=(0, 1)) - b_zq = exp(b_zn[None, :] - b_z) - b_dq = tl.zeros([BC, BK], dtype=tl.float32) - for i_j in range(0, i_i): - p_k = tl.make_block_ptr(k + i_bh * T*K, (T, K), (K, 1), (i_t * BT + i_j * BC, i_k * BK), (BC, BK), (1, 0)) - p_dA = tl.make_block_ptr(dA + i_bh * T * BT, (T, BT), (BT, 1), (i_t * BT + i_i * BC, i_j * BC), (BC, BC), (1, 0)) - # [BC, BK] - b_k = tl.load(p_k, boundary_check=(0, 1)) - b_kz = exp(b_k - b_zn[None, :]).to(b_k.dtype) - # [BC, BC] - b_dA = tl.load(p_dA, boundary_check=(0, 1)) - # [BC, BK] - b_dq += tl.dot(b_dA, b_kz, allow_tf32=False) - b_dq *= b_zq - - o_i = tl.arange(0, BC) - o_dA = i_bh * T * BT + (i_t * BT + i_i * BC + tl.arange(0, BC)) * BT + i_i * BC - m_dA = (i_t * BT + i_i * BC + tl.arange(0, BC)) < T - for j in range(0, BC): - p_kj = tl.make_block_ptr(k + i_bh * T*K, (T * K,), (1,), ((i_t * BT + i_i*BC+j) * K + i_k * BK,), (BK,), (0,)) - # [BC,] - b_dA = tl.load(dA + o_dA + j, mask=m_dA, other=0) - # [BK,] - b_kj = tl.load(p_kj, boundary_check=(0,)).to(tl.float32) - # [BC, BK] - m_i = o_i[:, None] >= j - # [BC, BK] - b_dq += tl.where(m_i, b_dA[:, None] * exp(b_kj[None, :] - b_z), 0.) - p_dq = tl.make_block_ptr(dq + i_bh * T*K, (T, K), (K, 1), (i_t * BT + i_i * BC, i_k * BK), (BC, BK), (1, 0)) - tl.store(p_dq, b_dq.to(p_dq.dtype.element_ty), boundary_check=(0, 1)) - - tl.debug_barrier() - p_k = tl.make_block_ptr(k + i_bh * T*K, (T, K), (K, 1), (i_t * BT + i_i * BC, i_k * BK), (BC, BK), (1, 0)) - p_zn = tl.make_block_ptr(z + i_bh * T*K, (T*K,), (1,), ((i_t * BT + i_i * BC + BC - 1) * K + i_k * BK,), (BK,), (0,)) - # [BK,] - b_zn = tl.load(p_zn, boundary_check=(0,)) - # [BC, BK] - b_k = tl.load(p_k, boundary_check=(0, 1)) - b_kz = exp(b_k - b_zn[None, :]) - b_dk = tl.zeros([BC, BK], dtype=tl.float32) - for i_j in range(i_i + 1, NC): - p_q = tl.make_block_ptr(q + i_bh * T*K, (T, K), (K, 1), (i_t * BT + i_j * BC, i_k * BK), (BC, BK), (1, 0)) - p_z = tl.make_block_ptr(z + i_bh * T*K, (T, K), (K, 1), (i_t * BT + i_j * BC, i_k * BK), (BC, BK), (1, 0)) - p_dA = tl.make_block_ptr(dA + i_bh * T * BT, (T, BT), (BT, 1), (i_t * BT + i_j * BC, i_i * BC), (BC, BC), (1, 0)) - # [BC, BK] - b_q = tl.load(p_q, boundary_check=(0, 1)) - b_z = tl.load(p_z, boundary_check=(0, 1)) - b_qz = (b_q * exp(b_zn[None, :] - b_z)).to(b_q.dtype) - # [BC, BC] - b_dA = tl.load(p_dA, boundary_check=(0, 1)) - # [BC, BK] - b_dk += tl.dot(tl.trans(b_dA), b_qz, allow_tf32=False) - b_dk *= b_kz - - o_dA = i_bh * T * BT + (i_t * BT + i_i * BC) * BT + i_i * BC + tl.arange(0, BC) - for j in range(0, BC): - p_qj = tl.make_block_ptr(q + i_bh * T*K, (T * K,), (1,), ((i_t * BT + i_i * BC + j) * K + i_k * BK,), (BK,), (0,)) - p_zj = tl.make_block_ptr(z + i_bh * T*K, (T * K,), (1,), ((i_t * BT + i_i * BC + j) * K + i_k * BK,), (BK,), (0,)) - # [BC,] - b_dA = tl.load(dA + o_dA + j * BT, mask=(i_t * BT + i_i * BC + j < T), other=0) - # [BK,] - b_qj = tl.load(p_qj, boundary_check=(0,)).to(tl.float32) - b_zj = tl.load(p_zj, boundary_check=(0,)).to(tl.float32) - # [BC, BK] - m_i = o_i[:, None] <= j - b_dk += tl.where(m_i, b_dA[:, None] * b_qj[None, :] * exp(b_k - b_zj[None, :]), 0.) - p_dk = tl.make_block_ptr(dk + i_bh * T*K, (T, K), (K, 1), (i_t * BT + i_i * BC, i_k * BK), (BC, BK), (1, 0)) - tl.store(p_dk, b_dk.to(p_dk.dtype.element_ty), boundary_check=(0, 1)) - - -@triton.jit(do_not_specialize=['T']) -def chunk_abc_bwd_kernel_intra_K( - v, - z, - do, - dA, - scale, - T, - V: tl.constexpr, - BT: tl.constexpr, - BC: tl.constexpr, - BV: tl.constexpr, - NC: tl.constexpr, -): - i_v, i_c, i_bh = tl.program_id(0), tl.program_id(1), tl.program_id(2) - i_t, i_i, i_j = i_c // (NC * NC), (i_c % (NC * NC)) // NC, (i_c % (NC * NC)) % NC - n_bh = tl.num_programs(2) - - if i_i > i_j: - p_v = tl.make_block_ptr(v + i_bh * T*V, (V, T), (1, V), (i_v * BV, i_t * BT + i_j * BC), (BV, BC), (0, 1)) - p_z = tl.make_block_ptr(z + i_bh * T*V, (T, V), (V, 1), (i_t * BT + i_i * BC, i_v * BV), (BC, BV), (1, 0)) - p_zn = tl.make_block_ptr(z + i_bh * T*V, (T * V,), (1,), ((i_t * BT + i_i * BC) * V + i_v * BV,), (BV,), (0,)) - p_do = tl.make_block_ptr(do + i_bh * T*V, (T, V), (V, 1), (i_t * BT + i_i * BC, i_v * BV), (BC, BV), (1, 0)) - p_dA = tl.make_block_ptr(dA+(i_bh+i_v*n_bh)*T*BT, (T, BT), (BT, 1), (i_t * BT + i_i * BC, i_j * BC), (BC, BC), (1, 0)) - # [BV,] - b_zn = tl.load(p_zn, boundary_check=(0,)) - # [BC, BV] - b_z = tl.load(p_z, boundary_check=(0, 1)) - b_do = tl.load(p_do, boundary_check=(0, 1)) - b_do = (b_do * exp(b_zn[None, :] - b_z) * scale).to(b_do.dtype) - # [BV, BC] - b_v = tl.load(p_v, boundary_check=(0, 1)) - b_v = exp(b_v - b_zn[:, None]).to(b_v.dtype) - # [BC, BC] - b_dA = tl.dot(b_do, b_v, allow_tf32=False) - tl.store(p_dA, b_dA.to(dA.dtype.element_ty), boundary_check=(0, 1)) - elif i_i == i_j: - p_v = tl.make_block_ptr(v + i_bh * T*V, (T * V,), (1,), ((i_t * BT + i_j * BC) * V + i_v * BV,), (BV,), (0,)) - p_z = tl.make_block_ptr(z + i_bh * T*V, (T, V), (V, 1), (i_t * BT + i_i * BC, i_v * BV), (BC, BV), (1, 0)) - p_do = tl.make_block_ptr(do + i_bh * T*V, (T, V), (V, 1), (i_t * BT + i_i * BC, i_v * BV), (BC, BV), (1, 0)) - # [BC, BV] - b_z = tl.load(p_z, boundary_check=(0, 1)) - b_do = tl.load(p_do, boundary_check=(0, 1)) * scale - - o_i = tl.arange(0, BC) - o_A = (i_bh + i_v * n_bh) * T * BT + (i_t * BT + i_i * BC + tl.arange(0, BC)) * BT + i_j * BC - m_A = (i_t * BT + i_i * BC + tl.arange(0, BC)) < T - for j in range(0, BC): - # [BV,] - b_v = tl.load(p_v, boundary_check=(0,)).to(tl.float32) - # [BC,] - b_dA = tl.sum(b_do * exp(b_v[None, :] - b_z), 1) - b_dA = tl.where(o_i >= j, b_dA, 0) - tl.store(dA + o_A + j, b_dA.to(b_do.dtype), mask=m_A) - - p_v = tl.advance(p_v, (V,)) - - -@triton.jit(do_not_specialize=['T']) -def chunk_abc_bwd_kernel_K( - q, - k, - v, - z, - h, - A, - do, - dh, - dq, - dk, - dv, - dA, - scale, - T, - K: tl.constexpr, - V: tl.constexpr, - BT: tl.constexpr, - BK: tl.constexpr, - BV: tl.constexpr, - NT: tl.constexpr, -): - i_k, i_t, i_bh = tl.program_id(0), tl.program_id(1), tl.program_id(2) - i_p = tl.maximum(i_t * BT - 1, 0) - n_bh = tl.num_programs(2) - - o_i = tl.arange(0, BT) - m_s = o_i[:, None] >= o_i[None, :] - - p_q = tl.make_block_ptr(q + i_bh * T*K, (T, K), (K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) - p_k = tl.make_block_ptr(k + i_bh * T*K, (T, K), (K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) - p_A = tl.make_block_ptr(A + (i_k*n_bh+i_bh) * T * BT, (T, BT), (BT, 1), (i_t * BT, 0), (BT, BT), (1, 0)) - - # [BT, BK] - b_q = tl.load(p_q, boundary_check=(0, 1)) - b_k = tl.load(p_k, boundary_check=(0, 1)) - # [BT, BT] - b_A = tl.dot((b_q * scale).to(b_q.dtype), tl.trans(b_k), allow_tf32=False) - b_A = tl.where(m_s, b_A, 0.) - tl.store(p_A, b_A.to(p_A.dtype.element_ty), boundary_check=(0, 1)) - - b_dq = tl.zeros([BT, BK], dtype=tl.float32) - b_dk = tl.zeros([BT, BK], dtype=tl.float32) - for i_v in range(tl.cdiv(V, BV)): - p_v = tl.make_block_ptr(v + i_bh * T*V, (T, V), (V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) - p_z = tl.make_block_ptr(z + i_bh * T*V, (T, V), (V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) - p_zp = tl.make_block_ptr(z + i_bh * T*V, (T * V,), (1,), (i_p * V + i_v * BV,), (BV,), (0,)) - p_zc = tl.make_block_ptr(z + i_bh * T*V, (T * V,), (1,), ((i_t * BT + BT - 1) * V + i_v * BV,), (BV,), (0,)) - p_h = tl.make_block_ptr(h + i_bh * NT*K*V + i_t * K*V, (V, K), (1, V), (i_v * BV, i_k * BK), (BV, BK), (0, 1)) - - p_do = tl.make_block_ptr(do + i_bh * T*V, (T, V), (V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) - p_dh = tl.make_block_ptr(dh + i_bh * NT*K*V + i_t * K*V, (K, V), (V, 1), (i_k * BK, i_v * BV), (BK, BV), (1, 0)) - p_dv = tl.make_block_ptr(dv + (i_k*n_bh+i_bh) * T*V, (T, V), (V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) - - # [BV,] - b_zp = tl.load(p_zp, boundary_check=(0,)) - b_zc = tl.load(p_zc, boundary_check=(0,)) - # [BT, BV] - b_v = tl.load(p_v, boundary_check=(0, 1)) - b_v = exp(b_v - b_zc[None, :]).to(b_v.dtype) - b_z = tl.load(p_z, boundary_check=(0, 1)) - b_z = exp(b_zp[None, :] - b_z) - # [BV, BK] - b_h = tl.load(p_h, boundary_check=(0, 1)) - # [BT, BV] - b_do = tl.load(p_do, boundary_check=(0, 1)) - b_do = (b_do * b_z * scale).to(b_do.dtype) - # [BK, BV] - b_dh = tl.load(p_dh, boundary_check=(0, 1)) - - # [BT, BK] - b_dq += tl.dot(b_do, b_h, allow_tf32=False) - b_dk += tl.dot(b_v, tl.trans(b_dh), allow_tf32=False) - # [BT, BV] - b_dv = b_v * tl.dot(b_k, b_dh, allow_tf32=False) - tl.store(p_dv, b_dv.to(p_dv.dtype.element_ty), boundary_check=(0, 1)) - p_dA = tl.make_block_ptr(dA + i_bh * T * BT, (T, BT), (BT, 1), (i_t * BT, 0), (BT, BT), (1, 0)) - # [BT, BT] - b_dA = tl.load(p_dA, boundary_check=(0, 1)) - # [BT, BK] - b_dq += tl.dot(b_dA, b_k, allow_tf32=False) - b_dk += tl.dot(tl.trans(b_dA).to(b_k.dtype), b_q, allow_tf32=False) - - p_dq = tl.make_block_ptr(dq + i_bh * T*K, (T, K), (K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) - p_dk = tl.make_block_ptr(dk + i_bh * T*K, (T, K), (K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) - tl.store(p_dq, b_dq.to(p_dq.dtype.element_ty), boundary_check=(0, 1)) - tl.store(p_dk, b_dk.to(p_dk.dtype.element_ty), boundary_check=(0, 1)) - - -@triton.jit(do_not_specialize=['T']) -def chunk_abc_bwd_kernel_intra_KV( - v, - z, - A, - do, - dv, - T, - V: tl.constexpr, - BT: tl.constexpr, - BC: tl.constexpr, - BV: tl.constexpr, - NC: tl.constexpr, -): - i_v, i_c, i_bh = tl.program_id(0), tl.program_id(1), tl.program_id(2) - i_t, i_i = i_c // NC, i_c % NC - - p_v = tl.make_block_ptr(v + i_bh * T*V, (T, V), (V, 1), (i_t * BT + i_i * BC, i_v * BV), (BC, BV), (1, 0)) - p_zn = tl.make_block_ptr(z + i_bh * T*V, (T*V,), (1,), ((i_t * BT + i_i * BC + BC - 1) * V + i_v * BV,), (BV,), (0,)) - # [BV,] - b_zn = tl.load(p_zn, boundary_check=(0,)) - # [BC, BV] - b_v = tl.load(p_v, boundary_check=(0, 1)) - b_dv = tl.zeros([BC, BV], dtype=tl.float32) - for i_j in range(i_i + 1, NC): - p_z = tl.make_block_ptr(z + i_bh * T*V, (T, V), (V, 1), (i_t * BT + i_j * BC, i_v * BV), (BC, BV), (1, 0)) - p_A = tl.make_block_ptr(A + i_bh * T * BT, (BT, T), (1, BT), (i_i * BC, i_t * BT + i_j * BC), (BC, BC), (0, 1)) - p_do = tl.make_block_ptr(do + i_bh * T*V, (T, V), (V, 1), (i_t * BT + i_j * BC, i_v * BV), (BC, BV), (1, 0)) - # [BC, BV] - b_z = tl.load(p_z, boundary_check=(0, 1)) - b_do = tl.load(p_do, boundary_check=(0, 1)) - b_do = (b_do * exp(b_zn[None, :] - b_z)).to(b_do.dtype) - # [BC, BC] - b_A = tl.load(p_A, boundary_check=(0, 1)) - b_dv += tl.dot(b_A, b_do, allow_tf32=False) - b_dv *= exp(b_v - b_zn[None, :]) - - o_i = tl.arange(0, BC) - for j in range(0, BC): - p_z = tl.make_block_ptr(z + i_bh * T*V, (T * V,), (1,), ((i_t * BT + i_i * BC + j) * V + i_v * BV,), (BV,), (0,)) - p_A = tl.make_block_ptr(A + i_bh * T * BT, (T * BT,), (1,), ((i_t * BT + i_i * BC + j) * BT + i_i * BC,), (BC,), (0,)) - p_do = tl.make_block_ptr(do + i_bh * T*V, (T * V,), (1,), ((i_t * BT + i_i * BC + j) * V + i_v * BV,), (BV,), (0,)) - # [BC,] - b_A = tl.load(p_A, boundary_check=(0,)) - # [BV,] - b_z = tl.load(p_z, boundary_check=(0,)) - b_do = tl.load(p_do, boundary_check=(0,)) - # [BC, BV] - m_i = o_i[:, None] <= j - b_dv += tl.where(m_i, exp(b_v - b_z[None, :]) * b_A[:, None] * b_do[None, :], 0.) - p_dv = tl.make_block_ptr(dv + i_bh * T*V, (T, V), (V, 1), (i_t * BT + i_i * BC, i_v * BV), (BC, BV), (1, 0)) - tl.store(p_dv, b_dv.to(p_dv.dtype.element_ty), boundary_check=(0, 1)) - - -@triton.jit(do_not_specialize=['T']) -def chunk_abc_bwd_kernel_rcum_inter( - s, - z, - ss, - doo, - T, - S: tl.constexpr, - BT: tl.constexpr, - BS: tl.constexpr, - NT: tl.constexpr, -): - i_m, i_bh = tl.program_id(0), tl.program_id(1) - - b_sp = tl.zeros([BS], dtype=tl.float32) - b_zp = tl.full([BS], float('inf'), dtype=tl.float32) - for i_t in range(NT - 1, -1, -1): - p_s = tl.make_block_ptr(s + i_bh * T*S, (T, S), (S, 1), (i_t * BT, i_m * BS), (BT, BS), (1, 0)) - p_z = tl.make_block_ptr(z + i_bh * T*S, (T, S), (S, 1), (i_t * BT, i_m * BS), (BT, BS), (1, 0)) - p_zc = tl.make_block_ptr(z + i_bh * T*S, (T*S,), (1,), ((i_t * BT) * S + i_m * BS,), (BS,), (0,)) - p_ss = tl.make_block_ptr(ss + i_bh * T*S, (T, S), (S, 1), (i_t * BT, i_m * BS), (BT, BS), (1, 0)) - p_doo = tl.make_block_ptr(doo + i_bh * T*S, (T, S), (S, 1), (i_t * BT, i_m * BS), (BT, BS), (1, 0)) - # [BS,] - b_zc = tl.load(p_zc, boundary_check=(0,)) - # [BT, BS] - b_s = tl.load(p_s, boundary_check=(0, 1)) - b_z = tl.load(p_z, boundary_check=(0, 1)) - b_ss = tl.load(p_ss, boundary_check=(0, 1)) - - b_doo = exp(b_s - b_zp[None, :]) * b_sp[None, :] - tl.store(p_doo, b_doo.to(p_doo.dtype.element_ty), boundary_check=(0, 1)) - # [BS,] - b_sp = b_sp * exp(b_zc - b_zp) + tl.sum(b_ss * exp(b_zc[None, :] - b_z), 0) - b_zp = b_zc - - -@triton.jit(do_not_specialize=['T']) -def chunk_abc_bwd_kernel_rcum_intra( - s, - z, - ss, - doo, - T, - S: tl.constexpr, - BT: tl.constexpr, - BC: tl.constexpr, - BS: tl.constexpr, - NC: tl.constexpr, -): - i_s, i_c, i_bh = tl.program_id(0), tl.program_id(1), tl.program_id(2) - i_t, i_i = i_c // NC, i_c % NC - - o_i = tl.arange(0, BC) - m_o = tl.full([BC, BC], 1., dtype=tl.float32) - - p_s = tl.make_block_ptr(s + i_bh * T*S, (T, S), (S, 1), (i_t * BT + i_i * BC, i_s * BS), (BC, BS), (1, 0)) - p_zn = tl.make_block_ptr(z + i_bh * T*S, (T*S,), (1,), ((i_t * BT + i_i * BC + BC - 1) * S + i_s * BS,), (BS,), (0,)) - p_doo = tl.make_block_ptr(doo + i_bh * T*S, (T, S), (S, 1), (i_t * BT + i_i * BC, i_s * BS), (BC, BS), (1, 0)) - # [BC, BS] - b_s = tl.load(p_s, boundary_check=(0, 1)) - # [BS,] - b_zn = tl.load(p_zn, boundary_check=(0,)) - - b_doo = tl.zeros([BC, BS], dtype=tl.float32) - for i_j in range(i_i + 1, NC): - p_z = tl.make_block_ptr(z + i_bh * T*S, (T, S), (S, 1), (i_t * BT + i_j * BC, i_s * BS), (BC, BS), (1, 0)) - p_ss = tl.make_block_ptr(ss + i_bh * T*S, (T, S), (S, 1), (i_t * BT + i_j * BC, i_s * BS), (BC, BS), (1, 0)) - # [BC, BS] - b_z = tl.load(p_z, boundary_check=(0, 1)) - b_ss = tl.load(p_ss, boundary_check=(0, 1)) - # [BC, BS] - b_doo += b_ss * exp(b_zn[None, :] - b_z) - b_doo = exp(b_s - b_zn[None, :]) * tl.dot(m_o.to(b_s.dtype), b_doo.to(b_s.dtype), allow_tf32=False) - - for j in range(0, BC): - p_z = tl.make_block_ptr(z + i_bh * T*S, (T*S,), (1,), ((i_t * BT + i_i * BC + j) * S + i_s * BS,), (BS,), (0,)) - p_ss = tl.make_block_ptr(ss + i_bh * T*S, (T*S,), (1,), ((i_t * BT + i_i * BC + j) * S + i_s * BS,), (BS,), (0,)) - # [BS,] - b_z = tl.load(p_z, boundary_check=(0,)) - b_ss = tl.load(p_ss, boundary_check=(0,)) - # [BC, BS] - m_i = o_i[:, None] <= j - b_doo += tl.where(m_i, exp(b_s - b_z[None, :]) * b_ss[None, :], 0.) - b_doo += tl.load(p_doo, boundary_check=(0, 1)) - tl.store(p_doo, b_doo.to(p_doo.dtype.element_ty), boundary_check=(0, 1)) - - -class ChunkABCFunction(torch.autograd.Function): - - @staticmethod - @input_guard - def forward(ctx, q, k, v, s, initial_state, output_final_state): - B, H, T, K, V, M = *q.shape, v.shape[-1], s.shape[-1] - BT, BC = 64, 16 - BK = min(64, triton.next_power_of_2(K)) - BV = min(64, triton.next_power_of_2(V)) - BM = min(64, triton.next_power_of_2(M)) - NT, NC = triton.cdiv(T, BT), triton.cdiv(BT, BC) - NV, NM = triton.cdiv(V, BV), triton.cdiv(M, BM) - num_warps = 4 if BK == 64 else 2 - num_stages = 1 - - def fwd_pre(s, B, H, T, S): - # keep cummulative normalizer in fp32 - z = torch.empty_like(s, dtype=torch.float) - grid = (B * H,) - logcumsumexp_fwd_kernel[grid]( - s, z, - T=T, S=S, - ) - return z - - def fwd_inner(q, k, v, z, B, H, T, K, V, BT, BK, BV, NT, normk=False, h0=None, ht=None): - NK, NV = triton.cdiv(K, BK), triton.cdiv(V, BV) - h = q.new_empty(B, H, NT * K, V) - grid = (NV, NK, B * H) - chunk_abc_fwd_kernel_h[grid]( - k, v, z, h, h0, ht, - T=T, K=K, V=V, BT=BT, BK=BK, BV=BV, NT=NT, - NORMK=normk, - USE_INITIAL_STATE=h0 is not None, - STORE_FINAL_STATE=ht is not None, - num_warps=num_warps, - num_stages=num_stages, - ) - return h - - final_state = None - if output_final_state: - final_state = (q.new_empty(B, H, K, M, dtype=torch.float), - q.new_empty(B, H, M, V, dtype=torch.float)) - - z = fwd_pre(s, B, H, T, M) - scale = K ** -0.5 - hk = fwd_inner( - q=q, k=k, v=s, z=z, - B=B, H=H, T=T, K=K, V=M, BT=BT, BK=BK, BV=BM, NT=NT, - normk=False, - h0=initial_state[0] if initial_state is not None else None, - ht=final_state[0] if final_state is not None else None, - ) - ok1 = torch.empty_like(s) - Ak = q.new_empty(B, H, T, BT) - grid = (NM, NT, B * H) - chunk_abc_fwd_kernel_K[grid]( - q, k, z, hk, ok1, Ak, - scale=scale, - T=T, K=K, V=M, BT=BT, BK=BK, BV=BM, NT=NT, - num_warps=num_warps, - num_stages=num_stages, - ) - ok0 = torch.empty_like(s) - grid = (NM, NT * NC, B * H) - chunk_abc_fwd_kernel_intra_K[grid]( - s, z, ok0, Ak, - T=T, V=M, BT=BT, BC=BC, BV=BM, NC=NC, - num_warps=2, - num_stages=num_stages, - ) - ok = ok0.add_(ok1) - - scale = 1. - # p is kept in fp32 for safe softmax backward - p = softmax_fwd(ok, dtype=torch.float) - qv = p.to(q.dtype) - - scale = 1. - hv = fwd_inner( - q=qv, k=s, v=v, z=z, - B=B, H=H, T=T, K=M, V=V, BT=BT, BK=BM, BV=BV, NT=NT, - normk=True, - h0=initial_state[1] if initial_state is not None else None, - ht=final_state[1] if final_state is not None else None, - ) - Av = q.new_zeros(NM, B, H, T, BT) - grid = (NM, NT * NC * NC, B * H) - chunk_abc_fwd_kernel_intra_V[grid]( - qv, s, z, Av, - scale=scale, - T=T, K=M, BT=BT, BC=BC, BK=BM, NC=NC, - num_warps=2, - num_stages=num_stages, - ) - Av = Av.sum(0) - ov = torch.empty_like(v) - grid = (NV, NT, B * H) - chunk_abc_fwd_kernel_V[grid]( - qv, v, z, hv, ov, Av, - scale=scale, - T=T, - K=M, - V=V, - BT=BT, - BK=BM, - BV=BV, - NT=NT, - num_warps=num_warps, - num_stages=num_stages, - ) - ctx.save_for_backward(q, k, v, s, z, ok, p, hk, hv, Av) - ctx.BT = BT - return ov, final_state - - @staticmethod - @input_guard - def backward(ctx, dov, dht=None): - q, k, v, s, z, ok, p, hk, hv, Av = ctx.saved_tensors - B, H, T, K, V, M = *q.shape, v.shape[-1], s.shape[-1] - BT, BC = ctx.BT, 16 - BK = min(64, triton.next_power_of_2(K)) - BV = min(64, triton.next_power_of_2(V)) - BM = min(64, triton.next_power_of_2(M)) - NT, NC = triton.cdiv(T, BT), triton.cdiv(BT, BC) - NK, NM = triton.cdiv(K, BK), triton.cdiv(M, BM) - num_warps = 4 if BK == 64 else 2 - num_stages = 1 - - def bwd_inner(q, z, do, B, H, T, K, V, BT, BK, BV, NT, scale, normk=False): - NK, NV = triton.cdiv(K, BK), triton.cdiv(V, BV) - dh = q.new_empty(B, H, NT * K, V) - grid = (NK, NV, B * H) - chunk_abc_bwd_kernel_dh[grid]( - q, z, do, dh, - scale=scale, - T=T, K=K, V=V, BT=BT, BK=BK, BV=BV, NT=NT, - NORMK=normk, - num_warps=num_warps, - num_stages=num_stages, - ) - return dh - - def bwd_post(s, z, ss, B, H, T, S, BT, BC, BS, NT, NC, NS): - doo = torch.empty_like(s) - grid = (NS, B * H) - chunk_abc_bwd_kernel_rcum_inter[grid]( - s, z, ss, doo, - T=T, S=S, BT=BT, BS=BS, NT=NT, - num_warps=num_warps, - num_stages=num_stages, - ) - grid = (NS, NT * NC, B * H) - chunk_abc_bwd_kernel_rcum_intra[grid]( - s, z, ss, doo, - T=T, S=S, BT=BT, BC=BC, BS=BS, NC=NC, - num_warps=num_warps, - num_stages=num_stages, - ) - return doo - - scale = 1. - qv = p.to(q.dtype) - dhv = bwd_inner( - qv, z, dov, - B=B, H=H, T=T, K=M, V=V, BT=BT, BK=BM, BV=BV, NT=NT, - scale=scale, - normk=True, - ) - dp1 = torch.empty_like(p) - dsv1 = torch.empty_like(s, dtype=torch.float) - dv = v.new_empty(NM, *v.shape) - dAv = q.new_zeros(B, H, T, BT) - grid = (NM, NT, B * H) - chunk_abc_bwd_kernel_V[grid]( - s, v, z, hv, Av, dov, dhv, dp1, dsv1, dv, dAv, - scale=scale, - T=T, K=M, V=V, BT=BT, BK=BM, BV=BV, NT=NT, - num_warps=num_warps, - num_stages=num_stages, - ) - dv = dv.sum(0) - dp0 = torch.empty_like(p) - dsv0 = s.new_zeros(s.shape, dtype=torch.float) - grid = (NM, NT * NC, B * H) - chunk_abc_bwd_kernel_intra_V[grid]( - qv, s, z, dAv, dp0, dsv0, - T=T, K=M, BT=BT, BC=BC, BK=BM, NC=NC, - num_warps=2, - num_stages=num_stages, - ) - dp = dp1.add_(dp0) - dsv = dsv1.add_(dsv0) - - # softmax gradient, equivalent to: - # dok = p * (dp - (p * dp).sum(-1, True)) - dok = softmax_bwd(p, dp, dtype=ok.dtype) - - scale = K ** -0.5 - dhk = bwd_inner( - q, z, dok, - B=B, H=H, T=T, K=K, V=M, BT=BT, BK=BK, BV=BM, NT=NT, - scale=scale, - normk=False, - ) - dAk = q.new_zeros(NM, B, H, T, BT) - grid = (NM, NT * NC * NC, B * H) - chunk_abc_bwd_kernel_intra_K[grid]( - s, z, dok, dAk, - scale=scale, - T=T, V=M, BT=BT, BC=BC, BV=BM, NC=NC, - num_warps=2, - num_stages=num_stages, - ) - dAk = dAk.sum(0) - - Ak = q.new_zeros(NK, B, H, T, BT) - dq = torch.empty_like(q) - dk = torch.empty_like(k) - dsk1 = s.new_empty(NK, *s.shape, dtype=torch.float) - grid = (NK, NT, B * H) - chunk_abc_bwd_kernel_K[grid]( - q, k, s, z, hk, Ak, dok, dhk, dq, dk, dsk1, dAk, - scale=scale, - T=T, K=K, V=M, BT=BT, BK=BK, BV=BM, NT=NT, - num_warps=num_warps, - num_stages=num_stages, - ) - Ak = Ak.sum(0) - dsk1 = dsk1.sum(0) - dsk0 = torch.empty_like(s, dtype=torch.float) - grid = (NM, NT * NC, B * H) - chunk_abc_bwd_kernel_intra_KV[grid]( - s, z, Ak, dok, dsk0, - T=T, V=M, BT=BT, BC=BC, BV=BM, NC=NC, - num_warps=2, - num_stages=num_stages, - ) - ds = dsv.add_(dsk1.add_(dsk0)) - ds -= bwd_post(s, z, ok * dok + p * dp, B, H, T, M, BT, BC, BM, NT, NC, NM) - ds = ds.to(s.dtype) - return dq, dk, dv, ds, None, None - - -@torch.compiler.disable -def chunk_abc( - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - s: torch.Tensor, - initial_state: tuple[torch.Tensor] | None = None, - output_final_state: bool = False, - head_first: bool = False, -) -> tuple[torch.Tensor, torch.Tensor]: - r""" - Args: - q (torch.Tensor): - queries of shape `[B, T, H, K]`. - k (torch.Tensor): - keys of shape `[B, T, H, K]`. - v (torch.Tensor): - values of shape `[B, T, H, V]`. - s (torch.Tensor): - slot representations of shape `[B, T, H, M]`. - initial_state (Optional[Tuple[torch.Tensor, torch.Tensor]]): - Initial states of shape `[B, H, K, M]` and `[B, H, M, V]`. Default: `None`. - output_final_state (Optional[bool]): - Whether to output the final state of shape `[B, H, K, M]` and `[B, H, M, V]`. Default: `False`. - head_first (Optional[bool]): - Whether the inputs are in the head-first format. Default: `False`. - This argument has been deprecated. - - Returns: - o (torch.Tensor): - Outputs of shape `[B, T, H, V]`. - final_state (torch.Tensor): - Final state of shape `[B, H, K, M]` and `[B, H, M, V]` if `output_final_state=True` else `None`. - """ - if not head_first: - q, k, v, s = map(lambda x: x.transpose(1, 2), (q, k, v, s)) - o, final_state = ChunkABCFunction.apply(q, k, v, s, initial_state, output_final_state) - if not head_first: - o = o.transpose(1, 2) - return o, final_state diff --git a/build/torch-xpu/ops/abc/naive.py b/build/torch-xpu/ops/abc/naive.py deleted file mode 100644 index d307e49192d9b51224bfb28815e9270987cdca0c..0000000000000000000000000000000000000000 --- a/build/torch-xpu/ops/abc/naive.py +++ /dev/null @@ -1,99 +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 -from einops import repeat - - -def naive_recurrent_abc( - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - s: torch.Tensor, - g: torch.Tensor | None = None, - scale: int | None = None, - initial_state: torch.Tensor | None = None, - output_final_state: bool | None = False, -) -> torch.Tensor: - dtype = q.dtype - - NG = q.shape[1]//k.shape[1] - # [batch_size, n_heads, seq_len, n_slots] - if g is None: - z = s.float().logcumsumexp(2) - g = torch.cat((z[:, :, :1], z[:, :, :-1]), 2) - z - s = torch.exp(s - z) - q, k, v, s, g = map(lambda x: x.float(), (q, k, v, s, g)) - k, v, s, g = map(lambda x: repeat(x, 'b h t d -> b (h g) t d', g=NG), (k, v, s, g)) - if initial_state is not None: - initial_state = tuple(map(lambda x: repeat(x, 'b h k v -> b (h g) k v', g=NG), initial_state)) - - B, H, T, K, V, M = *q.shape, v.shape[-1], s.shape[-1] - - hk = torch.zeros(B, H, K, M, dtype=torch.float, device=q.device) - ok = torch.zeros_like(s) - - if scale is None: - scale = q.shape[-1] ** -0.5 - - final_state = None - if initial_state is not None: - hk += initial_state[0] - - for i in range(T): - q_i = q[:, :, i] * scale - k_i = k[:, :, i] - v_i = s[:, :, i] - g_i = g[:, :, i].exp() - hk = hk * g_i[..., None, :] + k_i[..., None] * v_i[..., None, :] - ok[:, :, i] = (q_i[..., None] * hk).sum(-2) - - qv = ok.softmax(-1) - hv = torch.zeros(B, H, M, V, dtype=torch.float, device=q.device) - ov = torch.zeros_like(v) - if initial_state is not None: - hv += initial_state[1] - - for i in range(T): - q_i = qv[:, :, i] - k_i = s[:, :, i] - v_i = v[:, :, i] - g_i = g[:, :, i].exp() - hv = hv * g_i[..., :, None] + k_i[..., None] * v_i[..., None, :] - ov[:, :, i] = (q_i[..., None] * hv).sum(-2) - - if output_final_state: - final_state = (hk.view(B, -1, NG, K, M)[:, :, 0], hv.view(B, -1, NG, M, V)[:, :, 0]) - return ov.to(dtype), final_state - - -def naive_cumsum_abc( - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - s: torch.Tensor, -) -> torch.Tensor: - """ - A simple implementation of vanilla ABC that is more aligned with the descriptions in the paper. - This is just for demonstration purposes, with no numerical stabilities guaranteed. - """ - - dtype = q.dtype - q, k, v, s = map(lambda x: x.float(), (q, k, v, s)) - - scale = q.shape[-1] ** -0.5 - # [batch_size, n_heads, seq_len, n_slots] - s = (s - s.max(2, True)[0]).exp() - z = s.cumsum(2) - # [batch_size, n_heads, seq_len, n_slots, d_head] - K = (s.unsqueeze(-1) * k.unsqueeze(-2)).cumsum(2) / z.unsqueeze(-1) - V = (s.unsqueeze(-1) * v.unsqueeze(-2)).cumsum(2) / z.unsqueeze(-1) - # [batch_size, n_heads, seq_len, n_slots] - p = torch.einsum('...d,...md->...m', q * scale, K).softmax(-1) - # [batch_size, n_heads, seq_len, d_head] - o = torch.einsum('...m,...md->...d', p, V) - return o.to(dtype), None diff --git a/build/torch-xpu/ops/attn/__init__.py b/build/torch-xpu/ops/attn/__init__.py deleted file mode 100644 index 4f86e570433d825740b342b4606e9bd18a422984..0000000000000000000000000000000000000000 --- a/build/torch-xpu/ops/attn/__init__.py +++ /dev/null @@ -1,14 +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 .naive import naive_parallel_attn -from .parallel import parallel_attn - -__all__ = [ - 'naive_parallel_attn', - 'parallel_attn', -] diff --git a/build/torch-xpu/ops/attn/backends/__init__.py b/build/torch-xpu/ops/attn/backends/__init__.py deleted file mode 100644 index 76179f136c9df857ad5bbfda568412dc06812e33..0000000000000000000000000000000000000000 --- a/build/torch-xpu/ops/attn/backends/__init__.py +++ /dev/null @@ -1,11 +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 ....ops.backends import BackendRegistry -from ....ops.common.backends.tilelang import TileLangBackend - -BackendRegistry("attn").register(TileLangBackend()) diff --git a/build/torch-xpu/ops/attn/decoding.py b/build/torch-xpu/ops/attn/decoding.py deleted file mode 100644 index c9a8aea4c294e5f9e5158335abeaa77b0da005cd..0000000000000000000000000000000000000000 --- a/build/torch-xpu/ops/attn/decoding.py +++ /dev/null @@ -1,209 +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 -import triton.language as tl - -from ...ops.utils.cumsum import chunk_global_cumsum -from ...ops.utils.op import exp -from ...utils import autotune_cache_kwargs, check_shared_mem - - -@triton.heuristics({ - 'USE_G': lambda args: args['g_cumsum'] is not None, - 'USE_SINK_BIAS': lambda args: args['sink_bias'] is not None, -}) -@triton.autotune( - configs=[ - triton.Config({}, num_warps=num_warps, num_stages=num_stages) - for num_warps in [1, 2, 4] + ([] if check_shared_mem('hopper') else [8]) - for num_stages in [2, 3, 4, 5] - ], - key=['H', 'G', 'K', 'V', 'BK', 'BV', 'USE_G', 'USE_SINK_BIAS'], - **autotune_cache_kwargs, -) -@triton.jit -def naive_attn_decoding_kernel( - q, - k, - v, - o, - g_cumsum, - sink_bias, - scale, - cu_seqlens, - T, - B: tl.constexpr, - H: tl.constexpr, - HQ: tl.constexpr, - G: tl.constexpr, - K: tl.constexpr, - V: tl.constexpr, - BS: tl.constexpr, - BK: tl.constexpr, - BV: tl.constexpr, - USE_G: tl.constexpr, - USE_SINK_BIAS: tl.constexpr, -): - i_v, i_bh = tl.program_id(0), tl.program_id(1) - i_b, i_hq = i_bh // HQ, i_bh % HQ - i_h = i_hq // G - - bos, eos = tl.load(cu_seqlens + i_b).to(tl.int32), tl.load(cu_seqlens + i_b + 1).to(tl.int32) - T = eos - bos - - p_q = tl.make_block_ptr(q + i_bh * K, (K,), (1, ), (0, ), (BK,), (0,)) - p_o = tl.make_block_ptr(o + i_bh * V, (V,), (1, ), (0, ), (BV,), (0,)) - - b_q = tl.load(p_q, boundary_check=(0,)) - b_q = (b_q * scale).to(b_q.dtype) - - b_o = tl.zeros([BV], dtype=tl.float32) - - b_m = tl.full([1], float('-inf'), dtype=tl.float32) - b_acc = tl.zeros([1], dtype=tl.float32) - - if USE_G: - p_g = tl.make_block_ptr(g_cumsum + bos * HQ + i_hq, (T,), (HQ,), (T-1,), (1,), (0,)) - b_gq = tl.load(p_g, boundary_check=(0,)).to(tl.float32) - else: - b_gq = None - - if USE_SINK_BIAS: - b_sink_bias = tl.load(sink_bias + i_hq).to(tl.float32) - else: - b_sink_bias = None - - for i_s in range(0, T, BS): - p_k = tl.make_block_ptr(k + (bos * H + i_h) * K, (T, K), (H*K, 1), (i_s, 0), (BS, BK), (1, 0)) - p_v = tl.make_block_ptr(v + (bos * H + i_h) * V, (T, V), (H*V, 1), (i_s, i_v * BV), (BS, BV), (1, 0)) - # [BK, BS] - b_k = tl.load(p_k, boundary_check=(0, 1)) - # [BS, BV] - b_v = tl.load(p_v, boundary_check=(0, 1)) - # [BT, BS] - b_s = tl.sum(b_q[None, :] * b_k, 1) - - mask = i_s + tl.arange(0, BS) < T - b_s = tl.where(mask, b_s, float('-inf')) - - if USE_G: - p_gk = tl.make_block_ptr(g_cumsum + bos * HQ + i_hq, (T,), (HQ,), (i_s,), (BS,), (0,)) - b_gk = tl.load(p_gk, boundary_check=(0,)).to(tl.float32) - b_s += b_gq - b_gk - # [BT, BS] - b_m, b_mp = tl.maximum(b_m, tl.max(b_s)), b_m - b_r = exp(b_mp - b_m) - # [BT, BS] - b_p = exp(b_s - b_m) - - # [BT] - b_acc = b_acc * b_r + tl.sum(b_p, 0) - # [BT, BV] - b_o = b_o * b_r + tl.sum(b_p[:, None] * b_v, 0) - b_mp = b_m - - if USE_SINK_BIAS: - # keep the sink-bias merge finite when masking leaves a row with no valid key. - b_m = tl.where(b_m == float('-inf'), 0., b_m) - b_acc += exp(b_sink_bias - b_m) - b_o = b_o / b_acc - tl.store(p_o, b_o.to(p_o.dtype.element_ty), boundary_check=(0, )) - - -def attn_decoding_one_step( - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - g: torch.Tensor | None = None, - scale: float | None = None, - cu_seqlens: torch.LongTensor = None, - do_gate_scale: bool = False, - *, - sink_bias: torch.Tensor | None = None, -): - r""" - Args: - q (torch.Tensor): - query of shape `[1, B, HQ, K]`. - k (torch.Tensor): - keys of shape `[1, T, H, K]`. - GQA will be applied if HQ is divisible by H. T is the cumulative length for all batch. - v (torch.Tensor): - values of shape `[1, T, H, V]`. - g (Optional[torch.Tensor]): - log decay factors of shape `[1, T, HQ]`. Default: `None`. - scale (Optional[float]): - Scale factor for attention scores. - If not provided, it will default to `1 / sqrt(K)`. Default: `None`. - cu_seqlens (torch.LongTensor): - Cumulative sequence lengths of shape `[N+1]` used for variable-length training, - consistent with the FlashAttention API. - do_gate_scale (bool): - Whether to apply gate scale. Default: `False`. If `True`, the attention scale will also be applied - to the gating bias term in Forgetting Transformer or PaTH-FoX. - sink_bias (Optional[torch.Tensor]): - Per-query-head attention-sink bias logits of shape `[HQ]` — one - learnable scalar per query head, as introduced by GPT-OSS. - Augments the softmax denominator without contributing to the output. - - Returns: - o (torch.Tensor): - Outputs of shape `[1, B, HQ, V]`. - """ - assert cu_seqlens is not None, "The cu_seqlens must be provided for varlen decoding" - B, T, H, K, V = *k.shape, v.shape[-1] - N = len(cu_seqlens) - 1 - HQ = q.shape[2] - G = HQ // H - if scale is None: - scale = K ** -0.5 - if sink_bias is not None: - assert sink_bias.shape == (HQ,), "sink_bias must have shape [HQ]" - - BK = max(triton.next_power_of_2(K), 16) - if check_shared_mem('hopper', q.device.index): - BS = min(64, max(16, triton.next_power_of_2(T))) - BV = min(256, max(16, triton.next_power_of_2(V))) - elif check_shared_mem('ampere', q.device.index): - BS = min(32, max(16, triton.next_power_of_2(T))) - BV = min(128, max(16, triton.next_power_of_2(V))) - else: - BS = min(32, max(16, triton.next_power_of_2(T))) - BV = min(64, max(16, triton.next_power_of_2(V))) - g_cumsum = chunk_global_cumsum( - g, - cu_seqlens=cu_seqlens, - scale=scale if do_gate_scale else None, - output_dtype=torch.float32, - ) if g is not None else None - NV = triton.cdiv(V, BV) - o = torch.empty(*q.shape[:-1], V, dtype=v.dtype, device=q.device) - - grid = (NV, N * HQ) - naive_attn_decoding_kernel[grid]( - q=q, - k=k, - v=v, - o=o, - g_cumsum=g_cumsum, - sink_bias=sink_bias, - scale=scale, - cu_seqlens=cu_seqlens, - B=B, - T=T, - H=H, - HQ=HQ, - G=G, - K=K, - V=V, - BS=BS, - BK=BK, - BV=BV, - ) - return o diff --git a/build/torch-xpu/ops/attn/naive.py b/build/torch-xpu/ops/attn/naive.py deleted file mode 100644 index 5bdebf1e60779bb6cfab34afb40edd91320de1d4..0000000000000000000000000000000000000000 --- a/build/torch-xpu/ops/attn/naive.py +++ /dev/null @@ -1,176 +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 torch.nn.functional as F - - -def naive_parallel_attn( - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - scale: float | None = None, - window_size: int | None = None, - causal: bool = True, - *, - g: torch.Tensor | None = None, - sink_bias: torch.Tensor | None = None, -): - """ - Reference PyTorch implementation of parallel attention that returns both output and max_logits. - - Args: - q: [B, T, HQ, D] - k: [B, T, H, D] - v: [B, T, H, D] - scale: float, optional. If None, defaults to 1 / sqrt(D) - window_size: int, optional. If provided, each query at position i only attends to - keys in [i - window_size + 1, i]. If None, full causal attention is used. - causal: bool, default True - g: [B, T, HQ], optional per-query-head gating logits. Any scaling - (e.g. matching `attn_decoding_one_step(do_gate_scale=True)`) should - be applied by the caller before passing `g` in. - sink_bias: [HQ], optional per-query-head attention-sink bias logits - (GPT-OSS style). One scalar per query head added to the softmax - denominator without a corresponding key/value — absorbs probability - mass but does not contribute to the output. - - Returns: - output: [B, T, HQ, D] - max_logits: [B, T, HQ] - """ - B, T, HQ, D = q.shape - H = k.shape[2] - G = HQ // H - - if scale is None: - scale = D ** -0.5 - - # reshape q to separate groups: [B, T, HQ, D] -> [B, T, H, G, D] - q = q.reshape(B, T, H, G, D) - - # compute attention scores via einsum: [B, H, G, T, T] - # k is [B, T, H, D] — no group dim, so each group shares the same k - scores = torch.einsum('bqhgd,bkhd->bhgqk', q, k) * scale - - # apply causal mask - if causal: - row_idx = torch.arange(T, device=q.device)[None, :, None] - col_idx = torch.arange(T, device=q.device)[None, None, :] - mask = col_idx > row_idx - if window_size is not None: - mask = mask | (row_idx - col_idx >= window_size) - scores = scores.masked_fill(mask[:, None, None], float('-inf')) - - if g is not None: - assert g.shape == (B, T, HQ), "g must have shape [B, T, HQ]" - g_cumsum = g.float().cumsum(1).reshape(B, T, H, G).permute(0, 2, 3, 1) - scores = scores + (g_cumsum[..., :, None] - g_cumsum[..., None, :]) - - # max_logits: [B, H, G, T] -> [B, T, HQ] - max_logits = scores.max(dim=-1).values - if sink_bias is not None: - assert sink_bias.shape == (HQ,), "sink_bias must have shape [HQ]" - sink_bias_logits = sink_bias.reshape(H, G)[None, :, :, None] - max_logits = torch.maximum(max_logits, sink_bias_logits) - - if sink_bias is None: - # compute output via einsum: [B, H, G, T, T] x [B, T, H, D] -> [B, T, H, G, D] - attn_weights = F.softmax(scores, dim=-1) - output = torch.einsum('bhgqk,bkhd->bqhgd', attn_weights, v).reshape(B, T, HQ, D) - else: - probs_unnorm = torch.exp(scores - max_logits[..., None]) - sink_bias_unnorm = torch.exp(sink_bias_logits - max_logits) - denom = probs_unnorm.sum(dim=-1) + sink_bias_unnorm - output = torch.einsum('bhgqk,bkhd->bqhgd', probs_unnorm, v) - output = (output / denom.permute(0, 3, 1, 2)[..., None]).reshape(B, T, HQ, D) - - return output, max_logits.permute(0, 3, 1, 2).reshape(B, T, HQ) - - -def naive_attn_decoding( - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - g: torch.Tensor | None = None, - scale: float | None = None, - cu_seqlens: torch.LongTensor | None = None, - do_gate_scale: bool = False, - *, - sink_bias: torch.Tensor | None = None, -): - """ - Reference PyTorch implementation of packed-varlen decoding attention, - mirroring `attn_decoding_one_step`. A single query per sequence attends - to all of its KV (no causal mask needed — query is at the last position). - - Args: - q: [1, B, HQ, D] — one query token per sequence - k: [1, T_total, H, D] - v: [1, T_total, H, D] - g: [1, T_total, HQ], optional log decay factors - scale: float, defaults to 1/sqrt(D) - cu_seqlens: [B+1] - do_gate_scale: bool, if True scales `g` by `scale` before use - (matches PaTH / Forgetting Transformer convention). - sink_bias: [HQ], optional GPT-OSS-style sink bias logits - - Returns: - o: [1, B, HQ, V] - """ - assert cu_seqlens is not None, "cu_seqlens must be provided for varlen decoding" - HQ, D = q.shape[-2], q.shape[-1] - V = v.shape[-1] - H = k.shape[2] - G = HQ // H - if scale is None: - scale = D ** -0.5 - if sink_bias is not None: - assert sink_bias.shape == (HQ,), "sink_bias must have shape [HQ]" - - outputs = [] - for i in range(len(cu_seqlens) - 1): - bos, eos = int(cu_seqlens[i]), int(cu_seqlens[i + 1]) - T_i = eos - bos - qi = q[:, i:i + 1] # [1, 1, HQ, D] - - if T_i == 0: - # no KV for this sequence → output zeros (sink_bias has no value to contribute) - outputs.append(torch.zeros((1, 1, HQ, V), dtype=q.dtype, device=q.device)) - continue - - ki = k[:, bos:eos] # [1, T_i, H, D] - vi = v[:, bos:eos] # [1, T_i, H, D] - - # scores: [1, H, G, 1, T_i] - qi_g = qi.reshape(1, 1, H, G, D) - scores = torch.einsum('bqhgd,bkhd->bhgqk', qi_g, ki) * scale - - if g is not None: - gi = g[:, bos:eos].float() - if do_gate_scale: - gi = gi * scale - # query is at position T_i - 1 (end of sequence) - gi_cumsum = gi.cumsum(1).reshape(1, T_i, H, G).permute(0, 2, 3, 1) # [1, H, G, T_i] - g_q = gi_cumsum[..., -1:].unsqueeze(-1) # [1, H, G, 1, 1] - scores = scores + (g_q - gi_cumsum[..., None, :]) - - if sink_bias is None: - attn_weights = F.softmax(scores, dim=-1) - oi = torch.einsum('bhgqk,bkhd->bqhgd', attn_weights, vi).reshape(1, 1, HQ, V) - else: - sink_bias_logits = sink_bias.reshape(H, G)[None, :, :, None] # [1, H, G, 1] - max_logits = torch.maximum(scores.max(dim=-1).values, sink_bias_logits) # [1, H, G, 1] - probs_unnorm = torch.exp(scores - max_logits[..., None]) - sink_bias_unnorm = torch.exp(sink_bias_logits - max_logits) - denom = probs_unnorm.sum(dim=-1) + sink_bias_unnorm # [1, H, G, 1] - oi_pack = torch.einsum('bhgqk,bkhd->bqhgd', probs_unnorm, vi) - oi = (oi_pack / denom.permute(0, 3, 1, 2)[..., None]).reshape(1, 1, HQ, V) - - outputs.append(oi) - - return torch.cat(outputs, dim=1) # [1, B, HQ, V] diff --git a/build/torch-xpu/ops/attn/parallel.py b/build/torch-xpu/ops/attn/parallel.py deleted file mode 100644 index 3adef692899e644d29db235b9c069b50ad79abb4..0000000000000000000000000000000000000000 --- a/build/torch-xpu/ops/attn/parallel.py +++ /dev/null @@ -1,839 +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 -import triton.language as tl -from einops import reduce - -from ...ops.backends import dispatch -from ...ops.utils import prepare_chunk_indices -from ...ops.utils.constant import RCP_LN2 -from ...ops.utils.cumsum import chunk_global_cumsum -from ...ops.utils.op import exp2, log2 -from ...utils import autocast_custom_bwd, autocast_custom_fwd, check_shared_mem, contiguous - - -@triton.heuristics({ - 'USE_G': lambda args: args['g_cumsum'] is not None, - 'USE_SINK_BIAS': lambda args: args['sink_bias'] is not None, - 'USE_WINDOW': lambda args: args['W'] is not None, - 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None, -}) -@triton.jit -def parallel_attn_fwd_kernel( - q, - k, - v, - o, - g_cumsum, - sink_bias, - lse, - scale, - cu_seqlens, - chunk_indices, - T, - W: tl.constexpr, - B: tl.constexpr, - H: tl.constexpr, - HQ: tl.constexpr, - G: tl.constexpr, - K: tl.constexpr, - V: tl.constexpr, - BT: tl.constexpr, - BS: tl.constexpr, - BK: tl.constexpr, - BV: tl.constexpr, - USE_G: tl.constexpr, - USE_SINK_BIAS: tl.constexpr, - USE_WINDOW: tl.constexpr, - IS_VARLEN: tl.constexpr, -): - i_v, i_t, i_bh = tl.program_id(0), tl.program_id(1), tl.program_id(2) - i_b, i_hq = i_bh // HQ, i_bh % HQ - i_h = i_hq // G - - 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).to(tl.int32) - else: - i_n = i_b - bos, eos = (i_n * T).to(tl.int64), (i_n * T + T).to(tl.int64) - RCP_LN2: tl.constexpr = 1.4426950216 - - p_q = tl.make_block_ptr(q + (bos * HQ + i_hq) * K, (T, K), (HQ*K, 1), (i_t * BT, 0), (BT, BK), (1, 0)) - p_o = tl.make_block_ptr(o + (bos * HQ + i_hq) * V, (T, V), (HQ*V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) - p_lse = tl.make_block_ptr(lse + bos * HQ + i_hq, (T,), (HQ,), (i_t * BT,), (BT,), (0,)) - - # the Q block is kept in the shared memory throughout the whole kernel - # [BT, BK] - b_q = tl.load(p_q, boundary_check=(0, 1)) - # [BT, BV] - b_o = tl.zeros([BT, BV], dtype=tl.float32) - - b_m = tl.full([BT], float('-inf'), dtype=tl.float32) - b_acc = tl.zeros([BT], dtype=tl.float32) - - if USE_G: - p_g = tl.make_block_ptr(g_cumsum + bos * HQ + i_hq, (T,), (HQ,), (i_t * BT,), (BT,), (0,)) - b_gq = tl.load(p_g, boundary_check=(0,)).to(tl.float32) - else: - b_gq = None - - if USE_SINK_BIAS: - b_sink_bias = tl.load(sink_bias + i_hq).to(tl.float32) - else: - b_sink_bias = None - - # [BT] - o_q = i_t * BT + tl.arange(0, BT) - - # for sliding window, skip key blocks that are entirely outside the window. - # the earliest key position any query in this block needs: max(0, i_t*BT - W + 1) - i_start = tl.maximum((i_t * BT - W + 1) // BS * BS, 0) if USE_WINDOW else 0 - - for i_s in range(i_start, i_t * BT, BS): - p_k = tl.make_block_ptr(k + (bos * H + i_h) * K, (K, T), (1, H*K), (0, i_s), (BK, BS), (0, 1)) - p_v = tl.make_block_ptr(v + (bos * H + i_h) * V, (T, V), (H*V, 1), (i_s, i_v * BV), (BS, BV), (1, 0)) - # [BK, BS] - b_k = tl.load(p_k, boundary_check=(0, 1)) - # [BS, BV] - b_v = tl.load(p_v, boundary_check=(0, 1)) - # [BT, BS] - b_s = tl.dot(b_q, b_k) * scale * RCP_LN2 - - o_k = i_s + tl.arange(0, BS) - m_k = o_k < T - if USE_G: - b_gk = tl.load(g_cumsum + (bos + o_k) * HQ + i_hq, mask=m_k, other=0).to(tl.float32) - b_s += b_gq[:, None] - b_gk[None, :] - - if USE_WINDOW: - b_s = tl.where((o_q[:, None] - o_k[None, :] < W) & m_k[None, :], b_s, float('-inf')) - - # [BT, BS] - b_m, b_mp = tl.maximum(b_m, tl.max(b_s, 1)), b_m - # keep the online softmax pivot finite for rows that still have no valid key. - # this matches sglang's masked-row stabilization and avoids -inf - (-inf) = NaN. - b_mw = tl.where(b_m == float('-inf'), 0., b_m) - b_r = exp2(b_mp - b_mw) - b_p = exp2(b_s - b_mw[:, None]) - # [BT] - b_acc = b_acc * b_r + tl.sum(b_p, 1) - # [BT, BV] - b_o = b_o * b_r[:, None] + tl.dot(b_p.to(b_q.dtype), b_v) - - b_mp = b_m - - for i_s in range(i_t * BT, min((i_t + 1) * BT, T), BS): - p_k = tl.make_block_ptr(k + (bos * H + i_h) * K, (K, T), (1, H*K), (0, i_s), (BK, BS), (0, 1)) - p_v = tl.make_block_ptr(v + (bos * H + i_h) * V, (T, V), (H*V, 1), (i_s, i_v * BV), (BS, BV), (1, 0)) - - # [BS] - o_k = i_s + tl.arange(0, BS) - m_k = o_k < T - # [BK, BS] - b_k = tl.load(p_k, boundary_check=(0, 1)) - # [BS, BV] - b_v = tl.load(p_v, boundary_check=(0, 1)) - # [BT, BS] - b_s = tl.dot(b_q, b_k) * scale * RCP_LN2 - - if USE_G: - b_gk = tl.load(g_cumsum + (bos + o_k) * HQ + i_hq, mask=m_k, other=0).to(tl.float32) - b_s += b_gq[:, None] - b_gk[None, :] - - m_s = (o_q[:, None] >= o_k[None, :]) & m_k[None, :] - if USE_WINDOW: - m_s = m_s & (o_q[:, None] - o_k[None, :] < W) - b_s = tl.where(m_s, b_s, float('-inf')) - - # [BT] - b_m, b_mp = tl.maximum(b_m, tl.max(b_s, 1)), b_m - b_mw = tl.where(b_m == float('-inf'), 0., b_m) - b_r = exp2(b_mp - b_mw) - b_p = exp2(b_s - b_mw[:, None]) - # [BT] - b_acc = b_acc * b_r + tl.sum(b_p, 1) - # [BT, BV] - b_o = b_o * b_r[:, None] + tl.dot(b_p.to(b_q.dtype), b_v) - b_mp = b_m - - if USE_SINK_BIAS: - # when a row has no valid key at all, b_m is still -inf here. - # use a finite pivot before merging the sink-bias mass so lse becomes - # the sink-bias logit instead of hitting the -inf + inf = NaN path. - b_m = tl.where(b_m == float('-inf'), 0., b_m) - # denominator-only sink-bias update (matches GPT-OSS / sglang): - # the bias logit augments the softmax normalizer without contributing - # to the value matmul. - b_acc += exp2(b_sink_bias - b_m) - - b_o = b_o / b_acc[:, None] - b_m += log2(b_acc) - tl.store(p_o, b_o.to(p_o.dtype.element_ty), boundary_check=(0, 1)) - tl.store(p_lse, b_m.to(p_lse.dtype.element_ty), boundary_check=(0,)) - - -@triton.jit -def parallel_attn_bwd_kernel_preprocess( - o, - do, - delta, - B: tl.constexpr, - V: tl.constexpr, -): - i_n = tl.program_id(0).to(tl.int64) - o_d = tl.arange(0, B) - m_d = o_d < V - - b_o = tl.load(o + i_n * V + o_d, mask=m_d, other=0) - b_do = tl.load(do + i_n * V + o_d, mask=m_d, other=0).to(tl.float32) - b_delta = tl.sum(b_o * b_do) - - tl.store(delta + i_n, b_delta.to(delta.dtype.element_ty)) - - -@triton.heuristics({ - 'USE_G': lambda args: args['g_cumsum'] is not None, - 'USE_WINDOW': lambda args: args['W'] is not None, - 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None, -}) -@triton.jit(do_not_specialize=['T']) -def parallel_attn_bwd_kernel_dq( - q, - k, - v, - lse, - delta, - do, - dq, - dg_cumsum, - g_cumsum, - scale, - cu_seqlens, - chunk_indices, - T, - W: tl.constexpr, - B: tl.constexpr, - H: tl.constexpr, - HQ: tl.constexpr, - G: tl.constexpr, - K: tl.constexpr, - V: tl.constexpr, - BT: tl.constexpr, - BS: tl.constexpr, - BK: tl.constexpr, - BV: tl.constexpr, - USE_G: tl.constexpr, - USE_WINDOW: tl.constexpr, - IS_VARLEN: tl.constexpr, -): - i_v, i_t, i_bh = tl.program_id(0), tl.program_id(1), tl.program_id(2) - i_b, i_hq = i_bh // HQ, i_bh % HQ - i_h = i_hq // G - - 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).to(tl.int32) - else: - i_n = i_b - bos, eos = (i_n * T).to(tl.int64), (i_n * T + T).to(tl.int64) - # NOTE: we must multiply RCP_LN2 after tl.dot for high precision - RCP_LN2: tl.constexpr = 1.4426950216 - - p_q = tl.make_block_ptr(q + (bos * HQ + i_hq) * K, (T, K), (HQ*K, 1), (i_t * BT, 0), (BT, BK), (1, 0)) - p_dq = tl.make_block_ptr(dq + (bos * HQ + i_hq) * K, (T, K), (HQ*K, 1), (i_t * BT, 0), (BT, BK), (1, 0)) - p_do = tl.make_block_ptr(do + (bos * HQ + i_hq) * V, (T, V), (HQ*V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) - p_lse = tl.make_block_ptr(lse + bos * HQ + i_hq, (T,), (HQ,), (i_t * BT,), (BT,), (0,)) - p_delta = tl.make_block_ptr(delta + bos * HQ + i_hq, (T,), (HQ,), (i_t * BT,), (BT,), (0,)) - - # [BT, BK] - b_q = tl.load(p_q, boundary_check=(0, 1)) - # [BT, BV] - b_do = tl.load(p_do, boundary_check=(0, 1)) - # [BT] - b_lse = tl.load(p_lse, boundary_check=(0,)) - b_delta = tl.load(p_delta, boundary_check=(0,)) - - # [BT, BK] - b_dq = tl.zeros([BT, BK], dtype=tl.float32) - if USE_G: - b_dg = tl.zeros([BT], dtype=tl.float32) - p_gq = tl.make_block_ptr(g_cumsum + bos * HQ + i_hq, (T,), (HQ,), (i_t * BT,), (BT,), (0,)) - b_gq = tl.load(p_gq, boundary_check=(0,)).to(tl.float32) - else: - b_gq = None - b_dg = None - - o_q = i_t * BT + tl.arange(0, BT) - - i_start = tl.maximum((i_t * BT - W + 1) // BS * BS, 0) if USE_WINDOW else 0 - - for i_s in range(i_start, i_t * BT, BS): - p_k = tl.make_block_ptr(k + (bos * H + i_h) * K, (K, T), (1, H*K), (0, i_s), (BK, BS), (0, 1)) - p_v = tl.make_block_ptr(v + (bos * H + i_h) * V, (V, T), (1, H*V), (i_v * BV, i_s), (BV, BS), (0, 1)) - - o_k = i_s + tl.arange(0, BS) - m_k = o_k < T - # [BK, BS] - b_k = tl.load(p_k, boundary_check=(0, 1)) - # [BV, BS] - b_v = tl.load(p_v, boundary_check=(0, 1)) - # [BT, BS] - b_s = tl.dot(b_q, b_k) * scale * RCP_LN2 - if USE_G: - b_gk = tl.load(g_cumsum + (bos + o_k) * HQ + i_hq, mask=m_k, other=0).to(tl.float32) - b_s += b_gq[:, None] - b_gk[None, :] - - if USE_WINDOW: - b_s = tl.where((o_q[:, None] - o_k[None, :] < W) & m_k[None, :], b_s, float('-inf')) - b_p = exp2(b_s - b_lse[:, None]) - # [BT, BV] @ [BV, BS] -> [BT, BS] - b_dp = tl.dot(b_do, b_v) - b_ds = b_p * (b_dp.to(tl.float32) - b_delta[:, None]) - # [BT, BS] @ [BS, BK] -> [BT, BK] - b_dq += tl.dot(b_ds.to(b_k.dtype), tl.trans(b_k)) - if USE_G: - b_dg += tl.sum(b_ds, 1) - - for i_s in range(i_t * BT, min((i_t + 1) * BT, T), BS): - p_k = tl.make_block_ptr(k + (bos * H + i_h) * K, (K, T), (1, H*K), (0, i_s), (BK, BS), (0, 1)) - p_v = tl.make_block_ptr(v + (bos * H + i_h) * V, (V, T), (1, H*V), (i_v * BV, i_s), (BV, BS), (0, 1)) - - # [BS] - o_k = i_s + tl.arange(0, BS) - m_k = o_k < T - # [BK, BS] - b_k = tl.load(p_k, boundary_check=(0, 1)) - # [BV, BS] - b_v = tl.load(p_v, boundary_check=(0, 1)) - # [BT, BS] - b_s = tl.dot(b_q, b_k) * scale * RCP_LN2 - - if USE_G: - p_gk = tl.make_block_ptr(g_cumsum + bos * HQ + i_hq, (T,), (HQ,), (i_s,), (BS,), (0,)) - b_gk = tl.load(p_gk, boundary_check=(0,)).to(tl.float32) - b_s += b_gq[:, None] - b_gk[None, :] - if USE_WINDOW: - b_p = tl.where( - (o_q[:, None] >= o_k[None, :]) & (o_q[:, None] - o_k[None, :] < W) & m_k[None, :], - exp2(b_s - b_lse[:, None]), 0 - ) - else: - b_p = tl.where((o_q[:, None] >= o_k[None, :]) & m_k[None, :], exp2(b_s - b_lse[:, None]), 0) - - # [BT, BV] @ [BV, BS] -> [BT, BS] - b_dp = tl.dot(b_do, b_v) - b_ds = b_p * (b_dp.to(tl.float32) - b_delta[:, None]) - # [BT, BS] @ [BS, BK] -> [BT, BK] - b_dq += tl.dot(b_ds.to(b_k.dtype), tl.trans(b_k)) - if USE_G: - b_dg += tl.sum(b_ds, 1) - - b_dq *= scale - tl.store(p_dq, b_dq.to(p_dq.dtype.element_ty), boundary_check=(0, 1)) - if USE_G: - p_dg = tl.make_block_ptr(dg_cumsum + bos * HQ + i_hq, (T,), (HQ,), (i_t * BT,), (BT,), (0,)) - tl.store(p_dg, b_dg.to(p_dg.dtype.element_ty), boundary_check=(0,)) - - -@triton.heuristics({ - 'USE_G': lambda args: args['g_cumsum'] is not None, - 'USE_WINDOW': lambda args: args['W'] is not None, - 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None, -}) -@triton.jit(do_not_specialize=['T']) -def parallel_attn_bwd_kernel_dkv( - q, - k, - v, - g_cumsum, - lse, - delta, - do, - dk, - dv, - dg_cumsum, - cu_seqlens, - chunk_indices, - scale, - T, - W: tl.constexpr, - B: tl.constexpr, - H: tl.constexpr, - HQ: tl.constexpr, - G: tl.constexpr, - K: tl.constexpr, - V: tl.constexpr, - BT: tl.constexpr, - BS: tl.constexpr, - BK: tl.constexpr, - BV: tl.constexpr, - USE_G: tl.constexpr, - USE_WINDOW: tl.constexpr, - IS_VARLEN: tl.constexpr, -): - i_v, i_t, i_bh = tl.program_id(0), tl.program_id(1), tl.program_id(2) - i_b, i_hq = i_bh // HQ, i_bh % HQ - i_h = i_hq // G - - 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).to(tl.int32) - else: - i_n = i_b - bos, eos = (i_n * T).to(tl.int64), (i_n * T + T).to(tl.int64) - RCP_LN2: tl.constexpr = 1.4426950216 - - p_k = tl.make_block_ptr(k + (bos * H + i_h) * K, (T, K), (H*K, 1), (i_t * BT, 0), (BT, BK), (1, 0)) - p_v = tl.make_block_ptr(v + (bos * H + i_h) * V, (T, V), (H*V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) - p_dk = tl.make_block_ptr(dk + (bos * HQ + i_hq) * K, (T, K), (HQ*K, 1), (i_t * BT, 0), (BT, BK), (1, 0)) - p_dv = tl.make_block_ptr(dv + (bos * HQ + i_hq) * V, (T, V), (HQ*V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) - - # [BT, BK] - b_k = tl.load(p_k, boundary_check=(0, 1)) - b_dk = tl.zeros([BT, BK], dtype=tl.float32) - # [BT, BV] - b_v = tl.load(p_v, boundary_check=(0, 1)) - b_dv = tl.zeros([BT, BV], dtype=tl.float32) - - o_k = i_t * BT + tl.arange(0, BT) - - if USE_G: - p_gk = tl.make_block_ptr(g_cumsum + bos * HQ + i_hq, (T,), (HQ,), (i_t * BT,), (BT,), (0,)) - b_gk = tl.load(p_gk, boundary_check=(0,)).to(tl.float32) - b_dg = tl.zeros([BT], dtype=tl.float32) - else: - b_gk = None - b_dg = None - - for i_s in range(i_t * BT, min((i_t + 1) * BT, T), BS): - p_q = tl.make_block_ptr(q + (bos * HQ + i_hq) * K, (T, K), (HQ*K, 1), (i_s, 0), (BS, BK), (1, 0)) - p_do = tl.make_block_ptr(do + (bos * HQ + i_hq) * V, (T, V), (HQ*V, 1), (i_s, i_v * BV), (BS, BV), (1, 0)) - p_lse = tl.make_block_ptr(lse + bos * HQ + i_hq, (T,), (HQ,), (i_s,), (BS,), (0,)) - p_delta = tl.make_block_ptr(delta + bos * HQ + i_hq, (T,), (HQ,), (i_s,), (BS,), (0,)) - - # [BS] - o_q = i_s + tl.arange(0, BS) - m_q = o_q < T - # [BS, BK] - b_q = tl.load(p_q, boundary_check=(0, 1)) - # [BS, BV] - b_do = tl.load(p_do, boundary_check=(0, 1)) - # [BS] - b_lse = tl.load(p_lse, boundary_check=(0,)) - b_delta = tl.load(p_delta, boundary_check=(0,)) - # [BT, BS] - b_s = tl.dot(b_k, tl.trans(b_q)) * scale * RCP_LN2 - if USE_G: - p_gq = tl.make_block_ptr(g_cumsum + bos * HQ + i_hq, (T,), (HQ,), (i_s,), (BS,), (0,)) - b_gq = tl.load(p_gq, boundary_check=(0,)).to(tl.float32) - b_s += b_gq[None, :] - b_gk[:, None] - if USE_WINDOW: - b_p = tl.where( - (o_k[:, None] <= o_q[None, :]) & (o_q[None, :] - o_k[:, None] < W) & m_q[None, :], - exp2(b_s - b_lse[None, :]), 0 - ) - else: - b_p = tl.where((o_k[:, None] <= o_q[None, :]) & m_q[None, :], exp2(b_s - b_lse[None, :]), 0) - # [BT, BS] @ [BS, BV] -> [BT, BV] - b_dv += tl.dot(b_p.to(b_do.dtype), b_do) - # [BT, BV] @ [BV, BS] -> [BT, BS] - b_dp = tl.dot(b_v, tl.trans(b_do)) - # [BT, BS] - b_ds = b_p * (b_dp - b_delta[None, :]) - # [BT, BS] @ [BS, BK] -> [BT, BK] - b_dk += tl.dot(b_ds.to(b_q.dtype), b_q) - if USE_G: - b_dg -= tl.sum(b_ds, 1) - - # for sliding window, limit the range of future query blocks to process. - # a key at position k_pos can only be attended to by queries at positions [k_pos, k_pos + W - 1]. - # so we only need queries up to (i_t + 1) * BT - 1 + W - 1. - i_end = min(tl.cdiv(T, BS) * BS, (i_t + 1) * BT + W - 1) if USE_WINDOW else tl.cdiv(T, BS) * BS - - for i_s in range((i_t + 1) * BT, i_end, BS): - p_q = tl.make_block_ptr(q + (bos * HQ + i_hq) * K, (T, K), (HQ*K, 1), (i_s, 0), (BS, BK), (1, 0)) - p_do = tl.make_block_ptr(do + (bos * HQ + i_hq) * V, (T, V), (HQ*V, 1), (i_s, i_v * BV), (BS, BV), (1, 0)) - p_lse = tl.make_block_ptr(lse + bos * HQ + i_hq, (T,), (HQ,), (i_s,), (BS,), (0,)) - p_delta = tl.make_block_ptr(delta + bos * HQ + i_hq, (T,), (HQ,), (i_s,), (BS,), (0,)) - - # [BS] - o_q = i_s + tl.arange(0, BS) - m_q = o_q < T - # [BS, BK] - b_q = tl.load(p_q, boundary_check=(0, 1)) - # [BS, BV] - b_do = tl.load(p_do, boundary_check=(0, 1)) - # [BS] - b_lse = tl.load(p_lse, boundary_check=(0,)) - b_delta = tl.load(p_delta, boundary_check=(0,)) - # [BT, BS] - b_s = tl.dot(b_k, tl.trans(b_q)) * scale * RCP_LN2 - if USE_G: - p_gq = tl.make_block_ptr(g_cumsum + bos * HQ + i_hq, (T,), (HQ,), (i_s,), (BS,), (0,)) - b_gq = tl.load(p_gq, boundary_check=(0,)).to(tl.float32) - b_s += b_gq[None, :] - b_gk[:, None] - if USE_WINDOW: - b_p = tl.where((o_q[None, :] - o_k[:, None] < W) & m_q[None, :], exp2(b_s - b_lse[None, :]), 0) - else: - b_p = tl.where(m_q[None, :], exp2(b_s - b_lse[None, :]), 0) - # [BT, BS] @ [BS, BV] -> [BT, BV] - b_dv += tl.dot(b_p.to(b_do.dtype), b_do) - # [BT, BV] @ [BV, BS] -> [BT, BS] - b_dp = tl.dot(b_v, tl.trans(b_do)) - # [BT, BS] - b_ds = b_p * (b_dp - b_delta[None, :]) - # [BT, BS] @ [BS, BK] -> [BT, BK] - b_dk += tl.dot(b_ds.to(b_q.dtype), b_q) - if USE_G: - b_dg -= tl.sum(b_ds, 1) - - b_dk = b_dk * scale - tl.store(p_dk, b_dk.to(p_dk.dtype.element_ty), boundary_check=(0, 1)) - tl.store(p_dv, b_dv.to(p_dv.dtype.element_ty), boundary_check=(0, 1)) - if USE_G: - p_dg = tl.make_block_ptr(dg_cumsum + bos * HQ + i_hq, (T,), (HQ,), (i_t * BT,), (BT,), (0,)) - tl.store(p_dg, b_dg.to(p_dg.dtype.element_ty), boundary_check=(0,)) - - -@dispatch('attn') -def parallel_attn_fwd( - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - g_cumsum: torch.Tensor, - sink_bias: torch.Tensor | None, - scale: float, - window_size: int | None = None, - cu_seqlens: torch.LongTensor | None = None, - chunk_indices: torch.LongTensor | None = None, -): - B, T, H, K, V = *k.shape, v.shape[-1] - HQ = q.shape[2] - G = HQ // H - BT = 128 - if check_shared_mem('hopper', q.device.index): - BS = min(64, max(16, triton.next_power_of_2(T))) - BK = min(256, max(16, triton.next_power_of_2(K))) - BV = min(256, max(16, triton.next_power_of_2(V))) - num_warps = 8 - elif check_shared_mem('ampere', q.device.index): - BS = min(32, max(16, triton.next_power_of_2(T))) - BK = min(256, max(16, triton.next_power_of_2(K))) - BV = min(128, max(16, triton.next_power_of_2(V))) - num_warps = 4 - else: - BS = min(32, max(16, triton.next_power_of_2(T))) - BK = min(256, max(16, triton.next_power_of_2(K))) - BV = min(64, max(16, triton.next_power_of_2(V))) - num_warps = 2 - NK = triton.cdiv(K, BK) - NV = triton.cdiv(V, BV) - - if chunk_indices is None and cu_seqlens is not None: - chunk_indices = prepare_chunk_indices(cu_seqlens, BT) - NT = triton.cdiv(T, BT) if cu_seqlens is None else len(chunk_indices) - assert NK == 1, "The key dimension can not be larger than 256" - - o = torch.empty(B, T, HQ, V, dtype=v.dtype, device=q.device) - lse = torch.empty(B, T, HQ, dtype=torch.float, device=q.device) - grid = (NV, NT, B * HQ) - parallel_attn_fwd_kernel[grid]( - q=q, - k=k, - v=v, - o=o, - g_cumsum=g_cumsum, - sink_bias=sink_bias, - lse=lse, - scale=scale, - cu_seqlens=cu_seqlens, - chunk_indices=chunk_indices, - B=B, - T=T, - W=window_size, - H=H, - HQ=HQ, - G=G, - K=K, - V=V, - BT=BT, - BS=BS, - BK=BK, - BV=BV, - num_warps=num_warps, - ) - return o, lse - - -def parallel_attn_bwd_preprocess( - o: torch.Tensor, - do: torch.Tensor, -): - V = o.shape[-1] - delta = torch.empty_like(o[..., 0], dtype=torch.float) - parallel_attn_bwd_kernel_preprocess[(delta.numel(),)]( - o=o, - do=do, - delta=delta, - B=triton.next_power_of_2(V), - V=V, - ) - return delta - - -@dispatch('attn') -def parallel_attn_bwd( - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - o: torch.Tensor, - g_cumsum: torch.Tensor, - lse: torch.Tensor, - do: torch.Tensor, - sink_bias: torch.Tensor | None = None, - scale: float = None, - window_size: int | None = None, - chunk_size: int = 128, - cu_seqlens: torch.LongTensor | None = None, - chunk_indices: torch.LongTensor | None = None, -): - B, T, H, K, V = *k.shape, v.shape[-1] - HQ = q.shape[2] - G = HQ // H - # dq/dk are reduced over the full value dim in one program (no cross-program accumulation), - # so BV must span all of V (NV == 1). Don't cap it here -- the forward can, the backward can't. - if check_shared_mem('hopper'): - BT = 128 - BS = 64 - BK = max(triton.next_power_of_2(K), 16) - BV = max(triton.next_power_of_2(V), 16) - num_warps = 8 - elif check_shared_mem('ampere'): - BS = 32 - BK = max(triton.next_power_of_2(K), 16) - BV = max(triton.next_power_of_2(V), 16) - BT = 128 if K <= 64 else 64 - num_warps = 4 - else: - BT = 64 - BS = 32 - BK = max(triton.next_power_of_2(K), 16) - BV = max(triton.next_power_of_2(V), 16) - num_warps = 2 - - if chunk_indices is None and cu_seqlens is not None: - chunk_indices = prepare_chunk_indices(cu_seqlens, BT) - NT = triton.cdiv(T, BT) if cu_seqlens is None else len(chunk_indices) - NV = triton.cdiv(V, BV) - - delta = parallel_attn_bwd_preprocess(o, do) - - dq = torch.empty(B, T, HQ, K, dtype=k.dtype if H == HQ else torch.float, device=q.device) - dk = torch.empty(B, T, HQ, K, dtype=k.dtype if H == HQ else torch.float, device=q.device) - dv = torch.empty(B, T, HQ, V, dtype=v.dtype if H == HQ else torch.float, device=q.device) - grid = (NV, NT, B * HQ) - - dg_cumsum, dg_cumsum_k = None, None - if g_cumsum is not None: - dg_cumsum = torch.empty(B, T, HQ, dtype=torch.float, device=q.device) - dg_cumsum_k = torch.empty(B, T, HQ, dtype=torch.float, device=q.device) - - parallel_attn_bwd_kernel_dq[grid]( - q=q, - k=k, - v=v, - g_cumsum=g_cumsum, - lse=lse, - delta=delta, - do=do, - dq=dq, - dg_cumsum=dg_cumsum, - cu_seqlens=cu_seqlens, - chunk_indices=chunk_indices, - scale=scale, - T=T, - W=window_size, - B=B, - H=H, - HQ=HQ, - G=G, - K=K, - V=V, - BT=BT, - BS=BS, - BK=BK, - BV=BV, - num_warps=num_warps, - ) - parallel_attn_bwd_kernel_dkv[grid]( - q=q, - k=k, - v=v, - g_cumsum=g_cumsum, - lse=lse, - delta=delta, - do=do, - dk=dk, - dv=dv, - dg_cumsum=dg_cumsum_k, - cu_seqlens=cu_seqlens, - chunk_indices=chunk_indices, - scale=scale, - T=T, - W=window_size, - B=B, - H=H, - HQ=HQ, - G=G, - K=K, - V=V, - BT=BT, - BS=BS, - BK=BK, - BV=BV, - num_warps=num_warps, - ) - dk = reduce(dk, 'b t (h g) k -> b t h k', g=G, reduction='sum') - dv = reduce(dv, 'b t (h g) v -> b t h v', g=G, reduction='sum') - if g_cumsum is not None: - dg_cumsum.add_(dg_cumsum_k) - - dsink_bias = None - if sink_bias is not None: - p_sink_bias = torch.exp2(sink_bias[None, None, :] - lse) - dsink_bias = -(p_sink_bias * delta).sum((0, 1)) - - return dq, dk, dv, dg_cumsum, dsink_bias - - -@torch.compile -class ParallelAttentionFunction(torch.autograd.Function): - - @staticmethod - @contiguous - @autocast_custom_fwd - def forward(ctx, q, k, v, g, sink_bias, scale, window_size, cu_seqlens, chunk_indices=None): - ctx.dtype = q.dtype - - g_cumsum = chunk_global_cumsum(g, cu_seqlens=cu_seqlens, scale=RCP_LN2) if g is not None else None - sink_bias = sink_bias * RCP_LN2 if sink_bias is not None else None - o, lse = parallel_attn_fwd( - q=q, - k=k, - v=v, - g_cumsum=g_cumsum, - sink_bias=sink_bias, - scale=scale, - window_size=window_size, - cu_seqlens=cu_seqlens, - chunk_indices=chunk_indices, - ) - ctx.save_for_backward(q, k, v, o, g_cumsum, lse, sink_bias) - ctx.scale = scale - ctx.window_size = window_size - ctx.cu_seqlens = cu_seqlens - return o.to(q.dtype) - - @staticmethod - @contiguous - @autocast_custom_bwd - def backward(ctx, do): - q, k, v, o, g_cumsum, lse, sink_bias = ctx.saved_tensors - dq, dk, dv, dg, dsink_bias = parallel_attn_bwd( - q=q, - k=k, - v=v, - o=o, - g_cumsum=g_cumsum, - lse=lse, - do=do, - sink_bias=sink_bias, - scale=ctx.scale, - window_size=ctx.window_size, - cu_seqlens=ctx.cu_seqlens, - ) - if dg is not None: - dg = chunk_global_cumsum(dg, cu_seqlens=ctx.cu_seqlens, reverse=True) - - return dq.to(q), dk.to(k), dv.to(v), dg, dsink_bias, None, None, None, None - - -def parallel_attn( - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - g: torch.Tensor | None = None, - scale: float | None = None, - window_size: int | None = None, - cu_seqlens: torch.LongTensor | None = None, - chunk_indices: torch.LongTensor | None = None, - *, - sink_bias: torch.Tensor | None = None, - **kwargs -) -> torch.Tensor: - r""" - Args: - q (torch.Tensor): - queries of shape `[B, T, HQ, K]`. - k (torch.Tensor): - keys of shape `[B, T, H, K]`. - GQA will be applied if HQ is divisible by H. - v (torch.Tensor): - values of shape `[B, T, H, V]`. - g (Optional[torch.Tensor]): - log decay factors of shape `[B, T, HQ]`. - scale (Optional[float]): - Scale factor for attention scores. - If not provided, it will default to `1 / sqrt(K)`. Default: `None`. - window_size (Optional[int]): - Sliding window size. If provided, each query at position i only attends to - keys in `[i - window_size + 1, i]`. If `None`, full causal attention is used. - Default: `None`. - cu_seqlens (torch.LongTensor): - Cumulative sequence lengths of shape `[N+1]` used for variable-length training, - consistent with the FlashAttention API. - sink_bias (Optional[torch.Tensor]): - Per-query-head attention-sink bias logits of shape `[HQ]` — one - learnable scalar per query head, as introduced by GPT-OSS. - - Augments the softmax denominator with `exp(sink_bias[h])` without - adding a corresponding key/value entry, so the model can route - attention mass to a learnable "no-op" target: - p_i = exp(s_i) / (sum_j exp(s_j) + exp(sink_bias[h])) - o = sum_i p_i * v_i # sink slot contributes no value - When `None`, standard softmax is used. Reserved name: the future - `sink_tokens_*` kwargs will support Xiao 2024-style K/V sink tokens - and may be combined with `sink_bias`. - - Returns: - o (torch.Tensor): - Outputs of shape `[B, T, HQ, V]`. - """ - if 'head_first' in kwargs: - raise DeprecationWarning( - "head_first has been removed. Inputs must be in `[B, T, H, ...]` format.", - ) - if scale is None: - scale = k.shape[-1] ** -0.5 - if cu_seqlens is not None and q.shape[0] != 1: - raise ValueError( - f"The batch size is expected to be 1 rather than {q.shape[0]} when using `cu_seqlens`. " - f"Please flatten variable-length inputs before processing.", - ) - if sink_bias is not None: - assert sink_bias.shape == (q.shape[2],), "sink_bias must have shape [HQ]" - - o = ParallelAttentionFunction.apply( - q, k, v, g, sink_bias, scale, window_size, cu_seqlens, chunk_indices - ) - return o diff --git a/build/torch-xpu/ops/attnres/__init__.py b/build/torch-xpu/ops/attnres/__init__.py deleted file mode 100644 index 95ddd9b74d8148bd094c2ea281fa9611957c2b92..0000000000000000000000000000000000000000 --- a/build/torch-xpu/ops/attnres/__init__.py +++ /dev/null @@ -1,14 +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 ...ops.attnres.fused import fused_attnres -from ...ops.attnres.naive import naive_attnres - -__all__ = [ - 'fused_attnres', - 'naive_attnres', -] diff --git a/build/torch-xpu/ops/attnres/fused.py b/build/torch-xpu/ops/attnres/fused.py deleted file mode 100644 index ea0389b68b6abce3dd143aa460db5e86f043eced..0000000000000000000000000000000000000000 --- a/build/torch-xpu/ops/attnres/fused.py +++ /dev/null @@ -1,568 +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 - -from collections.abc import Sequence - -import torch -import triton -import triton.language as tl - -from ...ops.utils.cache import fla_cache_autotune -from ...ops.utils.op import exp -from ...utils import ( - autocast_custom_bwd, - autocast_custom_fwd, - autotune_cache_kwargs, - input_guard, -) - - -@fla_cache_autotune( - configs=[ - triton.Config({'BL': BL}, num_warps=num_warps, num_stages=num_stages) - for BL in [1, 2, 4, 8] - for num_warps in [4, 8, 16] - for num_stages in [2, 3] - ], - key=['L2', 'D', 'HAS_ONORM', 'SAVE_OPRE'], - **autotune_cache_kwargs, -) -@triton.jit(do_not_specialize=['L']) -def attnres_fwd_kernel( - q, - res, - w, - ow, - o, - o_pre, - rstd, - logit, - lse, - N, - L, - L2: tl.constexpr, - D: tl.constexpr, - eps: tl.constexpr, - scale: tl.constexpr, - BL: tl.constexpr, - BD: tl.constexpr, - HAS_ONORM: tl.constexpr, - SAVE_OPRE: tl.constexpr, -): - i_n = tl.program_id(0).to(tl.int64) - - # [BD] - o_d = tl.max_contiguous(tl.multiple_of(tl.arange(0, BD), BD), BD) - m_d = o_d < D - # [BD] q * w, reused across all residual-source tiles - b_qw = tl.load(q + o_d, mask=m_d, other=0.).to(tl.float32) * tl.load(w + o_d, mask=m_d, other=0.).to(tl.float32) - - # online softmax over L; b_o accumulates in registers so each v tile is read once (logit + weighted sum) - b_m = tl.full([], float('-inf'), dtype=tl.float32) - b_acc = tl.zeros([], dtype=tl.float32) - b_o = tl.zeros([BD], dtype=tl.float32) - for i_l in range(tl.cdiv(L, BL)): - # [BL] - o_l = i_l * BL + tl.arange(0, BL) - m_l = o_l < L - # per-tile base pointers from the length-L2 padded tuple; OOB rows keep res[0] and are masked by m_l - p_v = res[0] + o_l * 0 - for i in tl.static_range(1, L2): - p_v = tl.where(o_l == i, res[i], p_v) - p_v = tl.multiple_of(p_v, 16) - - # [BL, BD] gather: row l from source l at offset i_n*D + o_d - b_v = tl.load( - tl.multiple_of(p_v[:, None] + (i_n * D + o_d[None, :]), (1, 16)), - mask=m_l[:, None] & m_d[None, :], - other=0.0, - eviction_policy="evict_first", - ).to(tl.float32) - - # [BL] - b_rstd = tl.rsqrt(tl.sum(b_v * b_v, axis=1) / D + eps) - b_logit = tl.sum(b_v * b_qw[None, :], axis=1) * b_rstd - b_s = tl.where(m_l, b_logit * scale, float('-inf')) - - b_m, b_mp = tl.maximum(b_m, tl.max(b_s, axis=0)), b_m - b_r = exp(b_mp - b_m) - # [BL] - b_p = exp(b_s - b_m) - b_acc = b_acc * b_r + tl.sum(b_p, axis=0) - # [BD] - b_o = b_o * b_r + tl.sum(b_p[:, None] * b_v, axis=0) - - # rstd and logit saved for bwd_dv - p_rstd = tl.make_block_ptr(rstd + i_n, (L,), (N,), (i_l * BL,), (BL,), (0,)) - p_logit = tl.make_block_ptr(logit + i_n, (L,), (N,), (i_l * BL,), (BL,), (0,)) - tl.store(p_rstd, b_rstd.to(rstd.dtype.element_ty), boundary_check=(0,)) - tl.store(p_logit, b_logit.to(logit.dtype.element_ty), boundary_check=(0,)) - - tl.store(lse + i_n, b_m + tl.log(b_acc)) - - # [BD] pre-norm mixed residual sum_l p_l * v_l - b_o = b_o / b_acc - if SAVE_OPRE: - p_o_pre = tl.make_block_ptr(o_pre + i_n * D, (D,), (1,), (0,), (BD,), (0,)) - tl.store(p_o_pre, b_o.to(p_o_pre.dtype.element_ty), boundary_check=(0,)) - # fold the optional output RMSNorm into the returned output o (o_rstd is recomputed from o_pre in bwd, not stored) - if HAS_ONORM: - b_o_rstd = tl.rsqrt(tl.sum(tl.where(m_d, b_o * b_o, 0.0), axis=0) / D + eps) - b_ow = tl.load(ow + o_d, mask=m_d, other=0.).to(tl.float32) - b_o = b_o * b_o_rstd * b_ow - p_o = tl.make_block_ptr(o + i_n * D, (D,), (1,), (0,), (BD,), (0,)) - tl.store(p_o, b_o.to(p_o.dtype.element_ty), boundary_check=(0,)) - - -@fla_cache_autotune( - configs=[ - triton.Config({'BL': BL}, num_warps=num_warps, num_stages=num_stages) - for BL in [1, 2, 4, 8] - for num_warps in [4, 8, 16] - for num_stages in [2, 3] - ], - key=['L2', 'D', 'HAS_ONORM', 'SAVE_OPRE'], - **autotune_cache_kwargs, -) -@triton.jit(do_not_specialize=['L']) -def attnres_bwd_kernel_dv( - q, - res, - w, - ow, - o_pre, - rstd, - logit, - lse, - do, - dres, - dqw, - dow_partial, - N, - L, - L2: tl.constexpr, - D: tl.constexpr, - eps: tl.constexpr, - scale: tl.constexpr, - BL: tl.constexpr, - BD: tl.constexpr, - HAS_ONORM: tl.constexpr, - SAVE_OPRE: tl.constexpr, -): - i_n = tl.program_id(0).to(tl.int64) - - # [BD] - o_d = tl.max_contiguous(tl.multiple_of(tl.arange(0, BD), BD), BD) - m_d = o_d < D - b_qw = tl.load(q + o_d, mask=m_d, other=0.).to(tl.float32) * tl.load(w + o_d, mask=m_d, other=0.).to(tl.float32) - b_lse = tl.load(lse + i_n).to(tl.float32) - p_do = tl.make_block_ptr(do + i_n * D, (D,), (1,), (0,), (BD,), (0,)) - b_do = tl.load(p_do, boundary_check=(0,), padding_option="zero").to(tl.float32) - if SAVE_OPRE: - p_o_pre = tl.make_block_ptr(o_pre + i_n * D, (D,), (1,), (0,), (BD,), (0,)) - b_o_pre = tl.load(p_o_pre, boundary_check=(0,), padding_option="zero").to(tl.float32) - else: - # level 1: recompute the mix sum_l p_l * v_l from V - b_o_pre = tl.zeros([BD], dtype=tl.float32) - for i_l in range(tl.cdiv(L, BL)): - o_l = i_l * BL + tl.arange(0, BL) - m_l = o_l < L - p_v = res[0] + o_l * 0 - for i in tl.static_range(1, L2): - p_v = tl.where(o_l == i, res[i], p_v) - p_v = tl.multiple_of(p_v, 16) - b_v = tl.load( - tl.multiple_of(p_v[:, None] + (i_n * D + o_d[None, :]), (1, 16)), - mask=m_l[:, None] & m_d[None, :], - other=0.0, - ).to(tl.float32) - p_logit = tl.make_block_ptr(logit + i_n, (L,), (N,), (i_l * BL,), (BL,), (0,)) - b_logit = tl.load(p_logit, boundary_check=(0,), padding_option="zero").to(tl.float32) - b_p = tl.where(m_l, exp(b_logit * scale - b_lse), 0.0) - b_o_pre += tl.sum(b_p[:, None] * b_v, axis=0) - - # output RMSNorm bwd: turn b_do into the gradient w.r.t. the pre-norm output and stage dow_partial - if HAS_ONORM: - b_o_rstd = tl.rsqrt(tl.sum(tl.where(m_d, b_o_pre * b_o_pre, 0.0), axis=0) / D + eps) - b_ow = tl.load(ow + o_d, mask=m_d, other=0.).to(tl.float32) - b_xhat = b_o_pre * b_o_rstd - b_c1 = tl.sum(tl.where(m_d, b_xhat * b_ow * b_do, 0.0), axis=0) / D - p_dow = tl.make_block_ptr(dow_partial + i_n * D, (D,), (1,), (0,), (BD,), (0,)) - tl.store(p_dow, (b_xhat * b_do).to(p_dow.dtype.element_ty), boundary_check=(0,)) - b_do = (b_ow * b_do - b_xhat * b_c1) * b_o_rstd - # delta = sum_l p*dp = - b_delta = tl.sum(tl.where(m_d, b_do * b_o_pre, 0.0), axis=0) - - # [BD] dqw accumulates over the L tiles - b_dqw = tl.zeros([BD], dtype=tl.float32) - for i_l in range(tl.cdiv(L, BL)): - # [BL] - o_l = i_l * BL + tl.arange(0, BL) - m_l = o_l < L - m_v = m_l[:, None] & m_d[None, :] - # per-tile source / dv base pointers from the length-L2 padded tuple - p_v = res[0] + o_l * 0 - p_dv = dres[0] + o_l * 0 - for i in tl.static_range(1, L2): - p_v = tl.where(o_l == i, res[i], p_v) - p_dv = tl.where(o_l == i, dres[i], p_dv) - p_v = tl.multiple_of(p_v, 16) - p_dv = tl.multiple_of(p_dv, 16) - # [BL, BD] v tile, read once and reused for dp / dv / dqw - b_v = tl.load( - tl.multiple_of(p_v[:, None] + (i_n * D + o_d[None, :]), (1, 16)), - mask=m_v, - other=0.0, - ).to(tl.float32) - - p_rstd = tl.make_block_ptr(rstd + i_n, (L,), (N,), (i_l * BL,), (BL,), (0,)) - p_logit = tl.make_block_ptr(logit + i_n, (L,), (N,), (i_l * BL,), (BL,), (0,)) - # [BL]; recompute probs from logit + lse, OOB rows masked to 0 - b_rstd = tl.load(p_rstd, boundary_check=(0,), padding_option="zero").to(tl.float32) - b_logit = tl.load(p_logit, boundary_check=(0,), padding_option="zero").to(tl.float32) - b_p = tl.where(m_l, exp(b_logit * scale - b_lse), 0.0) - - # softmax bwd with delta already known - b_dp = tl.sum(b_v * b_do[None, :], axis=1) - b_ds = b_p * (b_dp - b_delta) * scale - # [BL, BD] - b_k = b_v * b_rstd[:, None] - b_dv = b_p[:, None] * b_do[None, :] + (b_ds * b_rstd)[:, None] * (b_qw[None, :] - b_k * (b_logit / D)[:, None]) - tl.store( - tl.multiple_of(p_dv[:, None] + (i_n * D + o_d[None, :]), (1, 16)), - b_dv.to(dres[0].dtype.element_ty), - mask=m_v, - ) - # [BD] - b_dqw += tl.sum(b_ds[:, None] * b_k, axis=0) - - p_dqw = tl.make_block_ptr(dqw + i_n * D, (D,), (1,), (0,), (BD,), (0,)) - tl.store(p_dqw, b_dqw, boundary_check=(0,)) - - -@fla_cache_autotune( - configs=[ - triton.Config({'BN': BN, 'BD': BD}, num_warps=num_warps, num_stages=num_stages) - for BN, BD, num_warps in [(1024, 16, 4), (2048, 32, 4), (2048, 32, 8), (4096, 32, 8), (4096, 64, 8)] - for num_stages in [3, 4] - ], - key=['N', 'D', 'HAS_ONORM'], - **autotune_cache_kwargs, -) -@triton.jit -def attnres_bwd_kernel_dqdw( - q, - w, - dqw, - dow_partial, - dq, - dw, - dow, - N, - D: tl.constexpr, - BN: tl.constexpr, - BD: tl.constexpr, - HAS_ONORM: tl.constexpr, -): - i_d = tl.program_id(0).to(tl.int32) - - # [BD] - o_d = i_d * BD + tl.arange(0, BD) - m_d = o_d < D - - # column-sum dqw (and dow_partial) over the N axis - # [BD] - b_dqw = tl.zeros([BD], dtype=tl.float32) - b_dow = tl.zeros([BD], dtype=tl.float32) - for i_n in range(0, N, BN): - p_dqw = tl.make_block_ptr(dqw, (N, D), (D, 1), (i_n, i_d * BD), (BN, BD), (1, 0)) - b_dqw += tl.sum(tl.load(p_dqw, boundary_check=(0, 1), padding_option="zero").to(tl.float32), axis=0) - if HAS_ONORM: - p_dow = tl.make_block_ptr(dow_partial, (N, D), (D, 1), (i_n, i_d * BD), (BN, BD), (1, 0)) - b_dow += tl.sum(tl.load(p_dow, boundary_check=(0, 1), padding_option="zero").to(tl.float32), axis=0) - - # the logit uses the q * w product, so dq = (sum_n dqw) * w and dw = (sum_n dqw) * q - # [BD] - b_q = tl.load(q + o_d, mask=m_d, other=0.).to(tl.float32) - b_w = tl.load(w + o_d, mask=m_d, other=0.).to(tl.float32) - tl.store(dq + o_d, b_dqw * b_w, mask=m_d) - tl.store(dw + o_d, b_dqw * b_q, mask=m_d) - if HAS_ONORM: - tl.store(dow + o_d, b_dow, mask=m_d) - - -def _build_ptr_table(tensors: Sequence[torch.Tensor]) -> tuple[torch.Tensor, ...]: - # pad the per-source tensor tuple to a fixed length so Triton can compile a single kernel per L2 bucket. - # the tuple length is part of the kernel's compile signature; padded slots are address-only (never read/written). - L2 = max(8, triton.next_power_of_2(len(tensors))) - assert 1 <= len(tensors) <= L2 - for t in tensors: - assert t.data_ptr() % 16 == 0, "attnres residual sources must be 16-byte aligned" - return tuple(tensors) + (tensors[0],) * (L2 - len(tensors)) - - -def fused_attnres_fwd( - q: torch.Tensor, - residuals: Sequence[torch.Tensor], - res: tuple[torch.Tensor, ...], - w: torch.Tensor, - ow: torch.Tensor | None, - eps: float, - scale: float, - checkpoint_level: int, -) -> tuple[torch.Tensor, torch.Tensor | None, torch.Tensor, torch.Tensor, torch.Tensor]: - if not residuals[0].is_cuda: - raise ValueError("Triton attnres requires CUDA tensors") - - output_shape = residuals[0].shape - L, N, D = len(residuals), residuals[0].numel() // output_shape[-1], output_shape[-1] - - dtype = residuals[0].dtype - stats_shape = (L, *output_shape[:-1]) - - has_onorm = ow is not None - save_opre = checkpoint_level == 0 - o = torch.empty(output_shape, device=residuals[0].device, dtype=dtype) - o_pre = torch.empty(output_shape, device=residuals[0].device, dtype=dtype) if save_opre else None - lse = torch.empty(output_shape[:-1], device=residuals[0].device, dtype=torch.float32) - rstd = torch.empty(stats_shape, device=residuals[0].device, dtype=torch.float32) - logit = torch.empty_like(rstd) - - L2 = max(8, triton.next_power_of_2(L)) - attnres_fwd_kernel[(N,)]( - q=q, - res=res, - w=w, - ow=ow, - o=o, - o_pre=o_pre, - rstd=rstd, - logit=logit, - lse=lse, - N=N, - L=L, - L2=L2, - D=D, - eps=eps, - scale=scale, - BD=triton.next_power_of_2(D), - HAS_ONORM=has_onorm, - SAVE_OPRE=save_opre, - ) - - return o, o_pre, rstd, logit, lse - - -def fused_attnres_bwd( - do: torch.Tensor, - q: torch.Tensor, - residuals: Sequence[torch.Tensor], - res: tuple[torch.Tensor, ...], - w: torch.Tensor, - ow: torch.Tensor | None, - o_pre: torch.Tensor | None, - rstd: torch.Tensor, - logit: torch.Tensor, - lse: torch.Tensor, - eps: float, - scale: float, - checkpoint_level: int, -) -> tuple[list[torch.Tensor], torch.Tensor, torch.Tensor, torch.Tensor | None]: - L, N, D = len(residuals), do.numel() // do.shape[-1], do.shape[-1] - - # bwd_dv produces dvs + dqw (and dow_partial when fusing output RMSNorm); bwd_dqdw reduces dqw / dow_partial over N - # into dq / dw / dow. - has_onorm = ow is not None - if has_onorm: - dow_partial = torch.empty_like(do, dtype=torch.float32) - dow = torch.empty_like(ow) - else: - dow_partial = dow = None - - dvs = [torch.empty_like(r) for r in residuals] - dres = _build_ptr_table(dvs) - dqw = torch.empty_like(do, dtype=torch.float32) - dq = torch.empty_like(q) - dw = torch.empty_like(w) - - L2 = max(8, triton.next_power_of_2(L)) - attnres_bwd_kernel_dv[(N,)]( - q=q, - res=res, - w=w, - ow=ow, - o_pre=o_pre, - rstd=rstd, - logit=logit, - lse=lse, - do=do, - dres=dres, - dqw=dqw, - dow_partial=dow_partial, - N=N, - L=L, - L2=L2, - D=D, - eps=eps, - scale=scale, - BD=triton.next_power_of_2(D), - HAS_ONORM=has_onorm, - SAVE_OPRE=checkpoint_level == 0, - ) - - def grid(meta): return (triton.cdiv(D, meta['BD']),) - attnres_bwd_kernel_dqdw[grid]( - q=q, - w=w, - dqw=dqw, - dow_partial=dow_partial, - dq=dq, - dw=dw, - dow=dow, - N=N, - D=D, - HAS_ONORM=has_onorm, - ) - - return dvs, dq, dw, dow - - -class FusedAttnresFunction(torch.autograd.Function): - - @staticmethod - @input_guard - @autocast_custom_fwd - def forward( - ctx, - query: torch.Tensor, - rms_weight: torch.Tensor, - output_rms_weight: torch.Tensor | None, - rms_eps: float, - scale: float, - return_weights: bool, - checkpoint_level: int, - *residuals: torch.Tensor, - ) -> tuple[torch.Tensor, torch.Tensor]: - # `res` is built once here and threaded through fwd/bwd so neither internal wrapper rebuilds it. - # the tuple is pure Python, no H2D copy. - res = _build_ptr_table(residuals) - o, o_pre, rstd, logit, lse = fused_attnres_fwd( - q=query, - residuals=residuals, - res=res, - w=rms_weight, - ow=output_rms_weight, - eps=rms_eps, - scale=scale, - checkpoint_level=checkpoint_level, - ) - ctx.save_for_backward(query, rms_weight, output_rms_weight, o_pre, rstd, logit, lse, *residuals) - ctx.eps = rms_eps - ctx.scale = scale - ctx.checkpoint_level = checkpoint_level - ctx.res = res - # probs are materialized only when requested; bwd recomputes them from logit + lse - p = (logit * scale - lse).exp() if return_weights else o.new_empty(0) - ctx.mark_non_differentiable(p) - return o, p - - @staticmethod - @input_guard - @autocast_custom_bwd - def backward( - ctx, - do: torch.Tensor, - dp: torch.Tensor | None = None, - ): - del dp - query, rms_weight, output_rms_weight, o_pre, rstd, logit, lse, *residuals = ctx.saved_tensors - dvs, dq, dw, dow = fused_attnres_bwd( - do=do, - q=query, - residuals=residuals, - res=ctx.res, - w=rms_weight, - ow=output_rms_weight, - o_pre=o_pre, - rstd=rstd, - logit=logit, - lse=lse, - eps=ctx.eps, - scale=ctx.scale, - checkpoint_level=ctx.checkpoint_level, - ) - return (dq, dw, dow, None, None, None, None, *dvs) - - -def fused_attnres( - query: torch.Tensor, - residuals: Sequence[torch.Tensor], - rms_weight: torch.Tensor, - output_rms_weight: torch.Tensor | None = None, - rms_eps: float = 1e-6, - scale: float = 1.0, - return_weights: bool = False, - checkpoint_level: int = 1, -) -> torch.Tensor | tuple[torch.Tensor, ...]: - r""" - Apply AttnRes residual aggregation. - - AttnRes normalizes each residual source with RMSNorm, scores it against `query`, applies softmax over the - residual-source dimension, and returns the weighted sum of residual sources. - See `Attention Residuals `_. - - Args: - query (torch.Tensor): - Per-layer pseudo-query of shape `[D]` or `[D, 1]`, where `D` is the hidden size. - residuals (Sequence[torch.Tensor]): - Non-empty sequence of same-dtype, same-`D` residual sources, each of shape `[..., D]`. - rms_weight (torch.Tensor): - RMSNorm scale for key normalization of shape `[D]`. - output_rms_weight (torch.Tensor, optional): - If set, an extra RMSNorm with this weight is applied to the mixed residual before returning, fusing the - prenorm that would otherwise follow the AttnRes call (e.g. `attn_norm` / `mlp_norm`). Default: `None`. - rms_eps (float): - RMSNorm epsilon (also used for `output_rms_weight` when set). Default: `1e-6`. - scale (float): - Scale factor applied to AttnRes logits before softmax. Default: `1.0`. - return_weights (bool): - Whether to return depth softmax probabilities. Default: `False`. - checkpoint_level (int): - Backward memory/recompute trade-off. - `0` keeps the mixed residual; - `1` drops it and recomputes it from the sources in backward (less memory, one extra read). - Default: `1`. - - Returns: - o (torch.Tensor): - Mixed residual of shape `[..., D]`. - p (torch.Tensor): - Depth softmax probabilities of shape `[L, ...]` if `return_weights=True`, otherwise not returned. - """ - if len(residuals) == 0: - raise ValueError("residuals must contain at least one source") - if checkpoint_level not in (0, 1): - raise ValueError(f"checkpoint_level must be 0 or 1, got {checkpoint_level}") - - output_shape = residuals[0].shape - D = output_shape[-1] - flat_residuals = tuple(r.reshape(-1, D).contiguous() for r in residuals) - - o, p = FusedAttnresFunction.apply( - query, rms_weight, output_rms_weight, rms_eps, scale, return_weights, checkpoint_level, *flat_residuals, - ) - o = o.view(output_shape) - - if return_weights: - p = p.view(len(residuals), *output_shape[:-1]) - return o, p - return o - - -__all__ = ["fused_attnres"] diff --git a/build/torch-xpu/ops/attnres/naive.py b/build/torch-xpu/ops/attnres/naive.py deleted file mode 100644 index 2d83ae4f1a008917d3c20e8a06f0595f589dbada..0000000000000000000000000000000000000000 --- a/build/torch-xpu/ops/attnres/naive.py +++ /dev/null @@ -1,83 +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 - -from collections.abc import Sequence - -import torch -import torch.nn.functional as F -from einops import einsum - - -def naive_attnres( - query: torch.Tensor, - residuals: Sequence[torch.Tensor], - rms_weight: torch.Tensor, - output_rms_weight: torch.Tensor | None = None, - rms_eps: float = 1e-6, - scale: float = 1.0, - return_weights: bool = False, -) -> torch.Tensor | tuple[torch.Tensor, ...]: - r""" - Apply AttnRes residual aggregation. - - AttnRes normalizes each residual source with RMSNorm, scores it against `query`, applies softmax over the - residual-source dimension, and returns the weighted sum of residual sources. - See `Attention Residuals `_. - - Args: - query (torch.Tensor): - Per-layer pseudo-query of shape `[D]` or `[D, 1]`, where `D` is the hidden size. - residuals (Sequence[torch.Tensor]): - Non-empty sequence of same-dtype, same-`D` residual sources, each of shape `[..., D]`. - rms_weight (torch.Tensor): - RMSNorm scale for key normalization of shape `[D]`. - output_rms_weight (torch.Tensor, optional): - If set, an extra RMSNorm with this weight is applied to the mixed residual before returning, fusing the - prenorm that would otherwise follow the AttnRes call (e.g. `attn_norm` / `mlp_norm`). Default: `None`. - rms_eps (float): - RMSNorm epsilon (also used for `output_rms_weight` when set). Default: `1e-6`. - scale (float): - Scale factor applied to AttnRes logits before softmax. Default: `1.0`. - return_weights (bool): - Whether to return depth softmax probabilities. Default: `False`. - - Returns: - o (torch.Tensor): - Mixed residual of shape `[..., D]`. - p (torch.Tensor): - Depth softmax probabilities of shape `[L, ...]` if `return_weights=True`, otherwise not returned. - """ - if len(residuals) == 0: - raise ValueError("residuals must contain at least one source") - - output_shape = residuals[0].shape - D = output_shape[-1] - # the stack here is part of the *reference* impl, not the API; it lets einsum-style math run on one tensor while autograd - # still routes per-source gradients back to each leaf through the inner views. - stacked = torch.stack(tuple(residual.view(-1, D) for residual in residuals), dim=0) - - # all math runs in fp32 end-to-end; the final downcast back to the residual dtype happens once, just before returning. - v = stacked.float() - k = F.rms_norm(v, (D,), rms_weight.flatten().float(), rms_eps) - p = (einsum(k, query.flatten().float() * scale, "l ... d, d -> l ...")).softmax(dim=0) - o = einsum(p, v, "l ..., l ... d -> ... d") - o = o.view(output_shape) - - if output_rms_weight is not None: - o = F.rms_norm(o, (D,), output_rms_weight.float(), rms_eps) - - o = o.to(stacked.dtype) - - if return_weights: - p = p.view(len(residuals), *output_shape[:-1]) - return o, p - return o - - -__all__ = ["naive_attnres"] diff --git a/build/torch-xpu/ops/backends/__init__.py b/build/torch-xpu/ops/backends/__init__.py deleted file mode 100644 index 3d9cd4a2770287cfb5dcbecd5aa8dbc2ef53743c..0000000000000000000000000000000000000000 --- a/build/torch-xpu/ops/backends/__init__.py +++ /dev/null @@ -1,223 +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 - -"""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") - - -class BaseBackend: - """Base class for operation-specific backends. - - 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. - `None` indicates no external dependency. Default: `None`. - env_var (str, Optional): - Environment variable name that controls whether the backend is enabled. - `None` means always enabled. Default: `None`. - default_enable (bool, Optional): - Whether the backend is enabled by default when `env_var` is not set. - Set to `False` to require explicit user opt-in. Default: `True`. - priority (int, Optional): - Backend priority. Lower values indicate higher priority. Default: 5. - """ - - 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: - if cls.package_name is None: - 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'] diff --git a/build/torch-xpu/ops/based/__init__.py b/build/torch-xpu/ops/based/__init__.py deleted file mode 100644 index 6206f5a95016cbfeb4b217c527ab084294ea78b2..0000000000000000000000000000000000000000 --- a/build/torch-xpu/ops/based/__init__.py +++ /dev/null @@ -1,14 +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 .fused_chunk import fused_chunk_based -from .parallel import parallel_based - -__all__ = [ - 'fused_chunk_based', - 'parallel_based', -] diff --git a/build/torch-xpu/ops/based/fused_chunk.py b/build/torch-xpu/ops/based/fused_chunk.py deleted file mode 100644 index f1652f6a6023e6f04ad580d93474523137ea9e13..0000000000000000000000000000000000000000 --- a/build/torch-xpu/ops/based/fused_chunk.py +++ /dev/null @@ -1,375 +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 -import triton.language as tl - -from ...utils import autocast_custom_bwd, autocast_custom_fwd, input_guard - - -@triton.jit(do_not_specialize=['T']) -def fused_chunk_based_fwd_kernel( - q, - k, - v, - o, - z, - scale, # K ** -0.5 - T, - B: tl.constexpr, - H: tl.constexpr, - K: tl.constexpr, - V: tl.constexpr, - BT: tl.constexpr, - BK: tl.constexpr, - BV: tl.constexpr, -): - i_v, i_k, i_bh = tl.program_id(0), tl.program_id(1), tl.program_id(2) - - o_i = tl.arange(0, BT) - - # [BT, BT] - m_s = o_i[:, None] >= o_i[None, :] - - # [BV], zero-order taylor expansion - b_h_0o = tl.zeros([BV], dtype=tl.float32) - # [BK, BV], first-order taylor expansion - b_h_1o = tl.zeros([BK, BV], dtype=tl.float32) - # [BK, BK, BV] second-order taylor expansion - b_h_2o = tl.zeros([BK*BK, BV], dtype=tl.float32) - - # make block pointers - p_q = tl.make_block_ptr(q + i_bh * T*K, (T, K), (K, 1), (0, i_k * BK), (BT, BK), (1, 0)) - p_k = tl.make_block_ptr(k + i_bh * T*K, (K, T), (1, K), (i_k * BK, 0), (BK, BT), (0, 1)) - p_v = tl.make_block_ptr(v + i_bh * T*V, (T, V), (V, 1), (0, i_v * BV), (BT, BV), (1, 0)) - p_o = tl.make_block_ptr(o + (i_bh + i_k*B*H) * T*V, (T, V), (V, 1), (0, i_v * BV), (BT, BV), (1, 0)) - - p_z = z + (i_bh + i_k * B * H) * T + tl.arange(0, BT) - k_2o = tl.zeros([1, BK * BK], dtype=tl.float32) - k_1o = tl.zeros([1, BK], dtype=tl.float32) - k_0o = 0 - - for i in range(0, tl.cdiv(T, BT)): - # [BK, BT] - b_k = tl.load(p_k, boundary_check=(0, 1)) - # [BK*BK, BT] - b_k_2o = b_k[:, None, :] * b_k[None, :, :] - b_k_2o = tl.reshape(b_k_2o, [BK * BK, BT]).to(b_k.dtype) - # [BT, BV] - b_v = tl.load(p_v, boundary_check=(0, 1)) - # [BT, BK] - b_q = (tl.load(p_q, boundary_check=(0, 1)) * scale).to(b_k.dtype) - b_o = tl.zeros([BT, BV], dtype=tl.float32) - b_z = tl.zeros([BT], dtype=tl.float32) - - # interchunk - # zero-order - b_o += b_h_0o - b_z += k_0o - # first-order - b_o += tl.dot(b_q, b_h_1o.to(b_q.dtype), allow_tf32=False) - b_z += tl.sum(b_q * k_1o, axis=1) - # second-order - b_q_2o = b_q[:, :, None] * b_q[:, None, :] - b_q_2o = tl.reshape(b_q_2o, [BT, BK * BK]).to(b_k.dtype) - b_o += tl.dot(b_q_2o, b_h_2o.to(b_q_2o.dtype), allow_tf32=False) * 0.5 - b_z += tl.sum(b_q_2o * k_2o, axis=1) * 0.5 - - # update running statistics - k_1o += tl.sum(b_k, axis=1)[None, :] - k_2o += tl.sum(b_k_2o, axis=1)[None, :] - k_0o += BT - - # intrachunk - # [BT, BT] - b_s = tl.dot(b_q, b_k, allow_tf32=False) - b_s = 1 + b_s + 0.5 * b_s * b_s - b_s = tl.where(m_s, b_s, 0) - b_z += tl.sum(b_s, axis=1) - b_o += tl.dot(b_s.to(b_q.dtype), b_v, allow_tf32=False) - # [TB, BV] - tl.store(p_o, b_o.to(p_o.dtype.element_ty), boundary_check=(0, 1)) - tl.store(p_z, b_z.to(p_z.dtype.element_ty), mask=(i * BT + tl.arange(0, BT)) < T) - - # update hidden state - # [BK, BV] - b_h_2o = b_h_2o + tl.dot(b_k_2o.to(b_v.dtype), b_v, allow_tf32=False) - b_h_1o = b_h_1o + tl.dot(b_k, b_v, allow_tf32=False) - b_h_0o = b_h_0o + tl.sum(b_v, axis=0) - - p_q = tl.advance(p_q, (BT, 0)) - p_k = tl.advance(p_k, (0, BT)) - p_v = tl.advance(p_v, (BT, 0)) - p_o = tl.advance(p_o, (BT, 0)) - p_z += BT - - -# Similar to Algorithm1 of https://arxiv.org/abs/2006.16236 -@triton.jit -def fused_chunk_based_bwd_kernel( - # NV: number of split in the V dimension. NK: number of split in the K dimension - q, - k, - v, - do, - dz, - dq, - dk, - dv, - scale, # K ** -0.5 - T, - B: tl.constexpr, - H: tl.constexpr, - K: tl.constexpr, - V: tl.constexpr, - BT: tl.constexpr, - BK: tl.constexpr, - BV: tl.constexpr, -): - i_v, i_k, i_bh = tl.program_id(0), tl.program_id(1), tl.program_id(2) - - o_i = tl.arange(0, BT) - m_s = o_i[:, None] >= o_i[None, :] - - # [BV], zero-order taylor expansion - # b_h_0o = tl.zeros([BV], dtype=tl.float32) - # [BK, BV], first-order taylor expansion - b_h_1o = tl.zeros([BV, BK], dtype=tl.float32) - # [BK, BK, BV] second-order taylor expansion - b_h_2o = tl.zeros([BV, BK*BK], dtype=tl.float32) - - k_1o = tl.zeros([1, BK], dtype=tl.float32) - k_2o = tl.zeros([1, BK * BK], dtype=tl.float32) - - for i in range(0, tl.cdiv(T, BT)): - p_q = tl.make_block_ptr(q + i_bh * T*K, (T, K), (K, 1), (i * BT, i_k * BK), (BT, BK), (1, 0)) - p_k = tl.make_block_ptr(k + i_bh * T*K, (T, K), (K, 1), (i * BT, i_k * BK), (BT, BK), (1, 0)) - p_v = tl.make_block_ptr(v + i_bh * T*V, (V, T), (1, V), (i_v * BV, i * BT), (BV, BT), (0, 1)) - p_do = tl.make_block_ptr(do + i_bh * T*V, (T, V), (V, 1), (i * BT, i_v * BV), (BT, BV), (1, 0)) - p_dq = tl.make_block_ptr(dq + (i_bh + i_v*B*H) * T*K, (T, K), (K, 1), (i*BT, i_k*BK), (BT, BK), (1, 0)) - p_dz = dz + (i_bh) * T + tl.arange(0, BT) + i * BT - b_dq = tl.zeros([BT, BK], dtype=tl.float32) - - # load tensors - # [BT, BK] - b_q = tl.load(p_q, boundary_check=(0, 1)) - b_q = (b_q * scale).to(b_q.dtype) - b_k = tl.load(p_k, boundary_check=(0, 1)) - b_do = tl.load(p_do, boundary_check=(0, 1)).to(b_q.dtype) - b_dz = tl.load(p_dz, mask=(tl.arange(0, BT) + i * BT) < T) - # [BV, BT] - b_v = tl.load(p_v, boundary_check=(0, 1)) - - # inter-chunk - b_dq += tl.dot(b_do, (b_h_1o).to(b_do.dtype), allow_tf32=False) - if i_v == 0: - b_dq += b_dz[:, None] * k_1o - b_dq_2o = tl.dot(b_do, (b_h_2o).to(b_do.dtype), allow_tf32=False) * 0.5 - if i_v == 0: - b_dq_2o += (b_dz[:, None] * k_2o) * 0.5 - b_dq_2o = tl.reshape(b_dq_2o, [BT, BK, BK]) - b_dq += tl.sum(b_dq_2o * b_q[:, :, None], axis=1) - b_dq += tl.sum(b_dq_2o * b_q[:, None, :], axis=2) - b_dq *= scale - - # intra-chunk - # [BT, BT] - b_ds = tl.dot(b_do, b_v, allow_tf32=False) - if i_v == 0: - b_ds += b_dz[:, None] - b_ds = tl.where(m_s, b_ds, 0) * scale - b_s = tl.dot(b_q, tl.trans(b_k), allow_tf32=False) - b_s = tl.where(m_s, b_s, 0) - b_dq += tl.dot((b_ds * (1 + b_s)).to(b_q.dtype), b_k, allow_tf32=False) - - # store - tl.store(p_dq, b_dq.to(p_dq.dtype.element_ty), boundary_check=(0, 1)) - - # update hidden state - # [BT, BK*BK] - b_k_2o = b_k[:, :, None] * b_k[:, None, :] - b_k_2o = tl.reshape(b_k_2o, [BT, BK * BK]).to(b_k.dtype) - # [BV, BK*BK] - b_h_2o = b_h_2o + tl.dot(b_v, b_k_2o.to(b_v.dtype), allow_tf32=False) - # [BV, BK] - b_h_1o = b_h_1o + tl.dot(b_v, b_k, allow_tf32=False) - - if i_v == 0: - # update running statistics - k_1o += tl.sum(b_k, axis=0)[None, :] - k_2o += tl.sum(b_k_2o, axis=0)[None, :] - - tl.debug_barrier() - b_h_1o = None - b_h_2o = None - - # [BK, BV], first-order taylor expansion - b_dh_1o = tl.zeros([BK, BV], dtype=tl.float32) - # [BK, BK, BV] second-order taylor expansion - b_dh_2o = tl.zeros([BK*BK, BV], dtype=tl.float32) - b_dh_0o = tl.zeros([BV], dtype=tl.float32) - m_s = tl.arange(0, BT)[:, None] <= tl.arange(0, BT)[None, :] - - dq_1o = tl.zeros([1, BK], dtype=tl.float32) - dq_2o = tl.zeros([BK * BK, 1], dtype=tl.float32) - - for i in range(tl.cdiv(T, BT) * BT - BT, -BT, -BT): - p_q = tl.make_block_ptr(q + i_bh * T*K, (K, T), (1, K), (i_k * BK, i), (BK, BT), (0, 1)) - p_k = tl.make_block_ptr(k + i_bh * T*K, (T, K), (K, 1), (i, i_k * BK), (BT, BK), (1, 0)) - p_v = tl.make_block_ptr(v + i_bh * T*V, (T, V), (V, 1), (i, i_v * BV), (BT, BV), (1, 0)) - p_do = tl.make_block_ptr(do + i_bh * T*V, (T, V), (V, 1), (i, i_v * BV), (BT, BV), (1, 0)) - p_dk = tl.make_block_ptr(dk + (i_bh+i_v*B*H) * T*K, (T, K), (K, 1), (i, i_k*BK), (BT, BK), (1, 0)) - p_dv = tl.make_block_ptr(dv + (i_bh+i_k*B*H) * T*V, (T, V), (V, 1), (i, i_v*BV), (BT, BV), (1, 0)) - p_dz = dz + (i_bh) * T + tl.arange(0, BT) + i - - b_dk = tl.zeros([BT, BK], dtype=tl.float32) - b_dv = tl.zeros([BT, BV], dtype=tl.float32) - - b_q = tl.load(p_q, boundary_check=(0, 1)) - b_k = tl.load(p_k, boundary_check=(0, 1)) - b_v = tl.load(p_v, boundary_check=(0, 1)) - b_do = tl.load(p_do, boundary_check=(0, 1)).to(b_q.dtype) - b_dz = tl.load(p_dz, mask=(tl.arange(0, BT)+i) < T) - b_q = (b_q * scale).to(b_k.dtype) - - # intra chunk - b_ds = tl.dot(b_v, tl.trans(b_do), allow_tf32=False) - if i_v == 0: - b_ds += b_dz[None, :] - b_ds = tl.where(m_s, b_ds, 0) - b_s = tl.dot(b_k, b_q, allow_tf32=False) - b_s2 = 1 + b_s + 0.5 * b_s * b_s - b_s = tl.where(m_s, b_s, 0) - b_s2 = tl.where(m_s, b_s2, 0) - b_ds *= (1+b_s) - - b_dk += tl.dot(b_ds.to(b_k.dtype), tl.trans(b_q), allow_tf32=False) - b_dv += tl.dot(b_s2.to(b_do.dtype), b_do, allow_tf32=False) - - # inter chunk - b_k_2o = b_k[:, :, None] * b_k[:, None, :] - b_k_2o = tl.reshape(b_k_2o, [BT, BK * BK]).to(b_k.dtype) - - b_dv += tl.dot(b_k, b_dh_1o.to(b_k.dtype), allow_tf32=False) - b_dv += tl.dot(b_k_2o, b_dh_2o.to(b_k.dtype), allow_tf32=False) - b_dv += b_dh_0o - - b_dk += tl.dot(b_v, tl.trans(b_dh_1o).to(b_k.dtype), allow_tf32=False) - - if i_v == 0: - b_dk += dq_1o - - b_dk_2o = tl.dot(b_dh_2o.to(b_k.dtype), tl.trans(b_v), allow_tf32=False) - if i_v == 0: - b_dk_2o += dq_2o - b_dk_2o = tl.reshape(b_dk_2o, [BK, BK, BT]) - b_k_fp32 = tl.trans(b_k.to(tl.float32)) - b_dk2 = tl.sum(b_dk_2o * b_k_fp32[:, None, :], axis=0) - b_dk2 += tl.sum(b_dk_2o * b_k_fp32[None, :, :], axis=1) - b_dk += tl.trans(b_dk2) - - # hidden state update - b_dh_0o += tl.sum(b_do, axis=0) - b_dh_1o = b_dh_1o + tl.dot(b_q, b_do, allow_tf32=False) - b_q_2o = b_q[None, :, :] * b_q[:, None, :] - b_q_2o = tl.reshape(b_q_2o, [BK * BK, BT]).to(b_k.dtype) - b_dh_2o = b_dh_2o + tl.dot(b_q_2o, b_do, allow_tf32=False) * 0.5 - - if i_v == 0: - dq_1o += (tl.sum(b_dz[None, :] * b_q, axis=1))[None, :] - dq_2o += (tl.sum(b_dz[None, :] * b_q_2o, axis=1) * 0.5)[:, None] - - tl.store(p_dk, b_dk.to(p_dk.dtype.element_ty), boundary_check=(0, 1)) - tl.store(p_dv, b_dv.to(p_dv.dtype.element_ty), boundary_check=(0, 1)) - - -class FusedChunkBasedFunction(torch.autograd.Function): - - @staticmethod - @input_guard - @autocast_custom_fwd - def forward(ctx, q, k, v, scale=1): - B, H, T, K, V = *k.shape, v.shape[-1] - - scale = scale - BT = 16 - BK, BV = min(K, 16), min(V, 32) - BK, BV = max(BK, 16), max(BV, 16) - NK, NV = triton.cdiv(K, BK), triton.cdiv(V, BV) - - num_warps = 4 - - # the norm of o might explode, so we need to use float32 here - o = q.new_empty(NK, B, H, T, V, dtype=torch.float32) - z = q.new_empty(NK, B, H, T, dtype=torch.float32) - - grid = (NV, NK, B * H) - fused_chunk_based_fwd_kernel[grid]( - q, k, v, o, z, - scale, - T=T, B=B, H=H, K=K, V=V, BT=BT, BK=BK, BV=BV, - num_warps=num_warps, - ) - o = o.sum(0) - z = z.sum(0) - ctx.save_for_backward(q, k, v) - ctx.scale = scale - return o.to(q.dtype), z.to(z.dtype) - - @staticmethod - @input_guard - @autocast_custom_bwd - def backward(ctx, do, dz): - q, k, v = ctx.saved_tensors - B, H, T, K, V = *k.shape, v.shape[-1] - scale = ctx.scale - - BT = 16 - BK, BV = min(K, 16), min(V, 32) - BK, BV = max(BK, 16), max(BV, 16) - NK, NV = triton.cdiv(K, BK), triton.cdiv(V, BV) - num_stages = 1 - num_warps = 4 - - dq = q.new_empty(NV, B, H, T, K) - dk = q.new_empty(NV, B, H, T, K) - dv = q.new_empty(NK, B, H, T, V) - grid = (NV, NK, B * H) - - fused_chunk_based_bwd_kernel[grid]( - q, k, v, do, dz, dq, dk, dv, - scale, - T=T, B=B, H=H, K=K, V=V, BT=BT, BK=BK, BV=BV, - num_warps=num_warps, - num_stages=num_stages, - ) - dq = dq.sum(0) - dk = dk.sum(0) - dv = dv.sum(0) - return dq.to(q.dtype), dk.to(k.dtype), dv.to(v.dtype), None - - -def fused_chunk_based( - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - scale: float | None = None, - use_norm: bool = True, - head_first: bool = False, -): - assert q.shape[-1] <= 16, 'only support feature dimension up to 16.' - if scale is None: - scale = q.shape[-1] ** -0.5 - if not head_first: - q, k, v = map(lambda x: x.transpose(1, 2), (q, k, v)) - o, z = FusedChunkBasedFunction.apply(q, k, v, scale) - if use_norm: - o = o / (z[..., None] + 1e-6) - if not head_first: - o = o.transpose(1, 2) - return o.to(q.dtype) diff --git a/build/torch-xpu/ops/based/naive.py b/build/torch-xpu/ops/based/naive.py deleted file mode 100644 index 2773af0b13d8da77956f64de92d7ce82fcee7437..0000000000000000000000000000000000000000 --- a/build/torch-xpu/ops/based/naive.py +++ /dev/null @@ -1,75 +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 -from einops import rearrange - - -def naive_parallel_based( - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - scale: float | None = None, - use_norm: bool = True, -): - if scale is None: - scale = q.shape[-1] ** -0.5 - q = q * scale - attn = q @ k.transpose(-2, -1) - attn = 1 + attn + 1/2 * (attn ** 2) - attn.masked_fill_(~torch.tril(torch.ones( - q.shape[-2], q.shape[-2], dtype=torch.bool, device=q.device)), 0) - o = attn @ v - if use_norm: - z = attn.sum(-1) - return o / (z[..., None] + 1e-6) - else: - return o - - -def naive_chunk_based(q, k, v, chunk_size=256): - q = q * (q.shape[-1] ** -0.5) - # compute normalizer. - k_cumsum = torch.cumsum(k, dim=-2) - kk_cumsum = torch.cumsum(k.unsqueeze(-1) * k.unsqueeze(-2), dim=-3) - # first - z = (q * k_cumsum).sum(-1) - # second order - z += (q.unsqueeze(-1) * q.unsqueeze(-2) * kk_cumsum).sum((-1, -2)) * 0.5 - # zero-th order - z += (torch.arange(0, q.shape[-2]).to(z.device) * 1.0 + 1.0)[None, None, :] - - # compute o - # constant term - _o = v.cumsum(-2) - - q = rearrange(q, 'b h (n c) d -> b h n c d', c=chunk_size) - - k = rearrange(k, 'b h (n c) d -> b h n c d', c=chunk_size) - v = rearrange(v, 'b h (n c) d -> b h n c d', c=chunk_size) - - intra_chunk_attn = q @ k.transpose(-2, -1) - intra_chunk_attn = intra_chunk_attn + 1/2 * (intra_chunk_attn ** 2) - intra_chunk_attn.masked_fill_(~torch.tril(torch.ones(chunk_size, chunk_size, dtype=torch.bool, device=q.device)), 0) - o = intra_chunk_attn @ v - - # quadractic term - kv = torch.einsum('b h n c x, b h n c y, b h n c z -> b h n x y z', k, k, v) - kv = kv.cumsum(2) - kv = torch.cat([torch.zeros_like(kv[:, :, :1]), kv[:, :, :-1]], dim=2) - - o += 0.5 * torch.einsum('b h n x y z, b h n c x, b h n c y -> b h n c z', kv, q, q) - - # linear term - kv = torch.einsum('b h n c x, b h n c y -> b h n x y', k, v) - kv = kv.cumsum(2) - kv = torch.cat([torch.zeros_like(kv[:, :, :1]), kv[:, :, :-1]], dim=2) - o += torch.einsum('b h n x y, b h n c x -> b h n c y', kv, q) - - o = rearrange(o, 'b h n c d -> b h (n c) d') - o = o + _o - return o / (z[..., None] + 1e-6) diff --git a/build/torch-xpu/ops/based/parallel.py b/build/torch-xpu/ops/based/parallel.py deleted file mode 100644 index a693d3837d02ef7ad299ad9980c0bb47b41c6808..0000000000000000000000000000000000000000 --- a/build/torch-xpu/ops/based/parallel.py +++ /dev/null @@ -1,410 +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 -import triton.language as tl - -from ...utils import autocast_custom_bwd, autocast_custom_fwd, input_guard - -# Based: An Educational and Effective Sequence Mixer -# https://hazyresearch.stanford.edu/blog/2023-12-11-zoology2-based - - -@triton.jit(do_not_specialize=['T']) -def parallel_based_fwd_kernel( - q, - k, - v, - o, - z, - scale, - T, - B: tl.constexpr, - H: tl.constexpr, - K: tl.constexpr, - V: tl.constexpr, - BTL: tl.constexpr, - BTS: tl.constexpr, - BK: tl.constexpr, - BV: tl.constexpr, -): - # i_c: chunk index. used for sequence parallelism - i_kv, i_c, i_bh = tl.program_id(0), tl.program_id(1), tl.program_id(2) - NV = tl.cdiv(V, BV) - i_k = i_kv // (NV) - i_v = i_kv % (NV) - - p_q = tl.make_block_ptr(q + i_bh * T*K, (T, K), (K, 1), (i_c * BTL, i_k * BK), (BTL, BK), (1, 0)) - p_k = tl.make_block_ptr(k + i_bh * T*K, (K, T), (1, K), (i_k * BK, 0), (BK, BTS), (0, 1)) - p_v = tl.make_block_ptr(v + i_bh * T*V, (T, V), (V, 1), (0, i_v * BV), (BTS, BV), (1, 0)) - - # [BQ, BD] block Q, in the shared memory throughout the whole kernel - b_q = tl.load(p_q, boundary_check=(0, 1)) - b_q = (b_q * scale).to(b_q.dtype) - b_o = tl.zeros([BTL, BV], dtype=tl.float32) - b_z = tl.zeros([BTL], dtype=tl.float32) - - # Q block and K block have no overlap - # no need for mask, thereby saving flops - for _ in range(0, i_c * BTL, BTS): - # [BK, BTS] - b_k = tl.load(p_k, boundary_check=(0, 1)) - - # [BTS, BV] - b_v = tl.load(p_v, boundary_check=(0, 1)) - # [BTL, BTS] - b_s = tl.dot(b_q, (b_k), allow_tf32=False) - b_s = 1 + b_s + 0.5 * b_s * b_s - b_z += tl.sum(b_s, axis=1) - - # [BQ, BD] - b_o = b_o + tl.dot(b_s.to(b_v.dtype), b_v, allow_tf32=False) - p_k = tl.advance(p_k, (0, BTS)) - p_v = tl.advance(p_v, (BTS, 0)) - - # # rescale interchunk output - tl.debug_barrier() - o_q = tl.arange(0, BTL) - # # sync threads, easy for compiler to optimize - # tl.debug_barrier() - - o_k = tl.arange(0, BTS) - p_k = tl.make_block_ptr(k + i_bh * T*K, (K, T), (1, K), (i_k * BK, i_c * BTL), (BK, BTS), (0, 1)) - p_v = tl.make_block_ptr(v + i_bh * T*V, (T, V), (V, 1), (i_c * BTL, i_v * BV), (BTS, BV), (1, 0)) - # Q block and K block have overlap. masks required - for _ in range(i_c * BTL, (i_c + 1) * BTL, BTS): - # [BK, BTS] - b_k = tl.load(p_k, boundary_check=(0, 1)) - # [BTS, BV] - b_v = tl.load(p_v, boundary_check=(0, 1)) - # [BTL, BTS] - m_s = o_q[:, None] >= o_k[None, :] - b_s = tl.dot(b_q, b_k, allow_tf32=False) - b_s = 1 + b_s + 0.5 * b_s * b_s - b_s = tl.where(m_s, b_s, 0) - b_z += tl.sum(b_s, axis=1) - # [BTL, BV] - b_o += tl.dot(b_s.to(b_q.dtype), b_v, allow_tf32=False) - - p_k = tl.advance(p_k, (0, BTS)) - p_v = tl.advance(p_v, (BTS, 0)) - o_k += BTS - - p_o = tl.make_block_ptr(o + (i_bh + B * H * i_k) * T*V, (T, V), (V, 1), (i_c*BTL, i_v*BV), (BTL, BV), (1, 0)) - p_z = z + (i_bh + B * H * i_k) * T + i_c * BTL + tl.arange(0, BTL) - tl.store(p_o, b_o.to(p_o.dtype.element_ty), boundary_check=(0, 1)) - tl.store(p_z, b_z.to(p_z.dtype.element_ty), mask=((i_c * BTL + tl.arange(0, BTL)) < T)) - - -@triton.jit -def _parallel_based_bwd_dq( - i_bh, - i_c, - i_k, - i_v, - q, - k, - v, - do, - dz, - dq, - scale, - T, - B: tl.constexpr, - H: tl.constexpr, - BTL: tl.constexpr, - BTS: tl.constexpr, - BK: tl.constexpr, - BV: tl.constexpr, - K: tl.constexpr, - V: tl.constexpr, -): - p_do = tl.make_block_ptr(do + i_bh * T*V, (T, V), (V, 1), (i_c * BTL, i_v * BV), (BTL, BV), (1, 0)) - p_q = tl.make_block_ptr(q + (i_bh) * T*K, (T, K), (K, 1), (i_c*BTL, i_k*BK), (BTL, BK), (1, 0)) - b_q = tl.load(p_q, boundary_check=(0, 1)) - b_q = (b_q * scale).to(b_q.dtype) - - b_do = tl.load(p_do, boundary_check=(0, 1)).to(b_q.dtype) - b_dq = tl.zeros([BTL, BK], dtype=tl.float32) - p_k = tl.make_block_ptr(k + i_bh * T*K, (T, K), (K, 1), (0, i_k * BK), (BTS, BK), (1, 0)) - p_v = tl.make_block_ptr(v + i_bh * T*V, (V, T), (1, V), (i_v * BV, 0), (BV, BTS), (0, 1)) - p_dz = dz + i_bh * T + i_c * BTL + tl.arange(0, BTL) - b_dz = tl.load(p_dz, mask=(i_c * BTL + tl.arange(0, BTL)) < T) - - for _ in range(0, i_c * BTL, BTS): - # [BTS, BK] - b_k = tl.load(p_k, boundary_check=(0, 1)) - # [BV, BTS] - b_v = tl.load(p_v, boundary_check=(0, 1)) - # [BTL, BTS] - b_ds = tl.dot(b_do, b_v, allow_tf32=False) - if i_v == 0: - b_ds += b_dz[:, None] - else: - b_ds = b_ds - b_s = tl.dot(b_q, tl.trans(b_k), allow_tf32=False) - # [BQ, BD] - b_dq += tl.dot((b_ds * (1 + b_s)).to(b_v.dtype), b_k, allow_tf32=False) - p_k = tl.advance(p_k, (BTS, 0)) - p_v = tl.advance(p_v, (0, BTS)) - - b_dq *= scale - o_q = tl.arange(0, BTL) - o_k = tl.arange(0, BTS) - p_k = tl.make_block_ptr(k + i_bh * T*K, (T, K), (K, 1), (i_c * BTL, i_k * BK), (BTS, BK), (1, 0)) - p_v = tl.make_block_ptr(v + i_bh * T*V, (V, T), (1, V), (i_v * BV, i_c * BTL), (BV, BTS), (0, 1)) - # Q block and K block have overlap. masks required - for _ in range(i_c * BTL, (i_c + 1) * BTL, BTS): - # [BTS, BK] - b_k = tl.load(p_k, boundary_check=(0, 1)) - # [BV, BTS] - b_v = tl.load(p_v, boundary_check=(0, 1)) - # [BTL, BTS] - m_s = o_q[:, None] >= o_k[None, :] - b_ds = tl.dot(b_do, b_v, allow_tf32=False) - if i_v == 0: - b_ds += b_dz[:, None] - else: - b_ds = b_ds - b_ds = tl.where(m_s, b_ds, 0) * scale - b_s = tl.dot(b_q, tl.trans(b_k), allow_tf32=False) - b_s = tl.where(m_s, b_s, 0) - # [BTL, BK] - b_dq += tl.dot((b_ds + b_ds * b_s).to(b_k.dtype), b_k, allow_tf32=False) - p_k = tl.advance(p_k, (BTS, 0)) - p_v = tl.advance(p_v, (0, BTS)) - o_k += BTS - p_dq = tl.make_block_ptr(dq + (i_bh + B * H * i_v) * T*K, (T, K), (K, 1), (i_c*BTL, i_k*BK), (BTL, BK), (1, 0)) - tl.store(p_dq, b_dq.to(p_dq.dtype.element_ty), boundary_check=(0, 1)) - return - - -@triton.jit -def _parallel_based_bwd_dkv( - i_bh, - i_c, - i_k, - i_v, - q, - k, - v, - do, - dz, - dk, - dv, - scale, - T, - B: tl.constexpr, - H: tl.constexpr, - BTL: tl.constexpr, - BTS: tl.constexpr, - BK: tl.constexpr, - BV: tl.constexpr, - K: tl.constexpr, - V: tl.constexpr, -): - # compute dk dv - p_k = tl.make_block_ptr(k + i_bh * T*K, (T, K), (K, 1), (i_c * BTL, i_k * BK), (BTL, BK), (1, 0)) - p_v = tl.make_block_ptr(v + i_bh * T*V, (T, V), (V, 1), (i_c * BTL, i_v * BV), (BTL, BV), (1, 0)) - b_k, b_v = tl.load(p_k, boundary_check=(0, 1)), tl.load(p_v, boundary_check=(0, 1)) - b_dk, b_dv = tl.zeros([BTL, BK], dtype=tl.float32), tl.zeros([BTL, BV], dtype=tl.float32) - - for i in range((tl.cdiv(T, BTS) * BTS)-BTS, (i_c + 1) * BTL - BTS, -BTS): - p_q = tl.make_block_ptr(q + i_bh * T*K, (K, T), (1, K), (i_k * BK, i), (BK, BTS), (0, 1)) - p_do = tl.make_block_ptr(do + i_bh * T*V, (V, T), (1, V), (i_v * BV, i), (BV, BTS), (0, 1)) - p_dz = dz + i_bh * T + i + tl.arange(0, BTS) - b_q = tl.load(p_q, boundary_check=(0, 1)) # [BK, BTS] - b_do = tl.load(p_do, boundary_check=(0, 1)).to(b_q.dtype) # [BV, BTS] - b_dz = tl.load(p_dz, mask=(i + tl.arange(0, BTS)) < T) - b_s = tl.dot(b_k.to(b_q.dtype), b_q, allow_tf32=False) * scale # [BTL, BTS] - b_s2 = 1 + b_s + 0.5 * b_s * b_s - b_dv += tl.dot(b_s2.to(b_q.dtype), tl.trans(b_do), allow_tf32=False) - b_ds = tl.dot(b_v, b_do, allow_tf32=False) * scale - if i_v == 0: - b_ds += b_dz[None, :] * scale - else: - b_ds = b_ds - b_dk += tl.dot((b_ds + b_ds * b_s).to(b_q.dtype), tl.trans(b_q), allow_tf32=False) - - tl.debug_barrier() - o_q, o_k = tl.arange(0, BTS), tl.arange(0, BTL) - for i in range(i_c*BTL, (i_c+1)*BTL, BTS): - p_q = tl.make_block_ptr(q + i_bh * T*K, (K, T), (1, K), (i_k * BK, i), (BK, BTS), (0, 1)) - p_do = tl.make_block_ptr(do + i_bh * T*V, (V, T), (1, V), (i_v * BV, i), (BV, BTS), (0, 1)) - p_dz = dz + i_bh * T + i + tl.arange(0, BTS) - b_q = tl.load(p_q, boundary_check=(0, 1)) # [BD, BQ] - b_do = tl.load(p_do, boundary_check=(0, 1)).to(b_q.dtype) - b_dz = tl.load(p_dz, mask=(i + tl.arange(0, BTS)) < T) - # [BK, BQ] - m_s = o_k[:, None] <= o_q[None, :] - b_s = tl.dot(b_k, b_q, allow_tf32=False) * scale - b_s2 = 1 + b_s + 0.5 * b_s * b_s - b_s = tl.where(m_s, b_s, 0) - b_s2 = tl.where(m_s, b_s2, 0) - - b_ds = tl.dot(b_v, b_do, allow_tf32=False) - if i_v == 0: - b_ds += b_dz[None, :] - else: - b_ds = b_ds - b_ds = tl.where(m_s, b_ds, 0) * scale - # [BK, BD] - b_dv += tl.dot(b_s2.to(b_q.dtype), tl.trans(b_do), allow_tf32=False) - b_dk += tl.dot((b_ds + b_ds * b_s).to(b_q.dtype), tl.trans(b_q), allow_tf32=False) - o_q += BTS - - p_dk = tl.make_block_ptr(dk + (i_bh + B * H * i_v) * T*K, (T, K), (K, 1), (i_c*BTL, i_k*BK), (BTL, BK), (1, 0)) - p_dv = tl.make_block_ptr(dv + (i_bh + B * H * i_k) * T*V, (T, V), (V, 1), (i_c*BTL, i_v*BV), (BTL, BV), (1, 0)) - tl.store(p_dk, b_dk.to(p_dk.dtype.element_ty), boundary_check=(0, 1)) - tl.store(p_dv, b_dv.to(p_dv.dtype.element_ty), boundary_check=(0, 1)) - return - - -@triton.jit(do_not_specialize=['T']) -def parallel_based_bwd_kernel( - q, - k, - v, - do, - dz, - dq, - dk, - dv, - scale, - T, - B: tl.constexpr, - H: tl.constexpr, - K: tl.constexpr, - V: tl.constexpr, - BTL: tl.constexpr, - BTS: tl.constexpr, - BK: tl.constexpr, - BV: tl.constexpr, -): - i_kv, i_c, i_bh = tl.program_id(0), tl.program_id(1), tl.program_id(2) - NV = tl.cdiv(V, BV) - i_k = i_kv // (NV) - i_v = i_kv % NV - _parallel_based_bwd_dq( - i_bh, i_c, i_k, i_v, - q, k, v, do, dz, dq, - scale, T, B, H, BTL, BTS, BK, BV, K, V, - ) - tl.debug_barrier() - _parallel_based_bwd_dkv( - i_bh, i_c, i_k, i_v, - q, k, v, do, dz, dk, dv, - scale, T, B, H, BTL, BTS, BK, BV, K, V, - ) - - -class ParallelBasedFunction(torch.autograd.Function): - - @staticmethod - @input_guard - @autocast_custom_fwd - def forward(ctx, q, k, v, scale): - BTL, BTS = 128, 32 - assert BTL % BTS == 0 - # assert q.shape[-1] % 16 == 0 - BK = min(128, max(triton.next_power_of_2(k.shape[-1]), 16)) - BV = min(128, max(triton.next_power_of_2(v.shape[-1]), 16)) - B, H, T, K, V = *k.shape, v.shape[-1] - num_stages = 2 - num_warps = 4 - NK = triton.cdiv(K, BK) - NV = triton.cdiv(V, BV) - grid = (NK * NV, triton.cdiv(T, BTL), B * H) - - assert NK == 1, "will encounter some synchronization issue if not." - - o = torch.empty(NK, B, H, T, V, device=q.device) - z = torch.empty(NK, B, H, T, device=q.device) - parallel_based_fwd_kernel[grid]( - q, k, v, o, z, - scale, - B=B, - H=H, - T=T, - K=K, - V=V, - BTL=BTL, - BTS=BTS, - BK=BK, - BV=BV, - num_warps=num_warps, - num_stages=num_stages, - ) - ctx.save_for_backward(q, k, v) - ctx.scale = scale - return o.sum(0).to(q.dtype), z.sum(0).to(q.dtype) - - @staticmethod - @input_guard - @autocast_custom_bwd - def backward(ctx, do, dz): - q, k, v = ctx.saved_tensors - scale = ctx.scale - BTL, BTS = 64, 32 - assert BTL % BTS == 0 - BK = min(128, max(triton.next_power_of_2(k.shape[-1]), 16)) - BV = min(128, max(triton.next_power_of_2(v.shape[-1]), 16)) - B, H, T, K, V = *k.shape, v.shape[-1] - num_stages = 2 - num_warps = 4 - NK = triton.cdiv(K, BK) - NV = triton.cdiv(V, BV) - grid = (NK * NV, triton.cdiv(T, BTL), B * H) - - assert NK == 1, "will encounter some synchronization issue if not" - - dq = torch.empty(NV, B, H, T, K, dtype=q.dtype, device=q.device) - dk = torch.empty(NV, B, H, T, K, dtype=q.dtype, device=q.device) - dv = torch.empty(NK, B, H, T, V, dtype=q.dtype, device=q.device) - - parallel_based_bwd_kernel[grid]( - q, k, v, do, dz, dq, dk, dv, - scale, - B=B, - H=H, - T=T, - K=K, - V=V, - BTL=BTL, - BTS=BTS, - BK=BK, - BV=BV, - num_warps=num_warps, - num_stages=num_stages, - ) - - return dq.sum(0).to(q.dtype), dk.sum(0).to(k.dtype), dv.sum(0).to(v.dtype), None - - -triton_parallel_based = ParallelBasedFunction.apply - - -def parallel_based( - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - scale: float | None = None, - use_norm: bool = True, - head_first: bool = False, -): - assert q.shape[-1] <= 128, "only support feature dim up to 128" - if scale is None: - scale = q.shape[-1] ** -0.5 - if not head_first: - q, k, v = map(lambda x: x.transpose(1, 2), (q, k, v)) - o, z = triton_parallel_based(q, k, v, scale) - if use_norm: - o = o / (z[..., None] + 1e-6) - if not head_first: - o = o.transpose(1, 2) - return o.to(q.dtype) diff --git a/build/torch-xpu/ops/comba/__init__.py b/build/torch-xpu/ops/comba/__init__.py deleted file mode 100644 index 29b92ee6be00526affce32d92813930fb00fd017..0000000000000000000000000000000000000000 --- a/build/torch-xpu/ops/comba/__init__.py +++ /dev/null @@ -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 - -from .chunk import chunk_comba -from .fused_recurrent import fused_recurrent_comba -from .naive import naive_chunk_comba, naive_recurrent_comba - -__all__ = [ - "chunk_comba", - "fused_recurrent_comba", - "naive_chunk_comba", - "naive_recurrent_comba", -] diff --git a/build/torch-xpu/ops/comba/chunk.py b/build/torch-xpu/ops/comba/chunk.py deleted file mode 100644 index e2497babafeb13b7a751e3462b794a034356f78f..0000000000000000000000000000000000000000 --- a/build/torch-xpu/ops/comba/chunk.py +++ /dev/null @@ -1,403 +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 - -from ...modules.l2norm import l2norm_bwd, l2norm_fwd -from ...ops.comba.utils import chunk_comba_cumsum_scalar_bwd, chunk_comba_cumsum_scalar_fwd -from ...ops.comba.wy_fast import chunk_scaled_dot_comba_pkt_fwd, prepare_wy_repr_bwd, recompute_w_u_fwd -from ...ops.common.chunk_delta_h import chunk_gated_delta_rule_bwd_dhu, chunk_gated_delta_rule_fwd_h -from ...ops.common.chunk_o import chunk_bwd_dqkwg, chunk_bwd_dv_local, chunk_fwd_o -from ...ops.utils import chunk_local_cumsum, prepare_chunk_indices, solve_tril -from ...ops.utils.constant import RCP_LN2 -from ...utils import autocast_custom_bwd, autocast_custom_fwd, input_guard - - -def chunk_comba_fwd( - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - p: torch.Tensor, - g: torch.Tensor, - beta: torch.Tensor, - scale: float, - initial_state: torch.Tensor, - output_final_state: bool, - cu_seqlens: torch.LongTensor | None = None, - chunk_indices: torch.LongTensor | None = None, -): - g0, g = chunk_comba_cumsum_scalar_fwd( - g, - chunk_size=64, - cu_seqlens=cu_seqlens, - chunk_indices=chunk_indices, - scale=RCP_LN2, - ) - # obtain WY representation. u is actually the new v. - A = chunk_scaled_dot_comba_pkt_fwd( - k=k, - p=p, - beta=beta, - g0=g0, - g=g, - cu_seqlens=cu_seqlens, - output_dtype=torch.float32, - chunk_indices=chunk_indices, - ) - A = solve_tril( - A=A, - cu_seqlens=cu_seqlens, - chunk_indices=chunk_indices, - output_dtype=k.dtype, - ) - w, u = recompute_w_u_fwd( - k=p, - v=v, - beta=beta, - A=A, - g_cumsum=g0, - cu_seqlens=cu_seqlens, - chunk_indices=chunk_indices, - ) - h, v_new, final_state = chunk_gated_delta_rule_fwd_h( - k=k, - w=w, - u=u, - g=g, - initial_state=initial_state, - output_final_state=output_final_state, - cu_seqlens=cu_seqlens, - chunk_indices=chunk_indices, - ) - o = chunk_fwd_o( - q=q, - k=k, - v=v_new, - h=h, - g=g, - scale=scale, - cu_seqlens=cu_seqlens, - chunk_indices=chunk_indices, - ) - return g0, g, o, A, final_state - - -def chunk_comba_bwd( - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - p: torch.Tensor, - g0: torch.Tensor, - g: torch.Tensor, - beta: torch.Tensor, - A: torch.Tensor, - scale: float, - initial_state: torch.Tensor, - do: torch.Tensor, - dht: torch.Tensor, - cu_seqlens: torch.LongTensor | None = None, - chunk_indices: torch.LongTensor | None = None, -): - w, u = recompute_w_u_fwd( - k=p, - v=v, - beta=beta, - A=A, - g_cumsum=g0, - cu_seqlens=cu_seqlens, - chunk_indices=chunk_indices, - ) - h, v_new, _ = chunk_gated_delta_rule_fwd_h( - k=k, - w=w, - u=u, - g=g, - initial_state=initial_state, - output_final_state=False, - cu_seqlens=cu_seqlens, - chunk_indices=chunk_indices, - ) - dv = chunk_bwd_dv_local( - q=q, - k=k, - g=g, - do=do, - scale=scale, - cu_seqlens=cu_seqlens, - chunk_indices=chunk_indices, - ) - dh, dh0, dv = chunk_gated_delta_rule_bwd_dhu( - q=q, - k=k, - w=w, - g=g, - h0=initial_state, - dht=dht, - do=do, - dv=dv, - scale=scale, - cu_seqlens=cu_seqlens, - chunk_indices=chunk_indices, - ) - dq, dk, dw, dg = chunk_bwd_dqkwg( - q=q, - k=k, - v=v_new, - w=w, - g=g, - h=h, - dv=dv, - do=do, - dh=dh, - scale=scale, - cu_seqlens=cu_seqlens, - chunk_indices=chunk_indices, - ) - dk2, dv, dp, db, dg0, dg2 = prepare_wy_repr_bwd( - k=k, - v=v, - p=p, - beta=beta, - g0=g0, - g=g, - A=A, - dw=dw, - du=dv, - cu_seqlens=cu_seqlens, - chunk_indices=chunk_indices, - ) - dk.add_(dk2) - dg.add_(dg2) - assert dg.dtype == torch.float32, "dg should be fp32" - dg = chunk_local_cumsum(dg, chunk_size=64, reverse=True, cu_seqlens=cu_seqlens, chunk_indices=chunk_indices) - # dg0 = d(g_cumsum - g) - dg += chunk_comba_cumsum_scalar_bwd(dg0, chunk_size=64, cu_seqlens=cu_seqlens, chunk_indices=chunk_indices) - return dq, dk, dv, dp, db, dg, dh0 - - -class ChunkCombaFunction(torch.autograd.Function): - - @staticmethod - @input_guard - @autocast_custom_fwd - def forward( - ctx, - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - p: torch.Tensor, - g: torch.Tensor, - beta: torch.Tensor, - scale: float, - initial_state: torch.Tensor, - output_final_state: bool, - use_qk_l2norm_in_kernel: bool = False, - cu_seqlens: torch.LongTensor | None = None, - cu_seqlens_cpu: torch.LongTensor | None = None, - ): - if use_qk_l2norm_in_kernel: - q, q_rstd = l2norm_fwd(q) - k, k_rstd = l2norm_fwd(k) - p, p_rstd = l2norm_fwd(p) - else: - q_rstd, k_rstd, p_rstd = None, None, None - - chunk_indices = None - if cu_seqlens is not None: - chunk_indices = prepare_chunk_indices(cu_seqlens, 64, cu_seqlens_cpu=cu_seqlens_cpu) - - g0, g, o, A, final_state = chunk_comba_fwd( - q=q, - k=k, - v=v, - p=p, - g=g, - beta=beta, - scale=scale, - initial_state=initial_state, - output_final_state=output_final_state, - cu_seqlens=cu_seqlens, - chunk_indices=chunk_indices, - ) - ctx.save_for_backward( - q, - q_rstd, - k, - k_rstd, - p, - p_rstd, - v, - g0, - g, - beta, - A, - initial_state, - cu_seqlens, - chunk_indices, - ) - ctx.scale = scale - ctx.use_qk_l2norm_in_kernel = use_qk_l2norm_in_kernel - return o.to(q.dtype), final_state - - @staticmethod - @input_guard - @autocast_custom_bwd - def backward( - ctx, - do: torch.Tensor, - dht: torch.Tensor, - ): - ( - q, - q_rstd, - k, - k_rstd, - p, - p_rstd, - v, - g0, - g, - beta, - A, - initial_state, - cu_seqlens, - chunk_indices, - ) = ctx.saved_tensors - dq, dk, dv, dp, db, dg, dh0 = chunk_comba_bwd( - q=q, - k=k, - v=v, - p=p, - g0=g0, - g=g, - beta=beta, - A=A, - scale=ctx.scale, - initial_state=initial_state, - do=do, - dht=dht, - cu_seqlens=cu_seqlens, - chunk_indices=chunk_indices, - ) - if ctx.use_qk_l2norm_in_kernel: - dq = l2norm_bwd(q, q_rstd, dq) - dk = l2norm_bwd(k, k_rstd, dk) - dp = l2norm_bwd(p, p_rstd, dp) - return dq.to(q), dk.to(k), dv.to(v), dp.to(p), dg.to(g), db.to(beta), None, dh0, None, None, None, None - - -@torch.compiler.disable -def chunk_comba( - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - p: torch.Tensor, - g: torch.Tensor, - beta: torch.Tensor = None, - scale: float = None, - initial_state: torch.Tensor = None, - output_final_state: bool = False, - use_qk_l2norm_in_kernel: bool = False, - cu_seqlens: torch.LongTensor | None = None, - cu_seqlens_cpu: torch.LongTensor | None = None, -): - r""" - Args: - q (torch.Tensor): - queries of shape `[B, T, H, K]`. - k (torch.Tensor): - keys of shape `[B, T, H, K]`. - v (torch.Tensor): - values of shape `[B, T, H, V]`. - p (torch.Tensor): - auxiliary keys of shape `[B, T, H, K]`. - g (torch.Tensor): - (forget) gating tensor (in log space!) of shape `[B, T, H]`. - beta (torch.Tensor): - betas of shape `[B, T, H]`. - scale (Optional[int]): - Scale factor for the RetNet attention scores. - If not provided, it will default to `1 / sqrt(K)`. Default: `None`. - initial_state (Optional[torch.Tensor]): - Initial state of shape `[N, H, K, V]` for `N` input sequences. - For equal-length input sequences, `N` equals the batch size `B`. - Default: `None`. - output_final_state (Optional[bool]): - Whether to output the final state of shape `[N, H, K, V]`. Default: `False`. - use_qk_l2norm_in_kernel (bool): - Whether to apply L2norm to the q/k tensor internally. Default: `False`. - cu_seqlens (torch.LongTensor): - Cumulative sequence lengths of shape `[N+1]` used for variable-length training, - consistent with the FlashAttention API. - - Returns: - o (torch.Tensor): - Outputs of shape `[B, T, H, V]`. - final_state (torch.Tensor): - Final state of shape `[N, H, K, V]` if `output_final_state=True` else `None`. - - Examples:: - >>> import torch - >>> import torch.nn.functional as F - >>> from einops import rearrange - >>> from ...ops.comba import chunk_comba - # inputs with equal lengths - >>> B, T, H, K, V = 4, 2048, 4, 512, 512 - >>> q = torch.randn(B, T, H, K, dtype=torch.bfloat16, device='cuda') - >>> k = F.normalize(torch.randn(B, T, H, K, dtype=torch.bfloat16, device='cuda'), p=2, dim=-1) - >>> v = torch.randn(B, T, H, V, dtype=torch.bfloat16, device='cuda') - >>> b = torch.rand(H, dtype=torch.bfloat16, device='cuda').sigmoid() - >>> p = k * b[:, None] - >>> beta = torch.rand(B, T, H, dtype=torch.bfloat16, device='cuda').sigmoid() - >>> g = F.logsigmoid(torch.rand(B, T, H, dtype=torch.bfloat16, device='cuda')) - >>> h0 = torch.randn(B, H, K, V, dtype=torch.bfloat16, device='cuda') - >>> o, ht = chunk_comba( - q, k, v, p, g, beta, - initial_state=h0, - output_final_state=True - ) - # for variable-length inputs, the batch size `B` is expected to be 1 and `cu_seqlens` is required - >>> q, k, v, beta, g = map(lambda x: rearrange(x, 'b t ... -> 1 (b t) ...'), (q, k, v, beta, g)) - # for a batch with 4 sequences, `cu_seqlens` with 5 start/end positions are expected - >>> cu_seqlens = q.new_tensor([0, 2048, 4096, 6144, 8192], dtype=torch.long) - >>> o_var, ht_var = chunk_comba( - q, k, v, p, g, beta, - initial_state=h0, - output_final_state=True, - cu_seqlens=cu_seqlens - ) - """ - if p is None: - p = k - if cu_seqlens is not None: - if q.shape[0] != 1: - raise ValueError( - f"The batch size is expected to be 1 rather than {q.shape[0]} when using `cu_seqlens`." - f"Please flatten variable-length inputs before processing.", - ) - if initial_state is not None and initial_state.shape[0] != len(cu_seqlens) - 1: - raise ValueError( - f"The number of initial states is expected to be equal to the number of input sequences, " - f"i.e., {len(cu_seqlens) - 1} rather than {initial_state.shape[0]}.", - ) - if scale is None: - scale = k.shape[-1] ** -0.5 - o, final_state = ChunkCombaFunction.apply( - q, - k, - v, - p, - g, - beta, - scale, - initial_state, - output_final_state, - use_qk_l2norm_in_kernel, - cu_seqlens, - cu_seqlens_cpu, - ) - return o, final_state diff --git a/build/torch-xpu/ops/comba/fused_recurrent.py b/build/torch-xpu/ops/comba/fused_recurrent.py deleted file mode 100644 index bd773d7bf6d5b83c949898c1a7445f476aab3e52..0000000000000000000000000000000000000000 --- a/build/torch-xpu/ops/comba/fused_recurrent.py +++ /dev/null @@ -1,334 +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 -import triton.language as tl - -from ...ops.utils.op import exp -from ...utils import input_guard - - -@triton.heuristics({ - 'USE_INITIAL_STATE': lambda args: args['h0'] is not None, - 'STORE_FINAL_STATE': lambda args: args['ht'] is not None, - 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None, -}) -@triton.jit(do_not_specialize=['T']) -def fused_recurrent_comba_fwd_kernel( - q, - k, - p, - v, - g, - beta, - o, - h0, - ht, - cu_seqlens, - scale, - T, - B: tl.constexpr, - H: tl.constexpr, - HV: tl.constexpr, - K: tl.constexpr, - V: tl.constexpr, - BK: tl.constexpr, - BV: tl.constexpr, - USE_INITIAL_STATE: tl.constexpr, # whether to use initial state - STORE_FINAL_STATE: tl.constexpr, # whether to store final state - IS_BETA_HEADWISE: tl.constexpr, # whether beta is headwise vector or scalar, - USE_QK_L2NORM_IN_KERNEL: tl.constexpr, - IS_VARLEN: tl.constexpr, -): - i_k, i_v, i_nh = tl.program_id(0), tl.program_id(1), tl.program_id(2) - i_n, i_hv = i_nh // HV, i_nh % HV - i_h = i_hv // (HV // H) - if IS_VARLEN: - bos, eos = tl.load(cu_seqlens + i_n).to(tl.int64), tl.load(cu_seqlens + i_n + 1).to(tl.int64) - all = T - T = eos - bos - else: - bos, eos = i_n * T, i_n * T + T - all = B * T - o_k = i_k * BK + tl.arange(0, BK) - o_v = i_v * BV + tl.arange(0, BV) - - p_q = q + (bos * H + i_h) * K + o_k - p_k = k + (bos * H + i_h) * K + o_k - p_v = v + (bos * HV + i_hv) * V + o_v - p_p = p + (bos * H + i_h) * K + o_k - if IS_BETA_HEADWISE: - p_beta = beta + (bos * HV + i_hv) * V + o_v - else: - p_beta = beta + bos * HV + i_hv - p_g = g + bos * HV + i_hv - p_o = o + ((i_k * all + bos) * HV + i_hv) * V + o_v - - mask_k = o_k < K - mask_v = o_v < V - mask_h = mask_k[:, None] & mask_v[None, :] - - b_h = tl.zeros([BK, BV], dtype=tl.float32) - if USE_INITIAL_STATE: - p_h0 = h0 + i_nh * K*V + o_k[:, None] * V + o_v[None, :] - b_h += tl.load(p_h0, mask=mask_h, other=0).to(tl.float32) - - for _ in range(0, T): - b_q = tl.load(p_q, mask=mask_k, other=0).to(tl.float32) - b_k = tl.load(p_k, mask=mask_k, other=0).to(tl.float32) - b_v = tl.load(p_v, mask=mask_v, other=0).to(tl.float32) - b_p = tl.load(p_p, mask=mask_k, other=0).to(tl.float32) - b_g = tl.load(p_g).to(tl.float32) - - if USE_QK_L2NORM_IN_KERNEL: - b_q = b_q / tl.sqrt(tl.sum(b_q * b_q) + 1e-6) - b_k = b_k / tl.sqrt(tl.sum(b_k * b_k) + 1e-6) - b_p = b_p / tl.sqrt(tl.sum(b_p * b_p) + 1e-6) - b_q = b_q * scale - # [BV] - b_v -= tl.sum(b_h * b_p[:, None], 0) - # [BK, BV] - b_h *= exp(b_g) - if IS_BETA_HEADWISE: - b_beta = tl.load(p_beta, mask=mask_v, other=0).to(tl.float32) - else: - b_beta = tl.load(p_beta).to(tl.float32) - b_v *= b_beta - # [BK, BV] - b_h += b_k[:, None] * b_v[None, :] - # [BV] - b_o = tl.sum(b_h * b_q[:, None], 0) - tl.store(p_o, b_o.to(p_o.dtype.element_ty), mask=mask_v) - - p_q += H*K - p_k += H*K - p_o += HV*V - p_v += HV*V - p_p += H*K - p_g += HV - p_beta += HV * (V if IS_BETA_HEADWISE else 1) - - if STORE_FINAL_STATE: - p_ht = ht + i_nh * K*V + o_k[:, None] * V + o_v[None, :] - tl.store(p_ht, b_h.to(p_ht.dtype.element_ty), mask=mask_h) - - -def fused_recurrent_comba_fwd( - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - p: torch.Tensor, - g: torch.Tensor, - beta: torch.Tensor, - scale: float, - initial_state: torch.Tensor, - output_final_state: bool, - use_qk_l2norm_in_kernel: bool = False, - cu_seqlens: torch.LongTensor | None = None, -) -> tuple[torch.Tensor, torch.Tensor]: - B, T, H, K, V = *k.shape, v.shape[-1] - HV = v.shape[2] - N = B if cu_seqlens is None else len(cu_seqlens) - 1 - BK, BV = triton.next_power_of_2(K), min(triton.next_power_of_2(V), 8) - NK, NV = triton.cdiv(K, BK), triton.cdiv(V, BV) - assert NK == 1, "NK > 1 is not supported yet" - num_stages = 3 - num_warps = 1 - - o = q.new_empty(NK, *v.shape) - if output_final_state: - final_state = q.new_empty(N, HV, K, V, dtype=torch.float32) - else: - final_state = None - - grid = (NK, NV, N * HV) - fused_recurrent_comba_fwd_kernel[grid]( - q=q, - k=k, - p=p, - v=v, - g=g, - beta=beta, - o=o, - h0=initial_state, - ht=final_state, - cu_seqlens=cu_seqlens, - scale=scale, - T=T, - B=B, - H=H, - HV=HV, - K=K, - V=V, - BK=BK, - BV=BV, - IS_BETA_HEADWISE=beta.ndim == v.ndim, - USE_QK_L2NORM_IN_KERNEL=use_qk_l2norm_in_kernel, - num_warps=num_warps, - num_stages=num_stages, - ) - o = o.squeeze(0) - return o, final_state - - -class FusedRecurrentCombaFunction(torch.autograd.Function): - - @staticmethod - @input_guard - def forward( - ctx, - q: torch.Tensor, - k: torch.Tensor, - p: torch.Tensor, - v: torch.Tensor, - g: torch.Tensor, - beta: torch.Tensor, - scale: float, - initial_state: torch.Tensor, - output_final_state: bool, - use_qk_l2norm_in_kernel: bool = False, - cu_seqlens: torch.LongTensor | None = None, - ): - o, final_state = fused_recurrent_comba_fwd( - q=q, - k=k, - p=p, - v=v, - g=g, - beta=beta, - scale=scale, - initial_state=initial_state, - output_final_state=output_final_state, - use_qk_l2norm_in_kernel=use_qk_l2norm_in_kernel, - cu_seqlens=cu_seqlens, - ) - - return o, final_state - - @staticmethod - @input_guard - def backward(ctx, do, dht): - raise NotImplementedError( - "Backward pass is not implemented yet and we do not have plans to implement it " - "because we haven't figured out how to compute dg without materializing the full " - "hidden states for all time steps.", - ) - - -def fused_recurrent_comba( - q: torch.Tensor, - k: torch.Tensor, - p: torch.Tensor, - v: torch.Tensor, - g: torch.Tensor, - beta: torch.Tensor = None, - scale: float = None, - initial_state: torch.Tensor = None, - output_final_state: bool = False, - use_qk_l2norm_in_kernel: bool = False, - cu_seqlens: torch.LongTensor | None = None, -) -> tuple[torch.Tensor, torch.Tensor]: - r""" - Args: - q (torch.Tensor): - queries of shape `[B, T, H, K]`. - k (torch.Tensor): - keys of shape `[B, T, H, K]`. - p (torch.Tensor): - auxiliary keys of shape `[B, T, H, K]`. - v (torch.Tensor): - values of shape `[B, T, HV, V]`. - GVA is applied if `HV > H`. - g (torch.Tensor): - g (decays) of shape `[B, T, HV]`. - beta (torch.Tensor): - betas of shape `[B, T, HV]`. - scale (Optional[int]): - Scale factor for the RetNet attention scores. - If not provided, it will default to `1 / sqrt(K)`. Default: `None`. - initial_state (Optional[torch.Tensor]): - Initial state of shape `[N, HV, K, V]` for `N` input sequences. - For equal-length input sequences, `N` equals the batch size `B`. - Default: `None`. - output_final_state (Optional[bool]): - Whether to output the final state of shape `[N, HV, K, V]`. Default: `False`. - use_qk_l2norm_in_kernel (Optional[bool]): - Whether to use qk l2norm within the kernel for saving GPU memory. - Default: `False`. - cu_seqlens (torch.LongTensor): - Cumulative sequence lengths of shape `[N+1]` used for variable-length training, - consistent with the FlashAttention API. - - Returns: - o (torch.Tensor): - Outputs of shape `[B, T, HV, V]`. - final_state (torch.Tensor): - Final state of shape `[N, HV, K, V]` if `output_final_state=True` else `None`. - - Examples:: - >>> import torch - >>> import torch.nn.functional as F - >>> from einops import rearrange - >>> from ...ops.comba import fused_recurrent_comba - # inputs with equal lengths - >>> B, T, H, HV, K, V = 4, 2048, 4, 8, 512, 512 - >>> q = torch.randn(B, T, H, K, device='cuda') - >>> k = F.normalize(torch.randn(B, T, H, K, device='cuda'), p=2, dim=-1) - >>> v = torch.randn(B, T, HV, V, device='cuda') - >>> b = torch.rand(H, dtype=torch.bfloat16, device='cuda').sigmoid() - >>> p = k * b[:, None] - >>> g = F.logsigmoid(torch.rand(B, T, HV, device='cuda')) - >>> beta = torch.rand(B, T, HV, device='cuda').sigmoid() - >>> h0 = torch.randn(B, HV, K, V, device='cuda') - >>> o, ht = fused_recurrent_comba( - q, k, v, p, g, beta, - initial_state=h0, - output_final_state=True - ) - # for variable-length inputs, the batch size `B` is expected to be 1 and `cu_seqlens` is required - >>> q, k, v, p, g, beta = map(lambda x: rearrange(x, 'b t ... -> 1 (b t) ...'), (q, k, v, p, g, beta)) - # for a batch with 4 sequences, `cu_seqlens` with 5 start/end positions are expected - >>> cu_seqlens = q.new_tensor([0, 2048, 4096, 6144, 8192], dtype=torch.long) - >>> o_var, ht_var = fused_recurrent_comba( - q, k, p, v, g, beta, - initial_state=h0, - output_final_state=True, - cu_seqlens=cu_seqlens - ) - """ - if cu_seqlens is not None: - if q.shape[0] != 1: - raise ValueError( - f"The batch size is expected to be 1 rather than {q.shape[0]} when using `cu_seqlens`." - f"Please flatten variable-length inputs before processing.", - ) - if initial_state is not None and initial_state.shape[0] != len(cu_seqlens) - 1: - raise ValueError( - f"The number of initial states is expected to be equal to the number of input sequences, " - f"i.e., {len(cu_seqlens) - 1} rather than {initial_state.shape[0]}.", - ) - if scale is None: - scale = k.shape[-1] ** -0.5 - if beta is None: - beta = torch.ones_like(q[..., 0]) - if p is None: - p = k - o, final_state = FusedRecurrentCombaFunction.apply( - q, - k, - p, - v, - g, - beta, - scale, - initial_state, - output_final_state, - use_qk_l2norm_in_kernel, - cu_seqlens, - ) - return o, final_state diff --git a/build/torch-xpu/ops/comba/naive.py b/build/torch-xpu/ops/comba/naive.py deleted file mode 100644 index 3f5072904e946d3eb00bb63756dda6aff396d929..0000000000000000000000000000000000000000 --- a/build/torch-xpu/ops/comba/naive.py +++ /dev/null @@ -1,171 +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 torch.nn.functional as F -from einops import rearrange - - -def naive_recurrent_comba( - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - p: torch.Tensor, - beta: torch.Tensor, - g: torch.Tensor, - scale: float = None, - initial_state: torch.Tensor = None, - output_final_state: bool = False, -): - """ - Reference PyTorch implementation of recurrent COMBA. - - Args: - q: [B, T, H, K] - k: [B, T, H, K] - v: [B, T, H, V] - p: [B, T, H, K] - beta: [B, T, H] - g: [B, T, H] - scale: float, optional - initial_state: [B, H, K, V], optional - output_final_state: bool - - Returns: - o: [B, T, H, V] - final_state: [B, H, K, V] if output_final_state else None - """ - q, k, v, p, beta, g = map(lambda x: x.transpose(1, 2).contiguous().to(torch.float32), [q, k, v, p, beta, g]) - B, H, T, K, V = *k.shape, v.shape[-1] - o = torch.zeros(B, H, T, V).to(v) - h = torch.zeros(B, H, K, V).to(v) - if initial_state is not None: - h = initial_state.to(torch.float32) - if scale is None: - scale = 1 / (q.shape[-1] ** 0.5) - q = q * scale - - for i in range(T): - b_q = q[:, :, i] - b_k = k[:, :, i] - b_v = v[:, :, i].clone() - b_p = p[:, :, i] - h = h.clone() * g[:, :, i].exp()[..., None, None] - b_beta = beta[:, :, i] - b_v = b_v - (h.clone() * b_p[..., None]).sum(-2) - b_v = b_v * b_beta[..., None] - h = h.clone() + b_k.unsqueeze(-1) * b_v.unsqueeze(-2) - o[:, :, i] = torch.einsum('bhd,bhdm->bhm', b_q, h) - - if not output_final_state: - h = None - o = o.transpose(1, 2).contiguous() - return o, h - - -def naive_chunk_comba( - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - p: torch.Tensor, - g: torch.Tensor, - beta: torch.Tensor, - chunk_size: int = 64, - scale: float = None, - initial_state: torch.Tensor = None, - output_final_state: bool = False, -): - """ - Reference PyTorch implementation of chunk COMBA. - - Args: - q: [B, T, H, K] - k: [B, T, H, K] - v: [B, T, H, V] - p: [B, T, H, K] - g: [B, T, H] - beta: [B, T, H] - chunk_size: int - scale: float, optional - initial_state: [B, H, K, V], optional - output_final_state: bool - - Returns: - o: [B, T, H, V] - final_state: [B, H, K, V] if output_final_state else None - """ - BT = chunk_size - if scale is None: - scale = 1 / (q.shape[-1] ** 0.5) - - q, k, v, p, beta, g = map(lambda x: x.transpose(1, 2).contiguous().to(torch.float32), [q, k, v, p, beta, g]) - - T = q.shape[-2] - pad_len = (BT - (T % BT)) % BT - if pad_len > 0: - q = F.pad(q, (0, 0, 0, pad_len)) - k = F.pad(k, (0, 0, 0, pad_len)) - v = F.pad(v, (0, 0, 0, pad_len)) - p = F.pad(p, (0, 0, 0, pad_len)) - beta = F.pad(beta, (0, pad_len)) - g = F.pad(g, (0, pad_len)) - - decay = g - chunk_size = BT - b, h, l, d_k = q.shape - d_v = v.shape[-1] - q = q * scale - v = v * beta[..., None] - p_beta = p * beta[..., None] - assert l % chunk_size == 0 - - # note that diagonal is masked. - mask = torch.triu(torch.ones(chunk_size, chunk_size, dtype=torch.bool, device=q.device), diagonal=0) - q, k, v, p_beta, decay, g = map( - lambda x: rearrange(x, 'b h (n c) d -> b h n c d', c=chunk_size), - [q, k, v, p_beta, decay.unsqueeze(-1), g.unsqueeze(-1)], - ) - decay = decay.squeeze(-1).cumsum(-1) # [B, H, n, c] - decay_0 = decay - g.squeeze(-1) # [B, H, n, c] - L_mask = ((decay.unsqueeze(-1) - decay.unsqueeze(-2)).tril().exp().float()).tril() - L_mask_0 = ((decay_0.unsqueeze(-1) - decay.unsqueeze(-2)).tril().exp().float()).tril() - - # [B, H, n, c, d] @ [B, H, n, d, c] -> [B, H, n, c, c] - attn = -((p_beta @ k.transpose(-1, -2)) * L_mask_0).masked_fill(mask, 0) - for i in range(1, chunk_size): - attn[..., i, :i] = attn[..., i, :i].clone() + (attn[..., i, :i, None].clone() * attn[..., :i, :i].clone()).sum(-2) - attn = attn + torch.eye(chunk_size, dtype=torch.float, device=q.device) - - # for U - k_cumsum = attn @ v - # for W - k_cumdecay = attn @ (p_beta * decay_0[..., None].exp()) - v = k_cumsum - - S = k.new_zeros(b, h, d_k, d_v) - if initial_state is not None: - S = initial_state.to(torch.float32) - - o = torch.zeros_like(v) - mask = torch.triu(torch.ones(chunk_size, chunk_size, dtype=torch.bool, device=q.device), diagonal=1) - for i in range(0, l // chunk_size): - q_i, k_i, v_i = q[:, :, i], k[:, :, i], v[:, :, i] - attn = (q_i @ k_i.transpose(-1, -2) * L_mask[:, :, i]).masked_fill_(mask, 0) - v_prime = k_cumdecay[:, :, i] @ S - v_new = v_i - v_prime - o_inter = (q_i * decay[:, :, i, :, None].exp()) @ S - o[:, :, i] = o_inter + attn @ v_new - S = S * decay[:, :, i, -1, None, None].exp() + (k_i * (decay[:, :, i, -1, None] - decay[:, :, i]).exp() - [..., None]).transpose(-1, -2) @ v_new - if not output_final_state: - S = None - - # unpad - o = rearrange(o, 'b h n c d -> b h (n c) d') - o = o[:, :, :T] - o = o.transpose(1, 2) - return o, S diff --git a/build/torch-xpu/ops/comba/utils.py b/build/torch-xpu/ops/comba/utils.py deleted file mode 100644 index 2126de39cb9e272df0ed8891b2c3c34ef5900368..0000000000000000000000000000000000000000 --- a/build/torch-xpu/ops/comba/utils.py +++ /dev/null @@ -1,170 +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 -import triton.language as tl - -from ...ops.utils.index import prepare_chunk_indices -from ...utils import autotune_cache_kwargs - - -@triton.heuristics({ - 'HAS_SCALE': lambda args: args['scale'] is not None, - 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None, -}) -@triton.autotune( - configs=[ - triton.Config({}, num_warps=num_warps) - for num_warps in [1, 2, 4, 8] - ], - key=['B', 'H', 'BT', 'IS_VARLEN'], - **autotune_cache_kwargs, -) -@triton.jit(do_not_specialize=['T']) -def chunk_comba_cumsum_scalar_fwd_kernel( - g, - g0, - g1, - scale, - cu_seqlens, - chunk_indices, - T, - B: tl.constexpr, - H: tl.constexpr, - BT: tl.constexpr, - HAS_SCALE: tl.constexpr, - IS_VARLEN: tl.constexpr, -): - i_t, i_bh = tl.program_id(0), tl.program_id(1) - i_b, i_h = i_bh // H, i_bh % H - 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.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) - T = eos - bos - else: - bos, eos = i_b * T, i_b * T + T - - p_g = tl.make_block_ptr(g + bos*H + i_h, (T,), (H,), (i_t * BT,), (BT,), (0,)) - p_g0 = tl.make_block_ptr(g0 + bos*H + i_h, (T,), (H,), (i_t * BT,), (BT,), (0,)) - p_g1 = tl.make_block_ptr(g1 + bos*H + i_h, (T,), (H,), (i_t * BT,), (BT,), (0,)) - # [BT] - b_g = tl.load(p_g, boundary_check=(0,)).to(tl.float32) - if HAS_SCALE: - b_g = b_g * scale - b_g1 = tl.cumsum(b_g, axis=0) - b_g0 = b_g1 - b_g - tl.store(p_g0, b_g0.to(p_g0.dtype.element_ty), boundary_check=(0,)) - tl.store(p_g1, b_g1.to(p_g1.dtype.element_ty), boundary_check=(0,)) - - -def chunk_comba_cumsum_scalar_fwd( - g: torch.Tensor, - chunk_size: int, - cu_seqlens: torch.Tensor | None = None, - output_dtype: torch.dtype | None = torch.float, - chunk_indices: torch.LongTensor | None = None, - scale: float | None = None, -) -> torch.Tensor: - B, T, H = g.shape - assert chunk_size == 2**(chunk_size.bit_length()-1), "chunk_size must be a power of 2" - BT = chunk_size - if chunk_indices is None and cu_seqlens is not None: - chunk_indices = prepare_chunk_indices(cu_seqlens, BT) - NT = triton.cdiv(T, BT) if cu_seqlens is None else len(chunk_indices) - g0, g1 = torch.empty_like(g, dtype=output_dtype or g.dtype), torch.empty_like(g, dtype=output_dtype or g.dtype) - grid = (NT, B * H) - chunk_comba_cumsum_scalar_fwd_kernel[grid]( - g, - g0, - g1, - scale, - cu_seqlens, - chunk_indices, - T=T, - B=B, - H=H, - BT=BT, - ) - return g0, g1 - - -@triton.heuristics({ - 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None, -}) -@triton.autotune( - configs=[ - triton.Config({}, num_warps=num_warps) - for num_warps in [1, 2, 4, 8] - ], - key=['B', 'H', 'BT', 'IS_VARLEN'], - **autotune_cache_kwargs, -) -@triton.jit(do_not_specialize=['T']) -def chunk_comba_cumsum_scalar_bwd_kernel( - dg0, - dgr, - cu_seqlens, - chunk_indices, - T, - B: tl.constexpr, - H: tl.constexpr, - BT: tl.constexpr, - IS_VARLEN: tl.constexpr, -): - i_t, i_bh = tl.program_id(0), tl.program_id(1) - i_b, i_h = i_bh // H, i_bh % H - 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.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) - T = eos - bos - else: - bos, eos = i_b * T, i_b * T + T - - p_dg0 = tl.make_block_ptr(dg0 + bos*H + i_h, (T,), (H,), (i_t * BT,), (BT,), (0,)) - p_dgr = tl.make_block_ptr(dgr + bos*H + i_h, (T,), (H,), (i_t * BT,), (BT,), (0,)) - # [BT] - """ - b_dg: 1,2,3,4 - b_dg0: 0,1,2,3 - b_temp: 0,1,3,6 - b_dz: 6 - b_dgr: 6,5,3,0 - """ - b_dg0 = tl.load(p_dg0, boundary_check=(0,)).to(tl.float32) - b_temp = tl.cumsum(b_dg0, axis=0) - b_dz = tl.sum(b_dg0, axis=0) - b_dgr = -b_temp + b_dz[None] - tl.store(p_dgr, b_dgr.to(p_dgr.dtype.element_ty), boundary_check=(0,)) - - -def chunk_comba_cumsum_scalar_bwd( - dg0: torch.Tensor, - chunk_size: int, - cu_seqlens: torch.Tensor | None = None, - output_dtype: torch.dtype | None = torch.float, - chunk_indices: torch.LongTensor | None = None, -) -> torch.Tensor: - B, T, H = dg0.shape - assert chunk_size == 2**(chunk_size.bit_length()-1), "chunk_size must be a power of 2" - BT = chunk_size - if chunk_indices is None and cu_seqlens is not None: - chunk_indices = prepare_chunk_indices(cu_seqlens, BT) - NT = triton.cdiv(T, BT) if cu_seqlens is None else len(chunk_indices) - dg = torch.empty_like(dg0, dtype=output_dtype or dg0.dtype) - grid = (NT, B * H) - chunk_comba_cumsum_scalar_bwd_kernel[grid]( - dg0, - dg, - cu_seqlens, - chunk_indices, - T=T, - B=B, - H=H, - BT=BT, - ) - return dg diff --git a/build/torch-xpu/ops/comba/wy_fast.py b/build/torch-xpu/ops/comba/wy_fast.py deleted file mode 100644 index 431dbb7cfcd4fc9c36e78bc0f460415ada743638..0000000000000000000000000000000000000000 --- a/build/torch-xpu/ops/comba/wy_fast.py +++ /dev/null @@ -1,434 +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 -import triton.language as tl - -from ...ops.utils import prepare_chunk_indices -from ...ops.utils.op import exp2 -from ...utils import autotune_cache_kwargs, check_shared_mem - - -@triton.heuristics({ - 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None, - 'USE_G': lambda args: args['g'] is not None, -}) -@triton.autotune( - configs=[ - triton.Config({'BK': BK}, num_warps=num_warps, num_stages=num_stages) - for BK in [32, 64, 128] - for num_warps in [2, 4, 8] - for num_stages in [2, 3, 4] - ], - key=['H', 'K', 'BT', 'IS_VARLEN', 'USE_G'], - **autotune_cache_kwargs, -) -@triton.jit(do_not_specialize=['T']) -def chunk_scaled_dot_comba_pkt_fwd_kernel( - k, - p, - beta, - g0, - g, - A, - cu_seqlens, - chunk_indices, - T, - H: tl.constexpr, - K: tl.constexpr, - BT: tl.constexpr, - BK: tl.constexpr, - IS_VARLEN: tl.constexpr, - USE_G: tl.constexpr, -): - i_t, i_bh = tl.program_id(0), tl.program_id(1) - i_b, i_h = i_bh // H, i_bh % H - 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.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) - T = eos - bos - else: - bos, eos = i_b * T, i_b * T + T - o_t = i_t * BT + tl.arange(0, BT) - m_t = o_t < T - - p_beta = tl.make_block_ptr(beta + bos*H + i_h, (T,), (H,), (i_t * BT,), (BT,), (0,)) - b_beta = tl.load(p_beta, boundary_check=(0,)) - - b_A = tl.zeros([BT, BT], dtype=tl.float32) - for i_k in range(tl.cdiv(K, BK)): - p_k = tl.make_block_ptr(k + (bos*H + i_h) * K, (T, K), (H*K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) - p_p = tl.make_block_ptr(p + (bos*H + i_h) * K, (T, K), (H*K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) - b_k = tl.load(p_k, boundary_check=(0, 1)) - b_p = tl.load(p_p, boundary_check=(0, 1)) - b_pb = b_p * b_beta[:, None] - b_A += tl.dot(b_pb.to(b_k.dtype), tl.trans(b_k)) - - if USE_G: - p_g0 = tl.make_block_ptr(g0 + bos*H + i_h, (T,), (H,), (i_t * BT,), (BT,), (0,)) - p_g = tl.make_block_ptr(g + bos*H + i_h, (T,), (H,), (i_t * BT,), (BT,), (0,)) - b_g0 = tl.load(p_g0, boundary_check=(0,)) - b_g = tl.load(p_g, boundary_check=(0,)) - b_A = b_A * exp2(b_g0[:, None] - b_g[None, :]) - - m_A = (o_t[:, None] > o_t[None, :]) & (m_t[:, None] & m_t) - b_A = tl.where(m_A, b_A, 0) - p_A = tl.make_block_ptr(A + (bos*H + i_h) * BT, (T, BT), (BT*H, 1), (i_t * BT, 0), (BT, BT), (1, 0)) - tl.store(p_A, b_A.to(p_A.dtype.element_ty), boundary_check=(0, 1)) - - -def chunk_scaled_dot_comba_pkt_fwd( - k: torch.Tensor, - p: torch.Tensor, - beta: torch.Tensor, - g0: torch.Tensor | None = None, - g: torch.Tensor | None = None, - cu_seqlens: torch.LongTensor | None = None, - chunk_size: int = 64, - output_dtype: torch.dtype = torch.float32, - chunk_indices: torch.LongTensor | None = None, -) -> torch.Tensor: - r""" - Compute beta \mathcal{A}(i-1/j) * P * K^T. - - Args: - k (torch.Tensor): - The key tensor of shape `[B, T, H, K]`. - p (torch.Tensor): - The auxiliary key tensor of shape `[B, T, H, K]`. - beta (torch.Tensor): - The beta tensor of shape `[B, T, H]`. - g0 (torch.Tensor): - The cumulative sum minus the original one of the gate tensor of shape `[B, T, H]`. - Default: None - g (torch.Tensor): - The cumulative sum of the gate tensor of shape `[B, T, H]`. - Default: None - cu_seqlens (torch.LongTensor): - The cumulative sequence lengths of the input tensor. - Default: None - chunk_size (int): - The chunk size. Default: 64. - output_dtype (torch.dtype): - The dtype of the output tensor. Default: `torch.float32` - - Returns: - beta * K * K^T of shape `[B, T, H, BT]` where `BT` is the chunk size. - """ - B, T, H, K = k.shape - BT = chunk_size - if chunk_indices is None and cu_seqlens is not None: - chunk_indices = prepare_chunk_indices(cu_seqlens, BT) - NT = triton.cdiv(T, BT) if cu_seqlens is None else len(chunk_indices) - A = torch.empty(B, T, H, BT, device=k.device, dtype=output_dtype) - chunk_scaled_dot_comba_pkt_fwd_kernel[(NT, B * H)]( - k=k, - p=p, - beta=beta, - g0=g0, - g=g, - A=A, - cu_seqlens=cu_seqlens, - chunk_indices=chunk_indices, - T=T, - H=H, - K=K, - BT=BT, - ) - return A - - -@triton.heuristics({ - 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None, -}) -@triton.autotune( - configs=[ - triton.Config({}, num_warps=num_warps, num_stages=num_stages) - for num_warps in [2, 4] - for num_stages in [2, 3, 4] - ], - key=['H', 'K', 'V', 'BT', 'BK', 'BV', 'IS_VARLEN'], - **autotune_cache_kwargs, -) -@triton.jit(do_not_specialize=['T']) -def prepare_wy_repr_bwd_kernel( - k, - v, - p, - beta, - g0, - g, - A, - dw, - du, - dk, - dv, - dp, - dbeta, - dg0, - dg, - cu_seqlens, - chunk_indices, - T, - H: tl.constexpr, - K: tl.constexpr, - V: tl.constexpr, - BT: tl.constexpr, - BK: tl.constexpr, - BV: tl.constexpr, - IS_VARLEN: tl.constexpr, -): - i_t, i_bh = tl.program_id(0), tl.program_id(1) - i_b, i_h = i_bh // H, i_bh % H - 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.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) - T = eos - bos - else: - bos, eos = i_b * T, i_b * T + T - - p_beta = tl.make_block_ptr(beta + (bos*H + i_h), (T,), (H,), (i_t * BT,), (BT,), (0,)) - p_g0 = tl.make_block_ptr(g0 + (bos*H + i_h), (T,), (H,), (i_t * BT,), (BT,), (0,)) - p_g = tl.make_block_ptr(g + (bos*H + i_h), (T,), (H,), (i_t * BT,), (BT,), (0,)) - p_A = tl.make_block_ptr(A + (bos*H + i_h) * BT, (BT, T), (1, H*BT), (0, i_t * BT), (BT, BT), (0, 1)) - - b_A = tl.load(p_A, boundary_check=(0, 1)) - b_beta = tl.load(p_beta, boundary_check=(0,)) - b_g0 = tl.load(p_g0, boundary_check=(0,)) - b_g0_exp = exp2(b_g0) - b_g = tl.load(p_g, boundary_check=(0,)) - - b_dbeta = tl.zeros([BT], dtype=tl.float32) - b_dA = tl.zeros([BT, BT], dtype=tl.float32) - b_dg0 = tl.zeros([BT], dtype=tl.float32) - - for i_k in range(tl.cdiv(K, BK)): - p_p = tl.make_block_ptr(p + (bos*H + i_h) * K, (T, K), (H*K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) - p_dp = tl.make_block_ptr(dp + (bos*H + i_h) * K, (T, K), (H*K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) - p_dw = tl.make_block_ptr(dw + (bos*H + i_h) * K, (T, K), (H*K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) - b_p = tl.load(p_p, boundary_check=(0, 1)) - b_p_beta_g0 = (b_p * b_beta[:, None] * b_g0_exp[:, None]).to(b_p.dtype) - b_dw = tl.load(p_dw, boundary_check=(0, 1)) - b_dA += tl.dot(b_dw, tl.trans(b_p_beta_g0)) - b_dp_beta_g0 = tl.dot(b_A, b_dw) - b_dp = b_dp_beta_g0 * b_beta[:, None] * b_g0_exp[:, None] - b_dbeta += tl.sum(b_dp_beta_g0 * b_p * b_g0_exp[:, None], 1) - b_dg0 += tl.sum(b_dp * b_p, 1) - tl.store(p_dp, b_dp.to(p_dp.dtype.element_ty), boundary_check=(0, 1)) - - for i_v in range(tl.cdiv(V, BV)): - p_v = tl.make_block_ptr(v + (bos*H + i_h) * V, (T, V), (H*V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) - p_dv = tl.make_block_ptr(dv + (bos*H + i_h) * V, (T, V), (H*V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) - p_du = tl.make_block_ptr(du + (bos*H + i_h) * V, (T, V), (H*V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) - b_v = tl.load(p_v, boundary_check=(0, 1)) - b_v_beta = (b_v * b_beta[:, None]).to(b_v.dtype) - b_du = tl.load(p_du, boundary_check=(0, 1)) - b_dA += tl.dot(b_du, tl.trans(b_v_beta)) - b_dv_beta = tl.dot(b_A, b_du) - b_dv = b_dv_beta * b_beta[:, None] - b_dbeta += tl.sum(b_dv_beta * b_v, 1) - tl.store(p_dv, b_dv.to(p_dv.dtype.element_ty), boundary_check=(0, 1)) - - o_t = i_t * BT + tl.arange(0, BT) - m_t = o_t < T - m_A = (o_t[:, None] > o_t[None, :]) & (m_t[:, None] & m_t) - b_dA = tl.where(m_A, b_dA, 0) - b_dA = tl.dot(b_dA.to(b_A.dtype), b_A) - b_dA = tl.dot(b_A, b_dA.to(b_A.dtype)) - b_dA = tl.where(m_A, -b_dA * exp2(b_g0[:, None] - b_g[None, :]), 0).to(k.dtype.element_ty) - b_dA = b_dA.to(k.dtype.element_ty) - b_A = tl.zeros([BT, BT], dtype=tl.float32) - - for i_k in range(tl.cdiv(K, BK)): - p_k = tl.make_block_ptr(k + (bos*H + i_h) * K, (T, K), (H*K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) - p_p = tl.make_block_ptr(p + (bos*H + i_h) * K, (T, K), (H*K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) - p_dk = tl.make_block_ptr(dk + (bos*H + i_h) * K, (T, K), (H*K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) - p_dp = tl.make_block_ptr(dp + (bos*H + i_h) * K, (T, K), (H*K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) - b_k = tl.load(p_k, boundary_check=(0, 1)) - b_p = tl.load(p_p, boundary_check=(0, 1)) - b_dp = tl.load(p_dp, boundary_check=(0, 1)) - b_p_beta = (b_p * b_beta[:, None]).to(b_p.dtype) - b_A += tl.dot(b_p_beta, tl.trans(b_k)) - b_dp_beta = tl.dot(b_dA, b_k) - b_dbeta += tl.sum(b_dp_beta * b_p, 1) - b_dk = tl.dot(tl.trans(b_dA), b_p_beta) - b_dp += b_dp_beta * b_beta[:, None] - tl.store(p_dk, b_dk.to(p_dk.dtype.element_ty), boundary_check=(0, 1)) - tl.store(p_dp, b_dp.to(p_dp.dtype.element_ty), boundary_check=(0, 1)) - - b_dA_A = b_dA * b_A - b_dg0 += tl.sum(b_dA_A, axis=1) - b_dg = - tl.sum(b_dA_A, axis=0) - p_dg = tl.make_block_ptr(dg + (bos*H + i_h), (T,), (H,), (i_t * BT,), (BT,), (0,)) - p_dg0 = tl.make_block_ptr(dg0 + (bos*H + i_h), (T,), (H,), (i_t * BT,), (BT,), (0,)) - p_dbeta = tl.make_block_ptr(dbeta + (bos*H + i_h), (T,), (H,), (i_t * BT,), (BT,), (0,)) - tl.store(p_dg, b_dg.to(p_dg.dtype.element_ty), boundary_check=(0,)) - tl.store(p_dg0, b_dg0.to(p_dg0.dtype.element_ty), boundary_check=(0,)) - tl.store(p_dbeta, b_dbeta.to(p_dbeta.dtype.element_ty), boundary_check=(0,)) - - -@triton.heuristics({ - 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None, -}) -@triton.autotune( - configs=[ - triton.Config({}, num_warps=num_warps, num_stages=num_stages) - for num_warps in [2, 4, 8] - for num_stages in [2, 3, 4] - ], - key=['H', 'K', 'V', 'BT', 'BK', 'BV', 'IS_VARLEN'], - **autotune_cache_kwargs, -) -@triton.jit(do_not_specialize=['T']) -def recompute_w_u_fwd_kernel( - k, - v, - beta, - w, - u, - A, - g, - cu_seqlens, - chunk_indices, - T, - H: tl.constexpr, - K: tl.constexpr, - V: tl.constexpr, - BT: tl.constexpr, - BK: tl.constexpr, - BV: tl.constexpr, - IS_VARLEN: tl.constexpr, -): - i_t, i_bh = tl.program_id(0), tl.program_id(1) - i_b, i_h = i_bh // H, i_bh % H - 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.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) - T = eos - bos - else: - bos, eos = i_b * T, i_b * T + T - p_beta = tl.make_block_ptr(beta + bos*H + i_h, (T,), (H,), (i_t * BT,), (BT,), (0,)) - p_g = tl.make_block_ptr(g + (bos*H + i_h), (T,), (H,), (i_t * BT,), (BT,), (0,)) - p_A = tl.make_block_ptr(A + (bos*H + i_h) * BT, (T, BT), (H*BT, 1), (i_t * BT, 0), (BT, BT), (1, 0)) - b_beta = tl.load(p_beta, boundary_check=(0,)) - b_A = tl.load(p_A, boundary_check=(0, 1)) - b_g = exp2(tl.load(p_g, boundary_check=(0,))) - - for i_v in range(tl.cdiv(V, BV)): - p_v = tl.make_block_ptr(v + (bos*H + i_h) * V, (T, V), (H*V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) - p_u = tl.make_block_ptr(u + (bos*H + i_h) * V, (T, V), (H*V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) - b_v = tl.load(p_v, boundary_check=(0, 1)) - b_vb = (b_v * b_beta[:, None]).to(b_v.dtype) - b_u = tl.dot(b_A, b_vb, allow_tf32=False) - tl.store(p_u, b_u.to(p_u.dtype.element_ty), boundary_check=(0, 1)) - - for i_k in range(tl.cdiv(K, BK)): - p_k = tl.make_block_ptr(k + (bos*H + i_h) * K, (T, K), (H*K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) - p_w = tl.make_block_ptr(w + (bos*H + i_h) * K, (T, K), (H*K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) - b_k = tl.load(p_k, boundary_check=(0, 1)) - b_kb = (b_k * b_beta[:, None] * b_g[:, None]).to(b_k.dtype) - b_w = tl.dot(b_A, b_kb) - tl.store(p_w, b_w.to(p_w.dtype.element_ty), boundary_check=(0, 1)) - - -def recompute_w_u_fwd( - k: torch.Tensor, - v: torch.Tensor, - beta: torch.Tensor, - g_cumsum: torch.Tensor, - A: torch.Tensor, - cu_seqlens: torch.LongTensor | None, - chunk_indices: torch.LongTensor | None = None, -) -> tuple[torch.Tensor, torch.Tensor]: - B, T, H, K, V = *k.shape, v.shape[-1] - BT = A.shape[-1] - - if chunk_indices is None and cu_seqlens is not None: - chunk_indices = prepare_chunk_indices(cu_seqlens, BT) - NT = triton.cdiv(T, BT) if cu_seqlens is None else len(chunk_indices) - BK = 64 - BV = 64 - - u = torch.empty_like(v) - w = torch.empty_like(k) - recompute_w_u_fwd_kernel[(NT, B*H)]( - k=k, - v=v, - beta=beta, - w=w, - u=u, - A=A, - g=g_cumsum, - cu_seqlens=cu_seqlens, - chunk_indices=chunk_indices, - T=T, - H=H, - K=K, - V=V, - BT=BT, - BK=BK, - BV=BV, - ) - return w, u - - -def prepare_wy_repr_bwd( - k: torch.Tensor, - v: torch.Tensor, - p: torch.Tensor, - g0: torch.Tensor, - g: torch.Tensor, - beta: torch.Tensor, - A: torch.Tensor, - dw: torch.Tensor, - du: torch.Tensor, - cu_seqlens: torch.LongTensor | None, - chunk_indices: torch.LongTensor | None = None, -) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: - B, T, H, K, V = *k.shape, v.shape[-1] - BT = A.shape[-1] - if chunk_indices is None and cu_seqlens is not None: - chunk_indices = prepare_chunk_indices(cu_seqlens, BT) - NT = triton.cdiv(T, BT) if cu_seqlens is None else len(chunk_indices) - CONST_TILING = 64 if check_shared_mem() else 32 - BK = min(max(triton.next_power_of_2(K), 16), CONST_TILING) - BV = min(max(triton.next_power_of_2(V), 16), CONST_TILING) - - dk = torch.empty_like(k) - dv = torch.empty_like(v) - dp = torch.empty_like(p) - dbeta = torch.empty_like(beta) - dg0 = torch.empty_like(g0) - dg = torch.empty_like(g) - prepare_wy_repr_bwd_kernel[(NT, B * H)]( - k=k, - v=v, - p=p, - beta=beta, - g0=g0, - g=g, - A=A, - dw=dw, - du=du, - dk=dk, - dv=dv, - dp=dp, - dbeta=dbeta, - dg0=dg0, - dg=dg, - cu_seqlens=cu_seqlens, - chunk_indices=chunk_indices, - T=T, - H=H, - K=K, - V=V, - BT=BT, - BK=BK, - BV=BV, - ) - return dk, dv, dp, dbeta, dg0, dg diff --git a/build/torch-xpu/ops/common/__init__.py b/build/torch-xpu/ops/common/__init__.py deleted file mode 100644 index e69de29bb2d1d6434b8b29ae775ad8c2e48c5391..0000000000000000000000000000000000000000 diff --git a/build/torch-xpu/ops/common/backends/__init__.py b/build/torch-xpu/ops/common/backends/__init__.py deleted file mode 100644 index 18ea07b2dee4fa20beb9809b751e5b86f222db2a..0000000000000000000000000000000000000000 --- a/build/torch-xpu/ops/common/backends/__init__.py +++ /dev/null @@ -1,21 +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 - -"""Common backends for shared operations like chunk_gated_delta_rule_fwd_h.""" - -from ....ops.backends import BackendRegistry, dispatch -from ....ops.common.backends.intracard import IntraCardCPBackend -from ....ops.common.backends.tilelang import TileLangBackend - -common_registry = BackendRegistry("common") - - -common_registry.register(IntraCardCPBackend()) -common_registry.register(TileLangBackend()) - - -__all__ = ['common_registry', 'dispatch'] diff --git a/build/torch-xpu/ops/common/backends/intracard.py b/build/torch-xpu/ops/common/backends/intracard.py deleted file mode 100644 index 17c47ad52841ca4aa6cc6d201c689490eae338a7..0000000000000000000000000000000000000000 --- a/build/torch-xpu/ops/common/backends/intracard.py +++ /dev/null @@ -1,98 +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 - -"""Intra-card CP backend for shared delta rule operations. - -Accelerates prefill by splitting long sequences into sub-sequences -and processing them in parallel across SMs. - -Only active under torch.inference_mode() with varlen (cu_seqlens != None). -""" - -from __future__ import annotations - -import os - -import torch - -from ....ops.backends import BaseBackend - -# Maximum number of sub-sequences per original sequence -# Limits merge chain depth to control precision loss -MAX_SUBSEQS = int(os.environ.get('FLA_INTRACARD_MAX_SPLITS', 32)) - - -class IntraCardCPBackend(BaseBackend): - """Intra-card context parallel backend for chunk_gated_delta_rule_fwd_h.""" - - backend_type = "intracard_cp" - package_name = None # No external package needed - env_var = "FLA_INTRACARD_CP" - default_enable = False - - @classmethod - def is_available(cls) -> bool: - return True - - def chunk_gated_delta_rule_fwd_h_verifier( - self, - k: torch.Tensor, - w: torch.Tensor, - u: torch.Tensor, - g: torch.Tensor | None = None, - gk: torch.Tensor | None = None, - initial_state: torch.Tensor | None = None, - output_final_state: bool = False, - chunk_size: int = 64, - save_new_value: bool = True, - state_v_first: bool = False, - cu_seqlens: torch.LongTensor | None = None, - cu_seqlens_cpu: torch.LongTensor | None = None, - chunk_indices: torch.LongTensor | None = None, - ) -> tuple[bool, str | None]: - """Check if intracard CP should handle this call.""" - # Only in inference mode - if not torch.is_inference_mode_enabled(): - return False, "Not in inference mode" - - # Only for varlen - if cu_seqlens is None: - return False, "cu_seqlens is None" - - return True, None - - def chunk_gated_delta_rule_fwd_h( - self, - k: torch.Tensor, - w: torch.Tensor, - u: torch.Tensor, - g: torch.Tensor | None = None, - gk: torch.Tensor | None = None, - initial_state: torch.Tensor | None = None, - output_final_state: bool = False, - chunk_size: int = 64, - save_new_value: bool = True, - state_v_first: bool = False, - cu_seqlens: torch.LongTensor | None = None, - cu_seqlens_cpu: torch.LongTensor | None = None, - chunk_indices: torch.LongTensor | None = None, - ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor | None]: - """Intra-card CP implementation of chunk_gated_delta_rule_fwd_h.""" - from ....ops.common.intracard_cp import intracard_fwd_h - - return intracard_fwd_h( - k=k, w=w, u=u, g=g, gk=gk, - initial_state=initial_state, - output_final_state=output_final_state, - chunk_size=chunk_size, - save_new_value=save_new_value, - cu_seqlens=cu_seqlens, - cu_seqlens_cpu=cu_seqlens_cpu, - chunk_indices=chunk_indices, - max_splits=MAX_SUBSEQS, - state_v_first=state_v_first, - ) diff --git a/build/torch-xpu/ops/common/backends/tilelang/__init__.py b/build/torch-xpu/ops/common/backends/tilelang/__init__.py deleted file mode 100644 index 088f5ac608a1a8ac368eda0cdc17594083ce999b..0000000000000000000000000000000000000000 --- a/build/torch-xpu/ops/common/backends/tilelang/__init__.py +++ /dev/null @@ -1,201 +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 - -"""TileLang backend for common chunk operations. - -Enabled by default on Hopper (sm90+) with Triton >= 3.4.0 to work around -hardware-specific regressions (see #640). Can also be forced via FLA_TILELANG=1. -""" - -from __future__ import annotations - -import os - -import torch - -from .....ops.backends import BaseBackend -from .....utils import IS_NVIDIA_HOPPER, TRITON_ABOVE_3_4_0, find_spec_cached - -_TILELANG_AVAILABLE = find_spec_cached("tilelang") is not None - - -class TileLangBackend(BaseBackend): - - backend_type = "tilelang" - package_name = "tilelang" - env_var = "FLA_TILELANG" - - @classmethod - def is_available(cls) -> bool: - return _TILELANG_AVAILABLE - - @classmethod - def is_enabled(cls) -> bool: - # Explicit opt-in / opt-out always wins. - val = os.environ.get(cls.env_var) - if val is not None: - return val != "0" - # Default on only where the Triton path is known to be broken: - # Hopper (sm90) with Triton >= 3.4.0 (see #640). Everywhere else the - # Triton backend stays the default unless the user forces FLA_TILELANG=1. - return IS_NVIDIA_HOPPER and TRITON_ABOVE_3_4_0 - - def chunk_bwd_dqkwg_verifier( - self, - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - do: torch.Tensor, - h: torch.Tensor, - dh: torch.Tensor, - w: torch.Tensor | None = None, - g: torch.Tensor | None = None, - g_gamma: torch.Tensor | None = None, - dv: torch.Tensor | None = None, - scale: float | None = None, - state_v_first: bool = False, - cu_seqlens: torch.LongTensor | None = None, - chunk_size: int = 64, - chunk_indices: torch.LongTensor | None = None, - ) -> tuple[bool, str | None]: - if g is None: - return False, "TileLang backend only supports gated case (g != None)" - if g_gamma is not None: - return False, "TileLang backend does not support g_gamma" - if v.shape[2] % k.shape[2] != 0: - return False, ( - f"TileLang backend requires num_v_heads (HV={v.shape[2]}) to be divisible by " - f"num_qk_heads (H={k.shape[2]}); HV % H must be 0 for GVA" - ) - if h.dtype != q.dtype: - return False, ( - f"TileLang backend requires h.dtype == q.dtype (got h={h.dtype}, q={q.dtype}); " - "e.g. simple_gla's bwd keeps h/dh in fp32 for h·dh reduction precision" - ) - return True, None - - def chunk_bwd_dqkwg( - self, - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - do: torch.Tensor, - h: torch.Tensor, - dh: torch.Tensor, - w: torch.Tensor | None = None, - g: torch.Tensor | None = None, - g_gamma: torch.Tensor | None = None, - dv: torch.Tensor | None = None, - scale: float | None = None, - state_v_first: bool = False, - cu_seqlens: torch.LongTensor | None = None, - chunk_size: int = 64, - chunk_indices: torch.LongTensor | None = None, - ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor | None, torch.Tensor | None]: - from .....ops.common.backends.tilelang.chunk_bwd import ( - chunk_bwd_dqkwg_tilelang, - ) - return chunk_bwd_dqkwg_tilelang( - q=q, - k=k, - v=v, - do=do, - h=h, - dh=dh, - w=w, - g=g, - g_gamma=g_gamma, - dv=dv, - scale=scale, - cu_seqlens=cu_seqlens, - chunk_size=chunk_size, - chunk_indices=chunk_indices, - state_v_first=state_v_first, - ) - - def parallel_attn_fwd_verifier( - self, - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - g_cumsum: torch.Tensor | None, - sink_bias: torch.Tensor | None, - scale: float, - window_size: int | None = None, - cu_seqlens: torch.LongTensor | None = None, - chunk_indices: torch.LongTensor | None = None, - ) -> tuple[bool, str | None]: - if q.dtype not in (torch.float16, torch.bfloat16, torch.float32): - return False, f"TileLang backend does not support dtype {q.dtype}; fall back to Triton" - return True, None - - def parallel_attn_fwd( - self, - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - g_cumsum: torch.Tensor | None, - sink_bias: torch.Tensor | None, - scale: float, - window_size: int | None = None, - cu_seqlens: torch.LongTensor | None = None, - chunk_indices: torch.LongTensor | None = None, - ) -> tuple[torch.Tensor, torch.Tensor]: - from .....ops.common.backends.tilelang.parallel_attn_fwd import ( - parallel_attn_fwd_tilelang, - ) - return parallel_attn_fwd_tilelang( - q=q, k=k, v=v, g_cumsum=g_cumsum, sink_bias=sink_bias, - scale=scale, window_size=window_size, cu_seqlens=cu_seqlens, - chunk_indices=chunk_indices, - ) - - def parallel_attn_bwd_verifier( - self, - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - o: torch.Tensor, - g_cumsum: torch.Tensor | None, - lse: torch.Tensor, - do: torch.Tensor, - sink_bias: torch.Tensor | None = None, - scale: float | None = None, - window_size: int | None = None, - chunk_size: int = 128, - cu_seqlens: torch.LongTensor | None = None, - chunk_indices: torch.LongTensor | None = None, - ) -> tuple[bool, str | None]: - if q.dtype not in (torch.float16, torch.bfloat16, torch.float32): - return False, f"TileLang backend does not support dtype {q.dtype}; fall back to Triton" - return True, None - - def parallel_attn_bwd( - self, - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - o: torch.Tensor, - g_cumsum: torch.Tensor | None, - lse: torch.Tensor, - do: torch.Tensor, - sink_bias: torch.Tensor | None = None, - scale: float | None = None, - window_size: int | None = None, - chunk_size: int = 128, - cu_seqlens: torch.LongTensor | None = None, - chunk_indices: torch.LongTensor | None = None, - ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor | None, torch.Tensor | None]: - from .....ops.common.backends.tilelang.parallel_attn_bwd import ( - parallel_attn_bwd_tilelang, - ) - return parallel_attn_bwd_tilelang( - q=q, k=k, v=v, o=o, g_cumsum=g_cumsum, lse=lse, do=do, - sink_bias=sink_bias, scale=scale, window_size=window_size, - chunk_size=chunk_size, cu_seqlens=cu_seqlens, - chunk_indices=chunk_indices, - ) diff --git a/build/torch-xpu/ops/common/backends/tilelang/chunk_bwd.py b/build/torch-xpu/ops/common/backends/tilelang/chunk_bwd.py deleted file mode 100644 index 3c89a544206dfdd92fb9a87b4daf2092f5162f53..0000000000000000000000000000000000000000 --- a/build/torch-xpu/ops/common/backends/tilelang/chunk_bwd.py +++ /dev/null @@ -1,395 +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 tilelang -import tilelang.language as T -import torch -import triton - -from .....ops.utils import prepare_chunk_indices -from .....utils import check_shared_mem - - -@tilelang.jit(pass_configs={ - tilelang.PassConfigKey.TL_DISABLE_TMA_LOWER: True, - tilelang.PassConfigKey.TL_ENABLE_FAST_MATH: True, -}) -def _build_kernel( - B, - H, - HV, - K, - V, - BT, - BK, - BV, - NK, - hD1, - hD2, - dtype_str, - USE_G, - USE_DW, - STATE_V_FIRST, - IS_VARLEN=False, - num_warps=4, -): - dtype_map = {'float16': T.float16, 'bfloat16': T.bfloat16, 'float32': T.float32} - _dtype = dtype_map[dtype_str] - NV = tilelang.cdiv(V, BV) - threads = num_warps * 32 - tile_hD1, tile_hD2 = (BV, BK) if STATE_V_FIRST else (BK, BV) - - # Rebind to underscore-prefixed locals so the kernel body (a closure below) - # stays identical to the previous nested layout. tilelang caches _build_kernel - # by the outer (B, H, K, V, BT, BK, BV, NK, hD1, hD2, dtype_str, ...) tuple. - _B, _H, _HV, _K, _V = B, H, HV, K, V - _G = HV // H - _BT, _BK, _BV, _NK = BT, BK, BV, NK - _NV = NV - _hD1, _hD2, _thD1, _thD2 = hD1, hD2, tile_hD1, tile_hD2 - _threads = threads - _USE_G, _USE_DW = USE_G, USE_DW - _TS, _VAR = STATE_V_FIRST, IS_VARLEN - - # T, NT, total_h are dynamic (vary with sequence length, no recompilation). - # B, H are compile-time (stable across batches, enables fast integer division). - T_d, NT_d, total_h_d, Ncu_d = T.dynamic("T, NT, total_h, Ncu") - - # 4D tensor shapes using dynamic T + compile-time B, H, K, V. - # q/k carry H qk-heads; dq/dk/dw and v/do/g/dg carry HV value-heads - # (HV >= H for GVA). dq/dk/dw are written at value-head granularity and - # reduced to qk-heads by the caller when HV > H. - qk_s = (_B, T_d, _H, _K) - dqk_s = (_B, T_d, _HV, _K) - v_s = (_B, T_d, _HV, _V) - h_s = (total_h_d, _hD1, _hD2) - g_s = (_B, T_d, _HV) - dg_s = (_NK, _B, T_d, _HV) - - @T.macro - def kernel_body(q, k, v, g, h, do, dh, dq, dk, dw, dv, dg, scale, - i_b, i_h, i_k, t_s, T_seq, i_t_local, h_idx, k_off): - # Map the value-head program id to the qk-head it belongs to. For GVA - # (HV > H) multiple value-heads share the same q/k head. - i_hqk = i_h // _G - # -- accumulators -- - b_dq = T.alloc_fragment((_BT, _BK), T.float32) - b_dk = T.alloc_fragment((_BT, _BK), T.float32) - b_ds = T.alloc_fragment((_BT, _BT), T.float32) - T.clear(b_dq) - T.clear(b_dk) - T.clear(b_ds) - - if _USE_DW: - b_dw = T.alloc_fragment((_BT, _BK), T.float32) - T.clear(b_dw) - - # -- shared tiles -- - s_v = T.alloc_shared((_BT, _BV), _dtype) - s_do = T.alloc_shared((_BT, _BV), _dtype) - s_h = T.alloc_shared((_thD1, _thD2), _dtype) - s_dh = T.alloc_shared((_thD1, _thD2), _dtype) - - # dg_last accumulator (shared scalar, accumulated across V-loop) - if _USE_G: - s_dg_last_acc = T.alloc_shared((1,), T.float32) - for _i in T.Parallel(1): - s_dg_last_acc[0] = 0.0 - T.sync_threads() - - # ========== V-loop ========== - for i_v_py in T.Pipelined(_NV, num_stages=2): - v_off_c = i_v_py * _BV - - T.copy(v[i_b, t_s:t_s + _BT, i_h, v_off_c:v_off_c + _BV], s_v) - T.copy(do[i_b, t_s:t_s + _BT, i_h, v_off_c:v_off_c + _BV], s_do) - - if _TS: - T.copy(h[h_idx, v_off_c:v_off_c + _BV, k_off:k_off + _BK], s_h) - T.copy(dh[h_idx, v_off_c:v_off_c + _BV, k_off:k_off + _BK], s_dh) - else: - T.copy(h[h_idx, k_off:k_off + _BK, v_off_c:v_off_c + _BV], s_h) - T.copy(dh[h_idx, k_off:k_off + _BK, v_off_c:v_off_c + _BV], s_dh) - - T.gemm(s_do, s_v, b_ds, transpose_B=True) - - # h·dh reduction must precede the gemms that consume s_h/s_dh: - # the downstream gemms are what the pipeline recognizes as consumers, - # so they act as the barrier against the next iter's prefetch. - if _USE_G: - f_hdh = T.alloc_fragment((_thD1, _thD2), T.float32) - for _i, _j in T.Parallel(_thD1, _thD2): - f_hdh[_i, _j] = T.cast(s_h[_i, _j], T.float32) * T.cast(s_dh[_i, _j], T.float32) - f_hdh_row = T.alloc_fragment((_thD1,), T.float32) - T.reduce_sum(f_hdh, f_hdh_row, dim=1) - f_hdh_scalar = T.alloc_fragment((1,), T.float32) - T.reduce_sum(f_hdh_row, f_hdh_scalar, dim=0) - s_dg_last_acc[0] = s_dg_last_acc[0] + f_hdh_scalar[0] - - if _TS: - T.gemm(s_do, s_h, b_dq) - T.gemm(s_v, s_dh, b_dk) - else: - T.gemm(s_do, s_h, b_dq, transpose_B=True) - T.gemm(s_v, s_dh, b_dk, transpose_B=True) - - if _USE_DW: - s_dv = T.alloc_shared((_BT, _BV), _dtype) - T.copy(dv[i_b, t_s:t_s + _BT, i_h, v_off_c:v_off_c + _BV], s_dv) - if _TS: - T.gemm(s_dv, s_h, b_dw) - else: - T.gemm(s_dv, s_h, b_dw, transpose_B=True) - - # ========== store dw (negated, with varlen boundary mask) ========== - if _USE_DW: - s_dw_out = T.alloc_shared((_BT, _BK), _dtype) - for _i, _j in T.Parallel(_BT, _BK): - s_dw_out[_i, _j] = T.cast(-b_dw[_i, _j], _dtype) - T.sync_threads() - for _i, _j in T.Parallel(_BT, _BK): - if (i_t_local * _BT + _i) < T_seq: - dw[i_b, t_s + _i, i_h, k_off + _j] = s_dw_out[_i, _j] - - # dg_last is now in s_dg_last_acc[0] (shared memory, visible to all threads) - - # ========== load q, k ========== - # q/k are indexed by the shared qk-head, not the value-head. - s_q = T.alloc_shared((_BT, _BK), _dtype) - s_k = T.alloc_shared((_BT, _BK), _dtype) - T.copy(q[i_b, t_s:t_s + _BT, i_hqk, k_off:k_off + _BK], s_q) - T.copy(k[i_b, t_s:t_s + _BT, i_hqk, k_off:k_off + _BK], s_k) - - # ========== USE_G path ========== - if _USE_G: - # dg_last_acc already reduced to `red` scalar above - - # g in shared memory (all threads can cross-read any element) - s_g = T.alloc_shared((_BT,), T.float32) - T.copy(g[i_b, t_s:t_s + _BT, i_h], s_g, disable_tma=True) - - last_pos = T.max(0, T.min(_BT, T_seq - i_t_local * _BT) - 1) - g_last = s_g[last_pos] - b_dg_last = T.alloc_var(T.float32) - b_dg_last = s_dg_last_acc[0] * T.exp2(g_last) - - # b_dq *= exp2(g) * scale (inline, no extra fragment) - for _i, _j in T.Parallel(_BT, _BK): - b_dq[_i, _j] = b_dq[_i, _j] * T.exp2(s_g[_i]) * scale - - # Gate b_dk with m_t mask: zero out OOB positions - for _i, _j in T.Parallel(_BT, _BK): - m_t = (i_t_local * _BT + _i) < T_seq - b_dk[_i, _j] = T.if_then_else( - m_t, - b_dk[_i, _j] * T.exp2(-s_g[_i] + g_last), - 0.0) - - # dk*k reduction (before ds gemms): compute product directly in - # fragments via fragment×shared. b_dg_last uses the gated-only b_dk. - f_prod2 = T.alloc_fragment((_BT, _BK), T.float32) - for _i, _j in T.Parallel(_BT, _BK): - f_prod2[_i, _j] = b_dk[_i, _j] * s_k[_i, _j] - f_dg2 = T.alloc_fragment((_BT,), T.float32) - T.reduce_sum(f_prod2, f_dg2, dim=1) - # b_dg_last += sum(dk*k) via fragment reduce (f_dg2 → scalar) - f_dkk_scalar = T.alloc_fragment((1,), T.float32) - T.reduce_sum(f_dg2, f_dkk_scalar, dim=0) - b_dg_last = b_dg_last + f_dkk_scalar[0] - - # b_ds = where(causal, b_ds * exp2(g_i - g_j), 0) * scale - for _i, _j in T.Parallel(_BT, _BT): - causal = (_i >= _j) & ((i_t_local * _BT + _i) < T_seq) & ((i_t_local * _BT + _j) < T_seq) - b_ds[_i, _j] = T.if_then_else( - causal, - b_ds[_i, _j] * T.exp2(s_g[_i] - s_g[_j]) * scale, - 0.0) - - # cast ds for final gemms - s_ds = T.alloc_shared((_BT, _BT), _dtype) - f_ds = T.alloc_fragment((_BT, _BT), _dtype) - for _i, _j in T.Parallel(_BT, _BT): - f_ds[_i, _j] = T.cast(b_ds[_i, _j], _dtype) - T.copy(f_ds, s_ds) - - T.gemm(s_ds, s_k, b_dq) # dq += ds @ k - T.gemm(s_ds, s_q, b_dk, transpose_A=True) # dk += ds^T @ q - - # b_dg = sum(b_dq * q) - sum(b_dk * k) using the fully-updated b_dq/b_dk - f_prod1 = T.alloc_fragment((_BT, _BK), T.float32) - for _i, _j in T.Parallel(_BT, _BK): - f_prod1[_i, _j] = b_dq[_i, _j] * s_q[_i, _j] - f_prod2[_i, _j] = b_dk[_i, _j] * s_k[_i, _j] - f_dg1 = T.alloc_fragment((_BT,), T.float32) - T.reduce_sum(f_prod1, f_dg1, dim=1) - T.reduce_sum(f_prod2, f_dg2, dim=1) - f_dg_diff = T.alloc_fragment((_BT,), T.float32) - for _i in T.Parallel(_BT): - f_dg_diff[_i] = f_dg1[_i] - f_dg2[_i] - s_dg = T.alloc_shared((_BT,), T.float32) - T.copy(f_dg_diff, s_dg) - - # store dq, dk via shared (fragment has MMA layout, can't index directly) - f_out = T.alloc_fragment((_BT, _BK), _dtype) - s_out = T.alloc_shared((_BT, _BK), _dtype) - for _i, _j in T.Parallel(_BT, _BK): - f_out[_i, _j] = T.cast(b_dq[_i, _j], _dtype) - T.copy(f_out, s_out) - T.sync_threads() - for _i, _j in T.Parallel(_BT, _BK): - if (i_t_local * _BT + _i) < T_seq: - dq[i_b, t_s + _i, i_h, k_off + _j] = s_out[_i, _j] - for _i, _j in T.Parallel(_BT, _BK): - f_out[_i, _j] = T.cast(b_dk[_i, _j], _dtype) - T.copy(f_out, s_out) - T.sync_threads() - for _i, _j in T.Parallel(_BT, _BK): - if (i_t_local * _BT + _i) < T_seq: - dk[i_b, t_s + _i, i_h, k_off + _j] = s_out[_i, _j] - - # store dg: merge dg_last into last valid position - for _i in T.Parallel(_BT): - if (i_t_local * _BT + _i) < T_seq: - val = T.if_then_else(_i == last_pos, s_dg[_i] + b_dg_last, s_dg[_i]) - dg[i_k, i_b, t_s + _i, i_h] = val - - if _VAR: - @T.prim_func - def kernel( - q: T.Tensor(qk_s, _dtype), k: T.Tensor(qk_s, _dtype), - v: T.Tensor(v_s, _dtype), g: T.Tensor(g_s, T.float32), - h: T.Tensor(h_s, _dtype), do: T.Tensor(v_s, _dtype), - dh: T.Tensor(h_s, _dtype), dq: T.Tensor(dqk_s, _dtype), - dk: T.Tensor(dqk_s, _dtype), dw: T.Tensor(dqk_s, _dtype), - dv: T.Tensor(v_s, _dtype), dg: T.Tensor(dg_s, T.float32), - cu_seqlens: T.Tensor((Ncu_d,), T.int32), - chunk_indices: T.Tensor((NT_d, 2), T.int32), - scale: T.float32, - ): - with T.Kernel(_NK, NT_d, _HV, threads=_threads) as (i_k, i_t, i_h): - i_n = chunk_indices[i_t, 0] - i_t_local = chunk_indices[i_t, 1] - bos = cu_seqlens[i_n] - T_seq = cu_seqlens[i_n + 1] - bos - h_idx = i_t * _HV + i_h - t_s = bos + i_t_local * _BT - kernel_body(q, k, v, g, h, do, dh, dq, dk, dw, dv, dg, - scale, 0, i_h, i_k, t_s, T_seq, i_t_local, - h_idx, i_k * _BK) - else: - @T.prim_func - def kernel( - q: T.Tensor(qk_s, _dtype), k: T.Tensor(qk_s, _dtype), - v: T.Tensor(v_s, _dtype), g: T.Tensor(g_s, T.float32), - h: T.Tensor(h_s, _dtype), do: T.Tensor(v_s, _dtype), - dh: T.Tensor(h_s, _dtype), dq: T.Tensor(dqk_s, _dtype), - dk: T.Tensor(dqk_s, _dtype), dw: T.Tensor(dqk_s, _dtype), - dv: T.Tensor(v_s, _dtype), dg: T.Tensor(dg_s, T.float32), - scale: T.float32, - ): - with T.Kernel(_NK, T.ceildiv(T_d, _BT), _B * _HV, threads=_threads) as (i_k, i_t, i_bh): - i_b = i_bh // _HV - i_h = i_bh % _HV - NT_local = T.ceildiv(T_d, _BT) - h_idx = (i_b * NT_local + i_t) * _HV + i_h - t_s = i_t * _BT - kernel_body(q, k, v, g, h, do, dh, dq, dk, dw, dv, dg, - scale, i_b, i_h, i_k, t_s, T_d, i_t, - h_idx, i_k * _BK) - - return kernel - - -def chunk_bwd_dqkwg_tilelang( - q, - k, - v, - do, - h, - dh, - w=None, - g=None, - g_gamma=None, - dv=None, - scale=None, - state_v_first=False, - cu_seqlens=None, - chunk_size=64, - chunk_indices=None, -): - B, T, H, K = k.shape - HV, V = v.shape[2], v.shape[-1] - BT = chunk_size - if chunk_indices is None and cu_seqlens is not None: - chunk_indices = prepare_chunk_indices(cu_seqlens, BT) - IS_VARLEN = cu_seqlens is not None - - CONST_TILING = 64 if check_shared_mem() else 32 - BK = min(max(triton.next_power_of_2(K), 16), CONST_TILING) - BV = min(max(triton.next_power_of_2(V), 16), CONST_TILING) - NK = triton.cdiv(K, BK) - if scale is None: - scale = K ** -0.5 - - USE_G = g is not None - USE_DW = w is not None - - # Outputs — kernel writes dq/dk at value-head (HV) granularity; reduce to - # qk-head (H) below when GVA is active. dg stays at HV (per value-head). - dq = torch.empty(B, T, HV, K, dtype=q.dtype, device=q.device) - dk = torch.empty(B, T, HV, K, dtype=k.dtype, device=k.device) - dw_out = torch.empty_like(w) if USE_DW else None - dg = torch.zeros(NK, B, T, HV, dtype=torch.float32, device=q.device) if USE_G else None - - h_flat = h.reshape(-1, h.shape[-2], h.shape[-1]) - dh_flat = dh.reshape(-1, dh.shape[-2], dh.shape[-1]) - hD1, hD2 = h_flat.shape[-2], h_flat.shape[-1] - dtype_str = {torch.float16: 'float16', torch.bfloat16: 'bfloat16', torch.float32: 'float32'}[q.dtype] - - # Cache key: B, H, HV, tile sizes, flags. T is dynamic (no recompilation for different seq lengths). - # Small head dims (< 64) cannot be warp-partitioned across 4 warps by TileLang's - # GEMM (tile too small); drop to 2 warps so the kernel still compiles. - num_warps = 4 if min(K, V) >= 64 else 2 - kernel = _build_kernel( - B, - H, - HV, - K, - V, - BT, - BK, - BV, - NK, - hD1, - hD2, - dtype_str, - USE_G, - USE_DW, - state_v_first, - IS_VARLEN, - num_warps=num_warps, - ) - - # Unused optional params still need shape-matching tensors for TileLang - g_kern = g if USE_G else q.new_empty(B, T, HV) - dw_kern = dw_out if USE_DW else q.new_empty(B, T, HV, K) - dv_kern = dv if USE_DW else q.new_empty(B, T, HV, V) - dg_kern = dg if USE_G else q.new_empty(NK, B, T, HV, dtype=torch.float32) - - if IS_VARLEN: - kernel(q, k, v, g_kern, h_flat, do, dh_flat, dq, dk, dw_kern, dv_kern, dg_kern, - cu_seqlens.int(), chunk_indices.int(), scale) - else: - kernel(q, k, v, g_kern, h_flat, do, dh_flat, dq, dk, dw_kern, dv_kern, dg_kern, scale) - - if dg is not None: - dg = dg.sum(0) - if H != HV: - G = HV // H - dq = dq.view(B, T, H, G, K).sum(3) - dk = dk.view(B, T, H, G, K).sum(3) - return dq, dk, dw_out, dg diff --git a/build/torch-xpu/ops/common/backends/tilelang/parallel_attn_bwd.py b/build/torch-xpu/ops/common/backends/tilelang/parallel_attn_bwd.py deleted file mode 100644 index efab180746a6a3558c464d73b0d3c9c607ac85e1..0000000000000000000000000000000000000000 --- a/build/torch-xpu/ops/common/backends/tilelang/parallel_attn_bwd.py +++ /dev/null @@ -1,300 +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 tilelang -import tilelang.language as T -import torch - -from .....ops.utils import prepare_chunk_indices -from .....ops.utils.constant import RCP_LN2 - - -@tilelang.jit(pass_configs={ - tilelang.PassConfigKey.TL_DISABLE_TMA_LOWER: True, - tilelang.PassConfigKey.TL_ENABLE_FAST_MATH: True, -}) -def _build_parallel_attn_bwd_kernel( - B, HQ, H, K, V, BT, sm_scale, log2e_scale, dtype_str, - USE_G=False, USE_WINDOW=False, WINDOW_SIZE=0, IS_VARLEN=False, num_warps=4, -): - """Build the TileLang JIT kernel for parallel attention backward. - - Note: B (batch size) is baked into the compiled kernel as a compile-time - constant because it appears in static tensor shape annotations. Varying B - at runtime will trigger a recompilation. Callers should either keep B fixed - per process or pad batches to a fixed size. - """ - - dtype_map = {'float16': T.float16, 'bfloat16': T.bfloat16, 'float32': T.float32} - _dtype = dtype_map[dtype_str] - accum_dtype = T.float32 - threads = num_warps * 32 - _BT = BT - _B, _HQ, _H = B, HQ, H - _K_orig, _V_orig = K, V - # Pad to next multiple of 16 for WGMMA inner-dim constraints (shared mem only) - _K = (K + 15) // 16 * 16 - _V = (V + 15) // 16 * 16 - _G = HQ // H - _USE_G = USE_G - _USE_WINDOW = USE_WINDOW - _W = WINDOW_SIZE - _sm_scale = sm_scale - _log2e_scale = log2e_scale - - T_d, NT_d, Ncu_d = T.dynamic("T, NT, Ncu") - - q_s = (_B, T_d, _HQ, _K_orig) - kv_s = (_B, T_d, _H, _K_orig) - vv_s = (_B, T_d, _H, _V_orig) - lse_s = (_B, T_d, _HQ) - dq_s = (_B, T_d, _HQ, _K_orig) - dkv_s = (_B, T_d, _H, _K_orig) - dvv_s = (_B, T_d, _H, _V_orig) - - @T.macro - def kernel_body(q, k, v, g, lse, delta, do, dq_out, dk_out, dv_out, dg, - i_b, i_h, i_t, t_s, T_seq, bos): - i_hkv = i_h // _G - - K_shared = T.alloc_shared((_BT, _K), _dtype) - V_shared = T.alloc_shared((_BT, _V), _dtype) - T.copy(k[i_b, t_s:t_s + _BT, i_hkv, :], K_shared) - T.copy(v[i_b, t_s:t_s + _BT, i_hkv, :], V_shared) - - if _USE_G: - g_k_shared = T.alloc_shared((_BT,), accum_dtype) - T.copy(g[i_b, t_s:t_s + _BT, i_h], g_k_shared) - - qkT = T.alloc_fragment((_BT, _BT), accum_dtype) - qkT_cast = T.alloc_fragment((_BT, _BT), _dtype) - dsT = T.alloc_fragment((_BT, _BT), accum_dtype) - dsT_cast = T.alloc_fragment((_BT, _BT), _dtype) - ds = T.alloc_fragment((_BT, _BT), _dtype) - - dv_local = T.alloc_fragment((_BT, _V), accum_dtype) - dk_local = T.alloc_fragment((_BT, _K), accum_dtype) - T.clear(dv_local) - T.clear(dk_local) - - if _USE_G: - dg_k_acc = T.alloc_fragment((_BT,), accum_dtype) - T.clear(dg_k_acc) - - # Allocate shared/fragment buffers outside the loop so they can be reused - # across query-tile iterations. - q_shared = T.alloc_shared((_BT, _K), _dtype) - do_shared = T.alloc_shared((_BT, _V), _dtype) - lse_shared = T.alloc_shared((_BT,), accum_dtype) - delta_shared = T.alloc_shared((_BT,), accum_dtype) - dsT_shared = T.alloc_shared((_BT, _BT), _dtype) - dq_local = T.alloc_fragment((_BT, _K), accum_dtype) - if _USE_G: - g_q_shared = T.alloc_shared((_BT,), accum_dtype) - f_dg_q = T.alloc_fragment((_BT,), accum_dtype) - f_dg_k = T.alloc_fragment((_BT,), accum_dtype) - - i_t_local = (t_s - bos) // _BT - loop_st = i_t_local - if _USE_WINDOW: - # For a K tile at positions [t_s, t_s+BT), the farthest query - # that can still attend (under window W) is at position - # (t_s+BT-1) + (W-1). Round up to the enclosing Q tile. - loop_ed_swa = i_t_local + 1 + (_W + _BT - 2) // _BT - loop_ed = T.min(T.ceildiv(T_seq, _BT), loop_ed_swa) - else: - loop_ed = T.ceildiv(T_seq, _BT) - - # Note: T.Pipelined is intentionally avoided here because combining it - # with standard-layout T.atomic_add into dq_out produces garbage values - # (observed as inf/575+ diffs) for non-power-of-2 dimensions. - for k_local in T.serial(loop_st, loop_ed): - q_t_s = bos + k_local * _BT - - T.copy(q[i_b, q_t_s:q_t_s + _BT, i_h, :], q_shared) - T.copy(do[i_b, q_t_s:q_t_s + _BT, i_h, :], do_shared) - T.copy(lse[i_b, q_t_s:q_t_s + _BT, i_h], lse_shared) - T.copy(delta[i_b, q_t_s:q_t_s + _BT, i_h], delta_shared) - - if _USE_G: - T.copy(g[i_b, q_t_s:q_t_s + _BT, i_h], g_q_shared) - - T.clear(qkT) - T.gemm(K_shared, q_shared, qkT, transpose_B=True, policy=T.GemmWarpPolicy.FullRow) - - for i, j in T.Parallel(_BT, _BT): - if _USE_G: - qkT[i, j] = T.exp2(qkT[i, j] * _log2e_scale + g_q_shared[j] - g_k_shared[i] - lse_shared[j]) - else: - qkT[i, j] = T.exp2(qkT[i, j] * _log2e_scale - lse_shared[j]) - - # causal + boundary (+ optional sliding window) mask - for i, j in T.Parallel(_BT, _BT): - key_pos = t_s + i - q_pos = q_t_s + j - causal = (key_pos <= q_pos) & (key_pos < T_seq + bos) & (q_pos < T_seq + bos) - if _USE_WINDOW: - in_window = (q_pos - key_pos < _W) - mask = causal & in_window - else: - mask = causal - qkT[i, j] = T.if_then_else(mask, qkT[i, j], T.cast(0, accum_dtype)) - - # dv += p @ do - T.copy(qkT, qkT_cast) - T.gemm(qkT_cast, do_shared, dv_local, policy=T.GemmWarpPolicy.FullRow) - - # ds = p * (dp - delta) - T.clear(dsT) - T.gemm(V_shared, do_shared, dsT, transpose_B=True, policy=T.GemmWarpPolicy.FullRow) - - for i, j in T.Parallel(_BT, _BT): - dsT[i, j] = qkT[i, j] * (dsT[i, j] - delta_shared[j]) - for i, j in T.Parallel(_BT, _BT): - dsT_cast[i, j] = dsT[i, j] * _sm_scale - - # dk += ds^T @ q - T.copy(dsT_cast, ds) - T.gemm(ds, q_shared, dk_local, policy=T.GemmWarpPolicy.FullRow) - - # dq += ds @ k (atomic_add into standard-layout output) - T.copy(dsT_cast, dsT_shared) - T.clear(dq_local) - T.gemm(dsT_shared, K_shared, dq_local, transpose_A=True, policy=T.GemmWarpPolicy.FullRow) - for i, d in T.Parallel(_BT, _K_orig): - if q_t_s + i < T_seq + bos: - T.atomic_add(dq_out[i_b, q_t_s + i, i_h, d], dq_local[i, d]) - - if _USE_G: - # dg for query positions (positive) - T.reduce_sum(dsT, f_dg_q, dim=0) - for j in T.Parallel(_BT): - if q_t_s + j < T_seq + bos: - T.atomic_add(dg[i_b, q_t_s + j, i_h], f_dg_q[j]) - - # accumulate dg for key positions (negative) - T.reduce_sum(dsT, f_dg_k, dim=1) - for i in T.Parallel(_BT): - dg_k_acc[i] += f_dg_k[i] - - # store dv and dk via element-wise atomic_add with explicit boundary checks. - # atomic_add is required because in GQA multiple query-head blocks map to - # the same KV head and accumulate into the same dk/dv rows. - for i, d in T.Parallel(_BT, _V_orig): - if t_s + i < T_seq + bos: - T.atomic_add(dv_out[i_b, t_s + i, i_hkv, d], dv_local[i, d]) - for i, d in T.Parallel(_BT, _K_orig): - if t_s + i < T_seq + bos: - T.atomic_add(dk_out[i_b, t_s + i, i_hkv, d], dk_local[i, d]) - - if _USE_G: - for i in T.Parallel(_BT): - if t_s + i < T_seq + bos: - T.atomic_add(dg[i_b, t_s + i, i_h], -dg_k_acc[i]) - - if IS_VARLEN: - @T.prim_func - def kernel( - q: T.Tensor(q_s, _dtype), k: T.Tensor(kv_s, _dtype), - v: T.Tensor(vv_s, _dtype), g: T.Tensor(lse_s, accum_dtype), - lse: T.Tensor(lse_s, accum_dtype), delta: T.Tensor(lse_s, accum_dtype), - do: T.Tensor(q_s, _dtype), dq_out: T.Tensor(dq_s, accum_dtype), - dk_out: T.Tensor(dkv_s, accum_dtype), dv_out: T.Tensor(dvv_s, accum_dtype), - dg: T.Tensor(lse_s, accum_dtype), - cu_seqlens: T.Tensor((Ncu_d,), T.int32), - chunk_indices: T.Tensor((NT_d, 2), T.int32), - ): - with T.Kernel(_HQ, NT_d, 1, threads=threads) as (i_h, i_t, _): - i_n = chunk_indices[i_t, 0] - i_t_local = chunk_indices[i_t, 1] - bos = cu_seqlens[i_n] - T_seq = cu_seqlens[i_n + 1] - bos - t_s = bos + i_t_local * _BT - kernel_body(q, k, v, g, lse, delta, do, dq_out, dk_out, dv_out, dg, - 0, i_h, i_t, t_s, T_seq, bos) - else: - @T.prim_func - def kernel( - q: T.Tensor(q_s, _dtype), k: T.Tensor(kv_s, _dtype), - v: T.Tensor(vv_s, _dtype), g: T.Tensor(lse_s, accum_dtype), - lse: T.Tensor(lse_s, accum_dtype), delta: T.Tensor(lse_s, accum_dtype), - do: T.Tensor(q_s, _dtype), dq_out: T.Tensor(dq_s, accum_dtype), - dk_out: T.Tensor(dkv_s, accum_dtype), dv_out: T.Tensor(dvv_s, accum_dtype), - dg: T.Tensor(lse_s, accum_dtype), - ): - with T.Kernel(_HQ, T.ceildiv(T_d, _BT), _B, threads=threads) as (i_h, i_t, i_b): - t_s = i_t * _BT - kernel_body(q, k, v, g, lse, delta, do, dq_out, dk_out, dv_out, dg, - i_b, i_h, i_t, t_s, T_d, 0) - - return kernel - - -def parallel_attn_bwd_tilelang( - q, k, v, o, g_cumsum, lse, do, - sink_bias=None, scale=None, window_size=None, - chunk_size=128, cu_seqlens=None, chunk_indices=None, -): - B, T, HQ, K = q.shape - H = k.shape[2] - V = v.shape[-1] - BT = 64 - sm_scale = scale if scale is not None else K ** -0.5 - log2e_scale = sm_scale * RCP_LN2 - - USE_G = g_cumsum is not None - USE_WINDOW = window_size is not None - IS_VARLEN = cu_seqlens is not None - - if chunk_indices is None and IS_VARLEN: - chunk_indices = prepare_chunk_indices(cu_seqlens, BT) - - # reuse Triton preprocess for delta (works with unpadded tensors) - from .....ops.attn.parallel import parallel_attn_bwd_preprocess - delta = parallel_attn_bwd_preprocess(o, do) - - # No padding on the last dim: TileLang T.copy auto-zero-fills OOB columns - # when loading to shared memory, and auto-skips OOB writes. dq/dk/dv all - # need float32 buffers because GQA has multiple query-head blocks writing - # to the same KV head (dk/dv) and query rows receive contributions from - # multiple key tiles (dq). - dq = torch.zeros(B, T, HQ, K, dtype=torch.float32, device=q.device) - dk = torch.zeros(B, T, H, K, dtype=torch.float32, device=k.device) - dv = torch.zeros(B, T, H, V, dtype=torch.float32, device=v.device) - dg = torch.zeros(B, T, HQ, dtype=torch.float32, device=q.device) if USE_G else None - - g_kern = g_cumsum.float() if USE_G else torch.zeros(B, T, HQ, dtype=torch.float32, device=q.device) - dg_kern = dg if USE_G else torch.zeros(B, T, HQ, dtype=torch.float32, device=q.device) - - dtype_str = {torch.float16: 'float16', torch.bfloat16: 'bfloat16', torch.float32: 'float32'}.get(q.dtype) - if dtype_str is None: - raise ValueError(f"Unsupported dtype {q.dtype} for TileLang backend") - - kernel = _build_parallel_attn_bwd_kernel( - B, HQ, H, K, V, BT, sm_scale, log2e_scale, dtype_str, - USE_G=USE_G, USE_WINDOW=USE_WINDOW, WINDOW_SIZE=window_size or 0, - IS_VARLEN=IS_VARLEN, - ) - - if IS_VARLEN: - kernel(q, k, v, g_kern, lse, delta, do, dq, dk, dv, dg_kern, - cu_seqlens.int(), chunk_indices.int()) - else: - kernel(q, k, v, g_kern, lse, delta, do, dq, dk, dv, dg_kern) - - # sink_bias gradient is derived from the saved lse (which already includes - # the sink mass in the softmax normalizer) and delta — no kernel change - # needed. Matches the Triton reference. - dsink_bias = None - if sink_bias is not None: - p_sink = torch.exp2(sink_bias.float()[None, None, :] - lse.float()) - dsink_bias = -(p_sink * delta.float()).sum((0, 1)) - - dq = dq.to(q.dtype) - dk = dk.to(k.dtype) - dv = dv.to(v.dtype) - return dq, dk, dv, dg, dsink_bias diff --git a/build/torch-xpu/ops/common/backends/tilelang/parallel_attn_fwd.py b/build/torch-xpu/ops/common/backends/tilelang/parallel_attn_fwd.py deleted file mode 100644 index 9ff1bbe7978d9275513b9ccfef584f809c18e17b..0000000000000000000000000000000000000000 --- a/build/torch-xpu/ops/common/backends/tilelang/parallel_attn_fwd.py +++ /dev/null @@ -1,298 +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 - -"""TileLang forward kernel for parallel (causal) attention. - -Supports: GQA, gating (g_cumsum), sliding-window attention, sink_bias, -variable-length (cu_seqlens). Output format matches the Triton reference -at `fla.ops.attn.parallel.parallel_attn_fwd` so LSE is directly consumable -by either backend during backward. -""" - -import tilelang -import tilelang.language as T -import torch - -from .....ops.utils import prepare_chunk_indices -from .....ops.utils.constant import RCP_LN2 - - -def _pick_tile(K: int, V: int, is_hopper_plus: bool) -> tuple[int, int, int]: - """Return (BT, BS, num_warps) tuned for head dim.""" - if is_hopper_plus: - if max(K, V) <= 64: - return 128, 64, 8 - if max(K, V) <= 128: - return 64, 64, 4 - return 64, 64, 4 - # Ampere / pre-Hopper - if max(K, V) <= 64: - return 128, 64, 4 - if max(K, V) <= 128: - return 128, 64, 4 - return 64, 32, 4 - - -@tilelang.jit(pass_configs={ - tilelang.PassConfigKey.TL_DISABLE_TMA_LOWER: True, - tilelang.PassConfigKey.TL_ENABLE_FAST_MATH: True, -}) -def _build_parallel_attn_fwd_kernel( - B, HQ, H, K, V, BT, BS, sm_scale, log2e_scale, dtype_str, - USE_G=False, USE_WINDOW=False, WINDOW_SIZE=0, USE_SINK=False, - IS_VARLEN=False, num_warps=8, num_stages=1, -): - """Build the TileLang JIT kernel for parallel attention forward. - - Note: B (batch size) is baked into the compiled kernel as a compile-time - constant because it appears in static tensor shape annotations. Varying - B at runtime triggers a recompilation. Callers should keep B fixed per - process or pad batches to a fixed size. - """ - dtype_map = {'float16': T.float16, 'bfloat16': T.bfloat16, 'float32': T.float32} - _dtype = dtype_map[dtype_str] - accum_dtype = T.float32 - threads = num_warps * 32 - _BT, _BS = BT, BS - _B, _HQ, _H = B, HQ, H - _K_orig, _V_orig = K, V - # Pad shared-mem last dim to a multiple of 16 for WGMMA/MMA inner-dim. - # Same approach as the bwd kernel — tensors keep their original shape; - # only the shared buffers are padded. - _K = (K + 15) // 16 * 16 - _V = (V + 15) // 16 * 16 - _G = HQ // H - _USE_G = USE_G - _USE_WINDOW = USE_WINDOW - _USE_SINK = USE_SINK - _W = WINDOW_SIZE - _sm_scale = sm_scale - _log2e_scale = log2e_scale - - T_d, NT_d, Ncu_d = T.dynamic("T, NT, Ncu") - - q_s = (_B, T_d, _HQ, _K_orig) - kv_s = (_B, T_d, _H, _K_orig) - vv_s = (_B, T_d, _H, _V_orig) - o_s = (_B, T_d, _HQ, _V_orig) - lse_s = (_B, T_d, _HQ) - sink_s = (_HQ,) - - @T.macro - def kernel_body(q, k, v, g, sink_bias, o, lse, - i_b, i_h, i_t, t_s, T_seq, bos): - i_hkv = i_h // _G - - Q_shared = T.alloc_shared((_BT, _K), _dtype) - K_shared = T.alloc_shared((_BS, _K), _dtype) - V_shared = T.alloc_shared((_BS, _V), _dtype) - - acc_s = T.alloc_fragment((_BT, _BS), accum_dtype) - acc_s_cast = T.alloc_fragment((_BT, _BS), _dtype) - acc_o = T.alloc_fragment((_BT, _V), accum_dtype) - - m_f = T.alloc_fragment((_BT,), accum_dtype) - m_prev = T.alloc_fragment((_BT,), accum_dtype) - l_f = T.alloc_fragment((_BT,), accum_dtype) - p_sum = T.alloc_fragment((_BT,), accum_dtype) - scale_row = T.alloc_fragment((_BT,), accum_dtype) - - if _USE_G: - g_q_shared = T.alloc_shared((_BT,), accum_dtype) - g_k_shared = T.alloc_shared((_BS,), accum_dtype) - T.copy(g[i_b, t_s:t_s + _BT, i_h], g_q_shared) - - T.copy(q[i_b, t_s:t_s + _BT, i_h, :], Q_shared) - T.fill(acc_o, 0) - T.fill(l_f, 0) - T.fill(m_f, -T.infinity(accum_dtype)) - - i_t_local = (t_s - bos) // _BT - if _USE_WINDOW: - # earliest key the last query in this tile can see - # (q_pos - k_pos < W) => k_pos > q_pos - W. - # Round down to the enclosing K tile start. - s0_raw = i_t_local * _BT - _W + 1 - s0 = T.max(s0_raw, 0) - loop_st = (s0 // _BS) - else: - loop_st = 0 - loop_ed_raw = T.ceildiv((i_t_local + 1) * _BT, _BS) - loop_ed = T.min(loop_ed_raw, T.ceildiv(T_seq, _BS)) - - for k_idx in T.Pipelined(loop_st, loop_ed, num_stages=num_stages): - k_t_s = bos + k_idx * _BS - - T.copy(k[i_b, k_t_s:k_t_s + _BS, i_hkv, :], K_shared) - T.copy(v[i_b, k_t_s:k_t_s + _BS, i_hkv, :], V_shared) - if _USE_G: - T.copy(g[i_b, k_t_s:k_t_s + _BS, i_h], g_k_shared) - - T.clear(acc_s) - T.gemm(Q_shared, K_shared, acc_s, - transpose_B=True, policy=T.GemmWarpPolicy.FullRow) - - # Apply mask (causal + boundary + optional SWA), add log-gates, - # and convert scores to the log2 domain. Masked entries get - # large-negative so exp2(...) -> 0. - for i, j in T.Parallel(_BT, _BS): - q_pos = t_s + i - key_pos = k_t_s + j - causal = (key_pos <= q_pos) & (key_pos < T_seq + bos) & (q_pos < T_seq + bos) - if _USE_WINDOW: - mask = causal & (q_pos - key_pos < _W) - else: - mask = causal - if _USE_G: - s = acc_s[i, j] * _log2e_scale + g_q_shared[i] - g_k_shared[j] - else: - s = acc_s[i, j] * _log2e_scale - acc_s[i, j] = T.if_then_else(mask, s, -T.infinity(accum_dtype)) - - # Online softmax update. - T.copy(m_f, m_prev) - T.reduce_max(acc_s, m_f, dim=1, clear=False) - for i in T.Parallel(_BT): - m_f[i] = T.max(m_f[i], m_prev[i]) - # Guard against all-masked rows: use a finite pivot so we - # don't produce NaN via (-inf) - (-inf) anywhere downstream. - m_stable = T.alloc_fragment((_BT,), accum_dtype) - for i in T.Parallel(_BT): - m_stable[i] = T.if_then_else( - m_f[i] == -T.infinity(accum_dtype), 0., m_f[i]) - for i in T.Parallel(_BT): - scale_row[i] = T.if_then_else( - m_prev[i] == -T.infinity(accum_dtype), 0., - T.exp2(m_prev[i] - m_stable[i])) - for i, j in T.Parallel(_BT, _V): - acc_o[i, j] *= scale_row[i] - for i, j in T.Parallel(_BT, _BS): - acc_s[i, j] = T.exp2(acc_s[i, j] - m_stable[i]) - T.reduce_sum(acc_s, p_sum, dim=1) - for i in T.Parallel(_BT): - l_f[i] = l_f[i] * scale_row[i] + p_sum[i] - - T.copy(acc_s, acc_s_cast) - T.gemm(acc_s_cast, V_shared, acc_o, - policy=T.GemmWarpPolicy.FullRow) - - # sink-bias contributes to the normalizer but not to the value matmul - if _USE_SINK: - # If a row has no valid keys yet, m_f is still -inf — pin it to 0 - # so `sink_bias - m_f` is finite. - for i in T.Parallel(_BT): - m_f[i] = T.if_then_else( - m_f[i] == -T.infinity(accum_dtype), 0., m_f[i]) - for i in T.Parallel(_BT): - l_f[i] += T.exp2(sink_bias[i_h] - m_f[i]) - - # finalize O and LSE - O_shared = T.alloc_shared((_BT, _V), _dtype) - lse_shared = T.alloc_shared((_BT,), accum_dtype) - for i, d in T.Parallel(_BT, _V): - acc_o[i, d] = acc_o[i, d] / l_f[i] - T.copy(acc_o, O_shared) - for i in T.Parallel(_BT): - lse_shared[i] = m_f[i] + T.log2(l_f[i]) - for i, d in T.Parallel(_BT, _V_orig): - if t_s + i < T_seq + bos: - o[i_b, t_s + i, i_h, d] = O_shared[i, d] - for i in T.Parallel(_BT): - if t_s + i < T_seq + bos: - lse[i_b, t_s + i, i_h] = lse_shared[i] - - if IS_VARLEN: - @T.prim_func - def kernel( - q: T.Tensor(q_s, _dtype), k: T.Tensor(kv_s, _dtype), - v: T.Tensor(vv_s, _dtype), g: T.Tensor(lse_s, accum_dtype), - sink_bias: T.Tensor(sink_s, accum_dtype), - o: T.Tensor(o_s, _dtype), lse: T.Tensor(lse_s, accum_dtype), - cu_seqlens: T.Tensor((Ncu_d,), T.int32), - chunk_indices: T.Tensor((NT_d, 2), T.int32), - ): - with T.Kernel(_HQ, NT_d, 1, threads=threads) as (i_h, i_t, _): - i_n = chunk_indices[i_t, 0] - i_t_local = chunk_indices[i_t, 1] - bos = cu_seqlens[i_n] - eos = cu_seqlens[i_n + 1] - T_seq = eos - bos - t_s = bos + i_t_local * _BT - kernel_body(q, k, v, g, sink_bias, o, lse, - 0, i_h, i_t, t_s, T_seq, bos) - else: - @T.prim_func - def kernel( - q: T.Tensor(q_s, _dtype), k: T.Tensor(kv_s, _dtype), - v: T.Tensor(vv_s, _dtype), g: T.Tensor(lse_s, accum_dtype), - sink_bias: T.Tensor(sink_s, accum_dtype), - o: T.Tensor(o_s, _dtype), lse: T.Tensor(lse_s, accum_dtype), - ): - with T.Kernel(_HQ, T.ceildiv(T_d, _BT), _B, threads=threads) as (i_h, i_t, i_b): - t_s = i_t * _BT - kernel_body(q, k, v, g, sink_bias, o, lse, - i_b, i_h, i_t, t_s, T_d, 0) - - return kernel - - -def parallel_attn_fwd_tilelang( - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - g_cumsum: torch.Tensor | None, - sink_bias: torch.Tensor | None, - scale: float, - window_size: int | None = None, - cu_seqlens: torch.LongTensor | None = None, - chunk_indices: torch.LongTensor | None = None, -): - B, T, HQ, K = q.shape - H = k.shape[2] - V = v.shape[-1] - - USE_G = g_cumsum is not None - USE_WINDOW = window_size is not None - USE_SINK = sink_bias is not None - IS_VARLEN = cu_seqlens is not None - - is_hopper_plus = torch.cuda.get_device_capability()[0] >= 9 - BT, BS, num_warps = _pick_tile(K, V, is_hopper_plus) - - if chunk_indices is None and IS_VARLEN: - chunk_indices = prepare_chunk_indices(cu_seqlens, BT) - - sm_scale = scale if scale is not None else K ** -0.5 - log2e_scale = sm_scale * RCP_LN2 - - dtype_str = {torch.float16: 'float16', torch.bfloat16: 'bfloat16', - torch.float32: 'float32'}.get(q.dtype) - if dtype_str is None: - raise ValueError(f"Unsupported dtype {q.dtype} for TileLang backend") - - o = torch.empty(B, T, HQ, V, dtype=q.dtype, device=q.device) - lse = torch.empty(B, T, HQ, dtype=torch.float32, device=q.device) - - g_kern = g_cumsum.float() if USE_G else torch.zeros( - B, T, HQ, dtype=torch.float32, device=q.device) - sink_kern = sink_bias.float() if USE_SINK else torch.zeros( - HQ, dtype=torch.float32, device=q.device) - - kernel = _build_parallel_attn_fwd_kernel( - B, HQ, H, K, V, BT, BS, sm_scale, log2e_scale, dtype_str, - USE_G=USE_G, USE_WINDOW=USE_WINDOW, WINDOW_SIZE=window_size or 0, - USE_SINK=USE_SINK, IS_VARLEN=IS_VARLEN, num_warps=num_warps, - num_stages=1 if USE_WINDOW else 2, - ) - - if IS_VARLEN: - kernel(q, k, v, g_kern, sink_kern, o, lse, - cu_seqlens.int(), chunk_indices.int()) - else: - kernel(q, k, v, g_kern, sink_kern, o, lse) - - return o, lse diff --git a/build/torch-xpu/ops/common/chunk_delta_h.py b/build/torch-xpu/ops/common/chunk_delta_h.py deleted file mode 100644 index 89a556630064549d9d4156d559ea3feed8b2832b..0000000000000000000000000000000000000000 --- a/build/torch-xpu/ops/common/chunk_delta_h.py +++ /dev/null @@ -1,734 +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 -import triton.language as tl - -from ...ops.backends import dispatch -from ...ops.utils import prepare_chunk_indices, prepare_chunk_offsets -from ...ops.utils.cache import fla_cache_autotune -from ...ops.utils.op import exp2 -from ...utils import ( - IS_INTEL, - IS_NVIDIA_BLACKWELL, - IS_NVIDIA_HOPPER, - autotune_cache_kwargs, - check_shared_mem, -) - -NUM_WARPS = [2, 4] if IS_NVIDIA_HOPPER else [2, 4, 8, 16] - -# TODO: Triton mainline fixes a Blackwell tl.dot recurrence race. -# Keep this kernel on num_warps=2 for Blackwell until Triton 3.8 is released -# and we re-validate the wider config space. -# Intel needs more warps than NVIDIA here: 8 warps is ~1.5x faster than the best -# config reachable under the [2, 4] cap. -if IS_NVIDIA_BLACKWELL: - GATED_DELTA_RULE_FWD_H_NUM_WARPS = [2] -elif IS_INTEL: - GATED_DELTA_RULE_FWD_H_NUM_WARPS = [2, 4, 8, 16] -else: - GATED_DELTA_RULE_FWD_H_NUM_WARPS = [2, 4] - - -@triton.heuristics({ - 'USE_G': lambda args: args['g'] is not None, - 'USE_GK': lambda args: args['gk'] is not None, - 'USE_INITIAL_STATE': lambda args: args['h0'] is not None, - 'STORE_FINAL_STATE': lambda args: args['ht'] is not None, - 'SAVE_NEW_VALUE': lambda args: args['v_new'] is not None, - 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None, -}) -@fla_cache_autotune( - configs=[ - triton.Config({'BV': BV}, num_warps=num_warps, num_stages=num_stages) - for num_warps in GATED_DELTA_RULE_FWD_H_NUM_WARPS - for num_stages in ([2, 3, 4] if check_shared_mem('ampere') else [2, 1]) - for BV in ([32, 64] if check_shared_mem('ada') else [32]) - ], - key=['H', 'HV', 'K', 'V', 'BT', 'STATE_V_FIRST'], - **autotune_cache_kwargs, -) -@triton.jit(do_not_specialize=['T']) -def chunk_gated_delta_rule_fwd_kernel_h_blockdim64( - k, - v, - w, - v_new, - g, - gk, - h, - h0, - ht, - cu_seqlens, - chunk_offsets, - T, - H: tl.constexpr, - HV: tl.constexpr, - K: tl.constexpr, - V: tl.constexpr, - BT: tl.constexpr, - BV: tl.constexpr, - USE_G: tl.constexpr, - USE_GK: tl.constexpr, - USE_INITIAL_STATE: tl.constexpr, - STORE_FINAL_STATE: tl.constexpr, - SAVE_NEW_VALUE: tl.constexpr, - STATE_V_FIRST: tl.constexpr, - IS_VARLEN: tl.constexpr, -): - i_v, i_nh = tl.program_id(0), tl.program_id(1) - i_n, i_h = i_nh // HV, i_nh % HV - if IS_VARLEN: - bos, eos = tl.load(cu_seqlens + i_n).to(tl.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) - T = eos - bos - NT = tl.cdiv(T, BT) - boh = tl.load(chunk_offsets + i_n).to(tl.int32) - else: - bos, eos = i_n * T, i_n * T + T - NT = tl.cdiv(T, BT) - boh = i_n * NT - - if STATE_V_FIRST: - b_h1 = tl.zeros([BV, 64], dtype=tl.float32) - if K > 64: - b_h2 = tl.zeros([BV, 64], dtype=tl.float32) - if K > 128: - b_h3 = tl.zeros([BV, 64], dtype=tl.float32) - if K > 192: - b_h4 = tl.zeros([BV, 64], dtype=tl.float32) - else: - b_h1 = tl.zeros([64, BV], dtype=tl.float32) - if K > 64: - b_h2 = tl.zeros([64, BV], dtype=tl.float32) - if K > 128: - b_h3 = tl.zeros([64, BV], dtype=tl.float32) - if K > 192: - b_h4 = tl.zeros([64, BV], dtype=tl.float32) - - # calculate offset - h += (boh * HV + i_h).to(tl.int64) * K*V - v += (bos * HV + i_h).to(tl.int64) * V - k += (bos * H + i_h // (HV // H)).to(tl.int64) * K - w += (bos * HV + i_h).to(tl.int64) * K - if SAVE_NEW_VALUE: - v_new += (bos * HV + i_h).to(tl.int64) * V - - if USE_INITIAL_STATE: - h0 = h0 + i_nh * K*V - if STORE_FINAL_STATE: - ht = ht + i_nh * K*V - - # load initial state - if USE_INITIAL_STATE: - if STATE_V_FIRST: - p_h0_1 = tl.make_block_ptr(h0, (V, K), (K, 1), (i_v * BV, 0), (BV, 64), (1, 0)) - else: - p_h0_1 = tl.make_block_ptr(h0, (K, V), (V, 1), (0, i_v * BV), (64, BV), (1, 0)) - b_h1 += tl.load(p_h0_1, boundary_check=(0, 1)).to(tl.float32) - if K > 64: - if STATE_V_FIRST: - p_h0_2 = tl.make_block_ptr(h0, (V, K), (K, 1), (i_v * BV, 64), (BV, 64), (1, 0)) - else: - p_h0_2 = tl.make_block_ptr(h0, (K, V), (V, 1), (64, i_v * BV), (64, BV), (1, 0)) - b_h2 += tl.load(p_h0_2, boundary_check=(0, 1)).to(tl.float32) - if K > 128: - if STATE_V_FIRST: - p_h0_3 = tl.make_block_ptr(h0, (V, K), (K, 1), (i_v * BV, 128), (BV, 64), (1, 0)) - else: - p_h0_3 = tl.make_block_ptr(h0, (K, V), (V, 1), (128, i_v * BV), (64, BV), (1, 0)) - b_h3 += tl.load(p_h0_3, boundary_check=(0, 1)).to(tl.float32) - if K > 192: - if STATE_V_FIRST: - p_h0_4 = tl.make_block_ptr(h0, (V, K), (K, 1), (i_v * BV, 192), (BV, 64), (1, 0)) - else: - p_h0_4 = tl.make_block_ptr(h0, (K, V), (V, 1), (192, i_v * BV), (64, BV), (1, 0)) - b_h4 += tl.load(p_h0_4, boundary_check=(0, 1)).to(tl.float32) - - # main recurrence - for i_t in range(NT): - i_t_int64 = i_t.to(tl.int64) - if STATE_V_FIRST: - p_h1 = tl.make_block_ptr(h + i_t_int64 * HV*K*V, (V, K), (K, 1), (i_v * BV, 0), (BV, 64), (1, 0)) - else: - p_h1 = tl.make_block_ptr(h + i_t_int64 * HV*K*V, (K, V), (V, 1), (0, i_v * BV), (64, BV), (1, 0)) - tl.store(p_h1, b_h1.to(p_h1.dtype.element_ty), boundary_check=(0, 1)) - if K > 64: - if STATE_V_FIRST: - p_h2 = tl.make_block_ptr(h + i_t_int64 * HV*K*V, (V, K), (K, 1), (i_v * BV, 64), (BV, 64), (1, 0)) - else: - p_h2 = tl.make_block_ptr(h + i_t_int64 * HV*K*V, (K, V), (V, 1), (64, i_v * BV), (64, BV), (1, 0)) - tl.store(p_h2, b_h2.to(p_h2.dtype.element_ty), boundary_check=(0, 1)) - if K > 128: - if STATE_V_FIRST: - p_h3 = tl.make_block_ptr(h + i_t_int64 * HV*K*V, (V, K), (K, 1), (i_v * BV, 128), (BV, 64), (1, 0)) - else: - p_h3 = tl.make_block_ptr(h + i_t_int64 * HV*K*V, (K, V), (V, 1), (128, i_v * BV), (64, BV), (1, 0)) - tl.store(p_h3, b_h3.to(p_h3.dtype.element_ty), boundary_check=(0, 1)) - if K > 192: - if STATE_V_FIRST: - p_h4 = tl.make_block_ptr(h + i_t_int64 * HV*K*V, (V, K), (K, 1), (i_v * BV, 192), (BV, 64), (1, 0)) - else: - p_h4 = tl.make_block_ptr(h + i_t_int64 * HV*K*V, (K, V), (V, 1), (192, i_v * BV), (64, BV), (1, 0)) - tl.store(p_h4, b_h4.to(p_h4.dtype.element_ty), boundary_check=(0, 1)) - - p_w = tl.make_block_ptr(w, (T, K), (HV*K, 1), (i_t * BT, 0), (BT, 64), (1, 0)) - b_w = tl.load(p_w, boundary_check=(0, 1)) - if STATE_V_FIRST: - b_v = tl.dot(b_w, tl.trans(b_h1).to(b_w.dtype)) - else: - b_v = tl.dot(b_w, b_h1.to(b_w.dtype)) - if K > 64: - p_w = tl.make_block_ptr(w, (T, K), (HV*K, 1), (i_t * BT, 64), (BT, 64), (1, 0)) - b_w = tl.load(p_w, boundary_check=(0, 1)) - if STATE_V_FIRST: - b_v += tl.dot(b_w, tl.trans(b_h2).to(b_w.dtype)) - else: - b_v += tl.dot(b_w, b_h2.to(b_w.dtype)) - if K > 128: - p_w = tl.make_block_ptr(w, (T, K), (HV*K, 1), (i_t * BT, 128), (BT, 64), (1, 0)) - b_w = tl.load(p_w, boundary_check=(0, 1)) - if STATE_V_FIRST: - b_v += tl.dot(b_w, tl.trans(b_h3).to(b_w.dtype)) - else: - b_v += tl.dot(b_w, b_h3.to(b_w.dtype)) - if K > 192: - p_w = tl.make_block_ptr(w, (T, K), (HV*K, 1), (i_t * BT, 192), (BT, 64), (1, 0)) - b_w = tl.load(p_w, boundary_check=(0, 1)) - if STATE_V_FIRST: - b_v += tl.dot(b_w, tl.trans(b_h4).to(b_w.dtype)) - else: - b_v += tl.dot(b_w, b_h4.to(b_w.dtype)) - p_v = tl.make_block_ptr(v, (T, V), (HV*V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) - b_v = tl.load(p_v, boundary_check=(0, 1)) - b_v - - if SAVE_NEW_VALUE: - p_v = tl.make_block_ptr(v_new, (T, V), (HV*V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) - tl.store(p_v, b_v.to(p_v.dtype.element_ty), boundary_check=(0, 1)) - - last_idx = min((i_t + 1) * BT, T) - 1 - if USE_G: - m_t = (i_t * BT + tl.arange(0, BT)) < T - b_g_last = tl.load(g + (bos * HV + last_idx * HV + i_h).to(tl.int64)).to(tl.float32) - p_g = tl.make_block_ptr(g + (bos * HV + i_h).to(tl.int64), (T,), (HV,), (i_t * BT,), (BT,), (0,)) - b_g = tl.load(p_g, boundary_check=(0,)).to(tl.float32) - b_v = b_v * tl.where(m_t, exp2(b_g_last - b_g), 0)[:, None] - b_g_last = exp2(b_g_last) - b_h1 *= b_g_last - if K > 64: - b_h2 *= b_g_last - if K > 128: - b_h3 *= b_g_last - if K > 192: - b_h4 *= b_g_last - - if USE_GK: - o_k1 = tl.arange(0, 64) - b_gk_last1 = tl.load(gk + (bos + last_idx) * HV*K + i_h * K + o_k1, mask=(o_k1 < K), other=0.).to(tl.float32) - if STATE_V_FIRST: - b_h1 *= exp2(b_gk_last1)[None, :] - else: - b_h1 *= exp2(b_gk_last1)[:, None] - if K > 64: - o_k2 = 64 + o_k1 - b_gk_last2 = tl.load(gk + (bos + last_idx) * HV*K + i_h * K + o_k2, mask=(o_k2 < K), other=0.).to(tl.float32) - if STATE_V_FIRST: - b_h2 *= exp2(b_gk_last2)[None, :] - else: - b_h2 *= exp2(b_gk_last2)[:, None] - if K > 128: - o_k3 = 128 + o_k1 - b_gk_last3 = tl.load(gk + (bos + last_idx) * HV*K + i_h * K + o_k3, mask=(o_k3 < K), other=0.).to(tl.float32) - if STATE_V_FIRST: - b_h3 *= exp2(b_gk_last3)[None, :] - else: - b_h3 *= exp2(b_gk_last3)[:, None] - if K > 192: - o_k4 = 192 + o_k1 - b_gk_last4 = tl.load(gk + (bos + last_idx) * HV*K + i_h * K + o_k4, mask=(o_k4 < K), other=0.).to(tl.float32) - if STATE_V_FIRST: - b_h4 *= exp2(b_gk_last4)[None, :] - else: - b_h4 *= exp2(b_gk_last4)[:, None] - b_v = b_v.to(k.dtype.element_ty) - - p_k = tl.make_block_ptr(k, (K, T), (1, H*K), (0, i_t * BT), (64, BT), (0, 1)) - b_k = tl.load(p_k, boundary_check=(0, 1)) - if STATE_V_FIRST: - b_h1 += tl.trans(tl.dot(b_k, b_v)) - else: - b_h1 += tl.dot(b_k, b_v) - if K > 64: - p_k = tl.make_block_ptr(k, (K, T), (1, H*K), (64, i_t * BT), (64, BT), (0, 1)) - b_k = tl.load(p_k, boundary_check=(0, 1)) - if STATE_V_FIRST: - b_h2 += tl.trans(tl.dot(b_k, b_v)) - else: - b_h2 += tl.dot(b_k, b_v) - if K > 128: - p_k = tl.make_block_ptr(k, (K, T), (1, H*K), (128, i_t * BT), (64, BT), (0, 1)) - b_k = tl.load(p_k, boundary_check=(0, 1)) - if STATE_V_FIRST: - b_h3 += tl.trans(tl.dot(b_k, b_v)) - else: - b_h3 += tl.dot(b_k, b_v) - if K > 192: - p_k = tl.make_block_ptr(k, (K, T), (1, H*K), (192, i_t * BT), (64, BT), (0, 1)) - b_k = tl.load(p_k, boundary_check=(0, 1)) - if STATE_V_FIRST: - b_h4 += tl.trans(tl.dot(b_k, b_v)) - else: - b_h4 += tl.dot(b_k, b_v) - - if STORE_FINAL_STATE: - if STATE_V_FIRST: - p_ht = tl.make_block_ptr(ht, (V, K), (K, 1), (i_v * BV, 0), (BV, 64), (1, 0)) - else: - p_ht = tl.make_block_ptr(ht, (K, V), (V, 1), (0, i_v * BV), (64, BV), (1, 0)) - tl.store(p_ht, b_h1.to(p_ht.dtype.element_ty), boundary_check=(0, 1)) - if K > 64: - if STATE_V_FIRST: - p_ht = tl.make_block_ptr(ht, (V, K), (K, 1), (i_v * BV, 64), (BV, 64), (1, 0)) - else: - p_ht = tl.make_block_ptr(ht, (K, V), (V, 1), (64, i_v * BV), (64, BV), (1, 0)) - tl.store(p_ht, b_h2.to(p_ht.dtype.element_ty), boundary_check=(0, 1)) - if K > 128: - if STATE_V_FIRST: - p_ht = tl.make_block_ptr(ht, (V, K), (K, 1), (i_v * BV, 128), (BV, 64), (1, 0)) - else: - p_ht = tl.make_block_ptr(ht, (K, V), (V, 1), (128, i_v * BV), (64, BV), (1, 0)) - tl.store(p_ht, b_h3.to(p_ht.dtype.element_ty), boundary_check=(0, 1)) - if K > 192: - if STATE_V_FIRST: - p_ht = tl.make_block_ptr(ht, (V, K), (K, 1), (i_v * BV, 192), (BV, 64), (1, 0)) - else: - p_ht = tl.make_block_ptr(ht, (K, V), (V, 1), (192, i_v * BV), (64, BV), (1, 0)) - tl.store(p_ht, b_h4.to(p_ht.dtype.element_ty), boundary_check=(0, 1)) - - -@triton.heuristics({ - 'USE_G': lambda args: args['g'] is not None, - 'USE_GK': lambda args: args['gk'] is not None, - 'USE_INITIAL_STATE': lambda args: args['dh0'] is not None, - 'USE_FINAL_STATE_GRADIENT': lambda args: args['dht'] is not None, - 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None, -}) -@fla_cache_autotune( - configs=[ - triton.Config({'BV': BV}, num_warps=num_warps, num_stages=num_stages) - for num_warps in [2, 4] - for num_stages in ([2, 3, 4] if check_shared_mem('ampere') else [1]) - for BV in ([32, 64] if check_shared_mem('ada') else [32]) - ], - key=['H', 'HV', 'K', 'V', 'BT', 'BV', 'USE_G', 'STATE_V_FIRST'], - **autotune_cache_kwargs, -) -@triton.jit(do_not_specialize=['T']) -def chunk_gated_delta_rule_bwd_kernel_dhu_blockdim64( - q, - k, - w, - g, - gk, - dht, - dh0, - do, - dh, - dv, - dv2, - cu_seqlens, - chunk_offsets, - scale, - T, - H: tl.constexpr, - HV: tl.constexpr, - K: tl.constexpr, - V: tl.constexpr, - BT: tl.constexpr, - BV: tl.constexpr, - USE_G: tl.constexpr, - USE_GK: tl.constexpr, - USE_INITIAL_STATE: tl.constexpr, - USE_FINAL_STATE_GRADIENT: tl.constexpr, - STATE_V_FIRST: tl.constexpr, - IS_VARLEN: tl.constexpr, -): - i_v, i_nh = tl.program_id(0), tl.program_id(1) - i_n, i_h = i_nh // HV, i_nh % HV - if IS_VARLEN: - bos, eos = tl.load(cu_seqlens + i_n).to(tl.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) - T = eos - bos - NT = tl.cdiv(T, BT) - boh = tl.load(chunk_offsets + i_n).to(tl.int32) - else: - bos, eos = i_n * T, i_n * T + T - NT = tl.cdiv(T, BT) - boh = i_n * NT - - if STATE_V_FIRST: - b_dh1 = tl.zeros([BV, 64], dtype=tl.float32) - if K > 64: - b_dh2 = tl.zeros([BV, 64], dtype=tl.float32) - if K > 128: - b_dh3 = tl.zeros([BV, 64], dtype=tl.float32) - if K > 192: - b_dh4 = tl.zeros([BV, 64], dtype=tl.float32) - else: - b_dh1 = tl.zeros([64, BV], dtype=tl.float32) - if K > 64: - b_dh2 = tl.zeros([64, BV], dtype=tl.float32) - if K > 128: - b_dh3 = tl.zeros([64, BV], dtype=tl.float32) - if K > 192: - b_dh4 = tl.zeros([64, BV], dtype=tl.float32) - - # calculate offset - q += (bos * H + i_h // (HV // H)).to(tl.int64) * K - k += (bos * H + i_h // (HV // H)).to(tl.int64) * K - w += (bos * HV + i_h).to(tl.int64) * K - do += (bos * HV + i_h).to(tl.int64) * V - dv += (bos * HV + i_h).to(tl.int64) * V - dv2 += (bos * HV + i_h).to(tl.int64) * V - dh += (boh * HV + i_h).to(tl.int64) * K*V - if USE_GK: - gk += (bos * HV + i_h).to(tl.int64) * K - - if USE_INITIAL_STATE: - dh0 += i_nh * K*V - if USE_FINAL_STATE_GRADIENT: - dht += i_nh * K*V - - if USE_FINAL_STATE_GRADIENT: - if STATE_V_FIRST: - p_dht1 = tl.make_block_ptr(dht, (V, K), (K, 1), (i_v * BV, 0), (BV, 64), (1, 0)) - else: - p_dht1 = tl.make_block_ptr(dht, (K, V), (V, 1), (0, i_v * BV), (64, BV), (1, 0)) - b_dh1 += tl.load(p_dht1, boundary_check=(0, 1)) - if K > 64: - if STATE_V_FIRST: - p_dht2 = tl.make_block_ptr(dht, (V, K), (K, 1), (i_v * BV, 64), (BV, 64), (1, 0)) - else: - p_dht2 = tl.make_block_ptr(dht, (K, V), (V, 1), (64, i_v * BV), (64, BV), (1, 0)) - b_dh2 += tl.load(p_dht2, boundary_check=(0, 1)) - if K > 128: - if STATE_V_FIRST: - p_dht3 = tl.make_block_ptr(dht, (V, K), (K, 1), (i_v * BV, 128), (BV, 64), (1, 0)) - else: - p_dht3 = tl.make_block_ptr(dht, (K, V), (V, 1), (128, i_v * BV), (64, BV), (1, 0)) - b_dh3 += tl.load(p_dht3, boundary_check=(0, 1)) - if K > 192: - if STATE_V_FIRST: - p_dht4 = tl.make_block_ptr(dht, (V, K), (K, 1), (i_v * BV, 192), (BV, 64), (1, 0)) - else: - p_dht4 = tl.make_block_ptr(dht, (K, V), (V, 1), (192, i_v * BV), (64, BV), (1, 0)) - b_dh4 += tl.load(p_dht4, boundary_check=(0, 1)) - - for i_t in range(NT - 1, -1, -1): - i_t_int64 = i_t.to(tl.int64) - if STATE_V_FIRST: - p_dh1 = tl.make_block_ptr(dh + i_t_int64*HV*K*V, (V, K), (K, 1), (i_v * BV, 0), (BV, 64), (1, 0)) - else: - p_dh1 = tl.make_block_ptr(dh + i_t_int64*HV*K*V, (K, V), (V, 1), (0, i_v * BV), (64, BV), (1, 0)) - tl.store(p_dh1, b_dh1.to(p_dh1.dtype.element_ty), boundary_check=(0, 1)) - if K > 64: - if STATE_V_FIRST: - p_dh2 = tl.make_block_ptr(dh + i_t_int64*HV*K*V, (V, K), (K, 1), (i_v * BV, 64), (BV, 64), (1, 0)) - else: - p_dh2 = tl.make_block_ptr(dh + i_t_int64*HV*K*V, (K, V), (V, 1), (64, i_v * BV), (64, BV), (1, 0)) - tl.store(p_dh2, b_dh2.to(p_dh2.dtype.element_ty), boundary_check=(0, 1)) - if K > 128: - if STATE_V_FIRST: - p_dh3 = tl.make_block_ptr(dh + i_t_int64*HV*K*V, (V, K), (K, 1), (i_v * BV, 128), (BV, 64), (1, 0)) - else: - p_dh3 = tl.make_block_ptr(dh + i_t_int64*HV*K*V, (K, V), (V, 1), (128, i_v * BV), (64, BV), (1, 0)) - tl.store(p_dh3, b_dh3.to(p_dh3.dtype.element_ty), boundary_check=(0, 1)) - if K > 192: - if STATE_V_FIRST: - p_dh4 = tl.make_block_ptr(dh + i_t_int64*HV*K*V, (V, K), (K, 1), (i_v * BV, 192), (BV, 64), (1, 0)) - else: - p_dh4 = tl.make_block_ptr(dh + i_t_int64*HV*K*V, (K, V), (V, 1), (192, i_v * BV), (64, BV), (1, 0)) - tl.store(p_dh4, b_dh4.to(p_dh4.dtype.element_ty), boundary_check=(0, 1)) - - last_idx = min((i_t + 1) * BT, T) - 1 - if USE_G: - bg_last = tl.load(g + (bos + last_idx) * HV + i_h).to(tl.float32) - p_g = tl.make_block_ptr(g + bos * HV + i_h, (T,), (HV,), (i_t * BT,), (BT,), (0,)) - b_g = tl.load(p_g, boundary_check=(0,)).to(tl.float32) - bg_last_exp = exp2(bg_last) - b_g_exp = exp2(b_g) - p_dv = tl.make_block_ptr(dv, (T, V), (HV*V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) - p_dv2 = tl.make_block_ptr(dv2, (T, V), (HV*V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) - p_do = tl.make_block_ptr(do, (T, V), (HV*V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) - - b_do = tl.load(p_do, boundary_check=(0, 1)) - - # Update dv - p_k = tl.make_block_ptr(k, (T, K), (H*K, 1), (i_t * BT, 0), (BT, 64), (1, 0)) - b_k = tl.load(p_k, boundary_check=(0, 1)) - if USE_GK: - o_k1 = tl.arange(0, 64) - b_gk_last1 = tl.load(gk + last_idx * HV*K + o_k1, mask=(o_k1 < K), other=0.).to(tl.float32) - if STATE_V_FIRST: - b_dv = tl.dot(b_k, tl.trans(b_dh1).to(b_k.dtype)) - else: - b_dv = tl.dot(b_k, b_dh1.to(b_k.dtype)) - - if K > 64: - p_k = tl.make_block_ptr(k, (T, K), (H*K, 1), (i_t * BT, 64), (BT, 64), (1, 0)) - b_k = tl.load(p_k, boundary_check=(0, 1)) - if USE_GK: - o_k2 = 64 + o_k1 - b_gk_last2 = tl.load(gk + last_idx * HV*K + o_k2, mask=(o_k2 < K), other=0.).to(tl.float32) - if STATE_V_FIRST: - b_dv += tl.dot(b_k, tl.trans(b_dh2).to(b_k.dtype)) - else: - b_dv += tl.dot(b_k, b_dh2.to(b_k.dtype)) - - if K > 128: - p_k = tl.make_block_ptr(k, (T, K), (H*K, 1), (i_t * BT, 128), (BT, 64), (1, 0)) - b_k = tl.load(p_k, boundary_check=(0, 1)) - if USE_GK: - o_k3 = 128 + o_k1 - b_gk_last3 = tl.load(gk + last_idx * HV*K + o_k3, mask=(o_k3 < K), other=0.).to(tl.float32) - if STATE_V_FIRST: - b_dv += tl.dot(b_k, tl.trans(b_dh3).to(b_k.dtype)) - else: - b_dv += tl.dot(b_k, b_dh3.to(b_k.dtype)) - - if K > 192: - p_k = tl.make_block_ptr(k, (T, K), (H*K, 1), (i_t * BT, 192), (BT, 64), (1, 0)) - b_k = tl.load(p_k, boundary_check=(0, 1)) - if USE_GK: - o_k4 = 192 + o_k1 - b_gk_last4 = tl.load(gk + last_idx * HV*K + o_k4, mask=(o_k4 < K), other=0.).to(tl.float32) - if STATE_V_FIRST: - b_dv += tl.dot(b_k, tl.trans(b_dh4).to(b_k.dtype)) - else: - b_dv += tl.dot(b_k, b_dh4.to(b_k.dtype)) - - if USE_G: - m_t = (i_t * BT + tl.arange(0, BT)) < T - b_dv *= tl.where(m_t, exp2(bg_last - b_g), 0)[:, None] - b_dv += tl.load(p_dv, boundary_check=(0, 1)) - - tl.store(p_dv2, b_dv.to(p_dv.dtype.element_ty), boundary_check=(0, 1)) - # Update dh - p_w = tl.make_block_ptr(w, (K, T), (1, HV*K), (0, i_t * BT), (64, BT), (0, 1)) - p_q = tl.make_block_ptr(q, (K, T), (1, H*K), (0, i_t * BT), (64, BT), (0, 1)) - b_w = tl.load(p_w, boundary_check=(0, 1)) - b_q = tl.load(p_q, boundary_check=(0, 1)) - if USE_G: - b_dh1 *= bg_last_exp - b_q = b_q * b_g_exp[None, :] - if USE_GK: - if STATE_V_FIRST: - b_dh1 *= exp2(b_gk_last1)[None, :] - else: - b_dh1 *= exp2(b_gk_last1[:, None]) - if STATE_V_FIRST: - b_dh1 += tl.trans(tl.dot(b_q.to(b_q.dtype), b_do.to(b_q.dtype)) * scale - tl.dot(b_w, b_dv.to(b_w.dtype))) - else: - b_dh1 += tl.dot(b_q.to(b_q.dtype), b_do.to(b_q.dtype)) * scale - tl.dot(b_w, b_dv.to(b_w.dtype)) - if K > 64: - p_q = tl.make_block_ptr(q, (K, T), (1, H*K), (64, i_t * BT), (64, BT), (0, 1)) - p_w = tl.make_block_ptr(w, (K, T), (1, HV*K), (64, i_t * BT), (64, BT), (0, 1)) - b_q = tl.load(p_q, boundary_check=(0, 1)) - b_w = tl.load(p_w, boundary_check=(0, 1)) - if USE_G: - b_dh2 *= bg_last_exp - b_q = b_q * b_g_exp[None, :] - if USE_GK: - if STATE_V_FIRST: - b_dh2 *= exp2(b_gk_last2)[None, :] - else: - b_dh2 *= exp2(b_gk_last2[:, None]) - if STATE_V_FIRST: - b_dh2 += tl.trans(tl.dot(b_q.to(b_q.dtype), b_do.to(b_q.dtype)) * scale - tl.dot(b_w, b_dv.to(b_w.dtype))) - else: - b_dh2 += tl.dot(b_q.to(b_q.dtype), b_do.to(b_q.dtype)) * scale - tl.dot(b_w, b_dv.to(b_w.dtype)) - if K > 128: - p_q = tl.make_block_ptr(q, (K, T), (1, H*K), (128, i_t * BT), (64, BT), (0, 1)) - p_w = tl.make_block_ptr(w, (K, T), (1, HV*K), (128, i_t * BT), (64, BT), (0, 1)) - b_q = tl.load(p_q, boundary_check=(0, 1)) - b_w = tl.load(p_w, boundary_check=(0, 1)) - if USE_G: - b_dh3 *= bg_last_exp - b_q = b_q * b_g_exp[None, :] - if USE_GK: - if STATE_V_FIRST: - b_dh3 *= exp2(b_gk_last3)[None, :] - else: - b_dh3 *= exp2(b_gk_last3[:, None]) - if STATE_V_FIRST: - b_dh3 += tl.trans(tl.dot(b_q.to(b_q.dtype), b_do.to(b_q.dtype)) * scale - tl.dot(b_w, b_dv.to(b_w.dtype))) - else: - b_dh3 += tl.dot(b_q.to(b_q.dtype), b_do.to(b_q.dtype)) * scale - tl.dot(b_w, b_dv.to(b_w.dtype)) - if K > 192: - p_q = tl.make_block_ptr(q, (K, T), (1, H*K), (192, i_t * BT), (64, BT), (0, 1)) - p_w = tl.make_block_ptr(w, (K, T), (1, HV*K), (192, i_t * BT), (64, BT), (0, 1)) - b_q = tl.load(p_q, boundary_check=(0, 1)) - b_w = tl.load(p_w, boundary_check=(0, 1)) - if USE_G: - b_dh4 *= bg_last_exp - b_q = b_q * b_g_exp[None, :] - if USE_GK: - if STATE_V_FIRST: - b_dh4 *= exp2(b_gk_last4)[None, :] - else: - b_dh4 *= exp2(b_gk_last4[:, None]) - if STATE_V_FIRST: - b_dh4 += tl.trans(tl.dot(b_q.to(b_q.dtype), b_do.to(b_q.dtype)) * scale - tl.dot(b_w, b_dv.to(b_w.dtype))) - else: - b_dh4 += tl.dot(b_q.to(b_q.dtype), b_do.to(b_q.dtype)) * scale - tl.dot(b_w, b_dv.to(b_w.dtype)) - - if USE_INITIAL_STATE: - if STATE_V_FIRST: - p_dh0 = tl.make_block_ptr(dh0, (V, K), (K, 1), (i_v * BV, 0), (BV, 64), (1, 0)) - else: - p_dh0 = tl.make_block_ptr(dh0, (K, V), (V, 1), (0, i_v * BV), (64, BV), (1, 0)) - tl.store(p_dh0, b_dh1.to(p_dh0.dtype.element_ty), boundary_check=(0, 1)) - if K > 64: - if STATE_V_FIRST: - p_dh1 = tl.make_block_ptr(dh0, (V, K), (K, 1), (i_v * BV, 64), (BV, 64), (1, 0)) - else: - p_dh1 = tl.make_block_ptr(dh0, (K, V), (V, 1), (64, i_v * BV), (64, BV), (1, 0)) - tl.store(p_dh1, b_dh2.to(p_dh1.dtype.element_ty), boundary_check=(0, 1)) - if K > 128: - if STATE_V_FIRST: - p_dh2 = tl.make_block_ptr(dh0, (V, K), (K, 1), (i_v * BV, 128), (BV, 64), (1, 0)) - else: - p_dh2 = tl.make_block_ptr(dh0, (K, V), (V, 1), (128, i_v * BV), (64, BV), (1, 0)) - tl.store(p_dh2, b_dh3.to(p_dh2.dtype.element_ty), boundary_check=(0, 1)) - if K > 192: - if STATE_V_FIRST: - p_dh3 = tl.make_block_ptr(dh0, (V, K), (K, 1), (i_v * BV, 192), (BV, 64), (1, 0)) - else: - p_dh3 = tl.make_block_ptr(dh0, (K, V), (V, 1), (192, i_v * BV), (64, BV), (1, 0)) - tl.store(p_dh3, b_dh4.to(p_dh3.dtype.element_ty), boundary_check=(0, 1)) - - -@dispatch('common') -def chunk_gated_delta_rule_fwd_h( - k: torch.Tensor, - w: torch.Tensor, - u: torch.Tensor, - g: torch.Tensor | None = None, - gk: torch.Tensor | None = None, - initial_state: torch.Tensor | None = None, - output_final_state: bool = False, - chunk_size: int = 64, - save_new_value: bool = True, - state_v_first: bool = False, - cu_seqlens: torch.LongTensor | None = None, - cu_seqlens_cpu: torch.LongTensor | None = None, - chunk_indices: torch.LongTensor | None = None, -) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor | None]: - B, T, H, K, V, HV = *k.shape, u.shape[-1], u.shape[2] - BT = chunk_size - - if chunk_indices is None and cu_seqlens is not None: - chunk_indices = prepare_chunk_indices(cu_seqlens, chunk_size) - # N: the actual number of sequences in the batch with either equal or variable lengths - if cu_seqlens is None: - N, NT, chunk_offsets = B, triton.cdiv(T, BT), None - else: - N, NT, chunk_offsets = len(cu_seqlens) - 1, len(chunk_indices), prepare_chunk_offsets(cu_seqlens, BT) - assert K <= 256, "current kernel does not support head dimension larger than 256." - - if state_v_first: - h = k.new_empty(B, NT, HV, V, K) - final_state = k.new_zeros(N, HV, V, K, dtype=torch.float32) if output_final_state else None - else: - h = k.new_empty(B, NT, HV, K, V) - final_state = k.new_zeros(N, HV, K, V, dtype=torch.float32) if output_final_state else None - - v_new = torch.empty_like(u) if save_new_value else None - def grid(meta): return (triton.cdiv(V, meta['BV']), N*HV) - chunk_gated_delta_rule_fwd_kernel_h_blockdim64[grid]( - k=k, - v=u, - w=w, - v_new=v_new, - g=g, - gk=gk, - h=h, - h0=initial_state, - ht=final_state, - cu_seqlens=cu_seqlens, - chunk_offsets=chunk_offsets, - T=T, - H=H, - HV=HV, - K=K, - V=V, - BT=BT, - STATE_V_FIRST=state_v_first, - ) - return h, v_new, final_state - - -def chunk_gated_delta_rule_bwd_dhu( - q: torch.Tensor, - k: torch.Tensor, - w: torch.Tensor, - do: torch.Tensor, - dv: torch.Tensor, - g: torch.Tensor | None = None, - gk: torch.Tensor | None = None, - h0: torch.Tensor | None = None, - dht: torch.Tensor | None = None, - scale: float | None = None, - state_v_first: bool = False, - cu_seqlens: torch.LongTensor | None = None, - chunk_size: int = 64, - chunk_indices: torch.LongTensor | None = None, -) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: - B, T, H, K, V, HV = *q.shape, do.shape[-1], do.shape[2] - # N: the actual number of sequences in the batch with either equal or variable lengths - BT = chunk_size - assert K <= 256, "current kernel does not support head dimension being larger than 256." - - if chunk_indices is None and cu_seqlens is not None: - chunk_indices = prepare_chunk_indices(cu_seqlens, chunk_size) - if cu_seqlens is None: - N, NT, chunk_offsets = B, triton.cdiv(T, BT), None - else: - N, NT, chunk_offsets = len(cu_seqlens) - 1, len(chunk_indices), prepare_chunk_offsets(cu_seqlens, BT) - - if state_v_first: - dh = q.new_empty(B, NT, HV, V, K) - else: - dh = q.new_empty(B, NT, HV, K, V) - dh0 = torch.empty_like(h0, dtype=torch.float32) if h0 is not None else None - dv2 = torch.empty_like(dv) - - def grid(meta): return (triton.cdiv(V, meta['BV']), N*HV) - chunk_gated_delta_rule_bwd_kernel_dhu_blockdim64[grid]( - q=q, - k=k, - w=w, - g=g, - gk=gk, - dht=dht, - dh0=dh0, - do=do, - dh=dh, - dv=dv, - dv2=dv2, - cu_seqlens=cu_seqlens, - chunk_offsets=chunk_offsets, - scale=scale, - T=T, - H=H, - HV=HV, - K=K, - V=V, - BT=BT, - STATE_V_FIRST=state_v_first, - ) - return dh, dh0, dv2 diff --git a/build/torch-xpu/ops/common/chunk_h.py b/build/torch-xpu/ops/common/chunk_h.py deleted file mode 100644 index e4d485ed6df46ddbf77973e2e1d73cf03dfb8c96..0000000000000000000000000000000000000000 --- a/build/torch-xpu/ops/common/chunk_h.py +++ /dev/null @@ -1,417 +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 -import triton.language as tl - -from ...ops.utils import prepare_chunk_offsets -from ...ops.utils.op import exp2 -from ...utils import autotune_cache_kwargs, check_shared_mem - -BKV_LIST = [32, 64] if check_shared_mem() else [16, 32] - - -@triton.heuristics({ - 'USE_INITIAL_STATE': lambda args: args['h0'] is not None, - 'STORE_FINAL_STATE': lambda args: args['ht'] is not None, - 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None, -}) -@triton.autotune( - configs=[ - triton.Config({'BK': BK, 'BV': BV}, num_warps=num_warps, num_stages=num_stages) - for BK in BKV_LIST - for BV in BKV_LIST - for num_warps in [1, 2, 4, 8] - for num_stages in [2, 3, 4] - ], - key=['BT', 'USE_G', 'USE_GK', 'USE_GV', 'STATE_V_FIRST'], - **autotune_cache_kwargs, -) -@triton.jit(do_not_specialize=['T']) -def chunk_fwd_kernel_h( - k, - v, - h, - g, - g_gamma, - gk, - gv, - h0, - ht, - cu_seqlens, - split_offsets, - T, - H: tl.constexpr, - K: tl.constexpr, - V: tl.constexpr, - BT: tl.constexpr, - BS: tl.constexpr, - BK: tl.constexpr, - BV: tl.constexpr, - USE_G: tl.constexpr, - USE_G_GAMMA: tl.constexpr, - USE_GK: tl.constexpr, - USE_GV: tl.constexpr, - USE_INITIAL_STATE: tl.constexpr, - STORE_FINAL_STATE: tl.constexpr, - IS_VARLEN: tl.constexpr, - STATE_V_FIRST: tl.constexpr, -): - i_k, i_v, i_nh = tl.program_id(0), tl.program_id(1), tl.program_id(2) - i_n, i_h = i_nh // H, i_nh % H - if IS_VARLEN: - bos, eos = tl.load(cu_seqlens + i_n).to(tl.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) - T = eos - bos - NT, NS = tl.cdiv(T, BT), tl.cdiv(T, BS) - boh = tl.load(split_offsets + i_n).to(tl.int32) - else: - bos, eos = i_n * T, i_n * T + T - NT, NS = tl.cdiv(T, BT), tl.cdiv(T, BS) - boh = i_n * NS - NTS = BS // BT - - if USE_G_GAMMA: - # decay rate given the head index - b_gamma = tl.load(g_gamma + i_h) - b_g = b_gamma * (tl.arange(0, BT) + 1) - - # [BK, BV] accumulator; STATE_V_FIRST only flips the stored state's HBM layout to [V, K], applied at the load/store below. - b_h = tl.zeros([BK, BV], dtype=tl.float32) - if USE_INITIAL_STATE: - if STATE_V_FIRST: - p_h0 = tl.make_block_ptr(h0 + i_nh * K*V, (V, K), (K, 1), (i_v * BV, i_k * BK), (BV, BK), (1, 0)) - b_h = tl.trans(tl.load(p_h0, boundary_check=(0, 1))).to(tl.float32) - else: - p_h0 = tl.make_block_ptr(h0 + i_nh * K*V, (K, V), (V, 1), (i_k * BK, i_v * BV), (BK, BV), (1, 0)) - b_h = tl.load(p_h0, boundary_check=(0, 1)).to(tl.float32) - - for i_t in range(NT): - i_s = i_t // NTS - p_k = tl.make_block_ptr(k + (bos*H + i_h) * K, (K, T), (1, H*K), (i_k * BK, i_t * BT), (BK, BT), (0, 1)) - p_v = tl.make_block_ptr(v + (bos*H + i_h) * V, (T, V), (H*V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) - - o_h = ((boh + i_s) * H + i_h).to(tl.int64) * K*V - if STATE_V_FIRST: - p_h = tl.make_block_ptr(h + o_h, (V, K), (K, 1), (i_v * BV, i_k * BK), (BV, BK), (1, 0)) - else: - p_h = tl.make_block_ptr(h + o_h, (K, V), (V, 1), (i_k * BK, i_v * BV), (BK, BV), (1, 0)) - - if i_t % NTS == 0: - tl.store(p_h, (tl.trans(b_h) if STATE_V_FIRST else b_h).to(p_h.dtype.element_ty), boundary_check=(0, 1)) - # [BK, BT] - b_k = tl.load(p_k, boundary_check=(0, 1)) - # [BT, BV] - b_v = tl.load(p_v, boundary_check=(0, 1)) - last_idx = min((i_t + 1) * BT, T) - 1 - - # scalar decay - if USE_G: - b_g_last = tl.load(g + bos * H + last_idx * H + i_h) - p_g = g + bos*H + (i_t * BT + tl.arange(0, BT)) * H + i_h - b_g = tl.load(p_g, mask=(i_t * BT + tl.arange(0, BT) < T), other=0.) - b_h *= exp2(b_g_last) - b_v = (b_v * exp2(b_g_last - b_g)[:, None]).to(b_v.dtype) - - if USE_G_GAMMA: - b_g_last = b_gamma * min(BT, T - i_t * BT) - b_h *= exp2(b_g_last) - b_v = (b_v * exp2(b_g_last - b_g)[:, None]).to(b_v.dtype) - - # vector decay, h = Diag(gk) @ h - if USE_GK: - p_gk = tl.make_block_ptr(gk + (bos*H + i_h) * K, (K, T), (1, H*K), (i_k * BK, i_t * BT), (BK, BT), (0, 1)) - p_gk_last = gk + (bos + last_idx) * H*K + i_h * K + i_k * BK + tl.arange(0, BK) - - b_gk_last = tl.load(p_gk_last, mask=(i_k * BK + tl.arange(0, BK) < K), other=0.) - b_gk = tl.load(p_gk, boundary_check=(0, 1)) - b_h *= exp2(b_gk_last)[:, None] - b_k = (b_k * exp2(b_gk_last[:, None] - b_gk)).to(b_k.dtype) - - # vector decay, h = h @ Diag(gv) - if USE_GV: - p_gv = tl.make_block_ptr(gv + (bos*H + i_h) * V, (T, V), (H*V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) - p_gv_last = gv + (bos + last_idx) * H*V + i_h * V + i_v * BV + tl.arange(0, BV) - - b_gv_last = tl.load(p_gv_last, mask=(i_v * BV + tl.arange(0, BV) < V), other=0.) - b_gv = tl.load(p_gv, boundary_check=(0, 1)) - b_h *= exp2(b_gv_last)[None, :] - b_v = (b_v * exp2(b_gv_last[None, :] - b_gv)).to(b_v.dtype) - - b_h += tl.dot(b_k, b_v) - - if STORE_FINAL_STATE: - if STATE_V_FIRST: - p_ht = tl.make_block_ptr(ht + i_nh * K*V, (V, K), (K, 1), (i_v * BV, i_k * BK), (BV, BK), (1, 0)) - tl.store(p_ht, tl.trans(b_h).to(p_ht.dtype.element_ty), boundary_check=(0, 1)) - else: - p_ht = tl.make_block_ptr(ht + i_nh * K*V, (K, V), (V, 1), (i_k * BK, i_v * BV), (BK, BV), (1, 0)) - tl.store(p_ht, b_h.to(p_ht.dtype.element_ty), boundary_check=(0, 1)) - - -@triton.heuristics({ - 'STORE_INITIAL_STATE_GRADIENT': lambda args: args['dh0'] is not None, - 'USE_FINAL_STATE_GRADIENT': lambda args: args['dht'] is not None, - 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None, -}) -@triton.autotune( - configs=[ - triton.Config({'BK': BK, 'BV': BV}, num_warps=num_warps, num_stages=num_stages) - for BK in BKV_LIST - for BV in BKV_LIST - for num_warps in [1, 2, 4, 8] - for num_stages in [2, 3, 4] - ], - key=['BT', 'USE_G', 'USE_GK', 'USE_GV', 'STATE_V_FIRST'], - **autotune_cache_kwargs, -) -@triton.jit(do_not_specialize=['T']) -def chunk_bwd_kernel_dh( - q, - g, - g_gamma, - gk, - gv, - do, - dh, - dht, - dh0, - cu_seqlens, - split_offsets, - scale, - T, - HQ: tl.constexpr, - H: tl.constexpr, - K: tl.constexpr, - V: tl.constexpr, - BT: tl.constexpr, - BS: tl.constexpr, - BK: tl.constexpr, - BV: tl.constexpr, - NG: tl.constexpr, - USE_G: tl.constexpr, - USE_G_GAMMA: tl.constexpr, - USE_GK: tl.constexpr, - USE_GV: tl.constexpr, - STORE_INITIAL_STATE_GRADIENT: tl.constexpr, - USE_FINAL_STATE_GRADIENT: tl.constexpr, - IS_VARLEN: tl.constexpr, - STATE_V_FIRST: tl.constexpr, -): - i_k, i_v, i_nh = tl.program_id(0), tl.program_id(1), tl.program_id(2) - i_n, i_hq = i_nh // HQ, i_nh % HQ - i_h = i_hq // NG - if IS_VARLEN: - bos, eos = tl.load(cu_seqlens + i_n).to(tl.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) - T = eos - bos - NT = tl.cdiv(T, BT) - NS = tl.cdiv(T, BS) - boh = tl.load(split_offsets + i_n).to(tl.int32) - else: - bos, eos = i_n * T, i_n * T + T - NT = tl.cdiv(T, BT) - NS = tl.cdiv(T, BS) - boh = i_n * NS - - if USE_G_GAMMA: - b_gamma = tl.load(g_gamma + i_h) - b_g = b_gamma * (tl.arange(0, BT) + 1) - - # [BK, BV] accumulator; STATE_V_FIRST only flips the stored state's HBM layout to [V, K], applied at the load/store below. - b_dh = tl.zeros([BK, BV], dtype=tl.float32) - if USE_FINAL_STATE_GRADIENT: - if STATE_V_FIRST: - p_dht = tl.make_block_ptr(dht + i_nh * K*V, (V, K), (K, 1), (i_v * BV, i_k * BK), (BV, BK), (1, 0)) - b_dh += tl.trans(tl.load(p_dht, boundary_check=(0, 1))).to(tl.float32) - else: - p_dht = tl.make_block_ptr(dht + i_nh * K*V, (K, V), (V, 1), (i_k * BK, i_v * BV), (BK, BV), (1, 0)) - b_dh += tl.load(p_dht, boundary_check=(0, 1)).to(tl.float32) - - for i_t in range(NT - 1, -1, -1): - i_s = i_t // (BS // BT) - o_dh = ((boh + i_s) * H + i_h).to(tl.int64) * K*V - if STATE_V_FIRST: - p_dh = tl.make_block_ptr(dh + o_dh, (V, K), (K, 1), (i_v * BV, i_k * BK), (BV, BK), (1, 0)) - else: - p_dh = tl.make_block_ptr(dh + o_dh, (K, V), (V, 1), (i_k * BK, i_v * BV), (BK, BV), (1, 0)) - - if i_t % (BS // BT) == 0: - tl.store(p_dh, (tl.trans(b_dh) if STATE_V_FIRST else b_dh).to(p_dh.dtype.element_ty), boundary_check=(0, 1)) - last_idx = min(i_t * BT + BT, T) - 1 - # [BK, BT] - p_q = tl.make_block_ptr(q + (bos*HQ + i_hq) * K, (K, T), (1, HQ*K), (i_k * BK, i_t * BT), (BK, BT), (0, 1)) - p_do = tl.make_block_ptr(do + (bos*HQ + i_hq) * V, (T, V), (HQ*V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) - b_q = tl.load(p_q, boundary_check=(0, 1)) - b_q = (b_q * scale).to(b_q.dtype) - # [BT, BV] - b_do = tl.load(p_do, boundary_check=(0, 1)) - - if USE_G: - p_g = g + (bos + i_t * BT + tl.arange(0, BT)) * H + i_h - b_g_last = tl.load(g + (bos + last_idx) * H + i_h) - b_g = tl.load(p_g, mask=(i_t * BT + tl.arange(0, BT) < T), other=0.) - b_q = (b_q * exp2(b_g)[None, :]).to(b_q.dtype) - b_dh *= exp2(b_g_last) - - if USE_G_GAMMA: - b_g_last = b_gamma * min(BT, T - i_t * BT) - b_q = (b_q * exp2(b_g)[None, :]).to(b_q.dtype) - b_dh *= exp2(b_g_last) - - if USE_GK: - p_gk = tl.make_block_ptr(gk + (bos*H + i_h) * K, (K, T), (1, H*K), (i_k * BK, i_t * BT), (BK, BT), (0, 1)) - p_gk_last = gk + (bos + last_idx) * H*K + i_h * K + i_k * BK + tl.arange(0, BK) - - b_gk = tl.load(p_gk, boundary_check=(0, 1)) - b_gk_last = tl.load(p_gk_last, mask=(i_k * BK + tl.arange(0, BK) < K), other=0.) - b_q = (b_q * exp2(b_gk)).to(b_q.dtype) - b_dh *= exp2(b_gk_last)[:, None] - - if USE_GV: - p_gv = tl.make_block_ptr(gv + (bos*H + i_h) * V, (T, V), (H*V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) - p_gv_last = gv + (bos + last_idx) * H*V + i_h * V + i_v * BV + tl.arange(0, BV) - - b_gv = tl.load(p_gv, boundary_check=(0, 1)) - b_gv_last = tl.load(p_gv_last, mask=(i_v * BV + tl.arange(0, BV) < V), other=0.) - b_do = (b_do * exp2(b_gv)) - b_dh *= exp2(b_gv_last)[None, :] - - b_dh += tl.dot(b_q, b_do.to(b_q.dtype)) - - if STORE_INITIAL_STATE_GRADIENT: - if STATE_V_FIRST: - p_dh0 = tl.make_block_ptr(dh0 + i_nh * K*V, (V, K), (K, 1), (i_v * BV, i_k * BK), (BV, BK), (1, 0)) - tl.store(p_dh0, tl.trans(b_dh).to(p_dh0.dtype.element_ty), boundary_check=(0, 1)) - else: - p_dh0 = tl.make_block_ptr(dh0 + i_nh * K*V, (K, V), (V, 1), (i_k * BK, i_v * BV), (BK, BV), (1, 0)) - tl.store(p_dh0, b_dh.to(p_dh0.dtype.element_ty), boundary_check=(0, 1)) - - -def chunk_fwd_h( - k: torch.Tensor, - v: torch.Tensor, - g: torch.Tensor | None = None, - g_gamma: torch.Tensor | None = None, - gk: torch.Tensor | None = None, - gv: torch.Tensor | None = None, - h0: torch.Tensor | None = None, - output_final_state: bool = False, - state_v_first: bool = False, - cu_seqlens: torch.Tensor | None = None, - chunk_size: int = 64, - split_size: int | None = None, - states_in_fp32: bool = False, -) -> tuple[torch.Tensor, torch.Tensor]: - B, T, H, K, V = *k.shape, v.shape[-1] - BT = chunk_size - BS = BT if split_size is None else split_size - assert BS % BT == 0, f"The `split_size` (got {BS}) must be a multiple of `chunk_size` {BT}" - # N: the actual number of sequences in the batch with either equal or variable lengths - if cu_seqlens is None: - N, NS, split_offsets = B, triton.cdiv(T, BS), None - else: - split_offsets = prepare_chunk_offsets(cu_seqlens, BS) - N, NS = len(cu_seqlens) - 1, split_offsets[-1].item() - - # `state_v_first` stores the states in V-first `[V, K]` layout instead of `[K, V]` - state_shape = (V, K) if state_v_first else (K, V) - h = k.new_empty(B, NS, H, *state_shape, dtype=k.dtype if not states_in_fp32 else torch.float) - ht = k.new_empty(N, H, *state_shape, dtype=torch.float) if output_final_state else None - def grid(meta): return (triton.cdiv(K, meta['BK']), triton.cdiv(V, meta['BV']), N * H) - chunk_fwd_kernel_h[grid]( - k=k, - v=v, - h=h, - g=g, - g_gamma=g_gamma, - gk=gk, - gv=gv, - h0=h0, - ht=ht, - cu_seqlens=cu_seqlens, - split_offsets=split_offsets, - T=T, - H=H, - K=K, - V=V, - BT=BT, - BS=BS, - USE_G=g is not None, - USE_G_GAMMA=g_gamma is not None, - USE_GK=gk is not None, - USE_GV=gv is not None, - STATE_V_FIRST=state_v_first, - ) - return h, ht - - -def chunk_bwd_dh( - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - do: torch.Tensor, - h0: torch.Tensor, - dht: torch.Tensor, - scale: float, - g: torch.Tensor | None = None, - g_gamma: torch.Tensor | None = None, - gk: torch.Tensor | None = None, - gv: torch.Tensor | None = None, - state_v_first: bool = False, - cu_seqlens: torch.Tensor | None = None, - chunk_size: int = 64, - split_size: int | None = None, - states_in_fp32: bool = False, -) -> tuple[torch.Tensor, torch.Tensor]: - B, T, H, K, V = *k.shape, v.shape[-1] - HQ = q.shape[2] - BT = chunk_size - BS = BT if split_size is None else split_size - assert BS % BT == 0, f"The `split_size` (got {BS}) must be a multiple of `chunk_size` {BT}" - # N: the actual number of sequences in the batch with either equal or variable lengths - # NG: number of groups in GQA - if cu_seqlens is None: - N, NS, split_offsets = B, triton.cdiv(T, BS), None - else: - split_offsets = prepare_chunk_offsets(cu_seqlens, BS) - N, NS = len(cu_seqlens) - 1, split_offsets[-1].item() - NG = HQ // H - - # `state_v_first` stores the states in V-first `[V, K]` layout instead of `[K, V]` - state_shape = (V, K) if state_v_first else (K, V) - dh = k.new_empty(B, NS, HQ, *state_shape, dtype=k.dtype if not states_in_fp32 else torch.float) - dh0 = torch.empty_like(h0, dtype=torch.float) if h0 is not None else None - - def grid(meta): return (triton.cdiv(K, meta['BK']), triton.cdiv(V, meta['BV']), N * H) - chunk_bwd_kernel_dh[grid]( - q=q, - g=g, - g_gamma=g_gamma, - gk=gk, - gv=gv, - do=do, - dh=dh, - dht=dht, - dh0=dh0, - cu_seqlens=cu_seqlens, - split_offsets=split_offsets, - scale=scale, - T=T, - HQ=HQ, - H=H, - K=K, - V=V, - BT=BT, - BS=BS, - NG=NG, - USE_G=g is not None, - USE_G_GAMMA=g_gamma is not None, - USE_GK=gk is not None, - USE_GV=gv is not None, - STATE_V_FIRST=state_v_first, - ) - return dh, dh0 diff --git a/build/torch-xpu/ops/common/chunk_h_parallel.py b/build/torch-xpu/ops/common/chunk_h_parallel.py deleted file mode 100644 index 2d864f5a710a0ddd563a628fde675b6cee99db86..0000000000000000000000000000000000000000 --- a/build/torch-xpu/ops/common/chunk_h_parallel.py +++ /dev/null @@ -1,563 +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 - -""" -Fully parallelized state passing. -""" - - -import torch -import triton -import triton.language as tl - -from ...ops.utils import prepare_chunk_indices, prepare_chunk_offsets -from ...ops.utils.op import exp -from ...utils import autotune_cache_kwargs - - -@triton.heuristics({ - 'USE_INITIAL_STATE': lambda args: args['h0'] is not None, - 'STORE_FINAL_STATE': lambda args: args['ht'] is not None, - 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None, -}) -@triton.autotune( - configs=[ - triton.Config({'BK': BK, 'BV': BV}, num_warps=num_warps, num_stages=num_stages) - for BK in [32, 64, 128] - for BV in [32, 64, 128] - for num_warps in [2, 4, 8] - for num_stages in [2, 3, 4] - ], - key=['BT', 'USE_G', 'USE_GK', 'USE_GV'], - **autotune_cache_kwargs, -) -@triton.jit(do_not_specialize=['T']) -def chunk_fwd_kernel_h_parallel( - k, - v, - h, - g, - gk, - gv, - h0, - ht, - cu_seqlens, - chunk_indices, - T, - H: tl.constexpr, - K: tl.constexpr, - V: tl.constexpr, - BT: tl.constexpr, - BK: tl.constexpr, - BV: tl.constexpr, - USE_G: tl.constexpr, - USE_GK: tl.constexpr, - USE_GV: tl.constexpr, - USE_INITIAL_STATE: tl.constexpr, - STORE_FINAL_STATE: tl.constexpr, - IS_VARLEN: tl.constexpr, -): - i_kv, i_t, i_bh = tl.program_id(0), tl.program_id(1), tl.program_id(2) - - NV = tl.cdiv(V, BV) - # i_b: batch index - # i_h: head index - # i_n: sequence index - # i_t: chunk index within current sequence - # i_tg: (global) chunk index across all sequences - i_k, i_v = i_kv // NV, i_kv % NV - i_b, i_h = i_bh // H, i_bh % H - 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.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) - T = eos - bos - NT = tl.cdiv(T, BT) - else: - bos, eos = i_b * T, i_b * T + T - NT = tl.cdiv(T, BT) - i_n, i_tg = i_b, i_b * NT + i_t - i_nh = i_n * H + i_h - - p_k = tl.make_block_ptr(k + (bos*H + i_h) * K, (K, T), (1, H*K), (i_k * BK, i_t * BT), (BK, BT), (0, 1)) - p_v = tl.make_block_ptr(v + (bos*H + i_h) * V, (T, V), (H*V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) - p_h = tl.make_block_ptr(h + (i_tg * H + i_h) * K*V, (K, V), (V, 1), (i_k * BK, i_v * BV), (BK, BV), (1, 0)) - - if i_t == 0: - if USE_INITIAL_STATE: - p_h0 = tl.make_block_ptr(h0 + i_nh * K*V, (K, V), (V, 1), (i_k * BK, i_v * BV), (BK, BV), (1, 0)) - b_h = tl.load(p_h0, boundary_check=(0, 1)).to(tl.float32) - else: - b_h = tl.zeros([BK, BV], dtype=tl.float32) - tl.store(p_h, b_h.to(p_h.dtype.element_ty), boundary_check=(0, 1)) - - # [BK, BT] - b_k = tl.load(p_k, boundary_check=(0, 1)) - # [BT, BV] - b_v = tl.load(p_v, boundary_check=(0, 1)) - - last_idx = min(i_t * BT + BT, T) - 1 - # scalar decay - if USE_G: - b_g_last = tl.load(g + bos * H + last_idx * H + i_h) - p_g = g + bos*H + (i_t * BT + tl.arange(0, BT)) * H + i_h - b_g = tl.load(p_g, mask=(i_t * BT + tl.arange(0, BT) < T), other=0.) - b_v = (b_v * exp(b_g_last - b_g)[:, None]).to(b_v.dtype) - - # vector decay, h = Diag(gk) @ h - if USE_GK: - p_gk = tl.make_block_ptr(gk + (bos*H + i_h) * K, (K, T), (1, H*K), (i_k * BK, i_t * BT), (BK, BT), (0, 1)) - p_gk_last = gk + (bos + last_idx) * H*K + i_h * K + i_k * BK + tl.arange(0, BK) - b_gk_last = tl.load(p_gk_last, mask=(i_k * BK + tl.arange(0, BK) < K), other=0.) - - b_gk = tl.load(p_gk, boundary_check=(0, 1)) - b_k = (b_k * exp(b_gk_last[:, None] - b_gk)).to(b_k.dtype) - - # vector decay, h = h @ Diag(gv) - if USE_GV: - p_gv = tl.make_block_ptr(gv + (bos*H + i_h) * V, (T, V), (H*V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) - p_gv_last = gv + (bos + last_idx) * H*V + i_h * V + i_v * BV + tl.arange(0, BV) - - b_gv_last = tl.load(p_gv_last, mask=(i_v * BV + tl.arange(0, BV) < V), other=0.) - b_gv = tl.load(p_gv, boundary_check=(0, 1)) - b_v = (b_v * exp(b_gv_last[None, :] - b_gv)).to(b_v.dtype) - - b_h = tl.dot(b_k, b_v) - if i_t < NT - 1: - p_h = tl.make_block_ptr(h + ((i_tg + 1) * H + i_h) * K*V, (K, V), (V, 1), (i_k * BK, i_v * BV), (BK, BV), (1, 0)) - tl.store(p_h, b_h.to(p_h.dtype.element_ty), boundary_check=(0, 1)) - elif STORE_FINAL_STATE: - p_ht = tl.make_block_ptr(ht + i_nh * K*V, (K, V), (V, 1), (i_k * BK, i_v * BV), (BK, BV), (1, 0)) - tl.store(p_ht, b_h.to(p_ht.dtype.element_ty), boundary_check=(0, 1)) - - -@triton.heuristics({ - 'STORE_FINAL_STATE': lambda args: args['ht'] is not None, - 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None, -}) -@triton.autotune( - configs=[ - triton.Config({'BK': BK, 'BV': BV}, num_warps=num_warps, num_stages=num_stages) - for BK in [32, 64, 128] - for BV in [32, 64, 128] - for num_warps in [2, 4, 8, 16] - for num_stages in [2, 3] - ], - key=['BT', 'USE_G', 'USE_GK', 'USE_GV'], - **autotune_cache_kwargs, -) -@triton.jit(do_not_specialize=['T']) -def chunk_fwd_kernel_h_reduction( - h, - g, - gk, - gv, - kvt, - ht, - cu_seqlens, - chunk_offsets, - T, - H: tl.constexpr, - K: tl.constexpr, - V: tl.constexpr, - BT: tl.constexpr, - BK: tl.constexpr, - BV: tl.constexpr, - USE_G: tl.constexpr, - USE_GK: tl.constexpr, - USE_GV: tl.constexpr, - STORE_FINAL_STATE: tl.constexpr, - IS_VARLEN: tl.constexpr, -): - i_k, i_v, i_nh = tl.program_id(0), tl.program_id(1), tl.program_id(2) - i_n, i_h = i_nh // H, i_nh % H - if IS_VARLEN: - bos, eos = tl.load(cu_seqlens + i_n).to(tl.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) - T = eos - bos - NT = tl.cdiv(T, BT) - boh = tl.load(chunk_offsets + i_n).to(tl.int32) - else: - bos, eos = i_n * T, i_n * T + T - NT = tl.cdiv(T, BT) - boh = i_n * NT - - # [BK, BV] - b_h = tl.zeros([BK, BV], dtype=tl.float32) - for i_t in range(NT): - p_h = tl.make_block_ptr(h + ((boh + i_t) * H + i_h) * K*V, (K, V), (V, 1), (i_k * BK, i_v * BV), (BK, BV), (1, 0)) - b_h += tl.load(p_h, boundary_check=(0, 1)).to(tl.float32) - if i_t > 0: - tl.store(p_h, b_h.to(p_h.dtype.element_ty), boundary_check=(0, 1)) - - last_idx = min(i_t * BT + BT, T) - 1 - # scalar decay - if USE_G: - b_g_last = tl.load(g + bos * H + last_idx * H + i_h) - b_h *= exp(b_g_last) - - # vector decay, h = Diag(gk) @ h - if USE_GK: - p_gk_last = gk + (bos + last_idx) * H*K + i_h * K + i_k * BK + tl.arange(0, BK) - - b_gk_last = tl.load(p_gk_last, mask=(i_k * BK + tl.arange(0, BK) < K), other=0.) - b_h *= exp(b_gk_last)[:, None] - - # vector decay, h = h @ Diag(gv) - if USE_GV: - p_gv_last = gv + (bos + last_idx) * H*V + i_h * V + i_v * BV + tl.arange(0, BV) - - b_gv_last = tl.load(p_gv_last, mask=(i_v * BV + tl.arange(0, BV) < V), other=0.) - b_h *= exp(b_gv_last)[None, :] - - if STORE_FINAL_STATE: - p_kvt = tl.make_block_ptr(kvt + i_nh * K*V, (K, V), (V, 1), (i_k * BK, i_v * BV), (BK, BV), (1, 0)) - p_ht = tl.make_block_ptr(ht + i_nh * K*V, (K, V), (V, 1), (i_k * BK, i_v * BV), (BK, BV), (1, 0)) - b_h += tl.load(p_kvt, boundary_check=(0, 1)).to(tl.float32) - tl.store(p_ht, b_h.to(p_ht.dtype.element_ty), boundary_check=(0, 1)) - - -@triton.heuristics({ - 'STORE_INITIAL_STATE_GRADIENT': lambda args: args['dh0'] is not None, - 'USE_FINAL_STATE_GRADIENT': lambda args: args['dht'] is not None, - 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None, -}) -@triton.autotune( - configs=[ - triton.Config({'BK': BK, 'BV': BV}, num_warps=num_warps, num_stages=num_stages) - for BK in [32, 64, 128] - for BV in [32, 64, 128] - for num_warps in [2, 4, 8] - for num_stages in [2, 3, 4] - ], - key=['BT', 'USE_G', 'USE_GK', 'USE_GV'], - **autotune_cache_kwargs, -) -@triton.jit(do_not_specialize=['T']) -def chunk_bwd_kernel_dh_parallel( - q, - g, - gk, - gv, - do, - dh, - dht, - dh0, - cu_seqlens, - chunk_indices, - scale, - T, - HQ: tl.constexpr, - H: tl.constexpr, - K: tl.constexpr, - V: tl.constexpr, - BT: tl.constexpr, - BK: tl.constexpr, - BV: tl.constexpr, - NG: tl.constexpr, - USE_G: tl.constexpr, - USE_GK: tl.constexpr, - USE_GV: tl.constexpr, - STORE_INITIAL_STATE_GRADIENT: tl.constexpr, - USE_FINAL_STATE_GRADIENT: tl.constexpr, - IS_VARLEN: tl.constexpr, -): - i_kv, i_t, i_bh = tl.program_id(0), tl.program_id(1), tl.program_id(2) - - NV = tl.cdiv(V, BV) - i_k, i_v = i_kv // NV, i_kv % NV - i_b, i_hq = i_bh // HQ, i_bh % HQ - i_h = i_hq // NG - 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.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) - T = eos - bos - NT = tl.cdiv(T, BT) - else: - bos, eos = i_b * T, i_b * T + T - NT = tl.cdiv(T, BT) - i_n, i_tg = i_b, i_b * NT + i_t - i_nh = i_n * HQ + i_hq - - p_q = tl.make_block_ptr(q + (bos*HQ + i_hq) * K, (K, T), (1, HQ*K), (i_k * BK, i_t * BT), (BK, BT), (0, 1)) - p_do = tl.make_block_ptr(do + (bos*HQ + i_hq) * V, (T, V), (HQ*V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) - p_dh = tl.make_block_ptr(dh + (i_tg * H + i_h) * K*V, (K, V), (V, 1), (i_k * BK, i_v * BV), (BK, BV), (1, 0)) - - if i_t == NT - 1: - if USE_FINAL_STATE_GRADIENT: - p_dht = tl.make_block_ptr(dht + i_nh * K*V, (K, V), (V, 1), (i_k * BK, i_v * BV), (BK, BV), (1, 0)) - b_dh = tl.load(p_dht, boundary_check=(0, 1)).to(tl.float32) - else: - b_dh = tl.zeros([BK, BV], dtype=tl.float32) - tl.store(p_dh, b_dh.to(p_dh.dtype.element_ty), boundary_check=(0, 1)) - - # [BK, BT] - b_q = tl.load(p_q, boundary_check=(0, 1)) - b_q = (b_q * scale).to(b_q.dtype) - # [BT, BV] - b_do = tl.load(p_do, boundary_check=(0, 1)) - - if USE_G: - p_g = g + (bos + i_t * BT + tl.arange(0, BT)) * H + i_h - b_g = tl.load(p_g, mask=(i_t * BT + tl.arange(0, BT) < T), other=0.) - b_q = (b_q * exp(b_g)[None, :]).to(b_q.dtype) - - if USE_GK: - p_gk = tl.make_block_ptr(gk + (bos*H + i_h) * K, (K, T), (1, H*K), (i_k * BK, i_t * BT), (BK, BT), (0, 1)) - b_gk = tl.load(p_gk, boundary_check=(0, 1)) - b_q = (b_q * exp(b_gk)).to(b_q.dtype) - - if USE_GV: - p_gv = tl.make_block_ptr(gv + (bos*H + i_h) * V, (T, V), (H*V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) - b_gv = tl.load(p_gv, boundary_check=(0, 1)) - b_do = (b_do * exp(b_gv)).to(b_do.dtype) - - b_dh = tl.dot(b_q, b_do) - if i_t > 0: - p_dh = tl.make_block_ptr(dh + ((i_tg - 1) * H + i_h) * K*V, (K, V), (V, 1), (i_k * BK, i_v * BV), (BK, BV), (1, 0)) - tl.store(p_dh, b_dh.to(p_dh.dtype.element_ty), boundary_check=(0, 1)) - elif STORE_INITIAL_STATE_GRADIENT: - p_dh0 = tl.make_block_ptr(dh0 + i_nh * K*V, (K, V), (V, 1), (i_k * BK, i_v * BV), (BK, BV), (1, 0)) - tl.store(p_dh0, b_dh.to(p_dh0.dtype.element_ty), boundary_check=(0, 1)) - - -@triton.heuristics({ - 'STORE_INITIAL_STATE_GRADIENT': lambda args: args['dh0'] is not None, - 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None, -}) -@triton.autotune( - configs=[ - triton.Config({'BK': BK, 'BV': BV}, num_warps=num_warps, num_stages=num_stages) - for BK in [32, 64, 128] - for BV in [32, 64, 128] - for num_warps in [2, 4, 8, 16] - for num_stages in [2, 3] - ], - key=['BT', 'USE_G', 'USE_GK', 'USE_GV'], - **autotune_cache_kwargs, -) -@triton.jit(do_not_specialize=['T']) -def chunk_bwd_kernel_dh_reduction( - g, - gk, - gv, - dh, - doq0, - dh0, - cu_seqlens, - chunk_offsets, - T, - HQ: tl.constexpr, - H: tl.constexpr, - K: tl.constexpr, - V: tl.constexpr, - BT: tl.constexpr, - BK: tl.constexpr, - BV: tl.constexpr, - NG: tl.constexpr, - USE_G: tl.constexpr, - USE_GK: tl.constexpr, - USE_GV: tl.constexpr, - STORE_INITIAL_STATE_GRADIENT: tl.constexpr, - IS_VARLEN: tl.constexpr, -): - i_k, i_v, i_nh = tl.program_id(0), tl.program_id(1), tl.program_id(2) - i_n, i_hq = i_nh // HQ, i_nh % HQ - i_h = i_hq // NG - if IS_VARLEN: - bos, eos = tl.load(cu_seqlens + i_n).to(tl.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) - T = eos - bos - NT = tl.cdiv(T, BT) - boh = tl.load(chunk_offsets + i_n).to(tl.int32) - else: - bos, eos = i_n * T, i_n * T + T - NT = tl.cdiv(T, BT) - boh = i_n * NT - - # [BK, BV] - b_dh = tl.zeros([BK, BV], dtype=tl.float32) - for i_t in range(NT - 1, -1, -1): - p_dh = tl.make_block_ptr(dh + ((boh+i_t) * H + i_h) * K*V, (K, V), (V, 1), (i_k * BK, i_v * BV), (BK, BV), (1, 0)) - b_dh += tl.load(p_dh, boundary_check=(0, 1)).to(tl.float32) - if i_t < NT - 1: - tl.store(p_dh, b_dh.to(p_dh.dtype.element_ty), boundary_check=(0, 1)) - - last_idx = min(i_t * BT + BT, T) - 1 - if USE_G: - b_g_last = tl.load(g + (bos + last_idx) * H + i_h) - b_dh *= exp(b_g_last) - - if USE_GK: - p_gk_last = gk + (bos + last_idx) * H*K + i_h * K + i_k * BK + tl.arange(0, BK) - b_gk_last = tl.load(p_gk_last, mask=(i_k * BK + tl.arange(0, BK) < K), other=0.) - b_dh *= exp(b_gk_last)[:, None] - - if USE_GV: - p_gv_last = gv + (bos + last_idx) * H*V + i_h * V + i_v * BV + tl.arange(0, BV) - b_gv_last = tl.load(p_gv_last, mask=(i_v * BV + tl.arange(0, BV) < V), other=0.) - b_dh *= exp(b_gv_last)[None, :] - - if STORE_INITIAL_STATE_GRADIENT: - p_doq0 = tl.make_block_ptr(doq0 + i_nh * K*V, (K, V), (V, 1), (i_k * BK, i_v * BV), (BK, BV), (1, 0)) - p_dh0 = tl.make_block_ptr(dh0 + i_nh * K*V, (K, V), (V, 1), (i_k * BK, i_v * BV), (BK, BV), (1, 0)) - b_dh += tl.load(p_doq0, boundary_check=(0, 1)).to(tl.float32) - tl.store(p_dh0, b_dh.to(p_dh0.dtype.element_ty), boundary_check=(0, 1)) - - -def chunk_fwd_h( - k: torch.Tensor, - v: torch.Tensor, - g: torch.Tensor, - gk: torch.Tensor, - gv: torch.Tensor, - h0: torch.Tensor, - output_final_state: bool, - states_in_fp32: bool = False, - cu_seqlens: torch.Tensor | None = None, - chunk_size: int = 64, - chunk_indices: torch.LongTensor | None = None, -) -> tuple[torch.Tensor, torch.Tensor]: - B, T, H, K, V = *k.shape, v.shape[-1] - BT = chunk_size - - if chunk_indices is None and cu_seqlens is not None: - chunk_indices = prepare_chunk_indices(cu_seqlens, BT) - # N: the actual number of sequences in the batch with either equal or variable lengths - if cu_seqlens is None: - N, NT, chunk_offsets = B, triton.cdiv(T, BT), None - else: - N, NT, chunk_offsets = len(cu_seqlens) - 1, len(chunk_indices), prepare_chunk_offsets(cu_seqlens, BT) - - h = k.new_empty(B, NT, H, K, V, dtype=torch.float) - ht = k.new_empty(N, H, K, V, dtype=torch.float) if output_final_state else None - def grid(meta): return (triton.cdiv(K, meta['BK']) * triton.cdiv(V, meta['BV']), NT, B * H) - chunk_fwd_kernel_h_parallel[grid]( - k=k, - v=v, - h=h, - g=g, - gk=gk, - gv=gv, - h0=h0, - ht=ht, - cu_seqlens=cu_seqlens, - chunk_indices=chunk_indices, - T=T, - H=H, - K=K, - V=V, - BT=BT, - USE_G=g is not None, - USE_GK=gk is not None, - USE_GV=gv is not None, - ) - kvt, ht = ht, (torch.empty_like(ht) if output_final_state else None) - def grid(meta): return (triton.cdiv(K, meta['BK']), triton.cdiv(V, meta['BV']), N * H) - chunk_fwd_kernel_h_reduction[grid]( - h=h, - g=g, - gk=gk, - gv=gv, - kvt=kvt, - ht=ht, - cu_seqlens=cu_seqlens, - chunk_offsets=chunk_offsets, - T=T, - H=H, - K=K, - V=V, - BT=BT, - USE_G=g is not None, - USE_GK=gk is not None, - USE_GV=gv is not None, - ) - h = h.to(k.dtype) if not states_in_fp32 else h - return h, ht - - -def chunk_bwd_dh( - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - g: torch.Tensor, - gk: torch.Tensor, - gv: torch.Tensor, - do: torch.Tensor, - h0: torch.Tensor, - dht: torch.Tensor, - scale: float, - states_in_fp32: bool = False, - cu_seqlens: torch.Tensor | None = None, - chunk_size: int = 64, - chunk_indices: torch.LongTensor | None = None, -) -> tuple[torch.Tensor, torch.Tensor]: - B, T, H, K, V = *k.shape, v.shape[-1] - HQ = q.shape[2] - BT = chunk_size - - if chunk_indices is None and cu_seqlens is not None: - chunk_indices = prepare_chunk_indices(cu_seqlens, BT) - # N: the actual number of sequences in the batch with either equal or variable lengths - # NG: number of groups in GQA - if cu_seqlens is None: - N, NT, chunk_offsets = B, triton.cdiv(T, BT), None - else: - N, NT, chunk_offsets = len(cu_seqlens) - 1, len(chunk_indices), prepare_chunk_offsets(cu_seqlens, BT) - NG = HQ // H - - dh = k.new_empty(B, NT, HQ, K, V, dtype=k.dtype if not states_in_fp32 else torch.float) - dh0 = torch.empty_like(h0, dtype=torch.float) if h0 is not None else None - - def grid(meta): return (triton.cdiv(K, meta['BK']) * triton.cdiv(V, meta['BV']), NT, B * HQ) - chunk_bwd_kernel_dh_parallel[grid]( - q=q, - g=g, - gk=gk, - gv=gv, - do=do, - dh=dh, - dht=dht, - dh0=dh0, - cu_seqlens=cu_seqlens, - chunk_indices=chunk_indices, - scale=scale, - T=T, - HQ=HQ, - H=H, - K=K, - V=V, - BT=BT, - NG=NG, - USE_G=g is not None, - USE_GK=gk is not None, - USE_GV=gv is not None, - ) - - doq0, dh0 = dh0, (torch.empty_like(dh0) if dh0 is not None else None) - def grid(meta): return (triton.cdiv(K, meta['BK']), triton.cdiv(V, meta['BV']), N * HQ) - chunk_bwd_kernel_dh_reduction[grid]( - g=g, - gk=gk, - gv=gv, - dh=dh, - doq0=doq0, - dh0=dh0, - cu_seqlens=cu_seqlens, - chunk_offsets=chunk_offsets, - T=T, - HQ=HQ, - H=H, - K=K, - V=V, - BT=BT, - NG=NG, - USE_G=g is not None, - USE_GK=gk is not None, - USE_GV=gv is not None, - ) - dh = dh.to(q.dtype) if not states_in_fp32 else dh - return dh, dh0 diff --git a/build/torch-xpu/ops/common/chunk_h_split.py b/build/torch-xpu/ops/common/chunk_h_split.py deleted file mode 100644 index 3956afc77407c1a2ae77fc5b1d73fa657399bd62..0000000000000000000000000000000000000000 --- a/build/torch-xpu/ops/common/chunk_h_split.py +++ /dev/null @@ -1,603 +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 -import triton.language as tl - -from ...ops.utils.op import exp -from ...utils import autotune_cache_kwargs - - -@triton.heuristics({ - 'USE_INITIAL_STATE': lambda args: args['h0'] is not None, - 'STORE_FINAL_STATE': lambda args: args['ht'] is not None, - 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None, -}) -@triton.autotune( - configs=[ - triton.Config({'BK': BK, 'BV': BV}, num_warps=num_warps, num_stages=num_stages) - for BK in [32, 64] - for BV in [32, 64] - for num_warps in [2, 4, 8] - for num_stages in [2, 3] - ], - key=['BT', 'USE_G', 'USE_GK', 'USE_GV'], - **autotune_cache_kwargs, -) -@triton.jit(do_not_specialize=['T']) -def chunk_fwd_kernel_h_split( - k, - v, - g, - gk, - gv, - hs, - hr, - h0, - ht, - cu_seqlens, - split_indices, - T, - S: tl.constexpr, - H: tl.constexpr, - K: tl.constexpr, - V: tl.constexpr, - BT: tl.constexpr, - BK: tl.constexpr, - BV: tl.constexpr, - USE_G: tl.constexpr, - USE_GK: tl.constexpr, - USE_GV: tl.constexpr, - USE_INITIAL_STATE: tl.constexpr, - STORE_FINAL_STATE: tl.constexpr, - IS_VARLEN: tl.constexpr, -): - # handle one split at a time - # i_h: head index - # i_n: sequence index - # i_s: local split index inside a sequence - i_k, i_v, i_sh = tl.program_id(0), tl.program_id(1), tl.program_id(2) - i_ss, i_h = i_sh // H, i_sh % H - if IS_VARLEN: - i_n, i_s = tl.load(split_indices + i_ss * 2).to(tl.int32), tl.load(split_indices + i_ss * 2 + 1).to(tl.int32) - bos, eos = tl.load(cu_seqlens + i_n).to(tl.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) - T = eos - bos - NS = tl.cdiv(T, S) - else: - NS = tl.cdiv(T, S) - i_n, i_s = i_ss // NS, i_ss % NS - bos, eos = i_n * T, i_n * T + T - i_nh = i_n * H + i_h - - # [BK, BV] - b_h = tl.zeros([BK, BV], dtype=tl.float32) - # for the first split, we directly store the state as the final result - if i_s == 0: - if USE_INITIAL_STATE: - p_h0 = tl.make_block_ptr(h0 + i_nh * K*V, (K, V), (V, 1), (i_k * BK, i_v * BV), (BK, BV), (1, 0)) - b_h += tl.load(p_h0, boundary_check=(0, 1)).to(tl.float32) - p_hr = tl.make_block_ptr(hr + i_sh * K*V, (K, V), (V, 1), (i_k * BK, i_v * BV), (BK, BV), (1, 0)) - tl.store(p_hr, b_h.to(p_hr.dtype.element_ty), boundary_check=(0, 1)) - for i_t in range(tl.cdiv(i_s * S, BT), tl.cdiv(min(i_s * S + S, T), BT)): - p_k = tl.make_block_ptr(k + (bos*H + i_h) * K, (K, T), (1, H*K), (i_k * BK, i_t * BT), (BK, BT), (0, 1)) - p_v = tl.make_block_ptr(v + (bos*H + i_h) * V, (T, V), (H*V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) - # [BK, BT] - b_k = tl.load(p_k, boundary_check=(0, 1)) - # [BT, BV] - b_v = tl.load(p_v, boundary_check=(0, 1)) - last_idx = min(i_t * BT + BT, T) - 1 - - # scalar decay - if USE_G: - b_g_last = tl.load(g + bos * H + last_idx * H + i_h) - p_g = g + bos*H + (i_t * BT + tl.arange(0, BT)) * H + i_h - b_h *= exp(b_g_last) - b_g = tl.load(p_g, mask=(i_t * BT + tl.arange(0, BT) < T), other=0.) - b_v = (b_v * exp(b_g_last - b_g)[:, None]).to(b_v.dtype) - - # vector decay, h = Diag(gk) @ h - if USE_GK: - p_gk = tl.make_block_ptr(gk + (bos*H + i_h) * K, (K, T), (1, H*K), (i_k * BK, i_t * BT), (BK, BT), (0, 1)) - p_gk_last = gk + (bos + last_idx) * H*K + i_h * K + i_k * BK + tl.arange(0, BK) - - b_gk_last = tl.load(p_gk_last, mask=(i_k * BK + tl.arange(0, BK) < K), other=0.) - b_h *= exp(b_gk_last)[:, None] - - b_gk = tl.load(p_gk, boundary_check=(0, 1)) - b_k = (b_k * exp(b_gk_last[:, None] - b_gk)).to(b_k.dtype) - - # vector decay, h = h @ Diag(gv) - if USE_GV: - p_gv = tl.make_block_ptr(gv + (bos*H + i_h) * V, (T, V), (H*V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) - p_gv_last = gv + (bos + last_idx) * H*V + i_h * V + i_v * BV + tl.arange(0, BV) - - b_gv_last = tl.load(p_gv_last, mask=(i_v * BV + tl.arange(0, BV) < V), other=0.) - b_h *= exp(b_gv_last)[None, :] - - b_gv = tl.load(p_gv, boundary_check=(0, 1)) - b_v = (b_v * exp(b_gv_last[None, :] - b_gv)).to(b_v.dtype) - - b_h += tl.dot(b_k, b_v) - - # if there are more than one splits, we store the result to (unreduced) hs - # otherwise, we store the result to ht as the final state - if NS > 1: - p_hs = tl.make_block_ptr(hs + i_sh * K*V, (K, V), (V, 1), (i_k * BK, i_v * BV), (BK, BV), (1, 0)) - tl.store(p_hs, b_h.to(p_hs.dtype.element_ty), boundary_check=(0, 1)) - elif STORE_FINAL_STATE: - p_ht = tl.make_block_ptr(ht + i_nh * K*V, (K, V), (V, 1), (i_k * BK, i_v * BV), (BK, BV), (1, 0)) - tl.store(p_ht, b_h.to(p_ht.dtype.element_ty), boundary_check=(0, 1)) - - -@triton.heuristics({ - 'STORE_FINAL_STATE': lambda args: args['ht'] is not None, - 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None, -}) -@triton.autotune( - configs=[ - triton.Config({'BK': BK, 'BV': BV}, num_warps=num_warps, num_stages=num_stages) - for BK in [32, 64] - for BV in [32, 64] - for num_warps in [2, 4, 8] - for num_stages in [2, 3, 4] - ], - key=['BT', 'USE_G', 'USE_GK', 'USE_GV'], - **autotune_cache_kwargs, -) -@triton.jit(do_not_specialize=['T']) -def chunk_fwd_kernel_h_reduction( - g, - gk, - gv, - hs, - hr, - ht, - cu_seqlens, - split_offsets, - T, - S: tl.constexpr, - H: tl.constexpr, - K: tl.constexpr, - V: tl.constexpr, - BT: tl.constexpr, - BK: tl.constexpr, - BV: tl.constexpr, - USE_G: tl.constexpr, - USE_GK: tl.constexpr, - USE_GV: tl.constexpr, - STORE_FINAL_STATE: tl.constexpr, - IS_VARLEN: tl.constexpr, -): - i_k, i_v, i_nh = tl.program_id(0), tl.program_id(1), tl.program_id(2) - i_n, i_h = i_nh // H, i_nh % H - if IS_VARLEN: - bos, eos = tl.load(cu_seqlens + i_n).to(tl.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) - T = eos - bos - NS = tl.cdiv(T, S) - boh = tl.load(split_offsets + i_n).to(tl.int32) - else: - bos, eos = i_n * T, i_n * T + T - NS = tl.cdiv(T, S) - boh = i_n * NS - - b_h = tl.zeros([BK, BV], dtype=tl.float32) - # skip the first split - for i_s in range(1, NS): - p_hs = tl.make_block_ptr(hs + ((boh + i_s-1) * H + i_h) * K*V, (K, V), (V, 1), (i_k * BK, i_v * BV), (BK, BV), (1, 0)) - p_hr = tl.make_block_ptr(hr + ((boh + i_s) * H + i_h) * K*V, (K, V), (V, 1), (i_k * BK, i_v * BV), (BK, BV), (1, 0)) - b_h += tl.load(p_hs, boundary_check=(0, 1)).to(tl.float32) - tl.store(p_hr, b_h.to(p_hr.dtype.element_ty), boundary_check=(0, 1)) - - for i_t in range(tl.cdiv(i_s * S, BT), tl.cdiv(min(i_s * S + S, T), BT)): - last_idx = min(i_t * BT + BT, T) - 1 - # scalar decay - if USE_G: - b_g_last = tl.load(g + bos * H + last_idx * H + i_h) - b_h *= exp(b_g_last) - - # vector decay, h = Diag(gk) @ h - if USE_GK: - p_gk_last = gk + (bos + last_idx) * H*K + i_h * K + i_k * BK + tl.arange(0, BK) - b_gk_last = tl.load(p_gk_last, mask=(i_k * BK + tl.arange(0, BK) < K), other=0.) - b_h *= exp(b_gk_last)[:, None] - - # vector decay, h = h @ Diag(gv) - if USE_GV: - p_gv_last = gv + (bos + last_idx) * H*V + i_h * V + i_v * BV + tl.arange(0, BV) - b_gv_last = tl.load(p_gv_last, mask=(i_v * BV + tl.arange(0, BV) < V), other=0.) - b_h *= exp(b_gv_last)[None, :] - - if NS > 1: - if STORE_FINAL_STATE: - p_hs = tl.make_block_ptr(hs + ((boh + NS-1) * H + i_h)*K*V, (K, V), (V, 1), (i_k * BK, i_v * BV), (BK, BV), (1, 0)) - p_ht = tl.make_block_ptr(ht + i_nh * K*V, (K, V), (V, 1), (i_k * BK, i_v * BV), (BK, BV), (1, 0)) - b_h += tl.load(p_hs, boundary_check=(0, 1)).to(tl.float32) - tl.store(p_ht, b_h.to(p_ht.dtype.element_ty), boundary_check=(0, 1)) - - -@triton.heuristics({ - 'USE_FINAL_STATE_GRADIENT': lambda args: args['dht'] is not None, - 'STORE_INITIAL_STATE_GRADIENT': lambda args: args['dh0'] is not None, - 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None, -}) -@triton.autotune( - configs=[ - triton.Config({'BK': BK, 'BV': BV}, num_warps=num_warps, num_stages=num_stages) - for BK in [32, 64] - for BV in [32, 64] - for num_warps in [2, 4, 8] - for num_stages in [2, 3] - ], - key=['BT', 'USE_G', 'USE_GK', 'USE_GV'], - **autotune_cache_kwargs, -) -@triton.jit(do_not_specialize=['T']) -def chunk_bwd_kernel_dh_split( - q, - g, - gk, - gv, - do, - dht, - dhs, - dhr, - dh0, - cu_seqlens, - split_indices, - scale, - T, - S: tl.constexpr, - HQ: tl.constexpr, - H: tl.constexpr, - K: tl.constexpr, - V: tl.constexpr, - BT: tl.constexpr, - BK: tl.constexpr, - BV: tl.constexpr, - NG: tl.constexpr, - USE_G: tl.constexpr, - USE_GK: tl.constexpr, - USE_GV: tl.constexpr, - USE_FINAL_STATE_GRADIENT: tl.constexpr, - STORE_INITIAL_STATE_GRADIENT: tl.constexpr, - IS_VARLEN: tl.constexpr, -): - # handle one split at a time - # i_h: head index - # i_n: sequence index - # i_s: local split index inside a sequence - i_k, i_v, i_sh = tl.program_id(0), tl.program_id(1), tl.program_id(2) - i_ss, i_hq = i_sh // HQ, i_sh % HQ - if IS_VARLEN: - i_n, i_s = tl.load(split_indices + i_ss * 2).to(tl.int32), tl.load(split_indices + i_ss * 2 + 1).to(tl.int32) - bos, eos = tl.load(cu_seqlens + i_n).to(tl.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) - T = eos - bos - NS = tl.cdiv(T, S) - else: - NS = tl.cdiv(T, S) - i_n, i_s = i_ss // NS, i_ss % NS - bos, eos = i_n * T, i_n * T + T - i_nh = i_n * HQ + i_hq - i_h = i_hq // NG - - # [BK, BV] - b_dh = tl.zeros([BK, BV], dtype=tl.float32) - if i_s == NS - 1: - if USE_FINAL_STATE_GRADIENT: - p_dht = tl.make_block_ptr(dht + i_nh * K*V, (K, V), (V, 1), (i_k * BK, i_v * BV), (BK, BV), (1, 0)) - b_dh += tl.load(p_dht, boundary_check=(0, 1)).to(tl.float32) - p_dhr = tl.make_block_ptr(dhr + i_sh * K*V, (K, V), (V, 1), (i_k * BK, i_v * BV), (BK, BV), (1, 0)) - tl.store(p_dhr, b_dh.to(p_dhr.dtype.element_ty), boundary_check=(0, 1)) - - for i_t in range(tl.cdiv(min(i_s * S + S, T), BT) - 1, tl.cdiv(i_s * S, BT) - 1, -1): - p_q = tl.make_block_ptr(q + (bos*HQ + i_hq) * K, (K, T), (1, HQ*K), (i_k * BK, i_t * BT), (BK, BT), (0, 1)) - p_do = tl.make_block_ptr(do + (bos*HQ + i_hq) * V, (T, V), (HQ*V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) - - b_q = tl.load(p_q, boundary_check=(0, 1)) - b_q = (b_q * scale).to(b_q.dtype) - # [BT, BV] - b_do = tl.load(p_do, boundary_check=(0, 1)) - - last_idx = min(i_t * BT + BT, T) - 1 - if USE_G: - p_g = g + (bos + i_t * BT + tl.arange(0, BT)) * H + i_h - b_g_last = tl.load(g + (bos + last_idx) * H + i_h) - b_g = tl.load(p_g, mask=(i_t * BT + tl.arange(0, BT) < T), other=0.) - b_q = (b_q * exp(b_g)[None, :]).to(b_q.dtype) - b_dh *= exp(b_g_last) - - if USE_GK: - p_gk = tl.make_block_ptr(gk + (bos*H + i_h) * K, (K, T), (1, H*K), (i_k * BK, i_t * BT), (BK, BT), (0, 1)) - p_gk_last = gk + (bos + last_idx) * H*K + i_h * K + i_k * BK + tl.arange(0, BK) - - b_gk = tl.load(p_gk, boundary_check=(0, 1)) - b_q = (b_q * exp(b_gk)).to(b_q.dtype) - b_gk_last = tl.load(p_gk_last, mask=(i_k * BK + tl.arange(0, BK) < K), other=0.) - b_dh *= exp(b_gk_last)[:, None] - - if USE_GV: - p_gv = tl.make_block_ptr(gv + (bos*H + i_h) * V, (T, V), (H*V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) - p_gv_last = gv + (bos + last_idx) * H*V + i_h * V + i_v * BV + tl.arange(0, BV) - - b_gv = tl.load(p_gv, boundary_check=(0, 1)) - b_do = (b_do * exp(b_gv)).to(b_do.dtype) - - b_gv_last = tl.load(p_gv_last, mask=(i_v * BV + tl.arange(0, BV) < V), other=0.) - b_dh *= exp(b_gv_last)[None, :] - - b_dh += tl.dot(b_q, b_do) - - if NS > 1: - p_dhs = tl.make_block_ptr(dhs + i_sh * K*V, (K, V), (V, 1), (i_k * BK, i_v * BV), (BK, BV), (1, 0)) - tl.store(p_dhs, b_dh.to(p_dhs.dtype.element_ty), boundary_check=(0, 1)) - elif STORE_INITIAL_STATE_GRADIENT: - p_dh0 = tl.make_block_ptr(dh0 + i_nh * K*V, (K, V), (V, 1), (i_k * BK, i_v * BV), (BK, BV), (1, 0)) - tl.store(p_dh0, b_dh.to(p_dh0.dtype.element_ty), boundary_check=(0, 1)) - - -@triton.heuristics({ - 'STORE_INITIAL_STATE_GRADIENT': lambda args: args['dh0'] is not None, - 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None, -}) -@triton.autotune( - configs=[ - triton.Config({'BK': BK, 'BV': BV}, num_warps=num_warps, num_stages=num_stages) - for BK in [32, 64] - for BV in [32, 64] - for num_warps in [2, 4, 8] - for num_stages in [2, 3, 4] - ], - key=['BT', 'USE_G', 'USE_GK', 'USE_GV'], - **autotune_cache_kwargs, -) -@triton.jit(do_not_specialize=['T']) -def chunk_bwd_kernel_dh_reduction( - g, - gk, - gv, - dhs, - dhr, - dh0, - cu_seqlens, - split_offsets, - T, - S: tl.constexpr, - H: tl.constexpr, - HQ: tl.constexpr, - K: tl.constexpr, - V: tl.constexpr, - BT: tl.constexpr, - BK: tl.constexpr, - BV: tl.constexpr, - NG: tl.constexpr, - USE_G: tl.constexpr, - USE_GK: tl.constexpr, - USE_GV: tl.constexpr, - STORE_INITIAL_STATE_GRADIENT: tl.constexpr, - IS_VARLEN: tl.constexpr, -): - i_k, i_v, i_nh = tl.program_id(0), tl.program_id(1), tl.program_id(2) - i_n, i_hq = i_nh // HQ, i_nh % HQ - i_h = i_hq // NG - if IS_VARLEN: - bos, eos = tl.load(cu_seqlens + i_n).to(tl.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) - T = eos - bos - NS = tl.cdiv(T, S) - boh = tl.load(split_offsets + i_n).to(tl.int32) - else: - bos, eos = i_n * T, i_n * T + T - NS = tl.cdiv(T, S) - boh = i_n * NS - - b_dh = tl.zeros([BK, BV], dtype=tl.float32) - for i_s in range(NS - 2, -1, -1): - p_dhs = tl.make_block_ptr(dhs + ((boh+i_s+1) * H + i_h) * K*V, (K, V), (V, 1), (i_k * BK, i_v * BV), (BK, BV), (1, 0)) - p_dhr = tl.make_block_ptr(dhr + ((boh+i_s) * H + i_h) * K*V, (K, V), (V, 1), (i_k * BK, i_v * BV), (BK, BV), (1, 0)) - b_dh += tl.load(p_dhs, boundary_check=(0, 1)).to(tl.float32) - tl.store(p_dhr, b_dh.to(p_dhr.dtype.element_ty), boundary_check=(0, 1)) - - for i_t in range(tl.cdiv(min(i_s * S + S, T), BT) - 1, tl.cdiv(i_s * S, BT) - 1, -1): - last_idx = min(i_t * BT + BT, T) - 1 - # scalar decay - if USE_G: - b_g_last = tl.load(g + (bos + last_idx) * H + i_h) - b_dh *= exp(b_g_last) - - if USE_GK: - p_gk_last = gk + (bos + last_idx) * H*K + i_h * K + i_k * BK + tl.arange(0, BK) - b_gk_last = tl.load(p_gk_last, mask=(i_k * BK + tl.arange(0, BK) < K), other=0.) - b_dh *= exp(b_gk_last)[:, None] - - if USE_GV: - p_gv_last = gv + (bos + last_idx) * H*V + i_h * V + i_v * BV + tl.arange(0, BV) - b_gv_last = tl.load(p_gv_last, mask=(i_v * BV + tl.arange(0, BV) < V), other=0.) - b_dh *= exp(b_gv_last)[None, :] - - if NS > 1: - if STORE_INITIAL_STATE_GRADIENT: - p_dhs = tl.make_block_ptr(dhs + (boh * H + i_h)*K*V, (K, V), (V, 1), (i_k * BK, i_v * BV), (BK, BV), (1, 0)) - p_dh0 = tl.make_block_ptr(dh0 + i_nh * K*V, (K, V), (V, 1), (i_k * BK, i_v * BV), (BK, BV), (1, 0)) - b_dh += tl.load(p_dhs, boundary_check=(0, 1)).to(tl.float32) - tl.store(p_dh0, b_dh.to(p_dh0.dtype.element_ty), boundary_check=(0, 1)) - - -def chunk_fwd_h( - k: torch.Tensor, - v: torch.Tensor, - g: torch.Tensor, - gk: torch.Tensor, - gv: torch.Tensor, - h0: torch.Tensor, - output_final_state: bool, - cu_seqlens: torch.LongTensor | None = None, - split_offsets: torch.LongTensor | None = None, - split_indices: torch.LongTensor | None = None, - chunk_size: int = 64, - split_size: int = 256, - states_in_fp32: bool = True, -) -> tuple[torch.Tensor, torch.Tensor]: - B, T, H, K, V = *k.shape, v.shape[-1] - # B: batch size - # N: the actual number of sequences in the batch - # H: number of heads - # T: sequence length, can be variable across sequences - # S: split size, a multiple of chunk size - # BT: chunk size - S, BT = split_size, chunk_size - assert S % BT == 0, f"The `split_size` (got {S}) must be a multiple of `chunk_size` {BT}" - if cu_seqlens is None: - N = B - NS = N * triton.cdiv(T, S) - else: - N = len(cu_seqlens) - 1 - NS = split_offsets[-1] - - # unreduced kv states per split - hs = k.new_empty(NS, H, K, V, dtype=torch.float) - # reduced states per split - hr = k.new_empty(NS, H, K, V, dtype=torch.float if states_in_fp32 else k.dtype) - ht = k.new_empty(N, H, K, V, dtype=torch.float) if output_final_state else None - # parallelized over splits - def grid(meta): return (triton.cdiv(K, meta['BK']), triton.cdiv(V, meta['BV']), NS * H) - chunk_fwd_kernel_h_split[grid]( - k=k, - v=v, - g=g, - gk=gk, - gv=gv, - hs=hs, - hr=hr, - h0=h0, - ht=ht, - cu_seqlens=cu_seqlens, - split_indices=split_indices, - T=T, - S=S, - H=H, - K=K, - V=V, - BT=BT, - USE_G=g is not None, - USE_GK=gk is not None, - USE_GV=gv is not None, - ) - def grid(meta): return (triton.cdiv(K, meta['BK']), triton.cdiv(V, meta['BV']), N * H) - chunk_fwd_kernel_h_reduction[grid]( - g=g, - gk=gk, - gv=gv, - hs=hs, - hr=hr, - ht=ht, - cu_seqlens=cu_seqlens, - split_offsets=split_offsets, - T=T, - S=S, - H=H, - K=K, - V=V, - BT=BT, - USE_G=g is not None, - USE_GK=gk is not None, - USE_GV=gv is not None, - ) - return hr, ht - - -def chunk_bwd_dh( - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - g: torch.Tensor, - gk: torch.Tensor, - gv: torch.Tensor, - do: torch.Tensor, - h0: torch.Tensor, - dht: torch.Tensor, - scale: float, - cu_seqlens: torch.Tensor | None = None, - split_offsets: torch.Tensor | None = None, - split_indices: torch.Tensor | None = None, - chunk_size: int = 64, - split_size: int = 256, - states_in_fp32: bool = True, -) -> tuple[torch.Tensor, torch.Tensor]: - B, T, H, K, V = *k.shape, v.shape[-1] - HQ = q.shape[2] - # B: batch size - # N: the actual number of sequences in the batch - # H: number of heads - # T: sequence length, can be variable across sequences - # S: split size, a multiple of chunk size - # BT: chunk size - S, BT = max(chunk_size, min(split_size, triton.next_power_of_2(T))), chunk_size - assert S % BT == 0, f"The `split_size` (got {S}) must be a multiple of `chunk_size` {BT}" - if cu_seqlens is None: - N = B - NS = N * triton.cdiv(T, S) - else: - N = len(cu_seqlens) - 1 - NS = split_offsets[-1] - # number of groups in GQA - NG = HQ // H - - dhs = q.new_empty(NS, HQ, K, V, dtype=torch.float) - dhr = q.new_empty(NS, HQ, K, V, dtype=torch.float if states_in_fp32 else k.dtype) - dh0 = torch.empty_like(h0, dtype=torch.float) if h0 is not None else None - - # parallelized over splits - def grid(meta): return (triton.cdiv(K, meta['BK']), triton.cdiv(V, meta['BV']), NS * HQ) - chunk_bwd_kernel_dh_split[grid]( - q=q, - g=g, - gk=gk, - gv=gv, - do=do, - dht=dht, - dhs=dhs, - dhr=dhr, - dh0=dh0, - cu_seqlens=cu_seqlens, - split_indices=split_indices, - scale=scale, - T=T, - S=S, - HQ=HQ, - H=H, - K=K, - V=V, - BT=BT, - NG=NG, - USE_G=g is not None, - USE_GK=gk is not None, - USE_GV=gv is not None, - ) - - def grid(meta): return (triton.cdiv(K, meta['BK']), triton.cdiv(V, meta['BV']), N * HQ) - chunk_bwd_kernel_dh_reduction[grid]( - g=g, - gk=gk, - gv=gv, - dhs=dhs, - dhr=dhr, - dh0=dh0, - cu_seqlens=cu_seqlens, - split_offsets=split_offsets, - T=T, - S=S, - HQ=HQ, - H=H, - K=K, - V=V, - BT=BT, - NG=NG, - USE_G=g is not None, - USE_GK=gk is not None, - USE_GV=gv is not None, - ) - return dhr, dh0 diff --git a/build/torch-xpu/ops/common/chunk_o.py b/build/torch-xpu/ops/common/chunk_o.py deleted file mode 100644 index 868d012b9733cfa156e3429c5ee696e24f998792..0000000000000000000000000000000000000000 --- a/build/torch-xpu/ops/common/chunk_o.py +++ /dev/null @@ -1,758 +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 -import triton.language as tl - -from ...ops.common.backends import dispatch -from ...ops.utils import prepare_chunk_indices -from ...ops.utils.cache import fla_cache_autotune -from ...ops.utils.op import exp2 -from ...utils import ( - IS_INTEL, - IS_NVIDIA_HOPPER, - TRITON_ABOVE_3_4_0, - TRITON_ABOVE_3_7_1, - autotune_cache_kwargs, - check_shared_mem, -) - -BKV_LIST = [64, 128] if check_shared_mem() else ([32, 64] if check_shared_mem('ada') else [32]) -NUM_WARPS = [2, 4] if IS_NVIDIA_HOPPER else [2, 4, 8] - -# Upstream pairs each tile size with a single warp count, which is tuned for NVIDIA. -# On Intel that pairing is off by a factor of two: BK=BV=64 is fastest at 8 warps but is -# only offered at 4, so the autotuner falls back to the 32x32 config and leaves ~2.2x on -# the table. Widen the space there instead of changing the defaults for other vendors. -_O_CONFIGS = [ - triton.Config({'BK': 128, 'BV': 128}, num_warps=8, num_stages=3), - triton.Config({'BK': 64, 'BV': 64}, num_warps=4, num_stages=3), - triton.Config({'BK': 32, 'BV': 32}, num_warps=2, num_stages=3), -] -if IS_INTEL: - _O_CONFIGS += [ - triton.Config({'BK': BK, 'BV': BV}, num_warps=w, num_stages=2) - for BK, BV in [(128, 64), (64, 64), (64, 32), (32, 64), (32, 32)] - for w in (4, 8) - ] - - -@triton.heuristics({ - 'USE_G': lambda args: args['g'] is not None, - 'USE_G_GAMMA': lambda args: args['g_gamma'] is not None, - 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None, -}) -@fla_cache_autotune( - configs=_O_CONFIGS, - key=['H', 'HV', 'K', 'V', 'BT', 'STATE_V_FIRST'], - **autotune_cache_kwargs, -) -@triton.jit(do_not_specialize=['T']) -def chunk_fwd_kernel_o( - q, - k, - v, - h, - g, - g_gamma, - o, - cu_seqlens, - chunk_indices, - scale, - T, - H: tl.constexpr, - HV: tl.constexpr, - K: tl.constexpr, - V: tl.constexpr, - BT: tl.constexpr, - BK: tl.constexpr, - BV: tl.constexpr, - USE_G: tl.constexpr, - USE_G_GAMMA: tl.constexpr, - STATE_V_FIRST: tl.constexpr, - IS_VARLEN: tl.constexpr, -): - i_v, i_t, i_bh = tl.program_id(0), tl.program_id(1), tl.program_id(2).to(tl.int64) - i_b, i_h = i_bh // HV, i_bh % HV - - 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.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) - T = eos - bos - NT = tl.cdiv(T, BT) - else: - NT = tl.cdiv(T, BT) - i_tg = i_b * NT + i_t - bos, eos = i_b * T, i_b * T + T - - # offset calculation - q += (bos * H + i_h // (HV // H)) * K - k += (bos * H + i_h // (HV // H)) * K - v += (bos * HV + i_h) * V - o += (bos * HV + i_h) * V - h += (i_tg * HV + i_h).to(tl.int64) * K*V - - b_o = tl.zeros([BT, BV], dtype=tl.float32) - b_A = tl.zeros([BT, BT], dtype=tl.float32) - - for i_k in range(tl.cdiv(K, BK)): - p_q = tl.make_block_ptr(q, (T, K), (H*K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) - p_k = tl.make_block_ptr(k, (K, T), (1, H*K), (i_k * BK, i_t * BT), (BK, BT), (0, 1)) - if STATE_V_FIRST: - p_h = tl.make_block_ptr(h, (V, K), (K, 1), (i_v * BV, i_k * BK), (BV, BK), (1, 0)) - else: - p_h = tl.make_block_ptr(h, (K, V), (V, 1), (i_k * BK, i_v * BV), (BK, BV), (1, 0)) - # [BT, BK] - b_q = tl.load(p_q, boundary_check=(0, 1)) - # [BK, BT] - b_k = tl.load(p_k, boundary_check=(0, 1)) - b_h = tl.load(p_h, boundary_check=(0, 1)) - - # [BT, BK] @ [BK, BV] -> [BT, BV] - if STATE_V_FIRST: - b_o += tl.dot(b_q, tl.trans(b_h)) - else: - b_o += tl.dot(b_q, b_h) - # [BT, BK] @ [BK, BT] -> [BT, BT] - b_A += tl.dot(b_q, b_k) - - if USE_G: - g += bos * HV + i_h - p_g = tl.make_block_ptr(g, (T,), (HV,), (i_t * BT,), (BT,), (0,)) - b_g = tl.load(p_g, boundary_check=(0,)) - b_o = b_o * exp2(b_g)[:, None] - b_A = b_A * exp2(b_g[:, None] - b_g[None, :]) - if USE_G_GAMMA: - b_gamma = tl.load(g_gamma + i_h) - b_g = b_gamma * (tl.arange(0, BT) + 1) - b_o = b_o * exp2(b_g)[:, None] - b_A = b_A * exp2(b_g[:, None] - b_g[None, :]) - o_t = i_t * BT + tl.arange(0, BT) - m_t = o_t < T - m_A = (o_t[:, None] >= o_t[None, :]) & (m_t[:, None] & m_t) - b_A = tl.where(m_A, b_A, 0) - - p_v = tl.make_block_ptr(v, (T, V), (HV*V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) - p_o = tl.make_block_ptr(o, (T, V), (HV*V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) - - b_v = tl.load(p_v, boundary_check=(0, 1)) - # to fix mma -> mma layout conversion - # already solved by triton v3.2 or higher - b_o = b_o * scale + tl.dot(b_A.to(b_v.dtype), b_v) * scale - tl.store(p_o, b_o.to(p_o.dtype.element_ty), boundary_check=(0, 1)) - - -@triton.heuristics({ - 'USE_G': lambda args: args['g'] is not None, - 'USE_G_GAMMA': lambda args: args['g_gamma'] is not None, - 'USE_DW': lambda args: args['dw'] is not None, - 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None, -}) -@fla_cache_autotune( - configs=[ - triton.Config({}, num_warps=num_warps, num_stages=num_stages) - for num_warps in NUM_WARPS - for num_stages in [2, 3, 4] - ], - key=['H', 'HV', 'K', 'V', 'BT', 'BK', 'BV', 'USE_G', 'USE_G_GAMMA', 'USE_DW', 'STATE_V_FIRST'], - **autotune_cache_kwargs, -) -@triton.jit(do_not_specialize=['T']) -def chunk_bwd_kernel_dqkwg( - q, - k, - v, - g, - g_gamma, - h, - do, - dh, - dq, - dk, - dw, - dv, - dg, - cu_seqlens, - chunk_indices, - scale, - B: tl.constexpr, - T, - H: tl.constexpr, - HV: tl.constexpr, - K: tl.constexpr, - V: tl.constexpr, - BT: tl.constexpr, - BK: tl.constexpr, - BV: tl.constexpr, - USE_G: tl.constexpr, - USE_G_GAMMA: tl.constexpr, - USE_DW: tl.constexpr, - STATE_V_FIRST: tl.constexpr, - IS_VARLEN: tl.constexpr, -): - i_k, i_t, i_bh = tl.program_id(0), tl.program_id(1), tl.program_id(2).to(tl.int64) - i_b, i_h = i_bh // HV, i_bh % HV - - all = B * T - 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.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) - T = eos - bos - NT = tl.cdiv(T, BT) - else: - NT = tl.cdiv(T, BT) - i_tg = i_b * NT + i_t - bos, eos = i_b * T, i_b * T + T - - # offset calculation - v += (bos * HV + i_h) * V - do += (bos * HV + i_h) * V - h += (i_tg * HV + i_h).to(tl.int64) * K*V - dh += (i_tg * HV + i_h).to(tl.int64) * K*V - q += (bos * H + i_h // (HV // H)) * K - k += (bos * H + i_h // (HV // H)) * K - dq += (bos * HV + i_h) * K - dk += (bos * HV + i_h) * K - - # for delta rule only - if USE_DW: - dw += (bos * HV + i_h) * K - dv += (bos * HV + i_h) * V - - if USE_G: - dg += i_k * all * HV - b_dg_last = tl.zeros([1], dtype=tl.float32) if USE_G else None - if USE_G_GAMMA: - b_gamma = tl.load(g_gamma + i_h) - b_g = b_gamma * (tl.arange(0, BT) + 1) - b_g_last = b_gamma * min(BT, T - i_t * BT) - b_dq = tl.zeros([BT, BK], dtype=tl.float32) - b_dk = tl.zeros([BT, BK], dtype=tl.float32) - b_ds = tl.zeros([BT, BT], dtype=tl.float32) - b_dw = tl.zeros([BT, BK], dtype=tl.float32) if USE_DW else None - - for i_v in range(tl.cdiv(V, BV)): - p_v = tl.make_block_ptr(v, (T, V), (HV*V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) - p_do = tl.make_block_ptr(do, (T, V), (HV*V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) - if STATE_V_FIRST: - p_h = tl.make_block_ptr(h, (V, K), (K, 1), (i_v * BV, i_k * BK), (BV, BK), (1, 0)) - p_dh = tl.make_block_ptr(dh, (V, K), (K, 1), (i_v * BV, i_k * BK), (BV, BK), (1, 0)) - else: - p_h = tl.make_block_ptr(h, (V, K), (1, V), (i_v * BV, i_k * BK), (BV, BK), (0, 1)) - p_dh = tl.make_block_ptr(dh, (V, K), (1, V), (i_v * BV, i_k * BK), (BV, BK), (0, 1)) - # [BT, BV] - b_v = tl.load(p_v, boundary_check=(0, 1)) - b_do = tl.load(p_do, boundary_check=(0, 1)) - # [BV, BK] - b_h = tl.load(p_h, boundary_check=(0, 1)) - b_dh = tl.load(p_dh, boundary_check=(0, 1)) - if USE_G: - b_dg_last += (tl.sum(b_h * b_dh)) - # [BT, BV] @ [BV, BT] -> [BT, BT] - b_ds += tl.dot(b_do, tl.trans(b_v)) - # [BT, BV] @ [BV, BK] -> [BT, BK] - b_dq += tl.dot(b_do, b_h.to(b_do.dtype)) - # [BT, BV] @ [BV, BK] -> [BT, BK] - b_dk += tl.dot(b_v, b_dh.to(b_v.dtype)) - if USE_DW: - p_dv = tl.make_block_ptr(dv, (T, V), (HV*V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) - b_dv = tl.load(p_dv, boundary_check=(0, 1)) - b_dw += tl.dot(b_dv.to(b_v.dtype), b_h.to(b_v.dtype)) - - if USE_DW: - p_dw = tl.make_block_ptr(dw, (T, K), (HV*K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) - tl.store(p_dw, -b_dw.to(p_dw.dtype.element_ty), boundary_check=(0, 1)) - - tl.debug_barrier() - p_q = tl.make_block_ptr(q, (T, K), (H*K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) - p_k = tl.make_block_ptr(k, (T, K), (H*K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) - b_q = tl.load(p_q, boundary_check=(0, 1)) - b_k = tl.load(p_k, boundary_check=(0, 1)) - - p_dq = tl.make_block_ptr(dq, (T, K), (HV*K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) - p_dk = tl.make_block_ptr(dk, (T, K), (HV*K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) - - o_t = i_t * BT + tl.arange(0, BT) - m_t = o_t < T - m_A = (o_t[:, None] >= o_t[None, :]) & (m_t[:, None] & m_t) - if USE_G: - g += bos * HV + i_h - dg += bos * HV + i_h - p_g = tl.make_block_ptr(g, (T,), (HV,), (i_t * BT,), (BT,), (0,)) - b_g = tl.load(p_g, boundary_check=(0,)) - b_g_last = tl.load(g + (min(i_t * BT + BT, T) - 1) * HV) - b_dg_last *= exp2(b_g_last) - b_dq = b_dq * exp2(b_g)[:, None] * scale - b_dk = b_dk * tl.where(m_t, exp2(-b_g + b_g_last), 0)[:, None] - b_dg_last += tl.sum(b_dk * b_k) - - b_ds = tl.where(m_A, b_ds * exp2(b_g[:, None] - b_g[None, :]), 0) * scale - b_ds = b_ds.to(b_k.dtype) - # [BT, BK] - b_dq += tl.dot(b_ds, b_k) - b_dk += tl.dot(tl.trans(b_ds), b_q) - - b_dg = tl.sum(b_dq * b_q, axis=1) - tl.sum(b_dk * b_k, axis=1) - - p_dg = tl.make_block_ptr(dg, (T,), (HV,), (i_t * BT,), (BT,), (0,)) - # (SY 09/21) revcumsum in a separate kernel due to strange triton compiler issue - # b_dg = tl.dot(tl.where(o_t[:, None] <= o_t[None, :], 1., 0.), b_dg, allow_tf32=False) + b_dg_last) - b_dg = tl.where(o_t < min(i_t * BT + BT, T) - 1, b_dg, b_dg + b_dg_last) - tl.store(p_dq, b_dq.to(p_dq.dtype.element_ty), boundary_check=(0, 1)) - tl.store(p_dk, b_dk.to(p_dk.dtype.element_ty), boundary_check=(0, 1)) - tl.store(p_dg, b_dg.to(p_dg.dtype.element_ty), boundary_check=(0,)) - - elif USE_G_GAMMA: - b_dq = b_dq * exp2(b_g)[:, None] * scale - b_dk = b_dk * tl.where(m_t, exp2(-b_g + b_g_last), 0)[:, None] - b_ds = tl.where(m_A, b_ds * exp2(b_g[:, None] - b_g[None, :]), 0) * scale - b_ds = b_ds.to(b_k.dtype) - # [BT, BK] - b_dq += tl.dot(b_ds, b_k) - b_dk += tl.dot(tl.trans(b_ds), b_q) - tl.store(p_dq, b_dq.to(p_dq.dtype.element_ty), boundary_check=(0, 1)) - tl.store(p_dk, b_dk.to(p_dk.dtype.element_ty), boundary_check=(0, 1)) - - else: - b_ds = tl.where(m_A, b_ds, 0) - b_ds = b_ds.to(b_k.dtype) - b_dq += tl.dot(b_ds, b_k) - b_dk += tl.dot(tl.trans(b_ds), b_q) * scale - b_dq *= scale - tl.store(p_dq, b_dq.to(p_dq.dtype.element_ty), boundary_check=(0, 1)) - tl.store(p_dk, b_dk.to(p_dk.dtype.element_ty), boundary_check=(0, 1)) - - -@triton.heuristics({ - 'USE_G': lambda args: args['g'] is not None, - 'USE_G_GAMMA': lambda args: args['g_gamma'] is not None, - 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None, -}) -@triton.autotune( - configs=[ - triton.Config({}, num_warps=num_warps, num_stages=num_stages) - for num_warps in NUM_WARPS - for num_stages in [2, 3, 4] - ], - key=['H', 'HV', 'K', 'V', 'BT', 'BK', 'BV', 'USE_G', 'USE_G_GAMMA', 'STATE_V_FIRST'], - **autotune_cache_kwargs, -) -@triton.jit(do_not_specialize=['T']) -def chunk_bwd_kernel_dv( - q, - k, - g, - g_gamma, - do, - dv, - dh, - cu_seqlens, - chunk_indices, - scale, - T, - H: tl.constexpr, - HV: tl.constexpr, - K: tl.constexpr, - V: tl.constexpr, - BT: tl.constexpr, - BK: tl.constexpr, - BV: tl.constexpr, - USE_G: tl.constexpr, - USE_G_GAMMA: tl.constexpr, - IS_VARLEN: tl.constexpr, - STATE_V_FIRST: tl.constexpr, -): - i_v, i_t, i_bh = tl.program_id(0), tl.program_id(1), tl.program_id(2).to(tl.int64) - i_b, i_h = i_bh // HV, i_bh % HV - 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.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) - T = eos - bos - NT = tl.cdiv(T, BT) - else: - NT = tl.cdiv(T, BT) - i_tg = i_b * NT + i_t - bos, eos = i_b * T, i_b * T + T - - b_dv = tl.zeros([BT, BV], dtype=tl.float32) - - # offset calculation - q += (bos * H + i_h // (HV // H)) * K - k += (bos * H + i_h // (HV // H)) * K - do += (bos * HV + i_h) * V - dv += (bos * HV + i_h) * V - dh += (i_tg * HV + i_h).to(tl.int64) * K*V - - b_A = tl.zeros([BT, BT], dtype=tl.float32) - for i_k in range(tl.cdiv(K, BK)): - p_k = tl.make_block_ptr(k, (T, K), (H*K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) - p_q = tl.make_block_ptr(q, (K, T), (1, H*K), (i_k * BK, i_t * BT), (BK, BT), (0, 1)) - b_q = tl.load(p_q, boundary_check=(0, 1)) - b_k = tl.load(p_k, boundary_check=(0, 1)) - b_A += tl.dot(b_k, b_q) - if STATE_V_FIRST: - p_dh = tl.make_block_ptr(dh, (V, K), (K, 1), (i_v * BV, i_k * BK), (BV, BK), (1, 0)) - b_dh = tl.trans(tl.load(p_dh, boundary_check=(0, 1))) - else: - p_dh = tl.make_block_ptr(dh, (K, V), (V, 1), (i_k * BK, i_v * BV), (BK, BV), (1, 0)) - b_dh = tl.load(p_dh, boundary_check=(0, 1)) - b_dv += tl.dot(b_k, b_dh.to(b_k.dtype)) - - o_t = i_t * BT + tl.arange(0, BT) - m_t = o_t < T - if USE_G: - g += bos * HV + i_h - p_g = tl.make_block_ptr(g, (T,), (HV,), (i_t * BT,), (BT,), (0,)) - b_g = tl.load(p_g, boundary_check=(0,)) - b_g_last = tl.load(g + (min(i_t * BT + BT, T) - 1) * HV) - if USE_G_GAMMA: - b_gamma = tl.load(g_gamma + i_h) - b_g = b_gamma * (tl.arange(0, BT) + 1) - b_g_last = b_gamma * min(BT, T - i_t * BT) - - m_A = (o_t[:, None] <= o_t[None, :]) & (m_t[:, None] & m_t) - if USE_G or USE_G_GAMMA: - b_A = tl.where(m_A, b_A * exp2(b_g[None, :] - b_g[:, None]) * scale, 0).to(do.dtype.element_ty) - b_dv *= tl.where(m_t, exp2(-b_g + b_g_last), 0)[:, None] - else: - b_A = tl.where(m_A, b_A * scale, 0).to(do.dtype.element_ty) - p_do = tl.make_block_ptr(do, (T, V), (HV*V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) - p_dv = tl.make_block_ptr(dv, (T, V), (HV*V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) - b_do = tl.load(p_do, boundary_check=(0, 1)) - b_dv += tl.dot(b_A.to(b_do.dtype), b_do) - tl.store(p_dv, b_dv.to(p_dv.dtype.element_ty), boundary_check=(0, 1)) - - -@triton.heuristics({ - 'USE_G': lambda args: args['g'] is not None, - 'USE_G_GAMMA': lambda args: args['g_gamma'] is not None, - 'USE_A': lambda args: args['A'] is not None, - 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None, -}) -@fla_cache_autotune( - configs=[ - triton.Config({}, num_warps=num_warps, num_stages=num_stages) - for num_warps in NUM_WARPS - for num_stages in [2, 3, 4] - ], - key=['H', 'HV', 'K', 'V', 'BT', 'BK', 'BV', 'USE_G'], - **autotune_cache_kwargs, -) -@triton.jit(do_not_specialize=['T']) -def chunk_bwd_kernel_dv_local( - q, - k, - g, - g_gamma, - A, - do, - dv, - cu_seqlens, - chunk_indices, - scale, - T, - H: tl.constexpr, - HV: tl.constexpr, - K: tl.constexpr, - V: tl.constexpr, - BT: tl.constexpr, - BK: tl.constexpr, - BV: tl.constexpr, - USE_G: tl.constexpr, - USE_G_GAMMA: tl.constexpr, - USE_A: tl.constexpr, - IS_VARLEN: tl.constexpr, -): - i_t, i_bh = tl.program_id(0), tl.program_id(1).to(tl.int64) - i_b, i_h = i_bh // HV, i_bh % HV - 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.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) - T = eos - bos - else: - bos, eos = i_b * T, i_b * T + T - - # offset calculation - q += (bos * H + i_h // (HV // H)) * K - k += (bos * H + i_h // (HV // H)) * K - do += (bos * HV + i_h) * V - dv += (bos * HV + i_h) * V - - if USE_A: - p_A = tl.make_block_ptr(A + (bos * HV + i_h) * BT, (BT, T), (1, HV*BT), (0, i_t * BT), (BT, BT), (0, 1)) - b_A = tl.load(p_A, boundary_check=(0, 1)) - else: - if USE_G: - g += bos * HV + i_h - p_g = tl.make_block_ptr(g, (T,), (HV,), (i_t * BT,), (BT,), (0,)) - b_g = tl.load(p_g, boundary_check=(0,)) - if USE_G_GAMMA: - b_gamma = tl.load(g_gamma + i_h) - b_g = b_gamma * (tl.arange(0, BT) + 1) - - b_A = tl.zeros([BT, BT], dtype=tl.float32) - for i_k in range(tl.cdiv(K, BK)): - p_k = tl.make_block_ptr(k, (T, K), (H*K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) - p_q = tl.make_block_ptr(q, (K, T), (1, H*K), (i_k * BK, i_t * BT), (BK, BT), (0, 1)) - - b_k = tl.load(p_k, boundary_check=(0, 1)) - b_q = tl.load(p_q, boundary_check=(0, 1)) - b_A += tl.dot(b_k, b_q) * scale - if USE_G or USE_G_GAMMA: - b_A *= exp2(b_g[None, :] - b_g[:, None]) - o_t = i_t * BT + tl.arange(0, BT) - m_t = o_t < T - m_A = (o_t[:, None] <= o_t[None, :]) & (m_t[:, None] & m_t) - b_A = tl.where(m_A, b_A, 0).to(do.dtype.element_ty) - - for i_v in range(tl.cdiv(V, BV)): - p_do = tl.make_block_ptr(do, (T, V), (HV*V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) - p_dv = tl.make_block_ptr(dv, (T, V), (HV*V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) - b_do = tl.load(p_do, boundary_check=(0, 1)) - b_dv = tl.dot(b_A.to(b_do.dtype), b_do) - tl.store(p_dv, b_dv.to(p_dv.dtype.element_ty), boundary_check=(0, 1)) - - -def chunk_fwd_o( - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - h: torch.Tensor, - g: torch.Tensor | None = None, - g_gamma: torch.Tensor | None = None, - scale: float | None = None, - state_v_first: bool = False, - cu_seqlens: torch.LongTensor | None = None, - chunk_size: int = 64, - chunk_indices: torch.LongTensor | None = None, -) -> torch.Tensor: - B, T, H, K, V, HV = *q.shape, v.shape[-1], v.shape[2] - BT = chunk_size - if chunk_indices is None and cu_seqlens is not None: - chunk_indices = prepare_chunk_indices(cu_seqlens, BT) - NT = triton.cdiv(T, BT) if cu_seqlens is None else len(chunk_indices) - if scale is None: - scale = k.shape[-1] ** -0.5 - - o = torch.empty_like(v) - def grid(meta): return (triton.cdiv(V, meta['BV']), NT, B * HV) - chunk_fwd_kernel_o[grid]( - q=q, - k=k, - v=v, - h=h, - g=g, - g_gamma=g_gamma, - o=o, - cu_seqlens=cu_seqlens, - chunk_indices=chunk_indices, - scale=scale, - T=T, - H=H, - HV=HV, - K=K, - V=V, - BT=BT, - STATE_V_FIRST=state_v_first, - ) - return o - - -def chunk_bwd_dv( - q: torch.Tensor, - k: torch.Tensor, - do: torch.Tensor, - dh: torch.Tensor, - g: torch.Tensor | None = None, - g_gamma: torch.Tensor | None = None, - scale: float | None = None, - state_v_first: bool = False, - cu_seqlens: torch.LongTensor | None = None, - chunk_size: int = 64, - chunk_indices: torch.LongTensor | None = None, -) -> torch.Tensor: - B, T, H, K, V, HV = *k.shape, do.shape[-1], do.shape[2] - BT = chunk_size - if chunk_indices is None and cu_seqlens is not None: - chunk_indices = prepare_chunk_indices(cu_seqlens, BT) - # H100 can have larger block size - if check_shared_mem('hopper', k.device.index): - CONST_TILING = 128 - elif check_shared_mem('ada', k.device.index): - CONST_TILING = 64 - else: - CONST_TILING = 32 - BK = min(max(triton.next_power_of_2(K), 16), CONST_TILING) - BV = min(max(triton.next_power_of_2(V), 16), CONST_TILING) - NT = triton.cdiv(T, BT) if cu_seqlens is None else len(chunk_indices) - NV = triton.cdiv(V, BV) - if scale is None: - scale = k.shape[-1] ** -0.5 - - dv = torch.empty_like(do) - grid = (NV, NT, B * HV) - chunk_bwd_kernel_dv[grid]( - q=q, - k=k, - g=g, - g_gamma=g_gamma, - do=do, - dv=dv, - dh=dh, - cu_seqlens=cu_seqlens, - chunk_indices=chunk_indices, - scale=scale, - T=T, - H=H, - HV=HV, - K=K, - V=V, - BT=BT, - BK=BK, - BV=BV, - STATE_V_FIRST=state_v_first, - ) - return dv - - -def chunk_bwd_dv_local( - q: torch.Tensor, - k: torch.Tensor, - do: torch.Tensor, - g: torch.Tensor | None = None, - g_gamma: torch.Tensor | None = None, - A: torch.Tensor | None = None, - scale: float = None, - cu_seqlens: torch.LongTensor | None = None, - chunk_size: int = 64, - chunk_indices: torch.LongTensor | None = None, -) -> torch.Tensor: - B, T, H, K, V, HV = *k.shape, do.shape[-1], do.shape[2] - BT = chunk_size - if chunk_indices is None and cu_seqlens is not None: - chunk_indices = prepare_chunk_indices(cu_seqlens, BT) - # H100 can have larger block size - if check_shared_mem('hopper', k.device.index): - CONST_TILING = 128 - elif check_shared_mem('ada', k.device.index): - CONST_TILING = 64 - else: - CONST_TILING = 32 - BK = min(max(triton.next_power_of_2(K), 16), CONST_TILING) - BV = min(max(triton.next_power_of_2(V), 16), CONST_TILING) - NT = triton.cdiv(T, BT) if cu_seqlens is None else len(chunk_indices) - - dv = torch.empty_like(do) - grid = (NT, B * HV) - chunk_bwd_kernel_dv_local[grid]( - q=q, - k=k, - g=g, - g_gamma=g_gamma, - A=A, - do=do, - dv=dv, - cu_seqlens=cu_seqlens, - chunk_indices=chunk_indices, - scale=scale, - T=T, - H=H, - HV=HV, - K=K, - V=V, - BT=BT, - BK=BK, - BV=BV, - ) - return dv - - -@dispatch('common') -def chunk_bwd_dqkwg( - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - do: torch.Tensor, - h: torch.Tensor, - dh: torch.Tensor, - w: torch.Tensor | None = None, - g: torch.Tensor | None = None, - g_gamma: torch.Tensor | None = None, - dv: torch.Tensor | None = None, - scale: float | None = None, - state_v_first: bool = False, - cu_seqlens: torch.LongTensor | None = None, - chunk_size: int = 64, - chunk_indices: torch.LongTensor | None = None, -) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: - if g is not None and IS_NVIDIA_HOPPER and TRITON_ABOVE_3_4_0 and not TRITON_ABOVE_3_7_1: - raise RuntimeError( - "Triton >= 3.4.0 and < 3.7.1 on Hopper GPUs produces incorrect results for " - "gated chunk_bwd_dqkwg (see #640). Please upgrade Triton to >= 3.7.1 or " - "install tilelang: `pip install tilelang`" - ) - - B, T, H, K, V, HV = *k.shape, v.shape[-1], v.shape[2] - BT = chunk_size - if chunk_indices is None and cu_seqlens is not None: - chunk_indices = prepare_chunk_indices(cu_seqlens, BT) - NT = triton.cdiv(T, BT) if cu_seqlens is None else len(chunk_indices) - - if check_shared_mem('hopper', k.device.index): - CONST_TILING = 128 - elif check_shared_mem('ada', k.device.index): - CONST_TILING = 64 - else: - CONST_TILING = 32 - BK = min(max(triton.next_power_of_2(K), 16), CONST_TILING) - BV = min(max(triton.next_power_of_2(V), 16), CONST_TILING) - NK = triton.cdiv(K, BK) - dq = q.new_empty(B, T, HV, K) - dk = k.new_empty(B, T, HV, K) - dg = torch.empty(NK, *g.shape, dtype=torch.float32, device=g.device) if g is not None else None - dw = torch.empty_like(w) if w is not None else None - - grid = (NK, NT, B * HV) - chunk_bwd_kernel_dqkwg[grid]( - q=q, - k=k, - v=v, - g=g, - g_gamma=g_gamma, - h=h, - do=do, - dh=dh, - dw=dw, - dq=dq, - dk=dk, - dv=dv, - dg=dg, - cu_seqlens=cu_seqlens, - chunk_indices=chunk_indices, - scale=scale, - B=B, - T=T, - H=H, - HV=HV, - K=K, - V=V, - BT=BT, - BK=BK, - BV=BV, - STATE_V_FIRST=state_v_first, - ) - - if H != HV: - dq = dq.view(B, T, H, HV // H, K).sum(3) - dk = dk.view(B, T, H, HV // H, K).sum(3) - if dg is not None: - dg = dg.sum(0) - return dq, dk, dw, dg diff --git a/build/torch-xpu/ops/common/chunk_scaled_dot_kkt.py b/build/torch-xpu/ops/common/chunk_scaled_dot_kkt.py deleted file mode 100644 index 41dee5853e23f989ba83049f3cb4e51cb51f1b76..0000000000000000000000000000000000000000 --- a/build/torch-xpu/ops/common/chunk_scaled_dot_kkt.py +++ /dev/null @@ -1,132 +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 -import triton.language as tl - -from ...ops.utils import prepare_chunk_indices -from ...ops.utils.op import exp2 -from ...utils import autotune_cache_kwargs - - -@triton.heuristics({ - 'USE_G': lambda args: args['g'] is not None, - 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None, -}) -@triton.autotune( - configs=[ - triton.Config({'BK': BK}, num_warps=num_warps, num_stages=num_stages) - for BK in [32, 64, 128] - for num_warps in [2, 4, 8] - for num_stages in [2, 3, 4] - ], - key=['H', 'HV', 'K', 'BT', 'IS_VARLEN'], - **autotune_cache_kwargs, -) -@triton.jit(do_not_specialize=['T']) -def chunk_scaled_dot_kkt_fwd_kernel( - k, - g, - beta, - A, - cu_seqlens, - chunk_indices, - T, - H: tl.constexpr, - HV: tl.constexpr, - K: tl.constexpr, - BT: tl.constexpr, - BK: tl.constexpr, - IS_VARLEN: tl.constexpr, - USE_G: tl.constexpr, -): - i_t, i_bh = tl.program_id(0), tl.program_id(1).to(tl.int64) - i_b, i_h = i_bh // HV, i_bh % HV - 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.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) - T = eos - bos - else: - bos, eos = i_b * T, i_b * T + T - o_t = i_t * BT + tl.arange(0, BT) - m_t = o_t < T - - p_b = tl.make_block_ptr(beta + bos*HV + i_h, (T,), (HV,), (i_t * BT,), (BT,), (0,)) - b_b = tl.load(p_b, boundary_check=(0,)) - - b_A = tl.zeros([BT, BT], dtype=tl.float32) - for i_k in range(tl.cdiv(K, BK)): - p_k = tl.make_block_ptr(k + (bos*H + i_h // (HV // H)) * K, (T, K), (H*K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) - b_k = tl.load(p_k, boundary_check=(0, 1)) - b_A += tl.dot(b_k, tl.trans(b_k)) - - if USE_G: - p_g = tl.make_block_ptr(g + bos*HV + i_h, (T,), (HV,), (i_t * BT,), (BT,), (0,)) - b_g = tl.load(p_g, boundary_check=(0,)) - b_g_diff = b_g[:, None] - b_g[None, :] - b_A *= exp2(b_g_diff) - b_A *= b_b[:, None] - - m_A = (o_t[:, None] > o_t[None, :]) & (m_t[:, None] & m_t) - b_A = tl.where(m_A, b_A, 0) - p_A = tl.make_block_ptr(A + (bos*HV + i_h) * BT, (T, BT), (BT*HV, 1), (i_t * BT, 0), (BT, BT), (1, 0)) - tl.store(p_A, b_A.to(p_A.dtype.element_ty), boundary_check=(0, 1)) - - -def chunk_scaled_dot_kkt_fwd( - k: torch.Tensor, - g: torch.Tensor | None = None, - beta: torch.Tensor | None = None, - cu_seqlens: torch.LongTensor | None = None, - chunk_size: int = 64, - output_dtype: torch.dtype = torch.float32, - chunk_indices: torch.LongTensor | None = None, -) -> torch.Tensor: - r""" - Compute beta * K * K^T. - - Args: - k (torch.Tensor): - The key tensor of shape `[B, T, H, K]` where `H` is the number of query/key heads. - g (torch.Tensor): - The cumulative sum of the gate tensor of shape `[B, T, HV]`. Default: `None`. - beta (torch.Tensor): - The beta tensor of shape `[B, T, HV]` where `HV` is the number of value/output heads. - cu_seqlens (torch.LongTensor): - The cumulative sequence lengths of the input tensor. - Default: None - chunk_size (int): - The chunk size. Default: 64. - output_dtype (torch.dtype): - The dtype of the output tensor. Default: `torch.float32` - chunk_indices (torch.LongTensor): - The chunk indices of the input tensor. Default: None. - Returns: - beta * K * K^T of shape `[B, T, HV, BT]` where `BT` is the chunk size. - For GVA, H < HV and HV % H == 0. For standard attention, H == HV. - """ - B, T, H, K, HV = *k.shape, beta.shape[2] - BT = chunk_size - if chunk_indices is None and cu_seqlens is not None: - chunk_indices = prepare_chunk_indices(cu_seqlens, BT) - NT = triton.cdiv(T, BT) if cu_seqlens is None else len(chunk_indices) - A = torch.empty(B, T, HV, BT, device=k.device, dtype=output_dtype) - chunk_scaled_dot_kkt_fwd_kernel[(NT, B * HV)]( - k=k, - g=g, - beta=beta, - A=A, - cu_seqlens=cu_seqlens, - chunk_indices=chunk_indices, - T=T, - H=H, - HV=HV, - K=K, - BT=BT, - ) - return A diff --git a/build/torch-xpu/ops/common/fused_chunk.py b/build/torch-xpu/ops/common/fused_chunk.py deleted file mode 100644 index fbb21b11cd5d3e12521c0dd7e0da0ac9f78039af..0000000000000000000000000000000000000000 --- a/build/torch-xpu/ops/common/fused_chunk.py +++ /dev/null @@ -1,650 +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 -import triton.language as tl - -from ...ops.utils import chunk_local_cumsum -from ...ops.utils.constant import RCP_LN2 -from ...ops.utils.op import exp2 -from ...utils import ( - IS_NVIDIA_HOPPER, - autocast_custom_bwd, - autocast_custom_fwd, - autotune_cache_kwargs, - check_shared_mem, - input_guard, -) - -BKV_LIST = [64, 128] if check_shared_mem() else [32, 64] -NUM_WARPS = [2, 4] if IS_NVIDIA_HOPPER else [2, 4, 8] - - -@triton.heuristics({ - 'USE_G': lambda args: args['g'] is not None, - 'USE_G_GAMMA': lambda args: args['g_gamma'] is not None, - 'USE_INITIAL_STATE': lambda args: args['h0'] is not None, - 'STORE_FINAL_STATE': lambda args: args['ht'] is not None, - 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None, -}) -@triton.autotune( - configs=[ - triton.Config({'BV': BV}, num_warps=num_warps, num_stages=num_stages) - for BV in BKV_LIST - for num_warps in NUM_WARPS - for num_stages in [2, 3, 4] - ], - key=['H', 'K', 'V', 'BT'], - **autotune_cache_kwargs, -) -@triton.jit(do_not_specialize=['T']) -def fused_chunk_fwd_kernel( - q, - k, - v, - g, - g_gamma, - o, - h0, - ht, - cu_seqlens, - scale, - T, - B: tl.constexpr, - H: tl.constexpr, - K: tl.constexpr, - V: tl.constexpr, - BT: tl.constexpr, - BK: tl.constexpr, - BV: tl.constexpr, - USE_G: tl.constexpr, - USE_G_GAMMA: tl.constexpr, - USE_INITIAL_STATE: tl.constexpr, - STORE_FINAL_STATE: tl.constexpr, - IS_VARLEN: tl.constexpr, -): - i_v, i_k, i_nh = tl.program_id(0), tl.program_id(1), tl.program_id(2) - i_n, i_h = i_nh // H, i_nh % H - - all = B * T - if IS_VARLEN: - bos, eos = tl.load(cu_seqlens + i_n).to(tl.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) - T = eos - bos - else: - bos, eos = i_n * T, i_n * T + T - NT = tl.cdiv(T, BT) - - o_i = tl.arange(0, BT) - - if USE_G_GAMMA: - # decay rate given the head index - b_gamma = tl.load(g_gamma + i_h) - b_g = b_gamma * (o_i + 1) - b_g_last = b_gamma * BT - b_gq = exp2(b_g) - b_gk = exp2(b_g_last - b_g) - b_gn = exp2(b_g_last) - - # [BT, BT] - m_s = o_i[:, None] >= o_i[None, :] - - q = q + (bos*H + i_h) * K - k = k + (bos*H + i_h) * K - v = v + (bos*H + i_h) * V - o = o + (i_k * all + bos).to(tl.int64) * H*V + i_h * V - - # [BK, BV] - b_h = tl.zeros([BK, BV], dtype=tl.float32) - if USE_INITIAL_STATE: - p_h = tl.make_block_ptr(h0 + i_nh * K*V, (K, V), (V, 1), (i_k * BK, i_v * BV), (BK, BV), (1, 0)) - b_h = tl.load(p_h, boundary_check=(0, 1)).to(tl.float32) - - for i_t in range(0, NT): - p_q = tl.make_block_ptr(q, (T, K), (H*K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) - p_k = tl.make_block_ptr(k, (K, T), (1, H*K), (i_k * BK, i_t * BT), (BK, BT), (0, 1)) - p_v = tl.make_block_ptr(v, (T, V), (H*V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) - p_o = tl.make_block_ptr(o, (T, V), (H*V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) - - o_t = i_t * BT + tl.arange(0, BT) - m_t = o_t < T - # [BT, BK] - b_q = tl.load(p_q, boundary_check=(0, 1)) - b_q = (b_q * scale).to(b_q.dtype) - # [BK, BT] - b_k = tl.load(p_k, boundary_check=(0, 1)) - # [BT, BV] - b_v = tl.load(p_v, boundary_check=(0, 1)) - last_idx = min(i_t * BT + BT, T) - 1 - - # [BT, BT] - b_s = tl.dot(b_q, b_k) - - # scalar decay - if USE_G: - p_g = g + (bos + o_t) * H + i_h - b_g = tl.load(p_g, mask=(o_t < T), other=0.) - b_g_last = tl.load(g + (bos + last_idx) * H + i_h) - - b_gq = exp2(b_g) - b_gk = exp2(b_g_last - b_g) - b_gn = exp2(b_g_last) - if USE_G_GAMMA: - b_g_last = b_gamma * min(BT, T - i_t * BT) - b_gk = exp2(b_g_last - b_g) - b_gn = exp2(b_g_last) - if USE_G or USE_G_GAMMA: - b_gs = tl.where(m_s & m_t, exp2(b_g[:, None] - b_g[None, :]), 0) - # [BT, BT] - b_s *= b_gs - # [BT, BV] - b_o = tl.dot(b_s.to(b_q.dtype), b_v) + tl.dot(b_q, b_h.to(b_q.dtype)) * b_gq[:, None] - b_v = (b_v * b_gk[:, None]).to(b_v.dtype) - b_h *= b_gn - else: - # [BT, BT] - b_s *= m_s & m_t - # [BT, BV] - b_o = tl.dot(b_s.to(b_q.dtype), b_v) + tl.dot(b_q, b_h.to(b_q.dtype)) - - b_h += tl.dot(b_k, b_v) - - tl.store(p_o, b_o.to(p_o.dtype.element_ty), boundary_check=(0, 1)) - - if STORE_FINAL_STATE: - p_ht = tl.make_block_ptr(ht + i_nh * K*V, (K, V), (V, 1), (i_k * BK, i_v * BV), (BK, BV), (1, 0)) - tl.store(p_ht, b_h.to(p_ht.dtype.element_ty), boundary_check=(0, 1)) - - -@triton.heuristics({ - 'USE_G': lambda args: args['g'] is not None, - 'USE_G_GAMMA': lambda args: args['g_gamma'] is not None, - 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None, - 'USE_INITIAL_STATE': lambda args: args['dh0'] is not None, - 'USE_FINAL_STATE': lambda args: args['dht'] is not None, -}) -@triton.autotune( - configs=[ - triton.Config({}, num_warps=num_warps, num_stages=num_stages) - for num_warps in NUM_WARPS - for num_stages in [2, 3, 4] - ], - key=['H', 'K', 'V', 'BT'], - **autotune_cache_kwargs, -) -@triton.jit(do_not_specialize=['T']) -def fused_chunk_bwd_kernel( - q, - k, - v, - g, - g_gamma, - do, - dq, - dk, - dv, - dg, - h0, - dht, - dh0, - cu_seqlens, - scale, - T, - B: tl.constexpr, - H: tl.constexpr, - K: tl.constexpr, - V: tl.constexpr, - BT: tl.constexpr, - BK: tl.constexpr, - BV: tl.constexpr, - USE_G: tl.constexpr, - USE_G_GAMMA: tl.constexpr, - IS_VARLEN: tl.constexpr, - USE_INITIAL_STATE: tl.constexpr, - USE_FINAL_STATE: tl.constexpr, -): - i_v, i_k, i_nh = tl.program_id(0), tl.program_id(1), tl.program_id(2) - i_n, i_h = i_nh // H, i_nh % H - - all = B * T - if IS_VARLEN: - bos, eos = tl.load(cu_seqlens + i_n).to(tl.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) - T = eos - bos - else: - bos, eos = i_n * T, i_n * T + T - NT = tl.cdiv(T, BT) - NV = tl.cdiv(V, BV) - - o_i = tl.arange(0, BT) - if USE_G_GAMMA: - b_gamma = tl.load(g_gamma + i_h) - b_g = b_gamma * (o_i + 1) - b_g_last = b_gamma * BT - b_gq = exp2(b_g) - b_gk = exp2(b_g_last - b_g) - b_gn = exp2(b_g_last) - - m_s = o_i[:, None] >= o_i[None, :] - - q = q + (bos*H + i_h) * K - k = k + (bos*H + i_h) * K - v = v + (bos*H + i_h) * V - do = do + (bos*H + i_h) * V - dq = dq + (i_v * all + bos).to(tl.int64) * H*K + i_h * K - dk = dk + (i_v * all + bos).to(tl.int64) * H*K + i_h * K - dv = dv + (i_k * all + bos).to(tl.int64) * H*V + i_h * V - - # [BV, BK] - b_h = tl.zeros([BV, BK], dtype=tl.float32) - if USE_INITIAL_STATE: - p_h = tl.make_block_ptr(h0 + i_nh * K*V, (V, K), (1, V), (i_v * BV, i_k * BK), (BV, BK), (0, 1)) - b_h = tl.load(p_h, boundary_check=(0, 1)).to(tl.float32) - - for i_t in range(0, NT): - p_q = tl.make_block_ptr(q, (T, K), (H*K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) - p_k = tl.make_block_ptr(k, (T, K), (H*K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) - p_v = tl.make_block_ptr(v, (V, T), (1, H*V), (i_v * BV, i_t * BT), (BV, BT), (0, 1)) - p_do = tl.make_block_ptr(do, (T, V), (H*V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) - p_dq = tl.make_block_ptr(dq, (T, K), (H*K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) - - o_t = i_t * BT + tl.arange(0, BT) - m_t = o_t < T - # [BT, BK] - b_k = tl.load(p_k, boundary_check=(0, 1)) - # [BV, BT] - b_v = tl.load(p_v, boundary_check=(0, 1)) - # [BT, BV] - b_do = tl.load(p_do, boundary_check=(0, 1)) - last_idx = min(i_t * BT + BT, T) - 1 - - # [BT, BT] - b_ds = tl.dot(b_do, b_v) * scale - - # scalar decay - if USE_G: - p_g = g + (bos + o_t) * H + i_h - b_g = tl.load(p_g, mask=(o_t < T), other=0.) - b_g_last = tl.load(g + (bos + last_idx) * H + i_h) - - b_gq = exp2(b_g) - b_gk = exp2(b_g_last - b_g) - b_gn = exp2(b_g_last) - - p_dg = dg + ((i_k * NV + i_v) * all + (bos + o_t)).to(tl.int64) * H + i_h - # [BT, BT] - b_gs = tl.where(m_s & m_t, exp2(b_g[:, None] - b_g[None, :]), 0) - b_ds = b_ds * b_gs - # [BT, BK] - b_q = tl.load(p_q, boundary_check=(0, 1)) - b_dq = tl.dot(b_ds.to(b_k.dtype), b_k) + tl.dot((b_do * b_gq[:, None] * scale).to(b_k.dtype), b_h.to(b_k.dtype)) - # [BT] - b_dg_t = tl.sum(b_q * b_dq, 1) - tl.store(p_dg, b_dg_t.to(p_dg.dtype.element_ty), mask=m_t) - # [BV, BK] - b_h = b_h * b_gn + tl.dot(b_v, (b_k * b_gk[:, None]).to(b_k.dtype)) - - elif USE_G_GAMMA: - b_g_last = b_gamma * min(BT, T - i_t * BT) - b_gk = exp2(b_g_last - b_g) - b_gn = exp2(b_g_last) - - # [BT, BT] - b_gs = tl.where(m_s & m_t, exp2(b_g[:, None] - b_g[None, :]), 0) - b_ds = b_ds * b_gs - # [BT, BK] - b_q = tl.load(p_q, boundary_check=(0, 1)) - b_dq = tl.dot(b_ds.to(b_k.dtype), b_k) + tl.dot((b_do * b_gq[:, None] * scale).to(b_k.dtype), b_h.to(b_k.dtype)) - # [BV, BK] - b_h = b_h * b_gn + tl.dot(b_v, (b_k * b_gk[:, None]).to(b_k.dtype)) - - else: - # [BT, BT] - b_ds *= m_s & m_t - # [BT, BK] - b_dq = tl.dot(b_ds.to(b_k.dtype), b_k) + tl.dot((b_do * scale).to(b_k.dtype), b_h.to(b_k.dtype)) - # [BV, BK] - b_h += tl.dot(b_v, b_k) - - tl.store(p_dq, b_dq.to(p_dq.dtype.element_ty), boundary_check=(0, 1)) - - # [BK, BV] - b_dh = tl.zeros([BK, BV], dtype=tl.float32) - if USE_FINAL_STATE: - p_dh = tl.make_block_ptr(dht + i_nh * K*V, (K, V), (V, 1), (i_k * BK, i_v * BV), (BK, BV), (1, 0)) - b_dh += tl.load(p_dh, boundary_check=(0, 1)).to(tl.float32) - - if USE_G: - b_dg = tl.zeros([BT], dtype=tl.float32) - b_dg_last = tl.sum(tl.trans(b_h) * b_dh) - - # sync threads - b_h = None - tl.debug_barrier() - - for i_t in range(NT - 1, -1, -1): - p_q = tl.make_block_ptr(q, (K, T), (1, H*K), (i_k * BK, i_t * BT), (BK, BT), (0, 1)) - p_k = tl.make_block_ptr(k, (T, K), (H*K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) - p_v = tl.make_block_ptr(v, (T, V), (H*V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) - p_do = tl.make_block_ptr(do, (T, V), (H*V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) - p_dk = tl.make_block_ptr(dk, (T, K), (H*K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) - p_dv = tl.make_block_ptr(dv, (T, V), (H*V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) - # [BK, BT] - b_q = tl.load(p_q, boundary_check=(0, 1)) - # [BT, BK] - b_k = tl.load(p_k, boundary_check=(0, 1)) - # [BT, BV] - b_v = tl.load(p_v, boundary_check=(0, 1)) - b_do = tl.load(p_do, boundary_check=(0, 1)) - last_idx = min(i_t * BT + BT, T) - 1 - - o_t = i_t * BT + tl.arange(0, BT) - m_t = o_t < T - # [BT, BT] - b_s = tl.dot(b_k, b_q) - b_ds = tl.dot(b_v, tl.trans(b_do)) - - if USE_G: - p_g = g + (bos + o_t) * H + i_h - p_dg = dg + ((i_k * NV + i_v) * all + (bos + o_t)).to(tl.int64) * H + i_h - b_g = tl.load(p_g, mask=m_t, other=0.) - b_g_last = tl.load(g + (bos + last_idx) * H + i_h) - - b_gq = exp2(b_g) - b_gk = exp2(b_g_last - b_g) - b_gn = exp2(b_g_last) - b_gs = tl.trans(tl.where(m_s & (m_t[:, None] & m_t), exp2(b_g[:, None] - b_g[None, :]), 0)) * scale - - b_s = b_s * b_gs - b_ds = b_ds * b_gs - - # [BT, BK] - b_dk = tl.dot(b_ds.to(b_k.dtype), tl.trans(b_q)) + tl.dot(b_v, tl.trans(b_dh).to(b_v.dtype)) * b_gk[:, None] - - # [BT] - b_dg_t = tl.where(m_t, tl.load(p_dg, mask=m_t, other=0.) - tl.sum(b_k * b_dk, 1), 0) - b_dg_last += tl.sum(b_dg_t, 0) - b_dg = b_dg_last + b_dg_t - tl.cumsum(b_dg_t, 0) - - # [BT, BV] - b_dv = tl.dot(b_s.to(b_do.dtype), b_do) + tl.dot(b_k, b_dh.to(b_k.dtype)) * b_gk[:, None] - # [BK, BV] - b_dh = b_dh * b_gn + tl.dot(b_q, (b_do * b_gq[:, None] * scale).to(b_do.dtype)) - - tl.store(p_dg, b_dg.to(p_dg.dtype.element_ty), mask=m_t) - - elif USE_G_GAMMA: - b_g_last = b_gamma * min(BT, T - i_t * BT) - b_gk = exp2(b_g_last - b_g) - b_gn = exp2(b_g_last) - b_gs = tl.trans(tl.where(m_s & (m_t[:, None] & m_t), exp2(b_g[:, None] - b_g[None, :]), 0)) * scale - - b_s = b_s * b_gs - b_ds = b_ds * b_gs - - b_dk = tl.dot(b_ds.to(b_k.dtype), tl.trans(b_q)) + tl.dot(b_v, tl.trans(b_dh).to(b_v.dtype)) * b_gk[:, None] - # [BT, BV] - b_dv = tl.dot(b_s.to(b_do.dtype), b_do) + tl.dot(b_k, b_dh.to(b_k.dtype)) * b_gk[:, None] - # [BK, BV] - b_dh = b_dh * b_gn + tl.dot(b_q, (b_do * b_gq[:, None] * scale).to(b_do.dtype)) - - else: - mask = tl.trans(m_s & (m_t[:, None] & m_t)) - b_s = tl.where(mask, b_s * scale, 0).to(b_do.dtype) - b_ds = tl.where(mask, b_ds * scale, 0).to(b_q.dtype) - - b_dk = tl.dot(b_ds, tl.trans(b_q)) + tl.dot(b_v, tl.trans(b_dh).to(b_v.dtype)) - # [BT, BV] - b_dv = tl.dot(b_s.to(b_do.dtype), b_do) + tl.dot(b_k, b_dh.to(b_k.dtype)) - # [BK, BV] - b_dh += tl.dot(b_q, (b_do * scale).to(b_do.dtype)) - - tl.store(p_dk, b_dk.to(p_dk.dtype.element_ty), boundary_check=(0, 1)) - tl.store(p_dv, b_dv.to(p_dv.dtype.element_ty), boundary_check=(0, 1)) - - if USE_INITIAL_STATE: - p_dh0 = tl.make_block_ptr(dh0 + i_nh * K*V, (K, V), (V, 1), (i_k * BK, i_v * BV), (BK, BV), (1, 0)) - tl.store(p_dh0, b_dh.to(p_dh0.dtype.element_ty), boundary_check=(0, 1)) - - -def fused_chunk_fwd( - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - g: torch.Tensor | None = None, - g_gamma: torch.Tensor | None = None, - scale: float | None = None, - initial_state: torch.Tensor | None = None, - output_final_state: bool = False, - cu_seqlens: torch.LongTensor | None = None, - chunk_size: int = 64, -): - B, T, H, K, V = *q.shape, v.shape[-1] - BT = chunk_size - BK = min(max(triton.next_power_of_2(K), 16), 64) - N = B if cu_seqlens is None else len(cu_seqlens) - 1 - NK = triton.cdiv(K, BK) - - o = v.new_empty(NK, *v.shape, dtype=torch.float) if NK > 1 else torch.empty_like(v) - ht = k.new_empty(N, H, K, V, dtype=torch.float) if output_final_state else None - def grid(meta): return (triton.cdiv(V, meta['BV']), NK, N * H) - fused_chunk_fwd_kernel[grid]( - q=q, - k=k, - v=v, - g=g, - g_gamma=g_gamma, - o=o, - h0=initial_state, - ht=ht, - cu_seqlens=cu_seqlens, - scale=scale, - B=B, - T=T, - H=H, - K=K, - V=V, - BT=BT, - BK=BK, - ) - if NK > 1: - o = o.sum(0).to(v) - return o, ht - - -def fused_chunk_bwd( - q, - k, - v, - g, - g_gamma, - do, - scale, - initial_state: torch.Tensor, - dht: torch.Tensor, - cu_seqlens: torch.LongTensor | None = None, - chunk_size: int = 64, -): - B, T, H, K, V = *q.shape, v.shape[-1] - N = B if cu_seqlens is None else len(cu_seqlens) - 1 - BT = chunk_size - BK = min(max(triton.next_power_of_2(K), 16), 64) - BV = min(max(triton.next_power_of_2(V), 16), 64) - NK, NV = triton.cdiv(K, BK), triton.cdiv(V, BV) - - dq = q.new_empty(NV, *q.shape, dtype=torch.float) if NV > 1 else torch.empty_like(q) - dk = k.new_empty(NV, *k.shape, dtype=torch.float) if NV > 1 else torch.empty_like(k) - dv = v.new_empty(NK, *v.shape, dtype=torch.float) if NK > 1 else torch.empty_like(v) - dg = g.new_empty(NK*NV, *g.shape, dtype=torch.float) if g is not None else None - dh0 = torch.empty_like(initial_state) if initial_state is not None else None - - grid = (NV, NK, N * H) - fused_chunk_bwd_kernel[grid]( - q=q, - k=k, - v=v, - g=g, - g_gamma=g_gamma, - do=do, - dq=dq, - dk=dk, - dv=dv, - dg=dg, - h0=initial_state, - dht=dht, - dh0=dh0, - cu_seqlens=cu_seqlens, - scale=scale, - T=T, - B=B, - H=H, - K=K, - V=V, - BT=BT, - BK=BK, - BV=BV, - ) - dq = dq.sum(0) if NV > 1 else dq - dk = dk.sum(0) if NV > 1 else dk - dv = dv.sum(0) if NK > 1 else dv - if dg is not None: - dg = dg.sum(0).to(g) - - return dq, dk, dv, dg, dh0 - - -class FusedChunkFunction(torch.autograd.Function): - - @staticmethod - @input_guard - @autocast_custom_fwd - def forward( - ctx, - q, - k, - v, - g, - g_gamma, - scale, - initial_state, - output_final_state, - cu_seqlens, - ): - chunk_size = min(64, max(16, triton.next_power_of_2(q.shape[1]))) - # Pre-scale by RCP_LN2 so downstream kernels can use exp2 directly. - if g is not None: - g = chunk_local_cumsum( - g, - chunk_size=chunk_size, - scale=RCP_LN2, - cu_seqlens=cu_seqlens, - ) - if g_gamma is not None: - g_gamma = g_gamma * RCP_LN2 - o, ht = fused_chunk_fwd( - q=q, - k=k, - v=v, - g=g, - g_gamma=g_gamma, - scale=scale, - initial_state=initial_state, - output_final_state=output_final_state, - cu_seqlens=cu_seqlens, - chunk_size=chunk_size, - ) - - ctx.save_for_backward(q, k, v, g, g_gamma, initial_state) - ctx.chunk_size = chunk_size - ctx.scale = scale - ctx.cu_seqlens = cu_seqlens - return o.to(q.dtype), ht - - @staticmethod - @input_guard - @autocast_custom_bwd - def backward(ctx, do, dht=None): - q, k, v, g, g_gamma, initial_state = ctx.saved_tensors - - dq, dk, dv, dg, dh0 = fused_chunk_bwd( - q=q, - k=k, - v=v, - g=g, - g_gamma=g_gamma, - do=do, - scale=ctx.scale, - initial_state=initial_state, - dht=dht, - cu_seqlens=ctx.cu_seqlens, - chunk_size=ctx.chunk_size, - ) - if g is not None: - dg = dg.to(g) - return dq.to(q), dk.to(k), dv.to(v), dg, None, None, dh0, None, None - - -def fused_chunk( - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - g: torch.Tensor | None = None, - g_gamma: torch.Tensor | None = None, - scale: float | None = None, - initial_state: torch.Tensor | None = None, - output_final_state: bool = False, - cu_seqlens: torch.LongTensor | None = None, -) -> tuple[torch.Tensor, torch.Tensor]: - r""" - Args: - q (torch.Tensor): - queries of shape `[B, T, H, K]`. - k (torch.Tensor): - keys of shape `[B, T, H, K]`. - v (torch.Tensor): - values of shape `[B, T, H, V]`. - g (torch.Tensor): - Forget gates of shape `[B, T, H]`. - Compared to GLA, the gating is head-wise instead of elementwise. - g_gamma (torch.Tensor): - Log decay of shape `[H]`. - Head-wise data-independent decay is used if `g_gamma` is provided. - Only one of `g` or `g_gamma` should be provided. - scale (Optional[int]): - Scale factor for the attention scores. - If not provided, it will default to `1 / sqrt(K)`. Default: `None`. - initial_state (Optional[torch.Tensor]): - Initial state of shape `[N, H, K, V]` for `N` input sequences. - For equal-length input sequences, `N` equals the batch size `B`. - Default: `None`. - output_final_state (Optional[bool]): - Whether to output the final state of shape `[N, H, K, V]`. Default: `False`. - cu_seqlens (torch.LongTensor): - Cumulative sequence lengths of shape `[N+1]` used for variable-length training, - consistent with the FlashAttention API. - - Returns: - o (torch.Tensor): - Outputs of shape `[B, T, H, V]`. - final_state (torch.Tensor): - Final state of shape `[N, H, K, V]` if `output_final_state=True` else `None`. - """ - if g is not None and g_gamma is not None: - raise ValueError("Only one of `g` or `g_gamma` should be provided.") - if scale is None: - scale = k.shape[-1] ** -0.5 - o, final_state = FusedChunkFunction.apply( - q, - k, - v, - g, - g_gamma, - scale, - initial_state, - output_final_state, - cu_seqlens, - ) - return o, final_state diff --git a/build/torch-xpu/ops/common/fused_recurrent.py b/build/torch-xpu/ops/common/fused_recurrent.py deleted file mode 100644 index 396295f62bf701324832b6cd9e374f96bea43a94..0000000000000000000000000000000000000000 --- a/build/torch-xpu/ops/common/fused_recurrent.py +++ /dev/null @@ -1,652 +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 -import triton.language as tl - -from ...ops.utils.op import exp -from ...utils import autocast_custom_bwd, autocast_custom_fwd, autotune_cache_kwargs, input_guard - - -@triton.heuristics({ - 'USE_INITIAL_STATE': lambda args: args['h0'] is not None, - 'STORE_FINAL_STATE': lambda args: args['ht'] is not None, - 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None, -}) -@triton.autotune( - configs=[ - triton.Config({}, num_warps=num_warps) - for num_warps in [4, 8] - ], - key=['BK', 'BV', 'USE_G', 'USE_G_GAMMA', 'USE_GK', 'USE_GV', 'STATE_V_FIRST'], - **autotune_cache_kwargs, -) -@triton.jit(do_not_specialize=['B', 'T']) -def fused_recurrent_fwd_kernel( - q, - k, - v, - g, - g_gamma, - gk, - gv, - o, - h0, - ht, - cu_seqlens, - scale, - B, - T, - H: tl.constexpr, - K: tl.constexpr, - V: tl.constexpr, - BK: tl.constexpr, - BV: tl.constexpr, - REVERSE: tl.constexpr, - USE_G: tl.constexpr, - USE_G_GAMMA: tl.constexpr, - USE_GK: tl.constexpr, - USE_GV: tl.constexpr, - USE_INITIAL_STATE: tl.constexpr, - STORE_FINAL_STATE: tl.constexpr, - IS_VARLEN: tl.constexpr, - STATE_V_FIRST: tl.constexpr = False, -): - i_v, i_k, i_nh = tl.program_id(0).to(tl.int64), tl.program_id(1).to(tl.int64), tl.program_id(2).to(tl.int64) - i_n, i_h = i_nh // H, i_nh % H - - all = B * T - if IS_VARLEN: - bos, eos = tl.load(cu_seqlens + i_n).to(tl.int64), tl.load(cu_seqlens + i_n + 1).to(tl.int64) - T = eos - bos - else: - bos, eos = i_n * T, i_n * T + T - - o_k = i_k * BK + tl.arange(0, BK) - o_v = i_v * BV + tl.arange(0, BV) - p_q = q + (bos + ((T-1) if REVERSE else 0)) * H*K + i_h * K + o_k - p_k = k + (bos + ((T-1) if REVERSE else 0)) * H*K + i_h * K + o_k - p_v = v + (bos + ((T-1) if REVERSE else 0)) * H*V + i_h * V + o_v - p_o = o + ((i_k * all + bos) + ((T-1) if REVERSE else 0)) * H*V + i_h * V + o_v - if USE_G: - p_g = g + (bos + ((T-1) if REVERSE else 0)) * H + i_h - if USE_GK: - p_gk = gk + (bos + ((T-1) if REVERSE else 0)) * H*K + i_h * K + o_k - if USE_GV: - p_gv = gv + (bos + ((T-1) if REVERSE else 0)) * H*V + i_h * V + o_v - if USE_G_GAMMA: - b_g_gamma = tl.load(g_gamma + i_h) - - m_k = o_k < K - m_v = o_v < V - if STATE_V_FIRST: - m_h = m_v[:, None] & m_k[None, :] - b_h = tl.zeros([BV, BK], dtype=tl.float32) - else: - m_h = m_k[:, None] & m_v[None, :] - b_h = tl.zeros([BK, BV], dtype=tl.float32) - - if USE_INITIAL_STATE: - if STATE_V_FIRST: - p_h0 = h0 + i_nh * K*V + o_v[:, None] * K + o_k[None, :] - else: - p_h0 = h0 + i_nh * K*V + o_k[:, None] * V + o_v[None, :] - b_h += tl.load(p_h0, mask=m_h, other=0).to(tl.float32) - - for _ in range(0, T): - b_q = tl.load(p_q, mask=m_k, other=0).to(tl.float32) * scale - b_k = tl.load(p_k, mask=m_k, other=0).to(tl.float32) - b_v = tl.load(p_v, mask=m_v, other=0).to(tl.float32) - if USE_G: - b_g = tl.load(p_g).to(tl.float32) - b_h = b_h * exp(b_g) - if USE_G_GAMMA: - b_h = b_h * exp(b_g_gamma) - if USE_GK: - b_gk = tl.load(p_gk, mask=m_k, other=0).to(tl.float32) - if STATE_V_FIRST: - b_h = b_h * exp(b_gk[None, :]) - else: - b_h = b_h * exp(b_gk[:, None]) - if USE_GV: - b_gv = tl.load(p_gv, mask=m_v, other=0).to(tl.float32) - if STATE_V_FIRST: - b_h = b_h * exp(b_gv[:, None]) - else: - b_h = b_h * exp(b_gv[None, :]) - if STATE_V_FIRST: - b_h += b_v[:, None] * b_k[None, :] - b_o = tl.sum(b_h * b_q[None, :], axis=1) - else: - b_h += b_k[:, None] * b_v[None, :] - b_o = tl.sum(b_h * b_q[:, None], axis=0) - tl.store(p_o, b_o.to(p_o.dtype.element_ty), mask=m_v) - p_q += (-1 if REVERSE else 1) * H*K - p_k += (-1 if REVERSE else 1) * H*K - p_v += (-1 if REVERSE else 1) * H*V - p_o += (-1 if REVERSE else 1) * H*V - if USE_G: - p_g += (-1 if REVERSE else 1) * H - if USE_GK: - p_gk += (-1 if REVERSE else 1) * H*K - if USE_GV: - p_gv += (-1 if REVERSE else 1) * H*V - - if STORE_FINAL_STATE: - if STATE_V_FIRST: - p_ht = ht + i_nh * K*V + o_v[:, None] * K + o_k[None, :] - else: - p_ht = ht + i_nh * K*V + o_k[:, None] * V + o_v[None, :] - tl.store(p_ht, b_h.to(p_ht.dtype.element_ty), mask=m_h) - - -@triton.heuristics({ - 'USE_INITIAL_STATE': lambda args: args['h0'] is not None, - 'STORE_INITIAL_STATE_GRADIENT': lambda args: args['dh0'] is not None, - 'USE_FINAL_STATE_GRADIENT': lambda args: args['dht'] is not None, - 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None, -}) -@triton.autotune( - configs=[ - triton.Config({}, num_warps=num_warps) - for num_warps in [4] - ], - key=['BK', 'BV', 'USE_G', 'USE_G_GAMMA', 'USE_GK', 'USE_GV', 'STATE_V_FIRST'], - **autotune_cache_kwargs, -) -@triton.jit(do_not_specialize=['B', 'T']) -def fused_recurrent_bwd_kernel( - q, - k, - v, - g, - g_gamma, - gk, - gv, - o, - h0, - do, - dq, - dk, - dv, - dg, - dgk, - dgv, - dht, - dh0, - cu_seqlens, - scale, - B, - T, - H: tl.constexpr, - K: tl.constexpr, - V: tl.constexpr, - BK: tl.constexpr, - BV: tl.constexpr, - REVERSE: tl.constexpr, - USE_G: tl.constexpr, - USE_G_GAMMA: tl.constexpr, - USE_GK: tl.constexpr, - USE_GV: tl.constexpr, - USE_INITIAL_STATE: tl.constexpr, - STORE_INITIAL_STATE_GRADIENT: tl.constexpr, - USE_FINAL_STATE_GRADIENT: tl.constexpr, - IS_VARLEN: tl.constexpr, - STATE_V_FIRST: tl.constexpr = False, -): - i_v, i_k, i_nh = tl.program_id(0).to(tl.int64), tl.program_id(1).to(tl.int64), tl.program_id(2).to(tl.int64) - i_n, i_h = i_nh // H, i_nh % H - - all = B * T - if IS_VARLEN: - bos, eos = tl.load(cu_seqlens + i_n).to(tl.int64), tl.load(cu_seqlens + i_n + 1).to(tl.int64) - T = eos - bos - else: - bos, eos = i_n * T, i_n * T + T - NV = tl.cdiv(V, BV) - - o_k = i_k * BK + tl.arange(0, BK) - o_v = i_v * BV + tl.arange(0, BV) - m_k = o_k < K - m_v = o_v < V - if STATE_V_FIRST: - m_h = m_v[:, None] & m_k[None, :] - else: - m_h = m_k[:, None] & m_v[None, :] - - p_k = k + (bos + ((T-1) if REVERSE else 0)) * H*K + i_h * K + o_k - p_v = v + (bos + ((T-1) if REVERSE else 0)) * H*V + i_h * V + o_v - p_do = do + (bos + ((T-1) if REVERSE else 0)) * H*V + i_h * V + o_v - p_dq = dq + ((i_v * all + bos) + ((T-1) if REVERSE else 0)) * H*K + i_h * K + o_k - if USE_G: - p_g = g + (bos + ((T-1) if REVERSE else 0)) * H + i_h - if USE_GK: - p_gk = gk + (bos + ((T-1) if REVERSE else 0)) * H*K + i_h * K + o_k - if USE_GV: - p_gv = gv + (bos + ((T-1) if REVERSE else 0)) * H*V + i_h * V + o_v - if USE_G_GAMMA: - b_g_gamma = tl.load(g_gamma + i_h) - - if STATE_V_FIRST: - b_h = tl.zeros([BV, BK], dtype=tl.float32) - else: - b_h = tl.zeros([BK, BV], dtype=tl.float32) - if USE_INITIAL_STATE: - if STATE_V_FIRST: - p_h0 = h0 + i_nh * K*V + o_v[:, None] * K + o_k[None, :] - else: - p_h0 = h0 + i_nh * K*V + o_k[:, None] * V + o_v[None, :] - b_h += tl.load(p_h0, mask=m_h, other=0).to(tl.float32) - - for _ in range(0, T): - b_k = tl.load(p_k, mask=m_k, other=0).to(tl.float32) - b_v = tl.load(p_v, mask=m_v, other=0).to(tl.float32) - b_do = tl.load(p_do, mask=m_v, other=0).to(tl.float32) - if USE_G: - b_g = tl.load(p_g).to(tl.float32) - b_h = b_h * exp(b_g) - if USE_G_GAMMA: - b_h = b_h * exp(b_g_gamma) - if USE_GK: - b_gk = tl.load(p_gk, mask=m_k, other=0).to(tl.float32) - if STATE_V_FIRST: - b_h = b_h * exp(b_gk[None, :]) - else: - b_h = b_h * exp(b_gk[:, None]) - if USE_GV: - b_gv = tl.load(p_gv, mask=m_v, other=0).to(tl.float32) - if STATE_V_FIRST: - b_h = b_h * exp(b_gv[:, None]) - else: - b_h = b_h * exp(b_gv[None, :]) - if STATE_V_FIRST: - b_h += b_v[:, None] * b_k[None, :] - b_dq = tl.sum(b_h * b_do[:, None], axis=0) * scale - else: - b_h += b_k[:, None] * b_v[None, :] - b_dq = tl.sum(b_h * b_do[None, :], axis=1) * scale - tl.store(p_dq, b_dq.to(p_dq.dtype.element_ty), mask=m_k) - - p_k += (-1 if REVERSE else 1) * H*K - p_v += (-1 if REVERSE else 1) * H*V - p_do += (-1 if REVERSE else 1) * H*V - p_dq += (-1 if REVERSE else 1) * H*K - if USE_G: - p_g += (-1 if REVERSE else 1) * H - if USE_GK: - p_gk += (-1 if REVERSE else 1) * H*K - if USE_GV: - p_gv += (-1 if REVERSE else 1) * H*V - - # sync threads - tl.debug_barrier() - - p_q = q + (bos + ((T - 1) if not REVERSE else 0)) * H*K + i_h * K + o_k - p_k = k + (bos + ((T - 1) if not REVERSE else 0)) * H*K + i_h * K + o_k - p_v = v + (bos + ((T - 1) if not REVERSE else 0)) * H*V + i_h * V + o_v - - p_do = do + (bos + ((T - 1) if not REVERSE else 0)) * H*V + i_h * V + o_v - p_dq = dq + ((i_v * all + bos) + ((T - 1) if not REVERSE else 0)) * H*K + i_h * K + o_k - p_dk = dk + ((i_v * all + bos) + ((T - 1) if not REVERSE else 0)) * H*K + i_h * K + o_k - p_dv = dv + ((i_k * all + bos) + ((T - 1) if not REVERSE else 0)) * H*V + i_h * V + o_v - if USE_G: - p_g = g + (bos + ((T - 1) if not REVERSE else 0)) * H + i_h - p_dg = dg + ((i_k * NV + i_v) * all + bos + ((T - 1) if not REVERSE else 0)) * H + i_h - if USE_GK: - p_gk = gk + (bos + ((T - 1) if not REVERSE else 0)) * H*K + i_h * K + o_k - p_dgk = dgk + ((i_v * all + bos) + ((T - 1) if not REVERSE else 0)) * H*K + i_h * K + o_k - if USE_GV: - p_o = o + (bos + ((T - 1) if not REVERSE else 0)) * H*V + i_h * V + o_v - p_gv = gv + (bos + ((T - 1) if not REVERSE else 0)) * H*V + i_h * V + o_v - p_dgv = dgv + ((i_k * all + bos) + ((T - 1) if not REVERSE else 0)) * H*V + i_h * V + o_v - - if STATE_V_FIRST: - b_dh = tl.zeros([BV, BK], dtype=tl.float32) - else: - b_dh = tl.zeros([BK, BV], dtype=tl.float32) - if USE_FINAL_STATE_GRADIENT: - if STATE_V_FIRST: - p_dht = dht + i_nh * K*V + o_v[:, None] * K + o_k[None, :] - else: - p_dht = dht + i_nh * K*V + o_k[:, None] * V + o_v[None, :] - b_dh += tl.load(p_dht, mask=m_h, other=0).to(tl.float32) - - if USE_G: - b_dg = tl.sum(b_h * b_dh) - if USE_GK: - if STATE_V_FIRST: - b_dgk = tl.sum(b_h * b_dh, 0) - else: - b_dgk = tl.sum(b_h * b_dh, 1) - if USE_GV: - if STATE_V_FIRST: - b_dgv = tl.sum(b_h * b_dh, 1) - else: - b_dgv = tl.sum(b_h * b_dh, 0) - - for _ in range(T): - b_q = tl.load(p_q, mask=m_k, other=0).to(tl.float32) - b_k = tl.load(p_k, mask=m_k, other=0).to(tl.float32) - b_v = tl.load(p_v, mask=m_v, other=0).to(tl.float32) - b_do = tl.load(p_do, mask=m_v, other=0).to(tl.float32) - if STATE_V_FIRST: - b_dh += b_do[:, None] * (b_q * scale)[None, :] - b_dk = tl.sum(b_dh * b_v[:, None], axis=0) - b_dv = tl.sum(b_dh * b_k[None, :], axis=1) - else: - b_dh += (b_q * scale)[:, None] * b_do[None, :] - b_dk = tl.sum(b_dh * b_v[None, :], axis=1) - b_dv = tl.sum(b_dh * b_k[:, None], axis=0) - - if USE_G: - b_g = tl.load(p_g).to(tl.float32) - b_dq = tl.load(p_dq, mask=m_k, other=0).to(tl.float32) - b_dg += tl.sum(b_q * b_dq - b_k * b_dk) - b_dh *= exp(b_g) - tl.store(p_dg, b_dg.to(p_dg.dtype.element_ty)) - if USE_G_GAMMA: - b_dh *= exp(b_g_gamma) - if USE_GK: - b_gk = tl.load(p_gk, mask=m_k, other=0).to(tl.float32) - b_dq = tl.load(p_dq, mask=m_k, other=0).to(tl.float32) - b_dgk += b_q * b_dq - b_k * b_dk - if STATE_V_FIRST: - b_dh *= exp(b_gk)[None, :] - else: - b_dh *= exp(b_gk)[:, None] - tl.store(p_dgk, b_dgk.to(p_dgk.dtype.element_ty), mask=m_k) - if USE_GV: - b_o = tl.load(p_o, mask=m_v, other=0).to(tl.float32) - b_gv = tl.load(p_gv, mask=m_v, other=0).to(tl.float32) - if i_k == 0: - b_dgv += b_o * b_do - b_dgv -= b_v * b_dv - if STATE_V_FIRST: - b_dh *= exp(b_gv)[:, None] - else: - b_dh *= exp(b_gv)[None, :] - tl.store(p_dgv, b_dgv.to(p_dgv.dtype.element_ty), mask=m_v) - - tl.store(p_dk, b_dk.to(p_dk.dtype.element_ty), mask=m_k) - tl.store(p_dv, b_dv.to(p_dv.dtype.element_ty), mask=m_v) - - p_q += (1 if REVERSE else -1) * H*K - p_k += (1 if REVERSE else -1) * H*K - p_v += (1 if REVERSE else -1) * H*V - - p_do += (1 if REVERSE else -1) * H*V - p_dq += (1 if REVERSE else -1) * H*K - p_dk += (1 if REVERSE else -1) * H*K - p_dv += (1 if REVERSE else -1) * H*V - if USE_G: - p_g += (1 if REVERSE else -1) * H - p_dg += (1 if REVERSE else -1) * H - if USE_GK: - p_gk += (1 if REVERSE else -1) * H*K - p_dgk += (1 if REVERSE else -1) * H*K - if USE_GV: - p_o += (1 if REVERSE else -1) * H*V - p_gv += (1 if REVERSE else -1) * H*V - p_dgv += (1 if REVERSE else -1) * H*V - - if STORE_INITIAL_STATE_GRADIENT: - if STATE_V_FIRST: - p_dh0 = dh0 + i_nh * K*V + o_v[:, None] * K + o_k[None, :] - else: - p_dh0 = dh0 + i_nh * K*V + o_k[:, None] * V + o_v[None, :] - tl.store(p_dh0, b_dh.to(p_dh0.dtype.element_ty), mask=m_h) - - -def fused_recurrent_fwd( - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - g: torch.Tensor | None = None, - g_gamma: torch.Tensor | None = None, - gk: torch.Tensor | None = None, - gv: torch.Tensor | None = None, - scale: float | None = None, - initial_state: torch.Tensor | None = None, - output_final_state: bool = False, - reverse: bool = False, - state_v_first: bool = False, - cu_seqlens: torch.LongTensor | None = None, -): - B, T, H, K, V = *k.shape, v.shape[-1] - N = B if cu_seqlens is None else len(cu_seqlens) - 1 - BK, BV = min(triton.next_power_of_2(K), 64), min(triton.next_power_of_2(V), 64) - NK, NV = triton.cdiv(K, BK), triton.cdiv(V, BV) - - h0 = initial_state - if output_final_state: - if state_v_first: - ht = q.new_empty(N, H, V, K, dtype=torch.float32) - else: - ht = q.new_empty(N, H, K, V, dtype=torch.float32) - else: - ht = None - o = q.new_empty(NK, *v.shape, dtype=torch.float32) - - grid = (NV, NK, N * H) - fused_recurrent_fwd_kernel[grid]( - q=q, - k=k, - v=v, - g=g, - g_gamma=g_gamma, - gk=gk, - gv=gv, - o=o, - h0=h0, - ht=ht, - cu_seqlens=cu_seqlens, - scale=scale, - T=T, - B=B, - H=H, - K=K, - V=V, - BK=BK, - BV=BV, - USE_G=g is not None, - USE_G_GAMMA=g_gamma is not None, - USE_GK=gk is not None, - USE_GV=gv is not None, - REVERSE=reverse, - STATE_V_FIRST=state_v_first, - ) - o = o.sum(0) - return o, ht - - -def fused_recurrent_bwd( - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - g: torch.Tensor | None = None, - g_gamma: torch.Tensor | None = None, - gk: torch.Tensor | None = None, - gv: torch.Tensor | None = None, - o: torch.Tensor | None = None, - do: torch.Tensor | None = None, - dht: torch.Tensor | None = None, - scale: float | None = None, - initial_state: torch.Tensor | None = None, - reverse: bool = False, - state_v_first: bool = False, - cu_seqlens: torch.LongTensor | None = None, -): - B, T, H, K, V = *k.shape, v.shape[-1] - N = B if cu_seqlens is None else len(cu_seqlens) - 1 - - BK, BV = min(triton.next_power_of_2(K), 64), min(triton.next_power_of_2(V), 64) - NK, NV = triton.cdiv(K, BK), triton.cdiv(V, BV) - - h0 = initial_state - dq = q.new_empty(NV, *q.shape, dtype=torch.float32) - dk = q.new_empty(NV, *k.shape, dtype=torch.float32) - dv = q.new_empty(NK, *v.shape, dtype=torch.float32) - dh0 = torch.empty_like(h0) if h0 is not None else None - - dg, dgk, dgv = None, None, None - if g is not None: - dg = g.new_empty(NK*NV, *g.shape, dtype=torch.float32) - if gk is not None: - dgk = gk.new_empty(NV, *gk.shape, dtype=torch.float32) - if gv is not None: - dgv = gv.new_empty(NK, *gv.shape, dtype=torch.float32) - - grid = (NV, NK, N * H) - fused_recurrent_bwd_kernel[grid]( - q=q, - k=k, - v=v, - g=g, - g_gamma=g_gamma, - gk=gk, - gv=gv, - o=o, - h0=h0, - do=do, - dq=dq, - dk=dk, - dv=dv, - dg=dg, - dgk=dgk, - dgv=dgv, - dht=dht, - dh0=dh0, - cu_seqlens=cu_seqlens, - scale=scale, - B=B, - T=T, - H=H, - K=K, - V=V, - BK=BK, - BV=BV, - USE_G=g is not None, - USE_G_GAMMA=g_gamma is not None, - USE_GK=gk is not None, - USE_GV=gv is not None, - REVERSE=reverse, - STATE_V_FIRST=state_v_first, - ) - dq = dq.sum(0) - dk = dk.sum(0) - dv = dv.sum(0) - if g is not None: - dg = dg.sum(0).to(g) - if gk is not None: - dgk = dgk.sum(0).to(gk) - if gv is not None: - dgv = dgv.sum(0).to(gv) - - return dq, dk, dv, dg, dgk, dgv, dh0 - - -class FusedRecurrentFunction(torch.autograd.Function): - - @staticmethod - @input_guard - @autocast_custom_fwd - def forward( - ctx, - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - g: torch.Tensor | None = None, - g_gamma: torch.Tensor | None = None, - gk: torch.Tensor | None = None, - gv: torch.Tensor | None = None, - scale: float | None = None, - initial_state: torch.Tensor | None = None, - output_final_state: bool = False, - reverse: bool = False, - state_v_first: bool = False, - cu_seqlens: torch.LongTensor | None = None, - ): - o, ht = fused_recurrent_fwd( - q=q, - k=k, - v=v, - g=g, - g_gamma=g_gamma, - gk=gk, - gv=gv, - scale=scale, - initial_state=initial_state, - output_final_state=output_final_state, - reverse=reverse, - cu_seqlens=cu_seqlens, - state_v_first=state_v_first, - ) - ctx.save_for_backward(q, k, v, g, g_gamma, gk, gv, initial_state, o) - ctx.scale = scale - ctx.reverse = reverse - ctx.cu_seqlens = cu_seqlens - ctx.state_v_first = state_v_first - return o.to(q.dtype), ht - - @staticmethod - @input_guard - @autocast_custom_bwd - def backward(ctx, do, dht): - q, k, v, g, g_gamma, gk, gv, initial_state, o = ctx.saved_tensors - dq, dk, dv, dg, dgk, dgv, dh0 = fused_recurrent_bwd( - q=q, - k=k, - v=v, - g=g, - g_gamma=g_gamma, - gk=gk, - gv=gv, - o=o, - do=do, - dht=dht, - scale=ctx.scale, - initial_state=initial_state, - reverse=ctx.reverse, - cu_seqlens=ctx.cu_seqlens, - state_v_first=ctx.state_v_first, - ) - return dq.to(q.dtype), dk.to(k.dtype), dv.to(v.dtype), dg, None, dgk, dgv, None, dh0, None, None, None, None - - -def fused_recurrent( - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - g: torch.Tensor | None = None, - g_gamma: torch.Tensor | None = None, - gk: torch.Tensor | None = None, - gv: torch.Tensor | None = None, - scale: float | None = None, - initial_state: torch.Tensor | None = None, - output_final_state: bool = False, - reverse: bool = False, - state_v_first: bool = False, - cu_seqlens: torch.LongTensor | None = None, -): - if scale is None: - scale = k.shape[-1] ** -0.5 - return FusedRecurrentFunction.apply( - q, - k, - v, - g, - g_gamma, - gk, - gv, - scale, - initial_state, - output_final_state, - reverse, - state_v_first, - cu_seqlens, - ) diff --git a/build/torch-xpu/ops/common/gate.py b/build/torch-xpu/ops/common/gate.py deleted file mode 100644 index 3f60b99a59145c2c4189aa19b5607e1d0c5e67ac..0000000000000000000000000000000000000000 --- a/build/torch-xpu/ops/common/gate.py +++ /dev/null @@ -1,107 +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 - -# Shared gate helpers reused across delta-rule family ops (KDA, GDN, ...). - -import torch -import triton -import triton.language as tl - -from ...utils import autocast_custom_bwd, autocast_custom_fwd, input_guard - - -@triton.jit -def fused_beta_sigmoid_fwd_kernel( - x, - y, - scale, - n_elements, - BLOCK_SIZE: tl.constexpr, -): - pid = tl.program_id(0).to(tl.int64) - offs = pid * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE).to(tl.int64) - mask = offs < n_elements - b_x = tl.load(x + offs, mask=mask, other=0).to(tl.float32) - b_y = scale * tl.sigmoid(b_x) - tl.store(y + offs, b_y.to(y.dtype.element_ty), mask=mask) - - -@triton.jit -def fused_beta_sigmoid_bwd_kernel( - x, - dy, - dx, - scale, - n_elements, - BLOCK_SIZE: tl.constexpr, -): - pid = tl.program_id(0).to(tl.int64) - offs = pid * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE).to(tl.int64) - mask = offs < n_elements - b_x = tl.load(x + offs, mask=mask, other=0).to(tl.float32) - b_dy = tl.load(dy + offs, mask=mask, other=0).to(tl.float32) - b_y = tl.sigmoid(b_x) - b_dx = b_dy * scale * b_y * (1.0 - b_y) - tl.store(dx + offs, b_dx.to(dx.dtype.element_ty), mask=mask) - - -_BETA_SIGMOID_BLOCK_SIZE = 2048 -_BETA_SIGMOID_NUM_WARPS = 8 - - -def fused_beta_sigmoid_fwd(x: torch.Tensor, scale: float = 1.0) -> torch.Tensor: - y = torch.empty_like(x, dtype=torch.float32) - n_elements = x.numel() - grid = (triton.cdiv(n_elements, _BETA_SIGMOID_BLOCK_SIZE),) - fused_beta_sigmoid_fwd_kernel[grid]( - x, - y, - scale, - n_elements, - BLOCK_SIZE=_BETA_SIGMOID_BLOCK_SIZE, - num_warps=_BETA_SIGMOID_NUM_WARPS, - ) - return y - - -def fused_beta_sigmoid_bwd(x: torch.Tensor, dy: torch.Tensor, scale: float = 1.0) -> torch.Tensor: - dx = torch.empty_like(x) - n_elements = x.numel() - grid = (triton.cdiv(n_elements, _BETA_SIGMOID_BLOCK_SIZE),) - fused_beta_sigmoid_bwd_kernel[grid]( - x, - dy, - dx, - scale, - n_elements, - BLOCK_SIZE=_BETA_SIGMOID_BLOCK_SIZE, - num_warps=_BETA_SIGMOID_NUM_WARPS, - ) - return dx - - -class BetaSigmoidFunction(torch.autograd.Function): - @staticmethod - @input_guard - @autocast_custom_fwd - def forward(ctx, x: torch.Tensor, scale: float = 1.0) -> torch.Tensor: - y = fused_beta_sigmoid_fwd(x, scale) - ctx.save_for_backward(x) - ctx.scale = scale - return y - - @staticmethod - @input_guard - @autocast_custom_bwd - def backward(ctx, dy: torch.Tensor): - (x,) = ctx.saved_tensors - dx = fused_beta_sigmoid_bwd(x, dy, ctx.scale) - return dx.type_as(x), None - - -def fused_beta_sigmoid(x: torch.Tensor, scale: float = 1.0) -> torch.Tensor: - return BetaSigmoidFunction.apply(x, scale) diff --git a/build/torch-xpu/ops/common/intracard_cp.py b/build/torch-xpu/ops/common/intracard_cp.py deleted file mode 100644 index fb82b50041d010cba28828b756789e91fb3307c8..0000000000000000000000000000000000000000 --- a/build/torch-xpu/ops/common/intracard_cp.py +++ /dev/null @@ -1,624 +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 - -"""Intra-Card Context Parallel for KDA inference (varlen mode only). - -Optimized: all CPU-side index computation uses pure Python loops instead of -torch tensor operations (repeat_interleave, arange, cumsum, etc.) to eliminate -per-op overhead on tiny arrays. GPU tensors are created directly from Python -lists to minimize cudaStreamSynchronize calls. -""" - -from __future__ import annotations - -import logging -import weakref -from collections import OrderedDict -from typing import NamedTuple - -import torch -import triton - -from ...ops.common.chunk_delta_h import chunk_gated_delta_rule_fwd_kernel_h_blockdim64 -from ...ops.cp.chunk_delta_h import pre_process_fwd_kernel_merged -from ...ops.utils.index import prepare_chunk_indices, prepare_chunk_offsets -from ...utils import get_multiprocessor_count - -logger = logging.getLogger(__name__) - - -# Cache for intracard_fwd_h precomputation (Python results + GPU tensors) -# Key: object id of cu_seqlens (consistent with tensor_cache philosophy) -_intracard_cache: OrderedDict[tuple, _CacheEntry] = OrderedDict() -_INTRACARD_CACHE_MAXSIZE = 32 - - -class _CacheEntry(NamedTuple): - """Cache entry for intracard_fwd_h precomputation. - - Caches both Python computation results and GPU tensors to eliminate - redundant CPU→GPU transfers and Python loop computation. - """ - # Keep a weak reference to validate id-based key safety. - # If Python reuses an object id after GC, this guard prevents stale hits. - cu_seqlens_ref: weakref.ReferenceType[torch.Tensor] - # From prepare_subseq_cu_seqlens - cu_seqlens_subseq_values: list[int] - split_info: SplitSeqInfo - total_subseqs: int - # From _precompute_intracard_indices - cu_seqlens_split_values: list[int] - S_split_total: int - non_first_indices: list[int] - first_subseq_indices: list[int] - last_subseq_indices: list[int] - num_non_first: int - merge_seq_offsets: list[int] - merge_init_offsets: list[int] - # GPU tensors (cached to avoid H2D transfer) - cu_seqlens_subseq_gpu: torch.Tensor - cu_seqlens_split_flat: torch.Tensor - - -class SplitSeqInfo(NamedTuple): - """Information about split sequences (Python lists for zero-overhead access).""" - split_seq_ids: list[int] # [num_split_seqs] original sequence indices - start_subseq_idx: list[int] # [num_split_seqs] start index in subseq array - num_subseqs: list[int] # [num_split_seqs] number of sub-sequences per split - - @property - def num_split_seqs(self) -> int: - return len(self.split_seq_ids) - - def __bool__(self) -> bool: - return self.num_split_seqs > 0 - - -def _raw_chunk_gated_delta_rule_fwd_h( - k: torch.Tensor, - w: torch.Tensor, - u: torch.Tensor, - g: torch.Tensor | None = None, - gk: torch.Tensor | None = None, - initial_state: torch.Tensor | None = None, - output_final_state: bool = False, - chunk_size: int = 64, - save_new_value: bool = True, - state_v_first: bool = False, - cu_seqlens: torch.LongTensor | None = None, - chunk_indices: torch.LongTensor | None = None, -) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor | None]: - B, T, H, K, V, HV = *k.shape, u.shape[-1], u.shape[2] - BT = chunk_size - - if chunk_indices is None and cu_seqlens is not None: - chunk_indices = prepare_chunk_indices(cu_seqlens, chunk_size) - if cu_seqlens is None: - N, NT, chunk_offsets = B, triton.cdiv(T, BT), None - else: - N, NT, chunk_offsets = len(cu_seqlens) - 1, len(chunk_indices), prepare_chunk_offsets(cu_seqlens, BT) - - if state_v_first: - h = k.new_empty(B, NT, HV, V, K) - final_state = k.new_zeros(N, HV, V, K, dtype=torch.float32) if output_final_state else None - else: - h = k.new_empty(B, NT, HV, K, V) - final_state = k.new_zeros(N, HV, K, V, dtype=torch.float32) if output_final_state else None - v_new = torch.empty_like(u) if save_new_value else None - - def grid(meta): - return (triton.cdiv(V, meta['BV']), N * HV) - - chunk_gated_delta_rule_fwd_kernel_h_blockdim64[grid]( - k=k, - v=u, - w=w, - v_new=v_new, - g=g, - gk=gk, - h=h, - h0=initial_state, - ht=final_state, - cu_seqlens=cu_seqlens, - chunk_offsets=chunk_offsets, - T=T, - HV=HV, - H=H, - K=K, - V=V, - BT=BT, - STATE_V_FIRST=state_v_first, - ) - return h, v_new, final_state - - -def compute_subseq_len( - seq_len: int, - num_sms: int, - num_heads: int, - chunk_size: int = 64, -) -> int: - """Compute sub-sequence length for intracard splitting. - - For linear recurrence (fwd_h), the sequential scan is the bottleneck. - Splitting always reduces the critical path and helps, as long as the - sequence is long enough to amortize the pre_scan + merge overhead. - - The fwd_h kernel grid is (num_v_blocks, N*HV) where num_v_blocks ≈ 2. - Each sub-sequence contributes 2*HV blocks. We target enough splits so - that even a single long sequence can saturate all SMs. - - A floor on subseq_chunks (MIN_SUBSEQ_CHUNKS) prevents subseq_len from - being too small, which would cause prepare_subseq_cu_seqlens to - unnecessarily split shorter sequences in mixed-length batches - (split threshold = 2 * subseq_len). - """ - seq_chunks = (seq_len + chunk_size - 1) // chunk_size - - if seq_chunks < 8: - return seq_len - - # Target splits: saturate SMs with the longest sequence alone. - # Each sub-seq contributes NUM_V_BLOCKS * num_heads blocks. - # Always at least 4 — for linear recurrence, CP4 always helps. - NUM_V_BLOCKS = 2 - target_splits = max(4, num_sms // (NUM_V_BLOCKS * num_heads)) - - subseq_chunks = (seq_chunks + target_splits - 1) // target_splits - - # Floor: prevent subseq_len from being too small. - # With chunk_size=64, MIN_SUBSEQ_CHUNKS=128 → subseq_len >= 8192 tokens, - # split threshold (3 * subseq_len) = 24576 tokens. - # Sequences shorter than it won't be split. - MIN_SUBSEQ_CHUNKS = 128 - subseq_chunks = max(subseq_chunks, MIN_SUBSEQ_CHUNKS) - - return subseq_chunks * chunk_size - - -def prepare_subseq_cu_seqlens( - cu_seqlens_cpu: torch.Tensor, - subseq_len: int, - chunk_size: int = 64, - max_splits: int = 32, -) -> tuple[list[int], SplitSeqInfo | bool, int]: - """Insert subseq split points into original cu_seqlens. - - Optimized: uses pure Python loops instead of torch tensor operations - for the small index arrays (typically 1-32 elements). - - Returns: - boundaries: List of cu_seqlens boundaries (can be used directly by _precompute_intracard_indices) - split_info: SplitSeqInfo for sequences that need splitting, or False if no splitting needed - total_subseqs: Total number of sub-sequences after splitting - """ - N = len(cu_seqlens_cpu) - 1 - if N == 0: - return cu_seqlens_cpu.tolist(), False, 0 - - subseq_chunks = (subseq_len + chunk_size - 1) // chunk_size - threshold_subseq_len = 3 * subseq_len - - split_seq_ids: list[int] = [] - start_subseq_idxs: list[int] = [] - num_subseqs_list: list[int] = [] - - # Build boundaries using pure Python loop - boundaries: list[int] = [0] - cumsum_offset = 0 - - for i in range(N): - seq_start = int(cu_seqlens_cpu[i].item()) - seq_end = int(cu_seqlens_cpu[i + 1].item()) - seq_len_i = seq_end - seq_start - seq_chunks_i = (seq_len_i + chunk_size - 1) // chunk_size - - if seq_len_i >= threshold_subseq_len: - # This sequence needs splitting - num_ss = min(max_splits, (seq_chunks_i + subseq_chunks - 1) // subseq_chunks) - chunks_per = (seq_chunks_i + num_ss - 1) // num_ss - actual_ssl = chunks_per * chunk_size - - split_seq_ids.append(i) - start_subseq_idxs.append(cumsum_offset) - num_subseqs_list.append(num_ss) - - for j in range(num_ss): - boundary = min(seq_start + (j + 1) * actual_ssl, seq_end) - boundaries.append(boundary) - cumsum_offset += num_ss - else: - # No split needed, single sub-sequence - boundaries.append(seq_end) - cumsum_offset += 1 - - if not split_seq_ids: - return cu_seqlens_cpu.tolist(), False, 0 - - total_subseqs = cumsum_offset - - split_info = SplitSeqInfo( - split_seq_ids=split_seq_ids, - start_subseq_idx=start_subseq_idxs, - num_subseqs=num_subseqs_list, - ) - - return boundaries, split_info, total_subseqs - - -def intracard_pre_scan( - kg: torch.Tensor, - w: torch.Tensor, - u: torch.Tensor, - g: torch.Tensor | None, - gk: torch.Tensor | None, - cu_seqlens_subseq_split: torch.Tensor, - S_split: int, - chunk_size: int = 64, -): - H, K, V, HV = kg.shape[2], kg.shape[3], u.shape[3], u.shape[2] - BK = triton.next_power_of_2(K) - BLOCK_SIZE = 32 if K <= 64 else 64 - - hm = kg.new_empty(S_split, HV, K, V + K, dtype=torch.float32) - - grid = (triton.cdiv(V, BLOCK_SIZE) + triton.cdiv(K, BLOCK_SIZE), HV, S_split) - pre_process_fwd_kernel_merged[grid]( - k=kg, - v=u, - w=w, - g=g, - gk=gk, - bg=None, - u=u, - hm=hm, - cu_seqlens=cu_seqlens_subseq_split, - T=0, - H=H, - HV=HV, - K=K, - V=V, - BT=chunk_size, - BLOCK_SIZE=BLOCK_SIZE, - BK1=BK, - MULTI_SEQS=True, - ) - - return hm - - -def intracard_merge( - hm: torch.Tensor, - split_info: SplitSeqInfo, - num_non_first: int, - merge_seq_offsets: list[int], - merge_init_offsets: list[int], - device: torch.device, - initial_state: torch.Tensor | None = None, - state_v_first: bool = False, -) -> tuple[torch.Tensor | None, int]: - """Merge sub-sequence states using pre-computed parameters. - - All CPU-side preparation (cumsum, offset lists) is done in the caller - using pure Python loops. This function only creates GPU tensors and - launches the merge kernel. - """ - from ...ops.cp.chunk_delta_h import merge_fwd_bwd_kernel - - if num_non_first == 0: - return None, 0 - - HV = hm.shape[1] - K = hm.shape[2] - V = hm.shape[3] - K - BK = triton.next_power_of_2(K) - - num_split_seqs = split_info.num_split_seqs - - # Create all small GPU tensors from Python lists in one batch - # Merge into a single CPU→GPU transfer to minimize cudaStreamSynchronize - all_int_data = merge_seq_offsets + merge_init_offsets + split_info.split_seq_ids - all_tensor = torch.tensor(all_int_data, dtype=torch.int32, device=device) - n_so = len(merge_seq_offsets) - n_io = len(merge_init_offsets) - seq_offsets = all_tensor[:n_so] - init_offsets = all_tensor[n_so:n_so+n_io] - h0_seq_ids = all_tensor[n_so+n_io:] - - if state_v_first: - initial_states_merge = hm.new_empty(num_non_first, HV, V, K, dtype=torch.float32) - else: - initial_states_merge = hm.new_empty(num_non_first, HV, K, V, dtype=torch.float32) - - def grid(meta): - return (triton.cdiv(V, meta['BV']), num_split_seqs, HV) - - merge_fwd_bwd_kernel[grid]( - h=initial_states_merge, - ag_hm=hm, - pre_or_post_num_ranks=num_split_seqs, - rank=0, - seq_offsets=seq_offsets, - init_offsets=init_offsets, - h0_seq_ids=h0_seq_ids, - h0=initial_state, - HV=HV, - K=K, - V=V, - BK=BK, - FORWARD=True, - INTRACARD_MODE=True, - NUM_SEQ_ENTRIES=num_split_seqs, - STATE_V_FIRST=state_v_first, - ) - - return initial_states_merge, num_non_first - - -def _precompute_intracard_indices( - split_info: SplitSeqInfo, - cu_seqlens_subseq_values: list[int], - N_orig: int, -) -> tuple[list[int], int, list[int], list[int], list[int], int, list[int], list[int]]: - """Pre-compute all derived indices using pure Python loops. - - Returns: - cu_seqlens_split_values: flattened cu_seqlens boundaries for split seqs (for pre_scan) - S_split_total: total number of sub-sequences from splits - non_first_indices: indices for scattering merge results into initial_state_expanded - first_subseq_indices: indices of first sub-sequence for each original sequence - last_subseq_indices: indices of last sub-sequence for each original sequence - num_non_first: total non-first sub-sequences (merge work) - merge_seq_offsets: cumulative sub-sequence counts for merge kernel - merge_init_offsets: cumulative non-first counts for merge kernel - """ - starts = split_info.start_subseq_idx - num_ss = split_info.num_subseqs - split_ids = split_info.split_seq_ids - - # cu_seqlens_split_values: for each split seq, extract [start:start+n+1] boundaries - cu_seqlens_split_values: list[int] = [] - S_split_total = 0 - for s, n in zip(starts, num_ss): - cu_seqlens_split_values.extend(cu_seqlens_subseq_values[s:s+n+1]) - S_split_total += n - - # num_subseqs_per_seq: [N_orig], default 1 for unsplit sequences - num_subseqs_per_seq = [1] * N_orig - for sid, nss in zip(split_ids, num_ss): - num_subseqs_per_seq[sid] = nss - - # non_first_indices: for scattering merged initial states - non_first_indices: list[int] = [] - for s, n in zip(starts, num_ss): - for j in range(1, n): - non_first_indices.append(s + j) - - # first_subseq_indices: for scattering original initial states - first_subseq_indices: list[int] = [0] - running = 0 - for i in range(N_orig - 1): - running += num_subseqs_per_seq[i] - first_subseq_indices.append(running) - - # last_subseq_indices: for gathering final states - last_subseq_indices: list[int] = [] - running = 0 - for n in num_subseqs_per_seq: - running += n - last_subseq_indices.append(running - 1) - - # merge parameters - merge_seq_offsets: list[int] = [0] - merge_init_offsets: list[int] = [0] - for n in num_ss: - merge_seq_offsets.append(merge_seq_offsets[-1] + n) - merge_init_offsets.append(merge_init_offsets[-1] + n - 1) - num_non_first = merge_init_offsets[-1] - - return ( - cu_seqlens_split_values, - S_split_total, - non_first_indices, - first_subseq_indices, - last_subseq_indices, - num_non_first, - merge_seq_offsets, - merge_init_offsets, - ) - - -def intracard_fwd_h( - k: torch.Tensor, - w: torch.Tensor, - u: torch.Tensor, - g: torch.Tensor | None = None, - gk: torch.Tensor | None = None, - initial_state: torch.Tensor | None = None, - output_final_state: bool = False, - chunk_size: int = 64, - save_new_value: bool = True, - state_v_first: bool = False, - cu_seqlens: torch.LongTensor | None = None, - cu_seqlens_cpu: torch.LongTensor | None = None, - chunk_indices: torch.LongTensor | None = None, - max_splits: int = 32, -) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor | None]: - assert cu_seqlens is not None, "intracard_fwd_h requires cu_seqlens" - - _, _, _Hq, K = k.shape - V = u.shape[-1] - HV = u.shape[2] - device = k.device - - if cu_seqlens_cpu is None: - cu_seqlens_cpu = cu_seqlens.cpu() - - seq_lens = torch.diff(cu_seqlens_cpu) - max_seq_len = int(seq_lens.max().item()) - num_sms = get_multiprocessor_count() - subseq_len = compute_subseq_len(max_seq_len, num_sms, HV, chunk_size) - - early_return = (seq_lens < 2 * subseq_len).all() - - cached = None - cache_key = None - - if not early_return: - # Use object identity (id) for cache key, consistent with tensor_cache philosophy - # vLLM slice creates new Python objects per batch, so id(cu_seqlens) is safe - cache_key = ( - id(cu_seqlens), # Object identity, not content hash - subseq_len, - chunk_size, - max_splits, - str(device), - ) - cached = _intracard_cache.get(cache_key) - if cached is not None: - # Guard against rare Python id reuse after original tensor is GC-ed. - # We only consider it a hit when the weakref points to the current object. - if cached.cu_seqlens_ref() is cu_seqlens: - _intracard_cache.move_to_end(cache_key) - else: - _intracard_cache.pop(cache_key, None) - cached = None - - if cached is not None: - # Cache hit: reuse all precomputed results including GPU tensors - cu_seqlens_subseq_values = cached.cu_seqlens_subseq_values - split_info = cached.split_info - total_subseqs = cached.total_subseqs - cu_seqlens_split_values = cached.cu_seqlens_split_values - S_split_total = cached.S_split_total - non_first_indices = cached.non_first_indices - first_subseq_indices = cached.first_subseq_indices - last_subseq_indices = cached.last_subseq_indices - num_non_first = cached.num_non_first - merge_seq_offsets = cached.merge_seq_offsets - merge_init_offsets = cached.merge_init_offsets - cu_seqlens_subseq_gpu = cached.cu_seqlens_subseq_gpu - cu_seqlens_split_flat = cached.cu_seqlens_split_flat - else: - # Cache miss: compute Python lists - cu_seqlens_subseq_values, split_info, total_subseqs = prepare_subseq_cu_seqlens( - cu_seqlens_cpu, subseq_len, chunk_size, max_splits=max_splits - ) - - if early_return or not split_info: - return _raw_chunk_gated_delta_rule_fwd_h( - k=k, - w=w, - u=u, - g=g, - gk=gk, - initial_state=initial_state, - output_final_state=output_final_state, - chunk_size=chunk_size, - save_new_value=save_new_value, - cu_seqlens=cu_seqlens, - chunk_indices=chunk_indices, - state_v_first=state_v_first, - ) - - N_orig = len(cu_seqlens_cpu) - 1 - - if cached is None: - # Cache miss: continue Python computation and create GPU tensors - ( - cu_seqlens_split_values, - S_split_total, - non_first_indices, - first_subseq_indices, - last_subseq_indices, - num_non_first, - merge_seq_offsets, - merge_init_offsets, - ) = _precompute_intracard_indices(split_info, cu_seqlens_subseq_values, N_orig) - - # Create GPU tensors (will be cached for reuse) - dtype = cu_seqlens_cpu.dtype - cu_seqlens_subseq_gpu = torch.tensor(cu_seqlens_subseq_values, dtype=dtype, device=device) - cu_seqlens_split_flat = torch.tensor(cu_seqlens_split_values, dtype=dtype, device=device) - - # Store all results in cache (including GPU tensors to avoid H2D) - _intracard_cache[cache_key] = _CacheEntry( - cu_seqlens_ref=weakref.ref(cu_seqlens), - cu_seqlens_subseq_values=cu_seqlens_subseq_values, - split_info=split_info, - total_subseqs=total_subseqs, - cu_seqlens_split_values=cu_seqlens_split_values, - S_split_total=S_split_total, - non_first_indices=non_first_indices, - first_subseq_indices=first_subseq_indices, - last_subseq_indices=last_subseq_indices, - num_non_first=num_non_first, - merge_seq_offsets=merge_seq_offsets, - merge_init_offsets=merge_init_offsets, - cu_seqlens_subseq_gpu=cu_seqlens_subseq_gpu, - cu_seqlens_split_flat=cu_seqlens_split_flat, - ) - # Evict oldest entries if over capacity - while len(_intracard_cache) > _INTRACARD_CACHE_MAXSIZE: - _intracard_cache.popitem(last=False) - - hm = intracard_pre_scan( - kg=k, - w=w, - u=u, - g=g, - gk=gk, - cu_seqlens_subseq_split=cu_seqlens_split_flat, - S_split=S_split_total, - chunk_size=chunk_size, - ) - - initial_states_merge, num_non_first = intracard_merge( - hm=hm, - split_info=split_info, - num_non_first=num_non_first, - merge_seq_offsets=merge_seq_offsets, - merge_init_offsets=merge_init_offsets, - device=device, - initial_state=initial_state, - state_v_first=state_v_first, - ) - - if state_v_first: - initial_state_expanded = k.new_zeros(total_subseqs, HV, V, K, dtype=torch.float32) - else: - initial_state_expanded = k.new_zeros(total_subseqs, HV, K, V, dtype=torch.float32) - - if initial_state is not None: - initial_state_expanded[first_subseq_indices] = initial_state - - if initial_states_merge is not None and num_non_first > 0: - initial_state_expanded[non_first_indices] = initial_states_merge - - chunk_indices_subseq = prepare_chunk_indices(cu_seqlens_subseq_gpu, chunk_size) - - h, v_new, final_state_subseq = _raw_chunk_gated_delta_rule_fwd_h( - k=k, - w=w, - u=u, - g=g, - gk=gk, - initial_state=initial_state_expanded, - output_final_state=output_final_state, - chunk_size=chunk_size, - save_new_value=save_new_value, - cu_seqlens=cu_seqlens_subseq_gpu, - chunk_indices=chunk_indices_subseq, - state_v_first=state_v_first, - ) - - if output_final_state and final_state_subseq is not None: - final_state = final_state_subseq[last_subseq_indices] - else: - final_state = final_state_subseq - - return h, v_new, final_state diff --git a/build/torch-xpu/ops/cp/__init__.py b/build/torch-xpu/ops/cp/__init__.py deleted file mode 100644 index 5d1f3f4575bd6c453a7769332fed3687114d7e4b..0000000000000000000000000000000000000000 --- a/build/torch-xpu/ops/cp/__init__.py +++ /dev/null @@ -1,32 +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 - -# 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", -] diff --git a/build/torch-xpu/ops/cp/chunk_delta_h.py b/build/torch-xpu/ops/cp/chunk_delta_h.py deleted file mode 100644 index dffa38c93f0862505e6ab854a72d390ed54889c1..0000000000000000000000000000000000000000 --- a/build/torch-xpu/ops/cp/chunk_delta_h.py +++ /dev/null @@ -1,922 +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 - -from typing import TYPE_CHECKING - -import torch -import torch.distributed as dist -import triton -import triton.language as tl - -from ...ops.cp.comm import all_gather_into_tensor -from ...ops.utils.op import exp2 -from ...utils import autotune_cache_kwargs, check_shared_mem - -if TYPE_CHECKING: - from ...ops.cp.context import FLACPContext - - -@triton.heuristics({ - 'USE_G': lambda args: args['g'] is not None, - 'USE_GK': lambda args: args['gk'] is not None, - 'USE_BG': lambda args: args['bg'] is not None, - 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None, -}) -@triton.autotune( - configs=[ - triton.Config({}, num_warps=num_warps, num_stages=num_stages) - for num_warps in [2, 4] - for num_stages in [2, 3, 4] - ], - key=['H', 'HV', 'K', 'V', 'BT'], - **autotune_cache_kwargs, -) -@triton.jit(do_not_specialize=['T']) -def pre_process_fwd_kernel_merged( - k, - v, - w, - g, - gk, - bg, - u, - hm, - cu_seqlens, - T, - H: tl.constexpr, - HV: tl.constexpr, - K: tl.constexpr, - V: tl.constexpr, - BT: tl.constexpr, - BLOCK_SIZE: tl.constexpr, - BK1: tl.constexpr, - USE_G: tl.constexpr, - USE_GK: tl.constexpr, - USE_BG: tl.constexpr, - IS_VARLEN: tl.constexpr, - MULTI_SEQS: tl.constexpr, -): - i_col, i_h = tl.program_id(0), tl.program_id(1) - if MULTI_SEQS: - i_n = tl.program_id(2) - # Offset hm for this subseq: hm[i_n, h, k, v+k] - hm += i_n * HV * K * (K + V) + i_h * K * (K + V) - else: - i_n = 0 - hm += i_h * K * (K + V) - if IS_VARLEN: - bos, eos = tl.load(cu_seqlens + i_n).to(tl.int64), tl.load(cu_seqlens + i_n + 1).to(tl.int64) - T = (eos - bos).to(tl.int32) - NT = tl.cdiv(T, BT) - else: - bos, eos = (i_n * T).to(tl.int64), (i_n * T + T).to(tl.int64) - NT = tl.cdiv(T, BT) - - # Determine if this block handles h (V part) or m (K part) - # i_col is in range [0, cdiv(V + K, BLOCK_SIZE)) - # Columns [0, V) are for h, columns [V, V+K) are for m - is_h_part = i_col * BLOCK_SIZE < V - # For DPLR (USE_BG), w and bg share the same head dim H as k/ag. - # For GDN/KDA, w has head dim HV (same as v). - k += ((bos * H + i_h // (HV // H)) * K).to(tl.int64) - if USE_BG: - w += ((bos * H + i_h // (HV // H)) * K).to(tl.int64) - bg += ((bos * H + i_h // (HV // H)) * K).to(tl.int64) - else: - w += ((bos * HV + i_h) * K).to(tl.int64) - if USE_G: - g += (bos * HV + i_h).to(tl.int64) - if USE_GK: - gk += ((bos * HV + i_h) * K).to(tl.int64) - stride_k = H * K - stride_w = H * K if USE_BG else HV * K - - if is_h_part: - # ====== Stage 1: Compute h (K x V) ====== - v += ((bos * HV + i_h) * V).to(tl.int64) - if USE_BG: - # DPLR keeps u and v as separate tensors; both need the per-head offset. - # For GDN/KDA, u is aliased to v at the Python wrapper level. - u += ((bos * HV + i_h) * V).to(tl.int64) - stride_v = HV * V - i_v = i_col - - # Initialize h accumulators - b_h1 = tl.zeros([64, BLOCK_SIZE], dtype=tl.float32) - if K > 64: - b_h2 = tl.zeros([64, BLOCK_SIZE], dtype=tl.float32) - if K > 128: - b_h3 = tl.zeros([64, BLOCK_SIZE], dtype=tl.float32) - if K > 192: - b_h4 = tl.zeros([64, BLOCK_SIZE], dtype=tl.float32) - - # Main recurrence for h - for i_t in range(NT): - # Compute decayed v - p_w = tl.make_block_ptr(w, (T, K), (stride_w, 1), (i_t * BT, 0), (BT, 64), (1, 0)) - b_w = tl.load(p_w, boundary_check=(0, 1)) - b_v_decay = tl.dot(b_w, b_h1.to(b_w.dtype)) - if K > 64: - p_w = tl.make_block_ptr(w, (T, K), (stride_w, 1), (i_t * BT, 64), (BT, 64), (1, 0)) - b_w = tl.load(p_w, boundary_check=(0, 1)) - b_v_decay += tl.dot(b_w, b_h2.to(b_w.dtype)) - if K > 128: - p_w = tl.make_block_ptr(w, (T, K), (stride_w, 1), (i_t * BT, 128), (BT, 64), (1, 0)) - b_w = tl.load(p_w, boundary_check=(0, 1)) - b_v_decay += tl.dot(b_w, b_h3.to(b_w.dtype)) - if K > 192: - p_w = tl.make_block_ptr(w, (T, K), (stride_w, 1), (i_t * BT, 192), (BT, 64), (1, 0)) - b_w = tl.load(p_w, boundary_check=(0, 1)) - b_v_decay += tl.dot(b_w, b_h4.to(b_w.dtype)) - - p_v = tl.make_block_ptr(v, (T, V), (stride_v, 1), (i_t * BT, i_v * BLOCK_SIZE), (BT, BLOCK_SIZE), (1, 0)) - if USE_BG: - # DPLR mode: v2 = w @ h + u, h += kg^T @ v + bg^T @ v2 - b_v_orig = tl.load(p_v, boundary_check=(0, 1)) - p_u = tl.make_block_ptr(u, (T, V), (stride_v, 1), (i_t * BT, i_v * BLOCK_SIZE), (BT, BLOCK_SIZE), (1, 0)) - b_v = b_v_decay + tl.load(p_u, boundary_check=(0, 1)) - else: - # GDN/KDA mode: v_new = v - w @ h - b_v = tl.load(p_v, boundary_check=(0, 1)) - b_v_decay - - last_idx = min((i_t + 1) * BT, T) - 1 - - # Apply g decay - if USE_G: - m_t = (i_t * BT + tl.arange(0, BT)) < T - b_g_last = tl.load(g + last_idx * HV).to(tl.float32) - p_g = tl.make_block_ptr(g, (T,), (HV,), (i_t * BT,), (BT,), (0,)) - b_g = tl.load(p_g, boundary_check=(0,)).to(tl.float32) - b_v = b_v * tl.where(m_t, exp2(b_g_last - b_g), 0)[:, None] - b_g_last = exp2(b_g_last) - b_h1 *= b_g_last - if K > 64: - b_h2 *= b_g_last - if K > 128: - b_h3 *= b_g_last - if K > 192: - b_h4 *= b_g_last - - # Apply gk decay - if USE_GK: - o_k1 = tl.arange(0, 64) - p_gk_last = gk + last_idx * HV * K - b_gk_last1 = tl.load(p_gk_last + o_k1, mask=(o_k1 < K), other=0.).to(tl.float32) - b_h1 *= exp2(b_gk_last1)[:, None] - if K > 64: - o_k2 = 64 + o_k1 - b_gk_last2 = tl.load(p_gk_last + o_k2, mask=(o_k2 < K), other=0.).to(tl.float32) - b_h2 *= exp2(b_gk_last2)[:, None] - if K > 128: - o_k3 = 128 + o_k1 - b_gk_last3 = tl.load(p_gk_last + o_k3, mask=(o_k3 < K), other=0.).to(tl.float32) - b_h3 *= exp2(b_gk_last3)[:, None] - if K > 192: - o_k4 = 192 + o_k1 - b_gk_last4 = tl.load(p_gk_last + o_k4, mask=(o_k4 < K), other=0.).to(tl.float32) - b_h4 *= exp2(b_gk_last4)[:, None] - b_v = b_v.to(k.dtype.element_ty) - - # Update h - p_k = tl.make_block_ptr(k, (K, T), (1, stride_k), (0, i_t * BT), (64, BT), (0, 1)) - b_k = tl.load(p_k, boundary_check=(0, 1)) - if USE_BG: - # DPLR mode: h += kg^T @ v + bg^T @ v2 - p_bg = tl.make_block_ptr(bg, (K, T), (1, stride_k), (0, i_t * BT), (64, BT), (0, 1)) - b_bg = tl.load(p_bg, boundary_check=(0, 1)) - b_h1 += tl.dot(b_k, b_v_orig.to(b_k.dtype)) + tl.dot(b_bg, b_v) - if K > 64: - p_k = tl.make_block_ptr(k, (K, T), (1, stride_k), (64, i_t * BT), (64, BT), (0, 1)) - b_k = tl.load(p_k, boundary_check=(0, 1)) - p_bg = tl.make_block_ptr(bg, (K, T), (1, stride_k), (64, i_t * BT), (64, BT), (0, 1)) - b_bg = tl.load(p_bg, boundary_check=(0, 1)) - b_h2 += tl.dot(b_k, b_v_orig.to(b_k.dtype)) + tl.dot(b_bg, b_v) - if K > 128: - p_k = tl.make_block_ptr(k, (K, T), (1, stride_k), (128, i_t * BT), (64, BT), (0, 1)) - b_k = tl.load(p_k, boundary_check=(0, 1)) - p_bg = tl.make_block_ptr(bg, (K, T), (1, stride_k), (128, i_t * BT), (64, BT), (0, 1)) - b_bg = tl.load(p_bg, boundary_check=(0, 1)) - b_h3 += tl.dot(b_k, b_v_orig.to(b_k.dtype)) + tl.dot(b_bg, b_v) - if K > 192: - p_k = tl.make_block_ptr(k, (K, T), (1, stride_k), (192, i_t * BT), (64, BT), (0, 1)) - b_k = tl.load(p_k, boundary_check=(0, 1)) - p_bg = tl.make_block_ptr(bg, (K, T), (1, stride_k), (192, i_t * BT), (64, BT), (0, 1)) - b_bg = tl.load(p_bg, boundary_check=(0, 1)) - b_h4 += tl.dot(b_k, b_v_orig.to(b_k.dtype)) + tl.dot(b_bg, b_v) - else: - # GDN/KDA mode: h += k^T @ v_new - b_h1 += tl.dot(b_k, b_v) - if K > 64: - p_k = tl.make_block_ptr(k, (K, T), (1, stride_k), (64, i_t * BT), (64, BT), (0, 1)) - b_k = tl.load(p_k, boundary_check=(0, 1)) - b_h2 += tl.dot(b_k, b_v) - if K > 128: - p_k = tl.make_block_ptr(k, (K, T), (1, stride_k), (128, i_t * BT), (64, BT), (0, 1)) - b_k = tl.load(p_k, boundary_check=(0, 1)) - b_h3 += tl.dot(b_k, b_v) - if K > 192: - p_k = tl.make_block_ptr(k, (K, T), (1, stride_k), (192, i_t * BT), (64, BT), (0, 1)) - b_k = tl.load(p_k, boundary_check=(0, 1)) - b_h4 += tl.dot(b_k, b_v) - - # Store h results - stride_hm_kv = K + V - p_h1 = tl.make_block_ptr(hm, (K, V), (stride_hm_kv, 1), (0, i_v * BLOCK_SIZE), (64, BLOCK_SIZE), (1, 0)) - tl.store(p_h1, b_h1.to(p_h1.dtype.element_ty), boundary_check=(0, 1)) - if K > 64: - p_h2 = tl.make_block_ptr(hm, (K, V), (stride_hm_kv, 1), (64, i_v * BLOCK_SIZE), (64, BLOCK_SIZE), (1, 0)) - tl.store(p_h2, b_h2.to(p_h2.dtype.element_ty), boundary_check=(0, 1)) - if K > 128: - p_h3 = tl.make_block_ptr(hm, (K, V), (stride_hm_kv, 1), (128, i_v * BLOCK_SIZE), (64, BLOCK_SIZE), (1, 0)) - tl.store(p_h3, b_h3.to(p_h3.dtype.element_ty), boundary_check=(0, 1)) - if K > 192: - p_h4 = tl.make_block_ptr(hm, (K, V), (stride_hm_kv, 1), (192, i_v * BLOCK_SIZE), (64, BLOCK_SIZE), (1, 0)) - tl.store(p_h4, b_h4.to(p_h4.dtype.element_ty), boundary_check=(0, 1)) - else: - # ====== Stage 2: Compute m (K x K) ====== - # i_col is for m part, map to K dimension - # m starts at column V, so offset = i_col * BLOCK_SIZE - V - # Use tl.cdiv to correctly compute the number of blocks for V dimension - i_k_col = i_col - tl.cdiv(V, BLOCK_SIZE) - - # Following stage2 kernel design: - # - BK1 is the full K dimension (next_power_of_2(K)) - # - BLOCK_SIZE is the column block size (like BK2=32 in stage2) - # Each block computes a (BK1, BLOCK_SIZE) sub-matrix of m - row = tl.arange(0, BK1) - col = tl.arange(0, BLOCK_SIZE) + i_k_col * BLOCK_SIZE - - # Initialize as identity matrix: M_0 = I - b_m = tl.where(row[:, None] == col[None, :], 1.0, 0.0) - - for i_t in range(NT): - # Load k and w with full BK1 rows - if USE_BG: - # DPLR mode: use bg for transition matrix - # bg was already offset at the beginning of the kernel - p_k = tl.make_block_ptr(bg, (T, K), (stride_k, 1), (i_t * BT, 0), (BT, BK1), (1, 0)) - else: - p_k = tl.make_block_ptr(k, (T, K), (stride_k, 1), (i_t * BT, 0), (BT, BK1), (1, 0)) - b_k = tl.load(p_k, boundary_check=(0, 1)) - p_w = tl.make_block_ptr(w, (T, K), (stride_w, 1), (i_t * BT, 0), (BT, BK1), (1, 0)) - b_w = tl.load(p_w, boundary_check=(0, 1)) - - last_idx = min((i_t + 1) * BT, T) - 1 - - if USE_G: - m_t = (i_t * BT + tl.arange(0, BT)) < T - b_g_last = tl.load(g + last_idx * HV).to(tl.float32) - p_g = tl.make_block_ptr(g, (T,), (HV,), (i_t * BT,), (BT,), (0,)) - b_g = tl.load(p_g, boundary_check=(0,)).to(tl.float32) - b_k = b_k * tl.where(m_t, exp2(b_g_last - b_g), 0)[:, None] - b_g_last = exp2(b_g_last) - b_diag = tl.where(row[:, None] == row[None, :], b_g_last, 0.0) - elif USE_GK: - b_gk_last = tl.load(gk + last_idx * HV * K + row, mask=(row < K), other=0.).to(tl.float32) - b_gk_last = exp2(b_gk_last) - b_diag = tl.where(row[:, None] == row[None, :], b_gk_last[:, None], 0.0) - else: - b_diag = tl.where(row[:, None] == row[None, :], 1.0, 0.0) - - # Compute m update - if USE_BG: - # DPLR mode: M = (diag + bg^T @ w) @ M - # bg was already offset at the beginning of the kernel - b_kw = tl.dot(tl.trans(b_k.to(b_w.dtype)), b_w) - b_m_i = b_diag + b_kw - else: - # GDN/KDA mode: M = (diag - k^T @ w) @ M - b_kw = tl.dot(tl.trans(b_k.to(b_w.dtype)), b_w) - b_m_i = b_diag - b_kw - b_m = tl.dot(b_m_i.to(tl.float32), b_m.to(tl.float32)) - - # Store m result - stride_hm_kv = K + V - p_m = tl.make_block_ptr(hm + V, (K, K), (stride_hm_kv, 1), (0, i_k_col * BLOCK_SIZE), (BK1, BLOCK_SIZE), (1, 0)) - tl.store(p_m, b_m.to(p_m.dtype.element_ty), boundary_check=(0, 1)) - - -@triton.heuristics({ - 'HAS_H0': lambda args: args['h0'] is not None, -}) -@triton.autotune( - configs=[ - triton.Config({'BV': BV}, num_warps=num_warps, num_stages=num_stages) - for num_warps in [2, 4] - for num_stages in [2, 3, 4] - for BV in [32, 64] - ], - key=['HV', 'K', 'V', 'BT'], - **autotune_cache_kwargs, -) -@triton.jit(do_not_specialize=['pre_or_post_num_ranks', 'rank', 'NUM_SEQ_ENTRIES']) -def merge_fwd_bwd_kernel( - h, # [HV, K, V] or [num_non_first, HV, K, V] for intracard (or [V, K] when transposed) - ag_hm, # [HV, K, K+V] or [S_split, HV, K, K+V] for intracard (always [K, V+K]) - pre_or_post_num_ranks, # num_ranks for CP, NUM_SPLIT_SEQS for intracard - rank, # rank for CP, not used for intracard - seq_offsets, # None for CP, [num_split_seqs+1] for intracard - init_offsets, # None for CP, [num_split_seqs+1] for intracard - h0_seq_ids, # None for CP, [num_split_seqs] for intracard - h0, # None or [N_orig, HV, K, V] for intracard (or [V, K] when transposed) - HV: tl.constexpr, - K: tl.constexpr, - V: tl.constexpr, - BV: tl.constexpr, - BK: tl.constexpr, - FORWARD: tl.constexpr, # True for FWD, False for BWD - INTRACARD_MODE: tl.constexpr, # True: intracard mode, False: CP mode - NUM_SEQ_ENTRIES, # num_split_seqs for intracard - HAS_H0: tl.constexpr, # Heuristic: whether h0 is provided - STATE_V_FIRST: tl.constexpr = False, # When True, h0/h use [V, K] layout; ag_hm always [K, V+K] -): - """ - Unified merge kernel for both CP and Intra-card modes. - - CP mode (INTRACARD_MODE=False): - Grid: (V/BV, HV) - Merges across ranks for context parallel. - - Intra-card mode (INTRACARD_MODE=True): - Grid: (V/BV, NUM_SEQ_ENTRIES, HV) - Merges across subseqs within card for intra-card context parallel. - - When STATE_V_FIRST=True, h0 and output h use [V, K] layout. - ag_hm always uses [K, V+K] layout (from pre_scan). - The recurrence h' = M @ h + he becomes h_T' = h_T @ M^T + he^T. - """ - i_v = tl.program_id(0) - if INTRACARD_MODE: - i_seq = tl.program_id(1) - i_h = tl.program_id(2) - - if i_seq >= NUM_SEQ_ENTRIES: - return - - # Load offsets for this sequence - ss_start = tl.load(seq_offsets + i_seq).to(tl.int32) - ss_end = tl.load(seq_offsets + i_seq + 1).to(tl.int32) - init_base = tl.load(init_offsets + i_seq).to(tl.int32) - num_subseqs = ss_end - ss_start - - stride_hm_s = HV * K * (V + K) - stride_hm_h = K * (V + K) - - # Initialize from h0 if provided - if HAS_H0: - orig_seq_id = tl.load(h0_seq_ids + i_seq).to(tl.int32) - if STATE_V_FIRST: - p_h0 = tl.make_block_ptr( - h0 + (orig_seq_id * HV + i_h) * V * K, - (V, K), (K, 1), (i_v * BV, 0), (BV, BK), (1, 0) - ) - b_h = tl.load(p_h0, boundary_check=(0, 1)).to(tl.float32) - else: - p_h0 = tl.make_block_ptr( - h0 + (orig_seq_id * HV + i_h) * K * V, - (K, V), (V, 1), (0, i_v * BV), (BK, BV), (1, 0) - ) - b_h = tl.load(p_h0, boundary_check=(0, 1)).to(tl.float32) - else: - if STATE_V_FIRST: - b_h = tl.zeros([BV, BK], dtype=tl.float32) - else: - b_h = tl.zeros([BK, BV], dtype=tl.float32) - - # Merge loop over subseqs - for idx in range(num_subseqs): - i_ss = ss_start + idx - base = i_ss * stride_hm_s + i_h * stride_hm_h - - # he and m are always in [K, V+K] layout from pre_scan - p_he = tl.make_block_ptr( - ag_hm + base, (K, V), (V + K, 1), (0, i_v * BV), (BK, BV), (1, 0) - ) - b_he = tl.load(p_he, boundary_check=(0, 1)).to(tl.float32) - p_m = tl.make_block_ptr( - ag_hm + base + V, (K, K), (V + K, 1), (0, 0), (BK, BK), (1, 0) - ) - b_m = tl.load(p_m, boundary_check=(0, 1)).to(tl.float32) - if STATE_V_FIRST: - # h_T' = h_T @ M^T + he^T - b_h = tl.dot(b_h.to(tl.float32), tl.trans(b_m)) + tl.trans(b_he) - else: - b_h = tl.dot(b_m.to(tl.float32), b_h.to(tl.float32)) + b_he.to(tl.float32) - - # Store for non-first subseqs - if idx < num_subseqs - 1: - init_idx = init_base + idx - stride_init = HV * K * V - if STATE_V_FIRST: - p_out = tl.make_block_ptr( - h + init_idx * stride_init + i_h * V * K, - (V, K), (K, 1), (i_v * BV, 0), (BV, BK), (1, 0) - ) - else: - p_out = tl.make_block_ptr( - h + init_idx * stride_init + i_h * K * V, - (K, V), (V, 1), (0, i_v * BV), (BK, BV), (1, 0) - ) - tl.store(p_out, b_h.to(p_out.dtype.element_ty), boundary_check=(0, 1)) - else: - # CP mode - i_h = tl.program_id(1) - num_ranks = pre_or_post_num_ranks.to(tl.int32) - h += i_h * K * V - ag_hm += i_h * K * (K + V) - stride = HV * K * (K + V) - if STATE_V_FIRST: - b_h = tl.zeros([BV, BK], dtype=tl.float32) - else: - b_h = tl.zeros([BK, BV], dtype=tl.float32) - for idx in range(num_ranks): - if FORWARD: - cur_rank = rank - num_ranks + idx - else: - cur_rank = rank + num_ranks - idx - p_ag_h = tl.make_block_ptr(ag_hm + cur_rank * stride, (K, V), (K + V, 1), (0, i_v * BV), (BK, BV), (1, 0)) - b_ag_h = tl.load(p_ag_h, boundary_check=(0, 1)) - p_ag_m = tl.make_block_ptr(ag_hm + cur_rank * stride + V, (K, K), (K + V, 1), (0, 0), (BK, BK), (1, 0)) - b_ag_m = tl.load(p_ag_m, boundary_check=(0, 1)) - if STATE_V_FIRST: - b_h = tl.dot(b_h.to(tl.float32), tl.trans(b_ag_m).to(tl.float32)) + tl.trans(b_ag_h).to(tl.float32) - else: - b_h = tl.dot(b_ag_m.to(tl.float32), b_h.to(tl.float32)) + b_ag_h.to(tl.float32) - if STATE_V_FIRST: - p_h = tl.make_block_ptr(h, (V, K), (K, 1), (i_v * BV, 0), (BV, BK), (1, 0)) - else: - p_h = tl.make_block_ptr(h, (K, V), (V, 1), (0, i_v * BV), (BK, BV), (1, 0)) - tl.store(p_h, b_h.to(p_h.dtype.element_ty), boundary_check=(0, 1)) - - -@triton.heuristics({ - 'USE_G': lambda args: args['g'] is not None, - 'USE_GK': lambda args: args['gk'] is not None, - 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None, -}) -@triton.autotune( - configs=[ - triton.Config({}, num_warps=num_warps, num_stages=num_stages) - for num_warps in [2, 4] - for num_stages in ([4, 3, 2] if check_shared_mem('ampere') else [1]) - ], - key=['H', 'HV', 'K', 'V', 'BT'], - **autotune_cache_kwargs, -) -@triton.jit(do_not_specialize=['T']) -def pre_process_bwd_kernel_merged( - q, - k, - w, - g, - gk, - do, - dhm, - dv, - cu_seqlens, - scale, - T, - H: tl.constexpr, - HV: tl.constexpr, - K: tl.constexpr, - V: tl.constexpr, - BT: tl.constexpr, - BLOCK_SIZE: tl.constexpr, - BK1: tl.constexpr, - USE_G: tl.constexpr, - USE_GK: tl.constexpr, - USE_BG: tl.constexpr, - IS_VARLEN: tl.constexpr, -): - """ - Merged backward kernel that computes both dh (K x V) and dm (K x K) in a single kernel. - - Similar to pre_process_fwd_kernel_merged, this kernel uses a unified grid where: - - Columns [0, V) are for computing dh (stage 1) - - Columns [V, V+K) are for computing dm (stage 2) - """ - i_col, i_h = tl.program_id(0), tl.program_id(1) - i_n = 0 - if IS_VARLEN: - bos, eos = tl.load(cu_seqlens + i_n).to(tl.int64), tl.load(cu_seqlens + i_n + 1).to(tl.int64) - T = (eos - bos).to(tl.int32) - NT = tl.cdiv(T, BT) - else: - bos, eos = (i_n * T).to(tl.int64), (i_n * T + T).to(tl.int64) - NT = tl.cdiv(T, BT) - - # Determine if this block handles dh (V part) or dm (K part) - is_dh_part = i_col * BLOCK_SIZE < V - - # Calculate offsets - # For DPLR (USE_BG), w shares the same head dim H as q/k/ag. - # For GDN/KDA, w has head dim HV (same as do/dv). - q += ((bos * H + i_h // (HV // H)) * K).to(tl.int64) - k += ((bos * H + i_h // (HV // H)) * K).to(tl.int64) - if USE_BG: - w += ((bos * H + i_h // (HV // H)) * K).to(tl.int64) - else: - w += ((bos * HV + i_h) * K).to(tl.int64) - if USE_G: - g += (bos * HV + i_h).to(tl.int64) - if USE_GK: - gk += ((bos * HV + i_h) * K).to(tl.int64) - dhm += i_h * K * (V + K) - stride_qk = H * K - stride_w = H * K if USE_BG else HV * K - - if is_dh_part: - # ====== Stage 1: Compute dh (K x V) ====== - do += ((bos * HV + i_h) * V).to(tl.int64) - dv += ((bos * HV + i_h) * V).to(tl.int64) - stride_v = HV * V - i_v = i_col - - # Initialize dh accumulators - b_dh1 = tl.zeros([64, BLOCK_SIZE], dtype=tl.float32) - if K > 64: - b_dh2 = tl.zeros([64, BLOCK_SIZE], dtype=tl.float32) - if K > 128: - b_dh3 = tl.zeros([64, BLOCK_SIZE], dtype=tl.float32) - if K > 192: - b_dh4 = tl.zeros([64, BLOCK_SIZE], dtype=tl.float32) - - # Main recurrence for dh (reverse order) - for i_t in range(NT - 1, -1, -1): - last_idx = min((i_t + 1) * BT, T) - 1 - - if USE_G: - bg_last = tl.load(g + last_idx * HV).to(tl.float32) - p_g = tl.make_block_ptr(g, (T,), (HV,), (i_t * BT,), (BT,), (0,)) - b_g = tl.load(p_g, boundary_check=(0,)).to(tl.float32) - bg_last_exp = exp2(bg_last) - b_g_exp = exp2(b_g) - - p_dv = tl.make_block_ptr(dv, (T, V), (stride_v, 1), (i_t * BT, i_v * BLOCK_SIZE), (BT, BLOCK_SIZE), (1, 0)) - p_do = tl.make_block_ptr(do, (T, V), (stride_v, 1), (i_t * BT, i_v * BLOCK_SIZE), (BT, BLOCK_SIZE), (1, 0)) - b_do = tl.load(p_do, boundary_check=(0, 1)) - - # Update dv - p_k = tl.make_block_ptr(k, (T, K), (stride_qk, 1), (i_t * BT, 0), (BT, 64), (1, 0)) - b_k = tl.load(p_k, boundary_check=(0, 1)) - if USE_GK: - o_k1 = tl.arange(0, 64) - p_gk_last = gk + last_idx * HV * K - b_gk_last1 = tl.load(p_gk_last + o_k1, mask=(o_k1 < K), other=0.).to(tl.float32) - b_dv = tl.dot(b_k, b_dh1.to(b_k.dtype)) - - if K > 64: - p_k = tl.make_block_ptr(k, (T, K), (stride_qk, 1), (i_t * BT, 64), (BT, 64), (1, 0)) - b_k = tl.load(p_k, boundary_check=(0, 1)) - if USE_GK: - o_k2 = 64 + o_k1 - b_gk_last2 = tl.load(p_gk_last + o_k2, mask=(o_k2 < K), other=0.).to(tl.float32) - b_dv += tl.dot(b_k, b_dh2.to(b_k.dtype)) - - if K > 128: - p_k = tl.make_block_ptr(k, (T, K), (stride_qk, 1), (i_t * BT, 128), (BT, 64), (1, 0)) - b_k = tl.load(p_k, boundary_check=(0, 1)) - if USE_GK: - o_k3 = 128 + o_k1 - b_gk_last3 = tl.load(p_gk_last + o_k3, mask=(o_k3 < K), other=0.).to(tl.float32) - b_dv += tl.dot(b_k, b_dh3.to(b_k.dtype)) - - if K > 192: - p_k = tl.make_block_ptr(k, (T, K), (stride_qk, 1), (i_t * BT, 192), (BT, 64), (1, 0)) - b_k = tl.load(p_k, boundary_check=(0, 1)) - if USE_GK: - o_k4 = 192 + o_k1 - b_gk_last4 = tl.load(p_gk_last + o_k4, mask=(o_k4 < K), other=0.).to(tl.float32) - b_dv += tl.dot(b_k, b_dh4.to(b_k.dtype)) - - if USE_G: - m_t = (i_t * BT + tl.arange(0, BT)) < T - b_dv *= tl.where(m_t, exp2(bg_last - b_g), 0)[:, None] - b_dv += tl.load(p_dv, boundary_check=(0, 1)) - - # Update dh - p_w = tl.make_block_ptr(w, (K, T), (1, stride_w), (0, i_t * BT), (64, BT), (0, 1)) - p_q = tl.make_block_ptr(q, (K, T), (1, stride_qk), (0, i_t * BT), (64, BT), (0, 1)) - b_w = tl.load(p_w, boundary_check=(0, 1)) - b_q = tl.load(p_q, boundary_check=(0, 1)) - if USE_G: - b_dh1 *= bg_last_exp - b_q = b_q * b_g_exp[None, :] - if USE_GK: - b_dh1 *= exp2(b_gk_last1[:, None]) - if USE_BG: - # DPLR mode: dh += q^T @ do + w^T @ dv2 (no scale) - b_dh1 += tl.dot(b_q.to(b_q.dtype), b_do.to(b_q.dtype)) + tl.dot(b_w, b_dv.to(b_w.dtype)) - else: - b_dh1 += tl.dot(b_q.to(b_q.dtype), b_do.to(b_q.dtype)) * scale - tl.dot(b_w, b_dv.to(b_w.dtype)) - - if K > 64: - p_q = tl.make_block_ptr(q, (K, T), (1, stride_qk), (64, i_t * BT), (64, BT), (0, 1)) - p_w = tl.make_block_ptr(w, (K, T), (1, stride_w), (64, i_t * BT), (64, BT), (0, 1)) - b_q = tl.load(p_q, boundary_check=(0, 1)) - b_w = tl.load(p_w, boundary_check=(0, 1)) - if USE_G: - b_dh2 *= bg_last_exp - b_q = b_q * b_g_exp[None, :] - if USE_GK: - b_dh2 *= exp2(b_gk_last2[:, None]) - if USE_BG: - b_dh2 += tl.dot(b_q.to(b_q.dtype), b_do.to(b_q.dtype)) + tl.dot(b_w, b_dv.to(b_w.dtype)) - else: - b_dh2 += tl.dot(b_q.to(b_q.dtype), b_do.to(b_q.dtype)) * scale - tl.dot(b_w, b_dv.to(b_w.dtype)) - - if K > 128: - p_q = tl.make_block_ptr(q, (K, T), (1, stride_qk), (128, i_t * BT), (64, BT), (0, 1)) - p_w = tl.make_block_ptr(w, (K, T), (1, stride_w), (128, i_t * BT), (64, BT), (0, 1)) - b_q = tl.load(p_q, boundary_check=(0, 1)) - b_w = tl.load(p_w, boundary_check=(0, 1)) - if USE_G: - b_dh3 *= bg_last_exp - b_q = b_q * b_g_exp[None, :] - if USE_GK: - b_dh3 *= exp2(b_gk_last3[:, None]) - if USE_BG: - b_dh3 += tl.dot(b_q.to(b_q.dtype), b_do.to(b_q.dtype)) + tl.dot(b_w, b_dv.to(b_w.dtype)) - else: - b_dh3 += tl.dot(b_q.to(b_q.dtype), b_do.to(b_q.dtype)) * scale - tl.dot(b_w, b_dv.to(b_w.dtype)) - - if K > 192: - p_q = tl.make_block_ptr(q, (K, T), (1, stride_qk), (192, i_t * BT), (64, BT), (0, 1)) - p_w = tl.make_block_ptr(w, (K, T), (1, stride_w), (192, i_t * BT), (64, BT), (0, 1)) - b_q = tl.load(p_q, boundary_check=(0, 1)) - b_w = tl.load(p_w, boundary_check=(0, 1)) - if USE_G: - b_dh4 *= bg_last_exp - b_q = b_q * b_g_exp[None, :] - if USE_GK: - b_dh4 *= exp2(b_gk_last4[:, None]) - if USE_BG: - b_dh4 += tl.dot(b_q.to(b_q.dtype), b_do.to(b_q.dtype)) + tl.dot(b_w, b_dv.to(b_w.dtype)) - else: - b_dh4 += tl.dot(b_q.to(b_q.dtype), b_do.to(b_q.dtype)) * scale - tl.dot(b_w, b_dv.to(b_w.dtype)) - - # Store dh results - p_dh1 = tl.make_block_ptr(dhm, (K, V), (V + K, 1), (0, i_v * BLOCK_SIZE), (64, BLOCK_SIZE), (1, 0)) - tl.store(p_dh1, b_dh1.to(p_dh1.dtype.element_ty), boundary_check=(0, 1)) - if K > 64: - p_dh2 = tl.make_block_ptr(dhm, (K, V), (V + K, 1), (64, i_v * BLOCK_SIZE), (64, BLOCK_SIZE), (1, 0)) - tl.store(p_dh2, b_dh2.to(p_dh2.dtype.element_ty), boundary_check=(0, 1)) - if K > 128: - p_dh3 = tl.make_block_ptr(dhm, (K, V), (V + K, 1), (128, i_v * BLOCK_SIZE), (64, BLOCK_SIZE), (1, 0)) - tl.store(p_dh3, b_dh3.to(p_dh3.dtype.element_ty), boundary_check=(0, 1)) - if K > 192: - p_dh4 = tl.make_block_ptr(dhm, (K, V), (V + K, 1), (192, i_v * BLOCK_SIZE), (64, BLOCK_SIZE), (1, 0)) - tl.store(p_dh4, b_dh4.to(p_dh4.dtype.element_ty), boundary_check=(0, 1)) - else: - # ====== Stage 2: Compute dm (K x K) ====== - # i_col is for dm part, map to K dimension - i_k_col = i_col - tl.cdiv(V, BLOCK_SIZE) - - # Following stage2 kernel design for backward (FORWARD=False) - # - BK1 is the full K dimension (next_power_of_2(K)) - # - BLOCK_SIZE is the column block size - row = tl.arange(0, BK1) - col = tl.arange(0, BLOCK_SIZE) + i_k_col * BLOCK_SIZE - - # Initialize as identity matrix: M_0 = I - b_m = tl.where(row[:, None] == col[None, :], 1.0, 0.0) - - for _i_t in range(NT): - # Reverse order for backward - i_t = NT - 1 - _i_t - - # Load k and w with full BK1 rows - p_k = tl.make_block_ptr(k, (T, K), (stride_qk, 1), (i_t * BT, 0), (BT, BK1), (1, 0)) - b_k = tl.load(p_k, boundary_check=(0, 1)) - p_w = tl.make_block_ptr(w, (T, K), (stride_w, 1), (i_t * BT, 0), (BT, BK1), (1, 0)) - b_w = tl.load(p_w, boundary_check=(0, 1)) - - last_idx = min((i_t + 1) * BT, T) - 1 - - if USE_G: - m_t = (i_t * BT + tl.arange(0, BT)) < T - b_g_last = tl.load(g + last_idx * HV).to(tl.float32) - p_g = tl.make_block_ptr(g, (T,), (HV,), (i_t * BT,), (BT,), (0,)) - b_g = tl.load(p_g, boundary_check=(0,)).to(tl.float32) - b_k = b_k * tl.where(m_t, exp2(b_g_last - b_g), 0)[:, None] - b_g_last = exp2(b_g_last) - b_diag = tl.where(row[:, None] == row[None, :], b_g_last, 0.0) - elif USE_GK: - b_gk_last = tl.load(gk + last_idx * HV * K + row, mask=(row < K), other=0.).to(tl.float32) - b_gk_last = exp2(b_gk_last) - b_diag = tl.where(row[:, None] == row[None, :], b_gk_last[:, None], 0.0) - else: - b_diag = tl.where(row[:, None] == row[None, :], 1.0, 0.0) - - # Compute dm update for backward - b_kw = tl.dot(tl.trans(b_w), b_k.to(b_w.dtype)) - if USE_BG: - # DPLR mode: dM = (diag + w^T @ bg) @ dM - b_m_i = b_diag + b_kw - else: - # GDN/KDA mode: dM = (diag - w^T @ k) @ dM - b_m_i = b_diag - b_kw - # Keep m chain in fp32 to avoid precision loss from repeated bf16 casting - b_m = tl.dot(b_m_i.to(tl.float32), b_m.to(tl.float32)) - - # Store dm result - p_m = tl.make_block_ptr(dhm + V, (K, K), (V + K, 1), (0, i_k_col * BLOCK_SIZE), (BK1, BLOCK_SIZE), (1, 0)) - tl.store(p_m, b_m.to(p_m.dtype.element_ty), boundary_check=(0, 1)) - - -def chunk_gated_delta_rule_fwd_h_pre_process( - k: torch.Tensor, - w: torch.Tensor, - u: torch.Tensor, - g: torch.Tensor | None = None, - gk: torch.Tensor | None = None, - bg: torch.Tensor | None = None, - v: torch.Tensor | None = None, - chunk_size: int = 64, # SY: remove this argument and force chunk size 64? - state_v_first: bool = False, - cu_seqlens: torch.LongTensor | None = None, - initial_state: torch.Tensor | None = None, - context: FLACPContext = None, -) -> tuple[torch.Tensor, torch.Tensor]: - if context is None or context.group is None: - return initial_state - assert initial_state is None, "When enable CP, the provided initial_state must be None." - rank = dist.get_rank(group=context.group) - - B, T, H, K, V, HV = *k.shape, u.shape[-1], u.shape[2] - BT = chunk_size - BK = triton.next_power_of_2(K) - - # N: the actual number of sequences in the batch with either equal or variable lengths - if cu_seqlens is None: - N = B - else: - N = len(cu_seqlens) - 1 - assert K <= 256, "current kernel does not support head dimension larger than 256." - - hm = k.new_zeros(HV, K, (V + K), dtype=torch.float32) - if state_v_first: - initial_state = k.new_zeros(N, HV, V, K, dtype=torch.float32) - else: - initial_state = k.new_zeros(N, HV, K, V, dtype=torch.float32) - if not context.is_last_rank: - BLOCK_SIZE = 32 if K <= 64 else 64 - grid = (triton.cdiv(V, BLOCK_SIZE) + triton.cdiv(K, BLOCK_SIZE), HV) - # For DPLR, v provides the original v for computing h contributions, - # while u remains the WY-processed values (A_ab @ A_ak @ v) for v_new = w @ h + u. - pre_process_fwd_kernel_merged[grid]( - k=k, - v=u if v is None else v, - w=w, - g=g, - gk=gk, - bg=bg, - u=u, - hm=hm, - cu_seqlens=cu_seqlens[-2:], - T=T, - H=H, - HV=HV, - K=K, - V=V, - BT=BT, - BK1=BK, - BLOCK_SIZE=BLOCK_SIZE, - MULTI_SEQS=False, - ) - ag_hm, _ = all_gather_into_tensor(hm, group=context.group) - if not context.is_first_rank: - def grid(meta): return (triton.cdiv(V, meta['BV']), HV) - merge_fwd_bwd_kernel[grid]( - h=initial_state[0], - ag_hm=ag_hm, - pre_or_post_num_ranks=context.pre_num_ranks, - rank=rank, - seq_offsets=None, - init_offsets=None, - h0_seq_ids=None, - h0=None, - HV=HV, - K=K, - V=V, - BK=BK, - FORWARD=True, - INTRACARD_MODE=False, - NUM_SEQ_ENTRIES=0, - STATE_V_FIRST=state_v_first, - ) - return initial_state - - -def chunk_gated_delta_rule_bwd_dhu_pre_process( - q: torch.Tensor, - k: torch.Tensor, - w: torch.Tensor, - do: torch.Tensor, - dv: torch.Tensor, - g: torch.Tensor | None = None, - gk: torch.Tensor | None = None, - bg: torch.Tensor | None = None, - scale: float | None = None, - state_v_first: bool = False, - cu_seqlens: torch.LongTensor | None = None, - dht: torch.Tensor | None = None, - initial_state: torch.Tensor | None = None, - context: FLACPContext | None = None, - chunk_size: int = 64, -) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: - if context is None or context.group is None: - return dht, initial_state - assert dht is None, "When enable CP, the provided dht must be None." - rank = dist.get_rank(context.group) - - B, T, H, K, V, HV = *q.shape, do.shape[-1], do.shape[2] - # N: the actual number of sequences in the batch with either equal or variable lengths - BT = chunk_size - assert K <= 256, "current kernel does not support head dimension being larger than 256." - BK = triton.next_power_of_2(K) - - if cu_seqlens is None: - N = B - else: - N = len(cu_seqlens) - 1 - - dhm = q.new_zeros(HV, K, V + K, dtype=torch.float32) - if state_v_first: - dht = q.new_zeros(N, HV, V, K, dtype=torch.float32) - else: - dht = q.new_zeros(N, HV, K, V, dtype=torch.float32) - - if not context.is_first_rank: - BLOCK_SIZE = 32 if K <= 64 else 64 - grid = (triton.cdiv(V, BLOCK_SIZE) + triton.cdiv(K, BLOCK_SIZE), HV) - pre_process_bwd_kernel_merged[grid]( - q=q, - k=k if bg is None else bg, - w=w, - g=g, - gk=gk, - do=do, - dhm=dhm, - dv=dv, - cu_seqlens=cu_seqlens[:2], - scale=scale, - T=T, - H=H, - HV=HV, - K=K, - V=V, - BT=BT, - BK1=BK, - BLOCK_SIZE=BLOCK_SIZE, - USE_BG=bg is not None, - ) - - ag_dhm, _ = all_gather_into_tensor(dhm, group=context.group) - - if not context.is_last_rank: - def grid(meta): return (triton.cdiv(V, meta['BV']), HV) - merge_fwd_bwd_kernel[grid]( - h=dht[-1], - ag_hm=ag_dhm, - pre_or_post_num_ranks=context.post_num_ranks, - rank=rank, - seq_offsets=None, - init_offsets=None, - h0_seq_ids=None, - h0=None, - HV=HV, - K=K, - V=V, - BK=BK, - FORWARD=False, - INTRACARD_MODE=False, - NUM_SEQ_ENTRIES=0, - STATE_V_FIRST=state_v_first, - ) - - # initial_state is None in the CP mode - # We only need to compute dht of current rank and pass it to the backward kernel - return dht, None - - -def compress_h0(h0: torch.Tensor, context: FLACPContext): - if h0 is None or len(context.cu_seqlens) == 2: - return h0 - # Here must use clone op or the full tensor will be saved for backward - return h0[:1].clone() - - -def expand_h0(h0: torch.Tensor, context: FLACPContext): - if h0 is None or len(context.cu_seqlens) == 2: - return h0 - B = len(context.cu_seqlens) - 1 - expand_h0 = h0.new_zeros(B, *h0.shape[1:]) - expand_h0[:1] = h0 - return expand_h0 diff --git a/build/torch-xpu/ops/cp/comm.py b/build/torch-xpu/ops/cp/comm.py deleted file mode 100644 index 89a1cac0006862f840eb9527f83fc7f997130932..0000000000000000000000000000000000000000 --- a/build/torch-xpu/ops/cp/comm.py +++ /dev/null @@ -1,169 +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 - -from typing import TYPE_CHECKING - -import torch -import torch.distributed as dist - -if TYPE_CHECKING: - from torch.distributed import ProcessGroup - - -def all_gather_into_tensor( - inp: torch.Tensor, - out: torch.Tensor | None = None, - group: ProcessGroup | None = None, - async_op: bool = False -) -> tuple[torch.Tensor, dist.Work | None]: - """ - All-gather a tensor across ranks. - - Args: - inp: Input tensor to gather - out: Optional output tensor of shape [world_size, *inp.shape] - group: Process group - async_op: Whether to perform async operation - - Returns: - Tuple of (output tensor, handle if async_op else None) - """ - world_size = dist.get_world_size(group=group) - if out is None: - out = torch.empty(world_size, *inp.shape, device=inp.device, dtype=inp.dtype) - handle = dist.all_gather_into_tensor(out, inp, group=group, async_op=async_op) - return out, handle - - -def all_reduce_sum( - inp: torch.Tensor, - group: ProcessGroup | None = None, - async_op: bool = False -) -> tuple[torch.Tensor, dist.Work | None]: - """ - All-reduce sum a tensor across ranks. - - Args: - inp: Input tensor to reduce (modified in-place) - group: Process group - async_op: Whether to perform async operation - - Returns: - Tuple of (reduced tensor, handle if async_op else None) - """ - handle = dist.all_reduce(inp, op=dist.ReduceOp.SUM, group=group, async_op=async_op) - return inp, handle - - -def send_recv_fwd( - send_tensor: torch.Tensor, - group: ProcessGroup, - recv_from_prev: bool = True -) -> torch.Tensor: - """ - Forward pass communication: send tensor to next rank, receive from previous rank. - - Uses all_gather for simplicity and to ensure all ranks participate. - - Args: - send_tensor: Tensor to send (e.g., tails for conv1d) - group: Process group - recv_from_prev: If True, receive from previous rank; if False, receive from next rank - - Returns: - Received tensor from the specified rank (zeros if no valid source) - """ - rank = dist.get_rank(group) - world_size = dist.get_world_size(group) - - # All-gather to ensure all ranks participate - gathered, _ = all_gather_into_tensor(send_tensor, group=group, async_op=False) - - if recv_from_prev: - # Receive from previous rank - if rank == 0: - return torch.zeros_like(send_tensor) - else: - return gathered[rank - 1].clone() - else: - # Receive from next rank - if rank == world_size - 1: - return torch.zeros_like(send_tensor) - else: - return gathered[rank + 1].clone() - - -def send_recv_bwd( - send_tensor: torch.Tensor, - group: ProcessGroup, - recv_from_next: bool = True -) -> torch.Tensor: - """ - Backward pass communication: send gradient to previous rank, receive from next rank. - - Uses all_gather for simplicity and to ensure all ranks participate. - - Args: - send_tensor: Gradient tensor to send - group: Process group - recv_from_next: If True, receive from next rank; if False, receive from previous rank - - Returns: - Received gradient tensor from the specified rank (zeros if no valid source) - """ - rank = dist.get_rank(group) - world_size = dist.get_world_size(group) - - # All-gather to ensure all ranks participate - gathered, _ = all_gather_into_tensor(send_tensor, group=group, async_op=False) - - if recv_from_next: - # Receive from next rank - if rank == world_size - 1: - return torch.zeros_like(send_tensor) - else: - return gathered[rank + 1].clone() - else: - # Receive from previous rank - if rank == 0: - return torch.zeros_like(send_tensor) - else: - return gathered[rank - 1].clone() - - -# ============ Convenience aliases for conv1d CP ============ - -def conv_cp_send_recv_fwd(tails: torch.Tensor, group: ProcessGroup) -> torch.Tensor: - """ - Conv1d CP forward: each rank sends its tails, receives previous rank's tails as heads. - - Args: - tails: [W-1, D] or [N, D, W-1] - tail tokens from current rank - group: Process group - - Returns: - heads: Same shape as tails - head tokens from previous rank (zeros for rank 0) - """ - return send_recv_fwd(tails, group, recv_from_prev=True) - - -def conv_cp_send_recv_bwd(d_initial_state: torch.Tensor, group: ProcessGroup) -> torch.Tensor: - """ - Conv1d CP backward: each rank sends d_initial_state, receives from next rank. - - The received gradient should be added to the last W-1 tokens' gradient. - - Args: - d_initial_state: [W-1, D] or [N, D, W-1] - gradient w.r.t. initial state - group: Process group - - Returns: - recv_grad: Same shape - gradient from next rank (zeros for last rank) - """ - return send_recv_bwd(d_initial_state, group, recv_from_next=True) diff --git a/build/torch-xpu/ops/cp/context.py b/build/torch-xpu/ops/cp/context.py deleted file mode 100644 index 0fd0bd5dd5eb63e9f47e24234e3d9ef7fbcd18d2..0000000000000000000000000000000000000000 --- a/build/torch-xpu/ops/cp/context.py +++ /dev/null @@ -1,172 +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 - -from dataclasses import dataclass -from typing import TYPE_CHECKING - -import torch -import torch.distributed as dist - -from ...utils import tensor_cache - -if TYPE_CHECKING: - from torch.distributed import ProcessGroup - - -@dataclass -class FLACPContext: - """FLA Context Parallel Context - Operator-level context management.""" - group: ProcessGroup | None = None - cu_seqlens: torch.Tensor | None = None - cu_seqlens_cpu: torch.Tensor | None = None - is_last_rank: bool | None = None - pre_num_ranks: int | None = None - is_first_rank: bool | None = None - post_num_ranks: int | None = None - conv1d_kernel_size: int | None = None - pre_num_conv_tokens: int | None = None - - def copy_for_backward(self) -> FLACPContext: - """Create a copy for backward pass (useful when PP_SIZE > 1).""" - return FLACPContext( - group=self.group, - cu_seqlens=self.cu_seqlens.clone() if self.cu_seqlens is not None else None, - cu_seqlens_cpu=self.cu_seqlens_cpu.clone() if self.cu_seqlens_cpu is not None else None, - is_last_rank=self.is_last_rank, - pre_num_ranks=self.pre_num_ranks, - is_first_rank=self.is_first_rank, - post_num_ranks=self.post_num_ranks, - conv1d_kernel_size=self.conv1d_kernel_size, - pre_num_conv_tokens=self.pre_num_conv_tokens, - ) - - @property - def num_seqs(self) -> int: - """Number of sequences in this rank.""" - return 0 if self.cu_seqlens is None else len(self.cu_seqlens) - 1 - - @property - def is_cp_enabled(self) -> bool: - """Whether context parallel is enabled.""" - return self.group is not None - - -@tensor_cache -def get_cp_cu_seqlens( - cu_seqlens: torch.LongTensor, - cu_seqlens_cpu: torch.LongTensor | None = None, - world_size: int | None = None, - rank: int | None = None, - group: dist.ProcessGroup | None = None, - conv1d_kernel_size: int | None = None -) -> FLACPContext: - # 1. Initialize environment info - if world_size is None: - assert group is not None - world_size = dist.get_world_size(group=group) - rank = dist.get_rank(group=group) - - # 2. Operate on CPU to avoid D2H sync and leverage vectorization (int64/long) - if cu_seqlens_cpu is None: - cu_seqlens_cpu = cu_seqlens.cpu() - cu_seqlens_cpu = cu_seqlens_cpu.to(dtype=torch.long) - - # Get total tokens and current rank's responsible range - # Assume cu_seqlens is [0, s1, s1+s2, ..., total] - total_tokens = cu_seqlens_cpu[-1].item() - part_len = total_tokens // world_size - rank_start = part_len * rank - rank_end = rank_start + part_len - - # 3. Vectorized search: find sequences overlapping with current rank's interval [rank_start, rank_end) - # We need to find idx such that: global_ends[idx] > rank_start AND global_starts[idx] < rank_end - - # Optimization: cu_seqlens is sorted, use searchsorted to quickly locate boundaries - # Find first sequence whose end > rank_start - # cu_seqlens_cpu[1:] contains all sequence end points - start_seq_idx = torch.searchsorted(cu_seqlens_cpu[1:], rank_start, side='right') - - # Find first sequence whose start >= rank_end, sequences before this may overlap - # cu_seqlens_cpu[:-1] contains all sequence start points - end_seq_idx = torch.searchsorted(cu_seqlens_cpu[:-1], rank_end, side='left') - - # Slice cu_seqlens_cpu[start_seq_idx : end_seq_idx + 1] to get relevant global cu_seqlens nodes - # +1 because end_seq_idx is an open boundary, and cu_seqlens length is num_seqs + 1 - subset_cu_seqlens = cu_seqlens_cpu[start_seq_idx: end_seq_idx + 1] - - # 4. Compute local cu_seqlens on CPU (int32) - # Clamp global coordinates to [rank_start, rank_end], subtract rank_start to get local coordinates - # unique_consecutive removes duplicates from clamping (e.g., sequences entirely outside this rank) - local_cu_seqlens_cpu = ( - subset_cu_seqlens.clamp(min=rank_start, max=rank_end) - rank_start - ).unique_consecutive().to(torch.int32) - - # Transfer to GPU (int32, small tensor, fast transfer) - # non_blocking=True can further hide latency in CUDA streams - local_cu_seqlens_gpu = local_cu_seqlens_cpu.to( - device=cu_seqlens.device, non_blocking=True - ) - - # 5. Compute Context Parallel metadata (first/last rank info) - # Use slice endpoints directly, avoiding loops - - # Get global info for the first sequence that has data on current rank - first_seq_global_start = cu_seqlens_cpu[start_seq_idx].item() - # Get global info for the last sequence that has data on current rank - last_seq_global_end = cu_seqlens_cpu[end_seq_idx].item() - - # Number of tokens current rank needs from previous ranks for conv - pre_num_conv_tokens = max(0, rank_start - first_seq_global_start) - - # Compute first sequence's starting rank - first_rank_of_first_seq = first_seq_global_start // part_len - # Number of previous ranks current rank needs to receive state from - pre_num_ranks = rank - first_rank_of_first_seq - # Whether current rank is the first in the sequence's processing chain - is_first_rank = (rank == first_rank_of_first_seq) - - # Compute last sequence's ending rank - # (last_seq_global_end - 1) is the index of the last token - last_rank_of_last_seq = (last_seq_global_end - 1) // part_len - # Number of subsequent ranks current rank needs to send state to - post_num_ranks = last_rank_of_last_seq - rank - # Whether current rank is the last in the sequence's processing chain - is_last_rank = (rank == last_rank_of_last_seq) - - return FLACPContext( - group=group, - cu_seqlens=local_cu_seqlens_gpu, - cu_seqlens_cpu=local_cu_seqlens_cpu, - is_last_rank=is_last_rank, - pre_num_ranks=pre_num_ranks, - is_first_rank=is_first_rank, - post_num_ranks=post_num_ranks, - conv1d_kernel_size=conv1d_kernel_size, - pre_num_conv_tokens=pre_num_conv_tokens - ) - - -def build_cp_context( - cu_seqlens: torch.Tensor, - group: ProcessGroup, - conv1d_kernel_size: int | None = None, - cu_seqlens_cpu: torch.Tensor | None = None, -) -> FLACPContext: - """Build a CP context for the given cu_seqlens and process group. - - Args: - cu_seqlens: Cumulative sequence lengths tensor (before partition). - group: Process group for CP communication. - conv1d_kernel_size: Kernel size for convolution (optional). - cu_seqlens_cpu: CPU version of cu_seqlens to avoid d2h transfer (optional). - - Returns: - FLACPContext with computed cu_seqlens and rank information. - """ - return get_cp_cu_seqlens(cu_seqlens, cu_seqlens_cpu=cu_seqlens_cpu, group=group, conv1d_kernel_size=conv1d_kernel_size) diff --git a/build/torch-xpu/ops/delta_rule/__init__.py b/build/torch-xpu/ops/delta_rule/__init__.py deleted file mode 100644 index 8c8459c46fbc1de6dd3647a38d9870500d44aad2..0000000000000000000000000000000000000000 --- a/build/torch-xpu/ops/delta_rule/__init__.py +++ /dev/null @@ -1,16 +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 .chunk import chunk_delta_rule -from .fused_chunk import fused_chunk_delta_rule -from .fused_recurrent import fused_recurrent_delta_rule - -__all__ = [ - 'chunk_delta_rule', - 'fused_chunk_delta_rule', - 'fused_recurrent_delta_rule', -] diff --git a/build/torch-xpu/ops/delta_rule/chunk.py b/build/torch-xpu/ops/delta_rule/chunk.py deleted file mode 100644 index 0bd00789a1412e317bd7ea297d64c009f5841e09..0000000000000000000000000000000000000000 --- a/build/torch-xpu/ops/delta_rule/chunk.py +++ /dev/null @@ -1,338 +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 - -from ...modules.l2norm import l2norm_bwd, l2norm_fwd -from ...ops.common.chunk_delta_h import chunk_gated_delta_rule_bwd_dhu, chunk_gated_delta_rule_fwd_h -from ...ops.common.chunk_o import chunk_bwd_dqkwg, chunk_bwd_dv_local, chunk_fwd_o -from ...ops.delta_rule.wy_fast import prepare_wy_repr_bwd, prepare_wy_repr_fwd, recompute_w_u_fwd -from ...ops.utils.index import prepare_chunk_indices -from ...utils import autocast_custom_bwd, autocast_custom_fwd, input_guard - - -def chunk_delta_rule_fwd( - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - beta: torch.Tensor, - scale: float, - initial_state: torch.Tensor, - output_final_state: bool, - cu_seqlens: torch.LongTensor | None = None, - chunk_indices: torch.LongTensor | None = None, - chunk_size: int = 64, -): - # obtain WY representation. u is actually the new v. - w, u, A = prepare_wy_repr_fwd( - k=k, - v=v, - beta=beta, - cu_seqlens=cu_seqlens, - chunk_indices=chunk_indices, - chunk_size=chunk_size, - ) - h, v_new, final_state = chunk_gated_delta_rule_fwd_h( - k=k, - w=w, - u=u, - g=None, - initial_state=initial_state, - output_final_state=output_final_state, - cu_seqlens=cu_seqlens, - chunk_indices=chunk_indices, - chunk_size=chunk_size, - ) - - o = chunk_fwd_o( - q=q, - k=k, - v=v_new, - h=h, - g=None, - scale=scale, - cu_seqlens=cu_seqlens, - chunk_indices=chunk_indices, - chunk_size=chunk_size, - ) - return o, A, final_state - - -def chunk_delta_rule_bwd( - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - beta: torch.Tensor, - A: torch.Tensor, - scale: float, - initial_state: torch.Tensor, - do: torch.Tensor, - dht: torch.Tensor, - cu_seqlens: torch.LongTensor | None = None, - chunk_indices: torch.LongTensor | None = None, - chunk_size: int = 64, -): - w, u = recompute_w_u_fwd( - k=k, - v=v, - beta=beta, - A=A, - cu_seqlens=cu_seqlens, - chunk_indices=chunk_indices, - ) - h, v_new, _ = chunk_gated_delta_rule_fwd_h( - k=k, - w=w, - u=u, - g=None, - initial_state=initial_state, - output_final_state=False, - cu_seqlens=cu_seqlens, - chunk_indices=chunk_indices, - chunk_size=chunk_size, - ) - dv = chunk_bwd_dv_local( - q=q, - k=k, - do=do, - g=None, - scale=scale, - cu_seqlens=cu_seqlens, - chunk_indices=chunk_indices, - chunk_size=chunk_size, - ) - dh, dh0, dv = chunk_gated_delta_rule_bwd_dhu( - q=q, - k=k, - w=w, - g=None, - h0=initial_state, - dht=dht, - do=do, - dv=dv, - scale=scale, - cu_seqlens=cu_seqlens, - chunk_indices=chunk_indices, - chunk_size=chunk_size, - ) - dq, dk, dw, _ = chunk_bwd_dqkwg( - q=q, - k=k, - v=v_new, - h=h, - w=w, - dv=dv, - do=do, - dh=dh, - g=None, - scale=scale, - cu_seqlens=cu_seqlens, - chunk_indices=chunk_indices, - chunk_size=chunk_size, - ) - dk2, dv, db = prepare_wy_repr_bwd( - k=k, - v=v, - beta=beta, - A=A, - dw=dw, - du=dv, - cu_seqlens=cu_seqlens, - chunk_indices=chunk_indices, - ) - dk.add_(dk2) - return dq, dk, dv, db, dh0 - - -class ChunkDeltaRuleFunction(torch.autograd.Function): - - @staticmethod - @input_guard - @autocast_custom_fwd - def forward( - ctx, - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - beta: torch.Tensor, - scale: float, - initial_state: torch.Tensor, - output_final_state: bool, - use_qk_l2norm_in_kernel: bool = False, - cu_seqlens: torch.LongTensor | None = None, - cu_seqlens_cpu: torch.LongTensor | None = None, - chunk_size: int = 64, - ): - if use_qk_l2norm_in_kernel: - q, q_rstd = l2norm_fwd(q) - k, k_rstd = l2norm_fwd(k) - else: - q_rstd, k_rstd = None, None - - chunk_indices = prepare_chunk_indices( - cu_seqlens, chunk_size, cu_seqlens_cpu=cu_seqlens_cpu) if cu_seqlens is not None else None - o, A, final_state = chunk_delta_rule_fwd( - q=q, - k=k, - v=v, - beta=beta, - scale=scale, - initial_state=initial_state, - output_final_state=output_final_state, - cu_seqlens=cu_seqlens, - chunk_indices=chunk_indices, - chunk_size=chunk_size, - ) - ctx.save_for_backward(q, q_rstd, k, k_rstd, v, beta, A, initial_state, cu_seqlens, chunk_indices) - ctx.scale = scale - ctx.chunk_size = chunk_size - ctx.use_qk_l2norm_in_kernel = use_qk_l2norm_in_kernel - return o.to(q.dtype), final_state - - @staticmethod - @input_guard - @autocast_custom_bwd - def backward( - ctx, - do: torch.Tensor, - dht: torch.Tensor, - ): - q, q_rstd, k, k_rstd, v, beta, A, initial_state, cu_seqlens, chunk_indices = ctx.saved_tensors - - dq, dk, dv, db, dh0 = chunk_delta_rule_bwd( - q=q, - k=k, - v=v, - beta=beta, - A=A, - scale=ctx.scale, - initial_state=initial_state, - do=do, - dht=dht, - cu_seqlens=cu_seqlens, - chunk_indices=chunk_indices, - chunk_size=ctx.chunk_size, - ) - if ctx.use_qk_l2norm_in_kernel: - dq = l2norm_bwd(q, q_rstd, dq) - dk = l2norm_bwd(k, k_rstd, dk) - return dq.to(q.dtype), dk.to(k.dtype), dv.to(v.dtype), db.to(beta.dtype), None, dh0, None, None, None, None, None - - -@torch.compiler.disable -def chunk_delta_rule( - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - beta: torch.Tensor, - scale: float | None = None, - initial_state: torch.Tensor | None = None, - output_final_state: bool = False, - use_qk_l2norm_in_kernel: bool = False, - cu_seqlens: torch.LongTensor | None = None, - cu_seqlens_cpu: torch.LongTensor | None = None, - **kwargs, -): - r""" - Args: - q (torch.Tensor): - queries of shape `[B, T, H, K]`. - k (torch.Tensor): - keys of shape `[B, T, H, K]`. - v (torch.Tensor): - values of shape `[B, T, H, V]`. - beta (torch.Tensor): - betas of shape `[B, T, H]`. - scale (Optional[float]): - Scale factor for the attention scores. - If not provided, it will default to `1 / sqrt(K)`. Default: `None`. - initial_state (Optional[torch.Tensor]): - Initial state of shape `[N, H, K, V]` for `N` input sequences. - For equal-length input sequences, `N` equals the batch size `B`. - Default: `None`. - output_final_state (Optional[bool]): - Whether to output the final state of shape `[N, H, K, V]`. Default: `False`. - use_qk_l2norm_in_kernel (Optional[bool]): - Whether to use qk l2norm within the kernel for saving GPU memory. - Default: `False`. - cu_seqlens (torch.LongTensor): - Cumulative sequence lengths of shape `[N+1]` used for variable-length training, - consistent with the FlashAttention API. - cu_seqlens_cpu (torch.LongTensor): - CPU copy of `cu_seqlens` to avoid unnecessary device synchronization. Default: `None`. - - Returns: - o (torch.Tensor): - Outputs of shape `[B, T, H, V]`. - final_state (torch.Tensor): - Final state of shape `[N, H, K, V]` if `output_final_state=True` else `None`. - - Examples:: - >>> import torch - >>> import torch.nn.functional as F - >>> from einops import rearrange - >>> from ...ops.delta_rule import chunk_delta_rule - # inputs with equal lengths - >>> B, T, H, K, V = 4, 2048, 4, 512, 512 - >>> q = torch.randn(B, T, H, K, dtype=torch.bfloat16, device='cuda') - >>> k = F.normalize(torch.randn(B, T, H, K, dtype=torch.bfloat16, device='cuda'), p=2, dim=-1) - >>> v = torch.randn(B, T, H, V, dtype=torch.bfloat16, device='cuda') - >>> beta = torch.rand(B, T, H, dtype=torch.bfloat16, device='cuda').sigmoid() - >>> h0 = torch.randn(B, H, K, V, dtype=torch.bfloat16, device='cuda') - >>> o, ht = chunk_delta_rule( - q, k, v, beta, - initial_state=h0, - output_final_state=True - ) - # for variable-length inputs, the batch size `B` is expected to be 1 and `cu_seqlens` is required - >>> q, k, v, beta = map(lambda x: rearrange(x, 'b t ... -> 1 (b t) ...'), (q, k, v, beta)) - # for a batch with 4 sequences, `cu_seqlens` with 5 start/end positions are expected - >>> cu_seqlens = q.new_tensor([0, 2048, 4096, 6144, 8192], dtype=torch.long) - >>> o, ht = chunk_delta_rule( - q, k, v, beta, - initial_state=h0, - output_final_state=True, - cu_seqlens=cu_seqlens - ) - """ - assert q.dtype == k.dtype == v.dtype - assert q.dtype != torch.float32, "ChunkDeltaRuleFunction does not support float32. Please use bfloat16." - assert len(beta.shape) == 3, "beta must be of shape (batch size, num of head, seq len)." - - if 'head_first' in kwargs: - raise DeprecationWarning( - "head_first has been removed. Inputs must be in `[B, T, H, ...]` format.", - ) - chunk_size = kwargs.pop('chunk_size', 64) - if chunk_size not in (16, 32, 64): - raise ValueError(f"`chunk_size` must be 16, 32, or 64, got {chunk_size}.") - if cu_seqlens is not None: - if q.shape[0] != 1: - raise ValueError( - f"The batch size is expected to be 1 rather than {q.shape[0]} when using `cu_seqlens`. " - f"Please flatten variable-length inputs before processing.", - ) - if initial_state is not None and initial_state.shape[0] != len(cu_seqlens) - 1: - raise ValueError( - f"The number of initial states is expected to be equal to the number of input sequences, " - f"i.e., {len(cu_seqlens) - 1} rather than {initial_state.shape[0]}.", - ) - scale = k.shape[-1] ** -0.5 if scale is None else scale - o, final_state = ChunkDeltaRuleFunction.apply( - q, - k, - v, - beta, - scale, - initial_state, - output_final_state, - use_qk_l2norm_in_kernel, - cu_seqlens, - cu_seqlens_cpu, - chunk_size, - ) - return o, final_state diff --git a/build/torch-xpu/ops/delta_rule/fused_chunk.py b/build/torch-xpu/ops/delta_rule/fused_chunk.py deleted file mode 100644 index 87118c9c5d28b67aeef7079454dba7dc46a55560..0000000000000000000000000000000000000000 --- a/build/torch-xpu/ops/delta_rule/fused_chunk.py +++ /dev/null @@ -1,11 +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 - -def fused_chunk_delta_rule( - **kwargs, -): - raise NotImplementedError("fused_chunk_delta_rule is deprecated. Please use chunk_delta_rule instead.") diff --git a/build/torch-xpu/ops/delta_rule/fused_recurrent.py b/build/torch-xpu/ops/delta_rule/fused_recurrent.py deleted file mode 100644 index aa12c025915ea9d54860979722128ef98db025f0..0000000000000000000000000000000000000000 --- a/build/torch-xpu/ops/delta_rule/fused_recurrent.py +++ /dev/null @@ -1,537 +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 -import triton.language as tl - -from ...modules.l2norm import l2norm_bwd, l2norm_fwd -from ...utils import input_guard - - -@triton.heuristics({ - 'USE_INITIAL_STATE': lambda args: args['h0'] is not None, - 'STORE_FINAL_STATE': lambda args: args['ht'] is not None, - 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None, -}) -@triton.jit(do_not_specialize=['T']) -def fused_recurrent_delta_rule_fwd_kernel( - q, - k, - v, - u, - beta, - o, - h0, - ht, - cu_seqlens, - scale, - T, - B: tl.constexpr, - H: tl.constexpr, - K: tl.constexpr, - V: tl.constexpr, - BK: tl.constexpr, - BV: tl.constexpr, - USE_INITIAL_STATE: tl.constexpr, - STORE_FINAL_STATE: tl.constexpr, - IS_BETA_HEADWISE: tl.constexpr, - IS_VARLEN: tl.constexpr, -): - i_v, i_k, i_nh = tl.program_id(0), tl.program_id(1), tl.program_id(2) - i_n, i_h = i_nh // H, i_nh % H - if IS_VARLEN: - bos, eos = tl.load(cu_seqlens + i_n).to(tl.int64), tl.load(cu_seqlens + i_n + 1).to(tl.int64) - all = T - T = eos - bos - else: - bos, eos = i_n * T, i_n * T + T - all = B * T - - p_q = q + (bos * H + i_h) * K + i_k * BK + tl.arange(0, BK) - p_k = k + (bos * H + i_h) * K + i_k * BK + tl.arange(0, BK) - p_v = v + (bos * H + i_h) * V + i_v * BV + tl.arange(0, BV) - p_u = u + (bos * H + i_h) * V + i_v * BV + tl.arange(0, BV) - if IS_BETA_HEADWISE: - p_beta = beta + (bos * H + i_h) * V + i_v * BV + tl.arange(0, BV) - else: - p_beta = beta + bos * H + i_h - p_o = o + ((i_k * all + bos) * H + i_h) * V + i_v * BV + tl.arange(0, BV) - - mask_k = (i_k * BK + tl.arange(0, BK)) < K - mask_v = (i_v * BV + tl.arange(0, BV)) < V - mask_h = mask_k[None, :] & mask_v[:, None] - - b_h = tl.zeros([BV, BK], dtype=tl.float32) - if USE_INITIAL_STATE: - p_h0 = h0 + i_nh * K * V + (i_k * BK + tl.arange(0, BK)[None, :]) * V + (i_v * BV + tl.arange(0, BV)[:, None]) - b_h += tl.load(p_h0, mask=mask_h, other=0).to(tl.float32) - - for _ in range(0, T): - b_k = tl.load(p_k, mask=mask_k, other=0).to(tl.float32) - b_v = tl.load(p_v, mask=mask_v, other=0).to(tl.float32) - b_q = tl.load(p_q, mask=mask_k, other=0).to(tl.float32) * scale - b_v_minus = tl.sum(b_h * b_k[None, :], axis=1) - b_v -= b_v_minus - if IS_BETA_HEADWISE: - b_beta = tl.load(p_beta, mask=mask_v, other=0).to(tl.float32) - else: - b_beta = tl.load(p_beta).to(tl.float32) - tl.store(p_u, b_v.to(p_v.dtype.element_ty), mask=mask_v) - b_v *= b_beta - b_h += b_k[None, :] * b_v[:, None] - b_o = b_h * b_q[None, :] - b_o = tl.sum(b_o, axis=1) - tl.store(p_o, b_o.to(p_o.dtype.element_ty), mask=mask_v) - - p_q += H*K - p_k += H*K - p_o += H*V - p_v += H*V - p_u += H*V - p_beta += H * (V if IS_BETA_HEADWISE else 1) - - if STORE_FINAL_STATE: - p_ht = ht + i_nh * K * V + (i_k * BK + tl.arange(0, BK)[None, :]) * V + (i_v * BV + tl.arange(0, BV)[:, None]) - tl.store(p_ht, b_h.to(p_ht.dtype.element_ty), mask=mask_h) - - -@triton.heuristics({ - 'USE_INITIAL_STATE': lambda args: args['h0'] is not None, - 'USE_FINAL_STATE_GRADIENT': lambda args: args['dht'] is not None, - 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None, -}) -@triton.jit(do_not_specialize=['T']) -def fused_recurrent_delta_rule_bwd_kernel( - q, - k, - v, - beta, - h0, - dh0, - dht, - do, - dq, - dk, - dv, - db, - cu_seqlens, - scale, - B: tl.constexpr, - T, - H: tl.constexpr, - K: tl.constexpr, - V: tl.constexpr, - BK: tl.constexpr, - BV: tl.constexpr, - NK: tl.constexpr, - IS_BETA_HEADWISE: tl.constexpr, # whether beta is headwise vector or scalar - USE_INITIAL_STATE: tl.constexpr, # whether to use dh0 - USE_FINAL_STATE_GRADIENT: tl.constexpr, # whether to use dht - IS_VARLEN: tl.constexpr, -): - i_v, i_k, i_nh = tl.program_id(0), tl.program_id(1), tl.program_id(2) - i_n, i_h = i_nh // H, i_nh % H - if IS_VARLEN: - bos, eos = tl.load(cu_seqlens + i_n).to(tl.int64), tl.load(cu_seqlens + i_n + 1).to(tl.int64) - all = T - T = eos - bos - else: - bos, eos = i_n * T, i_n * T + T - all = B * T - - mask_k = i_k * BK + tl.arange(0, BK) < K - mask_v = i_v * BV + tl.arange(0, BV) < V - - p_q = q + (bos * H + i_h) * K + i_k * BK + tl.arange(0, BK) + (T - 1) * H*K - p_k = k + (bos * H + i_h) * K + i_k * BK + tl.arange(0, BK) + (T - 1) * H*K - p_v = v + (bos * H + i_h) * V + i_v * BV + tl.arange(0, BV) + (T - 1) * H*V - p_do = do + (bos * H + i_h) * V + i_v * BV + tl.arange(0, BV) + (T - 1) * H*V - p_dk = dk + ((i_v * all + bos) * H + i_h) * K + i_k * BK + tl.arange(0, BK) + (T - 1) * H*K - p_dv = dv + ((i_k * all + bos) * H + i_h) * V + i_v * BV + tl.arange(0, BV) + (T - 1) * H*V - if IS_BETA_HEADWISE: - p_beta = beta + (bos + T - 1) * H*V + i_h * V + i_v * BV + tl.arange(0, BV) - p_dbeta = db + ((i_v * NK + i_k) * all + bos + T - 1) * H*V + i_h * V + tl.arange(0, BV) - else: - p_beta = beta + (bos + T - 1) * H + i_h - p_dbeta = db + (i_v * all + bos + T - 1) * H + i_h - - b_dh = tl.zeros([BK, BV], dtype=tl.float32) - if USE_FINAL_STATE_GRADIENT: - p_ht = dht + i_nh * K * V + (i_k * BK + tl.arange(0, BK)[:, None]) * V + (i_v * BV + tl.arange(0, BV)[None, :]) - b_dh += tl.load(p_ht, mask=mask_k[:, None] & mask_v[None, :], other=0).to(tl.float32) - - for _ in range(T): - b_q = tl.load(p_q, mask=mask_k, other=0).to(tl.float32) * scale - b_k = tl.load(p_k, mask=mask_k, other=0).to(tl.float32) - b_v = tl.load(p_v, mask=mask_v, other=0).to(tl.float32) - b_do = tl.load(p_do, mask=mask_v, other=0).to(tl.float32) - if IS_BETA_HEADWISE: - b_beta = tl.load(p_beta, mask=mask_v, other=0).to(tl.float32) - else: - b_beta = tl.load(p_beta).to(tl.float32) - b_dh += b_q[:, None] * b_do[None, :] - b_dk = tl.sum(b_dh * (b_v * b_beta)[None, :], axis=1) - b_dv = tl.sum(b_dh * b_k[:, None], axis=0) - - b_db = b_dv * b_v if IS_BETA_HEADWISE else tl.sum(b_dv * b_v) - b_dv = b_dv * b_beta - - tl.store(p_dk, b_dk.to(p_dk.dtype.element_ty), mask=mask_k) - tl.store(p_dv, b_dv.to(p_dv.dtype.element_ty), mask=mask_v) - if IS_BETA_HEADWISE: - tl.store(p_dbeta, b_db.to(p_dbeta.dtype.element_ty), mask=mask_v) - else: - tl.store(p_dbeta, b_db.to(p_dbeta.dtype.element_ty)) - - b_dh -= b_k[:, None] * b_dv[None, :] - - p_q -= H*K - p_k -= H*K - p_v -= H*V - p_do -= H*V - p_dk -= H*K - p_dv -= H*V - p_dbeta -= H * (V if IS_BETA_HEADWISE else 1) - p_beta -= H * (V if IS_BETA_HEADWISE else 1) - - if USE_INITIAL_STATE: - p_dh0 = dh0 + i_nh * K * V + (i_k * BK + tl.arange(0, BK)[:, None]) * V + (i_v * BV + tl.arange(0, BV)[None, :]) - tl.store(p_dh0, b_dh.to(p_dh0.dtype.element_ty), mask=mask_k[:, None] & mask_v[None, :]) - - tl.debug_barrier() - - b_h = tl.zeros([BK, BV], dtype=tl.float32) - - p_q = q + (bos * H + i_h) * K + i_k * BK + tl.arange(0, BK) - p_k = k + (bos * H + i_h) * K + i_k * BK + tl.arange(0, BK) - p_v = v + (bos * H + i_h) * V + i_v * BV + tl.arange(0, BV) - if IS_BETA_HEADWISE: - p_beta = beta + (bos * H + i_h) * V + i_v * BV + tl.arange(0, BV) - else: - p_beta = beta + bos * H + i_h - p_do = do + (bos * H + i_h) * V + i_v * BV + tl.arange(0, BV) - p_dq = dq + ((i_v * all + bos) * H + i_h) * K + i_k * BK + tl.arange(0, BK) - p_dk = dk + ((i_v * all + bos) * H + i_h) * K + i_k * BK + tl.arange(0, BK) - p_dv = dv + ((i_k * all + bos) * H + i_h) * V + i_v * BV + tl.arange(0, BV) - - if USE_INITIAL_STATE: - mask_h = mask_k[:, None] & mask_v[None, :] - p_h0 = h0 + i_nh * K * V + (i_k * BK + tl.arange(0, BK)[:, None]) * V + (i_v * BV + tl.arange(0, BV)[None, :]) - b_h += tl.load(p_h0, mask=mask_h, other=0).to(tl.float32) - - for _ in range(0, T): - b_dk = tl.load(p_dk, mask=mask_k, other=0).to(tl.float32) - b_dv = tl.load(p_dv, mask=mask_v, other=0).to(tl.float32) - b_dk -= tl.sum(b_dv[None, :] * b_h, axis=1) - tl.store(p_dk, b_dk.to(p_dk.dtype.element_ty), mask=mask_k) - - b_k = tl.load(p_k, mask=mask_k, other=0).to(tl.float32) - b_v = tl.load(p_v, mask=mask_v, other=0).to(tl.float32) - b_do = tl.load(p_do, mask=mask_v, other=0).to(tl.float32) - if IS_BETA_HEADWISE: - b_beta = tl.load(p_beta, mask=mask_v, other=0).to(tl.float32) - else: - b_beta = tl.load(p_beta).to(tl.float32) - b_v *= b_beta - - b_h += b_k[:, None] * b_v[None, :] - b_dq = b_h * b_do[None, :] - d_q = tl.sum(b_dq, axis=1) * scale - tl.store(p_dq, d_q.to(p_dq.dtype.element_ty), mask=mask_k) - - p_k += H*K - p_v += H*V - p_do += H*V - p_dq += H*K - p_dk += H*K - p_dv += H*V - p_beta += H * (V if IS_BETA_HEADWISE else 1) - - -def fused_recurrent_delta_rule_fwd( - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - beta: torch.Tensor, - scale: float, - initial_state: torch.Tensor, - output_final_state: bool, - cu_seqlens: torch.LongTensor | None = None, -) -> tuple[torch.Tensor, torch.Tensor]: - B, T, H, K, V = *k.shape, v.shape[-1] - N = B if cu_seqlens is None else len(cu_seqlens) - 1 - BK, BV = triton.next_power_of_2(K), min(triton.next_power_of_2(V), 8) - NK, NV = triton.cdiv(K, BK), triton.cdiv(V, BV) - assert NK == 1, "NK > 1 is not supported yet" - num_stages = 1 - num_warps = 1 - - o = q.new_empty(NK, *v.shape) - if output_final_state: - final_state = q.new_empty(N, H, K, V, dtype=torch.float32) - else: - final_state = None - - grid = (NV, NK, N * H) - u = torch.empty_like(v) - fused_recurrent_delta_rule_fwd_kernel[grid]( - q, - k, - v, - u, - beta, - o, - initial_state, - final_state, - cu_seqlens, - scale, - T=T, - B=B, - H=H, - K=K, - V=V, - BK=BK, - BV=BV, - IS_BETA_HEADWISE=beta.ndim == v.ndim, - num_warps=num_warps, - num_stages=num_stages, - ) - o = o.squeeze(0) - return o, u, final_state - - -def fused_recurrent_delta_rule_bwd( - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - beta: torch.Tensor, - dht: torch.Tensor, - do: torch.Tensor, - scale: float, - initial_state: torch.Tensor, - cu_seqlens: torch.LongTensor | None = None, -) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: - B, T, H, K, V = *k.shape, v.shape[-1] - N = B if cu_seqlens is None else len(cu_seqlens) - 1 - BK, BV = triton.next_power_of_2(K), min(triton.next_power_of_2(V), 32) - NK, NV = triton.cdiv(K, BK), triton.cdiv(V, BV) - assert NK == 1, "NK > 1 is not supported yet" - num_stages = 1 - num_warps = 2 - - beta_vector = beta.ndim == v.ndim - - dq = q.new_empty(NV, *q.shape) - dk = q.new_empty(NV, *k.shape) - dv = q.new_empty(NK, *v.shape) - if beta_vector: - db = q.new_empty(NV, NK, B, T, H, V) - else: - db = q.new_empty(NV, B, T, H) - grid = (NV, NK, N * H) - - if initial_state is not None and initial_state.requires_grad: - dh0 = torch.empty_like(initial_state, dtype=torch.float32) - else: - dh0 = None - - fused_recurrent_delta_rule_bwd_kernel[grid]( - q, - k, - v, - beta, - initial_state, - dh0, - dht, - do, - dq, - dk, - dv, - db, - cu_seqlens, - scale, - T=T, - B=B, - H=H, - K=K, - V=V, - BK=BK, - BV=BV, - NK=NK, - IS_BETA_HEADWISE=beta_vector, - num_warps=num_warps, - num_stages=num_stages, - ) - dq = dq.sum(0) - dk = dk.sum(0) - dv = dv.sum(0) - db = db.sum((0, 1)) if beta_vector else db.sum(0) - - return dq, dk, dv, db, dh0 - - -class FusedRecurrentFunction(torch.autograd.Function): - - @staticmethod - @input_guard - def forward( - ctx, - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - beta: torch.Tensor, - scale: float, - initial_state: torch.Tensor, - output_final_state: bool, - use_qk_l2norm_in_kernel: bool = False, - cu_seqlens: torch.LongTensor | None = None, - ): - if use_qk_l2norm_in_kernel: - q, q_rstd = l2norm_fwd(q) - k, k_rstd = l2norm_fwd(k) - else: - q_rstd, k_rstd = None, None - - o, u, final_state = fused_recurrent_delta_rule_fwd( - q=q, - k=k, - v=v, - beta=beta, - scale=scale, - initial_state=initial_state, - output_final_state=output_final_state, - cu_seqlens=cu_seqlens, - ) - - ctx.save_for_backward(q, q_rstd, k, k_rstd, u, beta, initial_state) - ctx.scale = scale - ctx.cu_seqlens = cu_seqlens - ctx.use_qk_l2norm_in_kernel = use_qk_l2norm_in_kernel - return o, final_state - - @staticmethod - @input_guard - def backward(ctx, do, dht): - q, q_rstd, k, k_rstd, v, beta, initial_state = ctx.saved_tensors - dq, dk, dv, db, dh0 = fused_recurrent_delta_rule_bwd( - q=q, - k=k, - v=v, - beta=beta, - dht=dht, - do=do, - scale=ctx.scale, - initial_state=initial_state, - cu_seqlens=ctx.cu_seqlens, - ) - if ctx.use_qk_l2norm_in_kernel: - dq = l2norm_bwd(q, q_rstd, dq) - dk = l2norm_bwd(k, k_rstd, dk) - return dq.to(q), dk.to(k), dv.to(v), db.to(beta), None, dh0, None, None, None - - -@torch.compiler.disable -def fused_recurrent_delta_rule( - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - beta: torch.Tensor = None, - scale: float = None, - initial_state: torch.Tensor = None, - output_final_state: bool = False, - use_qk_l2norm_in_kernel: bool = False, - cu_seqlens: torch.LongTensor | None = None, -) -> tuple[torch.Tensor, torch.Tensor]: - r""" - Args: - q (torch.Tensor): - queries of shape `[B, T, H, K]`. - k (torch.Tensor): - keys of shape `[B, T, H, K]`. - v (torch.Tensor): - values of shape `[B, T, H, V]`. - beta (torch.Tensor): - betas of shape `[B, T, H]`. - scale (Optional[float]): - Scale factor for the RetNet attention scores. - If not provided, it will default to `1 / sqrt(K)`. Default: `None`. - initial_state (Optional[torch.Tensor]): - Initial state of shape `[N, H, K, V]` for `N` input sequences. - For equal-length input sequences, `N` equals the batch size `B`. - Default: `None`. - output_final_state (Optional[bool]): - Whether to output the final state of shape `[N, H, K, V]`. Default: `False`. - use_qk_l2norm_in_kernel (Optional[bool]): - Whether to use L2 normalization in the kernel. Default: `False`. - cu_seqlens (torch.LongTensor): - Cumulative sequence lengths of shape `[N+1]` used for variable-length training, - consistent with the FlashAttention API. - - Returns: - o (torch.Tensor): - Outputs of shape `[B, T, H, V]`. - final_state (torch.Tensor): - Final state of shape `[N, H, K, V]` if `output_final_state=True` else `None`. - - Examples:: - >>> import torch - >>> import torch.nn.functional as F - >>> from einops import rearrange - >>> from ...ops.delta_rule import fused_recurrent_delta_rule - # inputs with equal lengths - >>> B, T, H, K, V = 4, 2048, 4, 512, 512 - >>> q = torch.randn(B, T, H, K, device='cuda') - >>> k = F.normalize(torch.randn(B, T, H, K, device='cuda'), p=2, dim=-1) - >>> v = torch.randn(B, T, H, V, device='cuda') - >>> beta = torch.rand(B, T, H, device='cuda').sigmoid() - >>> h0 = torch.randn(B, H, K, V, device='cuda') - >>> o, ht = fused_recurrent_delta_rule( - q, k, v, beta, - initial_state=h0, - output_final_state=True - ) - # for variable-length inputs, the batch size `B` is expected to be 1 and `cu_seqlens` is required - >>> q, k, v, beta = map(lambda x: rearrange(x, 'b t ... -> 1 (b t) ...'), (q, k, v, beta)) - # for a batch with 4 sequences, `cu_seqlens` with 5 start/end positions are expected - >>> cu_seqlens = q.new_tensor([0, 2048, 4096, 6144, 8192], dtype=torch.long) - >>> o, ht = fused_recurrent_delta_rule( - q, k, v, beta, - initial_state=h0, - output_final_state=True, - cu_seqlens=cu_seqlens - ) - """ - if cu_seqlens is not None: - if q.shape[0] != 1: - raise ValueError( - f"The batch size is expected to be 1 rather than {q.shape[0]} when using `cu_seqlens`." - f"Please flatten variable-length inputs before processing.", - ) - if initial_state is not None and initial_state.shape[0] != len(cu_seqlens) - 1: - raise ValueError( - f"The number of initial states is expected to be equal to the number of input sequences, " - f"i.e., {len(cu_seqlens) - 1} rather than {initial_state.shape[0]}.", - ) - if scale is None: - scale = k.shape[-1] ** -0.5 - else: - assert scale > 0, "scale must be positive" - if beta is None: - beta = torch.ones_like(q[..., 0]) - o, final_state = FusedRecurrentFunction.apply( - q, - k, - v, - beta, - scale, - initial_state, - output_final_state, - use_qk_l2norm_in_kernel, - cu_seqlens, - ) - return o, final_state diff --git a/build/torch-xpu/ops/delta_rule/naive.py b/build/torch-xpu/ops/delta_rule/naive.py deleted file mode 100644 index d1f65a859a9cf74b98e71fbd9eb18c2165f9822a..0000000000000000000000000000000000000000 --- a/build/torch-xpu/ops/delta_rule/naive.py +++ /dev/null @@ -1,125 +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 -from einops import rearrange - - -def delta_rule_recurrence(q, k, v, beta, initial_state=None, output_final_state=True): - orig_dtype = q.dtype - b, h, l, d_k = q.shape - q, k, v, beta = map(lambda x: x.float(), [q, k, v, beta]) - d_v = v.shape[-1] - o = torch.zeros_like(v) - S = torch.zeros(b, h, d_k, d_v).to(v) - q = q * (d_k ** -0.5) - - if beta.ndim < v.ndim: - beta = beta[..., None] - - if initial_state is not None: - S += initial_state - - for i in range(l): - _k = k[:, :, i] - _q = q[:, :, i] - _v = v[:, :, i].clone() - beta_i = beta[:, :, i] - _v = _v - (S.clone() * _k[..., None]).sum(-2) - _v = _v * beta_i - S = S.clone() + _k.unsqueeze(-1) * _v.unsqueeze(-2) - o[:, :, i] = torch.einsum('bhd,bhdm->bhm', _q, S) - S = None if output_final_state is False else S - return o.to(orig_dtype), S - - -def delta_rule_chunkwise(q, k, v, beta, chunk_size=32): - b, h, l, d_k = q.shape - d_v = v.shape[-1] - q = q * (d_k ** -0.5) - v = v * beta[..., None] - k_beta = k * beta[..., None] - - assert l % chunk_size == 0 - - # compute (I - tri(diag(beta) KK^T))^{-1} - mask = torch.triu(torch.ones(chunk_size, chunk_size, dtype=torch.bool, device=q.device), diagonal=0) - q, k, v, k_beta = map(lambda x: rearrange(x, 'b h (n c) d -> b h n c d', c=chunk_size), [q, k, v, k_beta]) - attn = -(k_beta @ k.transpose(-1, -2)).masked_fill(mask, 0) - for i in range(1, chunk_size): - attn[..., i, :i] = attn[..., i, :i] + (attn[..., i, :, None].clone() * attn[..., :, :i].clone()).sum(-2) - attn = attn + torch.eye(chunk_size, dtype=torch.float, device=q.device) - - u = attn @ v - w = attn @ k_beta - S = k.new_zeros(b, h, d_k, d_v) - o = torch.zeros_like(v) - mask = torch.triu(torch.ones(chunk_size, chunk_size, dtype=torch.bool, device=q.device), diagonal=1) - for i in range(0, l // chunk_size): - q_i, k_i = q[:, :, i], k[:, :, i] - attn = (q_i @ k_i.transpose(-1, -2)).masked_fill_(mask, 0) - u_i = u[:, :, i] - w[:, :, i] @ S - o_inter = q_i @ S - o[:, :, i] = o_inter + attn @ u_i - S = S + k_i.transpose(-1, -2) @ u_i - - return rearrange(o, 'b h n c d -> b h (n c) d'), S - - -def delta_rule_parallel(q, k, v, beta, BM=128, BN=32): - b, h, l, d_k = q.shape - # d_v = v.shape[-1] - q = q * (d_k ** -0.5) - v = v * beta[..., None] - k_beta = k * beta[..., None] - # compute (I - tri(diag(beta) KK^T))^{-1} - q, k, v, k_beta = map(lambda x: rearrange(x, 'b h (n c) d -> b h n c d', c=BN), [q, k, v, k_beta]) - mask = torch.triu(torch.ones(BN, BN, dtype=torch.bool, device=q.device), diagonal=0) - T = -(k_beta @ k.transpose(-1, -2)).masked_fill(mask, 0) - for i in range(1, BN): - T[..., i, :i] = T[..., i, :i].clone() + (T[..., i, :, None].clone() * T[..., :, :i].clone()).sum(-2) - T = T + torch.eye(BN, dtype=torch.float, device=q.device) - - mask2 = torch.triu(torch.ones(BN, BN, dtype=torch.bool, device=q.device), diagonal=1) - A_local = (q @ k.transpose(-1, -2)).masked_fill(mask2, 0) @ T - o_intra = A_local @ v - - # apply cumprod transition matrices on k to the last position within the chunk - k = k - ((k @ k.transpose(-1, -2)).masked_fill(mask, 0) @ T).transpose(-1, -2) @ k_beta - # apply cumprod transition matrices on q to the first position within the chunk - q = q - A_local @ k_beta - o_intra = A_local @ v - - A = torch.zeros(b, h, l, l, device=q.device) - - q, k, v, k_beta, o_intra = map(lambda x: rearrange(x, 'b h n c d -> b h (n c) d'), [q, k, v, k_beta, o_intra]) - o = torch.empty_like(v) - for i in range(0, l, BM): - q_i = q[:, :, i:i+BM] - o_i = o_intra[:, :, i:i+BM] - # intra block - for j in range(i + BM - 2 * BN, i-BN, -BN): - k_j = k[:, :, j:j+BN] - A_ij = q_i @ k_j.transpose(-1, -2) - mask = torch.arange(i, i+BM) >= (j + BN) - A_ij = A_ij.masked_fill_(~mask[:, None].to(A_ij.device), 0) - A[:, :, i:i+BM, j:j+BN] = A_ij - q_i = q_i - A_ij @ k_beta[:, :, j:j+BN] - o_i += A_ij @ v[:, :, j:j+BN] - # inter block - for j in range(i - BN, -BN, -BN): - k_j = k[:, :, j:j+BN] - A_ij = q_i @ k_j.transpose(-1, -2) - A[:, :, i:i+BM, j:j+BN] = A_ij - q_i = q_i - A_ij @ k_beta[:, :, j:j+BN] - o_i += A_ij @ v[:, :, j:j+BN] - o[:, :, i:i+BM] = o_i - - for i in range(0, l//BN): - A[:, :, i*BN:i*BN+BN, i*BN:i*BN+BN] = A_local[:, :, i] - - return o, A diff --git a/build/torch-xpu/ops/delta_rule/parallel.py b/build/torch-xpu/ops/delta_rule/parallel.py deleted file mode 100644 index fdb47543a81a1f781556108e178e394ec487ac5e..0000000000000000000000000000000000000000 --- a/build/torch-xpu/ops/delta_rule/parallel.py +++ /dev/null @@ -1,395 +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 -import triton.language as tl -from einops import rearrange - -from ...ops.delta_rule.wy_fast import fwd_prepare_T -from ...utils import autocast_custom_bwd, autocast_custom_fwd, autotune_cache_kwargs, input_guard - - -@triton.autotune( - configs=[ - triton.Config({}, num_warps=num_warps) - for num_warps in [1, 2, 4] - ], - key=['BT', 'K', 'V'], - **autotune_cache_kwargs, -) -@triton.jit(do_not_specialize=['T']) -def chunk_transform_qk_fwd_kernel( - q, - k, - v, - beta, - o, - A, - q_new, - k_new, - A_local, - scale, - T, - K: tl.constexpr, - V: tl.constexpr, - BK: tl.constexpr, - BV: tl.constexpr, - BT: tl.constexpr, - OUTPUT_ATTENTIONS: tl.constexpr, -): - i_t, i_bh = tl.program_id(0), tl.program_id(1) - - p_q = tl.make_block_ptr(q + i_bh * T*K, (T, K), (K, 1), (i_t * BT, 0), (BT, BK), (1, 0)) - p_k = tl.make_block_ptr(k + i_bh * T*K, (T, K), (K, 1), (i_t * BT, 0), (BT, BK), (1, 0)) - p_v = tl.make_block_ptr(v + i_bh * T*V, (T, V), (V, 1), (i_t * BT, 0), (BT, BV), (1, 0)) - b_q = (tl.load(p_q, boundary_check=(0, 1)) * scale).to(p_q.dtype.element_ty) - b_k = tl.load(p_k, boundary_check=(0, 1)) - b_v = tl.load(p_v, boundary_check=(0, 1)) - - p_T = tl.make_block_ptr(A + i_bh * T * BT, (T, BT), (BT, 1), (i_t * BT, 0), (BT, BT), (1, 0)) - b_T = tl.load(p_T, boundary_check=(0, 1)) - - o_i = tl.arange(0, BT) - m_t = o_i[:, None] >= o_i[None, :] - b_qk = tl.where(m_t, tl.dot(b_q, tl.trans(b_k), allow_tf32=False), 0).to(b_q.dtype) - m_t = o_i[:, None] > o_i[None, :] - b_kk = tl.where(m_t, tl.dot(b_k, tl.trans(b_k), allow_tf32=False), 0).to(b_k.dtype) - - p_beta = tl.make_block_ptr(beta + i_bh * T, (T, ), (1, ), (i_t * BT, ), (BT, ), (0, )) - b_beta = tl.load(p_beta, boundary_check=(0, )) - b_k_beta = (b_k * b_beta[:, None]).to(b_k.dtype) - - b_qkT = tl.dot(b_qk, b_T, allow_tf32=False).to(b_k.dtype) - - if OUTPUT_ATTENTIONS: - p_a = tl.make_block_ptr(A_local + i_bh * T * BT, (T, BT), (BT, 1), (i_t * BT, 0), (BT, BT), (1, 0)) - tl.store(p_a, b_qkT.to(p_a.dtype.element_ty), boundary_check=(0, 1)) - - b_kkT = tl.dot(b_kk, b_T, allow_tf32=False).to(b_k.dtype) - p_o = tl.make_block_ptr(o + i_bh * T*V, (T, V), (V, 1), (i_t * BT, 0), (BT, BV), (1, 0)) - tl.store(p_o, tl.dot(b_qkT, b_v).to(p_o.dtype.element_ty), boundary_check=(0, 1)) - - p_q_new = tl.make_block_ptr(q_new + i_bh * T*K, (T, K), (K, 1), (i_t * BT, 0), (BT, BK), (1, 0)) - tl.store(p_q_new, (b_q - tl.dot(b_qkT, b_k_beta, allow_tf32=False)).to(p_q_new.dtype.element_ty), boundary_check=(0, 1)) - - p_k_new = tl.make_block_ptr(k_new + i_bh * T*K, (T, K), (K, 1), (i_t * BT, 0), (BT, BK), (1, 0)) - b_k_new = b_k - tl.dot(tl.trans(b_kkT), b_k_beta, allow_tf32=False) - tl.store(p_k_new, b_k_new.to(p_k_new.dtype.element_ty), boundary_check=(0, 1)) - - -def chunk_transform_qk_fwd( - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - beta: torch.Tensor, - A: torch.Tensor, - scale: float, - chunk_size: int, - output_attentions: bool, -): - B, H, T, K = k.shape - BT = chunk_size - q_new = torch.empty_like(q) - k_new = torch.empty_like(k) - o = torch.empty_like(v) - grid = (triton.cdiv(T, BT), B*H) - V = v.shape[-1] - A_local = torch.empty_like(A) if output_attentions else None - chunk_transform_qk_fwd_kernel[grid]( - q, - k, - v, - beta, - o, - A, - q_new, - k_new, - A_local, - scale=scale, - T=T, - K=K, - V=V, - BT=BT, - BK=triton.next_power_of_2(K), - BV=triton.next_power_of_2(V), - OUTPUT_ATTENTIONS=output_attentions, - ) - return q_new, k_new, o, A_local - - -@triton.autotune( - configs=[ - triton.Config({}, num_warps=1), - triton.Config({}, num_warps=2), - ], - key=['BT'], - **autotune_cache_kwargs, -) -@triton.jit(do_not_specialize=['T']) -def save_intra_chunk_attn( - A, - A_local, - T, - BT: tl.constexpr, -): - i_t, i_bh = tl.program_id(0), tl.program_id(1) - p_A = tl.make_block_ptr(A + i_bh * T * T, (T, T), (T, 1), (i_t * BT, i_t * BT), (BT, BT), (1, 0)) - p_A_local = tl.make_block_ptr(A_local + i_bh * T * BT, (T, BT), (BT, 1), (i_t * BT, 0), (BT, BT), (1, 0)) - b_A_local = tl.load(p_A_local, boundary_check=(0, 1)) - tl.store(p_A, b_A_local.to(p_A.dtype.element_ty), boundary_check=(0, 1)) - - -@triton.heuristics({ - 'OUTPUT_ATTENTIONS': lambda args: args['attn'] is not None, -}) -@triton.jit(do_not_specialize=['T']) -def parallel_delta_rule_fwd_kernel( - q, - k, - k2, # original k - v, - beta, - o, - o_new, - attn, - T, - K: tl.constexpr, - V: tl.constexpr, - BT: tl.constexpr, - BS: tl.constexpr, - BK: tl.constexpr, - BV: tl.constexpr, - OUTPUT_ATTENTIONS: tl.constexpr, -): - i_t, i_bh = tl.program_id(0), tl.program_id(1) - p_q = tl.make_block_ptr(q + i_bh * T*K, (T, K), (K, 1), (i_t * BT, 0), (BT, BK), (1, 0)) - - # the Q block is kept in the shared memory throughout the whole kernel - # [BT, BK] - b_q = tl.zeros([BT, BK], dtype=tl.float32) - b_q += tl.load(p_q, boundary_check=(0, 1)) - - b_o = tl.zeros([BT, BV], dtype=tl.float32) - p_o = tl.make_block_ptr(o + i_bh * T*V, (T, V), (V, 1), (i_t * BT, 0), (BT, BV), (1, 0)) - b_o += tl.load(p_o, boundary_check=(0, 1)) - - # As opposed to Flashattention, this kernel requires scanning the KV blocks from right to left - # Q block and K block have overlap. - # masks required - for offset in range((i_t + 1) * BT - 2 * BS, i_t * BT - BS, -BS): - p_k = tl.make_block_ptr(k + i_bh * T*K, (K, T), (1, K), (0, offset), (BK, BS), (0, 1)) - p_k2 = tl.make_block_ptr(k2 + i_bh * T*K, (T, K), (K, 1), (offset, 0), (BS, BK), (1, 0)) - p_v = tl.make_block_ptr(v + i_bh * T*V, (T, V), (V, 1), (offset, 0), (BS, BV), (1, 0)) - p_beta = tl.make_block_ptr(beta + i_bh * T, (T, ), (1, ), (offset, ), (BS, ), (0,)) - # [BK, BS] - b_k = tl.load(p_k, boundary_check=(0, 1)) - # [BS, BV] - b_v = tl.load(p_v, boundary_check=(0, 1)) - # [BS] - b_beta = tl.load(p_beta, boundary_check=(0,)) - # [BT, BS] - m_s = tl.arange(0, BT) >= (offset - i_t*BT + BS) - b_s = tl.dot(b_q.to(b_k.dtype), b_k, allow_tf32=False) - b_s = tl.where(m_s[:, None], b_s, 0) - - b_o += tl.dot(b_s.to(b_v.dtype), b_v, allow_tf32=False) - b_k2 = (tl.load(p_k2, boundary_check=(0, 1)) * b_beta[:, None]).to(b_v.dtype) - b_q -= tl.dot(b_s.to(b_v.dtype), b_k2, allow_tf32=False) - - if OUTPUT_ATTENTIONS: - p_a = tl.make_block_ptr(attn + i_bh * T * T, (T, T), (T, 1), (i_t * BT, offset), (BT, BS), (1, 0)) - tl.store(p_a, b_s.to(p_a.dtype.element_ty), boundary_check=(0, 1)) - - # Q block and K block have no overlap - # no need for mask, thereby saving flops - for offset in range(i_t * BT - BS, -BS, -BS): - p_k = tl.make_block_ptr(k + i_bh * T*K, (K, T), (1, K), (0, offset), (BK, BS), (0, 1)) - p_v = tl.make_block_ptr(v + i_bh * T*V, (T, V), (V, 1), (offset, 0), (BS, BV), (1, 0)) - p_beta = tl.make_block_ptr(beta + i_bh * T, (T, ), (1, ), (offset, ), (BS, ), (0,)) - p_k2 = tl.make_block_ptr(k2 + i_bh * T*K, (T, K), (K, 1), (offset, 0), (BS, BK), (1, 0)) - - # [BK, BS] - b_k = tl.load(p_k, boundary_check=(0, 1)) - # [BS, BV] - b_v = tl.load(p_v, boundary_check=(0, 1)) - # [BS] - b_beta = tl.load(p_beta, boundary_check=(0,)) - # [BT, BS] - b_s = (tl.dot(b_q.to(b_k.dtype), b_k, allow_tf32=False)) - # [BT, BV] - b_o += tl.dot(b_s.to(b_v.dtype), b_v, allow_tf32=False) - b_k2 = (tl.load(p_k2, boundary_check=(0, 1)) * b_beta[:, None]).to(b_v.dtype) - b_q -= tl.dot(b_s.to(b_v.dtype), b_k2, allow_tf32=False).to(b_q.dtype) - - if OUTPUT_ATTENTIONS: - p_a = tl.make_block_ptr(attn + i_bh * T * T, (T, T), (T, 1), (i_t * BT, offset), (BT, BS), (1, 0)) - tl.store(p_a, b_s.to(p_a.dtype.element_ty), boundary_check=(0, 1)) - - p_o_new = tl.make_block_ptr(o_new + i_bh * T*V, (T, V), (V, 1), (i_t*BT, 0), (BT, BV), (1, 0)) - tl.store(p_o_new, b_o.to(p_o.dtype.element_ty), boundary_check=(0, 1)) - - -class ParallelDeltaRuleFunction(torch.autograd.Function): - - @staticmethod - @input_guard - @autocast_custom_fwd - def forward(ctx, q, k, v, beta, scale, output_attentions): - B, H, T, K, V = *k.shape, v.shape[-1] - assert q.shape[-1] <= 128, 'The maximum supported sequence length is 128.' - BT, BS = 128, 32 - BK = triton.next_power_of_2(k.shape[-1]) - BV = triton.next_power_of_2(v.shape[-1]) - assert BT % BS == 0 - - A = fwd_prepare_T(k, beta, BS) - attn = q.new_zeros(B, H, T, T) if output_attentions else None - q_new, k_new, o, A_local = chunk_transform_qk_fwd( - q, - k, - v, - beta, - A, - scale, - BS, - output_attentions, - ) - - num_stages = 3 if K <= 64 else 2 - num_warps = 4 - grid = (triton.cdiv(T, BT), B * H) - o_new = torch.empty_like(o) - - parallel_delta_rule_fwd_kernel[grid]( - q=q_new, - k=k_new, - k2=k, - v=v, - beta=beta, - o=o, - o_new=o_new, - attn=attn, - T=T, - K=K, - V=V, - BT=BT, - BS=BS, - BK=BK, - BV=BV, - num_stages=num_stages, - num_warps=num_warps, - ) - - if output_attentions: - grid = (triton.cdiv(T, BS), B * H) - save_intra_chunk_attn[grid]( - A=attn, - A_local=A_local, - T=T, - BT=BS, - ) - return o_new.to(q.dtype), attn - - @staticmethod - @input_guard - @autocast_custom_bwd - def backward(ctx, do, d_attn=None): - raise NotImplementedError('Backward pass is not implemented. Stay tuned!') - - -def parallel_delta_rule( - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - beta: torch.Tensor, - scale: float | None = None, - output_attentions: bool = False, - **kwargs, -) -> tuple[torch.Tensor, torch.Tensor]: - r""" - Args: - q (torch.Tensor): - queries of shape `[B, T, H, K]`. - k (torch.Tensor): - keys of shape `[B, T, H, K]`. - v (torch.Tensor): - values of shape `[B, T, H, V]`. - beta (torch.Tensor): - betas of shape `[B, T, H]`. - scale (Optional[float]): - Scale factor for attention scores. - If not provided, it will default to `1 / sqrt(K)`. Default: `None`. - output_attentions (bool): - Whether to output the materialized attention scores of shape `[B, H, T, T]`. Default: `False`. - - Returns: - o (torch.Tensor): - Outputs of shape `[B, T, H, V]`. - attn (torch.Tensor): - Attention scores of shape `[B, H, T, T]` if `output_attentions=True` else `None`. - """ - if 'head_first' in kwargs: - raise DeprecationWarning( - "head_first has been removed. Inputs must be in `[B, T, H, ...]` format.", - ) - o, attn = ParallelDeltaRuleFunction.apply(q, k, v, beta, scale, output_attentions) - return o, attn - - -def naive_delta_rule_parallel(q, k, v, beta, BM=128, BN=32): - b, h, l, d_k = q.shape - q = q * (d_k ** -0.5) - v = v * beta[..., None] - k_beta = k * beta[..., None] - # compute (I - tri(diag(beta) KK^T))^{-1} - q, k, v, k_beta = map(lambda x: rearrange(x, 'b h (n c) d -> b h n c d', c=BN), [q, k, v, k_beta]) - mask = torch.triu(torch.ones(BN, BN, dtype=torch.bool, device=q.device), diagonal=0) - T = -(k_beta @ k.transpose(-1, -2)).masked_fill(mask, 0) - for i in range(1, BN): - T[..., i, :i] = T[..., i, :i].clone() + (T[..., i, :, None].clone() * T[..., :, :i].clone()).sum(-2) - T = T + torch.eye(BN, dtype=q.dtype, device=q.device) - - mask2 = torch.triu(torch.ones(BN, BN, dtype=torch.bool, device=q.device), diagonal=1) - A_local = (q @ k.transpose(-1, -2)).masked_fill(mask2, 0) @ T - o_intra = A_local @ v - - # apply cumprod transition matrices on k to the last position within the chunk - k = k - ((k @ k.transpose(-1, -2)).masked_fill(mask, 0) @ T).transpose(-1, -2) @ k_beta - # apply cumprod transition matrices on q to the first position within the chunk - q = q - A_local @ k_beta - o_intra = A_local @ v - - A = torch.zeros(b, h, l, l, device=q.device) - - q, k, v, k_beta, o_intra = map(lambda x: rearrange(x, 'b h n c d -> b h (n c) d'), [q, k, v, k_beta, o_intra]) - o = torch.empty_like(v) - for i in range(0, l, BM): - q_i = q[:, :, i:i+BM] - o_i = o_intra[:, :, i:i+BM] - # intra block - for j in range(i + BM - 2 * BN, i-BN, -BN): - k_j = k[:, :, j:j+BN] - A_ij = q_i @ k_j.transpose(-1, -2) - mask = torch.arange(i, i+BM) >= (j + BN) - A_ij = A_ij.masked_fill_(~mask[:, None].to(A_ij.device), 0) - A[:, :, i:i+BM, j:j+BN] = A_ij - q_i = q_i - A_ij @ k_beta[:, :, j:j+BN] - o_i += A_ij @ v[:, :, j:j+BN] - # inter block - for j in range(i - BN, -BN, -BN): - k_j = k[:, :, j:j+BN] - A_ij = q_i @ k_j.transpose(-1, -2) - A[:, :, i:i+BM, j:j+BN] = A_ij - q_i = q_i - A_ij @ k_beta[:, :, j:j+BN] - o_i += A_ij @ v[:, :, j:j+BN] - o[:, :, i:i+BM] = o_i - - for i in range(0, l//BN): - A[:, :, i*BN:i*BN+BN, i*BN:i*BN+BN] = A_local[:, :, i] - - return o, A diff --git a/build/torch-xpu/ops/delta_rule/wy_fast.py b/build/torch-xpu/ops/delta_rule/wy_fast.py deleted file mode 100644 index 52f7af2641052cee489eec5d53c8dcb9628dfd11..0000000000000000000000000000000000000000 --- a/build/torch-xpu/ops/delta_rule/wy_fast.py +++ /dev/null @@ -1,312 +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 -import triton.language as tl - -from ...ops.common.chunk_scaled_dot_kkt import chunk_scaled_dot_kkt_fwd -from ...ops.utils import prepare_chunk_indices -from ...ops.utils.op import safe_dot -from ...ops.utils.solve_tril import solve_tril -from ...utils import IS_NVIDIA_HOPPER, autotune_cache_kwargs, check_shared_mem - -# Triton miscompiles `prepare_wy_repr_bwd_kernel` with num_warps=4 on Hopper -# (sm_90) for BT=64: it produces incorrect dk/dbeta (and can raise an illegal -# memory access), while num_warps=2 is correct (see #984). Restrict the Hopper -# autotune configs to num_warps=2 until the upstream compiler issue is resolved. -NUM_WARPS = [2] if IS_NVIDIA_HOPPER else [2, 4, 8] - - -@triton.heuristics({ - 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None, -}) -@triton.autotune( - configs=[ - triton.Config({}, num_warps=num_warps, num_stages=num_stages) - for num_warps in [2, 4, 8] - for num_stages in [2, 3, 4] - ], - key=['H', 'K', 'V', 'BT', 'BK', 'BV', 'IS_VARLEN'], - **autotune_cache_kwargs, -) -@triton.jit(do_not_specialize=['T']) -def recompute_w_u_fwd_kernel( - k, - v, - beta, - w, - u, - A, - cu_seqlens, - chunk_indices, - T, - H: tl.constexpr, - K: tl.constexpr, - V: tl.constexpr, - BT: tl.constexpr, - BK: tl.constexpr, - BV: tl.constexpr, - IS_VARLEN: tl.constexpr, -): - i_t, i_bh = tl.program_id(0), tl.program_id(1) - i_b, i_h = i_bh // H, i_bh % H - 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.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) - T = eos - bos - else: - bos, eos = i_b * T, i_b * T + T - - p_beta = tl.make_block_ptr(beta + bos*H + i_h, (T,), (H,), (i_t * BT,), (BT,), (0,)) - p_A = tl.make_block_ptr(A + (bos*H + i_h) * BT, (T, BT), (H*BT, 1), (i_t * BT, 0), (BT, BT), (1, 0)) - b_beta = tl.load(p_beta, boundary_check=(0,)) - b_A = tl.load(p_A, boundary_check=(0, 1)) - - for i_v in range(tl.cdiv(V, BV)): - p_v = tl.make_block_ptr(v + (bos*H + i_h) * V, (T, V), (H*V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) - p_u = tl.make_block_ptr(u + (bos*H + i_h) * V, (T, V), (H*V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) - b_v = tl.load(p_v, boundary_check=(0, 1)) - b_vb = (b_v * b_beta[:, None]).to(b_v.dtype) - b_u = tl.dot(b_A.to(b_vb.dtype), b_vb, allow_tf32=False) - tl.store(p_u, (b_u).to(p_u.dtype.element_ty), boundary_check=(0, 1)) - - for i_k in range(tl.cdiv(K, BK)): - p_k = tl.make_block_ptr(k + (bos*H + i_h) * K, (T, K), (H*K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) - p_w = tl.make_block_ptr(w + (bos*H + i_h) * K, (T, K), (H*K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) - b_k = tl.load(p_k, boundary_check=(0, 1)) - b_kb = (b_k * b_beta[:, None]).to(b_k.dtype) - b_w = tl.dot(b_A.to(b_kb.dtype), b_kb, allow_tf32=False) - tl.store(p_w, b_w.to(p_w.dtype.element_ty), boundary_check=(0, 1)) - - -@triton.heuristics({ - 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None, -}) -@triton.autotune( - configs=[ - triton.Config({}, num_warps=num_warps, num_stages=num_stages) - for num_warps in NUM_WARPS - for num_stages in [2, 3, 4] - ], - key=['H', 'K', 'V', 'BT', 'BK', 'BV', 'IS_VARLEN'], - **autotune_cache_kwargs, -) -@triton.jit(do_not_specialize=['T']) -def prepare_wy_repr_bwd_kernel( - k, - v, - beta, - A, - dw, - du, - dk, - dv, - dbeta, - cu_seqlens, - chunk_indices, - T, - H: tl.constexpr, - K: tl.constexpr, - V: tl.constexpr, - BT: tl.constexpr, - BK: tl.constexpr, - BV: tl.constexpr, - IS_VARLEN: tl.constexpr, -): - i_t, i_bh = tl.program_id(0), tl.program_id(1), tl.program_id(2) - i_b, i_h = i_bh // H, i_bh % H - 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.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) - T = eos - bos - else: - bos, eos = i_b * T, i_b * T + T - - p_beta = tl.make_block_ptr(beta + bos*H + i_h, (T,), (H,), (i_t * BT,), (BT,), (0,)) - p_A = tl.make_block_ptr(A + (bos*H + i_h) * BT, (BT, T), (1, H*BT), (0, i_t * BT), (BT, BT), (0, 1)) - - b_beta = tl.load(p_beta, boundary_check=(0,)) - b_A = tl.load(p_A, boundary_check=(0, 1)) - - b_dbeta = tl.zeros([BT], dtype=tl.float32) - b_dA = tl.zeros([BT, BT], dtype=tl.float32) - for i_v in range(tl.cdiv(V, BV)): - p_v = tl.make_block_ptr(v + (bos*H + i_h) * V, (T, V), (H*V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) - p_dv = tl.make_block_ptr(dv + (bos*H + i_h) * V, (T, V), (H*V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) - p_du = tl.make_block_ptr(du + (bos*H + i_h) * V, (T, V), (H*V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) - - b_v = tl.load(p_v, boundary_check=(0, 1)) - b_v_beta = (b_v * b_beta[:, None]).to(b_v.dtype) - b_du = tl.load(p_du, boundary_check=(0, 1)) - b_dA += tl.dot(b_du, tl.trans(b_v_beta), allow_tf32=False) - b_dv_beta = tl.dot(b_A, b_du, allow_tf32=False) - b_dv = b_dv_beta * b_beta[:, None] - b_dbeta += tl.sum(b_dv_beta * b_v, 1) - - tl.store(p_dv, b_dv.to(p_dv.dtype.element_ty), boundary_check=(0, 1)) - - for i_k in range(tl.cdiv(K, BK)): - p_k = tl.make_block_ptr(k + (bos*H + i_h) * K, (T, K), (H*K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) - p_dk = tl.make_block_ptr(dk + (bos*H + i_h) * K, (T, K), (H*K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) - p_dw = tl.make_block_ptr(dw + (bos*H + i_h) * K, (T, K), (H*K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) - b_k = tl.load(p_k, boundary_check=(0, 1)) - b_k_beta = (b_k * b_beta[:, None]).to(b_k.dtype) - b_dw = tl.load(p_dw, boundary_check=(0, 1)) - b_dA += tl.dot(b_dw, tl.trans(b_k_beta), allow_tf32=False) - b_dk_beta = tl.dot(b_A, b_dw, allow_tf32=False) - b_dk = b_dk_beta * b_beta[:, None] - b_dbeta += tl.sum(b_dk_beta * b_k, 1) - - tl.store(p_dk, b_dk.to(p_dk.dtype.element_ty), boundary_check=(0, 1)) - - b_dA = tl.where(tl.arange(0, BT)[:, None] > tl.arange(0, BT)[None, :], b_dA, 0) - b_dA = tl.dot(b_dA.to(b_A.dtype), b_A) - b_dA = tl.dot(b_A, b_dA.to(b_A.dtype)) - b_dA = tl.where(tl.arange(0, BT)[:, None] > tl.arange(0, BT)[None, :], -b_dA, 0).to(k.dtype.element_ty) - - for i_k in range(tl.cdiv(K, BK)): - p_k = tl.make_block_ptr(k + (bos*H + i_h) * K, (T, K), (H*K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) - p_dk = tl.make_block_ptr(dk + (bos*H + i_h) * K, (T, K), (H*K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) - b_k = tl.load(p_k, boundary_check=(0, 1)) - b_dk = tl.load(p_dk, boundary_check=(0, 1)) - b_k_beta = (b_k * b_beta[:, None]).to(b_k.dtype) - - b_dk_beta = tl.dot(b_dA, b_k, allow_tf32=False) - b_dbeta += tl.sum(b_dk_beta * b_k, 1) - b_dk += safe_dot(tl.trans(b_dA), b_k_beta, allow_tf32=False) - b_dk += b_dk_beta * b_beta[:, None] - tl.store(p_dk, b_dk.to(p_dk.dtype.element_ty), boundary_check=(0, 1)) - - p_dbeta = tl.make_block_ptr(dbeta + bos*H + i_h, (T,), (H,), (i_t * BT,), (BT,), (0,)) - tl.store(p_dbeta, b_dbeta.to(p_dbeta.dtype.element_ty), boundary_check=(0,)) - - -def prepare_wy_repr_fwd( - k: torch.Tensor, - v: torch.Tensor, - beta: torch.Tensor, - cu_seqlens: torch.LongTensor | None, - chunk_indices: torch.LongTensor | None = None, - chunk_size: int = 64, -) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: - A = chunk_scaled_dot_kkt_fwd( - k=k, - beta=beta, - cu_seqlens=cu_seqlens, - chunk_size=chunk_size, - output_dtype=torch.float32, - chunk_indices=chunk_indices, - ) - A = solve_tril( - A=A, - cu_seqlens=cu_seqlens, - chunk_indices=chunk_indices, - output_dtype=k.dtype, - ) - w, u = recompute_w_u_fwd( - k=k, - v=v, - beta=beta, - A=A, - cu_seqlens=cu_seqlens, - chunk_indices=chunk_indices, - ) - return w, u, A - - -def recompute_w_u_fwd( - k: torch.Tensor, - v: torch.Tensor, - beta: torch.Tensor, - A: torch.Tensor, - cu_seqlens: torch.LongTensor | None, - chunk_indices: torch.LongTensor | None = None, -) -> tuple[torch.Tensor, torch.Tensor]: - B, T, H, K, V = *k.shape, v.shape[-1] - BT = A.shape[-1] - CONST_TILING = 64 if check_shared_mem() else 32 - BK = min(max(triton.next_power_of_2(K), 16), CONST_TILING) - BV = min(max(triton.next_power_of_2(V), 16), CONST_TILING) - - if chunk_indices is None and cu_seqlens is not None: - chunk_indices = prepare_chunk_indices(cu_seqlens, BT) - NT = triton.cdiv(T, BT) if cu_seqlens is None else len(chunk_indices) - - u = torch.empty_like(v) - w = torch.empty_like(k) - recompute_w_u_fwd_kernel[(NT, B*H)]( - k, - v, - beta, - w, - u, - A, - cu_seqlens=cu_seqlens, - chunk_indices=chunk_indices, - T=T, - H=H, - K=K, - V=V, - BT=BT, - BK=BK, - BV=BV, - ) - return w, u - - -def prepare_wy_repr_bwd( - k: torch.Tensor, - v: torch.Tensor, - beta: torch.Tensor, - A: torch.Tensor, - dw: torch.Tensor, - du: torch.Tensor, - cu_seqlens: torch.LongTensor | None, - chunk_indices: torch.LongTensor | None = None, -) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: - B, T, H, K, V = *k.shape, v.shape[-1] - BT = A.shape[-1] - CONST_TILING = 64 if check_shared_mem() else 32 - BK = min(max(triton.next_power_of_2(K), 16), CONST_TILING) - BV = min(max(triton.next_power_of_2(V), 16), CONST_TILING) - - if chunk_indices is None and cu_seqlens is not None: - chunk_indices = prepare_chunk_indices(cu_seqlens, BT) - NT = triton.cdiv(T, BT) if cu_seqlens is None else len(chunk_indices) - - dk = torch.empty_like(k) - dv = torch.empty_like(v) - dbeta = torch.empty_like(beta) - prepare_wy_repr_bwd_kernel[(NT, B * H)]( - k, - v, - beta, - A, - dw, - du, - dk, - dv, - dbeta, - cu_seqlens=cu_seqlens, - chunk_indices=chunk_indices, - T=T, - H=H, - K=K, - V=V, - BT=BT, - BK=BK, - BV=BV, - ) - return dk, dv, dbeta - - -fwd_prepare_wy_repr = prepare_wy_repr_fwd - -bwd_prepare_wy_repr = prepare_wy_repr_bwd - -fwd_recompute_w_u = recompute_w_u_fwd diff --git a/build/torch-xpu/ops/deltaformer/__init__.py b/build/torch-xpu/ops/deltaformer/__init__.py deleted file mode 100644 index a993bbb5c532fbeca00f5a285651789ed4a01dc5..0000000000000000000000000000000000000000 --- a/build/torch-xpu/ops/deltaformer/__init__.py +++ /dev/null @@ -1,14 +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 .naive import naive_deltaformer_attn -from .parallel import deltaformer_attn - -__all__ = [ - 'deltaformer_attn', - 'naive_deltaformer_attn', -] diff --git a/build/torch-xpu/ops/deltaformer/invcum.py b/build/torch-xpu/ops/deltaformer/invcum.py deleted file mode 100644 index 537f8a89bb18ae3ba0d2254f41cf27cc277e803a..0000000000000000000000000000000000000000 --- a/build/torch-xpu/ops/deltaformer/invcum.py +++ /dev/null @@ -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 - -import torch - - -def forward(u, w): - return torch.linalg.solve_triangular( - w.float(), - u.float(), - upper=False, - unitriangular=True, - ).to(u.dtype) - - -def forward_inplace(u, w): - u.copy_(forward(u, w)) - - -def backward_x(do, w): - return torch.linalg.solve_triangular( - w.tril(-1).mH.float(), - do.float(), - upper=True, - unitriangular=True, - ).to(do.dtype) - - -def backward(do, w, x): - du = torch.linalg.solve_triangular( - w.tril(-1).mH.float(), - do.float(), - upper=True, - unitriangular=True, - ).to(do.dtype) - dw = torch.bmm(-du, x.mH) - dw = dw.tril(-1) - return du, dw diff --git a/build/torch-xpu/ops/deltaformer/naive.py b/build/torch-xpu/ops/deltaformer/naive.py deleted file mode 100644 index e1398c31e4d52c53aeb5e62e5091ff23110a6e16..0000000000000000000000000000000000000000 --- a/build/torch-xpu/ops/deltaformer/naive.py +++ /dev/null @@ -1,155 +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 math - -import torch - - -def tril_softmax(scores: torch.Tensor, strict: bool = True) -> torch.Tensor: - """ - Row-wise causal softmax over strictly lower-triangular (j < i) positions. - - Args: - scores: [B, H, T, T] raw attention scores (q @ k^T). - strict: if True, mask out diagonal as well (strictly causal). Otherwise include diagonal. - - Returns: - probs: [B, H, T, T] with probabilities on j < i (or j <= i if strict=False), zeros elsewhere. - """ - T = scores.size(-1) - device = scores.device - i = torch.arange(T, device=device).view(1, 1, T, 1) - j = torch.arange(T, device=device).view(1, 1, 1, T) - if strict: - mask = (j < i) - else: - mask = (j <= i) - - masked = scores.masked_fill(~mask, float('-inf')) - max_per_row = masked.max(dim=-1, keepdim=True).values - exp = (masked - max_per_row).exp() - exp = exp.masked_fill(~mask, 0.0) - denom = exp.sum(dim=-1, keepdim=True).clamp_min_(1e-20) - probs = exp / denom - return probs - - -def naive_causal_attention_bhtd( - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, -) -> torch.Tensor: - B, H, T, D = q.shape - qk_scale = 1.0 / math.sqrt(D) - scores = torch.matmul(q, k.transpose(-1, -2)) * qk_scale # [B, H, T, T] - causal_mask = torch.triu(torch.ones(T, T, device=q.device), diagonal=1).bool() - scores = scores.masked_fill(causal_mask, float('-inf')) - attn_weights = torch.softmax(scores, dim=-1) # [B, H, T, T] - o = torch.matmul(attn_weights, v) # [B, H, T, D] - - return o - - -def naive_deltaformer_attn_head_first( - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - beta: torch.Tensor | None = None, -) -> torch.Tensor: - """ - Naive reference implementation of DeltaFormer attention for head-first format. - - Two-stage process: - 1. Computes u[i] = v[i] - beta[i] * sum_{j torch.Tensor: - """ - Naive reference implementation of DeltaFormer attention for sequence-first format. - - Args: - q: [B, T, H, D] - k: [B, T, H, D] - v: [B, T, H, D] - beta: [B, T, H] or None (defaults to ones) - - Returns: - o: [B, T, H, D] - """ - assert q.dim() == 4 and k.dim() == 4 and v.dim() == 4, "q,k,v must be [B,T,H,D]" - B, T, H, D = q.shape - assert k.shape == (B, T, H, D) and v.shape == (B, T, H, D) - - q_bhtd = q.transpose(1, 2) # [B, T, H, D] -> [B, H, T, D] - k_bhtd = k.transpose(1, 2) # [B, T, H, D] -> [B, H, T, D] - v_bhtd = v.transpose(1, 2) # [B, T, H, D] -> [B, H, T, D] - - if beta is not None: - assert beta.shape == (B, T, H) - beta_bhtd = beta.transpose(1, 2) # [B, T, H] -> [B, H, T] - else: - beta_bhtd = None - - o_bhtd = naive_deltaformer_attn_head_first(q_bhtd, k_bhtd, v_bhtd, beta_bhtd) - - o_bthd = o_bhtd.transpose(1, 2) # [B, H, T, D] -> [B, T, H, D] - - return o_bthd - - -__all__ = [ - 'naive_deltaformer_attn', - 'tril_softmax', -] diff --git a/build/torch-xpu/ops/deltaformer/parallel.py b/build/torch-xpu/ops/deltaformer/parallel.py deleted file mode 100644 index eb60bee81d5b82cd666d2fef0913dc9f5a0c1cea..0000000000000000000000000000000000000000 --- a/build/torch-xpu/ops/deltaformer/parallel.py +++ /dev/null @@ -1,1002 +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 math -import warnings - -import torch -import triton -import triton.language as tl - -from . import invcum - -try: - from flash_attn import flash_attn_func, flash_attn_varlen_func -except ImportError: - warnings.warn( - "Flash Attention is not installed. Please install it via `pip install flash-attn --no-build-isolation`", - category=ImportWarning, - ) - flash_attn_func = None - -from ...utils.layers import pad_input, unpad_input - -BLOCK_SIZE_C = 512 - - -def parallel_deltaformer_chunk_fwd( - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - u: torch.Tensor, - qk_scale: float, - beta: torch.Tensor, -): - C, H, D = q.size() - T, _H, _D = k.size() - __C, __H = beta.size() - assert H == _H and D == _D and H == __H and __C == C - w = torch.empty(C, H, C, device=q.device, dtype=q.dtype) - lse = torch.empty(C, H, device=q.device, dtype=torch.float) - parallel_deltaformer_kernel(q, k, v, u, w, lse, qk_scale, beta) - return w, lse - - -def parallel_deltaformer_bwd_u_chunk( - q: torch.Tensor, - k: torch.Tensor, - lse: torch.Tensor, - grad_v: torch.Tensor, - fa_scale: float, - beta: torch.Tensor, -): - C, H, D = q.size() - T, _H, _D = k.size() - grad_u = torch.empty_like(q) - - def grid(META): - return (triton.cdiv(C, META['BLOCK_C']), H) - - parallel_deltaformer_bwd_kernel_u[grid]( - grad_u, q, k, grad_v, lse, beta, - H, T, C, D, fa_scale, - ) - return grad_u - - -def parallel_deltaformer_bwd_qk( - q: torch.Tensor, - k: torch.Tensor, - u: torch.Tensor, - lse: torch.Tensor, - grad_v: torch.Tensor, - qk_scale: float, - fa_scale: float, - beta: torch.Tensor, -): - T, H, D = k.size() - row_dot_sum = torch.empty_like(lse) - - def grid_bp(META): - return (triton.cdiv(T, META['BLOCK_C']), H) - - parallel_deltaformer_bwd_kernel_row_sum[grid_bp]( - row_dot_sum, q, k, grad_v, u, lse, - H, T, D, - fa_scale, - ) - grad_k = torch.empty_like(k) - grad_q = torch.empty_like(q) - - parallel_deltaformer_bwd_kernel_qk[grid_bp]( - grad_q, grad_k, q, k, grad_v, u, lse, beta, row_dot_sum, - H, T, D, - fa_scale, qk_scale, - ) - return grad_q, grad_k, row_dot_sum - - -def parallel_deltaformer_kernel( - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - u: torch.Tensor, - w: torch.Tensor, - lse: torch.Tensor, - qk_scale: float, - beta: torch.Tensor, -) -> None: - C, H, D = q.size() - T, _H, _D = k.size() - - def grid(META): - return (triton.cdiv(C, META['BLOCK_C']), H) - - parallel_deltaformer_fwd_kernel[grid]( - q, k, v, u, w, lse, beta, - H, T, C, D, qk_scale, - ) - - -def _config_deltaformer(): - return [ - triton.Config({'BLOCK_C': BC, 'BLOCK_T': BT}, num_stages=ns, num_warps=nw) - for BC in [128, 64] - for BT in [64, 32] - for ns in [3, 2] - for nw in [8, 4] - ] - - -@triton.autotune(configs=_config_deltaformer(), key=['C', 'D']) -@triton.jit -def parallel_deltaformer_fwd_kernel( - q_ptr, - k_ptr, - v_ptr, - u_ptr, - w_ptr, - lse_ptr, - beta_ptr, - H, - T, - C, - D: tl.constexpr, - qk_scale: float, - BLOCK_C: tl.constexpr, - BLOCK_T: tl.constexpr, -): - pid_c = tl.program_id(axis=0) - pid_h = tl.program_id(axis=1) - - rowid_block = tl.arange(0, BLOCK_C) + pid_c * BLOCK_C - colid_block = tl.arange(0, BLOCK_T) - - rowmax = tl.zeros([BLOCK_C], dtype=tl.float32) - float('inf') - rowsum = tl.zeros([BLOCK_C], dtype=tl.float32) + 1 - acc = tl.zeros([BLOCK_C, D], dtype=tl.float32) - - q_blk_ptr = tl.make_block_ptr( - base=q_ptr + pid_h * D, - shape=(C, D), - strides=(H * D, 1), - offsets=(pid_c * BLOCK_C, 0), - block_shape=(BLOCK_C, D), - order=(1, 0), - ) - q = tl.load(q_blk_ptr, boundary_check=(0,)) - - for kv_i in range(0, T, BLOCK_T): - k_blk_ptr = tl.make_block_ptr( - base=k_ptr + pid_h * D, - shape=(D, T), - strides=(1, H * D), - offsets=(0, kv_i), - block_shape=(D, BLOCK_T), - order=(0, 1), - ) - k = tl.load(k_blk_ptr, boundary_check=(1,)) - qk = tl.dot(q, k) * qk_scale - - if kv_i >= T - C: - mask = (T - C - kv_i + rowid_block[:, None] - colid_block[None, :] < 1) - qk = tl.where(mask, -1e6, qk) - - rowmax_i = tl.maximum(rowmax, tl.max(qk, axis=1)) - qk -= rowmax_i[:, None] - p = tl.math.exp2(qk) - - rowsum_i = tl.sum(p, axis=1) - alpha = tl.math.exp2(rowmax - rowmax_i) - rowsum = rowsum * alpha + rowsum_i - acc = acc * alpha[:, None] - rowmax = rowmax_i - - if kv_i < T - C: - u_blk_ptr = tl.make_block_ptr( - base=u_ptr + pid_h * D, - shape=(T, D), - strides=(H * D, 1), - offsets=(kv_i, 0), - block_shape=(BLOCK_T, D), - order=(1, 0), - ) - u = tl.load(u_blk_ptr, boundary_check=(0,)) - acc = tl.dot(p.to(u_ptr.dtype.element_ty), u, acc) - - lse = rowmax + tl.math.log2(rowsum) - lse_block_ptr = lse_ptr + pid_h + rowid_block * H - lse_mask = rowid_block < C - tl.store(lse_block_ptr, lse, mask=lse_mask) - - v_ptr = tl.make_block_ptr( - base=v_ptr + pid_h * D, - shape=(C, D), - strides=(H * D, 1), - offsets=(pid_c * BLOCK_C, 0), - block_shape=(BLOCK_C, D), - order=(1, 0), - ) - acc = acc / rowsum[:, None] - - beta_ptr = tl.make_block_ptr( - base=beta_ptr + pid_h, - shape=(C,), - strides=(H,), - offsets=(pid_c * BLOCK_C,), - block_shape=(BLOCK_C,), - order=(0,), - ) - beta = tl.load(beta_ptr, boundary_check=(0,)) - acc = acc * beta[:, None] - - v = tl.load(v_ptr, boundary_check=(0,)) - u = v - acc.to(v_ptr.dtype.element_ty) - u_block_ptr = tl.make_block_ptr( - base=u_ptr + pid_h * D, - shape=(T, D), - strides=(H * D, 1), - offsets=(T - C + pid_c * BLOCK_C, 0), - block_shape=(BLOCK_C, D), - order=(1, 0), - ) - tl.store(u_block_ptr, u, boundary_check=(0, 1)) - - for kv_i in range(T - C, T, BLOCK_T): - k_blk_ptr = tl.make_block_ptr( - base=k_ptr + pid_h * D, - shape=(D, T), - strides=(1, H * D), - offsets=(0, kv_i), - block_shape=(D, BLOCK_T), - order=(0, 1), - ) - k = tl.load(k_blk_ptr, boundary_check=(1,)) - qk = tl.dot(q, k) * qk_scale - - mask = (T - C - kv_i + rowid_block[:, None] - colid_block[None, :] < 1) - qk -= rowmax[:, None] - p = tl.math.exp2(qk) / rowsum[:, None] - p = tl.where(mask, 0, p) - w_blk_ptr = tl.make_block_ptr( - base=w_ptr + pid_h * C, - shape=(C, C), - strides=(H * C, 1), - offsets=(pid_c * BLOCK_C, kv_i - (T - C)), - block_shape=(BLOCK_C, BLOCK_T), - order=(1, 0), - ) - tl.store(w_blk_ptr, p.to(w_ptr.dtype.element_ty), boundary_check=(0, 1)) - - -@triton.autotune(configs=_config_deltaformer(), key=['C', 'D']) -@triton.jit -def parallel_deltaformer_bwd_kernel_u( - o_ptr, - q_ptr, - k_ptr, - v_ptr, - lse_ptr, - beta_ptr, - H, - T, - C, - D: tl.constexpr, - fa_scale, - BLOCK_C: tl.constexpr, - BLOCK_T: tl.constexpr, -): - pid_c = tl.program_id(axis=0) - pid_h = tl.program_id(axis=1) - - acc = tl.zeros([BLOCK_C, D], dtype=tl.float32) - - q_blk_ptr = tl.make_block_ptr( - base=q_ptr + pid_h * D, - shape=(C, D), - strides=(H * D, 1), - offsets=(pid_c * BLOCK_C, 0), - block_shape=(BLOCK_C, D), - order=(1, 0), - ) - q = tl.load(q_blk_ptr, boundary_check=(0,)) - - for kv_i in range(0, T, BLOCK_T): - k_blk_ptr = tl.make_block_ptr( - base=k_ptr + pid_h * D, - shape=(D, T), - strides=(1, H * D), - offsets=(0, kv_i), - block_shape=(D, BLOCK_T), - order=(0, 1), - ) - k = tl.load(k_blk_ptr, boundary_check=(1,)) - qk = tl.dot(q, k) * fa_scale - - lse_blk_ptr = tl.make_block_ptr( - base=lse_ptr + pid_h, - shape=(T,), - strides=(H,), - offsets=(kv_i,), - block_shape=(BLOCK_T,), - order=(0,), - ) - lse = tl.load(lse_blk_ptr, boundary_check=(0,)) - beta_blk_ptr = tl.make_block_ptr( - base=beta_ptr + pid_h, - shape=(T,), - strides=(H,), - offsets=(kv_i,), - block_shape=(BLOCK_T,), - order=(0,), - ) - beta = tl.load(beta_blk_ptr, boundary_check=(0,)) - - p = tl.math.exp2(qk - lse[None, :]) * beta[None, :] - - v_blk_ptr = tl.make_block_ptr( - base=v_ptr + pid_h * D, - shape=(T, D), - strides=(H * D, 1), - offsets=(kv_i, 0), - block_shape=(BLOCK_T, D), - order=(1, 0), - ) - v = tl.load(v_blk_ptr, boundary_check=(0,)) - acc = tl.dot(p.to(v_ptr.dtype.element_ty), v, acc) - - o_blk_ptr = tl.make_block_ptr( - base=o_ptr + pid_h * D, - shape=(C, D), - strides=(H * D, 1), - offsets=(pid_c * BLOCK_C, 0), - block_shape=(BLOCK_C, D), - order=(1, 0), - ) - tl.store(o_blk_ptr, acc.to(o_ptr.dtype.element_ty), boundary_check=(0,)) - - -@triton.autotune(configs=_config_deltaformer(), key=['T', 'D']) -@triton.jit -def parallel_deltaformer_bwd_kernel_row_sum( - row_dot_ptr, - q_ptr, - k_ptr, - grad_v_ptr, - u_ptr, - lse_ptr, - H, - T, - D: tl.constexpr, - fa_scale, - BLOCK_C: tl.constexpr, - BLOCK_T: tl.constexpr, -): - pid_c = tl.program_id(axis=0) - pid_h = tl.program_id(axis=1) - - rowid_block = tl.arange(0, BLOCK_C) + pid_c * BLOCK_C - colid_block = tl.arange(0, BLOCK_T) - - acc = tl.zeros([BLOCK_C], dtype=tl.float32) - - k_row_blk_ptr = tl.make_block_ptr( - base=q_ptr + pid_h * D, - shape=(T, D), - strides=(H * D, 1), - offsets=(pid_c * BLOCK_C, 0), - block_shape=(BLOCK_C, D), - order=(1, 0), - ) - k_row = tl.load(k_row_blk_ptr, boundary_check=(0,)) - lse_blk_ptr = tl.make_block_ptr( - base=lse_ptr + pid_h, - shape=(T,), - strides=(H,), - offsets=(pid_c * BLOCK_C,), - block_shape=(BLOCK_C,), - order=(0,), - ) - lse = tl.load(lse_blk_ptr, boundary_check=(0,)) - grad_v_blk_ptr = tl.make_block_ptr( - base=grad_v_ptr + pid_h * D, - shape=(T, D), - strides=(H * D, 1), - offsets=(pid_c * BLOCK_C, 0), - block_shape=(BLOCK_C, D), - order=(1, 0), - ) - grad_v_row = -tl.load(grad_v_blk_ptr, boundary_check=(0,)) - - for kv_i in range(0, (pid_c + 1) * BLOCK_C, BLOCK_T): - k_blk_ptr = tl.make_block_ptr( - base=k_ptr + pid_h * D, - shape=(D, T), - strides=(1, H * D), - offsets=(0, kv_i), - block_shape=(D, BLOCK_T), - order=(0, 1), - ) - k = tl.load(k_blk_ptr, boundary_check=(1,)) - qk = tl.dot(k_row, k) * fa_scale - p = tl.math.exp2(qk - lse[:, None]) - - u_blk_ptr = tl.make_block_ptr( - base=u_ptr + pid_h * D, - shape=(D, T), - strides=(1, H * D), - offsets=(0, kv_i), - block_shape=(D, BLOCK_T), - order=(0, 1), - ) - ut = tl.load(u_blk_ptr, boundary_check=(1,)) - dp = tl.dot(grad_v_row, ut) - if kv_i + BLOCK_T >= pid_c * BLOCK_C: - mask = (rowid_block[:, None] <= colid_block[None, :] + kv_i) - p = tl.where(mask, 0., p) - dp = tl.where(mask, 0., dp) - acc += tl.sum(p * dp, axis=1) - row_dot_block_ptr = tl.make_block_ptr( - base=row_dot_ptr + pid_h, - shape=(T,), - strides=(H,), - offsets=(pid_c * BLOCK_C,), - block_shape=(BLOCK_C,), - order=(0,), - ) - tl.store(row_dot_block_ptr, acc, boundary_check=(0,)) - - -@triton.autotune(configs=[triton.Config({'BLOCK_C': BC}, num_stages=ns, num_warps=nw) - for BC in [64, 32] - for ns in [4, 3] - for nw in [4]], key=['T', 'D']) -@triton.jit -def parallel_deltaformer_bwd_kernel_qk( - grad_q_ptr, - grad_k_ptr, - q_ptr, - k_ptr, - grad_v_ptr, - u_ptr, - lse_ptr, - beta_ptr, - row_dot_ptr, - H, - T, - D: tl.constexpr, - fa_scale: tl.constexpr, - qk_scale: tl.constexpr, - BLOCK_C: tl.constexpr, -): - pid_c = tl.program_id(axis=0) - pid_h = tl.program_id(axis=1) - block_i = tl.arange(0, BLOCK_C) - - acc = tl.zeros([BLOCK_C, D], dtype=tl.float32) - - k_row_blk_ptr = tl.make_block_ptr( - base=q_ptr + pid_h * D, - shape=(T, D), - strides=(H * D, 1), - offsets=(pid_c * BLOCK_C, 0), - block_shape=(BLOCK_C, D), - order=(1, 0), - ) - k_row = tl.load(k_row_blk_ptr, boundary_check=(0,)) - lse_blk_ptr = tl.make_block_ptr( - base=lse_ptr + pid_h, - shape=(T,), - strides=(H,), - offsets=(pid_c * BLOCK_C,), - block_shape=(BLOCK_C,), - order=(0,), - ) - lse = tl.load(lse_blk_ptr, boundary_check=(0,)) - beta_blk_ptr = tl.make_block_ptr( - base=beta_ptr + pid_h, - shape=(T,), - strides=(H,), - offsets=(pid_c * BLOCK_C,), - block_shape=(BLOCK_C,), - order=(0,), - ) - beta = tl.load(beta_blk_ptr, boundary_check=(0,)) - grad_v_blk_ptr = tl.make_block_ptr( - base=grad_v_ptr + pid_h * D, - shape=(T, D), - strides=(H * D, 1), - offsets=(pid_c * BLOCK_C, 0), - block_shape=(BLOCK_C, D), - order=(1, 0), - ) - grad_v_row = -tl.load(grad_v_blk_ptr, boundary_check=(0,)) - row_dot_blk_ptr = tl.make_block_ptr( - base=row_dot_ptr + pid_h, - shape=(T,), - strides=(H,), - offsets=(pid_c * BLOCK_C,), - block_shape=(BLOCK_C,), - order=(0,), - ) - row_dot_row = tl.load(row_dot_blk_ptr, boundary_check=(0,)).to(k_ptr.dtype.element_ty) - - for kv_i in range(0, pid_c * BLOCK_C, BLOCK_C): - k_blk_ptr = tl.make_block_ptr( - base=k_ptr + pid_h * D, - shape=(D, T), - strides=(1, H * D), - offsets=(0, kv_i), - block_shape=(D, BLOCK_C), - order=(0, 1), - ) - kt = tl.load(k_blk_ptr, boundary_check=(1,)) - qk = tl.dot(k_row, kt) * fa_scale - p = tl.math.exp2(qk - lse[:, None]) * beta[:, None] - - u_blk_ptr = tl.make_block_ptr( - base=u_ptr + pid_h * D, - shape=(D, T), - strides=(1, H * D), - offsets=(0, kv_i), - block_shape=(D, BLOCK_C), - order=(0, 1), - ) - ut = tl.load(u_blk_ptr) - dp = tl.dot(grad_v_row, ut) - da = p * (dp - row_dot_row[:, None]) - k = tl.trans(kt, 1, 0) - acc = tl.dot(da.to(k.dtype), k, acc) - - k_row_blk_ptr = tl.make_block_ptr( - base=k_ptr + pid_h * D, - shape=(T, D), - strides=(H * D, 1), - offsets=(pid_c * BLOCK_C, 0), - block_shape=(BLOCK_C, D), - order=(1, 0), - ) - k_row_true = tl.load(k_row_blk_ptr, boundary_check=(0,)) - qk = tl.dot(k_row, tl.trans(k_row_true, 1, 0)) * fa_scale - p = tl.math.exp2(qk - lse[:, None]) * beta[:, None] - u_blk_ptr = tl.make_block_ptr( - base=u_ptr + pid_h * D, - shape=(D, T), - strides=(1, H * D), - offsets=(0, pid_c * BLOCK_C), - block_shape=(D, BLOCK_C), - order=(0, 1), - ) - ut = tl.load(u_blk_ptr) - dp = tl.dot(grad_v_row, ut) - dpm = dp - row_dot_row[:, None] - mask = block_i[None, :] < block_i[:, None] - p = tl.where(mask, p, 0.) - dpm = tl.where(mask, dpm, 0.) - da = p * dpm - daat = da - acc = tl.dot(daat.to(k_row.dtype), k_row_true, acc) - - grad_q_blk_ptr = tl.make_block_ptr( - base=grad_q_ptr + pid_h * D, - shape=(T, D), - strides=(H * D, 1), - offsets=(BLOCK_C * pid_c, 0), - block_shape=(BLOCK_C, D), - order=(1, 0), - ) - acc = acc * qk_scale - tl.store(grad_q_blk_ptr, acc.to(grad_q_ptr.dtype.element_ty), boundary_check=(0,)) - - daat = tl.trans(da, 1, 0) - acc = tl.dot(daat.to(k_row.dtype), k_row) - k_row = k_row_true - nu = -tl.trans(ut, 1, 0) - for kv_i in range((pid_c + 1) * BLOCK_C, T, BLOCK_C): - k_blk_ptr = tl.make_block_ptr( - base=q_ptr + pid_h * D, - shape=(D, T), - strides=(1, H * D), - offsets=(0, kv_i), - block_shape=(D, BLOCK_C), - order=(0, 1), - ) - kt = tl.load(k_blk_ptr, boundary_check=(1,)) - lse_blk_ptr = tl.make_block_ptr( - base=lse_ptr + pid_h, - shape=(T,), - strides=(H,), - offsets=(kv_i,), - block_shape=(BLOCK_C,), - order=(0,), - ) - lse = tl.load(lse_blk_ptr, boundary_check=(0,)) - beta_blk_ptr = tl.make_block_ptr( - base=beta_ptr + pid_h, - shape=(T,), - strides=(H,), - offsets=(kv_i,), - block_shape=(BLOCK_C,), - order=(0,), - ) - beta = tl.load(beta_blk_ptr, boundary_check=(0,)) - qk = tl.dot(k_row, kt) * fa_scale - p = tl.math.exp2(qk - lse[None, :]) * beta[None, :] - - grad_vt_blk_ptr = tl.make_block_ptr( - base=grad_v_ptr + pid_h * D, - shape=(D, T), - strides=(1, H * D), - offsets=(0, kv_i), - block_shape=(D, BLOCK_C), - order=(0, 1), - ) - grad_vt = tl.load(grad_vt_blk_ptr, boundary_check=(1,)) - row_dot_blk_ptr = tl.make_block_ptr( - base=row_dot_ptr + pid_h, - shape=(T,), - strides=(H,), - offsets=(kv_i,), - block_shape=(BLOCK_C,), - order=(0,), - ) - row_dot = tl.load(row_dot_blk_ptr, boundary_check=(0,)).to(k_ptr.dtype.element_ty) - dp = tl.dot(nu, grad_vt) - da = p * (dp - row_dot[None, :]) - k = tl.trans(kt, 1, 0) - acc = tl.dot(da.to(k.dtype), k, acc) - - grad_k_blk_ptr = tl.make_block_ptr( - base=grad_k_ptr + pid_h * D, - shape=(T, D), - strides=(H * D, 1), - offsets=(BLOCK_C * pid_c, 0), - block_shape=(BLOCK_C, D), - order=(1, 0), - ) - acc = acc * qk_scale - tl.store(grad_k_blk_ptr, acc.to(grad_k_ptr.dtype.element_ty), boundary_check=(0,)) - - -class ParallelDeltaformerFunction(torch.autograd.Function): - @staticmethod - def forward( - ctx, - qo: torch.Tensor, - ko: torch.Tensor, - vo: torch.Tensor, - betao: torch.Tensor | None = None, - C: int = BLOCK_SIZE_C, - cu_seqlens: torch.LongTensor | None = None, - ): - B, T, H, D = ko.size() - C = min(C, T) - ctx.C = C - ctx.cu_seqlens = cu_seqlens - - if cu_seqlens is not None: - need_aux = qo.requires_grad or ko.requires_grad or vo.requires_grad or (betao is not None and betao.requires_grad) - u, ws, lses = ParallelDeltaformerFunction._forward_impl( - qo, ko, vo, betao, C, need_aux=need_aux, cu_seqlens=cu_seqlens) - saved_beta = betao if betao is not None else torch.ones(B, T, H, device=ko.device, dtype=ko.dtype) - ctx.beta_is_none = betao is None - if need_aux: - ctx.save_for_backward(qo, ko, vo, u, ws, lses, saved_beta) - else: - ctx.save_for_backward() - return u - - u, ws, lses = ParallelDeltaformerFunction._forward_impl(qo, ko, vo, betao, C, need_aux=True) - saved_beta = betao if betao is not None else torch.ones(B, T, H, device=ko.device, dtype=ko.dtype) - ctx.save_for_backward(qo, ko, vo, u, ws, lses, saved_beta) - ctx.beta_is_none = betao is None - return u - - @staticmethod - def backward( - ctx, - grad_u: torch.Tensor, - ): - if getattr(ctx, 'cu_seqlens', None) is not None: - cu = ctx.cu_seqlens - qo, ko, vo, u_full, ws, lses, betao = ctx.saved_tensors - B, T_max, H, D = ko.size() - qk_scale = 1.0 / math.sqrt(D) - fa_scale = qk_scale / math.log(2) - - dq = torch.zeros_like(qo) - dk = torch.zeros_like(ko) - dv = torch.zeros_like(vo) - dbeta = None if ctx.beta_is_none else torch.zeros_like(betao) - - C = ctx.C - N = len(cu) - 1 - chunk_bases = [] - total = 0 - lengths = [] - for b in range(N): - L = int(cu[b + 1].item() - cu[b].item()) - lengths.append(L) - chunk_bases.append(total) - if L > 0: - total += (L + C - 1) // C - - for b in range(N): - L = lengths[b] - if L == 0: - continue - base = chunk_bases[b] - seq_start = int(cu[b].item()) - - seq_end = seq_start + L - q_seq = qo[0, seq_start:seq_end, :, :] - k_seq = ko[0, seq_start:seq_end, :, :] - u_seq = u_full[0, seq_start:seq_end, :, :] - beta_seq = betao[0, seq_start:seq_end, :] - lse_seq = lses[0, seq_start:seq_end, :] - go_seq = grad_u[0, seq_start:seq_end, :, :] - - gv_seq = torch.zeros_like(u_seq) - start = ((L - 1) // C) * C - for i_local in range(start, -1, -C): - Ci = min(C, L - i_local) - i0 = i_local - i1 = i_local + Ci - do = go_seq[i0:i1, :, :] - if i_local < L - C: - qi = k_seq[i0:i1, :, :] - ki = q_seq[i1:L, :, :] - lse_tail = lse_seq[i1:L, :] - beta_tail = beta_seq[i1:L, :] - du_tail = parallel_deltaformer_bwd_u_chunk(qi, ki, lse_tail, gv_seq[i1:L, :, :], fa_scale, beta_tail) - do = do - du_tail - Wpad = ws[base + (i_local // C)] - W = Wpad[:Ci, :, :Ci] - W_t = W.transpose(0, 1).contiguous() - du_chunk = invcum.backward_x(do.transpose(0, 1).contiguous(), W_t).transpose(0, 1).contiguous() - gv_seq[i0:i1, :, :].copy_(du_chunk) - - gq, gk, gbeta = parallel_deltaformer_bwd_qk(q_seq, k_seq, u_seq, lse_seq, gv_seq, qk_scale, fa_scale, beta_seq) - dq[0, seq_start:seq_end, :, :].copy_(gq) - dk[0, seq_start:seq_end, :, :].copy_(gk) - dv[0, seq_start:seq_end, :, :].copy_(gv_seq) - if dbeta is not None: - dbeta[0, seq_start:seq_end, :].copy_(gbeta) - - return dq, dk, dv, dbeta, None, None - qo, ko, vo, u, ws, lses, betao = ctx.saved_tensors - C = ctx.C - B, T, H, D = ko.size() - - grad_q = torch.zeros_like(qo) - grad_k = torch.zeros_like(ko) - grad_v = torch.zeros_like(vo) - grad_beta_out = None if ctx.beta_is_none else torch.zeros_like(betao) - - qk_scale = 1.0 / math.sqrt(D) - fa_scale = qk_scale / math.log(2) - - chunk_base = 0 - for b in range(B): - grad_v_seq = torch.empty(T, H, D, device=ko.device, dtype=ko.dtype) - for i in range(T - C, -1, -C): - Ci = min(C, T - i) - do = grad_u[b, i:i + Ci, :, :] - - if i < T - C: - qi = ko[b, i:i + Ci, :, :] - ki = qo[b, i + Ci:, :, :] - lse = lses[b, i + Ci:, :] - if not ctx.beta_is_none: - beta_single = betao[b, i + Ci:, :] - else: - beta_single = torch.ones(T - i - Ci, H, device=ko.device, dtype=ko.dtype) - du = parallel_deltaformer_bwd_u_chunk(qi, ki, lse, grad_v_seq[i + Ci:, :, :], fa_scale, beta_single) - do = grad_u[b, i:i + Ci, :, :] - du - - W = ws[chunk_base + (i // C)][:Ci, :, :Ci] - W_t = W.transpose(0, 1).contiguous() - du = invcum.backward_x(do.transpose(0, 1).contiguous(), W_t).transpose(0, 1).contiguous() - grad_v_seq[i:i + Ci, :, :].copy_(du) - - q_seq = qo[b] - k_seq = ko[b] - u_seq = u[b] - lse_seq = lses[b] - beta_seq = betao[b] if not ctx.beta_is_none else torch.ones(T, H, device=ko.device, dtype=ko.dtype) - - gq, gk, gbeta = parallel_deltaformer_bwd_qk(q_seq, k_seq, u_seq, lse_seq, grad_v_seq, qk_scale, fa_scale, beta_seq) - - grad_q[b].copy_(gq) - grad_k[b].copy_(gk) - grad_v[b].copy_(grad_v_seq) - if not ctx.beta_is_none: - grad_beta_out[b].copy_(gbeta) - - chunk_base += (T + C - 1) // C - - return grad_q, grad_k, grad_v, grad_beta_out, None, None - - @staticmethod - def _forward_impl( - qo: torch.Tensor, - ko: torch.Tensor, - vo: torch.Tensor, - betao: torch.Tensor | None, - C: int, - need_aux: bool, - cu_seqlens: torch.LongTensor | None = None, - ) -> tuple[torch.Tensor, torch.Tensor | None, torch.Tensor | None]: - B, T_max, H, D = ko.size() - C = min(C, T_max) - qk_scale = 1.0 / math.sqrt(D) - fa_scale = qk_scale / math.log(2) - - if cu_seqlens is None: - if betao is None: - beta_full = torch.ones(B, T_max, H, device=ko.device, dtype=ko.dtype) - else: - beta_full = betao - - u_full = torch.empty_like(vo) - if need_aux: - total_chunks = B * ((T_max + C - 1) // C) - ws = torch.empty(total_chunks, C, H, C, device=ko.device, dtype=ko.dtype) - lses = torch.empty(B, T_max, H, device=ko.device, dtype=torch.float) - chunk_base = 0 - else: - ws = None - lses = None - chunk_base = 0 - - for b in range(B): - for i in range(0, T_max, C): - Ci = min(C, T_max - i) - - qi = qo[b, i:i + Ci, :, :] - ki = ko[b, :i + Ci, :, :] - vi = vo[b, i:i + Ci, :, :] - ui_prev = u_full[b, :i + Ci, :, :] - betai = beta_full[b, i:i + Ci, :] - - w, lse_chunk = parallel_deltaformer_chunk_fwd(qi, ki, vi, ui_prev, fa_scale, betai) - w = w * betai.unsqueeze(-1) - if need_aux: - wpad = torch.zeros(C, H, C, device=ko.device, dtype=ko.dtype) - wpad[:Ci, :, :Ci].copy_(w) - ws[chunk_base + (i // C)].copy_(wpad) - lses[b, i:i + Ci, :].copy_(lse_chunk) - - u_chunk_view = u_full[b, i:i + Ci, :, :] - w_t = w.transpose(0, 1).contiguous() - u_chunk_view_t = u_chunk_view.transpose(0, 1).contiguous() - invcum.forward_inplace(u_chunk_view_t, w_t) - u_chunk_view.copy_(u_chunk_view_t.transpose(0, 1)) - - chunk_base += (T_max + C - 1) // C - - return u_full, ws, lses - - N = len(cu_seqlens) - 1 - assert cu_seqlens.dim() == 1 and cu_seqlens.size(0) == N + 1, "cu_seqlens must be [N+1]" - device = ko.device - dtype_k = ko.dtype - if betao is None: - beta_full = torch.ones(B, T_max, H, device=device, dtype=dtype_k) - else: - beta_full = betao - - u_full = torch.empty_like(vo) - if need_aux: - total_chunks = sum((max(0, int(cu_seqlens[b + 1].item() - cu_seqlens[b].item())) + C - 1) // C - for b in range(N)) - ws = torch.empty(total_chunks, C, H, C, device=device, dtype=dtype_k) - lses = torch.empty(B, T_max, H, device=device, dtype=torch.float) - chunk_base = 0 - else: - ws = None - lses = None - chunk_base = 0 - - for b in range(N): - seq_start = int(cu_seqlens[b].item()) - seq_end = int(cu_seqlens[b + 1].item()) - L = max(0, seq_end - seq_start) - if L == 0: - continue - - for i_local in range(0, L, C): - Ci = min(C, L - i_local) - li0 = i_local - li1 = i_local + Ci - - abs_start = seq_start + li0 - abs_end = seq_start + li1 - abs_context_end = seq_start + li1 - - qi = qo[0, abs_start:abs_end, :, :] - ki = ko[0, seq_start:abs_context_end, :, :] - vi = vo[0, abs_start:abs_end, :, :] - ui_prev = u_full[0, seq_start:abs_context_end, :, :] - betai = beta_full[0, abs_start:abs_end, :] - - w, lse_chunk = parallel_deltaformer_chunk_fwd(qi, ki, vi, ui_prev, fa_scale, betai) - w = w * betai.unsqueeze(-1) - if need_aux: - wpad = torch.zeros(C, H, C, device=device, dtype=dtype_k) - wpad[:Ci, :, :Ci].copy_(w) - ws[chunk_base + (i_local // C)].copy_(wpad) - lses[0, abs_start:abs_end, :].copy_(lse_chunk) - - u_chunk_view = u_full[0, abs_start:abs_end, :, :] - w_t = w.transpose(0, 1).contiguous() - u_chunk_view_t = u_chunk_view.transpose(0, 1).contiguous() - invcum.forward_inplace(u_chunk_view_t, w_t) - u_chunk_view.copy_(u_chunk_view_t.transpose(0, 1)) - - chunk_base += (L + C - 1) // C - - return u_full, ws, lses - - -def deltaformer_attn( - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - beta: torch.Tensor | None = None, - attention_mask: torch.LongTensor | None = None, - cu_seqlens: torch.LongTensor | None = None, - C: int = BLOCK_SIZE_C, -) -> torch.Tensor: - if flash_attn_func is None: - raise ImportError("Please install Flash Attention via `pip install flash-attn --no-build-isolation` first") - - if cu_seqlens is not None and q.shape[0] != 1: - raise ValueError( - f"The batch size is expected to be 1 rather than {q.shape[0]} when using `cu_seqlens`. " - f"Please flatten variable-length inputs before processing.", - ) - - B, T, H, D = k.shape - C = min(C, T) - - u = ParallelDeltaformerFunction.apply(q, k, v, beta, C, cu_seqlens) - - if attention_mask is not None: - q_padded, (k_padded, u_padded), indices_q, cu_seqlens_lens, max_seq_lens = unpad_input(q, (k, u), attention_mask, T) - cu_seqlens_q, cu_seqlens_k = cu_seqlens_lens - max_seqlen_q, max_seqlen_k = max_seq_lens - o = flash_attn_varlen_func( - q_padded, k_padded, u_padded, - cu_seqlens_q=cu_seqlens_q, - cu_seqlens_k=cu_seqlens_k, - max_seqlen_q=max_seqlen_q, - max_seqlen_k=max_seqlen_k, - causal=True, - window_size=(-1, -1), - ) - o = pad_input(o, indices_q, B, T) - elif cu_seqlens is not None: - max_seqlen = int((cu_seqlens[1:] - cu_seqlens[:-1]).max().item()) - o = flash_attn_varlen_func( - q.squeeze(0), k.squeeze(0), u.squeeze(0), - cu_seqlens_q=cu_seqlens, - cu_seqlens_k=cu_seqlens, - max_seqlen_q=max_seqlen, - max_seqlen_k=max_seqlen, - causal=True, - window_size=(-1, -1), - ).unsqueeze(0) - else: - o = flash_attn_func(q, k, u, causal=True, window_size=(-1, -1)) - - return o - - -__all__ = [ - 'deltaformer_attn', -] diff --git a/build/torch-xpu/ops/dsa/naive.py b/build/torch-xpu/ops/dsa/naive.py deleted file mode 100644 index 0f9fb890100858efecb477f0179a18503829bd81..0000000000000000000000000000000000000000 --- a/build/torch-xpu/ops/dsa/naive.py +++ /dev/null @@ -1,191 +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 warnings - -import torch -import torch.nn.functional as F -from einops import repeat - - -def naive_dsa_indexer( - q_idx: torch.Tensor, - k_idx: torch.Tensor, - w_idx: torch.Tensor | None = None, - topk: int = 2048, - scale: float | None = None, - activation: str | None = None, - cu_seqlens: torch.LongTensor | None = None, -) -> torch.LongTensor: - r""" - Lightning indexer of DeepSeek Sparse Attention (DSA), scoring causal keys and selecting the top-k. - - For a query at position `t` and a preceding token at position `s`, - the index score is the weighted sum of head-wise similarities (MQA: a single index key is shared across heads): - - I[t, s] = \sum_j w_idx[t, j] * act(scale * q_idx[t, j] . k_idx[s]) - - The `topk` highest-scoring causal keys are then kept per query. - A positive `scale`, a monotone `activation`, and the weights do not change the ranking, - so they only matter when the raw scores are consumed directly rather than for selection. - - Args: - q_idx (torch.Tensor): - Index queries of shape `[B, T, HI, DI]`, where `HI` is the number of indexer heads. - k_idx (torch.Tensor): - Index keys of shape `[B, T, DI]`. A single key per token is shared across all heads. - w_idx (torch.Tensor, Optional): - Per-head index weights of shape `[B, T, HI]`. If `None`, heads are summed with equal weight. Default: `None`. - topk (int, Optional): - Number of keys selected per query. Default: `2048`. - scale (float, Optional): - Scale factor applied to the dot product. If `None`, defaults to `1 / sqrt(DI)`. Default: `None`. - activation (str, Optional): - Name of the activation (a `torch.nn.functional` name) applied to head-wise similarities before aggregation, - e.g. `'relu'`. If `None`, no activation is applied. Default: `None`. - cu_seqlens (torch.LongTensor, Optional): - Cumulative sequence lengths of shape `[N+1]` for variable-length inputs. - When provided, the batch size must be 1. Default: `None`. - - Returns: - indices (torch.LongTensor): - Selected key indices of shape `[B, T, topk]`, padded with `-1`. - For varlen inputs, the indices are local offsets within each sequence. - """ - if scale is None: - scale = q_idx.shape[-1] ** -0.5 - q_idx, k_idx = q_idx.float(), k_idx.float() - if w_idx is not None: - w_idx = w_idx.float() - - def select(q_b, k_b, w_b): - # q_b: [T, HI, DI], k_b: [T, DI], w_b: [T, HI] or None -> indices [T, topk] - T = q_b.shape[0] - # [HI, T, T] - score = torch.einsum('m h d, n d -> h m n', q_b, k_b) * scale - if activation is not None: - score = getattr(F, activation)(score) - # aggregate over indexer heads -> [T, T] - logits = score.sum(0) if w_b is None else torch.einsum('h m n, m h -> m n', score, w_b) - i_t = torch.arange(T, device=q_b.device) - logits = logits.masked_fill(i_t[:, None] < i_t[None, :], float('-inf')) - S = min(topk, T) - values, indices = torch.topk(logits, S, dim=-1) - # drop slots backed by invisible (`-inf`) positions, then pad up to `topk` - indices = indices.masked_fill(torch.isinf(values), -1) - if topk > S: - indices = torch.cat([indices, indices.new_full((T, topk - S), -1)], dim=-1) - return indices - - if cu_seqlens is None: - indices = torch.stack([ - select(q_idx[i], k_idx[i], None if w_idx is None else w_idx[i]) for i in range(q_idx.shape[0]) - ], 0) - else: - assert q_idx.shape[0] == 1, "batch size must be 1 when cu_seqlens are provided" - indices = q_idx.new_full((1, q_idx.shape[1], topk), -1, dtype=torch.long) - for i in range(len(cu_seqlens) - 1): - bos, eos = cu_seqlens[i], cu_seqlens[i + 1] - w_b = None if w_idx is None else w_idx[0, bos:eos] - indices[0, bos:eos] = select(q_idx[0, bos:eos], k_idx[0, bos:eos], w_b) - return indices - - -def naive_dsa( - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - q_idx: torch.Tensor, - k_idx: torch.Tensor, - w_idx: torch.Tensor | None = None, - indices: torch.LongTensor | None = None, - topk: int = 2048, - scale: float | None = None, - cu_seqlens: torch.LongTensor | None = None, - return_indices: bool = False, -) -> torch.Tensor | tuple[torch.Tensor, torch.LongTensor]: - r""" - Naive reference for DeepSeek Sparse Attention (DSA). - - DSA scores every causal key with a lightweight indexer, keeps the top-k per query, - and attends over those keys only. The selection is shared across all query heads. - - Reference: - DeepSeek-V3.2 (https://arxiv.org/abs/2512.02556). - - Args: - q (torch.Tensor): - Queries of shape `[B, T, HQ, K]`. - k (torch.Tensor): - Keys of shape `[B, T, H, K]`. GQA is supported with group size `G = HQ // H`. - v (torch.Tensor): - Values of shape `[B, T, H, V]`. - q_idx (torch.Tensor): - Index queries of shape `[B, T, HI, DI]`. - k_idx (torch.Tensor): - Index keys of shape `[B, T, DI]`, shared across indexer heads. - w_idx (torch.Tensor, Optional): - Per-head index weights of shape `[B, T, HI]`. - If `None`, indexer heads are summed with equal weight. Default: `None`. - indices (torch.LongTensor, Optional): - Precomputed selection of shape `[B, T, topk]`, padded with `-1`. - If provided, the indexer is skipped and `q_idx`/`k_idx`/`w_idx` are ignored. Default: `None`. - topk (int, Optional): - Number of keys selected per query. Default: `2048`. - scale (float, Optional): - Scale factor for attention scores. If `None`, defaults to `1 / sqrt(K)`. Default: `None`. - cu_seqlens (torch.LongTensor, Optional): - Cumulative sequence lengths of shape `[N+1]` for variable-length inputs, consistent with the FlashAttention API. - When provided, the batch size must be 1. Default: `None`. - return_indices (bool, Optional): - Whether to also return the selected indices. Default: `False`. - - Returns: - o (torch.Tensor): - Outputs of shape `[B, T, HQ, V]`. - When `return_indices=True`, also returns the selected indices of shape `[B, T, topk]`. - """ - if scale is None: - scale = q.shape[-1] ** -0.5 - if cu_seqlens is not None: - assert q.shape[0] == 1, "batch size must be 1 when cu_seqlens are provided" - - dtype = q.dtype - G = q.shape[2] // k.shape[2] - q, k, v = (x.float() for x in (q, k, v)) - k, v = (repeat(x, 'b t h d -> b t (h g) d', g=G) for x in (k, v)) - - if indices is None: - indices = naive_dsa_indexer(q_idx, k_idx, w_idx, topk, activation='relu', cu_seqlens=cu_seqlens) - elif q_idx is not None: - warnings.warn("`indices` is provided, the lightning indexer is skipped") - - def attend(q_b, k_b, v_b, i_b): - T = q_b.shape[0] - # scatter selection into a boolean mask, with the dummy column `T` absorbing `-1` padding - mask = q_b.new_zeros(T, T + 1, dtype=torch.bool) - mask.scatter_(1, i_b.masked_fill(i_b < 0, T), True) - mask = mask[:, :T] - i_t = torch.arange(T, device=q_b.device) - mask = mask & (i_t[:, None] >= i_t[None, :]) - # [HQ, T, T] - score = torch.einsum('m h d, n h d -> h m n', q_b, k_b) * scale - score = score.masked_fill(~mask[None], float('-inf')) - p = torch.softmax(score, dim=-1) - return torch.einsum('h m n, n h d -> m h d', p, v_b) - - o = torch.empty_like(v) - if cu_seqlens is None: - for i in range(q.shape[0]): - o[i] = attend(q[i], k[i], v[i], indices[i]) - else: - for i in range(len(cu_seqlens) - 1): - bos, eos = cu_seqlens[i], cu_seqlens[i + 1] - o[0, bos:eos] = attend(q[0, bos:eos], k[0, bos:eos], v[0, bos:eos], indices[0, bos:eos]) - - o = o.to(dtype) - return (o, indices) if return_indices else o diff --git a/build/torch-xpu/ops/forgetting_attn/__init__.py b/build/torch-xpu/ops/forgetting_attn/__init__.py deleted file mode 100644 index 2b60577b96dfdfa17a48a3e858bf391114ba52f4..0000000000000000000000000000000000000000 --- a/build/torch-xpu/ops/forgetting_attn/__init__.py +++ /dev/null @@ -1,14 +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 .naive import naive_forgetting_attn -from .parallel import parallel_forgetting_attn - -__all__ = [ - 'naive_forgetting_attn', - 'parallel_forgetting_attn', -] diff --git a/build/torch-xpu/ops/forgetting_attn/naive.py b/build/torch-xpu/ops/forgetting_attn/naive.py deleted file mode 100644 index cf5612ad941075eda96004929aa4cf94d2038c19..0000000000000000000000000000000000000000 --- a/build/torch-xpu/ops/forgetting_attn/naive.py +++ /dev/null @@ -1,53 +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 torch.nn.functional as F -from einops import rearrange, repeat - - -def naive_forgetting_attn( - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - g: torch.Tensor, - scale: float | None = None, - window_size: int | None = None, -): - """ - Reference PyTorch implementation of forgetting attention. - - Args: - q: [B, T, HQ, D] - k: [B, T, H, D] - v: [B, T, H, D] - g: [B, T, HQ] - scale: float, optional. If None, defaults to 1 / sqrt(D) - window_size: int, optional. If provided, each query only attends to keys - in [i - window_size + 1, i]. If None, full causal attention is used. - - Returns: - output: [B, T, HQ, D] - """ - _, T, HQ, D = q.shape - H = k.shape[2] - G = HQ // H - - if scale is None: - scale = D ** -0.5 - - gc = g.float().cumsum(1) - mask = torch.tril(torch.ones((T, T), dtype=torch.bool, device=q.device)) - if window_size is not None: - mask = mask & torch.triu(torch.ones((T, T), dtype=torch.bool, device=q.device), diagonal=-(window_size - 1)) - - ref = torch.einsum("bqhd,bkhd->bhqk", q.float() * scale, repeat(k, "b t h d -> b t (h g) d", g=G).float()) - ref = ref + rearrange(gc, "b t h -> b h t 1") - rearrange(gc, "b t h -> b h 1 t") - ref = ref.masked_fill(~mask.unsqueeze(0).unsqueeze(0), -float('inf')) - ref = torch.einsum("bhqk,bkhd->bqhd", F.softmax(ref, dim=-1), repeat(v, "b t h d -> b t (h g) d", g=G).float()) - - return ref diff --git a/build/torch-xpu/ops/forgetting_attn/parallel.py b/build/torch-xpu/ops/forgetting_attn/parallel.py deleted file mode 100644 index 594263a317ab6f8323fbe67d7689db68809dfbd8..0000000000000000000000000000000000000000 --- a/build/torch-xpu/ops/forgetting_attn/parallel.py +++ /dev/null @@ -1,62 +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 - -from ...ops.attn.parallel import parallel_attn - - -def parallel_forgetting_attn( - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - g: torch.Tensor, - scale: float | None = None, - window_size: int | None = None, - cu_seqlens: torch.LongTensor | None = None, - **kwargs -) -> torch.Tensor: - r""" - Args: - q (torch.Tensor): - queries of shape `[B, T, HQ, K]`. - k (torch.Tensor): - keys of shape `[B, T, H, K]`. - GQA will be applied if HQ is divisible by H. - v (torch.Tensor): - values of shape `[B, T, H, V]`. - g (torch.Tensor): - log decay factors of shape `[B, T, HQ]`. - scale (Optional[float]): - Scale factor for attention scores. - If not provided, it will default to `1 / sqrt(K)`. Default: `None`. - window_size (Optional[int]): - Sliding window size. If provided, each query at position i only attends to - keys in `[i - window_size + 1, i]`. If `None`, full causal attention is used. - Default: `None`. - cu_seqlens (torch.LongTensor): - Cumulative sequence lengths of shape `[N+1]` used for variable-length training, - consistent with the FlashAttention API. - - Returns: - o (torch.Tensor): - Outputs of shape `[B, T, HQ, V]`. - """ - if 'head_first' in kwargs: - raise DeprecationWarning( - "head_first has been removed. Inputs must be in `[B, T, H, ...]` format.", - ) - if scale is None: - scale = k.shape[-1] ** -0.5 - if cu_seqlens is not None and q.shape[0] != 1: - raise ValueError( - f"The batch size is expected to be 1 rather than {q.shape[0]} when using `cu_seqlens`. " - f"Please flatten variable-length inputs before processing.", - ) - - o = parallel_attn(q, k, v, g, scale, window_size=window_size, cu_seqlens=cu_seqlens) - return o diff --git a/build/torch-xpu/ops/gated_delta_product/chunk.py b/build/torch-xpu/ops/gated_delta_product/chunk.py deleted file mode 100644 index 43a28158f6a04babbb02ba3bbbd4870e5c9be144..0000000000000000000000000000000000000000 --- a/build/torch-xpu/ops/gated_delta_product/chunk.py +++ /dev/null @@ -1,373 +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 -from einops import rearrange - -from ...modules.l2norm import l2norm_bwd, l2norm_fwd -from ...ops.common.chunk_scaled_dot_kkt import chunk_scaled_dot_kkt_fwd -from ...ops.delta_rule.chunk import chunk_delta_rule_bwd -from ...ops.delta_rule.wy_fast import recompute_w_u_fwd as dn_recompute_w_u_fwd -from ...ops.gated_delta_product.chunk_deltaproduct_h import chunk_gated_delta_product_fwd_h -from ...ops.gated_delta_product.chunk_deltaproduct_o import chunk_gated_delta_product_fwd_o -from ...ops.gated_delta_rule.chunk import chunk_gated_delta_rule_bwd -from ...ops.gated_delta_rule.wy_fast import recompute_w_u_fwd as gdn_recompute_w_u_fwd -from ...ops.utils import chunk_local_cumsum, solve_tril -from ...ops.utils.constant import RCP_LN2 -from ...ops.utils.index import prepare_chunk_indices -from ...utils import autocast_custom_bwd, autocast_custom_fwd, input_guard - - -def chunk_gated_delta_product_fwd( - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - g: torch.Tensor, - beta: torch.Tensor, - scale: float, - cu_seqlens: torch.LongTensor | None = None, - initial_state: torch.Tensor | None = None, - output_final_state: bool = False, - num_householder: int = 1, - chunk_indices: torch.LongTensor | None = None, - chunk_indices_dp: torch.LongTensor | None = None, -): - cu_seqlens_dp = cu_seqlens * num_householder if cu_seqlens is not None else None - if g is not None: - g_interleaved = g.new_zeros(g.shape[0], g.shape[1], num_householder, g.shape[2], dtype=torch.float32) - g_interleaved[:, :, 0] = g - g_interleaved = rearrange(g_interleaved, 'b l n h -> b (l n) h').contiguous() - g = chunk_local_cumsum( - g, - chunk_size=64, - scale=RCP_LN2, - cu_seqlens=cu_seqlens, - output_dtype=torch.float32, - chunk_indices=chunk_indices, - ) - g_interleaved = chunk_local_cumsum( - g_interleaved, - chunk_size=64, - scale=RCP_LN2, - cu_seqlens=cu_seqlens_dp, - output_dtype=torch.float32, - chunk_indices=chunk_indices_dp, - ) - else: - g_interleaved = None - g = None - # obtain WY representation. u is actually the new v. - A = chunk_scaled_dot_kkt_fwd( - k=k, - g=g_interleaved, - beta=beta, - cu_seqlens=cu_seqlens_dp, - output_dtype=torch.float32, - chunk_indices=chunk_indices_dp, - ) - A = solve_tril( - A=A, - cu_seqlens=cu_seqlens_dp, - chunk_indices=chunk_indices_dp, - output_dtype=k.dtype, - ) - if g is not None: - w, u = gdn_recompute_w_u_fwd( - k=k, - v=v, - beta=beta, - A=A, - g=g_interleaved, - cu_seqlens=cu_seqlens_dp, - chunk_indices=chunk_indices_dp, - ) - else: - w, u = dn_recompute_w_u_fwd( - k=k, - v=v, - beta=beta, - A=A, - cu_seqlens=cu_seqlens_dp, - chunk_indices=chunk_indices_dp, - ) - h, v_new, final_state = chunk_gated_delta_product_fwd_h( - k=k, - w=w, - u=u, - g=g_interleaved, - initial_state=initial_state, - output_final_state=output_final_state, - cu_seqlens=cu_seqlens_dp, - num_householder=num_householder, - chunk_indices=chunk_indices, - ) - o = chunk_gated_delta_product_fwd_o( - q=q, - k=k, - v=v_new, - h=h, - g=g, - scale=scale, - cu_seqlens=cu_seqlens, - num_householder=num_householder, - chunk_indices=chunk_indices, - ) - return g, g_interleaved, o, A, final_state - - -class ChunkGatedDeltaProductFunction(torch.autograd.Function): - - @staticmethod - @input_guard - @autocast_custom_fwd - def forward( - ctx, - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - g: torch.Tensor, - beta: torch.Tensor, - scale: float, - num_householder: int, - initial_state: torch.Tensor, - output_final_state: bool, - use_qk_l2norm_in_kernel: bool = False, - cu_seqlens: torch.LongTensor | None = None, - cu_seqlens_cpu: torch.LongTensor | None = None, - ): - if use_qk_l2norm_in_kernel: - q, q_rstd = l2norm_fwd(q) - k, k_rstd = l2norm_fwd(k) - else: - q_rstd, k_rstd = None, None - - chunk_indices = None - chunk_indices_dp = None - if cu_seqlens is not None: - chunk_indices = prepare_chunk_indices(cu_seqlens, 64, cu_seqlens_cpu=cu_seqlens_cpu) - cu_seqlens_cpu_dp = None - if cu_seqlens_cpu is not None: - cu_seqlens_cpu_dp = cu_seqlens_cpu * num_householder - chunk_indices_dp = prepare_chunk_indices( - cu_seqlens * num_householder, 64, cu_seqlens_cpu=cu_seqlens_cpu_dp - ) - - g, g_interleaved, o, A, final_state = chunk_gated_delta_product_fwd( - q=q, - k=k, - v=v, - g=g, - beta=beta, - scale=scale, - initial_state=initial_state, - output_final_state=output_final_state, - cu_seqlens=cu_seqlens, - num_householder=num_householder, - chunk_indices=chunk_indices, - chunk_indices_dp=chunk_indices_dp, - ) - ctx.save_for_backward( - q, - q_rstd, - k, - k_rstd, - v, - g_interleaved, - beta, - A, - initial_state, - cu_seqlens, - chunk_indices_dp, - ) - ctx.scale = scale - ctx.use_qk_l2norm_in_kernel = use_qk_l2norm_in_kernel - ctx.num_householder = num_householder - return o.to(q.dtype), final_state - - @staticmethod - @input_guard - @autocast_custom_bwd - def backward( - ctx, - do: torch.Tensor, - dht: torch.Tensor, - ): - ( - q, - q_rstd, - k, - k_rstd, - v, - g, - beta, - A, - initial_state, - cu_seqlens, - chunk_indices_dp, - ) = ctx.saved_tensors - q_new = q.new_zeros(q.shape[0], q.shape[1], ctx.num_householder, q.shape[2], q.shape[3]) - q_new[:, :, -1] = q - do_new = do.new_zeros(do.shape[0], do.shape[1], ctx.num_householder, do.shape[2], do.shape[3]) - do_new[:, :, -1] = do - q_org, q = q, rearrange(q_new, 'b t n h d -> b (t n) h d') - do = rearrange(do_new, 'b t n h d -> b (t n) h d') - # call the gated deltanet kernel for now. - # TODO: optimize the backward pass like the forward pass. - if g is not None: - dq, dk, dv, db, dg, dh0, _, _ = chunk_gated_delta_rule_bwd( - q=q, - k=k, - v=v, - g=g, - beta=beta, - A=A, - scale=ctx.scale, - initial_state=initial_state, - do=do, - dht=dht, - cu_seqlens=cu_seqlens * ctx.num_householder if cu_seqlens is not None else None, - chunk_indices=chunk_indices_dp, - ) - dg = rearrange(dg, 'b (l n) h -> b l n h ', n=ctx.num_householder)[:, :, 0].contiguous().to(g) - else: - dq, dk, dv, db, dh0 = chunk_delta_rule_bwd( - q=q, - k=k, - v=v, - beta=beta, - A=A, - scale=ctx.scale, - initial_state=initial_state, - do=do, - dht=dht, - cu_seqlens=cu_seqlens * ctx.num_householder if cu_seqlens is not None else None, - chunk_indices=chunk_indices_dp, - ) - dg = None - dq = rearrange(dq, 'b (l n) h d -> b l n h d', n=ctx.num_householder)[:, :, -1].contiguous() - if ctx.use_qk_l2norm_in_kernel: - dq = l2norm_bwd(q_org, q_rstd, dq) - dk = l2norm_bwd(k, k_rstd, dk) - return dq.to(q), dk.to(k), dv.to(v), dg, db.to(beta), None, None, dh0, None, None, None, None - - -@torch.compiler.disable -def chunk_gated_delta_product( - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - g: torch.Tensor, - beta: torch.Tensor, - num_householder: int, - scale: float = None, - initial_state: torch.Tensor = None, - output_final_state: bool = False, - use_qk_l2norm_in_kernel: bool = False, - cu_seqlens: torch.LongTensor | None = None, - cu_seqlens_cpu: torch.LongTensor | None = None, -): - r""" - Args: - q (torch.Tensor): - queries of shape `[B, T, H, K]`. - k (torch.Tensor): - keys of shape `[B, T, H, K]`. - v (torch.Tensor): - values of shape `[B, T, H, V]`. - g (torch.Tensor): - (forget) gating tensor (in log space!) of shape `[B, T, H]`. - beta (torch.Tensor): - betas of shape `[B, T, H]`. - num_householder (int): - Number of householder transformations to apply. Default: `1`. - scale (Optional[float]): - Scale factor for the RetNet attention scores. - If not provided, it will default to `1 / sqrt(K)`. Default: `None`. - initial_state (Optional[torch.Tensor]): - Initial state of shape `[N, H, K, V]` for `N` input sequences. - For equal-length input sequences, `N` equals the batch size `B`. - Default: `None`. - output_final_state (Optional[bool]): - Whether to output the final state of shape `[N, H, K, V]`. Default: `False`. - use_qk_l2norm_in_kernel (Optional[bool]): - Whether to use qk l2norm within the kernel for saving GPU memory. - Default: `False`. - cu_seqlens (torch.LongTensor): - Cumulative sequence lengths of shape `[N+1]` used for variable-length training, - consistent with the FlashAttention API. - - Returns: - o (torch.Tensor): - Outputs of shape `[B, T, H, V]`. - final_state (torch.Tensor): - Final state of shape `[N, H, K, V]` if `output_final_state=True` else `None`. - - Examples:: - >>> import torch - >>> import torch.nn.functional as F - >>> from einops import rearrange - >>> from ...ops.gated_delta_rule import chunk_gated_delta_product - # inputs with equal lengths - >>> B, T, H, K, V = 4, 2048, 4, 512, 512 - >>> q = torch.randn(B, T, H, K, dtype=torch.bfloat16, device='cuda') - >>> k = F.normalize(torch.randn(B, T, H, K, dtype=torch.bfloat16, device='cuda'), p=2, dim=-1) - >>> v = torch.randn(B, T, H, V, dtype=torch.bfloat16, device='cuda') - >>> beta = torch.rand(B, T, H, dtype=torch.bfloat16, device='cuda').sigmoid() - >>> g = F.logsigmoid(torch.rand(B, T, H, dtype=torch.bfloat16, device='cuda')) - >>> h0 = torch.randn(B, H, K, V, dtype=torch.bfloat16, device='cuda') - >>> o, ht = chunk_gated_delta_product( - q, k, v, g, beta, - initial_state=h0, - output_final_state=True - ) - # for variable-length inputs, the batch size `B` is expected to be 1 and `cu_seqlens` is required - >>> q, k, v, beta, g = map(lambda x: rearrange(x, 'b t ... -> 1 (b t) ...'), (q, k, v, beta, g)) - # for a batch with 4 sequences, `cu_seqlens` with 5 start/end positions are expected - >>> cu_seqlens = q.new_tensor([0, 2048, 4096, 6144, 8192], dtype=torch.long) - >>> o, ht = chunk_gated_delta_product( - q, k, v, g, beta, - initial_state=h0, - output_final_state=True, - cu_seqlens=cu_seqlens - ) - """ - assert q.dtype != torch.float32, "ChunkGatedDeltaProductFunction does not support float32. Please use bfloat16." - B, T, H, K, V = *q.shape, v.shape[-1] - assert k.shape == (B, T*num_householder, H, K) - assert v.shape == (B, T*num_householder, H, V) - assert beta.shape == (B, T*num_householder, H) - if g is not None: - assert g.shape == (B, T, H) - - if cu_seqlens is not None: - if q.shape[0] != 1: - raise ValueError( - f"The batch size is expected to be 1 rather than {q.shape[0]} when using `cu_seqlens`." - f"Please flatten variable-length inputs before processing.", - ) - if initial_state is not None and initial_state.shape[0] != len(cu_seqlens) - 1: - raise ValueError( - f"The number of initial states is expected to be equal to the number of input sequences, " - f"i.e., {len(cu_seqlens) - 1} rather than {initial_state.shape[0]}.", - ) - if scale is None: - scale = k.shape[-1] ** -0.5 - o, final_state = ChunkGatedDeltaProductFunction.apply( - q, - k, - v, - g, - beta, - scale, - num_householder, - initial_state, - output_final_state, - use_qk_l2norm_in_kernel, - cu_seqlens, - cu_seqlens_cpu, - ) - return o, final_state diff --git a/build/torch-xpu/ops/gated_delta_product/chunk_deltaproduct_h.py b/build/torch-xpu/ops/gated_delta_product/chunk_deltaproduct_h.py deleted file mode 100644 index 6ffdc626f0c6dc03a87883c4aa98538377578dfc..0000000000000000000000000000000000000000 --- a/build/torch-xpu/ops/gated_delta_product/chunk_deltaproduct_h.py +++ /dev/null @@ -1,508 +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 -import triton.language as tl - -from ...ops.utils import prepare_chunk_indices, prepare_chunk_offsets -from ...ops.utils.op import exp2 -from ...utils import IS_NVIDIA_HOPPER, autotune_cache_kwargs - -NUM_WARPS = [2, 4] if IS_NVIDIA_HOPPER else [2, 4, 8, 16] - - -@triton.heuristics({ - 'USE_G': lambda args: args['g'] is not None, - 'USE_INITIAL_STATE': lambda args: args['h0'] is not None, - 'STORE_FINAL_STATE': lambda args: args['ht'] is not None, - 'SAVE_NEW_VALUE': lambda args: args['v_new'] is not None, - 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None, -}) -@triton.autotune( - configs=[ - triton.Config({'BV': BV}, num_warps=num_warps, num_stages=num_stages) - for num_warps in [2, 4] - for num_stages in [2, 3, 4] - for BV in [32, 64] - ], - key=['H', 'K', 'V', 'BT', 'USE_G'], - **autotune_cache_kwargs, -) -@triton.jit(do_not_specialize=['T']) -def chunk_gated_delta_product_fwd_kernel_h_blockdim64( - k, - v, - w, - v_new, - g, - h, - h0, - ht, - cu_seqlens, - chunk_offsets, - T, - num_householder: tl.constexpr, # number of delta products - H: tl.constexpr, - K: tl.constexpr, - V: tl.constexpr, - BT: tl.constexpr, - BV: tl.constexpr, - USE_G: tl.constexpr, - USE_INITIAL_STATE: tl.constexpr, - STORE_FINAL_STATE: tl.constexpr, - SAVE_NEW_VALUE: tl.constexpr, - IS_VARLEN: tl.constexpr, -): - i_v, i_nh = tl.program_id(0), tl.program_id(1) - i_n, i_h = i_nh // H, i_nh % H - if IS_VARLEN: - bos, eos = tl.load(cu_seqlens + i_n).to(tl.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) - T = eos - bos - NT = tl.cdiv(T, BT) - boh = tl.load(chunk_offsets + i_n).to(tl.int32) - else: - bos, eos = i_n * T, i_n * T + T - NT = tl.cdiv(T, BT) - boh = i_n * tl.cdiv(T // num_householder, BT) - - # [BK, BV] - b_h1 = tl.zeros([64, BV], dtype=tl.float32) - if K > 64: - b_h2 = tl.zeros([64, BV], dtype=tl.float32) - if K > 128: - b_h3 = tl.zeros([64, BV], dtype=tl.float32) - if K > 192: - b_h4 = tl.zeros([64, BV], dtype=tl.float32) - - # calculate offset - h += (boh * H + i_h) * K*V - v += (bos * H + i_h) * V - k += (bos * H + i_h) * K - w += (bos * H + i_h) * K - if SAVE_NEW_VALUE: - v_new += (bos * H + i_h) * V - stride_v = H*V - stride_h = H*K*V - stride_k = H*K - if USE_INITIAL_STATE: - h0 = h0 + i_nh * K*V - if STORE_FINAL_STATE: - ht = ht + i_nh * K*V - - # load initial state - if USE_INITIAL_STATE: - p_h0_1 = tl.make_block_ptr(h0, (K, V), (V, 1), (0, i_v * BV), (64, BV), (1, 0)) - b_h1 += tl.load(p_h0_1, boundary_check=(0, 1)).to(tl.float32) - if K > 64: - p_h0_2 = tl.make_block_ptr(h0, (K, V), (V, 1), (64, i_v * BV), (64, BV), (1, 0)) - b_h2 += tl.load(p_h0_2, boundary_check=(0, 1)).to(tl.float32) - if K > 128: - p_h0_3 = tl.make_block_ptr(h0, (K, V), (V, 1), (128, i_v * BV), (64, BV), (1, 0)) - b_h3 += tl.load(p_h0_3, boundary_check=(0, 1)).to(tl.float32) - if K > 192: - p_h0_4 = tl.make_block_ptr(h0, (K, V), (V, 1), (192, i_v * BV), (64, BV), (1, 0)) - b_h4 += tl.load(p_h0_4, boundary_check=(0, 1)).to(tl.float32) - - # main recurrence - for i_t in range(NT): - if i_t % num_householder == 0: - i_t_true = i_t // num_householder - p_h1 = tl.make_block_ptr(h + i_t_true * stride_h, (K, V), (V, 1), (0, i_v * BV), (64, BV), (1, 0)) - tl.store(p_h1, b_h1.to(p_h1.dtype.element_ty), boundary_check=(0, 1)) - if K > 64: - p_h2 = tl.make_block_ptr(h + i_t_true * stride_h, (K, V), (V, 1), (64, i_v * BV), (64, BV), (1, 0)) - tl.store(p_h2, b_h2.to(p_h2.dtype.element_ty), boundary_check=(0, 1)) - if K > 128: - p_h3 = tl.make_block_ptr(h + i_t_true * stride_h, (K, V), (V, 1), (128, i_v * BV), (64, BV), (1, 0)) - tl.store(p_h3, b_h3.to(p_h3.dtype.element_ty), boundary_check=(0, 1)) - if K > 192: - p_h4 = tl.make_block_ptr(h + i_t_true * stride_h, (K, V), (V, 1), (192, i_v * BV), (64, BV), (1, 0)) - tl.store(p_h4, b_h4.to(p_h4.dtype.element_ty), boundary_check=(0, 1)) - - p_v = tl.make_block_ptr(v, (T, V), (stride_v, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) - p_v_new = tl.make_block_ptr(v_new, (T, V), (stride_v, 1), (i_t * BT, i_v * BV), - (BT, BV), (1, 0)) if SAVE_NEW_VALUE else None - b_v_new = tl.zeros([BT, BV], dtype=tl.float32) - p_w = tl.make_block_ptr(w, (T, K), (stride_k, 1), (i_t * BT, 0), (BT, 64), (1, 0)) - b_w = tl.load(p_w, boundary_check=(0, 1)) - b_v_new += tl.dot(b_w, b_h1.to(b_w.dtype)) - if K > 64: - p_w = tl.make_block_ptr(w, (T, K), (stride_k, 1), (i_t * BT, 64), (BT, 64), (1, 0)) - b_w = tl.load(p_w, boundary_check=(0, 1)) - b_v_new += tl.dot(b_w, b_h2.to(b_w.dtype)) - if K > 128: - p_w = tl.make_block_ptr(w, (T, K), (stride_k, 1), (i_t * BT, 128), (BT, 64), (1, 0)) - b_w = tl.load(p_w, boundary_check=(0, 1)) - b_v_new += tl.dot(b_w, b_h3.to(b_w.dtype)) - if K > 192: - p_w = tl.make_block_ptr(w, (T, K), (stride_k, 1), (i_t * BT, 192), (BT, 64), (1, 0)) - b_w = tl.load(p_w, boundary_check=(0, 1)) - b_v_new += tl.dot(b_w, b_h4.to(b_w.dtype)) - b_v_new = -b_v_new + tl.load(p_v, boundary_check=(0, 1)) - - if SAVE_NEW_VALUE: - p_v_new = tl.make_block_ptr(v_new, (T, V), (stride_v, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) - tl.store(p_v_new, b_v_new.to(p_v_new.dtype.element_ty), boundary_check=(0, 1)) - - if USE_G: - m_t = (i_t * BT + tl.arange(0, BT)) < T - last_idx = min((i_t + 1) * BT, T) - 1 - b_g_last = tl.load(g + bos * H + last_idx * H + i_h) - p_g = tl.make_block_ptr(g + bos * H + i_h, (T,), (H,), (i_t * BT,), (BT,), (0,)) - b_g = tl.load(p_g, boundary_check=(0,)) - b_v_new = b_v_new * tl.where(m_t, exp2(b_g_last - b_g), 0)[:, None] - b_g_last = exp2(b_g_last) - b_h1 = b_h1 * b_g_last - if K > 64: - b_h2 = b_h2 * b_g_last - if K > 128: - b_h3 = b_h3 * b_g_last - if K > 192: - b_h4 = b_h4 * b_g_last - b_v_new = b_v_new.to(k.dtype.element_ty) - p_k = tl.make_block_ptr(k, (K, T), (1, stride_k), (0, i_t * BT), (64, BT), (0, 1)) - b_k = tl.load(p_k, boundary_check=(0, 1)) - b_h1 += tl.dot(b_k, b_v_new) - if K > 64: - p_k = tl.make_block_ptr(k, (K, T), (1, stride_k), (64, i_t * BT), (64, BT), (0, 1)) - b_k = tl.load(p_k, boundary_check=(0, 1)) - b_h2 += tl.dot(b_k, b_v_new) - if K > 128: - p_k = tl.make_block_ptr(k, (K, T), (1, stride_k), (128, i_t * BT), (64, BT), (0, 1)) - b_k = tl.load(p_k, boundary_check=(0, 1)) - b_h3 += tl.dot(b_k, b_v_new) - if K > 192: - p_k = tl.make_block_ptr(k, (K, T), (1, stride_k), (192, i_t * BT), (64, BT), (0, 1)) - b_k = tl.load(p_k, boundary_check=(0, 1)) - b_h4 += tl.dot(b_k, b_v_new) - # epilogue - if STORE_FINAL_STATE: - p_ht = tl.make_block_ptr(ht, (K, V), (V, 1), (0, i_v * BV), (64, BV), (1, 0)) - tl.store(p_ht, b_h1.to(p_ht.dtype.element_ty), boundary_check=(0, 1)) - if K > 64: - p_ht = tl.make_block_ptr(ht, (K, V), (V, 1), (64, i_v * BV), (64, BV), (1, 0)) - tl.store(p_ht, b_h2.to(p_ht.dtype.element_ty), boundary_check=(0, 1)) - if K > 128: - p_ht = tl.make_block_ptr(ht, (K, V), (V, 1), (128, i_v * BV), (64, BV), (1, 0)) - tl.store(p_ht, b_h3.to(p_ht.dtype.element_ty), boundary_check=(0, 1)) - if K > 192: - p_ht = tl.make_block_ptr(ht, (K, V), (V, 1), (192, i_v * BV), (64, BV), (1, 0)) - tl.store(p_ht, b_h4.to(p_ht.dtype.element_ty), boundary_check=(0, 1)) - - -@triton.heuristics({ - 'USE_G': lambda args: args['g'] is not None, - 'USE_INITIAL_STATE': lambda args: args['dh0'] is not None, - 'USE_FINAL_STATE_GRADIENT': lambda args: args['dht'] is not None, - 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None, -}) -@triton.autotune( - configs=[ - triton.Config({'BV': BV}, num_warps=num_warps, num_stages=num_stages) - for num_warps in [2, 4] - for num_stages in [4, 3, 2] - for BV in [64, 32] - ], - key=['H', 'K', 'V', 'BT', 'BV', 'USE_G'], - **autotune_cache_kwargs, -) -@triton.jit(do_not_specialize=['T']) -def chunk_gated_delta_product_bwd_kernel_dhu_blockdim64( - q, - k, - w, - g, - dht, - dh0, - do, - dh, - dv, - dv2, - cu_seqlens, - chunk_offsets, - scale, - T, - H: tl.constexpr, - K: tl.constexpr, - V: tl.constexpr, - BT: tl.constexpr, - BV: tl.constexpr, - USE_G: tl.constexpr, - USE_INITIAL_STATE: tl.constexpr, - USE_FINAL_STATE_GRADIENT: tl.constexpr, - IS_VARLEN: tl.constexpr, -): - i_v, i_nh = tl.program_id(0), tl.program_id(1) - i_n, i_h = i_nh // H, i_nh % H - if IS_VARLEN: - bos, eos = tl.load(cu_seqlens + i_n).to(tl.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) - T = eos - bos - NT = tl.cdiv(T, BT) - boh = tl.load(chunk_offsets + i_n).to(tl.int32) - else: - bos, eos = i_n * T, i_n * T + T - NT = tl.cdiv(T, BT) - boh = i_n * NT - - # [BK, BV] - b_dh1 = tl.zeros([64, BV], dtype=tl.float32) - if K > 64: - b_dh2 = tl.zeros([64, BV], dtype=tl.float32) - if K > 128: - b_dh3 = tl.zeros([64, BV], dtype=tl.float32) - if K > 192: - b_dh4 = tl.zeros([64, BV], dtype=tl.float32) - - # calculate offset - dh += (boh * H + i_h) * K*V - dv += (bos * H + i_h) * V - dv2 += (bos * H + i_h) * V - q += (bos * H + i_h) * K - k += (bos * H + i_h) * K - w += (bos * H + i_h) * K - do += (bos * H + i_h) * V - stride_v = H*V - stride_h = H*K*V - stride_k = H*K - if USE_INITIAL_STATE: - dh0 += i_nh * K*V - if USE_FINAL_STATE_GRADIENT: - dht += i_nh * K*V - - if USE_FINAL_STATE_GRADIENT: - p_dht1 = tl.make_block_ptr(dht, (K, V), (V, 1), (0, i_v * BV), (64, BV), (1, 0)) - b_dh1 += tl.load(p_dht1, boundary_check=(0, 1)) - if K > 64: - p_dht2 = tl.make_block_ptr(dht, (K, V), (V, 1), (64, i_v * BV), (64, BV), (1, 0)) - b_dh2 += tl.load(p_dht2, boundary_check=(0, 1)) - if K > 128: - p_dht3 = tl.make_block_ptr(dht, (K, V), (V, 1), (128, i_v * BV), (64, BV), (1, 0)) - b_dh3 += tl.load(p_dht3, boundary_check=(0, 1)) - if K > 192: - p_dht4 = tl.make_block_ptr(dht, (K, V), (V, 1), (192, i_v * BV), (64, BV), (1, 0)) - b_dh4 += tl.load(p_dht4, boundary_check=(0, 1)) - - for i_t in range(NT - 1, -1, -1): - p_dh1 = tl.make_block_ptr(dh + i_t*stride_h, (K, V), (V, 1), (0, i_v * BV), (64, BV), (1, 0)) - tl.store(p_dh1, b_dh1.to(p_dh1.dtype.element_ty), boundary_check=(0, 1)) - if K > 64: - p_dh2 = tl.make_block_ptr(dh + i_t*stride_h, (K, V), (V, 1), (64, i_v * BV), (64, BV), (1, 0)) - tl.store(p_dh2, b_dh2.to(p_dh2.dtype.element_ty), boundary_check=(0, 1)) - if K > 128: - p_dh3 = tl.make_block_ptr(dh + i_t*stride_h, (K, V), (V, 1), (128, i_v * BV), (64, BV), (1, 0)) - tl.store(p_dh3, b_dh3.to(p_dh3.dtype.element_ty), boundary_check=(0, 1)) - if K > 192: - p_dh4 = tl.make_block_ptr(dh + i_t*stride_h, (K, V), (V, 1), (192, i_v * BV), (64, BV), (1, 0)) - tl.store(p_dh4, b_dh4.to(p_dh4.dtype.element_ty), boundary_check=(0, 1)) - - if USE_G: - last_idx = min((i_t + 1) * BT, T) - 1 - bg_last = tl.load(g + (bos + last_idx) * H + i_h) - bg_last_exp = exp2(bg_last) - p_g = tl.make_block_ptr(g + bos * H + i_h, (T,), (H,), (i_t * BT,), (BT,), (0,)) - b_g = tl.load(p_g, boundary_check=(0,)) - b_g_exp = exp2(b_g) - else: - bg_last = None - last_idx = None - b_g = None - b_g_exp = None - - p_dv = tl.make_block_ptr(dv, (T, V), (stride_v, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) - p_wo = tl.make_block_ptr(do, (T, V), (stride_v, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) - p_dv2 = tl.make_block_ptr(dv2, (T, V), (stride_v, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) - - b_wo = tl.load(p_wo, boundary_check=(0, 1)) - b_dv = tl.zeros([BT, BV], dtype=tl.float32) - - # Update dv - p_k = tl.make_block_ptr(k, (T, K), (stride_k, 1), (i_t * BT, 0), (BT, 64), (1, 0)) - b_k = tl.load(p_k, boundary_check=(0, 1)) - b_dv += tl.dot(b_k, b_dh1.to(b_k.dtype)) - - if K > 64: - p_k = tl.make_block_ptr(k, (T, K), (stride_k, 1), (i_t * BT, 64), (BT, 64), (1, 0)) - b_k = tl.load(p_k, boundary_check=(0, 1)) - b_dv += tl.dot(b_k, b_dh2.to(b_k.dtype)) - - if K > 128: - p_k = tl.make_block_ptr(k, (T, K), (stride_k, 1), (i_t * BT, 128), (BT, 64), (1, 0)) - b_k = tl.load(p_k, boundary_check=(0, 1)) - b_dv += tl.dot(b_k, b_dh3.to(b_k.dtype)) - - if K > 192: - p_k = tl.make_block_ptr(k, (T, K), (stride_k, 1), (i_t * BT, 192), (BT, 64), (1, 0)) - b_k = tl.load(p_k, boundary_check=(0, 1)) - b_dv += tl.dot(b_k, b_dh4.to(b_k.dtype)) - - if USE_G: - m_t = (i_t * BT + tl.arange(0, BT)) < T - b_dv *= tl.where(m_t, exp2(bg_last - b_g), 0)[:, None] - b_dv += tl.load(p_dv, boundary_check=(0, 1)) - - tl.store(p_dv2, b_dv.to(p_dv.dtype.element_ty), boundary_check=(0, 1)) - # Update dh - p_w = tl.make_block_ptr(w, (K, T), (1, stride_k), (0, i_t * BT), (64, BT), (0, 1)) - p_q = tl.make_block_ptr(q, (K, T), (1, stride_k), (0, i_t * BT), (64, BT), (0, 1)) - b_w = tl.load(p_w, boundary_check=(0, 1)) - b_q = tl.load(p_q, boundary_check=(0, 1)) - if USE_G: - b_dh1 *= bg_last_exp - b_q = b_q * b_g_exp[None, :] - b_q = (b_q * scale).to(b_q.dtype) - b_dh1 += tl.dot(b_q, b_wo.to(b_q.dtype))-tl.dot(b_w, b_dv.to(b_w.dtype)) - if K > 64: - p_q = tl.make_block_ptr(q, (K, T), (1, stride_k), (64, i_t * BT), (64, BT), (0, 1)) - p_w = tl.make_block_ptr(w, (K, T), (1, stride_k), (64, i_t * BT), (64, BT), (0, 1)) - b_q = tl.load(p_q, boundary_check=(0, 1)) - b_w = tl.load(p_w, boundary_check=(0, 1)) - if USE_G: - b_dh2 *= bg_last_exp - b_q = b_q * b_g_exp[None, :] - b_q = (b_q * scale).to(b_q.dtype) - b_dh2 += tl.dot(b_q, b_wo.to(b_q.dtype))-tl.dot(b_w, b_dv.to(b_w.dtype)) - if K > 128: - p_q = tl.make_block_ptr(q, (K, T), (1, stride_k), (128, i_t * BT), (64, BT), (0, 1)) - p_w = tl.make_block_ptr(w, (K, T), (1, stride_k), (128, i_t * BT), (64, BT), (0, 1)) - b_q = tl.load(p_q, boundary_check=(0, 1)) - b_w = tl.load(p_w, boundary_check=(0, 1)) - if USE_G: - b_dh3 *= bg_last_exp - b_q = b_q * b_g_exp[None, :] - b_q = (b_q * scale).to(b_q.dtype) - b_dh3 += tl.dot(b_q, b_wo.to(b_q.dtype))-tl.dot(b_w, b_dv.to(b_w.dtype)) - if K > 192: - p_q = tl.make_block_ptr(q, (K, T), (1, stride_k), (192, i_t * BT), (64, BT), (0, 1)) - p_w = tl.make_block_ptr(w, (K, T), (1, stride_k), (192, i_t * BT), (64, BT), (0, 1)) - b_q = tl.load(p_q, boundary_check=(0, 1)) - b_w = tl.load(p_w, boundary_check=(0, 1)) - if USE_G: - b_dh4 *= bg_last_exp - b_q = b_q * b_g_exp[None, :] - b_q = (b_q * scale).to(b_q.dtype) - b_dh4 += tl.dot(b_q, b_wo.to(b_q.dtype))-tl.dot(b_w, b_dv.to(b_w.dtype)) - - if USE_INITIAL_STATE: - p_dh0 = tl.make_block_ptr(dh0, (K, V), (V, 1), (0, i_v * BV), (64, BV), (1, 0)) - tl.store(p_dh0, b_dh1.to(p_dh0.dtype.element_ty), boundary_check=(0, 1)) - if K > 64: - p_dh1 = tl.make_block_ptr(dh0, (K, V), (V, 1), (64, i_v * BV), (64, BV), (1, 0)) - tl.store(p_dh1, b_dh2.to(p_dh1.dtype.element_ty), boundary_check=(0, 1)) - if K > 128: - p_dh2 = tl.make_block_ptr(dh0, (K, V), (V, 1), (128, i_v * BV), (64, BV), (1, 0)) - tl.store(p_dh2, b_dh3.to(p_dh2.dtype.element_ty), boundary_check=(0, 1)) - if K > 192: - p_dh3 = tl.make_block_ptr(dh0, (K, V), (V, 1), (192, i_v * BV), (64, BV), (1, 0)) - tl.store(p_dh3, b_dh4.to(p_dh3.dtype.element_ty), boundary_check=(0, 1)) - - -def chunk_gated_delta_product_fwd_h( - k: torch.Tensor, - w: torch.Tensor, - u: torch.Tensor, - g: torch.Tensor | None = None, - initial_state: torch.Tensor | None = None, - output_final_state: bool = False, - chunk_size: int = 64, # SY: remove this argument and force chunk size 64? - save_new_value: bool = True, - cu_seqlens: torch.LongTensor | None = None, - num_householder: int = 1, - chunk_indices: torch.LongTensor | None = None, -) -> tuple[torch.Tensor, torch.Tensor]: - B, T, H, K, V = *k.shape, u.shape[-1] - assert T % num_householder == 0, "T must be divisible by num_householder" - T_true = T // num_householder - BT = chunk_size - if chunk_indices is None and cu_seqlens is not None: - chunk_indices = prepare_chunk_indices(cu_seqlens // num_householder, chunk_size) - # N: the actual number of sequences in the batch with either equal or variable lengths - if cu_seqlens is None: - N, NT, chunk_offsets = B, triton.cdiv(T_true, BT), None - else: - N, NT, chunk_offsets = len(cu_seqlens) - \ - 1, len(chunk_indices), prepare_chunk_offsets(cu_seqlens // num_householder, BT) - assert K <= 256, "current kernel does not support head dimension larger than 256." - h = k.new_empty(B, NT, H, K, V) - final_state = k.new_empty(N, H, K, V, dtype=torch.float32) if output_final_state else None - v_new = torch.empty_like(u) if save_new_value else None - - def grid(meta): return (triton.cdiv(V, meta['BV']), N*H) - chunk_gated_delta_product_fwd_kernel_h_blockdim64[grid]( - k=k, - v=u, - w=w, - v_new=v_new, - g=g, - h=h, - h0=initial_state, - ht=final_state, - cu_seqlens=cu_seqlens, - chunk_offsets=chunk_offsets, - num_householder=num_householder, - T=T, - H=H, - K=K, - V=V, - BT=BT, - ) - return h, v_new, final_state - - -def chunk_gated_delta_product_bwd_dhu( - q: torch.Tensor, - k: torch.Tensor, - w: torch.Tensor, - g: torch.Tensor, - h0: torch.Tensor, - dht: torch.Tensor | None, - do: torch.Tensor, - dv: torch.Tensor, - scale: float, - cu_seqlens: torch.LongTensor | None = None, - chunk_size: int = 64, # SY: remove this argument and force chunk size 64? - chunk_indices: torch.LongTensor | None = None, -) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: - B, T, H, K, V = *q.shape, do.shape[-1] - - # N: the actual number of sequences in the batch with either equal or variable lengths - BT = 64 - assert K <= 256, "current kernel does not support head dimension being larger than 256." - - if chunk_indices is None and cu_seqlens is not None: - chunk_indices = prepare_chunk_indices(cu_seqlens, chunk_size) - if cu_seqlens is None: - N, NT, chunk_offsets = B, triton.cdiv(T, BT), None - else: - N, NT, chunk_offsets = len(cu_seqlens) - 1, len(chunk_indices), prepare_chunk_offsets(cu_seqlens, BT) - - dh = q.new_empty(B, NT, H, K, V) - dh0 = torch.empty_like(h0, dtype=torch.float32) if h0 is not None else None - dv2 = torch.empty_like(dv) - - def grid(meta): return (triton.cdiv(V, meta['BV']), N*H) - chunk_gated_delta_product_bwd_kernel_dhu_blockdim64[grid]( - q=q, - k=k, - w=w, - g=g, - dht=dht, - dh0=dh0, - do=do, - dh=dh, - dv=dv, - dv2=dv2, - cu_seqlens=cu_seqlens, - chunk_offsets=chunk_offsets, - scale=scale, - T=T, - H=H, - K=K, - V=V, - BT=BT, - ) - return dh, dh0, dv2 diff --git a/build/torch-xpu/ops/gated_delta_product/chunk_deltaproduct_o.py b/build/torch-xpu/ops/gated_delta_product/chunk_deltaproduct_o.py deleted file mode 100644 index 975530c4c61d64006c1ee17886aa11d598823a98..0000000000000000000000000000000000000000 --- a/build/torch-xpu/ops/gated_delta_product/chunk_deltaproduct_o.py +++ /dev/null @@ -1,159 +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 -import triton.language as tl - -from ...ops.utils import prepare_chunk_indices -from ...ops.utils.op import exp2 -from ...utils import IS_NVIDIA_HOPPER, autotune_cache_kwargs, check_shared_mem - -BKV_LIST = [64, 128] if check_shared_mem() else [32, 64] -NUM_WARPS = [2, 4] if IS_NVIDIA_HOPPER else [2, 4, 8] - - -@triton.heuristics({ - 'USE_G': lambda args: args['g'] is not None, - 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None, -}) -@triton.autotune( - configs=[ - triton.Config({'BK': BK, 'BV': BV}, num_warps=num_warps, num_stages=num_stages) - for BK in BKV_LIST - for BV in BKV_LIST - for num_warps in NUM_WARPS - for num_stages in [2, 3, 4] - ], - key=['H', 'K', 'V', 'BT'], - **autotune_cache_kwargs, -) -@triton.jit(do_not_specialize=['T']) -def chunk_fwd_kernel_o( - q, - k, - v, - h, - g, - o, - cu_seqlens, - chunk_indices, - scale, - T, - num_householder: tl.constexpr, - H: tl.constexpr, - K: tl.constexpr, - V: tl.constexpr, - BT: tl.constexpr, - BK: tl.constexpr, - BV: tl.constexpr, - USE_G: tl.constexpr, - IS_VARLEN: tl.constexpr, -): - i_v, i_t, i_bh = tl.program_id(0), tl.program_id(1), tl.program_id(2) - i_b, i_h = i_bh // H, i_bh % H - - 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.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) - T = eos - bos - NT = tl.cdiv(T, BT) - else: - NT = tl.cdiv(T, BT) - i_tg = i_b * NT + i_t - bos, eos = i_b * T, i_b * T + T - - # offset calculation - q += (bos * H + i_h) * K - k += (bos * num_householder * H + i_h) * K - v += (bos * num_householder * H + i_h) * V - o += (bos * H + i_h) * V - h += (i_tg * H + i_h).to(tl.int64) * K*V - - b_o = tl.zeros([BT, BV], dtype=tl.float32) - - for i_k in range(tl.cdiv(K, BK)): - p_q = tl.make_block_ptr(q, (T, K), (H*K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) - p_h = tl.make_block_ptr(h, (K, V), (V, 1), (i_k * BK, i_v * BV), (BK, BV), (1, 0)) - # [BT, BK] - b_q = tl.load(p_q, boundary_check=(0, 1)) - # [BK, BV] - b_h = tl.load(p_h, boundary_check=(0, 1)) - # [BT, BK] @ [BK, BV] -> [BT, BV] - b_o += tl.dot(b_q, b_h) - - o_t = i_t * BT + tl.arange(0, BT) - m_t = o_t < T - if USE_G: - g += bos * H + i_h - p_g = tl.make_block_ptr(g, (T,), (H,), (i_t * BT,), (BT,), (0,)) - b_g = tl.load(p_g, boundary_check=(0,)) - m_A = (o_t[:, None] >= o_t[None, :]) & (m_t[:, None] & m_t) - b_m = tl.where(m_A, exp2(b_g[:, None] - b_g[None, :]), 0) - b_o = b_o * exp2(b_g)[:, None] - else: - b_m = ((o_t[:, None] >= o_t[None, :]) & (m_t[:, None] & m_t)).to(tl.float32) - - for i_dp in range(num_householder): - b_A = tl.zeros([BT, BT], dtype=tl.float32) - for i_k in range(tl.cdiv(K, BK)): - p_q = tl.make_block_ptr(q, (T, K), (H*K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) - p_k = tl.make_block_ptr(k+i_dp*H*K, (K, T), (1, num_householder*H*K), (i_k * BK, i_t * BT), (BK, BT), (0, 1)) - # [BT, BK] - b_q = tl.load(p_q, boundary_check=(0, 1)) - # [BK, BT] - b_k = tl.load(p_k, boundary_check=(0, 1)) - # [BT, BK] @ [BK, BT] -> [BT, BT] - b_A += tl.dot(b_q, b_k) - b_A = b_A * b_m - p_v = tl.make_block_ptr(v+i_dp*H*V, (T, V), (H*V*num_householder, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) - b_v = tl.load(p_v, boundary_check=(0, 1)) - b_o += tl.dot(b_A.to(b_v.dtype), b_v) - b_o = b_o * scale - p_o = tl.make_block_ptr(o, (T, V), (H*V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) - tl.store(p_o, b_o.to(p_o.dtype.element_ty), boundary_check=(0, 1)) - - -def chunk_gated_delta_product_fwd_o( - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - h: torch.Tensor, - g: torch.Tensor | None = None, # cumsum of log decay - scale: float | None = None, - cu_seqlens: torch.LongTensor | None = None, - chunk_size: int = 64, - num_householder: int = 1, - chunk_indices: torch.LongTensor | None = None, -) -> torch.Tensor: - assert q.shape[1] * num_householder == k.shape[1], "q.shape[1] * num_householder must be equal to k.shape[1]" - B, T, H, K, V = *q.shape, v.shape[-1] - BT = chunk_size - if chunk_indices is None and cu_seqlens is not None: - chunk_indices = prepare_chunk_indices(cu_seqlens, BT) - NT = triton.cdiv(T, BT) if cu_seqlens is None else len(chunk_indices) - o = v.new_empty(B, T, H, V).fill_(-float('inf')) - def grid(meta): return (triton.cdiv(V, meta['BV']), NT, B * H) - chunk_fwd_kernel_o[grid]( - q, - k, - v, - h, - g, - o, - cu_seqlens, - chunk_indices, - scale, - T=T, - num_householder=num_householder, - H=H, - K=K, - V=V, - BT=BT, - ) - return o diff --git a/build/torch-xpu/ops/gated_delta_product/chunk_ref.py b/build/torch-xpu/ops/gated_delta_product/chunk_ref.py deleted file mode 100644 index 81429aaa58674bc0aaba1169cc15821a33ca52a2..0000000000000000000000000000000000000000 --- a/build/torch-xpu/ops/gated_delta_product/chunk_ref.py +++ /dev/null @@ -1,70 +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 -from einops import rearrange - -from ...ops.delta_rule import chunk_delta_rule -from ...ops.gated_delta_rule import chunk_gated_delta_rule - - -@torch.compiler.disable -def chunk_gated_delta_product_ref( - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - g: torch.Tensor, - beta: torch.Tensor, - num_householder: int, - scale: float = None, - initial_state: torch.Tensor = None, - output_final_state: bool = False, - cu_seqlens: torch.LongTensor | None = None, - use_qk_l2norm_in_kernel: bool = False, -): - assert q.dtype != torch.float32, "ChunkGatedDeltaProductFunction does not support float32. Please use bfloat16." - B, T, H, K = q.shape - V = v.shape[-1] - assert k.shape == (B, T*num_householder, H, K) - assert v.shape == (B, T*num_householder, H, V) - assert beta.shape == (B, T*num_householder, H) - if g is not None: - assert g.shape == (B, T, H) - q_new = q.new_zeros(B, T, num_householder, H, K) - q_new[:, :, -1] = q - q = rearrange(q_new, 'b t n h d -> b (t n) h d') - - if g is not None: - g_new = g.new_zeros(B, T, num_householder, H, dtype=torch.float32) - g_new[:, :, 0] = g - g = rearrange(g_new, 'b t n h -> b (t n) h') - o, final_state = chunk_gated_delta_rule( - q=q, - k=k, - v=v, - g=g, - beta=beta, - initial_state=initial_state, - output_final_state=output_final_state, - cu_seqlens=cu_seqlens * num_householder if cu_seqlens is not None else None, - use_qk_l2norm_in_kernel=use_qk_l2norm_in_kernel, - scale=scale, - ) - else: - o, final_state = chunk_delta_rule( - q=q, - k=k, - v=v, - beta=beta, - initial_state=initial_state, - output_final_state=output_final_state, - cu_seqlens=cu_seqlens * num_householder if cu_seqlens is not None else None, - use_qk_l2norm_in_kernel=use_qk_l2norm_in_kernel, - scale=scale, - ) - o = rearrange(o, 'b (t n) h d -> b t n h d', n=num_householder) - return o[:, :, -1].contiguous(), final_state diff --git a/build/torch-xpu/ops/gated_delta_product/naive.py b/build/torch-xpu/ops/gated_delta_product/naive.py deleted file mode 100644 index 106f13d2954e5fad48f46be4526b92b876a7d8a0..0000000000000000000000000000000000000000 --- a/build/torch-xpu/ops/gated_delta_product/naive.py +++ /dev/null @@ -1,43 +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 - - -def naive_recurrent_gated_delta_product(q, k, v, g, beta, scale, cu_seqlens, - initial_state=None, output_final_state=False, - num_householder=1): - q_original_dtype = q.dtype - B, T, H, K = q.shape - V = v.shape[-1] - assert k.shape == (B, T*num_householder, H, K) - assert v.shape == (B, T*num_householder, H, V) - assert beta.shape == (B, T*num_householder, H) - if g is not None: - assert g.shape == (B, T, H) - q, k, v, beta = map(lambda x: x.float(), (q, k, v, beta)) - - h = torch.zeros(B, H, K, V, dtype=torch.float32, device=q.device) - if initial_state is not None: - h = initial_state - - o = torch.zeros(B, T, H, V, dtype=torch.float32, device=q.device) - - for i in range(T): - if g is not None: - h = h * g[:, i, :].exp()[..., None, None] - # multiple state transition - for j in range(num_householder): - k_ij = k[:, i*num_householder+j, :, :] - v_ij = v[:, i*num_householder+j, :, :] - beta_ij = beta[:, i*num_householder+j, :] - h = h + (v_ij - (h * k_ij[..., None]).sum(-2)).unsqueeze(-2) * k_ij[..., None] * beta_ij[..., None, None] - # memory readout - q_i = q[:, i, :, :] - o_i = (h * q_i[..., None]).sum(-2) - o[:, i] = o_i - return o.to(q_original_dtype), h diff --git a/build/torch-xpu/ops/gated_delta_rule/__init__.py b/build/torch-xpu/ops/gated_delta_rule/__init__.py deleted file mode 100644 index 7e65713b5e68ca8cd389b8d4f13d637dfcf19f1d..0000000000000000000000000000000000000000 --- a/build/torch-xpu/ops/gated_delta_rule/__init__.py +++ /dev/null @@ -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 - -from .chunk import chunk_gated_delta_rule, chunk_gdn -from .fused_recurrent import fused_recurrent_gated_delta_rule, fused_recurrent_gdn -from .naive import naive_chunk_gated_delta_rule, naive_recurrent_gated_delta_rule - -__all__ = [ - "chunk_gated_delta_rule", "chunk_gdn", - "fused_recurrent_gated_delta_rule", "fused_recurrent_gdn", - "naive_chunk_gated_delta_rule", - "naive_recurrent_gated_delta_rule", -] diff --git a/build/torch-xpu/ops/gated_delta_rule/chunk.py b/build/torch-xpu/ops/gated_delta_rule/chunk.py deleted file mode 100644 index 63fd1b099859f34dae89aad281cf435bab25c761..0000000000000000000000000000000000000000 --- a/build/torch-xpu/ops/gated_delta_rule/chunk.py +++ /dev/null @@ -1,589 +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 warnings - -import torch - -from ...modules.l2norm import l2norm_bwd, l2norm_fwd -from ...ops.common.chunk_delta_h import chunk_gated_delta_rule_bwd_dhu, chunk_gated_delta_rule_fwd_h -from ...ops.common.chunk_o import chunk_bwd_dqkwg, chunk_bwd_dv_local, chunk_fwd_o -from ...ops.common.gate import fused_beta_sigmoid, fused_beta_sigmoid_bwd -from ...ops.cp import FLACPContext -from ...ops.cp.chunk_delta_h import ( - chunk_gated_delta_rule_bwd_dhu_pre_process, - chunk_gated_delta_rule_fwd_h_pre_process, - compress_h0, - expand_h0, -) -from ...ops.gated_delta_rule.chunk_fwd import chunk_gated_delta_rule_fwd_intra -from ...ops.gated_delta_rule.gate import gdn_gate_bwd, gdn_gate_chunk_cumsum -from ...ops.gated_delta_rule.wy_fast import prepare_wy_repr_bwd, recompute_w_u_fwd -from ...ops.utils import chunk_local_cumsum -from ...ops.utils.constant import RCP_LN2 -from ...ops.utils.index import prepare_chunk_indices -from ...utils import autocast_custom_bwd, autocast_custom_fwd, input_guard - - -def chunk_gated_delta_rule_fwd( - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - g: torch.Tensor, - beta: torch.Tensor, - scale: float, - initial_state: torch.Tensor, - output_final_state: bool, - state_v_first: bool = False, - cu_seqlens: torch.LongTensor | None = None, - cp_context: FLACPContext | None = None, - chunk_indices: torch.LongTensor | None = None, - use_gate_in_kernel: bool = False, - A_log: torch.Tensor | None = None, - dt_bias: torch.Tensor | None = None, - chunk_size: int = 64, -): - g_input = g if use_gate_in_kernel else None - if use_gate_in_kernel: - g = gdn_gate_chunk_cumsum( - g=g, - A_log=A_log, - chunk_size=chunk_size, - scale=RCP_LN2, - dt_bias=dt_bias, - cu_seqlens=cu_seqlens, - chunk_indices=chunk_indices, - ) - else: - g = chunk_local_cumsum( - g, - chunk_size=chunk_size, - scale=RCP_LN2, - cu_seqlens=cu_seqlens, - chunk_indices=chunk_indices, - ) - # obtain WY representation. u is actually the new v. - # fused kkt + solve_tril + recompute_w_u - w, u, A = chunk_gated_delta_rule_fwd_intra( - k=k, - v=v, - g=g, - beta=beta, - cu_seqlens=cu_seqlens, - chunk_indices=chunk_indices, - chunk_size=chunk_size, - ) - - if cp_context is not None: - initial_state = chunk_gated_delta_rule_fwd_h_pre_process( - k=k, - w=w, - u=u, - g=g, - cu_seqlens=cu_seqlens, - initial_state=initial_state, - context=cp_context, - state_v_first=state_v_first, - chunk_size=chunk_size, - ) - - h, v_new, final_state = chunk_gated_delta_rule_fwd_h( - k=k, - w=w, - u=u, - g=g, - initial_state=initial_state, - output_final_state=output_final_state, - cu_seqlens=cu_seqlens, - chunk_indices=chunk_indices, - state_v_first=state_v_first, - chunk_size=chunk_size, - ) - - if cp_context is not None: - initial_state = compress_h0(initial_state, context=cp_context) - - o = chunk_fwd_o( - q=q, - k=k, - v=v_new, - h=h, - g=g, - scale=scale, - cu_seqlens=cu_seqlens, - chunk_indices=chunk_indices, - state_v_first=state_v_first, - chunk_size=chunk_size, - ) - return g, o, A, final_state, initial_state, g_input - - -def chunk_gated_delta_rule_bwd( - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - g: torch.Tensor, - beta: torch.Tensor, - A: torch.Tensor, - scale: float, - initial_state: torch.Tensor, - do: torch.Tensor, - dht: torch.Tensor, - state_v_first: bool = False, - cu_seqlens: torch.LongTensor | None = None, - cp_context: FLACPContext | None = None, - chunk_indices: torch.LongTensor | None = None, - use_gate_in_kernel: bool = False, - g_input: torch.Tensor | None = None, - A_log: torch.Tensor | None = None, - dt_bias: torch.Tensor | None = None, - chunk_size: int = 64, -): - w, u = recompute_w_u_fwd( - k=k, - v=v, - beta=beta, - A=A, - g=g, - cu_seqlens=cu_seqlens, - chunk_indices=chunk_indices, - ) - - if cp_context is not None: - initial_state = expand_h0(initial_state, context=cp_context) - - h, v_new, _ = chunk_gated_delta_rule_fwd_h( - k=k, - w=w, - u=u, - g=g, - initial_state=initial_state, - output_final_state=False, - cu_seqlens=cu_seqlens, - chunk_indices=chunk_indices, - state_v_first=state_v_first, - chunk_size=chunk_size, - ) - dv = chunk_bwd_dv_local( - q=q, - k=k, - g=g, - do=do, - scale=scale, - cu_seqlens=cu_seqlens, - chunk_indices=chunk_indices, - chunk_size=chunk_size, - ) - - if cp_context is not None: - # initial_state is None in the CP mode - # We only need to compute dht of current rank and pass it to the backward kernel - dht, initial_state = chunk_gated_delta_rule_bwd_dhu_pre_process( - q=q, - k=k, - w=w, - do=do, - dv=dv, - g=g, - scale=scale, - cu_seqlens=cu_seqlens, - dht=dht, - initial_state=initial_state, - context=cp_context, - state_v_first=state_v_first, - chunk_size=chunk_size, - ) - - dh, dh0, dv = chunk_gated_delta_rule_bwd_dhu( - q=q, - k=k, - w=w, - g=g, - h0=initial_state, - dht=dht, - do=do, - dv=dv, - scale=scale, - cu_seqlens=cu_seqlens, - chunk_indices=chunk_indices, - state_v_first=state_v_first, - chunk_size=chunk_size, - ) - dq, dk, dw, dg = chunk_bwd_dqkwg( - q=q, - k=k, - v=v_new, - w=w, - g=g, - h=h, - dv=dv, - do=do, - dh=dh, - scale=scale, - cu_seqlens=cu_seqlens, - chunk_indices=chunk_indices, - state_v_first=state_v_first, - chunk_size=chunk_size, - ) - dk2, dv, db, dg2 = prepare_wy_repr_bwd( - k=k, - v=v, - beta=beta, - g=g, - A=A, - dw=dw, - du=dv, - cu_seqlens=cu_seqlens, - chunk_indices=chunk_indices, - ) - dk.add_(dk2) - dg.add_(dg2) - dg = chunk_local_cumsum(dg, chunk_size=chunk_size, reverse=True, cu_seqlens=cu_seqlens, chunk_indices=chunk_indices) - dA_log, ddt_bias = None, None - if use_gate_in_kernel: - dg, dA_log, ddt_bias = gdn_gate_bwd(g=g_input, A_log=A_log, dt_bias=dt_bias, dyg=dg) - return dq, dk, dv, db, dg, dh0, dA_log, ddt_bias - - -class ChunkGatedDeltaRuleFunction(torch.autograd.Function): - - @staticmethod - @input_guard - @autocast_custom_fwd - def forward( - ctx, - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - g: torch.Tensor, - beta: torch.Tensor, - scale: float, - initial_state: torch.Tensor, - output_final_state: bool, - state_v_first: bool = False, - cu_seqlens: torch.LongTensor | None = None, - cu_seqlens_cpu: torch.LongTensor | None = None, - use_qk_l2norm_in_kernel: bool = False, - use_gate_in_kernel: bool = False, - A_log: torch.Tensor | None = None, - dt_bias: torch.Tensor | None = None, - use_beta_sigmoid_in_kernel: bool = False, - allow_neg_eigval: bool = False, - cp_context: FLACPContext | None = None, - chunk_size: int = 64, - ): - q_rstd, k_rstd = None, None - if use_qk_l2norm_in_kernel: - q, q_rstd = l2norm_fwd(q) - k, k_rstd = l2norm_fwd(k) - - beta_raw = beta - if use_beta_sigmoid_in_kernel: - beta = fused_beta_sigmoid(beta_raw, scale=2.0 if allow_neg_eigval else 1.0) - - chunk_indices = None - if cu_seqlens is not None: - chunk_indices = prepare_chunk_indices(cu_seqlens, chunk_size, cu_seqlens_cpu=cu_seqlens_cpu) - g, o, A, final_state, initial_state, g_input = chunk_gated_delta_rule_fwd( - q=q, - k=k, - v=v, - g=g, - beta=beta, - scale=scale, - initial_state=initial_state, - output_final_state=output_final_state, - cu_seqlens=cu_seqlens, - cp_context=cp_context, - chunk_indices=chunk_indices, - state_v_first=state_v_first, - use_gate_in_kernel=use_gate_in_kernel, - A_log=A_log, - dt_bias=dt_bias, - chunk_size=chunk_size, - ) - ctx.save_for_backward( - q, - q_rstd, - k, - k_rstd, - v, - g, - beta_raw, - beta, - A, - initial_state, - cu_seqlens, - chunk_indices, - g_input, - A_log, - dt_bias, - ) - ctx.scale = scale - ctx.chunk_size = chunk_size - ctx.use_qk_l2norm_in_kernel = use_qk_l2norm_in_kernel - ctx.use_beta_sigmoid_in_kernel = use_beta_sigmoid_in_kernel - ctx.allow_neg_eigval = allow_neg_eigval - ctx.cp_context = cp_context - ctx.state_v_first = state_v_first - ctx.use_gate_in_kernel = use_gate_in_kernel - return o.to(q.dtype), final_state - - @staticmethod - @input_guard - @autocast_custom_bwd - def backward( - ctx, - do: torch.Tensor, - dht: torch.Tensor, - ): - ( - q, - q_rstd, - k, - k_rstd, - v, - g, - beta_raw, - beta, - A, - initial_state, - cu_seqlens, - chunk_indices, - g_input, - A_log, - dt_bias, - ) = ctx.saved_tensors - dq, dk, dv, db, dg, dh0, dA_log, ddt_bias = chunk_gated_delta_rule_bwd( - q=q, - k=k, - v=v, - g=g, - beta=beta, - A=A, - scale=ctx.scale, - initial_state=initial_state, - do=do, - dht=dht, - cu_seqlens=cu_seqlens, - cp_context=ctx.cp_context, - chunk_indices=chunk_indices, - state_v_first=ctx.state_v_first, - use_gate_in_kernel=ctx.use_gate_in_kernel, - g_input=g_input, - A_log=A_log, - dt_bias=dt_bias, - chunk_size=ctx.chunk_size, - ) - if ctx.use_qk_l2norm_in_kernel: - dq = l2norm_bwd(q, q_rstd, dq) - dk = l2norm_bwd(k, k_rstd, dk) - if ctx.use_beta_sigmoid_in_kernel: - db = fused_beta_sigmoid_bwd(beta_raw, db, scale=2.0 if ctx.allow_neg_eigval else 1.0) - return ( - dq.to(q), dk.to(k), dv.to(v), dg.to(g), db.to(beta_raw), - None, dh0, None, None, None, None, None, None, dA_log, ddt_bias, - None, None, None, None, - ) - - -@torch.compiler.disable -def chunk_gated_delta_rule( - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - g: torch.Tensor, - beta: torch.Tensor, - scale: float | None = None, - initial_state: torch.Tensor | None = None, - output_final_state: bool = False, - use_qk_l2norm_in_kernel: bool = False, - use_beta_sigmoid_in_kernel: bool = False, - allow_neg_eigval: bool = False, - state_v_first: bool = False, - cu_seqlens: torch.LongTensor | None = None, - cu_seqlens_cpu: torch.LongTensor | None = None, - cp_context: FLACPContext | None = None, - **kwargs, -): - r""" - Args: - q (torch.Tensor): - queries of shape `[B, T, H, K]`. - k (torch.Tensor): - keys of shape `[B, T, H, K]`. - v (torch.Tensor): - values of shape `[B, T, HV, V]`. - GVA (Grouped Value Attention) is applied if `HV > H`, where `HV` must be divisible by `H`. - g (torch.Tensor): - (forget) gating tensor of shape `[B, T, HV]`. - When `use_gate_in_kernel=False` (default), `g` should be in log space (pre-computed decay). - When `use_gate_in_kernel=True`, `g` is the raw input before gate activation; - the kernel fuses `-exp(A_log) * softplus(g + dt_bias)` + chunk cumsum internally. - beta (torch.Tensor): - betas of shape `[B, T, HV]`. - scale (Optional[float]): - Scale factor for the RetNet attention scores. - If not provided, it will default to `1 / sqrt(K)`. Default: `None`. - initial_state (Optional[torch.Tensor]): - Initial state of shape `[N, HV, K, V]` for `N` input sequences. - For equal-length input sequences, `N` equals the batch size `B`. - Default: `None`. - output_final_state (Optional[bool]): - Whether to output the final state of shape `[N, HV, K, V]`. Default: `False`. - use_qk_l2norm_in_kernel (bool): - Whether to apply L2norm to the q/k tensor internally. Default: `False`. - use_gate_in_kernel (bool): - Whether to compute the log-space GDN decay internally. - When `True`, the passed `g` is the raw input, and `A_log` must be provided. - The kernel fuses gate activation + chunk cumsum in a single pass. - Default: `False`. - A_log (Optional[torch.Tensor]): - Decay parameter of shape `[HV]`. Required when `use_gate_in_kernel=True`. - dt_bias (Optional[torch.Tensor]): - Bias added to `g` before activation, of shape `[HV]`. - Only used when `use_gate_in_kernel=True`. - use_beta_sigmoid_in_kernel (bool): - Whether to apply `torch.sigmoid(beta)` before launching the chunk kernel. - - If `True`, the passed `beta` acts as the raw beta logits. - - If `False`, `beta` is expected to already be in post-sigmoid space. - Default: `False`. - allow_neg_eigval (bool): - Whether to allow negative eigenvalues by scaling `beta` to `[0, 2)`. - Only takes effect together with `use_beta_sigmoid_in_kernel=True`, in which case - the kernel computes `2 * sigmoid(beta)` instead of `sigmoid(beta)`. Default: `False`. - state_v_first (Optional[bool]): - Store the recurrent state in V-first ``[V, K]`` layout instead of the default ``[K, V]``. Default: ``False``. - cu_seqlens (torch.LongTensor): - Cumulative sequence lengths of shape `[N+1]` used for variable-length training, - consistent with the FlashAttention API. - cp_context (Optional[FLACPContext]): - Context parallel context for distributed training across multiple devices. - When provided, `initial_state` and `output_final_state` are not supported, - and `cu_seqlens` will be overridden by the context. Default: `None`. - - Returns: - o (torch.Tensor): - Outputs of shape `[B, T, HV, V]`. - final_state (torch.Tensor): - Final state of shape `[N, HV, K, V]` if `output_final_state=True` else `None`. - - Examples:: - >>> import torch - >>> import torch.nn.functional as F - >>> from einops import rearrange - >>> from ...ops.gated_delta_rule import chunk_gated_delta_rule - # inputs with equal lengths - >>> B, T, H, HV, K, V = 4, 2048, 4, 8, 512, 512 - >>> q = torch.randn(B, T, H, K, dtype=torch.bfloat16, device='cuda') - >>> k = F.normalize(torch.randn(B, T, H, K, dtype=torch.bfloat16, device='cuda'), p=2, dim=-1) - >>> v = torch.randn(B, T, HV, V, dtype=torch.bfloat16, device='cuda') - >>> beta = torch.rand(B, T, HV, dtype=torch.bfloat16, device='cuda').sigmoid() - >>> g = F.logsigmoid(torch.rand(B, T, HV, dtype=torch.bfloat16, device='cuda')) - >>> h0 = torch.randn(B, HV, K, V, dtype=torch.bfloat16, device='cuda') - >>> o, ht = chunk_gated_delta_rule( - q, k, v, g, beta, - initial_state=h0, - output_final_state=True - ) - # for variable-length inputs, the batch size `B` is expected to be 1 and `cu_seqlens` is required - >>> q, k, v, beta, g = map(lambda x: rearrange(x, 'b t ... -> 1 (b t) ...'), (q, k, v, beta, g)) - # for a batch with 4 sequences, `cu_seqlens` with 5 start/end positions are expected - >>> cu_seqlens = q.new_tensor([0, 2048, 4096, 6144, 8192], dtype=torch.long) - >>> o, ht = chunk_gated_delta_rule( - q, k, v, g, beta, - initial_state=h0, - output_final_state=True, - cu_seqlens=cu_seqlens - ) - """ - if 'transpose_state_layout' in kwargs: - if state_v_first: - raise ValueError("Cannot pass both `state_v_first` and the deprecated `transpose_state_layout`.") - warnings.warn( - "`transpose_state_layout` is deprecated and renamed to `state_v_first`.", - DeprecationWarning, - stacklevel=2, - ) - state_v_first = kwargs.pop('transpose_state_layout') - - # Validate head dimensions - if q.shape[2] != k.shape[2]: - raise ValueError( - f"q and k must have the same number of heads, " - f"but got q.shape[2]={q.shape[2]} and k.shape[2]={k.shape[2]}" - ) - H, HV = q.shape[2], v.shape[2] - if HV % H != 0: - raise ValueError( - f"For GVA, num_v_heads (HV={HV}) must be evenly divisible by " - f"num_heads (H={H}), but got HV % H = {HV % H}" - ) - - if 'head_first' in kwargs: - raise DeprecationWarning( - "head_first has been removed. Inputs must be in `[B, T, H, ...]` format.", - ) - - chunk_size = kwargs.pop('chunk_size', 64) - if chunk_size not in (16, 32, 64): - raise ValueError(f"`chunk_size` must be 16, 32, or 64 for Gated Delta Rule, got {chunk_size}.") - - 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" - assert cp_context.cu_seqlens is not None, "cu_seqlens is required for CP" - cu_seqlens = cp_context.cu_seqlens - if cp_context.cu_seqlens_cpu is not None: - cu_seqlens_cpu = cp_context.cu_seqlens_cpu - - if cu_seqlens is not None: - if q.shape[0] != 1: - raise ValueError( - f"The batch size is expected to be 1 rather than {q.shape[0]} when using `cu_seqlens`." - f"Please flatten variable-length inputs before processing.", - ) - if initial_state is not None and initial_state.shape[0] != len(cu_seqlens) - 1: - raise ValueError( - f"The number of initial states is expected to be equal to the number of input sequences, " - f"i.e., {len(cu_seqlens) - 1} rather than {initial_state.shape[0]}.", - ) - use_gate_in_kernel = kwargs.get('use_gate_in_kernel', False) - A_log = kwargs.get('A_log') - dt_bias = kwargs.get('dt_bias') - if use_gate_in_kernel: - assert A_log is not None, "A_log must be provided when use_gate_in_kernel=True." - if allow_neg_eigval and not use_beta_sigmoid_in_kernel: - raise ValueError("`allow_neg_eigval=True` requires `use_beta_sigmoid_in_kernel=True`.") - - if scale is None: - scale = k.shape[-1] ** -0.5 - o, final_state = ChunkGatedDeltaRuleFunction.apply( - q, - k, - v, - g, - beta, - scale, - initial_state, - output_final_state, - state_v_first, - cu_seqlens, - cu_seqlens_cpu, - use_qk_l2norm_in_kernel, - use_gate_in_kernel, - A_log, - dt_bias, - use_beta_sigmoid_in_kernel, - allow_neg_eigval, - cp_context, - chunk_size, - ) - return o, final_state - - -chunk_gdn = chunk_gated_delta_rule diff --git a/build/torch-xpu/ops/gated_delta_rule/chunk_fwd.py b/build/torch-xpu/ops/gated_delta_rule/chunk_fwd.py deleted file mode 100644 index 55febabc537687f7c71a5b97ff0eb76bdf90c067..0000000000000000000000000000000000000000 --- a/build/torch-xpu/ops/gated_delta_rule/chunk_fwd.py +++ /dev/null @@ -1,414 +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 -import triton.language as tl - -from ...ops.common.chunk_scaled_dot_kkt import chunk_scaled_dot_kkt_fwd -from ...ops.gated_delta_rule.wy_fast import recompute_w_u_fwd -from ...ops.utils import prepare_chunk_indices, solve_tril -from ...ops.utils.cache import fla_cache_autotune -from ...ops.utils.op import exp2 -from ...utils import IS_INTEL, IS_TF32_SUPPORTED, autotune_cache_kwargs - -if IS_TF32_SUPPORTED: - SOLVE_TRIL_DOT_PRECISION = tl.constexpr('tf32') -else: - SOLVE_TRIL_DOT_PRECISION = tl.constexpr('ieee') - - -@triton.heuristics({ - 'USE_G': lambda args: args['g'] is not None, - 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None, -}) -@fla_cache_autotune( - configs=[ - triton.Config({'BK': BK}, num_warps=num_warps) - for BK in [32, 64] - for num_warps in [1, 2, 4] - ], - key=['H', 'HV', 'K', 'BC'], - **autotune_cache_kwargs, -) -@triton.jit(do_not_specialize=['T']) -def chunk_gated_delta_rule_fwd_kkt_solve_kernel( - k, - g, - beta, - A, - cu_seqlens, - chunk_indices, - T, - H: tl.constexpr, - HV: tl.constexpr, - K: tl.constexpr, - BT: tl.constexpr, - BC: tl.constexpr, - BK: tl.constexpr, - USE_G: tl.constexpr, - IS_VARLEN: tl.constexpr, -): - """ - Fused kernel: compute beta * K @ K^T (lower triangular) + solve_tril (I+A)^{-1} in one pass. - - This kernel fuses chunk_scaled_dot_kkt_fwd and solve_tril into a single kernel, - avoiding the HBM round-trip for the intermediate A matrix. - - Steps: - 1. Compute all 10 lower-triangular [BC, BC] blocks of beta * K @ K^T in registers - 2. Apply gate and beta scaling - 3. Forward substitution on diagonal blocks - 4. Block merge to get full (I+A)^{-1} - 5. Write result to A (output) - """ - i_t, i_bh = tl.program_id(0), tl.program_id(1) - i_b, i_h = i_bh // HV, i_bh % HV - - 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.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) - T = eos - bos - else: - bos, eos = i_b * T, i_b * T + T - - if i_t * BT >= T: - return - - i_tc0 = i_t * BT - i_tc1 = i_t * BT + BC - i_tc2 = i_t * BT + 2 * BC - i_tc3 = i_t * BT + 3 * BC - - k += (bos * H + i_h // (HV // H)) * K - A += (bos * HV + i_h) * BT - - o_i = tl.arange(0, BC) - m_tc0 = (i_tc0 + o_i) < T - m_tc1 = (i_tc1 + o_i) < T - m_tc2 = (i_tc2 + o_i) < T - m_tc3 = (i_tc3 + o_i) < T - - # load beta for each sub-chunk - p_b0 = tl.make_block_ptr(beta + bos * HV + i_h, (T,), (HV,), (i_tc0,), (BC,), (0,)) - p_b1 = tl.make_block_ptr(beta + bos * HV + i_h, (T,), (HV,), (i_tc1,), (BC,), (0,)) - p_b2 = tl.make_block_ptr(beta + bos * HV + i_h, (T,), (HV,), (i_tc2,), (BC,), (0,)) - p_b3 = tl.make_block_ptr(beta + bos * HV + i_h, (T,), (HV,), (i_tc3,), (BC,), (0,)) - b_b0 = tl.load(p_b0, boundary_check=(0,)).to(tl.float32) - b_b1 = tl.load(p_b1, boundary_check=(0,)).to(tl.float32) - b_b2 = tl.load(p_b2, boundary_check=(0,)).to(tl.float32) - b_b3 = tl.load(p_b3, boundary_check=(0,)).to(tl.float32) - - # load gate if used - if USE_G: - p_g0 = tl.make_block_ptr(g + bos * HV + i_h, (T,), (HV,), (i_tc0,), (BC,), (0,)) - p_g1 = tl.make_block_ptr(g + bos * HV + i_h, (T,), (HV,), (i_tc1,), (BC,), (0,)) - p_g2 = tl.make_block_ptr(g + bos * HV + i_h, (T,), (HV,), (i_tc2,), (BC,), (0,)) - p_g3 = tl.make_block_ptr(g + bos * HV + i_h, (T,), (HV,), (i_tc3,), (BC,), (0,)) - - b_g0 = tl.load(p_g0, boundary_check=(0,)).to(tl.float32) - b_g1 = tl.load(p_g1, boundary_check=(0,)).to(tl.float32) - b_g2 = tl.load(p_g2, boundary_check=(0,)).to(tl.float32) - b_g3 = tl.load(p_g3, boundary_check=(0,)).to(tl.float32) - - ############################################################################ - # Step 1: compute all 10 lower-triangular [BC, BC] blocks of K @ K^T - ############################################################################ - - # 4 diagonal blocks - b_A00 = tl.zeros([BC, BC], dtype=tl.float32) - b_A11 = tl.zeros([BC, BC], dtype=tl.float32) - b_A22 = tl.zeros([BC, BC], dtype=tl.float32) - b_A33 = tl.zeros([BC, BC], dtype=tl.float32) - - # 6 off-diagonal blocks - b_A10 = tl.zeros([BC, BC], dtype=tl.float32) - b_A20 = tl.zeros([BC, BC], dtype=tl.float32) - b_A21 = tl.zeros([BC, BC], dtype=tl.float32) - b_A30 = tl.zeros([BC, BC], dtype=tl.float32) - b_A31 = tl.zeros([BC, BC], dtype=tl.float32) - b_A32 = tl.zeros([BC, BC], dtype=tl.float32) - - for i_k in range(tl.cdiv(K, BK)): - p_k0 = tl.make_block_ptr(k, (T, K), (H*K, 1), (i_tc0, i_k * BK), (BC, BK), (1, 0)) - b_k0 = tl.load(p_k0, boundary_check=(0, 1)) - # diagonal block 0 - b_A00 += tl.dot(b_k0, tl.trans(b_k0)) - - if i_tc1 < T: - p_k1 = tl.make_block_ptr(k, (T, K), (H*K, 1), (i_tc1, i_k * BK), (BC, BK), (1, 0)) - b_k1 = tl.load(p_k1, boundary_check=(0, 1)) - # diagonal block 1 - b_A11 += tl.dot(b_k1, tl.trans(b_k1)) - # off-diagonal (1,0) - b_A10 += tl.dot(b_k1, tl.trans(b_k0)) - - if i_tc2 < T: - p_k2 = tl.make_block_ptr(k, (T, K), (H*K, 1), (i_tc2, i_k * BK), (BC, BK), (1, 0)) - b_k2 = tl.load(p_k2, boundary_check=(0, 1)) - # diagonal block 2 - b_A22 += tl.dot(b_k2, tl.trans(b_k2)) - # off-diagonal (2,0), (2,1) - b_A20 += tl.dot(b_k2, tl.trans(b_k0)) - b_A21 += tl.dot(b_k2, tl.trans(b_k1)) - - if i_tc3 < T: - p_k3 = tl.make_block_ptr(k, (T, K), (H*K, 1), (i_tc3, i_k * BK), (BC, BK), (1, 0)) - b_k3 = tl.load(p_k3, boundary_check=(0, 1)) - # diagonal block 3 - b_A33 += tl.dot(b_k3, tl.trans(b_k3)) - # off-diagonal (3,0), (3,1), (3,2) - b_A30 += tl.dot(b_k3, tl.trans(b_k0)) - b_A31 += tl.dot(b_k3, tl.trans(b_k1)) - b_A32 += tl.dot(b_k3, tl.trans(b_k2)) - - ############################################################################ - # Step 2: apply gate and beta scaling - ############################################################################ - - # apply gate, beta scaling, and masking - # m_d: strictly lower triangular mask for diagonal blocks - # m_tc: boundary mask to prevent NaN from 0 * inf (IEEE 754) when - # out-of-bounds g loads as 0 via boundary_check and exp2(0 - g_inbounds) overflows - m_d = o_i[:, None] > o_i[None, :] - m_I = o_i[:, None] == o_i[None, :] - - if USE_G: - b_A00 *= tl.where(m_d & m_tc0[:, None] & m_tc0[None, :], exp2(b_g0[:, None] - b_g0[None, :]), 0.) - b_A11 *= tl.where(m_d & m_tc1[:, None] & m_tc1[None, :], exp2(b_g1[:, None] - b_g1[None, :]), 0.) - b_A22 *= tl.where(m_d & m_tc2[:, None] & m_tc2[None, :], exp2(b_g2[:, None] - b_g2[None, :]), 0.) - b_A33 *= tl.where(m_d & m_tc3[:, None] & m_tc3[None, :], exp2(b_g3[:, None] - b_g3[None, :]), 0.) - - b_A10 *= tl.where(m_tc1[:, None] & m_tc0[None, :], exp2(b_g1[:, None] - b_g0[None, :]), 0.) - b_A20 *= tl.where(m_tc2[:, None] & m_tc0[None, :], exp2(b_g2[:, None] - b_g0[None, :]), 0.) - b_A21 *= tl.where(m_tc2[:, None] & m_tc1[None, :], exp2(b_g2[:, None] - b_g1[None, :]), 0.) - b_A30 *= tl.where(m_tc3[:, None] & m_tc0[None, :], exp2(b_g3[:, None] - b_g0[None, :]), 0.) - b_A31 *= tl.where(m_tc3[:, None] & m_tc1[None, :], exp2(b_g3[:, None] - b_g1[None, :]), 0.) - b_A32 *= tl.where(m_tc3[:, None] & m_tc2[None, :], exp2(b_g3[:, None] - b_g2[None, :]), 0.) - else: - b_A00 = tl.where(m_d, b_A00, 0.) - b_A11 = tl.where(m_d, b_A11, 0.) - b_A22 = tl.where(m_d, b_A22, 0.) - b_A33 = tl.where(m_d, b_A33, 0.) - - # diagonal blocks: scaled by beta - b_A00 = b_A00 * b_b0[:, None] - b_A11 = b_A11 * b_b1[:, None] - b_A22 = b_A22 * b_b2[:, None] - b_A33 = b_A33 * b_b3[:, None] - - # off-diagonal blocks: full block, scaled by beta - b_A10 = b_A10 * b_b1[:, None] - b_A20 = b_A20 * b_b2[:, None] - b_A21 = b_A21 * b_b2[:, None] - b_A30 = b_A30 * b_b3[:, None] - b_A31 = b_A31 * b_b3[:, None] - b_A32 = b_A32 * b_b3[:, None] - - ############################################################################ - # Step 3: forward substitution on diagonal blocks -> (I + A_diag)^{-1} - # - # Same algorithm as solve_tril, but rows are extracted from in-register - # [BC, BC] tensor via tl.sum(tl.where(mask, tensor, 0), 0) instead of - # tl.load from HBM. - ############################################################################ - - b_Ai00 = -b_A00 - b_Ai11 = -b_A11 - b_Ai22 = -b_A22 - b_Ai33 = -b_A33 - - for i in range(2, min(BC, T - i_tc0)): - b_a00 = tl.sum(tl.where((o_i == i)[:, None], -b_A00, 0.), 0) - b_a00 = tl.where(o_i < i, b_a00, 0.) - b_a00 = b_a00 + tl.sum(b_a00[:, None] * b_Ai00, 0) - b_Ai00 = tl.where((o_i == i)[:, None], b_a00, b_Ai00) - for i in range(2, min(BC, T - i_tc1)): - b_a11 = tl.sum(tl.where((o_i == i)[:, None], -b_A11, 0.), 0) - b_a11 = tl.where(o_i < i, b_a11, 0.) - b_a11 = b_a11 + tl.sum(b_a11[:, None] * b_Ai11, 0) - b_Ai11 = tl.where((o_i == i)[:, None], b_a11, b_Ai11) - for i in range(2, min(BC, T - i_tc2)): - b_a22 = tl.sum(tl.where((o_i == i)[:, None], -b_A22, 0.), 0) - b_a22 = tl.where(o_i < i, b_a22, 0.) - b_a22 = b_a22 + tl.sum(b_a22[:, None] * b_Ai22, 0) - b_Ai22 = tl.where((o_i == i)[:, None], b_a22, b_Ai22) - for i in range(2, min(BC, T - i_tc3)): - b_a33 = tl.sum(tl.where((o_i == i)[:, None], -b_A33, 0.), 0) - b_a33 = tl.where(o_i < i, b_a33, 0.) - b_a33 = b_a33 + tl.sum(b_a33[:, None] * b_Ai33, 0) - b_Ai33 = tl.where((o_i == i)[:, None], b_a33, b_Ai33) - - b_Ai00 += m_I - b_Ai11 += m_I - b_Ai22 += m_I - b_Ai33 += m_I - - ############################################################################ - # Step 4: block merge -> full (I + A)^{-1} - ############################################################################ - - b_Ai10 = -tl.dot( - tl.dot(b_Ai11, b_A10, input_precision=SOLVE_TRIL_DOT_PRECISION), - b_Ai00, - input_precision=SOLVE_TRIL_DOT_PRECISION - ) - b_Ai21 = -tl.dot( - tl.dot(b_Ai22, b_A21, input_precision=SOLVE_TRIL_DOT_PRECISION), - b_Ai11, - input_precision=SOLVE_TRIL_DOT_PRECISION - ) - b_Ai32 = -tl.dot( - tl.dot(b_Ai33, b_A32, input_precision=SOLVE_TRIL_DOT_PRECISION), - b_Ai22, - input_precision=SOLVE_TRIL_DOT_PRECISION - ) - - b_Ai20 = -tl.dot( - b_Ai22, - tl.dot(b_A20, b_Ai00, input_precision=SOLVE_TRIL_DOT_PRECISION) + - tl.dot(b_A21, b_Ai10, input_precision=SOLVE_TRIL_DOT_PRECISION), - input_precision=SOLVE_TRIL_DOT_PRECISION, - ) - b_Ai31 = -tl.dot( - b_Ai33, - tl.dot(b_A31, b_Ai11, input_precision=SOLVE_TRIL_DOT_PRECISION) + - tl.dot(b_A32, b_Ai21, input_precision=SOLVE_TRIL_DOT_PRECISION), - input_precision=SOLVE_TRIL_DOT_PRECISION, - ) - b_Ai30 = -tl.dot( - b_Ai33, - tl.dot(b_A30, b_Ai00, input_precision=SOLVE_TRIL_DOT_PRECISION) + - tl.dot(b_A31, b_Ai10, input_precision=SOLVE_TRIL_DOT_PRECISION) + - tl.dot(b_A32, b_Ai20, input_precision=SOLVE_TRIL_DOT_PRECISION), - input_precision=SOLVE_TRIL_DOT_PRECISION, - ) - - ############################################################################ - # Step 5: store full (I + A)^{-1} to output A - ############################################################################ - - p_A00 = tl.make_block_ptr(A, (T, BT), (HV*BT, 1), (i_tc0, 0), (BC, BC), (1, 0)) - p_A10 = tl.make_block_ptr(A, (T, BT), (HV*BT, 1), (i_tc1, 0), (BC, BC), (1, 0)) - p_A11 = tl.make_block_ptr(A, (T, BT), (HV*BT, 1), (i_tc1, BC), (BC, BC), (1, 0)) - p_A20 = tl.make_block_ptr(A, (T, BT), (HV*BT, 1), (i_tc2, 0), (BC, BC), (1, 0)) - p_A21 = tl.make_block_ptr(A, (T, BT), (HV*BT, 1), (i_tc2, BC), (BC, BC), (1, 0)) - p_A22 = tl.make_block_ptr(A, (T, BT), (HV*BT, 1), (i_tc2, 2*BC), (BC, BC), (1, 0)) - p_A30 = tl.make_block_ptr(A, (T, BT), (HV*BT, 1), (i_tc3, 0), (BC, BC), (1, 0)) - p_A31 = tl.make_block_ptr(A, (T, BT), (HV*BT, 1), (i_tc3, BC), (BC, BC), (1, 0)) - p_A32 = tl.make_block_ptr(A, (T, BT), (HV*BT, 1), (i_tc3, 2*BC), (BC, BC), (1, 0)) - p_A33 = tl.make_block_ptr(A, (T, BT), (HV*BT, 1), (i_tc3, 3*BC), (BC, BC), (1, 0)) - - tl.store(p_A00, b_Ai00.to(A.dtype.element_ty), boundary_check=(0, 1)) - tl.store(p_A10, b_Ai10.to(A.dtype.element_ty), boundary_check=(0, 1)) - tl.store(p_A11, b_Ai11.to(A.dtype.element_ty), boundary_check=(0, 1)) - tl.store(p_A20, b_Ai20.to(A.dtype.element_ty), boundary_check=(0, 1)) - tl.store(p_A21, b_Ai21.to(A.dtype.element_ty), boundary_check=(0, 1)) - tl.store(p_A22, b_Ai22.to(A.dtype.element_ty), boundary_check=(0, 1)) - tl.store(p_A30, b_Ai30.to(A.dtype.element_ty), boundary_check=(0, 1)) - tl.store(p_A31, b_Ai31.to(A.dtype.element_ty), boundary_check=(0, 1)) - tl.store(p_A32, b_Ai32.to(A.dtype.element_ty), boundary_check=(0, 1)) - tl.store(p_A33, b_Ai33.to(A.dtype.element_ty), boundary_check=(0, 1)) - - -def chunk_gated_delta_rule_fwd_intra( - k: torch.Tensor, - v: torch.Tensor, - g: torch.Tensor | None = None, - beta: torch.Tensor | None = None, - cu_seqlens: torch.LongTensor | None = None, - chunk_size: int = 64, - chunk_indices: torch.LongTensor | None = None, -) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: - r""" - GDN intra-chunk forward: fused or unfused kkt + solve_tril + recompute_w_u. - - For ``chunk_size == 64``, this uses the fused kkt + solve_tril path. For - other supported chunk sizes, it computes the mathematically equivalent - representation with ``chunk_scaled_dot_kkt_fwd`` followed by ``solve_tril``. - - Args: - k (torch.Tensor): - The key tensor of shape `[B, T, H, K]`. - v (torch.Tensor): - The value tensor of shape `[B, T, HV, V]`. - g (torch.Tensor): - The cumulative sum of the gate tensor of shape `[B, T, HV]`. Default: `None`. - beta (torch.Tensor): - The beta tensor of shape `[B, T, HV]`. - cu_seqlens (torch.LongTensor): - The cumulative sequence lengths. Default: `None`. - chunk_size (int): - The chunk size. Default: 64. - chunk_indices (torch.LongTensor): - Precomputed chunk indices. Default: `None`. - - Returns: - w (torch.Tensor): shape `[B, T, HV, K]` - u (torch.Tensor): shape `[B, T, HV, V]` - A (torch.Tensor): shape `[B, T, HV, BT]`, the solved (I+A)^{-1} matrix - """ - if chunk_size not in (16, 32, 64): - raise ValueError(f"`chunk_size` must be 16, 32, or 64, got {chunk_size}.") - - B, T, H, K, HV = *k.shape, beta.shape[2] - BT = chunk_size - - if chunk_indices is None and cu_seqlens is not None: - chunk_indices = prepare_chunk_indices(cu_seqlens, BT) - - # The fused kernel keeps ten [BC, BC] fp32 accumulators live across the K loop. - # That fits NVIDIA's register file but spills on Intel GPUs, where the unfused - # two-kernel path measures 2.3-3.0x faster despite the extra HBM round-trip. - if BT == 64 and not IS_INTEL: - # Step 1: fused kkt + solve_tril - BC = 16 - NT = triton.cdiv(T, BT) if cu_seqlens is None else len(chunk_indices) - A = torch.zeros(B, T, HV, BT, device=k.device, dtype=k.dtype) - chunk_gated_delta_rule_fwd_kkt_solve_kernel[(NT, B * HV)]( - k=k, - g=g, - beta=beta, - A=A, - cu_seqlens=cu_seqlens, - chunk_indices=chunk_indices, - T=T, - H=H, - HV=HV, - K=K, - BT=BT, - BC=BC, - ) - else: - # Step 1: mathematically equivalent unfused kkt + solve_tril - A = chunk_scaled_dot_kkt_fwd( - k=k, - g=g, - beta=beta, - cu_seqlens=cu_seqlens, - chunk_indices=chunk_indices, - chunk_size=BT, - output_dtype=torch.float32, - ) - A = solve_tril( - A=A, - cu_seqlens=cu_seqlens, - chunk_indices=chunk_indices, - output_dtype=k.dtype, - ) - - # Step 2: recompute_w_u - w, u = recompute_w_u_fwd( - k=k, - v=v, - beta=beta, - A=A, - g=g, - cu_seqlens=cu_seqlens, - chunk_indices=chunk_indices, - ) - return w, u, A diff --git a/build/torch-xpu/ops/gated_delta_rule/gate.py b/build/torch-xpu/ops/gated_delta_rule/gate.py deleted file mode 100644 index bf478e67dc1efdeacf56b88d21b1dd0243124534..0000000000000000000000000000000000000000 --- a/build/torch-xpu/ops/gated_delta_rule/gate.py +++ /dev/null @@ -1,334 +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 torch.nn.functional as F -import triton -import triton.language as tl - -from ...ops.utils.cache import fla_cache_autotune -from ...ops.utils.index import prepare_chunk_indices -from ...ops.utils.op import exp -from ...ops.utils.softplus import softplus -from ...utils import autocast_custom_bwd, autocast_custom_fwd, autotune_cache_kwargs, input_guard - - -def naive_gdn_gate( - g: torch.Tensor, - A_log: torch.Tensor, - dt_bias: torch.Tensor | None = None, - output_dtype: torch.dtype = torch.float32, -) -> torch.Tensor: - """ - Torch reference implementation for GDN gate computation. - - Computes: ``g = -A_log.exp() * softplus(g + dt_bias)`` - - Args: - g (torch.Tensor): - Input tensor of shape `[..., HV]`. - A_log (torch.Tensor): - Decay parameter tensor with `HV` elements. - dt_bias (torch.Tensor | None): - Optional bias tensor added to `g` before activation, shape `[HV]`. - - Returns: - Output tensor of shape `[..., HV]`. - """ - g = g.float() - if dt_bias is not None: - g = g + dt_bias.float() - return (-A_log.float().exp() * F.softplus(g)).to(output_dtype) - - -@triton.heuristics({ - 'HAS_BIAS': lambda args: args['dt_bias'] is not None, - 'HAS_SCALE': lambda args: args['scale'] is not None, - 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None, -}) -@fla_cache_autotune( - configs=[ - triton.Config({}, num_warps=num_warps) - for num_warps in [1, 2, 4, 8] - ], - key=['H', 'BT', 'IS_VARLEN', 'REVERSE'], - **autotune_cache_kwargs, -) -@triton.jit(do_not_specialize=['T']) -def gdn_gate_chunk_cumsum_scalar_kernel( - g, - A_log, - dt_bias, - o, - scale, - cu_seqlens, - chunk_indices, - T, - H: tl.constexpr, - BT: tl.constexpr, - REVERSE: tl.constexpr, - HAS_BIAS: tl.constexpr, - HAS_SCALE: tl.constexpr, - IS_VARLEN: tl.constexpr, -): - i_t, i_bh = tl.program_id(0), tl.program_id(1) - i_b, i_h = i_bh // H, i_bh % H - - 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.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) - T = eos - bos - else: - bos, eos = i_b * T, i_b * T + T - - p_g = tl.make_block_ptr(g + bos * H + i_h, (T,), (H,), (i_t * BT,), (BT,), (0,)) - p_o = tl.make_block_ptr(o + bos * H + i_h, (T,), (H,), (i_t * BT,), (BT,), (0,)) - - b_g = tl.load(p_g, boundary_check=(0,)).to(tl.float32) - if HAS_BIAS: - b_g = b_g + tl.load(dt_bias + i_h).to(tl.float32) - b_A = tl.load(A_log + i_h).to(tl.float32) - b_gate = -exp(b_A) * softplus(b_g) - - b_o = tl.cumsum(b_gate, axis=0) - if REVERSE: - b_z = tl.sum(b_gate, axis=0) - b_o = -b_o + b_z[None] + b_gate - if HAS_SCALE: - b_o *= scale - tl.store(p_o, b_o.to(p_o.dtype.element_ty), boundary_check=(0,)) - - -@triton.heuristics({ - 'HAS_BIAS': lambda args: args['dt_bias'] is not None, -}) -@fla_cache_autotune( - configs=[ - triton.Config({}, num_warps=num_warps) - for num_warps in [1, 2, 4, 8] - ], - key=['H', 'BT'], - **autotune_cache_kwargs, -) -@triton.jit(do_not_specialize=['T']) -def gdn_gate_bwd_kernel( - g, - A_log, - dt_bias, - dyg, - dg, - dA, - T, - H: tl.constexpr, - BT: tl.constexpr, - HAS_BIAS: tl.constexpr, -): - i_t, i_h = tl.program_id(0), tl.program_id(1) - - b_A = tl.load(A_log + i_h).to(tl.float32) - - p_g = tl.make_block_ptr(g + i_h, (T,), (H,), (i_t * BT,), (BT,), (0,)) - p_dg = tl.make_block_ptr(dg + i_h, (T,), (H,), (i_t * BT,), (BT,), (0,)) - p_dyg = tl.make_block_ptr(dyg + i_h, (T,), (H,), (i_t * BT,), (BT,), (0,)) - - b_g = tl.load(p_g, boundary_check=(0,)).to(tl.float32) - b_dyg = tl.load(p_dyg, boundary_check=(0,)).to(tl.float32) - - if HAS_BIAS: - b_g = b_g + tl.load(dt_bias + i_h).to(tl.float32) - - # gate = -exp(A_log) * softplus(g + bias) - # d(gate)/d(g) = -exp(A_log) * sigmoid(g + bias) (softplus' = sigmoid) - # d(gate)/d(A_log) = -exp(A_log) * softplus(g + bias) = gate - b_neg_expA = -exp(b_A) - b_yg = b_neg_expA * softplus(b_g) - b_dg = b_neg_expA * (b_dyg * tl.sigmoid(b_g)) - b_dA = tl.sum(b_dyg * b_yg, 0) - - tl.store(p_dg, b_dg.to(p_dg.dtype.element_ty), boundary_check=(0,)) - tl.store(dA + i_t * H + i_h, b_dA) - - -@input_guard -def gdn_gate_chunk_cumsum( - g: torch.Tensor, - A_log: torch.Tensor, - chunk_size: int, - scale: float = None, - dt_bias: torch.Tensor | None = None, - cu_seqlens: torch.LongTensor | None = None, - chunk_indices: torch.LongTensor | None = None, - output_dtype: torch.dtype | None = torch.float, -) -> torch.Tensor: - B, T, H = g.shape - BT = chunk_size - if chunk_indices is None and cu_seqlens is not None: - chunk_indices = prepare_chunk_indices(cu_seqlens, BT) - NT = triton.cdiv(T, BT) if cu_seqlens is None else len(chunk_indices) - - o = torch.empty_like(g, dtype=output_dtype or g.dtype) - gdn_gate_chunk_cumsum_scalar_kernel[(NT, B * H)]( - g=g, - A_log=A_log, - dt_bias=dt_bias, - o=o, - scale=scale, - cu_seqlens=cu_seqlens, - chunk_indices=chunk_indices, - T=T, - H=H, - BT=BT, - REVERSE=False, - ) - return o - - -def gdn_gate_bwd( - g: torch.Tensor, - A_log: torch.Tensor, - dt_bias: torch.Tensor | None, - dyg: torch.Tensor, -) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor | None]: - H = g.shape[-1] - T = g.numel() // H - BT = 32 - NT = triton.cdiv(T, BT) - - dg = torch.empty_like(g, dtype=torch.float32) - dA = A_log.new_empty(NT, H, dtype=torch.float32) - - gdn_gate_bwd_kernel[(NT, H)]( - g=g, - A_log=A_log, - dt_bias=dt_bias, - dyg=dyg, - dg=dg, - dA=dA, - T=T, - H=H, - BT=BT, - ) - - dg = dg.view_as(g).type_as(g) - dA = dA.sum(0).view_as(A_log).type_as(A_log) - dbias = dg.view(-1, H).sum(0).to(dt_bias) if dt_bias is not None else None - - return dg, dA, dbias - - -@triton.heuristics({ - 'HAS_BIAS': lambda args: args['dt_bias'] is not None, -}) -@fla_cache_autotune( - configs=[ - triton.Config({'BT': BT}, num_warps=num_warps, num_stages=num_stages) - for BT in [32, 64, 128] - for num_warps in [1, 2, 4, 8] - for num_stages in [2, 3] - ], - key=['H'], - **autotune_cache_kwargs, -) -@triton.jit(do_not_specialize=['T']) -def gdn_gate_fwd_kernel( - g, - A_log, - dt_bias, - yg, - T, - H: tl.constexpr, - BT: tl.constexpr, - HAS_BIAS: tl.constexpr, -): - i_t, i_h = tl.program_id(0), tl.program_id(1) - - b_A = tl.load(A_log + i_h).to(tl.float32) - - p_g = tl.make_block_ptr(g + i_h, (T,), (H,), (i_t * BT,), (BT,), (0,)) - p_yg = tl.make_block_ptr(yg + i_h, (T,), (H,), (i_t * BT,), (BT,), (0,)) - b_g = tl.load(p_g, boundary_check=(0,)).to(tl.float32) - if HAS_BIAS: - b_g = b_g + tl.load(dt_bias + i_h).to(tl.float32) - b_yg = -exp(b_A) * softplus(b_g) - tl.store(p_yg, b_yg.to(p_yg.dtype.element_ty), boundary_check=(0,)) - - -def gdn_gate_fwd( - g: torch.Tensor, - A_log: torch.Tensor, - dt_bias: torch.Tensor | None = None, - output_dtype: torch.dtype = torch.float32, -) -> torch.Tensor: - H = g.shape[-1] - T = g.numel() // H - - yg = torch.empty_like(g, dtype=output_dtype) - - def grid(meta): - return (triton.cdiv(T, meta['BT']), H) - - gdn_gate_fwd_kernel[grid]( - g=g, - A_log=A_log, - dt_bias=dt_bias, - yg=yg, - T=T, - H=H, - ) - return yg - - -class GDNGateFunction(torch.autograd.Function): - @staticmethod - @input_guard - @autocast_custom_fwd - def forward( - ctx, - g: torch.Tensor, - A_log: torch.Tensor, - dt_bias: torch.Tensor | None = None, - output_dtype: torch.dtype = torch.float32, - ) -> torch.Tensor: - yg = gdn_gate_fwd(g=g, A_log=A_log, dt_bias=dt_bias, output_dtype=output_dtype) - ctx.save_for_backward(g, A_log, dt_bias) - return yg - - @staticmethod - @input_guard - @autocast_custom_bwd - def backward(ctx, dyg: torch.Tensor): - g, A_log, dt_bias = ctx.saved_tensors - dg, dA, dbias = gdn_gate_bwd(g=g, A_log=A_log, dt_bias=dt_bias, dyg=dyg) - return dg, dA, dbias, None - - -@torch.compiler.disable -def fused_gdn_gate( - g: torch.Tensor, - A_log: torch.Tensor, - dt_bias: torch.Tensor | None = None, - output_dtype: torch.dtype = torch.float32, -) -> torch.Tensor: - r""" - Fused GDN gate computation with autograd support. - - Computes: ``g = -A_log.exp() * softplus(g + dt_bias)`` - - Args: - g (torch.Tensor): - Input tensor of shape `[..., HV]`. - A_log (torch.Tensor): - Decay parameter tensor with `HV` elements. - dt_bias (torch.Tensor | None): - Optional bias tensor added to `g` before activation, shape `[HV]`. - output_dtype (torch.dtype): - The dtype of the output tensor. Default: `torch.float32`. - - Returns: - Output tensor of shape `[..., HV]`. - """ - return GDNGateFunction.apply(g, A_log, dt_bias, output_dtype) diff --git a/build/torch-xpu/ops/gated_delta_rule/naive.py b/build/torch-xpu/ops/gated_delta_rule/naive.py deleted file mode 100644 index cd0cf0d1b15c581e78b20a816601fdde2b6cc70f..0000000000000000000000000000000000000000 --- a/build/torch-xpu/ops/gated_delta_rule/naive.py +++ /dev/null @@ -1,161 +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 torch.nn.functional as F -from einops import rearrange - - -def naive_recurrent_gated_delta_rule( - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - beta: torch.Tensor, - g: torch.Tensor, - scale: float = None, - initial_state: torch.Tensor = None, - output_final_state: bool = False, -): - """ - Reference PyTorch implementation of recurrent gated delta rule. - - Args: - q: [B, T, H, K] - k: [B, T, H, K] - v: [B, T, H, V] - beta: [B, T, H] - g: [B, T, H] - scale: float, optional - initial_state: [B, H, K, V], optional - output_final_state: bool - - Returns: - o: [B, T, H, V] - final_state: [B, H, K, V] if output_final_state else None - """ - q, k, v, beta, g = map(lambda x: x.transpose(1, 2).contiguous().to(torch.float32), [q, k, v, beta, g]) - B, H, T, K, V = *k.shape, v.shape[-1] - o = torch.zeros(B, H, T, V).to(v) - h = torch.zeros(B, H, K, V).to(v) - if initial_state is not None: - h = initial_state.to(torch.float32) - if scale is None: - scale = 1 / (q.shape[-1] ** 0.5) - q = q * scale - - for i in range(T): - b_q = q[:, :, i] - b_k = k[:, :, i] - b_v = v[:, :, i].clone() - h = h.clone() * g[:, :, i].exp()[..., None, None] - b_beta = beta[:, :, i] - b_v = b_v - (h.clone() * b_k[..., None]).sum(-2) - b_v = b_v * b_beta[..., None] - h = h.clone() + b_k.unsqueeze(-1) * b_v.unsqueeze(-2) - o[:, :, i] = torch.einsum('bhd,bhdm->bhm', b_q, h) - - if not output_final_state: - h = None - o = o.transpose(1, 2).contiguous() - return o, h - - -def naive_chunk_gated_delta_rule( - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - g: torch.Tensor, - beta: torch.Tensor, - chunk_size: int = 64, - scale: float = None, - initial_state: torch.Tensor = None, - output_final_state: bool = False, -): - """ - Reference PyTorch implementation of chunk gated delta rule. - - Args: - q: [B, T, H, K] - k: [B, T, H, K] - v: [B, T, H, V] - g: [B, T, H] - beta: [B, T, H] - chunk_size: int - scale: float, optional - initial_state: [B, H, K, V], optional - output_final_state: bool - - Returns: - o: [B, T, H, V] - final_state: [B, H, K, V] if output_final_state else None - """ - BT = chunk_size - if scale is None: - scale = 1 / (q.shape[-1] ** 0.5) - - q, k, v, beta, g = map(lambda x: x.transpose(1, 2).contiguous().to(torch.float32), [q, k, v, beta, g]) - - T = q.shape[-2] - pad_len = (BT - (T % BT)) % BT - if pad_len > 0: - q = F.pad(q, (0, 0, 0, pad_len)) - k = F.pad(k, (0, 0, 0, pad_len)) - v = F.pad(v, (0, 0, 0, pad_len)) - beta = F.pad(beta, (0, pad_len)) - g = F.pad(g, (0, pad_len)) - - q, k, v, beta, g = map(lambda x: x.to(torch.float32), [q, k, v, beta, g]) - decay = g - chunk_size = BT - b, h, l, d_k = q.shape - d_v = v.shape[-1] - q = q * scale - v = v * beta[..., None] - k_beta = k * beta[..., None] - assert l % chunk_size == 0 - - # note that diagonal is masked. - mask = torch.triu(torch.ones(chunk_size, chunk_size, dtype=torch.bool, device=q.device), diagonal=0) - q, k, v, k_beta, decay = map( - lambda x: rearrange(x, 'b h (n c) d -> b h n c d', c=chunk_size), - [q, k, v, k_beta, decay.unsqueeze(-1)], - ) - decay = decay.squeeze(-1).cumsum(-1) - decay_exp = decay.exp()[..., None] - L_mask = ((decay.unsqueeze(-1) - decay.unsqueeze(-2)).tril().exp().float()).tril() - attn = -((k_beta @ k.transpose(-1, -2)) * L_mask).masked_fill(mask, 0) - for i in range(1, chunk_size): - attn[..., i, :i] = attn[..., i, :i].clone() + (attn[..., i, :i, None].clone() * attn[..., :i, :i].clone()).sum(-2) - attn = attn + torch.eye(chunk_size, dtype=torch.float, device=q.device) - attn = attn - k_cumsum = attn @ v - k_cumdecay = attn @ (k_beta * decay_exp) - v = k_cumsum - - S = k.new_zeros(b, h, d_k, d_v) - if initial_state is not None: - S = initial_state.to(torch.float32) - - o = torch.zeros_like(v) - mask = torch.triu(torch.ones(chunk_size, chunk_size, dtype=torch.bool, device=q.device), diagonal=1) - for i in range(0, l // chunk_size): - q_i, k_i, v_i = q[:, :, i], k[:, :, i], v[:, :, i] - attn = (q_i @ k_i.transpose(-1, -2) * L_mask[:, :, i]).masked_fill_(mask, 0) - v_prime = (k_cumdecay[:, :, i]) @ S - v_new = v_i - v_prime - o_inter = (q_i * decay[:, :, i, :, None].exp()) @ S - o[:, :, i] = o_inter + attn @ v_new - S = S * decay[:, :, i, -1, None, None].exp() + (k_i * (decay[:, :, i, -1, None] - decay[:, :, i]).exp() - [..., None]).transpose(-1, -2) @ v_new - if not output_final_state: - S = None - - # unpad - o = rearrange(o, 'b h n c d -> b h (n c) d') - o = o[:, :, :T] - o = o.transpose(1, 2) - return o, S diff --git a/build/torch-xpu/ops/gated_delta_rule/wy_fast.py b/build/torch-xpu/ops/gated_delta_rule/wy_fast.py deleted file mode 100644 index d212ce610dc7b34d36584ff0f7323f5ccb4f062e..0000000000000000000000000000000000000000 --- a/build/torch-xpu/ops/gated_delta_rule/wy_fast.py +++ /dev/null @@ -1,332 +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 -import triton.language as tl - -from ...ops.utils import prepare_chunk_indices -from ...ops.utils.cache import fla_cache_autotune -from ...ops.utils.op import exp2 -from ...utils import IS_INTEL, IS_NVIDIA_BLACKWELL, autotune_cache_kwargs, check_shared_mem - -# Blackwell can select unstable Triton configs for prepare_wy_repr_bwd_kernel -# during autotuning (see #913). Restrict it to the config that has been -# validated on B200 until the wider config space is re-validated. -PREPARE_WY_REPR_BWD_NUM_WARPS = [2] if IS_NVIDIA_BLACKWELL else [2, 4] -PREPARE_WY_REPR_BWD_NUM_STAGES = [4] if IS_NVIDIA_BLACKWELL else [2, 3, 4] - -# Intel keeps scaling past the warp counts NVIDIA prefers: 16 warps is ~1.3x faster -# than 8 for recompute_w_u. -RECOMPUTE_W_U_NUM_WARPS = [2, 4, 8, 16] if IS_INTEL else [2, 4, 8] - - -@triton.heuristics({ - 'USE_G': lambda args: args['g'] is not None, - 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None, -}) -@fla_cache_autotune( - configs=[ - triton.Config({}, num_warps=num_warps, num_stages=num_stages) - for num_warps in RECOMPUTE_W_U_NUM_WARPS - for num_stages in [2, 3, 4] - ], - key=['H', 'HV', 'K', 'V', 'BT', 'BK', 'BV', 'IS_VARLEN'], - **autotune_cache_kwargs, -) -@triton.jit(do_not_specialize=['T']) -def recompute_w_u_fwd_kernel( - k, - v, - beta, - w, - u, - A, - g, - cu_seqlens, - chunk_indices, - T, - H: tl.constexpr, - HV: tl.constexpr, - K: tl.constexpr, - V: tl.constexpr, - BT: tl.constexpr, - BK: tl.constexpr, - BV: tl.constexpr, - USE_G: tl.constexpr, - IS_VARLEN: tl.constexpr, -): - i_t, i_bh = tl.program_id(0), tl.program_id(1).to(tl.int64) - i_b, i_h = i_bh // HV, i_bh % HV - 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.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) - T = eos - bos - else: - bos, eos = i_b * T, i_b * T + T - p_b = tl.make_block_ptr(beta + bos*HV + i_h, (T,), (HV,), (i_t * BT,), (BT,), (0,)) - b_b = tl.load(p_b, boundary_check=(0,)) - - p_A = tl.make_block_ptr(A + (bos*HV + i_h) * BT, (T, BT), (HV*BT, 1), (i_t * BT, 0), (BT, BT), (1, 0)) - b_A = tl.load(p_A, boundary_check=(0, 1)) - - for i_v in range(tl.cdiv(V, BV)): - p_v = tl.make_block_ptr(v + (bos*HV + i_h) * V, (T, V), (HV*V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) - p_u = tl.make_block_ptr(u + (bos*HV + i_h) * V, (T, V), (HV*V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) - b_v = tl.load(p_v, boundary_check=(0, 1)) - b_vb = (b_v * b_b[:, None]).to(b_v.dtype) - b_u = tl.dot(b_A, b_vb, allow_tf32=False) - tl.store(p_u, b_u.to(p_u.dtype.element_ty), boundary_check=(0, 1)) - - if USE_G: - p_g = tl.make_block_ptr(g + (bos*HV + i_h), (T,), (HV,), (i_t * BT,), (BT,), (0,)) - b_g = exp2(tl.load(p_g, boundary_check=(0,))) - - for i_k in range(tl.cdiv(K, BK)): - p_k = tl.make_block_ptr(k + (bos*H + i_h // (HV // H)) * K, (T, K), (H*K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) - p_w = tl.make_block_ptr(w + (bos*HV + i_h) * K, (T, K), (HV*K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) - b_k = tl.load(p_k, boundary_check=(0, 1)) - b_kb = b_k * b_b[:, None] - if USE_G: - b_kb *= b_g[:, None] - b_w = tl.dot(b_A, b_kb.to(b_k.dtype)) - tl.store(p_w, b_w.to(p_w.dtype.element_ty), boundary_check=(0, 1)) - - -@triton.heuristics({ - 'USE_G': lambda args: args['g'] is not None, - 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None, -}) -@fla_cache_autotune( - configs=[ - triton.Config({}, num_warps=num_warps, num_stages=num_stages) - for num_warps in PREPARE_WY_REPR_BWD_NUM_WARPS - for num_stages in PREPARE_WY_REPR_BWD_NUM_STAGES - ], - key=['H', 'HV', 'K', 'V', 'BT', 'BK', 'BV', 'IS_VARLEN'], - **autotune_cache_kwargs, -) -@triton.jit(do_not_specialize=['T']) -def prepare_wy_repr_bwd_kernel( - k, - v, - beta, - g, - A, - dw, - du, - dk, - dv, - db, - dg, - cu_seqlens, - chunk_indices, - T, - H: tl.constexpr, - HV: tl.constexpr, - K: tl.constexpr, - V: tl.constexpr, - BT: tl.constexpr, - BK: tl.constexpr, - BV: tl.constexpr, - USE_G: tl.constexpr, - IS_VARLEN: tl.constexpr, -): - i_t, i_bh = tl.program_id(0), tl.program_id(1).to(tl.int64) - i_b, i_h = i_bh // HV, i_bh % HV - 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.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) - T = eos - bos - else: - bos, eos = i_b * T, i_b * T + T - - p_b = tl.make_block_ptr(beta + (bos*HV + i_h), (T,), (HV,), (i_t * BT,), (BT,), (0,)) - p_db = tl.make_block_ptr(db + (bos*HV + i_h), (T,), (HV,), (i_t * BT,), (BT,), (0,)) - p_A = tl.make_block_ptr(A + (bos*HV + i_h) * BT, (BT, T), (1, HV*BT), (0, i_t * BT), (BT, BT), (0, 1)) - - b_b = tl.load(p_b, boundary_check=(0,)) - b_db = tl.zeros([BT], dtype=tl.float32) - b_A = tl.load(p_A, boundary_check=(0, 1)) - b_dA = tl.zeros([BT, BT], dtype=tl.float32) - - if USE_G: - p_g = tl.make_block_ptr(g + (bos*HV + i_h), (T,), (HV,), (i_t * BT,), (BT,), (0,)) - b_g = tl.load(p_g, boundary_check=(0,)) - b_g_exp = exp2(b_g) - b_dg = tl.zeros([BT], dtype=tl.float32) - - for i_k in range(tl.cdiv(K, BK)): - p_k = tl.make_block_ptr(k + (bos*H + i_h // (HV // H)) * K, (T, K), (H*K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) - p_dk = tl.make_block_ptr(dk + (bos*HV + i_h) * K, (T, K), (HV*K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) - p_dw = tl.make_block_ptr(dw + (bos*HV + i_h) * K, (T, K), (HV*K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) - # [BT, BK] - b_k = tl.load(p_k, boundary_check=(0, 1)) - if USE_G: - b_kbg = b_k * (b_b * b_g_exp)[:, None] - else: - b_kbg = b_k * b_b[:, None] - b_dw = tl.load(p_dw, boundary_check=(0, 1)) - - b_dA += tl.dot(b_dw, tl.trans(b_kbg).to(b_dw.dtype)) - b_dkbg = tl.dot(b_A, b_dw) - if USE_G: - b_dk = b_dkbg * (b_g_exp * b_b)[:, None] - b_db += tl.sum(b_dkbg * b_k * b_g_exp[:, None], 1) - b_dg += tl.sum(b_dkbg * b_kbg, 1) - else: - b_dk = b_dkbg * b_b[:, None] - b_db += tl.sum(b_dkbg * b_k, 1) - tl.store(p_dk, b_dk.to(p_dk.dtype.element_ty), boundary_check=(0, 1)) - - for i_v in range(tl.cdiv(V, BV)): - p_v = tl.make_block_ptr(v + (bos*HV + i_h) * V, (T, V), (HV*V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) - p_dv = tl.make_block_ptr(dv + (bos*HV + i_h) * V, (T, V), (HV*V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) - p_du = tl.make_block_ptr(du + (bos*HV + i_h) * V, (T, V), (HV*V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) - b_v = tl.load(p_v, boundary_check=(0, 1)) - b_vb = (b_v * b_b[:, None]).to(b_v.dtype) - b_du = tl.load(p_du, boundary_check=(0, 1)) - b_dA += tl.dot(b_du, tl.trans(b_vb)) - b_dvb = tl.dot(b_A, b_du) - b_dv = b_dvb * b_b[:, None] - b_db += tl.sum(b_dvb * b_v, 1) - tl.store(p_dv, b_dv.to(p_dv.dtype.element_ty), boundary_check=(0, 1)) - - o_t = i_t * BT + tl.arange(0, BT) - m_t = o_t < T - m_A = (o_t[:, None] > o_t[None, :]) & (m_t[:, None] & m_t) - b_dA = tl.where(m_A, b_dA, 0) - b_dA = tl.dot(b_dA.to(b_A.dtype), b_A) - b_dA = tl.dot(b_A, b_dA.to(b_A.dtype)) - - if USE_G: - b_dA *= exp2(b_g[:, None] - b_g[None, :]) - - b_A = tl.zeros([BT, BT], dtype=tl.float32) - b_dA = tl.where(m_A, -b_dA, 0).to(k.dtype.element_ty) - - tl.debug_barrier() - for i_k in range(tl.cdiv(K, BK)): - p_k = tl.make_block_ptr(k + (bos*H + i_h // (HV // H)) * K, (T, K), (H*K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) - p_dk = tl.make_block_ptr(dk + (bos*HV + i_h) * K, (T, K), (HV*K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) - b_k = tl.load(p_k, boundary_check=(0, 1)) - b_kt = tl.trans(b_k) - b_kb = b_k * b_b[:, None] - - b_A += tl.dot(b_k, b_kt) - b_dkb = tl.dot(b_dA, b_k) - b_db += tl.sum(b_dkb * b_k, 1) - b_dk = b_dkb * b_b[:, None] + tl.trans(tl.dot(tl.trans(b_kb).to(b_dA.dtype), b_dA)) - b_dk += tl.load(p_dk, boundary_check=(0, 1)) - - tl.store(p_dk, b_dk.to(p_dk.dtype.element_ty), boundary_check=(0, 1)) - tl.store(p_db, b_db.to(p_db.dtype.element_ty), boundary_check=(0,)) - - b_A *= b_b[:, None] - if USE_G: - b_AdA = b_dA * b_A - p_dg = tl.make_block_ptr(dg + (bos*HV + i_h), (T,), (HV,), (i_t * BT,), (BT,), (0,)) - b_dg += tl.sum(b_AdA, axis=1) - tl.sum(b_AdA, axis=0) - tl.store(p_dg, b_dg.to(p_dg.dtype.element_ty), boundary_check=(0,)) - - -def recompute_w_u_fwd( - k: torch.Tensor, - v: torch.Tensor, - beta: torch.Tensor, - A: torch.Tensor, - g: torch.Tensor | None = None, - cu_seqlens: torch.LongTensor | None = None, - chunk_indices: torch.LongTensor | None = None, -) -> tuple[torch.Tensor, torch.Tensor]: - B, T, H, K, V, HV = *k.shape, v.shape[-1], v.shape[2] - BT = A.shape[-1] - BK = 64 - BV = 64 - - if chunk_indices is None and cu_seqlens is not None: - chunk_indices = prepare_chunk_indices(cu_seqlens, BT) - NT = triton.cdiv(T, BT) if cu_seqlens is None else len(chunk_indices) - - w = k.new_empty(B, T, HV, K) - u = torch.empty_like(v) - recompute_w_u_fwd_kernel[(NT, B*HV)]( - k=k, - v=v, - beta=beta, - w=w, - u=u, - A=A, - g=g, - cu_seqlens=cu_seqlens, - chunk_indices=chunk_indices, - T=T, - H=H, - HV=HV, - K=K, - V=V, - BT=BT, - BK=BK, - BV=BV, - ) - return w, u - - -def prepare_wy_repr_bwd( - k: torch.Tensor, - v: torch.Tensor, - beta: torch.Tensor, - A: torch.Tensor, - dw: torch.Tensor, - du: torch.Tensor, - g: torch.Tensor = None, - cu_seqlens: torch.LongTensor | None = None, - chunk_indices: torch.LongTensor | None = None, -) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: - B, T, H, K, V, HV = *k.shape, v.shape[-1], v.shape[2] - BT = A.shape[-1] - if chunk_indices is None and cu_seqlens is not None: - chunk_indices = prepare_chunk_indices(cu_seqlens, BT) - NT = triton.cdiv(T, BT) if cu_seqlens is None else len(chunk_indices) - CONST_TILING = 64 if check_shared_mem() else 32 - BK = min(max(triton.next_power_of_2(K), 16), CONST_TILING) - BV = min(max(triton.next_power_of_2(V), 16), CONST_TILING) - - dk = k.new_empty(B, T, HV, K) - dv = torch.empty_like(v) - dg = torch.empty_like(g) if g is not None else None - db = torch.empty_like(beta) - prepare_wy_repr_bwd_kernel[(NT, B * HV)]( - k=k, - v=v, - beta=beta, - g=g, - A=A, - dw=dw, - du=du, - dk=dk, - dv=dv, - db=db, - dg=dg, - cu_seqlens=cu_seqlens, - chunk_indices=chunk_indices, - T=T, - H=H, - HV=HV, - K=K, - V=V, - BT=BT, - BK=BK, - BV=BV, - ) - if H != HV: - dk = dk.view(B, T, H, HV // H, K).sum(3) - return dk, dv, db, dg - - -fwd_recompute_w_u = recompute_w_u_fwd -bwd_prepare_wy_repr = prepare_wy_repr_bwd diff --git a/build/torch-xpu/ops/gated_oja_rule/__init__.py b/build/torch-xpu/ops/gated_oja_rule/__init__.py deleted file mode 100644 index 6edc02706d325434d477ac279ee411f1e2adee11..0000000000000000000000000000000000000000 --- a/build/torch-xpu/ops/gated_oja_rule/__init__.py +++ /dev/null @@ -1,14 +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 .chunk import chunk_gated_oja_rule -from .fused_recurrent import fused_recurrent_gated_oja_rule - -__all__ = [ - "chunk_gated_oja_rule", - "fused_recurrent_gated_oja_rule" -] diff --git a/build/torch-xpu/ops/gated_oja_rule/chunk.py b/build/torch-xpu/ops/gated_oja_rule/chunk.py deleted file mode 100644 index 15ef812680f87a15a81a2144b4d16f31ab840fc1..0000000000000000000000000000000000000000 --- a/build/torch-xpu/ops/gated_oja_rule/chunk.py +++ /dev/null @@ -1,361 +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 - -from ...modules.l2norm import l2norm_bwd, l2norm_fwd -from ...ops.gated_oja_rule.chunk_h import chunk_oja_bwd_dhu, chunk_oja_bwd_dvwg_h, chunk_oja_fwd_h -from ...ops.gated_oja_rule.chunk_kkt import chunk_scaled_dot_kkt_bwd_gk, chunk_scaled_dot_kkt_fwd -from ...ops.gated_oja_rule.chunk_o import ( - chunk_oja_bwd_dA, - chunk_oja_bwd_dqk, - chunk_oja_bwd_dv_o, - chunk_oja_fwd_o, -) -from ...ops.gated_oja_rule.wy_fast import prepare_wy_repr_bwd, recompute_w_u_fwd -from ...ops.utils import chunk_local_cumsum, solve_tril -from ...ops.utils.index import prepare_chunk_indices -from ...utils import autocast_custom_bwd, autocast_custom_fwd, input_guard - - -def chunk_oja_fwd( - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - gv: torch.Tensor, - beta: torch.Tensor, - scale: float, - initial_state: torch.Tensor, - output_final_state: bool, - g_cumsum: bool = True, - cu_seqlens: torch.LongTensor | None = None, - chunk_indices: torch.LongTensor | None = None, - chunk_size: int = 64, -): - if g_cumsum: - gv = chunk_local_cumsum(gv, chunk_size=chunk_size, cu_seqlens=cu_seqlens, chunk_indices=chunk_indices) - A = chunk_scaled_dot_kkt_fwd( - k=v, - gk=gv, - beta=beta, - cu_seqlens=cu_seqlens, - chunk_indices=chunk_indices, - chunk_size=chunk_size, - output_dtype=torch.float32 - ) - A = solve_tril( - A=A, - cu_seqlens=cu_seqlens, - chunk_indices=chunk_indices, - output_dtype=k.dtype - ) - w, u, vg = recompute_w_u_fwd( - k=k, - v=v, - beta=beta, - A=A, - gv=gv, - cu_seqlens=cu_seqlens, - chunk_indices=chunk_indices, - ) - h, k_new, final_state = chunk_oja_fwd_h( - v=vg, - w=w, - u=u, - gv=gv, - initial_state=initial_state, - output_final_state=output_final_state, - cu_seqlens=cu_seqlens, - chunk_indices=chunk_indices, - chunk_size=chunk_size, - ) - _, o = chunk_oja_fwd_o( - q=q, - k=k_new, - v=v, - h=h, - gv=gv, - scale=scale, - cu_seqlens=cu_seqlens, - chunk_indices=chunk_indices, - chunk_size=chunk_size, - ) - return gv, o, A, final_state - - -def chunk_oja_bwd( - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - gv: torch.Tensor, - beta: torch.Tensor, - A: torch.Tensor, - o: torch.Tensor, - scale: float, - initial_state: torch.Tensor, - do: torch.Tensor, - dht: torch.Tensor, - dgk: torch.Tensor | None = None, - cu_seqlens: torch.LongTensor | None = None, - chunk_indices: torch.LongTensor | None = None, - chunk_size: int = 64, -): - w, u, vg = recompute_w_u_fwd( - k=k, - v=v, - beta=beta, - A=A, - gv=gv, - cu_seqlens=cu_seqlens, - chunk_indices=chunk_indices, - ) - - h, k_new, _ = chunk_oja_fwd_h( - v=vg, - w=w, - u=u, - gv=gv, - initial_state=initial_state, - output_final_state=False, - cu_seqlens=cu_seqlens, - chunk_indices=chunk_indices, - chunk_size=chunk_size, - ) - - dAqk = chunk_oja_bwd_dA( - v=v, - gv=gv, - do=do, - scale=scale, - cu_seqlens=cu_seqlens, - chunk_indices=chunk_indices, - chunk_size=chunk_size, - ) - - Aqk, dq, dk_new = chunk_oja_bwd_dqk( - q=q, - k=k_new, - h=h, - gv=gv, - dA=dAqk, - do=do, - scale=scale, - cu_seqlens=cu_seqlens, - chunk_indices=chunk_indices, - chunk_size=chunk_size, - ) - - dh, dh0, dk_new = chunk_oja_bwd_dhu( - q=q, - vg=vg, - w=w, - gv=gv, - h0=initial_state, - dht=dht, - do=do, - dk=dk_new, - scale=scale, - cu_seqlens=cu_seqlens, - chunk_indices=chunk_indices, - states_in_fp32=False, - chunk_size=chunk_size, - ) - - dv, dw, dgv_last = chunk_oja_bwd_dvwg_h( - k=k_new, - v=v, - gv=gv, - h=h, - dh=dh, - dk=dk_new, - dgk=dgk, - cu_seqlens=cu_seqlens, - chunk_indices=chunk_indices, - chunk_size=chunk_size, - ) - - dv, dgv1 = chunk_oja_bwd_dv_o( - v=v, - gv=gv, - o=o, - A=Aqk, - dv=dv, - do=do, - cu_seqlens=cu_seqlens, - chunk_indices=chunk_indices, - chunk_size=chunk_size, - ) - - dk, dv1, db, dgv2, dAvv = prepare_wy_repr_bwd( - k=k, - v=v, - beta=beta, - gv=gv, - A=A, - dw=dw, - du=dk_new, - cu_seqlens=cu_seqlens, - chunk_indices=chunk_indices, - ) - - dv2, dgv3, db2 = chunk_scaled_dot_kkt_bwd_gk( - k=v, - g=gv, - beta=beta, - dA=dAvv, - cu_seqlens=cu_seqlens, - chunk_indices=chunk_indices, - chunk_size=chunk_size, - ) - - dv = dv.add_(dv1).add_(dv2) - db = db.add_(db2) - dgv = dgv_last.add_(chunk_local_cumsum( - dgv1.add_(dgv2).add_(dgv3), chunk_size=chunk_size, reverse=True, cu_seqlens=cu_seqlens, chunk_indices=chunk_indices - )) - return dq, dk, dv, db, dgv, dh0 - - -class ChunkOJAFunction(torch.autograd.Function): - - @staticmethod - @input_guard - @autocast_custom_fwd - def forward( - ctx, - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - gv: torch.Tensor, - beta: torch.Tensor, - scale: float, - initial_state: torch.Tensor, - output_final_state: bool, - cu_seqlens: torch.LongTensor | None = None, - cu_seqlens_cpu: torch.LongTensor | None = None, - use_q_l2norm: bool = False, - use_k_l2norm: bool = False, - chunk_size: int = 64, - ): - q_rstd, k_rstd = None, None - if use_q_l2norm: - q, q_rstd = l2norm_fwd(q) - if use_k_l2norm: - k, k_rstd = l2norm_fwd(k) - - chunk_indices = prepare_chunk_indices( - cu_seqlens, chunk_size, cu_seqlens_cpu=cu_seqlens_cpu) if cu_seqlens is not None else None - gv, o, A, final_state = chunk_oja_fwd( - q=q, - k=k, - v=v, - gv=gv, - beta=beta, - scale=scale, - initial_state=initial_state, - output_final_state=output_final_state, - cu_seqlens=cu_seqlens, - chunk_indices=chunk_indices, - chunk_size=chunk_size, - ) - ctx.save_for_backward(q, q_rstd, k, k_rstd, v, gv, beta, A, o, initial_state, cu_seqlens, chunk_indices) - ctx.scale = scale - ctx.chunk_size = chunk_size - ctx.use_q_l2norm = use_q_l2norm - ctx.use_k_l2norm = use_k_l2norm - return o.to(q.dtype), final_state - - @staticmethod - @input_guard - @autocast_custom_bwd - def backward( - ctx, - do: torch.Tensor, - dht: torch.Tensor - ): - q, q_rstd, k, k_rstd, v, gv, beta, A, o, initial_state, cu_seqlens, chunk_indices = ctx.saved_tensors - dq, dk, dv, db, dg, dh0 = chunk_oja_bwd( - q=q, - k=k, - v=v, - gv=gv, - beta=beta, - A=A, - o=o, - scale=ctx.scale, - initial_state=initial_state, - do=do, - dht=dht, - cu_seqlens=cu_seqlens, - chunk_indices=chunk_indices, - chunk_size=ctx.chunk_size, - ) - - if ctx.use_q_l2norm: - dq = l2norm_bwd(q, q_rstd, dq) - if ctx.use_k_l2norm: - dk = l2norm_bwd(k, k_rstd, dk) - return dq.to(q), dk.to(k), dv.to(v), dg.to(gv), db.to(beta), None, dh0, None, None, None, None, None, None - - -@torch.compiler.disable -def chunk_gated_oja_rule( - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - gv: torch.Tensor, - beta: torch.Tensor, - scale: float | None = None, - initial_state: torch.Tensor | None = None, - output_final_state: bool = False, - use_q_l2norm: bool = False, - use_k_l2norm: bool = False, - cu_seqlens: torch.LongTensor | None = None, - cu_seqlens_cpu: torch.LongTensor | None = None, - **kwargs, -): - if 'head_first' in kwargs: - raise DeprecationWarning( - "head_first has been removed. Inputs must be in `[B, T, H, ...]` format.", - ) - if 'use_qk_l2norm_in_kernel' in kwargs and (not use_q_l2norm and not use_k_l2norm): - use_q_l2norm = True - use_k_l2norm = True - - chunk_size = kwargs.pop('chunk_size', 64) - if chunk_size not in (16, 32, 64): - raise ValueError(f"`chunk_size` must be 16, 32, or 64, got {chunk_size}.") - - if cu_seqlens is not None: - if q.shape[0] != 1: - raise ValueError( - f"The batch size is expected to be 1 rather than {q.shape[0]} when using `cu_seqlens`." - f"Please flatten variable-length inputs before processing." - ) - if initial_state is not None and initial_state.shape[0] != len(cu_seqlens) - 1: - raise ValueError( - f"The number of initial states is expected to be equal to the number of input sequences, " - f"i.e., {len(cu_seqlens) - 1} rather than {initial_state.shape[0]}." - ) - if scale is None: - scale = k.shape[-1] ** -0.5 - o, final_state = ChunkOJAFunction.apply( - q, - k, - v, - gv, - beta, - scale, - initial_state, - output_final_state, - cu_seqlens, - cu_seqlens_cpu, - use_q_l2norm, - use_k_l2norm, - chunk_size, - ) - return o, final_state diff --git a/build/torch-xpu/ops/gated_oja_rule/chunk_h.py b/build/torch-xpu/ops/gated_oja_rule/chunk_h.py deleted file mode 100644 index 7f21802673379180becd59948162534db18ad108..0000000000000000000000000000000000000000 --- a/build/torch-xpu/ops/gated_oja_rule/chunk_h.py +++ /dev/null @@ -1,811 +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 -import triton.language as tl - -from ...ops.utils import prepare_chunk_indices, prepare_chunk_offsets -from ...ops.utils.op import exp -from ...utils import check_shared_mem, is_nvidia_hopper - -BKV_LIST = [64, 128] if check_shared_mem() else [32, 64] -NUM_WARPS = [2, 4] if is_nvidia_hopper else [2, 4, 8, 16] - - -@triton.heuristics({ - 'USE_GV': lambda args: args['gv'] is not None, - 'USE_INITIAL_STATE': lambda args: args['h0'] is not None, - 'STORE_FINAL_STATE': lambda args: args['ht'] is not None, - 'SAVE_NEW_KEY': lambda args: args['k_new'] is not None, - 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None, -}) -@triton.autotune( - configs=[ - triton.Config({'BK': BK}, num_warps=num_warps, num_stages=num_stages) - for num_warps in [2, 4] - for num_stages in [2, 3, 4] - for BK in [32, 64] - ], - key=['H', 'K', 'V', 'BT'], -) -@triton.jit(do_not_specialize=['T']) -def chunk_oja_fwd_kernel_h_blockdim64( - v, - u, - w, - k_new, - gv, - h, - h0, - ht, - cu_seqlens, - chunk_offsets, - T, - H: tl.constexpr, - K: tl.constexpr, - V: tl.constexpr, - BT: tl.constexpr, - BK: tl.constexpr, - USE_GV: tl.constexpr, - USE_INITIAL_STATE: tl.constexpr, - STORE_FINAL_STATE: tl.constexpr, - SAVE_NEW_KEY: tl.constexpr, - IS_VARLEN: tl.constexpr, -): - # (triton.cdiv(K, meta['BK']), N*H) - i_k, i_nh = tl.program_id(0), tl.program_id(1) - i_n, i_h = i_nh // H, i_nh % H - if IS_VARLEN: - bos, eos = tl.load(cu_seqlens + i_n).to(tl.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) - T = eos - bos - NT = tl.cdiv(T, BT) - boh = tl.load(chunk_offsets + i_n).to(tl.int32) - else: - bos, eos = i_n * T, i_n * T + T - NT = tl.cdiv(T, BT) - boh = i_n * NT - - # [BK, BV] - b_h1 = tl.zeros([BK, 64], dtype=tl.float32) - if V > 64: - b_h2 = tl.zeros([BK, 64], dtype=tl.float32) - if V > 128: - b_h3 = tl.zeros([BK, 64], dtype=tl.float32) - if V > 192: - b_h4 = tl.zeros([BK, 64], dtype=tl.float32) - - # calculate offset - h += ((boh * H + i_h) * K*V).to(tl.int64) - v += ((bos * H + i_h) * V).to(tl.int64) - u += ((bos * H + i_h) * K).to(tl.int64) - w += ((bos * H + i_h) * V).to(tl.int64) - if SAVE_NEW_KEY: - k_new += ((bos * H + i_h) * K).to(tl.int64) - stride_v = H*V - stride_h = H*K*V - stride_k = H*K - if USE_INITIAL_STATE: - h0 = h0 + i_nh * K*V - if STORE_FINAL_STATE: - ht = ht + i_nh * K*V - - # load initial state - if USE_INITIAL_STATE: - p_h0_1 = tl.make_block_ptr(h0, (K, V), (V, 1), (i_k * BK, 0), (BK, 64), (1, 0)) - b_h1 += tl.load(p_h0_1, boundary_check=(0, 1)).to(tl.float32) - if V > 64: - p_h0_2 = tl.make_block_ptr(h0, (K, V), (V, 1), (i_k * BK, 64), (BK, 64), (1, 0)) - b_h2 += tl.load(p_h0_2, boundary_check=(0, 1)).to(tl.float32) - if V > 128: - p_h0_3 = tl.make_block_ptr(h0, (K, V), (V, 1), (i_k * BK, 128), (BK, 64), (1, 0)) - b_h3 += tl.load(p_h0_3, boundary_check=(0, 1)).to(tl.float32) - if V > 192: - p_h0_4 = tl.make_block_ptr(h0, (K, V), (V, 1), (i_k * BK, 192), (BK, 64), (1, 0)) - b_h4 += tl.load(p_h0_4, boundary_check=(0, 1)).to(tl.float32) - - # main recurrence - for i_t in range(NT): - p_h1 = tl.make_block_ptr(h + i_t * stride_h, (K, V), (V, 1), (i_k * BK, 0), (BK, 64), (1, 0)) - tl.store(p_h1, b_h1.to(p_h1.dtype.element_ty), boundary_check=(0, 1)) - if V > 64: - p_h2 = tl.make_block_ptr(h + i_t * stride_h, (K, V), (V, 1), (i_k * BK, 64), (BK, 64), (1, 0)) - tl.store(p_h2, b_h2.to(p_h2.dtype.element_ty), boundary_check=(0, 1)) - if V > 128: - p_h3 = tl.make_block_ptr(h + i_t * stride_h, (K, V), (V, 1), (i_k * BK, 128), (BK, 64), (1, 0)) - tl.store(p_h3, b_h3.to(p_h3.dtype.element_ty), boundary_check=(0, 1)) - if V > 192: - p_h4 = tl.make_block_ptr(h + i_t * stride_h, (K, V), (V, 1), (i_k * BK, 192), (BK, 64), (1, 0)) - tl.store(p_h4, b_h4.to(p_h4.dtype.element_ty), boundary_check=(0, 1)) - - p_w = tl.make_block_ptr(w, (T, V), (stride_v, 1), (i_t * BT, 0), (BT, 64), (1, 0)) - b_w = tl.load(p_w, boundary_check=(0, 1)) - b_k = tl.dot(b_w, tl.trans(b_h1).to(b_w.dtype)) # BT BK - if V > 64: - p_w = tl.make_block_ptr(w, (T, V), (stride_v, 1), (i_t * BT, 64), (BT, 64), (1, 0)) - b_w = tl.load(p_w, boundary_check=(0, 1)) - b_k += tl.dot(b_w, tl.trans(b_h2).to(b_w.dtype)) - if V > 128: - p_w = tl.make_block_ptr(w, (T, V), (stride_v, 1), (i_t * BT, 128), (BT, 64), (1, 0)) - b_w = tl.load(p_w, boundary_check=(0, 1)) - b_k += tl.dot(b_w, tl.trans(b_h3).to(b_w.dtype)) - if V > 192: - p_w = tl.make_block_ptr(w, (T, V), (stride_v, 1), (i_t * BT, 192), (BT, 64), (1, 0)) - b_w = tl.load(p_w, boundary_check=(0, 1)) - b_k += tl.dot(b_w, tl.trans(b_h4).to(b_w.dtype)) - - p_u = tl.make_block_ptr(u, (T, K), (stride_k, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) - b_k = tl.load(p_u, boundary_check=(0, 1)) - b_k - - if SAVE_NEW_KEY: - p_k = tl.make_block_ptr(k_new, (T, K), (stride_k, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) - tl.store(p_k, b_k.to(p_k.dtype.element_ty), boundary_check=(0, 1)) - - last_idx = min((i_t + 1) * BT, T) - 1 - - if USE_GV: - o_v1 = tl.arange(0, 64) - b_gk_last1 = tl.load(gv + (bos + last_idx) * H*V + i_h * V + o_v1, mask=(o_v1 < V), other=0.) - b_h1 *= exp(b_gk_last1)[None, :] - if V > 64: - o_v2 = 64 + o_v1 - b_gk_last2 = tl.load(gv + (bos + last_idx) * H*V + i_h * V + o_v2, mask=(o_v2 < V), other=0.) - b_h2 *= exp(b_gk_last2)[None, :] - if V > 128: - o_v3 = 128 + o_v1 - b_gk_last3 = tl.load(gv + (bos + last_idx) * H*V + i_h * V + o_v3, mask=(o_v3 < V), other=0.) - b_h3 *= exp(b_gk_last3)[None, :] - if V > 192: - o_v4 = 192 + o_v1 - b_gk_last4 = tl.load(gv + (bos + last_idx) * H*V + i_h * V + o_v4, mask=(o_v4 < V), other=0.) - b_h4 *= exp(b_gk_last4)[None, :] - - b_k = b_k.to(v.dtype.element_ty) # BT BK - - p_v = tl.make_block_ptr(v, (T, V), (stride_v, 1), (i_t * BT, 0), (BT, 64), (1, 0)) - b_v = tl.load(p_v, boundary_check=(0, 1)) # BT BV - b_h1 += tl.dot(tl.trans(b_k), b_v) - if V > 64: - p_v = tl.make_block_ptr(v, (T, V), (stride_v, 1), (i_t * BT, 64), (BT, 64), (1, 0)) - b_v = tl.load(p_v, boundary_check=(0, 1)) - b_h2 += tl.dot(tl.trans(b_k), b_v) - if V > 128: - p_v = tl.make_block_ptr(v, (T, V), (stride_v, 1), (i_t * BT, 128), (BT, 64), (1, 0)) - b_v = tl.load(p_v, boundary_check=(0, 1)) - b_h3 += tl.dot(tl.trans(b_k), b_v) - if V > 192: - p_v = tl.make_block_ptr(v, (T, V), (stride_v, 1), (i_t * BT, 192), (BT, 64), (1, 0)) - b_v = tl.load(p_v, boundary_check=(0, 1)) - b_h4 += tl.dot(tl.trans(b_k), b_v) - # epilogue - if STORE_FINAL_STATE: - p_ht = tl.make_block_ptr(ht, (K, V), (V, 1), (i_k * BK, 0), (BK, 64), (1, 0)) - tl.store(p_ht, b_h1.to(p_ht.dtype.element_ty), boundary_check=(0, 1)) - if V > 64: - p_ht = tl.make_block_ptr(ht, (K, V), (V, 1), (i_k * BK, 64), (BK, 64), (1, 0)) - tl.store(p_ht, b_h2.to(p_ht.dtype.element_ty), boundary_check=(0, 1)) - if V > 128: - p_ht = tl.make_block_ptr(ht, (K, V), (V, 1), (i_k * BK, 128), (BK, 64), (1, 0)) - tl.store(p_ht, b_h3.to(p_ht.dtype.element_ty), boundary_check=(0, 1)) - if V > 192: - p_ht = tl.make_block_ptr(ht, (K, V), (V, 1), (i_k * BK, 192), (BK, 64), (1, 0)) - tl.store(p_ht, b_h4.to(p_ht.dtype.element_ty), boundary_check=(0, 1)) - - -def chunk_oja_fwd_h( - v: torch.Tensor, - w: torch.Tensor, - u: torch.Tensor, - gv: torch.Tensor | None = None, - initial_state: torch.Tensor | None = None, - output_final_state: bool = False, - chunk_size: int = 64, # SY: remove this argument and force chunk size 64? - save_new_key: bool = True, - cu_seqlens: torch.LongTensor | None = None, - chunk_indices: torch.LongTensor | None = None, -) -> tuple[torch.Tensor, torch.Tensor | None, torch.Tensor | None]: - B, T, H, V, K = *v.shape, u.shape[-1] - BT = chunk_size - - if chunk_indices is None and cu_seqlens is not None: - chunk_indices = prepare_chunk_indices(cu_seqlens, chunk_size) - if cu_seqlens is None: - N, NT, chunk_offsets = B, triton.cdiv(T, BT), None - else: - N, NT, chunk_offsets = len(cu_seqlens) - 1, len(chunk_indices), prepare_chunk_offsets(cu_seqlens, BT) - assert V <= 256, "current kernel does not support head dimension larger than 256." - - h = v.new_empty(B, NT, H, K, V) - final_state = v.new_empty(N, H, K, V, dtype=torch.float32) if output_final_state else None - - k_new = torch.empty_like(u) if save_new_key else None - def grid(meta): return (triton.cdiv(K, meta['BK']), N*H) - chunk_oja_fwd_kernel_h_blockdim64[grid]( - v=v, - u=u, - w=w, - k_new=k_new, - gv=gv, - h=h, - h0=initial_state, - ht=final_state, - cu_seqlens=cu_seqlens, - chunk_offsets=chunk_offsets, - T=T, - H=H, - K=K, - V=V, - BT=BT - ) - return h, k_new, final_state - - -@triton.heuristics({ - 'USE_GV': lambda args: args['gv'] is not None, - 'USE_INITIAL_STATE': lambda args: args['dh0'] is not None, - 'USE_FINAL_STATE_GRADIENT': lambda args: args['dht'] is not None, - 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None, -}) -@triton.autotune( - configs=[ - triton.Config({'BK': BK}, num_warps=num_warps, num_stages=num_stages) - for num_warps in [2, 4] - for num_stages in [4, 3, 2] - for BK in [64, 32] - ], - key=['H', 'K', 'V', 'BT', 'BK', 'USE_GV'], -) -@triton.jit(do_not_specialize=['T']) -def chunk_oja_bwd_kernel_dhu_blockdim64( - q, - vg, - w, - gv, - dht, - dh0, - do, - dh, - dk, - dk2, - cu_seqlens, - chunk_offsets, - scale, - T, - H: tl.constexpr, - K: tl.constexpr, - V: tl.constexpr, - BT: tl.constexpr, - BK: tl.constexpr, - USE_GV: tl.constexpr, - USE_INITIAL_STATE: tl.constexpr, - USE_FINAL_STATE_GRADIENT: tl.constexpr, - IS_VARLEN: tl.constexpr -): - i_k, i_nh = tl.program_id(0), tl.program_id(1) - i_n, i_h = i_nh // H, i_nh % H - if IS_VARLEN: - bos, eos = tl.load(cu_seqlens + i_n).to(tl.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) - T = eos - bos - NT = tl.cdiv(T, BT) - boh = tl.load(chunk_offsets + i_n).to(tl.int32) - else: - bos, eos = i_n * T, i_n * T + T - NT = tl.cdiv(T, BT) - boh = i_n * NT - - # [BK, BV] - b_dh1 = tl.zeros([BK, 64], dtype=tl.float32) - if V > 64: - b_dh2 = tl.zeros([BK, 64], dtype=tl.float32) - if V > 128: - b_dh3 = tl.zeros([BK, 64], dtype=tl.float32) - if V > 192: - b_dh4 = tl.zeros([BK, 64], dtype=tl.float32) - - # calculate offset - q += ((bos * H + i_h) * K).to(tl.int64) - vg += ((bos * H + i_h) * V).to(tl.int64) - w += ((bos * H + i_h) * V).to(tl.int64) - do += ((bos * H + i_h) * V).to(tl.int64) - dk += ((bos * H + i_h) * K).to(tl.int64) - dk2 += ((bos * H + i_h) * K).to(tl.int64) - dh += ((boh * H + i_h) * K*V).to(tl.int64) - if USE_GV: - gv += ((bos * H + i_h) * V).to(tl.int64) - - stride_v = H*V - stride_h = H*K*V - stride_k = H*K - if USE_INITIAL_STATE: - dh0 += i_nh * K*V - if USE_FINAL_STATE_GRADIENT: - dht += i_nh * K*V - - if USE_FINAL_STATE_GRADIENT: - p_dht1 = tl.make_block_ptr(dht, (K, V), (V, 1), (i_k * BK, 0), (BK, 64), (1, 0)) # [BK, BV] - b_dh1 += tl.load(p_dht1, boundary_check=(0, 1)) - if V > 64: - p_dht2 = tl.make_block_ptr(dht, (K, V), (V, 1), (i_k * BK, 64), (BK, 64), (1, 0)) - b_dh2 += tl.load(p_dht2, boundary_check=(0, 1)) - if V > 128: - p_dht3 = tl.make_block_ptr(dht, (K, V), (V, 1), (i_k * BK, 128), (BK, 64), (1, 0)) - b_dh3 += tl.load(p_dht3, boundary_check=(0, 1)) - if V > 192: - p_dht4 = tl.make_block_ptr(dht, (K, V), (V, 1), (i_k * BK, 192), (BK, 64), (1, 0)) - b_dh4 += tl.load(p_dht4, boundary_check=(0, 1)) - - for i_t in range(NT - 1, -1, -1): - p_dh1 = tl.make_block_ptr(dh + i_t*stride_h, (K, V), (V, 1), (i_k * BK, 0), (BK, 64), (1, 0)) - tl.store(p_dh1, b_dh1.to(p_dh1.dtype.element_ty), boundary_check=(0, 1)) - if V > 64: - p_dh2 = tl.make_block_ptr(dh + i_t*stride_h, (K, V), (V, 1), (i_k * BK, 64), (BK, 64), (1, 0)) - tl.store(p_dh2, b_dh2.to(p_dh2.dtype.element_ty), boundary_check=(0, 1)) - if V > 128: - p_dh3 = tl.make_block_ptr(dh + i_t*stride_h, (K, V), (V, 1), (i_k * BK, 128), (BK, 64), (1, 0)) - tl.store(p_dh3, b_dh3.to(p_dh3.dtype.element_ty), boundary_check=(0, 1)) - if V > 192: - p_dh4 = tl.make_block_ptr(dh + i_t*stride_h, (K, V), (V, 1), (i_k * BK, 192), (BK, 64), (1, 0)) - tl.store(p_dh4, b_dh4.to(p_dh4.dtype.element_ty), boundary_check=(0, 1)) - - last_idx = min((i_t + 1) * BT, T) - 1 - - # Update dk_new, 按K切分 - p_dk = tl.make_block_ptr(dk, (T, K), (stride_k, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) # [BT, BK] - p_dk2 = tl.make_block_ptr(dk2, (T, K), (stride_k, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) # [BT, BK] - - if V > 0: - p_v = tl.make_block_ptr(vg, (T, V), (stride_v, 1), (i_t * BT, 0), (BT, 64), (1, 0)) - b_v = tl.load(p_v, boundary_check=(0, 1)) # [BT, BV] - b_dk = tl.dot(b_v, tl.trans(b_dh1).to(b_v.dtype)) # [BT, BV] @ [BV, BK] -> [BT, BK] - - if V > 64: - p_v = tl.make_block_ptr(vg, (T, V), (stride_v, 1), (i_t * BT, 64), (BT, 64), (1, 0)) - b_v = tl.load(p_v, boundary_check=(0, 1)) - b_dk += tl.dot(b_v, tl.trans(b_dh2).to(b_v.dtype)) - - if V > 128: - p_v = tl.make_block_ptr(vg, (T, V), (stride_v, 1), (i_t * BT, 128), (BT, 64), (1, 0)) - b_v = tl.load(p_v, boundary_check=(0, 1)) - b_dk += tl.dot(b_v, tl.trans(b_dh3).to(b_v.dtype)) - - if V > 192: - p_v = tl.make_block_ptr(vg, (T, V), (stride_v, 1), (i_t * BT, 192), (BT, 64), (1, 0)) - b_v = tl.load(p_v, boundary_check=(0, 1)) - b_dk += tl.dot(b_v, tl.trans(b_dh4).to(b_v.dtype)) - - b_dk += tl.load(p_dk, boundary_check=(0, 1)) - - tl.store(p_dk2, b_dk.to(p_dk.dtype.element_ty), boundary_check=(0, 1)) - - # Update dh, 按照K切分，收集所有V维度，q一次就好，wdo要收集所有 - - p_q = tl.make_block_ptr(q, (K, T), (1, stride_k), (i_k * BK, i_t * BT), (BK, BT), (0, 1)) # [BK, BT] - b_q = tl.load(p_q, boundary_check=(0, 1)) - - if V > 0: - p_do = tl.make_block_ptr(do, (T, V), (stride_v, 1), (i_t * BT, 0), (BT, 64), (1, 0)) # [BT, BV] - b_do = tl.load(p_do, boundary_check=(0, 1)) - p_w = tl.make_block_ptr(w, (T, V), (stride_v, 1), (i_t * BT, 0), (BT, 64), (1, 0)) # [BT, BV] - b_w = tl.load(p_w, boundary_check=(0, 1)) - p_gv = tl.make_block_ptr(gv, (T, V), (stride_v, 1), (i_t * BT, 0), (BT, 64), (1, 0)) # [BT, BV] - b_gv = tl.load(p_gv, boundary_check=(0, 1)) - if USE_GV: - o_v1 = tl.arange(0, 64) - b_gv_last1 = tl.load(gv + last_idx * H*V + o_v1, mask=(o_v1 < V), other=0.) - b_dh1 *= exp(b_gv_last1[None, :]) - b_do *= exp(b_gv) - b_dh1 += tl.dot(b_q.to(b_q.dtype), b_do.to(b_q.dtype)) * scale - \ - tl.dot(tl.trans(b_dk).to(b_w.dtype), b_w) # [BK, BT] @ [BT, BV] - [BK, BT] @ [BT, BV] - - if V > 64: - p_do = tl.make_block_ptr(do, (T, V), (stride_v, 1), (i_t * BT, 64), (BT, 64), (1, 0)) - b_do = tl.load(p_do, boundary_check=(0, 1)) - p_w = tl.make_block_ptr(w, (T, V), (stride_v, 1), (i_t * BT, 64), (BT, 64), (1, 0)) # [BT, BV] - b_w = tl.load(p_w, boundary_check=(0, 1)) - p_gv = tl.make_block_ptr(gv, (T, V), (stride_v, 1), (i_t * BT, 64), (BT, 64), (1, 0)) # [BT, BV] - b_gv = tl.load(p_gv, boundary_check=(0, 1)) - if USE_GV: - o_v2 = 64 + o_v1 - b_gv_last2 = tl.load(gv + last_idx * H*V + o_v2, mask=(o_v2 < V), other=0.) - b_dh2 *= exp(b_gv_last2[None, :]) - b_do *= exp(b_gv) - b_dh2 += tl.dot(b_q.to(b_q.dtype), b_do.to(b_q.dtype)) * scale - tl.dot(tl.trans(b_dk).to(b_w.dtype), b_w) - - if V > 128: - p_do = tl.make_block_ptr(do, (T, V), (stride_v, 1), (i_t * BT, 128), (BT, 64), (1, 0)) - b_do = tl.load(p_do, boundary_check=(0, 1)) - p_w = tl.make_block_ptr(w, (T, V), (stride_v, 1), (i_t * BT, 128), (BT, 64), (1, 0)) # [BT, BV] - b_w = tl.load(p_w, boundary_check=(0, 1)) - p_gv = tl.make_block_ptr(gv, (T, V), (stride_v, 1), (i_t * BT, 128), (BT, 64), (1, 0)) # [BT, BV] - b_gv = tl.load(p_gv, boundary_check=(0, 1)) - if USE_GV: - o_v3 = 128 + o_v1 - b_gv_last3 = tl.load(gv + last_idx * H*V + o_v3, mask=(o_v3 < V), other=0.) - b_dh3 *= exp(b_gv_last3[None, :]) - b_do *= exp(b_gv) - b_dh3 += tl.dot(b_q.to(b_q.dtype), b_do.to(b_q.dtype)) * scale - tl.dot(tl.trans(b_dk).to(b_w.dtype), b_w) - - if V > 192: - p_do = tl.make_block_ptr(do, (T, V), (stride_v, 1), (i_t * BT, 192), (BT, 64), (1, 0)) - b_do = tl.load(p_do, boundary_check=(0, 1)) - p_w = tl.make_block_ptr(w, (T, V), (stride_v, 1), (i_t * BT, 192), (BT, 64), (1, 0)) # [BT, BV] - b_w = tl.load(p_w, boundary_check=(0, 1)) - p_gv = tl.make_block_ptr(gv, (T, V), (stride_v, 1), (i_t * BT, 192), (BT, 64), (1, 0)) # [BT, BV] - b_gv = tl.load(p_gv, boundary_check=(0, 1)) - if USE_GV: - o_v4 = 192 + o_v1 - b_gv_last4 = tl.load(gv + last_idx * H*V + o_v4, mask=(o_v4 < V), other=0.) - b_dh4 *= exp(b_gv_last4[None, :]) - b_do *= exp(b_gv) - b_dh4 += tl.dot(b_q.to(b_q.dtype), b_do.to(b_q.dtype)) * scale - tl.dot(tl.trans(b_dk).to(b_w.dtype), b_w) - - if USE_INITIAL_STATE: - p_dh0 = tl.make_block_ptr(dh0, (K, V), (V, 1), (i_k * BK, 0), (BK, 64), (1, 0)) - tl.store(p_dh0, b_dh1.to(p_dh0.dtype.element_ty), boundary_check=(0, 1)) - if V > 64: - p_dh1 = tl.make_block_ptr(dh0, (K, V), (V, 1), (i_k * BK, 64), (BK, 64), (1, 0)) - tl.store(p_dh1, b_dh2.to(p_dh1.dtype.element_ty), boundary_check=(0, 1)) - if V > 128: - p_dh2 = tl.make_block_ptr(dh0, (K, V), (V, 1), (i_k * BK, 128), (BK, 64), (1, 0)) - tl.store(p_dh2, b_dh3.to(p_dh2.dtype.element_ty), boundary_check=(0, 1)) - if V > 192: - p_dh3 = tl.make_block_ptr(dh0, (K, V), (V, 1), (i_k * BK, 192), (BK, 64), (1, 0)) - tl.store(p_dh3, b_dh4.to(p_dh3.dtype.element_ty), boundary_check=(0, 1)) - - -def chunk_oja_bwd_dhu( - q: torch.Tensor, - vg: torch.Tensor, - w: torch.Tensor, - do: torch.Tensor, - dk: torch.Tensor, - gv: torch.Tensor | None = None, - h0: torch.Tensor | None = None, - dht: torch.Tensor | None = None, - scale: float | None = None, - cu_seqlens: torch.LongTensor | None = None, - chunk_indices: torch.LongTensor | None = None, - chunk_size: int = 64, # SY: remove this argument and force chunk size 64? - states_in_fp32: bool = False -) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: - B, T, H, K, V = *q.shape, do.shape[-1] - BT = chunk_size - assert K <= 256, "current kernel does not support head dimension being larger than 256." - - if chunk_indices is None and cu_seqlens is not None: - chunk_indices = prepare_chunk_indices(cu_seqlens, chunk_size) - if cu_seqlens is None: - N, NT, chunk_offsets = B, triton.cdiv(T, BT), None - else: - N, NT, chunk_offsets = len(cu_seqlens) - 1, len(chunk_indices), prepare_chunk_offsets(cu_seqlens, BT) - - dh = q.new_empty(B, NT, H, K, V, dtype=q.dtype if not states_in_fp32 else torch.float) - dh0 = torch.empty_like(h0, dtype=torch.float32) if h0 is not None else None - dk2 = torch.empty_like(dk) - - def grid(meta): return (triton.cdiv(K, meta['BK']), N*H) - chunk_oja_bwd_kernel_dhu_blockdim64[grid]( - q=q, - vg=vg, - w=w, - gv=gv, - dht=dht, - dh0=dh0, - do=do, - dh=dh, - dk=dk, - dk2=dk2, - cu_seqlens=cu_seqlens, - chunk_offsets=chunk_offsets, - scale=scale, - T=T, - H=H, - K=K, - V=V, - BT=BT, - ) - return dh, dh0, dk2 - - -@triton.heuristics({ - 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None -}) -@triton.autotune( - configs=[ - triton.Config({}, num_warps=num_warps, num_stages=num_stages) - for num_warps in [2, 4] - for num_stages in [2, 3, 4] - ], - key=['BT'] -) -@triton.jit(do_not_specialize=['T']) -def chunk_gsa_bwd_k_kernel_dqkvg( - q, - k, - v, - h, - g, - A, - do, - dh, - dq, - dk, - dv, - dg, - dgv, - dA, - cu_seqlens, - chunk_indices, - scale, - T, - B: tl.constexpr, - HQ: tl.constexpr, - H: tl.constexpr, - K: tl.constexpr, - V: tl.constexpr, - BT: tl.constexpr, - BK: tl.constexpr, - BV: tl.constexpr, - NG: tl.constexpr, - IS_VARLEN: tl.constexpr, -): - i_k, i_t, i_bh = tl.program_id(0), tl.program_id(1), tl.program_id(2) - i_b, i_hq = i_bh // HQ, i_bh % HQ - i_h = i_hq // NG - 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.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) - all = T - T = eos - bos - NT = tl.cdiv(T, BT) - else: - NT = tl.cdiv(T, BT) - i_tg = i_b * NT + i_t - bos, eos = i_b * T, i_b * T + T - all = B * T - - o_i = tl.arange(0, BT) - o_t = min(i_t * BT + BT, T) - m_s = o_i[:, None] >= o_i[None, :] - - p_q = tl.make_block_ptr(q + (bos*HQ+i_hq) * K, (T, K), (HQ*K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) - p_k = tl.make_block_ptr(k + (bos*H+i_h) * K, (T, K), (H*K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) - p_A = tl.make_block_ptr(A + ((i_k*all+bos)*HQ+i_hq)*BT, (T, BT), (HQ*BT, 1), (i_t * BT, 0), (BT, BT), (1, 0)) - - # [BT, BK] - b_q = tl.load(p_q, boundary_check=(0, 1)) - b_k = tl.load(p_k, boundary_check=(0, 1)) - # [BT, BT] - b_A = tl.dot((b_q * scale).to(b_q.dtype), tl.trans(b_k)) - b_A = tl.where(m_s, b_A, 0.) - tl.store(p_A, b_A.to(p_A.dtype.element_ty), boundary_check=(0, 1)) - - b_dq = tl.zeros([BT, BK], dtype=tl.float32) - b_dk = tl.zeros([BT, BK], dtype=tl.float32) - for i_v in range(tl.cdiv(V, BV)): - o_v = i_v * BV + tl.arange(0, BV) - p_v = tl.make_block_ptr(v + (bos*H+i_h)*V, (T, V), (H*V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) - p_g = tl.make_block_ptr(g + (bos*H+i_h)*V, (T, V), (H*V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) - p_gn = g + (bos + o_t - 1) * H*V + i_h * V + o_v - p_do = tl.make_block_ptr(do + (bos*HQ+i_hq)*V, (T, V), (HQ*V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) - p_dv = tl.make_block_ptr(dv + ((i_k*all+bos)*HQ+i_hq)*V, (T, V), (HQ*V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) - p_dg = tl.make_block_ptr(dg + (bos*HQ+i_hq)*V, (T, V), (HQ*V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) - p_dgv = tl.make_block_ptr(dgv+((i_k*all+bos)*HQ+i_hq)*V, (T, V), (HQ*V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) - p_h = tl.make_block_ptr(h + (i_tg * H + i_h) * K*V, (V, K), (1, V), (i_v * BV, i_k * BK), (BV, BK), (0, 1)) - p_dh = tl.make_block_ptr(dh + (i_tg * HQ + i_hq) * K*V, (K, V), (V, 1), (i_k * BK, i_v * BV), (BK, BV), (1, 0)) - m_v = o_v < V - - # [BV,] - b_gn = tl.load(p_gn, mask=m_v, other=0) - # [BT, BV] - b_v = tl.load(p_v, boundary_check=(0, 1)) - b_g = tl.load(p_g, boundary_check=(0, 1)) - b_gv = exp(b_gn[None, :] - b_g) - # [BV, BK] - b_h = tl.load(p_h, boundary_check=(0, 1)) - # [BT, BV] - b_do = tl.load(p_do, boundary_check=(0, 1)) - b_do = (b_do * exp(b_g) * scale).to(b_do.dtype) - # [BK, BV] - b_dh = tl.load(p_dh, boundary_check=(0, 1)) - # [BV] - b_dg = tl.sum(tl.trans(b_h) * b_dh, 0) * exp(b_gn) - - b_dh = b_dh.to(b_k.dtype) - # [BT, BK] - b_dq += tl.dot(b_do, b_h.to(b_k.dtype)) - b_dk += tl.dot((b_v * b_gv).to(b_v.dtype), tl.trans(b_dh)) - # [BT, BV] - b_dv = tl.dot(b_k, b_dh) * b_gv - # [BV] - b_dg += tl.sum(b_dv * b_v, 0) - - if i_k == 0: - b_dgv = tl.load(p_dg, boundary_check=(0, 1)) + b_dg[None, :] - else: - b_dgv = tl.zeros([BT, BV], dtype=tl.float32) + b_dg[None, :] - - tl.store(p_dgv, b_dgv.to(p_dgv.dtype.element_ty), boundary_check=(0, 1)) - tl.store(p_dv, b_dv.to(p_dv.dtype.element_ty), boundary_check=(0, 1)) - p_dA = tl.make_block_ptr(dA + (bos*HQ + i_hq) * BT, (T, BT), (HQ*BT, 1), (i_t * BT, 0), (BT, BT), (1, 0)) - p_dq = tl.make_block_ptr(dq + (bos*HQ + i_hq) * K, (T, K), (HQ*K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) - p_dk = tl.make_block_ptr(dk + (bos*HQ + i_hq) * K, (T, K), (HQ*K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) - # [BT, BT] - b_dA = tl.load(p_dA, boundary_check=(0, 1)) - # [BT, BK] - b_dq += tl.dot(b_dA, b_k) - b_dk += tl.dot(tl.trans(b_dA).to(b_k.dtype), b_q) - - tl.store(p_dq, b_dq.to(p_dq.dtype.element_ty), boundary_check=(0, 1)) - tl.store(p_dk, b_dk.to(p_dk.dtype.element_ty), boundary_check=(0, 1)) - - -@triton.heuristics({ - 'USE_GV': lambda args: args['gv'] is not None, - 'HAVE_GK': lambda args: args['dgk'] is not None, - 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None, -}) -@triton.autotune( - configs=[ - triton.Config({}, num_warps=num_warps, num_stages=num_stages) - for num_warps in NUM_WARPS - for num_stages in [2, 3, 4] - ], - key=['H', 'K', 'V', 'BT', 'BK', 'BV', 'USE_GV'], -) -@triton.jit(do_not_specialize=['T']) -def chunk_oja_bwd_kernel_dvwg_h( - k, - v, - gv, - h, - dh, - dk, - dw, - dv, - dgv_last, - dgk, - cu_seqlens, - chunk_indices, - T, - H: tl.constexpr, - K: tl.constexpr, - V: tl.constexpr, - BT: tl.constexpr, - BK: tl.constexpr, - BV: tl.constexpr, - USE_GV: tl.constexpr, - HAVE_GK: tl.constexpr, - IS_VARLEN: tl.constexpr, -): - i_v, i_t, i_bh = tl.program_id(0), tl.program_id(1), tl.program_id(2) - i_b, i_h = i_bh // H, i_bh % H - - 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.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) - T = eos - bos - NT = tl.cdiv(T, BT) - else: - NT = tl.cdiv(T, BT) - i_tg = i_b * NT + i_t - bos, eos = i_b * T, i_b * T + T - - # offset calculation - k += (bos * H + i_h) * K - v += (bos * H + i_h) * V - gv += (bos * H + i_h) * V - h += (i_tg * H + i_h).to(tl.int64) * K*V - dh += (i_tg * H + i_h).to(tl.int64) * K*V - dk += (bos * H + i_h) * K - dw += (bos * H + i_h) * V - dv += (bos * H + i_h) * V - dgv_last += (bos * H + i_h) * V - - b_dvg = tl.zeros([BT, BV], dtype=tl.float32) - b_dw = tl.zeros([BT, BV], dtype=tl.float32) - b_dgv_last = tl.zeros([BV,], dtype=tl.float32) - - if USE_GV: - o_v = i_v * BV + tl.arange(0, BV) - m_v = o_v < V - p_gn = gv + (min(T, i_t * BT + BT) - 1) * H*V + o_v - p_gv = tl.make_block_ptr(gv, (T, V), (H*V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) - b_gn = tl.load(p_gn, mask=m_v, other=0) - b_gv = tl.load(p_gv, boundary_check=(0, 1)) - - for i_k in range(tl.cdiv(K, BK)): - p_k = tl.make_block_ptr(k, (T, K), (H*K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) - p_dk = tl.make_block_ptr(dk, (T, K), (H*K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) - p_h = tl.make_block_ptr(h, (K, V), (V, 1), (i_k * BK, i_v * BV), (BK, BV), (1, 0)) - p_dh = tl.make_block_ptr(dh, (K, V), (V, 1), (i_k * BK, i_v * BV), (BK, BV), (1, 0)) - b_k = tl.load(p_k, boundary_check=(0, 1)) # BT BK - b_dk = tl.load(p_dk, boundary_check=(0, 1)) # BT BK - b_h = tl.load(p_h, boundary_check=(0, 1)) # BK BV - b_dh = tl.load(p_dh, boundary_check=(0, 1)) # BK BV - - b_dvg += tl.dot(b_k, b_dh.to(b_k.dtype)) # BT BK @ BK BV -> BT BV - b_dw += tl.dot(b_dk.to(b_k.dtype), b_h.to(b_k.dtype)) # BT BK @ BK BV -> BT BV - b_dgv_last += tl.sum((b_h * b_dh) * exp(b_gn), axis=0) - - if USE_GV: - b_dv = b_dvg * exp(b_gn[None, :] - b_gv) - - p_v = tl.make_block_ptr(v, (T, V), (H*V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) - p_dv = tl.make_block_ptr(dv, (T, V), (H*V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) - p_dw = tl.make_block_ptr(dw, (T, V), (H*V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) - p_dgv_last = tl.make_block_ptr(dgv_last, (T, V), (H*V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) - b_v = tl.load(p_v, boundary_check=(0, 1)) - - b_dgv_last += tl.sum(b_dv * b_v, axis=0) - - # 留给GSA2的接口 - if HAVE_GK: - dgk += (bos * H + i_h) * V - p_dgk = tl.make_block_ptr(dgk, (T, V), (H*V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) - b_dgk = tl.load(p_dgk, boundary_check=(0, 1)) - b_dgv_last = b_dgk + b_dgv_last[None, :] - else: - b_dgv_last = tl.zeros([BT, BV], dtype=tl.float32) + b_dgv_last[None, :] - - tl.store(p_dv, b_dv.to(p_dv.dtype.element_ty), boundary_check=(0, 1)) - tl.store(p_dw, -b_dw.to(p_dw.dtype.element_ty), boundary_check=(0, 1)) - tl.store(p_dgv_last, b_dgv_last.to(p_dgv_last.dtype.element_ty), boundary_check=(0, 1)) - - -def chunk_oja_bwd_dvwg_h( - k: torch.Tensor, - v: torch.Tensor, - h: torch.Tensor, - dh: torch.Tensor, - dk: torch.Tensor, - gv: torch.Tensor | None = None, - dgk: torch.Tensor | None = None, - cu_seqlens: torch.LongTensor | None = None, - chunk_indices: torch.LongTensor | None = None, - chunk_size: int = 64, -) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: - - B, T, H, K, V = *k.shape, v.shape[-1] - BT = min(chunk_size, max(16, triton.next_power_of_2(T))) - if chunk_indices is None and cu_seqlens is not None: - chunk_indices = prepare_chunk_indices(cu_seqlens, BT) - NT = triton.cdiv(T, BT) if cu_seqlens is None else len(chunk_indices) - - CONST_TILING = 64 if check_shared_mem() else 32 - BK = min(max(triton.next_power_of_2(K), 16), CONST_TILING) - BV = min(max(triton.next_power_of_2(V), 16), CONST_TILING) - NV = triton.cdiv(V, BV) - dv = torch.empty_like(v, dtype=torch.float) - dw = torch.empty_like(v) - dgv_last = torch.empty_like(gv) - - grid = (NV, NT, B * H) - chunk_oja_bwd_kernel_dvwg_h[grid]( - k=k, - v=v, - gv=gv, - h=h, - dh=dh, - dw=dw, - dk=dk, - dv=dv, - dgv_last=dgv_last, - dgk=dgk, - cu_seqlens=cu_seqlens, - chunk_indices=chunk_indices, - T=T, - H=H, - K=K, - V=V, - BT=BT, - BK=BK, - BV=BV, - ) - return dv, dw, dgv_last diff --git a/build/torch-xpu/ops/gated_oja_rule/chunk_kkt.py b/build/torch-xpu/ops/gated_oja_rule/chunk_kkt.py deleted file mode 100644 index f152bfdc15109b57a85117324372352bc4934742..0000000000000000000000000000000000000000 --- a/build/torch-xpu/ops/gated_oja_rule/chunk_kkt.py +++ /dev/null @@ -1,525 +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 -import triton.language as tl - -from ...ops.utils import prepare_chunk_indices -from ...ops.utils.op import exp - - -@triton.heuristics({ - 'USE_G': lambda args: args['g'] is not None, - 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None, -}) -@triton.autotune( - configs=[ - triton.Config({'BK': BK}, num_warps=num_warps, num_stages=num_stages) - for BK in [32, 64, 128] - for num_warps in [2, 4, 8] - for num_stages in [2, 3, 4] - ], - key=['H', 'K', 'BT', 'IS_VARLEN'], -) -@triton.jit(do_not_specialize=['T']) -def chunk_scaled_dot_kkt_fwd_kernel( - k, - g, - beta, - A, - cu_seqlens, - chunk_indices, - T, - H: tl.constexpr, - K: tl.constexpr, - BT: tl.constexpr, - BK: tl.constexpr, - IS_VARLEN: tl.constexpr, - USE_G: tl.constexpr, -): - i_t, i_bh = tl.program_id(0), tl.program_id(1) - i_b, i_h = i_bh // H, i_bh % H - 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.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) - T = eos - bos - else: - bos, eos = i_b * T, i_b * T + T - o_t = i_t * BT + tl.arange(0, BT) - m_t = o_t < T - - p_b = tl.make_block_ptr(beta + bos*H + i_h, (T,), (H,), (i_t * BT,), (BT,), (0,)) - b_b = tl.load(p_b, boundary_check=(0,)) - - b_A = tl.zeros([BT, BT], dtype=tl.float32) - for i_k in range(tl.cdiv(K, BK)): - p_k = tl.make_block_ptr(k + (bos*H + i_h) * K, (T, K), (H*K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) - b_k = tl.load(p_k, boundary_check=(0, 1)) - b_A += tl.dot(b_k, tl.trans(b_k)) - - if USE_G: - p_g = tl.make_block_ptr(g + bos*H + i_h, (T,), (H,), (i_t * BT,), (BT,), (0,)) - b_g = tl.load(p_g, boundary_check=(0,)) - b_g_diff = b_g[:, None] - b_g[None, :] - b_A *= exp(b_g_diff) - b_A *= b_b[:, None] - - m_A = (o_t[:, None] > o_t[None, :]) & (m_t[:, None] & m_t) - b_A = tl.where(m_A, b_A, 0) - p_A = tl.make_block_ptr(A + (bos*H + i_h) * BT, (T, BT), (BT*H, 1), (i_t * BT, 0), (BT, BT), (1, 0)) - tl.store(p_A, b_A.to(p_A.dtype.element_ty), boundary_check=(0, 1)) - - -@triton.heuristics({ - 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None -}) -@triton.autotune( - configs=[ - triton.Config({'BK': BK}, num_warps=num_warps, num_stages=num_stages) - for BK in [32, 64] - for num_warps in [1, 2, 4, 8] - for num_stages in [2, 3, 4] - ], - key=["BC"] -) -@triton.jit(do_not_specialize=['T']) -def chunk_scaled_dot_kkt_fwd_kernel_intra_sub_inter( - k, - g, - beta, - A, - cu_seqlens, - chunk_indices, - T, - H: tl.constexpr, - K: tl.constexpr, - BT: tl.constexpr, - BC: tl.constexpr, - BK: tl.constexpr, - NC: tl.constexpr, - IS_VARLEN: tl.constexpr, -): - i_t, i_c, i_bh = tl.program_id(0), tl.program_id(1), tl.program_id(2) - i_b, i_h = i_bh // H, i_bh % H - i_i, i_j = i_c // NC, i_c % NC - 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.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) - T = eos - bos - else: - bos, eos = i_b * T, i_b * T + T - - if i_t * BT + i_i * BC >= T: - return - if i_i <= i_j: - return - - k += (bos * H + i_h) * K - g += (bos * H + i_h) * K - A += (bos * H + i_h) * BT - - p_b = tl.make_block_ptr(beta + bos * H + i_h, (T,), (H,), (i_t * BT + i_i * BC,), (BC,), (0,)) - b_b = tl.load(p_b, boundary_check=(0,)) - - b_A = tl.zeros([BC, BC], dtype=tl.float32) - for i_k in range(tl.cdiv(K, BK)): - p_k = tl.make_block_ptr(k, (T, K), (H*K, 1), (i_t * BT + i_i * BC, i_k * BK), (BC, BK), (1, 0)) - p_g = tl.make_block_ptr(g, (T, K), (H*K, 1), (i_t * BT + i_i * BC, i_k * BK), (BC, BK), (1, 0)) - b_kt = tl.make_block_ptr(k, (K, T), (1, H*K), (i_k * BK, i_t * BT + i_j * BC), (BK, BC), (0, 1)) - p_gk = tl.make_block_ptr(g, (K, T), (1, H*K), (i_k * BK, i_t * BT + i_j * BC), (BK, BC), (0, 1)) - - o_k = i_k * BK + tl.arange(0, BK) - m_k = o_k < K - # [BK,] - b_gn = tl.load(g + (i_t * BT + i_i * BC) * H*K + o_k, mask=m_k, other=0) - # [BC, BK] - b_g = tl.load(p_g, boundary_check=(0, 1)) - b_k = tl.load(p_k, boundary_check=(0, 1)) * exp(b_g - b_gn[None, :]) - # [BK, BC] - b_gk = tl.load(p_gk, boundary_check=(0, 1)) - b_kt = tl.load(b_kt, boundary_check=(0, 1)) * exp(b_gn[:, None] - b_gk) - # [BC, BC] - b_A += tl.dot(b_k, b_kt) - b_A *= b_b[:, None] - - p_A = tl.make_block_ptr(A, (T, BT), (H*BT, 1), (i_t * BT + i_i * BC, i_j * BC), (BC, BC), (1, 0)) - tl.store(p_A, b_A.to(A.dtype.element_ty), boundary_check=(0, 1)) - - -@triton.heuristics({ - 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None -}) -@triton.autotune( - configs=[ - triton.Config({}, num_warps=1), - triton.Config({}, num_warps=2), - triton.Config({}, num_warps=4), - triton.Config({}, num_warps=8), - ], - key=["BK", "BT"] -) -@triton.jit(do_not_specialize=['T']) -def chunk_scaled_dot_kkt_fwd_kernel_intra_sub_intra( - k, - g, - beta, - A, - cu_seqlens, - chunk_indices, - T, - H: tl.constexpr, - K: tl.constexpr, - BT: tl.constexpr, - BC: tl.constexpr, - BK: tl.constexpr, - IS_VARLEN: tl.constexpr, -): - i_t, i_i, i_bh = tl.program_id(0), tl.program_id(1), tl.program_id(2) - i_b, i_h = i_bh // H, i_bh % H - 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.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) - T = eos - bos - else: - bos, eos = i_b * T, i_b * T + T - - if i_t * BT + i_i * BC >= T: - return - - o_i = tl.arange(0, BC) - o_k = tl.arange(0, BK) - m_k = o_k < K - m_A = (i_t * BT + i_i * BC + o_i) < T - o_A = (bos + i_t * BT + i_i * BC + o_i) * H*BT + i_h * BT + i_i * BC - - p_k = tl.make_block_ptr(k + (bos * H + i_h) * K, (T, K), (H*K, 1), (i_t * BT + i_i * BC, 0), (BC, BK), (1, 0)) - p_g = tl.make_block_ptr(g + (bos * H + i_h) * K, (T, K), (H*K, 1), (i_t * BT + i_i * BC, 0), (BC, BK), (1, 0)) - p_b = beta + (bos + i_t * BT + i_i * BC + o_i) * H + i_h - - b_k = tl.load(p_k, boundary_check=(0, 1)) * tl.load(p_b, mask=m_A, other=0)[:, None] - b_g = tl.load(p_g, boundary_check=(0, 1)) - - p_kt = k + (bos + i_t * BT + i_i * BC) * H*K + i_h * K + o_k - p_gk = g + (bos + i_t * BT + i_i * BC) * H*K + i_h * K + o_k - for j in range(0, min(BC, T - i_t * BT - i_i * BC)): - b_kt = tl.load(p_kt, mask=m_k, other=0).to(tl.float32) - b_gk = tl.load(p_gk, mask=m_k, other=0).to(tl.float32) - b_A = tl.sum(b_k * b_kt[None, :] * exp(b_g - b_gk[None, :]), 1) - b_A = tl.where(o_i > j, b_A, 0.) - - tl.store(A + o_A + j, b_A, mask=m_A) - p_kt += H*K - p_gk += H*K - - -@triton.heuristics({ - 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None -}) -@triton.autotune( - configs=[ - triton.Config({}, num_warps=num_warps) - for num_warps in [1, 2, 4, 8] - ], - key=['BK', 'NC', 'BT'], -) -@triton.jit(do_not_specialize=['B', 'T']) -def chunk_scaled_dot_kkt_bwd_kernel_gk( - k, - g, - beta, - dA, - dk, - dg, - db, - cu_seqlens, - chunk_indices, - B, - T, - H: tl.constexpr, - K: tl.constexpr, - BT: tl.constexpr, - BC: tl.constexpr, - BK: tl.constexpr, - NC: tl.constexpr, - IS_VARLEN: tl.constexpr, -): - i_k, i_c, i_bh = tl.program_id(0), tl.program_id(1), tl.program_id(2) - i_b, i_h = i_bh // H, i_bh % H - i_t, i_i = i_c // NC, i_c % NC - - all = B * T - 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.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) - else: - bos, eos = i_b * T, i_b * T + T - T = eos - bos - if i_t * BT + i_i * BC >= T: - return - - o_k = i_k * BK + tl.arange(0, BK) - m_k = o_k < K - - k += (bos * H + i_h) * K - g += (bos * H + i_h) * K - beta += bos * H + i_h - - dA += (bos * H + i_h) * BT - dk += (bos * H + i_h) * K - dg += (bos * H + i_h) * K - db += (i_k * all + bos) * H + i_h - - p_g = tl.make_block_ptr(g, (T, K), (H*K, 1), (i_t * BT + i_i * BC, i_k * BK), (BC, BK), (1, 0)) - p_b = tl.make_block_ptr(beta, (T,), (H,), (i_t * BT + i_i * BC,), (BC,), (0,)) - # [BC, BK] - b_g = tl.load(p_g, boundary_check=(0, 1)) - b_dk = tl.zeros([BC, BK], dtype=tl.float32) - # [BC] - b_b = tl.load(p_b, boundary_check=(0,)) - if i_i > 0: - p_gn = g + (i_t * BT + i_i * BC) * H*K + o_k - # [BK,] - b_gn = tl.load(p_gn, mask=m_k, other=0) - for i_j in range(0, i_i): - p_k = tl.make_block_ptr(k, (T, K), (H*K, 1), (i_t * BT + i_j * BC, i_k * BK), (BC, BK), (1, 0)) - p_gk = tl.make_block_ptr(g, (T, K), (H*K, 1), (i_t * BT + i_j * BC, i_k * BK), (BC, BK), (1, 0)) - p_dA = tl.make_block_ptr(dA, (T, BT), (H*BT, 1), (i_t * BT + i_i * BC, i_j * BC), (BC, BC), (1, 0)) - # [BC, BK] - b_k = tl.load(p_k, boundary_check=(0, 1)) - b_gk = tl.load(p_gk, boundary_check=(0, 1)) - b_kg = b_k * exp(b_gn[None, :] - b_gk) - # [BC, BC] - b_dA = tl.load(p_dA, boundary_check=(0, 1)) - # [BC, BK] - b_dkb = tl.dot(b_dA, b_kg) * exp(b_g - b_gn[None, :]) - b_dk += b_dkb - - o_i = tl.arange(0, BC) - m_dA = (i_t * BT + i_i * BC + o_i) < T - o_dA = (i_t * BT + i_i * BC + o_i) * H*BT + i_i * BC - p_kj = k + (i_t * BT + i_i * BC) * H*K + o_k - p_gkj = g + (i_t * BT + i_i * BC) * H*K + o_k - - p_k = tl.make_block_ptr(k, (T, K), (H*K, 1), (i_t * BT + i_i * BC, i_k * BK), (BC, BK), (1, 0)) - b_k = tl.load(p_k, boundary_check=(0, 1)) - for j in range(0, min(BC, T - i_t * BT - i_i * BC)): - # [BC] - b_dA = tl.load(dA + o_dA + j, mask=m_dA, other=0) - # [BK] - b_kj = tl.load(p_kj, mask=m_k, other=0).to(tl.float32) - b_gkj = tl.load(p_gkj, mask=m_k, other=0).to(tl.float32) - # [BC, BK] - m_i = o_i[:, None] >= j - # [BC, BK] - b_dkb = tl.where(m_i, b_dA[:, None] * b_kj[None, :] * exp(b_g - b_gkj[None, :]), 0.) - b_dk += b_dkb - - p_kj += H*K - p_gkj += H*K - b_db = tl.sum(b_dk * b_k, 1) - b_dk *= b_b[:, None] - p_db = tl.make_block_ptr(db, (T,), (H,), (i_t * BT + i_i * BC,), (BC,), (0,)) - tl.store(p_db, b_db.to(p_db.dtype.element_ty), boundary_check=(0,)) - - tl.debug_barrier() - # [BC, BK] - b_dkt = tl.zeros([BC, BK], dtype=tl.float32) - - NC = min(NC, tl.cdiv(T - i_t * BT, BC)) - if i_i < NC - 1: - p_gn = g + (min(i_t * BT + i_i * BC + BC, T) - 1) * H*K + o_k - # [BK,] - b_gn = tl.load(p_gn, mask=m_k, other=0) - for i_j in range(i_i + 1, NC): - p_k = tl.make_block_ptr(k, (T, K), (H*K, 1), (i_t * BT + i_j * BC, i_k * BK), (BC, BK), (1, 0)) - p_gk = tl.make_block_ptr(g, (T, K), (H*K, 1), (i_t * BT + i_j * BC, i_k*BK), (BC, BK), (1, 0)) - p_dA = tl.make_block_ptr(dA, (BT, T), (1, H*BT), (i_i * BC, i_t * BT + i_j * BC), (BC, BC), (0, 1)) - p_b = tl.make_block_ptr(beta, (T,), (H,), (i_t * BT + i_j * BC,), (BC,), (0,)) - - o_j = i_t * BT + i_j * BC + o_i - m_j = o_j < T - # [BC] - b_b = tl.load(p_b, boundary_check=(0,)) - # [BC, BK] - b_kb = tl.load(p_k, boundary_check=(0, 1)).to(tl.float32) * b_b[:, None] - b_gk = tl.load(p_gk, boundary_check=(0, 1)) - b_kbg = b_kb * tl.where(m_j[:, None], exp(b_gk - b_gn[None, :]), 0) - # [BC, BC] - b_dA = tl.load(p_dA, boundary_check=(0, 1)) - # [BC, BK] - # (SY 09/17) important to not use bf16 here to have a good precision. - b_dkt += tl.dot(b_dA, b_kbg) - b_dkt *= exp(b_gn[None, :] - b_g) - o_dA = (i_t * BT + i_i * BC) * H*BT + i_i * BC + o_i - p_kj = k + (i_t * BT + i_i * BC) * H*K + o_k - p_gkj = g + (i_t * BT + i_i * BC) * H*K + o_k - p_bj = beta + (i_t * BT + i_i * BC) * H - - for j in range(0, min(BC, T - i_t * BT - i_i * BC)): - # [BC,] - b_dA = tl.load(dA + o_dA + j * H*BT) - # [BK,] - b_kbj = tl.load(p_kj, mask=m_k, other=0).to(tl.float32) * tl.load(p_bj) - b_gkj = tl.load(p_gkj, mask=m_k, other=0).to(tl.float32) - b_kbgj = b_kbj[None, :] * exp(b_gkj[None, :] - b_g) - # [BC, BK] - m_i = o_i[:, None] <= j - b_dkt += tl.where(m_i, b_dA[:, None] * b_kbgj, 0.) - - p_kj += H*K - p_gkj += H*K - p_bj += H - b_dg = (b_dk - b_dkt) * b_k - b_dk += b_dkt - - p_dk = tl.make_block_ptr(dk, (T, K), (H*K, 1), (i_t * BT + i_i * BC, i_k * BK), (BC, BK), (1, 0)) - p_dg = tl.make_block_ptr(dg, (T, K), (H*K, 1), (i_t * BT + i_i * BC, i_k * BK), (BC, BK), (1, 0)) - tl.store(p_dk, b_dk.to(p_dk.dtype.element_ty), boundary_check=(0, 1)) - tl.store(p_dg, b_dg.to(p_dg.dtype.element_ty), boundary_check=(0, 1)) - - -def chunk_scaled_dot_kkt_fwd( - k: torch.Tensor, - g: torch.Tensor | None = None, - gk: torch.Tensor | None = None, - beta: torch.Tensor | None = None, - cu_seqlens: torch.LongTensor | None = None, - chunk_indices: torch.LongTensor | None = None, - chunk_size: int = 64, - output_dtype: torch.dtype = torch.float32 -) -> torch.Tensor: - r""" - Compute beta * K * K^T. - - Args: - k (torch.Tensor): - The key tensor of shape `[B, T, H, K]`. - beta (torch.Tensor): - The beta tensor of shape `[B, T, H]`. - g (torch.Tensor): - The cumulative sum of the gate tensor of shape `[B, T, H]`. Default: `None`. - gk (torch.Tensor): - The cumulative sum of the gate tensor of shape `[B, T, H, K]` applied to the key tensor. Default: `None`. - cu_seqlens (torch.LongTensor): - The cumulative sequence lengths of the input tensor. - Default: None - chunk_indices (torch.LongTensor): - Pre-computed chunk indices. Default: None - chunk_size (int): - The chunk size. Default: 64. - output_dtype (torch.dtype): - The dtype of the output tensor. Default: `torch.float32` - - Returns: - beta * K * K^T of shape `[B, T, H, BT]` where `BT` is the chunk size. - """ - B, T, H, K = k.shape - BT = chunk_size - if chunk_indices is None and cu_seqlens is not None: - chunk_indices = prepare_chunk_indices(cu_seqlens, BT) - NT = triton.cdiv(T, BT) if cu_seqlens is None else len(chunk_indices) - if gk is None: - A = torch.empty(B, T, H, BT, device=k.device, dtype=output_dtype) - chunk_scaled_dot_kkt_fwd_kernel[(NT, B * H)]( - k=k, - g=g, - beta=beta, - A=A, - cu_seqlens=cu_seqlens, - chunk_indices=chunk_indices, - T=T, - H=H, - K=K, - BT=BT, - ) - return A - - BC = min(16, BT) - NC = triton.cdiv(BT, BC) - BK = max(triton.next_power_of_2(K), 16) - A = torch.zeros(B, T, H, BT, device=k.device, dtype=output_dtype) - grid = (NT, NC * NC, B * H) - chunk_scaled_dot_kkt_fwd_kernel_intra_sub_inter[grid]( - k=k, - g=gk, - beta=beta, - A=A, - cu_seqlens=cu_seqlens, - chunk_indices=chunk_indices, - T=T, - H=H, - K=K, - BT=BT, - BC=BC, - NC=NC, - ) - - grid = (NT, NC, B * H) - chunk_scaled_dot_kkt_fwd_kernel_intra_sub_intra[grid]( - k=k, - g=gk, - beta=beta, - A=A, - cu_seqlens=cu_seqlens, - chunk_indices=chunk_indices, - T=T, - H=H, - K=K, - BT=BT, - BC=BC, - BK=BK, - ) - return A - - -def chunk_scaled_dot_kkt_bwd_gk( - k: torch.Tensor, - g: torch.Tensor, - beta: torch.Tensor, - dA: torch.Tensor, - cu_seqlens: torch.LongTensor | None = None, - chunk_indices: torch.LongTensor | None = None, - chunk_size: int = 64 -): - B, T, H, K = k.shape - BT = chunk_size - BC = min(16, BT) - BK = min(64, triton.next_power_of_2(K)) - - if chunk_indices is None and cu_seqlens is not None: - chunk_indices = prepare_chunk_indices(cu_seqlens, chunk_size) - NT = triton.cdiv(T, BT) if cu_seqlens is None else len(chunk_indices) - NC = triton.cdiv(BT, BC) - NK = triton.cdiv(K, BK) - - dk = torch.empty_like(k, dtype=torch.float) - dg = torch.empty_like(g, dtype=torch.float) - db = beta.new_empty(NK, *beta.shape, dtype=torch.float) - grid = (NK, NT * NC, B * H) - chunk_scaled_dot_kkt_bwd_kernel_gk[grid]( - k=k, - g=g, - beta=beta, - dA=dA, - dk=dk, - dg=dg, - db=db, - cu_seqlens=cu_seqlens, - chunk_indices=chunk_indices, - B=B, - T=T, - H=H, - K=K, - BT=BT, - BC=BC, - BK=BK, - NC=NC, - ) - db = db.sum(0) - - return dk, dg, db diff --git a/build/torch-xpu/ops/gated_oja_rule/chunk_o.py b/build/torch-xpu/ops/gated_oja_rule/chunk_o.py deleted file mode 100644 index 43dc94d230b70475817c59d93f361e98ab5bac65..0000000000000000000000000000000000000000 --- a/build/torch-xpu/ops/gated_oja_rule/chunk_o.py +++ /dev/null @@ -1,680 +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 -import triton.language as tl - -from ...ops.utils import prepare_chunk_indices -from ...ops.utils.op import exp -from ...utils import check_shared_mem, is_nvidia_hopper - -BKV_LIST = [64, 128] if check_shared_mem() else [32, 64] -NUM_WARPS = [2, 4] if is_nvidia_hopper else [2, 4, 8] - - -@triton.heuristics({ - 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None -}) -@triton.autotune( - configs=[ - triton.Config({'BK': BK, 'BV': BV}, num_warps=num_warps, num_stages=num_stages) - for BK in [32, 64] - for BV in [32, 64] - for num_warps in [2, 4, 8] - for num_stages in [2, 3, 4] - ], - key=['BT'] -) -@triton.jit(do_not_specialize=['T']) -def chunk_oja_fwd_inter( - q, - k, - h, - gv, - o, - A, - cu_seqlens, - chunk_indices, - scale, - T, - HQ: tl.constexpr, - H: tl.constexpr, - K: tl.constexpr, - V: tl.constexpr, - BT: tl.constexpr, - BK: tl.constexpr, - BV: tl.constexpr, - NG: tl.constexpr, - IS_VARLEN: tl.constexpr, -): - i_v, i_t, i_bh = tl.program_id(0), tl.program_id(1), tl.program_id(2) - i_b, i_hq = i_bh // HQ, i_bh % HQ - i_h = i_hq // NG - 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.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) - T = eos - bos - NT = tl.cdiv(T, BT) - else: - NT = tl.cdiv(T, BT) - i_tg = i_b * NT + i_t - bos, eos = i_b * T, i_b * T + T - - o_i = tl.arange(0, BT) - m_s = o_i[:, None] >= o_i[None, :] - - b_o = tl.zeros([BT, BV], dtype=tl.float32) - b_A = tl.zeros([BT, BT], dtype=tl.float32) - for i_k in range(tl.cdiv(K, BK)): - p_q = tl.make_block_ptr(q + (bos * HQ + i_hq) * K, (T, K), (HQ*K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) - p_k = tl.make_block_ptr(k + (bos * H + i_h) * K, (K, T), (1, H*K), (i_k * BK, i_t * BT), (BK, BT), (0, 1)) - p_h = tl.make_block_ptr(h + (i_tg * H + i_h) * K*V, (K, V), (V, 1), (i_k * BK, i_v * BV), (BK, BV), (1, 0)) - - # [BT, BK] - b_q = tl.load(p_q, boundary_check=(0, 1)) - b_q = (b_q * scale).to(b_q.dtype) - # [BK, BT] - b_k = tl.load(p_k, boundary_check=(0, 1)) - # [BK, BV] - b_h = tl.load(p_h, boundary_check=(0, 1)) - # [BT, BV] - b_o += tl.dot(b_q, b_h) - # [BT, BT] - b_A += tl.dot(b_q, b_k) - p_g = tl.make_block_ptr(gv + (bos * H + i_h) * V, (T, V), (H*V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) - p_o = tl.make_block_ptr(o + (bos * HQ + i_hq) * V, (T, V), (HQ*V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) - p_A = tl.make_block_ptr(A + (bos * HQ + i_hq) * BT, (T, BT), (HQ*BT, 1), (i_t * BT, 0), (BT, BT), (1, 0)) - # [BT, BV] - b_g = tl.load(p_g, boundary_check=(0, 1)) - b_o = b_o * exp(b_g) - tl.store(p_o, b_o.to(p_o.dtype.element_ty), boundary_check=(0, 1)) - - # [BT, BT] - b_A = tl.where(m_s, b_A, 0.) - if i_v == 0: - tl.store(p_A, b_A.to(p_A.dtype.element_ty), boundary_check=(0, 1)) - - -@triton.heuristics({ - 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None -}) -@triton.jit(do_not_specialize=['T']) -def chunk_oja_fwd_intra( - v, - gv, - o, - A, - cu_seqlens, - chunk_indices, - T, - HQ: tl.constexpr, - H: tl.constexpr, - V: tl.constexpr, - BT: tl.constexpr, - BC: tl.constexpr, - BV: tl.constexpr, - NC: tl.constexpr, - NG: tl.constexpr, - IS_VARLEN: tl.constexpr, -): - i_v, i_c, i_bh = tl.program_id(0), tl.program_id(1), tl.program_id(2) - i_b, i_hq = i_bh // HQ, i_bh % HQ - i_h = i_hq // NG - i_t, i_i = i_c // NC, i_c % NC - 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.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) - T = eos - bos - else: - bos, eos = i_b * T, i_b * T + T - - o_v = i_v * BV + tl.arange(0, BV) - m_v = o_v < V - - if i_t * BT + i_i * BC >= T: - return - - p_g = tl.make_block_ptr(gv + (bos * H + i_h) * V, (T, V), (H*V, 1), (i_t * BT + i_i * BC, i_v * BV), (BC, BV), (1, 0)) - p_gn = gv + (bos + min(i_t * BT + i_i * BC, T)) * H*V + i_h * V + o_v - # [BV,] - b_gn = tl.load(p_gn, mask=m_v, other=0) - # [BC, BV] - b_o = tl.zeros([BC, BV], dtype=tl.float32) - for i_j in range(0, i_i): - p_A = tl.make_block_ptr(A + (bos*HQ+i_hq) * BT, (T, BT), (HQ*BT, 1), (i_t*BT+i_i*BC, i_j * BC), (BC, BC), (1, 0)) - p_v = tl.make_block_ptr(v + (bos*H+i_h) * V, (T, V), (H*V, 1), (i_t * BT + i_j * BC, i_v * BV), (BC, BV), (1, 0)) - p_gv = tl.make_block_ptr(gv + (bos*H+i_h) * V, (T, V), (H*V, 1), (i_t * BT + i_j * BC, i_v * BV), (BC, BV), (1, 0)) - # [BC, BV] - b_v = tl.load(p_v, boundary_check=(0, 1)) - b_gv = tl.load(p_gv, boundary_check=(0, 1)) - b_vg = (b_v * exp(b_gn[None, :] - b_gv)).to(b_v.dtype) - # [BC, BC] - b_A = tl.load(p_A, boundary_check=(0, 1)) - b_o += tl.dot(b_A, b_vg) - # [BC, BV] - b_g = tl.load(p_g, boundary_check=(0, 1)) - b_o *= exp(b_g - b_gn[None, :]) - - o_i = tl.arange(0, BC) - o_A = (bos + i_t * BT + i_i * BC + tl.arange(0, BC)) * HQ*BT + i_hq * BT + i_i * BC - m_A = (i_t * BT + i_i * BC + tl.arange(0, BC)) < T - for j in range(0, min(BC, T - i_t * BT - i_i * BC)): - p_v = v + (bos + i_t * BT + i_i * BC + j) * H*V + i_h * V + o_v - p_gv = gv + (bos + i_t * BT + i_i * BC + j) * H*V + i_h * V + o_v - # [BC,] - b_A = tl.load(A + o_A + j, mask=m_A, other=0) - # [BV,] - b_v = tl.load(p_v, mask=m_v, other=0).to(tl.float32) - b_gv = tl.load(p_gv, mask=m_v, other=0).to(tl.float32) - # [BC, BV] - b_vg = b_v[None, :] * exp(b_g - b_gv[None, :]) - # avoid 0 * inf = inf - b_o += tl.where(o_i[:, None] >= j, b_A[:, None] * b_vg, 0.) - p_o = tl.make_block_ptr(o + (bos*HQ + i_hq) * V, (T, V), (HQ*V, 1), (i_t * BT + i_i * BC, i_v * BV), (BC, BV), (1, 0)) - b_o += tl.load(p_o, boundary_check=(0, 1)) - tl.store(p_o, b_o.to(p_o.dtype.element_ty), boundary_check=(0, 1)) - - -def chunk_oja_fwd_o( - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - gv: torch.Tensor, - h: torch.Tensor, - scale: float = 1., - cu_seqlens: torch.LongTensor | None = None, - chunk_indices: torch.LongTensor | None = None, - chunk_size: int = 64 -) -> tuple[torch.Tensor, torch.Tensor]: - B, T, H, K, V = *k.shape, v.shape[-1] - BT = min(chunk_size, max(16, triton.next_power_of_2(T))) - BC = min(16, BT) - BV = min(64, triton.next_power_of_2(V)) - HQ = q.shape[2] - - if chunk_indices is None and cu_seqlens is not None: - chunk_indices = prepare_chunk_indices(cu_seqlens, BT) - NT = triton.cdiv(T, BT) if cu_seqlens is None else len(chunk_indices) - NC = triton.cdiv(BT, BC) - NG = HQ // H - - o = v.new_empty(B, T, HQ, V) - A = q.new_empty(B, T, HQ, BT) - def grid(meta): return (triton.cdiv(V, meta['BV']), NT, B * HQ) - chunk_oja_fwd_inter[grid]( - q, - k, - h, - gv, - o, - A, - cu_seqlens=cu_seqlens, - chunk_indices=chunk_indices, - scale=scale, - T=T, - HQ=HQ, - H=H, - K=K, - V=V, - BT=BT, - NG=NG, - ) - - def grid(meta): return (triton.cdiv(V, meta['BV']), NT * NC, B * HQ) - chunk_oja_fwd_intra[grid]( - v, - gv, - o, - A, - cu_seqlens=cu_seqlens, - chunk_indices=chunk_indices, - T=T, - HQ=HQ, - H=H, - V=V, - BT=BT, - BC=BC, - BV=BV, - NC=NC, - NG=NG, - num_warps=4, - num_stages=2 - ) - return A, o - - -@triton.heuristics({ - 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None -}) -@triton.autotune( - configs=[ - triton.Config({}, num_warps=num_warps) - for num_warps in [2, 4, 8] - ], - key=["BT"] -) -@triton.jit(do_not_specialize=['T']) -def chunk_oja_bwd_kernel_dA( - v, - gv, - do, - dA, - chunk_indices, - cu_seqlens, - scale, - T, - B: tl.constexpr, - H: tl.constexpr, - V: tl.constexpr, - BT: tl.constexpr, - BC: tl.constexpr, - BV: tl.constexpr, - NC: tl.constexpr, - IS_VARLEN: tl.constexpr, -): - i_v, i_c, i_bh = tl.program_id(0), tl.program_id(1), tl.program_id(2) - i_b, i_h = i_bh // H, i_bh % H - i_t, i_i, i_j = i_c // (NC * NC), (i_c % (NC * NC)) // NC, (i_c % (NC * NC)) % NC - 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.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) - all = T - T = eos - bos - else: - bos, eos = i_b * T, i_b * T + T - all = B * T - - o_v = i_v * BV + tl.arange(0, BV) - m_v = o_v < V - - if i_t * BT + i_i * BC >= T: - return - - # [BC, BC] - b_dA = tl.zeros([BC, BC], dtype=tl.float32) - if i_i > i_j: - p_v = tl.make_block_ptr(v + (bos*H+i_h) * V, (V, T), (1, H*V), (i_v * BV, i_t*BT + i_j*BC), (BV, BC), (0, 1)) - p_gv = tl.make_block_ptr(gv + (bos*H+i_h) * V, (V, T), (1, H*V), (i_v * BV, i_t*BT + i_j*BC), (BV, BC), (0, 1)) - p_gn = gv + (bos + i_t*BT + i_i*BC) * H*V + i_h * V + o_v - p_g = tl.make_block_ptr(gv + (bos*H+i_h) * V, (T, V), (H*V, 1), (i_t*BT + i_i*BC, i_v*BV), (BC, BV), (1, 0)) - p_do = tl.make_block_ptr(do + (bos*H+i_h) * V, (T, V), (H*V, 1), (i_t*BT + i_i*BC, i_v*BV), (BC, BV), (1, 0)) - # [BV,] - b_gn = tl.load(p_gn, mask=m_v, other=0.) - # [BC, BV] - b_g = tl.load(p_g, boundary_check=(0, 1)) - b_do = tl.load(p_do, boundary_check=(0, 1)) - b_do = (b_do * exp(b_g - b_gn[None, :]) * scale).to(b_do.dtype) - # [BV, BC] - b_v = tl.load(p_v, boundary_check=(0, 1)) - b_gv = tl.load(p_gv, boundary_check=(0, 1)) - b_vg = (b_v * exp(b_gn[:, None] - b_gv)).to(b_v.dtype) - # [BC, BC] - b_dA = tl.dot(b_do, b_vg) - elif i_i == i_j: - p_g = tl.make_block_ptr(gv + (bos*H + i_h) * V, (T, V), (H*V, 1), (i_t*BT + i_i*BC, i_v*BV), (BC, BV), (1, 0)) - p_do = tl.make_block_ptr(do + (bos*H + i_h) * V, (T, V), (H*V, 1), (i_t*BT + i_i*BC, i_v*BV), (BC, BV), (1, 0)) - p_v = v + (bos + i_t*BT + i_j*BC) * H*V + i_h * V + o_v - p_gv = gv + (bos + i_t*BT + i_j*BC) * H*V + i_h * V + o_v - # [BC, BV] - b_g = tl.load(p_g, boundary_check=(0, 1)) - b_do = tl.load(p_do, boundary_check=(0, 1)) * scale - m_v = o_v < V - - o_i = tl.arange(0, BC) - # [BC, BC] - m_dA = o_i[:, None] >= o_i[None, :] - for j in range(0, min(BC, T - i_t * BT - i_j * BC)): - # [BV,] - b_v = tl.load(p_v, mask=m_v, other=0).to(tl.float32) - b_gv = tl.load(p_gv, mask=m_v, other=0).to(tl.float32) - # [BC,] - b_dAj = tl.sum(b_do * b_v[None, :] * exp(b_g - b_gv[None, :]), 1) - b_dA = tl.where((o_i == j)[None, :], b_dAj[:, None], b_dA) - - p_v += H*V - p_gv += H*V - b_dA = tl.where(m_dA, b_dA, 0.) - - p_dA = tl.make_block_ptr(dA+((i_v*all+bos)*H+i_h)*BT, (T, BT), (H*BT, 1), (i_t*BT+i_i*BC, i_j*BC), (BC, BC), (1, 0)) - tl.store(p_dA, b_dA.to(dA.dtype.element_ty), boundary_check=(0, 1)) - - -def chunk_oja_bwd_dA( - v: torch.Tensor, - gv: torch.Tensor, - do: torch.Tensor, - scale: float = 1., - cu_seqlens: torch.LongTensor | None = None, - chunk_indices: torch.LongTensor | None = None, - chunk_size: int = 64 -): - B, T, H, V = v.shape - BT = min(chunk_size, max(16, triton.next_power_of_2(T))) - BC = min(16, BT) - BV = min(64, triton.next_power_of_2(V)) - - if chunk_indices is None and cu_seqlens is not None: - chunk_indices = prepare_chunk_indices(cu_seqlens, BT) - NT = triton.cdiv(T, BT) if cu_seqlens is None else len(chunk_indices) - NC = triton.cdiv(BT, BC) - NV = triton.cdiv(V, BV) - - dA = v.new_empty(NV, B, T, H, BT) - # 计算dA - grid = (NV, NT * NC * NC, B * H) - chunk_oja_bwd_kernel_dA[grid]( - v, - gv, - do, - dA, - cu_seqlens=cu_seqlens, - chunk_indices=chunk_indices, - scale=scale, - T=T, - B=B, - H=H, - V=V, - BT=BT, - BC=BC, - BV=BV, - NC=NC, - ) - dA = dA.sum(0, dtype=dA.dtype) - - return dA - - -@triton.heuristics({ - 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None -}) -@triton.autotune( - configs=[ - triton.Config({}, num_warps=num_warps, num_stages=num_stages) - for num_warps in [2, 4] - for num_stages in [2, 3, 4] - ], - key=['BT'] -) -@triton.jit(do_not_specialize=['T']) -def chunk_oja_bwd_kernel_dqk( - q, - k, - h, - gv, - A, - dq, - dk, - dA, - do, - scale, - cu_seqlens, - chunk_indices, - B, - T, - H: tl.constexpr, - K: tl.constexpr, - V: tl.constexpr, - BT: tl.constexpr, - BK: tl.constexpr, - BV: tl.constexpr, - IS_VARLEN: tl.constexpr, -): - i_k, i_t, i_bh = tl.program_id(0), tl.program_id(1), tl.program_id(2) - i_b, i_h = i_bh // H, i_bh % H - 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.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) - all = T - T = eos - bos - NT = tl.cdiv(T, BT) - else: - NT = tl.cdiv(T, BT) - i_tg = i_b * NT + i_t - bos, eos = i_b * T, i_b * T + T - all = B * T - - o_i = tl.arange(0, BT) - m_s = o_i[:, None] >= o_i[None, :] - - # [B, T, H, BT] - p_q = tl.make_block_ptr(q + (bos*H+i_h) * K, (T, K), (H*K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) - p_k = tl.make_block_ptr(k + (bos*H+i_h) * K, (T, K), (H*K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) - p_A = tl.make_block_ptr(A + ((i_k*all+bos)*H+i_h)*BT, (T, BT), (H*BT, 1), (i_t * BT, 0), (BT, BT), (1, 0)) - b_q = tl.load(p_q, boundary_check=(0, 1)) - b_k = tl.load(p_k, boundary_check=(0, 1)) - - b_A = tl.dot((b_q * scale).to(b_q.dtype), tl.trans(b_k)) - b_A = tl.where(m_s, b_A, 0.) - tl.store(p_A, b_A.to(p_A.dtype.element_ty), boundary_check=(0, 1)) - - b_dq = tl.zeros([BT, BK], dtype=tl.float32) - - # 先计算do对应的dq - for i_v in range(tl.cdiv(V, BV)): - p_h = tl.make_block_ptr(h + (i_tg * H + i_h) * K*V, (V, K), (1, V), (i_v * BV, i_k * BK), (BV, BK), (0, 1)) - p_do = tl.make_block_ptr(do + (bos*H+i_h)*V, (T, V), (H*V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) - p_gv = tl.make_block_ptr(gv + (bos*H+i_h)*V, (T, V), (H*V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) - b_h = tl.load(p_h, boundary_check=(0, 1)) - b_do = tl.load(p_do, boundary_check=(0, 1)) - b_gv = tl.load(p_gv, boundary_check=(0, 1)) - b_do = (b_do * exp(b_gv) * scale).to(b_do.dtype) - b_dq += tl.dot(b_do, b_h.to(b_do.dtype)) - - # 接着计算dA对应的dq, dk - p_dA = tl.make_block_ptr(dA + (bos*H + i_h) * BT, (T, BT), (H*BT, 1), (i_t * BT, 0), (BT, BT), (1, 0)) - p_dq = tl.make_block_ptr(dq + (bos*H + i_h) * K, (T, K), (H*K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) - p_dk = tl.make_block_ptr(dk + (bos*H + i_h) * K, (T, K), (H*K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) - # [BT, BT] - b_dA = tl.load(p_dA, boundary_check=(0, 1)) - # [BT, BK] - b_dq += tl.dot(b_dA.to(b_q.dtype), b_k) - b_dk = tl.dot(tl.trans(b_dA).to(b_q.dtype), b_q) - - tl.store(p_dq, b_dq.to(p_dq.dtype.element_ty), boundary_check=(0, 1)) - tl.store(p_dk, b_dk.to(p_dk.dtype.element_ty), boundary_check=(0, 1)) - - -def chunk_oja_bwd_dqk( - q: torch.Tensor, - k: torch.Tensor, - h: torch.Tensor, - gv: torch.Tensor, - dA: torch.Tensor, - do: torch.Tensor, - scale: float = 1., - cu_seqlens: torch.LongTensor | None = None, - chunk_indices: torch.LongTensor | None = None, - chunk_size: int = 64 -): - B, T, H, K, V = *q.shape, gv.shape[-1] - BT = min(chunk_size, max(16, triton.next_power_of_2(T))) - BK = min(64, triton.next_power_of_2(K)) - BV = min(64, triton.next_power_of_2(V)) - - if chunk_indices is None and cu_seqlens is not None: - chunk_indices = prepare_chunk_indices(cu_seqlens, BT) - NT = triton.cdiv(T, BT) if cu_seqlens is None else len(chunk_indices) - NK = triton.cdiv(K, BK) - - dq = torch.empty_like(q) - dk = torch.empty_like(k) - A = dA.new_empty(NK, B, T, H, BT) - # 计算dA - grid = (NK, NT, B * H) - chunk_oja_bwd_kernel_dqk[grid]( - q, - k, - h, - gv, - A, - dq, - dk, - dA, - do, - scale, - cu_seqlens=cu_seqlens, - chunk_indices=chunk_indices, - B=B, - T=T, - H=H, - K=K, - V=V, - BT=BT, - BK=BK, - BV=BV - ) - - A = A.sum(0, dtype=A.dtype) - - return A, dq, dk - - -@triton.heuristics({ - 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None -}) -@triton.jit(do_not_specialize=['T']) -def chunk_oja_bwd_kernel_dv_o( - v, - g, - o, - A, - do, - dv, - dv2, - dg, - cu_seqlens, - chunk_indices, - T, - H: tl.constexpr, - V: tl.constexpr, - BT: tl.constexpr, - BC: tl.constexpr, - BV: tl.constexpr, - NC: tl.constexpr, - IS_VARLEN: tl.constexpr, -): - i_v, i_c, i_bh = tl.program_id(0), tl.program_id(1), tl.program_id(2) - i_b, i_h = i_bh // H, i_bh % H - i_t, i_i = i_c // NC, i_c % NC - 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.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) - T = eos - bos - else: - bos, eos = i_b * T, i_b * T + T - - o_v = i_v * BV + tl.arange(0, BV) - m_v = o_v < V - - if i_t * BT + i_i * BC >= T: - return - - p_gv = tl.make_block_ptr(g + (bos*H+i_h)*V, (T, V), (H*V, 1), (i_t * BT + i_i * BC, i_v * BV), (BC, BV), (1, 0)) - p_gn = g + (bos + min(i_t * BT + i_i * BC + BC, T)-1)*H*V + i_h*V + o_v - # [BV,] - b_gn = tl.load(p_gn, mask=m_v, other=0) - # [BC, BV] - b_gv = tl.load(p_gv, boundary_check=(0, 1)) - b_dvg = tl.zeros([BC, BV], dtype=tl.float32) - for i_j in range(i_i + 1, NC): - p_g = tl.make_block_ptr(g + (bos*H+i_h) * V, (T, V), (H*V, 1), (i_t * BT + i_j * BC, i_v * BV), (BC, BV), (1, 0)) - p_A = tl.make_block_ptr(A + (bos*H+i_h) * BT, (BT, T), (1, H*BT), (i_i*BC, i_t*BT + i_j*BC), (BC, BC), (0, 1)) - p_do = tl.make_block_ptr(do + (bos*H+i_h) * V, (T, V), (H*V, 1), (i_t*BT + i_j*BC, i_v*BV), (BC, BV), (1, 0)) - # [BC, BV] - b_g = tl.load(p_g, boundary_check=(0, 1)) - b_do = tl.load(p_do, boundary_check=(0, 1)) * exp(b_g - b_gn[None, :]) - # [BC, BC] - b_A = tl.load(p_A, boundary_check=(0, 1)) - # [BC, BV] - b_dvg += tl.dot(b_A, b_do.to(b_A.dtype)) - b_dv = b_dvg * exp(b_gn[None, :] - b_gv) - - o_i = tl.arange(0, BC) - o_c = i_i * BC + tl.arange(0, BC) - - p_g = g + (bos + i_t * BT + i_i * BC) * H*V + i_h * V + o_v - p_A = A + (bos + i_t*BT + i_i*BC) * H*BT + i_h * BT + o_c - p_do = do + (bos + i_t*BT + i_i*BC) * H*V + i_h * V + o_v - for j in range(0, min(BC, T - i_t * BT - i_i * BC)): - # [BC,] - b_A = tl.load(p_A) - # [BV,] - b_g = tl.load(p_g, mask=m_v, other=0) - b_do = tl.load(p_do, mask=m_v, other=0) - # [BC, BV] - m_i = o_i[:, None] <= j - b_dv += tl.where(m_i, exp(b_g[None, :] - b_gv) * b_A[:, None] * b_do[None, :], 0.) - - p_g += H * V - p_A += H * BT - p_do += H * V - p_o = tl.make_block_ptr(o + (bos*H+i_h)*V, (T, V), (H*V, 1), (i_t*BT + i_i*BC, i_v*BV), (BC, BV), (1, 0)) - p_v = tl.make_block_ptr(v + (bos*H+i_h)*V, (T, V), (H*V, 1), (i_t*BT + i_i*BC, i_v*BV), (BC, BV), (1, 0)) - p_do = tl.make_block_ptr(do + (bos*H+i_h)*V, (T, V), (H*V, 1), (i_t*BT + i_i*BC, i_v*BV), (BC, BV), (1, 0)) - p_dv = tl.make_block_ptr(dv + (bos*H+i_h)*V, (T, V), (H*V, 1), (i_t*BT + i_i*BC, i_v*BV), (BC, BV), (1, 0)) - p_dv2 = tl.make_block_ptr(dv2 + (bos*H+i_h)*V, (T, V), (H*V, 1), (i_t*BT + i_i*BC, i_v*BV), (BC, BV), (1, 0)) - p_dg = tl.make_block_ptr(dg + (bos*H+i_h)*V, (T, V), (H*V, 1), (i_t*BT + i_i*BC, i_v*BV), (BC, BV), (1, 0)) - - b_o = tl.load(p_o, boundary_check=(0, 1)).to(tl.float32) - b_v = tl.load(p_v, boundary_check=(0, 1)).to(tl.float32) - b_do = tl.load(p_do, boundary_check=(0, 1)).to(tl.float32) - b_dv = b_dv + tl.load(p_dv, boundary_check=(0, 1)).to(tl.float32) - b_dg = b_o * b_do - b_v * b_dv - tl.store(p_dv2, b_dv.to(p_dv.dtype.element_ty), boundary_check=(0, 1)) - tl.store(p_dg, b_dg.to(p_dg.dtype.element_ty), boundary_check=(0, 1)) - - -def chunk_oja_bwd_dv_o( - v: torch.Tensor, - gv: torch.Tensor, - o: torch.Tensor, - A: torch.Tensor, - dv: torch.Tensor, - do: torch.Tensor, - cu_seqlens: torch.LongTensor | None = None, - chunk_indices: torch.LongTensor | None = None, - chunk_size: int = 64 -): - B, T, H, V = v.shape - BT = min(chunk_size, max(16, triton.next_power_of_2(T))) - BC = min(16, BT) - BV = min(64, triton.next_power_of_2(V)) - - if chunk_indices is None and cu_seqlens is not None: - chunk_indices = prepare_chunk_indices(cu_seqlens, BT) - NT = triton.cdiv(T, BT) if cu_seqlens is None else len(chunk_indices) - NC = triton.cdiv(BT, BC) - - dv2 = torch.empty_like(v, dtype=torch.float) - dgv = torch.empty_like(gv) - # 计算dA - def grid(meta): return (triton.cdiv(V, meta['BV']), NT * NC, B * H) - chunk_oja_bwd_kernel_dv_o[grid]( - v=v, - g=gv, - o=o, - A=A, - do=do, - dv=dv, - dv2=dv2, - dg=dgv, - cu_seqlens=cu_seqlens, - chunk_indices=chunk_indices, - T=T, - H=H, - V=V, - BT=BT, - BC=BC, - BV=BV, - NC=NC, - num_warps=4, - num_stages=2 - ) - return dv2, dgv diff --git a/build/torch-xpu/ops/gated_oja_rule/fused_recurrent.py b/build/torch-xpu/ops/gated_oja_rule/fused_recurrent.py deleted file mode 100644 index 4535fdc40137e59f53e36beea8620d03d4f7223c..0000000000000000000000000000000000000000 --- a/build/torch-xpu/ops/gated_oja_rule/fused_recurrent.py +++ /dev/null @@ -1,267 +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 -import triton.language as tl - -from ...ops.utils.op import exp -from ...utils import input_guard - - -@triton.heuristics({ - 'USE_GV': lambda args: args['gv'] is not None, - 'USE_INITIAL_STATE': lambda args: args['h0'] is not None, - 'STORE_FINAL_STATE': lambda args: args['ht'] is not None, - 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None -}) -@triton.jit(do_not_specialize=['T']) -def fused_recurrent_oja_fwd_kernel( - q, - k, - v, - gv, - beta, - o, - h0, - ht, - cu_seqlens, - scale, - T, - B: tl.constexpr, - H: tl.constexpr, - HV: tl.constexpr, - K: tl.constexpr, - V: tl.constexpr, - BK: tl.constexpr, - BV: tl.constexpr, - USE_GV: tl.constexpr, - USE_Q_L2NORM: tl.constexpr, - USE_K_L2NORM: tl.constexpr, - IS_BETA_HEADWISE: tl.constexpr, - USE_INITIAL_STATE: tl.constexpr, - STORE_FINAL_STATE: tl.constexpr, - IS_VARLEN: tl.constexpr, -): - i_v, i_nh = tl.program_id(0), tl.program_id(1) - i_n, i_hv = i_nh // HV, i_nh % HV - i_h = i_hv // (HV // H) - if IS_VARLEN: - bos, eos = tl.load(cu_seqlens + i_n).to(tl.int64), tl.load(cu_seqlens + i_n + 1).to(tl.int64) - T = eos - bos - else: - bos, eos = i_n * T, i_n * T + T - o_k = tl.arange(0, BK) - o_v = i_v * BV + tl.arange(0, BV) - - p_q = q + (bos * H + i_h) * K + o_k - p_k = k + (bos * H + i_h) * K + o_k - p_v = v + (bos * HV + i_hv) * V + o_v - if USE_GV: - p_gv = gv + (bos * HV + i_hv) * V + o_v - if IS_BETA_HEADWISE: - p_beta = beta + bos * HV + i_hv - else: - p_beta = beta + (bos * HV + i_hv) * V + o_v - - p_o = o + (bos * HV + i_hv) * V + o_v - - mask_k = o_k < K - mask_v = o_v < V - mask_h = mask_k[:, None] & mask_v[None, :] - - b_h = tl.zeros([BK, BV], dtype=tl.float32) - if USE_INITIAL_STATE: - p_h0 = h0 + i_nh * K*V + o_k[:, None] * V + o_v[None, :] - b_h += tl.load(p_h0, mask=mask_h, other=0).to(tl.float32) - - for _ in range(0, T): - b_q = tl.load(p_q, mask=mask_k, other=0).to(tl.float32) - b_k = tl.load(p_k, mask=mask_k, other=0).to(tl.float32) - b_v = tl.load(p_v, mask=mask_v, other=0).to(tl.float32) - if USE_Q_L2NORM: - b_q = b_q / tl.sqrt(tl.sum(b_q * b_q) + 1e-6) - if USE_K_L2NORM: - b_k = b_k / tl.sqrt(tl.sum(b_k * b_k) + 1e-6) - b_q = b_q * scale - if IS_BETA_HEADWISE: - b_beta = tl.load(p_beta).to(tl.float32) - else: - b_beta = tl.load(p_beta, mask=mask_v, other=0).to(tl.float32) - - # [BK, BV] - if USE_GV: - b_gv = tl.load(p_gv, mask=mask_v, other=0).to(tl.float32) - b_h *= exp(b_gv[None, :]) - - b_k = b_beta * (b_k - tl.sum(b_h * b_v[None, :], 1)) - b_h += b_k[:, None] * b_v - - # [BV] - b_o = tl.sum(b_h * b_q[:, None], 0) - tl.store(p_o, b_o.to(p_o.dtype.element_ty), mask=mask_v) - - p_q += H*K - p_k += H*K - p_v += HV*V - if USE_GV: - p_gv += HV*V - p_beta += HV * (1 if IS_BETA_HEADWISE else V) - p_o += HV*V - - if STORE_FINAL_STATE: - p_ht = ht + i_nh * K*V + o_k[:, None] * V + o_v[None, :] - tl.store(p_ht, b_h.to(p_ht.dtype.element_ty), mask=mask_h) - - -def fused_recurrent_oja_fwd( - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - gv: torch.Tensor | None = None, - beta: torch.Tensor | None = None, - scale: float = None, - initial_state: torch.Tensor = None, - output_final_state: bool = False, - use_q_l2norm: bool = False, - use_k_l2norm: bool = False, - cu_seqlens: torch.LongTensor | None = None, -) -> tuple[torch.Tensor, torch.Tensor]: - B, T, H, K, V = *k.shape, v.shape[-1] - assert V <= 128 - HV = v.shape[2] - N = B if cu_seqlens is None else len(cu_seqlens) - 1 - BK, BV = triton.next_power_of_2(K), min(triton.next_power_of_2(V), 256) - NV = triton.cdiv(V, BV) - num_stages = 3 - num_warps = 1 - - o = torch.empty_like(v) - final_state = q.new_empty(N, HV, K, V, dtype=torch.float32) if output_final_state else None - - grid = (NV, N * HV) - fused_recurrent_oja_fwd_kernel[grid]( - q=q, - k=k, - v=v, - gv=gv, - beta=beta, - o=o, - h0=initial_state, - ht=final_state, - cu_seqlens=cu_seqlens, - scale=scale, - T=T, - B=B, - H=H, - HV=HV, - K=K, - V=V, - BK=BK, - BV=BV, - IS_BETA_HEADWISE=beta.ndim != v.ndim, - USE_Q_L2NORM=use_q_l2norm, - USE_K_L2NORM=use_k_l2norm, - num_warps=num_warps, - num_stages=num_stages, - ) - return o, final_state - - -class FusedRecurrentFunction(torch.autograd.Function): - - @staticmethod - @input_guard - def forward( - ctx, - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - gv: torch.Tensor | None = None, - beta: torch.Tensor | None = None, - scale: float = None, - initial_state: torch.Tensor = None, - output_final_state: bool = False, - use_q_l2norm: bool = False, - use_k_l2norm: bool = False, - cu_seqlens: torch.LongTensor | None = None, - ): - o, final_state = fused_recurrent_oja_fwd( - q=q, - k=k, - v=v, - gv=gv, - beta=beta, - scale=scale, - initial_state=initial_state, - output_final_state=output_final_state, - use_q_l2norm=use_q_l2norm, - use_k_l2norm=use_k_l2norm, - cu_seqlens=cu_seqlens, - ) - - return o, final_state - - @staticmethod - @input_guard - def backward(ctx, do, dht): - raise NotImplementedError( - "Backward pass is not implemented yet and we do not have plans to implement it " - "because we haven't figured out how to compute dg without materializing the full " - "hidden states for all time steps." - ) - - -def fused_recurrent_gated_oja_rule( - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - gv: torch.Tensor | None = None, - beta: torch.Tensor | None = None, - scale: float = None, - initial_state: torch.Tensor = None, - output_final_state: bool = False, - use_q_l2norm: bool = False, - use_k_l2norm: bool = False, - cu_seqlens: torch.LongTensor | None = None, - **kwargs, -) -> tuple[torch.Tensor, torch.Tensor]: - - if 'use_qk_l2norm_in_kernel' in kwargs and (not use_q_l2norm and not use_k_l2norm): - use_q_l2norm = True - use_k_l2norm = True - - if cu_seqlens is not None: - if q.shape[0] != 1: - raise ValueError( - f"The batch size is expected to be 1 rather than {q.shape[0]} when using `cu_seqlens`." - f"Please flatten variable-length inputs before processing." - ) - if initial_state is not None and initial_state.shape[0] != len(cu_seqlens) - 1: - raise ValueError( - f"The number of initial states is expected to be equal to the number of input sequences, " - f"i.e., {len(cu_seqlens) - 1} rather than {initial_state.shape[0]}." - ) - if scale is None: - scale = k.shape[-1] ** -0.5 - if beta is None: - beta = torch.ones_like(q[..., 0]) - - o, final_state = FusedRecurrentFunction.apply( - q, - k, - v, - gv, - beta, - scale, - initial_state, - output_final_state, - use_q_l2norm, - use_k_l2norm, - cu_seqlens, - ) - return o, final_state diff --git a/build/torch-xpu/ops/gated_oja_rule/wy_fast.py b/build/torch-xpu/ops/gated_oja_rule/wy_fast.py deleted file mode 100644 index d93a873f411cc297baceeb6c0642b7285d4ef5de..0000000000000000000000000000000000000000 --- a/build/torch-xpu/ops/gated_oja_rule/wy_fast.py +++ /dev/null @@ -1,294 +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 -import triton.language as tl - -from ...ops.utils import prepare_chunk_indices -from ...ops.utils.op import exp -from ...utils import check_shared_mem - - -@triton.heuristics({ - 'STORE_VG': lambda args: args['vg'] is not None, - 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None -}) -@triton.autotune( - configs=[ - triton.Config({}, num_warps=num_warps, num_stages=num_stages) - for num_warps in [2, 4, 8] - for num_stages in [2, 3, 4] - ], - key=['H', 'K', 'V', 'BT', 'BK', 'BV', 'IS_VARLEN'], -) -@triton.jit(do_not_specialize=['T']) -def recompute_w_u_fwd_kernel( - k, - v, - vg, - beta, - w, - u, - A, - gv, - cu_seqlens, - chunk_indices, - T, - H: tl.constexpr, - K: tl.constexpr, - V: tl.constexpr, - BT: tl.constexpr, - BK: tl.constexpr, - BV: tl.constexpr, - STORE_VG: tl.constexpr, - IS_VARLEN: tl.constexpr -): - i_t, i_bh = tl.program_id(0), tl.program_id(1) - i_b, i_h = i_bh // H, i_bh % H - 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.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) - T = eos - bos - else: - bos, eos = i_b * T, i_b * T + T - p_b = tl.make_block_ptr(beta + bos*H + i_h, (T,), (H,), (i_t * BT,), (BT,), (0,)) - b_b = tl.load(p_b, boundary_check=(0,)) - - p_A = tl.make_block_ptr(A + (bos*H + i_h) * BT, (T, BT), (H*BT, 1), (i_t * BT, 0), (BT, BT), (1, 0)) - b_A = tl.load(p_A, boundary_check=(0, 1)) - - for i_v in range(tl.cdiv(V, BV)): - p_v = tl.make_block_ptr(v + (bos*H + i_h) * V, (T, V), (H*V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) - p_w = tl.make_block_ptr(w + (bos*H + i_h) * V, (T, V), (H*V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) - b_v = tl.load(p_v, boundary_check=(0, 1)) - b_vb = b_v * b_b[:, None] - - p_gv = tl.make_block_ptr(gv + (bos*H + i_h) * V, (T, V), (H*V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) - b_gv = tl.load(p_gv, boundary_check=(0, 1)) - b_vb *= exp(b_gv) - if STORE_VG: - last_idx = min(i_t * BT + BT, T) - 1 - - o_v = i_v * BV + tl.arange(0, BV) - m_v = o_v < V - b_gn = tl.load(gv + ((bos + last_idx) * H + i_h) * V + o_v, mask=m_v, other=0.) - b_vg = b_v * exp(b_gn - b_gv) - - p_vg = tl.make_block_ptr(vg + (bos * H + i_h) * V, (T, V), (H*V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) - tl.store(p_vg, b_vg.to(p_vg.dtype.element_ty), boundary_check=(0, 1)) - - b_w = tl.dot(b_A, b_vb.to(b_v.dtype)) - tl.store(p_w, b_w.to(p_w.dtype.element_ty), boundary_check=(0, 1)) - - for i_k in range(tl.cdiv(K, BK)): - p_k = tl.make_block_ptr(k + (bos*H + i_h) * K, (T, K), (H*K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) - p_u = tl.make_block_ptr(u + (bos*H + i_h) * K, (T, K), (H*K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) - b_k = tl.load(p_k, boundary_check=(0, 1)) - b_kb = (b_k * b_b[:, None]).to(b_k.dtype) - b_u = tl.dot(b_A, b_kb, allow_tf32=False) - tl.store(p_u, b_u.to(p_u.dtype.element_ty), boundary_check=(0, 1)) - - -@triton.heuristics({ - 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None -}) -@triton.autotune( - configs=[ - triton.Config({}, num_warps=num_warps, num_stages=num_stages) - for num_warps in [2, 4] - for num_stages in [2, 3, 4] - ], - key=['H', 'K', 'V', 'BT', 'BK', 'BV', 'IS_VARLEN'] -) -@triton.jit(do_not_specialize=['T']) -def prepare_wy_repr_bwd_kernel( - k, - v, - beta, - gv, - A, - dA, - dw, - du, - dk, - dv, - db, - dgv, - cu_seqlens, - chunk_indices, - T, - H: tl.constexpr, - K: tl.constexpr, - V: tl.constexpr, - BT: tl.constexpr, - BK: tl.constexpr, - BV: tl.constexpr, - IS_VARLEN: tl.constexpr -): - i_t, i_bh = tl.program_id(0), tl.program_id(1) - i_b, i_h = i_bh // H, i_bh % H - 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.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) - T = eos - bos - else: - bos, eos = i_b * T, i_b * T + T - - p_b = tl.make_block_ptr(beta + (bos*H + i_h), (T,), (H,), (i_t * BT,), (BT,), (0,)) - p_db = tl.make_block_ptr(db + (bos*H + i_h), (T,), (H,), (i_t * BT,), (BT,), (0,)) - p_A = tl.make_block_ptr(A + (bos*H + i_h) * BT, (BT, T), (1, H*BT), (0, i_t * BT), (BT, BT), (0, 1)) - - b_b = tl.load(p_b, boundary_check=(0,)) - b_db = tl.zeros([BT], dtype=tl.float32) - b_A = tl.load(p_A, boundary_check=(0, 1)) - b_dA = tl.zeros([BT, BT], dtype=tl.float32) - - for i_v in range(tl.cdiv(V, BV)): - p_v = tl.make_block_ptr(v + (bos*H + i_h) * V, (T, V), (H*V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) - p_dv = tl.make_block_ptr(dv + (bos*H + i_h) * V, (T, V), (H*V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) - p_dw = tl.make_block_ptr(dw + (bos*H + i_h) * V, (T, V), (H*V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) - b_v = tl.load(p_v, boundary_check=(0, 1)) - p_gv = tl.make_block_ptr(gv + (bos*H + i_h) * V, (T, V), (H*V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) - b_gv_exp = exp(tl.load(p_gv, boundary_check=(0, 1))) - b_vbg = b_v * b_b[:, None] * b_gv_exp - b_dw = tl.load(p_dw, boundary_check=(0, 1)) - - b_dA += tl.dot(b_dw, tl.trans(b_vbg).to(b_dw.dtype)) - b_dvbg = tl.dot(b_A, b_dw) - b_dv = b_dvbg * b_gv_exp * b_b[:, None] - b_db += tl.sum(b_dvbg * b_v * b_gv_exp, 1) - b_dgv = b_dvbg * b_vbg - - p_dgv = tl.make_block_ptr(dgv + (bos*H + i_h) * V, (T, V), (H*V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) - tl.store(p_dgv, b_dgv.to(p_dgv.dtype.element_ty), boundary_check=(0, 1)) - tl.store(p_dv, b_dv.to(p_dv.dtype.element_ty), boundary_check=(0, 1)) - - for i_k in range(tl.cdiv(K, BK)): - p_k = tl.make_block_ptr(k + (bos*H + i_h) * K, (T, K), (H*K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) - p_dk = tl.make_block_ptr(dk + (bos*H + i_h) * K, (T, K), (H*K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) - p_du = tl.make_block_ptr(du + (bos*H + i_h) * K, (T, K), (H*K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) - # [BT, BK] - b_k = tl.load(p_k, boundary_check=(0, 1)) - b_kb = (b_k * b_b[:, None]).to(b_k.dtype) # BT BK - b_du = tl.load(p_du, boundary_check=(0, 1)) # BT BK - b_dA += tl.dot(b_du, tl.trans(b_kb)) # BT BT - b_dkb = tl.dot(b_A, b_du) # BT BK - b_dk = b_dkb * b_b[:, None] - b_db += tl.sum(b_dkb * b_k, 1) - tl.store(p_dk, b_dk.to(p_dk.dtype.element_ty), boundary_check=(0, 1)) - - o_t = i_t * BT + tl.arange(0, BT) - m_t = o_t < T - m_A = (o_t[:, None] > o_t[None, :]) & (m_t[:, None] & m_t) - b_dA = tl.where(m_A, b_dA, 0) - b_dA = tl.dot(b_dA.to(b_A.dtype), b_A) - b_dA = tl.dot(b_A, b_dA.to(b_A.dtype)) - - b_dA = tl.where(m_A, -b_dA, 0) - - # if USE_GV: - p_dA = tl.make_block_ptr(dA + (bos*H + i_h) * BT, (T, BT), (H*BT, 1), (i_t * BT, 0), (BT, BT), (1, 0)) - tl.store(p_dA, b_dA.to(p_dA.dtype.element_ty), boundary_check=(0, 1)) - tl.store(p_db, b_db.to(p_db.dtype.element_ty), boundary_check=(0,)) - - -def recompute_w_u_fwd( - k: torch.Tensor, - v: torch.Tensor, - beta: torch.Tensor, - A: torch.Tensor, - gv: torch.Tensor | None = None, - cu_seqlens: torch.LongTensor | None = None, - chunk_indices: torch.LongTensor | None = None, -) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor | None]: - B, T, H, K, V = *k.shape, v.shape[-1] - BT = A.shape[-1] - BK = 64 - BV = 64 - - if chunk_indices is None and cu_seqlens is not None: - chunk_indices = prepare_chunk_indices(cu_seqlens, BT) - NT = triton.cdiv(T, BT) if cu_seqlens is None else len(chunk_indices) - - w = torch.empty_like(v) - u = torch.empty_like(k) - vg = torch.empty_like(v) if gv is not None else None - recompute_w_u_fwd_kernel[(NT, B*H)]( - k=k, - v=v, - vg=vg, - beta=beta, - w=w, - u=u, - A=A, - gv=gv, - cu_seqlens=cu_seqlens, - chunk_indices=chunk_indices, - T=T, - H=H, - K=K, - V=V, - BT=BT, - BK=BK, - BV=BV, - ) - return w, u, vg - - -def prepare_wy_repr_bwd( - k: torch.Tensor, - v: torch.Tensor, - beta: torch.Tensor, - A: torch.Tensor, - dw: torch.Tensor, - du: torch.Tensor, - gv: torch.Tensor = None, - cu_seqlens: torch.LongTensor | None = None, - chunk_indices: torch.LongTensor | None = None, -) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: - B, T, H, K, V = *k.shape, v.shape[-1] - BT = A.shape[-1] - if chunk_indices is None and cu_seqlens is not None: - chunk_indices = prepare_chunk_indices(cu_seqlens, BT) - NT = triton.cdiv(T, BT) if cu_seqlens is None else len(chunk_indices) - CONST_TILING = 64 if check_shared_mem() else 32 - BK = min(max(triton.next_power_of_2(K), 16), CONST_TILING) - BV = min(max(triton.next_power_of_2(V), 16), CONST_TILING) - - dk = torch.empty_like(k) - dv = torch.empty_like(v, dtype=torch.float) - - dgv = torch.empty_like(gv, dtype=torch.float) - dA = torch.empty_like(A, dtype=torch.float) - db = torch.empty_like(beta, dtype=torch.float) - - prepare_wy_repr_bwd_kernel[(NT, B * H)]( - k=k, - v=v, - beta=beta, - gv=gv, - A=A, - dA=dA, - dw=dw, - du=du, - dk=dk, - dv=dv, - db=db, - dgv=dgv, - cu_seqlens=cu_seqlens, - chunk_indices=chunk_indices, - T=T, - H=H, - K=K, - V=V, - BT=BT, - BK=BK, - BV=BV, - ) - - return dk, dv, db, dgv, dA diff --git a/build/torch-xpu/ops/gdn2/__init__.py b/build/torch-xpu/ops/gdn2/__init__.py deleted file mode 100644 index f7fd83239f65c0077e6a4d197ff6485a871fc009..0000000000000000000000000000000000000000 --- a/build/torch-xpu/ops/gdn2/__init__.py +++ /dev/null @@ -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 - -from .chunk import chunk_gdn2 -from .fused_recurrent import fused_recurrent_gdn2 -from .naive import naive_chunk_gdn2, naive_recurrent_gdn2 - -__all__ = [ - "chunk_gdn2", - "fused_recurrent_gdn2", - "naive_chunk_gdn2", - "naive_recurrent_gdn2", -] diff --git a/build/torch-xpu/ops/gdn2/chunk.py b/build/torch-xpu/ops/gdn2/chunk.py deleted file mode 100644 index 2d97a9333c9e74c66ed2273e8aec341ff7f752c3..0000000000000000000000000000000000000000 --- a/build/torch-xpu/ops/gdn2/chunk.py +++ /dev/null @@ -1,363 +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 - -r""" -Chunkwise training kernel for GDN-2 (Gated DeltaNet 2). - -GDN-2 extends the gated delta rule with two independent channel-wise gates: an -erase gate ``b in R^K`` on the key axis and a write gate ``w in R^V`` on the -value axis. The per-token recurrence on the matrix state ``S in R^{K x V}`` is - - S_t = (I - k_t (b_t * k_t)^T) Diag(exp(g_t)) S_{t-1} + k_t (w_t * v_t)^T - -where ``*`` is the elementwise product. Setting ``b_t = w_t = beta`` (scalar) -recovers KDA; further collapsing ``g_t`` to a scalar recovers Gated DeltaNet v1. - -This file is the public entry point and the autograd wrapper. Forward and -backward orchestration live in ``chunk_fwd.py`` and ``chunk_bwd.py``; the -GDN-2-specific Triton kernels live in ``chunk_intra.py``, -``chunk_intra_token_parallel.py``, ``wy_fast.py``, and ``chunk_bwd.py``. -""" - -from __future__ import annotations - -import warnings - -import torch - -from ...modules.l2norm import l2norm_bwd, l2norm_fwd -from ...ops.gdn2.chunk_bwd import chunk_gdn2_bwd -from ...ops.gdn2.chunk_fwd import chunk_gdn2_fwd -from ...ops.utils import prepare_chunk_indices -from ...utils import autocast_custom_bwd, autocast_custom_fwd, input_guard - - -class ChunkGDN2Function(torch.autograd.Function): - """Autograd-compatible wrapper around GDN-2 forward / backward.""" - - @staticmethod - @input_guard - @autocast_custom_fwd - def forward( - ctx, - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - g: torch.Tensor, - b: torch.Tensor, - w: torch.Tensor, - A_log: torch.Tensor | None, - dt_bias: torch.Tensor | None, - scale: float, - initial_state: torch.Tensor | None, - output_final_state: bool, - use_qk_l2norm_in_kernel: bool, - use_gate_in_kernel: bool, - cu_seqlens: torch.LongTensor | None, - cu_seqlens_cpu: torch.LongTensor | None, - safe_gate: bool, - lower_bound: float | None, - chunk_size: int, - disable_recompute: bool, - return_intermediate_states: bool, - state_v_first: bool, - ): - q_rstd, k_rstd = None, None - if use_qk_l2norm_in_kernel: - q, q_rstd = l2norm_fwd(q) - k, k_rstd = l2norm_fwd(k) - - chunk_indices = ( - prepare_chunk_indices(cu_seqlens, chunk_size) - if cu_seqlens is not None else None - ) - - g_input = g - - (o, final_state, g_cumsum, Aqk, Akk, - w_wy, u_wy, qg, kg, v_new, h, initial_state) = chunk_gdn2_fwd( - q=q, - k=k, - v=v, - g=g_input, - b=b, - w_gate=w, - scale=scale, - initial_state=initial_state, - output_final_state=output_final_state, - cu_seqlens=cu_seqlens, - cu_seqlens_cpu=cu_seqlens_cpu, - chunk_indices=chunk_indices, - chunk_size=chunk_size, - safe_gate=safe_gate, - lower_bound=lower_bound, - use_gate_in_kernel=use_gate_in_kernel, - A_log=A_log, - dt_bias=dt_bias, - disable_recompute=disable_recompute, - return_intermediate_states=return_intermediate_states, - state_v_first=state_v_first, - ) - - if return_intermediate_states: - assert torch.is_inference_mode_enabled(), ( - "return_intermediate_states is only allowed in inference mode" - ) - assert disable_recompute is False, ( - "return_intermediate_states must be used with disable_recompute=False" - ) - return o.type_as(q), final_state, h - - ctx.save_for_backward( - q, q_rstd, k, k_rstd, v, g_cumsum, g_input, b, w, A_log, dt_bias, - Aqk, Akk, w_wy, u_wy, qg, kg, v_new, h, - initial_state, cu_seqlens, chunk_indices, - ) - ctx.chunk_size = chunk_size - ctx.safe_gate = safe_gate - ctx.scale = scale - ctx.lower_bound = lower_bound - ctx.use_qk_l2norm_in_kernel = use_qk_l2norm_in_kernel - ctx.use_gate_in_kernel = use_gate_in_kernel - ctx.disable_recompute = disable_recompute - ctx.state_v_first = state_v_first - return o.type_as(q), final_state - - @staticmethod - @input_guard - @autocast_custom_bwd - def backward(ctx, do: torch.Tensor, dht: torch.Tensor): - (q, q_rstd, k, k_rstd, v, g_cumsum, g_input, b, w, A_log, dt_bias, - Aqk, Akk, w_wy, u_wy, qg, kg, v_new, h, - initial_state, cu_seqlens, chunk_indices) = ctx.saved_tensors - - dq, dk, dv, db, dw, dg, dh0, dA_log, dt_bias_grad = chunk_gdn2_bwd( - q=q, - k=k, - v=v, - b=b, - w_gate=w, - Aqk=Aqk, - Akk=Akk, - scale=ctx.scale, - initial_state=initial_state, - do=do, - dht=dht, - g=g_cumsum, - g_org=g_input if ctx.use_gate_in_kernel else None, - cu_seqlens=cu_seqlens, - chunk_indices=chunk_indices, - chunk_size=ctx.chunk_size, - safe_gate=ctx.safe_gate, - lower_bound=ctx.lower_bound, - use_gate_in_kernel=ctx.use_gate_in_kernel, - A_log=A_log, - dt_bias=dt_bias, - state_v_first=ctx.state_v_first, - w_wy=w_wy, u_wy=u_wy, qg=qg, kg=kg, v_new=v_new, h=h, - disable_recompute=ctx.disable_recompute, - ) - - if ctx.use_qk_l2norm_in_kernel: - dq = l2norm_bwd(q, q_rstd, dq) - dk = l2norm_bwd(k, k_rstd, dk) - - return ( - dq.to(q.dtype), # q - dk.to(k.dtype), # k - dv.to(v.dtype), # v - dg.to(g_input.dtype), # g - db.to(b.dtype), # b - dw.to(w.dtype), # w - dA_log, # A_log - dt_bias_grad, # dt_bias - None, # scale - dh0, # initial_state - None, # output_final_state - None, # use_qk_l2norm_in_kernel - None, # use_gate_in_kernel - None, # cu_seqlens - None, # cu_seqlens_cpu - None, # safe_gate - None, # lower_bound - None, # chunk_size - None, # disable_recompute - None, # return_intermediate_states - None, # state_v_first - ) - - -@torch.compiler.disable -def chunk_gdn2( - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - g: torch.Tensor, - b: torch.Tensor, - w: torch.Tensor, - scale: float | None = None, - initial_state: torch.Tensor | None = None, - output_final_state: bool = False, - use_qk_l2norm_in_kernel: bool = False, - use_gate_in_kernel: bool = False, - cu_seqlens: torch.LongTensor | None = None, - cu_seqlens_cpu: torch.LongTensor | None = None, - safe_gate: bool = False, - lower_bound: float | None = None, - disable_recompute: bool = False, - return_intermediate_states: bool = False, - state_v_first: bool = False, - **kwargs, -): - r""" - Chunkwise forward for Gated DeltaNet 2 (GDN-2). - - The matrix-state recurrence on ``S in R^{K x V}`` is - - S_t = (I - k_t (b_t * k_t)^T) Diag(exp(g_t)) S_{t-1} + k_t (w_t * v_t)^T - - where ``b in R^K`` and ``w in R^V`` are channel-wise erase and write gates. - Setting ``b = w = beta`` (scalar) recovers KDA. - - Args: - q: queries of shape ``[B, T, H, K]``. - k: keys of shape ``[B, T, H, K]``. - v: values of shape ``[B, T, H, V]``. - g: (forget) log-decay of shape ``[B, T, H, K]``. With - ``use_gate_in_kernel=True`` this is the raw pre-activation and the - kernel computes ``-exp(A_log) * softplus(g + dt_bias)`` (or the - bounded variant if ``lower_bound`` is set). - b: channel-wise ERASE gate of shape ``[B, T, H, K]``. Replaces KDA's - scalar beta. Typical range: ``[0, 2]``. - w: channel-wise WRITE gate of shape ``[B, T, H, V]``. New for GDN-2. - Typical range: ``[0, 1]``. - scale: attention scale. Defaults to ``1 / sqrt(K)``. - initial_state: optional ``[N, H, K, V]`` initial state in float32 (or - ``[N, H, V, K]`` if ``state_v_first=True``). - output_final_state: whether to output the final recurrent state. - use_qk_l2norm_in_kernel: L2-normalize q and k inside the kernel. - use_gate_in_kernel: compute the gate activation inside the kernel from - raw ``g`` and (required) ``A_log`` (plus optional ``dt_bias``, - ``lower_bound``). - cu_seqlens: ``[N+1]`` packed-sequence offsets. Requires batch size 1. - cu_seqlens_cpu: optional CPU mirror of ``cu_seqlens``, forwarded to the - state-recurrence kernel. - safe_gate: use the safe-gate intra kernel variant (M=16 TensorCore - path; requires gate values in ``[-5, 0)`` if combined with - ``use_gate_in_kernel=True``). - lower_bound: when ``safe_gate=True`` and ``use_gate_in_kernel=True``, - use the bounded gate activation - ``lower_bound * sigmoid(exp(A_log) * g)``. Must lie in ``[-5, 0)``. - disable_recompute: retain forward intermediates for a faster backward - at the cost of memory. Default: ``False``. - return_intermediate_states: when ``True``, also returns the per-chunk - pre-states ``h`` (shape ``[B, NT, H, K, V]``). Must be used inside - ``torch.inference_mode()``. - state_v_first: store the recurrent state in ``[V, K]`` layout instead - of the default ``[K, V]``. - - Returns: - - Normal mode: ``(o, final_state)``. - - Intermediate mode (``return_intermediate_states=True``): - ``(o, final_state, h)``. - - Examples:: - - >>> import torch - >>> import torch.nn.functional as F - >>> from ...ops.gdn2 import chunk_gdn2 - >>> B, T, H, K, V = 2, 1024, 4, 128, 128 - >>> q = torch.randn(B, T, H, K, dtype=torch.bfloat16, device='cuda') - >>> k = torch.randn(B, T, H, K, dtype=torch.bfloat16, device='cuda') - >>> v = torch.randn(B, T, H, V, dtype=torch.bfloat16, device='cuda') - >>> g = F.logsigmoid(torch.rand(B, T, H, K, device='cuda')).to(torch.bfloat16) - >>> b = torch.rand(B, T, H, K, dtype=torch.bfloat16, device='cuda') - >>> w = torch.rand(B, T, H, V, dtype=torch.bfloat16, device='cuda') - >>> o, ht = chunk_gdn2(q, k, v, g, b, w, output_final_state=True) - """ - if 'transpose_state_layout' in kwargs: - if state_v_first: - raise ValueError("Cannot pass both `state_v_first` and the deprecated `transpose_state_layout`.") - warnings.warn( - "`transpose_state_layout` is deprecated and renamed to `state_v_first`.", - DeprecationWarning, - stacklevel=2, - ) - state_v_first = kwargs.pop('transpose_state_layout') - - if cu_seqlens is not None: - if q.shape[0] != 1: - raise ValueError( - f"The batch size is expected to be 1 rather than {q.shape[0]} when using `cu_seqlens`. " - f"Please flatten variable-length inputs before processing." - ) - if initial_state is not None and initial_state.shape[0] != len(cu_seqlens) - 1: - raise ValueError( - f"The number of initial states is expected to be equal to the number of input sequences, " - f"i.e., {len(cu_seqlens) - 1} rather than {initial_state.shape[0]}." - ) - if initial_state is not None: - assert initial_state.dtype == torch.float32, "initial_state must be in float32." - - A_log, dt_bias = None, None - if use_gate_in_kernel: - assert "A_log" in kwargs, "A_log must be provided when use_gate_in_kernel=True." - A_log, dt_bias = kwargs["A_log"], kwargs.get("dt_bias") - - chunk_size = kwargs.pop("chunk_size", 64) - if chunk_size != 64: - # The blocked solver in chunk_intra is hardcoded to NC=4 sub-chunks of - # size 16, which assumes BT == 64. - raise ValueError(f"`chunk_size` must be 64 for GDN-2, got {chunk_size}.") - - if safe_gate and use_gate_in_kernel: - if lower_bound is None: - raise ValueError("`lower_bound` must be specified when `safe_gate=True` and `use_gate_in_kernel=True`.") - if not (-5 <= lower_bound < 0): - raise ValueError(f"`lower_bound` must be in the safe range [-5, 0), got {lower_bound}.") - - assert q.shape == k.shape == g.shape, "q, k, g must have the same shape." - assert k.shape[-1] <= 256, f"GDN-2 only supports key headdim <= 256, got {k.shape[-1]}." - assert b.shape == q.shape, ( - f"b (channel-wise erase gate) must have shape [B, T, H, K] matching q; " - f"got {tuple(b.shape)} vs q {tuple(q.shape)}." - ) - assert w.shape == v.shape, ( - f"w (channel-wise write gate) must have shape [B, T, H, V] matching v; " - f"got {tuple(w.shape)} vs v {tuple(v.shape)}." - ) - assert v.shape == (*q.shape[:3], v.shape[-1]), ( - f"v must have shape [B, T, H, V] matching q on (B, T, H); " - f"got {tuple(v.shape)} vs q {tuple(q.shape)}." - ) - - if scale is None: - scale = k.shape[-1] ** -0.5 - - return ChunkGDN2Function.apply( - q, - k, - v, - g, - b, - w, - A_log, - dt_bias, - scale, - initial_state, - output_final_state, - use_qk_l2norm_in_kernel, - use_gate_in_kernel, - cu_seqlens, - cu_seqlens_cpu, - safe_gate, - lower_bound, - chunk_size, - disable_recompute, - return_intermediate_states, - state_v_first, - ) diff --git a/build/torch-xpu/ops/gdn2/chunk_bwd.py b/build/torch-xpu/ops/gdn2/chunk_bwd.py deleted file mode 100644 index a75b316ad2e61167ae7316a851528120efcb732c..0000000000000000000000000000000000000000 --- a/build/torch-xpu/ops/gdn2/chunk_bwd.py +++ /dev/null @@ -1,467 +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 - -r""" -Backward pipeline for GDN-2 chunkwise training. - -This file owns the GDN-2-specific WY-aware backward kernel and the end-to-end -backward orchestration. The inter-chunk state-recurrence backward -(``chunk_gated_delta_rule_bwd_dhu``) and the dAv backward -(``chunk_kda_bwd_dAv``) are shared with KDA and imported as-is. - -The crux of the GDN-2 backward is the WY-aware vector-Jacobian product through -``A = (I + tril(T,-1))^{-1}``. For a scalar write strength ``beta`` the -contribution of ``u`` to ``dA`` factors as ``dU @ (beta * V)^T = beta * (dU @ -V^T)``, so KDA can apply ``beta`` as a row-scalar post-multiply. GDN-2 replaces -``beta`` with two channel-wise gates ``b`` (key axis) and ``w_gate`` (value -axis) that live on different axes; no scalar post-scale can recover them. The -kernel therefore bakes the gates directly into the dA accumulation: - - dA += dU @ (w_gate * V)^T - dA += dW @ (b * exp(gk) * K)^T - -and produces the partial gradients db (channel-wise, K-dim) and dw (channel- -wise, V-dim) at the same time. -""" - -from __future__ import annotations - -import torch -import triton -import triton.language as tl - -from ...ops.common.chunk_delta_h import chunk_gated_delta_rule_bwd_dhu, chunk_gated_delta_rule_fwd_h -from ...ops.gdn2.chunk_intra import chunk_gdn2_bwd_intra -from ...ops.gdn2.wy_fast import recompute_w_u_fwd_gdn2 -from ...ops.kda.chunk_bwd import chunk_kda_bwd_dAv -from ...ops.kda.gate import kda_gate_bwd, kda_gate_chunk_cumsum -from ...ops.utils import chunk_local_cumsum, prepare_chunk_indices -from ...ops.utils.cache import fla_cache_autotune -from ...ops.utils.constant import RCP_LN2 -from ...ops.utils.op import exp2 -from ...utils import IS_NVIDIA_HOPPER, autotune_cache_kwargs - -NUM_WARPS_WY = [2, 4] if IS_NVIDIA_HOPPER else [2, 4, 8] - - -@triton.heuristics({ - 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None, -}) -@fla_cache_autotune( - configs=[ - triton.Config({'BK': BK, 'BV': BV}, num_warps=num_warps, num_stages=num_stages) - for BK in [32, 64] - for BV in [32, 64] - for num_warps in NUM_WARPS_WY - for num_stages in [2, 3, 4] - if not (IS_NVIDIA_HOPPER and BK == 32 and num_warps == 4) - ], - key=['BT', 'STATE_V_FIRST'], - **autotune_cache_kwargs, -) -@triton.jit(do_not_specialize=['T']) -def chunk_gdn2_bwd_kernel_wy_dqkg_fused( - q, - k, - v, - v_new, - g, - b, - w_gate, - A, - h, - do, - dh, - dq, - dk, - dv, - dv2, - dg, - db, - dw, - dA, - cu_seqlens, - chunk_indices, - scale, - T, - H: tl.constexpr, - K: tl.constexpr, - V: tl.constexpr, - BT: tl.constexpr, - BK: tl.constexpr, - BV: tl.constexpr, - STATE_V_FIRST: tl.constexpr, - IS_VARLEN: tl.constexpr, -): - i_t, i_bh = tl.program_id(0), tl.program_id(1) - i_b, i_h = i_bh // H, i_bh % H - - if IS_VARLEN: - i_tg = i_t.to(tl.int64) - 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).to(tl.int32) - NT = tl.cdiv(T, BT) - else: - NT = tl.cdiv(T, BT) - i_tg = (i_b * NT + i_t).to(tl.int64) - bos, eos = (i_b * T).to(tl.int64), (i_b * T + T).to(tl.int64) - - o_t = i_t * BT + tl.arange(0, BT) - m_t = o_t < T - m_last = (o_t == min(T, i_t * BT + BT) - 1) - - q += (bos * H + i_h) * K - k += (bos * H + i_h) * K - v += (bos * H + i_h) * V - v_new += (bos * H + i_h) * V - g += (bos * H + i_h) * K - b += (bos * H + i_h) * K - w_gate += (bos * H + i_h) * V - A += (bos * H + i_h) * BT - h += (i_tg * H + i_h) * K * V - do += (bos * H + i_h) * V - dh += (i_tg * H + i_h) * K * V - dq += (bos * H + i_h) * K - dk += (bos * H + i_h) * K - dv += (bos * H + i_h) * V - dv2 += (bos * H + i_h) * V - dg += (bos * H + i_h) * K - db += (bos * H + i_h) * K - dw += (bos * H + i_h) * V - dA += (bos * H + i_h) * BT - - p_A = tl.make_block_ptr(A, (BT, T), (1, H * BT), (0, i_t * BT), (BT, BT), (0, 1)) - b_A = tl.load(p_A, boundary_check=(0, 1)) - - b_dA = tl.zeros([BT, BT], dtype=tl.float32) - - for i_k in range(tl.cdiv(K, BK)): - o_k = i_k * BK + tl.arange(0, BK) - m_k = o_k < K - - p_k = tl.make_block_ptr(k, (T, K), (H * K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) - p_g = tl.make_block_ptr(g, (T, K), (H * K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) - p_b = tl.make_block_ptr(b, (T, K), (H * K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) - b_k = tl.load(p_k, boundary_check=(0, 1)) - b_g = tl.load(p_g, boundary_check=(0, 1)).to(tl.float32) - b_b = tl.load(p_b, boundary_check=(0, 1)) - - p_gn = g + (min(T, i_t * BT + BT) - 1).to(tl.int64) * H * K + o_k - b_gn = tl.load(p_gn, mask=m_k, other=0).to(tl.float32) - - b_dq = tl.zeros([BT, BK], dtype=tl.float32) - b_dk = tl.zeros([BT, BK], dtype=tl.float32) - b_dw_flow = tl.zeros([BT, BK], dtype=tl.float32) - b_dgk = tl.zeros([BK], dtype=tl.float32) - - for i_v in range(tl.cdiv(V, BV)): - p_v_new = tl.make_block_ptr(v_new, (T, V), (H * V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) - p_do = tl.make_block_ptr(do, (T, V), (H * V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) - if STATE_V_FIRST: - p_h = tl.make_block_ptr(h, (V, K), (K, 1), (i_v * BV, i_k * BK), (BV, BK), (1, 0)) - p_dh = tl.make_block_ptr(dh, (V, K), (K, 1), (i_v * BV, i_k * BK), (BV, BK), (1, 0)) - else: - p_h = tl.make_block_ptr(h, (V, K), (1, V), (i_v * BV, i_k * BK), (BV, BK), (0, 1)) - p_dh = tl.make_block_ptr(dh, (V, K), (1, V), (i_v * BV, i_k * BK), (BV, BK), (0, 1)) - p_dv = tl.make_block_ptr(dv, (T, V), (H * V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) - b_v_new = tl.load(p_v_new, boundary_check=(0, 1)) - b_do = tl.load(p_do, boundary_check=(0, 1)) - b_h = tl.load(p_h, boundary_check=(0, 1)) - b_dh = tl.load(p_dh, boundary_check=(0, 1)) - b_dv = tl.load(p_dv, boundary_check=(0, 1)) - - b_dgk += tl.sum(b_h * b_dh, axis=0) - b_dq += tl.dot(b_do, b_h.to(b_do.dtype)) - b_dk += tl.dot(b_v_new, b_dh.to(b_v_new.dtype)) - b_dw_flow += tl.dot(b_dv.to(b_v_new.dtype), b_h.to(b_v_new.dtype)) - tl.debug_barrier() - - if i_k == 0: - p_v = tl.make_block_ptr(v, (T, V), (H * V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) - p_dv2 = tl.make_block_ptr(dv2, (T, V), (H * V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) - p_wg = tl.make_block_ptr(w_gate, (T, V), (H * V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) - p_dw_gate = tl.make_block_ptr(dw, (T, V), (H * V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) - - b_v = tl.load(p_v, boundary_check=(0, 1)) - b_wg = tl.load(p_wg, boundary_check=(0, 1)) - # dA gets (w_gate * v) on the value side - the GDN-2 channel-wise twist. - b_dA += tl.dot(b_dv, tl.trans(b_v * b_wg)) - - b_dvb = tl.dot(b_A, b_dv) - b_dv2 = b_dvb * b_wg - b_dw_gate = b_dvb * b_v - - tl.store(p_dv2, b_dv2.to(p_dv2.dtype.element_ty), boundary_check=(0, 1)) - tl.store(p_dw_gate, b_dw_gate.to(p_dw_gate.dtype.element_ty), boundary_check=(0, 1)) - - b_gk_exp = exp2(b_g) - b_gb = b_gk_exp * b_b - b_dgk *= exp2(b_gn) - b_dq = b_dq * b_gk_exp * scale - b_dk = b_dk * tl.where(m_t[:, None], exp2(b_gn[None, :] - b_g), 0) - - b_kg = b_k * b_gk_exp - - b_dw_flow = -b_dw_flow.to(b_A.dtype) - # dA gets (b * exp(gk) * k) on the key side - the GDN-2 channel-wise twist. - b_dA += tl.dot(b_dw_flow, tl.trans((b_kg * b_b).to(b_A.dtype))) - - b_dkgb = tl.dot(b_A, b_dw_flow) - p_db = tl.make_block_ptr(db, (T, K), (H * K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) - b_db_partial = b_dkgb * b_kg - tl.store(p_db, b_db_partial.to(p_db.dtype.element_ty), boundary_check=(0, 1)) - - p_q = tl.make_block_ptr(q, (T, K), (H * K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) - b_q = tl.load(p_q, boundary_check=(0, 1)) - b_kdk = b_k * b_dk - b_dgk += tl.sum(b_kdk, axis=0) - b_dg = b_q * b_dq - b_kdk + m_last[:, None] * b_dgk + b_kg * b_dkgb * b_b - b_dk = b_dk + b_dkgb * b_gb - - p_dq = tl.make_block_ptr(dq, (T, K), (H * K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) - p_dk = tl.make_block_ptr(dk, (T, K), (H * K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) - p_dg = tl.make_block_ptr(dg, (T, K), (H * K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) - tl.store(p_dq, b_dq.to(p_dq.dtype.element_ty), boundary_check=(0, 1)) - tl.store(p_dk, b_dk.to(p_dk.dtype.element_ty), boundary_check=(0, 1)) - tl.store(p_dg, b_dg.to(p_dg.dtype.element_ty), boundary_check=(0, 1)) - - m_A = (o_t[:, None] > o_t[None, :]) & (m_t[:, None] & m_t) - b_dA = tl.where(m_A, b_dA, 0) - b_dA = tl.dot(b_dA.to(b_A.dtype), b_A) - b_dA = tl.dot(b_A, b_dA.to(b_A.dtype)) - b_dA = tl.where(m_A, -b_dA, 0) - p_dA = tl.make_block_ptr(dA, (T, BT), (H * BT, 1), (i_t * BT, 0), (BT, BT), (1, 0)) - tl.store(p_dA, b_dA.to(p_dA.dtype.element_ty), boundary_check=(0, 1)) - - -def chunk_gdn2_bwd_wy_dqkg_fused( - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - v_new: torch.Tensor, - g: torch.Tensor, - b: torch.Tensor, - w_gate: torch.Tensor, - A: torch.Tensor, - h: torch.Tensor, - do: torch.Tensor, - dh: torch.Tensor, - dv: torch.Tensor, - scale: float | None = None, - cu_seqlens: torch.LongTensor | None = None, - chunk_size: int = 64, - chunk_indices: torch.LongTensor | None = None, - state_v_first: bool = False, -): - """Fused WY backward producing dq, dk, dv, db (K-dim), dw (V-dim), dg, dA.""" - B, T, H, K, V = *k.shape, v.shape[-1] - BT = chunk_size - - if chunk_indices is None and cu_seqlens is not None: - chunk_indices = prepare_chunk_indices(cu_seqlens, BT) - NT = triton.cdiv(T, BT) if cu_seqlens is None else len(chunk_indices) - - dq = torch.empty_like(q, dtype=torch.float32) - dk = torch.empty_like(k, dtype=torch.float32) - dv2 = torch.empty_like(v) - dg = torch.empty_like(g, dtype=torch.float32) - db = torch.empty_like(b, dtype=torch.float32) - dw = torch.empty_like(w_gate, dtype=torch.float32) - dA = torch.empty_like(A, dtype=torch.float32) - - grid = (NT, B * H) - chunk_gdn2_bwd_kernel_wy_dqkg_fused[grid]( - q=q, - k=k, - v=v, - v_new=v_new, - g=g, - b=b, - w_gate=w_gate, - A=A, - h=h, - do=do, - dh=dh, - dq=dq, - dk=dk, - dv=dv, - dv2=dv2, - dg=dg, - db=db, - dw=dw, - dA=dA, - cu_seqlens=cu_seqlens, - chunk_indices=chunk_indices, - scale=scale, - T=T, - H=H, - K=K, - V=V, - BT=BT, - STATE_V_FIRST=state_v_first, - ) - dv = dv2 - return dq, dk, dv, db, dw, dg, dA - - -def chunk_gdn2_bwd( - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - b: torch.Tensor, - w_gate: torch.Tensor, - Aqk: torch.Tensor, - Akk: torch.Tensor, - scale: float, - initial_state: torch.Tensor | None, - do: torch.Tensor, - dht: torch.Tensor | None, - g: torch.Tensor | None = None, - g_org: torch.Tensor | None = None, - cu_seqlens: torch.LongTensor | None = None, - chunk_indices: torch.LongTensor | None = None, - chunk_size: int = 64, - safe_gate: bool = False, - lower_bound: float | None = None, - use_gate_in_kernel: bool = False, - A_log: torch.Tensor | None = None, - dt_bias: torch.Tensor | None = None, - state_v_first: bool = False, - w_wy: torch.Tensor | None = None, - u_wy: torch.Tensor | None = None, - qg: torch.Tensor | None = None, - kg: torch.Tensor | None = None, - v_new: torch.Tensor | None = None, - h: torch.Tensor | None = None, - disable_recompute: bool = False, -): - """End-to-end GDN-2 backward. - - Returns (dq, dk, dv, db, dw, dg, dh0, dA_log, dt_bias_grad). ``db`` has - shape [B, T, H, K] (channel-wise erase gate); ``dw`` has shape [B, T, H, V] - (channel-wise write gate). - """ - if not disable_recompute: - if use_gate_in_kernel: - g = kda_gate_chunk_cumsum( - g=g_org, - A_log=A_log, - dt_bias=dt_bias, - scale=RCP_LN2, - chunk_size=chunk_size, - cu_seqlens=cu_seqlens, - chunk_indices=chunk_indices, - lower_bound=lower_bound, - ) - w_wy, u_wy, qg, kg = recompute_w_u_fwd_gdn2( - k=k, - v=v, - b=b, - w_gate=w_gate, - A=Akk, - q=q, - gk=g, - cu_seqlens=cu_seqlens, - chunk_indices=chunk_indices, - ) - h, v_new, _ = chunk_gated_delta_rule_fwd_h( - k=kg, - w=w_wy, - u=u_wy, - gk=g, - initial_state=initial_state, - output_final_state=False, - cu_seqlens=cu_seqlens, - chunk_indices=chunk_indices, - chunk_size=chunk_size, - state_v_first=state_v_first, - ) - - dAqk, dv = chunk_kda_bwd_dAv( - q=q, - k=k, - v=v_new, - do=do, - A=Aqk, - scale=scale, - cu_seqlens=cu_seqlens, - chunk_size=chunk_size, - chunk_indices=chunk_indices, - ) - - dh, dh0, dv = chunk_gated_delta_rule_bwd_dhu( - q=qg, - k=kg, - w=w_wy, - gk=g, - h0=initial_state, - dht=dht, - do=do, - dv=dv, - scale=scale, - cu_seqlens=cu_seqlens, - chunk_indices=chunk_indices, - chunk_size=chunk_size, - state_v_first=state_v_first, - ) - - dq, dk, dv, db, dw, dg, dAkk = chunk_gdn2_bwd_wy_dqkg_fused( - q=q, - k=k, - v=v, - v_new=v_new, - g=g, - b=b, - w_gate=w_gate, - A=Akk, - h=h, - do=do, - dh=dh, - dv=dv, - scale=scale, - cu_seqlens=cu_seqlens, - chunk_size=chunk_size, - chunk_indices=chunk_indices, - state_v_first=state_v_first, - ) - - dq, dk, db, dg = chunk_gdn2_bwd_intra( - q=q, - k=k, - g=g, - b=b, - dAqk=dAqk, - dAkk=dAkk, - dq=dq, - dk=dk, - db=db, - dg=dg, - cu_seqlens=cu_seqlens, - chunk_size=chunk_size, - chunk_indices=chunk_indices, - safe_gate=safe_gate, - ) - - dA_log, dt_bias_grad = None, None - dg = chunk_local_cumsum( - dg, - chunk_size=chunk_size, - reverse=True, - cu_seqlens=cu_seqlens, - chunk_indices=chunk_indices, - ) - if use_gate_in_kernel: - dg, dA_log, dt_bias_grad = kda_gate_bwd( - g=g_org, - A_log=A_log, - dt_bias=dt_bias, - dyg=dg, - lower_bound=lower_bound, - ) - - return dq, dk, dv, db, dw, dg, dh0, dA_log, dt_bias_grad diff --git a/build/torch-xpu/ops/gdn2/chunk_fwd.py b/build/torch-xpu/ops/gdn2/chunk_fwd.py deleted file mode 100644 index 0b913f748f26a7eec931c82f6cfefaebe4877ced..0000000000000000000000000000000000000000 --- a/build/torch-xpu/ops/gdn2/chunk_fwd.py +++ /dev/null @@ -1,126 +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 - -# Forward orchestration for GDN-2 chunkwise training. - -import torch - -from ...ops.common.chunk_delta_h import chunk_gated_delta_rule_fwd_h -from ...ops.gdn2.chunk_intra import chunk_gdn2_fwd_intra -from ...ops.gla.chunk import chunk_gla_fwd_o_gk -from ...ops.kda.gate import kda_gate_chunk_cumsum -from ...ops.utils import chunk_local_cumsum -from ...ops.utils.constant import RCP_LN2 - - -def chunk_gdn2_fwd( - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - g: torch.Tensor, - b: torch.Tensor, - w_gate: torch.Tensor, - scale: float, - initial_state: torch.Tensor | None, - output_final_state: bool, - cu_seqlens: torch.LongTensor | None = None, - cu_seqlens_cpu: torch.LongTensor | None = None, - chunk_indices: torch.LongTensor | None = None, - chunk_size: int = 64, - safe_gate: bool = False, - lower_bound: float | None = None, - use_gate_in_kernel: bool = False, - A_log: torch.Tensor | None = None, - dt_bias: torch.Tensor | None = None, - disable_recompute: bool = False, - return_intermediate_states: bool = False, - state_v_first: bool = False, -): - """Top-level GDN-2 forward pipeline. - - The pipeline is: - 1. Compute the base-2 log-decay cumsum within each chunk - (``kda_gate_chunk_cumsum`` if ``use_gate_in_kernel`` else - ``chunk_local_cumsum``). - 2. Build the intra-chunk score matrices (Aqk, Akk_inv) and the WY - auxiliaries (w_wy, u_wy, qg, kg) via ``chunk_gdn2_fwd_intra``. - 3. Run the inter-chunk state recurrence (shared with KDA / GDN v1). - 4. Compose the output via the GLA-style ``chunk_gla_fwd_o_gk``. - - Returns ``(o, final_state, g_cumsum, Aqk, Akk, w_wy, u_wy, qg, kg, v_new, - h, initial_state)``. - """ - if use_gate_in_kernel: - g = kda_gate_chunk_cumsum( - g=g, - A_log=A_log, - dt_bias=dt_bias, - scale=RCP_LN2, - chunk_size=chunk_size, - cu_seqlens=cu_seqlens, - chunk_indices=chunk_indices, - lower_bound=lower_bound, - ) - else: - g = chunk_local_cumsum( - g=g, - scale=RCP_LN2, - chunk_size=chunk_size, - cu_seqlens=cu_seqlens, - chunk_indices=chunk_indices, - ) - - w_wy, u_wy, qg, kg, Aqk, Akk = chunk_gdn2_fwd_intra( - q=q, - k=k, - v=v, - gk=g, - b=b, - w_gate=w_gate, - scale=scale, - cu_seqlens=cu_seqlens, - chunk_size=chunk_size, - chunk_indices=chunk_indices, - safe_gate=safe_gate, - disable_recompute=disable_recompute, - ) - - h, v_new, final_state = chunk_gated_delta_rule_fwd_h( - k=kg, - w=w_wy, - u=u_wy, - gk=g, - initial_state=initial_state, - output_final_state=output_final_state, - cu_seqlens=cu_seqlens, - cu_seqlens_cpu=cu_seqlens_cpu, - chunk_indices=chunk_indices, - chunk_size=chunk_size, - state_v_first=state_v_first, - ) - - o = chunk_gla_fwd_o_gk( - q=q, - v=v_new, - g=g, - A=Aqk, - h=h, - scale=scale, - cu_seqlens=cu_seqlens, - chunk_size=chunk_size, - chunk_indices=chunk_indices, - state_v_first=state_v_first, - ) - - if disable_recompute is False: - # Free intermediates that the backward will recompute. - w_wy, u_wy, qg, kg, v_new = None, None, None, None, None - if not return_intermediate_states: - h = None - if use_gate_in_kernel: - g = None - return o, final_state, g, Aqk, Akk, w_wy, u_wy, qg, kg, v_new, h, initial_state diff --git a/build/torch-xpu/ops/gdn2/chunk_intra.py b/build/torch-xpu/ops/gdn2/chunk_intra.py deleted file mode 100644 index 510114d1ee4abe108e246ba530bed0a2caae8efa..0000000000000000000000000000000000000000 --- a/build/torch-xpu/ops/gdn2/chunk_intra.py +++ /dev/null @@ -1,843 +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 - -# Intra-chunk Triton kernels for GDN-2. -# -# Three kernels live here: -# * chunk_gdn2_fwd_kernel_intra_sub_chunk - sub-chunk diagonal builder used -# by the safe-gate path (alternative to the token-parallel diagonal kernel). -# * chunk_gdn2_fwd_kernel_inter_solve_fused - off-diagonal Akk blocks + the -# blocked WY triangular solve, fused. Reads diagonal blocks from Akkd, emits -# the full lower-triangular inverse to Akk. -# * chunk_gdn2_bwd_kernel_intra - backward over the within-chunk -# score matrices, producing dq, dk, db, dg contributions. -# -# Compared with KDA, the GDN-2 variant takes a channel-wise erase gate ``b`` on -# the K axis instead of a per-head scalar ``beta``; the b tensor is folded into -# the relevant tiles directly rather than applied as a row-scalar multiply. - -import torch -import triton -import triton.language as tl - -from ...ops.gdn2.chunk_intra_token_parallel import chunk_gdn2_fwd_intra_token_parallel -from ...ops.gdn2.wy_fast import recompute_w_u_fwd_gdn2 -from ...ops.utils import prepare_chunk_indices -from ...ops.utils.cache import fla_cache_autotune -from ...ops.utils.op import exp2, gather -from ...utils import IS_GATHER_SUPPORTED, IS_TF32_SUPPORTED, autotune_cache_kwargs - -if IS_TF32_SUPPORTED: - SOLVE_TRIL_DOT_PRECISION = tl.constexpr('tf32') -else: - SOLVE_TRIL_DOT_PRECISION = tl.constexpr('ieee') - - -@triton.heuristics({ - 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None, -}) -@fla_cache_autotune( - configs=[ - triton.Config({}, num_warps=num_warps, num_stages=num_stages) - for num_warps in [1, 2, 4, 8] - for num_stages in [2, 3, 4] - ], - key=["BT", "BC"], - **autotune_cache_kwargs, -) -@triton.jit(do_not_specialize=['T']) -def chunk_gdn2_fwd_kernel_intra_sub_chunk( - q, - k, - g, - b, - Aqk, - Akk, - scale, - cu_seqlens, - chunk_indices, - T, - H: tl.constexpr, - K: tl.constexpr, - BT: tl.constexpr, - BC: tl.constexpr, - BK: tl.constexpr, - IS_VARLEN: tl.constexpr, - USE_GATHER: tl.constexpr, -): - i_t, i_i, i_bh = tl.program_id(0), tl.program_id(1), tl.program_id(2) - i_b, i_h = i_bh // H, i_bh % H - - 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.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) - T = eos - bos - else: - bos, eos = i_b * T, i_b * T + T - - i_ti = i_t * BT + i_i * BC - if i_ti >= T: - return - - o_c = i_ti + tl.arange(0, BC) - m_c = o_c < T - - q = q + (bos * H + i_h) * K - k = k + (bos * H + i_h) * K - g = g + (bos * H + i_h) * K - b = b + (bos * H + i_h) * K - Aqk = Aqk + (bos * H + i_h) * BT - Akk = Akk + (bos * H + i_h) * BC - - p_q = tl.make_block_ptr(q, (T, K), (H * K, 1), (i_ti, 0), (BC, BK), (1, 0)) - p_k = tl.make_block_ptr(k, (T, K), (H * K, 1), (i_ti, 0), (BC, BK), (1, 0)) - p_g = tl.make_block_ptr(g, (T, K), (H * K, 1), (i_ti, 0), (BC, BK), (1, 0)) - p_b = tl.make_block_ptr(b, (T, K), (H * K, 1), (i_ti, 0), (BC, BK), (1, 0)) - - b_q = tl.load(p_q, boundary_check=(0, 1)) - b_k = tl.load(p_k, boundary_check=(0, 1)) - b_g = tl.load(p_g, boundary_check=(0, 1)) - b_b = tl.load(p_b, boundary_check=(0, 1)) - - if USE_GATHER: - b_gn = gather(b_g, tl.full([1, BK], min(BC // 2, T - i_ti - 1), dtype=tl.int16), axis=0) - else: - p_gn = g + (i_ti + min(BC // 2, T - i_ti - 1)) * H * K + tl.arange(0, BK) - b_gn = tl.load(p_gn, mask=tl.arange(0, BK) < K, other=0.0) - b_gn = b_gn[None, :] - - b_gm = (b_g - b_gn).to(tl.float32) - b_gq = tl.where(m_c[:, None], exp2(b_gm), 0.) - b_gk = tl.where(m_c[:, None], exp2(-b_gm), 0.) - - b_kgt = tl.trans(b_k * b_gk) - - b_bk = (b_b.to(tl.float32) * b_k.to(tl.float32)).to(b_k.dtype) - - b_Aqk = tl.dot(b_q * b_gq, b_kgt) * scale - b_Akk = tl.dot(b_bk * b_gq, b_kgt) - - o_i = tl.arange(0, BC) - m_Aqk = o_i[:, None] >= o_i[None, :] - m_Akk = o_i[:, None] > o_i[None, :] - m_I = o_i[:, None] == o_i[None, :] - - b_Aqk = tl.where(m_Aqk, b_Aqk, 0.0) - b_Akk = tl.where(m_Akk, b_Akk, 0.0) - - p_Aqk = tl.make_block_ptr(Aqk, (T, BT), (H * BT, 1), (i_ti, i_i * BC), (BC, BC), (1, 0)) - p_Akk = tl.make_block_ptr(Akk, (T, BC), (H * BC, 1), (i_ti, 0), (BC, BC), (1, 0)) - tl.store(p_Aqk, b_Aqk.to(Aqk.dtype.element_ty), boundary_check=(0, 1)) - tl.store(p_Akk, b_Akk.to(Akk.dtype.element_ty), boundary_check=(0, 1)) - - tl.debug_barrier() - - b_Ai = -b_Akk - for i in range(2, min(BC, T - i_ti)): - b_a = -tl.load(Akk + (i_ti + i) * H * BC + o_i) - b_a = tl.where(o_i < i, b_a, 0.) - b_a += tl.sum(b_a[:, None] * b_Ai, 0) - b_Ai = tl.where((o_i == i)[:, None], b_a, b_Ai) - b_Ai += m_I - tl.store(p_Akk, b_Ai.to(Akk.dtype.element_ty), boundary_check=(0, 1)) - - -@triton.heuristics({ - 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None, -}) -@fla_cache_autotune( - configs=[ - triton.Config({'BK': BK}, num_warps=num_warps) - for BK in [32, 64] - for num_warps in [1, 2, 4] - ], - key=["H", "K", "BC"], - **autotune_cache_kwargs, -) -@triton.jit(do_not_specialize=['T']) -def chunk_gdn2_fwd_kernel_inter_solve_fused( - q, - k, - g, - b, - Aqk, - Akkd, - Akk, - scale, - cu_seqlens, - chunk_indices, - T, - H: tl.constexpr, - K: tl.constexpr, - BT: tl.constexpr, - BC: tl.constexpr, - BK: tl.constexpr, - IS_VARLEN: tl.constexpr, - USE_SAFE_GATE: tl.constexpr, -): - i_t, i_bh = tl.program_id(0), tl.program_id(1) - i_b, i_h = i_bh // H, i_bh % H - - 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.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) - T = eos - bos - else: - bos, eos = i_b * T, i_b * T + T - - if i_t * BT >= T: - return - - i_tc0 = i_t * BT - i_tc1 = i_t * BT + BC - i_tc2 = i_t * BT + 2 * BC - i_tc3 = i_t * BT + 3 * BC - - q += (bos * H + i_h) * K - k += (bos * H + i_h) * K - g += (bos * H + i_h) * K - b += (bos * H + i_h) * K - Aqk += (bos * H + i_h) * BT - Akk += (bos * H + i_h) * BT - Akkd += (bos * H + i_h) * BC - - o_i = tl.arange(0, BC) - m_tc1 = (i_tc1 + o_i) < T - m_tc2 = (i_tc2 + o_i) < T - m_tc3 = (i_tc3 + o_i) < T - - b_Aqk10 = tl.zeros([BC, BC], dtype=tl.float32) - b_Akk10 = tl.zeros([BC, BC], dtype=tl.float32) - - b_Aqk20 = tl.zeros([BC, BC], dtype=tl.float32) - b_Akk20 = tl.zeros([BC, BC], dtype=tl.float32) - b_Aqk21 = tl.zeros([BC, BC], dtype=tl.float32) - b_Akk21 = tl.zeros([BC, BC], dtype=tl.float32) - - b_Aqk30 = tl.zeros([BC, BC], dtype=tl.float32) - b_Akk30 = tl.zeros([BC, BC], dtype=tl.float32) - b_Aqk31 = tl.zeros([BC, BC], dtype=tl.float32) - b_Akk31 = tl.zeros([BC, BC], dtype=tl.float32) - b_Aqk32 = tl.zeros([BC, BC], dtype=tl.float32) - b_Akk32 = tl.zeros([BC, BC], dtype=tl.float32) - - for i_k in range(tl.cdiv(K, BK)): - o_k = i_k * BK + tl.arange(0, BK) - m_k = o_k < K - - p_k0 = tl.make_block_ptr(k, (T, K), (H * K, 1), (i_tc0, i_k * BK), (BC, BK), (1, 0)) - p_g0 = tl.make_block_ptr(g, (T, K), (H * K, 1), (i_tc0, i_k * BK), (BC, BK), (1, 0)) - b_k0 = tl.load(p_k0, boundary_check=(0, 1)).to(tl.float32) - b_g0 = tl.load(p_g0, boundary_check=(0, 1)).to(tl.float32) - - if i_tc1 < T: - p_q1 = tl.make_block_ptr(q, (T, K), (H * K, 1), (i_tc1, i_k * BK), (BC, BK), (1, 0)) - p_k1 = tl.make_block_ptr(k, (T, K), (H * K, 1), (i_tc1, i_k * BK), (BC, BK), (1, 0)) - p_g1 = tl.make_block_ptr(g, (T, K), (H * K, 1), (i_tc1, i_k * BK), (BC, BK), (1, 0)) - p_b1 = tl.make_block_ptr(b, (T, K), (H * K, 1), (i_tc1, i_k * BK), (BC, BK), (1, 0)) - b_q1 = tl.load(p_q1, boundary_check=(0, 1)).to(tl.float32) - b_k1 = tl.load(p_k1, boundary_check=(0, 1)).to(tl.float32) - b_g1 = tl.load(p_g1, boundary_check=(0, 1)).to(tl.float32) - b_b1 = tl.load(p_b1, boundary_check=(0, 1)).to(tl.float32) - b_gn1 = tl.load(g + i_tc1 * H * K + o_k, mask=m_k, other=0).to(tl.float32) - b_gqn = tl.where(m_tc1[:, None], exp2(b_g1 - b_gn1[None, :]), 0) - b_kgt = tl.trans(b_k0 * exp2(b_gn1[None, :] - b_g0)) - b_bk1 = b_b1 * b_k1 - b_Aqk10 += tl.dot(b_q1 * b_gqn, b_kgt) - b_Akk10 += tl.dot(b_bk1 * b_gqn, b_kgt) - - if i_tc2 < T: - p_q2 = tl.make_block_ptr(q, (T, K), (H * K, 1), (i_tc2, i_k * BK), (BC, BK), (1, 0)) - p_k2 = tl.make_block_ptr(k, (T, K), (H * K, 1), (i_tc2, i_k * BK), (BC, BK), (1, 0)) - p_g2 = tl.make_block_ptr(g, (T, K), (H * K, 1), (i_tc2, i_k * BK), (BC, BK), (1, 0)) - p_b2 = tl.make_block_ptr(b, (T, K), (H * K, 1), (i_tc2, i_k * BK), (BC, BK), (1, 0)) - b_q2 = tl.load(p_q2, boundary_check=(0, 1)).to(tl.float32) - b_k2 = tl.load(p_k2, boundary_check=(0, 1)).to(tl.float32) - b_g2 = tl.load(p_g2, boundary_check=(0, 1)).to(tl.float32) - b_b2 = tl.load(p_b2, boundary_check=(0, 1)).to(tl.float32) - b_gn2 = tl.load(g + i_tc2 * H * K + o_k, mask=m_k, other=0).to(tl.float32) - b_gqn2 = tl.where(m_tc2[:, None], exp2(b_g2 - b_gn2[None, :]), 0) - b_qg2 = b_q2 * b_gqn2 - b_bkg2 = (b_b2 * b_k2) * b_gqn2 - b_kgt = tl.trans(b_k0 * exp2(b_gn2[None, :] - b_g0)) - b_Aqk20 += tl.dot(b_qg2, b_kgt) - b_Akk20 += tl.dot(b_bkg2, b_kgt) - b_kgt = tl.trans(b_k1 * exp2(b_gn2[None, :] - b_g1)) - b_Aqk21 += tl.dot(b_qg2, b_kgt) - b_Akk21 += tl.dot(b_bkg2, b_kgt) - - if i_tc3 < T: - p_q3 = tl.make_block_ptr(q, (T, K), (H * K, 1), (i_tc3, i_k * BK), (BC, BK), (1, 0)) - p_k3 = tl.make_block_ptr(k, (T, K), (H * K, 1), (i_tc3, i_k * BK), (BC, BK), (1, 0)) - p_g3 = tl.make_block_ptr(g, (T, K), (H * K, 1), (i_tc3, i_k * BK), (BC, BK), (1, 0)) - p_b3 = tl.make_block_ptr(b, (T, K), (H * K, 1), (i_tc3, i_k * BK), (BC, BK), (1, 0)) - b_q3 = tl.load(p_q3, boundary_check=(0, 1)).to(tl.float32) - b_k3 = tl.load(p_k3, boundary_check=(0, 1)).to(tl.float32) - b_g3 = tl.load(p_g3, boundary_check=(0, 1)).to(tl.float32) - b_b3 = tl.load(p_b3, boundary_check=(0, 1)).to(tl.float32) - b_gn3 = tl.load(g + i_tc3 * H * K + o_k, mask=m_k, other=0).to(tl.float32) - b_gqn3 = tl.where(m_tc3[:, None], exp2(b_g3 - b_gn3[None, :]), 0) - b_qg3 = b_q3 * b_gqn3 - b_bkg3 = (b_b3 * b_k3) * b_gqn3 - b_kgt = tl.trans(b_k0 * exp2(b_gn3[None, :] - b_g0)) - b_Aqk30 += tl.dot(b_qg3, b_kgt) - b_Akk30 += tl.dot(b_bkg3, b_kgt) - b_kgt = tl.trans(b_k1 * exp2(b_gn3[None, :] - b_g1)) - b_Aqk31 += tl.dot(b_qg3, b_kgt) - b_Akk31 += tl.dot(b_bkg3, b_kgt) - b_kgt = tl.trans(b_k2 * exp2(b_gn3[None, :] - b_g2)) - b_Aqk32 += tl.dot(b_qg3, b_kgt) - b_Akk32 += tl.dot(b_bkg3, b_kgt) - - if i_tc1 < T: - p_Aqk10 = tl.make_block_ptr(Aqk, (T, BT), (H * BT, 1), (i_tc1, 0), (BC, BC), (1, 0)) - tl.store(p_Aqk10, (b_Aqk10 * scale).to(Aqk.dtype.element_ty), boundary_check=(0, 1)) - if i_tc2 < T: - p_Aqk20 = tl.make_block_ptr(Aqk, (T, BT), (H * BT, 1), (i_tc2, 0), (BC, BC), (1, 0)) - p_Aqk21 = tl.make_block_ptr(Aqk, (T, BT), (H * BT, 1), (i_tc2, BC), (BC, BC), (1, 0)) - tl.store(p_Aqk20, (b_Aqk20 * scale).to(Aqk.dtype.element_ty), boundary_check=(0, 1)) - tl.store(p_Aqk21, (b_Aqk21 * scale).to(Aqk.dtype.element_ty), boundary_check=(0, 1)) - if i_tc3 < T: - p_Aqk30 = tl.make_block_ptr(Aqk, (T, BT), (H * BT, 1), (i_tc3, 0), (BC, BC), (1, 0)) - p_Aqk31 = tl.make_block_ptr(Aqk, (T, BT), (H * BT, 1), (i_tc3, BC), (BC, BC), (1, 0)) - p_Aqk32 = tl.make_block_ptr(Aqk, (T, BT), (H * BT, 1), (i_tc3, 2 * BC), (BC, BC), (1, 0)) - tl.store(p_Aqk30, (b_Aqk30 * scale).to(Aqk.dtype.element_ty), boundary_check=(0, 1)) - tl.store(p_Aqk31, (b_Aqk31 * scale).to(Aqk.dtype.element_ty), boundary_check=(0, 1)) - tl.store(p_Aqk32, (b_Aqk32 * scale).to(Aqk.dtype.element_ty), boundary_check=(0, 1)) - - p_Akk00 = tl.make_block_ptr(Akkd, (T, BC), (H * BC, 1), (i_tc0, 0), (BC, BC), (1, 0)) - p_Akk11 = tl.make_block_ptr(Akkd, (T, BC), (H * BC, 1), (i_tc1, 0), (BC, BC), (1, 0)) - p_Akk22 = tl.make_block_ptr(Akkd, (T, BC), (H * BC, 1), (i_tc2, 0), (BC, BC), (1, 0)) - p_Akk33 = tl.make_block_ptr(Akkd, (T, BC), (H * BC, 1), (i_tc3, 0), (BC, BC), (1, 0)) - b_Ai00 = tl.load(p_Akk00, boundary_check=(0, 1)).to(tl.float32) - b_Ai11 = tl.load(p_Akk11, boundary_check=(0, 1)).to(tl.float32) - b_Ai22 = tl.load(p_Akk22, boundary_check=(0, 1)).to(tl.float32) - b_Ai33 = tl.load(p_Akk33, boundary_check=(0, 1)).to(tl.float32) - - if not USE_SAFE_GATE: - m_A = o_i[:, None] > o_i[None, :] - m_I = o_i[:, None] == o_i[None, :] - - b_Ai00 = -tl.where(m_A, b_Ai00, 0) - b_Ai11 = -tl.where(m_A, b_Ai11, 0) - b_Ai22 = -tl.where(m_A, b_Ai22, 0) - b_Ai33 = -tl.where(m_A, b_Ai33, 0) - - for i in range(2, min(BC, T - i_tc0)): - b_a00 = -tl.load(Akkd + (i_tc0 + i) * H * BC + o_i) - b_a00 = tl.where(o_i < i, b_a00, 0.) - b_a00 += tl.sum(b_a00[:, None] * b_Ai00, 0) - b_Ai00 = tl.where((o_i == i)[:, None], b_a00, b_Ai00) - for i in range(BC + 2, min(2 * BC, T - i_tc0)): - b_a11 = -tl.load(Akkd + (i_tc0 + i) * H * BC + o_i) - b_a11 = tl.where(o_i < i - BC, b_a11, 0.) - b_a11 += tl.sum(b_a11[:, None] * b_Ai11, 0) - b_Ai11 = tl.where((o_i == i - BC)[:, None], b_a11, b_Ai11) - for i in range(2 * BC + 2, min(3 * BC, T - i_tc0)): - b_a22 = -tl.load(Akkd + (i_tc0 + i) * H * BC + o_i) - b_a22 = tl.where(o_i < i - 2 * BC, b_a22, 0.) - b_a22 += tl.sum(b_a22[:, None] * b_Ai22, 0) - b_Ai22 = tl.where((o_i == i - 2 * BC)[:, None], b_a22, b_Ai22) - for i in range(3 * BC + 2, min(4 * BC, T - i_tc0)): - b_a33 = -tl.load(Akkd + (i_tc0 + i) * H * BC + o_i) - b_a33 = tl.where(o_i < i - 3 * BC, b_a33, 0.) - b_a33 += tl.sum(b_a33[:, None] * b_Ai33, 0) - b_Ai33 = tl.where((o_i == i - 3 * BC)[:, None], b_a33, b_Ai33) - - b_Ai00 += m_I - b_Ai11 += m_I - b_Ai22 += m_I - b_Ai33 += m_I - - b_Ai10 = -tl.dot( - tl.dot(b_Ai11, b_Akk10, input_precision=SOLVE_TRIL_DOT_PRECISION), - b_Ai00, - input_precision=SOLVE_TRIL_DOT_PRECISION, - ) - b_Ai21 = -tl.dot( - tl.dot(b_Ai22, b_Akk21, input_precision=SOLVE_TRIL_DOT_PRECISION), - b_Ai11, - input_precision=SOLVE_TRIL_DOT_PRECISION, - ) - b_Ai32 = -tl.dot( - tl.dot(b_Ai33, b_Akk32, input_precision=SOLVE_TRIL_DOT_PRECISION), - b_Ai22, - input_precision=SOLVE_TRIL_DOT_PRECISION, - ) - - b_Ai20 = -tl.dot( - b_Ai22, - tl.dot(b_Akk20, b_Ai00, input_precision=SOLVE_TRIL_DOT_PRECISION) - + tl.dot(b_Akk21, b_Ai10, input_precision=SOLVE_TRIL_DOT_PRECISION), - input_precision=SOLVE_TRIL_DOT_PRECISION, - ) - b_Ai31 = -tl.dot( - b_Ai33, - tl.dot(b_Akk31, b_Ai11, input_precision=SOLVE_TRIL_DOT_PRECISION) - + tl.dot(b_Akk32, b_Ai21, input_precision=SOLVE_TRIL_DOT_PRECISION), - input_precision=SOLVE_TRIL_DOT_PRECISION, - ) - b_Ai30 = -tl.dot( - b_Ai33, - tl.dot(b_Akk30, b_Ai00, input_precision=SOLVE_TRIL_DOT_PRECISION) - + tl.dot(b_Akk31, b_Ai10, input_precision=SOLVE_TRIL_DOT_PRECISION) - + tl.dot(b_Akk32, b_Ai20, input_precision=SOLVE_TRIL_DOT_PRECISION), - input_precision=SOLVE_TRIL_DOT_PRECISION, - ) - - p_Akk00 = tl.make_block_ptr(Akk, (T, BT), (H * BT, 1), (i_tc0, 0), (BC, BC), (1, 0)) - p_Akk10 = tl.make_block_ptr(Akk, (T, BT), (H * BT, 1), (i_tc1, 0), (BC, BC), (1, 0)) - p_Akk11 = tl.make_block_ptr(Akk, (T, BT), (H * BT, 1), (i_tc1, BC), (BC, BC), (1, 0)) - p_Akk20 = tl.make_block_ptr(Akk, (T, BT), (H * BT, 1), (i_tc2, 0), (BC, BC), (1, 0)) - p_Akk21 = tl.make_block_ptr(Akk, (T, BT), (H * BT, 1), (i_tc2, BC), (BC, BC), (1, 0)) - p_Akk22 = tl.make_block_ptr(Akk, (T, BT), (H * BT, 1), (i_tc2, 2 * BC), (BC, BC), (1, 0)) - p_Akk30 = tl.make_block_ptr(Akk, (T, BT), (H * BT, 1), (i_tc3, 0), (BC, BC), (1, 0)) - p_Akk31 = tl.make_block_ptr(Akk, (T, BT), (H * BT, 1), (i_tc3, BC), (BC, BC), (1, 0)) - p_Akk32 = tl.make_block_ptr(Akk, (T, BT), (H * BT, 1), (i_tc3, 2 * BC), (BC, BC), (1, 0)) - p_Akk33 = tl.make_block_ptr(Akk, (T, BT), (H * BT, 1), (i_tc3, 3 * BC), (BC, BC), (1, 0)) - - tl.store(p_Akk00, b_Ai00.to(Akk.dtype.element_ty), boundary_check=(0, 1)) - tl.store(p_Akk10, b_Ai10.to(Akk.dtype.element_ty), boundary_check=(0, 1)) - tl.store(p_Akk11, b_Ai11.to(Akk.dtype.element_ty), boundary_check=(0, 1)) - tl.store(p_Akk20, b_Ai20.to(Akk.dtype.element_ty), boundary_check=(0, 1)) - tl.store(p_Akk21, b_Ai21.to(Akk.dtype.element_ty), boundary_check=(0, 1)) - tl.store(p_Akk22, b_Ai22.to(Akk.dtype.element_ty), boundary_check=(0, 1)) - tl.store(p_Akk30, b_Ai30.to(Akk.dtype.element_ty), boundary_check=(0, 1)) - tl.store(p_Akk31, b_Ai31.to(Akk.dtype.element_ty), boundary_check=(0, 1)) - tl.store(p_Akk32, b_Ai32.to(Akk.dtype.element_ty), boundary_check=(0, 1)) - tl.store(p_Akk33, b_Ai33.to(Akk.dtype.element_ty), boundary_check=(0, 1)) - - -@triton.heuristics({ - 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None, -}) -@fla_cache_autotune( - configs=[ - triton.Config({}, num_warps=num_warps, num_stages=num_stages) - for num_warps in [1, 2, 4, 8] - for num_stages in [2, 3, 4] - ], - key=['BK', 'NC', 'BT'], - **autotune_cache_kwargs, -) -@triton.jit(do_not_specialize=['B', 'T']) -def chunk_gdn2_bwd_kernel_intra( - q, - k, - g, - b, - dAqk, - dAkk, - dq, - dq2, - dk, - dk2, - dg, - dg2, - db, - cu_seqlens, - chunk_indices, - B, - T, - H: tl.constexpr, - K: tl.constexpr, - BT: tl.constexpr, - BC: tl.constexpr, - BK: tl.constexpr, - NC: tl.constexpr, - IS_VARLEN: tl.constexpr, - SAFE_GATE: tl.constexpr, - USE_GATHER: tl.constexpr, -): - i_kc, i_t, i_bh = tl.program_id(0), tl.program_id(1), tl.program_id(2) - i_b, i_h = i_bh // H, i_bh % H - i_k, i_i = i_kc // NC, i_kc % NC - - all = B * T - 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.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) - else: - bos, eos = i_b * T, i_b * T + T - T = eos - bos - - i_ti = i_t * BT + i_i * BC - if i_ti >= T: - return - - o_k = i_k * BK + tl.arange(0, BK) - m_k = o_k < K - - q += (bos * H + i_h) * K - k += (bos * H + i_h) * K - g += (bos * H + i_h) * K - b += (bos * H + i_h) * K - - dAqk += (bos * H + i_h) * BT - dAkk += (bos * H + i_h) * BT - dq += (bos * H + i_h) * K - dq2 += (bos * H + i_h) * K - dk += (bos * H + i_h) * K - dk2 += (bos * H + i_h) * K - dg += (bos * H + i_h) * K - dg2 += (bos * H + i_h) * K - db += (i_k * all + bos) * H * BK + i_h * BK - - p_g = tl.make_block_ptr(g, (T, K), (H * K, 1), (i_ti, i_k * BK), (BC, BK), (1, 0)) - b_g = tl.load(p_g, boundary_check=(0, 1)).to(tl.float32) - - p_b = tl.make_block_ptr(b, (T, K), (H * K, 1), (i_ti, i_k * BK), (BC, BK), (1, 0)) - b_b = tl.load(p_b, boundary_check=(0, 1)) - - b_dq2 = tl.zeros([BC, BK], dtype=tl.float32) - b_dk2 = tl.zeros([BC, BK], dtype=tl.float32) - - if i_i > 0: - p_gn = g + i_ti * H * K + o_k - b_gn = tl.load(p_gn, mask=m_k, other=0).to(tl.float32)[None, :] - for i_j in range(0, i_i): - p_k = tl.make_block_ptr(k, (T, K), (H * K, 1), (i_t * BT + i_j * BC, i_k * BK), (BC, BK), (1, 0)) - p_gk = tl.make_block_ptr(g, (T, K), (H * K, 1), (i_t * BT + i_j * BC, i_k * BK), (BC, BK), (1, 0)) - p_dAqk = tl.make_block_ptr(dAqk, (T, BT), (H * BT, 1), (i_ti, i_j * BC), (BC, BC), (1, 0)) - p_dAkk = tl.make_block_ptr(dAkk, (T, BT), (H * BT, 1), (i_ti, i_j * BC), (BC, BC), (1, 0)) - b_kj = tl.load(p_k, boundary_check=(0, 1)) - b_gk = tl.load(p_gk, boundary_check=(0, 1)) - b_kg = b_kj * exp2(b_gn - b_gk) - b_dAqk = tl.load(p_dAqk, boundary_check=(0, 1)) - b_dAkk = tl.load(p_dAkk, boundary_check=(0, 1)) - b_dq2 += tl.dot(b_dAqk, b_kg) - b_dk2 += tl.dot(b_dAkk, b_kg) - b_gqn = exp2(b_g - b_gn) - b_dq2 *= b_gqn - b_dk2 *= b_gqn - - o_i = tl.arange(0, BC) - m_dA = (i_ti + o_i) < T - o_dA = (i_ti + o_i) * H * BT + i_i * BC - p_kj = k + i_ti * H * K + o_k - p_gkj = g + i_ti * H * K + o_k - - p_q = tl.make_block_ptr(q, (T, K), (H * K, 1), (i_ti, i_k * BK), (BC, BK), (1, 0)) - p_k = tl.make_block_ptr(k, (T, K), (H * K, 1), (i_ti, i_k * BK), (BC, BK), (1, 0)) - b_q = tl.load(p_q, boundary_check=(0, 1)) - b_k = tl.load(p_k, boundary_check=(0, 1)) - - if SAFE_GATE: - if USE_GATHER: - b_gn = gather(b_g, tl.full([1, BK], min(BC // 2, T - i_ti - 1), dtype=tl.int16), axis=0) - else: - p_gn = g + (i_ti + min(BC // 2, T - i_ti - 1)) * H * K + o_k - b_gn = tl.load(p_gn, mask=m_k, other=0)[None, :] - - p_dAqk = tl.make_block_ptr(dAqk, (T, BT), (H * BT, 1), (i_ti, i_i * BC), (BC, BC), (1, 0)) - p_dAkk = tl.make_block_ptr(dAkk, (T, BT), (H * BT, 1), (i_ti, i_i * BC), (BC, BC), (1, 0)) - b_dAqk_diag_qk = tl.load(p_dAqk, boundary_check=(0, 1)).to(tl.float32) - b_dAkk_diag_qk = tl.load(p_dAkk, boundary_check=(0, 1)).to(tl.float32) - - m_i_diag_qk = (o_i[:, None] >= o_i[None, :]) & ((i_ti + o_i[:, None]) < T) & ((i_ti + o_i[None, :]) < T) - m_j_diag_qk = (i_ti + o_i[:, None]) < T - - b_dAqk_diag_qk = tl.where(m_i_diag_qk, b_dAqk_diag_qk, 0.) - b_dAkk_diag_qk = tl.where(m_i_diag_qk, b_dAkk_diag_qk, 0.) - b_g_diag_qk = tl.where(m_j_diag_qk, b_g - b_gn, 0.) - exp_b_g_diag_qk = tl.where(m_j_diag_qk, exp2(b_g_diag_qk), 0.) - exp_neg_b_g_diag_qk = tl.where(m_j_diag_qk, exp2(-b_g_diag_qk), 0.) - - b_k_exp_diag_qk = b_k * exp_neg_b_g_diag_qk - b_dq2 += tl.dot(b_dAqk_diag_qk, b_k_exp_diag_qk) * exp_b_g_diag_qk - b_dk2 += tl.dot(b_dAkk_diag_qk, b_k_exp_diag_qk) * exp_b_g_diag_qk - else: - for j in range(0, min(BC, T - i_t * BT - i_i * BC)): - b_dAqk = tl.load(dAqk + o_dA + j, mask=m_dA, other=0) - b_dAkk = tl.load(dAkk + o_dA + j, mask=m_dA, other=0) - b_kj = tl.load(p_kj, mask=m_k, other=0).to(tl.float32) - b_gkj = tl.load(p_gkj, mask=m_k, other=0).to(tl.float32) - m_i = o_i[:, None] >= j - b_gqk = exp2(b_g - b_gkj[None, :]) - b_dq2 += tl.where(m_i, b_dAqk[:, None] * b_kj[None, :] * b_gqk, 0.) - b_dk2 += tl.where(m_i, b_dAkk[:, None] * b_kj[None, :] * b_gqk, 0.) - - p_kj += H * K - p_gkj += H * K - - # db gets the elementwise b * k contribution from dk2 (before b is folded - # back into dk2 itself); then dk2 is rescaled by the channel-wise gate. - b_db_tile = b_dk2 * b_k - b_dk2 = b_dk2 * b_b - - p_dq = tl.make_block_ptr(dq, (T, K), (H * K, 1), (i_ti, i_k * BK), (BC, BK), (1, 0)) - p_dq2 = tl.make_block_ptr(dq2, (T, K), (H * K, 1), (i_ti, i_k * BK), (BC, BK), (1, 0)) - p_db = tl.make_block_ptr(db, (T, BK), (H * BK, 1), (i_ti, 0), (BC, BK), (1, 0)) - - b_dg2 = b_q * b_dq2 - b_dq2 = b_dq2 + tl.load(p_dq, boundary_check=(0, 1)) - tl.store(p_dq2, b_dq2.to(p_dq2.dtype.element_ty), boundary_check=(0, 1)) - tl.store(p_db, b_db_tile.to(p_db.dtype.element_ty), boundary_check=(0, 1)) - - tl.debug_barrier() - b_dkt = tl.zeros([BC, BK], dtype=tl.float32) - - NC = min(NC, tl.cdiv(T - i_t * BT, BC)) - if i_i < NC - 1: - p_gn = g + (min(i_ti + BC, T) - 1) * H * K + o_k - b_gn = tl.load(p_gn, mask=m_k, other=0).to(tl.float32)[None, :] - for i_j in range(i_i + 1, NC): - p_q = tl.make_block_ptr(q, (T, K), (H * K, 1), (i_t * BT + i_j * BC, i_k * BK), (BC, BK), (1, 0)) - p_k = tl.make_block_ptr(k, (T, K), (H * K, 1), (i_t * BT + i_j * BC, i_k * BK), (BC, BK), (1, 0)) - p_gk = tl.make_block_ptr(g, (T, K), (H * K, 1), (i_t * BT + i_j * BC, i_k * BK), (BC, BK), (1, 0)) - p_bj = tl.make_block_ptr(b, (T, K), (H * K, 1), (i_t * BT + i_j * BC, i_k * BK), (BC, BK), (1, 0)) - p_dAqk = tl.make_block_ptr(dAqk, (BT, T), (1, H * BT), (i_i * BC, i_t * BT + i_j * BC), (BC, BC), (0, 1)) - p_dAkk = tl.make_block_ptr(dAkk, (BT, T), (1, H * BT), (i_i * BC, i_t * BT + i_j * BC), (BC, BC), (0, 1)) - b_bj = tl.load(p_bj, boundary_check=(0, 1)) - b_qj = tl.load(p_q, boundary_check=(0, 1)) - b_kbj = tl.load(p_k, boundary_check=(0, 1)) * b_bj - b_gk = tl.load(p_gk, boundary_check=(0, 1)).to(tl.float32) - b_dAqk = tl.load(p_dAqk, boundary_check=(0, 1)) - b_dAkk = tl.load(p_dAkk, boundary_check=(0, 1)) - - o_j = i_t * BT + i_j * BC + o_i - m_j = o_j < T - b_gkn = exp2(b_gk - b_gn) - b_qg = b_qj * tl.where(m_j[:, None], b_gkn, 0) - b_kbg = b_kbj * tl.where(m_j[:, None], b_gkn, 0) - b_dkt += tl.dot(b_dAqk, b_qg) - b_dkt += tl.dot(b_dAkk, b_kbg) - b_dkt *= exp2(b_gn - b_g) - - o_dA = i_ti * H * BT + i_i * BC + o_i - p_qj = q + i_ti * H * K + o_k - p_kj = k + i_ti * H * K + o_k - p_gkj = g + i_ti * H * K + o_k - p_bj_ptr = b + i_ti * H * K + o_k - - if SAFE_GATE: - if USE_GATHER: - b_gn = gather(b_g, tl.full([1, BK], min(BC // 2, T - i_ti - 1), dtype=tl.int16), axis=0) - else: - p_gn = g + (i_ti + min(BC // 2, T - i_ti - 1)) * H * K + o_k - b_gn = tl.load(p_gn, mask=m_k, other=0).to(tl.float32)[None, :] - p_q = tl.make_block_ptr(q, (T, K), (H * K, 1), (i_ti, i_k * BK), (BC, BK), (1, 0)) - b_q = tl.load(p_q, boundary_check=(0, 1)) - p_dAqk = tl.make_block_ptr(dAqk, (BT, T), (1, H * BT), (i_i * BC, i_ti), (BC, BC), (0, 1)) - p_dAkk = tl.make_block_ptr(dAkk, (BT, T), (1, H * BT), (i_i * BC, i_ti), (BC, BC), (0, 1)) - b_dAqk_diag_kk = tl.load(p_dAqk, boundary_check=(0, 1)).to(tl.float32) - b_dAkk_diag_kk = tl.load(p_dAkk, boundary_check=(0, 1)).to(tl.float32) - - m_i_diag_kk = (o_i[:, None] <= o_i[None, :]) & ((i_ti + o_i[:, None]) < T) & ((i_ti + o_i[None, :]) < T) - m_j_diag_kk = (i_ti + o_i[:, None]) < T - - b_dAqk_diag_kk = tl.where(m_i_diag_kk, b_dAqk_diag_kk, 0.) - b_dAkk_diag_kk = tl.where(m_i_diag_kk, b_dAkk_diag_kk, 0.) - b_g_diag_kk = tl.where(m_j_diag_kk, b_g - b_gn, 0.) - exp_b_g_diag_kk = tl.where(m_j_diag_kk, exp2(b_g_diag_kk), 0.) - exp_neg_b_g_diag_kk = tl.where(m_j_diag_kk, exp2(-b_g_diag_kk), 0.) - - b_q_exp = b_q * exp_b_g_diag_kk - b_kb_exp = b_k * b_b * exp_b_g_diag_kk - - b_dkt += tl.dot(b_dAqk_diag_kk, b_q_exp) * exp_neg_b_g_diag_kk - b_dkt += tl.dot(b_dAkk_diag_kk, b_kb_exp) * exp_neg_b_g_diag_kk - else: - for j in range(0, min(BC, T - i_t * BT - i_i * BC)): - b_dAqk = tl.load(dAqk + o_dA + j * H * BT) - b_dAkk = tl.load(dAkk + o_dA + j * H * BT) - b_qj = tl.load(p_qj, mask=m_k, other=0).to(tl.float32) - b_kj = tl.load(p_kj, mask=m_k, other=0).to(tl.float32) - b_bj = tl.load(p_bj_ptr, mask=m_k, other=0).to(tl.float32) - b_kbj = b_kj * b_bj - b_gkj = tl.load(p_gkj, mask=m_k, other=0).to(tl.float32) - m_i = o_i[:, None] <= j - b_gkq = exp2(b_gkj[None, :] - b_g) - b_dkt += tl.where(m_i, b_dAqk[:, None] * b_qj[None, :] * b_gkq, 0.) - b_dkt += tl.where(m_i, b_dAkk[:, None] * b_kbj[None, :] * b_gkq, 0.) - - p_qj += H * K - p_kj += H * K - p_gkj += H * K - p_bj_ptr += H * K - - p_dk = tl.make_block_ptr(dk, (T, K), (H * K, 1), (i_ti, i_k * BK), (BC, BK), (1, 0)) - p_dk2 = tl.make_block_ptr(dk2, (T, K), (H * K, 1), (i_ti, i_k * BK), (BC, BK), (1, 0)) - p_dg = tl.make_block_ptr(dg, (T, K), (H * K, 1), (i_ti, i_k * BK), (BC, BK), (1, 0)) - p_dg2 = tl.make_block_ptr(dg2, (T, K), (H * K, 1), (i_ti, i_k * BK), (BC, BK), (1, 0)) - - b_dg2 += (b_dk2 - b_dkt) * b_k + tl.load(p_dg, boundary_check=(0, 1)) - b_dk2 += tl.load(p_dk, boundary_check=(0, 1)) - b_dk2 += b_dkt - - tl.store(p_dk2, b_dk2.to(p_dk2.dtype.element_ty), boundary_check=(0, 1)) - tl.store(p_dg2, b_dg2.to(p_dg2.dtype.element_ty), boundary_check=(0, 1)) - - -def chunk_gdn2_fwd_intra( - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - gk: torch.Tensor, - b: torch.Tensor, - w_gate: torch.Tensor, - scale: float, - cu_seqlens: torch.LongTensor | None = None, - chunk_size: int = 64, - chunk_indices: torch.LongTensor | None = None, - safe_gate: bool = False, - disable_recompute: bool = False, -): - """Intra-chunk forward: build Aqk, Akk_inv, and the WY auxiliaries (w, u).""" - B, T, H, K = k.shape - BT = chunk_size - BC = 16 - if chunk_indices is None and cu_seqlens is not None: - chunk_indices = prepare_chunk_indices(cu_seqlens, BT) - NT = triton.cdiv(T, BT) if cu_seqlens is None else len(chunk_indices) - - Aqk = torch.empty(B, T, H, BT, device=k.device, dtype=k.dtype) - # Akk must be zero-initialized - kernel only writes lower triangular. - Akk = torch.zeros(B, T, H, BT, device=k.device, dtype=k.dtype) - Akkd = torch.empty(B, T, H, BC, device=k.device, dtype=torch.float32) - - if safe_gate: - grid = (NT, triton.cdiv(BT, BC), B * H) - BK = triton.next_power_of_2(K) - chunk_gdn2_fwd_kernel_intra_sub_chunk[grid]( - q=q, - k=k, - g=gk, - b=b, - Aqk=Aqk, - Akk=Akkd, - scale=scale, - cu_seqlens=cu_seqlens, - chunk_indices=chunk_indices, - T=T, - H=H, - K=K, - BT=BT, - BC=BC, - BK=BK, - USE_GATHER=IS_GATHER_SUPPORTED, - ) - else: - chunk_gdn2_fwd_intra_token_parallel( - q=q, - k=k, - gk=gk, - b=b, - Aqk=Aqk, - Akk=Akkd, - scale=scale, - cu_seqlens=cu_seqlens, - chunk_size=BT, - sub_chunk_size=BC, - ) - - # Fused inter-subchunk Akk + solve_tril → full Akk_inv. - grid = (NT, B * H) - chunk_gdn2_fwd_kernel_inter_solve_fused[grid]( - q=q, - k=k, - g=gk, - b=b, - Aqk=Aqk, - Akkd=Akkd, - Akk=Akk, - scale=scale, - cu_seqlens=cu_seqlens, - chunk_indices=chunk_indices, - T=T, - H=H, - K=K, - BT=BT, - BC=BC, - USE_SAFE_GATE=safe_gate, - ) - - w, u, qg, kg = recompute_w_u_fwd_gdn2( - k=k, - v=v, - b=b, - w_gate=w_gate, - A=Akk, - q=q if disable_recompute else None, - gk=gk, - cu_seqlens=cu_seqlens, - chunk_indices=chunk_indices, - ) - return w, u, qg, kg, Aqk, Akk - - -def chunk_gdn2_bwd_intra( - q: torch.Tensor, - k: torch.Tensor, - g: torch.Tensor, - b: torch.Tensor, - dAqk: torch.Tensor, - dAkk: torch.Tensor, - dq: torch.Tensor, - dk: torch.Tensor, - db: torch.Tensor, - dg: torch.Tensor, - cu_seqlens: torch.LongTensor | None = None, - chunk_indices: torch.LongTensor | None = None, - chunk_size: int = 64, - safe_gate: bool = False, -): - """Intra-chunk backward: q, k, g, b contributions from dAqk, dAkk. - - Adds the intra-chunk piece of db (channel-wise, K-dim) to the running ``db`` - accumulator passed in. - """ - B, T, H, K = k.shape - BT = chunk_size - BC = min(16, BT) - BK = min(32, triton.next_power_of_2(K)) - - if chunk_indices is None and cu_seqlens is not None: - chunk_indices = prepare_chunk_indices(cu_seqlens, BT) - NT = triton.cdiv(T, BT) if cu_seqlens is None else len(chunk_indices) - NC = triton.cdiv(BT, BC) - NK = triton.cdiv(K, BK) - - dq2 = torch.empty_like(q) - dk2 = torch.empty_like(k) - db2 = q.new_empty(NK, B, T, H, BK, dtype=torch.float32) - dg2 = torch.empty_like(dg, dtype=torch.float32) - - grid = (NK * NC, NT, B * H) - chunk_gdn2_bwd_kernel_intra[grid]( - q=q, - k=k, - g=g, - b=b, - dAqk=dAqk, - dAkk=dAkk, - dq=dq, - dq2=dq2, - dk=dk, - dk2=dk2, - dg=dg, - dg2=dg2, - db=db2, - cu_seqlens=cu_seqlens, - chunk_indices=chunk_indices, - B=B, - T=T, - H=H, - K=K, - BT=BT, - BC=BC, - BK=BK, - NC=NC, - SAFE_GATE=safe_gate, - USE_GATHER=IS_GATHER_SUPPORTED, - ) - dq = dq2 - dk = dk2 - db2_combined = db2.permute(1, 2, 3, 0, 4).contiguous().reshape(B, T, H, NK * BK)[..., :K] - db = db.add_(db2_combined) - dg = dg2 - return dq, dk, db, dg diff --git a/build/torch-xpu/ops/gdn2/chunk_intra_token_parallel.py b/build/torch-xpu/ops/gdn2/chunk_intra_token_parallel.py deleted file mode 100644 index 90e091618d3c2efa997f587614c084942efdccb7..0000000000000000000000000000000000000000 --- a/build/torch-xpu/ops/gdn2/chunk_intra_token_parallel.py +++ /dev/null @@ -1,172 +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 - -# Token-parallel intra-chunk kernel for GDN-2. -# -# Builds the diagonal blocks of the within-chunk score matrices used by the WY -# representation: -# Aqk - causal query-key scores, decay-weighted, for the output path. -# Akk - gated key-key scores (the strictly-lower matrix T whose inverse -# (I + T)^{-1} defines the WY representation). -# -# Difference vs KDA: the channel-wise erase gate ``b`` is folded into the key -# tile before the dot product, instead of multiplying the row by a scalar -# ``beta`` after the dot. The rest is identical. - -import torch -import triton -import triton.language as tl - -from ...ops.utils.cache import fla_cache_autotune -from ...ops.utils.op import exp2 -from ...utils import autotune_cache_kwargs - - -@triton.heuristics({ - 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None, -}) -@fla_cache_autotune( - configs=[ - triton.Config({'BH': BH}, num_warps=num_warps) - for BH in [1, 2, 4, 8] - for num_warps in [1, 2, 4, 8] - ], - key=["K", "H"], - **autotune_cache_kwargs, -) -@triton.jit(do_not_specialize=['T', 'N']) -def chunk_gdn2_fwd_kernel_intra_token_parallel( - q, - k, - g, - b, - Aqk, - Akk, - scale, - cu_seqlens, - N, - T, - H: tl.constexpr, - K: tl.constexpr, - BT: tl.constexpr, - BC: tl.constexpr, - BH: tl.constexpr, - IS_VARLEN: tl.constexpr, -): - i_tg, i_hg = tl.program_id(0), tl.program_id(1) - - if IS_VARLEN: - i_n = 0 - left, right = 0, N - for _ in range(20): - if left < right: - mid = (left + right) // 2 - if i_tg < tl.load(cu_seqlens + mid + 1).to(tl.int32): - right = mid - else: - left = mid + 1 - i_n = left - bos, eos = tl.load(cu_seqlens + i_n).to(tl.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) - T = eos - bos - i_t = i_tg - bos - else: - bos = (i_tg // T) * T - i_t = i_tg % T - - if i_t >= T: - return - - i_c = i_t // BT - i_s = (i_t % BT) // BC - i_tc = i_c * BT - i_ts = i_tc + i_s * BC - - q += bos * H * K - k += bos * H * K - g += bos * H * K - b += bos * H * K - Aqk += bos * H * BT - Akk += bos * H * BC - - BK: tl.constexpr = triton.next_power_of_2(K) - o_h = tl.arange(0, BH) - o_k = tl.arange(0, BK) - m_h = (i_hg * BH + o_h) < H - m_k = o_k < K - - p_q = tl.make_block_ptr(q + i_t * H * K, (H, K), (K, 1), (i_hg * BH, 0), (BH, BK), (1, 0)) - p_k = tl.make_block_ptr(k + i_t * H * K, (H, K), (K, 1), (i_hg * BH, 0), (BH, BK), (1, 0)) - p_g = tl.make_block_ptr(g + i_t * H * K, (H, K), (K, 1), (i_hg * BH, 0), (BH, BK), (1, 0)) - p_b = tl.make_block_ptr(b + i_t * H * K, (H, K), (K, 1), (i_hg * BH, 0), (BH, BK), (1, 0)) - - b_q = tl.load(p_q, boundary_check=(0, 1)).to(tl.float32) - b_k = tl.load(p_k, boundary_check=(0, 1)).to(tl.float32) - b_g = tl.load(p_g, boundary_check=(0, 1)).to(tl.float32) - b_b = tl.load(p_b, boundary_check=(0, 1)).to(tl.float32) - - # Fold channel-wise erase gate into the key tile. - b_k = b_k * b_b - - for j in range(i_ts, min(i_t + 1, min(T, i_ts + BC))): - p_kj = tl.make_block_ptr(k + j * H * K, (H, K), (K, 1), (i_hg * BH, 0), (BH, BK), (1, 0)) - p_gj = tl.make_block_ptr(g + j * H * K, (H, K), (K, 1), (i_hg * BH, 0), (BH, BK), (1, 0)) - b_kj = tl.load(p_kj, boundary_check=(0, 1)).to(tl.float32) - b_gj = tl.load(p_gj, boundary_check=(0, 1)).to(tl.float32) - - b_kgj = b_kj * exp2(b_g - b_gj) - b_kgj = tl.where(m_k[None, :], b_kgj, 0.0) - b_Aqk = tl.sum(b_q * b_kgj, axis=1) * scale - b_Akk = tl.sum(b_k * b_kgj, axis=1) * tl.where(j < i_t, 1.0, 0.0) - - tl.store( - Aqk + i_t * H * BT + (i_hg * BH + o_h) * BT + j % BT, - b_Aqk.to(Aqk.dtype.element_ty), mask=m_h, - ) - tl.store( - Akk + i_t * H * BC + (i_hg * BH + o_h) * BC + j - i_ts, - b_Akk.to(Akk.dtype.element_ty), mask=m_h, - ) - - -def chunk_gdn2_fwd_intra_token_parallel( - q: torch.Tensor, - k: torch.Tensor, - gk: torch.Tensor, - b: torch.Tensor, - Aqk: torch.Tensor, - Akk: torch.Tensor, - scale: float, - cu_seqlens: torch.LongTensor | None = None, - chunk_size: int = 64, - sub_chunk_size: int = 16, -): - """Token-parallel diagonal kernel: builds Aqk plus diagonal Akk in fp32.""" - B, T, H, K = q.shape - N = len(cu_seqlens) - 1 if cu_seqlens is not None else B - BT = chunk_size - BC = sub_chunk_size - - def grid(meta): - return (B * T, triton.cdiv(H, meta['BH'])) - - chunk_gdn2_fwd_kernel_intra_token_parallel[grid]( - q=q, - k=k, - g=gk, - b=b, - Aqk=Aqk, - Akk=Akk, - scale=scale, - cu_seqlens=cu_seqlens, - N=N, - T=T, - H=H, - K=K, - BT=BT, - BC=BC, - ) - return Aqk, Akk diff --git a/build/torch-xpu/ops/gdn2/fused_recurrent.py b/build/torch-xpu/ops/gdn2/fused_recurrent.py deleted file mode 100644 index 3e7b2cd267feeb9b123332b7ed83252b0edfbcc8..0000000000000000000000000000000000000000 --- a/build/torch-xpu/ops/gdn2/fused_recurrent.py +++ /dev/null @@ -1,468 +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 - -r""" -Token-by-token recurrent kernel for GDN-2 (Gated DeltaNet 2). - -This is the inference-time counterpart of the chunkwise training kernels. It -runs the GDN-2 recurrence one token at a time, with no chunk padding, which -makes it the preferred path for autoregressive decoding at short sequence -lengths. It is forward-only and does not track gradients; training always uses -the chunkwise kernel. - -Per token, the matrix state ``S`` in ``R^{K x V}`` is updated by - - S <- Diag(exp(g)) * S # channel-wise decay - v_new = (w * v) - (b * k)^T S # gated write minus gated read - S <- S + k (v_new)^T # rank-one write - o = S^T q # output read - -where ``*`` is the elementwise product, ``b`` is the channel-wise erase gate on -the key axis, ``w`` is the channel-wise write gate on the value axis, and -``exp(g)`` is the channel-wise decay. Setting ``b = w = beta`` (scalar) recovers -KDA. - -The kernel mirrors the KDA recurrent kernel's varlen / continuous-batching / -speculative-decoding feature set so it slots into the same serving paths. - -Public entry point: ``fused_recurrent_gdn2``. -""" - -from __future__ import annotations - -import warnings - -import torch -import triton -import triton.language as tl - -from ...ops.utils.op import exp -from ...ops.utils.softplus import softplus -from ...utils import input_guard - - -@triton.heuristics( - { - "USE_INITIAL_STATE": lambda args: args["h0"] is not None, - "STORE_FINAL_STATE": lambda args: args["ht"] is not None, - "IS_VARLEN": lambda args: args["cu_seqlens"] is not None, - "IS_CONTINUOUS_BATCHING": lambda args: args["ssm_state_indices"] is not None, - "IS_SPEC_DECODING": lambda args: args["num_accepted_tokens"] is not None, - "HAS_DT_BIAS": lambda args: args["dt_bias"] is not None, - "USE_LOWER_BOUND": lambda args: args["lower_bound"] is not None, - } -) -@triton.jit(do_not_specialize=["N", "T"]) -def fused_recurrent_gdn2_fwd_kernel( - q, - k, - v, - g, - b, - w, - A_log, - dt_bias, - o, - h0, - ht, - cu_seqlens, - ssm_state_indices, - num_accepted_tokens, - lower_bound, - scale: tl.constexpr, - N: tl.int64, - T: tl.int64, - H: tl.constexpr, - HV: tl.constexpr, - K: tl.constexpr, - V: tl.constexpr, - BK: tl.constexpr, - BV: tl.constexpr, - stride_init_state_token: tl.constexpr, - stride_final_state_token: tl.constexpr, - stride_indices_seq: tl.constexpr, - stride_indices_tok: tl.constexpr, - USE_INITIAL_STATE: tl.constexpr, - INPLACE_FINAL_STATE: tl.constexpr, - USE_QK_L2NORM_IN_KERNEL: tl.constexpr, - IS_VARLEN: tl.constexpr, - IS_CONTINUOUS_BATCHING: tl.constexpr, - IS_SPEC_DECODING: tl.constexpr, - STORE_FINAL_STATE: tl.constexpr, - HAS_DT_BIAS: tl.constexpr, - USE_GATE_IN_KERNEL: tl.constexpr, - USE_LOWER_BOUND: tl.constexpr, - STATE_V_FIRST: tl.constexpr, - num_stages: tl.constexpr, -): - pid = tl.program_id(0) - NV = tl.cdiv(V, BV) - NK = tl.cdiv(K, BK) - i_k = pid % NK - pid_rest = pid // NK - - i_v = pid_rest % NV - i_nh = pid_rest // NV - i_n, i_hv = i_nh // HV, i_nh % HV - i_h = i_hv // (HV // H) - if IS_VARLEN: - bos, eos = ( - tl.load(cu_seqlens + i_n).to(tl.int64), - tl.load(cu_seqlens + i_n + 1).to(tl.int64), - ) - T = eos - bos - else: - bos, eos = i_n * T, i_n * T + T - - if T == 0: - return - - o_k = i_k * BK + tl.arange(0, BK) - o_v = i_v * BV + tl.arange(0, BV) - - p_q = q + (bos * H + i_h) * K + o_k - p_k = k + (bos * H + i_h) * K + o_k - p_v = v + (bos * HV + i_hv) * V + o_v - p_b = b + (bos * HV + i_hv) * K + o_k - p_w = w + (bos * HV + i_hv) * V + o_v - p_g = g + (bos * HV + i_hv) * K + o_k - p_o = o + (bos * HV + i_hv) * V + o_v - - mask_k = o_k < K - mask_v = o_v < V - if STATE_V_FIRST: - mask_h = mask_v[:, None] & mask_k[None, :] - else: - mask_h = mask_k[:, None] & mask_v[None, :] - - if STATE_V_FIRST: - b_h = tl.zeros([BV, BK], dtype=tl.float32) - else: - b_h = tl.zeros([BK, BV], dtype=tl.float32) - if USE_INITIAL_STATE: - if IS_CONTINUOUS_BATCHING: - if IS_SPEC_DECODING: - i_t = tl.load(num_accepted_tokens + i_n).to(tl.int64) - 1 - else: - i_t = 0 - p_h0 = ( - h0 - + tl.load(ssm_state_indices + i_n * stride_indices_seq + i_t).to(tl.int64) - * stride_init_state_token - ) - if STATE_V_FIRST: - p_h0 = p_h0 + i_hv * K * V + o_v[:, None] * K + o_k[None, :] - else: - p_h0 = p_h0 + i_hv * K * V + o_k[:, None] * V + o_v[None, :] - else: - if STATE_V_FIRST: - p_h0 = h0 + (i_n * HV + i_hv) * K * V + o_v[:, None] * K + o_k[None, :] - else: - p_h0 = h0 + (i_n * HV + i_hv) * K * V + o_k[:, None] * V + o_v[None, :] - b_h += tl.load(p_h0, mask=mask_h, other=0).to(tl.float32) - - for i_t in tl.range(0, T, num_stages=num_stages): - b_q = tl.load(p_q, mask=mask_k, other=0, eviction_policy='evict_last').to(tl.float32) - b_k = tl.load(p_k, mask=mask_k, other=0, eviction_policy='evict_last').to(tl.float32) - b_v = tl.load(p_v, mask=mask_v, other=0, eviction_policy='evict_first').to(tl.float32) - - if USE_QK_L2NORM_IN_KERNEL: - b_q = b_q / tl.sqrt(tl.sum(b_q * b_q) + 1e-6) - b_k = b_k / tl.sqrt(tl.sum(b_k * b_k) + 1e-6) - b_q = b_q * scale - b_g = tl.load(p_g, eviction_policy='evict_last').to(tl.float32) - - if USE_GATE_IN_KERNEL: - b_A = tl.load(A_log + i_hv).to(tl.float32) - - if HAS_DT_BIAS: - b_bias = tl.load(dt_bias + i_hv * K + o_k, mask=mask_k, other=0).to(tl.float32) - b_g = b_g + b_bias - - if USE_LOWER_BOUND: - b_gk = lower_bound * tl.sigmoid(exp(b_A) * b_g) - else: - b_gk = -exp(b_A) * softplus(b_g) - else: - b_gk = b_g - - # Per-channel decay on the running state. - if STATE_V_FIRST: - b_h *= exp(b_gk[None, :]) - else: - b_h *= exp(b_gk[:, None]) - - b_b = tl.load(p_b, mask=mask_k, other=0, eviction_policy='evict_last').to(tl.float32) - b_bk = b_b * b_k - - # erase contribution along V: (b * k)^T S - if STATE_V_FIRST: - erase_d = tl.sum(b_h * b_bk[None, :], 1) - else: - erase_d = tl.sum(b_h * b_bk[:, None], 0) - - b_w = tl.load(p_w, mask=mask_v, other=0, eviction_policy='evict_first').to(tl.float32) - b_v_new = b_w * b_v - erase_d - - if STATE_V_FIRST: - b_h += b_v_new[:, None] * b_k[None, :] - b_o = tl.sum(b_h * b_q[None, :], 1) - else: - b_h += b_k[:, None] * b_v_new[None, :] - b_o = tl.sum(b_h * b_q[:, None], 0) - tl.store(p_o, b_o.to(p_o.dtype.element_ty), mask=mask_v, eviction_policy='evict_first') - - if IS_CONTINUOUS_BATCHING: - if INPLACE_FINAL_STATE: - p_ht = ( - ht - + tl.load(ssm_state_indices + i_n * stride_indices_seq + i_t).to(tl.int64) - * stride_final_state_token - ) - else: - p_ht = ht + (bos + i_t) * stride_final_state_token - if STATE_V_FIRST: - p_ht = p_ht + i_hv * K * V + o_v[:, None] * K + o_k[None, :] - else: - p_ht = p_ht + i_hv * K * V + o_k[:, None] * V + o_v[None, :] - tl.store(p_ht, b_h.to(p_ht.dtype.element_ty), mask=mask_h) - - p_q += H * K - p_k += H * K - p_o += HV * V - p_v += HV * V - p_g += HV * K - p_b += HV * K - p_w += HV * V - - if not IS_CONTINUOUS_BATCHING: - if STORE_FINAL_STATE: - if STATE_V_FIRST: - p_ht = ht + (i_n * HV + i_hv) * K * V + o_v[:, None] * K + o_k[None, :] - else: - p_ht = ht + (i_n * HV + i_hv) * K * V + o_k[:, None] * V + o_v[None, :] - tl.store(p_ht, b_h.to(p_ht.dtype.element_ty), mask=mask_h) - - -@torch.compiler.disable -def fused_recurrent_gdn2_fwd( - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - g: torch.Tensor, - b: torch.Tensor, - w: torch.Tensor, - A_log: torch.Tensor | None = None, - dt_bias: torch.Tensor | None = None, - initial_state: torch.Tensor | None = None, - scale: float | None = None, - output_final_state: bool = False, - inplace_final_state: bool = False, - state_v_first: bool = False, - cu_seqlens: torch.LongTensor | None = None, - ssm_state_indices: torch.Tensor | None = None, - num_accepted_tokens: torch.Tensor | None = None, - use_qk_l2norm_in_kernel: bool = False, - use_gate_in_kernel: bool = False, - lower_bound: float | None = None, - out: torch.Tensor | None = None, - **kwargs, -) -> tuple[torch.Tensor, torch.Tensor]: - if scale is None: - scale = k.shape[-1] ** -0.5 - - B, T, H, K, V = *k.shape, v.shape[-1] - HV = v.shape[2] - N = B if cu_seqlens is None else len(cu_seqlens) - 1 - BK = triton.next_power_of_2(K) - BV = 32 - - if out is None: - out = torch.zeros_like(v) - else: - assert out.shape == v.shape - if inplace_final_state: - assert initial_state is not None, "inplace_final_state=True requires an initial_state" - final_state = initial_state - elif output_final_state: - if state_v_first: - final_state = q.new_empty(N, HV, V, K, dtype=torch.float32) - else: - final_state = q.new_empty(N, HV, K, V, dtype=torch.float32) - else: - final_state = None - - stride_init_state_token = initial_state.stride(0) if initial_state is not None else 1 - stride_final_state_token = final_state.stride(0) if final_state is not None else 1 - - if ssm_state_indices is None: - stride_indices_seq, stride_indices_tok = 1, 1 - elif ssm_state_indices.ndim == 1: - stride_indices_seq, stride_indices_tok = ssm_state_indices.stride(0), 1 - else: - stride_indices_seq, stride_indices_tok = ssm_state_indices.stride() - - grid = (triton.cdiv(V, BV) * N * HV, ) - fused_recurrent_gdn2_fwd_kernel[grid]( - q=q, - k=k, - v=v, - g=g, - b=b, - w=w, - A_log=A_log, - dt_bias=dt_bias, - o=out, - h0=initial_state, - ht=final_state, - cu_seqlens=cu_seqlens, - ssm_state_indices=ssm_state_indices, - num_accepted_tokens=num_accepted_tokens, - lower_bound=lower_bound, - scale=scale, - N=N, - T=T, - H=H, - HV=HV, - K=K, - V=V, - BK=BK, - BV=BV, - stride_init_state_token=stride_init_state_token, - stride_final_state_token=stride_final_state_token, - stride_indices_seq=stride_indices_seq, - stride_indices_tok=stride_indices_tok, - USE_QK_L2NORM_IN_KERNEL=use_qk_l2norm_in_kernel, - INPLACE_FINAL_STATE=inplace_final_state, - USE_GATE_IN_KERNEL=use_gate_in_kernel, - STATE_V_FIRST=state_v_first, - num_warps=4, - num_stages=2, - ) - - return out, final_state - - -@input_guard -def fused_recurrent_gdn2( - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - g: torch.Tensor, - b: torch.Tensor, - w: torch.Tensor, - A_log: torch.Tensor | None = None, - dt_bias: torch.Tensor | None = None, - scale: float | None = None, - initial_state: torch.Tensor | None = None, - output_final_state: bool = False, - use_qk_l2norm_in_kernel: bool = False, - use_gate_in_kernel: bool = False, - lower_bound: float | None = None, - state_v_first: bool = False, - cu_seqlens: torch.LongTensor | None = None, - **kwargs, -) -> tuple[torch.Tensor, torch.Tensor]: - r""" - Token-by-token forward for GDN-2 (Gated DeltaNet 2). Inference-only. - - The recurrence on the matrix state ``S in R^{K x V}`` is - - S <- (I - k (b * k)^T) Diag(exp(g)) S + k (w * v)^T - - Setting ``b = w = beta`` (scalar) recovers KDA exactly. - - Args: - q (torch.Tensor): queries of shape ``[B, T, H, K]``. - k (torch.Tensor): keys of shape ``[B, T, H, K]``. - v (torch.Tensor): values of shape ``[B, T, HV, V]`` (GVA if HV > H). - g (torch.Tensor): channel-wise log-decay of shape ``[B, T, HV, K]``. - When ``use_gate_in_kernel=True``, this is the raw pre-activation and - the kernel computes ``-exp(A_log) * softplus(g + dt_bias)`` (or the - bounded variant if ``lower_bound`` is set). - b (torch.Tensor): channel-wise erase gate of shape ``[B, T, HV, K]``. - w (torch.Tensor): channel-wise write gate of shape ``[B, T, HV, V]``. - A_log (Optional[torch.Tensor]): decay parameter of shape ``[HV]``. - Required when ``use_gate_in_kernel=True``. - dt_bias (Optional[torch.Tensor]): per-channel bias of shape ``[HV * K]``, - only used when ``use_gate_in_kernel=True``. - scale (Optional[float]): attention scale, defaults to ``1/sqrt(K)``. - initial_state (Optional[torch.Tensor]): ``[N, HV, K, V]`` in float32 - (or ``[N, HV, V, K]`` if ``state_v_first=True``). - output_final_state (bool): whether to output the final state. - use_qk_l2norm_in_kernel (bool): L2-normalize q and k inside the kernel. - use_gate_in_kernel (bool): compute the gate activation from raw ``g``. - lower_bound (Optional[float]): when ``use_gate_in_kernel=True`` and set, - use ``lower_bound * sigmoid(exp(A_log) * g)`` instead of the softplus - activation. - state_v_first (bool): store state as ``[V, K]`` instead of the default - ``[K, V]``. - cu_seqlens (Optional[torch.LongTensor]): ``[N+1]`` packed-sequence offsets. - - Returns: - o (torch.Tensor): outputs of shape ``[B, T, HV, V]``. - final_state (Optional[torch.Tensor]): ``[N, HV, K, V]`` if requested. - """ - if 'transpose_state_layout' in kwargs: - if state_v_first: - raise ValueError("Cannot pass both `state_v_first` and the deprecated `transpose_state_layout`.") - warnings.warn( - "`transpose_state_layout` is deprecated and renamed to `state_v_first`.", - DeprecationWarning, - stacklevel=2, - ) - state_v_first = kwargs.pop('transpose_state_layout') - - if cu_seqlens is not None: - if q.shape[0] != 1: - raise ValueError( - f"The batch size is expected to be 1 rather than {q.shape[0]} when using `cu_seqlens`. " - f"Please flatten variable-length inputs before processing." - ) - if initial_state is not None and initial_state.shape[0] != len(cu_seqlens) - 1: - raise ValueError( - f"The number of initial states is expected to be equal to the number of input sequences, " - f"i.e., {len(cu_seqlens) - 1} rather than {initial_state.shape[0]}." - ) - assert b.shape == (*q.shape[:3], k.shape[-1]), ( - f"b must have shape [B, T, HV, K]; got {tuple(b.shape)} " - f"vs expected {(*q.shape[:3], k.shape[-1])}." - ) - assert w.shape == v.shape, ( - f"w must have shape [B, T, HV, V] matching v; got {tuple(w.shape)} " - f"vs v {tuple(v.shape)}." - ) - if scale is None: - scale = k.shape[-1] ** -0.5 - - o, final_state = fused_recurrent_gdn2_fwd( - q=q, - k=k, - v=v, - g=g, - b=b, - w=w, - A_log=A_log, - dt_bias=dt_bias, - scale=scale, - initial_state=initial_state, - inplace_final_state=False, - output_final_state=output_final_state, - use_qk_l2norm_in_kernel=use_qk_l2norm_in_kernel, - use_gate_in_kernel=use_gate_in_kernel, - lower_bound=lower_bound, - cu_seqlens=cu_seqlens, - state_v_first=state_v_first, - ) - return o, final_state - - -__all__ = [ - "fused_recurrent_gdn2", - "fused_recurrent_gdn2_fwd", - "fused_recurrent_gdn2_fwd_kernel", -] diff --git a/build/torch-xpu/ops/gdn2/naive.py b/build/torch-xpu/ops/gdn2/naive.py deleted file mode 100644 index 1a4f46f5dde3eab9abd361de9136ede65520c07e..0000000000000000000000000000000000000000 --- a/build/torch-xpu/ops/gdn2/naive.py +++ /dev/null @@ -1,205 +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 - -# Reference PyTorch implementation of Gated DeltaNet 2 (GDN-2). -# -# Update rule (per token, on the matrix state S in R^{K x V}): -# -# S_t = (I - k_t (b_t * k_t)^T) Diag(exp(g_t)) S_{t-1} + k_t (w_t * v_t)^T -# -# where `*` is the elementwise (Hadamard) product, `b_t in R^K` is the -# channel-wise erase gate on the key axis, `w_t in R^V` is the channel-wise -# write gate on the value axis, and `g_t in R^K` is the channel-wise log-decay. -# Collapsing `b_t = w_t = beta` to a scalar recovers KDA. - -import torch -import torch.nn.functional as F - - -def naive_recurrent_gdn2( - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - g: torch.Tensor, - b: torch.Tensor, - w: torch.Tensor, - scale: float | None = None, - initial_state: torch.Tensor | None = None, - output_final_state: bool = False, -): - """Token-by-token reference forward pass for GDN-2. - - Args: - q: queries of shape [B, T, H, K]. - k: keys of shape [B, T, H, K]. - v: values of shape [B, T, H, V]. - g: log-decay of shape [B, T, H, K] (already in the natural-log base). - b: channel-wise erase gate of shape [B, T, H, K] (range typically [0, 1]). - w: channel-wise write gate of shape [B, T, H, V] (range typically [0, 1]). - scale: attention scale; defaults to 1 / sqrt(K). - initial_state: optional [B, H, K, V] initial state in float32. - output_final_state: whether to return the final state. - - Returns: - o: outputs of shape [B, T, H, V]. - final_state: [B, H, K, V] if output_final_state else None. - """ - if scale is None: - scale = q.shape[-1] ** -0.5 - q, k, v, g, b, w = (x.transpose(1, 2).contiguous().float() for x in (q, k, v, g, b, w)) - B, H, T, K = k.shape - V = v.shape[-1] - o = torch.zeros(B, H, T, V, device=v.device, dtype=torch.float32) - h = torch.zeros(B, H, K, V, device=v.device, dtype=torch.float32) - if initial_state is not None: - h = initial_state.to(torch.float32).clone() - q = q * scale - - for t in range(T): - b_q = q[:, :, t] # [B, H, K] - b_k = k[:, :, t] # [B, H, K] - b_v = v[:, :, t] # [B, H, V] - b_g = g[:, :, t] # [B, H, K] - b_b = b[:, :, t] # [B, H, K] - b_w = w[:, :, t] # [B, H, V] - - # Per-channel decay applied along the K axis of the state. - h = h * b_g.exp().unsqueeze(-1) - # Gated read at key (b * k): pulls the existing contribution at this key. - erase = ((b_b * b_k).unsqueeze(-1) * h).sum(-2) # [B, H, V] - # Gated write: commit (w * v) with the erase contribution subtracted. - b_v_new = b_w * b_v - erase # [B, H, V] - # Rank-1 outer-product update. - h = h + b_k.unsqueeze(-1) * b_v_new.unsqueeze(-2) - o[:, :, t] = (b_q.unsqueeze(-1) * h).sum(-2) - - o = o.transpose(1, 2).contiguous().to(v.dtype) - if not output_final_state: - h = None - return o, h - - -def naive_chunk_gdn2( - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - g: torch.Tensor, - b: torch.Tensor, - w: torch.Tensor, - scale: float | None = None, - initial_state: torch.Tensor | None = None, - output_final_state: bool = False, - chunk_size: int = 64, -): - """Chunkwise PyTorch reference for GDN-2. - - The within-chunk recurrence is expanded via the WY representation: - `A = (I + tril(K' D K^T, -1))^{-1}` where `K' = b * k * exp(g)`, then - `u = A (w * v)` and `w_wy = A (b * k * exp(g))`. The inter-chunk state - recurrence carries the rank-`chunk_size` update across chunks. - """ - if scale is None: - scale = q.shape[-1] ** -0.5 - BT = chunk_size - - q, k, v, g, b, w = (x.transpose(1, 2).contiguous().float() for x in (q, k, v, g, b, w)) - B, H, T, K = k.shape - V = v.shape[-1] - pad_len = (BT - (T % BT)) % BT - if pad_len > 0: - q = F.pad(q, (0, 0, 0, pad_len)) - k = F.pad(k, (0, 0, 0, pad_len)) - v = F.pad(v, (0, 0, 0, pad_len)) - g = F.pad(g, (0, 0, 0, pad_len)) - b = F.pad(b, (0, 0, 0, pad_len)) - w = F.pad(w, (0, 0, 0, pad_len)) - T_pad = q.shape[2] - NT = T_pad // BT - - q = q * scale - # Reshape to chunks: [B, H, NT, BT, D] - - def chunk(x): - return x.view(B, H, NT, BT, -1) - q, k, v, g, b, w = (chunk(x) for x in (q, k, v, g, b, w)) - - # Cumulative log-decay within each chunk (per channel). - g_cum = g.cumsum(-2) # [B, H, NT, BT, K] - g_last = g_cum[..., -1:, :] # [B, H, NT, 1, K] — final state decay - # Pre-decayed key for the WY block. - k_g = k * g_cum.exp() # k * exp(g) - k_g_b = k_g * b # (b * k) * exp(g) - - # Strictly lower-triangular gated kk score matrix used by the WY inverse. - # Each (i, j) entry has decay factor exp(g_cum[i] - g_cum[j]). - decay_ij = (g_cum.unsqueeze(-2) - g_cum.unsqueeze(-3)) # [B, H, NT, BT, BT, K] - decay_ij_exp = decay_ij.exp() - tril_mask = torch.tril( - torch.ones(BT, BT, device=q.device, dtype=torch.bool), - diagonal=-1, - ) - # The strictly lower-triangular K^T K product, decay-weighted, b-modulated. - # The contraction is along the K axis of k and (b * k). - bk = b * k # [B, H, NT, BT, K] - # T_ij = sum_d (b*k)_i_d * k_j_d * exp(g_i_d - g_j_d) - decay_for_T = decay_ij_exp # [..., i, j, K] - T_lower = torch.einsum('bhnik,bhnjk,bhnijk->bhnij', - bk, k, decay_for_T) - T_lower = T_lower.masked_fill(~tril_mask, 0.0) - # Blocked forward substitution to build A_inv = (I + T_lower)^{-1} in place, - # mirroring the WY-block solve in fla.ops.gated_delta_rule.naive. - A_inv = -T_lower - for i in range(1, BT): - A_inv[..., i, :i] = A_inv[..., i, :i].clone() + ( - A_inv[..., i, :i, None].clone() * A_inv[..., :i, :i].clone() - ).sum(-2) - A_inv = A_inv + torch.eye(BT, device=q.device, dtype=torch.float32) - - # Within-chunk auxiliaries. - u_wy = A_inv @ (w * v) # [B, H, NT, BT, V] - w_wy = A_inv @ k_g_b # [B, H, NT, BT, K] - # Decayed key for the inter-chunk recurrence: k * exp(g_last - g_cum). - k_tail = k * (g_last - g_cum).exp() # [B, H, NT, BT, K] - - # Output path: causal intra-chunk QK^T (with decay) times v_new, plus - # the contribution from the carried state. - decay_qk = (g_cum.unsqueeze(-2) - g_cum.unsqueeze(-3)).exp() # [B, H, NT, BT, BT, K] - causal_mask = torch.tril( - torch.ones(BT, BT, device=q.device, dtype=torch.bool), - diagonal=0, - ) - - S = torch.zeros(B, H, K, V, device=v.device, dtype=torch.float32) - if initial_state is not None: - S = initial_state.to(torch.float32).clone() - o = torch.zeros_like(v) - for n in range(NT): - q_n = q[:, :, n] # [B, H, BT, K] - k_n = k[:, :, n] - g_n = g_cum[:, :, n] - g_last_n = g_last[:, :, n].squeeze(-2) # [B, H, K] - w_n = w_wy[:, :, n] # [B, H, BT, K] - u_n = u_wy[:, :, n] # [B, H, BT, V] - k_tail_n = k_tail[:, :, n] # [B, H, BT, K] - - # v_new = u - w_wy @ S (delta-rule write minus carried contribution). - v_new = u_n - w_n @ S # [B, H, BT, V] - # Intra-chunk attention with channel-wise decay on the K axis. - # A_qk[i, j] = sum_d q_i_d * k_j_d * exp(g_i_d - g_j_d) - A_qk = torch.einsum('bhik,bhjk,bhijk->bhij', - q_n, k_n, decay_qk[:, :, n]).masked_fill(~causal_mask, 0.0) - o_intra = A_qk @ v_new - # Inter-chunk contribution: q decayed up to chunk start times the carried state. - o_inter = (q_n * g_n.exp()) @ S - o[:, :, n] = o_intra + o_inter - # State update: S <- Diag(exp(g_last)) S + k_tail^T v_new. - S = S * g_last_n.unsqueeze(-1).exp() + k_tail_n.transpose(-1, -2) @ v_new - - o = o.reshape(B, H, T_pad, V)[:, :, :T].transpose(1, 2).contiguous().to(v.dtype) - if not output_final_state: - S = None - return o, S diff --git a/build/torch-xpu/ops/gdn2/wy_fast.py b/build/torch-xpu/ops/gdn2/wy_fast.py deleted file mode 100644 index 7f758624d39ae01d20b0f026fc7fcbd9d647f77c..0000000000000000000000000000000000000000 --- a/build/torch-xpu/ops/gdn2/wy_fast.py +++ /dev/null @@ -1,221 +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 - -# WY-representation recompute for GDN-2. -# -# Given the solved chunk inverse ``A = (I + tril(T,-1))^{-1}``, this kernel -# produces the two within-chunk auxiliaries consumed by the inter-chunk state -# recurrence: -# -# u = A @ (w * v) # write-side, V-axis -# w = A @ (b * exp(gk) * k) # erase-side, K-axis -# -# It optionally emits ``qg = q * exp(gk)`` and the tail-decayed key -# ``kg = k * exp(gn - gk)`` used by the inter-chunk update. -# -# Difference vs the KDA WY recompute: KDA uses a single scalar ``beta`` for both -# the u and w blocks. GDN-2 uses the channel-wise ``w_gate`` on the V axis for u -# and the channel-wise ``b`` on the K axis for w. These two gates live on -# different axes and cannot be reduced to a shared scalar; the kernel loads -# them per BV / BK block respectively. - -import torch -import triton -import triton.language as tl - -from ...ops.utils import prepare_chunk_indices -from ...ops.utils.cache import fla_cache_autotune -from ...ops.utils.op import exp2 -from ...utils import autotune_cache_kwargs - - -@triton.heuristics({ - 'STORE_QG': lambda args: args['qg'] is not None, - 'STORE_KG': lambda args: args['kg'] is not None, - 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None, -}) -@fla_cache_autotune( - configs=[ - triton.Config({}, num_warps=num_warps, num_stages=num_stages) - for num_warps in [2, 4, 8] - for num_stages in [2, 3, 4] - ], - key=['H', 'K', 'V', 'BT', 'BK', 'BV', 'IS_VARLEN'], - **autotune_cache_kwargs, -) -@triton.jit(do_not_specialize=['T']) -def recompute_w_u_fwd_gdn2_kernel( - q, - k, - qg, - kg, - v, - b, # channel-wise erase gate, [B, T, H, K] - w_gate, # channel-wise write gate, [B, T, H, V] - w, # output: A @ (b * exp(gk) * k), [B, T, H, K] - u, # output: A @ (w_gate * v), [B, T, H, V] - A, - gk, - cu_seqlens, - chunk_indices, - T, - H: tl.constexpr, - K: tl.constexpr, - V: tl.constexpr, - BT: tl.constexpr, - BK: tl.constexpr, - BV: tl.constexpr, - STORE_QG: tl.constexpr, - STORE_KG: tl.constexpr, - IS_VARLEN: tl.constexpr, -): - i_t, i_bh = tl.program_id(0), tl.program_id(1) - i_b, i_h = i_bh // H, i_bh % H - 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.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) - T = eos - bos - else: - bos, eos = i_b * T, i_b * T + T - - p_A = tl.make_block_ptr( - A + (bos * H + i_h) * BT, (T, BT), (H * BT, 1), - (i_t * BT, 0), (BT, BT), (1, 0), - ) - b_A = tl.load(p_A, boundary_check=(0, 1)) - - for i_v in range(tl.cdiv(V, BV)): - p_v = tl.make_block_ptr( - v + (bos * H + i_h) * V, (T, V), (H * V, 1), - (i_t * BT, i_v * BV), (BT, BV), (1, 0), - ) - p_u = tl.make_block_ptr( - u + (bos * H + i_h) * V, (T, V), (H * V, 1), - (i_t * BT, i_v * BV), (BT, BV), (1, 0), - ) - p_wg = tl.make_block_ptr( - w_gate + (bos * H + i_h) * V, (T, V), (H * V, 1), - (i_t * BT, i_v * BV), (BT, BV), (1, 0), - ) - b_v = tl.load(p_v, boundary_check=(0, 1)) - b_wg = tl.load(p_wg, boundary_check=(0, 1)) - b_vb = (b_v * b_wg).to(b_v.dtype) - b_u = tl.dot(b_A, b_vb) - tl.store(p_u, b_u.to(p_u.dtype.element_ty), boundary_check=(0, 1)) - - for i_k in range(tl.cdiv(K, BK)): - p_w = tl.make_block_ptr( - w + (bos * H + i_h) * K, (T, K), (H * K, 1), - (i_t * BT, i_k * BK), (BT, BK), (1, 0), - ) - p_k = tl.make_block_ptr( - k + (bos * H + i_h) * K, (T, K), (H * K, 1), - (i_t * BT, i_k * BK), (BT, BK), (1, 0), - ) - p_b = tl.make_block_ptr( - b + (bos * H + i_h) * K, (T, K), (H * K, 1), - (i_t * BT, i_k * BK), (BT, BK), (1, 0), - ) - b_k = tl.load(p_k, boundary_check=(0, 1)) - b_b = tl.load(p_b, boundary_check=(0, 1)) - b_kb = b_k * b_b - - p_gk = tl.make_block_ptr( - gk + (bos * H + i_h) * K, (T, K), (H * K, 1), - (i_t * BT, i_k * BK), (BT, BK), (1, 0), - ) - b_gk = tl.load(p_gk, boundary_check=(0, 1)).to(tl.float32) - b_kb *= exp2(b_gk) - - if STORE_QG: - p_q = tl.make_block_ptr( - q + (bos * H + i_h) * K, (T, K), (H * K, 1), - (i_t * BT, i_k * BK), (BT, BK), (1, 0), - ) - p_qg = tl.make_block_ptr( - qg + (bos * H + i_h) * K, (T, K), (H * K, 1), - (i_t * BT, i_k * BK), (BT, BK), (1, 0), - ) - b_q = tl.load(p_q, boundary_check=(0, 1)) - b_qg = b_q * exp2(b_gk) - tl.store(p_qg, b_qg.to(p_qg.dtype.element_ty), boundary_check=(0, 1)) - - if STORE_KG: - last_idx = min(i_t * BT + BT, T) - 1 - o_k = i_k * BK + tl.arange(0, BK) - m_k = o_k < K - b_gn = tl.load( - gk + ((bos + last_idx) * H + i_h) * K + o_k, - mask=m_k, other=0., - ).to(tl.float32) - b_kg = b_k * tl.where( - (i_t * BT + tl.arange(0, BT) < T)[:, None], - exp2(b_gn[None, :] - b_gk), - 0, - ) - p_kg = tl.make_block_ptr( - kg + (bos * H + i_h) * K, (T, K), (H * K, 1), - (i_t * BT, i_k * BK), (BT, BK), (1, 0), - ) - tl.store(p_kg, b_kg.to(p_kg.dtype.element_ty), boundary_check=(0, 1)) - - b_w = tl.dot(b_A, b_kb.to(b_k.dtype)) - tl.store(p_w, b_w.to(p_w.dtype.element_ty), boundary_check=(0, 1)) - - -def recompute_w_u_fwd_gdn2( - k: torch.Tensor, - v: torch.Tensor, - b: torch.Tensor, - w_gate: torch.Tensor, - A: torch.Tensor, - q: torch.Tensor | None = None, - gk: torch.Tensor | None = None, - cu_seqlens: torch.LongTensor | None = None, - chunk_indices: torch.LongTensor | None = None, -) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor | None, torch.Tensor | None]: - """Produce GDN-2 WY auxiliaries from the solved chunk inverse ``A``. - - Returns ``(w, u, qg, kg)`` with shapes matching the inputs. ``qg`` is - produced only when ``q`` is provided; ``kg`` only when ``gk`` is provided. - """ - B, T, H, K, V = *k.shape, v.shape[-1] - BT = A.shape[-1] - BK = 64 - BV = 64 - - if chunk_indices is None and cu_seqlens is not None: - chunk_indices = prepare_chunk_indices(cu_seqlens, BT) - NT = triton.cdiv(T, BT) if cu_seqlens is None else len(chunk_indices) - - w = torch.empty_like(k) - u = torch.empty_like(v) - qg = torch.empty_like(q) if q is not None else None - kg = torch.empty_like(k) if gk is not None else None - recompute_w_u_fwd_gdn2_kernel[(NT, B * H)]( - q=q, - k=k, - qg=qg, - kg=kg, - v=v, - b=b, - w_gate=w_gate, - w=w, - u=u, - A=A, - gk=gk, - cu_seqlens=cu_seqlens, - chunk_indices=chunk_indices, - T=T, - H=H, - K=K, - V=V, - BT=BT, - BK=BK, - BV=BV, - ) - return w, u, qg, kg diff --git a/build/torch-xpu/ops/generalized_delta_rule/__init__.py b/build/torch-xpu/ops/generalized_delta_rule/__init__.py deleted file mode 100644 index ff4c7656f940b562b781bcefcb60b1676ea1f53b..0000000000000000000000000000000000000000 --- a/build/torch-xpu/ops/generalized_delta_rule/__init__.py +++ /dev/null @@ -1,16 +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 .dplr import chunk_dplr_delta_rule, fused_recurrent_dplr_delta_rule -from .iplr import chunk_iplr_delta_rule, fused_recurrent_iplr_delta_rule - -__all__ = [ - 'chunk_dplr_delta_rule', - 'chunk_iplr_delta_rule', - 'fused_recurrent_dplr_delta_rule', - 'fused_recurrent_iplr_delta_rule', -] diff --git a/build/torch-xpu/ops/generalized_delta_rule/dplr/__init__.py b/build/torch-xpu/ops/generalized_delta_rule/dplr/__init__.py deleted file mode 100644 index 7e4a0b5e5615e9c7fce594a90f548d0b9bf64a9a..0000000000000000000000000000000000000000 --- a/build/torch-xpu/ops/generalized_delta_rule/dplr/__init__.py +++ /dev/null @@ -1,14 +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 .chunk import chunk_dplr_delta_rule -from .fused_recurrent import fused_recurrent_dplr_delta_rule - -__all__ = [ - 'chunk_dplr_delta_rule', - 'fused_recurrent_dplr_delta_rule', -] diff --git a/build/torch-xpu/ops/generalized_delta_rule/dplr/chunk.py b/build/torch-xpu/ops/generalized_delta_rule/dplr/chunk.py deleted file mode 100644 index 240d0f71ffc59774bcf079e88622520692715898..0000000000000000000000000000000000000000 --- a/build/torch-xpu/ops/generalized_delta_rule/dplr/chunk.py +++ /dev/null @@ -1,554 +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 warnings - -import torch - -from ....ops.cp import FLACPContext -from ....ops.cp.chunk_delta_h import ( - chunk_gated_delta_rule_bwd_dhu_pre_process, - chunk_gated_delta_rule_fwd_h_pre_process, - compress_h0, - expand_h0, -) -from ....ops.generalized_delta_rule.dplr.chunk_A_bwd import chunk_dplr_bwd_dqk_intra -from ....ops.generalized_delta_rule.dplr.chunk_A_fwd import chunk_dplr_fwd_intra -from ....ops.generalized_delta_rule.dplr.chunk_h_bwd import chunk_dplr_bwd_dhu -from ....ops.generalized_delta_rule.dplr.chunk_h_fwd import chunk_dplr_fwd_h -from ....ops.generalized_delta_rule.dplr.chunk_o_bwd import chunk_dplr_bwd_dAu, chunk_dplr_bwd_dv, chunk_dplr_bwd_o -from ....ops.generalized_delta_rule.dplr.chunk_o_fwd import chunk_dplr_fwd_o -from ....ops.generalized_delta_rule.dplr.wy_fast_bwd import chunk_dplr_bwd_wy -from ....ops.generalized_delta_rule.dplr.wy_fast_fwd import prepare_wy_repr_fwd -from ....ops.rwkv6.chunk import chunk_rwkv6_fwd_cumsum -from ....ops.utils import prepare_chunk_indices -from ....ops.utils.constant import RCP_LN2 -from ....utils import TRITON_ABOVE_3_4_0, autocast_custom_bwd, autocast_custom_fwd, input_guard - - -def chunk_dplr_fwd( - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - a: torch.Tensor, - b: torch.Tensor, - gk: torch.Tensor, - scale: float, - initial_state: torch.Tensor, - output_final_state: bool, - cu_seqlens: torch.LongTensor | None = None, - chunk_size: int = 16, - safe_gate: bool = False, - chunk_indices: torch.LongTensor | None = None, - disable_recompute: bool = False, - cp_context: FLACPContext | None = None, -): - gi, ge = chunk_rwkv6_fwd_cumsum( - gk, - chunk_size, - scale=RCP_LN2, - cu_seqlens=cu_seqlens, - chunk_indices=chunk_indices, - ) - - A_ab, A_qk, A_ak, A_qb, qg, kg, ag, bg = chunk_dplr_fwd_intra( - q=q, - k=k, - a=a, - b=b, - gi=gi, - ge=ge, - scale=scale, - cu_seqlens=cu_seqlens, - safe_gate=safe_gate, - chunk_size=chunk_size, - chunk_indices=chunk_indices, - ) - - # A_ab, A_ak, gi, ge torch.float32 - # A_qk, A_qb, qg, kg, ag, bg, dtype=q.dtype, eg: bf16 - w, u, A_ab_inv = prepare_wy_repr_fwd( - ag=ag, - A_ab=A_ab, - A_ak=A_ak, - v=v, - cu_seqlens=cu_seqlens, - chunk_size=chunk_size, - chunk_indices=chunk_indices, - ) - - if cp_context is not None: - initial_state = chunk_gated_delta_rule_fwd_h_pre_process( - k=kg, - w=w, - u=u, - gk=gi, - bg=bg, - v=v, - cu_seqlens=cu_seqlens, - initial_state=initial_state, - context=cp_context, - chunk_size=chunk_size, - ) - - h, v_new, final_state = chunk_dplr_fwd_h( - kg=kg, - bg=bg, - v=v, - w=w, - u=u, - gk=gi, - initial_state=initial_state, - output_final_state=output_final_state, - cu_seqlens=cu_seqlens, - chunk_size=chunk_size, - chunk_indices=chunk_indices, - ) - - if cp_context is not None: - initial_state = compress_h0(initial_state, context=cp_context) - - o = chunk_dplr_fwd_o( - qg=qg, - v=v, - v_new=v_new, - A_qk=A_qk, - A_qb=A_qb, - h=h, - cu_seqlens=cu_seqlens, - chunk_size=chunk_size, - chunk_indices=chunk_indices, - ) - - if disable_recompute: - return o, final_state, initial_state, (gi, ge, A_qk, A_qb, A_ak, qg, kg, ag, bg, w, h, v_new, A_ab_inv) - else: - return o, final_state, initial_state, None - - -class ChunkDPLRDeltaRuleFunction(torch.autograd.Function): - - @staticmethod - @input_guard - @autocast_custom_fwd - def forward( - ctx, - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - a: torch.Tensor, - b: torch.Tensor, - gk: torch.Tensor, - scale: float, - initial_state: torch.Tensor, - output_final_state: bool, - cu_seqlens: torch.LongTensor | None = None, - cu_seqlens_cpu: torch.LongTensor | None = None, - safe_gate: bool = False, - chunk_size: int | None = None, - disable_recompute: bool = False, - cp_context: FLACPContext | None = None, - ): - # Due to gate numerical stability consideration, we only support chunk_size=16 when safe_gate=True - # And in practice, chunk_size=16 is sufficient for no safe gate situations. - # It's different from the other chunk implementations. - if chunk_size is None: - chunk_size = 16 - elif TRITON_ABOVE_3_4_0: - chunk_size = chunk_size - else: - # Avoid Triton Compiler error - warnings.warn( - "Set chunk_size to 16, to avoid triton compiler erorr. " - f"original chunk_size {chunk_size}", - category=RuntimeWarning, - stacklevel=2, - ) - chunk_size = 16 - chunk_indices = None - if cu_seqlens is not None: - chunk_indices = prepare_chunk_indices(cu_seqlens, chunk_size, cu_seqlens_cpu=cu_seqlens_cpu) - - # chunk_dplr_fwd returns the possibly CP-merged + compressed initial_state; - # for CP this is the [1, HV, K, V] state we must save for backward so the - # forward recomputation can rebuild the correct per-rank h. - o, final_state, initial_state, cache = chunk_dplr_fwd( - q=q, - k=k, - v=v, - a=a, - b=b, - gk=gk, - scale=scale, - initial_state=initial_state, - output_final_state=output_final_state, - cu_seqlens=cu_seqlens, - chunk_size=chunk_size, - safe_gate=safe_gate, - chunk_indices=chunk_indices, - disable_recompute=disable_recompute, - cp_context=cp_context, - ) - - if disable_recompute: - gi, ge, A_qk, A_qb, A_ak, qg, kg, ag, bg, w, h, v_new, A_ab_inv = cache - ctx.save_for_backward( - q, - k, - v, - a, - b, - gk, - initial_state, - gi, - ge, - A_qk, - A_qb, - A_ak, - qg, - kg, - ag, - bg, - w, - h, - v_new, - A_ab_inv, - ) - else: - ctx.save_for_backward(q, k, v, a, b, gk, initial_state) - ctx.cu_seqlens = cu_seqlens - ctx.scale = scale - ctx.chunk_size = chunk_size - ctx.chunk_indices = chunk_indices - ctx.safe_gate = safe_gate - ctx.disable_recompute = disable_recompute - ctx.cp_context = cp_context - return o.to(q.dtype), final_state - - @staticmethod - @input_guard - @autocast_custom_bwd - def backward( - ctx, - do: torch.Tensor, - dht: torch.Tensor, - ): - if ctx.disable_recompute: - ( - q, - k, - v, - a, - b, - gk, - initial_state, - gi, - ge, - A_qk, - A_qb, - A_ak, - qg, - kg, - ag, - bg, - w, - h, - v_new, - A_ab_inv, - ) = ctx.saved_tensors - else: - q, k, v, a, b, gk, initial_state = ctx.saved_tensors - chunk_size = ctx.chunk_size - cu_seqlens = ctx.cu_seqlens - scale = ctx.scale - cp_context = ctx.cp_context - - # When CP compressed the saved initial_state to a single [1, HV, K, V] entry - # (for the first local sequence only), expand it back to [N_local, HV, K, V] - # BEFORE the forward recomputation. Otherwise chunk_dplr_fwd_h indexes past - # the buffer for the non-first sequences on this rank and reads garbage. - if cp_context is not None and initial_state is not None and initial_state.shape[0] == 1: - initial_state = expand_h0(initial_state, context=cp_context) - - if not ctx.disable_recompute: - # ******* start recomputing everything, otherwise i believe the gpu memory will be exhausted ******* - gi, ge = chunk_rwkv6_fwd_cumsum( - gk, - chunk_size, - scale=RCP_LN2, - cu_seqlens=cu_seqlens, - chunk_indices=ctx.chunk_indices, - ) - - A_ab, A_qk, A_ak, A_qb, qg, kg, ag, bg = chunk_dplr_fwd_intra( - q=q, - k=k, - a=a, - b=b, - gi=gi, - ge=ge, - scale=scale, - cu_seqlens=cu_seqlens, - safe_gate=ctx.safe_gate, - chunk_size=chunk_size, - chunk_indices=ctx.chunk_indices, - ) - w, u, A_ab_inv = prepare_wy_repr_fwd( - ag=ag, - A_ab=A_ab, - A_ak=A_ak, - v=v, - cu_seqlens=cu_seqlens, - chunk_size=chunk_size, - chunk_indices=ctx.chunk_indices, - ) - del A_ab - h, v_new, _ = chunk_dplr_fwd_h( - kg=kg, - bg=bg, - v=v, - w=w, - u=u, - gk=gi, - initial_state=initial_state, - cu_seqlens=cu_seqlens, - chunk_size=chunk_size, - chunk_indices=ctx.chunk_indices, - ) - del u - # ******* end of recomputation ******* - # A_ak, A_ab_inv, gi, ge torch.float32 - # A_qk, A_qb, qg, kg, ag, bg, v_new dtype=q.dtype, eg: bf16 - - dv_new_intra, dA_qk, dA_qb = chunk_dplr_bwd_dAu( - v=v, - v_new=v_new, - do=do, - A_qb=A_qb, - scale=scale, - cu_seqlens=cu_seqlens, - chunk_size=chunk_size, - chunk_indices=ctx.chunk_indices, - ) - - if cp_context is not None: - dht, initial_state = chunk_gated_delta_rule_bwd_dhu_pre_process( - q=qg, - k=kg, - w=w, - do=do, - dv=dv_new_intra, - gk=gi, - bg=bg, - scale=1.0, - cu_seqlens=cu_seqlens, - dht=dht, - initial_state=initial_state, - context=cp_context, - chunk_size=chunk_size, - ) - - dh, dh0, dv_new = chunk_dplr_bwd_dhu( - qg=qg, - bg=bg, - w=w, - gk=gi, - h0=initial_state, - dht=dht, - do=do, - dv=dv_new_intra, - cu_seqlens=cu_seqlens, - chunk_size=chunk_size, - chunk_indices=ctx.chunk_indices, - ) - - dv = chunk_dplr_bwd_dv( - A_qk=A_qk, - kg=kg, - do=do, - dh=dh, - cu_seqlens=cu_seqlens, - chunk_size=chunk_size, - chunk_indices=ctx.chunk_indices, - ) - del A_qk - - dqg, dkg, dw, dbg, dgk_last = chunk_dplr_bwd_o( - k=kg, - b=bg, - v=v, - v_new=v_new, - do=do, - h=h, - dh=dh, - dv=dv_new, - w=w, - gk=gi, - cu_seqlens=cu_seqlens, - chunk_size=chunk_size, - scale=scale, - chunk_indices=ctx.chunk_indices, - ) - del v_new - - dA_ab, dA_ak, dv, dag = chunk_dplr_bwd_wy( - A_ab_inv=A_ab_inv, - A_ak=A_ak, - v=v, - ag=ag, - dw=dw, - du=dv_new, - dv0=dv, - cu_seqlens=cu_seqlens, - chunk_size=chunk_size, - chunk_indices=ctx.chunk_indices, - ) - del A_ak - - dq, dk, da, db, dgk = chunk_dplr_bwd_dqk_intra( - q=q, - k=k, - a=a, - b=b, - gi=gi, - ge=ge, - dAqk=dA_qk, - dAqb=dA_qb, - dAak=dA_ak, - dAab=dA_ab, - dgk_last=dgk_last, - dqg=dqg, - dkg=dkg, - dag=dag, - dbg=dbg, - chunk_size=chunk_size, - scale=scale, - cu_seqlens=cu_seqlens, - safe_gate=ctx.safe_gate, - chunk_indices=ctx.chunk_indices, - ) - - return ( - dq.to(q), dk.to(k), dv.to(v), da.to(a), db.to(b), dgk.to(gk), - None, dh0, None, None, None, None, None, None, None, - ) - - -@torch.compiler.disable -def chunk_dplr_delta_rule( - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - a: torch.Tensor, - b: torch.Tensor, - gk: torch.Tensor, - scale: float | None = None, - initial_state: torch.Tensor | None = None, - output_final_state: bool = False, - cu_seqlens: torch.LongTensor | None = None, - cu_seqlens_cpu: torch.LongTensor | None = None, - safe_gate: bool = False, - chunk_size: int | None = None, - disable_recompute: bool = False, - cp_context: FLACPContext | None = None, - **kwargs, -): - r""" - Args: - q (torch.Tensor): - queries of shape `[B, T, H, K]`. - k (torch.Tensor): - keys of shape `[B, T, H, K]`. - v (torch.Tensor): - values of shape `[B, T, H, V]`. - a (torch.Tensor): - activations of shape `[B, T, H, K]`. - b (torch.Tensor): - betas of shape `[B, T, H, K]`. - gk (torch.Tensor): - gk of shape `[B, T, H, K]`. Decay term in log space. - scale (Optional[float]): - Scale factor for the attention scores. - If not provided, it will default to `1 / sqrt(K)`. Default: `None`. - initial_state (Optional[torch.Tensor]): - Initial state of shape `[N, H, K, V]` for `N` input sequences. - For equal-length input sequences, `N` equals the batch size `B`. - Default: `None`. - output_final_state (Optional[bool]): - Whether to output the final state of shape `[N, H, K, V]`. Default: `False`. - cu_seqlens (torch.LongTensor): - Cumulative sequence lengths of shape `[N+1]` used for variable-length training, - consistent with the FlashAttention API. - cu_seqlens_cpu (torch.LongTensor): - CPU copy of `cu_seqlens` to avoid unnecessary device synchronization. Default: `None`. - safe_gate (Optional[bool]): - Whether the kernel can assume the input gate values `g` are in a safe range. - When `True`, the kernel can use M=16 TensorCore acceleration. - The safe range is approximately `[-5, 0)`. Default: `False`. - chunk_size (Optional[int]): - Chunk size for the chunked computation. Default: `None`, which means 16. - disable_recompute (Optional[bool]): - Whether to disable gradient recomputation in the kernel. Default: `False`. - cp_context (Optional[FLACPContext]): - Context parallel context for distributed training across multiple devices. - When provided, `initial_state` and `output_final_state` are not supported. - Default: `None`. - - Returns: - o (torch.Tensor): - Outputs of shape `[B, T, H, V]`. - final_state (torch.Tensor): - Final state of shape `[N, H, K, V]` if `output_final_state=True` else `None`. - """ - if q.dtype == torch.float32: - warnings.warn( - """ChunkDeltaRuleFunction does not support float32 on some platforms. Please use bfloat16/float16. - If you want to use float32, please solve the issue by yourself.""", - category=RuntimeWarning, - stacklevel=2, - ) - if 'head_first' in kwargs: - raise DeprecationWarning( - "head_first has been removed. Inputs must be in `[B, T, H, ...]` format.", - ) - 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" - assert cp_context.cu_seqlens is not None, "cu_seqlens is required for CP" - cu_seqlens = cp_context.cu_seqlens - if cp_context.cu_seqlens_cpu is not None: - cu_seqlens_cpu = cp_context.cu_seqlens_cpu - if cu_seqlens is not None: - if q.shape[0] != 1: - raise ValueError( - f"The batch size is expected to be 1 rather than {q.shape[0]} when using `cu_seqlens`." - f"Please flatten variable-length inputs before processing.", - ) - if initial_state is not None and initial_state.shape[0] != len(cu_seqlens) - 1: - raise ValueError( - f"The number of initial states is expected to be equal to the number of input sequences, " - f"i.e., {len(cu_seqlens) - 1} rather than {initial_state.shape[0]}.", - ) - scale = k.shape[-1] ** -0.5 if scale is None else scale - o, final_state = ChunkDPLRDeltaRuleFunction.apply( - q, - k, - v, - a, - b, - gk, - scale, - initial_state, - output_final_state, - cu_seqlens, - cu_seqlens_cpu, - safe_gate, - chunk_size, - disable_recompute, - cp_context, - ) - return o, final_state diff --git a/build/torch-xpu/ops/generalized_delta_rule/dplr/chunk_A_bwd.py b/build/torch-xpu/ops/generalized_delta_rule/dplr/chunk_A_bwd.py deleted file mode 100644 index fa6c97b7d5c5e1f8926ab19d2a1c486ece9e7bec..0000000000000000000000000000000000000000 --- a/build/torch-xpu/ops/generalized_delta_rule/dplr/chunk_A_bwd.py +++ /dev/null @@ -1,544 +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 -import triton.language as tl - -from ....ops.utils import prepare_chunk_indices -from ....ops.utils.op import exp2, gather -from ....utils import IS_AMD, IS_GATHER_SUPPORTED, autotune_cache_kwargs, check_shared_mem - -NUM_WARPS_AUTOTUNE = [2, 4, 8, 16] if IS_AMD else [2, 4, 8, 16, 32] - - -@triton.heuristics({ - 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None, -}) -@triton.autotune( - configs=[ - triton.Config({}, num_warps=num_warps, num_stages=num_stages) - for num_warps in NUM_WARPS_AUTOTUNE - for num_stages in [2, 3, 4] - ], - key=['BK', 'BT', 'K'], - **autotune_cache_kwargs, -) -@triton.jit(do_not_specialize=['T']) -def chunk_dplr_bwd_kernel_intra( - q, - k, - a, - b, - gi, - ge, - dAqk, - dAqb, - dAak, - dAab, - dq, - dk, - da, - db, - dqg, - dkg, - dag, - dbg, - dgk, - dgk_offset, - cu_seqlens, - chunk_indices, - scale: tl.constexpr, - T, - H: tl.constexpr, - K: tl.constexpr, - BT: tl.constexpr, - BC: tl.constexpr, - BK: tl.constexpr, - IS_VARLEN: tl.constexpr, - GATHER_SUPPORTED: tl.constexpr, -): - i_k, i_t, i_bh = tl.program_id(0), tl.program_id(1), tl.program_id(2) - i_b, i_h = i_bh // H, i_bh % H - 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.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) - T = eos - bos - else: - bos, eos = (i_b * T).to(tl.int32), (i_b * T + T).to(tl.int32) - - if i_t * BT >= T: - return - - # offset calculation - ge += (bos*H + i_h) * K - gi += (bos*H + i_h) * K - q += (bos*H + i_h) * K - a += (bos*H + i_h) * K - b += (bos*H + i_h) * K - k += (bos*H + i_h) * K - dq += (bos*H + i_h) * K - dk += (bos*H + i_h) * K - da += (bos*H + i_h) * K - db += (bos*H + i_h) * K - dqg += (bos*H + i_h) * K - dag += (bos*H + i_h) * K - dkg += (bos*H + i_h) * K - dbg += (bos*H + i_h) * K - dgk += (bos*H + i_h) * K - dgk_offset += (bos*H + i_h) * K - dAqk += (bos*H + i_h) * BT - dAqb += (bos*H + i_h) * BT - dAak += (bos*H + i_h) * BT - dAab += (bos*H + i_h) * BT - - stride_qk = H*K - stride_A = H*BT - - p_ge = tl.make_block_ptr(ge, (T, K), (stride_qk, 1), (i_t * BT, i_k * BK), (BC, BK), (1, 0)) - p_gi = tl.make_block_ptr(gi, (T, K), (stride_qk, 1), (i_t * BT, i_k * BK), (BC, BK), (1, 0)) - # [BC, BK] - b_ge = tl.load(p_ge, boundary_check=(0, 1)) - b_gi = tl.load(p_gi, boundary_check=(0, 1)) - b_dq = tl.zeros([BC, BK], dtype=tl.float32) - b_da = tl.zeros([BC, BK], dtype=tl.float32) - b_dk = tl.zeros([BC, BK], dtype=tl.float32) - b_db = tl.zeros([BC, BK], dtype=tl.float32) - # intra chunk gradient calculation - p_dAqk = tl.make_block_ptr(dAqk, (T, BT), (stride_A, 1), (i_t*BT, 0), (BC, BC), (1, 0)) - p_dAab = tl.make_block_ptr(dAab, (T, BT), (stride_A, 1), (i_t*BT, 0), (BC, BC), (1, 0)) - p_dAqb = tl.make_block_ptr(dAqb, (T, BT), (stride_A, 1), (i_t*BT, 0), (BC, BC), (1, 0)) - p_dAak = tl.make_block_ptr(dAak, (T, BT), (stride_A, 1), (i_t*BT, 0), (BC, BC), (1, 0)) - o_i = tl.arange(0, BC) - p_k = tl.make_block_ptr(k, (T, K), (stride_qk, 1), (i_t*BT, i_k*BK), (BC, BK), (1, 0)) - p_b = tl.make_block_ptr(b, (T, K), (stride_qk, 1), (i_t*BT, i_k*BK), (BC, BK), (1, 0)) - p_a = tl.make_block_ptr(a, (T, K), (stride_qk, 1), (i_t*BT, i_k*BK), (BC, BK), (1, 0)) - p_q = tl.make_block_ptr(q, (T, K), (stride_qk, 1), (i_t*BT, i_k*BK), (BC, BK), (1, 0)) - b_k = tl.load(p_k, boundary_check=(0, 1)) - b_b = tl.load(p_b, boundary_check=(0, 1)) - b_q = tl.load(p_q, boundary_check=(0, 1)) - b_a = tl.load(p_a, boundary_check=(0, 1)) - b_dAqk = tl.load(p_dAqk, boundary_check=(0, 1)) - b_dAab = tl.load(p_dAab, boundary_check=(0, 1)) - b_dAqb = tl.load(p_dAqb, boundary_check=(0, 1)) - b_dAak = tl.load(p_dAak, boundary_check=(0, 1)) - - # inter chunk gradient calculation - o_k = i_k * BK + tl.arange(0, BK) - m_k = o_k < K - # intra chunk gradient calculation - for j in range(0, min(BC, T - i_t * BT)): - # trick to index the block - if GATHER_SUPPORTED: - row_idx = tl.full([1, BK], j, dtype=tl.int16) - col_idx = tl.full([BC, 1], j, dtype=tl.int16) - row_idx_bc = tl.full([1, BC], j, dtype=tl.int16) - # [1, BK] - b_kj = gather(b_k, row_idx, axis=0) - b_bj = gather(b_b, row_idx, axis=0) - b_gij = gather(b_gi, row_idx, axis=0) - b_gej = gather(b_ge, row_idx, axis=0) - b_qj = gather(b_q, row_idx, axis=0) - b_aj = gather(b_a, row_idx, axis=0) - # [BC, 1] - b_dAqk_j = gather(b_dAqk, col_idx, axis=1) - b_dAab_j = gather(b_dAab, col_idx, axis=1) - b_dAqb_j = gather(b_dAqb, col_idx, axis=1) - b_dAak_j = gather(b_dAak, col_idx, axis=1) - # [1, BC] -> [BC, 1] - b_dA_qk_j = tl.sum(gather(b_dAqk, row_idx_bc, axis=0), 0)[:, None] - b_dA_qk_j = tl.sum(gather(b_dAqk, row_idx_bc, axis=0), 0)[:, None] - b_dA_ab_j = tl.sum(gather(b_dAab, row_idx_bc, axis=0), 0)[:, None] - b_dA_qb_j = tl.sum(gather(b_dAqb, row_idx_bc, axis=0), 0)[:, None] - b_dA_ak_j = tl.sum(gather(b_dAak, row_idx_bc, axis=0), 0)[:, None] - else: - mask_idx = tl.arange(0, BC) == j - b_kj = tl.sum(tl.where(mask_idx[:, None], b_k, 0), 0)[None, :] - b_bj = tl.sum(tl.where(mask_idx[:, None], b_b, 0), 0)[None, :] - b_gij = tl.sum(tl.where(mask_idx[:, None], b_gi, 0), 0)[None, :] - b_gej = tl.sum(tl.where(mask_idx[:, None], b_ge, 0), 0)[None, :] - b_dAqk_j = tl.sum(tl.where(mask_idx[None, :], b_dAqk, 0), 1)[:, None] - b_dAab_j = tl.sum(tl.where(mask_idx[None, :], b_dAab, 0), 1)[:, None] - b_dAqb_j = tl.sum(tl.where(mask_idx[None, :], b_dAqb, 0), 1)[:, None] - b_dAak_j = tl.sum(tl.where(mask_idx[None, :], b_dAak, 0), 1)[:, None] - b_dA_qk_j = tl.sum(tl.where(mask_idx[:, None], b_dAqk, 0), 0)[:, None] - b_dA_ab_j = tl.sum(tl.where(mask_idx[:, None], b_dAab, 0), 0)[:, None] - b_dA_qb_j = tl.sum(tl.where(mask_idx[:, None], b_dAqb, 0), 0)[:, None] - b_dA_ak_j = tl.sum(tl.where(mask_idx[:, None], b_dAak, 0), 0)[:, None] - # [1, BK] b_qj, b_aj - b_qj = tl.sum(tl.where(mask_idx[:, None], b_q, 0), 0)[None, :] - b_aj = tl.sum(tl.where(mask_idx[:, None], b_a, 0), 0)[None, :] - - m_e = o_i[:, None] > j - m_i = o_i[:, None] >= j - tmp1 = exp2(b_gi - b_gij) - tmp2 = exp2(b_ge - b_gij) - b_dq += tl.where(m_i, b_dAqk_j * b_kj * tmp1, 0.) - b_dq += tl.where(m_i, b_dAqb_j * b_bj * tmp1, 0.) - b_da += tl.where(m_e, b_dAab_j * b_bj * tmp2, 0.) - b_da += tl.where(m_e, b_dAak_j * b_kj * tmp2, 0.) - - m_i = o_i[:, None] <= j - m_e = o_i[:, None] < j - tmp1 = exp2(b_gij - b_gi) - tmp2 = exp2(b_gej - b_gi) - b_dk += tl.where(m_i, b_dA_qk_j * b_qj * tmp1, 0.) - b_dk += tl.where(m_e, b_dA_ak_j * b_aj * tmp2, 0.) - b_db += tl.where(m_i, b_dA_qb_j * b_qj * tmp1, 0.) - b_db += tl.where(m_e, b_dA_ab_j * b_aj * tmp2, 0.) - - # post processing - p_dq = tl.make_block_ptr(dq, (T, K), (stride_qk, 1), (i_t * BT, i_k * BK), (BC, BK), (1, 0)) - p_dk = tl.make_block_ptr(dk, (T, K), (stride_qk, 1), (i_t * BT, i_k * BK), (BC, BK), (1, 0)) - p_da = tl.make_block_ptr(da, (T, K), (stride_qk, 1), (i_t * BT, i_k * BK), (BC, BK), (1, 0)) - p_db = tl.make_block_ptr(db, (T, K), (stride_qk, 1), (i_t * BT, i_k * BK), (BC, BK), (1, 0)) - p_dgk = tl.make_block_ptr(dgk, (T, K), (stride_qk, 1), (i_t * BT, i_k * BK), (BC, BK), (1, 0)) - p_dgk_offset = tl.make_block_ptr(dgk_offset, (T, K), (stride_qk, 1), (i_t * BT, i_k * BK), (BC, BK), (1, 0)) - p_dqg = tl.make_block_ptr(dqg, (T, K), (stride_qk, 1), (i_t * BT, i_k * BK), (BC, BK), (1, 0)) - p_dkg = tl.make_block_ptr(dkg, (T, K), (stride_qk, 1), (i_t * BT, i_k * BK), (BC, BK), (1, 0)) - p_dag = tl.make_block_ptr(dag, (T, K), (stride_qk, 1), (i_t * BT, i_k * BK), (BC, BK), (1, 0)) - p_dbg = tl.make_block_ptr(dbg, (T, K), (stride_qk, 1), (i_t * BT, i_k * BK), (BC, BK), (1, 0)) - p_gn = gi + (min(i_t * BT + BT, T) - 1)*stride_qk + o_k - p_gn = tl.max_contiguous(tl.multiple_of(p_gn, BK), BK) - b_gn = tl.load(p_gn, mask=m_k, other=0) - b_da += tl.load(p_dag, boundary_check=(0, 1)) * exp2(b_ge) - b_dq += tl.load(p_dqg, boundary_check=(0, 1)) * exp2(b_gi) * scale - tmp = exp2(b_gn[None, :] - b_gi) - b_dk += tl.load(p_dkg, boundary_check=(0, 1)).to(tl.float32) * tmp - b_db += tl.load(p_dbg, boundary_check=(0, 1)).to(tl.float32) * tmp - tl.store(p_dq, (b_dq).to(p_dq.dtype.element_ty), boundary_check=(0, 1)) - tl.store(p_dk, b_dk.to(p_dk.dtype.element_ty), boundary_check=(0, 1)) - tl.store(p_da, b_da.to(p_da.dtype.element_ty), boundary_check=(0, 1)) - tl.store(p_db, b_db.to(p_db.dtype.element_ty), boundary_check=(0, 1)) - b_dgk = (b_dq * b_q + b_da * b_a - b_dk * b_k - b_db * b_b).to(tl.float32) - b_dgk_offset = b_da * b_a - tl.store(p_dgk, b_dgk.to(p_dgk.dtype.element_ty), boundary_check=(0, 1)) - tl.store(p_dgk_offset, b_dgk_offset.to(p_dgk_offset.dtype.element_ty), boundary_check=(0, 1)) - - -@triton.heuristics({ - 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None, -}) -@triton.autotune( - configs=[ - triton.Config({}, num_warps=num_warps, num_stages=num_stages) - for num_warps in [4, 8] - for num_stages in [2, 3, 4] - ], - key=['BK', 'BT'], - **autotune_cache_kwargs, -) -@triton.jit(do_not_specialize=['T']) -def chunk_dplr_bwd_kernel_intra_tensorcore( - q, - k, - a, - b, - gi, - ge, - dAqk, - dAqb, - dAak, - dAab, - dq, - dk, - da, - db, - dqg, - dkg, - dag, - dbg, - dgk, - dgk_offset, - cu_seqlens, - chunk_indices, - scale: tl.constexpr, - T, - H: tl.constexpr, - K: tl.constexpr, - BT: tl.constexpr, - BC: tl.constexpr, - BK: tl.constexpr, - IS_VARLEN: tl.constexpr, - GATHER_SUPPORTED: tl.constexpr, -): - i_k, i_t, i_bh = tl.program_id(0), tl.program_id(1), tl.program_id(2) - i_b, i_h = i_bh // H, i_bh % H - - 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.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) - T_len = eos - bos - else: - bos, eos = (i_b * T).to(tl.int32), (i_b * T + T).to(tl.int32) - T_len = T - - if i_t * BT >= T_len: - return - - offset_base_k = (bos * H + i_h) * K - offset_base_attn = (bos * H + i_h) * BT - - valid_len = min(T_len - i_t * BT, BT) - mid_idx = valid_len // 2 - m_k = tl.arange(0, BK) + i_k * BK < K - p_offset = gi + offset_base_k + (i_t * BT + mid_idx) * K + tl.arange(0, BK) + i_k * BK - b_offset = tl.load(p_offset, mask=m_k, other=0.0).to(tl.float32) - - # Q, K, A, B, Gates: [BT, BK] - p_q = tl.make_block_ptr(q + offset_base_k, (T_len, K), (H*K, 1), (i_t*BT, i_k*BK), (BT, BK), (1, 0)) - p_k = tl.make_block_ptr(k + offset_base_k, (T_len, K), (H*K, 1), (i_t*BT, i_k*BK), (BT, BK), (1, 0)) - p_a = tl.make_block_ptr(a + offset_base_k, (T_len, K), (H*K, 1), (i_t*BT, i_k*BK), (BT, BK), (1, 0)) - p_b = tl.make_block_ptr(b + offset_base_k, (T_len, K), (H*K, 1), (i_t*BT, i_k*BK), (BT, BK), (1, 0)) - p_gi = tl.make_block_ptr(gi + offset_base_k, (T_len, K), (H*K, 1), (i_t*BT, i_k*BK), (BT, BK), (1, 0)) - p_ge = tl.make_block_ptr(ge + offset_base_k, (T_len, K), (H*K, 1), (i_t*BT, i_k*BK), (BT, BK), (1, 0)) - - b_q = tl.load(p_q, boundary_check=(0, 1)) - b_k = tl.load(p_k, boundary_check=(0, 1)) - b_a = tl.load(p_a, boundary_check=(0, 1)) - b_b = tl.load(p_b, boundary_check=(0, 1)) - b_gi_val = tl.load(p_gi, boundary_check=(0, 1)).to(tl.float32) - b_ge_val = tl.load(p_ge, boundary_check=(0, 1)).to(tl.float32) - - b_gi_shifted = b_gi_val - b_offset[None, :] - b_ge_shifted = b_ge_val - b_offset[None, :] - exp_gi_shifted = exp2(b_gi_shifted) - inv_exp_gi_shifted = exp2(-b_gi_shifted) - exp_ge_shifted = exp2(b_ge_shifted) - - q_ops = (b_q * exp_gi_shifted).to(tl.float32) - k_ops = (b_k * inv_exp_gi_shifted).to(tl.float32) - b_ops = (b_b * inv_exp_gi_shifted).to(tl.float32) - a_ops = (b_a * exp_ge_shifted).to(tl.float32) - - p_dAqk = tl.make_block_ptr(dAqk + offset_base_attn, (T_len, BT), (H*BT, 1), (i_t*BT, 0), (BT, BT), (1, 0)) - p_dAqb = tl.make_block_ptr(dAqb + offset_base_attn, (T_len, BT), (H*BT, 1), (i_t*BT, 0), (BT, BT), (1, 0)) - p_dAak = tl.make_block_ptr(dAak + offset_base_attn, (T_len, BT), (H*BT, 1), (i_t*BT, 0), (BT, BT), (1, 0)) - p_dAab = tl.make_block_ptr(dAab + offset_base_attn, (T_len, BT), (H*BT, 1), (i_t*BT, 0), (BT, BT), (1, 0)) - - b_dAqk = tl.load(p_dAqk, boundary_check=(0, 1)) - b_dAqb = tl.load(p_dAqb, boundary_check=(0, 1)) - b_dAak = tl.load(p_dAak, boundary_check=(0, 1)) - b_dAab = tl.load(p_dAab, boundary_check=(0, 1)) - - offs_n = tl.arange(0, BT) - offs_m = tl.arange(0, BT) - mask_inclusive = offs_m[:, None] >= offs_n[None, :] - mask_strict = offs_m[:, None] > offs_n[None, :] - - b_dAqk = tl.where(mask_inclusive, b_dAqk, 0.0).to(b_dAqk.dtype) - b_dAqb = tl.where(mask_inclusive, b_dAqb, 0.0).to(b_dAqb.dtype) - b_dAak = tl.where(mask_strict, b_dAak, 0.0).to(tl.float32) - b_dAab = tl.where(mask_strict, b_dAab, 0.0).to(tl.float32) - - # Intra-chunk gradients calculation - k_ops_h = k_ops.to(b_dAqk.dtype) - b_ops_h = b_ops.to(b_dAqb.dtype) - b_dq_intra = tl.dot(b_dAqk, k_ops_h) + tl.dot(b_dAqb, b_ops_h) - b_da_intra = tl.dot(b_dAak, k_ops) + tl.dot(b_dAab, b_ops) - - b_dAqk_T = tl.trans(b_dAqk) - b_dAqb_T = tl.trans(b_dAqb) - b_dAak_T = tl.trans(b_dAak) - b_dAab_T = tl.trans(b_dAab) - q_ops_h = q_ops.to(b_dAqk.dtype) - b_dk_intra = tl.dot(b_dAqk_T, q_ops_h) + tl.dot(b_dAak_T, a_ops) - b_db_intra = tl.dot(b_dAqb_T, q_ops_h) + tl.dot(b_dAab_T, a_ops) - - # Inter-chunk gradients loading - p_dqg = tl.make_block_ptr(dqg + offset_base_k, (T_len, K), (H*K, 1), (i_t*BT, i_k*BK), (BT, BK), (1, 0)) - p_dkg = tl.make_block_ptr(dkg + offset_base_k, (T_len, K), (H*K, 1), (i_t*BT, i_k*BK), (BT, BK), (1, 0)) - p_dag = tl.make_block_ptr(dag + offset_base_k, (T_len, K), (H*K, 1), (i_t*BT, i_k*BK), (BT, BK), (1, 0)) - p_dbg = tl.make_block_ptr(dbg + offset_base_k, (T_len, K), (H*K, 1), (i_t*BT, i_k*BK), (BT, BK), (1, 0)) - - b_dqg = tl.load(p_dqg, boundary_check=(0, 1)) - b_dag = tl.load(p_dag, boundary_check=(0, 1)) - b_dkg = tl.load(p_dkg, boundary_check=(0, 1)) - b_dbg = tl.load(p_dbg, boundary_check=(0, 1)) - - last_idx = min((i_t+1) * BT, T_len) - 1 - p_g_last = gi + offset_base_k + last_idx * H * K + tl.arange(0, BK) + i_k * BK - b_g_last = tl.load(p_g_last, mask=m_k, other=0.0) - - b_dq = b_dq_intra * exp_gi_shifted + b_dqg * exp2(b_gi_val) * scale - b_da = b_da_intra * exp_ge_shifted + b_dag * exp2(b_ge_val) - term_inter_stabilized = exp2(b_g_last[None, :] - b_gi_val) - b_dk = b_dk_intra * inv_exp_gi_shifted + b_dkg * term_inter_stabilized - b_db = b_db_intra * inv_exp_gi_shifted + b_dbg * term_inter_stabilized - - b_dgk = (b_dq * b_q + b_da * b_a - b_dk * b_k - b_db * b_b) - b_dgk_offset = b_da * b_a - - p_dq = tl.make_block_ptr(dq + offset_base_k, (T_len, K), (H*K, 1), (i_t*BT, i_k*BK), (BT, BK), (1, 0)) - p_dk = tl.make_block_ptr(dk + offset_base_k, (T_len, K), (H*K, 1), (i_t*BT, i_k*BK), (BT, BK), (1, 0)) - p_da = tl.make_block_ptr(da + offset_base_k, (T_len, K), (H*K, 1), (i_t*BT, i_k*BK), (BT, BK), (1, 0)) - p_db = tl.make_block_ptr(db + offset_base_k, (T_len, K), (H*K, 1), (i_t*BT, i_k*BK), (BT, BK), (1, 0)) - p_dgk = tl.make_block_ptr(dgk + offset_base_k, (T_len, K), (H*K, 1), (i_t*BT, i_k*BK), (BT, BK), (1, 0)) - p_dgk_offset = tl.make_block_ptr(dgk_offset + offset_base_k, (T_len, K), (H*K, 1), (i_t*BT, i_k*BK), (BT, BK), (1, 0)) - - tl.store(p_dq, b_dq.to(p_dq.dtype.element_ty), boundary_check=(0, 1)) - tl.store(p_dk, b_dk.to(p_dk.dtype.element_ty), boundary_check=(0, 1)) - tl.store(p_da, b_da.to(p_da.dtype.element_ty), boundary_check=(0, 1)) - tl.store(p_db, b_db.to(p_db.dtype.element_ty), boundary_check=(0, 1)) - tl.store(p_dgk, b_dgk.to(p_dgk.dtype.element_ty), boundary_check=(0, 1)) - tl.store(p_dgk_offset, b_dgk_offset.to(p_dgk_offset.dtype.element_ty), boundary_check=(0, 1)) - - -@triton.heuristics({ - 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None, -}) -@triton.autotune( - configs=[ - triton.Config({'BK': BK}, num_warps=num_warps, num_stages=num_stages) - for num_warps in NUM_WARPS_AUTOTUNE - for num_stages in [2, 3, 4] - for BK in [32, 64] - ], - key=['BK', 'BT', 'K'], - **autotune_cache_kwargs, -) -@triton.jit(do_not_specialize=['T']) -def chunk_dplr_bwd_dgk_kernel( - dgk, - dgk_offset, - dgk_last, - dgk_output, - cu_seqlens, - chunk_indices, - T, - H: tl.constexpr, - K: tl.constexpr, - BT: tl.constexpr, - BK: tl.constexpr, - IS_VARLEN: tl.constexpr, -): - i_t, i_k, i_bh = tl.program_id(0), tl.program_id(1), tl.program_id(2) - i_b, i_h = i_bh // H, i_bh % H - 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.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) - T = eos - bos - NT = tl.cdiv(T, BT) - else: - NT = tl.cdiv(T, BT) - i_tg = (i_b * NT + i_t).to(tl.int32) - bos, eos = (i_b * T).to(tl.int32), (i_b * T + T).to(tl.int32) - - stride_qk = H * K - dgk += (bos * H + i_h) * K - dgk_offset += (bos * H + i_h) * K - dgk_last += (i_tg * H + i_h) * K - dgk_output += (bos * H + i_h) * K - p_dgk_last = dgk_last + tl.arange(0, BK) + i_k * BK - m_k = tl.arange(0, BK) + i_k * BK < K - b_dgk_last = tl.load(p_dgk_last, mask=m_k, other=0) - p_dgk_offset = tl.make_block_ptr(dgk_offset, (T, K), (stride_qk, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) - p_dgk = tl.make_block_ptr(dgk, (T, K), (stride_qk, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) - b_dgk = tl.load(p_dgk, boundary_check=(0, 1)) - b_dgk_offset = tl.load(p_dgk_offset, boundary_check=(0, 1)) - # m_inv_cumsum = (tl.arange(0, BT)[:, None] <= tl.arange(0, BT)[None, :]).to(tl.float32) - # b_dgk_cumsum = tl.dot(m_inv_cumsum, b_dgk, allow_tf32=False) - b_dgk_cumsum = tl.cumsum(b_dgk, 0, reverse=True) - b_dgk_cumsum += b_dgk_last[None, :] - b_dgk_cumsum -= b_dgk_offset - p_dgk_output = tl.make_block_ptr(dgk_output, (T, K), (stride_qk, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) - tl.store(p_dgk_output, b_dgk_cumsum.to(p_dgk_output.dtype.element_ty), boundary_check=(0, 1)) - - -def chunk_dplr_bwd_dqk_intra( - q: torch.Tensor, - k: torch.Tensor, - a: torch.Tensor, - b: torch.Tensor, - gi: torch.Tensor, - ge: torch.Tensor, - dAqk: torch.Tensor, - dAqb: torch.Tensor, - dAak: torch.Tensor, - dAab: torch.Tensor, - dqg: torch.Tensor, - dkg: torch.Tensor, - dag: torch.Tensor, - dbg: torch.Tensor, - dgk_last: torch.Tensor, - scale: float = 1.0, - cu_seqlens: torch.LongTensor | None = None, - chunk_size: int = 64, - safe_gate: bool = False, - chunk_indices: torch.LongTensor | None = None, -): - B, T, H, K = q.shape - BT = chunk_size - BK = min(64, triton.next_power_of_2(K)) if check_shared_mem() else min(32, triton.next_power_of_2(K)) - - if chunk_indices is None: - chunk_indices = prepare_chunk_indices(cu_seqlens, BT) if cu_seqlens is not None else None - NT = triton.cdiv(T, BT) if cu_seqlens is None else len(chunk_indices) - NK = triton.cdiv(K, BK) - - dq = torch.empty_like(q) - dk = torch.empty_like(k) - da = torch.empty_like(a) - db = torch.empty_like(b) - dgk = torch.empty_like(gi, dtype=torch.float) - dgk_offset = torch.empty_like(gi, dtype=torch.float) - - grid = (NK, NT, B * H) - if safe_gate: - chunk_dplr_bwd_kernel_intra_func = chunk_dplr_bwd_kernel_intra_tensorcore - else: - chunk_dplr_bwd_kernel_intra_func = chunk_dplr_bwd_kernel_intra - chunk_dplr_bwd_kernel_intra_func[grid]( - q=q, - k=k, - a=a, - b=b, - gi=gi, - ge=ge, - dAqk=dAqk, - dAqb=dAqb, - dAak=dAak, - dAab=dAab, - dq=dq, - dk=dk, - dgk=dgk, - dgk_offset=dgk_offset, - dqg=dqg, - dkg=dkg, - dag=dag, - dbg=dbg, - da=da, - db=db, - cu_seqlens=cu_seqlens, - chunk_indices=chunk_indices, - scale=scale, - T=T, - H=H, - K=K, - BT=BT, - BC=BT, - BK=BK, - GATHER_SUPPORTED=IS_GATHER_SUPPORTED, - ) - - dgk_output = torch.empty_like(dgk) - - def grid(meta): return (NT, triton.cdiv(K, meta['BK']), B * H) - chunk_dplr_bwd_dgk_kernel[grid]( - dgk=dgk, - dgk_offset=dgk_offset, - dgk_last=dgk_last, - dgk_output=dgk_output, - cu_seqlens=cu_seqlens, - chunk_indices=chunk_indices, - T=T, - H=H, - K=K, - BT=BT, - ) - return dq, dk, da, db, dgk_output diff --git a/build/torch-xpu/ops/generalized_delta_rule/dplr/chunk_A_fwd.py b/build/torch-xpu/ops/generalized_delta_rule/dplr/chunk_A_fwd.py deleted file mode 100644 index b37433aad93359578bb2708ae36b8a2a02dd89b2..0000000000000000000000000000000000000000 --- a/build/torch-xpu/ops/generalized_delta_rule/dplr/chunk_A_fwd.py +++ /dev/null @@ -1,371 +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 -import triton.language as tl - -from ....ops.utils import prepare_chunk_indices -from ....ops.utils.op import exp2, gather -from ....utils import IS_AMD, IS_GATHER_SUPPORTED, autotune_cache_kwargs - -NUM_WARPS_AUTOTUNE = [2, 4, 8, 16] if IS_AMD else [2, 4, 8, 16, 32] - - -@triton.heuristics({ - 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None, -}) -@triton.autotune( - configs=[ - triton.Config({}, num_warps=num_warps, num_stages=num_stages) - for num_warps in NUM_WARPS_AUTOTUNE - for num_stages in [2, 3, 4] - ], - key=['BK', 'BT'], - **autotune_cache_kwargs, -) -@triton.jit(do_not_specialize=['T']) -def chunk_dplr_fwd_A_kernel_intra_sub_intra( - q, - k, - a, - b, - gi, - ge, - qg, - kg, - ag, - bg, - Aqk, - Aqb, - Aab, - Aak, - cu_seqlens, - chunk_indices, - scale: tl.constexpr, - T, - H: tl.constexpr, - K: tl.constexpr, - BT: tl.constexpr, - BC: tl.constexpr, - BK: tl.constexpr, - IS_VARLEN: tl.constexpr, - GATHER_SUPPORTED: tl.constexpr, -): - i_t, i_b, i_h = 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.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) - T = eos - bos - else: - bos, eos = i_b * T, i_b * T + T - - if i_t * BT >= T: - return - - o_i = tl.arange(0, BC) - o_k = tl.arange(0, BK) - m_k = o_k < K - m_A = (i_t * BT + tl.arange(0, BC)) < T - last_idx = min((i_t+1) * BT, T) - 1 - o_A = (bos + i_t * BT + tl.arange(0, BC)) * H*BT + i_h * BT - p_q = tl.make_block_ptr(q + (bos * H + i_h) * K, (T, K), (H*K, 1), (i_t * BT, 0), (BC, BK), (1, 0)) - p_k = tl.make_block_ptr(k + (bos * H + i_h) * K, (T, K), (H*K, 1), (i_t * BT, 0), (BC, BK), (1, 0)) - p_a = tl.make_block_ptr(a + (bos * H + i_h) * K, (T, K), (H*K, 1), (i_t * BT, 0), (BC, BK), (1, 0)) - p_b = tl.make_block_ptr(b + (bos * H + i_h) * K, (T, K), (H*K, 1), (i_t * BT, 0), (BC, BK), (1, 0)) - p_gi = tl.make_block_ptr(gi + (bos * H + i_h) * K, (T, K), (H*K, 1), (i_t * BT, 0), (BC, BK), (1, 0)) - p_ge = tl.make_block_ptr(ge + (bos * H + i_h) * K, (T, K), (H*K, 1), (i_t * BT, 0), (BC, BK), (1, 0)) - p_g_last = gi + (bos * H + i_h) * K + last_idx * H * K + tl.arange(0, BK) - b_g_last = tl.load(p_g_last, mask=m_k, other=0) - p_qg = tl.make_block_ptr(qg + (bos * H + i_h) * K, (T, K), (H*K, 1), (i_t * BT, 0), (BC, BK), (1, 0)) - p_kg = tl.make_block_ptr(kg + (bos * H + i_h) * K, (T, K), (H*K, 1), (i_t * BT, 0), (BC, BK), (1, 0)) - p_ag = tl.make_block_ptr(ag + (bos * H + i_h) * K, (T, K), (H*K, 1), (i_t * BT, 0), (BC, BK), (1, 0)) - p_bg = tl.make_block_ptr(bg + (bos * H + i_h) * K, (T, K), (H*K, 1), (i_t * BT, 0), (BC, BK), (1, 0)) - - b_q = tl.load(p_q, boundary_check=(0, 1)) - b_q = b_q * scale - b_k = tl.load(p_k, boundary_check=(0, 1)) - b_a = tl.load(p_a, boundary_check=(0, 1)) - b_b = tl.load(p_b, boundary_check=(0, 1)) - b_gi = tl.load(p_gi, boundary_check=(0, 1)).to(tl.float32) - b_ge = tl.load(p_ge, boundary_check=(0, 1)).to(tl.float32) - - # deal with decay term. - g_exp = exp2(b_gi) - g_exp_inv = exp2(-b_gi + b_g_last[None, :]) - b_qg = b_q * g_exp - b_kg = b_k * g_exp_inv - b_bg = b_b * g_exp_inv - b_ag = b_a * exp2(b_ge) - tl.store(p_qg, b_qg.to(p_qg.dtype.element_ty, fp_downcast_rounding="rtne"), boundary_check=(0, 1)) - tl.store(p_bg, b_bg.to(p_bg.dtype.element_ty, fp_downcast_rounding="rtne"), boundary_check=(0, 1)) - tl.store(p_ag, b_ag.to(p_ag.dtype.element_ty, fp_downcast_rounding="rtne"), boundary_check=(0, 1)) - tl.store(p_kg, b_kg.to(p_kg.dtype.element_ty, fp_downcast_rounding="rtne"), boundary_check=(0, 1)) - # tl.debug_barrier() - - b_q = b_q.to(b_k.dtype) - # inner attn - for j in range(0, min(BC, T - i_t * BT)): - # a trick to index the j-th row of b_k, b_g, b_b - if GATHER_SUPPORTED: - row_idx = tl.full([1, BK], j, dtype=tl.int16) - # [1, BK] - b_k_j = gather(b_k, row_idx, axis=0) - b_gk_j = gather(b_gi, row_idx, axis=0) - b_b_j = gather(b_b, row_idx, axis=0) - else: - mask = tl.arange(0, BC) == j - b_k_j = tl.sum(tl.where(mask[:, None], b_k, 0), 0)[None, :] - b_gk_j = tl.sum(tl.where(mask[:, None], b_gi, 0), 0)[None, :] - b_b_j = tl.sum(tl.where(mask[:, None], b_b, 0), 0)[None, :] - tmp = exp2(b_gi - b_gk_j) - b_A_qk = tl.sum(b_q * b_k_j * tmp, 1) - m_i = (o_i >= j).to(tl.float32) - b_A_qk = b_A_qk * m_i - b_A_qb = tl.sum(b_q * b_b_j * tmp, 1) - b_A_qb = b_A_qb * m_i - tmp2 = exp2(b_ge - b_gk_j) - b_A_ak = tl.sum(b_a * b_k_j * tmp2, 1) - m_i2 = (o_i > j).to(tl.float32) - b_A_ak = b_A_ak * m_i2 - b_A_ab = tl.sum(b_a * b_b_j * tmp2, 1) - b_A_ab = b_A_ab * m_i2 - - tl.store(Aqk + o_A + j, b_A_qk.to(dtype=Aqk.dtype.element_ty, fp_downcast_rounding="rtne"), mask=m_A) - tl.store(Aqb + o_A + j, b_A_qb.to(dtype=Aqb.dtype.element_ty, fp_downcast_rounding="rtne"), mask=m_A) - tl.store(Aab + o_A + j, b_A_ab.to(dtype=Aqb.dtype.element_ty, fp_downcast_rounding="rtne"), mask=m_A) - tl.store(Aak + o_A + j, b_A_ak.to(dtype=Aqk.dtype.element_ty, fp_downcast_rounding="rtne"), mask=m_A) - - -@triton.heuristics({ - 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None, -}) -@triton.autotune( - configs=[ - triton.Config({}, num_warps=num_warps, num_stages=num_stages) - for num_warps in [4, 8] - for num_stages in [2, 3] - ], - key=['BK', 'BT'], - **autotune_cache_kwargs, -) -@triton.jit(do_not_specialize=['T']) -def chunk_dplr_fwd_A_kernel_intra_tensorcore( - q, - k, - a, - b, - gi, - ge, - qg, - kg, - ag, - bg, - Aqk, - Aqb, - Aab, - Aak, - cu_seqlens, - chunk_indices, - scale: tl.constexpr, - T, - H: tl.constexpr, - K: tl.constexpr, - BT: tl.constexpr, - BC: tl.constexpr, - BK: tl.constexpr, - IS_VARLEN: tl.constexpr, - GATHER_SUPPORTED: tl.constexpr, -): - i_t, i_b, i_h = 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.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) - T_len = eos - bos - else: - bos = i_b * T - T_len = T - - if i_t * BT >= T_len: - return - - # Compute base offset for all tensors - offset_base = (bos * H + i_h) * K - - # Load the current chunk of Q, K, A, B and their gates - p_q = tl.make_block_ptr(q + offset_base, (T_len, K), (H*K, 1), (i_t * BT, 0), (BT, BK), (1, 0)) - p_k = tl.make_block_ptr(k + offset_base, (T_len, K), (H*K, 1), (i_t * BT, 0), (BT, BK), (1, 0)) - p_a = tl.make_block_ptr(a + offset_base, (T_len, K), (H*K, 1), (i_t * BT, 0), (BT, BK), (1, 0)) - p_b = tl.make_block_ptr(b + offset_base, (T_len, K), (H*K, 1), (i_t * BT, 0), (BT, BK), (1, 0)) - p_gi = tl.make_block_ptr(gi + offset_base, (T_len, K), (H*K, 1), (i_t * BT, 0), (BT, BK), (1, 0)) - p_ge = tl.make_block_ptr(ge + offset_base, (T_len, K), (H*K, 1), (i_t * BT, 0), (BT, BK), (1, 0)) - - b_q = tl.load(p_q, boundary_check=(0, 1)) - b_k = tl.load(p_k, boundary_check=(0, 1)) - b_a = tl.load(p_a, boundary_check=(0, 1)) - b_b = tl.load(p_b, boundary_check=(0, 1)) - b_gi_val = tl.load(p_gi, boundary_check=(0, 1)).to(tl.float32) - b_ge_val = tl.load(p_ge, boundary_check=(0, 1)).to(tl.float32) - - # Calculate the index of the middle element of the valid part of the chunk - valid_len = min(T_len - i_t * BT, BT) - mid_idx = valid_len // 2 - - # Load the offset vector from Global Memory - # p_offset points to gi[i_t*BT + mid_idx, :] - m_k = tl.arange(0, BK) < K - p_offset = gi + offset_base + (i_t * BT + mid_idx) * K + tl.arange(0, BK) - b_offset = tl.load(p_offset, mask=m_k, other=0.0).to(tl.float32) - - # Apply offset to gate values - # These operations broadcast [BK] to [BT, BK] - b_gi_val = b_gi_val - b_offset[None, :] - b_ge_val = b_ge_val - b_offset[None, :] - - # Apply decay factors (now numerically safe) - # q_term ~ exp2(gi - offset) - # k_term ~ exp2(-gi + offset) - exp_gi = exp2(b_gi_val) - inv_exp_gi = exp2(-b_gi_val) - exp_ge = exp2(b_ge_val) - - b_q = (b_q * scale).to(tl.float32) - - # Compute gated operands for matrix multiplication - # q, k in l2 norm, < 1, i - q_ops = (b_q * exp_gi).to(tl.float32) - k_ops = (b_k * inv_exp_gi).to(tl.float32) - b_ops = (b_b * inv_exp_gi).to(tl.float32) - a_ops = (b_a * exp_ge).to(tl.float32) - - # Load gate values at the last position for inter-chunk decay - last_idx = min((i_t+1) * BT, T_len) - 1 - p_g_last = gi + offset_base + last_idx * H * K + tl.arange(0, BK) - b_g_last = tl.load(p_g_last, mask=m_k, other=0.0) - - exp_offset = exp2(b_offset) - b_g_centered = b_g_last - b_offset - exp_g_centered = exp2(b_g_centered) - - # Create pointers for writing - p_qg = tl.make_block_ptr(qg + offset_base, (T_len, K), (H*K, 1), (i_t * BT, 0), (BT, BK), (1, 0)) - p_kg = tl.make_block_ptr(kg + offset_base, (T_len, K), (H*K, 1), (i_t * BT, 0), (BT, BK), (1, 0)) - p_ag = tl.make_block_ptr(ag + offset_base, (T_len, K), (H*K, 1), (i_t * BT, 0), (BT, BK), (1, 0)) - p_bg = tl.make_block_ptr(bg + offset_base, (T_len, K), (H*K, 1), (i_t * BT, 0), (BT, BK), (1, 0)) - - # Store gated Q and A - tl.store(p_qg, (q_ops * exp_offset[None, :]).to(p_qg.dtype.element_ty), boundary_check=(0, 1)) - tl.store(p_ag, (a_ops * exp_offset[None, :]).to(p_ag.dtype.element_ty), boundary_check=(0, 1)) - - # Store gated K and B - b_kg_g = k_ops * exp_g_centered[None, :] - b_bg_g = b_ops * exp_g_centered[None, :] - tl.store(p_kg, b_kg_g.to(p_kg.dtype.element_ty), boundary_check=(0, 1)) - tl.store(p_bg, b_bg_g.to(p_bg.dtype.element_ty), boundary_check=(0, 1)) - - # Transpose K and B for dot product - k_ops_t = tl.trans(k_ops) - b_ops_t = tl.trans(b_ops) - - # Compute intra-chunk attention using TensorCores - q_ops_h = q_ops.to(b_q.dtype) - b_A_qk = tl.dot(q_ops_h, k_ops_t.to(b_q.dtype)) - b_A_qb = tl.dot(q_ops_h, b_ops_t.to(b_q.dtype)) - b_A_ak = tl.dot(a_ops, k_ops_t) - b_A_ab = tl.dot(a_ops, b_ops_t) - - # Create causal masks - offs_n = tl.arange(0, BT) - offs_m = tl.arange(0, BT) - mask_inclusive = offs_m[:, None] >= offs_n[None, :] - mask_strict = offs_m[:, None] > offs_n[None, :] - - # Apply causal masking - b_A_qk = tl.where(mask_inclusive, b_A_qk, 0.0) - b_A_qb = tl.where(mask_inclusive, b_A_qb, 0.0) - b_A_ak = tl.where(mask_strict, b_A_ak, 0.0) - b_A_ab = tl.where(mask_strict, b_A_ab, 0.0) - - # Store the intra-chunk attention matrices - offset_out_base = (bos * H + i_h) * BT - - p_Aqk = tl.make_block_ptr(Aqk + offset_out_base, (T_len, BT), (H*BT, 1), (i_t * BT, 0), (BT, BT), (1, 0)) - p_Aqb = tl.make_block_ptr(Aqb + offset_out_base, (T_len, BT), (H*BT, 1), (i_t * BT, 0), (BT, BT), (1, 0)) - p_Aak = tl.make_block_ptr(Aak + offset_out_base, (T_len, BT), (H*BT, 1), (i_t * BT, 0), (BT, BT), (1, 0)) - p_Aab = tl.make_block_ptr(Aab + offset_out_base, (T_len, BT), (H*BT, 1), (i_t * BT, 0), (BT, BT), (1, 0)) - - tl.store(p_Aqk, b_A_qk.to(p_Aqk.dtype.element_ty), boundary_check=(0, 1)) - tl.store(p_Aqb, b_A_qb.to(p_Aqb.dtype.element_ty), boundary_check=(0, 1)) - tl.store(p_Aak, b_A_ak.to(p_Aak.dtype.element_ty), boundary_check=(0, 1)) - tl.store(p_Aab, b_A_ab.to(p_Aab.dtype.element_ty), boundary_check=(0, 1)) - - -def chunk_dplr_fwd_intra( - q: torch.Tensor, - k: torch.Tensor, - a: torch.Tensor, - b: torch.Tensor, - gi: torch.Tensor, - ge: torch.Tensor, - scale: float, - chunk_size: int, - safe_gate: bool = False, - cu_seqlens: torch.LongTensor | None = None, - chunk_indices: torch.LongTensor | None = None, -): - B, T, H, K = k.shape - BT = chunk_size - - if chunk_indices is None: - chunk_indices = prepare_chunk_indices(cu_seqlens, BT) if cu_seqlens is not None else None - NT = triton.cdiv(T, BT) if cu_seqlens is None else len(chunk_indices) - - Aqk = q.new_empty(B, T, H, BT, dtype=q.dtype) - Aqb = q.new_empty(B, T, H, BT, dtype=q.dtype) - # involving matrix inverse and it'd be better to use float here. - Aab = q.new_empty(B, T, H, BT, dtype=torch.float) - Aak = q.new_empty(B, T, H, BT, dtype=torch.float) - - grid = (NT, B, H) - BK = max(triton.next_power_of_2(K), 16) - qg = torch.empty_like(q) - kg = torch.empty_like(k, dtype=q.dtype) - ag = torch.empty_like(a, dtype=q.dtype) - bg = torch.empty_like(b, dtype=q.dtype) - if safe_gate: - chunk_dplr_fwd_A_kernel_intra_func = chunk_dplr_fwd_A_kernel_intra_tensorcore - else: - chunk_dplr_fwd_A_kernel_intra_func = chunk_dplr_fwd_A_kernel_intra_sub_intra - chunk_dplr_fwd_A_kernel_intra_func[grid]( - q=q, - k=k, - a=a, - b=b, - gi=gi, - ge=ge, - Aqk=Aqk, - Aqb=Aqb, - Aab=Aab, - Aak=Aak, - qg=qg, - kg=kg, - ag=ag, - bg=bg, - cu_seqlens=cu_seqlens, - chunk_indices=chunk_indices, - scale=scale, - T=T, - H=H, - K=K, - BT=BT, - BC=BT, - BK=BK, - GATHER_SUPPORTED=IS_GATHER_SUPPORTED, - ) - return Aab, Aqk, Aak, Aqb, qg, kg, ag, bg diff --git a/build/torch-xpu/ops/generalized_delta_rule/dplr/chunk_h_bwd.py b/build/torch-xpu/ops/generalized_delta_rule/dplr/chunk_h_bwd.py deleted file mode 100644 index 167a651114c69f1992ccedd6fcda7c46737e09c7..0000000000000000000000000000000000000000 --- a/build/torch-xpu/ops/generalized_delta_rule/dplr/chunk_h_bwd.py +++ /dev/null @@ -1,179 +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 -import triton.language as tl - -from ....ops.utils import prepare_chunk_indices, prepare_chunk_offsets -from ....ops.utils.op import exp2 -from ....utils import IS_AMD, autotune_cache_kwargs, check_shared_mem - -NUM_WARPS_AUTOTUNE = [2, 4, 8, 16] if IS_AMD else [2, 4, 8, 16, 32] - - -@triton.heuristics({ - 'USE_FINAL_STATE_GRADIENT': lambda args: args['dht'] is not None, - 'USE_INITIAL_STATE': lambda args: args['dh0'] is not None, - 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None, -}) -@triton.autotune( - configs=[ - triton.Config({}, num_warps=num_warps, num_stages=num_stages) - for num_warps in NUM_WARPS_AUTOTUNE - for num_stages in [2, 3, 4] - ], - key=['BT', 'BK', 'BV', 'V'], - **autotune_cache_kwargs, -) -@triton.jit(do_not_specialize=['T']) -def chunk_dplr_bwd_kernel_dhu( - qg, - bg, - w, - gk, - dht, - dh0, - do, - dh, - dv, - dv2, - cu_seqlens, - chunk_offsets, - T, - H: tl.constexpr, - K: tl.constexpr, - V: tl.constexpr, - BT: tl.constexpr, - BC: tl.constexpr, - BK: tl.constexpr, - BV: tl.constexpr, - USE_FINAL_STATE_GRADIENT: tl.constexpr, - USE_INITIAL_STATE: tl.constexpr, - IS_VARLEN: tl.constexpr, -): - i_k, i_v, i_nh = tl.program_id(0), tl.program_id(1), tl.program_id(2) - i_n, i_h = i_nh // H, i_nh % H - if IS_VARLEN: - bos, eos = tl.load(cu_seqlens + i_n).to(tl.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) - T = eos - bos - NT = tl.cdiv(T, BT) - boh = tl.load(chunk_offsets + i_n).to(tl.int32) - else: - bos, eos = i_n * T, i_n * T + T - NT = tl.cdiv(T, BT) - boh = i_n * NT - - # [BK, BV] - b_dh = tl.zeros([BK, BV], dtype=tl.float32) - if USE_FINAL_STATE_GRADIENT: - p_dht = tl.make_block_ptr(dht + i_nh * K*V, (K, V), (V, 1), (i_k * BK, i_v * BV), (BK, BV), (1, 0)) - b_dh += tl.load(p_dht, boundary_check=(0, 1)) - - mask_k = tl.arange(0, BK) < K - for i_t in range(NT - 1, -1, -1): - p_dh = tl.make_block_ptr(dh + ((boh+i_t) * H + i_h) * K*V, (K, V), (V, 1), (i_k * BK, i_v * BV), (BK, BV), (1, 0)) - tl.store(p_dh, b_dh.to(p_dh.dtype.element_ty), boundary_check=(0, 1)) - b_dh_tmp = tl.zeros([BK, BV], dtype=tl.float32) - for i_c in range(tl.cdiv(BT, BC) - 1, -1, -1): - p_qg = tl.make_block_ptr(qg+(bos*H+i_h)*K, (K, T), (1, H*K), (i_k * BK, i_t * BT + i_c * BC), (BK, BC), (0, 1)) - p_bg = tl.make_block_ptr(bg+(bos*H+i_h)*K, (T, K), (H*K, 1), (i_t * BT + i_c * BC, i_k * BK), (BC, BK), (1, 0)) - p_w = tl.make_block_ptr(w+(bos*H+i_h)*K, (K, T), (1, H*K), (i_k * BK, i_t * BT + i_c * BC), (BK, BC), (0, 1)) - p_dv = tl.make_block_ptr(dv+(bos*H+i_h)*V, (T, V), (H*V, 1), (i_t*BT + i_c * BC, i_v * BV), (BC, BV), (1, 0)) - p_do = tl.make_block_ptr(do+(bos*H+i_h)*V, (T, V), (H*V, 1), (i_t*BT + i_c * BC, i_v * BV), (BC, BV), (1, 0)) - p_dv2 = tl.make_block_ptr(dv2+(bos*H+i_h)*V, (T, V), (H*V, 1), (i_t*BT + i_c * BC, i_v * BV), (BC, BV), (1, 0)) - # [BK, BT] - b_qg = tl.load(p_qg, boundary_check=(0, 1)) - # [BT, BK] - b_bg = tl.load(p_bg, boundary_check=(0, 1)) - b_w = tl.load(p_w, boundary_check=(0, 1)) - # [BT, V] - b_do = tl.load(p_do, boundary_check=(0, 1)) - b_dv = tl.load(p_dv, boundary_check=(0, 1)) - b_dv2 = b_dv + tl.dot(b_bg, b_dh.to(b_bg.dtype)) - tl.store(p_dv2, b_dv2.to(p_dv.dtype.element_ty), boundary_check=(0, 1)) - # [BK, BV] - b_dh_tmp += tl.dot(b_qg, b_do.to(b_qg.dtype)) - b_dh_tmp += tl.dot(b_w, b_dv2.to(b_qg.dtype)) - last_idx = min((i_t + 1) * BT, T) - 1 - bg_last = tl.load(gk + ((bos + last_idx) * H + i_h) * K + tl.arange(0, BK), mask=mask_k) - b_dh *= exp2(bg_last)[:, None] - b_dh += b_dh_tmp - - if USE_INITIAL_STATE: - p_dh0 = tl.make_block_ptr(dh0 + i_nh * K*V, (K, V), (V, 1), (i_k * BK, i_v * BV), (BK, BV), (1, 0)) - tl.store(p_dh0, b_dh.to(p_dh0.dtype.element_ty), boundary_check=(0, 1)) - - -def chunk_dplr_bwd_dhu( - qg: torch.Tensor, - bg: torch.Tensor, - w: torch.Tensor, - gk: torch.Tensor, - h0: torch.Tensor, - dht: torch.Tensor | None, - do: torch.Tensor, - dv: torch.Tensor, - cu_seqlens: torch.LongTensor | None = None, - chunk_size: int = 64, - chunk_indices: torch.LongTensor | None = None, -) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: - B, T, H, K, V = *qg.shape, do.shape[-1] - BT = chunk_size - BK = max(triton.next_power_of_2(K), 16) - assert BK <= 256, "current kernel does not support head dimension being larger than 256." - # H100 - if check_shared_mem('hopper', qg.device.index): - BV = 64 - BC = 64 if K <= 128 else 32 - elif check_shared_mem('ampere', qg.device.index): # A100 - BV = 32 - BC = 32 - else: # Etc: 4090 - BV = 16 - BC = 16 - - if chunk_indices is None: - chunk_indices = prepare_chunk_indices(cu_seqlens, BT) if cu_seqlens is not None else None - # N: the actual number of sequences in the batch with either equal or variable lengths - if cu_seqlens is None: - N, NT, chunk_offsets = B, triton.cdiv(T, BT), None - else: - N, NT, chunk_offsets = len(cu_seqlens) - 1, len(chunk_indices), prepare_chunk_offsets(cu_seqlens, BT) - - BC = min(BT, BC) - NK, NV = triton.cdiv(K, BK), triton.cdiv(V, BV) - assert NK == 1, 'NK > 1 is not supported because it involves time-consuming synchronization' - - dh = qg.new_empty(B, NT, H, K, V) - dh0 = torch.empty_like(h0, dtype=torch.float32) if h0 is not None else None - dv2 = torch.zeros_like(dv) - - grid = (NK, NV, N * H) - chunk_dplr_bwd_kernel_dhu[grid]( - qg=qg, - bg=bg, - w=w, - gk=gk, - dht=dht, - dh0=dh0, - do=do, - dh=dh, - dv=dv, - dv2=dv2, - cu_seqlens=cu_seqlens, - chunk_offsets=chunk_offsets, - T=T, - H=H, - K=K, - V=V, - BT=BT, - BC=BC, - BK=BK, - BV=BV, - ) - return dh, dh0, dv2 diff --git a/build/torch-xpu/ops/generalized_delta_rule/dplr/chunk_h_fwd.py b/build/torch-xpu/ops/generalized_delta_rule/dplr/chunk_h_fwd.py deleted file mode 100644 index 050ae739a549bc02830bba1e6f97f607d2ac5406..0000000000000000000000000000000000000000 --- a/build/torch-xpu/ops/generalized_delta_rule/dplr/chunk_h_fwd.py +++ /dev/null @@ -1,179 +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 -import triton.language as tl - -from ....ops.utils import prepare_chunk_indices, prepare_chunk_offsets -from ....ops.utils.op import exp2 -from ....utils import IS_AMD, autotune_cache_kwargs, check_shared_mem - -NUM_WARPS_AUTOTUNE = [2, 4, 8, 16] if IS_AMD else [2, 4, 8, 16, 32] - - -@triton.heuristics({ - 'USE_INITIAL_STATE': lambda args: args['h0'] is not None, - 'STORE_FINAL_STATE': lambda args: args['ht'] is not None, - 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None, -}) -@triton.autotune( - configs=[ - triton.Config({}, num_warps=num_warps, num_stages=num_stages) - for num_warps in NUM_WARPS_AUTOTUNE - for num_stages in [2, 3, 4] - ], - key=['BT', 'BK', 'BV'], - **autotune_cache_kwargs, -) -@triton.jit(do_not_specialize=['T']) -def chunk_dplr_fwd_kernel_h( - kg, - v, - w, - bg, - u, - v_new, - gk, - h, - h0, - ht, - cu_seqlens, - chunk_offsets, - T, - H: tl.constexpr, - K: tl.constexpr, - V: tl.constexpr, - BT: tl.constexpr, - BC: tl.constexpr, - BK: tl.constexpr, - BV: tl.constexpr, - USE_INITIAL_STATE: tl.constexpr, - STORE_FINAL_STATE: tl.constexpr, - IS_VARLEN: tl.constexpr, -): - i_k, i_v, i_nh = tl.program_id(0), tl.program_id(1), tl.program_id(2) - i_n, i_h = i_nh // H, i_nh % H - if IS_VARLEN: - bos, eos = tl.load(cu_seqlens + i_n).to(tl.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) - T = eos - bos - NT = tl.cdiv(T, BT) - boh = tl.load(chunk_offsets + i_n).to(tl.int32) - else: - bos, eos = i_n * T, i_n * T + T - NT = tl.cdiv(T, BT) - boh = i_n * NT - o_k = i_k * BK + tl.arange(0, BK) - - # [BK, BV] - b_h = tl.zeros([BK, BV], dtype=tl.float32) - if USE_INITIAL_STATE: - p_h0 = tl.make_block_ptr(h0 + i_nh * K*V, (K, V), (V, 1), (i_k * BK, i_v * BV), (BK, BV), (1, 0)) - b_h = tl.load(p_h0, boundary_check=(0, 1)).to(tl.float32) - - for i_t in range(NT): - p_h = tl.make_block_ptr(h + ((boh + i_t) * H + i_h) * K*V, (K, V), (V, 1), (i_k * BK, i_v * BV), (BK, BV), (1, 0)) - tl.store(p_h, b_h.to(p_h.dtype.element_ty), boundary_check=(0, 1)) - - b_hc = tl.zeros([BK, BV], dtype=tl.float32) - # since we need to make all DK in the SRAM. we face serve SRAM memory burden. By subchunking we allievate such burden - for i_c in range(tl.cdiv(min(BT, T - i_t * BT), BC)): - p_kg = tl.make_block_ptr(kg+(bos*H+i_h)*K, (K, T), (1, H*K), (i_k * BK, i_t * BT + i_c * BC), (BK, BC), (0, 1)) - p_bg = tl.make_block_ptr(bg+(bos*H+i_h)*K, (K, T), (1, H*K), (i_k * BK, i_t * BT + i_c * BC), (BK, BC), (0, 1)) - p_w = tl.make_block_ptr(w+(bos*H+i_h)*K, (T, K), (H*K, 1), (i_t * BT + i_c * BC, i_k * BK), (BC, BK), (1, 0)) - p_v = tl.make_block_ptr(v+(bos*H+i_h)*V, (T, V), (H*V, 1), (i_t * BT + i_c * BC, i_v * BV), (BC, BV), (1, 0)) - p_u = tl.make_block_ptr(u+(bos*H+i_h)*V, (T, V), (H*V, 1), (i_t * BT + i_c * BC, i_v * BV), (BC, BV), (1, 0)) - p_v_new = tl.make_block_ptr(v_new+(bos*H+i_h)*V, (T, V), (H*V, 1), (i_t*BT+i_c*BC, i_v * BV), (BC, BV), (1, 0)) - # [BK, BC] - b_kg = tl.load(p_kg, boundary_check=(0, 1)) - b_v = tl.load(p_v, boundary_check=(0, 1)) - b_w = tl.load(p_w, boundary_check=(0, 1)) - b_bg = tl.load(p_bg, boundary_check=(0, 1)) - b_v2 = tl.dot(b_w, b_h.to(b_w.dtype)) + tl.load(p_u, boundary_check=(0, 1)) - b_hc += tl.dot(b_kg, b_v) - b_hc += tl.dot(b_bg.to(b_hc.dtype), b_v2) - tl.store(p_v_new, b_v2.to(p_v_new.dtype.element_ty), boundary_check=(0, 1)) - - last_idx = min((i_t + 1) * BT, T) - 1 - b_g_last = tl.load(gk + (bos + last_idx) * H*K + i_h * K + o_k, mask=o_k < K).to(tl.float32) - b_h *= exp2(b_g_last[:, None]) - b_h += b_hc - - if STORE_FINAL_STATE: - p_ht = tl.make_block_ptr(ht + i_nh * K*V, (K, V), (V, 1), (i_k * BK, i_v * BV), (BK, BV), (1, 0)) - tl.store(p_ht, b_h.to(p_ht.dtype.element_ty, fp_downcast_rounding="rtne"), boundary_check=(0, 1)) - - -def chunk_dplr_fwd_h( - kg: torch.Tensor, - v: torch.Tensor, - w: torch.Tensor, - u: torch.Tensor, - bg: torch.Tensor, - gk: torch.Tensor, - initial_state: torch.Tensor | None = None, - output_final_state: bool = False, - cu_seqlens: torch.LongTensor | None = None, - chunk_size: int = 64, - chunk_indices: torch.LongTensor | None = None, -) -> tuple[torch.Tensor, torch.Tensor]: - B, T, H, K, V = *kg.shape, u.shape[-1] - BT = chunk_size - - if chunk_indices is None: - chunk_indices = prepare_chunk_indices(cu_seqlens, BT) if cu_seqlens is not None else None - # N: the actual number of sequences in the batch with either equal or variable lengths - if cu_seqlens is None: - N, NT, chunk_offsets = B, triton.cdiv(T, BT), None - else: - N, NT, chunk_offsets = len(cu_seqlens) - 1, len(chunk_indices), prepare_chunk_offsets(cu_seqlens, BT) - BK = max(triton.next_power_of_2(K), 16) - assert BK <= 256, "current kernel does not support head dimension larger than 256." - # H100 can have larger block size - - if check_shared_mem('hopper', kg.device.index): - BV = 64 - BC = 64 if K <= 128 else 32 - elif check_shared_mem('ampere', kg.device.index): # A100 - BV = 32 - BC = 32 - else: - BV = 16 - BC = 16 - - BC = min(BT, BC) - NK = triton.cdiv(K, BK) - NV = triton.cdiv(V, BV) - assert NK == 1, 'NK > 1 is not supported because it involves time-consuming synchronization' - - h = kg.new_empty(B, NT, H, K, V) - final_state = kg.new_empty(N, H, K, V, dtype=torch.float32) if output_final_state else None - v_new = torch.empty_like(u) - grid = (NK, NV, N * H) - chunk_dplr_fwd_kernel_h[grid]( - kg=kg, - v=v, - w=w, - bg=bg, - u=u, - v_new=v_new, - h=h, - gk=gk, - h0=initial_state, - ht=final_state, - cu_seqlens=cu_seqlens, - chunk_offsets=chunk_offsets, - T=T, - H=H, - K=K, - V=V, - BT=BT, - BC=BC, - BK=BK, - BV=BV, - ) - return h, v_new, final_state diff --git a/build/torch-xpu/ops/generalized_delta_rule/dplr/chunk_o_bwd.py b/build/torch-xpu/ops/generalized_delta_rule/dplr/chunk_o_bwd.py deleted file mode 100644 index be0b1b270bde765e42864198ead92903ce2189ed..0000000000000000000000000000000000000000 --- a/build/torch-xpu/ops/generalized_delta_rule/dplr/chunk_o_bwd.py +++ /dev/null @@ -1,438 +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 -import triton.language as tl - -from ....ops.utils import prepare_chunk_indices -from ....ops.utils.op import exp2 -from ....utils import IS_AMD, autotune_cache_kwargs, check_shared_mem - -NUM_WARPS_AUTOTUNE = [2, 4, 8, 16] if IS_AMD else [2, 4, 8, 16, 32] - -BK_LIST = [32, 64, 128] if check_shared_mem() else [16, 32] - - -@triton.heuristics({ - 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None, -}) -@triton.autotune( - configs=[ - triton.Config({}, num_warps=num_warps, num_stages=num_stages) - for num_warps in NUM_WARPS_AUTOTUNE - for num_stages in [2, 3, 4] - ], - key=['BV', 'BT'], - **autotune_cache_kwargs, -) -@triton.jit(do_not_specialize=['T']) -def chunk_dplr_bwd_kernel_dAu( - v, - do, - v_new, - A_qb, - dA_qk, - dA_qb, - dv_new, - cu_seqlens, - chunk_indices, - scale: tl.constexpr, - T, - H: tl.constexpr, - V: tl.constexpr, - BT: tl.constexpr, - BV: tl.constexpr, - IS_VARLEN: tl.constexpr, -): - i_t, i_bh = tl.program_id(0), tl.program_id(1) - i_b, i_h = i_bh // H, i_bh % H - 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.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) - else: - bos, eos = i_b * T, i_b * T + T - T = eos - bos - - b_dA_qk = tl.zeros([BT, BT], dtype=tl.float32) - b_dA_qb = tl.zeros([BT, BT], dtype=tl.float32) - - p_A_qb = tl.make_block_ptr(A_qb + (bos * H + i_h) * BT, (T, BT), (H*BT, 1), (i_t * BT, 0), (BT, BT), (1, 0)) - - b_A_qb = tl.load(p_A_qb, boundary_check=(0, 1)) - # causal mask - b_A_qb = tl.where(tl.arange(0, BT)[:, None] >= tl.arange(0, BT)[None, :], b_A_qb, 0.).to(b_A_qb.dtype) - - for i_v in range(tl.cdiv(V, BV)): - p_do = tl.make_block_ptr(do + (bos*H + i_h) * V, (T, V), (H*V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) - p_v = tl.make_block_ptr(v + (bos*H + i_h) * V, (V, T), (1, H*V), (i_v * BV, i_t * BT), (BV, BT), (0, 1)) - p_v_new = tl.make_block_ptr(v_new + (bos*H + i_h) * V, (V, T), (1, H*V), (i_v * BV, i_t * BT), (BV, BT), (0, 1)) - p_dv_new = tl.make_block_ptr(dv_new + (bos*H + i_h) * V, (T, V), (H*V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) - b_v = tl.load(p_v, boundary_check=(0, 1)) - b_do = tl.load(p_do, boundary_check=(0, 1)) - b_v_new = tl.load(p_v_new, boundary_check=(0, 1)) - b_dA_qk += tl.dot(b_do, b_v) - b_dA_qb += tl.dot(b_do, b_v_new) - b_dv_new = tl.dot(tl.trans(b_A_qb), b_do) - # for recurrent - tl.store(p_dv_new, b_dv_new.to(p_dv_new.dtype.element_ty), boundary_check=(0, 1)) - - p_dA_qk = tl.make_block_ptr(dA_qk + (bos * H + i_h) * BT, (T, BT), (H*BT, 1), (i_t * BT, 0), (BT, BT), (1, 0)) - p_dA_qb = tl.make_block_ptr(dA_qb + (bos * H + i_h) * BT, (T, BT), (H*BT, 1), (i_t * BT, 0), (BT, BT), (1, 0)) - m_s = tl.arange(0, BT)[:, None] >= tl.arange(0, BT)[None, :] - b_dA_qk = tl.where(m_s, b_dA_qk * scale, 0.) - tl.store(p_dA_qk, b_dA_qk.to(p_dA_qk.dtype.element_ty), boundary_check=(0, 1)) - b_dA_qb = tl.where(m_s, b_dA_qb * scale, 0.) - tl.store(p_dA_qb, b_dA_qb.to(p_dA_qb.dtype.element_ty), boundary_check=(0, 1)) - - -@triton.heuristics({ - 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None, -}) -@triton.autotune( - configs=[ - triton.Config({}, num_warps=num_warps, num_stages=num_stages) - for num_warps in NUM_WARPS_AUTOTUNE - for num_stages in [2, 3, 4] - ], - key=['BT', 'BK', 'BV'], - **autotune_cache_kwargs, -) -@triton.jit -def chunk_dplr_bwd_o_kernel( - v, - v_new, - h, - do, - dh, - dk, - db, - w, - dq, - dv, - dw, - gk, - dgk_last, - k, - b, - cu_seqlens, - chunk_indices, - T, - H: tl.constexpr, - K: tl.constexpr, - V: tl.constexpr, - BT: tl.constexpr, - BK: tl.constexpr, - BV: tl.constexpr, - IS_VARLEN: tl.constexpr, -): - i_k, i_t, i_bh = tl.program_id(0), tl.program_id(1), tl.program_id(2) - i_b, i_h = i_bh // H, i_bh % H - - 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.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) - T = eos - bos - NT = tl.cdiv(T, BT) - else: - NT = tl.cdiv(T, BT) - i_tg = i_b * NT + i_t - bos, eos = i_b * T, i_b * T + T - - # offset calculation - v += (bos * H + i_h) * V - v_new += (bos * H + i_h) * V - do += (bos * H + i_h) * V - h += (i_tg * H + i_h) * K * V - dh += (i_tg * H + i_h) * K * V - dk += (bos * H + i_h) * K - k += (bos * H + i_h) * K - db += (bos * H + i_h) * K - b += (bos * H + i_h) * K - dw += (bos * H + i_h) * K - dv += (bos * H + i_h) * V - dq += (bos * H + i_h) * K - w += (bos * H + i_h) * K - - dgk_last += (i_tg * H + i_h) * K - gk += (bos * H + i_h) * K - - stride_qk = H*K - stride_vo = H*V - - b_dq = tl.zeros([BT, BK], dtype=tl.float32) - b_dk = tl.zeros([BT, BK], dtype=tl.float32) - b_dw = tl.zeros([BT, BK], dtype=tl.float32) - b_db = tl.zeros([BT, BK], dtype=tl.float32) - b_dgk_last = tl.zeros([BK], dtype=tl.float32) - - for i_v in range(tl.cdiv(V, BV)): - p_v = tl.make_block_ptr(v, (T, V), (stride_vo, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) - p_v_new = tl.make_block_ptr(v_new, (T, V), (stride_vo, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) - p_do = tl.make_block_ptr(do, (T, V), (stride_vo, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) - p_h = tl.make_block_ptr(h, (V, K), (1, V), (i_v * BV, i_k * BK), (BV, BK), (0, 1)) - p_dh = tl.make_block_ptr(dh, (V, K), (1, V), (i_v * BV, i_k * BK), (BV, BK), (0, 1)) - # [BT, BV] - b_v = tl.load(p_v, boundary_check=(0, 1)) - b_v_new = tl.load(p_v_new, boundary_check=(0, 1)) - b_do = tl.load(p_do, boundary_check=(0, 1)) - # [BV, BK] - b_h = tl.load(p_h, boundary_check=(0, 1)) - b_dh = tl.load(p_dh, boundary_check=(0, 1)) - b_dgk_last += tl.sum((b_h * b_dh).to(tl.float32), axis=0) - - # [BT, BV] @ [BV, BK] -> [BT, BK] - b_dq += tl.dot(b_do, b_h.to(b_do.dtype)) - # [BT, BV] @ [BV, BK] -> [BT, BK] - b_dk += tl.dot(b_v, b_dh.to(b_v.dtype)) - b_db += tl.dot(b_v_new, b_dh.to(b_v_new.dtype)) - p_dv = tl.make_block_ptr(dv, (T, V), (stride_vo, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) - b_dv = tl.load(p_dv, boundary_check=(0, 1)) - b_dw += tl.dot(b_dv.to(b_v.dtype), b_h.to(b_v.dtype)) - - m_k = (i_k*BK+tl.arange(0, BK)) < K - last_idx = min(i_t * BT + BT, T) - 1 - b_gk_last = tl.load(gk + last_idx * stride_qk + i_k*BK + tl.arange(0, BK), mask=m_k, other=float('-inf')) - b_dgk_last *= exp2(b_gk_last) - p_k = tl.make_block_ptr(k, (T, K), (stride_qk, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) - p_b = tl.make_block_ptr(b, (T, K), (stride_qk, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) - b_k = tl.load(p_k, boundary_check=(0, 1)) - b_b = tl.load(p_b, boundary_check=(0, 1)) - b_dgk_last += tl.sum(b_k * b_dk, axis=0) - b_dgk_last += tl.sum(b_b * b_db, axis=0) - tl.store(dgk_last + tl.arange(0, BK) + i_k * BK, b_dgk_last, mask=m_k) - - p_dw = tl.make_block_ptr(dw, (T, K), (stride_qk, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) - p_dk = tl.make_block_ptr(dk, (T, K), (stride_qk, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) - p_db = tl.make_block_ptr(db, (T, K), (stride_qk, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) - p_dq = tl.make_block_ptr(dq, (T, K), (stride_qk, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) - tl.store(p_dw, b_dw.to(p_dw.dtype.element_ty), boundary_check=(0, 1)) - tl.store(p_dk, b_dk.to(p_dk.dtype.element_ty), boundary_check=(0, 1)) - tl.store(p_db, b_db.to(p_db.dtype.element_ty), boundary_check=(0, 1)) - tl.store(p_dq, b_dq.to(p_dq.dtype.element_ty), boundary_check=(0, 1)) - - -@triton.heuristics({ - 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None, -}) -@triton.autotune( - configs=[ - triton.Config({'BK': BK, 'BV': BV}, num_warps=num_warps, num_stages=num_stages) - for num_warps in NUM_WARPS_AUTOTUNE - for num_stages in [2, 3, 4] - for BK in BK_LIST - for BV in BK_LIST - ], - key=['BT'], - **autotune_cache_kwargs, -) -@triton.jit -def chunk_dplr_bwd_kernel_dv( - A_qk, - kg, - do, - dv, - dh, - cu_seqlens, - chunk_indices, - T, - H: tl.constexpr, - K: tl.constexpr, - V: tl.constexpr, - BT: tl.constexpr, - BK: tl.constexpr, - BV: tl.constexpr, - IS_VARLEN: tl.constexpr, -): - i_v, i_t, i_bh = tl.program_id(0), tl.program_id(1), tl.program_id(2) - i_b, i_h = i_bh // H, i_bh % H - 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.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) - T = eos - bos - NT = tl.cdiv(T, BT) - else: - NT = tl.cdiv(T, BT) - i_tg = i_b * NT + i_t - bos, eos = i_b * T, i_b * T + T - - b_dv = tl.zeros([BT, BV], dtype=tl.float32) - - # offset calculation - A_qk += (bos * H + i_h) * BT - do += (bos * H + i_h) * V - dv += (bos * H + i_h) * V - kg += (bos * H + i_h) * K - dh += (i_tg * H + i_h) * K*V - - stride_qk = H*K - stride_vo = H*V - stride_A = H*BT - - for i_k in range(tl.cdiv(K, BK)): - p_dh = tl.make_block_ptr(dh, (K, V), (V, 1), (i_k * BK, i_v * BV), (BK, BV), (1, 0)) - p_kg = tl.make_block_ptr(kg, (T, K), (stride_qk, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) - b_dh = tl.load(p_dh, boundary_check=(0, 1)) - b_kg = tl.load(p_kg, boundary_check=(0, 1)) - b_dv += tl.dot(b_kg, b_dh.to(b_kg.dtype)) - - p_Aqk = tl.make_block_ptr(A_qk, (BT, T), (1, stride_A), (0, i_t * BT), (BT, BT), (0, 1)) - b_A = tl.where(tl.arange(0, BT)[:, None] <= tl.arange(0, BT)[None, :], tl.load(p_Aqk, boundary_check=(0, 1)), 0) - p_do = tl.make_block_ptr(do, (T, V), (stride_vo, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) - p_dv = tl.make_block_ptr(dv, (T, V), (stride_vo, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) - b_do = tl.load(p_do, boundary_check=(0, 1)) - b_dv += tl.dot(b_A.to(b_do.dtype), b_do) - tl.store(p_dv, b_dv.to(p_dv.dtype.element_ty), boundary_check=(0, 1)) - - -def chunk_dplr_bwd_dv( - A_qk: torch.Tensor, - kg: torch.Tensor, - do: torch.Tensor, - dh: torch.Tensor, - cu_seqlens: torch.LongTensor | None = None, - chunk_size: int = 64, - chunk_indices: torch.LongTensor | None = None, -) -> torch.Tensor: - B, T, H, K, V = *kg.shape, do.shape[-1] - BT = chunk_size - - if chunk_indices is None: - chunk_indices = prepare_chunk_indices(cu_seqlens, BT) if cu_seqlens is not None else None - NT = triton.cdiv(T, BT) if cu_seqlens is None else len(chunk_indices) - - dv = torch.empty_like(do) - - def grid(meta): return (triton.cdiv(V, meta['BV']), NT, B * H) - chunk_dplr_bwd_kernel_dv[grid]( - A_qk=A_qk, - kg=kg, - do=do, - dv=dv, - dh=dh, - cu_seqlens=cu_seqlens, - chunk_indices=chunk_indices, - T=T, - H=H, - K=K, - V=V, - BT=BT, - ) - return dv - - -def chunk_dplr_bwd_o( - k: torch.Tensor, - b: torch.Tensor, - v: torch.Tensor, - v_new: torch.Tensor, - gk: torch.Tensor, - do: torch.Tensor, - h: torch.Tensor, - dh: torch.Tensor, - dv: torch.Tensor, - w: torch.Tensor, - cu_seqlens: torch.LongTensor | None = None, - chunk_size: int = 64, - scale: float = 1.0, - chunk_indices: torch.LongTensor | None = None, -) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: - - B, T, H, K, V = *w.shape, v.shape[-1] - - BT = chunk_size - if chunk_indices is None: - chunk_indices = prepare_chunk_indices(cu_seqlens, BT) if cu_seqlens is not None else None - NT = triton.cdiv(T, BT) if cu_seqlens is None else len(chunk_indices) - - BK = min(max(triton.next_power_of_2(K), 16), 64) if check_shared_mem() else min(triton.next_power_of_2(K), 32) - BV = min(max(triton.next_power_of_2(V), 16), 64) if check_shared_mem() else min(triton.next_power_of_2(K), 32) - NK = triton.cdiv(K, BK) - dq = torch.empty_like(k) - dk = torch.empty_like(k) - dw = torch.empty_like(w) - db = torch.empty_like(b) - grid = (NK, NT, B * H) - - dgk_last = torch.empty(B, NT, H, K, dtype=torch.float, device=w.device) - - chunk_dplr_bwd_o_kernel[grid]( - k=k, - b=b, - v=v, - v_new=v_new, - h=h, - do=do, - dh=dh, - dq=dq, - dk=dk, - db=db, - dgk_last=dgk_last, - w=w, - dv=dv, - dw=dw, - gk=gk, - cu_seqlens=cu_seqlens, - chunk_indices=chunk_indices, - T=T, - H=H, - K=K, - V=V, - BT=BT, - BK=BK, - BV=BV, - ) - return dq, dk, dw, db, dgk_last - - -def chunk_dplr_bwd_dAu( - v: torch.Tensor, - v_new: torch.Tensor, - do: torch.Tensor, - A_qb: torch.Tensor, - scale: float, - cu_seqlens: torch.LongTensor | None = None, - chunk_size: int = 64, - chunk_indices: torch.LongTensor | None = None, -) -> torch.Tensor: - B, T, H, V = v.shape - BT = chunk_size - if chunk_indices is None: - chunk_indices = prepare_chunk_indices(cu_seqlens, BT) if cu_seqlens is not None else None - NT = triton.cdiv(T, BT) if cu_seqlens is None else len(chunk_indices) - - if check_shared_mem('ampere'): # A100 - BV = min(triton.next_power_of_2(V), 128) - elif check_shared_mem('ada'): # 4090 - BV = min(max(triton.next_power_of_2(V), 16), 64) - else: - BV = min(triton.next_power_of_2(V), 32) - - grid = (NT, B * H) - dA_qk = torch.empty(B, T, H, BT, dtype=torch.float, device=v.device) - dA_qb = torch.empty(B, T, H, BT, dtype=torch.float, device=v.device) - dv_new = torch.empty_like(v_new) - chunk_dplr_bwd_kernel_dAu[grid]( - v=v, - do=do, - v_new=v_new, - A_qb=A_qb, - dA_qk=dA_qk, - dA_qb=dA_qb, - dv_new=dv_new, - cu_seqlens=cu_seqlens, - chunk_indices=chunk_indices, - scale=scale, - T=T, - H=H, - V=V, - BT=BT, - BV=BV, - ) - return dv_new, dA_qk, dA_qb diff --git a/build/torch-xpu/ops/generalized_delta_rule/dplr/chunk_o_fwd.py b/build/torch-xpu/ops/generalized_delta_rule/dplr/chunk_o_fwd.py deleted file mode 100644 index 2de1977e78885a4e432ab0ce3c3323efe80adfa1..0000000000000000000000000000000000000000 --- a/build/torch-xpu/ops/generalized_delta_rule/dplr/chunk_o_fwd.py +++ /dev/null @@ -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 - -import torch -import triton -import triton.language as tl - -from ....ops.utils import prepare_chunk_indices -from ....utils import IS_AMD, autotune_cache_kwargs, check_shared_mem - -NUM_WARPS_AUTOTUNE = [2, 4, 8, 16] if IS_AMD else [2, 4, 8, 16, 32] - -BK_LIST = [32, 64, 128] if check_shared_mem() else [16, 32] - - -@triton.heuristics({ - 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None, -}) -@triton.autotune( - configs=[ - triton.Config({'BK': BK, 'BV': BV}, num_warps=num_warps, num_stages=num_stages) - for BK in BK_LIST - for BV in BK_LIST - for num_warps in NUM_WARPS_AUTOTUNE - for num_stages in [2, 3, 4] - ], - key=['BT'], - **autotune_cache_kwargs, -) -@triton.jit(do_not_specialize=['T']) -def chunk_dplr_fwd_kernel_o( - qg, - v, - v_new, - A_qk, - A_qb, - h, - o, - cu_seqlens, - chunk_indices, - T, - H: tl.constexpr, - K: tl.constexpr, - V: tl.constexpr, - BT: tl.constexpr, - BK: tl.constexpr, - BV: tl.constexpr, - IS_VARLEN: tl.constexpr, -): - i_v, i_t, i_bh = tl.program_id(0), tl.program_id(1), tl.program_id(2) - i_b, i_h = i_bh // H, i_bh % H - - 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.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) - T = eos - bos - NT = tl.cdiv(T, BT) - else: - NT = tl.cdiv(T, BT) - i_tg = i_b * NT + i_t - bos, eos = i_b * T, i_b * T + T - - b_o = tl.zeros([BT, BV], dtype=tl.float32) - for i_k in range(tl.cdiv(K, BK)): - p_qg = tl.make_block_ptr(qg + (bos * H + i_h) * K, (T, K), (H*K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) - p_h = tl.make_block_ptr(h + (i_tg * H + i_h) * K*V, (K, V), (V, 1), (i_k * BK, i_v * BV), (BK, BV), (1, 0)) - b_qg = tl.load(p_qg, boundary_check=(0, 1)) - b_h = tl.load(p_h, boundary_check=(0, 1)) - b_o += tl.dot(b_qg, b_h) - - p_Aqk = tl.make_block_ptr(A_qk + (bos * H + i_h) * BT, (T, BT), (H*BT, 1), (i_t * BT, 0), (BT, BT), (1, 0)) - p_Aqb = tl.make_block_ptr(A_qb + (bos * H + i_h) * BT, (T, BT), (H*BT, 1), (i_t * BT, 0), (BT, BT), (1, 0)) - p_v = tl.make_block_ptr(v + (bos * H + i_h) * V, (T, V), (H*V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) - p_v_new = tl.make_block_ptr(v_new + (bos * H + i_h) * V, (T, V), (H*V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) - p_o = tl.make_block_ptr(o + (bos * H + i_h) * V, (T, V), (H*V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) - - m_s = tl.arange(0, BT)[:, None] >= tl.arange(0, BT)[None, :] - b_Aqk = tl.load(p_Aqk, boundary_check=(0, 1)) - b_Aqb = tl.load(p_Aqb, boundary_check=(0, 1)) - b_Aqk = tl.where(m_s, b_Aqk, 0) - b_Aqb = tl.where(m_s, b_Aqb, 0) - b_v = tl.load(p_v, boundary_check=(0, 1)) - b_v_new = tl.load(p_v_new, boundary_check=(0, 1)) - b_o = b_o + tl.dot(b_Aqk.to(b_v.dtype), b_v) + tl.dot(b_Aqb.to(b_v_new.dtype), b_v_new) - tl.store(p_o, b_o.to(p_o.dtype.element_ty), boundary_check=(0, 1)) - - -def chunk_dplr_fwd_o( - qg: torch.Tensor, - v: torch.Tensor, - v_new: torch.Tensor, - A_qk: torch.Tensor, - A_qb: torch.Tensor, - h: torch.Tensor, - cu_seqlens: torch.LongTensor | None = None, - chunk_size: int = 64, - chunk_indices: torch.LongTensor | None = None, -) -> torch.Tensor: - B, T, H, K, V = *qg.shape, v.shape[-1] - BT = chunk_size - - if chunk_indices is None: - chunk_indices = prepare_chunk_indices(cu_seqlens, BT) if cu_seqlens is not None else None - NT = triton.cdiv(T, BT) if cu_seqlens is None else len(chunk_indices) - - o = torch.empty_like(v) - def grid(meta): return (triton.cdiv(V, meta['BV']), NT, B * H) - chunk_dplr_fwd_kernel_o[grid]( - qg=qg, - v=v, - v_new=v_new, - A_qk=A_qk, - A_qb=A_qb, - h=h, - o=o, - cu_seqlens=cu_seqlens, - chunk_indices=chunk_indices, - T=T, - H=H, - K=K, - V=V, - BT=BT, - ) - return o diff --git a/build/torch-xpu/ops/generalized_delta_rule/dplr/fused_recurrent.py b/build/torch-xpu/ops/generalized_delta_rule/dplr/fused_recurrent.py deleted file mode 100644 index ccaaf7652486de50eebd246cf5a6b7640823ab42..0000000000000000000000000000000000000000 --- a/build/torch-xpu/ops/generalized_delta_rule/dplr/fused_recurrent.py +++ /dev/null @@ -1,270 +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 -import triton.language as tl - -from ....ops.utils.op import exp -from ....utils import autocast_custom_bwd, autocast_custom_fwd, autotune_cache_kwargs, input_guard - - -@triton.heuristics({ - 'USE_INITIAL_STATE': lambda args: args['h0'] is not None, - 'STORE_FINAL_STATE': lambda args: args['ht'] is not None, - 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None, -}) -@triton.autotune( - configs=[ - triton.Config({'BV': BV}, num_warps=num_warps, num_stages=num_stages) - for BV in [16, 32, 64] - for num_warps in [2, 4, 8, 16] - for num_stages in [2, 3, 4] - ], - key=['BK'], - **autotune_cache_kwargs, -) -@triton.jit(do_not_specialize=['T']) -def fused_recurrent_dplr_delta_rule_fwd_kernel( - q, - k, - v, - a, - b, - gk, - o, - h0, - ht, - cu_seqlens, - scale, - T, - B: tl.constexpr, - H: tl.constexpr, - K: tl.constexpr, - V: tl.constexpr, - BK: tl.constexpr, - BV: tl.constexpr, - REVERSE: tl.constexpr, - USE_INITIAL_STATE: tl.constexpr, - STORE_FINAL_STATE: tl.constexpr, - IS_VARLEN: tl.constexpr, -): - i_v, i_nh = tl.program_id(0).to(tl.int64), tl.program_id(1).to(tl.int64) - i_n, i_h = i_nh // H, i_nh % H - - if IS_VARLEN: - bos, eos = tl.load(cu_seqlens + i_n).to(tl.int64), tl.load(cu_seqlens + i_n + 1).to(tl.int64) - T = eos - bos - else: - bos, eos = i_n * T, i_n * T + T - - o_k = tl.arange(0, BK) - o_v = i_v * BV + tl.arange(0, BV) - p_q = q + (bos + ((T - 1) if REVERSE else 0)) * H*K + i_h * K + o_k - p_k = k + (bos + ((T - 1) if REVERSE else 0)) * H*K + i_h * K + o_k - p_a = a + (bos + ((T - 1) if REVERSE else 0)) * H*K + i_h * K + o_k - p_b = b + (bos + ((T - 1) if REVERSE else 0)) * H*K + i_h * K + o_k - p_gk = gk + (bos + ((T - 1) if REVERSE else 0)) * H*K + i_h * K + o_k - p_v = v + (bos + ((T - 1) if REVERSE else 0)) * H*V + i_h * V + o_v - p_o = o + (bos + ((T - 1) if REVERSE else 0)) * H*V + i_h * V + o_v - - mask_k = o_k < K - mask_v = o_v < V - mask_h = mask_k[:, None] & mask_v[None, :] - b_h = tl.zeros([BK, BV], dtype=tl.float32) - - if USE_INITIAL_STATE: - p_h0 = h0 + i_nh * K*V + o_k[:, None] * V + o_v - b_h += tl.load(p_h0, mask=mask_h, other=0).to(tl.float32) - - for _ in range(0, T): - b_q = tl.load(p_q, mask=mask_k, other=0).to(tl.float32) * scale - b_k = tl.load(p_k, mask=mask_k, other=0).to(tl.float32) - b_a = tl.load(p_a, mask=mask_k, other=0).to(tl.float32) - b_b = tl.load(p_b, mask=mask_k, other=0).to(tl.float32) - b_gk = tl.load(p_gk, mask=mask_k, other=0).to(tl.float32) - b_v = tl.load(p_v, mask=mask_v, other=0).to(tl.float32) - - b_h = exp(b_gk)[:, None] * b_h + b_b[:, None] * tl.sum(b_a[:, None] * b_h, 0)[None, :] - b_h += b_k[:, None] * b_v[None, :] - b_o = tl.sum(b_h * b_q[:, None], 0) - - tl.store(p_o, b_o.to(p_o.dtype.element_ty), mask=mask_v) - p_q += (-1 if REVERSE else 1) * H*K - p_k += (-1 if REVERSE else 1) * H*K - p_a += (-1 if REVERSE else 1) * H*K - p_b += (-1 if REVERSE else 1) * H*K - p_gk += (-1 if REVERSE else 1) * H*K - p_v += (-1 if REVERSE else 1) * H*V - p_o += (-1 if REVERSE else 1) * H*V - - if STORE_FINAL_STATE: - p_ht = ht + i_nh * K*V + o_k[:, None] * V + o_v - tl.store(p_ht, b_h.to(p_ht.dtype.element_ty), mask=mask_h) - - -def fused_recurrent_dplr_delta_rule_fwd( - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - a: torch.Tensor, - b: torch.Tensor, - gk: torch.Tensor, - scale: float | None = 1.0, - initial_state: torch.Tensor | None = None, - output_final_state: bool = False, - reverse: bool = False, - cu_seqlens: torch.LongTensor | None = None, -): - B, T, H, K, V = *k.shape, v.shape[-1] - N = B if cu_seqlens is None else len(cu_seqlens) - 1 - BK = triton.next_power_of_2(K) - - h0 = initial_state - ht = q.new_empty(N, H, K, V, dtype=torch.float32) if output_final_state else None - o = torch.empty_like(v) - - def grid(meta): return (triton.cdiv(V, meta['BV']), N * H) - fused_recurrent_dplr_delta_rule_fwd_kernel[grid]( - q=q, - k=k, - v=v, - a=a, - b=b, - gk=gk, - o=o, - h0=h0, - ht=ht, - cu_seqlens=cu_seqlens, - scale=scale, - T=T, - B=B, - H=H, - K=K, - V=V, - BK=BK, - REVERSE=reverse, - ) - return o, ht - - -class FusedRecurrentDPLRDeltaRuleFunction(torch.autograd.Function): - - @staticmethod - @input_guard - @autocast_custom_fwd - def forward( - ctx, - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - a: torch.Tensor, - b: torch.Tensor, - gk: torch.Tensor, - scale: float | None = None, - initial_state: torch.Tensor | None = None, - output_final_state: bool = False, - reverse: bool = False, - cu_seqlens: torch.LongTensor | None = None, - ): - o, ht = fused_recurrent_dplr_delta_rule_fwd( - q=q, - k=k, - v=v, - a=a, - b=b, - gk=gk, - scale=scale, - initial_state=initial_state, - output_final_state=output_final_state, - reverse=reverse, - cu_seqlens=cu_seqlens, - ) - return o, ht - - @staticmethod - @input_guard - @autocast_custom_bwd - def backward(ctx, do, dht): - raise NotImplementedError( - "Backward pass for fused_recurrent_dplr_delta_rule is not implemented and will not be supported. " - "This kernel is only for inference. " - "For training, please use `chunk_dplr_delta_rule`.", - ) - - -def fused_recurrent_dplr_delta_rule( - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - a: torch.Tensor, - b: torch.Tensor, - gk: torch.Tensor, - scale: float | None = None, - initial_state: torch.Tensor | None = None, - output_final_state: bool = False, - reverse: bool = False, - cu_seqlens: torch.Tensor | None = None, -) -> tuple[torch.Tensor, torch.Tensor]: - r""" - This function computes the recurrence S_t = S_t @ (Diag(g_t) + a_t b_t^T) + v_t k_t^T in a recurrent manner. - - Args: - q (torch.Tensor): - queries of shape `[B, T, H, K]`. - k (torch.Tensor): - keys of shape `[B, T, H, K]`. - v (torch.Tensor): - values of shape `[B, T, H, V]`. - a (torch.Tensor): - a of shape `[B, T, H, K]`. - b (torch.Tensor): - b of shape `[B, T, H, K]`. - gk (torch.Tensor): - gk of shape `[B, T, H, K]`. decay term in log space! - scale (Optional[float]): - Scale factor for the RetNet attention scores. - If not provided, it will default to `1 / sqrt(K)`. Default: `None`. - initial_state (Optional[torch.Tensor]): - Initial state of shape `[N, H, K, V]` for `N` input sequences. - For equal-length input sequences, `N` equals the batch size `B`. - Default: `None`. - output_final_state (Optional[bool]): - Whether to output the final state of shape `[N, H, K, V]`. Default: `False`. - reverse (Optional[bool]): - If `True`, process the state passing in reverse order. Default: `False`. - cu_seqlens (Optional[torch.Tensor]): - Cumulative sequence lengths of shape `[N + 1]` used for variable-length training, - consistent with the FlashAttention API. - """ - if cu_seqlens is not None: - if q.shape[0] != 1: - raise ValueError( - f"The batch size is expected to be 1 rather than {q.shape[0]} when using `cu_seqlens`." - f"Please flatten variable-length inputs before processing.", - ) - if initial_state is not None and initial_state.shape[0] != len(cu_seqlens) - 1: - raise ValueError( - f"The number of initial states is expected to be equal to the number of input sequences, " - f"i.e., {len(cu_seqlens) - 1} rather than {initial_state.shape[0]}.", - ) - if scale is None: - scale = q.shape[-1] ** -0.5 - o, final_state = FusedRecurrentDPLRDeltaRuleFunction.apply( - q, - k, - v, - a, - b, - gk, - scale, - initial_state, - output_final_state, - reverse, - cu_seqlens, - ) - return o, final_state diff --git a/build/torch-xpu/ops/generalized_delta_rule/dplr/naive.py b/build/torch-xpu/ops/generalized_delta_rule/dplr/naive.py deleted file mode 100644 index 0bf902a24a4a91b3bde909902d1923f714fb839a..0000000000000000000000000000000000000000 --- a/build/torch-xpu/ops/generalized_delta_rule/dplr/naive.py +++ /dev/null @@ -1,101 +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 -from einops import rearrange - -# S_t = S_t @ (I + alpha_t beta_t^T) + v_t k_t^T -# q, k, alpha, beta [B, H, L, D_K] -# v [B, H, L, D_V] - - -def dplr_recurrence(q, k, v, alpha, beta, gk, initial_state=None, output_final_state=True): - orig_dtype = q.dtype - b, h, l, d_k = q.shape - q, k, v, beta, gk = map(lambda x: x.float(), [q, k, v, beta, gk]) - d_v = v.shape[-1] - o = torch.zeros_like(v) - S = torch.zeros(b, h, d_k, d_v).to(v) - q = q * (d_k ** -0.5) - - if initial_state is not None: - S += initial_state - - for i in range(l): - _k = k[:, :, i] - _q = q[:, :, i] - _v = v[:, :, i] - _alpha = alpha[:, :, i].clone() - _beta = beta[:, :, i].clone() - _kv = _k[..., None] * _v[..., None, :] + (S.clone() * _alpha[..., None]).sum(-2, keepdim=True) * _beta[..., None] - S = S.clone() * gk[:, :, i].exp()[..., None] + _kv - o[:, :, i] = torch.einsum('bhd,bhdm->bhm', _q, S) - S = None if output_final_state is False else S - return o.to(orig_dtype), S - - -def dplr_chunkwise(q, k, v, alpha, beta, gk, initial_state=None, output_final_state=True, chunk_size=32): - b, h, l, d_k = q.shape - d_v = v.shape[-1] - q = q * (d_k ** -0.5) - v = v - assert l % chunk_size == 0 - - S = k.new_zeros(b, h, d_k, d_v).to(q) - if initial_state is not None: - S += initial_state - - # note that diagonal is masked. - mask = torch.triu(torch.ones(chunk_size, chunk_size, dtype=torch.bool, device=q.device), diagonal=0) - q, k, v, alpha, beta, gk = map(lambda x: rearrange(x, 'b h (n c) d -> b h n c d', - c=chunk_size).float(), [q, k, v, alpha, beta, gk]) - - gk_cumsum = gk.cumsum(-2) - - # v2 = (alpha @ k.transpose(-1, -2)).masked_fill_(mask, 0) @ v - A_ab = torch.zeros(b, h, l // chunk_size, chunk_size, chunk_size).to(q.device) - A_qk = torch.zeros(b, h, l // chunk_size, chunk_size, chunk_size).to(q.device) - A_ak = torch.zeros(b, h, l // chunk_size, chunk_size, chunk_size).to(q.device) - A_qb = torch.zeros(b, h, l // chunk_size, chunk_size, chunk_size).to(q.device) - - for i in range(chunk_size): - alpha_i = alpha[:, :, :, i, None] - q_i = q[:, :, :, i, None] - gk_i = gk_cumsum[:, :, :, i, None] - mask = (torch.arange(chunk_size) <= i).to(q.device) - attn_i = (gk_i - gk_cumsum).masked_fill(~mask.unsqueeze(-1), float('-inf')).exp() - A_qk[:, :, :, i, :] = (q_i * k * attn_i).sum(-1).clone() - A_qb[:, :, :, i, :] = (q_i * beta * attn_i).sum(-1).clone() - mask = (torch.arange(chunk_size) < i).to(q.device) - # shift by one. - attn_i = (gk_i - gk[:, :, :, i, None] - gk_cumsum).masked_fill(~mask.unsqueeze(-1), float('-inf')).exp() - A_ab[:, :, :, i, :] = (alpha_i * beta * attn_i).sum(-1).clone() - A_ak[:, :, :, i, :] = (alpha_i * k * attn_i).sum(-1).clone() - - A_ab = A_ab - for i in range(1, chunk_size): - A_ab[..., i, :i] = A_ab[..., i, :i].clone() + (A_ab[..., i, :, None].clone() * A_ab[..., :, :i].clone()).sum(-2) - - A_ab = A_ab + torch.eye(chunk_size, dtype=torch.float, device=q.device) - u = A_ab @ (A_ak @ v) - w = A_ab @ ((gk_cumsum-gk).exp() * alpha) - - o = torch.zeros_like(v) - mask = torch.triu(torch.ones(chunk_size, chunk_size, dtype=torch.bool, device=q.device), diagonal=1) - for i in range(0, l // chunk_size): - q_i, k_i, v_i, u_i, w_i, beta_i = q[:, :, i], k[:, :, i], v[:, :, i], u[:, :, i], w[:, :, i], beta[:, :, i] - v2_i = u_i + w_i @ S - - o_1 = A_qk[:, :, i] @ v_i - o_2 = A_qb[:, :, i] @ v2_i - o_3 = (q_i * gk_cumsum[:, :, i].exp()) @ S - o[:, :, i] = o_1 + o_2 + o_3 - decay = (gk_cumsum[:, :, i, -1, None] - gk_cumsum[:, :, i]).exp() - S = S*gk_cumsum[:, :, i, -1, :, None].exp() + (k_i * decay).transpose(-1, -2) @ v_i + \ - (beta_i * decay).transpose(-1, -2) @ v2_i - S = None if output_final_state is False else S - return rearrange(o, 'b h n c d -> b h (n c) d'), S diff --git a/build/torch-xpu/ops/generalized_delta_rule/dplr/wy_fast_bwd.py b/build/torch-xpu/ops/generalized_delta_rule/dplr/wy_fast_bwd.py deleted file mode 100644 index bf7b53ae84c415219f76a24f376ceb3159e04898..0000000000000000000000000000000000000000 --- a/build/torch-xpu/ops/generalized_delta_rule/dplr/wy_fast_bwd.py +++ /dev/null @@ -1,168 +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 -import triton.language as tl - -from ....ops.utils import prepare_chunk_indices -from ....utils import IS_INTEL_ALCHEMIST, autotune_cache_kwargs, check_shared_mem - -# https://github.com/intel/intel-xpu-backend-for-triton/issues/3449 -triton_config = {'grf_mode': 'large'} if IS_INTEL_ALCHEMIST else {} - - -@triton.heuristics({ - 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None, -}) -@triton.autotune( - configs=[ - triton.Config(triton_config, num_warps=num_warps, num_stages=num_stages) - for num_warps in [2, 4, 8, 16] - for num_stages in [2, 3, 4] - ], - key=['BT', 'BK', 'BV'], - **autotune_cache_kwargs, -) -@triton.jit(do_not_specialize=['T']) -def prepare_wy_repr_bwd_kernel( - A_ab_inv, - A_ak, - ag, - v, - dw, - du, - dv, - dv0, - dag, - dAak, - dAab, - cu_seqlens, - chunk_indices, - T, - H: tl.constexpr, - K: tl.constexpr, - V: tl.constexpr, - BT: tl.constexpr, - BK: tl.constexpr, - BV: tl.constexpr, - IS_VARLEN: tl.constexpr, -): - i_t, i_bh = tl.program_id(0), tl.program_id(1) - i_b, i_h = i_bh // H, i_bh % H - 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.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) - T = eos - bos - else: - bos, eos = i_b * T, i_b * T + T - - p_Aak_t = tl.make_block_ptr(A_ak + (bos*H + i_h) * BT, (BT, T), (1, H*BT), (0, i_t * BT), (BT, BT), (0, 1)) - p_Aab_inv_t = tl.make_block_ptr(A_ab_inv + (bos*H + i_h) * BT, (BT, T), (1, H*BT), (0, i_t * BT), (BT, BT), (0, 1)) - p_dAak = tl.make_block_ptr(dAak + (bos*H + i_h) * BT, (T, BT), (H*BT, 1), (i_t * BT, 0), (BT, BT), (1, 0)) - p_dAab = tl.make_block_ptr(dAab + (bos*H + i_h) * BT, (T, BT), (H*BT, 1), (i_t * BT, 0), (BT, BT), (1, 0)) - - b_A_ab_inv_t = tl.load(p_Aab_inv_t, boundary_check=(0, 1)) - b_A_ak_t = tl.load(p_Aak_t, boundary_check=(0, 1)) - b_A_ak_t = tl.where(tl.arange(0, BT)[:, None] < tl.arange(0, BT)[None, :], b_A_ak_t, 0) - b_A_ab_inv_t = tl.where(tl.arange(0, BT)[:, None] <= tl.arange(0, BT)[None, :], b_A_ab_inv_t, 0) - b_A_tmp_t = tl.dot(b_A_ak_t, b_A_ab_inv_t).to(v.dtype.element_ty) - b_dA_tmp = tl.zeros([BT, BT], dtype=tl.float32) - - for i_v in range(tl.cdiv(V, BV)): - p_v = tl.make_block_ptr(v + (bos*H + i_h) * V, (T, V), (H*V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) - p_dv = tl.make_block_ptr(dv + (bos*H + i_h) * V, (T, V), (H*V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) - p_dv0 = tl.make_block_ptr(dv0 + (bos*H + i_h) * V, (T, V), (H*V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) - p_du = tl.make_block_ptr(du + (bos*H + i_h) * V, (T, V), (H*V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) - b_v = tl.load(p_v, boundary_check=(0, 1)) - b_du = tl.load(p_du, boundary_check=(0, 1)) - b_dA_tmp += tl.dot(b_du.to(b_v.dtype), tl.trans(b_v)) - b_dv0 = tl.load(p_dv0, boundary_check=(0, 1)) - b_dv = b_dv0 + tl.dot(b_A_tmp_t, b_du) - tl.store(p_dv, b_dv.to(p_dv.dtype.element_ty), boundary_check=(0, 1)) - - m_i = tl.arange(0, BT)[:, None] > tl.arange(0, BT)[None, :] - b_dA_tmp = tl.where(m_i, b_dA_tmp, 0) - b_dA_ak = tl.dot(b_A_ab_inv_t, b_dA_tmp) - b_dA_ak = tl.where(m_i, b_dA_ak, 0) - tl.store(p_dAak, b_dA_ak, boundary_check=(0, 1)) - b_dA_ab_inv = tl.dot(b_dA_tmp, b_A_ak_t) - - for i_k in range(tl.cdiv(K, BK)): - p_ag = tl.make_block_ptr(ag + (bos * H + i_h) * K, (T, K), (H*K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) - p_dag = tl.make_block_ptr(dag + (bos * H + i_h) * K, (T, K), (H*K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) - p_dw = tl.make_block_ptr(dw + (bos * H + i_h) * K, (T, K), (H*K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) - b_ag = tl.load(p_ag, boundary_check=(0, 1)) - b_dw = tl.load(p_dw, boundary_check=(0, 1)) - b_dA_ab_inv += tl.dot(b_dw, tl.trans(b_ag)) - b_dag = tl.dot(b_A_ab_inv_t.to(b_dw.dtype), b_dw) - tl.store(p_dag, b_dag.to(p_dag.dtype.element_ty), boundary_check=(0, 1)) - - # if we know dL/dA^(-1), for dL/dA, we can use the following formula: - # dL/dA = -(A^(-1))^T @ (dL/dA^(-1)) @ (A^(-1))^T - # in the fwd pass we use fwd substitution to calculate (I-lower(A_ab))^-1. - # denote A = I - lower(A_ab), B = A^-1 - # in the backward pass. - # dL/dA = -(B)^T @ (dL/dB) @ B^T - # dL/dA_ab = lower(B^T @ dL/dB @ B^T) - b_dA_ab_inv = tl.where(tl.arange(0, BT)[:, None] >= tl.arange(0, BT)[None, :], b_dA_ab_inv, 0) - b_dA_ab_inv = tl.dot(b_A_ab_inv_t, b_dA_ab_inv) - b_dA_ab_inv = tl.dot(b_dA_ab_inv, b_A_ab_inv_t) - b_dA_ab_inv = tl.where(m_i, b_dA_ab_inv, 0) - tl.store(p_dAab, b_dA_ab_inv, boundary_check=(0, 1)) - - -def chunk_dplr_bwd_wy( - A_ab_inv: torch.Tensor, - A_ak: torch.Tensor, - v: torch.Tensor, - ag: torch.Tensor, - dw: torch.Tensor, - du: torch.Tensor, - dv0: torch.Tensor, - cu_seqlens: torch.LongTensor | None, - chunk_size: int, - chunk_indices: torch.LongTensor | None = None, -) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: - A_ab_inv, A_ak, v, ag, dw, du = map(lambda x: x.contiguous(), [A_ab_inv, A_ak, v, ag, dw, du]) - B, T, H, K, V = *dw.shape, du.shape[-1] - BT = chunk_size - - if chunk_indices is None: - chunk_indices = prepare_chunk_indices(cu_seqlens, BT) if cu_seqlens is not None else None - NT = triton.cdiv(T, BT) if cu_seqlens is None else len(chunk_indices) - BK = min(max(triton.next_power_of_2(K), 16), 64) - BV = min(max(triton.next_power_of_2(V), 16), 64) if check_shared_mem() else min(max(triton.next_power_of_2(V), 16), 32) - - dA_ab = torch.empty_like(A_ab_inv, dtype=torch.float) - dA_ak = torch.empty_like(A_ak, dtype=torch.float) - dv = torch.empty_like(v) - dag = torch.empty_like(ag) - - prepare_wy_repr_bwd_kernel[(NT, B * H)]( - A_ab_inv=A_ab_inv, - A_ak=A_ak, - ag=ag, - v=v, - dw=dw, - du=du, - dv=dv, - dv0=dv0, - dag=dag, - dAak=dA_ak, - dAab=dA_ab, - cu_seqlens=cu_seqlens, - chunk_indices=chunk_indices, - T=T, - H=H, - K=K, - V=V, - BT=BT, - BK=BK, - BV=BV, - ) - return dA_ab, dA_ak, dv, dag diff --git a/build/torch-xpu/ops/generalized_delta_rule/dplr/wy_fast_fwd.py b/build/torch-xpu/ops/generalized_delta_rule/dplr/wy_fast_fwd.py deleted file mode 100644 index 1d8c91f6b106be2412e6b07895c90f0dd0ab0bda..0000000000000000000000000000000000000000 --- a/build/torch-xpu/ops/generalized_delta_rule/dplr/wy_fast_fwd.py +++ /dev/null @@ -1,291 +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 -import triton.language as tl - -from ....ops.utils import prepare_chunk_indices -from ....ops.utils.op import gather -from ....utils import IS_GATHER_SUPPORTED, autotune_cache_kwargs - - -@triton.heuristics({ - 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None, -}) -@triton.autotune( - configs=[ - triton.Config({}, num_warps=num_warps) - for num_warps in [1, 2, 4, 8, 16] - ], - key=['BT'], - **autotune_cache_kwargs, -) -@triton.jit(do_not_specialize=['T']) -def prepare_wy_repr_fwd_kernel_chunk32( - A_ab, - A_ab_inv, - cu_seqlens, - chunk_indices, - T, - H: tl.constexpr, - BT: tl.constexpr, - BC: tl.constexpr, # placeholder, do not delete - IS_VARLEN: tl.constexpr, -): - i_t, i_bh = tl.program_id(0), tl.program_id(1) - i_b, i_h = i_bh // H, i_bh % H - 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.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) - T = eos - bos - else: - bos, eos = i_b * T, i_b * T + T - p_Aab = tl.make_block_ptr(A_ab + (bos*H + i_h) * BT, (T, BT), (H*BT, 1), (i_t * BT, 0), (BT, BT), (1, 0)) - p_Aab_inv = tl.make_block_ptr(A_ab_inv + (bos*H + i_h) * BT, (T, BT), (H*BT, 1), (i_t * BT, 0), (BT, BT), (1, 0)) - b_A_ab = tl.load(p_Aab, boundary_check=(0, 1)) - b_A_ab = tl.where(tl.arange(0, BT)[:, None] > tl.arange(0, BT)[None, :], b_A_ab, 0) - for i in range(1, BT): - mask = tl.arange(0, BT) == i - b_a = tl.sum(tl.where(mask[:, None], b_A_ab, 0), 0) - b_a = b_a + tl.sum(b_a[:, None] * b_A_ab, 0) * (tl.arange(0, BT) < i) - b_A_ab = tl.where(mask[:, None], b_a, b_A_ab) - b_A_ab += tl.arange(0, BT)[:, None] == tl.arange(0, BT)[None, :] - tl.store(p_Aab_inv, b_A_ab.to(p_Aab_inv.dtype.element_ty), boundary_check=(0, 1)) - - -@triton.heuristics({ - 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None, -}) -@triton.autotune( - configs=[ - triton.Config({}, num_warps=num_warps, num_stages=num_stages) - for num_warps in [2, 4, 8] - for num_stages in [2, 3, 4] - ], - key=['BC'], - **autotune_cache_kwargs, -) -@triton.jit(do_not_specialize=['T']) -def prepare_wy_repr_fwd_kernel_chunk64( - A_ab, - A_ab_inv, - cu_seqlens, - chunk_indices, - T, - H: tl.constexpr, - BT: tl.constexpr, - BC: tl.constexpr, - IS_VARLEN: tl.constexpr, - GATHER_SUPPORTED: tl.constexpr = IS_GATHER_SUPPORTED, -): - i_t, i_bh = tl.program_id(0), tl.program_id(1) - i_b, i_h = i_bh // H, i_bh % H - 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.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) - T = eos - bos - else: - bos, eos = i_b * T, i_b * T + T - - p_A1 = tl.make_block_ptr(A_ab + (bos*H + i_h) * BT, (T, BT), (H*BT, 1), (i_t * BT, 0), (BC, BC), (1, 0)) - p_A2 = tl.make_block_ptr(A_ab + (bos*H + i_h) * BT, (T, BT), (H*BT, 1), (i_t * BT + BC, BC), (BC, BC), (1, 0)) - p_A3 = tl.make_block_ptr(A_ab + (bos*H + i_h) * BT, (T, BT), (H*BT, 1), (i_t * BT + BC, 0), (BC, BC), (1, 0)) - p_A_inv1 = tl.make_block_ptr(A_ab_inv + (bos*H + i_h) * BT, (T, BT), (H*BT, 1), (i_t * BT, 0), (BC, BC), (1, 0)) - p_A_inv2 = tl.make_block_ptr(A_ab_inv + (bos*H + i_h) * BT, (T, BT), (H*BT, 1), (i_t * BT + BC, BC), (BC, BC), (1, 0)) - p_A_inv3 = tl.make_block_ptr(A_ab_inv + (bos*H + i_h) * BT, (T, BT), (H*BT, 1), (i_t * BT + BC, 0), (BC, BC), (1, 0)) - p_A_inv4 = tl.make_block_ptr(A_ab_inv + (bos*H + i_h) * BT, (T, BT), (H*BT, 1), (i_t * BT, BC), (BC, BC), (1, 0)) - - b_A = tl.load(p_A1, boundary_check=(0, 1)) - b_A2 = tl.load(p_A2, boundary_check=(0, 1)) - b_A3 = tl.load(p_A3, boundary_check=(0, 1)) - b_A = tl.where(tl.arange(0, BC)[:, None] > tl.arange(0, BC)[None, :], b_A, 0) - b_A2 = tl.where(tl.arange(0, BC)[:, None] > tl.arange(0, BC)[None, :], b_A2, 0) - - for i in range(1, BC): - if GATHER_SUPPORTED: - row_idx = tl.full([1, BC], i, dtype=tl.int16) - # [1, BK] -> [BK] - b_a = tl.sum(gather(b_A, row_idx, axis=0), 0) - b_a2 = tl.sum(gather(b_A2, row_idx, axis=0), 0) - else: - mask = tl.arange(0, BC) == i - b_a = tl.sum(tl.where(mask[:, None], b_A, 0), 0) - b_a2 = tl.sum(tl.where(mask[:, None], b_A2, 0), 0) - mask = tl.arange(0, BC) == i - # b_a = tl.sum(tl.where(mask[:, None], b_A, 0), 0) - # b_a2 = tl.sum(tl.where(mask[:, None], b_A2, 0), 0) - b_a = b_a + tl.sum(b_a[:, None] * b_A, 0) * (tl.arange(0, BC) < i) - b_a2 = b_a2 + tl.sum(b_a2[:, None] * b_A2, 0) * (tl.arange(0, BC) < i) - b_A = tl.where(mask[:, None], b_a, b_A) - b_A2 = tl.where(mask[:, None], b_a2, b_A2) - - # blockwise computation of lower triangular matrix's inverse - # i.e., [A11, 0; A21, A22]^-1 = [A11^-1, 0; -A22^-1 A21 A11^-1, A22^-1] - b_A += tl.arange(0, BC)[:, None] == tl.arange(0, BC)[None, :] - b_A2 += tl.arange(0, BC)[:, None] == tl.arange(0, BC)[None, :] - b_A3 = tl.dot(tl.dot(b_A2, b_A3), b_A) - # tl.debug_barrier() - tl.store(p_A_inv1, b_A.to(p_A_inv1.dtype.element_ty, fp_downcast_rounding="rtne"), boundary_check=(0, 1)) - tl.store(p_A_inv2, b_A2.to(p_A_inv2.dtype.element_ty, fp_downcast_rounding="rtne"), boundary_check=(0, 1)) - tl.store(p_A_inv3, b_A3.to(p_A_inv3.dtype.element_ty, fp_downcast_rounding="rtne"), boundary_check=(0, 1)) - # causal mask - tl.store(p_A_inv4, tl.zeros([BC, BC], dtype=tl.float32).to(p_A_inv4.dtype.element_ty), boundary_check=(0, 1)) - - -@triton.heuristics({ - 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None, -}) -@triton.autotune( - configs=[ - triton.Config({}, num_warps=num_warps, num_stages=num_stages) - for num_warps in [2, 4, 8, 16] - for num_stages in [2, 3, 4] - ], - key=['H', 'K', 'V', 'BT', 'BK', 'BV', 'IS_VARLEN'], - **autotune_cache_kwargs, -) -@triton.jit(do_not_specialize=['T']) -def wu_fwd_kernel( - w, - u, - ag, - v, - A_ab_inv, - A_ak, - cu_seqlens, - chunk_indices, - T, - H: tl.constexpr, - K: tl.constexpr, - V: tl.constexpr, - BT: tl.constexpr, - BK: tl.constexpr, - BV: tl.constexpr, - IS_VARLEN: tl.constexpr, -): - i_t, i_bh = tl.program_id(0), tl.program_id(1) - i_b, i_h = i_bh // H, i_bh % H - 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.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) - T = eos - bos - else: - bos, eos = i_b * T, i_b * T + T - o_s = tl.arange(0, BT) - - p_A_ab_inv = tl.make_block_ptr(A_ab_inv + (bos*H + i_h) * BT, (T, BT), (H*BT, 1), (i_t * BT, 0), (BT, BT), (1, 0)) - p_A_ak = tl.make_block_ptr(A_ak + (bos*H + i_h) * BT, (T, BT), (H*BT, 1), (i_t * BT, 0), (BT, BT), (1, 0)) - - b_Aab_inv = tl.load(p_A_ab_inv, boundary_check=(0, 1)) - b_Aak = tl.load(p_A_ak, boundary_check=(0, 1)) - b_Aab_inv = tl.where(o_s[:, None] >= o_s[None, :], b_Aab_inv, 0) - b_Aak = tl.where(o_s[:, None] > o_s[None, :], b_Aak, 0) - # let's use tf32 here - b_Aak = tl.dot(b_Aab_inv, b_Aak) - # (SY 01/04) should be bf16 or tf32? To verify. - b_Aak = b_Aak.to(v.dtype.element_ty, fp_downcast_rounding="rtne") - b_Aab_inv = b_Aab_inv.to(ag.dtype.element_ty, fp_downcast_rounding="rtne") - - for i_k in range(tl.cdiv(K, BK)): - p_ag = tl.make_block_ptr(ag + (bos*H + i_h) * K, (T, K), (H*K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) - p_w = tl.make_block_ptr(w + (bos*H + i_h) * K, (T, K), (H*K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) - b_ag = tl.load(p_ag, boundary_check=(0, 1)) - b_w = tl.dot(b_Aab_inv, b_ag) # both bf16 or fp16 - tl.store(p_w, b_w.to(p_w.dtype.element_ty, fp_downcast_rounding="rtne"), boundary_check=(0, 1)) - - for i_v in range(tl.cdiv(V, BV)): - p_v = tl.make_block_ptr(v + (bos*H + i_h) * V, (T, V), (H*V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) - p_u = tl.make_block_ptr(u + (bos*H + i_h) * V, (T, V), (H*V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) - b_v = tl.load(p_v, boundary_check=(0, 1)) - b_u = tl.dot(b_Aak, b_v) # both bf16 or fp16 - tl.store(p_u, b_u.to(p_u.dtype.element_ty, fp_downcast_rounding="rtne"), boundary_check=(0, 1)) - - -def wu_fwd( - ag: torch.Tensor, - v: torch.Tensor, - A_ak: torch.Tensor, - A_ab_inv: torch.Tensor, - cu_seqlens: torch.LongTensor | None, - chunk_size: int, - chunk_indices: torch.LongTensor | None = None, -) -> tuple[torch.Tensor, torch.Tensor]: - B, T, H, K, V = *ag.shape, v.shape[-1] - BT = chunk_size - - if chunk_indices is None: - chunk_indices = prepare_chunk_indices(cu_seqlens, BT) if cu_seqlens is not None else None - NT = triton.cdiv(T, BT) if cu_seqlens is None else len(chunk_indices) - BK = min(max(triton.next_power_of_2(K), 16), 64) - BV = min(max(triton.next_power_of_2(V), 16), 64) - - w = torch.empty_like(ag) - u = torch.empty_like(v) - wu_fwd_kernel[(NT, B * H)]( - ag=ag, - v=v, - A_ak=A_ak, - A_ab_inv=A_ab_inv, - w=w, - u=u, - cu_seqlens=cu_seqlens, - chunk_indices=chunk_indices, - T=T, - H=H, - K=K, - V=V, - BT=BT, - BK=BK, - BV=BV, - ) - return w, u - - -def prepare_wy_repr_fwd( - ag: torch.Tensor, - v: torch.Tensor, - A_ak: torch.Tensor, - A_ab: torch.Tensor, - cu_seqlens: torch.LongTensor | None, - chunk_size: int = 64, - chunk_indices: torch.LongTensor | None = None, -) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: - B, T, H, _ = ag.shape - BT = chunk_size - - if chunk_indices is None: - chunk_indices = prepare_chunk_indices(cu_seqlens, BT) if cu_seqlens is not None else None - NT = triton.cdiv(T, BT) if cu_seqlens is None else len(chunk_indices) - BC = min(BT, 32) - fwd_fn = prepare_wy_repr_fwd_kernel_chunk64 if BT == 64 else prepare_wy_repr_fwd_kernel_chunk32 - A_ab_inv = torch.empty_like(A_ab) - fwd_fn[(NT, B * H)]( - A_ab=A_ab, - A_ab_inv=A_ab_inv, - cu_seqlens=cu_seqlens, - chunk_indices=chunk_indices, - T=T, - H=H, - BT=BT, - BC=BC, - ) - w, u = wu_fwd( - ag=ag, - v=v, - A_ak=A_ak, - A_ab_inv=A_ab_inv, - cu_seqlens=cu_seqlens, - chunk_size=BT, - chunk_indices=chunk_indices, - ) - return w, u, A_ab_inv - - -fwd_prepare_wy_repr = prepare_wy_repr_fwd - -fwd_wu = wu_fwd diff --git a/build/torch-xpu/ops/generalized_delta_rule/iplr/__init__.py b/build/torch-xpu/ops/generalized_delta_rule/iplr/__init__.py deleted file mode 100644 index 2fe6d4f9c0e92bdbb08e4b5eebe148506d075b89..0000000000000000000000000000000000000000 --- a/build/torch-xpu/ops/generalized_delta_rule/iplr/__init__.py +++ /dev/null @@ -1,14 +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 .chunk import chunk_iplr_delta_rule -from .fused_recurrent import fused_recurrent_iplr_delta_rule - -__all__ = [ - 'chunk_iplr_delta_rule', - 'fused_recurrent_iplr_delta_rule', -] diff --git a/build/torch-xpu/ops/generalized_delta_rule/iplr/chunk.py b/build/torch-xpu/ops/generalized_delta_rule/iplr/chunk.py deleted file mode 100644 index f4bb021597aed7559ccf47ca373ad62ebc5bc456..0000000000000000000000000000000000000000 --- a/build/torch-xpu/ops/generalized_delta_rule/iplr/chunk.py +++ /dev/null @@ -1,507 +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 -import triton.language as tl - -from ....ops.generalized_delta_rule.iplr.wy_fast import prepare_wy_repr_fwd -from ....ops.utils import prepare_chunk_indices, prepare_chunk_offsets -from ....utils import ( - autocast_custom_bwd, - autocast_custom_fwd, - autotune_cache_kwargs, - check_shared_mem, - input_guard, -) - -BKV_LIST = [64, 128] if check_shared_mem() else [32, 64] - - -@triton.heuristics({ - 'USE_INITIAL_STATE': lambda args: args['h0'] is not None, - 'STORE_FINAL_STATE': lambda args: args['ht'] is not None, - 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None, -}) -@triton.autotune( - configs=[ - triton.Config({}, num_warps=num_warps) - for num_warps in [2, 4] + ([] if check_shared_mem('hopper') else [8]) - ], - key=['BT', 'BK', 'BV'], - **autotune_cache_kwargs, -) -@triton.jit(do_not_specialize=['T']) -def chunk_generalized_iplr_delta_rule_fwd_kernel_h( - k, - v, - d, - b, - u, - v_new, - h, - h0, - ht, - cu_seqlens, - chunk_offsets, - T, - H: tl.constexpr, - K: tl.constexpr, - V: tl.constexpr, - BT: tl.constexpr, - BC: tl.constexpr, - BK: tl.constexpr, - BV: tl.constexpr, - USE_INITIAL_STATE: tl.constexpr, - STORE_FINAL_STATE: tl.constexpr, - IS_VARLEN: tl.constexpr, -): - i_k, i_v, i_nh = tl.program_id(0), tl.program_id(1), tl.program_id(2) - i_n, i_h = i_nh // H, i_nh % H - if IS_VARLEN: - bos, eos = tl.load(cu_seqlens + i_n).to(tl.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) - T = eos - bos - NT = tl.cdiv(T, BT) - boh = tl.load(chunk_offsets + i_n).to(tl.int32) - else: - bos, eos = i_n * T, i_n * T + T - NT = tl.cdiv(T, BT) - boh = i_n * NT - - # [BK, BV] - b_h = tl.zeros([BK, BV], dtype=tl.float32) - if USE_INITIAL_STATE: - p_h0 = tl.make_block_ptr(h0 + i_nh * K*V, (K, V), (V, 1), (i_k * BK, i_v * BV), (BK, BV), (1, 0)) - b_h = tl.load(p_h0, boundary_check=(0, 1)).to(tl.float32) - - for i_t in range(NT): - p_h = tl.make_block_ptr(h + ((boh + i_t) * H + i_h) * K*V, (K, V), (V, 1), (i_k * BK, i_v * BV), (BK, BV), (1, 0)) - tl.store(p_h, b_h.to(p_h.dtype.element_ty), boundary_check=(0, 1)) - b_hc = tl.zeros([BK, BV], dtype=tl.float32) - # since we need to make all DK in the SRAM. we face serve SRAM memory burden. By subchunking we allievate such burden - for i_c in range(tl.cdiv(min(BT, T - i_t * BT), BC)): - p_k = tl.make_block_ptr(k+(bos*H+i_h)*K, (K, T), (1, H*K), (i_k * BK, i_t * BT + i_c * BC), (BK, BC), (0, 1)) - p_b = tl.make_block_ptr(b+(bos*H+i_h)*K, (K, T), (1, H*K), (i_k * BK, i_t * BT + i_c * BC), (BK, BC), (0, 1)) - p_d = tl.make_block_ptr(d+(bos*H+i_h)*K, (T, K), (H*K, 1), (i_t * BT + i_c * BC, i_k * BK), (BC, BK), (1, 0)) - p_v = tl.make_block_ptr(v+(bos*H+i_h)*V, (T, V), (H*V, 1), (i_t * BT + i_c * BC, i_v * BV), (BC, BV), (1, 0)) - p_u = tl.make_block_ptr(u+(bos*H+i_h)*V, (T, V), (H*V, 1), (i_t * BT + i_c * BC, i_v * BV), (BC, BV), (1, 0)) - p_v_new = tl.make_block_ptr(v_new+(bos*H+i_h)*V, (T, V), (H*V, 1), (i_t*BT+i_c*BC, i_v * BV), (BC, BV), (1, 0)) - # [BK, BC] - b_k = tl.load(p_k, boundary_check=(0, 1)) - b_v = tl.load(p_v, boundary_check=(0, 1)) - b_d = tl.load(p_d, boundary_check=(0, 1)) - b_b = tl.load(p_b, boundary_check=(0, 1)) - b_v2 = tl.dot(b_d, b_h.to(b_d.dtype)) + tl.load(p_u, boundary_check=(0, 1)) - b_hc += tl.dot(b_k, b_v) - b_hc += tl.dot(b_b, b_v2.to(b_k.dtype)) - tl.store(p_v_new, b_v2.to(p_v_new.dtype.element_ty), boundary_check=(0, 1)) - b_h += b_hc - - if STORE_FINAL_STATE: - p_ht = tl.make_block_ptr(ht + i_nh * K*V, (K, V), (V, 1), (i_k * BK, i_v * BV), (BK, BV), (1, 0)) - tl.store(p_ht, b_h.to(p_ht.dtype.element_ty), boundary_check=(0, 1)) - - -@triton.heuristics({ - 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None, -}) -@triton.autotune( - configs=[ - triton.Config({'BK': BK, 'BV': BV}, num_warps=num_warps) - for BK in BKV_LIST - for BV in BKV_LIST - for num_warps in [2, 4, 8] - ], - key=['BT'], - **autotune_cache_kwargs, -) -@triton.jit(do_not_specialize=['T']) -def chunk_generalized_iplr_delta_rule_fwd_kernel_o( - q, - k, - v, - u, - b, - h, - o, - cu_seqlens, - chunk_indices, - scale, - T, - H: tl.constexpr, - K: tl.constexpr, - V: tl.constexpr, - BT: tl.constexpr, - BK: tl.constexpr, - BV: tl.constexpr, - IS_VARLEN: tl.constexpr, -): - i_v, i_t, i_bh = tl.program_id(0), tl.program_id(1), tl.program_id(2) - i_b, i_h = i_bh // H, i_bh % H - - 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.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) - T = eos - bos - NT = tl.cdiv(T, BT) - else: - NT = tl.cdiv(T, BT) - i_tg = i_b * NT + i_t - bos, eos = i_b * T, i_b * T + T - - # offset calculation - q += (bos * H + i_h) * K - k += (bos * H + i_h) * K - b += (bos * H + i_h) * K - v += (bos * H + i_h) * V - u += (bos * H + i_h) * V - o += (bos * H + i_h) * V - h += (i_tg * H + i_h) * K * V - stride_qk = H*K - stride_vo = H*V - - b_o = tl.zeros([BT, BV], dtype=tl.float32) - b_Aqk = tl.zeros([BT, BT], dtype=tl.float32) - b_Aqb = tl.zeros([BT, BT], dtype=tl.float32) - - for i_k in range(tl.cdiv(K, BK)): - p_q = tl.make_block_ptr(q, (T, K), (stride_qk, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) - p_k = tl.make_block_ptr(k, (K, T), (1, stride_qk), (i_k * BK, i_t * BT), (BK, BT), (0, 1)) - p_h = tl.make_block_ptr(h, (K, V), (V, 1), (i_k * BK, i_v * BV), (BK, BV), (1, 0)) - p_b = tl.make_block_ptr(b, (K, T), (1, stride_qk), (i_k * BK, i_t * BT), (BK, BT), (0, 1)) - # [BT, BK] - b_q = tl.load(p_q, boundary_check=(0, 1)) - # [BK, BT] - b_k = tl.load(p_k, boundary_check=(0, 1)) - b_b = tl.load(p_b, boundary_check=(0, 1)) - # [BK, BV] - b_h = tl.load(p_h, boundary_check=(0, 1)) - # [BT, BK] @ [BK, BV] -> [BT, BV] - b_o += tl.dot(b_q, b_h) - # [BT, BK] @ [BK, BT] -> [BT, BT] - b_Aqk += tl.dot(b_q, b_k) - # [BT, BK] @ [BK, BT] -> [BT, BT] - b_Aqb += tl.dot(b_q, b_b) - - o_i = tl.arange(0, BT) - m_A = o_i[:, None] >= o_i[None, :] - b_Aqk = tl.where(m_A, b_Aqk, 0) - b_Aqb = tl.where(m_A, b_Aqb, 0) - - p_v = tl.make_block_ptr(v, (T, V), (stride_vo, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) - p_u = tl.make_block_ptr(u, (T, V), (stride_vo, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) - p_o = tl.make_block_ptr(o, (T, V), (stride_vo, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) - b_v = tl.load(p_v, boundary_check=(0, 1)) - b_u = tl.load(p_u, boundary_check=(0, 1)) - b_o = (b_o + tl.dot(b_Aqk.to(b_v.dtype), b_v) + tl.dot(b_Aqb.to(b_u.dtype), b_u)) * scale - tl.store(p_o, b_o.to(p_o.dtype.element_ty), boundary_check=(0, 1)) - - -def chunk_generalized_iplr_delta_rule_fwd_o( - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - v_new: torch.Tensor, - b: torch.Tensor, - h: torch.Tensor, - scale: float | None = None, - cu_seqlens: torch.LongTensor | None = None, - chunk_size: int = 64, - chunk_indices: torch.LongTensor | None = None, -) -> torch.Tensor: - B, T, H, K, V = *q.shape, v.shape[-1] - if scale is None: - scale = k.shape[-1] ** -0.5 - BT = chunk_size - - if chunk_indices is None: - chunk_indices = prepare_chunk_indices(cu_seqlens, BT) if cu_seqlens is not None else None - NT = triton.cdiv(T, BT) if cu_seqlens is None else len(chunk_indices) - - o = torch.empty_like(v) - - def grid(meta): return ( - triton.cdiv(V, meta['BV']), - NT, - B * H, - ) - chunk_generalized_iplr_delta_rule_fwd_kernel_o[grid]( - q=q, - k=k, - v=v, - u=v_new, - b=b, - h=h, - o=o, - cu_seqlens=cu_seqlens, - chunk_indices=chunk_indices, - scale=scale, - T=T, - H=H, - K=K, - V=V, - BT=BT, - ) - return o - - -def chunk_generalized_iplr_delta_rule_fwd_h( - k: torch.Tensor, - v: torch.Tensor, - w: torch.Tensor, - u: torch.Tensor, - b: torch.Tensor, - initial_state: torch.Tensor | None = None, - output_final_state: bool = False, - cu_seqlens: torch.LongTensor | None = None, - chunk_size: int = 64, - chunk_indices: torch.LongTensor | None = None, -) -> tuple[torch.Tensor, torch.Tensor]: - B, T, H, K, V = *k.shape, u.shape[-1] - BT = chunk_size - - if chunk_indices is None: - chunk_indices = prepare_chunk_indices(cu_seqlens, BT) if cu_seqlens is not None else None - # N: the actual number of sequences in the batch with either equal or variable lengths - if cu_seqlens is None: - N, NT, chunk_offsets = B, triton.cdiv(T, BT), None - else: - N, NT, chunk_offsets = len(cu_seqlens) - 1, len(chunk_indices), prepare_chunk_offsets(cu_seqlens, BT) - - BK = max(triton.next_power_of_2(K), 16) - assert BK <= 256, "current kernel does not support head dimension larger than 256." - # H100 can have larger block size - - if check_shared_mem('hopper', k.device.index): - BV = 64 - BC = 64 if K <= 128 else 32 - elif check_shared_mem('ampere', k.device.index): # A100 - BV = 32 - BC = 32 - else: - BV = 16 - BC = 16 - - BC = min(BT, BC) - NK = triton.cdiv(K, BK) - NV = triton.cdiv(V, BV) - - assert NK == 1, 'NK > 1 is not supported because it involves time-consuming synchronization' - - h = k.new_empty(B, NT, H, K, V) - final_state = k.new_empty(N, H, K, V, dtype=torch.float32) if output_final_state else None - - v_new = torch.empty_like(u) - grid = (NK, NV, N * H) - - chunk_generalized_iplr_delta_rule_fwd_kernel_h[grid]( - k=k, - v=v, - d=w, - b=b, - u=u, - v_new=v_new, - h=h, - h0=initial_state, - ht=final_state, - cu_seqlens=cu_seqlens, - chunk_offsets=chunk_offsets, - T=T, - H=H, - K=K, - V=V, - BT=BT, - BC=BC, - BK=BK, - BV=BV, - ) - return h, v_new, final_state - - -def chunk_generalized_iplr_delta_rule_fwd( - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - a: torch.Tensor, - b: torch.Tensor, - scale: float, - initial_state: torch.Tensor, - output_final_state: bool, - cu_seqlens: torch.LongTensor | None = None, - chunk_size: int = 64, - chunk_indices: torch.LongTensor | None = None, -): - w, u, _ = prepare_wy_repr_fwd( - a=a, - b=b, - k=k, - v=v, - cu_seqlens=cu_seqlens, - chunk_size=chunk_size, - chunk_indices=chunk_indices, - ) - - h, v_new, final_state = chunk_generalized_iplr_delta_rule_fwd_h( - k=k, - v=v, - b=b, - w=w, - u=u, - initial_state=initial_state, - output_final_state=output_final_state, - cu_seqlens=cu_seqlens, - chunk_size=chunk_size, - chunk_indices=chunk_indices, - ) - o = chunk_generalized_iplr_delta_rule_fwd_o( - q=q, - k=k, - v=v, - v_new=v_new, - b=b, - h=h, - scale=scale, - cu_seqlens=cu_seqlens, - chunk_size=chunk_size, - chunk_indices=chunk_indices, - ) - return o, final_state - - -class ChunkGeneralizedIPLRDeltaRuleFunction(torch.autograd.Function): - - @staticmethod - @input_guard - @autocast_custom_fwd - def forward( - ctx, - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - a: torch.Tensor, - b: torch.Tensor, - scale: float, - initial_state: torch.Tensor, - output_final_state: bool, - cu_seqlens: torch.LongTensor | None = None, - cu_seqlens_cpu: torch.LongTensor | None = None, - chunk_size: int | None = None, - ): - if chunk_size is None: - chunk_size = 64 - chunk_indices = prepare_chunk_indices( - cu_seqlens, chunk_size, cu_seqlens_cpu=cu_seqlens_cpu) if cu_seqlens is not None else None - o, final_state = chunk_generalized_iplr_delta_rule_fwd( - q=q, - k=k, - v=v, - a=a, - b=b, - scale=scale, - initial_state=initial_state, - output_final_state=output_final_state, - cu_seqlens=cu_seqlens, - chunk_size=chunk_size, - chunk_indices=chunk_indices, - ) - return o.to(q.dtype), final_state - - @staticmethod - @input_guard - @autocast_custom_bwd - def backward( - ctx, - do: torch.Tensor, - dht: torch.Tensor, - ): - raise NotImplementedError( - "Backward pass for ChunkGeneralizedIPLRDeltaRuleFunction is not implemented yet. " - "Stay tuned!", - ) - - -@torch.compiler.disable -def chunk_iplr_delta_rule( - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - a: torch.Tensor, - b: torch.Tensor, - scale: float | None = None, - initial_state: torch.Tensor | None = None, - output_final_state: bool = False, - cu_seqlens: torch.LongTensor | None = None, - cu_seqlens_cpu: torch.LongTensor | None = None, - **kwargs, -): - r""" - Args: - q (torch.Tensor): - queries of shape `[B, T, H, K]`. - k (torch.Tensor): - keys of shape `[B, T, H, K]`. - v (torch.Tensor): - values of shape `[B, T, H, V]`. - a (torch.Tensor): - activations of shape `[B, T, H, K]`. - b (torch.Tensor): - betas of shape `[B, T, H, K]`. - scale (Optional[float]): - Scale factor for the attention scores. - If not provided, it will default to `1 / sqrt(K)`. Default: `None`. - initial_state (Optional[torch.Tensor]): - Initial state of shape `[N, H, K, V]` for `N` input sequences. - For equal-length input sequences, `N` equals the batch size `B`. - Default: `None`. - output_final_state (Optional[bool]): - Whether to output the final state of shape `[N, H, K, V]`. Default: `False`. - cu_seqlens (torch.LongTensor): - Cumulative sequence lengths of shape `[N+1]` used for variable-length training, - consistent with the FlashAttention API. - cu_seqlens_cpu (torch.LongTensor): - CPU copy of `cu_seqlens` to avoid unnecessary device synchronization. Default: `None`. - - Returns: - o (torch.Tensor): - Outputs of shape `[B, T, H, V]`. - final_state (torch.Tensor): - Final state of shape `[N, H, K, V]` if `output_final_state=True` else `None`. - """ - if 'head_first' in kwargs: - raise DeprecationWarning( - "head_first has been removed. Inputs must be in `[B, T, H, ...]` format.", - ) - chunk_size = kwargs.pop('chunk_size', None) - if chunk_size is not None and chunk_size != 2 ** (chunk_size.bit_length() - 1): - raise ValueError(f"`chunk_size` must be a power of 2, got {chunk_size}.") - if cu_seqlens is not None: - if q.shape[0] != 1: - raise ValueError( - f"The batch size is expected to be 1 rather than {q.shape[0]} when using `cu_seqlens`. " - f"Please flatten variable-length inputs before processing.", - ) - if initial_state is not None and initial_state.shape[0] != len(cu_seqlens) - 1: - raise ValueError( - f"The number of initial states is expected to be equal to the number of input sequences, " - f"i.e., {len(cu_seqlens) - 1} rather than {initial_state.shape[0]}.", - ) - scale = k.shape[-1] ** -0.5 if scale is None else scale - o, final_state = ChunkGeneralizedIPLRDeltaRuleFunction.apply( - q, - k, - v, - a, - b, - scale, - initial_state, - output_final_state, - cu_seqlens, - cu_seqlens_cpu, - chunk_size, - ) - return o, final_state diff --git a/build/torch-xpu/ops/generalized_delta_rule/iplr/fused_recurrent.py b/build/torch-xpu/ops/generalized_delta_rule/iplr/fused_recurrent.py deleted file mode 100644 index c7046da862bbc927cd0b762d78236a611fdc4d63..0000000000000000000000000000000000000000 --- a/build/torch-xpu/ops/generalized_delta_rule/iplr/fused_recurrent.py +++ /dev/null @@ -1,456 +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 -import triton.language as tl - -from ....utils import autotune_cache_kwargs, input_guard - - -@triton.heuristics({ - 'USE_INITIAL_STATE': lambda args: args['h0'] is not None, - 'STORE_FINAL_STATE': lambda args: args['ht'] is not None, - 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None, -}) -@triton.autotune( - configs=[ - triton.Config({'BV': BV}, num_warps=num_warps, num_stages=num_stages) - for BV in [32, 64] - for num_warps in [2, 4, 8, 16] - for num_stages in [2, 3, 4] - ], - key=["BK"], - **autotune_cache_kwargs, -) -@triton.jit -def fused_recurrent_fwd_kernel( - q, # query [B, H, L, K] - k, # key [B, H, L, V] - v, # value [B, H, L, V]. - a, # a [B, H, L, K] - b, # b [B, H, L, K] - o, # output [B, H, L, V] - ha, # tmp variable [B, H, L, V] for storing intermediate results of (h * a[None, :]).sum(0) - h0, # initial hidden state [B, H, K, V] - ht, # final hidden state [B, H, K, V] - cu_seqlens, # varlen cu_seqlens - scale, # K ** -0.5 - H, # n_heads - T, # seq_len - K: tl.constexpr, # K - V: tl.constexpr, # V - BK: tl.constexpr, # BLOCK SIZE along the K dimension - BV: tl.constexpr, # BLOCK SIZE along the V dimension - USE_INITIAL_STATE: tl.constexpr, # whether to use initial state - STORE_FINAL_STATE: tl.constexpr, # whether to store final state - IS_VARLEN: tl.constexpr, -): - i_v, i_nh = tl.program_id(0), tl.program_id(1) - i_n, i_h = i_nh // H, i_nh % H - - if IS_VARLEN: - bos, eos = tl.load(cu_seqlens + i_n).to(tl.int64), tl.load(cu_seqlens + i_n + 1).to(tl.int64) - T = eos - bos - else: - bos, eos = i_n * T, i_n * T + T - - p_q = q + (bos * H + i_h) * K + tl.arange(0, BK) - p_k = k + (bos * H + i_h) * K + tl.arange(0, BK) - p_a = a + (bos * H + i_h) * K + tl.arange(0, BK) - p_b = b + (bos * H + i_h) * K + tl.arange(0, BK) - p_ha = ha + (bos * H + i_h) * V + i_v * BV + tl.arange(0, BV) - p_v = v + (bos * H + i_h) * V + i_v * BV + tl.arange(0, BV) - p_o = o + (bos * H + i_h) * V + i_v * BV + tl.arange(0, BV) - - mask_k = tl.arange(0, BK) < K - mask_v = (i_v * BV + tl.arange(0, BV)) < V - mask_h = mask_k[None, :] & mask_v[:, None] - - b_h = tl.zeros([BV, BK], dtype=tl.float32) - - if USE_INITIAL_STATE: - p_h0 = h0 + i_nh * K * V + (tl.arange(0, BK)[None, :]) * V + ((i_v * BV + tl.arange(0, BV))[:, None]) - b_h += tl.load(p_h0, mask=mask_h, other=0).to(tl.float32) - - for _ in range(0, T): - b_k = tl.load(p_k, mask=mask_k, other=0).to(tl.float32) - b_v = tl.load(p_v, mask=mask_v, other=0).to(tl.float32) - b_q = tl.load(p_q, mask=mask_k, other=0).to(tl.float32) * scale - b_a = tl.load(p_a, mask=mask_k, other=0).to(tl.float32) - b_b = tl.load(p_b, mask=mask_k, other=0).to(tl.float32) - # to store - tmp = tl.sum(b_h * b_a[None, :], axis=1) - b_h += (tmp[:, None] * b_b[None, :] + b_k[None, :] * b_v[:, None]) - b_o = b_h * b_q[None, :] - b_o = tl.sum(b_o, axis=1) - tl.store(p_o, b_o.to(p_o.dtype.element_ty), mask=mask_v) - tl.store(p_ha, tmp.to(p_ha.dtype.element_ty), mask=mask_v) - p_q += K*H - p_k += K*H - p_o += V*H - p_v += V*H - p_ha += V*H - p_a += K*H - p_b += K*H - - if STORE_FINAL_STATE: - p_ht = ht + i_nh * K * V + (tl.arange(0, BK)[None, :]) * V + ((i_v * BV + tl.arange(0, BV))[:, None]) - tl.store(p_ht, b_h.to(p_ht.dtype.element_ty), mask=mask_h) - - -@triton.heuristics({ - 'USE_INITIAL_STATE': lambda args: args['h0'] is not None, - 'USE_DHT': lambda args: args['dht'] is not None, - 'USE_DH0': lambda args: args['dh0'] is not None, - 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None, -}) -@triton.autotune( - configs=[ - triton.Config({}, num_warps=num_warps, num_stages=num_stages) - for num_warps in [2, 4, 8, 16] - for num_stages in [2, 3] - ], - key=["BK", "BV"], - **autotune_cache_kwargs, -) -@triton.jit -def fused_recurrent_bwd_kernel( - # B: batch_size, H: n_heads, T: seq_len, D: b_dhead - # NV: number of split in the V dimension. NK: number of split in the K dimension - q, # query [B, H, L, K] - k, # key [B, H, L, V] - v, # value [B, H, L, V] - a, # a [B, H, L, K] - b, # b [B, H, L, K] - ha, # ha [B, H, L, V] - dht, # gradient of final state [B, H, K, V] - dh0, # gradient of initial state [B, H, K, V] - do, # gradient of output [B, H, L, V] - dq, # gradient of query [NV, B, H, L, K] - dk, # gradient of key [NV, B, H, L, K] - dv, # gradient of value [NK, B, H, L, V] - da, # gradient of a [NV, B, H, L, K] - db, # gradient of b [NV, B, H, L, K] - dha, # gradient of ha [NK, B, H, L, V] - h0, # initial state [B, H, K, V] - scale, # K ** -0.5 - cu_seqlens, # cu_seqlens - B, # batch_size - H, # n_heads - T, # seq_len - K: tl.constexpr, # K - V: tl.constexpr, # V - BK: tl.constexpr, # BLOCK SIZE along the K dimension - BV: tl.constexpr, # BLOCK SIZE along the V dimension - USE_INITIAL_STATE: tl.constexpr, # whether to use initial state h0 - USE_DH0: tl.constexpr, # whether to use dh0 - USE_DHT: tl.constexpr, # whether to use dht - IS_VARLEN: tl.constexpr, -): - i_v, i_nh = tl.program_id(0), tl.program_id(1) - i_n, i_h = i_nh // H, i_nh % H - dk += i_v * B * H * K * T - db += i_v * B * H * K * T - dq += i_v * B * H * K * T - da += i_v * B * H * K * T - if IS_VARLEN: - bos, eos = tl.load(cu_seqlens + i_n).to(tl.int64), tl.load(cu_seqlens + i_n + 1).to(tl.int64) - T = eos - bos - else: - bos, eos = i_n * T, i_n * T + T - mask_k = tl.arange(0, BK) < K - mask_v = (tl.arange(0, BV) + i_v * BV) < V - - q += (bos * H + i_h) * K - k += (bos * H + i_h) * K - v += (bos * H + i_h) * V + i_v * BV - ha += (bos * H + i_h) * V + i_v * BV - a += (bos * H + i_h) * K - b += (bos * H + i_h) * K - do += (bos * H + i_h) * V + i_v * BV - dq += (bos * H + i_h) * K - dk += (bos * H + i_h) * K - dv += (bos * H + i_h) * V + i_v * BV - da += (bos * H + i_h) * K - db += (bos * H + i_h) * K - dha += (bos * H + i_h) * V + i_v * BV - - p_q = q + tl.arange(0, BK) + (T - 1) * H*K - p_k = k + tl.arange(0, BK) + (T - 1) * H*K - p_v = v + tl.arange(0, BV) + (T - 1) * H*V - p_ha = ha + tl.arange(0, BV) + (T - 1) * H*V - p_a = a + tl.arange(0, BK) + (T - 1) * H*K - p_b = b + tl.arange(0, BK) + (T - 1) * H*K - p_do = do + tl.arange(0, BV) + (T - 1) * H*V - p_dk = dk + tl.arange(0, BK) + (T - 1) * H*K - p_dv = dv + tl.arange(0, BV) + (T - 1) * H*V - p_dha = dha + tl.arange(0, BV) + (T - 1) * H*V - p_db = db + tl.arange(0, BK) + (T - 1) * H*K - p_da = da + tl.arange(0, BK) + (T - 1) * H*K - p_dq = dq + tl.arange(0, BK) + (T - 1) * H*K - - b_dh = tl.zeros([BK, BV], dtype=tl.float32) - if USE_DHT: - p_ht = dht + i_nh * K * V + (tl.arange(0, BK)[:, None]) * V + ((i_v * BV + tl.arange(0, BV))[None, :]) - b_dh += tl.load(p_ht, mask=mask_k[:, None] & mask_v[None, :], other=0).to(tl.float32) - - for _ in range(T): - b_q = tl.load(p_q, mask=mask_k, other=0).to(tl.float32) * scale - b_k = tl.load(p_k, mask=mask_k, other=0).to(tl.float32) - b_v = tl.load(p_v, mask=mask_v, other=0).to(tl.float32) - b_do = tl.load(p_do, mask=mask_v, other=0).to(tl.float32) - b_b = tl.load(p_b, mask=mask_k, other=0).to(tl.float32) - b_a = tl.load(p_a, mask=mask_k, other=0).to(tl.float32) - b_ha = tl.load(p_ha, mask=mask_v, other=0).to(tl.float32) - - b_dh += b_q[:, None] * b_do[None, :] - d_k = tl.sum(b_dh * b_v[None, :], axis=1) - d_v = tl.sum(b_dh * b_k[:, None], axis=0) - tl.store(p_dk, d_k.to(p_dk.dtype.element_ty), mask=mask_k) - tl.store(p_dv, d_v.to(p_dv.dtype.element_ty), mask=mask_v) - - b_dha = tl.sum(b_dh * b_b[:, None], axis=0) - tl.store(p_dha, b_dha.to(p_dha.dtype.element_ty), mask=mask_v) - b_db = tl.sum(b_dh * b_ha[None, :], axis=1) - tl.store(p_db, b_db.to(p_db.dtype.element_ty), mask=mask_k) - - b_dh += b_dha[None, :] * b_a[:, None] - p_do -= H*V - p_q -= H*K - p_k -= H*K - p_v -= H*V - p_dk -= H*K - p_dv -= H*V - p_b -= H*K - p_db -= H*K - p_a -= H*K - p_dha -= H*V - p_ha -= H*V - - if USE_DH0: - p_dh0 = dh0 + i_nh * K * V + (tl.arange(0, BK)[:, None]) * V + (i_v * BV + tl.arange(0, BV)[None, :]) - tl.store(p_dh0, b_dh.to(p_dh0.dtype.element_ty), mask=mask_k[:, None] & mask_v[None, :]) - - tl.debug_barrier() - - b_h = tl.zeros([BK, BV], dtype=tl.float32) - - if USE_INITIAL_STATE: - mask_kv = mask_k[:, None] & mask_v[None, :] - p_h0 = h0 + i_nh * K * V + (tl.arange(0, BK)[:, None]) * V + ((i_v * BV + tl.arange(0, BV))[None, :]) - b_h += tl.load(p_h0, mask=mask_kv, other=0).to(tl.float32) - - p_k = k + tl.arange(0, BK) - p_v = v + tl.arange(0, BV) - p_ha = ha + tl.arange(0, BV) - p_do = do + tl.arange(0, BV) - p_dha = dha + tl.arange(0, BV) - p_da = da + tl.arange(0, BK) - p_dq = dq + tl.arange(0, BK) - p_b = b + tl.arange(0, BK) - - for i in range(0, T): - b_dha = tl.load(p_dha, mask=mask_v, other=0).to(tl.float32) - d_a = tl.sum(b_dha[None, :] * b_h, axis=1) - tl.store(p_da, d_a.to(p_da.dtype.element_ty), mask=mask_k) - b_k = tl.load(p_k, mask=mask_k, other=0).to(tl.float32) - b_v = tl.load(p_v, mask=mask_v, other=0).to(tl.float32) - b_do = tl.load(p_do, mask=mask_v, other=0).to(tl.float32) - b_b = tl.load(p_b, mask=mask_k, other=0).to(tl.float32) - b_ha = tl.load(p_ha, mask=mask_v, other=0).to(tl.float32) - b_h += b_k[:, None] * b_v[None, :] + b_b[:, None] * b_ha[None, :] - _d_q = b_h * b_do[None, :] - d_q = tl.sum(_d_q, axis=1) * scale - tl.store(p_dq, d_q.to(p_dq.dtype.element_ty), mask=mask_k) - - p_k += H*K - p_do += H*V - p_v += H*V - p_da += H*K - p_dha += H*V - p_ha += H*V - p_dq += H*K - p_b += H*K - - -class FusedRecurrentIPLRDeltaRuleFunction(torch.autograd.Function): - - @staticmethod - @input_guard - def forward( - ctx, - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - a: torch.Tensor, - b: torch.Tensor, - scale: float | None = None, - initial_state: torch.Tensor | None = None, - output_final_state: bool = False, - cu_seqlens: torch.LongTensor | None = None, - ): - B, T, H, K, V = *k.shape, v.shape[-1] - N = B if cu_seqlens is None else len(cu_seqlens) - 1 - - BK = triton.next_power_of_2(K) - if output_final_state: - final_state = q.new_empty(B, H, K, V, dtype=torch.float32) - else: - final_state = None - - ha = torch.empty_like(v, dtype=torch.float32) - - def grid(meta): return ( - triton.cdiv(V, meta['BV']), - N * H, - ) - o = torch.empty_like(v) - fused_recurrent_fwd_kernel[grid]( - q=q, - k=k, - v=v, - a=a, - b=b, - o=o, - ha=ha, - h0=initial_state, - ht=final_state, - scale=scale, - cu_seqlens=cu_seqlens, - H=H, - T=T, - K=K, - V=V, - BK=BK, - ) - ctx.save_for_backward(q, k, v, a, b, ha, initial_state) - ctx.scale = scale - ctx.cu_seqlens = cu_seqlens - return o, final_state - - @staticmethod - @input_guard - def backward(ctx, do, dht): - q, k, v, a, b, ha, initial_state = ctx.saved_tensors - B, T, H, K, V = *q.shape, v.shape[-1] - N = B if ctx.cu_seqlens is None else len(ctx.cu_seqlens) - 1 - BK, BV = triton.next_power_of_2(K), min(triton.next_power_of_2(V), 64) - NV = triton.cdiv(V, BV) - scale = ctx.scale - - dq = q.new_empty(NV, *q.shape) - dk = k.new_empty(NV, *k.shape) - da = a.new_empty(NV, *a.shape) - db = b.new_empty(NV, *b.shape) - dv = torch.empty_like(v) - dha = torch.empty_like(ha) - grid = (NV, N * H) - - if initial_state is not None and initial_state.requires_grad: - dh0 = torch.empty_like(initial_state, dtype=torch.float32) - else: - dh0 = None - - fused_recurrent_bwd_kernel[grid]( - q=q, - k=k, - v=v, - a=a, - b=b, - ha=ha, - dht=dht, - dh0=dh0, - do=do, - dq=dq, - dk=dk, - dv=dv, - da=da, - db=db, - dha=dha, - h0=initial_state, - scale=scale, - cu_seqlens=ctx.cu_seqlens, - B=B, - H=H, - T=T, - K=K, - V=V, - BK=BK, - BV=BV, - ) - dq = dq.sum(0) - dk = dk.sum(0) - da = da.sum(0) - db = db.sum(0) - return dq.to(q), dk.to(k), dv.to(v), da.to(a), db.to(b), None, dh0, None, None - - -def fused_recurrent_iplr_delta_rule( - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - a: torch.Tensor, - b: torch.Tensor, - scale: float = None, - initial_state: torch.Tensor = None, - output_final_state: bool = False, - cu_seqlens: torch.Tensor | None = None, -) -> tuple[torch.Tensor, torch.Tensor]: - r""" - This function computes the recurrence S_t = S_t @ (I + a_t b_t^T) + v_t k_t^T in a recurrent manner. - - Args: - q (torch.Tensor): - queries of shape `[B, T, H, K]` - k (torch.Tensor): - keys of shape `[B, T, H, K]` - v (torch.Tensor): - values of shape `[B, T, H, V]` - a (torch.Tensor): - as of shape `[B, T, H, K]` - b (torch.Tensor): - bs of shape `[B, T, H, K]` - scale (Optional[float]): - Scale factor for the RetNet attention scores. - If not provided, it will default to `1 / sqrt(K)`. Default: `None`. - initial_state (Optional[torch.Tensor]): - Initial state of shape `[B, H, K, V]`. Default: `None`. - output_final_state (Optional[bool]): - Whether to output the final state of shape `[B, H, K, V]`. Default: `False`. - cu_seqlens (torch.LongTensor): - Cumulative sequence lengths of shape `[N+1]` used for variable-length training, - consistent with the FlashAttention API. - - """ - if cu_seqlens is not None: - if q.shape[0] != 1: - raise ValueError( - f"The batch size is expected to be 1 rather than {q.shape[0]} when using `cu_seqlens`." - f"Please flatten variable-length inputs before processing.", - ) - if initial_state is not None and initial_state.shape[0] != len(cu_seqlens) - 1: - raise ValueError( - f"The number of initial states is expected to be equal to the number of input sequences, " - f"i.e., {len(cu_seqlens) - 1} rather than {initial_state.shape[0]}.", - ) - if scale is None: - scale = q.shape[-1] ** -0.5 - else: - assert scale > 0, "scale must be positive" - o, final_state = FusedRecurrentIPLRDeltaRuleFunction.apply( - q, - k, - v, - a, - b, - scale, - initial_state, - output_final_state, - cu_seqlens, - ) - return o, final_state diff --git a/build/torch-xpu/ops/generalized_delta_rule/iplr/naive.py b/build/torch-xpu/ops/generalized_delta_rule/iplr/naive.py deleted file mode 100644 index a8c54e70bf8f4369abf7ad9c277409ba77ee05c5..0000000000000000000000000000000000000000 --- a/build/torch-xpu/ops/generalized_delta_rule/iplr/naive.py +++ /dev/null @@ -1,74 +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 -from einops import rearrange - - -# S_t = S_t @ (I + alpha_t beta_t^T) + v_t k_t^T -# q, k, alpha, beta [B, H, L, D_K] -# v [B, H, L, D_V] -def iplr_recurrence(q, k, v, alpha, beta, initial_state=None, output_final_state=True): - orig_dtype = q.dtype - b, h, l, d_k = q.shape - q, k, v, beta = map(lambda x: x.float(), [q, k, v, beta]) - d_v = v.shape[-1] - o = torch.zeros_like(v) - S = torch.zeros(b, h, d_k, d_v).to(v) - q = q * (d_k ** -0.5) - - if initial_state is not None: - S += initial_state - - for i in range(l): - _k = k[:, :, i] - _q = q[:, :, i] - _v = v[:, :, i] - _alpha = alpha[:, :, i] - _beta = beta[:, :, i] - _kv = _k[..., None] * _v[..., None, :] + (S.clone() * _alpha[..., None]).sum(-2, keepdim=True) * _beta[..., None] - S = S + _kv - o[:, :, i] = torch.einsum('bhd,bhdm->bhm', _q, S) - S = None if output_final_state is False else S - return o.to(orig_dtype), S - - -def iplr_chunkwise(q, k, v, alpha, beta, initial_state=None, output_final_state=True, chunk_size=32): - b, h, l, d_k = q.shape - d_v = v.shape[-1] - q = q * (d_k ** -0.5) - v = v - assert l % chunk_size == 0 - - S = k.new_zeros(b, h, d_k, d_v) - if initial_state is not None: - S += initial_state - - # note that diagonal is masked. - mask = torch.triu(torch.ones(chunk_size, chunk_size, dtype=torch.bool, device=q.device), diagonal=0) - q, k, v, alpha, beta = map(lambda x: rearrange(x, 'b h (n c) d -> b h n c d', c=chunk_size), [q, k, v, alpha, beta]) - - v2 = (alpha @ k.transpose(-1, -2)).masked_fill_(mask, 0) @ v - attn = (alpha @ beta.transpose(-1, -2)).masked_fill(mask, 0) - for i in range(1, chunk_size): - attn[..., i, :i] = attn[..., i, :i] + (attn[..., i, :, None].clone() * attn[..., :, :i].clone()).sum(-2) - - attn = attn + torch.eye(chunk_size, dtype=torch.float, device=q.device) - u = attn @ v2 - w = attn @ alpha - o = torch.zeros_like(v) - mask = torch.triu(torch.ones(chunk_size, chunk_size, dtype=torch.bool, device=q.device), diagonal=1) - for i in range(0, l // chunk_size): - q_i, k_i, v_i, u_i, w_i, beta_i = q[:, :, i], k[:, :, i], v[:, :, i], u[:, :, i], w[:, :, i], beta[:, :, i] - o_1 = (q_i @ k_i.transpose(-1, -2)).masked_fill_(mask, 0) @ v_i - v2_i = u_i + w_i @ S - o_2 = (q_i @ beta_i.transpose(-1, -2)).masked_fill_(mask, 0) @ (v2_i) - o_3 = q_i @ S - o[:, :, i] = o_1 + o_2 + o_3 - S = S + k_i.transpose(-1, -2) @ v_i + beta_i.transpose(-1, -2) @ v2_i - S = None if output_final_state is False else S - return rearrange(o, 'b h n c d -> b h (n c) d'), S diff --git a/build/torch-xpu/ops/generalized_delta_rule/iplr/wy_fast.py b/build/torch-xpu/ops/generalized_delta_rule/iplr/wy_fast.py deleted file mode 100644 index 32c1ef769305fed000ff73fe9624be55ef07f106..0000000000000000000000000000000000000000 --- a/build/torch-xpu/ops/generalized_delta_rule/iplr/wy_fast.py +++ /dev/null @@ -1,309 +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 -import triton.language as tl - -from ....ops.utils import prepare_chunk_indices -from ....utils import IS_NVIDIA_HOPPER, autotune_cache_kwargs, check_shared_mem - -NUM_WARPS = [2, 4] if IS_NVIDIA_HOPPER else [2, 4, 8] - - -@triton.heuristics({ - 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None, -}) -@triton.autotune( - configs=[ - triton.Config({}, num_warps=num_warps) - for num_warps in [1, 2, 4, 8, 16] - ], - key=['BK'], - **autotune_cache_kwargs, -) -@triton.jit(do_not_specialize=['T']) -def prepare_wy_repr_fwd_kernel_chunk32( - a, - b, - A, - cu_seqlens, - chunk_indices, - T, - H: tl.constexpr, - K: tl.constexpr, - BT: tl.constexpr, - BK: tl.constexpr, - BC: tl.constexpr, # dummy placeholder - IS_VARLEN: tl.constexpr, -): - i_t, i_bh = tl.program_id(0), tl.program_id(1) - i_b, i_h = i_bh // H, i_bh % H - 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.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) - T = eos - bos - else: - bos, eos = i_b * T, i_b * T + T - - b_A = tl.zeros([BT, BT], dtype=tl.float32) - for i_k in range(tl.cdiv(K, BK)): - p_a = tl.make_block_ptr(a + (bos * H + i_h) * K, (T, K), (H*K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) - p_b = tl.make_block_ptr(b + (bos * H + i_h) * K, (K, T), (1, K*H), (i_k * BK, i_t * BT), (BK, BT), (0, 1)) - b_a = tl.load(p_a, boundary_check=(0, 1)) - b_b = tl.load(p_b, boundary_check=(0, 1)) - b_A += tl.dot(b_a, b_b) - - b_A = tl.where(tl.arange(0, BT)[:, None] > tl.arange(0, BT)[None, :], b_A, 0) - for i in range(1, BT): - mask = tl.arange(0, BT) == i - b_a = tl.sum(tl.where(mask[:, None], b_A, 0), 0) - b_a = b_a + tl.sum(b_a[:, None] * b_A, 0) * (tl.arange(0, BT) < i) - b_A = tl.where(mask[:, None], b_a, b_A) - b_A += tl.arange(0, BT)[:, None] == tl.arange(0, BT)[None, :] - - p_A = tl.make_block_ptr(A + (bos*H + i_h) * BT, (T, BT), (H*BT, 1), (i_t * BT, 0), (BT, BT), (1, 0)) - tl.store(p_A, b_A.to(p_A.dtype.element_ty), boundary_check=(0, 1)) - - -@triton.heuristics({ - 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None, -}) -@triton.autotune( - configs=[ - triton.Config({}, num_warps=num_warps) - for num_warps in [1, 2, 4, 8, 16] - ], - key=['BK'], - **autotune_cache_kwargs, -) -@triton.jit(do_not_specialize=['T']) -def prepare_wy_repr_fwd_kernel_chunk64( - a, - b, - A, - cu_seqlens, - chunk_indices, - T, - H: tl.constexpr, - K: tl.constexpr, - BT: tl.constexpr, - BK: tl.constexpr, - BC: tl.constexpr, - IS_VARLEN: tl.constexpr, -): - i_t, i_bh = tl.program_id(0), tl.program_id(1) - i_b, i_h = i_bh // H, i_bh % H - 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.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) - T = eos - bos - else: - bos, eos = i_b * T, i_b * T + T - - b_A = tl.zeros([BC, BC], dtype=tl.float32) - b_A2 = tl.zeros([BC, BC], dtype=tl.float32) - b_A3 = tl.zeros([BC, BC], dtype=tl.float32) - - for i_k in range(tl.cdiv(K, BK)): - p_a1 = tl.make_block_ptr(a + (bos * H + i_h) * K, (T, K), (H*K, 1), (i_t * BT, i_k * BK), (BC, BK), (1, 0)) - p_a2 = tl.make_block_ptr(a + (bos * H + i_h) * K, (T, K), (H*K, 1), (i_t * BT + BC, i_k * BK), (BC, BK), (1, 0)) - p_b1 = tl.make_block_ptr(b + (bos * H + i_h) * K, (K, T), (1, K*H), (i_k * BK, i_t * BT), (BK, BC), (0, 1)) - p_b2 = tl.make_block_ptr(b + (bos * H + i_h) * K, (K, T), (1, K*H), (i_k * BK, i_t * BT + BC), (BK, BC), (0, 1)) - b_a1 = tl.load(p_a1, boundary_check=(0, 1)) - b_a2 = tl.load(p_a2, boundary_check=(0, 1)) - b_b1 = tl.load(p_b1, boundary_check=(0, 1)) - b_b2 = tl.load(p_b2, boundary_check=(0, 1)) - b_A += tl.dot(b_a1, b_b1, allow_tf32=False) - b_A2 += tl.dot(b_a2, b_b2, allow_tf32=False) - b_A3 += tl.dot(b_a2, b_b1, allow_tf32=False) - - b_A = tl.where(tl.arange(0, BC)[:, None] > tl.arange(0, BC)[None, :], b_A, 0) - b_A2 = tl.where(tl.arange(0, BC)[:, None] > tl.arange(0, BC)[None, :], b_A2, 0) - - for i in range(1, BC): - mask = tl.arange(0, BC) == i - b_a = tl.sum(tl.where(mask[:, None], b_A, 0), 0) - b_a2 = tl.sum(tl.where(mask[:, None], b_A2, 0), 0) - b_a = b_a + tl.sum(b_a[:, None] * b_A, 0) * (tl.arange(0, BC) < i) - b_a2 = b_a2 + tl.sum(b_a2[:, None] * b_A2, 0) * (tl.arange(0, BC) < i) - b_A = tl.where(mask[:, None], b_a, b_A) - b_A2 = tl.where(mask[:, None], b_a2, b_A2) - - # blockwise computation of lower triangular matrix's inverse - # i.e., [A11, 0; A21, A22]^-1 = [A11^-1, 0; -A22^-1 A21 A11^-1, A22^-1] - b_A += tl.arange(0, BC)[:, None] == tl.arange(0, BC)[None, :] - b_A2 += tl.arange(0, BC)[:, None] == tl.arange(0, BC)[None, :] - b_A3 = tl.dot(tl.dot(b_A2, b_A3, allow_tf32=False), b_A, allow_tf32=False) - - p_A1 = tl.make_block_ptr(A + (bos*H + i_h) * BT, (T, BT), (H*BT, 1), (i_t * BT, 0), (BC, BC), (1, 0)) - p_A2 = tl.make_block_ptr(A + (bos*H + i_h) * BT, (T, BT), (H*BT, 1), (i_t * BT + BC, BC), (BC, BC), (1, 0)) - p_A3 = tl.make_block_ptr(A + (bos*H + i_h) * BT, (T, BT), (H*BT, 1), (i_t * BT + BC, 0), (BC, BC), (1, 0)) - p_A4 = tl.make_block_ptr(A + (bos*H + i_h) * BT, (T, BT), (H*BT, 1), (i_t * BT, BC), (BC, BC), (1, 0)) - tl.store(p_A1, b_A.to(p_A1.dtype.element_ty), boundary_check=(0, 1)) - tl.store(p_A2, b_A2.to(p_A2.dtype.element_ty), boundary_check=(0, 1)) - tl.store(p_A3, b_A3.to(p_A3.dtype.element_ty), boundary_check=(0, 1)) - # causal mask - tl.store(p_A4, tl.zeros([BC, BC], dtype=tl.float32).to(p_A4.dtype.element_ty), boundary_check=(0, 1)) - - -@triton.heuristics({ - 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None, -}) -@triton.autotune( - configs=[ - triton.Config({}, num_warps=num_warps) - for num_warps in NUM_WARPS - ], - key=['BT', 'BK', 'BV'], - **autotune_cache_kwargs, -) -@triton.jit(do_not_specialize=['T']) -def wu_fwd_kernel( - w, - u, - a, - k, - v, - A, - cu_seqlens, - chunk_indices, - T, - H: tl.constexpr, - K: tl.constexpr, - V: tl.constexpr, - BT: tl.constexpr, - BK: tl.constexpr, - BV: tl.constexpr, - IS_VARLEN: tl.constexpr, -): - i_t, i_bh = tl.program_id(0), tl.program_id(1) - i_b, i_h = i_bh // H, i_bh % H - 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.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) - T = eos - bos - else: - bos, eos = i_b * T, i_b * T + T - - p_A = tl.make_block_ptr(A + (bos*H + i_h) * BT, (T, BT), (H*BT, 1), (i_t * BT, 0), (BT, BT), (1, 0)) - - b_A = tl.load(p_A, boundary_check=(0, 1)) - b_Aak = tl.zeros([BT, BT], dtype=tl.float32) - - for i_k in range(tl.cdiv(K, BK)): - p_k = tl.make_block_ptr(k + (bos * H + i_h) * K, (T, K), (H*K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) - p_a = tl.make_block_ptr(a + (bos * H + i_h) * K, (T, K), (H*K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) - p_w = tl.make_block_ptr(w + (bos * H + i_h) * K, (T, K), (H*K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) - b_k = tl.load(p_k, boundary_check=(0, 1)) - b_a = tl.load(p_a, boundary_check=(0, 1)) - b_w = tl.dot(b_A, b_a) - b_Aak += tl.dot(b_a, tl.trans(b_k)) - tl.store(p_w, b_w.to(p_w.dtype.element_ty), boundary_check=(0, 1)) - - b_Aak = tl.where(tl.arange(0, BT)[:, None] > tl.arange(0, BT)[None, :], b_Aak, 0) - b_Aak = b_Aak.to(k.dtype.element_ty) - - for i_v in range(tl.cdiv(V, BV)): - p_v = tl.make_block_ptr(v + (bos*H + i_h) * V, (T, V), (H*V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) - p_u = tl.make_block_ptr(u + (bos*H + i_h) * V, (T, V), (H*V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) - b_v = tl.load(p_v, boundary_check=(0, 1)) - b_v = tl.dot(b_Aak, b_v).to(v.dtype.element_ty) - b_u = tl.dot(b_A, b_v) - tl.store(p_u, b_u.to(p_u.dtype.element_ty), boundary_check=(0, 1)) - - -def prepare_wy_repr_fwd( - a: torch.Tensor, - b: torch.Tensor, - v: torch.Tensor, - k: torch.Tensor, - cu_seqlens: torch.LongTensor | None, - chunk_size: int = 64, - chunk_indices: torch.LongTensor | None = None, -) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: - B, T, H, K = a.shape - BT = chunk_size - - if chunk_indices is None: - chunk_indices = prepare_chunk_indices(cu_seqlens, BT) if cu_seqlens is not None else None - NT = triton.cdiv(T, BT) if cu_seqlens is None else len(chunk_indices) - BC = min(BT, 32) - BK = min(max(triton.next_power_of_2(K), 16), 64) - - A = torch.empty(B, T, H, BT, device=a.device, dtype=a.dtype) - fwd_fn = prepare_wy_repr_fwd_kernel_chunk64 if BT == 64 else prepare_wy_repr_fwd_kernel_chunk32 - - fwd_fn[(NT, B * H)]( - a=a, - b=b, - A=A, - cu_seqlens=cu_seqlens, - chunk_indices=chunk_indices, - T=T, - H=H, - K=K, - BT=BT, - BK=BK, - BC=BC, - ) - w, u = wu_fwd( - a=a, - v=v, - k=k, - A=A, - cu_seqlens=cu_seqlens, - chunk_size=chunk_size, - chunk_indices=chunk_indices, - ) - return w, u, A - - -def wu_fwd( - a: torch.Tensor, - v: torch.Tensor, - k: torch.Tensor, - A: torch.Tensor, - cu_seqlens: torch.LongTensor | None, - chunk_size: int, - chunk_indices: torch.LongTensor | None = None, -) -> tuple[torch.Tensor, torch.Tensor]: - B, T, H, K, V = *a.shape, v.shape[-1] - BT = chunk_size - - if chunk_indices is None: - chunk_indices = prepare_chunk_indices(cu_seqlens, BT) if cu_seqlens is not None else None - NT = triton.cdiv(T, BT) if cu_seqlens is None else len(chunk_indices) - CONST_TILING = 64 if check_shared_mem() else 32 - BK = min(max(triton.next_power_of_2(K), 16), CONST_TILING) - BV = min(max(triton.next_power_of_2(V), 16), CONST_TILING) - - u = torch.empty_like(v) - w = torch.empty_like(a) - wu_fwd_kernel[(NT, B*H)]( - a=a, - v=v, - w=w, - u=u, - A=A, - k=k, - cu_seqlens=cu_seqlens, - chunk_indices=chunk_indices, - T=T, - H=H, - K=K, - V=V, - BT=BT, - BK=BK, - BV=BV, - ) - return w, u - - -fwd_prepare_wy_repr = prepare_wy_repr_fwd - -fwd_wu = wu_fwd diff --git a/build/torch-xpu/ops/gla/__init__.py b/build/torch-xpu/ops/gla/__init__.py deleted file mode 100644 index c8d92231609102511c83ceca484f2b3d5fdbb286..0000000000000000000000000000000000000000 --- a/build/torch-xpu/ops/gla/__init__.py +++ /dev/null @@ -1,16 +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 .chunk import chunk_gla -from .fused_chunk import fused_chunk_gla -from .fused_recurrent import fused_recurrent_gla - -__all__ = [ - 'chunk_gla', - 'fused_chunk_gla', - 'fused_recurrent_gla', -] diff --git a/build/torch-xpu/ops/gla/chunk.py b/build/torch-xpu/ops/gla/chunk.py deleted file mode 100644 index e39b6fec24c07f381d04c26e545224e030471cce..0000000000000000000000000000000000000000 --- a/build/torch-xpu/ops/gla/chunk.py +++ /dev/null @@ -1,1461 +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 -import triton.language as tl - -from ...ops.common.chunk_h import chunk_bwd_dh, chunk_fwd_h -from ...ops.utils import prepare_chunk_indices -from ...ops.utils.cache import fla_cache_autotune -from ...ops.utils.constant import RCP_LN2 -from ...ops.utils.cumsum import chunk_local_cumsum -from ...ops.utils.op import exp2 -from ...utils import autotune_cache_kwargs, check_shared_mem, input_guard - -BK_LIST = [32, 64] if check_shared_mem() else [16, 32] -BV_LIST = [64, 128] if check_shared_mem('ampere') else [16, 32] - - -def _prune_gla_bwd_configs(configs, nargs, **kwargs): - # Keep a tile only if it leaves headroom below its dim, or is the smallest - # option (so small dims still autotune); this avoids a software-pipelined - # make_block_ptr prefetching past the tensor. K/V arrive as launch kwargs. - args = {**(nargs or {}), **kwargs} - K, V = args['K'], args['V'] - min_bk = min(c.kwargs['BK'] for c in configs) - min_bv = min(c.kwargs['BV'] for c in configs) - return [ - c for c in configs - if (c.kwargs['BK'] < K or c.kwargs['BK'] == min_bk) - and (c.kwargs['BV'] < V or c.kwargs['BV'] == min_bv) - ] - - -@triton.heuristics({ - 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None, -}) -@fla_cache_autotune( - configs=[ - triton.Config({'BK': BK}, num_warps=num_warps, num_stages=num_stages) - for BK in [32, 64] - for num_warps in [1, 2, 4, 8] - for num_stages in [2, 3, 4] - ], - key=['BC'], - **autotune_cache_kwargs, -) -@triton.jit(do_not_specialize=['T']) -def chunk_gla_fwd_A_kernel_intra_sub_inter( - q, - k, - g, - A, - cu_seqlens, - chunk_indices, - scale, - T, - H: tl.constexpr, - K: tl.constexpr, - BT: tl.constexpr, - BC: tl.constexpr, - BK: tl.constexpr, - NC: tl.constexpr, - IS_VARLEN: tl.constexpr, -): - i_t, i_c, i_bh = tl.program_id(0), tl.program_id(1), tl.program_id(2) - i_b, i_h = i_bh // H, i_bh % H - i_i, i_j = i_c // NC, i_c % NC - 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.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) - T = eos - bos - else: - bos, eos = i_b * T, i_b * T + T - - if i_t * BT + i_i * BC >= T: - return - if i_i <= i_j: - return - - b_A = tl.zeros([BC, BC], dtype=tl.float32) - for i_k in range(tl.cdiv(K, BK)): - o_k = i_k * BK + tl.arange(0, BK) - m_k = o_k < K - - p_q = tl.make_block_ptr(q + (bos*H+i_h)*K, (T, K), (H*K, 1), (i_t * BT + i_i * BC, i_k * BK), (BC, BK), (1, 0)) - p_g = tl.make_block_ptr(g + (bos*H+i_h)*K, (T, K), (H*K, 1), (i_t * BT + i_i * BC, i_k * BK), (BC, BK), (1, 0)) - p_k = tl.make_block_ptr(k + (bos*H+i_h)*K, (K, T), (1, H*K), (i_k * BK, i_t * BT + i_j * BC), (BK, BC), (0, 1)) - p_gk = tl.make_block_ptr(g + (bos*H+i_h)*K, (K, T), (1, H*K), (i_k * BK, i_t * BT + i_j * BC), (BK, BC), (0, 1)) - p_gn = g + (bos + i_t * BT + i_i * BC) * H*K + i_h * K + o_k - - # [BK,] - b_gn = tl.load(p_gn, mask=m_k, other=0) - # [BC, BK] - b_q = tl.load(p_q, boundary_check=(0, 1)) - b_g = tl.load(p_g, boundary_check=(0, 1)) - b_qg = b_q * exp2(b_g - b_gn[None, :]) * scale - # [BK, BC] - b_k = tl.load(p_k, boundary_check=(0, 1)) - b_gk = tl.load(p_gk, boundary_check=(0, 1)) - b_kg = b_k * exp2(b_gn[:, None] - b_gk) - # [BC, BC] using tf32 to improve precision here. - b_A += tl.dot(b_qg, b_kg) - - p_A = tl.make_block_ptr(A + (bos*H + i_h)*BT, (T, BT), (H*BT, 1), (i_t * BT + i_i * BC, i_j * BC), (BC, BC), (1, 0)) - tl.store(p_A, b_A.to(A.dtype.element_ty), boundary_check=(0, 1)) - - -@triton.heuristics({ - 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None, -}) -@fla_cache_autotune( - configs=[ - triton.Config({}, num_warps=num_warps, num_stages=num_stages) - for num_warps in [1, 2, 4, 8] - for num_stages in [2, 3] - ], - key=['BK', 'BT'], - **autotune_cache_kwargs, -) -@triton.jit(do_not_specialize=['T']) -def chunk_gla_fwd_A_kernel_intra_sub_intra( - q, - k, - g, - A, - cu_seqlens, - chunk_indices, - scale, - T, - H: tl.constexpr, - K: tl.constexpr, - BT: tl.constexpr, - BC: tl.constexpr, - BK: tl.constexpr, - IS_VARLEN: tl.constexpr, -): - i_t, i_i, i_bh = tl.program_id(0), tl.program_id(1), tl.program_id(2) - i_b, i_h = i_bh // H, i_bh % H - i_j = i_i - 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.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) - T = eos - bos - else: - bos, eos = i_b * T, i_b * T + T - - if i_t * BT + i_i * BC >= T: - return - - o_i = tl.arange(0, BC) - o_k = tl.arange(0, BK) - o_A = (i_t * BT + i_i * BC + tl.arange(0, BC)) * H*BT + i_j * BC - m_k = o_k < K - m_A = (i_t * BT + i_i * BC + tl.arange(0, BC)) < T - - q += (bos * H + i_h) * K - k += (bos * H + i_h) * K - g += (bos * H + i_h) * K - A += (bos * H + i_h) * BT - - p_q = tl.make_block_ptr(q, (T, K), (H*K, 1), (i_t * BT + i_i * BC, 0), (BC, BK), (1, 0)) - p_g = tl.make_block_ptr(g, (T, K), (H*K, 1), (i_t * BT + i_i * BC, 0), (BC, BK), (1, 0)) - b_q = tl.load(p_q, boundary_check=(0, 1)) - b_g = tl.load(p_g, boundary_check=(0, 1)) - - p_k = k + (i_t * BT + i_j * BC) * H*K + o_k - p_gk = g + (i_t * BT + i_j * BC) * H*K + o_k - - for j in range(0, min(BC, T - i_t * BT - i_i * BC)): - b_k = tl.load(p_k, mask=m_k, other=0).to(tl.float32) - b_gk = tl.load(p_gk, mask=m_k, other=0).to(tl.float32) - b_A = tl.sum(b_q * b_k[None, :] * exp2(b_g - b_gk[None, :]), 1) * scale - tl.store(A + o_A + j, b_A, mask=m_A) - p_k += H*K - p_gk += H*K - - tl.debug_barrier() - b_A = tl.zeros([BC, BC], dtype=tl.float32) - tl.store(A + o_A[:, None] + o_i, b_A, mask=m_A[:, None] & (o_i[:, None] < o_i)) - - -@triton.heuristics({ - 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None, -}) -@fla_cache_autotune( - configs=[ - triton.Config({}, num_warps=num_warps) - for num_warps in [1, 2, 4, 8] - ], - key=['BC', 'BK'], - **autotune_cache_kwargs, -) -@triton.jit(do_not_specialize=['T']) -def chunk_gla_fwd_A_kernel_intra_sub_intra_split( - q, - k, - g, - A, - cu_seqlens, - chunk_indices, - scale, - T, - B: tl.constexpr, - H: tl.constexpr, - K: tl.constexpr, - BT: tl.constexpr, - BC: tl.constexpr, - BK: tl.constexpr, - NC: tl.constexpr, - IS_VARLEN: tl.constexpr, -): - i_k, i_tc, i_bh = tl.program_id(0), tl.program_id(1), tl.program_id(2) - i_b, i_h = i_bh // H, i_bh % H - i_t, i_i = i_tc // NC, i_tc % NC - i_j = i_i - 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.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) - all = T - T = eos - bos - else: - bos, eos = i_b * T, i_b * T + T - all = B * T - - if i_t * BT + i_i * BC >= T: - return - - o_i = tl.arange(0, BC) - o_k = i_k * BK + tl.arange(0, BK) - o_A = (i_t * BT + i_i * BC + tl.arange(0, BC)) * H*BC - m_k = o_k < K - m_A = (i_t * BT + i_i * BC + tl.arange(0, BC)) < T - - q += (bos * H + i_h) * K - k += (bos * H + i_h) * K - g += (bos * H + i_h) * K - A += ((i_k * all + bos) * H + i_h) * BC - - p_q = tl.make_block_ptr(q, (T, K), (H*K, 1), (i_t * BT + i_i * BC, i_k * BK), (BC, BK), (1, 0)) - p_g = tl.make_block_ptr(g, (T, K), (H*K, 1), (i_t * BT + i_i * BC, i_k * BK), (BC, BK), (1, 0)) - b_q = tl.load(p_q, boundary_check=(0, 1)) - b_g = tl.load(p_g, boundary_check=(0, 1)) - - p_k = k + (i_t * BT + i_j * BC) * H*K + o_k - p_gk = g + (i_t * BT + i_j * BC) * H*K + o_k - for j in range(0, min(BC, T - i_t * BT - i_i * BC)): - b_k = tl.load(p_k, mask=m_k, other=0).to(tl.float32) - b_gk = tl.load(p_gk, mask=m_k, other=0).to(tl.float32) - b_A = tl.sum(b_q * b_k[None, :] * exp2(b_g - b_gk[None, :]), 1) * scale - tl.store(A + o_A + j, b_A, mask=m_A) - p_k += H*K - p_gk += H*K - - tl.debug_barrier() - b_A = tl.zeros([BC, BC], dtype=tl.float32) - tl.store(A + o_A[:, None] + o_i, b_A, mask=m_A[:, None] & (o_i[:, None] < o_i)) - - -@triton.heuristics({ - 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None, -}) -@fla_cache_autotune( - configs=[ - triton.Config({}, num_warps=1), - triton.Config({}, num_warps=2), - triton.Config({}, num_warps=4), - triton.Config({}, num_warps=8), - ], - key=['BC'], - **autotune_cache_kwargs, -) -@triton.jit(do_not_specialize=['T']) -def chunk_gla_fwd_A_kernel_intra_sub_intra_merge( - A, - A2, - cu_seqlens, - chunk_indices, - T, - B: tl.constexpr, - H: tl.constexpr, - BT: tl.constexpr, - BC: tl.constexpr, - NK: tl.constexpr, - IS_VARLEN: tl.constexpr, -): - i_t, i_c, i_bh = tl.program_id(0), tl.program_id(1), tl.program_id(2) - i_b, i_h = i_bh // H, i_bh % H - 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.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) - all = T - T = eos - bos - else: - bos, eos = i_b * T, i_b * T + T - all = B * T - - if i_t * BT + i_c * BC >= T: - return - - b_A = tl.zeros([BC, BC], dtype=tl.float32) - for i_k in range(0, NK): - p_A = tl.make_block_ptr(A + (i_k*all+bos)*H*BC+i_h*BC, (T, BC), (H*BC, 1), (i_t*BT + i_c*BC, 0), (BC, BC), (1, 0)) - b_A += tl.load(p_A, boundary_check=(0, 1)) - p_A2 = tl.make_block_ptr(A2 + (bos*H+i_h)*BT, (T, BT), (H*BT, 1), (i_t * BT + i_c * BC, i_c * BC), (BC, BC), (1, 0)) - tl.store(p_A2, b_A.to(A2.dtype.element_ty), boundary_check=(0, 1)) - - -@triton.heuristics({ - 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None, -}) -@fla_cache_autotune( - configs=[ - triton.Config({'BK': BK, 'BV': BV}, num_warps=num_warps, num_stages=num_stages) - for BK in [32, 64] - for BV in [64, 128] - for num_warps in [2, 4, 8] - for num_stages in [2, 3, 4] - ], - key=['BT', 'HV', 'STATE_V_FIRST'], - **autotune_cache_kwargs, -) -@triton.jit(do_not_specialize=['T']) -def chunk_gla_fwd_kernel_o( - q, - v, - g, - h, - o, - A, - cu_seqlens, - chunk_indices, - scale, - T, - H: tl.constexpr, - HV: tl.constexpr, - K: tl.constexpr, - V: tl.constexpr, - BT: tl.constexpr, - BK: tl.constexpr, - BV: tl.constexpr, - STATE_V_FIRST: tl.constexpr, - IS_VARLEN: tl.constexpr, -): - i_v, i_t, i_bh = tl.program_id(0), tl.program_id(1), tl.program_id(2) - i_b, i_hv = i_bh // HV, i_bh % HV - i_h = i_hv // (HV // H) - if IS_VARLEN: - i_tg = i_t.to(tl.int64) - 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 - NT = tl.cdiv(T, BT) - else: - NT = tl.cdiv(T, BT) - i_tg = (i_b * NT + i_t).to(tl.int64) - bos, eos = (i_b * T).to(tl.int64), (i_b * T + T).to(tl.int64) - - m_s = tl.arange(0, BT)[:, None] >= tl.arange(0, BT)[None, :] - - q += (bos * H + i_h) * K - g += (bos * HV + i_hv) * K - v += (bos * HV + i_hv) * V - o += (bos * HV + i_hv) * V - h += (i_tg * HV + i_hv).to(tl.int64) * K * V - A += (bos * HV + i_hv) * BT - - b_o = tl.zeros([BT, BV], dtype=tl.float32) - for i_k in range(tl.cdiv(K, BK)): - p_q = tl.make_block_ptr(q, (T, K), (H*K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) - p_g = tl.make_block_ptr(g, (T, K), (HV*K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) - if STATE_V_FIRST: - p_h = tl.make_block_ptr(h, (V, K), (K, 1), (i_v * BV, i_k * BK), (BV, BK), (1, 0)) - else: - p_h = tl.make_block_ptr(h, (K, V), (V, 1), (i_k * BK, i_v * BV), (BK, BV), (1, 0)) - - # [BT, BK] - b_q = tl.load(p_q, boundary_check=(0, 1)) - # [BT, BK] - b_g = tl.load(p_g, boundary_check=(0, 1)).to(tl.float32) - # [BT, BK] - b_qg = (b_q * exp2(b_g)).to(b_q.dtype) - b_h = tl.load(p_h, boundary_check=(0, 1)) - if i_k >= 0: - if STATE_V_FIRST: - b_o += tl.dot(b_qg, tl.trans(b_h).to(b_qg.dtype)) - else: - b_o += tl.dot(b_qg, b_h.to(b_qg.dtype)) - b_o *= scale - p_v = tl.make_block_ptr(v, (T, V), (HV*V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) - p_o = tl.make_block_ptr(o, (T, V), (HV*V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) - p_A = tl.make_block_ptr(A, (T, BT), (HV*BT, 1), (i_t * BT, 0), (BT, BT), (1, 0)) - # [BT, BV] - b_v = tl.load(p_v, boundary_check=(0, 1)) - # [BT, BT] - b_A = tl.load(p_A, boundary_check=(0, 1)) - b_A = tl.where(m_s, b_A, 0.).to(b_v.dtype) - b_o += tl.dot(b_A, b_v) - tl.store(p_o, b_o.to(p_o.dtype.element_ty), boundary_check=(0, 1)) - - -@triton.heuristics({ - 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None, -}) -@fla_cache_autotune( - configs=[ - triton.Config({}, num_warps=num_warps, num_stages=num_stages) - for num_warps in [1, 2, 4, 8] - for num_stages in [2, 3, 4] - ], - key=['BK', 'NC', 'BT'], - **autotune_cache_kwargs, -) -@triton.jit(do_not_specialize=['T']) -def chunk_gla_bwd_kernel_intra( - q, - k, - g, - dA, - dq, - dk, - cu_seqlens, - chunk_indices, - T, - H: tl.constexpr, - K: tl.constexpr, - BT: tl.constexpr, - BC: tl.constexpr, - BK: tl.constexpr, - NC: tl.constexpr, - IS_VARLEN: tl.constexpr, -): - i_kc, i_t, i_bh = tl.program_id(0), tl.program_id(1), tl.program_id(2) - i_b, i_h = i_bh // H, i_bh % H - i_k, i_i = i_kc // NC, i_kc % NC - 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.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) - else: - bos, eos = i_b * T, i_b * T + T - T = eos - bos - if i_t * BT + i_i * BC >= T: - return - - o_k = i_k * BK + tl.arange(0, BK) - m_k = o_k < K - - p_g = tl.make_block_ptr(g + (bos*H + i_h) * K, (T, K), (H*K, 1), (i_t * BT + i_i * BC, i_k * BK), (BC, BK), (1, 0)) - # [BC, BK] - b_g = tl.load(p_g, boundary_check=(0, 1)) - - b_dq = tl.zeros([BC, BK], dtype=tl.float32) - if i_i > 0: - p_gn = g + (bos + i_t * BT + i_i * BC) * H*K + i_h*K + o_k - # [BK,] - b_gn = tl.load(p_gn, mask=m_k, other=0) - for i_j in range(0, i_i): - p_k = tl.make_block_ptr(k+(bos*H+i_h)*K, (T, K), (H*K, 1), (i_t*BT+i_j*BC, i_k * BK), (BC, BK), (1, 0)) - p_gk = tl.make_block_ptr(g+(bos*H+i_h)*K, (T, K), (H*K, 1), (i_t*BT+i_j*BC, i_k * BK), (BC, BK), (1, 0)) - p_dA = tl.make_block_ptr(dA+(bos*H+i_h)*BT, (T, BT), (H*BT, 1), (i_t*BT+i_i*BC, i_j * BC), (BC, BC), (1, 0)) - # [BC, BK] - b_k = tl.load(p_k, boundary_check=(0, 1)) - b_gk = tl.load(p_gk, boundary_check=(0, 1)) - b_kg = b_k * exp2(b_gn[None, :] - b_gk) - # [BC, BC] - b_dA = tl.load(p_dA, boundary_check=(0, 1)) - - b_dq += tl.dot(b_dA, b_kg) - b_dq *= exp2(b_g - b_gn[None, :]) - o_i = tl.arange(0, BC) - m_dA = (i_t * BT + i_i * BC + tl.arange(0, BC)) < T - o_dA = bos*H*BT + (i_t * BT + i_i * BC + tl.arange(0, BC)) * H*BT + i_h * BT + i_i * BC - p_kj = k + (bos + i_t * BT + i_i * BC) * H*K + i_h * K + o_k - p_gkj = g + (bos + i_t * BT + i_i * BC) * H*K + i_h * K + o_k - p_dq = tl.make_block_ptr(dq + (bos*H + i_h) * K, (T, K), (H*K, 1), (i_t * BT + i_i * BC, i_k * BK), (BC, BK), (1, 0)) - - for j in range(0, min(BC, T - i_t * BT - i_i * BC)): - # [BC,] - b_dA = tl.load(dA + o_dA + j, mask=m_dA, other=0) - # [BK,] - b_kj = tl.load(p_kj, mask=m_k, other=0).to(tl.float32) - b_gkj = tl.load(p_gkj, mask=m_k, other=0).to(tl.float32) - # [BC, BK] - m_i = o_i[:, None] >= j - # [BC, BK] - # (SY 09/17) important to not use bf16 here to have a good precision. - b_dq += tl.where(m_i, b_dA[:, None] * b_kj[None, :] * exp2(b_g - b_gkj[None, :]), 0.) - p_kj += H*K - p_gkj += H*K - tl.store(p_dq, b_dq.to(p_dq.dtype.element_ty), boundary_check=(0, 1)) - - tl.debug_barrier() - # [BC, BK] - b_dk = tl.zeros([BC, BK], dtype=tl.float32) - - NC = min(NC, tl.cdiv(T - i_t * BT, BC)) - if i_i < NC - 1: - p_gn = g + (bos + min(i_t * BT + i_i * BC + BC, T) - 1) * H*K + i_h * K + o_k - - # [BK,] - b_gn = tl.load(p_gn, mask=m_k, other=0) - for i_j in range(i_i + 1, NC): - p_q = tl.make_block_ptr(q + (bos*H+i_h)*K, (T, K), (H*K, 1), (i_t*BT+i_j*BC, i_k*BK), (BC, BK), (1, 0)) - p_gq = tl.make_block_ptr(g + (bos*H+i_h)*K, (T, K), (H*K, 1), (i_t*BT+i_j*BC, i_k*BK), (BC, BK), (1, 0)) - p_dA = tl.make_block_ptr(dA + (bos*H+i_h)*BT, (BT, T), (1, H*BT), (i_i*BC, i_t*BT + i_j*BC), (BC, BC), (0, 1)) - - o_j = i_t * BT + i_j * BC + o_i - m_j = o_j < T - # [BC, BK] - b_q = tl.load(p_q, boundary_check=(0, 1)) - b_gq = tl.load(p_gq, boundary_check=(0, 1)) - b_qg = b_q * tl.where(m_j[:, None], exp2(b_gq - b_gn[None, :]), 0) - # [BC, BC] - b_dA = tl.load(p_dA, boundary_check=(0, 1)) - # [BC, BK] - # (SY 09/17) important to not use bf16 here to have a good precision. - b_dk += tl.dot(b_dA, b_qg) - b_dk *= exp2(b_gn[None, :] - b_g) - o_dA = bos*H*BT + (i_t * BT + i_i * BC) * H*BT + i_h * BT + i_i * BC + tl.arange(0, BC) - p_qj = q + (bos + i_t * BT + i_i * BC) * H*K + i_h * K + o_k - p_gqj = g + (bos + i_t * BT + i_i * BC) * H*K + i_h * K + o_k - p_dk = tl.make_block_ptr(dk + (bos*H+i_h)*K, (T, K), (H*K, 1), (i_t * BT + i_i * BC, i_k * BK), (BC, BK), (1, 0)) - for j in range(0, min(BC, T - i_t * BT - i_i * BC)): - # [BC,] - b_dA = tl.load(dA + o_dA + j * H*BT) - # [BK,] - b_qj = tl.load(p_qj, mask=m_k, other=0).to(tl.float32) - b_gqj = tl.load(p_gqj, mask=m_k, other=0).to(tl.float32) - # [BC, BK] - m_i = o_i[:, None] <= j - b_dk += tl.where(m_i, b_dA[:, None] * b_qj[None, :] * exp2(b_gqj[None, :] - b_g), 0.) - p_qj += H*K - p_gqj += H*K - tl.store(p_dk, b_dk.to(p_dk.dtype.element_ty), boundary_check=(0, 1)) - - -@triton.heuristics({ - 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None, -}) -@fla_cache_autotune( - configs=[ - triton.Config({}, num_warps=num_warps, num_stages=num_stages) - for num_warps in [1, 2, 4, 8] - for num_stages in [2, 3, 4] - ], - key=['BV', 'BT'], - **autotune_cache_kwargs, -) -@triton.jit(do_not_specialize=['T']) -def chunk_gla_bwd_kernel_dA( - v, - do, - dA, - cu_seqlens, - chunk_indices, - scale, - T, - H: tl.constexpr, - V: tl.constexpr, - BT: tl.constexpr, - BV: tl.constexpr, - IS_VARLEN: tl.constexpr, -): - i_t, i_bh = tl.program_id(0), tl.program_id(1) - i_b, i_h = i_bh // H, i_bh % H - 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.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) - else: - bos, eos = i_b * T, i_b * T + T - T = eos - bos - - b_dA = tl.zeros([BT, BT], dtype=tl.float32) - for i_v in range(tl.cdiv(V, BV)): - p_do = tl.make_block_ptr(do + (bos*H + i_h) * V, (T, V), (H*V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) - p_v = tl.make_block_ptr(v + (bos*H + i_h) * V, (V, T), (1, H*V), (i_v * BV, i_t * BT), (BV, BT), (0, 1)) - b_v = tl.load(p_v, boundary_check=(0, 1)) - b_do = tl.load(p_do, boundary_check=(0, 1)) - - b_dA += tl.dot(b_do, b_v) - - p_dA = tl.make_block_ptr(dA + (bos * H + i_h) * BT, (T, BT), (H*BT, 1), (i_t * BT, 0), (BT, BT), (1, 0)) - m_s = tl.arange(0, BT)[:, None] >= tl.arange(0, BT)[None, :] - b_dA = tl.where(m_s, b_dA * scale, 0.) - tl.store(p_dA, b_dA.to(p_dA.dtype.element_ty), boundary_check=(0, 1)) - - -@triton.heuristics({ - 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None, -}) -@fla_cache_autotune( - configs=[ - triton.Config({'BK': BK, 'BV': BV}, num_warps=num_warps, num_stages=num_stages) - for BK in BK_LIST - for BV in BV_LIST - for num_warps in [2, 4, 8] - for num_stages in [2, 3, 4] - ], - key=['BT', 'STATE_V_FIRST', 'K', 'V'], - prune_configs_by={'early_config_prune': _prune_gla_bwd_configs}, - **autotune_cache_kwargs, -) -@triton.jit(do_not_specialize=['T']) -def chunk_gla_bwd_kernel_dv( - k, - g, - A, - do, - dh, - dv, - cu_seqlens, - chunk_indices, - T, - H: tl.constexpr, - K: tl.constexpr, - V: tl.constexpr, - BT: tl.constexpr, - BK: tl.constexpr, - BV: tl.constexpr, - IS_VARLEN: tl.constexpr, - STATE_V_FIRST: tl.constexpr, -): - i_v, i_t, i_bh = tl.program_id(0), tl.program_id(1), tl.program_id(2) - i_b, i_h = i_bh // H, i_bh % H - 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.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) - T = eos - bos - NT = tl.cdiv(T, BT) - else: - NT = tl.cdiv(T, BT) - i_tg = i_b * NT + i_t - bos, eos = i_b * T, i_b * T + T - - p_A = tl.make_block_ptr(A + (bos * H + i_h) * BT, (BT, T), (1, H*BT), (0, i_t * BT), (BT, BT), (0, 1)) - p_do = tl.make_block_ptr(do + (bos * H + i_h) * V, (T, V), (H*V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) - p_dv = tl.make_block_ptr(dv + (bos * H + i_h) * V, (T, V), (H*V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) - b_A = tl.load(p_A, boundary_check=(0, 1)) - b_do = tl.load(p_do, boundary_check=(0, 1)) - - b_A = tl.where(tl.arange(0, BT)[:, None] <= tl.arange(0, BT)[None, :], b_A, 0.) - # (SY 09/17) important to disallow tf32 here to maintain a good precision. - b_dv = tl.dot(b_A, b_do.to(b_A.dtype), allow_tf32=False) - - for i_k in range(tl.cdiv(K, BK)): - o_k = i_k * BK + tl.arange(0, BK) - m_k = o_k < K - - p_k = tl.make_block_ptr(k + (bos * H + i_h) * K, (T, K), (H*K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) - p_gk = tl.make_block_ptr(g + (bos * H + i_h) * K, (T, K), (H*K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) - p_gn = g + (bos + min(i_t * BT + BT, T) - 1)*H*K + i_h * K + o_k - if STATE_V_FIRST: - # dh stored as [V, K]; read a logical [BK, BV] tile via on-the-fly transpose - p_dh = tl.make_block_ptr(dh + (i_tg * H + i_h) * K*V, (K, V), (1, K), (i_k * BK, i_v * BV), (BK, BV), (0, 1)) - else: - p_dh = tl.make_block_ptr(dh + (i_tg * H + i_h) * K*V, (K, V), (V, 1), (i_k * BK, i_v * BV), (BK, BV), (1, 0)) - - b_k = tl.load(p_k, boundary_check=(0, 1)) - b_gk = tl.load(p_gk, boundary_check=(0, 1)) - b_dh = tl.load(p_dh, boundary_check=(0, 1)) - - b_gn = exp2(tl.load(p_gn, mask=m_k, other=0)[None, :] - b_gk) - b_k = (b_k * b_gn).to(b_k.dtype) - # [BT, BV] - # (SY 09/17) it is ok to have bf16 interchunk gradient contribution here - b_dv += tl.dot(b_k, b_dh.to(b_k.dtype)) - - tl.store(p_dv, b_dv.to(p_dv.dtype.element_ty), boundary_check=(0, 1)) - - -@triton.heuristics({ - 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None, -}) -@fla_cache_autotune( - configs=[ - triton.Config({'BK': BK, 'BV': BV}, num_warps=num_warps, num_stages=num_stages) - for BK in BK_LIST - for BV in BV_LIST - for num_warps in [2, 4, 8] - for num_stages in [2, 3, 4] - ], - key=['BT', 'STATE_V_FIRST', 'K', 'V'], - prune_configs_by={'early_config_prune': _prune_gla_bwd_configs}, - **autotune_cache_kwargs, -) -@triton.jit(do_not_specialize=['T']) -def chunk_gla_bwd_kernel_inter( - q, - k, - v, - g, - h, - do, - dh, - dq, - dk, - dq2, - dk2, - dg, - cu_seqlens, - chunk_indices, - scale, - T, - H: tl.constexpr, - K: tl.constexpr, - V: tl.constexpr, - BT: tl.constexpr, - BK: tl.constexpr, - BV: tl.constexpr, - IS_VARLEN: tl.constexpr, - STATE_V_FIRST: tl.constexpr, -): - i_k, i_t, i_bh = tl.program_id(0), tl.program_id(1), tl.program_id(2) - i_b, i_h = i_bh // H, i_bh % H - 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.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) - T = eos - bos - NT = tl.cdiv(T, BT) - else: - NT = tl.cdiv(T, BT) - i_tg = i_b * NT + i_t - bos, eos = i_b * T, i_b * T + T - o_k = i_k * BK + tl.arange(0, BK) - m_k = o_k < K - - q += (bos * H + i_h) * K - k += (bos * H + i_h) * K - v += (bos * H + i_h) * V - g += (bos * H + i_h) * K - h += (i_tg * H + i_h) * K*V - do += (bos * H + i_h) * V - dh += (i_tg * H + i_h) * K*V - dq += (bos * H + i_h) * K - dk += (bos * H + i_h) * K - dq2 += (bos * H + i_h) * K - dk2 += (bos * H + i_h) * K - dg += (bos * H + i_h) * K - - p_gk = tl.make_block_ptr(g, (T, K), (H*K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) - b_gk = tl.load(p_gk, boundary_check=(0, 1)) - p_gn = g + (min(T, i_t * BT + BT) - 1) * H*K + o_k - b_gn = tl.load(p_gn, mask=m_k, other=0) - b_dq = tl.zeros([BT, BK], dtype=tl.float32) - b_dk = tl.zeros([BT, BK], dtype=tl.float32) - b_dgk = tl.zeros([BK], dtype=tl.float32) - - for i_v in range(tl.cdiv(V, BV)): - p_v = tl.make_block_ptr(v, (T, V), (H*V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) - p_do = tl.make_block_ptr(do, (T, V), (H*V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) - if STATE_V_FIRST: - # h / dh stored as [V, K] -- the [BV, BK] tile is now a contiguous read - p_h = tl.make_block_ptr(h, (V, K), (K, 1), (i_v * BV, i_k * BK), (BV, BK), (1, 0)) - p_dh = tl.make_block_ptr(dh, (V, K), (K, 1), (i_v * BV, i_k * BK), (BV, BK), (1, 0)) - else: - p_h = tl.make_block_ptr(h, (V, K), (1, V), (i_v * BV, i_k * BK), (BV, BK), (0, 1)) - p_dh = tl.make_block_ptr(dh, (V, K), (1, V), (i_v * BV, i_k * BK), (BV, BK), (0, 1)) - # [BT, BV] - b_v = tl.load(p_v, boundary_check=(0, 1)) - b_do = tl.load(p_do, boundary_check=(0, 1)) - # [BV, BK] - b_h = tl.load(p_h, boundary_check=(0, 1)) - b_dh = tl.load(p_dh, boundary_check=(0, 1)) - - # [BK] - b_dgk += tl.sum(b_h * b_dh, axis=0) - # [BT, BK] - b_dq += tl.dot(b_do, b_h.to(b_do.dtype)) - b_dk += tl.dot(b_v, b_dh.to(b_v.dtype)) - - b_dgk *= exp2(b_gn) - b_dq *= scale - b_dq = b_dq * exp2(b_gk) - b_dk = b_dk * exp2(b_gn[None, :] - b_gk) - p_q = tl.make_block_ptr(q, (T, K), (H*K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) - p_k = tl.make_block_ptr(k, (T, K), (H*K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) - p_dq = tl.make_block_ptr(dq, (T, K), (H*K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) - p_dk = tl.make_block_ptr(dk, (T, K), (H*K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) - b_q = tl.load(p_q, boundary_check=(0, 1)) - b_k = tl.load(p_k, boundary_check=(0, 1)) - b_dgk += tl.sum(b_dk * b_k, axis=0) - b_dq += tl.load(p_dq, boundary_check=(0, 1)) - b_dk += tl.load(p_dk, boundary_check=(0, 1)) - b_dg = b_q * b_dq - b_k * b_dk - # tl.debug_barrier() - b_dg = b_dg - tl.cumsum(b_dg, axis=0) + tl.sum(b_dg, axis=0)[None, :] + b_dgk[None, :] - # Buggy due to strange triton compiler issue. - # m_s = tl.where(tl.arange(0, BT)[:, None] <= tl.arange(0, BT)[None, :], 1., 0.) - # b_dg = tl.dot(m_s, b_dg, allow_tf32=False) + b_dgk[None, :] - p_dq = tl.make_block_ptr(dq2, (T, K), (H*K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) - p_dk = tl.make_block_ptr(dk2, (T, K), (H*K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) - p_dg = tl.make_block_ptr(dg, (T, K), (H*K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) - tl.store(p_dq, b_dq.to(p_dq.dtype.element_ty), boundary_check=(0, 1)) - tl.store(p_dk, b_dk.to(p_dk.dtype.element_ty), boundary_check=(0, 1)) - tl.store(p_dg, b_dg.to(p_dg.dtype.element_ty), boundary_check=(0, 1)) - - -def chunk_gla_fwd_intra_gk( - q: torch.Tensor, - k: torch.Tensor, - g: torch.Tensor, - scale: float, - cu_seqlens: torch.LongTensor | None = None, - chunk_size: int = 64, - chunk_indices: torch.LongTensor | None = None, -): - B, T, H, K = k.shape - BT = chunk_size - - if chunk_indices is None and cu_seqlens is not None: - chunk_indices = prepare_chunk_indices(cu_seqlens, chunk_size) - NT = triton.cdiv(T, BT) if cu_seqlens is None else len(chunk_indices) - BC = min(16, BT) - NC = triton.cdiv(BT, BC) - - A = q.new_empty(B, T, H, BT, dtype=torch.float) - grid = (NT, NC * NC, B * H) - chunk_gla_fwd_A_kernel_intra_sub_inter[grid]( - q=q, - k=k, - g=g, - A=A, - cu_seqlens=cu_seqlens, - chunk_indices=chunk_indices, - scale=scale, - T=T, - H=H, - K=K, - BT=BT, - BC=BC, - NC=NC, - ) - - grid = (NT, NC, B * H) - # load the entire [BC, K] blocks into SRAM at once - if K <= 256: - BK = max(triton.next_power_of_2(K), 16) - chunk_gla_fwd_A_kernel_intra_sub_intra[grid]( - q=q, - k=k, - g=g, - A=A, - cu_seqlens=cu_seqlens, - chunk_indices=chunk_indices, - scale=scale, - T=T, - H=H, - K=K, - BT=BT, - BC=BC, - BK=BK, - ) - # split then merge - else: - BK = min(128, triton.next_power_of_2(K)) - NK = triton.cdiv(K, BK) - A_intra = q.new_empty(NK, B, T, H, BC, dtype=torch.float) - - grid = (NK, NT * NC, B * H) - chunk_gla_fwd_A_kernel_intra_sub_intra_split[grid]( - q=q, - k=k, - g=g, - A=A_intra, - cu_seqlens=cu_seqlens, - chunk_indices=chunk_indices, - scale=scale, - T=T, - B=B, - H=H, - K=K, - BT=BT, - BC=BC, - BK=BK, - NC=NC, - ) - - grid = (NT, NC, B * H) - chunk_gla_fwd_A_kernel_intra_sub_intra_merge[grid]( - A=A_intra, - A2=A, - cu_seqlens=cu_seqlens, - chunk_indices=chunk_indices, - T=T, - B=B, - H=H, - BT=BT, - BC=BC, - NK=NK, - ) - return A - - -def chunk_gla_fwd_o_gk( - q: torch.Tensor, - v: torch.Tensor, - g: torch.Tensor, - A: torch.Tensor, - h: torch.Tensor, - scale: float, - state_v_first: bool = False, - cu_seqlens: torch.LongTensor | None = None, - chunk_size: int = 64, - chunk_indices: torch.LongTensor | None = None, -): - B, T, H, K, HV, V = *q.shape, v.shape[2], v.shape[-1] - BT = chunk_size - - if chunk_indices is None and cu_seqlens is not None: - chunk_indices = prepare_chunk_indices(cu_seqlens, chunk_size) - NT = triton.cdiv(T, BT) if cu_seqlens is None else len(chunk_indices) - - # Please ensure zeros, since vllm will use padding v - o = torch.zeros_like(v) - def grid(meta): return (triton.cdiv(V, meta['BV']), NT, B * HV) - chunk_gla_fwd_kernel_o[grid]( - q=q, - v=v, - g=g, - h=h, - o=o, - A=A, - cu_seqlens=cu_seqlens, - chunk_indices=chunk_indices, - scale=scale, - T=T, - H=H, - HV=HV, - K=K, - V=V, - BT=BT, - STATE_V_FIRST=state_v_first, - ) - return o - - -def chunk_gla_bwd_dA( - v: torch.Tensor, - do: torch.Tensor, - scale: float, - cu_seqlens: torch.LongTensor | None = None, - chunk_size: int = 64, - chunk_indices: torch.LongTensor | None = None, -): - B, T, H, V = v.shape - BT = chunk_size - - if chunk_indices is None and cu_seqlens is not None: - chunk_indices = prepare_chunk_indices(cu_seqlens, chunk_size) - NT = triton.cdiv(T, BT) if cu_seqlens is None else len(chunk_indices) - BV = min(64, triton.next_power_of_2(V)) - - dA = v.new_empty(B, T, H, BT, dtype=torch.float) - grid = (NT, B * H) - chunk_gla_bwd_kernel_dA[grid]( - v=v, - do=do, - dA=dA, - cu_seqlens=cu_seqlens, - chunk_indices=chunk_indices, - scale=scale, - T=T, - H=H, - V=V, - BT=BT, - BV=BV, - ) - return dA - - -def chunk_gla_bwd_dv( - k: torch.Tensor, - g: torch.Tensor, - A: torch.Tensor, - do: torch.Tensor, - dh: torch.Tensor, - state_v_first: bool = False, - cu_seqlens: torch.LongTensor | None = None, - chunk_size: int = 64, - chunk_indices: torch.LongTensor | None = None, -): - B, T, H, K, V = *k.shape, do.shape[-1] - BT = chunk_size - - if chunk_indices is None and cu_seqlens is not None: - chunk_indices = prepare_chunk_indices(cu_seqlens, chunk_size) - NT = triton.cdiv(T, BT) if cu_seqlens is None else len(chunk_indices) - - dv = torch.empty_like(do) - def grid(meta): return (triton.cdiv(V, meta['BV']), NT, B * H) - chunk_gla_bwd_kernel_dv[grid]( - k=k, - g=g, - A=A, - do=do, - dh=dh, - dv=dv, - cu_seqlens=cu_seqlens, - chunk_indices=chunk_indices, - T=T, - H=H, - K=K, - V=V, - BT=BT, - STATE_V_FIRST=state_v_first, - ) - return dv - - -def chunk_gla_bwd_dqk_intra( - q: torch.Tensor, - k: torch.Tensor, - g: torch.Tensor, - dA: torch.Tensor, - cu_seqlens: torch.LongTensor | None = None, - chunk_size: int = 64, - chunk_indices: torch.LongTensor | None = None, -): - B, T, H, K = q.shape - BT = chunk_size - BC = min(16, BT) - BK = min(64, triton.next_power_of_2(K)) - - if chunk_indices is None and cu_seqlens is not None: - chunk_indices = prepare_chunk_indices(cu_seqlens, chunk_size) - NT = triton.cdiv(T, BT) if cu_seqlens is None else len(chunk_indices) - NC = triton.cdiv(BT, BC) - NK = triton.cdiv(K, BK) - - dq = torch.empty_like(q, dtype=torch.float) - dk = torch.empty_like(k, dtype=torch.float) - grid = (NK * NC, NT, B * H) - chunk_gla_bwd_kernel_intra[grid]( - q=q, - k=k, - g=g, - dA=dA, - dq=dq, - dk=dk, - cu_seqlens=cu_seqlens, - chunk_indices=chunk_indices, - T=T, - H=H, - K=K, - BT=BT, - BC=BC, - BK=BK, - NC=NC, - ) - return dq, dk - - -def chunk_gla_bwd_dqkg( - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - h: torch.Tensor, - g: torch.Tensor, - do: torch.Tensor, - dh: torch.Tensor, - dq: torch.Tensor, - dk: torch.Tensor, - scale: float | None = None, - state_v_first: bool = False, - cu_seqlens: torch.LongTensor | None = None, - chunk_size: int = 64, - chunk_indices: torch.LongTensor | None = None, -): - B, T, H, K, V = *k.shape, v.shape[-1] - BT = chunk_size - - if chunk_indices is None and cu_seqlens is not None: - chunk_indices = prepare_chunk_indices(cu_seqlens, chunk_size) - NT = triton.cdiv(T, BT) if cu_seqlens is None else len(chunk_indices) - - dg = torch.empty_like(g) - dq2 = torch.empty_like(dq) - dk2 = torch.empty_like(dk) - def grid(meta): return (triton.cdiv(K, meta['BK']), NT, B * H) - chunk_gla_bwd_kernel_inter[grid]( - q=q, - k=k, - v=v, - g=g, - h=h, - do=do, - dh=dh, - dq=dq, - dk=dk, - dq2=dq2, - dk2=dk2, - dg=dg, - cu_seqlens=cu_seqlens, - chunk_indices=chunk_indices, - scale=scale, - T=T, - H=H, - K=K, - V=V, - BT=BT, - STATE_V_FIRST=state_v_first, - ) - return dq2, dk2, dg - - -def chunk_gla_fwd( - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - g: torch.Tensor, - g_cumsum: torch.Tensor | None, - scale: float, - initial_state: torch.Tensor, - output_final_state: bool, - state_v_first: bool = False, - cu_seqlens: torch.LongTensor | None = None, - chunk_size: int = 64, - chunk_indices: torch.LongTensor | None = None, -) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: - if g_cumsum is None: - g_cumsum = chunk_local_cumsum( - g, - chunk_size, - scale=RCP_LN2, - cu_seqlens=cu_seqlens, - ) - - h, ht = chunk_fwd_h( - k=k, - v=v, - g=None, - gk=g_cumsum, - gv=None, - h0=initial_state, - output_final_state=output_final_state, - cu_seqlens=cu_seqlens, - chunk_size=chunk_size, - states_in_fp32=False, - state_v_first=state_v_first, - ) - - # the intra A is kept in fp32 - # the computation has very marginal effect on the entire throughput - A = chunk_gla_fwd_intra_gk( - q=q, - k=k, - g=g_cumsum, - scale=scale, - cu_seqlens=cu_seqlens, - chunk_size=chunk_size, - chunk_indices=chunk_indices, - ) - o = chunk_gla_fwd_o_gk( - q=q, - v=v, - g=g_cumsum, - A=A, - h=h, - scale=scale, - cu_seqlens=cu_seqlens, - chunk_size=chunk_size, - chunk_indices=chunk_indices, - state_v_first=state_v_first, - ) - return g_cumsum, A, h, ht, o - - -def chunk_gla_bwd( - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - g: torch.Tensor, - g_cumsum: torch.Tensor | None, - scale: float, - initial_state: torch.Tensor, - h: torch.Tensor, - A: torch.Tensor, - do: torch.Tensor, - dht: torch.Tensor, - state_v_first: bool = False, - cu_seqlens: torch.LongTensor | None = None, - chunk_size: int = 64, - chunk_indices: torch.LongTensor | None = None, -): - if g_cumsum is None: - g_cumsum = chunk_local_cumsum( - g, - chunk_size, - scale=RCP_LN2, - cu_seqlens=cu_seqlens, - ) - - if h is None: - h, _ = chunk_fwd_h( - k=k, - v=v, - g=None, - gk=g_cumsum, - gv=None, - h0=initial_state, - output_final_state=False, - cu_seqlens=cu_seqlens, - chunk_size=chunk_size, - states_in_fp32=True, - state_v_first=state_v_first, - ) - dh, dh0 = chunk_bwd_dh( - q=q, - k=k, - v=v, - g=None, - gk=g_cumsum, - gv=None, - do=do, - h0=initial_state, - dht=dht, - scale=scale, - cu_seqlens=cu_seqlens, - chunk_size=chunk_size, - states_in_fp32=True, - state_v_first=state_v_first, - ) - - dv = chunk_gla_bwd_dv( - k=k, - g=g_cumsum, - A=A, - do=do, - dh=dh, - cu_seqlens=cu_seqlens, - chunk_size=chunk_size, - chunk_indices=chunk_indices, - state_v_first=state_v_first, - ) - - # dq dk in fp32 - dA = chunk_gla_bwd_dA( - v=v, - do=do, - scale=scale, - cu_seqlens=cu_seqlens, - chunk_size=chunk_size, - chunk_indices=chunk_indices, - ) - dq, dk = chunk_gla_bwd_dqk_intra( - q=q, - k=k, - g=g_cumsum, - dA=dA, - cu_seqlens=cu_seqlens, - chunk_size=chunk_size, - chunk_indices=chunk_indices, - ) - dq, dk, dg = chunk_gla_bwd_dqkg( - q=q, - k=k, - v=v, - h=h, - g=g_cumsum, - do=do, - dh=dh, - dq=dq, - dk=dk, - scale=scale, - cu_seqlens=cu_seqlens, - chunk_size=chunk_size, - chunk_indices=chunk_indices, - state_v_first=state_v_first, - ) - return dq, dk, dv, dg, dh0 - - -class ChunkGLAFunction(torch.autograd.Function): - - @staticmethod - @input_guard - def forward( - ctx, - q, - k, - v, - g, - scale, - initial_state, - output_final_state, - state_v_first, - cu_seqlens, - cu_seqlens_cpu, - ): - chunk_size = min(64, max(16, triton.next_power_of_2(q.shape[1]))) - if cu_seqlens is not None: - chunk_indices = prepare_chunk_indices( - cu_seqlens, - chunk_size, - cu_seqlens_cpu=cu_seqlens_cpu, - ) - else: - chunk_indices = None - - g_cumsum, A, _, ht, o = chunk_gla_fwd( - q=q, - k=k, - v=v, - g=g, - g_cumsum=None, - scale=scale, - initial_state=initial_state, - output_final_state=output_final_state, - cu_seqlens=cu_seqlens, - chunk_size=chunk_size, - chunk_indices=chunk_indices, - state_v_first=state_v_first, - ) - # recompute g_cumsum in bwd pass - if g.dtype != torch.float: - g_cumsum = None - else: - g = None - ctx.save_for_backward(q, k, v, g, g_cumsum, initial_state, A, chunk_indices) - ctx.chunk_size = chunk_size - ctx.scale = scale - ctx.cu_seqlens = cu_seqlens - ctx.state_v_first = state_v_first - return o, ht - - @staticmethod - @input_guard - def backward(ctx, do, dht): - q, k, v, g, g_cumsum, initial_state, A, chunk_indices = ctx.saved_tensors - chunk_size, scale, cu_seqlens = ctx.chunk_size, ctx.scale, ctx.cu_seqlens - dq, dk, dv, dg, dh0 = chunk_gla_bwd( - q=q, - k=k, - v=v, - g=g, - g_cumsum=g_cumsum, - scale=scale, - h=None, - A=A, - initial_state=initial_state, - do=do, - dht=dht, - cu_seqlens=cu_seqlens, - chunk_size=chunk_size, - chunk_indices=chunk_indices, - state_v_first=ctx.state_v_first, - ) - return dq.to(q), dk.to(k), dv.to(v), dg, None, dh0, None, None, None, None - - -@torch.compiler.disable -def chunk_gla( - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - g: torch.Tensor, - scale: int | None = None, - initial_state: torch.Tensor = None, - output_final_state: bool = False, - state_v_first: bool = False, - cu_seqlens: torch.LongTensor | None = None, - cu_seqlens_cpu: torch.LongTensor | None = None, -) -> tuple[torch.Tensor, torch.Tensor]: - r""" - Args: - q (torch.Tensor): - queries of shape `[B, T, H, K]`. - k (torch.Tensor): - keys of shape `[B, T, H, K]`. - v (torch.Tensor): - values of shape `[B, T, H, V]`. - g (torch.Tensor): - Forget gates of shape `[B, T, H, K]`. - scale (Optional[float]): - Scale factor for the attention scores. - If not provided, it will default to `1 / sqrt(K)`. Default: `None`. - initial_state (Optional[torch.Tensor]): - Initial state of shape `[N, H, K, V]` (or `[N, H, V, K]` if `state_v_first=True`) - for `N` input sequences. - For equal-length input sequences, `N` equals the batch size `B`. - Default: `None`. - output_final_state (Optional[bool]): - Whether to output the final state of shape `[N, H, K, V]` - (or `[N, H, V, K]` if `state_v_first=True`). Default: `False`. - state_v_first (Optional[bool]): - Store the recurrent state in V-first `[V, K]` layout instead of the default `[K, V]`. Default: `False`. - cu_seqlens (torch.LongTensor): - Cumulative sequence lengths of shape `[N+1]` used for variable-length training, - consistent with the FlashAttention API. - - Returns: - o (torch.Tensor): - Outputs of shape `[B, T, H, V]`. - final_state (torch.Tensor): - Final state of shape `[N, H, K, V]` (or `[N, H, V, K]` if `state_v_first=True`) - if `output_final_state=True` else `None`. - - Examples:: - >>> import torch - >>> import torch.nn.functional as F - >>> from einops import rearrange - >>> from ...ops.gla import chunk_gla - # inputs with equal lengths - >>> B, T, H, K, V = 4, 2048, 4, 512, 512 - >>> q = torch.randn(B, T, H, K, device='cuda') - >>> k = torch.randn(B, T, H, K, device='cuda') - >>> v = torch.randn(B, T, H, V, device='cuda') - >>> g = F.logsigmoid(torch.randn(B, T, H, K, device='cuda')) - >>> h0 = torch.randn(B, H, K, V, device='cuda') - >>> o, ht = chunk_gla( - q, k, v, g, - initial_state=h0, - output_final_state=True - ) - # for variable-length inputs, the batch size `B` is expected to be 1 and `cu_seqlens` is required - >>> q, k, v, g = map(lambda x: rearrange(x, 'b t h d -> 1 (b t) h d'), (q, k, v, g)) - # for a batch with 4 sequences, `cu_seqlens` with 5 start/end positions are expected - >>> cu_seqlens = q.new_tensor([0, 2048, 4096, 6144, 8192], dtype=torch.long) - >>> o, ht = chunk_gla( - q, k, v, g, - initial_state=h0, - output_final_state=True, - cu_seqlens=cu_seqlens - ) - """ - if cu_seqlens is not None: - if q.shape[0] != 1: - raise ValueError( - f"The batch size is expected to be 1 rather than {q.shape[0]} when using `cu_seqlens`." - f"Please flatten variable-length inputs before processing.", - ) - if initial_state is not None and initial_state.shape[0] != len(cu_seqlens) - 1: - raise ValueError( - f"The number of initial states is expected to be equal to the number of input sequences, " - f"i.e., {len(cu_seqlens) - 1} rather than {initial_state.shape[0]}.", - ) - if scale is None: - scale = q.shape[-1] ** -0.5 - if initial_state is not None: - assert initial_state.dtype == torch.float32, "initial_state must be in float32." - assert q.shape == k.shape == g.shape, "q, k, g must have the same shape." - assert v.shape == (*q.shape[:3], v.shape[-1]), "v must be of shape (batch size, seq len, num of head, head dim)." - o, final_state = ChunkGLAFunction.apply( - q, - k, - v, - g, - scale, - initial_state, - output_final_state, - state_v_first, - cu_seqlens, - cu_seqlens_cpu, - ) - return o, final_state diff --git a/build/torch-xpu/ops/gla/fused_chunk.py b/build/torch-xpu/ops/gla/fused_chunk.py deleted file mode 100644 index 21a5a3de173011cd7acb3c50a35743244fb40b94..0000000000000000000000000000000000000000 --- a/build/torch-xpu/ops/gla/fused_chunk.py +++ /dev/null @@ -1,11 +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 - -def fused_chunk_gla( - **kwargs, -): - raise NotImplementedError("`fused_chunk_gla` is deprecated. Please use `chunk_gla` instead.") diff --git a/build/torch-xpu/ops/gla/fused_recurrent.py b/build/torch-xpu/ops/gla/fused_recurrent.py deleted file mode 100644 index d6d695a9ec95bab3912634ccf3903eb9f932ebd5..0000000000000000000000000000000000000000 --- a/build/torch-xpu/ops/gla/fused_recurrent.py +++ /dev/null @@ -1,132 +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 warnings - -import torch - -from ...ops.common.fused_recurrent import fused_recurrent - - -def fused_recurrent_gla( - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - gk: torch.Tensor | None = None, - gv: torch.Tensor | None = None, - scale: int | None = None, - initial_state: torch.Tensor | None = None, - output_final_state: bool = False, - reverse: bool = False, - state_v_first: bool = False, - cu_seqlens: torch.LongTensor | None = None, - **kwargs, -) -> tuple[torch.Tensor, torch.Tensor]: - r""" - Args: - q (torch.Tensor): - queries of shape `[B, T, H, K]`. - k (torch.Tensor): - keys of shape `[B, T, H, K]`. - v (torch.Tensor): - values of shape `[B, T, H, V]`. - gk (torch.Tensor): - Forget gates of shape `[B, T, H, K]`. - gv (torch.Tensor): - Forget gates of shape `[B, T, H, V]` applied to values. - scale (Optional[float]): - Scale factor for the attention scores. - If not provided, it will default to `1 / sqrt(K)`. Default: `None`. - initial_state (Optional[torch.Tensor]): - Initial state of shape `[N, H, K, V]` for `N` input sequences. - For equal-length input sequences, `N` equals the batch size `B`. - Default: `None`. - output_final_state (Optional[bool]): - Whether to output the final state of shape `[N, H, K, V]`. Default: `False`. - reverse (Optional[bool]): - If `True`, process the state passing in reverse order. Default: `False`. - state_v_first (Optional[bool]): - Store the recurrent state in V-first ``[V, K]`` layout instead of the default ``[K, V]``. Default: ``False``. - cu_seqlens (torch.LongTensor): - Cumulative sequence lengths of shape `[N+1]` used for variable-length training, - consistent with the FlashAttention API. - `initial_state`, `final_state`, and `dh0`/`dht` all follow the chosen layout. - Default: `False`. - - Returns: - o (torch.Tensor): - Outputs of shape `[B, T, H, V]`. - final_state (torch.Tensor): - Final state of shape `[N, H, K, V]` if `output_final_state=True` else `None`. - - Examples:: - >>> import torch - >>> import torch.nn.functional as F - >>> from einops import rearrange - >>> from ...ops.gla import fused_recurrent_gla - # inputs with equal lengths - >>> B, T, H, K, V = 4, 2048, 4, 512, 512 - >>> q = torch.randn(B, T, H, K, device='cuda') - >>> k = torch.randn(B, T, H, K, device='cuda') - >>> v = torch.randn(B, T, H, V, device='cuda') - >>> g = F.logsigmoid(torch.randn(B, T, H, K, device='cuda')) - >>> h0 = torch.randn(B, H, K, V, device='cuda') - >>> o, ht = fused_recurrent_gla( - q, k, v, g, - initial_state=h0, - output_final_state=True - ) - # for variable-length inputs, the batch size `B` is expected to be 1 and `cu_seqlens` is required - >>> q, k, v, g = map(lambda x: rearrange(x, 'b t h d -> 1 (b t) h d'), (q, k, v, g)) - # for a batch with 4 sequences, `cu_seqlens` with 5 start/end positions are expected - >>> cu_seqlens = q.new_tensor([0, 2048, 4096, 6144, 8192], dtype=torch.long) - >>> o_var, ht_var = fused_recurrent_gla( - q, k, v, g, - initial_state=h0, - output_final_state=True, - cu_seqlens=cu_seqlens - ) - >>> assert o.allclose(o_var.view(o.shape)) - """ - if 'transpose_state_layout' in kwargs: - if state_v_first: - raise ValueError("Cannot pass both `state_v_first` and the deprecated `transpose_state_layout`.") - warnings.warn( - "`transpose_state_layout` is deprecated and renamed to `state_v_first`.", - DeprecationWarning, - stacklevel=2, - ) - state_v_first = kwargs.pop('transpose_state_layout') - - if cu_seqlens is not None: - if q.shape[0] != 1: - raise ValueError( - f"The batch size is expected to be 1 rather than {q.shape[0]} when using `cu_seqlens`." - f"Please flatten variable-length inputs before processing.", - ) - if initial_state is not None and initial_state.shape[0] != len(cu_seqlens) - 1: - raise ValueError( - f"The number of initial states is expected to be equal to the number of input sequences, " - f"i.e., {len(cu_seqlens) - 1} rather than {initial_state.shape[0]}.", - ) - if scale is None: - scale = k.shape[-1] ** -0.5 - o, final_state = fused_recurrent( - q=q, - k=k, - v=v, - g=None, - gk=gk, - gv=gv, - scale=scale, - initial_state=initial_state, - output_final_state=output_final_state, - reverse=reverse, - cu_seqlens=cu_seqlens, - state_v_first=state_v_first, - ) - return o, final_state diff --git a/build/torch-xpu/ops/gla/naive.py b/build/torch-xpu/ops/gla/naive.py deleted file mode 100644 index 7608f5cdc3b13aff27a18c81ede79c1c7f3bec96..0000000000000000000000000000000000000000 --- a/build/torch-xpu/ops/gla/naive.py +++ /dev/null @@ -1,44 +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 - - -def ceildiv(a, b): - return -(a // -b) - - -def naive_recurrent_gla( - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - gk: torch.Tensor, - initial_state: torch.Tensor | None = None, - output_final_state: bool = False, -): - dtype = q.dtype - q, k, v, gk = map(lambda x: x.transpose(1, 2).float(), (q, k, v, gk)) - B, H, T, K, V = *q.shape, v.shape[-1] - o = torch.zeros_like(v) - scale = K ** -0.5 - - h = q.new_zeros(B, H, K, V, dtype=torch.float32) - if initial_state is not None: - h += initial_state.float() - - for i in range(T): - q_i = q[:, :, i] * scale - k_i = k[:, :, i] - v_i = v[:, :, i] - gk_i = gk[:, :, i].exp() - kv_i = k_i[..., None] * v_i[..., None, :] - h = h * gk_i[..., None] + kv_i - o[:, :, i] = (q_i[..., None] * h).sum(-2) - - if not output_final_state: - h = None - return o.transpose(1, 2).to(dtype), h diff --git a/build/torch-xpu/ops/gsa/__init__.py b/build/torch-xpu/ops/gsa/__init__.py deleted file mode 100644 index e1e1dabb5ffdfcd96672d5d3132f6a0a4f2e698f..0000000000000000000000000000000000000000 --- a/build/torch-xpu/ops/gsa/__init__.py +++ /dev/null @@ -1,14 +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 .chunk import chunk_gsa -from .fused_recurrent import fused_recurrent_gsa - -__all__ = [ - 'chunk_gsa', - 'fused_recurrent_gsa', -] diff --git a/build/torch-xpu/ops/gsa/chunk.py b/build/torch-xpu/ops/gsa/chunk.py deleted file mode 100644 index 21e68ac66504807274ce4e718e83cd0150a5775f..0000000000000000000000000000000000000000 --- a/build/torch-xpu/ops/gsa/chunk.py +++ /dev/null @@ -1,1218 +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 -import triton.language as tl -from einops import reduce - -from ...ops.common.chunk_h import chunk_bwd_dh, chunk_fwd_h -from ...ops.gla.chunk import chunk_gla_bwd, chunk_gla_fwd -from ...ops.utils import prepare_chunk_indices -from ...ops.utils.constant import RCP_LN2 -from ...ops.utils.cumsum import chunk_local_cumsum -from ...ops.utils.op import exp2 -from ...ops.utils.softmax import softmax_bwd, softmax_fwd -from ...utils import autotune_cache_kwargs, input_guard - - -@triton.heuristics({ - 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None, -}) -@triton.autotune( - configs=[ - triton.Config({'BK': BK, 'BV': BV}, num_warps=num_warps, num_stages=num_stages) - for BK in [32, 64] - for BV in [32, 64] - for num_warps in [2, 4, 8] - for num_stages in [2, 3, 4] - ], - key=['BT'], - **autotune_cache_kwargs, -) -@triton.jit(do_not_specialize=['T']) -def chunk_gsa_fwd_k_kernel_inter( - q, - k, - h, - g, - o, - A, - cu_seqlens, - chunk_indices, - scale, - T, - HQ: tl.constexpr, - H: tl.constexpr, - K: tl.constexpr, - V: tl.constexpr, - BT: tl.constexpr, - BK: tl.constexpr, - BV: tl.constexpr, - NG: tl.constexpr, - IS_VARLEN: tl.constexpr, -): - i_v, i_t, i_bh = tl.program_id(0), tl.program_id(1), tl.program_id(2) - i_b, i_hq = i_bh // HQ, i_bh % HQ - i_h = i_hq // NG - 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.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) - T = eos - bos - NT = tl.cdiv(T, BT) - else: - NT = tl.cdiv(T, BT) - i_tg = i_b * NT + i_t - bos, eos = i_b * T, i_b * T + T - - o_i = tl.arange(0, BT) - m_s = o_i[:, None] >= o_i[None, :] - - b_o = tl.zeros([BT, BV], dtype=tl.float32) - b_A = tl.zeros([BT, BT], dtype=tl.float32) - for i_k in range(tl.cdiv(K, BK)): - p_q = tl.make_block_ptr(q + (bos * HQ + i_hq) * K, (T, K), (HQ*K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) - p_k = tl.make_block_ptr(k + (bos * H + i_h) * K, (K, T), (1, H*K), (i_k * BK, i_t * BT), (BK, BT), (0, 1)) - p_h = tl.make_block_ptr(h + (i_tg * H + i_h) * K*V, (K, V), (V, 1), (i_k * BK, i_v * BV), (BK, BV), (1, 0)) - - # [BT, BK] - b_q = tl.load(p_q, boundary_check=(0, 1)) - b_q = (b_q * scale).to(b_q.dtype) - # [BK, BT] - b_k = tl.load(p_k, boundary_check=(0, 1)) - # [BK, BV] - b_h = tl.load(p_h, boundary_check=(0, 1)) - # [BT, BV] - b_o += tl.dot(b_q, b_h) - # [BT, BT] - b_A += tl.dot(b_q, b_k) - p_g = tl.make_block_ptr(g + (bos * H + i_h) * V, (T, V), (H*V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) - p_o = tl.make_block_ptr(o + (bos * HQ + i_hq) * V, (T, V), (HQ*V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) - p_A = tl.make_block_ptr(A + (bos * HQ + i_hq) * BT, (T, BT), (HQ*BT, 1), (i_t * BT, 0), (BT, BT), (1, 0)) - # [BT, BV] - b_g = tl.load(p_g, boundary_check=(0, 1)) - b_o = b_o * exp2(b_g) - tl.store(p_o, b_o.to(p_o.dtype.element_ty), boundary_check=(0, 1)) - - # [BT, BT] - b_A = tl.where(m_s, b_A, 0.) - if i_v == 0: - tl.store(p_A, b_A.to(p_A.dtype.element_ty), boundary_check=(0, 1)) - - -@triton.heuristics({ - 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None, -}) -@triton.jit(do_not_specialize=['T']) -def chunk_gsa_fwd_k_kernel_intra( - v, - g, - o, - A, - cu_seqlens, - chunk_indices, - T, - HQ: tl.constexpr, - H: tl.constexpr, - V: tl.constexpr, - BT: tl.constexpr, - BC: tl.constexpr, - BV: tl.constexpr, - NC: tl.constexpr, - NG: tl.constexpr, - IS_VARLEN: tl.constexpr, -): - i_v, i_c, i_bh = tl.program_id(0), tl.program_id(1), tl.program_id(2) - i_b, i_hq = i_bh // HQ, i_bh % HQ - i_h = i_hq // NG - i_t, i_i = i_c // NC, i_c % NC - 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.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) - T = eos - bos - else: - bos, eos = i_b * T, i_b * T + T - - o_i = tl.arange(0, BC) - o_v = i_v * BV + tl.arange(0, BV) - m_v = o_v < V - - if i_t * BT + i_i * BC >= T: - return - - p_g = tl.make_block_ptr(g + (bos * H + i_h) * V, (T, V), (H*V, 1), (i_t * BT + i_i * BC, i_v * BV), (BC, BV), (1, 0)) - p_gn = g + (bos + min(i_t * BT + i_i * BC, T)) * H*V + i_h * V + o_v - # [BV,] - b_gn = tl.load(p_gn, mask=m_v, other=0) - # [BC, BV] - b_o = tl.zeros([BC, BV], dtype=tl.float32) - for i_j in range(0, i_i): - p_A = tl.make_block_ptr(A + (bos*HQ+i_hq) * BT, (T, BT), (HQ*BT, 1), (i_t*BT+i_i*BC, i_j * BC), (BC, BC), (1, 0)) - p_v = tl.make_block_ptr(v + (bos*H+i_h) * V, (T, V), (H*V, 1), (i_t * BT + i_j * BC, i_v * BV), (BC, BV), (1, 0)) - p_gv = tl.make_block_ptr(g + (bos*H+i_h) * V, (T, V), (H*V, 1), (i_t * BT + i_j * BC, i_v * BV), (BC, BV), (1, 0)) - # [BC, BV] - b_v = tl.load(p_v, boundary_check=(0, 1)) - b_gv = tl.load(p_gv, boundary_check=(0, 1)) - b_vg = (b_v * exp2(b_gn[None, :] - b_gv)).to(b_v.dtype) - # [BC, BC] - b_A = tl.load(p_A, boundary_check=(0, 1)) - b_o += tl.dot(b_A, b_vg) - # [BC, BV] - b_g = tl.load(p_g, boundary_check=(0, 1)) - b_o *= exp2(b_g - b_gn[None, :]) - - o_A = (bos + i_t * BT + i_i * BC + tl.arange(0, BC)) * HQ*BT + i_hq * BT + i_i * BC - m_A = (i_t * BT + i_i * BC + tl.arange(0, BC)) < T - for j in range(0, min(BC, T - i_t * BT - i_i * BC)): - p_v = v + (bos + i_t * BT + i_i * BC + j) * H*V + i_h * V + o_v - p_gv = g + (bos + i_t * BT + i_i * BC + j) * H*V + i_h * V + o_v - # [BC,] - b_A = tl.load(A + o_A + j, mask=m_A, other=0) - # [BV,] - b_v = tl.load(p_v, mask=m_v, other=0).to(tl.float32) - b_gv = tl.load(p_gv, mask=m_v, other=0).to(tl.float32) - # [BC, BV] - b_vg = b_v[None, :] * exp2(b_g - b_gv[None, :]) - # avoid 0 * inf = inf - b_o += tl.where(o_i[:, None] >= j, b_A[:, None] * b_vg, 0.) - p_o = tl.make_block_ptr(o + (bos*HQ + i_hq) * V, (T, V), (HQ*V, 1), (i_t * BT + i_i * BC, i_v * BV), (BC, BV), (1, 0)) - b_o += tl.load(p_o, boundary_check=(0, 1)) - tl.store(p_o, b_o.to(p_o.dtype.element_ty), boundary_check=(0, 1)) - - -@triton.heuristics({ - 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None, -}) -@triton.autotune( - configs=[ - triton.Config({}, num_warps=num_warps) - for num_warps in [2, 4, 8] - ], - key=['BT'], - **autotune_cache_kwargs, -) -@triton.jit(do_not_specialize=['T']) -def chunk_gsa_bwd_k_kernel_dA( - v, - g, - do, - dA, - chunk_indices, - cu_seqlens, - scale, - T, - B: tl.constexpr, - HQ: tl.constexpr, - H: tl.constexpr, - V: tl.constexpr, - BT: tl.constexpr, - BC: tl.constexpr, - BV: tl.constexpr, - NC: tl.constexpr, - NG: tl.constexpr, - IS_VARLEN: tl.constexpr, -): - i_v, i_c, i_bh = tl.program_id(0), tl.program_id(1), tl.program_id(2) - i_b, i_hq = i_bh // HQ, i_bh % HQ - i_h = i_hq // NG - i_t, i_i, i_j = i_c // (NC * NC), (i_c % (NC * NC)) // NC, (i_c % (NC * NC)) % NC - 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.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) - all = T - T = eos - bos - else: - bos, eos = i_b * T, i_b * T + T - all = B * T - - o_v = i_v * BV + tl.arange(0, BV) - m_v = o_v < V - - if i_t * BT + i_i * BC >= T: - return - - p_dA = tl.make_block_ptr(dA+((i_v*all+bos)*HQ+i_hq)*BT, (T, BT), (HQ*BT, 1), (i_t*BT+i_i*BC, i_j*BC), (BC, BC), (1, 0)) - - # [BC, BC] - b_dA = tl.zeros([BC, BC], dtype=tl.float32) - if i_i > i_j: - p_v = tl.make_block_ptr(v + (bos*H+i_h) * V, (V, T), (1, H*V), (i_v * BV, i_t*BT + i_j*BC), (BV, BC), (0, 1)) - p_gv = tl.make_block_ptr(g + (bos*H+i_h) * V, (V, T), (1, H*V), (i_v * BV, i_t*BT + i_j*BC), (BV, BC), (0, 1)) - p_gn = g + (bos + i_t*BT + i_i*BC) * H*V + i_h * V + o_v - p_g = tl.make_block_ptr(g + (bos*H+i_h) * V, (T, V), (H*V, 1), (i_t*BT + i_i*BC, i_v*BV), (BC, BV), (1, 0)) - p_do = tl.make_block_ptr(do + (bos*HQ+i_hq) * V, (T, V), (HQ*V, 1), (i_t*BT + i_i*BC, i_v*BV), (BC, BV), (1, 0)) - # [BV,] - b_gn = tl.load(p_gn, mask=m_v, other=0.) - # [BC, BV] - b_g = tl.load(p_g, boundary_check=(0, 1)) - b_do = tl.load(p_do, boundary_check=(0, 1)) - b_do = (b_do * exp2(b_g - b_gn[None, :])).to(b_do.dtype) - # [BV, BC] - b_v = tl.load(p_v, boundary_check=(0, 1)) - b_gv = tl.load(p_gv, boundary_check=(0, 1)) - b_vg = (b_v * exp2(b_gn[:, None] - b_gv)).to(b_v.dtype) - # [BC, BC] - b_dA = tl.dot(b_do, b_vg) * scale - elif i_i == i_j: - p_g = tl.make_block_ptr(g + (bos*H + i_h) * V, (T, V), (H*V, 1), (i_t*BT + i_i*BC, i_v*BV), (BC, BV), (1, 0)) - p_do = tl.make_block_ptr(do + (bos*HQ + i_hq) * V, (T, V), (HQ*V, 1), (i_t*BT + i_i*BC, i_v*BV), (BC, BV), (1, 0)) - p_v = v + (bos + i_t*BT + i_j*BC) * H*V + i_h * V + o_v - p_gv = g + (bos + i_t*BT + i_j*BC) * H*V + i_h * V + o_v - # [BC, BV] - b_g = tl.load(p_g, boundary_check=(0, 1)) - b_do = tl.load(p_do, boundary_check=(0, 1)) * scale - m_v = o_v < V - - o_i = tl.arange(0, BC) - # [BC, BC] - m_dA = o_i[:, None] >= o_i[None, :] - for j in range(0, min(BC, T - i_t * BT - i_j * BC)): - # [BV,] - b_v = tl.load(p_v, mask=m_v, other=0).to(tl.float32) - b_gv = tl.load(p_gv, mask=m_v, other=0).to(tl.float32) - # [BC,] - b_dAj = tl.sum(b_do * b_v[None, :] * exp2(b_g - b_gv[None, :]), 1) - b_dA = tl.where((o_i == j)[None, :], b_dAj[:, None], b_dA) - - p_v += H*V - p_gv += H*V - b_dA = tl.where(m_dA, b_dA, 0.) - tl.store(p_dA, b_dA.to(dA.dtype.element_ty), boundary_check=(0, 1)) - - -@triton.heuristics({ - 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None, -}) -@triton.autotune( - configs=[ - triton.Config({}, num_warps=num_warps, num_stages=num_stages) - for num_warps in [2, 4] - for num_stages in [2, 3, 4] - ], - key=['BT'], - **autotune_cache_kwargs, -) -@triton.jit(do_not_specialize=['T']) -def chunk_gsa_bwd_k_kernel_dqkvg( - q, - k, - v, - h, - g, - A, - do, - dh, - dq, - dk, - dv, - dg, - dgv, - dA, - cu_seqlens, - chunk_indices, - scale, - T, - B: tl.constexpr, - HQ: tl.constexpr, - H: tl.constexpr, - K: tl.constexpr, - V: tl.constexpr, - BT: tl.constexpr, - BK: tl.constexpr, - BV: tl.constexpr, - NG: tl.constexpr, - IS_VARLEN: tl.constexpr, -): - i_k, i_t, i_bh = tl.program_id(0), tl.program_id(1), tl.program_id(2) - i_b, i_hq = i_bh // HQ, i_bh % HQ - i_h = i_hq // NG - 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.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) - all = T - T = eos - bos - NT = tl.cdiv(T, BT) - else: - NT = tl.cdiv(T, BT) - i_tg = i_b * NT + i_t - bos, eos = i_b * T, i_b * T + T - all = B * T - - o_i = tl.arange(0, BT) - o_t = min(i_t * BT + BT, T) - m_s = o_i[:, None] >= o_i[None, :] - - p_q = tl.make_block_ptr(q + (bos*HQ+i_hq) * K, (T, K), (HQ*K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) - p_k = tl.make_block_ptr(k + (bos*H+i_h) * K, (T, K), (H*K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) - p_A = tl.make_block_ptr(A + ((i_k*all+bos)*HQ+i_hq)*BT, (T, BT), (HQ*BT, 1), (i_t * BT, 0), (BT, BT), (1, 0)) - - # [BT, BK] - b_q = tl.load(p_q, boundary_check=(0, 1)) - b_k = tl.load(p_k, boundary_check=(0, 1)) - # [BT, BT] - b_A = tl.dot((b_q * scale).to(b_q.dtype), tl.trans(b_k)) - b_A = tl.where(m_s, b_A, 0.) - tl.store(p_A, b_A.to(p_A.dtype.element_ty), boundary_check=(0, 1)) - - b_dq = tl.zeros([BT, BK], dtype=tl.float32) - b_dk = tl.zeros([BT, BK], dtype=tl.float32) - for i_v in range(tl.cdiv(V, BV)): - o_v = i_v * BV + tl.arange(0, BV) - p_v = tl.make_block_ptr(v + (bos*H+i_h)*V, (T, V), (H*V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) - p_g = tl.make_block_ptr(g + (bos*H+i_h)*V, (T, V), (H*V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) - p_gn = g + (bos + o_t - 1) * H*V + i_h * V + o_v - p_do = tl.make_block_ptr(do + (bos*HQ+i_hq)*V, (T, V), (HQ*V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) - p_dv = tl.make_block_ptr(dv + ((i_k*all+bos)*HQ+i_hq)*V, (T, V), (HQ*V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) - p_dg = tl.make_block_ptr(dg + (bos*HQ+i_hq)*V, (T, V), (HQ*V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) - p_dgv = tl.make_block_ptr(dgv+((i_k*all+bos)*HQ+i_hq)*V, (T, V), (HQ*V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) - p_h = tl.make_block_ptr(h + (i_tg * H + i_h) * K*V, (V, K), (1, V), (i_v * BV, i_k * BK), (BV, BK), (0, 1)) - p_dh = tl.make_block_ptr(dh + (i_tg * HQ + i_hq) * K*V, (K, V), (V, 1), (i_k * BK, i_v * BV), (BK, BV), (1, 0)) - m_v = o_v < V - - # [BV,] - b_gn = tl.load(p_gn, mask=m_v, other=0) - # [BT, BV] - b_v = tl.load(p_v, boundary_check=(0, 1)) - b_g = tl.load(p_g, boundary_check=(0, 1)) - b_gv = exp2(b_gn[None, :] - b_g) - # [BV, BK] - b_h = tl.load(p_h, boundary_check=(0, 1)) - # [BT, BV] - b_do = tl.load(p_do, boundary_check=(0, 1)) - b_do = (b_do * exp2(b_g)).to(b_do.dtype) - # [BK, BV] - b_dh = tl.load(p_dh, boundary_check=(0, 1)) - # [BV] - b_dg = tl.sum(tl.trans(b_h) * b_dh, 0) * exp2(b_gn) - - b_dh = b_dh.to(b_k.dtype) - # [BT, BK] - b_dq += tl.dot(b_do, b_h.to(b_k.dtype)) * scale - b_dk += tl.dot((b_v * b_gv).to(b_v.dtype), tl.trans(b_dh)) - # [BT, BV] - b_dv = tl.dot(b_k, b_dh) * b_gv - # [BV] - b_dg += tl.sum(b_dv * b_v, 0) - - if i_k == 0: - b_dgv = tl.load(p_dg, boundary_check=(0, 1)) + b_dg[None, :] - else: - b_dgv = tl.zeros([BT, BV], dtype=tl.float32) + b_dg[None, :] - - tl.store(p_dgv, b_dgv.to(p_dgv.dtype.element_ty), boundary_check=(0, 1)) - tl.store(p_dv, b_dv.to(p_dv.dtype.element_ty), boundary_check=(0, 1)) - p_dA = tl.make_block_ptr(dA + (bos*HQ + i_hq) * BT, (T, BT), (HQ*BT, 1), (i_t * BT, 0), (BT, BT), (1, 0)) - p_dq = tl.make_block_ptr(dq + (bos*HQ + i_hq) * K, (T, K), (HQ*K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) - p_dk = tl.make_block_ptr(dk + (bos*HQ + i_hq) * K, (T, K), (HQ*K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) - # [BT, BT] - b_dA = tl.load(p_dA, boundary_check=(0, 1)) - # [BT, BK] - b_dq += tl.dot(b_dA, b_k) - b_dk += tl.dot(tl.trans(b_dA).to(b_k.dtype), b_q) - - tl.store(p_dq, b_dq.to(p_dq.dtype.element_ty), boundary_check=(0, 1)) - tl.store(p_dk, b_dk.to(p_dk.dtype.element_ty), boundary_check=(0, 1)) - - -@triton.heuristics({ - 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None, -}) -@triton.jit(do_not_specialize=['T']) -def chunk_gsa_bwd_k_kernel_intra_dvg( - v, - g, - o, - A, - do, - dv, - dg, - cu_seqlens, - chunk_indices, - T, - HQ: tl.constexpr, - H: tl.constexpr, - V: tl.constexpr, - BT: tl.constexpr, - BC: tl.constexpr, - BV: tl.constexpr, - NC: tl.constexpr, - NG: tl.constexpr, - IS_VARLEN: tl.constexpr, -): - i_v, i_c, i_bh = tl.program_id(0), tl.program_id(1), tl.program_id(2) - i_b, i_hq = i_bh // HQ, i_bh % HQ - i_h = i_hq // NG - i_t, i_i = i_c // NC, i_c % NC - 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.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) - T = eos - bos - else: - bos, eos = i_b * T, i_b * T + T - - o_i = tl.arange(0, BC) - o_v = i_v * BV + tl.arange(0, BV) - m_v = o_v < V - - if i_t * BT + i_i * BC >= T: - return - - p_gv = tl.make_block_ptr(g + (bos*H+i_h)*V, (T, V), (H*V, 1), (i_t * BT + i_i * BC, i_v * BV), (BC, BV), (1, 0)) - p_gn = g + (bos + min(i_t * BT + i_i * BC + BC, T)-1)*H*V + i_h*V + o_v - # [BV,] - b_gn = tl.load(p_gn, mask=m_v, other=0) - # [BC, BV] - b_gv = tl.load(p_gv, boundary_check=(0, 1)) - b_dv = tl.zeros([BC, BV], dtype=tl.float32) - for i_j in range(i_i + 1, NC): - p_g = tl.make_block_ptr(g + (bos*H+i_h) * V, (T, V), (H*V, 1), (i_t * BT + i_j * BC, i_v * BV), (BC, BV), (1, 0)) - p_A = tl.make_block_ptr(A + (bos*HQ+i_hq) * BT, (BT, T), (1, HQ*BT), (i_i*BC, i_t*BT + i_j*BC), (BC, BC), (0, 1)) - p_do = tl.make_block_ptr(do + (bos*HQ+i_hq) * V, (T, V), (HQ*V, 1), (i_t*BT + i_j*BC, i_v*BV), (BC, BV), (1, 0)) - - m_j = (i_t * BT + i_j * BC + o_i) < T - # [BC, BV] - b_g = tl.load(p_g, boundary_check=(0, 1)) - b_do = tl.load(p_do, boundary_check=(0, 1)) * tl.where(m_j[:, None], exp2(b_g - b_gn[None, :]), 0) - # [BC, BC] - b_A = tl.load(p_A, boundary_check=(0, 1)) - # [BC, BV] - b_dv += tl.dot(b_A, b_do.to(b_A.dtype)) - b_dv *= exp2(b_gn[None, :] - b_gv) - - p_g = g + (bos + i_t * BT + i_i * BC) * H*V + i_h * V + o_v - p_A = A + (bos + i_t*BT + i_i*BC) * HQ*BT + i_hq * BT + i_i * BC + o_i - p_do = do + (bos + i_t*BT + i_i*BC) * HQ*V + i_hq * V + o_v - for j in range(0, min(BC, T - i_t * BT - i_i * BC)): - # [BC,] - b_A = tl.load(p_A) - # [BV,] - b_g = tl.load(p_g, mask=m_v, other=0) - b_do = tl.load(p_do, mask=m_v, other=0) - # [BC, BV] - m_i = o_i[:, None] <= j - b_dv += tl.where(m_i, exp2(b_g[None, :] - b_gv) * b_A[:, None] * b_do[None, :], 0.) - - p_g += H * V - p_A += HQ * BT - p_do += HQ * V - p_o = tl.make_block_ptr(o + (bos*HQ+i_hq)*V, (T, V), (HQ*V, 1), (i_t*BT + i_i*BC, i_v*BV), (BC, BV), (1, 0)) - p_v = tl.make_block_ptr(v + (bos*H+i_h)*V, (T, V), (H*V, 1), (i_t*BT + i_i*BC, i_v*BV), (BC, BV), (1, 0)) - p_do = tl.make_block_ptr(do + (bos*HQ+i_hq)*V, (T, V), (HQ*V, 1), (i_t*BT + i_i*BC, i_v*BV), (BC, BV), (1, 0)) - p_dv = tl.make_block_ptr(dv + (bos*HQ+i_hq)*V, (T, V), (HQ*V, 1), (i_t*BT + i_i*BC, i_v*BV), (BC, BV), (1, 0)) - p_dg = tl.make_block_ptr(dg + (bos*HQ+i_hq)*V, (T, V), (HQ*V, 1), (i_t*BT + i_i*BC, i_v*BV), (BC, BV), (1, 0)) - - b_o = tl.load(p_o, boundary_check=(0, 1)).to(tl.float32) - b_v = tl.load(p_v, boundary_check=(0, 1)).to(tl.float32) - b_do = tl.load(p_do, boundary_check=(0, 1)).to(tl.float32) - b_dv = b_dv + tl.load(p_dv, boundary_check=(0, 1)).to(tl.float32) - b_dg = b_o * b_do - b_v * b_dv - tl.store(p_dv, b_dv.to(p_dv.dtype.element_ty), boundary_check=(0, 1)) - tl.store(p_dg, b_dg.to(p_dg.dtype.element_ty), boundary_check=(0, 1)) - - -def chunk_gsa_fwd_v( - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - g: torch.Tensor, - scale: float = 1., - initial_state: torch.Tensor | None = None, - output_final_state: bool = False, - cu_seqlens: torch.LongTensor | None = None, - chunk_size: int = 64, - chunk_indices: torch.LongTensor | None = None, -) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: - _, A, h, ht, o = chunk_gla_fwd( - q=q, - k=k, - v=v, - g=None, - g_cumsum=g, - scale=scale, - initial_state=initial_state, - output_final_state=output_final_state, - cu_seqlens=cu_seqlens, - chunk_size=chunk_size, - chunk_indices=chunk_indices, - ) - return A, h, ht, o - - -def chunk_gsa_fwd_k( - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - g: torch.Tensor, - h0: torch.Tensor | None = None, - output_final_state: bool = False, - scale: float = 1., - cu_seqlens: torch.LongTensor | None = None, - chunk_size: int = 64, - chunk_indices: torch.LongTensor | None = None, -) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: - B, T, H, K, V = *k.shape, v.shape[-1] - BT = chunk_size - BC = min(16, BT) - BV = min(64, triton.next_power_of_2(V)) - HQ = q.shape[2] - - if chunk_indices is None: - if cu_seqlens is not None: - chunk_indices = prepare_chunk_indices(cu_seqlens, BT) - else: - chunk_indices = None - NT = triton.cdiv(T, BT) if cu_seqlens is None else len(chunk_indices) - NC = triton.cdiv(BT, BC) - NG = HQ // H - - h, ht = chunk_fwd_h( - k=k, - v=v, - g=None, - gk=None, - gv=g, - h0=h0, - output_final_state=output_final_state, - cu_seqlens=cu_seqlens, - chunk_size=BT, - states_in_fp32=False, - ) - o = v.new_empty(B, T, HQ, V) - A = q.new_empty(B, T, HQ, BT) - def grid(meta): return (triton.cdiv(V, meta['BV']), NT, B * HQ) - chunk_gsa_fwd_k_kernel_inter[grid]( - q, - k, - h, - g, - o, - A, - cu_seqlens=cu_seqlens, - chunk_indices=chunk_indices, - scale=scale, - T=T, - HQ=HQ, - H=H, - K=K, - V=V, - BT=BT, - NG=NG, - ) - - def grid(meta): return (triton.cdiv(V, meta['BV']), NT * NC, B * HQ) - chunk_gsa_fwd_k_kernel_intra[grid]( - v, - g, - o, - A, - cu_seqlens=cu_seqlens, - chunk_indices=chunk_indices, - T=T, - HQ=HQ, - H=H, - V=V, - BT=BT, - BC=BC, - BV=BV, - NC=NC, - NG=NG, - num_warps=4, - num_stages=2, - ) - return A, h, ht, o - - -def chunk_gsa_bwd_v( - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - g: torch.Tensor, - h0: torch.Tensor, - h: torch.Tensor, - A: torch.Tensor, - do: torch.Tensor, - dht: torch.Tensor, - dg: torch.Tensor, - scale: float = 1., - cu_seqlens: torch.LongTensor | None = None, - chunk_size: int = 64, - chunk_indices: torch.LongTensor | None = None, -): - dq, dk, dv, dg, dh0 = chunk_gla_bwd( - q=q, - k=k, - v=v, - g=None, - g_cumsum=g, - scale=scale, - initial_state=h0, - h=h, - A=A, - do=do, - dht=dht, - cu_seqlens=cu_seqlens, - chunk_size=chunk_size, - chunk_indices=chunk_indices, - ) - return dq, dk, dv, dg, dh0 - - -def chunk_gsa_bwd_k( - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - g: torch.Tensor, - h: torch.Tensor, - h0: torch.Tensor, - o: torch.Tensor, - do: torch.Tensor, - dht: torch.Tensor, - dg: torch.Tensor, - scale: float = 1., - cu_seqlens: torch.LongTensor | None = None, - chunk_size: int = 64, - chunk_indices: torch.LongTensor | None = None, -): - B, T, H, K, V = *k.shape, v.shape[-1] - BT = chunk_size - BC = min(16, BT) - BK = min(64, triton.next_power_of_2(K)) - BV = min(64, triton.next_power_of_2(V)) - HQ = q.shape[2] - - if chunk_indices is None: - if cu_seqlens is not None: - chunk_indices = prepare_chunk_indices(cu_seqlens, BT) - else: - chunk_indices = None - NT = triton.cdiv(T, BT) if cu_seqlens is None else len(chunk_indices) - NC = triton.cdiv(BT, BC) - NK = triton.cdiv(K, BK) - NV = triton.cdiv(V, BV) - NG = HQ // H - - if h is None: - h, _ = chunk_fwd_h( - k=k, - v=v, - g=None, - gk=None, - gv=g, - h0=h0, - output_final_state=False, - cu_seqlens=cu_seqlens, - chunk_size=BT, - states_in_fp32=False, - ) - dh, dh0 = chunk_bwd_dh( - q=q, - k=k, - v=v, - g=None, - gk=None, - gv=g, - do=do, - h0=h0, - dht=dht, - scale=scale, - cu_seqlens=cu_seqlens, - chunk_size=BT, - states_in_fp32=True, - ) - dA = q.new_empty(NV, B, T, HQ, BT) - grid = (NV, NT * NC * NC, B * HQ) - chunk_gsa_bwd_k_kernel_dA[grid]( - v, - g, - do, - dA, - cu_seqlens=cu_seqlens, - chunk_indices=chunk_indices, - scale=scale, - T=T, - B=B, - HQ=HQ, - H=H, - V=V, - BT=BT, - BC=BC, - BV=BV, - NC=NC, - NG=NG, - ) - dA = dA.sum(0, dtype=dA.dtype) - - A = do.new_empty(NK, B, T, HQ, BT) - dq = torch.empty_like(q) - dk = k.new_empty(B, T, HQ, K) - dv = v.new_empty(NK, B, T, HQ, V) - dgv = g.new_empty(NK, B, T, HQ, V, dtype=torch.float) - grid = (NK, NT, B * HQ) - chunk_gsa_bwd_k_kernel_dqkvg[grid]( - q, - k, - v, - h, - g, - A, - do, - dh, - dq, - dk, - dv, - dg, - dgv, - dA, - cu_seqlens=cu_seqlens, - chunk_indices=chunk_indices, - scale=scale, - T=T, - B=B, - HQ=HQ, - H=H, - K=K, - V=V, - BT=BT, - BK=BK, - BV=BV, - NG=NG, - ) - A = A.sum(0, dtype=A.dtype) - dv = dv.sum(0, dtype=dv.dtype) - dgv = dgv.sum(0, dtype=dgv.dtype) - - def grid(meta): return (triton.cdiv(V, meta['BV']), NT * NC, B * HQ) - chunk_gsa_bwd_k_kernel_intra_dvg[grid]( - v, - g, - o, - A, - do, - dv, - dg, - cu_seqlens=cu_seqlens, - chunk_indices=chunk_indices, - T=T, - HQ=HQ, - H=H, - V=V, - BT=BT, - BC=BC, - BV=BV, - NC=NC, - NG=NG, - num_warps=4, - num_stages=2, - ) - dg = dgv.add_( - chunk_local_cumsum( - dg, - chunk_size=BT, - reverse=True, - cu_seqlens=cu_seqlens, - chunk_indices=chunk_indices, - ) - ) - - return dq, dk, dv, dg, dh0 - - -def chunk_gsa_fwd( - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - s: torch.Tensor, - g: torch.Tensor, - initial_state: tuple[torch.Tensor, torch.Tensor] | None = None, - output_final_state: bool = False, - scale: float = 1., - cu_seqlens: torch.LongTensor | None = None, - chunk_size: int = 64, - chunk_indices: torch.LongTensor | None = None, -) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: - hk0, hv0 = None, None - if initial_state is not None: - hk0, hv0 = initial_state - Ak, hk, hkt, ok = chunk_gsa_fwd_k( - q=q, - k=k, - v=s, - g=g, - h0=hk0, - output_final_state=output_final_state, - scale=scale, - cu_seqlens=cu_seqlens, - chunk_size=chunk_size, - chunk_indices=chunk_indices, - ) - - # p is kept in fp32 for safe softmax backward - p = softmax_fwd(ok, dtype=torch.float) - - qv = p.to(q.dtype) - Av, hv, hvt, ov = chunk_gsa_fwd_v( - q=qv, - k=s, - v=v, - g=g, - scale=1., - initial_state=hv0, - output_final_state=output_final_state, - cu_seqlens=cu_seqlens, - chunk_size=chunk_size, - chunk_indices=chunk_indices, - ) - return Ak, hk, hkt, ok, p, Av, hv, hvt, ov - - -def chunk_gsa_bwd( - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - s: torch.Tensor, - g: torch.Tensor, - ok: torch.Tensor, - p: torch.Tensor, - A: tuple[torch.Tensor, torch.Tensor], - h: tuple[torch.Tensor, torch.Tensor], - initial_state: tuple[torch.Tensor, torch.Tensor] | None, - scale: float, - do: torch.Tensor, - dht: tuple[torch.Tensor, torch.Tensor], - cu_seqlens: torch.LongTensor | None = None, - chunk_size: int = 64, - chunk_indices: torch.LongTensor | None = None, -): - hk0, hv0 = None, None - if initial_state is not None: - hk0, hv0 = initial_state - - _, Av = A - hk, hv = h - dhkt, dhvt = dht - - qv = p.to(q.dtype) - dqv, dsv, dv, dg, dhv0 = chunk_gsa_bwd_v( - q=qv, - k=s, - v=v, - g=g, - h0=hv0, - h=hv, - A=Av, - do=do, - dht=dhvt, - dg=None, - scale=1., - cu_seqlens=cu_seqlens, - chunk_size=chunk_size, - chunk_indices=chunk_indices, - ) - - # softmax gradient, equivalent to: - # dok = qv * (dqv - (qv * dqv).sum(-1, True)) - dok = softmax_bwd(p, dqv, dtype=ok.dtype) - - dq, dk, dsk, dg, dhk0 = chunk_gsa_bwd_k( - q=q, - k=k, - v=s, - g=g, - h0=hk0, - h=hk, - o=ok, - do=dok, - dht=dhkt, - dg=dg, - scale=scale, - cu_seqlens=cu_seqlens, - chunk_size=chunk_size, - chunk_indices=chunk_indices, - ) - - ds = dsv.add_(dsk) - if q.shape[1] != k.shape[1]: - dk, dv, ds, dg = map(lambda x: reduce(x, 'b (h g) ... -> b h ...', 'sum', h=k.shape[1]), (dk, dv, ds, dg)) - dg = dg.to(s.dtype) - return dq, dk, dv, ds, dg, dhk0, dhv0 - - -class ChunkGSAFunction(torch.autograd.Function): - - @staticmethod - @input_guard - def forward( - ctx, - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - s: torch.Tensor, - g: torch.Tensor, - scale: float, - hk0: torch.Tensor | None, - hv0: torch.Tensor | None, - output_final_state: bool, - checkpoint_level: int, - cu_seqlens: torch.LongTensor | None, - cu_seqlens_cpu: torch.LongTensor | None, - chunk_size: int | None = None, - ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: - if chunk_size is None: - chunk_size = 64 - - if cu_seqlens is not None: - chunk_indices = prepare_chunk_indices( - cu_seqlens, - chunk_size, - cu_seqlens_cpu=cu_seqlens_cpu, - ) - else: - chunk_indices = None - - # Pre-scale by RCP_LN2 so downstream kernels can use exp2 directly. - g_org = g - g = chunk_local_cumsum( - g, - chunk_size, - scale=RCP_LN2, - cu_seqlens=cu_seqlens, - chunk_indices=chunk_indices, - ) - Ak, hk, hkt, ok, p, Av, hv, hvt, ov = chunk_gsa_fwd( - q=q, - k=k, - v=v, - s=s, - g=g, - initial_state=(hk0, hv0), - output_final_state=output_final_state, - scale=scale, - cu_seqlens=cu_seqlens, - chunk_size=chunk_size, - chunk_indices=chunk_indices, - ) - - if checkpoint_level >= 1: - del g - g = g_org - if checkpoint_level > 1: - del hk - del hv - hk, hv = None, None - else: - hk0, hv0 = None, None - - ctx.save_for_backward( - q, - k, - v, - s, - g, - ok, - p, - Av, - hk0, - hv0, - hk, - hv, - chunk_indices, - ) - ctx.checkpoint_level = checkpoint_level - ctx.scale = scale - ctx.cu_seqlens = cu_seqlens - ctx.chunk_size = chunk_size - return ov, hkt, hvt - - @staticmethod - @input_guard - def backward(ctx, dov, dhkt=None, dhvt=None): - ( - q, - k, - v, - s, - g, - ok, - p, - Av, - hk0, - hv0, - hk, - hv, - chunk_indices, - ) = ctx.saved_tensors - scale = ctx.scale - cu_seqlens = ctx.cu_seqlens - chunk_size = ctx.chunk_size - - if ctx.checkpoint_level >= 1: - g = chunk_local_cumsum( - g, - chunk_size, - scale=RCP_LN2, - cu_seqlens=cu_seqlens, - chunk_indices=chunk_indices, - ) - dq, dk, dv, ds, dg, dhk0, dhv0 = chunk_gsa_bwd( - q=q, - k=k, - v=v, - s=s, - g=g, - ok=ok, - p=p, - A=(None, Av), - h=(hk, hv), - initial_state=(hk0, hv0), - scale=scale, - do=dov, - dht=(dhkt, dhvt), - cu_seqlens=cu_seqlens, - chunk_size=chunk_size, - chunk_indices=chunk_indices, - ) - return dq, dk, dv, ds, dg, None, dhk0, dhv0, None, None, None, None, None - - -@torch.compiler.disable -def chunk_gsa( - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - s: torch.Tensor, - g: torch.Tensor | None = None, - scale: float | None = None, - initial_state: tuple[torch.Tensor] | None = None, - output_final_state: bool | None = False, - checkpoint_level: int | None = 2, - cu_seqlens: torch.LongTensor | None = None, - cu_seqlens_cpu: torch.LongTensor | None = None, - **kwargs, -) -> tuple[torch.Tensor, torch.Tensor]: - r""" - Args: - q (torch.Tensor): - queries of shape `[B, T, HQ, K]`. - k (torch.Tensor): - keys of shape `[B, T, H, K]`. - GQA is performed if `H` is not equal to `HQ`. - v (torch.Tensor): - values of shape `[B, T, H, V]`. - s (torch.Tensor): - slot representations of shape `[B, T, H, M]`. - g (Optional[torch.Tensor]): - Forget gates of shape `[B, T, H, M]` applied to keys. - If not provided, this function is equivalent to vanilla ABC. Default: `None`. - scale (Optional[float]): - Scale factor for attention scores. - If not provided, it will default to `1 / sqrt(K)`. Default: `None`. - initial_state (Optional[Tuple[torch.Tensor]]): - Initial state tuple having tensors of shape `[N, H, K, M]` and `[N, H, M, V]` for `N` input sequences. - For equal-length input sequences, `N` equals the batch size `B`. - Default: `None`. - output_final_state (Optional[bool]): - Whether to output the final state tuple, having tensors of shape `[N, H, K, M]` and `[N, H, M, V]`. - Default: `False`. - checkpoint_level (Optional[int]): - Checkpointing level; higher values will save more memory and do more recomputations during backward. - Default: `2`. - - Level `0`: no memory saved, no recomputation. - - Level `1`: recompute the fp32 cumulative values during backward. - - Level `2`: recompute the fp32 cumulative values and forward hidden states during backward. - cu_seqlens (torch.LongTensor): - Cumulative sequence lengths of shape `[N+1]` used for variable-length training, - consistent with the FlashAttention API. - cu_seqlens_cpu (torch.LongTensor): - CPU copy of `cu_seqlens` to avoid unnecessary device synchronization. Default: `None`. - - Returns: - o (torch.Tensor): - Outputs of shape `[B, T, H, V]`. - final_state (Tuple[torch.Tensor]): - Final state tuple having tensors of shape `[N, H, K, M]` and `[N, H, M, V]` if `output_final_state=True`. - `None` otherwise. - - Examples:: - >>> import torch - >>> import torch.nn.functional as F - >>> from einops import rearrange - >>> from ...ops.gsa import fused_recurrent_gsa - # inputs with equal lengths - >>> B, T, H, K, V, M = 4, 2048, 4, 512, 512, 64 - >>> q = torch.randn(B, T, H, K, device='cuda') - >>> k = torch.randn(B, T, H, K, device='cuda') - >>> v = torch.randn(B, T, H, V, device='cuda') - >>> s = torch.randn(B, T, H, M, device='cuda') - >>> g = F.logsigmoid(torch.randn(B, T, H, M, device='cuda')) - >>> h0 = (torch.randn(B, H, K, M, device='cuda'), torch.randn(B, H, M, V, device='cuda')) - >>> o, (hk, hv) = chunk_gsa( - q, k, v, s, g, - initial_state=h0, - output_final_state=True - ) - # for variable-length inputs, the batch size `B` is expected to be 1 and `cu_seqlens` is required - >>> q, k, v, s, g = map(lambda x: rearrange(x, 'b t h d -> 1 (b t) h d'), (q, k, v, s, g)) - # for a batch with 4 sequences, `cu_seqlens` with 5 start/end positions are expected - >>> cu_seqlens = q.new_tensor([0, 2048, 4096, 6144, 8192], dtype=torch.long) - >>> o_var, (hk_var, hv_var) = chunk_gsa( - q, k, v, s, g, - initial_state=h0, - output_final_state=True, - cu_seqlens=cu_seqlens - ) - >>> assert o.allclose(o_var.view(o.shape)) - >>> assert hk.allclose(hk_var) - >>> assert hv.allclose(hv_var) - """ - if 'head_first' in kwargs: - raise DeprecationWarning( - "head_first has been removed. Inputs must be in `[B, T, H, ...]` format.", - ) - chunk_size = kwargs.pop('chunk_size', None) - if chunk_size is not None and chunk_size != 2 ** (chunk_size.bit_length() - 1): - raise ValueError(f"`chunk_size` must be a power of 2, got {chunk_size}.") - if cu_seqlens is not None: - if q.shape[0] != 1: - raise ValueError( - f"The batch size is expected to be 1 rather than {q.shape[0]} when using `cu_seqlens`. " - f"Please flatten variable-length inputs before processing.", - ) - if initial_state is not None and initial_state[0].shape[0] != len(cu_seqlens) - 1: - raise ValueError( - f"The number of initial states is expected to be equal to the number of input sequences, " - f"i.e., {len(cu_seqlens) - 1} rather than {initial_state[0].shape[0]}.", - ) - assert checkpoint_level in [0, 1, 2] - if g is None: - # TODO: this 3 steps took huge amount of time, ought to be optimized - z = s.float().logcumsumexp(2) - g = torch.cat((z[:, :, :1], z[:, :, :-1]), 1) - z - s = torch.exp(s - z).to(k.dtype) - if scale is None: - scale = q.shape[-1] ** -0.5 - - hk0, hv0 = None, None - if initial_state is not None: - hk0, hv0 = initial_state - o, *final_state = ChunkGSAFunction.apply( - q, - k, - v, - s, - g, - scale, - hk0, - hv0, - output_final_state, - checkpoint_level, - cu_seqlens, - cu_seqlens_cpu, - chunk_size, - ) - return o, final_state diff --git a/build/torch-xpu/ops/gsa/fused_recurrent.py b/build/torch-xpu/ops/gsa/fused_recurrent.py deleted file mode 100644 index 686b1836c9224cc6f97784cfe658d2de812c95cf..0000000000000000000000000000000000000000 --- a/build/torch-xpu/ops/gsa/fused_recurrent.py +++ /dev/null @@ -1,538 +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 -import triton.language as tl - -from ...ops.common.fused_recurrent import fused_recurrent_bwd_kernel, fused_recurrent_fwd_kernel -from ...ops.utils.op import exp -from ...utils import autocast_custom_bwd, autocast_custom_fwd, input_guard - - -@triton.jit -def fused_recurrent_gsa_inference_kernel( - q, - k, - v, - s, - g, - o, - hk0, - hv0, - hkt, - hvt, - scale, - K: tl.constexpr, - V: tl.constexpr, - M: tl.constexpr, - BK: tl.constexpr, - BV: tl.constexpr, - NG: tl.constexpr, -): - i_bh = tl.program_id(0) - i_bg = i_bh // NG - - b_s = tl.load(s + i_bg * M + tl.arange(0, M)).to(tl.float32) - b_g = tl.load(g + i_bg * M + tl.arange(0, M)).to(tl.float32) - b_g = exp(b_g) - - b_ok = tl.zeros([M], dtype=tl.float32) - for i_k in range(tl.cdiv(K, BK)): - o_k = i_k * BK + tl.arange(0, BK) - - p_hk0 = hk0 + i_bg * K * M + (o_k[None, :]) * M + tl.arange(0, M)[:, None] - # [BK,] - mask_k = o_k < K - # [M, BK] - mask_hk = (tl.arange(0, M) < M)[:, None] & mask_k[None, :] - # [M, BK] - b_hk = tl.load(p_hk0, mask=mask_hk, other=0.).to(tl.float32) - # [BK,] - b_q = tl.load(q + i_bh * K + o_k, mask=mask_k, other=0.).to(tl.float32) * scale - b_k = tl.load(k + i_bg * K + o_k, mask=mask_k, other=0.).to(tl.float32) - b_hk = b_hk * b_g[:, None] + b_k[None, :] * b_s[:, None] - b_ok += tl.sum(b_hk * b_q[None, :], axis=1) - - if i_bh % NG == 0: - p_hkt = hkt + i_bg * K * M + o_k[None, :] * M + tl.arange(0, M)[:, None] - tl.store(p_hkt, b_hk.to(p_hkt.dtype.element_ty), mask=mask_hk) - - b_qv = tl.softmax(b_ok) - for i_v in range(tl.cdiv(V, BV)): - o_v = i_v * BV + tl.arange(0, BV) - - p_hv0 = hv0 + i_bg * M * V + tl.arange(0, M)[None, :] * V + o_v[:, None] - # [BV,] - mask_v = o_v < V - # [BV, M] - mask_hv = mask_v[:, None] & (tl.arange(0, M) < M)[None, :] - # [BV, M] - b_hv = tl.load(p_hv0, mask=mask_hv, other=0).to(tl.float32) - # [BV,] - b_v = tl.load(v + i_bg * V + o_v, mask=mask_v, other=0).to(tl.float32) - b_hv = b_hv * b_g[None, :] + b_s[None, :] * b_v[:, None] - b_ov = tl.sum(b_hv * b_qv[None, :], axis=1) - - tl.store(o + i_bh * V + o_v, b_ov.to(o.dtype.element_ty), mask=mask_v) - - if i_bh % NG == 0: - p_hvt = hvt + i_bg * M * V + tl.arange(0, M)[None, :] * V + o_v[:, None] - tl.store(p_hvt, b_hv.to(p_hvt.dtype.element_ty), mask=mask_hv) - - -def fused_recurrent_gsa_inference( - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - s: torch.Tensor, - g: torch.Tensor, - initial_state: tuple[torch.Tensor, torch.Tensor] | None = None, - output_final_state: bool = False, - scale: float = 1., -) -> torch.Tensor: - B, T, H, K, V, M = *k.shape, v.shape[-1], s.shape[-1] - HQ = q.shape[2] - BK, BV = min(triton.next_power_of_2(K), 64), min(triton.next_power_of_2(V), 64) - NG = HQ // H - - if initial_state != (None, None) and initial_state is not None: - hk0, hv0 = initial_state - else: - hk0, hv0 = q.new_zeros(B, H, K, M, dtype=torch.float), q.new_zeros(B, H, M, V, dtype=torch.float) - - hkt, hvt = None, None - if output_final_state: - if NG == 1: - hkt, hvt = hk0, hv0 - else: - hkt, hvt = q.new_empty(B, H, K, M, dtype=torch.float), q.new_empty(B, H, M, V, dtype=torch.float) - - o = v.new_empty(B, T, HQ, V) - grid = (B * HQ,) - fused_recurrent_gsa_inference_kernel[grid]( - q, - k, - v, - s, - g, - o, - hk0, - hv0, - hkt, - hvt, - scale=scale, - K=K, - V=V, - M=M, - BK=BK, - BV=BV, - NG=NG, - ) - return o, (hkt, hvt) - - -def fused_recurrent_gsa_fwd( - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - s: torch.Tensor, - g: torch.Tensor, - initial_state: tuple[torch.Tensor, torch.Tensor] | None = None, - output_final_state: bool = False, - scale: float = 1., - reverse: bool = False, - cu_seqlens: torch.LongTensor | None = None, -) -> tuple[torch.Tensor, tuple[torch.Tensor]]: - B, T, H, K, V, M = *k.shape, v.shape[-1], s.shape[-1] - N = B if cu_seqlens is None else len(cu_seqlens) - 1 - HQ = q.shape[2] - if HQ != H: - raise ValueError("GQA not supported yet.") - - BK, BV, BM = min(triton.next_power_of_2(K), 64), min(triton.next_power_of_2(V), 64), min(triton.next_power_of_2(M), 64) - NK, NV, NM = triton.cdiv(K, BK), triton.cdiv(V, BV), triton.cdiv(M, BM) - - hk0, hv0 = None, None - if initial_state != (None, None) and initial_state is not None: - hk0, hv0 = initial_state - hkt, hvt = None, None - if output_final_state: - hkt, hvt = q.new_empty(N, H, K, M, dtype=torch.float), q.new_empty(N, H, M, V, dtype=torch.float) - - ok = q.new_empty(NK, *s.shape, dtype=torch.float) - gk, gv = None, g - grid = (NM, NK, N * H) - fused_recurrent_fwd_kernel[grid]( - q=q, - k=k, - v=s, - g=None, - g_gamma=None, - gk=gk, - gv=gv, - o=ok, - h0=hk0, - ht=hkt, - cu_seqlens=cu_seqlens, - scale=scale, - B=B, - T=T, - H=H, - K=K, - V=M, - BK=BK, - BV=BM, - USE_G=False, - USE_G_GAMMA=False, - USE_GK=False, - USE_GV=True, - REVERSE=reverse, - ) - ok = ok.sum(0) - - qv = ok.softmax(-1, dtype=torch.float) - ov = q.new_empty(NM, *v.shape, dtype=torch.float) - gk, gv = g, None - grid = (NV, NM, N * H) - fused_recurrent_fwd_kernel[grid]( - q=qv, - k=s, - v=v, - g=None, - g_gamma=None, - gk=gk, - gv=gv, - o=ov, - h0=hv0, - ht=hvt, - cu_seqlens=cu_seqlens, - scale=1., - B=B, - T=T, - H=H, - K=M, - V=V, - BK=BM, - BV=BV, - USE_G=False, - USE_G_GAMMA=False, - USE_GK=True, - USE_GV=False, - REVERSE=reverse, - ) - ov = ov.sum(0) - return ok, hkt, qv, ov, hvt - - -def fused_recurrent_gsa_bwd( - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - s: torch.Tensor, - g: torch.Tensor, - qv: torch.Tensor, - hk0: torch.Tensor | None = None, - hv0: torch.Tensor | None = None, - ok: torch.Tensor | None = None, - do: torch.Tensor | None = None, - dhkt: torch.Tensor | None = None, - dhvt: torch.Tensor | None = None, - scale: float = 1., - reverse: bool = False, - cu_seqlens: torch.LongTensor | None = None, -) -> tuple[torch.Tensor]: - B, T, H, K, V, M = *q.shape, v.shape[-1], s.shape[-1] - N = B if cu_seqlens is None else len(cu_seqlens) - 1 - - BK, BV, BM = min(triton.next_power_of_2(K), 64), min(triton.next_power_of_2(V), 64), min(triton.next_power_of_2(M), 64) - NK, NV, NM = triton.cdiv(K, BK), triton.cdiv(V, BV), triton.cdiv(M, BM) - - dqv = q.new_empty(NV, B, T, H, M, dtype=torch.float) - dsv = q.new_empty(NV, B, T, H, M, dtype=torch.float) - dv = q.new_empty(NM, B, T, H, V, dtype=torch.float) - dgv = q.new_empty(NV, B, T, H, M, dtype=torch.float) - dhv0 = torch.empty_like(hv0)if hv0 is not None else None - - grid = (NV, NM, N * H) - fused_recurrent_bwd_kernel[grid]( - q=qv, - k=s, - v=v, - g=None, - g_gamma=None, - gk=g, - gv=None, - o=None, - h0=hv0, - do=do, - dq=dqv, - dk=dsv, - dv=dv, - dg=None, - dgk=dgv, - dgv=None, - dht=dhvt, - dh0=dhv0, - cu_seqlens=cu_seqlens, - scale=1., - B=B, - T=T, - H=H, - K=M, - V=V, - BK=BM, - BV=BV, - USE_G=False, - USE_G_GAMMA=False, - USE_GK=True, - USE_GV=False, - REVERSE=reverse, - ) - dqv = dqv.sum(0) - dsv = dsv.sum(0) - dv = dv.sum(0) - dgv = dgv.sum(0) - - dok = qv * (dqv - (qv * dqv).sum(-1, True)) - dq = q.new_empty(NM, B, T, H, K, dtype=torch.float) - dk = q.new_empty(NM, B, T, H, K, dtype=torch.float) - dsk = q.new_empty(NK, B, T, H, M, dtype=torch.float) - dgk = q.new_empty(NK, B, T, H, M, dtype=torch.float) - dhk0 = torch.empty_like(hk0)if hk0 is not None else None - - grid = (NM, NK, N * H) - fused_recurrent_bwd_kernel[grid]( - q=q, - k=k, - v=s, - g=None, - g_gamma=None, - gk=None, - gv=g, - o=ok, - h0=hk0, - do=dok, - dq=dq, - dk=dk, - dv=dsk, - dg=None, - dgk=None, - dgv=dgk, - dht=dhkt, - dh0=dhk0, - cu_seqlens=cu_seqlens, - scale=scale, - B=B, - T=T, - H=H, - K=K, - V=M, - BK=BK, - BV=BM, - USE_G=False, - USE_G_GAMMA=False, - USE_GK=False, - USE_GV=True, - REVERSE=reverse, - ) - dq = dq.sum(0) - dk = dk.sum(0) - dsk = dsk.sum(0) - dgk = dgk.sum(0) - - ds = dsk.add_(dsv) - dg = dgk.add_(dgv) - - return dq, dk, dv, ds, dg, dhk0, dhv0 - - -class FusedRecurrentGSAFunction(torch.autograd.Function): - - @staticmethod - @input_guard - @autocast_custom_fwd - def forward( - ctx, - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - s: torch.Tensor, - g: torch.Tensor, - scale: float | None = None, - hk0: torch.Tensor | None = None, - hv0: torch.Tensor | None = None, - output_final_state: bool = False, - reverse: bool = False, - cu_seqlens: torch.LongTensor | None = None, - ) -> tuple[torch.Tensor, tuple[torch.Tensor]]: - T = q.shape[1] - if T == 1 and not q.requires_grad: - o, (hkt, hvt) = fused_recurrent_gsa_inference( - q=q, - k=k, - v=v, - s=s, - g=g, - initial_state=(hk0, hv0), - output_final_state=output_final_state, - scale=scale, - ) - return o, hkt, hvt - ok, hkt, qv, ov, hvt = fused_recurrent_gsa_fwd( - q=q, - k=k, - v=v, - s=s, - g=g, - initial_state=(hk0, hv0), - output_final_state=output_final_state, - scale=scale, - reverse=reverse, - cu_seqlens=cu_seqlens, - ) - ctx.save_for_backward(q, k, v, s, g, qv, hk0, hv0, ok) - ctx.scale = scale - ctx.reverse = reverse - ctx.cu_seqlens = cu_seqlens - return ov.to(q.dtype), hkt, hvt - - @staticmethod - @input_guard - @autocast_custom_bwd - def backward(ctx, do, dhkt=None, dhvt=None): - q, k, v, s, g, qv, hk0, hv0, ok = ctx.saved_tensors - scale = ctx.scale - reverse = ctx.reverse - cu_seqlens = ctx.cu_seqlens - - dq, dk, dv, ds, dg, dhk0, dhv0 = fused_recurrent_gsa_bwd( - q=q, - k=k, - v=v, - s=s, - g=g, - qv=qv, - hk0=hk0, - hv0=hv0, - ok=ok, - do=do, - dhkt=dhkt, - dhvt=dhvt, - scale=scale, - reverse=reverse, - cu_seqlens=cu_seqlens, - ) - return dq.to(q), dk.to(k), dv.to(v), ds.to(s), dg.to(g), None, dhk0, dhv0, None, None, None - - -def fused_recurrent_gsa( - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - s: torch.Tensor, - g: torch.Tensor | None = None, - scale: int | None = None, - initial_state: tuple[torch.Tensor] | None = None, - output_final_state: bool | None = False, - reverse: bool | None = False, - cu_seqlens: torch.LongTensor | None = None, -) -> tuple[torch.Tensor, torch.Tensor]: - r""" - Args: - q (torch.Tensor): - queries of shape `[B, T, H, K]`. - k (torch.Tensor): - keys of shape `[B, T, H, K]`. - v (torch.Tensor): - values of shape `[B, T, H, V]`. - s (torch.Tensor): - slot representations of shape `[B, T, H, M]`. - g (torch.Tensor): - Forget gates of shape `[B, H, T, M]` applied to keys. - scale (Optional[float]): - Scale factor for the attention scores. - If not provided, it will default to `1 / sqrt(K)`. Default: `None`. - initial_state (Optional[Tuple[torch.Tensor]]): - Initial state tuple having tensors of shape `[N, H, K, M]` and `[N, H, M, V]` for `N` input sequences. - For equal-length input sequences, `N` equals the batch size `B`. - Default: `None`. - output_final_state (Optional[bool]): - Whether to output the final state of shape `[N, H, K, V]` and `[N, H, M, V]`. - Default: `False`. - reverse (Optional[bool]): - If `True`, process the state passing in reverse order. Default: `False`. - cu_seqlens (torch.LongTensor): - Cumulative sequence lengths of shape `[N+1]` used for variable-length training, - consistent with the FlashAttention API. - - Returns: - o (torch.Tensor): - Outputs of shape `[B, T, H, V]`. - final_state (Tuple[torch.Tensor]): - Final state tuple having tensors of shape `[N, H, K, M]` and `[N, H, M, V]`. - - Examples:: - >>> import torch - >>> import torch.nn.functional as F - >>> from einops import rearrange - >>> from ...ops.gsa import fused_recurrent_gsa - # inputs with equal lengths - >>> B, T, H, K, V, M = 4, 2048, 4, 512, 512, 64 - >>> q = torch.randn(B, T, H, K, device='cuda') - >>> k = torch.randn(B, T, H, K, device='cuda') - >>> v = torch.randn(B, T, H, V, device='cuda') - >>> s = torch.randn(B, T, H, M, device='cuda') - >>> g = F.logsigmoid(torch.randn(B, T, H, M, device='cuda')) - >>> h0 = (torch.randn(B, H, K, M, device='cuda'), torch.randn(B, H, M, V, device='cuda')) - >>> o, (hk, hv) = fused_recurrent_gsa( - q, k, v, s, g, - initial_state=h0, - output_final_state=True - ) - # for variable-length inputs, the batch size `B` is expected to be 1 and `cu_seqlens` is required - >>> q, k, v, s, g = map(lambda x: rearrange(x, 'b t h d -> 1 (b t) h d'), (q, k, v, s, g)) - # for a batch with 4 sequences, `cu_seqlens` with 5 start/end positions are expected - >>> cu_seqlens = q.new_tensor([0, 2048, 4096, 6144, 8192], dtype=torch.long) - >>> o_var, (hk_var, hv_var) = fused_recurrent_gsa( - q, k, v, s, g, - initial_state=h0, - output_final_state=True, - cu_seqlens=cu_seqlens - ) - >>> assert o.allclose(o_var.view(o.shape)) - >>> assert hk.allclose(hk_var) - >>> assert hv.allclose(hv_var) - """ - if cu_seqlens is not None: - if q.shape[0] != 1: - raise ValueError( - f"The batch size is expected to be 1 rather than {q.shape[0]} when using `cu_seqlens`." - f"Please flatten variable-length inputs before processing.", - ) - if initial_state is not None and initial_state[0].shape[0] != len(cu_seqlens) - 1: - raise ValueError( - f"The number of initial states is expected to be equal to the number of input sequences, " - f"i.e., {len(cu_seqlens) - 1} rather than {initial_state[0].shape[0]}.", - ) - if scale is None: - scale = k.shape[-1] ** -0.5 - if initial_state is None: - initial_state = (None, None) - o, *final_state = FusedRecurrentGSAFunction.apply( - q, - k, - v, - s, - g, - scale, - *initial_state, - output_final_state, - reverse, - cu_seqlens, - ) - return o, final_state diff --git a/build/torch-xpu/ops/gsa/naive.py b/build/torch-xpu/ops/gsa/naive.py deleted file mode 100644 index 73173bf394349092bce3833aa0d9dce022d0f693..0000000000000000000000000000000000000000 --- a/build/torch-xpu/ops/gsa/naive.py +++ /dev/null @@ -1,72 +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 -from einops import repeat - - -def naive_recurrent_gsa( - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - s: torch.Tensor, - g: torch.Tensor | None = None, - scale: int | None = None, - initial_state: torch.Tensor | None = None, - output_final_state: bool | None = False, -) -> torch.Tensor: - dtype = q.dtype - q, k, v, s, g = map(lambda x: x.transpose(1, 2).contiguous().float(), (q, k, v, s, g)) - - NG = q.shape[1]//k.shape[1] - # [batch_size, n_heads, seq_len, n_slots] - if g is None: - z = s.float().logcumsumexp(2) - g = torch.cat((z[:, :, :1], z[:, :, :-1]), 2) - z - s = torch.exp(s - z) - k, v, s, g = map(lambda x: repeat(x, 'b h t d -> b (h g) t d', g=NG), (k, v, s, g)) - if initial_state is not None: - initial_state = tuple(map(lambda x: repeat(x, 'b h k v -> b (h g) k v', g=NG), initial_state)) - - B, H, T, K, V, M = *q.shape, v.shape[-1], s.shape[-1] - - hk = torch.zeros(B, H, K, M, dtype=torch.float, device=q.device) - ok = torch.zeros_like(s) - - if scale is None: - scale = q.shape[-1] ** -0.5 - - final_state = None - if initial_state is not None: - hk += initial_state[0] - - for i in range(T): - q_i = q[:, :, i] * scale - k_i = k[:, :, i] - v_i = s[:, :, i] - g_i = g[:, :, i].exp() - hk = hk * g_i[..., None, :] + k_i[..., None] * v_i[..., None, :] - ok[:, :, i] = (q_i[..., None] * hk).sum(-2) - - qv = ok.softmax(-1) - hv = torch.zeros(B, H, M, V, dtype=torch.float, device=q.device) - ov = torch.zeros_like(v) - if initial_state is not None: - hv += initial_state[1] - - for i in range(T): - q_i = qv[:, :, i] - k_i = s[:, :, i] - v_i = v[:, :, i] - g_i = g[:, :, i].exp() - hv = hv * g_i[..., :, None] + k_i[..., None] * v_i[..., None, :] - ov[:, :, i] = (q_i[..., None] * hv).sum(-2) - - if output_final_state: - final_state = (hk.view(B, -1, NG, K, M)[:, :, 0], hv.view(B, -1, NG, M, V)[:, :, 0]) - ov = ov.transpose(1, 2).contiguous() - return ov.to(dtype), final_state diff --git a/build/torch-xpu/ops/hgrn/__init__.py b/build/torch-xpu/ops/hgrn/__init__.py deleted file mode 100644 index a757f6144c213b0587a48bc79cb004985ac5e021..0000000000000000000000000000000000000000 --- a/build/torch-xpu/ops/hgrn/__init__.py +++ /dev/null @@ -1,14 +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 .chunk import chunk_hgrn -from .fused_recurrent import fused_recurrent_hgrn - -__all__ = [ - 'chunk_hgrn', - 'fused_recurrent_hgrn', -] diff --git a/build/torch-xpu/ops/hgrn/chunk.py b/build/torch-xpu/ops/hgrn/chunk.py deleted file mode 100644 index 746de2c73c397e0ad964cdf8d19807e94f6e1202..0000000000000000000000000000000000000000 --- a/build/torch-xpu/ops/hgrn/chunk.py +++ /dev/null @@ -1,287 +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 - -# this function implements the chunkwise form of HGRN, inspired by -# [Volodymyr Kyrylov in his blog post](https://proger.github.io/posts/scan/chunk.html) -# also refer to the `accelerated-scan` lib: https://github.com/proger/accelerated-scan - -# from tests on H800, with B, D = 16, 128, we see that the chunk can be greatly faster than the recurrent: -# -# Performance: -# seq_len chunk recurrent chunk_bwd recurrent_bwd -# 0 128.0 0.039360 0.061056 0.312160 0.205008 -# 1 256.0 0.045824 0.123712 0.308784 0.297696 -# 2 512.0 0.058688 0.241952 0.310720 0.626528 -# 3 1024.0 0.088288 0.476992 0.313184 1.333152 -# 4 2048.0 0.169472 0.943264 0.452464 2.724864 -# 5 4096.0 0.329920 1.886144 0.881600 5.551520 -# 6 8192.0 0.647872 3.755040 1.740496 11.117184 -# 7 16384.0 1.272064 7.520576 3.446608 22.362528 - - -import torch -import triton -import triton.language as tl - -from ...ops.utils.op import exp -from ...utils import autotune_cache_kwargs, input_guard - - -@triton.autotune( - configs=[ - triton.Config({'BD': 32}, num_warps=1), - triton.Config({'BD': 32}, num_warps=2), - triton.Config({'BD': 32}, num_warps=4), - triton.Config({'BD': 32}, num_warps=8), - triton.Config({'BD': 64}, num_warps=1), - triton.Config({'BD': 64}, num_warps=2), - triton.Config({'BD': 64}, num_warps=4), - triton.Config({'BD': 64}, num_warps=8), - triton.Config({'BD': 128}, num_warps=1), - triton.Config({'BD': 128}, num_warps=2), - triton.Config({'BD': 128}, num_warps=4), - triton.Config({'BD': 128}, num_warps=8), - ], - key=['D'], - **autotune_cache_kwargs, -) -@triton.jit(do_not_specialize=['T']) -def chunk_hgrn_fwd_kernel_h( - x, - g, - gc, - o, - h0, - T, - D: tl.constexpr, - BT: tl.constexpr, - BD: tl.constexpr, - USE_INITIAL_STATE: tl.constexpr, -): - i_d, i_t, i_b = tl.program_id(0), tl.program_id(1), tl.program_id(2) - o_d = i_d * BD + tl.arange(0, BD) - mask = o_d < D - - p_x = x + i_b * T * D + i_t * BT * D + o_d - p_g = g + i_b * T * D + i_t * BT * D + o_d - p_gc = gc + i_b * T * D + i_t * BT * D + o_d - p_o = o + i_b * T * D + i_t * BT * D + o_d - - b_h = tl.zeros([BD], dtype=tl.float32) - b_gc = tl.zeros([BD], dtype=tl.float32) - if USE_INITIAL_STATE: - if i_t == 0: - b_h += tl.load(h0 + i_b * D + o_d, mask=mask, other=0).to(tl.float32) - for i in range(0, BT): - mask_t = mask & ((i_t * BT + i) < T) - b_x = tl.load(p_x, mask=mask_t, other=0).to(tl.float32) - b_g = tl.load(p_g, mask=mask_t, other=0).to(tl.float32) - b_h = exp(b_g) * b_h + b_x - b_gc = b_gc + b_g - tl.store(p_gc, b_gc.to(p_o.dtype.element_ty), mask=mask_t) - tl.store(p_o, b_h.to(p_o.dtype.element_ty), mask=mask_t) - - p_x += D - p_g += D - p_gc += D - p_o += D - - -@triton.jit(do_not_specialize=['T']) -def chunk_hgrn_fwd_kernel_o( - gc, - o, - s_b, - s_t, - s_d, - T, - D: tl.constexpr, - BT: tl.constexpr, - BD: tl.constexpr, -): - i_d, i_b = tl.program_id(0), tl.program_id(1) - o_d = i_d * BD + tl.arange(0, BD) - mask = o_d < D - - for i_t in range(1, tl.cdiv(T, BT)): - p_gc = tl.make_block_ptr(gc + i_b * s_b, (T, D), (s_t, s_d), (i_t * BT, i_d * BD), (BT, BD), (1, 0)) - p_o = tl.make_block_ptr(o + i_b * s_b, (T, D), (s_t, s_d), (i_t * BT, i_d * BD), (BT, BD), (1, 0)) - - # [BD,] - b_h0 = tl.load(o + i_b * T * D + i_t * BT * D - D + o_d, mask=mask, other=0).to(tl.float32) - # [BT, BD] - b_gc = tl.load(p_gc, boundary_check=(0, 1)).to(tl.float32) - b_o = tl.load(p_o, boundary_check=(0, 1)).to(tl.float32) - b_o = b_o + exp(b_gc) * b_h0[None, :] - tl.store(p_o, b_o.to(p_o.dtype.element_ty), boundary_check=(0, 1)) - - -@triton.autotune( - configs=[ - triton.Config({'BD': BD}, num_warps=num_warps) - for BD in [32, 64, 128] - for num_warps in [1, 2, 4, 8] - ], - key=['D'], - **autotune_cache_kwargs, -) -@triton.jit(do_not_specialize=['T']) -def chunk_hgrn_bwd_kernel_h( - g, - gc, - dx, - do, - T, - D: tl.constexpr, - BT: tl.constexpr, - BD: tl.constexpr, -): - i_d, i_t, i_b = tl.program_id(0), tl.program_id(1), tl.program_id(2) - o_d = i_d * BD + tl.arange(0, BD) - mask = o_d < D - BC = min(BT, T - i_t * BT) - NT = tl.num_programs(1) - - p_g = g + (i_b * T + i_t * BT + BC - 1) * D + o_d - p_gc = gc + (i_b * T + i_t * BT + BC - 1) * D + o_d - p_dx = dx + (i_b * T + i_t * BT + BC - 1) * D + o_d - p_do = do + (i_b * T + i_t * BT + BC - 1) * D + o_d - - if i_t == NT - 1: - b_gc = tl.zeros([BD], dtype=tl.float32) - else: - b_gc = tl.load(g + (i_b * T + i_t * BT + BT) * D + o_d, mask=mask, other=0).to(tl.float32) - b_dh = tl.zeros([BD], dtype=tl.float32) - for _ in range(BC - 1, -1, -1): - tl.store(p_gc, b_gc.to(p_gc.dtype.element_ty), mask=mask) - - b_g = tl.load(p_g, mask=mask, other=0).to(tl.float32) - b_do = tl.load(p_do, mask=mask, other=0).to(tl.float32) - - b_gc = b_gc + b_g - b_dh = b_dh + b_do - b_dx = b_dh - b_dh = b_dh * exp(b_g) - - tl.store(p_dx, b_dx.to(p_dx.dtype.element_ty), mask=mask) - - p_g -= D - p_gc -= D - p_dx -= D - p_do -= D - - -@triton.jit(do_not_specialize=['T']) -def chunk_hgrn_bwd_kernel_o( - g, - gc, - o, - dx, - dg, - s_b, - s_t, - s_d, - T, - D: tl.constexpr, - BT: tl.constexpr, - BD: tl.constexpr, -): - i_d, i_b = tl.program_id(0), tl.program_id(1) - o_d = i_d * BD + tl.arange(0, BD) - mask = o_d < D - - for i_t in range(tl.cdiv(T, BT) - 1, -1, -1): - p_g = tl.make_block_ptr(g + i_b * s_b, (T, D), (s_t, s_d), (i_t * BT, i_d * BD), (BT, BD), (1, 0)) - p_gc = tl.make_block_ptr(gc + i_b * s_b, (T, D), (s_t, s_d), (i_t * BT, i_d * BD), (BT, BD), (1, 0)) - p_o = tl.make_block_ptr(o + i_b * s_b, (T, D), (s_t, s_d), (i_t * BT - 1, i_d * BD), (BT, BD), (1, 0)) - p_dx = tl.make_block_ptr(dx + i_b * s_b, (T, D), (s_t, s_d), (i_t * BT, i_d * BD), (BT, BD), (1, 0)) - p_dg = tl.make_block_ptr(dg + i_b * s_b, (T, D), (s_t, s_d), (i_t * BT, i_d * BD), (BT, BD), (1, 0)) - - # [BD,] - mask_t = mask & ((i_t + 1) * BT < T) - b_ht = tl.load(dx + i_b * T * D + (i_t + 1) * BT * D + o_d, mask=mask_t, other=0).to(tl.float32) - # [BT, BD] - b_g = tl.load(p_g, boundary_check=(0, 1)).to(tl.float32) - b_gc = tl.load(p_gc, boundary_check=(0, 1)).to(tl.float32) - b_o = tl.load(p_o, boundary_check=(0, 1)).to(tl.float32) - b_dx = tl.load(p_dx, boundary_check=(0, 1)).to(tl.float32) - - b_dx = b_dx + exp(b_gc) * b_ht[None, :] - b_dg = b_o * b_dx * exp(b_g) - 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)) - - -class ChunkHGRNFunction(torch.autograd.Function): - - @staticmethod - @input_guard - def forward(ctx, x, g, initial_state=None, output_final_state=False): - B, T, D = x.shape - BT, BD = 128, min(64, triton.next_power_of_2(D)) - num_warps = 8 if BD == 64 else 4 - - gc = torch.empty_like(g, dtype=torch.float) - o = torch.empty_like(x, dtype=torch.float) - def grid(meta): return (triton.cdiv(D, meta['BD']), triton.cdiv(T, meta['BT']), B) - chunk_hgrn_fwd_kernel_h[grid]( - x, g, gc, o, initial_state, - T=T, D=D, BT=BT, - USE_INITIAL_STATE=initial_state is not None, - ) - def grid(meta): return (triton.cdiv(D, meta['BD']), B) - chunk_hgrn_fwd_kernel_o[grid]( - gc, o, - o.stride(-3), o.stride(-2), o.stride(-1), - T=T, D=D, BT=BT, BD=BD, - num_warps=num_warps, - ) - final_state = None - if output_final_state: - final_state = o[:, -1].clone() - o = o.to(x.dtype) - ctx.save_for_backward(g, o, initial_state) - return o, final_state - - @staticmethod - @input_guard - def backward(ctx, do, dht=None): - g, o, initial_state = ctx.saved_tensors - B, T, D = do.shape - BT, BD = 128, min(64, triton.next_power_of_2(D)) - num_warps = 8 if BD == 64 else 4 - - gc = torch.empty_like(g, dtype=torch.float) - dx = torch.empty_like(o, dtype=torch.float) - def grid(meta): return (triton.cdiv(D, meta['BD']), triton.cdiv(T, meta['BT']), B) - chunk_hgrn_bwd_kernel_h[grid]( - g, gc, dx, do, - T=T, D=D, BT=BT, - ) - - dg = torch.empty_like(g, dtype=torch.float) - def grid(meta): return (triton.cdiv(D, meta['BD']), B) - chunk_hgrn_bwd_kernel_o[grid]( - g, gc, o, dx, dg, - o.stride(-3), o.stride(-2), o.stride(-1), - T=T, D=D, BT=BT, BD=BD, - num_warps=num_warps, - ) - if initial_state is not None: - dg[:, 0] = (initial_state * dx[:, 0] * g[:, 0].float().exp()).to(dg.dtype) - - return dx.to(o.dtype), dg, None, None - - -@torch.compiler.disable -def chunk_hgrn( - x: torch.Tensor, - g: torch.Tensor, - initial_state: torch.Tensor = None, - output_final_state: bool = False, -) -> tuple[torch.Tensor, torch.Tensor]: - return ChunkHGRNFunction.apply(x, g, initial_state, output_final_state) diff --git a/build/torch-xpu/ops/hgrn/fused_recurrent.py b/build/torch-xpu/ops/hgrn/fused_recurrent.py deleted file mode 100644 index 0992dc29c8a1b588e9dc135d1fbaaeb7299d208e..0000000000000000000000000000000000000000 --- a/build/torch-xpu/ops/hgrn/fused_recurrent.py +++ /dev/null @@ -1,323 +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 -import triton.language as tl - -from ...ops.utils.op import exp -from ...utils import autotune_cache_kwargs, input_guard - - -@triton.heuristics({ - 'USE_INITIAL_STATE': lambda args: args['h0'] is not None, - 'STORE_FINAL_STATE': lambda args: args['ht'] is not None, - 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None, -}) -@triton.autotune( - configs=[ - triton.Config({'BD': BD}, num_warps=num_warps) - for BD in [32, 64, 128] - for num_warps in [1, 2, 4, 8] - ], - key=['D'], - **autotune_cache_kwargs, -) -@triton.jit(do_not_specialize=['T']) -def fused_recurrent_hgrn_fwd_kernel( - x, - g, - o, - h0, - ht, - cu_seqlens, - T, - D: tl.constexpr, - BD: tl.constexpr, - USE_INITIAL_STATE: tl.constexpr, - STORE_FINAL_STATE: tl.constexpr, - IS_VARLEN: tl.constexpr, -): - i_d, i_n = tl.program_id(0), tl.program_id(1) - if IS_VARLEN: - bos, eos = tl.load(cu_seqlens + i_n).to(tl.int64), tl.load(cu_seqlens + i_n + 1).to(tl.int64) - T = eos - bos - else: - bos, eos = i_n * T, i_n * T + T - - o_d = i_d * BD + tl.arange(0, BD) - mask = o_d < D - - p_x = x + bos * D + o_d - p_g = g + bos * D + o_d - p_o = o + bos * D + o_d - - b_h = tl.zeros([BD], dtype=tl.float32) - if USE_INITIAL_STATE: - p_h0 = h0 + i_n * D + o_d - b_h += tl.load(p_h0, mask=mask, other=0).to(tl.float32) - for _ in range(0, T): - b_x = tl.load(p_x, mask=mask, other=0).to(tl.float32) - b_g = tl.load(p_g, mask=mask, other=0).to(tl.float32) - b_h = exp(b_g) * b_h + b_x - tl.store(p_o, b_h.to(p_o.dtype.element_ty), mask=mask) - - p_x += D - p_g += D - p_o += D - - if STORE_FINAL_STATE: - p_ht = ht + i_n * D + o_d - tl.store(p_ht, b_h.to(p_ht.dtype.element_ty), mask=mask) - - -@triton.heuristics({ - 'USE_INITIAL_STATE': lambda args: args['h0'] is not None, - 'USE_FINAL_STATE_GRADIENT': lambda args: args['dht'] is not None, - 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None, -}) -@triton.autotune( - configs=[ - triton.Config({'BD': BD}, num_warps=num_warps) - for BD in [32, 64, 128] - for num_warps in [1, 2, 4, 8] - ], - key=['D'], - **autotune_cache_kwargs, -) -@triton.jit(do_not_specialize=['T']) -def fused_recurrent_hgrn_bwd_kernel( - g, - o, - h0, - dx, - dg, - do, - dht, - dh0, - cu_seqlens, - T, - D: tl.constexpr, - BD: tl.constexpr, - USE_INITIAL_STATE: tl.constexpr, - USE_FINAL_STATE_GRADIENT: tl.constexpr, - IS_VARLEN: tl.constexpr, -): - i_d, i_n = tl.program_id(0), tl.program_id(1) - if IS_VARLEN: - bos, eos = tl.load(cu_seqlens + i_n).to(tl.int64), tl.load(cu_seqlens + i_n + 1).to(tl.int64) - T = eos - bos - else: - bos, eos = i_n * T, i_n * T + T - - o_d = i_d * BD + tl.arange(0, BD) - mask = o_d < D - - p_g = g + (bos + T - 1) * D + o_d - p_o = o + (bos + T - 2) * D + o_d - p_dx = dx + (bos + T - 1) * D + o_d - p_dg = dg + (bos + T - 1) * D + o_d - p_do = do + (bos + T - 1) * D + o_d - - b_dh = tl.zeros([BD], dtype=tl.float32) - if USE_FINAL_STATE_GRADIENT: - p_dht = dht + i_n * D + o_d - b_dh += tl.load(p_dht, mask=mask, other=0).to(tl.float32) - - for i in range(T - 1, -1, -1): - b_g = tl.load(p_g, mask=mask, other=0).to(tl.float32) - b_do = tl.load(p_do, mask=mask, other=0).to(tl.float32) - if i > 0: - b_o = tl.load(p_o, mask=mask, other=0).to(tl.float32) - elif USE_INITIAL_STATE: - b_o = tl.load(h0 + i_n * D + o_d, mask=mask, other=0).to(tl.float32) - else: - b_o = tl.zeros([BD], dtype=tl.float32) - - b_dh = b_dh + b_do - b_dx = b_dh - b_dh = b_dh * exp(b_g) - b_dg = b_dh * b_o - tl.store(p_dx, b_dx.to(p_dx.dtype.element_ty), mask=mask) - tl.store(p_dg, b_dg.to(p_dg.dtype.element_ty), mask=mask) - - p_g -= D - p_o -= D - p_dx -= D - p_dg -= D - p_do -= D - - if USE_INITIAL_STATE: - p_dh0 = dh0 + i_n * D + o_d - tl.store(p_dh0, b_dh.to(p_dh0.dtype.element_ty), mask=mask) - - -def fused_recurrent_hgrn_fwd( - x: torch.Tensor, - g: torch.Tensor, - initial_state: torch.Tensor = None, - output_final_state: bool = False, - cu_seqlens: torch.LongTensor | None = None, -) -> tuple[torch.Tensor, torch.Tensor]: - B, T, D = x.shape - N = B if cu_seqlens is None else len(cu_seqlens) - 1 - - o = torch.empty_like(x) - final_state = x.new_empty(N, D) if output_final_state else None - - def grid(meta): return (triton.cdiv(D, meta['BD']), N) - fused_recurrent_hgrn_fwd_kernel[grid]( - x=x, - g=g, - o=o, - h0=initial_state, - ht=final_state, - cu_seqlens=cu_seqlens, - T=T, - D=D, - ) - return o, final_state - - -def fused_recurrent_hgrn_bwd( - g: torch.Tensor, - o: torch.Tensor, - do: torch.Tensor, - dht: torch.Tensor = None, - initial_state: torch.Tensor = None, - cu_seqlens: torch.LongTensor | None = None, -) -> tuple[torch.Tensor, torch.Tensor]: - B, T, D = do.shape - N = B if cu_seqlens is None else len(cu_seqlens) - 1 - - dx = torch.empty_like(o, dtype=torch.float) - dg = torch.empty_like(g, dtype=torch.float) - dh0 = torch.empty_like(initial_state, dtype=torch.float) if initial_state is not None else None - def grid(meta): return (triton.cdiv(D, meta['BD']), N) - fused_recurrent_hgrn_bwd_kernel[grid]( - g=g, - o=o, - h0=initial_state, - dx=dx, - dg=dg, - do=do, - dht=dht, - dh0=dh0, - cu_seqlens=cu_seqlens, - T=T, - D=D, - ) - return dx, dg, dh0 - - -class FusedRecurrentHGRNFunction(torch.autograd.Function): - - @staticmethod - @input_guard - def forward( - ctx, - x: torch.Tensor, - g: torch.Tensor, - initial_state: torch.Tensor = None, - output_final_state: bool = False, - cu_seqlens: torch.LongTensor | None = None, - ): - o, ht = fused_recurrent_hgrn_fwd( - x=x, - g=g, - initial_state=initial_state, - output_final_state=output_final_state, - cu_seqlens=cu_seqlens, - ) - ctx.save_for_backward(g, o, initial_state) - ctx.cu_seqlens = cu_seqlens - return o, ht - - @staticmethod - @input_guard - def backward(ctx, do, dht=None): - g, o, initial_state = ctx.saved_tensors - cu_seqlens = ctx.cu_seqlens - - dx, dg, dh0 = fused_recurrent_hgrn_bwd( - g=g, - o=o, - do=do, - dht=dht, - initial_state=initial_state, - cu_seqlens=cu_seqlens, - ) - return dx, dg, dh0, None, None - - -@torch.compiler.disable -def fused_recurrent_hgrn( - x: torch.Tensor, - g: torch.Tensor, - initial_state: torch.Tensor = None, - output_final_state: bool = False, - cu_seqlens: torch.LongTensor | None = None, -) -> tuple[torch.Tensor, torch.Tensor]: - r""" - Args: - x (torch.Tensor): - inputs of shape `[B, T, D]. - g (torch.Tensor): - Forget gates of shape `[B, T, D]`. - initial_state (Optional[torch.Tensor]): - Initial state of shape `[N, D]` for `N` input sequences. - For equal-length input sequences, `N` equals the batch size `B`. - Default: `None`. - output_final_state (Optional[bool]): - Whether to output the final state of shape `[N, D]`. Default: `False`. - cu_seqlens (torch.LongTensor): - Cumulative sequence lengths of shape `[N+1]` used for variable-length training, - consistent with the FlashAttention API. - - Returns: - o (torch.Tensor): - Outputs of shape `[B, T, D]`. - final_state (torch.Tensor): - Final state of shape `[N, D]` if `output_final_state=True` else `None`. - - Examples:: - >>> import torch - >>> import torch.nn.functional as F - >>> from einops import rearrange - >>> from ...ops.hgrn import fused_recurrent_hgrn - # inputs with equal lengths - >>> B, T, D = 4, 2048, 512 - >>> x = torch.randn(B, T, D, device='cuda') - >>> g = F.logsigmoid(torch.randn(B, T, D, device='cuda')) - >>> h0 = torch.randn(B, D, device='cuda') - >>> o, ht = fused_recurrent_hgrn(x, g, initial_state=h0, output_final_state=True) - # for variable-length inputs, the batch size `B` is expected to be 1 and `cu_seqlens` is required - >>> x, g = map(lambda x: rearrange(x, 'b t d -> 1 (b t) d'), (x, g)) - # for a batch with 4 sequences, `cu_seqlens` with 5 start/end positions are expected - >>> cu_seqlens = x.new_tensor([0, 2048, 4096, 6144, 8192], dtype=torch.long) - >>> o_var, ht_var = fused_recurrent_hgrn(x, g, initial_state=h0, output_final_state=True, cu_seqlens=cu_seqlens) - >>> assert o.allclose(o_var.view(o.shape)) - >>> assert ht.allclose(ht_var) - """ - if cu_seqlens is not None: - if x.shape[0] != 1: - raise ValueError( - f"The batch size is expected to be 1 rather than {x.shape[0]} when using `cu_seqlens`. " - f"Please flatten variable-length inputs before processing.", - ) - if initial_state is not None and initial_state.shape[0] != len(cu_seqlens) - 1: - raise ValueError( - f"The number of initial states is expected to be equal to the number of input sequences, " - f"i.e., {len(cu_seqlens) - 1} rather than {initial_state.shape[0]}.", - ) - return FusedRecurrentHGRNFunction.apply( - x, - g, - initial_state, - output_final_state, - cu_seqlens, - ) diff --git a/build/torch-xpu/ops/hgrn/naive.py b/build/torch-xpu/ops/hgrn/naive.py deleted file mode 100644 index 40f3012fb9cdd50367a573f0e77b1a649114379c..0000000000000000000000000000000000000000 --- a/build/torch-xpu/ops/hgrn/naive.py +++ /dev/null @@ -1,66 +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 - - -def naive_recurrent_hgrn( - x: torch.Tensor, - g: torch.Tensor, - initial_state: torch.Tensor | None = None, - output_final_state: bool | None = False, -) -> torch.Tensor: - dtype = x.dtype - x, g = map(lambda i: i.float(), (x, g)) - B, T, D = x.shape - - h = torch.zeros(B, D, dtype=torch.float, device=x.device) - o = torch.zeros_like(x) - - final_state = None - if initial_state is not None: - h += initial_state - - for i in range(T): - h = g[:, i].exp() * h + x[:, i] - o[:, i] = h - - if output_final_state: - final_state = h - return o.to(dtype), final_state - - -def naive_chunk_hgrn( - x: torch.Tensor, - g: torch.Tensor, - initial_state: torch.Tensor | None = None, - output_final_state: bool | None = False, - chunk_size: int = 64, -) -> torch.Tensor: - dtype = x.dtype - x, g = map(lambda i: i.float(), (x, g)) - B, T, D = x.shape - - gc = g.view(B, chunk_size, D).cumsum(-2).view_as(g) - h = torch.zeros(B, D, dtype=torch.float, device=x.device) - o = torch.zeros_like(x) - - final_state = None - if initial_state is not None: - h += initial_state - - for i in range(0, T, chunk_size): - hp = h - h = torch.zeros(B, D, dtype=torch.float, device=x.device) - for j in range(i, i + chunk_size): - h = g[:, j].exp() * h + x[:, j] - o[:, j] = hp * gc[:, j].exp() + h - h = o[:, j].clone() - - if output_final_state: - final_state = h - return o.to(dtype), final_state diff --git a/build/torch-xpu/ops/kda/__init__.py b/build/torch-xpu/ops/kda/__init__.py deleted file mode 100644 index f896d2b9bc3e7b8db0ac0912224358b27e154179..0000000000000000000000000000000000000000 --- a/build/torch-xpu/ops/kda/__init__.py +++ /dev/null @@ -1,14 +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 .chunk import chunk_kda -from .fused_recurrent import fused_recurrent_kda - -__all__ = [ - "chunk_kda", - "fused_recurrent_kda", -] diff --git a/build/torch-xpu/ops/kda/backends/__init__.py b/build/torch-xpu/ops/kda/backends/__init__.py deleted file mode 100644 index 2b5d14c9e3bf8784d3fbc38b135ac392060c3fb3..0000000000000000000000000000000000000000 --- a/build/torch-xpu/ops/kda/backends/__init__.py +++ /dev/null @@ -1,19 +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 - -"""KDA backends.""" - -from ....ops.backends import BackendRegistry, dispatch -from ....ops.kda.backends.flashkda import FlashKDABackend -from ....ops.kda.backends.tilelang import KDATileLangBackend - -kda_registry = BackendRegistry("kda") -kda_registry.register(FlashKDABackend()) -kda_registry.register(KDATileLangBackend()) - - -__all__ = ['dispatch', 'kda_registry'] diff --git a/build/torch-xpu/ops/kda/backends/flashkda.py b/build/torch-xpu/ops/kda/backends/flashkda.py deleted file mode 100644 index 16b6cfba000efddb7f44c7f9ab8d9fef65243549..0000000000000000000000000000000000000000 --- a/build/torch-xpu/ops/kda/backends/flashkda.py +++ /dev/null @@ -1,151 +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 -# -# Copyright (c) 2026 Moonshot AI - -"""FlashKDA CUTLASS forward backend for chunk_kda (inference only).""" - -from __future__ import annotations - -from typing import TYPE_CHECKING - -import torch - -from ....ops.backends import BaseBackend - -if TYPE_CHECKING: - from ....ops.cp import FLACPContext - - -class FlashKDABackend(BaseBackend): - """Copyright (c) 2026 Moonshot AI - - Fused CUTLASS forward (replaces the multi-kernel Triton path). - https://github.com/MoonshotAI/FlashKDA - - Enabled only under ``torch.inference_mode()``; disable with ``FLA_FLASH_KDA=0``. - The kernel fuses q/k L2 norm, beta sigmoid, and the KDA gate, so callers must pass - raw tensors and set all three ``*_in_kernel`` flags. - """ - - backend_type = "flashkda" - package_name = "flash_kda" - env_var = "FLA_FLASH_KDA" - default_enable = True - priority = 3 - - def chunk_kda_verifier( - self, - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - g: torch.Tensor, - beta: torch.Tensor, - scale: float | None = None, - initial_state: torch.Tensor | None = None, - output_final_state: bool = False, - use_qk_l2norm_in_kernel: bool = False, - use_gate_in_kernel: bool = False, - use_beta_sigmoid_in_kernel: bool = False, - state_v_first: bool = False, - cu_seqlens: torch.LongTensor | None = None, - cu_seqlens_cpu: torch.LongTensor | None = None, - safe_gate: bool = False, - lower_bound: float | None = None, - disable_recompute: bool = False, - return_intermediate_states: bool = False, - cp_context: FLACPContext | None = None, - **kwargs, - ) -> tuple[bool, str | None]: - if torch.is_grad_enabled(): - return False, "FlashKDA only supports inference mode" - if q.dtype != torch.bfloat16: - return False, f"FlashKDA requires bfloat16, got {q.dtype}" - if q.shape[-1] != 128: - return False, f"FlashKDA requires K=128, got {q.shape[-1]}" - if v.shape[-1] != 128: - return False, f"FlashKDA requires V=128, got {v.shape[-1]}" - if v.shape[2] != q.shape[2]: - return False, f"FlashKDA does not support GVA (HV={v.shape[2]} != H={q.shape[2]})" - if not use_gate_in_kernel: - return False, "FlashKDA requires use_gate_in_kernel=True" - if not use_qk_l2norm_in_kernel: - return False, "FlashKDA requires use_qk_l2norm_in_kernel=True" - if not use_beta_sigmoid_in_kernel: - return False, "FlashKDA requires use_beta_sigmoid_in_kernel=True" - if not state_v_first: - return False, "FlashKDA requires state_v_first=True" - if cp_context is not None: - return False, "FlashKDA does not support context parallel" - if return_intermediate_states: - return False, "FlashKDA does not support return_intermediate_states" - if not safe_gate: - return False, "FlashKDA requires safe_gate=True" - return True, None - - def chunk_kda( - self, - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - g: torch.Tensor, - beta: torch.Tensor, - scale: float | None = None, - initial_state: torch.Tensor | None = None, - output_final_state: bool = False, - use_qk_l2norm_in_kernel: bool = False, - use_gate_in_kernel: bool = False, - use_beta_sigmoid_in_kernel: bool = False, - state_v_first: bool = False, - cu_seqlens: torch.LongTensor | None = None, - cu_seqlens_cpu: torch.LongTensor | None = None, - safe_gate: bool = False, - lower_bound: float | None = None, - disable_recompute: bool = False, - return_intermediate_states: bool = False, - cp_context: FLACPContext | None = None, - A_log: torch.Tensor | None = None, - dt_bias: torch.Tensor | None = None, - **kwargs, - ): - import flash_kda - - if scale is None: - scale = q.shape[-1] ** -0.5 - - K = q.shape[-1] - HV, V = v.shape[2], v.shape[-1] - N = len(cu_seqlens) - 1 if cu_seqlens is not None else q.shape[0] - - if dt_bias is not None and dt_bias.ndim == 1: - dt_bias = dt_bias.view(HV, -1) - - out_buf = torch.empty_like(v) - - if initial_state is not None: - initial_state = initial_state.contiguous() - - final_state = None - if output_final_state: - final_state = torch.empty(N, HV, K, V, dtype=torch.float32, device=q.device) - - if cu_seqlens is not None and cu_seqlens.dtype != torch.long: - cu_seqlens = cu_seqlens.to(torch.long) - - flash_kda.fwd( - q, k, v, g, beta, - scale, - out_buf, - A_log=A_log, - dt_bias=dt_bias, - lower_bound=lower_bound, - initial_state=initial_state, - final_state=final_state, - cu_seqlens=cu_seqlens, - ) - - return out_buf, final_state diff --git a/build/torch-xpu/ops/kda/backends/tilelang/__init__.py b/build/torch-xpu/ops/kda/backends/tilelang/__init__.py deleted file mode 100644 index 3bc44e67242d147cca00e39effd8b2667566be02..0000000000000000000000000000000000000000 --- a/build/torch-xpu/ops/kda/backends/tilelang/__init__.py +++ /dev/null @@ -1,84 +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 - -"""TileLang backend for KDA operations.""" - -from __future__ import annotations - -import torch - -from .....ops.backends import BaseBackend -from .....utils import find_spec_cached - -_TILELANG_AVAILABLE = find_spec_cached("tilelang") is not None - - -class KDATileLangBackend(BaseBackend): - - backend_type = "tilelang" - package_name = "tilelang" - env_var = "FLA_TILELANG" - - @classmethod - def is_available(cls) -> bool: - return _TILELANG_AVAILABLE - - def chunk_kda_bwd_wy_dqkg_fused_verifier( - self, - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - v_new: torch.Tensor, - g: torch.Tensor, - beta: torch.Tensor, - A: torch.Tensor, - h: torch.Tensor, - do: torch.Tensor, - dh: torch.Tensor, - dv: torch.Tensor, - scale: float | None = None, - state_v_first: bool = False, - cu_seqlens: torch.LongTensor | None = None, - chunk_size: int = 64, - chunk_indices: torch.LongTensor | None = None, - ) -> tuple[bool, str | None]: - if v.shape[2] != k.shape[2]: - return False, ( - "TileLang backend does not support GQA (v has more heads than k); " - "use repeat_interleave on k/q to match v's head count, or fall back to Triton" - ) - return True, None - - def chunk_kda_bwd_wy_dqkg_fused( - self, - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - v_new: torch.Tensor, - g: torch.Tensor, - beta: torch.Tensor, - A: torch.Tensor, - h: torch.Tensor, - do: torch.Tensor, - dh: torch.Tensor, - dv: torch.Tensor, - scale: float | None = None, - state_v_first: bool = False, - cu_seqlens: torch.LongTensor | None = None, - chunk_size: int = 64, - chunk_indices: torch.LongTensor | None = None, - ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: - from .....ops.kda.backends.tilelang.chunk_bwd_dqkg import ( - chunk_kda_bwd_wy_dqkg_fused_tilelang, - ) - return chunk_kda_bwd_wy_dqkg_fused_tilelang( - q=q, k=k, v=v, v_new=v_new, g=g, beta=beta, A=A, - h=h, do=do, dh=dh, dv=dv, - scale=scale, cu_seqlens=cu_seqlens, - chunk_size=chunk_size, chunk_indices=chunk_indices, - state_v_first=state_v_first, - ) diff --git a/build/torch-xpu/ops/kda/backends/tilelang/chunk_bwd_dqkg.py b/build/torch-xpu/ops/kda/backends/tilelang/chunk_bwd_dqkg.py deleted file mode 100644 index 19ed77ac0395afe869b42f331f504255f58927d6..0000000000000000000000000000000000000000 --- a/build/torch-xpu/ops/kda/backends/tilelang/chunk_bwd_dqkg.py +++ /dev/null @@ -1,471 +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 - -"""TileLang implementation of chunk_kda_bwd_wy_dqkg_fused. - -Ported from the Triton kernel in fla/ops/kda/chunk_bwd.py. -Key structural differences from the common chunk_bwd_dqkwg kernel: - - K-loop is INTERNAL (no NK grid dimension): grid is (NT, B*H) or (NT, H) for varlen - - Gate g is 2D: (B, T, H, K), not (B, T, H) - - Extra inputs: v_new, beta, A (attention matrix) - - Extra outputs: dv2, db, dA - - dgk is per-K accumulator (BK,), not scalar - - dA post-processing: causal mask -> dA@A -> A@(dA@A) -> negate - - exp2 is always used (no exp/exp2 option) -""" - - -import tilelang -import tilelang.language as T -import torch -import triton - -from .....ops.utils import prepare_chunk_indices -from .....utils import check_shared_mem - - -@tilelang.jit(pass_configs={ - tilelang.PassConfigKey.TL_DISABLE_TMA_LOWER: True, - tilelang.PassConfigKey.TL_ENABLE_FAST_MATH: True, -}) -def _build_kda_bwd_kernel( - B, H, K, V, BT, BK, BV, - hD1, hD2, - dtype_str, - STATE_V_FIRST, IS_VARLEN=False, - num_warps=4, -): - dtype_map = {'float16': T.float16, 'bfloat16': T.bfloat16, 'float32': T.float32} - dtype = dtype_map[dtype_str] - NK = tilelang.cdiv(K, BK) - NV = tilelang.cdiv(V, BV) - threads = num_warps * 32 - tile_hD1, tile_hD2 = (BV, BK) if STATE_V_FIRST else (BK, BV) - - # Rebind to underscore-prefixed locals for closure capture - _B, _H, _K, _V = B, H, K, V - _BT, _BK, _BV, _NK, _NV = BT, BK, BV, NK, NV - _hD1, _hD2, _thD1, _thD2 = hD1, hD2, tile_hD1, tile_hD2 - _dtype = dtype - _threads = threads - _TS = STATE_V_FIRST - _VAR = IS_VARLEN - - T_d, NT_d, total_h_d, N_seqs_d = T.dynamic("T, NT, total_h, N_seqs") - - qk_s = (_B, T_d, _H, _K) - v_s = (_B, T_d, _H, _V) - g_s = (_B, T_d, _H, _K) # 2D gate: per-K dimension - beta_s = (_B, T_d, _H) - A_s = (_B, T_d, _H, _BT) # attention matrix: (B, T, H, BT) - h_s = (total_h_d, _hD1, _hD2) - db_s = (_B, T_d, _H) - dA_s = (_B, T_d, _H, _BT) - - @T.macro - def kernel_body(q, k, v, v_new, g, beta, A, h, do, dh, dq, dk, dv, dv2, dg, db, dA, scale, - i_b, i_h, t_s, T_seq, i_t_local, h_idx): - # ========== 1. Load beta and A for this chunk ========== - s_beta = T.alloc_shared((_BT,), T.float32) - T.copy(beta[i_b, t_s:t_s + _BT, i_h], s_beta, disable_tma=True) - # Zero OOB positions for varlen safety - for _i in T.Parallel(_BT): - if (i_t_local * _BT + _i) >= T_seq: - s_beta[_i] = 0.0 - - # A in Triton is loaded transposed: b_A[bt, t] = A[i_b, t_s+t, i_h, bt] - # Load, transpose, and zero OOB rows/cols for varlen safety - s_A = T.alloc_shared((_BT, _BT), _dtype) - T.copy(A[i_b, t_s:t_s + _BT, i_h, 0:_BT], s_A) - f_A_tmp = T.alloc_fragment((_BT, _BT), _dtype) - for _i, _j in T.Parallel(_BT, _BT): - valid = ((i_t_local * _BT + _i) < T_seq) & ((i_t_local * _BT + _j) < T_seq) - f_A_tmp[_i, _j] = T.if_then_else(valid, s_A[_j, _i], T.cast(0, _dtype)) - T.copy(f_A_tmp, s_A) - - # ========== 2. Init dA, db accumulators ========== - b_dA = T.alloc_fragment((_BT, _BT), T.float32) - T.clear(b_dA) - b_db = T.alloc_fragment((_BT,), T.float32) - T.clear(b_db) - - # ========== 3. Pre-K V-loop (dA_v contribution, dv2, db_v) ========== - # These computations are K-independent: dA += dv@v^T, dv2 = A@dv*beta, db += sum((A@dv)*v) - s_dv_pre = T.alloc_shared((_BT, _BV), _dtype) - s_v_orig = T.alloc_shared((_BT, _BV), _dtype) - b_dvb = T.alloc_fragment((_BT, _BV), T.float32) - s_dv2 = T.alloc_shared((_BT, _BV), _dtype) - f_dvv = T.alloc_fragment((_BT, _BV), T.float32) - f_dvv_row = T.alloc_fragment((_BT,), T.float32) - - for i_v in T.serial(_NV): - v_off = i_v * _BV - - T.copy(dv[i_b, t_s:t_s + _BT, i_h, v_off:v_off + _BV], s_dv_pre) - T.copy(v[i_b, t_s:t_s + _BT, i_h, v_off:v_off + _BV], s_v_orig) - - # b_dA += dv @ v^T (BT, BT) - T.gemm(s_dv_pre, s_v_orig, b_dA, transpose_B=True) - - # b_dvb = A @ dv (BT, BV) - T.clear(b_dvb) - T.gemm(s_A, s_dv_pre, b_dvb) - - # Store dv2 = dvb * beta[:, None] - for _i, _j in T.Parallel(_BT, _BV): - s_dv2[_i, _j] = T.cast(b_dvb[_i, _j] * s_beta[_i], _dtype) - for _i, _j in T.Parallel(_BT, _BV): - if (i_t_local * _BT + _i) < T_seq: - dv2[i_b, t_s + _i, i_h, v_off + _j] = s_dv2[_i, _j] - - # db += sum(dvb * v, dim=1) - for _i, _j in T.Parallel(_BT, _BV): - f_dvv[_i, _j] = b_dvb[_i, _j] * T.cast(s_v_orig[_i, _j], T.float32) - T.reduce_sum(f_dvv, f_dvv_row, dim=1) - for _i in T.Parallel(_BT): - b_db[_i] = b_db[_i] + f_dvv_row[_i] - - # ========== 4. K-loop ========== - for i_k in T.serial(_NK): - k_off = i_k * _BK - - # Load k, g for this K-block (zero OOB rows for varlen cross-row safety) - s_k = T.alloc_shared((_BT, _BK), _dtype) - T.copy(k[i_b, t_s:t_s + _BT, i_h, k_off:k_off + _BK], s_k) - for _i, _j in T.Parallel(_BT, _BK): - if (i_t_local * _BT + _i) >= T_seq: - s_k[_i, _j] = T.cast(0, _dtype) - - s_g = T.alloc_shared((_BT, _BK), T.float32) - T.copy(g[i_b, t_s:t_s + _BT, i_h, k_off:k_off + _BK], s_g) - for _i, _j in T.Parallel(_BT, _BK): - if (i_t_local * _BT + _i) >= T_seq: - s_g[_i, _j] = 0.0 - - # g at last valid position: (BK,) - last_pos = T.max(0, T.min(_BT, T_seq - i_t_local * _BT) - 1) - s_gn = T.alloc_shared((_BK,), T.float32) - for _j in T.Parallel(_BK): - s_gn[_j] = s_g[last_pos, _j] - - # Per-K accumulators - b_dq = T.alloc_fragment((_BT, _BK), T.float32) - b_dk = T.alloc_fragment((_BT, _BK), T.float32) - b_dw = T.alloc_fragment((_BT, _BK), T.float32) - T.clear(b_dq) - T.clear(b_dk) - T.clear(b_dw) - - # dgk: per-K accumulator for sum(h*dh, axis=0) -- shape (BK,) - # Use shared memory (avoid 1D fragment/reducer issues) - s_dgk = T.alloc_shared((_BK,), T.float32) - for _j in T.Parallel(_BK): - s_dgk[_j] = 0.0 - - # ========== 4a. V-loop (dq, dk, dw, dgk) ========== - s_v_new = T.alloc_shared((_BT, _BV), _dtype) - s_do = T.alloc_shared((_BT, _BV), _dtype) - s_h = T.alloc_shared((_thD1, _thD2), _dtype) - s_dh = T.alloc_shared((_thD1, _thD2), _dtype) - s_dv_tile = T.alloc_shared((_BT, _BV), _dtype) - f_hdh = T.alloc_fragment((_thD1, _thD2), T.float32) - f_hdh_k_nots = T.alloc_fragment((_thD1,), T.float32) - f_hdh_t = T.alloc_fragment((_thD2, _thD1), T.float32) - f_hdh_k_ts = T.alloc_fragment((_thD2,), T.float32) - - for i_v in T.serial(_NV): - v_off = i_v * _BV - - T.copy(v_new[i_b, t_s:t_s + _BT, i_h, v_off:v_off + _BV], s_v_new) - T.copy(do[i_b, t_s:t_s + _BT, i_h, v_off:v_off + _BV], s_do) - - if _TS: - T.copy(h[h_idx, v_off:v_off + _BV, k_off:k_off + _BK], s_h) - T.copy(dh[h_idx, v_off:v_off + _BV, k_off:k_off + _BK], s_dh) - else: - T.copy(h[h_idx, k_off:k_off + _BK, v_off:v_off + _BV], s_h) - T.copy(dh[h_idx, k_off:k_off + _BK, v_off:v_off + _BV], s_dh) - - # GEMMs first (consume s_h/s_dh before any alloc might reuse their memory) - # b_dq += do @ h, b_dk += v_new @ dh - if _TS: - T.gemm(s_do, s_h, b_dq) - T.gemm(s_v_new, s_dh, b_dk) - else: - T.gemm(s_do, s_h, b_dq, transpose_B=True) - T.gemm(s_v_new, s_dh, b_dk, transpose_B=True) - - # b_dw += dv @ h (dv feeds into dw → cross-row dA/dkgb, needs OOB mask) - T.copy(dv[i_b, t_s:t_s + _BT, i_h, v_off:v_off + _BV], s_dv_tile) - - if _TS: - T.gemm(s_dv_tile, s_h, b_dw) - else: - T.gemm(s_dv_tile, s_h, b_dw, transpose_B=True) - - # dgk += sum(h * dh, axis=0) -- per-column sum -> (BK,) - for _i, _j in T.Parallel(_thD1, _thD2): - f_hdh[_i, _j] = T.cast(s_h[_i, _j], T.float32) * T.cast(s_dh[_i, _j], T.float32) - - if not _TS: - T.reduce_sum(f_hdh, f_hdh_k_nots, dim=1) - for _j in T.Parallel(_BK): - s_dgk[_j] = s_dgk[_j] + f_hdh_k_nots[_j] - - else: - for _i, _j in T.Parallel(_thD2, _thD1): - f_hdh_t[_i, _j] = f_hdh[_j, _i] - - T.reduce_sum(f_hdh_t, f_hdh_k_ts, dim=1) - for _j in T.Parallel(_BK): - s_dgk[_j] = s_dgk[_j] + f_hdh_k_ts[_j] - - # ========== 4b. Gate dq, dk ========== - # b_dq = b_dq * exp2(g) * scale - s_dq = T.alloc_shared((_BT, _BK), T.float32) - T.copy(b_dq, s_dq) - for _i, _j in T.Parallel(_BT, _BK): - s_dq[_i, _j] = s_dq[_i, _j] * T.exp2(s_g[_i, _j]) * scale - - # b_dk = b_dk * where(m_t, exp2(gn - g), 0) - s_dk = T.alloc_shared((_BT, _BK), T.float32) - T.copy(b_dk, s_dk) - for _i, _j in T.Parallel(_BT, _BK): - m_t = (i_t_local * _BT + _i) < T_seq - s_dk[_i, _j] = T.if_then_else( - m_t, - s_dk[_i, _j] * T.exp2(s_gn[_j] - s_g[_i, _j]), - 0.0) - - # ========== 4c. kg = k * exp2(g), dgk *= exp2(gn) ========== - s_kg = T.alloc_shared((_BT, _BK), _dtype) - for _i, _j in T.Parallel(_BT, _BK): - s_kg[_i, _j] = T.cast(T.cast(s_k[_i, _j], T.float32) * T.exp2(s_g[_i, _j]), _dtype) - - for _j in T.Parallel(_BK): - s_dgk[_j] = s_dgk[_j] * T.exp2(s_gn[_j]) - - # ========== 4d. dw = -dw ========== - s_dw = T.alloc_shared((_BT, _BK), _dtype) - T.copy(b_dw, s_dw) - for _i, _j in T.Parallel(_BT, _BK): - s_dw[_i, _j] = T.cast(-T.cast(s_dw[_i, _j], T.float32), _dtype) - - # ========== 4e. dA += dw @ kg^T ========== - T.gemm(s_dw, s_kg, b_dA, transpose_B=True) - - # ========== 4f. dkgb = A @ dw ========== - b_dkgb = T.alloc_fragment((_BT, _BK), T.float32) - T.clear(b_dkgb) - T.gemm(s_A, s_dw, b_dkgb) - - # ========== 4g. db += sum(dkgb * kg, dim=1) ========== - f_dkgb_kg = T.alloc_fragment((_BT, _BK), T.float32) - for _i, _j in T.Parallel(_BT, _BK): - f_dkgb_kg[_i, _j] = b_dkgb[_i, _j] * T.cast(s_kg[_i, _j], T.float32) - f_db_row = T.alloc_fragment((_BT,), T.float32) - T.reduce_sum(f_dkgb_kg, f_db_row, dim=1) - for _i in T.Parallel(_BT): - b_db[_i] = b_db[_i] + f_db_row[_i] - - # ========== 4h. dgk += col_sum(k * dk) ========== - f_kdk = T.alloc_fragment((_BT, _BK), T.float32) - for _i, _j in T.Parallel(_BT, _BK): - f_kdk[_i, _j] = T.cast(s_k[_i, _j], T.float32) * s_dk[_i, _j] - # Column sum of kdk -> (BK,) via transpose then row-reduce - f_kdk_t = T.alloc_fragment((_BK, _BT), T.float32) - for _i, _j in T.Parallel(_BK, _BT): - f_kdk_t[_i, _j] = f_kdk[_j, _i] - f_kdk_col = T.alloc_fragment((_BK,), T.float32) - T.reduce_sum(f_kdk_t, f_kdk_col, dim=1) - s_kdk_col = T.alloc_shared((_BK,), T.float32) - T.copy(f_kdk_col, s_kdk_col) - for _j in T.Parallel(_BK): - s_dgk[_j] = s_dgk[_j] + s_kdk_col[_j] - - # ========== 4i. dg = q*dq - kdk + m_last*dgk + kg*dkgb*beta ========== - s_dg_out = T.alloc_shared((_BT, _BK), T.float32) - m_last_pos = T.max(0, T.min(_BT, T_seq - i_t_local * _BT) - 1) - for _i, _j in T.Parallel(_BT, _BK): - q_dq = T.cast(q[i_b, t_s + _i, i_h, k_off + _j], T.float32) * s_dq[_i, _j] - last_term = T.if_then_else(_i == m_last_pos, s_dgk[_j], 0.0) - beta_term = T.cast(s_kg[_i, _j], T.float32) * b_dkgb[_i, _j] * s_beta[_i] - s_dg_out[_i, _j] = q_dq - f_kdk[_i, _j] + last_term + beta_term - - # ========== 4j. dk += dkgb * gb where gb = exp2(g) * beta ========== - for _i, _j in T.Parallel(_BT, _BK): - s_dk[_i, _j] = s_dk[_i, _j] + b_dkgb[_i, _j] * T.exp2(s_g[_i, _j]) * s_beta[_i] - - # ========== 4k. Store dq, dk, dg for this K-block ========== - for _i, _j in T.Parallel(_BT, _BK): - if (i_t_local * _BT + _i) < T_seq: - dq[i_b, t_s + _i, i_h, k_off + _j] = s_dq[_i, _j] - for _i, _j in T.Parallel(_BT, _BK): - if (i_t_local * _BT + _i) < T_seq: - dk[i_b, t_s + _i, i_h, k_off + _j] = s_dk[_i, _j] - for _i, _j in T.Parallel(_BT, _BK): - if (i_t_local * _BT + _i) < T_seq: - dg[i_b, t_s + _i, i_h, k_off + _j] = s_dg_out[_i, _j] - - # ========== 5. dA post-processing ========== - # m_A = (i > j) & valid - # b_dA = where(m_A, b_dA * beta[j], 0) - # b_dA = b_dA @ A - # b_dA = A @ b_dA - # b_dA = where(m_A, -b_dA, 0) - - s_dA = T.alloc_shared((_BT, _BT), T.float32) - T.copy(b_dA, s_dA) - - # Apply causal mask and beta weighting - for _i, _j in T.Parallel(_BT, _BT): - m_valid = ((i_t_local * _BT + _i) < T_seq) & ((i_t_local * _BT + _j) < T_seq) - m_upper = (_i > _j) & m_valid - s_dA[_i, _j] = T.if_then_else(m_upper, s_dA[_i, _j] * s_beta[_j], 0.0) - - # dA = dA @ A - s_dA_dtype = T.alloc_shared((_BT, _BT), _dtype) - for _i, _j in T.Parallel(_BT, _BT): - s_dA_dtype[_i, _j] = T.cast(s_dA[_i, _j], _dtype) - - b_dA2 = T.alloc_fragment((_BT, _BT), T.float32) - T.clear(b_dA2) - T.gemm(s_dA_dtype, s_A, b_dA2) - - # dA = A @ dA - # Reuse s_dA_dtype for the second cast - for _i, _j in T.Parallel(_BT, _BT): - s_dA_dtype[_i, _j] = T.cast(b_dA2[_i, _j], _dtype) - - b_dA3 = T.alloc_fragment((_BT, _BT), T.float32) - T.clear(b_dA3) - T.gemm(s_A, s_dA_dtype, b_dA3) - - # Apply mask and negate, reuse s_dA buffer - for _i, _j in T.Parallel(_BT, _BT): - m_valid = ((i_t_local * _BT + _i) < T_seq) & ((i_t_local * _BT + _j) < T_seq) - m_upper = (_i > _j) & m_valid - s_dA[_i, _j] = T.if_then_else(m_upper, -b_dA3[_i, _j], 0.0) - - for _i, _j in T.Parallel(_BT, _BT): - if (i_t_local * _BT + _i) < T_seq: - dA[i_b, t_s + _i, i_h, _j] = s_dA[_i, _j] - - # ========== 6. Store db ========== - for _i in T.Parallel(_BT): - if (i_t_local * _BT + _i) < T_seq: - db[i_b, t_s + _i, i_h] = b_db[_i] - - # ========== Kernel entry points ========== - if _VAR: - @T.prim_func - def kernel( - q: T.Tensor(qk_s, _dtype), k: T.Tensor(qk_s, _dtype), - v: T.Tensor(v_s, _dtype), v_new: T.Tensor(v_s, _dtype), - g: T.Tensor(g_s, T.float32), beta: T.Tensor(beta_s, T.float32), - A: T.Tensor(A_s, _dtype), - h: T.Tensor(h_s, _dtype), do: T.Tensor(v_s, _dtype), - dh: T.Tensor(h_s, _dtype), - dq: T.Tensor(qk_s, T.float32), dk: T.Tensor(qk_s, T.float32), - dv: T.Tensor(v_s, _dtype), dv2: T.Tensor(v_s, _dtype), - dg: T.Tensor(g_s, T.float32), db: T.Tensor(db_s, T.float32), - dA: T.Tensor(dA_s, T.float32), - cu_seqlens: T.Tensor((N_seqs_d,), T.int32), - chunk_indices: T.Tensor((NT_d, 2), T.int32), - scale: T.float32, - ): - with T.Kernel(NT_d, _H, threads=_threads) as (i_t, i_h): - i_n = chunk_indices[i_t, 0] - i_t_local = chunk_indices[i_t, 1] - bos = cu_seqlens[i_n] - T_seq = cu_seqlens[i_n + 1] - bos - h_idx = i_t * _H + i_h - t_s = bos + i_t_local * _BT - kernel_body(q, k, v, v_new, g, beta, A, h, do, dh, - dq, dk, dv, dv2, dg, db, dA, - scale, 0, i_h, t_s, T_seq, i_t_local, h_idx) - else: - @T.prim_func - def kernel( - q: T.Tensor(qk_s, _dtype), k: T.Tensor(qk_s, _dtype), - v: T.Tensor(v_s, _dtype), v_new: T.Tensor(v_s, _dtype), - g: T.Tensor(g_s, T.float32), beta: T.Tensor(beta_s, T.float32), - A: T.Tensor(A_s, _dtype), - h: T.Tensor(h_s, _dtype), do: T.Tensor(v_s, _dtype), - dh: T.Tensor(h_s, _dtype), - dq: T.Tensor(qk_s, T.float32), dk: T.Tensor(qk_s, T.float32), - dv: T.Tensor(v_s, _dtype), dv2: T.Tensor(v_s, _dtype), - dg: T.Tensor(g_s, T.float32), db: T.Tensor(db_s, T.float32), - dA: T.Tensor(dA_s, T.float32), - scale: T.float32, - ): - with T.Kernel(T.ceildiv(T_d, _BT), _B * _H, threads=_threads) as (i_t, i_bh): - i_b = i_bh // _H - i_h = i_bh % _H - NT_local = T.ceildiv(T_d, _BT) - h_idx = (i_b * NT_local + i_t) * _H + i_h - t_s = i_t * _BT - kernel_body(q, k, v, v_new, g, beta, A, h, do, dh, - dq, dk, dv, dv2, dg, db, dA, - scale, i_b, i_h, t_s, T_d, i_t, h_idx) - - return kernel - - -def chunk_kda_bwd_wy_dqkg_fused_tilelang( - q, k, v, v_new, g, beta, A, h, do, dh, dv, - scale=None, cu_seqlens=None, chunk_size=64, - chunk_indices=None, state_v_first=False, -): - B, _, H, K, V = *k.shape, v.shape[-1] - BT = chunk_size - # Ensure float32 inputs match kernel signature - g = g.float() - beta = beta.float() - if chunk_indices is None and cu_seqlens is not None: - chunk_indices = prepare_chunk_indices(cu_seqlens, BT) - IS_VARLEN = cu_seqlens is not None - - CONST_TILING = 64 if check_shared_mem() else 32 - if scale is None: - scale = K ** -0.5 - - # Outputs - dq = torch.empty_like(q, dtype=torch.float) - dk = torch.empty_like(k, dtype=torch.float) - dv2 = torch.empty_like(v) - dg = torch.empty_like(g, dtype=torch.float) - db = torch.empty_like(beta, dtype=torch.float) - dA = torch.empty_like(A, dtype=torch.float) - - h_flat = h.reshape(-1, h.shape[-2], h.shape[-1]) - dh_flat = dh.reshape(-1, dh.shape[-2], dh.shape[-1]) - hD1, hD2 = h_flat.shape[-2], h_flat.shape[-1] - dtype_str = {torch.float16: 'float16', torch.bfloat16: 'bfloat16', torch.float32: 'float32'}[q.dtype] - - # Static tiling config - BK = min(max(triton.next_power_of_2(K), 16), CONST_TILING) - BV = min(max(triton.next_power_of_2(V), 16), CONST_TILING) - num_warps = 8 if K <= 64 else 4 - - kernel = _build_kda_bwd_kernel( - B, H, K, V, BT, BK, BV, - hD1, hD2, dtype_str, - state_v_first, IS_VARLEN, - num_warps=num_warps, - ) - - if IS_VARLEN: - kernel(q, k, v, v_new, g, beta, A, h_flat, do, dh_flat, - dq, dk, dv, dv2, dg, db, dA, - cu_seqlens.int(), chunk_indices.int(), scale) - else: - kernel(q, k, v, v_new, g, beta, A, h_flat, do, dh_flat, - dq, dk, dv, dv2, dg, db, dA, scale) - - dv = dv2 - return dq, dk, dv, db, dg, dA diff --git a/build/torch-xpu/ops/kda/chunk.py b/build/torch-xpu/ops/kda/chunk.py deleted file mode 100644 index b11ef0a8b8becb69e828413ac7ae37c3dd7fedaf..0000000000000000000000000000000000000000 --- a/build/torch-xpu/ops/kda/chunk.py +++ /dev/null @@ -1,440 +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 - -# Related files are modified and supported by the Moonshot AI Team - -import warnings - -import torch - -from ...modules.l2norm import l2norm_bwd, l2norm_fwd -from ...ops.backends import dispatch -from ...ops.common.gate import fused_beta_sigmoid, fused_beta_sigmoid_bwd -from ...ops.cp import FLACPContext -from ...ops.kda.chunk_bwd import chunk_kda_bwd -from ...ops.kda.chunk_fwd import chunk_kda_fwd -from ...ops.utils.index import prepare_chunk_indices -from ...utils import autocast_custom_bwd, autocast_custom_fwd, input_guard - - -class ChunkKDAFunction(torch.autograd.Function): - @staticmethod - @input_guard - @autocast_custom_fwd - def forward( - ctx, - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - g: torch.Tensor, - beta: torch.Tensor, - A_log: torch.Tensor, - dt_bias: torch.Tensor, - scale: float, - initial_state: torch.Tensor, - output_final_state: bool = False, - use_qk_l2norm_in_kernel: bool = False, - use_gate_in_kernel: bool = False, - use_beta_sigmoid_in_kernel: bool = False, - allow_neg_eigval: bool = False, - state_v_first: bool = False, - cu_seqlens: torch.LongTensor | None = None, - cu_seqlens_cpu: torch.LongTensor | None = None, - safe_gate: bool = False, - lower_bound: float | None = None, - chunk_size: int = 64, - disable_recompute: bool = False, - return_intermediate_states: bool = False, - cp_context: FLACPContext | None = None, - ): - # Apply l2norm - q_rstd, k_rstd = None, None - if use_qk_l2norm_in_kernel: - q, q_rstd = l2norm_fwd(q) - k, k_rstd = l2norm_fwd(k) - - beta_raw = beta - if use_beta_sigmoid_in_kernel: - beta = fused_beta_sigmoid(beta_raw, scale=2.0 if allow_neg_eigval else 1.0) - - chunk_indices = None - if cu_seqlens is not None: - chunk_indices = prepare_chunk_indices( - cu_seqlens, - chunk_size, - cu_seqlens_cpu=cu_seqlens_cpu, - ) - - g_input = g - - (o, final_state, g_cumsum, Aqk, Akk, w, u, qg, kg, v_new, h, initial_state) = chunk_kda_fwd( - q=q, - k=k, - v=v, - g=g_input, - beta=beta, - scale=scale, - initial_state=initial_state, - output_final_state=output_final_state, - cu_seqlens=cu_seqlens, - cu_seqlens_cpu=cu_seqlens_cpu, - chunk_indices=chunk_indices, - safe_gate=safe_gate, - lower_bound=lower_bound, - use_gate_in_kernel=use_gate_in_kernel, - A_log=A_log, - dt_bias=dt_bias, - chunk_size=chunk_size, - disable_recompute=disable_recompute, - return_intermediate_states=return_intermediate_states, - cp_context=cp_context, - state_v_first=state_v_first, - ) - - if return_intermediate_states: - assert torch.is_inference_mode_enabled(), "return_intermediate_states is only allowed in inference mode" - assert disable_recompute is False, "return_intermediate_states must be used with disable_recompute=False" - return o.type_as(q), final_state, h - - ctx.save_for_backward( - q, q_rstd, k, k_rstd, v, g_cumsum, g_input, beta_raw, beta, A_log, dt_bias, Aqk, Akk, - w, u, qg, kg, v_new, h, - initial_state, cu_seqlens, chunk_indices - ) - ctx.chunk_size = chunk_size - ctx.safe_gate = safe_gate - ctx.scale = scale - ctx.lower_bound = lower_bound - ctx.use_qk_l2norm_in_kernel = use_qk_l2norm_in_kernel - ctx.use_gate_in_kernel = use_gate_in_kernel - ctx.use_beta_sigmoid_in_kernel = use_beta_sigmoid_in_kernel - ctx.allow_neg_eigval = allow_neg_eigval - ctx.disable_recompute = disable_recompute - ctx.cp_context = cp_context - ctx.state_v_first = state_v_first - return o.type_as(q), final_state - - @staticmethod - @input_guard - @autocast_custom_bwd - def backward( - ctx, - do: torch.Tensor, - dht: torch.Tensor, - ): - (q, q_rstd, k, k_rstd, v, g_cumsum, g_input, beta_raw, beta, A_log, dt_bias, Aqk, Akk, - w, u, qg, kg, v_new, h, - initial_state, cu_seqlens, chunk_indices) = ( - ctx.saved_tensors - ) - - dq, dk, dv, db, dg, dh0, dA, dbias = chunk_kda_bwd( - q=q, - k=k, - v=v, - beta=beta, - Aqk=Aqk, - Akk=Akk, - scale=ctx.scale, - initial_state=initial_state, - do=do, - dht=dht, - g=g_cumsum, - g_org=g_input if ctx.use_gate_in_kernel else None, - state_v_first=ctx.state_v_first, - cu_seqlens=cu_seqlens, - chunk_indices=chunk_indices, - chunk_size=ctx.chunk_size, - safe_gate=ctx.safe_gate, - lower_bound=ctx.lower_bound, - use_gate_in_kernel=ctx.use_gate_in_kernel, - A_log=A_log, - dt_bias=dt_bias, - disable_recompute=ctx.disable_recompute, - cp_context=ctx.cp_context, - w=w, - u=u, - qg=qg, - kg=kg, - v_new=v_new, - h=h, - ) - if ctx.use_qk_l2norm_in_kernel: - dq = l2norm_bwd(q, q_rstd, dq) - dk = l2norm_bwd(k, k_rstd, dk) - if ctx.use_beta_sigmoid_in_kernel: - db = fused_beta_sigmoid_bwd(beta_raw, db, scale=2.0 if ctx.allow_neg_eigval else 1.0) - - return (dq.to(q), dk.to(k), dv.to(v), dg.to(g_input), db.to(beta_raw), dA, dbias, None, dh0, - None, None, None, None, None, None, None, None, None, None, None, None, None, None) - - -@dispatch('kda') -@torch.compiler.disable -def chunk_kda( - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - g: torch.Tensor, - beta: torch.Tensor, - scale: float | None = None, - initial_state: torch.Tensor | None = None, - output_final_state: bool = False, - use_qk_l2norm_in_kernel: bool = False, - use_gate_in_kernel: bool = False, - use_beta_sigmoid_in_kernel: bool = False, - allow_neg_eigval: bool = False, - safe_gate: bool = False, - lower_bound: float | None = None, - disable_recompute: bool = False, - return_intermediate_states: bool = False, - state_v_first: bool = False, - cu_seqlens: torch.LongTensor | None = None, - cu_seqlens_cpu: torch.LongTensor | None = None, - cp_context: FLACPContext = None, - **kwargs, -): - r""" - Args: - q (torch.Tensor): - queries of shape ``[B, T, H, K]``. - k (torch.Tensor): - keys of shape ``[B, T, H, K]``. - v (torch.Tensor): - values of shape ``[B, T, HV, V]``. - GVA (Grouped Value Attention) is applied if ``HV > H``, where ``HV`` must be divisible by ``H``. - g (torch.Tensor): - (forget) gating tensor (in log space!) of shape ``[B, T, HV, K]``. - When ``use_gate_in_kernel=False`` (default), ``g`` should be the pre-computed decay value. - When ``use_gate_in_kernel=True``, ``g`` is the raw input before gate activation; - the kernel fuses ``-exp(A_log) * softplus(g + dt_bias)`` + chunk cumsum internally. - beta (torch.Tensor): - betas of shape ``[B, T, HV]``. - scale (Optional[float]): - Scale factor for the KDA attention scores. - If not provided, it will default to ``1 / sqrt(K)``. Default: ``None``. - initial_state (Optional[torch.Tensor]): - Initial state of shape ``[N, HV, K, V]`` for ``N`` input sequences. - For equal-length input sequences, ``N`` equals the batch size ``B``. - Default: ``None``. - output_final_state (Optional[bool]): - Whether to output the final state of shape ``[N, HV, K, V]``. Default: ``False``. - use_qk_l2norm_in_kernel (bool): - Whether to apply L2norm to the q,k tensor internally. Default: ``False``. - use_gate_in_kernel (bool): - Whether to compute the log-space KDA decay internally. - - If ``True``: - The passed ``g`` acts as the raw input for ``-exp(A_log) * softplus(g + dt_bias.view(HV, K))``. - Note that as part of the input arguments, - ``A_log`` (shape ``[HV]``) and the optional ``dt_bias`` (shape ``[HV * K]``) should be provided. - - If ``False``, ``g`` is expected to be the pre-computed decay value. - Default: ``False``. - use_beta_sigmoid_in_kernel (bool): - Whether to apply ``torch.sigmoid(beta)`` before launching the chunk kernel. - - If ``True``, the passed ``beta`` acts as the raw beta logits. - - If ``False``, ``beta`` is expected to already be in post-sigmoid space. - Default: ``False``. - allow_neg_eigval (bool): - Whether to allow negative eigenvalues by scaling ``beta`` to ``[0, 2)``. - Only takes effect together with ``use_beta_sigmoid_in_kernel=True``, in which case - the kernel computes ``2 * sigmoid(beta)`` instead of ``sigmoid(beta)``. - Default: ``False``. - safe_gate (bool): - Whether to clamp the gate to ``[lower_bound, 0)`` and enable M=16 TensorCore - acceleration for higher throughput. Requires ``lower_bound`` to be set. - Default: ``False``. - lower_bound (Optional[float]): - Lower bound for the forget gate (in log space). When set together with - ``safe_gate=True``, changes the gate activation from - ``-exp(A_log) * softplus(g + dt_bias)`` to - ``lower_bound * sigmoid(exp(A_log) * (g + dt_bias))``, - which naturally clamps the output to ``[lower_bound, 0)``. - Recommended value: ``-5`` (i.e., per-step decay ``exp(-5) ≈ 0.0067``). - Default: ``None``. - disable_recompute (bool): - Whether to disable gradient recomputation in the kernel. When ``True``, the kernel - will save all intermediate activations for backward pass, which is beneficial - for training small models at the cost of increased memory usage. Default: ``False``. - return_intermediate_states (bool): - If True, returns intermediate state ``h`` for inference scenarios (e.g., vLLM). - Must be used within ``torch.inference_mode()`` and will return a 3-tuple instead of 2-tuple. - This is not intended for training as it bypasses autograd. Default: ``False``. - state_v_first (Optional[bool]): - Store the recurrent state in V-first ``[V, K]`` layout instead of the default ``[K, V]``. Default: ``False``. - cu_seqlens (torch.LongTensor): - Cumulative sequence lengths of shape ``[N+1]`` used for variable-length training, - consistent with the FlashAttention API. - cu_seqlens_cpu (torch.LongTensor): - Cumulative sequence lengths of shape ``[N+1]`` used for variable-length training, - consistent with the FlashAttention API. - cp_context (Optional[FLACPContext]): - Context parallel context for distributed training across multiple devices. - When provided, ``initial_state`` and ``output_final_state`` are not supported, - and ``cu_seqlens`` will be overridden by the context. Default: ``None``. - - Returns: - - Normal mode (return_intermediate_states=False): A tuple (o, final_state) - o (torch.Tensor): - Outputs of shape ``[B, T, HV, V]``. - final_state (torch.Tensor): - Final state of shape ``[N, HV, K, V]`` if ``output_final_state=True`` else ``None``. - - Inference mode (return_intermediate_states=True): A tuple (o, final_state, h) - o (torch.Tensor): - Outputs of shape ``[B, T, HV, V]``. - final_state (torch.Tensor): - Final state of shape ``[N, HV, K, V]`` if ``output_final_state=True`` else ``None``. - h (torch.Tensor): - Intermediate states of shape ``[B, NT, HV, K, V]`` and dtype ``bfloat16``. - - For equal-length sequences: ``NT = ceil(T / chunk_size)`` - - For variable-length sequences (cu_seqlens): B is always 1 (flattened), - NT is the total number of chunks across all sequences. - - Examples:: - >>> import torch - >>> import torch.nn.functional as F - >>> from einops import rearrange - >>> from ...ops.kda import chunk_kda - # inputs with equal lengths (no GVA, HV == H) - >>> B, T, H, K, V = 4, 2048, 4, 512, 512 - >>> q = torch.randn(B, T, H, K, dtype=torch.bfloat16, device='cuda') - >>> k = torch.randn(B, T, H, K, dtype=torch.bfloat16, device='cuda') - >>> v = torch.randn(B, T, H, V, dtype=torch.bfloat16, device='cuda') - >>> beta = torch.rand(B, T, H, dtype=torch.bfloat16, device='cuda') - >>> g = torch.rand(B, T, H, K, dtype=torch.bfloat16, device='cuda') - >>> h0 = torch.randn(B, H, K, V, dtype=torch.bfloat16, device='cuda') - >>> A_log = torch.randn(H, dtype=torch.float32, device='cuda') - >>> dt_bias = torch.randn(H * K, dtype=torch.float32, device='cuda') - >>> o, ht = chunk_kda( - q, k, v, g, beta, - A_log=A_log, - dt_bias=dt_bias, - use_qk_l2norm_in_kernel=True, - use_gate_in_kernel=True, - initial_state=h0, - output_final_state=True - ) - # GVA mode (HV > H) - >>> HV = 8 # 2x more value heads than qk heads - >>> v = torch.randn(B, T, HV, V, dtype=torch.bfloat16, device='cuda') - >>> g = torch.rand(B, T, HV, K, dtype=torch.bfloat16, device='cuda') - >>> beta = torch.rand(B, T, HV, dtype=torch.bfloat16, device='cuda') - >>> h0 = torch.randn(B, HV, K, V, dtype=torch.bfloat16, device='cuda') - >>> A_log = torch.randn(HV, dtype=torch.float32, device='cuda') - >>> dt_bias = torch.randn(HV * K, dtype=torch.float32, device='cuda') - >>> o, ht = chunk_kda( - q, k, v, g, beta, - A_log=A_log, - dt_bias=dt_bias, - use_qk_l2norm_in_kernel=True, - use_gate_in_kernel=True, - initial_state=h0, - output_final_state=True - ) - # for variable-length inputs, the batch size `B` is expected to be 1 and `cu_seqlens` is required - >>> q, k, v, beta, g = map(lambda x: rearrange(x, 'b t ... -> 1 (b t) ...'), (q, k, v, beta, g)) - # for a batch with 4 sequences, `cu_seqlens` with 5 start/end positions are expected - >>> cu_seqlens = q.new_tensor([0, 2048, 4096, 6144, 8192], dtype=torch.long) - >>> o, ht = chunk_kda( - q, k, v, g, beta, - A_log=A_log, - dt_bias=dt_bias, - use_qk_l2norm_in_kernel=True, - use_gate_in_kernel=True, - initial_state=h0, - output_final_state=True, - cu_seqlens=cu_seqlens - ) - """ - if 'transpose_state_layout' in kwargs: - if state_v_first: - raise ValueError("Cannot pass both `state_v_first` and the deprecated `transpose_state_layout`.") - warnings.warn( - "`transpose_state_layout` is deprecated and renamed to `state_v_first`.", - DeprecationWarning, - stacklevel=2, - ) - state_v_first = kwargs.pop('transpose_state_layout') - - 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" - assert cp_context.cu_seqlens is not None, "cu_seqlens is required for CP" - # Override cu_seqlens and cu_seqlens_cpu with the ones from the context - cu_seqlens = cp_context.cu_seqlens - if cp_context.cu_seqlens_cpu is not None: - cu_seqlens_cpu = cp_context.cu_seqlens_cpu - - if cu_seqlens is not None: - if q.shape[0] != 1: - raise ValueError( - f"The batch size is expected to be 1 rather than {q.shape[0]} when using `cu_seqlens`." - f"Please flatten variable-length inputs before processing.", - ) - if initial_state is not None and initial_state.shape[0] != len(cu_seqlens) - 1: - raise ValueError( - f"The number of initial states is expected to be equal to the number of input sequences, " - f"i.e., {len(cu_seqlens) - 1} rather than {initial_state.shape[0]}.", - ) - if initial_state is not None: - assert initial_state.dtype == torch.float32, "initial_state must be in float32." - - A_log, dt_bias = None, None - if use_gate_in_kernel: - assert "A_log" in kwargs, "A_log must be provided when use_gate_in_kernel=True." - A_log, dt_bias = kwargs["A_log"], kwargs.get("dt_bias") - - chunk_size = kwargs.pop("chunk_size", 64) - if chunk_size not in (32, 64): - raise ValueError(f"`chunk_size` must be either 32 or 64 for KDA, got {chunk_size}.") - - if safe_gate and use_gate_in_kernel: - if lower_bound is None: - raise ValueError("`lower_bound` must be specified when `safe_gate=True` and `use_gate_in_kernel=True`.") - if not (-5 <= lower_bound < 0): - raise ValueError(f"`lower_bound` must be in the safe range [-5, 0), got {lower_bound}.") - - if allow_neg_eigval and not use_beta_sigmoid_in_kernel: - raise ValueError("`allow_neg_eigval=True` requires `use_beta_sigmoid_in_kernel=True`.") - - # Validate head dimensions for GVA - B, T, H, K, HV = *q.shape, v.shape[2] - assert q.shape == k.shape, f"q and k must have the same shape, got q={q.shape} vs k={k.shape}" - assert K <= 256, f"Currently we only support key headdim <=256 for KDA, got {K}." - assert HV % H == 0, ( - f"For GVA, num_v_heads (HV={HV}) must be evenly divisible by num_qk_heads (H={H}), " - f"but got HV % H = {HV % H}" - ) - assert g.shape == (B, T, HV, K), f"g must have shape [B, T, HV, K]={[B, T, HV, K]}, got {list(g.shape)}" - assert beta.shape == (B, T, HV), f"beta must have shape [B, T, HV]={[B, T, HV]}, got {list(beta.shape)}" - - if scale is None: - scale = K ** -0.5 - return ChunkKDAFunction.apply( - q, - k, - v, - g, - beta, - A_log, - dt_bias, - scale, - initial_state, - output_final_state, - use_qk_l2norm_in_kernel, - use_gate_in_kernel, - use_beta_sigmoid_in_kernel, - allow_neg_eigval, - state_v_first, - cu_seqlens, - cu_seqlens_cpu, - safe_gate, - lower_bound, - chunk_size, - disable_recompute, - return_intermediate_states, - cp_context, - ) diff --git a/build/torch-xpu/ops/kda/chunk_bwd.py b/build/torch-xpu/ops/kda/chunk_bwd.py deleted file mode 100644 index 13dba60d48cba5123cf72c82b2ded9930666c523..0000000000000000000000000000000000000000 --- a/build/torch-xpu/ops/kda/chunk_bwd.py +++ /dev/null @@ -1,600 +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 -import triton.language as tl - -from ...ops.backends import dispatch -from ...ops.common.chunk_delta_h import chunk_gated_delta_rule_bwd_dhu, chunk_gated_delta_rule_fwd_h -from ...ops.cp import FLACPContext -from ...ops.cp.chunk_delta_h import chunk_gated_delta_rule_bwd_dhu_pre_process, expand_h0 -from ...ops.kda.chunk_intra import chunk_kda_bwd_intra -from ...ops.kda.gate import kda_gate_bwd, kda_gate_chunk_cumsum -from ...ops.kda.wy_fast import recompute_w_u_fwd -from ...ops.utils import chunk_local_cumsum, prepare_chunk_indices -from ...ops.utils.cache import fla_cache_autotune -from ...ops.utils.constant import RCP_LN2 -from ...ops.utils.op import exp2 -from ...utils import IS_NVIDIA_HOPPER, autotune_cache_kwargs, check_shared_mem - -BK_LIST = [32, 64] if check_shared_mem() else [16, 32] -BV_LIST = [64, 128] if check_shared_mem('ampere') else [16, 32] -NUM_WARPS = [2, 4] if IS_NVIDIA_HOPPER else [2, 4, 8] - - -@triton.heuristics({ - 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None, -}) -@fla_cache_autotune( - configs=[ - triton.Config({}, num_warps=num_warps, num_stages=num_stages) - for num_warps in NUM_WARPS - for num_stages in [2, 3, 4] - ], - key=['H', 'HV', 'K', 'V', 'BT', 'BK', 'BV'], - **autotune_cache_kwargs, -) -@triton.jit(do_not_specialize=['T']) -def chunk_kda_bwd_kernel_dAv( - q, - k, - v, - A, - do, - dv, - dA, - cu_seqlens, - chunk_indices, - scale, - T, - H: tl.constexpr, - HV: tl.constexpr, - K: tl.constexpr, - V: tl.constexpr, - BT: tl.constexpr, - BK: tl.constexpr, - BV: tl.constexpr, - IS_VARLEN: tl.constexpr, -): - i_t, i_bh = tl.program_id(0), tl.program_id(1) - i_b, i_hv = i_bh // HV, i_bh % HV - i_h = i_hv // (HV // H) - 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.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) - T = eos - bos - else: - bos, eos = i_b * T, i_b * T + T - - # offset calculation - q += (bos * H + i_h) * K - k += (bos * H + i_h) * K - v += (bos * HV + i_hv) * V - do += (bos * HV + i_hv) * V - dv += (bos * HV + i_hv) * V - dA += (bos * HV + i_hv) * BT - - p_A = tl.make_block_ptr(A + (bos * HV + i_hv) * BT, (BT, T), (1, HV*BT), (0, i_t * BT), (BT, BT), (0, 1)) - b_A = tl.load(p_A, boundary_check=(0, 1)) - - o_t = i_t * BT + tl.arange(0, BT) - m_t = o_t < T - m_A = (o_t[:, None] <= o_t[None, :]) & (m_t[:, None] & m_t) - b_A = tl.where(m_A, b_A, 0).to(do.dtype.element_ty) - - b_dA = tl.zeros([BT, BT], dtype=tl.float32) - for i_v in range(tl.cdiv(V, BV)): - p_v = tl.make_block_ptr(v, (V, T), (1, HV*V), (i_v * BV, i_t * BT), (BV, BT), (0, 1)) - p_do = tl.make_block_ptr(do, (T, V), (HV*V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) - p_dv = tl.make_block_ptr(dv, (T, V), (HV*V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) - # [BV, BT] - b_v = tl.load(p_v, boundary_check=(0, 1)) - # [BT, BV] - b_do = tl.load(p_do, boundary_check=(0, 1)) - # [BT, BT] - b_dA += tl.dot(b_do, b_v) - # [BT, BV] - b_dv = tl.dot(b_A.to(b_do.dtype), b_do) - tl.store(p_dv, b_dv.to(p_dv.dtype.element_ty), boundary_check=(0, 1)) - - p_dA = tl.make_block_ptr(dA, (T, BT), (HV*BT, 1), (i_t * BT, 0), (BT, BT), (1, 0)) - b_dA = tl.where(o_t[:, None] >= o_t, b_dA * scale, 0.) - tl.store(p_dA, b_dA.to(p_dA.dtype.element_ty), boundary_check=(0, 1)) - - -@triton.heuristics({ - 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None, -}) -@fla_cache_autotune( - configs=[ - triton.Config({'BK': BK, 'BV': BV}, num_warps=num_warps, num_stages=num_stages) - for BK in BK_LIST - for BV in BV_LIST - for num_warps in NUM_WARPS - for num_stages in [2, 3, 4] - if not (IS_NVIDIA_HOPPER and BK == 32 and num_warps == 4) - ], - key=['BT', 'HV', 'STATE_V_FIRST'], - **autotune_cache_kwargs, -) -@triton.jit(do_not_specialize=['T']) -def chunk_kda_bwd_kernel_wy_dqkg_fused( - q, - k, - v, - v_new, - g, - beta, - A, - h, - do, - dh, - dq, - dk, - dv, - dv2, - dg, - db, - dA, - cu_seqlens, - chunk_indices, - scale, - T, - H: tl.constexpr, - HV: tl.constexpr, - K: tl.constexpr, - V: tl.constexpr, - BT: tl.constexpr, - BK: tl.constexpr, - BV: tl.constexpr, - STATE_V_FIRST: tl.constexpr, - IS_VARLEN: tl.constexpr, -): - i_t, i_bh = tl.program_id(0), tl.program_id(1) - i_b, i_hv = i_bh // HV, i_bh % HV - i_h = i_hv // (HV // H) - - if IS_VARLEN: - i_tg = i_t.to(tl.int64) - 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).to(tl.int32) - NT = tl.cdiv(T, BT) - else: - NT = tl.cdiv(T, BT) - i_tg = (i_b * NT + i_t).to(tl.int64) - bos, eos = (i_b * T).to(tl.int64), (i_b * T + T).to(tl.int64) - - o_t = i_t * BT + tl.arange(0, BT) - m_t = o_t < T - m_last = (o_t == min(T, i_t * BT + BT) - 1) - - q += (bos * H + i_h) * K - k += (bos * H + i_h) * K - v += (bos * HV + i_hv) * V - v_new += (bos * HV + i_hv) * V - g += (bos * HV + i_hv) * K - beta += bos * HV + i_hv - A += (bos * HV + i_hv) * BT - h += (i_tg * HV + i_hv) * K*V - do += (bos * HV + i_hv) * V - dh += (i_tg * HV + i_hv) * K*V - dq += (bos * HV + i_hv) * K - dk += (bos * HV + i_hv) * K - dv += (bos * HV + i_hv) * V - dv2 += (bos * HV + i_hv) * V - dg += (bos * HV + i_hv) * K - db += bos * HV + i_hv - dA += (bos * HV + i_hv) * BT - - p_beta = tl.make_block_ptr(beta, (T,), (HV,), (i_t * BT,), (BT,), (0,)) - b_beta = tl.load(p_beta, boundary_check=(0,)) - - p_A = tl.make_block_ptr(A, (BT, T), (1, HV * BT), (0, i_t * BT), (BT, BT), (0, 1)) - b_A = tl.load(p_A, boundary_check=(0, 1)) - - b_dA = tl.zeros([BT, BT], dtype=tl.float32) - b_db = tl.zeros([BT], dtype=tl.float32) - - for i_k in range(tl.cdiv(K, BK)): - o_k = i_k * BK + tl.arange(0, BK) - m_k = o_k < K - - p_k = tl.make_block_ptr(k, (T, K), (H*K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) - p_g = tl.make_block_ptr(g, (T, K), (HV*K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) - b_k = tl.load(p_k, boundary_check=(0, 1)) - b_g = tl.load(p_g, boundary_check=(0, 1)).to(tl.float32) - - p_gn = g + (min(T, i_t * BT + BT) - 1).to(tl.int64) * HV*K + o_k - b_gn = tl.load(p_gn, mask=m_k, other=0).to(tl.float32) - - b_dq = tl.zeros([BT, BK], dtype=tl.float32) - b_dk = tl.zeros([BT, BK], dtype=tl.float32) - b_dw = tl.zeros([BT, BK], dtype=tl.float32) - b_dgk = tl.zeros([BK], dtype=tl.float32) - - for i_v in range(tl.cdiv(V, BV)): - p_v_new = tl.make_block_ptr(v_new, (T, V), (HV*V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) - p_do = tl.make_block_ptr(do, (T, V), (HV*V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) - if STATE_V_FIRST: - p_h = tl.make_block_ptr(h, (V, K), (K, 1), (i_v * BV, i_k * BK), (BV, BK), (1, 0)) - p_dh = tl.make_block_ptr(dh, (V, K), (K, 1), (i_v * BV, i_k * BK), (BV, BK), (1, 0)) - else: - p_h = tl.make_block_ptr(h, (V, K), (1, V), (i_v * BV, i_k * BK), (BV, BK), (0, 1)) - p_dh = tl.make_block_ptr(dh, (V, K), (1, V), (i_v * BV, i_k * BK), (BV, BK), (0, 1)) - p_dv = tl.make_block_ptr(dv, (T, V), (HV*V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) - # [BT, BV] - b_v_new = tl.load(p_v_new, boundary_check=(0, 1)) - b_do = tl.load(p_do, boundary_check=(0, 1)) - # [BV, BK] - b_h = tl.load(p_h, boundary_check=(0, 1)) - b_dh = tl.load(p_dh, boundary_check=(0, 1)) - # [BT, BV] - b_dv = tl.load(p_dv, boundary_check=(0, 1)) - - b_dgk += tl.sum(b_h * b_dh, axis=0) - b_dq += tl.dot(b_do, b_h.to(b_do.dtype)) - b_dk += tl.dot(b_v_new, b_dh.to(b_v_new.dtype)) - b_dw += tl.dot(b_dv.to(b_v_new.dtype), b_h.to(b_v_new.dtype)) - tl.debug_barrier() # DO NOT REMOVE THIS LINE! - if i_k == 0: - p_v = tl.make_block_ptr(v, (T, V), (HV*V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) - p_dv2 = tl.make_block_ptr(dv2, (T, V), (HV*V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) - - b_v = tl.load(p_v, boundary_check=(0, 1)) - - b_dA += tl.dot(b_dv, tl.trans(b_v)) - - b_dvb = tl.dot(b_A, b_dv) - b_dv2 = b_dvb * b_beta[:, None] - b_db += tl.sum(b_dvb * b_v, 1) - - tl.store(p_dv2, b_dv2.to(p_dv2.dtype.element_ty), boundary_check=(0, 1)) - - b_gk_exp = exp2(b_g) - b_gb = b_gk_exp * b_beta[:, None] - b_dgk *= exp2(b_gn) - b_dq = b_dq * b_gk_exp * scale - b_dk = b_dk * tl.where(m_t[:, None], exp2(b_gn[None, :] - b_g), 0) - - b_kg = b_k * b_gk_exp - - b_dw = -b_dw.to(b_A.dtype) - b_dA += tl.dot(b_dw, tl.trans(b_kg.to(b_A.dtype))) - - b_dkgb = tl.dot(b_A, b_dw) - b_db += tl.sum(b_dkgb * b_kg, 1) - - p_q = tl.make_block_ptr(q, (T, K), (H*K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) - b_q = tl.load(p_q, boundary_check=(0, 1)) - b_kdk = b_k * b_dk - b_dgk += tl.sum(b_kdk, axis=0) - b_dg = b_q * b_dq - b_kdk + m_last[:, None] * b_dgk + b_kg * b_dkgb * b_beta[:, None] - b_dk = b_dk + b_dkgb * b_gb - - p_dq = tl.make_block_ptr(dq, (T, K), (HV*K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) - p_dk = tl.make_block_ptr(dk, (T, K), (HV*K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) - p_dg = tl.make_block_ptr(dg, (T, K), (HV*K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) - tl.store(p_dq, b_dq.to(p_dq.dtype.element_ty), boundary_check=(0, 1)) - tl.store(p_dk, b_dk.to(p_dk.dtype.element_ty), boundary_check=(0, 1)) - tl.store(p_dg, b_dg.to(p_dg.dtype.element_ty), boundary_check=(0, 1)) - - m_A = (o_t[:, None] > o_t[None, :]) & (m_t[:, None] & m_t) - b_dA = tl.where(m_A, b_dA * b_beta[None, :], 0) - b_dA = tl.dot(b_dA.to(b_A.dtype), b_A) - b_dA = tl.dot(b_A, b_dA.to(b_A.dtype)) - b_dA = tl.where(m_A, -b_dA, 0) - - p_dA = tl.make_block_ptr(dA, (T, BT), (HV * BT, 1), (i_t * BT, 0), (BT, BT), (1, 0)) - p_db = tl.make_block_ptr(db, (T,), (HV,), (i_t * BT,), (BT,), (0,)) - tl.store(p_dA, b_dA.to(p_dA.dtype.element_ty), boundary_check=(0, 1)) - tl.store(p_db, b_db.to(p_db.dtype.element_ty), boundary_check=(0,)) - - -def chunk_kda_bwd_dAv( - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - do: torch.Tensor, - A: torch.Tensor | None = None, - scale: float = None, - cu_seqlens: torch.LongTensor | None = None, - chunk_size: int = 64, - chunk_indices: torch.LongTensor | None = None, -) -> tuple[torch.Tensor, torch.Tensor]: - B, T, H, K, HV, V = *k.shape, do.shape[2], do.shape[-1] - BT = chunk_size - if chunk_indices is None and cu_seqlens is not None: - chunk_indices = prepare_chunk_indices(cu_seqlens, BT) - # H100 can have larger block size - if check_shared_mem('hopper', k.device.index): - CONST_TILING = 128 - elif check_shared_mem: - CONST_TILING = 64 - else: - CONST_TILING = 32 - BK = min(max(triton.next_power_of_2(K), 16), CONST_TILING) - BV = min(max(triton.next_power_of_2(V), 16), CONST_TILING) - NT = triton.cdiv(T, BT) if cu_seqlens is None else len(chunk_indices) - - dA = v.new_empty(B, T, HV, BT, dtype=torch.float) - dv = torch.empty_like(do) - grid = (NT, B * HV) - chunk_kda_bwd_kernel_dAv[grid]( - q=q, - k=k, - v=v, - A=A, - do=do, - dv=dv, - dA=dA, - cu_seqlens=cu_seqlens, - chunk_indices=chunk_indices, - scale=scale, - T=T, - H=H, - HV=HV, - K=K, - V=V, - BT=BT, - BK=BK, - BV=BV, - ) - return dA, dv - - -@dispatch('kda') -def chunk_kda_bwd_wy_dqkg_fused( - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - v_new: torch.Tensor, - g: torch.Tensor, - beta: torch.Tensor, - A: torch.Tensor, - h: torch.Tensor, - do: torch.Tensor, - dh: torch.Tensor, - dv: torch.Tensor, - scale: float | None = None, - state_v_first: bool = False, - cu_seqlens: torch.LongTensor | None = None, - chunk_size: int = 64, - chunk_indices: torch.LongTensor | None = None, -): - B, T, H, K, HV, V = *k.shape, v.shape[2], v.shape[-1] - BT = chunk_size - - if chunk_indices is None and cu_seqlens is not None: - chunk_indices = prepare_chunk_indices(cu_seqlens, BT) - NT = triton.cdiv(T, BT) if cu_seqlens is None else len(chunk_indices) - - # dq, dk are allocated at HV dimension; caller reduces to H if GVA - dq = g.new_empty(B, T, HV, K, dtype=torch.float) - dk = g.new_empty(B, T, HV, K, dtype=torch.float) - dv2 = torch.empty_like(v) - dg = torch.empty_like(g, dtype=torch.float) - db = torch.empty_like(beta, dtype=torch.float) - dA = torch.empty_like(A, dtype=torch.float) - - grid = (NT, B * HV) - chunk_kda_bwd_kernel_wy_dqkg_fused[grid]( - q=q, - k=k, - v=v, - v_new=v_new, - g=g, - beta=beta, - A=A, - h=h, - do=do, - dh=dh, - dq=dq, - dk=dk, - dv=dv, - dv2=dv2, - dg=dg, - db=db, - dA=dA, - cu_seqlens=cu_seqlens, - chunk_indices=chunk_indices, - scale=scale, - T=T, - H=H, - HV=HV, - K=K, - V=V, - BT=BT, - STATE_V_FIRST=state_v_first, - ) - dv = dv2 - return dq, dk, dv, db, dg, dA - - -def chunk_kda_bwd( - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - beta: torch.Tensor, - Aqk: torch.Tensor, - Akk: torch.Tensor, - scale: float, - initial_state: torch.Tensor, - do: torch.Tensor, - dht: torch.Tensor, - g: torch.Tensor | None = None, - g_org: torch.Tensor | None = None, - state_v_first: bool = False, - cu_seqlens: torch.LongTensor | None = None, - chunk_indices: torch.LongTensor | None = None, - chunk_size: int = 64, - safe_gate: bool = False, - lower_bound: float | None = None, - use_gate_in_kernel: bool = False, - A_log: torch.Tensor | None = None, - dt_bias: torch.Tensor | None = None, - disable_recompute: bool = False, - cp_context: FLACPContext | None = None, - **kwargs, -): - H, HV = q.shape[2], v.shape[2] - G = HV // H - - if disable_recompute is False: - if use_gate_in_kernel: - g = kda_gate_chunk_cumsum( - g=g_org, - A_log=A_log, - dt_bias=dt_bias, - scale=RCP_LN2, - chunk_size=chunk_size, - cu_seqlens=cu_seqlens, - chunk_indices=chunk_indices, - lower_bound=lower_bound - ) - w, u, qg, kg = recompute_w_u_fwd( - q=q, - k=k, - v=v, - beta=beta, - A=Akk, - gk=g, - cu_seqlens=cu_seqlens, - chunk_indices=chunk_indices, - ) - if cp_context is not None: - # Restore the full initial_state tensor from the compressed version. - # Only the first sequence's state is non-zero as it's the only one that could be cross-rank. - initial_state = expand_h0(initial_state, context=cp_context) - h, v_new, _ = chunk_gated_delta_rule_fwd_h( - k=kg, - w=w, - u=u, - gk=g, - initial_state=initial_state, - output_final_state=False, - cu_seqlens=cu_seqlens, - chunk_indices=chunk_indices, - chunk_size=chunk_size, - state_v_first=state_v_first, - ) - else: - w, u, qg, kg, v_new, h = kwargs["w"], kwargs["u"], kwargs["qg"], kwargs["kg"], kwargs["v_new"], kwargs["h"] - if cp_context is not None: - # Restore the full initial_state tensor from the compressed version. - # Only the first sequence's state is non-zero as it's the only one that could be cross-rank. - initial_state = expand_h0(initial_state, context=cp_context) - - # dAqk = do @ v.T - # dv = A @ do - dAqk, dv = chunk_kda_bwd_dAv( - q=q, - k=k, - v=v_new, - do=do, - A=Aqk, - scale=scale, - cu_seqlens=cu_seqlens, - chunk_size=chunk_size, - chunk_indices=chunk_indices, - ) - - if cp_context is not None: - # initial_state is None in the CP mode - # We only need to compute dht of current rank and pass it to the backward kernel - dht, initial_state = chunk_gated_delta_rule_bwd_dhu_pre_process( - q=qg, - k=kg, - w=w, - do=do, - dv=dv, - gk=g, - scale=scale, - cu_seqlens=cu_seqlens, - dht=dht, - initial_state=initial_state, - context=cp_context, - chunk_size=chunk_size, - state_v_first=state_v_first, - ) - - dh, dh0, dv = chunk_gated_delta_rule_bwd_dhu( - q=qg, - k=kg, - w=w, - gk=g, - h0=initial_state, - dht=dht, - do=do, - dv=dv, - scale=scale, - cu_seqlens=cu_seqlens, - chunk_size=chunk_size, - chunk_indices=chunk_indices, - state_v_first=state_v_first, - ) - - dq, dk, dv, db, dg, dAkk = chunk_kda_bwd_wy_dqkg_fused( - q=q, - k=k, - v=v, - v_new=v_new, - g=g, - beta=beta, - A=Akk, - h=h, - do=do, - dh=dh, - dv=dv, - scale=scale, - cu_seqlens=cu_seqlens, - chunk_size=chunk_size, - chunk_indices=chunk_indices, - state_v_first=state_v_first, - ) - - dq, dk, db, dg = chunk_kda_bwd_intra( - q=q, - k=k, - g=g, - beta=beta, - dAqk=dAqk, - dAkk=dAkk, - dq=dq, - dk=dk, - db=db, - dg=dg, - cu_seqlens=cu_seqlens, - chunk_size=chunk_size, - chunk_indices=chunk_indices, - safe_gate=safe_gate - ) - - # For GVA, reduce dq and dk from [B, T, HV, K] back to [B, T, H, K] - if HV > H: - dq = dq.view(*dq.shape[:2], H, G, dq.shape[-1]).sum(dim=3) - dk = dk.view(*dk.shape[:2], H, G, dk.shape[-1]).sum(dim=3) - - dA, dbias = None, None - dg = chunk_local_cumsum( - dg, - chunk_size=chunk_size, - reverse=True, - cu_seqlens=cu_seqlens, - chunk_indices=chunk_indices, - ) - if use_gate_in_kernel: - dg, dA, dbias = kda_gate_bwd( - g=g_org, - A_log=A_log, - dt_bias=dt_bias, - dyg=dg, - lower_bound=lower_bound - ) - - return dq, dk, dv, db, dg, dh0, dA, dbias diff --git a/build/torch-xpu/ops/kda/chunk_fwd.py b/build/torch-xpu/ops/kda/chunk_fwd.py deleted file mode 100644 index 2dbd234be98339d63f0c7419fef0312d39aac86b..0000000000000000000000000000000000000000 --- a/build/torch-xpu/ops/kda/chunk_fwd.py +++ /dev/null @@ -1,134 +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 - -from ...ops.common.chunk_delta_h import chunk_gated_delta_rule_fwd_h -from ...ops.cp import FLACPContext -from ...ops.cp.chunk_delta_h import chunk_gated_delta_rule_fwd_h_pre_process, compress_h0 -from ...ops.gla.chunk import chunk_gla_fwd_o_gk -from ...ops.kda.chunk_intra import chunk_kda_fwd_intra -from ...ops.kda.gate import kda_gate_chunk_cumsum -from ...ops.utils import chunk_local_cumsum -from ...ops.utils.constant import RCP_LN2 - - -def chunk_kda_fwd( - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - g: torch.Tensor, - beta: torch.Tensor, - scale: float, - initial_state: torch.Tensor, - output_final_state: bool, - state_v_first: bool = False, - cu_seqlens: torch.LongTensor | None = None, - cu_seqlens_cpu: torch.LongTensor | None = None, - chunk_indices: torch.LongTensor | None = None, - chunk_size: int = 64, - safe_gate: bool = False, - lower_bound: float | None = None, - use_gate_in_kernel: bool = False, - A_log: torch.Tensor | None = None, - dt_bias: torch.Tensor | None = None, - disable_recompute: bool = False, - return_intermediate_states: bool = False, - cp_context: FLACPContext | None = None, -): - # Apply gate activation - g_org = None - if use_gate_in_kernel: - g_org = g - g = kda_gate_chunk_cumsum( - g=g_org, - A_log=A_log, - dt_bias=dt_bias, - scale=RCP_LN2, - chunk_size=chunk_size, - cu_seqlens=cu_seqlens, - chunk_indices=chunk_indices, - lower_bound=lower_bound, - ) - else: - g = chunk_local_cumsum( - g=g, - scale=RCP_LN2, - chunk_size=chunk_size, - cu_seqlens=cu_seqlens, - chunk_indices=chunk_indices - ) - - # qg = None if disable_recompute is False - w, u, qg, kg, Aqk, Akk = chunk_kda_fwd_intra( - q=q, - k=k, - v=v, - gk=g, - beta=beta, - scale=scale, - cu_seqlens=cu_seqlens, - chunk_size=chunk_size, - chunk_indices=chunk_indices, - safe_gate=safe_gate, - disable_recompute=disable_recompute - ) - - if cp_context is not None: - initial_state = chunk_gated_delta_rule_fwd_h_pre_process( - k=kg, - w=w, - u=u, - gk=g, - cu_seqlens=cu_seqlens, - initial_state=initial_state, - context=cp_context, - chunk_size=chunk_size, - state_v_first=state_v_first, - ) - - h, v_new, final_state = chunk_gated_delta_rule_fwd_h( - k=kg, - w=w, - u=u, - gk=g, - initial_state=initial_state, - output_final_state=output_final_state, - cu_seqlens=cu_seqlens, - cu_seqlens_cpu=cu_seqlens_cpu, - chunk_indices=chunk_indices, - chunk_size=chunk_size, - state_v_first=state_v_first, - ) - - if cp_context is not None: - # In Context Parallel (CP) mode, global initial states are not supported at the entry point. - # The `initial_state` here is computed internally via inter-rank communication. - # Since only the first sequence in the local batch can be a continuation of a cross-rank sequence, - # only the first state in the tensor is relevant. We compress it to optimize memory for `save_for_backward`. - initial_state = compress_h0(initial_state, context=cp_context) - - o = chunk_gla_fwd_o_gk( - q=q, - v=v_new, - g=g, - A=Aqk, - h=h, - scale=scale, - cu_seqlens=cu_seqlens, - chunk_size=chunk_size, - chunk_indices=chunk_indices, - state_v_first=state_v_first, - ) - if disable_recompute is False: - # Delete to save memory - w, u, qg, kg, v_new = None, None, None, None, None - if not return_intermediate_states: - h = None - if use_gate_in_kernel: - g = None - return o, final_state, g, Aqk, Akk, w, u, qg, kg, v_new, h, initial_state diff --git a/build/torch-xpu/ops/kda/chunk_intra.py b/build/torch-xpu/ops/kda/chunk_intra.py deleted file mode 100644 index 795e5996a4d870cc5b813ac6eb81036a4b270651..0000000000000000000000000000000000000000 --- a/build/torch-xpu/ops/kda/chunk_intra.py +++ /dev/null @@ -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 - -import torch -import triton -import triton.language as tl - -from ...ops.kda.chunk_intra_token_parallel import chunk_kda_fwd_intra_token_parallel -from ...ops.kda.wy_fast import recompute_w_u_fwd -from ...ops.utils import prepare_chunk_indices -from ...ops.utils.cache import fla_cache_autotune -from ...ops.utils.op import exp2, gather -from ...utils import IS_GATHER_SUPPORTED, IS_TF32_SUPPORTED, autotune_cache_kwargs - -if IS_TF32_SUPPORTED: - SOLVE_TRIL_DOT_PRECISION = tl.constexpr('tf32') -else: - SOLVE_TRIL_DOT_PRECISION = tl.constexpr('ieee') - -################################################################################ -# Fused inter + solve_tril kernel: compute off-diagonal Akk and solve in one pass -################################################################################ - - -@triton.heuristics({ - 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None, -}) -@fla_cache_autotune( - configs=[ - triton.Config({'BK': BK}, num_warps=num_warps) - for BK in [32, 64] - for num_warps in [1, 2, 4] - ], - key=["H", "HV", "K", "BT", "BC", "NC"], - **autotune_cache_kwargs, -) -@triton.jit(do_not_specialize=['T']) -def chunk_kda_fwd_kernel_inter_solve_fused( - q, - k, - g, - beta, - Aqk, - Akkd, - Akk, - scale, - cu_seqlens, - chunk_indices, - T, - H: tl.constexpr, - HV: tl.constexpr, - K: tl.constexpr, - BT: tl.constexpr, - BC: tl.constexpr, - NC: tl.constexpr, - BK: tl.constexpr, - IS_VARLEN: tl.constexpr, - USE_SAFE_GATE: tl.constexpr, -): - """ - Fused kernel: compute inter-subchunk Akk + solve_tril in one pass. - Prerequisite: token_parallel has already computed diagonal Akk blocks in Akkd. - - This kernel: - 1. Computes off-diagonal Aqk blocks -> writes to global - 2. Computes off-diagonal Akk blocks -> keeps in registers - 3. Loads diagonal Akk blocks from Akkd (fp32) - 4. Does forward substitution on diagonals - 5. Computes merged Akk_inv - 6. Writes Akk_inv to Akk - """ - i_t, i_bh = tl.program_id(0), tl.program_id(1) - i_b, i_hv = i_bh // HV, i_bh % HV - i_h = i_hv // (HV // H) - - 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.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) - T = eos - bos - else: - bos, eos = i_b * T, i_b * T + T - - if i_t * BT >= T: - return - - i_tc0 = i_t * BT - i_tc1 = i_t * BT + BC - i_tc2 = i_t * BT + 2 * BC - i_tc3 = i_t * BT + 3 * BC - - q += (bos * H + i_h) * K - k += (bos * H + i_h) * K - g += (bos * HV + i_hv) * K - Aqk += (bos * HV + i_hv) * BT - Akk += (bos * HV + i_hv) * BT - Akkd += (bos * HV + i_hv) * BC - - o_i = tl.arange(0, BC) - m_tc1 = (i_tc1 + o_i) < T - m_tc2 = (i_tc2 + o_i) < T - m_tc3 = (i_tc3 + o_i) < T - - b_Aqk10 = tl.zeros([BC, BC], dtype=tl.float32) - b_Akk10 = tl.zeros([BC, BC], dtype=tl.float32) - - b_Aqk20 = tl.zeros([BC, BC], dtype=tl.float32) - b_Akk20 = tl.zeros([BC, BC], dtype=tl.float32) - b_Aqk21 = tl.zeros([BC, BC], dtype=tl.float32) - b_Akk21 = tl.zeros([BC, BC], dtype=tl.float32) - - b_Aqk30 = tl.zeros([BC, BC], dtype=tl.float32) - b_Akk30 = tl.zeros([BC, BC], dtype=tl.float32) - b_Aqk31 = tl.zeros([BC, BC], dtype=tl.float32) - b_Akk31 = tl.zeros([BC, BC], dtype=tl.float32) - b_Aqk32 = tl.zeros([BC, BC], dtype=tl.float32) - b_Akk32 = tl.zeros([BC, BC], dtype=tl.float32) - - ################################################################################ - # off-diagonal blocks - ################################################################################ - for i_k in range(tl.cdiv(K, BK)): - o_k = i_k * BK + tl.arange(0, BK) - m_k = o_k < K - - p_k0 = tl.make_block_ptr(k, (T, K), (H*K, 1), (i_tc0, i_k * BK), (BC, BK), (1, 0)) - p_g0 = tl.make_block_ptr(g, (T, K), (HV*K, 1), (i_tc0, i_k * BK), (BC, BK), (1, 0)) - b_k0 = tl.load(p_k0, boundary_check=(0, 1)).to(tl.float32) - b_g0 = tl.load(p_g0, boundary_check=(0, 1)).to(tl.float32) - - if i_tc1 < T: - p_q1 = tl.make_block_ptr(q, (T, K), (H*K, 1), (i_tc1, i_k * BK), (BC, BK), (1, 0)) - p_k1 = tl.make_block_ptr(k, (T, K), (H*K, 1), (i_tc1, i_k * BK), (BC, BK), (1, 0)) - p_g1 = tl.make_block_ptr(g, (T, K), (HV*K, 1), (i_tc1, i_k * BK), (BC, BK), (1, 0)) - # [BC, BK] - b_q1 = tl.load(p_q1, boundary_check=(0, 1)).to(tl.float32) - b_k1 = tl.load(p_k1, boundary_check=(0, 1)).to(tl.float32) - b_g1 = tl.load(p_g1, boundary_check=(0, 1)).to(tl.float32) - # [BK] - b_gn1 = tl.load(g + i_tc1 * HV*K + o_k, mask=m_k, other=0).to(tl.float32) - # [BC, BK] - b_gqn = tl.where(m_tc1[:, None], exp2(b_g1 - b_gn1[None, :]), 0) - # [BK, BC] - b_kgt = tl.trans(b_k0 * exp2(b_gn1[None, :] - b_g0)) - # [BC, BC] - b_Aqk10 += tl.dot(b_q1 * b_gqn, b_kgt) - b_Akk10 += tl.dot(b_k1 * b_gqn, b_kgt) - - if NC >= 3 and i_tc2 < T: - p_q2 = tl.make_block_ptr(q, (T, K), (H*K, 1), (i_tc2, i_k * BK), (BC, BK), (1, 0)) - p_k2 = tl.make_block_ptr(k, (T, K), (H*K, 1), (i_tc2, i_k * BK), (BC, BK), (1, 0)) - p_g2 = tl.make_block_ptr(g, (T, K), (HV*K, 1), (i_tc2, i_k * BK), (BC, BK), (1, 0)) - # [BC, BK] - b_q2 = tl.load(p_q2, boundary_check=(0, 1)).to(tl.float32) - b_k2 = tl.load(p_k2, boundary_check=(0, 1)).to(tl.float32) - b_g2 = tl.load(p_g2, boundary_check=(0, 1)).to(tl.float32) - # [BK] - b_gn2 = tl.load(g + i_tc2 * HV*K + o_k, mask=m_k, other=0).to(tl.float32) - # [BC, BK] - b_gqn2 = tl.where(m_tc2[:, None], exp2(b_g2 - b_gn2[None, :]), 0) - b_qg2 = b_q2 * b_gqn2 - b_kg2 = b_k2 * b_gqn2 - # [BK, BC] - b_kgt = tl.trans(b_k0 * exp2(b_gn2[None, :] - b_g0)) - b_Aqk20 += tl.dot(b_qg2, b_kgt) - b_Akk20 += tl.dot(b_kg2, b_kgt) - # [BC, BC] - b_kgt = tl.trans(b_k1 * exp2(b_gn2[None, :] - b_g1)) - # [BC, BC] - b_Aqk21 += tl.dot(b_qg2, b_kgt) - b_Akk21 += tl.dot(b_kg2, b_kgt) - - if NC >= 4 and i_tc3 < T: - p_q3 = tl.make_block_ptr(q, (T, K), (H*K, 1), (i_tc3, i_k * BK), (BC, BK), (1, 0)) - p_k3 = tl.make_block_ptr(k, (T, K), (H*K, 1), (i_tc3, i_k * BK), (BC, BK), (1, 0)) - p_g3 = tl.make_block_ptr(g, (T, K), (HV*K, 1), (i_tc3, i_k * BK), (BC, BK), (1, 0)) - # [BC, BK] - b_q3 = tl.load(p_q3, boundary_check=(0, 1)).to(tl.float32) - b_k3 = tl.load(p_k3, boundary_check=(0, 1)).to(tl.float32) - b_g3 = tl.load(p_g3, boundary_check=(0, 1)).to(tl.float32) - # [BK] - b_gn3 = tl.load(g + i_tc3 * HV*K + o_k, mask=m_k, other=0).to(tl.float32) - # [BC, BK] - b_gqn3 = tl.where(m_tc3[:, None], exp2(b_g3 - b_gn3[None, :]), 0) - b_qg3 = b_q3 * b_gqn3 - b_kg3 = b_k3 * b_gqn3 - # [BK, BC] - b_kgt = tl.trans(b_k0 * exp2(b_gn3[None, :] - b_g0)) - # [BC, BC] - b_Aqk30 += tl.dot(b_qg3, b_kgt) - b_Akk30 += tl.dot(b_kg3, b_kgt) - # [BK, BC] - b_kgt = tl.trans(b_k1 * exp2(b_gn3[None, :] - b_g1)) - # [BC, BC] - b_Aqk31 += tl.dot(b_qg3, b_kgt) - b_Akk31 += tl.dot(b_kg3, b_kgt) - # [BK, BC] - b_kgt = tl.trans(b_k2 * exp2(b_gn3[None, :] - b_g2)) - # [BC, BC] - b_Aqk32 += tl.dot(b_qg3, b_kgt) - b_Akk32 += tl.dot(b_kg3, b_kgt) - - ################################################################################ - # save off-diagonal Aqk blocks and prepare Akk - ################################################################################ - if i_tc1 < T: - p_Aqk10 = tl.make_block_ptr(Aqk, (T, BT), (HV*BT, 1), (i_tc1, 0), (BC, BC), (1, 0)) - tl.store(p_Aqk10, (b_Aqk10 * scale).to(Aqk.dtype.element_ty), boundary_check=(0, 1)) - - p_b1 = tl.make_block_ptr(beta + bos * HV + i_hv, (T,), (HV,), (i_tc1,), (BC,), (0,)) - b_b1 = tl.load(p_b1, boundary_check=(0,)).to(tl.float32) - b_Akk10 = b_Akk10 * b_b1[:, None] - if NC >= 3 and i_tc2 < T: - p_Aqk20 = tl.make_block_ptr(Aqk, (T, BT), (HV*BT, 1), (i_tc2, 0), (BC, BC), (1, 0)) - p_Aqk21 = tl.make_block_ptr(Aqk, (T, BT), (HV*BT, 1), (i_tc2, BC), (BC, BC), (1, 0)) - tl.store(p_Aqk20, (b_Aqk20 * scale).to(Aqk.dtype.element_ty), boundary_check=(0, 1)) - tl.store(p_Aqk21, (b_Aqk21 * scale).to(Aqk.dtype.element_ty), boundary_check=(0, 1)) - - p_b2 = tl.make_block_ptr(beta + bos * HV + i_hv, (T,), (HV,), (i_tc2,), (BC,), (0,)) - b_b2 = tl.load(p_b2, boundary_check=(0,)).to(tl.float32) - b_Akk20 = b_Akk20 * b_b2[:, None] - b_Akk21 = b_Akk21 * b_b2[:, None] - if NC >= 4 and i_tc3 < T: - p_Aqk30 = tl.make_block_ptr(Aqk, (T, BT), (HV*BT, 1), (i_tc3, 0), (BC, BC), (1, 0)) - p_Aqk31 = tl.make_block_ptr(Aqk, (T, BT), (HV*BT, 1), (i_tc3, BC), (BC, BC), (1, 0)) - p_Aqk32 = tl.make_block_ptr(Aqk, (T, BT), (HV*BT, 1), (i_tc3, 2*BC), (BC, BC), (1, 0)) - tl.store(p_Aqk30, (b_Aqk30 * scale).to(Aqk.dtype.element_ty), boundary_check=(0, 1)) - tl.store(p_Aqk31, (b_Aqk31 * scale).to(Aqk.dtype.element_ty), boundary_check=(0, 1)) - tl.store(p_Aqk32, (b_Aqk32 * scale).to(Aqk.dtype.element_ty), boundary_check=(0, 1)) - - p_b3 = tl.make_block_ptr(beta + bos * HV + i_hv, (T,), (HV,), (i_tc3,), (BC,), (0,)) - b_b3 = tl.load(p_b3, boundary_check=(0,)).to(tl.float32) - b_Akk30 = b_Akk30 * b_b3[:, None] - b_Akk31 = b_Akk31 * b_b3[:, None] - b_Akk32 = b_Akk32 * b_b3[:, None] - - p_Akk00 = tl.make_block_ptr(Akkd, (T, BC), (HV*BC, 1), (i_tc0, 0), (BC, BC), (1, 0)) - p_Akk11 = tl.make_block_ptr(Akkd, (T, BC), (HV*BC, 1), (i_tc1, 0), (BC, BC), (1, 0)) - b_Ai00 = tl.load(p_Akk00, boundary_check=(0, 1)).to(tl.float32) - b_Ai11 = tl.load(p_Akk11, boundary_check=(0, 1)).to(tl.float32) - if NC >= 3: - p_Akk22 = tl.make_block_ptr(Akkd, (T, BC), (HV*BC, 1), (i_tc2, 0), (BC, BC), (1, 0)) - b_Ai22 = tl.load(p_Akk22, boundary_check=(0, 1)).to(tl.float32) - if NC >= 4: - p_Akk33 = tl.make_block_ptr(Akkd, (T, BC), (HV*BC, 1), (i_tc3, 0), (BC, BC), (1, 0)) - b_Ai33 = tl.load(p_Akk33, boundary_check=(0, 1)).to(tl.float32) - - ################################################################################ - # forward substitution on diagonals - ################################################################################ - - if not USE_SAFE_GATE: - m_A = o_i[:, None] > o_i[None, :] - m_I = o_i[:, None] == o_i[None, :] - - b_Ai00 = -tl.where(m_A, b_Ai00, 0) - b_Ai11 = -tl.where(m_A, b_Ai11, 0) - if NC >= 3: - b_Ai22 = -tl.where(m_A, b_Ai22, 0) - if NC >= 4: - b_Ai33 = -tl.where(m_A, b_Ai33, 0) - - for i in range(2, min(BC, T - i_tc0)): - b_a00 = -tl.load(Akkd + (i_tc0 + i) * HV*BC + o_i) - b_a00 = tl.where(o_i < i, b_a00, 0.) - b_a00 += tl.sum(b_a00[:, None] * b_Ai00, 0) - b_Ai00 = tl.where((o_i == i)[:, None], b_a00, b_Ai00) - for i in range(BC + 2, min(2*BC, T - i_tc0)): - b_a11 = -tl.load(Akkd + (i_tc0 + i) * HV*BC + o_i) - b_a11 = tl.where(o_i < i - BC, b_a11, 0.) - b_a11 += tl.sum(b_a11[:, None] * b_Ai11, 0) - b_Ai11 = tl.where((o_i == i - BC)[:, None], b_a11, b_Ai11) - if NC >= 3: - for i in range(2*BC + 2, min(3*BC, T - i_tc0)): - b_a22 = -tl.load(Akkd + (i_tc0 + i) * HV*BC + o_i) - b_a22 = tl.where(o_i < i - 2*BC, b_a22, 0.) - b_a22 += tl.sum(b_a22[:, None] * b_Ai22, 0) - b_Ai22 = tl.where((o_i == i - 2*BC)[:, None], b_a22, b_Ai22) - if NC >= 4: - for i in range(3*BC + 2, min(4*BC, T - i_tc0)): - b_a33 = -tl.load(Akkd + (i_tc0 + i) * HV*BC + o_i) - b_a33 = tl.where(o_i < i - 3*BC, b_a33, 0.) - b_a33 += tl.sum(b_a33[:, None] * b_Ai33, 0) - b_Ai33 = tl.where((o_i == i - 3*BC)[:, None], b_a33, b_Ai33) - - b_Ai00 += m_I - b_Ai11 += m_I - if NC >= 3: - b_Ai22 += m_I - if NC >= 4: - b_Ai33 += m_I - - ################################################################################ - # compute merged inverse using off-diagonals - ################################################################################ - - # we used tf32 to maintain matrix inverse's precision whenever possible. - b_Ai10 = -tl.dot( - tl.dot(b_Ai11, b_Akk10, input_precision=SOLVE_TRIL_DOT_PRECISION), - b_Ai00, - input_precision=SOLVE_TRIL_DOT_PRECISION - ) - - if NC >= 3: - b_Ai21 = -tl.dot( - tl.dot(b_Ai22, b_Akk21, input_precision=SOLVE_TRIL_DOT_PRECISION), - b_Ai11, - input_precision=SOLVE_TRIL_DOT_PRECISION - ) - b_Ai20 = -tl.dot( - b_Ai22, - tl.dot(b_Akk20, b_Ai00, input_precision=SOLVE_TRIL_DOT_PRECISION) + - tl.dot(b_Akk21, b_Ai10, input_precision=SOLVE_TRIL_DOT_PRECISION), - input_precision=SOLVE_TRIL_DOT_PRECISION - ) - if NC >= 4: - b_Ai32 = -tl.dot( - tl.dot(b_Ai33, b_Akk32, input_precision=SOLVE_TRIL_DOT_PRECISION), - b_Ai22, - input_precision=SOLVE_TRIL_DOT_PRECISION - ) - b_Ai31 = -tl.dot( - b_Ai33, - tl.dot(b_Akk31, b_Ai11, input_precision=SOLVE_TRIL_DOT_PRECISION) + - tl.dot(b_Akk32, b_Ai21, input_precision=SOLVE_TRIL_DOT_PRECISION), - input_precision=SOLVE_TRIL_DOT_PRECISION - ) - b_Ai30 = -tl.dot( - b_Ai33, - tl.dot(b_Akk30, b_Ai00, input_precision=SOLVE_TRIL_DOT_PRECISION) + - tl.dot(b_Akk31, b_Ai10, input_precision=SOLVE_TRIL_DOT_PRECISION) + - tl.dot(b_Akk32, b_Ai20, input_precision=SOLVE_TRIL_DOT_PRECISION), - input_precision=SOLVE_TRIL_DOT_PRECISION - ) - - ################################################################################ - # store full Akk_inv to Akk - ################################################################################ - - p_Akk00 = tl.make_block_ptr(Akk, (T, BT), (HV*BT, 1), (i_tc0, 0), (BC, BC), (1, 0)) - p_Akk10 = tl.make_block_ptr(Akk, (T, BT), (HV*BT, 1), (i_tc1, 0), (BC, BC), (1, 0)) - p_Akk11 = tl.make_block_ptr(Akk, (T, BT), (HV*BT, 1), (i_tc1, BC), (BC, BC), (1, 0)) - - tl.store(p_Akk00, b_Ai00.to(Akk.dtype.element_ty), boundary_check=(0, 1)) - tl.store(p_Akk10, b_Ai10.to(Akk.dtype.element_ty), boundary_check=(0, 1)) - tl.store(p_Akk11, b_Ai11.to(Akk.dtype.element_ty), boundary_check=(0, 1)) - if NC >= 3: - p_Akk20 = tl.make_block_ptr(Akk, (T, BT), (HV*BT, 1), (i_tc2, 0), (BC, BC), (1, 0)) - p_Akk21 = tl.make_block_ptr(Akk, (T, BT), (HV*BT, 1), (i_tc2, BC), (BC, BC), (1, 0)) - p_Akk22 = tl.make_block_ptr(Akk, (T, BT), (HV*BT, 1), (i_tc2, 2*BC), (BC, BC), (1, 0)) - tl.store(p_Akk20, b_Ai20.to(Akk.dtype.element_ty), boundary_check=(0, 1)) - tl.store(p_Akk21, b_Ai21.to(Akk.dtype.element_ty), boundary_check=(0, 1)) - tl.store(p_Akk22, b_Ai22.to(Akk.dtype.element_ty), boundary_check=(0, 1)) - if NC >= 4: - p_Akk30 = tl.make_block_ptr(Akk, (T, BT), (HV*BT, 1), (i_tc3, 0), (BC, BC), (1, 0)) - p_Akk31 = tl.make_block_ptr(Akk, (T, BT), (HV*BT, 1), (i_tc3, BC), (BC, BC), (1, 0)) - p_Akk32 = tl.make_block_ptr(Akk, (T, BT), (HV*BT, 1), (i_tc3, 2*BC), (BC, BC), (1, 0)) - p_Akk33 = tl.make_block_ptr(Akk, (T, BT), (HV*BT, 1), (i_tc3, 3*BC), (BC, BC), (1, 0)) - tl.store(p_Akk30, b_Ai30.to(Akk.dtype.element_ty), boundary_check=(0, 1)) - tl.store(p_Akk31, b_Ai31.to(Akk.dtype.element_ty), boundary_check=(0, 1)) - tl.store(p_Akk32, b_Ai32.to(Akk.dtype.element_ty), boundary_check=(0, 1)) - tl.store(p_Akk33, b_Ai33.to(Akk.dtype.element_ty), boundary_check=(0, 1)) - - -@triton.heuristics({ - 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None, -}) -@fla_cache_autotune( - configs=[ - triton.Config({}, num_warps=num_warps, num_stages=num_stages) - for num_warps in [1, 2, 4, 8] - for num_stages in [2, 3, 4] - ], - key=['BK', 'NC', 'BT', 'HV'], - **autotune_cache_kwargs, -) -@triton.jit(do_not_specialize=['B', 'T']) -def chunk_kda_bwd_kernel_intra( - q, - k, - g, - beta, - dAqk, - dAkk, - dq, - dq2, - dk, - dk2, - dg, - dg2, - db, - cu_seqlens, - chunk_indices, - B, - T, - H: tl.constexpr, - HV: tl.constexpr, - K: tl.constexpr, - BT: tl.constexpr, - BC: tl.constexpr, - BK: tl.constexpr, - NC: tl.constexpr, - IS_VARLEN: tl.constexpr, - SAFE_GATE: tl.constexpr, - USE_GATHER: tl.constexpr, -): - i_kc, i_t, i_bh = tl.program_id(0), tl.program_id(1), tl.program_id(2) - i_b, i_hv = i_bh // HV, i_bh % HV - i_h = i_hv // (HV // H) - i_k, i_i = i_kc // NC, i_kc % NC - - all = B * T - 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.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) - else: - bos, eos = i_b * T, i_b * T + T - T = eos - bos - - i_ti = i_t * BT + i_i * BC - if i_ti >= T: - return - - o_k = i_k * BK + tl.arange(0, BK) - m_k = o_k < K - - q += (bos * H + i_h) * K - k += (bos * H + i_h) * K - g += (bos * HV + i_hv) * K - beta += bos * HV + i_hv - - dAqk += (bos * HV + i_hv) * BT - dAkk += (bos * HV + i_hv) * BT - dq += (bos * HV + i_hv) * K - dq2 += (bos * HV + i_hv) * K - dk += (bos * HV + i_hv) * K - dk2 += (bos * HV + i_hv) * K - dg += (bos * HV + i_hv) * K - dg2 += (bos * HV + i_hv) * K - db += (i_k * all + bos) * HV + i_hv - - p_g = tl.make_block_ptr(g, (T, K), (HV*K, 1), (i_ti, i_k * BK), (BC, BK), (1, 0)) - b_g = tl.load(p_g, boundary_check=(0, 1)).to(tl.float32) - - p_b = tl.make_block_ptr(beta, (T,), (HV,), (i_ti,), (BC,), (0,)) - b_b = tl.load(p_b, boundary_check=(0,)) - - b_dq2 = tl.zeros([BC, BK], dtype=tl.float32) - b_dk2 = tl.zeros([BC, BK], dtype=tl.float32) - if i_i > 0: - p_gn = g + i_ti * HV*K + o_k - # [BK,] - b_gn = tl.load(p_gn, mask=m_k, other=0).to(tl.float32)[None, :] - for i_j in range(0, i_i): - p_k = tl.make_block_ptr(k, (T, K), (H*K, 1), (i_t * BT + i_j * BC, i_k * BK), (BC, BK), (1, 0)) - p_gk = tl.make_block_ptr(g, (T, K), (HV*K, 1), (i_t * BT + i_j * BC, i_k * BK), (BC, BK), (1, 0)) - p_dAqk = tl.make_block_ptr(dAqk, (T, BT), (HV*BT, 1), (i_ti, i_j * BC), (BC, BC), (1, 0)) - p_dAkk = tl.make_block_ptr(dAkk, (T, BT), (HV*BT, 1), (i_ti, i_j * BC), (BC, BC), (1, 0)) - # [BC, BK] - b_k = tl.load(p_k, boundary_check=(0, 1)) - b_gk = tl.load(p_gk, boundary_check=(0, 1)) - b_kg = b_k * exp2(b_gn - b_gk) - # [BC, BC] - b_dAqk = tl.load(p_dAqk, boundary_check=(0, 1)) - b_dAkk = tl.load(p_dAkk, boundary_check=(0, 1)) - # [BC, BK] - b_dq2 += tl.dot(b_dAqk, b_kg) - b_dk2 += tl.dot(b_dAkk, b_kg) - b_gqn = exp2(b_g - b_gn) - b_dq2 *= b_gqn - b_dk2 *= b_gqn - - o_i = tl.arange(0, BC) - m_dA = (i_ti + o_i) < T - o_dA = (i_ti + o_i) * HV*BT + i_i * BC - p_kj = k + i_ti * H*K + o_k - p_gkj = g + i_ti * HV*K + o_k - - p_q = tl.make_block_ptr(q, (T, K), (H*K, 1), (i_ti, i_k * BK), (BC, BK), (1, 0)) - p_k = tl.make_block_ptr(k, (T, K), (H*K, 1), (i_ti, i_k * BK), (BC, BK), (1, 0)) - b_q = tl.load(p_q, boundary_check=(0, 1)) - b_k = tl.load(p_k, boundary_check=(0, 1)) - - if SAFE_GATE: - if USE_GATHER: - b_gn = gather(b_g, tl.full([1, BK], min(BC//2, T - i_ti - 1), dtype=tl.int16), axis=0) - else: - p_gn = g + (i_ti + min(BC // 2, T - i_ti - 1)) * HV*K + o_k - b_gn = tl.load(p_gn, mask=m_k, other=0)[None, :] - - p_dAqk = tl.make_block_ptr(dAqk, (T, BT), (HV*BT, 1), (i_ti, i_i * BC), (BC, BC), (1, 0)) - p_dAkk = tl.make_block_ptr(dAkk, (T, BT), (HV*BT, 1), (i_ti, i_i * BC), (BC, BC), (1, 0)) - b_dAqk_diag_qk = tl.load(p_dAqk, boundary_check=(0, 1)).to(tl.float32) - b_dAkk_diag_qk = tl.load(p_dAkk, boundary_check=(0, 1)).to(tl.float32) - - m_i_diag_qk = (o_i[:, None] >= o_i[None, :]) & ((i_ti + o_i[:, None]) < T) & ((i_ti + o_i[None, :]) < T) - m_j_diag_qk = (i_ti + o_i[:, None]) < T - - b_dAqk_diag_qk = tl.where(m_i_diag_qk, b_dAqk_diag_qk, 0.) - b_dAkk_diag_qk = tl.where(m_i_diag_qk, b_dAkk_diag_qk, 0.) - b_g_diag_qk = tl.where(m_j_diag_qk, b_g - b_gn, 0.) - exp_b_g_diag_qk = tl.where(m_j_diag_qk, exp2(b_g_diag_qk), 0.) - exp_neg_b_g_diag_qk = tl.where(m_j_diag_qk, exp2(-b_g_diag_qk), 0.) - - b_k_exp_diag_qk = b_k * exp_neg_b_g_diag_qk - b_dq2 += tl.dot(b_dAqk_diag_qk, b_k_exp_diag_qk) * exp_b_g_diag_qk - b_dk2 += tl.dot(b_dAkk_diag_qk, b_k_exp_diag_qk) * exp_b_g_diag_qk - else: - for j in range(0, min(BC, T - i_t * BT - i_i * BC)): - # [BC] - b_dAqk = tl.load(dAqk + o_dA + j, mask=m_dA, other=0) - b_dAkk = tl.load(dAkk + o_dA + j, mask=m_dA, other=0) - # [BK] - b_kj = tl.load(p_kj, mask=m_k, other=0).to(tl.float32) - b_gkj = tl.load(p_gkj, mask=m_k, other=0).to(tl.float32) - # [BC, BK] - m_i = o_i[:, None] >= j - # [BC, BK] - b_gqk = exp2(b_g - b_gkj[None, :]) - b_dq2 += tl.where(m_i, b_dAqk[:, None] * b_kj[None, :] * b_gqk, 0.) - b_dk2 += tl.where(m_i, b_dAkk[:, None] * b_kj[None, :] * b_gqk, 0.) - - p_kj += H*K - p_gkj += HV*K - - b_db = tl.sum(b_dk2 * b_k, 1) - b_dk2 *= b_b[:, None] - - p_dq = tl.make_block_ptr(dq, (T, K), (HV*K, 1), (i_ti, i_k * BK), (BC, BK), (1, 0)) - p_dq2 = tl.make_block_ptr(dq2, (T, K), (HV*K, 1), (i_ti, i_k * BK), (BC, BK), (1, 0)) - p_db = tl.make_block_ptr(db, (T,), (HV,), (i_ti,), (BC,), (0,)) - - b_dg2 = b_q * b_dq2 - b_dq2 = b_dq2 + tl.load(p_dq, boundary_check=(0, 1)) - tl.store(p_dq2, b_dq2.to(p_dq2.dtype.element_ty), boundary_check=(0, 1)) - tl.store(p_db, b_db.to(p_db.dtype.element_ty), boundary_check=(0,)) - - tl.debug_barrier() - b_dkt = tl.zeros([BC, BK], dtype=tl.float32) - - NC = min(NC, tl.cdiv(T - i_t * BT, BC)) - if i_i < NC - 1: - p_gn = g + (min(i_ti + BC, T) - 1) * HV*K + o_k - # [BK,] - b_gn = tl.load(p_gn, mask=m_k, other=0).to(tl.float32)[None, :] - for i_j in range(i_i + 1, NC): - p_q = tl.make_block_ptr(q, (T, K), (H*K, 1), (i_t*BT+i_j*BC, i_k*BK), (BC, BK), (1, 0)) - p_k = tl.make_block_ptr(k, (T, K), (H*K, 1), (i_t * BT + i_j * BC, i_k * BK), (BC, BK), (1, 0)) - p_gk = tl.make_block_ptr(g, (T, K), (HV*K, 1), (i_t * BT + i_j * BC, i_k*BK), (BC, BK), (1, 0)) - p_b = tl.make_block_ptr(beta, (T,), (HV,), (i_t * BT + i_j * BC,), (BC,), (0,)) - p_dAqk = tl.make_block_ptr(dAqk, (BT, T), (1, HV*BT), (i_i * BC, i_t * BT + i_j * BC), (BC, BC), (0, 1)) - p_dAkk = tl.make_block_ptr(dAkk, (BT, T), (1, HV*BT), (i_i * BC, i_t * BT + i_j * BC), (BC, BC), (0, 1)) - # [BC] - b_b = tl.load(p_b, boundary_check=(0,)) - # [BC, BK] - b_q = tl.load(p_q, boundary_check=(0, 1)) - b_kb = tl.load(p_k, boundary_check=(0, 1)) * b_b[:, None] - b_gk = tl.load(p_gk, boundary_check=(0, 1)).to(tl.float32) - # [BC, BC] - b_dAqk = tl.load(p_dAqk, boundary_check=(0, 1)) - b_dAkk = tl.load(p_dAkk, boundary_check=(0, 1)) - - o_j = i_t * BT + i_j * BC + o_i - m_j = o_j < T - # [BC, BK] - b_gkn = exp2(b_gk - b_gn) - b_qg = b_q * tl.where(m_j[:, None], b_gkn, 0) - b_kbg = b_kb * tl.where(m_j[:, None], b_gkn, 0) - # [BC, BK] - # (SY 09/17) important to not use bf16 here to have a good precision. - b_dkt += tl.dot(b_dAqk, b_qg) - b_dkt += tl.dot(b_dAkk, b_kbg) - b_dkt *= exp2(b_gn - b_g) - o_dA = i_ti * HV*BT + i_i * BC + o_i - p_qj = q + i_ti * H*K + o_k - p_kj = k + i_ti * H*K + o_k - p_gkj = g + i_ti * HV*K + o_k - p_bj = beta + i_ti * HV - - if SAFE_GATE: - if USE_GATHER: - b_gn = gather(b_g, tl.full([1, BK], min(BC//2, T - i_ti - 1), dtype=tl.int16), axis=0) - else: - p_gn = g + (i_ti + min(BC // 2, T - i_ti - 1)) * HV*K + o_k - b_gn = tl.load(p_gn, mask=m_k, other=0).to(tl.float32)[None, :] - p_q = tl.make_block_ptr(q, (T, K), (H*K, 1), (i_ti, i_k * BK), (BC, BK), (1, 0)) - b_q = tl.load(p_q, boundary_check=(0, 1)) - p_b = tl.make_block_ptr(beta, (T,), (HV,), (i_ti,), (BC,), (0,)) - b_b = tl.load(p_b, boundary_check=(0,)) - - p_dAqk = tl.make_block_ptr(dAqk, (BT, T), (1, HV*BT), (i_i * BC, i_ti), (BC, BC), (0, 1)) - p_dAkk = tl.make_block_ptr(dAkk, (BT, T), (1, HV*BT), (i_i * BC, i_ti), (BC, BC), (0, 1)) - b_dAqk_diag_kk = tl.load(p_dAqk, boundary_check=(0, 1)).to(tl.float32) - b_dAkk_diag_kk = tl.load(p_dAkk, boundary_check=(0, 1)).to(tl.float32) - - m_i_diag_kk = (o_i[:, None] <= o_i[None, :]) & ((i_ti + o_i[:, None]) < T) & ((i_ti + o_i[None, :]) < T) - m_j_diag_kk = (i_ti + o_i[:, None]) < T - - b_dAqk_diag_kk = tl.where(m_i_diag_kk, b_dAqk_diag_kk, 0.) - b_dAkk_diag_kk = tl.where(m_i_diag_kk, b_dAkk_diag_kk, 0.) - # ensure numerical stability - b_g_diag_kk = tl.where(m_j_diag_kk, b_g - b_gn, 0.) - exp_b_g_diag_kk = tl.where(m_j_diag_kk, exp2(b_g_diag_kk), 0.) - exp_neg_b_g_diag_kk = tl.where(m_j_diag_kk, exp2(-b_g_diag_kk), 0.) - - b_q_exp = b_q * exp_b_g_diag_kk - b_kb_exp = b_k * b_b[:, None] * exp_b_g_diag_kk - - b_dkt += tl.dot(b_dAqk_diag_kk, b_q_exp) * exp_neg_b_g_diag_kk - b_dkt += tl.dot(b_dAkk_diag_kk, b_kb_exp) * exp_neg_b_g_diag_kk - else: - for j in range(0, min(BC, T - i_t * BT - i_i * BC)): - # [BC,] - b_dAqk = tl.load(dAqk + o_dA + j * HV*BT) - b_dAkk = tl.load(dAkk + o_dA + j * HV*BT) - # [BK,] - b_qj = tl.load(p_qj, mask=m_k, other=0).to(tl.float32) - b_kbj = tl.load(p_kj, mask=m_k, other=0).to(tl.float32) * tl.load(p_bj) - b_gkj = tl.load(p_gkj, mask=m_k, other=0).to(tl.float32) - # [BC, BK] - m_i = o_i[:, None] <= j - b_gkq = exp2(b_gkj[None, :] - b_g) - b_dkt += tl.where(m_i, b_dAqk[:, None] * b_qj[None, :] * b_gkq, 0.) - b_dkt += tl.where(m_i, b_dAkk[:, None] * b_kbj[None, :] * b_gkq, 0.) - - p_qj += H*K - p_kj += H*K - p_gkj += HV*K - p_bj += HV - p_dk = tl.make_block_ptr(dk, (T, K), (HV*K, 1), (i_ti, i_k * BK), (BC, BK), (1, 0)) - p_dk2 = tl.make_block_ptr(dk2, (T, K), (HV*K, 1), (i_ti, i_k * BK), (BC, BK), (1, 0)) - p_dg = tl.make_block_ptr(dg, (T, K), (HV*K, 1), (i_ti, i_k * BK), (BC, BK), (1, 0)) - p_dg2 = tl.make_block_ptr(dg2, (T, K), (HV*K, 1), (i_ti, i_k * BK), (BC, BK), (1, 0)) - - b_dg2 += (b_dk2 - b_dkt) * b_k + tl.load(p_dg, boundary_check=(0, 1)) - b_dk2 += tl.load(p_dk, boundary_check=(0, 1)) - b_dk2 += b_dkt - - tl.store(p_dk2, b_dk2.to(p_dk2.dtype.element_ty), boundary_check=(0, 1)) - tl.store(p_dg2, b_dg2.to(p_dg2.dtype.element_ty), boundary_check=(0, 1)) - - -@triton.heuristics({ - 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None, -}) -@fla_cache_autotune( - configs=[ - triton.Config({}, num_warps=num_warps, num_stages=num_stages) - for num_warps in [1, 2, 4, 8] - for num_stages in [2, 3, 4] - ], - key=["BT", "BC", "HV"], - **autotune_cache_kwargs, -) -@triton.jit(do_not_specialize=['T']) -def chunk_kda_fwd_kernel_intra_sub_chunk( - q, - k, - g, - beta, - Aqk, - Akk, - scale, - cu_seqlens, - chunk_indices, - T, - H: tl.constexpr, - HV: tl.constexpr, - K: tl.constexpr, - BT: tl.constexpr, - BC: tl.constexpr, - BK: tl.constexpr, - IS_VARLEN: tl.constexpr, - USE_GATHER: tl.constexpr, -): - i_t, i_i, i_bh = tl.program_id(0), tl.program_id(1), tl.program_id(2) - i_b, i_hv = i_bh // HV, i_bh % HV - i_h = i_hv // (HV // H) - - 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.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) - T = eos - bos - else: - bos, eos = i_b * T, i_b * T + T - - i_ti = i_t * BT + i_i * BC - if i_ti >= T: - return - - o_c = i_ti + tl.arange(0, BC) - m_c = o_c < T - - q = q + (bos * H + i_h) * K - k = k + (bos * H + i_h) * K - g = g + (bos * HV + i_hv) * K - beta = beta + bos * HV + i_hv - Aqk = Aqk + (bos * HV + i_hv) * BT - Akk = Akk + (bos * HV + i_hv) * BC - - p_q = tl.make_block_ptr(q, (T, K), (H*K, 1), (i_ti, 0), (BC, BK), (1, 0)) - p_k = tl.make_block_ptr(k, (T, K), (H*K, 1), (i_ti, 0), (BC, BK), (1, 0)) - p_g = tl.make_block_ptr(g, (T, K), (HV*K, 1), (i_ti, 0), (BC, BK), (1, 0)) - - p_beta = tl.make_block_ptr(beta, (T,), (HV,), (i_ti,), (BC,), (0,)) - - b_q = tl.load(p_q, boundary_check=(0, 1)) - b_k = tl.load(p_k, boundary_check=(0, 1)) - b_g = tl.load(p_g, boundary_check=(0, 1)) - b_beta = tl.load(p_beta, boundary_check=(0,)) - - if USE_GATHER: - b_gn = gather(b_g, tl.full([1, BK], min(BC//2, T - i_ti - 1), dtype=tl.int16), axis=0) - else: - # caculate offset - p_gn = g + (i_ti + min(BC // 2, T - i_ti - 1)) * HV*K + tl.arange(0, BK) - b_gn = tl.load(p_gn, mask=tl.arange(0, BK) < K, other=0.0) - b_gn = b_gn[None, :] - - # current block, keep numerical stability by subtracting the left boundary - # less than 85 to avoid overflow in exp2 - b_gm = (b_g - b_gn).to(tl.float32) - - b_gq = tl.where(m_c[:, None], exp2(b_gm), 0.) - b_gk = tl.where(m_c[:, None], exp2(-b_gm), 0.) - - b_kgt = tl.trans(b_k * b_gk) - - b_Aqk = tl.dot(b_q * b_gq, b_kgt) * scale - b_Akk = tl.dot(b_k * b_gq, b_kgt) * b_beta[:, None] - - o_i = tl.arange(0, BC) - m_Aqk = o_i[:, None] >= o_i[None, :] - m_Akk = o_i[:, None] > o_i[None, :] - m_I = o_i[:, None] == o_i[None, :] - - b_Aqk = tl.where(m_Aqk, b_Aqk, 0.0) - b_Akk = tl.where(m_Akk, b_Akk, 0.0) - - p_Aqk = tl.make_block_ptr(Aqk, (T, BT), (HV*BT, 1), (i_ti, i_i * BC), (BC, BC), (1, 0)) - p_Akk = tl.make_block_ptr(Akk, (T, BC), (HV*BC, 1), (i_ti, 0), (BC, BC), (1, 0)) - tl.store(p_Aqk, b_Aqk.to(Aqk.dtype.element_ty), boundary_check=(0, 1)) - tl.store(p_Akk, b_Akk.to(Akk.dtype.element_ty), boundary_check=(0, 1)) - - tl.debug_barrier() - - ################################################################################ - # forward substitution - ################################################################################ - - b_Ai = -b_Akk - for i in range(2, min(BC, T - i_ti)): - b_a = -tl.load(Akk + (i_ti + i) * HV*BC + o_i) - b_a = tl.where(o_i < i, b_a, 0.) - b_a += tl.sum(b_a[:, None] * b_Ai, 0) - b_Ai = tl.where((o_i == i)[:, None], b_a, b_Ai) - b_Ai += m_I - tl.store(p_Akk, b_Ai.to(Akk.dtype.element_ty), boundary_check=(0, 1)) - - -def chunk_kda_fwd_intra( - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - gk: torch.Tensor | None = None, - beta: torch.Tensor | None = None, - scale: float | None = None, - cu_seqlens: torch.LongTensor | None = None, - chunk_size: int = 64, - chunk_indices: torch.LongTensor | None = None, - safe_gate: bool = False, - disable_recompute: bool = False, -): - B, T, H, K, HV = *k.shape, gk.shape[2] - BT = chunk_size - if BT not in (32, 64): - raise ValueError(f"KDA intra chunk kernel only supports chunk_size 32 or 64, got {BT}.") - BC = 16 - if chunk_indices is None and cu_seqlens is not None: - chunk_indices = prepare_chunk_indices(cu_seqlens, BT) - NT = triton.cdiv(T, BT) if cu_seqlens is None else len(chunk_indices) - NC = triton.cdiv(BT, BC) - - Aqk = torch.empty(B, T, HV, BT, device=k.device, dtype=k.dtype) - # Akk must be zero-initialized - kernel only writes lower triangular - Akk = torch.zeros(B, T, HV, BT, device=k.device, dtype=k.dtype) - # Separate fp32 buffer for diagonal 16x16 blocks (for precision in solve_tril) - Akkd = torch.empty(B, T, HV, BC, device=k.device, dtype=torch.float32) - - # Step 1: Run token_parallel first to compute diagonal blocks into Akkd (fp32) - # Step 1: compute diagonal blocks into Akk_diag (fp32) - if safe_gate: - grid = (NT, NC, B * HV) - BK = triton.next_power_of_2(K) - chunk_kda_fwd_kernel_intra_sub_chunk[grid]( - q=q, - k=k, - g=gk, - beta=beta, - Aqk=Aqk, - Akk=Akkd, - scale=scale, - cu_seqlens=cu_seqlens, - chunk_indices=chunk_indices, - T=T, - H=H, - HV=HV, - K=K, - BT=BT, - BC=BC, - BK=BK, - USE_GATHER=IS_GATHER_SUPPORTED, - ) - else: - Aqk, Akkd = chunk_kda_fwd_intra_token_parallel( - q=q, - k=k, - gk=gk, - beta=beta, - Aqk=Aqk, - Akk=Akkd, - scale=scale, - cu_seqlens=cu_seqlens, - chunk_size=BT, - sub_chunk_size=BC, - ) - - # Step 2: Fused inter + solve_tril (works for both fixed-len and varlen) - grid = (NT, B * HV) - chunk_kda_fwd_kernel_inter_solve_fused[grid]( - q=q, - k=k, - g=gk, - beta=beta, - Aqk=Aqk, - Akkd=Akkd, - Akk=Akk, - scale=scale, - cu_seqlens=cu_seqlens, - chunk_indices=chunk_indices, - T=T, - H=H, - HV=HV, - K=K, - BT=BT, - BC=BC, - NC=NC, - USE_SAFE_GATE=safe_gate, - ) - w, u, qg, kg = recompute_w_u_fwd( - k=k, - v=v, - beta=beta, - A=Akk, - q=q if disable_recompute else None, - gk=gk, - cu_seqlens=cu_seqlens, - chunk_indices=chunk_indices, - ) - return w, u, qg, kg, Aqk, Akk - - -def chunk_kda_bwd_intra( - q: torch.Tensor, - k: torch.Tensor, - g: torch.Tensor, - beta: torch.Tensor, - dAqk: torch.Tensor, - dAkk: torch.Tensor, - dq: torch.Tensor, - dk: torch.Tensor, - db: torch.Tensor, - dg: torch.Tensor, - cu_seqlens: torch.LongTensor | None = None, - chunk_indices: torch.LongTensor | None = None, - chunk_size: int = 64, - safe_gate: bool = False, -): - B, T, H, K, HV = *k.shape, g.shape[2] - BT = chunk_size - BC = min(16, BT) - BK = min(32, triton.next_power_of_2(K)) - - if chunk_indices is None and cu_seqlens is not None: - chunk_indices = prepare_chunk_indices(cu_seqlens, BT) - NT = triton.cdiv(T, BT) if cu_seqlens is None else len(chunk_indices) - NC = triton.cdiv(BT, BC) - NK = triton.cdiv(K, BK) - - dq2 = torch.empty_like(dq) - dk2 = torch.empty_like(dk) - db2 = beta.new_empty(NK, *beta.shape, dtype=torch.float) - dg2 = torch.empty_like(dg, dtype=torch.float) - grid = (NK * NC, NT, B * HV) - chunk_kda_bwd_kernel_intra[grid]( - q=q, - k=k, - g=g, - beta=beta, - dAqk=dAqk, - dAkk=dAkk, - dq=dq, - dq2=dq2, - dk=dk, - dk2=dk2, - dg=dg, - dg2=dg2, - db=db2, - cu_seqlens=cu_seqlens, - chunk_indices=chunk_indices, - B=B, - T=T, - H=H, - HV=HV, - K=K, - BT=BT, - BC=BC, - BK=BK, - NC=NC, - SAFE_GATE=safe_gate, - USE_GATHER=IS_GATHER_SUPPORTED, - ) - dq = dq2 - dk = dk2 - db = db2.sum(0).add_(db) - dg = dg2 - - return dq, dk, db, dg diff --git a/build/torch-xpu/ops/kda/chunk_intra_token_parallel.py b/build/torch-xpu/ops/kda/chunk_intra_token_parallel.py deleted file mode 100644 index d1ce022245551716e6477c06fd011d8d3ed4c21e..0000000000000000000000000000000000000000 --- a/build/torch-xpu/ops/kda/chunk_intra_token_parallel.py +++ /dev/null @@ -1,178 +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 - -# Token-parallel implementation of KDA intra chunk kernel - -import torch -import triton -import triton.language as tl - -from ...ops.utils.cache import fla_cache_autotune -from ...ops.utils.op import exp2 -from ...utils import autotune_cache_kwargs - - -@triton.heuristics({ - 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None, -}) -@fla_cache_autotune( - configs=[ - triton.Config({'BH': BH}, num_warps=num_warps) - for BH in [1, 2, 4, 8] - for num_warps in [1, 2, 4, 8] - ], - key=["K", "H", "HV"], - **autotune_cache_kwargs, -) -@triton.jit(do_not_specialize=['T', 'N']) -def chunk_kda_fwd_kernel_intra_token_parallel( - q, - k, - g, - beta, - Aqk, - Akk, - scale, - cu_seqlens, - N, - T, - H: tl.constexpr, - HV: tl.constexpr, - K: tl.constexpr, - BT: tl.constexpr, - BC: tl.constexpr, - BH: tl.constexpr, - IS_VARLEN: tl.constexpr, -): - i_tg, i_hg = tl.program_id(0), tl.program_id(1) - - if IS_VARLEN: - i_n = 0 - left, right = 0, N - - # Unrolled binary search (max B=2^32) - # We can limit iterations based on expected max batch size if needed - # 20 iterations covers B=1M, usually enough - for _ in range(20): - if left < right: - mid = (left + right) // 2 - if i_tg < tl.load(cu_seqlens + mid + 1).to(tl.int32): - right = mid - else: - left = mid + 1 - i_n = left - - bos, eos = tl.load(cu_seqlens + i_n).to(tl.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) - T = eos - bos - i_t = i_tg - bos - else: - bos = (i_tg // T) * T - i_t = i_tg % T - - if i_t >= T: - return - - i_c = i_t // BT - i_s = (i_t % BT) // BC - i_tc = i_c * BT - i_ts = i_tc + i_s * BC - - G: tl.constexpr = HV // H - - q += bos * H*K - k += bos * H*K - g += bos * HV*K - Aqk += bos * HV*BT - Akk += bos * HV*BC - beta += bos * HV - - BK: tl.constexpr = triton.next_power_of_2(K) - o_hv = i_hg * BH + tl.arange(0, BH) - o_h = o_hv // G - o_k = tl.arange(0, BK) - m_hv = o_hv < HV - m_k = o_k < K - m_hk = m_hv[:, None] & m_k[None, :] - - # q/k: [B, T, H, K], manual load via mapped qk head index - p_qk = o_h[:, None] * K + o_k[None, :] - b_q = tl.load(q + i_t * H * K + p_qk, mask=m_hk, other=0).to(tl.float32) - b_k = tl.load(k + i_t * H * K + p_qk, mask=m_hk, other=0).to(tl.float32) - - # g: [B, T, HV, K], beta: [B, T, HV] - p_g = tl.make_block_ptr(g + i_t * HV * K, (HV, K), (K, 1), (i_hg * BH, 0), (BH, BK), (1, 0)) - p_beta = tl.make_block_ptr(beta + i_t * HV, (HV,), (1,), (i_hg * BH,), (BH,), (0,)) - b_g = tl.load(p_g, boundary_check=(0, 1)).to(tl.float32) - b_k = b_k * tl.load(p_beta, boundary_check=(0,)).to(tl.float32)[:, None] - - for j in range(i_ts, min(i_t + 1, min(T, i_ts + BC))): - b_kj = tl.load(k + j * H * K + p_qk, mask=m_hk, other=0).to(tl.float32) - p_gj = tl.make_block_ptr(g + j * HV * K, (HV, K), (K, 1), (i_hg * BH, 0), (BH, BK), (1, 0)) - b_gj = tl.load(p_gj, boundary_check=(0, 1)).to(tl.float32) - - b_kgj = tl.where(m_k[None, :], b_kj * exp2(b_g - b_gj), 0.0) - b_Aqk = tl.sum(b_q * b_kgj, axis=1) * scale - b_Akk = tl.sum(b_k * b_kgj, axis=1) * tl.where(j < i_t, 1.0, 0.0) - - tl.store(Aqk + i_t * HV * BT + o_hv * BT + j % BT, b_Aqk.to(Aqk.dtype.element_ty), mask=m_hv) - tl.store(Akk + i_t * HV * BC + o_hv * BC + j - i_ts, b_Akk.to(Akk.dtype.element_ty), mask=m_hv) - - -def chunk_kda_fwd_intra_token_parallel( - q: torch.Tensor, - k: torch.Tensor, - gk: torch.Tensor, - beta: torch.Tensor, - Aqk: torch.Tensor, - Akk: torch.Tensor, - scale: float, - cu_seqlens: torch.LongTensor | None = None, - chunk_size: int = 64, - sub_chunk_size: int = 16, -) -> None: - """ - Token-parallel implementation: each token gets its own thread block. - Supports both fixed-length and variable-length sequences. - Reduces wasted computation on padding. - - Writes directly to Aqk and Akk tensors (in-place). - - Args: - q: [B, T, H, K] - k: [B, T, H, K] - gk: [B, T, HV, K] cumsum of gates (HV >= H for GVA) - beta: [B, T, HV] - Aqk: [B, T, HV, BT] output tensor to write to - Akk: [B, T, HV, BC] output tensor for diagonal blocks (fp32) - scale: attention scale - chunk_size: BT (default 64) - sub_chunk_size: BC (default 16) - """ - B, T, H, K, HV = *q.shape, gk.shape[2] - N = len(cu_seqlens) - 1 if cu_seqlens is not None else B - BT = chunk_size - BC = sub_chunk_size - - def grid(meta): return (B * T, triton.cdiv(HV, meta['BH'])) - chunk_kda_fwd_kernel_intra_token_parallel[grid]( - q=q, - k=k, - g=gk, - beta=beta, - Aqk=Aqk, - Akk=Akk, - scale=scale, - cu_seqlens=cu_seqlens, - N=N, - T=T, - H=H, - HV=HV, - K=K, - BT=BT, - BC=BC, - ) - return Aqk, Akk diff --git a/build/torch-xpu/ops/kda/gate.py b/build/torch-xpu/ops/kda/gate.py deleted file mode 100644 index bf38ddb6ee9ffa2d50ecbe1dd8699a42d561ded2..0000000000000000000000000000000000000000 --- a/build/torch-xpu/ops/kda/gate.py +++ /dev/null @@ -1,461 +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 - -# This file is modified and supported by the Moonshot AI Team - -import torch -import torch.nn.functional as F -import triton -import triton.language as tl - -from ...ops.utils.cache import fla_cache_autotune -from ...ops.utils.index import prepare_chunk_indices -from ...ops.utils.op import exp -from ...ops.utils.softplus import softplus -from ...utils import IS_AMD, autocast_custom_bwd, autocast_custom_fwd, autotune_cache_kwargs, check_shared_mem, input_guard - -BS_LIST = [32, 64] if check_shared_mem() else [16, 32] -BT_LIST_AUTOTUNE = [32, 64, 128] -NUM_WARPS_AUTOTUNE = [2, 4, 8, 16] if IS_AMD else [4, 8, 16, 32] - - -def naive_kda_gate( - g: torch.Tensor, - A_log: torch.Tensor, - dt_bias: torch.Tensor | None = None, - output_dtype: torch.dtype = torch.float32, -) -> torch.Tensor: - """ - Torch reference implementation for KDA gate computation. - - Computes: g = -A_log.exp().unsqueeze(-1) * softplus(g + dt_bias.view(g.shape[-2:])) - - Args: - g (torch.Tensor): - Input tensor of shape `[..., H, K]`. - A_log (torch.Tensor): - Parameter tensor with `H` elements. - dt_bias (torch.Tensor | None): - Optional bias tensor added to `g` before activation, shape `[H * K]`. - - Returns: - Output tensor of shape `[..., H, K]` . - """ - H, _ = g.shape[-2:] - g = g.float() - if dt_bias is not None: - g = g + dt_bias.view(H, -1) - - g = (-A_log.view(H, 1).float().exp() * F.softplus(g.float())).to(output_dtype) - return g - - -def naive_kda_lowerbound_gate( - g: torch.Tensor, - A_log: torch.Tensor, - dt_bias: torch.Tensor | None = None, - lower_bound: float = -5.0, - output_dtype: torch.dtype = torch.float32, -) -> torch.Tensor: - H, _ = g.shape[-2:] - g = g.float() - if dt_bias is not None: - g = g + dt_bias.view(H, -1) - g = lower_bound * F.sigmoid(A_log.view(H, 1).exp() * g) - return g.to(output_dtype) - - -@triton.heuristics({ - "HAS_BIAS": lambda args: args["dt_bias"] is not None, - "HAS_BETA": lambda args: args["beta"] is not None, - 'USE_LOWER_BOUND': lambda args: args['lower_bound'] is not None, -}) -@fla_cache_autotune( - configs=[ - triton.Config({"BT": BT}, num_warps=num_warps, num_stages=num_stages) - for BT in BT_LIST_AUTOTUNE - for num_warps in NUM_WARPS_AUTOTUNE - for num_stages in [2, 3] - ], - key=["H", "D"], - **autotune_cache_kwargs, -) -@triton.jit(do_not_specialize=['T']) -def kda_gate_fwd_kernel( - g, - A_log, - dt_bias, - beta, - yg, - yb, - lower_bound, - T, - H: tl.constexpr, - D: tl.constexpr, - BT: tl.constexpr, - BD: tl.constexpr, - HAS_BIAS: tl.constexpr, - HAS_BETA: tl.constexpr, - USE_LOWER_BOUND: tl.constexpr, -): - i_t, i_h = tl.program_id(0), tl.program_id(1) - - b_A = tl.load(A_log + i_h).to(tl.float32) - - p_g = tl.make_block_ptr(g + i_h * D, (T, D), (H * D, 1), (i_t * BT, 0), (BT, BD), (1, 0)) - p_yg = tl.make_block_ptr(yg + i_h * D, (T, D), (H * D, 1), (i_t * BT, 0), (BT, BD), (1, 0)) - # [BT, BD] - b_g = tl.load(p_g, boundary_check=(0, 1)).to(tl.float32) - if HAS_BIAS: - p_b = tl.make_block_ptr(dt_bias, (H * D,), (1,), (i_h * D,), (BD,), (0,)) - b_g = b_g + tl.load(p_b, boundary_check=(0,)).to(tl.float32) - if not USE_LOWER_BOUND: - b_yg = -exp(b_A) * softplus(b_g) - else: - b_yg = lower_bound * tl.sigmoid(exp(b_A) * b_g) - tl.store(p_yg, b_yg.to(p_yg.dtype.element_ty), boundary_check=(0, 1)) - - if HAS_BETA: - p_b = tl.make_block_ptr(beta + i_h, (T,), (H,), (i_t * BT,), (BT,), (0,)) - p_yb = tl.make_block_ptr(yb + i_h, (T,), (H,), (i_t * BT,), (BT,), (0,)) - b_yb = tl.sigmoid(tl.load(p_b, boundary_check=(0,)).to(tl.float32)) - tl.store(p_yb, b_yb.to(p_yb.dtype.element_ty), boundary_check=(0,)) - - -@triton.heuristics({ - "HAS_BIAS": lambda args: args["dt_bias"] is not None, - "HAS_BETA": lambda args: args["beta"] is not None, - 'USE_LOWER_BOUND': lambda args: args['lower_bound'] is not None, -}) -@fla_cache_autotune( - configs=[ - triton.Config({}, num_warps=num_warps, num_stages=num_stages) - for num_warps in NUM_WARPS_AUTOTUNE - for num_stages in [2, 3] - ], - key=["H", "D"], - **autotune_cache_kwargs, -) -@triton.jit(do_not_specialize=['T']) -def kda_gate_bwd_kernel( - g, - A_log, - dt_bias, - beta, - dyg, - dyb, - dg, - dA, - dbeta, - lower_bound, - T, - H: tl.constexpr, - D: tl.constexpr, - BT: tl.constexpr, - BD: tl.constexpr, - HAS_BIAS: tl.constexpr, - HAS_BETA: tl.constexpr, - USE_LOWER_BOUND: tl.constexpr, -): - i_t, i_h = tl.program_id(0), tl.program_id(1) - - b_A = tl.load(A_log + i_h).to(tl.float32) - - p_g = tl.make_block_ptr(g + i_h * D, (T, D), (H * D, 1), (i_t * BT, 0), (BT, BD), (1, 0)) - p_dg = tl.make_block_ptr(dg + i_h * D, (T, D), (H * D, 1), (i_t * BT, 0), (BT, BD), (1, 0)) - p_dyg = tl.make_block_ptr(dyg + i_h * D, (T, D), (H * D, 1), (i_t * BT, 0), (BT, BD), (1, 0)) - - # [BT, BD] - b_g = tl.load(p_g, boundary_check=(0, 1)).to(tl.float32) - b_dyg = tl.load(p_dyg, boundary_check=(0, 1)).to(tl.float32) - - if HAS_BIAS: - p_b = tl.make_block_ptr(dt_bias, (H * D,), (1,), (i_h * D,), (BD,), (0,)) - b_g = b_g + tl.load(p_b, boundary_check=(0,)).to(tl.float32) - - # [BT, BD] - if not USE_LOWER_BOUND: - b_A = -exp(b_A) - b_yg = b_A * softplus(b_g) - b_dg = b_A * (b_dyg * tl.sigmoid(b_g)) - b_dA = tl.sum(tl.sum(b_dyg * b_yg, 1), 0) - else: - b_A = exp(b_A) - b_inner = b_A * b_g - b_sig = tl.sigmoid(b_inner) - b_dsig = b_sig * (1.0 - b_sig) - # Common term: dy * (LB * dsig) - b_d_inner_term = b_dyg * (lower_bound * b_dsig) - # dg = d_inner_term * A - b_dg = b_d_inner_term * b_A - b_dA = tl.sum(tl.sum(b_dg * b_g, 1), 0) - - tl.store(p_dg, b_dg.to(p_dg.dtype.element_ty), boundary_check=(0, 1)) - tl.store(dA + i_t * H + i_h, b_dA) - - if HAS_BETA: - p_b = tl.make_block_ptr(beta + i_h, (T,), (H,), (i_t * BT,), (BT,), (0,)) - p_db = tl.make_block_ptr(dbeta + i_h, (T,), (H,), (i_t * BT,), (BT,), (0,)) - p_dyb = tl.make_block_ptr(dyb + i_h, (T,), (H,), (i_t * BT,), (BT,), (0,)) - - b_b = tl.load(p_b, boundary_check=(0,)).to(tl.float32) - b_db = tl.load(p_dyb, boundary_check=(0,)).to(tl.float32) * b_b * (1.0 - b_b) - tl.store(p_db, b_db.to(p_db.dtype.element_ty), boundary_check=(0,)) - - -def kda_gate_fwd( - g: torch.Tensor, - A_log: torch.Tensor, - dt_bias: torch.Tensor | None = None, - lower_bound: float | None = None, - output_dtype: torch.dtype = torch.float32, -) -> torch.Tensor: - H, K = g.shape[-2:] - T = g.numel() // (H * K) - - yg = torch.empty_like(g, dtype=output_dtype) - - def grid(meta): - return (triton.cdiv(T, meta["BT"]), H) - - kda_gate_fwd_kernel[grid]( - g=g, - A_log=A_log, - dt_bias=dt_bias, - beta=None, - yg=yg, - yb=None, - T=T, - H=H, - D=K, - BD=triton.next_power_of_2(K), - lower_bound=lower_bound, - ) - return yg - - -def kda_gate_bwd( - g: torch.Tensor, - A_log: torch.Tensor, - dt_bias: torch.Tensor | None = None, - dyg: torch.Tensor | None = None, - lower_bound: float | None = None, -) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor | None]: - H, K = g.shape[-2:] - T = g.numel() // (H * K) - BT = 32 - NT = triton.cdiv(T, BT) - - dg = torch.empty_like(g, dtype=torch.float32) - dA = A_log.new_empty(NT, H, dtype=torch.float32) - - grid = (triton.cdiv(T, BT), H) - kda_gate_bwd_kernel[grid]( - g=g, - A_log=A_log, - dt_bias=dt_bias, - beta=None, - dyg=dyg, - dyb=None, - dg=dg, - dA=dA, - dbeta=None, - T=T, - H=H, - D=K, - BT=BT, - BD=triton.next_power_of_2(K), - lower_bound=lower_bound, - ) - - dg = dg.view_as(g).type_as(g) - dA = dA.sum(0).view_as(A_log).type_as(A_log) - dbias = dg.view(-1, H * K).sum(0).to(dt_bias) if dt_bias is not None else None - - return dg, dA, dbias - - -class KDAGateFunction(torch.autograd.Function): - @staticmethod - @input_guard - @autocast_custom_fwd - def forward( - ctx, - g: torch.Tensor, - A_log: torch.Tensor, - dt_bias: torch.Tensor | None = None, - lower_bound: float | None = None, - output_dtype: torch.dtype = torch.float32, - ) -> torch.Tensor: - yg = kda_gate_fwd( - g=g, - A_log=A_log, - dt_bias=dt_bias, - lower_bound=lower_bound, - output_dtype=output_dtype - ) - ctx.save_for_backward(g, A_log, dt_bias) - ctx.lower_bound = lower_bound - return yg - - @staticmethod - @input_guard - @autocast_custom_bwd - def backward(ctx, dyg: torch.Tensor): - g, A_log, dt_bias = ctx.saved_tensors - dg, dA, dbias = kda_gate_bwd( - g=g, - A_log=A_log, - dt_bias=dt_bias, - dyg=dyg, - lower_bound=ctx.lower_bound - ) - return dg, dA, dbias, None, None - - -@torch.compiler.disable -def fused_kda_gate( - g: torch.Tensor, - A_log: torch.Tensor, - dt_bias: torch.Tensor | None = None, - lower_bound: float | None = None, - output_dtype: torch.dtype = torch.float32, -) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]: - """ - Fused KDA gate computation with autograd support. - - Computes: g = -A_log.exp().unsqueeze(-1) * softplus(g + dt_bias.view(g.shape[-2:])) - - Args: - g (torch.Tensor): - Input tensor of shape `[..., H, K]`. - A_log (torch.Tensor): - Parameter tensor with `H` elements. - dt_bias (torch.Tensor | None): - Optional bias tensor added to `g` before activation, shape `[H * K]`. - - Returns: - Output tensor of shape `[..., H, K]`. - """ - return KDAGateFunction.apply(g, A_log, dt_bias, lower_bound, output_dtype) - - -@triton.heuristics({ - "HAS_BIAS": lambda args: args["dt_bias"] is not None, - 'HAS_SCALE': lambda args: args['scale'] is not None, - 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None, - 'USE_LOWER_BOUND': lambda args: args['lower_bound'] is not None, -}) -@fla_cache_autotune( - configs=[ - triton.Config({'BS': BS}, num_warps=num_warps) - for BS in BS_LIST - for num_warps in [2, 4, 8] - ], - key=['H', 'S', 'BT', 'IS_VARLEN', 'REVERSE'], - **autotune_cache_kwargs, -) -@triton.jit(do_not_specialize=['T']) -def kda_gate_chunk_cumsum_vector_kernel( - s, - A_log, - dt_bias, - o, - scale, - cu_seqlens, - chunk_indices, - lower_bound, - T, - H: tl.constexpr, - S: tl.constexpr, - BT: tl.constexpr, - BS: tl.constexpr, - REVERSE: tl.constexpr, - HAS_BIAS: tl.constexpr, - HAS_SCALE: tl.constexpr, - IS_VARLEN: tl.constexpr, - USE_LOWER_BOUND: tl.constexpr, -): - i_s, i_t, i_bh = tl.program_id(0), tl.program_id(1), tl.program_id(2) - i_b, i_h = i_bh // H, i_bh % H - 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.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) - T = eos - bos - else: - bos, eos = i_b * T, i_b * T + T - - p_s = tl.make_block_ptr(s + (bos * H + i_h) * S, (T, S), (H*S, 1), (i_t * BT, i_s * BS), (BT, BS), (1, 0)) - p_o = tl.make_block_ptr(o + (bos * H + i_h) * S, (T, S), (H*S, 1), (i_t * BT, i_s * BS), (BT, BS), (1, 0)) - # [BT, BS] - b_s = tl.load(p_s, boundary_check=(0, 1)).to(tl.float32) - - # Apply dt_bias if exists - if HAS_BIAS: - p_b = tl.make_block_ptr(dt_bias + i_h * S, (S,), (1,), (i_s * BS,), (BS,), (0,)) - b_bias = tl.load(p_b, boundary_check=(0,)).to(tl.float32) - b_s = b_s + b_bias[None, :] - - b_A = tl.load(A_log + i_h).to(tl.float32) - if not USE_LOWER_BOUND: - # Apply gate: -exp(A_log) * softplus(g + bias) - b_gate = -exp(b_A) * softplus(b_s) - else: - b_gate = lower_bound * tl.sigmoid(exp(b_A) * b_s) - - # Apply chunk local cumsum - if REVERSE: - b_o = tl.cumsum(b_gate, axis=0, reverse=True) - else: - b_o = tl.cumsum(b_gate, axis=0) - - if HAS_SCALE: - b_o *= scale - tl.store(p_o, b_o.to(p_o.dtype.element_ty), boundary_check=(0, 1)) - - -@input_guard -def kda_gate_chunk_cumsum( - g: torch.Tensor, - A_log: torch.Tensor, - chunk_size: int, - scale: float = None, - dt_bias: torch.Tensor | None = None, - cu_seqlens: torch.Tensor | None = None, - output_dtype: torch.dtype | None = torch.float, - chunk_indices: torch.LongTensor | None = None, - lower_bound: float | None = None, - **kwargs, -) -> torch.Tensor: - if cu_seqlens is not None: - assert g.shape[0] == 1, "Only batch size 1 is supported when cu_seqlens are provided" - assert len(g.shape) == 4 - B, T, H, S = g.shape - BT = chunk_size - if chunk_indices is None and cu_seqlens is not None: - chunk_indices = prepare_chunk_indices(cu_seqlens, BT) - NT = triton.cdiv(T, BT) if cu_seqlens is None else len(chunk_indices) - assert chunk_size == 2**(chunk_size.bit_length()-1), "chunk_size must be a power of 2" - - g_org, g = g, torch.empty_like(g, dtype=output_dtype or g.dtype) - def grid(meta): return (triton.cdiv(meta['S'], meta['BS']), NT, B * H) - kda_gate_chunk_cumsum_vector_kernel[grid]( - s=g_org, - A_log=A_log, - dt_bias=dt_bias, - o=g, - scale=scale, - cu_seqlens=cu_seqlens, - chunk_indices=chunk_indices, - lower_bound=lower_bound, - T=T, - H=H, - S=S, - BT=BT, - REVERSE=False, - ) - return g diff --git a/build/torch-xpu/ops/kda/naive.py b/build/torch-xpu/ops/kda/naive.py deleted file mode 100644 index 51ff7076b26244366cf7ed2999ed13052da19cd5..0000000000000000000000000000000000000000 --- a/build/torch-xpu/ops/kda/naive.py +++ /dev/null @@ -1,166 +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 -from einops import rearrange - - -def naive_recurrent_kda( - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - g: torch.Tensor, - beta: torch.Tensor, - scale: float | None = None, - initial_state: torch.Tensor | None = None, - output_final_state: bool = False, -): - r""" - Args: - q (torch.Tensor): - Queries of shape ``[B, T, H, K]``. - k (torch.Tensor): - Keys of shape ``[B, T, H, K]``. - v (torch.Tensor): - Values of shape ``[B, T, HV, V]``. ``HV`` must be divisible by ``H``. - g (torch.Tensor): - Per-dimension decay gates (log-space) of shape ``[B, T, HV, K]``. - beta (torch.Tensor): - Beta scalars of shape ``[B, T, HV]``. - scale (Optional[float]): - Scale factor. Defaults to ``1 / sqrt(K)``. - initial_state (Optional[torch.Tensor]): - Initial state of shape ``[B, HV, K, V]``. - output_final_state (bool): - Whether to return the final state. - - Returns: - A tuple ``(o, S)`` where ``o`` has shape ``[B, T, HV, V]`` and - ``S`` has shape ``[B, HV, K, V]`` if ``output_final_state`` else ``None``. - """ - dtype = v.dtype - B, T, H, K, HV, V = *q.shape, v.shape[2], v.shape[-1] - G = HV // H - if scale is None: - scale = K ** -0.5 - - q, k, v, g, beta = map(lambda x: x.to(torch.float), [q, k, v, g, beta]) - q = q.repeat_interleave(G, dim=2) * scale # [B, T, HV, K] - k = k.repeat_interleave(G, dim=2) # [B, T, HV, K] - - S = k.new_zeros(B, HV, K, V).to(q) - if initial_state is not None: - S += initial_state - o = torch.zeros_like(v) - for i in range(0, T): - q_i, k_i, v_i, g_i, b_i = q[:, i], k[:, i], v[:, i], g[:, i], beta[:, i] - S = S * g_i[..., None].exp() - S = S + torch.einsum('b h k, b h v -> b h k v', b_i[..., None] * k_i, v_i - (k_i[..., None] * S).sum(-2)) - o[:, i] = torch.einsum('b h k, b h k v -> b h v', q_i, S) - if not output_final_state: - S = None - return o.to(dtype), S - - -def naive_chunk_kda( - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - g: torch.Tensor, - beta: torch.Tensor, - scale: float | None = None, - initial_state: torch.Tensor | None = None, - output_final_state: bool = False, - chunk_size: int = 64, -): - r""" - Args: - q (torch.Tensor): - Queries of shape ``[B, T, H, K]``. - k (torch.Tensor): - Keys of shape ``[B, T, H, K]``. - v (torch.Tensor): - Values of shape ``[B, T, HV, V]``. ``HV`` must be divisible by ``H``. - g (torch.Tensor): - Per-dimension decay gates (log-space) of shape ``[B, T, HV, K]``. - beta (torch.Tensor): - Beta scalars of shape ``[B, T, HV]``. - scale (Optional[float]): - Scale factor. Defaults to ``1 / sqrt(K)``. - initial_state (Optional[torch.Tensor]): - Initial state of shape ``[B, HV, K, V]``. - output_final_state (bool): - Whether to return the final state. - chunk_size (int): - Chunk size for the chunked computation. Default: 64. - - Returns: - A tuple ``(o, S)`` where ``o`` has shape ``[B, T, HV, V]`` and - ``S`` has shape ``[B, HV, K, V]`` if ``output_final_state`` else ``None``. - """ - dtype = v.dtype - B, T, H, K, HV, V = *q.shape, v.shape[2], v.shape[-1] - G = HV // H - BT = chunk_size - NT = T // BT - if scale is None: - scale = K ** -0.5 - assert T % BT == 0 - - # Rearrange into chunks: [B, head, NT, BT, ...] - q, k = [rearrange(x, 'b (n c) h ... -> b h n c ...', c=BT).to(torch.float) for x in [q, k]] - v, g, beta = [rearrange(x, 'b (n c) h ... -> b h n c ...', c=BT).to(torch.float) for x in [v, g, beta]] - # Expand q/k to value head dim for GVA: [B, H, ...] -> [B, HV, ...] - q = q.repeat_interleave(G, dim=1) * scale # [B, HV, NT, BT, K] - k = k.repeat_interleave(G, dim=1) # [B, HV, NT, BT, K] - g = g.cumsum(-2) - - # note that diagonal is masked. - mask = torch.triu(torch.ones(BT, BT, dtype=torch.bool, device=q.device), diagonal=0) - - # Akk uses k (expanded to HV) and g (per value head) - A = torch.zeros(*g.shape[:-1], BT, dtype=torch.float, device=q.device) - for i in range(BT): - k_i = k[..., i, :] - g_i = g[..., i:i+1, :] - A[..., i] = torch.einsum('... c d, ... d -> ... c', k * (g - g_i).exp(), k_i) - A = A * beta[..., None] - - A = -A.masked_fill(mask, 0) - for i in range(1, BT): - A[..., i, :i] = A[..., i, :i].clone() + (A[..., i, :, None].clone() * A[..., :, :i].clone()).sum(-2) - A = (A + torch.eye(BT, dtype=torch.float, device=q.device)) * beta[..., None, :] - - w = A @ (g.exp() * k) - u = A @ v - - S = k.new_zeros(B, HV, K, V).to(q) - if initial_state is not None: - S += initial_state - o = torch.zeros_like(v) - mask = torch.triu(torch.ones(BT, BT, dtype=torch.bool, device=q.device), diagonal=1) - for i in range(0, NT): - # [B, HV, BT, ...] - q_i = q[:, :, i] # [B, HV, BT, K] - k_i = k[:, :, i] # [B, HV, BT, K] - u_i = u[:, :, i] # [B, HV, BT, V] - g_i = g[:, :, i] # [B, HV, BT, K] - w_i = w[:, :, i] # [B, HV, BT, K] - # Aqk: per value head (q from qk head, g from value head, k from qk head) - Aqk = torch.zeros(B, HV, BT, BT, dtype=torch.float, device=q.device) - for j in range(BT): - k_j = k[:, :, i, j] - g_j = g[:, :, i, j:j+1, :] - Aqk[..., j] = torch.einsum('... c d, ... d -> ... c', q_i * (g_i - g_j).exp(), k_j) - Aqk = Aqk.masked_fill(mask, 0) - v_i = u_i - w_i @ S - o[:, :, i] = (q_i * g_i.exp()) @ S + Aqk @ v_i - S = S * rearrange(g_i[:, :, -1].exp(), 'b h k -> b h k 1') - S += rearrange((g_i[:, :, -1:] - g_i).exp() * k_i, 'b h c k -> b h k c') @ v_i - if not output_final_state: - S = None - return rearrange(o, 'b h n c d -> b (n c) h d').to(dtype), S diff --git a/build/torch-xpu/ops/kda/wy_fast.py b/build/torch-xpu/ops/kda/wy_fast.py deleted file mode 100644 index 521e5784e0c724addbd65fcc8d46cb4d8792d9bd..0000000000000000000000000000000000000000 --- a/build/torch-xpu/ops/kda/wy_fast.py +++ /dev/null @@ -1,353 +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 -import triton.language as tl - -from ...ops.utils import prepare_chunk_indices -from ...ops.utils.cache import fla_cache_autotune -from ...ops.utils.op import exp2 -from ...utils import autotune_cache_kwargs, check_shared_mem - - -@triton.heuristics({ - 'STORE_QG': lambda args: args['qg'] is not None, - 'STORE_KG': lambda args: args['kg'] is not None, - 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None, -}) -@fla_cache_autotune( - configs=[ - triton.Config({}, num_warps=num_warps, num_stages=num_stages) - for num_warps in [2, 4, 8] - for num_stages in [2, 3, 4] - ], - key=['H', 'HV', 'K', 'V', 'BT', 'BK', 'BV', 'IS_VARLEN'], - **autotune_cache_kwargs, -) -@triton.jit(do_not_specialize=['T']) -def recompute_w_u_fwd_kda_kernel( - q, - k, - qg, - kg, - v, - beta, - w, - u, - A, - gk, - cu_seqlens, - chunk_indices, - T, - H: tl.constexpr, - HV: tl.constexpr, - K: tl.constexpr, - V: tl.constexpr, - BT: tl.constexpr, - BK: tl.constexpr, - BV: tl.constexpr, - STORE_QG: tl.constexpr, - STORE_KG: tl.constexpr, - IS_VARLEN: tl.constexpr, -): - i_t, i_bh = tl.program_id(0), tl.program_id(1) - i_b, i_hv = i_bh // HV, i_bh % HV - i_h = i_hv // (HV // H) - 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.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) - T = eos - bos - else: - bos, eos = i_b * T, i_b * T + T - - k += (bos * H + i_h) * K - v += (bos * HV + i_hv) * V - u += (bos * HV + i_hv) * V - w += (bos * HV + i_hv) * K - gk += (bos * HV + i_hv) * K - beta += bos * HV + i_hv - A += (bos * HV + i_hv) * BT - if STORE_QG: - q += (bos * H + i_h) * K - qg += (bos * HV + i_hv) * K - if STORE_KG: - kg += (bos * HV + i_hv) * K - - p_b = tl.make_block_ptr(beta, (T,), (HV,), (i_t * BT,), (BT,), (0,)) - b_b = tl.load(p_b, boundary_check=(0,)) - - p_A = tl.make_block_ptr(A, (T, BT), (HV*BT, 1), (i_t * BT, 0), (BT, BT), (1, 0)) - b_A = tl.load(p_A, boundary_check=(0, 1)) - - for i_v in range(tl.cdiv(V, BV)): - p_v = tl.make_block_ptr(v, (T, V), (HV*V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) - p_u = tl.make_block_ptr(u, (T, V), (HV*V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) - b_v = tl.load(p_v, boundary_check=(0, 1)) - b_vb = (b_v * b_b[:, None]).to(b_v.dtype) - b_u = tl.dot(b_A, b_vb) - tl.store(p_u, b_u.to(p_u.dtype.element_ty), boundary_check=(0, 1)) - - for i_k in range(tl.cdiv(K, BK)): - p_w = tl.make_block_ptr(w, (T, K), (HV*K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) - p_k = tl.make_block_ptr(k, (T, K), (H*K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) - b_k = tl.load(p_k, boundary_check=(0, 1)) - b_kb = b_k * b_b[:, None] - - p_gk = tl.make_block_ptr(gk, (T, K), (HV*K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) - b_gk = tl.load(p_gk, boundary_check=(0, 1)).to(tl.float32) - b_kb *= exp2(b_gk) - if STORE_QG: - p_q = tl.make_block_ptr(q, (T, K), (H*K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) - p_qg = tl.make_block_ptr(qg, (T, K), (HV*K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) - b_q = tl.load(p_q, boundary_check=(0, 1)) - b_qg = b_q * exp2(b_gk) - tl.store(p_qg, b_qg.to(p_qg.dtype.element_ty), boundary_check=(0, 1)) - if STORE_KG: - last_idx = min(i_t * BT + BT, T) - 1 - o_k = i_k * BK + tl.arange(0, BK) - m_k = o_k < K - b_gn = tl.load(gk + last_idx * HV*K + o_k, mask=m_k, other=0.).to(tl.float32) - b_kg = b_k * tl.where((i_t * BT + tl.arange(0, BT) < T)[:, None], exp2(b_gn[None, :] - b_gk), 0) - p_kg = tl.make_block_ptr(kg, (T, K), (HV*K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) - tl.store(p_kg, b_kg.to(p_kg.dtype.element_ty), boundary_check=(0, 1)) - - b_w = tl.dot(b_A, b_kb.to(b_k.dtype)) - tl.store(p_w, b_w.to(p_w.dtype.element_ty), boundary_check=(0, 1)) - - -@triton.heuristics({ - 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None, -}) -@fla_cache_autotune( - configs=[ - triton.Config({}, num_warps=num_warps, num_stages=num_stages) - for num_warps in [2, 4] - for num_stages in [2, 3, 4] - ], - key=['H', 'HV', 'K', 'V', 'BT', 'BK', 'BV', 'IS_VARLEN'], - **autotune_cache_kwargs, -) -@triton.jit(do_not_specialize=['T']) -def prepare_wy_repr_bwd_kda_kernel( - k, - v, - beta, - gk, - A, - dA, - dw, - du, - dk, - dk2, - dv, - db, - dg, - dg2, - cu_seqlens, - chunk_indices, - T, - H: tl.constexpr, - HV: tl.constexpr, - K: tl.constexpr, - V: tl.constexpr, - BT: tl.constexpr, - BK: tl.constexpr, - BV: tl.constexpr, - IS_VARLEN: tl.constexpr, -): - i_t, i_bh = tl.program_id(0), tl.program_id(1) - i_b, i_hv = i_bh // HV, i_bh % HV - i_h = i_hv // (HV // H) - 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.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) - T = eos - bos - else: - bos, eos = i_b * T, i_b * T + T - - k += (bos * H + i_h) * K - v += (bos * HV + i_hv) * V - beta += bos * HV + i_hv - gk += (bos * HV + i_hv) * K - A += (bos * HV + i_hv) * BT - dA += (bos * HV + i_hv) * BT - dk += (bos * HV + i_hv) * K - dk2 += (bos * HV + i_hv) * K - dw += (bos * HV + i_hv) * K - du += (bos * HV + i_hv) * V - dv += (bos * HV + i_hv) * V - db += bos * HV + i_hv - dg += (bos * HV + i_hv) * K - dg2 += (bos * HV + i_hv) * K - - p_b = tl.make_block_ptr(beta, (T,), (HV,), (i_t * BT,), (BT,), (0,)) - p_db = tl.make_block_ptr(db, (T,), (HV,), (i_t * BT,), (BT,), (0,)) - p_A = tl.make_block_ptr(A, (BT, T), (1, HV*BT), (0, i_t * BT), (BT, BT), (0, 1)) - - b_b = tl.load(p_b, boundary_check=(0,)) - b_db = tl.zeros([BT], dtype=tl.float32) - b_A = tl.load(p_A, boundary_check=(0, 1)) - b_dA = tl.zeros([BT, BT], dtype=tl.float32) - - for i_k in range(tl.cdiv(K, BK)): - p_k = tl.make_block_ptr(k, (T, K), (H*K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) - p_dk = tl.make_block_ptr(dk, (T, K), (HV*K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) - p_dk2 = tl.make_block_ptr(dk2, (T, K), (HV*K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) - p_dw = tl.make_block_ptr(dw, (T, K), (HV*K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) - p_dg = tl.make_block_ptr(dg, (T, K), (HV*K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) - p_dg2 = tl.make_block_ptr(dg2, (T, K), (HV*K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) - - # [BT, BK] - b_k = tl.load(p_k, boundary_check=(0, 1)) - p_gk = tl.make_block_ptr(gk, (T, K), (HV*K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) - b_gk_exp = exp2(tl.load(p_gk, boundary_check=(0, 1))) - b_kbg = b_k * b_b[:, None] * b_gk_exp - b_dw = tl.load(p_dw, boundary_check=(0, 1)) - - b_dA += tl.dot(b_dw, tl.trans(b_kbg).to(b_dw.dtype)) - b_dkbg = tl.dot(b_A, b_dw) - b_dk = b_dkbg * b_gk_exp * b_b[:, None] + tl.load(p_dk, boundary_check=(0, 1)) - b_db += tl.sum(b_dkbg * b_k * b_gk_exp, 1) - b_dg = b_kbg * b_dkbg + tl.load(p_dg, boundary_check=(0, 1)) - - tl.store(p_dk2, b_dk.to(p_dk2.dtype.element_ty), boundary_check=(0, 1)) - tl.store(p_dg2, b_dg.to(p_dg2.dtype.element_ty), boundary_check=(0, 1)) - - for i_v in range(tl.cdiv(V, BV)): - p_v = tl.make_block_ptr(v, (T, V), (HV*V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) - p_dv = tl.make_block_ptr(dv, (T, V), (HV*V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) - p_du = tl.make_block_ptr(du, (T, V), (HV*V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) - b_v = tl.load(p_v, boundary_check=(0, 1)) - b_vb = (b_v * b_b[:, None]).to(b_v.dtype) - b_du = tl.load(p_du, boundary_check=(0, 1)) - b_dA += tl.dot(b_du, tl.trans(b_vb)) - b_dvb = tl.dot(b_A, b_du) - b_dv = b_dvb * b_b[:, None] - b_db += tl.sum(b_dvb * b_v, 1) - tl.store(p_dv, b_dv.to(p_dv.dtype.element_ty), boundary_check=(0, 1)) - - o_t = i_t * BT + tl.arange(0, BT) - m_t = o_t < T - m_A = (o_t[:, None] > o_t[None, :]) & (m_t[:, None] & m_t) - b_dA = tl.where(m_A, b_dA, 0) - b_dA = tl.dot(b_dA.to(b_A.dtype), b_A) - b_dA = tl.dot(b_A, b_dA.to(b_A.dtype)) - - b_dA = tl.where(m_A, -b_dA, 0) - - p_dA = tl.make_block_ptr(dA, (T, BT), (HV*BT, 1), (i_t * BT, 0), (BT, BT), (1, 0)) - tl.store(p_dA, b_dA.to(p_dA.dtype.element_ty), boundary_check=(0, 1)) - tl.store(p_db, b_db.to(p_db.dtype.element_ty), boundary_check=(0,)) - - -def recompute_w_u_fwd( - k: torch.Tensor, - v: torch.Tensor, - beta: torch.Tensor, - A: torch.Tensor, - gk: torch.Tensor, - q: torch.Tensor | None = None, - cu_seqlens: torch.LongTensor | None = None, - chunk_indices: torch.LongTensor | None = None, -) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor | None, torch.Tensor | None]: - B, T, H, K, V = *k.shape, v.shape[-1] - HV = v.shape[2] - BT = A.shape[-1] - BK = 64 - BV = 64 - - if chunk_indices is None and cu_seqlens is not None: - chunk_indices = prepare_chunk_indices(cu_seqlens, BT) - NT = triton.cdiv(T, BT) if cu_seqlens is None else len(chunk_indices) - - w = torch.empty(B, T, HV, K, device=k.device, dtype=k.dtype) - u = torch.empty_like(v) - qg = torch.empty(B, T, HV, K, device=k.device, dtype=k.dtype) if q is not None else None - kg = torch.empty(B, T, HV, K, device=k.device, dtype=k.dtype) - recompute_w_u_fwd_kda_kernel[(NT, B*HV)]( - q=q, - k=k, - qg=qg, - kg=kg, - v=v, - beta=beta, - w=w, - u=u, - A=A, - gk=gk, - cu_seqlens=cu_seqlens, - chunk_indices=chunk_indices, - T=T, - H=H, - HV=HV, - K=K, - V=V, - BT=BT, - BK=BK, - BV=BV, - ) - return w, u, qg, kg - - -def prepare_wy_repr_bwd( - k: torch.Tensor, - v: torch.Tensor, - beta: torch.Tensor, - gk: torch.Tensor, - A: torch.Tensor, - dk: torch.Tensor, - dw: torch.Tensor, - du: torch.Tensor, - dg: torch.Tensor, - cu_seqlens: torch.LongTensor | None = None, - chunk_indices: torch.LongTensor | None = None, -) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: - B, T, H, K, V = *k.shape, v.shape[-1] - HV = v.shape[2] - BT = A.shape[-1] - if chunk_indices is None and cu_seqlens is not None: - chunk_indices = prepare_chunk_indices(cu_seqlens, BT) - NT = triton.cdiv(T, BT) if cu_seqlens is None else len(chunk_indices) - CONST_TILING = 64 if check_shared_mem() else 32 - BK = min(max(triton.next_power_of_2(K), 16), CONST_TILING) - BV = min(max(triton.next_power_of_2(V), 16), CONST_TILING) - - dk2 = torch.empty_like(dk, dtype=torch.float) - dv = torch.empty_like(v) - dg2 = torch.empty_like(gk, dtype=torch.float) - dA = torch.empty_like(A, dtype=torch.float) - db = torch.empty_like(beta, dtype=torch.float) - prepare_wy_repr_bwd_kda_kernel[(NT, B * HV)]( - k=k, - v=v, - beta=beta, - gk=gk, - A=A, - dA=dA, - dw=dw, - du=du, - dk=dk, - dk2=dk2, - dv=dv, - db=db, - dg=dg, - dg2=dg2, - cu_seqlens=cu_seqlens, - chunk_indices=chunk_indices, - T=T, - H=H, - HV=HV, - K=K, - V=V, - BT=BT, - BK=BK, - BV=BV, - ) - dk = dk2 - dg = dg2 - return dk, dv, db, dg, dA diff --git a/build/torch-xpu/ops/lightning_attn/__init__.py b/build/torch-xpu/ops/lightning_attn/__init__.py deleted file mode 100644 index ba7e5c6caa33ad2ad6ee2ba48f28618d613c5f51..0000000000000000000000000000000000000000 --- a/build/torch-xpu/ops/lightning_attn/__init__.py +++ /dev/null @@ -1,14 +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 .chunk import chunk_lightning_attn -from .fused_recurrent import fused_recurrent_lightning_attn - -__all__ = [ - 'chunk_lightning_attn', - 'fused_recurrent_lightning_attn', -] diff --git a/build/torch-xpu/ops/lightning_attn/chunk.py b/build/torch-xpu/ops/lightning_attn/chunk.py deleted file mode 100644 index 728e9a8111fb9d4ba5d18c8ec8f2ec5a31053c7f..0000000000000000000000000000000000000000 --- a/build/torch-xpu/ops/lightning_attn/chunk.py +++ /dev/null @@ -1,85 +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 - -from ...ops.simple_gla.chunk import chunk_simple_gla - - -@torch.compiler.disable -def chunk_lightning_attn( - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - layer_idx: int, - num_layers: int, - scale: float | None = None, - initial_state: torch.Tensor | None = None, - output_final_state: bool = False, - cu_seqlens: torch.LongTensor | None = None, - **kwargs, -) -> tuple[torch.Tensor, torch.Tensor]: - r""" - Args: - q (torch.Tensor): - queries of shape `[B, T, H, K]`. - k (torch.Tensor): - keys of shape `[B, T, H, K]`. - v (torch.Tensor): - values of shape `[B, T, H, V]`. - layer_idx (int): - The index of the current layer. - num_layers (int): - The total number of layers. Both `layer_idx` and `num_layers` are used to compute the decay factor. - scale (Optional[float]): - Scale factor for the attention scores. - If not provided, it will default to `1 / sqrt(K)`. Default: `None`. - initial_state (Optional[torch.Tensor]): - Initial state of shape `[N, H, K, V]` for `N` input sequences. - For equal-length input sequences, `N` equals the batch size `B`. - Default: `None`. - output_final_state (Optional[bool]): - Whether to output the final state of shape `[N, H, K, V]`. Default: `False`. - cu_seqlens (torch.LongTensor): - Cumulative sequence lengths of shape `[N+1]` used for variable-length training, - consistent with the FlashAttention API. - - Returns: - o (torch.Tensor): - Outputs of shape `[B, T, H, V]`. - final_state (torch.Tensor): - Final state of shape `[N, H, K, V]` if `output_final_state=True` else `None`. - """ - if 'head_first' in kwargs: - raise DeprecationWarning( - "head_first has been removed. Inputs must be in `[B, T, H, ...]` format.", - ) - chunk_size = kwargs.pop('chunk_size', None) - if cu_seqlens is not None: - if q.shape[0] != 1: - raise ValueError( - f"The batch size is expected to be 1 rather than {q.shape[0]} when using `cu_seqlens`. " - f"Please flatten variable-length inputs before processing.", - ) - if initial_state is not None and initial_state.shape[0] != len(cu_seqlens) - 1: - raise ValueError( - f"The number of initial states is expected to be equal to the number of input sequences, " - f"i.e., {len(cu_seqlens) - 1} rather than {initial_state.shape[0]}.", - ) - H = q.shape[2] - g_gamma = -(8 / H * (1 - layer_idx / num_layers)) * q.new_tensor(range(H), dtype=torch.float) - return chunk_simple_gla( - q=q, - k=k, - v=v, - scale=scale, - g_gamma=g_gamma, - initial_state=initial_state, - output_final_state=output_final_state, - cu_seqlens=cu_seqlens, - chunk_size=chunk_size, - ) diff --git a/build/torch-xpu/ops/lightning_attn/fused_recurrent.py b/build/torch-xpu/ops/lightning_attn/fused_recurrent.py deleted file mode 100644 index a8efaefcc46371b9dc979bdb908bd6554bbcc196..0000000000000000000000000000000000000000 --- a/build/torch-xpu/ops/lightning_attn/fused_recurrent.py +++ /dev/null @@ -1,86 +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 - -from ...ops.simple_gla.fused_recurrent import fused_recurrent_simple_gla - - -def fused_recurrent_lightning_attn( - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - layer_idx: int, - num_layers: int, - scale: float | None = None, - initial_state: torch.Tensor | None = None, - output_final_state: bool = False, - reverse: bool = False, - cu_seqlens: torch.LongTensor | None = None, - **kwargs, -) -> tuple[torch.Tensor, torch.Tensor]: - r""" - Args: - q (torch.Tensor): - queries of shape `[B, T, H, K]`. - k (torch.Tensor): - keys of shape `[B, T, H, K]`. - v (torch.Tensor): - values of shape `[B, T, H, V]`. - layer_idx (int): - The index of the current layer. - num_layers (int): - The total number of layers. Both `layer_idx` and `num_layers` are used to compute the decay factor. - scale (Optional[float]): - Scale factor for the attention scores. - If not provided, it will default to `1 / sqrt(K)`. Default: `None`. - initial_state (Optional[torch.Tensor]): - Initial state of shape `[N, H, K, V]` for `N` input sequences. - For equal-length input sequences, `N` equals the batch size `B`. - Default: `None`. - output_final_state (Optional[bool]): - Whether to output the final state of shape `[N, H, K, V]`. Default: `False`. - reverse (Optional[bool]): - If `True`, process the state passing in reverse order. Default: `False`. - cu_seqlens (torch.LongTensor): - Cumulative sequence lengths of shape `[N+1]` used for variable-length training, - consistent with the FlashAttention API. - - Returns: - o (torch.Tensor): - Outputs of shape `[B, T, H, V]`. - final_state (torch.Tensor): - Final state of shape `[N, H, K, V]` if `output_final_state=True` else `None`. - """ - if 'head_first' in kwargs: - raise DeprecationWarning( - "head_first has been removed. Inputs must be in `[B, T, H, ...]` format.", - ) - if cu_seqlens is not None: - if q.shape[0] != 1: - raise ValueError( - f"The batch size is expected to be 1 rather than {q.shape[0]} when using `cu_seqlens`. " - f"Please flatten variable-length inputs before processing.", - ) - if initial_state is not None and initial_state.shape[0] != len(cu_seqlens) - 1: - raise ValueError( - f"The number of initial states is expected to be equal to the number of input sequences, " - f"i.e., {len(cu_seqlens) - 1} rather than {initial_state.shape[0]}.", - ) - H = q.shape[2] - g_gamma = -(8 / H * (1 - layer_idx / num_layers)) * q.new_tensor(range(H), dtype=torch.float) - return fused_recurrent_simple_gla( - q=q, - k=k, - v=v, - g_gamma=g_gamma, - scale=scale, - initial_state=initial_state, - output_final_state=output_final_state, - reverse=reverse, - cu_seqlens=cu_seqlens, - ) diff --git a/build/torch-xpu/ops/linear_attn/__init__.py b/build/torch-xpu/ops/linear_attn/__init__.py deleted file mode 100644 index 45ab72080353b5ae73472b8c4cf5e510a3e7c81c..0000000000000000000000000000000000000000 --- a/build/torch-xpu/ops/linear_attn/__init__.py +++ /dev/null @@ -1,16 +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 .chunk import chunk_linear_attn -from .fused_chunk import fused_chunk_linear_attn -from .fused_recurrent import fused_recurrent_linear_attn - -__all__ = [ - 'chunk_linear_attn', - 'fused_chunk_linear_attn', - 'fused_recurrent_linear_attn', -] diff --git a/build/torch-xpu/ops/linear_attn/chunk.py b/build/torch-xpu/ops/linear_attn/chunk.py deleted file mode 100644 index 5f43a16ee45aa1f5958cfb94235675c68676dd10..0000000000000000000000000000000000000000 --- a/build/torch-xpu/ops/linear_attn/chunk.py +++ /dev/null @@ -1,85 +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 - -from ...ops.linear_attn.utils import normalize_with_z_state -from ...ops.simple_gla import chunk_simple_gla - - -@torch.compiler.disable -def chunk_linear_attn( - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - scale: float | None = None, - initial_state: torch.Tensor | tuple | None = None, - output_final_state: bool = False, - normalize: bool = True, - cu_seqlens: torch.LongTensor | None = None, -) -> tuple[torch.Tensor, torch.Tensor]: - r""" - Args: - q (torch.Tensor): - queries of shape `[B, T, H, K]`. - k (torch.Tensor): - keys of shape `[B, T, H, K]`. - v (torch.Tensor): - values of shape `[B, T, H, V]`. - scale (Optional[float]): - Scale factor for the linear attention scores. - If not provided, it will default to `1 / sqrt(K)`. Default: `None`. - initial_state (Optional[torch.Tensor | tuple]): - Initial recurrent state. When `normalize=False`, a tensor of shape `[B, H, K, V]`. - When `normalize=True`, a tuple `(kv_state, z_state)` where `z_state` has shape - `[B, 1, H, K]` (running cumulative key for the denominator). Default: `None`. - output_final_state (Optional[bool]): - Whether to output the final state. Default: `False`. - normalize (bool): - Whether to normalize the output. Default: `True`. - cu_seqlens (torch.LongTensor): - Cumulative sequence lengths of shape `[N+1]` used for variable-length training, - consistent with the FlashAttention API. - - Returns: - o (torch.Tensor): - Outputs of shape `[B, T, H, V]`. - final_state (torch.Tensor | tuple | None): - Mirrors the shape of `initial_state`: a tensor when `normalize=False`, a - `(kv_state, z_state)` tuple when `normalize=True`, or `None` when - `output_final_state=False`. - """ - if normalize and isinstance(initial_state, tuple): - kv_init, z_init = initial_state - else: - kv_init, z_init = initial_state, None - if cu_seqlens is not None: - if q.shape[0] != 1: - raise ValueError( - f"The batch size is expected to be 1 rather than {q.shape[0]} when using `cu_seqlens`. " - f"Please flatten variable-length inputs before processing.", - ) - if kv_init is not None and kv_init.shape[0] != len(cu_seqlens) - 1: - raise ValueError( - f"The number of initial states is expected to be equal to the number of input sequences, " - f"i.e., {len(cu_seqlens) - 1} rather than {kv_init.shape[0]}.", - ) - if scale is None: - scale = k.shape[-1] ** -0.5 - o, final_state = chunk_simple_gla( - q=q, - k=k, - v=v, - scale=scale, - initial_state=kv_init, - output_final_state=output_final_state, - cu_seqlens=cu_seqlens, - ) - if normalize: - o, z_state = normalize_with_z_state(o, q, k, scale, z_init, reverse=False, cu_seqlens=cu_seqlens) - return o, (final_state, z_state) - return o, final_state diff --git a/build/torch-xpu/ops/linear_attn/fused_chunk.py b/build/torch-xpu/ops/linear_attn/fused_chunk.py deleted file mode 100644 index 1ed1c81a5e256fcc7fb185467618c11518b87656..0000000000000000000000000000000000000000 --- a/build/torch-xpu/ops/linear_attn/fused_chunk.py +++ /dev/null @@ -1,85 +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 - -from ...ops.linear_attn.utils import normalize_with_z_state -from ...ops.simple_gla import fused_chunk_simple_gla - - -@torch.compiler.disable -def fused_chunk_linear_attn( - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - scale: float | None = None, - initial_state: torch.Tensor | tuple | None = None, - output_final_state: bool = False, - normalize: bool = True, - cu_seqlens: torch.LongTensor | None = None, -) -> tuple[torch.Tensor, torch.Tensor]: - r""" - Args: - q (torch.Tensor): - queries of shape `[B, T, H, K]`. - k (torch.Tensor): - keys of shape `[B, T, H, K]`. - v (torch.Tensor): - values of shape `[B, T, H, V]`. - scale (Optional[float]): - Scale factor for linear attention scores. - If not provided, it will default to `1 / sqrt(K)`. Default: `None`. - initial_state (Optional[torch.Tensor | tuple]): - Initial recurrent state. When `normalize=False`, a tensor of shape `[B, H, K, V]`. - When `normalize=True`, a tuple `(kv_state, z_state)` where `z_state` has shape - `[B, 1, H, K]` (running cumulative key for the denominator). Default: `None`. - output_final_state (Optional[bool]): - Whether to output the final state. Default: `False`. - normalize (bool): - Whether to normalize the output. Default: `True`. - cu_seqlens (torch.LongTensor): - Cumulative sequence lengths of shape `[N+1]` used for variable-length training, - consistent with the FlashAttention API. - - Returns: - o (torch.Tensor): - Outputs of shape `[B, T, H, V]`. - final_state (torch.Tensor | tuple | None): - Mirrors the shape of `initial_state`: a tensor when `normalize=False`, a - `(kv_state, z_state)` tuple when `normalize=True`, or `None` when - `output_final_state=False`. - """ - if normalize and isinstance(initial_state, tuple): - kv_init, z_init = initial_state - else: - kv_init, z_init = initial_state, None - if cu_seqlens is not None: - if q.shape[0] != 1: - raise ValueError( - f"The batch size is expected to be 1 rather than {q.shape[0]} when using `cu_seqlens`. " - f"Please flatten variable-length inputs before processing.", - ) - if kv_init is not None and kv_init.shape[0] != len(cu_seqlens) - 1: - raise ValueError( - f"The number of initial states is expected to be equal to the number of input sequences, " - f"i.e., {len(cu_seqlens) - 1} rather than {kv_init.shape[0]}.", - ) - if scale is None: - scale = q.shape[-1] ** -0.5 - o, final_state = fused_chunk_simple_gla( - q=q, - k=k, - v=v, - scale=scale, - initial_state=kv_init, - output_final_state=output_final_state, - cu_seqlens=cu_seqlens, - ) - if normalize: - o, z_state = normalize_with_z_state(o, q, k, scale, z_init, reverse=False, cu_seqlens=cu_seqlens) - return o, (final_state, z_state) - return o, final_state diff --git a/build/torch-xpu/ops/linear_attn/fused_recurrent.py b/build/torch-xpu/ops/linear_attn/fused_recurrent.py deleted file mode 100644 index ca01a22a332481935d7d54b30699a6e8f2cb04d2..0000000000000000000000000000000000000000 --- a/build/torch-xpu/ops/linear_attn/fused_recurrent.py +++ /dev/null @@ -1,88 +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 - -from ...ops.linear_attn.utils import normalize_with_z_state -from ...ops.simple_gla.fused_recurrent import fused_recurrent_simple_gla - - -def fused_recurrent_linear_attn( - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - scale: float | None = None, - initial_state: torch.Tensor | tuple | None = None, - output_final_state: bool = False, - reverse: bool = False, - normalize: bool = False, - cu_seqlens: torch.LongTensor | None = None, -) -> tuple[torch.Tensor, torch.Tensor]: - r""" - Args: - q (torch.Tensor): - queries of shape `[B, T, H, K]`. - k (torch.Tensor): - keys of shape `[B, T, H, K]`. - v (torch.Tensor): - values of shape `[B, T, H, V]`. - scale (Optional[float]): - Scale factor for linear attention scores. - If not provided, it will default to `1 / sqrt(K)`. Default: `None`. - initial_state (Optional[torch.Tensor | tuple]): - Initial recurrent state. When `normalize=False`, a tensor of shape `[B, H, K, V]`. - When `normalize=True`, a tuple `(kv_state, z_state)` where `z_state` has shape - `[B, 1, H, K]` (running cumulative key for the denominator). Default: `None`. - output_final_state (Optional[bool]): - Whether to output the final state. Default: `False`. - reverse (bool): - Whether to scan in reverse order. Default: `False`. - normalize (bool): - Whether to normalize the output. Default: `False`. - cu_seqlens (torch.LongTensor): - Cumulative sequence lengths of shape `[N+1]` used for variable-length training, - consistent with the FlashAttention API. - - Returns: - o (torch.Tensor): - Outputs of shape `[B, T, H, V]`. - final_state (torch.Tensor | tuple | None): - Mirrors the shape of `initial_state`: a tensor when `normalize=False`, a - `(kv_state, z_state)` tuple when `normalize=True`, or `None` when - `output_final_state=False`. - """ - if normalize and isinstance(initial_state, tuple): - kv_init, z_init = initial_state - else: - kv_init, z_init = initial_state, None - if cu_seqlens is not None: - if q.shape[0] != 1: - raise ValueError( - f"The batch size is expected to be 1 rather than {q.shape[0]} when using `cu_seqlens`. " - f"Please flatten variable-length inputs before processing.", - ) - if kv_init is not None and kv_init.shape[0] != len(cu_seqlens) - 1: - raise ValueError( - f"The number of initial states is expected to be equal to the number of input sequences, " - f"i.e., {len(cu_seqlens) - 1} rather than {kv_init.shape[0]}.", - ) - if scale is None: - scale = q.shape[-1] ** -0.5 - o, final_state = fused_recurrent_simple_gla( - q=q, - k=k, - v=v, - scale=scale, - initial_state=kv_init, - output_final_state=output_final_state, - reverse=reverse, - cu_seqlens=cu_seqlens, - ) - if normalize: - o, z_state = normalize_with_z_state(o, q, k, scale, z_init, reverse, cu_seqlens) - return o, (final_state, z_state) - return o, final_state diff --git a/build/torch-xpu/ops/linear_attn/naive.py b/build/torch-xpu/ops/linear_attn/naive.py deleted file mode 100644 index 806f15942e0083f0631f42e4fffe39b3362130e3..0000000000000000000000000000000000000000 --- a/build/torch-xpu/ops/linear_attn/naive.py +++ /dev/null @@ -1,76 +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 -from einops import rearrange - - -def naive_recurrent_linear_attn( - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - scale: float | None = None, - initial_state: torch.Tensor | tuple | None = None, - output_final_state: bool = False, - normalize: bool = False, -) -> tuple[torch.Tensor, torch.Tensor]: - dtype = q.dtype - if scale is None: - scale = q.shape[-1] ** -0.5 - B, T, H, K, V = *q.shape, v.shape[-1] - if normalize and isinstance(initial_state, tuple): - kv_init, z_init = initial_state - else: - kv_init, z_init = initial_state, None - q, k, v = map(lambda x: x.to(torch.float32), (q, k, v)) - o = torch.empty_like(v) - - S = torch.zeros((B, H, K, V), device=q.device, dtype=torch.float32) - if kv_init is not None: - S = S + kv_init - for t in range(T): - S = S + torch.einsum('b h k, b h v -> b h k v', k[:, t], v[:, t]) - o[:, t] = torch.einsum('b h k v, b h k -> b h v', S, q[:, t] * scale) - if normalize: - k_cum = k.cumsum(1) if z_init is None else k.cumsum(1) + z_init - o = o / ((q * scale * k_cum).sum(-1, keepdim=True) + 1e-10) - if output_final_state: - return o.to(dtype), (S, k_cum[:, -1:]) - return o.to(dtype), None - return o.to(dtype), S if output_final_state else None - - -def naive_chunk_linear_attn( - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - scale: float | None = None, - normalize: bool = False, -) -> tuple[torch.Tensor, torch.Tensor]: - if scale is None: - scale = q.shape[-1] ** -0.5 - chunk_size = 64 - q = rearrange(q, 'b (n c) h d -> b h n c d', c=chunk_size) * scale - k = rearrange(k, 'b (n c) h d -> b h n c d', c=chunk_size) - v = rearrange(v, 'b (n c) h d -> b h n c d', c=chunk_size) - kv = k.transpose(-1, -2) @ v - kv = kv.cumsum(2) - kv = torch.cat([torch.zeros_like(kv[:, :, :1]), kv[:, :, :-1]], dim=2) - inter = q @ kv - intra = (( - q @ k.transpose(-1, -2)).masked_fill_( - torch.triu(torch.ones(chunk_size, chunk_size, dtype=bool, device=q.device), diagonal=1), - 0, - )) @ v - o = inter + intra - o = rearrange(o, 'b h n c d -> b (n c) h d') - if normalize: - q = rearrange(q, 'b h n c d -> b (n c) h d') - k = rearrange(k, 'b h n c d -> b (n c) h d') - k_cum = k.cumsum(1) - o = o / ((q * k_cum).sum(-1, keepdim=True) + 1e-10) - return o diff --git a/build/torch-xpu/ops/linear_attn/utils.py b/build/torch-xpu/ops/linear_attn/utils.py deleted file mode 100644 index 9175d78f8ba17bf6d7710fdfab6979af089f7488..0000000000000000000000000000000000000000 --- a/build/torch-xpu/ops/linear_attn/utils.py +++ /dev/null @@ -1,106 +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 - -from ...ops.utils.cumsum import chunk_global_cumsum - - -class StatefulNormalizeFunction(torch.autograd.Function): - """Normalized linear-attention divide step with a hand-written backward. - - Forward computes the running key-cumsum z_t and divides ``o`` by ````. - Backward mirrors the pattern used by simple_gla / gla / gated_delta_rule - """ - - @staticmethod - def forward(ctx, o, q, k, scale, z_init, reverse, cu_seqlens): - k_cum = chunk_global_cumsum(k, reverse=reverse, cu_seqlens=cu_seqlens, output_dtype=k.dtype) - if z_init is not None: - if cu_seqlens is not None: - seq_lens = (cu_seqlens[1:] - cu_seqlens[:-1]).to(torch.long) - z_init_b = z_init.squeeze(1).repeat_interleave(seq_lens, dim=0).unsqueeze(0) - else: - z_init_b = z_init - k_cum = k_cum + z_init_b - denom = (q * scale * k_cum).sum(-1, keepdim=True) + 1e-10 - o_out = o / denom - if cu_seqlens is not None: - idx = (cu_seqlens[:-1] if reverse else cu_seqlens[1:] - 1).to(torch.long) - z_state_out = k_cum[0, idx].unsqueeze(1).contiguous() - else: - z_state_out = (k_cum[:, :1] if reverse else k_cum[:, -1:]).contiguous() - - ctx.save_for_backward(o, q, k_cum, denom) - ctx.scale = scale - ctx.reverse = reverse - ctx.cu_seqlens = cu_seqlens - ctx.has_z_init = z_init is not None - ctx.k_dtype = k.dtype - return o_out, z_state_out - - @staticmethod - def backward(ctx, do, dz_state): - o, q, k_cum, denom = ctx.saved_tensors - scale = ctx.scale - reverse = ctx.reverse - cu_seqlens = ctx.cu_seqlens - - d_o_in = do / denom - ddenom = -(do * o).sum(-1, keepdim=True) / (denom * denom) - - dq = ddenom * scale * k_cum - dk_cum = ddenom * scale * q - - if dz_state is not None: - if cu_seqlens is not None: - idx = (cu_seqlens[:-1] if reverse else cu_seqlens[1:] - 1).to(torch.long) - dk_cum[0].index_add_(0, idx, dz_state.squeeze(1).to(dk_cum.dtype)) - elif reverse: - dk_cum[:, :1].add_(dz_state.to(dk_cum.dtype)) - else: - dk_cum[:, -1:].add_(dz_state.to(dk_cum.dtype)) - - dz_init = None - if ctx.has_z_init: - if cu_seqlens is not None: - seq_lens = (cu_seqlens[1:] - cu_seqlens[:-1]).to(torch.long) - seg_id = torch.repeat_interleave( - torch.arange(seq_lens.numel(), device=dk_cum.device), seq_lens, - ) - flat = dk_cum[0] - dz_init = torch.zeros( - seq_lens.numel(), flat.shape[1], flat.shape[2], - device=flat.device, dtype=flat.dtype, - ) - dz_init.index_add_(0, seg_id, flat) - dz_init = dz_init.unsqueeze(1) - else: - dz_init = dk_cum.sum(dim=1, keepdim=True) - - dk = chunk_global_cumsum( - dk_cum, reverse=not reverse, cu_seqlens=cu_seqlens, output_dtype=ctx.k_dtype, - ) - return d_o_in, dq, dk, None, dz_init, None, None - - -def normalize_with_z_state( - o: torch.Tensor, - q: torch.Tensor, - k: torch.Tensor, - scale: float, - z_init: torch.Tensor | None, - reverse: bool, - cu_seqlens: torch.LongTensor | None, -) -> tuple[torch.Tensor, torch.Tensor]: - """Apply running normalization Z_t = sum_{s<=t} k_s (or sum_{s>=t} for reverse). - - z_init carries the cumulative key from prior chunks; the returned z_state is the - boundary needed to chain into the next chunk. Varlen path broadcasts z_init and - extracts the boundary per-sequence. - """ - return StatefulNormalizeFunction.apply(o, q, k, scale, z_init, reverse, cu_seqlens) diff --git a/build/torch-xpu/ops/log_linear_attn/chunk.py b/build/torch-xpu/ops/log_linear_attn/chunk.py deleted file mode 100644 index 35991283935d5c72ed8fb563134908ac6de080ff..0000000000000000000000000000000000000000 --- a/build/torch-xpu/ops/log_linear_attn/chunk.py +++ /dev/null @@ -1,1923 +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 math -import warnings -from dataclasses import dataclass - -import torch -import torch.nn.functional as F -import triton -import triton.language as tl -from einops import reduce - -from ...ops.utils import chunk_local_cumsum -from ...ops.utils.op import exp -from ...utils import autocast_custom_bwd, autocast_custom_fwd, autotune_cache_kwargs, input_guard - -BLOCK_K = 64 - - -@triton.heuristics( - { - "IS_VARLEN": lambda args: args["cu_seqlens"] is not None, - "USE_INITIAL_STATE": lambda args: args["h0"] is not None, - "STORE_FINAL_STATE": lambda args: args["ht"] is not None, - }, -) -@triton.autotune( - configs=[ - triton.Config({"BK": BLOCK_K}, num_warps=num_warps, num_stages=num_stages) - for num_warps in [4] - for num_stages in [2, 3, 4] - ], - key=["H", "K", "V"], - **autotune_cache_kwargs, -) -@triton.jit(do_not_specialize=["T"]) -def chunkwise_fwd_kernel( - q, - k, - v, - g, - level_scales, - llut, - o, - h0, - ht, - offsets, - new_offsets, - cu_seqlens, - T, - H: tl.constexpr, - K: tl.constexpr, - V: tl.constexpr, - L: tl.constexpr, - BT: tl.constexpr, - BK: tl.constexpr, - L_IN: tl.constexpr, - L_OUT: tl.constexpr, - MIN_LEVEL: tl.constexpr, - MAX_LEVEL: tl.constexpr, - IS_VARLEN: tl.constexpr, - USE_INITIAL_STATE: tl.constexpr, - STORE_FINAL_STATE: tl.constexpr, -): - p_llut = tl.make_block_ptr(llut, (BT, BT), (BT, 1), (0, 0), (BT, BT), (1, 0)) - b_llut = tl.load(p_llut, boundary_check=(0, 1)) - # parallel over sequences and heads - i_k = tl.program_id(0) - i_nh = tl.program_id(1) - i_n, i_h = i_nh // H, i_nh % H - - if IS_VARLEN: - bos, eos = ( - tl.load(cu_seqlens + i_n).to(tl.int64), - tl.load(cu_seqlens + i_n + 1).to(tl.int64), - ) - T = (eos - bos).to(tl.int32) - else: - bos, eos = (i_n * T).to(tl.int64), (i_n * T + T).to(tl.int64) - - o_i = tl.arange(0, BT) - - # For hierarchical masking - num_intra_levels = (tl.log2(float(BT))).to(tl.int32) + 1 - i_idx = o_i[:, None] # BT x 1 - j_idx = o_i[None, :] # 1 x BT - - # This is not great. - # See issue: https://github.com/triton-lang/triton/discussions/1313 - KV_0_CREATED = MIN_LEVEL <= 1 and MAX_LEVEL >= 0 - KV_1_CREATED = MIN_LEVEL <= 2 and MAX_LEVEL >= 0 - KV_2_CREATED = MIN_LEVEL <= 3 and MAX_LEVEL >= 1 - KV_3_CREATED = MIN_LEVEL <= 4 and MAX_LEVEL >= 2 - KV_4_CREATED = MIN_LEVEL <= 5 and MAX_LEVEL >= 3 - KV_5_CREATED = MIN_LEVEL <= 6 and MAX_LEVEL >= 4 - KV_6_CREATED = MIN_LEVEL <= 7 and MAX_LEVEL >= 5 - KV_7_CREATED = MIN_LEVEL <= 8 and MAX_LEVEL >= 6 - KV_8_CREATED = MIN_LEVEL <= 9 and MAX_LEVEL >= 7 - KV_9_CREATED = MIN_LEVEL <= 10 and MAX_LEVEL >= 8 - KV_10_CREATED = MIN_LEVEL <= 11 and MAX_LEVEL >= 9 - KV_11_CREATED = MIN_LEVEL <= 12 and MAX_LEVEL >= 10 - - kv_0 = tl.zeros([BK, V], dtype=tl.float32) - kv_1 = tl.zeros([BK, V], dtype=tl.float32) - kv_2 = tl.zeros([BK, V], dtype=tl.float32) - kv_3 = tl.zeros([BK, V], dtype=tl.float32) - kv_4 = tl.zeros([BK, V], dtype=tl.float32) - kv_5 = tl.zeros([BK, V], dtype=tl.float32) - kv_6 = tl.zeros([BK, V], dtype=tl.float32) - kv_7 = tl.zeros([BK, V], dtype=tl.float32) - kv_8 = tl.zeros([BK, V], dtype=tl.float32) - kv_9 = tl.zeros([BK, V], dtype=tl.float32) - kv_10 = tl.zeros([BK, V], dtype=tl.float32) - kv_11 = tl.zeros([BK, V], dtype=tl.float32) - - offset = 0 # total number to cached tokens - first_chunk_index = 0 # next chunk index to compute - if USE_INITIAL_STATE: - offset = tl.load(offsets + i_n) - - first_chunk_index = offset // BT - - if KV_0_CREATED and (first_chunk_index & 1 > 0): - p_kv_0 = tl.make_block_ptr( - h0 + ((i_n * L_IN + 0) * H + i_h) * K * V, - (K, V), - (V, 1), - (i_k * BK, 0), - (BK, V), - (1, 0), - ) - kv_0 = tl.load(p_kv_0, boundary_check=(0, 1)) - if KV_1_CREATED and (first_chunk_index & 2 > 0): - p_kv_1 = tl.make_block_ptr( - h0 + ((i_n * L_IN + 1) * H + i_h) * K * V, - (K, V), - (V, 1), - (i_k * BK, 0), - (BK, V), - (1, 0), - ) - kv_1 = tl.load(p_kv_1, boundary_check=(0, 1)) - if KV_2_CREATED and (first_chunk_index & 4 > 0): - p_kv_2 = tl.make_block_ptr( - h0 + ((i_n * L_IN + 2) * H + i_h) * K * V, - (K, V), - (V, 1), - (i_k * BK, 0), - (BK, V), - (1, 0), - ) - kv_2 = tl.load(p_kv_2, boundary_check=(0, 1)) - if KV_3_CREATED and (first_chunk_index & 8 > 0): - p_kv_3 = tl.make_block_ptr( - h0 + ((i_n * L_IN + 3) * H + i_h) * K * V, - (K, V), - (V, 1), - (i_k * BK, 0), - (BK, V), - (1, 0), - ) - kv_3 = tl.load(p_kv_3, boundary_check=(0, 1)) - if KV_4_CREATED and (first_chunk_index & 16 > 0): - p_kv_4 = tl.make_block_ptr( - h0 + ((i_n * L_IN + 4) * H + i_h) * K * V, - (K, V), - (V, 1), - (i_k * BK, 0), - (BK, V), - (1, 0), - ) - kv_4 = tl.load(p_kv_4, boundary_check=(0, 1)) - if KV_5_CREATED and (first_chunk_index & 32 > 0): - p_kv_5 = tl.make_block_ptr( - h0 + ((i_n * L_IN + 5) * H + i_h) * K * V, - (K, V), - (V, 1), - (i_k * BK, 0), - (BK, V), - (1, 0), - ) - kv_5 = tl.load(p_kv_5, boundary_check=(0, 1)) - if KV_6_CREATED and (first_chunk_index & 64 > 0): - p_kv_6 = tl.make_block_ptr( - h0 + ((i_n * L_IN + 6) * H + i_h) * K * V, - (K, V), - (V, 1), - (i_k * BK, 0), - (BK, V), - (1, 0), - ) - kv_6 = tl.load(p_kv_6, boundary_check=(0, 1)) - if KV_7_CREATED and (first_chunk_index & 128 > 0): - p_kv_7 = tl.make_block_ptr( - h0 + ((i_n * L_IN + 7) * H + i_h) * K * V, - (K, V), - (V, 1), - (i_k * BK, 0), - (BK, V), - (1, 0), - ) - kv_7 = tl.load(p_kv_7, boundary_check=(0, 1)) - if KV_8_CREATED and (first_chunk_index & 256 > 0): - p_kv_8 = tl.make_block_ptr( - h0 + ((i_n * L_IN + 8) * H + i_h) * K * V, - (K, V), - (V, 1), - (i_k * BK, 0), - (BK, V), - (1, 0), - ) - kv_8 = tl.load(p_kv_8, boundary_check=(0, 1)) - if KV_9_CREATED and (first_chunk_index & 512 > 0): - p_kv_9 = tl.make_block_ptr( - h0 + ((i_n * L_IN + 9) * H + i_h) * K * V, - (K, V), - (V, 1), - (i_k * BK, 0), - (BK, V), - (1, 0), - ) - kv_9 = tl.load(p_kv_9, boundary_check=(0, 1)) - if KV_10_CREATED and (first_chunk_index & 1024 > 0): - p_kv_10 = tl.make_block_ptr( - h0 + ((i_n * L_IN + 10) * H + i_h) * K * V, - (K, V), - (V, 1), - (i_k * BK, 0), - (BK, V), - (1, 0), - ) - kv_10 = tl.load(p_kv_10, boundary_check=(0, 1)) - if KV_11_CREATED and (first_chunk_index & 2048 > 0): - p_kv_11 = tl.make_block_ptr( - h0 + ((i_n * L_IN + 11) * H + i_h) * K * V, - (K, V), - (V, 1), - (i_k * BK, 0), - (BK, V), - (1, 0), - ) - kv_11 = tl.load(p_kv_11, boundary_check=(0, 1)) - - NT = tl.cdiv(T, BT) - output_offset = -1 * (offset % BT) - for i_t in range(NT): - b_h_ptrs = level_scales + ((bos + i_t * BT + i_idx) * H + i_h) * L + b_llut - b_h = tl.load(b_h_ptrs, mask=i_idx >= j_idx) - - p_g = tl.make_block_ptr(g + bos * H + i_h, (T,), (H,), (i_t * BT,), (BT,), (0,)) - p_q = tl.make_block_ptr( - q + bos * K, (T, K), (K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0), - ) - p_k = tl.make_block_ptr( - k + bos * K, (K, T), (1, K), (i_k * BK, i_t * BT), (BK, BT), (0, 1), - ) - p_v = tl.make_block_ptr( - v + (bos * H + i_h) * V, - (T, V), - (H * V, 1), - (i_t * BT, 0), - (BT, V), - (1, 0), - ) - p_o = tl.make_block_ptr( - o + ((bos * H + i_h) * (K // BK) + i_k) * V, - (T, V), - (H * (K // BK) * V, 1), - (i_t * BT + output_offset, 0), - (BT, V), - (1, 0), - ) - - b_g = tl.load(p_g, boundary_check=(0,)) - b_q = tl.load(p_q, boundary_check=(0, 1)) - b_k = tl.load(p_k, boundary_check=(0, 1)) - - m_t = i_t * BT + o_i < T - b_s = tl.where((i_idx >= j_idx) & m_t[:, None] & m_t[None, :], tl.exp(b_g[:, None] - b_g[None, :]), 0) - b_s = (tl.dot(b_q, b_k) * b_s).to(b_q.dtype) * b_h - - b_v = tl.load(p_v, boundary_check=(0, 1)) - b_o = tl.zeros((BT, V), dtype=tl.float32) - if MIN_LEVEL == 0: - b_o += tl.dot(b_s, b_v) - - chunk_index = ( - first_chunk_index + i_t - ) # index of the chunk over the entire sequence, including the offset - - if MIN_LEVEL <= 0 and MAX_LEVEL >= 0: - if chunk_index & 1: - p_l = tl.make_block_ptr( - level_scales + (bos * H + i_h) * L, - (T, L), - (H * L, 1), - (i_t * BT, num_intra_levels), - (BT, 1), - (1, 0), - ) - b_l = tl.load(p_l, boundary_check=(0, 1)) - b_o += tl.dot((b_l * b_q), kv_0.to(b_q.dtype)) * tl.exp(b_g)[:, None] - if MIN_LEVEL <= 1 and MAX_LEVEL >= 1: - if chunk_index & 2: - p_l = tl.make_block_ptr( - level_scales + (bos * H + i_h) * L, - (T, L), - (H * L, 1), - (i_t * BT, num_intra_levels + 1), - (BT, 1), - (1, 0), - ) - b_l = tl.load(p_l, boundary_check=(0, 1)) - b_o += tl.dot((b_l * b_q), kv_1.to(b_q.dtype)) * tl.exp(b_g)[:, None] - if MIN_LEVEL <= 2 and MAX_LEVEL >= 2: - if chunk_index & 4: - p_l = tl.make_block_ptr( - level_scales + (bos * H + i_h) * L, - (T, L), - (H * L, 1), - (i_t * BT, num_intra_levels + 2), - (BT, 1), - (1, 0), - ) - b_l = tl.load(p_l, boundary_check=(0, 1)) - b_o += tl.dot((b_l * b_q), kv_2.to(b_q.dtype)) * tl.exp(b_g)[:, None] - if MIN_LEVEL <= 3 and MAX_LEVEL >= 3: - if chunk_index & 8: - p_l = tl.make_block_ptr( - level_scales + (bos * H + i_h) * L, - (T, L), - (H * L, 1), - (i_t * BT, num_intra_levels + 3), - (BT, 1), - (1, 0), - ) - b_l = tl.load(p_l, boundary_check=(0, 1)) - b_o += tl.dot((b_l * b_q), kv_3.to(b_q.dtype)) * tl.exp(b_g)[:, None] - if MIN_LEVEL <= 4 and MAX_LEVEL >= 4: - if chunk_index & 16: - p_l = tl.make_block_ptr( - level_scales + (bos * H + i_h) * L, - (T, L), - (H * L, 1), - (i_t * BT, num_intra_levels + 4), - (BT, 1), - (1, 0), - ) - b_l = tl.load(p_l, boundary_check=(0, 1)) - b_o += tl.dot((b_l * b_q), kv_4.to(b_q.dtype)) * tl.exp(b_g)[:, None] - if MIN_LEVEL <= 5 and MAX_LEVEL >= 5: - if chunk_index & 32: - p_l = tl.make_block_ptr( - level_scales + (bos * H + i_h) * L, - (T, L), - (H * L, 1), - (i_t * BT, num_intra_levels + 5), - (BT, 1), - (1, 0), - ) - b_l = tl.load(p_l, boundary_check=(0, 1)) - b_o += tl.dot((b_l * b_q), kv_5.to(b_q.dtype)) * tl.exp(b_g)[:, None] - if MIN_LEVEL <= 6 and MAX_LEVEL >= 6: - if chunk_index & 64: - p_l = tl.make_block_ptr( - level_scales + (bos * H + i_h) * L, - (T, L), - (H * L, 1), - (i_t * BT, num_intra_levels + 6), - (BT, 1), - (1, 0), - ) - b_l = tl.load(p_l, boundary_check=(0, 1)) - b_o += tl.dot((b_l * b_q), kv_6.to(b_q.dtype)) * tl.exp(b_g)[:, None] - if MIN_LEVEL <= 7 and MAX_LEVEL >= 7: - if chunk_index & 128: # 8192 - 16384 - p_l = tl.make_block_ptr( - level_scales + (bos * H + i_h) * L, - (T, L), - (H * L, 1), - (i_t * BT, num_intra_levels + 7), - (BT, 1), - (1, 0), - ) - b_l = tl.load(p_l, boundary_check=(0, 1)) - b_o += tl.dot((b_l * b_q), kv_7.to(b_q.dtype)) * tl.exp(b_g)[:, None] - if MIN_LEVEL <= 8 and MAX_LEVEL >= 8: - if chunk_index & 256: - p_l = tl.make_block_ptr( - level_scales + (bos * H + i_h) * L, - (T, L), - (H * L, 1), - (i_t * BT, num_intra_levels + 8), - (BT, 1), - (1, 0), - ) - b_l = tl.load(p_l, boundary_check=(0, 1)) - b_o += tl.dot((b_l * b_q), kv_8.to(b_q.dtype)) * tl.exp(b_g)[:, None] - if MIN_LEVEL <= 9 and MAX_LEVEL >= 9: - if chunk_index & 512: - p_l = tl.make_block_ptr( - level_scales + (bos * H + i_h) * L, - (T, L), - (H * L, 1), - (i_t * BT, num_intra_levels + 9), - (BT, 1), - (1, 0), - ) - b_l = tl.load(p_l, boundary_check=(0, 1)) - b_o += tl.dot((b_l * b_q), kv_9.to(b_q.dtype)) * tl.exp(b_g)[:, None] - if MIN_LEVEL <= 10 and MAX_LEVEL >= 10: - if chunk_index & 1024: - p_l = tl.make_block_ptr( - level_scales + (bos * H + i_h) * L, - (T, L), - (H * L, 1), - (i_t * BT, num_intra_levels + 10), - (BT, 1), - (1, 0), - ) - b_l = tl.load(p_l, boundary_check=(0, 1)) - b_o += tl.dot((b_l * b_q), kv_10.to(b_q.dtype)) * tl.exp(b_g)[:, None] - if MIN_LEVEL <= 11 and MAX_LEVEL >= 11: - if chunk_index & 2048: - p_l = tl.make_block_ptr( - level_scales + (bos * H + i_h) * L, - (T, L), - (H * L, 1), - (i_t * BT, num_intra_levels + 11), - (BT, 1), - (1, 0), - ) - b_l = tl.load(p_l, boundary_check=(0, 1)) - b_o += tl.dot((b_l * b_q), kv_11.to(b_q.dtype)) * tl.exp(b_g)[:, None] - - tl.store(p_o, b_o.to(p_o.dtype.element_ty), boundary_check=(0, 1)) - - if i_t < NT - 1 or T % BT == 0: - # Only apply the state update if the last chunk is a full chunk. - # Otherwise, it needs to be included in the next kernel call. - - # update the recurrent states - last_idx = min((i_t + 1) * BT, T) - 1 - b_g_last = tl.load(g + bos * H + last_idx * H + i_h) - if KV_0_CREATED: - kv_0 *= tl.exp(b_g_last) - if KV_1_CREATED: - kv_1 *= tl.exp(b_g_last) - if KV_2_CREATED: - kv_2 *= tl.exp(b_g_last) - if KV_3_CREATED: - kv_3 *= tl.exp(b_g_last) - if KV_4_CREATED: - kv_4 *= tl.exp(b_g_last) - if KV_5_CREATED: - kv_5 *= tl.exp(b_g_last) - if KV_6_CREATED: - kv_6 *= tl.exp(b_g_last) - if KV_7_CREATED: - kv_7 *= tl.exp(b_g_last) - if KV_8_CREATED: - kv_8 *= tl.exp(b_g_last) - if KV_9_CREATED: - kv_9 *= tl.exp(b_g_last) - if KV_10_CREATED: - kv_10 *= tl.exp(b_g_last) - if KV_11_CREATED: - kv_11 *= tl.exp(b_g_last) - - b_v = (b_v * tl.exp(b_g_last - b_g)[:, None]).to(b_v.dtype) - if MIN_LEVEL <= 1: - kv_0 += tl.dot(b_k, b_v) - elif MIN_LEVEL == 2: - kv_1 += tl.dot(b_k, b_v) - elif MIN_LEVEL == 3: - kv_2 += tl.dot(b_k, b_v) - elif MIN_LEVEL == 4: - kv_3 += tl.dot(b_k, b_v) - elif MIN_LEVEL == 5: - kv_4 += tl.dot(b_k, b_v) - elif MIN_LEVEL == 6: - kv_5 += tl.dot(b_k, b_v) - elif MIN_LEVEL == 7: - kv_6 += tl.dot(b_k, b_v) - elif MIN_LEVEL == 8: - kv_7 += tl.dot(b_k, b_v) - elif MIN_LEVEL == 9: - kv_8 += tl.dot(b_k, b_v) - elif MIN_LEVEL == 10: - kv_9 += tl.dot(b_k, b_v) - elif MIN_LEVEL == 11: - kv_10 += tl.dot(b_k, b_v) - - check_value = (~chunk_index & (chunk_index + 1)) - 1 - - if MIN_LEVEL <= 1 and MAX_LEVEL >= 0: - if check_value & 1: - kv_1 += kv_0 - kv_0 = tl.zeros([BK, V], dtype=tl.float32) - if MIN_LEVEL <= 2 and MAX_LEVEL >= 1: - if check_value & 2: - kv_2 += kv_1 - kv_1 = tl.zeros([BK, V], dtype=tl.float32) - if MIN_LEVEL <= 3 and MAX_LEVEL >= 2: - if check_value & 4: - kv_3 += kv_2 - kv_2 = tl.zeros([BK, V], dtype=tl.float32) - if MIN_LEVEL <= 4 and MAX_LEVEL >= 3: - if check_value & 8: - kv_4 += kv_3 - kv_3 = tl.zeros([BK, V], dtype=tl.float32) - if MIN_LEVEL <= 5 and MAX_LEVEL >= 4: - if check_value & 16: - kv_5 += kv_4 - kv_4 = tl.zeros([BK, V], dtype=tl.float32) - if MIN_LEVEL <= 6 and MAX_LEVEL >= 5: - if check_value & 32: - kv_6 += kv_5 - kv_5 = tl.zeros([BK, V], dtype=tl.float32) - if MIN_LEVEL <= 7 and MAX_LEVEL >= 6: - if check_value & 64: - kv_7 += kv_6 - kv_6 = tl.zeros([BK, V], dtype=tl.float32) - if MIN_LEVEL <= 8 and MAX_LEVEL >= 7: - if check_value & 128: - kv_8 += kv_7 - kv_7 = tl.zeros([BK, V], dtype=tl.float32) - if MIN_LEVEL <= 9 and MAX_LEVEL >= 8: - if check_value & 256: - kv_9 += kv_8 - kv_8 = tl.zeros([BK, V], dtype=tl.float32) - if MIN_LEVEL <= 10 and MAX_LEVEL >= 9: - if check_value & 512: - kv_10 += kv_9 - kv_9 = tl.zeros([BK, V], dtype=tl.float32) - if MIN_LEVEL <= 11 and MAX_LEVEL >= 10: - if check_value & 1024: - kv_11 += kv_10 - kv_10 = tl.zeros([BK, V], dtype=tl.float32) - - chunk_index = offset // BT + T // BT - - if STORE_FINAL_STATE: - if (MIN_LEVEL <= 0 and MAX_LEVEL >= 0) and (chunk_index & 1 > 0): - p_kv = tl.make_block_ptr( - ht + ((i_n * L_OUT + 0) * H + i_h) * K * V, - (K, V), - (V, 1), - (i_k * BK, 0), - (BK, V), - (1, 0), - ) - tl.store(p_kv, kv_0, boundary_check=(0, 1)) - if (MIN_LEVEL <= 1 and MAX_LEVEL >= 1) and (chunk_index & 2 > 0): - p_kv = tl.make_block_ptr( - ht + ((i_n * L_OUT + 1) * H + i_h) * K * V, - (K, V), - (V, 1), - (i_k * BK, 0), - (BK, V), - (1, 0), - ) - tl.store(p_kv, kv_1, boundary_check=(0, 1)) - if (MIN_LEVEL <= 2 and MAX_LEVEL >= 2) and (chunk_index & 4 > 0): - p_kv = tl.make_block_ptr( - ht + ((i_n * L_OUT + 2) * H + i_h) * K * V, - (K, V), - (V, 1), - (i_k * BK, 0), - (BK, V), - (1, 0), - ) - tl.store(p_kv, kv_2, boundary_check=(0, 1)) - if (MIN_LEVEL <= 3 and MAX_LEVEL >= 3) and (chunk_index & 8 > 0): - p_kv = tl.make_block_ptr( - ht + ((i_n * L_OUT + 3) * H + i_h) * K * V, - (K, V), - (V, 1), - (i_k * BK, 0), - (BK, V), - (1, 0), - ) - tl.store(p_kv, kv_3, boundary_check=(0, 1)) - if (MIN_LEVEL <= 4 and MAX_LEVEL >= 4) and (chunk_index & 16 > 0): - p_kv = tl.make_block_ptr( - ht + ((i_n * L_OUT + 4) * H + i_h) * K * V, - (K, V), - (V, 1), - (i_k * BK, 0), - (BK, V), - (1, 0), - ) - tl.store(p_kv, kv_4, boundary_check=(0, 1)) - if (MIN_LEVEL <= 5 and MAX_LEVEL >= 5) and (chunk_index & 32 > 0): - p_kv = tl.make_block_ptr( - ht + ((i_n * L_OUT + 5) * H + i_h) * K * V, - (K, V), - (V, 1), - (i_k * BK, 0), - (BK, V), - (1, 0), - ) - tl.store(p_kv, kv_5, boundary_check=(0, 1)) - if (MIN_LEVEL <= 6 and MAX_LEVEL >= 6) and (chunk_index & 64 > 0): - p_kv = tl.make_block_ptr( - ht + ((i_n * L_OUT + 6) * H + i_h) * K * V, - (K, V), - (V, 1), - (i_k * BK, 0), - (BK, V), - (1, 0), - ) - tl.store(p_kv, kv_6, boundary_check=(0, 1)) - if (MIN_LEVEL <= 7 and MAX_LEVEL >= 7) and (chunk_index & 128 > 0): - p_kv = tl.make_block_ptr( - ht + ((i_n * L_OUT + 7) * H + i_h) * K * V, - (K, V), - (V, 1), - (i_k * BK, 0), - (BK, V), - (1, 0), - ) - tl.store(p_kv, kv_7, boundary_check=(0, 1)) - if (MIN_LEVEL <= 8 and MAX_LEVEL >= 8) and (chunk_index & 256 > 0): - p_kv = tl.make_block_ptr( - ht + ((i_n * L_OUT + 8) * H + i_h) * K * V, - (K, V), - (V, 1), - (i_k * BK, 0), - (BK, V), - (1, 0), - ) - tl.store(p_kv, kv_8, boundary_check=(0, 1)) - if (MIN_LEVEL <= 9 and MAX_LEVEL >= 9) and (chunk_index & 512 > 0): - p_kv = tl.make_block_ptr( - ht + ((i_n * L_OUT + 9) * H + i_h) * K * V, - (K, V), - (V, 1), - (i_k * BK, 0), - (BK, V), - (1, 0), - ) - tl.store(p_kv, kv_9, boundary_check=(0, 1)) - if (MIN_LEVEL <= 10 and MAX_LEVEL >= 10) and (chunk_index & 1024 > 0): - p_kv = tl.make_block_ptr( - ht + ((i_n * L_OUT + 10) * H + i_h) * K * V, - (K, V), - (V, 1), - (i_k * BK, 0), - (BK, V), - (1, 0), - ) - tl.store(p_kv, kv_10, boundary_check=(0, 1)) - if (MIN_LEVEL <= 11 and MAX_LEVEL >= 11) and (chunk_index & 2048 > 0): - p_kv = tl.make_block_ptr( - ht + ((i_n * L_OUT + 11) * H + i_h) * K * V, - (K, V), - (V, 1), - (i_k * BK, 0), - (BK, V), - (1, 0), - ) - tl.store(p_kv, kv_11, boundary_check=(0, 1)) - - tl.store(new_offsets + i_n, (offset // BT) * BT + T) - - -@triton.heuristics({ - "IS_VARLEN": lambda args: args["cu_seqlens"] is not None, -}) -@triton.jit(do_not_specialize=["T"]) -def copy_input_kernel( - q, - k, - v, - g, - level_scales, - cu_seqlens, - q_prev, - k_prev, - v_prev, - g_prev, - level_scales_prev, - offsets, - q_new, - k_new, - v_new, - g_new, - level_scales_new, - T, - H: tl.constexpr, - K: tl.constexpr, - V: tl.constexpr, - L: tl.constexpr, - BT: tl.constexpr, - IS_VARLEN: tl.constexpr, -): - # parallel over sequences and heads - i_nh = tl.program_id(0) - i_n, i_h = i_nh // H, i_nh % H - - if IS_VARLEN: - bos, eos = ( - tl.load(cu_seqlens + i_n).to(tl.int64), - tl.load(cu_seqlens + i_n + 1).to(tl.int64), - ) - T = (eos - bos).to(tl.int32) - else: - bos, eos = (i_n * T).to(tl.int64), (i_n * T + T).to(tl.int64) - - offset = tl.load(offsets + i_n) - input_offset = -1 * (offset % BT) - - NT = tl.cdiv(T, BT) - - for i_t in range(NT): - p_g = tl.make_block_ptr( - g + bos * H + i_h, (T,), (H,), (i_t * BT + input_offset,), (BT,), (0,), - ) - p_q = tl.make_block_ptr( - q + bos * K, (T, K), (K, 1), (i_t * BT + input_offset, 0), (BT, K), (1, 0), - ) - p_k = tl.make_block_ptr( - k + bos * K, (T, K), (K, 1), (i_t * BT + input_offset, 0), (BT, K), (1, 0), - ) - p_v = tl.make_block_ptr( - v + (bos * H + i_h) * V, - (T, V), - (H * V, 1), - (i_t * BT + input_offset, 0), - (BT, V), - (1, 0), - ) - p_g_new = tl.make_block_ptr( - g_new + bos * H + i_h, (T,), (H,), (i_t * BT,), (BT,), (0,), - ) - p_q_new = tl.make_block_ptr( - q_new + bos * K, (T, K), (K, 1), (i_t * BT, 0), (BT, K), (1, 0), - ) - p_k_new = tl.make_block_ptr( - k_new + bos * K, (T, K), (K, 1), (i_t * BT, 0), (BT, K), (1, 0), - ) - p_v_new = tl.make_block_ptr( - v_new + (bos * H + i_h) * V, - (T, V), - (H * V, 1), - (i_t * BT, 0), - (BT, V), - (1, 0), - ) - - b_g = tl.load(p_g, boundary_check=(0,)) - b_q = tl.load(p_q, boundary_check=(0, 1)) - b_k = tl.load(p_k, boundary_check=(0, 1)) - b_v = tl.load(p_v, boundary_check=(0, 1)) - - if i_t == 0: - p_g_prev = tl.make_block_ptr( - g_prev + i_n * BT * H + i_h, (BT,), (H,), (0,), (BT,), (0,), - ) - p_q_prev = tl.make_block_ptr( - q_prev + i_n * BT * K, (BT, K), (K, 1), (0, 0), (BT, K), (1, 0), - ) - p_k_prev = tl.make_block_ptr( - k_prev + i_n * BT * K, (BT, K), (K, 1), (0, 0), (BT, K), (1, 0), - ) - p_v_prev = tl.make_block_ptr( - v_prev + (i_n * BT * H + i_h) * V, - (BT, V), - (H * V, 1), - (0, 0), - (BT, V), - (1, 0), - ) - - b_g += tl.load(p_g_prev, boundary_check=(0,)) - b_q += tl.load(p_q_prev, boundary_check=(0, 1)) - b_k += tl.load(p_k_prev, boundary_check=(0, 1)) - b_v += tl.load(p_v_prev, boundary_check=(0, 1)) - - tl.store(p_g_new, b_g, boundary_check=(0,)) - tl.store(p_q_new, b_q, boundary_check=(0, 1)) - tl.store(p_k_new, b_k, boundary_check=(0, 1)) - tl.store(p_v_new, b_v, boundary_check=(0, 1)) - - for i in range(L): - p_l = tl.make_block_ptr( - level_scales + (bos * H + i_h) * L, - (T, L), - (H * L, 1), - (i_t * BT + input_offset, i), - (BT, 1), - (1, 0), - ) - p_l_new = tl.make_block_ptr( - level_scales_new + (bos * H + i_h) * L, - (T, L), - (H * L, 1), - (i_t * BT, i), - (BT, 1), - (1, 0), - ) - b_l = tl.load(p_l, boundary_check=(0,)) - if i_t == 0: - p_l_prev = tl.make_block_ptr( - level_scales_prev + (i_n * BT * H + i_h) * L, - (BT, L), - (H * L, 1), - (0, i), - (BT, 1), - (1, 0), - ) - b_l += tl.load(p_l_prev, boundary_check=(0,)) - tl.store(p_l_new, b_l, boundary_check=(0,)) - - -@triton.heuristics( - { - "IS_VARLEN": lambda args: args["cu_seqlens"] is not None, - }, -) -@triton.jit(do_not_specialize=["T"]) -def copy_last_chunk_kernel( - q, - k, - v, - g, - level_scales, - cu_seqlens, - q_prev, - k_prev, - v_prev, - g_prev, - level_scales_prev, - offsets, - T, - H: tl.constexpr, - K: tl.constexpr, - V: tl.constexpr, - L: tl.constexpr, - BT: tl.constexpr, - IS_VARLEN: tl.constexpr, -): - # parallel over sequences and heads - i_nh = tl.program_id(0) - i_n, i_h = i_nh // H, i_nh % H - - if IS_VARLEN: - bos, eos = ( - tl.load(cu_seqlens + i_n).to(tl.int64), - tl.load(cu_seqlens + i_n + 1).to(tl.int64), - ) - T = (eos - bos).to(tl.int32) - else: - bos, eos = (i_n * T).to(tl.int64), (i_n * T + T).to(tl.int64) - - seq_offset = (T // BT) * BT - - p_g = tl.make_block_ptr(g + bos * H + i_h, (T,), (H,), (seq_offset,), (BT,), (0,)) - p_q = tl.make_block_ptr( - q + bos * K, (T, K), (K, 1), (seq_offset, 0), (BT, K), (1, 0), - ) - p_k = tl.make_block_ptr( - k + bos * K, (T, K), (K, 1), (seq_offset, 0), (BT, K), (1, 0), - ) - p_v = tl.make_block_ptr( - v + (bos * H + i_h) * V, - (T, V), - (H * V, 1), - (seq_offset, 0), - (BT, V), - (1, 0), - ) - p_g_prev = tl.make_block_ptr( - g_prev + i_n * BT * H + i_h, (BT,), (H,), (0,), (BT,), (0,), - ) - p_q_prev = tl.make_block_ptr( - q_prev + i_n * BT * K, (BT, K), (K, 1), (0, 0), (BT, K), (1, 0), - ) - p_k_prev = tl.make_block_ptr( - k_prev + i_n * BT * K, (BT, K), (K, 1), (0, 0), (BT, K), (1, 0), - ) - p_v_prev = tl.make_block_ptr( - v_prev + (i_n * BT * H + i_h) * V, (BT, V), (H * V, 1), (0, 0), (BT, V), (1, 0), - ) - - tl.store(p_g_prev, tl.load(p_g, boundary_check=(0,)), boundary_check=(0,)) - tl.store(p_q_prev, tl.load(p_q, boundary_check=(0, 1)), boundary_check=(0, 1)) - tl.store(p_k_prev, tl.load(p_k, boundary_check=(0, 1)), boundary_check=(0, 1)) - tl.store(p_v_prev, tl.load(p_v, boundary_check=(0, 1)), boundary_check=(0, 1)) - - for i in range(L): - p_l = tl.make_block_ptr( - level_scales + (bos * H + i_h) * L, - (T, L), - (H * L, 1), - (seq_offset, i), - (BT, 1), - (1, 0), - ) - p_l_prev = tl.make_block_ptr( - level_scales_prev + (i_n * BT * H + i_h) * L, - (BT, L), - (H * L, 1), - (0, i), - (BT, 1), - (1, 0), - ) - tl.store(p_l_prev, tl.load(p_l, boundary_check=(0,)), boundary_check=(0,)) - - -@triton.heuristics({"IS_VARLEN": lambda args: args["cu_seqlens"] is not None}) -@triton.autotune( - configs=[ - triton.Config({"BK": BK}, num_warps=num_warps, num_stages=num_stages) - for BK in [32, 64, 128] - for num_warps in [4] - for num_stages in [2, 3, 4] - ], - key=["H", "K", "V"], - restore_value=["dh", "dg_last"], - **autotune_cache_kwargs, -) -@triton.jit(do_not_specialize=["T"]) -def chunkwise_bwd_kernel_dhg( - do, - q, - g, - l, - h_l, - dh, - dg_last, - ell, - T, - cu_seqlens, - H: tl.constexpr, - K: tl.constexpr, - V: tl.constexpr, - L: tl.constexpr, - BT: tl.constexpr, - BK: tl.constexpr, - NT: tl.constexpr, - IS_VARLEN: tl.constexpr, -): - # parallel over batches and heads - i_k = tl.program_id(0) - i_nh = tl.program_id(1) - i_n, i_h = i_nh // H, i_nh % H - - if IS_VARLEN: - bos, eos = ( - tl.load(cu_seqlens + i_n).to(tl.int64), - tl.load(cu_seqlens + i_n + 1).to(tl.int64), - ) - T = (eos - bos).to(tl.int32) - else: - bos, eos = (i_n * T).to(tl.int64), (i_n * T + T).to(tl.int64) - - b_dh = tl.zeros([BK, V], dtype=tl.float32) - - num_intra_levels = (tl.log2(float(BT))).to(tl.int32) + 1 - - for i_t in range(tl.cdiv(T, BT) - 1, -1, -1): - p_dh = tl.make_block_ptr( - dh + ((i_n * NT + i_t) * H + i_h) * K * V, - (K, V), - (V, 1), - (i_k * BK, 0), - (BK, V), - (1, 0), - ) - b_dh_old = tl.load(p_dh, boundary_check=(0, 1)) - - if (i_t & (1 << ell)) == 0: # store the chunk - tl.store( - p_dh, b_dh.to(p_dh.dtype.element_ty) + b_dh_old, boundary_check=(0, 1), - ) - # if you are about the transition to compute, reset to zeros - if i_t > 0 and ((i_t - 1) & (1 << ell)) > 0: - b_dh = tl.zeros([BK, V], dtype=tl.float32) - if i_t & (1 << ell): - p_h = tl.make_block_ptr( - h_l + ((i_n * NT + i_t) * H + i_h) * K * V, - (K, V), - (V, 1), - (i_k * BK, 0), - (BK, V), - (1, 0), - ) - - b_h = tl.load(p_h, boundary_check=(0, 1)) - p_dg_last = dg_last + i_n * NT * H + i_t * H + i_h - tl.atomic_add(p_dg_last, tl.sum(b_h * (b_dh + b_dh_old))) - last_idx = min((i_t + 1) * BT, T) - 1 - b_g_last = tl.exp(tl.load(g + bos * H + last_idx * H + i_h)) - b_dh *= b_g_last - if i_t & (1 << ell): # compute this chunk - p_g = tl.make_block_ptr( - g + bos * H + i_h, (T,), (H,), (i_t * BT,), (BT,), (0,), - ) - p_q = tl.make_block_ptr( - q + bos * K, (K, T), (1, K), (i_k * BK, i_t * BT), (BK, BT), (0, 1), - ) - p_do = tl.make_block_ptr( - do + (bos * H + i_h) * V, - (T, V), - (H * V, 1), - (i_t * BT, 0), - (BT, V), - (1, 0), - ) - p_l = tl.make_block_ptr( - l + (bos * H + i_h) * L + num_intra_levels + ell, - (T,), - (H * L,), - (i_t * BT,), - (BT,), - (0,), - ) - b_l = tl.load(p_l, boundary_check=(0,)) - b_g = tl.load(p_g, boundary_check=(0,)) - b_q = tl.load(p_q, boundary_check=(0, 1)) - b_q = (b_q * (tl.exp(b_g) * b_l)[None, :]).to(b_q.dtype) - b_do = tl.load(p_do, boundary_check=(0, 1)) - - b_s = tl.dot(b_q, b_do).to(b_q.dtype) - b_dh += b_s - - -@triton.heuristics({"IS_VARLEN": lambda args: args["cu_seqlens"] is not None}) -@triton.autotune( - configs=[ - triton.Config({}, num_warps=num_warps, num_stages=num_stages) - for num_warps in [4] - for num_stages in [2, 3, 4] - ], - key=["H", "K", "V"], - restore_value=["dq", "dg"], - **autotune_cache_kwargs, -) -@triton.jit(do_not_specialize=["T"]) -def chunkwise_bwd_kernel_hdqgl( - do, - q, - k, - v, - g, - l, - h_l, - dq, - dg, - dl, - ell, - T, - cu_seqlens, - H: tl.constexpr, - K: tl.constexpr, - V: tl.constexpr, - L: tl.constexpr, - BT: tl.constexpr, - NT: tl.constexpr, - IS_VARLEN: tl.constexpr, -): - # parallel over batches and heads - i_nh = tl.program_id(0) - i_n, i_h = i_nh // H, i_nh % H - - if IS_VARLEN: - bos, eos = ( - tl.load(cu_seqlens + i_n).to(tl.int64), - tl.load(cu_seqlens + i_n + 1).to(tl.int64), - ) - T = (eos - bos).to(tl.int32) - else: - bos, eos = (i_n * T).to(tl.int64), (i_n * T + T).to(tl.int64) - - b_h = tl.zeros([V, K], dtype=tl.float32) - - num_intra_levels = (tl.log2(float(BT))).to(tl.int32) + 1 - - for i_t in range(tl.cdiv(T, BT)): - p_g = tl.make_block_ptr(g + bos * H + i_h, (T,), (H,), (i_t * BT,), (BT,), (0,)) - if i_t & (1 << ell): # compute and store derivatives - p_do = tl.make_block_ptr( - do + (bos * H + i_h) * V, - (T, V), - (H * V, 1), - (i_t * BT, 0), - (BT, V), - (1, 0), - ) - p_q = tl.make_block_ptr( - q + bos * K, (T, K), (K, 1), (i_t * BT, 0), (BT, K), (1, 0), - ) - p_l = tl.make_block_ptr( - l + (bos * H + i_h) * L + num_intra_levels + ell, - (T,), - (H * L,), - (i_t * BT,), - (BT,), - (0,), - ) - p_dq = tl.make_block_ptr( - dq + (bos * H + i_h) * K, - (T, K), - (H * K, 1), - (i_t * BT, 0), - (BT, K), - (1, 0), - ) - p_dg = tl.make_block_ptr( - dg + bos * H + i_h, (T,), (H,), (i_t * BT,), (BT,), (0,), - ) - p_dl = tl.make_block_ptr( - dl + (bos * H + i_h) * L + num_intra_levels + ell, - (T,), - (H * L,), - (i_t * BT,), - (BT,), - (0,), - ) - p_h = tl.make_block_ptr( - h_l + ((i_n * NT + i_t) * H + i_h) * K * V, - (V, K), - (1, V), - (0, 0), - (V, K), - (0, 1), - ) - - b_do = tl.load(p_do, boundary_check=(0, 1)) - b_q = tl.load(p_q, boundary_check=(0, 1)) - b_g = tl.load(p_g, boundary_check=(0,)) - b_l = tl.load(p_l, boundary_check=(0,)) - - b_dlq = tl.exp(b_g)[:, None] * tl.dot(b_do, b_h.to(b_do.dtype)) - - b_dl = tl.sum(b_dlq * b_q, axis=1) - b_dg = b_l * b_dl - - tl.store(p_h, b_h, boundary_check=(0, 1)) - b_dq_old = tl.load(p_dq, boundary_check=(0, 1)) - tl.store( - p_dq, - (b_l[:, None] * b_dlq).to(p_dq.dtype.element_ty) + b_dq_old, - boundary_check=(0, 1), - ) - tl.store(p_dl, b_dl.to(p_dl.dtype.element_ty), boundary_check=(0,)) - b_dg_old = tl.load(p_dg, boundary_check=(0,)) - tl.store( - p_dg, b_dg.to(p_dg.dtype.element_ty) + b_dg_old, boundary_check=(0,), - ) - if ((i_t + 1) & (1 << ell)) == 0: - b_h = tl.zeros([V, K], dtype=tl.float32) - - last_idx = min((i_t + 1) * BT, T) - 1 - b_g_last = tl.load(g + bos * H + last_idx * H + i_h) - b_h *= tl.exp(b_g_last) - if (i_t & (1 << ell)) == 0: # update the state - p_k = tl.make_block_ptr( - k + bos * K, (T, K), (K, 1), (i_t * BT, 0), (BT, K), (1, 0), - ) - p_v = tl.make_block_ptr( - v + (bos * H + i_h) * V, - (V, T), - (1, H * V), - (0, i_t * BT), - (V, BT), - (0, 1), - ) - b_g = tl.load(p_g, boundary_check=(0,)) - b_k = tl.load(p_k, boundary_check=(0, 1)) - b_v = tl.load(p_v, boundary_check=(0, 1)) - b_k = (b_k * tl.exp(b_g_last - b_g)[:, None]).to(b_k.dtype) - b_h += tl.dot(b_v, b_k) - - -@triton.heuristics({"IS_VARLEN": lambda args: args["cu_seqlens"] is not None}) -@triton.autotune( - configs=[ - triton.Config({}, num_warps=num_warps, num_stages=num_stages) - for num_warps in [4] - for num_stages in [2, 3, 4] - ], - key=["H", "K", "V"], - restore_value=["dk", "dg", "dg_last"], - **autotune_cache_kwargs, -) -@triton.jit(do_not_specialize=["T"]) -def chunkwise_bwd_kernel_dkg( - dh, - k, - v, - g, - dg_last, - dk, - dg, - cu_seqlens, - T, - H: tl.constexpr, - K: tl.constexpr, - V: tl.constexpr, - L: tl.constexpr, - BT: tl.constexpr, - NT: tl.constexpr, - IS_VARLEN: tl.constexpr, -): - i_t, i_nh = tl.program_id(0), tl.program_id(1) - i_n, i_h = i_nh // H, i_nh % H - - if IS_VARLEN: - bos, eos = ( - tl.load(cu_seqlens + i_n).to(tl.int64), - tl.load(cu_seqlens + i_n + 1).to(tl.int64), - ) - T = (eos - bos).to(tl.int32) - else: - bos, eos = (i_n * T).to(tl.int64), (i_n * T + T).to(tl.int64) - - o_i = tl.arange(0, BT) - o_t = i_t * BT + o_i - m_t = o_t < T - - p_dh = tl.make_block_ptr( - dh + ((i_n * NT + i_t) * H + i_h) * K * V, - (V, K), - (1, V), - (0, 0), - (V, K), - (0, 1), - ) - p_g = tl.make_block_ptr(g + bos * H + i_h, (T,), (H,), (i_t * BT,), (BT,), (0,)) - p_k = tl.make_block_ptr(k + bos * K, (T, K), (K, 1), (i_t * BT, 0), (BT, K), (1, 0)) - p_v = tl.make_block_ptr( - v + (bos * H + i_h) * V, - (T, V), - (H * V, 1), - (i_t * BT, 0), - (BT, V), - (1, 0), - ) - p_dk = tl.make_block_ptr( - dk + (bos * H + i_h) * K, (T, K), (H * K, 1), (i_t * BT, 0), (BT, K), (1, 0), - ) - p_dg = tl.make_block_ptr(dg + bos * H + i_h, (T,), (H,), (i_t * BT,), (BT,), (0,)) - - b_dh = tl.load(p_dh, boundary_check=(0, 1)) - b_g = tl.load(p_g, boundary_check=(0,)) - b_v = tl.load(p_v, boundary_check=(0, 1)) - b_k = tl.load(p_k, boundary_check=(0, 1)) - last_idx = min((i_t + 1) * BT, T) - 1 - b_g_last = tl.load(g + bos * H + last_idx * H + i_h) - p_dg_last = dg_last + i_n * NT * H + i_t * H + i_h - b_dg_last = tl.load(p_dg_last) - - b_dg_last *= tl.exp(b_g_last) - b_dk = tl.where(m_t, exp(b_g_last - b_g), 0)[:, None] * tl.dot(b_v, b_dh).to(b_v.dtype) - b_dg = tl.load(p_dg, boundary_check=(0,)) - b_dg -= tl.sum(b_k * b_dk, axis=1) - b_dg_last += tl.sum(b_dk * b_k) - - b_dg = tl.where(o_i < BT - 1, b_dg, b_dg + b_dg_last) - - tl.store(p_dg, b_dg, boundary_check=(0,)) - tl.store(p_dk, b_dk.to(p_dk.dtype.element_ty), boundary_check=(0, 1)) - - -@triton.heuristics({"IS_VARLEN": lambda args: args["cu_seqlens"] is not None}) -@triton.autotune( - configs=[ - triton.Config({}, num_warps=num_warps, num_stages=num_stages) - for num_warps in [4] - for num_stages in [2, 3, 4] - ], - key=["H", "K", "V"], - restore_value=["dv"], - **autotune_cache_kwargs, -) -@triton.jit(do_not_specialize=["T"]) -def chunkwise_bwd_kernel_dv( - dh, - k, - g, - dv, - T, - cu_seqlens, - H: tl.constexpr, - K: tl.constexpr, - V: tl.constexpr, - L: tl.constexpr, - BT: tl.constexpr, - NT: tl.constexpr, - IS_VARLEN: tl.constexpr, -): - i_t, i_nh = tl.program_id(0), tl.program_id(1) - i_n, i_h = i_nh // H, i_nh % H - - if IS_VARLEN: - bos, eos = ( - tl.load(cu_seqlens + i_n).to(tl.int64), - tl.load(cu_seqlens + i_n + 1).to(tl.int64), - ) - T = (eos - bos).to(tl.int32) - else: - bos, eos = (i_n * T).to(tl.int64), (i_n * T + T).to(tl.int64) - - o_t = i_t * BT + tl.arange(0, BT) - m_t = o_t < T - - p_dh = tl.make_block_ptr( - dh + ((i_n * NT + i_t) * H + i_h) * K * V, - (K, V), - (V, 1), - (0, 0), - (K, V), - (1, 0), - ) - p_g = tl.make_block_ptr(g + bos * H + i_h, (T,), (H,), (i_t * BT,), (BT,), (0,)) - p_k = tl.make_block_ptr(k + bos * K, (T, K), (K, 1), (i_t * BT, 0), (BT, K), (1, 0)) - p_dv = tl.make_block_ptr( - dv + (bos * H + i_h) * V, (T, V), (H * V, 1), (i_t * BT, 0), (BT, V), (1, 0), - ) - - last_idx = min((i_t + 1) * BT, T) - 1 - b_g_last = tl.load(g + bos * H + last_idx * H + i_h) - - b_dh = tl.load(p_dh, boundary_check=(0, 1)) - b_g = tl.load(p_g, boundary_check=(0,)) - b_k = tl.load(p_k, boundary_check=(0, 1)) - b_dv = tl.where(m_t, exp(-b_g + b_g_last), 0)[:, None] * tl.dot(b_k, b_dh).to(b_k.dtype) - tl.store(p_dv, b_dv.to(p_dv.dtype.element_ty), boundary_check=(0, 1)) - - -@triton.heuristics({"IS_VARLEN": lambda args: args["cu_seqlens"] is not None}) -@triton.autotune( - configs=[ - triton.Config({}, num_warps=num_warps, num_stages=num_stages) - for num_warps in [4] - for num_stages in [2, 3, 4] - ], - key=["H", "K", "V"], - restore_value=["dl", "dq", "dk", "dv", "dg"], - **autotune_cache_kwargs, -) -@triton.jit(do_not_specialize=["T"]) -def chunkwise_bwd_kernel_diag( - do, - q, - k, - v, - g, - l, - llut, - mask, - dq, - dk, - dv, - dg, - dl, - cu_seqlens, - T, - H: tl.constexpr, - K: tl.constexpr, - V: tl.constexpr, - L: tl.constexpr, - BT: tl.constexpr, - IS_VARLEN: tl.constexpr, -): - p_llut = tl.make_block_ptr(llut, (BT, BT), (BT, 1), (0, 0), (BT, BT), (1, 0)) - b_llut = tl.load(p_llut, boundary_check=(0, 1)) - i_t, i_nh = tl.program_id(0), tl.program_id(1) - i_n, i_h = i_nh // H, i_nh % H - - if IS_VARLEN: - bos, eos = ( - tl.load(cu_seqlens + i_n).to(tl.int64), - tl.load(cu_seqlens + i_n + 1).to(tl.int64), - ) - T = (eos - bos).to(tl.int32) - else: - bos, eos = (i_n * T).to(tl.int64), (i_n * T + T).to(tl.int64) - - o_i = tl.arange(0, BT) - i_idx = o_i[:, None] # BT x 1 - j_idx = o_i[None, :] # 1 x BT - - b_h_ptrs = l + ((bos + i_t * BT + i_idx) * H + i_h) * L + b_llut - b_h = tl.load(b_h_ptrs, mask=i_idx >= j_idx) - - p_g = tl.make_block_ptr(g + bos * H + i_h, (T,), (H,), (i_t * BT,), (BT,), (0,)) - p_q = tl.make_block_ptr(q + bos * K, (K, T), (1, K), (0, i_t * BT), (K, BT), (0, 1)) - p_k = tl.make_block_ptr(k + bos * K, (T, K), (K, 1), (i_t * BT, 0), (BT, K), (1, 0)) - p_v = tl.make_block_ptr( - v + (bos * H + i_h) * V, (V, T), (1, H * V), (0, i_t * BT), (V, BT), (0, 1), - ) - p_do = tl.make_block_ptr( - do + (bos * H + i_h) * V, (T, V), (H * V, 1), (i_t * BT, 0), (BT, V), (1, 0), - ) - p_dg = tl.make_block_ptr(dg + bos * H + i_h, (T,), (H,), (i_t * BT,), (BT,), (0,)) - p_dq = tl.make_block_ptr( - dq + (bos * H + i_h) * K, (T, K), (H * K, 1), (i_t * BT, 0), (BT, K), (1, 0), - ) - p_dk = tl.make_block_ptr( - dk + (bos * H + i_h) * K, (T, K), (H * K, 1), (i_t * BT, 0), (BT, K), (1, 0), - ) - p_dv = tl.make_block_ptr( - dv + (bos * H + i_h) * V, (T, V), (H * V, 1), (i_t * BT, 0), (BT, V), (1, 0), - ) - - b_g = tl.load(p_g, boundary_check=(0,)) - b_q = tl.load(p_q, boundary_check=(0, 1)) - b_k = tl.load(p_k, boundary_check=(0, 1)) - b_v = tl.load(p_v, boundary_check=(0, 1)) - b_do = tl.load(p_do, boundary_check=(0, 1)) - b_dq = tl.load(p_dq, boundary_check=(0, 1)) - b_dk = tl.load(p_dk, boundary_check=(0, 1)) - b_dv = tl.load(p_dv, boundary_check=(0, 1)) - b_dg = tl.load(p_dg, boundary_check=(0,)) - - b_s = (tl.dot(b_k, b_q)).to(b_q.dtype) - # Apply causal and padding masks - m_t = i_t * BT + o_i < T - b_a = tl.where((i_idx >= j_idx) & m_t[:, None] & m_t[None, :], tl.exp(b_g[:, None] - b_g[None, :]), 0) - b_dv += tl.dot((b_s * tl.trans(b_a * b_h)).to(b_do.dtype), b_do) - b_ds = tl.dot(b_do, b_v) * b_a - b_dl = b_ds * tl.trans(b_s) - b_dg += tl.sum(b_dl * b_h, axis=1) - b_dg -= tl.sum(b_dl * b_h, axis=0) - b_ds = (b_ds * b_h).to(b_k.dtype) - b_dq += tl.dot(b_ds, b_k) - b_dk += tl.trans(tl.dot(b_q, b_ds)) - - tl.store(p_dv, b_dv.to(p_dv.dtype.element_ty), boundary_check=(0, 1)) - tl.store(p_dq, b_dq.to(p_dq.dtype.element_ty), boundary_check=(0, 1)) - tl.store(p_dk, b_dk.to(p_dk.dtype.element_ty), boundary_check=(0, 1)) - tl.store(p_dg, b_dg.to(p_dg.dtype.element_ty), boundary_check=(0,)) - - num_intra_levels = (tl.log2(float(BT))).to(tl.int32) + 1 - - for i in range(num_intra_levels): - p_mask = tl.make_block_ptr(mask + i * (BT * BT), (BT, BT), (BT, 1), (0, 0), (BT, BT), (1, 0)) - b_mask = tl.load(p_mask, boundary_check=(0, 1)) - dl_i = tl.sum(tl.where(b_mask == 1, b_dl, 0), axis=1) - p_dl_i = tl.make_block_ptr(dl + (bos * H + i_h) * L + i, (T,), (H * L,), (i_t * BT,), (BT,), (0,)) - tl.store(p_dl_i, dl_i, boundary_check=(0,)) - - -def construct_binary_level_mask(level, T): - if level == 0: - return torch.diag(torch.ones(T, dtype=torch.bool)) - - indices = torch.cartesian_prod(torch.arange(T), torch.arange(T)) - - mask = torch.where( - torch.logical_and( - torch.logical_and( - indices[:, 0] % (1 << level) >= (1 << (level - 1)), - indices[:, 1] + (1 << (level - 1)) - >= indices[:, 0] - (indices[:, 0] % (1 << (level - 1))), - ), - indices[:, 1] < indices[:, 0] - (indices[:, 0] % (1 << (level - 1))), - ).view(T, T), - 1, - 0, - ) - - return mask - - -def level_lut(BT, device): - lut = torch.zeros((BT, BT), dtype=torch.int32, device=device) - for level in range(1, ceil_log(BT, 2) + 1): - mask = construct_binary_level_mask(level, BT).to(device) - lut = torch.where(mask.to(torch.bool), level, lut) - return lut - - -def masks(BT, device): - masks = [] - for level in range(0, ceil_log(BT, 2) + 1): - mask = construct_binary_level_mask(level, BT).to(device).to(torch.int32) - masks.append(mask) - return torch.stack(masks) - - -def ceil_div(x: int, y: int) -> int: - return math.ceil(x / y) - - -def ceil_log(x: int, b: int) -> int: - return math.ceil(math.log(x, b)) - - -@dataclass -class LogLinearAttentionState: - ht: torch.Tensor - offsets: torch.Tensor - q_prev: torch.Tensor - k_prev: torch.Tensor - v_prev: torch.Tensor - g_prev: torch.Tensor - level_scales_prev: torch.Tensor - - -class ChunkLogLinearAttentionFunction(torch.autograd.Function): - - @staticmethod - @input_guard - @autocast_custom_fwd - def forward( - ctx, - q, - k, - v, - g, - level_scales, - initial_state, - output_final_state, - cu_seqlens, - ): - B, T, G, K = k.shape - _, _, H, V = v.shape - _, _, _, L = level_scales.shape - - if G != 1: - raise ValueError("Group dimension must be 1.") - - if not math.log2(V).is_integer(): - raise ValueError( - "Head dimension must be a power of two. Please pad the head dimension to the next power of two.", - ) - - if K % BLOCK_K != 0: - raise ValueError(f"State dimension must be divisible by {BLOCK_K}.") - - if triton.__version__ > "3.2.0": - warnings.warn("Triton>3.2.0 detected, which is known to have worse performance. " - "For optimal performance, it is recommended to install Triton==3.2.0 (if possible).") - - BT = 64 # chunk size - - h0 = initial_state.ht if initial_state is not None else None - offsets = initial_state.offsets if initial_state is not None else None - - if cu_seqlens is None: - NT = ceil_div(T + (torch.max(offsets) if offsets is not None else 0), BT) - MAX_LEVEL = ceil_log(NT, 2) - 1 - else: - NT = max( - [ - ceil_div( - cu_seqlens[i + 1] - - cu_seqlens[i] - + (offsets[i] if offsets is not None else 0), - BT, - ) - for i in range(len(cu_seqlens) - 1) - ], - ) - MAX_LEVEL = ceil_log(NT, 2) - 1 - B = len(cu_seqlens) - 1 - - if MAX_LEVEL > 10: - raise ValueError("Sequence length must be less than 2**17") - - S0 = B if cu_seqlens is None else 1 - o = torch.zeros( - (S0, T, H, (K // BLOCK_K), V), - dtype=v.dtype, - device=v.device, - ) - - if initial_state is not None: - if cu_seqlens is not None: - cu_seqlens = cu_seqlens + F.pad(torch.cumsum(offsets % BT), (1, 0)) - else: - assert (offsets == offsets[0]).all() - T += offsets[0].item() % BT - S1 = cu_seqlens[-1] if cu_seqlens is not None else T - q_new = torch.zeros((S0, S1, G, K), dtype=q.dtype, device=q.device) - k_new = torch.zeros((S0, S1, G, K), dtype=k.dtype, device=k.device) - v_new = torch.zeros((S0, S1, H, V), dtype=v.dtype, device=v.device) - g_new = torch.zeros((S0, S1, H), dtype=g.dtype, device=g.device) - level_scales_new = torch.zeros((S0, S1, H, L), dtype=level_scales.dtype, device=level_scales.device) - - copy_input_kernel[(B * H,)]( - q=q, - k=k, - v=v, - g=g, - level_scales=level_scales, - cu_seqlens=cu_seqlens, - q_prev=initial_state.q_prev, - k_prev=initial_state.k_prev, - v_prev=initial_state.v_prev, - g_prev=initial_state.g_prev, - level_scales_prev=initial_state.level_scales_prev, - q_new=q_new, - k_new=k_new, - v_new=v_new, - g_new=g_new, - level_scales_new=level_scales_new, - offsets=offsets, - T=T, - H=H, - K=K, - V=V, - L=L, - BT=BT, - ) - q = q_new - k = k_new - v = v_new - g = g_new - level_scales = level_scales_new - - # Store one extra level (MAX_LEVEL + 2) in case the length is multiple of 2 - ht = ( - torch.zeros((B, MAX_LEVEL + 2, H, K, V), dtype=torch.float, device=v.device) - if output_final_state - else None - ) - - new_offsets = torch.zeros((B,), dtype=torch.int32, device=v.device) - g = chunk_local_cumsum(g, chunk_size=BT, cu_seqlens=cu_seqlens) - - def grid(meta): - return (triton.cdiv(K, meta["BK"]), B * H) - - l_in = h0.shape[1] if initial_state is not None else None - l_out = ht.shape[1] if output_final_state else None - - ctx.llut = level_lut(BT, v.device) - - chunkwise_fwd_kernel[grid]( - q=q, - k=k, - v=v, - g=g, - level_scales=level_scales, - llut=ctx.llut, - o=o, - h0=h0, - ht=ht, - offsets=offsets, - new_offsets=new_offsets, - cu_seqlens=cu_seqlens, - T=T, - H=H, - K=K, - V=V, - L=L, - BT=BT, - L_IN=l_in, - L_OUT=l_out, - MIN_LEVEL=0, - MAX_LEVEL=MAX_LEVEL, - ) - - ctx.save_for_backward(q, k, v, g, level_scales, initial_state, cu_seqlens) - ctx.chunk_size = BT - - if output_final_state: - q_prev = torch.zeros((B, BT, G, K), dtype=q.dtype, device=q.device) - k_prev = torch.zeros((B, BT, G, K), dtype=k.dtype, device=k.device) - v_prev = torch.zeros((B, BT, H, V), dtype=v.dtype, device=v.device) - g_prev = torch.zeros((B, BT, H), dtype=g.dtype, device=g.device) - level_scales_prev = torch.zeros((B, BT, H, L), dtype=level_scales.dtype, device=level_scales.device) - - copy_last_chunk_kernel[(B * H,)]( - q=q, - k=k, - v=v, - g=g, - level_scales=level_scales, - cu_seqlens=cu_seqlens, - q_prev=q_prev, - k_prev=k_prev, - v_prev=v_prev, - g_prev=g_prev, - level_scales_prev=level_scales_prev, - offsets=new_offsets, - T=T, - H=H, - K=K, - V=V, - L=L, - BT=BT, - ) - - final_state = LogLinearAttentionState( - ht=ht, - offsets=new_offsets, - q_prev=q_prev, - k_prev=k_prev, - v_prev=v_prev, - g_prev=g_prev, - level_scales_prev=level_scales_prev, - ) - return o.sum(dim=-2), final_state - - return o.sum(dim=-2), None - - @staticmethod - @input_guard - @autocast_custom_bwd - def backward(ctx, do, dht): - if triton.__version__ < "3.1.0": - raise ValueError("Triton>=3.1.0 is required") - - q, k, v, g, level_scales, initial_state, cu_seqlens = ctx.saved_tensors - chunk_size = ctx.chunk_size - llut = ctx.llut - mask = masks(chunk_size, v.device) - - if initial_state is not None: - raise NotImplementedError( - "Backward pass is not implemented for log-linear attention with a prefilled kernel.", - ) - - B, T, G, K = k.shape - assert G == 1, "Multi-head attention is not supported" - _, _, H, V = v.shape - _, _, _, L = level_scales.shape - BT = chunk_size - if cu_seqlens is not None: - NT = max( - [ - ceil_div(cu_seqlens[i + 1] - cu_seqlens[i], BT) - for i in range(len(cu_seqlens) - 1) - ], - ) - else: - NT = ceil_div(T, BT) - - if cu_seqlens is not None: - B = len(cu_seqlens) - 1 - - dh = torch.zeros((B, NT, H, K, V), dtype=v.dtype, device=v.device) - dq = torch.zeros((B if cu_seqlens is None else 1, T, H, K), dtype=v.dtype, device=v.device) - dk = torch.zeros((B if cu_seqlens is None else 1, T, H, K), dtype=v.dtype, device=v.device) - dv = torch.zeros_like(v) - dg = torch.zeros(g.shape, dtype=torch.float, device=v.device) - dl = torch.zeros(level_scales.shape, dtype=torch.float, device=v.device) - h_l = torch.zeros((B, NT, H, K, V), dtype=torch.float, device=v.device) - dg_last = torch.zeros((B, NT, H), dtype=torch.float, device=v.device) - do = do.to(v.dtype) - - grid = (B * H,) - - def grid_f(meta): - return (triton.cdiv(K, meta["BK"]), B * H) - - grid_t = (NT, B * H) - - num_inter_chunk_levels = ceil_log(NT, 2) - for ell in range(num_inter_chunk_levels - 1, -1, -1): - chunkwise_bwd_kernel_hdqgl[grid]( - do=do, - q=q, - k=k, - v=v, - g=g, - l=level_scales, - h_l=h_l, - dq=dq, - dg=dg, - dl=dl, - cu_seqlens=cu_seqlens, - ell=ell, - T=T, - H=H, - K=K, - V=V, - L=L, - BT=BT, - NT=NT, - ) - chunkwise_bwd_kernel_dhg[grid_f]( - do=do, - q=q, - g=g, - l=level_scales, - h_l=h_l, - dh=dh, - dg_last=dg_last, - cu_seqlens=cu_seqlens, - ell=ell, - T=T, - H=H, - K=K, - V=V, - L=L, - BT=BT, - NT=NT, - ) - - chunkwise_bwd_kernel_dkg[grid_t]( - dh=dh, - k=k, - v=v, - g=g, - dg_last=dg_last, - dk=dk, - dg=dg, - cu_seqlens=cu_seqlens, - T=T, - H=H, - K=K, - V=V, - L=L, - BT=BT, - NT=NT, - ) - - chunkwise_bwd_kernel_dv[grid_t]( - dh=dh, - k=k, - g=g, - dv=dv, - cu_seqlens=cu_seqlens, - T=T, - H=H, - K=K, - V=V, - L=L, - BT=BT, - NT=NT, - ) - - chunkwise_bwd_kernel_diag[grid_t]( - do=do, - q=q, - k=k, - v=v, - g=g, - l=level_scales, - llut=llut, - mask=mask, - dq=dq, - dk=dk, - dv=dv, - dg=dg, - dl=dl, - cu_seqlens=cu_seqlens, - T=T, - H=H, - K=K, - V=V, - L=L, - BT=BT, - ) - - dg = chunk_local_cumsum(dg, chunk_size=chunk_size, reverse=True, cu_seqlens=cu_seqlens).to(g.dtype) - - dq = reduce(dq, "b t (g h) k -> b t g k", "sum", g=G, h=H // G) - dk = reduce(dk, "b t (g h) k -> b t g k", "sum", g=G, h=H // G) - return dq, dk, dv, dg, dl, None, None, None - - -@torch.compiler.disable -def chunk_log_linear_attn( - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - g: torch.Tensor, - level_scales: torch.Tensor, - initial_state: LogLinearAttentionState | None = None, - output_final_state: bool = False, - cu_seqlens: torch.LongTensor | None = None, -) -> tuple[torch.Tensor, torch.Tensor]: - r""" - Args: - q (torch.Tensor): - queries of shape `[B, T, H, K]`. - k (torch.Tensor): - keys of shape `[B, T, H, K]`. - v (torch.Tensor): - values of shape `[B, T, H, V]`. - g (torch.Tensor): - Forget gates of shape `[B, T, H]`. - level_scales (torch.Tensor): - Scales for each level of shape `[B, T, H, L]`. - initial_state (Optional[LogLinearAttentionState]): - Initial state of shape `[N, H, K, V]` for `N` input sequences. - For equal-length input sequences, `N` equals the batch size `B`. - Default: `None`. - output_final_state (Optional[bool]): - Whether to output the final state of shape `[N, H, K, V]`. Default: `False`. - cu_seqlens (torch.LongTensor): - Cumulative sequence lengths of shape `[N+1]` used for variable-length training, - consistent with the FlashAttention API. - - Returns: - o (torch.Tensor): - Outputs of shape `[B, T, H, V]`. - final_state (torch.Tensor): - Final state of type `LogLinearAttentionState` if `output_final_state=True` else `None`. - - """ - if cu_seqlens is not None: - if q.shape[0] != 1: - raise ValueError( - f"The batch size is expected to be 1 rather than {q.shape[0]} when using `cu_seqlens`." - f"Please flatten variable-length inputs before processing.", - ) - - o, final_state = ChunkLogLinearAttentionFunction.apply( - q, - k, - v, - g, - level_scales, - initial_state, - output_final_state, - cu_seqlens, - ) - return o, final_state diff --git a/build/torch-xpu/ops/log_linear_attn/naive.py b/build/torch-xpu/ops/log_linear_attn/naive.py deleted file mode 100644 index 9badd5605bd3bfab72d892315d695a51707ad03c..0000000000000000000000000000000000000000 --- a/build/torch-xpu/ops/log_linear_attn/naive.py +++ /dev/null @@ -1,57 +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 numpy as np -import torch - - -def segsum(x): - T = x.size(-1) - x_cumsum = torch.cumsum(x, dim=-1) - x_segsum = x_cumsum[..., :, None] - x_cumsum[..., None, :] - mask = torch.tril(torch.ones(T, T, device=x.device, dtype=bool)) - x_segsum = x_segsum.masked_fill(~mask, -torch.inf) - return x_segsum - - -def construct_level_mask(level, L): - T = L.size(-1) - if level == 0: - return torch.diag_embed(L[..., level, :]) - - indices = torch.cartesian_prod(torch.arange(T), torch.arange(T)).to(L.device) - - mask = torch.where( - torch.logical_and( - torch.logical_and( - indices[:, 0] % (1 << level) >= (1 << (level - 1)), - indices[:, 1] + (1 << (level - 1)) - >= indices[:, 0] - (indices[:, 0] % (1 << (level - 1))), - ), - indices[:, 1] < indices[:, 0] - (indices[:, 0] % (1 << (level - 1))), - ).view(T, T), - L[..., level, :].unsqueeze(-1).expand(*([-1] * (len(L.shape) - 2)), T, T), - 0, - ) - - return mask - - -def construct_H_matrix(a, L): - T = a.size(-1) - A = torch.exp(segsum(a)) - H = torch.zeros_like(A) - for level in range(int(np.ceil(np.log2(T))) + 1): - mask = construct_level_mask(level, L) - H += A * mask - return H - - -def naive_log_linear_attn(q, k, v, g, level_scales): - H = construct_H_matrix(g.permute(0, 2, 1), level_scales.permute(0, 2, 3, 1)) - M = torch.einsum("bhlc,blhn,bchn->bhlc", H, q, k) - return torch.einsum("bhlc,bchp->blhp", M, v) diff --git a/build/torch-xpu/ops/mesa_net/__init__.py b/build/torch-xpu/ops/mesa_net/__init__.py deleted file mode 100644 index 5d522aadf731d0d69c6c559372bd09905e5ca9fc..0000000000000000000000000000000000000000 --- a/build/torch-xpu/ops/mesa_net/__init__.py +++ /dev/null @@ -1,12 +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 .chunk import chunk_mesa_net -from .decoding_one_step import mesa_net_decoding_one_step -from .naive import naive_mesa_net_decoding_one_step, naive_mesa_net_exact - -__all__ = ['chunk_mesa_net', 'mesa_net_decoding_one_step', 'naive_mesa_net_decoding_one_step', 'naive_mesa_net_exact'] diff --git a/build/torch-xpu/ops/mesa_net/chunk.py b/build/torch-xpu/ops/mesa_net/chunk.py deleted file mode 100644 index 72031cc120808fb77e541ce9e40d89b6c36b3785..0000000000000000000000000000000000000000 --- a/build/torch-xpu/ops/mesa_net/chunk.py +++ /dev/null @@ -1,450 +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 - -from ...modules.l2norm import l2norm_bwd, l2norm_fwd -from ...ops.common.chunk_h import chunk_bwd_dh -from ...ops.mesa_net.chunk_cg_solver_bwd import chunk_mesa_cg_bwd -from ...ops.mesa_net.chunk_cg_solver_fwd import chunk_mesa_cg_fwd -from ...ops.mesa_net.chunk_h_fwd import chunk_mesa_fwd_h -from ...ops.mesa_net.chunk_h_kk_intra_bwd import chunk_mesa_net_h_kk_bwd_intra_fn -from ...ops.mesa_net.chunk_h_kv_intra_bwd import chunk_mesa_net_h_kv_bwd_intra_fn -from ...ops.utils import chunk_local_cumsum, prepare_chunk_indices -from ...ops.utils.constant import RCP_LN2 -from ...utils import autocast_custom_bwd, autocast_custom_fwd, input_guard - - -def chunk_fwd_mesa_net_fwd( - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - g: torch.Tensor, - beta: torch.Tensor, - lamb: torch.Tensor, - cu_seqlens: torch.Tensor, - max_CG_iteration: int = 30, - chunk_size: int = 64, - h_kk_init: torch.Tensor | None = None, - h_kv_init: torch.Tensor | None = None, - output_final_state: bool = False, - chunk_indices: torch.LongTensor | None = None, -) -> torch.Tensor: - if g is not None: - g = chunk_local_cumsum( - g, - chunk_size=chunk_size, - scale=RCP_LN2, - cu_seqlens=cu_seqlens, - chunk_indices=chunk_indices, - ) - h_kk, h_kv, h_kk_final, h_kv_final = chunk_mesa_fwd_h( - k=k, - v=v, - g=g, - beta=beta, - h_init=h_kk_init, - h_kv_init=h_kv_init, - output_final_state=output_final_state, - states_in_fp32=False, - cu_seqlens=cu_seqlens, - chunk_size=chunk_size, - ) - q_star, o = chunk_mesa_cg_fwd( - q=q, - k=k, - h=h_kk, - h_kv=h_kv, - v=v, - g_local_cumsum=g, - beta=beta, - lamb=lamb, - cu_seqlens=cu_seqlens, - chunk_size=chunk_size, - max_CG_iteration=max_CG_iteration, - chunk_indices=chunk_indices, - ) - return g, q_star, o, (h_kk_final, h_kv_final) - - -def chunk_fwd_mesa_net_bwd( - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - g: torch.Tensor, - beta: torch.Tensor, - lamb: torch.Tensor, - q_star: torch.Tensor, # should be cached in the forward pass - do: torch.Tensor, - cu_seqlens: torch.Tensor, - max_CG_iteration: int = 30, - chunk_size: int = 64, - h_kk_init: torch.Tensor | None = None, - h_kv_init: torch.Tensor | None = None, - dh_kv_final: torch.Tensor | None = None, - dh_kk_final: torch.Tensor | None = None, - chunk_indices: torch.LongTensor | None = None, -) -> torch.Tensor: - # recompute the hidden states, which is quite cheap - h_kk, h_kv, _, _ = chunk_mesa_fwd_h( - k=k, - v=v, - g=g, - beta=beta, - h_init=h_kk_init, - h_kv_init=h_kv_init, - output_final_state=False, - states_in_fp32=False, - cu_seqlens=cu_seqlens, - chunk_size=chunk_size, - ) - dh_kv, dh0_kv = chunk_bwd_dh( - q=q_star, - k=k, - v=v, - g=g, - gk=None, - gv=None, - do=do, - h0=h_kv_init, - dht=dh_kv_final, - states_in_fp32=False, - cu_seqlens=cu_seqlens, - chunk_size=chunk_size, - scale=1, - ) - dq, dk_beta, dv, dg = chunk_mesa_net_h_kv_bwd_intra_fn( - q_star=q_star, - k=k, - v=v, - beta=beta, - h_kv=h_kv, - dh_kv=dh_kv, - g=g, - do=do, - cu_seqlens=cu_seqlens, - chunk_size=chunk_size, - chunk_indices=chunk_indices, - ) - dq = chunk_mesa_cg_bwd( - dq=dq, - k=k, - h=h_kk, - g_local_cumsum=g, - beta=beta, - lamb=lamb, - cu_seqlens=cu_seqlens, - chunk_size=chunk_size, - max_CG_iteration=max_CG_iteration, - output_dtype=torch.float16, - chunk_indices=chunk_indices, - ) - dh_kk, dh0_kk = chunk_bwd_dh( - q=dq, - k=k, - v=k, - g=g, - gk=None, - gv=None, - do=q_star, - h0=h_kk_init, - dht=-dh_kk_final if dh_kk_final is not None else None, - states_in_fp32=False, - cu_seqlens=cu_seqlens, - chunk_size=chunk_size, - scale=1, - ) - dk, dg2, dlamb, dbeta = chunk_mesa_net_h_kk_bwd_intra_fn( - k=k, - g=g, - beta=beta, - h=h_kk, - dh=dh_kk, - dk_beta=dk_beta, - q_star=q_star, - dq=dq, - cu_seqlens=cu_seqlens, - chunk_size=chunk_size, - chunk_indices=chunk_indices, - ) - dg.add_(dg2) - dg = chunk_local_cumsum( - dg, - chunk_size=chunk_size, - reverse=True, - cu_seqlens=cu_seqlens, - chunk_indices=chunk_indices, - ).to(g) - return dq, dk, dv, dg, dbeta, dlamb, -dh0_kk if dh0_kk is not None else None, dh0_kv if dh0_kv is not None else None - - -class ChunkMesaNetFunction(torch.autograd.Function): - @staticmethod - @input_guard - @autocast_custom_fwd - def forward( - ctx, - q, - k, - v, - g, - beta, - lamb, - cu_seqlens, - cu_seqlens_cpu, - max_CG_iteration, - h_kk_init, - h_kv_init, - output_final_state, - use_qk_l2norm_in_kernel, - ): - chunk_size = 64 - if cu_seqlens is not None: - chunk_indices = prepare_chunk_indices( - cu_seqlens, - chunk_size, - cu_seqlens_cpu=cu_seqlens_cpu, - ) - else: - chunk_indices = None - - if use_qk_l2norm_in_kernel: - q, q_rstd = l2norm_fwd(q, output_dtype=torch.float16) - k, k_rstd = l2norm_fwd(k, output_dtype=torch.float16) - else: - q_rstd, k_rstd = None, None - q = q.to(torch.float16) - k = k.to(torch.float16) - - g_cumsum, q_star, o, (h_kk_final, h_kv_final) = chunk_fwd_mesa_net_fwd( - q=q, - k=k, - v=v, - g=g, - beta=beta, - lamb=lamb, - cu_seqlens=cu_seqlens, - max_CG_iteration=max_CG_iteration, - chunk_size=chunk_size, - h_kk_init=h_kk_init, - h_kv_init=h_kv_init, - output_final_state=output_final_state, - chunk_indices=chunk_indices, - ) - ctx.max_CG_iteration = max_CG_iteration - ctx.chunk_size = chunk_size - ctx.cu_seqlens = cu_seqlens - ctx.use_qk_l2norm_in_kernel = use_qk_l2norm_in_kernel - ctx.save_for_backward( - q, - q_rstd, - k, - k_rstd, - v, - g_cumsum, - beta, - lamb, - h_kk_init, - h_kv_init, - q_star, - o, - chunk_indices, - ) - return o, h_kk_final, h_kv_final - - @staticmethod - @input_guard - @autocast_custom_bwd - def backward(ctx, do, dh_kk_final=None, dh_kv_final=None): - ( - q, - q_rstd, - k, - k_rstd, - v, - g, - beta, - lamb, - h_kk_init, - h_kv_init, - q_star, - o, - chunk_indices, - ) = ctx.saved_tensors - - max_CG_iteration = ctx.max_CG_iteration - chunk_size = ctx.chunk_size - cu_seqlens = ctx.cu_seqlens - dq, dk, dv, dg, dbeta, dlamb, dh0_kk, dh0_kv = chunk_fwd_mesa_net_bwd( - q=q, - k=k, - v=v, - g=g, - beta=beta, - lamb=lamb, - q_star=q_star, - do=do, - cu_seqlens=cu_seqlens, - max_CG_iteration=max_CG_iteration, - chunk_size=chunk_size, - h_kk_init=h_kk_init, - h_kv_init=h_kv_init, - dh_kv_final=dh_kv_final, - dh_kk_final=dh_kk_final, - chunk_indices=chunk_indices, - ) - if ctx.use_qk_l2norm_in_kernel: - dq = l2norm_bwd(q, q_rstd, dq) - dk = l2norm_bwd(k, k_rstd, dk) - return dq, dk, dv.to(v), dg.to(g), dbeta.to(beta), dlamb.to(lamb), None, None, None, dh0_kk, dh0_kv, None, None - - -@torch.compiler.disable -def chunk_mesa_net( - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - g: torch.Tensor, - beta: torch.Tensor, - lamb: torch.Tensor, - h_kk_init: torch.Tensor | None = None, - h_kv_init: torch.Tensor | None = None, - output_final_state: bool = False, - max_CG_iteration: int = 30, - use_qk_l2norm_in_kernel: bool = False, - cu_seqlens: torch.LongTensor | None = None, - cu_seqlens_cpu: torch.LongTensor | None = None, -): - r""" - Args: - q (torch.Tensor): - queries of shape `[B, T, H, K]` - k (torch.Tensor): - keys of shape `[B, T, H, K]`. Should be l2-normalized before passing in. - v (torch.Tensor): - values of shape `[B, T, H, V]`. - g (torch.Tensor): - decay factors of shape `[B, T, H]`. Note that `g` should be in log space, - that is, `g = log(decay_factor) < 0`. - Recommended input dtype: `torch.float32`. - beta (torch.Tensor): - betas of shape `[B, T, H]`. Recommended input dtype: `torch.float32`. - lamb (torch.Tensor): - lambdas of shape `[B, T, H]`. Recommended input dtype: `torch.float32`. - h_kk_init (Optional[torch.Tensor]): - Initial state of shape `[N, H, K, V]` for `N` input sequences. - For equal-length input sequences, `N` equals the batch size `B`. - Default: `None`. - h_kv_init (Optional[torch.Tensor]): - Initial state of shape `[N, H, K, V]` for `N` input sequences. - For equal-length input sequences, `N` equals the batch size `B`. - Default: `None`. - max_CG_iteration (int): - Maximum number of conjugate gradient iterations for solving the linear system. Default: `30`. - output_final_state (Optional[bool]): - Whether to output the final state of shape `[N, H, K, V]`. Default: `False`. - use_qk_l2norm_in_kernel (bool): - Do l2 normalization on Q and K in the kernel for saving GPU memory. Default: `False`. - cu_seqlens (torch.LongTensor): - Cumulative sequence lengths of shape `[N+1]` used for variable-length training, - consistent with the FlashAttention API. - - Returns: - o (torch.Tensor): - Outputs of shape `[B, T, H, V]`. - (final_states_kk, final_states_kv) (Tuple[torch.Tensor, torch.Tensor]): - Final states of shape `[N, H, K, K]` and `[N, H, K, V]` if `output_final_state=True` else `(None, None)`. - Recall that MesaNet has two states, `h_kk` and `h_kv`! - - Examples:: - >>> import torch - >>> import torch.nn.functional as F - >>> from einops import rearrange - >>> from ...ops.mesa_net import chunk_mesa_net - # inputs with equal lengths - >>> B, T, H, K, V = 4, 2048, 16, 128, 128 - >>> q = torch.randn(B, T, H, K, dtype=torch.bfloat16, device='cuda') - >>> k = torch.randn(B, T, H, K, dtype=torch.bfloat16, device='cuda') - >>> v = torch.randn(B, T, H, V, dtype=torch.bfloat16, device='cuda') - >>> g = F.logsigmoid(torch.randn(B, T, H, dtype=torch.float32, device='cuda')) - >>> beta = torch.rand(B, T, H, dtype=torch.float32, device='cuda').sigmoid() - # lower bound is 0.25 for numerical stability - >>> lamb = F.softplus(torch.rand(H, K, dtype=torch.float32, device='cuda')) + 0.25 - >>> init_state_kk = torch.randn(B, H, K, V, dtype=torch.float32, device='cuda') - >>> init_state_kv = torch.randn(B, H, K, V, dtype=torch.float32, device='cuda') - >>> o, (final_state_kk, final_state_kv) = chunk_mesa_net( - q, k, v, beta, lamb, - h_kk_init=init_state_kk, - h_kv_init=init_state_kv, - max_CG_iteration=30, - output_final_state=True - ) - # for variable-length inputs, the batch size `B` is expected to be 1 and `cu_seqlens` is required - >>> q, k, v, beta = map(lambda x: rearrange(x, 'b t ... -> 1 (b t) ...'), (q, k, v, beta)) - # for a batch with 4 sequences, `cu_seqlens` with 5 start/end positions are expected - >>> cu_seqlens = q.new_tensor([0, 2048, 4096, 6144, 8192], dtype=torch.long) - >>> o_var, (final_state_kk_var, final_state_kv_var) = chunk_mesa_net( - q, k, v, beta, lamb, - h_kk_init=init_state_kk, - h_kv_init=init_state_kv, - max_CG_iteration=30, - output_final_state=True, - cu_seqlens=cu_seqlens - ) - """ - B, T, H, K = q.shape - assert k.shape == (B, T, H, K), "k must be of shape (batch size, seq len, num head, head dim)." - assert v.shape == (B, T, H, K), "v must be of shape (batch size, seq len, num head, head dim)." - assert g.shape == (B, T, H), "g must be of shape (batch size, seq len, num head)." - assert beta.shape == (B, T, H), "beta must be of shape (batch size, seq len, num head)." - assert lamb.shape == (H, K), "lamb must be of shape (num head, key dim)." - - if h_kv_init is not None: - assert h_kv_init.dtype == torch.float32, "h_kv_init must be in float32." - if cu_seqlens is None: - assert h_kv_init.shape == (B, H, K, K), ( - "h_kv_init must be of shape (batch size, num head, head dim, head dim)." - ) - if h_kk_init is not None: - assert h_kk_init.dtype == torch.float32, "h_kk_init must be in float32." - if cu_seqlens is None: - assert h_kk_init.shape == (B, H, K, K), ( - "h_kk_init must be of shape (batch size, num head, head dim, head dim)." - ) - - if cu_seqlens is not None: - if q.shape[0] != 1: - raise ValueError( - f"The batch size is expected to be 1 rather than {q.shape[0]} when using `cu_seqlens`." - f"Please flatten variable-length inputs before processing.", - ) - if h_kk_init is not None and h_kk_init.shape[0] != len(cu_seqlens) - 1: - raise ValueError( - f"The number of initial states is expected to be equal to the number of input sequences, " - f"i.e., {len(cu_seqlens) - 1} rather than {h_kk_init.shape[0]}.", - ) - if h_kv_init is not None and h_kv_init.shape[0] != len(cu_seqlens) - 1: - raise ValueError( - f"The number of initial states is expected to be equal to the number of input sequences, " - f"i.e., {len(cu_seqlens) - 1} rather than {h_kv_init.shape[0]}.", - ) - o, final_state_kk, final_state_kv = ChunkMesaNetFunction.apply( - q, - k, - v, - g, - beta, - lamb, - cu_seqlens, - cu_seqlens_cpu, - max_CG_iteration, - h_kk_init, - h_kv_init, - output_final_state, - use_qk_l2norm_in_kernel, - ) - return o, final_state_kk, final_state_kv diff --git a/build/torch-xpu/ops/mesa_net/chunk_cg_solver_bwd.py b/build/torch-xpu/ops/mesa_net/chunk_cg_solver_bwd.py deleted file mode 100644 index a08425bb369573341a69e6c8c03e0f815b374c55..0000000000000000000000000000000000000000 --- a/build/torch-xpu/ops/mesa_net/chunk_cg_solver_bwd.py +++ /dev/null @@ -1,165 +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 -import triton.language as tl - -from ...ops.utils import prepare_chunk_indices -from ...ops.utils.op import exp2 - - -@triton.jit() -def chunk_update_once( - b_p, - b_k, - b_v, - b_m, - b_g_exp_q, - b_h, - b_lamb, -): - b_o = tl.dot((tl.dot(b_p.to(b_k.dtype), tl.trans(b_k)) * b_m).to(b_v.dtype), b_v) - b_o += tl.dot((b_p * b_g_exp_q).to(b_h.dtype), b_h) - if b_lamb is not None: - b_o += b_lamb[None, :] * b_p - return b_o - - -@triton.heuristics({ - 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None, -}) -@triton.jit(do_not_specialize=['T']) -def chunk_fwd_mesa_cg_dim64_kernel( - dq, - dq_final, - k, - h, - g, - beta, - lamb, - cu_seqlens, - chunk_indices, - T, - max_CG_iteration: tl.constexpr, - H: tl.constexpr, - K: tl.constexpr, - BT: tl.constexpr, - BK: tl.constexpr, - IS_VARLEN: tl.constexpr, -): - i_t, i_bh = tl.program_id(0), tl.program_id(1) - i_b, i_h = i_bh // H, i_bh % H - - 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.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) - T = eos - bos - NT = tl.cdiv(T, BT) - else: - NT = tl.cdiv(T, BT) - i_tg = i_b * NT + i_t - bos, eos = i_b * T, i_b * T + T - - o_t = i_t * BT + tl.arange(0, BT) - m_t = o_t < T - - # offset calculation - dq += (bos * H + i_h) * K - dq_final += (bos * H + i_h) * K - k += (bos * H + i_h) * K - h += (i_tg * H + i_h).to(tl.int64) * K * K - - g += bos * H + i_h - beta += bos * H + i_h - lamb += i_h * K - - p_q = tl.make_block_ptr(dq, (T, K), (H*K, 1), (i_t * BT, 0), (BT, BK), (1, 0)) - p_k = tl.make_block_ptr(k, (T, K), (H*K, 1), (i_t * BT, 0), (BT, BK), (1, 0)) - p_h = tl.make_block_ptr(h, (K, K), (K, 1), (0, 0), (BK, BK), (1, 0)) - - b_h = tl.load(p_h, boundary_check=(0, 1)) - b_k = tl.load(p_k, boundary_check=(0, 1)) - b_q = tl.load(p_q, boundary_check=(0, 1)).to(tl.float32) - - p_g = tl.make_block_ptr(g, (T,), (H,), (i_t * BT,), (BT,), (0,)) - b_g = tl.load(p_g, boundary_check=(0,)).to(tl.float32) - p_beta = tl.make_block_ptr(beta, (T,), (H,), (i_t * BT,), (BT,), (0,)) - b_beta = tl.load(p_beta, boundary_check=(0,)).to(tl.float32) - p_lamb = tl.make_block_ptr(lamb, (K,), (1,), (0,), (BK,), (0,)) - b_lamb = tl.load(p_lamb, boundary_check=(0,)).to(tl.float32) - - b_m = exp2(b_g[:, None] - b_g[None, :]) * b_beta[None, :] - b_m = tl.where((o_t[:, None] >= o_t[None, :]) & (m_t[:, None] & m_t[None, :]), b_m, 0) - b_g_exp_q = exp2(b_g)[:, None] - - b_x = tl.zeros([BT, BK], dtype=tl.float32) - b_p = tl.zeros([BT, BK], dtype=tl.float32) - b_r = tl.zeros([BT, BK], dtype=tl.float32) - - b_x += b_q * 0. - b_r += b_q - b_p += b_q - b_delta_old = tl.sum(b_r*b_r, axis=1) - for _ in range(max_CG_iteration): - b_o = chunk_update_once(b_p, b_k, b_k, b_m, b_g_exp_q, b_h, b_lamb) - alpha = b_delta_old / (tl.sum(b_p*b_o, axis=1) + 1e-5) - b_x += alpha[:, None] * b_p - b_r = b_r - alpha[:, None] * b_o - b_delta_new = tl.sum(b_r*b_r, axis=1) - b_p = b_r + (b_delta_new / (b_delta_old + 1e-5))[:, None] * b_p - b_delta_old = b_delta_new - - p_q_final = tl.make_block_ptr(dq_final, (T, K), (H*K, 1), (i_t * BT, 0), (BT, BK), (1, 0)) - tl.store(p_q_final, b_x.to(p_q_final.dtype.element_ty), boundary_check=(0, 1)) - - -def chunk_mesa_cg_bwd( - dq: torch.Tensor, - k: torch.Tensor, - h: torch.Tensor, - g_local_cumsum: torch.Tensor, - beta: torch.Tensor, - lamb: torch.Tensor, # lambda - cu_seqlens: torch.Tensor | None = None, - chunk_size: int = 64, - max_CG_iteration: int = 30, - output_dtype: torch.dtype | None = None, - chunk_indices: torch.LongTensor | None = None, -) -> torch.Tensor: - B, T, H, K = dq.shape - assert K <= 128, "head dimension must be less than 128" - assert chunk_size <= 64 or K <= 64, "either chunk size or head dimension must be no greater than 64" - dq_final = torch.empty_like(dq, dtype=dq.dtype if output_dtype is None else output_dtype) - - if chunk_indices is None and cu_seqlens is not None: - chunk_indices = prepare_chunk_indices(cu_seqlens, chunk_size) - NT = triton.cdiv(T, chunk_size) if cu_seqlens is None else len(chunk_indices) - BK = max(triton.next_power_of_2(K), 16) - grid = (NT, H*B) - - chunk_fwd_mesa_cg_dim64_kernel[grid]( - dq=dq, - dq_final=dq_final, - k=k, - h=h, - g=g_local_cumsum, - beta=beta, - lamb=lamb, - cu_seqlens=cu_seqlens, - chunk_indices=chunk_indices, - max_CG_iteration=max_CG_iteration, - T=T, - H=H, - K=K, - BT=chunk_size, - BK=BK, - num_warps=4, - num_stages=1, - ) - return dq_final diff --git a/build/torch-xpu/ops/mesa_net/chunk_cg_solver_fwd.py b/build/torch-xpu/ops/mesa_net/chunk_cg_solver_fwd.py deleted file mode 100644 index 75f474462185c8d69427b03db6b0d07e9160d34e..0000000000000000000000000000000000000000 --- a/build/torch-xpu/ops/mesa_net/chunk_cg_solver_fwd.py +++ /dev/null @@ -1,189 +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 -import triton.language as tl - -from ...ops.utils import prepare_chunk_indices -from ...ops.utils.op import exp2 - - -@triton.jit() -def chunk_update_once( - b_p, - b_k, - b_v, - b_m, - b_g_exp_q, - b_h, - b_lamb, -): - b_o = tl.dot((tl.dot(b_p.to(b_k.dtype), tl.trans(b_k)) * b_m).to(b_v.dtype), b_v) - b_o += tl.dot((b_p * b_g_exp_q).to(b_h.dtype), b_h) - if b_lamb is not None: - b_o += b_lamb[None, :] * b_p - return b_o - - -@triton.heuristics({ - 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None, -}) -@triton.jit(do_not_specialize=['T']) -def chunk_fwd_mesa_cg_dim64_kernel( - q, - q_final, - k, - h, - o, - v, - h_kv, - g, - beta, - lamb, - cu_seqlens, - chunk_indices, - T, - max_CG_iteration: tl.constexpr, - H: tl.constexpr, - K: tl.constexpr, - BT: tl.constexpr, - BK: tl.constexpr, - IS_VARLEN: tl.constexpr, -): - i_t, i_bh = tl.program_id(0), tl.program_id(1) - i_b, i_h = i_bh // H, i_bh % H - - 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.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) - T = eos - bos - NT = tl.cdiv(T, BT) - else: - NT = tl.cdiv(T, BT) - i_tg = i_b * NT + i_t - bos, eos = i_b * T, i_b * T + T - - o_t = i_t * BT + tl.arange(0, BT) - m_t = o_t < T - - # offset calculation - q += (bos * H + i_h) * K - q_final += (bos * H + i_h) * K - k += (bos * H + i_h) * K - h += (i_tg * H + i_h).to(tl.int64) * K * K - g += bos * H + i_h - beta += bos * H + i_h - lamb += i_h * K - - o += (bos * H + i_h) * K - v += (bos * H + i_h) * K - h_kv += (i_tg * H + i_h).to(tl.int64) * K * K - - p_q = tl.make_block_ptr(q, (T, K), (H*K, 1), (i_t * BT, 0), (BT, BK), (1, 0)) - p_k = tl.make_block_ptr(k, (T, K), (H*K, 1), (i_t * BT, 0), (BT, BK), (1, 0)) - p_h = tl.make_block_ptr(h, (K, K), (K, 1), (0, 0), (BK, BK), (1, 0)) - - b_h = tl.load(p_h, boundary_check=(0, 1)) - b_k = tl.load(p_k, boundary_check=(0, 1)) - b_q = tl.load(p_q, boundary_check=(0, 1)).to(tl.float32) - - p_g = tl.make_block_ptr(g, (T,), (H,), (i_t * BT,), (BT,), (0,)) - b_g = tl.load(p_g, boundary_check=(0,)).to(tl.float32) - p_beta = tl.make_block_ptr(beta, (T,), (H,), (i_t * BT,), (BT,), (0,)) - b_beta = tl.load(p_beta, boundary_check=(0,)).to(tl.float32) - p_lamb = tl.make_block_ptr(lamb, (K,), (1,), (0,), (BK,), (0,)) - - b_lamb = tl.load(p_lamb, boundary_check=(0,)).to(tl.float32) - - b_m = exp2(b_g[:, None] - b_g[None, :]) * b_beta[None, :] - b_m = tl.where((o_t[:, None] >= o_t[None, :]) & (m_t[:, None] & m_t[None, :]), b_m, 0) - b_g_exp_q = exp2(b_g)[:, None] - - b_x = tl.zeros([BT, BK], dtype=tl.float32) - b_p = tl.zeros([BT, BK], dtype=tl.float32) - b_r = tl.zeros([BT, BK], dtype=tl.float32) - - b_x += b_q * 0. - b_r += b_q - b_p += b_r - b_delta_old = tl.sum(b_r*b_r, axis=1) - for i in range(max_CG_iteration): - b_o = chunk_update_once(b_p, b_k, b_k, b_m, b_g_exp_q, b_h, b_lamb) - alpha = b_delta_old / (tl.sum(b_p*b_o, axis=1) + 1e-5) - b_x += alpha[:, None] * b_p - b_r = b_r - alpha[:, None] * b_o - b_delta_new = tl.sum(b_r*b_r, axis=1) - b_p = b_r + (b_delta_new / (b_delta_old + 1e-5))[:, None] * b_p - b_delta_old = b_delta_new - - p_q_final = tl.make_block_ptr(q_final, (T, K), (H*K, 1), (i_t * BT, 0), (BT, BK), (1, 0)) - tl.store(p_q_final, b_x.to(p_q_final.dtype.element_ty), boundary_check=(0, 1)) - - p_h_kv = tl.make_block_ptr(h_kv, (K, K), (K, 1), (0, 0), (BK, BK), (1, 0)) - b_h_kv = tl.load(p_h_kv, boundary_check=(0, 1)) - p_v = tl.make_block_ptr(v, (T, K), (H*K, 1), (i_t * BT, 0), (BT, BK), (1, 0)) - b_v = tl.load(p_v, boundary_check=(0, 1)) - b_o = chunk_update_once(b_x, b_k, b_v, b_m, b_g_exp_q, b_h_kv, None) - p_o = tl.make_block_ptr(o, (T, K), (H*K, 1), (i_t * BT, 0), (BT, BK), (1, 0)) - tl.store(p_o, b_o.to(p_o.dtype.element_ty), boundary_check=(0, 1)) - - -def chunk_mesa_cg_fwd( - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - h: torch.Tensor, - h_kv: torch.Tensor, - g_local_cumsum: torch.Tensor, - beta: torch.Tensor, - lamb: torch.Tensor, # lambda - cu_seqlens: torch.Tensor | None = None, - chunk_size: int = 64, - max_CG_iteration: int = 30, - output_dtype: torch.dtype | None = None, - chunk_indices: torch.LongTensor | None = None, -) -> torch.Tensor: - B, T, H, K = q.shape - assert K <= 128, "head dimension must be less than 128" - assert chunk_size <= 64 or K <= 64, "either chunk size or head dimension must be no greater than 64" - q_final = torch.empty_like(q, dtype=q.dtype if output_dtype is None else output_dtype) - - assert v is not None, "v must be provided if calculate_output is True" - assert h_kv is not None, "h_kv must be provided if calculate_output is True" - o = torch.empty_like(v) - - if chunk_indices is None and cu_seqlens is not None: - chunk_indices = prepare_chunk_indices(cu_seqlens, chunk_size) - NT = triton.cdiv(T, chunk_size) if cu_seqlens is None else len(chunk_indices) - BK = max(triton.next_power_of_2(K), 16) - grid = (NT, H*B) - - chunk_fwd_mesa_cg_dim64_kernel[grid]( - q=q, - q_final=q_final, - o=o, - v=v, - h_kv=h_kv, - k=k, - h=h, - g=g_local_cumsum, - beta=beta, - lamb=lamb, - cu_seqlens=cu_seqlens, - chunk_indices=chunk_indices, - max_CG_iteration=max_CG_iteration, - T=T, - H=H, - K=K, - BT=chunk_size, - BK=BK, - num_warps=4, - num_stages=1, - ) - return q_final, o diff --git a/build/torch-xpu/ops/mesa_net/chunk_h_fwd.py b/build/torch-xpu/ops/mesa_net/chunk_h_fwd.py deleted file mode 100644 index 88bfa11b11698295a6c835694efce8e887fea887..0000000000000000000000000000000000000000 --- a/build/torch-xpu/ops/mesa_net/chunk_h_fwd.py +++ /dev/null @@ -1,174 +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 -import triton.language as tl - -from ...ops.utils import prepare_chunk_offsets -from ...ops.utils.op import exp2, safe_dot -from ...utils import autotune_cache_kwargs - - -@triton.heuristics({ - 'USE_INITIAL_STATE': lambda args: args['h_init'] is not None, - 'STORE_FINAL_STATE': lambda args: args['h_final'] is not None, - 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None, -}) -@triton.autotune( - configs=[ - triton.Config({}, num_warps=num_warps, num_stages=num_stages) - for num_warps in [1, 2, 4, 8] - for num_stages in [2, 3, 4] - ], - key=['BT'], - **autotune_cache_kwargs, -) -@triton.jit(do_not_specialize=['T']) -def chunk_mesa_net_fwd_kernel_h( - k, - v, - beta, - g, - h, - h_kv, - h_init, - h_kv_init, - h_final, - h_kv_final, - cu_seqlens, - split_offsets, - T, - H: tl.constexpr, - K: tl.constexpr, - V: tl.constexpr, - BT: tl.constexpr, - BS: tl.constexpr, - BK: tl.constexpr, - BV: tl.constexpr, - USE_INITIAL_STATE: tl.constexpr, - STORE_FINAL_STATE: tl.constexpr, - IS_VARLEN: tl.constexpr, -): - i_k, i_v, i_nh = tl.program_id(0), tl.program_id(1), tl.program_id(2) - i_n, i_h = i_nh // H, i_nh % H - if IS_VARLEN: - bos, eos = tl.load(cu_seqlens + i_n).to(tl.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) - T = eos - bos - NT = tl.cdiv(T, BT) - NS = tl.cdiv(T, BS) - boh = tl.load(split_offsets + i_n).to(tl.int32) - else: - bos, eos = i_n * T, i_n * T + T - NT = tl.cdiv(T, BT) - NS = tl.cdiv(T, BS) - boh = i_n * NS - - # [BK, BV] - b_h = tl.zeros([BK, BV], dtype=tl.float32) - b_h_kv = tl.zeros([BK, BV], dtype=tl.float32) - if USE_INITIAL_STATE: - p_h0 = tl.make_block_ptr(h_init + i_nh * K*V, (K, V), (V, 1), (i_k * BK, i_v * BV), (BK, BV), (1, 0)) - b_h = tl.load(p_h0, boundary_check=(0, 1)).to(tl.float32) - p_h_kv0 = tl.make_block_ptr(h_kv_init + i_nh * K*V, (K, V), (V, 1), (i_k * BK, i_v * BV), (BK, BV), (1, 0)) - b_h_kv = tl.load(p_h_kv0, boundary_check=(0, 1)).to(tl.float32) - - for i_t in range(NT): - i_s = i_t // (BS // BT) - p_k = tl.make_block_ptr(k + (bos*H + i_h) * K, (T, K), (H*K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) - p_k2 = tl.make_block_ptr(k + (bos*H + i_h) * V, (T, V), (H*V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) - p_v = tl.make_block_ptr(v + (bos*H + i_h) * V, (T, V), (H*V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) - p_beta = tl.make_block_ptr(beta + (bos*H + i_h), (T,), (H,), (i_t * BT,), (BT, ), (0,)) - b_beta = tl.load(p_beta, boundary_check=(0,)) - - o_h = ((boh + i_s) * H + i_h).to(tl.int64) * K*V - p_h = tl.make_block_ptr(h + o_h, (K, V), (V, 1), (i_k * BK, i_v * BV), (BK, BV), (1, 0)) - p_h_kv = tl.make_block_ptr(h_kv + o_h, (K, V), (V, 1), (i_k * BK, i_v * BV), (BK, BV), (1, 0)) - - if i_t % (BS // BT) == 0: - tl.store(p_h, b_h.to(p_h.dtype.element_ty), boundary_check=(0, 1)) - tl.store(p_h_kv, b_h_kv.to(p_h_kv.dtype.element_ty), boundary_check=(0, 1)) - - # [BK, BT] - b_k = tl.load(p_k, boundary_check=(0, 1)) - b_k2 = tl.load(p_k2, boundary_check=(0, 1)) - # [BT, BV] - b_v = tl.load(p_v, boundary_check=(0, 1)) - last_idx = min((i_t + 1) * BT, T) - 1 - - # scalar decay - b_g_last = tl.load(g + bos * H + last_idx * H + i_h) - p_g = g + bos*H + (i_t * BT + tl.arange(0, BT)) * H + i_h - b_h *= exp2(b_g_last) - b_h_kv *= exp2(b_g_last) - b_g = tl.load(p_g, mask=(i_t * BT + tl.arange(0, BT) < T), other=0.) - b_k_decay = ((b_k * exp2(b_g_last - b_g)[:, None]) * b_beta[:, None]).to(b_k2.dtype) - b_h += safe_dot(tl.trans(b_k_decay), b_k2) - b_h_kv += safe_dot(tl.trans(b_k_decay), b_v.to(b_k2.dtype)) - - if STORE_FINAL_STATE: - p_ht = tl.make_block_ptr(h_final + i_nh * K*V, (K, V), (V, 1), (i_k * BK, i_v * BV), (BK, BV), (1, 0)) - tl.store(p_ht, b_h.to(p_ht.dtype.element_ty), boundary_check=(0, 1)) - p_h_kv_final = tl.make_block_ptr(h_kv_final + i_nh * K*V, (K, V), (V, 1), (i_k * BK, i_v * BV), (BK, BV), (1, 0)) - tl.store(p_h_kv_final, b_h_kv.to(p_h_kv_final.dtype.element_ty), boundary_check=(0, 1)) - - -def chunk_mesa_fwd_h( - k: torch.Tensor, - v: torch.Tensor, - g: torch.Tensor, - beta: torch.Tensor, - h_init: torch.Tensor, - h_kv_init: torch.Tensor, - output_final_state: bool, - cu_seqlens: torch.Tensor | None = None, - chunk_size: int = 64, - split_size: int | None = None, - states_in_fp32: bool = False, -) -> tuple[torch.Tensor, torch.Tensor]: - B, T, H, K, V = *k.shape, v.shape[-1] - assert K == V, "K must be equal to V for now" - BT = chunk_size - BS = BT if split_size is None else split_size - assert BS % BT == 0, f"The `split_size` (got {BS}) must be a multiple of `chunk_size` {BT}" - # N: the actual number of sequences in the batch with either equal or variable lengths - if cu_seqlens is None: - N, NS, split_offsets = B, triton.cdiv(T, BS), None - else: - split_offsets = prepare_chunk_offsets(cu_seqlens, BS) - N, NS = len(cu_seqlens) - 1, split_offsets[-1].item() - - h = k.new_empty(B, NS, H, K, V, dtype=k.dtype if not states_in_fp32 else torch.float) - h_kv = k.new_empty(B, NS, H, K, V, dtype=k.dtype if not states_in_fp32 else torch.float) - h_final = k.new_empty(N, H, K, V, dtype=torch.float) if output_final_state else None - h_kv_final = k.new_empty(N, H, K, V, dtype=torch.float) - - def grid(meta): return (triton.cdiv(K, 64), triton.cdiv(V, 64), N * H) - - chunk_mesa_net_fwd_kernel_h[grid]( - k=k, - v=v, - beta=beta, - g=g, - h=h, - h_kv=h_kv, - h_init=h_init, - h_kv_init=h_kv_init, - h_final=h_final, - h_kv_final=h_kv_final, - cu_seqlens=cu_seqlens, - split_offsets=split_offsets, - T=T, - H=H, - K=K, - V=V, - BT=BT, - BS=BS, - BK=64, - BV=64, - ) - return h, h_kv, h_final, h_kv_final diff --git a/build/torch-xpu/ops/mesa_net/chunk_h_kk_intra_bwd.py b/build/torch-xpu/ops/mesa_net/chunk_h_kk_intra_bwd.py deleted file mode 100644 index 172bdddaba53dd99a5662133b644420139949907..0000000000000000000000000000000000000000 --- a/build/torch-xpu/ops/mesa_net/chunk_h_kk_intra_bwd.py +++ /dev/null @@ -1,195 +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 -import triton.language as tl - -from ...ops.utils import prepare_chunk_indices -from ...ops.utils.op import exp2 - - -@triton.heuristics({ - 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None, -}) -@triton.jit(do_not_specialize=['T']) -def chunk_mesa_net_h_kk_bwd_intra_kernel( - k, - beta, - h, - dh, - g, - q_star, - dq, - dk, - dg, - dbeta, - dk_beta, - dlamb, - cu_seqlens, - chunk_indices, - B: tl.constexpr, - T, - H: tl.constexpr, - K: tl.constexpr, - V: tl.constexpr, - BT: tl.constexpr, - BK: tl.constexpr, - BV: tl.constexpr, - IS_VARLEN: tl.constexpr, -): - i_t, i_bh = tl.program_id(0), tl.program_id(1) - i_b, i_h = i_bh // H, i_bh % H - - 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.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) - T = eos - bos - NT = tl.cdiv(T, BT) - else: - NT = tl.cdiv(T, BT) - i_tg = i_b * NT + i_t - bos, eos = i_b * T, i_b * T + T - - o_t = i_t * BT + tl.arange(0, BT) - m_t = o_t < T - - # offset calculation - q_star += (bos * H + i_h) * V - dq += (bos * H + i_h) * V - h += (i_tg * H + i_h).to(tl.int64) * K*V - dh += (i_tg * H + i_h).to(tl.int64) * K*V - k += (bos * H + i_h) * K - dk += (bos * H + i_h) * K - dk_beta += (bos * H + i_h) * K - dlamb += (i_tg * H + i_h).to(tl.int64) * K - beta += (bos * H + i_h) - dbeta += (bos * H + i_h) - g += bos * H + i_h - dg += bos * H + i_h - - b_dk = tl.zeros([BT, BK], dtype=tl.float32) - b_dv = tl.zeros([BT, BK], dtype=tl.float32) - b_dbeta = tl.zeros([BT], dtype=tl.float32) - b_dg_last = tl.zeros([1], dtype=tl.float32) - b_dg = tl.zeros([BT], dtype=tl.float32) - - p_g = tl.make_block_ptr(g, (T,), (H,), (i_t * BT,), (BT,), (0,)) - b_g = tl.load(p_g, boundary_check=(0,)) - b_g_last = tl.load(g + (min(i_t * BT + BT, T) - 1) * H) - b_gk = tl.where(m_t, exp2(b_g_last - b_g), 0) - - p_q_star = tl.make_block_ptr(q_star, (T, V), (H*V, 1), (i_t * BT, 0), (BT, BV), (1, 0)) - b_q_star = tl.load(p_q_star, boundary_check=(0, 1)) - p_dq = tl.make_block_ptr(dq, (T, V), (H*V, 1), (i_t * BT, 0), (BT, BV), (1, 0)) - b_dq = tl.load(p_dq, boundary_check=(0, 1)) - b_dlamb = -tl.sum(b_q_star * b_dq, axis=0) - p_dlamb = tl.make_block_ptr(dlamb, (K,), (1,), (0,), (BK,), (0,)) - tl.store(p_dlamb, b_dlamb.to(p_dlamb.dtype.element_ty), boundary_check=(0,)) - - p_h = tl.make_block_ptr(h, (V, K), (1, V), (0, 0), (BV, BK), (0, 1)) - p_dh = tl.make_block_ptr(dh, (V, K), (1, V), (0, 0), (BV, BK), (0, 1)) - p_beta = tl.make_block_ptr(beta, (T,), (H,), (i_t * BT,), (BT,), (0,)) - b_beta = tl.load(p_beta, boundary_check=(0,)) - p_k = tl.make_block_ptr(k, (T, K), (H*K, 1), (i_t * BT, 0), (BT, BK), (1, 0)) - b_v = tl.load(p_k, boundary_check=(0, 1)) - b_k = (b_v * b_beta[:, None]).to(b_v.dtype) - - b_m = tl.where((o_t[:, None] >= o_t[None, :]) & (m_t[:, None] & m_t[None, :]), exp2(b_g[:, None] - b_g[None, :]), 0) - b_s = tl.dot(b_q_star, tl.trans(b_k)) * b_m - b_ds = tl.dot(b_dq, tl.trans(b_v)) - b_dv += tl.dot(tl.trans(b_s.to(b_dq.dtype)), b_dq) - b_dm = b_s * b_ds - b_dm = tl.where(tl.arange(0, BT)[:, None] >= tl.arange(0, BT)[None, :], b_dm, 0) - b_dg += tl.sum(b_dm, axis=1) - b_dg -= tl.sum(b_dm, axis=0) - b_ds = b_ds * b_m - b_dk += tl.dot(tl.trans(b_ds.to(b_q_star.dtype)), b_q_star) - - b_h = tl.load(p_h, boundary_check=(0, 1)) - b_dg += tl.sum(tl.dot(b_dq, tl.trans(b_h)) * exp2(b_g)[:, None] * b_q_star, axis=1) - b_dh = tl.load(p_dh, boundary_check=(0, 1)) - b_dk2 = tl.dot(b_v, b_dh.to(b_v.dtype)) * b_gk[:, None] - b_dg -= tl.sum(b_dk2 * b_k, axis=1) - b_dg_last += tl.sum(b_dk2 * b_k) - b_dk += b_dk2 - b_dv += tl.dot(b_k, tl.trans(b_dh).to(b_k.dtype)) * b_gk[:, None] - b_dh = b_dh * b_h - b_dg_last += tl.sum(b_dh) * exp2(b_g_last) - - p_dk_beta = tl.make_block_ptr(dk_beta, (T, K), (H*K, 1), (i_t * BT, 0), (BT, BK), (1, 0)) - b_dk -= tl.load(p_dk_beta, boundary_check=(0, 1)) - b_dbeta = tl.sum(b_dk * b_v, axis=1) - b_dk = b_dk * b_beta[:, None] + b_dv - b_dk = -b_dk - - b_dg = tl.where(o_t < min(i_t * BT + BT, T) - 1, b_dg, b_dg + b_dg_last) - p_dk = tl.make_block_ptr(dk, (T, K), (H*K, 1), (i_t * BT, 0), (BT, BK), (1, 0)) - tl.store(p_dk, b_dk.to(p_dk.dtype.element_ty), boundary_check=(0, 1)) - p_dg = tl.make_block_ptr(dg, (T,), (H,), (i_t * BT,), (BT,), (0,)) - tl.store(p_dg, -b_dg.to(p_dg.dtype.element_ty), boundary_check=(0,)) - p_dbeta = tl.make_block_ptr(dbeta, (T,), (H,), (i_t * BT,), (BT,), (0,)) - tl.store(p_dbeta, -b_dbeta.to(p_dbeta.dtype.element_ty), boundary_check=(0,)) - - -def chunk_mesa_net_h_kk_bwd_intra_fn( - k: torch.Tensor, - beta: torch.Tensor, - g: torch.Tensor, - h: torch.Tensor, - dh: torch.Tensor, - q_star: torch.Tensor, - dq: torch.Tensor, - dk_beta: torch.Tensor, - cu_seqlens: torch.LongTensor | None = None, - chunk_size: int = 64, - chunk_indices: torch.LongTensor | None = None, -) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: - - B, T, H, K = k.shape - V = K - BT = min(chunk_size, max(16, triton.next_power_of_2(T))) - if chunk_indices is None and cu_seqlens is not None: - chunk_indices = prepare_chunk_indices(cu_seqlens, BT) - NT = triton.cdiv(T, BT) if cu_seqlens is None else len(chunk_indices) - - # CONST_TILING = 64 - BK = max(triton.next_power_of_2(K), 16) - BV = max(triton.next_power_of_2(V), 16) - dk = torch.empty_like(k) - dg = torch.empty_like(g) - dbeta = torch.empty_like(beta) - dlamb = torch.empty(B, NT, H, K, dtype=torch.float32, device=k.device) - grid = (NT, B * H) - - chunk_mesa_net_h_kk_bwd_intra_kernel[grid]( - k=k, - h=h, - dh=dh, - g=g, - q_star=q_star, - beta=beta, - dbeta=dbeta, - dq=dq, - dk=dk, - dk_beta=dk_beta, - dg=dg, - dlamb=dlamb, - cu_seqlens=cu_seqlens, - chunk_indices=chunk_indices, - B=B, - T=T, - H=H, - K=K, - V=V, - BT=BT, - BK=BK, - BV=BV, - ) - dlamb = dlamb.sum([0, 1]) - return dk, dg, dlamb, dbeta diff --git a/build/torch-xpu/ops/mesa_net/chunk_h_kv_intra_bwd.py b/build/torch-xpu/ops/mesa_net/chunk_h_kv_intra_bwd.py deleted file mode 100644 index 9767ed1c2b8bb24e8268250bbee5e13131b276a6..0000000000000000000000000000000000000000 --- a/build/torch-xpu/ops/mesa_net/chunk_h_kv_intra_bwd.py +++ /dev/null @@ -1,217 +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 -import triton.language as tl - -from ...ops.mesa_net.chunk_h_kv_intra_bwd_separate import chunk_mesa_net_h_kv_bwd_intra_separate_fn -from ...ops.utils import prepare_chunk_indices -from ...ops.utils.op import exp2 -from ...utils import IS_NVIDIA_HOPPER, autotune_cache_kwargs, check_shared_mem - -NUM_WARPS = [2, 4] if IS_NVIDIA_HOPPER else [2, 4, 8] - - -@triton.heuristics({ - 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None, -}) -@triton.autotune( - configs=[ - triton.Config({}, num_warps=num_warps, num_stages=num_stages) - for num_warps in NUM_WARPS - for num_stages in [2, 3, 4] - ], - key=['H', 'K', 'V', 'BT', 'BK', 'BV'], - **autotune_cache_kwargs, -) -@triton.jit(do_not_specialize=['T']) -def chunk_mesa_net_h_kv_bwd_intra_kernel( - q_star, - k, - v, - beta, - h_kv, - g, - do, - dh_kv, - dq, - dk_beta, - dg, - dv, - cu_seqlens, - chunk_indices, - B: tl.constexpr, - T, - H: tl.constexpr, - K: tl.constexpr, - V: tl.constexpr, - BT: tl.constexpr, - BK: tl.constexpr, - BV: tl.constexpr, - IS_VARLEN: tl.constexpr, -): - i_t, i_bh = tl.program_id(0), tl.program_id(1) - i_b, i_h = i_bh // H, i_bh % H - 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.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) - T = eos - bos - NT = tl.cdiv(T, BT) - else: - NT = tl.cdiv(T, BT) - i_tg = i_b * NT + i_t - bos, eos = i_b * T, i_b * T + T - - o_t = i_t * BT + tl.arange(0, BT) - m_t = o_t < T - - # offset calculation - v += (bos * H + i_h) * V - do += (bos * H + i_h) * V - h_kv += (i_tg * H + i_h).to(tl.int64) * K*V - dh_kv += (i_tg * H + i_h).to(tl.int64) * K*V - q_star += (bos * H + i_h) * K - k += (bos * H + i_h) * K - beta += (bos * H + i_h) - g += bos * H + i_h - dg += bos * H + i_h - dq += (bos * H + i_h) * K - dk_beta += (bos * H + i_h) * K - dv += (bos * H + i_h) * V - - b_dq = tl.zeros([BT, BK], dtype=tl.float32) - b_dk = tl.zeros([BT, BK], dtype=tl.float32) - b_ds = tl.zeros([BT, BT], dtype=tl.float32) - b_dv = tl.zeros([BT, BK], dtype=tl.float32) - b_dg_last = tl.zeros([1], dtype=tl.float32) - b_dg = tl.zeros([BT], dtype=tl.float32) - - p_q = tl.make_block_ptr(q_star, (T, K), (H*K, 1), (i_t * BT, 0), (BT, BK), (1, 0)) - p_k = tl.make_block_ptr(k, (T, K), (H*K, 1), (i_t * BT, 0), (BT, BK), (1, 0)) - p_v = tl.make_block_ptr(v, (T, V), (H*V, 1), (i_t * BT, 0), (BT, BV), (1, 0)) - p_g = tl.make_block_ptr(g, (T,), (H,), (i_t * BT,), (BT,), (0,)) - p_beta = tl.make_block_ptr(beta, (T, ), (H, ), (i_t * BT,), (BT,), (0,)) - p_do = tl.make_block_ptr(do, (T, V), (H*V, 1), (i_t * BT, 0), (BT, BV), (1, 0)) - p_h = tl.make_block_ptr(h_kv, (V, K), (1, V), (0, 0), (BV, BK), (0, 1)) - p_dh = tl.make_block_ptr(dh_kv, (V, K), (1, V), (0, 0), (BV, BK), (0, 1)) - - b_q = tl.load(p_q, boundary_check=(0, 1)) - b_k = tl.load(p_k, boundary_check=(0, 1)) - b_v = tl.load(p_v, boundary_check=(0, 1)) - b_g = tl.load(p_g, boundary_check=(0,)) - b_beta = tl.load(p_beta, boundary_check=(0, )) - b_do = tl.load(p_do, boundary_check=(0, 1)) - b_h = tl.load(p_h, boundary_check=(0, 1)) - b_dh = tl.load(p_dh, boundary_check=(0, 1)) - b_g_last = tl.load(g + (min(i_t * BT + BT, T) - 1) * H) - - # calculation - b_dg_last += tl.sum(b_h * b_dh) - b_dg_last *= exp2(b_g_last) - - b_m = tl.where((o_t[:, None] >= o_t[None, :]) & (m_t[:, None] & m_t[None, :]), exp2(b_g[:, None] - b_g[None, :]), 0) - b_k = (b_k * b_beta[:, None]).to(b_k.dtype) - b_s = tl.dot(b_q, tl.trans(b_k)) * b_m - - b_ds = tl.dot(b_do, tl.trans(b_v)) - b_dm = b_s * b_ds - b_dm = tl.where(tl.arange(0, BT)[:, None] >= tl.arange(0, BT)[None, :], b_dm, 0) - - b_dg += tl.sum(b_dm, axis=1) - b_dg -= tl.sum(b_dm, axis=0) - - b_g_exp_q = exp2(b_g) - b_g_exp_k = tl.where(m_t, exp2(-b_g + b_g_last), 0) - b_ds = b_ds * b_m - b_dq += tl.dot(b_do, b_h.to(b_do.dtype)) * b_g_exp_q[:, None] - b_dk += tl.dot(b_v, b_dh.to(b_v.dtype)) * b_g_exp_k[:, None] - b_dg_last += tl.sum(b_dk * b_k) - b_dg -= tl.sum(b_dk * b_k, axis=1) - b_dg += tl.sum(b_dq * b_q, axis=1) - b_dq += tl.dot(b_ds.to(b_k.dtype), b_k) - b_dv += tl.dot(b_k, tl.trans(b_dh).to(b_k.dtype)) * b_g_exp_k[:, None] + tl.dot(tl.trans(b_s.to(b_do.dtype)), b_do) - b_dk += tl.dot(tl.trans(b_ds.to(b_q.dtype)), b_q) - - b_dg = tl.where(o_t < min(i_t * BT + BT, T) - 1, b_dg, b_dg + b_dg_last) - p_dq = tl.make_block_ptr(dq, (T, K), (H*K, 1), (i_t * BT, 0), (BT, BK), (1, 0)) - p_dk = tl.make_block_ptr(dk_beta, (T, K), (H*K, 1), (i_t * BT, 0), (BT, BK), (1, 0)) - p_dv = tl.make_block_ptr(dv, (T, V), (H*V, 1), (i_t * BT, 0), (BT, BV), (1, 0)) - p_dg = tl.make_block_ptr(dg, (T,), (H,), (i_t * BT,), (BT,), (0,)) - tl.store(p_dq, b_dq.to(p_dq.dtype.element_ty), boundary_check=(0, 1)) - tl.store(p_dk, b_dk.to(p_dk.dtype.element_ty), boundary_check=(0, 1)) - tl.store(p_dv, b_dv.to(p_dv.dtype.element_ty), boundary_check=(0, 1)) - tl.store(p_dg, b_dg.to(p_dg.dtype.element_ty), boundary_check=(0,)) - - -def chunk_mesa_net_h_kv_bwd_intra_fn( - q_star, - k, - v, - beta, - h_kv, - dh_kv, - g, - do, - cu_seqlens, - chunk_size=64, - chunk_indices: torch.LongTensor | None = None, -): - # share memory is not large enough for a single fused kernel - if not check_shared_mem('ampere'): - return chunk_mesa_net_h_kv_bwd_intra_separate_fn( - q_star=q_star, - k=k, - v=v, - beta=beta, - h_kv=h_kv, - dh_kv=dh_kv, - g=g, - do=do, - cu_seqlens=cu_seqlens, - chunk_size=chunk_size, - chunk_indices=chunk_indices, - ) - B, T, H, K, V = *k.shape, v.shape[-1] - BT = chunk_size - if chunk_indices is None and cu_seqlens is not None: - chunk_indices = prepare_chunk_indices(cu_seqlens, BT) - NT = triton.cdiv(T, BT) if cu_seqlens is None else len(chunk_indices) - - BK = max(triton.next_power_of_2(K), 16) - BV = max(triton.next_power_of_2(V), 16) - dq = torch.empty_like(q_star, dtype=torch.float32) - dk = torch.empty_like(k) - dv = torch.empty_like(v) - dg = torch.empty_like(g) - grid = (NT, B * H) - chunk_mesa_net_h_kv_bwd_intra_kernel[grid]( - q_star=q_star, - k=k, - v=v, - beta=beta, - h_kv=h_kv, - g=g, - do=do, - dh_kv=dh_kv, - dq=dq, - dk_beta=dk, - dg=dg, - dv=dv, - cu_seqlens=cu_seqlens, - chunk_indices=chunk_indices, - B=B, - T=T, - H=H, - K=K, - V=V, - BT=BT, - BK=BK, - BV=BV, - ) - return dq, dk, dv, dg diff --git a/build/torch-xpu/ops/mesa_net/chunk_h_kv_intra_bwd_separate.py b/build/torch-xpu/ops/mesa_net/chunk_h_kv_intra_bwd_separate.py deleted file mode 100644 index 9f352a6265869004b7d745040fff7a3340a95bd1..0000000000000000000000000000000000000000 --- a/build/torch-xpu/ops/mesa_net/chunk_h_kv_intra_bwd_separate.py +++ /dev/null @@ -1,308 +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 -import triton.language as tl - -from ...ops.utils import prepare_chunk_indices -from ...ops.utils.op import exp2 -from ...utils import IS_NVIDIA_HOPPER, autotune_cache_kwargs - -NUM_WARPS = [2, 4] if IS_NVIDIA_HOPPER else [2, 4, 8] - - -@triton.heuristics({ - 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None, -}) -@triton.autotune( - configs=[ - triton.Config({}, num_warps=num_warps, num_stages=num_stages) - for num_warps in NUM_WARPS - for num_stages in [2, 3, 4] - ], - key=['H', 'K', 'V', 'BT', 'BK', 'BV'], - **autotune_cache_kwargs, -) -@triton.jit(do_not_specialize=['T']) -def chunk_mesa_net_h_kv_bwd_intra_kernel_dkv( - q_star, - k, - v, - beta, - h_kv, - g, - do, - dh_kv, - dk_beta, - dg, - dv, - cu_seqlens, - chunk_indices, - B: tl.constexpr, - T, - H: tl.constexpr, - K: tl.constexpr, - V: tl.constexpr, - BT: tl.constexpr, - BK: tl.constexpr, - BV: tl.constexpr, - IS_VARLEN: tl.constexpr, -): - i_t, i_bh = tl.program_id(0), tl.program_id(1) - i_b, i_h = i_bh // H, i_bh % H - 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.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) - T = eos - bos - NT = tl.cdiv(T, BT) - else: - NT = tl.cdiv(T, BT) - i_tg = i_b * NT + i_t - bos, eos = i_b * T, i_b * T + T - - o_t = i_t * BT + tl.arange(0, BT) - m_t = o_t < T - - # offset calculation - v += (bos * H + i_h) * V - do += (bos * H + i_h) * V - h_kv += (i_tg * H + i_h).to(tl.int64) * K*V - dh_kv += (i_tg * H + i_h).to(tl.int64) * K*V - q_star += (bos * H + i_h) * K - k += (bos * H + i_h) * K - beta += (bos * H + i_h) - g += bos * H + i_h - dg += bos * H + i_h - dk_beta += (bos * H + i_h) * K - dv += (bos * H + i_h) * V - - b_dk = tl.zeros([BT, BK], dtype=tl.float32) - b_ds = tl.zeros([BT, BT], dtype=tl.float32) - b_dv = tl.zeros([BT, BK], dtype=tl.float32) - b_dg_last = tl.zeros([1], dtype=tl.float32) - b_dg = tl.zeros([BT], dtype=tl.float32) - - p_v = tl.make_block_ptr(v, (T, V), (H*V, 1), (i_t * BT, 0), (BT, BV), (1, 0)) - p_k = tl.make_block_ptr(k, (T, K), (H*K, 1), (i_t * BT, 0), (BT, BK), (1, 0)) - p_beta = tl.make_block_ptr(beta, (T, ), (H, ), (i_t * BT,), (BT,), (0,)) - p_do = tl.make_block_ptr(do, (T, V), (H*V, 1), (i_t * BT, 0), (BT, BV), (1, 0)) - p_h = tl.make_block_ptr(h_kv, (V, K), (1, V), (0, 0), (BV, BK), (0, 1)) - p_dh = tl.make_block_ptr(dh_kv, (V, K), (1, V), (0, 0), (BV, BK), (0, 1)) - p_q = tl.make_block_ptr(q_star, (T, K), (H*K, 1), (i_t * BT, 0), (BT, BK), (1, 0)) - p_g = tl.make_block_ptr(g, (T,), (H,), (i_t * BT,), (BT,), (0,)) - - b_q = tl.load(p_q, boundary_check=(0, 1)) - b_v = tl.load(p_v, boundary_check=(0, 1)) - b_k = tl.load(p_k, boundary_check=(0, 1)) - b_beta = tl.load(p_beta, boundary_check=(0, )) - b_g = tl.load(p_g, boundary_check=(0,)) - b_do = tl.load(p_do, boundary_check=(0, 1)) - b_h = tl.load(p_h, boundary_check=(0, 1)) - b_dh = tl.load(p_dh, boundary_check=(0, 1)) - b_g_last = tl.load(g + (min(i_t * BT + BT, T) - 1) * H) - - # calculation - b_dg_last += tl.sum(b_h * b_dh) - b_dg_last *= exp2(b_g_last) - - b_m = tl.where((o_t[:, None] >= o_t[None, :]) & (m_t[:, None] & m_t[None, :]), exp2(b_g[:, None] - b_g[None, :]), 0) - b_k = (b_k * b_beta[:, None]).to(b_k.dtype) - b_s = tl.dot(b_q, tl.trans(b_k)) * b_m - b_ds = tl.dot(b_do, tl.trans(b_v)) - b_dm = b_s * b_ds - b_dm = tl.where(tl.arange(0, BT)[:, None] >= tl.arange(0, BT)[None, :], b_dm, 0) - b_dg += tl.sum(b_dm, axis=1) - b_dg -= tl.sum(b_dm, axis=0) - b_g_exp_k = tl.where(m_t, exp2(-b_g + b_g_last), 0) - b_ds = b_ds * b_m - b_dk += tl.dot(b_v, b_dh.to(b_v.dtype)) * b_g_exp_k[:, None] - b_dg_last += tl.sum(b_dk * b_k) - b_dg -= tl.sum(b_dk * b_k, axis=1) - b_dv += tl.dot(b_k, tl.trans(b_dh).to(b_k.dtype)) * b_g_exp_k[:, None] + tl.dot(tl.trans(b_s.to(b_do.dtype)), b_do) - b_dk += tl.dot(tl.trans(b_ds.to(b_q.dtype)), b_q) - b_dg = tl.where(o_t < min(i_t * BT + BT, T) - 1, b_dg, b_dg + b_dg_last) - p_dk = tl.make_block_ptr(dk_beta, (T, K), (H*K, 1), (i_t * BT, 0), (BT, BK), (1, 0)) - p_dv = tl.make_block_ptr(dv, (T, V), (H*V, 1), (i_t * BT, 0), (BT, BV), (1, 0)) - p_dg = tl.make_block_ptr(dg, (T,), (H,), (i_t * BT,), (BT,), (0,)) - tl.store(p_dk, b_dk.to(p_dk.dtype.element_ty), boundary_check=(0, 1)) - tl.store(p_dv, b_dv.to(p_dv.dtype.element_ty), boundary_check=(0, 1)) - tl.store(p_dg, b_dg.to(p_dg.dtype.element_ty), boundary_check=(0,)) - - -@triton.heuristics({ - 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None, -}) -@triton.autotune( - configs=[ - triton.Config({}, num_warps=num_warps, num_stages=num_stages) - for num_warps in NUM_WARPS - for num_stages in [2, 3, 4] - ], - key=['H', 'K', 'V', 'BT', 'BK', 'BV'], - **autotune_cache_kwargs, -) -@triton.jit(do_not_specialize=['T']) -def chunk_mesa_net_h_kv_bwd_intra_kernel_dq( - q_star, - k, - v, - beta, - h_kv, - g, - do, - dq, - dg_prev, - dg, - cu_seqlens, - chunk_indices, - B: tl.constexpr, - T, - H: tl.constexpr, - K: tl.constexpr, - V: tl.constexpr, - BT: tl.constexpr, - BK: tl.constexpr, - BV: tl.constexpr, - IS_VARLEN: tl.constexpr, -): - i_t, i_bh = tl.program_id(0), tl.program_id(1) - i_b, i_h = i_bh // H, i_bh % H - 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.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) - T = eos - bos - NT = tl.cdiv(T, BT) - else: - NT = tl.cdiv(T, BT) - i_tg = i_b * NT + i_t - bos, eos = i_b * T, i_b * T + T - - o_t = i_t * BT + tl.arange(0, BT) - m_t = o_t < T - - # offset calculation - v += (bos * H + i_h) * V - do += (bos * H + i_h) * V - h_kv += (i_tg * H + i_h).to(tl.int64) * K*V - q_star += (bos * H + i_h) * K - k += (bos * H + i_h) * K - beta += (bos * H + i_h) - g += bos * H + i_h - dg_prev += bos * H + i_h - dg += bos * H + i_h - dq += (bos * H + i_h) * K - - b_dq = tl.zeros([BT, BK], dtype=tl.float32) - - p_v = tl.make_block_ptr(v, (T, V), (H*V, 1), (i_t * BT, 0), (BT, BV), (1, 0)) - p_k = tl.make_block_ptr(k, (T, K), (H*K, 1), (i_t * BT, 0), (BT, BK), (1, 0)) - p_beta = tl.make_block_ptr(beta, (T, ), (H, ), (i_t * BT,), (BT,), (0,)) - p_do = tl.make_block_ptr(do, (T, V), (H*V, 1), (i_t * BT, 0), (BT, BV), (1, 0)) - p_h = tl.make_block_ptr(h_kv, (V, K), (1, V), (0, 0), (BV, BK), (0, 1)) - p_q = tl.make_block_ptr(q_star, (T, K), (H*K, 1), (i_t * BT, 0), (BT, BK), (1, 0)) - p_g = tl.make_block_ptr(g, (T,), (H,), (i_t * BT,), (BT,), (0,)) - p_dg_prev = tl.make_block_ptr(dg_prev, (T,), (H,), (i_t * BT,), (BT,), (0,)) - p_dg = tl.make_block_ptr(dg, (T,), (H,), (i_t * BT,), (BT,), (0,)) - p_dq = tl.make_block_ptr(dq, (T, K), (H*K, 1), (i_t * BT, 0), (BT, BK), (1, 0)) - - b_q = tl.load(p_q, boundary_check=(0, 1)) - b_v = tl.load(p_v, boundary_check=(0, 1)) - b_k = tl.load(p_k, boundary_check=(0, 1)) - b_beta = tl.load(p_beta, boundary_check=(0, )) - b_g = tl.load(p_g, boundary_check=(0,)) - b_do = tl.load(p_do, boundary_check=(0, 1)) - b_h = tl.load(p_h, boundary_check=(0, 1)) - - b_m = tl.where((o_t[:, None] >= o_t[None, :]) & (m_t[:, None] & m_t[None, :]), exp2(b_g[:, None] - b_g[None, :]), 0) - b_k = (b_k * b_beta[:, None]).to(b_k.dtype) - - b_ds = tl.dot(b_do, tl.trans(b_v)) * b_m - b_g_exp_q = exp2(b_g) - b_dq = tl.dot(b_do, b_h.to(b_do.dtype)) * b_g_exp_q[:, None] - b_dg = tl.sum(b_dq * b_q, axis=1) + tl.load(p_dg_prev, boundary_check=(0,)) - b_dq += tl.dot(b_ds.to(b_k.dtype), b_k) - - tl.store(p_dq, b_dq.to(p_dq.dtype.element_ty), boundary_check=(0, 1)) - tl.store(p_dg, b_dg.to(p_dg.dtype.element_ty), boundary_check=(0,)) - - -def chunk_mesa_net_h_kv_bwd_intra_separate_fn( - q_star, - k, - v, - beta, - h_kv, - dh_kv, - g, - do, - cu_seqlens, - chunk_size=64, - chunk_indices: torch.LongTensor | None = None, -): - B, T, H, K, V = *k.shape, v.shape[-1] - BT = chunk_size - if chunk_indices is None and cu_seqlens is not None: - chunk_indices = prepare_chunk_indices(cu_seqlens, BT) - NT = triton.cdiv(T, BT) if cu_seqlens is None else len(chunk_indices) - - BK = max(triton.next_power_of_2(K), 16) - BV = max(triton.next_power_of_2(V), 16) - dq = torch.empty_like(q_star, dtype=torch.float32) - dk = torch.empty_like(k) - dv = torch.empty_like(v) - dg = torch.empty_like(g) - grid = (NT, B * H) - chunk_mesa_net_h_kv_bwd_intra_kernel_dkv[grid]( - q_star=q_star, - k=k, - v=v, - beta=beta, - h_kv=h_kv, - g=g, - do=do, - dh_kv=dh_kv, - dk_beta=dk, - dg=dg, - dv=dv, - cu_seqlens=cu_seqlens, - chunk_indices=chunk_indices, - B=B, - T=T, - H=H, - K=K, - V=V, - BT=BT, - BK=BK, - BV=BV, - ) - dg_final = torch.empty_like(dg) - chunk_mesa_net_h_kv_bwd_intra_kernel_dq[grid]( - q_star=q_star, - k=k, - v=v, - beta=beta, - h_kv=h_kv, - g=g, - do=do, - dg=dg_final, - dg_prev=dg, - dq=dq, - cu_seqlens=cu_seqlens, - chunk_indices=chunk_indices, - B=B, - T=T, - H=H, - K=K, - V=V, - BT=BT, - BK=BK, - BV=BV, - ) - return dq, dk, dv, dg_final diff --git a/build/torch-xpu/ops/mesa_net/decoding_one_step.py b/build/torch-xpu/ops/mesa_net/decoding_one_step.py deleted file mode 100644 index a2277af73ac435297524367beacba0698889722c..0000000000000000000000000000000000000000 --- a/build/torch-xpu/ops/mesa_net/decoding_one_step.py +++ /dev/null @@ -1,179 +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 -import triton.language as tl - -from ...ops.utils.op import exp -from ...utils import input_guard - - -@triton.jit -def mesa_net_decoding_one_step_kernel( - q, - k, - v, - g, - o, - lamb, - beta, - prev_h_kk, - prev_h_kv, - curr_h_kk, - curr_h_kv, - B: tl.constexpr, - H: tl.constexpr, - K: tl.constexpr, - V: tl.constexpr, - BK: tl.constexpr, - BV: tl.constexpr, - MAX_CG_STEP: tl.constexpr, -): - i_nh = tl.program_id(0) - i_h = i_nh % H - - o_k = tl.arange(0, BK) - o_v = tl.arange(0, BV) - - p_q = q + i_nh * K + o_k - p_k = k + i_nh * K + o_k - p_v = v + i_nh * V + o_v - p_beta = beta + i_nh - p_g = g + i_nh - p_lamb = lamb + i_h * K + o_k - - b_g = exp(tl.load(p_g).to(tl.float32)) - b_beta = tl.load(p_beta).to(tl.float32) - - mask_k = o_k < K - mask_v = o_v < V - mask_kk = mask_k[:, None] & mask_k[None, :] - mask_kv = mask_k[:, None] & mask_v[None, :] - - b_k = tl.load(p_k, mask=mask_k, other=0).to(tl.float32) - b_v = tl.load(p_v, mask=mask_v, other=0).to(tl.float32) - b_q = tl.load(p_q, mask=mask_k, other=0).to(tl.float32) - b_lamb = tl.load(p_lamb, mask=mask_k, other=0).to(tl.float32) - - p_hkk_prev = prev_h_kk + i_nh * K * K + o_k[:, None] * K + o_k[None, :] - b_h_kk = tl.load(p_hkk_prev, mask=mask_kk, other=0).to(tl.float32) - - b_h_kk = b_h_kk * b_g + (b_k * b_beta)[:, None] * b_k[None, :] - - p_hkk_curr = curr_h_kk + i_nh * K * K + o_k[:, None] * K + o_k[None, :] - tl.store(p_hkk_curr, b_h_kk.to(p_hkk_curr.dtype.element_ty), mask=mask_kk) - - p_hkv_prev = prev_h_kv + i_nh * K * V + o_k[:, None] * V + o_v[None, :] - b_h_kv = tl.load(p_hkv_prev, mask=mask_kv, other=0).to(tl.float32) - b_h_kv = b_h_kv * b_g + (b_k * b_beta)[:, None] * b_v[None, :] - p_hkv_curr = curr_h_kv + i_nh * K * V + o_k[:, None] * V + o_v[None, :] - tl.store(p_hkv_curr, b_h_kv.to(p_hkv_curr.dtype.element_ty), mask=mask_kv) - - diag_mask = tl.arange(0, BK)[:, None] == tl.arange(0, BK)[None, :] - diag_mask = diag_mask & mask_kk - b_h_kk_diag = tl.sum(tl.where(diag_mask, b_h_kk, 0.0), axis=1) - - b_x = b_q / (b_h_kk_diag + b_lamb + 1e-5) - b_Hx = tl.sum(b_h_kk * b_x[:, None], axis=0) - b_r = b_q - b_Hx - b_lamb * b_x - b_p = tl.zeros([BK], dtype=tl.float32) - b_p += b_r - delta_old = tl.sum(b_r * b_r) - - for i_iter in range(MAX_CG_STEP): - b_Ap = tl.sum(b_h_kk * b_p[:, None], axis=0) + b_lamb * b_p - pAp = tl.sum(b_p * b_Ap) - alpha = delta_old / (pAp + 1e-5) - b_x = b_x + alpha * b_p - b_r = b_r - alpha * b_Ap - delta_new = tl.sum(b_r * b_r) - beta_cg = delta_new / (delta_old + 1e-5) - b_p = b_r + beta_cg * b_p - delta_old = delta_new - b_o = tl.sum(b_h_kv * b_x[:, None], axis=0) - p_o = o + i_nh * V + o_v - tl.store(p_o, b_o.to(p_o.dtype.element_ty), mask=mask_v) - - -@input_guard -def mesa_net_decoding_one_step( - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - g: torch.Tensor, - lamb: torch.Tensor, - beta: torch.Tensor, - prev_h_kk: torch.Tensor, - prev_h_kv: torch.Tensor, - max_CG_iteration: int = 30, -): - """ - Triton implementation of Mesa Net CG one step - - Args: - q (torch.Tensor): - query tensor [B, H, K] - k (torch.Tensor): - key tensor [B, H, K] - v (torch.Tensor): - value tensor [B, H, V] - g (torch.Tensor): - gate tensor [B, H] - lamb (torch.Tensor): - lambda tensor [H, K] - beta (torch.Tensor): - beta tensor [B, H] - prev_h_kk (torch.Tensor): - previous hidden state KK [B, H, K, K] - prev_h_kv (torch.Tensor): - previous hidden state KV [B, H, K, V] - max_CG_iteration (int): - maximum CG iterations - - Returns: - o (torch.Tensor): - output tensor [B, H, V] - h_kk_new (torch.Tensor): - updated hidden state KK [B, H, K, K] - h_kv_new (torch.Tensor): - updated hidden state KV [B, H, K, V] - """ - B, H, K, V = *q.shape, v.shape[-1] - - o = torch.empty((B, H, V), dtype=q.dtype, device=q.device) - curr_h_kk = torch.empty_like(prev_h_kk) - curr_h_kv = torch.empty_like(prev_h_kv) - - BK = max(triton.next_power_of_2(K), 16) - BV = max(triton.next_power_of_2(V), 16) - - assert BK <= 128 and BV <= 128, "BK and BV must be less than or equal to 128" - - grid = (B * H,) - mesa_net_decoding_one_step_kernel[grid]( - q=q, - k=k, - v=v, - g=g, - o=o, - lamb=lamb, - beta=beta, - prev_h_kk=prev_h_kk, - prev_h_kv=prev_h_kv, - curr_h_kk=curr_h_kk, - curr_h_kv=curr_h_kv, - B=B, - H=H, - K=K, - V=V, - BK=BK, - BV=BV, - MAX_CG_STEP=max_CG_iteration, - num_warps=4 if BK <= 64 else 8, - ) - return o, curr_h_kk, curr_h_kv diff --git a/build/torch-xpu/ops/mesa_net/naive.py b/build/torch-xpu/ops/mesa_net/naive.py deleted file mode 100644 index 0b4649670c35bd0d4a31a5064494f3109c6dbdfd..0000000000000000000000000000000000000000 --- a/build/torch-xpu/ops/mesa_net/naive.py +++ /dev/null @@ -1,135 +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 -from einops import rearrange - - -def naive_mesa_net_decoding_one_step(q, k, v, g, lamb, beta, prev_h_kk, prev_h_kv, max_CG_iteration=30): - q = q.float().clone() - k = k.float().clone() - v = v.float().clone() - g = g.float().clone() - lamb = lamb.float().clone() - beta = beta.float().clone() - B, h, d = q.shape - k_beta = k * beta.unsqueeze(-1) - - h_kk = prev_h_kk * g.exp()[..., None, None] + k_beta.unsqueeze(-1) * k.unsqueeze(-2) - h_kv = prev_h_kv * g.exp()[..., None, None] + k_beta.unsqueeze(-1) * v.unsqueeze(-2) - diag_H = torch.diagonal(h_kk, dim1=-2, dim2=-1) - lamb = lamb.unsqueeze(0) - x = q / (diag_H + lamb) - r = q - (x.unsqueeze(-1) * h_kk).sum(-2) - (lamb * x) - p = r.clone() - delta_old = (r * r).sum(-1) - # CG iteration - for i in range(max_CG_iteration): - q = (p.unsqueeze(-1) * h_kk).sum(-2) + (lamb * p) - alpha = (delta_old / ((p * q).sum(-1) + 1e-5)) - x = x + (alpha[..., None] * p) - r = r - (alpha[..., None] * q) - delta_new = (r * r).sum(-1) - beta = delta_new / (delta_old + 1e-5) - p = r + (beta[..., None] * p) - delta_old = delta_new - o = (x.unsqueeze(-1) * h_kv).sum(-2) - return o, h_kk, h_kv - - -def naive_mesa_net_exact(q, k, v, g, lamb, beta, h_kk_init=None, h_kv_init=None): - B, L, h, d = q.shape - q = q.float() - k = k.float() - v = v.float() - g = g.float() - lamb = lamb.float() - beta = beta.float() - - h_kk = h_kk_init.clone() if h_kk_init is not None else torch.zeros(B, h, d, d, device=q.device) - h_kv = h_kv_init.clone() if h_kv_init is not None else torch.zeros(B, h, d, d, device=q.device) - - h_kk_all = torch.zeros(B, L, h, d, d, device=q.device) - h_kv_all = torch.zeros(B, L, h, d, d, device=q.device) - for i in range(L): - h_kk = h_kk * g[:, i, :, None, None].exp() + (k[:, i, :, :] * beta[:, i, :, None] - )[..., None] * k[:, i, :, None, :] - h_kv = h_kv * g[:, i, :, None, None].exp() + (k[:, i, :, :] * beta[:, i, :, None] - )[..., None] * v[:, i, :, None, :] - h_kk_all[:, i] = h_kk - h_kv_all[:, i] = h_kv - - q_star_gold = torch.linalg.solve(h_kk_all + torch.diag_embed(lamb)[None, None, ...], q) - o_gold = (q_star_gold[..., :, None] * h_kv_all).sum(-2) - return o_gold, h_kk, h_kv - - -def naive_mesa_net_CG(q, k, v, g, lamb, beta, chunk_size, max_CG_iteration=30, h_kk_init=None, h_kv_init=None): - B, L, h, d = q.shape - C = chunk_size - - def chunk_fn(x): return rearrange(x, 'b (n c) h ... -> b h n c ...', c=C).float() - - q_chunk, k_chunk, v_chunk, g_chunk, beta_chunk = map(chunk_fn, [q, k, v, g, beta]) - - g_chunk = g_chunk.cumsum(dim=-1) - - pairwise_decay = (g_chunk[..., None] - g_chunk[..., None, :]).exp().tril() * beta_chunk[..., None, :] - - num_chunks = q_chunk.shape[2] - - h_kv_all = torch.zeros(B, h, num_chunks, d, d, device=q.device) - h_kk_all = torch.zeros(B, h, num_chunks, d, d, device=q.device) - - h_kv = torch.zeros(B, h, d, d, device=q.device) - h_kk = torch.zeros(B, h, d, d, device=q.device) - - if h_kk_init is not None: - h_kk += h_kk_init - if h_kv_init is not None: - h_kv += h_kv_init - - chunk_decay_k = (g_chunk[..., -1, None] - g_chunk).exp() - chunk_decay_q = g_chunk.exp() - - k_chunk_processed = k_chunk * chunk_decay_k[..., None] * beta_chunk[..., None] - - for i in range(num_chunks): - h_kv_all[:, :, i, :, :] = h_kv - h_kk_all[:, :, i, :, :] = h_kk - - k_chunk_i = k_chunk[:, :, i, :, :] - v_chunk_i = v_chunk[:, :, i, :, :] - k_chunk_i_processed = k_chunk_processed[:, :, i, :, :] - - h_kk = h_kk * g_chunk[:, :, i, -1, None, None].exp() + (k_chunk_i_processed).transpose(-2, -1) @ k_chunk_i - h_kv = h_kv * g_chunk[:, :, i, -1, None, None].exp() + (k_chunk_i_processed).transpose(-2, -1) @ v_chunk_i - - # CG solver to approximate the matrix inverse solution. - # diag_H = torch.diagonal(h_kk_all, dim1=-2, dim2=-1) - lamb = lamb[None, :, None, None, :] - x = torch.zeros_like(q_chunk) - r = q_chunk - (x * chunk_decay_q[..., None]) @ h_kk_all - ((x @ k_chunk.transpose(-2, -1)) - * pairwise_decay) @ k_chunk - (lamb * x) - p = r.clone() - delta_old = (r * r).sum(-1) - - # CG iteration - for i in range(max_CG_iteration): - q = (p * chunk_decay_q[..., None]) @ h_kk_all + ((p @ k_chunk.transpose(-1, -2)) - * pairwise_decay) @ k_chunk + (lamb * p) - alpha = (delta_old / ((p * q).sum(-1) + 1e-5)) - x = x + (alpha[..., None] * p) - r = r - (alpha[..., None] * q) - delta_new = (r * r).sum(-1) - beta = delta_new / (delta_old + 1e-5) - p = r + (beta[..., None] * p) - delta_old = delta_new - - o = (x * chunk_decay_q[..., None]) @ h_kv_all + ((x @ k_chunk.transpose(-1, -2)) - * pairwise_decay) @ v_chunk - return rearrange(o, 'b h n c d -> b (n c) h d'), h_kk, h_kv diff --git a/build/torch-xpu/ops/moba/parallel.py b/build/torch-xpu/ops/moba/parallel.py deleted file mode 100644 index 054b96394b30e82fd7b0fc46de7c9eff78a767b0..0000000000000000000000000000000000000000 --- a/build/torch-xpu/ops/moba/parallel.py +++ /dev/null @@ -1,516 +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 -from einops import rearrange - -from ...ops.utils.index import ( - prepare_chunk_indices, - prepare_chunk_offsets, - prepare_cu_seqlens_from_lens, - prepare_lens, - prepare_split_cu_seqlens, -) - -try: - from flash_attn import flash_attn_varlen_func - from flash_attn.flash_attn_interface import _flash_attn_varlen_backward, _flash_attn_varlen_forward -except ImportError: - flash_attn_varlen_func = None - _flash_attn_varlen_backward = None - _flash_attn_varlen_forward = None - - -def prepare_moba_chunks( - cu_seqlens: torch.LongTensor, - chunk_size: int, -) -> tuple[torch.Tensor, torch.Tensor, int, torch.Tensor]: - r"""Split a packed variable-length batch into fixed-size MoBA chunks and emit - the bookkeeping needed by `parallel_moba`. - - Each sample in `cu_seqlens` is split along the token axis into contiguous - windows of length `chunk_size`. A chunk never crosses a sample boundary, - so short samples may produce tail chunks shorter than `chunk_size`. The - **last chunk of every sample** is excluded from the MoBA *target* pool: - no future query token can attend to it (causality), and its own tokens - are still served by the chunk-local self-attn path. - - Args: - cu_seqlens (torch.LongTensor): - Cumulative sequence lengths of shape `[N+1]`, packed-varlen style. - chunk_size (int): - MoBA block size. - - Returns: - A 4-tuple `(cu_chunks, target_chunks, num_target_chunks, chunk_to_batch)`: - - - **cu_chunks** (`torch.Tensor` of shape `[num_chunks + 1]`): - Cumulative token offsets per chunk (a finer-grained `cu_seqlens`). - - **target_chunks** (`torch.Tensor` of shape `[num_target_chunks]`): - Chunk indices that are eligible MoBA targets (all except the last - chunk of each sample). - - **num_target_chunks** (`int`): - Length of `target_chunks`. - - **chunk_to_batch** (`torch.Tensor` of shape `[num_chunks]`): - Sample index each chunk belongs to. - - Example: - Three packed samples of lengths 9, 6, 10 with `chunk_size=4`: - - - Sample 0 → chunks [0,4), [4,8), [8,9) (tail of length 1) - - Sample 1 → chunks [9,13), [13,15) (tail of length 2) - - Sample 2 → chunks [15,19), [19,23), [23,25) (tail of length 2) - - Tails (chunks 2, 4, 7) are dropped from the MoBA target pool, - leaving 5 eligible target chunks. - - >>> cu_seqlens = torch.tensor([0, 9, 15, 25], dtype=torch.int32) - >>> cu_chunks, target_chunks, n_targets, c2b = prepare_moba_chunks(cu_seqlens, chunk_size=4) - >>> cu_chunks # token offsets of all 8 chunks - tensor([ 0, 4, 8, 9, 13, 15, 19, 23, 25], dtype=torch.int32) - >>> target_chunks # MoBA-eligible targets (exclude each sample's tail) - tensor([0, 1, 3, 5, 6]) - >>> n_targets - 5 - >>> c2b # chunk → sample mapping - tensor([0, 0, 0, 1, 1, 2, 2, 2], dtype=torch.int32) - """ - if torch.any(prepare_lens(cu_seqlens) == 0): - raise ValueError("parallel_moba does not support empty sequences in cu_seqlens") - - # chunk_offsets[b] = first chunk id of sample b - chunk_offsets = prepare_chunk_offsets(cu_seqlens, chunk_size) - num_chunks = int(chunk_offsets[-1]) - # cu_chunks[c] = start token offset of chunk c (within the packed tensor) - cu_chunks = prepare_split_cu_seqlens( - split_size=chunk_size, - cu_seqlens=cu_seqlens, - dtype=torch.int32, - device=cu_seqlens.device, - ) - # chunk_to_batch[c] = sample id of chunk c - chunk_to_batch = prepare_chunk_indices(cu_seqlens, chunk_size)[:, 0].to(torch.int32) - - # Drop each sample's last chunk from the MoBA target pool (causality). - is_target = torch.ones(num_chunks, dtype=torch.bool, device=cu_seqlens.device) - is_target[chunk_offsets[1:] - 1] = False - target_chunks = is_target.nonzero(as_tuple=True)[0] - - return cu_chunks, target_chunks, len(target_chunks), chunk_to_batch - - -class ParallelMoBAFunction(torch.autograd.Function): - - @staticmethod - def forward( - ctx, - q, - k, - v, - self_attn_cu_seqlens, - moba_q, - moba_kv, - moba_cu_seqlens_q, - moba_cu_seqlens_k, - max_seqlen, - chunk_size, - moba_q_sh_indices, - ): - ctx.max_seqlen = max_seqlen - ctx.chunk_size = chunk_size - ctx.softmax_scale = softmax_scale = q.shape[-1] ** (-0.5) - - # self attn - self_attn_out_sh, self_attn_lse_hs, _, _ = ( - _flash_attn_varlen_forward( - q=q, - k=k, - v=v, - cu_seqlens_q=self_attn_cu_seqlens, - cu_seqlens_k=self_attn_cu_seqlens, - max_seqlen_q=max_seqlen, - max_seqlen_k=max_seqlen, - softmax_scale=softmax_scale, - causal=True, - dropout_p=0.0, - ) - ) - - # moba attn - moba_attn_out, moba_attn_lse_hs, _, _ = _flash_attn_varlen_forward( - q=moba_q, - k=moba_kv[:, 0], - v=moba_kv[:, 1], - cu_seqlens_q=moba_cu_seqlens_q, - cu_seqlens_k=moba_cu_seqlens_k, - max_seqlen_q=max_seqlen, - max_seqlen_k=chunk_size, - softmax_scale=softmax_scale, - causal=False, - dropout_p=0.0, - ) - - # convert lse shape hs -> sh ( follow the legacy mix attn logic ) - self_attn_lse_sh = self_attn_lse_hs.t().contiguous() - moba_attn_lse = moba_attn_lse_hs.t().contiguous() - - # output buffer [T, H, K], same shape as q - output = torch.zeros_like(q, dtype=torch.float32) - - # flatten T & H for index ops - output_2d = output.view(-1, q.shape[-1]) - - # calc mixed_lse - # minus max lse to avoid exp explosion - max_lse_1d = self_attn_lse_sh.view(-1) - max_lse_1d = max_lse_1d.index_reduce( - 0, moba_q_sh_indices, moba_attn_lse.view(-1), "amax" - ) - self_attn_lse_sh = self_attn_lse_sh - max_lse_1d.view_as(self_attn_lse_sh) - moba_attn_lse = ( - moba_attn_lse.view(-1) - .sub(max_lse_1d.index_select(0, moba_q_sh_indices)) - .reshape_as(moba_attn_lse) - ) - - mixed_attn_se_sh = self_attn_lse_sh.exp() - moba_attn_se = moba_attn_lse.exp() - - mixed_attn_se_sh.view(-1).index_add_( - 0, moba_q_sh_indices, moba_attn_se.view(-1) - ) - mixed_attn_lse_sh = mixed_attn_se_sh.log() - - # add attn output - factor = (self_attn_lse_sh - mixed_attn_lse_sh).exp() # [ T, H ] - self_attn_out_sh = self_attn_out_sh * factor.unsqueeze(-1) - output_2d += self_attn_out_sh.reshape_as(output_2d) - - # add moba output - mixed_attn_lse = ( - mixed_attn_lse_sh.view(-1) - .index_select(0, moba_q_sh_indices) - .view_as(moba_attn_lse) - ) - factor = (moba_attn_lse - mixed_attn_lse).exp() # [ T, H ] - moba_attn_out = moba_attn_out * factor.unsqueeze(-1) - raw_attn_out = moba_attn_out.view(-1, moba_attn_out.shape[-1]) - output_2d.index_add_(0, moba_q_sh_indices, raw_attn_out) - output = output.to(q.dtype) - # add back max lse - mixed_attn_lse_sh = mixed_attn_lse_sh + max_lse_1d.view_as(mixed_attn_se_sh) - ctx.save_for_backward( - output, - mixed_attn_lse_sh, - q, - k, - v, - self_attn_cu_seqlens, - moba_q, - moba_kv, - moba_cu_seqlens_q, - moba_cu_seqlens_k, - moba_q_sh_indices, - ) - - return output - - @staticmethod - def backward(ctx, d_output): - max_seqlen, chunk_size, softmax_scale = ctx.max_seqlen, ctx.chunk_size, ctx.softmax_scale - ( - output, mixed_attn_vlse_sh, q, k, v, self_attn_cu_seqlens, moba_q, - moba_kv, moba_cu_seqlens_q, moba_cu_seqlens_k, moba_q_sh_indices, - ) = ctx.saved_tensors - d_output = d_output.contiguous() - - dq = torch.empty_like(q) - dk = torch.empty_like(k) - dv = torch.empty_like(v) - - _ = _flash_attn_varlen_backward( - dout=d_output, - q=q, - k=k, - v=v, - out=output, - softmax_lse=mixed_attn_vlse_sh.t().contiguous(), - dq=dq, - dk=dk, - dv=dv, - cu_seqlens_q=self_attn_cu_seqlens, - cu_seqlens_k=self_attn_cu_seqlens, - max_seqlen_q=max_seqlen, - max_seqlen_k=max_seqlen, - softmax_scale=softmax_scale, - causal=True, - dropout_p=0.0, - window_size_left=-1, - window_size_right=-1, - softcap=0.0, - alibi_slopes=None, - deterministic=True, - ) - - K = q.shape[-1] - d_moba_output = d_output.view(-1, K).index_select(0, moba_q_sh_indices).unsqueeze(1) - moba_output = output.view(-1, K).index_select(0, moba_q_sh_indices).unsqueeze(1) - mixed_attn_vlse = mixed_attn_vlse_sh.view(-1).index_select(0, moba_q_sh_indices).view(1, -1) - - dmq = torch.empty_like(moba_q) - dmk = torch.empty_like(moba_kv[:, 0]) - dmv = torch.empty_like(moba_kv[:, 1]) - - _ = _flash_attn_varlen_backward( - dout=d_moba_output, - q=moba_q, - k=moba_kv[:, 0], - v=moba_kv[:, 1], - out=moba_output, - softmax_lse=mixed_attn_vlse, - dq=dmq, - dk=dmk, - dv=dmv, - cu_seqlens_q=moba_cu_seqlens_q, - cu_seqlens_k=moba_cu_seqlens_k, - max_seqlen_q=max_seqlen, - max_seqlen_k=chunk_size, - softmax_scale=softmax_scale, - causal=False, - dropout_p=0.0, - window_size_left=-1, - window_size_right=-1, - softcap=0.0, - alibi_slopes=None, - deterministic=True, - ) - - dmkv = torch.stack((dmk, dmv), dim=1) - return dq, dk, dv, None, dmq, dmkv, None, None, None, None, None - - -def parallel_moba( - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - cu_seqlens: torch.LongTensor, - max_seqlen: int, - chunk_size: int, - topk: int, -) -> torch.Tensor: - r"""Flash-attn based MoBA implementation. - - MoBA (Mixture of Block Attention, https://arxiv.org/abs/2502.13189) makes - each query token attend only to a small subset of KV blocks, rather than - the full causal history. The computation is split into two flash-attn - passes that are merged via an online-softmax LSE combine: - - 1. **Chunking.** Split the packed sequence into fixed-size blocks of length - `chunk_size` using `prepare_moba_chunks`. Each sample's last chunk is - reserved for self-attn (it cannot be a MoBA target — causality). - 2. **Gate scoring.** Summarize every candidate block by the mean of its - keys (one representative key per block), then score every query token - against every representative via a dot product. Mask out future blocks - and blocks outside the query's own sample. - 3. **Top-k selection.** For each (query token, head) pair keep the - `topk - 1` highest-scoring blocks (the local block is always served - by self-attn, so we only need `topk - 1` extra blocks from MoBA). - 4. **Two-stream attention + online merge.** - - *Self-attn stream*: chunk-local causal flash-attn within each block. - - *MoBA stream*: gathered queries attend to their selected blocks' - full KV via varlen flash-attn. - - The per-position outputs from the two streams are combined with a - stable LSE-based online softmax, yielding the same result as a - single softmax over the union of attended keys. - - Edge cases: - - If `topk - 1 <= 0` after capping at the number of filtered chunks, - the op short-circuits to a plain causal `flash_attn_varlen_func` - over the full sequence (no block selection needed). - - Args: - q (torch.Tensor): - Queries of shape `[B, T, H, K]`. When `cu_seqlens` is provided, B - must be 1 and all sequences are packed along `T`, following the - FlashAttention varlen convention. - k (torch.Tensor): - Keys of shape `[B, T, H, K]`. - v (torch.Tensor): - Values of shape `[B, T, H, V]`. - cu_seqlens (torch.LongTensor): - Cumulative sequence lengths of shape `[N+1]` used for variable-length - training, consistent with the FlashAttention API. - max_seqlen (int): - Max sequence length across the packed batch, consistent with the - FlashAttention API. - chunk_size (int): - Size of each MoBA block. - topk (int): - Number of blocks each query attends to, counting the local block. - With `topk=1` the op degenerates to full causal self-attention. - - Returns: - torch.Tensor: Output of shape `[B, T, H, V]`, same dtype as `q`. - - Raises: - ImportError: If `flash-attn` is not installed. - ValueError: If `q.shape[0] != 1` while `cu_seqlens` is provided. - - Example: - Single packed sample of 1024 tokens, 4 heads, head-dim 64, chunked - into 128-token blocks with each query attending to 4 blocks: - - >>> import torch - >>> from ...ops.moba import parallel_moba - >>> B, T, H, D = 1, 1024, 4, 64 - >>> q = torch.randn(B, T, H, D, dtype=torch.float16, device='cuda') - >>> k = torch.randn(B, T, H, D, dtype=torch.float16, device='cuda') - >>> v = torch.randn(B, T, H, D, dtype=torch.float16, device='cuda') - >>> cu_seqlens = torch.tensor([0, T], dtype=torch.int32, device='cuda') - >>> o = parallel_moba(q, k, v, cu_seqlens, max_seqlen=T, chunk_size=128, topk=4) - >>> o.shape - torch.Size([1, 1024, 4, 64]) - - Two packed samples of lengths 512 and 256: - - >>> cu_seqlens = torch.tensor([0, 512, 768], dtype=torch.int32, device='cuda') - >>> q = torch.randn(1, 768, H, D, dtype=torch.float16, device='cuda') - >>> k = torch.randn(1, 768, H, D, dtype=torch.float16, device='cuda') - >>> v = torch.randn(1, 768, H, D, dtype=torch.float16, device='cuda') - >>> o = parallel_moba(q, k, v, cu_seqlens, max_seqlen=512, chunk_size=128, topk=4) - >>> o.shape - torch.Size([1, 768, 4, 64]) - """ - if flash_attn_varlen_func is None: - raise ImportError( - "`parallel_moba` requires `flash-attn`. Install it via `pip install flash-attn`." - ) - if cu_seqlens is not None and q.shape[0] != 1: - raise ValueError( - f"The batch size is expected to be 1 rather than {q.shape[0]} when using `cu_seqlens`. " - f"Please flatten variable-length inputs before processing.", - ) - - # The underlying `_flash_attn_varlen_forward/backward` kernels expect packed - # 3-D `[total_T, H, D]`; squeeze the leading batch dim here and restore it - # on the output. - q, k, v = q.squeeze(0), k.squeeze(0), v.squeeze(0) - T, H, K = q.shape - - # prepare chunk meta - cu_chunks, target_chunks, num_target_chunks, chunk_to_batch = ( - prepare_moba_chunks(cu_seqlens, chunk_size) - ) - - # the last chunk is always chosen by self-attn, so we only need `topk - 1` from MoBA - topk = min(topk - 1, num_target_chunks) - - # corner case: no MoBA chunks selectable, fall back to plain causal self-attn - if topk <= 0: - return flash_attn_varlen_func( - q, k, v, cu_seqlens, cu_seqlens, max_seqlen, max_seqlen, causal=True - ).unsqueeze(0) - - kv = torch.stack((k, v), dim=1) - - self_attn_cu_seqlens = cu_chunks - - # filtered_kv is a dense matrix that only contains filtered chunk of kv - filtered_kv_indices = torch.arange( - 0, chunk_size, dtype=torch.int32, device=q.device - )[None, :].repeat(num_target_chunks, 1) - filtered_kv_indices += cu_chunks[target_chunks][:, None] - filtered_kv = kv.index_select(0, filtered_kv_indices.view(-1)) - - # key_gate_weight [ F_N_CHUNK, H, K ], float32 for better gate logit perception - key_gate_weight = ( - filtered_kv[:, 0] - .view(num_target_chunks, chunk_size, H, K) - .mean(dim=1) - .float() - ) - q = q.float() - # gate [ F_N_CHUNK, H, T ] - gate = torch.einsum("nhk,thk->nht", key_gate_weight, q) - key_gate_weight = key_gate_weight.type_as(k) - q = q.type_as(k) - - # mask out chunks that lie outside the current sequence, and the current chunk itself - gate_seq_idx = torch.arange(0, T, device=q.device, dtype=torch.int32)[ - None, : - ].repeat(num_target_chunks, 1) - chunk_end = cu_chunks[target_chunks + 1] - batch_end = cu_seqlens[chunk_to_batch[target_chunks] + 1] - gate_chunk_end_mask = gate_seq_idx < chunk_end[:, None] - gate_batch_end_mask = gate_seq_idx >= batch_end[:, None] - gate_inf_mask = gate_chunk_end_mask | gate_batch_end_mask - gate.masked_fill_(gate_inf_mask.unsqueeze(1), -float("inf")) - - # find topk chunks per (head, token), then AND with the causal mask - # gate_mask [ N_CHUNK, H, T ], True means the (chunk, head, token) triple participates in MoBA attn - _, gate_top_k_idx = torch.topk(gate, k=topk, dim=0, largest=True, sorted=False) - gate_mask = torch.logical_not(gate.isinf()) - gate_idx_mask = torch.zeros_like(gate_mask).scatter_(dim=0, index=gate_top_k_idx, value=True) - gate_mask = torch.logical_and(gate_mask, gate_idx_mask) - - # varlen trick: combining all q index that needs moba attn - # the result will be like [ C0H0 ][ C0H1 ][ C0H2 ][ ... ][ CnHm ] - moba_q_indices = gate_mask.reshape(gate_mask.shape[0], -1).nonzero(as_tuple=True)[-1] # [HT] * N - # moba_seqlens_q[i]: number of q tokens selected for the i-th kv (chunk, head) pair - moba_seqlens_q = gate_mask.sum(dim=-1).flatten() - # gather the selected q tokens, shape [ selected_T, K ] - moba_q = rearrange(q, "t h k -> (h t) k").index_select(0, moba_q_indices) - moba_q = moba_q.unsqueeze(1) - # moba_q_sh_indices: position of each gathered q token inside the original q tensor - moba_q_sh_indices = moba_q_indices % T * H + moba_q_indices // T - - # reorganize kv to align with moba_q (grouped as (H, chunk) pairs) - - # cut off (chunk, head) pairs whose q selection is empty - q_zero_mask = moba_seqlens_q == 0 - valid_expert_mask = ~q_zero_mask - zero_expert_count = q_zero_mask.sum() - if zero_expert_count > 0: - moba_seqlens_q = moba_seqlens_q[valid_expert_mask] - # moba cu_seqlens for flash-attn varlen - moba_cu_seqlens_q = prepare_cu_seqlens_from_lens(moba_seqlens_q) - moba_kv = rearrange(filtered_kv, "t x h k -> h t x k") - moba_kv = moba_kv.split(chunk_size, dim=1) - moba_kv = torch.cat(moba_kv, dim=0) - if zero_expert_count > 0: - assert valid_expert_mask.sum() == moba_kv.shape[0] - zero_expert_count - # drop (chunk, head) pairs with zero q, otherwise grads may be nan - moba_kv = moba_kv[valid_expert_mask] - moba_kv = moba_kv.flatten(start_dim=0, end_dim=1).unsqueeze(2) - moba_cu_seqlens_k = ( - torch.arange( - 0, - num_target_chunks * H + 1 - zero_expert_count, - dtype=torch.int32, - device=q.device, - ) - * chunk_size - ) - - assert moba_cu_seqlens_k.shape == moba_cu_seqlens_q.shape, ( - f"moba_cu_seqlens_k.shape != moba_cu_seqlens_q.shape " - f"{moba_cu_seqlens_k.shape} != {moba_cu_seqlens_q.shape}" - ) - - return ParallelMoBAFunction.apply( - q, - k, - v, - self_attn_cu_seqlens, - moba_q, - moba_kv, - moba_cu_seqlens_q, - moba_cu_seqlens_k, - max_seqlen, - chunk_size, - moba_q_sh_indices, - ).unsqueeze(0) diff --git a/build/torch-xpu/ops/nsa/__init__.py b/build/torch-xpu/ops/nsa/__init__.py deleted file mode 100644 index f5c634b15f2bc3abe0d1ca963ce0b205573446ec..0000000000000000000000000000000000000000 --- a/build/torch-xpu/ops/nsa/__init__.py +++ /dev/null @@ -1,14 +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 .naive import naive_nsa -from .parallel import parallel_nsa - -__all__ = [ - 'naive_nsa', - 'parallel_nsa', -] diff --git a/build/torch-xpu/ops/nsa/compression.py b/build/torch-xpu/ops/nsa/compression.py deleted file mode 100644 index 6f51d36d41ddcd43595c85ef6783cefeffcbdbf2..0000000000000000000000000000000000000000 --- a/build/torch-xpu/ops/nsa/compression.py +++ /dev/null @@ -1,572 +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 -import triton.language as tl - -from ...ops.attn.parallel import parallel_attn_bwd_preprocess -from ...ops.utils import prepare_chunk_indices, prepare_chunk_offsets, prepare_token_indices -from ...ops.utils.op import exp, log -from ...utils import autocast_custom_bwd, autocast_custom_fwd, autotune_cache_kwargs, check_shared_mem, contiguous - - -@triton.heuristics({ - 'IS_VARLEN': lambda args: args['cu_seqlens_q'] is not None, -}) -@triton.autotune( - configs=[ - triton.Config({}, num_warps=num_warps) - for num_warps in [1, 2, 4] - ], - key=['BS', 'BK', 'BV'], - **autotune_cache_kwargs, -) -@triton.jit -def parallel_nsa_compression_fwd_kernel( - q, - k, - v, - o, - lse, - scale, - cu_seqlens_q, - cu_seqlens_k, - token_indices_q, - chunk_offsets, - TQ, - TK, - H: tl.constexpr, - HQ: tl.constexpr, - G: tl.constexpr, - K: tl.constexpr, - V: tl.constexpr, - BC: tl.constexpr, - BS: tl.constexpr, - BK: tl.constexpr, - BV: tl.constexpr, - IS_VARLEN: tl.constexpr, -): - i_t, i_v, i_bh = tl.program_id(0), tl.program_id(1), tl.program_id(2) - i_b, i_h = i_bh // H, i_bh % H - - if IS_VARLEN: - i_n, i_t = tl.load(token_indices_q + i_t * 2).to(tl.int32), tl.load(token_indices_q + i_t * 2 + 1).to(tl.int32) - bos_q, eos_q = tl.load(cu_seqlens_q + i_n).to(tl.int64), tl.load(cu_seqlens_q + i_n + 1).to(tl.int64) - bos_k, eos_k = tl.load(cu_seqlens_k + i_n).to(tl.int64), tl.load(cu_seqlens_k + i_n + 1).to(tl.int64) - TQ = (eos_q - bos_q).to(tl.int32) - TK = (eos_k - bos_k).to(tl.int32) - TC = tl.cdiv(TK, BS) - boc = tl.load(chunk_offsets + i_n).to(tl.int32) - else: - bos_q, eos_q = (i_b * TQ).to(tl.int64), (i_b * TQ + TQ).to(tl.int64) - TC = tl.cdiv(TK, BS) - boc = i_b * TC - - p_q = tl.make_block_ptr(q + (bos_q + i_t) * HQ*K, (HQ, K), (K, 1), (i_h * G, 0), (G, BK), (1, 0)) - Q_OFFSET = TK - TQ - # the Q block is kept in the shared memory throughout the whole kernel - # [G, BK] - b_q = tl.load(p_q, boundary_check=(0, 1)) - b_q = (b_q * scale).to(b_q.dtype) - - # number of complete compression blocks visible to the query (q tokens are the last TQ of the sequence) - NC = (i_t + Q_OFFSET + 1) // BS - - p_o = tl.make_block_ptr(o + (bos_q + i_t) * HQ*V, (HQ, V), (V, 1), (i_h * G, i_v * BV), (G, BV), (1, 0)) - # [G, BV] - b_o = tl.zeros([G, BV], dtype=tl.float32) - # max scores for the current block - b_m = tl.full([G], float('-inf'), dtype=tl.float32) - # lse = log(acc) + m - b_acc = tl.zeros([G], dtype=tl.float32) - - for i_c in range(0, NC, BC): - o_c = i_c + tl.arange(0, BC) - - p_k = tl.make_block_ptr(k + (boc * H + i_h) * K, (K, TC), (1, H*K), (0, i_c), (BK, BC), (0, 1)) - p_v = tl.make_block_ptr(v + (boc * H + i_h) * V, (TC, V), (H*V, 1), (i_c, i_v * BV), (BC, BV), (1, 0)) - # [BK, BC] - b_k = tl.load(p_k, boundary_check=(0, 1)) - # [BC, BV] - b_v = tl.load(p_v, boundary_check=(0, 1)) - # [G, BC] - b_s = tl.dot(b_q, b_k) - # causal mask: only the NC complete blocks are visible - b_s = tl.where((o_c < NC)[None, :], b_s, float('-inf')) - - # [G] - b_m, b_mp = tl.maximum(b_m, tl.max(b_s, 1)), b_m - b_r = exp(b_mp - b_m) - # [G, BC] - b_p = exp(b_s - b_m[:, None]) - # [G] - b_acc = b_acc * b_r + tl.sum(b_p, 1) - # [G, BV] - b_o = b_o * b_r[:, None] + tl.dot(b_p.to(b_q.dtype), b_v) - if NC == 0: - b_lse = tl.zeros([G], dtype=tl.float32) - else: - b_o = b_o / b_acc[:, None] - b_lse = b_m + log(b_acc) - - tl.store(p_o, b_o.to(p_o.dtype.element_ty), boundary_check=(0, 1)) - if i_v == 0: - tl.store(lse + (bos_q + i_t) * HQ + i_h * G + tl.arange(0, G), b_lse.to(lse.dtype.element_ty)) - - -@triton.heuristics({ - 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None, -}) -@triton.autotune( - configs=[ - triton.Config({}, num_warps=num_warps) - for num_warps in [1, 2, 4] - ], - key=['BS', 'BK', 'BV'], - **autotune_cache_kwargs, -) -@triton.jit(do_not_specialize=['T']) -def parallel_nsa_compression_bwd_kernel_dq( - q, - k, - v, - lse, - delta, - do, - dq, - scale, - cu_seqlens, - token_indices, - chunk_offsets, - T, - B: tl.constexpr, - H: tl.constexpr, - HQ: tl.constexpr, - G: tl.constexpr, - K: tl.constexpr, - V: tl.constexpr, - BC: tl.constexpr, - BS: tl.constexpr, - BK: tl.constexpr, - BV: tl.constexpr, - IS_VARLEN: tl.constexpr, -): - i_t, i_v, i_bh = tl.program_id(0), tl.program_id(1), tl.program_id(2) - i_b, i_h = i_bh // H, i_bh % H - - all = B * T - if IS_VARLEN: - i_n, i_t = tl.load(token_indices + i_t * 2).to(tl.int32), tl.load(token_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).to(tl.int32) - boc = tl.load(chunk_offsets + i_n).to(tl.int32) - else: - bos, eos = (i_b * T).to(tl.int64), (i_b * T + T).to(tl.int64) - boc = i_b * tl.cdiv(T, BS) - - q += (bos + i_t) * HQ*K - do += (bos + i_t) * HQ*V - lse += (bos + i_t) * HQ - delta += (bos + i_t) * HQ - dq += (i_v * all + bos + i_t) * HQ*K - - p_q = tl.make_block_ptr(q, (HQ, K), (K, 1), (i_h * G, 0), (G, BK), (1, 0)) - p_dq = tl.make_block_ptr(dq, (HQ, K), (K, 1), (i_h * G, 0), (G, BK), (1, 0)) - - # [G, BK] - b_q = tl.load(p_q, boundary_check=(0, 1)) - b_q = (b_q * scale).to(b_q.dtype) - - p_do = tl.make_block_ptr(do, (HQ, V), (V, 1), (i_h * G, i_v * BV), (G, BV), (1, 0)) - p_lse = lse + i_h * G + tl.arange(0, G) - p_delta = delta + i_h * G + tl.arange(0, G) - - # the number of compression representations in total - TC = tl.cdiv(T, BS) - # the number of compression representations required to iterate over - # incomplete compression blocks are not included - NC = (i_t + 1) // BS - - # [G, BV] - b_do = tl.load(p_do, boundary_check=(0, 1)) - # [G] - b_lse = tl.load(p_lse) - b_delta = tl.load(p_delta) - - # [G, BK] - b_dq = tl.zeros([G, BK], dtype=tl.float32) - for i_c in range(0, NC, BC): - o_c = i_c + tl.arange(0, BC) - p_k = tl.make_block_ptr(k + (boc * H + i_h) * K, (K, TC), (1, H*K), (0, i_c), (BK, BC), (0, 1)) - p_v = tl.make_block_ptr(v + (boc * H + i_h) * V, (V, TC), (1, H*V), (i_v * BV, i_c), (BV, BC), (0, 1)) - # [BK, BC] - b_k = tl.load(p_k, boundary_check=(0, 1)) - # [BV, BC] - b_v = tl.load(p_v, boundary_check=(0, 1)) - - # [G, BC] - b_s = tl.dot(b_q, b_k) - b_p = exp(b_s - b_lse[:, None]) - b_p = tl.where((o_c < NC)[None, :], b_p, 0) - - # [G, BV] @ [BV, BC] -> [G, BC] - b_dp = tl.dot(b_do, b_v) - b_ds = b_p * (b_dp.to(tl.float32) - b_delta[:, None]) - # [G, BC] @ [BC, BK] -> [G, BK] - b_dq += tl.dot(b_ds.to(b_k.dtype), tl.trans(b_k)) - b_dq *= scale - tl.store(p_dq, b_dq.to(p_dq.dtype.element_ty), boundary_check=(0, 1)) - - -@triton.heuristics({ - 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None, -}) -@triton.autotune( - configs=[ - triton.Config({}, num_warps=num_warps) - for num_warps in [1, 2, 4] - ], - key=['BS', 'BK', 'BV'], - **autotune_cache_kwargs, -) -@triton.jit(do_not_specialize=['T', 'TC']) -def parallel_nsa_compression_bwd_kernel_dkv( - q, - k, - v, - lse, - delta, - do, - dk, - dv, - cu_seqlens, - chunk_indices, - chunk_offsets, - scale, - T, - TC, - B: tl.constexpr, - H: tl.constexpr, - HQ: tl.constexpr, - G: tl.constexpr, - K: tl.constexpr, - V: tl.constexpr, - BC: tl.constexpr, - BS: tl.constexpr, - BK: tl.constexpr, - BV: tl.constexpr, - IS_VARLEN: tl.constexpr, -): - i_v, i_c, i_bh = tl.program_id(0), tl.program_id(1), tl.program_id(2) - i_b, i_h = i_bh // H, i_bh % H - - all = B * TC - - if IS_VARLEN: - i_n, i_c = tl.load(chunk_indices + i_c * 2).to(tl.int32), tl.load(chunk_indices + i_c * 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).to(tl.int32) - # the number of compression representations in total - TC = tl.cdiv(T, BS) - boc = tl.load(chunk_offsets + i_n).to(tl.int32) - else: - bos, eos = (i_b * T).to(tl.int64), (i_b * T + T).to(tl.int64) - boc = i_b * tl.cdiv(T, BS) - - p_k = tl.make_block_ptr(k + (boc * H + i_h) * K, (TC, K), (H*K, 1), (i_c * BC, 0), (BC, BK), (1, 0)) - p_v = tl.make_block_ptr(v + (boc * H + i_h) * V, (TC, V), (H*V, 1), (i_c * BC, i_v * BV), (BC, BV), (1, 0)) - p_dk = tl.make_block_ptr(dk + (i_v * all*H + boc * H + i_h) * K, (TC, K), (H*K, 1), (i_c * BC, 0), (BC, BK), (1, 0)) - p_dv = tl.make_block_ptr(dv + (boc * H + i_h) * V, (TC, V), (H*V, 1), (i_c * BC, i_v * BV), (BC, BV), (1, 0)) - - # [BC, BK] - b_k = tl.load(p_k, boundary_check=(0, 1)) - b_dk = tl.zeros([BC, BK], dtype=tl.float32) - # [BC, BV] - b_v = tl.load(p_v, boundary_check=(0, 1)) - b_dv = tl.zeros([BC, BV], dtype=tl.float32) - - for i in range(i_c * BC * BS, T): - o_c = i_c * BC + tl.arange(0, BC) - - p_q = tl.make_block_ptr(q + (bos + i) * HQ*K, (HQ, K), (K, 1), (i_h * G, 0), (G, BK), (1, 0)) - # [G, BK] - b_q = tl.load(p_q, boundary_check=(0, 1)) - b_q = (b_q * scale).to(b_q.dtype) - - p_do = tl.make_block_ptr(do + (bos + i) * HQ*V, (HQ, V), (V, 1), (i_h * G, i_v * BV), (G, BV), (1, 0)) - p_lse = lse + (bos + i) * HQ + i_h * G + tl.arange(0, G) - p_delta = delta + (bos + i) * HQ + i_h * G + tl.arange(0, G) - # [G, BV] - b_do = tl.load(p_do, boundary_check=(0, 1)) - # [G] - b_lse = tl.load(p_lse) - b_delta = tl.load(p_delta) - # [BC, G] - b_s = tl.dot(b_k, tl.trans(b_q)) - b_p = exp(b_s - b_lse[None, :]) - b_p = tl.where((i >= max(0, (o_c + 1) * BS - 1))[:, None], b_p, 0) - # [BC, G] @ [G, BV] -> [BC, BV] - b_dv += tl.dot(b_p.to(b_do.dtype), b_do) - # [BC, BV] @ [BV, G] -> [BC, G] - b_dp = tl.dot(b_v, tl.trans(b_do)) - # [BC, G] - b_ds = b_p * (b_dp - b_delta[None, :]) - # [BC, G] @ [G, BK] -> [BC, BK] - b_dk += tl.dot(b_ds.to(b_q.dtype), b_q) - - tl.store(p_dk, b_dk.to(p_dk.dtype.element_ty), boundary_check=(0, 1)) - tl.store(p_dv, b_dv.to(p_dv.dtype.element_ty), boundary_check=(0, 1)) - - -def parallel_nsa_compression_fwd( - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - TK: int, - block_size: int, - scale: float, - cu_seqlens_q: torch.LongTensor | None = None, - cu_seqlens_k: torch.LongTensor | None = None, - token_indices_q: torch.LongTensor | None = None, -): - B, TQ, HQ, K, V = *q.shape, v.shape[-1] - H = k.shape[2] - G = HQ // H - BC = BS = block_size - if check_shared_mem('hopper', q.device.index): - BK = min(256, triton.next_power_of_2(K)) - BV = min(256, triton.next_power_of_2(V)) - else: - BK = min(128, triton.next_power_of_2(K)) - BV = min(128, triton.next_power_of_2(V)) - NK = triton.cdiv(K, BK) - NV = triton.cdiv(V, BV) - assert NK == 1, "The key dimension can not be larger than 256" - - chunk_offsets = prepare_chunk_offsets(cu_seqlens_k, BS) if cu_seqlens_k is not None else None - - grid = (TQ, NV, B * H) - o = torch.empty(B, TQ, HQ, V, dtype=v.dtype, device=q.device) - lse = torch.empty(B, TQ, HQ, dtype=torch.float, device=q.device) - - parallel_nsa_compression_fwd_kernel[grid]( - q=q, - k=k, - v=v, - o=o, - lse=lse, - scale=scale, - cu_seqlens_q=cu_seqlens_q, - cu_seqlens_k=cu_seqlens_k, - token_indices_q=token_indices_q, - chunk_offsets=chunk_offsets, - TQ=TQ, - TK=TK, - H=H, - HQ=HQ, - G=G, - K=K, - V=V, - BC=BC, - BS=BS, - BK=BK, - BV=BV, - ) - return o, lse - - -def parallel_nsa_compression_bwd( - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - o: torch.Tensor, - lse: torch.Tensor, - do: torch.Tensor, - block_size: int = 64, - scale: float = None, - cu_seqlens: torch.LongTensor | None = None, - token_indices: torch.LongTensor | None = None, - chunk_indices: torch.LongTensor | None = None, -): - B, T, HQ, K, V = *q.shape, v.shape[-1] - TC = k.shape[1] - H = k.shape[2] - G = HQ // H - BC = BS = block_size - BK = max(triton.next_power_of_2(K), 16) - BV = min(128, max(triton.next_power_of_2(v.shape[-1]), 16)) - NV = triton.cdiv(V, BV) - if cu_seqlens is not None: - chunk_offsets = prepare_chunk_offsets(cu_seqlens, BS) - if chunk_indices is None: - chunk_indices = prepare_chunk_indices(chunk_offsets, BC) - NC = len(chunk_indices) - else: - chunk_indices, chunk_offsets = None, None - NC = triton.cdiv(triton.cdiv(T, BS), BC) - - delta = parallel_attn_bwd_preprocess(o, do) - - dq = torch.empty(NV, *q.shape, dtype=q.dtype if NV == 1 else torch.float, device=q.device) - grid = (T, NV, B * H) - parallel_nsa_compression_bwd_kernel_dq[grid]( - q=q, - k=k, - v=v, - lse=lse, - delta=delta, - do=do, - dq=dq, - scale=scale, - cu_seqlens=cu_seqlens, - token_indices=token_indices, - chunk_offsets=chunk_offsets, - T=T, - B=B, - H=H, - HQ=HQ, - G=G, - K=K, - V=V, - BC=BC, - BS=BS, - BK=BK, - BV=BV, - ) - dq = dq.sum(0) - - dk = torch.empty(NV, *k.shape, dtype=k.dtype if NV == 1 else torch.float, device=q.device) - dv = torch.empty(v.shape, dtype=v.dtype, device=q.device) - - grid = (NV, NC, B * H) - parallel_nsa_compression_bwd_kernel_dkv[grid]( - q=q, - k=k, - v=v, - lse=lse, - delta=delta, - do=do, - dk=dk, - dv=dv, - cu_seqlens=cu_seqlens, - chunk_indices=chunk_indices, - chunk_offsets=chunk_offsets, - scale=scale, - T=T, - TC=TC, - B=B, - H=H, - HQ=HQ, - G=G, - K=K, - V=V, - BC=BC, - BS=BS, - BK=BK, - BV=BV, - ) - dk = dk.sum(0) - return dq, dk, dv - - -class ParallelNSACompressionFunction(torch.autograd.Function): - - @staticmethod - @contiguous - @autocast_custom_fwd - def forward( - ctx, - q, - k, - v, - TK, - block_size, - scale, - cu_seqlens, - ): - ctx.dtype = q.dtype - - # 2-d sequence indices denoting the cu_seqlens of tokens in each sequence - # for example, if the passed `cu_seqlens` is [0, 2, 6], - # then there are 2 and 4 tokens in the 1st and 2nd sequences respectively, and `token_indices` will be - # [[0, 0], [0, 1], [1, 0], [1, 1], [1, 2], [1, 3]] - if cu_seqlens is not None: - if isinstance(cu_seqlens, tuple): - cu_seqlens_q, cu_seqlens_k = cu_seqlens - else: - cu_seqlens_q = cu_seqlens_k = cu_seqlens - token_indices_q = prepare_token_indices(cu_seqlens_q) - else: - cu_seqlens_q = cu_seqlens_k = token_indices_q = None - - o, lse = parallel_nsa_compression_fwd( - q=q, - k=k, - v=v, - TK=TK, - block_size=block_size, - scale=scale, - cu_seqlens_q=cu_seqlens_q, - cu_seqlens_k=cu_seqlens_k, - token_indices_q=token_indices_q - ) - ctx.save_for_backward(q, k, v, o, lse) - # Use cu_seqlens of q in backward, as cu_seqlens for q & k are different only for inference - ctx.cu_seqlens = cu_seqlens_q - ctx.token_indices = token_indices_q - ctx.block_size = block_size - ctx.scale = scale - # q/k cu_seqlens differ only in cached inference (TQ != TK), where backward is not supported - ctx.tq_ne_tk = isinstance(cu_seqlens, tuple) - return o.to(q.dtype), lse - - @staticmethod - @contiguous - @autocast_custom_bwd - def backward(ctx, do, *args): - if ctx.tq_ne_tk: - raise NotImplementedError( - "Backward is not supported when `cu_seqlens` differs for queries and keys (cached inference). " - "Run the forward under `torch.no_grad()`." - ) - q, k, v, o, lse = ctx.saved_tensors - dq, dk, dv = parallel_nsa_compression_bwd( - q=q, - k=k, - v=v, - o=o, - lse=lse, - do=do, - block_size=ctx.block_size, - scale=ctx.scale, - cu_seqlens=ctx.cu_seqlens, - token_indices=ctx.token_indices, - ) - return dq.to(q), dk.to(k), dv.to(v), None, None, None, None - - -def parallel_nsa_compression( - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - TK: int, - block_size: int = 64, - scale: float = None, - cu_seqlens: torch.LongTensor | tuple[torch.LongTensor, torch.LongTensor] | None = None -): - if scale is None: - scale = k.shape[-1] ** -0.5 - return ParallelNSACompressionFunction.apply( - q, - k, - v, - TK, - block_size, - scale, - cu_seqlens, - ) diff --git a/build/torch-xpu/ops/nsa/naive.py b/build/torch-xpu/ops/nsa/naive.py deleted file mode 100644 index 9c155db6809518535c8ca245b9dbb2c7e8b5ae58..0000000000000000000000000000000000000000 --- a/build/torch-xpu/ops/nsa/naive.py +++ /dev/null @@ -1,426 +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 warnings - -import torch -from einops import repeat - -from ...ops.utils import prepare_chunk_offsets -from ...ops.utils.pooling import mean_pooling - -try: - from flash_attn import flash_attn_func, flash_attn_varlen_func -except ImportError: - warnings.warn( - "Flash Attention is not installed. Please install it via `pip install flash-attn --no-build-isolation`", - category=ImportWarning - ) - flash_attn_func = flash_attn_varlen_func = None - - -def naive_nsa_selection( - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - block_indices: torch.LongTensor, - block_size: int = 64, - scale: float | None = None, - cu_seqlens: torch.LongTensor | tuple[torch.LongTensor, torch.LongTensor] | None = None, - **kwargs, -) -> torch.Tensor: - r""" - Args: - q (torch.Tensor): - queries of shape `[B, TQ, HQ, K]`. - k (torch.Tensor): - keys of shape `[B, T, H, K]`. - GQA is enforced here: the ratio of query heads (HQ) to key/value heads (H) must be a power of 2 and >= 16. - This is a kernel tile dimension, which Triton requires to be a power-of-2 block shape. - v (torch.Tensor): - values of shape `[B, T, H, V]`. - block_indices (torch.LongTensor): - Block indices of shape `[B, TQ, H, S]`. - `S` is the number of selected blocks for each query token, which is set to 16 in the paper. - block_size (int): - Selected block size. Default: `64`. - scale (Optional[float]): - Scale factor for attention scores. - If not provided, it will default to `1 / sqrt(K)`. Default: `None`. - cu_seqlens (torch.LongTensor, Tuple[torch.LongTensor, torch.LongTensor] or None): - Cumulative sequence lengths of shape `[N+1]` used for variable-length training, - consistent with the FlashAttention API. - When a tuple is provided, it should contain two tensors: `(cu_seqlens_q, cu_seqlens_k)`. - - Returns: - o (torch.Tensor): - Outputs of shape `[B, TQ, HQ, V]`. - """ - if 'head_first' in kwargs: - raise DeprecationWarning( - "head_first has been removed. Inputs must be in `[B, T, H, ...]` format.", - ) - if scale is None: - scale = k.shape[-1] ** -0.5 - - dtype = q.dtype - G = q.shape[2] // k.shape[2] - BS = block_size - k, v, block_indices = (repeat(x, 'b t h d -> b t (h g) d', g=G) for x in (k, v, block_indices)) - q, k, v = map(lambda x: x.float(), (q, k, v)) - B = q.shape[0] - - o = torch.zeros_like(v) - varlen = True - if cu_seqlens is None: - varlen = False - TQ = TK = q.shape[1] - cu_q = torch.cat([ - block_indices.new_tensor(range(0, B * TQ, TQ)), block_indices.new_tensor([B * TQ]) - ]).to(device=q.device) - cu_k = torch.cat([ - block_indices.new_tensor(range(0, B * TK, TK)), block_indices.new_tensor([B * TK]) - ]).to(device=q.device) - else: - if isinstance(cu_seqlens, tuple): - cu_q, cu_k = cu_seqlens - else: - cu_q = cu_k = cu_seqlens - - for i in range(len(cu_q) - 1): - if not varlen: - q_b, k_b, v_b, i_b = q[i], k[i], v[i], block_indices[i] - else: - TQ = cu_q[i+1] - cu_q[i] - TK = cu_k[i+1] - cu_k[i] - q_b, k_b, v_b, i_b = (q[0][cu_q[i]:cu_q[i+1]], k[0][cu_k[i]:cu_k[i+1]], - v[0][cu_k[i]:cu_k[i+1]], block_indices[0][cu_q[i]:cu_q[i+1]]) - assert TQ == TK, "TQ != TK case is not supported in naive_nsa_selection" - i_b = i_b.unsqueeze(-1) * BS + i_b.new_tensor(range(BS)) - # [T, S*BS, HQ] - i_b = i_b.view(TQ, block_indices.shape[2], -1).transpose(1, 2) - for i_q in range(TQ): - # [HQ, D] - q_i = q_b[i_q] * scale - # [S*BS, HQ] - i_i = i_b[i_q] - # [S*BS, HQ, -1] - k_i, v_i = map(lambda x: x.gather(0, i_i.clamp(0, TK-1).unsqueeze(-1).expand(*i_i.shape, x.shape[-1])), - (k_b, v_b)) - # [S*BS, HQ] - attn = torch.einsum('h d, n h d -> n h', q_i, k_i).masked_fill( - torch.logical_or(i_i > i_q, i_i < 0), float('-inf')).softmax(0) - if not varlen: - o[i, i_q] = torch.einsum('n h, n h v -> h v', attn, v_i) - else: - o[0][cu_q[i] + i_q] = torch.einsum('n h, n h v -> h v', attn, v_i) - - return o.to(dtype) - - -def naive_nsa_compression( - q: torch.Tensor, - k_cmp: torch.Tensor, - v_cmp: torch.Tensor, - block_size: int, - scale: float, - cu_seqlens: torch.LongTensor | None = None, -) -> tuple[torch.Tensor, torch.Tensor]: - r""" - Naive reference for the compressed-attention branch of NSA. - - A query at position `t` attends to compressed block `c` only if the block is fully in the past, - i.e. `(c + 1) * block_size - 1 <= t`. - - Args: - q (torch.Tensor): - Queries of shape `[B, TQ, HQ, K]`. - k_cmp (torch.Tensor): - Compressed keys of shape `[B, TC, H, K]`, where `TC` is the number of compressed blocks. - v_cmp (torch.Tensor): - Compressed values of shape `[B, TC, H, V]`. - block_size (int): - Compression block size. - scale (float): - Scale factor for attention scores. - cu_seqlens (torch.LongTensor, Optional): - Cumulative sequence lengths of shape `[N+1]` for variable-length inputs. Default: `None`. - - Returns: - o (torch.Tensor): - Compressed-attention outputs of shape `[B, TQ, HQ, V]`. - lse (torch.Tensor): - Log-sum-exp of attention scores of shape `[B, TQ, HQ]`, `-inf` where no block is visible yet. - """ - dtype = q.dtype - H = k_cmp.shape[2] - G = q.shape[2] // H - q, k_cmp, v_cmp = (x.float() for x in (q, k_cmp, v_cmp)) - k_cmp, v_cmp = (repeat(x, 'b t h d -> b t (h g) d', g=G) for x in (k_cmp, v_cmp)) - - def attend(q_b, k_b, v_b): - # q_b: [TQ, HQ, K], k_b/v_b: [TC, HQ, *] - TQ, TC = q_b.shape[0], k_b.shape[0] - attn = torch.einsum('t h d, c h d -> t h c', q_b, k_b) * scale - i_t = torch.arange(TQ, device=q_b.device)[:, None] - i_c = torch.arange(TC, device=q_b.device)[None, :] - allow = ((i_c + 1) * block_size - 1) <= i_t - attn = attn.masked_fill(~allow[:, None, :], float('-inf')) - lse_b = torch.logsumexp(attn, dim=-1) - # queries with no visible block softmax to nan -> 0 - attn = torch.nan_to_num(torch.softmax(attn, dim=-1), nan=0.0) - o_b = torch.einsum('t h c, c h v -> t h v', attn, v_b) - return o_b, lse_b - - if cu_seqlens is None: - o, lse = zip(*(attend(q[i], k_cmp[i], v_cmp[i]) for i in range(q.shape[0]))) - o, lse = torch.stack(o, 0), torch.stack(lse, 0) - else: - cu_q, cu_k = cu_seqlens, prepare_chunk_offsets(cu_seqlens, block_size) - o, lse = zip(*( - attend(q[0, cu_q[i]:cu_q[i + 1]], k_cmp[0, cu_k[i]:cu_k[i + 1]], v_cmp[0, cu_k[i]:cu_k[i + 1]]) - for i in range(len(cu_q) - 1) - )) - o, lse = torch.cat(o, 0).unsqueeze(0), torch.cat(lse, 0).unsqueeze(0) - return o.to(dtype), lse - - -def naive_nsa_topk( - q: torch.Tensor, - k_cmp: torch.Tensor, - block_counts: int | torch.Tensor, - block_size: int, - scale: float, - cu_seqlens: torch.LongTensor | tuple[torch.LongTensor, torch.LongTensor] | None = None, -) -> torch.LongTensor: - r""" - Naive reference for NSA top-k block selection. - - For each query, blocks are ranked by their attention probability averaged over the query group, - and the top `block_counts` blocks are kept. The first block and the current/previous blocks are - always selected; causally-invisible or surplus slots are padded with `-1`. - - Args: - q (torch.Tensor): - Queries of shape `[B, TQ, HQ, K]`. - k_cmp (torch.Tensor): - Compressed keys of shape `[B, TC, H, K]`, where `TC` is the number of compressed blocks. - block_counts (int or torch.Tensor): - Number of blocks to select per query. Either an int, or a tensor of shape `[B, TQ, H]`. - block_size (int): - Compression block size. - scale (float): - Scale factor for attention scores. - cu_seqlens (torch.LongTensor, tuple or None, Optional): - Cumulative sequence lengths of shape `[N+1]` for variable-length inputs. - A tuple holds `(cu_seqlens_q, cu_seqlens_k)`. Default: `None`. - - Returns: - block_indices (torch.LongTensor): - Selected block indices of shape `[B, TQ, H, S]`, padded with `-1`. - """ - B, TQ, HQ, _ = q.shape - H = k_cmp.shape[2] - G = HQ // H - k_cmp = repeat(k_cmp, 'b t h d -> b t (h g) d', g=G) - - device = q.device - varlen = True - if cu_seqlens is None: - varlen = False - TQ = q.shape[1] - TC = k_cmp.shape[1] - cu_q = torch.cat([torch.arange(0, B * TQ, TQ), torch.tensor([B * TQ])]) - cu_k = torch.cat([torch.arange(0, B * TC, TC), torch.tensor([B * TC])]) - else: - assert B == 1 - if isinstance(cu_seqlens, tuple): - cu_q, cu_k = cu_seqlens - else: - cu_q = cu_k = cu_seqlens - cu_k = prepare_chunk_offsets(cu_k, block_size) - - if isinstance(block_counts, int): - S = int(block_counts) - assert S >= 0, "block_counts (int) must be >= 0" - elif torch.is_tensor(block_counts): - S = int(block_counts.max().item()) - block_indices = torch.full((B, TQ, H, S), -1, device=device, dtype=torch.long) - - for i in range(len(cu_q) - 1): - if not varlen: - q_b, k_b = q[i], k_cmp[i] - else: - TQ = (cu_q[i+1] - cu_q[i]).item() - TC = (cu_k[i+1] - cu_k[i]).item() - q_b, k_b = q[0][cu_q[i]:cu_q[i+1]], k_cmp[0][cu_k[i]:cu_k[i+1]] - - # [TQ, H, G, TC] - attn = torch.einsum('t h d, c h d -> t h c', q_b, k_b).reshape(TQ, H, G, TC) * scale - i_t = torch.arange(TQ, device=device).unsqueeze(1) - i_c = torch.arange(TC, device=device).unsqueeze(0) - # block c is causally visible once its last token (c + 1) * block_size - 1 is in the past - allow_causal = ((i_c + 1) * block_size - 1) <= i_t # [TQ, TC] - # the first block and the query's own/previous block are always selected - i_blk = i_t // block_size - forced = (i_c == i_blk) | (i_c == 0) | (i_c == i_blk - 1) # [TQ, TC] - attn = attn.masked_fill(~allow_causal[:, None, None, :], float('-inf')) - allow = allow_causal | forced - - probs = torch.nan_to_num(torch.softmax(attn, dim=-1), nan=0.0) # [TQ, H, G, TC] - scores = probs.mean(dim=2) # [TQ, H, TC] - scores = torch.where(forced[:, None, :], 1.0, scores) - - if isinstance(block_counts, int): - n_sel = torch.full((TQ, H), S, dtype=torch.long, device=device) - elif torch.is_tensor(block_counts): - if varlen: - assert block_counts.shape == (1, TQ, H) - n_sel = block_counts[0].to(device=device, dtype=torch.long) - else: - assert block_counts.shape == (B, TQ, H) - n_sel = block_counts[i].to(device=device, dtype=torch.long) - else: - raise TypeError("block_counts must be int or torch.Tensor") - - _, i_top = torch.topk(scores, k=min(S, TC), dim=-1) # [TQ, H, min(S, TC)] - - # keep a selected block only if it is allowed and within the per-query quota, else pad with -1 - top_allowed = torch.gather(allow[:, None, :].expand(TQ, H, TC).long(), dim=-1, index=i_top).bool() - i_s = torch.arange(S, device=device).view(1, 1, S) - in_quota = (i_s < n_sel.unsqueeze(-1))[:, :, :TC] - sel = torch.where(top_allowed & in_quota, i_top, torch.full_like(i_top, -1)) - - if S > TC: - sel = torch.cat((sel, torch.full((TQ, H, S - TC), -1, device=device, dtype=i_top.dtype)), dim=-1) - if varlen: - block_indices[0, cu_q[i]:cu_q[i+1]] = sel - else: - block_indices[i] = sel - return block_indices - - -def naive_nsa( - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - g_cmp: torch.Tensor | None = None, - g_slc: torch.Tensor | None = None, - g_swa: torch.Tensor | None = None, - block_indices: torch.LongTensor | None = None, - block_counts: torch.LongTensor | int = 16, - block_size: int = 64, - window_size: int = 0, - scale: float | None = None, - cu_seqlens: torch.LongTensor | tuple[torch.LongTensor, torch.LongTensor] | None = None, - return_block_indices: bool = False, -) -> torch.Tensor | tuple[torch.Tensor, torch.LongTensor]: - r""" - Naive reference for NSA, combining the compressed, selected, and sliding-window branches. - - Args: - q (torch.Tensor): - Queries of shape `[B, TQ, HQ, K]`. - k (torch.Tensor): - Keys of shape `[B, T, H, K]`. - GQA is enforced here: the ratio of query heads (HQ) to key/value heads (H) must be a power of 2 and >= 16. - This is a kernel tile dimension, which Triton requires to be a power-of-2 block shape. - v (torch.Tensor): - Values of shape `[B, T, H, V]`. - g_cmp (torch.Tensor, Optional): - Gate score for compressed attention of shape `[B, TQ, HQ]`. Default: `None`. - g_slc (torch.Tensor, Optional): - Gate score for selected attention of shape `[B, TQ, HQ]`. Default: `None`. - g_swa (torch.Tensor, Optional): - Gate score for sliding-window attention of shape `[B, TQ, HQ]`. Default: `None`. - block_indices (torch.LongTensor, Optional): - Block indices of shape `[B, TQ, H, S]`, where `S` is the number of selected blocks per query. - If provided, overrides the selection computed from compression when `g_cmp` is given. Default: `None`. - block_counts (torch.LongTensor or int, Optional): - Number of selected blocks per query. A tensor of shape `[B, TQ, H]`, or an int. Default: 16. - block_size (int, Optional): - Selected block size. Default: 64. - window_size (int, Optional): - Sliding window size. Default: 0. - scale (float, Optional): - Scale factor for attention scores. If `None`, defaults to `1 / sqrt(K)`. Default: `None`. - cu_seqlens (torch.LongTensor, tuple or None, Optional): - Cumulative sequence lengths of shape `[N+1]` for variable-length inputs, consistent with the - FlashAttention API. A tuple holds `(cu_seqlens_q, cu_seqlens_k)`. Default: `None`. - return_block_indices (bool, Optional): - Whether to also return the selected block indices. Default: `False`. - - Returns: - o (torch.Tensor): - Outputs of shape `[B, TQ, HQ, V]`. When `return_block_indices=True`, also returns the - selected block indices of shape `[B, TQ, H, S]`. - """ - assert block_counts is not None, "block counts must be provided for selection" - if scale is None: - scale = k.shape[-1] ** -0.5 - if cu_seqlens is not None: - assert q.shape[0] == 1, "batch size must be 1 when cu_seqlens are provided" - G = q.shape[2] // k.shape[2] - assert G >= 16 and (G & (G - 1)) == 0, "Group size (HQ/H) must be a power of 2 and >= 16 in NSA" - - if cu_seqlens is not None: - if isinstance(cu_seqlens, tuple): - cu_seqlens_q, cu_seqlens_k = cu_seqlens - else: - cu_seqlens_q = cu_seqlens_k = cu_seqlens - else: - cu_seqlens_q = cu_seqlens_k = None - - k_cmp, v_cmp = mean_pooling(k, block_size, cu_seqlens), mean_pooling(v, block_size, cu_seqlens) - o_cmp = None - if g_cmp is not None: - o_cmp, _ = naive_nsa_compression( - q=q, - k_cmp=k_cmp, - v_cmp=v_cmp, - block_size=block_size, - scale=scale, - cu_seqlens=cu_seqlens - ) - if block_indices is None: - block_indices = naive_nsa_topk( - q=q, - k_cmp=k_cmp, - block_counts=block_counts, - block_size=block_size, - scale=scale, - cu_seqlens=cu_seqlens - ) - else: - warnings.warn("`block_indices` is provided, overriding the selection computed from compression") - o = o_slc = naive_nsa_selection(q, k, v, block_indices, block_size, scale, cu_seqlens) - if g_slc is not None: - o = o_slc * g_slc.unsqueeze(-1) - if o_cmp is not None: - o = torch.addcmul(o, o_cmp, g_cmp.unsqueeze(-1)) - if window_size > 0: - if cu_seqlens is not None: - o_swa = flash_attn_varlen_func( - q.squeeze(0), k.squeeze(0), v.squeeze(0), - cu_seqlens_q=cu_seqlens_q, - cu_seqlens_k=cu_seqlens_k, - max_seqlen_q=q.shape[1], - max_seqlen_k=k.shape[1], - causal=True, - window_size=(window_size-1, 0) - ).unsqueeze(0) - else: - o_swa = flash_attn_func( - q, k, v, - causal=True, - window_size=(window_size-1, 0) - ) - o = torch.addcmul(o, o_swa, g_swa.unsqueeze(-1)) - if return_block_indices: - return o, block_indices - else: - return o diff --git a/build/torch-xpu/ops/nsa/parallel.py b/build/torch-xpu/ops/nsa/parallel.py deleted file mode 100644 index c58ebf1238d261f7c034ac435a7713b00bf74ffd..0000000000000000000000000000000000000000 --- a/build/torch-xpu/ops/nsa/parallel.py +++ /dev/null @@ -1,936 +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 warnings - -import torch -import triton -import triton.language as tl - -from ...ops.attn.parallel import parallel_attn_bwd_preprocess -from ...ops.nsa.compression import parallel_nsa_compression -from ...ops.nsa.utils import _bitonic_merge -from ...ops.utils import prepare_block_csr, prepare_chunk_indices, prepare_chunk_offsets, prepare_lens, prepare_token_indices -from ...ops.utils.op import exp, log -from ...ops.utils.pooling import mean_pooling -from ...utils import autocast_custom_bwd, autocast_custom_fwd, autotune_cache_kwargs, check_shared_mem, contiguous - -try: - from flash_attn import flash_attn_func, flash_attn_varlen_func -except ImportError: - warnings.warn( - "Flash Attention is not installed. Please install it via `pip install flash-attn --no-build-isolation`", - category=ImportWarning, - ) - flash_attn_func = flash_attn_varlen_func = None - - -@triton.heuristics({ - 'IS_VARLEN': lambda args: args['cu_seqlens_q'] is not None, -}) -@triton.autotune( - configs=[ - triton.Config({}, num_warps=num_warps) - for num_warps in [1, 2, 4] - ], - key=['BS', 'BK'], - **autotune_cache_kwargs, -) -@triton.jit -def parallel_nsa_kernel_topk( - q, - k, - lse, - scale, - block_indices, - cu_seqlens_q, - cu_seqlens_k, - token_indices_q, - chunk_offsets, - TQ, - TK, - H: tl.constexpr, - HQ: tl.constexpr, - G: tl.constexpr, - K: tl.constexpr, - S: tl.constexpr, - BC: tl.constexpr, - BS: tl.constexpr, - BK: tl.constexpr, - IS_VARLEN: tl.constexpr, -): - i_t, i_bh = tl.program_id(0), tl.program_id(1) - i_b, i_h = i_bh // H, i_bh % H - - if IS_VARLEN: - i_n, i_t = tl.load(token_indices_q + i_t * 2).to(tl.int32), tl.load(token_indices_q + i_t * 2 + 1).to(tl.int32) - bos_q, eos_q = tl.load(cu_seqlens_q + i_n).to(tl.int64), tl.load(cu_seqlens_q + i_n + 1).to(tl.int64) - bos_k, eos_k = tl.load(cu_seqlens_k + i_n).to(tl.int64), tl.load(cu_seqlens_k + i_n + 1).to(tl.int64) - TQ = (eos_q - bos_q).to(tl.int32) - TK = (eos_k - bos_k).to(tl.int32) - TC = tl.cdiv(TK, BS) - boc = tl.load(chunk_offsets + i_n).to(tl.int32) - else: - bos_q, eos_q = (i_b * TQ).to(tl.int64), (i_b * TQ + TQ).to(tl.int64) - TC = tl.cdiv(TK, BS) - boc = i_b * TC - # boc is the start of the current sequence at [B, TC] dimensions - p_q = tl.make_block_ptr(q + (bos_q + i_t) * HQ * K, (HQ, K), (K, 1), (i_h * G, 0), (G, BK), (1, 0)) - Q_OFFSET = TK - TQ - - # the Q block is kept in the shared memory throughout the whole kernel - # [G, BK] - b_q = tl.load(p_q, boundary_check=(0, 1)) - b_q = (b_q * scale).to(b_q.dtype) - - # number of complete compression blocks visible to the query (q tokens are the last TQ of the sequence) - NC = (i_t + Q_OFFSET + 1) // BS - ################################ - # 1. lse computation - ################################ - if lse is not None: - b_lse = tl.load(lse + (bos_q + i_t) * HQ + i_h * G + tl.arange(0, G)) - else: - # max scores for the current block - b_m = tl.full([G], float('-inf'), dtype=tl.float32) - # lse = log(acc) + m - b_acc = tl.zeros([G], dtype=tl.float32) - for i_c in range(0, NC, BC): - o_c = i_c + tl.arange(0, BC) - - p_k = tl.make_block_ptr(k + (boc * H + i_h) * K, (K, TC), (1, H*K), (0, i_c), (BK, BC), (0, 1)) - # [BK, BC] - b_k = tl.load(p_k, boundary_check=(0, 1)) - - # [G, BC] - b_s = tl.dot(b_q, b_k) - b_s = tl.where((o_c < NC)[None, :], b_s, float('-inf')) - - # [G] - b_m, b_mp = tl.maximum(b_m, tl.max(b_s, 1)), b_m - b_r = exp(b_mp - b_m) - # [G, BC] - b_p = exp(b_s - b_m[:, None]) - # [G] - b_acc = b_acc * b_r + tl.sum(b_p, 1) - - b_mp = b_m - if NC == 0: - b_lse = tl.zeros([G], dtype=tl.float32) - else: - b_lse = b_m + log(b_acc) - - ################################ - # 2. topk selection - ################################ - # [BC] - b_i = tl.full([BC], -1, dtype=tl.float32) - o_i = tl.zeros([BC], dtype=tl.int32) - m_i = tl.arange(0, BC) < BC//2 - - IC = (i_t + Q_OFFSET) // BS # Idx of the current query block - for i_c in range(0, IC + 1, BC): # +1, because the current block might be also included - o_c = i_c + tl.arange(0, BC) - p_k = tl.make_block_ptr(k + (boc * H + i_h) * K, (K, TC), (1, H*K), (0, i_c), (BK, BC), (0, 1)) - # [BK, BC] - b_k = tl.load(p_k, boundary_check=(0, 1)) - # [G, BC] - b_s = tl.dot(b_q, b_k) - b_s = tl.where(o_c < IC, b_s, float('-inf')) - # the 1st and the last 2 blocks are always selected, with normalized score set to 1.0 - b_p = tl.where((o_c == 0) | ((o_c == IC - 1) | (o_c == IC)), 1., exp(b_s - b_lse[:, None])) - # [BC] importance score = sum over the group's heads - b_i, b_ip = tl.sum(b_p, 0), b_i - # blocks with index < 0 will be skipped - o_i, o_ip = tl.where(o_c <= IC, o_c, -1), o_i - - n_dims: tl.constexpr = tl.standard._log2(b_i.shape[0]) - for i in tl.static_range(1, n_dims): - b_i, o_i = _bitonic_merge(b_i, o_i.to(tl.int32), i, 2, n_dims) - - if i_c != 0: - b_i, o_i = _bitonic_merge(b_i, o_i.to(tl.int32), n_dims, False, n_dims) - b_i_new = b_ip * m_i + b_i * (1 - m_i) - o_i_new = o_ip * m_i + o_i * (1 - m_i) - b_i, o_i = _bitonic_merge(b_i_new, o_i_new.to(tl.int32), n_dims, True, n_dims) - else: - b_i, o_i = _bitonic_merge(b_i, o_i.to(tl.int32), n_dims, True, n_dims) - - m_top = tl.arange(0, BC // S) == 0 - b_top = tl.sum(m_top[:, None] * tl.reshape(o_i, [BC // S, S]), 0) - - p_b = tl.make_block_ptr(block_indices + (bos_q + i_t) * H*S, (H*S,), (1,), (i_h * S,), (S,), (0,)) - tl.store(p_b, b_top.to(p_b.dtype.element_ty)) - - -@triton.heuristics({ - 'IS_VARLEN': lambda args: args['cu_seqlens_q'] is not None, - 'USE_BLOCK_COUNTS': lambda args: isinstance(args['block_counts'], torch.Tensor), -}) -@triton.autotune( - configs=[ - triton.Config({}, num_warps=num_warps, num_stages=num_stages) - for num_warps in [1, 2, 4, 8] - for num_stages in [1, 2, 3] - ], - key=['BS', 'BK', 'BV', 'G'], - **autotune_cache_kwargs, -) -@triton.jit -def parallel_nsa_fwd_kernel( - q, - k, - v, - o, - lse, - scale, - block_indices, - block_counts, - cu_seqlens_q, - cu_seqlens_k, - token_indices_q, - TQ, - TK, - H: tl.constexpr, - HQ: tl.constexpr, - G: tl.constexpr, - K: tl.constexpr, - V: tl.constexpr, - S: tl.constexpr, - BS: tl.constexpr, - BK: tl.constexpr, - BV: tl.constexpr, - IS_VARLEN: tl.constexpr, - USE_BLOCK_COUNTS: tl.constexpr, -): - i_t, i_v, i_bh = tl.program_id(0), tl.program_id(1), tl.program_id(2) - i_b, i_h = i_bh // H, i_bh % H - # k/v: [B, TK, H, *], q: [B, TQ, HQ, K], block_indices: [B, TQ, H, S], lse: [B, TQ, HQ]; G = HQ // H - - if IS_VARLEN: - # token_indices_q maps a flattened query to its (sequence index, in-sequence position) - i_n, i_t = tl.load(token_indices_q + i_t * 2).to(tl.int32), tl.load(token_indices_q + i_t * 2 + 1).to(tl.int32) - bos_q, eos_q = tl.load(cu_seqlens_q + i_n).to(tl.int64), tl.load(cu_seqlens_q + i_n + 1).to(tl.int64) - bos_k, eos_k = tl.load(cu_seqlens_k + i_n).to(tl.int64), tl.load(cu_seqlens_k + i_n + 1).to(tl.int64) - TQ = (eos_q - bos_q).to(tl.int32) - TK = (eos_k - bos_k).to(tl.int32) - else: - bos_q, eos_q = (i_b * TQ).to(tl.int64), (i_b * TQ + TQ).to(tl.int64) - bos_k, eos_k = (i_b * TK).to(tl.int64), (i_b * TK + TK).to(tl.int64) - - # q tokens are assumed to be the last TQ tokens of the sequence (cached decoding) - Q_OFFSET = TK - TQ - - k += (bos_k * H + i_h) * K - v += (bos_k * H + i_h) * V - block_indices += (bos_q + i_t) * H * S + i_h * S - - # block_counts: [B, TQ, H] - if USE_BLOCK_COUNTS: - NS = tl.load(block_counts + (bos_q + i_t) * H + i_h) - else: - NS = S - - p_q = tl.make_block_ptr(q + (bos_q + i_t) * HQ * K, (HQ, K), (K, 1), (i_h * G, 0), (G, BK), (1, 0)) - # the Q block is kept in shared memory throughout the kernel - # [G, BK] - b_q = tl.load(p_q, boundary_check=(0, 1)) - b_q = (b_q * scale).to(b_q.dtype) - - p_o = tl.make_block_ptr(o + (bos_q + i_t) * HQ * V, (HQ, V), (V, 1), (i_h * G, i_v * BV), (G, BV), (1, 0)) - p_lse = lse + (bos_q + i_t) * HQ + i_h * G + tl.arange(0, G) - - # [G, BV] - b_o = tl.zeros([G, BV], dtype=tl.float32) - - b_m = tl.full([G], float('-inf'), dtype=tl.float32) - b_acc = tl.zeros([G], dtype=tl.float32) - for i in range(NS): - i_s = tl.load(block_indices + i).to(tl.int32) * BS # start token index of the i-th selected KV block - if i_s <= Q_OFFSET + i_t and i_s >= 0: - p_k = tl.make_block_ptr(k, (K, TK), (1, H*K), (0, i_s), (BK, BS), (0, 1)) - p_v = tl.make_block_ptr(v, (TK, V), (H*V, 1), (i_s, i_v * BV), (BS, BV), (1, 0)) - # [BK, BS] - b_k = tl.load(p_k, boundary_check=(0, 1)) - # [BS, BV] - b_v = tl.load(p_v, boundary_check=(0, 1)) - # [G, BS] - b_s = tl.dot(b_q, b_k) - # causal mask against the absolute query position Q_OFFSET + i_t - b_s = tl.where((Q_OFFSET + i_t >= (i_s + tl.arange(0, BS)))[None, :], b_s, float('-inf')) - - # [G] - b_m, b_mp = tl.maximum(b_m, tl.max(b_s, 1)), b_m - b_r = exp(b_mp - b_m) - # [G, BS] - b_p = exp(b_s - b_m[:, None]) - # [G] - b_acc = b_acc * b_r + tl.sum(b_p, 1) - # [G, BV] - b_o = b_o * b_r[:, None] + tl.dot(b_p.to(b_q.dtype), b_v) - - b_o = b_o / b_acc[:, None] - b_m += log(b_acc) - tl.store(p_o, b_o.to(p_o.dtype.element_ty), boundary_check=(0, 1)) - tl.store(p_lse, b_m.to(p_lse.dtype.element_ty)) - - -@triton.heuristics({ - 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None, - 'USE_BLOCK_COUNTS': lambda args: isinstance(args['block_counts'], torch.Tensor), -}) -@triton.autotune( - configs=[ - triton.Config({}, num_warps=num_warps, num_stages=num_stages) - for num_warps in [1, 2, 4, 8] - for num_stages in [1, 2, 3] - ], - key=['BS', 'BK', 'BV', 'G'], - **autotune_cache_kwargs, -) -@triton.jit(do_not_specialize=['T']) -def parallel_nsa_bwd_kernel_dq( - q, - k, - v, - lse, - delta, - do, - dq, - scale, - block_indices, - block_counts, - cu_seqlens, - token_indices, - T, - B: tl.constexpr, - H: tl.constexpr, - HQ: tl.constexpr, - G: tl.constexpr, - K: tl.constexpr, - V: tl.constexpr, - S: tl.constexpr, - BS: tl.constexpr, - BK: tl.constexpr, - BV: tl.constexpr, - IS_VARLEN: tl.constexpr, - USE_BLOCK_COUNTS: tl.constexpr, -): - i_t, i_v, i_bh = tl.program_id(0), tl.program_id(1), tl.program_id(2) - i_b, i_h = i_bh // H, i_bh % H - - all = B * T - if IS_VARLEN: - i_n, i_t = tl.load(token_indices + i_t * 2).to(tl.int32), tl.load(token_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).to(tl.int32) - else: - bos, eos = (i_b * T).to(tl.int64), (i_b * T + T).to(tl.int64) - - q += (bos + i_t) * HQ*K - do += (bos + i_t) * HQ*V - lse += (bos + i_t) * HQ - delta += (bos + i_t) * HQ - dq += (i_v * all + bos + i_t) * HQ*K - block_indices += (bos + i_t) * H*S + i_h * S - - if USE_BLOCK_COUNTS: - NS = tl.load(block_counts + (bos + i_t) * H + i_h) - else: - NS = S - - k += (bos * H + i_h) * K - v += (bos * H + i_h) * V - - p_q = tl.make_block_ptr(q, (HQ, K), (K, 1), (i_h * G, 0), (G, BK), (1, 0)) - p_dq = tl.make_block_ptr(dq, (HQ, K), (K, 1), (i_h * G, 0), (G, BK), (1, 0)) - - # [G, BK] - b_q = tl.load(p_q, boundary_check=(0, 1)) - b_q = (b_q * scale).to(b_q.dtype) - - p_do = tl.make_block_ptr(do, (HQ, V), (V, 1), (i_h * G, i_v * BV), (G, BV), (1, 0)) - p_lse = lse + i_h * G + tl.arange(0, G) - p_delta = delta + i_h * G + tl.arange(0, G) - - # [G, BV] - b_do = tl.load(p_do, boundary_check=(0, 1)) - # [G] - b_lse = tl.load(p_lse) - b_delta = tl.load(p_delta) - - # [G, BK] - b_dq = tl.zeros([G, BK], dtype=tl.float32) - for i in range(NS): - i_s = tl.load(block_indices + i).to(tl.int32) * BS - if i_s <= i_t and i_s >= 0: - p_k = tl.make_block_ptr(k, (K, T), (1, H*K), (0, i_s), (BK, BS), (0, 1)) - p_v = tl.make_block_ptr(v, (V, T), (1, H*V), (i_v * BV, i_s), (BV, BS), (0, 1)) - # [BK, BS] - b_k = tl.load(p_k, boundary_check=(0, 1)) - # [BV, BS] - b_v = tl.load(p_v, boundary_check=(0, 1)) - - # [G, BS] - b_s = tl.dot(b_q, b_k) - b_p = exp(b_s - b_lse[:, None]) - b_p = tl.where((i_t >= (i_s + tl.arange(0, BS)))[None, :], b_p, 0) - - # [G, BV] @ [BV, BS] -> [G, BS] - b_dp = tl.dot(b_do, b_v) - b_ds = b_p * (b_dp.to(tl.float32) - b_delta[:, None]) - # [G, BS] @ [BS, BK] -> [G, BK] - b_dq += tl.dot(b_ds.to(b_k.dtype), tl.trans(b_k)) - b_dq *= scale - - tl.store(p_dq, b_dq.to(p_dq.dtype.element_ty), boundary_check=(0, 1)) - - -def _prune_dkv_bq(configs, nargs, **kwargs): - # the gather stacks BQ queries x G group heads into a [BQ*G]-wide tile; past ~1024 columns the - # Triton compiler fails outright (which autotune cannot catch), so cap BQ*G here. Shared-memory - # fit is left to autotune, which prunes configs that don't fit the device. Always keep BQ=1. - G = {**(nargs or {}), **kwargs}.get('G', 1) - return [c for c in configs if c.kwargs['BQ'] == 1 or c.kwargs['BQ'] * G <= 256] - - -@triton.heuristics({ - 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None, -}) -@triton.autotune( - configs=[ - triton.Config({'BQ': BQ}, num_warps=num_warps, num_stages=num_stages) - for BQ in [1, 2, 4, 8] - for num_warps in [4, 8] - for num_stages in [1, 2] - ], - key=['BS', 'BK', 'BV', 'G'], - prune_configs_by={'early_config_prune': _prune_dkv_bq}, - **autotune_cache_kwargs, -) -@triton.jit(do_not_specialize=['T']) -def parallel_nsa_bwd_kernel_dkv( - q, - k, - v, - lse, - delta, - do, - dk, - dv, - scale, - csr_indices, - csr_offsets, - cu_seqlens, - chunk_indices, - T, - B: tl.constexpr, - H: tl.constexpr, - HQ: tl.constexpr, - G: tl.constexpr, - K: tl.constexpr, - V: tl.constexpr, - TC: tl.constexpr, - BS: tl.constexpr, - BK: tl.constexpr, - BV: tl.constexpr, - BQ: tl.constexpr, - IS_VARLEN: tl.constexpr, -): - i_v, i_blk = tl.program_id(0), tl.program_id(1) - all = B * T - if IS_VARLEN: - i_c, i_h = i_blk // H, i_blk % H - i_n = tl.load(chunk_indices + i_c * 2).to(tl.int32) - i_s = tl.load(chunk_indices + i_c * 2 + 1).to(tl.int32) - bos = tl.load(cu_seqlens + i_n).to(tl.int64) - eos = tl.load(cu_seqlens + i_n + 1).to(tl.int64) - else: - i_s = i_blk % TC - i_b, i_h = (i_blk // TC) // H, (i_blk // TC) % H - bos = (i_b * T).to(tl.int64) - eos = bos + T - o_t = bos + i_s * BS + tl.arange(0, BS) - o_d = tl.arange(0, BK) - o_v = i_v * BV + tl.arange(0, BV) - o_h = tl.arange(0, BQ * G) % G - m_t, m_d, m_v = o_t < eos, o_d < K, o_v < V - - # this block's CSR slice: the queries that selected it - i_q0, i_q1 = tl.load(csr_offsets + i_blk).to(tl.int64), tl.load(csr_offsets + i_blk + 1).to(tl.int64) - NQ = i_q1 - i_q0 - - q += i_h * G*K - k += i_h * K - v += i_h * V - lse += i_h * G - delta += i_h * G - do += i_h * G*V - dk += i_h * K - dv += i_h * V - - # [BS, BK] / [BS, BV] k/v tile for this kv block - b_k = tl.load(k + o_t[:, None] * H*K + o_d[None, :], mask=m_t[:, None] & m_d[None, :], other=0.) - b_v = tl.load(v + o_t[:, None] * H*V + o_v[None, :], mask=m_t[:, None] & m_v[None, :], other=0.) - b_dk = tl.zeros([BS, BK], dtype=tl.float32) - b_dv = tl.zeros([BS, BV], dtype=tl.float32) - - for i in range(tl.cdiv(NQ, BQ)): - # BQ queries x G group heads flattened to [BQ*G] columns; row r -> query slot i*BQ + r//G, head o_h[r] - o_q = i * BQ + tl.arange(0, BQ * G) // G - m_q = o_q < NQ - o_q = tl.load(csr_indices + i_q0 + o_q, mask=m_q, other=0).to(tl.int64) - - # gather q/do/lse/delta for the [BQ*G] (query, group-head) columns - # [BQ*G, BK] - b_q = tl.load(q + o_q[:, None] * HQ*K + o_h[:, None] * K + o_d[None, :], mask=m_q[:, None] & m_d[None, :], other=0.) - b_q = (b_q * scale).to(b_k.dtype) - # [BQ*G, BV] - b_do = tl.load(do + o_q[:, None] * HQ*V + o_h[:, None] * V + o_v[None, :], mask=m_q[:, None] & m_v[None, :], other=0.) - # [BQ*G] - b_lse = tl.load(lse + o_q * HQ + o_h, mask=m_q, other=0.) - b_delta = tl.load(delta + o_q * HQ + o_h, mask=m_q, other=0.) - - # [BS, BK] @ [BK, BQ*G] -> [BS, BQ*G] - b_s = tl.dot(b_k, tl.trans(b_q)) - b_p = exp(b_s - b_lse[None, :]) - # causal: a key contributes only to queries at or after its position - b_p = tl.where((o_q[None, :] >= o_t[:, None]) & m_q[None, :], b_p, 0.) - - # [BS, BQ*G] @ [BQ*G, BV] -> [BS, BV] - b_dv += tl.dot(b_p.to(b_do.dtype), b_do) - # [BS, BV] @ [BV, BQ*G] -> [BS, BQ*G] - b_dp = tl.dot(b_v, tl.trans(b_do)) - b_ds = b_p * (b_dp.to(tl.float32) - b_delta[None, :]) - # [BS, BQ*G] @ [BQ*G, BK] -> [BS, BK] - b_dk += tl.dot(b_ds.to(b_q.dtype), b_q) - - o_dk = (i_v * all + o_t)[:, None] * H*K + o_d[None, :] - o_dv = o_t[:, None] * H*V + o_v[None, :] - tl.store(dk + o_dk, b_dk.to(dk.dtype.element_ty), mask=m_t[:, None] & m_d[None, :]) - tl.store(dv + o_dv, b_dv.to(dv.dtype.element_ty), mask=m_t[:, None] & m_v[None, :]) - - -@contiguous -def parallel_nsa_topk( - q: torch.Tensor, - k: torch.Tensor, - TK: int, - lse: torch.Tensor | None, - block_counts: torch.LongTensor | int, - block_size: int = 64, - scale: float = None, - cu_seqlens: torch.LongTensor | tuple[torch.LongTensor, torch.LongTensor] | None = None, -) -> torch.LongTensor: - B, TQ, HQ, K, H = *q.shape, k.shape[2] - - assert k.shape[0] == q.shape[0] and k.shape[-1] == q.shape[-1], "The last dimension of k and q must match" - assert lse is None or lse.shape == (B, TQ, HQ), "The shape of lse must be (B, TQ, HQ)" - - if cu_seqlens is not None: - if isinstance(cu_seqlens, tuple): - cu_seqlens_q, cu_seqlens_k = cu_seqlens - else: - cu_seqlens_q = cu_seqlens_k = cu_seqlens - token_indices_q = prepare_token_indices(cu_seqlens_q) - else: - cu_seqlens_q = cu_seqlens_k = token_indices_q = None - - G = HQ // H - # the number of selected blocks for each token - S = block_counts if isinstance(block_counts, int) else block_counts.max().item() - S = triton.next_power_of_2(S) - # here we set BC = BS, but beware that they can be chosen separately if required - BC = BS = block_size - BK = max(triton.next_power_of_2(K), 16) - assert BC >= 2 * S, f"BC ({BC}) must be greater than or equal to 2 * S ({S})" - - block_indices = torch.zeros(B, TQ, H, S, dtype=torch.int32, device=q.device) - chunk_offsets = prepare_chunk_offsets(cu_seqlens_k, BS) if cu_seqlens_k is not None else None - grid = (TQ, B * H) - # the 1st and the last 2 blocks are always selected - parallel_nsa_kernel_topk[grid]( - q=q, - k=k, - lse=lse, - scale=scale, - block_indices=block_indices, - cu_seqlens_q=cu_seqlens_q, - cu_seqlens_k=cu_seqlens_k, - token_indices_q=token_indices_q, - chunk_offsets=chunk_offsets, - TQ=TQ, - TK=TK, - H=H, - HQ=HQ, - G=G, - K=K, - S=S, - BC=BC, - BS=BS, - BK=BK, - ) - return block_indices - - -@contiguous -def parallel_nsa_fwd( - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - block_indices: torch.LongTensor, - block_counts: torch.LongTensor | int, - block_size: int, - scale: float, - cu_seqlens_q: torch.LongTensor | None = None, - cu_seqlens_k: torch.LongTensor | None = None, - token_indices_q: torch.LongTensor | None = None, -): - B, TK, H, K, V, S = *k.shape, v.shape[-1], block_indices.shape[-1] - _, TQ, HQ, _ = q.shape - G = HQ // H - BS = block_size - if check_shared_mem('hopper', q.device.index): - BK = min(256, triton.next_power_of_2(K)) - BV = min(256, triton.next_power_of_2(V)) - else: - BK = min(128, triton.next_power_of_2(K)) - BV = min(128, triton.next_power_of_2(V)) - NK = triton.cdiv(K, BK) - NV = triton.cdiv(V, BV) - assert NK == 1, "The key dimension can not be larger than 256" - - grid = (TQ, NV, B * H) - o = torch.empty(B, TQ, HQ, V, dtype=v.dtype, device=q.device) - lse = torch.empty(B, TQ, HQ, dtype=torch.float, device=q.device) - - parallel_nsa_fwd_kernel[grid]( - q=q, - k=k, - v=v, - o=o, - lse=lse, - scale=scale, - block_indices=block_indices, - block_counts=block_counts, - cu_seqlens_q=cu_seqlens_q, - cu_seqlens_k=cu_seqlens_k, - token_indices_q=token_indices_q, - TQ=TQ, - TK=TK, - H=H, - HQ=HQ, - G=G, - K=K, - V=V, - S=S, - BS=BS, - BK=BK, - BV=BV, - ) - return o, lse - - -def parallel_nsa_bwd( - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - o: torch.Tensor, - lse: torch.Tensor, - do: torch.Tensor, - block_indices: torch.Tensor, - block_counts: torch.LongTensor | int, - block_size: int = 64, - scale: float = None, - cu_seqlens: torch.LongTensor | None = None, - token_indices: torch.LongTensor | None = None, - chunk_indices: torch.LongTensor | None = None, -): - B, T, H, K, V, S = *k.shape, v.shape[-1], block_indices.shape[-1] - HQ = q.shape[2] - G = HQ // H - BS = block_size - BK = max(triton.next_power_of_2(K), 16) - BV = min(128, max(triton.next_power_of_2(v.shape[-1]), 16)) - NV = triton.cdiv(V, BV) - - delta = parallel_attn_bwd_preprocess(o, do) - - dq = torch.empty(NV, *q.shape, dtype=q.dtype if NV == 1 else torch.float, device=q.device) - grid = (T, NV, B * H) - parallel_nsa_bwd_kernel_dq[grid]( - q=q, - k=k, - v=v, - lse=lse, - delta=delta, - do=do, - dq=dq, - block_indices=block_indices, - block_counts=block_counts, - cu_seqlens=cu_seqlens, - token_indices=token_indices, - scale=scale, - T=T, - B=B, - H=H, - HQ=HQ, - G=G, - K=K, - V=V, - S=S, - BS=BS, - BK=BK, - BV=BV, - ) - dq = dq.sum(0) - - if cu_seqlens is not None and chunk_indices is None: - chunk_indices = prepare_chunk_indices(cu_seqlens, BS) - - M = triton.cdiv(T, BS) if cu_seqlens is None else triton.cdiv(prepare_lens(cu_seqlens).max().item(), BS) - dk = torch.empty(NV, *k.shape, dtype=k.dtype if NV == 1 else torch.float, device=q.device) - dv = torch.empty(v.shape, dtype=v.dtype, device=q.device) - - # invert the selection into the per-block list of selecting queries (CSR), so the kernel - # gathers and batches them into one MMA instead of scanning every query tile. - csr_indices, csr_offsets = prepare_block_csr( - block_indices=block_indices, - block_counts=block_counts, - cu_seqlens=cu_seqlens, - chunk_indices=chunk_indices, - num_blocks=M, - block_size=block_size, - ) - NB = chunk_indices.shape[0] * H if cu_seqlens is not None else B * H * M - grid = (NV, NB) - parallel_nsa_bwd_kernel_dkv[grid]( - q=q, - k=k, - v=v, - lse=lse, - delta=delta, - do=do, - dk=dk, - dv=dv, - scale=scale, - csr_indices=csr_indices, - csr_offsets=csr_offsets, - cu_seqlens=cu_seqlens, - chunk_indices=chunk_indices, - T=T, - B=B, - H=H, - HQ=HQ, - G=G, - K=K, - V=V, - TC=M, - BS=BS, - BK=BK, - BV=BV, - ) - dk = dk.sum(0) - return dq, dk, dv - - -@torch.compile -class ParallelNSAFunction(torch.autograd.Function): - - @staticmethod - @contiguous - @autocast_custom_fwd - def forward(ctx, q, k, v, block_indices, block_counts, block_size, scale, cu_seqlens): - ctx.dtype = q.dtype - - # 2-d sequence indices denoting the cu_seqlens of tokens in each sequence - # for example, if the passed `cu_seqlens` is [0, 2, 6], - # then there are 2 and 4 tokens in the 1st and 2nd sequences respectively, and `token_indices` will be - # [[0, 0], [0, 1], [1, 0], [1, 1], [1, 2], [1, 3]] - if cu_seqlens is not None: - if isinstance(cu_seqlens, tuple): - cu_seqlens_q, cu_seqlens_k = cu_seqlens - else: - cu_seqlens_q = cu_seqlens_k = cu_seqlens - token_indices_q = prepare_token_indices(cu_seqlens_q) - else: - cu_seqlens_q = cu_seqlens_k = token_indices_q = None - - o, lse = parallel_nsa_fwd( - q=q, - k=k, - v=v, - block_indices=block_indices, - block_counts=block_counts, - block_size=block_size, - scale=scale, - cu_seqlens_q=cu_seqlens_q, - cu_seqlens_k=cu_seqlens_k, - token_indices_q=token_indices_q - ) - ctx.save_for_backward(q, k, v, o, lse) - ctx.block_indices = block_indices - ctx.block_counts = block_counts - # Use cu_seqlens of q in backward, as cu_seqlens for q & k are different only for inference - ctx.cu_seqlens = cu_seqlens_q - ctx.token_indices = token_indices_q - ctx.block_size = block_size - ctx.scale = scale - # q/k cu_seqlens differ only in cached inference (TQ != TK), where backward is not supported - ctx.tq_ne_tk = isinstance(cu_seqlens, tuple) - return o.to(q.dtype) - - @staticmethod - @contiguous - @autocast_custom_bwd - def backward(ctx, do): - if ctx.tq_ne_tk: - raise NotImplementedError( - "Backward is not supported when `cu_seqlens` differs for queries and keys (cached inference). " - "Run the forward under `torch.no_grad()`." - ) - q, k, v, o, lse = ctx.saved_tensors - dq, dk, dv = parallel_nsa_bwd( - q=q, - k=k, - v=v, - o=o, - lse=lse, - do=do, - block_indices=ctx.block_indices, - block_counts=ctx.block_counts, - block_size=ctx.block_size, - scale=ctx.scale, - cu_seqlens=ctx.cu_seqlens, - token_indices=ctx.token_indices, - ) - return dq.to(q), dk.to(k), dv.to(v), None, None, None, None, None, None, None, None - - -@contiguous -def parallel_nsa( - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - g_cmp: torch.Tensor | None = None, - g_slc: torch.Tensor | None = None, - g_swa: torch.Tensor | None = None, - block_indices: torch.LongTensor | None = None, - block_counts: torch.LongTensor | int = 16, - block_size: int = 64, - window_size: int = 0, - scale: float | None = None, - cu_seqlens: torch.LongTensor | tuple[torch.LongTensor, torch.LongTensor] | None = None, -) -> torch.Tensor: - r""" - Args: - q (torch.Tensor): - queries of shape `[B, TQ, HQ, K]`. - k (torch.Tensor): - keys of shape `[B, T, H, K]`. - GQA is enforced here: the ratio of query heads (HQ) to key/value heads (H) must be a power of 2 and >= 16. - This is a kernel tile dimension, which Triton requires to be a power-of-2 block shape. - v (torch.Tensor): - values of shape `[B, T, H, V]`. - g_cmp (torch.Tensor): - Gate score for compressed attention of shape `[B, TQ, HQ]`. - g_slc (torch.Tensor): - Gate score for selected attention of shape `[B, TQ, HQ]`. - g_swa (torch.Tensor): - Gate score for sliding attentionof shape `[B, TQ, HQ]`. - block_indices (torch.LongTensor): - Block indices of shape `[B, TQ, H, S]`. - `S` is the number of selected blocks for each query token, which is set to 16 in the paper. - Will override the computed block indices from compression if provided. - block_counts (Optional[Union[torch.LongTensor, int]]): - Number of selected blocks for each query. - If a tensor is provided, with shape `[B, TQ, H]`, - each query can select the same number of blocks. - If not provided, it will default to 16. - block_size (int): - Selected block size. Default: 64. - window_size (int): - Sliding window size. Default: 0. - scale (Optional[float]): - Scale factor for attention scores. - If not provided, it will default to `1 / sqrt(K)`. Default: `None`. - cu_seqlens (torch.LongTensor, Tuple[torch.LongTensor, torch.LongTensor] or None): - Cumulative sequence lengths of shape `[N+1]` used for variable-length training, - consistent with the FlashAttention API. - When a tuple is provided, it should contain two tensors: `(cu_seqlens_q, cu_seqlens_k)`. - - Returns: - o (torch.Tensor): - Outputs of shape `[B, T, HQ, V]`. - """ - assert block_counts is not None, "block counts must be provided for selection" - if scale is None: - scale = k.shape[-1] ** -0.5 - if cu_seqlens is not None and q.shape[0] != 1: - raise ValueError( - f"The batch size is expected to be 1 rather than {q.shape[0]} when using `cu_seqlens`. " - f"Please flatten variable-length inputs before processing.", - ) - G = q.shape[2] // k.shape[2] - assert G >= 16 and (G & (G - 1)) == 0, "Group size (HQ/H) must be a power of 2 and >= 16 in NSA" - - if cu_seqlens is not None: - if isinstance(cu_seqlens, tuple): - cu_seqlens_q, cu_seqlens_k = cu_seqlens - else: - cu_seqlens_q = cu_seqlens_k = cu_seqlens - else: - cu_seqlens_q = cu_seqlens_k = None - k_cmp, v_cmp = mean_pooling(k, block_size, cu_seqlens_k), mean_pooling(v, block_size, cu_seqlens_k) - o_cmp, lse_cmp = None, None - if g_cmp is not None: - o_cmp, lse_cmp = parallel_nsa_compression( - q=q, - k=k_cmp, - v=v_cmp, - TK=k.shape[1], - block_size=block_size, - scale=scale, - cu_seqlens=cu_seqlens, - ) - if block_indices is None: - block_indices = parallel_nsa_topk( - q=q, - k=k_cmp, - lse=lse_cmp, - TK=k.shape[1], - block_counts=block_counts, - block_size=block_size, - scale=scale, - cu_seqlens=cu_seqlens, - ) - else: - warnings.warn("`block_indices` is provided, overriding the selection computed from compression") - o = o_slc = ParallelNSAFunction.apply(q, k, v, block_indices, block_counts, block_size, scale, cu_seqlens) - if g_slc is not None: - o = o_slc * g_slc.unsqueeze(-1) - if o_cmp is not None: - o = torch.addcmul(o, o_cmp, g_cmp.unsqueeze(-1)) - if window_size > 0: - if cu_seqlens is not None: - o_swa = flash_attn_varlen_func( - q.squeeze(0), k.squeeze(0), v.squeeze(0), - cu_seqlens_q=cu_seqlens_q, - cu_seqlens_k=cu_seqlens_k, - max_seqlen_q=q.shape[1], - max_seqlen_k=k.shape[1], - causal=True, - window_size=(window_size-1, 0), - ).unsqueeze(0) - else: - o_swa = flash_attn_func( - q, k, v, - causal=True, - window_size=(window_size-1, 0), - ) - o = torch.addcmul(o, o_swa, g_swa.unsqueeze(-1)) - return o diff --git a/build/torch-xpu/ops/nsa/utils.py b/build/torch-xpu/ops/nsa/utils.py deleted file mode 100644 index 36f9372371a5ae2d65b9b5a226a22ec1a5d41565..0000000000000000000000000000000000000000 --- a/build/torch-xpu/ops/nsa/utils.py +++ /dev/null @@ -1,96 +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 - -# Implements argsort based on bitonic sort. -# [What is bitonic sort?](https://en.wikipedia.org/wiki/Bitonic_sorter) - -# Code adapted from https://github.com/triton-lang/triton/issues/3698#issuecomment-2067681396 - - -import triton -import triton.language as tl - -from ...ops.utils.op import log2 - - -@triton.jit -def _compare_and_swap( - x, - ids, - flip, - i: tl.constexpr, - n_dims: tl.constexpr, -): - n_outer: tl.constexpr = x.numel >> n_dims - shape: tl.constexpr = [n_outer * 2**i, 2, 2**(n_dims - i - 1)] - y = tl.reshape(x, shape) - # slice left/right with 'stride' 2**(n_dims - i - 1) - mask = tl.arange(0, 2)[None, :, None] - left = tl.broadcast_to(tl.sum(y * (1 - mask), 1)[:, None, :], shape).to(y.dtype) - right = tl.broadcast_to(tl.sum(y * mask, 1)[:, None, :], shape).to(y.dtype) - left = tl.reshape(left, x.shape) - right = tl.reshape(right, x.shape) - # idx - y_idx = tl.reshape(ids, shape) - left_idx = tl.broadcast_to(tl.sum(y_idx * (1 - mask), 1)[:, None, :], shape) - right_idx = tl.broadcast_to(tl.sum(y_idx * mask, 1)[:, None, :], shape) - left_idx = tl.reshape(left_idx, x.shape).to(y_idx.dtype) - right_idx = tl.reshape(right_idx, x.shape).to(y_idx.dtype) - # actual compare-and-swap - idtype = tl.core.get_int_dtype(bitwidth=x.dtype.primitive_bitwidth, signed=True) - ileft = left.to(idtype, bitcast=True) - iright = right.to(idtype, bitcast=True) - ix = x.to(idtype, bitcast=True) - - cond = (left > right) != flip - ret = ix ^ tl.where(cond, ileft ^ iright, tl.zeros_like(ix)) - new_ids = ids ^ tl.where(cond, left_idx ^ right_idx, tl.zeros_like(ids)) - return ret.to(x.dtype, bitcast=True), new_ids - - -@triton.jit -def _bitonic_merge( - x, - ids, - stage: tl.constexpr, - order: tl.constexpr, - n_dims: tl.constexpr, -): - n_outer: tl.constexpr = x.numel >> n_dims - tl.static_assert(stage <= n_dims) - # flip denotes whether to re-arrange sub-sequences of elements in ascending or - # descending order. - # if flip = 00000000... then all elements will be re-arranged ascendingly at this stage - # if flip = 00110011... then all the elements will be re-arranged alternatingly (with - # a stride of 2) at this stage - if order == 2: - shape: tl.constexpr = [n_outer * 2**(n_dims - 1 - stage), 2, 2**stage] - flip = tl.reshape(tl.broadcast_to(tl.arange(0, 2)[None, :, None], shape), x.shape) - else: - flip = order - # perform `stage` rounds of `compare-and-swap` - for i in tl.static_range(stage): - x, ids = _compare_and_swap(x, ids, flip, i + (n_dims - stage), n_dims) - return x, ids - - -@triton.jit -def argsort( - x, - ids, - dim: tl.constexpr = None, - descending: tl.constexpr = tl.core.CONSTEXPR_0, -): - # handle default dimension or check that it is the most minor dim - _dim: tl.constexpr = len(x.shape) - 1 if dim is None else dim - tl.static_assert(_dim == len(x.shape) - 1, "only minor dimension is currently supported") - # iteratively run bitonic merge-sort steps - n_dims: tl.constexpr = log2(x.shape[_dim]) - - for i in tl.static_range(1, n_dims + 1): - x, ids = _bitonic_merge(x, ids, i, 2 if i < n_dims else descending, n_dims) - return x, ids diff --git a/build/torch-xpu/ops/parallax/__init__.py b/build/torch-xpu/ops/parallax/__init__.py deleted file mode 100644 index 60de44ec854bf48c0fe168b605a76bbc82359f9c..0000000000000000000000000000000000000000 --- a/build/torch-xpu/ops/parallax/__init__.py +++ /dev/null @@ -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 -# -# Parallax (parameterized local linear attention), contributed by -# Yifei Zuo et al. (https://arxiv.org/abs/2605.29157). - -from .naive import naive_parallax -from .parallel import parallel_parallax - -__all__ = [ - 'naive_parallax', - 'parallel_parallax', -] diff --git a/build/torch-xpu/ops/parallax/decode.py b/build/torch-xpu/ops/parallax/decode.py deleted file mode 100644 index 4be585faa1e95dbd754ea408343a4f62486e50a9..0000000000000000000000000000000000000000 --- a/build/torch-xpu/ops/parallax/decode.py +++ /dev/null @@ -1,337 +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 -import triton.language as tl - -from ...ops.parallax.parallel import _block_size -from ...ops.utils.op import exp2 - - -@triton.jit(do_not_specialize=['Sq', 'Skv']) -def parallax_decode_kernel( - q, - r, - k, - v, - o, - scale, - cache_start, - Sq, - Skv, - HQ: tl.constexpr, - H: tl.constexpr, - G: tl.constexpr, - K: tl.constexpr, - BK: tl.constexpr, - WINDOW_SIZE_LEFT: tl.constexpr, - USE_CACHE_START: tl.constexpr, - BT: tl.constexpr, - BS: tl.constexpr, -): - """Forward-only Parallax over a cached KV (prefill / chunked decode). - - The ``Sq`` query tokens are the last ``Sq`` positions of a length-``Skv`` - sequence, so query row ``i`` is at absolute position ``Skv - Sq + i`` and - attends to keys ``[0, Skv - Sq + i]`` (causal), further restricted to the - sliding window and to ``[cache_start, Skv)`` when set. One program owns a - ``BT``-row query block; see ``naive_parallax`` for the output formula. - """ - i_t, i_bh = tl.program_id(0), tl.program_id(1) - i_b, i_hq = i_bh // HQ, i_bh % HQ - i_h = i_hq // G # kv head shared by this q head (GQA) - RCP_LN2: tl.constexpr = 1.4426950216 - - # First valid key index (left-padding); 0 when the whole cache is valid. - if USE_CACHE_START: - kv_lo = tl.load(cache_start + i_b).to(tl.int32) - else: - kv_lo = 0 - - q_off = i_t * BT - kv_offset = Skv - Sq # absolute position of query row 0 - rows = q_off + tl.arange(0, BT) - abs_q = (kv_offset + rows)[:, None] # [BT, 1] absolute position per query row - row_mask = (rows < Sq)[:, None] # [BT, 1] real (non-padded) query rows - - # Restrict the key-block loop to what this query block can reach: up to the - # causal diagonal (KV_END_BLOCK) and down to the window's left edge (else 0). - max_abs = kv_offset + tl.minimum(Sq, q_off + BT) - 1 - KV_END_BLOCK = tl.cdiv(tl.minimum(Skv, max_abs + 1), BS) - if WINDOW_SIZE_LEFT >= 0: - leftmost = tl.maximum(kv_lo, kv_offset + q_off - WINDOW_SIZE_LEFT + 1) - else: - leftmost = kv_lo - KV_START_BLOCK = leftmost // BS - - p_q = tl.make_block_ptr(q + (i_b * Sq * HQ + i_hq) * K, (Sq, K), (HQ * K, 1), (q_off, 0), (BT, BK), (1, 0)) - p_r = tl.make_block_ptr(r + (i_b * Sq * HQ + i_hq) * K, (Sq, K), (HQ * K, 1), (q_off, 0), (BT, BK), (1, 0)) - p_k = tl.make_block_ptr(k + (i_b * Skv * H + i_h) * K, (Skv, K), (H * K, 1), (KV_START_BLOCK * BS, 0), (BS, BK), (1, 0)) - p_v = tl.make_block_ptr(v + (i_b * Skv * H + i_h) * K, (Skv, K), (H * K, 1), (KV_START_BLOCK * BS, 0), (BS, BK), (1, 0)) - p_o = tl.make_block_ptr(o + (i_b * Sq * HQ + i_hq) * K, (Sq, K), (HQ * K, 1), (q_off, 0), (BT, BK), (1, 0)) - - b_q = tl.load(p_q, boundary_check=(0, 1), padding_option="zero") - b_r = tl.load(p_r, boundary_check=(0, 1), padding_option="zero") - m_acc = tl.zeros((BT, 1), dtype=tl.float32) - float("inf") - d1_acc = tl.zeros((BT, 1), dtype=tl.float32) - d2_acc = tl.zeros((BT, 1), dtype=tl.float32) - barv_acc = tl.zeros((BT, BK), dtype=tl.float32) - Rv_acc = tl.zeros((BT, BK), dtype=tl.float32) - scale_log2 = scale * RCP_LN2 - - for col_block_id in range(KV_START_BLOCK, KV_END_BLOCK): - col = (col_block_id * BS + tl.arange(0, BS))[None, :] # [1, BS] key positions - b_k = tl.load(p_k, boundary_check=(0, 1), padding_option="zero") # [BS, BK] - b_v = tl.load(p_v, boundary_check=(0, 1), padding_option="zero") # [BS, BK] - # [BT, BS]: causal, in-cache, past left-padding; optionally inside the window. - mask = (abs_q >= col) & row_mask & (col < Skv) & (col >= kv_lo) - if WINDOW_SIZE_LEFT >= 0: - mask = mask & (col >= abs_q - WINDOW_SIZE_LEFT + 1) - qk = tl.dot(b_q, tl.trans(b_k), out_dtype=tl.float32) * scale_log2 # [BT, BS], base-2 logits - qk = tl.where(mask, qk, -float("inf")) - m_new = tl.maximum(m_acc, tl.max(qk, axis=1, keep_dims=True)) - # finite pivot so a row with no valid key yet doesn't hit exp2(-inf - -inf) = NaN - safe_m = tl.where(m_new == -float("inf"), 0.0, m_new) - alpha = exp2(m_acc - safe_m) # online-softmax rescale of running state - w = exp2(qk - safe_m) # [BT, BS] = p1 (unnormalized softmax) - rk = tl.dot(b_r, tl.trans(b_k), out_dtype=tl.float32) # [BT, BS] = r @ k^T (unscaled) - wr = w * rk # [BT, BS] = p2 (unnormalized) - d1_acc = alpha * d1_acc + tl.sum(w, axis=1, keep_dims=True) # running sum(p1) - d2_acc = alpha * d2_acc + tl.sum(wr, axis=1, keep_dims=True) # running sum(p2) - barv_acc = alpha * barv_acc # running p1 @ v - Rv_acc = alpha * Rv_acc # running p2 @ v - barv_acc = tl.dot(w.to(b_v.dtype), b_v, out_dtype=tl.float32, acc=barv_acc) - Rv_acc = tl.dot(wr.to(b_v.dtype), b_v, out_dtype=tl.float32, acc=Rv_acc) - m_acc = m_new - p_k = tl.advance(p_k, (BS, 0)) - p_v = tl.advance(p_v, (BS, 0)) - - # Rows that see no valid key (e.g. left-padded query positions) have d1 == 0; - # emit a finite zero instead of inf/NaN so padding can't poison valid rows. - inv_d1 = tl.where(row_mask & (d1_acc > 0.0), 1.0 / d1_acc, 0.0) - b_barv = barv_acc * inv_d1 # O1 / d1 - b_bart = d2_acc * inv_d1 # d2 / d1 - b_o = b_barv + b_bart * b_barv - Rv_acc * inv_d1 # O1/d1 * (1 + d2/d1) - O2/d1 - - tl.store(p_o, b_o.to(p_o.dtype.element_ty), boundary_check=(0, 1)) - - -def parallax_decode( - q: torch.Tensor, - r: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - scale: float | None = None, - window_size: int | None = None, - cache_start: torch.LongTensor | None = None, -) -> torch.Tensor: - r""" - Forward-only Parallax decode/prefill over a cached KV (inference; no autograd). - - The ``Sq`` query tokens are treated as the *last* ``Sq`` positions of a - length-``Skv`` sequence: query ``i`` sits at absolute position ``Skv - Sq + i`` - and attends causally to keys ``[0, Skv - Sq + i]`` (and, with `window_size`, - only the most recent `window_size` of them). `Sq == Skv` reduces to a full - causal prefill; `Sq == 1` is a single decode step. - - Args: - q (torch.Tensor): - new queries of shape `[B, Sq, HQ, D]`. - r (torch.Tensor): - new secondary queries of shape `[B, Sq, HQ, D]`. - k (torch.Tensor): - cached keys of shape `[B, Skv, H, D]` (`Skv >= Sq`). - v (torch.Tensor): - cached values of shape `[B, Skv, H, D]`. - scale (float, Optional): - Scale applied to `q @ k^T` only. Defaults to `1 / sqrt(D)`. Default: `None`. - window_size (int, Optional): - Sliding-window length. Default: `None`. - cache_start (torch.LongTensor, Optional): - Per-batch first valid key index of shape `[B]` (to mask left-padding - in the cache). `None` means the whole cache is valid. Default: `None`. - - Returns: - o (torch.Tensor): - output of shape `[B, Sq, HQ, D]`. - """ - B, Sq, HQ, K = q.shape - Skv, H = k.shape[1], k.shape[2] - G = HQ // H - if scale is None: - scale = K ** -0.5 - window_size_left = -1 if window_size is None else window_size - - q, r, k, v = (x.contiguous() for x in (q, r, k, v)) - if cache_start is not None: - cache_start = cache_start.to(device=q.device, dtype=torch.int32).contiguous() - BK = triton.next_power_of_2(K) - BT = _block_size(K, q.device.index) - o = torch.empty_like(q) - grid = (triton.cdiv(Sq, BT), B * HQ) - parallax_decode_kernel[grid]( - q, r, k, v, o, float(scale), cache_start, Sq, Skv, - HQ=HQ, H=H, G=G, K=K, BK=BK, - WINDOW_SIZE_LEFT=window_size_left, - USE_CACHE_START=cache_start is not None, - BT=BT, BS=BT, - num_warps=8, num_stages=2, - ) - return o - - -@triton.heuristics({ - 'USE_CACHE_START': lambda args: args['cache_start'] is not None, -}) -@triton.jit(do_not_specialize=['Skv']) -def parallax_decode_one_step_kernel( - q, - r, - k, - v, - o, - scale, - cache_start, - Skv, - HQ: tl.constexpr, - H: tl.constexpr, - G: tl.constexpr, - K: tl.constexpr, - BK: tl.constexpr, - WINDOW_SIZE_LEFT: tl.constexpr, - USE_CACHE_START: tl.constexpr, - BS: tl.constexpr, -): - """Single-token Parallax decode: one query per (batch, head) over its cached KV. - - The query is the sequence's last position, so causality is implicit (it sees - every key) and only the lower bound from the window / left-padding matters. - The query is held as a *vector* and reduced against the cache with an online - softmax, so there is none of the wasted-tile compute the prefill-shaped - ``parallax_decode_kernel`` incurs at ``Sq == 1``. One program per - (batch, head); see ``naive_parallax`` for the output formula. - """ - i_bh = tl.program_id(0) - i_b, i_hq = i_bh // HQ, i_bh % HQ - i_h = i_hq // G # kv head shared by this q head (GQA) - RCP_LN2: tl.constexpr = 1.4426950216 - - # Lowest key the query attends to: left-padding start, raised to the window edge. - kv_lo = 0 - if USE_CACHE_START: - kv_lo = tl.load(cache_start + i_b).to(tl.int32) - if WINDOW_SIZE_LEFT >= 0: - kv_lo = tl.maximum(kv_lo, Skv - WINDOW_SIZE_LEFT) - kv_lo = tl.maximum(kv_lo, 0) - - p_q = tl.make_block_ptr(q + i_bh * K, (K,), (1,), (0,), (BK,), (0,)) - p_r = tl.make_block_ptr(r + i_bh * K, (K,), (1,), (0,), (BK,), (0,)) - p_o = tl.make_block_ptr(o + i_bh * K, (K,), (1,), (0,), (BK,), (0,)) - b_q = tl.load(p_q, boundary_check=(0,), padding_option="zero").to(tl.float32) # [BK] query vector - b_r = tl.load(p_r, boundary_check=(0,), padding_option="zero").to(tl.float32) # [BK] secondary query - scale_log2 = scale * RCP_LN2 - - # Running online-softmax state for the single query: pivot m, denominators - # d1/d2 = sum(p1)/sum(p2), and unnormalized outputs o1/o2 = p1@v / p2@v. - m = tl.full((1,), -float("inf"), dtype=tl.float32) - d1 = tl.zeros((1,), dtype=tl.float32) - d2 = tl.zeros((1,), dtype=tl.float32) - o1 = tl.zeros((BK,), dtype=tl.float32) - o2 = tl.zeros((BK,), dtype=tl.float32) - - start_block = kv_lo // BS - p_k = tl.make_block_ptr(k + (i_b * Skv * H + i_h) * K, (Skv, K), (H * K, 1), (start_block * BS, 0), (BS, BK), (1, 0)) - p_v = tl.make_block_ptr(v + (i_b * Skv * H + i_h) * K, (Skv, K), (H * K, 1), (start_block * BS, 0), (BS, BK), (1, 0)) - for i_s in range(start_block * BS, tl.cdiv(Skv, BS) * BS, BS): - col = i_s + tl.arange(0, BS) - mask = (col >= kv_lo) & (col < Skv) # [BS] valid keys - b_k = tl.load(p_k, boundary_check=(0, 1), padding_option="zero") # [BS, BK] - b_v = tl.load(p_v, boundary_check=(0, 1), padding_option="zero") # [BS, BK] - s1 = tl.sum(b_q[None, :] * b_k, axis=1) * scale_log2 # [BS] = scale * (q . k), base-2 - s2 = tl.sum(b_r[None, :] * b_k, axis=1) # [BS] = r . k (unscaled) - s1 = tl.where(mask, s1, -float("inf")) - m_new = tl.maximum(m, tl.max(s1)) - m_safe = tl.where(m_new == -float("inf"), 0.0, m_new) # finite pivot (empty cache -> 0) - alpha = exp2(m - m_safe) # rescale running state - p1 = exp2(s1 - m_safe) # [BS] - p2 = p1 * s2 # [BS] - d1 = d1 * alpha + tl.sum(p1) - d2 = d2 * alpha + tl.sum(p2) - o1 = o1 * alpha + tl.sum(p1[:, None] * b_v, axis=0) # [BK] - o2 = o2 * alpha + tl.sum(p2[:, None] * b_v, axis=0) - m = m_new - p_k = tl.advance(p_k, (BS, 0)) - p_v = tl.advance(p_v, (BS, 0)) - - inv_d1 = tl.where(d1 > 0.0, 1.0 / d1, 0.0) # 0 when no valid key (avoid NaN) - out = o1 * inv_d1 * (1.0 + d2 * inv_d1) - o2 * inv_d1 # [BK] O1/d1*(1 + d2/d1) - O2/d1 - tl.store(p_o, out.to(p_o.dtype.element_ty), boundary_check=(0,)) - - -def parallax_decode_one_step( - q: torch.Tensor, - r: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - scale: float | None = None, - window_size: int | None = None, - cache_start: torch.LongTensor | None = None, -) -> torch.Tensor: - r""" - Single-token Parallax decode (one query per sequence over the cached KV). - - Optimized for the `Sq == 1` autoregressive step: the query is loaded as a - vector (not a `[BT, D]` tile) and reduced against the cache with an online - softmax, avoiding the wasted-tile compute of the prefill-shaped - :func:`parallax_decode`. Forward-only (inference). - - Args: - q (torch.Tensor): - new query of shape `[B, 1, HQ, D]`. - r (torch.Tensor): - new secondary query of shape `[B, 1, HQ, D]`. - k (torch.Tensor): - cached keys of shape `[B, Skv, H, D]`. - v (torch.Tensor): - cached values of shape `[B, Skv, H, D]`. - scale (float, Optional): - Scale applied to `q @ k^T` only. Defaults to `1 / sqrt(D)`. Default: `None`. - window_size (int, Optional): - Sliding-window length; the query attends to the most recent `window_size` keys. - Default: `None`. - cache_start (torch.LongTensor, Optional): - Per-batch first valid key index of shape `[B]` (left-padding). Default: `None`. - - Returns: - o (torch.Tensor): - output of shape `[B, 1, HQ, D]`. - """ - B, Sq, HQ, K = q.shape - if Sq != 1: - raise ValueError(f"parallax_decode_one_step expects a single query (Sq=1), got Sq={Sq}") - Skv, H = k.shape[1], k.shape[2] - G = HQ // H - if scale is None: - scale = K ** -0.5 - window_size_left = -1 if window_size is None else window_size - - q, r, k, v = (x.contiguous() for x in (q, r, k, v)) - if cache_start is not None: - cache_start = cache_start.to(device=q.device, dtype=torch.int32).contiguous() - BK = triton.next_power_of_2(K) - o = torch.empty_like(q) - grid = (B * HQ,) - parallax_decode_one_step_kernel[grid]( - q, r, k, v, o, float(scale), cache_start, Skv, - HQ=HQ, H=H, G=G, K=K, BK=BK, - WINDOW_SIZE_LEFT=window_size_left, - USE_CACHE_START=cache_start is not None, - BS=128, - num_warps=4, num_stages=2, - ) - return o diff --git a/build/torch-xpu/ops/parallax/naive.py b/build/torch-xpu/ops/parallax/naive.py deleted file mode 100644 index 174d859ae2f758712572c037ee347c07fc860956..0000000000000000000000000000000000000000 --- a/build/torch-xpu/ops/parallax/naive.py +++ /dev/null @@ -1,96 +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 - - -def naive_parallax( - q: torch.Tensor, - r: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - scale: float | None = None, - window_size: int | None = None, - causal: bool = True, -) -> torch.Tensor: - r""" - Reference PyTorch implementation of Parallax (parameterized local - linear attention, Algorithm 1 of https://arxiv.org/abs/2605.29157). - - With ``s1 = scale * (q @ k^T)`` and ``s2 = r @ k^T`` (note: ``s2`` is NOT - scaled), ``p1 = softmax(s1)``, ``d1 = sum(p1)``, ``p2 = p1 * s2``, - ``d2 = sum(p2)``, ``O1 = p1 @ v``, ``O2 = p2 @ v``, the output is - - out = O1 / d1 * (1 + d2 / d1) - O2 / d1 - - so ``r`` injects a first-order multiplicative correction onto the - softmax-weighted values. Computation is performed in fp32 and returned in - the input dtype. - - Args: - q (torch.Tensor): - queries of shape `[B, T, HQ, D]`. - r (torch.Tensor): - secondary queries of shape `[B, T, HQ, D]` (same shape as `q`). - k (torch.Tensor): - keys of shape `[B, T, H, D]`. GQA is applied when `HQ` is divisible by `H`. - v (torch.Tensor): - values of shape `[B, T, H, D]`. - scale (float, Optional): - Scale applied to the `q @ k^T` logits only. If `None`, defaults to `1 / sqrt(D)`. - Default: `None`. - window_size (int, Optional): - Sliding-window length. If provided, each query at position `i` only attends to - keys in `[i - window_size + 1, i]`. If `None`, full causal attention is used. - Default: `None`. - causal (bool, Optional): - Whether to apply the causal mask. Default: `True`. - - Returns: - o (torch.Tensor): - output of shape `[B, T, HQ, D]`. - """ - B, T, HQ, D = q.shape - H = k.shape[2] - G = HQ // H - - if scale is None: - scale = D ** -0.5 - - dtype = q.dtype - q = q.float().reshape(B, T, H, G, D) - r = r.float().reshape(B, T, H, G, D) - k = k.float() - v = v.float() - - # [B, H, G, T, T] - s1 = torch.einsum('bqhgd,bkhd->bhgqk', q, k) * scale - s2 = torch.einsum('bqhgd,bkhd->bhgqk', r, k) - - if causal: - row_idx = torch.arange(T, device=q.device)[:, None] - col_idx = torch.arange(T, device=q.device)[None, :] - mask = col_idx > row_idx - if window_size is not None: - mask = mask | (row_idx - col_idx >= window_size) - s1 = s1.masked_fill(mask[None, None, None], float('-inf')) - - m = s1.amax(dim=-1, keepdim=True) - # Fully-masked rows cannot occur with a causal mask (the diagonal is always - # in-window), but guard the `-inf` pivot defensively. - m = torch.where(torch.isneginf(m), torch.zeros_like(m), m) - p1 = (s1 - m).exp() # [B, H, G, Tq, Tk] - d1 = p1.sum(dim=-1) # [B, H, G, Tq] - p2 = p1 * s2 - d2 = p2.sum(dim=-1) # [B, H, G, Tq] - o1 = torch.einsum('bhgqk,bkhd->bqhgd', p1, v) # [B, Tq, H, G, D] - o2 = torch.einsum('bhgqk,bkhd->bqhgd', p2, v) - - c_norm = (d2 / d1).permute(0, 3, 1, 2) # [B, Tq, H, G] - inv_d1 = (1.0 / d1).permute(0, 3, 1, 2) # [B, Tq, H, G] - out = o1 * inv_d1[..., None] * (1.0 + c_norm[..., None]) - o2 * inv_d1[..., None] - return out.reshape(B, T, HQ, D).to(dtype) diff --git a/build/torch-xpu/ops/parallax/parallel.py b/build/torch-xpu/ops/parallax/parallel.py deleted file mode 100644 index 40c3c425b5cb0028596126543a6408ab52d485dc..0000000000000000000000000000000000000000 --- a/build/torch-xpu/ops/parallax/parallel.py +++ /dev/null @@ -1,807 +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 -import triton.language as tl -from einops import reduce - -from ...ops.utils import prepare_chunk_indices -from ...ops.utils.op import exp2 -from ...utils import IS_NVIDIA_BLACKWELL, autocast_custom_bwd, autocast_custom_fwd, check_shared_mem, contiguous - - -def _block_size(head_dim: int, device_index: int) -> int: - # A single square tile size shared by all kernels so one `chunk_indices` - # (built host-side for varlen) matches every grid. Kept modest to bound the - # fp32 accumulator footprint Parallax carries (barv/Rv/grad accumulators). - if check_shared_mem('hopper', device_index) and not IS_NVIDIA_BLACKWELL and head_dim <= 64: - return 128 - return 64 - - -@triton.heuristics({ - 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None, -}) -@triton.jit(do_not_specialize=['T']) -def parallel_parallax_fwd_kernel( - q, - r, - k, - v, - o, - barv, - d1, - bart, - m, - scale, - cu_seqlens, - chunk_indices, - T, - HQ: tl.constexpr, - H: tl.constexpr, - G: tl.constexpr, - K: tl.constexpr, - BK: tl.constexpr, - WINDOW_SIZE_LEFT: tl.constexpr, - BT: tl.constexpr, - BS: tl.constexpr, - IS_VARLEN: tl.constexpr, -): - i_t, i_bh = tl.program_id(0), tl.program_id(1) - i_b, i_hq = i_bh // HQ, i_bh % HQ - i_h = i_hq // G - - 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 - else: - bos = (i_b * T).to(tl.int64) - RCP_LN2: tl.constexpr = 1.4426950216 - - row_offset = i_t * BT - row_indices = row_offset + tl.arange(0, BT) - row_mask = row_indices[:, None] < T - NUM_TOTAL_BLOCKS = tl.cdiv(tl.minimum(T, row_offset + BT), BS) - NUM_SAFE_BLOCKS = tl.minimum(row_offset, T) // BS - - # SWA col-block boundaries. WINDOW_SIZE_LEFT < 0 disables SWA. - if WINDOW_SIZE_LEFT >= 0: - leftmost_valid = tl.maximum(0, row_offset - WINDOW_SIZE_LEFT + 1) - FIRST_COL_BLOCK = leftmost_valid // BS - # Phase A is unmasked, so the safe zone must clear the window's left edge for - # the tile's LAST row (row_offset + BT - 1), not its first. - safe_left_valid = tl.maximum(0, row_offset + BT - WINDOW_SIZE_LEFT) - SAFE_LEFT_START = (safe_left_valid + BS - 1) // BS - else: - FIRST_COL_BLOCK = 0 - SAFE_LEFT_START = 0 - LEFT_BORDER_END = tl.minimum(SAFE_LEFT_START, NUM_SAFE_BLOCKS) - SAFE_MIDDLE_START = tl.maximum(FIRST_COL_BLOCK, SAFE_LEFT_START) - RIGHT_BORDER_START = tl.maximum(FIRST_COL_BLOCK, NUM_SAFE_BLOCKS) - - p_q = tl.make_block_ptr(q + (bos * HQ + i_hq) * K, (T, K), (HQ * K, 1), (row_offset, 0), (BT, BK), (1, 0)) - p_r = tl.make_block_ptr(r + (bos * HQ + i_hq) * K, (T, K), (HQ * K, 1), (row_offset, 0), (BT, BK), (1, 0)) - p_k = tl.make_block_ptr(k + (bos * H + i_h) * K, (T, K), (H * K, 1), (FIRST_COL_BLOCK * BS, 0), (BS, BK), (1, 0)) - p_v = tl.make_block_ptr(v + (bos * H + i_h) * K, (T, K), (H * K, 1), (FIRST_COL_BLOCK * BS, 0), (BS, BK), (1, 0)) - p_o = tl.make_block_ptr(o + (bos * HQ + i_hq) * K, (T, K), (HQ * K, 1), (row_offset, 0), (BT, BK), (1, 0)) - p_barv = tl.make_block_ptr(barv + (bos * HQ + i_hq) * K, (T, K), (HQ * K, 1), (row_offset, 0), (BT, BK), (1, 0)) - p_d1 = tl.make_block_ptr(d1 + bos * HQ + i_hq, (T, 1), (HQ, 1), (row_offset, 0), (BT, 1), (1, 0)) - p_bart = tl.make_block_ptr(bart + bos * HQ + i_hq, (T, 1), (HQ, 1), (row_offset, 0), (BT, 1), (1, 0)) - p_m = tl.make_block_ptr(m + bos * HQ + i_hq, (T, 1), (HQ, 1), (row_offset, 0), (BT, 1), (1, 0)) - - b_q = tl.load(p_q, boundary_check=(0, 1), padding_option="zero") - b_r = tl.load(p_r, boundary_check=(0, 1), padding_option="zero") - m_acc = tl.zeros((BT, 1), dtype=tl.float32) - float("inf") - d1_acc = tl.zeros((BT, 1), dtype=tl.float32) - d2_acc = tl.zeros((BT, 1), dtype=tl.float32) - barv_acc = tl.zeros((BT, BK), dtype=tl.float32) - Rv_acc = tl.zeros((BT, BK), dtype=tl.float32) - scale_log2 = scale * RCP_LN2 - - # Phase 0: left-border blocks (SWA only). Window mask only. - for col_block_id in range(FIRST_COL_BLOCK, LEFT_BORDER_END): - col_indices = col_block_id * BS + tl.arange(0, BS) - b_k = tl.load(p_k, boundary_check=(0, 1), padding_option="zero") - b_v = tl.load(p_v, boundary_check=(0, 1), padding_option="zero") - mask = ( - (col_indices[None, :] >= row_indices[:, None] - WINDOW_SIZE_LEFT + 1) - & row_mask - & (col_indices[None, :] < T) - ) - qk = tl.dot(b_q, tl.trans(b_k), out_dtype=tl.float32) * scale_log2 - qk = tl.where(mask, qk, -float("inf")) - m_new = tl.maximum(m_acc, tl.max(qk, axis=1, keep_dims=True)) - safe_m = tl.where(m_new == -float("inf"), 0.0, m_new) - alpha = exp2(m_acc - safe_m) - w = exp2(qk - safe_m) - rk = tl.dot(b_r, tl.trans(b_k), out_dtype=tl.float32) - wr = w * rk - d1_acc = alpha * d1_acc + tl.sum(w, axis=1, keep_dims=True) - d2_acc = alpha * d2_acc + tl.sum(wr, axis=1, keep_dims=True) - barv_acc = alpha * barv_acc - Rv_acc = alpha * Rv_acc - barv_acc = tl.dot(w.to(b_v.dtype), b_v, out_dtype=tl.float32, acc=barv_acc) - Rv_acc = tl.dot(wr.to(b_v.dtype), b_v, out_dtype=tl.float32, acc=Rv_acc) - m_acc = m_new - p_k = tl.advance(p_k, (BS, 0)) - p_v = tl.advance(p_v, (BS, 0)) - - # Phase A: safe blocks (no mask). - for _safe in range(SAFE_MIDDLE_START, NUM_SAFE_BLOCKS): - b_k = tl.load(p_k, boundary_check=(0, 1), padding_option="zero") - b_v = tl.load(p_v, boundary_check=(0, 1), padding_option="zero") - qk = tl.dot(b_q, tl.trans(b_k), out_dtype=tl.float32) * scale_log2 - m_new = tl.maximum(m_acc, tl.max(qk, axis=1, keep_dims=True)) - safe_m = tl.where(m_new == -float("inf"), 0.0, m_new) - alpha = exp2(m_acc - safe_m) - w = exp2(qk - safe_m) - rk = tl.dot(b_r, tl.trans(b_k), out_dtype=tl.float32) - wr = w * rk - d1_acc = alpha * d1_acc + tl.sum(w, axis=1, keep_dims=True) - d2_acc = alpha * d2_acc + tl.sum(wr, axis=1, keep_dims=True) - barv_acc = alpha * barv_acc - Rv_acc = alpha * Rv_acc - barv_acc = tl.dot(w.to(b_v.dtype), b_v, out_dtype=tl.float32, acc=barv_acc) - Rv_acc = tl.dot(wr.to(b_v.dtype), b_v, out_dtype=tl.float32, acc=Rv_acc) - m_acc = m_new - p_k = tl.advance(p_k, (BS, 0)) - p_v = tl.advance(p_v, (BS, 0)) - - # Phase B: right-border blocks (causal + boundary + window mask). - for col_block_id in range(RIGHT_BORDER_START, NUM_TOTAL_BLOCKS): - col_indices = col_block_id * BS + tl.arange(0, BS) - b_k = tl.load(p_k, boundary_check=(0, 1), padding_option="zero") - b_v = tl.load(p_v, boundary_check=(0, 1), padding_option="zero") - if WINDOW_SIZE_LEFT >= 0: - mask = ( - (row_indices[:, None] >= col_indices[None, :]) - & (col_indices[None, :] >= row_indices[:, None] - WINDOW_SIZE_LEFT + 1) - & row_mask - & (col_indices[None, :] < T) - ) - else: - mask = ( - (row_indices[:, None] >= col_indices[None, :]) - & row_mask - & (col_indices[None, :] < T) - ) - qk = tl.dot(b_q, tl.trans(b_k), out_dtype=tl.float32) * scale_log2 - qk = tl.where(mask, qk, -float("inf")) - m_new = tl.maximum(m_acc, tl.max(qk, axis=1, keep_dims=True)) - safe_m = tl.where(m_new == -float("inf"), 0.0, m_new) - alpha = exp2(m_acc - safe_m) - w = exp2(qk - safe_m) - rk = tl.dot(b_r, tl.trans(b_k), out_dtype=tl.float32) - wr = w * rk - d1_acc = alpha * d1_acc + tl.sum(w, axis=1, keep_dims=True) - d2_acc = alpha * d2_acc + tl.sum(wr, axis=1, keep_dims=True) - barv_acc = alpha * barv_acc - Rv_acc = alpha * Rv_acc - barv_acc = tl.dot(w.to(b_v.dtype), b_v, out_dtype=tl.float32, acc=barv_acc) - Rv_acc = tl.dot(wr.to(b_v.dtype), b_v, out_dtype=tl.float32, acc=Rv_acc) - m_acc = m_new - p_k = tl.advance(p_k, (BS, 0)) - p_v = tl.advance(p_v, (BS, 0)) - - inv_d1 = tl.where(row_mask, 1.0 / d1_acc, 0.0) - b_barv = barv_acc * inv_d1 - b_bart = d2_acc * inv_d1 - b_o = b_barv + b_bart * b_barv - Rv_acc * inv_d1 - - tl.store(p_o, b_o.to(p_o.dtype.element_ty), boundary_check=(0, 1)) - tl.store(p_barv, b_barv.to(p_barv.dtype.element_ty), boundary_check=(0, 1)) - tl.store(p_d1, d1_acc, boundary_check=(0, 1)) - tl.store(p_bart, b_bart, boundary_check=(0, 1)) - tl.store(p_m, m_acc, boundary_check=(0, 1)) - - -@triton.heuristics({ - 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None, -}) -@triton.jit(do_not_specialize=['T']) -def parallel_parallax_bwd_kernel_preprocess( - grad_o, - o, - barv, - delta_t, - delta_b, - cu_seqlens, - chunk_indices, - T, - HQ: tl.constexpr, - K: tl.constexpr, - BK: tl.constexpr, - BT: tl.constexpr, - IS_VARLEN: tl.constexpr, -): - i_t, i_bh = tl.program_id(0), tl.program_id(1) - i_b, i_hq = i_bh // HQ, i_bh % HQ - - 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 - else: - bos = (i_b * T).to(tl.int64) - - row_offset = i_t * BT - p_grad_o = tl.make_block_ptr(grad_o + (bos * HQ + i_hq) * K, (T, K), (HQ * K, 1), (row_offset, 0), (BT, BK), (1, 0)) - p_o = tl.make_block_ptr(o + (bos * HQ + i_hq) * K, (T, K), (HQ * K, 1), (row_offset, 0), (BT, BK), (1, 0)) - p_barv = tl.make_block_ptr(barv + (bos * HQ + i_hq) * K, (T, K), (HQ * K, 1), (row_offset, 0), (BT, BK), (1, 0)) - p_t = tl.make_block_ptr(delta_t + bos * HQ + i_hq, (T, 1), (HQ, 1), (row_offset, 0), (BT, 1), (1, 0)) - p_b = tl.make_block_ptr(delta_b + bos * HQ + i_hq, (T, 1), (HQ, 1), (row_offset, 0), (BT, 1), (1, 0)) - - b_grad_o = tl.load(p_grad_o, boundary_check=(0, 1), padding_option="zero").to(tl.float32) - b_o = tl.load(p_o, boundary_check=(0, 1), padding_option="zero").to(tl.float32) - b_barv = tl.load(p_barv, boundary_check=(0, 1), padding_option="zero").to(tl.float32) - - b_t = tl.sum(b_grad_o * b_o, axis=1, keep_dims=True) - b_b = tl.sum(b_grad_o * b_barv, axis=1, keep_dims=True) - - tl.store(p_t, b_t, boundary_check=(0, 1)) - tl.store(p_b, b_b, boundary_check=(0, 1)) - - -@triton.heuristics({ - 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None, -}) -@triton.jit(do_not_specialize=['T']) -def parallel_parallax_bwd_kernel_dqr( - q, - r, - k, - v, - d1, - bart, - m, - delta_t, - delta_b, - grad_o, - grad_q, - grad_r, - scale, - cu_seqlens, - chunk_indices, - T, - HQ: tl.constexpr, - H: tl.constexpr, - G: tl.constexpr, - K: tl.constexpr, - BK: tl.constexpr, - WINDOW_SIZE_LEFT: tl.constexpr, - BT: tl.constexpr, - BS: tl.constexpr, - IS_VARLEN: tl.constexpr, -): - i_t, i_bh = tl.program_id(0), tl.program_id(1) - i_b, i_hq = i_bh // HQ, i_bh % HQ - i_h = i_hq // G - - 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 - else: - bos = (i_b * T).to(tl.int64) - RCP_LN2: tl.constexpr = 1.4426950216 - - row_offset = i_t * BT - row_indices = row_offset + tl.arange(0, BT) - row_mask = row_indices[:, None] < T - NUM_TOTAL_BLOCKS = tl.cdiv(tl.minimum(T, row_offset + BT), BS) - NUM_SAFE_BLOCKS = tl.minimum(row_offset, T) // BS - - if WINDOW_SIZE_LEFT >= 0: - leftmost_valid = tl.maximum(0, row_offset - WINDOW_SIZE_LEFT + 1) - FIRST_COL_BLOCK = leftmost_valid // BS - # Phase A is unmasked, so the safe zone must clear the window's left edge for - # the tile's LAST row (row_offset + BT - 1), not its first. - safe_left_valid = tl.maximum(0, row_offset + BT - WINDOW_SIZE_LEFT) - SAFE_LEFT_START = (safe_left_valid + BS - 1) // BS - else: - FIRST_COL_BLOCK = 0 - SAFE_LEFT_START = 0 - LEFT_BORDER_END = tl.minimum(SAFE_LEFT_START, NUM_SAFE_BLOCKS) - SAFE_MIDDLE_START = tl.maximum(FIRST_COL_BLOCK, SAFE_LEFT_START) - RIGHT_BORDER_START = tl.maximum(FIRST_COL_BLOCK, NUM_SAFE_BLOCKS) - - p_q = tl.make_block_ptr(q + (bos * HQ + i_hq) * K, (T, K), (HQ * K, 1), (row_offset, 0), (BT, BK), (1, 0)) - p_r = tl.make_block_ptr(r + (bos * HQ + i_hq) * K, (T, K), (HQ * K, 1), (row_offset, 0), (BT, BK), (1, 0)) - p_k = tl.make_block_ptr(k + (bos * H + i_h) * K, (T, K), (H * K, 1), (FIRST_COL_BLOCK * BS, 0), (BS, BK), (1, 0)) - p_v = tl.make_block_ptr(v + (bos * H + i_h) * K, (T, K), (H * K, 1), (FIRST_COL_BLOCK * BS, 0), (BS, BK), (1, 0)) - p_d1 = tl.make_block_ptr(d1 + bos * HQ + i_hq, (T, 1), (HQ, 1), (row_offset, 0), (BT, 1), (1, 0)) - p_bart = tl.make_block_ptr(bart + bos * HQ + i_hq, (T, 1), (HQ, 1), (row_offset, 0), (BT, 1), (1, 0)) - p_m = tl.make_block_ptr(m + bos * HQ + i_hq, (T, 1), (HQ, 1), (row_offset, 0), (BT, 1), (1, 0)) - p_t = tl.make_block_ptr(delta_t + bos * HQ + i_hq, (T, 1), (HQ, 1), (row_offset, 0), (BT, 1), (1, 0)) - p_b = tl.make_block_ptr(delta_b + bos * HQ + i_hq, (T, 1), (HQ, 1), (row_offset, 0), (BT, 1), (1, 0)) - p_grad_o = tl.make_block_ptr(grad_o + (bos * HQ + i_hq) * K, (T, K), (HQ * K, 1), (row_offset, 0), (BT, BK), (1, 0)) - p_grad_q = tl.make_block_ptr(grad_q + (bos * HQ + i_hq) * K, (T, K), (HQ * K, 1), (row_offset, 0), (BT, BK), (1, 0)) - p_grad_r = tl.make_block_ptr(grad_r + (bos * HQ + i_hq) * K, (T, K), (HQ * K, 1), (row_offset, 0), (BT, BK), (1, 0)) - - b_q = tl.load(p_q, boundary_check=(0, 1), padding_option="zero") - b_r = tl.load(p_r, boundary_check=(0, 1), padding_option="zero") - b_d1 = tl.load(p_d1, boundary_check=(0, 1), padding_option="zero") - b_bart = tl.load(p_bart, boundary_check=(0, 1), padding_option="zero") - b_m = tl.load(p_m, boundary_check=(0, 1), padding_option="zero") - b_t = tl.load(p_t, boundary_check=(0, 1), padding_option="zero") - b_b = tl.load(p_b, boundary_check=(0, 1), padding_option="zero") - grad_o_tile = tl.load(p_grad_o, boundary_check=(0, 1), padding_option="zero") - grad_q_acc = tl.zeros((BT, BK), dtype=tl.float32) - grad_r_acc = tl.zeros((BT, BK), dtype=tl.float32) - scale_log2 = scale * RCP_LN2 - - inv_d1 = tl.where(row_mask, 1.0 / b_d1, 0.0) - - # Phase 0: left-border blocks (SWA only). - for col_block_id in range(FIRST_COL_BLOCK, LEFT_BORDER_END): - col_indices = col_block_id * BS + tl.arange(0, BS) - b_k = tl.load(p_k, boundary_check=(0, 1), padding_option="zero") - b_v = tl.load(p_v, boundary_check=(0, 1), padding_option="zero") - mask = ( - (col_indices[None, :] >= row_indices[:, None] - WINDOW_SIZE_LEFT + 1) - & row_mask - & (col_indices[None, :] < T) - ) - qk = tl.dot(b_q, tl.trans(b_k), out_dtype=tl.float32) * scale_log2 - qk = tl.where(mask, qk, -float("inf")) - w = exp2(qk - b_m) - a = tl.dot(grad_o_tile, tl.trans(b_v), out_dtype=tl.float32) - rk = tl.dot(b_r, tl.trans(b_k), out_dtype=tl.float32) - p = w * inv_d1 - bart_minus_rk = b_bart - rk - delta = a - b_b - gl = p * (a - b_t + bart_minus_rk * delta) - gu = -p * delta - grad_q_acc = tl.dot(gl.to(b_k.dtype), b_k, out_dtype=tl.float32, acc=grad_q_acc) - grad_r_acc = tl.dot(gu.to(b_k.dtype), b_k, out_dtype=tl.float32, acc=grad_r_acc) - p_k = tl.advance(p_k, (BS, 0)) - p_v = tl.advance(p_v, (BS, 0)) - - # Phase A: safe blocks (no mask). - for _ in range(SAFE_MIDDLE_START, NUM_SAFE_BLOCKS): - b_k = tl.load(p_k, boundary_check=(0, 1), padding_option="zero") - b_v = tl.load(p_v, boundary_check=(0, 1), padding_option="zero") - qk = tl.dot(b_q, tl.trans(b_k), out_dtype=tl.float32) * scale_log2 - w = exp2(qk - b_m) - a = tl.dot(grad_o_tile, tl.trans(b_v), out_dtype=tl.float32) - rk = tl.dot(b_r, tl.trans(b_k), out_dtype=tl.float32) - p = w * inv_d1 - bart_minus_rk = b_bart - rk - delta = a - b_b - gl = p * (a - b_t + bart_minus_rk * delta) - gu = -p * delta - grad_q_acc = tl.dot(gl.to(b_k.dtype), b_k, out_dtype=tl.float32, acc=grad_q_acc) - grad_r_acc = tl.dot(gu.to(b_k.dtype), b_k, out_dtype=tl.float32, acc=grad_r_acc) - p_k = tl.advance(p_k, (BS, 0)) - p_v = tl.advance(p_v, (BS, 0)) - - # Phase B: right-border blocks (causal + boundary + window mask). - for col_block_id in range(RIGHT_BORDER_START, NUM_TOTAL_BLOCKS): - col_indices = col_block_id * BS + tl.arange(0, BS) - b_k = tl.load(p_k, boundary_check=(0, 1), padding_option="zero") - b_v = tl.load(p_v, boundary_check=(0, 1), padding_option="zero") - if WINDOW_SIZE_LEFT >= 0: - mask = ( - (row_indices[:, None] >= col_indices[None, :]) - & (col_indices[None, :] >= row_indices[:, None] - WINDOW_SIZE_LEFT + 1) - & row_mask - & (col_indices[None, :] < T) - ) - else: - mask = ( - (row_indices[:, None] >= col_indices[None, :]) - & row_mask - & (col_indices[None, :] < T) - ) - qk = tl.dot(b_q, tl.trans(b_k), out_dtype=tl.float32) * scale_log2 - qk = tl.where(mask, qk, -float("inf")) - w = exp2(qk - b_m) - a = tl.dot(grad_o_tile, tl.trans(b_v), out_dtype=tl.float32) - rk = tl.dot(b_r, tl.trans(b_k), out_dtype=tl.float32) - p = w * inv_d1 - bart_minus_rk = b_bart - rk - delta = a - b_b - gl = p * (a - b_t + bart_minus_rk * delta) - gu = -p * delta - grad_q_acc = tl.dot(gl.to(b_k.dtype), b_k, out_dtype=tl.float32, acc=grad_q_acc) - grad_r_acc = tl.dot(gu.to(b_k.dtype), b_k, out_dtype=tl.float32, acc=grad_r_acc) - p_k = tl.advance(p_k, (BS, 0)) - p_v = tl.advance(p_v, (BS, 0)) - - grad_q_acc = scale * grad_q_acc - - tl.store(p_grad_q, grad_q_acc.to(p_grad_q.dtype.element_ty), boundary_check=(0, 1)) - tl.store(p_grad_r, grad_r_acc.to(p_grad_r.dtype.element_ty), boundary_check=(0, 1)) - - -@triton.heuristics({ - 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None, -}) -@triton.jit(do_not_specialize=['T']) -def parallel_parallax_bwd_kernel_dkv( - q, - r, - k, - v, - d1, - bart, - m, - delta_t, - delta_b, - grad_o, - grad_k, - grad_v, - scale, - cu_seqlens, - chunk_indices, - T, - HQ: tl.constexpr, - H: tl.constexpr, - G: tl.constexpr, - K: tl.constexpr, - BK: tl.constexpr, - WINDOW_SIZE_LEFT: tl.constexpr, - BT: tl.constexpr, - BS: tl.constexpr, - IS_VARLEN: tl.constexpr, -): - i_t, i_bh = tl.program_id(0), tl.program_id(1) - i_b, i_hq = i_bh // HQ, i_bh % HQ - i_h = i_hq // G - - 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 - else: - bos = (i_b * T).to(tl.int64) - RCP_LN2: tl.constexpr = 1.4426950216 - - col_offset = i_t * BS - col_indices = col_offset + tl.arange(0, BS) - - start_row_block = col_offset // BT - start_row_offset = start_row_block * BT - - num_row_blocks_qbound = tl.cdiv(T, BT) - if WINDOW_SIZE_LEFT >= 0: - last_row_window = tl.cdiv(col_offset + BS + WINDOW_SIZE_LEFT - 1, BT) - num_row_blocks = tl.minimum(num_row_blocks_qbound, last_row_window) - WINDOW_SAFE_END = (col_offset + WINDOW_SIZE_LEFT) // BT - else: - num_row_blocks = num_row_blocks_qbound - WINDOW_SAFE_END = num_row_blocks - - p_q = tl.make_block_ptr(q + (bos * HQ + i_hq) * K, (T, K), (HQ * K, 1), (start_row_offset, 0), (BT, BK), (1, 0)) - p_r = tl.make_block_ptr(r + (bos * HQ + i_hq) * K, (T, K), (HQ * K, 1), (start_row_offset, 0), (BT, BK), (1, 0)) - p_k = tl.make_block_ptr(k + (bos * H + i_h) * K, (T, K), (H * K, 1), (col_offset, 0), (BS, BK), (1, 0)) - p_v = tl.make_block_ptr(v + (bos * H + i_h) * K, (T, K), (H * K, 1), (col_offset, 0), (BS, BK), (1, 0)) - p_d1 = tl.make_block_ptr(d1 + bos * HQ + i_hq, (T, 1), (HQ, 1), (start_row_offset, 0), (BT, 1), (1, 0)) - p_bart = tl.make_block_ptr(bart + bos * HQ + i_hq, (T, 1), (HQ, 1), (start_row_offset, 0), (BT, 1), (1, 0)) - p_m = tl.make_block_ptr(m + bos * HQ + i_hq, (T, 1), (HQ, 1), (start_row_offset, 0), (BT, 1), (1, 0)) - p_t = tl.make_block_ptr(delta_t + bos * HQ + i_hq, (T, 1), (HQ, 1), (start_row_offset, 0), (BT, 1), (1, 0)) - p_b = tl.make_block_ptr(delta_b + bos * HQ + i_hq, (T, 1), (HQ, 1), (start_row_offset, 0), (BT, 1), (1, 0)) - p_grad_o = tl.make_block_ptr(grad_o + (bos * HQ + i_hq) * K, (T, K), (HQ * K, 1), (start_row_offset, 0), (BT, BK), (1, 0)) - p_grad_k = tl.make_block_ptr(grad_k + (bos * HQ + i_hq) * K, (T, K), (HQ * K, 1), (col_offset, 0), (BS, BK), (1, 0)) - p_grad_v = tl.make_block_ptr(grad_v + (bos * HQ + i_hq) * K, (T, K), (HQ * K, 1), (col_offset, 0), (BS, BK), (1, 0)) - - b_k = tl.load(p_k, boundary_check=(0, 1), padding_option="zero") - b_v = tl.load(p_v, boundary_check=(0, 1), padding_option="zero") - grad_k_acc = tl.zeros((BS, BK), dtype=tl.float32) - grad_v_acc = tl.zeros((BS, BK), dtype=tl.float32) - scale_log2 = scale * RCP_LN2 - - first_safe_row_block = tl.cdiv(col_offset + BS, BT) - SAFE_MIDDLE_END = tl.minimum(WINDOW_SAFE_END, num_row_blocks) - WINDOW_BORDER_START = tl.maximum(first_safe_row_block, WINDOW_SAFE_END) - - # Phase A: causal-border row blocks. - causal_end = tl.minimum(first_safe_row_block, num_row_blocks) - for row_block_id in range(start_row_block, causal_end): - row_offset = row_block_id * BT - row_indices = row_offset + tl.arange(0, BT) - row_mask = row_indices[:, None] < T - b_q = tl.load(p_q, boundary_check=(0, 1), padding_option="zero") - b_r = tl.load(p_r, boundary_check=(0, 1), padding_option="zero") - b_d1 = tl.load(p_d1, boundary_check=(0, 1), padding_option="zero") - b_bart = tl.load(p_bart, boundary_check=(0, 1), padding_option="zero") - b_m = tl.load(p_m, boundary_check=(0, 1), padding_option="zero") - b_t = tl.load(p_t, boundary_check=(0, 1), padding_option="zero") - b_b = tl.load(p_b, boundary_check=(0, 1), padding_option="zero") - grad_o_tile = tl.load(p_grad_o, boundary_check=(0, 1), padding_option="zero") - - qk = tl.dot(b_q, tl.trans(b_k), out_dtype=tl.float32) * scale_log2 - rk = tl.dot(b_r, tl.trans(b_k), out_dtype=tl.float32) - inv_d1 = tl.where(row_mask, 1.0 / b_d1, 0.0) - if WINDOW_SIZE_LEFT >= 0: - mask = ( - (row_indices[:, None] >= col_indices[None, :]) - & (col_indices[None, :] >= row_indices[:, None] - WINDOW_SIZE_LEFT + 1) - & row_mask - & (col_indices[None, :] < T) - ) - else: - mask = ( - (row_indices[:, None] >= col_indices[None, :]) - & row_mask - & (col_indices[None, :] < T) - ) - qk = tl.where(mask, qk, -float("inf")) - w = exp2(qk - b_m) - p = w * inv_d1 - a = tl.dot(grad_o_tile, tl.trans(b_v), out_dtype=tl.float32) - delta = a - b_b - bart_minus_rk = b_bart - rk - gl = p * (a - b_t + bart_minus_rk * delta) * scale - gu = -p * delta - grad_k_acc = tl.dot(tl.trans(gl).to(b_q.dtype), b_q, out_dtype=tl.float32, acc=grad_k_acc) - grad_k_acc = tl.dot(tl.trans(gu).to(b_r.dtype), b_r, out_dtype=tl.float32, acc=grad_k_acc) - weights = p * (1 + bart_minus_rk) - grad_v_acc = tl.dot(tl.trans(weights).to(grad_o_tile.dtype), grad_o_tile, out_dtype=tl.float32, acc=grad_v_acc) - - p_q = tl.advance(p_q, (BT, 0)) - p_r = tl.advance(p_r, (BT, 0)) - p_d1 = tl.advance(p_d1, (BT, 0)) - p_bart = tl.advance(p_bart, (BT, 0)) - p_m = tl.advance(p_m, (BT, 0)) - p_t = tl.advance(p_t, (BT, 0)) - p_b = tl.advance(p_b, (BT, 0)) - p_grad_o = tl.advance(p_grad_o, (BT, 0)) - - # Phase B: safe row blocks (no causal/col/window mask). - safe_b_start = tl.maximum(first_safe_row_block, start_row_block) - for row_block_id in range(safe_b_start, SAFE_MIDDLE_END): - row_offset = row_block_id * BT - row_indices = row_offset + tl.arange(0, BT) - row_mask = row_indices[:, None] < T - b_q = tl.load(p_q, boundary_check=(0, 1), padding_option="zero") - b_r = tl.load(p_r, boundary_check=(0, 1), padding_option="zero") - b_d1 = tl.load(p_d1, boundary_check=(0, 1), padding_option="zero") - b_bart = tl.load(p_bart, boundary_check=(0, 1), padding_option="zero") - b_m = tl.load(p_m, boundary_check=(0, 1), padding_option="zero") - b_t = tl.load(p_t, boundary_check=(0, 1), padding_option="zero") - b_b = tl.load(p_b, boundary_check=(0, 1), padding_option="zero") - grad_o_tile = tl.load(p_grad_o, boundary_check=(0, 1), padding_option="zero") - - qk = tl.dot(b_q, tl.trans(b_k), out_dtype=tl.float32) * scale_log2 - rk = tl.dot(b_r, tl.trans(b_k), out_dtype=tl.float32) - inv_d1 = tl.where(row_mask, 1.0 / b_d1, 0.0) - w = exp2(qk - b_m) - p = w * inv_d1 - a = tl.dot(grad_o_tile, tl.trans(b_v), out_dtype=tl.float32) - delta = a - b_b - bart_minus_rk = b_bart - rk - gl = p * (a - b_t + bart_minus_rk * delta) * scale - gu = -p * delta - grad_k_acc = tl.dot(tl.trans(gl).to(b_q.dtype), b_q, out_dtype=tl.float32, acc=grad_k_acc) - grad_k_acc = tl.dot(tl.trans(gu).to(b_r.dtype), b_r, out_dtype=tl.float32, acc=grad_k_acc) - weights = p * (1 + bart_minus_rk) - grad_v_acc = tl.dot(tl.trans(weights).to(grad_o_tile.dtype), grad_o_tile, out_dtype=tl.float32, acc=grad_v_acc) - - p_q = tl.advance(p_q, (BT, 0)) - p_r = tl.advance(p_r, (BT, 0)) - p_d1 = tl.advance(p_d1, (BT, 0)) - p_bart = tl.advance(p_bart, (BT, 0)) - p_m = tl.advance(p_m, (BT, 0)) - p_t = tl.advance(p_t, (BT, 0)) - p_b = tl.advance(p_b, (BT, 0)) - p_grad_o = tl.advance(p_grad_o, (BT, 0)) - - # Phase C: window-border row blocks (SWA only). - window_border_start = tl.maximum(WINDOW_BORDER_START, start_row_block) - for row_block_id in range(window_border_start, num_row_blocks): - row_offset = row_block_id * BT - row_indices = row_offset + tl.arange(0, BT) - row_mask = row_indices[:, None] < T - b_q = tl.load(p_q, boundary_check=(0, 1), padding_option="zero") - b_r = tl.load(p_r, boundary_check=(0, 1), padding_option="zero") - b_d1 = tl.load(p_d1, boundary_check=(0, 1), padding_option="zero") - b_bart = tl.load(p_bart, boundary_check=(0, 1), padding_option="zero") - b_m = tl.load(p_m, boundary_check=(0, 1), padding_option="zero") - b_t = tl.load(p_t, boundary_check=(0, 1), padding_option="zero") - b_b = tl.load(p_b, boundary_check=(0, 1), padding_option="zero") - grad_o_tile = tl.load(p_grad_o, boundary_check=(0, 1), padding_option="zero") - - qk = tl.dot(b_q, tl.trans(b_k), out_dtype=tl.float32) * scale_log2 - rk = tl.dot(b_r, tl.trans(b_k), out_dtype=tl.float32) - inv_d1 = tl.where(row_mask, 1.0 / b_d1, 0.0) - mask = ( - (col_indices[None, :] >= row_indices[:, None] - WINDOW_SIZE_LEFT + 1) - & row_mask - & (col_indices[None, :] < T) - ) - qk = tl.where(mask, qk, -float("inf")) - w = exp2(qk - b_m) - p = w * inv_d1 - a = tl.dot(grad_o_tile, tl.trans(b_v), out_dtype=tl.float32) - delta = a - b_b - bart_minus_rk = b_bart - rk - gl = p * (a - b_t + bart_minus_rk * delta) * scale - gu = -p * delta - grad_k_acc = tl.dot(tl.trans(gl).to(b_q.dtype), b_q, out_dtype=tl.float32, acc=grad_k_acc) - grad_k_acc = tl.dot(tl.trans(gu).to(b_r.dtype), b_r, out_dtype=tl.float32, acc=grad_k_acc) - weights = p * (1 + bart_minus_rk) - grad_v_acc = tl.dot(tl.trans(weights).to(grad_o_tile.dtype), grad_o_tile, out_dtype=tl.float32, acc=grad_v_acc) - - p_q = tl.advance(p_q, (BT, 0)) - p_r = tl.advance(p_r, (BT, 0)) - p_d1 = tl.advance(p_d1, (BT, 0)) - p_bart = tl.advance(p_bart, (BT, 0)) - p_m = tl.advance(p_m, (BT, 0)) - p_t = tl.advance(p_t, (BT, 0)) - p_b = tl.advance(p_b, (BT, 0)) - p_grad_o = tl.advance(p_grad_o, (BT, 0)) - - tl.store(p_grad_k, grad_k_acc.to(p_grad_k.dtype.element_ty), boundary_check=(0, 1)) - tl.store(p_grad_v, grad_v_acc.to(p_grad_v.dtype.element_ty), boundary_check=(0, 1)) - - -def parallel_parallax_fwd(q, r, k, v, scale, cu_seqlens=None, chunk_indices=None, window_size_left=-1): - """Parallax forward (Triton). `(B, T, HQ, D)` / packed `(1, T_total, HQ, D)` inputs. - - Returns `(o, barv, d1, bart, m)`: `o`/`barv` in the input dtype and layout; - `d1`/`bart`/`m` are fp32 per-(position, query-head) scalars `(B, T, HQ)`. - """ - B, T, HQ, K = q.shape - H = k.shape[2] - G = HQ // H - BK = triton.next_power_of_2(K) - BT = _block_size(K, q.device.index) - o = torch.empty_like(q) - barv = torch.empty_like(q) - d1 = torch.empty((B, T, HQ), device=q.device, dtype=torch.float32) - bart = torch.empty((B, T, HQ), device=q.device, dtype=torch.float32) - m = torch.empty((B, T, HQ), device=q.device, dtype=torch.float32) - - NT = triton.cdiv(T, BT) if cu_seqlens is None else len(chunk_indices) - grid = (NT, B * HQ) - parallel_parallax_fwd_kernel[grid]( - q, r, k, v, o, barv, d1, bart, m, - scale, cu_seqlens, chunk_indices, T, - HQ=HQ, H=H, G=G, K=K, BK=BK, - WINDOW_SIZE_LEFT=window_size_left, BT=BT, BS=BT, - num_warps=8, num_stages=2, - ) - return o, barv, d1, bart, m - - -def parallel_parallax_bwd(q, r, k, v, o, barv, d1, bart, m, grad_o, scale, cu_seqlens=None, chunk_indices=None, window_size_left=-1): - """Parallax backward (Triton). Returns grads matching `q, r, k, v`.""" - B, T, HQ, K = q.shape - H = k.shape[2] - G = HQ // H - BK = triton.next_power_of_2(K) - BT = _block_size(K, q.device.index) - - grad_q = torch.empty_like(q) - grad_r = torch.empty_like(r) - # dK/dV are written per q-head, then folded back to the kv-head axis. - grad_k_buf = torch.empty((B, T, HQ, K), device=q.device, dtype=q.dtype) - grad_v_buf = torch.empty((B, T, HQ, K), device=q.device, dtype=q.dtype) - delta_t = torch.empty((B, T, HQ), device=q.device, dtype=torch.float32) - delta_b = torch.empty((B, T, HQ), device=q.device, dtype=torch.float32) - - NT = triton.cdiv(T, BT) if cu_seqlens is None else len(chunk_indices) - grid = (NT, B * HQ) - - parallel_parallax_bwd_kernel_preprocess[grid]( - grad_o, o, barv, delta_t, delta_b, - cu_seqlens, chunk_indices, T, - HQ=HQ, K=K, BK=BK, BT=BT, - num_warps=4, num_stages=2, - ) - parallel_parallax_bwd_kernel_dqr[grid]( - q, r, k, v, d1, bart, m, delta_t, delta_b, grad_o, grad_q, grad_r, - scale, cu_seqlens, chunk_indices, T, - HQ=HQ, H=H, G=G, K=K, BK=BK, - WINDOW_SIZE_LEFT=window_size_left, BT=BT, BS=BT, - num_warps=8, num_stages=2, - ) - parallel_parallax_bwd_kernel_dkv[grid]( - q, r, k, v, d1, bart, m, delta_t, delta_b, grad_o, grad_k_buf, grad_v_buf, - scale, cu_seqlens, chunk_indices, T, - HQ=HQ, H=H, G=G, K=K, BK=BK, - WINDOW_SIZE_LEFT=window_size_left, BT=BT, BS=BT, - num_warps=8, num_stages=2, - ) - - if G == 1: - grad_k = grad_k_buf - grad_v = grad_v_buf - else: - grad_k = reduce(grad_k_buf, 'b t (h g) k -> b t h k', g=G, reduction='sum') - grad_v = reduce(grad_v_buf, 'b t (h g) k -> b t h k', g=G, reduction='sum') - return grad_q, grad_r, grad_k, grad_v - - -class ParallaxFunction(torch.autograd.Function): - - @staticmethod - @contiguous - @autocast_custom_fwd - def forward(ctx, q, r, k, v, scale, window_size_left, cu_seqlens): - chunk_indices = prepare_chunk_indices(cu_seqlens, _block_size(q.shape[-1], q.device.index)) \ - if cu_seqlens is not None else None - o, barv, d1, bart, m = parallel_parallax_fwd(q, r, k, v, scale, cu_seqlens, chunk_indices, window_size_left) - ctx.save_for_backward(q, r, k, v, o, barv, d1, bart, m) - ctx.scale = scale - ctx.window_size_left = window_size_left - ctx.cu_seqlens = cu_seqlens - ctx.chunk_indices = chunk_indices - return o - - @staticmethod - @contiguous - @autocast_custom_bwd - def backward(ctx, do): - q, r, k, v, o, barv, d1, bart, m = ctx.saved_tensors - gq, gr, gk, gv = parallel_parallax_bwd( - q, r, k, v, o, barv, d1, bart, m, do, - ctx.scale, ctx.cu_seqlens, ctx.chunk_indices, ctx.window_size_left, - ) - return gq.to(q), gr.to(r), gk.to(k), gv.to(v), None, None, None - - -def parallel_parallax( - q: torch.Tensor, - r: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - scale: float | None = None, - window_size: int | None = None, - cu_seqlens: torch.LongTensor | None = None, - **kwargs, -) -> torch.Tensor: - r""" - Causal Parallax (parameterized local linear attention) with autograd, - backed by Triton kernels. See `fla.ops.parallax.naive.naive_parallax` for - the reference math. - - Args: - q (torch.Tensor): - queries of shape `[B, T, HQ, D]`. - r (torch.Tensor): - secondary queries of shape `[B, T, HQ, D]` (same shape as `q`). NOTE: - `r` is *not* scaled by `scale`; pass it un-pre-scaled. - k (torch.Tensor): - keys of shape `[B, T, H, D]`. GQA is applied when `HQ` is divisible by `H`. - v (torch.Tensor): - values of shape `[B, T, H, D]`. - scale (float, Optional): - Scale applied to the `q @ k^T` logits only. If `None`, defaults to `1 / sqrt(D)`. - Default: `None`. - window_size (int, Optional): - Sliding-window length. If provided, each query at position `i` only attends to - keys in `[i - window_size + 1, i]`. If `None`, full causal attention is used. - Default: `None`. - cu_seqlens (torch.LongTensor, Optional): - Cumulative sequence lengths of shape `[N+1]` for variable-length training - (FlashAttention convention). The batch size must be 1 when packing. Default: `None`. - - Returns: - o (torch.Tensor): - output of shape `[B, T, HQ, D]`. - """ - if 'head_first' in kwargs: - raise DeprecationWarning( - "head_first has been removed. Inputs must be in `[B, T, H, ...]` format.", - ) - if q.dtype not in (torch.bfloat16, torch.float16): - raise TypeError(f"parallel_parallax requires bf16 or fp16 inputs, got q.dtype={q.dtype}") - if scale is None: - scale = k.shape[-1] ** -0.5 - if cu_seqlens is not None and q.shape[0] != 1: - raise ValueError( - f"The batch size is expected to be 1 rather than {q.shape[0]} when using `cu_seqlens`. " - f"Please flatten variable-length inputs before processing.", - ) - # The kernel keeps cols [i - W + 1, i] (W keys total, diagonal included), - # matching FLA's `window_size=W` semantics exactly (no off-by-one). - window_size_left = -1 if window_size is None else window_size - return ParallaxFunction.apply(q, r, k, v, float(scale), window_size_left, cu_seqlens) diff --git a/build/torch-xpu/ops/path_attn/__init__.py b/build/torch-xpu/ops/path_attn/__init__.py deleted file mode 100644 index bb3e57da3de190497201c2913eec2381b74740f0..0000000000000000000000000000000000000000 --- a/build/torch-xpu/ops/path_attn/__init__.py +++ /dev/null @@ -1,14 +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 .naive import naive_path_attn -from .parallel import parallel_path_attn - -__all__ = [ - 'naive_path_attn', - 'parallel_path_attn', -] diff --git a/build/torch-xpu/ops/path_attn/cumprod_householder_bwd.py b/build/torch-xpu/ops/path_attn/cumprod_householder_bwd.py deleted file mode 100644 index b7bd259ca925d3b4d0598139591f18a9a8a8b7d6..0000000000000000000000000000000000000000 --- a/build/torch-xpu/ops/path_attn/cumprod_householder_bwd.py +++ /dev/null @@ -1,149 +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 -import triton.language as tl - -from ...ops.utils import prepare_chunk_indices, prepare_chunk_offsets -from ...utils import check_shared_mem - - -@triton.heuristics({ - 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None, -}) -@triton.jit(do_not_specialize=['T']) -def chunk_cumprod_householder_bwd_kernel( - hc_suffix, dhc_whole, - k, dk, w1, w2, dw1, dw2, dk_new, - cu_seqlens, split_indices, chunk_offsets, split_offsets, - BT: tl.constexpr, # previous small chunk size - K: tl.constexpr, - BK: tl.constexpr, - T, - S: tl.constexpr, - G: tl.constexpr, - H: tl.constexpr, - HQ: tl.constexpr, - IS_VARLEN: tl.constexpr, -): - i_ss, i_hq = tl.program_id(0), tl.program_id(1) - i_h = i_hq // G - - if IS_VARLEN: - i_n, i_s = tl.load(split_indices + i_ss * 2).to(tl.int32), tl.load(split_indices + i_ss * 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).to(tl.int32) - NS = tl.cdiv(T, S) - boh = tl.load(chunk_offsets + i_n).to(tl.int32) - boh_large = tl.load(split_offsets + i_n).to(tl.int32) - else: - NS = tl.cdiv(T, S) - i_n, i_s = i_ss // NS, i_ss % NS - bos, eos = (i_n * T).to(tl.int64), (i_n * T + T).to(tl.int64) - boh = i_n * tl.cdiv(T, BT) - boh_large = i_n * tl.cdiv(T, S) - - # offset calculations - dhc_whole += ((boh_large + i_s) * HQ + i_hq) * K * K - hc_suffix += ((boh + tl.cdiv(i_s * S, BT)) * H + i_h) * K * K - k += (bos * H + i_h) * K - w1 += (bos * H + i_h) * K - w2 += (bos * H + i_h) * K - dw1 += (bos * HQ + i_hq) * K - dw2 += (bos * HQ + i_hq) * K - - # dh += ((boh + tl.cdiv(i_s * S, BT)) * HQ + i_hq) * K * K - dk += (bos * HQ + i_hq) * K - dk_new += (bos * HQ + i_hq) * K - - stride_h = H * K * K - NT_small = tl.cdiv(min(S, T-i_s*S), BT) - p_dhc_whole = tl.make_block_ptr(dhc_whole, (K, K), (K, 1), (0, 0), (BK, BK), (1, 0)) - b_dhc = tl.zeros([BK, BK], dtype=tl.float32) - b_dhc += tl.load(p_dhc_whole, boundary_check=(0, 1)) - - # calculate dh - for i_t_small in range(0, NT_small): - p_k = tl.make_block_ptr(k, (T, K), (H*K, 1), (i_s*S + i_t_small*BT, 0), (BT, BK), (1, 0)) - p_dk = tl.make_block_ptr(dk, (T, K), (HQ*K, 1), (i_s*S + i_t_small*BT, 0), (BT, BK), (1, 0)) - p_dk_new = tl.make_block_ptr(dk_new, (T, K), (HQ*K, 1), (i_s*S + i_t_small*BT, 0), (BT, BK), (1, 0)) - p_hc = tl.make_block_ptr(hc_suffix + i_t_small * stride_h, (K, K), (K, 1), (0, 0), (BK, BK), (1, 0)) - - b_k = tl.load(p_k, boundary_check=(0, 1)) - b_dk = tl.load(p_dk, boundary_check=(0, 1)) - - p_w1 = tl.make_block_ptr(w1, (T, K), (H*K, 1), (i_s*S + i_t_small*BT, 0), (BT, BK), (1, 0)) - p_w2 = tl.make_block_ptr(w2, (T, K), (H*K, 1), (i_s*S + i_t_small*BT, 0), (BT, BK), (1, 0)) - - b_w1 = tl.load(p_w1, boundary_check=(0, 1)) - b_w2 = tl.load(p_w2, boundary_check=(0, 1)) - b_hc = tl.load(p_hc, boundary_check=(0, 1)) - - b_dk_new = b_dk - tl.dot(b_dk.to(b_hc.dtype), b_hc) - tl.store(p_dk_new, b_dk_new.to(dk_new.dtype.element_ty), boundary_check=(0, 1)) - - b_dh = b_dhc - tl.dot(tl.trans(b_hc), b_dhc.to(b_hc.dtype)) - b_dw2 = tl.dot(b_w1, b_dh.to(b_w1.dtype)) - b_dw1 = tl.dot(b_w2, tl.trans(b_dh.to(b_w2.dtype))) - - p_dw1 = tl.make_block_ptr(dw1, (T, K), (HQ*K, 1), (i_s*S + i_t_small*BT, 0), (BT, BK), (1, 0)) - p_dw2 = tl.make_block_ptr(dw2, (T, K), (HQ*K, 1), (i_s*S + i_t_small*BT, 0), (BT, BK), (1, 0)) - - tl.store(p_dw1, b_dw1.to(dw1.dtype.element_ty), boundary_check=(0, 1)) - tl.store(p_dw2, b_dw2.to(dw2.dtype.element_ty), boundary_check=(0, 1)) - - b_dhc = b_dhc - tl.dot(tl.dot(b_dhc.to(b_w2.dtype), tl.trans(b_w2)).to(b_w1.dtype), b_w1) - b_dhc -= tl.dot(tl.trans(b_dk).to(b_k.dtype), b_k) - - -def chunk_cumprod_householder_bwd_fn( - w1: torch.Tensor, - w2: torch.Tensor, - hc_suffix: torch.Tensor, - dhc_whole: torch.Tensor, - k: torch.Tensor, - dk: torch.Tensor, - S: int, # split size, aka large chunk size - BT: int, # small chunk size - cu_seqlens: torch.Tensor = None, - chunk_indices: torch.LongTensor | None = None, -): - B, T, HQ, K = dk.shape - H = k.shape[2] - G = HQ // H - - if chunk_indices is None and cu_seqlens is not None: - chunk_indices = prepare_chunk_indices(cu_seqlens, S) - split_indices = chunk_indices - chunk_offsets = prepare_chunk_offsets(cu_seqlens, BT) if cu_seqlens is not None else None - split_offsets = prepare_chunk_offsets(cu_seqlens, S) if cu_seqlens is not None else None - - if cu_seqlens is None: - N = B - NS = N * triton.cdiv(T, S) - else: - N = len(cu_seqlens) - 1 - NS = split_offsets[-1].item() - - grid = (NS, HQ) - dw1 = torch.empty_like(dk, dtype=torch.float32) - dw2 = torch.empty_like(dk, dtype=torch.float32) - dk_new = torch.empty_like(dk, dtype=torch.float32) - - chunk_cumprod_householder_bwd_kernel[grid]( - hc_suffix=hc_suffix, dhc_whole=dhc_whole, - k=k, dk=dk, w1=w1, w2=w2, dw1=dw1, dw2=dw2, dk_new=dk_new, - cu_seqlens=cu_seqlens, - split_indices=split_indices, chunk_offsets=chunk_offsets, split_offsets=split_offsets, - BT=BT, K=K, G=G, H=H, HQ=HQ, BK=K, - T=T, S=S, - # SY (2025/07/08): I don't know why when K == 128 if I set num_warps=4 the result would be completely wrong - num_warps=8 if K == 128 else 4, - num_stages=2 if check_shared_mem('ampere') else 1, - ) - return dw1, dw2, dk_new diff --git a/build/torch-xpu/ops/path_attn/cumprod_householder_fwd.py b/build/torch-xpu/ops/path_attn/cumprod_householder_fwd.py deleted file mode 100644 index 1ced41573e5352b81472bcb492911891e126fb68..0000000000000000000000000000000000000000 --- a/build/torch-xpu/ops/path_attn/cumprod_householder_fwd.py +++ /dev/null @@ -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 - -import torch -import triton -import triton.language as tl - -from ...ops.utils import prepare_chunk_indices, prepare_chunk_offsets -from ...utils import check_shared_mem - - -@triton.heuristics({ - 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None, -}) -@triton.jit -def chunk_cumprod_householder_fwd_kernel( - k, - k_new, - w1, - w2, - hc_suffix, - hc_whole, - cu_seqlens, - split_indices, - chunk_offsets, - split_offsets, - BT: tl.constexpr, # small chunk size - K: tl.constexpr, - H: tl.constexpr, - BK: tl.constexpr, - T: tl.constexpr, - S: tl.constexpr, # split size, aka large chunk size - IS_VARLEN: tl.constexpr, -): - i_ss, i_h = tl.program_id(0), tl.program_id(1) - - if IS_VARLEN: - i_n, i_s = tl.load(split_indices + i_ss * 2).to(tl.int32), tl.load(split_indices + i_ss * 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).to(tl.int32) - NS = tl.cdiv(T, S) - - boh = tl.load(chunk_offsets + i_n).to(tl.int32) - boh_large = tl.load(split_offsets + i_n).to(tl.int32) - else: - NS = tl.cdiv(T, S) - i_n, i_s = i_ss // NS, i_ss % NS - bos, eos = (i_n * T).to(tl.int64), (i_n * T + T).to(tl.int64) - - boh = i_n * tl.cdiv(T, BT) - boh_large = i_n * tl.cdiv(T, S) - - NT_small = tl.cdiv(min(S, T-i_s*S), BT) - stride_h = H*K*K - - # offset calculations - hc_whole += ((boh_large + i_s) * H + i_h) * K * K - hc_suffix += ((boh + tl.cdiv(i_s * S, BT)) * H + i_h) * K * K - - k += (bos * H + i_h) * K - k_new += (bos * H + i_h) * K - w1 += (bos * H + i_h) * K - w2 += (bos * H + i_h) * K - - b_h = tl.zeros([BK, BK], dtype=tl.float32) - for i_t_small in range(NT_small-1, -1, -1): - p_hc_suffix = tl.make_block_ptr(hc_suffix + i_t_small * stride_h, (K, K), (K, 1), (0, 0), (BK, BK), (1, 0)) - tl.store(p_hc_suffix, b_h.to(hc_suffix.dtype.element_ty), boundary_check=(0, 1)) - p_k = tl.make_block_ptr(k, (T, K), (H*K, 1), (i_s * S + i_t_small * BT, 0), (BT, BK), (1, 0)) - b_k = tl.load(p_k, boundary_check=(0, 1)) - b_k = (b_k - tl.dot(b_k, tl.trans(b_h.to(b_k.dtype)))) - p_w1 = tl.make_block_ptr(w1, (K, T), (1, H*K), (0, i_s * S + i_t_small * BT), (BK, BT), (0, 1)) - p_w2 = tl.make_block_ptr(w2, (T, K), (H*K, 1), (i_s * S + i_t_small * BT, 0), (BT, BK), (1, 0)) - b_w1 = tl.load(p_w1, boundary_check=(0, 1)) - b_w2 = tl.load(p_w2, boundary_check=(0, 1)) - b_v_new = (b_w1 - tl.dot(b_h.to(b_w1.dtype), b_w1)).to(b_w2.dtype) - b_h += tl.dot(b_v_new, b_w2) - p_k_new = tl.make_block_ptr(k_new, (T, K), (H*K, 1), (i_s * S + i_t_small * BT, 0), (BT, BK), (1, 0)) - tl.store(p_k_new, b_k.to(k_new.dtype.element_ty), boundary_check=(0, 1)) - - p_hc_whole = tl.make_block_ptr(hc_whole, (K, K), (K, 1), (0, 0), (BK, BK), (1, 0)) - tl.store(p_hc_whole, b_h.to(hc_whole.dtype.element_ty), boundary_check=(0, 1)) - - -def chunk_cumprod_householder_fwd_fn( - k: torch.Tensor, - w1: torch.Tensor, - w2: torch.Tensor, - S: int, # split size, aka large chunk size - BT: int, # small chunk size - cu_seqlens: torch.Tensor = None, - chunk_indices: torch.LongTensor | None = None, -): - B, T, H, K = k.shape - - if chunk_indices is None and cu_seqlens is not None: - chunk_indices = prepare_chunk_indices(cu_seqlens, S) - split_indices = chunk_indices - chunk_offsets = prepare_chunk_offsets(cu_seqlens, BT) if cu_seqlens is not None else None - split_offsets = prepare_chunk_offsets(cu_seqlens, S) if cu_seqlens is not None else None - - if cu_seqlens is None: - N = B - NS = N * triton.cdiv(T, S) - NT = N * triton.cdiv(T, BT) - else: - N = len(cu_seqlens) - 1 - NS = split_offsets[-1] - NT = chunk_offsets[-1] - - grid = (NS, H) - hc_whole = torch.empty((NS, H, K, K), device=k.device, dtype=w1.dtype) - k_new = torch.empty_like(k, dtype=k.dtype) - hc_suffix = torch.empty((NT, H, K, K), device=k.device, dtype=w1.dtype) - chunk_cumprod_householder_fwd_kernel[grid]( - k=k, k_new=k_new, w1=w1, w2=w2, hc_whole=hc_whole, hc_suffix=hc_suffix, - cu_seqlens=cu_seqlens, - split_indices=split_indices, chunk_offsets=chunk_offsets, split_offsets=split_offsets, - BT=BT, K=K, H=H, BK=K, - T=T, S=S, - # SY (2025/07/08): I don't know why when K == 128 if I set num_warps=4 the result would be completely wrong - num_warps=8 if K == 128 else 4, - num_stages=3 if check_shared_mem('ampere') else 1, - ) - return k_new, hc_suffix, hc_whole diff --git a/build/torch-xpu/ops/path_attn/intra_chunk_preprocess_bwd.py b/build/torch-xpu/ops/path_attn/intra_chunk_preprocess_bwd.py deleted file mode 100644 index d7b3674d2740c0a1f323b4e51dbf8ed42962942f..0000000000000000000000000000000000000000 --- a/build/torch-xpu/ops/path_attn/intra_chunk_preprocess_bwd.py +++ /dev/null @@ -1,151 +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 -import triton.language as tl - -from ...ops.utils import prepare_chunk_indices -from ...utils import check_shared_mem - - -# episold -@triton.heuristics({ - 'IS_VARLEN': lambda args: args['offsets'] is not None, -}) -@triton.jit(do_not_specialize=['T']) -def intra_chunk_preprocess_bwd_kernel( - q, k, w, w2, beta, - AT, - dA_local, dq, dq_new, dk, dk_new, dw, dbeta, dw1, dw2, T, - offsets, indices, - HQ: tl.constexpr, G: tl.constexpr, H: tl.constexpr, - K: tl.constexpr, BT: tl.constexpr, BK: tl.constexpr, - IS_VARLEN: tl.constexpr, -): - i_t, i_nh = tl.program_id(0), tl.program_id(1) - i_n, i_hq = i_nh // HQ, i_nh % HQ - i_h = i_hq // G - - if IS_VARLEN: - i_n, i_t = tl.load(indices + i_t * 2).to(tl.int32), tl.load(indices + i_t * 2 + 1).to(tl.int32) - bos, eos = tl.load(offsets + i_n).to(tl.int64), tl.load(offsets + i_n + 1).to(tl.int64) - T = (eos - bos).to(tl.int32) - else: - bos, eos = (i_n * T).to(tl.int64), (i_n * T + T).to(tl.int64) - - b_dk = tl.zeros([BT, BK], dtype=tl.float32) - b_dw_beta = tl.zeros([BT, BK], dtype=tl.float32) - b_dw = tl.zeros([BT, BK], dtype=tl.float32) - b_dT = tl.zeros([BT, BT], dtype=tl.float32) - - p_q = tl.make_block_ptr(q + (bos * HQ + i_hq) * K, (T, K), (K*HQ, 1), (i_t * BT, 0), (BT, BK), (1, 0)) - p_k = tl.make_block_ptr(k + (bos * H + i_h) * K, (T, K), (K*H, 1), (i_t * BT, 0), (BT, BK), (1, 0)) - p_w = tl.make_block_ptr(w + (bos * H + i_h) * K, (T, K), (K*H, 1), (i_t * BT, 0), (BT, BK), (1, 0)) - p_w2 = tl.make_block_ptr(w2 + (bos * H + i_h) * K, (T, K), (K*H, 1), (i_t * BT, 0), (BT, BK), (1, 0)) - p_beta = tl.make_block_ptr(beta + (bos * H + i_h), (T, ), (H, ), (i_t * BT, ), (BT, ), (0, )) - p_T = tl.make_block_ptr(AT + (bos * H + i_h) * BT, (T, BT), (BT*H, 1), (i_t * BT, 0), (BT, BT), (1, 0)) - b_w = tl.load(p_w, boundary_check=(0, 1)) - b_Twb = tl.load(p_w2, boundary_check=(0, 1)) - b_beta = tl.load(p_beta, boundary_check=(0, )) - b_q = tl.load(p_q, boundary_check=(0, 1)) - b_k = tl.load(p_k, boundary_check=(0, 1)) - b_T = tl.load(p_T, boundary_check=(0, 1)) - b_w_beta = (b_w * b_beta[:, None]).to(b_w.dtype) - - o_i = tl.arange(0, BT) - b_qw = tl.where(o_i[:, None] >= o_i[None, :], tl.dot(b_q, tl.trans(b_w)), 0).to(b_q.dtype) - b_wbk = tl.where(o_i[:, None] > o_i[None, :], tl.dot(b_w_beta, tl.trans(b_k)), 0).to(b_k.dtype) - b_Twbk = tl.dot(b_T, b_wbk).to(b_w.dtype) - - p_dA_local = tl.make_block_ptr(dA_local + (bos * HQ + i_hq) * BT, (T, BT), (BT*HQ, 1), (i_t * BT, 0), (BT, BT), (1, 0)) - b_dA_local = tl.load(p_dA_local, boundary_check=(0, 1)) - - # # Twb part qw part. - p_dq = tl.make_block_ptr(dq + (bos * HQ + i_hq) * K, (T, K), (K*HQ, 1), (i_t * BT, 0), (BT, BK), (1, 0)) - b_dq = tl.load(p_dq, boundary_check=(0, 1)) - - p_dw1 = tl.make_block_ptr(dw1 + (bos * HQ + i_hq) * K, (T, K), (K*HQ, 1), (i_t * BT, 0), (BT, BK), (1, 0)) - b_dw += tl.load(p_dw1, boundary_check=(0, 1)) - - b_dqw = -tl.dot(b_dA_local, tl.trans(b_Twbk)) - tl.dot(b_dq.to(b_Twb.dtype), tl.trans(b_Twb)) - p_dw2 = tl.make_block_ptr(dw2 + (bos * HQ + i_hq) * K, (T, K), (K*HQ, 1), (i_t * BT, 0), (BT, BK), (1, 0)) - b_dTwb = -tl.dot(tl.trans(b_qw), b_dq) + tl.load(p_dw2, boundary_check=(0, 1)) - b_dT += tl.dot(b_dTwb.to(b_w_beta.dtype), tl.trans(b_w_beta)) - b_dw_beta += tl.dot(tl.trans(b_T), b_dTwb.to(b_T.dtype)) - - b_dqw = tl.where(tl.arange(0, BT)[:, None] >= tl.arange(0, BT)[None, :], b_dqw, 0) - b_dq += tl.dot(b_dA_local.to(b_k.dtype), b_k) - b_dq += tl.dot(b_dqw.to(b_w.dtype), b_w) - b_dw += tl.dot(tl.trans(b_dqw.to(b_q.dtype)), b_q) - p_q_new = tl.make_block_ptr(dq_new + (bos * HQ + i_hq) * K, (T, K), (K*HQ, 1), (i_t * BT, 0), (BT, BK), (1, 0)) - tl.store(p_q_new, b_dq.to(dq_new.dtype.element_ty), boundary_check=(0, 1)) - - # Twbk part - p_dk = tl.make_block_ptr(dk + (bos * HQ + i_hq) * K, (T, K), (K*HQ, 1), (i_t * BT, 0), (BT, BK), (1, 0)) - b_dk = tl.load(p_dk, boundary_check=(0, 1)) - b_dTwbk = -tl.dot(tl.trans(b_qw), b_dA_local.to(b_qw.dtype)) - tl.dot(b_w, tl.trans(b_dk.to(b_w.dtype))) - b_dw -= tl.dot(b_Twbk, b_dk.to(b_w.dtype)) - b_dT += tl.dot(b_dTwbk.to(b_wbk.dtype), tl.trans(b_wbk)) - b_dwbk = tl.where(o_i[:, None] > o_i[None, :], tl.dot(tl.trans(b_T), b_dTwbk.to(b_T.dtype)), 0).to(b_w.dtype) - b_dw_beta += tl.dot(b_dwbk, b_k) - - b_dk += tl.dot(tl.trans(b_dwbk), b_w_beta) - b_dk += tl.dot(tl.trans(b_dA_local), b_q) - p_dk_new = tl.make_block_ptr(dk_new + (bos * HQ + i_hq) * K, (T, K), (K*HQ, 1), (i_t * BT, 0), (BT, BK), (1, 0)) - tl.store(p_dk_new, b_dk.to(dk_new.dtype.element_ty), boundary_check=(0, 1)) - - # matrix inverse's gradient - p_T = tl.make_block_ptr(AT + (bos * H + i_h) * BT, (BT, T), (1, BT*H), (0, i_t * BT), (BT, BT), (0, 1)) - b_Tt = tl.load(p_T, boundary_check=(0, 1)) - b_dT = tl.where(tl.arange(0, BT)[:, None] > tl.arange(0, BT)[None, :], b_dT, 0).to(b_w.dtype) - b_dT = tl.dot(b_Tt, b_dT).to(b_w.dtype) - b_dT = tl.dot(b_dT, b_Tt) - b_dT = tl.where(tl.arange(0, BT)[:, None] > tl.arange(0, BT)[None, :], -b_dT, 0).to(b_k.dtype) - - b_dw_beta += tl.dot(b_dT, b_w) - b_dw += tl.dot(tl.trans(b_dT), b_w_beta) - b_dw += b_dw_beta * b_beta[:, None] - b_dbeta = tl.sum(b_dw_beta * b_w, axis=1) - - p_dw = tl.make_block_ptr(dw + (bos * HQ + i_hq) * K, (T, K), (K*HQ, 1), (i_t * BT, 0), (BT, BK), (1, 0)) - tl.store(p_dw, b_dw.to(dw.dtype.element_ty), boundary_check=(0, 1)) - p_dbeta = tl.make_block_ptr(dbeta + (bos * HQ + i_hq), (T, ), (HQ, ), (i_t * BT, ), (BT, ), (0, )) - tl.store(p_dbeta, b_dbeta.to(dbeta.dtype.element_ty), boundary_check=(0, )) - - -def intra_chunk_preprocess_bwd_fn(q, k, w, w2, beta, - dq, dk, dA_local, - dw1, dw2, - A, L, D, do, scale, cu_seqlens=None, - chunk_indices: torch.LongTensor | None = None): - BT = A.shape[-1] - HQ = q.shape[-2] - B, T, H, K = k.shape - G = HQ//H - if chunk_indices is None and cu_seqlens is not None: - chunk_indices = prepare_chunk_indices(cu_seqlens, BT) - indices = chunk_indices - NT = triton.cdiv(T, BT) if cu_seqlens is None else len(indices) - grid = (NT, B*HQ) - # better precision because h would be of norm smaller than 1 anyways - - dbeta = torch.empty(B, T, HQ, device=q.device, dtype=k.dtype if G == 1 else torch.float32) - dw = torch.empty(B, T, HQ, K, device=q.device, dtype=k.dtype if G == 1 else torch.float32) - dk_new = torch.empty_like(dk, dtype=k.dtype if G == 1 else torch.float32) # float32 reduction - dq_new = torch.empty_like(dq, dtype=q.dtype) - - intra_chunk_preprocess_bwd_kernel[grid]( - q=q, k=k, w=w, w2=w2, beta=beta, - AT=A, - dA_local=dA_local, dq=dq, dq_new=dq_new, dk=dk, dk_new=dk_new, dw=dw, dbeta=dbeta, dw1=dw1, dw2=dw2, T=T, - offsets=cu_seqlens, indices=indices, - HQ=HQ, G=G, H=H, - K=K, BT=BT, BK=triton.next_power_of_2(K), - num_stages=3 if check_shared_mem('hopper') else 1, - ) - return dq_new, dk_new, dbeta, dw diff --git a/build/torch-xpu/ops/path_attn/intra_chunk_preprocess_bwd_prepare.py b/build/torch-xpu/ops/path_attn/intra_chunk_preprocess_bwd_prepare.py deleted file mode 100644 index f39242a4522290eb721ebf4b9ce0237ea99a2bb6..0000000000000000000000000000000000000000 --- a/build/torch-xpu/ops/path_attn/intra_chunk_preprocess_bwd_prepare.py +++ /dev/null @@ -1,207 +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 -import triton.language as tl - -from ...ops.utils import prepare_chunk_indices, prepare_chunk_offsets - - -@triton.heuristics({ - "USE_GATE": lambda args: args['g_cumsum'] is not None, - "IS_VARLEN": lambda args: args['offsets'] is not None, -}) -@triton.jit(do_not_specialize=['T']) -def chunk_transform_qk_bwd_kernel_prepare( - q, - k, - v, - w, - beta, - g_cumsum, - L, - D, - h, - q_new, - k_new, - AT, - dA_local, - dv, - do, - dg_cumsum, - scale, - indices, # varlen helper - offsets, # varlen helper - chunk_offsets, # varlen helper - T, - G: tl.constexpr, - HQ: tl.constexpr, - H: tl.constexpr, - K: tl.constexpr, - V: tl.constexpr, - BK: tl.constexpr, - BV: tl.constexpr, - BT: tl.constexpr, - IS_VARLEN: tl.constexpr, - USE_GATE: tl.constexpr, - RETURN_H: tl.constexpr, -): - i_t, i_nh = tl.program_id(0), tl.program_id(1) - i_n, i_hq = i_nh // HQ, i_nh % HQ - i_h = i_hq // G - - if IS_VARLEN: - i_n, i_t = tl.load(indices + i_t * 2).to(tl.int32), tl.load(indices + i_t * 2 + 1).to(tl.int32) - bos, eos = tl.load(offsets + i_n).to(tl.int64), tl.load(offsets + i_n + 1).to(tl.int64) - T = (eos - bos).to(tl.int32) - boh = tl.load(chunk_offsets + i_n).to(tl.int32) - else: - bos, eos = (i_n * T).to(tl.int64), (i_n * T + T).to(tl.int64) - NT = tl.cdiv(T, BT) - boh = i_n * NT - - sm_scale = scale * 1.44269504 - # offset calculations - dA_local += (bos*HQ + i_hq) * BT - AT += (bos*H + i_h) * BT - q += (bos*HQ + i_hq) * K - q_new += (bos*HQ + i_hq) * K - k += (bos*H + i_h) * K - k_new += (bos*H + i_h) * K - w += (bos*H + i_h) * K - v += (bos*H + i_h) * V - do += (bos*HQ + i_hq) * V - dv += (bos*HQ + i_hq) * V - beta += (bos*H + i_h) - if RETURN_H: - h += ((boh + i_t) * H + i_h) * K * K - else: - h += (bos*H + i_h) * K - if USE_GATE: - g_cumsum += (bos*HQ + i_hq) - dg_cumsum += (bos*HQ + i_hq) - L += (bos*HQ + i_hq) - D += (bos*HQ + i_hq) - - p_q = tl.make_block_ptr(q, (T, K), (HQ*K, 1), (i_t * BT, 0), (BT, BK), (1, 0)) - p_k = tl.make_block_ptr(k, (K, T), (1, H*K), (0, i_t * BT), (BK, BT), (0, 1)) - p_w = tl.make_block_ptr(w, (T, K), (H*K, 1), (i_t * BT, 0), (BT, BK), (1, 0)) - p_beta = tl.make_block_ptr(beta, (T, ), (H, ), (i_t * BT, ), (BT, ), (0, )) - - b_q = tl.load(p_q, boundary_check=(0, 1)) - b_kt = tl.load(p_k, boundary_check=(0, 1)) - b_w = tl.load(p_w, boundary_check=(0, 1)) - b_beta = tl.load(p_beta, boundary_check=(0, )) - p_T = tl.make_block_ptr(AT, (T, BT), (BT*H, 1), (i_t * BT, 0), (BT, BT), (1, 0)) - b_T = tl.load(p_T, boundary_check=(0, 1)) * b_beta[None, :] - - o_i = tl.arange(0, BT) - m_t = o_i[:, None] >= o_i[None, :] - b_qw = tl.where(m_t, tl.dot(b_q, tl.trans(b_w.to(b_q.dtype))), 0).to(b_q.dtype) - b_qwT = tl.dot(b_qw, b_T.to(b_q.dtype)).to(b_q.dtype) - b_wbk = tl.where(o_i[:, None] > o_i[None, :], tl.dot(b_w.to(b_kt.dtype), b_kt), 0).to(b_q.dtype) - b_A = tl.where(m_t, tl.dot(b_q, b_kt) - tl.dot(b_qwT, b_wbk), 0) - - b_q = b_q.to(tl.float32) - tl.dot(b_qwT, b_w.to(b_qwT.dtype)) - p_q_new = tl.make_block_ptr(q_new, (T, K), (K*HQ, 1), (i_t * BT, 0), (BT, K), (1, 0)) - tl.store(p_q_new, b_q.to(p_q_new.dtype.element_ty), boundary_check=(0, 1)) - - if i_hq % G == 0: - b_Twb = tl.dot(b_T, b_w) # tf32 - p_h = tl.make_block_ptr(h, (T, K), (K * H, 1), (i_t * BT, 0), (BT, BK), (1, 0)) - tl.store(p_h, b_Twb.to(p_h.dtype.element_ty), boundary_check=(0, 1)) - b_T_wbk = tl.dot(b_T.to(b_wbk.dtype), b_wbk).to(b_kt.dtype) - p_k_new = tl.make_block_ptr(k_new, (K, T), (1, K*H), (0, i_t * BT), (BK, BT), (0, 1)) - tl.store(p_k_new, (b_kt - tl.dot(tl.trans(b_w.to(b_kt.dtype)), b_T_wbk) - ).to(p_k_new.dtype.element_ty), boundary_check=(0, 1)) - - if USE_GATE: - p_g_cumsum = tl.make_block_ptr(g_cumsum, (T, ), (HQ, ), (i_t * BT, ), (BT, ), (0, )) - b_g_cumsum = tl.load(p_g_cumsum, boundary_check=(0, )) - b_A = b_A + (b_g_cumsum[:, None] - b_g_cumsum[None, :]) - b_A = tl.where((i_t * BT + tl.arange(0, BT) < T)[:, None], b_A, float("-inf")) # avoid nan - - p_l = tl.make_block_ptr(L, (T, ), (HQ, ), (i_t * BT, ), (BT, ), (0, )) - b_l = tl.load(p_l, boundary_check=(0, )) - p_delta = tl.make_block_ptr(D, (T, ), (HQ, ), (i_t * BT, ), (BT, ), (0, )) - delta = tl.load(p_delta, boundary_check=(0, )) - - b_A_softmax = tl.exp2(tl.where(o_i[:, None] >= o_i[None, :], b_A * sm_scale - b_l[:, None], float("-inf"))) - p_do = tl.make_block_ptr(do, (T, V), (HQ*V, 1), (i_t * BT, 0), (BT, BV), (1, 0)) - b_do = tl.load(p_do, boundary_check=(0, 1)) - b_dv = tl.dot(tl.trans(b_A_softmax.to(b_do.dtype)), b_do) - p_dv = tl.make_block_ptr(dv, (T, V), (HQ*V, 1), (i_t * BT, 0), (BT, BV), (1, 0)) - tl.store(p_dv, b_dv.to(p_dv.dtype.element_ty), boundary_check=(0, 1)) - - p_v = tl.make_block_ptr(v, (V, T), (1, H*V), (0, i_t * BT), (BV, BT), (0, 1)) - b_v = tl.load(p_v, boundary_check=(0, 1)) - b_dp = tl.dot(b_do, b_v) - b_dA = ((b_dp - delta[:, None]) * b_A_softmax * scale) - if USE_GATE: - b_dgq = tl.sum(b_dA, axis=1) - tl.sum(b_dA, axis=0) - p_dg = tl.make_block_ptr(dg_cumsum, (T, ), (HQ, ), (i_t * BT, ), (BT, ), (0, )) - tl.store(p_dg, b_dgq.to(p_dg.dtype.element_ty), boundary_check=(0,)) - p_dA = tl.make_block_ptr(dA_local, (T, BT), (BT*HQ, 1), (i_t * BT, 0), (BT, BT), (1, 0)) - tl.store(p_dA, b_dA.to(p_dA.dtype.element_ty), boundary_check=(0, 1)) - - -def intra_chunk_preprocess_bwd_prepare_fn(q, k, v, w, beta, g_cumsum, A, L, D, do, scale, return_h=True, cu_seqlens=None, - chunk_indices: torch.LongTensor | None = None): - BT = A.shape[-1] - HQ = q.shape[-2] - B, T, H, K = k.shape - G = HQ//H - - V = v.shape[-1] - q_new = torch.empty_like(q) - k_new = torch.empty_like(k) - - if chunk_indices is None and cu_seqlens is not None: - chunk_indices = prepare_chunk_indices(cu_seqlens, BT) - indices = chunk_indices - chunk_offsets = prepare_chunk_offsets(cu_seqlens, BT) if cu_seqlens is not None else None - NT = triton.cdiv(T, BT) if cu_seqlens is None else len(indices) - grid = (NT, B*HQ) - h = torch.empty_like(w) - dA_local = torch.empty(B, T, HQ, BT, dtype=q.dtype, device=q.device) - dv = torch.empty(B, T, HQ, V, device=q.device, dtype=torch.float32) - dg_cumsum = torch.empty_like(g_cumsum) if g_cumsum is not None else None - - chunk_transform_qk_bwd_kernel_prepare[grid]( - q=q, - k=k, - v=v, - w=w, - beta=beta, - g_cumsum=g_cumsum, - AT=A, - dA_local=dA_local, - dv=dv, - dg_cumsum=dg_cumsum, - do=do, - L=L, - D=D, - h=h, - q_new=q_new, - k_new=k_new, - scale=scale, - offsets=cu_seqlens, - indices=indices, - chunk_offsets=chunk_offsets, - T=T, - H=H, - G=G, - HQ=HQ, - K=K, - V=V, - BK=triton.next_power_of_2(K), - BV=triton.next_power_of_2(V), - BT=BT, - RETURN_H=return_h, - ) - return q_new, k_new, h, dA_local, dv, dg_cumsum diff --git a/build/torch-xpu/ops/path_attn/intra_chunk_preprocess_fwd.py b/build/torch-xpu/ops/path_attn/intra_chunk_preprocess_fwd.py deleted file mode 100644 index c1f49ecf64cf23f2f3ebe9454e3daa77bf11c1d8..0000000000000000000000000000000000000000 --- a/build/torch-xpu/ops/path_attn/intra_chunk_preprocess_fwd.py +++ /dev/null @@ -1,179 +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 -import triton.language as tl - -from ...ops.utils import prepare_chunk_indices - - -@triton.heuristics({ - "USE_G": lambda args: args['g_cumsum'] is not None, - "IS_VARLEN": lambda args: args['offsets'] is not None, -}) -@triton.jit(do_not_specialize=['T']) -def intra_chunk_preprocess_fwd_kernel( - q, - k, - v, - w, - beta, - g_cumsum, - o, - A, - L, - M, - w2, - q_new, - k_new, - scale, - indices, # varlen helper - offsets, # varlen helper - T, - H: tl.constexpr, - G: tl.constexpr, - HQ: tl.constexpr, - K: tl.constexpr, - V: tl.constexpr, - BK: tl.constexpr, - BV: tl.constexpr, - BT: tl.constexpr, - IS_VARLEN: tl.constexpr, - USE_G: tl.constexpr, -): - i_t, i_nh = tl.program_id(0), tl.program_id(1) - i_n, i_hq = i_nh // HQ, i_nh % HQ - i_h = i_hq // G - - if IS_VARLEN: - i_n, i_t = tl.load(indices + i_t * 2).to(tl.int32), tl.load(indices + i_t * 2 + 1).to(tl.int32) - bos, eos = tl.load(offsets + i_n).to(tl.int64), tl.load(offsets + i_n + 1).to(tl.int64) - T = (eos - bos).to(tl.int32) - else: - bos, eos = (i_n * T).to(tl.int64), (i_n * T + T).to(tl.int64) - - sm_scale = scale * 1.44269504 - # offset calculations - A += (bos*H + i_h) * BT - q += (bos*HQ + i_hq) * K - q_new += (bos*HQ + i_hq) * K - k += (bos*H + i_h) * K - k_new += (bos*H + i_h) * K - w2 += (bos*H + i_h) * K - w += (bos*H + i_h) * K - v += (bos*H + i_h) * V - o += (bos*HQ + i_hq) * V - beta += (bos*H + i_h) - if USE_G: - g_cumsum += (bos*HQ + i_hq) - L += (bos*HQ + i_hq) - M += (bos*HQ + i_hq) - - p_q = tl.make_block_ptr(q, (T, K), (HQ*K, 1), (i_t * BT, 0), (BT, BK), (1, 0)) - p_k = tl.make_block_ptr(k, (K, T), (1, H*K), (0, i_t * BT), (BK, BT), (0, 1)) - p_w = tl.make_block_ptr(w, (T, K), (H*K, 1), (i_t * BT, 0), (BT, BK), (1, 0)) - p_v = tl.make_block_ptr(v, (T, V), (H*V, 1), (i_t * BT, 0), (BT, BV), (1, 0)) - p_beta = tl.make_block_ptr(beta, (T, ), (H, ), (i_t * BT, ), (BT, ), (0, )) - p_T = tl.make_block_ptr(A, (T, BT), (BT*H, 1), (i_t * BT, 0), (BT, BT), (1, 0)) - - b_beta = tl.load(p_beta, boundary_check=(0, )) - b_q = tl.load(p_q, boundary_check=(0, 1)) - b_kt = tl.load(p_k, boundary_check=(0, 1)) - b_v = tl.load(p_v, boundary_check=(0, 1)) - b_w = tl.load(p_w, boundary_check=(0, 1)) - b_T = tl.load(p_T, boundary_check=(0, 1)) - b_T = b_T * b_beta[None, :] - - o_i = tl.arange(0, BT) - m_t = o_i[:, None] >= o_i[None, :] - - b_qw = tl.where(m_t, tl.dot(b_q, tl.trans(b_w.to(b_q.dtype))), 0).to(b_q.dtype) - b_qwT = tl.dot(b_qw, b_T.to(b_q.dtype)).to(b_q.dtype) - b_wbk = tl.where(o_i[:, None] > o_i[None, :], tl.dot(b_w.to(b_q.dtype), b_kt), 0).to(b_q.dtype) - b_A = tl.where(m_t, tl.dot(b_q, b_kt) - tl.dot(b_qwT.to(b_q.dtype), b_wbk), 0) - - b_q = b_q.to(tl.float32) - tl.dot(b_qwT, b_w.to(b_q.dtype)) - p_q_new = tl.make_block_ptr(q_new, (T, K), (K*HQ, 1), (i_t * BT, 0), (BT, K), (1, 0)) - tl.store(p_q_new, b_q.to(p_q_new.dtype.element_ty), boundary_check=(0, 1)) - - if i_hq % G == 0: - b_Twb = tl.dot(b_T, b_w) - p_w2 = tl.make_block_ptr(w2, (T, K), (K*H, 1), (i_t * BT, 0), (BT, BK), (1, 0)) - tl.store(p_w2, b_Twb.to(p_w2.dtype.element_ty), boundary_check=(0, 1)) - b_T_wbk = tl.dot(b_T.to(b_kt.dtype), b_wbk).to(b_kt.dtype) - p_k_new = tl.make_block_ptr(k_new, (K, T), (1, K*H), (0, i_t * BT), (BK, BT), (0, 1)) - tl.store(p_k_new, (b_kt - tl.dot(tl.trans(b_w.to(b_kt.dtype)), b_T_wbk) - ).to(p_k_new.dtype.element_ty), boundary_check=(0, 1)) - - if USE_G: - p_g_cumsum = tl.make_block_ptr(g_cumsum, (T, ), (HQ, ), (i_t * BT, ), (BT, ), (0, )) - b_g_cumsum = tl.load(p_g_cumsum, boundary_check=(0, )) - b_A = b_A + (b_g_cumsum[:, None] - b_g_cumsum[None, :]) - b_A = tl.where((i_t * BT + tl.arange(0, BT) < T)[:, None], b_A, float("-inf")) # avoid nan - - b_qkT_softmax = tl.where(o_i[:, None] >= o_i[None, :], b_A * sm_scale, float("-inf")) - m_i = tl.max(b_qkT_softmax, 1) - b_qkT_softmax = tl.math.exp2(b_qkT_softmax - m_i[:, None]) - l_i = tl.sum(b_qkT_softmax, 1) - b_o = tl.dot(b_qkT_softmax.to(b_v.dtype), b_v) - p_o = tl.make_block_ptr(o, (T, V), (V*HQ, 1), (i_t * BT, 0), (BT, BV), (1, 0)) - tl.store(p_o, b_o.to(p_o.dtype.element_ty), boundary_check=(0, 1)) - p_l = tl.make_block_ptr(L, (T, ), (HQ, ), (i_t * BT, ), (BT, ), (0, )) - p_m = tl.make_block_ptr(M, (T, ), (HQ, ), (i_t * BT, ), (BT, ), (0, )) - tl.store(p_m, m_i.to(p_m.dtype.element_ty), boundary_check=(0,)) - tl.store(p_l, l_i.to(p_l.dtype.element_ty), boundary_check=(0,)) - - -def intra_chunk_preprocess_fwd_fn(q, k, v, w, beta, g_cumsum, A, scale, BT, cu_seqlens, - chunk_indices: torch.LongTensor | None = None): - HQ = q.shape[-2] - B, T, H, K = k.shape - V = v.shape[-1] - q_new = torch.empty_like(q, dtype=torch.float32) # for stability - k_new = torch.empty_like(k) - o = torch.empty(B, T, HQ, V, device=q.device, dtype=torch.float32) - - if chunk_indices is None and cu_seqlens is not None: - chunk_indices = prepare_chunk_indices(cu_seqlens, BT) - indices = chunk_indices - NT = triton.cdiv(T, BT) if cu_seqlens is None else len(indices) - grid = (NT, B*HQ) - L = torch.empty(B, T, HQ, dtype=torch.float32, device=q.device) - M = torch.empty(B, T, HQ, dtype=torch.float32, device=q.device) - w2 = torch.empty_like(w) - G = HQ//H - - intra_chunk_preprocess_fwd_kernel[grid]( - q=q, - k=k, - v=v, - w=w, - beta=beta, - g_cumsum=g_cumsum, - o=o, - A=A, - L=L, - M=M, - w2=w2, - q_new=q_new, - k_new=k_new, - scale=scale, - offsets=cu_seqlens, - indices=indices, - T=T, - H=H, - G=G, - HQ=HQ, - K=K, - V=V, - BK=triton.next_power_of_2(K), - BV=triton.next_power_of_2(V), - BT=BT, - num_warps=4 if BT == 64 else 2, - ) - return q_new, k_new, w2, o, L, M diff --git a/build/torch-xpu/ops/path_attn/naive.py b/build/torch-xpu/ops/path_attn/naive.py deleted file mode 100644 index 9607dc6f0c75f428ff9cf0b17e831fc8a6bc02db..0000000000000000000000000000000000000000 --- a/build/torch-xpu/ops/path_attn/naive.py +++ /dev/null @@ -1,97 +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 torch.nn.functional as F -from einops import rearrange - - -def naive_path_attn( - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - w: torch.Tensor, - beta: torch.Tensor, - g: torch.Tensor, - scale: float, - chunk_size: int = 64, -): - """ - Reference PyTorch implementation of path attention. - - Args: - q: [B, T, HQ, D] - k: [B, T, H, D] - v: [B, T, H, D] - w: [B, T, H, D] - beta: [B, T, H] - g: [B, T, HQ] - scale: float - chunk_size: int, default 64 - - Returns: - output: [B, T, HQ, D] - """ - original_dtype = q.dtype - HQ = q.shape[2] - H = k.shape[2] - BT = chunk_size - - q, k, v, w, beta, g = map(lambda x: x.to(torch.float).transpose(1, 2), [q, k, v, w, beta, g]) - g_cumsum = g.cumsum(-1) - - q = q.unsqueeze(2).expand(-1, -1, HQ // HQ, -1, -1).flatten(1, 2) - k = k.unsqueeze(2).expand(-1, -1, HQ // H, -1, -1).flatten(1, 2) - v = v.unsqueeze(2).expand(-1, -1, HQ // H, -1, -1).flatten(1, 2) - w = w.unsqueeze(2).expand(-1, -1, HQ // H, -1, -1).flatten(1, 2) - beta = beta.unsqueeze(2).expand(-1, -1, HQ // H, -1).flatten(1, 2) - - b, h, l, _ = q.shape - if l % BT != 0: - padding_size = BT - l % BT - q, k, w = map(lambda x: F.pad(x, (0, 0, 0, padding_size)), [q, k, w]) - beta = F.pad(beta, (0, padding_size)) - - seq_len = q.shape[2] - w_beta = w * beta[..., None] - q, k, w, w_beta = map(lambda x: rearrange(x, 'b h (n c) d -> b h n c d', c=BT), [q, k, w, w_beta]) - - mask = torch.triu(torch.ones(BT, BT, dtype=torch.bool, device=q.device), diagonal=0) - T_mat = -(w_beta @ w.transpose(-1, -2)).masked_fill(mask, 0) - - for i in range(1, BT): - T_mat[..., i, :i] = T_mat[..., i, :i].clone() + (T_mat[..., i, :, None].clone() * T_mat[..., :, :i].clone()).sum(-2) - - T_mat = T_mat + torch.eye(BT, dtype=q.dtype, device=q.device) - Twbk = T_mat @ (w_beta @ k.transpose(-1, -2)).masked_fill(mask, 0) - qw = (q @ w.transpose(-1, -2)).tril() - Twb = T_mat @ w_beta - A_local = (q @ k.transpose(-1, -2)).tril() - qw @ Twbk - q = q - qw @ Twb - k = k - Twbk.transpose(-1, -2) @ w - H_mat = w.transpose(-1, -2) @ Twb - - A = torch.zeros(b, h, seq_len, seq_len, device=q.device) - q, k, w, w_beta = map(lambda x: rearrange(x, 'b h n c d -> b h (n c) d'), [q, k, w, w_beta]) - - for i in range(0, seq_len, BT): - q_i = q[:, :, i:i+BT].clone() - for j in range(i - BT, -BT, -BT): - k_j = k[:, :, j:j+BT] - A_ij = q_i @ k_j.transpose(-1, -2) - A[:, :, i:i+BT, j:j+BT] = A_ij - q_i = q_i - q_i @ H_mat[:, :, j // BT] - - for i in range(0, seq_len // BT): - A[:, :, i*BT:i*BT+BT, i*BT:i*BT+BT] = A_local[:, :, i] - - A = A.masked_fill_(~torch.tril(torch.ones(seq_len, seq_len, device=q.device, dtype=torch.bool)), float("-inf")) - A = A[:, :, :l, :l] - A = A + g_cumsum[..., None] - g_cumsum[..., None, :] - ref_o = (A * scale).softmax(-1).to(v) @ v - - return ref_o.to(original_dtype).transpose(1, 2) diff --git a/build/torch-xpu/ops/path_attn/parallel.py b/build/torch-xpu/ops/path_attn/parallel.py deleted file mode 100644 index 71f015b5a750784c893ac43f7688911f025f9ed6..0000000000000000000000000000000000000000 --- a/build/torch-xpu/ops/path_attn/parallel.py +++ /dev/null @@ -1,287 +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 -from einops import reduce - -from ...ops.attn.parallel import parallel_attn_bwd_preprocess -from ...ops.common.chunk_scaled_dot_kkt import chunk_scaled_dot_kkt_fwd -from ...ops.path_attn.cumprod_householder_bwd import chunk_cumprod_householder_bwd_fn -from ...ops.path_attn.cumprod_householder_fwd import chunk_cumprod_householder_fwd_fn -from ...ops.path_attn.intra_chunk_preprocess_bwd import intra_chunk_preprocess_bwd_fn -from ...ops.path_attn.intra_chunk_preprocess_bwd_prepare import intra_chunk_preprocess_bwd_prepare_fn -from ...ops.path_attn.intra_chunk_preprocess_fwd import intra_chunk_preprocess_fwd_fn -from ...ops.path_attn.parallel_path_bwd_inter_dkv import parallel_path_bwd_dkv_fn -from ...ops.path_attn.parallel_path_bwd_inter_dqh import parallel_path_bwd_dq_fn -from ...ops.path_attn.parallel_path_bwd_intra import parallel_path_bwd_intra_chunk_fn -from ...ops.path_attn.parallel_path_fwd import parallel_path_fwd_fn -from ...ops.path_attn.prepare_k_cache import prepare_k_cache_fn -from ...ops.path_attn.transform_q import transform_q_fwd_fn -from ...ops.utils.cumsum import chunk_global_cumsum -from ...ops.utils.solve_tril import solve_tril -from ...utils import autocast_custom_bwd, autocast_custom_fwd, check_shared_mem, input_guard - - -class ParallelPATHAttentionFunction(torch.autograd.Function): - @staticmethod - @input_guard - @autocast_custom_fwd - def forward(ctx, q, k, v, w, beta, g, scale, cu_seqlens, use_cache=False): - - g_cumsum = chunk_global_cumsum(g, cu_seqlens=cu_seqlens, output_dtype=torch.float32) if g is not None else None - BS = 64 if check_shared_mem('hopper') else 32 - BT = 128 if check_shared_mem('ampere') else 64 - - A = chunk_scaled_dot_kkt_fwd( - k=w, - beta=beta, - cu_seqlens=cu_seqlens, - chunk_size=BS, - output_dtype=torch.float32, - ) - A = solve_tril( - A=A, - cu_seqlens=cu_seqlens, - output_dtype=w.dtype, # force fp32? - ) - q_new, k_new, w2, o, L, M = intra_chunk_preprocess_fwd_fn( - q=q, - k=k, - v=v, - w=w, - beta=beta, - g_cumsum=g_cumsum, - A=A, - scale=scale, - BT=BS, - cu_seqlens=cu_seqlens, - ) - w_fp16 = w.to(torch.float16) - w2_fp16 = w2.to(torch.float16) - o, L = parallel_path_fwd_fn( - q=q_new, - k=k_new, - v=v, - L=L, - w1=w_fp16, - w2=w2_fp16, - M=M, - o=o, - g_cumsum=g_cumsum, - scale=scale, - cu_seqlens=cu_seqlens, - BT=BT, - BS=BS, - ) - k_cache = prepare_k_cache_fn(k=k_new, w1=w, w2=w2, cu_seqlens=cu_seqlens, BS=BS, use_cache=use_cache) - ctx.save_for_backward(q, k, v, w, g_cumsum, o, beta, L, A) - ctx.scale = scale - ctx.cu_seqlens = cu_seqlens - return o, k_cache - - @staticmethod - @input_guard - @autocast_custom_bwd - def backward(ctx, do, dk_new): - q, k, v, w, g_cumsum, o, beta, L, A = ctx.saved_tensors - BT = 128 if check_shared_mem('ampere') else 64 - BS = 64 if check_shared_mem('hopper') else 32 - S = 512 - cu_seqlens = ctx.cu_seqlens - delta = parallel_attn_bwd_preprocess(o, do) - - q_new, k_new, h, dA_local, dv, dg_cumsum = intra_chunk_preprocess_bwd_prepare_fn( - q=q, - k=k, - v=v, - w=w, - beta=beta, - g_cumsum=g_cumsum, - A=A, - L=L, - D=delta, - do=do, - scale=ctx.scale, - cu_seqlens=cu_seqlens, - return_h=False, - ) - w_fp16 = w.to(torch.float16) - h_fp16 = h.to(torch.float16) - k_new_large, hc_suffix, hc_whole = chunk_cumprod_householder_fwd_fn( - k=k_new, - w1=w_fp16, - w2=h_fp16, - S=S, - BT=BS, - cu_seqlens=cu_seqlens, - ) - q_new_large = transform_q_fwd_fn(q=q_new, w1=w_fp16, w2=h_fp16, cu_seqlens=cu_seqlens, BT=BT, BS=BS, S=S) - w = w.to(q.dtype) - h = h.to(q.dtype) - A = A.to(q.dtype) - dk, dv, _ = parallel_path_bwd_dkv_fn( - q=q_new_large, - k=k_new_large, - v=v, - g_cumsum=g_cumsum, - do=do, - dv=dv, - dg_cumsum=dg_cumsum, - hc_whole=hc_whole, - scale=ctx.scale, - cu_seqlens=cu_seqlens, - L=L, - D=delta, - S=S, - BT=BT, - BS=BS, - ) - dq, dhc_whole, dg_cumsum = parallel_path_bwd_dq_fn( - q=q_new_large, - k=k_new_large, - v=v, - g_cumsum=g_cumsum, - do=do, - dg_cumsum=dg_cumsum, - hc_whole=hc_whole, - scale=ctx.scale, - cu_seqlens=cu_seqlens, - L=L, - D=delta, - S=S, - BT=BT, - BS=BS, - ) - dw1, dw2, dk = chunk_cumprod_householder_bwd_fn( - w1=w, - w2=h, - k=k_new, - dk=dk, - hc_suffix=hc_suffix, - dhc_whole=dhc_whole, - cu_seqlens=cu_seqlens, - S=S, - BT=BS, - ) - dq, dk, dv, dw1, dw2, dg_cumsum = parallel_path_bwd_intra_chunk_fn( - q=q_new, - k=k_new, - v=v, - g_cumsum=g_cumsum, - w1=w, - w2=h, - L=L, - D=delta, - scale=ctx.scale, - dw1=dw1, - dw2=dw2, - dq=dq, - dk=dk, - dv=dv, - do=do, - dg_cumsum=dg_cumsum, - cu_seqlens=cu_seqlens, - S=S, - BT=BS, - ) - dq, dk, dbeta, dw = intra_chunk_preprocess_bwd_fn( - q=q, - k=k, - w=w, - w2=h, - beta=beta, - dq=dq, - dk=dk, - dw1=dw1, - dw2=dw2, - dA_local=dA_local, - A=A, - L=L, - D=delta, - do=do, - scale=ctx.scale, - cu_seqlens=cu_seqlens, - ) - G = q.shape[-2] // k.shape[-2] - if G > 1: - assert dk.dtype == dv.dtype == dw.dtype == dbeta.dtype == torch.float32, 'reduction requires float32' - dk = reduce(dk, 'b t (h g) k -> b t h k', g=G, reduction='sum') - dv = reduce(dv, 'b t (h g) k -> b t h k', g=G, reduction='sum') - dw = reduce(dw, 'b t (h g) k -> b t h k', g=G, reduction='sum') - dbeta = reduce(dbeta, 'b t (h g) -> b t h', g=G, reduction='sum') - if dg_cumsum is not None: - dg_cumsum = chunk_global_cumsum(dg_cumsum, cu_seqlens=cu_seqlens, reverse=True) - return (dq.to(q.dtype), dk.to(k.dtype), dv.to(v.dtype), dw.to(w.dtype), - dbeta.to(beta.dtype), - dg_cumsum.to(g_cumsum.dtype) if g_cumsum is not None else None, - None, None, None, None) - - -@torch.compiler.disable -def parallel_path_attn( - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - w: torch.Tensor, - beta: torch.Tensor, - g: torch.Tensor | None = None, - scale: float = None, - cu_seqlens: torch.Tensor | None = None, - use_cache: bool = False, -) -> tuple[torch.Tensor, torch.Tensor]: - r""" - Args: - q (torch.Tensor): - queries of shape `[B, T, HQ, K]` - k (torch.Tensor): - keys of shape `[B, T, H, K]` - v (torch.Tensor): - values of shape `[B, T, H, V]` - w (torch.Tensor): - weights of shape `[B, T, H, K]` - beta (torch.Tensor): - beta of shape `[B, T, H]` - g (torch.Tensor): - g of shape `[B, T, HQ]` - scale (float): - Scale factor for attention scores. - If not provided, it will default to `1 / sqrt(K)`. Default: `None`. - cu_seqlens (torch.LongTensor): - Cumulative sequence lengths of shape `[N+1]` used for variable-length training, - consistent with the FlashAttention API. - use_cache (bool): - Whether to transform and cache the key values for decoding. Default: `False`. - - Returns: - o (torch.Tensor): - output of shape `[B, T, HQ, V]` - k_cache (torch.Tensor): - k_cache of shape `[B, T, H, K]` - """ - if scale is None: - scale = k.shape[-1]**-0.5 - if cu_seqlens is not None and q.shape[0] != 1: - raise ValueError( - f"The batch size is expected to be 1 rather than {q.shape[0]} when using `cu_seqlens`. " - f"Please flatten variable-length inputs before processing.", - ) - assert w.dtype == beta.dtype == torch.float32, 'w, beta should be float32 to preserve precision.' - if g is not None: - assert g.dtype == torch.float32, 'g should be float32 to preserve precision.' - assert q.shape[-1] in [16, 32, 64, 128], "only support head_dim in [16, 32, 64, 128] for now. Stay tuned!" - assert v.shape[-1] in [16, 32, 64, 128], "only support head_dim in [16, 32, 64, 128] for now. Stay tuned!" - assert q.shape[-1] == k.shape[-1], 'q, k should have the same head_dim.' - assert k.shape == w.shape, 'k, w should have the same shape.' - assert beta.shape[:3] == k.shape[:3], 'beta should have the same number of heads as k' - if g is not None: - assert g.shape[:3] == q.shape[:3], 'g should have the same number of heads as q' - assert q.shape[-2] % k.shape[-2] == 0, 'the number of query heads should be divisible by the number of key heads' - o, k_cache = ParallelPATHAttentionFunction.apply(q, k, v, w, beta, g, scale, cu_seqlens, use_cache) - return o, k_cache - - -parallel_path_attention = parallel_path_attn diff --git a/build/torch-xpu/ops/path_attn/parallel_path_bwd_inter_dkv.py b/build/torch-xpu/ops/path_attn/parallel_path_bwd_inter_dkv.py deleted file mode 100644 index af070c93e13417da7292dfc94ffb4602155255d4..0000000000000000000000000000000000000000 --- a/build/torch-xpu/ops/path_attn/parallel_path_bwd_inter_dkv.py +++ /dev/null @@ -1,199 +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 -import triton.language as tl - -from ...ops.utils import prepare_chunk_indices, prepare_chunk_offsets - - -@triton.heuristics({ - 'USE_GATE': lambda args: args['g_cumsum'] is not None, - 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None, -}) -@triton.jit(do_not_specialize=['T']) -def parallel_path_bwd_dkv_kernel( - q, - k, - v, - g_cumsum, - hc_whole, - scale, - L, - D, - dk, - dv, - do, - dg_cumsum, - cu_seqlens, - indices, - split_offsets, - T, - G: tl.constexpr, - HQ: tl.constexpr, - H: tl.constexpr, - K: tl.constexpr, - V: tl.constexpr, - BT: tl.constexpr, - BS: tl.constexpr, - BK: tl.constexpr, - BV: tl.constexpr, - S: tl.constexpr, - IS_VARLEN: tl.constexpr, - USE_GATE: tl.constexpr, - NUM_BLOCKS: tl.constexpr, -): - i_t, i_bh = tl.program_id(0), tl.program_id(1) - i_b, i_hq = i_bh // HQ, i_bh % HQ - i_h = i_hq // G - - if IS_VARLEN: - i_n, i_t = tl.load(indices + i_t * 2).to(tl.int32), tl.load(indices + i_t * 2 + 1).to(tl.int32) - boh_large = tl.load(split_offsets + i_n).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).to(tl.int32) - else: - i_n = i_b - bos, eos = (i_n * T).to(tl.int64), (i_n * T + T).to(tl.int64) - boh_large = i_n * tl.cdiv(T, S) - - # offset calculations - - do += (bos * HQ + i_hq) * V - dk += (bos * HQ + i_hq) * K - dv += (bos * HQ + i_hq) * K - L += (bos * HQ + i_hq) - D += (bos * HQ + i_hq) - - k += (bos * H + i_h) * K # GQA when H!=HQ - v += (bos * H + i_h) * V # GQA when H!=HQ - hc_whole += (boh_large * H + i_h) * K * K - - if USE_GATE: - g_cumsum += (bos * HQ + i_hq) - dg_cumsum += (bos * HQ + i_hq) - - # constants - sm_scale = scale * 1.44269504 - - # load query - p_k = tl.make_block_ptr(k, (T, K), (H*K, 1), (i_t * BT, 0), (BT, BK), (1, 0)) - b_k = tl.load(p_k, boundary_check=(0, 1)) - p_v = tl.make_block_ptr(v, (T, V), (H*V, 1), (i_t * BT, 0), (BT, BV), (1, 0)) - b_v = tl.load(p_v, boundary_check=(0, 1)) - - if USE_GATE: - b_g_cumsum_k = tl.zeros([BT], dtype=tl.float32) - p_g_cumsum_k = tl.make_block_ptr(g_cumsum, (T, ), (HQ, ), (i_t * BT, ), (BT, ), (0, )) - b_g_cumsum_k += tl.load(p_g_cumsum_k, boundary_check=(0, )) - b_dg_cumsum_k = tl.zeros([BT], dtype=tl.float32) - else: - b_g_cumsum_k = None - b_dg_cumsum_k = None - - b_dk = tl.zeros([BT, BK], dtype=tl.float32) - b_dv = tl.zeros([BT, BV], dtype=tl.float32) - - last_chunk_start = tl.floor(i_t*BT / S).to(tl.int32) * S - idx_j = (tl.floor(i_t * BT / S).to(tl.int32) + 1).to(tl.int32) - - last_chunk_end = tl.ceil(T / BS).to(tl.int32) * BS - BS - - for offset in range(last_chunk_end, last_chunk_start+S-BS, -BS): - p_delta = tl.make_block_ptr(D, (T, ), (HQ, ), (offset, ), (BS, ), (0, )) - p_l = tl.make_block_ptr(L, (T, ), (HQ, ), (offset, ), (BS, ), (0, )) - b_delta = tl.load(p_delta, boundary_check=(0, )) - b_l = tl.load(p_l, boundary_check=(0, )) - - p_q = tl.make_block_ptr(q + ((bos.to(tl.int64) * NUM_BLOCKS + idx_j) * HQ + i_hq) * K, (T, K), - (HQ*K*NUM_BLOCKS, 1), (offset, 0), (BS, BK), (1, 0)) - b_q = tl.load(p_q, boundary_check=(0, 1)) - b_A = tl.dot(b_k, tl.trans(b_q).to(b_k.dtype)) - if USE_GATE: - p_g_cumsum_q = tl.make_block_ptr(g_cumsum, (T, ), (HQ, ), (offset, ), (BS, ), (0, )) - b_g_cumsum_q = tl.load(p_g_cumsum_q, boundary_check=(0, )) - b_A = b_A + b_g_cumsum_q[None, :] - b_g_cumsum_k[:, None] - b_A = tl.where((offset + tl.arange(0, BS) < T)[None, :], b_A, float("-inf")) # avoid nan - b_A_softmax = tl.math.exp2(b_A * sm_scale - b_l[None, :]) - p_do = tl.make_block_ptr(do, (T, V), (HQ*V, 1), (offset, 0), (BS, BV), (1, 0)) - b_do = tl.load(p_do, boundary_check=(0, 1)) - b_dv += tl.dot(b_A_softmax.to(b_do.dtype), b_do) - b_dp = tl.dot(b_v, tl.trans(b_do)) - - b_dA = ((b_dp - b_delta[None, :]) * b_A_softmax * scale) - if USE_GATE: - b_dg_cumsum_k -= tl.sum(b_dA, axis=1) - b_dk += tl.dot(b_dA.to(b_q.dtype), b_q) - - p_dk = tl.make_block_ptr(dk, (T, K), (HQ*K, 1), (i_t * BT, 0), (BT, BK), (1, 0)) - tl.store(p_dk, b_dk.to(dk.dtype.element_ty), boundary_check=(0, 1)) - mask = i_t * BT + tl.arange(0, BT) < T - tl.atomic_add( - dv + (i_t * BT + tl.arange(0, BT))[:, None] * HQ * V + tl.arange(0, BV)[None, :], - b_dv, - mask=mask[:, None], - sem='relaxed', - ) - if USE_GATE: - tl.atomic_add(dg_cumsum + (i_t * BT + tl.arange(0, BT)) * HQ, b_dg_cumsum_k, mask=mask, sem='relaxed') - - -def parallel_path_bwd_dkv_fn( - q, k, v, g_cumsum, do, dv, dg_cumsum, - hc_whole, scale, L, D, - cu_seqlens, - S, BT, BS, - chunk_indices: torch.LongTensor | None = None, -): - B, T, num_blocks, HQ, K = q.shape - V = v.shape[-1] - H = k.shape[-2] - G = HQ // H - - if chunk_indices is None and cu_seqlens is not None: - chunk_indices = prepare_chunk_indices(cu_seqlens, BT) - indices = chunk_indices - split_offsets = prepare_chunk_offsets(cu_seqlens, S) if cu_seqlens is not None else None - NT = triton.cdiv(T, BT) if cu_seqlens is None else len(indices) - - if cu_seqlens is not None: - assert split_offsets[-1] == hc_whole.shape[0] - - dk = torch.empty(B, T, HQ, K, dtype=torch.float32, device=q.device) - - parallel_path_bwd_dkv_kernel[(NT, B*HQ)]( - q=q, - k=k, - v=v, - g_cumsum=g_cumsum, - hc_whole=hc_whole, - scale=scale, - L=L, - D=D, - dk=dk, - dv=dv, - do=do, - dg_cumsum=dg_cumsum, - cu_seqlens=cu_seqlens, - indices=indices, - split_offsets=split_offsets, - T=T, - S=S, - BT=BT, - BS=BS, - G=G, - HQ=HQ, - H=H, - K=K, - V=V, - BK=triton.next_power_of_2(K), - BV=triton.next_power_of_2(V), - num_warps=8 if (BT == 128 and K == 128) else 4, - NUM_BLOCKS=num_blocks, - ) - return dk, dv, dg_cumsum diff --git a/build/torch-xpu/ops/path_attn/parallel_path_bwd_inter_dqh.py b/build/torch-xpu/ops/path_attn/parallel_path_bwd_inter_dqh.py deleted file mode 100644 index 81c6096840c3990600ffd5ee1e677c442765e974..0000000000000000000000000000000000000000 --- a/build/torch-xpu/ops/path_attn/parallel_path_bwd_inter_dqh.py +++ /dev/null @@ -1,210 +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 -import triton.language as tl - -from ...ops.utils import prepare_chunk_indices, prepare_chunk_offsets -from ...ops.utils.op import exp2 -from ...utils import check_shared_mem - - -@triton.heuristics({ - 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None, - 'USE_GATE': lambda args: args['g_cumsum'] is not None, -}) -@triton.jit(do_not_specialize=['T']) -def parallel_path_bwd_dq_kernel( - q, - k, - v, - g_cumsum, - hc_whole, - scale, - L, - D, - dq, - do, - dhc_whole, - dg_cumsum, - cu_seqlens, - indices, - split_offsets, # varlen specific - T, - G: tl.constexpr, - HQ: tl.constexpr, - H: tl.constexpr, - K: tl.constexpr, - V: tl.constexpr, - BT: tl.constexpr, - BS: tl.constexpr, - BK: tl.constexpr, - BV: tl.constexpr, - S: tl.constexpr, # aka larger chunk size - NUM_BLOCKS: tl.constexpr, - IS_VARLEN: tl.constexpr, - USE_GATE: tl.constexpr, -): - i_t, i_nh = tl.program_id(0), tl.program_id(1) - i_n, i_hq = i_nh // HQ, i_nh % HQ - i_h = i_hq // G - - if IS_VARLEN: - i_n, i_t = tl.load(indices + i_t * 2).to(tl.int32), tl.load(indices + i_t * 2 + 1).to(tl.int32) - boh_large = tl.load(split_offsets + i_n).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).to(tl.int32) - else: - bos, eos = (i_n * T).to(tl.int64), (i_n * T + T).to(tl.int64) - boh_large = i_n * tl.cdiv(T, S) - o_t = i_t * BT + tl.arange(0, BT) - m_t = o_t < T - - k += (bos * H + i_h) * K # GQA when H!=HQ - v += (bos * H + i_h) * V # GQA when H!=HQ - do += (bos * HQ + i_hq) * V - dq += (bos * HQ + i_hq) * K - hc_whole += (boh_large * H + i_h) * K * K - dhc_whole += (boh_large * HQ + i_hq) * K * K - L += (bos * HQ + i_hq) - D += (bos * HQ + i_hq) - if USE_GATE: - g_cumsum += (bos * HQ + i_hq) - dg_cumsum += (bos * HQ + i_hq) - - # constants - stride_h = H * K * K - stride_hq = HQ * K * K - sm_scale = scale * 1.44269504 - - # load query - p_do = tl.make_block_ptr(do, (T, V), (HQ*V, 1), (i_t * BT, 0), (BT, BV), (1, 0)) - b_do = tl.load(p_do, boundary_check=(0, 1)) - - p_l = tl.make_block_ptr(L, (T,), (HQ,), (i_t * BT,), (BT,), (0,)) - p_d = tl.make_block_ptr(D, (T,), (HQ,), (i_t * BT,), (BT,), (0,)) - b_l = tl.load(p_l, boundary_check=(0,)) - b_delta = tl.load(p_d, boundary_check=(0,)) - - if USE_GATE: - p_g_cumsum_q = tl.make_block_ptr(g_cumsum, (T,), (HQ,), (i_t * BT,), (BT,), (0,)) - b_g_cumsum_q = tl.load(p_g_cumsum_q, boundary_check=(0,)).to(tl.float32) - b_dg_cumsum_q = tl.zeros([BT], dtype=tl.float32) - else: - b_g_cumsum_q = None - b_dg_cumsum_q = None - - curr_end = ((i_t * BT // S) * S).to(tl.int32) - b_dq = tl.zeros([BT, K], dtype=tl.float32) - - for offset_outer in range(0, curr_end, S): - idx_j = offset_outer // S - p_q = tl.make_block_ptr(q + ((bos.to(tl.int64) * NUM_BLOCKS + idx_j + 1) * HQ + i_hq) * K, (T, K), - (HQ*K*NUM_BLOCKS, 1), (i_t * BT, 0), (BT, BK), (1, 0)) - b_q = tl.load(p_q, boundary_check=(0, 1)) - - b_dh = -tl.dot(tl.trans(b_q), b_dq.to(b_q.dtype)) - tl.atomic_add(dhc_whole + idx_j * stride_hq + tl.arange(0, K) - [:, None] * K + tl.arange(0, K)[None, :], b_dh, sem='relaxed') - p_h = tl.make_block_ptr(hc_whole + idx_j * stride_h, (K, K), (K, 1), (0, 0), (BK, BK), (1, 0)) - b_h = tl.load(p_h, boundary_check=(0, 1)) - b_dq = b_dq - tl.dot(b_dq.to(b_h.dtype), tl.trans(b_h)) - - for offset in range(offset_outer, min(offset_outer+S, i_t*BT), BS): - p_k = tl.make_block_ptr(k, (T, K), (H * K, 1), (offset, 0), (BS, BK), (1, 0)) - b_k = tl.load(p_k, boundary_check=(0, 1)) - b_A = tl.dot(b_q, tl.trans(b_k).to(b_q.dtype)) - if USE_GATE: - p_g_cumsum_k = tl.make_block_ptr(g_cumsum, (T,), (HQ,), (offset,), (BS,), (0,)) - b_g_cumsum_k = tl.load(p_g_cumsum_k, boundary_check=(0,)).to(tl.float32) - b_A = b_A + b_g_cumsum_q[:, None] - b_g_cumsum_k[None, :] - b_A = exp2(b_A * sm_scale - b_l[:, None]) - b_A = tl.where(m_t[:, None], b_A, 0) - p_v = tl.make_block_ptr(v, (T, V), (H*V, 1), (offset, 0), (BS, BV), (1, 0)) - b_v = tl.load(p_v, boundary_check=(0, 1)) - b_dp = tl.dot(b_do, tl.trans(b_v).to(b_do.dtype)) - b_dA = (b_dp - b_delta[:, None]) * b_A * scale - b_dq += tl.dot(b_dA.to(b_k.dtype), b_k) - if USE_GATE: - b_dg_cumsum_q += tl.sum(b_dA, axis=1) - - p_dq = tl.make_block_ptr(dq, (T, K), (K * HQ, 1), (i_t * BT, 0), (BT, BK), (1, 0)) - tl.store(p_dq, b_dq.to(dq.dtype.element_ty), boundary_check=(0, 1)) - if USE_GATE: - tl.atomic_add(dg_cumsum + o_t * HQ, b_dg_cumsum_q, mask=m_t, sem='relaxed') - - -def parallel_path_bwd_dq_fn( - q, - k, - v, - g_cumsum, - do, - dg_cumsum, - hc_whole, - scale, - L, - D, - cu_seqlens, - S, - BT, - BS, - chunk_indices: torch.LongTensor | None = None, -): - B, T, num_blocks, HQ, K = q.shape - H, V = v.shape[-2:] - G = HQ // H - BK, BV = triton.next_power_of_2(K), triton.next_power_of_2(V) - - if chunk_indices is None and cu_seqlens is not None: - chunk_indices = prepare_chunk_indices(cu_seqlens, BT) - indices = chunk_indices - split_offsets = prepare_chunk_offsets(cu_seqlens, S) if cu_seqlens is not None else None - NT = triton.cdiv(T, BT) if cu_seqlens is None else len(indices) - - # should be NS - if cu_seqlens is not None: - assert split_offsets[-1] == hc_whole.shape[0] - dq = torch.empty(B, T, HQ, K, dtype=torch.float32, device=q.device) - - # [NS, HQ, K, K] instead of [NS, H, K, K] - # atomic add must be initialized to 0 - dhc_whole = torch.zeros(hc_whole.shape[0], HQ, K, K, dtype=torch.float32, device=q.device) - - parallel_path_bwd_dq_kernel[(NT, B*HQ)]( - q=q, - k=k, - v=v, - g_cumsum=g_cumsum, - hc_whole=hc_whole, - scale=scale, - L=L, - D=D, - dq=dq, - do=do, - dhc_whole=dhc_whole, - dg_cumsum=dg_cumsum, - cu_seqlens=cu_seqlens, - indices=indices, - split_offsets=split_offsets, - T=T, - S=S, - BT=BT, - BS=BS, - G=G, - HQ=HQ, - H=H, - K=K, - V=V, - BK=BK, - BV=BV, - NUM_BLOCKS=num_blocks, - num_warps=8 if (BT == 128 and K == 128) else 4, - num_stages=3 if check_shared_mem('ampere') else 2, - ) - return dq, dhc_whole, dg_cumsum diff --git a/build/torch-xpu/ops/path_attn/parallel_path_bwd_intra.py b/build/torch-xpu/ops/path_attn/parallel_path_bwd_intra.py deleted file mode 100644 index 314bd08730fc5df992bb3e48c62d3b8cd7508c4a..0000000000000000000000000000000000000000 --- a/build/torch-xpu/ops/path_attn/parallel_path_bwd_intra.py +++ /dev/null @@ -1,182 +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 -import triton.language as tl - -from ...ops.utils import prepare_chunk_indices - - -@triton.heuristics({ - 'IS_VARLEN': lambda args: args['offsets'] is not None, - 'USE_GATE': lambda args: args['g_cumsum'] is not None, -}) -@triton.jit(do_not_specialize=['T']) -def parallel_path_bwd_intra_chunk_kernel( - q, k, v, g_cumsum, w1, w2, - L, D, - dq, dq_new, dk, dv, dw1, dw2, do, dg_cumsum, - offsets, indices, - T, scale, - G: tl.constexpr, HQ: tl.constexpr, H: tl.constexpr, - K: tl.constexpr, V: tl.constexpr, BK: tl.constexpr, BV: tl.constexpr, - BT: tl.constexpr, S: tl.constexpr, - IS_VARLEN: tl.constexpr, USE_GATE: tl.constexpr, -): - i_t, i_bh = tl.program_id(0), tl.program_id(1) - i_b, i_hq = i_bh // HQ, i_bh % HQ - i_h = i_hq // G - - if IS_VARLEN: - i_n, i_t = tl.load(indices + i_t * 2).to(tl.int32), tl.load(indices + i_t * 2 + 1).to(tl.int32) - bos, eos = tl.load(offsets + i_n).to(tl.int64), tl.load(offsets + i_n + 1).to(tl.int64) - T = (eos - bos).to(tl.int32) - else: - i_n = i_b - bos, eos = (i_n * T).to(tl.int64), (i_n * T + T).to(tl.int64) - - # offset calculations - k += (bos * H + i_h) * K # GQA when H!=HQ - v += (bos * H + i_h) * V # GQA when H!=HQ - w1 += (bos * H + i_h) * K - w2 += (bos * H + i_h) * K - - q += (bos * HQ + i_hq) * K - dq += (bos * HQ + i_hq) * K - dq_new += (bos * HQ + i_hq) * K - dk += (bos * HQ + i_hq) * K - dv += (bos * HQ + i_hq) * V - do += (bos * HQ + i_hq) * V - dw1 += (bos * HQ + i_hq) * K - dw2 += (bos * HQ + i_hq) * K - L += (bos * HQ + i_hq) - D += (bos * HQ + i_hq) - if USE_GATE: - g_cumsum += (bos * HQ + i_hq) - dg_cumsum += (bos * HQ + i_hq) - - # constants - sm_scale = scale * 1.44269504 - - p_do = tl.make_block_ptr(do, (T, V), (HQ*V, 1), (i_t * BT, 0), (BT, BV), (1, 0)) - # [BT, BV] - b_do = tl.load(p_do, boundary_check=(0, 1)) - - p_delta = tl.make_block_ptr(D, (T, ), (HQ, ), (i_t * BT, ), (BT, ), (0, )) - b_delta = tl.load(p_delta, boundary_check=(0, )) - p_l = tl.make_block_ptr(L, (T, ), (HQ, ), (i_t * BT, ), (BT, ), (0, )) - b_l = tl.load(p_l, boundary_check=(0, )) - - b_dq = tl.zeros([BT, BK], dtype=tl.float32) - p_dq = tl.make_block_ptr(dq, (T, K), (HQ*K, 1), (i_t * BT, 0), (BT, BK), (1, 0)) - b_dq += tl.load(p_dq, boundary_check=(0, 1)) - p_q = tl.make_block_ptr(q, (T, K), (HQ*K, 1), (i_t * BT, 0), (BT, BK), (1, 0)) - b_q = tl.load(p_q, boundary_check=(0, 1)) - - if USE_GATE: - p_gq_cumsum = tl.make_block_ptr(g_cumsum, (T, ), (HQ, ), (i_t * BT, ), (BT, ), (0, )) - b_gq_cumsum = tl.load(p_gq_cumsum, boundary_check=(0, )) - b_dgq = tl.zeros([BT], dtype=tl.float32) - else: - b_dgq = None - - curr_start = (tl.floor(i_t * BT / S).to(tl.int32) * S).to(tl.int32) - - for offset in range(curr_start, i_t * BT, BT): - mask = offset + tl.arange(0, BT) < T - p_k = tl.make_block_ptr(k, (T, K), (H*K, 1), (offset, 0), (BT, BK), (1, 0)) - b_k = tl.load(p_k, boundary_check=(0, 1)) - b_q_tmp = tl.zeros([BT, BK], dtype=tl.float32) - b_q_tmp += b_q - for i_t_small in range(i_t * BT - BT, offset, -BT): - p_w1 = tl.make_block_ptr(w1, (T, K), (H*K, 1), (i_t_small, 0), (BT, BK), (1, 0)) - b_w1 = tl.load(p_w1, boundary_check=(0, 1)) - p_w2 = tl.make_block_ptr(w2, (T, K), (H*K, 1), (i_t_small, 0), (BT, BK), (1, 0)) - b_w2 = tl.load(p_w2, boundary_check=(0, 1)) - b_A_tmp = tl.dot(b_q_tmp.to(b_w1.dtype), tl.trans(b_w1)) - b_q_tmp -= tl.dot(b_A_tmp.to(b_w1.dtype), b_w2) - b_q2 = b_q_tmp.to(b_k.dtype) - b_A = tl.dot(b_q2, tl.trans(b_k)) - if USE_GATE: - p_gk_cumsum = tl.make_block_ptr(g_cumsum, (T, ), (HQ, ), (offset, ), (BT, ), (0, )) - b_gk_cumsum = tl.load(p_gk_cumsum, boundary_check=(0, )) - b_A = b_A + b_gq_cumsum[:, None] - b_gk_cumsum[None, :] - b_A = tl.where((i_t * BT + tl.arange(0, BT) < T)[:, None], b_A, float("-inf")) # avoid nan - b_A_softmax = tl.math.exp2(b_A * sm_scale - b_l[:, None]) - b_dv = tl.dot(tl.trans(b_A_softmax.to(b_do.dtype)), b_do) - tl.atomic_add( - dv + ((offset + tl.arange(0, BT)) * HQ * V)[:, None] + tl.arange(0, BV)[None, :], - b_dv.to(dv.dtype.element_ty), - mask=mask[:, None], - sem='relaxed', - ) - p_v = tl.make_block_ptr(v, (T, V), (V*H, 1), (offset, 0), (BT, BV), (1, 0)) - b_v = tl.load(p_v, boundary_check=(0, 1)) - b_dp = tl.dot(b_do, tl.trans(b_v)) - b_dA = ((b_dp - b_delta[:, None]) * b_A_softmax * scale) - if USE_GATE: - b_dgk = -tl.sum(b_dA, axis=0) - tl.atomic_add(dg_cumsum + (offset + tl.arange(0, BT)) * HQ, b_dgk, mask=mask, sem='relaxed') - b_dgq += tl.sum(b_dA, axis=1) - b_dA = b_dA.to(b_v.dtype) - b_dk = tl.dot(tl.trans(b_dA), b_q2) - tl.atomic_add(dk + (offset + tl.arange(0, BT))[:, None] * HQ*K + tl.arange(0, - BK)[None, :], b_dk, mask=mask[:, None], sem='relaxed') - p_w1 = tl.make_block_ptr(w1, (T, K), (H*K, 1), (offset, 0), (BT, BK), (1, 0)) - b_w1 = tl.load(p_w1, boundary_check=(0, 1)) - p_w2 = tl.make_block_ptr(w2, (T, K), (H*K, 1), (offset, 0), (BT, BK), (1, 0)) - b_w2 = tl.load(p_w2, boundary_check=(0, 1)) - b_dA2 = tl.dot(b_dq.to(b_w2.dtype), tl.trans(b_w2)).to(b_v.dtype) - b_A2 = tl.dot(b_q2.to(b_w1.dtype), tl.trans(b_w1)).to(b_v.dtype) - b_dw2 = -tl.dot(tl.trans(b_A2), b_dq.to(b_v.dtype)) - tl.atomic_add(dw2 + (offset + tl.arange(0, BT))[:, None] * HQ*K + tl.arange(0, - BK)[None, :], b_dw2, mask=mask[:, None], sem='relaxed') - b_dw1 = -tl.dot(tl.trans(b_dA2), b_q2.to(b_v.dtype)) - tl.atomic_add(dw1 + (offset + tl.arange(0, BT))[:, None] * HQ*K + tl.arange(0, - BK)[None, :], b_dw1, mask=mask[:, None], sem='relaxed') - b_dq -= tl.dot(b_dA2, b_w1.to(b_v.dtype)) - b_dq += tl.dot(b_dA.to(b_k.dtype), b_k) - - p_dq_new = tl.make_block_ptr(dq_new, (T, K), (HQ*K, 1), (i_t * BT, 0), (BT, BK), (1, 0)) - tl.store(p_dq_new, b_dq.to(dq_new.dtype.element_ty), boundary_check=(0, 1)) - mask = i_t * BT + tl.arange(0, BT) < T - if USE_GATE: - tl.atomic_add(dg_cumsum + (i_t * BT + tl.arange(0, BT)) * HQ, b_dgq, mask=mask, sem='relaxed') - - -def parallel_path_bwd_intra_chunk_fn( - q, k, v, g_cumsum, w1, w2, - dq, dk, dv, dg_cumsum, dw1, dw2, do, - scale, L, D, - cu_seqlens, - S, BT, - chunk_indices: torch.LongTensor | None = None, -): - assert dk.dtype == dv.dtype == dw1.dtype == dw2.dtype == torch.float32, 'atomic_add requires float32' - B, T, HQ, K = q.shape - assert dk.shape == dq.shape - - V = v.shape[-1] - H = k.shape[-2] - G = HQ // H - if chunk_indices is None and cu_seqlens is not None: - chunk_indices = prepare_chunk_indices(cu_seqlens, BT) - indices = chunk_indices - NT = triton.cdiv(T, BT) if cu_seqlens is None else len(indices) - dq_new = torch.empty_like(dq, dtype=q.dtype) - parallel_path_bwd_intra_chunk_kernel[(NT, B*HQ)]( - q=q, k=k, v=v, g_cumsum=g_cumsum, - w1=w1, w2=w2, L=L, D=D, - dq=dq, dq_new=dq_new, dk=dk, dv=dv, dw1=dw1, dw2=dw2, - do=do, dg_cumsum=dg_cumsum, - offsets=cu_seqlens, indices=indices, - T=T, S=S, BT=BT, scale=scale, - G=G, HQ=HQ, H=H, K=K, V=V, - BK=triton.next_power_of_2(K), BV=triton.next_power_of_2(V), - ) - return dq_new, dk, dv, dw1, dw2, dg_cumsum diff --git a/build/torch-xpu/ops/path_attn/parallel_path_fwd.py b/build/torch-xpu/ops/path_attn/parallel_path_fwd.py deleted file mode 100644 index 7fe37babba80d3a86bb64b3f9fbfc201438e7440..0000000000000000000000000000000000000000 --- a/build/torch-xpu/ops/path_attn/parallel_path_fwd.py +++ /dev/null @@ -1,205 +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 -import triton.language as tl - -from ...ops.utils import prepare_chunk_indices - - -@triton.heuristics({ - 'USE_GATE': lambda args: args['g_cumsum'] is not None, - 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None, -}) -@triton.jit(do_not_specialize=['T']) -def parallel_path_fwd_kernel( - q, - k, - v, - o, - o_new, - g_cumsum, - w1, - w2, - scale, - L, - L_new, - M, - cu_seqlens, - indices, - T, - G: tl.constexpr, - HQ: tl.constexpr, - H: tl.constexpr, - K: tl.constexpr, - V: tl.constexpr, - BT: tl.constexpr, - BS: tl.constexpr, - BK: tl.constexpr, - BV: tl.constexpr, - USE_GATE: tl.constexpr, - IS_VARLEN: tl.constexpr, -): - i_t, i_bh = tl.program_id(0), tl.program_id(1) - i_b, i_hq = i_bh // HQ, i_bh % HQ - i_h = i_hq // G - - if IS_VARLEN: - i_n, i_t = tl.load(indices + i_t * 2).to(tl.int32), tl.load(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).to(tl.int32) - else: - i_n = i_b - bos, eos = (i_n * T).to(tl.int64), (i_n * T + T).to(tl.int64) - - p_q = tl.make_block_ptr(q + (bos * HQ + i_hq) * K, (T, K), (HQ*K, 1), (i_t * BT, 0), (BT, BK), (1, 0)) - b_q = tl.zeros([BT, BK], dtype=tl.float32) - b_q += tl.load(p_q, boundary_check=(0, 1)) - sm_scale = scale * 1.44269504 - b_o = tl.zeros([BT, BV], dtype=tl.float32) - p_o = tl.make_block_ptr(o + (bos * HQ + i_hq) * V, (T, V), (HQ*V, 1), (i_t * BT, 0), (BT, BV), (1, 0)) - b_o += tl.load(p_o, boundary_check=(0, 1)) - - p_L = tl.make_block_ptr(L + bos * HQ + i_hq, (T, ), (HQ, ), (i_t * BT, ), (BT, ), (0,)) - p_M = tl.make_block_ptr(M + bos * HQ + i_hq, (T, ), (HQ, ), (i_t * BT, ), (BT, ), (0,)) - b_l = tl.load(p_L, boundary_check=(0,)) - b_m = tl.load(p_M, boundary_check=(0,)) - - if USE_GATE: - p_g_cumsum_q = tl.make_block_ptr(g_cumsum + bos * HQ + i_hq, (T, ), (HQ, ), (i_t * BT, ), (BT, ), (0,)) - b_g_cumsum_q = tl.load(p_g_cumsum_q, boundary_check=(0,)) - else: - b_g_cumsum_q = None - - for offset in range((i_t + 1) * BT - 2 * BS, i_t*BT-BS, -BS): - p_k = tl.make_block_ptr(k + (bos * H + i_h) * K, (K, T), (1, K*H), (0, offset), (BK, BS), (0, 1)) # GQA when H!=HQ - p_v = tl.make_block_ptr(v + (bos * H + i_h) * V, (T, V), (V*H, 1), (offset, 0), (BS, BV), (1, 0)) # GQA when H!=HQ - p_w1 = tl.make_block_ptr(w1 + (bos * H + i_h) * K, (K, T), (1, K*H), (0, offset), (BK, BS), (0, 1)) - p_w2 = tl.make_block_ptr(w2 + (bos * H + i_h) * K, (T, K), (K*H, 1), (offset, 0), (BS, BK), (1, 0)) - # [BK, BS] - - b_k = tl.load(p_k, boundary_check=(0, 1)) - # [BS, BV] - b_v = tl.load(p_v, boundary_check=(0, 1)) - # [BK, BK] - b_w1 = tl.load(p_w1, boundary_check=(0, 1)) - b_w2 = tl.load(p_w2, boundary_check=(0, 1)) - # [BT, BS] - m_s = i_t * BT + tl.arange(0, BT) >= (offset + BS) - b_s = tl.dot(b_q.to(b_k.dtype), b_k) - - if USE_GATE: - p_g_cumsum_k = tl.make_block_ptr(g_cumsum + (bos * HQ + i_hq), (T, ), (HQ, ), (offset, ), (BS, ), (0,)) - b_g_cumsum_k = tl.load(p_g_cumsum_k, boundary_check=(0,)) - b_s = b_s + b_g_cumsum_q[:, None] - b_g_cumsum_k[None, :] - b_s = tl.where(m_s[:, None], b_s * sm_scale, float("-inf")) - b_m_new = tl.maximum(b_m, tl.max(b_s, 1)) - alpha = tl.math.exp2(b_m - b_m_new) - b_s = tl.math.exp2(b_s - b_m_new[:, None]) - b_o *= alpha[:, None] - b_l = b_l * alpha + tl.sum(b_s, 1) - b_m = b_m_new - b_o += tl.dot(b_s.to(b_v.dtype), b_v) - b_s2 = tl.dot(b_q.to(b_w1.dtype), b_w1) - b_s2 = tl.where(m_s[:, None], b_s2, 0) - b_q -= tl.dot(b_s2.to(b_w2.dtype), b_w2) - - tl.debug_barrier() - - for offset in range(i_t * BT - BS, -BS, -BS): - p_k = tl.make_block_ptr(k + (bos * H + i_h) * K, (K, T), (1, K*H), (0, offset), (BK, BS), (0, 1)) # GQA when H!=HQ - p_v = tl.make_block_ptr(v + (bos * H + i_h) * V, (T, V), (V*H, 1), (offset, 0), (BS, BV), (1, 0)) # GQA when H!=HQ - p_w1 = tl.make_block_ptr(w1 + (bos * H + i_h) * K, (K, T), (1, K*H), (0, offset), (BK, BS), (0, 1)) - p_w2 = tl.make_block_ptr(w2 + (bos * H + i_h) * K, (T, K), (K*H, 1), (offset, 0), (BS, BK), (1, 0)) - # [BK, BS] - b_k = tl.load(p_k, boundary_check=(0, 1)) - # [BS, BV] - b_v = tl.load(p_v, boundary_check=(0, 1)) - b_w1 = tl.load(p_w1, boundary_check=(0, 1)) - b_w2 = tl.load(p_w2, boundary_check=(0, 1)) - # [BT, BS] - b_s = tl.dot(b_q.to(b_k.dtype), b_k) - if USE_GATE: - p_g_cumsum_k = tl.make_block_ptr(g_cumsum + (bos * HQ + i_hq), (T, ), (HQ, ), (offset, ), (BS, ), (0,)) - b_g_cumsum_k = tl.load(p_g_cumsum_k, boundary_check=(0,)) - b_s = b_s + b_g_cumsum_q[:, None] - b_g_cumsum_k[None, :] - b_s = b_s * sm_scale - b_m_new = tl.maximum(b_m, tl.max(b_s, 1)) - alpha = tl.math.exp2(b_m - b_m_new) - b_s = tl.math.exp2(b_s - b_m_new[:, None]) - b_o *= alpha[:, None] - b_l = b_l * alpha + tl.sum(b_s, 1) - b_m = b_m_new - b_o += tl.dot(b_s.to(b_v.dtype), b_v) - b_s2 = tl.dot(b_q.to(b_w1.dtype), b_w1) - b_q -= tl.dot(b_s2.to(b_w2.dtype), b_w2) - - b_o = b_o / b_l[:, None] - p_o_new = tl.make_block_ptr(o_new + (bos * HQ + i_hq) * V, (T, V), (HQ*V, 1), (i_t*BT, 0), (BT, BV), (1, 0)) - tl.store(p_o_new, b_o.to(p_o_new.dtype.element_ty), boundary_check=(0, 1)) - b_l = tl.math.log2(b_l) + b_m - p_L_new = tl.make_block_ptr(L_new + (bos * HQ + i_hq), (T, ), (HQ, ), (i_t * BT, ), (BT, ), (0,)) - tl.store(p_L_new, b_l.to(p_L_new.dtype.element_ty), boundary_check=(0,)) - - -def parallel_path_fwd_fn( - q, - k, - v, - o, - g_cumsum, - w1, - w2, - scale, - L, - M, - cu_seqlens, - BT, - BS, - chunk_indices: torch.LongTensor | None = None, -): - B, T, HQ, K = q.shape - V = v.shape[-1] - H = k.shape[-2] - G = HQ // H - if chunk_indices is None and cu_seqlens is not None: - chunk_indices = prepare_chunk_indices(cu_seqlens, BT) - indices = chunk_indices - NT = triton.cdiv(T, BT) if cu_seqlens is None else len(indices) - grid = (NT, B * HQ) - o_new = torch.empty_like(o, dtype=v.dtype) - L_new = torch.empty_like(L) - - parallel_path_fwd_kernel[grid]( - q=q, - k=k, - v=v, - o=o, - o_new=o_new, - w1=w1, - w2=w2, - g_cumsum=g_cumsum, - scale=scale, - cu_seqlens=cu_seqlens, - indices=indices, - L=L, - L_new=L_new, - M=M, - T=T, - K=K, - V=V, - BK=triton.next_power_of_2(K), - BV=triton.next_power_of_2(V), - G=G, - HQ=HQ, - H=H, - BS=BS, - BT=BT, - num_warps=8 if (BT == 128 and K == 128) else 4, - ) - return o_new, L_new diff --git a/build/torch-xpu/ops/path_attn/prepare_k_cache.py b/build/torch-xpu/ops/path_attn/prepare_k_cache.py deleted file mode 100644 index 2fe4d40c191af498ce0e56a321e654040b27f442..0000000000000000000000000000000000000000 --- a/build/torch-xpu/ops/path_attn/prepare_k_cache.py +++ /dev/null @@ -1,83 +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 -import triton.language as tl - -from ...ops.utils import prepare_chunk_indices - - -@triton.heuristics({ - 'IS_VARLEN': lambda args: args['offsets'] is not None, -}) -@triton.jit(do_not_specialize=['T']) -def parallel_path_fwd_kernel_prepare_k_cache( - k, k_new, w1, w2, - offsets, indices, - T, - H: tl.constexpr, - K: tl.constexpr, - BT: tl.constexpr, BK: tl.constexpr, - IS_VARLEN: tl.constexpr, -): - i_t, i_bh = tl.program_id(0), tl.program_id(1) - i_b, i_h = i_bh // H, i_bh % H - - if IS_VARLEN: - i_n, i_t = tl.load(indices + i_t * 2).to(tl.int32), tl.load(indices + i_t * 2 + 1).to(tl.int32) - bos, eos = tl.load(offsets + i_n).to(tl.int64), tl.load(offsets + i_n + 1).to(tl.int64) - T = (eos - bos).to(tl.int32) - else: - i_n = i_b - bos, eos = (i_n * T).to(tl.int64), (i_n * T + T).to(tl.int64) - - k += (bos * H + i_h) * K - k_new += (bos * H + i_h) * K - w1 += (bos * H + i_h) * K - w2 += (bos * H + i_h) * K - # constants - p_k = tl.make_block_ptr(k, (T, K), (H*K, 1), (i_t * BT, 0), (BT, BK), (1, 0)) - b_k = tl.zeros([BT, BK], dtype=tl.float32) - b_k += tl.load(p_k, boundary_check=(0, 1)) - for k_block_idx in range(i_t + 1, tl.cdiv(T, BT)): - p_w1 = tl.make_block_ptr(w1, (T, K), (H*K, 1), (k_block_idx * BT, 0), (BT, BK), (1, 0)) - p_w2 = tl.make_block_ptr(w2, (T, K), (H*K, 1), (k_block_idx * BT, 0), (BT, BK), (1, 0)) - b_w1 = tl.load(p_w1, boundary_check=(0, 1)) - b_w2 = tl.load(p_w2, boundary_check=(0, 1)) - b_A = tl.dot(b_k.to(b_w2.dtype), tl.trans(b_w2)) - b_k = b_k - tl.dot(b_A.to(b_w1.dtype), b_w1) - - p_k_new = tl.make_block_ptr(k_new, (T, K), (H*K, 1), (i_t * BT, 0), (BT, BK), (1, 0)) - tl.store(p_k_new, b_k.to(p_k_new.dtype.element_ty), boundary_check=(0, 1)) - - -def prepare_k_cache_fn(k, w1, w2, cu_seqlens, BS, use_cache=False, chunk_indices: torch.LongTensor | None = None): - if not use_cache: - return None - else: - B, T, H, K = k.shape - k_new = torch.empty_like(k) - if chunk_indices is None and cu_seqlens is not None: - chunk_indices = prepare_chunk_indices(cu_seqlens, BS) - indices = chunk_indices - NT = triton.cdiv(T, BS) if cu_seqlens is None else len(indices) - grid = (NT, B * H) - parallel_path_fwd_kernel_prepare_k_cache[grid]( - k=k, - k_new=k_new, - w1=w1, - w2=w2, - offsets=cu_seqlens, - indices=indices, - H=H, - T=T, - K=K, - BT=BS, - BK=triton.next_power_of_2(K), - ) - return k_new diff --git a/build/torch-xpu/ops/path_attn/transform_q.py b/build/torch-xpu/ops/path_attn/transform_q.py deleted file mode 100644 index 303d7a6322168f377daee064b852d5acea588801..0000000000000000000000000000000000000000 --- a/build/torch-xpu/ops/path_attn/transform_q.py +++ /dev/null @@ -1,115 +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 -import triton.language as tl - -from ...ops.utils import get_max_num_splits, prepare_chunk_indices - - -@triton.heuristics({ - 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None, -}) -@triton.jit(do_not_specialize=['T']) -def transform_q_fwd_kernel( - q, - q_new, - w1, - w2, - cu_seqlens, - indices, - T, - S: tl.constexpr, - G: tl.constexpr, - HQ: tl.constexpr, - H: tl.constexpr, - K: tl.constexpr, - BT: tl.constexpr, - BS: tl.constexpr, - BK: tl.constexpr, - NUM_BLOCKS: tl.constexpr, - IS_VARLEN: tl.constexpr, -): - i_t, i_bh = tl.program_id(0), tl.program_id(1) - i_b, i_hq = i_bh // HQ, i_bh % HQ - i_h = i_hq // G - - if IS_VARLEN: - i_n, i_t = tl.load(indices + i_t * 2).to(tl.int32), tl.load(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).to(tl.int32) - else: - i_n = i_b - bos, eos = (i_n * T).to(tl.int64), (i_n * T + T).to(tl.int64) - # boh = i_n * tl.cdiv(T, BS) - p_q = tl.make_block_ptr(q + (bos * HQ + i_hq) * K, (T, K), (HQ*K, 1), (i_t * BT, 0), (BT, BK), (1, 0)) - b_q = tl.zeros([BT, BK], dtype=tl.float32) - b_q += tl.load(p_q, boundary_check=(0, 1)) - - if BS == BT: - if (i_t * BT) % S == 0: - p_q_new = tl.make_block_ptr(q_new + ((bos.to(tl.int64) * NUM_BLOCKS + (i_t * BT // S)) * HQ + i_hq) * K, - (T, K), (HQ*K*NUM_BLOCKS, 1), (i_t * BT, 0), (BT, BK), (1, 0)) - tl.store(p_q_new, b_q.to(q_new.dtype.element_ty), boundary_check=(0, 1)) - - for offset in range((i_t + 1) * BT - 2 * BS, S-BS, -BS): - p_w1 = tl.make_block_ptr(w1 + (bos * H + i_h) * K, (K, T), (1, K*H), (0, offset), (BK, BS), (0, 1)) - p_w2 = tl.make_block_ptr(w2 + (bos * H + i_h) * K, (T, K), (K*H, 1), (offset, 0), (BS, BK), (1, 0)) - b_w1 = tl.load(p_w1, boundary_check=(0, 1)) - b_w2 = tl.load(p_w2, boundary_check=(0, 1)) - m_s = i_t * BT + tl.arange(0, BT) >= (offset + BS) - b_s2 = tl.dot(b_q.to(b_w1.dtype), b_w1) - b_s2 = tl.where(m_s[:, None], b_s2, 0) - b_q -= tl.dot(b_s2.to(b_w2.dtype), b_w2) - - if offset % S == 0: - p_q_new = tl.make_block_ptr(q_new + ((bos.to(tl.int64) * NUM_BLOCKS + (offset // S)) * HQ + i_hq) * K, - (T, K), (HQ*K*NUM_BLOCKS, 1), (i_t * BT, 0), (BT, BK), (1, 0)) - tl.store(p_q_new, b_q.to(q_new.dtype.element_ty), boundary_check=(0, 1)) - - -def transform_q_fwd_fn( - q, - w1, - w2, - cu_seqlens, - BT, - BS, - S, - chunk_indices: torch.LongTensor | None = None, -): - B, T, HQ, K = q.shape - H = w1.shape[-2] - G = HQ // H - if chunk_indices is None and cu_seqlens is not None: - chunk_indices = prepare_chunk_indices(cu_seqlens, BT) - indices = chunk_indices - NT = triton.cdiv(T, BT) if cu_seqlens is None else len(indices) - - num_blocks = triton.cdiv(T, S) if cu_seqlens is None else get_max_num_splits(cu_seqlens, S) - q_new = torch.zeros(B, T, num_blocks, HQ, K, dtype=q.dtype, device=q.device) - transform_q_fwd_kernel[(NT, B * HQ)]( - q=q, - q_new=q_new, - w1=w1, - w2=w2, - cu_seqlens=cu_seqlens, - indices=indices, - T=T, - K=K, - BK=triton.next_power_of_2(K), - G=G, - HQ=HQ, - H=H, - BS=BS, - BT=BT, - S=S, - NUM_BLOCKS=num_blocks, - num_warps=8 if (BT == 128 and K == 128) else 4, - ) - return q_new diff --git a/build/torch-xpu/ops/rebased/naive.py b/build/torch-xpu/ops/rebased/naive.py deleted file mode 100644 index 5120655d651e3602b2d7bb22a99fdf95965c8351..0000000000000000000000000000000000000000 --- a/build/torch-xpu/ops/rebased/naive.py +++ /dev/null @@ -1,29 +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 - - -def naive_parallel_rebased( - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - scale: float | None = None, - use_norm: bool = True, -) -> torch.Tensor: - if scale is None: - scale = q.shape[-1] ** -0.5 - q = q * scale - attn = q @ k.transpose(-2, -1) - attn = attn ** 2 - attn.masked_fill_(~torch.tril(torch.ones(q.shape[-2], q.shape[-2], dtype=torch.bool, device=q.device)), 0) - o = attn @ v - if use_norm: - z = attn.sum(-1) - return o / (z[..., None] + 1e-6) - else: - return o diff --git a/build/torch-xpu/ops/rebased/parallel.py b/build/torch-xpu/ops/rebased/parallel.py deleted file mode 100644 index e4945e8f2c32ec859acdeffcea3fcff2d7076ef3..0000000000000000000000000000000000000000 --- a/build/torch-xpu/ops/rebased/parallel.py +++ /dev/null @@ -1,467 +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 -import triton.language as tl - -from ...utils import autocast_custom_bwd, autocast_custom_fwd, input_guard - -# Rebased: Linear Transformers with Learnable Kernel Functions are Better In-Context Models -# https://github.com/corl-team/rebased/blob/main/flash_linear_attention/fla/ops/triton/rebased_fast/parallel.py - - -@triton.jit(do_not_specialize=['T']) -def parallel_rebased_fwd_kernel( - q, - k, - v, - o, - z, - scale, - T, - B: tl.constexpr, - H: tl.constexpr, - K: tl.constexpr, - V: tl.constexpr, - BTL: tl.constexpr, - BTS: tl.constexpr, - BK: tl.constexpr, - BV: tl.constexpr, -): - # i_c: chunk index. used for sequence parallelism - i_kv, i_c, i_bh = tl.program_id(0), tl.program_id(1), tl.program_id(2) - NV = tl.cdiv(V, BV) - i_k = i_kv // (NV) - i_v = i_kv % (NV) - - p_q = tl.make_block_ptr(q + i_bh * T*K, (T, K), (K, 1), (i_c*BTL, i_k*BK), (BTL, BK), (1, 0)) - p_k = tl.make_block_ptr(k + i_bh * T*K, (K, T), (1, K), (i_k*BK, 0), (BK, BTS), (0, 1)) - p_v = tl.make_block_ptr(v + i_bh * T*V, (T, V), (V, 1), (0, i_v*BV), (BTS, BV), (1, 0)) - - # [BQ, BD] block Q, in the shared memory throughout the whole kernel - b_q = tl.load(p_q, boundary_check=(0, 1)) - b_q = (b_q * scale).to(b_q.dtype) - b_o = tl.zeros([BTL, BV], dtype=tl.float32) - b_z = tl.zeros([BTL], dtype=tl.float32) - - # Q block and K block have no overlap - # no need for mask, thereby saving flops - for _ in range(0, i_c*BTL, BTS): - # [BK, BTS] - b_k = tl.load(p_k, boundary_check=(0, 1)) - - # [BTS, BV] - b_v = tl.load(p_v, boundary_check=(0, 1)) - # [BTL, BTS] - b_s = tl.dot(b_q, (b_k), allow_tf32=False) - b_s = b_s * b_s - b_z += tl.sum(b_s, axis=1) - - # [BQ, BD] - b_o = b_o + tl.dot(b_s.to(b_v.dtype), b_v, allow_tf32=False) - p_k = tl.advance(p_k, (0, BTS)) - p_v = tl.advance(p_v, (BTS, 0)) - - # # rescale interchunk output - tl.debug_barrier() - o_q = tl.arange(0, BTL) - # # sync threads, easy for compiler to optimize - # tl.debug_barrier() - - o_k = tl.arange(0, BTS) - p_k = tl.make_block_ptr(k + i_bh * T*K, (K, T), (1, K), (i_k*BK, i_c*BTL), (BK, BTS), (0, 1)) - p_v = tl.make_block_ptr(v + i_bh * T*V, (T, V), (V, 1), (i_c*BTL, i_v*BV), (BTS, BV), (1, 0)) - # Q block and K block have overlap. masks required - for _ in range(i_c*BTL, (i_c + 1) * BTL, BTS): - # [BK, BTS] - b_k = tl.load(p_k, boundary_check=(0, 1)) - # [BTS, BV] - b_v = tl.load(p_v, boundary_check=(0, 1)) - # [BTL, BTS] - m_s = o_q[:, None] >= o_k[None, :] - b_s = tl.dot(b_q, b_k, allow_tf32=False) - b_s = b_s * b_s - b_s = tl.where(m_s, b_s, 0) - b_z += tl.sum(b_s, axis=1) - # [BTL, BV] - b_o += tl.dot(b_s.to(b_q.dtype), b_v, allow_tf32=False) - p_k = tl.advance(p_k, (0, BTS)) - p_v = tl.advance(p_v, (BTS, 0)) - o_k += BTS - - p_o = tl.make_block_ptr(o + (i_bh + B * H * i_k) * T*V, (T, V), (V, 1), (i_c*BTL, i_v*BV), (BTL, BV), (1, 0)) - p_z = z + (i_bh + B * H * i_k) * T + i_c*BTL + tl.arange(0, BTL) - tl.store(p_o, b_o.to(p_o.dtype.element_ty), boundary_check=(0, 1)) - tl.store(p_z, b_z.to(p_z.dtype.element_ty), mask=((i_c*BTL + tl.arange(0, BTL)) < T)) - - -@triton.jit(do_not_specialize=['T']) -def _parallel_rebased_bwd_dq( - i_bh, - i_c, - i_k, - i_v, - i_h, - q, - k, - v, - do, - dz, - dq, - scale, - T, - B: tl.constexpr, - H: tl.constexpr, - K: tl.constexpr, - V: tl.constexpr, - BTL: tl.constexpr, - BTS: tl.constexpr, - BK: tl.constexpr, - BV: tl.constexpr, -): - p_do = tl.make_block_ptr(do + i_bh * T*V, (T, V), (V, 1), (i_c*BTL, i_v*BV), (BTL, BV), (1, 0)) - p_q = tl.make_block_ptr(q + (i_bh) * T*K, (T, K), (K, 1), (i_c*BTL, i_k*BK), (BTL, BK), (1, 0)) - b_q = tl.load(p_q, boundary_check=(0, 1)) - b_do = tl.load(p_do, boundary_check=(0, 1)).to(b_q.dtype) - b_q = (b_q * scale).to(b_q.dtype) - b_dq = tl.zeros([BTL, BK], dtype=tl.float32) - p_k = tl.make_block_ptr(k + i_bh * T*K, (T, K), (K, 1), (0, i_k*BK), (BTS, BK), (1, 0)) - p_v = tl.make_block_ptr(v + i_bh * T*V, (V, T), (1, V), (i_v*BV, 0), (BV, BTS), (0, 1)) - p_dz = dz + i_bh * T + i_c*BTL + tl.arange(0, BTL) - b_dz = tl.load(p_dz, mask=(i_c*BTL + tl.arange(0, BTL)) < T) - - for _ in range(0, i_c*BTL, BTS): - # [BTS, BK] - b_k = tl.load(p_k, boundary_check=(0, 1)) - # [BV, BTS] - b_v = tl.load(p_v, boundary_check=(0, 1)) - # [BTL, BTS] - b_ds = tl.dot(b_do, b_v, allow_tf32=False) - if i_v == 0: - b_ds += b_dz[:, None] - else: - b_ds = b_ds - b_s = tl.dot(b_q, tl.trans(b_k), allow_tf32=False) - # [BQ, BD] - b_dq += tl.dot((2 * b_ds * b_s).to(b_v.dtype), b_k, allow_tf32=False) - p_k = tl.advance(p_k, (BTS, 0)) - p_v = tl.advance(p_v, (0, BTS)) - - b_dq *= scale - o_q = tl.arange(0, BTL) - o_k = tl.arange(0, BTS) - p_k = tl.make_block_ptr(k + i_bh * T*K, (T, K), (K, 1), (i_c*BTL, i_k*BK), (BTS, BK), (1, 0)) - p_v = tl.make_block_ptr(v + i_bh * T*V, (V, T), (1, V), (i_v*BV, i_c*BTL), (BV, BTS), (0, 1)) - # Q block and K block have overlap. masks required - for _ in range(i_c*BTL, (i_c + 1) * BTL, BTS): - # [BTS, BK] - b_k = tl.load(p_k, boundary_check=(0, 1)) - # [BV, BTS] - b_v = tl.load(p_v, boundary_check=(0, 1)) - # [BTL, BTS] - m_s = o_q[:, None] >= o_k[None, :] - b_ds = tl.dot(b_do, b_v, allow_tf32=False) - if i_v == 0: - b_ds += b_dz[:, None] - else: - b_ds = b_ds - b_ds = tl.where(m_s, b_ds, 0) * scale - b_s = tl.dot(b_q, tl.trans(b_k), allow_tf32=False) - b_s = tl.where(m_s, b_s, 0) - # [BTL, BK] - b_dq += tl.dot((2 * b_ds * b_s).to(b_k.dtype), - b_k, allow_tf32=False) - p_k = tl.advance(p_k, (BTS, 0)) - p_v = tl.advance(p_v, (0, BTS)) - o_k += BTS - p_dq = tl.make_block_ptr(dq + (i_bh + B * H * i_v) * T*K, (T, K), (K, 1), (i_c*BTL, i_k*BK), (BTL, BK), (1, 0)) - tl.store(p_dq, b_dq.to(p_dq.dtype.element_ty), boundary_check=(0, 1)) - return - - -@triton.jit(do_not_specialize=['T']) -def _parallel_rebased_bwd_dkv( - i_bh, - i_c, - i_k, - i_v, - i_h, - q, - k, - v, - do, - dz, - dk, - dv, - scale, - T, - B: tl.constexpr, - H: tl.constexpr, - K: tl.constexpr, - V: tl.constexpr, - BTL: tl.constexpr, - BTS: tl.constexpr, - BK: tl.constexpr, - BV: tl.constexpr, -): - # compute dk dv - p_k = tl.make_block_ptr(k + i_bh * T*K, (T, K), (K, 1), (i_c*BTL, i_k*BK), (BTL, BK), (1, 0)) - p_v = tl.make_block_ptr(v + i_bh * T*V, (T, V), (V, 1), (i_c*BTL, i_v*BV), (BTL, BV), (1, 0)) - b_k, b_v = tl.load(p_k, boundary_check=(0, 1)), tl.load(p_v, boundary_check=(0, 1)) - b_dk, b_dv = tl.zeros([BTL, BK], dtype=tl.float32), tl.zeros( - [BTL, BV], dtype=tl.float32) - - for i in range((tl.cdiv(T, BTS) * BTS)-BTS, (i_c + 1) * BTL - BTS, -BTS): - p_q = tl.make_block_ptr(q + i_bh * T*K, (K, T), (1, K), (i_k*BK, i), (BK, BTS), (0, 1)) - p_do = tl.make_block_ptr(do + i_bh * T*V, (V, T), (1, V), (i_v*BV, i), (BV, BTS), (0, 1)) - p_dz = dz + i_bh * T + i + tl.arange(0, BTS) - # [BK, BTS] - b_q = tl.load(p_q, boundary_check=(0, 1)) - # [BV, BTS] - b_do = tl.load(p_do, boundary_check=(0, 1)).to(b_q.dtype) - b_dz = tl.load(p_dz, mask=(i + tl.arange(0, BTS)) < T) - # [BTL, BTS] - b_s = tl.dot(b_k.to(b_q.dtype), b_q, allow_tf32=False) * scale - b_s2 = b_s * b_s - b_dv += tl.dot(b_s2.to(b_q.dtype), tl.trans(b_do), allow_tf32=False) - b_ds = tl.dot(b_v, b_do, allow_tf32=False) * scale - if i_v == 0: - b_ds += b_dz[None, :] * scale - else: - b_ds = b_ds - b_dk += tl.dot((2 * b_ds * b_s).to(b_q.dtype), tl.trans(b_q), allow_tf32=False) - - tl.debug_barrier() - o_q, o_k = tl.arange(0, BTS), tl.arange(0, BTL) - for i in range(i_c*BTL, (i_c+1)*BTL, BTS): - p_q = tl.make_block_ptr(q + i_bh * T*K, (K, T), (1, K), (i_k*BK, i), (BK, BTS), (0, 1)) - p_do = tl.make_block_ptr(do + i_bh * T*V, (V, T), (1, V), (i_v*BV, i), (BV, BTS), (0, 1)) - p_dz = dz + i_bh * T + i + tl.arange(0, BTS) - b_q = tl.load(p_q, boundary_check=(0, 1)) # [BD, BQ] - b_do = tl.load(p_do, boundary_check=(0, 1)).to(b_q.dtype) - b_dz = tl.load(p_dz, mask=(i + tl.arange(0, BTS)) < T) - # [BK, BQ] - m_s = o_k[:, None] <= o_q[None, :] - b_s = tl.dot(b_k, b_q, allow_tf32=False) * scale - b_s2 = b_s * b_s - b_s = tl.where(m_s, b_s, 0) - b_s2 = tl.where(m_s, b_s2, 0) - - b_ds = tl.dot(b_v, b_do, allow_tf32=False) - if i_v == 0: - b_ds += b_dz[None, :] - else: - b_ds = b_ds - b_ds = tl.where(m_s, b_ds, 0) * scale - # [BK, BD] - b_dv += tl.dot(b_s2.to(b_q.dtype), tl.trans(b_do), allow_tf32=False) - b_dk += tl.dot((2 * b_ds * b_s).to(b_q.dtype), tl.trans(b_q), allow_tf32=False) - o_q += BTS - - p_dk = tl.make_block_ptr(dk + (i_bh + B * H * i_v) * T*K, (T, K), (K, 1), (i_c*BTL, i_k*BK), (BTL, BK), (1, 0)) - p_dv = tl.make_block_ptr(dv + (i_bh + B * H * i_k) * T*V, (T, V), (V, 1), (i_c*BTL, i_v*BV), (BTL, BV), (1, 0)) - tl.store(p_dk, b_dk.to(p_dk.dtype.element_ty), boundary_check=(0, 1)) - tl.store(p_dv, b_dv.to(p_dv.dtype.element_ty), boundary_check=(0, 1)) - return - - -@triton.jit(do_not_specialize=['T']) -def parallel_rebased_bwd_kernel( - q, - k, - v, - do, - dz, - dq, - dk, - dv, - scale, - T, - B: tl.constexpr, - H: tl.constexpr, - K: tl.constexpr, - V: tl.constexpr, - BTL: tl.constexpr, - BTS: tl.constexpr, - BK: tl.constexpr, - BV: tl.constexpr, -): - i_kv, i_c, i_bh = tl.program_id(0), tl.program_id(1), tl.program_id(2) - NV = tl.cdiv(V, BV) - i_k = i_kv // (NV) - i_v = i_kv % (NV) - i_h = i_bh % H - _parallel_rebased_bwd_dq( - i_bh, - i_c, - i_k, - i_v, - i_h, - q, - k, - v, - do, - dz, - dq, - scale, - B=B, - H=H, - T=T, - K=K, - V=V, - BTL=BTL, - BTS=BTS, - BK=BK, - BV=BV, - ) - tl.debug_barrier() - _parallel_rebased_bwd_dkv( - i_bh, - i_c, - i_k, - i_v, - i_h, - q, - k, - v, - do, - dz, - dk, - dv, - scale, - B=B, - H=H, - T=T, - K=K, - V=V, - BTL=BTL, - BTS=BTS, - BK=BK, - BV=BV, - ) - - -class ParallelBasedFunction(torch.autograd.Function): - - @staticmethod - @input_guard - @autocast_custom_fwd - def forward(ctx, q, k, v, scale): - BTL, BTS = 128, 32 - assert BTL % BTS == 0 - # assert q.shape[-1] % 16 == 0 - BK = min(128, max(triton.next_power_of_2(k.shape[-1]), 16)) - BV = min(128, max(triton.next_power_of_2(v.shape[-1]), 16)) - B, H, T, K, V = *k.shape, v.shape[-1] - num_stages = 2 - num_warps = 4 - NK = triton.cdiv(K, BK) - NV = triton.cdiv(V, BV) - grid = (NK * NV, triton.cdiv(T, BTL), B * H) - - assert NK == 1, "will encounter some synchronization issue if not." - - o = torch.empty(NK, B, H, T, V, device=q.device) - z = torch.empty(NK, B, H, T, device=q.device) - parallel_rebased_fwd_kernel[grid]( - q, - k, - v, - o, - z, - scale, - T=T, - B=B, - H=H, - K=K, - V=V, - BTL=BTL, - BTS=BTS, - BK=BK, - BV=BV, - num_warps=num_warps, - num_stages=num_stages, - ) - ctx.save_for_backward(q, k, v) - ctx.scale = scale - return o.sum(0).to(q.dtype), z.sum(0).to(q.dtype) - - @staticmethod - @input_guard - @autocast_custom_bwd - def backward(ctx, do, dz): - q, k, v = ctx.saved_tensors - scale = ctx.scale - BTL, BTS = 64, 32 - assert BTL % BTS == 0 - BK = min(128, max(triton.next_power_of_2(k.shape[-1]), 16)) - BV = min(128, max(triton.next_power_of_2(v.shape[-1]), 16)) - B, H, T, K, V = *k.shape, v.shape[-1] - num_stages = 2 - num_warps = 4 - NK = triton.cdiv(K, BK) - NV = triton.cdiv(V, BV) - grid = (NK * NV, triton.cdiv(T, BTL), B * H) - - assert NK == 1, "will encounter some synchronization issue if not" - - dq = torch.empty(NV, B, H, T, K, dtype=q.dtype, device=q.device) - dk = torch.empty(NV, B, H, T, K, dtype=q.dtype, device=q.device) - dv = torch.empty(NK, B, H, T, V, dtype=q.dtype, device=q.device) - - parallel_rebased_bwd_kernel[grid]( - q, - k, - v, - do, - dz, - dq, - dk, - dv, - scale, - T=T, - B=B, - H=H, - K=K, - V=V, - BTL=BTL, - BTS=BTS, - BK=BK, - BV=BV, - num_warps=num_warps, - num_stages=num_stages, - ) - - return dq.sum(0).to(q.dtype), dk.sum(0).to(k.dtype), dv.sum(0).to(v.dtype), None - - -def parallel_rebased( - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - eps: float = 1e-5, - use_scale: bool = True, - use_normalize: bool = True, - return_both: bool = False, - head_first: bool = False, -): - assert q.shape[-1] <= 128, "only support feature dim up to 128" - if use_scale: - scale = q.shape[-1] ** -0.5 - else: - scale = 1 - if not head_first: - q, k, v = map(lambda x: x.transpose(1, 2), (q, k, v)) - o, z = ParallelBasedFunction.apply(q, k, v, scale) - if return_both: - return o, z - if use_normalize: - o = o / (z[..., None] + eps) - if not head_first: - o = o.transpose(1, 2) - return o.to(q.dtype) diff --git a/build/torch-xpu/ops/retention/__init__.py b/build/torch-xpu/ops/retention/__init__.py deleted file mode 100644 index 35907b70566eda7c6e1e142dda6c759e10bcf6e0..0000000000000000000000000000000000000000 --- a/build/torch-xpu/ops/retention/__init__.py +++ /dev/null @@ -1,18 +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 .chunk import chunk_retention -from .fused_chunk import fused_chunk_retention -from .fused_recurrent import fused_recurrent_retention -from .parallel import parallel_retention - -__all__ = [ - 'chunk_retention', - 'fused_chunk_retention', - 'fused_recurrent_retention', - 'parallel_retention', -] diff --git a/build/torch-xpu/ops/retention/chunk.py b/build/torch-xpu/ops/retention/chunk.py deleted file mode 100644 index 9afc986508228684ddb89e312e9b50fe24437f2e..0000000000000000000000000000000000000000 --- a/build/torch-xpu/ops/retention/chunk.py +++ /dev/null @@ -1,80 +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 - -from ...ops.simple_gla.chunk import chunk_simple_gla - - -@torch.compiler.disable -def chunk_retention( - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - scale: float | None = None, - initial_state: torch.Tensor | None = None, - output_final_state: bool = False, - cu_seqlens: torch.LongTensor | None = None, - **kwargs, -) -> tuple[torch.Tensor, torch.Tensor]: - r""" - Args: - q (torch.Tensor): - queries of shape `[B, T, H, K]`. - k (torch.Tensor): - keys of shape `[B, T, H, K]`. - v (torch.Tensor): - values of shape `[B, T, H, V]`. - scale (Optional[float]): - Scale factor for the attention scores. - If not provided, it will default to `1 / sqrt(K)`. Default: `None`. - initial_state (Optional[torch.Tensor]): - Initial state of shape `[N, H, K, V]` for `N` input sequences. - For equal-length input sequences, `N` equals the batch size `B`. - Default: `None`. - output_final_state (Optional[bool]): - Whether to output the final state of shape `[N, H, K, V]`. Default: `False`. - cu_seqlens (torch.LongTensor): - Cumulative sequence lengths of shape `[N+1]` used for variable-length training, - consistent with the FlashAttention API. - - Returns: - o (torch.Tensor): - Outputs of shape `[B, T, H, V]`. - final_state (torch.Tensor): - Final state of shape `[N, H, K, V]` if `output_final_state=True` else `None`. - - """ - if 'head_first' in kwargs: - raise DeprecationWarning( - "head_first has been removed. Inputs must be in `[B, T, H, ...]` format.", - ) - chunk_size = kwargs.pop('chunk_size', None) - if cu_seqlens is not None: - if q.shape[0] != 1: - raise ValueError( - f"The batch size is expected to be 1 rather than {q.shape[0]} when using `cu_seqlens`. " - f"Please flatten variable-length inputs before processing.", - ) - if initial_state is not None and initial_state.shape[0] != len(cu_seqlens) - 1: - raise ValueError( - f"The number of initial states is expected to be equal to the number of input sequences, " - f"i.e., {len(cu_seqlens) - 1} rather than {initial_state.shape[0]}.", - ) - g_gamma = (1 - q.new_tensor(2., dtype=torch.float).pow(-5. - q.new_tensor(range(q.shape[2]), dtype=torch.float))).log() - o, final_state = chunk_simple_gla( - q=q, - k=k, - v=v, - scale=scale, - g_gamma=g_gamma, - initial_state=initial_state, - output_final_state=output_final_state, - cu_seqlens=cu_seqlens, - chunk_size=chunk_size, - ) - return o, final_state diff --git a/build/torch-xpu/ops/retention/fused_chunk.py b/build/torch-xpu/ops/retention/fused_chunk.py deleted file mode 100644 index 603b4f5188c6f9a362907734821082ebf4409937..0000000000000000000000000000000000000000 --- a/build/torch-xpu/ops/retention/fused_chunk.py +++ /dev/null @@ -1,78 +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 - -from ...ops.simple_gla import fused_chunk_simple_gla - - -@torch.compiler.disable -def fused_chunk_retention( - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - scale: float | None = None, - initial_state: torch.Tensor | None = None, - output_final_state: bool = False, - cu_seqlens: torch.LongTensor | None = None, - **kwargs, -) -> tuple[torch.Tensor, torch.Tensor]: - r""" - Args: - q (torch.Tensor): - queries of shape `[B, T, H, K]`. - k (torch.Tensor): - keys of shape `[B, T, H, K]`. - v (torch.Tensor): - values of shape `[B, T, H, V]`. - scale (Optional[float]): - Scale factor for the attention scores. - If not provided, it will default to `1 / sqrt(K)`. Default: `None`. - initial_state (Optional[torch.Tensor]): - Initial state of shape `[N, H, K, V]` for `N` input sequences. - For equal-length input sequences, `N` equals the batch size `B`. - Default: `None`. - output_final_state (Optional[bool]): - Whether to output the final state of shape `[N, H, K, V]`. Default: `False`. - cu_seqlens (torch.LongTensor): - Cumulative sequence lengths of shape `[N+1]` used for variable-length training, - consistent with the FlashAttention API. - - Returns: - o (torch.Tensor): - Outputs of shape `[B, T, H, V]`. - final_state (torch.Tensor): - Final state of shape `[N, H, K, V]` if `output_final_state=True` else `None`. - - """ - if 'head_first' in kwargs: - raise DeprecationWarning( - "head_first has been removed. Inputs must be in `[B, T, H, ...]` format.", - ) - if cu_seqlens is not None: - if q.shape[0] != 1: - raise ValueError( - f"The batch size is expected to be 1 rather than {q.shape[0]} when using `cu_seqlens`. " - f"Please flatten variable-length inputs before processing.", - ) - if initial_state is not None and initial_state.shape[0] != len(cu_seqlens) - 1: - raise ValueError( - f"The number of initial states is expected to be equal to the number of input sequences, " - f"i.e., {len(cu_seqlens) - 1} rather than {initial_state.shape[0]}.", - ) - g_gamma = (1 - q.new_tensor(2., dtype=torch.float).pow(-5. - q.new_tensor(range(q.shape[2]), dtype=torch.float))).log() - o, final_state = fused_chunk_simple_gla( - q=q, - k=k, - v=v, - g_gamma=g_gamma, - scale=scale, - initial_state=initial_state, - output_final_state=output_final_state, - cu_seqlens=cu_seqlens, - ) - return o, final_state diff --git a/build/torch-xpu/ops/retention/fused_recurrent.py b/build/torch-xpu/ops/retention/fused_recurrent.py deleted file mode 100644 index 5f141640e74ce953133fbf3378fe054bd99094b3..0000000000000000000000000000000000000000 --- a/build/torch-xpu/ops/retention/fused_recurrent.py +++ /dev/null @@ -1,61 +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 warnings - -import torch - -from ...ops.simple_gla.fused_recurrent import fused_recurrent_simple_gla - - -def fused_recurrent_retention( - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - scale: float | None = None, - initial_state: torch.Tensor | None = None, - output_final_state: bool = False, - reverse: bool = False, - state_v_first: bool = False, - cu_seqlens: torch.LongTensor | None = None, - **kwargs, -) -> tuple[torch.Tensor, torch.Tensor]: - if 'transpose_state_layout' in kwargs: - if state_v_first: - raise ValueError("Cannot pass both `state_v_first` and the deprecated `transpose_state_layout`.") - warnings.warn( - "`transpose_state_layout` is deprecated and renamed to `state_v_first`.", - DeprecationWarning, - stacklevel=2, - ) - state_v_first = kwargs.pop('transpose_state_layout') - - if cu_seqlens is not None: - if q.shape[0] != 1: - raise ValueError( - f"The batch size is expected to be 1 rather than {q.shape[0]} when using `cu_seqlens`. " - f"Please flatten variable-length inputs before processing.", - ) - if initial_state is not None and initial_state.shape[0] != len(cu_seqlens) - 1: - raise ValueError( - f"The number of initial states is expected to be equal to the number of input sequences, " - f"i.e., {len(cu_seqlens) - 1} rather than {initial_state.shape[0]}.", - ) - g_gamma = (1 - q.new_tensor(2., dtype=torch.float).pow(-5. - q.new_tensor(range(q.shape[2]), dtype=torch.float))).log() - o, final_state = fused_recurrent_simple_gla( - q=q, - k=k, - v=v, - g_gamma=g_gamma, - scale=scale, - initial_state=initial_state, - output_final_state=output_final_state, - reverse=reverse, - cu_seqlens=cu_seqlens, - state_v_first=state_v_first, - ) - return o, final_state diff --git a/build/torch-xpu/ops/retention/naive.py b/build/torch-xpu/ops/retention/naive.py deleted file mode 100644 index b81ea21f82f06eeb7e303ae1c4e6e38fa5a14115..0000000000000000000000000000000000000000 --- a/build/torch-xpu/ops/retention/naive.py +++ /dev/null @@ -1,20 +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 - - -def naive_retention(q, k, v): - orig_type = q.dtype - q, k, v = q.float(), k.float(), v.float() - _, n_heads, seq_len, d_head = q.shape - s = (1 - q.new_tensor(2., dtype=torch.float).pow(-5. - q.new_tensor(range(n_heads), dtype=torch.float))).log2() - n = q.new_tensor(range(seq_len), dtype=torch.float) - n = torch.exp2((n.unsqueeze(-1) - n) * s.view(-1, 1, 1)) * n.unsqueeze(-1).ge(n) - s = torch.einsum('bhqd,bhkd,hqk->bhqk', q * d_head ** -0.5, k, n.to(q.dtype)) - o = torch.einsum('bhqk,bhkd->bhqd', s, v) - return o.to(orig_type) diff --git a/build/torch-xpu/ops/retention/parallel.py b/build/torch-xpu/ops/retention/parallel.py deleted file mode 100644 index 2c65b409b8d82bd6cd0a37ba09ac1ae36272a3a1..0000000000000000000000000000000000000000 --- a/build/torch-xpu/ops/retention/parallel.py +++ /dev/null @@ -1,66 +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 - -from ...ops.simple_gla.parallel import parallel_simple_gla - - -def parallel_retention( - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - scale: float | None = None, - output_attentions: bool = False, - cu_seqlens: torch.LongTensor | None = None, - **kwargs, -) -> tuple[torch.Tensor, torch.Tensor]: - r""" - Args: - q (torch.Tensor): - queries of shape `[B, T, H, K]`. - k (torch.Tensor): - keys of shape `[B, T, H, K]`. - v (torch.Tensor): - values of shape `[B, T, H, V]`. - scale (Optional[float]): - Scale factor for attention scores. - If not provided, it will default to `1 / sqrt(K)`. Default: `None`. - output_attentions (bool): - Whether to output the materialized attention scores of shape `[B, H, T, T]`. Default: `False`. - cu_seqlens (torch.LongTensor): - Cumulative sequence lengths of shape `[N+1]` used for variable-length training, - consistent with the FlashAttention API. - - Returns: - o (torch.Tensor): - Outputs of shape `[B, T, H, V]`. - attn (torch.Tensor): - Attention scores of shape `[B, H, T, T]` if `output_attentions=True` else `None`. - """ - if 'head_first' in kwargs: - raise DeprecationWarning( - "head_first has been removed. Inputs must be in `[B, T, H, ...]` format.", - ) - if cu_seqlens is not None and q.shape[0] != 1: - raise ValueError( - f"The batch size is expected to be 1 rather than {q.shape[0]} when using `cu_seqlens`. " - f"Please flatten variable-length inputs before processing.", - ) - s = (1 - q.new_tensor(2., dtype=torch.float).pow(-5. - q.new_tensor(range(q.shape[2]), dtype=torch.float))).log() - g = s[None, None, :].expand(q.shape[0], q.shape[1], q.shape[2]) - - o, attn = parallel_simple_gla( - q=q, - k=k, - v=v, - scale=scale, - g=g, - output_attentions=output_attentions, - cu_seqlens=cu_seqlens, - ) - return o, attn diff --git a/build/torch-xpu/ops/rwkv4/__init__.py b/build/torch-xpu/ops/rwkv4/__init__.py deleted file mode 100644 index 5fb25ccf1b9a0dad25827856f76b1bc684673e43..0000000000000000000000000000000000000000 --- a/build/torch-xpu/ops/rwkv4/__init__.py +++ /dev/null @@ -1,12 +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 .fused_recurrent import fused_recurrent_rwkv4 - -__all__ = [ - 'fused_recurrent_rwkv4', -] diff --git a/build/torch-xpu/ops/rwkv4/fused_recurrent.py b/build/torch-xpu/ops/rwkv4/fused_recurrent.py deleted file mode 100644 index c88f1ec0d7f56a916e33ea011571086e583c9ce0..0000000000000000000000000000000000000000 --- a/build/torch-xpu/ops/rwkv4/fused_recurrent.py +++ /dev/null @@ -1,476 +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, cast - -import torch -import triton -import triton.language as tl -from torch import Tensor -from torch.autograd.function import Function, FunctionCtx, once_differentiable - -from ...ops.utils.op import exp -from ...utils import input_guard - - -def get_block_size_c(chans: int) -> int: - if chans < 32: - return 32 - if chans < 64: - return 64 - return 128 - - -@triton.jit -def fused_recurrent_rwkv4_forward_kernel( - # W - w_ptr, - w_s_c, - # U - u_ptr, - u_s_c, - # K - k_ptr, - k_s_b, - k_s_t, - k_s_c, - # V - v_ptr, - v_s_b, - v_s_t, - v_s_c, - # State - state_ptr, - state_s_b, - state_s_abe, - state_s_c, - # WKV - wkv_ptr, - wkv_s_b, - wkv_s_t, - wkv_s_c, - # Output state - state_out_ptr, - state_out_s_b, - state_out_s_abe, - state_out_s_t, - state_out_s_c, - # Params - chans, - tsz, - BLOCK_SIZE_C: tl.constexpr, -): - # Parallelize over the batch dimension. - b_idx = tl.program_id(0) - c_idx = tl.program_id(1) - - cs = (c_idx * BLOCK_SIZE_C) + tl.arange(0, BLOCK_SIZE_C) - cmask = cs < chans - - # Pointers to the batch (and possibly channel) for the input tensors. - k_ptr = k_ptr + b_idx * k_s_b - v_ptr = v_ptr + b_idx * v_s_b - alpha_ptr = state_ptr + b_idx * state_s_b - beta_ptr = state_ptr + b_idx * state_s_b + state_s_abe - eps_ptr = state_ptr + b_idx * state_s_b + 2 * state_s_abe - - # Pointers to the batch (and possibly channel) for the output tensors. - wkv_ptr = wkv_ptr + b_idx * wkv_s_b - alpha_out_ptr = state_out_ptr + b_idx * state_out_s_b - beta_out_ptr = state_out_ptr + b_idx * state_out_s_b + state_out_s_abe - eps_out_ptr = state_out_ptr + b_idx * state_out_s_b + 2 * state_out_s_abe - - # Loads parameters. - alpha = tl.load(alpha_ptr + cs * state_s_c, mask=cmask).to(tl.float32) - beta = tl.load(beta_ptr + cs * state_s_c, mask=cmask).to(tl.float32) - eps = tl.load(eps_ptr + cs * state_s_c, mask=cmask).to(tl.float32) - w = tl.load(w_ptr + cs * w_s_c, mask=cmask).to(tl.float32) - u = tl.load(u_ptr + cs * u_s_c, mask=cmask).to(tl.float32) - - for t in range(tsz): - kt = tl.load(k_ptr + t * k_s_t + cs * k_s_c, mask=cmask).to(tl.float32) - vt = tl.load(v_ptr + t * v_s_t + cs * v_s_c, mask=cmask).to(tl.float32) - - ukt = u + kt - tau = tl.maximum(ukt, eps) - e1a = exp(eps - tau) - e2a = exp(ukt - tau) - wkv = (e1a * alpha + e2a * vt) / (e1a * beta + e2a) - tl.store(wkv_ptr + t * wkv_s_t + cs * wkv_s_c, wkv, mask=cmask) - - w_eps = w + eps - eps = tl.maximum(w_eps, kt) - e1b = exp(w_eps - eps) - e2b = exp(kt - eps) - alpha = e1b * alpha + e2b * vt - beta = e1b * beta + e2b - tl.store(alpha_out_ptr + t * state_out_s_t + cs * state_out_s_c, alpha, mask=cmask) - tl.store(beta_out_ptr + t * state_out_s_t + cs * state_out_s_c, beta, mask=cmask) - tl.store(eps_out_ptr + t * state_out_s_t + cs * state_out_s_c, eps, mask=cmask) - - -def fused_recurrent_rwkv4_forward( - w: Tensor, - u: Tensor, - k: Tensor, - v: Tensor, - state: Tensor, -) -> tuple[Tensor, Tensor]: - (bsz, tsz, chans) = k.shape - - # New tensors to output. - wkvs = k.new_empty(bsz, tsz, chans) - state_out = k.new_empty(bsz, 3, tsz, chans) - - # Constants. - block_size_c = get_block_size_c(chans) - - def grid(meta: dict[str, Any]) -> tuple[int, ...]: - return (bsz, triton.cdiv(chans, meta["BLOCK_SIZE_C"])) - - fused_recurrent_rwkv4_forward_kernel[grid]( - # W - w, - w.stride(0), - # U - u, - u.stride(0), - # K - k, - k.stride(0), - k.stride(1), - k.stride(2), - # V - v, - v.stride(0), - v.stride(1), - v.stride(2), - # State - state, - state.stride(0), - state.stride(1), - state.stride(3), - # WKV - wkvs, - wkvs.stride(0), - wkvs.stride(1), - wkvs.stride(2), - # Output state - state_out, - state_out.stride(0), - state_out.stride(1), - state_out.stride(2), - state_out.stride(3), - # Params - chans, - tsz, - BLOCK_SIZE_C=block_size_c, - ) - - state_out = torch.cat((state, state_out), dim=2) - - return wkvs, state_out - - -@triton.jit -def fused_recurrent_rwkv4_backward_kernel( - # W - w_ptr, - w_s_c, - # U - u_ptr, - u_s_c, - # K - k_ptr, - k_s_b, - k_s_t, - k_s_c, - # V - v_ptr, - v_s_b, - v_s_t, - v_s_c, - # State - state_ptr, - state_s_b, - state_s_abe, - state_s_t, - state_s_c, - # WKV grad - gwkv_ptr, - gwkv_s_b, - gwkv_s_t, - gwkv_s_c, - # Output state grad - gstate_out_ptr, - gstate_out_s_b, - gstate_out_s_abe, - gstate_out_s_c, - # W grad - gw_ptr, - gw_s_b, - gw_s_c, - # U grad - gu_ptr, - gu_s_b, - gu_s_c, - # K grad - gk_ptr, - gk_s_b, - gk_s_t, - gk_s_c, - # V grad - gv_ptr, - gv_s_b, - gv_s_t, - gv_s_c, - # State grad - gstate_ptr, - gstate_s_b, - gstate_s_abe, - gstate_s_c, - # Params - tsz, - chans, - BLOCK_SIZE_C: tl.constexpr, -): - # Parallelize over the batch dimension. - b_idx = tl.program_id(0) - c_idx = tl.program_id(1) - - cs = (c_idx * BLOCK_SIZE_C) + tl.arange(0, BLOCK_SIZE_C) - cmask = cs < chans - - # Pointers to the batch (and possibly channel) for the input tensors. - k_ptr = k_ptr + b_idx * k_s_b - v_ptr = v_ptr + b_idx * v_s_b - alpha_ptr = state_ptr + b_idx * state_s_b - beta_ptr = state_ptr + b_idx * state_s_b + state_s_abe - eps_ptr = state_ptr + b_idx * state_s_b + 2 * state_s_abe - - # Pointers to the batch (and possibly channel) for the output tensors. - gk_ptr = gk_ptr + b_idx * gk_s_b - gv_ptr = gv_ptr + b_idx * gv_s_b - - # Pointers to gradients which were received by the function. - gwkv_ptr = gwkv_ptr + b_idx * gwkv_s_b - galpha_out_ptr = gstate_out_ptr + b_idx * gstate_out_s_b - gbeta_out_ptr = gstate_out_ptr + b_idx * gstate_out_s_b + gstate_out_s_abe - geps_out_ptr = gstate_out_ptr + b_idx * gstate_out_s_b + 2 * gstate_out_s_abe - - # Loads parameters. - galpha = tl.load(galpha_out_ptr + gstate_out_s_c * cs, mask=cmask).to(tl.float32) - gbeta = tl.load(gbeta_out_ptr + gstate_out_s_c * cs, mask=cmask).to(tl.float32) - geps = tl.load(geps_out_ptr + gstate_out_s_c * cs, mask=cmask).to(tl.float32) - w = tl.load(w_ptr + w_s_c * cs, mask=cmask).to(tl.float32) - u = tl.load(u_ptr + u_s_c * cs, mask=cmask).to(tl.float32) - - # Gradient accumulators. - gw = tl.zeros_like(w) - gu = tl.zeros_like(u) - - alpha_prev = tl.load(alpha_ptr + tsz * state_s_t + state_s_c * cs, mask=cmask).to(tl.float32) - beta_prev = tl.load(beta_ptr + tsz * state_s_t + state_s_c * cs, mask=cmask).to(tl.float32) - eps_prev = tl.load(eps_ptr + tsz * state_s_t + state_s_c * cs, mask=cmask).to(tl.float32) - - for t in range(tsz): - tc = tsz - t - 1 - - kt = tl.load(k_ptr + tc * k_s_t + k_s_c * cs, mask=cmask).to(tl.float32) - vt = tl.load(v_ptr + tc * v_s_t + v_s_c * cs, mask=cmask).to(tl.float32) - - alpha_curr = alpha_prev - beta_curr = beta_prev - eps_curr = eps_prev - - alpha_prev = tl.load(alpha_ptr + tc * state_s_t + state_s_c * cs, mask=cmask).to(tl.float32) - beta_prev = tl.load(beta_ptr + tc * state_s_t + state_s_c * cs, mask=cmask).to(tl.float32) - eps_prev = tl.load(eps_ptr + tc * state_s_t + state_s_c * cs, mask=cmask).to(tl.float32) - - ukt = u + kt - tau = tl.maximum(ukt, eps_prev) - e1 = exp(eps_prev - tau) - e2 = exp(ukt - tau) - - euke = exp(ukt + eps_prev - 2 * tau) - - denom = e1 * beta_prev + e2 - denom_sq = denom * denom - - gwkvt = tl.load(gwkv_ptr + tc * gwkv_s_t + gwkv_s_c * cs, mask=cmask).to(tl.float32) - - # Backpropagates wkv gradients. - guk = gwkvt * e2 * (e1 * beta_prev * vt - e1 * alpha_prev) / denom_sq - gu += guk - gk = guk - gv = gwkvt * e2 / denom - - galpha_wkv = gwkvt * e1 / denom - gbeta_wkv = -gwkvt * e1 * (e2 * vt + e1 * alpha_prev) / denom_sq - geps_wkv_denom = e1 * beta_prev + e2 - geps_wkv = gwkvt * euke * (alpha_prev - vt * beta_prev) / (geps_wkv_denom * geps_wkv_denom) - - e1 = exp(w + eps_prev - eps_curr) - e2 = exp(kt - eps_curr) - - # Backpropagates alpha gradients. - galpha_we = galpha * e1 * alpha_prev - gw += galpha_we - gk += galpha * e2 * vt - gv += galpha * e2 - geps += galpha * -alpha_curr - - # Backpropagates beta gradients. - gbeta_we = gbeta * e1 * beta_prev - gw += gbeta_we - gk += gbeta * e2 - geps += gbeta * -beta_curr - - # Backpropagates epsilon gradients. - geps_mask = w + eps_prev > kt - geps_we = tl.where(geps_mask, geps, tl.zeros_like(geps)) - gw += geps_we - gk += tl.where(geps_mask, tl.zeros_like(geps), geps) - - # Stores the gradients for k and v. - tl.store(gk_ptr + tc * gk_s_t + gk_s_c * cs, gk, mask=cmask) - tl.store(gv_ptr + tc * gv_s_t + gv_s_c * cs, gv, mask=cmask) - - # Computes new gradients for alpha and beta. - galpha = galpha * e1 + galpha_wkv - gbeta = gbeta * e1 + gbeta_wkv - geps = galpha_we + gbeta_we + geps_we + geps_wkv - - # Stores final gradients for alpha and beta. - galpha_ptr = gstate_ptr + b_idx * gstate_s_b - gbeta_ptr = gstate_ptr + b_idx * gstate_s_b + gstate_s_abe - geps_ptr = gstate_ptr + b_idx * gstate_s_b + 2 * gstate_s_abe - tl.store(galpha_ptr + gstate_s_c * cs, galpha, mask=cmask) - tl.store(gbeta_ptr + gstate_s_c * cs, gbeta, mask=cmask) - tl.store(geps_ptr + gstate_s_c * cs, geps, mask=cmask) - - # Stores final gradients for w and u. - tl.store(gw_ptr + gw_s_b * b_idx + gw_s_c * cs, gw*w, mask=cmask) - tl.store(gu_ptr + gu_s_b * b_idx + gu_s_c * cs, gu, mask=cmask) - - -def fused_recurrent_rwkv4_backward( - w: Tensor, - u: Tensor, - k: Tensor, - v: Tensor, - state: Tensor, - grad_wkv: Tensor, - grad_state: Tensor, -) -> tuple[Tensor, Tensor, Tensor, Tensor, Tensor]: - bsz, tsz, chans = k.shape - - gw = w.new_empty(bsz, chans, dtype=torch.float) # New tensors to output. - gu = u.new_empty(bsz, chans, dtype=torch.float) - gk = torch.empty_like(k) - gv = torch.empty_like(v) - gstate = k.new_empty(bsz, 3, 1, chans) - - block_size_c = get_block_size_c(chans) # Constants. - - def grid(meta: dict[str, Any]) -> tuple[int, ...]: - return (bsz, triton.cdiv(chans, meta["BLOCK_SIZE_C"])) - - fused_recurrent_rwkv4_backward_kernel[grid]( - # W - w, - w.stride(0), - # U - u, - u.stride(0), - # K - k, - k.stride(0), - k.stride(1), - k.stride(2), - # V - v, - v.stride(0), - v.stride(1), - v.stride(2), - # State - state, - state.stride(0), - state.stride(1), - state.stride(2), - state.stride(3), - # WKV grad - grad_wkv, - grad_wkv.stride(0), - grad_wkv.stride(1), - grad_wkv.stride(2), - # Output state grad - grad_state, - grad_state.stride(0), - grad_state.stride(1), - grad_state.stride(3), - # W grad - gw, - gw.stride(0), - gw.stride(1), - # U grad - gu, - gu.stride(0), - gu.stride(1), - # K grad - gk, - gk.stride(0), - gk.stride(1), - gk.stride(2), - # V grad - gv, - gv.stride(0), - gv.stride(1), - gv.stride(2), - # State grad - gstate, - gstate.stride(0), - gstate.stride(1), - gstate.stride(3), - # Params - tsz, - chans, - BLOCK_SIZE_C=block_size_c, - ) - - return gw.sum(0), gu.sum(0), gk, gv, gstate - - -class FusedRecurrentRWKV4Function(Function): - - @staticmethod - @input_guard - def forward( - ctx: FunctionCtx, - w: Tensor, - u: Tensor, - k: Tensor, - v: Tensor, - state: Tensor, - ) -> tuple[Tensor, Tensor]: - ctx.w_dtype = w.dtype - w = -torch.exp(w.float()) - wkv, state_out = fused_recurrent_rwkv4_forward(w, u, k, v, state) - ctx.save_for_backward(w, u, k, v, state_out[:, :, :-1]) - return wkv, state_out[:, :, -1:] - - @staticmethod - @once_differentiable - @input_guard - def backward(ctx: FunctionCtx, gwkv: Tensor, gstate: Tensor) -> tuple[Tensor, Tensor, Tensor, Tensor, Tensor]: - w, u, k, v, state = cast("tuple[Tensor, ...]", ctx.saved_tensors) - gw, gu, gk, gv, gstate = fused_recurrent_rwkv4_backward(w, u, k, v, state, gwkv, gstate) - return gw.to(ctx.w_dtype), gu.to(u), gk.to(k), gv.to(v), gstate.to(state) - - -def fused_recurrent_rwkv4(w: Tensor, u: Tensor, k: Tensor, v: Tensor, state: Tensor) -> tuple[Tensor, Tensor]: - return FusedRecurrentRWKV4Function.apply(w, u, k, v, state) diff --git a/build/torch-xpu/ops/rwkv6/__init__.py b/build/torch-xpu/ops/rwkv6/__init__.py deleted file mode 100644 index 7362a8b287d052b3f1867840695f21c4005f0e7d..0000000000000000000000000000000000000000 --- a/build/torch-xpu/ops/rwkv6/__init__.py +++ /dev/null @@ -1,14 +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 .chunk import chunk_rwkv6 -from .fused_recurrent import fused_recurrent_rwkv6 - -__all__ = [ - 'chunk_rwkv6', - 'fused_recurrent_rwkv6', -] diff --git a/build/torch-xpu/ops/rwkv6/chunk.py b/build/torch-xpu/ops/rwkv6/chunk.py deleted file mode 100644 index 7f930426b7d64bd68503bd807e65d9c625a0bb8d..0000000000000000000000000000000000000000 --- a/build/torch-xpu/ops/rwkv6/chunk.py +++ /dev/null @@ -1,1344 +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 -import triton.language as tl - -from ...ops.common.chunk_h import chunk_fwd_h -from ...ops.gla.chunk import chunk_gla_bwd_dA, chunk_gla_bwd_dv, chunk_gla_fwd_o_gk -from ...ops.utils import prepare_chunk_indices, prepare_chunk_offsets -from ...ops.utils.constant import RCP_LN2 -from ...ops.utils.op import exp2 -from ...utils import ( - autocast_custom_bwd, - autocast_custom_fwd, - autotune_cache_kwargs, - check_shared_mem, - input_guard, -) - -BK_LIST = [32, 64] if check_shared_mem() else [16, 32] -BV_LIST = [32, 64] if check_shared_mem() else [16, 32] - - -@triton.heuristics({ - 'HAS_SCALE': lambda args: args['scale'] is not None, - 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None, -}) -@triton.autotune( - configs=[ - triton.Config({'BS': BS}, num_warps=num_warps, num_stages=num_stages) - for BS in [16, 32, 64] - for num_warps in [4, 8, 16] - for num_stages in [2, 3, 4] - ], - key=['S', 'BT', 'HAS_SCALE'], - **autotune_cache_kwargs, -) -@triton.jit(do_not_specialize=['T']) -def chunk_rwkv6_fwd_cumsum_kernel( - s, - oi, - oe, - scale, - cu_seqlens, - chunk_indices, - T, - H: tl.constexpr, - S: tl.constexpr, - BT: tl.constexpr, - BS: tl.constexpr, - HAS_SCALE: tl.constexpr, - IS_VARLEN: tl.constexpr, -): - i_s, i_t, i_bh = tl.program_id(0), tl.program_id(1), tl.program_id(2) - i_b, i_h = i_bh // H, i_bh % H - 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.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) - T = eos - bos - else: - bos, eos = i_b * T, i_b * T + T - - o_i = tl.arange(0, BT) - m_i = tl.where(o_i[:, None] >= o_i[None, :], 1., 0.).to(tl.float32) - m_e = tl.where(o_i[:, None] > o_i[None, :], 1., 0.).to(tl.float32) - - p_s = tl.make_block_ptr(s + (bos * H + i_h) * S, (T, S), (H*S, 1), (i_t * BT, i_s * BS), (BT, BS), (1, 0)) - p_oi = tl.make_block_ptr(oi + (bos * H + i_h) * S, (T, S), (H*S, 1), (i_t * BT, i_s * BS), (BT, BS), (1, 0)) - p_oe = tl.make_block_ptr(oe + (bos * H + i_h) * S, (T, S), (H*S, 1), (i_t * BT, i_s * BS), (BT, BS), (1, 0)) - # [BT, BS] - b_s = tl.load(p_s, boundary_check=(0, 1)).to(tl.float32) - b_oi = tl.dot(m_i, b_s) - b_oe = tl.dot(m_e, b_s) - if HAS_SCALE: - # Pre-scale by RCP_LN2 so downstream kernels can use exp2 directly. - b_oi = b_oi * scale - b_oe = b_oe * scale - tl.store(p_oi, b_oi.to(p_oi.dtype.element_ty, fp_downcast_rounding="rtne"), boundary_check=(0, 1)) - tl.store(p_oe, b_oe.to(p_oe.dtype.element_ty, fp_downcast_rounding="rtne"), boundary_check=(0, 1)) - - -def chunk_rwkv6_fwd_cumsum( - g: torch.Tensor, - chunk_size: int, - scale: float | None = None, - cu_seqlens: torch.Tensor | None = None, - chunk_indices: torch.Tensor | None = None, -) -> torch.Tensor: - B, T, H, S = g.shape - BT = chunk_size - if chunk_indices is None: - chunk_indices = prepare_chunk_indices(cu_seqlens, chunk_size, cu_seqlens_cpu=None) if cu_seqlens is not None else None - NT = triton.cdiv(T, BT) if cu_seqlens is None else len(chunk_indices) - - gi, ge = torch.empty_like(g, dtype=torch.float), torch.empty_like(g, dtype=torch.float) - def grid(meta): return (triton.cdiv(meta['S'], meta['BS']), NT, B * H) - # keep cummulative normalizer in fp32 - chunk_rwkv6_fwd_cumsum_kernel[grid]( - g, - gi, - ge, - scale, - cu_seqlens, - chunk_indices, - T=T, - H=H, - S=S, - BT=BT, - ) - return gi, ge - - -@triton.heuristics({ - 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None, -}) -@triton.autotune( - configs=[ - triton.Config({'BK': BK}, num_warps=num_warps, num_stages=num_stages) - for BK in [32, 64] - for num_warps in [1, 2, 4, 8] - for num_stages in [2, 3, 4] - ], - key=['BC'], - **autotune_cache_kwargs, -) -@triton.jit(do_not_specialize=['T']) -def chunk_rwkv6_fwd_A_kernel_intra_sub_inter( - q, - k, - gi, # cumulative decay inclusive - ge, # cumulative decay exclusive - A, - cu_seqlens, - chunk_indices, - scale, - T, - H: tl.constexpr, - K: tl.constexpr, - BT: tl.constexpr, - BC: tl.constexpr, - BK: tl.constexpr, - NC: tl.constexpr, - IS_VARLEN: tl.constexpr, -): - i_t, i_c, i_bh = tl.program_id(0), tl.program_id(1), tl.program_id(2) - i_b, i_h = i_bh // H, i_bh % H - i_i, i_j = i_c // NC, i_c % NC - 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.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) - T = eos - bos - else: - bos, eos = i_b * T, i_b * T + T - - if i_t * BT + i_i * BC >= T: - return - if i_i <= i_j: - return - - m_i = i_t * BT + i_i * BC + tl.arange(0, BC) < T - - b_A = tl.zeros([BC, BC], dtype=tl.float32) - for i_k in range(tl.cdiv(K, BK)): - o_k = i_k * BK + tl.arange(0, BK) - m_k = o_k < K - - p_q = tl.make_block_ptr(q + (bos*H+i_h)*K, (T, K), (H*K, 1), (i_t * BT + i_i * BC, i_k * BK), (BC, BK), (1, 0)) - p_gq = tl.make_block_ptr(ge + (bos*H+i_h)*K, (T, K), (H*K, 1), (i_t * BT + i_i * BC, i_k * BK), (BC, BK), (1, 0)) - p_k = tl.make_block_ptr(k + (bos*H+i_h)*K, (K, T), (1, H*K), (i_k * BK, i_t * BT + i_j * BC), (BK, BC), (0, 1)) - p_gk = tl.make_block_ptr(gi + (bos*H+i_h)*K, (K, T), (1, H*K), (i_k * BK, i_t * BT + i_j * BC), (BK, BC), (0, 1)) - p_gn = gi + (bos + i_t * BT + i_i * BC - 1) * H*K + i_h * K + o_k - - # [BK,] - b_gn = tl.load(p_gn, mask=m_k, other=0) - # [BC, BK] - b_q = tl.load(p_q, boundary_check=(0, 1)) - b_gq = tl.where(m_i[:, None] & m_k, tl.load(p_gq, boundary_check=(0, 1)), float('-inf')) - b_qg = b_q * exp2(b_gq - b_gn[None, :]) * scale - # [BK, BC] - b_k = tl.load(p_k, boundary_check=(0, 1)) - b_gk = tl.load(p_gk, boundary_check=(0, 1)) - b_kg = b_k * exp2(b_gn[:, None] - b_gk) - # [BC, BC] using tf32 to improve precision here. - b_A += tl.dot(b_qg, b_kg) - - p_A = tl.make_block_ptr(A + (bos*H + i_h)*BT, (T, BT), (H*BT, 1), (i_t * BT + i_i * BC, i_j * BC), (BC, BC), (1, 0)) - tl.store(p_A, b_A.to(A.dtype.element_ty), boundary_check=(0, 1)) - - -@triton.heuristics({ - 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None, -}) -@triton.autotune( - configs=[ - triton.Config({}, num_warps=num_warps) - for num_warps in [1, 2, 4, 8] - ], - key=['BK', 'BT'], - **autotune_cache_kwargs, -) -@triton.jit(do_not_specialize=['T']) -def chunk_rwkv6_fwd_A_kernel_intra_sub_intra( - q, - k, - gi, - ge, - u, - A, - cu_seqlens, - chunk_indices, - scale, - T, - H: tl.constexpr, - K: tl.constexpr, - BT: tl.constexpr, - BC: tl.constexpr, - BK: tl.constexpr, - IS_VARLEN: tl.constexpr, -): - i_t, i_i, i_bh = tl.program_id(0), tl.program_id(1), tl.program_id(2) - i_b, i_h = i_bh // H, i_bh % H - i_j = i_i - 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.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) - T = eos - bos - else: - bos, eos = i_b * T, i_b * T + T - - if i_t * BT + i_i * BC >= T: - return - - o_i = tl.arange(0, BC) - o_k = tl.arange(0, BK) - m_k = o_k < K - m_A = (i_t * BT + i_i * BC + tl.arange(0, BC)) < T - o_A = (bos + i_t * BT + i_i * BC + tl.arange(0, BC)) * H*BT + i_h * BT + i_j * BC - p_q = tl.make_block_ptr(q + (bos * H + i_h) * K, (T, K), (H*K, 1), (i_t * BT + i_i * BC, 0), (BC, BK), (1, 0)) - p_g = tl.make_block_ptr(ge + (bos * H + i_h) * K, (T, K), (H*K, 1), (i_t * BT + i_i * BC, 0), (BC, BK), (1, 0)) - p_qj = q + (bos + i_t * BT + i_j * BC) * H*K + i_h * K + o_k - p_kj = k + (bos + i_t * BT + i_j * BC) * H*K + i_h * K + o_k - p_gk = gi + (bos + i_t * BT + i_j * BC) * H*K + i_h * K + o_k - - b_q = tl.load(p_q, boundary_check=(0, 1)) - b_g = tl.load(p_g, boundary_check=(0, 1)) - - p_u = tl.make_block_ptr(u + i_h * K, (K,), (1,), (0,), (BK,), (0,)) - b_u = tl.load(p_u, boundary_check=(0,)) - for j in range(0, min(BC, T - i_t * BT - i_i * BC)): - b_qj = tl.load(p_qj, mask=m_k, other=0).to(tl.float32) - b_kj = tl.load(p_kj, mask=m_k, other=0).to(tl.float32) - b_gk = tl.load(p_gk, mask=m_k, other=0).to(tl.float32) - b_A = tl.sum(b_q * b_kj[None, :] * exp2(b_g - b_gk[None, :]), 1) - b_A = tl.where(o_i > j, b_A * scale, 0.) - b_A = tl.where(o_i != j, b_A, tl.sum(b_qj * b_kj * b_u * scale)) - tl.store(A + o_A + j, b_A, mask=m_A) - p_qj += H*K - p_kj += H*K - p_gk += H*K - - -@triton.heuristics({ - 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None, -}) -@triton.autotune( - configs=[ - triton.Config({}, num_warps=1), - triton.Config({}, num_warps=2), - triton.Config({}, num_warps=4), - triton.Config({}, num_warps=8), - ], - key=['BC', 'BK'], - **autotune_cache_kwargs, -) -@triton.jit(do_not_specialize=['T']) -def chunk_rwkv6_fwd_A_kernel_intra_sub_intra_split( - q, - k, - gi, - ge, - u, - A, - cu_seqlens, - chunk_indices, - scale, - B: tl.constexpr, - T, - H: tl.constexpr, - K: tl.constexpr, - BT: tl.constexpr, - BC: tl.constexpr, - BK: tl.constexpr, - NC: tl.constexpr, - IS_VARLEN: tl.constexpr, -): - i_k, i_tc, i_bh = tl.program_id(0), tl.program_id(1), tl.program_id(2) - i_b, i_h = i_bh // H, i_bh % H - i_t, i_i = i_tc // NC, i_tc % NC - i_j = i_i - 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.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) - all = T - T = eos - bos - else: - bos, eos = i_b * T, i_b * T + T - all = B * T - - if i_t * BT + i_i * BC >= T: - return - - o_i = tl.arange(0, BC) - o_k = i_k * BK + tl.arange(0, BK) - m_k = o_k < K - m_A = (i_t * BT + i_i * BC + tl.arange(0, BC)) < T - - o_A = (i_k * all + bos + i_t * BT + i_i * BC + tl.arange(0, BC)) * H*BC + i_h * BC - p_q = tl.make_block_ptr(q + (bos * H + i_h) * K, (T, K), (H*K, 1), (i_t * BT + i_i * BC, i_k * BK), (BC, BK), (1, 0)) - p_g = tl.make_block_ptr(ge + (bos * H + i_h) * K, (T, K), (H*K, 1), (i_t * BT + i_i * BC, i_k * BK), (BC, BK), (1, 0)) - p_qj = q + (bos + i_t * BT + i_j * BC) * H*K + i_h * K + o_k - p_kj = k + (bos + i_t * BT + i_j * BC) * H*K + i_h * K + o_k - p_gk = gi + (bos + i_t * BT + i_j * BC) * H*K + i_h * K + o_k - - b_q = tl.load(p_q, boundary_check=(0, 1)) - b_g = tl.load(p_g, boundary_check=(0, 1)) - - p_u = tl.make_block_ptr(u + i_h * K, (K,), (1,), (i_k * BK), (BK,), (0,)) - b_u = tl.load(p_u, boundary_check=(0,)) - for j in range(0, min(BC, T - i_t * BT - i_i * BC)): - b_qj = tl.load(p_qj, mask=m_k, other=0).to(tl.float32) - b_kj = tl.load(p_kj, mask=m_k, other=0).to(tl.float32) - b_gk = tl.load(p_gk, mask=m_k, other=0).to(tl.float32) - b_A = tl.sum(b_q * b_kj[None, :] * exp2(b_g - b_gk[None, :]), 1) - b_A = tl.where(o_i > j, b_A * scale, 0.) - b_A = tl.where(o_i != j, b_A, tl.sum(b_qj * b_kj * b_u * scale)) - tl.store(A + o_A + j, b_A, mask=m_A) - p_qj += H*K - p_kj += H*K - p_gk += H*K - - -@triton.heuristics({ - 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None, -}) -@triton.autotune( - configs=[ - triton.Config({}, num_warps=1), - triton.Config({}, num_warps=2), - triton.Config({}, num_warps=4), - triton.Config({}, num_warps=8), - ], - key=['BC'], - **autotune_cache_kwargs, -) -@triton.jit(do_not_specialize=['T']) -def chunk_rwkv6_fwd_A_kernel_intra_sub_intra_merge( - A, - A2, - cu_seqlens, - chunk_indices, - T, - B: tl.constexpr, - H: tl.constexpr, - BT: tl.constexpr, - BC: tl.constexpr, - NK: tl.constexpr, - IS_VARLEN: tl.constexpr, -): - i_t, i_c, i_bh = tl.program_id(0), tl.program_id(1), tl.program_id(2) - i_b, i_h = i_bh // H, i_bh % H - 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.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) - all = T - T = eos - bos - else: - bos, eos = i_b * T, i_b * T + T - all = B * T - - if i_t * BT + i_c * BC >= T: - return - - b_A = tl.zeros([BC, BC], dtype=tl.float32) - for i_k in range(0, NK): - p_A = tl.make_block_ptr(A + (i_k*all+bos)*H*BC+i_h*BC, (T, BC), (H*BC, 1), (i_t*BT + i_c*BC, 0), (BC, BC), (1, 0)) - b_A += tl.load(p_A, boundary_check=(0, 1)) - p_A2 = tl.make_block_ptr(A2 + (bos*H+i_h)*BT, (T, BT), (H*BT, 1), (i_t * BT + i_c * BC, i_c * BC), (BC, BC), (1, 0)) - tl.store(p_A2, b_A.to(A2.dtype.element_ty), boundary_check=(0, 1)) - - -@triton.heuristics({ - 'STORE_INITIAL_STATE_GRADIENT': lambda args: args['dh0'] is not None, - 'USE_FINAL_STATE_GRADIENT': lambda args: args['dht'] is not None, - 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None, -}) -@triton.autotune( - configs=[ - triton.Config({'BK': BK, 'BV': BV}, num_warps=num_warps, num_stages=num_stages) - for BK in BK_LIST - for BV in BV_LIST - for num_warps in [1, 2, 4, 8] - for num_stages in [2, 3, 4] - ], - key=['BT'], - **autotune_cache_kwargs, -) -@triton.jit(do_not_specialize=['T']) -def chunk_rwkv6_bwd_kernel_dh( - q, - gi, - ge, - do, - dh, - dht, - dh0, - cu_seqlens, - chunk_offsets, - scale, - T, - HQ: tl.constexpr, - H: tl.constexpr, - K: tl.constexpr, - V: tl.constexpr, - BT: tl.constexpr, - BK: tl.constexpr, - BV: tl.constexpr, - NG: tl.constexpr, - STORE_INITIAL_STATE_GRADIENT: tl.constexpr, - USE_FINAL_STATE_GRADIENT: tl.constexpr, - IS_VARLEN: tl.constexpr, -): - i_k, i_v, i_nh = tl.program_id(0), tl.program_id(1), tl.program_id(2) - i_n, i_hq = i_nh // HQ, i_nh % HQ - i_h = i_hq // NG - if IS_VARLEN: - bos, eos = tl.load(cu_seqlens + i_n).to(tl.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) - T = eos - bos - NT = tl.cdiv(T, BT) - boh = tl.load(chunk_offsets + i_n).to(tl.int32) - else: - bos, eos = i_n * T, i_n * T + T - NT = tl.cdiv(T, BT) - boh = i_n * NT - - # [BK, BV] - b_dh = tl.zeros([BK, BV], dtype=tl.float32) - if USE_FINAL_STATE_GRADIENT: - p_dht = tl.make_block_ptr(dht + i_nh * K*V, (K, V), (V, 1), (i_k * BK, i_v * BV), (BK, BV), (1, 0)) - b_dh += tl.load(p_dht, boundary_check=(0, 1)).to(tl.float32) - - for i_t in range(NT - 1, -1, -1): - p_dh = tl.make_block_ptr(dh + ((boh+i_t) * H + i_h) * K*V, (K, V), (V, 1), (i_k * BK, i_v * BV), (BK, BV), (1, 0)) - tl.store(p_dh, b_dh.to(p_dh.dtype.element_ty), boundary_check=(0, 1)) - last_idx = min(i_t * BT + BT, T) - 1 - # [BK, BT] - p_q = tl.make_block_ptr(q + (bos*HQ + i_hq) * K, (K, T), (1, HQ*K), (i_k * BK, i_t * BT), (BK, BT), (0, 1)) - p_do = tl.make_block_ptr(do + (bos*HQ + i_hq) * V, (T, V), (HQ*V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) - b_q = tl.load(p_q, boundary_check=(0, 1)) - # [BT, BV] - b_do = tl.load(p_do, boundary_check=(0, 1)) - - p_gk = tl.make_block_ptr(ge + (bos*H + i_h) * K, (K, T), (1, H*K), (i_k * BK, i_t * BT), (BK, BT), (0, 1)) - p_gk_last = gi + (bos + last_idx) * H*K + i_h * K + i_k * BK + tl.arange(0, BK) - - b_gk = tl.load(p_gk, boundary_check=(0, 1)) - b_q = (b_q * exp2(b_gk) * scale).to(b_q.dtype) - b_gk_last = tl.load(p_gk_last, mask=(i_k * BK + tl.arange(0, BK) < K), other=0.) - b_dh *= exp2(b_gk_last)[:, None] - b_dh += tl.dot(b_q, b_do) - - if STORE_INITIAL_STATE_GRADIENT: - p_dh0 = tl.make_block_ptr(dh0 + i_nh * K*V, (K, V), (V, 1), (i_k * BK, i_v * BV), (BK, BV), (1, 0)) - tl.store(p_dh0, b_dh.to(p_dh0.dtype.element_ty), boundary_check=(0, 1)) - - -@triton.heuristics({ - 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None, -}) -@triton.autotune( - configs=[ - triton.Config({}, num_warps=num_warps) - for num_warps in [1, 2, 4, 8] - ], - key=['BK', 'NC', 'BT'], - **autotune_cache_kwargs, -) -@triton.jit(do_not_specialize=['T']) -def chunk_rwkv6_bwd_kernel_intra( - q, - k, - gi, - ge, - dA, - dq, - dk, - cu_seqlens, - chunk_indices, - T, - H: tl.constexpr, - K: tl.constexpr, - BT: tl.constexpr, - BC: tl.constexpr, - BK: tl.constexpr, - NC: tl.constexpr, - IS_VARLEN: tl.constexpr, -): - i_k, i_c, i_bh = tl.program_id(0), tl.program_id(1), tl.program_id(2) - i_b, i_h = i_bh // H, i_bh % H - i_t, i_i = i_c // NC, i_c % NC - 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.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) - else: - bos, eos = i_b * T, i_b * T + T - T = eos - bos - if i_t * BT + i_i * BC >= T: - return - - o_k = i_k * BK + tl.arange(0, BK) - m_k = o_k < K - - p_ge = tl.make_block_ptr(ge + (bos*H + i_h) * K, (T, K), (H*K, 1), (i_t * BT + i_i * BC, i_k * BK), (BC, BK), (1, 0)) - # [BC, BK] - b_ge = tl.load(p_ge, boundary_check=(0, 1)) - b_dq = tl.zeros([BC, BK], dtype=tl.float32) - if i_i > 0: - p_gn = gi + (bos + i_t * BT + i_i * BC - 1) * H*K + i_h*K + o_k - # [BK,] - b_gn = tl.load(p_gn, mask=m_k, other=0) - for i_j in range(0, i_i): - p_k = tl.make_block_ptr(k+(bos*H+i_h)*K, (T, K), (H*K, 1), (i_t*BT+i_j*BC, i_k * BK), (BC, BK), (1, 0)) - p_gk = tl.make_block_ptr(gi+(bos*H+i_h)*K, (T, K), (H*K, 1), (i_t*BT+i_j*BC, i_k * BK), (BC, BK), (1, 0)) - p_dA = tl.make_block_ptr(dA+(bos*H+i_h)*BT, (T, BT), (H*BT, 1), (i_t*BT+i_i*BC, i_j * BC), (BC, BC), (1, 0)) - # [BC, BK] - b_k = tl.load(p_k, boundary_check=(0, 1)) - b_gk = tl.load(p_gk, boundary_check=(0, 1)) - b_kg = b_k * exp2(b_gn[None, :] - b_gk) - # [BC, BC] - b_dA = tl.load(p_dA, boundary_check=(0, 1)) - # [BC, BK] - b_dq += tl.dot(b_dA, b_kg) - b_dq *= exp2(b_ge - b_gn[None, :]) - - o_i = tl.arange(0, BC) - m_dA = (i_t * BT + i_i * BC + tl.arange(0, BC)) < T - o_dA = bos*H*BT + (i_t * BT + i_i * BC + tl.arange(0, BC)) * H*BT + i_h * BT + i_i * BC - p_kj = k + (bos + i_t * BT + i_i * BC) * H*K + i_h * K + o_k - p_gkj = gi + (bos + i_t * BT + i_i * BC) * H*K + i_h * K + o_k - p_dq = tl.make_block_ptr(dq + (bos*H + i_h) * K, (T, K), (H*K, 1), (i_t * BT + i_i * BC, i_k * BK), (BC, BK), (1, 0)) - - for j in range(0, min(BC, T - i_t * BT - i_i * BC)): - # [BC,] - b_dA = tl.load(dA + o_dA + j, mask=m_dA, other=0) - # [BK,] - b_kj = tl.load(p_kj, mask=m_k, other=0).to(tl.float32) - b_gkj = tl.load(p_gkj, mask=m_k, other=0).to(tl.float32) - # [BC, BK] - m_i = o_i[:, None] > j - # [BC, BK] - # (SY 09/17) important to not use bf16 here to have a good precision. - b_dq += tl.where(m_i, b_dA[:, None] * b_kj[None, :] * exp2(b_ge - b_gkj[None, :]), 0.) - p_kj += H*K - p_gkj += H*K - tl.store(p_dq, b_dq.to(p_dq.dtype.element_ty), boundary_check=(0, 1)) - - tl.debug_barrier() - p_k = tl.make_block_ptr(k + (bos*H + i_h) * K, (T, K), (H*K, 1), (i_t * BT + i_i * BC, i_k * BK), (BC, BK), (1, 0)) - p_gk = tl.make_block_ptr(gi + (bos*H + i_h) * K, (T, K), (H*K, 1), (i_t * BT + i_i * BC, i_k * BK), (BC, BK), (1, 0)) - - # [BC, BK] - b_k = tl.load(p_k, boundary_check=(0, 1)) - b_gk = tl.load(p_gk, boundary_check=(0, 1)) - b_dk = tl.zeros([BC, BK], dtype=tl.float32) - - NC = min(NC, tl.cdiv(T - i_t * BT, BC)) - if i_i < NC - 1: - p_gn = gi + (bos + min(i_t * BT + i_i * BC + BC, T) - 1) * H*K + i_h*K + o_k - - # [BK,] - b_gn = tl.load(p_gn, mask=m_k, other=0) - for i_j in range(i_i + 1, NC): - m_j = (i_t * BT + i_j * BC + tl.arange(0, BC)) < T - p_q = tl.make_block_ptr(q + (bos*H+i_h)*K, (T, K), (H*K, 1), (i_t * BT + i_j * BC, i_k*BK), (BC, BK), (1, 0)) - p_gq = tl.make_block_ptr(ge + (bos*H+i_h)*K, (T, K), (H*K, 1), (i_t * BT + i_j * BC, i_k*BK), (BC, BK), (1, 0)) - p_dA = tl.make_block_ptr(dA + (bos*H+i_h)*BT, (BT, T), (1, H*BT), (i_i*BC, i_t*BT + i_j*BC), (BC, BC), (0, 1)) - # [BC, BK] - b_q = tl.load(p_q, boundary_check=(0, 1)) - b_gq = tl.where(m_j[:, None] & m_k, tl.load(p_gq, boundary_check=(0, 1)), float('-inf')) - b_qg = b_q * exp2(b_gq - b_gn[None, :]) - # [BC, BC] - b_dA = tl.load(p_dA, boundary_check=(0, 1)) - # [BC, BK] - # (SY 09/17) important to not use bf16 here to have a good precision. - b_dk += tl.dot(b_dA, b_qg) - b_dk *= exp2(b_gn[None, :] - b_gk) - o_dA = bos*H*BT + (i_t * BT + i_i * BC) * H*BT + i_h * BT + i_i * BC + tl.arange(0, BC) - p_qj = q + (bos + i_t * BT + i_i * BC) * H*K + i_h * K + o_k - p_gqj = ge + (bos + i_t * BT + i_i * BC) * H*K + i_h * K + o_k - p_dk = tl.make_block_ptr(dk + (bos*H+i_h)*K, (T, K), (H*K, 1), (i_t * BT + i_i * BC, i_k * BK), (BC, BK), (1, 0)) - for j in range(0, min(BC, T - i_t * BT - i_i * BC)): - # [BC,] - b_dA = tl.load(dA + o_dA + j * H*BT) - # [BK,] - b_qj = tl.load(p_qj, mask=m_k, other=0).to(tl.float32) - b_gqj = tl.load(p_gqj, mask=m_k, other=0).to(tl.float32) - # [BC, BK] - m_i = o_i[:, None] < j - b_dk += tl.where(m_i, b_dA[:, None] * b_qj[None, :] * exp2(b_gqj[None, :] - b_gk), 0.) - p_qj += H*K - p_gqj += H*K - tl.store(p_dk, b_dk.to(p_dk.dtype.element_ty), boundary_check=(0, 1)) - - -@triton.heuristics({ - 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None, -}) -@triton.autotune( - configs=[ - triton.Config({'BK': BK, 'BV': BV}, num_warps=num_warps) - for BK in BK_LIST - for BV in BV_LIST - for num_warps in [2, 4, 8] - ], - key=['BT'], - **autotune_cache_kwargs, -) -@triton.jit(do_not_specialize=['T']) -def chunk_rwkv6_bwd_kernel_inter( - q, - k, - v, - h, - gi, - ge, - u, - do, - dh, - dA, - dq, - dk, - dq2, - dk2, - dg, - du, - cu_seqlens, - chunk_indices, - scale, - T, - H: tl.constexpr, - K: tl.constexpr, - V: tl.constexpr, - BT: tl.constexpr, - BK: tl.constexpr, - BV: tl.constexpr, - IS_VARLEN: tl.constexpr, -): - i_k, i_t, i_bh = tl.program_id(0), tl.program_id(1), tl.program_id(2) - i_b, i_h = i_bh // H, i_bh % H - - 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.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) - T = eos - bos - NT = tl.cdiv(T, BT) - else: - NT = tl.cdiv(T, BT) - i_tg = i_b * NT + i_t - bos, eos = i_b * T, i_b * T + T - o_k = i_k * BK + tl.arange(0, BK) - m_k = o_k < K - - p_gk = tl.make_block_ptr(ge + (bos*H+i_h)*K, (T, K), (H*K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) - p_gi = tl.make_block_ptr(gi + (bos*H+i_h)*K, (T, K), (H*K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) - p_gn = gi + (bos + min(T, i_t * BT + BT)-1) * H*K + i_h * K + o_k - b_gn = tl.load(p_gn, mask=m_k, other=0) - b_dq = tl.zeros([BT, BK], dtype=tl.float32) - b_dk = tl.zeros([BT, BK], dtype=tl.float32) - b_dgk = tl.zeros([BK], dtype=tl.float32) - - for i_v in range(tl.cdiv(V, BV)): - p_v = tl.make_block_ptr(v + (bos*H + i_h) * V, (T, V), (H*V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) - p_do = tl.make_block_ptr(do + (bos*H + i_h) * V, (T, V), (H*V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) - p_h = tl.make_block_ptr(h + (i_tg * H + i_h) * K*V, (V, K), (1, V), (i_v * BV, i_k * BK), (BV, BK), (0, 1)) - p_dh = tl.make_block_ptr(dh + (i_tg * H + i_h) * K*V, (V, K), (1, V), (i_v * BV, i_k * BK), (BV, BK), (0, 1)) - # [BT, BV] - b_v = tl.load(p_v, boundary_check=(0, 1)) - b_do = tl.load(p_do, boundary_check=(0, 1)) - # [BV, BK] - b_h = tl.load(p_h, boundary_check=(0, 1)) - b_dh = tl.load(p_dh, boundary_check=(0, 1)) - # [BK] - b_dgk += tl.sum(b_h * b_dh, axis=0) - # [BT, BK] - b_dq += tl.dot(b_do, b_h.to(b_do.dtype)) - b_dk += tl.dot(b_v, b_dh.to(b_v.dtype)) - b_dgk *= exp2(b_gn) - b_dq *= scale - b_gk = tl.load(p_gk, boundary_check=(0, 1)) - b_gi = tl.load(p_gi, boundary_check=(0, 1)) - b_dq = b_dq * exp2(b_gk) - b_dk = b_dk * exp2(b_gn[None, :] - b_gi) - - o_i = tl.arange(0, BT) - p_q = tl.make_block_ptr(q + (bos*H+i_h)*K, (T, K), (H*K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) - p_k = tl.make_block_ptr(k + (bos*H+i_h)*K, (T, K), (H*K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) - p_dq = tl.make_block_ptr(dq + (bos*H+i_h)*K, (T, K), (H*K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) - p_dk = tl.make_block_ptr(dk + (bos*H+i_h)*K, (T, K), (H*K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) - p_dA_dig = dA + ((bos + i_t * BT + o_i) * H + i_h) * BT + o_i - b_q = tl.load(p_q, boundary_check=(0, 1)) - b_k = tl.load(p_k, boundary_check=(0, 1)) - b_dgk += tl.sum(b_dk * b_k, axis=0) - - b_dq += tl.load(p_dq, boundary_check=(0, 1)) - b_dk += tl.load(p_dk, boundary_check=(0, 1)) - b_dg = b_q * b_dq - b_k * b_dk - b_dg = b_dg - tl.cumsum(b_dg, axis=0) + tl.sum(b_dg, axis=0)[None, :] + b_dgk[None, :] - b_q * b_dq - # [BT,] - b_dA_dig = tl.load(p_dA_dig, mask=(i_t * BT + o_i) < T, other=0) - - p_u = tl.make_block_ptr(u + i_h * K, (K,), (1,), (i_k * BK,), (BK,), (0,)) - b_u = tl.load(p_u, boundary_check=(0,)) - # scale is already applied to b_dA_diag - b_dq += (b_dA_dig[:, None] * b_u[None, :] * b_k) - b_dk += (b_dA_dig[:, None] * b_u[None, :] * b_q) - b_du = tl.sum(b_dA_dig[:, None] * b_q * b_k, axis=0) - p_du = tl.make_block_ptr(du + (i_tg * H + i_h) * K, (K,), (1,), (i_k * BK,), (BK,), (0,)) - tl.store(p_du, b_du, boundary_check=(0,)) - - p_dq = tl.make_block_ptr(dq2 + (bos * H + i_h) * K, (T, K), (H*K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) - p_dk = tl.make_block_ptr(dk2 + (bos * H + i_h) * K, (T, K), (H*K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) - p_dg = tl.make_block_ptr(dg + (bos * H + i_h) * K, (T, K), (H*K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) - tl.store(p_dq, b_dq.to(p_dq.dtype.element_ty), boundary_check=(0, 1)) - tl.store(p_dk, b_dk.to(p_dk.dtype.element_ty), boundary_check=(0, 1)) - tl.store(p_dg, b_dg.to(p_dg.dtype.element_ty), boundary_check=(0, 1)) - - -def chunk_rwkv6_fwd_intra( - q: torch.Tensor, - k: torch.Tensor, - gi: torch.Tensor, - ge: torch.Tensor, - u: torch.Tensor, - scale: float, - cu_seqlens: torch.LongTensor | None = None, - chunk_size: int = 64, - chunk_indices: torch.LongTensor | None = None, -): - B, T, H, K = k.shape - BT = chunk_size - - if chunk_indices is None and cu_seqlens is not None: - chunk_indices = prepare_chunk_indices(cu_seqlens, chunk_size) - NT = triton.cdiv(T, BT) if cu_seqlens is None else len(chunk_indices) - BC = min(16, BT) - NC = triton.cdiv(BT, BC) - - A = q.new_empty(B, T, H, BT, dtype=torch.float) - grid = (NT, NC * NC, B * H) - chunk_rwkv6_fwd_A_kernel_intra_sub_inter[grid]( - q, - k, - gi, - ge, - A, - cu_seqlens, - chunk_indices, - scale, - T=T, - H=H, - K=K, - BT=BT, - BC=BC, - NC=NC, - ) - - grid = (NT, NC, B * H) - # load the entire [BC, K] blocks into SRAM at once - if K <= 256: - BK = max(triton.next_power_of_2(K), 16) - chunk_rwkv6_fwd_A_kernel_intra_sub_intra[grid]( - q, - k, - gi, - ge, - u, - A, - cu_seqlens, - chunk_indices, - scale, - T=T, - H=H, - K=K, - BT=BT, - BC=BC, - BK=BK, - ) - # split then merge - else: - BK = min(128, triton.next_power_of_2(K)) - NK = triton.cdiv(K, BK) - A_intra = q.new_empty(NK, B, T, H, BC, dtype=torch.float) - - grid = (NK, NT * NC, B * H) - chunk_rwkv6_fwd_A_kernel_intra_sub_intra_split[grid]( - q, - k, - gi, - ge, - u, - A_intra, - cu_seqlens, - chunk_indices, - scale, - B=B, - T=T, - H=H, - K=K, - BT=BT, - BC=BC, - BK=BK, - NC=NC, - ) - - grid = (NT, NC, B * H) - chunk_rwkv6_fwd_A_kernel_intra_sub_intra_merge[grid]( - A_intra, - A, - cu_seqlens, - chunk_indices, - B=B, - T=T, - H=H, - BT=BT, - BC=BC, - NK=NK, - ) - return A - - -def chunk_rwkv6_bwd_dh( - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - gi: torch.Tensor, - ge: torch.Tensor, - do: torch.Tensor, - h0: torch.Tensor, - dht: torch.Tensor, - scale: float, - cu_seqlens: torch.Tensor | None = None, - chunk_size: int = 64, - states_in_fp32: bool = False, - chunk_indices: torch.LongTensor | None = None, -) -> tuple[torch.Tensor, torch.Tensor]: - B, T, H, K, V = *k.shape, v.shape[-1] - HQ = q.shape[2] - BT = chunk_size - if chunk_indices is None and cu_seqlens is not None: - chunk_indices = prepare_chunk_indices(cu_seqlens, chunk_size) - if cu_seqlens is None: - N, NT, chunk_offsets = B, triton.cdiv(T, BT), None - else: - N, NT = len(cu_seqlens) - 1, len(chunk_indices) - chunk_offsets = prepare_chunk_offsets(cu_seqlens, BT) - NG = HQ // H - - dh = k.new_empty(B, NT, HQ, K, V, dtype=k.dtype if not states_in_fp32 else torch.float) - dh0 = torch.empty_like(h0, dtype=torch.float) if h0 is not None else None - - def grid(meta): return (triton.cdiv(K, meta['BK']), triton.cdiv(V, meta['BV']), N * H) - chunk_rwkv6_bwd_kernel_dh[grid]( - q=q, - gi=gi, - ge=ge, - do=do, - dh=dh, - dht=dht, - dh0=dh0, - cu_seqlens=cu_seqlens, - chunk_offsets=chunk_offsets, - scale=scale, - T=T, - HQ=HQ, - H=H, - K=K, - V=V, - BT=BT, - NG=NG, - ) - return dh, dh0 - - -def chunk_rwkv6_bwd_dqk_intra( - q: torch.Tensor, - k: torch.Tensor, - gi: torch.Tensor, - ge: torch.Tensor, - dA: torch.Tensor, - cu_seqlens: torch.LongTensor | None = None, - chunk_size: int = 64, - chunk_indices: torch.LongTensor | None = None, -): - B, T, H, K = q.shape - BT = chunk_size - BC = min(16, BT) - BK = min(64, triton.next_power_of_2(K)) - - if chunk_indices is None and cu_seqlens is not None: - chunk_indices = prepare_chunk_indices(cu_seqlens, chunk_size) - NT = triton.cdiv(T, BT) if cu_seqlens is None else len(chunk_indices) - NC = triton.cdiv(BT, BC) - NK = triton.cdiv(K, BK) - - dq = torch.empty_like(q, dtype=torch.float) - dk = torch.empty_like(k, dtype=torch.float) - grid = (NK, NT * NC, B * H) - chunk_rwkv6_bwd_kernel_intra[grid]( - q, - k, - gi, - ge, - dA, - dq, - dk, - cu_seqlens, - chunk_indices, - T=T, - H=H, - K=K, - BT=BT, - BC=BC, - BK=BK, - NC=NC, - ) - return dq, dk - - -def chunk_rwkv6_bwd_dqkgu( - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - h: torch.Tensor, - g: torch.Tensor, - gi: torch.Tensor, - ge: torch.Tensor, - u: torch.Tensor, - do: torch.Tensor, - dh: torch.Tensor, - dA: torch.Tensor, - dq: torch.Tensor, - dk: torch.Tensor, - scale: float, - cu_seqlens: torch.LongTensor | None = None, - chunk_size: int = 64, - chunk_indices: torch.LongTensor | None = None, -): - B, T, H, K, V = *k.shape, v.shape[-1] - BT = chunk_size - - if chunk_indices is None and cu_seqlens is not None: - chunk_indices = prepare_chunk_indices(cu_seqlens, chunk_size) - NT = triton.cdiv(T, BT) if cu_seqlens is None else len(chunk_indices) - - dq2 = torch.empty_like(dq) - dk2 = torch.empty_like(dk) - dg = torch.empty_like(g) - du = u.new_empty(B * NT, H, K, dtype=torch.float) - def grid(meta): return (triton.cdiv(K, meta['BK']), NT, B * H) - chunk_rwkv6_bwd_kernel_inter[grid]( - q, - k, - v, - h, - gi, - ge, - u, - do, - dh, - dA, - dq, - dk, - dq2, - dk2, - dg, - du, - cu_seqlens, - chunk_indices, - scale, - T=T, - H=H, - K=K, - V=V, - BT=BT, - ) - du = du.sum(0) - return dq2, dk2, dg, du - - -def chunk_rwkv6_fwd( - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - g: torch.Tensor, - u: torch.Tensor, - scale: float, - initial_state: torch.Tensor, - output_final_state: bool, - cu_seqlens: torch.LongTensor | None = None, - chunk_size: int = 64, - chunk_indices: torch.LongTensor | None = None, -) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: - # gi/ge are pre-scaled by RCP_LN2 inside the cumsum kernel; downstream uses exp2. - gi, ge = chunk_rwkv6_fwd_cumsum( - g, - chunk_size=chunk_size, - scale=RCP_LN2, - cu_seqlens=cu_seqlens, - chunk_indices=chunk_indices, - ) - h, ht = chunk_fwd_h( - k=k, - v=v, - g=None, - gk=gi, - gv=None, - h0=initial_state, - output_final_state=output_final_state, - cu_seqlens=cu_seqlens, - chunk_size=chunk_size, - states_in_fp32=True, - ) - A = chunk_rwkv6_fwd_intra( - q=q, - k=k, - gi=gi, - ge=ge, - u=u, - scale=scale, - cu_seqlens=cu_seqlens, - chunk_size=chunk_size, - chunk_indices=chunk_indices, - ) - - o = chunk_gla_fwd_o_gk( - q=q, - v=v, - g=ge, - A=A, - h=h, - scale=scale, - cu_seqlens=cu_seqlens, - chunk_size=chunk_size, - chunk_indices=chunk_indices, - ) - return A, h, ht, o - - -def chunk_rwkv6_bwd( - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - g: torch.Tensor, - u: torch.Tensor, - scale: float, - initial_state: torch.Tensor, - A: torch.Tensor, - do: torch.Tensor, - dht: torch.Tensor, - cu_seqlens: torch.LongTensor | None = None, - chunk_size: int = 64, - chunk_indices: torch.LongTensor | None = None, -): - gi, ge = chunk_rwkv6_fwd_cumsum( - g, - chunk_size=chunk_size, - scale=RCP_LN2, - cu_seqlens=cu_seqlens, - chunk_indices=chunk_indices, - ) - h, _ = chunk_fwd_h( - k=k, - v=v, - g=None, - gk=gi, - gv=None, - h0=initial_state, - output_final_state=False, - cu_seqlens=cu_seqlens, - chunk_size=chunk_size, - states_in_fp32=True, - ) - dh, dh0 = chunk_rwkv6_bwd_dh( - q=q, - k=k, - v=v, - gi=gi, - ge=ge, - do=do, - h0=initial_state, - dht=dht, - scale=scale, - cu_seqlens=cu_seqlens, - chunk_size=chunk_size, - states_in_fp32=True, - chunk_indices=chunk_indices, - ) - - dA = chunk_gla_bwd_dA( - v=v, - do=do, - scale=scale, - cu_seqlens=cu_seqlens, - chunk_size=chunk_size, - chunk_indices=chunk_indices, - ) - dv = chunk_gla_bwd_dv( - k=k, - g=gi, - A=A, - do=do, - dh=dh, - cu_seqlens=cu_seqlens, - chunk_size=chunk_size, - ) - dq, dk = chunk_rwkv6_bwd_dqk_intra( - q=q, - k=k, - gi=gi, - ge=ge, - dA=dA, - cu_seqlens=cu_seqlens, - chunk_size=chunk_size, - chunk_indices=chunk_indices, - ) - dq, dk, dg, du = chunk_rwkv6_bwd_dqkgu( - q=q, - k=k, - v=v, - h=h, - g=g, - gi=gi, - ge=ge, - u=u, - do=do, - dh=dh, - dA=dA, - dq=dq, - dk=dk, - scale=scale, - cu_seqlens=cu_seqlens, - chunk_size=chunk_size, - chunk_indices=chunk_indices, - ) - return dq, dk, dv, dg, du, dh0 - - -class ChunkRWKV6Function(torch.autograd.Function): - - @staticmethod - @input_guard - @autocast_custom_fwd - def forward( - ctx, - q, - k, - v, - g, - u, - scale, - initial_state, - output_final_state, - cu_seqlens, - cu_seqlens_cpu, - chunk_size: int | None = None, - ): - if chunk_size is None: - chunk_size = 64 - - chunk_indices = prepare_chunk_indices( - cu_seqlens, chunk_size, cu_seqlens_cpu=cu_seqlens_cpu) if cu_seqlens is not None else None - - A, h, ht, o = chunk_rwkv6_fwd( - q=q, - k=k, - v=v, - g=g, - u=u, - scale=scale, - initial_state=initial_state, - output_final_state=output_final_state, - cu_seqlens=cu_seqlens, - chunk_size=chunk_size, - chunk_indices=chunk_indices, - ) - - ctx.save_for_backward(q, k, v, g, initial_state, A, u, chunk_indices) - - ctx.chunk_size = chunk_size - ctx.scale = scale - ctx.cu_seqlens = cu_seqlens - return o, ht - - @staticmethod - @input_guard - @autocast_custom_bwd - def backward(ctx, do, dht): - q, k, v, g, initial_state, A, u, chunk_indices = ctx.saved_tensors - chunk_size, scale, cu_seqlens = ctx.chunk_size, ctx.scale, ctx.cu_seqlens - dq, dk, dv, dg, du, dh0 = chunk_rwkv6_bwd( - q=q, - k=k, - v=v, - g=g, - u=u, - scale=scale, - initial_state=initial_state, - A=A, - do=do, - dht=dht, - cu_seqlens=cu_seqlens, - chunk_size=chunk_size, - chunk_indices=chunk_indices, - ) - return dq.to(q), dk.to(k), dv.to(v), dg.to(g), du.to(u), None, dh0, None, None, None, None - - -@torch.compiler.disable -def chunk_rwkv6( - r: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - w: torch.Tensor, - u: torch.Tensor, - scale: float | None = None, - initial_state: torch.Tensor | None = None, - output_final_state: bool = False, - cu_seqlens: torch.LongTensor | None = None, - cu_seqlens_cpu: torch.LongTensor | None = None, - **kwargs, -) -> tuple[torch.Tensor, torch.Tensor]: - r""" - Args: - r (torch.Tensor): - queries of shape `[B, T, H, K]`. - k (torch.Tensor): - keys of shape `[B, T, H, K]`. - v (torch.Tensor): - values of shape `[B, T, H, V]`. - w (torch.Tensor): - Forget gates of shape `[B, T, H, K]` applied to keys. - u (torch.Tensor): - bonus representations of shape `[H, K]`. - scale (Optional[float]): - Scale factor for the attention scores. - If not provided, it will default to `1 / sqrt(K)`. Default: `None`. - initial_state (Optional[torch.Tensor]): - Initial state of shape `[N, H, K, V]` for `N` input sequences. - For equal-length input sequences, `N` equals the batch size `B`. - Default: `None`. - output_final_state (Optional[bool]): - Whether to output the final state of shape `[N, H, K, V]`. Default: `False`. - cu_seqlens (torch.LongTensor): - Cumulative sequence lengths of shape `[N+1]` used for variable-length training, - consistent with the FlashAttention API. - cu_seqlens_cpu (torch.LongTensor): - CPU copy of `cu_seqlens` to avoid unnecessary device synchronization. Default: `None`. - - Returns: - o (torch.Tensor): - Outputs of shape `[B, T, H, V]`. - final_state (Optional[torch.Tensor]): - Final state of shape `[N, H, K, V]` if `output_final_state=True` else `None`. - - Examples:: - >>> import torch - >>> import torch.nn.functional as F - >>> from einops import rearrange - >>> from ...ops.rwkv6 import chunk_rwkv6 - # inputs with equal lengths - >>> B, T, H, K, V = 4, 2048, 4, 512, 512 - >>> r = torch.randn(B, T, H, K, device='cuda') - >>> k = torch.randn(B, T, H, K, device='cuda') - >>> v = torch.randn(B, T, H, V, device='cuda') - >>> w = F.logsigmoid(torch.randn(B, T, H, K, device='cuda')) - >>> u = torch.randn(H, K, device='cuda') - >>> h0 = torch.randn(B, H, K, V, device='cuda') - >>> o, ht = chunk_rwkv6( - r, k, v, w, u, - initial_state=h0, - output_final_state=True - ) - # for variable-length inputs, the batch size `B` is expected to be 1 and `cu_seqlens` is required - >>> r, k, v, w = map(lambda x: rearrange(x, 'b t h d -> 1 (b t) h d'), (r, k, v, w)) - # for a batch with 4 sequences, `cu_seqlens` with 5 start/end positions are expected - >>> cu_seqlens = r.new_tensor([0, 2048, 4096, 6144, 8192], dtype=torch.long) - >>> o_var, ht_var = chunk_rwkv6( - r, k, v, w, u, - initial_state=h0, - output_final_state=True, - cu_seqlens=cu_seqlens - ) - >>> assert o.allclose(o_var.view(o.shape)) - >>> assert ht.allclose(ht_var) - """ - if 'head_first' in kwargs: - raise DeprecationWarning( - "head_first has been removed. Inputs must be in `[B, T, H, ...]` format.", - ) - chunk_size = kwargs.pop('chunk_size', None) - if chunk_size is not None and chunk_size != 2 ** (chunk_size.bit_length() - 1): - raise ValueError(f"`chunk_size` must be a power of 2, got {chunk_size}.") - if cu_seqlens is not None: - if r.shape[0] != 1: - raise ValueError( - f"The batch size is expected to be 1 rather than {r.shape[0]} when using `cu_seqlens`. " - f"Please flatten variable-length inputs before processing.", - ) - if initial_state is not None and initial_state.shape[0] != len(cu_seqlens) - 1: - raise ValueError( - f"The number of initial states is expected to be equal to the number of input sequences, " - f"i.e., {len(cu_seqlens) - 1} rather than {initial_state.shape[0]}.", - ) - if scale is None: - scale = r.shape[-1] ** -0.5 - o, final_state = ChunkRWKV6Function.apply( - r, - k, - v, - w, - u, - scale, - initial_state, - output_final_state, - cu_seqlens, - cu_seqlens_cpu, - chunk_size, - ) - return o, final_state diff --git a/build/torch-xpu/ops/rwkv6/chunk_naive.py b/build/torch-xpu/ops/rwkv6/chunk_naive.py deleted file mode 100644 index 351deb63de4b4b8c5de4c653a480c136536965f8..0000000000000000000000000000000000000000 --- a/build/torch-xpu/ops/rwkv6/chunk_naive.py +++ /dev/null @@ -1,48 +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 -from einops import rearrange - - -def naive_chunk_rwkv6( - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - w: torch.Tensor, - u: torch.Tensor, - chunk_size: int = 32, -): - assert q.shape[-2] % chunk_size == 0 - orig_dtype = q.dtype - num_chunk = q.shape[-2] // chunk_size - u = u.unsqueeze(0) - - q, k, v, w = map(lambda x: rearrange(x, 'b h (n c) d -> b h n c d', c=chunk_size).float(), (q, k, v, w)) - - w_cumsum = w.cumsum(-2) - - kw = k * (w_cumsum[..., -1, None, :] - w_cumsum).exp() - wkv = kw.transpose(-1, -2) @ v - - wkv_new = torch.zeros_like(wkv) - - for i in range(num_chunk - 1): - wkv_new[:, :, i+1] = (wkv_new[:, :, i] * w_cumsum[:, :, i, -1, :, None].exp()) + wkv[:, :, i] - - o_inter = torch.einsum('b h n d p, b h n c d -> b h n c p', wkv_new, (q * (w_cumsum - w).exp())) - - o_intra = torch.zeros_like(o_inter) - for i in range(chunk_size): - attn = (q[:, :, :, i, None] * k * (w_cumsum[:, :, :, i, None] - w[:, :, :, i, None] - w_cumsum).exp()).sum(-1) - mask = (torch.arange(0, chunk_size) < i).to(attn.device) - attn.masked_fill_(~mask, 0) - intra_inter_o = (attn.unsqueeze(-1) * v).sum(-2) - intra_intra_o = (q[:, :, :, i] * u.unsqueeze(2) * k[:, :, :, i]).sum(-1).unsqueeze(-1) * v[:, :, :, i] - o_intra[:, :, :, i] = intra_inter_o + intra_intra_o - o = o_inter + o_intra - return rearrange(o, 'b h n c d -> b h (n c) d').to(orig_dtype) diff --git a/build/torch-xpu/ops/rwkv6/fused_recurrent.py b/build/torch-xpu/ops/rwkv6/fused_recurrent.py deleted file mode 100644 index e9437cb8c090c0ce0808d94d8821d1afa2e989f7..0000000000000000000000000000000000000000 --- a/build/torch-xpu/ops/rwkv6/fused_recurrent.py +++ /dev/null @@ -1,667 +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 -import triton.language as tl - -from ...ops.utils.op import exp -from ...utils import autocast_custom_bwd, autocast_custom_fwd, autotune_cache_kwargs, input_guard - - -@triton.heuristics({ - 'USE_INITIAL_STATE': lambda args: args['h0'] is not None, - 'STORE_FINAL_STATE': lambda args: args['ht'] is not None, - 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None, -}) -@triton.autotune( - configs=[ - triton.Config({}, num_warps=num_warps) - for num_warps in [1, 2, 4, 8, 16] - ], - key=['BK', 'BV'], - **autotune_cache_kwargs, -) -@triton.jit(do_not_specialize=['T']) -def fused_recurrent_rwkv6_fwd_kernel( - q, # query [B, H, T, K]/[B, T, H, K] - k, # key [B, H, T, K]/[B, T, H, K] - v, # value [B, H, T, V]/[B, T, H, V] - w, # log gate [B, H, T]/[B, T, H] or None - u, # bonus [B, H, K] - o, # output [NK, B, H, T, V]/[NK, B, T, H, V] - h0, # initial hidden state [B, H, K, V] - ht, # final hidden state [B, H, K, V] - cu_seqlens, - scale, - T, - B: tl.constexpr, - H: tl.constexpr, - K: tl.constexpr, - V: tl.constexpr, - BK: tl.constexpr, - BV: tl.constexpr, - REVERSE: tl.constexpr, # whether to reverse the recurrence - USE_INITIAL_STATE: tl.constexpr, # whether to use initial state - STORE_FINAL_STATE: tl.constexpr, # whether to store final state - IS_VARLEN: tl.constexpr, -): - i_v, i_k, i_nh = tl.program_id(0).to(tl.int64), tl.program_id(1).to(tl.int64), tl.program_id(2).to(tl.int64) - i_n, i_h = i_nh // H, i_nh % H - if IS_VARLEN: - bos, eos = tl.load(cu_seqlens + i_n).to(tl.int64), tl.load(cu_seqlens + i_n + 1).to(tl.int64) - all = T - T = eos - bos - else: - bos, eos = i_n * T, i_n * T + T - all = B * T - - o_k = i_k * BK + tl.arange(0, BK) - o_v = i_v * BV + tl.arange(0, BV) - p_q = q + (bos + ((T-1) if REVERSE else 0)) * H*K + i_h * K + o_k - p_k = k + (bos + ((T-1) if REVERSE else 0)) * H*K + i_h * K + o_k - p_v = v + (bos + ((T-1) if REVERSE else 0)) * H*V + i_h * V + o_v - p_w = w + (bos + ((T-1) if REVERSE else 0)) * H*K + i_h * K + o_k - p_o = o + ((i_k * all + bos) + ((T-1) if REVERSE else 0)) * H*V + i_h * V + o_v - p_u = u + i_h * K + o_k - - mask_k = o_k < K - mask_v = o_v < V - mask_h = mask_k[:, None] & mask_v[None, :] - - b_u = tl.load(p_u, mask=mask_k, other=0).to(tl.float32) - - b_h = tl.zeros([BK, BV], dtype=tl.float32) - if USE_INITIAL_STATE: - p_h0 = h0 + i_nh * K*V + o_k[:, None] * V + o_v[None, :] - b_h += tl.load(p_h0, mask=mask_h, other=0).to(tl.float32) - - for _ in range(0, T): - b_q = tl.load(p_q, mask=mask_k, other=0).to(tl.float32) * scale - b_k = tl.load(p_k, mask=mask_k, other=0).to(tl.float32) - b_v = tl.load(p_v, mask=mask_v, other=0).to(tl.float32) - b_w = tl.load(p_w, mask=mask_k, other=0).to(tl.float32) - b_kv = b_k[:, None] * b_v[None, :] - b_o = tl.sum((b_h + b_kv * b_u[:, None]) * b_q[:, None], 0) - b_h = b_h * exp(b_w)[:, None] + b_kv - tl.store(p_o, b_o.to(p_o.dtype.element_ty), mask=mask_v) - p_q += (-1 if REVERSE else 1) * H*K - p_k += (-1 if REVERSE else 1) * H*K - p_v += (-1 if REVERSE else 1) * H*V - p_w += (-1 if REVERSE else 1) * H*K - p_o += (-1 if REVERSE else 1) * H*V - - if STORE_FINAL_STATE: - p_ht = ht + i_nh * K*V + o_k[:, None] * V + o_v[None, :] - tl.store(p_ht, b_h.to(p_ht.dtype.element_ty), mask=mask_h) - - -@triton.heuristics({ - 'USE_INITIAL_STATE': lambda args: args['h0'] is not None, - 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None, -}) -@triton.autotune( - configs=[ - triton.Config({}, num_warps=1), - triton.Config({}, num_warps=2), - triton.Config({}, num_warps=4), - ], - key=['BK', 'BV'], - **autotune_cache_kwargs, -) -@triton.jit(do_not_specialize=['T']) -def fused_recurrent_rwkv6_bwd_kernel_dq( - k, # key [B, H, T, V]/[B, T, H, V] - v, # value [B, H, T, V]/[B, T, H, V] - w, # log gate [B, H, T]/[B, T, H] - u, # bonus [B, H, K] - do, # gradient of output [B, H, T, V]/[B, T, H, V] - dq, # gradient of query [NV, B, H, T, K]/[NV, B, T, H, K] - dq1, # gradient of query_aux [NV, B, H, T, K]/[NV, B, T, H, K] - h0, - cu_seqlens, - scale, - T, - B: tl.constexpr, - H: tl.constexpr, - K: tl.constexpr, - V: tl.constexpr, - BK: tl.constexpr, - BV: tl.constexpr, - REVERSE: tl.constexpr, - USE_INITIAL_STATE: tl.constexpr, - IS_VARLEN: tl.constexpr, -): - i_v, i_k, i_nh = tl.program_id(0).to(tl.int64), tl.program_id(1).to(tl.int64), tl.program_id(2).to(tl.int64) - i_n, i_h = i_nh // H, i_nh % H - if IS_VARLEN: - bos, eos = tl.load(cu_seqlens + i_n).to(tl.int64), tl.load(cu_seqlens + i_n + 1).to(tl.int64) - all = T - T = eos - bos - else: - bos, eos = i_n * T, i_n * T + T - all = B * T - - o_k = i_k * BK + tl.arange(0, BK) - o_v = i_v * BV + tl.arange(0, BV) - p_k = k + (bos + ((T-1) if REVERSE else 0)) * H*K + i_h * K + o_k - p_v = v + (bos + ((T-1) if REVERSE else 0)) * H*V + i_h * V + o_v - p_w = w + (bos + ((T-1) if REVERSE else 0)) * H*K + i_h * K + o_k - p_do = do + (bos + ((T-1) if REVERSE else 0)) * H*V + i_h * V + o_v - p_dq = dq + ((i_v * all + bos) + ((T-1) if REVERSE else 0)) * H*K + i_h * K + o_k - p_dq1 = dq1 + ((i_v * all + bos) + ((T-1) if REVERSE else 0)) * H*K + i_h * K + o_k - p_u = u + i_h * K + o_k - - mask_k = o_k < K - mask_v = o_v < V - mask_h = mask_k[:, None] & mask_v[None, :] - - b_u = tl.load(p_u, mask=mask_k, other=0).to(tl.float32) - - b_h = tl.zeros([BK, BV], dtype=tl.float32) - if USE_INITIAL_STATE: - p_h0 = h0 + i_nh * K*V + o_k[:, None] * V + o_v[None, :] - b_h += tl.load(p_h0, mask=mask_h, other=0).to(tl.float32) - - for _ in range(0, T): - b_k = tl.load(p_k, mask=mask_k, other=0).to(tl.float32) - b_v = tl.load(p_v, mask=mask_v, other=0).to(tl.float32) - b_w = tl.load(p_w, mask=mask_k, other=0).to(tl.float32) - b_do = tl.load(p_do, mask=mask_v, other=0).to(tl.float32) - b_kv = b_k[:, None] * b_v[None, :] - - b_hq = b_h * b_do[None, :] - b_dq = tl.sum(b_hq + b_kv * b_u[:, None] * b_do[None, :], 1) * scale - b_dq1 = tl.sum(b_hq, 1) - b_h = b_h * exp(b_w)[:, None] - b_h += b_kv - tl.store(p_dq, b_dq.to(p_dq.dtype.element_ty), mask=mask_k) - tl.store(p_dq1, b_dq1.to(p_dq1.dtype.element_ty), mask=mask_k) - - p_k += (-1 if REVERSE else 1) * H*K - p_v += (-1 if REVERSE else 1) * H*V - p_w += (-1 if REVERSE else 1) * H*K - p_do += (-1 if REVERSE else 1) * H*V - p_dq += (-1 if REVERSE else 1) * H*K - p_dq1 += (-1 if REVERSE else 1) * H*K - - -@triton.heuristics({ - 'USE_INITIAL_STATE': lambda args: args['dh0'] is not None, - 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None, -}) -@triton.autotune( - configs=[ - triton.Config({}, num_warps=1), - triton.Config({}, num_warps=2), - triton.Config({}, num_warps=4), - ], - key=['BK', 'BV'], - **autotune_cache_kwargs, -) -@triton.jit(do_not_specialize=['T']) -def fused_recurrent_rwkv6_bwd_kernel_dkv( - q, # query [B, H, T, K]/[B, T, H, K] - k, # key [B, H, T, V]/[B, T, H, V] - v, # value [B, H, T, V]/[B, T, H, V] - w, # log gate [B, H, T]/[B, T, H] - u, # bonus [B, H, K] - do, # gradient of output [B, H, T, V]/[B, T, H, V] - dk, # gradient of key [NV, B, H, T, K]/[NK, B, T, H, K] - dk1, # gradient of key_aux [NV, B, H, T, K]/[NK, B, T, H, K] - dv, # gradient of value [NK, B, H, T, V]/[NV, B, T, H, V] - dh0, # gradient of initial hidden state [N, H, K, V] - cu_seqlens, - scale, - T, - B: tl.constexpr, - H: tl.constexpr, - K: tl.constexpr, - V: tl.constexpr, - BK: tl.constexpr, - BV: tl.constexpr, - REVERSE: tl.constexpr, - USE_INITIAL_STATE: tl.constexpr, - IS_VARLEN: tl.constexpr, -): - i_v, i_k, i_nh = tl.program_id(0).to(tl.int64), tl.program_id(1).to(tl.int64), tl.program_id(2).to(tl.int64) - i_n, i_h = i_nh // H, i_nh % H - if IS_VARLEN: - bos, eos = tl.load(cu_seqlens + i_n).to(tl.int64), tl.load(cu_seqlens + i_n + 1).to(tl.int64) - all = T - T = eos - bos - else: - bos, eos = i_n * T, i_n * T + T - all = B * T - - o_k = i_k * BK + tl.arange(0, BK) - o_v = i_v * BV + tl.arange(0, BV) - p_q = q + (bos + ((T-1) if not REVERSE else 0)) * H*K + i_h * K + o_k - p_k = k + (bos + ((T-1) if not REVERSE else 0)) * H*K + i_h * K + o_k - p_v = v + (bos + ((T-1) if not REVERSE else 0)) * H*V + i_h * V + o_v - p_w = w + (bos + ((T-1) if not REVERSE else 0)) * H*K + i_h * K + o_k - p_do = do + (bos + ((T-1) if not REVERSE else 0)) * H*V + i_h * V + o_v - p_dk = dk + ((i_v * all + bos) + ((T-1) if not REVERSE else 0)) * H*K + i_h * K + o_k - p_dk1 = dk1 + ((i_v * all + bos) + ((T-1) if not REVERSE else 0)) * H*K + i_h * K + o_k - p_dv = dv + ((i_k * all + bos) + ((T-1) if not REVERSE else 0)) * H*V + i_h * V + o_v - p_u = u + i_h * K + o_k - - mask_k = o_k < K - mask_v = o_v < V - mask_h = mask_k[:, None] & mask_v[None, :] - - b_u = tl.load(p_u, mask=mask_k, other=0).to(tl.float32) - - b_dh = tl.zeros([BK, BV], dtype=tl.float32) - for _ in range(T - 1, -1, -1): - b_q = tl.load(p_q, mask=mask_k, other=0).to(tl.float32) * scale - b_k = tl.load(p_k, mask=mask_k, other=0).to(tl.float32) - b_v = tl.load(p_v, mask=mask_v, other=0).to(tl.float32) - b_w = tl.load(p_w, mask=mask_k, other=0).to(tl.float32) - b_do = tl.load(p_do, mask=mask_v, other=0).to(tl.float32) - b_dkv = b_q[:, None] * b_do[None, :] - b_dk = tl.sum(b_dh * b_v[None, :], 1) - tl.store(p_dk1, b_dk.to(p_dk1.dtype.element_ty), mask=mask_k) - b_dk += tl.sum(b_dkv * b_u[:, None] * b_v[None, :], 1) - b_dv = tl.sum((b_dh + (b_dkv * b_u[:, None])) * b_k[:, None], 0) - - tl.store(p_dk, b_dk.to(p_dk.dtype.element_ty), mask=mask_k) - tl.store(p_dv, b_dv.to(p_dv.dtype.element_ty), mask=mask_v) - b_dh *= exp(b_w)[:, None] - b_dh += b_dkv - - p_q += (-1 if not REVERSE else 1) * H*K - p_k += (-1 if not REVERSE else 1) * H*K - p_v += (-1 if not REVERSE else 1) * H*V - p_w += (-1 if not REVERSE else 1) * H*K - p_do += (-1 if not REVERSE else 1) * H*V - p_dk += (-1 if not REVERSE else 1) * H*K - p_dk1 += (-1 if not REVERSE else 1) * H*K - p_dv += (-1 if not REVERSE else 1) * H*V - - if USE_INITIAL_STATE: - p_dh0 = dh0 + i_nh * K*V + o_k[:, None] * V + o_v[None, :] - tl.store(p_dh0, b_dh.to(p_dh0.dtype.element_ty), mask=mask_h) - - -@triton.heuristics({ - 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None, -}) -@triton.autotune( - configs=[ - triton.Config({'BT': BT, 'BK': BK}, num_warps=num_warps) - for BT in [16, 32, 64] - for BK in [32, 64] - for num_warps in [1, 2, 4, 8] - ], - key=['K'], - **autotune_cache_kwargs, -) -@triton.jit(do_not_specialize=['T']) -def fused_recurrent_rwkv6_bwd_kernel_dw( - q, - k, - dq, - dk, - dw, - cu_seqlens, - scale, - T, - H: tl.constexpr, - K: tl.constexpr, - BT: tl.constexpr, - BK: tl.constexpr, - REVERSE: tl.constexpr, - IS_VARLEN: tl.constexpr, -): - i_k, i_nh = tl.program_id(0), tl.program_id(1) - i_n, i_h = i_nh // H, i_nh % H - if IS_VARLEN: - bos, eos = tl.load(cu_seqlens + i_n).to(tl.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) - else: - bos, eos = i_n * T, i_n * T + T - T = eos - bos - NT = tl.cdiv(T, BT) - - o_i = tl.arange(0, BT) - m_i = tl.where(o_i[:, None] >= o_i[None, :], 1., 0.) if not REVERSE else tl.where(o_i[:, None] <= o_i[None, :], 1., 0.) - - b_z = tl.zeros([BK], dtype=tl.float32) - - i_t = 0 if not REVERSE else NT - 1 - for _ in range(NT): - p_q = tl.make_block_ptr(q + (bos*H + i_h) * K, (T, K), (H*K, 1), (i_t * BT + 1, i_k * BK), (BT, BK), (1, 0)) - p_k = tl.make_block_ptr(k + (bos*H + i_h) * K, (T-1, K), (H*K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) - p_dq = tl.make_block_ptr(dq + (bos*H + i_h) * K, (T, K), (H*K, 1), (i_t * BT + 1, i_k * BK), (BT, BK), (1, 0)) - p_dk = tl.make_block_ptr(dk + (bos*H + i_h) * K, (T-1, K), (H*K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) - p_dw = tl.make_block_ptr(dw + (bos*H + i_h) * K, (T, K), (H*K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) - # [BT, BK] - b_q = tl.load(p_q, boundary_check=(0, 1)).to(tl.float32) - b_dq = tl.load(p_dq, boundary_check=(0, 1)).to(tl.float32) - b_k = tl.load(p_k, boundary_check=(0, 1)).to(tl.float32) - b_dk = tl.load(p_dk, boundary_check=(0, 1)).to(tl.float32) - b_dw = (b_q * b_dq * scale) - b_k * b_dk - b_c = b_z[None, :] + tl.dot(m_i, b_dw, allow_tf32=False) - tl.store(p_dw, b_c.to(p_dw.dtype.element_ty), boundary_check=(0, 1)) - if i_t >= 0: - b_z += tl.sum(b_dw, 0) - - i_t += (1 if not REVERSE else -1) - - -def fused_recurrent_rwkv6_fwd( - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - w: torch.Tensor, - u: torch.Tensor, - scale: float | None = None, - initial_state: torch.Tensor | None = None, - output_final_state: bool = False, - reverse: bool = False, - cu_seqlens: torch.LongTensor | None = None, -): - B, T, H, K, V = *k.shape, v.shape[-1] - N = B if cu_seqlens is None else len(cu_seqlens) - 1 - BK, BV = min(triton.next_power_of_2(K), 32), min(triton.next_power_of_2(V), 32) - NK, NV = triton.cdiv(K, BK), triton.cdiv(V, BV) - - h0 = initial_state - ht = q.new_empty(N, H, K, V, dtype=torch.float) if output_final_state else None - o = q.new_empty(NK, *v.shape, dtype=torch.float) - - grid = (NV, NK, N * H) - fused_recurrent_rwkv6_fwd_kernel[grid]( - q, - k, - v, - w, - u, - o, - h0, - ht, - cu_seqlens, - scale, - T=T, - B=B, - H=H, - K=K, - V=V, - BK=BK, - BV=BV, - REVERSE=reverse, - ) - o = o.sum(0) - return o, ht - - -def fused_recurrent_rwkv6_bwd( - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - w: torch.Tensor, - u: torch.Tensor, - do: torch.Tensor, - scale: float | None = None, - initial_state: torch.Tensor | None = None, - reverse: bool = False, - cu_seqlens: torch.LongTensor | None = None, -): - B, T, H, K, V = *k.shape, v.shape[-1] - N = B if cu_seqlens is None else len(cu_seqlens) - 1 - - BK, BV = min(triton.next_power_of_2(K), 16), min(triton.next_power_of_2(V), 64) - NK, NV = triton.cdiv(K, BK), triton.cdiv(V, BV) - - dq = q.new_empty(NV, *q.shape, dtype=torch.float) - dq1 = torch.empty_like(dq) - - grid = (NV, NK, N * H) - fused_recurrent_rwkv6_bwd_kernel_dq[grid]( - k, - v, - w, - u, - do, - dq, - dq1, - initial_state, - cu_seqlens, - scale, - T=T, - B=B, - H=H, - K=K, - V=V, - BK=BK, - BV=BV, - REVERSE=reverse, - ) - dq = dq.sum(0) - dq1 = dq1.sum(0) - - BK, BV = min(triton.next_power_of_2(K), 32), min(triton.next_power_of_2(V), 32) - NK, NV = triton.cdiv(K, BK), triton.cdiv(V, BV) - - dk = q.new_empty(NV, *k.shape, dtype=torch.float) - dk1 = q.new_empty(NV, *k.shape, dtype=torch.float) - dv = q.new_empty(NK, *v.shape, dtype=torch.float) - - dh0 = torch.empty_like(initial_state) if initial_state is not None else None - grid = (NV, NK, N * H) - fused_recurrent_rwkv6_bwd_kernel_dkv[grid]( - q, - k, - v, - w, - u, - do, - dk, - dk1, - dv, - dh0, - cu_seqlens, - scale, - T=T, - B=B, - H=H, - K=K, - V=V, - BK=BK, - BV=BV, - REVERSE=reverse, - ) - dk = dk.sum(0) - dk1 = dk1.sum(0) - dv = dv.sum(0) - - dw = torch.empty_like(w) - def grid(meta): return (triton.cdiv(meta['K'], meta['BK']), N * H) - fused_recurrent_rwkv6_bwd_kernel_dw[grid]( - q, - k, - dq1, - dk1, - dw, - cu_seqlens, - scale, - T=T, - H=H, - K=K, - REVERSE=not reverse, - ) - du = (do.float() * v).sum(-1, True, dtype=torch.float) * q * k * scale - du = du.sum((0, 1)) - return dq, dk, dv, dw, du, dh0 - - -class FusedRecurrentRWKV6Function(torch.autograd.Function): - - @staticmethod - @input_guard - @autocast_custom_fwd - def forward( - ctx, - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - w: torch.Tensor, - u: torch.Tensor, - scale: float | None = None, - initial_state: torch.Tensor | None = None, - output_final_state: bool = False, - reverse: bool = False, - cu_seqlens: torch.LongTensor | None = None, - ): - o, ht = fused_recurrent_rwkv6_fwd( - q=q, - k=k, - v=v, - w=w, - u=u, - scale=scale, - initial_state=initial_state, - output_final_state=output_final_state, - reverse=reverse, - cu_seqlens=cu_seqlens, - ) - ctx.save_for_backward(q, k, v, w, u, initial_state) - ctx.scale = scale - ctx.reverse = reverse - ctx.cu_seqlens = cu_seqlens - return o.to(v), ht - - @staticmethod - @input_guard - @autocast_custom_bwd - def backward(ctx, do, dht): - q, k, v, w, u, initial_state = ctx.saved_tensors - - dq, dk, dv, dw, du, dh0 = fused_recurrent_rwkv6_bwd( - q=q, - k=k, - v=v, - w=w, - u=u, - do=do, - scale=ctx.scale, - initial_state=initial_state, - reverse=ctx.reverse, - cu_seqlens=ctx.cu_seqlens, - ) - dh0 = dh0.to(initial_state) if dh0 is not None else dh0 - return dq.to(q), dk.to(k), dv.to(v), dw.to(w), du.to(u), None, dh0, None, None, None - - -def fused_recurrent_rwkv6( - r: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - w: torch.Tensor, - u: torch.Tensor, - scale: float | None = None, - initial_state: torch.Tensor | None = None, - output_final_state: bool = False, - reverse: bool = False, - cu_seqlens: torch.LongTensor | None = None, - **kwargs, -) -> tuple[torch.Tensor, torch.Tensor]: - r""" - Args: - r (torch.Tensor): - reception of shape `[B, T, H, K]`. - Alias: q, query in linear attention. - k (torch.Tensor): - keys of shape `[B, T, H, K]`. - v (torch.Tensor): - values of shape `[B, T, H, V]`. - w (torch.Tensor): - data-dependent decays of shape `[B, T, H, K]` in log space. Alias: g. - u (torch.Tensor): - bonus of shape `[H, K]`. - scale (Optional[float]): - Scale factor for the attention scores. - If not provided, it will default to `1 / sqrt(K)`. Default: `None`. - initial_state (Optional[torch.Tensor]): - Initial state of shape `[N, H, K, V]` for `N` input sequences. - For equal-length input sequences, `N` equals the batch size `B`. - Default: `None`. - output_final_state (Optional[bool]): - Whether to output the final state of shape `[N, H, K, V]`. Default: `False`. - reverse (Optional[bool]): - If `True`, process the state passing in reverse order. Default: `False`. - cu_seqlens (torch.LongTensor): - Cumulative sequence lengths of shape `[N+1]` used for variable-length training, - consistent with the FlashAttention API. - - Returns: - o (torch.Tensor): - Outputs of shape `[B, T, H, V]`. - final_state (Optional[torch.Tensor]): - Final state of shape `[N, H, K, V]` if `output_final_state=True` else `None`. - - Examples:: - >>> import torch - >>> import torch.nn.functional as F - >>> from einops import rearrange - >>> from ...ops.rwkv6 import fused_recurrent_rwkv6 - # inputs with equal lengths - >>> B, T, H, K, V = 4, 2048, 4, 512, 512 - >>> q = torch.randn(B, T, H, K, device='cuda') - >>> k = torch.randn(B, T, H, K, device='cuda') - >>> v = torch.randn(B, T, H, V, device='cuda') - >>> g = F.logsigmoid(torch.randn(B, T, H, K, device='cuda')) - >>> u = torch.randn(H, K, device='cuda') - >>> h0 = torch.randn(B, H, K, V, device='cuda') - >>> o, ht = fused_recurrent_rwkv6( - q, k, v, g, u, - initial_state=h0, - output_final_state=True - ) - # for variable-length inputs, the batch size `B` is expected to be 1 and `cu_seqlens` is required - >>> q, k, v, g = map(lambda x: rearrange(x, 'b t h d -> 1 (b t) h d'), (q, k, v, g)) - # for a batch with 4 sequences, `cu_seqlens` with 5 start/end positions are expected - >>> cu_seqlens = q.new_tensor([0, 2048, 4096, 6144, 8192], dtype=torch.long) - >>> o_var, ht_var = fused_recurrent_rwkv6( - q, k, v, g, u, - initial_state=h0, - output_final_state=True, - cu_seqlens=cu_seqlens - ) - >>> assert o.allclose(o_var.view(o.shape)) - >>> assert ht.allclose(ht_var) - """ - if 'head_first' in kwargs: - raise DeprecationWarning( - "head_first has been removed. Inputs must be in `[B, T, H, ...]` format.", - ) - if cu_seqlens is not None: - if r.shape[0] != 1: - raise ValueError( - f"The batch size is expected to be 1 rather than {r.shape[0]} when using `cu_seqlens`. " - f"Please flatten variable-length inputs before processing.", - ) - if initial_state is not None and initial_state.shape[0] != len(cu_seqlens) - 1: - raise ValueError( - f"The number of initial states is expected to be equal to the number of input sequences, " - f"i.e., {len(cu_seqlens) - 1} rather than {initial_state.shape[0]}.", - ) - if scale is None: - scale = k.shape[-1] ** -0.5 - o, final_state = FusedRecurrentRWKV6Function.apply( - r, - k, - v, - w, - u, - scale, - initial_state, - output_final_state, - reverse, - cu_seqlens, - ) - return o, final_state diff --git a/build/torch-xpu/ops/rwkv6/recurrent_naive.py b/build/torch-xpu/ops/rwkv6/recurrent_naive.py deleted file mode 100644 index 4759c14d007a0e44389e8afbfe32988c77dcfba6..0000000000000000000000000000000000000000 --- a/build/torch-xpu/ops/rwkv6/recurrent_naive.py +++ /dev/null @@ -1,105 +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 - - -def naive_recurrent_rwkv6( - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - w: torch.Tensor, - u: torch.Tensor, - scale: float | None = None, - initial_state: torch.Tensor | None = None, - output_final_state: bool | None = False, -): - orig_dtype = q.dtype - B, H, T, K, V = *q.shape, v.shape[-1] - q, k, v, w, u = map(lambda x: x.float(), (q, k, v, w, u)) - h = torch.zeros(B, H, K, V, dtype=torch.float32, device=q.device) - o = torch.zeros_like(v) - - if scale is None: - scale = K ** -0.5 - - if initial_state is not None: - h += initial_state - - for i in range(T): - q_i = q[:, :, i, :] * scale - k_i = k[:, :, i] - v_i = v[:, :, i, :] - w_i = w[:, :, i].exp() - kv_i = k_i[..., None] * v_i[..., None, :] - o_i = (h + u[None, ..., None] * kv_i) * q_i[..., None] - o[:, :, i] = o_i.sum(-2) - h = h * w_i[..., None] + kv_i - ht = h if output_final_state else None - return o.to(orig_dtype), ht - - -@torch.no_grad -def naive_recurrent_rwkv6_bwd( - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - w: torch.Tensor, - u: torch.Tensor, - o: torch.Tensor, - do: torch.Tensor, - initial_state: torch.Tensor | None = None, -): - q, k, v, w, u, o, do = (x.to(dtype=torch.float32) for x in (q, k, v, w, u, o, do)) - B, H, T, K, V = q.shape[0], q.shape[1], q.shape[2], q.shape[3], v.shape[-1] - h = torch.zeros(B, H, K, V, dtype=torch.float32, device=q.device) - dq = torch.zeros_like(q) - dq_aux = torch.zeros_like(q) - - if initial_state is not None: - h += initial_state - - for i in range(T): - k_i = k[:, :, i] - v_i = v[:, :, i] - w_i = w[:, :, i].exp() - kv_i = k_i[..., None] * v_i[..., None, :] - h_i = (h + u[None, ..., None] * kv_i) - dq_i = (do[:, :, i, None, :] * h_i).sum(-1) - dq_aux_i = (do[:, :, i, None, :] * h).sum(-1) - dq[:, :, i] = dq_i - dq_aux[:, :, i] = dq_aux_i - h = h * w_i[..., None] + kv_i - - du = torch.zeros_like(u) - dh = torch.zeros_like(h) - dk = torch.zeros_like(k) - dk_aux = torch.zeros_like(k) - dv = torch.zeros_like(v) - - for i in range(T - 1, -1, -1): - d_kv_i = do[:, :, i, None, :] * q[:, :, i, :, None] - k_i = k[:, :, i] - v_i = v[:, :, i] - du_i = (d_kv_i * k_i[..., None] * v_i[..., None, :]).sum(-1) - du += du_i.sum(0) - dk_i = (dh * v_i[..., None, :]).sum(-1) - dk_aux[:, :, i] = dk_i - dk_i += (d_kv_i * u[None, ..., None] * v_i[..., None, :]).sum(-1) - dv_i = (d_kv_i * u[None, ..., None] * k_i[..., None]).sum(-2) - dv_i += (dh * k_i[..., None]).sum(-2) - - dk[:, :, i] = dk_i - dv[:, :, i] = dv_i - dh = dh * w[:, :, i, :, None].exp() + d_kv_i - - # dw = q * dq_aux - k * dk_aux - dw = torch.zeros_like(w) - for i in range(T - 2, -1, -1): - dw[:, :, i] = dw[:, :, i+1] + dq_aux[:, :, i+1] * q[:, :, i+1] - dk_aux[:, :, i] * k[:, :, i] - - return dq, dk, dv, dw, du, dh diff --git a/build/torch-xpu/ops/rwkv7/__init__.py b/build/torch-xpu/ops/rwkv7/__init__.py deleted file mode 100644 index 9c31931cbe98b31f014bc42db4b406871639a4a4..0000000000000000000000000000000000000000 --- a/build/torch-xpu/ops/rwkv7/__init__.py +++ /dev/null @@ -1,15 +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 .chunk import chunk_rwkv7 -from .fused_recurrent import fused_mul_recurrent_rwkv7, fused_recurrent_rwkv7 - -__all__ = [ - 'chunk_rwkv7', - 'fused_mul_recurrent_rwkv7', - 'fused_recurrent_rwkv7', -] diff --git a/build/torch-xpu/ops/rwkv7/channel_mixing.py b/build/torch-xpu/ops/rwkv7/channel_mixing.py deleted file mode 100644 index 3facea3459a8f797ef2ec40fba9a7085fc874534..0000000000000000000000000000000000000000 --- a/build/torch-xpu/ops/rwkv7/channel_mixing.py +++ /dev/null @@ -1,336 +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 logging - -import torch -import triton -import triton.language as tl - -from ...utils import ( - autocast_custom_bwd, - autocast_custom_fwd, - autotune_cache_kwargs, - check_pytorch_version, - input_guard, -) - -logger = logging.getLogger(__name__) - -if not check_pytorch_version('2.4'): - logger.warning('PyTorch < 2.4 detected - computations may be slower due to lack of optimizations') - - -@triton.autotune( - configs=[ - triton.Config({'BLOCK_SIZE': block_size}) - for block_size in [128, 256, 512, 1024, 2048, 4096, 8192] - ], - key=['hidden_dim'], - **autotune_cache_kwargs, -) -@triton.jit -def rwkv_seq_mix_kernel( - x_ptr, - x_prev_ptr, - mix_k_ptr, - output_ptr, - batch_size: tl.constexpr, - token_length, - hidden_dim: tl.constexpr, - BLOCK_SIZE: tl.constexpr, -): - block_start = tl.program_id(0) * BLOCK_SIZE - block_idx = block_start + tl.arange(0, BLOCK_SIZE)[:] - - total_seq_dim = token_length * hidden_dim - batch_idx = block_idx // total_seq_dim - seq_and_feat = block_idx % total_seq_dim - seq_idx = seq_and_feat // hidden_dim - feat_idx = seq_and_feat % hidden_dim - - is_valid = (batch_idx < batch_size) & (seq_idx < token_length) - - x_idx = batch_idx * total_seq_dim + seq_idx * hidden_dim + feat_idx - - curr_x = tl.load(x_ptr + x_idx, mask=is_valid, other=0.0).to(tl.float32) - k_value = tl.load(mix_k_ptr + feat_idx).to(tl.float32) - - is_first = seq_idx < 1 - prev_state_idx = batch_idx * hidden_dim + feat_idx - prev_state = tl.load(x_prev_ptr + prev_state_idx, - mask=(is_first & is_valid), - other=0.0).to(tl.float32) - - prev_x_idx = x_idx - hidden_dim - prev_x = tl.load(x_ptr + prev_x_idx, - mask=(~is_first & is_valid), - other=0.0).to(tl.float32) - - prev_value = tl.where(is_first, prev_state, prev_x) - state_diff = prev_value - curr_x - mixed = state_diff * k_value - result = tl.cast(curr_x + mixed, dtype=output_ptr.dtype.element_ty, fp_downcast_rounding='rtne') - tl.store(output_ptr + x_idx, result, mask=is_valid) - - -@triton.jit -def rwkv_channel_mixing_pow_and_relu( - in_ptr, - out_ptr, - BLOCK_SIZE: tl.constexpr, -): - """Fused ReLU and Power operation: x = ReLU(x)^2""" - xoffset = tl.program_id(0) * BLOCK_SIZE - xindex = xoffset + tl.arange(0, BLOCK_SIZE) - x0 = xindex - x = tl.load(in_ptr + (x0), None) - x = tl.maximum(x, 0.0).to(tl.float32) - x = tl.cast(x * x, dtype=out_ptr.dtype.element_ty, fp_downcast_rounding='rtne') - tl.store(out_ptr + (x0), x, None) - - -def rwkv_mix_torch(x: torch.Tensor, x_prev: torch.Tensor, x_k: torch.Tensor): - if x_prev.dim() == 2: - x_prev = x_prev.unsqueeze(1) # (batch_size, 1, hidden_dim) - xx = torch.cat((x_prev, x[:, :-1, :]), dim=1) - x - k = x.addcmul(xx, x_k) - return k - - -def rwkv_relu_and_square_torch(x: torch.Tensor): - return torch.relu(x) ** 2 - - -def rwkv_mix_fwd(x, x_prev, x_k): - has_batch = x.dim() == 3 - - if has_batch: - batch_size, token_length, hidden_dim = x.shape - else: - token_length, hidden_dim = x.shape - batch_size = 1 - x = x.unsqueeze(0) - x_prev = x_prev.unsqueeze(0) - - token_length = x.shape[1] - hidden_dim = x.shape[2] - total_elements = batch_size * token_length * hidden_dim - - output = torch.empty_like(x) - - def grid(meta): return ( - (total_elements + meta['BLOCK_SIZE'] - 1) // meta['BLOCK_SIZE'], # grid_0 - 1, # grid_1 - 1, # grid_2 - ) - - rwkv_seq_mix_kernel[grid]( - x.contiguous(), - x_prev.contiguous(), - x_k.squeeze(), - output, - batch_size=batch_size, - token_length=token_length, - hidden_dim=hidden_dim, - ) - if not has_batch: - output = output.squeeze(0) - return output - - -def rwkv_relu_and_square_fwd(x: torch.Tensor, inplace: bool = True): - """ - Triton implementation of RWKV's ReLU and square operation - Args: - x: Input tensor - Returns: - Tensor after ReLU and square operations - """ - x = x.contiguous() - output = x if inplace else torch.empty_like(x) - - def grid(meta): return ( - (output.numel() + meta['BLOCK_SIZE'] - 1) // meta['BLOCK_SIZE'], # grid_0 - 1, # grid_1 - 1, # grid_2 - ) - rwkv_channel_mixing_pow_and_relu[grid]( - x, - output, - BLOCK_SIZE=4096, - ) - - return output - - -@triton.jit -def relu_square_bwd_kernel( - out_ptr, - forward_input_ptr, - BLOCK_SIZE: tl.constexpr, -): - """ReLU(x)^2 backward kernel - grad_input = grad_output * 2 * x if x > 0 else 0 - """ - pid = tl.program_id(0) - block_start = pid * BLOCK_SIZE - offsets = block_start + tl.arange(0, BLOCK_SIZE) - - x = tl.load(forward_input_ptr + offsets).to(tl.float32) - grad = tl.load(out_ptr + offsets).to(tl.float32) - - x = tl.maximum(x, 0.0) - - grad_input = grad * 2 * x - - tl.store(out_ptr + offsets, grad_input.to(out_ptr.dtype.element_ty)) - - -@triton.autotune( - configs=[ - triton.Config({'BLOCK_SIZE': block_size}) - for block_size in [128, 256, 512, 1024, 2048, 4096, 8192] - ], - key=['hidden_dim'], - **autotune_cache_kwargs, -) -@triton.jit -def rwkv_mix_bwd_kenel( - dk1_ptr0, - xk_ptr, - dx_ptr, - dx_prev_ptr, - batch_size, - token_length, - hidden_dim: tl.constexpr, - BLOCK_SIZE: tl.constexpr, -): - pid = tl.program_id(0) - offsets = pid * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE) - - batch_idx = offsets // (token_length * hidden_dim) - seq_feat = offsets % (token_length * hidden_dim) - seq_idx = seq_feat // hidden_dim - feat_idx = seq_feat % hidden_dim - - is_valid = offsets < (batch_size * token_length * hidden_dim) - - dk1 = tl.load(dk1_ptr0 + offsets, mask=is_valid) - xk = tl.load(xk_ptr + feat_idx, mask=is_valid) - prod = dk1 * xk - - mask_next = seq_idx < (token_length - 1) - next_offset = offsets + hidden_dim - dk1_next = tl.load(dk1_ptr0 + next_offset, mask=mask_next & is_valid, other=0.0) - prod_next = dk1_next * xk - dx_val = dk1 - prod + tl.where(mask_next, prod_next, 0.0) - dx_val = tl.cast(dx_val, dtype=dx_ptr.dtype.element_ty, fp_downcast_rounding='rtne') - tl.store(dx_ptr + offsets, dx_val, mask=is_valid) - - dx_prev_offset = batch_idx * hidden_dim + feat_idx - is_first_step = seq_idx == 0 - - tl.store( - dx_prev_ptr + dx_prev_offset, - tl.cast(prod, dtype=dx_prev_ptr.dtype.element_ty), - mask=is_first_step, - ) - - -@torch.compile(fullgraph=True) -def compute_x_k_grad(dk1, x, x_prev): - """ - Args: - dk1: (batch*seq_len, hidden_dim) - x: (batch, seq_len, hidden_dim) - x_prev: (batch, hidden_dim) or (batch, 1, hidden_dim) - """ - - if x_prev.dim() == 2: - x_prev = x_prev.unsqueeze(1) # (batch, 1, hidden_dim) - xx = torch.cat((x_prev, x[:, :-1, :]), dim=1) - x # (batch, seq_len, hidden_dim) - - # (hidden_dim,) --> (1, 1, hidden_dim) - grad_x_k = (dk1 * xx.reshape(-1, x.shape[2])).sum(dim=0).view(1, 1, -1) - return grad_x_k - - -def rwkv_channel_mixing_bwd(grad_output, x, x_prev, x_k, key_weight, value_weight, k1, k1_K, k, inplace=True): - batch_size = x.shape[0] if x.dim() == 3 else 1 - seq_len, n_embd = x.shape[-2], x.shape[-1] - - dV = k.transpose(-2, -1) @ grad_output - dk = grad_output @ value_weight.transpose(-2, -1) - - BLOCK_SIZE = 4096 - grid = ((dk.numel() + BLOCK_SIZE - 1) // BLOCK_SIZE,) - relu_square_bwd_kernel[grid]( - dk, - k1_K, - BLOCK_SIZE=BLOCK_SIZE, - ) - - dK = k1.transpose(-2, -1) @ dk - dk1 = dk @ key_weight.transpose(-2, -1) - dk1 = dk1.view(-1, n_embd).contiguous() - - dk_reduced = compute_x_k_grad(dk1, x, x_prev) - dx_prev = torch.empty_like(x_prev) if not inplace else x_prev - dx = torch.empty_like(x) if not inplace else x - - def grid(meta): return ((batch_size * seq_len * n_embd + meta['BLOCK_SIZE'] - 1) // meta['BLOCK_SIZE'], 1, 1) - rwkv_mix_bwd_kenel[grid]( - dk1, - x_k.squeeze(), - dx, - dx_prev, - batch_size, - seq_len, - n_embd, - ) - # dx_prev.shape batch_size, seq_len, n_embd - return dx, dx_prev, dk_reduced, dK, dV - - -class Rwkv7ChannelMixing(torch.autograd.Function): - @staticmethod - @input_guard - @autocast_custom_fwd - def forward(ctx, x, x_prev, x_k, key_weight, value_weight, inplace: bool = True): - k1 = rwkv_mix_fwd(x, x_prev, x_k) - k1_K = k1 @ key_weight - k = rwkv_relu_and_square_fwd(k1_K, inplace=True) - ctx.save_for_backward(x, x_prev, x_k, key_weight, value_weight) - ctx.inplace = inplace - return k @ value_weight - - @staticmethod - @input_guard - @autocast_custom_bwd - def backward(ctx, dkv): - x, x_prev, x_k, key_weight, value_weight = ctx.saved_tensors - k1 = rwkv_mix_fwd(x, x_prev, x_k) - k1_K = k1 @ key_weight - k = rwkv_relu_and_square_fwd(k1_K, inplace=False) - dx, dx_prev, dk_reduced, dK, dV = rwkv_channel_mixing_bwd( - dkv, x, x_prev, x_k, key_weight, value_weight, k1, k1_K, k, ctx.inplace) - return dx, dx_prev, dk_reduced, dK, dV, None - - -def channel_mixing_rwkv7(x: torch.Tensor, x_prev: torch.Tensor, x_k: torch.Tensor, - key_weight: torch.Tensor, value_weight: torch.Tensor, inplace: bool = True): - assert x.dim() == 3 - - return Rwkv7ChannelMixing.apply(x, x_prev, x_k, key_weight, value_weight, inplace), x[-1, :] - - -def channel_mixing_rwkv7_torch(x, x_prev, x_k, key_weight, value_weight): - k1 = rwkv_mix_torch(x, x_prev, x_k) - k1_K = k1 @ key_weight - k = rwkv_relu_and_square_torch(k1_K) - return k @ value_weight, x[-1, :] diff --git a/build/torch-xpu/ops/rwkv7/chunk.py b/build/torch-xpu/ops/rwkv7/chunk.py deleted file mode 100644 index 598472ea5c3070ee87f947f75199d65e326c0fa9..0000000000000000000000000000000000000000 --- a/build/torch-xpu/ops/rwkv7/chunk.py +++ /dev/null @@ -1,81 +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 - -from ...ops.generalized_delta_rule import chunk_dplr_delta_rule - - -def chunk_rwkv7( - r: torch.Tensor, - w: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - a: torch.Tensor, - b: torch.Tensor, - scale: float = 1.0, - initial_state: torch.Tensor | None = None, - output_final_state: bool = False, - cu_seqlens: torch.LongTensor | None = None, - cu_seqlens_cpu: torch.LongTensor | None = None, - safe_gate: bool = False, - chunk_size: int | None = None, - **kwargs, -): - """ - Args: - r (torch.Tensor): - r of shape `[B, T, H, K]`. - w (torch.Tensor): - log decay of shape `[B, T, H, K]`. - k (torch.Tensor): - k of shape `[B, T, H, K]`. - v (torch.Tensor): - v of shape `[B, T, H, V]`. - a (torch.Tensor): - a of shape `[B, T, H, K]`. - b (torch.Tensor): - b of shape `[B, T, H, K]`. - scale (float): - Scale factor for the attention scores. Default: `1.0`. - initial_state (Optional[torch.Tensor]): - Initial state of shape `[N, H, K, V]` for `N` input sequences. - For equal-length input sequences, `N` equals the batch size `B`. - Default: `None`. - output_final_state (Optional[bool]): - Whether to output the final state of shape `[N, H, K, V]`. Default: `False`. - cu_seqlens (torch.LongTensor): - Cumulative sequence lengths of shape `[N+1]` used for variable-length training, - consistent with the FlashAttention API. - cu_seqlens_cpu (torch.LongTensor): - CPU copy of `cu_seqlens` to avoid unnecessary device synchronization. Default: `None`. - safe_gate (Optional[bool]): - Whether the kernel can assume the input gate values `g` are in a safe range. - When `True`, the kernel can use M=16 TensorCore acceleration. - The safe range is approximately `[-5, 0)`. Default: `False`. - chunk_size (Optional[int]): - Chunk size for the chunked computation. Default: `None`, which means 16. - """ - if 'head_first' in kwargs: - raise DeprecationWarning( - "head_first has been removed. Inputs must be in `[B, T, H, ...]` format.", - ) - return chunk_dplr_delta_rule( - q=r, - k=k, - v=v, - a=a, - b=b, - gk=w, - scale=scale, - initial_state=initial_state, - output_final_state=output_final_state, - cu_seqlens=cu_seqlens, - cu_seqlens_cpu=cu_seqlens_cpu, - safe_gate=safe_gate, - chunk_size=chunk_size, - ) diff --git a/build/torch-xpu/ops/rwkv7/fused_addcmul.py b/build/torch-xpu/ops/rwkv7/fused_addcmul.py deleted file mode 100644 index ba18ba322282cd191b913c623e15d2011b4e7e4f..0000000000000000000000000000000000000000 --- a/build/torch-xpu/ops/rwkv7/fused_addcmul.py +++ /dev/null @@ -1,299 +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 logging -import os -import sys - -import torch -import triton -import triton.language as tl -from packaging.version import Version - -from ...utils import IS_AMD, IS_ARM, autotune_cache_kwargs, check_pytorch_version, input_guard - -logger = logging.getLogger(__name__) - -if not check_pytorch_version('2.4'): - logger.warning('PyTorch < 2.4 detected - computations may be slower due to lack of optimizations') - - -def identity_decorator(fn): - return fn - - -current_python_version = Version(f"{sys.version_info.major}.{sys.version_info.minor}.{sys.version_info.micro}") -min_torch_compile_version = Version("3.11" if IS_ARM else "3.10") -fla_use_compile = os.getenv('FLA_USE_COMPILE', '1').lower() in ('1', 'true', 'yes') - -if current_python_version >= min_torch_compile_version and fla_use_compile: - torch_compile = torch.compile(fullgraph=True) -else: - if fla_use_compile: - logger.warning(f'torch.compile requires Python >= {min_torch_compile_version}, using identity decorator instead') - torch_compile = identity_decorator - -NUM_WARPS_AUTOTUNE = [2, 4, 8, 16] if IS_AMD else [2, 4, 8, 16, 32] - - -@triton.autotune( - configs=[ - triton.Config({'BT': BT}, num_warps=num_warps, num_stages=num_stages) - for num_warps in NUM_WARPS_AUTOTUNE - for num_stages in [1, 2, 3] - for BT in [2, 4, 8] - ], - key=['BD'], - **autotune_cache_kwargs, -) -@triton.jit -def fused_addcmul_fwd_kernel( - hidden, - delta, - ixr, ixw, ixk, ixv, ixa, ixg, - oxr, oxw, oxk, oxv, oxa, oxg, - use_xg: tl.constexpr, - T, - T_OFFSET, - BT: tl.constexpr, - D: tl.constexpr, - BD: tl.constexpr, -): - i_b, i_t = tl.program_id(0), tl.program_id(1) * BT - - bos = i_b * (T + T_OFFSET) - t_vec = i_t + T_OFFSET + tl.arange(0, BT) - mask_t = t_vec < (T + T_OFFSET) - o_d = tl.arange(0, BD)[None, :] - off_vec = (bos + t_vec)[:, None] * D + o_d - m_d = o_d < D - mask = mask_t[:, None] & m_d - - b_h = tl.load(hidden + off_vec, mask=mask, other=0.) - b_x = tl.load(delta + off_vec, mask=mask, other=0.) - b_r = tl.load(ixr + o_d, mask=m_d) - b_w = tl.load(ixw + o_d, mask=m_d) - b_k = tl.load(ixk + o_d, mask=m_d) - b_v = tl.load(ixv + o_d, mask=m_d) - b_a = tl.load(ixa + o_d, mask=m_d) - - o_r = tl.fma(b_x, b_r, b_h) - o_w = tl.fma(b_x, b_w, b_h) - o_k = tl.fma(b_x, b_k, b_h) - o_v = tl.fma(b_x, b_v, b_h) - o_a = tl.fma(b_x, b_a, b_h) - - tl.store(oxr + off_vec, o_r.to(oxr.dtype.element_ty), mask=mask) - tl.store(oxw + off_vec, o_w.to(oxw.dtype.element_ty), mask=mask) - tl.store(oxk + off_vec, o_k.to(oxk.dtype.element_ty), mask=mask) - tl.store(oxv + off_vec, o_v.to(oxv.dtype.element_ty), mask=mask) - tl.store(oxa + off_vec, o_a.to(oxa.dtype.element_ty), mask=mask) - - if use_xg: - b_g = tl.load(ixg + o_d, mask=m_d) - o_g = tl.fma(b_x, b_g, b_h) - tl.store(oxg + off_vec, o_g.to(oxg.dtype.element_ty), mask=mask) - - -@triton.autotune( - configs=[ - triton.Config({'BT': BT}, num_warps=num_warps, num_stages=num_stages) - for num_warps in NUM_WARPS_AUTOTUNE - for num_stages in [1, 2, 3] - for BT in [2, 4, 8] - ], - key=['BD'], - **autotune_cache_kwargs, -) -@triton.jit -def addcmul_bwd_kernel1( - ixr, - ixw, - ixk, - ixv, - ixa, - ixg, - dxr, - dxw, - dxk, - dxv, - dxa, - dxg, - ghidden, - gx, - use_xg: tl.constexpr, - T, - T_OFFSET, - BT: tl.constexpr, - D: tl.constexpr, - BD: tl.constexpr, - DTYPE: tl.constexpr, -): - i_b, i_t = tl.program_id(0), tl.program_id(1) - - t_idx = T_OFFSET + i_t * BT + tl.arange(0, BT)[:, None] - mask_t = t_idx < (T + T_OFFSET) - - d_idx = tl.arange(0, BD)[None, :] - mask_d = d_idx < D - mask = mask_t & mask_d - - offset_base = i_b * (T + T_OFFSET) * D - x_idx = (offset_base + t_idx * D + d_idx).to(tl.uint32) - - b_dxr = tl.load(dxr + x_idx, mask=mask).to(DTYPE) - b_dxw = tl.load(dxw + x_idx, mask=mask).to(DTYPE) - b_dxk = tl.load(dxk + x_idx, mask=mask).to(DTYPE) - b_dxv = tl.load(dxv + x_idx, mask=mask).to(DTYPE) - b_dxa = tl.load(dxa + x_idx, mask=mask).to(DTYPE) - - b_ixr = tl.load(ixr + d_idx, mask=mask_d).to(DTYPE) - b_ixw = tl.load(ixw + d_idx, mask=mask_d).to(DTYPE) - b_ixk = tl.load(ixk + d_idx, mask=mask_d).to(DTYPE) - b_ixv = tl.load(ixv + d_idx, mask=mask_d).to(DTYPE) - b_ixa = tl.load(ixa + d_idx, mask=mask_d).to(DTYPE) - - g_hidden = b_dxr + b_dxw + b_dxk + b_dxv + b_dxa - g_x = b_dxr * b_ixr + b_dxw * b_ixw + b_dxk * b_ixk + b_dxv * b_ixv + b_dxa * b_ixa - - if use_xg: - b_dxg = tl.load(dxg + x_idx, mask=mask).to(DTYPE) - b_ixg = tl.load(ixg + d_idx, mask=mask_d).to(DTYPE) - g_hidden += b_dxg - g_x += b_dxg * b_ixg - - tl.store(ghidden + x_idx, g_hidden.to(ghidden.dtype.element_ty), mask=mask) - tl.store(gx + x_idx, g_x.to(gx.dtype.element_ty), mask=mask) - - -def addcmul_bwd1(d_xr, d_xw, d_xk, d_xv, d_xa, d_xg, - x_r, x_w, x_k, x_v, x_a, x_g, hidden_states, delta, use_xg, inplace=True): - B, T, D = hidden_states.size() - g_hiddn = hidden_states if inplace else torch.empty_like(hidden_states) - g_delta = torch.empty_like(delta) - for t in range(0, T, 65536): - T_OFFSET = t - T_SIZE = min(65536, T - t) - def grid(meta): return (B, triton.cdiv(T_SIZE, meta['BT'])) - addcmul_bwd_kernel1[grid]( - ixr=x_r, - ixw=x_w, - ixk=x_k, - ixv=x_v, - ixa=x_a, - ixg=x_g, - dxr=d_xr, - dxw=d_xw, - dxk=d_xk, - dxv=d_xv, - dxa=d_xa, - dxg=d_xg, - ghidden=g_hiddn, - gx=g_delta, - use_xg=use_xg, - T=T_SIZE, - T_OFFSET=T_OFFSET, - D=D, - BD=triton.next_power_of_2(D), - DTYPE=tl.float16 if hidden_states.dtype == torch.float16 else tl.float32, - ) - return g_hiddn, g_delta - - -@torch_compile -def addcmul_bwd2(d_oxr, d_xw, d_xk, d_xv, d_xa, d_xg, delta, use_xg: bool): - g_xr = (d_oxr * delta).sum(dim=(0, 1), keepdim=True, dtype=torch.float32) - g_xw = (d_xw * delta).sum(dim=(0, 1), keepdim=True, dtype=torch.float32) - g_xk = (d_xk * delta).sum(dim=(0, 1), keepdim=True, dtype=torch.float32) - g_xv = (d_xv * delta).sum(dim=(0, 1), keepdim=True, dtype=torch.float32) - g_xa = (d_xa * delta).sum(dim=(0, 1), keepdim=True, dtype=torch.float32) - g_xg = (d_xg * delta).sum(dim=(0, 1), keepdim=True, dtype=torch.float32) if use_xg else None - return g_xr, g_xw, g_xk, g_xv, g_xa, g_xg - - -class Rwkv7FusedAddcmul(torch.autograd.Function): - @staticmethod - @input_guard - def forward( - ctx, hidden_states, delta, - x_r, x_w, x_k, x_v, x_a, x_g, - ): - B, T, D = hidden_states.size() - oxr = torch.empty_like(hidden_states) - oxw = torch.empty_like(hidden_states) - oxk = torch.empty_like(hidden_states) - oxv = torch.empty_like(hidden_states) - oxa = torch.empty_like(hidden_states) - if x_g is not None: - use_xg = True - oxg = torch.empty_like(hidden_states) - else: - use_xg = False - oxg = None - - for t in range(0, T, 65536): - T_OFFSET = t - T_SIZE = min(65536, T - t) - def grid(meta): return (B, triton.cdiv(T_SIZE, meta['BT'])) - fused_addcmul_fwd_kernel[grid]( - hidden_states, delta, - x_r, x_w, x_k, x_v, x_a, x_g, - oxr, oxw, oxk, oxv, oxa, oxg, - use_xg, - T=T_SIZE, - T_OFFSET=T_OFFSET, - D=D, - BD=triton.next_power_of_2(D), - ) - - ctx.save_for_backward(hidden_states, delta, - x_r, x_w, x_k, x_v, x_a, x_g) - ctx.use_xg = use_xg - return oxr, oxw, oxk, oxv, oxa, oxg - - @staticmethod - @input_guard - def backward(ctx, dxr, - dxw, dxk, dxv, dxa, dxg): - hidden_states, delta, x_r, x_w, x_k, x_v, x_a, x_g = ctx.saved_tensors - - d_hiddn, d_xx = addcmul_bwd1(dxr, dxw, dxk, dxv, dxa, dxg, x_r, x_w, x_k, x_v, x_a, x_g, - hidden_states, delta, ctx.use_xg) - - d_ixr, d_ixw, d_ixk, d_ixv, d_ixa, d_ixg = addcmul_bwd2(dxr, dxw, dxk, dxv, dxa, dxg, delta, ctx.use_xg) - - return d_hiddn, d_xx, d_ixr, d_ixw, d_ixk, d_ixv, d_ixa, d_ixg - - -def fused_addcmul_rwkv7( - hidden_states: torch.Tensor, - delta: torch.Tensor, - xr: torch.Tensor, - xw: torch.Tensor, - xk: torch.Tensor, - xv: torch.Tensor, - xa: torch.Tensor, - xg: torch.Tensor | None = None, -): - if hidden_states.shape[1] == 1: - # Special case for decode - return torch_addcmul_rwkv7(hidden_states, delta, xr, xw, xk, xv, xa, xg) - return Rwkv7FusedAddcmul.apply(hidden_states, delta, xr, xw, xk, xv, xa, xg) - - -def torch_addcmul_rwkv7(hidden_states, delta, xr, xw, xk, xv, xa, xg=None): - oxr = torch.addcmul(hidden_states, delta, xr) - oxw = torch.addcmul(hidden_states, delta, xw) - oxk = torch.addcmul(hidden_states, delta, xk) - oxv = torch.addcmul(hidden_states, delta, xv) - oxa = torch.addcmul(hidden_states, delta, xa) - if xg is not None: - oxg = torch.addcmul(hidden_states, delta, xg) - return oxr, oxw, oxk, oxv, oxa, oxg - else: - return oxr, oxw, oxk, oxv, oxa, None - return oxr, oxw, oxk, oxv, oxa, None diff --git a/build/torch-xpu/ops/rwkv7/fused_k_update.py b/build/torch-xpu/ops/rwkv7/fused_k_update.py deleted file mode 100644 index 4a6942675c9b39521f459917697500ab1e0c1922..0000000000000000000000000000000000000000 --- a/build/torch-xpu/ops/rwkv7/fused_k_update.py +++ /dev/null @@ -1,356 +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 -import triton.language as tl - -from ...ops.utils import prepare_chunk_indices -from ...utils import IS_AMD, autotune_cache_kwargs, get_multiprocessor_count, input_guard - -NUM_WARPS_AUTOTUNE = [2, 4, 8, 16] if IS_AMD else [2, 4, 8, 16, 32] - - -def k_update_ref(k: torch.Tensor, a: torch.Tensor, ka: torch.Tensor) -> torch.Tensor: - return k.addcmul(k * (a - 1), ka) - - -@triton.heuristics({'IS_VARLEN': lambda args: args['cu_seqlens'] is not None}) -@triton.autotune( - configs=[ - triton.Config({}, num_warps=w, num_stages=s) - for w in NUM_WARPS_AUTOTUNE - for s in [1, 2, 3] - ], - key=['BD'], - **autotune_cache_kwargs, -) -@triton.jit -def k_update_fwd_kernel_short( - k, a, ka, out, - cu_seqlens, - T, D, - BD: tl.constexpr, - IS_VARLEN: tl.constexpr, -): - i_b, i_t = tl.program_id(0), tl.program_id(1) - - if IS_VARLEN: - bos = tl.load(cu_seqlens + i_b).to(tl.int32) - eos = tl.load(cu_seqlens + i_b + 1).to(tl.int32) - g_t = bos + i_t - if g_t >= eos: - return - offset = g_t * D - else: - g_t = i_t - offset = i_b * T * D + g_t * D - - o_d = tl.arange(0, BD) - m_d = o_d < D - off = offset + o_d - - b_k = tl.load(k + off, mask=m_d, other=0.).to(tl.float32) - b_a = tl.load(a + off, mask=m_d, other=0.).to(tl.float32) - b_ka = tl.load(ka + o_d, mask=m_d, eviction_policy='evict_last').to(tl.float32) - - out_val = b_k * (1 + (b_a - 1) * b_ka) - tl.store(out + off, out_val.to(out.dtype.element_ty), mask=m_d) - - -@triton.heuristics({'IS_VARLEN': lambda args: args['cu_seqlens'] is not None}) -@triton.autotune( - configs=[ - triton.Config({}, num_warps=w, num_stages=s) - for w in NUM_WARPS_AUTOTUNE - for s in [1, 2, 3] - ], - key=['BD', 'BT'], - **autotune_cache_kwargs, -) -@triton.jit -def k_update_fwd_kernel_long( - k, a, ka, out, - cu_seqlens, chunk_indices, - T, D, - BD: tl.constexpr, BT: tl.constexpr, - IS_VARLEN: tl.constexpr, -): - i_d, i_t_blk, i_b = tl.program_id(0), tl.program_id(1), tl.program_id(2) - - if IS_VARLEN: - i_n, i_t_blk = tl.load(chunk_indices + i_t_blk * 2).to(tl.int32), \ - tl.load(chunk_indices + i_t_blk * 2 + 1).to(tl.int32) - bos = tl.load(cu_seqlens + i_n).to(tl.int32) - eos = tl.load(cu_seqlens + i_n + 1).to(tl.int32) - t_start = i_t_blk * BT - t_end = tl.minimum(t_start + BT, eos - bos) - else: - bos = i_b * T - eos = (i_b + 1) * T - t_start = i_t_blk * BT - t_end = tl.minimum(t_start + BT, T) - - o_d = i_d * BD + tl.arange(0, BD) - m_d = o_d < D - - for t in range(t_start, t_end): - global_t = bos + t - off = global_t * D + o_d - b_k = tl.load(k + off, mask=m_d, other=0.).to(tl.float32) - b_a = tl.load(a + off, mask=m_d, other=0.).to(tl.float32) - b_ka = tl.load(ka + o_d, mask=m_d, eviction_policy='evict_last').to(tl.float32) - out_val = b_k * (1 + (b_a - 1) * b_ka) - tl.store(out + off, out_val.to(out.dtype.element_ty), mask=m_d) - - -@triton.heuristics({'IS_VARLEN': lambda args: args['cu_seqlens'] is not None}) -@triton.autotune( - configs=[ - triton.Config({'BT': BT}, num_warps=w, num_stages=s) - for w in NUM_WARPS_AUTOTUNE - for s in [1, 2, 3] - for BT in [2, 4, 8] - ], - key=['BD'], - **autotune_cache_kwargs, -) -@triton.jit -def k_update_bwd_kernel_short( - grad_out, k, a, ka, - dk, da, dka, - cu_seqlens, - T, D, - BT: tl.constexpr, - BD: tl.constexpr, - IS_VARLEN: tl.constexpr, -): - i_b, i_t_base = tl.program_id(0), tl.program_id(1) * BT - - if IS_VARLEN: - bos = tl.load(cu_seqlens + i_b).to(tl.int32) - eos = tl.load(cu_seqlens + i_b + 1).to(tl.int32) - seq_len = eos - bos - else: - bos = i_b * T - eos = (i_b + 1) * T - seq_len = T - - t_vec = i_t_base + tl.arange(0, BT) - mask_t = t_vec < seq_len - global_t_vec = bos + t_vec - - o_d = tl.arange(0, BD)[None, :] - m_d = o_d < D - off = global_t_vec[:, None] * D + o_d - - b_go = tl.load(grad_out + off, mask=mask_t[:, None] & m_d, other=0.).to(tl.float32) - b_k = tl.load(k + off, mask=mask_t[:, None] & m_d, other=0.).to(tl.float32) - b_a = tl.load(a + off, mask=mask_t[:, None] & m_d, other=0.).to(tl.float32) - b_ka = tl.load(ka + o_d, mask=m_d, eviction_policy='evict_last').to(tl.float32) # [1, BD] - - dk_vec = b_go * (1 + (b_a - 1) * b_ka) - da_vec = b_go * b_k * b_ka - dka_vec = b_go * b_k * (b_a - 1) - tl.store(dk + off, dk_vec.to(dk.dtype.element_ty), mask=mask_t[:, None] & m_d) - tl.store(da + off, da_vec.to(da.dtype.element_ty), mask=mask_t[:, None] & m_d) - tl.store(dka + off, dka_vec.to(dka.dtype.element_ty), mask=mask_t[:, None] & m_d) - - -@triton.heuristics({'IS_VARLEN': lambda args: args['cu_seqlens'] is not None}) -@triton.autotune( - configs=[ - triton.Config({}, num_warps=w, num_stages=s) - for w in NUM_WARPS_AUTOTUNE - for s in [1, 2, 3] - ], - key=['BD', 'BT'], - **autotune_cache_kwargs, -) -@triton.jit -def k_update_bwd_kernel_long( - grad_out, k, a, ka, - dk, da, dka, - cu_seqlens, chunk_indices, - T, D, - BD: tl.constexpr, BT: tl.constexpr, - IS_VARLEN: tl.constexpr, -): - i_d, i_t_blk, i_b = tl.program_id(0), tl.program_id(1), tl.program_id(2) - - if IS_VARLEN: - i_n, i_t_blk = tl.load(chunk_indices + i_t_blk * 2).to(tl.int32), \ - tl.load(chunk_indices + i_t_blk * 2 + 1).to(tl.int32) - bos = tl.load(cu_seqlens + i_n).to(tl.int32) - eos = tl.load(cu_seqlens + i_n + 1).to(tl.int32) - t_start = i_t_blk * BT - t_end = tl.minimum(t_start + BT, eos - bos) - else: - bos = i_b * T - eos = (i_b + 1) * T - t_start = i_t_blk * BT - t_end = tl.minimum(t_start + BT, T) - - o_d = i_d * BD + tl.arange(0, BD) - m_d = o_d < D - - for t in range(t_start, t_end): - global_t = bos + t - off = global_t * D + o_d - - b_go = tl.load(grad_out + off, mask=m_d, other=0.).to(tl.float32) - b_k = tl.load(k + off, mask=m_d, other=0.).to(tl.float32) - b_a = tl.load(a + off, mask=m_d, other=0.).to(tl.float32) - b_ka = tl.load(ka + o_d, mask=m_d, eviction_policy='evict_last').to(tl.float32) - - tl.store(dk + off, (b_go * (1 + (b_a - 1) * b_ka)).to(dk.dtype.element_ty), mask=m_d) - tl.store(da + off, (b_go * b_k * b_ka).to(da.dtype.element_ty), mask=m_d) - tl.store(dka + off, (b_go * b_k * (b_a - 1)).to(dka.dtype.element_ty), mask=m_d) - - -def k_update_fwd( - k: torch.Tensor, - a: torch.Tensor, - ka: torch.Tensor, - cu_seqlens: torch.Tensor | None = None, - cu_seqlens_cpu: torch.LongTensor | None = None, -) -> torch.Tensor: - B, T, D = k.shape - out = torch.empty_like(k) - use_short = T <= 512 - - if use_short: - if cu_seqlens is not None: - N = len(cu_seqlens) - 1 - else: - N = B - BD = triton.next_power_of_2(D) - grid = (N, T) - k_update_fwd_kernel_short[grid]( - k, a, ka, out, - cu_seqlens, - T, D, - BD=BD, - ) - else: - BT = min(64, triton.next_power_of_2( - triton.cdiv(max(16, B * T), get_multiprocessor_count(k.device.index)), - )) - if cu_seqlens is not None: - chunk_idx = prepare_chunk_indices(cu_seqlens, BT, cu_seqlens_cpu=cu_seqlens_cpu) - NT = len(chunk_idx) - N = len(cu_seqlens) - 1 - else: - chunk_idx = None - NT = triton.cdiv(T, BT) - N = B - - BD = triton.next_power_of_2(D) - - def grid(meta): - return (triton.cdiv(D, meta['BD']), NT, N) - - k_update_fwd_kernel_long[grid]( - k, a, ka, out, - cu_seqlens, chunk_idx, - T, D, - BD=BD, BT=BT, - ) - - return out, use_short, N, T - - -def k_update_bwd( - grad_out: torch.Tensor, - k: torch.Tensor, - a: torch.Tensor, - ka: torch.Tensor, - cu_seqlens: torch.Tensor | None, - use_short: bool, - N: int, - T: int, - cu_seqlens_cpu: torch.LongTensor | None = None, -): - B, _, D = grad_out.shape - dk = torch.empty_like(k) - da = torch.empty_like(a) - dka_tmp = torch.empty_like(k, dtype=torch.float32) - - if use_short: - BD = triton.next_power_of_2(D) - def grid(meta): return (N, triton.cdiv(T, meta['BT'])) - k_update_bwd_kernel_short[grid]( - grad_out, k, a, ka, - dk, da, dka_tmp, - cu_seqlens, - T, D, - BD=BD, - ) - else: - BT = min(64, triton.next_power_of_2( - triton.cdiv(max(16, B * T), get_multiprocessor_count(grad_out.device.index)), - )) - if cu_seqlens is not None: - chunk_idx = prepare_chunk_indices(cu_seqlens, BT, cu_seqlens_cpu=cu_seqlens_cpu) - NT = len(chunk_idx) - else: - chunk_idx = None - NT = triton.cdiv(T, BT) - - BD = triton.next_power_of_2(D) - - def grid(meta): - return (triton.cdiv(D, meta['BD']), NT, N) - - k_update_bwd_kernel_long[grid]( - grad_out, k, a, ka, - dk, da, dka_tmp, - cu_seqlens, chunk_idx, - T, D, - BD=BD, BT=BT, - ) - - if dka_tmp.dim() == 3: - dka = dka_tmp.sum(dim=(0, 1), keepdim=True).type_as(ka) - else: - dka = dka_tmp.sum(dim=(0, 1)).type_as(ka) - - return dk, da, dka - - -class KUpdateFunction(torch.autograd.Function): - @staticmethod - @input_guard - def forward(ctx, k, a, ka, cu_seqlens=None, cu_seqlens_cpu=None): - out, use_short, N, T = k_update_fwd(k, a, ka, cu_seqlens, cu_seqlens_cpu=cu_seqlens_cpu) - ctx.save_for_backward(k, a, ka) - ctx.use_short = use_short - ctx.N = N - ctx.T = T - ctx.cu_seqlens = cu_seqlens - ctx.cu_seqlens_cpu = cu_seqlens_cpu - return out - - @staticmethod - @input_guard - def backward(ctx, grad_output): - k, a, ka = ctx.saved_tensors - dk, da, dka = k_update_bwd( - grad_output, k, a, ka, - ctx.cu_seqlens, - ctx.use_short, - ctx.N, - ctx.T, - cu_seqlens_cpu=ctx.cu_seqlens_cpu, - ) - return dk, da, dka, None, None - - -def fused_k_rwkv7(k, a, ka, cu_seqlens=None, cu_seqlens_cpu=None): - if k.shape[1] == 1: - return k_update_ref(k, a, ka) - return KUpdateFunction.apply(k, a, ka, cu_seqlens, cu_seqlens_cpu) diff --git a/build/torch-xpu/ops/rwkv7/fused_recurrent.py b/build/torch-xpu/ops/rwkv7/fused_recurrent.py deleted file mode 100644 index 59a33027840b7a19a6dd7135240f2f92729fa8b1..0000000000000000000000000000000000000000 --- a/build/torch-xpu/ops/rwkv7/fused_recurrent.py +++ /dev/null @@ -1,313 +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 -import triton.language as tl - -from ...ops.generalized_delta_rule import fused_recurrent_dplr_delta_rule -from ...ops.utils.op import exp -from ...utils import autotune_cache_kwargs, input_guard - - -@triton.heuristics({ - 'USE_INITIAL_STATE': lambda args: args['h0'] is not None, - 'STORE_FINAL_STATE': lambda args: args['ht'] is not None, - 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None, -}) -@triton.autotune( - configs=[ - triton.Config({'BV': BV}, num_warps=num_warps, num_stages=num_stages) - for BV in [16, 32, 64] - for num_warps in [2, 4, 8, 16] - for num_stages in [2, 3, 4] - ], - key=['BK'], - **autotune_cache_kwargs, -) -@triton.jit(do_not_specialize=['T']) -def fused_recurrent_rwkv7_fwd_kernel( - r, - w, - k, - v, - kk, - a, - o, - h0, - ht, - cu_seqlens, - scale, - T, - B: tl.constexpr, - H: tl.constexpr, - K: tl.constexpr, - V: tl.constexpr, - BK: tl.constexpr, - BV: tl.constexpr, - REVERSE: tl.constexpr, - USE_INITIAL_STATE: tl.constexpr, - STORE_FINAL_STATE: tl.constexpr, - IS_VARLEN: tl.constexpr, - IS_DECODE: tl.constexpr, -): - i_v, i_nh = tl.program_id(0).to(tl.int64), tl.program_id(1).to(tl.int64) - i_n, i_h = i_nh // H, i_nh % H - - if IS_VARLEN: - bos, eos = tl.load(cu_seqlens + i_n).to(tl.int64), tl.load(cu_seqlens + i_n + 1).to(tl.int64) - T = eos - bos - else: - bos, eos = i_n * T, i_n * T + T - - o_k = tl.arange(0, BK) - o_v = i_v * BV + tl.arange(0, BV) - p_r = r + (bos + ((T - 1) if REVERSE else 0)) * H*K + i_h * K + o_k - p_w = w + (bos + ((T - 1) if REVERSE else 0)) * H*K + i_h * K + o_k - p_k = k + (bos + ((T - 1) if REVERSE else 0)) * H*K + i_h * K + o_k - p_v = v + (bos + ((T - 1) if REVERSE else 0)) * H*V + i_h * V + o_v - p_a = a + (bos + ((T - 1) if REVERSE else 0)) * H*K + i_h * K + o_k - p_kk = kk + (bos + ((T - 1) if REVERSE else 0)) * H*K + i_h * K + o_k - - p_o = o + (bos + ((T - 1) if REVERSE else 0)) * H*V + i_h * V + o_v - - mask_k = o_k < K - mask_v = o_v < V - mask_h = mask_k[:, None] & mask_v[None, :] - b_h = tl.zeros([BK, BV], dtype=tl.float32) - - if USE_INITIAL_STATE: - p_h0 = h0 + i_nh * K*V + o_k[:, None] * V + o_v - b_h += tl.load(p_h0, mask=mask_h, other=0).to(tl.float32) - - if IS_DECODE: - b_r = tl.load(p_r, mask=mask_k, other=0).to(tl.float32) * scale - b_w = tl.load(p_w, mask=mask_k, other=0).to(tl.float32) - b_k = tl.load(p_k, mask=mask_k, other=0).to(tl.float32) - b_v = tl.load(p_v, mask=mask_v, other=0).to(tl.float32) - b_a = tl.load(p_a, mask=mask_k, other=0).to(tl.float32) - b_kk = tl.load(p_kk, mask=mask_k, other=0).to(tl.float32) - b_act_a = -b_kk - b_b = b_kk * b_a - - b_h = exp(b_w)[:, None] * b_h + b_b[:, None] * tl.sum(b_act_a[:, None] * b_h, 0)[None, :] - b_h += b_k[:, None] * b_v[None, :] - b_o = tl.sum(b_h * b_r[:, None], 0) - - tl.store(p_o, b_o.to(p_o.dtype.element_ty), mask=mask_v) - else: - for _ in range(0, T): - b_r = tl.load(p_r, mask=mask_k, other=0).to(tl.float32) * scale - b_w = tl.load(p_w, mask=mask_k, other=0).to(tl.float32) - b_k = tl.load(p_k, mask=mask_k, other=0).to(tl.float32) - b_v = tl.load(p_v, mask=mask_v, other=0).to(tl.float32) - b_a = tl.load(p_a, mask=mask_k, other=0).to(tl.float32) - b_kk = tl.load(p_kk, mask=mask_k, other=0).to(tl.float32) - b_act_a = -b_kk - b_b = b_kk * b_a - - b_h = exp(b_w)[:, None] * b_h + b_b[:, None] * tl.sum(b_act_a[:, None] * b_h, 0)[None, :] - b_h += b_k[:, None] * b_v[None, :] - b_o = tl.sum(b_h * b_r[:, None], 0) - - tl.store(p_o, b_o.to(p_o.dtype.element_ty), mask=mask_v) - p_r += (-1 if REVERSE else 1) * H*K - p_w += (-1 if REVERSE else 1) * H*K - p_k += (-1 if REVERSE else 1) * H*K - p_v += (-1 if REVERSE else 1) * H*V - p_a += (-1 if REVERSE else 1) * H*K - p_kk += (-1 if REVERSE else 1) * H*K - p_o += (-1 if REVERSE else 1) * H*V - - if STORE_FINAL_STATE: - p_ht = ht + i_nh * K*V + o_k[:, None] * V + o_v - tl.store(p_ht, b_h.to(p_ht.dtype.element_ty), mask=mask_h) - - -@input_guard -def fused_recurrent_rwkv7_fwd( - r: torch.Tensor, - w: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - kk: torch.Tensor, - a: torch.Tensor, - scale: float | None = 1.0, - initial_state: torch.Tensor | None = None, - output_final_state: bool = False, - reverse: bool = False, - cu_seqlens: torch.LongTensor | None = None, -): - B, T, H, K, V = *k.shape, v.shape[-1] - N = B if cu_seqlens is None else len(cu_seqlens) - 1 - BK = triton.next_power_of_2(K) - IS_DECODE = (T == 1) - - h0 = initial_state - if not output_final_state: - ht = None - else: - ht = r.new_empty(N, H, K, V, dtype=torch.float32) - o = torch.empty_like(v) - - def grid(meta): return (triton.cdiv(V, meta['BV']), N * H) - fused_recurrent_rwkv7_fwd_kernel[grid]( - r, - w, - k, - v, - kk, - a, - o, - h0, - ht, - cu_seqlens, - scale, - T=T, - B=B, - H=H, - K=K, - V=V, - BK=BK, - REVERSE=reverse, - IS_DECODE=IS_DECODE, - ) - return o, ht - - -def fused_recurrent_rwkv7( - r: torch.Tensor, - w: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - a: torch.Tensor, - b: torch.Tensor, - scale: float | None = None, - initial_state: torch.Tensor | None = None, - output_final_state: bool = True, - cu_seqlens: torch.LongTensor | None = None, - **kwargs, -): - """ - Args: - r (torch.Tensor): - r of shape `[B, T, H, K]`. - w (torch.Tensor): - log decay of shape `[B, T, H, K]`. - k (torch.Tensor): - k of shape `[B, T, H, K]`. - v (torch.Tensor): - v of shape `[B, T, H, V]`. - a (torch.Tensor): - a of shape `[B, T, H, K]`. - b (torch.Tensor): - b of shape `[B, T, H, K]`. - scale (Optional[float]): - Scale factor for the attention scores. - If not provided, it will default to `1 / sqrt(K)`. Default: `None`. - initial_state (Optional[torch.Tensor]): - Initial state of shape `[B, H, K, V]` if `cu_seqlens` is `None`, - else `[N, H, K, V]` where `N = len(cu_seqlens) - 1`. Default: `None`. - output_final_state (Optional[bool]): - Whether to output the final state. Default: `True`. - cu_seqlens (torch.LongTensor): - Cumulative sequence lengths of shape `[N+1]` used for variable-length training, - consistent with the FlashAttention API. - """ - if 'head_first' in kwargs: - raise DeprecationWarning( - "head_first has been removed. Inputs must be in `[B, T, H, ...]` format.", - ) - return fused_recurrent_dplr_delta_rule( - q=r, - k=k, - v=v, - a=a, - b=b, - gk=w, - scale=scale, - initial_state=initial_state, - output_final_state=output_final_state, - cu_seqlens=cu_seqlens, - ) - - -def fused_mul_recurrent_rwkv7( - r: torch.Tensor, - w: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - kk: torch.Tensor, - a: torch.Tensor, - scale: float = 1.0, - initial_state: torch.Tensor | None = None, - output_final_state: bool = False, - reverse: bool = False, - cu_seqlens: torch.Tensor | None = None, - **kwargs, -) -> tuple[torch.Tensor, torch.Tensor]: - r""" - This function computes the recurrence S_t = S_t @ (I + a_t b_t^T) + v_t k_t^T in a recurrent manner. - - Args: - r (torch.Tensor): - queries of shape `[B, T, H, K]`. - w (torch.Tensor): - log decay of shape `[B, T, H, K]`. - k (torch.Tensor): - keys of shape `[B, T, H, K]`. - v (torch.Tensor): - values of shape `[B, T, H, V]`. - kk (torch.Tensor): - kk of shape `[B, T, H, K]`. - a (torch.Tensor): - a of shape `[B, T, H, K]`. - scale (float): - Scale factor for the attention scores. Default: `1.0`. - initial_state (Optional[torch.Tensor]): - Initial state of shape `[N, H, K, V]` for `N` input sequences. - For equal-length input sequences, `N` equals the batch size `B`. - Default: `None`. - output_final_state (Optional[bool]): - Whether to output the final state of shape `[N, H, K, V]`. Default: `False`. - reverse (Optional[bool]): - If `True`, process the state passing in reverse order. Default: `False`. - cu_seqlens (Optional[torch.Tensor]): - Cumulative sequence lengths of shape `[N+1]` used for variable-length training, - consistent with the FlashAttention API. - """ - if 'head_first' in kwargs: - raise DeprecationWarning( - "head_first has been removed. Inputs must be in `[B, T, H, ...]` format.", - ) - if cu_seqlens is not None: - if r.shape[0] != 1: - raise ValueError( - f"The batch size is expected to be 1 rather than {r.shape[0]} when using `cu_seqlens`. " - f"Please flatten variable-length inputs before processing.", - ) - if initial_state is not None and initial_state.shape[0] != len(cu_seqlens) - 1: - raise ValueError( - f"The number of initial states is expected to be equal to the number of input sequences, " - f"i.e., {len(cu_seqlens) - 1} rather than {initial_state.shape[0]}.", - ) - if scale is None: - scale = r.shape[-1] ** -0.5 - o, final_state = fused_recurrent_rwkv7_fwd( - r, - w, - k, - v, - kk, - a, - scale, - initial_state, - output_final_state, - reverse, - cu_seqlens, - ) - return o, final_state diff --git a/build/torch-xpu/ops/rwkv7/gate_output_correction.py b/build/torch-xpu/ops/rwkv7/gate_output_correction.py deleted file mode 100644 index 59ca4ca68db21184f9fa47e61b681b21617f9de0..0000000000000000000000000000000000000000 --- a/build/torch-xpu/ops/rwkv7/gate_output_correction.py +++ /dev/null @@ -1,278 +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 -import triton.language as tl - -from ...utils import autocast_custom_bwd, autocast_custom_fwd, autotune_cache_kwargs, input_guard - - -def gate_output_correction_ref( - o: torch.Tensor, - r: torch.Tensor, - k: torch.Tensor, - r_k: torch.Tensor, - v: torch.Tensor, - g: torch.Tensor, -): - """ - o: [B, T, H*D] - r: [B, T, H, D] - k: [B, T, H, D] - r_k: [H, D] - v: [B, T, H, D] - g: [B, T, H*D] - """ - # Unsqueeze r_k for broadcasting across batch and time - correction_term = ((r * k * r_k.unsqueeze(0).unsqueeze(0)).sum(-1, keepdim=True) * v).view(o.shape) - output = (o + correction_term) * g - return output - - -def gate_output_correction_backward_ref(grad_output, o, r, k, r_k, v, g): - """ - Reference backward pass implementation in pure PyTorch. - """ - B, T, HD = o.shape - H, D = r.shape[-2], r.shape[-1] - - # Unsqueeze r_k for broadcasting - r_k_b = r_k.unsqueeze(0).unsqueeze(0) - correction_scalar = (r * k * r_k_b).sum(-1, keepdim=True) - gated_input = o + (correction_scalar * v).view(B, T, HD) - - grad_g = grad_output * gated_input - grad_gated_input = grad_output * g - grad_o = grad_gated_input - grad_correction = grad_gated_input - grad_correction_reshaped = grad_correction.view(B, T, H, D) - grad_v = grad_correction_reshaped * correction_scalar - grad_correction_scalar = (grad_correction_reshaped * v).sum(-1, keepdim=True) - grad_r_mul_k_mul_rk = grad_correction_scalar.expand_as(r) - grad_r = grad_r_mul_k_mul_rk * k * r_k_b - grad_k = grad_r_mul_k_mul_rk * r * r_k_b - # Sum over batch and time, keep the head dimension - grad_r_k = (grad_r_mul_k_mul_rk * r * k).sum(dim=(0, 1)) - return grad_o, grad_r, grad_k, grad_r_k, grad_v, grad_g - - -@triton.autotune( - configs=[ - triton.Config({'BT': BT}, num_warps=num_warps) - for num_warps in [2, 4, 8] - for BT in [2, 4, 8] - ], - key=['num_heads', 'head_dim', 'BLOCK_SIZE_D'], - **autotune_cache_kwargs, -) -@triton.jit -def gate_output_correction_fwd_kernel( - o_ptr, r_ptr, k_ptr, r_k_ptr, v_ptr, g_ptr, output_ptr, - o_b_stride, o_t_stride, - r_b_stride, r_t_stride, r_h_stride, - v_b_stride, v_t_stride, v_h_stride, - r_k_h_stride, - T, - T_OFFSET, - num_heads: tl.constexpr, - head_dim: tl.constexpr, - BLOCK_SIZE_D: tl.constexpr, - BT: tl.constexpr, -): - pid_b, pid_t_block = tl.program_id(0), tl.program_id(1) - pid_h = tl.program_id(2) - t_local = pid_t_block * BT + tl.arange(0, BT)[:, None] - t_idx = T_OFFSET + t_local - mask_t = t_local < T - - d_idx = tl.arange(0, BLOCK_SIZE_D)[None, :] - mask_d = d_idx < head_dim - mask = mask_t & mask_d - - offset_rk_h = pid_h * r_k_h_stride - vec_r_k = tl.load(r_k_ptr + offset_rk_h + d_idx, mask=mask_d, other=0.0).to(tl.float32) - - offset_rh = pid_b * r_b_stride + t_idx * r_t_stride + pid_h * r_h_stride - offset_vh = pid_b * v_b_stride + t_idx * v_t_stride + pid_h * v_h_stride - vec_r = tl.load(r_ptr + offset_rh + d_idx, mask=mask, other=0.0).to(tl.float32) - vec_k = tl.load(k_ptr + offset_rh + d_idx, mask=mask, other=0.0).to(tl.float32) - vec_v = tl.load(v_ptr + offset_vh + d_idx, mask=mask, other=0.0).to(tl.float32) - correction = tl.sum(vec_r * vec_k * vec_r_k, axis=1)[:, None] * vec_v - - offset_o = pid_b * o_b_stride + t_idx * o_t_stride + pid_h * head_dim - vec_o = tl.load(o_ptr + offset_o + d_idx, mask=mask, other=0.0).to(tl.float32) - vec_g = tl.load(g_ptr + offset_o + d_idx, mask=mask, other=0.0).to(tl.float32) - final_output = (vec_o + correction) * vec_g - - tl.store(output_ptr + offset_o + d_idx, final_output.to(output_ptr.dtype.element_ty), mask=mask) - - -@triton.autotune( - configs=[ - triton.Config({'BT': BT}, num_warps=num_warps, num_stages=num_stages) - for num_warps in [2, 4, 8] - for num_stages in [1, 2, 4] - for BT in [2, 4, 8] - ], - key=['num_heads', 'head_dim', 'BLOCK_SIZE_D'], - **autotune_cache_kwargs, -) -@triton.jit -def gate_output_correction_bwd_kernel( - grad_output_ptr, o_ptr, r_ptr, k_ptr, r_k_ptr, v_ptr, g_ptr, - grad_o_ptr, grad_r_ptr, grad_k_ptr, grad_r_k_intermediate_ptr, grad_v_ptr, grad_g_ptr, - r_b_stride, r_t_stride, r_h_stride, o_b_stride, o_t_stride, r_k_h_stride, - T, - T_OFFSET, - num_heads: tl.constexpr, - head_dim: tl.constexpr, - BLOCK_SIZE_D: tl.constexpr, - BT: tl.constexpr, -): - pid_b, pid_t_block = tl.program_id(0), tl.program_id(1) - pid_h = tl.program_id(2) - - t_local = pid_t_block * BT + tl.arange(0, BT)[:, None] - t_idx = T_OFFSET + t_local - mask_t = t_local < T - - d_idx = tl.arange(0, BLOCK_SIZE_D)[None, :] - mask_d = d_idx < head_dim - mask = mask_t & mask_d - - rkv_offset = pid_b * r_b_stride + t_idx * r_t_stride + pid_h * r_h_stride - vec_r = tl.load(r_ptr + rkv_offset + d_idx, mask=mask, other=0.0).to(tl.float32) - vec_k = tl.load(k_ptr + rkv_offset + d_idx, mask=mask, other=0.0).to(tl.float32) - vec_v = tl.load(v_ptr + rkv_offset + d_idx, mask=mask, other=0.0).to(tl.float32) - - og_offset = pid_b * o_b_stride + t_idx * o_t_stride + pid_h * head_dim - vec_o = tl.load(o_ptr + og_offset + d_idx, mask=mask, other=0.0).to(tl.float32) - vec_g = tl.load(g_ptr + og_offset + d_idx, mask=mask, other=0.0).to(tl.float32) - vec_grad_output = tl.load(grad_output_ptr + og_offset + d_idx, mask=mask, other=0.0).to(tl.float32) - - offset_rk_h = pid_h * r_k_h_stride - vec_r_k = tl.load(r_k_ptr + offset_rk_h + d_idx, mask=mask_d, other=0.0).to(tl.float32) - - prod_r_k_rk = vec_r * vec_k * vec_r_k - corr_scalar = tl.sum(prod_r_k_rk, axis=1) - corr_vec = corr_scalar[:, None] * vec_v - gated_input = vec_o + corr_vec - - vec_grad_g = vec_grad_output * gated_input - vec_grad_gate = vec_grad_output * vec_g - vec_grad_o = vec_grad_gate - vec_grad_corr = vec_grad_gate - vec_grad_v = vec_grad_corr * corr_scalar[:, None] - grad_corr_s = tl.sum(vec_grad_corr * vec_v, axis=1)[:, None] - vec_grad_r = grad_corr_s * vec_k * vec_r_k - vec_grad_k = grad_corr_s * vec_r * vec_r_k - local_grad_rk = grad_corr_s * vec_r * vec_k - - tl.store(grad_o_ptr + og_offset + d_idx, vec_grad_o.to(grad_o_ptr.dtype.element_ty), mask=mask) - tl.store(grad_g_ptr + og_offset + d_idx, vec_grad_g.to(grad_g_ptr.dtype.element_ty), mask=mask) - tl.store(grad_r_ptr + rkv_offset + d_idx, vec_grad_r.to(grad_r_ptr.dtype.element_ty), mask=mask) - tl.store(grad_k_ptr + rkv_offset + d_idx, vec_grad_k.to(grad_k_ptr.dtype.element_ty), mask=mask) - tl.store(grad_v_ptr + rkv_offset + d_idx, vec_grad_v.to(grad_v_ptr.dtype.element_ty), mask=mask) - tl.store(grad_r_k_intermediate_ptr + rkv_offset + d_idx, - local_grad_rk.to(grad_r_k_intermediate_ptr.dtype.element_ty), mask=mask) - - -def gate_output_correction_backward_triton(grad_output, o, r, k, r_k, v, g): - batch_size, seq_len, _ = o.shape - num_heads, head_dim = r.shape[-2], r.shape[-1] - - grad_o = torch.empty_like(o) - grad_r = torch.empty_like(r) - grad_k = torch.empty_like(k) - grad_v = torch.empty_like(v) - grad_g = torch.empty_like(g) - # Keep intermediate in float32 for precision - grad_r_k = torch.empty_like(r, dtype=torch.float32) - - BLOCK_SIZE_D = triton.next_power_of_2(head_dim) - - for t_offset in range(0, seq_len, 65536): - T_SIZE = min(65536, seq_len - t_offset) - def grid(meta): return (batch_size, triton.cdiv(T_SIZE, meta['BT']), num_heads) - - gate_output_correction_bwd_kernel[grid]( - grad_output, o, r, k, r_k, v, g, - grad_o, grad_r, grad_k, grad_r_k, grad_v, grad_g, - r.stride(0), r.stride(1), r.stride(2), - o.stride(0), o.stride(1), - r_k.stride(0), - T_SIZE, t_offset, - num_heads=num_heads, head_dim=head_dim, BLOCK_SIZE_D=BLOCK_SIZE_D, - ) - # Sum over batch and time to get the final gradient for r_k - grad_r_k = grad_r_k.sum(dim=(0, 1)).type_as(r_k) - return grad_o, grad_r, grad_k, grad_r_k, grad_v, grad_g - - -class GateOutputCorrection(torch.autograd.Function): - @staticmethod - @autocast_custom_fwd - @input_guard - def forward(ctx, o, r, k, r_k, v, g): - assert r_k.dim() == 2 and r_k.shape[0] == r.shape[-2] and r_k.shape[1] == r.shape[-1] - - batch_size, seq_len, _ = o.shape - num_heads, head_dim = r.shape[-2], r.shape[-1] - output = torch.empty_like(o) - ctx.save_for_backward(o, r, k, r_k, v, g) - for t in range(0, seq_len, 65536): - T_OFFSET = t - T_SIZE = min(65536, seq_len - t) - def grid(meta): return (batch_size, triton.cdiv(T_SIZE, meta['BT']), num_heads) - - gate_output_correction_fwd_kernel[grid]( - o, r, k, r_k, v, g, output, - o.stride(0), o.stride(1), - r.stride(0), r.stride(1), r.stride(2), - v.stride(0), v.stride(1), v.stride(2), - r_k.stride(0), - T_SIZE, T_OFFSET, - num_heads, head_dim, BLOCK_SIZE_D=triton.next_power_of_2(head_dim), - ) - return output - - @staticmethod - @autocast_custom_bwd - @input_guard - def backward(ctx, grad_output): - o, r, k, r_k, v, g = ctx.saved_tensors - return gate_output_correction_backward_triton(grad_output, o, r, k, r_k, v, g) - - -@torch.compiler.disable -def gate_output_correction( - o: torch.Tensor, - r: torch.Tensor, - k: torch.Tensor, - r_k: torch.Tensor, - v: torch.Tensor, - g: torch.Tensor, -) -> torch.Tensor: - """ - Apply the RWKV7 output correction and output gate. - - Computes ``(o + correction) * g``, where - ``correction = ((r * k * r_k).sum(-1, keepdim=True) * v).view_as(o)``. - - Args: - o: Output tensor before correction, shape `[B, T, H * D]`. - r: Receptance tensor, shape `[B, T, H, D]`. - k: Key tensor, shape `[B, T, H, D]`. - r_k: Per-head correction weight, shape `[H, D]`. - v: Value tensor, shape `[B, T, H, D]`. - g: Output gate tensor, shape `[B, T, H * D]`. - - Returns: - Corrected and gated output tensor with the same shape as `o`. - """ - return GateOutputCorrection.apply(o, r, k, r_k, v, g) diff --git a/build/torch-xpu/ops/simple_gla/__init__.py b/build/torch-xpu/ops/simple_gla/__init__.py deleted file mode 100644 index ea5bb8eb53bff143f76c3728f1f556d9c3d2591d..0000000000000000000000000000000000000000 --- a/build/torch-xpu/ops/simple_gla/__init__.py +++ /dev/null @@ -1,18 +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 .chunk import chunk_simple_gla -from .fused_chunk import fused_chunk_simple_gla -from .fused_recurrent import fused_recurrent_simple_gla -from .parallel import parallel_simple_gla - -__all__ = [ - 'chunk_simple_gla', - 'fused_chunk_simple_gla', - 'fused_recurrent_simple_gla', - 'parallel_simple_gla', -] diff --git a/build/torch-xpu/ops/simple_gla/chunk.py b/build/torch-xpu/ops/simple_gla/chunk.py deleted file mode 100644 index 08313e35b8d2993f4dd05c0f83d85161b47fb82c..0000000000000000000000000000000000000000 --- a/build/torch-xpu/ops/simple_gla/chunk.py +++ /dev/null @@ -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 - -import torch - -from ...ops.common.chunk_h import chunk_bwd_dh, chunk_fwd_h -from ...ops.common.chunk_o import chunk_bwd_dqkwg, chunk_bwd_dv, chunk_fwd_o -from ...ops.utils import chunk_local_cumsum, prepare_chunk_indices -from ...ops.utils.constant import RCP_LN2 -from ...utils import autocast_custom_bwd, autocast_custom_fwd, input_guard - - -def chunk_simple_gla_fwd( - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - g: torch.Tensor | None = None, - g_gamma: torch.Tensor | None = None, - scale: float | None = None, - initial_state: torch.Tensor | None = None, - output_final_state: bool = False, - state_v_first: bool = False, - cu_seqlens: torch.LongTensor | None = None, - chunk_size: int = 64, - chunk_indices: torch.LongTensor | None = None, -) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: - h, ht = chunk_fwd_h( - k=k, - v=v, - g=g, - g_gamma=g_gamma, - gk=None, - gv=None, - h0=initial_state, - output_final_state=output_final_state, - cu_seqlens=cu_seqlens, - chunk_size=chunk_size, - states_in_fp32=False, - state_v_first=state_v_first, - ) - o = chunk_fwd_o( - q=q, - k=k, - v=v, - g=g, - g_gamma=g_gamma, - h=h, - scale=scale, - cu_seqlens=cu_seqlens, - chunk_size=chunk_size, - chunk_indices=chunk_indices, - state_v_first=state_v_first, - ) - return o, ht - - -def chunk_simple_gla_bwd( - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - g: torch.Tensor, - g_gamma: torch.Tensor, - initial_state: torch.Tensor, - do: torch.Tensor, - dht: torch.Tensor, - scale: float, - state_v_first: bool = False, - cu_seqlens: torch.LongTensor | None = None, - chunk_size: int = 64, - chunk_indices: torch.LongTensor | None = None, -) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: - # (SY 09/22) states_in_fp32 seems not affecting the error of dg but for safety, set to True - h, _ = chunk_fwd_h( - k=k, - v=v, - g=g, - g_gamma=g_gamma, - gk=None, - gv=None, - h0=initial_state, - output_final_state=False, - cu_seqlens=cu_seqlens, - chunk_size=chunk_size, - states_in_fp32=True, - state_v_first=state_v_first, - ) - dh, dh0 = chunk_bwd_dh( - q=q, - k=k, - v=v, - g=g, - g_gamma=g_gamma, - gk=None, - gv=None, - do=do, - h0=initial_state, - dht=dht, - scale=scale, - cu_seqlens=cu_seqlens, - chunk_size=chunk_size, - states_in_fp32=True, - state_v_first=state_v_first, - ) - dq, dk, _, dg = chunk_bwd_dqkwg( - q=q, - k=k, - v=v, - g=g, - g_gamma=g_gamma, - h=h, - do=do, - dh=dh, - scale=scale, - cu_seqlens=cu_seqlens, - chunk_size=chunk_size, - chunk_indices=chunk_indices, - state_v_first=state_v_first, - ) - dv = chunk_bwd_dv( - q=q, - k=k, - g=g, - g_gamma=g_gamma, - do=do, - dh=dh, - scale=scale, - cu_seqlens=cu_seqlens, - chunk_size=chunk_size, - chunk_indices=chunk_indices, - state_v_first=state_v_first, - ) - return dq, dk, dv, dg, dh0 - - -class ChunkSimpleGLAFunction(torch.autograd.Function): - - @staticmethod - @input_guard - @autocast_custom_fwd - def forward( - ctx, - q, - k, - v, - g, - g_gamma, - scale, - initial_state, - output_final_state, - state_v_first, - cu_seqlens, - cu_seqlens_cpu, - chunk_size: int | None = None, - ): - if chunk_size is None: - chunk_size = 64 - - if cu_seqlens is not None: - chunk_indices = prepare_chunk_indices( - cu_seqlens, - chunk_size, - cu_seqlens_cpu=cu_seqlens_cpu, - ) - else: - chunk_indices = None - - # Pre-scale by RCP_LN2 so downstream kernels can use exp2 directly. - if g is not None: - g = chunk_local_cumsum( - g, - chunk_size=chunk_size, - scale=RCP_LN2, - cu_seqlens=cu_seqlens, - chunk_indices=chunk_indices, - ) - if g_gamma is not None: - g_gamma = g_gamma * RCP_LN2 - o, ht = chunk_simple_gla_fwd( - q=q, - k=k, - v=v, - g=g, - g_gamma=g_gamma, - scale=scale, - initial_state=initial_state, - output_final_state=output_final_state, - cu_seqlens=cu_seqlens, - chunk_size=chunk_size, - chunk_indices=chunk_indices, - state_v_first=state_v_first, - ) - ctx.save_for_backward(q, k, v, g, g_gamma, initial_state, chunk_indices) - ctx.chunk_size = chunk_size - ctx.scale = scale - ctx.cu_seqlens = cu_seqlens - ctx.state_v_first = state_v_first - return o.to(q.dtype), ht - - @staticmethod - @input_guard - @autocast_custom_bwd - def backward(ctx, do, dht): - chunk_size, scale, cu_seqlens = ctx.chunk_size, ctx.scale, ctx.cu_seqlens - q, k, v, g, g_gamma, initial_state, chunk_indices = ctx.saved_tensors - dq, dk, dv, dg, dh0 = chunk_simple_gla_bwd( - q=q, - k=k, - v=v, - g=g, - g_gamma=g_gamma, - initial_state=initial_state, - do=do, - dht=dht, - scale=scale, - cu_seqlens=cu_seqlens, - chunk_size=chunk_size, - chunk_indices=chunk_indices, - state_v_first=ctx.state_v_first, - ) - if g is not None: - dg = chunk_local_cumsum( - dg, - chunk_size=chunk_size, - reverse=True, - cu_seqlens=cu_seqlens, - chunk_indices=chunk_indices, - ).to(g) - else: - dg = None - return dq.to(q), dk.to(k), dv.to(v), dg, None, None, dh0, None, None, None, None, None - - -@torch.compiler.disable -def chunk_simple_gla( - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - g: torch.Tensor | None = None, - g_gamma: torch.Tensor | None = None, - scale: float | None = None, - initial_state: torch.Tensor | None = None, - output_final_state: bool = False, - state_v_first: bool = False, - cu_seqlens: torch.LongTensor | None = None, - cu_seqlens_cpu: torch.LongTensor | None = None, - **kwargs, -) -> tuple[torch.Tensor, torch.Tensor]: - r""" - Args: - q (torch.Tensor): - queries of shape `[B, T, H, K]`. - k (torch.Tensor): - keys of shape `[B, T, H, K]`. - v (torch.Tensor): - values of shape `[B, T, H, V]`. - g (Optional[torch.Tensor]): - Forget gates of shape `[B, T, H]`. - Compared to GLA, the gating is head-wise instead of elementwise. - Default: `None`. - g_gamma (Optional[torch.Tensor]): - Log decay of shape `[H]`. - Head-wise data-independent decay is used if `g_gamma` is provided. - Only one of `g` or `g_gamma` should be provided. Default: `None`. - scale (Optional[float]): - Scale factor for the attention scores. - If not provided, it will default to `1 / sqrt(K)`. Default: `None`. - initial_state (Optional[torch.Tensor]): - Initial state of shape `[N, H, K, V]` (or `[N, H, V, K]` if `state_v_first=True`) - for `N` input sequences. - For equal-length input sequences, `N` equals the batch size `B`. - Default: `None`. - output_final_state (Optional[bool]): - Whether to output the final state of shape `[N, H, K, V]` - (or `[N, H, V, K]` if `state_v_first=True`). Default: `False`. - state_v_first (Optional[bool]): - Store the recurrent state in V-first `[V, K]` layout instead of the default `[K, V]`. Default: `False`. - cu_seqlens (torch.LongTensor): - Cumulative sequence lengths of shape `[N+1]` used for variable-length training, - consistent with the FlashAttention API. - cu_seqlens_cpu (torch.LongTensor): - CPU copy of `cu_seqlens` to avoid unnecessary device synchronization. Default: `None`. - - Returns: - o (torch.Tensor): - Outputs of shape `[B, T, H, V]`. - final_state (torch.Tensor): - Final state of shape `[N, H, K, V]` (or `[N, H, V, K]` if `state_v_first=True`) - if `output_final_state=True` else `None`. - - Examples:: - >>> import torch - >>> import torch.nn.functional as F - >>> from einops import rearrange - >>> from ...ops.simple_gla import chunk_simple_gla - # inputs with equal lengths - >>> B, T, H, K, V = 4, 2048, 4, 512, 512 - >>> q = torch.randn(B, T, H, K, device='cuda') - >>> k = torch.randn(B, T, H, K, device='cuda') - >>> v = torch.randn(B, T, H, V, device='cuda') - >>> g = F.logsigmoid(torch.randn(B, T, H, device='cuda')) - >>> o, ht = chunk_simple_gla( - q, k, v, g, - initial_state=None, - output_final_state=True - ) - # for variable-length inputs, the batch size `B` is expected to be 1 and `cu_seqlens` is required - >>> q, k, v, g = map(lambda x: rearrange(x, 'b t ... -> 1 (b t) ...'), (q, k, v, g)) - # for a batch with 4 sequences, `cu_seqlens` with 5 start/end positions are expected - >>> cu_seqlens = q.new_tensor([0, 2048, 4096, 6144, 8192], dtype=torch.long) - >>> o_var, ht_var = chunk_simple_gla( - q, k, v, g, - initial_state=None, - output_final_state=True, - cu_seqlens=cu_seqlens - ) - """ - if 'head_first' in kwargs: - raise DeprecationWarning( - "head_first has been removed. Inputs must be in `[B, T, H, ...]` format.", - ) - chunk_size = kwargs.pop('chunk_size', None) - if chunk_size is not None and chunk_size != 2 ** (chunk_size.bit_length() - 1): - raise ValueError(f"`chunk_size` must be a power of 2, got {chunk_size}.") - if cu_seqlens is not None: - if q.shape[0] != 1: - raise ValueError( - f"The batch size is expected to be 1 rather than {q.shape[0]} when using `cu_seqlens`. " - f"Please flatten variable-length inputs before processing.", - ) - if initial_state is not None and initial_state.shape[0] != len(cu_seqlens) - 1: - raise ValueError( - f"The number of initial states is expected to be equal to the number of input sequences, " - f"i.e., {len(cu_seqlens) - 1} rather than {initial_state.shape[0]}.", - ) - if scale is None: - scale = k.shape[-1] ** -0.5 - o, final_state = ChunkSimpleGLAFunction.apply( - q, - k, - v, - g, - g_gamma, - scale, - initial_state, - output_final_state, - state_v_first, - cu_seqlens, - cu_seqlens_cpu, - chunk_size, - ) - return o, final_state diff --git a/build/torch-xpu/ops/simple_gla/fused_chunk.py b/build/torch-xpu/ops/simple_gla/fused_chunk.py deleted file mode 100644 index 6e88e9c604b354c1714d5ab55029df7ab9bd44ae..0000000000000000000000000000000000000000 --- a/build/torch-xpu/ops/simple_gla/fused_chunk.py +++ /dev/null @@ -1,110 +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 - -from ...ops.common.fused_chunk import fused_chunk - - -def fused_chunk_simple_gla( - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - g: torch.Tensor = None, - g_gamma: torch.Tensor = None, - scale: float | None = None, - initial_state: torch.Tensor | None = None, - output_final_state: bool = False, - cu_seqlens: torch.LongTensor | None = None, -) -> tuple[torch.Tensor, torch.Tensor]: - r""" - Args: - q (torch.Tensor): - queries of shape `[B, T, H, K]`. - k (torch.Tensor): - keys of shape `[B, T, H, K]`. - v (torch.Tensor): - values of shape `[B, T, H, V]`. - g (torch.Tensor): - Forget gates of shape `[B, T, H]`. - Compared to GLA, the gating is head-wise instead of elementwise. - g_gamma (torch.Tensor): - Log decay of shape `[H]`. - Head-wise data-independent decay is used if `g_gamma` is provided. - Only one of `g` or `g_gamma` should be provided. - scale (Optional[int]): - Scale factor for the attention scores. - If not provided, it will default to `1 / sqrt(K)`. Default: `None`. - initial_state (Optional[torch.Tensor]): - Initial state of shape `[N, H, K, V]` for `N` input sequences. - For equal-length input sequences, `N` equals the batch size `B`. - Default: `None`. - output_final_state (Optional[bool]): - Whether to output the final state of shape `[N, H, K, V]`. Default: `False`. - cu_seqlens (torch.LongTensor): - Cumulative sequence lengths of shape `[N+1]` used for variable-length training, - consistent with the FlashAttention API. - - Returns: - o (torch.Tensor): - Outputs of shape `[B, T, H, V]`. - final_state (torch.Tensor): - Final state of shape `[N, H, K, V]` if `output_final_state=True` else `None`. - - Examples:: - >>> import torch - >>> import torch.nn.functional as F - >>> from einops import rearrange - >>> from ...ops.simple_gla import fused_chunk_simple_gla - # inputs with equal lengths - >>> B, T, H, K, V = 4, 2048, 4, 512, 512 - >>> q = torch.randn(B, T, H, K, device='cuda') - >>> k = torch.randn(B, T, H, K, device='cuda') - >>> v = torch.randn(B, T, H, V, device='cuda') - >>> g = F.logsigmoid(torch.randn(B, T, H, device='cuda')) - >>> h0 = torch.randn(B, H, K, V, device='cuda') - >>> o, ht = fused_chunk_simple_gla( - q, k, v, g, - initial_state=h0, - output_final_state=True - ) - # for variable-length inputs, the batch size `B` is expected to be 1 and `cu_seqlens` is required - >>> q, k, v, g = map(lambda x: rearrange(x, 'b t h d -> 1 (b t) h d'), (q, k, v, g)) - # for a batch with 4 sequences, `cu_seqlens` with 5 start/end positions are expected - >>> cu_seqlens = q.new_tensor([0, 2048, 4096, 6144, 8192], dtype=torch.long) - >>> o, ht = fused_chunk_simple_gla( - q, k, v, g, - initial_state=h0, - output_final_state=True, - cu_seqlens=cu_seqlens - ) - """ - if cu_seqlens is not None: - if q.shape[0] != 1: - raise ValueError( - f"The batch size is expected to be 1 rather than {q.shape[0]} when using `cu_seqlens`." - f"Please flatten variable-length inputs before processing.", - ) - if initial_state is not None and initial_state.shape[0] != len(cu_seqlens) - 1: - raise ValueError( - f"The number of initial states is expected to be equal to the number of input sequences, " - f"i.e., {len(cu_seqlens) - 1} rather than {initial_state.shape[0]}.", - ) - if scale is None: - scale = k.shape[-1] ** -0.5 - o, final_state = fused_chunk( - q=q, - k=k, - v=v, - g=g, - g_gamma=g_gamma, - scale=scale, - initial_state=initial_state, - output_final_state=output_final_state, - cu_seqlens=cu_seqlens, - ) - return o, final_state diff --git a/build/torch-xpu/ops/simple_gla/fused_recurrent.py b/build/torch-xpu/ops/simple_gla/fused_recurrent.py deleted file mode 100644 index dbc209dbd8c36eabdc01de640531c545a8c7f885..0000000000000000000000000000000000000000 --- a/build/torch-xpu/ops/simple_gla/fused_recurrent.py +++ /dev/null @@ -1,133 +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 warnings - -import torch - -from ...ops.common.fused_recurrent import fused_recurrent - - -def fused_recurrent_simple_gla( - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - g: torch.Tensor = None, - g_gamma: torch.Tensor = None, - scale: float | None = None, - initial_state: torch.Tensor | None = None, - output_final_state: bool = False, - reverse: bool = False, - state_v_first: bool = False, - cu_seqlens: torch.LongTensor | None = None, - **kwargs, -) -> tuple[torch.Tensor, torch.Tensor]: - r""" - Args: - q (torch.Tensor): - queries of shape `[B, T, H, K]`. - k (torch.Tensor): - keys of shape `[B, T, H, K]`. - v (torch.Tensor): - values of shape `[B, T, H, V]`. - g (torch.Tensor): - Forget gates of shape `[B, T, H]`. - Compared to GLA, the gating is head-wise instead of elementwise. - g_gamma (torch.Tensor): - Log decay of shape `[H]`. - Head-wise data-independent decay is used if `g_gamma` is provided. - Only one of `g` or `g_gamma` should be provided. - scale (Optional[float]): - Scale factor for the attention scores. - If not provided, it will default to `1 / sqrt(K)`. Default: `None`. - initial_state (Optional[torch.Tensor]): - Initial state of shape `[N, H, K, V]` for `N` input sequences. - For equal-length input sequences, `N` equals the batch size `B`. - Default: `None`. - output_final_state (Optional[bool]): - Whether to output the final state of shape `[N, H, K, V]`. Default: `False`. - reverse (Optional[bool]): - If `True`, process the state passing in reverse order. Default: `False`. - state_v_first (Optional[bool]): - Store the recurrent state in V-first ``[V, K]`` layout instead of the default ``[K, V]``. Default: ``False``. - cu_seqlens (torch.LongTensor): - Cumulative sequence lengths of shape `[N+1]` used for variable-length training, - consistent with the FlashAttention API. - `initial_state`, `final_state`, and `dh0`/`dht` all follow the chosen layout. - Default: `False`. - - Returns: - o (torch.Tensor): - Outputs of shape `[B, T, H, V]`. - final_state (torch.Tensor): - Final state of shape `[N, H, K, V]` if `output_final_state=True` else `None`. - - Examples:: - >>> import torch - >>> import torch.nn.functional as F - >>> from einops import rearrange - >>> from ...ops.simple_gla import fused_recurrent_simple_gla - # inputs with equal lengths - >>> B, T, H, K, V = 4, 2048, 4, 512, 512 - >>> q = torch.randn(B, T, H, K, device='cuda') - >>> k = torch.randn(B, T, H, K, device='cuda') - >>> v = torch.randn(B, T, H, V, device='cuda') - >>> g = F.logsigmoid(torch.randn(B, T, H, K, device='cuda')) - >>> h0 = torch.randn(B, H, K, V, device='cuda') - >>> o, ht = fused_recurrent_simple_gla( - q, k, v, g, - initial_state=h0, - output_final_state=True - ) - # for variable-length inputs, the batch size `B` is expected to be 1 and `cu_seqlens` is required - >>> q, k, v, g = map(lambda x: rearrange(x, 'b t h d -> 1 (b t) h d'), (q, k, v, g)) - # for a batch with 4 sequences, `cu_seqlens` with 5 start/end positions are expected - >>> cu_seqlens = q.new_tensor([0, 2048, 4096, 6144, 8192], dtype=torch.long) - >>> o_var, ht_var = fused_recurrent_simple_gla( - q, k, v, g, - initial_state=h0, - output_final_state=True, - cu_seqlens=cu_seqlens - ) - """ - if 'transpose_state_layout' in kwargs: - if state_v_first: - raise ValueError("Cannot pass both `state_v_first` and the deprecated `transpose_state_layout`.") - warnings.warn( - "`transpose_state_layout` is deprecated and renamed to `state_v_first`.", - DeprecationWarning, - stacklevel=2, - ) - state_v_first = kwargs.pop('transpose_state_layout') - - if cu_seqlens is not None: - if q.shape[0] != 1: - raise ValueError( - f"The batch size is expected to be 1 rather than {q.shape[0]} when using `cu_seqlens`." - f"Please flatten variable-length inputs before processing.", - ) - if initial_state is not None and initial_state.shape[0] != len(cu_seqlens) - 1: - raise ValueError( - f"The number of initial states is expected to be equal to the number of input sequences, " - f"i.e., {len(cu_seqlens) - 1} rather than {initial_state.shape[0]}.", - ) - if scale is None: - scale = k.shape[-1] ** -0.5 - o, final_state = fused_recurrent( - q=q, - k=k, - v=v, - g=g, - g_gamma=g_gamma, - scale=scale, - initial_state=initial_state, - output_final_state=output_final_state, - reverse=reverse, - cu_seqlens=cu_seqlens, - state_v_first=state_v_first, - ) - return o, final_state diff --git a/build/torch-xpu/ops/simple_gla/naive.py b/build/torch-xpu/ops/simple_gla/naive.py deleted file mode 100644 index 4b3f9c463eda9562fd12824afa19247e72a407e0..0000000000000000000000000000000000000000 --- a/build/torch-xpu/ops/simple_gla/naive.py +++ /dev/null @@ -1,117 +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 torch.nn.functional as F -from einops import rearrange - - -def naive_chunk_simple_gla( - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - g: torch.Tensor, - initial_state: torch.Tensor | None = None, - output_final_state: bool = False, - chunk_size: int = 64, - scale: float | None = None, -): - q, k, v, g = map(lambda x: rearrange(x, 'b t h ... -> b h t ...').to(torch.float32), [q, k, v, g]) - if scale is None: - scale = 1.0 / q.shape[-1] ** 0.5 - - T = q.shape[-2] - BT = chunk_size - pad_len = (BT - (T % BT)) % BT - if pad_len > 0: - # Pad all tensors - q = F.pad(q, (0, 0, 0, pad_len)) - k = F.pad(k, (0, 0, 0, pad_len)) - v = F.pad(v, (0, 0, 0, pad_len)) - g = F.pad(g, (0, pad_len)) - decay = g - B, H, T1, K = q.shape - V = v.shape[-1] - q = q * scale - q, k, v, decay = map(lambda x: rearrange(x, 'b h (n c) d -> b h n c d', c=chunk_size), [q, k, v, decay.unsqueeze(-1)]) - decay = decay.squeeze(-1).cumsum(-1) - L_mask = ((decay.unsqueeze(-1) - decay.unsqueeze(-2)).tril().exp().float()).tril() - S = k.new_zeros(B, H, K, V) - if initial_state is not None: - S = initial_state - o = torch.zeros_like(v) - for i in range(0, T1 // chunk_size): - q_i, k_i, v_i = q[:, :, i], k[:, :, i], v[:, :, i] - attn = (q_i @ k_i.transpose(-1, -2) * L_mask[:, :, i]) - o_inter = (q_i * decay[:, :, i, :, None].exp()) @ S - o[:, :, i] = o_inter + attn @ v_i - S = S * decay[:, :, i, -1, None, None].exp() + \ - (k_i * (decay[:, :, i, -1, None] - decay[:, :, i]).exp()[..., None]).transpose(-1, -2) @ v_i - if not output_final_state: - S = None - # unpad - o = rearrange(o, 'b h n c d -> b (n c) h d')[:, :T] - return o, S - - -def naive_recurrent_simple_gla( - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - g: torch.Tensor, - scale: float | None = None, - initial_state: torch.Tensor | None = None, - output_final_state: bool = True, -): - dtype = q.dtype - q, k, v, g = map(lambda x: x.transpose(1, 2).float(), (q, k, v, g)) - B, H, T, K = q.shape - V = v.shape[-1] - if scale is None: - scale = K ** -0.5 - q = q * scale - o = v.new_zeros(B, H, T, V) - - S = q.new_zeros(B, H, K, V) - if initial_state is not None: - S += initial_state - - for i in range(T): - gate = g[:, :, i].exp() - key = k[:, :, i] - value = v[:, :, i] - kv = key.unsqueeze(-1) * value.unsqueeze(-2) - S = S * gate.unsqueeze(-1).unsqueeze(-1) + kv - q_i = q[:, :, i, :] - o_i = (q_i.unsqueeze(-1) * S).sum(-2) - o[:, :, i] = o_i - if not output_final_state: - S = None - return o.transpose(1, 2).to(dtype), S - - -def naive_parallel_simple_gla( - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - g: torch.Tensor, - scale: float | None = None, -): - q, k, v, g = map(lambda x: rearrange(x, 'b t h ... -> b h t ...').to(torch.float32), [q, k, v, g]) - if scale is None: - scale = 1.0 / q.shape[-1] ** 0.5 - dtype = q.dtype - A = (q @ k.transpose(-1, -2) * scale) - if g is not None: - g = g.cumsum(-1) - D = (g.unsqueeze(-1) - g.unsqueeze(-2)).tril().exp().tril() - A = A * D - else: - A = A.tril() - o = A @ v - o = o.transpose(1, 2) - return o.to(dtype), A diff --git a/build/torch-xpu/ops/simple_gla/parallel.py b/build/torch-xpu/ops/simple_gla/parallel.py deleted file mode 100644 index e25f69bbaee7fbfd0817f55c70dd23833043ffc6..0000000000000000000000000000000000000000 --- a/build/torch-xpu/ops/simple_gla/parallel.py +++ /dev/null @@ -1,734 +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 -import triton.language as tl - -from ...ops.utils import prepare_chunk_indices -from ...ops.utils.constant import RCP_LN2 -from ...ops.utils.cumsum import chunk_global_cumsum, chunk_local_cumsum -from ...ops.utils.op import exp2 -from ...utils import ( - IS_INTEL_ALCHEMIST, - IS_NVIDIA_HOPPER, - autocast_custom_bwd, - autocast_custom_fwd, - autotune_cache_kwargs, - check_shared_mem, - input_guard, -) - -# https://github.com/intel/intel-xpu-backend-for-triton/issues/3449 -triton_config = {'grf_mode': 'large'} if IS_INTEL_ALCHEMIST else {} -NUM_WARPS = [2, 4, 8] if IS_NVIDIA_HOPPER else [2, 4, 8, 16] - - -@triton.heuristics({ - 'NV': lambda args: triton.cdiv(args['V'], args['BV']), - 'OUTPUT_ATTENTIONS': lambda args: args['attn'] is not None, - 'USE_G': lambda args: args['g'] is not None, - 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None, -}) -@triton.autotune( - configs=[ - triton.Config({}, num_warps=num_warps, num_stages=num_stages) - for num_warps in [2, 4, 8, 16] - for num_stages in [2, 3, 4] - ], - key=["BT", "BS", "BK", "BV", "USE_G"], - **autotune_cache_kwargs, -) -@triton.jit -def parallel_simple_gla_fwd_kernel( - q, - k, - v, - g, - o, - attn, - scale, - cu_seqlens, - chunk_indices, - T, - B: tl.constexpr, - H: tl.constexpr, - K: tl.constexpr, - V: tl.constexpr, - BT: tl.constexpr, - BS: tl.constexpr, - BK: tl.constexpr, - BV: tl.constexpr, - NV: tl.constexpr, - OUTPUT_ATTENTIONS: tl.constexpr, - IS_VARLEN: tl.constexpr, - USE_G: tl.constexpr, -): - i_kv, i_t, i_bh = tl.program_id(0), tl.program_id(1), tl.program_id(2) - i_k, i_v = i_kv // NV, i_kv % NV - i_b, i_h = i_bh // H, i_bh % H - - all = B * T - 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.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) - T = eos - bos - else: - bos, eos = i_b * T, i_b * T + T - - q += (bos * H + i_h) * K - k += (bos * H + i_h) * K - v += (bos * H + i_h) * V - o += ((i_k * all + bos) * H + i_h) * V - if USE_G: - g += bos * H + i_h - if OUTPUT_ATTENTIONS: - attn += i_k * B * H * T * T + (bos * H + i_h * T) * T - - p_q = tl.make_block_ptr(q, (T, K), (H*K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) - - # the Q block is kept in the shared memory throughout the whole kernel - # [BT, BK] - b_q = tl.load(p_q, boundary_check=(0, 1)) - b_q = (b_q * scale).to(b_q.dtype) - b_o = tl.zeros([BT, BV], dtype=tl.float32) - - # [BT] - o_q = i_t * BT + tl.arange(0, BT) - m_q = o_q < T - # Q block and K block have overlap. - # masks required - if USE_G: - # [BT,] - b_gq = tl.load(g + o_q * H, mask=m_q, other=float('-inf')).to(tl.float32) - # rescale interchunk output - else: - b_gq = None - - for i_s in range(i_t * BT, min((i_t + 1) * BT, T), BS): - p_k = tl.make_block_ptr(k, (K, T), (1, H*K), (i_k * BK, i_s), (BK, BS), (0, 1)) - p_v = tl.make_block_ptr(v, (T, V), (H*V, 1), (i_s, i_v * BV), (BS, BV), (1, 0)) - - o_k = i_s + tl.arange(0, BS) - m_k = o_k < T - # [BK, BS] - b_k = tl.load(p_k, boundary_check=(0, 1)) - # [BS, BV] - b_v = tl.load(p_v, boundary_check=(0, 1)) - # [BT, BS] - m_s = (o_q[:, None] >= o_k[None, :]) & (m_q[:, None] & m_k[None, :]) - b_s = tl.dot(b_q, b_k) - if USE_G: - b_gk = tl.load(g + o_k * H, mask=m_k, other=0) - b_s *= exp2(b_gq[:, None] - b_gk[None, :]) - b_s = tl.where(m_s, b_s, 0) - # [BT, BV] - if i_s >= 0: - b_o += tl.dot(b_s.to(b_q.dtype), b_v) - if OUTPUT_ATTENTIONS: - p_a = tl.make_block_ptr(attn, (T, T), (T, 1), (i_t * BT, i_s), (BT, BS), (1, 0)) - tl.store(p_a, b_s.to(p_a.dtype.element_ty), boundary_check=(0, 1)) - for i_s in range(i_t * BT - BS, -BS, -BS): - p_k = tl.make_block_ptr(k, (K, T), (1, H*K), (i_k * BK, i_s), (BK, BS), (0, 1)) - p_v = tl.make_block_ptr(v, (T, V), (H*V, 1), (i_s, i_v * BV), (BS, BV), (1, 0)) - - o_k = i_s + tl.arange(0, BS) - m_k = o_k < T - # [BK, BS] - b_k = tl.load(p_k, boundary_check=(0, 1)) - # [BS, BV] - b_v = tl.load(p_v, boundary_check=(0, 1)) - # [BT, BS] - m_s = m_q[:, None] & m_k[None, :] - b_s = tl.dot(b_q, b_k) - if USE_G: - b_g = tl.load(g + o_k * H, mask=m_k, other=0) - b_gn = tl.load(g + (min(i_s + BS, T) - 1) * H) - b_gp = tl.load(g + (i_s-1) * H) if i_s % BT > 0 else 0. - # No concrete meaning. Just to avoid some layout bugs. - b_s *= exp2(b_gq[:, None] + (b_gn - b_g)[None, :]) - b_gq += b_gn - b_gp - b_s = tl.where(m_s, b_s, 0) - if OUTPUT_ATTENTIONS: - p_a = tl.make_block_ptr(attn, (T, T), (T, 1), (i_t * BT, i_s), (BT, BS), (1, 0)) - tl.store(p_a, b_s.to(p_a.dtype.element_ty), boundary_check=(0, 1)) - if i_s >= 0: - b_o += tl.dot(b_s.to(b_v.dtype), b_v) - p_o = tl.make_block_ptr(o, (T, V), (H*V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) - tl.store(p_o, b_o.to(p_o.dtype.element_ty), boundary_check=(0, 1)) - - -@triton.jit(do_not_specialize=['T']) -def parallel_simple_gla_bwd_kernel_dq( - i_t, - i_k, - i_v, - q, - k, - v, - g, - do, - dq, - dg, - scale, - T, - H: tl.constexpr, - K: tl.constexpr, - V: tl.constexpr, - BT: tl.constexpr, - BS: tl.constexpr, - BK: tl.constexpr, - BV: tl.constexpr, - USE_G: tl.constexpr, -): - p_do = tl.make_block_ptr(do, (T, V), (H*V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) - # [BT, BV] - b_do = tl.load(p_do, boundary_check=(0, 1)) - # [BT, BK] - b_dq = tl.zeros([BT, BK], dtype=tl.float32) - - # [BT] - o_q = i_t * BT + tl.arange(0, BT) - m_q = o_q < T - for i_s in range(0, i_t * BT, BS): - p_k = tl.make_block_ptr(k, (T, K), (H*K, 1), (i_s, i_k * BK), (BS, BK), (1, 0)) - p_v = tl.make_block_ptr(v, (V, T), (1, H*V), (i_v * BV, i_s), (BV, BS), (0, 1)) - - o_k = i_s + tl.arange(0, BS) - m_k = o_k < T - # [BS, BK] - b_k = tl.load(p_k, boundary_check=(0, 1)) - # [BV, BS] - b_v = tl.load(p_v, boundary_check=(0, 1)) - # [BT, BV] @ [BV, BS] = [BT, BS] - b_ds = tl.dot(b_do, b_v) - if USE_G: - b_g = tl.load(g + o_k * H, mask=m_k, other=0) - b_gn = tl.load(g + (min(i_s + BS, T) - 1) * H) - b_gp = tl.load(g + (i_s - 1) * H) if i_s % BT > 0 else 0. - b_ds *= tl.where(m_k, exp2(b_gn - b_g), 0)[None, :] - if i_s > 0: - b_dq *= exp2(b_gn - b_gp) - # [BT, BS] @ [BS, BK] = [BT, BK] - b_dq += tl.dot(b_ds.to(b_v.dtype), b_k) - - if USE_G: - # [BT,] - b_gq = tl.load(g + o_q * H, mask=m_q, other=float('-inf')) - # [BT, BK] - b_dq *= exp2(b_gq)[:, None] - # Q block and K block have overlap. masks required - for i_s in range(i_t * BT, min((i_t + 1) * BT, T), BS): - p_k = tl.make_block_ptr(k, (T, K), (H*K, 1), (i_s, i_k * BK), (BS, BK), (1, 0)) - p_v = tl.make_block_ptr(v, (V, T), (1, H*V), (i_v * BV, i_s), (BV, BS), (0, 1)) - - o_k = i_s + tl.arange(0, BS) - m_k = o_k < T - # [BS, BK] - b_k = tl.load(p_k, boundary_check=(0, 1)) - # [BV, BS] - b_v = tl.load(p_v, boundary_check=(0, 1)) - # [BT, BV] @ [BV, BS] = [BT, BS] - b_ds = tl.dot(b_do, b_v) - if USE_G: - b_gk = tl.load(g + o_k * H, mask=m_k, other=0) - b_ds *= exp2(b_gq[:, None] - b_gk[None, :]) - m_s = (o_q[:, None] >= o_k[None, :]) & (m_q[:, None] & m_k[None, :]) - b_ds = tl.where(m_s, b_ds, 0) - # [BT, BK] - b_dq += tl.dot(b_ds.to(b_k.dtype), b_k) - - b_dq *= scale - p_dq = tl.make_block_ptr(dq, (T, K), (H*K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) - tl.store(p_dq, b_dq.to(p_dq.dtype.element_ty), boundary_check=(0, 1)) - if USE_G: - p_q = tl.make_block_ptr(q, (T, K), (H*K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) - b_q = tl.load(p_q, boundary_check=(0, 1)) - b_dg = tl.sum(b_dq * b_q, 1) - p_dg = tl.make_block_ptr(dg, (T,), (H,), (i_t * BT,), (BT,), (0,)) - tl.store(p_dg, b_dg.to(p_dg.dtype.element_ty), boundary_check=(0,)) - - -@triton.jit(do_not_specialize=['T']) -def parallel_simple_gla_bwd_kernel_dkv( - i_t, - i_k, - i_v, - q, - k, - v, - g, - do, - dk, - dv, - dg, - scale, - T, - H: tl.constexpr, - K: tl.constexpr, - V: tl.constexpr, - BT: tl.constexpr, - BS: tl.constexpr, - BK: tl.constexpr, - BV: tl.constexpr, - USE_G: tl.constexpr, -): - o_k = i_t * BT + tl.arange(0, BT) - m_k = o_k < T - # [BT, BK] - p_k = tl.make_block_ptr(k, (T, K), (H*K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) - b_k = tl.load(p_k, boundary_check=(0, 1)) - b_dk = tl.zeros([BT, BK], dtype=tl.float32) - # [BT, BV] - p_v = tl.make_block_ptr(v, (T, V), (H*V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) - b_v = tl.load(p_v, boundary_check=(0, 1)) - b_dv = tl.zeros([BT, BV], dtype=tl.float32) - if USE_G: - b_gk = tl.load(g + o_k * H, mask=m_k, other=0) - NTS = tl.cdiv(T, BS) - # [BT, BK] - for i_s in range(NTS * BS - BS, (i_t + 1) * BT - BS, -BS): - p_q = tl.make_block_ptr(q, (T, K), (H*K, 1), (i_s, i_k * BK), (BS, BK), (1, 0)) - p_do = tl.make_block_ptr(do, (T, V), (H*V, 1), (i_s, i_v * BV), (BS, BV), (1, 0)) - - o_q = i_s + tl.arange(0, BS) - m_q = o_q < T - # [BS, BK] - b_q = tl.load(p_q, boundary_check=(0, 1)) - # [BS, BV] - b_do = tl.load(p_do, boundary_check=(0, 1)) - # [BT, BS] - b_ds = tl.dot(b_v, tl.trans(b_do)) - b_s = tl.dot(b_k, tl.trans(b_q)) - if USE_G: - b_gq = tl.load(g + o_q * H, mask=m_q, other=float('-inf')) - b_gp = tl.load(g + (min(i_s + BS, T) - 1) * H) - b_gn = tl.load(g + (i_s - 1) * H) if i_s % BT > 0 else 0. - if i_s >= 0: - b_gpn = exp2(b_gp - b_gn) - b_dk *= b_gpn - b_dv *= b_gpn - b_gqn = exp2(b_gq - b_gn) - b_ds *= b_gqn[None, :] - b_s *= b_gqn[None, :] - # [BT, BK] - b_dk += tl.dot(b_ds.to(b_q.dtype), b_q) - # [BT, BV] - b_dv += tl.dot(b_s.to(b_do.dtype), b_do) - - if USE_G: - b_gn = tl.load(g + (min(i_t * BT + BT, T) - 1) * H) - if i_t >= 0: - b_gpn = exp2(b_gn - b_gk)[:, None] - b_dk *= b_gpn - b_dv *= b_gpn - - for i_s in range(i_t * BT, min((i_t + 1) * BT, T), BS): - p_q = tl.make_block_ptr(q, (T, K), (H*K, 1), (i_s, i_k * BK), (BS, BK), (1, 0)) - p_do = tl.make_block_ptr(do, (T, V), (H*V, 1), (i_s, i_v * BV), (BS, BV), (1, 0)) - - o_q = i_s + tl.arange(0, BS) - m_q = o_q < T - # [BS, BK] - b_q = tl.load(p_q, boundary_check=(0, 1)) - # [BS, BV] - b_do = tl.load(p_do, boundary_check=(0, 1)) - # [BS] - b_s = tl.dot(b_k, tl.trans(b_q)) - b_ds = tl.dot(b_v, tl.trans(b_do)) - if USE_G: - b_gq = tl.load(g + o_q * H, mask=m_q, other=float('-inf')) - if i_s >= 0: - b_gkq = exp2(-b_gk[:, None] + b_gq[None, :]) - b_ds *= b_gkq - b_s *= b_gkq - m_s = o_k[:, None] <= o_q[None, :] - b_s = tl.where(m_s, b_s, 0) - b_ds = tl.where(m_s, b_ds, 0) - # [BT, BK] - b_dk += tl.dot(b_ds.to(b_q.dtype), b_q) - b_dv += tl.dot(b_s.to(b_do.dtype), b_do) - b_dk *= scale - b_dv *= scale - p_dk = tl.make_block_ptr(dk, (T, K), (H*K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) - p_dv = tl.make_block_ptr(dv, (T, V), (H*V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) - tl.store(p_dk, b_dk.to(p_dk.dtype.element_ty), boundary_check=(0, 1)) - tl.store(p_dv, b_dv.to(p_dv.dtype.element_ty), boundary_check=(0, 1)) - if USE_G: - b_dg = tl.load(dg + o_k * H, mask=m_k, other=0) - b_dg -= tl.sum(b_dk * b_k, 1) - tl.store(dg + o_k * H, b_dg.to(dg.dtype.element_ty), mask=m_k) - - -@triton.heuristics({ - 'NV': lambda args: triton.cdiv(args['V'], args['BV']), - 'USE_G': lambda args: args['g'] is not None, - 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None, -}) -@triton.autotune( - configs=[ - triton.Config(triton_config, num_warps=num_warps) - for num_warps in NUM_WARPS - ], - key=['BT', 'BS', 'BK', 'BV', 'USE_G'], - **autotune_cache_kwargs, -) -@triton.jit(do_not_specialize=['T']) -def parallel_simple_gla_bwd_kernel( - q, - k, - v, - g, - do, - dq, - dk, - dv, - dg, - scale, - cu_seqlens, - chunk_indices, - T, - B: tl.constexpr, - H: tl.constexpr, - K: tl.constexpr, - V: tl.constexpr, - BT: tl.constexpr, - BS: tl.constexpr, - BK: tl.constexpr, - BV: tl.constexpr, - NV: tl.constexpr, - IS_VARLEN: tl.constexpr, - USE_G: tl.constexpr, -): - i_kv, i_t, i_bh = tl.program_id(0), tl.program_id(1), tl.program_id(2) - i_k, i_v = i_kv // NV, i_kv % NV - i_b, i_h = i_bh // H, i_bh % H - dq += i_v * B * H * T * K - dk += i_v * B * H * T * K - dv += i_k * B * H * T * V - if USE_G: - dg += i_kv * B * H * T - - 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.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) - T = eos - bos - else: - bos, eos = i_b * T, i_b * T + T - - q += (bos * H + i_h) * K - k += (bos * H + i_h) * K - v += (bos * H + i_h) * V - do += (bos * H + i_h) * V - dq += (bos * H + i_h) * K - dk += (bos * H + i_h) * K - dv += (bos * H + i_h) * V - if USE_G: - g += bos * H + i_h - dg += bos * H + i_h - - parallel_simple_gla_bwd_kernel_dq( - i_t=i_t, - i_k=i_k, - i_v=i_v, - q=q, - k=k, - v=v, - g=g, - do=do, - dq=dq, - dg=dg, - scale=scale, - T=T, - H=H, - K=K, - V=V, - BT=BT, - BS=BS, - BK=BK, - BV=BV, - USE_G=USE_G, - ) - tl.debug_barrier() - parallel_simple_gla_bwd_kernel_dkv( - i_t=i_t, - i_k=i_k, - i_v=i_v, - q=q, - k=k, - v=v, - g=g, - do=do, - dk=dk, - dv=dv, - dg=dg, - scale=scale, - T=T, - H=H, - K=K, - V=V, - BT=BT, - BS=BS, - BK=BK, - BV=BV, - USE_G=USE_G, - ) - - -def parallel_simple_gla_fwd( - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - g: torch.Tensor, - scale: float, - output_attentions: bool = False, - chunk_size: int = 128, - cu_seqlens: torch.LongTensor | None = None, - chunk_indices: torch.LongTensor | None = None, -): - B, T, H, K, V = *k.shape, v.shape[-1] - BT, BS = chunk_size, 32 - if check_shared_mem('hopper', k.device.index): - BK = min(256, triton.next_power_of_2(K)) - BV = min(256, triton.next_power_of_2(V)) - elif check_shared_mem('ampere', k.device.index): - BK = min(128, triton.next_power_of_2(K)) - BV = min(128, triton.next_power_of_2(V)) - else: - BK = min(64, triton.next_power_of_2(K)) - BV = min(64, triton.next_power_of_2(V)) - - NK = triton.cdiv(K, BK) - NV = triton.cdiv(V, BV) - assert BT % BS == 0 - - if chunk_indices is None: - chunk_indices = prepare_chunk_indices(cu_seqlens, chunk_size) if cu_seqlens is not None else None - NT = triton.cdiv(T, BT) if cu_seqlens is None else len(chunk_indices) - - # local cumulative decay in log space - if g is not None: - g = chunk_local_cumsum( - g, - chunk_size, - scale=RCP_LN2, - cu_seqlens=cu_seqlens, - chunk_indices=chunk_indices, - ) - grid = (NK * NV, NT, B * H) - o = torch.empty(NK, *v.shape, dtype=v.dtype if NK == 1 else torch.float, device=q.device) - attn = q.new_zeros(NK, B, H, T, T) if output_attentions else None - - parallel_simple_gla_fwd_kernel[grid]( - q=q, - k=k, - v=v, - g=g, - o=o, - attn=attn, - scale=scale, - cu_seqlens=cu_seqlens, - chunk_indices=chunk_indices, - B=B, - H=H, - T=T, - K=K, - V=V, - BT=BT, - BS=BS, - BK=BK, - BV=BV, - ) - o = o.sum(0) - - if output_attentions: - attn = attn.sum(0) - return o, g, attn - - -def parallel_simple_gla_bwd( - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - g: torch.Tensor, - do: torch.Tensor, - scale: float, - chunk_size: int = 128, - cu_seqlens: torch.LongTensor | None = None, - chunk_indices: torch.LongTensor | None = None, -): - B, T, H, K, V = *k.shape, v.shape[-1] - BT, BS = chunk_size, 32 - if check_shared_mem('hopper', k.device.index): - BK = min(256, triton.next_power_of_2(K)) - BV = min(256, triton.next_power_of_2(V)) - elif check_shared_mem('ampere', k.device.index): - BK = min(128, triton.next_power_of_2(K)) - BV = min(128, triton.next_power_of_2(V)) - elif check_shared_mem('ada', k.device.index): - BK = min(64, triton.next_power_of_2(K)) - BV = min(64, triton.next_power_of_2(V)) - else: - BK = min(32, triton.next_power_of_2(K)) - BV = min(32, triton.next_power_of_2(V)) - - NK = triton.cdiv(K, BK) - NV = triton.cdiv(V, BV) - assert BT % BS == 0 - - dq = torch.empty(NV, * q.shape, dtype=q.dtype if NV == 1 else torch.float, device=q.device) - dk = torch.empty(NV, * k.shape, dtype=k.dtype if NV == 1 else torch.float, device=q.device) - dv = torch.empty(NK, * v.shape, dtype=v.dtype if NK == 1 else torch.float, device=q.device) - dg = torch.empty(NK*NV, *g.shape, dtype=torch.float, device=q.device) if g is not None else None - - if chunk_indices is None: - chunk_indices = prepare_chunk_indices(cu_seqlens, chunk_size) if cu_seqlens is not None else None - NT = triton.cdiv(T, BT) if cu_seqlens is None else len(chunk_indices) - - grid = (NK * NV, NT, B * H) - parallel_simple_gla_bwd_kernel[grid]( - q=q, - k=k, - v=v, - g=g, - do=do, - dq=dq, - dk=dk, - dv=dv, - dg=dg, - cu_seqlens=cu_seqlens, - chunk_indices=chunk_indices, - scale=scale, - T=T, - B=B, - H=H, - K=K, - V=V, - BT=BT, - BS=BS, - BK=BK, - BV=BV, - ) - dq = dq.sum(0) - dk = dk.sum(0) - dv = dv.sum(0) - dg = chunk_global_cumsum(dg.sum(0), reverse=True, cu_seqlens=cu_seqlens) if g is not None else None - return dq, dk, dv, dg - - -class ParallelSimpleGLAFunction(torch.autograd.Function): - - @staticmethod - @input_guard - @autocast_custom_fwd - def forward(ctx, q, k, v, g, scale, output_attentions, cu_seqlens, cu_seqlens_cpu): - chunk_size = 128 - ctx.dtype = q.dtype - - chunk_indices = prepare_chunk_indices( - cu_seqlens, chunk_size, cu_seqlens_cpu=cu_seqlens_cpu) if cu_seqlens is not None else None - - o, g, attn = parallel_simple_gla_fwd( - q=q, - k=k, - v=v, - g=g, - scale=scale, - output_attentions=output_attentions, - chunk_size=chunk_size, - cu_seqlens=cu_seqlens, - chunk_indices=chunk_indices, - ) - ctx.save_for_backward(q, k, v, g, cu_seqlens, chunk_indices) - ctx.scale = scale - ctx.chunk_size = chunk_size - return o.to(q.dtype), attn - - @staticmethod - @input_guard - @autocast_custom_bwd - def backward(ctx, do, da=None): - q, k, v, g, cu_seqlens, chunk_indices = ctx.saved_tensors - dq, dk, dv, dg = parallel_simple_gla_bwd( - q=q, - k=k, - v=v, - g=g, - do=do, - scale=ctx.scale, - chunk_size=ctx.chunk_size, - cu_seqlens=cu_seqlens, - chunk_indices=chunk_indices, - ) - return dq.to(q), dk.to(k), dv.to(v), dg.to(ctx.dtype) if dg is not None else None, None, None, None, None - - -def parallel_simple_gla( - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - g: torch.Tensor | None = None, - scale: float | None = None, - output_attentions: bool = False, - cu_seqlens: torch.LongTensor | None = None, - cu_seqlens_cpu: torch.LongTensor | None = None, - **kwargs, -) -> tuple[torch.Tensor, torch.Tensor]: - r""" - Args: - q (torch.Tensor): - queries of shape `[B, T, H, K]`. - k (torch.Tensor): - keys of shape `[B, T, H, K]`. - v (torch.Tensor): - values of shape `[B, T, H, V]`. - g (Optional[torch.Tensor]): - Forget gates of shape `[B, T, H]`. - Compared to GLA, the gating is head-wise instead of elementwise. Default: `None`. - scale (Optional[float]): - Scale factor for attention scores. - If not provided, it will default to `1 / sqrt(K)`. Default: `None`. - output_attentions (bool): - Whether to output the materialized attention scores of shape `[B, H, T, T]`. Default: `False`. - cu_seqlens (torch.LongTensor): - Cumulative sequence lengths of shape `[N+1]` used for variable-length training, - consistent with the FlashAttention API. - cu_seqlens_cpu (torch.LongTensor): - CPU copy of `cu_seqlens` to avoid unnecessary device synchronization. Default: `None`. - - Returns: - o (torch.Tensor): - Outputs of shape `[B, T, H, V]`. - attn (torch.Tensor): - Attention scores of shape `[B, H, T, T]` if `output_attentions=True` else `None`. - """ - if 'head_first' in kwargs: - raise DeprecationWarning( - "head_first has been removed. Inputs must be in `[B, T, H, ...]` format.", - ) - if cu_seqlens is not None: - if q.shape[0] != 1: - raise ValueError( - f"The batch size is expected to be 1 rather than {q.shape[0]} when using `cu_seqlens`. " - f"Please flatten variable-length inputs before processing.", - ) - if output_attentions: - assert cu_seqlens is None, "output_attentions=True is not supported with variable-length sequences" - - if scale is None: - scale = k.shape[-1] ** -0.5 - o, attn = ParallelSimpleGLAFunction.apply( - q, - k, - v, - g, - scale, - output_attentions, - cu_seqlens, - cu_seqlens_cpu, - ) - return o, attn diff --git a/build/torch-xpu/ops/titans/__init__.py b/build/torch-xpu/ops/titans/__init__.py deleted file mode 100644 index a72421146f4875819a09561f95f89406c683b58e..0000000000000000000000000000000000000000 --- a/build/torch-xpu/ops/titans/__init__.py +++ /dev/null @@ -1,12 +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 .naive import chunk_titans_linear - -__all__ = [ - 'chunk_titans_linear', -] diff --git a/build/torch-xpu/ops/titans/log_impl.py b/build/torch-xpu/ops/titans/log_impl.py deleted file mode 100644 index 70af2a4ff29d0da4d9326690f68d022228c71470..0000000000000000000000000000000000000000 --- a/build/torch-xpu/ops/titans/log_impl.py +++ /dev/null @@ -1,160 +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 - - -def cal_n_log(log_theta, log_eta, seq_len): - """ - calculate n_{i,j} in log space - log(n_{i,j}) = log(θ_j) + sum_{k=j+1}^i log(η_k) - """ - # create log(n) - log_n = torch.zeros(*log_theta.shape, seq_len, dtype=log_eta.dtype).to( - log_eta.device, - ) # [batch_size, num_heads, seq_len, seq_len] - for i in range(seq_len): - for j in range(i + 1): - if i == j: - log_n[..., j, i] = log_theta[..., j] - else: - log_n[..., j, i] = log_theta[..., j] + torch.sum( - log_eta[..., j + 1: i + 1], dim=-1, - ) - - return log_n - - -def cal_f_log(log_beta, seq_len, log_m): - """ - cal_f_log(log_beta, seq_len, log_m) -> f - log(f_t) = log(sum_{i=1}^t exp(sum_{k=i+1}^t log(1-α_k) + sum_{k=1}^i log(η_k))) - """ - # create f - # f = torch.zeros_like(log_beta) - # for t in range(seq_len): - # for i in range(t + 1): - # f[..., t] += torch.exp(log_beta[..., t] - log_beta[..., i] + log_m[..., i]) - log_f = torch.zeros_like(log_beta) - for t in range(seq_len): - a_i = log_beta[..., t: t + 1] - log_beta[..., : t + 1] + log_m[..., : t + 1] - log_f[..., t] = torch.logsumexp(a_i, dim=-1) - f = torch.exp(log_f) - - # this version overflow and even slower - # t_indices = torch.arange(seq_len, device=log_beta.device) - # i_indices = torch.arange(seq_len, device=log_beta.device) - # - # mask = i_indices.unsqueeze(0) <= t_indices.unsqueeze(1) - # log_beta_t = log_beta.unsqueeze(-1) # [..., seq_len, 1] - # log_beta_i = log_beta.unsqueeze(-2) # [..., 1, seq_len] - # log_m_i = log_m.unsqueeze(-2) - # a_i = log_beta_t - log_beta_i + log_m_i - # masked_a_i = torch.where(mask, a_i, torch.tensor(-float('inf'), device=a_i.device, dtype=a_i.dtype)) - # log_f = torch.logsumexp(masked_a_i, dim=-1) # [..., seq_len] - # - # f = torch.exp(log_f) - return f - - -def cal_G_log(log_beta, log_n, seq_len): - """ - calculate G_{i,j} - log(G_{i,j}) = log(sum_{k=j}^i exp(log(β_i/β_k) + log(n_{k,j}))) - """ - # G = torch.zeros(*log_beta.shape[:-1], seq_len, seq_len, device = log_beta.device) - # # Fill in the lower triangular part - # for i in range(seq_len): # row - # for j in range(i + 1): # column - # # Sum from k=j to i - # for k in range(j, i + 1): - # G[..., i, j] += torch.exp(log_beta[..., i] - log_beta[..., k] + log_n[..., j, k]) - - log_G = torch.full( - (*log_beta.shape[:-1], seq_len, seq_len), float("-inf"), device=log_beta.device, - ) - # fill in the lower triangular part - for i in range(seq_len): # row - for j in range(i + 1): # column - terms = ( - log_beta[..., i: i + 1] - - log_beta[..., j: i + 1] - + log_n[..., j: j + 1, j: i + 1].squeeze(-2) - ) - # use logsumexp to avoid overflow - log_G[..., i, j] = torch.logsumexp(terms, dim=-1) - - G = torch.exp(log_G) - return G - - -def _combine_params_log(log_theta, log_alpha_complement, log_eta, seq_len): - """ - Update rule for Titans in log space - - Parameters: - - log_theta: log(θ) - - log_alpha_complement: log(1-α) - - log_eta: log(η) - - seq_len: sequence length - - Returns: - - log_beta, beta_T, log_f, f_T, log_g, log_G, m_T, n_T - """ - # calculate log(β_t) = sum_{k=1}^t log(1-α_k) - log_beta = torch.cumsum(log_alpha_complement, dim=-1) - - # get β_T - beta_T = torch.exp(log_beta[..., -1]) - - # calculate log(m_i) = sum_{k=1}^i log(η_k) - log_m = torch.cumsum(log_eta, dim=-1) - m_T = torch.exp(log_m[..., -1]) - - # cal log(n_{i,j}) - log_n = cal_n_log(log_theta, log_eta, seq_len) - n_T = torch.exp(log_n[..., -1]) - - # cal log(f_t) - f = cal_f_log(log_beta, seq_len, log_m) - f_T = f[..., -1] - - # cal log(G_{i,j}) - G = cal_G_log(log_beta, log_n, seq_len) - # get log(g_j) = log(G_{T,j}) - g = G[..., -1, :] - - return log_beta, beta_T, f, f_T, g, G, m_T, n_T - - -def combine_params_log(theta, alpha, eta, seq_len): - """ - log space Titians - - Parameters: - - theta: θ - - alpha: α - - eta: η - - seq_len: sequence length - - Returns: - - beta, beta_T, f, f_T, g, G, m_T, n_T - """ - # convert to log space - log_theta = torch.log(theta.squeeze(-1)) - log_alpha_complement = torch.log(1 - alpha.squeeze(-1)) - log_eta = torch.log(eta.squeeze(-1)) - - # combine params in log space - log_beta, beta_T, f, f_T, g, G, m_T, n_T = _combine_params_log( - log_theta, log_alpha_complement, log_eta, seq_len, - ) - - # convert back to normal space - beta = torch.exp(log_beta) - - return beta, beta_T, f, f_T, g, G, m_T, n_T diff --git a/build/torch-xpu/ops/titans/naive.py b/build/torch-xpu/ops/titans/naive.py deleted file mode 100644 index 926404ac30cdc46c9addb1cb957469020aff451f..0000000000000000000000000000000000000000 --- a/build/torch-xpu/ops/titans/naive.py +++ /dev/null @@ -1,380 +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 torch.nn.functional as F - -from ...ops.titans.log_impl import combine_params_log - - -def cal_n(theta, eta, seq_len): - n = torch.zeros(*theta.shape, seq_len, dtype=theta.dtype).to( - theta.device, - ) # [batch_size, num_heads, seq_len, seq_len] - - # 1. deal with diagonal elements - indices = torch.arange(seq_len, device=theta.device) - n[..., indices, indices] = theta[..., indices] - - # 2. Create a cumulative product matrix - # First create a mask to mark the positions where eta needs to be multiplied - mask = torch.triu(torch.ones(seq_len, seq_len), diagonal=1).to(theta.device) - # Convert mask to boolean type - mask = mask.bool() - # Expand eta to match the target shape - eta_expanded = eta.unsqueeze(-2).expand(*theta.shape[:-1], seq_len, seq_len) - # Create a matrix filled with 1s for cumulative product - cumulative = torch.ones_like(eta_expanded) - cumulative = torch.where(mask, eta_expanded, cumulative) - # Calculate the cumulative product - cumulative_prod = torch.cumprod(cumulative, dim=-1) - - # 3. Calculate non-diagonal elements - # Create an expanded version of theta - theta_expanded = theta.unsqueeze(-1).expand(*theta.shape[:-1], seq_len, seq_len) - # Create a mask to keep only the upper triangular part (excluding the diagonal) - upper_triangular = torch.triu(torch.ones_like(n), diagonal=1).bool() - # Combine theta and cumulative product - n = torch.where(upper_triangular, theta_expanded * cumulative_prod, n) - return n - - -def cal_f(beta, seq_len, m): - a = torch.tril(beta.to(torch.float32).unsqueeze(-1).expand(*beta.shape, seq_len), 0) - ratio = (m.to(torch.float32) / beta.to(torch.float32)).unsqueeze(-1) - f = torch.matmul(a, ratio).squeeze(-1) - return f.to(beta.dtype) - - -def cal_G(beta, n, seq_len): - i_indices = torch.arange(seq_len, device=beta.device) - j_indices = torch.arange(seq_len, device=beta.device) - k_indices = torch.arange(seq_len, device=beta.device) - beta_ratio = beta[..., :, None] / beta[..., None, :] # [..., i, k] - - # create mask - k_mask = (k_indices[None, None, :] >= j_indices[None, :, None]) & ( - k_indices[None, None, :] <= i_indices[:, None, None] - ) - - # use mask to filter out invalid values - masked_beta_ratio = beta_ratio[..., :, None, :] * k_mask # [..., i, j, k] - masked_n = n[..., None, :, :] * k_mask # [..., i, j, k] - # calculate G - G = torch.sum(masked_beta_ratio * masked_n, dim=-1) # [..., i, j] - return G - - -def combine_params(theta, alpha, eta, seq_len): - theta = theta.squeeze(-1) - eta = eta.squeeze(-1) - alpha = alpha.squeeze(-1) - beta = torch.cumprod(1 - alpha, dim=-1) # β_t = ∏(1 - α_t) in titans paper - beta_T = beta[..., -1] # β_T - # Calculate m_i = ∏(k=1 to i) η_k - m = torch.cumprod(eta, dim=-1) # [batch_size, num_heads, seq_len] - m_T = m[..., -1] # m_T - # Calculate n_{i,j} - # We need to calculate ∏(k=j+1 to i) η_k for each i,j pair - # # this may be optimized - # n = torch.zeros(*theta.shape, seq_len, dtype = theta.dtype).to( - # theta.device) # [batch_size, num_heads, seq_len, seq_len] - # for i in range(seq_len): - # for j in range(i + 1): - # if i == j: - # n[..., j, i] = theta[..., j] - # else: - # # Calculate product of eta from j+1 to i - # eta_product = torch.prod(eta[..., j + 1:i + 1], dim = -1) - # n[..., j, i] = theta[..., j] * eta_product - - n = cal_n(theta, eta, seq_len) - n_T = n[..., -1] # [batch_size, num_heads, seq_len] - # Calculate f_t = ∑(i=1 to t) (β_t/β_i) m_i - # f = torch.zeros_like(theta) - # for t in range(seq_len): - # for i in range(t + 1): - # f[..., t] += (beta[..., t] / beta[..., i]) * m[..., i] - f = cal_f(beta, seq_len, m) - f_T = f[..., -1] # [batch_size, num_heads, seq_len] - # Calculate g_j = ∑(i=j to t) (β_t/β_i) n_{i,j} - # g = torch.zeros_like(theta) # [batch_size, num_heads, seq_len] - # for j in range(seq_len): - # for i in range(j, seq_len): - # g[..., j] += (beta[..., -1] / beta[..., i]) * n[..., j, i] - # G = torch.zeros(*beta.shape[:-1], seq_len, seq_len, device = beta.device) - # # Fill in the lower triangular part - # for i in range(seq_len): # row - # for j in range(i + 1): # column - # # Sum from k=j to i - # for k in range(j, i + 1): - # G[..., i, j] += (beta[..., i] / beta[..., k]) * n[..., j, k] - G = cal_G(beta, n, seq_len) - g = G[:, :, -1, :] # [batch_size, num_heads, seq_len] - # g2, G2 = compute_g_and_G(beta, n, seq_len) - return beta, beta_T, f, f_T, g, G, m_T, n_T - - -def titans_linear( - q, k, v, w, b, theta, alpha, eta, eps, chunk_size, initial_state, output_final_state, -): - """ - Implementation of Titans Linear function based on the update rules: - M_t = (1 - alpha_t) * M_{t-1} + S_t - S_t = eta_t * S_{t-1} - theta_t * nabla_l(M_{t-1}; x_t) - - Args: - q: Query tensor - k: Key tensor - v: Value tensor - w: Weight tensor - b: Bias tensor - theta: Learning rate tensor - alpha: Momentum decay tensor - eta: Step size tensor - eps: Epsilon for numerical stability - initial_state: Initial state M_0 - output_final_state: Whether to output the final state - - Returns: - Tuple of (output tensor, final state) - """ - B, H, T, D = q.shape - device = q.device - w = w.reshape(H, 1, D).to(torch.float32) - b = b.reshape(H, 1, D).to(torch.float32) - # Initialize states - if initial_state is None: - M_prev = torch.zeros(B, H, D, D, device=device) - else: - M_prev = initial_state - M_prev_nabla = M_prev.clone() - S_prev = torch.zeros_like(M_prev) - outputs = [] - - # Process sequence step by step - for t in range(T): - # Get current step inputs - q_t = q[:, :, t: t + 1, :] # (batch_size, num_heads, 1, dim) - k_t = k[:, :, t: t + 1, :] # (batch_size, num_heads, 1, dim) - v_t = v[:, :, t: t + 1, :] # (batch_size, num_heads, 1, dim) - theta_t = theta[:, :, t: t + 1, :] # (batch_size, num_heads, 1, dim) - alpha_t = alpha[:, :, t: t + 1, :] # (batch_size, num_heads, 1, dim) - eta_t = eta[:, :, t: t + 1, :] # (batch_size, num_heads, 1, dim) - - # Compute gradient - km = k_t @ M_prev_nabla # (batch_size, num_heads, 1, dim) - reconstruction_target = v_t - k_t - mean = km.mean(-1, keepdim=True) - var = km.var(-1, unbiased=False, keepdim=True).to(torch.float32) - rstd = torch.sqrt(var + eps).to(torch.float32) - km_hat = (km - mean) / rstd - - grad = w * km_hat + b - reconstruction_target - grad = grad * w - # v_new = (D * grad - grad.sum(-1, keepdim = True) - km_hat * (grad * km_hat).sum(-1, keepdim = True)) / ( - # rstd * D) - v_new = D * grad - grad.sum(-1, keepdim=True) / (rstd * D) - proj_term = km_hat * (grad * km_hat).sum(-1, keepdim=True) / (rstd * D) - v_new = v_new - proj_term - # v_new = grad - - # Update S_t - S_t = eta_t * S_prev - 2 * theta_t * k_t.transpose(-2, -1) @ v_new - - # Update M_t - M_t = (1 - alpha_t) * M_prev + S_t - - # Store output - output_t = q_t @ M_t # (batch_size, num_heads, seq_len, dim) - mean = output_t.mean(dim=-1, keepdim=True) - var = output_t.var(dim=-1, unbiased=False, keepdim=True).to(torch.float32) - rstd = torch.sqrt(var + eps).to(torch.float32) - output_t = output_t + (output_t - mean) / rstd * w + b - outputs.append(output_t) - - # Update states for next step - if (t + 1) % chunk_size == 0: - M_prev_nabla = M_t.clone() - M_prev = M_t - S_prev = S_t - - # Stack outputs along sequence dimension - output = torch.stack(outputs, dim=-2).squeeze( - -3, - ) # (batch_size, num_heads, seq_len, dim) - - if output_final_state: - return output, M_prev - return output, None - - -def chunk_titans_linear( - q, k, v, w, b, theta, alpha, eta, eps, chunk_size, initial_state, output_final_state, -): - B, H, T, D = q.shape - num_batch = T // chunk_size - # [num_batch, B, num_heads, mini_batch_size, head_dim] - _q = q.reshape(B, H, num_batch, chunk_size, D).permute(2, 0, 1, 3, 4) - _k = k.reshape(B, H, num_batch, chunk_size, D).permute(2, 0, 1, 3, 4) - _v = v.reshape(B, H, num_batch, chunk_size, D).permute(2, 0, 1, 3, 4) - # [num_batch, B, num_heads, mini_batch_size, 1] - _eta = eta.reshape(B, H, num_batch, chunk_size, 1).permute(2, 0, 1, 3, 4) - _theta = theta.reshape(B, H, num_batch, chunk_size, 1).permute(2, 0, 1, 3, 4) - _alpha = alpha.reshape(B, H, num_batch, chunk_size, 1).permute(2, 0, 1, 3, 4) - # [H, 1, D] - w = w.reshape(H, 1, D).to(torch.float32) - b = b.reshape(H, 1, D).to(torch.float32) - # [num_heads, 1, head_dim] - if initial_state is None: - M_prev = torch.zeros((B, H, D, D), device=v.device, dtype=v.dtype).to( - torch.float32, - ) - else: - M_prev = initial_state - - S_prev = torch.zeros_like(M_prev) - - # [num_batch, B, num_heads, mini_batch_size, head_dim] - o = torch.empty_like(_v) - - for i in range(num_batch): - q_i, k_i, v_i, eta_i, theta_i, alpha_i = [ - x[i] for x in [_q, _k, _v, _eta, _theta, _alpha] - ] - - # beta, beta_T, f, f_T, g, G, m_T, n = combine_params(theta_i, alpha_i, eta_i, chunk_size) - beta, beta_T, f, f_T, g, G, m_T, n = combine_params_log( - theta_i, alpha_i, eta_i, chunk_size, - ) - - m_T = m_T.unsqueeze(-1).unsqueeze(-1) - beta_T = beta_T.unsqueeze(-1).unsqueeze(-1) - f_T = f_T.unsqueeze(-1).unsqueeze(-1) - g_diag = torch.diag_embed(g).to(q_i.dtype) - n = torch.diag_embed(n).to(q_i.dtype) - beta = torch.diag_embed(beta).to(q_i.dtype) - f = torch.diag_embed(f).to(q_i.dtype) - km = k_i @ M_prev - reconstruction_target = v_i - k_i - - mean = km.mean(-1, True) - var = km.var(-1, unbiased=False, keepdim=True).to(torch.float32) - rstd = torch.sqrt(var + eps).to(torch.float32) - km_hat = (km - mean) / rstd - - grad = w * km_hat + b - reconstruction_target - grad *= w - v_new = D * grad - grad.sum(-1, keepdim=True) / (rstd * D) - proj_term = km_hat * (grad * km_hat).sum(-1, keepdim=True) / (rstd * D) - v_new = v_new - proj_term - # v_new = (D * grad - grad.sum(-1, True)) - # print(f"Projection term stats: min={torch.abs(beta_T).min()}") - - # v_new = grad - - Attn = torch.tril(q_i @ k_i.transpose(-2, -1)) * G - - # o_i - output_t = beta @ q_i @ M_prev + f @ q_i @ S_prev - 2 * Attn @ v_new - - M_t = ( - beta_T * M_prev - + f_T * S_prev - - 2 * (g_diag @ k_i).transpose(-1, -2) @ v_new - ) - # cal S_T from S_0 - S_t = m_T * S_prev - 2 * (n @ k_i).transpose(-1, -2) @ v_new - # layer norm with residuals - mean = output_t.mean(dim=-1, keepdim=True) - var = output_t.var(dim=-1, unbiased=False, keepdim=True).to(torch.float32) - rstd = torch.sqrt(var + eps).to(torch.float32) - output_t = output_t + (output_t - mean) / rstd * w + b - o[i] = output_t - S_prev = S_t - M_prev = M_t - - # [B, num_mini_batch, mini_batch_size, num_heads, head_dim] - o = o.permute(1, 2, 0, 3, 4).reshape(B, H, T, D) - M_prev = M_prev if output_final_state else None - return o, M_prev - - -# most of the code is copied from ttt -def chunk_titans_linear_ref( - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - w: torch.Tensor, - b: torch.Tensor, - theta: torch.Tensor, - alpha: torch.Tensor, - eta: torch.Tensor, - eps: float = 1e-6, - chunk_size: int = 16, # chunk size - initial_state: torch.Tensor = None, - output_final_state: bool = False, - head_first: bool = False, - use_chunk: bool = True, -): - assert q.dtype == k.dtype == v.dtype - assert k.shape[-1] == v.shape[-1], "DK must equal to DV." - if not head_first: - q = q.transpose(1, 2) - k = k.transpose(1, 2) - v = v.transpose(1, 2) - eta = eta.transpose(1, 2) - alpha = alpha.transpose(1, 2) - theta = theta.transpose(1, 2) - seq_len = q.shape[-2] - pad_len = (chunk_size - (seq_len % chunk_size)) % chunk_size - if pad_len > 0: - q = F.pad(q, (0, 0, 0, pad_len)) - k = F.pad(k, (0, 0, 0, pad_len)) - v = F.pad(v, (0, 0, 0, pad_len)) - theta = F.pad(theta, (0, 0, 0, pad_len)) - alpha = F.pad(alpha, (0, 0, 0, pad_len)) - eta = F.pad(eta, (0, 0, 0, pad_len)) - theta[:, :, -1, :] = theta[:, :, -(pad_len + 1), :] - alpha[:, :, -1, :] = alpha[:, :, -(pad_len + 1), :] - eta[:, :, -1, :] = eta[:, :, -(pad_len + 1), :] - assert q.shape[-2] % chunk_size == 0, "Sequence length should be a multiple of BT." - q, k, v, w, b = map(lambda x: x.to(torch.float32), [q, k, v, w, b]) - if use_chunk: - o, final_state = chunk_titans_linear( - q, - k, - v, - w, - b, - theta, - alpha, - eta, - eps, - chunk_size, - initial_state, - output_final_state, - ) - else: - o, final_state = titans_linear( - q, - k, - v, - w, - b, - theta, - alpha, - eta, - eps, - chunk_size, - initial_state, - output_final_state, - ) - o = o[:, :, :seq_len, :] - if not head_first: - o = o.transpose(1, 2) - return o, final_state diff --git a/build/torch-xpu/ops/ttt/__init__.py b/build/torch-xpu/ops/ttt/__init__.py deleted file mode 100644 index a1e88509195dc1e2a4420432a884610434b88951..0000000000000000000000000000000000000000 --- a/build/torch-xpu/ops/ttt/__init__.py +++ /dev/null @@ -1,14 +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 .chunk import chunk_ttt_linear -from .fused_chunk import fused_chunk_ttt_linear - -__all__ = [ - 'chunk_ttt_linear', - 'fused_chunk_ttt_linear', -] diff --git a/build/torch-xpu/ops/ttt/chunk.py b/build/torch-xpu/ops/ttt/chunk.py deleted file mode 100644 index 3e5234997035702273550a2ee92f1c72b4655a89..0000000000000000000000000000000000000000 --- a/build/torch-xpu/ops/ttt/chunk.py +++ /dev/null @@ -1,1473 +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 torch.nn.functional as F -import triton -import triton.language as tl - -from ...modules.layernorm import group_norm -from ...ops.utils import prepare_chunk_indices, prepare_chunk_offsets -from ...utils import autocast_custom_bwd, autocast_custom_fwd, autotune_cache_kwargs, input_guard - - -@triton.heuristics({ - 'USE_INITIAL_STATE': lambda args: args['h0'] is not None, - 'USE_INITIAL_STATE_B': lambda args: args['hb0'] is not None, - 'STORE_FINAL_STATE': lambda args: args['ht'] is not None, - 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None, -}) -@triton.autotune( - configs=[ - triton.Config({}, num_warps=num_warps) - for num_warps in [1, 2, 4, 8] - ], - key=['BT', 'BK', 'BV'], - **autotune_cache_kwargs, -) -@triton.jit(do_not_specialize=['T']) -def chunk_ttt_linear_fwd_kernel_h( - k, - v, - v_new, - eta, - w, - b, - eps, - h, - hb, - h0, - hb0, - ht, - hbt, - cu_seqlens, - chunk_offsets, - T, - H: tl.constexpr, - K: tl.constexpr, - V: tl.constexpr, - BT: tl.constexpr, - BK: tl.constexpr, - BV: tl.constexpr, - USE_INITIAL_STATE: tl.constexpr, - USE_INITIAL_STATE_B: tl.constexpr, - STORE_FINAL_STATE: tl.constexpr, - IS_VARLEN: tl.constexpr, -): - i_k, i_v, i_nh = tl.program_id(0), tl.program_id(1), tl.program_id(2) - i_n, i_h = i_nh // H, i_nh % H - if IS_VARLEN: - bos, eos = tl.load(cu_seqlens + i_n).to(tl.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) - T = eos - bos - NT = tl.cdiv(T, BT) - boh = tl.load(chunk_offsets + i_n).to(tl.int32) - else: - bos, eos = i_n * T, i_n * T + T - NT = tl.cdiv(T, BT) - boh = i_n * NT - - # [BK, BV] - b_h = tl.zeros([BK, BV], dtype=tl.float32) - # [BV] - b_hb = tl.zeros([BV], dtype=tl.float32) - if USE_INITIAL_STATE: - p_h0 = tl.make_block_ptr(h0 + i_nh * K * V, (K, V), (V, 1), (i_k * BK, i_v * BV), (BK, BV), (1, 0)) - b_h = tl.load(p_h0, boundary_check=(0, 1), padding_option="zero").to(tl.float32) - if USE_INITIAL_STATE_B: - p_hb0 = tl.make_block_ptr(hb0 + i_nh * V, (V,), (1,), (i_v * BV,), (BV,), (0,)) - b_hb = tl.load(p_hb0, boundary_check=(0,), padding_option="zero").to(tl.float32) - - offs = tl.arange(0, BV) - b_w = tl.load(w + i_h * V + offs, mask=offs < V, other=0.) - b_b = tl.load(b + i_h * V + offs, mask=offs < V, other=0.) - - for i_t in range(NT): - p_h = tl.make_block_ptr(h + ((boh + i_t) * H + i_h) * K*V, (K, V), (V, 1), (i_k * BK, i_v * BV), (BK, BV), (1, 0)) - p_hb = tl.make_block_ptr(hb + ((boh + i_t) * H + i_h) * V, (V,), (1,), (i_v * BV,), (BV,), (0,)) - tl.store(p_h, b_h.to(p_h.dtype.element_ty), boundary_check=(0, 1)) - tl.store(p_hb, b_hb.to(p_hb.dtype.element_ty), boundary_check=(0,)) - p_k = tl.make_block_ptr(k+(bos*H+i_h)*K, (K, T), (1, H*K), (i_k * BK, i_t * BT), (BK, BT), (0, 1)) - p_v = tl.make_block_ptr(v+(bos*H+i_h)*V, (T, V), (H*V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) - p_v_new = tl.make_block_ptr(v_new+(bos*H+i_h)*V, (T, V), (H*V, 1), (i_t*BT, i_v * BV), (BT, BV), (1, 0)) - p_eta_last = eta+bos*H+i_h + (T-1)*H if i_t == NT-1 else eta+bos*H+i_h + (i_t*BT+BT-1)*H - b_k = tl.load(p_k, boundary_check=(0, 1), padding_option="zero") - b_v = tl.load(p_v, boundary_check=(0, 1), padding_option="zero") - - b_kh = tl.dot(tl.trans(b_k), b_h.to(b_k.dtype), allow_tf32=False).to(tl.float32) + b_hb[None, :] - b_kh = tl.where((offs < V)[None, :], b_kh, 0.) - mean = tl.sum(b_kh, axis=1, keep_dims=True) / V - xbar = tl.where((offs < V)[None, :], b_kh - mean, 0.) - var = tl.sum(xbar * xbar, axis=1, keep_dims=True) / V - rstd = 1 / tl.sqrt(var.to(tl.float32) + eps) - b_kh_hat = (b_kh - mean) * rstd - - b_v = b_kh_hat.to(b_k.dtype) * b_w[None, :].to(b_k.dtype) + \ - b_b[None, :].to(b_k.dtype) - b_v.to(b_k.dtype) + tl.trans(b_k) - b_v = tl.where((offs < V)[None, :], b_v * b_w[None, :].to(b_k.dtype), 0.) - b_v2 = rstd * (V * b_v - tl.sum(b_v, axis=1, keep_dims=True) - b_kh_hat.to(b_k.dtype) - * tl.sum(b_v * b_kh_hat.to(b_k.dtype), axis=1, keep_dims=True)) / V - tl.store(p_v_new, b_v2.to(p_v_new.dtype.element_ty), boundary_check=(0, 1)) - b_eta_last = tl.load(p_eta_last) - b_h = b_h - tl.dot(b_eta_last * b_k, b_v2.to(b_k.dtype), allow_tf32=False) - b_hb = b_hb - tl.sum(b_eta_last * b_v2.to(b_k.dtype), axis=0) - - if STORE_FINAL_STATE: - p_ht = tl.make_block_ptr(ht + i_nh * K*V, (K, V), (V, 1), (i_k * BK, i_v * BV), (BK, BV), (1, 0)) - p_hbt = tl.make_block_ptr(hbt + i_nh * V, (V,), (1,), (i_v * BV,), (BV,), (0,)) - tl.store(p_ht, b_h.to(p_ht.dtype.element_ty), boundary_check=(0, 1)) - tl.store(p_hbt, b_hb.to(p_hbt.dtype.element_ty), boundary_check=(0,)) - - -@triton.heuristics({ - 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None, -}) -@triton.autotune( - configs=[ - triton.Config({}, num_warps=num_warps, num_stages=num_stages) - for num_warps in [2, 4, 8] - for num_stages in [2, 3] - ], - key=['BT'], - **autotune_cache_kwargs, -) -@triton.jit(do_not_specialize=['T']) -def chunk_ttt_linear_fwd_kernel_o( - q, - k, - v, - eta, - h, - hb, - o, - cu_seqlens, - chunk_indices, - scale, - T, - H: tl.constexpr, - K: tl.constexpr, - V: tl.constexpr, - BT: tl.constexpr, - BK: tl.constexpr, - BV: tl.constexpr, - IS_VARLEN: tl.constexpr, -): - i_v, i_t, i_bh = tl.program_id(0), tl.program_id(1), tl.program_id(2) - i_b, i_h = i_bh // H, i_bh % H - - 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.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) - T = eos - bos - NT = tl.cdiv(T, BT) - else: - NT = tl.cdiv(T, BT) - i_tg = i_b * NT + i_t - bos, eos = i_b * T, i_b * T + T - - # offset calculation - q += (bos * H + i_h) * K - k += (bos * H + i_h) * K - v += (bos * H + i_h) * V - eta += bos * H + i_h - o += (bos * H + i_h) * V - h += (i_tg * H + i_h) * K * V - hb += (i_tg * H + i_h) * V - stride_qk = H*K - stride_vo = H*V - stride_eta = H - - p_q = tl.make_block_ptr(q, (T, K), (stride_qk, 1), (i_t * BT, 0), (BT, BK), (1, 0)) - p_k = tl.make_block_ptr(k, (K, T), (1, stride_qk), (0, i_t * BT), (BK, BT), (0, 1)) - p_eta = tl.make_block_ptr(eta, (T,), (stride_eta,), (i_t * BT,), (BT,), (0,)) - p_h = tl.make_block_ptr(h, (K, V), (V, 1), (0, i_v * BV), (BK, BV), (1, 0)) - p_hb = tl.make_block_ptr(hb, (V,), (1,), (i_v * BV,), (BV,), (0,)) - # [BT, BK] - b_q = tl.load(p_q, boundary_check=(0, 1), padding_option="zero") - # [BK, BT] - b_k = tl.load(p_k, boundary_check=(0, 1), padding_option="zero") - # [BT, 1] - b_eta = tl.load(p_eta, boundary_check=(0,), padding_option="zero") - # [BK, BV] - b_h = tl.load(p_h, boundary_check=(0, 1), padding_option="zero") - # [BV] - b_hb = tl.load(p_hb, boundary_check=(0,), padding_option="zero") - # [BT, BK] @ [BK, BV] -> [BT, BV] - b_o = tl.dot(b_q, b_h, allow_tf32=False) - # [BT, BK] @ [BK, BT] -> [BT, BT] - b_A = tl.dot(b_q, b_k, allow_tf32=False) - - o_i = tl.arange(0, BT) - m_A = o_i[:, None] >= o_i[None, :] - b_A = tl.where(m_A, b_A, 0) - b_Ae = tl.where(m_A, b_eta[:, None], 0.0) - - p_v = tl.make_block_ptr(v, (T, V), (stride_vo, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) - p_o = tl.make_block_ptr(o, (T, V), (stride_vo, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) - b_v = tl.load(p_v, boundary_check=(0, 1), padding_option="zero") - b_o = (b_o - tl.dot(b_eta[:, None] * b_A.to(b_v.dtype), b_v, allow_tf32=False)) * scale - b_o += b_hb[None, :] - tl.dot(b_Ae.to(b_v.dtype), b_v, allow_tf32=False) - tl.store(p_o, b_o.to(p_o.dtype.element_ty), boundary_check=(0, 1)) - - -@triton.heuristics({ - 'USE_INITIAL_STATE': lambda args: args['h0'] is not None, - 'USE_INITIAL_STATE_B': lambda args: args['hb0'] is not None, - 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None, -}) -@triton.autotune( - configs=[ - triton.Config({}, num_warps=num_warps) - for num_warps in [1, 2, 4, 8] - ], - key=['BT', 'BK', 'BV'], - **autotune_cache_kwargs, -) -@triton.jit(do_not_specialize=['T']) -def chunk_ttt_linear_bwd_kernel_h( - k, - v, - v_new, - eta, - w, - b, - eps, - h, - h0, - hb0, - x, - y, - r, - cu_seqlens, - chunk_offsets, - T, - H: tl.constexpr, - K: tl.constexpr, - V: tl.constexpr, - BT: tl.constexpr, - BK: tl.constexpr, - BV: tl.constexpr, - NT: tl.constexpr, - USE_INITIAL_STATE: tl.constexpr, - USE_INITIAL_STATE_B: tl.constexpr, - IS_VARLEN: tl.constexpr, -): - i_k, i_v, i_nh = tl.program_id(0), tl.program_id(1), tl.program_id(2) - i_n, i_h = i_nh // H, i_nh % H - if IS_VARLEN: - bos, eos = tl.load(cu_seqlens + i_n).to(tl.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) - T = eos - bos - NT = tl.cdiv(T, BT) - boh = tl.load(chunk_offsets + i_n).to(tl.int32) - else: - bos, eos = i_n * T, i_n * T + T - NT = tl.cdiv(T, BT) - boh = i_n * NT - - # [BK, BV] - b_h = tl.zeros([BK, BV], dtype=tl.float32) - # [BV] - b_hb = tl.zeros([BV], dtype=tl.float32) - if USE_INITIAL_STATE: - p_h0 = tl.make_block_ptr(h0 + i_nh * K * V, (K, V), (V, 1), (i_k * BK, i_v * BV), (BK, BV), (1, 0)) - b_h = tl.load(p_h0, boundary_check=(0, 1), padding_option="zero").to(tl.float32) - if USE_INITIAL_STATE_B: - p_hb0 = tl.make_block_ptr(hb0 + i_nh * V, (V,), (1,), (i_v * BV,), (BV,), (0,)) - b_hb = tl.load(p_hb0, boundary_check=(0,), padding_option="zero").to(tl.float32) - - offs = tl.arange(0, BV) - b_w = tl.load(w + i_h * V + offs, mask=offs < V, other=0.) - b_b = tl.load(b + i_h * V + offs, mask=offs < V, other=0.) - - for i_t in range(NT): - p_h = tl.make_block_ptr(h + ((boh + i_t) * H + i_h) * K*V, (K, V), (V, 1), (i_k * BK, i_v * BV), (BK, BV), (1, 0)) - tl.store(p_h, b_h.to(p_h.dtype.element_ty), boundary_check=(0, 1)) - p_k = tl.make_block_ptr(k+(bos*H+i_h)*K, (K, T), (1, H*K), (i_k * BK, i_t * BT), (BK, BT), (0, 1)) - p_v = tl.make_block_ptr(v+(bos*H+i_h)*V, (T, V), (H*V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) - p_v_new = tl.make_block_ptr(v_new+(bos*H+i_h)*V, (T, V), (H*V, 1), (i_t*BT, i_v * BV), (BT, BV), (1, 0)) - p_x = tl.make_block_ptr(x+(bos*H+i_h)*V, (T, V), (H*V, 1), (i_t*BT, i_v * BV), (BT, BV), (1, 0)) - p_y = tl.make_block_ptr(y+(bos*H+i_h)*V, (T, V), (H*V, 1), (i_t*BT, i_v * BV), (BT, BV), (1, 0)) - p_r = tl.make_block_ptr(r+bos*H+i_h, (T, 1), (H, 1), (i_t*BT, 0), (BT, 1), (1, 0)) - p_eta_last = eta+bos*H+i_h + (T-1)*H if i_t == NT-1 else eta+bos*H+i_h + (i_t*BT+BT-1)*H - b_k = tl.load(p_k, boundary_check=(0, 1), padding_option="zero") - b_v = tl.load(p_v, boundary_check=(0, 1), padding_option="zero") - - b_kh = tl.dot(tl.trans(b_k), b_h.to(b_k.dtype), allow_tf32=False).to(tl.float32) + b_hb[None, :] - b_kh = tl.where((offs < V)[None, :], b_kh, 0.) - mean = tl.sum(b_kh, axis=1, keep_dims=True) / V - xbar = tl.where((offs < V)[None, :], b_kh - mean, 0.) - var = tl.sum(xbar * xbar, axis=1, keep_dims=True) / V - rstd = 1 / tl.sqrt(var.to(tl.float32) + eps) - b_kh_hat = (b_kh - mean) * rstd - - b_v = b_kh_hat.to(b_k.dtype) * b_w[None, :].to(b_k.dtype) + \ - b_b[None, :].to(b_k.dtype) - b_v.to(b_k.dtype) + tl.trans(b_k) - b_v = tl.where((offs < V)[None, :], b_v * b_w[None, :].to(b_k.dtype), 0.) - b_v2 = rstd * (V * b_v - tl.sum(b_v, axis=1, keep_dims=True) - b_kh_hat.to(b_k.dtype) - * tl.sum(b_v * b_kh_hat.to(b_k.dtype), axis=1, keep_dims=True)) / V - tl.store(p_x, b_kh_hat.to(p_x.dtype.element_ty), boundary_check=(0, 1)) - tl.store(p_y, b_v.to(p_y.dtype.element_ty), boundary_check=(0, 1)) - tl.store(p_r, rstd.to(p_r.dtype.element_ty), boundary_check=(0, 1)) - tl.store(p_v_new, b_v2.to(p_v_new.dtype.element_ty), boundary_check=(0, 1)) - b_eta_last = tl.load(p_eta_last) - b_h = b_h - tl.dot(b_eta_last * b_k, b_v2.to(b_k.dtype), allow_tf32=False) - b_hb = b_hb - tl.sum(b_eta_last * b_v2.to(b_k.dtype), axis=0) - - -@triton.heuristics({ - 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None, -}) -@triton.autotune( - configs=[ - triton.Config({}, num_warps=num_warps) - for num_warps in [4] - ], - key=['BT', 'BK', 'BV'], - **autotune_cache_kwargs, -) -@triton.jit(do_not_specialize=['T']) -def chunk_ttt_linear_bwd_kernel_dv_local( - q, - k, - eta, - do, - dv, - cu_seqlens, - chunk_indices, - scale, - T, - H: tl.constexpr, - K: tl.constexpr, - V: tl.constexpr, - BT: tl.constexpr, - BK: tl.constexpr, - BV: tl.constexpr, - IS_VARLEN: tl.constexpr, -): - i_t, i_bh = tl.program_id(0), tl.program_id(1) - i_b, i_h = i_bh // H, i_bh % H - 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.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) - T = eos - bos - else: - bos, eos = i_b * T, i_b * T + T - - # offset calculation - q += (bos * H + i_h) * K - k += (bos * H + i_h) * K - eta += bos * H + i_h - do += (bos * H + i_h) * V - dv += (bos * H + i_h) * V - stride_qk = H*K - stride_vo = H*V - stride_eta = H - - b_A = tl.zeros([BT, BT], dtype=tl.float32) - for i_k in range(tl.cdiv(K, BK)): - p_k = tl.make_block_ptr(k, (T, K), (stride_qk, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) - p_q = tl.make_block_ptr(q, (K, T), (1, stride_qk), (i_k * BK, i_t * BT), (BK, BT), (0, 1)) - b_q = tl.load(p_q, boundary_check=(0, 1)) - b_k = tl.load(p_k, boundary_check=(0, 1)) - b_A += tl.dot(b_k, b_q) - - p_eta = tl.make_block_ptr(eta, (T,), (stride_eta,), (i_t * BT,), (BT,), (0,)) - b_eta = tl.load(p_eta, boundary_check=(0,)) - mask = (tl.arange(0, BT)[:, None] <= tl.arange(0, BT)[None, :]) - b_A = - tl.where(mask, b_A * scale * b_eta[None, :], 0).to(do.dtype.element_ty) - b_Ae = - tl.where(mask, b_eta[None, :], 0).to(do.dtype.element_ty) - - for i_v in range(tl.cdiv(V, BV)): - p_do = tl.make_block_ptr(do, (T, V), (stride_vo, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) - p_dv = tl.make_block_ptr(dv, (T, V), (stride_vo, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) - b_do = tl.load(p_do, boundary_check=(0, 1)) - b_dv = tl.dot(b_A.to(b_do.dtype), b_do) + tl.dot(b_Ae.to(b_do.dtype), b_do) - tl.store(p_dv, b_dv.to(p_dv.dtype.element_ty), boundary_check=(0, 1)) - - -@triton.heuristics({ - 'USE_FINAL_STATE_GRADIENT': lambda args: args['dht'] is not None, - 'USE_FINAL_STATE_GRADIENT_B': lambda args: args['dhbt'] is not None, - 'USE_INITIAL_STATE': lambda args: args['dh0'] is not None, - 'USE_INITIAL_STATE_B': lambda args: args['dhb0'] is not None, - 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None, -}) -@triton.autotune( - configs=[ - triton.Config({}, num_warps=num_warps) - for num_warps in [2, 4, 8, 16] - ], - key=['BT', 'BK', 'BV'], - **autotune_cache_kwargs, -) -@triton.jit(do_not_specialize=['T']) -def chunk_ttt_linear_bwd_kernel_norm( - q, - k, - v, - v_new, - x, - y, - r, - w, - b, - eta, - h, - dht, - dhbt, - dh0, - dhb0, - do, - dh, - dhb, - dv, - dv_new, - dk, - dw, - db, - cu_seqlens, - chunk_offsets, - scale, - T, - H: tl.constexpr, - K: tl.constexpr, - V: tl.constexpr, - BT: tl.constexpr, - BK: tl.constexpr, - BV: tl.constexpr, - USE_FINAL_STATE_GRADIENT: tl.constexpr, - USE_FINAL_STATE_GRADIENT_B: tl.constexpr, - USE_INITIAL_STATE: tl.constexpr, - USE_INITIAL_STATE_B: tl.constexpr, - IS_VARLEN: tl.constexpr, -): - i_k, i_v, i_nh = tl.program_id(0), tl.program_id(1), tl.program_id(2) - i_n, i_h = i_nh // H, i_nh % H - if IS_VARLEN: - bos, eos = tl.load(cu_seqlens + i_n).to(tl.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) - T = eos - bos - NT = tl.cdiv(T, BT) - boh = tl.load(chunk_offsets + i_n).to(tl.int32) - else: - bos, eos = i_n * T, i_n * T + T - NT = tl.cdiv(T, BT) - boh = i_n * NT - - # [BK, BV] - b_dh = tl.zeros([BK, BV], dtype=tl.float32) - # [BV] - b_dhb = tl.zeros([BV], dtype=tl.float32) - if USE_FINAL_STATE_GRADIENT: - p_dht = tl.make_block_ptr(dht + i_nh * K*V, (K, V), (V, 1), (i_k * BK, i_v * BV), (BK, BV), (1, 0)) - b_dh += tl.load(p_dht, boundary_check=(0, 1), padding_option="zero") - if USE_FINAL_STATE_GRADIENT_B: - p_dhbt = tl.make_block_ptr(dhbt + i_nh * V, (V,), (1,), (i_v * BV,), (BV,), (0,)) - b_dhb += tl.load(p_dhbt, boundary_check=(0,), padding_option="zero") - - # [BV] - offs_v = tl.arange(0, BV) - offs_t = tl.arange(0, BT) - b_w = tl.load(w + i_h * V + offs_v, mask=offs_v < V, other=0.) - b_b = tl.load(b + i_h * V + offs_v, mask=offs_v < V, other=0.) - b_dw = tl.zeros([BV], dtype=b_w.dtype) - b_db = tl.zeros([BV], dtype=b_b.dtype) - p_dw = tl.make_block_ptr(dw + i_nh * V, (V,), (1,), (i_v * BV,), (BV,), (0,)) - p_db = tl.make_block_ptr(db + i_nh * V, (V,), (1,), (i_v * BV,), (BV,), (0,)) - - for i_t in range(NT - 1, -1, -1): - p_h = tl.make_block_ptr(h + ((boh+i_t) * H + i_h) * K*V, (V, K), (1, V), (i_v * BV, i_k * BK), (BV, BK), (0, 1)) - p_dh = tl.make_block_ptr(dh + ((boh+i_t) * H + i_h) * K*V, (K, V), (V, 1), (i_k * BK, i_v * BV), (BK, BV), (1, 0)) - p_dhb = tl.make_block_ptr(dhb + ((boh+i_t) * H + i_h) * V, (V,), (1,), (i_v * BV,), (BV,), (0,)) - tl.store(p_dh, b_dh.to(p_dh.dtype.element_ty), boundary_check=(0, 1)) - tl.store(p_dhb, b_dhb.to(p_dhb.dtype.element_ty), boundary_check=(0,)) - p_q = tl.make_block_ptr(q+(bos*H+i_h)*K, (K, T), (1, H*K), (i_k * BK, i_t * BT), (BK, BT), (0, 1)) - p_k = tl.make_block_ptr(k+(bos*H+i_h)*K, (T, K), (H*K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) - p_v = tl.make_block_ptr(v+(bos*H+i_h)*V, (T, V), (H*V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) - p_v_new = tl.make_block_ptr(v_new+(bos*H+i_h)*V, (T, V), (H*V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) - p_x = tl.make_block_ptr(x+(bos*H+i_h)*V, (T, V), (H*V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) - p_y = tl.make_block_ptr(y+(bos*H+i_h)*V, (T, V), (H*V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) - p_dv_new = tl.make_block_ptr(dv_new+(bos*H+i_h)*V, (T, V), (H*V, 1), (i_t*BT, i_v * BV), (BT, BV), (1, 0)) - p_dv = tl.make_block_ptr(dv+(bos*H+i_h)*V, (T, V), (H*V, 1), (i_t*BT, i_v * BV), (BT, BV), (1, 0)) - p_dk = tl.make_block_ptr(dk+(bos*H+i_h)*K, (T, K), (H*K, 1), (i_t*BT, i_k * BK), (BT, BK), (1, 0)) - p_do = tl.make_block_ptr(do+(bos*H+i_h)*V, (T, V), (H*V, 1), (i_t*BT, i_v * BV), (BT, BV), (1, 0)) - p_r = tl.make_block_ptr(r+bos*H+i_h, (T, 1), (H, 1), (i_t*BT, 0), (BT, 1), (1, 0)) - p_eta_last = eta+bos*H+i_h + (T-1)*H if i_t == NT-1 else eta+bos*H+i_h + (i_t*BT+BT-1)*H - b_k = tl.load(p_k, boundary_check=(0, 1), padding_option="zero") - b_dv_new = tl.load(p_dv_new, boundary_check=(0, 1), padding_option="zero").to(b_k.dtype) - b_eta_last = tl.load(p_eta_last) - b_dv_new -= tl.dot(b_eta_last * b_k, b_dh.to(b_k.dtype)) - b_dv_new -= b_eta_last * b_dhb.to(b_k.dtype)[None, :] - - b_v_new = tl.load(p_v_new, boundary_check=(0, 1), padding_option="zero") - b_x = tl.load(p_x, boundary_check=(0, 1), padding_option="zero").to(b_k.dtype) - b_y = tl.load(p_y, boundary_check=(0, 1), padding_option="zero").to(b_k.dtype) - b_rstd = tl.load(p_r, boundary_check=(0, 1), padding_option="zero").to(tl.float32) - b_dy = b_rstd * (b_dv_new * V - tl.sum(b_dv_new, axis=1, keep_dims=True) - - b_x * tl.sum(b_dv_new * b_x, axis=1, keep_dims=True)) / V - b_dx = -b_rstd * (b_dv_new * tl.sum(b_x * b_y, axis=1, keep_dims=True) + - b_y * tl.sum(b_dv_new * b_x, axis=1, keep_dims=True)) / V - b_drstd = tl.sum(b_dv_new.to(b_rstd.dtype) * b_v_new.to(b_rstd.dtype) / b_rstd, axis=1, keep_dims=True) - - b_v = tl.load(p_v, boundary_check=(0, 1), padding_option="zero") - b_w = b_w.to(b_k.dtype) - b_b = b_b.to(b_k.dtype) - b_dv = -b_w * b_dy.to(b_k.dtype) - b_dk = b_w * b_dy.to(b_k.dtype) - b_dw += tl.sum(2 * b_w * b_x * b_dy.to(b_k.dtype) + - (b_b - b_v.to(b_k.dtype) + b_k) * b_dy.to(b_k.dtype), axis=0).to(b_dw.dtype) - b_db += tl.sum(b_w * b_dy.to(b_k.dtype), axis=0).to(b_db.dtype) - b_dx = b_dx.to(b_k.dtype) + b_w * b_w * b_dy.to(b_k.dtype) - - # d_rstd, dx --> dkh --> dk, dh - b_q = tl.load(p_q, boundary_check=(0, 1), padding_option="zero") - b_h = tl.load(p_h, boundary_check=(0, 1), padding_option="zero") - b_do = tl.load(p_do, boundary_check=(0, 1), padding_option="zero") - b_q = (b_q * scale).to(b_q.dtype) - b_dkh = b_rstd * (V * b_dx - tl.sum(b_dx, axis=1, keep_dims=True) - - b_x * tl.sum(b_x * b_dx, axis=1, keep_dims=True)) / V - b_dkh -= b_rstd * b_rstd * b_drstd * b_x / V - b_dkh = tl.where((offs_v < V)[None, :] * (offs_t < T-i_t*BT)[:, None], b_dkh, 0.) - b_dk += tl.dot(b_dkh, b_h.to(b_dkh.dtype)).to(b_k.dtype) - b_dh += tl.dot(b_q, b_do.to(b_q.dtype)) + tl.dot(tl.trans(b_k).to(b_dkh.dtype), b_dkh) - b_dhb += tl.sum(b_do + b_dkh, axis=0) - b_dh = tl.where((offs_v < V)[None, :], b_dh, 0.) - b_dhb = tl.where((offs_v < V), b_dhb, 0.) - - tl.store(p_dv, b_dv.to(p_dv.dtype.element_ty), boundary_check=(0, 1)) - tl.store(p_dk, b_dk.to(p_dk.dtype.element_ty), boundary_check=(0, 1)) - tl.store(p_dw, b_dw.to(p_dw.dtype.element_ty), boundary_check=(0,)) - tl.store(p_db, b_db.to(p_db.dtype.element_ty), boundary_check=(0,)) - - if USE_INITIAL_STATE: - p_dh0 = tl.make_block_ptr(dh0 + i_nh * K*V, (K, V), (V, 1), (i_k * BK, i_v * BV), (BK, BV), (1, 0)) - tl.store(p_dh0, b_dh.to(p_dh0.dtype.element_ty), boundary_check=(0, 1)) - if USE_INITIAL_STATE_B: - p_dhb0 = tl.make_block_ptr(dhb0+i_nh*V, (V,), (1,), (i_v * BV,), (BV,), (0,)) - tl.store(p_dhb0, b_dhb.to(p_dhb0.dtype.element_ty), boundary_check=(0,)) - - -@triton.heuristics({ - 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None, -}) -@triton.autotune( - configs=[ - triton.Config({}, num_warps=num_warps, num_stages=num_stages) - for num_warps in [2, 4, 8] - for num_stages in [2, 3] - ], - key=['BT', 'BK', 'BV'], - **autotune_cache_kwargs, -) -@triton.jit(do_not_specialize=['T']) -def chunk_bwd_kernel_dqke( - q, - k, - v, - e, - h, - do, - dh, - dhb, - dq, - dk, - de, - cu_seqlens, - chunk_indices, - scale, - T, - B: tl.constexpr, - H: tl.constexpr, - K: tl.constexpr, - V: tl.constexpr, - BT: tl.constexpr, - BK: tl.constexpr, - BV: tl.constexpr, - IS_VARLEN: tl.constexpr, -): - i_k, i_t, i_bh = tl.program_id(0), tl.program_id(1), tl.program_id(2) - i_b, i_h = i_bh // H, i_bh % H - 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.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) - T = eos - bos - NT = tl.cdiv(T, BT) - else: - NT = tl.cdiv(T, BT) - i_tg = i_b * NT + i_t - bos, eos = i_b * T, i_b * T + T - - # offset calculation - v += (bos * H + i_h) * V - do += (bos * H + i_h) * V - h += (i_tg * H + i_h) * K * V - dh += (i_tg * H + i_h) * K * V - dhb += (i_tg * H + i_h) * V - q += (bos * H + i_h) * K - k += (bos * H + i_h) * K - dq += (bos * H + i_h) * K - dk += (bos * H + i_h) * K - e += bos * H + i_h - de += bos * H + i_h - stride_qk = H*K - stride_vo = H*V - stride_e = H - - b_dq = tl.zeros([BT, BK], dtype=tl.float32) - b_dk = tl.zeros([BT, BK], dtype=tl.float32) - b_ds = tl.zeros([BT, BT], dtype=tl.float32) - b_de = tl.zeros([BT], dtype=tl.float32) - - p_k = tl.make_block_ptr(k, (T, K), (stride_qk, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) - b_k = tl.load(p_k, boundary_check=(0, 1)) - p_e_last = (e + (i_t*BT+BT-1)*stride_e) if (i_t*BT+BT) <= T else (e + (T-1)*stride_e) - i_last = (BT-1) if (i_t*BT+BT) <= T else (T % BT-1) - mask = (tl.arange(0, BT) == i_last) - b_e_last = tl.load(p_e_last) - - for i_v in range(tl.cdiv(V, BV)): - p_v = tl.make_block_ptr(v, (T, V), (stride_vo, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) - p_do = tl.make_block_ptr(do, (T, V), (stride_vo, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) - p_h = tl.make_block_ptr(h, (V, K), (1, V), (i_v * BV, i_k * BK), (BV, BK), (0, 1)) - p_dh = tl.make_block_ptr(dh, (V, K), (1, V), (i_v * BV, i_k * BK), (BV, BK), (0, 1)) - p_dhb = tl.make_block_ptr(dhb, (V,), (1,), (i_v * BV,), (BV,), (0,)) - # [BT, BV] - b_v = tl.load(p_v, boundary_check=(0, 1)) - b_do = tl.load(p_do, boundary_check=(0, 1)) - # [BV, BK] - b_h = tl.load(p_h, boundary_check=(0, 1)) - b_dh = tl.load(p_dh, boundary_check=(0, 1)) - # [BV] - b_dhb = tl.load(p_dhb, boundary_check=(0,)) - # [BT, BV] @ [BV, BT] -> [BT, BT] - b_ds += tl.dot(b_do, tl.trans(b_v)) - # [BT, BV] @ [BV, BK] -> [BT, BK] - b_dq += tl.dot(b_do, b_h.to(b_do.dtype)) - # [BT, BV] @ [BV, BK] -> [BT, BK] - b_dk -= b_e_last * tl.dot(b_v, b_dh.to(b_v.dtype)) - b_de -= mask * tl.sum(tl.trans(b_dh) * tl.dot(tl.trans(b_k), b_v.to(b_k.dtype))) - b_de -= mask * tl.sum(b_dhb * tl.sum(b_v, axis=0).to(b_k.dtype)) - - o_i = tl.arange(0, BT) - p_q = tl.make_block_ptr(q, (T, K), (stride_qk, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) - p_e = tl.make_block_ptr(e, (T,), (stride_e,), (i_t * BT,), (BT,), (0,)) - b_q = tl.load(p_q, boundary_check=(0, 1)) - b_e = tl.load(p_e, boundary_check=(0,)) - - p_dq = tl.make_block_ptr(dq, (T, K), (stride_qk, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) - p_dk = tl.make_block_ptr(dk, (T, K), (stride_qk, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)) - p_de = tl.make_block_ptr(de, (T,), (stride_e,), (i_t * BT,), (BT,), (0,)) - - b_ds = tl.where(o_i[:, None] >= o_i[None, :], b_ds, 0) - b_ds = b_ds.to(b_k.dtype) - b_dq -= tl.dot(b_ds, b_k) * b_e[:, None] - b_dk -= tl.dot(tl.trans(b_ds), b_q * b_e[:, None]) * scale - b_de -= tl.sum(scale * tl.dot(b_ds, b_k) * b_q, axis=1) - b_de -= tl.sum(b_ds, axis=1) - b_dq *= scale - tl.store(p_dq, b_dq.to(p_dq.dtype.element_ty), boundary_check=(0, 1)) - tl.store(p_dk, b_dk.to(p_dk.dtype.element_ty), boundary_check=(0, 1)) - tl.store(p_de, b_de.to(p_de.dtype.element_ty), boundary_check=(0,)) - - -def chunk_ttt_linear_fwd_h( - k: torch.Tensor, - v: torch.Tensor, - w: torch.Tensor, - b: torch.Tensor, - eta: torch.Tensor, - eps: float, - initial_state: torch.Tensor | None = None, - initial_state_bias: torch.Tensor | None = None, - output_final_state: bool = False, - cu_seqlens: torch.LongTensor | None = None, - chunk_size: int = 16, - chunk_indices: torch.LongTensor | None = None, -) -> tuple[torch.Tensor, torch.Tensor]: - B, T, H, K, V = *k.shape, v.shape[-1] - BT = chunk_size - - if chunk_indices is None and cu_seqlens is not None: - chunk_indices = prepare_chunk_indices(cu_seqlens, BT) - # N: the actual number of sequences in the batch with either equal or variable lengths - if cu_seqlens is None: - N, NT, chunk_offsets = B, triton.cdiv(T, BT), None - else: - N, NT, chunk_offsets = len(cu_seqlens) - 1, len(chunk_indices), prepare_chunk_offsets(cu_seqlens, BT) - BK = max(triton.next_power_of_2(K), 16) - BV = max(triton.next_power_of_2(V), 16) - assert max(BK, BV) <= 128, "current kernel does not support head dimension larger than 128." - NK = triton.cdiv(K, BK) - NV = triton.cdiv(V, BV) - assert NK == 1, 'NK > 1 is not supported because it involves time-consuming synchronization' - assert NV == 1, 'NV > 1 is not supported by TTT update rule.' - - h = k.new_empty(B, NT, H, K, V) - hb = k.new_empty(B, NT, H, 1, V) - final_state = k.new_empty(N, H, K, V, dtype=torch.float32) if output_final_state else None - final_state_bias = k.new_empty(N, H, 1, V, dtype=torch.float32) if output_final_state else None - - v_new = torch.empty_like(v) - grid = (NK, NV, N * H) - - chunk_ttt_linear_fwd_kernel_h[grid]( - k=k, - v=v, - v_new=v_new, - eta=eta, - w=w, - b=b, - eps=eps, - h=h, - hb=hb, - h0=initial_state, - hb0=initial_state_bias, - ht=final_state, - hbt=final_state_bias, - cu_seqlens=cu_seqlens, - chunk_offsets=chunk_offsets, - T=T, - H=H, - K=K, - V=V, - BT=BT, - BK=BK, - BV=BV, - ) - return h, hb, v_new, final_state, final_state_bias - - -def chunk_ttt_linear_fwd_o( - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - eta: torch.Tensor, - h: torch.Tensor, - hb: torch.Tensor, - scale: float | None = None, - cu_seqlens: torch.LongTensor | None = None, - chunk_size: int = 64, - chunk_indices: torch.LongTensor | None = None, -) -> torch.Tensor: - B, T, H, K, V = *q.shape, v.shape[-1] - if scale is None: - scale = k.shape[-1] ** -0.5 - BT = chunk_size - - if chunk_indices is None and cu_seqlens is not None: - chunk_indices = prepare_chunk_indices(cu_seqlens, BT) - NT = triton.cdiv(T, BT) if cu_seqlens is None else len(chunk_indices) - BK = max(triton.next_power_of_2(K), 16) - BV = max(triton.next_power_of_2(V), 16) - NK = triton.cdiv(K, BK) - NV = triton.cdiv(V, BV) - assert NK == 1, 'NK > 1 is not supported because it involves time-consuming synchronization' - assert NV == 1, 'NV > 1 is not supported by TTT update rule.' - - o = torch.empty_like(v) - - grid = (NV, NT, B * H) - chunk_ttt_linear_fwd_kernel_o[grid]( - q, - k, - v, - eta, - h, - hb, - o, - cu_seqlens, - chunk_indices, - scale, - T=T, - H=H, - K=K, - V=V, - BT=BT, - BK=BK, - BV=BV, - ) - return o - - -def chunk_ttt_linear_bwd_h( - k: torch.Tensor, - v: torch.Tensor, - w: torch.Tensor, - b: torch.Tensor, - eta: torch.Tensor, - eps: float, - initial_state: torch.Tensor | None = None, - initial_state_bias: torch.Tensor | None = None, - cu_seqlens: torch.LongTensor | None = None, - chunk_size: int = 16, - chunk_indices: torch.LongTensor | None = None, -) -> tuple[torch.Tensor, torch.Tensor]: - B, T, H, K, V = *k.shape, v.shape[-1] - BT = chunk_size - - if chunk_indices is None and cu_seqlens is not None: - chunk_indices = prepare_chunk_indices(cu_seqlens, BT) - # N: the actual number of sequences in the batch with either equal or variable lengths - if cu_seqlens is None: - N, NT, chunk_offsets = B, triton.cdiv(T, BT), None - else: - N, NT, chunk_offsets = len(cu_seqlens) - 1, len(chunk_indices), prepare_chunk_offsets(cu_seqlens, BT) - BK = max(triton.next_power_of_2(K), 16) - BV = max(triton.next_power_of_2(V), 16) - assert max(BK, BV) <= 128, "current kernel does not support head dimension larger than 128." - NK = triton.cdiv(K, BK) - NV = triton.cdiv(V, BV) - assert NK == 1, 'NK > 1 is not supported because it involves time-consuming synchronization' - assert NV == 1, 'NV > 1 is not supported by TTT update rule.' - - h = k.new_empty(B, NT, H, K, V) - rstd = v.new_empty(B, T, H, 1, dtype=torch.float32) - x = torch.empty_like(v) - y = torch.empty_like(v) - - v_new = torch.empty_like(v) - grid = (NK, NV, N * H) - - chunk_ttt_linear_bwd_kernel_h[grid]( - k=k, - v=v, - v_new=v_new, - eta=eta, - w=w, - b=b, - eps=eps, - h=h, - h0=initial_state, - hb0=initial_state_bias, - x=x, - y=y, - r=rstd, - cu_seqlens=cu_seqlens, - chunk_offsets=chunk_offsets, - T=T, - H=H, - K=K, - V=V, - BT=BT, - BK=BK, - BV=BV, - NT=NT, - ) - return h, v_new, x, y, rstd - - -def chunk_ttt_linear_bwd_dv_local( - q: torch.Tensor, - k: torch.Tensor, - eta: torch.Tensor, - do: torch.Tensor, - scale: float, - cu_seqlens: torch.LongTensor | None = None, - chunk_size: int = 16, - chunk_indices: torch.LongTensor | None = None, -) -> torch.Tensor: - B, T, H, K, V = *k.shape, do.shape[-1] - BT = chunk_size - - if chunk_indices is None and cu_seqlens is not None: - chunk_indices = prepare_chunk_indices(cu_seqlens, BT) - NT = triton.cdiv(T, BT) if cu_seqlens is None else len(chunk_indices) - BK = min(max(triton.next_power_of_2(K), 16), 128) - BV = min(max(triton.next_power_of_2(V), 16), 128) - - dv = torch.empty_like(do) - grid = (NT, B * H) - chunk_ttt_linear_bwd_kernel_dv_local[grid]( - q, - k, - eta, - do, - dv, - cu_seqlens, - chunk_indices, - scale, - T=T, - H=H, - K=K, - V=V, - BT=BT, - BK=BK, - BV=BV, - ) - return dv - - -def chunk_ttt_linear_bwd_norm( - q: torch.Tensor, # [B, H, L, D] - k: torch.Tensor, # [B, H, L, D] - v: torch.Tensor, # [B, H, L, D] - v_new: torch.Tensor, # [B, H, L, D] - x: torch.Tensor, # [B, H, L, D] - y: torch.Tensor, # [B, H, L, D] - rstd: torch.Tensor, # [B, H, L, 1] - w: torch.Tensor, # [H, D] - b: torch.Tensor, # [H, D] - eta: torch.Tensor, # [B, H, L, 1] - h0: torch.Tensor, # [B, H, D, D] - hb0: torch.Tensor, # [B, H, 1, D] - h: torch.Tensor, # [B, H, NT, D, D] - dht: torch.Tensor | None, # [B, H, D, D] - dhbt: torch.Tensor | None, # [B, H, 1, D] - dv_new: torch.Tensor | None, # [B, H, L, D] - do: torch.Tensor, # [B, H, L, D] - scale: float, - cu_seqlens: torch.LongTensor | None = None, - chunk_size: int = 16, - chunk_indices: torch.LongTensor | None = None, -) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: - # torch implementation of `dkh, dw, db, dk, dv` for LN^2 - assert cu_seqlens is None, "bwd of varlen is not implemented yet." - B, T, H, K, V = *q.shape, do.shape[-1] - BT = chunk_size - - if chunk_indices is None and cu_seqlens is not None: - chunk_indices = prepare_chunk_indices(cu_seqlens, BT) - if cu_seqlens is None: - N, NT, chunk_offsets = B, triton.cdiv(T, BT), None - else: - N, NT, chunk_offsets = len(cu_seqlens) - 1, len(chunk_indices), prepare_chunk_offsets(cu_seqlens, BT) - - BK = max(triton.next_power_of_2(K), 16) - BV = max(triton.next_power_of_2(V), 16) - NK = triton.cdiv(K, BK) - NV = triton.cdiv(V, BV) - assert NK == 1, 'NK > 1 is not supported by TTT.' - assert NV == 1, 'NV > 1 is not supported by TTT.' - - dh = q.new_empty(B, NT, H, K, V) - dhb = q.new_empty(B, NT, H, 1, V) - dh0 = torch.empty_like(h0, dtype=torch.float32) if h0 is not None else None - dhb0 = torch.empty_like(hb0, dtype=torch.float32) if hb0 is not None else None - dv = torch.empty_like(v) - dk = torch.empty_like(k) - dw = w.new_empty(B, H, V) - db = b.new_empty(B, H, V) - - grid = (NK, NV, N * H) - chunk_ttt_linear_bwd_kernel_norm[grid]( - q=q, - k=k, - v=v, - v_new=v_new, - x=x, - y=y, - r=rstd, - w=w, - b=b, - eta=eta, - h=h, - dht=dht, - dhbt=dhbt, - dh0=dh0, - dhb0=dhb0, - do=do, - dh=dh, - dhb=dhb, - dv=dv, - dv_new=dv_new, - dk=dk, - dw=dw, - db=db, - cu_seqlens=cu_seqlens, - chunk_offsets=chunk_offsets, - scale=scale, - T=T, - H=H, - K=K, - V=V, - BT=BT, - BK=BK, - BV=BV, - ) - dw = dw.sum(dim=0) - db = db.sum(dim=0) - return dh, dhb, dh0, dhb0, dv, dk, dw, db - - -def chunk_ttt_linear_bwd_norm_ref( - q: torch.Tensor, # [B, H, L, D] - k: torch.Tensor, # [B, H, L, D] - v: torch.Tensor, # [B, H, L, D] - v_new: torch.Tensor, # [B, H, L, D] - kh: torch.Tensor, # [B, H, L, D] - y: torch.Tensor, # [B, H, L, D] - w: torch.Tensor, # [H, D] - b: torch.Tensor, # [H, D] - eta: torch.Tensor, # [B, H, L, 1] - h0: torch.Tensor, # [B, H, D, D] - h: torch.Tensor, # [B, H, NT, D, D] - dht: torch.Tensor | None, # [B, H, D, D] - dv_new: torch.Tensor | None, # [B, H, L, D] - do: torch.Tensor, # [B, H, L, D] - scale: float, - eps: float, - cu_seqlens: torch.LongTensor | None = None, - chunk_size: int = 16, - chunk_indices: torch.LongTensor | None = None, -) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: - # torch implementation of `dkh, dw, db, dk, dv` for LN^2 - assert cu_seqlens is None, "bwd of varlen is not implemented yet." - B, T, H, K, V = *q.shape, do.shape[-1] - # [B, L, H, D] -> [B, H, L, D] - q, k, v, v_new, kh, y, h, eta, dv_new, do = [ - x.transpose(1, 2) for x in - [q, k, v, v_new, kh, y, h, eta, dv_new, do] - ] - BT = chunk_size - - if chunk_indices is None and cu_seqlens is not None: - chunk_indices = prepare_chunk_indices(cu_seqlens, BT) - NT = triton.cdiv(T, BT) if cu_seqlens is None else len(chunk_indices) - pad_len = (BT - (T % BT)) % BT - if pad_len > 0: - q, k, v, v_new, kh, y, eta, dv_new, do = [ - F.pad(x, (0, 0, 0, pad_len)) for x in - [q, k, v, v_new, kh, y, eta, dv_new, do] - ] - eta[:, :, -1, :] = eta[:, :, -(pad_len+1), :] - # [NT, B, H, BT, D] - q, k, v, v_new, kh, y, eta, dv_new, do = [ - x.reshape(B, H, NT, BT, -1).permute(2, 0, 1, 3, 4) for x in - [q, k, v, v_new, kh, y, eta, dv_new, do] - ] - h = h.permute(2, 0, 1, 3, 4) - - # allocate - dh = q.new_zeros(NT, B, H, K, V) - dv = torch.zeros_like(v) - dk = torch.zeros_like(k) - dw = torch.zeros_like(w) - db = torch.zeros_like(b) - # recurrent state - b_dh = dht if dht is not None else torch.zeros_like(dh[0]) - b_dh = b_dh.to(torch.float32) - - # [H, 1, D] - _w = w.reshape(H, 1, V).to(torch.float32) - _b = b.reshape(H, 1, V).to(torch.float32) - - # d_state passing - for i_t in range(NT - 1, -1, -1): - dh[i_t] = b_dh.to(dh.dtype) - # [B, H, BT, D] - _q, _k, _v, _v_new, _kh, _y, _h, _eta, _dv_new, _do = [ - x[i_t].to(torch.float32) for x in - (q, k, v, v_new, kh, y, h, eta, dv_new, do) - ] - _dv_new -= (_eta[:, :, -1, :, None] * _k) @ b_dh - - mean = _kh.mean(dim=-1, keepdim=True) - var = _kh.var(dim=-1, unbiased=False, keepdim=True).to(torch.float32) - rstd = 1 / torch.sqrt(var + eps).to(torch.float32) - x = (_kh - mean) * rstd - # [B, H, BT, D] - dy = rstd * (_dv_new*V - _dv_new.sum(dim=-1, keepdim=True) - x*(x*_dv_new).sum(dim=-1, keepdim=True)) / V - dx = -rstd * (_dv_new*(x*_y).sum(dim=-1, keepdim=True) + _y*(x*_dv_new).sum(dim=-1, keepdim=True)) / V - d_rstd = (_dv_new * _v_new / rstd).sum(dim=-1, keepdim=True) - - dv[i_t] = (-_w*dy).to(dv.dtype) - dk[i_t] += (_w*dy).to(dk.dtype) - dw += (2*_w*x*dy+(_b-_v+_k)*dy).sum(dim=(0, 2)).to(dw.dtype) - db += (_w*dy).sum(dim=(0, 2)).to(db.dtype) - dx += _w*_w*dy - - # d_rstd, dx --> dkh --> dk, dh - dkh = rstd * (V * dx - dx.sum(dim=-1, keepdim=True) - x * (x * dx).sum(dim=-1, keepdim=True)) / V - dkh -= rstd**2 * d_rstd * x / V - dk[i_t] += (dkh @ _h.transpose(-2, -1)).to(dk.dtype) - b_dh += (_q.transpose(-2, -1) * scale) @ _do + _k.transpose(-2, -1) @ dkh - dh0 = b_dh.to(torch.float32) if h0 is not None else None - - # [NT, B, H, BT, D] -> [B, H, T, D] - dv = dv.permute(1, 2, 0, 3, 4).reshape(B, H, -1, V)[:, :, :T, :] - dk = dk.permute(1, 2, 0, 3, 4).reshape(B, H, -1, K)[:, :, :T, :] - # [B, H, NT, D, D] - dh = dh.permute(1, 2, 0, 3, 4) - dv, dk, dh = [x.transpose(1, 2) for x in (dv, dk, dh)] - dh, dv, dk, dw, db = [x.contiguous() for x in (dh, dv, dk, dw, db)] - dh0 = dh0.contiguous() if h0 is not None else None - return dh, dh0, dv, dk, dw, db - - -def chunk_ttt_linear_bwd_dqke( - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - eta: torch.Tensor, - h: torch.Tensor, - do: torch.Tensor, - dh: torch.Tensor, - dhb: torch.Tensor, - scale: float, - cu_seqlens: torch.LongTensor | None = None, - chunk_size: int = 16, - chunk_indices: torch.LongTensor | None = None, -) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: - B, T, H, K, V = *k.shape, v.shape[-1] - BT = chunk_size - - if chunk_indices is None and cu_seqlens is not None: - chunk_indices = prepare_chunk_indices(cu_seqlens, BT) - NT = triton.cdiv(T, BT) if cu_seqlens is None else len(chunk_indices) - - BK = max(triton.next_power_of_2(K), 16) - BV = min(max(triton.next_power_of_2(V), 16), 64) - NK = triton.cdiv(K, BK) - assert NK == 1, "NK > 1 is not supported." - - dq = torch.empty_like(q) - dk = torch.empty_like(k) - de = torch.empty_like(eta) - grid = (NK, NT, B * H) - - chunk_bwd_kernel_dqke[grid]( - q=q, - k=k, - v=v, - e=eta, - h=h, - do=do, - dh=dh, - dhb=dhb, - dq=dq, - dk=dk, - de=de, - cu_seqlens=cu_seqlens, - chunk_indices=chunk_indices, - scale=scale, - B=B, - T=T, - H=H, - K=K, - V=V, - BT=BT, - BK=BK, - BV=BV, - ) - return dq, dk, de - - -def chunk_ttt_linear_fwd( - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - w: torch.Tensor, - b: torch.Tensor, - eta: torch.Tensor, - scale: float, - eps: float, - initial_state: torch.Tensor, - initial_state_bias: torch.Tensor, - output_final_state: bool, - cu_seqlens: torch.LongTensor | None = None, - chunk_size: int = 16, - chunk_indices: torch.LongTensor | None = None, -): - BT = chunk_size - h, hb, v_new, final_state, final_state_bias = chunk_ttt_linear_fwd_h( - k=k, - v=v, - w=w, - b=b, - eta=eta, - eps=eps, - initial_state=initial_state, - initial_state_bias=initial_state_bias, - output_final_state=output_final_state, - cu_seqlens=cu_seqlens, - chunk_size=BT, - chunk_indices=chunk_indices, - ) - o = chunk_ttt_linear_fwd_o( - q=q, - k=k, - v=v_new, - eta=eta, - h=h, - hb=hb, - scale=scale, - cu_seqlens=cu_seqlens, - chunk_size=BT, - chunk_indices=chunk_indices, - ) - return o, final_state, final_state_bias - - -def chunk_ttt_linear_bwd( - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - w: torch.Tensor, - b: torch.Tensor, - eta: torch.Tensor, - scale: float, - eps: float, - do: torch.Tensor, - dht: torch.Tensor, - dhbt: torch.Tensor, - chunk_size: int = 16, - initial_state: torch.Tensor = None, - initial_state_bias: torch.Tensor = None, - cu_seqlens: torch.LongTensor | None = None, - chunk_indices: torch.LongTensor | None = None, -): - BT = chunk_size - h, v_new, x, y, rstd = chunk_ttt_linear_bwd_h( - k=k, - v=v, - w=w, - b=b, - eta=eta, - eps=eps, - initial_state=initial_state, - initial_state_bias=initial_state_bias, - cu_seqlens=cu_seqlens, - chunk_size=BT, - chunk_indices=chunk_indices, - ) - dv_new = chunk_ttt_linear_bwd_dv_local( - q=q, - k=k, - eta=eta, - do=do, - scale=scale, - cu_seqlens=cu_seqlens, - chunk_size=BT, - chunk_indices=chunk_indices, - ) - dh, dhb, dh0, dhb0, dv, dk, dw, db = chunk_ttt_linear_bwd_norm( - q=q, - k=k, - v=v, - v_new=v_new, - x=x, - y=y, - rstd=rstd, - w=w, - b=b, - eta=eta, - h0=initial_state, - hb0=initial_state_bias, - h=h, - dht=dht, - dhbt=dhbt, - dv_new=dv_new, - do=do, - scale=scale, - cu_seqlens=cu_seqlens, - chunk_size=BT, - chunk_indices=chunk_indices, - ) - dq, dk2, de = chunk_ttt_linear_bwd_dqke( - q=q, - k=k, - v=v_new, - eta=eta, - h=h, - do=do, - dh=dh, - dhb=dhb, - scale=scale, - cu_seqlens=cu_seqlens, - chunk_size=BT, - chunk_indices=chunk_indices, - ) - dk.add_(dk2) - return dq, dk, dv, de, dw, db, dh0, dhb0 - - -class ChunkTTTLinearFunction(torch.autograd.Function): - - @staticmethod - @input_guard - @autocast_custom_fwd - def forward( - ctx, - q, - k, - v, - w, - b, - chunk_size, - eta, - scale, - eps, - initial_state, - initial_state_bias, - output_final_state, - cu_seqlens, - cu_seqlens_cpu, - ): - chunk_indices = prepare_chunk_indices( - cu_seqlens, chunk_size, cu_seqlens_cpu=cu_seqlens_cpu) if cu_seqlens is not None else None - o, final_state, final_state_bias = chunk_ttt_linear_fwd( - q=q, - k=k, - v=v, - w=w, - b=b, - eta=eta, - scale=scale, - eps=eps, - chunk_size=chunk_size, - initial_state=initial_state, - initial_state_bias=initial_state_bias, - output_final_state=output_final_state, - cu_seqlens=cu_seqlens, - chunk_indices=chunk_indices, - ) - ctx.save_for_backward(q, k, v, eta, w, b, initial_state, initial_state_bias, chunk_indices) - ctx.chunk_size = chunk_size - ctx.scale = scale - ctx.eps = eps - ctx.cu_seqlens = cu_seqlens - return o.to(q.dtype), final_state, final_state_bias - - @staticmethod - @input_guard - @autocast_custom_bwd - def backward(ctx, do, dht, dhbt): - q, k, v, eta, w, b, initial_state, initial_state_bias, chunk_indices = ctx.saved_tensors - dq, dk, dv, de, dw, db, dh0, dhb0 = chunk_ttt_linear_bwd( - q=q, - k=k, - v=v, - w=w, - b=b, - eta=eta, - scale=ctx.scale, - eps=ctx.eps, - do=do, - dht=dht, - dhbt=dhbt, - chunk_size=ctx.chunk_size, - initial_state=initial_state, - initial_state_bias=initial_state_bias, - cu_seqlens=ctx.cu_seqlens, - chunk_indices=chunk_indices, - ) - return ( - dq.to(q), dk.to(k), dv.to(v), dw.to(w), db.to(b), - None, de.to(eta), None, None, dh0, dhb0, None, None, None, - ) - - -def norm_residual(x, weight, bias, eps): - # GroupNorm and Residual - B, T, H, D = x.shape - x += group_norm( - x.reshape(B, T, -1).clone(), - weight=weight.reshape(-1).clone(), - bias=bias.reshape(-1).clone(), - eps=eps, - num_groups=H, - ).reshape(x.shape) - return x - - -def chunk_ttt_linear( - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - w: torch.Tensor, - b: torch.Tensor, - eta: torch.Tensor, - scale: float | None = None, - eps: float = 1e-6, - chunk_size: int = 16, - initial_state: torch.Tensor | None = None, - initial_state_bias: torch.Tensor | None = None, - output_final_state: bool = False, - cu_seqlens: torch.LongTensor | None = None, - cu_seqlens_cpu: torch.LongTensor | None = None, - **kwargs, -): - r""" - Args: - q (torch.Tensor): - queries of shape `[B, T, H, K]`. - k (torch.Tensor): - keys of shape `[B, T, H, K]`. - v (torch.Tensor): - values of shape `[B, T, H, V]`. - w (torch.Tensor): - layer norm weight of shape `[H, V]`. - b (torch.Tensor): - layer norm bias of shape `[H, V]`. - eta (torch.Tensor): - Learning rate for hidden state, of shape `[B, T, H, 1]`. - scale (Optional[float]): - Scale factor for the attention scores. - If not provided, it will default to `1 / sqrt(K)`. Default: `None`. - chunk_size (int): - Chunk size. Default: `16`. - initial_state (Optional[torch.Tensor]): - Initial state of shape `[N, H, K, V]`. Default: `None`. - initial_state_bias (Optional[torch.Tensor]): - Initial state bias of shape `[N, H, 1, V]`. Default: `None`. - output_final_state (Optional[bool]): - Whether to output the final state of shape `[N, H, K, V]`. Default: `False`. - cu_seqlens (torch.LongTensor): - Cumulative sequence lengths of shape `[N+1]` used for variable-length training, - consistent with the FlashAttention API. - cu_seqlens_cpu (torch.LongTensor): - CPU copy of `cu_seqlens` to avoid unnecessary device synchronization. Default: `None`. - - Returns: - o (torch.Tensor): - Outputs of shape `[B, T, H, V]`. - final_state (torch.Tensor): - Final state of shape `[N, H, K, V]` if `output_final_state=True` else `None`. - final_state_bias (torch.Tensor): - Final state bias of shape `[N, H, 1, V]` if `output_final_state=True` else `None`. - """ - assert q.dtype == k.dtype == v.dtype - assert k.shape[-1] == v.shape[-1], "DK must equal to DV." - if isinstance(eta, float): - eta = torch.full_like(q[:, :, :, :1], eta) - if 'head_first' in kwargs: - raise DeprecationWarning( - "head_first has been removed. Inputs must be in `[B, T, H, ...]` format.", - ) - if cu_seqlens is not None: - if q.shape[0] != 1: - raise ValueError( - f"The batch size is expected to be 1 rather than {q.shape[0]} when using `cu_seqlens`. " - f"Please flatten variable-length inputs before processing.", - ) - if initial_state is not None and initial_state.shape[0] != len(cu_seqlens) - 1: - raise ValueError( - f"The number of initial states is expected to be equal to the number of input sequences, " - f"i.e., {len(cu_seqlens) - 1} rather than {initial_state.shape[0]}.", - ) - if scale is None: - scale = k.shape[-1] ** -0.5 - else: - assert scale > 0, "Scale must be positive." - o, final_state, final_state_bias = ChunkTTTLinearFunction.apply( - q, - k, - v, - w, - b, - chunk_size, - eta, - scale, - eps, - initial_state, - initial_state_bias, - output_final_state, - cu_seqlens, - cu_seqlens_cpu, - ) - o = norm_residual(o, w, b, eps) - return o, final_state, final_state_bias diff --git a/build/torch-xpu/ops/ttt/fused_chunk.py b/build/torch-xpu/ops/ttt/fused_chunk.py deleted file mode 100644 index 200a76f4b0ec8668f116799e791cd346ed5f2519..0000000000000000000000000000000000000000 --- a/build/torch-xpu/ops/ttt/fused_chunk.py +++ /dev/null @@ -1,824 +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 -import triton.language as tl - -from ...modules.layernorm import group_norm -from ...utils import IS_NVIDIA_HOPPER, autocast_custom_bwd, autocast_custom_fwd, autotune_cache_kwargs, input_guard - -NUM_WARPS = [1, 2] if IS_NVIDIA_HOPPER else [1, 2, 4, 8] - - -@triton.heuristics({ - 'USE_INITIAL_STATE': lambda args: args['h0'] is not None, - 'USE_INITIAL_STATE_B': lambda args: args['hb0'] is not None, - 'STORE_FINAL_STATE': lambda args: args['ht'] is not None, - 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None, -}) -@triton.autotune( - configs=[ - triton.Config({}, num_warps=1), - triton.Config({}, num_warps=2), - triton.Config({}, num_warps=4), - ], - key=['BT', 'BK', 'BV'], - **autotune_cache_kwargs, -) -@triton.jit(do_not_specialize=['T']) -def fused_chunk_ttt_linear_fwd_kernel( - q, - k, - v, - eta, - w, - b, - o, - scale, - eps, - h0, - hb0, - ht, - hbt, - cu_seqlens, - T, - H: tl.constexpr, - K: tl.constexpr, - V: tl.constexpr, - BT: tl.constexpr, - BK: tl.constexpr, - BV: tl.constexpr, - USE_INITIAL_STATE: tl.constexpr, - USE_INITIAL_STATE_B: tl.constexpr, - STORE_FINAL_STATE: tl.constexpr, - IS_VARLEN: tl.constexpr, -): - i_nh = tl.program_id(0) - i_n, i_h = i_nh // H, i_nh % H - if IS_VARLEN: - bos, eos = tl.load(cu_seqlens + i_n).to(tl.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) - T = eos - bos - NT = tl.cdiv(T, BT) - else: - bos, eos = i_n * T, i_n * T + T - NT = tl.cdiv(T, BT) - - o_i = tl.arange(0, BT) - v_i = tl.arange(0, BV) - m_A = o_i[:, None] >= o_i[None, :] - b_w = tl.load(w + i_h * V + v_i, mask=v_i < V, other=0.) - b_b = tl.load(b + i_h * V + v_i, mask=v_i < V, other=0.) - - # [BK, BV] - b_h = tl.zeros([BK, BV], dtype=tl.float32) - # [BV] - b_hb = tl.zeros([BV], dtype=tl.float32) - if USE_INITIAL_STATE: - p_h0 = tl.make_block_ptr(h0 + i_nh * K * V, (K, V), (V, 1), (0, 0), (BK, BV), (1, 0)) - b_h = tl.load(p_h0, boundary_check=(0, 1), padding_option="zero").to(tl.float32) - if USE_INITIAL_STATE_B: - p_hb0 = tl.make_block_ptr(hb0 + i_nh * V, (V,), (1,), (0,), (BV,), (0,)) - b_hb = tl.load(p_hb0, boundary_check=(0,), padding_option="zero").to(tl.float32) - - for i_t in range(NT): - p_q = tl.make_block_ptr(q+(bos*H+i_h)*K, (T, K), (H*K, 1), (i_t*BT, 0), (BT, BK), (1, 0)) - p_k = tl.make_block_ptr(k+(bos*H+i_h)*K, (K, T), (1, H*K), (0, i_t*BT), (BK, BT), (0, 1)) - p_v = tl.make_block_ptr(v+(bos*H+i_h)*V, (T, V), (H*V, 1), (i_t*BT, 0), (BT, BV), (1, 0)) - p_o = tl.make_block_ptr(o+(bos*H+i_h)*V, (T, V), (H*V, 1), (i_t*BT, 0), (BT, BV), (1, 0)) - p_e = tl.make_block_ptr(eta+(bos*H+i_h), (T,), (H,), (i_t*BT,), (BT,), (0,)) - p_e_last = eta+bos*H+i_h + (T-1)*H if i_t == NT-1 else eta+bos*H+i_h + (i_t*BT+BT-1)*H - # [BK, BT] - b_k = tl.load(p_k, boundary_check=(0, 1), padding_option="zero") - # [BT, BV] - b_v = tl.load(p_v, boundary_check=(0, 1), padding_option="zero") - - # [BT, BV] - b_kh = tl.dot(tl.trans(b_k), b_h.to(b_k.dtype), allow_tf32=False).to(tl.float32) + b_hb[None, :] - b_kh = tl.where((v_i < V)[None, :], b_kh, 0.) - mean = tl.sum(b_kh, axis=1, keep_dims=True) / V - xbar = tl.where((v_i < V)[None, :], b_kh - mean, 0.) - var = tl.sum(xbar * xbar, axis=1, keep_dims=True) / V - rstd = 1 / tl.sqrt(var.to(tl.float32) + eps) - b_kh_hat = (b_kh - mean) * rstd - - b_v = b_kh_hat.to(b_k.dtype) * b_w[None, :].to(b_k.dtype) + \ - b_b[None, :].to(b_k.dtype) - b_v.to(b_k.dtype) + tl.trans(b_k) - b_v = tl.where((v_i < V)[None, :], b_v * b_w[None, :].to(b_k.dtype), 0.) - b_v2 = rstd * (V * b_v - tl.sum(b_v, axis=1, keep_dims=True) - b_kh_hat.to(b_k.dtype) - * tl.sum(b_v * b_kh_hat.to(b_k.dtype), axis=1, keep_dims=True)) / V - - # [BT, BK] - b_q = tl.load(p_q, boundary_check=(0, 1), padding_option="zero") - # [BT] - b_e = tl.load(p_e, boundary_check=(0,), padding_option="zero") - b_q = (b_q * scale).to(b_k.dtype) - - # [BT, BT] - b_A = tl.dot(b_q, b_k, allow_tf32=False) - b_A = tl.where(m_A, b_A, 0) - b_Ae = tl.where(m_A, b_e[:, None], 0.0) - - b_o = - tl.dot(b_e[:, None] * b_A.to(b_v2.dtype), b_v2, allow_tf32=False) - b_o += b_hb[None, :] - tl.dot(b_Ae.to(b_v2.dtype), b_v2, allow_tf32=False) - b_o += tl.dot(b_q, b_h.to(b_q.dtype), allow_tf32=False) - b_e_last = tl.load(p_e_last) - b_h = b_h - tl.dot(b_e_last * b_k, b_v2.to(b_k.dtype), allow_tf32=False) - b_hb = b_hb - tl.sum(b_e_last * b_v2.to(b_k.dtype), axis=0) - b_h = tl.where((v_i < V)[None, :], b_h, 0.) - b_hb = tl.where((v_i < V), b_hb, 0.) - tl.store(p_o, b_o.to(p_o.dtype.element_ty), boundary_check=(0, 1)) - - if STORE_FINAL_STATE: - p_ht = tl.make_block_ptr(ht + i_nh * K*V, (K, V), (V, 1), (0, 0), (BK, BV), (1, 0)) - p_hbt = tl.make_block_ptr(hbt + i_nh * V, (V,), (1,), (0,), (BV,), (0,)) - tl.store(p_ht, b_h.to(p_ht.dtype.element_ty), boundary_check=(0, 1)) - tl.store(p_hbt, b_hb.to(p_hbt.dtype.element_ty), boundary_check=(0,)) - - -@triton.heuristics({ - 'USE_INITIAL_STATE': lambda args: args['h0'] is not None, - 'USE_INITIAL_STATE_B': lambda args: args['hb0'] is not None, -}) -@triton.autotune( - configs=[ - triton.Config({}, num_warps=1), - triton.Config({}, num_warps=2), - triton.Config({}, num_warps=4), - ], - key=['BT', 'BK', 'BV'], - **autotune_cache_kwargs, -) -@triton.jit(do_not_specialize=['T']) -def fused_chunk_ttt_linear_bwd_kernel_h( - k, - v, - v2, - x, - y, - r, - w, - b, - eta, - h0, - hb0, - h, - do, - dq, - scale, - eps, - T, - H: tl.constexpr, - K: tl.constexpr, - V: tl.constexpr, - BT: tl.constexpr, - BK: tl.constexpr, - BV: tl.constexpr, - USE_INITIAL_STATE: tl.constexpr, - USE_INITIAL_STATE_B: tl.constexpr, -): - i_nh = tl.program_id(0) - i_n, i_h = i_nh // H, i_nh % H - bos, _ = i_n * T, i_n * T + T - NT = tl.cdiv(T, BT) - boh = i_n * NT - - o_i = tl.arange(0, BT) - v_i = tl.arange(0, BV) - m_A = o_i[:, None] >= o_i[None, :] - b_w = tl.load(w + i_h * V + v_i, mask=v_i < V, other=0.) - b_b = tl.load(b + i_h * V + v_i, mask=v_i < V, other=0.) - - # [BK, BV] - b_h = tl.zeros([BK, BV], dtype=tl.float32) - # [BV] - b_hb = tl.zeros([BV], dtype=tl.float32) - if USE_INITIAL_STATE: - p_h0 = tl.make_block_ptr(h0 + i_nh * K * V, (K, V), (V, 1), (0, 0), (BK, BV), (1, 0)) - b_h = tl.load(p_h0, boundary_check=(0, 1), padding_option="zero").to(tl.float32) - if USE_INITIAL_STATE_B: - p_hb0 = tl.make_block_ptr(hb0 + i_nh * V, (V,), (1,), (0,), (BV,), (0,)) - b_hb = tl.load(p_hb0, boundary_check=(0,), padding_option="zero").to(tl.float32) - - for i_t in range(NT): - p_h = tl.make_block_ptr(h+((boh+i_t)*H+i_h)*K*V, (K, V), (V, 1), (0, 0), (BK, BV), (1, 0)) - p_k = tl.make_block_ptr(k+(bos*H+i_h)*K, (K, T), (1, H*K), (0, i_t*BT), (BK, BT), (0, 1)) - p_v = tl.make_block_ptr(v+(bos*H+i_h)*V, (T, V), (H*V, 1), (i_t*BT, 0), (BT, BV), (1, 0)) - p_v2 = tl.make_block_ptr(v2+(bos*H+i_h)*V, (T, V), (H*V, 1), (i_t*BT, 0), (BT, BV), (1, 0)) - p_x = tl.make_block_ptr(x+(bos*H+i_h)*V, (T, V), (H*V, 1), (i_t*BT, 0), (BT, BV), (1, 0)) - p_y = tl.make_block_ptr(y+(bos*H+i_h)*V, (T, V), (H*V, 1), (i_t*BT, 0), (BT, BV), (1, 0)) - p_r = tl.make_block_ptr(r+bos*H+i_h, (T, 1), (H, 1), (i_t*BT, 0), (BT, 1), (1, 0)) - p_e = tl.make_block_ptr(eta+(bos*H+i_h), (T,), (H,), (i_t*BT,), (BT,), (0,)) - p_dq = tl.make_block_ptr(dq+(bos*H+i_h)*K, (T, K), (H*K, 1), (i_t*BT, 0), (BT, BK), (1, 0)) - p_do = tl.make_block_ptr(do+(bos*H+i_h)*V, (T, V), (H*V, 1), (i_t*BT, 0), (BT, BV), (1, 0)) - p_e_last = eta+bos*H+i_h + (T-1)*H if i_t == NT-1 else eta+bos*H+i_h + (i_t*BT+BT-1)*H - tl.store(p_h, b_h.to(p_h.dtype.element_ty), boundary_check=(0, 1)) - # [BK, BT] - b_k = tl.load(p_k, boundary_check=(0, 1), padding_option="zero") - # [BT, BV] - b_v = tl.load(p_v, boundary_check=(0, 1), padding_option="zero") - - b_kh = tl.dot(tl.trans(b_k), b_h.to(b_k.dtype), allow_tf32=False).to(tl.float32) + b_hb[None, :] - b_kh = tl.where((v_i < V)[None, :], b_kh, 0.) - mean = tl.sum(b_kh, axis=1, keep_dims=True) / V - xbar = tl.where((v_i < V)[None, :], b_kh - mean, 0.) - var = tl.sum(xbar * xbar, axis=1, keep_dims=True) / V - rstd = 1 / tl.sqrt(var.to(tl.float32) + eps) - b_kh_hat = (b_kh - mean) * rstd - - b_v = b_kh_hat.to(b_k.dtype) * b_w[None, :].to(b_k.dtype) + \ - b_b[None, :].to(b_k.dtype) - b_v.to(b_k.dtype) + tl.trans(b_k) - b_v = tl.where((v_i < V)[None, :], b_v * b_w[None, :].to(b_k.dtype), 0.) - b_v2 = rstd * (V * b_v - tl.sum(b_v, axis=1, keep_dims=True) - b_kh_hat.to(b_k.dtype) - * tl.sum(b_v * b_kh_hat.to(b_k.dtype), axis=1, keep_dims=True)) / V - tl.store(p_x, b_kh_hat.to(p_x.dtype.element_ty), boundary_check=(0, 1)) - tl.store(p_y, b_v.to(p_y.dtype.element_ty), boundary_check=(0, 1)) - tl.store(p_r, rstd.to(p_r.dtype.element_ty), boundary_check=(0, 1)) - tl.store(p_v2, b_v2.to(p_v2.dtype.element_ty), boundary_check=(0, 1)) - - b_e = tl.load(p_e, boundary_check=(0,), padding_option="zero") - b_do = tl.load(p_do, boundary_check=(0, 1), padding_option="zero") - - b_v2 = tl.where((v_i < V)[None, :], b_v2, 0.) - b_ds = tl.dot(b_do, tl.trans(b_v2).to(b_do.dtype)) - b_ds = tl.where(m_A, b_ds, 0) - b_ds = b_ds.to(b_k.dtype) - b_dq = tl.dot(b_do, tl.trans(b_h).to(b_do.dtype)) - b_dq -= tl.dot(b_ds, tl.trans(b_k)) * b_e[:, None] - b_dq *= scale - - b_e_last = tl.load(p_e_last) - b_h = b_h - tl.dot(b_e_last * b_k, b_v2.to(b_k.dtype), allow_tf32=False) - b_hb = b_hb - tl.sum(b_e_last * b_v2.to(b_k.dtype), axis=0) - b_h = tl.where((v_i < V)[None, :], b_h, 0.) - b_hb = tl.where((v_i < V), b_hb, 0.) - tl.store(p_dq, b_dq.to(p_dq.dtype.element_ty), boundary_check=(0, 1)) - - -@triton.heuristics({ - 'USE_INITIAL_STATE': lambda args: args['dh0'] is not None, - 'USE_INITIAL_STATE_B': lambda args: args['dhb0'] is not None, - 'USE_FINAL_STATE_GRADIENT': lambda args: args['dht'] is not None, - 'USE_FINAL_STATE_GRADIENT_B': lambda args: args['dhbt'] is not None, -}) -@triton.autotune( - configs=[ - triton.Config({}, num_warps=num_warps) - for num_warps in NUM_WARPS - ], - key=['BT', 'BK', 'BV'], - **autotune_cache_kwargs, -) -@triton.jit(do_not_specialize=['T']) -def fused_chunk_ttt_linear_bwd_kernel_dh( - q, - k, - v, - v2, - x, - y, - r, - w, - b, - eta, - h, - dht, - dhbt, - dh0, - dhb0, - do, - dk, - dv, - de, - dw, - db, - scale, - T, - H: tl.constexpr, - K: tl.constexpr, - V: tl.constexpr, - BT: tl.constexpr, - BK: tl.constexpr, - BV: tl.constexpr, - USE_INITIAL_STATE: tl.constexpr, - USE_INITIAL_STATE_B: tl.constexpr, - USE_FINAL_STATE_GRADIENT: tl.constexpr, - USE_FINAL_STATE_GRADIENT_B: tl.constexpr, -): - i_nh = tl.program_id(0) - i_n, i_h = i_nh // H, i_nh % H - bos, _ = i_n * T, i_n * T + T - NT = tl.cdiv(T, BT) - boh = i_n * NT - - # [BK, BV] - b_dh = tl.zeros([BK, BV], dtype=tl.float32) - # [BV] - b_dhb = tl.zeros([BV], dtype=tl.float32) - if USE_FINAL_STATE_GRADIENT: - p_dht = tl.make_block_ptr(dht + i_nh * K*V, (K, V), (V, 1), (0, 0), (BK, BV), (1, 0)) - b_dh += tl.load(p_dht, boundary_check=(0, 1), padding_option="zero") - if USE_FINAL_STATE_GRADIENT_B: - p_dhbt = tl.make_block_ptr(dhbt + i_nh * V, (V,), (1,), (0,), (BV,), (0,)) - b_dhb += tl.load(p_dhbt, boundary_check=(0,), padding_option="zero") - - # [BV] - o_i = tl.arange(0, BT) - v_i = tl.arange(0, BV) - m_A = o_i[:, None] >= o_i[None, :] - m_A_t = o_i[:, None] <= o_i[None, :] - b_w = tl.load(w + i_h * V + v_i, mask=v_i < V, other=0.) - b_b = tl.load(b + i_h * V + v_i, mask=v_i < V, other=0.) - b_dw = tl.zeros([BV], dtype=b_w.dtype) - b_db = tl.zeros([BV], dtype=b_b.dtype) - p_dw = tl.make_block_ptr(dw + i_nh * V, (V,), (1,), (0,), (BV,), (0,)) - p_db = tl.make_block_ptr(db + i_nh * V, (V,), (1,), (0,), (BV,), (0,)) - - for i_t in range(NT - 1, -1, -1): - p_h = tl.make_block_ptr(h+((boh+i_t)*H+i_h)*K*V, (V, K), (1, V), (0, 0), (BV, BK), (0, 1)) - p_q = tl.make_block_ptr(q+(bos*H+i_h)*K, (K, T), (1, H*K), (0, i_t*BT), (BK, BT), (0, 1)) - p_k = tl.make_block_ptr(k+(bos*H+i_h)*K, (T, K), (H*K, 1), (i_t*BT, 0), (BT, BK), (1, 0)) - p_v = tl.make_block_ptr(v+(bos*H+i_h)*V, (T, V), (H*V, 1), (i_t*BT, 0), (BT, BV), (1, 0)) - p_v2 = tl.make_block_ptr(v2+(bos*H+i_h)*V, (T, V), (H*V, 1), (i_t*BT, 0), (BT, BV), (1, 0)) - p_x = tl.make_block_ptr(x+(bos*H+i_h)*V, (T, V), (H*V, 1), (i_t*BT, 0), (BT, BV), (1, 0)) - p_y = tl.make_block_ptr(y+(bos*H+i_h)*V, (T, V), (H*V, 1), (i_t*BT, 0), (BT, BV), (1, 0)) - p_r = tl.make_block_ptr(r+bos*H+i_h, (T, 1), (H, 1), (i_t*BT, 0), (BT, 1), (1, 0)) - p_e = tl.make_block_ptr(eta+(bos*H+i_h), (T,), (H,), (i_t*BT,), (BT,), (0,)) - p_dv = tl.make_block_ptr(dv+(bos*H+i_h)*V, (T, V), (H*V, 1), (i_t*BT, 0), (BT, BV), (1, 0)) - p_dk = tl.make_block_ptr(dk+(bos*H+i_h)*K, (T, K), (H*K, 1), (i_t*BT, 0), (BT, BK), (1, 0)) - p_do = tl.make_block_ptr(do+(bos*H+i_h)*V, (T, V), (H*V, 1), (i_t*BT, 0), (BT, BV), (1, 0)) - p_de = tl.make_block_ptr(de+(bos*H+i_h), (T,), (H,), (i_t*BT,), (BT,), (0,)) - p_e_last = eta+bos*H+i_h + (T-1)*H if i_t == NT-1 else eta+bos*H+i_h + (i_t*BT+BT-1)*H - b_q = tl.load(p_q, boundary_check=(0, 1), padding_option="zero") - b_k = tl.load(p_k, boundary_check=(0, 1), padding_option="zero") - b_e = tl.load(p_e, boundary_check=(0,), padding_option="zero") - b_do = tl.load(p_do, boundary_check=(0, 1), padding_option="zero") - b_e_last = tl.load(p_e_last) - b_A = tl.dot(b_k, b_q) - b_A = - tl.where(m_A_t, b_A * scale * b_e[None, :], 0).to(do.dtype.element_ty) - b_Ae = - tl.where(m_A_t, b_e[None, :], 0).to(do.dtype.element_ty) - b_dv_new = tl.dot(b_A.to(b_do.dtype), b_do) + tl.dot(b_Ae.to(b_do.dtype), b_do) - b_dv_new -= tl.dot(b_e_last * b_k, b_dh.to(b_k.dtype)) - b_dv_new -= b_e_last * b_dhb.to(b_k.dtype)[None, :] - - b_v2 = tl.load(p_v2, boundary_check=(0, 1), padding_option="zero").to(b_k.dtype) - b_x = tl.load(p_x, boundary_check=(0, 1), padding_option="zero").to(b_k.dtype) - b_y = tl.load(p_y, boundary_check=(0, 1), padding_option="zero").to(b_k.dtype) - b_rstd = tl.load(p_r, boundary_check=(0, 1), padding_option="zero").to(tl.float32) - b_dy = b_rstd * (b_dv_new * V - tl.sum(b_dv_new, axis=1, keep_dims=True) - - b_x * tl.sum(b_dv_new * b_x, axis=1, keep_dims=True)) / V - b_dx = -b_rstd * (b_dv_new * tl.sum(b_x * b_y, axis=1, keep_dims=True) + - b_y * tl.sum(b_dv_new * b_x, axis=1, keep_dims=True)) / V - b_drstd = tl.sum(b_dv_new.to(b_rstd.dtype) * b_v2.to(b_rstd.dtype) / b_rstd, axis=1, keep_dims=True) - - b_v = tl.load(p_v, boundary_check=(0, 1), padding_option="zero") - b_w = b_w.to(b_k.dtype) - b_b = b_b.to(b_k.dtype) - b_dv = -b_w * b_dy.to(b_k.dtype) - b_dk = b_w * b_dy.to(b_k.dtype) - b_dw += tl.sum(2 * b_w * b_x * b_dy.to(b_k.dtype) + - (b_b - b_v.to(b_k.dtype) + b_k) * b_dy.to(b_k.dtype), axis=0).to(b_dw.dtype) - b_db += tl.sum(b_w * b_dy.to(b_k.dtype), axis=0).to(b_db.dtype) - b_dx = b_dx.to(b_k.dtype) + b_w * b_w * b_dy.to(b_k.dtype) - - b_h = tl.load(p_h, boundary_check=(0, 1), padding_option="zero") - b_q = (b_q * scale).to(b_q.dtype) - b_dkh = b_rstd * (V * b_dx - tl.sum(b_dx, axis=1, keep_dims=True) - - b_x * tl.sum(b_x * b_dx, axis=1, keep_dims=True)) / V - b_dkh -= b_rstd * b_rstd * b_drstd * b_x / V - b_dkh = tl.where((v_i < V)[None, :] * (o_i < T-i_t*BT)[:, None], b_dkh, 0.) - b_dk += tl.dot(b_dkh, b_h.to(b_dkh.dtype)).to(b_k.dtype) - - b_ds = tl.dot(b_do, tl.trans(b_v2)) - b_ds = tl.where(m_A, b_ds, 0) - b_ds = b_ds.to(b_k.dtype) - i_last = (BT-1) if (i_t*BT+BT) <= T else (T % BT-1) - mask = (o_i == i_last) - b_dk -= b_e_last * tl.dot(b_v2, tl.trans(b_dh).to(b_v2.dtype)) - b_dk -= tl.dot(tl.trans(b_ds), tl.trans(b_q) * b_e[:, None]) - b_de = mask * tl.sum(- b_dh * tl.trans(tl.dot(tl.trans(b_v2), b_k))).to(b_k.dtype) - b_de -= mask * tl.sum(b_dhb * tl.sum(b_v2, axis=0)).to(b_k.dtype) - b_de -= tl.sum(tl.dot(b_ds, b_k) * tl.trans(b_q).to(b_k.dtype), axis=1) - b_de -= tl.sum(b_ds, axis=1) - b_dh += tl.dot(b_q, b_do.to(b_q.dtype)) + tl.dot(tl.trans(b_k).to(b_dkh.dtype), b_dkh) - b_dhb += tl.sum(b_do + b_dkh, axis=0) - b_dh = tl.where((v_i < V)[None, :], b_dh, 0.) - b_dhb = tl.where((v_i < V), b_dhb, 0.) - - tl.store(p_dk, b_dk.to(p_dk.dtype.element_ty), boundary_check=(0, 1)) - tl.store(p_dv, b_dv.to(p_dv.dtype.element_ty), boundary_check=(0, 1)) - tl.store(p_de, b_de.to(p_de.dtype.element_ty), boundary_check=(0,)) - tl.store(p_dw, b_dw.to(p_dw.dtype.element_ty), boundary_check=(0,)) - tl.store(p_db, b_db.to(p_db.dtype.element_ty), boundary_check=(0,)) - - if USE_INITIAL_STATE: - p_dh0 = tl.make_block_ptr(dh0+i_nh*K*V, (K, V), (V, 1), (0, 0), (BK, BV), (1, 0)) - tl.store(p_dh0, b_dh.to(p_dh0.dtype.element_ty), boundary_check=(0, 1)) - if USE_INITIAL_STATE_B: - p_dhb0 = tl.make_block_ptr(dhb0+i_nh*V, (V,), (1,), (0,), (BV,), (0,)) - tl.store(p_dhb0, b_dhb.to(p_dhb0.dtype.element_ty), boundary_check=(0,)) - - -def fused_chunk_ttt_linear_bwd_h( - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - w: torch.Tensor, - b: torch.Tensor, - eta: torch.Tensor, - scale: float, - eps: float, - do: torch.Tensor, - BT: int = 16, - initial_state: torch.Tensor = None, - initial_state_bias: torch.Tensor = None, - cu_seqlens: torch.LongTensor | None = None, -): - assert cu_seqlens is None, "bwd of varlen is not implemented yet." - B, T, H, K, V = *k.shape, v.shape[-1] - # N: the actual number of sequences in the batch with either equal or variable lengths - N, NT = B, triton.cdiv(T, BT) - BK, BV = max(triton.next_power_of_2(K), 16), max(triton.next_power_of_2(V), 16) - assert max(BK, BV) <= 128, "current kernel does not support head dimension larger than 128." - - h = k.new_empty(B, NT, H, K, V) - r = v.new_empty(B, T, H, 1, dtype=torch.float32) - v2 = torch.empty_like(v) - x = torch.empty_like(v) - y = torch.empty_like(v) - dq = torch.empty_like(q) - - grid = (N * H,) - fused_chunk_ttt_linear_bwd_kernel_h[grid]( - k=k, - v=v, - v2=v2, - x=x, - y=y, - r=r, - w=w, - b=b, - eta=eta, - h0=initial_state, - hb0=initial_state_bias, - h=h, - do=do, - dq=dq, - scale=scale, - eps=eps, - T=T, - H=H, - K=K, - V=V, - BT=BT, - BK=BK, - BV=BV, - ) - return dq, h, v2, x, y, r - - -def fused_chunk_ttt_linear_bwd_dh( - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - v2: torch.Tensor, - x: torch.Tensor, - y: torch.Tensor, - r: torch.Tensor, - w: torch.Tensor, - b: torch.Tensor, - eta: torch.Tensor, - scale: float, - h: torch.Tensor, - do: torch.Tensor, - dht: torch.Tensor, - dhbt: torch.Tensor, - BT: int = 16, - initial_state: torch.Tensor = None, - initial_state_bias: torch.Tensor = None, - cu_seqlens: torch.LongTensor | None = None, -): - assert cu_seqlens is None, "bwd of varlen is not implemented yet." - B, T, H, K, V = *k.shape, v.shape[-1] - # N: the actual number of sequences in the batch with either equal or variable lengths - N = B - BK, BV = max(triton.next_power_of_2(K), 16), max(triton.next_power_of_2(V), 16) - assert max(BK, BV) <= 128, "current kernel does not support head dimension larger than 128." - - dh0 = torch.empty_like(initial_state, dtype=torch.float32) if initial_state is not None else None - dhb0 = torch.empty_like(initial_state_bias, dtype=torch.float32) if initial_state_bias is not None else None - dk = torch.empty_like(k) - dv = torch.empty_like(v) - de = torch.empty_like(eta) - dw = w.new_empty(B, H, V) - db = b.new_empty(B, H, V) - - grid = (N * H,) - fused_chunk_ttt_linear_bwd_kernel_dh[grid]( - q=q, - k=k, - v=v, - v2=v2, - x=x, - y=y, - r=r, - w=w, - b=b, - eta=eta, - h=h, - dht=dht, - dhbt=dhbt, - dh0=dh0, - dhb0=dhb0, - do=do, - dk=dk, - dv=dv, - de=de, - dw=dw, - db=db, - scale=scale, - T=T, - H=H, - K=K, - V=V, - BT=BT, - BK=BK, - BV=BV, - ) - dw = dw.sum(dim=0) - db = db.sum(dim=0) - return dk, dv, de, dw, db, dh0, dhb0 - - -def fused_chunk_ttt_linear_fwd( - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - w: torch.Tensor, - b: torch.Tensor, - eta: torch.Tensor, - scale: float, - eps: float, - initial_state: torch.Tensor, - initial_state_bias: torch.Tensor, - output_final_state: bool, - cu_seqlens: torch.LongTensor | None = None, - BT: int = 16, -): - B, T, H, K, V = *k.shape, v.shape[-1] - # N: the actual number of sequences in the batch with either equal or variable lengths - N = B if cu_seqlens is None else len(cu_seqlens) - 1 - BK, BV = max(triton.next_power_of_2(K), 16), max(triton.next_power_of_2(V), 16) - assert max(BK, BV) <= 128, "current kernel does not support head dimension larger than 128." - o = torch.empty_like(v) - final_state = k.new_empty(N, H, K, V, dtype=torch.float32) if output_final_state else None - final_state_bias = k.new_empty(N, H, 1, V, dtype=torch.float32) if output_final_state else None - - grid = (N * H,) - fused_chunk_ttt_linear_fwd_kernel[grid]( - q=q, - k=k, - v=v, - eta=eta, - w=w, - b=b, - o=o, - scale=scale, - eps=eps, - h0=initial_state, - hb0=initial_state_bias, - ht=final_state, - hbt=final_state_bias, - cu_seqlens=cu_seqlens, - T=T, - H=H, - K=K, - V=V, - BT=BT, - BK=BK, - BV=BV, - ) - return o, final_state, final_state_bias - - -def fused_chunk_ttt_linear_bwd( - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - w: torch.Tensor, - b: torch.Tensor, - eta: torch.Tensor, - scale: float, - eps: float, - do: torch.Tensor, - dht: torch.Tensor, - dhbt: torch.Tensor, - BT: int = 16, - initial_state: torch.Tensor = None, - initial_state_bias: torch.Tensor = None, - cu_seqlens: torch.LongTensor | None = None, -): - assert cu_seqlens is None, "bwd of varlen is not implemented yet." - dq, h, v2, x, y, rstd = fused_chunk_ttt_linear_bwd_h( - q=q, - k=k, - v=v, - w=w, - b=b, - eta=eta, - scale=scale, - eps=eps, - do=do, - BT=BT, - initial_state=initial_state, - initial_state_bias=initial_state_bias, - cu_seqlens=cu_seqlens, - ) - dk, dv, de, dw, db, dh0, dhb0 = fused_chunk_ttt_linear_bwd_dh( - q=q, - k=k, - v=v, - v2=v2, - x=x, - y=y, - r=rstd, - w=w, - b=b, - eta=eta, - scale=scale, - h=h, - do=do, - dht=dht, - dhbt=dhbt, - BT=BT, - initial_state=initial_state, - initial_state_bias=initial_state_bias, - cu_seqlens=cu_seqlens, - ) - return dq, dk, dv, de, dw, db, dh0, dhb0 - - -class FusedChunkTTTLinearFunction(torch.autograd.Function): - - @staticmethod - @input_guard - @autocast_custom_fwd - def forward(ctx, q, k, v, w, b, BT, eta, scale, eps, initial_state, - initial_state_bias, output_final_state, cu_seqlens): - o, final_state, final_state_bias = fused_chunk_ttt_linear_fwd( - q=q, - k=k, - v=v, - w=w, - b=b, - eta=eta, - scale=scale, - eps=eps, - BT=BT, - initial_state=initial_state, - initial_state_bias=initial_state_bias, - output_final_state=output_final_state, - cu_seqlens=cu_seqlens, - ) - ctx.save_for_backward(q, k, v, eta, w, b, initial_state, initial_state_bias) - ctx.BT = BT - ctx.scale = scale - ctx.eps = eps - ctx.cu_seqlens = cu_seqlens - return o.to(q.dtype), final_state, final_state_bias - - @staticmethod - @input_guard - @autocast_custom_bwd - def backward(ctx, do, dht, dhbt): - q, k, v, eta, w, b, initial_state, initial_state_bias = ctx.saved_tensors - dq, dk, dv, de, dw, db, dh0, dhb0 = fused_chunk_ttt_linear_bwd( - q=q, - k=k, - v=v, - w=w, - b=b, - eta=eta, - scale=ctx.scale, - eps=ctx.eps, - do=do, - dht=dht, - dhbt=dhbt, - BT=ctx.BT, - initial_state=initial_state, - initial_state_bias=initial_state_bias, - cu_seqlens=ctx.cu_seqlens, - ) - return dq.to(q), dk.to(k), dv.to(v), dw.to(w), db.to(b), None, de.to(eta), None, None, dh0, dhb0, None, None - - -def norm_residual(x, weight, bias, eps): - # GroupNorm and Residual - B, T, H, D = x.shape - x += group_norm( - x.reshape(B, T, -1).clone(), - weight=weight.reshape(-1).clone(), - bias=bias.reshape(-1).clone(), - eps=eps, - num_groups=H, - ).reshape(x.shape) - return x - - -def fused_chunk_ttt_linear( - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - w: torch.Tensor, - b: torch.Tensor, - eta: torch.Tensor, - scale: float | None = None, - eps: float = 1e-6, - chunk_size: int = 16, - initial_state: torch.Tensor | None = None, - initial_state_bias: torch.Tensor | None = None, - output_final_state: bool = False, - cu_seqlens: torch.LongTensor | None = None, - **kwargs, -): - r""" - Args: - q (torch.Tensor): - queries of shape `[B, T, H, K]`. - k (torch.Tensor): - keys of shape `[B, T, H, K]`. - v (torch.Tensor): - values of shape `[B, T, H, V]`. - w (torch.Tensor): - layer norm weight of shape `[H, V]`. - b (torch.Tensor): - layer norm bias of shape `[H, V]`. - eta (torch.Tensor): - Learning rate for hidden state, of shape `[B, T, H, 1]`. - scale (Optional[float]): - Scale factor for the attention scores. - If not provided, it will default to `1 / sqrt(K)`. Default: `None`. - chunk_size (int): - Chunk size. Default: `16`. - initial_state (Optional[torch.Tensor]): - Initial state of shape `[N, H, K, V]`. Default: `None`. - initial_state_bias (Optional[torch.Tensor]): - Initial state bias of shape `[N, H, 1, V]`. Default: `None`. - output_final_state (Optional[bool]): - Whether to output the final state of shape `[N, H, K, V]`. Default: `False`. - cu_seqlens (torch.LongTensor): - Cumulative sequence lengths of shape `[N+1]` used for variable-length training, - consistent with the FlashAttention API. - - Returns: - o (torch.Tensor): - Outputs of shape `[B, T, H, V]`. - final_state (torch.Tensor): - Final state of shape `[N, H, K, V]` if `output_final_state=True` else `None`. - final_state_bias (torch.Tensor): - Final state bias of shape `[N, H, 1, V]` if `output_final_state=True` else `None`. - """ - assert q.dtype == k.dtype == v.dtype - assert k.shape[-1] == v.shape[-1], "DK must equal to DV." - if isinstance(eta, float): - eta = torch.full_like(q[:, :, :, :1], eta) - if 'head_first' in kwargs: - raise DeprecationWarning( - "head_first has been removed. Inputs must be in `[B, T, H, ...]` format.", - ) - if cu_seqlens is not None: - if q.shape[0] != 1: - raise ValueError( - f"The batch size is expected to be 1 rather than {q.shape[0]} when using `cu_seqlens`. " - f"Please flatten variable-length inputs before processing.", - ) - if initial_state is not None and initial_state.shape[0] != len(cu_seqlens) - 1: - raise ValueError( - f"The number of initial states is expected to be equal to the number of input sequences, " - f"i.e., {len(cu_seqlens) - 1} rather than {initial_state.shape[0]}.", - ) - if scale is None: - scale = k.shape[-1] ** -0.5 - else: - assert scale > 0, "Scale must be positive." - o, final_state, final_state_bias = FusedChunkTTTLinearFunction.apply( - q, - k, - v, - w, - b, - chunk_size, - eta, - scale, - eps, - initial_state, - initial_state_bias, - output_final_state, - cu_seqlens, - ) - o = norm_residual(o, w, b, eps) - return o, final_state, final_state_bias diff --git a/build/torch-xpu/ops/ttt/naive.py b/build/torch-xpu/ops/ttt/naive.py deleted file mode 100644 index 55a543534970f6c2bb038f8de4979a65f68bea31..0000000000000000000000000000000000000000 --- a/build/torch-xpu/ops/ttt/naive.py +++ /dev/null @@ -1,130 +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 torch.nn.functional as F - - -def ttt_linear( - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - w: torch.Tensor, - b: torch.Tensor, - eta: torch.Tensor, - scale: float, - eps: float, - mini_batch_size: int, - initial_state: torch.Tensor, - initial_state_bias: torch.Tensor, - output_final_state: bool, -): - B, H, T, D = q.shape - BT = mini_batch_size - NT = T // BT - # [NT, B, H, mini_batch_size, D] - _q = q.reshape(B, H, NT, BT, D).permute(2, 0, 1, 3, 4) - _k = k.reshape(B, H, NT, BT, D).permute(2, 0, 1, 3, 4) - _v = v.reshape(B, H, NT, BT, D).permute(2, 0, 1, 3, 4) - # [NT, B, H, BT, 1] - _eta = eta.reshape(B, H, NT, BT, 1).permute(2, 0, 1, 3, 4) - # [H, 1, D] - w = w.reshape(H, 1, D).to(torch.float32) - b = b.reshape(H, 1, D).to(torch.float32) - - h = torch.zeros((B, H, D, D), device=v.device, dtype=torch.float32) if initial_state is None else initial_state - hb = torch.zeros((B, H, 1, D), device=v.device, dtype=torch.float32) if initial_state_bias is None else initial_state_bias - q *= scale - # [NT, B, H, BT, D] - o = torch.empty_like(_v) - - for i in range(NT): - q_i, k_i, v_i, eta_i = [x[i] for x in [_q, _k, _v, _eta]] - kh = k_i @ h + hb - reconstruction_target = v_i - k_i - - mean = kh.mean(-1, True) - var = kh.var(-1, unbiased=False, keepdim=True).to(torch.float32) - rstd = torch.sqrt(var + eps).to(torch.float32) - kh_hat = (kh - mean) / rstd - - g = w * kh_hat + b - reconstruction_target - g *= w - v_new = (D * g - g.sum(-1, True) - kh_hat * (g * kh_hat).sum(-1, True)) / (rstd * D) - - Attn = torch.tril(q_i @ k_i.transpose(-2, -1)) - o_i = q_i @ h - (eta_i * Attn) @ v_new + hb - torch.tril(eta_i.expand_as(Attn)) @ v_new - h = h - (eta_i[:, :, -1, :, None] * k_i).transpose(-1, -2) @ v_new - hb = hb - torch.sum(eta_i[:, :, -1, :, None] * v_new, dim=-2, keepdim=True) - # layer norm with residuals - - mean = o_i.mean(dim=-1, keepdim=True) - var = o_i.var(dim=-1, unbiased=False, keepdim=True).to(torch.float32) - rstd = torch.sqrt(var + eps).to(torch.float32) - o[i] = o_i + (o_i - mean) / rstd * w + b - - # [B, H, T, D] - o = o.permute(1, 2, 0, 3, 4).reshape(B, H, T, D) - h = h if output_final_state else None - hb = hb if output_final_state else None - return o, h, hb - - -def chunk_ttt_linear_ref( - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - w: torch.Tensor, - b: torch.Tensor, - eta: torch.Tensor, - scale: float = None, - eps: float = 1e-6, - mini_batch_size: int = 16, - initial_state: torch.Tensor = None, - initial_state_bias: torch.Tensor = None, - output_final_state: bool = False, - head_first: bool = False, -): - assert q.dtype == k.dtype == v.dtype - assert k.shape[-1] == v.shape[-1], "The key and value dimension must be the same." - if isinstance(eta, float): - eta = torch.full_like(q[:, :, :, :1], eta) - if scale is None: - scale = k.shape[-1] ** -0.5 - if not head_first: - q = q.transpose(1, 2) - k = k.transpose(1, 2) - v = v.transpose(1, 2) - eta = eta.transpose(1, 2) - T = q.shape[-2] - padded = (mini_batch_size - (T % mini_batch_size)) % mini_batch_size - if padded > 0: - q = F.pad(q, (0, 0, 0, padded)) - k = F.pad(k, (0, 0, 0, padded)) - v = F.pad(v, (0, 0, 0, padded)) - eta = F.pad(eta, (0, 0, 0, padded)) - eta[:, :, -1, :] = eta[:, :, -(padded+1), :] - assert q.shape[-2] % mini_batch_size == 0, "Sequence length should be a multiple of mini_batch_size." - q, k, v, eta, w, b = map(lambda x: x.to(torch.float32), [q, k, v, eta, w, b]) - o, final_state, final_state_bias = ttt_linear( - q, - k, - v, - w, - b, - eta, - scale, - eps, - mini_batch_size, - initial_state, - initial_state_bias, - output_final_state, - ) - o = o[:, :, :T, :].contiguous() - if not head_first: - o = o.transpose(1, 2) - return o, final_state, final_state_bias diff --git a/build/torch-xpu/ops/utils/__init__.py b/build/torch-xpu/ops/utils/__init__.py deleted file mode 100644 index 88acd8b9cb58d763c2b81ea68337ac445e6658cf..0000000000000000000000000000000000000000 --- a/build/torch-xpu/ops/utils/__init__.py +++ /dev/null @@ -1,65 +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 .csr import prepare_block_csr -from .cumsum import ( - chunk_global_cumsum, - chunk_global_cumsum_scalar, - chunk_global_cumsum_vector, - chunk_local_cumsum, - chunk_local_cumsum_scalar, - chunk_local_cumsum_vector, -) -from .index import ( - get_max_num_splits, - prepare_chunk_indices, - prepare_chunk_offsets, - prepare_cu_seqlens_from_lens, - prepare_cu_seqlens_from_mask, - prepare_lens, - prepare_lens_from_mask, - prepare_position_ids, - prepare_sequence_ids, - prepare_token_indices, -) -from .logsumexp import logsumexp_fwd -from .matmul import addmm, matmul -from .pack import pack_sequence, unpack_sequence -from .pooling import mean_pooling -from .softmax import softmax_bwd, softmax_fwd -from .softplus import softplus -from .solve_tril import solve_tril - -__all__ = [ - "addmm", - "chunk_global_cumsum", - "chunk_global_cumsum_scalar", - "chunk_global_cumsum_vector", - "chunk_local_cumsum", - "chunk_local_cumsum_scalar", - "chunk_local_cumsum_vector", - "get_max_num_splits", - "logsumexp_fwd", - "matmul", - "mean_pooling", - "pack_sequence", - "prepare_block_csr", - "prepare_chunk_indices", - "prepare_chunk_offsets", - "prepare_cu_seqlens_from_lens", - "prepare_cu_seqlens_from_mask", - "prepare_lens", - "prepare_lens_from_mask", - "prepare_position_ids", - "prepare_sequence_ids", - "prepare_token_indices", - "softmax_bwd", - "softmax_fwd", - "softplus", - "solve_tril", - "unpack_sequence", -] diff --git a/build/torch-xpu/ops/utils/backends/triton_ascend/__init__.py b/build/torch-xpu/ops/utils/backends/triton_ascend/__init__.py deleted file mode 100644 index e7f66dd814b9a23286971b95d6acf10f19c9459d..0000000000000000000000000000000000000000 --- a/build/torch-xpu/ops/utils/backends/triton_ascend/__init__.py +++ /dev/null @@ -1,38 +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 NPU backend for utils ops.""" - -from __future__ import annotations - -from .....ops.backends import BaseBackend - - -class TritonAscendUtilsBackend(BaseBackend): - 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 chunk_global_cumsum_verifier(self, *args, **kwargs): - return True, None - - def chunk_global_cumsum(self, *args, **kwargs): - from .cumsum import chunk_global_cumsum_npu - return chunk_global_cumsum_npu(*args, **kwargs) - - def chunk_local_cumsum_verifier(self, *args, **kwargs): - return True, None - - def chunk_local_cumsum(self, *args, **kwargs): - from .cumsum import chunk_local_cumsum_npu - return chunk_local_cumsum_npu(*args, **kwargs) diff --git a/build/torch-xpu/ops/utils/backends/triton_ascend/cumsum.py b/build/torch-xpu/ops/utils/backends/triton_ascend/cumsum.py deleted file mode 100644 index 514b413500fec396b63546b39772b3aa2596009d..0000000000000000000000000000000000000000 --- a/build/torch-xpu/ops/utils/backends/triton_ascend/cumsum.py +++ /dev/null @@ -1,631 +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 torch -import triton -import triton.language as tl - -from .....ops.utils.index import prepare_chunk_indices -from .....utils import input_guard -from .....utils.ascend_ub_manager import ( - ASCEND_MAX_GRID_DIM, - compute_grid_limited_tile_size, - compute_row_tile_block_size, - compute_ub_block_size, - iter_axis_launch_chunks, - max_grid_axis_chunks, -) - -_NUM_WARPS = 4 -# Peak live fp32 tiles: b_s, b_o (and b_z / partial sums for vector/global). -_CUMSUM_SCALAR_MEM_MULT = 3.0 -# b_s, b_c, b_z, plus tl.cumsum multi-buffer on [BT, BS] tiles. -_CUMSUM_VECTOR_MEM_MULT = 8.0 -_CUMSUM_SAFETY_MARGIN = 0.85 -_FALLBACK_BT_GLOBAL = 32 -_FALLBACK_BS_LOCAL = 32 -_FALLBACK_BS_GLOBAL = 16 -_MAX_BT_GLOBAL = 256 - - -def _get_global_scalar_bt(T: int) -> int: - """UB-safe time chunk size for global scalar cumsum.""" - desired = min(triton.next_power_of_2(T), _MAX_BT_GLOBAL) - return compute_ub_block_size( - T, - _CUMSUM_SCALAR_MEM_MULT, - safety_margin=_CUMSUM_SAFETY_MARGIN, - dtype_size=4, - fallback=_FALLBACK_BT_GLOBAL, - desired=desired, - ) - - -def _get_vector_bs(BT: int, S: int, *, fallback: int, max_block: int | None = None) -> int: - """UB-safe feature tile size for vector cumsum with fixed time chunk BT.""" - return compute_row_tile_block_size( - BT, - S, - _CUMSUM_VECTOR_MEM_MULT, - tiling_row=False, - safety_margin=_CUMSUM_SAFETY_MARGIN, - dtype_size=4, - fallback=fallback, - min_block=1, - max_block=max_block, - ) - - -def _get_global_vector_tile_config(T: int, S: int) -> tuple[int, int]: - """UB-safe (BT, BS) for global vector cumsum.""" - bs_cap = min(_FALLBACK_BS_GLOBAL, triton.next_power_of_2(S)) - BS = _get_vector_bs( - _FALLBACK_BT_GLOBAL, - S, - fallback=_FALLBACK_BS_GLOBAL, - max_block=bs_cap, - ) - BT = compute_row_tile_block_size( - T, - BS, - _CUMSUM_VECTOR_MEM_MULT, - tiling_row=True, - safety_margin=_CUMSUM_SAFETY_MARGIN, - dtype_size=4, - fallback=_FALLBACK_BT_GLOBAL, - min_block=1, - max_block=_FALLBACK_BT_GLOBAL, - ) - BS = _get_vector_bs(BT, S, fallback=_FALLBACK_BS_GLOBAL, max_block=bs_cap) - return BT, BS - - -def _launch_local_cumsum_scalar( - *, - g_org, - g, - scale, - cu_seqlens, - chunk_indices, - T, - B, - H, - BT, - NT, - head_first, - reverse, -): - bh_total = B * H - kernel_kwargs = dict( - s=g_org, - o=g, - scale=scale, - cu_seqlens=cu_seqlens, - T=T, - B=B, - H=H, - BT=BT, - HEAD_FIRST=head_first, - REVERSE=reverse, - num_warps=_NUM_WARPS, - ) - max_nt = max_grid_axis_chunks(NT, bh_total, max_grid=ASCEND_MAX_GRID_DIM) - for nt_off in range(0, NT, max_nt): - nt_len = min(max_nt, NT - nt_off) - if cu_seqlens is not None: - 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 - max_bh = max_grid_axis_chunks(bh_total, nt_len, max_grid=ASCEND_MAX_GRID_DIM) - for bh_off in range(0, bh_total, max_bh): - bh_len = min(max_bh, bh_total - bh_off) - kernel_kwargs['BH_OFFSET'] = bh_off - chunk_local_cumsum_scalar_kernel_npu[(nt_len, bh_len)](**kernel_kwargs) - - -def _launch_local_cumsum_vector( - *, - g_org, - g, - scale, - cu_seqlens, - chunk_indices, - T, - B, - H, - S, - BT, - BS, - NT, - head_first, - reverse, -): - bh_total = B * H - ns = triton.cdiv(S, BS) - kernel_kwargs = dict( - s=g_org, - o=g, - scale=scale, - cu_seqlens=cu_seqlens, - T=T, - B=B, - H=H, - S=S, - BT=BT, - BS=BS, - HEAD_FIRST=head_first, - REVERSE=reverse, - num_warps=_NUM_WARPS, - ) - max_nt = max_grid_axis_chunks(NT, ns * bh_total, max_grid=ASCEND_MAX_GRID_DIM) - for nt_off in range(0, NT, max_nt): - nt_len = min(max_nt, NT - nt_off) - if cu_seqlens is not None: - 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 - max_bh = max_grid_axis_chunks(bh_total, ns * nt_len, max_grid=ASCEND_MAX_GRID_DIM) - for bh_off in range(0, bh_total, max_bh): - bh_len = min(max_bh, bh_total - bh_off) - kernel_kwargs['BH_OFFSET'] = bh_off - chunk_local_cumsum_vector_kernel_npu[(ns, nt_len, bh_len)](**kernel_kwargs) - - -@triton.heuristics({ - 'HAS_SCALE': lambda args: args['scale'] is not None, - 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None, -}) -@triton.jit(do_not_specialize=['T']) -def chunk_local_cumsum_scalar_kernel_npu( - s, - o, - scale, - cu_seqlens, - chunk_indices, - T, - B: tl.constexpr, - H: tl.constexpr, - BT: tl.constexpr, - REVERSE: tl.constexpr, - HAS_SCALE: tl.constexpr, - IS_VARLEN: tl.constexpr, - HEAD_FIRST: tl.constexpr, - NT_OFFSET: tl.constexpr, - BH_OFFSET: tl.constexpr, -): - i_t, i_bh = tl.program_id(0), tl.program_id(1) - i_t += NT_OFFSET - i_bh += BH_OFFSET - i_b, i_h = i_bh // H, i_bh % H - 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.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) - T = eos - bos - else: - bos, eos = i_b * T, i_b * T + T - - if HEAD_FIRST: - p_s = tl.make_block_ptr(s + bos*H + i_h*T, (T,), (1,), (i_t * BT,), (BT,), (0,)) - p_o = tl.make_block_ptr(o + bos*H + i_h*T, (T,), (1,), (i_t * BT,), (BT,), (0,)) - else: - p_s = tl.make_block_ptr(s + bos*H + i_h, (T,), (H,), (i_t * BT,), (BT,), (0,)) - p_o = tl.make_block_ptr(o + bos*H + i_h, (T,), (H,), (i_t * BT,), (BT,), (0,)) - b_s = tl.load(p_s, boundary_check=(0,)).to(tl.float32) - b_o = tl.cumsum(b_s, axis=0) - if REVERSE: - b_z = tl.sum(b_s, axis=0) - b_o = -b_o + b_z[None] + b_s - if HAS_SCALE: - b_o *= scale - tl.store(p_o, b_o.to(p_o.dtype.element_ty), boundary_check=(0,)) - - -@triton.heuristics({ - 'HAS_SCALE': lambda args: args['scale'] is not None, - 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None, -}) -@triton.jit(do_not_specialize=['T']) -def chunk_local_cumsum_vector_kernel_npu( - s, - o, - scale, - cu_seqlens, - chunk_indices, - T, - B: tl.constexpr, - H: tl.constexpr, - S: tl.constexpr, - BT: tl.constexpr, - BS: tl.constexpr, - REVERSE: tl.constexpr, - HAS_SCALE: tl.constexpr, - IS_VARLEN: tl.constexpr, - HEAD_FIRST: tl.constexpr, - NT_OFFSET: tl.constexpr, - BH_OFFSET: tl.constexpr, -): - i_s, i_t, i_bh = tl.program_id(0), tl.program_id(1), tl.program_id(2) - i_t += NT_OFFSET - i_bh += BH_OFFSET - i_b, i_h = i_bh // H, i_bh % H - 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.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) - T = eos - bos - else: - bos, eos = i_b * T, i_b * T + T - - if HEAD_FIRST: - p_s = tl.make_block_ptr(s + (bos * H + i_h*T)*S, (T, S), (S, 1), (i_t * BT, i_s * BS), (BT, BS), (1, 0)) - p_o = tl.make_block_ptr(o + (bos * H + i_h*T)*S, (T, S), (S, 1), (i_t * BT, i_s * BS), (BT, BS), (1, 0)) - else: - p_s = tl.make_block_ptr(s + (bos * H + i_h) * S, (T, S), (H*S, 1), (i_t * BT, i_s * BS), (BT, BS), (1, 0)) - p_o = tl.make_block_ptr(o + (bos * H + i_h) * S, (T, S), (H*S, 1), (i_t * BT, i_s * BS), (BT, BS), (1, 0)) - b_s = tl.load(p_s, boundary_check=(0, 1)).to(tl.float32) - if REVERSE: - b_o = tl.cumsum(b_s, axis=0, reverse=True) - else: - b_o = tl.cumsum(b_s, axis=0) - if HAS_SCALE: - b_o *= scale - tl.store(p_o, b_o.to(p_o.dtype.element_ty), boundary_check=(0, 1)) - - -@triton.heuristics({ - 'HAS_SCALE': lambda args: args['scale'] is not None, - 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None, -}) -@triton.jit(do_not_specialize=['T']) -def chunk_global_cumsum_scalar_kernel_npu( - s, - o, - scale, - cu_seqlens, - T, - B: tl.constexpr, - H: tl.constexpr, - BT: tl.constexpr, - REVERSE: tl.constexpr, - HAS_SCALE: tl.constexpr, - IS_VARLEN: tl.constexpr, - HEAD_FIRST: tl.constexpr, - BH_OFFSET: tl.constexpr, -): - i_nh = tl.program_id(0) + BH_OFFSET - i_n, i_h = i_nh // H, i_nh % H - if IS_VARLEN: - bos, eos = tl.load(cu_seqlens + i_n).to(tl.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) - else: - bos, eos = i_n * T, i_n * T + T - T = eos - bos - - b_z = tl.zeros([], dtype=tl.float32) - NT = tl.cdiv(T, BT) - for i_c in range(NT): - i_t = NT - 1 - i_c if REVERSE else i_c - if HEAD_FIRST: - p_s = tl.make_block_ptr(s + bos*H + i_h*T, (T,), (1,), (i_t * BT,), (BT,), (0,)) - p_o = tl.make_block_ptr(o + bos*H + i_h*T, (T,), (1,), (i_t * BT,), (BT,), (0,)) - else: - p_s = tl.make_block_ptr(s + bos*H + i_h, (T,), (H,), (i_t * BT,), (BT,), (0,)) - p_o = tl.make_block_ptr(o + bos*H + i_h, (T,), (H,), (i_t * BT,), (BT,), (0,)) - b_s = tl.load(p_s, boundary_check=(0,)).to(tl.float32) - b_o = tl.cumsum(b_s, axis=0) - b_ss = tl.sum(b_s, 0) - if REVERSE: - b_o = -b_o + b_ss + b_s - b_o += b_z - if i_c >= 0: - b_z += b_ss - if HAS_SCALE: - b_o *= scale - tl.store(p_o, b_o.to(p_o.dtype.element_ty), boundary_check=(0,)) - - -@triton.heuristics({ - 'HAS_SCALE': lambda args: args['scale'] is not None, - 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None, -}) -@triton.jit(do_not_specialize=['T']) -def chunk_global_cumsum_vector_kernel_npu( - s, - o, - scale, - cu_seqlens, - T, - B: tl.constexpr, - H: tl.constexpr, - S: tl.constexpr, - BT: tl.constexpr, - BS: tl.constexpr, - REVERSE: tl.constexpr, - HAS_SCALE: tl.constexpr, - IS_VARLEN: tl.constexpr, - HEAD_FIRST: tl.constexpr, - BH_OFFSET: tl.constexpr, -): - i_s, i_nh = tl.program_id(0), tl.program_id(1) + BH_OFFSET - i_n, i_h = i_nh // H, i_nh % H - if IS_VARLEN: - bos, eos = tl.load(cu_seqlens + i_n).to(tl.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) - else: - bos, eos = i_n * T, i_n * T + T - T = eos - bos - - b_z = tl.zeros([BS], dtype=tl.float32) - NT = tl.cdiv(T, BT) - for i_c in range(NT): - i_t = NT - 1 - i_c if REVERSE else i_c - if HEAD_FIRST: - p_s = tl.make_block_ptr(s + (bos * H + i_h*T)*S, (T, S), (S, 1), (i_t * BT, i_s * BS), (BT, BS), (1, 0)) - p_o = tl.make_block_ptr(o + (bos * H + i_h*T)*S, (T, S), (S, 1), (i_t * BT, i_s * BS), (BT, BS), (1, 0)) - else: - p_s = tl.make_block_ptr(s + (bos * H + i_h) * S, (T, S), (H*S, 1), (i_t * BT, i_s * BS), (BT, BS), (1, 0)) - p_o = tl.make_block_ptr(o + (bos * H + i_h) * S, (T, S), (H*S, 1), (i_t * BT, i_s * BS), (BT, BS), (1, 0)) - b_s = tl.load(p_s, boundary_check=(0, 1)).to(tl.float32) - if REVERSE: - b_c = b_z[None, :] + tl.cumsum(b_s, axis=0, reverse=True) - else: - b_c = b_z[None, :] + tl.cumsum(b_s, axis=0) - if HAS_SCALE: - b_c *= scale - tl.store(p_o, b_c.to(p_o.dtype.element_ty), boundary_check=(0, 1)) - b_z += tl.sum(b_s, 0) - - -def chunk_local_cumsum_scalar_npu( - g: torch.Tensor, - chunk_size: int, - reverse: bool = False, - scale: float = None, - cu_seqlens: torch.Tensor | None = None, - head_first: bool = False, - output_dtype: torch.dtype | None = torch.float, - chunk_indices: torch.LongTensor | None = None, -) -> torch.Tensor: - if head_first: - B, H, T = g.shape - else: - B, T, H = g.shape - assert chunk_size == 2**(chunk_size.bit_length()-1), "chunk_size must be a power of 2" - BT = chunk_size - if chunk_indices is None and cu_seqlens is not None: - chunk_indices = prepare_chunk_indices(cu_seqlens, BT) - NT = triton.cdiv(T, BT) if cu_seqlens is None else len(chunk_indices) - g_org, g = g, torch.empty_like(g, dtype=output_dtype or g.dtype) - _launch_local_cumsum_scalar( - g_org=g_org, - g=g, - scale=scale, - cu_seqlens=cu_seqlens, - chunk_indices=chunk_indices, - T=T, - B=B, - H=H, - BT=BT, - NT=NT, - head_first=head_first, - reverse=reverse, - ) - return g - - -def chunk_local_cumsum_vector_npu( - g: torch.Tensor, - chunk_size: int, - reverse: bool = False, - scale: float = None, - cu_seqlens: torch.Tensor | None = None, - head_first: bool = False, - output_dtype: torch.dtype | None = torch.float, - chunk_indices: torch.LongTensor | None = None, -) -> torch.Tensor: - if head_first: - B, H, T, S = g.shape - else: - B, T, H, S = g.shape - BT = chunk_size - if chunk_indices is None and cu_seqlens is not None: - chunk_indices = prepare_chunk_indices(cu_seqlens, BT) - NT = triton.cdiv(T, BT) if cu_seqlens is None else len(chunk_indices) - assert chunk_size == 2**(chunk_size.bit_length()-1), "chunk_size must be a power of 2" - - BS = _get_vector_bs(BT, S, fallback=_FALLBACK_BS_LOCAL) - BS = compute_grid_limited_tile_size( - S, - NT * B * H, - BS, - max_grid=ASCEND_MAX_GRID_DIM, - ) - g_org, g = g, torch.empty_like(g, dtype=output_dtype or g.dtype) - _launch_local_cumsum_vector( - g_org=g_org, - g=g, - scale=scale, - cu_seqlens=cu_seqlens, - chunk_indices=chunk_indices, - T=T, - B=B, - H=H, - S=S, - BT=BT, - BS=BS, - NT=NT, - head_first=head_first, - reverse=reverse, - ) - return g - - -@input_guard -def chunk_global_cumsum_scalar_npu( - s: torch.Tensor, - reverse: bool = False, - cu_seqlens: torch.Tensor | None = None, - scale: float = None, - head_first: bool = False, - output_dtype: torch.dtype | None = torch.float, -) -> torch.Tensor: - if head_first: - B, H, T = s.shape - else: - B, T, H = s.shape - N = len(cu_seqlens) - 1 if cu_seqlens is not None else B - - BT = _get_global_scalar_bt(T) - z = torch.empty_like(s, dtype=output_dtype or s.dtype) - bh_total = N * H - kernel_kwargs = dict( - s=s, - o=z, - scale=scale, - cu_seqlens=cu_seqlens, - T=T, - B=B, - H=H, - BT=BT, - HEAD_FIRST=head_first, - REVERSE=reverse, - num_warps=_NUM_WARPS, - ) - for bh_off, bh_len in iter_axis_launch_chunks(bh_total, 1, max_grid=ASCEND_MAX_GRID_DIM): - kernel_kwargs['BH_OFFSET'] = bh_off - chunk_global_cumsum_scalar_kernel_npu[(bh_len,)](**kernel_kwargs) - return z - - -@input_guard -def chunk_global_cumsum_vector_npu( - s: torch.Tensor, - reverse: bool = False, - cu_seqlens: torch.Tensor | None = None, - scale: float = None, - head_first: bool = False, - output_dtype: torch.dtype | None = torch.float, -) -> torch.Tensor: - if head_first: - B, H, T, S = s.shape - else: - B, T, H, S = s.shape - N = len(cu_seqlens) - 1 if cu_seqlens is not None else B - BT, BS = _get_global_vector_tile_config(T, S) - BS = compute_grid_limited_tile_size( - S, - N * H, - BS, - max_grid=ASCEND_MAX_GRID_DIM, - ) - ns = triton.cdiv(S, BS) - - z = torch.empty_like(s, dtype=output_dtype or s.dtype) - bh_total = N * H - kernel_kwargs = dict( - s=s, - o=z, - scale=scale, - cu_seqlens=cu_seqlens, - T=T, - B=B, - H=H, - S=S, - BT=BT, - BS=BS, - HEAD_FIRST=head_first, - REVERSE=reverse, - num_warps=_NUM_WARPS, - ) - max_bh = max_grid_axis_chunks(bh_total, ns, max_grid=ASCEND_MAX_GRID_DIM) - for bh_off in range(0, bh_total, max_bh): - bh_len = min(max_bh, bh_total - bh_off) - kernel_kwargs['BH_OFFSET'] = bh_off - chunk_global_cumsum_vector_kernel_npu[(ns, bh_len)](**kernel_kwargs) - return z - - -@input_guard -def chunk_global_cumsum_npu( - s: torch.Tensor, - reverse: bool = False, - cu_seqlens: torch.Tensor | None = None, - scale: float = None, - head_first: bool = False, - output_dtype: torch.dtype | None = torch.float, -) -> torch.Tensor: - if cu_seqlens is not None: - assert s.shape[0] == 1, "Only batch size 1 is supported when cu_seqlens are provided" - if len(s.shape) == 3: - return chunk_global_cumsum_scalar_npu( - s=s, - reverse=reverse, - cu_seqlens=cu_seqlens, - scale=scale, - head_first=head_first, - output_dtype=output_dtype, - ) - if len(s.shape) == 4: - return chunk_global_cumsum_vector_npu( - s=s, - reverse=reverse, - cu_seqlens=cu_seqlens, - scale=scale, - head_first=head_first, - output_dtype=output_dtype, - ) - raise ValueError( - f"Unsupported input shape {s.shape}, " - f"which should be [B, T, H]/[B, T, H, D] if `head_first=False` " - f"or [B, H, T]/[B, H, T, D] otherwise", - ) - - -@input_guard -def chunk_local_cumsum_npu( - g: torch.Tensor, - chunk_size: int, - reverse: bool = False, - scale: float = None, - cu_seqlens: torch.Tensor | None = None, - head_first: bool = False, - output_dtype: torch.dtype | None = torch.float, - chunk_indices: torch.LongTensor | None = None, - **kwargs, -) -> torch.Tensor: - if cu_seqlens is not None: - assert g.shape[0] == 1, "Only batch size 1 is supported when cu_seqlens are provided" - if len(g.shape) == 3: - return chunk_local_cumsum_scalar_npu( - g=g, - chunk_size=chunk_size, - reverse=reverse, - scale=scale, - cu_seqlens=cu_seqlens, - head_first=head_first, - output_dtype=output_dtype, - chunk_indices=chunk_indices, - ) - if len(g.shape) == 4: - return chunk_local_cumsum_vector_npu( - g=g, - chunk_size=chunk_size, - reverse=reverse, - scale=scale, - cu_seqlens=cu_seqlens, - head_first=head_first, - output_dtype=output_dtype, - chunk_indices=chunk_indices, - ) - raise ValueError( - f"Unsupported input shape {g.shape}, " - f"which should be (B, T, H, D) if `head_first=False` " - f"or (B, H, T, D) otherwise", - ) diff --git a/build/torch-xpu/ops/utils/cache.py b/build/torch-xpu/ops/utils/cache.py deleted file mode 100644 index 40ab2dbc1802212f1e3e3ef4c97c2bc3ac79148f..0000000000000000000000000000000000000000 --- a/build/torch-xpu/ops/utils/cache.py +++ /dev/null @@ -1,449 +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 dataclasses -import enum -import json -import logging -import os -import re -from functools import cache, lru_cache -from pathlib import Path -from typing import Any - -import torch -import triton -from packaging import version -from triton.runtime.autotuner import Autotuner - -TRITON_ABOVE_3_5_1 = version.parse(triton.__version__) >= version.parse("3.5.1") -TRITON_ABOVE_3_4_0 = version.parse(triton.__version__) >= version.parse("3.4.0") - - -class FlaCacheMode(enum.Enum): - """Controls how FLA loads kernel configs from its config cache (FLA_CACHE_MODE env var). - - DISABLED — skip all cache lookups, always fall back to Triton autotune (default when FLA_CACHE_MODE is unset) - STRICT — exact key match only; falls back to Triton autotune if no match - FUZZY — exact key match → fuzzy key match; falls back to Triton autotune if no match - FULL — exact key match → fuzzy key match → default_config fallback - DEFAULT — use only the top-level default_config field, skip key-based lookup - ALWAYS — like DEFAULT, but re-reads config files on every kernel call; - useful for debugging: edit default_config in a JSON file and the next - kernel call picks it up without restarting the process - """ - DISABLED = "disabled" - STRICT = "strict" - FUZZY = "fuzzy" - FULL = "full" - DEFAULT = "default" - ALWAYS = "always" - - def uses_default_config(self) -> bool: - """Return True for modes that may fall back to default_config (FULL, DEFAULT, ALWAYS).""" - return self in (FlaCacheMode.FULL, FlaCacheMode.DEFAULT, FlaCacheMode.ALWAYS) - - @classmethod - def from_env(cls) -> "FlaCacheMode": - mode_str = os.environ.get("FLA_CACHE_MODE", cls.DISABLED.value) - try: - return cls(mode_str) - except ValueError: - valid = [m.value for m in cls] - raise ValueError( - f"Invalid FLA_CACHE_MODE={mode_str!r}. Valid values: {valid}" - ) from None - - -FLA_CACHE_MODE: FlaCacheMode = FlaCacheMode.from_env() -logger = logging.getLogger(__name__) - - -def sanitize_gpu_name(gpu_name: str) -> str: - sanitized = re.sub(r"[^0-9A-Za-z]+", "_", gpu_name) - sanitized = sanitized.strip("_") - return sanitized or "unknown_gpu" - - -@lru_cache(maxsize=1) -def get_gpu_info(): - """Get GPU model information. - - This function detects the GPU model and returns a sanitized string identifier. - It prioritizes FLA_GPU_NAME environment variable if set, then detects from - available hardware (CUDA, ROCm, Intel GPU, or CPU). - """ - # Check if GPU name is overridden via environment variable - gpu_name = None - # Check if GPU name is overridden via environment variable - if "FLA_GPU_NAME" in os.environ: - gpu_name = os.environ["FLA_GPU_NAME"] - # Try to get device name based on availability - elif torch.cuda.is_available(): - # Works for both NVIDIA and AMD GPUs (ROCm) - gpu_name = torch.cuda.get_device_name(0) - elif hasattr(torch, 'xpu') and torch.xpu.is_available(): - gpu_name = torch.xpu.get_device_name(0) - - if gpu_name: - return sanitize_gpu_name(gpu_name) - - # Default to CPU if no GPU available - return "cpu" - - -def get_fla_config_dir() -> Path: - """Get FLA's configs directory. - - The directory can be overridden by setting the FLA_CONFIG_DIR environment variable. - If set, configs will be loaded directly from $FLA_CONFIG_DIR/. Otherwise FLA - falls back to the default fla/configs/{GPU}/ directory in the project. - """ - # Check if custom config dir is set via environment variable - if "FLA_CONFIG_DIR" in os.environ: - return Path(os.environ["FLA_CONFIG_DIR"]) - - # Default: project_dir/fla/configs/{GPU}/ - project_dir = Path(__file__).parent.parent.parent - return project_dir / "configs" / get_gpu_info() - - -@dataclasses.dataclass(frozen=True) -class AutotuneKey: - """Autotune key with exact/fuzzy matching, serialization, and construction helpers.""" - autotune_key: tuple[Any, ...] - - @staticmethod - def normalize_autotune_key(value: Any) -> Any: - if isinstance(value, (list, tuple)): - return [AutotuneKey.normalize_autotune_key(v) for v in value] - if isinstance(value, dict): - return {k: AutotuneKey.normalize_autotune_key(v) for k, v in value.items()} - return value - - @staticmethod - def serialize(key: Any) -> str: - return json.dumps(AutotuneKey.normalize_autotune_key(key), separators=(",", ":"), sort_keys=True) - - @staticmethod - def key_hash(key: Any) -> str: - import hashlib - return hashlib.md5(AutotuneKey.serialize(key).encode()).hexdigest() - - @staticmethod - def is_numeric(value: Any) -> bool: - return isinstance(value, (int, float)) and not isinstance(value, bool) - - @staticmethod - def keys_fuzzy_match(cached_key: Any, requested_key: Any) -> bool: - # Fuzzy match: numeric leaves are compatible regardless of their actual numeric values - # (e.g. a config tuned for seq_len=1024 can apply to seq_len=2048). - # Structure (type, length, dict keys) must still match exactly. - if AutotuneKey.is_numeric(cached_key) and AutotuneKey.is_numeric(requested_key): - return True - if isinstance(cached_key, (list, tuple)) and isinstance(requested_key, (list, tuple)): - return len(cached_key) == len(requested_key) and all( - AutotuneKey.keys_fuzzy_match(c, r) for c, r in zip(cached_key, requested_key) - ) - if isinstance(cached_key, dict) and isinstance(requested_key, dict): - return cached_key.keys() == requested_key.keys() and all( - AutotuneKey.keys_fuzzy_match(cached_key[k], requested_key[k]) for k in cached_key - ) - return cached_key == requested_key - - @classmethod - def build( - cls, - arg_names: list[str], - key_names: list[str], - positional_args: tuple[Any, ...], - runtime_kwargs: dict[str, Any], - ) -> "AutotuneKey": - named_args = dict(zip(arg_names, positional_args)) - all_args = {**named_args, **runtime_kwargs} - tracked_args = {k: v for (k, v) in all_args.items() if k in arg_names} - tuning_key = [tracked_args[name] for name in key_names if name in tracked_args] - for arg in tracked_args.values(): - if hasattr(arg, "dtype"): - tuning_key.append(str(arg.dtype)) - return cls(autotune_key=tuple(tuning_key)) - - def exact_matches(self, entry_key: Any) -> bool: - return self.serialize(self.autotune_key) == self.serialize(entry_key) - - def fuzzy_matches(self, entry_key: Any) -> bool: - self_normalized = self.normalize_autotune_key(self.autotune_key) - entry_normalized = self.normalize_autotune_key(entry_key) - return ( - isinstance(self_normalized, list) - and isinstance(entry_normalized, list) - and len(self_normalized) == len(entry_normalized) - and AutotuneKey.keys_fuzzy_match(self_normalized, entry_normalized) - ) - - -@dataclasses.dataclass(frozen=True) -class KernelConfigFile: - """Validated in-memory representation of a {kernel_name}.json config file.""" - kernel_name: str | None - triton_version: str | None - autotune_entries: dict[str, dict[str, Any]] | None - default_config: dict[str, Any] | None - - @classmethod - def from_dict(cls, config_file: Path, data: Any) -> "KernelConfigFile | None": - """Parse and validate a raw JSON dict. Returns None (with a warning) if malformed.""" - def fail(msg, *args): - logger.warning(msg, *args) - raise ValueError - - try: - if not isinstance(data, dict): - fail("Malformed config %s: root is %s, expected dict", config_file, type(data).__name__) - raw_entries = data.get("autotune_entries") - entries: dict[str, dict[str, Any]] | None = None - if raw_entries is not None: - if not isinstance(raw_entries, dict): - fail("Malformed config %s: 'autotune_entries' is %s, expected dict", - config_file, type(raw_entries).__name__) - for h, entry in raw_entries.items(): - if not isinstance(entry, dict): - fail("Malformed config %s: autotune_entries[%r] is %s, expected dict", - config_file, h, type(entry).__name__) - if not isinstance(entry.get("config"), dict): - fail("Malformed config %s: autotune_entries[%r] missing valid 'config' field", config_file, h) - entries = raw_entries - default_config = data.get("default_config") - if default_config is not None and not isinstance(default_config, dict): - fail("Malformed config %s: 'default_config' is %s, expected dict", config_file, type(default_config).__name__) - return cls( - kernel_name=data.get("kernel_name"), - triton_version=data.get("triton_version"), - autotune_entries=entries, - default_config=default_config, - ) - except ValueError: - return None - - @classmethod - def from_file(cls, config_file: Path) -> "KernelConfigFile | None": - """Read and validate a config file. Returns None if the file is missing or malformed.""" - config_data = read_config_file(config_file) - if config_data is None: - return None - return cls.from_dict(config_file, config_data) - - def lookup_exact(self, key: AutotuneKey) -> dict[str, Any] | None: - if self.autotune_entries is None: - return None - return self.autotune_entries.get(AutotuneKey.key_hash(key.autotune_key)) - - def lookup_fuzzy(self, key: AutotuneKey) -> dict[str, Any] | None: - if self.autotune_entries is None: - return None - for entry in self.autotune_entries.values(): - if key.fuzzy_matches(entry.get("autotune_key")): - return entry - return None - - -@cache -def load_config_file(config_file: Path) -> dict[str, Any] | None: - try: - with open(config_file) as f: - return json.load(f) - except Exception as e: - logger.warning("Error reading config file %s: %s", config_file, e) - return None - - -def read_config_file(config_file: Path) -> dict[str, Any] | None: - """Read a config file, bypassing the in-process cache in ALWAYS mode.""" - if FLA_CACHE_MODE is FlaCacheMode.ALWAYS: - return load_config_file.__wrapped__(config_file) - return load_config_file(config_file) - - -def load_cached_config(kernel_name: str, autotune_key: AutotuneKey | None = None) -> dict[str, Any] | None: - """ - Load cached best config for a kernel from FLA configs directory. - - This function loads the cached best configuration for a given kernel name - from get_fla_config_dir()/{kernel_name}.json. - - Cache files may contain multiple autotune entries keyed by Triton's - runtime tuning key plus a top-level default config. - - If the config file is not found or cannot be loaded, a warning is printed - and None is returned, allowing fallback to Triton's autotune. - - The lookup mode is controlled by the FLA_CACHE_MODE environment variable (see FlaCacheMode). - - Args: - kernel_name: Name of the kernel (e.g., "causal_conv1d_fwd_kernel") - autotune_key: Triton autotune key for the current invocation - - Returns: - Best config dictionary or None if not found or disabled - """ - if FLA_CACHE_MODE is FlaCacheMode.DISABLED: - return None - - config_dir = get_fla_config_dir() - config_file = config_dir / f"{kernel_name}.json" - - if not config_file.exists(): - return None - - config_data = read_config_file(config_file) - if config_data is None: - return None - config = KernelConfigFile.from_dict(config_file, config_data) - if config is None: - return None - - if FLA_CACHE_MODE is FlaCacheMode.DEFAULT or FLA_CACHE_MODE is FlaCacheMode.ALWAYS: - return config.default_config - - # STRICT mode: exact match only, no fuzzy fallback - if FLA_CACHE_MODE is FlaCacheMode.STRICT: - if autotune_key is not None: - entry = config.lookup_exact(autotune_key) - if entry is not None: - return entry["config"] - return None - - # FULL and FUZZY modes: try exact key match first, then fuzzy match - if autotune_key is not None: - entry = config.lookup_exact(autotune_key) or config.lookup_fuzzy(autotune_key) - if entry is not None: - return entry["config"] - - if FLA_CACHE_MODE is FlaCacheMode.FUZZY: - return None - - # FULL mode: fall back to default_config, then legacy raw config (no autotune_entries) - if config.default_config is not None: - return config.default_config - if config.autotune_entries is not None: - return None - return config_data - - -class CachedAutotuner(Autotuner): - """ - A modified autotuner that loads best config from FLA's config directory. - - This class extends Triton's Autotuner but overrides the run method to - try loading cached configuration first before falling back to autotune. - """ - - def __init__(self, fn, arg_names, configs, key, reset_to_zero, restore_value, **kwargs): - super().__init__(fn, arg_names, configs, key, reset_to_zero, restore_value, **kwargs) - self.kernel_name = fn.fn.__name__ if hasattr(fn, 'fn') else fn.__name__ - - # None-safe pre/post hooks: Triton's defaults crash when a restore_value / reset_to_zero arg - # is None (idiomatic for optional pointers gated by a tl.constexpr flag). - # Fixed upstream in triton-lang/triton#10295 — remove this override once FLA's minimum Triton version has it. - if not self.user_defined_pre_hook and (self.reset_to_zero or self.restore_value): - def _pre_hook(kw, reset_only=False): - for n in self.reset_to_zero: - if kw[n] is not None: - kw[n].zero_() - if not reset_only: - self.restore_copies = {n: kw[n].clone() for n in self.restore_value if kw[n] is not None} - self.pre_hook = _pre_hook - if not self.user_defined_post_hook and self.restore_value: - def _post_hook(kw, exception): - for n, copy in self.restore_copies.items(): - kw[n].copy_(copy) - self.restore_copies = {} - self.post_hook = _post_hook - - def should_check_fla_cache(self, key: AutotuneKey) -> bool: - if FLA_CACHE_MODE is FlaCacheMode.DISABLED: - return False - if FLA_CACHE_MODE is FlaCacheMode.ALWAYS: - return True - return key.autotune_key not in self.cache - - def run(self, *args, **kwargs): - key = AutotuneKey.build(self.arg_names, self.keys, args, kwargs) - if self.should_check_fla_cache(key): - self.maybe_load_cached_config(key) - return super().run(*args, **kwargs) - - def maybe_load_cached_config(self, key: AutotuneKey): - best_config = load_cached_config(self.kernel_name, key) - - if best_config is not None: - kw = best_config["kwargs"] - num_warps = best_config["num_warps"] - num_stages = best_config["num_stages"] - - extra = { - "num_ctas": best_config["num_ctas"], - "maxnreg": best_config.get("maxnreg"), - "pre_hook": None, - "ir_override": best_config.get("ir_override"), - } if TRITON_ABOVE_3_5_1 else {} - cfg = triton.Config(kw, num_warps=num_warps, num_stages=num_stages, **extra) - - self.cache[key.autotune_key] = cfg - else: - logger.debug( - "No cached config found for kernel %s and key %s; falling back to Triton autotune", - self.kernel_name, - list(key.autotune_key), - ) - - -def fla_cache_autotune(configs, key=None, prune_configs_by=None, reset_to_zero=None, restore_value=None, - pre_hook=None, post_hook=None, warmup=None, rep=None, use_cuda_graph=False, - do_bench=None, cache_results=False): - """ - Decorator for auto-tuning a :code:`triton.jit`'d function with FLA config support. - - Extends Triton's autotune to load best configurations from FLA's config directory - (default: fla/configs/{GPU}/, or FLA_CONFIG_DIR/ when overridden), keyed by kernel - name from {kernel_name}.json. Lookup behaviour is controlled by FLA_CACHE_MODE. - Falls back to normal Triton autotuning when no cached config is found. - """ - # key can be None when we want to use cache only (no fallback autotune) - if key is None: - key = [] - - def decorator(fn): - kwargs = {} - if TRITON_ABOVE_3_4_0: - kwargs = {"cache_results": cache_results} - - return CachedAutotuner(fn, fn.arg_names, configs, key, reset_to_zero, restore_value, - pre_hook=pre_hook, post_hook=post_hook, - prune_configs_by=prune_configs_by, warmup=warmup, rep=rep, - use_cuda_graph=use_cuda_graph, do_bench=do_bench, - **kwargs, - ) - - return decorator - - -def configure_fla_cache_autotune(): - triton.autotune = fla_cache_autotune - logger.info( - "configure_fla_cache_autotune() is enabling FLA fla_cache_autotune; " - "triton.autotune will be replaced with fla_cache_autotune." - ) - - -def restore_autotune_backend(): - from triton.runtime.autotuner import autotune as original_autotune - triton.autotune = original_autotune - logger.info( - "restore_autotune_backend() is restoring Triton's original autotune; " - "triton.autotune will be replaced with triton.runtime.autotuner.autotune." - ) diff --git a/build/torch-xpu/ops/utils/constant.py b/build/torch-xpu/ops/utils/constant.py deleted file mode 100644 index 3a6c837be3ee8bb204cf7331490decc1dfe2fe2c..0000000000000000000000000000000000000000 --- a/build/torch-xpu/ops/utils/constant.py +++ /dev/null @@ -1,10 +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 - -# Approximate value of 1/ln(2), used for log/exp base conversion -# Best FP32 approximation: 1.4426950216 (hex 0x3FB8AA3B) -RCP_LN2 = 1.4426950216 diff --git a/build/torch-xpu/ops/utils/csr.py b/build/torch-xpu/ops/utils/csr.py deleted file mode 100644 index dab4cbab96c252ca1700ea4ded462c1298ffbdb1..0000000000000000000000000000000000000000 --- a/build/torch-xpu/ops/utils/csr.py +++ /dev/null @@ -1,170 +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 -import triton.language as tl - -from ...ops.utils.index import prepare_chunk_offsets - - -@triton.heuristics({ - 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None, - 'USE_BLOCK_COUNTS': lambda args: isinstance(args['block_counts'], torch.Tensor), -}) -@triton.jit(do_not_specialize=['T', 'N']) -def prepare_block_csr_kernel( - block_indices, - block_counts, - cu_seqlens, - chunk_offsets, - cursor, - csr_indices, - csr_offsets, - N, - T, - H: tl.constexpr, - S: tl.constexpr, - BT: tl.constexpr, - BS: tl.constexpr, - TC: tl.constexpr, - COUNT_ONLY: tl.constexpr, - USE_BLOCK_COUNTS: tl.constexpr, - IS_VARLEN: tl.constexpr, -): - i_t, i_bh = tl.program_id(0), tl.program_id(1) - i_b, i_h = i_bh // H, i_bh % H - o_t = i_t * BT + tl.arange(0, BT) - o_s = tl.arange(0, S) - m_t = o_t < T - # [BT] flattened (batch, query, kv-head) index; int64 to keep address arithmetic safe at large T - i_qh = ((i_b * T).to(tl.int64) + o_t) * H + i_h - - # [BT, S] selected blocks, masked to each query's valid causal in-range slots - b_i = tl.load(block_indices + i_qh[:, None] * S + o_s[None, :], mask=m_t[:, None], other=-1).to(tl.int64) - if USE_BLOCK_COUNTS: - b_m = m_t[:, None] & (o_s[None, :] < tl.load(block_counts + i_qh, mask=m_t, other=0)[:, None]) - else: - b_m = m_t[:, None] & (o_s[None, :] < block_counts) - b_m = b_m & (b_i >= 0) & (b_i < TC) & (b_i * BS <= o_t[:, None]) - - if IS_VARLEN: - # vectorized binary search for the sequence holding each query (32 steps cover any num_seq) - lo, hi = tl.zeros([BT], dtype=tl.int32), tl.full([BT], N, dtype=tl.int32) - for _ in range(32): - mid = (lo + hi) // 2 - go = tl.load(cu_seqlens + mid + 1, mask=m_t, other=0) <= o_t - lo, hi = tl.where(go, mid + 1, lo), tl.where(go, hi, mid) - block_base = tl.load(chunk_offsets + lo, mask=m_t, other=0).to(tl.int64) - block_id = (block_base[:, None] + b_i) * H + i_h - else: - block_id = (i_b * H + i_h) * TC + b_i - - if COUNT_ONLY: - tl.atomic_add(csr_offsets + block_id + 1, 1, mask=b_m) - else: - dst = tl.load(csr_offsets + block_id, mask=b_m, other=0).to(tl.int64) + tl.atomic_add(cursor + block_id, 1, mask=b_m) - b_q = tl.broadcast_to((i_b * T + o_t)[:, None], (BT, S)) - tl.store(csr_indices + dst, b_q.to(csr_indices.dtype.element_ty), mask=b_m) - - -def prepare_block_csr( - block_indices: torch.LongTensor, - block_counts: torch.LongTensor | int, - cu_seqlens: torch.LongTensor | None, - chunk_indices: torch.LongTensor | None, - num_blocks: int, - block_size: int, -) -> tuple[torch.Tensor, torch.Tensor]: - r""" - Invert a per-query block selection into CSR (compressed sparse row) form. - - `block_indices[b, t, h, :]` lists the blocks query `t` (kv-head `h`) selects. - The inverse maps each block to the queries that selected it, - which a block-parallel backward (e.g. NSA `bwd_dkv`) needs. - The result is CSR over a `[block, query]` matrix: `csr_indices` holds the selecting query positions grouped by block, - and `csr_offsets` holds the per-block row offsets, - so block `i` owns `csr_indices[csr_offsets[i]:csr_offsets[i + 1]]`. - - Block ids follow the launching kernel's program-id layout: - `(b * H + h) * num_blocks + s` for dense, `(global_block + s) * H + h` for varlen, - where the varlen block base comes from an in-kernel binary search over `cu_seqlens` into `chunk_offsets`. - Short problems use a counting-then-scatter sort (`csr_indices` over-allocated to its upper bound, so its length - need not be read back to the host); long ones bucket the pairs with a radix sort. Both yield the same CSR. - - Example (dense, B = H = 1, block_size = 1 so block b covers token b; causal needs b <= t): - - # input: each query lists the blocks it selects, -1 is padding - block_indices[0, :, 0, :] = - [[ 0, -1], # query 0 selects block 0 - [ 0, 1], # query 1 selects blocks 0, 1 - [ 1, 2], # query 2 selects blocks 1, 2 - [ 0, 3]] # query 3 selects blocks 0, 3 - - # invert -> which queries selected each block: - # block 0: queries 0, 1, 3 - # block 1: queries 1, 2 - # block 2: query 2 - # block 3: query 3 - - csr_indices = [0, 1, 3, 1, 2, 2, 3] # 7 selections, grouped by block (order within a block is arbitrary) - csr_offsets = [0, 3, 5, 6, 7] # block i's queries = csr_indices[csr_offsets[i]:csr_offsets[i+1]] - - Args: - block_indices (torch.LongTensor): - Selected block ids of shape `[B, T, H, S]`, padded with `-1`. - block_counts (torch.LongTensor or int): - Number of valid slots per query, a `[B, T, H]` tensor or an int. - cu_seqlens (torch.LongTensor, Optional): - Cumulative sequence lengths for variable-length packing. Default: `None` (dense). - chunk_indices (torch.LongTensor): - Per-chunk `(sequence, local-block)` index pairs; read only to size the varlen block-id space. - num_blocks (int): - Number of blocks per `(batch, head)`, i.e. the dense kernel's `TC`. - block_size (int): - Selected block size, used for the causal check and varlen block ids. - - Returns: - csr_indices (torch.Tensor): - `int32` selecting query positions, grouped by block; absolute (`b * T + t`). - csr_offsets (torch.Tensor): - `int32` CSR row offsets of shape `[NB + 1]`, one per block plus a final end offset. - """ - B, T, H, S = block_indices.shape - N = 0 if cu_seqlens is None else cu_seqlens.numel() - 1 - NB = B * H * num_blocks if cu_seqlens is None else chunk_indices.shape[0] * H - chunk_offsets = prepare_chunk_offsets(cu_seqlens, block_size) if cu_seqlens is not None else None - - cursor = block_indices.new_zeros(NB, dtype=torch.int32) - csr_offsets = block_indices.new_zeros(NB + 1, dtype=torch.int32) - csr_indices = block_indices.new_empty(B * T * H * S, dtype=torch.int32) - - BT = max(1, min(128, triton.next_power_of_2(max(1, 2048 // S)))) - grid = (triton.cdiv(T, BT), B * H) - - # counting sort: tally per-block counts, prefix-sum them into start offsets, then scatter. - # the two kernel passes can't merge -- the scatter position needs csr_offsets ready (global prefix sum). - kwargs = dict( - block_indices=block_indices, - block_counts=block_counts, - cu_seqlens=cu_seqlens, - chunk_offsets=chunk_offsets, - cursor=cursor, - csr_indices=csr_indices, - csr_offsets=csr_offsets, - N=N, - T=T, - H=H, - S=S, - BT=BT, - BS=block_size, - TC=num_blocks, - ) - prepare_block_csr_kernel[grid](**kwargs, COUNT_ONLY=True) - csr_offsets.cumsum_(0) - prepare_block_csr_kernel[grid](**kwargs, COUNT_ONLY=False) - return csr_indices, csr_offsets diff --git a/build/torch-xpu/ops/utils/cumsum.py b/build/torch-xpu/ops/utils/cumsum.py deleted file mode 100644 index 6140e860c8e1b17951b7a1df47cf9a9b93af4b25..0000000000000000000000000000000000000000 --- a/build/torch-xpu/ops/utils/cumsum.py +++ /dev/null @@ -1,476 +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 -import triton.language as tl - -from ...ops.backends import dispatch -from ...ops.utils.cache import fla_cache_autotune -from ...ops.utils.index import prepare_chunk_indices -from ...utils import autotune_cache_kwargs, check_shared_mem, input_guard - -BS_LIST = [32, 64] if check_shared_mem() else [16, 32] - - -@triton.heuristics({ - 'HAS_SCALE': lambda args: args['scale'] is not None, - 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None, -}) -@fla_cache_autotune( - configs=[ - triton.Config({}, num_warps=num_warps) - for num_warps in [1, 2, 4, 8] - ], - key=['B', 'H', 'BT', 'IS_VARLEN', 'REVERSE'], - **autotune_cache_kwargs, -) -@triton.jit(do_not_specialize=['T']) -def chunk_local_cumsum_scalar_kernel( - s, - o, - scale, - cu_seqlens, - chunk_indices, - T, - B: tl.constexpr, - H: tl.constexpr, - BT: tl.constexpr, - REVERSE: tl.constexpr, - HAS_SCALE: tl.constexpr, - IS_VARLEN: tl.constexpr, - HEAD_FIRST: tl.constexpr, -): - i_t, i_bh = tl.program_id(0), tl.program_id(1) - i_b, i_h = i_bh // H, i_bh % H - 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.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) - T = eos - bos - else: - bos, eos = i_b * T, i_b * T + T - - if HEAD_FIRST: - p_s = tl.make_block_ptr(s + bos*H + i_h*T, (T,), (1,), (i_t * BT,), (BT,), (0,)) - p_o = tl.make_block_ptr(o + bos*H + i_h*T, (T,), (1,), (i_t * BT,), (BT,), (0,)) - else: - p_s = tl.make_block_ptr(s + bos*H + i_h, (T,), (H,), (i_t * BT,), (BT,), (0,)) - p_o = tl.make_block_ptr(o + bos*H + i_h, (T,), (H,), (i_t * BT,), (BT,), (0,)) - # [BT] - b_s = tl.load(p_s, boundary_check=(0,)).to(tl.float32) - b_o = tl.cumsum(b_s, axis=0) - if REVERSE: - b_z = tl.sum(b_s, axis=0) - b_o = -b_o + b_z[None] + b_s - if HAS_SCALE: - b_o *= scale - tl.store(p_o, b_o.to(p_o.dtype.element_ty), boundary_check=(0,)) - - -@triton.heuristics({ - 'HAS_SCALE': lambda args: args['scale'] is not None, - 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None, -}) -@fla_cache_autotune( - configs=[ - triton.Config({'BS': BS}, num_warps=num_warps) - for BS in BS_LIST - for num_warps in [2, 4, 8] - ], - key=['B', 'H', 'S', 'BT', 'IS_VARLEN', 'REVERSE'], - **autotune_cache_kwargs, -) -@triton.jit(do_not_specialize=['T']) -def chunk_local_cumsum_vector_kernel( - s, - o, - scale, - cu_seqlens, - chunk_indices, - T, - B: tl.constexpr, - H: tl.constexpr, - S: tl.constexpr, - BT: tl.constexpr, - BS: tl.constexpr, - REVERSE: tl.constexpr, - HAS_SCALE: tl.constexpr, - IS_VARLEN: tl.constexpr, - HEAD_FIRST: tl.constexpr, -): - i_s, i_t, i_bh = tl.program_id(0), tl.program_id(1), tl.program_id(2) - i_b, i_h = i_bh // H, i_bh % H - 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.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) - T = eos - bos - else: - bos, eos = i_b * T, i_b * T + T - - if HEAD_FIRST: - p_s = tl.make_block_ptr(s + (bos * H + i_h*T)*S, (T, S), (S, 1), (i_t * BT, i_s * BS), (BT, BS), (1, 0)) - p_o = tl.make_block_ptr(o + (bos * H + i_h*T)*S, (T, S), (S, 1), (i_t * BT, i_s * BS), (BT, BS), (1, 0)) - else: - p_s = tl.make_block_ptr(s + (bos * H + i_h) * S, (T, S), (H*S, 1), (i_t * BT, i_s * BS), (BT, BS), (1, 0)) - p_o = tl.make_block_ptr(o + (bos * H + i_h) * S, (T, S), (H*S, 1), (i_t * BT, i_s * BS), (BT, BS), (1, 0)) - # [BT, BS] - b_s = tl.load(p_s, boundary_check=(0, 1)).to(tl.float32) - if REVERSE: - b_o = tl.cumsum(b_s, axis=0, reverse=True) - else: - b_o = tl.cumsum(b_s, axis=0) - if HAS_SCALE: - b_o *= scale - tl.store(p_o, b_o.to(p_o.dtype.element_ty), boundary_check=(0, 1)) - - -@triton.heuristics({ - 'HAS_SCALE': lambda args: args['scale'] is not None, - 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None, -}) -@triton.autotune( - configs=[ - triton.Config({'BT': BT}, num_warps=num_warps, num_stages=num_stages) - for BT in [32, 64, 128, 256] - for num_warps in [2, 4, 8] - for num_stages in [1, 2, 3, 4] - ], - key=['B', 'H', 'IS_VARLEN', 'REVERSE'], - **autotune_cache_kwargs, -) -@triton.jit(do_not_specialize=['T']) -def chunk_global_cumsum_scalar_kernel( - s, - o, - scale, - cu_seqlens, - T, - B: tl.constexpr, - H: tl.constexpr, - BT: tl.constexpr, - REVERSE: tl.constexpr, - HAS_SCALE: tl.constexpr, - IS_VARLEN: tl.constexpr, - HEAD_FIRST: tl.constexpr, -): - i_nh = tl.program_id(0) - i_n, i_h = i_nh // H, i_nh % H - if IS_VARLEN: - bos, eos = tl.load(cu_seqlens + i_n).to(tl.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) - else: - bos, eos = i_n * T, i_n * T + T - T = eos - bos - - b_z = tl.zeros([], dtype=tl.float32) - NT = tl.cdiv(T, BT) - for i_c in range(NT): - i_t = NT - 1 - i_c if REVERSE else i_c - if HEAD_FIRST: - p_s = tl.make_block_ptr(s + bos*H + i_h*T, (T,), (1,), (i_t * BT,), (BT,), (0,)) - p_o = tl.make_block_ptr(o + bos*H + i_h*T, (T,), (1,), (i_t * BT,), (BT,), (0,)) - else: - p_s = tl.make_block_ptr(s + bos*H + i_h, (T,), (H,), (i_t * BT,), (BT,), (0,)) - p_o = tl.make_block_ptr(o + bos*H + i_h, (T,), (H,), (i_t * BT,), (BT,), (0,)) - b_s = tl.load(p_s, boundary_check=(0,)).to(tl.float32) - b_o = tl.cumsum(b_s, axis=0) - b_ss = tl.sum(b_s, 0) - if REVERSE: - b_o = -b_o + b_ss + b_s - b_o += b_z - if i_c >= 0: - b_z += b_ss - if HAS_SCALE: - b_o *= scale - tl.store(p_o, b_o.to(p_o.dtype.element_ty), boundary_check=(0,)) - - -@triton.heuristics({ - 'HAS_SCALE': lambda args: args['scale'] is not None, - 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None, -}) -@triton.autotune( - configs=[ - triton.Config({'BT': BT}, num_warps=num_warps, num_stages=num_stages) - for BT in [16, 32, 64, 128] - for num_warps in [2, 4, 8] - for num_stages in [1, 2, 3, 4] - ], - key=['B', 'H', 'S', 'IS_VARLEN', 'REVERSE'], - **autotune_cache_kwargs, -) -@triton.jit(do_not_specialize=['T']) -def chunk_global_cumsum_vector_kernel( - s, - o, - scale, - cu_seqlens, - T, - B: tl.constexpr, - H: tl.constexpr, - S: tl.constexpr, - BT: tl.constexpr, - BS: tl.constexpr, - REVERSE: tl.constexpr, - HAS_SCALE: tl.constexpr, - IS_VARLEN: tl.constexpr, - HEAD_FIRST: tl.constexpr, -): - i_s, i_nh = tl.program_id(0), tl.program_id(1) - i_n, i_h = i_nh // H, i_nh % H - if IS_VARLEN: - bos, eos = tl.load(cu_seqlens + i_n).to(tl.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) - else: - bos, eos = i_n * T, i_n * T + T - T = eos - bos - - b_z = tl.zeros([BS], dtype=tl.float32) - NT = tl.cdiv(T, BT) - for i_c in range(NT): - i_t = NT - 1 - i_c if REVERSE else i_c - if HEAD_FIRST: - p_s = tl.make_block_ptr(s + (bos * H + i_h*T)*S, (T, S), (S, 1), (i_t * BT, i_s * BS), (BT, BS), (1, 0)) - p_o = tl.make_block_ptr(o + (bos * H + i_h*T)*S, (T, S), (S, 1), (i_t * BT, i_s * BS), (BT, BS), (1, 0)) - else: - p_s = tl.make_block_ptr(s + (bos * H + i_h) * S, (T, S), (H*S, 1), (i_t * BT, i_s * BS), (BT, BS), (1, 0)) - p_o = tl.make_block_ptr(o + (bos * H + i_h) * S, (T, S), (H*S, 1), (i_t * BT, i_s * BS), (BT, BS), (1, 0)) - # [BT, BS] - b_s = tl.load(p_s, boundary_check=(0, 1)).to(tl.float32) - if REVERSE: - b_c = b_z[None, :] + tl.cumsum(b_s, axis=0, reverse=True) - else: - b_c = b_z[None, :] + tl.cumsum(b_s, axis=0) - if HAS_SCALE: - b_c *= scale - tl.store(p_o, b_c.to(p_o.dtype.element_ty), boundary_check=(0, 1)) - b_z += tl.sum(b_s, 0) - - -def chunk_local_cumsum_scalar( - g: torch.Tensor, - chunk_size: int, - reverse: bool = False, - scale: float = None, - cu_seqlens: torch.Tensor | None = None, - head_first: bool = False, - output_dtype: torch.dtype | None = torch.float, - chunk_indices: torch.LongTensor | None = None, -) -> torch.Tensor: - if head_first: - B, H, T = g.shape - else: - B, T, H = g.shape - assert chunk_size == 2**(chunk_size.bit_length()-1), "chunk_size must be a power of 2" - BT = chunk_size - if chunk_indices is None and cu_seqlens is not None: - chunk_indices = prepare_chunk_indices(cu_seqlens, BT) - NT = triton.cdiv(T, BT) if cu_seqlens is None else len(chunk_indices) - g_org, g = g, torch.empty_like(g, dtype=output_dtype or g.dtype) - grid = (NT, B * H) - chunk_local_cumsum_scalar_kernel[grid]( - s=g_org, - o=g, - scale=scale, - cu_seqlens=cu_seqlens, - chunk_indices=chunk_indices, - T=T, - B=B, - H=H, - BT=BT, - HEAD_FIRST=head_first, - REVERSE=reverse, - ) - return g - - -def chunk_local_cumsum_vector( - g: torch.Tensor, - chunk_size: int, - reverse: bool = False, - scale: float = None, - cu_seqlens: torch.Tensor | None = None, - head_first: bool = False, - output_dtype: torch.dtype | None = torch.float, - chunk_indices: torch.LongTensor | None = None, -) -> torch.Tensor: - if head_first: - B, H, T, S = g.shape - else: - B, T, H, S = g.shape - BT = chunk_size - if chunk_indices is None and cu_seqlens is not None: - chunk_indices = prepare_chunk_indices(cu_seqlens, BT) - NT = triton.cdiv(T, BT) if cu_seqlens is None else len(chunk_indices) - assert chunk_size == 2**(chunk_size.bit_length()-1), "chunk_size must be a power of 2" - - g_org, g = g, torch.empty_like(g, dtype=output_dtype or g.dtype) - def grid(meta): return (triton.cdiv(meta['S'], meta['BS']), NT, B * H) - # keep cummulative normalizer in fp32 - # this kernel is equivalent to - # g = g.view(B, H, NT, BT, -1).cumsum(-2).view(B, H, T, -1) - chunk_local_cumsum_vector_kernel[grid]( - s=g_org, - o=g, - scale=scale, - cu_seqlens=cu_seqlens, - chunk_indices=chunk_indices, - T=T, - B=B, - H=H, - S=S, - BT=BT, - HEAD_FIRST=head_first, - REVERSE=reverse, - ) - return g - - -@input_guard -def chunk_global_cumsum_scalar( - s: torch.Tensor, - reverse: bool = False, - cu_seqlens: torch.Tensor | None = None, - scale: float = None, - head_first: bool = False, - output_dtype: torch.dtype | None = torch.float, -) -> torch.Tensor: - if head_first: - B, H, T = s.shape - else: - B, T, H = s.shape - N = len(cu_seqlens) - 1 if cu_seqlens is not None else B - - z = torch.empty_like(s, dtype=output_dtype or s.dtype) - grid = (N * H,) - chunk_global_cumsum_scalar_kernel[grid]( - s=s, - o=z, - scale=scale, - cu_seqlens=cu_seqlens, - T=T, - B=B, - H=H, - HEAD_FIRST=head_first, - REVERSE=reverse, - ) - return z - - -@input_guard -def chunk_global_cumsum_vector( - s: torch.Tensor, - reverse: bool = False, - cu_seqlens: torch.Tensor | None = None, - scale: float = None, - head_first: bool = False, - output_dtype: torch.dtype | None = torch.float, -) -> torch.Tensor: - if head_first: - B, H, T, S = s.shape - else: - B, T, H, S = s.shape - N = len(cu_seqlens) - 1 if cu_seqlens is not None else B - BS = min(32, triton.next_power_of_2(S)) - - z = torch.empty_like(s, dtype=output_dtype or s.dtype) - grid = (triton.cdiv(S, BS), N * H) - chunk_global_cumsum_vector_kernel[grid]( - s=s, - o=z, - scale=scale, - cu_seqlens=cu_seqlens, - T=T, - B=B, - H=H, - S=S, - BS=BS, - HEAD_FIRST=head_first, - REVERSE=reverse, - ) - return z - - -@input_guard -@dispatch('utils') -def chunk_global_cumsum( - s: torch.Tensor, - reverse: bool = False, - cu_seqlens: torch.Tensor | None = None, - scale: float = None, - head_first: bool = False, - output_dtype: torch.dtype | None = torch.float, -) -> torch.Tensor: - if cu_seqlens is not None: - assert s.shape[0] == 1, "Only batch size 1 is supported when cu_seqlens are provided" - if len(s.shape) == 3: - return chunk_global_cumsum_scalar( - s=s, - reverse=reverse, - cu_seqlens=cu_seqlens, - scale=scale, - head_first=head_first, - output_dtype=output_dtype, - ) - elif len(s.shape) == 4: - return chunk_global_cumsum_vector( - s=s, - reverse=reverse, - cu_seqlens=cu_seqlens, - scale=scale, - head_first=head_first, - output_dtype=output_dtype, - ) - else: - raise ValueError( - f"Unsupported input shape {s.shape}, " - f"which should be [B, T, H]/[B, T, H, D] if `head_first=False` " - f"or [B, H, T]/[B, H, T, D] otherwise", - ) - - -@input_guard -@dispatch('utils') -def chunk_local_cumsum( - g: torch.Tensor, - chunk_size: int, - reverse: bool = False, - scale: float = None, - cu_seqlens: torch.Tensor | None = None, - head_first: bool = False, - output_dtype: torch.dtype | None = torch.float, - chunk_indices: torch.LongTensor | None = None, - **kwargs, -) -> torch.Tensor: - if cu_seqlens is not None: - assert g.shape[0] == 1, "Only batch size 1 is supported when cu_seqlens are provided" - if len(g.shape) == 3: - return chunk_local_cumsum_scalar( - g=g, - chunk_size=chunk_size, - reverse=reverse, - scale=scale, - cu_seqlens=cu_seqlens, - head_first=head_first, - output_dtype=output_dtype, - chunk_indices=chunk_indices, - ) - elif len(g.shape) == 4: - return chunk_local_cumsum_vector( - g=g, - chunk_size=chunk_size, - reverse=reverse, - scale=scale, - cu_seqlens=cu_seqlens, - head_first=head_first, - output_dtype=output_dtype, - chunk_indices=chunk_indices, - ) - else: - raise ValueError( - f"Unsupported input shape {g.shape}, " - f"which should be (B, T, H, D) if `head_first=False` " - f"or (B, H, T, D) otherwise", - ) diff --git a/build/torch-xpu/ops/utils/index.py b/build/torch-xpu/ops/utils/index.py deleted file mode 100644 index 61de8d4537652a62e6f8c2b7cfa6c28a0f9c69d2..0000000000000000000000000000000000000000 --- a/build/torch-xpu/ops/utils/index.py +++ /dev/null @@ -1,183 +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 torch.nn.functional as F -import triton -import triton.language as tl - -from ...utils import autotune_cache_kwargs, tensor_cache - - -@triton.autotune( - configs=[ - triton.Config({}, num_warps=num_warps) - for num_warps in [4, 8, 16, 32] - ], - key=['B'], - **autotune_cache_kwargs, -) -@triton.jit -def prepare_position_ids_kernel( - y, - cu_seqlens, - B: tl.constexpr, -): - i_n = tl.program_id(0) - bos, eos = tl.load(cu_seqlens + i_n).to(tl.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) - T = eos - bos - - o = tl.arange(0, B) - for i in range(0, tl.cdiv(T, B) * B, B): - o_i = o + i - tl.store(y + bos + o_i, o_i, o_i < T) - - -@tensor_cache -def prepare_lens(cu_seqlens: torch.LongTensor) -> torch.LongTensor: - return torch.diff(cu_seqlens) - - -@tensor_cache -def prepare_lens_from_mask(mask: torch.BoolTensor) -> torch.LongTensor: - return mask.sum(dim=-1, dtype=torch.int32) - - -@tensor_cache -def prepare_cu_seqlens_from_lens( - lens: torch.LongTensor, - dtype: torch.dtype | None = torch.int32, -) -> torch.LongTensor: - return F.pad(lens.cumsum(dim=0, dtype=dtype), (1, 0)) - - -@tensor_cache -def prepare_cu_seqlens_from_mask( - mask: torch.BoolTensor, - dtype: torch.dtype | None = torch.int32, -) -> torch.LongTensor: - return prepare_cu_seqlens_from_lens(prepare_lens_from_mask(mask), dtype) - - -@tensor_cache -def prepare_split_cu_seqlens( - batch_size: int | None = None, - seq_len: int | None = None, - split_size: int | None = None, - cu_seqlens: torch.LongTensor | None = None, - dtype: torch.dtype | None = torch.int32, - device: torch.device | None = torch.device('cpu'), -) -> torch.LongTensor: - """Sub-split a (optionally packed) batch along the token axis. - - Two calling modes: - - **Rectangular batch**: pass `batch_size` and `seq_len`, leave - `cu_seqlens=None`. Internally synthesizes `[0, L, 2L, ..., B*L]`. - - **Packed varlen**: pass `cu_seqlens`. `batch_size` and `seq_len` are - ignored (kept as optional kwargs for backward-compat with callers - that used to pass dummies). - - `split_size` is always required. - - The legacy positional signature `(batch_size, seq_len, split_size, ...)` - continues to work — the first two args retain their position but may now - be omitted when `cu_seqlens` is supplied. - """ - if split_size is None: - raise TypeError("prepare_split_cu_seqlens() requires `split_size`") - if cu_seqlens is None: - if batch_size is None or seq_len is None: - raise TypeError( - "prepare_split_cu_seqlens() requires either `cu_seqlens`, " - "or both `batch_size` and `seq_len`" - ) - total_tokens = batch_size * seq_len - cu_seqlens = list(range(0, total_tokens, seq_len)) + [total_tokens] - else: - cu_seqlens = cu_seqlens.tolist() - return torch.tensor( - [ - i - for bos, eos in zip(cu_seqlens[:-1], cu_seqlens[1:], strict=False) - for i in range(bos, eos, split_size) - ] + [cu_seqlens[-1]], - dtype=dtype, - device=device, - ) - - -def _segmented_arange(counts: torch.LongTensor) -> tuple[torch.LongTensor, torch.LongTensor]: - """Expand per-segment counts into flat per-slot index tensors. - - Given segment sizes ``counts = [c0, c1, ...]``, return two 1-D tensors of - length ``counts.sum()`` that together label every slot with its segment and - its position within that segment. - - Example -- ``counts = [2, 3]`` (segment 0 spans 2 slots, segment 1 spans 3):: - - seg_id = [0, 0, 1, 1, 1] # which segment each slot belongs to - intra_idx = [0, 1, 0, 1, 2] # running index within that segment - - With CUDA ``counts``, ``repeat_interleave`` reads ``counts.sum()`` on the - host (one device sync). Pass host-side counts to avoid it. - """ - seg_id = torch.repeat_interleave( - torch.arange(counts.numel(), device=counts.device, dtype=counts.dtype), - counts, - ) - seg_start = F.pad(counts.cumsum(0), (1, 0))[:-1] - intra_idx = torch.arange(seg_id.shape[0], device=counts.device, dtype=counts.dtype) - seg_start[seg_id] - return seg_id, intra_idx - - -@tensor_cache -def prepare_position_ids(cu_seqlens: torch.LongTensor, cu_seqlens_cpu: torch.LongTensor | None = None) -> torch.LongTensor: - src = cu_seqlens_cpu if cu_seqlens_cpu is not None else cu_seqlens - _, position_ids = _segmented_arange(prepare_lens(src)) - return position_ids.to(cu_seqlens) - - -@tensor_cache -def prepare_sequence_ids(cu_seqlens: torch.LongTensor, cu_seqlens_cpu: torch.LongTensor | None = None) -> torch.LongTensor: - return prepare_position_ids(cu_seqlens, cu_seqlens_cpu).eq(0).cumsum(0) - 1 - - -@tensor_cache -def prepare_token_indices(cu_seqlens: torch.LongTensor, cu_seqlens_cpu: torch.LongTensor | None = None) -> torch.LongTensor: - position_ids = prepare_position_ids(cu_seqlens, cu_seqlens_cpu) - return torch.stack([prepare_sequence_ids(cu_seqlens, cu_seqlens_cpu), position_ids], 1).to(cu_seqlens) - - -@tensor_cache -def prepare_chunk_indices( - cu_seqlens: torch.LongTensor, - chunk_size: int, - cu_seqlens_cpu: torch.LongTensor | None = None, -) -> torch.LongTensor: - src = cu_seqlens_cpu if cu_seqlens_cpu is not None else cu_seqlens - chunk_counts = (prepare_lens(src) + (chunk_size - 1)).div(chunk_size, rounding_mode='floor') - seg_id, intra_chunk_idx = _segmented_arange(chunk_counts) - return torch.stack([seg_id, intra_chunk_idx], 1).to(cu_seqlens) - - -@tensor_cache -def prepare_chunk_offsets( - cu_seqlens: torch.LongTensor, - chunk_size: int, -) -> torch.LongTensor: - return F.pad(triton.cdiv(prepare_lens(cu_seqlens), chunk_size), (1, 0), value=0).cumsum(-1) - - -@tensor_cache -def get_max_num_splits( - cu_seqlens: torch.LongTensor, - chunk_size: int, - cu_seqlens_cpu: torch.LongTensor | None = None -) -> int: - if cu_seqlens_cpu is not None: - return triton.cdiv(int(max(prepare_lens(cu_seqlens_cpu))), chunk_size) - return triton.cdiv(int(max(prepare_lens(cu_seqlens))), chunk_size) diff --git a/build/torch-xpu/ops/utils/logcumsumexp.py b/build/torch-xpu/ops/utils/logcumsumexp.py deleted file mode 100644 index ee6773f73d9af1a2ded3dfe4294e89bbf121925a..0000000000000000000000000000000000000000 --- a/build/torch-xpu/ops/utils/logcumsumexp.py +++ /dev/null @@ -1,58 +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 triton -import triton.language as tl - -from ...ops.utils.op import exp, log -from ...utils import autotune_cache_kwargs - - -@triton.autotune( - configs=[ - triton.Config({'BT': BT}, num_warps=num_warps) - for BT in [16, 32, 64] - for num_warps in [2, 4, 8] - ], - key=['S'], - **autotune_cache_kwargs, -) -@triton.jit(do_not_specialize=['T']) -def logcumsumexp_fwd_kernel( - s, - z, - T, - S: tl.constexpr, - BT: tl.constexpr, -): - i_bh = tl.program_id(0) - o_i = tl.arange(0, BT) - m_s = tl.where(o_i[:, None] >= o_i[None, :], 1., 0.) - - b_mp = tl.full([S], float('-inf'), dtype=tl.float32) - b_zp = tl.zeros([S], dtype=tl.float32) - for i_t in range(tl.cdiv(T, BT)): - p_s = tl.make_block_ptr(s + i_bh * T*S, (T, S), (S, 1), (i_t * BT, 0), (BT, S), (1, 0)) - p_z = tl.make_block_ptr(z + i_bh * T*S, (T, S), (S, 1), (i_t * BT, 0), (BT, S), (1, 0)) - - # [BT, S] - b_s = tl.load(p_s, boundary_check=(0, 1)).to(tl.float32) - # [S,] - b_mc = tl.max(b_s, 0) - b_mc = tl.maximum(b_mp, b_mc) - b_zp = b_zp * exp(b_mp - b_mc) - # [BT, S] - b_s = exp(b_s - b_mc) - b_z = tl.dot(m_s, b_s, allow_tf32=False) + b_zp - # [S,] - b_zc = tl.max(b_z, 0) - b_mp = b_mc - b_zp = b_zc - # [BT, BS] - # small eps to prevent underflows - b_z = log(tl.where(b_z != 0, b_z, 1e-20)) + b_mc - tl.store(p_z, b_z.to(p_z.dtype.element_ty), boundary_check=(0, 1)) diff --git a/build/torch-xpu/ops/utils/logsumexp.py b/build/torch-xpu/ops/utils/logsumexp.py deleted file mode 100644 index 0a9e0a5f53e4cde9a5d1ffe29709ab56f2aae200..0000000000000000000000000000000000000000 --- a/build/torch-xpu/ops/utils/logsumexp.py +++ /dev/null @@ -1,84 +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 -import triton.language as tl - -from ...ops.utils.op import exp, log -from ...utils import autotune_cache_kwargs - - -@triton.heuristics({ - 'HAS_SCALE': lambda args: args['scale'] is not None, -}) -@triton.autotune( - configs=[ - triton.Config({}, num_warps=num_warps) - for num_warps in [1, 2, 4, 8, 16, 32] - ], - key=['D'], - **autotune_cache_kwargs, -) -@triton.jit -def logsumexp_fwd_kernel( - x, - z, - scale, - D: tl.constexpr, - B: tl.constexpr, - HAS_SCALE: 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 - 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) - - -def logsumexp_fwd( - x, - scale: float | None = None, - dtype: torch.dtype | None = None, -): - r""" - Compute the logsumexp of the input tensor over the last dimension. - - Args: - x (Tensor): - The input tensor of any shape. - scale (Optional[float]): - The scale applied to the input tensor. Default: `None`. - dtype (Optional[torch.dtype]): - The data type of the output tensor. Default: `None`. - Returns: - Tensor: The logsumexp of the input tensor. - """ - - 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, - 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 diff --git a/build/torch-xpu/ops/utils/matmul.py b/build/torch-xpu/ops/utils/matmul.py deleted file mode 100644 index cd86847a5ffcbfbb7f779ae212e7ad79da0ab136..0000000000000000000000000000000000000000 --- a/build/torch-xpu/ops/utils/matmul.py +++ /dev/null @@ -1,249 +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 - -# code adapted from -# https://triton-lang.org/main/getting-started/tutorials/03-matrix-multiplication.html - - -import torch -import triton -import triton.language as tl - -from ...ops.utils.op import exp -from ...utils import autotune_cache_kwargs, input_guard - - -# `triton.jit`'ed functions can be auto-tuned by using the `triton.autotune` decorator, which consumes: -# - A list of `triton.Config` objects that define different configurations of -# meta-parameters (e.g., `BM`) and compilation options (e.g., `num_warps`) to try -# - An auto-tuning *key* whose change in values will trigger evaluation of all the -# provided configs -@triton.heuristics({ - 'HAS_ALPHA': lambda args: args['alpha'] is not None, - 'HAS_BETA': lambda args: args['beta'] is not None, -}) -@triton.autotune( - configs=[ - triton.Config({'BM': 128, 'BK': 64, 'BN': 256, 'G': 4}, num_stages=3, num_warps=8), - triton.Config({'BM': 64, 'BK': 32, 'BN': 256, 'G': 4}, num_stages=4, num_warps=4), - triton.Config({'BM': 128, 'BK': 32, 'BN': 128, 'G': 4}, num_stages=4, num_warps=4), - triton.Config({'BM': 128, 'BK': 32, 'BN': 64, 'G': 4}, num_stages=4, num_warps=4), - triton.Config({'BM': 64, 'BK': 32, 'BN': 128, 'G': 4}, num_stages=4, num_warps=4), - triton.Config({'BM': 128, 'BK': 32, 'BN': 32, 'G': 4}, num_stages=4, num_warps=4), - triton.Config({'BM': 64, 'BK': 32, 'BN': 32, 'G': 4}, num_stages=5, num_warps=2), - triton.Config({'BM': 32, 'BK': 32, 'BN': 64, 'G': 4}, num_stages=5, num_warps=2), - # Good config for fp8 inputs. - # triton.Config({'BM': 128, 'BK': 128, 'BN': 256, 'G': 4}, num_stages=3, num_warps=8), - # triton.Config({'BM': 256, 'BK': 128, 'BN': 128, 'G': 4}, num_stages=3, num_warps=8), - # triton.Config({'BM': 256, 'BK': 128, 'BN': 64, 'G': 4}, num_stages=4, num_warps=4), - # triton.Config({'BM': 64, 'BK': 128, 'BN': 256, 'G': 4}, num_stages=4, num_warps=4), - # triton.Config({'BM': 128, 'BK': 128, 'BN': 128, 'G': 4}, num_stages=4, num_warps=4), - # triton.Config({'BM': 128, 'BK': 64, 'BN': 64, 'G': 4}, num_stages=4, num_warps=4), - # triton.Config({'BM': 64, 'BK': 64, 'BN': 128, 'G': 4}, num_stages=4, num_warps=4), - # triton.Config({'BM': 128, 'BK': 64, 'BN': 32, 'G': 4}, num_stages=4, num_warps=4) - ], - key=['M', 'N', 'K'], - **autotune_cache_kwargs, -) -@triton.jit -def matmul_kernel( - # Pointers to matrices - a, - b, - c, - input, - alpha, - beta, - # Matrix dimensions - M, - N, - K, - # The stride variables represent how much to increase the ptr by when moving by 1 - # element in a particular dimension. E.g. `s_am` is how much to increase `a` - # by to get the element one row down (A has M rows). - stride_ab, stride_am, stride_ak, # a: batch, M, K - stride_bk, stride_bn, # b: K, N - stride_cb, stride_cm, stride_cn, # c: batch, M, N - # Meta-parameters - BM: tl.constexpr, - BK: tl.constexpr, - BN: tl.constexpr, - G: tl.constexpr, - ACTIVATION: tl.constexpr, - HAS_INPUT: tl.constexpr, - HAS_ALPHA: tl.constexpr, - HAS_BETA: tl.constexpr, - ALLOW_TF32: tl.constexpr, - X_DIM: tl.constexpr = 1, -): - """Kernel for computing the matmul C = A x B. - A has shape (M, K), B has shape (K, N) and C has shape (M, N) - """ - # ----------------------------------------------------------- - # Map program ids `pid` to the block of C it should compute. - # This is done in a grouped ordering to promote L2 data reuse. - # See above `L2 Cache Optimizations` section for details. - i_b, i_m, i_n = tl.program_id(0), tl.program_id(1), tl.program_id(2) - - NM, NN = tl.num_programs(1), tl.num_programs(2) - i_m, i_n = tl.swizzle2d(i_m, i_n, NM, NN, G) - - # ---------------------------------------------------------- - # Create pointers for the first blocks of A and B. - # We will advance this pointer as we move in the K direction - # and accumulate - # `p_a` is a block of [BM, BK] pointers - # `p_b` is a block of [BK, BN] pointers - # See above `Pointer Arithmetic` section for details - a_batch_ptr = a + i_b * stride_ab - o_am = (i_m * BM + tl.arange(0, BM)) % M - o_bn = (i_n * BN + tl.arange(0, BN)) % N - o_k = tl.arange(0, BK) - - p_a = a_batch_ptr + (o_am[:, None] * stride_am + o_k[None, :] * stride_ak) - p_b = b + (o_k[:, None] * stride_bk + o_bn[None, :] * stride_bn) - - b_acc = tl.zeros((BM, BN), dtype=tl.float32) - for k in range(0, tl.cdiv(K, BK)): - # Load the next block of A and B, generate a mask by checking the K dimension. - # If it is out of bounds, set it to 0. - b_a = tl.load(p_a, mask=o_k[None, :] < K - k * BK, other=0.0) - b_b = tl.load(p_b, mask=o_k[:, None] < K - k * BK, other=0.0) - # We accumulate along the K dimension. - b_acc = tl.dot(b_a, b_b, acc=b_acc, allow_tf32=ALLOW_TF32) - # Advance the ptrs to the next K block. - p_a += BK * stride_ak - p_b += BK * stride_bk - - o_cm = i_m * BM + tl.arange(0, BM) - o_cn = i_n * BN + tl.arange(0, BN) - mask = (o_cm[:, None] < M) & (o_cn[None, :] < N) - - b_c = b_acc - # You can fuse arbitrary activation functions here - # while the b_acc is still in FP32! - if ACTIVATION == "leaky_relu": - b_c = leaky_relu(b_c) - elif ACTIVATION == "relu": - b_c = relu(b_c) - elif ACTIVATION == "sigmoid": - b_c = sigmoid(b_c) - elif ACTIVATION == "tanh": - b_c = tanh(b_c) - - if HAS_ALPHA: - b_c *= tl.load(alpha) - - if HAS_INPUT: - p_i = input + (stride_cm * o_cm[:, None] if X_DIM == 2 else 0) + stride_cn * o_cn[None, :] - mask_p = (o_cn[None, :] < N) if X_DIM == 1 else mask - b_i = tl.load(p_i, mask=mask_p, other=0.0).to(tl.float32) - if HAS_BETA: - b_i *= tl.load(beta) - b_c += b_i - - # ----------------------------------------------------------- - # Write back the block of the output matrix C with masks. - c_batch_ptr = c + i_b * stride_cb - p_c = c_batch_ptr + stride_cm * o_cm[:, None] + stride_cn * o_cn[None, :] - tl.store(p_c, b_c.to(c.dtype.element_ty), mask=mask) - - -# We can fuse `leaky_relu` by providing it as an `ACTIVATION` meta-parameter in `matmul_kernel`. -@triton.jit -def leaky_relu(x): - return tl.where(x >= 0, x, 0.01 * x) - - -@triton.jit -def sigmoid(x): - # σ(x) = 1 / (1 + exp(-x)) - return 1.0 / (1.0 + exp(-x)) - - -@triton.jit -def tanh(x): - # tanh(x) = (exp(x) - exp(-x)) / (exp(x) + exp(-x)) - # 2 * sigmoid(2x) - 1 - return (exp(x) - exp(-x)) / (exp(x) + exp(-x)) - - -@triton.jit -def relu(x): - # ReLU(x) = max(0, x) - return tl.maximum(x, 0.0) - - -@input_guard -def matmul(a, b, activation=''): - assert a.dim() in [2, 3], "a must be 2D or 3D" - assert b.dim() == 2, "b must be 2D" - assert a.shape[-1] == b.shape[0], f"Incompatible dimensions: A {a.shape}, B {b.shape}" - - if a.dim() == 2: - a_dim = 2 - a = a.unsqueeze(0).contiguous() # (1, M, K) - else: - a_dim = 3 - allow_tf32 = False if a.dtype == torch.float32 else True - - B, M, K = a.shape[0], a.shape[1], a.shape[2] - K_b, N = b.shape - assert K_b == K, f"Incompatible K dimension: A {K} vs B {K_b}" - c = a.new_empty(B, M, N) - - def grid(meta): return (B, triton.cdiv(M, meta['BM']), triton.cdiv(N, meta['BN'])) - matmul_kernel[grid]( - a, b, c, None, None, None, - M, N, K, - a.stride(0), a.stride(1), a.stride(2), # stride_ab, stride_am, stride_ak - b.stride(0), b.stride(1), # stride_bk, stride_bn (b.dim() == 2) - c.stride(0), c.stride(1), c.stride(2), # stride_cb, stride_cm, stride_cn - ACTIVATION=activation, - ALLOW_TF32=allow_tf32, - HAS_INPUT=False, - ) - return c.squeeze(0) if a_dim == 2 else c - - -@input_guard -def addmm( - x: torch.Tensor, - a: torch.Tensor, - b: torch.Tensor, - alpha: float | None = None, - beta: float | None = None, -) -> torch.Tensor: - assert a.dim() in [2, 3], "a must be 2D or 3D" - assert b.dim() == 2, "b must be 2D" - assert a.shape[-1] == b.shape[0], f"Incompatible dimensions: A {a.shape}, B {b.shape}" - - if a.dim() == 2: - a_dim = 2 - a = a.unsqueeze(0).contiguous() # (1, M, K) - else: - a_dim = 3 - allow_tf32 = False if a.dtype == torch.float32 else True - - B, M, K = a.shape[0], a.shape[1], a.shape[2] - K_b, N = b.shape - assert K_b == K, f"Incompatible K dimension: A {K} vs B {K_b}" - c = a.new_empty(B, M, N) - - def grid(meta): return (B, triton.cdiv(M, meta['BM']), triton.cdiv(N, meta['BN'])) - matmul_kernel[grid]( - a, b, c, x, alpha, beta, - M, N, K, - a.stride(0), a.stride(1), a.stride(2), # stride_ab, stride_am, stride_ak - b.stride(0), b.stride(1), # stride_bk, stride_bn (b.dim() == 2) - c.stride(0), c.stride(1), c.stride(2), # stride_cb, stride_cm, stride_cn - ACTIVATION=None, - ALLOW_TF32=allow_tf32, - HAS_INPUT=True, - X_DIM=x.dim(), - ) - return c.squeeze(0) if a_dim == 2 else c diff --git a/build/torch-xpu/ops/utils/op.py b/build/torch-xpu/ops/utils/op.py deleted file mode 100644 index a4f5e69ce2728f5eab09b8566daa843da47c5efb..0000000000000000000000000000000000000000 --- a/build/torch-xpu/ops/utils/op.py +++ /dev/null @@ -1,101 +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 os - -import triton -import triton.language as tl -import triton.language.extra.libdevice as tldevice - -from ...utils import IS_GATHER_SUPPORTED, IS_NVIDIA_BLACKWELL - -if os.environ.get('FLA_USE_FAST_OPS', '0') == '1': - @triton.jit - def exp(x): return tldevice.fast_expf(x.to(tl.float32)) - @triton.jit - def exp2(x): return tldevice.exp2(x.to(tl.float32)) - @triton.jit - def log(x): return tldevice.fast_logf(x.to(tl.float32)) - @triton.jit - def log2(x): return tldevice.fast_log2f(x.to(tl.float32)) - @triton.jit - def tanh(x): return tldevice.fast_tanhf(x.to(tl.float32)) -else: - @triton.jit - def exp(x): return tl.exp(x.to(tl.float32)) - @triton.jit - def exp2(x): return tl.math.exp2(x.to(tl.float32)) - @triton.jit - def log(x): return tl.log(x.to(tl.float32)) - @triton.jit - def log2(x): return tl.log2(x.to(tl.float32)) - @triton.jit - def tanh(x): return tldevice.tanh(x.to(tl.float32)) - - -if IS_NVIDIA_BLACKWELL: - """ - Compute tl.dot with Blackwell workaround. - - On SM100 datacenter and SM120 consumer Blackwell GPUs, wraps the result in - inline assembly to prevent the TritonGPUHoistTMEMAlloc pass from incorrectly - fusing add and dot operations. - See: https://github.com/fla-org/flash-linear-attention/issues/638 - - TODO: Remove this workaround once the Triton compiler bug is fixed. - Track upstream issue at: https://github.com/triton-lang/triton/issues/8695 - """ - @triton.jit - def safe_dot(a, b, allow_tf32: tl.constexpr = None): - return tl.inline_asm_elementwise( - asm="mov.f32 $0, $1;", - constraints="=r,r", - args=[tl.dot(a, b, allow_tf32=allow_tf32)], - dtype=tl.float32, - is_pure=True, - pack=1, - ) -else: - @triton.jit - def safe_dot(a, b, allow_tf32: tl.constexpr = None): - return tl.dot(a, b, allow_tf32=allow_tf32) - - -if not IS_GATHER_SUPPORTED: - @triton.jit - def gather(src, index, axis, _builder=None): - """ - Gather operation that works when tl.gather is not supported. - This is a fallback implementation that returns None. - Just to make triton compiler happy. - """ - return None -else: - gather = tl.gather - - -if hasattr(triton.language, '_experimental_make_tensor_descriptor'): - # For Triton 3.3.x - make_tensor_descriptor = triton.language._experimental_make_tensor_descriptor -elif hasattr(triton.language, 'make_tensor_descriptor'): - # For Triton 3.4.x and later - make_tensor_descriptor = triton.language.make_tensor_descriptor -else: - """ - Fallback implementation when TMA is not supported. - Returns None to indicate TMA descriptors are unavailable. - Just make triton compiler happy. - """ - @triton.jit - def make_tensor_descriptor( - base, - shape, - strides, - block_shape, - _builder=None, - ): - return None diff --git a/build/torch-xpu/ops/utils/pack.py b/build/torch-xpu/ops/utils/pack.py deleted file mode 100644 index 2e2f77074f54f75136045e7838cc330fa93bb1f4..0000000000000000000000000000000000000000 --- a/build/torch-xpu/ops/utils/pack.py +++ /dev/null @@ -1,212 +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 - -# Code adapted from https://github.com/mayank31398/cute-kernels - - -import torch -import triton -import triton.language as tl - -from ...ops.utils.index import prepare_lens -from ...utils import autotune_cache_kwargs, input_guard - - -@triton.autotune( - configs=[ - triton.Config({}, num_warps=num_warps) - for num_warps in [4, 8, 16, 32] - ], - key=['D', 'PADDING_SIDE', 'PACK'], - **autotune_cache_kwargs, -) -@triton.jit -def packunpack_sequence_kernel( - x, - y, - cu_seqlens, - S, - D, - BD: tl.constexpr, - PADDING_SIDE: tl.constexpr, - PACK: tl.constexpr, -): - i_d, i_s, i_b = tl.program_id(0), tl.program_id(1), tl.program_id(2) - bos, eos = tl.load(cu_seqlens + i_b), tl.load(cu_seqlens + i_b + 1) - - T = eos - bos - if PADDING_SIDE == 'left': - NP = S - T - if i_s < NP: - return - i_t = bos + (i_s - NP) - else: - if i_s >= T: - return - i_t = bos + i_s - - o_d = i_d * BD + tl.arange(0, BD) - mask = o_d < D - - if PACK: - b_x = tl.load(x + (i_b * S + i_s) * D + o_d, mask=mask) - tl.store(y + i_t * D + o_d, b_x, mask=mask) - else: - b_x = tl.load(x + i_t * D + o_d, mask=mask) - tl.store(y + (i_b * S + i_s) * D + o_d, b_x, mask=mask) - - -def pack_sequence_fwdbwd( - x: torch.Tensor, - cu_seqlens: torch.Tensor, - padding_side: str, -) -> torch.Tensor: - B, S = x.shape[:2] - D = x.numel() // (B * S) - BD = min(triton.next_power_of_2(D), 4096) - ND = triton.cdiv(D, BD) - - y = torch.empty(cu_seqlens[-1].item(), *x.shape[2:], device=x.device, dtype=x.dtype) - packunpack_sequence_kernel[ND, S, B]( - x=x, - y=y, - cu_seqlens=cu_seqlens, - S=S, - D=D, - BD=BD, - PADDING_SIDE=padding_side, - PACK=True, - ) - return y - - -def unpack_sequence_fwdbwd( - x: torch.Tensor, - cu_seqlens: torch.Tensor, - padding_side: str, - desired_shape: torch.Size, -) -> torch.Tensor: - if desired_shape is None: - desired_shape = (len(cu_seqlens) - 1, prepare_lens(cu_seqlens).max().item(), *x.shape[1:]) - y = torch.zeros(desired_shape, device=x.device, dtype=x.dtype) - B, S = y.shape[:2] - D = y.numel() // (B * S) - BD = min(triton.next_power_of_2(D), 4096) - ND = triton.cdiv(D, BD) - - packunpack_sequence_kernel[ND, S, B]( - x=x, - y=y, - cu_seqlens=cu_seqlens, - S=S, - D=D, - BD=BD, - PADDING_SIDE=padding_side, - PACK=False, - ) - return y - - -class PackSequenceFunction(torch.autograd.Function): - - @staticmethod - @input_guard - def forward( - ctx, - x: torch.Tensor, - cu_seqlens: torch.Tensor, - padding_side: str, - ) -> torch.Tensor: - assert padding_side in ['left', 'right'] - assert x.ndim >= 2 - - ctx.cu_seqlens = cu_seqlens - ctx.padding_side = padding_side - ctx.desired_shape = x.shape - - y = pack_sequence_fwdbwd( - x=x, - cu_seqlens=cu_seqlens, - padding_side=padding_side, - ) - return y - - @staticmethod - @input_guard - def backward(ctx, dy: torch.Tensor) -> tuple[torch.Tensor | None]: - dx = unpack_sequence_fwdbwd( - x=dy, - cu_seqlens=ctx.cu_seqlens, - padding_side=ctx.padding_side, - desired_shape=ctx.desired_shape, - ) - return dx, *[None] * 10 - - -class UnpackSequenceFunction(torch.autograd.Function): - - @staticmethod - @input_guard - def forward( - ctx, - x: torch.Tensor, - cu_seqlens: torch.Tensor, - padding_side: str, - desired_shape: torch.Size | None = None, - ) -> torch.Tensor: - assert padding_side in ['left', 'right'] - assert x.ndim >= 2 - if desired_shape is not None: - assert desired_shape[0] == cu_seqlens.shape[0] - 1 - assert desired_shape[2:] == x.shape[1:] - - ctx.cu_seqlens = cu_seqlens - ctx.padding_side = padding_side - - y = unpack_sequence_fwdbwd( - x=x, - cu_seqlens=cu_seqlens, - padding_side=padding_side, - desired_shape=desired_shape, - ) - return y - - @staticmethod - @input_guard - def backward(ctx, dy: torch.Tensor) -> tuple[torch.Tensor | None]: - dx = pack_sequence_fwdbwd( - x=dy, - cu_seqlens=ctx.cu_seqlens, - padding_side=ctx.padding_side, - ) - return dx, None, None, None - - -def pack_sequence( - x: torch.Tensor, - cu_seqlens: torch.Tensor, - padding_side: str = 'left', -) -> torch.Tensor: - return PackSequenceFunction.apply( - x, - cu_seqlens, - padding_side, - ) - - -def unpack_sequence( - x: torch.Tensor, - cu_seqlens: torch.Tensor, - padding_side: str = 'left', - desired_shape: torch.Size | None = None, -) -> torch.Tensor: - return UnpackSequenceFunction.apply( - x, - cu_seqlens, - padding_side, - desired_shape, - ) diff --git a/build/torch-xpu/ops/utils/pooling.py b/build/torch-xpu/ops/utils/pooling.py deleted file mode 100644 index b8c805c928b5919443e45374972df8a30b2d4d5e..0000000000000000000000000000000000000000 --- a/build/torch-xpu/ops/utils/pooling.py +++ /dev/null @@ -1,215 +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 -import triton.language as tl - -from ...ops.utils.index import prepare_chunk_indices -from ...utils import autocast_custom_bwd, autocast_custom_fwd, autotune_cache_kwargs, input_guard - - -@triton.heuristics({ - 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None, -}) -@triton.autotune( - configs=[ - triton.Config({'BD': BD}, num_warps=num_warps) - for BD in [16, 32, 64, 128] - for num_warps in [1, 2, 4, 8] - ], - key=['BT'], - **autotune_cache_kwargs, -) -@triton.jit(do_not_specialize=['T']) -def mean_pooling_fwd_kernel( - x, - o, - cu_seqlens, - chunk_indices, - T, - H: tl.constexpr, - D: tl.constexpr, - BT: tl.constexpr, - BD: tl.constexpr, - IS_VARLEN: tl.constexpr, -): - i_d, i_t, i_bh = tl.program_id(0), tl.program_id(1), tl.program_id(2) - i_b, i_h = i_bh // H, i_bh % H - 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.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) - T = eos - bos - NT = tl.cdiv(T, BT) - else: - NT = tl.cdiv(T, BT) - i_tg = i_b * NT + i_t - bos, eos = i_b * T, i_b * T + T - - p_x = tl.make_block_ptr(x + (bos * H + i_h) * D, (T, D), (H*D, 1), (i_t * BT, i_d * BD), (BT, BD), (1, 0)) - p_o = tl.make_block_ptr(o + (i_tg * H + i_h) * D, (D,), (1,), (i_d * BD,), (BD,), (0,)) - # [BT, BD] - b_x = tl.load(p_x, boundary_check=(0, 1)).to(tl.float32) - # [BD] - b_o = tl.sum(b_x, axis=0) / min(BT, T - i_t * BT) - tl.store(p_o, b_o.to(p_o.dtype.element_ty), boundary_check=(0,)) - - -@triton.heuristics({ - 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None, -}) -@triton.autotune( - configs=[ - triton.Config({'BD': BD}, num_warps=num_warps) - for BD in [16, 32, 64, 128] - for num_warps in [1, 2, 4, 8] - ], - key=['BT'], - **autotune_cache_kwargs, -) -@triton.jit(do_not_specialize=['T']) -def mean_pooling_bwd_kernel( - do, - dx, - cu_seqlens, - chunk_indices, - T, - H: tl.constexpr, - D: tl.constexpr, - BT: tl.constexpr, - BD: tl.constexpr, - IS_VARLEN: tl.constexpr, -): - i_d, i_t, i_bh = tl.program_id(0), tl.program_id(1), tl.program_id(2) - i_b, i_h = i_bh // H, i_bh % H - 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.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) - T = eos - bos - NT = tl.cdiv(T, BT) - else: - NT = tl.cdiv(T, BT) - i_tg = i_b * NT + i_t - bos, eos = i_b * T, i_b * T + T - - p_dx = tl.make_block_ptr(dx + (bos * H + i_h) * D, (T, D), (H*D, 1), (i_t * BT, i_d * BD), (BT, BD), (1, 0)) - p_do = tl.make_block_ptr(do + (i_tg * H + i_h) * D, (D,), (1,), (i_d * BD,), (BD,), (0,)) - # [BD] - b_do = tl.load(p_do, boundary_check=(0,)).to(tl.float32) - # [BT, BD] - b_dx = b_do / tl.full((BT,), min(BT, T - i_t * BT), dtype=tl.float32)[:, None] - tl.store(p_dx, b_dx.to(p_dx.dtype.element_ty), boundary_check=(0, 1)) - - -def mean_pooling_fwd( - x: torch.Tensor, - chunk_size: int, - cu_seqlens: torch.LongTensor | None = None, - chunk_indices: torch.LongTensor | None = None, -) -> torch.Tensor: - B, T, H, D = x.shape - BT = chunk_size - if chunk_indices is None and cu_seqlens is not None: - chunk_indices = prepare_chunk_indices(cu_seqlens, chunk_size) - NT = triton.cdiv(T, BT) if cu_seqlens is None else len(chunk_indices) - - o = x.new_empty(B, NT, H, D) - def grid(meta): return (triton.cdiv(D, meta['BD']), NT, B * H) - mean_pooling_fwd_kernel[grid]( - x, - o, - cu_seqlens, - chunk_indices, - T=T, - H=H, - D=D, - BT=BT, - ) - return o - - -def mean_pooling_bwd( - do: torch.Tensor, - batch_size: int, - seq_len: int, - chunk_size: int, - cu_seqlens: torch.LongTensor | None = None, - chunk_indices: torch.LongTensor | None = None, -) -> torch.Tensor: - B, T, H, D = batch_size, seq_len, *do.shape[-2:] - BT = chunk_size - if chunk_indices is None and cu_seqlens is not None: - chunk_indices = prepare_chunk_indices(cu_seqlens, chunk_size) - NT = triton.cdiv(T, BT) if cu_seqlens is None else len(chunk_indices) - - dx = do.new_empty(B, T, H, D) - def grid(meta): return (triton.cdiv(D, meta['BD']), NT, B * H) - mean_pooling_bwd_kernel[grid]( - do, - dx, - cu_seqlens, - chunk_indices, - T=T, - H=H, - D=D, - BT=BT, - ) - return dx - - -class MeanPoolingFunction(torch.autograd.Function): - - @staticmethod - @input_guard - @autocast_custom_fwd - def forward( - ctx, - x: torch.Tensor, - chunk_size: int, - cu_seqlens: torch.LongTensor | None = None, - ) -> torch.Tensor: - o = mean_pooling_fwd(x, chunk_size, cu_seqlens) - ctx.batch_size = x.shape[0] - ctx.seq_len = x.shape[1] - ctx.chunk_size = chunk_size - ctx.cu_seqlens = cu_seqlens - return o - - @staticmethod - @input_guard - @autocast_custom_bwd - def backward( - ctx, do, - ) -> tuple[torch.Tensor, None, None]: - batch_size = ctx.batch_size - seq_len = ctx.seq_len - chunk_size = ctx.chunk_size - cu_seqlens = ctx.cu_seqlens - dx = mean_pooling_bwd(do, batch_size, seq_len, chunk_size, cu_seqlens) - return dx, None, None - - -def mean_pooling( - x: torch.Tensor, - chunk_size: int, - cu_seqlens: torch.LongTensor | None = None, - head_first: bool = False, -) -> torch.Tensor: - if head_first: - x = x.transpose(1, 2) - if cu_seqlens is not None: - if x.shape[0] != 1: - raise ValueError( - f"The batch size is expected to be 1 rather than {x.shape[0]} when using `cu_seqlens`." - f"Please flatten variable-length inputs before processing.", - ) - o = MeanPoolingFunction.apply(x, chunk_size, cu_seqlens) - if head_first: - o = o.transpose(1, 2) - return o diff --git a/build/torch-xpu/ops/utils/softmax.py b/build/torch-xpu/ops/utils/softmax.py deleted file mode 100644 index 55806ad9e7ac7fbccec07c69d9e18bace3b5e34e..0000000000000000000000000000000000000000 --- a/build/torch-xpu/ops/utils/softmax.py +++ /dev/null @@ -1,110 +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 -import triton.language as tl - -from ...ops.utils.op import exp -from ...utils import IS_AMD, autotune_cache_kwargs - -NUM_WARPS_AUTOTUNE = [1, 2, 4, 8, 16] if IS_AMD else [1, 2, 4, 8, 16, 32] - - -@triton.autotune( - configs=[ - triton.Config({}, num_warps=num_warps) - for num_warps in NUM_WARPS_AUTOTUNE - ], - key=['D'], - **autotune_cache_kwargs, -) -@triton.jit -def softmax_fwd_kernel( - x, - p, - D: tl.constexpr, - B: tl.constexpr, -): - i_n = tl.program_id(0) - o_d = 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')) - b_m = tl.max(b_x, 0) - b_x = exp(b_x - b_m) - b_p = b_x / tl.sum(b_x, 0) - - tl.store(p + i_n * D + o_d, b_p.to(p.dtype.element_ty), mask=m_d) - - -@triton.autotune( - configs=[ - triton.Config({}, num_warps=num_warps) - for num_warps in NUM_WARPS_AUTOTUNE - ], - key=['D'], - **autotune_cache_kwargs, -) -@triton.jit -def softmax_bwd_kernel( - p, - dp, - ds, - D: tl.constexpr, - B: tl.constexpr, -): - i_n = tl.program_id(0) - o_d = tl.arange(0, B) - m_d = o_d < D - - b_p = tl.load(p + i_n * D + o_d, mask=m_d, other=0.) - b_dp = tl.load(dp + i_n * D + o_d, mask=m_d, other=0.) - b_pp = tl.sum(b_p * b_dp, 0) - b_ds = b_p * b_dp - b_p * b_pp - tl.store(ds + i_n * D + o_d, b_ds.to(ds.dtype.element_ty), mask=m_d) - - -def softmax_fwd( - x: torch.Tensor, - dtype: torch.dtype | None = torch.float, -) -> torch.Tensor: - shape = x.shape - x = x.view(-1, x.shape[-1]) - - N, D = x.shape - B = triton.next_power_of_2(D) - - p = torch.empty_like(x, dtype=dtype) - softmax_fwd_kernel[(N,)]( - x=x, - p=p, - D=D, - B=B, - ) - return p.view(*shape) - - -def softmax_bwd( - p: torch.Tensor, - dp: torch.Tensor, - dtype: torch.dtype | None = torch.float, -) -> torch.Tensor: - shape = p.shape - p = p.view(-1, p.shape[-1]) - ds = torch.empty_like(p, dtype=dtype) - - N, D = p.shape - B = triton.next_power_of_2(D) - softmax_bwd_kernel[(N,)]( - p=p, - dp=dp, - ds=ds, - D=D, - B=B, - ) - return ds.view(*shape) diff --git a/build/torch-xpu/ops/utils/softplus.py b/build/torch-xpu/ops/utils/softplus.py deleted file mode 100644 index b8844e7fdb7127495dba8659684c37bd484f4134..0000000000000000000000000000000000000000 --- a/build/torch-xpu/ops/utils/softplus.py +++ /dev/null @@ -1,115 +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 - -# REVISED FROM -# https://github.com/shawntan/stickbreaking-attention/blob/main/stickbreaking_attention/sb_varlen/softplus.py - -import triton -from triton import language as tl - -from ...utils import IS_NVIDIA - - -def _generate_softplus(num_pack): - template = """ - .reg .pred p; - setp.gt.f32 p, ${in_reg}, 20.; - @p mov.f32 ${out_reg}, ${in_reg}; - @!p mul.f32 ${out_reg}, ${in_reg}, 1.4426950408889634; - @!p ex2.approx.ftz.f32 ${out_reg}, ${out_reg}; - @!p add.f32 ${out_reg}, ${out_reg}, 1.0; - @!p lg2.approx.ftz.f32 ${out_reg}, ${out_reg}; - @!p mul.f32 ${out_reg}, ${out_reg}, 0.6931471805599453; - """ - out_str = "" - - for i in range(num_pack): - inner_str = template.format(out_reg=i, in_reg=i + num_pack) - out_str += "{" + inner_str + "}\n" - # flatten out because torch.compile doesn't like newlines - out_str = " ".join(out_str.split("\n")) - return out_str - - -def _generate_softplus2(num_pack): - template = """ - .reg .pred p; - setp.gt.f32 p, ${in_reg}, 15.; - @p mov.f32 ${out_reg}, ${in_reg}; - @!p ex2.approx.ftz.f32 ${out_reg}, ${in_reg}; - @!p add.f32 ${out_reg}, ${out_reg}, 1.0; - @!p lg2.approx.ftz.f32 ${out_reg}, ${out_reg}; - """ - out_str = "" - - for i in range(num_pack): - inner_str = template.format(out_reg=i, in_reg=i + num_pack) - out_str += "{" + inner_str + "}\n" - # flatten out because torch.compile doesn't like newlines - out_str = " ".join(out_str.split("\n")) - return out_str - - -def _generate_constraints(num_pack): - return ",".join("=r" for i in range(num_pack)) + "," + ",".join("r" for i in range(num_pack)) - - -_NUM_REG = 1 -s_softplus: tl.constexpr = tl.constexpr(_generate_softplus(_NUM_REG)) -s_softplus2: tl.constexpr = tl.constexpr(_generate_softplus2(_NUM_REG)) -s_constraints: tl.constexpr = tl.constexpr(_generate_constraints(_NUM_REG)) -NUM_REG: tl.constexpr = tl.constexpr(_NUM_REG) - - -@triton.jit -def softplus_nv(x): - # equivalent to: - # return tl.where(x < 20.0, tl.math.log(1 + tl.math.exp(x)), x) - return tl.inline_asm_elementwise( - asm=s_softplus, - constraints=s_constraints, - pack=NUM_REG, - args=[ - x, - ], - dtype=tl.float32, - is_pure=True, - ) - - -@triton.jit -def softplus_triton(x): - return tl.where(x < 20.0, tl.math.log(1 + tl.math.exp(x)), x) - - -@triton.jit -def softplus2_nv(x): - # equivalent to: - # return tl.where(x < 15.0, tl.math.log2(1 + tl.math.exp2(x)), x) - return tl.inline_asm_elementwise( - asm=s_softplus2, - constraints=s_constraints, - pack=NUM_REG, - args=[ - x, - ], - dtype=tl.float32, - is_pure=True, - ) - - -@triton.jit -def softplus2_triton(x): - return tl.where(x < 15.0, tl.math.log2(1 + tl.math.exp2(x)), x) - - -if IS_NVIDIA: - softplus = softplus_nv - softplus2 = softplus2_nv -else: - softplus = softplus_triton - softplus2 = softplus2_triton diff --git a/build/torch-xpu/ops/utils/solve_tril.py b/build/torch-xpu/ops/utils/solve_tril.py deleted file mode 100644 index f61ce4a3030e81d2747a5eaca8ec98928a7084f4..0000000000000000000000000000000000000000 --- a/build/torch-xpu/ops/utils/solve_tril.py +++ /dev/null @@ -1,397 +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 os - -import torch -import triton -import triton.language as tl - -from ...ops.utils.index import prepare_chunk_indices -from ...ops.utils.op import make_tensor_descriptor -from ...utils import IS_TMA_SUPPORTED, autotune_cache_kwargs, input_guard - -FLA_TRIL_PRECISION = os.environ.get('FLA_TRIL_PRECISION', 'ieee') -assert FLA_TRIL_PRECISION in ['ieee', 'tf32', 'tf32x3'], \ - f"FLA_TRIL_PRECISION must be one of 'ieee', 'tf32', or 'tf32x3', but got {FLA_TRIL_PRECISION}" -DOT_PRECISION_AUTOTUNE_LIST = ["ieee"] if not IS_TMA_SUPPORTED else list({"ieee", FLA_TRIL_PRECISION}) - - -@triton.heuristics({ - 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None, -}) -@triton.autotune( - configs=[ - triton.Config({'DOT_PRECISION': 'ieee'}, num_warps=num_warps, num_stages=num_stages) - for num_warps in [1, 2, 4, 8] - for num_stages in [2, 3, 4, 5] - ], - key=['BT'], - **autotune_cache_kwargs, -) -@triton.jit(do_not_specialize=['T']) -def solve_tril_16x16_kernel( - A, - Ai, - cu_seqlens, - chunk_indices, - T, - H: tl.constexpr, - BT: tl.constexpr, - USE_TMA: tl.constexpr, - IS_VARLEN: tl.constexpr, - DOT_PRECISION: tl.constexpr, -): - i_t, i_bh = tl.program_id(0), tl.program_id(1).to(tl.int64) - i_b, i_h = i_bh // H, i_bh % H - 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.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) - T = eos - bos - else: - bos, eos = i_b * T, i_b * T + T - o_i = tl.arange(0, 16) - m_A = o_i[:, None] > o_i[None, :] - m_I = o_i[:, None] == o_i[None, :] - - A = A + (bos*H + i_h) * BT - Ai = Ai + (bos*H + i_h) * 16 - - offset = (i_t * 16) % BT - if not USE_TMA: - p_A = tl.make_block_ptr(A, (T, BT), (H*BT, 1), (i_t * 16, offset), (16, 16), (1, 0)) - # [16, 16] - b_A = tl.load(p_A, boundary_check=(0, 1)).to(tl.float32) - b_A = tl.where(m_A, b_A, 0) - else: - desc = make_tensor_descriptor(A, [T, BT], [H*BT, 1], [16, 16]) - desc_o = make_tensor_descriptor(Ai, [T, 16], [H*16, 1], [16, 16]) - b_A = desc.load([i_t * 16, offset]).to(tl.float32) - b_A = tl.where(m_A, b_A, 0) - b_A = -b_A - - for i in range(2, min(16, T - i_t * 16)): - # [16] - b_a = -tl.load(A + (i_t * 16 + i) * H*BT + o_i + offset) - b_a = tl.where(o_i < i, b_a, 0.) - b_a = b_a + tl.sum(b_a[:, None] * b_A, 0) - b_A = tl.where((o_i == i)[:, None], b_a, b_A) - b_A += m_I - if not USE_TMA: - p_Ai = tl.make_block_ptr(Ai, (T, 16), (H*16, 1), (i_t * 16, 0), (16, 16), (1, 0)) - tl.store(p_Ai, b_A.to(p_Ai.dtype.element_ty, fp_downcast_rounding="rtne"), boundary_check=(0, 1)) - else: - desc_o.store([i_t * 16, 0], b_A.to(desc_o.dtype, fp_downcast_rounding="rtne")) - - -@triton.heuristics({ - 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None, -}) -@triton.autotune( - configs=[ - triton.Config({'DOT_PRECISION': DOT_PRECISION}, num_warps=num_warps, num_stages=num_stages) - for num_warps in [1, 2, 4, 8] - for num_stages in [2, 3, 4, 5] - for DOT_PRECISION in DOT_PRECISION_AUTOTUNE_LIST - ], - key=['H', 'BT', 'IS_VARLEN'], - **autotune_cache_kwargs, -) -@triton.jit(do_not_specialize=['T']) -def merge_16x16_to_32x32_inverse_kernel( - A, - Ai, - cu_seqlens, - chunk_indices, - T, - H: tl.constexpr, - BT: tl.constexpr, - USE_TMA: tl.constexpr, - IS_VARLEN: tl.constexpr, - DOT_PRECISION: tl.constexpr, -): - i_t, i_bh = tl.program_id(0), tl.program_id(1).to(tl.int64) - i_b, i_h = i_bh // H, i_bh % H - 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.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) - T = eos - bos - else: - bos, eos = i_b * T, i_b * T + T - - o_i = tl.arange(0, 16) - m_A = o_i[:, None] > o_i[None, :] - m_I = o_i[:, None] == o_i[None, :] - A += (bos * H + i_h) * BT - Ai += (bos * H + i_h) * BT - - if not USE_TMA: - p_A_11 = tl.make_block_ptr(A, (T, BT), (H*BT, 1), (i_t * BT, 0), (16, 16), (1, 0)) - p_A_22 = tl.make_block_ptr(A, (T, BT), (H*BT, 1), (i_t * BT + 16, 16), (16, 16), (1, 0)) - b_Ai_11 = tl.load(p_A_11, boundary_check=(0, 1)).to(tl.float32) - b_Ai_22 = tl.load(p_A_22, boundary_check=(0, 1)).to(tl.float32) - else: - desc = make_tensor_descriptor(A, [T, BT], [H*BT, 1], [16, 16]) - desc_o = make_tensor_descriptor(Ai, [T, BT], [H*BT, 1], [16, 16]) - b_Ai_11 = desc.load([i_t * BT + 0, 0]).to(tl.float32) - b_Ai_22 = desc.load([i_t * BT + 16, 16]).to(tl.float32) - - # [16, 16] - b_Ai_11 = -tl.where(m_A, b_Ai_11, 0) - b_Ai_22 = -tl.where(m_A, b_Ai_22, 0) - - for i in range(2, min(16, T - i_t * BT)): - b_a_11 = -tl.load(A + (i_t * BT + i) * H*BT + o_i) - b_a_11 += tl.sum(b_a_11[:, None] * b_Ai_11, 0) - b_Ai_11 = tl.where((o_i == i)[:, None], b_a_11, b_Ai_11) - for i in range(16 + 2, min(32, T - i_t * BT)): - b_a_22 = -tl.load(A + (i_t * BT + i) * H*BT + o_i + 16) - b_a_22 += tl.sum(b_a_22[:, None] * b_Ai_22, 0) - b_Ai_22 = tl.where((o_i == i - 16)[:, None], b_a_22, b_Ai_22) - - b_Ai_11 += m_I - b_Ai_22 += m_I - - if not USE_TMA: - p_A_21 = tl.make_block_ptr(A, (T, BT), (H*BT, 1), (i_t * BT + 16, 0), (16, 16), (1, 0)) - b_A_21 = tl.load(p_A_21, boundary_check=(0, 1)).to(tl.float32) - else: - b_A_21 = desc.load([i_t * BT + 16, 0]).to(tl.float32) - - b_Ai_21 = -tl.dot(tl.dot(b_Ai_22, b_A_21, input_precision=DOT_PRECISION), b_Ai_11, input_precision=DOT_PRECISION) - - if not USE_TMA: - p_Ai_11 = tl.make_block_ptr(Ai, (T, BT), (H*BT, 1), (i_t * BT, 0), (16, 16), (1, 0)) - p_Ai_21 = tl.make_block_ptr(Ai, (T, BT), (H*BT, 1), (i_t * BT + 16, 0), (16, 16), (1, 0)) - p_Ai_22 = tl.make_block_ptr(Ai, (T, BT), (H*BT, 1), (i_t * BT + 16, 16), (16, 16), (1, 0)) - tl.store(p_Ai_11, b_Ai_11.to(p_Ai_11.dtype.element_ty, fp_downcast_rounding="rtne"), boundary_check=(0, 1)) - tl.store(p_Ai_22, b_Ai_22.to(p_Ai_22.dtype.element_ty, fp_downcast_rounding="rtne"), boundary_check=(0, 1)) - tl.store(p_Ai_21, b_Ai_21.to(p_Ai_21.dtype.element_ty, fp_downcast_rounding="rtne"), boundary_check=(0, 1)) - else: - desc_o.store([i_t * BT + 0, 0], b_Ai_11.to(desc_o.dtype, fp_downcast_rounding="rtne")) - desc_o.store([i_t * BT + 16, 0], b_Ai_21.to(desc_o.dtype, fp_downcast_rounding="rtne")) - desc_o.store([i_t * BT + 16, 16], b_Ai_22.to(desc_o.dtype, fp_downcast_rounding="rtne")) - - -@triton.heuristics({ - 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None, -}) -@triton.autotune( - configs=[ - triton.Config({'DOT_PRECISION': DOT_PRECISION}, num_warps=num_warps, num_stages=num_stages) - for num_warps in [2, 4, 8] - for num_stages in [2, 3, 4, 5] - for DOT_PRECISION in DOT_PRECISION_AUTOTUNE_LIST - ], - key=['H', 'BT', 'IS_VARLEN'], - **autotune_cache_kwargs, -) -@triton.jit(do_not_specialize=['T']) -def merge_16x16_to_64x64_inverse_kernel( - A, - Ai, - cu_seqlens, - chunk_indices, - T, - H: tl.constexpr, - BT: tl.constexpr, - USE_TMA: tl.constexpr, - IS_VARLEN: tl.constexpr, - DOT_PRECISION: tl.constexpr, -): - i_t, i_bh = tl.program_id(0), tl.program_id(1).to(tl.int64) - i_b, i_h = i_bh // H, i_bh % H - 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.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32) - T = eos - bos - else: - bos, eos = i_b * T, i_b * T + T - - o_i = tl.arange(0, 16) - m_A = o_i[:, None] > o_i[None, :] - m_I = o_i[:, None] == o_i[None, :] - A += (bos * H + i_h) * BT - Ai += (bos * H + i_h) * BT - - if not USE_TMA: - p_A_11 = tl.make_block_ptr(A, (T, BT), (H*BT, 1), (i_t * BT, 0), (16, 16), (1, 0)) - p_A_22 = tl.make_block_ptr(A, (T, BT), (H*BT, 1), (i_t * BT + 16, 16), (16, 16), (1, 0)) - p_A_33 = tl.make_block_ptr(A, (T, BT), (H*BT, 1), (i_t * BT + 32, 32), (16, 16), (1, 0)) - p_A_44 = tl.make_block_ptr(A, (T, BT), (H*BT, 1), (i_t * BT + 48, 48), (16, 16), (1, 0)) - b_Ai_11 = tl.load(p_A_11, boundary_check=(0, 1)).to(tl.float32) - b_Ai_22 = tl.load(p_A_22, boundary_check=(0, 1)).to(tl.float32) - b_Ai_33 = tl.load(p_A_33, boundary_check=(0, 1)).to(tl.float32) - b_Ai_44 = tl.load(p_A_44, boundary_check=(0, 1)).to(tl.float32) - else: - desc = make_tensor_descriptor(A, [T, BT], [H*BT, 1], [16, 16]) - desc_o = make_tensor_descriptor(Ai, [T, BT], [H*BT, 1], [16, 16]) - b_Ai_11 = desc.load([i_t * BT + 0, 0]).to(tl.float32) - b_Ai_22 = desc.load([i_t * BT + 16, 16]).to(tl.float32) - b_Ai_33 = desc.load([i_t * BT + 32, 32]).to(tl.float32) - b_Ai_44 = desc.load([i_t * BT + 48, 48]).to(tl.float32) - - # [16, 16] - b_Ai_11 = -tl.where(m_A, b_Ai_11, 0) - b_Ai_22 = -tl.where(m_A, b_Ai_22, 0) - b_Ai_33 = -tl.where(m_A, b_Ai_33, 0) - b_Ai_44 = -tl.where(m_A, b_Ai_44, 0) - - for i in range(2, min(16, T - i_t * BT)): - b_a_11 = -tl.load(A + (i_t * BT + i) * H*BT + o_i) - b_a_11 = tl.where(o_i < i, b_a_11, 0.) - b_a_11 += tl.sum(b_a_11[:, None] * b_Ai_11, 0) - b_Ai_11 = tl.where((o_i == i)[:, None], b_a_11, b_Ai_11) - for i in range(16 + 2, min(32, T - i_t * BT)): - b_a_22 = -tl.load(A + (i_t * BT + i) * H*BT + o_i + 16) - b_a_22 = tl.where(o_i < i - 16, b_a_22, 0.) - b_a_22 += tl.sum(b_a_22[:, None] * b_Ai_22, 0) - b_Ai_22 = tl.where((o_i == i - 16)[:, None], b_a_22, b_Ai_22) - for i in range(32 + 2, min(48, T - i_t * BT)): - b_a_33 = -tl.load(A + (i_t * BT + i) * H*BT + o_i + 32) - b_a_33 = tl.where(o_i < i - 32, b_a_33, 0.) - b_a_33 += tl.sum(b_a_33[:, None] * b_Ai_33, 0) - b_Ai_33 = tl.where((o_i == i - 32)[:, None], b_a_33, b_Ai_33) - for i in range(48 + 2, min(64, T - i_t * BT)): - b_a_44 = -tl.load(A + (i_t * BT + i) * H*BT + o_i + 48) - b_a_44 = tl.where(o_i < i - 48, b_a_44, 0.) - b_a_44 += tl.sum(b_a_44[:, None] * b_Ai_44, 0) - b_Ai_44 = tl.where((o_i == i - 48)[:, None], b_a_44, b_Ai_44) - b_Ai_11 += m_I - b_Ai_22 += m_I - b_Ai_33 += m_I - b_Ai_44 += m_I - - if not USE_TMA: - p_A_21 = tl.make_block_ptr(A, (T, BT), (H*BT, 1), (i_t * BT + 16, 0), (16, 16), (1, 0)) - p_A_31 = tl.make_block_ptr(A, (T, BT), (H*BT, 1), (i_t * BT + 32, 0), (16, 16), (1, 0)) - p_A_32 = tl.make_block_ptr(A, (T, BT), (H*BT, 1), (i_t * BT + 32, 16), (16, 16), (1, 0)) - p_A_41 = tl.make_block_ptr(A, (T, BT), (H*BT, 1), (i_t * BT + 48, 0), (16, 16), (1, 0)) - p_A_42 = tl.make_block_ptr(A, (T, BT), (H*BT, 1), (i_t * BT + 48, 16), (16, 16), (1, 0)) - p_A_43 = tl.make_block_ptr(A, (T, BT), (H*BT, 1), (i_t * BT + 48, 32), (16, 16), (1, 0)) - b_A_21 = tl.load(p_A_21, boundary_check=(0, 1)).to(tl.float32) - b_A_31 = tl.load(p_A_31, boundary_check=(0, 1)).to(tl.float32) - b_A_32 = tl.load(p_A_32, boundary_check=(0, 1)).to(tl.float32) - b_A_41 = tl.load(p_A_41, boundary_check=(0, 1)).to(tl.float32) - b_A_42 = tl.load(p_A_42, boundary_check=(0, 1)).to(tl.float32) - b_A_43 = tl.load(p_A_43, boundary_check=(0, 1)).to(tl.float32) - else: - b_A_21 = desc.load([i_t * BT + 16, 0]).to(tl.float32) - b_A_31 = desc.load([i_t * BT + 32, 0]).to(tl.float32) - b_A_32 = desc.load([i_t * BT + 32, 16]).to(tl.float32) - b_A_41 = desc.load([i_t * BT + 48, 0]).to(tl.float32) - b_A_42 = desc.load([i_t * BT + 48, 16]).to(tl.float32) - b_A_43 = desc.load([i_t * BT + 48, 32]).to(tl.float32) - - b_Ai_21 = -tl.dot(tl.dot(b_Ai_22, b_A_21, input_precision=DOT_PRECISION), b_Ai_11, input_precision=DOT_PRECISION) - b_Ai_32 = -tl.dot(tl.dot(b_Ai_33, b_A_32, input_precision=DOT_PRECISION), b_Ai_22, input_precision=DOT_PRECISION) - b_Ai_43 = -tl.dot(tl.dot(b_Ai_44, b_A_43, input_precision=DOT_PRECISION), b_Ai_33, input_precision=DOT_PRECISION) - - b_Ai_31 = -tl.dot( - b_Ai_33, - tl.dot(b_A_31, b_Ai_11, input_precision=DOT_PRECISION) + - tl.dot(b_A_32, b_Ai_21, input_precision=DOT_PRECISION), - input_precision=DOT_PRECISION, - ) - b_Ai_42 = -tl.dot( - b_Ai_44, - tl.dot(b_A_42, b_Ai_22, input_precision=DOT_PRECISION) + - tl.dot(b_A_43, b_Ai_32, input_precision=DOT_PRECISION), - input_precision=DOT_PRECISION, - ) - b_Ai_41 = -tl.dot( - b_Ai_44, - tl.dot(b_A_41, b_Ai_11, input_precision=DOT_PRECISION) + - tl.dot(b_A_42, b_Ai_21, input_precision=DOT_PRECISION) + - tl.dot(b_A_43, b_Ai_31, input_precision=DOT_PRECISION), - input_precision=DOT_PRECISION, - ) - - if not USE_TMA: - p_Ai_11 = tl.make_block_ptr(Ai, (T, BT), (H*BT, 1), (i_t * BT, 0), (16, 16), (1, 0)) - p_Ai_22 = tl.make_block_ptr(Ai, (T, BT), (H*BT, 1), (i_t * BT + 16, 16), (16, 16), (1, 0)) - p_Ai_33 = tl.make_block_ptr(Ai, (T, BT), (H*BT, 1), (i_t * BT + 32, 32), (16, 16), (1, 0)) - p_Ai_44 = tl.make_block_ptr(Ai, (T, BT), (H*BT, 1), (i_t * BT + 48, 48), (16, 16), (1, 0)) - p_Ai_21 = tl.make_block_ptr(Ai, (T, BT), (H*BT, 1), (i_t * BT + 16, 0), (16, 16), (1, 0)) - p_Ai_31 = tl.make_block_ptr(Ai, (T, BT), (H*BT, 1), (i_t * BT + 32, 0), (16, 16), (1, 0)) - p_Ai_32 = tl.make_block_ptr(Ai, (T, BT), (H*BT, 1), (i_t * BT + 32, 16), (16, 16), (1, 0)) - p_Ai_41 = tl.make_block_ptr(Ai, (T, BT), (H*BT, 1), (i_t * BT + 48, 0), (16, 16), (1, 0)) - p_Ai_42 = tl.make_block_ptr(Ai, (T, BT), (H*BT, 1), (i_t * BT + 48, 16), (16, 16), (1, 0)) - p_Ai_43 = tl.make_block_ptr(Ai, (T, BT), (H*BT, 1), (i_t * BT + 48, 32), (16, 16), (1, 0)) - tl.store(p_Ai_11, b_Ai_11.to(p_Ai_11.dtype.element_ty, fp_downcast_rounding="rtne"), boundary_check=(0, 1)) - tl.store(p_Ai_22, b_Ai_22.to(p_Ai_22.dtype.element_ty, fp_downcast_rounding="rtne"), boundary_check=(0, 1)) - tl.store(p_Ai_33, b_Ai_33.to(p_Ai_33.dtype.element_ty, fp_downcast_rounding="rtne"), boundary_check=(0, 1)) - tl.store(p_Ai_44, b_Ai_44.to(p_Ai_44.dtype.element_ty, fp_downcast_rounding="rtne"), boundary_check=(0, 1)) - tl.store(p_Ai_21, b_Ai_21.to(p_Ai_21.dtype.element_ty, fp_downcast_rounding="rtne"), boundary_check=(0, 1)) - tl.store(p_Ai_31, b_Ai_31.to(p_Ai_31.dtype.element_ty, fp_downcast_rounding="rtne"), boundary_check=(0, 1)) - tl.store(p_Ai_32, b_Ai_32.to(p_Ai_32.dtype.element_ty, fp_downcast_rounding="rtne"), boundary_check=(0, 1)) - tl.store(p_Ai_41, b_Ai_41.to(p_Ai_41.dtype.element_ty, fp_downcast_rounding="rtne"), boundary_check=(0, 1)) - tl.store(p_Ai_42, b_Ai_42.to(p_Ai_42.dtype.element_ty, fp_downcast_rounding="rtne"), boundary_check=(0, 1)) - tl.store(p_Ai_43, b_Ai_43.to(p_Ai_43.dtype.element_ty, fp_downcast_rounding="rtne"), boundary_check=(0, 1)) - else: - desc_o.store([i_t * BT + 0, 0], b_Ai_11.to(desc_o.dtype, fp_downcast_rounding="rtne")) - desc_o.store([i_t * BT + 16, 16], b_Ai_22.to(desc_o.dtype, fp_downcast_rounding="rtne")) - desc_o.store([i_t * BT + 32, 32], b_Ai_33.to(desc_o.dtype, fp_downcast_rounding="rtne")) - desc_o.store([i_t * BT + 48, 48], b_Ai_44.to(desc_o.dtype, fp_downcast_rounding="rtne")) - desc_o.store([i_t * BT + 16, 0], b_Ai_21.to(desc_o.dtype, fp_downcast_rounding="rtne")) - desc_o.store([i_t * BT + 32, 0], b_Ai_31.to(desc_o.dtype, fp_downcast_rounding="rtne")) - desc_o.store([i_t * BT + 32, 16], b_Ai_32.to(desc_o.dtype, fp_downcast_rounding="rtne")) - desc_o.store([i_t * BT + 48, 0], b_Ai_41.to(desc_o.dtype, fp_downcast_rounding="rtne")) - desc_o.store([i_t * BT + 48, 16], b_Ai_42.to(desc_o.dtype, fp_downcast_rounding="rtne")) - desc_o.store([i_t * BT + 48, 32], b_Ai_43.to(desc_o.dtype, fp_downcast_rounding="rtne")) - - -@input_guard -def solve_tril( - A: torch.Tensor, - cu_seqlens: torch.Tensor | None = None, - chunk_indices: torch.LongTensor | None = None, - output_dtype: torch.dtype = torch.float, -) -> torch.Tensor: - """ - Compute the inverse of the matrix I + A - A should be strictly lower triangular, i.e., A.triu() == 0. - - Args: - A (torch.Tensor): - [B, T, H, BT], where BT should only be 16, 32, or 64. - cu_seqlens (torch.Tensor): - The cumulative sequence lengths of the input tensor. Default: `None`. - output_dtype (torch.dtype): - The dtype of the output tensor. Default: `torch.float`. - If `None`, the output dtype will be the same as the input dtype. - - Returns: - (I + A)^-1 with the same shape as A - """ - assert A.shape[-1] in [16, 32, 64] - output_dtype = A.dtype if output_dtype is None else output_dtype - - B, T, H, BT = A.shape - if chunk_indices is None and cu_seqlens is not None: - chunk_indices = prepare_chunk_indices(cu_seqlens, BT) - NT = len(chunk_indices) if cu_seqlens is not None else triton.cdiv(T, BT) - - Ai = torch.zeros_like(A, dtype=output_dtype) - if BT == 16: - merge_fn = solve_tril_16x16_kernel - elif BT == 32: - merge_fn = merge_16x16_to_32x32_inverse_kernel - elif BT == 64: - merge_fn = merge_16x16_to_64x64_inverse_kernel - - merge_fn[NT, B * H]( - A=A, - Ai=Ai, - cu_seqlens=cu_seqlens, - chunk_indices=chunk_indices, - T=T, - H=H, - BT=BT, - USE_TMA=IS_TMA_SUPPORTED, - ) - return Ai diff --git a/build/torch-xpu/ops/wall_attn/__init__.py b/build/torch-xpu/ops/wall_attn/__init__.py deleted file mode 100644 index 9011dfc779710cb0db41502a4400201602e2ae51..0000000000000000000000000000000000000000 --- a/build/torch-xpu/ops/wall_attn/__init__.py +++ /dev/null @@ -1,19 +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 -# -# Wall attention, contributed by Tilde Research (Timor Averbuch, Dhruv Pai). - -from .decode import build_wall_kv_cache, parallel_wall_attn_decode -from .naive import naive_wall_attn -from .parallel import parallel_wall_attn - -__all__ = [ - 'build_wall_kv_cache', - 'naive_wall_attn', - 'parallel_wall_attn', - 'parallel_wall_attn_decode', -] diff --git a/build/torch-xpu/ops/wall_attn/decode.py b/build/torch-xpu/ops/wall_attn/decode.py deleted file mode 100644 index 10a940e040cbf92db7fd21c215bb12bb76247b1c..0000000000000000000000000000000000000000 --- a/build/torch-xpu/ops/wall_attn/decode.py +++ /dev/null @@ -1,301 +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 -# -# Wall attention, contributed by Tilde Research (Timor Averbuch, Dhruv Pai). -# Heavily modified from fla/ops/gated_oja_rule/chunk.py. -"""Wall single-step decode kernel and pre-rescaled KV-cache builder.""" - -import torch -import triton -import triton.language as tl - -from ...ops.utils.op import exp2, log2 -from ...utils import check_shared_mem - -WALL_DECODE_AUTOTUNE_CONFIGS = [ - triton.Config({'BT': BT}, num_warps=nw, num_stages=ns) - for BT in (16, 32) - for nw in (2, 4, 8) - for ns in (2, 3) -] - - -@triton.autotune( - configs=WALL_DECODE_AUTOTUNE_CONFIGS, - key=['T_kv', 'K', 'V', 'HQ', 'H', 'C'], -) -@triton.heuristics({ - 'USE_SINK_BIAS': lambda args: args['sink_bias'] is not None, - 'USE_SCALAR_G': lambda args: args['g_scalar_cumsum'] is not None, -}) -@triton.jit -def parallel_wall_attn_decode_kernel( - q, # [B, T_q, HQ, K] - k_tilde, # [B, T_kv, HQ, K] pre-rescaled keys (per-Q-head) - v, # [B, T_kv, H, V] - o, # [B, T_q, HQ, V] - p_curr, # [B, T_q, HQ, K] - r_cache, # [B, NC, HQ, K] per-chunk anchors - g_scalar_cumsum, # [B, T_kv, HQ] or None (cache-side scalar cumsum; query-side at tail) - sink_bias, # [HQ] or None - lse, # [B, T_q, HQ] - scale, - T_q, - T_kv, - NC, - B: tl.constexpr, - H: tl.constexpr, - HQ: tl.constexpr, - G: tl.constexpr, - K: tl.constexpr, - V: tl.constexpr, - BT: tl.constexpr, - BS: tl.constexpr, - BK: tl.constexpr, - BV: tl.constexpr, - C: tl.constexpr, - USE_SINK_BIAS: tl.constexpr, - USE_SCALAR_G: tl.constexpr, -): - i_v, i_t, i_bh = tl.program_id(0), tl.program_id(1), tl.program_id(2) - i_b, i_hq = i_bh // HQ, i_bh % HQ - i_h = i_hq // G - RCP_LN2: tl.constexpr = 1.4426950216 - - bos_q = (i_b * T_q).to(tl.int64) - bos_kv = (i_b * T_kv).to(tl.int64) - bos_nc = (i_b * NC).to(tl.int64) - - p_q = tl.make_block_ptr( - q + (bos_q * HQ + i_hq) * K, (T_q, K), (HQ * K, 1), - (i_t * BT, 0), (BT, BK), (1, 0), - ) - p_pq = tl.make_block_ptr( - p_curr + (bos_q * HQ + i_hq) * K, (T_q, K), (HQ * K, 1), - (i_t * BT, 0), (BT, BK), (1, 0), - ) - p_o = tl.make_block_ptr( - o + (bos_q * HQ + i_hq) * V, (T_q, V), (HQ * V, 1), - (i_t * BT, i_v * BV), (BT, BV), (1, 0), - ) - p_lse = tl.make_block_ptr( - lse + bos_q * HQ + i_hq, (T_q,), (HQ,), (i_t * BT,), (BT,), (0,), - ) - - b_q = tl.load(p_q, boundary_check=(0, 1)) - b_pq = tl.load(p_pq, boundary_check=(0, 1)).to(tl.float32) - - if USE_SCALAR_G: - o_q_global = (T_kv - T_q) + i_t * BT + tl.arange(0, BT) - m_q = o_q_global < T_kv - b_cq = tl.load( - g_scalar_cumsum + (bos_kv + o_q_global) * HQ + i_hq, - mask=m_q, other=0, - ).to(tl.float32) - - if USE_SINK_BIAS: - b_sink_bias = tl.load(sink_bias + i_hq).to(tl.float32) * RCP_LN2 - - b_o = tl.zeros([BT, BV], dtype=tl.float32) - b_m = tl.full([BT], float('-inf'), dtype=tl.float32) - b_acc = tl.zeros([BT], dtype=tl.float32) - - # Iterate over the KV cache one chunk at a time: BS == C - for i_s in range(0, T_kv, BS): - i_c = i_s // C - # Load anchor - p_r = tl.make_block_ptr( - r_cache + (bos_nc * HQ + i_hq) * K, (NC, K), (HQ * K, 1), - (i_c, 0), (1, BK), (1, 0), - ) - b_R = tl.load(p_r, boundary_check=(0, 1)).to(tl.float32) - b_q_til = (b_q.to(tl.float32) * exp2(b_pq - b_R)).to(b_q.dtype) - - p_kt = tl.make_block_ptr( - k_tilde + (bos_kv * HQ + i_hq) * K, (K, T_kv), (1, HQ * K), - (0, i_s), (BK, BS), (0, 1), - ) - p_v = tl.make_block_ptr( - v + (bos_kv * H + i_h) * V, (T_kv, V), (H * V, 1), - (i_s, i_v * BV), (BS, BV), (1, 0), - ) - b_kt = tl.load(p_kt, boundary_check=(0, 1)) - b_v = tl.load(p_v, boundary_check=(0, 1)) - - b_s = tl.dot(b_q_til, b_kt) * scale * RCP_LN2 - - o_k = i_s + tl.arange(0, BS) - m_k = o_k < T_kv - if USE_SCALAR_G: - b_ck = tl.load( - g_scalar_cumsum + (bos_kv + o_k) * HQ + i_hq, - mask=m_k, other=0, - ).to(tl.float32) - b_s += b_cq[:, None] - b_ck[None, :] - - b_s = tl.where(m_k[None, :], b_s, float('-inf')) - - b_m, b_mp = tl.maximum(b_m, tl.max(b_s, 1)), b_m - b_mw = tl.where(b_m == float('-inf'), 0., b_m) - b_r = exp2(b_mp - b_mw) - b_p = exp2(b_s - b_mw[:, None]) - b_acc = b_acc * b_r + tl.sum(b_p, 1) - b_o = b_o * b_r[:, None] + tl.dot(b_p.to(b_v.dtype), b_v) - - if USE_SINK_BIAS: - b_m = tl.where(b_m == float('-inf'), 0., b_m) - b_acc += exp2(b_sink_bias - b_m) - - b_o = b_o / b_acc[:, None] - b_m += log2(b_acc) - tl.store(p_o, b_o.to(p_o.dtype.element_ty), boundary_check=(0, 1)) - tl.store(p_lse, b_m.to(p_lse.dtype.element_ty), boundary_check=(0,)) - - -def parallel_wall_attn_decode( - q: torch.Tensor, - v: torch.Tensor, - p_curr: torch.Tensor, - k_tilde: torch.Tensor, - r_cache: torch.Tensor, - sink_bias: torch.Tensor | None, - scale: float, - cache_chunk_size: int, - g_scalar_cumsum: torch.Tensor | None = None, -): - r"""Wall decode with a pre-rescaled KV cache. - - Use this directly from modeling code; it does NOT share a signature with - `parallel_wall_attn_fwd`. The caller is responsible for picking which to - call (e.g. training vs. cached generation). - - Shapes: - q : ``[B, T_q, HQ, K]`` current queries - v : ``[B, T_kv, H, V]`` cached values (standard layout) - p_curr : ``[B, T_q, HQ, K]`` per-channel prefix at the current rows - k_tilde : ``[B, T_kv, HQ, K]`` pre-rescaled keys (see `build_wall_kv_cache`) - r_cache : ``[B, NC, HQ, K]`` per-chunk anchors, ``NC = ceil(T_kv / C)`` - g_scalar_cumsum : ``[B, T_kv, HQ]`` or ``None`` optional FoX-style scalar gate - - Returns ``(o, lse)`` with ``o.shape == [B, T_q, HQ, V]``. - """ - C = cache_chunk_size - B, T_q, HQ, K = q.shape - _, T_kv, H, V = v.shape - NC = r_cache.shape[1] - G = HQ // H - - if k_tilde.shape != (B, T_kv, HQ, K): - raise ValueError(f"k_tilde shape {k_tilde.shape} != expected {(B, T_kv, HQ, K)}") - if r_cache.shape != (B, NC, HQ, K): - raise ValueError(f"r_cache shape {r_cache.shape} != expected {(B, NC, HQ, K)}") - if p_curr.shape != (B, T_q, HQ, K): - raise ValueError(f"p_curr shape {p_curr.shape} != expected {(B, T_q, HQ, K)}") - if T_kv > NC * C: - raise ValueError(f"NC*C ({NC * C}) < T_kv ({T_kv}); cache anchors do not cover cache") - if HQ % H != 0: - raise ValueError(f"HQ ({HQ}) must be divisible by H ({H})") - - for t in (q, v, p_curr, k_tilde, r_cache): - if t.stride(-1) != 1: - raise ValueError("decode tensors must be contiguous in the last (K or V) dim") - if g_scalar_cumsum is not None and g_scalar_cumsum.stride(-1) != 1: - raise ValueError("g_scalar_cumsum must be contiguous in the last dim") - - if check_shared_mem('hopper', q.device.index): - BK = min(256, max(16, triton.next_power_of_2(K))) - BV = min(256, max(16, triton.next_power_of_2(V))) - elif check_shared_mem('ampere', q.device.index): - BK = min(256, max(16, triton.next_power_of_2(K))) - BV = min(128, max(16, triton.next_power_of_2(V))) - else: - BK = min(256, max(16, triton.next_power_of_2(K))) - BV = min(64, max(16, triton.next_power_of_2(V))) - NV = triton.cdiv(V, BV) - assert triton.cdiv(K, BK) == 1, "decode kernel requires K <= 256" - - o = torch.empty(B, T_q, HQ, V, dtype=v.dtype, device=q.device) - lse = torch.empty(B, T_q, HQ, dtype=torch.float, device=q.device) - - def grid(meta): - return (NV, triton.cdiv(T_q, meta['BT']), B * HQ) - parallel_wall_attn_decode_kernel[grid]( - q=q, - k_tilde=k_tilde, - v=v, - o=o, - p_curr=p_curr, - r_cache=r_cache, - g_scalar_cumsum=g_scalar_cumsum, - sink_bias=sink_bias, - lse=lse, - scale=scale, - T_q=T_q, - T_kv=T_kv, - NC=NC, - B=B, - H=H, - HQ=HQ, - G=G, - K=K, - V=V, - BS=C, - BK=BK, - BV=BV, - C=C, - ) - return o, lse - - -def build_wall_kv_cache( - k: torch.Tensor, # [B, T, H, K] - g_cumsum: torch.Tensor, # [B, T, HQ, K] per-channel prefix P = cumsum(log_2 g) - chunk_size: int, - *, - out_dtype: torch.dtype | None = None, -) -> tuple[torch.Tensor, torch.Tensor]: - r"""Build a pre-rescaled Wall KV cache. - - Returns ``(k_tilde, r_cache)`` where, for each chunk ``c`` of size ``C``: - ``R_c = P[chunk_start_c]`` ``[B, HQ, K]`` - ``k_tilde[j] = k[j] * exp2(R_{c(j)} - P[j])`` ``[B, T, HQ, K]`` - ``r_cache[c] = R_c`` ``[B, NC, HQ, K]`` - - ``k_tilde`` is per-Q-head (carries the gate). Values are unchanged and - should be cached as-is. ``g_cumsum`` is the same prefix tensor passed to - the training forward. - """ - if k.dim() != 4: - raise ValueError("k must be [B, T, H, K]") - if g_cumsum.dim() != 4: - raise ValueError("g_cumsum must be [B, T, HQ, K]") - B, T, H, K = k.shape - _, _, HQ, _ = g_cumsum.shape - if g_cumsum.shape[:2] != (B, T) or g_cumsum.shape[-1] != K: - raise ValueError(f"g_cumsum shape {g_cumsum.shape} incompatible with k {k.shape}") - if HQ % H != 0: - raise ValueError(f"HQ ({HQ}) must be divisible by H ({H})") - G = HQ // H - NC = triton.cdiv(T, chunk_size) - out_dtype = out_dtype if out_dtype is not None else k.dtype - - # Per-Q-head broadcast of k along the HQ axis (GQA: each KV head feeds G query heads). - k_q = k.repeat_interleave(G, dim=2) # [B, T, HQ, K] - - # Per-chunk reference R_c = P[chunk_start, :] for each batch / head. - # Pad T up to NC*chunk_size with the last row (rescaler will be exp2(0)=1 for tail). - chunk_starts = torch.arange(NC, device=k.device) * chunk_size - chunk_starts = chunk_starts.clamp(max=T - 1) - r_cache = g_cumsum[:, chunk_starts, :, :].contiguous() # [B, NC, HQ, K] - - # Per-token R(j) = R_{c(j)} - token_chunk = torch.arange(T, device=k.device) // chunk_size # [T] - r_per_token = r_cache[:, token_chunk, :, :] # [B, T, HQ, K] - - k_tilde = (k_q.to(torch.float32) * torch.exp2(r_per_token.to(torch.float32) - g_cumsum.to(torch.float32))) - - return k_tilde.to(out_dtype).contiguous(), r_cache diff --git a/build/torch-xpu/ops/wall_attn/naive.py b/build/torch-xpu/ops/wall_attn/naive.py deleted file mode 100644 index 47a932f482b5cb11b42751b4335b2775ca66f8ae..0000000000000000000000000000000000000000 --- a/build/torch-xpu/ops/wall_attn/naive.py +++ /dev/null @@ -1,81 +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 -# -# Wall attention, contributed by Tilde Research (Timor Averbuch, Dhruv Pai). -"""Eager PyTorch reference for Wall attention (correctness oracle).""" - -import torch - -from ...ops.utils.constant import RCP_LN2 -from ...ops.utils.cumsum import chunk_global_cumsum - - -def naive_wall_attn( - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - g: torch.Tensor, - *, - scale: float, - window_size: int | None = None, - cu_seqlens: torch.LongTensor | None = None, - sink_bias: torch.Tensor | None = None, - g_scalar: torch.Tensor | None = None, -) -> torch.Tensor: - r"""Eager Wall forward (exact logits): per pair :math:`(i,j)`, - :math:`s_{ij} = \sum_n q_{in} k_{jn} \exp_2(P_{in}-P_{jn}) \cdot \mathrm{scale}\cdot\mathrm{RCP\_LN2}` - with :math:`P` the prefix cumsum of ``g`` in log_2 (same as the Triton path).""" - if g.dim() != 4 or g.shape != (*q.shape[:-1], q.shape[-1]): - raise ValueError("`g` must be [B, T, HQ, K] matching `q`") - G = q.shape[2] // k.shape[2] - if q.shape[2] % k.shape[2] != 0: - raise ValueError("HQ must be divisible by H (GQA)") - - P = chunk_global_cumsum(g, cu_seqlens=cu_seqlens, scale=RCP_LN2) - k_exp = k.repeat_interleave(G, dim=2) - B, T, HQ, _K = q.shape - diff = P.unsqueeze(2).float() - P.unsqueeze(1).float() - scores = ( - q.unsqueeze(2).float() * k_exp.unsqueeze(1).float() * torch.exp2(diff) - ).sum(-1).permute(0, 3, 1, 2).contiguous() - - i_idx = torch.arange(T, device=q.device, dtype=torch.long).view(1, T, 1) - j_idx = torch.arange(T, device=q.device, dtype=torch.long).view(1, 1, T) - valid = j_idx <= i_idx - if window_size is not None: - valid = valid & (i_idx - j_idx < window_size) - if cu_seqlens is not None: - if B != 1: - raise ValueError("reference varlen expects batch size 1") - seg = torch.zeros(T, dtype=torch.long, device=q.device) - for s in range(int(cu_seqlens.numel()) - 1): - a = int(cu_seqlens[s].item()) - b = int(cu_seqlens[s + 1].item()) - seg[a:b] = s - seg_i = seg.view(1, T, 1) - seg_j = seg.view(1, 1, T) - valid = valid & (seg_i == seg_j) - - valid_hq = valid.unsqueeze(1).expand(B, HQ, T, T) - scores = scores.masked_fill(~valid_hq, float("-inf")) * (scale * RCP_LN2) - - if g_scalar is not None: - c = chunk_global_cumsum(g_scalar, cu_seqlens=cu_seqlens, scale=RCP_LN2) - c_hq = c.permute(0, 2, 1).float() - scores = scores + c_hq.unsqueeze(-1) - c_hq.unsqueeze(-2) - - m = scores.amax(dim=-1, keepdim=True) - m_stable = torch.where(torch.isfinite(m), m, torch.zeros_like(m)) - p = torch.exp2(scores - m_stable) - den = p.sum(dim=-1) - if sink_bias is not None: - sink_l2 = (sink_bias * RCP_LN2).view(1, HQ, 1) - den = den + torch.exp2(sink_l2 - m_stable.squeeze(-1)) - w = p / den.unsqueeze(-1) - v_h = v.repeat_interleave(G, dim=2).permute(0, 2, 1, 3).float().contiguous() - o = torch.matmul(w, v_h).transpose(1, 2).contiguous() - return o diff --git a/build/torch-xpu/ops/wall_attn/parallel.py b/build/torch-xpu/ops/wall_attn/parallel.py deleted file mode 100644 index 9443838331279604f6e1c57e836a8f20c0b3a834..0000000000000000000000000000000000000000 --- a/build/torch-xpu/ops/wall_attn/parallel.py +++ /dev/null @@ -1,971 +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 -# -# Wall attention, contributed by Tilde Research (Timor Averbuch, Dhruv Pai). -# Heavily modified from fla/ops/gated_oja_rule/chunk.py. -"""Wall training/prefill kernels: forward, backward, and the autograd Function.""" - -import os - -import torch -import triton -import triton.language as tl -from einops import reduce - -from ...ops.attn.parallel import parallel_attn_bwd_preprocess -from ...ops.backends import dispatch -from ...ops.utils import prepare_chunk_indices -from ...ops.utils.constant import RCP_LN2 -from ...ops.utils.cumsum import chunk_global_cumsum -from ...ops.utils.op import exp2, log2 -from ...utils import autocast_custom_bwd, autocast_custom_fwd, check_shared_mem, contiguous - -_DEBUG_ASSERTS = os.environ.get("WALL_ATTN_DEBUG", "0") == "1" - -# Set WALL_ATTN_DKV_DIAG_BF16=0 to force fp32 tensor cores in the diagonal loop -# of parallel_wall_attn_bwd_kernel_dkv (precision over speed). -_DKV_DIAG_BF16 = os.environ.get("WALL_ATTN_DKV_DIAG_BF16", "1") == "1" - -WALL_FWD_AUTOTUNE_CONFIGS = [ - triton.Config({'BT': BT, 'BS': BS}, num_warps=nw, num_stages=ns) - for BT in (64, 128) - for BS in (32, 64) - for nw in (2, 4, 8) - for ns in (2, 3) -] - -WALL_BWD_AUTOTUNE_CONFIGS = [ - triton.Config({'BT': BT, 'BS': BS}, num_warps=nw, num_stages=ns) - for BT in (64, 128) - for BS in (32, 64) - for nw in (2, 4, 8) - for ns in (2, 3) -] - - -@triton.autotune( - configs=WALL_FWD_AUTOTUNE_CONFIGS, - key=['T', 'K', 'V', 'HQ', 'H'], -) -@triton.heuristics({ - 'USE_SINK_BIAS': lambda args: args['sink_bias'] is not None, - 'USE_WINDOW': lambda args: args['W'] is not None, - 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None, - 'USE_SCALAR_G': lambda args: args['g_scalar_cumsum'] is not None, -}) -@triton.jit -def parallel_wall_attn_fwd_kernel( - q, - k, - v, - o, - g_cumsum, - g_scalar_cumsum, - sink_bias, - lse, - scale, - cu_seqlens, - chunk_indices, - T, - W: tl.constexpr, - B: tl.constexpr, - H: tl.constexpr, - HQ: tl.constexpr, - G: tl.constexpr, - K: tl.constexpr, - V: tl.constexpr, - BT: tl.constexpr, - BS: tl.constexpr, - BK: tl.constexpr, - BV: tl.constexpr, - USE_SINK_BIAS: tl.constexpr, - USE_WINDOW: tl.constexpr, - IS_VARLEN: tl.constexpr, - USE_SCALAR_G: tl.constexpr, -): - i_v, i_t, i_bh = tl.program_id(0), tl.program_id(1), tl.program_id(2) - i_b, i_hq = i_bh // HQ, i_bh % HQ - i_h = i_hq // G - - 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).to(tl.int32) - else: - i_n = i_b - bos, eos = (i_b * T).to(tl.int64), (i_b * T + T).to(tl.int64) - RCP_LN2: tl.constexpr = 1.4426950216 - - p_q = tl.make_block_ptr(q + (bos * HQ + i_hq) * K, (T, K), (HQ*K, 1), (i_t * BT, 0), (BT, BK), (1, 0)) - p_o = tl.make_block_ptr(o + (bos * HQ + i_hq) * V, (T, V), (HQ*V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) - p_lse = tl.make_block_ptr(lse + bos * HQ + i_hq, (T,), (HQ,), (i_t * BT,), (BT,), (0,)) - - p_pq = tl.make_block_ptr(g_cumsum + (bos * HQ + i_hq) * K, (T, K), (HQ*K, 1), (i_t * BT, 0), (BT, BK), (1, 0)) - p_R = tl.make_block_ptr(g_cumsum + (bos * HQ + i_hq) * K, (T, K), (HQ*K, 1), (i_t * BT, 0), (1, BK), (1, 0)) - - b_q = tl.load(p_q, boundary_check=(0, 1)) - b_pq = tl.load(p_pq, boundary_check=(0, 1)).to(tl.float32) - b_R = tl.load(p_R, boundary_check=(0, 1)).to(tl.float32) - # b_R is at i_t*BT; the same value is used in both off-diag and diag loops, - # since |b_pq - b_R| within a BT-chunk is bounded by BT*|g_max|*RCP_LN2. - # Off-diagonal: P_q - R <= 0 and R - P_k <= 0 -> exp2 <= 1 -> bf16 safe. - b_q_til = (b_q.to(tl.float32) * exp2(b_pq - b_R)).to(b_q.dtype) - - b_o = tl.zeros([BT, BV], dtype=tl.float32) - b_m = tl.full([BT], float('-inf'), dtype=tl.float32) - b_acc = tl.zeros([BT], dtype=tl.float32) - - if USE_SCALAR_G: - p_cq = tl.make_block_ptr(g_scalar_cumsum + bos * HQ + i_hq, (T,), (HQ,), (i_t * BT,), (BT,), (0,)) - b_cq = tl.load(p_cq, boundary_check=(0,)).to(tl.float32) - - if USE_SINK_BIAS: - b_sink_bias = tl.load(sink_bias + i_hq).to(tl.float32) - else: - b_sink_bias = None - - o_q = i_t * BT + tl.arange(0, BT) - i_start = tl.maximum((i_t * BT - W + 1) // BS * BS, 0) if USE_WINDOW else 0 - - b_R_t = tl.trans(b_R) - - for i_s in range(i_start, i_t * BT, BS): - p_k = tl.make_block_ptr(k + (bos * H + i_h) * K, (K, T), (1, H*K), (0, i_s), (BK, BS), (0, 1)) - p_v = tl.make_block_ptr(v + (bos * H + i_h) * V, (T, V), (H*V, 1), (i_s, i_v * BV), (BS, BV), (1, 0)) - p_pk = tl.make_block_ptr(g_cumsum + (bos * HQ + i_hq) * K, (K, T), (1, HQ*K), (0, i_s), (BK, BS), (0, 1)) - - b_k = tl.load(p_k, boundary_check=(0, 1)) - b_v = tl.load(p_v, boundary_check=(0, 1)) - b_pk = tl.load(p_pk, boundary_check=(0, 1)).to(tl.float32) - b_k_til = (b_k.to(tl.float32) * exp2(b_R_t - b_pk)).to(b_k.dtype) - b_s = tl.dot(b_q_til, b_k_til) * scale * RCP_LN2 - - o_k = i_s + tl.arange(0, BS) - m_k = o_k < T - if USE_SCALAR_G: - b_ck = tl.load(g_scalar_cumsum + (bos + o_k) * HQ + i_hq, mask=m_k, other=0).to(tl.float32) - b_s += b_cq[:, None] - b_ck[None, :] - if USE_WINDOW: - b_s = tl.where((o_q[:, None] - o_k[None, :] < W) & m_k[None, :], b_s, float('-inf')) - - b_m, b_mp = tl.maximum(b_m, tl.max(b_s, 1)), b_m - b_mw = tl.where(b_m == float('-inf'), 0., b_m) - b_r = exp2(b_mp - b_mw) - b_p = exp2(b_s - b_mw[:, None]) - b_acc = b_acc * b_r + tl.sum(b_p, 1) - b_o = b_o * b_r[:, None] + tl.dot(b_p.to(b_v.dtype), b_v) - b_mp = b_m - - for i_s in range(i_t * BT, min((i_t + 1) * BT, T), BS): - p_k = tl.make_block_ptr(k + (bos * H + i_h) * K, (K, T), (1, H*K), (0, i_s), (BK, BS), (0, 1)) - p_v = tl.make_block_ptr(v + (bos * H + i_h) * V, (T, V), (H*V, 1), (i_s, i_v * BV), (BS, BV), (1, 0)) - p_pk = tl.make_block_ptr(g_cumsum + (bos * HQ + i_hq) * K, (K, T), (1, HQ*K), (0, i_s), (BK, BS), (0, 1)) - - # Per-sub-block local reference to prevent exp2 overflow. - p_R_local = tl.make_block_ptr(g_cumsum + (bos * HQ + i_hq) * K, (T, K), (HQ*K, 1), (i_s, 0), (1, BK), (1, 0)) - b_R_local = tl.load(p_R_local, boundary_check=(0, 1)).to(tl.float32) - b_R_local_bc = tl.broadcast_to(b_R_local, (BT, BK)) - b_R_local_t = tl.trans(b_R_local) - - o_k = i_s + tl.arange(0, BS) - m_k = o_k < T - b_k = tl.load(p_k, boundary_check=(0, 1)) - b_v = tl.load(p_v, boundary_check=(0, 1)) - b_pk = tl.load(p_pk, boundary_check=(0, 1)).to(tl.float32) - - # k_til in local frame: exp2 bounded by BS positions <= 110 (fp32-safe with BK dot). - b_exp_k = b_R_local_t - b_pk - b_exp_k = tl.where(b_exp_k > 110.0, tl.zeros_like(b_exp_k) + 110.0, b_exp_k) - b_k_til = b_k.to(tl.float32) * exp2(b_exp_k) - - # q_til in local frame; clamp exp for Q before sub-block (masked by causality). - b_exp_q = tl.minimum(b_pq - b_R_local_bc, tl.zeros([BT, BK], dtype=tl.float32)) - b_q_til_local = b_q.to(tl.float32) * exp2(b_exp_q) - b_s = tl.dot(b_q_til_local, b_k_til) * scale * RCP_LN2 - - if USE_SCALAR_G: - b_ck = tl.load(g_scalar_cumsum + (bos + o_k) * HQ + i_hq, mask=m_k, other=0).to(tl.float32) - b_s += b_cq[:, None] - b_ck[None, :] - - m_s = (o_q[:, None] >= o_k[None, :]) & m_k[None, :] - if USE_WINDOW: - m_s = m_s & (o_q[:, None] - o_k[None, :] < W) - b_s = tl.where(m_s, b_s, float('-inf')) - - b_m, b_mp = tl.maximum(b_m, tl.max(b_s, 1)), b_m - b_mw = tl.where(b_m == float('-inf'), 0., b_m) - b_r = exp2(b_mp - b_mw) - b_p = exp2(b_s - b_mw[:, None]) - b_acc = b_acc * b_r + tl.sum(b_p, 1) - b_o = b_o * b_r[:, None] + tl.dot(b_p.to(b_v.dtype), b_v) - b_mp = b_m - - if USE_SINK_BIAS: - b_m = tl.where(b_m == float('-inf'), 0., b_m) - b_acc += exp2(b_sink_bias - b_m) - - b_o = b_o / b_acc[:, None] - b_m += log2(b_acc) - tl.store(p_o, b_o.to(p_o.dtype.element_ty), boundary_check=(0, 1)) - tl.store(p_lse, b_m.to(p_lse.dtype.element_ty), boundary_check=(0,)) - - -@triton.autotune( - configs=WALL_BWD_AUTOTUNE_CONFIGS, - key=['T', 'K', 'V', 'HQ', 'H'], -) -@triton.heuristics({ - 'USE_WINDOW': lambda args: args['W'] is not None, - 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None, - 'USE_SCALAR_G': lambda args: args['g_scalar_cumsum'] is not None, -}) -@triton.jit(do_not_specialize=['T']) -def parallel_wall_attn_bwd_kernel_dq( - q, - k, - v, - g_cumsum, - g_scalar_cumsum, - lse, - delta, - do, - dq, - dg_cumsum, - dg_scalar_cumsum, - scale, - cu_seqlens, - chunk_indices, - T, - W: tl.constexpr, - B: tl.constexpr, - H: tl.constexpr, - HQ: tl.constexpr, - G: tl.constexpr, - K: tl.constexpr, - V: tl.constexpr, - BT: tl.constexpr, - BS: tl.constexpr, - BK: tl.constexpr, - BV: tl.constexpr, - USE_WINDOW: tl.constexpr, - IS_VARLEN: tl.constexpr, - USE_SCALAR_G: tl.constexpr, -): - i_v, i_t, i_bh = tl.program_id(0), tl.program_id(1), tl.program_id(2) - i_b, i_hq = i_bh // HQ, i_bh % HQ - i_h = i_hq // G - - 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).to(tl.int32) - else: - i_n = i_b - bos, eos = (i_b * T).to(tl.int64), (i_b * T + T).to(tl.int64) - RCP_LN2: tl.constexpr = 1.4426950216 - LN2: tl.constexpr = 0.6931471805599453 - - p_q = tl.make_block_ptr(q + (bos * HQ + i_hq) * K, (T, K), (HQ*K, 1), (i_t * BT, 0), (BT, BK), (1, 0)) - p_pq = tl.make_block_ptr(g_cumsum + (bos * HQ + i_hq) * K, (T, K), (HQ*K, 1), (i_t * BT, 0), (BT, BK), (1, 0)) - p_R = tl.make_block_ptr(g_cumsum + (bos * HQ + i_hq) * K, (T, K), (HQ*K, 1), (i_t * BT, 0), (1, BK), (1, 0)) - p_dq = tl.make_block_ptr(dq + (bos * HQ + i_hq) * K, (T, K), (HQ*K, 1), (i_t * BT, 0), (BT, BK), (1, 0)) - p_dg = tl.make_block_ptr(dg_cumsum + (bos * HQ + i_hq) * K, (T, K), (HQ*K, 1), (i_t * BT, 0), (BT, BK), (1, 0)) - p_do = tl.make_block_ptr(do + (bos * HQ + i_hq) * V, (T, V), (HQ*V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) - p_lse = tl.make_block_ptr(lse + bos * HQ + i_hq, (T,), (HQ,), (i_t * BT,), (BT,), (0,)) - p_delta = tl.make_block_ptr(delta + bos * HQ + i_hq, (T,), (HQ,), (i_t * BT,), (BT,), (0,)) - - b_q = tl.load(p_q, boundary_check=(0, 1)) - b_pq = tl.load(p_pq, boundary_check=(0, 1)).to(tl.float32) - b_R = tl.load(p_R, boundary_check=(0, 1)).to(tl.float32) - # Off-diagonal q_til (bf16 tensor cores, exp2 <= 1). - b_q_til = (b_q.to(tl.float32) * exp2(b_pq - b_R)).to(b_q.dtype) - - b_do = tl.load(p_do, boundary_check=(0, 1)) - b_lse = tl.load(p_lse, boundary_check=(0,)) - b_delta = tl.load(p_delta, boundary_check=(0,)) - - b_dq_til = tl.zeros([BT, BK], dtype=tl.float32) - - if USE_SCALAR_G: - p_cq = tl.make_block_ptr(g_scalar_cumsum + bos * HQ + i_hq, (T,), (HQ,), (i_t * BT,), (BT,), (0,)) - b_cq = tl.load(p_cq, boundary_check=(0,)).to(tl.float32) - b_dc = tl.zeros([BT], dtype=tl.float32) - - o_q = i_t * BT + tl.arange(0, BT) - i_start = tl.maximum((i_t * BT - W + 1) // BS * BS, 0) if USE_WINDOW else 0 - b_R_t = tl.trans(b_R) - - for i_s in range(i_start, i_t * BT, BS): - p_k = tl.make_block_ptr(k + (bos * H + i_h) * K, (K, T), (1, H*K), (0, i_s), (BK, BS), (0, 1)) - p_v = tl.make_block_ptr(v + (bos * H + i_h) * V, (V, T), (1, H*V), (i_v * BV, i_s), (BV, BS), (0, 1)) - p_pk = tl.make_block_ptr(g_cumsum + (bos * HQ + i_hq) * K, (K, T), (1, HQ*K), (0, i_s), (BK, BS), (0, 1)) - - o_k = i_s + tl.arange(0, BS) - m_k = o_k < T - b_k = tl.load(p_k, boundary_check=(0, 1)) - b_pk = tl.load(p_pk, boundary_check=(0, 1)).to(tl.float32) - b_v = tl.load(p_v, boundary_check=(0, 1)) - b_k_til = (b_k.to(tl.float32) * exp2(b_R_t - b_pk)).to(b_k.dtype) - b_s = tl.dot(b_q_til, b_k_til) * scale * RCP_LN2 - - if USE_SCALAR_G: - b_ck = tl.load(g_scalar_cumsum + (bos + o_k) * HQ + i_hq, mask=m_k, other=0).to(tl.float32) - b_s += b_cq[:, None] - b_ck[None, :] - if USE_WINDOW: - b_s = tl.where((o_q[:, None] - o_k[None, :] < W) & m_k[None, :], b_s, float('-inf')) - b_p = exp2(b_s - b_lse[:, None]) - b_dp = tl.dot(b_do, b_v) - b_ds = b_p * (b_dp.to(tl.float32) - b_delta[:, None]) - b_dq_til += tl.dot(b_ds.to(b_k.dtype), tl.trans(b_k_til)) - if USE_SCALAR_G: - b_dc += tl.sum(b_ds, 1) - - # Diagonal: use per-sub-block local reference to prevent exp2 overflow. - # Accumulate dq_diag directly in output space to avoid large frame corrections. - b_dq_diag = tl.zeros([BT, BK], dtype=tl.float32) - - for i_s in range(i_t * BT, min((i_t + 1) * BT, T), BS): - p_k = tl.make_block_ptr(k + (bos * H + i_h) * K, (K, T), (1, H*K), (0, i_s), (BK, BS), (0, 1)) - p_v = tl.make_block_ptr(v + (bos * H + i_h) * V, (V, T), (1, H*V), (i_v * BV, i_s), (BV, BS), (0, 1)) - p_pk = tl.make_block_ptr(g_cumsum + (bos * HQ + i_hq) * K, (K, T), (1, HQ*K), (0, i_s), (BK, BS), (0, 1)) - - p_R_local = tl.make_block_ptr(g_cumsum + (bos * HQ + i_hq) * K, (T, K), (HQ*K, 1), (i_s, 0), (1, BK), (1, 0)) - b_R_local = tl.load(p_R_local, boundary_check=(0, 1)).to(tl.float32) - b_R_local_bc = tl.broadcast_to(b_R_local, (BT, BK)) - b_R_local_t = tl.trans(b_R_local) - - o_k = i_s + tl.arange(0, BS) - m_k = o_k < T - b_k = tl.load(p_k, boundary_check=(0, 1)) - b_pk = tl.load(p_pk, boundary_check=(0, 1)).to(tl.float32) - b_v = tl.load(p_v, boundary_check=(0, 1)) - - # Local-ref k_til: exp2 arg bounded by BS positions <= 110 (fp32-safe with BK dot). - b_exp_k = b_R_local_t - b_pk - b_exp_k = tl.where(b_exp_k > 110.0, tl.zeros_like(b_exp_k) + 110.0, b_exp_k) - b_k_til = b_k.to(tl.float32) * exp2(b_exp_k) - - # q_til in local frame for computing scores. - b_q_til_local = b_q.to(tl.float32) * exp2(b_pq - b_R_local_bc) - b_s = tl.dot(b_q_til_local, b_k_til) * scale * RCP_LN2 - - if USE_SCALAR_G: - b_ck = tl.load(g_scalar_cumsum + (bos + o_k) * HQ + i_hq, mask=m_k, other=0).to(tl.float32) - b_s += b_cq[:, None] - b_ck[None, :] - - if USE_WINDOW: - b_p = tl.where( - (o_q[:, None] >= o_k[None, :]) & (o_q[:, None] - o_k[None, :] < W) & m_k[None, :], - exp2(b_s - b_lse[:, None]), 0 - ) - else: - b_p = tl.where((o_q[:, None] >= o_k[None, :]) & m_k[None, :], exp2(b_s - b_lse[:, None]), 0) - - b_dp = tl.dot(b_do, b_v) - b_ds = b_p * (b_dp.to(tl.float32) - b_delta[:, None]) - # dq_sub in local frame; convert to output space: dq += dq_sub * scale * exp2(Pq - R_local). - # For Q < sub-block start: causal mask makes ds=0, but mask explicitly to kill fp noise. - b_dq_sub = tl.dot(b_ds.to(tl.float32), tl.trans(b_k_til)) - m_q_causal = (o_q >= i_s)[:, None] - b_dq_diag += tl.where(m_q_causal, b_dq_sub * exp2(b_pq - b_R_local_bc), 0.0) - if USE_SCALAR_G: - b_dc += tl.sum(b_ds, 1) - - # Off-diagonal (chunk-ref) + diagonal (direct output space). - b_dq = (b_dq_til * scale) * exp2(b_pq - b_R) + b_dq_diag * scale - - # Structural: b_dg derived from b_dq (saves [BT,BK] fp32 accumulator). - b_dg = LN2 * b_q.to(tl.float32) * b_dq - tl.store(p_dq, b_dq.to(p_dq.dtype.element_ty), boundary_check=(0, 1)) - tl.store(p_dg, b_dg.to(p_dg.dtype.element_ty), boundary_check=(0, 1)) - if USE_SCALAR_G: - p_dc = tl.make_block_ptr(dg_scalar_cumsum + bos * HQ + i_hq, (T,), (HQ,), (i_t * BT,), (BT,), (0,)) - tl.store(p_dc, b_dc.to(p_dc.dtype.element_ty), boundary_check=(0,)) - - -@triton.autotune( - configs=WALL_BWD_AUTOTUNE_CONFIGS, - key=['T', 'K', 'V', 'HQ', 'H'], -) -@triton.heuristics({ - 'USE_WINDOW': lambda args: args['W'] is not None, - 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None, - 'USE_SCALAR_G': lambda args: args['g_scalar_cumsum'] is not None, -}) -@triton.jit(do_not_specialize=['T']) -def parallel_wall_attn_bwd_kernel_dkv( - q, - k, - v, - g_cumsum, - g_scalar_cumsum, - lse, - delta, - do, - dk, - dv, - dg_cumsum, - dg_scalar_cumsum, - cu_seqlens, - chunk_indices, - scale, - T, - W: tl.constexpr, - B: tl.constexpr, - H: tl.constexpr, - HQ: tl.constexpr, - G: tl.constexpr, - K: tl.constexpr, - V: tl.constexpr, - BT: tl.constexpr, - BS: tl.constexpr, - BK: tl.constexpr, - BV: tl.constexpr, - USE_WINDOW: tl.constexpr, - IS_VARLEN: tl.constexpr, - USE_SCALAR_G: tl.constexpr, - DIAG_BF16: tl.constexpr, -): - i_v, i_t, i_bh = tl.program_id(0), tl.program_id(1), tl.program_id(2) - i_b, i_hq = i_bh // HQ, i_bh % HQ - i_h = i_hq // G - - 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).to(tl.int32) - else: - i_n = i_b - bos, eos = (i_b * T).to(tl.int64), (i_b * T + T).to(tl.int64) - RCP_LN2: tl.constexpr = 1.4426950216 - LN2: tl.constexpr = 0.6931471805599453 - - p_k = tl.make_block_ptr(k + (bos * H + i_h) * K, (T, K), (H*K, 1), (i_t * BT, 0), (BT, BK), (1, 0)) - p_pk = tl.make_block_ptr(g_cumsum + (bos * HQ + i_hq) * K, (T, K), (HQ*K, 1), (i_t * BT, 0), (BT, BK), (1, 0)) - p_v = tl.make_block_ptr(v + (bos * H + i_h) * V, (T, V), (H*V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) - p_dk = tl.make_block_ptr(dk + (bos * HQ + i_hq) * K, (T, K), (HQ*K, 1), (i_t * BT, 0), (BT, BK), (1, 0)) - p_dv = tl.make_block_ptr(dv + (bos * HQ + i_hq) * V, (T, V), (HQ*V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)) - p_dg = tl.make_block_ptr(dg_cumsum + (bos * HQ + i_hq) * K, (T, K), (HQ*K, 1), (i_t * BT, 0), (BT, BK), (1, 0)) - - b_k = tl.load(p_k, boundary_check=(0, 1)) - b_pk = tl.load(p_pk, boundary_check=(0, 1)).to(tl.float32) - b_dk = tl.zeros([BT, BK], dtype=tl.float32) - b_v = tl.load(p_v, boundary_check=(0, 1)) - b_dv = tl.zeros([BT, BV], dtype=tl.float32) - - o_k = i_t * BT + tl.arange(0, BT) - - if USE_SCALAR_G: - p_ck = tl.make_block_ptr(g_scalar_cumsum + bos * HQ + i_hq, (T,), (HQ,), (i_t * BT,), (BT,), (0,)) - b_ck = tl.load(p_ck, boundary_check=(0,)).to(tl.float32) - b_dc = tl.zeros([BT], dtype=tl.float32) - - for i_s in range(i_t * BT, min((i_t + 1) * BT, T), BS): - p_Rq = tl.make_block_ptr(g_cumsum + (bos * HQ + i_hq) * K, (T, K), (HQ*K, 1), (i_s, 0), (1, BK), (1, 0)) - p_q = tl.make_block_ptr(q + (bos * HQ + i_hq) * K, (T, K), (HQ*K, 1), (i_s, 0), (BS, BK), (1, 0)) - p_pq = tl.make_block_ptr(g_cumsum + (bos * HQ + i_hq) * K, (T, K), (HQ*K, 1), (i_s, 0), (BS, BK), (1, 0)) - p_do = tl.make_block_ptr(do + (bos * HQ + i_hq) * V, (T, V), (HQ*V, 1), (i_s, i_v * BV), (BS, BV), (1, 0)) - p_lse = tl.make_block_ptr(lse + bos * HQ + i_hq, (T,), (HQ,), (i_s,), (BS,), (0,)) - p_delta = tl.make_block_ptr(delta + bos * HQ + i_hq, (T,), (HQ,), (i_s,), (BS,), (0,)) - - o_q = i_s + tl.arange(0, BS) - m_q = o_q < T - b_Rq = tl.load(p_Rq, boundary_check=(0, 1)).to(tl.float32) - b_Rq_bc_k = tl.broadcast_to(b_Rq, (BT, BK)) - b_Rq_bc_q = tl.broadcast_to(b_Rq, (BS, BK)) - b_q = tl.load(p_q, boundary_check=(0, 1)) - b_pq = tl.load(p_pq, boundary_check=(0, 1)).to(tl.float32) - b_exp_k = b_Rq_bc_k - b_pk - # Non-causal keys (key_pos > last query in sub-block) can have - # exp_k >> 110, overflowing the gradient headroom. The causal mask - # zeros their contribution in b_p, but inf*0 = NaN poisons b_dk/b_dg. - # Clamp to 110 (fp32-safe with BK/BS dot accumulation) so products - # in dg accumulation stay well within fp32. - b_exp_k = tl.where(b_exp_k > 110.0, tl.zeros_like(b_exp_k) + 110.0, b_exp_k) - b_exp_k_val = exp2(b_exp_k) # CSE: reused for k_til and final dk scaling - if DIAG_BF16: - # Fast path: bf16 tensor cores, fp32 accumulation. - b_q_til = (b_q.to(tl.float32) * exp2(b_pq - b_Rq_bc_q)).to(b_q.dtype) - b_k_til = (b_k.to(tl.float32) * b_exp_k_val).to(b_k.dtype) - else: - # Precise path: fp32 tensor cores. - b_q_til = b_q.to(tl.float32) * exp2(b_pq - b_Rq_bc_q) - b_k_til = b_k.to(tl.float32) * b_exp_k_val - b_do = tl.load(p_do, boundary_check=(0, 1)) - b_lse = tl.load(p_lse, boundary_check=(0,)) - b_delta = tl.load(p_delta, boundary_check=(0,)) - b_s = tl.dot(b_k_til, tl.trans(b_q_til)) * scale * RCP_LN2 - if USE_SCALAR_G: - b_cq = tl.load(g_scalar_cumsum + (bos + o_q) * HQ + i_hq, mask=m_q, other=0).to(tl.float32) - b_s += b_cq[None, :] - b_ck[:, None] - if USE_WINDOW: - b_p = tl.where( - (o_k[:, None] <= o_q[None, :]) & (o_q[None, :] - o_k[:, None] < W) & m_q[None, :], - exp2(b_s - b_lse[None, :]), 0 - ) - else: - b_p = tl.where((o_k[:, None] <= o_q[None, :]) & m_q[None, :], exp2(b_s - b_lse[None, :]), 0) - b_dv += tl.dot(b_p.to(b_do.dtype), b_do) - b_dp = tl.dot(b_v, tl.trans(b_do)) - b_ds = b_p * (b_dp - b_delta[None, :]) - if DIAG_BF16: - b_dk_til = tl.dot(b_ds.to(b_q.dtype), b_q_til) - else: - b_dk_til = tl.dot(b_ds.to(tl.float32), b_q_til) - b_dk += (b_dk_til * scale) * b_exp_k_val - if USE_SCALAR_G: - b_dc -= tl.sum(b_ds, 1) - - i_end = min(tl.cdiv(T, BS) * BS, (i_t + 1) * BT + W - 1) if USE_WINDOW else tl.cdiv(T, BS) * BS - - for i_s in range((i_t + 1) * BT, i_end, BS): - p_Rq = tl.make_block_ptr(g_cumsum + (bos * HQ + i_hq) * K, (T, K), (HQ*K, 1), (i_s, 0), (1, BK), (1, 0)) - p_q = tl.make_block_ptr(q + (bos * HQ + i_hq) * K, (T, K), (HQ*K, 1), (i_s, 0), (BS, BK), (1, 0)) - p_pq = tl.make_block_ptr(g_cumsum + (bos * HQ + i_hq) * K, (T, K), (HQ*K, 1), (i_s, 0), (BS, BK), (1, 0)) - p_do = tl.make_block_ptr(do + (bos * HQ + i_hq) * V, (T, V), (HQ*V, 1), (i_s, i_v * BV), (BS, BV), (1, 0)) - p_lse = tl.make_block_ptr(lse + bos * HQ + i_hq, (T,), (HQ,), (i_s,), (BS,), (0,)) - p_delta = tl.make_block_ptr(delta + bos * HQ + i_hq, (T,), (HQ,), (i_s,), (BS,), (0,)) - - o_q = i_s + tl.arange(0, BS) - m_q = o_q < T - b_Rq = tl.load(p_Rq, boundary_check=(0, 1)).to(tl.float32) - b_Rq_bc_k = tl.broadcast_to(b_Rq, (BT, BK)) - b_Rq_bc_q = tl.broadcast_to(b_Rq, (BS, BK)) - b_q = tl.load(p_q, boundary_check=(0, 1)) - b_pq = tl.load(p_pq, boundary_check=(0, 1)).to(tl.float32) - b_q_til = (b_q.to(tl.float32) * exp2(b_pq - b_Rq_bc_q)).to(b_q.dtype) - b_exp_k_off = b_Rq_bc_k - b_pk - b_exp_k_off_val = exp2(b_exp_k_off) # CSE: reused for k_til and final dk scaling - b_k_til = (b_k.to(tl.float32) * b_exp_k_off_val).to(b_k.dtype) - b_do = tl.load(p_do, boundary_check=(0, 1)) - b_lse = tl.load(p_lse, boundary_check=(0,)) - b_delta = tl.load(p_delta, boundary_check=(0,)) - b_s = tl.dot(b_k_til, tl.trans(b_q_til)) * scale * RCP_LN2 - if USE_SCALAR_G: - b_cq = tl.load(g_scalar_cumsum + (bos + o_q) * HQ + i_hq, mask=m_q, other=0).to(tl.float32) - b_s += b_cq[None, :] - b_ck[:, None] - if USE_WINDOW: - b_p = tl.where((o_q[None, :] - o_k[:, None] < W) & m_q[None, :], exp2(b_s - b_lse[None, :]), 0) - else: - b_p = tl.where(m_q[None, :], exp2(b_s - b_lse[None, :]), 0) - b_dv += tl.dot(b_p.to(b_do.dtype), b_do) - b_dp = tl.dot(b_v, tl.trans(b_do)) - b_ds = b_p * (b_dp - b_delta[None, :]) - b_dk_til = tl.dot(b_ds.to(b_q.dtype), b_q_til) - b_dk += (b_dk_til * scale) * b_exp_k_off_val - if USE_SCALAR_G: - b_dc -= tl.sum(b_ds, 1) - - # Structural: b_dg derived from b_dk (saves [BT,BK] fp32 accumulator). - b_dg = -LN2 * b_k.to(tl.float32) * b_dk - tl.store(p_dk, b_dk.to(p_dk.dtype.element_ty), boundary_check=(0, 1)) - tl.store(p_dv, b_dv.to(p_dv.dtype.element_ty), boundary_check=(0, 1)) - tl.store(p_dg, b_dg.to(p_dg.dtype.element_ty), boundary_check=(0, 1)) - if USE_SCALAR_G: - p_dc = tl.make_block_ptr(dg_scalar_cumsum + bos * HQ + i_hq, (T,), (HQ,), (i_t * BT,), (BT,), (0,)) - tl.store(p_dc, b_dc.to(p_dc.dtype.element_ty), boundary_check=(0,)) - - -@dispatch('attn') -def parallel_wall_attn_fwd( - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - g_cumsum: torch.Tensor, - sink_bias: torch.Tensor | None, - scale: float, - g_scalar_cumsum: torch.Tensor | None = None, - window_size: int | None = None, - cu_seqlens: torch.LongTensor | None = None, - chunk_indices: torch.LongTensor | None = None, -): - if g_cumsum.dim() != 4 or g_cumsum.shape != (*q.shape[:-1], q.shape[-1]): - raise ValueError( - f"`g_cumsum` must be [B, T, HQ, K] matching `q`; got {g_cumsum.shape} vs q {q.shape}" - ) - if g_cumsum.stride(-1) != 1: - raise ValueError("`g_cumsum` must be contiguous in the last (K) dimension") - g_cumsum = g_cumsum.contiguous() - - B, T, H, K, V = *k.shape, v.shape[-1] - HQ = q.shape[2] - G = HQ // H - if check_shared_mem('hopper', q.device.index): - BK = min(256, max(16, triton.next_power_of_2(K))) - BV = min(256, max(16, triton.next_power_of_2(V))) - elif check_shared_mem('ampere', q.device.index): - BK = min(256, max(16, triton.next_power_of_2(K))) - BV = min(128, max(16, triton.next_power_of_2(V))) - else: - BK = min(256, max(16, triton.next_power_of_2(K))) - BV = min(64, max(16, triton.next_power_of_2(V))) - NK = triton.cdiv(K, BK) - NV = triton.cdiv(V, BV) - assert NK == 1, "The key dimension can not be larger than 256" - - if cu_seqlens is not None: - # Autotuning sweeps BT in {64,128}, but `chunk_indices` is precomputed for a - # single BT. Pin BT=128 and bypass the autotuner (calling the heuristics-wrapped - # jit directly) so the varlen chunk map stays valid. - BT = 128 - if check_shared_mem('hopper', q.device.index): - BS, num_warps, num_stages = min(64, max(16, triton.next_power_of_2(T))), 8, 2 - elif check_shared_mem('ampere', q.device.index): - BS, num_warps, num_stages = min(32, max(16, triton.next_power_of_2(T))), 4, 2 - else: - BS, num_warps, num_stages = min(32, max(16, triton.next_power_of_2(T))), 2, 2 - if chunk_indices is None: - chunk_indices = prepare_chunk_indices(cu_seqlens, BT) - NT = len(chunk_indices) - o = torch.empty(B, T, HQ, V, dtype=v.dtype, device=q.device) - lse = torch.empty(B, T, HQ, dtype=torch.float, device=q.device) - parallel_wall_attn_fwd_kernel.fn[(NV, NT, B * HQ)]( - q=q, k=k, v=v, o=o, - g_cumsum=g_cumsum, g_scalar_cumsum=g_scalar_cumsum, sink_bias=sink_bias, - lse=lse, scale=scale, cu_seqlens=cu_seqlens, chunk_indices=chunk_indices, - B=B, T=T, W=window_size, H=H, HQ=HQ, G=G, K=K, V=V, - BT=BT, BS=BS, BK=BK, BV=BV, num_warps=num_warps, num_stages=num_stages, - ) - return o, lse - - o = torch.empty(B, T, HQ, V, dtype=v.dtype, device=q.device) - lse = torch.empty(B, T, HQ, dtype=torch.float, device=q.device) - - def grid(meta): - return (NV, triton.cdiv(T, meta['BT']), B * HQ) - parallel_wall_attn_fwd_kernel[grid]( - q=q, - k=k, - v=v, - o=o, - g_cumsum=g_cumsum, - g_scalar_cumsum=g_scalar_cumsum, - sink_bias=sink_bias, - lse=lse, - scale=scale, - cu_seqlens=cu_seqlens, - chunk_indices=chunk_indices, - B=B, - T=T, - W=window_size, - H=H, - HQ=HQ, - G=G, - K=K, - V=V, - BK=BK, - BV=BV, - ) - return o, lse - - -@dispatch('attn') -def parallel_wall_attn_bwd( - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - o: torch.Tensor, - g_cumsum: torch.Tensor, - lse: torch.Tensor, - do: torch.Tensor, - sink_bias: torch.Tensor | None = None, - scale: float | None = None, - g_scalar_cumsum: torch.Tensor | None = None, - window_size: int | None = None, - cu_seqlens: torch.LongTensor | None = None, - chunk_indices: torch.LongTensor | None = None, -): - if g_cumsum.dim() != 4 or g_cumsum.shape != (*q.shape[:-1], q.shape[-1]): - raise ValueError("`g_cumsum` must be [B, T, HQ, K] matching `q`") - g_cumsum = g_cumsum.contiguous() - - B, T, H, K, V = *k.shape, v.shape[-1] - HQ = q.shape[2] - G = HQ // H - # dq/dk are reduced over the full value dim in one program (no cross-program accumulation), - # so BV must span all of V (NV == 1). Don't cap it here -- the forward can, the backward can't. - BK = max(triton.next_power_of_2(K), 16) - BV = max(triton.next_power_of_2(V), 16) - - NV = triton.cdiv(V, BV) - - # Varlen: pin BT=128 and bypass the autotuner (BT sweep would invalidate the - # precomputed chunk map). `extra` injects the fixed launch config via `.fn`. - is_varlen = cu_seqlens is not None - extra = {} - if is_varlen: - BT = 128 - if check_shared_mem('hopper'): - BS, num_warps = min(64, max(16, triton.next_power_of_2(T))), 8 - elif check_shared_mem('ampere'): - BS, num_warps = min(32, max(16, triton.next_power_of_2(T))), 4 - else: - BS, num_warps = min(32, max(16, triton.next_power_of_2(T))), 2 - if chunk_indices is None: - chunk_indices = prepare_chunk_indices(cu_seqlens, BT) - NT = len(chunk_indices) - grid = (NV, NT, B * HQ) - extra = dict(BT=BT, BS=BS, num_warps=num_warps, num_stages=2) - else: - def grid(meta): - return (NV, triton.cdiv(T, meta['BT']), B * HQ) - - delta = parallel_attn_bwd_preprocess(o, do) - - dq = torch.empty(B, T, HQ, K, dtype=k.dtype if H == HQ else torch.float, device=q.device) - dk = torch.empty(B, T, HQ, K, dtype=k.dtype if H == HQ else torch.float, device=q.device) - dv = torch.empty(B, T, HQ, V, dtype=v.dtype if H == HQ else torch.float, device=q.device) - dq_kernel = parallel_wall_attn_bwd_kernel_dq.fn if is_varlen else parallel_wall_attn_bwd_kernel_dq - dkv_kernel = parallel_wall_attn_bwd_kernel_dkv.fn if is_varlen else parallel_wall_attn_bwd_kernel_dkv - - dg_cumsum = torch.empty(B, T, HQ, K, dtype=torch.float, device=q.device) - dg_cumsum_k = torch.empty_like(dg_cumsum) - - if g_scalar_cumsum is not None: - dg_scalar_cumsum = torch.empty(B, T, HQ, dtype=torch.float, device=q.device) - dg_scalar_cumsum_k = torch.empty_like(dg_scalar_cumsum) - else: - dg_scalar_cumsum = None - dg_scalar_cumsum_k = None - - dq_kernel[grid]( - q=q, - k=k, - v=v, - g_cumsum=g_cumsum, - g_scalar_cumsum=g_scalar_cumsum, - lse=lse, - delta=delta, - do=do, - dq=dq, - dg_cumsum=dg_cumsum, - dg_scalar_cumsum=dg_scalar_cumsum, - cu_seqlens=cu_seqlens, - chunk_indices=chunk_indices, - scale=scale, - T=T, - W=window_size, - B=B, - H=H, - HQ=HQ, - G=G, - K=K, - V=V, - BK=BK, - BV=BV, - **extra, - ) - dkv_kernel[grid]( - q=q, - k=k, - v=v, - g_cumsum=g_cumsum, - g_scalar_cumsum=g_scalar_cumsum, - lse=lse, - delta=delta, - do=do, - dk=dk, - dv=dv, - dg_cumsum=dg_cumsum_k, - dg_scalar_cumsum=dg_scalar_cumsum_k, - cu_seqlens=cu_seqlens, - chunk_indices=chunk_indices, - scale=scale, - T=T, - W=window_size, - B=B, - H=H, - HQ=HQ, - G=G, - K=K, - V=V, - BK=BK, - BV=BV, - DIAG_BF16=_DKV_DIAG_BF16, - **extra, - ) - dk = reduce(dk, 'b t (h g) k -> b t h k', g=G, reduction='sum') - dv = reduce(dv, 'b t (h g) v -> b t h v', g=G, reduction='sum') - dg_cumsum.add_(dg_cumsum_k) - - if g_scalar_cumsum is not None: - dg_scalar_cumsum.add_(dg_scalar_cumsum_k) - - dsink_bias = None - if sink_bias is not None: - p_sink_bias = torch.exp2(sink_bias[None, None, :] - lse) - dsink_bias = -(p_sink_bias * delta).sum((0, 1)) - - return dq, dk, dv, dg_cumsum, dsink_bias, dg_scalar_cumsum - - -class WallParallelAttentionFunction(torch.autograd.Function): - r"""Wall Attention: per-channel prefix :math:`P` (log_2); scores use - :math:`\sum_n q_{in} k_{jn} \exp_2(P_{in}-P_{jn})` via per-Q-block reference - :math:`R_n=P_{i_t\cdot BT,n}` inside :func:`parallel_wall_attn_fwd_kernel`.""" - - @staticmethod - @contiguous - @autocast_custom_fwd - def forward(ctx, q, k, v, g, sink_bias, scale, window_size, cu_seqlens, g_scalar=None, chunk_indices=None): - if g.shape[-1] != q.shape[-1] or g.shape[2] != q.shape[2]: - raise ValueError( - f"`g` must be [B, T, HQ, K] with same HQ and K as `q`; got {g.shape} vs q {q.shape}" - ) - if q.shape[2] % k.shape[2] != 0: - raise ValueError("HQ must be divisible by `k.shape[2]` for GQA") - - if _DEBUG_ASSERTS: - assert torch.isfinite(g).all(), f"[wall-attn fwd] NaN/Inf in g: {g.abs().max()}" - assert torch.isfinite(q).all(), f"[wall-attn fwd] NaN/Inf in q: {q.abs().max()}" - assert torch.isfinite(k).all(), f"[wall-attn fwd] NaN/Inf in k: {k.abs().max()}" - assert torch.isfinite(v).all(), f"[wall-attn fwd] NaN/Inf in v: {v.abs().max()}" - - P = chunk_global_cumsum(g, cu_seqlens=cu_seqlens, scale=RCP_LN2) - if _DEBUG_ASSERTS: - assert torch.isfinite(P).all(), ( - f"[wall-attn fwd] NaN/Inf in P=cumsum(g): P abs max={P.abs().max()}, g abs max={g.abs().max()}" - ) - - if g_scalar is not None: - if _DEBUG_ASSERTS: - assert torch.isfinite(g_scalar).all(), f"[wall-attn fwd] NaN/Inf in g_scalar: {g_scalar.abs().max()}" - c = chunk_global_cumsum(g_scalar, cu_seqlens=cu_seqlens, scale=RCP_LN2) - if _DEBUG_ASSERTS: - assert torch.isfinite(c).all(), f"[wall-attn fwd] NaN/Inf in c=cumsum(g_scalar): c abs max={c.abs().max()}" - else: - c = None - - sink_bias_scaled = sink_bias * RCP_LN2 if sink_bias is not None else None - o, lse = parallel_wall_attn_fwd( - q=q, - k=k, - v=v, - g_cumsum=P, - sink_bias=sink_bias_scaled, - scale=scale, - g_scalar_cumsum=c, - window_size=window_size, - cu_seqlens=cu_seqlens, - chunk_indices=chunk_indices, - ) - if _DEBUG_ASSERTS: - assert torch.isfinite(o).all(), ( - f"[wall-attn fwd] NaN/Inf in output o: o abs max={o.abs().max()}, " - f"lse finite={torch.isfinite(lse).all()}, P abs max={P.abs().max()}" - ) - - ctx.save_for_backward(q, k, v, o, P, lse, sink_bias_scaled, c if c is not None else torch.empty(0)) - ctx.scale = scale - ctx.window_size = window_size - ctx.cu_seqlens = cu_seqlens - ctx.has_scalar_g = g_scalar is not None - return o.to(q.dtype) - - @staticmethod - @contiguous - @autocast_custom_bwd - def backward(ctx, do): - q, k, v, o, P, lse, sink_bias_scaled, c_or_empty = ctx.saved_tensors - c = c_or_empty if ctx.has_scalar_g else None - - if _DEBUG_ASSERTS: - assert torch.isfinite(do).all(), f"[wall-attn bwd] NaN/Inf in do: {do.abs().max()}" - - dq, dk, dv, dP, dsink_bias, dc = parallel_wall_attn_bwd( - q=q, - k=k, - v=v, - o=o, - g_cumsum=P, - lse=lse, - do=do, - sink_bias=sink_bias_scaled, - scale=ctx.scale, - g_scalar_cumsum=c, - window_size=ctx.window_size, - cu_seqlens=ctx.cu_seqlens, - ) - if _DEBUG_ASSERTS: - assert torch.isfinite(dq).all(), f"[wall-attn bwd] NaN/Inf in dq: {dq.abs().max()}" - assert torch.isfinite(dk).all(), f"[wall-attn bwd] NaN/Inf in dk: {dk.abs().max()}" - assert torch.isfinite(dv).all(), f"[wall-attn bwd] NaN/Inf in dv: {dv.abs().max()}" - assert torch.isfinite(dP).all(), ( - f"[wall-attn bwd] NaN/Inf in dP: dP abs max={dP.abs().max()}, " - f"P abs max={P.abs().max()}, lse finite={torch.isfinite(lse).all()}" - ) - - # The kernel emits dP with an LN2 factor (b_dg = LN2 * q * dq); the forward set - # P = cumsum(g) * RCP_LN2, so dL/dg = RCP_LN2 * reverse_cumsum(dP) (nets the LN2). - dg = chunk_global_cumsum(dP, cu_seqlens=ctx.cu_seqlens, reverse=True, scale=RCP_LN2) - if _DEBUG_ASSERTS: - assert torch.isfinite(dg).all(), ( - f"[wall-attn bwd] NaN/Inf in dg (after rev cumsum): dg abs max={dg.abs().max()}, " - f"dP abs max={dP.abs().max()}" - ) - - if dc is not None: - dg_scalar = chunk_global_cumsum(dc, cu_seqlens=ctx.cu_seqlens, reverse=True) - else: - dg_scalar = None - - return dq.to(q), dk.to(k), dv.to(v), dg, dsink_bias, None, None, None, dg_scalar, None - - -def parallel_wall_attn( - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - g: torch.Tensor, - *, - g_scalar: torch.Tensor | None = None, - sink_bias: torch.Tensor | None = None, - scale: float | None = None, - window_size: int | None = None, - cu_seqlens: torch.LongTensor | None = None, -) -> torch.Tensor: - r"""Wall parallel attention (training / prefill forward + backward). - - Scores use a per-channel multiplicative decay: with :math:`P = \mathrm{cumsum}(g)` - in :math:`\log_2` space, the logit for pair :math:`(i, j)` is - :math:`\mathrm{scale} \sum_n q_{in} k_{jn} \exp_2(P_{in} - P_{jn})`. - - Args: - q (torch.Tensor): - Queries of shape ``[B, T, HQ, K]``. - k (torch.Tensor): - Keys of shape ``[B, T, H, K]`` (``H`` may be < ``HQ`` for GQA). - v (torch.Tensor): - Values of shape ``[B, T, H, V]``. - g (torch.Tensor): - Per-channel log-decay of shape ``[B, T, HQ, K]`` (cumsum'd internally). - g_scalar (torch.Tensor, Optional): - FoX-style additive scalar gate of shape ``[B, T, HQ]``. Default: `None`. - sink_bias (torch.Tensor, Optional): - Attention-sink logit of shape ``[HQ]``. Default: `None`. - scale (float, Optional): - Softmax scale. If `None`, defaults to ``K ** -0.5``. Default: `None`. - window_size (int, Optional): - Sliding-window width (causal). Default: `None`. - cu_seqlens (torch.LongTensor, Optional): - Cumulative seqlens of shape ``[N + 1]`` for varlen packing - (requires ``B == 1``). Default: `None`. - - Returns: - Attention output of shape ``[B, T, HQ, V]``. - """ - if scale is None: - scale = k.shape[-1] ** -0.5 - if g_scalar is not None and g_scalar.shape != q.shape[:-1]: - raise ValueError(f"`g_scalar` must be [B, T, HQ] matching q.shape[:-1]; got {g_scalar.shape}") - if cu_seqlens is not None and q.shape[0] != 1: - raise ValueError("`cu_seqlens` (varlen) requires batch size 1") - if sink_bias is not None and sink_bias.shape != (q.shape[2],): - raise ValueError(f"`sink_bias` must be [HQ]; got {sink_bias.shape}") - return WallParallelAttentionFunction.apply( - q, k, v, g, sink_bias, scale, window_size, cu_seqlens, g_scalar, None - ) diff --git a/build/torch-xpu/symbols.json b/build/torch-xpu/symbols.json deleted file mode 100644 index 51659e8ff9bfa602b4e2ba9cd9f8291b2a997e41..0000000000000000000000000000000000000000 --- a/build/torch-xpu/symbols.json +++ /dev/null @@ -1,640 +0,0 @@ -{ - "schema_version": 1, - "module": "_fla_xpu_fe2bb66", - "functions": [ - { - "name": "chunk_kda", - "qualname": "chunk_kda", - "module": "_fla_xpu_fe2bb66.ops.kda.chunk", - "kind": "function", - "doc": "Args:\n q (torch.Tensor):\n queries of shape ``[B, T, H, K]``.\n k (torch.Tensor):\n keys of shape ``[B, T, H, K]``.\n v (torch.Tensor):\n values of shape ``[B, T, HV, V]``.\n GVA (Grouped Value Attention) is applied if ``HV > H``, where ``HV`` must be divisible by ``H``.\n g (torch.Tensor):\n (forget) gating tensor (in log space!) of shape ``[B, T, HV, K]``.\n When ``use_gate_in_kernel=False`` (default), ``g`` should be the pre-computed decay value.\n When ``use_gate_in_kernel=True``, ``g`` is the raw input before gate activation;\n the kernel fuses ``-exp(A_log) * softplus(g + dt_bias)`` + chunk cumsum internally.\n beta (torch.Tensor):\n betas of shape ``[B, T, HV]``.\n scale (Optional[float]):\n Scale factor for the KDA attention scores.\n If not provided, it will default to ``1 / sqrt(K)``. Default: ``None``.\n initial_state (Optional[torch.Tensor]):\n Initial state of shape ``[N, HV, K, V]`` for ``N`` input sequences.\n For equal-length input sequences, ``N`` equals the batch size ``B``.\n Default: ``None``.\n output_final_state (Optional[bool]):\n Whether to output the final state of shape ``[N, HV, K, V]``. Default: ``False``.\n use_qk_l2norm_in_kernel (bool):\n Whether to apply L2norm to the q,k tensor internally. Default: ``False``.\n use_gate_in_kernel (bool):\n Whether to compute the log-space KDA decay internally.\n - If ``True``:\n The passed ``g`` acts as the raw input for ``-exp(A_log) * softplus(g + dt_bias.view(HV, K))``.\n Note that as part of the input arguments,\n ``A_log`` (shape ``[HV]``) and the optional ``dt_bias`` (shape ``[HV * K]``) should be provided.\n - If ``False``, ``g`` is expected to be the pre-computed decay value.\n Default: ``False``.\n use_beta_sigmoid_in_kernel (bool):\n Whether to apply ``torch.sigmoid(beta)`` before launching the chunk kernel.\n - If ``True``, the passed ``beta`` acts as the raw beta logits.\n - If ``False``, ``beta`` is expected to already be in post-sigmoid space.\n Default: ``False``.\n allow_neg_eigval (bool):\n Whether to allow negative eigenvalues by scaling ``beta`` to ``[0, 2)``.\n Only takes effect together with ``use_beta_sigmoid_in_kernel=True``, in which case\n the kernel computes ``2 * sigmoid(beta)`` instead of ``sigmoid(beta)``.\n Default: ``False``.\n safe_gate (bool):\n Whether to clamp the gate to ``[lower_bound, 0)`` and enable M=16 TensorCore\n acceleration for higher throughput. Requires ``lower_bound`` to be set.\n Default: ``False``.\n lower_bound (Optional[float]):\n Lower bound for the forget gate (in log space). When set together with\n ``safe_gate=True``, changes the gate activation from\n ``-exp(A_log) * softplus(g + dt_bias)`` to\n ``lower_bound * sigmoid(exp(A_log) * (g + dt_bias))``,\n which naturally clamps the output to ``[lower_bound, 0)``.\n Recommended value: ``-5`` (i.e., per-step decay ``exp(-5) ≈ 0.0067``).\n Default: ``None``.\n disable_recompute (bool):\n Whether to disable gradient recomputation in the kernel. When ``True``, the kernel\n will save all intermediate activations for backward pass, which is beneficial\n for training small models at the cost of increased memory usage. Default: ``False``.\n return_intermediate_states (bool):\n If True, returns intermediate state ``h`` for inference scenarios (e.g., vLLM).\n Must be used within ``torch.inference_mode()`` and will return a 3-tuple instead of 2-tuple.\n This is not intended for training as it bypasses autograd. Default: ``False``.\n state_v_first (Optional[bool]):\n Store the recurrent state in V-first ``[V, K]`` layout instead of the default ``[K, V]``. Default: ``False``.\n cu_seqlens (torch.LongTensor):\n Cumulative sequence lengths of shape ``[N+1]`` used for variable-length training,\n consistent with the FlashAttention API.\n cu_seqlens_cpu (torch.LongTensor):\n Cumulative sequence lengths of shape ``[N+1]`` used for variable-length training,\n consistent with the FlashAttention API.\n cp_context (Optional[FLACPContext]):\n Context parallel context for distributed training across multiple devices.\n When provided, ``initial_state`` and ``output_final_state`` are not supported,\n and ``cu_seqlens`` will be overridden by the context. Default: ``None``.\n\nReturns:\n - Normal mode (return_intermediate_states=False): A tuple (o, final_state)\n o (torch.Tensor):\n Outputs of shape ``[B, T, HV, V]``.\n final_state (torch.Tensor):\n Final state of shape ``[N, HV, K, V]`` if ``output_final_state=True`` else ``None``.\n - Inference mode (return_intermediate_states=True): A tuple (o, final_state, h)\n o (torch.Tensor):\n Outputs of shape ``[B, T, HV, V]``.\n final_state (torch.Tensor):\n Final state of shape ``[N, HV, K, V]`` if ``output_final_state=True`` else ``None``.\n h (torch.Tensor):\n Intermediate states of shape ``[B, NT, HV, K, V]`` and dtype ``bfloat16``.\n - For equal-length sequences: ``NT = ceil(T / chunk_size)``\n - For variable-length sequences (cu_seqlens): B is always 1 (flattened),\n NT is the total number of chunks across all sequences.\n\nExamples::\n >>> import torch\n >>> import torch.nn.functional as F\n >>> from einops import rearrange\n >>> from ...ops.kda import chunk_kda\n # inputs with equal lengths (no GVA, HV == H)\n >>> B, T, H, K, V = 4, 2048, 4, 512, 512\n >>> q = torch.randn(B, T, H, K, dtype=torch.bfloat16, device='cuda')\n >>> k = torch.randn(B, T, H, K, dtype=torch.bfloat16, device='cuda')\n >>> v = torch.randn(B, T, H, V, dtype=torch.bfloat16, device='cuda')\n >>> beta = torch.rand(B, T, H, dtype=torch.bfloat16, device='cuda')\n >>> g = torch.rand(B, T, H, K, dtype=torch.bfloat16, device='cuda')\n >>> h0 = torch.randn(B, H, K, V, dtype=torch.bfloat16, device='cuda')\n >>> A_log = torch.randn(H, dtype=torch.float32, device='cuda')\n >>> dt_bias = torch.randn(H * K, dtype=torch.float32, device='cuda')\n >>> o, ht = chunk_kda(\n q, k, v, g, beta,\n A_log=A_log,\n dt_bias=dt_bias,\n use_qk_l2norm_in_kernel=True,\n use_gate_in_kernel=True,\n initial_state=h0,\n output_final_state=True\n )\n # GVA mode (HV > H)\n >>> HV = 8 # 2x more value heads than qk heads\n >>> v = torch.randn(B, T, HV, V, dtype=torch.bfloat16, device='cuda')\n >>> g = torch.rand(B, T, HV, K, dtype=torch.bfloat16, device='cuda')\n >>> beta = torch.rand(B, T, HV, dtype=torch.bfloat16, device='cuda')\n >>> h0 = torch.randn(B, HV, K, V, dtype=torch.bfloat16, device='cuda')\n >>> A_log = torch.randn(HV, dtype=torch.float32, device='cuda')\n >>> dt_bias = torch.randn(HV * K, dtype=torch.float32, device='cuda')\n >>> o, ht = chunk_kda(\n q, k, v, g, beta,\n A_log=A_log,\n dt_bias=dt_bias,\n use_qk_l2norm_in_kernel=True,\n use_gate_in_kernel=True,\n initial_state=h0,\n output_final_state=True\n )\n # for variable-length inputs, the batch size `B` is expected to be 1 and `cu_seqlens` is required\n >>> q, k, v, beta, g = map(lambda x: rearrange(x, 'b t ... -> 1 (b t) ...'), (q, k, v, beta, g))\n # for a batch with 4 sequences, `cu_seqlens` with 5 start/end positions are expected\n >>> cu_seqlens = q.new_tensor([0, 2048, 4096, 6144, 8192], dtype=torch.long)\n >>> o, ht = chunk_kda(\n q, k, v, g, beta,\n A_log=A_log,\n dt_bias=dt_bias,\n use_qk_l2norm_in_kernel=True,\n use_gate_in_kernel=True,\n initial_state=h0,\n output_final_state=True,\n cu_seqlens=cu_seqlens\n )", - "signature": { - "parameters": [ - { - "name": "q", - "kind": "POSITIONAL_OR_KEYWORD", - "default": null, - "annotation": "torch.Tensor" - }, - { - "name": "k", - "kind": "POSITIONAL_OR_KEYWORD", - "default": null, - "annotation": "torch.Tensor" - }, - { - "name": "v", - "kind": "POSITIONAL_OR_KEYWORD", - "default": null, - "annotation": "torch.Tensor" - }, - { - "name": "g", - "kind": "POSITIONAL_OR_KEYWORD", - "default": null, - "annotation": "torch.Tensor" - }, - { - "name": "beta", - "kind": "POSITIONAL_OR_KEYWORD", - "default": null, - "annotation": "torch.Tensor" - }, - { - "name": "scale", - "kind": "POSITIONAL_OR_KEYWORD", - "default": "None", - "annotation": "float | None" - }, - { - "name": "initial_state", - "kind": "POSITIONAL_OR_KEYWORD", - "default": "None", - "annotation": "torch.Tensor | None" - }, - { - "name": "output_final_state", - "kind": "POSITIONAL_OR_KEYWORD", - "default": "False", - "annotation": "bool" - }, - { - "name": "use_qk_l2norm_in_kernel", - "kind": "POSITIONAL_OR_KEYWORD", - "default": "False", - "annotation": "bool" - }, - { - "name": "use_gate_in_kernel", - "kind": "POSITIONAL_OR_KEYWORD", - "default": "False", - "annotation": "bool" - }, - { - "name": "use_beta_sigmoid_in_kernel", - "kind": "POSITIONAL_OR_KEYWORD", - "default": "False", - "annotation": "bool" - }, - { - "name": "allow_neg_eigval", - "kind": "POSITIONAL_OR_KEYWORD", - "default": "False", - "annotation": "bool" - }, - { - "name": "safe_gate", - "kind": "POSITIONAL_OR_KEYWORD", - "default": "False", - "annotation": "bool" - }, - { - "name": "lower_bound", - "kind": "POSITIONAL_OR_KEYWORD", - "default": "None", - "annotation": "float | None" - }, - { - "name": "disable_recompute", - "kind": "POSITIONAL_OR_KEYWORD", - "default": "False", - "annotation": "bool" - }, - { - "name": "return_intermediate_states", - "kind": "POSITIONAL_OR_KEYWORD", - "default": "False", - "annotation": "bool" - }, - { - "name": "state_v_first", - "kind": "POSITIONAL_OR_KEYWORD", - "default": "False", - "annotation": "bool" - }, - { - "name": "cu_seqlens", - "kind": "POSITIONAL_OR_KEYWORD", - "default": "None", - "annotation": "torch.LongTensor | None" - }, - { - "name": "cu_seqlens_cpu", - "kind": "POSITIONAL_OR_KEYWORD", - "default": "None", - "annotation": "torch.LongTensor | None" - }, - { - "name": "cp_context", - "kind": "POSITIONAL_OR_KEYWORD", - "default": "None", - "annotation": "_fla_xpu_fe2bb66.ops.cp.context.FLACPContext" - }, - { - "name": "kwargs", - "kind": "VAR_KEYWORD", - "default": null, - "annotation": null - } - ], - "return_annotation": null - } - }, - { - "name": "fused_recurrent_kda", - "qualname": "fused_recurrent_kda", - "module": "_fla_xpu_fe2bb66.ops.kda.fused_recurrent", - "kind": "function", - "doc": "Args:\n q (torch.Tensor):\n queries of shape `[B, T, H, K]`.\n k (torch.Tensor):\n keys of shape `[B, T, H, K]`.\n v (torch.Tensor):\n values of shape `[B, T, HV, V]`.\n GVA is applied if `HV > H`.\n g (torch.Tensor):\n g (decays) of shape `[B, T, HV, K]`.\n beta (torch.Tensor):\n betas of shape `[B, T, HV]`.\n A_log (Optional[torch.Tensor]):\n Decay parameter of shape `[HV]`. Required when `use_gate_in_kernel=True`.\n dt_bias (Optional[torch.Tensor]):\n Bias added to `g` before activation, of shape `[HV]`.\n Only used when `use_gate_in_kernel=True`.\n scale (Optional[float]):\n Scale factor for the RetNet attention scores.\n If not provided, it will default to `1 / sqrt(K)`. Default: `None`.\n initial_state (Optional[torch.Tensor]):\n Initial state of shape `[N, HV, K, V]` for `N` input sequences.\n For equal-length input sequences, `N` equals the batch size `B`.\n Default: `None`.\n output_final_state (Optional[bool]):\n Whether to output the final state of shape `[N, HV, K, V]`. Default: `False`.\n use_qk_l2norm_in_kernel (Optional[bool]):\n Whether to use L2 normalization in the kernel. Default: `False`.\n use_gate_in_kernel (Optional[bool]):\n Whether to compute the log-space KDA decay internally.\n When `True`, `g` is the raw input and `A_log` must be provided; the kernel fuses\n gate activation into the recurrence. Default: `False`.\n use_beta_sigmoid_in_kernel (Optional[bool]):\n Whether to apply `torch.sigmoid(beta)` inside the kernel.\n - If `True`, the passed `beta` acts as the raw beta logits.\n - If `False`, `beta` is expected to already be in post-sigmoid space.\n Default: `False`.\n allow_neg_eigval (Optional[bool]):\n Whether to allow negative eigenvalues by scaling `beta` to `[0, 2)`.\n Only takes effect together with `use_beta_sigmoid_in_kernel=True`, in which case\n the kernel computes `2 * sigmoid(beta)` instead of `sigmoid(beta)`. Default: `False`.\n lower_bound (Optional[float]):\n Lower bound for the forget gate (in log space). Only used when `use_gate_in_kernel=True`. Default: `None`.\n state_v_first (Optional[bool]):\n Store the recurrent state in V-first ``[V, K]`` layout instead of the default ``[K, V]``. Default: ``False``.\n cu_seqlens (torch.LongTensor):\n Cumulative sequence lengths of shape `[N+1]` used for variable-length training,\n consistent with the FlashAttention API.\n\nReturns:\n o (torch.Tensor):\n Outputs of shape `[B, T, HV, V]`.\n final_state (torch.Tensor):\n Final state of shape `[N, HV, K, V]` if `output_final_state=True` else `None`.\n\nExamples::\n >>> import torch\n >>> import torch.nn.functional as F\n >>> from einops import rearrange\n >>> from ...ops.kda import fused_recurrent_kda\n # inputs with equal lengths\n >>> B, T, H, HV, K, V = 4, 2048, 4, 8, 512, 512\n >>> q = torch.randn(B, T, H, K, device='cuda')\n >>> k = F.normalize(torch.randn(B, T, H, K, device='cuda'), p=2, dim=-1)\n >>> v = torch.randn(B, T, HV, V, device='cuda')\n >>> g = F.logsigmoid(torch.rand(B, T, HV, K, device='cuda'))\n >>> beta = torch.rand(B, T, HV, device='cuda').sigmoid()\n >>> h0 = torch.randn(B, HV, K, V, device='cuda')\n >>> o, ht = fused_recurrent_kda(\n q, k, v, g, beta,\n initial_state=h0,\n output_final_state=True\n )\n # for variable-length inputs, the batch size `B` is expected to be 1 and `cu_seqlens` is required\n >>> q, k, v, g, beta = map(lambda x: rearrange(x, 'b t ... -> 1 (b t) ...'), (q, k, v, g, beta))\n # for a batch with 4 sequences, `cu_seqlens` with 5 start/end positions are expected\n >>> cu_seqlens = q.new_tensor([0, 2048, 4096, 6144, 8192], dtype=torch.long)\n >>> o_var, ht_var = fused_recurrent_kda(\n q, k, v, g, beta,\n initial_state=h0,\n output_final_state=True,\n cu_seqlens=cu_seqlens\n )", - "signature": { - "parameters": [ - { - "name": "q", - "kind": "POSITIONAL_OR_KEYWORD", - "default": null, - "annotation": "torch.Tensor" - }, - { - "name": "k", - "kind": "POSITIONAL_OR_KEYWORD", - "default": null, - "annotation": "torch.Tensor" - }, - { - "name": "v", - "kind": "POSITIONAL_OR_KEYWORD", - "default": null, - "annotation": "torch.Tensor" - }, - { - "name": "g", - "kind": "POSITIONAL_OR_KEYWORD", - "default": null, - "annotation": "torch.Tensor" - }, - { - "name": "beta", - "kind": "POSITIONAL_OR_KEYWORD", - "default": null, - "annotation": "torch.Tensor" - }, - { - "name": "A_log", - "kind": "POSITIONAL_OR_KEYWORD", - "default": "None", - "annotation": "torch.Tensor | None" - }, - { - "name": "dt_bias", - "kind": "POSITIONAL_OR_KEYWORD", - "default": "None", - "annotation": "torch.Tensor | None" - }, - { - "name": "scale", - "kind": "POSITIONAL_OR_KEYWORD", - "default": "None", - "annotation": "float | None" - }, - { - "name": "initial_state", - "kind": "POSITIONAL_OR_KEYWORD", - "default": "None", - "annotation": "torch.Tensor" - }, - { - "name": "output_final_state", - "kind": "POSITIONAL_OR_KEYWORD", - "default": "False", - "annotation": "bool" - }, - { - "name": "use_qk_l2norm_in_kernel", - "kind": "POSITIONAL_OR_KEYWORD", - "default": "False", - "annotation": "bool" - }, - { - "name": "use_gate_in_kernel", - "kind": "POSITIONAL_OR_KEYWORD", - "default": "False", - "annotation": "bool" - }, - { - "name": "use_beta_sigmoid_in_kernel", - "kind": "POSITIONAL_OR_KEYWORD", - "default": "False", - "annotation": "bool" - }, - { - "name": "allow_neg_eigval", - "kind": "POSITIONAL_OR_KEYWORD", - "default": "False", - "annotation": "bool" - }, - { - "name": "lower_bound", - "kind": "POSITIONAL_OR_KEYWORD", - "default": "None", - "annotation": "float | None" - }, - { - "name": "state_v_first", - "kind": "POSITIONAL_OR_KEYWORD", - "default": "False", - "annotation": "bool" - }, - { - "name": "cu_seqlens", - "kind": "POSITIONAL_OR_KEYWORD", - "default": "None", - "annotation": "torch.LongTensor | None" - }, - { - "name": "kwargs", - "kind": "VAR_KEYWORD", - "default": null, - "annotation": null - } - ], - "return_annotation": "tuple[torch.Tensor, torch.Tensor]" - } - }, - { - "name": "chunk_gated_delta_rule", - "qualname": "chunk_gated_delta_rule", - "module": "_fla_xpu_fe2bb66.ops.gated_delta_rule.chunk", - "kind": "function", - "doc": "Args:\n q (torch.Tensor):\n queries of shape `[B, T, H, K]`.\n k (torch.Tensor):\n keys of shape `[B, T, H, K]`.\n v (torch.Tensor):\n values of shape `[B, T, HV, V]`.\n GVA (Grouped Value Attention) is applied if `HV > H`, where `HV` must be divisible by `H`.\n g (torch.Tensor):\n (forget) gating tensor of shape `[B, T, HV]`.\n When `use_gate_in_kernel=False` (default), `g` should be in log space (pre-computed decay).\n When `use_gate_in_kernel=True`, `g` is the raw input before gate activation;\n the kernel fuses `-exp(A_log) * softplus(g + dt_bias)` + chunk cumsum internally.\n beta (torch.Tensor):\n betas of shape `[B, T, HV]`.\n scale (Optional[float]):\n Scale factor for the RetNet attention scores.\n If not provided, it will default to `1 / sqrt(K)`. Default: `None`.\n initial_state (Optional[torch.Tensor]):\n Initial state of shape `[N, HV, K, V]` for `N` input sequences.\n For equal-length input sequences, `N` equals the batch size `B`.\n Default: `None`.\n output_final_state (Optional[bool]):\n Whether to output the final state of shape `[N, HV, K, V]`. Default: `False`.\n use_qk_l2norm_in_kernel (bool):\n Whether to apply L2norm to the q/k tensor internally. Default: `False`.\n use_gate_in_kernel (bool):\n Whether to compute the log-space GDN decay internally.\n When `True`, the passed `g` is the raw input, and `A_log` must be provided.\n The kernel fuses gate activation + chunk cumsum in a single pass.\n Default: `False`.\n A_log (Optional[torch.Tensor]):\n Decay parameter of shape `[HV]`. Required when `use_gate_in_kernel=True`.\n dt_bias (Optional[torch.Tensor]):\n Bias added to `g` before activation, of shape `[HV]`.\n Only used when `use_gate_in_kernel=True`.\n use_beta_sigmoid_in_kernel (bool):\n Whether to apply `torch.sigmoid(beta)` before launching the chunk kernel.\n - If `True`, the passed `beta` acts as the raw beta logits.\n - If `False`, `beta` is expected to already be in post-sigmoid space.\n Default: `False`.\n allow_neg_eigval (bool):\n Whether to allow negative eigenvalues by scaling `beta` to `[0, 2)`.\n Only takes effect together with `use_beta_sigmoid_in_kernel=True`, in which case\n the kernel computes `2 * sigmoid(beta)` instead of `sigmoid(beta)`. Default: `False`.\n state_v_first (Optional[bool]):\n Store the recurrent state in V-first ``[V, K]`` layout instead of the default ``[K, V]``. Default: ``False``.\n cu_seqlens (torch.LongTensor):\n Cumulative sequence lengths of shape `[N+1]` used for variable-length training,\n consistent with the FlashAttention API.\n cp_context (Optional[FLACPContext]):\n Context parallel context for distributed training across multiple devices.\n When provided, `initial_state` and `output_final_state` are not supported,\n and `cu_seqlens` will be overridden by the context. Default: `None`.\n\nReturns:\n o (torch.Tensor):\n Outputs of shape `[B, T, HV, V]`.\n final_state (torch.Tensor):\n Final state of shape `[N, HV, K, V]` if `output_final_state=True` else `None`.\n\nExamples::\n >>> import torch\n >>> import torch.nn.functional as F\n >>> from einops import rearrange\n >>> from ...ops.gated_delta_rule import chunk_gated_delta_rule\n # inputs with equal lengths\n >>> B, T, H, HV, K, V = 4, 2048, 4, 8, 512, 512\n >>> q = torch.randn(B, T, H, K, dtype=torch.bfloat16, device='cuda')\n >>> k = F.normalize(torch.randn(B, T, H, K, dtype=torch.bfloat16, device='cuda'), p=2, dim=-1)\n >>> v = torch.randn(B, T, HV, V, dtype=torch.bfloat16, device='cuda')\n >>> beta = torch.rand(B, T, HV, dtype=torch.bfloat16, device='cuda').sigmoid()\n >>> g = F.logsigmoid(torch.rand(B, T, HV, dtype=torch.bfloat16, device='cuda'))\n >>> h0 = torch.randn(B, HV, K, V, dtype=torch.bfloat16, device='cuda')\n >>> o, ht = chunk_gated_delta_rule(\n q, k, v, g, beta,\n initial_state=h0,\n output_final_state=True\n )\n # for variable-length inputs, the batch size `B` is expected to be 1 and `cu_seqlens` is required\n >>> q, k, v, beta, g = map(lambda x: rearrange(x, 'b t ... -> 1 (b t) ...'), (q, k, v, beta, g))\n # for a batch with 4 sequences, `cu_seqlens` with 5 start/end positions are expected\n >>> cu_seqlens = q.new_tensor([0, 2048, 4096, 6144, 8192], dtype=torch.long)\n >>> o, ht = chunk_gated_delta_rule(\n q, k, v, g, beta,\n initial_state=h0,\n output_final_state=True,\n cu_seqlens=cu_seqlens\n )", - "signature": { - "parameters": [ - { - "name": "q", - "kind": "POSITIONAL_OR_KEYWORD", - "default": null, - "annotation": "torch.Tensor" - }, - { - "name": "k", - "kind": "POSITIONAL_OR_KEYWORD", - "default": null, - "annotation": "torch.Tensor" - }, - { - "name": "v", - "kind": "POSITIONAL_OR_KEYWORD", - "default": null, - "annotation": "torch.Tensor" - }, - { - "name": "g", - "kind": "POSITIONAL_OR_KEYWORD", - "default": null, - "annotation": "torch.Tensor" - }, - { - "name": "beta", - "kind": "POSITIONAL_OR_KEYWORD", - "default": null, - "annotation": "torch.Tensor" - }, - { - "name": "scale", - "kind": "POSITIONAL_OR_KEYWORD", - "default": "None", - "annotation": "float | None" - }, - { - "name": "initial_state", - "kind": "POSITIONAL_OR_KEYWORD", - "default": "None", - "annotation": "torch.Tensor | None" - }, - { - "name": "output_final_state", - "kind": "POSITIONAL_OR_KEYWORD", - "default": "False", - "annotation": "bool" - }, - { - "name": "use_qk_l2norm_in_kernel", - "kind": "POSITIONAL_OR_KEYWORD", - "default": "False", - "annotation": "bool" - }, - { - "name": "use_beta_sigmoid_in_kernel", - "kind": "POSITIONAL_OR_KEYWORD", - "default": "False", - "annotation": "bool" - }, - { - "name": "allow_neg_eigval", - "kind": "POSITIONAL_OR_KEYWORD", - "default": "False", - "annotation": "bool" - }, - { - "name": "state_v_first", - "kind": "POSITIONAL_OR_KEYWORD", - "default": "False", - "annotation": "bool" - }, - { - "name": "cu_seqlens", - "kind": "POSITIONAL_OR_KEYWORD", - "default": "None", - "annotation": "torch.LongTensor | None" - }, - { - "name": "cu_seqlens_cpu", - "kind": "POSITIONAL_OR_KEYWORD", - "default": "None", - "annotation": "torch.LongTensor | None" - }, - { - "name": "cp_context", - "kind": "POSITIONAL_OR_KEYWORD", - "default": "None", - "annotation": "_fla_xpu_fe2bb66.ops.cp.context.FLACPContext | None" - }, - { - "name": "kwargs", - "kind": "VAR_KEYWORD", - "default": null, - "annotation": null - } - ], - "return_annotation": null - } - }, - { - "name": "fused_recurrent_gated_delta_rule", - "qualname": "fused_recurrent_gated_delta_rule", - "module": "_fla_xpu_fe2bb66.ops.gated_delta_rule.fused_recurrent", - "kind": "function", - "doc": "Args:\n q (torch.Tensor):\n queries of shape `[B, T, H, K]`.\n k (torch.Tensor):\n keys of shape `[B, T, H, K]`.\n v (torch.Tensor):\n values of shape `[B, T, HV, V]`.\n GVA (Grouped Value Attention) is applied if `HV > H`, where `HV` must be divisible by `H`.\n g (torch.Tensor):\n g (decays) of shape `[B, T, HV]`. Default: `None`.\n When `use_gate_in_kernel=False` (default), `g` must be in log space (pre-computed decay).\n When `use_gate_in_kernel=True`, `g` is the raw pre-activation input; the kernel fuses\n `-exp(A_log) * softplus(g + dt_bias)` internally per step.\n gk (torch.Tensor):\n gk (decays) of shape `[B, T, HV, K]`. Default: `None`.\n gv (torch.Tensor):\n gv (decays) of shape `[B, T, HV, V]`. Default: `None`.\n beta (torch.Tensor):\n betas of shape `[B, T, HV]`.\n scale (Optional[float]):\n Scale factor for the RetNet attention scores.\n If not provided, it will default to `1 / sqrt(K)`. Default: `None`.\n initial_state (Optional[torch.Tensor]):\n Initial state of shape `[N, HV, K, V]` for `N` input sequences.\n For equal-length input sequences, `N` equals the batch size `B`.\n Default: `None`.\n output_final_state (Optional[bool]):\n Whether to output the final state of shape `[N, HV, K, V]`. Default: `False`.\n use_qk_l2norm_in_kernel (Optional[bool]):\n Whether to use L2 normalization in the kernel. Default: `False`.\n use_gate_in_kernel (bool):\n Whether to compute the log-space GDN decay internally.\n When `True`, `g` is the raw input and `A_log` must be provided; the kernel fuses\n gate activation into the recurrence. Default: `False`.\n A_log (Optional[torch.Tensor]):\n Decay parameter of shape `[HV]`. Required when `use_gate_in_kernel=True`.\n dt_bias (Optional[torch.Tensor]):\n Bias added to `g` before activation, of shape `[HV]`.\n Only used when `use_gate_in_kernel=True`.\n use_beta_sigmoid_in_kernel (Optional[bool]):\n Whether to apply `torch.sigmoid(beta)` inside the kernel.\n - If `True`, the passed `beta` acts as the raw beta logits.\n - If `False`, `beta` is expected to already be in post-sigmoid space.\n Default: `False`.\n allow_neg_eigval (Optional[bool]):\n Whether to allow negative eigenvalues by scaling `beta` to `[0, 2)`.\n Only takes effect together with `use_beta_sigmoid_in_kernel=True`, in which case\n the kernel computes `2 * sigmoid(beta)` instead of `sigmoid(beta)`. Default: `False`.\n state_v_first (Optional[bool]):\n Store the recurrent state in V-first ``[V, K]`` layout instead of the default ``[K, V]``. Default: ``False``.\n cu_seqlens (torch.LongTensor):\n Cumulative sequence lengths of shape `[N+1]` used for variable-length training,\n consistent with the FlashAttention API.\n\nReturns:\n o (torch.Tensor):\n Outputs of shape `[B, T, HV, V]`.\n final_state (torch.Tensor):\n Final state of shape `[N, HV, K, V]` if `output_final_state=True` else `None`.\n\nExamples::\n >>> import torch\n >>> import torch.nn.functional as F\n >>> from einops import rearrange\n >>> from ...ops.gated_delta_rule import fused_recurrent_gated_delta_rule\n # inputs with equal lengths\n >>> B, T, H, HV, K, V = 4, 2048, 4, 8, 512, 512\n >>> q = torch.randn(B, T, H, K, device='cuda')\n >>> k = F.normalize(torch.randn(B, T, H, K, device='cuda'), p=2, dim=-1)\n >>> v = torch.randn(B, T, HV, V, device='cuda')\n >>> g = F.logsigmoid(torch.rand(B, T, HV, device='cuda'))\n >>> beta = torch.rand(B, T, HV, device='cuda').sigmoid()\n >>> h0 = torch.randn(B, HV, K, V, device='cuda')\n >>> o, ht = fused_gated_recurrent_delta_rule(\n q, k, v, g, beta,\n initial_state=h0,\n output_final_state=True\n )\n # for variable-length inputs, the batch size `B` is expected to be 1 and `cu_seqlens` is required\n >>> q, k, v, g, beta = map(lambda x: rearrange(x, 'b t ... -> 1 (b t) ...'), (q, k, v, g, beta))\n # for a batch with 4 sequences, `cu_seqlens` with 5 start/end positions are expected\n >>> cu_seqlens = q.new_tensor([0, 2048, 4096, 6144, 8192], dtype=torch.long)\n >>> o, ht = fused_gated_recurrent_delta_rule(\n q, k, v, g, beta,\n initial_state=h0,\n output_final_state=True,\n cu_seqlens=cu_seqlens\n )", - "signature": { - "parameters": [ - { - "name": "q", - "kind": "POSITIONAL_OR_KEYWORD", - "default": null, - "annotation": "torch.Tensor" - }, - { - "name": "k", - "kind": "POSITIONAL_OR_KEYWORD", - "default": null, - "annotation": "torch.Tensor" - }, - { - "name": "v", - "kind": "POSITIONAL_OR_KEYWORD", - "default": null, - "annotation": "torch.Tensor" - }, - { - "name": "g", - "kind": "POSITIONAL_OR_KEYWORD", - "default": "None", - "annotation": "torch.Tensor | None" - }, - { - "name": "gk", - "kind": "POSITIONAL_OR_KEYWORD", - "default": "None", - "annotation": "torch.Tensor | None" - }, - { - "name": "gv", - "kind": "POSITIONAL_OR_KEYWORD", - "default": "None", - "annotation": "torch.Tensor | None" - }, - { - "name": "beta", - "kind": "POSITIONAL_OR_KEYWORD", - "default": "None", - "annotation": "torch.Tensor | None" - }, - { - "name": "scale", - "kind": "POSITIONAL_OR_KEYWORD", - "default": "None", - "annotation": "float" - }, - { - "name": "initial_state", - "kind": "POSITIONAL_OR_KEYWORD", - "default": "None", - "annotation": "torch.Tensor" - }, - { - "name": "output_final_state", - "kind": "POSITIONAL_OR_KEYWORD", - "default": "False", - "annotation": "bool" - }, - { - "name": "use_qk_l2norm_in_kernel", - "kind": "POSITIONAL_OR_KEYWORD", - "default": "False", - "annotation": "bool" - }, - { - "name": "use_gate_in_kernel", - "kind": "POSITIONAL_OR_KEYWORD", - "default": "False", - "annotation": "bool" - }, - { - "name": "A_log", - "kind": "POSITIONAL_OR_KEYWORD", - "default": "None", - "annotation": "torch.Tensor | None" - }, - { - "name": "dt_bias", - "kind": "POSITIONAL_OR_KEYWORD", - "default": "None", - "annotation": "torch.Tensor | None" - }, - { - "name": "use_beta_sigmoid_in_kernel", - "kind": "POSITIONAL_OR_KEYWORD", - "default": "False", - "annotation": "bool" - }, - { - "name": "allow_neg_eigval", - "kind": "POSITIONAL_OR_KEYWORD", - "default": "False", - "annotation": "bool" - }, - { - "name": "state_v_first", - "kind": "POSITIONAL_OR_KEYWORD", - "default": "False", - "annotation": "bool" - }, - { - "name": "cu_seqlens", - "kind": "POSITIONAL_OR_KEYWORD", - "default": "None", - "annotation": "torch.LongTensor | None" - }, - { - "name": "kwargs", - "kind": "VAR_KEYWORD", - "default": null, - "annotation": null - } - ], - "return_annotation": "tuple[torch.Tensor, torch.Tensor]" - } - } - ], - "layers": [ - { - "name": "FusedRMSNormGated", - "qualname": "FusedRMSNormGated", - "module": "_fla_xpu_fe2bb66.layers", - "kind": "class", - "doc": "Base class for all neural network modules.\n\nYour models should also subclass this class.\n\nModules can also contain other Modules, allowing them to be nested in\na tree structure. You can assign the submodules as regular attributes::\n\n import torch.nn as nn\n import torch.nn.functional as F\n\n\n class Model(nn.Module):\n def __init__(self) -> None:\n super().__init__()\n self.conv1 = nn.Conv2d(1, 20, 5)\n self.conv2 = nn.Conv2d(20, 20, 5)\n\n def forward(self, x):\n x = F.relu(self.conv1(x))\n return F.relu(self.conv2(x))\n\nSubmodules assigned in this way will be registered, and will also have their\nparameters converted when you call :meth:`to`, etc.\n\n.. note::\n As per the example above, an ``__init__()`` call to the parent class\n must be made before assignment on the child.\n\n:ivar training: Boolean represents whether this module is in training or\n evaluation mode.\n:vartype training: bool", - "signature": { - "parameters": [ - { - "name": "args", - "kind": "VAR_POSITIONAL", - "default": null, - "annotation": "Any" - }, - { - "name": "kwargs", - "kind": "VAR_KEYWORD", - "default": null, - "annotation": "Any" - } - ], - "return_annotation": "None" - }, - "attributes": { - "has_backward": null, - "can_torch_compile": null - } - }, - { - "name": "chunk_kimi_delta_attention", - "qualname": "chunk_kimi_delta_attention", - "module": "_fla_xpu_fe2bb66.layers", - "kind": "class", - "doc": "Base class for all neural network modules.\n\nYour models should also subclass this class.\n\nModules can also contain other Modules, allowing them to be nested in\na tree structure. You can assign the submodules as regular attributes::\n\n import torch.nn as nn\n import torch.nn.functional as F\n\n\n class Model(nn.Module):\n def __init__(self) -> None:\n super().__init__()\n self.conv1 = nn.Conv2d(1, 20, 5)\n self.conv2 = nn.Conv2d(20, 20, 5)\n\n def forward(self, x):\n x = F.relu(self.conv1(x))\n return F.relu(self.conv2(x))\n\nSubmodules assigned in this way will be registered, and will also have their\nparameters converted when you call :meth:`to`, etc.\n\n.. note::\n As per the example above, an ``__init__()`` call to the parent class\n must be made before assignment on the child.\n\n:ivar training: Boolean represents whether this module is in training or\n evaluation mode.\n:vartype training: bool", - "signature": { - "parameters": [ - { - "name": "args", - "kind": "VAR_POSITIONAL", - "default": null, - "annotation": "Any" - }, - { - "name": "kwargs", - "kind": "VAR_KEYWORD", - "default": null, - "annotation": "Any" - } - ], - "return_annotation": "None" - }, - "attributes": { - "has_backward": null, - "can_torch_compile": null - } - }, - { - "name": "recurrent_kimi_delta_attention", - "qualname": "recurrent_kimi_delta_attention", - "module": "_fla_xpu_fe2bb66.layers", - "kind": "class", - "doc": "Base class for all neural network modules.\n\nYour models should also subclass this class.\n\nModules can also contain other Modules, allowing them to be nested in\na tree structure. You can assign the submodules as regular attributes::\n\n import torch.nn as nn\n import torch.nn.functional as F\n\n\n class Model(nn.Module):\n def __init__(self) -> None:\n super().__init__()\n self.conv1 = nn.Conv2d(1, 20, 5)\n self.conv2 = nn.Conv2d(20, 20, 5)\n\n def forward(self, x):\n x = F.relu(self.conv1(x))\n return F.relu(self.conv2(x))\n\nSubmodules assigned in this way will be registered, and will also have their\nparameters converted when you call :meth:`to`, etc.\n\n.. note::\n As per the example above, an ``__init__()`` call to the parent class\n must be made before assignment on the child.\n\n:ivar training: Boolean represents whether this module is in training or\n evaluation mode.\n:vartype training: bool", - "signature": { - "parameters": [ - { - "name": "args", - "kind": "VAR_POSITIONAL", - "default": null, - "annotation": "Any" - }, - { - "name": "kwargs", - "kind": "VAR_KEYWORD", - "default": null, - "annotation": "Any" - } - ], - "return_annotation": "None" - }, - "attributes": { - "has_backward": null, - "can_torch_compile": null - } - }, - { - "name": "chunk_gated_delta_rule", - "qualname": "chunk_gated_delta_rule", - "module": "_fla_xpu_fe2bb66.layers", - "kind": "class", - "doc": "Base class for all neural network modules.\n\nYour models should also subclass this class.\n\nModules can also contain other Modules, allowing them to be nested in\na tree structure. You can assign the submodules as regular attributes::\n\n import torch.nn as nn\n import torch.nn.functional as F\n\n\n class Model(nn.Module):\n def __init__(self) -> None:\n super().__init__()\n self.conv1 = nn.Conv2d(1, 20, 5)\n self.conv2 = nn.Conv2d(20, 20, 5)\n\n def forward(self, x):\n x = F.relu(self.conv1(x))\n return F.relu(self.conv2(x))\n\nSubmodules assigned in this way will be registered, and will also have their\nparameters converted when you call :meth:`to`, etc.\n\n.. note::\n As per the example above, an ``__init__()`` call to the parent class\n must be made before assignment on the child.\n\n:ivar training: Boolean represents whether this module is in training or\n evaluation mode.\n:vartype training: bool", - "signature": { - "parameters": [ - { - "name": "args", - "kind": "VAR_POSITIONAL", - "default": null, - "annotation": "Any" - }, - { - "name": "kwargs", - "kind": "VAR_KEYWORD", - "default": null, - "annotation": "Any" - } - ], - "return_annotation": "None" - }, - "attributes": { - "has_backward": null, - "can_torch_compile": null - } - }, - { - "name": "recurrent_gated_delta_rule", - "qualname": "recurrent_gated_delta_rule", - "module": "_fla_xpu_fe2bb66.layers", - "kind": "class", - "doc": "Base class for all neural network modules.\n\nYour models should also subclass this class.\n\nModules can also contain other Modules, allowing them to be nested in\na tree structure. You can assign the submodules as regular attributes::\n\n import torch.nn as nn\n import torch.nn.functional as F\n\n\n class Model(nn.Module):\n def __init__(self) -> None:\n super().__init__()\n self.conv1 = nn.Conv2d(1, 20, 5)\n self.conv2 = nn.Conv2d(20, 20, 5)\n\n def forward(self, x):\n x = F.relu(self.conv1(x))\n return F.relu(self.conv2(x))\n\nSubmodules assigned in this way will be registered, and will also have their\nparameters converted when you call :meth:`to`, etc.\n\n.. note::\n As per the example above, an ``__init__()`` call to the parent class\n must be made before assignment on the child.\n\n:ivar training: Boolean represents whether this module is in training or\n evaluation mode.\n:vartype training: bool", - "signature": { - "parameters": [ - { - "name": "args", - "kind": "VAR_POSITIONAL", - "default": null, - "annotation": "Any" - }, - { - "name": "kwargs", - "kind": "VAR_KEYWORD", - "default": null, - "annotation": "Any" - } - ], - "return_annotation": "None" - }, - "attributes": { - "has_backward": null, - "can_torch_compile": null - } - } - ] -} diff --git a/build/torch-xpu/utils/__init__.py b/build/torch-xpu/utils/__init__.py deleted file mode 100644 index 0d2fbe5e9295cb9b75e2eb8fc86e417468a9e632..0000000000000000000000000000000000000000 --- a/build/torch-xpu/utils/__init__.py +++ /dev/null @@ -1,85 +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 sys - -from ._compat import ( # noqa: F401 - SUPPORTS_AUTOTUNE_CACHE, - TRITON_ABOVE_3_4_0, - TRITON_ABOVE_3_5_1, - TRITON_ABOVE_3_7_1, - autotune_cache_kwargs, - find_spec_cached, -) -from ._config import ( # noqa: F401 - FLA_CACHE_RESULTS, - FLA_CI_ENV, - FLA_DISABLE_TENSOR_CACHE, - FLA_TENSOR_CACHE_SIZE, -) -from ._decorators import ( # noqa: F401 - Action, - checkpoint, - contiguous, - deprecate_kwarg, - input_guard, - require_version, - tensor_cache, -) -from ._device import ( # noqa: F401 - IS_AMD, - IS_ARM, - IS_GATHER_SUPPORTED, - IS_INTEL, - IS_INTEL_ALCHEMIST, - IS_NPU, - IS_NVIDIA, - IS_NVIDIA_BLACKWELL, - IS_NVIDIA_HOPPER, - IS_TF32_SUPPORTED, - IS_TMA_SUPPORTED, - Backend, - autocast_custom_bwd, - autocast_custom_fwd, - check_environments, - check_pytorch_version, - check_shared_mem, - custom_device_ctx, - device, - device_name, - device_platform, - device_torch_lib, - get_all_max_shared_mem, - get_available_device, - get_multiprocessor_count, - map_triton_backend_to_torch_device, -) -from ._testing import assert_close, get_abs_err, get_err_ratio # noqa: F401 - - -def _register_aliases(): - current_module = sys.modules[__name__] - for key in ( - 'IS_AMD', - 'IS_ARM', - 'IS_INTEL', - 'IS_INTEL_ALCHEMIST', - 'IS_NVIDIA', - 'IS_NPU', - 'IS_NVIDIA_BLACKWELL', - 'IS_NVIDIA_HOPPER', - 'IS_TF32_SUPPORTED', - 'IS_GATHER_SUPPORTED', - 'IS_TMA_SUPPORTED', - ): - if hasattr(current_module, key): - setattr(current_module, key.lower(), getattr(current_module, key)) - - -_register_aliases() - -del _register_aliases diff --git a/build/torch-xpu/utils/_compat.py b/build/torch-xpu/utils/_compat.py deleted file mode 100644 index a702e300e509b8d436f2c9349daa696ded2b5a9c..0000000000000000000000000000000000000000 --- a/build/torch-xpu/utils/_compat.py +++ /dev/null @@ -1,27 +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 functools -import inspect -from importlib.util import find_spec - -import triton -from packaging import version as package_version - -from ._config import FLA_CACHE_RESULTS - -TRITON_ABOVE_3_4_0 = package_version.parse(triton.__version__) >= package_version.parse("3.4.0") -TRITON_ABOVE_3_5_1 = package_version.parse(triton.__version__) >= package_version.parse("3.5.1") -TRITON_ABOVE_3_7_1 = package_version.parse(triton.__version__) >= package_version.parse("3.7.1") - -SUPPORTS_AUTOTUNE_CACHE = "cache_results" in inspect.signature(triton.autotune).parameters -autotune_cache_kwargs = {"cache_results": FLA_CACHE_RESULTS} if SUPPORTS_AUTOTUNE_CACHE else {} - - -@functools.cache -def find_spec_cached(name): - return find_spec(name) diff --git a/build/torch-xpu/utils/_config.py b/build/torch-xpu/utils/_config.py deleted file mode 100644 index 79b6ce8a784f457b847834a6f3f6128a06bcf3f3..0000000000000000000000000000000000000000 --- a/build/torch-xpu/utils/_config.py +++ /dev/null @@ -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 - -import os - -FLA_CI_ENV = os.getenv("FLA_CI_ENV") == "1" -FLA_CACHE_RESULTS = os.getenv('FLA_CACHE_RESULTS', '1') == '1' - -FLA_DISABLE_TENSOR_CACHE = os.getenv('FLA_DISABLE_TENSOR_CACHE', '0') == '1' -try: - FLA_TENSOR_CACHE_SIZE = int(os.getenv('FLA_TENSOR_CACHE_SIZE', "4")) -except ValueError: - FLA_TENSOR_CACHE_SIZE = 4 diff --git a/build/torch-xpu/utils/_device.py b/build/torch-xpu/utils/_device.py deleted file mode 100644 index 0064c1939421d3831978e0210a63afa94a9fb36a..0000000000000000000000000000000000000000 --- a/build/torch-xpu/utils/_device.py +++ /dev/null @@ -1,217 +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 contextlib -import functools -import logging -import os -import platform -import sys -import warnings -from enum import Enum -from functools import cache, lru_cache - -import torch -import triton -from packaging import version as package_version - -logger = logging.getLogger(__name__) - - -@lru_cache(maxsize=1) -def check_environments(): - """ - Checks the current operating system, Triton version, and Python version, - issuing warnings if they don't meet recommendations. - This function's body only runs once due to lru_cache. - """ - # Check Operating System - if sys.platform == 'win32': - # Check if triton-windows is installed - try: - from importlib.metadata import PackageNotFoundError, metadata - metadata('triton-windows') - # triton-windows is installed, no warning needed - except PackageNotFoundError: - logger.warning( - "Detected Windows operating system. Consider installing triton-windows " - "(https://github.com/triton-lang/triton-windows) for better compatibility. " - "Without it, some features may not work correctly.", - ) - - triton_version = package_version.parse(triton.__version__) - required_triton_version = package_version.parse("3.3.0") - - if triton_version < required_triton_version: - logger.warning( - f"Current Triton version {triton_version} is below the recommended 3.3.0 version. " - "Errors may occur and these issues will not be fixed. " - "Please consider upgrading Triton.", - ) - - # Check Python version - py_version = package_version.parse(f"{sys.version_info.major}.{sys.version_info.minor}") - required_py_version = package_version.parse("3.11") - - if py_version < required_py_version: - logger.warning( - f"Current Python version {py_version} is below the recommended 3.11 version. " - "It is recommended to upgrade to Python 3.11 or higher for the best experience.", - ) - - return None - - -check_environments() - - -def _cpu_device_warning(): - warnings.warn(('Triton is not supported on current platform, roll back to CPU.'), stacklevel=2) - - -@cache -def check_pytorch_version(version_s: str = '2.4') -> bool: - return package_version.parse(torch.__version__) >= package_version.parse(version_s) - - -@cache -def get_multiprocessor_count(tensor_idx: int = 0) -> int: - try: - return triton.runtime.driver.active.utils.get_device_properties(tensor_idx)['multiprocessor_count'] - except Exception: - # Maybe we use a NPU device. - try: - if triton.runtime.driver.active.get_current_target().backend == 'npu': - return triton.runtime.driver.active.utils.get_device_properties(tensor_idx)['num_vectorcore'] - except Exception: - logger.debug('Failed to get NPU multiprocessor count, falling back to 1.', exc_info=True) - return 1 - - -@cache -def get_available_device() -> str: - try: - return triton.runtime.driver.active.get_current_target().backend - except Exception: - _cpu_device_warning() - return 'cpu' - - -def map_triton_backend_to_torch_device() -> str: - backend = get_available_device() # 'cuda' | 'hip' | 'xpu' | 'cpu' | ... - return {'cuda': 'cuda', 'hip': 'cuda', 'xpu': 'xpu'}.get(backend, backend) - - -# For AMD GPUs, the triton backend is 'hip', while for Nvidia GPUs, the triton backend is 'cuda'. -# However, the torch backend is 'cuda' for both Nvidia and AMD GPUs. -# Therefore, we need to check the triton backend to determine the actual GPU vendor. -device = get_available_device() if get_available_device() != 'hip' else 'cuda' -device_torch_lib = getattr(torch, device) -device_platform = get_available_device() -device_name = map_triton_backend_to_torch_device() - -IS_AMD = (device_platform == 'hip') -IS_ARM = platform.machine().lower() in ('aarch64', 'arm64') -IS_INTEL = (device_platform == 'xpu') -IS_INTEL_ALCHEMIST = (IS_INTEL and 'Intel(R) Arc(TM) A' in torch.xpu.get_device_name(0)) -IS_NVIDIA = (device_platform == 'cuda') -IS_NPU = (device_platform == 'npu') -IS_NVIDIA_BLACKWELL = (IS_NVIDIA and torch.cuda.get_device_capability()[0] in (10, 12)) -IS_NVIDIA_HOPPER = (IS_NVIDIA and ('NVIDIA H' in torch.cuda.get_device_name(0) or torch.cuda.get_device_capability()[0] == 9)) - -# Nvidia Ampere or newer, haven't check AMD and intel yet. -IS_TF32_SUPPORTED = (IS_NVIDIA and torch.cuda.get_device_capability(0)[0] >= 8) -IS_GATHER_SUPPORTED = hasattr(triton.language, 'gather') -IS_TMA_SUPPORTED = ( - IS_NVIDIA - and torch.cuda.get_device_capability(0)[0] >= 9 - and os.environ.get('FLA_USE_TMA', '0') == '1' - and (hasattr(triton.language, '_experimental_make_tensor_descriptor') or hasattr(triton.language, 'make_tensor_descriptor')) -) - -if IS_NVIDIA and not IS_TF32_SUPPORTED: - # Make old card happy, since triton will use tf32 by default. - # This is a workaround for old nvidia card. - os.environ['TRITON_F32_DEFAULT'] = 'ieee' - - -def _default_alloc_fn(size: int, alignment: int, stream: int | None): - return torch.empty(size, device=torch.device(device_name, device_torch_lib.current_device()), dtype=torch.int8) - - -if IS_TMA_SUPPORTED: - logger.info('TMA is supported, using TMA by default.') - triton.set_allocator(_default_alloc_fn) -elif IS_NVIDIA_BLACKWELL: - # Blackwell (SM100 datacenter / SM120 consumer): Triton compiler may emit global_scratch for - # autotuned kernels even without TMA. Register a default allocator to - # prevent NullAllocator crashes. See triton-lang/triton#10002. - logger.info('Blackwell detected: registering default global_scratch allocator.') - triton.set_allocator(_default_alloc_fn) - - -def get_all_max_shared_mem(): - try: - return [ - triton.runtime.driver.active.utils.get_device_properties(i)['max_shared_mem'] - for i in range(device_torch_lib.device_count()) - ] - except Exception: - _cpu_device_warning() - return [-1] - - -class Backend(Enum): - ADA = 101376 # RTX 4090 - AMPERE = 166912 # A100 - HOPPER = 232448 # H100 - DEFAULT = 102400 # Default - - @classmethod - def get_shared_memory(cls, arch: str) -> int: - try: - return cls[arch.upper()].value - except KeyError: - return cls.DEFAULT.value - - -@cache -def check_shared_mem(arch: str = "none", tensor_idx: int = 0) -> bool: - try: - device_shared_mem_list = get_all_max_shared_mem() - max_shared_memory = device_shared_mem_list[tensor_idx] - return max_shared_memory >= Backend.get_shared_memory(arch) - except Exception: - return False - - -if check_pytorch_version('2.4'): - if device == 'cpu': - device = 'cuda' - device_torch_lib = getattr(torch, device) - autocast_custom_fwd = functools.partial(torch.amp.custom_fwd, device_type=device) - autocast_custom_bwd = functools.partial(torch.amp.custom_bwd, device_type=device) - - def custom_device_ctx(index: int): - if index is None: - return contextlib.nullcontext() - try: - return device_torch_lib.device(index) - except (AttributeError, AssertionError, RuntimeError): - return contextlib.nullcontext() -else: - assert device == 'cuda', 'Only cuda device is supported for PyTorch version < 2.4.0.' - autocast_custom_fwd = device_torch_lib.amp.custom_fwd - autocast_custom_bwd = device_torch_lib.amp.custom_bwd - - def custom_device_ctx(index: int): - if index is None: - return contextlib.nullcontext() - try: - return torch.cuda.device(index) - except (AttributeError, AssertionError, RuntimeError): - return contextlib.nullcontext() diff --git a/build/torch-xpu/utils/ascend_ub_manager.py b/build/torch-xpu/utils/ascend_ub_manager.py deleted file mode 100644 index e1760b4345e19798339058537e7bb0ba45024a1d..0000000000000000000000000000000000000000 --- a/build/torch-xpu/utils/ascend_ub_manager.py +++ /dev/null @@ -1,576 +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 -# copied from https://github.com/linkedin/Liger-Kernel/blob/main/src/liger_kernel/ops/backends/_ascend/ub_manager.py - -""" -Unified Buffer (UB) Manager for Ascend NPU. - -This module provides UB capacity detection and tiling strategy computation -for running Triton kernels on Ascend NPU. It automatically calculates -optimal block sizes based on UB capacity constraints to prevent UB overflow. -""" - -import os -import warnings - -import torch -import triton - -# Ascend Triton launch limits (see triton_ascend/activations.py). -ASCEND_MAX_GRID_DIM = 65535 -# Legacy fused kernel byte cap (65536 // fp16 element size). -_FALLBACK_MAX_FUSED_BLOCK = 65536 // 2 -# Conservative UB fallback when runtime detection is unavailable (64 KiB). -_FALLBACK_UB_CAPACITY_BITS = 65536 * 8 - - -def is_npu_available() -> bool: - """Detect Ascend NPU availability.""" - if hasattr(torch, 'npu'): - try: - return torch.npu.is_available() - except Exception: - pass - try: - from transformers.utils import is_torch_npu_available - - return is_torch_npu_available() - except ImportError: - return False - - -def _fallback_ub_capacity(reason: str) -> int: - warnings.warn( - f"Using fallback UB capacity ({_FALLBACK_UB_CAPACITY_BITS // 8} bytes): {reason}", - stacklevel=3, - ) - return _FALLBACK_UB_CAPACITY_BITS - - -def _normalize_tiling_dims(tiling_dim: int | tuple[int, ...]) -> set: - """ - Normalize tiling dimension specification to a set of dimension indices. - - Args: - tiling_dim: Either an int (single dimension) or tuple of ints (multiple dimensions). - - Returns: - Set of dimension indices that can be tiled. - """ - if isinstance(tiling_dim, int): - return {tiling_dim} - elif isinstance(tiling_dim, tuple): - return set(tiling_dim) - else: - return set() - - -def _default_strategy( - ub_capacity_bits: int, - safety_margin: float, - dtype_size: int, - memory_multiplier: float, - shapes: tuple[tuple[int, ...], ...], - tiling_dims: tuple[int | tuple[int, ...], ...], -) -> tuple[int, ...]: - """ - Default tiling strategy: calculate maximum safe block size based on UB capacity. - - This is a unified strategy function that works for all kernels by abstracting - the memory calculation as: memory_multiplier * BLOCK_SIZE * unit_param * dtype_size * 8 bits - - Args: - ub_capacity_bits: UB capacity in bits - safety_margin: Safety margin as a float (e.g., 0.80 for 80%) - dtype_size: Size of data type in bytes (e.g., 2 for float16, 4 for float32) - memory_multiplier: Memory multiplier for estimating peak memory usage - shapes: Tuple of full shapes. Each shape is a tuple of dimension sizes. - - For ROPE: ((n_q_head, hd), (n_kv_head, hd)) - - For GEGLU: ((n_cols,),) - tiling_dims: Tuple specifying which dimensions can be tiled for each shape. - Each element can be: - - int: single dimension index (e.g., 0 for first dimension) - - tuple of ints: multiple dimensions that can be tiled together - - For ROPE: (0, 0) means first dimension of each shape can be tiled - - For GEGLU: (0,) means first dimension of the shape can be tiled - Length must match len(shapes). - - Returns: - Tuple of maximum safe block sizes, one for each shape. - Each element is a power of 2. - - Note: - For each shape, fixed dimensions (non-tiling) are multiplied together to get unit_param. - The final block size is computed in compute_default_tiling_strategy by taking - min(desired_block_size, max_safe_block_size) where desired_block_size = triton.next_power_of_2(original_dim). - """ - if not shapes or not tiling_dims: - return () - - # Calculate max_safe_block_size for each tiling dimension - max_safe_sizes = [] - - for shape, tiling_dim in zip(shapes, tiling_dims): - # Normalize tiling_dim to a set of dimension indices - tiling_dim_set = _normalize_tiling_dims(tiling_dim) - - # Validate tiling dimensions are within shape bounds - if not tiling_dim_set: - raise ValueError( - f"Invalid tiling_dim: {tiling_dim}. tiling_dim must be an int or a non-empty tuple of ints." - ) - if any(dim_idx < 0 or dim_idx >= len(shape) for dim_idx in tiling_dim_set): - raise ValueError( - f"Invalid tiling_dim: {tiling_dim} for shape {shape}. " - f"All dimension indices must be in range [0, {len(shape)})." - ) - - # Calculate unit_param: product of fixed (non-tiling) dimensions - unit_param = 1.0 - for dim_idx, dim_size in enumerate(shape): - if dim_idx not in tiling_dim_set: - if dim_size <= 0: - # Invalid dimension size, use conservative default - unit_param = 1.0 - break - unit_param *= float(dim_size) - - # Ensure unit_param is at least 1.0 - if unit_param <= 0: - unit_param = 1.0 - - # Calculate maximum safe block size based on UB capacity - # Memory: memory_multiplier * BLOCK_SIZE * unit_param * dtype_size * 8 bits - SAFE_UB_CAPACITY_BITS = int(ub_capacity_bits * safety_margin) - - # Solve: memory_multiplier * BLOCK_SIZE * unit_param * dtype_size * 8 <= SAFE_UB_CAPACITY_BITS - # BLOCK_SIZE <= SAFE_UB_CAPACITY_BITS / (memory_multiplier * unit_param * dtype_size * 8) - max_block_size = int(SAFE_UB_CAPACITY_BITS // (memory_multiplier * unit_param * dtype_size * 8)) - max_block_size = max(1, max_block_size) - - # Find largest power of 2 <= max_block_size - # Use triton.next_power_of_2(max_block_size + 1) // 2 to get the largest power of 2 <= max_block_size - safe_block_size = triton.next_power_of_2(max_block_size + 1) // 2 - max_safe_sizes.append(safe_block_size) - - return tuple(max_safe_sizes) - - -class UBManager: - """ - Unified Buffer Manager for Ascend NPU. - - Provides UB capacity detection and management for Ascend NPU devices. - The UB capacity is used by tiling strategy functions to calculate optimal block sizes. - """ - - def __init__(self, ub_capacity_bits: int | None = None): - """ - Initialize UB Manager. - - Args: - ub_capacity_bits: UB capacity in bits. If None, will be detected automatically. - """ - self._npu_model = self._detect_npu_model() - self._ub_capacity_bits = ub_capacity_bits or self._detect_ub_capacity() - - @property - def ub_capacity_bits(self) -> int: - """Get UB capacity in bits.""" - return self._ub_capacity_bits - - @property - def ub_capacity_bytes(self) -> int: - """Get UB capacity in bytes.""" - return self._ub_capacity_bits // 8 - - @property - def npu_model(self) -> str: - """Get detected NPU model name.""" - return self._npu_model - - def _detect_npu_model(self) -> str: - """Detect NPU model from device properties.""" - if not is_npu_available(): - return "unknown" - - try: - dev_props = torch.npu.get_device_properties(0) - # Try to get model name from device properties - return dev_props.name - except Exception: - pass - - return "default" - - def _detect_ub_capacity(self) -> int: - """ - Detect UB capacity from environment variable or get_soc_spec. - - Returns: - UB capacity in bits. Falls back to a conservative default with a warning - when detection fails. - """ - # Check environment variable first (in bits) - env_capacity = os.getenv("ASCEND_UB_CAPACITY_BITS") - if env_capacity is not None: - try: - capacity_bits = int(env_capacity) - if capacity_bits > 0: - return capacity_bits - except ValueError: - pass - - # Try to get from get_soc_spec (returns bytes, convert to bits) - if is_npu_available(): - try: - from tbe.common.platform import get_soc_spec, set_current_compile_soc_info - - # Set current SOC info for get_soc_spec to work correctly - device = torch.npu - soc_info = device.get_device_name(device.current_device()) - set_current_compile_soc_info(soc_info) - - # Query UB size (get_soc_spec returns size in bytes) - ub_size_bytes = get_soc_spec("UB_SIZE") - - if ub_size_bytes is None or ub_size_bytes <= 0: - raise ValueError(f"Invalid UB_SIZE from get_soc_spec: {ub_size_bytes}") - - # Convert bytes to bits - ub_capacity_bits = ub_size_bytes * 8 - return ub_capacity_bits - - except ImportError: - return _fallback_ub_capacity( - "Cannot import tbe.common.platform.get_soc_spec. " - "Source CANN set_env.sh or set ASCEND_UB_CAPACITY_BITS." - ) - except Exception as e: - return _fallback_ub_capacity( - f"Failed to detect UB capacity from get_soc_spec: {e}. " - "Set ASCEND_UB_CAPACITY_BITS to override." - ) - - return _fallback_ub_capacity( - "NPU is not available. Set ASCEND_UB_CAPACITY_BITS to override." - ) - - -# Global singleton instance -_ub_manager: UBManager | None = None - - -def get_ub_manager() -> UBManager: - """Get global UB manager instance.""" - global _ub_manager - if _ub_manager is None: - _ub_manager = UBManager() - return _ub_manager - - -def compute_default_tiling_strategy( - safety_margin: float = 0.80, - dtype_size: int | None = None, - memory_multiplier: float | None = None, - shapes: tuple[tuple[int, ...], ...] | None = None, - tiling_dims: tuple[int | tuple[int, ...], ...] | None = None, -) -> tuple[tuple[int, ...], ...] | None: - """ - Compute tiling strategy using the default strategy function. - - This function directly calls the default strategy and computes the final - tiling result. All kernels use the same unified strategy function, so - there's no need for kernel_name-based lookup. - - Args: - safety_margin: Safety margin as a float (e.g., 0.80 for 80%). Default is 0.80. - dtype_size: Size of data type in bytes (e.g., 2 for float16, 4 for float32). - Must be provided. If None or <= 0, defaults to 4 (float32). - memory_multiplier: Memory multiplier for estimating peak memory usage. - - For GEGLU: typically 10.0 for backward, 4.0 for forward - - For ROPE: typically 3.0 - If None, defaults to 10.0 (conservative estimate). - shapes: Tuple of full shapes. Each shape is a tuple of dimension sizes. - - For ROPE: ((n_q_head, hd), (n_kv_head, hd)) - - For GEGLU: ((n_cols,),) - Can pass original shapes (will handle padding internally) or padded shapes. - tiling_dims: Tuple specifying which dimensions can be tiled for each shape. - Each element can be: - - int: single dimension index (e.g., 0 for first dimension) - - tuple of ints: multiple dimensions that can be tiled together - - For ROPE: (0, 0) means first dimension of each shape can be tiled - - For GEGLU: (0,) means first dimension of the shape can be tiled - Length must match len(shapes). Cannot be empty. - - Returns: - Tuple of tiled shapes with same structure as input shapes. - Tiling dimensions are replaced with computed block sizes (power of 2), - while non-tiling dimensions are padded to next power of 2. - - For ROPE: ((block_size_q, pad_hd), (block_size_kv, pad_hd)) - - For GEGLU: ((block_size,),) - Returns None if shapes or tiling_dims is None or empty. - - Examples: - >>> # ROPE forward - >>> strategy = compute_default_tiling_strategy( - ... safety_margin=0.90, - ... dtype_size=4, - ... memory_multiplier=3.0, - ... shapes=((32, 128), (32, 128)), - ... tiling_dims=(0, 0) - ... ) - >>> # Returns: ((block_size_q, 128), (block_size_kv, 128)) - >>> # GEGLU forward - >>> strategy = compute_default_tiling_strategy( - ... safety_margin=0.80, - ... dtype_size=2, - ... memory_multiplier=7.0, - ... shapes=((4096,),), - ... tiling_dims=(0,) - ... ) - >>> # Returns: ((block_size,),) - """ - ub_manager = get_ub_manager() - - if shapes is None or not shapes or tiling_dims is None or not tiling_dims: - return None - - if len(shapes) != len(tiling_dims): - return None - - if dtype_size is None or dtype_size <= 0: - dtype_size = 4 # Default to float32 - - if memory_multiplier is None or memory_multiplier <= 0: - memory_multiplier = 10.0 # Default conservative estimate - - # Call strategy to get max_safe_block_size for each shape - max_supported = _default_strategy( - ub_manager.ub_capacity_bits, - safety_margin, - dtype_size, - memory_multiplier, - shapes, - tiling_dims, - ) - - if not max_supported or len(max_supported) != len(shapes): - return None - - # Build result: same structure as shapes, with tiling dims replaced by computed block sizes - result = [] - for shape, tiling_dim, max_safe in zip(shapes, tiling_dims, max_supported): - result_shape = list(shape) - - # Normalize tiling_dim to a set of dimension indices - tiling_dim_set = _normalize_tiling_dims(tiling_dim) - - # Validate tiling dimensions are within shape bounds - if not tiling_dim_set: - raise ValueError( - f"Invalid tiling_dim: {tiling_dim}. tiling_dim must be an int or a non-empty tuple of ints." - ) - if any(dim_idx < 0 or dim_idx >= len(result_shape) for dim_idx in tiling_dim_set): - raise ValueError( - f"Invalid tiling_dim: {tiling_dim} for shape {shape}. " - f"All dimension indices must be in range [0, {len(result_shape)})." - ) - - # Replace tiling dimensions with computed block sizes - # For each tiling dimension, compute: min(desired, max_safe) - for dim_idx in tiling_dim_set: - original_dim = result_shape[dim_idx] - desired = triton.next_power_of_2(original_dim) - final_val = min(desired, max_safe) - final_val = max(1, final_val) # Ensure at least 1 - result_shape[dim_idx] = final_val - - # Pad non-tiling dimensions to next power of 2 - for dim_idx, dim_size in enumerate(result_shape): - if dim_idx not in tiling_dim_set: - result_shape[dim_idx] = triton.next_power_of_2(dim_size) - - result.append(tuple(result_shape)) - - return tuple(result) - - -def compute_ub_block_size( - dim_size: int, - memory_multiplier: float, - *, - safety_margin: float = 0.9, - dtype_size: int = 4, - fallback: int = 2048, - min_block: int = 1, - max_block: int | None = None, - desired: int | None = None, -) -> int: - """Compute UB-safe block size for a single tilable dimension.""" - if desired is None: - desired = triton.next_power_of_2(dim_size) - - tile_shapes = compute_default_tiling_strategy( - safety_margin=safety_margin, - dtype_size=dtype_size, - memory_multiplier=memory_multiplier, - shapes=((dim_size,),), - tiling_dims=(0,), - ) - block = min(desired, tile_shapes[0][0]) if tile_shapes else min(desired, fallback) - block = max(min_block, block) - if max_block is not None: - block = min(block, max_block) - return block - - -def compute_vocab_block_size( - vocab_size: int, - num_rows: int, - memory_multiplier: float, - *, - safety_margin: float = 0.9, - max_block: int = 8192, - fallback: int = 2048, -) -> int: - """UB-safe vocab tile size respecting Ascend grid dim0 limit.""" - ub_block = compute_ub_block_size( - vocab_size, - memory_multiplier, - safety_margin=safety_margin, - fallback=fallback, - ) - max_splits = max(1, ASCEND_MAX_GRID_DIM // max(num_rows, 1)) - grid_min = triton.next_power_of_2((vocab_size + max_splits - 1) // max_splits) - return min(max(ub_block, grid_min), max_block) - - -def compute_elementwise_block_size( - n_elements: int, - memory_multiplier: float = 2.5, - *, - safety_margin: float = 0.9, - min_block: int = 1024, - fallback: int | None = None, -) -> int: - """UB-safe block size for elementwise kernels under grid limit.""" - if fallback is None: - fallback = _FALLBACK_MAX_FUSED_BLOCK - ub_block = compute_ub_block_size( - n_elements, - memory_multiplier, - safety_margin=safety_margin, - fallback=fallback, - ) - grid_min = max( - min_block, - triton.next_power_of_2((n_elements + ASCEND_MAX_GRID_DIM - 1) // ASCEND_MAX_GRID_DIM), - ) - return min(ub_block, grid_min) - - -def compute_activation_block_size( - total_elements: int, - is_backward: bool, - *, - max_grid: int = ASCEND_MAX_GRID_DIM, - max_core_dim: int = 65535, - safety_margin: float = 0.9, - max_block: int = 2048, - memory_multiplier: float | None = None, -) -> int: - """UB-safe block size for flattened activation kernels.""" - if memory_multiplier is None: - memory_multiplier = 6.0 if is_backward else 3.0 - block = compute_ub_block_size( - total_elements, - memory_multiplier, - safety_margin=safety_margin, - fallback=2048, - min_block=256, - max_block=max_block, - ) - block = min(block, max_core_dim // 8) - if triton.cdiv(total_elements, block) > max_grid: - block = triton.cdiv(total_elements, max_grid) - block = max(triton.next_power_of_2(block), 1) - return block - - -def compute_row_tile_block_size( - row_dim: int, - fixed_dim: int, - memory_multiplier: float, - *, - tiling_row: bool = True, - safety_margin: float = 0.85, - dtype_size: int = 4, - fallback: int = 16, - min_block: int = 1, - max_block: int | None = None, -) -> int: - """UB-safe tile along one axis of a 2D [row, col] kernel tile.""" - if tiling_row: - shapes = ((row_dim, fixed_dim),) - tiling_dims = (0,) - desired = triton.next_power_of_2(row_dim) - else: - shapes = ((row_dim, fixed_dim),) - tiling_dims = (1,) - desired = triton.next_power_of_2(fixed_dim) - - tile_shapes = compute_default_tiling_strategy( - safety_margin=safety_margin, - dtype_size=dtype_size, - memory_multiplier=memory_multiplier, - shapes=shapes, - tiling_dims=tiling_dims, - ) - if tile_shapes: - block = tile_shapes[0][0] if tiling_row else tile_shapes[0][1] - block = min(desired, block) - else: - block = min(desired, fallback) - block = max(min_block, block) - if max_block is not None: - block = min(block, max_block) - return block - - -def max_grid_axis_chunks( - axis_size: int, - other_grid_product: int, - *, - max_grid: int = ASCEND_MAX_GRID_DIM, -) -> int: - """Max launch chunks along one grid axis while keeping the product <= max_grid.""" - return max(1, max_grid // max(other_grid_product, 1)) - - -def compute_grid_limited_tile_size( - axis_size: int, - other_grid_product: int, - ub_safe_block: int, - *, - max_grid: int = ASCEND_MAX_GRID_DIM, - min_block: int = 1, -) -> int: - """Pick a UB-safe tile size; host-side chunking handles grid overflow.""" - return max(min_block, min(ub_safe_block, axis_size)) - - -def iter_axis_launch_chunks( - axis_size: int, - other_grid_product: int, - *, - max_grid: int = ASCEND_MAX_GRID_DIM, -): - """Yield ``(offset, chunk_len)`` for host-side grid-axis tiling.""" - max_chunks = max_grid_axis_chunks(axis_size, other_grid_product, max_grid=max_grid) - for offset in range(0, axis_size, max_chunks): - yield offset, min(max_chunks, axis_size - offset) diff --git a/build/torch-xpu/utils/layers.py b/build/torch-xpu/utils/layers.py deleted file mode 100644 index cd945e917e12c3471e686940a23e0751982ac6e1..0000000000000000000000000000000000000000 --- a/build/torch-xpu/utils/layers.py +++ /dev/null @@ -1,332 +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 - -# Code is adapted from flash-attn.bert_padding.py - -# NOTE: This is a copy of `fla.layers.utils` but we do not sync the FLA layers in hf-kernels - -import torch -from einops import rearrange, repeat - -from ..ops.utils.index import ( - prepare_cu_seqlens_from_lens, - prepare_cu_seqlens_from_mask, - prepare_lens_from_mask, -) -from . import tensor_cache - -_LAYER_IDX_REQUIRED_MSG = "{cls} requires `layer_idx` when `past_key_values` is provided." - - -class IndexFirstAxis(torch.autograd.Function): - - @staticmethod - def forward(ctx, x, indices): - ctx.save_for_backward(indices) - assert x.ndim >= 2 - ctx.first_axis_dim, other_shape = x.shape[0], x.shape[1:] - second_dim = other_shape.numel() - # TD [2022-03-04] For some reason torch.gather is a bit faster than indexing. - # return x[indices] - return torch.gather( - rearrange(x, "b ... -> b (...)"), 0, repeat(indices, "z -> z d", d=second_dim), - ).reshape(-1, *other_shape) - - @staticmethod - def backward(ctx, do): - (indices,) = ctx.saved_tensors - assert do.ndim >= 2 - other_shape = do.shape[1:] - do = rearrange(do, "b ... -> b (...)") - dx = torch.zeros( - [ctx.first_axis_dim, do.shape[1]], - device=do.device, - dtype=do.dtype, - ) - # TD [2022-03-04] For some reason torch.scatter is a bit faster than indexing. - # dx[indices] = do - dx.scatter_(0, repeat(indices, "z -> z d", d=do.shape[1]), do) - return dx.reshape(ctx.first_axis_dim, *other_shape), None - - -index_first_axis = IndexFirstAxis.apply - - -class IndexPutFirstAxis(torch.autograd.Function): - - @staticmethod - def forward(ctx, x, indices, first_axis_dim): - ctx.save_for_backward(indices) - assert indices.ndim == 1 - assert x.ndim >= 2 - y = torch.zeros(first_axis_dim, *x.shape[1:], device=x.device, dtype=x.dtype) - # TODO [2022-03-04] For some reason torch.scatter is a bit faster than indexing. - y[indices] = x - # y.scatter_(0, repeat(indices, 'z -> z d', d=x.shape[1]), x) - return y - - @staticmethod - def backward(ctx, do): - (indices,) = ctx.saved_tensors - # TODO [2022-03-04] For some reason torch.gather is a bit faster than indexing. - dx = do[indices] - # dx = torch.gather(do, 0, repeat(indices, 'z -> z d', d=do.shape[1])) - return dx, None, None - - -index_put_first_axis = IndexPutFirstAxis.apply - - -@tensor_cache -def _get_unpad_indices_and_cu( - attention_mask: torch.Tensor, -) -> tuple[torch.Tensor, torch.Tensor]: - lens = prepare_lens_from_mask(attention_mask) - indices = torch.nonzero(attention_mask.flatten(), as_tuple=False).flatten() - cu_seqlens = prepare_cu_seqlens_from_lens(lens) - return indices, cu_seqlens - - -@tensor_cache -def _get_last_token_unpad_indices_and_cu( - attention_mask: torch.Tensor, -) -> tuple[torch.Tensor, torch.Tensor]: - return _get_unpad_indices_and_cu(attention_mask[:, -1:]) - - -def get_unpad_indices_and_cu( - attention_mask: torch.Tensor, - q_len: int | None = None, -) -> tuple[torch.Tensor, torch.Tensor]: - """Return unpadding indices and cumulative sequence lengths.""" - if q_len is None or attention_mask.shape[-1] == q_len: - return _get_unpad_indices_and_cu(attention_mask) - if q_len == 1: - return _get_last_token_unpad_indices_and_cu(attention_mask) - return _get_unpad_indices_and_cu(attention_mask[:, -q_len:]) - - -@tensor_cache -def get_unpad_data( - attention_mask: torch.Tensor, -) -> tuple[torch.Tensor, torch.Tensor, int]: - """ - Retrieves indexing data required to repad unpadded (ragged) tensors. - - Args: - attention_mask (`torch.Tensor`): - Boolean or int tensor of shape (batch_size, sequence_length), 1 means valid and 0 means not valid. - - Return: - indices (`torch.Tensor`): - The indices of non-masked tokens from the flattened input sequence. - cu_seqlens (`torch.Tensor`): - The cumulative sequence lengths, used to index into ragged (unpadded) tensors. - `cu_seqlens` shape is [batch_size + 1]. - max_seqlen_in_batch (`int`): - Maximum sequence length in batch. - """ - lens = prepare_lens_from_mask(attention_mask) - indices = torch.nonzero(attention_mask.flatten(), as_tuple=False).flatten() - max_seqlen_in_batch = lens.max().item() - cu_seqlens = prepare_cu_seqlens_from_mask(attention_mask) - return indices, cu_seqlens, max_seqlen_in_batch - - -def unpad_hidden_states( - hidden_states: torch.Tensor, - cu_seqlens: torch.LongTensor | None, - attention_mask: torch.Tensor | None, - q_len: int, -) -> tuple[torch.Tensor, torch.Tensor | None, torch.LongTensor | None]: - """Unpad hidden states via attention_mask, unless explicit cu_seqlens are given. - - Explicit `cu_seqlens` take precedence: when they are provided, the input is - returned unchanged and no unpadding happens. Only when `cu_seqlens` is None - and `attention_mask` is present are the hidden states unpadded via the mask. - - Arguments: - hidden_states (`torch.Tensor`): - Input hidden states with padding. Shape: [batch_size, q_len, ...]. - cu_seqlens (`torch.LongTensor | None`): - Explicit cumulative sequence lengths provided by the caller. - attention_mask (`torch.Tensor | None`): - Boolean or int tensor of shape (batch_size, sequence_length), 1 means valid and 0 means not valid. - q_len (`int`): - Target length. - - Return: - hidden_states (`torch.Tensor`): - The unpadded hidden states of shape [1, total_tokens, ...], - or the input unchanged when unpadding is skipped. - indices (`torch.Tensor | None`): - The indices of non-masked tokens, or None when unpadding is skipped. - cu_seqlens (`torch.LongTensor | None`): - The cumulative sequence lengths derived from the mask, - or the explicit value passed in when unpadding is skipped. - """ - if cu_seqlens is None and attention_mask is not None: - indices, cu_seqlens = get_unpad_indices_and_cu(attention_mask, q_len) - hidden_states = index_first_axis(rearrange(hidden_states, "b s ... -> (b s) ..."), indices).unsqueeze(0) - return hidden_states, indices, cu_seqlens - return hidden_states, None, cu_seqlens - - -def repad_hidden_states( - hidden_states: torch.Tensor, - indices: torch.Tensor | None, - batch_size: int, - q_len: int, -) -> torch.Tensor: - """Inverse of `unpad_hidden_states`; a no-op when `indices` is None. - - Arguments: - hidden_states (`torch.Tensor`): - Hidden states of shape [1, total_tokens, ...]. - indices (`torch.Tensor | None`): - The indices returned by `unpad_hidden_states`, or None. - batch_size (`int`): - Batch size for the padded sequence. - q_len (`int`): - Target length. - - Return: - hidden_states of shape [batch_size, q_len, ...], - or the input unchanged when `indices` is None. - """ - if indices is None: - return hidden_states - return pad_input(hidden_states.squeeze(0), indices, batch_size, q_len) - - -def unpad_input( - q: torch.Tensor | tuple[torch.Tensor, ...], - states: tuple[torch.Tensor], - attention_mask: torch.Tensor, - q_len: int, - keepdim: bool = False, -): - """ - Unpads query, key, and values tensors, using a single dimension for all tokens - even though they belong to different batches. - - - Arguments: - q (`torch.Tensor` or `Tuple[torch.Tensor]`): - Query state with padding. Shape: [batch_size, q_len, ...]. - When it is a tuple, do unpadding for each tensor in the tuple. - states (`Tuple[torch.Tensor]`): - Attention state with padding. Shape: [batch_size, seq_len, ...]. - attention_mask (`torch.Tensor`): - Boolean or int tensor of shape [batch_size, sequence_length], 1 means valid and 0 means not valid. - q_len (`int`): - Target length. - keepdim (`bool`): - Whether to keep the batch dimension. Default: `False`. - - Return: - q (`torch.Tensor` or `Tuple[torch.Tensor]`): - Query state without padding. - Shape: [1, total_target_length, ...] if `keepdim=True` else [total_target_length, ...]. - When the `q` passed in is a tuple, return a tuple of such unpadded tensors. - states (`Tuple[torch.Tensor]`): - Attention state without padding. - Shape: [1, total_source_length, ...] if `keepdim=True` else [total_source_length, ...]. - indices_q (`torch.Tensor`): - The indices of non-masked tokens from the flattened input target sequence. - (cu_seqlens_q, cu_seqlens_k) (`Tuple[torch.LongTensor, torch.LongTensor]`): - The cumulative sequence lengths for the target (query) and source (key, value), - used to index into ragged (unpadded) tensors. - `cu_seqlens` shape is [batch_size + 1]. - (max_seqlen_in_batch_q, max_seqlen_in_batch_k) (`Tuple[int, int]`): - Maximum sequence length in batch (`max_seqlen_in_batch_q` for the target sequence - i.e. query, `max_seqlen_in_batch_k` for the source sequence i.e. key/value). - """ - indices_k, cu_seqlens_k, max_seqlen_in_batch_k = get_unpad_data(attention_mask) - batch_size, seq_len, *_ = states[0].shape - - state = tuple( - index_first_axis(rearrange(s, "b s ... -> (b s) ..."), indices_k) - for s in states - ) - if isinstance(q, torch.Tensor): - q = (q,) - cast_tuple = True - else: - cast_tuple = False - - if q_len == seq_len: - q = tuple(index_first_axis(rearrange(q_, "b s ... -> (b s) ..."), indices_k) for q_ in q) - cu_seqlens_q = cu_seqlens_k - max_seqlen_in_batch_q = max_seqlen_in_batch_k - indices_q = indices_k - elif q_len == 1: - max_seqlen_in_batch_q = 1 - cu_seqlens_q = torch.arange(batch_size + 1, dtype=torch.int32, device=q[0].device) - indices_q = cu_seqlens_q[:-1] - q = tuple(q_.squeeze(1) for q_ in q) - else: - raise NotImplementedError("We only support either q_len == k_len (prefilling) or q_len == 1 (decoding)") - - if keepdim: - q = tuple(q_.unsqueeze(0) for q_ in q) - state = tuple(s.unsqueeze(0) for s in state) - if cast_tuple: - q = q[0] - - return ( - q, - state, - indices_q, - (cu_seqlens_q, cu_seqlens_k), - (max_seqlen_in_batch_q, max_seqlen_in_batch_k), - ) - - -def pad_input( - hidden_states: torch.Tensor, - indices: torch.LongTensor, - batch_size: int, - seq_len: int, -) -> torch.Tensor: - """ - Args: - hidden_states ([total_tokens, ...]): - where total_tokens denotes the number of tokens in selected in attention_mask. - indices ([total_tokens]): - the indices that represent the non-masked tokens of the original padded input sequence. - batch_size (int): - batch_size size for the padded sequence. - seq_len (int): - maximum sequence length for the padded sequence. - - Return: - hidden_states of shape [batch_size, seq_len, ...] - """ - output = index_put_first_axis(hidden_states, indices, batch_size * seq_len) - return rearrange(output, "(b s) ... -> b s ...", b=batch_size) - - -def require_cache_layer_idx(module, past_key_values): - layer_idx = getattr(module, "layer_idx", None) - if past_key_values is not None and layer_idx is None: - raise ValueError(_LAYER_IDX_REQUIRED_MSG.format(cls=module.__class__.__name__)) - return layer_idx - - -def get_layer_cache(module, past_key_values): - layer_idx = require_cache_layer_idx(module, past_key_values) - if past_key_values is not None and len(past_key_values) > layer_idx: - return past_key_values[layer_idx] - return None - - -def update_layer_cache(module, past_key_values, **kwargs): - layer_idx = require_cache_layer_idx(module, past_key_values) - if past_key_values is not None: - return past_key_values.update(layer_idx=layer_idx, **kwargs) - return None