diff --git a/fla/layers/rwkv7.py b/fla/layers/rwkv7.py index c6d73f512c..f7ae645998 100644 --- a/fla/layers/rwkv7.py +++ b/fla/layers/rwkv7.py @@ -317,6 +317,7 @@ def forward( output_final_state=use_cache, cu_seqlens=cu_seqlens, safe_gate=True, + lower_bound=-0.6065306597126334, chunk_size=64, ) else: diff --git a/fla/ops/generalized_delta_rule/dplr/backends/__init__.py b/fla/ops/generalized_delta_rule/dplr/backends/__init__.py new file mode 100644 index 0000000000..41a27dad63 --- /dev/null +++ b/fla/ops/generalized_delta_rule/dplr/backends/__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 + +"""DPLR backends.""" + +from fla.ops.backends import BackendRegistry, dispatch +from fla.ops.generalized_delta_rule.dplr.backends.tilelang import DPLRTileLangBackend + +dplr_registry = BackendRegistry("generalized_delta_rule.dplr") +dplr_registry.register(DPLRTileLangBackend()) + + +__all__ = ['dispatch', 'dplr_registry'] diff --git a/fla/ops/generalized_delta_rule/dplr/backends/tilelang/__init__.py b/fla/ops/generalized_delta_rule/dplr/backends/tilelang/__init__.py new file mode 100644 index 0000000000..1c2bea6399 --- /dev/null +++ b/fla/ops/generalized_delta_rule/dplr/backends/tilelang/__init__.py @@ -0,0 +1,256 @@ +# Copyright (c) 2023-2026, Songlin Yang, Yu Zhang, Zhiyuan Li +# +# This source code is licensed under the MIT license found in the +# LICENSE file in 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 DPLR operations.""" + +from __future__ import annotations + +import logging + +import torch + +from fla.ops.backends import BaseBackend +from fla.utils import ( + find_spec_cached, + get_device_capability, + get_device_smem_optin, + get_multiprocessor_count, + has_usable_nvcc, +) + +from .schedules import chunk64_schedule_or_none, stream_bwd_schedule_or_none + +logger = logging.getLogger(__name__) + +_TILELANG_AVAILABLE = find_spec_cached("tilelang") is not None + +_FALLBACK_LOGGED: set[str] = set() + + +class DPLRTileLangBackend(BaseBackend): + + backend_type = "tilelang" + package_name = "tilelang" + env_var = "FLA_TILELANG" + + @classmethod + def is_available(cls) -> bool: + return _TILELANG_AVAILABLE and has_usable_nvcc() + + def chunk_dplr_delta_rule_verifier( + self, + 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=None, + lower_bound: float | None = None, + **kwargs, + ) -> tuple[bool, str | None]: + if lower_bound is not None and lower_bound >= 0: + # rejected so the default implementation surfaces the ValueError + return False, "TileLang backend requires `lower_bound < 0`; fall back to Triton" + if cp_context is not None: + if initial_state is not None: + return False, "TileLang backend does not support initial_state with CP; fall back to Triton" + if output_final_state: + return False, "TileLang backend does not support output_final_state with CP; fall back to Triton" + if getattr(cp_context, "cu_seqlens", None) is None: + return False, "TileLang backend requires cu_seqlens for CP; fall back to Triton" + if q.dtype not in (torch.float16, torch.bfloat16): + return False, f"TileLang backend does not support dtype {q.dtype}; fall back to Triton" + if not all(t.dtype == q.dtype for t in (k, v, a, b)): + return False, ( + "TileLang backend requires k/v/a/b dtypes to match q.dtype " + f"(got q={q.dtype}, k={k.dtype}, v={v.dtype}, a={a.dtype}, b={b.dtype}); " + "fall back to Triton" + ) + if gk.dtype not in (torch.float16, torch.bfloat16, torch.float32): + return False, f"TileLang backend does not support gk dtype {gk.dtype}; fall back to Triton" + if k.shape[-1] != v.shape[-1]: + return False, ( + f"TileLang backend requires K == V (got K={k.shape[-1]}, V={v.shape[-1]}); " + "fall back to Triton" + ) + if k.shape[-1] not in (64, 128): + return False, f"TileLang backend supports head dim 64 or 128 (got {k.shape[-1]}); fall back to Triton" + chunk_size = 16 if chunk_size is None else chunk_size + if chunk_size not in (16, 32, 64): + return False, f"TileLang backend supports chunk_size 16/32/64, got {chunk_size}; fall back to Triton" + from fla.ops.generalized_delta_rule.dplr.chunk import gate_bound_is_safe + bound = lower_bound if lower_bound is not None else (-5.0 if safe_gate else None) + if bound is None or not gate_bound_is_safe(bound, chunk_size): + return False, ( + "TileLang backend requires safe_gate or a lower_bound that fits " + f"chunk_size {chunk_size}: its mid-chunk-centered tensor-core scheme " + "keeps exp2 operands in fp32 range only while " + "(chunk_size/2+1)*bound*log2(e) <= 124 (the documented safe_gate " + "range [-5, 0) fits chunk_size 16/32 but not 64); fall back to Triton" + ) + if not q.is_cuda: + return False, "TileLang backend is CUDA-only; fall back to Triton" + dev = q.device.index or 0 + cc_major, cc_minor = get_device_capability(dev) + smem_cap = get_device_smem_optin(dev) + K = k.shape[-1] + in_dtype = "float16" if q.dtype == torch.float16 else "bfloat16" + if chunk_size == 16 and K == 128: + # measured ~0.5x vs Triton (the non-vectorized A-stage and the + # 2-warp h+o path at BT=16 do not pay off at K=128) + return False, "TileLang backend is slower than Triton at chunk_size 16 with head dim 128; fall back to Triton" + if chunk_size == 16 and (cu_seqlens is not None or cp_context is not None) and cc_major < 12: + # measured 0.83-1.02x vs Triton on sm_90 (H800, D64, all varlen + # sizes): the BT=16 serial state pass does not amortize ragged + # chunks there. Rect cs16 and cc12x varlen cs16 win and stay + # accepted. + return False, ( + "TileLang backend is slower than Triton at chunk_size 16 on " + f"variable-length inputs on compute capability {cc_major}.{cc_minor}; " + "fall back to Triton" + ) + if chunk_size == 64: + # reject configs no BT=64 kernel schedule can launch on this + # device (e.g. the K=128 stream backward needs 167936B on A100's + # 166912B cap or cc120's 101376B cap, where K=64 fits via the + # low schedule); + # the arithmetic is shared with the launcher so acceptance + # implies schedulability + if chunk64_schedule_or_none(K=K, V=K, in_dtype=in_dtype, smem_cap=smem_cap, + cc=cc_major * 10 + cc_minor) is None: + return False, ( + f"TileLang backend has no launchable backward schedule for " + f"chunk_size 64 with head dim {K} on a device with {smem_cap}B " + "shared memory per block; fall back to Triton" + ) + if cc_major < 12: + # measured 0.68-1.05x vs Triton on sm_90 (H800, rect and + # varlen, both head dims): the BT=64 wy/wu UT-transform + # stages are latency-bound there, and the fused h+o pass + # does not amortize at twice the chunk serial length + return False, ( + "TileLang backend is slower than Triton at chunk_size 64 " + f"on compute capability {cc_major}.{cc_minor}; fall back to Triton" + ) + # The fused h+o state pass parallelizes only over (V/BV, N, H) blocks + # and walks chunks serially, so it loses to the split Triton kernels + # when the grid underfills the GPU. Measured crossover on PRO 6000 / + # H100 class parts is around half the SM count. + bv = 64 if v.shape[-1] <= 64 else 32 + if cp_context is not None: + n_seqs = len(cp_context.cu_seqlens) - 1 + elif cu_seqlens is not None: + n_seqs = len(cu_seqlens) - 1 + else: + n_seqs = q.shape[0] + grid = n_seqs * q.shape[2] * ((v.shape[-1] + bv - 1) // bv) + sm = get_multiprocessor_count(dev) + if grid < sm // 2: + return False, ( + f"TileLang backend is slower than Triton on small grids (N*H*(V/BV)={grid} " + f"< {sm // 2} SMs/2); fall back to Triton" + ) + stream_schedule = stream_bwd_schedule_or_none( + K=K, V=K, BT=chunk_size, in_dtype=in_dtype, + smem_cap=smem_cap, + ) + if stream_schedule is None: + return False, ( + f"TileLang backend has no launchable backward schedule for " + f"chunk_size {chunk_size} with head dim {K} on a device with {smem_cap}B " + "shared memory per block; fall back to Triton" + ) + if stream_schedule == "low" and n_seqs * q.shape[2] < sm // 2: + # the low-smem stream backward is one serial chunk scan per + # (seq, head) block with no V split, and its 97KB footprint leaves + # no room to prefetch; below half the SMs it cannot hide the + # serial chain and measurably loses to the split Triton kernels + return False, ( + f"TileLang backend is slower than Triton when the low-smem stream backward " + f"underfills the device (N*H={n_seqs * q.shape[2]} < {sm // 2} SMs/2); " + "fall back to Triton" + ) + return True, None + + def chunk_dplr_delta_rule( + self, + 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=None, + lower_bound: float | None = None, + **kwargs, + ) -> tuple[torch.Tensor, torch.Tensor | None]: + from fla.ops.generalized_delta_rule.dplr.backends.tilelang.chunk import ( + chunk_dplr_delta_rule_tilelang, + ) + try: + return chunk_dplr_delta_rule_tilelang( + 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, + cu_seqlens_cpu=cu_seqlens_cpu, + safe_gate=safe_gate, + lower_bound=lower_bound, + chunk_size=chunk_size, + disable_recompute=disable_recompute, + cp_context=cp_context, + ) + except Exception as exc: + # The verifier gates on the schedule arithmetic, but JIT/launch + # failures can still escape; honor the dispatch contract and fall + # back to the default Triton implementation. Only the forward call + # is guarded: once it succeeds, autograd is committed to the + # TileLang backward. + if cp_context is not None: + # a rank-local fallback would diverge the CP collectives + raise + key = f"{type(exc).__name__}: {exc}" + if key not in _FALLBACK_LOGGED: + _FALLBACK_LOGGED.add(key) + logger.warning( + f"[FLA Backend] TileLang DPLR forward failed ({key}); falling back to Triton" + ) + from fla.ops.generalized_delta_rule.dplr import chunk as dplr_chunk + fn = dplr_chunk.chunk_dplr_delta_rule + while hasattr(fn, "__wrapped__"): + fn = fn.__wrapped__ + return fn( + 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, + cu_seqlens_cpu=cu_seqlens_cpu, + safe_gate=safe_gate, + lower_bound=lower_bound, + chunk_size=chunk_size, + disable_recompute=disable_recompute, + cp_context=cp_context, + ) diff --git a/fla/ops/generalized_delta_rule/dplr/backends/tilelang/chunk.py b/fla/ops/generalized_delta_rule/dplr/backends/tilelang/chunk.py new file mode 100644 index 0000000000..6a46a0c2da --- /dev/null +++ b/fla/ops/generalized_delta_rule/dplr/backends/tilelang/chunk.py @@ -0,0 +1,1427 @@ +# Copyright (c) 2023-2026, Songlin Yang, Yu Zhang, Zhiyuan Li +# +# This source code is licensed under the MIT license found in the +# LICENSE file in the root directory of this source tree. +# For a list of all contributors, visit: +# https://github.com/fla-org/flash-linear-attention/graphs/contributors + +"""Low-retention DPLR path with FLA-compatible active-varlen semantics.""" + +from contextvars import ContextVar + +import torch +from torch import Tensor +from torch.utils._python_dispatch import TorchDispatchMode + +from fla.ops.cp.chunk_delta_h import ( + chunk_gated_delta_rule_bwd_dhu_pre_process, + chunk_gated_delta_rule_fwd_h_pre_process, +) +from fla.ops.generalized_delta_rule.dplr.chunk_o_bwd import chunk_dplr_bwd_dAu +from fla.ops.utils.constant import RCP_LN2 + +from .chunk_A_bwd import ( + chunk_dplr_bwd_dqk_intra_fused_qside_into, +) +from .chunk_A_fwd import ( + chunk_dplr_fwd_intra, + chunk_dplr_fwd_intra_from_gk, +) +from .chunk_h_fwd import chunk_dplr_fwd_h +from .chunk_ho_fwd import chunk_dplr_fwd_ho, chunk_dplr_fwd_ho_ctx +from .chunk_stream_bwd import chunk_dplr_bwd_stream_into +from .cumsum import chunk_local_cumsum +from .layout import ChunkLayout, build_varlen_chunk_layout +from .wy_fast_bwd import chunk_dplr_bwd_wy_repr_into +from .wy_fast_fwd import prepare_wy_repr_fwd + +_CHECKPOINT_PHASE_NORMAL = "normal" +_CHECKPOINT_PHASE_FORWARD = "forward" +_CHECKPOINT_PHASE_RECOMPUTE = "recompute" +_DPLR_CHECKPOINT_PHASE: ContextVar[str] = ContextVar( + "rwkv7_dplr_checkpoint_phase", + default=_CHECKPOINT_PHASE_NORMAL, +) + +# Active FLACPContext for the duration of a CP call. Custom ops cannot take +# the context object, so it is threaded through a ContextVar instead. +_DPLR_CP_CONTEXT: ContextVar = ContextVar("fla_dplr_cp_context", default=None) + + +class _DPLRCheckpointPhaseMode(TorchDispatchMode): + """Checkpoint phase marker that does not enter the dispatch stack.""" + + def __init__(self, phase: str): + super().__init__() + self.phase = phase + self._token = None + + def __enter__(self): + token = _DPLR_CHECKPOINT_PHASE.set(self.phase) + self._token = token + return self + + def __exit__(self, exc_type, exc_value, traceback): + token = self._token + self._token = None + if token is not None: + _DPLR_CHECKPOINT_PHASE.reset(token) + return False + + def __torch_dispatch__(self, func, types, args=(), kwargs=None): + return func(*args, **(kwargs or {})) + + +def dplr_checkpoint_context_fn() -> tuple[TorchDispatchMode, TorchDispatchMode]: + """Mark eager non-reentrant checkpoint forward and recompute phases.""" + return ( + _DPLRCheckpointPhaseMode(_CHECKPOINT_PHASE_FORWARD), + _DPLRCheckpointPhaseMode(_CHECKPOINT_PHASE_RECOMPUTE), + ) + + +def _checkpoint_phase_for_dispatch(is_compiling: bool) -> str: + if is_compiling: + return _CHECKPOINT_PHASE_NORMAL + return _DPLR_CHECKPOINT_PHASE.get() + + +def _prepare_cuda_cu_seqlens( + cu_seqlens: torch.Tensor, + device: torch.device, +) -> torch.Tensor: + if cu_seqlens.device.type != "cuda": + raise ValueError("DPLR varlen layout requires CUDA cu_seqlens") + if cu_seqlens.device != device: + raise ValueError( + f"cu_seqlens device {cu_seqlens.device} must match tensor device {device}" + ) + return cu_seqlens.to(dtype=torch.int32).contiguous() + + +def _layout_from_saved( + cu_seqlens: Tensor, + chunk_indices: Tensor, + chunk_offsets: Tensor, + is_varlen: bool, +) -> ChunkLayout | None: + if not is_varlen: + return None + return ChunkLayout( + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + chunk_offsets=chunk_offsets, + ) + + +def _empty_layout_metadata(q: Tensor) -> tuple[Tensor, Tensor]: + return ( + q.new_empty((0, 2), dtype=torch.int32), + q.new_empty((0,), dtype=torch.int32), + ) + + +def _chunk_dplr_delta_rule_bwd_core( + do: Tensor, + dht: Tensor, + q: Tensor, + k: Tensor, + v: Tensor, + a: Tensor, + b: Tensor, + gk: Tensor, + h0: Tensor, + cu_seqlens: Tensor, + chunk_indices: Tensor, + chunk_offsets: Tensor, + scale: float, + has_initial_state: bool, + is_varlen: bool, + chunk_size: int, + saved_h: Tensor | None = None, + saved_v_new: Tensor | None = None, + saved_wy: tuple[Tensor, ...] | None = None, +) -> tuple[Tensor, Tensor, Tensor, Tensor, Tensor, Tensor, Tensor]: + if saved_h is not None and saved_h.numel() == 0: + saved_h = None + if saved_v_new is not None and saved_v_new.numel() == 0: + saved_v_new = None + + cu = cu_seqlens if is_varlen else None + layout = _layout_from_saved(cu_seqlens, chunk_indices, chunk_offsets, is_varlen) + initial_state = h0 if has_initial_state else None + dht_arg = dht if dht.numel() > 0 else None + + cp_context = _DPLR_CP_CONTEXT.get() + bwd_derives_ge = q.shape[-1] == 64 + # Under disable_recompute the backward consumes the wy-side + # intermediates the forward saved, skipping the from_gk/intra A-stage, + # prepare_wy_repr, and wu entirely. This also makes the backward use + # the forward's own cumsum gi at K=64 instead of re-deriving it inside + # from_gk (rounding-level difference only). The CP branch keeps the + # recompute path: its boundary pre-process consumes u's values, which + # the save set deliberately omits. + if saved_wy is not None and cp_context is None: + qg, kg, ag, bg, w, A_ab_inv, A_ak, A_qk, A_qb = saved_wy + gi, ge = chunk_local_cumsum( + gk, + chunk_size, + scale=RCP_LN2, + cu_seqlens=cu, + chunk_layout=layout, + output_ge=not bwd_derives_ge, + ) + u = torch.empty_like(v) + elif bwd_derives_ge: + A_ab, A_qk, A_ak, A_qb, qg, kg, ag, bg, gi = chunk_dplr_fwd_intra_from_gk( + q=q, + k=k, + a=a, + b=b, + gk=gk, + scale=scale, + chunk_size=chunk_size, + cu_seqlens=cu, + chunk_layout=layout, + ) + ge = None + w, u, A_ab_inv = prepare_wy_repr_fwd( + ag=ag, + v=v, + A_ak=A_ak, + A_ab=A_ab, + cu_seqlens=cu, + chunk_size=chunk_size, + chunk_layout=layout, + ) + del A_ab + else: + gi, ge = chunk_local_cumsum( + gk, + chunk_size, + scale=RCP_LN2, + cu_seqlens=cu, + chunk_layout=layout, + ) + 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, + chunk_size=chunk_size, + cu_seqlens=cu, + chunk_layout=layout, + ) + w, u, A_ab_inv = prepare_wy_repr_fwd( + ag=ag, + v=v, + A_ak=A_ak, + A_ab=A_ab, + cu_seqlens=cu, + chunk_size=chunk_size, + chunk_layout=layout, + ) + del A_ab + if cp_context is not None: + # CP: rebuild the corrected initial state (boundary exchange across + # ranks) before recomputing the chunk states + 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, + initial_state=None, + context=cp_context, + chunk_size=chunk_size, + ) + if saved_h is not None and saved_v_new is not None: + h = saved_h + v_new = saved_v_new + else: + h, v_new, _ = chunk_dplr_fwd_h( + kg=kg, + v=v, + w=w, + u=u, + bg=bg, + gk=gi, + initial_state=initial_state, + output_final_state=False, + cu_seqlens=cu, + chunk_size=chunk_size, + chunk_layout=layout, + ) + + batch, tokens, heads, key_dim = q.shape + value_dim = v.shape[-1] + n_chunks = ( + chunk_indices.shape[0] + if is_varlen + else batch * ((tokens + chunk_size - 1) // chunk_size) + ) + n_seqs = cu_seqlens.shape[0] - 1 if is_varlen else batch + n_dh0 = n_seqs if has_initial_state else 1 + + dh0_workspace = torch.empty( + (n_dh0, heads, key_dim, value_dim), + dtype=torch.float32, + device=q.device, + ) + dgk_last = torch.empty( + (n_chunks, heads, key_dim), + dtype=torch.float32, + device=q.device, + ) + dv_full_workspace = torch.empty_like(v_new) + if cp_context is not None: + # CP: compute the local dh boundary contribution and fold the + # following ranks' dh into this rank's terminal dh + assert dht_arg is None, "When enable CP, the provided dht must be None." + dv_new_intra, _, _ = chunk_dplr_bwd_dAu( + v=v, + v_new=v_new, + do=do, + A_qb=A_qb, + scale=scale, + cu_seqlens=cu, + chunk_size=chunk_size, + chunk_indices=layout.chunk_indices if is_varlen else None, + ) + dht_arg, _ = 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, + dht=None, + initial_state=None, + context=cp_context, + chunk_size=chunk_size, + ) + # The recompute path recycles the forward intermediates as gradient + # buffers; the saved path leaves them read-only (they alias ctx-saved + # tensors and the custom-op contract forbids mutating inputs), so the + # gradients land in fresh buffers instead. + recycle = saved_wy is None or cp_context is not None + dq_buf = qg if recycle else torch.empty_like(qg) + dk_buf = kg if recycle else torch.empty_like(kg) + dw_buf = w if recycle else torch.empty_like(w) + db_buf = bg if recycle else torch.empty_like(bg) + dag_buf = ag if recycle else torch.empty_like(ag) + dqg, dkg, dw, dbg, dgk_last, dv2, dv_full, dh0 = ( + chunk_dplr_bwd_stream_into( + qg=qg, + bg=bg, + w=w, + kg=kg, + v=v, + v_new=v_new, + gk=gi, + h=h, + h0=initial_state, + dht=dht_arg, + do=do, + A_qb_for_dv=A_qb, + A_qk=A_qk, + dq_out=dq_buf, + dk_out=dk_buf, + dw_out=dw_buf, + db_out=db_buf, + dgk_last_out=dgk_last, + dv2_out=u, + dv_full_out=dv_full_workspace, + dh0_out=dh0_workspace, + cu_seqlens=cu, + chunk_size=chunk_size, + chunk_layout=layout, + ) + ) + del A_qb + del A_qk + + dA_ab_workspace = torch.empty(A_ak.shape, dtype=q.dtype, device=A_ak.device) + dA_ak_workspace = torch.empty(A_ak.shape, dtype=q.dtype, device=A_ak.device) + dA_ab, dA_ak, dv_out, dag = chunk_dplr_bwd_wy_repr_into( + A_ab_inv=A_ab_inv, + A_ak=A_ak, + v=v, + ag=ag, + dw=dw, + du=dv2, + dv0=dv_full, + dA_ab_out=dA_ab_workspace, + dA_ak_out=dA_ak_workspace, + dv_out=dv_full, + dag_out=dag_buf, + cu_seqlens=cu, + chunk_size=chunk_size, + chunk_layout=layout, + ) + del h + del dv2 + + # dgk is returned in gk's dtype; recycle the dw buffer only when they agree + dgk_out = dw if (recycle and dw.dtype == gk.dtype) else torch.empty_like(gk) + dq, dk, da, db, dgk = chunk_dplr_bwd_dqk_intra_fused_qside_into( + q=q, + k=k, + a=a, + b=b, + gi=gi, + ge=ge, + gk=gk if bwd_derives_ge else None, + do=do, + v=v, + v_new=v_new, + dAak=dA_ak, + dAab=dA_ab, + dqg=dqg, + dkg=dkg, + dag=dag, + dbg=dbg, + dgk_last=dgk_last, + dq_out=dqg, + dk_out=dkg, + da_out=dag, + db_out=dbg, + dgk_out=dgk_out, + cu_seqlens=cu, + chunk_size=chunk_size, + scale=scale, + chunk_layout=layout, + dgk_dtype=gk.dtype, + ) + + dh0_out = dh0 if (cp_context is None and has_initial_state and dh0 is not None) else h0.new_empty((0,)) + return dq, dk, dv_out, da, db, dgk, dh0_out + + +@torch.library.custom_op( + "fla::chunk_dplr_delta_rule_bwd", + mutates_args=(), +) +def _chunk_dplr_delta_rule_bwd_op( + do: Tensor, + dht: Tensor, + q: Tensor, + k: Tensor, + v: Tensor, + a: Tensor, + b: Tensor, + gk: Tensor, + h0: Tensor, + cu_seqlens: Tensor, + chunk_indices: Tensor, + chunk_offsets: Tensor, + scale: float, + has_initial_state: bool, + is_varlen: bool, + chunk_size: int, +) -> tuple[Tensor, Tensor, Tensor, Tensor, Tensor, Tensor, Tensor]: + return _chunk_dplr_delta_rule_bwd_core( + do, + dht, + q, + k, + v, + a, + b, + gk, + h0, + cu_seqlens, + chunk_indices, + chunk_offsets, + scale, + has_initial_state, + is_varlen, + chunk_size, + ) + + +@_chunk_dplr_delta_rule_bwd_op.register_fake +def _chunk_dplr_delta_rule_bwd_fake( + do, + dht, + q, + k, + v, + a, + b, + gk, + h0, + cu_seqlens, + chunk_indices, + chunk_offsets, + scale: float, + has_initial_state: bool, + is_varlen: bool, + chunk_size: int, +): + dh0 = h0.new_empty(h0.shape) if has_initial_state else h0.new_empty((0,)) + return ( + q.new_empty(q.shape), + k.new_empty(k.shape), + v.new_empty(v.shape), + a.new_empty(a.shape), + b.new_empty(b.shape), + gk.new_empty(gk.shape), + dh0, + ) + + +@torch.library.custom_op( + "fla::chunk_dplr_delta_rule_bwd_ctx", + mutates_args=(), +) +def _chunk_dplr_delta_rule_bwd_ctx_op( + do: Tensor, + dht: Tensor, + q: Tensor, + k: Tensor, + v: Tensor, + a: Tensor, + b: Tensor, + gk: Tensor, + h0: Tensor, + cu_seqlens: Tensor, + chunk_indices: Tensor, + chunk_offsets: Tensor, + h_ctx: Tensor, + v_new_ctx: Tensor, + qg_ctx: Tensor, + kg_ctx: Tensor, + ag_ctx: Tensor, + bg_ctx: Tensor, + w_ctx: Tensor, + A_ab_inv_ctx: Tensor, + A_ak_ctx: Tensor, + A_qk_ctx: Tensor, + A_qb_ctx: Tensor, + scale: float, + has_initial_state: bool, + is_varlen: bool, + chunk_size: int, +) -> tuple[Tensor, Tensor, Tensor, Tensor, Tensor, Tensor, Tensor]: + return _chunk_dplr_delta_rule_bwd_core( + do, + dht, + q, + k, + v, + a, + b, + gk, + h0, + cu_seqlens, + chunk_indices, + chunk_offsets, + scale, + has_initial_state, + is_varlen, + chunk_size, + saved_h=h_ctx, + saved_v_new=v_new_ctx, + saved_wy=( + qg_ctx, kg_ctx, ag_ctx, bg_ctx, + w_ctx, A_ab_inv_ctx, A_ak_ctx, A_qk_ctx, A_qb_ctx, + ), + ) + + +@_chunk_dplr_delta_rule_bwd_ctx_op.register_fake +def _chunk_dplr_delta_rule_bwd_ctx_fake( + do, + dht, + q, + k, + v, + a, + b, + gk, + h0, + cu_seqlens, + chunk_indices, + chunk_offsets, + h_ctx, + v_new_ctx, + qg_ctx, + kg_ctx, + ag_ctx, + bg_ctx, + w_ctx, + A_ab_inv_ctx, + A_ak_ctx, + A_qk_ctx, + A_qb_ctx, + scale: float, + has_initial_state: bool, + is_varlen: bool, + chunk_size: int, +): + del h_ctx, v_new_ctx, qg_ctx, kg_ctx, ag_ctx, bg_ctx, w_ctx + del A_ab_inv_ctx, A_ak_ctx, A_qk_ctx, A_qb_ctx + return _chunk_dplr_delta_rule_bwd_fake( + do, + dht, + q, + k, + v, + a, + b, + gk, + h0, + cu_seqlens, + chunk_indices, + chunk_offsets, + scale, + has_initial_state, + is_varlen, + chunk_size, + ) + + +@torch.library.custom_op( + "fla::chunk_dplr_delta_rule_fwd", + mutates_args=(), +) +def _chunk_dplr_delta_rule_fwd_op( + q: Tensor, + k: Tensor, + v: Tensor, + a: Tensor, + b: Tensor, + gk: Tensor, + h0: Tensor, + cu_seqlens: Tensor, + scale: float, + has_initial_state: bool, + output_final_state: bool, + is_varlen: bool, + chunk_size: int, +) -> tuple[Tensor, Tensor, Tensor, Tensor]: + cu = cu_seqlens if is_varlen else None + layout = ( + build_varlen_chunk_layout( + cu_seqlens, + chunk_size, + q.shape[0] * q.shape[1], + ) + if is_varlen + else None + ) + initial_state = h0 if has_initial_state else None + + # from_gk (in-CTA chunk-local cumsum from bf16 gk) is the default A-stage + # for head_dim 64, in training and eval, rectangular and varlen. It + # removes the standalone fp32 gi/ge cumsum kernels and the fp32 ge tensor. + use_from_gk_a = q.shape[-1] == 64 + if use_from_gk_a: + A_ab, A_qk, A_ak, A_qb, qg, kg, ag, bg, gi = chunk_dplr_fwd_intra_from_gk( + q=q, + k=k, + a=a, + b=b, + gk=gk, + scale=scale, + chunk_size=chunk_size, + cu_seqlens=cu, + chunk_layout=layout, + ) + else: + gi, ge = chunk_local_cumsum( + gk, + chunk_size, + scale=RCP_LN2, + cu_seqlens=cu, + chunk_layout=layout, + ) + 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, + chunk_size=chunk_size, + cu_seqlens=cu, + chunk_layout=layout, + ) + w, u, _ = prepare_wy_repr_fwd( + ag=ag, + v=v, + A_ak=A_ak, + A_ab=A_ab, + cu_seqlens=cu, + chunk_size=chunk_size, + chunk_layout=layout, + ) + cp_context = _DPLR_CP_CONTEXT.get() + if cp_context is not None: + # CP: exchange the compact local boundary state (h = M @ h_in + c) + # across ranks and fold it into the corrected initial state + 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, + initial_state=None, + context=cp_context, + chunk_size=chunk_size, + ) + o, final_state = chunk_dplr_fwd_ho( + qg=qg, + kg=kg, + v=v, + w=w, + u=u, + bg=bg, + gk=gi, + A_qk=A_qk, + A_qb=A_qb, + initial_state=initial_state, + output_final_state=output_final_state, + cu_seqlens=cu, + chunk_size=chunk_size, + chunk_layout=layout, + ) + + if final_state is None: + final_state = q.new_empty((0,), dtype=torch.float32) + if layout is None: + chunk_indices, chunk_offsets = _empty_layout_metadata(q) + else: + chunk_indices = layout.chunk_indices + chunk_offsets = layout.chunk_offsets + return o, final_state, chunk_indices, chunk_offsets + + +@_chunk_dplr_delta_rule_fwd_op.register_fake +def _chunk_dplr_delta_rule_fwd_fake( + q, + k, + v, + a, + b, + gk, + h0, + cu_seqlens, + scale: float, + has_initial_state: bool, + output_final_state: bool, + is_varlen: bool, + chunk_size: int, +): + batch, tokens, heads, key_dim = q.shape + value_dim = v.shape[-1] + if is_varlen: + n_seqs = cu_seqlens.shape[0] - 1 + n_chunks = ( + (batch * tokens + chunk_size - 1) // chunk_size + + cu_seqlens.shape[0] + - 2 + ) + indices_shape = (n_chunks, 2) + offsets_shape = (cu_seqlens.shape[0],) + else: + n_seqs = batch + indices_shape = (0, 2) + offsets_shape = (0,) + final_shape = ( + (n_seqs, heads, key_dim, value_dim) + if output_final_state + else (0,) + ) + return ( + v.new_empty(v.shape), + q.new_empty(final_shape, dtype=torch.float32), + q.new_empty(indices_shape, dtype=torch.int32), + q.new_empty(offsets_shape, dtype=torch.int32), + ) + + +def _chunk_dplr_delta_rule_fwd_ctx_core( + q: Tensor, + k: Tensor, + v: Tensor, + a: Tensor, + b: Tensor, + gk: Tensor, + h0: Tensor, + cu_seqlens: Tensor, + scale: float, + has_initial_state: bool, + output_final_state: bool, + is_varlen: bool, + chunk_size: int, + store_context: bool, +) -> tuple[Tensor, ...]: + cu = cu_seqlens if is_varlen else None + layout = ( + build_varlen_chunk_layout( + cu_seqlens, + chunk_size, + q.shape[0] * q.shape[1], + ) + if is_varlen + else None + ) + initial_state = h0 if has_initial_state else None + + # Same A-stage split as the plain forward: from_gk (in-CTA cumsum) at + # head_dim 64 avoids the standalone cumsum and the intra variant's + # torch-side gate preparation entirely. + if q.shape[-1] == 64: + A_ab, A_qk, A_ak, A_qb, qg, kg, ag, bg, gi = chunk_dplr_fwd_intra_from_gk( + q=q, + k=k, + a=a, + b=b, + gk=gk, + scale=scale, + chunk_size=chunk_size, + cu_seqlens=cu, + chunk_layout=layout, + ) + else: + gi, ge = chunk_local_cumsum( + gk, + chunk_size, + scale=RCP_LN2, + cu_seqlens=cu, + chunk_layout=layout, + ) + 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, + chunk_size=chunk_size, + cu_seqlens=cu, + chunk_layout=layout, + ) + w, u, A_ab_inv = prepare_wy_repr_fwd( + ag=ag, + v=v, + A_ak=A_ak, + A_ab=A_ab, + cu_seqlens=cu, + chunk_size=chunk_size, + chunk_layout=layout, + ) + cp_context = _DPLR_CP_CONTEXT.get() + if cp_context is not None: + # CP: exchange the compact local boundary state (h = M @ h_in + c) + # across ranks and fold it into the corrected initial state + 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, + initial_state=None, + context=cp_context, + chunk_size=chunk_size, + ) + o, final_state, h_ctx, v_new_ctx = chunk_dplr_fwd_ho_ctx( + qg=qg, + kg=kg, + v=v, + w=w, + u=u, + bg=bg, + gk=gi, + A_qk=A_qk, + A_qb=A_qb, + initial_state=initial_state, + output_final_state=output_final_state, + cu_seqlens=cu, + chunk_size=chunk_size, + chunk_layout=layout, + store_context=store_context, + ) + + if final_state is None: + final_state = q.new_empty((0,), dtype=torch.float32) + if layout is None: + chunk_indices, chunk_offsets = _empty_layout_metadata(q) + else: + chunk_indices = layout.chunk_indices + chunk_offsets = layout.chunk_offsets + # The wy-side intermediates ride along as outputs so the disable_recompute + # backward can consume them from ctx instead of recomputing (they are + # forward intermediates anyway; returning them costs no extra traffic). + return ( + o, final_state, chunk_indices, chunk_offsets, h_ctx, v_new_ctx, + qg, kg, ag, bg, w, A_ab_inv, A_ak, A_qk, A_qb, + ) + + +@torch.library.custom_op( + "fla::chunk_dplr_delta_rule_fwd_ctx", + mutates_args=(), +) +def _chunk_dplr_delta_rule_fwd_ctx_op( + q: Tensor, + k: Tensor, + v: Tensor, + a: Tensor, + b: Tensor, + gk: Tensor, + h0: Tensor, + cu_seqlens: Tensor, + scale: float, + has_initial_state: bool, + output_final_state: bool, + is_varlen: bool, + chunk_size: int, +) -> tuple[ + Tensor, Tensor, Tensor, Tensor, Tensor, Tensor, + Tensor, Tensor, Tensor, Tensor, Tensor, Tensor, Tensor, Tensor, Tensor, +]: + return _chunk_dplr_delta_rule_fwd_ctx_core( + q, + k, + v, + a, + b, + gk, + h0, + cu_seqlens, + scale, + has_initial_state, + output_final_state, + is_varlen, + chunk_size, + store_context=True, + ) + + +@_chunk_dplr_delta_rule_fwd_ctx_op.register_fake +def _chunk_dplr_delta_rule_fwd_ctx_fake( + q, + k, + v, + a, + b, + gk, + h0, + cu_seqlens, + scale: float, + has_initial_state: bool, + output_final_state: bool, + is_varlen: bool, + chunk_size: int, +): + del k, a, b, gk, scale, has_initial_state + batch, tokens, heads, key_dim = q.shape + value_dim = v.shape[-1] + if is_varlen: + n_seqs = cu_seqlens.shape[0] - 1 + n_chunks = ( + (batch * tokens + chunk_size - 1) // chunk_size + + cu_seqlens.shape[0] + - 2 + ) + indices_shape = (n_chunks, 2) + offsets_shape = (cu_seqlens.shape[0],) + else: + n_seqs = batch + n_chunks = batch * ((tokens + chunk_size - 1) // chunk_size) + indices_shape = (0, 2) + offsets_shape = (0,) + final_shape = ( + (n_seqs, heads, key_dim, value_dim) + if output_final_state + else (0,) + ) + return ( + v.new_empty(v.shape), + q.new_empty(final_shape, dtype=torch.float32), + q.new_empty(indices_shape, dtype=torch.int32), + q.new_empty(offsets_shape, dtype=torch.int32), + q.new_empty((n_chunks, heads, key_dim, value_dim)), + v.new_empty(v.shape), + # wy-side save set for the disable_recompute backward + q.new_empty(q.shape), + q.new_empty(q.shape), + q.new_empty(q.shape), + q.new_empty(q.shape), + q.new_empty(q.shape), + q.new_empty((batch, tokens, heads, chunk_size), dtype=torch.float32), + q.new_empty((batch, tokens, heads, chunk_size), dtype=torch.float16), + q.new_empty((batch, tokens, heads, chunk_size)), + q.new_empty((batch, tokens, heads, chunk_size)), + ) + + +@torch.library.custom_op( + "fla::chunk_dplr_delta_rule_fwd_ctx_elided", + mutates_args=(), +) +def _chunk_dplr_delta_rule_fwd_ctx_elided_op( + q: Tensor, + k: Tensor, + v: Tensor, + a: Tensor, + b: Tensor, + gk: Tensor, + h0: Tensor, + cu_seqlens: Tensor, + scale: float, + has_initial_state: bool, + output_final_state: bool, + is_varlen: bool, + chunk_size: int, +) -> tuple[ + Tensor, Tensor, Tensor, Tensor, Tensor, Tensor, + Tensor, Tensor, Tensor, Tensor, Tensor, Tensor, Tensor, Tensor, Tensor, +]: + return _chunk_dplr_delta_rule_fwd_ctx_core( + q, + k, + v, + a, + b, + gk, + h0, + cu_seqlens, + scale, + has_initial_state, + output_final_state, + is_varlen, + chunk_size, + store_context=False, + ) + + +@_chunk_dplr_delta_rule_fwd_ctx_elided_op.register_fake +def _chunk_dplr_delta_rule_fwd_ctx_elided_fake( + q, + k, + v, + a, + b, + gk, + h0, + cu_seqlens, + scale: float, + has_initial_state: bool, + output_final_state: bool, + is_varlen: bool, + chunk_size: int, +): + outputs = _chunk_dplr_delta_rule_fwd_ctx_fake( + q, + k, + v, + a, + b, + gk, + h0, + cu_seqlens, + scale, + has_initial_state, + output_final_state, + is_varlen, + chunk_size, + ) + return ( + *outputs[:4], + q.new_empty((1,)).expand(outputs[4].shape), + v.new_empty((1,)).expand(outputs[5].shape), + *outputs[6:], + ) + + +def _chunk_dplr_setup_context(ctx, inputs, output): + ( + q, + k, + v, + a, + b, + gk, + h0, + cu_seqlens, + scale, + has_initial_state, + output_final_state, + is_varlen, + chunk_size, + ) = inputs + _, final_state, chunk_indices, chunk_offsets = output + if not output_final_state: + ctx.mark_non_differentiable(final_state) + ctx.mark_non_differentiable(chunk_indices, chunk_offsets) + # the backward consumes grads of o/final_state only and already handles + # them being None, so skip the engine's zero materialization of the rest + ctx.set_materialize_grads(False) + ctx.save_for_backward( + q, + k, + v, + a, + b, + gk, + h0, + cu_seqlens, + chunk_indices, + chunk_offsets, + ) + ctx.scale = float(scale) + ctx.has_initial_state = bool(has_initial_state) + ctx.is_varlen = bool(is_varlen) + ctx.chunk_size = int(chunk_size) + ctx.cp_context = _DPLR_CP_CONTEXT.get() + + +def _chunk_dplr_backward( + ctx, + do, + dht, + _dchunk_indices, + _dchunk_offsets, +): + ( + q, + k, + v, + a, + b, + gk, + h0, + cu_seqlens, + chunk_indices, + chunk_offsets, + ) = ctx.saved_tensors + dht_arg = dht if dht is not None else q.new_empty((0,), dtype=torch.float32) + token = _DPLR_CP_CONTEXT.set(getattr(ctx, "cp_context", None)) + try: + dq, dk, dv, da, db, dgk, dh0 = _chunk_dplr_delta_rule_bwd_op( + do, + dht_arg, + q, + k, + v, + a, + b, + gk, + h0, + cu_seqlens, + chunk_indices, + chunk_offsets, + ctx.scale, + ctx.has_initial_state, + ctx.is_varlen, + ctx.chunk_size, + ) + finally: + _DPLR_CP_CONTEXT.reset(token) + return ( + dq, + dk, + dv, + da, + db, + dgk, + dh0 if ctx.has_initial_state else None, + None, + None, + None, + None, + None, + None, + ) + + +_chunk_dplr_delta_rule_fwd_op.register_autograd( + _chunk_dplr_backward, + setup_context=_chunk_dplr_setup_context, +) + + +def _chunk_dplr_ctx_setup_context(ctx, inputs, output): + ( + q, + k, + v, + a, + b, + gk, + h0, + cu_seqlens, + scale, + has_initial_state, + output_final_state, + is_varlen, + chunk_size, + ) = inputs + _, final_state, chunk_indices, chunk_offsets, h_ctx, v_new_ctx = output[:6] + wy_ctx = output[6:] + if not output_final_state: + ctx.mark_non_differentiable(final_state) + ctx.mark_non_differentiable(chunk_indices, chunk_offsets, *wy_ctx) + # 13 of the 15 outputs are pure save-for-backward intermediates whose + # grads the backward discards; materializing them as zeros would cost + # ~3GB of dead fills per step at h4096 + ctx.set_materialize_grads(False) + ctx.save_for_backward( + q, + k, + v, + a, + b, + gk, + h0, + cu_seqlens, + chunk_indices, + chunk_offsets, + h_ctx, + v_new_ctx, + *wy_ctx, + ) + ctx.scale = float(scale) + ctx.has_initial_state = bool(has_initial_state) + ctx.is_varlen = bool(is_varlen) + ctx.chunk_size = int(chunk_size) + ctx.cp_context = _DPLR_CP_CONTEXT.get() + + +def _chunk_dplr_ctx_backward_from_saved( + ctx, + saved_tensors, + do, + dht, +): + ( + q, + k, + v, + a, + b, + gk, + h0, + cu_seqlens, + chunk_indices, + chunk_offsets, + h_ctx, + v_new_ctx, + *wy_ctx + ) = saved_tensors + dht_arg = dht if dht is not None else q.new_empty((0,), dtype=torch.float32) + token = _DPLR_CP_CONTEXT.set(getattr(ctx, "cp_context", None)) + try: + dq, dk, dv, da, db, dgk, dh0 = _chunk_dplr_delta_rule_bwd_ctx_op( + do, + dht_arg, + q, + k, + v, + a, + b, + gk, + h0, + cu_seqlens, + chunk_indices, + chunk_offsets, + h_ctx, + v_new_ctx, + *wy_ctx, + ctx.scale, + ctx.has_initial_state, + ctx.is_varlen, + ctx.chunk_size, + ) + finally: + _DPLR_CP_CONTEXT.reset(token) + return ( + dq, + dk, + dv, + da, + db, + dgk, + dh0 if ctx.has_initial_state else None, + None, + None, + None, + None, + None, + None, + ) + + +def _chunk_dplr_ctx_backward( + ctx, + do, + dht, + _dchunk_indices, + _dchunk_offsets, + _dh_ctx, + _dv_new_ctx, + *_dwy_ctx, +): + return _chunk_dplr_ctx_backward_from_saved( + ctx, + ctx.saved_tensors, + do, + dht, + ) + + +def _checkpoint_context_is_materialized(tensor: Tensor) -> bool: + required_bytes = tensor.numel() * tensor.element_size() + return tensor.untyped_storage().nbytes() >= required_bytes + + +def _chunk_dplr_ctx_elided_backward( + ctx, + do, + dht, + _dchunk_indices, + _dchunk_offsets, + _dh_ctx, + _dv_new_ctx, + *_dwy_ctx, +): + saved_tensors = ctx.saved_tensors + h_ctx, v_new_ctx = saved_tensors[10:12] + if not ( + _checkpoint_context_is_materialized(h_ctx) + and _checkpoint_context_is_materialized(v_new_ctx) + ): + raise RuntimeError( + "DPLR checkpoint context-write elision requires non-reentrant " + "checkpoint recomputation before backward" + ) + return _chunk_dplr_ctx_backward_from_saved( + ctx, + saved_tensors, + do, + dht, + ) + + +_chunk_dplr_delta_rule_fwd_ctx_op.register_autograd( + _chunk_dplr_ctx_backward, + setup_context=_chunk_dplr_ctx_setup_context, +) + +_chunk_dplr_delta_rule_fwd_ctx_elided_op.register_autograd( + _chunk_dplr_ctx_elided_backward, + setup_context=_chunk_dplr_ctx_setup_context, +) + + +def chunk_dplr_delta_rule_tilelang( + 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.Tensor | None = None, + cu_seqlens_cpu: torch.Tensor | None = None, + safe_gate: bool = False, + lower_bound: float | None = None, + chunk_size: int | None = None, + disable_recompute: bool = False, + cp_context=None, + **kwargs, +) -> tuple[torch.Tensor, torch.Tensor | None]: + del cu_seqlens_cpu, safe_gate, lower_bound + if "head_first" in kwargs: + raise DeprecationWarning( + "head_first has been removed; inputs must use [B, T, H, ...]" + ) + if kwargs: + raise TypeError(f"unexpected DPLR kwargs: {', '.join(sorted(kwargs))}") + 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 + chunk_size = 16 if chunk_size is None else int(chunk_size) + scale_f = float(q.shape[-1] ** -0.5 if scale is None else scale) + n_seqs = len(cu_seqlens) - 1 if cu_seqlens is not None else q.shape[0] + if initial_state is not None: + # the kernels read h0 with dense strides; the Triton path gets this + # normalization from input_guard, which this backend bypasses + h0 = initial_state.contiguous() + elif cp_context is not None: + # the corrected initial state is produced inside the op by the CP + # boundary exchange; this buffer only carries its shape/ABI + h0 = q.new_empty((n_seqs, q.shape[2], q.shape[3], v.shape[3]), dtype=torch.float32) + else: + h0 = q.new_empty((0,), dtype=torch.float32) + cu = ( + _prepare_cuda_cu_seqlens(cu_seqlens, q.device) + if cu_seqlens is not None + else q.new_empty((0,), dtype=torch.int32) + ) + is_compiling = torch.compiler.is_compiling() + checkpoint_phase = _checkpoint_phase_for_dispatch(is_compiling) + op_args = ( + q, + k, + v, + a, + b, + gk, + h0, + cu, + scale_f, + initial_state is not None or cp_context is not None, + output_final_state, + cu_seqlens is not None, + chunk_size, + ) + + def call_fwd_op(): + if disable_recompute: + elide_ctx = ( + checkpoint_phase == _CHECKPOINT_PHASE_FORWARD + or ( + checkpoint_phase == _CHECKPOINT_PHASE_NORMAL + and not torch.is_grad_enabled() + ) + ) + if not is_compiling and elide_ctx: + return _chunk_dplr_delta_rule_fwd_ctx_elided_op(*op_args)[:2] + return _chunk_dplr_delta_rule_fwd_ctx_op(*op_args)[:2] + return _chunk_dplr_delta_rule_fwd_op(*op_args)[:2] + + if cp_context is None: + # keep the plain path free of ContextVar ops, which Dynamo cannot + # trace (torch.compile fullgraph support); the ContextVar is only a + # side channel for CP + o, final_state = call_fwd_op() + else: + token = _DPLR_CP_CONTEXT.set(cp_context) + try: + o, final_state = call_fwd_op() + finally: + _DPLR_CP_CONTEXT.reset(token) + return o, final_state if output_final_state else None diff --git a/fla/ops/generalized_delta_rule/dplr/backends/tilelang/chunk_A_bwd.py b/fla/ops/generalized_delta_rule/dplr/backends/tilelang/chunk_A_bwd.py new file mode 100644 index 0000000000..e7e0ba6de7 --- /dev/null +++ b/fla/ops/generalized_delta_rule/dplr/backends/tilelang/chunk_A_bwd.py @@ -0,0 +1,529 @@ +# Copyright (c) 2023-2026, Songlin Yang, Yu Zhang, Zhiyuan Li +# +# This source code is licensed under the MIT license found in the +# LICENSE file in the root directory of this source tree. +# For a list of all contributors, visit: +# https://github.com/fla-org/flash-linear-attention/graphs/contributors + +"""DPLR intra-chunk backward (dq, dk, da, db, dgk). + +Inputs: + q, k, a, b: (B, T, H, K) — raw activations + gi, ge: (B, T, H, K) fp32 — cumsum gates + dAqk, dAqb: (B, T, H, BT) in_dtype — gradients of the 2 q-side A-matrices + dAak, dAab: (B, T, H, BT) fp32 — gradients of the 2 a-side A-matrices + dqg, dkg, dag, dbg: (B, T, H, K) — gradients of pre-gated outputs from chunk_A_fwd + dgk_last: (chunk_rows, H, K) fp32 — from chunk_dplr_bwd_o + +Outputs: + dq, dk, da, db: (B, T, H, K) + dgk: (B, T, H, K) — cumulative final gate gradient, cast to original gk dtype + +Math: a per-row, per-j loop combining contributions from each j-th token +in the chunk into the 4 gradient accumulators, then a final mixing with +the chunk-pregated gradients (dqg/dkg/dag/dbg), then a reverse-cumsum +to compose dgk_last back into the per-token dgk. +""" + +import tilelang +import tilelang.language as T +import torch + +from fla.ops.utils.constant import RCP_LN2 + +from .layout import ChunkLayout, build_rect_chunk_layout, build_varlen_chunk_layout +from .schedules import device_cc + + +def _a_bwd_config(K: int, BT: int, in_dtype: str, device: torch.device) -> dict[str, int]: + # BK=64 with 256 threads on cc90 (fits 228KB smem, 1 CTA/SM there so the + # extra warps hide unpipelined load latency); BK=32 elsewhere (cc120's + # 99KB cap). 256 threads also pay at BT=64 on cc120 (bwd_intra -14%, + # but +16% at BT=32 where occupancy already hides latency). Bulk + # vectorized copies only pay at BT=64 on cc90 (measured flat to -4% at + # BT<=32; at BT=64 on cc120 the bulk staging exceeds the 99KB smem cap). + cc = device_cc(device) + return { + "BK": 64 if cc == 90 else 32, + "threads": 256 if (cc == 90 or BT >= 64) else 128, + "bulk_copy": cc == 90 and BT >= 64, + } + + +@tilelang.jit( + pass_configs={tilelang.PassConfigKey.TL_ENABLE_FAST_MATH: True, + tilelang.PassConfigKey.TL_DISABLE_DATA_RACE_CHECK: False, + }, +) +def _chunk_dplr_bwd_kernel_intra( + H, K, BT, in_dtype, out_dtype, + scale_value: float, + FUSE_QSIDE_DA: bool = False, + V: int = 64, + BV: int = 32, + BK: int = 32, + threads: int = 128, + DERIVE_GE: bool = False, + bulk_copy: bool = False, + cumsum_scale_value: float = RCP_LN2, + gk_dtype: str | None = None, +): + acc_dtype = "float32" + # DERIVE_GE loads raw gk, which may stay in its own dtype (e.g. fp32) + ge_in_dtype = (gk_dtype or in_dtype) if DERIVE_GE else acc_dtype + n_tokens, n_seq_plus_one, n_chunks, n_tokens_d = T.dynamic( + "n_tokens, n_seq_plus_one, n_chunks, n_tokens_d" + ) + + @T.prim_func + def chunk_dplr_bwd_intra_tl( + q: T.Tensor((n_tokens, H, K), in_dtype), + k: T.Tensor((n_tokens, H, K), in_dtype), + a: T.Tensor((n_tokens, H, K), in_dtype), + b: T.Tensor((n_tokens, H, K), in_dtype), + gi: T.Tensor((n_tokens, H, K), acc_dtype), + ge: T.Tensor((n_tokens, H, K), ge_in_dtype), + # FUSE_QSIDE_DA never reads dAqk/dAqb globals (the q-side dA tiles are + # recomputed in-CTA), so their leading extent gets its own symbol and + # callers may pass size-1 dummies instead of full [n_tokens, H, BT] + # workspaces. + dAqk: T.Tensor((n_tokens_d, H, BT), in_dtype), + dAqb: T.Tensor((n_tokens_d, H, BT), in_dtype), + dAak: T.Tensor((n_tokens, H, BT), in_dtype), + dAab: T.Tensor((n_tokens, H, BT), in_dtype), + do: T.Tensor((n_tokens, H, V), in_dtype), + v: T.Tensor((n_tokens, H, V), in_dtype), + v_new: T.Tensor((n_tokens, H, V), in_dtype), + dqg: T.Tensor((n_tokens, H, K), in_dtype), + dkg: T.Tensor((n_tokens, H, K), in_dtype), + dag: T.Tensor((n_tokens, H, K), in_dtype), + dbg: T.Tensor((n_tokens, H, K), in_dtype), + cu_seqlens: T.Tensor((n_seq_plus_one,), "int32"), + chunk_indices: T.Tensor((n_chunks, 2), "int32"), + dgk_last: T.Tensor((n_chunks, H, K), acc_dtype), + dq: T.Tensor((n_tokens, H, K), in_dtype), + dk: T.Tensor((n_tokens, H, K), in_dtype), + da: T.Tensor((n_tokens, H, K), in_dtype), + db: T.Tensor((n_tokens, H, K), in_dtype), + dgk_output: T.Tensor((n_tokens, H, K), out_dtype), + ): + with T.Kernel(T.ceildiv(K, BK), n_chunks, H, threads=threads) as (i_k, i_c, i_h): + i_n = chunk_indices[i_c, 0] + i_t = chunk_indices[i_c, 1] + safe_i_n = T.max(i_n, 0) + seq_bos = cu_seqlens[safe_i_n] + seq_eos = cu_seqlens[safe_i_n + 1] + bos_raw = seq_bos + i_t * BT + eos_raw = T.min(bos_raw + BT, seq_eos) + is_valid_row = i_n >= 0 + bos = T.if_then_else(is_valid_row, bos_raw, T.int32(0)) + eos = T.if_then_else(is_valid_row, eos_raw, T.int32(0)) + last_idx = T.max(eos - 1, 0) + is_valid_chunk = is_valid_row and bos < eos + + q_shared = T.alloc_shared((BT, BK), in_dtype) + k_shared = T.alloc_shared((BT, BK), in_dtype) + a_shared = T.alloc_shared((BT, BK), in_dtype) + b_shared = T.alloc_shared((BT, BK), in_dtype) + gi_shared = T.alloc_shared((BT, BK), acc_dtype) + ge_shared = T.alloc_shared((BT, BK), acc_dtype) + + # Two staging tiles serve both dA pairs: dAqk/dAqb for the q-side + # GEMM group, then overwritten with dAak/dAab for the a-side + # group. Per-accumulator GEMM order is unchanged (bit-exact with + # the four-tile layout), and the footprint drops by 2*(BT,BT) + # fp32 tiles — what makes the fused A-backward fit a 99KB cap at + # BT=64. + dA1_shared = T.alloc_shared((BT, BT), acc_dtype) + dA2_shared = T.alloc_shared((BT, BT), acc_dtype) + if FUSE_QSIDE_DA: + do_shared = T.alloc_shared((BT, BV), in_dtype) + v_shared = T.alloc_shared((BT, BV), in_dtype) + dAqk_frag = T.alloc_fragment((BT, BT), acc_dtype) + v_new_shared = T.alloc_shared((BT, BV), in_dtype) + dAqb_frag = T.alloc_fragment((BT, BT), acc_dtype) + + dqg_shared = T.alloc_shared((BT, BK), in_dtype) + dkg_shared = T.alloc_shared((BT, BK), in_dtype) + dag_shared = T.alloc_shared((BT, BK), in_dtype) + dbg_shared = T.alloc_shared((BT, BK), in_dtype) + g_last = T.alloc_shared((BK,), acc_dtype) + + # Load Q, K, A, B, gi, ge tiles. On cc90 (1 CTA/SM at BT=64) the + # kernel is latency-bound, so interior chunks take bulk vectorized + # copies and the scalar predicated path (~1.5TB/s cap) is kept for + # boundary chunks only. Off cc90 occupancy already hides the load + # latency and bulk copies measured ~4% slower, so the fast path is + # compiled out there. (cp_async/sync alternates measured slower + # than the scalar path on both devices.) + full_tile = (is_valid_chunk and (bos + BT <= eos)) if bulk_copy else False + if full_tile: + T.copy(q[bos: bos + BT, i_h, i_k * BK: i_k * BK + BK], q_shared) + T.copy(k[bos: bos + BT, i_h, i_k * BK: i_k * BK + BK], k_shared) + T.copy(a[bos: bos + BT, i_h, i_k * BK: i_k * BK + BK], a_shared) + T.copy(b[bos: bos + BT, i_h, i_k * BK: i_k * BK + BK], b_shared) + T.copy(gi[bos: bos + BT, i_h, i_k * BK: i_k * BK + BK], gi_shared) + if DERIVE_GE: + # ge_shared first holds raw gk, then derives in place. + T.copy(ge[bos: bos + BT, i_h, i_k * BK: i_k * BK + BK], ge_shared) + for r, c in T.Parallel(BT, BK): + ge_shared[r, c] = ( + gi_shared[r, c] - ge_shared[r, c] * T.Cast(acc_dtype, cumsum_scale_value) + ) + else: + T.copy(ge[bos: bos + BT, i_h, i_k * BK: i_k * BK + BK], ge_shared) + else: + for r, c in T.Parallel(BT, BK): + t = bos + r + k_idx = i_k * BK + c + if t < eos and k_idx < K: + q_shared[r, c] = q[t, i_h, k_idx] + k_shared[r, c] = k[t, i_h, k_idx] + a_shared[r, c] = a[t, i_h, k_idx] + b_shared[r, c] = b[t, i_h, k_idx] + if DERIVE_GE: + giv = gi[t, i_h, k_idx] + gi_shared[r, c] = giv + ge_shared[r, c] = giv - T.Cast(acc_dtype, ge[t, i_h, k_idx]) * \ + T.Cast(acc_dtype, cumsum_scale_value) + else: + gi_shared[r, c] = gi[t, i_h, k_idx] + ge_shared[r, c] = ge[t, i_h, k_idx] + else: + q_shared[r, c] = T.Cast(in_dtype, 0.0) + k_shared[r, c] = T.Cast(in_dtype, 0.0) + a_shared[r, c] = T.Cast(in_dtype, 0.0) + b_shared[r, c] = T.Cast(in_dtype, 0.0) + gi_shared[r, c] = T.float32(0.0) + ge_shared[r, c] = T.float32(0.0) + + if FUSE_QSIDE_DA: + T.clear(dAqk_frag) + T.clear(dAqb_frag) + for i_v in T.serial(T.ceildiv(V, BV)): + if full_tile: + T.copy(do[bos: bos + BT, i_h, i_v * BV: i_v * BV + BV], do_shared) + T.copy(v[bos: bos + BT, i_h, i_v * BV: i_v * BV + BV], v_shared) + T.copy(v_new[bos: bos + BT, i_h, i_v * BV: i_v * BV + BV], v_new_shared) + else: + for r, c in T.Parallel(BT, BV): + t = bos + r + g_v = i_v * BV + c + if (t < eos) and (g_v < V): + do_shared[r, c] = do[t, i_h, g_v] + v_shared[r, c] = v[t, i_h, g_v] + v_new_shared[r, c] = v_new[t, i_h, g_v] + else: + do_shared[r, c] = T.Cast(in_dtype, 0.0) + v_shared[r, c] = T.Cast(in_dtype, 0.0) + v_new_shared[r, c] = T.Cast(in_dtype, 0.0) + T.gemm(do_shared, v_shared, dAqk_frag, transpose_B=True) + T.gemm(do_shared, v_new_shared, dAqb_frag, transpose_B=True) + for r, c in T.Parallel(BT, BT): + if r >= c: + # Match the saved path's materialized bf16 boundary. + # Without this explicit round, fused q-side recompute + # keeps dAqk/dAqb in fp32 and can create one-ULP + # dq/dk spikes versus the default FLA/saved envelope. + dA1_shared[r, c] = T.Cast( + acc_dtype, + T.Cast(in_dtype, dAqk_frag[r, c] * T.Cast(acc_dtype, scale_value)), + ) + dA2_shared[r, c] = T.Cast( + acc_dtype, + T.Cast(in_dtype, dAqb_frag[r, c] * T.Cast(acc_dtype, scale_value)), + ) + else: + dA1_shared[r, c] = T.float32(0.0) + dA2_shared[r, c] = T.float32(0.0) + else: + for r, c in T.Parallel(BT, BT): + t = bos + r + if t < eos: + dA1_shared[r, c] = T.Cast(acc_dtype, dAqk[t, i_h, c]) + dA2_shared[r, c] = T.Cast(acc_dtype, dAqb[t, i_h, c]) + else: + dA1_shared[r, c] = T.float32(0.0) + dA2_shared[r, c] = T.float32(0.0) + + # q-side causal mask: keep the inclusive lower triangle + for r, c in T.Parallel(BT, BT): + if r < c: + dA1_shared[r, c] = T.float32(0.0) + dA2_shared[r, c] = T.float32(0.0) + + # Compute stabilization offset from gi at mid-point of valid tokens + valid_len = T.min(eos - bos, BT) + mid = valid_len // 2 + offset = T.alloc_shared((BK,), acc_dtype) + for c in T.Parallel(BK): + offset[c] = gi_shared[mid, c] + + # Compute stabilized ops + q_ops = T.alloc_shared((BT, BK), acc_dtype) + k_ops = T.alloc_shared((BT, BK), acc_dtype) + a_ops = T.alloc_shared((BT, BK), acc_dtype) + b_ops = T.alloc_shared((BT, BK), acc_dtype) + for r, c in T.Parallel(BT, BK): + q_ops[r, c] = T.Cast(acc_dtype, q_shared[r, c]) * T.exp2(gi_shared[r, c] - offset[c]) + k_ops[r, c] = T.Cast(acc_dtype, k_shared[r, c]) * T.exp2(-(gi_shared[r, c] - offset[c])) + b_ops[r, c] = T.Cast(acc_dtype, b_shared[r, c]) * T.exp2(-(gi_shared[r, c] - offset[c])) + a_ops[r, c] = T.Cast(acc_dtype, a_shared[r, c]) * T.exp2(ge_shared[r, c] - offset[c]) + + # Intra-chunk gradients via GEMMs + dq_intra = T.alloc_fragment((BT, BK), acc_dtype) + da_intra = T.alloc_fragment((BT, BK), acc_dtype) + dk_intra = T.alloc_fragment((BT, BK), acc_dtype) + db_intra = T.alloc_fragment((BT, BK), acc_dtype) + T.clear(dq_intra) + T.clear(da_intra) + T.clear(dk_intra) + T.clear(db_intra) + + # q-side group: dq += dAqk @ k_ops + dAqb @ b_ops, plus the q-side + # partials of dk/db + T.gemm(dA1_shared, k_ops, dq_intra) + T.gemm(dA2_shared, b_ops, dq_intra) + T.gemm(dA1_shared, q_ops, dk_intra, transpose_A=True) + T.gemm(dA2_shared, q_ops, db_intra, transpose_A=True) + + # Overwrite the staging tiles with the a-side pair, masked to the + # strict lower triangle + T.sync_threads() + if full_tile: + T.copy(dAak[bos: bos + BT, i_h, 0:BT], dA1_shared) + T.copy(dAab[bos: bos + BT, i_h, 0:BT], dA2_shared) + else: + for r, c in T.Parallel(BT, BT): + t = bos + r + if t < eos: + dA1_shared[r, c] = T.Cast(acc_dtype, dAak[t, i_h, c]) + dA2_shared[r, c] = T.Cast(acc_dtype, dAab[t, i_h, c]) + else: + dA1_shared[r, c] = T.float32(0.0) + dA2_shared[r, c] = T.float32(0.0) + for r, c in T.Parallel(BT, BT): + if r <= c: + dA1_shared[r, c] = T.float32(0.0) + dA2_shared[r, c] = T.float32(0.0) + + # a-side group: da += dAak @ k_ops + dAab @ b_ops, plus the a-side + # partials of dk/db + T.gemm(dA1_shared, k_ops, da_intra) + T.gemm(dA2_shared, b_ops, da_intra) + T.gemm(dA1_shared, a_ops, dk_intra, transpose_A=True) + T.gemm(dA2_shared, a_ops, db_intra, transpose_A=True) + + # Load inter-chunk gradients and g_last + if full_tile: + T.copy(dqg[bos: bos + BT, i_h, i_k * BK: i_k * BK + BK], dqg_shared) + T.copy(dkg[bos: bos + BT, i_h, i_k * BK: i_k * BK + BK], dkg_shared) + T.copy(dag[bos: bos + BT, i_h, i_k * BK: i_k * BK + BK], dag_shared) + T.copy(dbg[bos: bos + BT, i_h, i_k * BK: i_k * BK + BK], dbg_shared) + else: + for r, c in T.Parallel(BT, BK): + t = bos + r + k_idx = i_k * BK + c + if t < eos and k_idx < K: + dqg_shared[r, c] = dqg[t, i_h, k_idx] + dkg_shared[r, c] = dkg[t, i_h, k_idx] + dag_shared[r, c] = dag[t, i_h, k_idx] + dbg_shared[r, c] = dbg[t, i_h, k_idx] + else: + dqg_shared[r, c] = T.Cast(in_dtype, 0.0) + dkg_shared[r, c] = T.Cast(in_dtype, 0.0) + dag_shared[r, c] = T.Cast(in_dtype, 0.0) + dbg_shared[r, c] = T.Cast(in_dtype, 0.0) + for c in T.Parallel(BK): + k_idx = i_k * BK + c + if is_valid_chunk and k_idx < K: + g_last[c] = gi[last_idx, i_h, k_idx] + else: + g_last[c] = T.float32(0.0) + + # Combine intra + inter, un-stabilize, and store + scale_v = T.Cast(acc_dtype, scale_value) + if full_tile: + for r, c in T.Parallel(BT, BK): + t = bos + r + k_idx = i_k * BK + c + dq_val = dq_intra[r, c] * T.exp2(gi_shared[r, c] - offset[c]) + \ + T.Cast(acc_dtype, dqg_shared[r, c]) * T.exp2(gi_shared[r, c]) * scale_v + da_val = da_intra[r, c] * T.exp2(ge_shared[r, c] - offset[c]) + \ + T.Cast(acc_dtype, dag_shared[r, c]) * T.exp2(ge_shared[r, c]) + dk_val = dk_intra[r, c] * T.exp2(-(gi_shared[r, c] - offset[c])) + \ + T.Cast(acc_dtype, dkg_shared[r, c]) * T.exp2(g_last[c] - gi_shared[r, c]) + db_val = db_intra[r, c] * T.exp2(-(gi_shared[r, c] - offset[c])) + \ + T.Cast(acc_dtype, dbg_shared[r, c]) * T.exp2(g_last[c] - gi_shared[r, c]) + + dq[t, i_h, k_idx] = T.Cast(in_dtype, dq_val) + dk[t, i_h, k_idx] = T.Cast(in_dtype, dk_val) + da[t, i_h, k_idx] = T.Cast(in_dtype, da_val) + db[t, i_h, k_idx] = T.Cast(in_dtype, db_val) + q_ops[r, c] = dq_val * T.Cast(acc_dtype, q_shared[r, c]) + da_val * T.Cast(acc_dtype, a_shared[r, c]) - \ + dk_val * T.Cast(acc_dtype, k_shared[r, c]) - db_val * T.Cast(acc_dtype, b_shared[r, c]) + k_ops[r, c] = da_val * T.Cast(acc_dtype, a_shared[r, c]) + else: + for r, c in T.Parallel(BT, BK): + t = bos + r + k_idx = i_k * BK + c + if t < eos and k_idx < K: + dq_val = dq_intra[r, c] * T.exp2(gi_shared[r, c] - offset[c]) + \ + T.Cast(acc_dtype, dqg_shared[r, c]) * T.exp2(gi_shared[r, c]) * scale_v + da_val = da_intra[r, c] * T.exp2(ge_shared[r, c] - offset[c]) + \ + T.Cast(acc_dtype, dag_shared[r, c]) * T.exp2(ge_shared[r, c]) + dk_val = dk_intra[r, c] * T.exp2(-(gi_shared[r, c] - offset[c])) + \ + T.Cast(acc_dtype, dkg_shared[r, c]) * T.exp2(g_last[c] - gi_shared[r, c]) + db_val = db_intra[r, c] * T.exp2(-(gi_shared[r, c] - offset[c])) + \ + T.Cast(acc_dtype, dbg_shared[r, c]) * T.exp2(g_last[c] - gi_shared[r, c]) + + dq[t, i_h, k_idx] = T.Cast(in_dtype, dq_val) + dk[t, i_h, k_idx] = T.Cast(in_dtype, dk_val) + da[t, i_h, k_idx] = T.Cast(in_dtype, da_val) + db[t, i_h, k_idx] = T.Cast(in_dtype, db_val) + q_ops[r, c] = dq_val * T.Cast(acc_dtype, q_shared[r, c]) + da_val * T.Cast(acc_dtype, a_shared[r, c]) - \ + dk_val * T.Cast(acc_dtype, k_shared[r, c]) - db_val * T.Cast(acc_dtype, b_shared[r, c]) + k_ops[r, c] = da_val * T.Cast(acc_dtype, a_shared[r, c]) + else: + q_ops[r, c] = T.float32(0.0) + k_ops[r, c] = T.float32(0.0) + + # Reuse q_ops/k_ops as dgk_raw/dgk_offset scratch and finish the + # reverse cumsum inside this kernel, avoiding two fp32 full-tensor outputs. + # The suffix sum must not read-modify-write q_ops: when BK < threads + # the T.Parallel lowering shares one column across several threads + # (address masked to tid % BK), and a fast sibling can overwrite a + # row before the others read it (hd128 BK=32 dgk corruption). q_ops + # stays read-only here; suffixes land in dgk_cum, where duplicate + # same-value writes are benign. + T.sync_threads() + dgk_suffix = T.alloc_fragment((BK,), acc_dtype) + dgk_cum = T.alloc_shared((BT, BK), acc_dtype) + for c in T.Parallel(BK): + dgk_suffix[c] = T.float32(0.0) + for r_rev in T.serial(BT): + r = BT - 1 - r_rev + for c in T.Parallel(BK): + dgk_suffix[c] += q_ops[r, c] + dgk_cum[r, c] = dgk_suffix[c] + T.sync_threads() + if full_tile: + for r, c in T.Parallel(BT, BK): + k_idx = i_k * BK + c + dgk_output[bos + r, i_h, k_idx] = T.Cast( + out_dtype, dgk_cum[r, c] + dgk_last[i_c, i_h, k_idx] - k_ops[r, c]) + else: + for r, c in T.Parallel(BT, BK): + t = bos + r + k_idx = i_k * BK + c + if t < eos and k_idx < K: + dgk_output[t, i_h, k_idx] = T.Cast( + out_dtype, dgk_cum[r, c] + T.if_then_else(is_valid_chunk, dgk_last[i_c, i_h, k_idx], T.float32(0.0)) - k_ops[r, c]) + + return chunk_dplr_bwd_intra_tl + + +def chunk_dplr_bwd_dqk_intra_fused_qside_into( + q: torch.Tensor, + k: torch.Tensor, + a: torch.Tensor, + b: torch.Tensor, + gi: torch.Tensor, + ge: torch.Tensor | None, + do: torch.Tensor, + v: torch.Tensor, + v_new: torch.Tensor, + dAak: torch.Tensor, + dAab: torch.Tensor, + dqg: torch.Tensor, + dkg: torch.Tensor, + dag: torch.Tensor, + dbg: torch.Tensor, + dgk_last: torch.Tensor, + dq_out: torch.Tensor, + dk_out: torch.Tensor, + da_out: torch.Tensor, + db_out: torch.Tensor, + dgk_out: torch.Tensor, + gk: torch.Tensor | None = None, + scale: float = 1.0, + cu_seqlens: torch.Tensor | None = None, + chunk_size: int = 16, + chunk_layout: ChunkLayout | None = None, + dgk_dtype: torch.dtype | None = None, +): + """Recompute q-side dA tiles on-chip and consume them immediately. + + When ``gk`` is passed (``ge`` may then be None), the kernel runs its + DERIVE_GE specialization: it loads raw bf16 ``gk`` instead of fp32 ``ge`` + and derives ``ge = gi - gk*RCP_LN2`` in-CTA, removing the fp32 ``ge`` + tensor from the training backward. + """ + if ge is None and gk is None: + raise ValueError("chunk_dplr_bwd_dqk_intra_fused_qside_into needs ge or gk") + for out in (dq_out, dk_out, da_out, db_out, dgk_out): + assert out.is_contiguous(), "chunk_dplr_bwd_dqk_intra_fused_qside_into requires contiguous outputs" + derive_ge = gk is not None + B, T_, H, K = q.shape + V = v.shape[-1] + BT = chunk_size + is_varlen = cu_seqlens is not None + + if is_varlen: + assert B == 1 + layout = chunk_layout if chunk_layout is not None else build_varlen_chunk_layout(cu_seqlens, BT, T_) + else: + layout = build_rect_chunk_layout(B, T_, BT, q.device) + n_chunks = layout.chunk_indices.shape[0] + N_tokens = B * T_ + + in_dtype = str(q.dtype).split(".")[-1] + dgk_out_dtype = str(dgk_dtype or dgk_out.dtype).split(".")[-1] + + q_f = q.reshape(N_tokens, H, K).contiguous() + k_f = k.reshape(N_tokens, H, K).contiguous() + a_f = a.reshape(N_tokens, H, K).contiguous() + b_f = b.reshape(N_tokens, H, K).contiguous() + gi_f = gi.reshape(N_tokens, H, K).contiguous() + ge_f = (gk if derive_ge else ge).reshape(N_tokens, H, K).contiguous() + do_f = do.reshape(N_tokens, H, V).contiguous() + v_f = v.reshape(N_tokens, H, V).contiguous() + v_new_f = v_new.reshape(N_tokens, H, V).contiguous() + dAak_f = dAak.reshape(N_tokens, H, BT).contiguous() + dAab_f = dAab.reshape(N_tokens, H, BT).contiguous() + dqg_f = dqg.reshape(N_tokens, H, K).contiguous() + dkg_f = dkg.reshape(N_tokens, H, K).contiguous() + dag_f = dag.reshape(N_tokens, H, K).contiguous() + dbg_f = dbg.reshape(N_tokens, H, K).contiguous() + dgk_last_f = dgk_last.reshape(n_chunks, H, K).contiguous() + dq_f = dq_out.reshape(N_tokens, H, K).contiguous() + dk_f = dk_out.reshape(N_tokens, H, K).contiguous() + da_f = da_out.reshape(N_tokens, H, K).contiguous() + db_f = db_out.reshape(N_tokens, H, K).contiguous() + dgk_out_f = dgk_out.reshape(N_tokens, H, K).contiguous() + # the fused kernel specialization ignores dAqk/dAqb (their leading extent + # is a separate JIT symbol), so size-1 dummies satisfy the signature + dummy_dA = q.new_empty((1, H, BT)) + # Deterministic BV selection (env knob removed): full V when possible, + # else the largest power-of-two divisor of V not exceeding 64. + bv = min(64, V) + while V % bv != 0 and bv > 1: + bv //= 2 + config = _a_bwd_config(K, BT, in_dtype, q.device) + if BT < 32: + config["threads"] = min(config["threads"], 32) + + kernel = _chunk_dplr_bwd_kernel_intra( + H, K, BT, in_dtype, dgk_out_dtype, float(scale), + FUSE_QSIDE_DA=True, V=V, BV=bv, + DERIVE_GE=derive_ge, + gk_dtype=str(gk.dtype).split(".")[-1] if derive_ge else None, + **config, + ) + kernel( + q_f, k_f, a_f, b_f, gi_f, ge_f, dummy_dA, dummy_dA, dAak_f, dAab_f, + do_f, v_f, v_new_f, + dqg_f, dkg_f, dag_f, dbg_f, layout.cu_seqlens, layout.chunk_indices, dgk_last_f, + dq_f, dk_f, da_f, db_f, dgk_out_f, + ) + return dq_out, dk_out, da_out, db_out, dgk_out diff --git a/fla/ops/generalized_delta_rule/dplr/backends/tilelang/chunk_A_fwd.py b/fla/ops/generalized_delta_rule/dplr/backends/tilelang/chunk_A_fwd.py new file mode 100644 index 0000000000..f2ffe1010f --- /dev/null +++ b/fla/ops/generalized_delta_rule/dplr/backends/tilelang/chunk_A_fwd.py @@ -0,0 +1,552 @@ +# Copyright (c) 2023-2026, Songlin Yang, Yu Zhang, Zhiyuan Li +# +# This source code is licensed under the MIT license found in the +# LICENSE file in the root directory of this source tree. +# For a list of all contributors, visit: +# https://github.com/fla-org/flash-linear-attention/graphs/contributors + +"""DPLR intra-chunk A-matrix forward. + +Produces four (B, T, H, BT) attention matrices: + A_qk[i, j] = scale * sum_k q[i, k] * k[j, k] * exp2(gi[i, k] - gi[j, k]) i >= j + A_qb[i, j] = scale * sum_k q[i, k] * b[j, k] * exp2(gi[i, k] - gi[j, k]) i >= j + A_ak[i, j] = sum_k a[i, k] * k[j, k] * exp2(ge[i, k] - gi[j, k]) i > j + A_ab[i, j] = sum_k a[i, k] * b[j, k] * exp2(ge[i, k] - gi[j, k]) i > j + +Plus four pre-gated tensors: + qg = scale * q * exp2(gi) + kg = k * exp2(gi_last - gi) + ag = a * exp2(ge) + bg = b * exp2(gi_last - gi) + +The kernel uses a centered tensorcore factorization to recover target training +throughput. Unit-scale random stress is still documented as a diagnostic case, +but the target distribution and FLA's own pure-PyTorch reference show that +stress case is not a TileLang-only correctness signal. +""" + +import tilelang +import tilelang.language as T +import torch + +from fla.ops.utils.constant import RCP_LN2 +from fla.utils import get_device_capability + +from .layout import ChunkLayout, build_rect_chunk_layout, build_varlen_chunk_layout + + +def _select_a_fwd_threads(major: int, BT: int) -> int: + if BT < 32: + return 32 + if major >= 9 and BT >= 64: + return 256 + return 128 + + +def _chunk_dplr_fwd_intra_tensorcore_kernel_impl( + H, K, BT, in_dtype, + scale_value: float, + threads: int = 128, +): + acc_dtype = "float32" + # fp16 cannot hold the centered exp2 operands (|centered_gi| can reach + # ~115 log2 at BT=32); keep the q-side GEMM operands in fp32 there, as + # FLA's Triton kernel does for both dtypes. + qside_dtype = acc_dtype if in_dtype == "float16" else in_dtype + n_tokens, n_seq_plus_one, n_chunks = T.dynamic( + "n_tokens, n_seq_plus_one, n_chunks" + ) + + @T.prim_func + def chunk_dplr_fwd_intra_tensorcore_tl( + q: T.Tensor((n_tokens, H, K), in_dtype), + k: T.Tensor((n_tokens, H, K), in_dtype), + a: T.Tensor((n_tokens, H, K), in_dtype), + b: T.Tensor((n_tokens, H, K), in_dtype), + gi: T.Tensor((n_tokens, H, K), acc_dtype), + ge: T.Tensor((n_tokens, H, K), acc_dtype), + cu_seqlens: T.Tensor((n_seq_plus_one,), "int32"), + chunk_indices: T.Tensor((n_chunks, 2), "int32"), + qg: T.Tensor((n_tokens, H, K), in_dtype), + kg: T.Tensor((n_tokens, H, K), in_dtype), + ag: T.Tensor((n_tokens, H, K), in_dtype), + bg: T.Tensor((n_tokens, H, K), in_dtype), + Aqk: T.Tensor((n_tokens, H, BT), in_dtype), + Aqb: T.Tensor((n_tokens, H, BT), in_dtype), + Aab: T.Tensor((n_tokens, H, BT), "float16"), + Aak: T.Tensor((n_tokens, H, BT), "float16"), + ): + with T.Kernel(n_chunks, H, threads=threads) as (i_c, i_h): + i_n = chunk_indices[i_c, 0] + i_t = chunk_indices[i_c, 1] + safe_i_n = T.max(i_n, 0) + seq_bos = cu_seqlens[safe_i_n] + seq_eos = cu_seqlens[safe_i_n + 1] + bos_raw = seq_bos + i_t * BT + eos_raw = T.min(bos_raw + BT, seq_eos) + is_valid_chunk = i_n >= 0 + bos = T.if_then_else(is_valid_chunk, bos_raw, T.int32(0)) + eos = T.if_then_else(is_valid_chunk, eos_raw, T.int32(0)) + last_idx = T.max(eos - 1, 0) + valid_len = eos - bos + + q_mat = T.alloc_shared((BT, K), qside_dtype) + k_mat = T.alloc_shared((BT, K), qside_dtype) + a_mat = T.alloc_shared((BT, K), acc_dtype) + b_mat = T.alloc_shared((BT, K), qside_dtype) + k_mat_acc = T.alloc_shared((BT, K), acc_dtype) + b_mat_acc = T.alloc_shared((BT, K), acc_dtype) + g_last = T.alloc_shared((K,), acc_dtype) + gate_offset = T.alloc_shared((K,), acc_dtype) + + scale_v = T.Cast(acc_dtype, scale_value) + + for c in T.Parallel(K): + if bos < eos: + g_last[c] = gi[last_idx, i_h, c] + else: + g_last[c] = 0.0 + mid = valid_len // 2 + for c in T.Parallel(K): + if bos < eos: + gate_offset[c] = gi[bos + mid, i_h, c] + else: + gate_offset[c] = 0.0 + + # Build centered TensorCore operands while keeping the public + # pre-gated tensors at the validated dtype boundary. + # Centering follows FLA's tensorcore variant and avoids multiplying + # heavily under/over-scaled bf16 operands when gi drifts across a + # chunk. + for r, c in T.Parallel(BT, K): + t = bos + r + if t < eos: + qv = T.Cast(acc_dtype, q[t, i_h, c]) + kv = T.Cast(acc_dtype, k[t, i_h, c]) + av = T.Cast(acc_dtype, a[t, i_h, c]) + bv = T.Cast(acc_dtype, b[t, i_h, c]) + giv = gi[t, i_h, c] + gev = ge[t, i_h, c] + q_scaled = qv * scale_v + qg_v = T.Cast(in_dtype, q_scaled * T.exp2(giv)) + ag_v = T.Cast(in_dtype, av * T.exp2(gev)) + kg_v = T.Cast(in_dtype, kv * T.exp2(-giv + g_last[c])) + bg_v = T.Cast(in_dtype, bv * T.exp2(-giv + g_last[c])) + qg[t, i_h, c] = qg_v + ag[t, i_h, c] = ag_v + kg[t, i_h, c] = kg_v + bg[t, i_h, c] = bg_v + centered_gi = giv - gate_offset[c] + centered_ge = gev - gate_offset[c] + q_mat[r, c] = T.Cast(qside_dtype, q_scaled * T.exp2(centered_gi)) + k_mat[r, c] = T.Cast(qside_dtype, kv * T.exp2(-centered_gi)) + b_mat[r, c] = T.Cast(qside_dtype, bv * T.exp2(-centered_gi)) + a_mat[r, c] = av * T.exp2(centered_ge) + k_mat_acc[r, c] = kv * T.exp2(-centered_gi) + b_mat_acc[r, c] = bv * T.exp2(-centered_gi) + else: + q_mat[r, c] = T.Cast(qside_dtype, 0.0) + k_mat[r, c] = T.Cast(qside_dtype, 0.0) + b_mat[r, c] = T.Cast(qside_dtype, 0.0) + a_mat[r, c] = T.Cast(acc_dtype, 0.0) + k_mat_acc[r, c] = T.Cast(acc_dtype, 0.0) + b_mat_acc[r, c] = T.Cast(acc_dtype, 0.0) + + A_qk_frag = T.alloc_fragment((BT, BT), acc_dtype) + A_qb_frag = T.alloc_fragment((BT, BT), acc_dtype) + A_ak_frag = T.alloc_fragment((BT, BT), acc_dtype) + A_ab_frag = T.alloc_fragment((BT, BT), acc_dtype) + T.gemm(q_mat, k_mat, A_qk_frag, transpose_B=True, clear_accum=True) + T.gemm(q_mat, b_mat, A_qb_frag, transpose_B=True, clear_accum=True) + T.gemm(a_mat, k_mat_acc, A_ak_frag, transpose_B=True, clear_accum=True) + T.gemm(a_mat, b_mat_acc, A_ab_frag, transpose_B=True, clear_accum=True) + + # Pairwise fused masked stores: q-side (bf16-gemm fragments) and + # a-side (fp32-gemm fragments) stay separate because wgmma assigns + # them different fragment layouts. Aab/Aak are stored fp16 + # (|Aab| ~<1 from decay; probe shows WY inverse err ~1e-4). + for r, c in T.Parallel(BT, BT): + t = bos + r + if t < eos: + valid_q = (c < valid_len) and (r >= c) + Aqk[t, i_h, c] = T.Cast( + in_dtype, + T.if_then_else(valid_q, A_qk_frag[r, c], T.Cast(acc_dtype, 0.0)), + ) + Aqb[t, i_h, c] = T.Cast( + in_dtype, + T.if_then_else(valid_q, A_qb_frag[r, c], T.Cast(acc_dtype, 0.0)), + ) + for r, c in T.Parallel(BT, BT): + t = bos + r + if t < eos: + valid_a = (c < valid_len) and (r > c) + Aak[t, i_h, c] = T.Cast( + "float16", + T.if_then_else(valid_a, A_ak_frag[r, c], T.Cast(acc_dtype, 0.0)), + ) + Aab[t, i_h, c] = T.Cast( + "float16", + T.if_then_else(valid_a, A_ab_frag[r, c], T.Cast(acc_dtype, 0.0)), + ) + + return chunk_dplr_fwd_intra_tensorcore_tl + + +@tilelang.jit( + pass_configs={ + tilelang.PassConfigKey.TL_ENABLE_FAST_MATH: True, + tilelang.PassConfigKey.TL_DISABLE_DATA_RACE_CHECK: False, + }, +) +def _chunk_dplr_fwd_intra_tensorcore_kernel_vec( + H, K, BT, in_dtype, + scale_value: float, + threads: int = 128, +): + return _chunk_dplr_fwd_intra_tensorcore_kernel_impl(H, K, BT, in_dtype, scale_value, threads) + + +@tilelang.jit( + pass_configs={ + tilelang.PassConfigKey.TL_ENABLE_FAST_MATH: True, + tilelang.PassConfigKey.TL_DISABLE_DATA_RACE_CHECK: False, + tilelang.PassConfigKey.TIR_DISABLE_VECTORIZE: True, + }, +) +def _chunk_dplr_fwd_intra_tensorcore_kernel_novec( + H, K, BT, in_dtype, + scale_value: float, + threads: int = 128, +): + return _chunk_dplr_fwd_intra_tensorcore_kernel_impl(H, K, BT, in_dtype, scale_value, threads) + + +def _chunk_dplr_fwd_intra_tensorcore_kernel( + H, K, BT, in_dtype, + scale_value: float, + threads: int = 128, +): + # At BT <= 16 the vectorized plan leaves the intra GEMMs without a valid + # warp partition (M=16 tiles need each warp to own >= 16 rows); compile + # that shape without vectorization. + if BT <= 16: + return _chunk_dplr_fwd_intra_tensorcore_kernel_novec(H, K, BT, in_dtype, scale_value, threads) + return _chunk_dplr_fwd_intra_tensorcore_kernel_vec(H, K, BT, in_dtype, scale_value, threads) + + +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, + cu_seqlens: torch.Tensor | None = None, + chunk_layout: ChunkLayout | None = None, +) -> tuple[torch.Tensor, ...]: + B, T_, H, K = q.shape + BT = chunk_size + is_varlen = cu_seqlens is not None + if K not in (64, 128): + raise NotImplementedError("chunk_dplr_fwd_intra is validated for head_dim 64 and 128.") + + if is_varlen: + assert B == 1 + layout = chunk_layout if chunk_layout is not None else build_varlen_chunk_layout(cu_seqlens, BT, T_) + else: + layout = build_rect_chunk_layout(B, T_, BT, q.device) + N_tokens = B * T_ + + in_dtype = str(q.dtype).split(".")[-1] + + q_f = q.reshape(N_tokens, H, K).contiguous() + k_f = k.reshape(N_tokens, H, K).contiguous() + a_f = a.reshape(N_tokens, H, K).contiguous() + b_f = b.reshape(N_tokens, H, K).contiguous() + gi_f = gi.reshape(N_tokens, H, K).contiguous() + ge_f = ge.reshape(N_tokens, H, K).contiguous() + + major = get_device_capability(q.device.index)[0] if q.is_cuda else 0 + threads = _select_a_fwd_threads(major, BT) + kernel = _chunk_dplr_fwd_intra_tensorcore_kernel( + H, K, BT, in_dtype, float(scale), threads=threads, + ) + + qg_f = torch.empty((N_tokens, H, K), dtype=q.dtype, device=q.device) + kg_f = torch.empty((N_tokens, H, K), dtype=q.dtype, device=q.device) + ag_f = torch.empty((N_tokens, H, K), dtype=q.dtype, device=q.device) + bg_f = torch.empty((N_tokens, H, K), dtype=q.dtype, device=q.device) + Aqk_f = torch.empty((N_tokens, H, BT), dtype=q.dtype, device=q.device) + Aqb_f = torch.empty((N_tokens, H, BT), dtype=q.dtype, device=q.device) + Aab_f = torch.empty((N_tokens, H, BT), dtype=torch.float16, device=q.device) + Aak_f = torch.empty((N_tokens, H, BT), dtype=torch.float16, device=q.device) + kernel( + q_f, k_f, a_f, b_f, gi_f, ge_f, layout.cu_seqlens, layout.chunk_indices, + qg_f, kg_f, ag_f, bg_f, Aqk_f, Aqb_f, Aab_f, Aak_f, + ) + + qg = qg_f.view(B, T_, H, K) + kg = kg_f.view(B, T_, H, K) + ag = ag_f.view(B, T_, H, K) + bg = bg_f.view(B, T_, H, K) + Aqk = Aqk_f.view(B, T_, H, BT) + Aqb = Aqb_f.view(B, T_, H, BT) + Aab = Aab_f.view(B, T_, H, BT) + Aak = Aak_f.view(B, T_, H, BT) + + return Aab, Aqk, Aak, Aqb, qg, kg, ag, bg + + +@tilelang.jit( + pass_configs={ + tilelang.PassConfigKey.TL_ENABLE_FAST_MATH: True, + tilelang.PassConfigKey.TL_DISABLE_DATA_RACE_CHECK: False, + }, +) +def _chunk_dplr_fwd_intra_from_gk_tensorcore_kernel( + H, K, BT, in_dtype, + scale_value: float, + cumsum_scale_value: float, + threads: int = 128, + gk_dtype: str | None = None, +): + """K=64 A-stage with the chunk-local gi cumsum computed inside the CTA.""" + acc_dtype = "float32" + # fp16 cannot hold the centered exp2 operands (|centered_gi| can reach + # ~115 log2 at BT=32); keep the q-side GEMM operands in fp32 there, as + # FLA's Triton kernel does for both dtypes. + qside_dtype = acc_dtype if in_dtype == "float16" else in_dtype + # raw gk may stay in its own dtype (e.g. fp32); all gate math is fp32 + gk_dtype = gk_dtype or in_dtype + n_tokens, n_seq_plus_one, n_chunks = T.dynamic( + "n_tokens, n_seq_plus_one, n_chunks" + ) + + @T.prim_func + def chunk_dplr_fwd_intra_from_gk_tl( + q: T.Tensor((n_tokens, H, K), in_dtype), + k: T.Tensor((n_tokens, H, K), in_dtype), + a: T.Tensor((n_tokens, H, K), in_dtype), + b: T.Tensor((n_tokens, H, K), in_dtype), + gk: T.Tensor((n_tokens, H, K), gk_dtype), + cu_seqlens: T.Tensor((n_seq_plus_one,), "int32"), + chunk_indices: T.Tensor((n_chunks, 2), "int32"), + qg: T.Tensor((n_tokens, H, K), in_dtype), + kg: T.Tensor((n_tokens, H, K), in_dtype), + ag: T.Tensor((n_tokens, H, K), in_dtype), + bg: T.Tensor((n_tokens, H, K), in_dtype), + Aqk: T.Tensor((n_tokens, H, BT), in_dtype), + Aqb: T.Tensor((n_tokens, H, BT), in_dtype), + Aab: T.Tensor((n_tokens, H, BT), "float16"), + Aak: T.Tensor((n_tokens, H, BT), "float16"), + gi_out: T.Tensor((n_tokens, H, K), acc_dtype), + ): + with T.Kernel(n_chunks, H, threads=threads) as (i_c, i_h): + i_n = chunk_indices[i_c, 0] + i_t = chunk_indices[i_c, 1] + safe_i_n = T.max(i_n, 0) + seq_bos = cu_seqlens[safe_i_n] + seq_eos = cu_seqlens[safe_i_n + 1] + bos_raw = seq_bos + i_t * BT + eos_raw = T.min(bos_raw + BT, seq_eos) + is_valid_chunk = i_n >= 0 + bos = T.if_then_else(is_valid_chunk, bos_raw, T.int32(0)) + eos = T.if_then_else(is_valid_chunk, eos_raw, T.int32(0)) + valid_len = eos - bos + + q_mat = T.alloc_shared((BT, K), qside_dtype) + k_mat = T.alloc_shared((BT, K), qside_dtype) + a_mat = T.alloc_shared((BT, K), acc_dtype) + b_mat = T.alloc_shared((BT, K), qside_dtype) + k_mat_acc = T.alloc_shared((BT, K), acc_dtype) + b_mat_acc = T.alloc_shared((BT, K), acc_dtype) + gi_mat = T.alloc_shared((BT, K), acc_dtype) + g_last = T.alloc_shared((K,), acc_dtype) + gate_offset = T.alloc_shared((K,), acc_dtype) + prefix_acc = T.alloc_fragment((K,), acc_dtype) + + scale_v = T.Cast(acc_dtype, scale_value) + cumsum_scale = T.Cast(acc_dtype, cumsum_scale_value) + + # hoist the gk tile into shared for the serial scan; the gating + # loop below re-reads gk from global (L2-resident by then), which + # ends this tile's lifetime at the scan and lets the allocator + # overlap it with the fp32 operand tiles — needed to fit a 99KB + # smem cap at BT=64 with fp32 gates + gk_shared = T.alloc_shared((BT, K), gk_dtype) + for r, c in T.Parallel(BT, K): + t = bos + r + if t < eos: + gk_shared[r, c] = gk[t, i_h, c] + else: + gk_shared[r, c] = T.Cast(gk_dtype, 0.0) + + for c in T.Parallel(K): + prefix_acc[c] = T.Cast(acc_dtype, 0.0) + + for r in T.serial(BT): + t = bos + r + for c in T.Parallel(K): + if t < eos: + prefix_acc[c] += T.Cast(acc_dtype, gk_shared[r, c]) + gi_mat[r, c] = prefix_acc[c] * cumsum_scale + else: + gi_mat[r, c] = T.Cast(acc_dtype, 0.0) + + # one batched coalesced store of gi instead of one per serial step + for r, c in T.Parallel(BT, K): + t = bos + r + if t < eos: + gi_out[t, i_h, c] = gi_mat[r, c] + + for c in T.Parallel(K): + if bos < eos: + g_last[c] = prefix_acc[c] * cumsum_scale + else: + g_last[c] = T.Cast(acc_dtype, 0.0) + + mid = valid_len // 2 + for c in T.Parallel(K): + if bos < eos: + gate_offset[c] = gi_mat[mid, c] + else: + gate_offset[c] = T.Cast(acc_dtype, 0.0) + + for r, c in T.Parallel(BT, K): + t = bos + r + if t < eos: + qv = T.Cast(acc_dtype, q[t, i_h, c]) + kv = T.Cast(acc_dtype, k[t, i_h, c]) + av = T.Cast(acc_dtype, a[t, i_h, c]) + bv = T.Cast(acc_dtype, b[t, i_h, c]) + gkv = T.Cast(acc_dtype, gk[t, i_h, c]) + giv = gi_mat[r, c] + gev = giv - gkv * cumsum_scale + q_scaled = qv * scale_v + qg_v = T.Cast(in_dtype, q_scaled * T.exp2(giv)) + ag_v = T.Cast(in_dtype, av * T.exp2(gev)) + kg_v = T.Cast(in_dtype, kv * T.exp2(-giv + g_last[c])) + bg_v = T.Cast(in_dtype, bv * T.exp2(-giv + g_last[c])) + qg[t, i_h, c] = qg_v + ag[t, i_h, c] = ag_v + kg[t, i_h, c] = kg_v + bg[t, i_h, c] = bg_v + centered_gi = giv - gate_offset[c] + centered_ge = gev - gate_offset[c] + q_mat[r, c] = T.Cast(qside_dtype, q_scaled * T.exp2(centered_gi)) + k_mat[r, c] = T.Cast(qside_dtype, kv * T.exp2(-centered_gi)) + b_mat[r, c] = T.Cast(qside_dtype, bv * T.exp2(-centered_gi)) + a_mat[r, c] = av * T.exp2(centered_ge) + k_mat_acc[r, c] = kv * T.exp2(-centered_gi) + b_mat_acc[r, c] = bv * T.exp2(-centered_gi) + else: + q_mat[r, c] = T.Cast(qside_dtype, 0.0) + k_mat[r, c] = T.Cast(qside_dtype, 0.0) + b_mat[r, c] = T.Cast(qside_dtype, 0.0) + a_mat[r, c] = T.Cast(acc_dtype, 0.0) + k_mat_acc[r, c] = T.Cast(acc_dtype, 0.0) + b_mat_acc[r, c] = T.Cast(acc_dtype, 0.0) + + A_qk_frag = T.alloc_fragment((BT, BT), acc_dtype) + A_qb_frag = T.alloc_fragment((BT, BT), acc_dtype) + A_ak_frag = T.alloc_fragment((BT, BT), acc_dtype) + A_ab_frag = T.alloc_fragment((BT, BT), acc_dtype) + T.gemm(q_mat, k_mat, A_qk_frag, transpose_B=True, clear_accum=True) + T.gemm(q_mat, b_mat, A_qb_frag, transpose_B=True, clear_accum=True) + T.gemm(a_mat, k_mat_acc, A_ak_frag, transpose_B=True, clear_accum=True) + T.gemm(a_mat, b_mat_acc, A_ab_frag, transpose_B=True, clear_accum=True) + + # Pairwise fused masked stores (see chunk_dplr_fwd_intra_tensorcore_tl). + for r, c in T.Parallel(BT, BT): + t = bos + r + if t < eos: + valid_q = (c < valid_len) and (r >= c) + Aqk[t, i_h, c] = T.Cast( + in_dtype, + T.if_then_else(valid_q, A_qk_frag[r, c], T.Cast(acc_dtype, 0.0)), + ) + Aqb[t, i_h, c] = T.Cast( + in_dtype, + T.if_then_else(valid_q, A_qb_frag[r, c], T.Cast(acc_dtype, 0.0)), + ) + for r, c in T.Parallel(BT, BT): + t = bos + r + if t < eos: + valid_a = (c < valid_len) and (r > c) + Aak[t, i_h, c] = T.Cast( + "float16", + T.if_then_else(valid_a, A_ak_frag[r, c], T.Cast(acc_dtype, 0.0)), + ) + Aab[t, i_h, c] = T.Cast( + "float16", + T.if_then_else(valid_a, A_ab_frag[r, c], T.Cast(acc_dtype, 0.0)), + ) + + return chunk_dplr_fwd_intra_from_gk_tl + + +def chunk_dplr_fwd_intra_from_gk( + q: torch.Tensor, + k: torch.Tensor, + a: torch.Tensor, + b: torch.Tensor, + gk: torch.Tensor, + scale: float, + chunk_size: int, + cu_seqlens: torch.Tensor | None = None, + chunk_layout: ChunkLayout | None = None, +) -> tuple[torch.Tensor, ...]: + B, T_, H, K = q.shape + BT = int(chunk_size) + if K != 64: + raise NotImplementedError("chunk_dplr_fwd_intra_from_gk is currently validated only for K=64.") + + is_varlen = cu_seqlens is not None + if is_varlen: + assert B == 1 + layout = chunk_layout if chunk_layout is not None else build_varlen_chunk_layout(cu_seqlens, BT, T_) + else: + layout = chunk_layout if chunk_layout is not None else build_rect_chunk_layout(B, T_, BT, q.device) + n_tokens = B * T_ + in_dtype = str(q.dtype).split(".")[-1] + + q_f = q.reshape(n_tokens, H, K).contiguous() + k_f = k.reshape(n_tokens, H, K).contiguous() + a_f = a.reshape(n_tokens, H, K).contiguous() + b_f = b.reshape(n_tokens, H, K).contiguous() + gk_f = gk.reshape(n_tokens, H, K).contiguous() + + major = get_device_capability(q.device.index)[0] if q.is_cuda else 0 + threads = _select_a_fwd_threads(major, BT) + kernel = _chunk_dplr_fwd_intra_from_gk_tensorcore_kernel( + H, K, BT, in_dtype, float(scale), float(RCP_LN2), + threads=threads, + gk_dtype=str(gk.dtype).split(".")[-1], + ) + qg_f = torch.empty((n_tokens, H, K), dtype=q.dtype, device=q.device) + kg_f = torch.empty((n_tokens, H, K), dtype=q.dtype, device=q.device) + ag_f = torch.empty((n_tokens, H, K), dtype=q.dtype, device=q.device) + bg_f = torch.empty((n_tokens, H, K), dtype=q.dtype, device=q.device) + Aqk_f = torch.empty((n_tokens, H, BT), dtype=q.dtype, device=q.device) + Aqb_f = torch.empty((n_tokens, H, BT), dtype=q.dtype, device=q.device) + Aab_f = torch.empty((n_tokens, H, BT), dtype=torch.float16, device=q.device) + Aak_f = torch.empty((n_tokens, H, BT), dtype=torch.float16, device=q.device) + gi_f = torch.empty((n_tokens, H, K), dtype=torch.float32, device=q.device) + kernel( + q_f, k_f, a_f, b_f, gk_f, layout.cu_seqlens, layout.chunk_indices, + qg_f, kg_f, ag_f, bg_f, Aqk_f, Aqb_f, Aab_f, Aak_f, gi_f, + ) + + qg = qg_f.view(B, T_, H, K) + kg = kg_f.view(B, T_, H, K) + ag = ag_f.view(B, T_, H, K) + bg = bg_f.view(B, T_, H, K) + Aqk = Aqk_f.view(B, T_, H, BT) + Aqb = Aqb_f.view(B, T_, H, BT) + Aab = Aab_f.view(B, T_, H, BT) + Aak = Aak_f.view(B, T_, H, BT) + gi = gi_f.view(B, T_, H, K) + return Aab, Aqk, Aak, Aqb, qg, kg, ag, bg, gi diff --git a/fla/ops/generalized_delta_rule/dplr/backends/tilelang/chunk_h_fwd.py b/fla/ops/generalized_delta_rule/dplr/backends/tilelang/chunk_h_fwd.py new file mode 100644 index 0000000000..c5dee54ec9 --- /dev/null +++ b/fla/ops/generalized_delta_rule/dplr/backends/tilelang/chunk_h_fwd.py @@ -0,0 +1,286 @@ +# Copyright (c) 2023-2026, Songlin Yang, Yu Zhang, Zhiyuan Li +# +# This source code is licensed under the MIT license found in the +# LICENSE file in the root directory of this source tree. +# For a list of all contributors, visit: +# https://github.com/fla-org/flash-linear-attention/graphs/contributors + +"""DPLR chunk-recurrent state forward. + +Per chunk: + for each BC sub-chunk: + v2 = w @ b_h + u # write to v_new + b_hc += kg @ v + bg @ v2 + b_h = exp2(g_last) * b_h + b_hc + +Adapted from kda/chunk_delta_h_fwd.py (which does the equivalent single-GEMM +recurrence). Differences from KDA: +- KDA: `v_new = u - w@h; b_h += K^T @ v_new` (one outer-product) +- DPLR: `v2 = w@h + u; b_h += kg^T @ v + bg^T @ v2` (two outer-products, no sign flip) +- DPLR sub-chunks the BC inner loop to bound shared-memory pressure +- DPLR applies the decay AFTER accumulating sub-chunks; KDA applies before +""" + +import tilelang +import tilelang.language as T +import torch + +from fla.utils import get_device_capability + +from .layout import ChunkLayout, build_rect_chunk_layout, build_varlen_chunk_layout + + +def _chunk_h_fwd_config(K: int, V: int, device_index: int) -> dict[str, int]: + cap_major = get_device_capability(device_index)[0] + if cap_major == 9 and K <= 128: + return {"BV": 64 if V >= 64 else 32 if V >= 32 else 16, "threads": 256} + elif cap_major == 8: + return {"BV": 32 if V >= 32 else 16, "threads": 128} + return {"BV": 16, "threads": 64} + + +@tilelang.jit( + pass_configs={tilelang.PassConfigKey.TL_ENABLE_FAST_MATH: False, + tilelang.PassConfigKey.TL_DISABLE_DATA_RACE_CHECK: False, + }, +) +def _chunk_dplr_fwd_h_kernel( + H, K, V, BT, BC, + in_dtype, state_dtype, + USE_INITIAL_STATE: bool, + STORE_FINAL_STATE: bool, + BV: int = 32, + threads: int = 128, +): + acc_dtype = "float32" + n_tokens, n_seq_plus_one, n_chunks, n_h0, n_ht = T.dynamic( + "n_tokens, n_seq_plus_one, n_chunks, n_h0, n_ht" + ) + n_seqs = n_seq_plus_one - 1 + + @T.prim_func + def chunk_dplr_fwd_h_tl( + kg: T.Tensor((n_tokens, H, K), in_dtype), + v: T.Tensor((n_tokens, H, V), in_dtype), + w: T.Tensor((n_tokens, H, K), in_dtype), + bg: T.Tensor((n_tokens, H, K), in_dtype), + u: T.Tensor((n_tokens, H, V), in_dtype), + gk: T.Tensor((n_tokens, H, K), acc_dtype), + h0: T.Tensor((n_h0, H, K, V), acc_dtype), + cu_seqlens: T.Tensor((n_seq_plus_one,), "int32"), + chunk_offsets: T.Tensor((n_seq_plus_one,), "int32"), + chunk_indices: T.Tensor((n_chunks, 2), "int32"), + h: T.Tensor((n_chunks, H, K, V), in_dtype), + v_new: T.Tensor((n_tokens, H, V), in_dtype), + ht: T.Tensor((n_ht, H, K, V), state_dtype), + ): + with T.Kernel(T.ceildiv(V, BV), n_seqs, H, threads=threads) as (i_v, i_n, i_h): + bos = cu_seqlens[i_n] + eos = cu_seqlens[i_n + 1] + boh = chunk_offsets[i_n] + n_chunks = chunk_offsets[i_n + 1] - boh + + b_h = T.alloc_fragment((K, BV), acc_dtype) + b_h_shared = T.alloc_shared((K, BV), in_dtype) + b_hc = T.alloc_fragment((K, BV), acc_dtype) + b_hc_kg = T.alloc_fragment((K, BV), acc_dtype) + b_hc_bg = T.alloc_fragment((K, BV), acc_dtype) + b_hc_kg_shared = T.alloc_shared((K, BV), acc_dtype) + b_hc_bg_shared = T.alloc_shared((K, BV), acc_dtype) + v2_frag = T.alloc_fragment((BC, BV), acc_dtype) + v2_shared = T.alloc_shared((BC, BV), acc_dtype) + kg_shared = T.alloc_shared((BC, K), in_dtype) + bg_shared = T.alloc_shared((BC, K), acc_dtype) + w_shared = T.alloc_shared((BC, K), in_dtype) + v_shared = T.alloc_shared((BC, BV), in_dtype) + u_shared = T.alloc_shared((BC, BV), in_dtype) + gk_last_shared = T.alloc_shared((K,), acc_dtype) + + # Init state. + if USE_INITIAL_STATE: + for k_idx, vv in T.Parallel(K, BV): + g_v = i_v * BV + vv + if g_v < V: + b_h[k_idx, vv] = h0[i_n, i_h, k_idx, g_v] + else: + b_h[k_idx, vv] = 0.0 + else: + for k_idx, vv in T.Parallel(K, BV): + b_h[k_idx, vv] = 0.0 + + # Per-chunk loop. + for i_t in T.serial(n_chunks): + # FLA computes w @ h with h rounded to the input dtype, while the + # recurrent state itself remains fp32 across chunks. + T.copy(b_h, b_h_shared) + for k_idx, vv in T.Parallel(K, BV): + g_v = i_v * BV + vv + if g_v < V: + h[boh + i_t, i_h, k_idx, g_v] = b_h_shared[k_idx, vv] + + T.clear(b_hc) + + # Sub-chunk loop. + NC = T.ceildiv(BT, BC) + for i_c in T.serial(NC): + t_off = bos + i_t * BT + i_c * BC + + # Load tiles with zero-padding for partial varlen chunks. + for c, k_idx in T.Parallel(BC, K): + if t_off + c < eos: + kg_shared[c, k_idx] = kg[t_off + c, i_h, k_idx] + bg_shared[c, k_idx] = T.Cast(acc_dtype, bg[t_off + c, i_h, k_idx]) + w_shared[c, k_idx] = w[t_off + c, i_h, k_idx] + else: + kg_shared[c, k_idx] = T.Cast(in_dtype, 0.0) + bg_shared[c, k_idx] = 0.0 + w_shared[c, k_idx] = T.Cast(in_dtype, 0.0) + for c, vv in T.Parallel(BC, BV): + g_v = i_v * BV + vv + if (t_off + c < eos) and (g_v < V): + v_shared[c, vv] = v[t_off + c, i_h, g_v] + u_shared[c, vv] = u[t_off + c, i_h, g_v] + else: + v_shared[c, vv] = T.Cast(in_dtype, 0.0) + u_shared[c, vv] = T.Cast(in_dtype, 0.0) + + T.gemm(w_shared, b_h_shared, v2_frag, clear_accum=True) + for c, vv in T.Parallel(BC, BV): + v2_frag[c, vv] = T.ieee_add(v2_frag[c, vv], T.Cast(acc_dtype, u_shared[c, vv])) + T.copy(v2_frag, v2_shared) + + for c, vv in T.Parallel(BC, BV): + g_v = i_v * BV + vv + if (t_off + c < eos) and (g_v < V): + v_new[t_off + c, i_h, g_v] = T.Cast(in_dtype, v2_frag[c, vv]) + + # b_hc += kg^T @ v + bg^T @ v2. The second product stays in + # fp32 because FLA casts bg to fp32 and consumes fp32 v2 here. + T.gemm( + kg_shared, + v_shared, + b_hc_kg, + transpose_A=True, + clear_accum=True, + ) + T.copy(b_hc_kg, b_hc_kg_shared) + T.gemm( + bg_shared, + v2_shared, + b_hc_bg, + transpose_A=True, + clear_accum=True, + ) + T.copy(b_hc_bg, b_hc_bg_shared) + for k_idx, vv in T.Parallel(K, BV): + b_hc[k_idx, vv] = T.ieee_add( + T.ieee_add(b_hc[k_idx, vv], b_hc_kg_shared[k_idx, vv]), + b_hc_bg_shared[k_idx, vv], + ) + + # Apply decay and accumulate. Clamp last_idx to <= eos-1. + last_idx = T.min(bos + (i_t + 1) * BT - 1, eos - 1) + for k_idx in T.Parallel(K): + gk_last_shared[k_idx] = gk[last_idx, i_h, k_idx] + # Match FLA's two-step update order: + # b_h *= exp2(g_last) + # b_h += b_hc + # Keeping these as separate statements avoids contracting the + # recurrent state update into a different fp32 expression. + for k_idx, vv in T.Parallel(K, BV): + b_h[k_idx, vv] = T.ieee_mul(T.exp2(gk_last_shared[k_idx]), b_h[k_idx, vv]) + for k_idx, vv in T.Parallel(K, BV): + b_h[k_idx, vv] = T.ieee_add(b_h[k_idx, vv], b_hc[k_idx, vv]) + + # Store final state. + if STORE_FINAL_STATE: + for k_idx, vv in T.Parallel(K, BV): + g_v = i_v * BV + vv + if g_v < V: + ht[i_n, i_h, k_idx, g_v] = T.Cast(state_dtype, b_h[k_idx, vv]) + + return chunk_dplr_fwd_h_tl + + +def chunk_dplr_fwd_h( + kg: torch.Tensor, + v: torch.Tensor, + w: torch.Tensor, + u: torch.Tensor, + bg: torch.Tensor, + gk: torch.Tensor, # = gi from cumsum, fp32 + initial_state: torch.Tensor | None = None, + output_final_state: bool = False, + cu_seqlens: torch.Tensor | None = None, + chunk_size: int = 16, + chunk_layout: ChunkLayout | None = None, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + B, T_, H, K = kg.shape + V = v.shape[-1] + BT = chunk_size + is_varlen = cu_seqlens is not None + + if is_varlen: + assert B == 1 + layout = chunk_layout if chunk_layout is not None else build_varlen_chunk_layout(cu_seqlens, BT, T_) + else: + layout = build_rect_chunk_layout(B, T_, BT, kg.device) + chunk_rows = layout.chunk_indices.shape[0] + active_nseq = layout.cu_seqlens.shape[0] - 1 + token_rows = B * T_ + + # Match FLA's Hopper path for the target H800 shape: K=64 uses the whole + # BT=64 chunk in one H recurrence sub-block. Splitting into two BC=32 + # halves changes the fp32 accumulation order and was enough to create rare + # one-ULP bf16 output spikes at long varlen sequence lengths. + cap_major = get_device_capability(kg.device.index)[0] + if cap_major == 9: + BC = min(BT, 64 if K <= 128 else 32) + elif cap_major == 8: + BC = min(BT, 32) + else: + BC = min(BT, 16) + + in_dtype = str(kg.dtype).split(".")[-1] + state_dtype = "float32" + use_h0 = initial_state is not None + store_ht = output_final_state + n_ht = active_nseq if store_ht else 1 + + if use_h0: + h0 = initial_state + else: + h0 = torch.empty((1, H, K, V), dtype=torch.float32, device=kg.device) + if h0.dtype != torch.float32: + h0 = h0.to(torch.float32) + + kg_f = kg.reshape(token_rows, H, K).contiguous() + v_f = v.reshape(token_rows, H, V).contiguous() + w_f = w.reshape(token_rows, H, K).contiguous() + bg_f = bg.reshape(token_rows, H, K).contiguous() + u_f = u.reshape(token_rows, H, V).contiguous() + gk_f = gk.reshape(token_rows, H, K).contiguous() + + kernel = _chunk_dplr_fwd_h_kernel( + H, K, V, BT, BC, + in_dtype, state_dtype, use_h0, store_ht, + **_chunk_h_fwd_config(K, V, kg.device.index), + ) + + h_flat = torch.empty((chunk_rows, H, K, V), dtype=kg.dtype, device=kg.device) + v_new_flat = torch.empty((token_rows, H, V), dtype=v.dtype, device=v.device) + ht = torch.empty((n_ht, H, K, V), dtype=torch.float32, device=kg.device) + kernel( + kg_f, v_f, w_f, bg_f, u_f, gk_f, h0, layout.cu_seqlens, + layout.chunk_offsets, layout.chunk_indices, h_flat, v_new_flat, ht, + ) + + # FLA returns (B, NT, H, K, V); varlen packs all sequences' chunks flat. + h_out = ( + h_flat.view(1, chunk_rows, H, K, V) + if is_varlen + else h_flat.view(B, chunk_rows // B, H, K, V) + ) + v_new = v_new_flat.view(B, T_, H, V) + final_state = ht if store_ht else None + return h_out, v_new, final_state diff --git a/fla/ops/generalized_delta_rule/dplr/backends/tilelang/chunk_ho_fwd.py b/fla/ops/generalized_delta_rule/dplr/backends/tilelang/chunk_ho_fwd.py new file mode 100644 index 0000000000..337f52c6b0 --- /dev/null +++ b/fla/ops/generalized_delta_rule/dplr/backends/tilelang/chunk_ho_fwd.py @@ -0,0 +1,551 @@ +# Copyright (c) 2023-2026, Songlin Yang, Yu Zhang, Zhiyuan Li +# +# This source code is licensed under the MIT license found in the +# LICENSE file in the root directory of this source tree. +# For a list of all contributors, visit: +# https://github.com/fla-org/flash-linear-attention/graphs/contributors + +"""Default fused DPLR H/O forward. + +Computes the chunk start state, per-token ``v_new``, and final output in one +forward TileLang stage. One kernel serves both entries: + +- ``chunk_dplr_fwd_ho`` (recompute forward) returns only ``(o, final_state)``, + so it avoids the global forward write/read of ``h`` and ``v_new``; +- ``chunk_dplr_fwd_ho_ctx`` (disable_recompute forward) additionally stores + per-chunk ``h`` and per-token ``v_new`` for the backward to consume. + +Backward recompute still uses the split H/O kernels because its gradients +consume those intermediates. +""" + +import tilelang +import tilelang.language as T +import torch + +from fla.utils import get_device_capability, get_device_smem_optin + +from .layout import ChunkLayout, build_rect_chunk_layout, build_varlen_chunk_layout +from .schedules import dtype_nbytes + + +def _ho_smem_bytes(BT: int, K: int, BV: int, in_dtype: str, num_stages: int = 1) -> int: + elem = dtype_nbytes(in_dtype) + # Keep this in sync with the shared buffers in _chunk_dplr_fwd_ho_kernel. + base = ( + elem * (K * BV + BT * K + BT * K + BT * BV + BT * BV + BT * K + BT * BT + BT * BT) + + 4 * (K * BV + K * BV + BT * BV + BT * K + K) + ) + if num_stages < 2: + return base + # Double-buffered: a second version of the seven in_dtype operand tiles + # (3x (BT, K), 2x (BT, BV), 2x (BT, BT)). The fp32 bg tile stays + # single-buffered (cp.async cannot cast during the load). + return base + elem * (3 * BT * K + 2 * BT * BV + 2 * BT * BT) + + +def _ho_fwd_num_stages(BT: int, K: int, BV: int, in_dtype: str, device_index: int) -> int: + if not torch.cuda.is_available(): + return 0 + if _ho_smem_bytes(BT, K, BV, in_dtype, 2) <= get_device_smem_optin(device_index): + return 2 + return 0 + + +def _ho_fwd_config(K: int, V: int, BT: int, in_dtype: str, device_index: int) -> dict[str, int]: + if not torch.cuda.is_available(): + return {"BV": 16, "threads": 64} + major = get_device_capability(device_index)[0] + smem_limit = get_device_smem_optin(device_index) + + def pick(candidates: list[int]) -> int: + for candidate in candidates: + if candidate <= V and _ho_smem_bytes(BT, K, candidate, in_dtype) <= smem_limit: + return candidate + return 16 + + if BT <= 16: + # 16-row GEMMs only partition across <= 2 warps (m_warp x n_warp must + # equal num_warps with >= 16 rows and >= 8 cols per warp). + return {"BV": pick([64, 32, 16]), "threads": 64} + elif major >= 9 and K <= 128: + return {"BV": pick([64, 32, 16]), "threads": 256} + elif major == 8: + return {"BV": pick([32, 16]), "threads": 128} + return {"BV": 16, "threads": 64} + + +def _ho_fragment_merge_flags( + K: int, + V: int, + BT: int, + in_dtype: str, + store_context: bool, + config: dict[str, int], + device_index: int = 0, +) -> dict[str, bool]: + capability = get_device_capability(device_index) + common_shape = K == 128 and V == 128 and BT == 32 and in_dtype == "bfloat16" + if ( + common_shape + and capability[0] == 12 + and config == {"BV": 32, "threads": 256} + ): + return {"DIRECT_KG_FRAGMENT": True, "DIRECT_BG_FRAGMENT": True} + if ( + common_shape + and capability == (9, 0) + and not store_context + and config == {"BV": 64, "threads": 256} + ): + return {"DIRECT_KG_FRAGMENT": True, "DIRECT_BG_FRAGMENT": False} + return {"DIRECT_KG_FRAGMENT": False, "DIRECT_BG_FRAGMENT": False} + + +@tilelang.jit( + pass_configs={ + tilelang.PassConfigKey.TL_ENABLE_FAST_MATH: True, + tilelang.PassConfigKey.TL_DISABLE_DATA_RACE_CHECK: False, + }, +) +def _chunk_dplr_fwd_ho_kernel( + H, K, V, BT, + in_dtype, state_dtype, + USE_INITIAL_STATE: bool, + STORE_FINAL_STATE: bool, + STORE_H_CTX: bool = False, + STORE_V_NEW_CTX: bool = False, + BV: int = 64, + threads: int = 128, + DIRECT_KG_FRAGMENT: bool = False, + DIRECT_BG_FRAGMENT: bool = False, + num_stages: int = 0, +): + acc_dtype = "float32" + # num_stages >= 2 double-buffers the seven in_dtype operand tiles and + # prefetches the next chunk with cp.async during the current chunk's + # GEMMs. Only the last chunk of a sequence can be ragged; when the + # prefetch target is ragged it is staged with predicated scalar loads + # instead (cp.async cannot mask rows by `eos`). The fp32 bg tile needs a + # cast on load, so it always takes the scalar path. + PIPELINED = num_stages >= 2 + n_bufs = 2 if PIPELINED else 1 + n_tokens, n_seq_plus_one, n_h0, n_ht, n_h_ctx, n_v_ctx = T.dynamic( + "n_tokens, n_seq_plus_one, n_h0, n_ht, n_h_ctx, n_v_ctx" + ) + n_seqs = n_seq_plus_one - 1 + + @T.macro + def load_ho_chunk( + chunk_bos, buf, eos, i_h, i_v, + qg, kg, v, w, u, A_qk, A_qb, + kg_shared, w_shared, qg_shared, v_shared, u_shared, A_qk_shared, A_qb_shared, + ): + full_tile = chunk_bos + BT <= eos + if full_tile: + # Bulk vectorized copies for full chunks. + T.copy(kg[chunk_bos: chunk_bos + BT, i_h, 0:K], kg_shared[buf, :, :]) + T.copy(w[chunk_bos: chunk_bos + BT, i_h, 0:K], w_shared[buf, :, :]) + T.copy(qg[chunk_bos: chunk_bos + BT, i_h, 0:K], qg_shared[buf, :, :]) + T.copy(v[chunk_bos: chunk_bos + BT, i_h, i_v * BV: i_v * BV + BV], v_shared[buf, :, :]) + T.copy(u[chunk_bos: chunk_bos + BT, i_h, i_v * BV: i_v * BV + BV], u_shared[buf, :, :]) + T.copy(A_qk[chunk_bos: chunk_bos + BT, i_h, 0:BT], A_qk_shared[buf, :, :]) + T.copy(A_qb[chunk_bos: chunk_bos + BT, i_h, 0:BT], A_qb_shared[buf, :, :]) + else: + for c, k_idx in T.Parallel(BT, K): + t = chunk_bos + c + if t < eos: + kg_shared[buf, c, k_idx] = kg[t, i_h, k_idx] + w_shared[buf, c, k_idx] = w[t, i_h, k_idx] + qg_shared[buf, c, k_idx] = qg[t, i_h, k_idx] + else: + kg_shared[buf, c, k_idx] = T.Cast(in_dtype, 0.0) + w_shared[buf, c, k_idx] = T.Cast(in_dtype, 0.0) + qg_shared[buf, c, k_idx] = T.Cast(in_dtype, 0.0) + + for c, vv in T.Parallel(BT, BV): + t = chunk_bos + c + g_v = i_v * BV + vv + if (t < eos) and (g_v < V): + v_shared[buf, c, vv] = v[t, i_h, g_v] + u_shared[buf, c, vv] = u[t, i_h, g_v] + else: + v_shared[buf, c, vv] = T.Cast(in_dtype, 0.0) + u_shared[buf, c, vv] = T.Cast(in_dtype, 0.0) + + for r, c in T.Parallel(BT, BT): + t = chunk_bos + r + if (t < eos) and (r >= c): + A_qk_shared[buf, r, c] = A_qk[t, i_h, c] + A_qb_shared[buf, r, c] = A_qb[t, i_h, c] + else: + A_qk_shared[buf, r, c] = T.Cast(in_dtype, 0.0) + A_qb_shared[buf, r, c] = T.Cast(in_dtype, 0.0) + + @T.macro + def prefetch_ho_chunk( + chunk_bos, buf, i_h, i_v, + qg, kg, v, w, u, A_qk, A_qb, + kg_shared, w_shared, qg_shared, v_shared, u_shared, A_qk_shared, A_qb_shared, + ): + # Full tiles only: 7 operand tiles, one commit group per async_copy. + T.async_copy(kg[chunk_bos: chunk_bos + BT, i_h, 0:K], kg_shared[buf, :, :]) + T.async_copy(w[chunk_bos: chunk_bos + BT, i_h, 0:K], w_shared[buf, :, :]) + T.async_copy(qg[chunk_bos: chunk_bos + BT, i_h, 0:K], qg_shared[buf, :, :]) + T.async_copy(v[chunk_bos: chunk_bos + BT, i_h, i_v * BV: i_v * BV + BV], v_shared[buf, :, :]) + T.async_copy(u[chunk_bos: chunk_bos + BT, i_h, i_v * BV: i_v * BV + BV], u_shared[buf, :, :]) + # Stored A matrices are already causally masked. + T.async_copy(A_qk[chunk_bos: chunk_bos + BT, i_h, 0:BT], A_qk_shared[buf, :, :]) + T.async_copy(A_qb[chunk_bos: chunk_bos + BT, i_h, 0:BT], A_qb_shared[buf, :, :]) + # Leave the 7 just issued groups pending, wait for the 7 issued for + # the current buffer at the previous iteration. + T.ptx_wait_group(7) + + @T.prim_func + def chunk_dplr_fwd_ho_tl( + qg: T.Tensor((n_tokens, H, K), in_dtype), + kg: T.Tensor((n_tokens, H, K), in_dtype), + v: T.Tensor((n_tokens, H, V), in_dtype), + w: T.Tensor((n_tokens, H, K), in_dtype), + u: T.Tensor((n_tokens, H, V), in_dtype), + bg: T.Tensor((n_tokens, H, K), in_dtype), + gk: T.Tensor((n_tokens, H, K), acc_dtype), + A_qk: T.Tensor((n_tokens, H, BT), in_dtype), + A_qb: T.Tensor((n_tokens, H, BT), in_dtype), + h0: T.Tensor((n_h0, H, K, V), acc_dtype), + cu_seqlens: T.Tensor((n_seq_plus_one,), "int32"), + chunk_offsets: T.Tensor((n_seq_plus_one,), "int32"), + o: T.Tensor((n_tokens, H, V), in_dtype), + ht: T.Tensor((n_ht, H, K, V), state_dtype), + h_ctx: T.Tensor((n_h_ctx, H, K, V), in_dtype), + v_new_ctx: T.Tensor((n_v_ctx, H, V), in_dtype), + ): + with T.Kernel(T.ceildiv(V, BV), n_seqs, H, threads=threads) as (i_v, i_n, i_h): + bos = cu_seqlens[i_n] + eos = cu_seqlens[i_n + 1] + boh = chunk_offsets[i_n] + n_chunks = chunk_offsets[i_n + 1] - boh + + b_h = T.alloc_fragment((K, BV), acc_dtype) + b_h_shared = T.alloc_shared((K, BV), in_dtype) + b_hc = T.alloc_fragment((K, BV), acc_dtype) + b_hc_kg = T.alloc_fragment((K, BV), acc_dtype) + b_hc_bg = T.alloc_fragment((K, BV), acc_dtype) + if not DIRECT_KG_FRAGMENT: + b_hc_kg_shared = T.alloc_shared((K, BV), acc_dtype) + if not DIRECT_BG_FRAGMENT: + b_hc_bg_shared = T.alloc_shared((K, BV), acc_dtype) + v2_frag = T.alloc_fragment((BT, BV), acc_dtype) + v2_acc_shared = T.alloc_shared((BT, BV), acc_dtype) + v2_shared = T.alloc_shared((BT, BV), in_dtype) + kg_shared = T.alloc_shared((n_bufs, BT, K), in_dtype) + bg_shared = T.alloc_shared((BT, K), acc_dtype) + w_shared = T.alloc_shared((n_bufs, BT, K), in_dtype) + v_shared = T.alloc_shared((n_bufs, BT, BV), in_dtype) + u_shared = T.alloc_shared((n_bufs, BT, BV), in_dtype) + qg_shared = T.alloc_shared((n_bufs, BT, K), in_dtype) + A_qk_shared = T.alloc_shared((n_bufs, BT, BT), in_dtype) + A_qb_shared = T.alloc_shared((n_bufs, BT, BT), in_dtype) + o_frag = T.alloc_fragment((BT, BV), acc_dtype) + gk_last_shared = T.alloc_shared((K,), acc_dtype) + + if USE_INITIAL_STATE: + for k_idx, vv in T.Parallel(K, BV): + g_v = i_v * BV + vv + if g_v < V: + b_h[k_idx, vv] = h0[i_n, i_h, k_idx, g_v] + else: + b_h[k_idx, vv] = 0.0 + else: + for k_idx, vv in T.Parallel(K, BV): + b_h[k_idx, vv] = 0.0 + + if PIPELINED: + load_ho_chunk(bos, 0, eos, i_h, i_v, qg, kg, v, w, u, A_qk, A_qb, + kg_shared, w_shared, qg_shared, v_shared, u_shared, + A_qk_shared, A_qb_shared) + + for i_t in T.serial(n_chunks): + chunk_bos = bos + i_t * BT + chunk_row = boh + i_t + cur = i_t % 2 if PIPELINED else 0 + T.copy(b_h, b_h_shared) + if STORE_H_CTX: + for k_idx, vv in T.Parallel(K, BV): + g_v = i_v * BV + vv + if g_v < V: + h_ctx[chunk_row, i_h, k_idx, g_v] = T.Cast( + in_dtype, b_h[k_idx, vv] + ) + T.clear(b_hc) + + if PIPELINED: + # The current buffer must be fully consumed before the + # cp.async below overwrites the idle one. + T.sync_threads() + if i_t + 1 < n_chunks: + nxt = (i_t + 1) % 2 + nxt_bos = chunk_bos + BT + if nxt_bos + BT <= eos: + prefetch_ho_chunk(nxt_bos, nxt, i_h, i_v, qg, kg, v, w, u, A_qk, A_qb, + kg_shared, w_shared, qg_shared, v_shared, u_shared, + A_qk_shared, A_qb_shared) + else: + load_ho_chunk(nxt_bos, nxt, eos, i_h, i_v, qg, kg, v, w, u, A_qk, A_qb, + kg_shared, w_shared, qg_shared, v_shared, u_shared, + A_qk_shared, A_qb_shared) + T.ptx_wait_group(0) + else: + T.ptx_wait_group(0) + else: + load_ho_chunk(chunk_bos, 0, eos, i_h, i_v, qg, kg, v, w, u, A_qk, A_qb, + kg_shared, w_shared, qg_shared, v_shared, u_shared, + A_qk_shared, A_qb_shared) + + for c, k_idx in T.Parallel(BT, K): + t = chunk_bos + c + if t < eos: + bg_shared[c, k_idx] = T.Cast(acc_dtype, bg[t, i_h, k_idx]) + else: + bg_shared[c, k_idx] = 0.0 + + T.gemm(w_shared[cur, :, :], b_h_shared, v2_frag, clear_accum=True) + for c, vv in T.Parallel(BT, BV): + t = chunk_bos + c + g_v = i_v * BV + vv + v2 = T.ieee_add(v2_frag[c, vv], T.Cast(acc_dtype, u_shared[cur, c, vv])) + if (t < eos) and (g_v < V): + v2_acc_shared[c, vv] = v2 + v2_shared[c, vv] = T.Cast(in_dtype, v2) + if STORE_V_NEW_CTX: + v_new_ctx[t, i_h, g_v] = v2_shared[c, vv] + else: + v2_acc_shared[c, vv] = T.Cast(acc_dtype, 0.0) + v2_shared[c, vv] = T.Cast(in_dtype, 0.0) + + T.gemm( + kg_shared[cur, :, :], + v_shared[cur, :, :], + b_hc_kg, + transpose_A=True, + clear_accum=True, + ) + if not DIRECT_KG_FRAGMENT: + T.copy(b_hc_kg, b_hc_kg_shared) + T.gemm( + bg_shared, + v2_acc_shared, + b_hc_bg, + transpose_A=True, + clear_accum=True, + ) + if not DIRECT_BG_FRAGMENT: + T.copy(b_hc_bg, b_hc_bg_shared) + if DIRECT_KG_FRAGMENT and DIRECT_BG_FRAGMENT: + for k_idx, vv in T.Parallel(K, BV): + b_hc[k_idx, vv] = T.ieee_add( + b_hc_kg[k_idx, vv], + b_hc_bg[k_idx, vv], + ) + elif DIRECT_KG_FRAGMENT: + for k_idx, vv in T.Parallel(K, BV): + b_hc[k_idx, vv] = T.ieee_add( + b_hc_kg[k_idx, vv], + b_hc_bg_shared[k_idx, vv], + ) + elif DIRECT_BG_FRAGMENT: + for k_idx, vv in T.Parallel(K, BV): + b_hc[k_idx, vv] = T.ieee_add( + b_hc_kg_shared[k_idx, vv], + b_hc_bg[k_idx, vv], + ) + else: + for k_idx, vv in T.Parallel(K, BV): + b_hc[k_idx, vv] = T.ieee_add( + b_hc_kg_shared[k_idx, vv], + b_hc_bg_shared[k_idx, vv], + ) + + T.gemm(qg_shared[cur, :, :], b_h_shared, o_frag, clear_accum=True) + + T.gemm(A_qk_shared[cur, :, :], v_shared[cur, :, :], o_frag) + T.gemm(A_qb_shared[cur, :, :], v2_shared, o_frag) + + # Stage o through v2_shared (dead after its GEMM): accumulator + # fragments are thread-strided, so a direct store lowers to + # 2-byte writes; staging enables vectorized ones. + T.copy(o_frag, v2_shared) + if (chunk_bos + BT <= eos) and (i_v * BV + BV <= V): + T.copy(v2_shared, o[chunk_bos: chunk_bos + BT, i_h, i_v * BV: i_v * BV + BV]) + else: + for c, vv in T.Parallel(BT, BV): + t = chunk_bos + c + g_v = i_v * BV + vv + if (t < eos) and (g_v < V): + o[t, i_h, g_v] = v2_shared[c, vv] + + last_idx = T.min(chunk_bos + BT - 1, eos - 1) + for k_idx in T.Parallel(K): + gk_last_shared[k_idx] = gk[last_idx, i_h, k_idx] + for k_idx, vv in T.Parallel(K, BV): + b_h[k_idx, vv] = T.ieee_mul(T.exp2(gk_last_shared[k_idx]), b_h[k_idx, vv]) + for k_idx, vv in T.Parallel(K, BV): + b_h[k_idx, vv] = T.ieee_add(b_h[k_idx, vv], b_hc[k_idx, vv]) + + if STORE_FINAL_STATE: + for k_idx, vv in T.Parallel(K, BV): + g_v = i_v * BV + vv + if g_v < V: + ht[i_n, i_h, k_idx, g_v] = T.Cast(state_dtype, b_h[k_idx, vv]) + + return chunk_dplr_fwd_ho_tl + + +def _prepare_ho_inputs( + qg, kg, v, w, u, bg, gk, A_qk, A_qb, + initial_state, cu_seqlens, chunk_size, chunk_layout, +): + B, T_, H, K = kg.shape + V = v.shape[-1] + + is_varlen = cu_seqlens is not None + if is_varlen: + assert B == 1 + layout = chunk_layout if chunk_layout is not None else build_varlen_chunk_layout(cu_seqlens, chunk_size, T_) + else: + layout = build_rect_chunk_layout(B, T_, chunk_size, kg.device) + token_rows = B * T_ + in_dtype = str(kg.dtype).split(".")[-1] + use_h0 = initial_state is not None + if use_h0: + h0 = initial_state + else: + h0 = torch.empty((1, H, K, V), dtype=torch.float32, device=kg.device) + if h0.dtype != torch.float32: + h0 = h0.to(torch.float32) + + flats = ( + qg.reshape(token_rows, H, K).contiguous(), + kg.reshape(token_rows, H, K).contiguous(), + v.reshape(token_rows, H, V).contiguous(), + w.reshape(token_rows, H, K).contiguous(), + u.reshape(token_rows, H, V).contiguous(), + bg.reshape(token_rows, H, K).contiguous(), + gk.reshape(token_rows, H, K).contiguous(), + A_qk.reshape(token_rows, H, chunk_size).contiguous(), + A_qb.reshape(token_rows, H, chunk_size).contiguous(), + ) + return layout, h0, in_dtype, flats + + +def chunk_dplr_fwd_ho( + qg: torch.Tensor, + kg: torch.Tensor, + v: torch.Tensor, + w: torch.Tensor, + u: torch.Tensor, + bg: torch.Tensor, + gk: torch.Tensor, + A_qk: torch.Tensor, + A_qb: torch.Tensor, + initial_state: torch.Tensor | None = None, + output_final_state: bool = False, + cu_seqlens: torch.Tensor | None = None, + chunk_size: int = 64, + chunk_layout: ChunkLayout | None = None, +) -> tuple[torch.Tensor, torch.Tensor | None]: + B, T_, H, K = kg.shape + V = v.shape[-1] + layout, h0, in_dtype, flats = _prepare_ho_inputs( + qg, kg, v, w, u, bg, gk, A_qk, A_qb, + initial_state, cu_seqlens, chunk_size, chunk_layout, + ) + qg_f, kg_f, v_f, w_f, u_f, bg_f, gk_f, A_qk_f, A_qb_f = flats + + config = _ho_fwd_config(K, V, chunk_size, in_dtype, kg.device.index) + merge_flags = _ho_fragment_merge_flags( + K, V, chunk_size, in_dtype, False, config, + device_index=kg.device.index, + ) + config["num_stages"] = _ho_fwd_num_stages(chunk_size, K, config["BV"], in_dtype, kg.device.index) + kernel = _chunk_dplr_fwd_ho_kernel( + H, K, V, chunk_size, + in_dtype, "float32", initial_state is not None, output_final_state, + **merge_flags, + **config, + ) + o_f = torch.empty((B * T_, H, V), dtype=v.dtype, device=v.device) + ht = torch.empty((layout.cu_seqlens.shape[0] - 1 if output_final_state else 1, H, K, V), + dtype=torch.float32, device=kg.device) + # ctx outputs are never stored here; 1-row dummies only carry the ABI + h_ctx_dummy = torch.empty((1, H, K, V), dtype=kg.dtype, device=kg.device) + v_new_ctx_dummy = torch.empty((1, H, V), dtype=v.dtype, device=v.device) + kernel( + qg_f, kg_f, v_f, w_f, u_f, bg_f, gk_f, A_qk_f, A_qb_f, + h0, layout.cu_seqlens, layout.chunk_offsets, o_f, ht, + h_ctx_dummy, v_new_ctx_dummy, + ) + final_state = ht if output_final_state else None + return o_f.view(B, T_, H, V), final_state + + +def chunk_dplr_fwd_ho_ctx( + qg: torch.Tensor, + kg: torch.Tensor, + v: torch.Tensor, + w: torch.Tensor, + u: torch.Tensor, + bg: torch.Tensor, + gk: torch.Tensor, + A_qk: torch.Tensor, + A_qb: torch.Tensor, + initial_state: torch.Tensor | None = None, + output_final_state: bool = False, + cu_seqlens: torch.Tensor | None = None, + chunk_size: int = 64, + chunk_layout: ChunkLayout | None = None, + store_context: bool = True, +) -> tuple[torch.Tensor, torch.Tensor | None, torch.Tensor, torch.Tensor]: + B, T_, H, K = kg.shape + V = v.shape[-1] + layout, h0, in_dtype, flats = _prepare_ho_inputs( + qg, kg, v, w, u, bg, gk, A_qk, A_qb, + initial_state, cu_seqlens, chunk_size, chunk_layout, + ) + qg_f, kg_f, v_f, w_f, u_f, bg_f, gk_f, A_qk_f, A_qb_f = flats + chunk_rows = layout.chunk_indices.shape[0] + token_rows = B * T_ + + config = _ho_fwd_config(K, V, chunk_size, in_dtype, kg.device.index) + merge_flags = _ho_fragment_merge_flags( + K, V, chunk_size, in_dtype, store_context, config, + device_index=kg.device.index, + ) + config["num_stages"] = _ho_fwd_num_stages(chunk_size, K, config["BV"], in_dtype, kg.device.index) + kernel = _chunk_dplr_fwd_ho_kernel( + H, K, V, chunk_size, + in_dtype, "float32", initial_state is not None, output_final_state, + STORE_H_CTX=store_context, + STORE_V_NEW_CTX=store_context, + **merge_flags, + **config, + ) + o_f = torch.empty((token_rows, H, V), dtype=v.dtype, device=v.device) + ht = torch.empty((layout.cu_seqlens.shape[0] - 1 if output_final_state else 1, H, K, V), + dtype=torch.float32, device=kg.device) + if store_context: + h_ctx = torch.empty((chunk_rows, H, K, V), dtype=kg.dtype, device=kg.device) + v_new_ctx = torch.empty((token_rows, H, V), dtype=v.dtype, device=v.device) + h_ctx_arg = h_ctx + v_new_ctx_arg = v_new_ctx + else: + h_ctx_arg = torch.empty((1, H, K, V), dtype=kg.dtype, device=kg.device) + v_new_ctx_arg = torch.empty((1, H, V), dtype=v.dtype, device=v.device) + h_ctx = kg.new_empty((1,)).expand(chunk_rows, H, K, V) + v_new_ctx = v.new_empty((1,)).expand(token_rows, H, V) + kernel( + qg_f, kg_f, v_f, w_f, u_f, bg_f, gk_f, A_qk_f, A_qb_f, + h0, layout.cu_seqlens, layout.chunk_offsets, o_f, ht, h_ctx_arg, v_new_ctx_arg, + ) + final_state = ht if output_final_state else None + return o_f.view(B, T_, H, V), final_state, h_ctx, v_new_ctx.view(B, T_, H, V) diff --git a/fla/ops/generalized_delta_rule/dplr/backends/tilelang/chunk_stream_bwd.py b/fla/ops/generalized_delta_rule/dplr/backends/tilelang/chunk_stream_bwd.py new file mode 100644 index 0000000000..315f105a48 --- /dev/null +++ b/fla/ops/generalized_delta_rule/dplr/backends/tilelang/chunk_stream_bwd.py @@ -0,0 +1,913 @@ +# Copyright (c) 2023-2026, Songlin Yang, Yu Zhang, Zhiyuan Li +# +# This source code is licensed under the MIT license found in the +# LICENSE file in the root directory of this source tree. +# For a list of all contributors, visit: +# https://github.com/fla-org/flash-linear-attention/graphs/contributors + +"""Streaming DPLR backward. + +The reverse `dhu` scan is fused with the q/o-side backward consumer so the +per-chunk `dh` state stays inside one sequence/head program instead of being +materialized as a global `(n_chunks, H, K, V)` tensor. +""" + +import tilelang +import tilelang.language as T +import torch + +from fla.utils import get_device_capability, get_device_smem_optin + +from .layout import ChunkLayout, build_rect_chunk_layout, build_varlen_chunk_layout +from .schedules import ( + stream_bwd_num_stages, + stream_bwd_schedule_or_none, + stream_default_threads, + stream_high_smem_bytes, + stream_low_bwd_config, + stream_low_smem_bytes, + stream_mid_smem_bytes, +) + + +def _stream_bwd_config(BT: int, cc: int) -> dict[str, int]: + # Micro-autotune on H800 (cc90) shows the high-SMEM schedule is faster with + # 256 threads for all BT>=32 training shapes, not just K=V=128. + if cc >= 90 and BT >= 32: + threads = 256 + else: + threads = stream_default_threads(BT) + return {"threads": threads} + + +def _select_stream_bwd_schedule( + *, + K: int, + V: int, + BT: int, + in_dtype: str, + device: torch.device, +) -> tuple[str, dict[str, int]]: + index = device.index or 0 + smem_cap = get_device_smem_optin(index) + major, minor = get_device_capability(index) + cc = major * 10 + minor + selected = stream_bwd_schedule_or_none(K=K, V=V, BT=BT, in_dtype=in_dtype, smem_cap=smem_cap) + if selected == "high": + config = _stream_bwd_config(BT, cc) + config["num_stages"] = stream_bwd_num_stages( + selected, K=K, V=V, BT=BT, in_dtype=in_dtype, smem_cap=smem_cap, + ) + return "high", config + if selected == "mid": + config = _stream_bwd_config(BT, cc) + config["num_stages"] = stream_bwd_num_stages( + selected, K=K, V=V, BT=BT, in_dtype=in_dtype, smem_cap=smem_cap, + ) + return "mid", config + if selected == "low": + return "low", stream_low_bwd_config(BT, V) + name = torch.cuda.get_device_properties(index).name + raise RuntimeError( + f"No launchable DPLR stream backward schedule for K={K}, V={V}, BT={BT}, " + f"dtype={in_dtype} on {name} cc{cc}: high={stream_high_smem_bytes(K, V, BT, in_dtype)}B, " + f"mid={stream_mid_smem_bytes(K, V, BT, in_dtype)}B, " + f"low={stream_low_smem_bytes(K, V, BT, in_dtype)}B, device cap={smem_cap}B" + ) + + +@tilelang.jit( + pass_configs={ + tilelang.PassConfigKey.TL_ENABLE_FAST_MATH: True, + tilelang.PassConfigKey.TL_DISABLE_DATA_RACE_CHECK: False, + }, +) +def _chunk_dplr_bwd_stream_kernel( + H, K, V, BT, + in_dtype, state_dtype, + USE_FINAL_STATE_GRADIENT: bool, + USE_INITIAL_STATE: bool, + threads: int = 128, + alias_kv: bool = False, + num_stages: int = 0, +): + acc_dtype = "float32" + # num_stages >= 2 double-buffers the nine operand tiles: the next chunk is + # prefetched with cp.async while the current one is consumed. cp.async + # cannot mask rows by `eos`, so the last (only possibly ragged) chunk of + # each sequence is staged with predicated scalar loads before the loop and + # the loop itself only ever prefetches full interior chunks. + PIPELINED = num_stages >= 2 + n_bufs = 2 if PIPELINED else 1 + n_tokens, n_seq_plus_one, n_chunks, n_dht, n_dh0 = T.dynamic( + "n_tokens, n_seq_plus_one, n_chunks, n_dht, n_dh0" + ) + n_seqs = n_seq_plus_one - 1 + + @T.prim_func + def chunk_dplr_bwd_stream_tl( + qg: T.Tensor((n_tokens, H, K), in_dtype), + bg: T.Tensor((n_tokens, H, K), in_dtype), + w: T.Tensor((n_tokens, H, K), in_dtype), + kg: T.Tensor((n_tokens, H, K), in_dtype), + v: T.Tensor((n_tokens, H, V), in_dtype), + v_new: T.Tensor((n_tokens, H, V), in_dtype), + gk: T.Tensor((n_tokens, H, K), acc_dtype), + do: T.Tensor((n_tokens, H, V), in_dtype), + h: T.Tensor((n_chunks, H, K, V), in_dtype), + A_qb: T.Tensor((n_tokens, H, BT), in_dtype), + A_qk: T.Tensor((n_tokens, H, BT), in_dtype), + dht: T.Tensor((n_dht, H, K, V), state_dtype), + cu_seqlens: T.Tensor((n_seq_plus_one,), "int32"), + chunk_offsets: T.Tensor((n_seq_plus_one,), "int32"), + dq_out: T.Tensor((n_tokens, H, K), in_dtype), + dk_out: T.Tensor((n_tokens, H, K), in_dtype), + dw_out: T.Tensor((n_tokens, H, K), in_dtype), + db_out: T.Tensor((n_tokens, H, K), in_dtype), + dgk_last: T.Tensor((n_chunks, H, K), acc_dtype), + dv2: T.Tensor((n_tokens, H, V), in_dtype), + dv_full: T.Tensor((n_tokens, H, V), in_dtype), + dh0: T.Tensor((n_dh0, H, K, V), state_dtype), + ): + with T.Kernel(n_seqs, H, threads=threads) as (i_n, i_h): + bos = cu_seqlens[i_n] + eos = cu_seqlens[i_n + 1] + boh = chunk_offsets[i_n] + n_chunks = chunk_offsets[i_n + 1] - boh + + b_dh = T.alloc_fragment((K, V), acc_dtype) + b_dh_tmp = T.alloc_fragment((K, V), acc_dtype) + if alias_kv: + # One (K, V) tile holds dh first (dv2/dk/db/dv_full GEMMs), + # then h (dgk_h/dq/dw GEMMs); one (BT, K) tile stages the four + # output rows sequentially. Fits the 227KB cc90 cap at + # K=V=128, BT=64 where the high schedule needs 291KB. + kv_shared = T.alloc_shared((K, V), in_dtype) + out_shared = T.alloc_shared((BT, K), in_dtype) + else: + b_dh_shared = T.alloc_shared((K, V), in_dtype) + # h rows are defined for every chunk (no ragged edge), so the + # (K, V) h tile joins the cp.async prefetch set under + # pipelining instead of stalling each iteration on a + # synchronous global->shared copy. + h_shared = T.alloc_shared((n_bufs, K, V), in_dtype) + + qg_shared = T.alloc_shared((n_bufs, BT, K), in_dtype) + bg_shared = T.alloc_shared((n_bufs, BT, K), in_dtype) + w_shared = T.alloc_shared((n_bufs, BT, K), in_dtype) + kg_shared = T.alloc_shared((n_bufs, BT, K), in_dtype) + v_shared = T.alloc_shared((n_bufs, BT, V), in_dtype) + v_new_shared = T.alloc_shared((n_bufs, BT, V), in_dtype) + do_shared = T.alloc_shared((n_bufs, BT, V), in_dtype) + A_qb_shared = T.alloc_shared((n_bufs, BT, BT), in_dtype) + A_qk_shared = T.alloc_shared((n_bufs, BT, BT), in_dtype) + + dv_intra_frag = T.alloc_fragment((BT, V), acc_dtype) + dv2_frag = T.alloc_fragment((BT, V), acc_dtype) + dv2_shared = T.alloc_shared((BT, V), in_dtype) + dv_full_frag = T.alloc_fragment((BT, V), acc_dtype) + dv_full_shared = T.alloc_shared((BT, V), in_dtype) + + dq_frag = T.alloc_fragment((BT, K), acc_dtype) + dk_frag = T.alloc_fragment((BT, K), acc_dtype) + dw_frag = T.alloc_fragment((BT, K), acc_dtype) + db_frag = T.alloc_fragment((BT, K), acc_dtype) + if not alias_kv: + dq_shared = T.alloc_shared((BT, K), in_dtype) + dw_shared = T.alloc_shared((BT, K), in_dtype) + dk_shared = T.alloc_shared((BT, K), in_dtype) + db_shared = T.alloc_shared((BT, K), in_dtype) + dgk_last_frag = T.alloc_fragment((K,), acc_dtype) + hdh_frag = T.alloc_fragment((K, V), acc_dtype) + dgk_part = T.alloc_fragment((4, K), acc_dtype) + dgk_part_shared = T.alloc_shared((4, K), acc_dtype) + gk_last_shared = T.alloc_shared((K,), acc_dtype) + + if USE_FINAL_STATE_GRADIENT: + for k_idx, vv in T.Parallel(K, V): + b_dh[k_idx, vv] = dht[i_n, i_h, k_idx, vv] + else: + for k_idx, vv in T.Parallel(K, V): + b_dh[k_idx, vv] = T.float32(0.0) + + if PIPELINED: + t_pro = bos + (n_chunks - 1) * BT + for r, c in T.Parallel(BT, K): + t = t_pro + r + if t < eos: + qg_shared[0, r, c] = qg[t, i_h, c] + bg_shared[0, r, c] = bg[t, i_h, c] + w_shared[0, r, c] = w[t, i_h, c] + kg_shared[0, r, c] = kg[t, i_h, c] + else: + qg_shared[0, r, c] = T.Cast(in_dtype, 0.0) + bg_shared[0, r, c] = T.Cast(in_dtype, 0.0) + w_shared[0, r, c] = T.Cast(in_dtype, 0.0) + kg_shared[0, r, c] = T.Cast(in_dtype, 0.0) + + for r, c in T.Parallel(BT, V): + t = t_pro + r + if t < eos: + v_shared[0, r, c] = v[t, i_h, c] + v_new_shared[0, r, c] = v_new[t, i_h, c] + do_shared[0, r, c] = do[t, i_h, c] + else: + v_shared[0, r, c] = T.Cast(in_dtype, 0.0) + v_new_shared[0, r, c] = T.Cast(in_dtype, 0.0) + do_shared[0, r, c] = T.Cast(in_dtype, 0.0) + + for r, c in T.Parallel(BT, BT): + t = t_pro + r + if (t < eos) and (r >= c): + A_qb_shared[0, r, c] = A_qb[t, i_h, c] + A_qk_shared[0, r, c] = A_qk[t, i_h, c] + else: + A_qb_shared[0, r, c] = T.Cast(in_dtype, 0.0) + A_qk_shared[0, r, c] = T.Cast(in_dtype, 0.0) + + if not alias_kv: + T.copy(h[boh + n_chunks - 1, i_h, 0:K, 0:V], h_shared[0, :, :]) + + for i_t_rev in T.serial(n_chunks): + i_t = n_chunks - 1 - i_t_rev + t_off = bos + i_t * BT + chunk_row = boh + i_t + cur = i_t_rev % 2 if PIPELINED else 0 + + T.clear(dq_frag) + T.clear(dk_frag) + T.clear(dw_frag) + T.clear(db_frag) + T.clear(b_dh_tmp) + for k_idx in T.Parallel(K): + dgk_last_frag[k_idx] = T.float32(0.0) + for gg, c in T.Parallel(4, K): + dgk_part[gg, c] = T.float32(0.0) + + if PIPELINED: + # The current buffer must be fully consumed before the + # cp.async below overwrites the idle one. + T.sync_threads() + if i_t_rev + 1 < n_chunks: + nxt = (i_t_rev + 1) % 2 + t_nxt = t_off - BT + T.async_copy(qg[t_nxt: t_nxt + BT, i_h, 0:K], qg_shared[nxt, :, :]) + T.async_copy(bg[t_nxt: t_nxt + BT, i_h, 0:K], bg_shared[nxt, :, :]) + T.async_copy(w[t_nxt: t_nxt + BT, i_h, 0:K], w_shared[nxt, :, :]) + T.async_copy(kg[t_nxt: t_nxt + BT, i_h, 0:K], kg_shared[nxt, :, :]) + T.async_copy(v[t_nxt: t_nxt + BT, i_h, 0:V], v_shared[nxt, :, :]) + T.async_copy(v_new[t_nxt: t_nxt + BT, i_h, 0:V], v_new_shared[nxt, :, :]) + T.async_copy(do[t_nxt: t_nxt + BT, i_h, 0:V], do_shared[nxt, :, :]) + # Stored A matrices are already causally masked. + T.async_copy(A_qb[t_nxt: t_nxt + BT, i_h, 0:BT], A_qb_shared[nxt, :, :]) + T.async_copy(A_qk[t_nxt: t_nxt + BT, i_h, 0:BT], A_qk_shared[nxt, :, :]) + if not alias_kv: + T.async_copy(h[chunk_row - 1, i_h, 0:K, 0:V], h_shared[nxt, :, :]) + # One commit group per async_copy: leave the just + # issued groups pending, wait for the ones from the + # previous iteration (the current buffer). + T.ptx_wait_group(9 if alias_kv else 10) + else: + T.ptx_wait_group(0) + else: + full_tile = t_off + BT <= eos + if full_tile: + # Bulk vectorized copies for interior chunks (TIRx + # showed the scalar predicated loads cap at ~1.5TB/s). + T.copy(qg[t_off: t_off + BT, i_h, 0:K], qg_shared[0, :, :]) + T.copy(bg[t_off: t_off + BT, i_h, 0:K], bg_shared[0, :, :]) + T.copy(w[t_off: t_off + BT, i_h, 0:K], w_shared[0, :, :]) + T.copy(kg[t_off: t_off + BT, i_h, 0:K], kg_shared[0, :, :]) + T.copy(v[t_off: t_off + BT, i_h, 0:V], v_shared[0, :, :]) + T.copy(v_new[t_off: t_off + BT, i_h, 0:V], v_new_shared[0, :, :]) + T.copy(do[t_off: t_off + BT, i_h, 0:V], do_shared[0, :, :]) + # Stored A matrices are already causally masked. + T.copy(A_qb[t_off: t_off + BT, i_h, 0:BT], A_qb_shared[0, :, :]) + T.copy(A_qk[t_off: t_off + BT, i_h, 0:BT], A_qk_shared[0, :, :]) + else: + for r, c in T.Parallel(BT, K): + t = t_off + r + if t < eos: + qg_shared[0, r, c] = qg[t, i_h, c] + bg_shared[0, r, c] = bg[t, i_h, c] + w_shared[0, r, c] = w[t, i_h, c] + kg_shared[0, r, c] = kg[t, i_h, c] + else: + qg_shared[0, r, c] = T.Cast(in_dtype, 0.0) + bg_shared[0, r, c] = T.Cast(in_dtype, 0.0) + w_shared[0, r, c] = T.Cast(in_dtype, 0.0) + kg_shared[0, r, c] = T.Cast(in_dtype, 0.0) + + for r, c in T.Parallel(BT, V): + t = t_off + r + if t < eos: + v_shared[0, r, c] = v[t, i_h, c] + v_new_shared[0, r, c] = v_new[t, i_h, c] + do_shared[0, r, c] = do[t, i_h, c] + else: + v_shared[0, r, c] = T.Cast(in_dtype, 0.0) + v_new_shared[0, r, c] = T.Cast(in_dtype, 0.0) + do_shared[0, r, c] = T.Cast(in_dtype, 0.0) + + for r, c in T.Parallel(BT, BT): + t = t_off + r + if (t < eos) and (r >= c): + A_qb_shared[0, r, c] = A_qb[t, i_h, c] + A_qk_shared[0, r, c] = A_qk[t, i_h, c] + else: + A_qb_shared[0, r, c] = T.Cast(in_dtype, 0.0) + A_qk_shared[0, r, c] = T.Cast(in_dtype, 0.0) + + if alias_kv: + T.copy(b_dh, kv_shared) + else: + if not PIPELINED: + T.copy(h[chunk_row, i_h, 0:K, 0:V], h_shared[0, :, :]) + T.copy(b_dh, b_dh_shared) + + # dv2 = A_qb^T @ do + bg @ dh + T.gemm(A_qb_shared[cur, :, :], do_shared[cur, :, :], + dv_intra_frag, transpose_A=True, clear_accum=True) + T.gemm(bg_shared[cur, :, :], kv_shared if alias_kv else b_dh_shared, + dv2_frag, clear_accum=True) + for r, vv in T.Parallel(BT, V): + t = t_off + r + dv2_frag[r, vv] = dv2_frag[r, vv] + dv_intra_frag[r, vv] + if t < eos: + dv2[t, i_h, vv] = T.Cast(in_dtype, dv2_frag[r, vv]) + T.copy(dv2_frag, dv2_shared) + + if alias_kv: + # dh-side reductions/GEMMs while kv_shared holds dh. The + # dgk_h term reads h straight from global memory; fragment + # reads in a cross-layout reduction miscompile (wrong + # thread map), so dh must come from the shared tile. The + # product goes through a dedicated fragment so the loads + # vectorize and the row reduction stays lowerable. + for k_idx, vv in T.Parallel(K, V): + hdh_frag[k_idx, vv] = ( + T.Cast(acc_dtype, h[chunk_row, i_h, k_idx, vv]) + * T.Cast(acc_dtype, kv_shared[k_idx, vv]) + ) + T.reduce_sum(hdh_frag, dgk_last_frag, dim=1, clear=False) + T.gemm(v_shared[cur, :, :], kv_shared, dk_frag, transpose_B=True) + T.gemm(v_new_shared[cur, :, :], kv_shared, db_frag, transpose_B=True) + T.gemm(kg_shared[cur, :, :], kv_shared, dv_full_frag, clear_accum=True) + T.gemm(A_qk_shared[cur, :, :], do_shared[cur, :, :], + dv_full_frag, transpose_A=True) + T.copy(dv_full_frag, dv_full_shared) + for r, vv in T.Parallel(BT, V): + t = t_off + r + if t < eos: + dv_full[t, i_h, vv] = T.Cast(in_dtype, dv_full_shared[r, vv]) + + # Stage each output row through the single (BT, K) tile, + # folding the dgk reduction into the dk/db passes. + T.copy(dk_frag, out_shared) + for r, c in T.Parallel(BT, K): + t = t_off + r + if t < eos: + dk_out[t, i_h, c] = out_shared[r, c] + for r_local in T.serial(BT // 4): + for gg, c in T.Parallel(4, K): + r = gg * (BT // 4) + r_local + dgk_part[gg, c] = ( + dgk_part[gg, c] + + T.Cast(acc_dtype, kg_shared[cur, r, c]) + * T.Cast(acc_dtype, out_shared[r, c]) + ) + T.copy(db_frag, out_shared) + for r, c in T.Parallel(BT, K): + t = t_off + r + if t < eos: + db_out[t, i_h, c] = out_shared[r, c] + for r_local in T.serial(BT // 4): + for gg, c in T.Parallel(4, K): + r = gg * (BT // 4) + r_local + dgk_part[gg, c] = ( + dgk_part[gg, c] + + T.Cast(acc_dtype, bg_shared[cur, r, c]) + * T.Cast(acc_dtype, out_shared[r, c]) + ) + + # h-side GEMMs after overwriting the state tile with h. + T.copy(h[chunk_row, i_h, 0:K, 0:V], kv_shared) + T.gemm(do_shared[cur, :, :], kv_shared, dq_frag, transpose_B=True) + T.gemm(dv2_shared, kv_shared, dw_frag, transpose_B=True) + T.copy(dq_frag, out_shared) + for r, c in T.Parallel(BT, K): + t = t_off + r + if t < eos: + dq_out[t, i_h, c] = out_shared[r, c] + T.copy(dw_frag, out_shared) + for r, c in T.Parallel(BT, K): + t = t_off + r + if t < eos: + dw_out[t, i_h, c] = out_shared[r, c] + else: + # q/o-side consumer of current dh. The h*dh product goes + # through a dedicated fragment so the smem loads vectorize + # and the row reduction stays lowerable. + for k_idx, vv in T.Parallel(K, V): + hdh_frag[k_idx, vv] = ( + T.Cast(acc_dtype, h_shared[cur, k_idx, vv]) + * T.Cast(acc_dtype, b_dh_shared[k_idx, vv]) + ) + T.reduce_sum(hdh_frag, dgk_last_frag, dim=1, clear=False) + T.gemm(do_shared[cur, :, :], h_shared[cur, :, :], dq_frag, transpose_B=True) + T.gemm(v_shared[cur, :, :], b_dh_shared, dk_frag, transpose_B=True) + T.gemm(v_new_shared[cur, :, :], b_dh_shared, db_frag, transpose_B=True) + T.gemm(dv2_shared, h_shared[cur, :, :], dw_frag, transpose_B=True) + + T.gemm(kg_shared[cur, :, :], b_dh_shared, dv_full_frag, clear_accum=True) + T.gemm(A_qk_shared[cur, :, :], do_shared[cur, :, :], + dv_full_frag, transpose_A=True) + T.copy(dv_full_frag, dv_full_shared) + for r, vv in T.Parallel(BT, V): + t = t_off + r + if t < eos: + dv_full[t, i_h, vv] = T.Cast(in_dtype, dv_full_shared[r, vv]) + + T.copy(dk_frag, dk_shared) + T.copy(db_frag, db_shared) + T.copy(dq_frag, dq_shared) + T.copy(dw_frag, dw_shared) + for r, c in T.Parallel(BT, K): + t = t_off + r + if t < eos: + dq_out[t, i_h, c] = dq_shared[r, c] + dk_out[t, i_h, c] = dk_shared[r, c] + dw_out[t, i_h, c] = dw_shared[r, c] + db_out[t, i_h, c] = db_shared[r, c] + # Split the serial over-BT dgk reduction across 4 lane + # groups so every thread participates (was 64/256 lanes, + # 21.5% of the kernel in the TIRx profile). + for r_local in T.serial(BT // 4): + for gg, c in T.Parallel(4, K): + r = gg * (BT // 4) + r_local + dgk_part[gg, c] = ( + dgk_part[gg, c] + + T.Cast(acc_dtype, kg_shared[cur, r, c]) + * T.Cast(acc_dtype, dk_shared[r, c]) + + T.Cast(acc_dtype, bg_shared[cur, r, c]) + * T.Cast(acc_dtype, db_shared[r, c]) + ) + + last_idx = T.min(t_off + BT - 1, eos - 1) + for c in T.Parallel(K): + gk_last_shared[c] = gk[last_idx, i_h, c] + dgk_last_frag[c] = dgk_last_frag[c] * T.exp2(gk_last_shared[c]) + for gg, c in T.Parallel(4, K): + dgk_part_shared[gg, c] = dgk_part[gg, c] + T.sync_threads() + for c in T.Parallel(K): + dgk_last[chunk_row, i_h, c] = ( + dgk_last_frag[c] + + dgk_part_shared[0, c] + + dgk_part_shared[1, c] + + dgk_part_shared[2, c] + + dgk_part_shared[3, c] + ) + + # Update dh for the previous chunk. + T.gemm(qg_shared[cur, :, :], do_shared[cur, :, :], b_dh_tmp, transpose_A=True) + T.gemm(w_shared[cur, :, :], dv2_shared, b_dh_tmp, transpose_A=True) + for k_idx, vv in T.Parallel(K, V): + b_dh[k_idx, vv] = T.exp2(gk_last_shared[k_idx]) * b_dh[k_idx, vv] + b_dh_tmp[k_idx, vv] + + if USE_INITIAL_STATE: + for k_idx, vv in T.Parallel(K, V): + dh0[i_n, i_h, k_idx, vv] = T.Cast(state_dtype, b_dh[k_idx, vv]) + + return chunk_dplr_bwd_stream_tl + + +@tilelang.jit( + pass_configs={ + tilelang.PassConfigKey.TL_ENABLE_FAST_MATH: True, + tilelang.PassConfigKey.TL_DISABLE_DATA_RACE_CHECK: False, + }, +) +def _chunk_dplr_bwd_stream_low_smem_kernel( + H, K, V, BT, + in_dtype, state_dtype, + USE_FINAL_STATE_GRADIENT: bool, + USE_INITIAL_STATE: bool, + threads: int = 128, + qside_bv: int = 16, +): + acc_dtype = "float32" + # Exact qside tiling lets the h slices bulk-load; ragged Vs keep the + # scalar predicated staging. + qside_exact = V % qside_bv == 0 + n_tokens, n_seq_plus_one, n_chunks, n_dht, n_dh0 = T.dynamic( + "n_tokens, n_seq_plus_one, n_chunks, n_dht, n_dh0" + ) + n_seqs = n_seq_plus_one - 1 + + @T.prim_func + def chunk_dplr_bwd_stream_low_smem_tl( + qg: T.Tensor((n_tokens, H, K), in_dtype), + bg: T.Tensor((n_tokens, H, K), in_dtype), + w: T.Tensor((n_tokens, H, K), in_dtype), + kg: T.Tensor((n_tokens, H, K), in_dtype), + v: T.Tensor((n_tokens, H, V), in_dtype), + v_new: T.Tensor((n_tokens, H, V), in_dtype), + gk: T.Tensor((n_tokens, H, K), acc_dtype), + do: T.Tensor((n_tokens, H, V), in_dtype), + h: T.Tensor((n_chunks, H, K, V), in_dtype), + A_qb: T.Tensor((n_tokens, H, BT), in_dtype), + A_qk: T.Tensor((n_tokens, H, BT), in_dtype), + dht: T.Tensor((n_dht, H, K, V), state_dtype), + cu_seqlens: T.Tensor((n_seq_plus_one,), "int32"), + chunk_offsets: T.Tensor((n_seq_plus_one,), "int32"), + dq_out: T.Tensor((n_tokens, H, K), in_dtype), + dk_out: T.Tensor((n_tokens, H, K), in_dtype), + dw_out: T.Tensor((n_tokens, H, K), in_dtype), + db_out: T.Tensor((n_tokens, H, K), in_dtype), + dgk_last: T.Tensor((n_chunks, H, K), acc_dtype), + dv2: T.Tensor((n_tokens, H, V), in_dtype), + dv_full: T.Tensor((n_tokens, H, V), in_dtype), + dh0: T.Tensor((n_dh0, H, K, V), state_dtype), + ): + with T.Kernel(n_seqs, H, threads=threads) as (i_n, i_h): + bos = cu_seqlens[i_n] + eos = cu_seqlens[i_n + 1] + boh = chunk_offsets[i_n] + n_chunks = chunk_offsets[i_n + 1] - boh + + b_dh = T.alloc_fragment((K, V), acc_dtype) + state_shared = T.alloc_shared((K, V), in_dtype) + + qg_shared = T.alloc_shared((BT, K), in_dtype) + bg_shared = T.alloc_shared((BT, K), in_dtype) + w_shared = T.alloc_shared((BT, K), in_dtype) + kg_shared = T.alloc_shared((BT, K), in_dtype) + do_shared = T.alloc_shared((BT, V), in_dtype) + v_like_shared = T.alloc_shared((BT, V), in_dtype) + qside_do_shared = T.alloc_shared((BT, qside_bv), in_dtype) + qside_value_shared = T.alloc_shared((BT, qside_bv), in_dtype) + qside_state_shared = T.alloc_shared((K, qside_bv), in_dtype) + A_shared = T.alloc_shared((BT, BT), in_dtype) + + dv_intra_frag = T.alloc_fragment((BT, V), acc_dtype) + dv2_frag = T.alloc_fragment((BT, V), acc_dtype) + dv2_shared = T.alloc_shared((BT, V), in_dtype) + dv_full_frag = T.alloc_fragment((BT, V), acc_dtype) + + dq_frag = T.alloc_fragment((BT, K), acc_dtype) + dk_frag = T.alloc_fragment((BT, K), acc_dtype) + dw_frag = T.alloc_fragment((BT, K), acc_dtype) + db_frag = T.alloc_fragment((BT, K), acc_dtype) + dgk_part = T.alloc_fragment((4, K), acc_dtype) + dgk_part_shared = T.alloc_shared((4, K), acc_dtype) + dgk_h_frag = T.alloc_fragment((K,), acc_dtype) + gk_last_frag = T.alloc_fragment((K,), acc_dtype) + + if USE_FINAL_STATE_GRADIENT: + for k_idx, vv in T.Parallel(K, V): + b_dh[k_idx, vv] = dht[i_n, i_h, k_idx, vv] + else: + for k_idx, vv in T.Parallel(K, V): + b_dh[k_idx, vv] = T.float32(0.0) + + for i_t_rev in T.serial(n_chunks): + i_t = n_chunks - 1 - i_t_rev + t_off = bos + i_t * BT + chunk_row = boh + i_t + + T.clear(dq_frag) + T.clear(dk_frag) + T.clear(dw_frag) + T.clear(db_frag) + for gg, c in T.Parallel(4, K): + dgk_part[gg, c] = T.float32(0.0) + for k_idx in T.Parallel(K): + dgk_h_frag[k_idx] = T.float32(0.0) + gk_last_frag[k_idx] = T.float32(0.0) + + full_tile = t_off + BT <= eos + if full_tile: + # Bulk vectorized copies for interior chunks (TIRx showed + # the scalar predicated loads cap at ~1.5TB/s). + T.copy(qg[t_off: t_off + BT, i_h, 0:K], qg_shared) + T.copy(bg[t_off: t_off + BT, i_h, 0:K], bg_shared) + T.copy(w[t_off: t_off + BT, i_h, 0:K], w_shared) + T.copy(kg[t_off: t_off + BT, i_h, 0:K], kg_shared) + T.copy(do[t_off: t_off + BT, i_h, 0:V], do_shared) + # Stored A matrices are already causally masked. + T.copy(A_qb[t_off: t_off + BT, i_h, 0:BT], A_shared) + else: + for r, c in T.Parallel(BT, K): + t = t_off + r + if t < eos: + qg_shared[r, c] = qg[t, i_h, c] + bg_shared[r, c] = bg[t, i_h, c] + w_shared[r, c] = w[t, i_h, c] + kg_shared[r, c] = kg[t, i_h, c] + else: + qg_shared[r, c] = T.Cast(in_dtype, 0.0) + bg_shared[r, c] = T.Cast(in_dtype, 0.0) + w_shared[r, c] = T.Cast(in_dtype, 0.0) + kg_shared[r, c] = T.Cast(in_dtype, 0.0) + + for r, c in T.Parallel(BT, V): + t = t_off + r + if t < eos: + do_shared[r, c] = do[t, i_h, c] + else: + do_shared[r, c] = T.Cast(in_dtype, 0.0) + + for r, c in T.Parallel(BT, BT): + t = t_off + r + if (t < eos) and (r >= c): + A_shared[r, c] = A_qb[t, i_h, c] + else: + A_shared[r, c] = T.Cast(in_dtype, 0.0) + + # state_shared first holds the current reverse state dH. + T.copy(b_dh, state_shared) + + # dv2 = A_qb^T @ do + bg @ dh + T.gemm(A_shared, do_shared, dv_intra_frag, transpose_A=True, clear_accum=True) + T.gemm(bg_shared, state_shared, dv2_frag, clear_accum=True) + for r, vv in T.Parallel(BT, V): + t = t_off + r + dv2_frag[r, vv] = dv2_frag[r, vv] + dv_intra_frag[r, vv] + if t < eos: + dv2[t, i_h, vv] = T.Cast(in_dtype, dv2_frag[r, vv]) + T.copy(dv2_frag, dv2_shared) + + # Reuse A_shared for A_qk; write dv_full directly from registers. + if full_tile: + T.copy(A_qk[t_off: t_off + BT, i_h, 0:BT], A_shared) + else: + for r, c in T.Parallel(BT, BT): + t = t_off + r + if (t < eos) and (r >= c): + A_shared[r, c] = A_qk[t, i_h, c] + else: + A_shared[r, c] = T.Cast(in_dtype, 0.0) + T.gemm(kg_shared, state_shared, dv_full_frag, clear_accum=True) + T.gemm(A_shared, do_shared, dv_full_frag, transpose_A=True) + for r, vv in T.Parallel(BT, V): + t = t_off + r + if t < eos: + dv_full[t, i_h, vv] = T.Cast(in_dtype, dv_full_frag[r, vv]) + + last_idx = T.min(t_off + BT - 1, eos - 1) + for c in T.Parallel(K): + gk_last_frag[c] = gk[last_idx, i_h, c] + + # q/o-side consumers. Tile over V to match the saved path's + # accumulation order and avoid long-bf16 spike drift in dq/dk. + # v_new and v are staged like do so the slices below never + # touch global memory; boundary rows stay zero-padded. + if full_tile: + T.copy(v_new[t_off: t_off + BT, i_h, 0:V], v_like_shared) + else: + for r, c in T.Parallel(BT, V): + t = t_off + r + if t < eos: + v_like_shared[r, c] = v_new[t, i_h, c] + else: + v_like_shared[r, c] = T.Cast(in_dtype, 0.0) + for i_v in T.serial(T.ceildiv(V, qside_bv)): + # h tile: dgk_h, dq, and dw consume this tile. h rows are + # defined for every chunk, so exact slices bulk-load. + if qside_exact: + T.copy(h[chunk_row, i_h, 0:K, i_v * qside_bv: (i_v + 1) * qside_bv], + qside_state_shared) + for k_idx, vv in T.Parallel(K, qside_bv): + dgk_h_frag[k_idx] = ( + dgk_h_frag[k_idx] + + T.Cast(acc_dtype, qside_state_shared[k_idx, vv]) + * T.Cast(acc_dtype, state_shared[k_idx, i_v * qside_bv + vv]) + ) + else: + for k_idx, vv in T.Parallel(K, qside_bv): + g_v = i_v * qside_bv + vv + if g_v < V: + qside_state_shared[k_idx, vv] = h[chunk_row, i_h, k_idx, g_v] + dgk_h_frag[k_idx] = ( + dgk_h_frag[k_idx] + + T.Cast(acc_dtype, qside_state_shared[k_idx, vv]) + * T.Cast(acc_dtype, state_shared[k_idx, g_v]) + ) + else: + qside_state_shared[k_idx, vv] = T.Cast(in_dtype, 0.0) + for r, vv in T.Parallel(BT, qside_bv): + t = t_off + r + g_v = i_v * qside_bv + vv + if (t < eos) and (g_v < V): + qside_do_shared[r, vv] = do_shared[r, g_v] + else: + qside_do_shared[r, vv] = T.Cast(in_dtype, 0.0) + T.gemm(qside_do_shared, qside_state_shared, dq_frag, transpose_B=True) + for r, vv in T.Parallel(BT, qside_bv): + g_v = i_v * qside_bv + vv + if g_v < V: + qside_value_shared[r, vv] = dv2_shared[r, g_v] + else: + qside_value_shared[r, vv] = T.Cast(in_dtype, 0.0) + T.gemm(qside_value_shared, qside_state_shared, dw_frag, transpose_B=True) + + # dh tile: db consumes it with the staged v_new, dk with v + # (staged below). Reloading dh into the same K x BV + # scratch avoids the K x V shared-memory conflict that + # broke the previous low-smem schedule. + for k_idx, vv in T.Parallel(K, qside_bv): + g_v = i_v * qside_bv + vv + if g_v < V: + qside_state_shared[k_idx, vv] = state_shared[k_idx, g_v] + else: + qside_state_shared[k_idx, vv] = T.Cast(in_dtype, 0.0) + for r, vv in T.Parallel(BT, qside_bv): + g_v = i_v * qside_bv + vv + if g_v < V: + qside_value_shared[r, vv] = v_like_shared[r, g_v] + else: + qside_value_shared[r, vv] = T.Cast(in_dtype, 0.0) + T.gemm(qside_value_shared, qside_state_shared, db_frag, transpose_B=True) + + if full_tile: + T.copy(v[t_off: t_off + BT, i_h, 0:V], v_like_shared) + else: + for r, c in T.Parallel(BT, V): + t = t_off + r + if t < eos: + v_like_shared[r, c] = v[t, i_h, c] + else: + v_like_shared[r, c] = T.Cast(in_dtype, 0.0) + for i_v in T.serial(T.ceildiv(V, qside_bv)): + for k_idx, vv in T.Parallel(K, qside_bv): + g_v = i_v * qside_bv + vv + if g_v < V: + qside_state_shared[k_idx, vv] = state_shared[k_idx, g_v] + else: + qside_state_shared[k_idx, vv] = T.Cast(in_dtype, 0.0) + for r, vv in T.Parallel(BT, qside_bv): + g_v = i_v * qside_bv + vv + if g_v < V: + qside_value_shared[r, vv] = v_like_shared[r, g_v] + else: + qside_value_shared[r, vv] = T.Cast(in_dtype, 0.0) + T.gemm(qside_value_shared, qside_state_shared, dk_frag, transpose_B=True) + + # Reuse v_like_shared as a single K-shaped scratch for the + # dgk_last terms, with the over-BT reduction split across 4 + # lane groups so every thread participates (21.5% of the + # kernel in the high schedule's TIRx profile when serial). + T.copy(dk_frag, v_like_shared) + for r_local in T.serial(BT // 4): + for gg, c in T.Parallel(4, K): + r = gg * (BT // 4) + r_local + dgk_part[gg, c] = ( + dgk_part[gg, c] + + T.Cast(acc_dtype, kg_shared[r, c]) * T.Cast(acc_dtype, v_like_shared[r, c]) + ) + T.copy(db_frag, v_like_shared) + for r_local in T.serial(BT // 4): + for gg, c in T.Parallel(4, K): + r = gg * (BT // 4) + r_local + dgk_part[gg, c] = ( + dgk_part[gg, c] + + T.Cast(acc_dtype, bg_shared[r, c]) * T.Cast(acc_dtype, v_like_shared[r, c]) + ) + + for gg, c in T.Parallel(4, K): + dgk_part_shared[gg, c] = dgk_part[gg, c] + T.sync_threads() + for c in T.Parallel(K): + dgk_last[chunk_row, i_h, c] = ( + dgk_h_frag[c] * T.exp2(gk_last_frag[c]) + + dgk_part_shared[0, c] + + dgk_part_shared[1, c] + + dgk_part_shared[2, c] + + dgk_part_shared[3, c] + ) + + for r, c in T.Parallel(BT, K): + t = t_off + r + if t < eos: + dq_out[t, i_h, c] = T.Cast(in_dtype, dq_frag[r, c]) + dk_out[t, i_h, c] = T.Cast(in_dtype, dk_frag[r, c]) + dw_out[t, i_h, c] = T.Cast(in_dtype, dw_frag[r, c]) + db_out[t, i_h, c] = T.Cast(in_dtype, db_frag[r, c]) + + # Update dh for the previous chunk: scale the carried state + # in place so both GEMMs accumulate straight into it (mirrors + # the Triton dhu kernel; one less (K, V) fp32 fragment). + for k_idx, vv in T.Parallel(K, V): + b_dh[k_idx, vv] = b_dh[k_idx, vv] * T.exp2(gk_last_frag[k_idx]) + T.gemm(qg_shared, do_shared, b_dh, transpose_A=True) + T.gemm(w_shared, dv2_shared, b_dh, transpose_A=True) + + if USE_INITIAL_STATE: + for k_idx, vv in T.Parallel(K, V): + dh0[i_n, i_h, k_idx, vv] = T.Cast(state_dtype, b_dh[k_idx, vv]) + + return chunk_dplr_bwd_stream_low_smem_tl + + +def chunk_dplr_bwd_stream_into( + qg: torch.Tensor, + bg: torch.Tensor, + w: torch.Tensor, + kg: torch.Tensor, + v: torch.Tensor, + v_new: torch.Tensor, + gk: torch.Tensor, + h: torch.Tensor, + h0: torch.Tensor | None, + dht: torch.Tensor | None, + do: torch.Tensor, + A_qb_for_dv: torch.Tensor, + A_qk: torch.Tensor, + dq_out: torch.Tensor, + dk_out: torch.Tensor, + dw_out: torch.Tensor, + db_out: torch.Tensor, + dgk_last_out: torch.Tensor, + dv2_out: torch.Tensor, + dv_full_out: torch.Tensor, + dh0_out: torch.Tensor, + cu_seqlens: torch.Tensor | None = None, + chunk_size: int = 16, + chunk_layout: ChunkLayout | None = None, +) -> tuple[torch.Tensor, torch.Tensor | None, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + B, T_, H, K = qg.shape + V = do.shape[-1] + BT = int(chunk_size) + is_varlen = cu_seqlens is not None + for out in (dq_out, dk_out, dw_out, db_out, dgk_last_out, dv2_out, dv_full_out, dh0_out): + assert out.is_contiguous(), "chunk_dplr_bwd_stream_into requires contiguous outputs" + if V != K: + raise NotImplementedError("The fused DPLR stream backward requires K == V.") + if is_varlen: + assert B == 1 + layout = chunk_layout if chunk_layout is not None else build_varlen_chunk_layout(cu_seqlens, BT, T_) + else: + layout = build_rect_chunk_layout(B, T_, BT, qg.device) + n_chunks = layout.chunk_indices.shape[0] + n_seqs = layout.cu_seqlens.shape[0] - 1 + n_tokens = B * T_ + in_dtype = str(qg.dtype).split(".")[-1] + state_dtype = "float32" + use_dht = dht is not None + use_dh0 = h0 is not None + n_dh0 = n_seqs if use_dh0 else 1 + + qg_f = qg.reshape(n_tokens, H, K).contiguous() + bg_f = bg.reshape(n_tokens, H, K).contiguous() + w_f = w.reshape(n_tokens, H, K).contiguous() + kg_f = kg.reshape(n_tokens, H, K).contiguous() + v_f = v.reshape(n_tokens, H, V).contiguous() + v_new_f = v_new.reshape(n_tokens, H, V).contiguous() + gk_f = gk.reshape(n_tokens, H, K).contiguous() + do_f = do.reshape(n_tokens, H, V).contiguous() + h_f = h.reshape(n_chunks, H, K, V).contiguous() + A_qb_f = A_qb_for_dv.reshape(n_tokens, H, BT).contiguous() + A_qk_f = A_qk.reshape(n_tokens, H, BT).contiguous() + + if use_dht: + dht_f = dht.reshape(n_seqs, H, K, V).contiguous().to(torch.float32) + else: + dht_f = torch.empty((1, H, K, V), dtype=torch.float32, device=qg.device) + + dq_f = dq_out.reshape(n_tokens, H, K).contiguous() + dk_f = dk_out.reshape(n_tokens, H, K).contiguous() + dw_f = dw_out.reshape(n_tokens, H, K).contiguous() + db_f = db_out.reshape(n_tokens, H, K).contiguous() + dgk_last_f = dgk_last_out.reshape(n_chunks, H, K).contiguous() + dv2_f = dv2_out.reshape(n_tokens, H, V).contiguous() + dv_full_f = dv_full_out.reshape(n_tokens, H, V).contiguous() + dh0_f = dh0_out.reshape(n_dh0, H, K, V).contiguous() + + schedule, config = _select_stream_bwd_schedule( + K=K, + V=V, + BT=BT, + in_dtype=in_dtype, + device=qg.device, + ) + if schedule == "low": + kernel = _chunk_dplr_bwd_stream_low_smem_kernel( + H, K, V, BT, + in_dtype, state_dtype, use_dht, use_dh0, + **config, + ) + else: + kernel = _chunk_dplr_bwd_stream_kernel( + H, K, V, BT, + in_dtype, state_dtype, use_dht, use_dh0, + alias_kv=(schedule == "mid"), + **config, + ) + kernel( + qg_f, bg_f, w_f, kg_f, v_f, v_new_f, gk_f, do_f, h_f, + A_qb_f, A_qk_f, dht_f, layout.cu_seqlens, layout.chunk_offsets, + dq_f, dk_f, dw_f, db_f, dgk_last_f, dv2_f, dv_full_f, dh0_f, + ) + dh0 = dh0_out if use_dh0 else None + return dq_out, dk_out, dw_out, db_out, dgk_last_out, dv2_out, dv_full_out, dh0 diff --git a/fla/ops/generalized_delta_rule/dplr/backends/tilelang/cumsum.py b/fla/ops/generalized_delta_rule/dplr/backends/tilelang/cumsum.py new file mode 100644 index 0000000000..2af59c2532 --- /dev/null +++ b/fla/ops/generalized_delta_rule/dplr/backends/tilelang/cumsum.py @@ -0,0 +1,219 @@ +# Copyright (c) 2023-2026, Songlin Yang, Yu Zhang, Zhiyuan Li +# +# This source code is licensed under the MIT license found in the +# LICENSE file in the root directory of this source tree. +# For a list of all contributors, visit: +# https://github.com/fla-org/flash-linear-attention/graphs/contributors + +"""Chunk-local cumulative sum for log-decay gates. + +Produces two outputs: + gi[bos + t, h, k] = scale * sum_{j=0..t} g[bos + j, h, k] (inclusive) + ge[bos + t, h, k] = scale * sum_{j=0..t-1} g[bos + j, h, k] (exclusive) + +where each chunk of size BT cumsums independently. `scale = RCP_LN2` lets +downstream kernels use `T.exp2` directly. + +Rectangular batches compute the prefix as a matmul against constant +(strict-)lower-triangular ones matrices (FLA's chunk_rwkv6_fwd_cumsum +formulation): the gate tile is loaded once, two tensor-core GEMMs produce +the inclusive and exclusive prefixes, and results are scaled and stored. +Same global traffic as a scan with no serial phases; the GEMMs accumulate +in fp32 from tf32-rounded inputs, matching the Triton reference's numerics. +This replaces the earlier vectorized PyTorch chain (fp32 cast + pad + +cumsum + scale + shifted zeros), which cost ~2.5 ms of elementwise/scan +glue kernels per call at h4096, and a segmented serial-scan kernel that +measured 2.2x behind Triton on sm_90. + +Varlen batches keep the irregular-boundary kernel keyed by chunk_indices. +""" + +import tilelang +import tilelang.language as T +import torch + +from .layout import ChunkLayout, build_varlen_chunk_layout + +# --------------------------------------------------------------------------- +# TileLang segmented scan kernel for rectangular batches. +# --------------------------------------------------------------------------- + + +@tilelang.jit( + pass_configs={ + tilelang.PassConfigKey.TL_ENABLE_FAST_MATH: True, + tilelang.PassConfigKey.TL_DISABLE_DATA_RACE_CHECK: False, + }, +) +def _chunk_local_cumsum_rect_kernel_tl( + H, K, BT, BS, in_dtype, scale_value: float, OUTPUT_GE: bool, threads: int = 128 +): + acc_dtype = "float32" + B, Tt = T.dynamic("B, Tt") + + @T.prim_func + def chunk_local_cumsum_rect_tl( + g: T.Tensor((B, Tt, H, K), in_dtype), + gi: T.Tensor((B, Tt, H, K), acc_dtype), + ge: T.Tensor((B, Tt, H, K), acc_dtype), + ): + tpc = T.ceildiv(Tt, BT) + with T.Kernel(T.ceildiv(K, BS), B * tpc, H, threads=threads) as (i_sb, i_c, i_h): + i_n = i_c // tpc + bos = (i_c % tpc) * BT + + g_shared = T.alloc_shared((BT, BS), acc_dtype) + mask_i = T.alloc_shared((BT, BT), acc_dtype) + if OUTPUT_GE: + mask_e = T.alloc_shared((BT, BT), acc_dtype) + gi_frag = T.alloc_fragment((BT, BS), acc_dtype) + if OUTPUT_GE: + ge_frag = T.alloc_fragment((BT, BS), acc_dtype) + + # Interior chunks bulk-copy; the last chunk of a batch element + # (T % BT != 0) takes the predicated scalar path. + if bos + BT <= Tt: + T.copy(g[i_n, bos: bos + BT, i_h, i_sb * BS: i_sb * BS + BS], g_shared) + else: + for r, c in T.Parallel(BT, BS): + t = bos + r + k_idx = i_sb * BS + c + if (t < Tt) and (k_idx < K): + g_shared[r, c] = T.Cast(acc_dtype, g[i_n, t, i_h, k_idx]) + else: + g_shared[r, c] = T.Cast(acc_dtype, 0.0) + + # Mask entries are exact 0/1 values. + for r, c in T.Parallel(BT, BT): + mask_i[r, c] = T.if_then_else(r >= c, T.Cast(acc_dtype, 1.0), T.Cast(acc_dtype, 0.0)) + if OUTPUT_GE: + mask_e[r, c] = T.if_then_else(r > c, T.Cast(acc_dtype, 1.0), T.Cast(acc_dtype, 0.0)) + + T.gemm(mask_i, g_shared, gi_frag, clear_accum=True) + if OUTPUT_GE: + T.gemm(mask_e, g_shared, ge_frag, clear_accum=True) + + for r, c in T.Parallel(BT, BS): + t = bos + r + k_idx = i_sb * BS + c + if (t < Tt) and (k_idx < K): + gi[i_n, t, i_h, k_idx] = gi_frag[r, c] * scale_value + if OUTPUT_GE: + ge[i_n, t, i_h, k_idx] = ge_frag[r, c] * scale_value + + return chunk_local_cumsum_rect_tl + + +# --------------------------------------------------------------------------- +# TileLang chunk-local scan kernel — kept for varlen batches. +# --------------------------------------------------------------------------- + + +@tilelang.jit( + pass_configs={ + tilelang.PassConfigKey.TL_ENABLE_FAST_MATH: True, + tilelang.PassConfigKey.TL_DISABLE_DATA_RACE_CHECK: False, + }, +) +def _chunk_local_cumsum_kernel_tl( + H, K, BT, BS, in_dtype, scale_value: float, OUTPUT_GE: bool +): + acc_dtype = "float32" + n_tokens, n_seq_plus_one, n_chunks = T.dynamic( + "n_tokens, n_seq_plus_one, n_chunks" + ) + + @T.prim_func + def chunk_local_cumsum_tl( + g: T.Tensor((n_tokens, H, K), in_dtype), + gi: T.Tensor((n_tokens, H, K), acc_dtype), + ge: T.Tensor((n_tokens, H, K), acc_dtype), + cu_seqlens: T.Tensor((n_seq_plus_one,), "int32"), + chunk_indices: T.Tensor((n_chunks, 2), "int32"), + ): + with T.Kernel(T.ceildiv(K, BS), n_chunks, H, threads=128) as (i_sb, i_c, i_h): + i_n = chunk_indices[i_c, 0] + i_t = chunk_indices[i_c, 1] + safe_i_n = T.max(i_n, 0) + seq_bos = cu_seqlens[safe_i_n] + seq_eos = cu_seqlens[safe_i_n + 1] + bos_raw = seq_bos + i_t * BT + eos_raw = T.min(bos_raw + BT, seq_eos) + is_valid_chunk = i_n >= 0 + bos = T.if_then_else(is_valid_chunk, bos_raw, T.int32(0)) + eos = T.if_then_else(is_valid_chunk, eos_raw, T.int32(0)) + + s_scalar = T.Cast(acc_dtype, scale_value) + acc = T.alloc_fragment((BS,), acc_dtype) + + for s_local in T.Parallel(BS): + acc[s_local] = T.Cast(acc_dtype, 0.0) + + # For the target BT=64,K=64 path, one thread lane owns one channel + # and scans the chunk time dimension in registers. This avoids the + # heavier generic T.cumsum shared-memory scan while keeping each + # timestep's K-lane loads/stores coalesced. + for t_local in T.serial(BT): + t = bos + t_local + for s_local in T.Parallel(BS): + s = i_sb * BS + s_local + if (t < eos) and (s < K): + if OUTPUT_GE: + ge[t, i_h, s] = acc[s_local] * s_scalar + acc[s_local] += T.Cast(acc_dtype, g[t, i_h, s]) + gi[t, i_h, s] = acc[s_local] * s_scalar + + return chunk_local_cumsum_tl + + +def chunk_local_cumsum( + g: torch.Tensor, + chunk_size: int, + scale: float | None = None, + cu_seqlens: torch.Tensor | None = None, + chunk_layout: ChunkLayout | None = None, + output_ge: bool = True, +) -> tuple[torch.Tensor, torch.Tensor | None]: + """Chunk-local inclusive (gi) and exclusive (ge) cumsum over the time axis. + + With ``output_ge=False`` the kernels never touch the ge buffer, which is + then aliased to gi to skip the unused fp32 allocation, and None is + returned in its place. + """ + is_varlen = cu_seqlens is not None + scale_f = float(scale) if scale is not None else 1.0 + in_dtype = str(g.dtype).split(".")[-1] + + if not is_varlen: + B, T_, H, K = g.shape + kernel = _chunk_local_cumsum_rect_kernel_tl( + H, K, chunk_size, 64, in_dtype, scale_f, output_ge + ) + gi = torch.empty((B, T_, H, K), dtype=torch.float32, device=g.device) + ge = torch.empty((B, T_, H, K), dtype=torch.float32, device=g.device) if output_ge else gi + kernel(g.contiguous(), gi, ge) + return gi, ge if output_ge else None + + # Varlen path: the chunk_indices-keyed kernel handles irregular boundaries. + B, T_, H, K = g.shape + assert B == 1, "Varlen expects B==1" + BT = chunk_size + + N_tokens = B * T_ + layout = chunk_layout if chunk_layout is not None else build_varlen_chunk_layout(cu_seqlens, BT, N_tokens) + g_flat = g.reshape(N_tokens, H, K).contiguous() + + BS = K if K <= 64 else 64 + while BS > 16 and K % BS != 0: + BS //= 2 + if K < BS: + BS = K + + kernel = _chunk_local_cumsum_kernel_tl(H, K, BT, BS, in_dtype, scale_f, output_ge) + gi_flat = torch.empty((N_tokens, H, K), dtype=torch.float32, device=g.device) + ge_flat = torch.empty((N_tokens, H, K), dtype=torch.float32, device=g.device) if output_ge else gi_flat + kernel(g_flat, gi_flat, ge_flat, layout.cu_seqlens, layout.chunk_indices) + return gi_flat.view(B, T_, H, K), ge_flat.view(B, T_, H, K) if output_ge else None + + +__all__ = ["chunk_local_cumsum"] diff --git a/fla/ops/generalized_delta_rule/dplr/backends/tilelang/layout.py b/fla/ops/generalized_delta_rule/dplr/backends/tilelang/layout.py new file mode 100644 index 0000000000..da547d7666 --- /dev/null +++ b/fla/ops/generalized_delta_rule/dplr/backends/tilelang/layout.py @@ -0,0 +1,124 @@ +# Copyright (c) 2023-2026, Songlin Yang, Yu Zhang, Zhiyuan Li +# +# This source code is licensed under the MIT license found in the +# LICENSE file in 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 collections import deque +from dataclasses import dataclass + +import torch + +__all__ = [ + "ChunkLayout", + "build_rect_chunk_layout", + "build_varlen_chunk_layout", +] + + +@dataclass(frozen=True) +class ChunkLayout: + """FLA-style chunk layout shared by DPLR stages.""" + + cu_seqlens: torch.Tensor + chunk_indices: torch.Tensor + chunk_offsets: torch.Tensor + + +def varlen_chunk_count_upper(total_tokens: int, n_seqs: int, chunk_size: int) -> int: + """Host-shape upper bound for packed-varlen chunk slots. + + The exact number of chunks is data-dependent: + sum_i ceil((cu[i + 1] - cu[i]) / chunk_size) + + Reading that exact value on the host introduces a CUDA sync. For varlen + kernels we instead allocate a tight upper bound, + ceil(total_tokens / chunk_size) + n_seqs - 1, and mark any extra slots as + no-op sentinel rows in the CUDA-built layout. + """ + if total_tokens < 0: + raise ValueError("total_tokens must be non-negative") + if n_seqs < 0: + raise ValueError("n_seqs must be non-negative") + if chunk_size <= 0: + raise ValueError("chunk_size must be positive") + if total_tokens == 0 or n_seqs == 0: + return 0 + return (total_tokens + chunk_size - 1) // chunk_size + max(n_seqs - 1, 0) + + +def _as_cuda_cu_seqlens(cu_seqlens: torch.Tensor) -> torch.Tensor: + if cu_seqlens.device.type != "cuda": + raise ValueError("varlen DPLR requires CUDA cu_seqlens; CPU layout construction is not used") + return cu_seqlens.to(dtype=torch.int32).contiguous() + + +# Bounded cache keyed on cu_seqlens identity + int values. fla.utils.tensor_cache +# is not reusable here: it identity-keys every arg, which only works for small +# interned ints like chunk_size — total_tokens is a fresh object per call. +# Cached tensors are kept alive by the cache itself, so ids cannot recycle. +_VARLEN_LAYOUT_CACHE: deque = deque(maxlen=4) + + +def build_varlen_chunk_layout( + cu_seqlens: torch.Tensor, + chunk_size: int, + total_tokens: int, +) -> ChunkLayout: + """Build fixed-shape varlen chunk indices/offsets on CUDA.""" + for cu_cached, cs_cached, tt_cached, layout in _VARLEN_LAYOUT_CACHE: + if cu_cached is cu_seqlens and cs_cached == chunk_size and tt_cached == total_tokens: + return layout + + cu = _as_cuda_cu_seqlens(cu_seqlens) + n_seqs = cu.shape[0] - 1 + nt_alloc = varlen_chunk_count_upper(total_tokens, n_seqs, chunk_size) + + lengths = cu[1:] - cu[:-1] + chunks_per_seq = torch.div(lengths + chunk_size - 1, chunk_size, rounding_mode="floor") + chunk_offsets = torch.cat( + [ + torch.zeros((1,), device=cu.device, dtype=torch.int32), + chunks_per_seq.cumsum(dim=0, dtype=torch.int32), + ], + dim=0, + ).contiguous() + + rows = torch.arange(nt_alloc, device=cu.device, dtype=torch.int32) + exact_total = chunk_offsets[-1] + valid = rows < exact_total + seq = torch.searchsorted(chunk_offsets, rows, right=True, out_int32=True) - 1 + seq = torch.where(valid, seq, torch.full_like(seq, -1)) + safe_seq = seq.clamp(min=0, max=max(n_seqs - 1, 0)) + local_chunk = rows - chunk_offsets[safe_seq] + local_chunk = torch.where(valid, local_chunk, torch.zeros_like(local_chunk)) + + chunk_indices = torch.stack([seq, local_chunk], dim=1).contiguous() + layout = ChunkLayout(cu, chunk_indices, chunk_offsets) + _VARLEN_LAYOUT_CACHE.append((cu_seqlens, chunk_size, total_tokens, layout)) + return layout + + +# Rect layouts are pure functions of (B, T, BT, device): every stage builds the +# same one within a call, so memoize it like the varlen layout above. +_RECT_LAYOUT_CACHE: dict = {} + + +def build_rect_chunk_layout(B: int, T_: int, BT: int, device: torch.device) -> ChunkLayout: + """Build canonical chunk layout for a rectangular batch.""" + key = (B, T_, BT, device) + layout = _RECT_LAYOUT_CACHE.get(key) + if layout is not None: + return layout + n_chunks_per_seq = (T_ + BT - 1) // BT + cu = torch.arange(B + 1, device=device, dtype=torch.int32) * T_ + seq = torch.arange(B, device=device, dtype=torch.int32).repeat_interleave(n_chunks_per_seq) + local = torch.arange(n_chunks_per_seq, device=device, dtype=torch.int32).repeat(B) + chunk_indices = torch.stack([seq, local], dim=1).contiguous() + chunk_offsets = (torch.arange(B + 1, device=device, dtype=torch.int32) * n_chunks_per_seq).contiguous() + layout = ChunkLayout(cu, chunk_indices, chunk_offsets) + if len(_RECT_LAYOUT_CACHE) >= 8: + _RECT_LAYOUT_CACHE.clear() + _RECT_LAYOUT_CACHE[key] = layout + return layout diff --git a/fla/ops/generalized_delta_rule/dplr/backends/tilelang/schedules.py b/fla/ops/generalized_delta_rule/dplr/backends/tilelang/schedules.py new file mode 100644 index 0000000000..43a8bab330 --- /dev/null +++ b/fla/ops/generalized_delta_rule/dplr/backends/tilelang/schedules.py @@ -0,0 +1,185 @@ +# Copyright (c) 2023-2026, Songlin Yang, Yu Zhang, Zhiyuan Li +# +# This source code is licensed under the MIT license found in the +# LICENSE file in the root directory of this source tree. +# For a list of all contributors, visit: +# https://github.com/fla-org/flash-linear-attention/graphs/contributors + +"""Backward schedule arithmetic for the TileLang DPLR backend. + +Pure-python helpers shared by the kernel launcher (which selects a schedule +for the current device) and the backend verifier (which must reject configs +no schedule can launch, without importing tilelang). +""" + +from __future__ import annotations + +import torch + +from fla.utils import get_device_capability + + +def device_cc(device: torch.device) -> int: + major, minor = get_device_capability(device.index or 0) + return major * 10 + minor + + +# Compiler-reported dynamic shared memory of the fused A-backward kernel +# (chunk_dplr_bwd_intra_tl, FUSE_QSIDE_DA=True) at BT=64 with the BK=32 +# config used off cc90. All of its shared tiles are (BT, BK), (BT, BT) or +# (BT, BV) with BV=64, so the footprint is independent of K and identical +# for bf16/fp16 (accumulator tiles are fp32 either way). The two aliased +# dA staging pairs keep it under cc120's 101376B optin; the A-forward +# from_gk stage sits 128B below it, so this stays the binding footprint. +A_BWD_FUSED_BT64_SMEM_BYTES = 98432 + + +def stream_default_threads(BT: int) -> int: + return 128 if BT >= 32 else 64 + + +def stream_low_default_qside_bv(BT: int) -> int: + return 32 if BT >= 64 else 16 + + +def stream_low_bwd_config(BT: int, V: int) -> dict[str, int]: + # At K=V=128 the persistent (K, V) fp32 dh fragment alone is 128 regs per + # thread at 128 threads, so the serial chunk loop spills to local memory + # every iteration; 256 threads halves that to a spill-free 64. + return { + "threads": 256 if V >= 128 else stream_default_threads(BT), + "qside_bv": stream_low_default_qside_bv(BT), + } + + +def dtype_nbytes(dtype: str) -> int: + if dtype in {"float32", "float"}: + return 4 + if dtype in {"bfloat16", "float16", "half"}: + return 2 + raise ValueError(f"unsupported DPLR stream backward dtype {dtype!r}") + + +def stream_pipeline_extra_smem_bytes(K: int, V: int, BT: int, in_dtype: str) -> int: + """Extra smem of the double-buffered chunk loop: a second version of + every global->shared operand tile (4x (BT, K), 3x (BT, V), 2x (BT, BT)). + The fp32 gk_last row stays single-buffered (scalar per-chunk loads).""" + elem = dtype_nbytes(in_dtype) + return elem * (4 * BT * K + 3 * BT * V + 2 * BT * BT) + + +def stream_high_smem_bytes(K: int, V: int, BT: int, in_dtype: str, num_stages: int = 1) -> int: + elem = dtype_nbytes(in_dtype) + # fp32 tiles: gk_last (K,) plus the (4, K) dgk_part staging buffer. + base = elem * (2 * K * V + 8 * BT * K + 5 * BT * V + 2 * BT * BT) + 20 * K + if num_stages < 2: + return base + # Double-buffered: a second version of the nine operand tiles plus the + # (K, V) h tile. + return base + stream_pipeline_extra_smem_bytes(K, V, BT, in_dtype) + elem * K * V + + +def stream_mid_smem_bytes(K: int, V: int, BT: int, in_dtype: str, num_stages: int = 1) -> int: + """High schedule with the two (K, V) state tiles aliased into one and the + four (BT, K) output staging tiles merged into one sequential buffer. The + alias_kv path keeps the h tile single-buffered when pipelined.""" + elem = dtype_nbytes(in_dtype) + base = stream_high_smem_bytes(K, V, BT, in_dtype, 1) - elem * (K * V + 3 * BT * K) + if num_stages < 2: + return base + return base + stream_pipeline_extra_smem_bytes(K, V, BT, in_dtype) + + +def stream_reuse_smem_bytes(K: int, V: int, BT: int, in_dtype: str, qside_bv: int) -> int: + elem = dtype_nbytes(in_dtype) + # The trailing fp32 term is the (4, K) dgk_part lane-split staging buffer. + return ( + elem * ( + K * V + + 4 * BT * K + + 3 * BT * V + + BT * BT + + BT * qside_bv + + K * qside_bv + ) + + 16 * K + ) + + +def stream_low_smem_bytes(K: int, V: int, BT: int, in_dtype: str) -> int: + return stream_reuse_smem_bytes(K, V, BT, in_dtype, stream_low_default_qside_bv(BT)) + + +def stream_bwd_schedule_or_none( + *, + K: int, + V: int, + BT: int, + in_dtype: str, + smem_cap: int, +) -> str | None: + """Return the stream-backward schedule that fits `smem_cap`, or None. + + Mirrors `_select_stream_bwd_schedule`; keep the two in sync so the + verifier never accepts a config the launcher cannot schedule. + """ + high_smem = stream_high_smem_bytes(K, V, BT, in_dtype) + mid_smem = stream_mid_smem_bytes(K, V, BT, in_dtype) + low_smem = stream_low_smem_bytes(K, V, BT, in_dtype) + low_dtype_supported = in_dtype in {"bfloat16", "float16", "half"} + + if high_smem <= smem_cap: + return "high" + if mid_smem <= smem_cap: + return "mid" + if low_dtype_supported and low_smem <= smem_cap: + return "low" + return None + + +def stream_bwd_num_stages( + schedule: str, + *, + K: int, + V: int, + BT: int, + in_dtype: str, + smem_cap: int, +) -> int: + """Software-pipeline depth for the high/mid chunk loop. + + Double-buffering the operand tiles pays off wherever the versioned + footprint still fits the optin cap; the low schedule stays serial. + Returns 0 (a plain serial loop) or 2. + """ + if schedule == "high": + smem = stream_high_smem_bytes(K, V, BT, in_dtype, num_stages=2) + elif schedule == "mid": + smem = stream_mid_smem_bytes(K, V, BT, in_dtype, num_stages=2) + else: + return 0 + return 2 if smem <= smem_cap else 0 + + +def chunk64_schedule_or_none( + *, + K: int, + V: int, + in_dtype: str, + smem_cap: int, + cc: int, +) -> str | None: + """Stream-backward schedule for BT=64, or None if any stage overflows. + + Acceptance must imply launchability of every BT=64 kernel: on cc90 the + fused A-backward stage runs its BK=64 config and fits the 228KB cap, and + elsewhere it runs BK=32 at A_BWD_FUSED_BT64_SMEM_BYTES, which a 99KB + optin (e.g. cc120) just fits — there the A-forward from_gk stage is the + next-tightest at 98304B. At K=V=128 the stream backward still exceeds + every sub-228KB cap, so this returns None regardless of the A stages. + """ + if cc != 90 and smem_cap < A_BWD_FUSED_BT64_SMEM_BYTES: + return None + return stream_bwd_schedule_or_none( + K=K, V=V, BT=64, in_dtype=in_dtype, smem_cap=smem_cap, + ) diff --git a/fla/ops/generalized_delta_rule/dplr/backends/tilelang/wy_fast_bwd.py b/fla/ops/generalized_delta_rule/dplr/backends/tilelang/wy_fast_bwd.py new file mode 100644 index 0000000000..0ff8ab20d1 --- /dev/null +++ b/fla/ops/generalized_delta_rule/dplr/backends/tilelang/wy_fast_bwd.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 + +"""DPLR WY-representation backward. + +Forward (from wy_fast_fwd.py): + A_ak_processed = A_ab_inv @ A_ak (strict-lower-tri output) + w = A_ab_inv @ ag + u = A_ak_processed @ v + +Backward (this kernel): + 1. dA_ak_processed = du @ v^T (strict-lower-tri) + dv = dv0 + A_ak_processed^T @ du + 2. dA_ak = A_ab_inv^T @ dA_ak_processed (strict-lower-tri) + 3. dA_ab_inv = dA_ak_processed @ A_ak^T + 4. dag = A_ab_inv^T @ dw + dA_ab_inv += dw @ ag^T + 5. dA_ab = strict_lower(A_ab_inv^T @ inclusive_lower(dA_ab_inv) @ A_ab_inv^T) + (matrix-inverse sensitivity identity) +""" + +import tilelang +import tilelang.language as T +import torch + +from .layout import ChunkLayout, build_rect_chunk_layout, build_varlen_chunk_layout +from .schedules import device_cc + + +def _wy_bwd_config(BT: int, device: torch.device) -> dict[str, int]: + # 256 threads + bulk copies pay at BT=64 on cc90 and cc120 alike + # (kernel-level 1.5x over the scalar path on both); cc90 also wins at + # BT=32 with 128 threads, while cc120 measures flat there, so the + # cc120 gate stays BT>=64. + if BT <= 16: + return {"threads": 32, "bulk_copy": False} + cc = device_cc(device) + if cc == 90: + return {"threads": 128 if BT < 64 else 256, "bulk_copy": True} + if cc == 120 and BT >= 64: + return {"threads": 256, "bulk_copy": True} + return {"threads": 128, "bulk_copy": False} + + +@tilelang.jit( + pass_configs={tilelang.PassConfigKey.TL_ENABLE_FAST_MATH: True, + tilelang.PassConfigKey.TL_DISABLE_DATA_RACE_CHECK: False, + }, +) +def _wy_fast_bwd_kernel( + H, K, V, BT, in_dtype, + STORE_DV: bool = True, + threads: int = 128, + bulk_copy: bool = False, +): + acc_dtype = "float32" + n_tokens, n_seq_plus_one, n_chunks = T.dynamic( + "n_tokens, n_seq_plus_one, n_chunks" + ) + + @T.prim_func + def wy_fast_bwd_tl( + A_ab_inv: T.Tensor((n_tokens, H, BT), acc_dtype), + A_ak: T.Tensor((n_tokens, H, BT), "float16"), + ag: T.Tensor((n_tokens, H, K), in_dtype), + v: T.Tensor((n_tokens, H, V), in_dtype), + dw: T.Tensor((n_tokens, H, K), in_dtype), + du: T.Tensor((n_tokens, H, V), in_dtype), + dv0: T.Tensor((n_tokens, H, V), in_dtype), + cu_seqlens: T.Tensor((n_seq_plus_one,), "int32"), + chunk_indices: T.Tensor((n_chunks, 2), "int32"), + dA_ab: T.Tensor((n_tokens, H, BT), in_dtype), + dA_ak: T.Tensor((n_tokens, H, BT), in_dtype), + dv: T.Tensor((n_tokens, H, V), in_dtype), + dag: T.Tensor((n_tokens, H, K), in_dtype), + ): + with T.Kernel(n_chunks, H, threads=threads) as (i_c, i_h): + i_n = chunk_indices[i_c, 0] + i_t = chunk_indices[i_c, 1] + safe_i_n = T.max(i_n, 0) + seq_bos = cu_seqlens[safe_i_n] + seq_eos = cu_seqlens[safe_i_n + 1] + bos_raw = seq_bos + i_t * BT + eos_raw = T.min(bos_raw + BT, seq_eos) + is_valid_chunk = i_n >= 0 + bos = T.if_then_else(is_valid_chunk, bos_raw, T.int32(0)) + eos = T.if_then_else(is_valid_chunk, eos_raw, T.int32(0)) + + # A matrices in in_dtype for GEMM operands (cast from fp32 input). + A_ab_inv_shared = T.alloc_shared((BT, BT), in_dtype) + A_ak_shared = T.alloc_shared((BT, BT), in_dtype) + A_tmp_shared = T.alloc_shared((BT, BT), in_dtype) + dA_tmp_frag = T.alloc_fragment((BT, BT), acc_dtype) + dA_tmp_shared = T.alloc_shared((BT, BT), in_dtype) + dA_ak_frag = T.alloc_fragment((BT, BT), acc_dtype) + dA_ab_inv_frag = T.alloc_fragment((BT, BT), acc_dtype) + dA_ab_inv_shared = T.alloc_shared((BT, BT), in_dtype) + dA_ab_frag = T.alloc_fragment((BT, BT), acc_dtype) + tmp_frag = T.alloc_fragment((BT, BT), acc_dtype) + tmp_shared = T.alloc_shared((BT, BT), in_dtype) + A_tmp_frag = T.alloc_fragment((BT, BT), acc_dtype) + + v_shared = T.alloc_shared((BT, V), in_dtype) + du_shared = T.alloc_shared((BT, V), in_dtype) + if STORE_DV: + dv0_shared = T.alloc_shared((BT, V), in_dtype) + dv_frag = T.alloc_fragment((BT, V), acc_dtype) + dv_shared = T.alloc_shared((BT, V), in_dtype) + ag_shared = T.alloc_shared((BT, K), in_dtype) + dw_shared = T.alloc_shared((BT, K), in_dtype) + dag_frag = T.alloc_fragment((BT, K), acc_dtype) + dag_shared = T.alloc_shared((BT, K), in_dtype) + + # Load A_ab_inv with inclusive-lower mask, A_ak with strict-lower + # mask. Interior chunks bulk-copy and mask in shared; boundary + # chunks keep the scalar predicated path (see chunk_A_bwd). + full_tile = (is_valid_chunk and (bos + BT <= eos)) if bulk_copy else False + if full_tile: + T.copy(A_ab_inv[bos: bos + BT, i_h, 0:BT], A_ab_inv_shared) + T.copy(A_ak[bos: bos + BT, i_h, 0:BT], A_ak_shared) + for r, c in T.Parallel(BT, BT): + if r < c: + A_ab_inv_shared[r, c] = T.Cast(in_dtype, 0.0) + if r <= c: + A_ak_shared[r, c] = T.Cast(in_dtype, 0.0) + else: + for r, c in T.Parallel(BT, BT): + t = bos + r + if (t < eos) and (r >= c): + A_ab_inv_shared[r, c] = T.Cast(in_dtype, A_ab_inv[t, i_h, c]) + else: + A_ab_inv_shared[r, c] = T.Cast(in_dtype, 0.0) + if (t < eos) and (r > c): + A_ak_shared[r, c] = T.Cast(in_dtype, A_ak[t, i_h, c]) + else: + A_ak_shared[r, c] = T.Cast(in_dtype, 0.0) + + # A_tmp = A_ab_inv @ A_ak (strict-lower) + T.gemm(A_ab_inv_shared, A_ak_shared, A_tmp_frag, clear_accum=True) + T.copy(A_tmp_frag, A_tmp_shared) + + # Load v, du, dv0; compute dA_tmp = du @ v^T and dv = dv0 + A_tmp^T @ du + if full_tile: + T.copy(v[bos: bos + BT, i_h, 0:V], v_shared) + T.copy(du[bos: bos + BT, i_h, 0:V], du_shared) + if STORE_DV: + T.copy(dv0[bos: bos + BT, i_h, 0:V], dv0_shared) + else: + for r, c in T.Parallel(BT, V): + t = bos + r + if t < eos: + v_shared[r, c] = v[t, i_h, c] + du_shared[r, c] = du[t, i_h, c] + if STORE_DV: + dv0_shared[r, c] = dv0[t, i_h, c] + else: + v_shared[r, c] = T.Cast(in_dtype, 0.0) + du_shared[r, c] = T.Cast(in_dtype, 0.0) + if STORE_DV: + dv0_shared[r, c] = T.Cast(in_dtype, 0.0) + + T.gemm(du_shared, v_shared, dA_tmp_frag, transpose_B=True, clear_accum=True) + if STORE_DV: + # dv = dv0 + A_tmp^T @ du + T.gemm(A_tmp_shared, du_shared, dv_frag, transpose_A=True, clear_accum=True) + for r, c in T.Parallel(BT, V): + dv_frag[r, c] = dv_frag[r, c] + T.Cast(acc_dtype, dv0_shared[r, c]) + T.copy(dv_frag, dv_shared) + if full_tile: + for r, c in T.Parallel(BT, V): + dv[bos + r, i_h, c] = dv_shared[r, c] + else: + for r, c in T.Parallel(BT, V): + t = bos + r + if t < eos: + dv[t, i_h, c] = dv_shared[r, c] + + # dA_tmp = strict_lower(dA_tmp) + for r, c in T.Parallel(BT, BT): + if r > c: + dA_tmp_frag[r, c] = dA_tmp_frag[r, c] + else: + dA_tmp_frag[r, c] = 0.0 + T.copy(dA_tmp_frag, dA_tmp_shared) + + # dA_ak = A_ab_inv^T @ dA_tmp (strict-lower) + T.gemm(A_ab_inv_shared, dA_tmp_shared, dA_ak_frag, transpose_A=True, clear_accum=True) + if full_tile: + for r, c in T.Parallel(BT, BT): + if r > c: + dA_ak[bos + r, i_h, c] = T.Cast(in_dtype, dA_ak_frag[r, c]) + else: + dA_ak[bos + r, i_h, c] = T.Cast(in_dtype, 0.0) + else: + for r, c in T.Parallel(BT, BT): + t = bos + r + if t < eos: + if r > c: + dA_ak[t, i_h, c] = T.Cast(in_dtype, dA_ak_frag[r, c]) + else: + dA_ak[t, i_h, c] = T.Cast(in_dtype, 0.0) + + # dA_ab_inv = dA_tmp @ A_ak^T + T.gemm(dA_tmp_shared, A_ak_shared, dA_ab_inv_frag, transpose_B=True, clear_accum=True) + + # Load ag, dw; compute dA_ab_inv += dw @ ag^T; dag = A_ab_inv^T @ dw + if full_tile: + T.copy(ag[bos: bos + BT, i_h, 0:K], ag_shared) + T.copy(dw[bos: bos + BT, i_h, 0:K], dw_shared) + else: + for r, c in T.Parallel(BT, K): + t = bos + r + if t < eos: + ag_shared[r, c] = ag[t, i_h, c] + dw_shared[r, c] = dw[t, i_h, c] + else: + ag_shared[r, c] = T.Cast(in_dtype, 0.0) + dw_shared[r, c] = T.Cast(in_dtype, 0.0) + T.gemm(dw_shared, ag_shared, dA_ab_inv_frag, transpose_B=True) + T.gemm(A_ab_inv_shared, dw_shared, dag_frag, transpose_A=True, clear_accum=True) + T.copy(dag_frag, dag_shared) + if full_tile: + for r, c in T.Parallel(BT, K): + dag[bos + r, i_h, c] = dag_shared[r, c] + else: + for r, c in T.Parallel(BT, K): + t = bos + r + if t < eos: + dag[t, i_h, c] = dag_shared[r, c] + + # dA_ab = strict_lower(A_ab_inv^T @ inclusive_lower(dA_ab_inv) @ A_ab_inv^T) + for r, c in T.Parallel(BT, BT): + if r >= c: + dA_ab_inv_shared[r, c] = T.Cast(in_dtype, dA_ab_inv_frag[r, c]) + else: + dA_ab_inv_shared[r, c] = T.Cast(in_dtype, 0.0) + + # tmp = A_ab_inv^T @ dA_ab_inv + T.gemm(A_ab_inv_shared, dA_ab_inv_shared, tmp_frag, transpose_A=True, clear_accum=True) + T.copy(tmp_frag, tmp_shared) + # dA_ab = tmp @ A_ab_inv^T + T.gemm(tmp_shared, A_ab_inv_shared, dA_ab_frag, transpose_B=True, clear_accum=True) + if full_tile: + for r, c in T.Parallel(BT, BT): + if r > c: + dA_ab[bos + r, i_h, c] = T.Cast(in_dtype, dA_ab_frag[r, c]) + else: + dA_ab[bos + r, i_h, c] = T.Cast(in_dtype, 0.0) + else: + for r, c in T.Parallel(BT, BT): + t = bos + r + if t < eos: + if r > c: + dA_ab[t, i_h, c] = T.Cast(in_dtype, dA_ab_frag[r, c]) + else: + dA_ab[t, i_h, c] = T.Cast(in_dtype, 0.0) + + return wy_fast_bwd_tl + + +def chunk_dplr_bwd_wy_repr_into( + A_ab_inv: torch.Tensor, + A_ak: torch.Tensor, + v: torch.Tensor, + ag: torch.Tensor, + dw: torch.Tensor, + du: torch.Tensor, + dv0: torch.Tensor, + dA_ab_out: torch.Tensor, + dA_ak_out: torch.Tensor, + dv_out: torch.Tensor, + dag_out: torch.Tensor, + store_dv: bool = True, + cu_seqlens: torch.Tensor | None = None, + chunk_size: int = 16, + chunk_layout: ChunkLayout | None = None, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + """Mutable-output variant used by recompute backward workspace reuse.""" + for out in (dA_ab_out, dA_ak_out, dv_out, dag_out): + assert out.is_contiguous(), "chunk_dplr_bwd_wy_repr_into requires contiguous outputs" + B, T_, H, K = dw.shape + V = du.shape[-1] + BT = chunk_size + is_varlen = cu_seqlens is not None + + if is_varlen: + assert B == 1 + layout = chunk_layout if chunk_layout is not None else build_varlen_chunk_layout(cu_seqlens, BT, T_) + else: + layout = build_rect_chunk_layout(B, T_, BT, dw.device) + N_tokens = B * T_ + in_dtype = str(dw.dtype).split(".")[-1] + + A_ab_inv_f = A_ab_inv.reshape(N_tokens, H, BT).contiguous() + A_ak_f = A_ak.reshape(N_tokens, H, BT).contiguous() + ag_f = ag.reshape(N_tokens, H, K).contiguous() + v_f = v.reshape(N_tokens, H, V).contiguous() + dw_f = dw.reshape(N_tokens, H, K).contiguous() + du_f = du.reshape(N_tokens, H, V).contiguous() + dv0_f = dv0.reshape(N_tokens, H, V).contiguous() + dA_ab_f = dA_ab_out.reshape(N_tokens, H, BT).contiguous() + dA_ak_f = dA_ak_out.reshape(N_tokens, H, BT).contiguous() + dv_f = dv_out.reshape(N_tokens, H, V).contiguous() + dag_f = dag_out.reshape(N_tokens, H, K).contiguous() + + kernel = _wy_fast_bwd_kernel( + H, K, V, BT, in_dtype, + STORE_DV=bool(store_dv), + **_wy_bwd_config(BT, dw.device), + ) + kernel( + A_ab_inv_f, A_ak_f, ag_f, v_f, dw_f, du_f, dv0_f, + layout.cu_seqlens, layout.chunk_indices, + dA_ab_f, dA_ak_f, dv_f, dag_f, + ) + return dA_ab_out, dA_ak_out, dv_out, dag_out diff --git a/fla/ops/generalized_delta_rule/dplr/backends/tilelang/wy_fast_fwd.py b/fla/ops/generalized_delta_rule/dplr/backends/tilelang/wy_fast_fwd.py new file mode 100644 index 0000000000..ec3db78b09 --- /dev/null +++ b/fla/ops/generalized_delta_rule/dplr/backends/tilelang/wy_fast_fwd.py @@ -0,0 +1,401 @@ +# Copyright (c) 2023-2026, Songlin Yang, Yu Zhang, Zhiyuan Li +# +# This source code is licensed under the MIT license found in the +# LICENSE file in the root directory of this source tree. +# For a list of all contributors, visit: +# https://github.com/fla-org/flash-linear-attention/graphs/contributors + +"""DPLR WY representation forward. + +Two stages: +1. `_prepare_wy_repr_fwd_kernel`: invert the strictly lower-triangular A_ab + in-place per chunk via the iterative formula M_{k+1} = M_k + e_i (e_i^T M_k) + for i = 1..BT-1, then add identity. Yields fp32 A_ab_inv, matching FLA's + public dtype boundary. +2. `_wu_fwd_kernel`: compute `b_Aak = A_ab_inv @ A_ak` (causal masks applied), + then `w = A_ab_inv @ ag` and `u = b_Aak @ v`. + +Adapted from FLA's prepare_wy_repr_fwd_kernel_chunk32 + wu_fwd_kernel. +""" + +import tilelang +import tilelang.language as T +import torch + +from .layout import ChunkLayout, build_rect_chunk_layout, build_varlen_chunk_layout +from .schedules import device_cc + + +def _wu_fwd_config(BT: int, K: int, device: torch.device) -> dict[str, int]: + # Bulk copies + wide threads pay at BT=64 on cc90 and cc120 alike + # (kernel-level 1.7-2.1x over the scalar path on both); cc90 also wins + # at BT=32 with 128 threads, while cc120 measures flat there, so the + # cc120 gate stays BT>=64. At BT=32/K>=128 on cc90, 256 threads spread + # the wide ag/v tiles and N=128 GEMMs further (-21% kernel time); at + # K=64 they over-subscribe the narrow tiles (+3%). + if BT <= 16: + return {"threads": 32, "bulk_copy": False} + cc = device_cc(device) + if cc == 90: + return {"threads": 256 if BT >= 64 or K >= 128 else 128, "bulk_copy": True} + if cc == 120 and BT >= 64: + return {"threads": 256, "bulk_copy": True} + return {"threads": 128, "bulk_copy": False} + + +@tilelang.jit( + pass_configs={tilelang.PassConfigKey.TL_ENABLE_FAST_MATH: True, + tilelang.PassConfigKey.TL_DISABLE_DATA_RACE_CHECK: False, + }, +) +def _prepare_wy_repr_fwd_kernel(H, BT, in_dtype, threads: int = 32): + acc_dtype = "float32" + n_tokens, n_seq_plus_one, n_chunks = T.dynamic( + "n_tokens, n_seq_plus_one, n_chunks" + ) + + @T.prim_func + def prepare_wy_repr_fwd_tl( + A_ab: T.Tensor((n_tokens, H, BT), "float16"), + cu_seqlens: T.Tensor((n_seq_plus_one,), "int32"), + chunk_indices: T.Tensor((n_chunks, 2), "int32"), + A_ab_inv: T.Tensor((n_tokens, H, BT), acc_dtype), + ): + with T.Kernel(n_chunks, H, threads=threads) as (i_c, i_h): + i_n = chunk_indices[i_c, 0] + i_t = chunk_indices[i_c, 1] + safe_i_n = T.max(i_n, 0) + seq_bos = cu_seqlens[safe_i_n] + seq_eos = cu_seqlens[safe_i_n + 1] + bos_raw = seq_bos + i_t * BT + eos_raw = T.min(bos_raw + BT, seq_eos) + is_valid_chunk = i_n >= 0 + bos = T.if_then_else(is_valid_chunk, bos_raw, T.int32(0)) + eos = T.if_then_else(is_valid_chunk, eos_raw, T.int32(0)) + + M = T.alloc_shared((BT, BT), acc_dtype) + v = T.alloc_shared((BT,), acc_dtype) + v_new = T.alloc_fragment((BT,), acc_dtype) + T.clear(v_new) + + # Load A_ab, mask to strict lower triangular. Interior chunks + # bulk-copy and mask in shared; boundary chunks stay scalar. + if is_valid_chunk and (bos + BT <= eos): + T.copy(A_ab[bos: bos + BT, i_h, 0:BT], M) + for r, c in T.Parallel(BT, BT): + if r <= c: + M[r, c] = 0.0 + else: + for r, c in T.Parallel(BT, BT): + t = bos + r + if (r > c) and (t < eos): + M[r, c] = A_ab[t, i_h, c] + else: + M[r, c] = 0.0 + + # Iterative inversion: for i in 1..BT-1. The j-reduction runs + # serially so each c-lane accumulator stays private to one thread; + # the trip count stops at row_i since v[j>=row_i] is zero. + for i in T.serial(BT - 1): + row_i = i + 1 + for c in T.Parallel(BT): + v[c] = M[row_i, c] + + for j in T.serial(1, row_i): + for c in T.Parallel(BT): + if c < j: + v_new[c] = v_new[c] + v[j] * M[j, c] + + for c in T.Parallel(BT): + if c < row_i: + M[row_i, c] = v[c] + v_new[c] + v_new[c] = 0.0 + + # Add identity to diagonal + for r, c in T.Parallel(BT, BT): + if r == c: + M[r, c] = M[r, c] + 1.0 + + # Store + for r, c in T.Parallel(BT, BT): + t = bos + r + if t < eos: + A_ab_inv[t, i_h, c] = M[r, c] + + return prepare_wy_repr_fwd_tl + + +@tilelang.jit( + pass_configs={tilelang.PassConfigKey.TL_ENABLE_FAST_MATH: True, + tilelang.PassConfigKey.TL_DISABLE_DATA_RACE_CHECK: False, + }, +) +def _prepare_wy_repr_fwd_kernel_bt64(H, in_dtype, threads: int = 32): + acc_dtype = "float32" + BT = 64 + BC = 32 + n_tokens, n_seq_plus_one, n_chunks = T.dynamic( + "n_tokens, n_seq_plus_one, n_chunks" + ) + + @T.prim_func + def prepare_wy_repr_fwd_bt64_tl( + A_ab: T.Tensor((n_tokens, H, BT), "float16"), + cu_seqlens: T.Tensor((n_seq_plus_one,), "int32"), + chunk_indices: T.Tensor((n_chunks, 2), "int32"), + A_ab_inv: T.Tensor((n_tokens, H, BT), acc_dtype), + ): + with T.Kernel(n_chunks, H, threads=threads) as (i_c, i_h): + i_n = chunk_indices[i_c, 0] + i_t = chunk_indices[i_c, 1] + safe_i_n = T.max(i_n, 0) + seq_bos = cu_seqlens[safe_i_n] + seq_eos = cu_seqlens[safe_i_n + 1] + bos_raw = seq_bos + i_t * BT + eos_raw = T.min(bos_raw + BT, seq_eos) + is_valid_chunk = i_n >= 0 + bos = T.if_then_else(is_valid_chunk, bos_raw, T.int32(0)) + eos = T.if_then_else(is_valid_chunk, eos_raw, T.int32(0)) + + A1 = T.alloc_shared((BC, BC), acc_dtype) + A2 = T.alloc_shared((BC, BC), acc_dtype) + A3 = T.alloc_shared((BC, BC), acc_dtype) + tmp = T.alloc_fragment((BC, BC), acc_dtype) + tmp_shared = T.alloc_shared((BC, BC), acc_dtype) + A3_out = T.alloc_fragment((BC, BC), acc_dtype) + v = T.alloc_shared((BC,), acc_dtype) + v2 = T.alloc_shared((BC,), acc_dtype) + v_new = T.alloc_fragment((BC,), acc_dtype) + v_new2 = T.alloc_fragment((BC,), acc_dtype) + T.clear(v_new) + T.clear(v_new2) + + # FLA's chunk64 path inverts the two 32x32 diagonal blocks + # independently, then forms the bottom-left block with two GEMMs. + # The two inversions share one loop so the independent FMA chains + # interleave (same per-matrix accumulation order); only j < row_i + # can contribute since v[j>=row_i] is exactly zero above the + # strict lower triangle. + for r, c in T.Parallel(BC, BC): + t1 = bos + r + t2 = bos + BC + r + if (r > c) and (t1 < eos): + A1[r, c] = A_ab[t1, i_h, c] + else: + A1[r, c] = 0.0 + if (r > c) and (t2 < eos): + A2[r, c] = A_ab[t2, i_h, BC + c] + else: + A2[r, c] = 0.0 + if t2 < eos: + A3[r, c] = A_ab[t2, i_h, c] + else: + A3[r, c] = 0.0 + + for i in T.serial(BC - 1): + row_i = i + 1 + for c in T.Parallel(BC): + v[c] = A1[row_i, c] + v2[c] = A2[row_i, c] + for j in T.serial(1, row_i): + for c in T.Parallel(BC): + if c < j: + v_new[c] = v_new[c] + v[j] * A1[j, c] + v_new2[c] = v_new2[c] + v2[j] * A2[j, c] + for c in T.Parallel(BC): + if c < row_i: + A1[row_i, c] = v[c] + v_new[c] + A2[row_i, c] = v2[c] + v_new2[c] + v_new[c] = 0.0 + v_new2[c] = 0.0 + + for r, c in T.Parallel(BC, BC): + if r == c: + A1[r, c] = A1[r, c] + 1.0 + A2[r, c] = A2[r, c] + 1.0 + + T.gemm(A2, A3, tmp, clear_accum=True) + T.copy(tmp, tmp_shared) + T.gemm(tmp_shared, A1, A3_out, clear_accum=True) + + for r, c in T.Parallel(BC, BC): + t1 = bos + r + t2 = bos + BC + r + if t1 < eos: + A_ab_inv[t1, i_h, c] = A1[r, c] + A_ab_inv[t1, i_h, BC + c] = 0.0 + if t2 < eos: + A_ab_inv[t2, i_h, c] = A3_out[r, c] + A_ab_inv[t2, i_h, BC + c] = A2[r, c] + + return prepare_wy_repr_fwd_bt64_tl + + +@tilelang.jit( + pass_configs={tilelang.PassConfigKey.TL_ENABLE_FAST_MATH: True, + tilelang.PassConfigKey.TL_DISABLE_DATA_RACE_CHECK: False, + }, +) +def _wu_fwd_kernel(H, K, V, BT, in_dtype, threads: int = 32, bulk_copy: bool = False): + acc_dtype = "float32" + n_tokens, n_seq_plus_one, n_chunks = T.dynamic( + "n_tokens, n_seq_plus_one, n_chunks" + ) + + @T.prim_func + def wu_fwd_tl( + ag: T.Tensor((n_tokens, H, K), in_dtype), + v: T.Tensor((n_tokens, H, V), in_dtype), + A_ab_inv: T.Tensor((n_tokens, H, BT), acc_dtype), + A_ak: T.Tensor((n_tokens, H, BT), "float16"), + cu_seqlens: T.Tensor((n_seq_plus_one,), "int32"), + chunk_indices: T.Tensor((n_chunks, 2), "int32"), + w: T.Tensor((n_tokens, H, K), in_dtype), + u: T.Tensor((n_tokens, H, V), in_dtype), + ): + with T.Kernel(n_chunks, H, threads=threads) as (i_c, i_h): + i_n = chunk_indices[i_c, 0] + i_t = chunk_indices[i_c, 1] + safe_i_n = T.max(i_n, 0) + seq_bos = cu_seqlens[safe_i_n] + seq_eos = cu_seqlens[safe_i_n + 1] + bos_raw = seq_bos + i_t * BT + eos_raw = T.min(bos_raw + BT, seq_eos) + is_valid_chunk = i_n >= 0 + bos = T.if_then_else(is_valid_chunk, bos_raw, T.int32(0)) + eos = T.if_then_else(is_valid_chunk, eos_raw, T.int32(0)) + + # FLA keeps A_ab_inv @ A_ak in fp32/tf32, then downcasts the + # processed Aak before the u = Aak @ v GEMM. Keep separate + # fp32 and input-dtype views so w remains aligned with FLA's + # bf16 A_ab_inv @ ag path. + A_inv_acc_shared = T.alloc_shared((BT, BT), acc_dtype) + A_ak_acc_shared = T.alloc_shared((BT, BT), acc_dtype) + A_inv_shared = T.alloc_shared((BT, BT), in_dtype) + Aak_processed_frag = T.alloc_fragment((BT, BT), acc_dtype) + Aak_processed_shared = T.alloc_shared((BT, BT), in_dtype) + ag_shared = T.alloc_shared((BT, K), in_dtype) + v_shared = T.alloc_shared((BT, V), in_dtype) + w_frag = T.alloc_fragment((BT, K), acc_dtype) + u_frag = T.alloc_fragment((BT, V), acc_dtype) + + # Load A_ab_inv (inclusive lower-tri after diagonal-add) and A_ak. + # Interior chunks bulk-copy and mask in shared; boundary chunks + # keep the scalar predicated path (same gating as chunk_A_bwd). + full_tile = (is_valid_chunk and (bos + BT <= eos)) if bulk_copy else False + if full_tile: + T.copy(A_ab_inv[bos: bos + BT, i_h, 0:BT], A_inv_acc_shared) + T.copy(A_ak[bos: bos + BT, i_h, 0:BT], A_ak_acc_shared) + # The bf16 view is the same tile after masking, so derive it + # from the fp32 copy in shared instead of reading global twice. + for r, c in T.Parallel(BT, BT): + if r < c: + A_inv_acc_shared[r, c] = 0.0 + if r <= c: + A_ak_acc_shared[r, c] = 0.0 + A_inv_shared[r, c] = T.Cast(in_dtype, A_inv_acc_shared[r, c]) + else: + for r, c in T.Parallel(BT, BT): + t = bos + r + if (t < eos) and (r >= c): + A_inv_acc_shared[r, c] = A_ab_inv[t, i_h, c] + else: + A_inv_acc_shared[r, c] = 0.0 + A_inv_shared[r, c] = T.Cast(in_dtype, A_inv_acc_shared[r, c]) + if (t < eos) and (r > c): + A_ak_acc_shared[r, c] = A_ak[t, i_h, c] + else: + A_ak_acc_shared[r, c] = 0.0 + + T.gemm(A_inv_acc_shared, A_ak_acc_shared, Aak_processed_frag, clear_accum=True) + for r, c in T.Parallel(BT, BT): + Aak_processed_shared[r, c] = T.Cast(in_dtype, Aak_processed_frag[r, c]) + + if full_tile: + T.copy(ag[bos: bos + BT, i_h, 0:K], ag_shared) + else: + for r, c in T.Parallel(BT, K): + t = bos + r + if t < eos: + ag_shared[r, c] = ag[t, i_h, c] + else: + ag_shared[r, c] = T.Cast(in_dtype, 0.0) + T.gemm(A_inv_shared, ag_shared, w_frag, clear_accum=True) + if full_tile: + for r, c in T.Parallel(BT, K): + w[bos + r, i_h, c] = T.Cast(in_dtype, w_frag[r, c]) + else: + for r, c in T.Parallel(BT, K): + t = bos + r + if t < eos: + w[t, i_h, c] = T.Cast(in_dtype, w_frag[r, c]) + + if full_tile: + T.copy(v[bos: bos + BT, i_h, 0:V], v_shared) + else: + for r, c in T.Parallel(BT, V): + t = bos + r + if t < eos: + v_shared[r, c] = v[t, i_h, c] + else: + v_shared[r, c] = T.Cast(in_dtype, 0.0) + T.gemm(Aak_processed_shared, v_shared, u_frag, clear_accum=True) + if full_tile: + for r, c in T.Parallel(BT, V): + u[bos + r, i_h, c] = T.Cast(in_dtype, u_frag[r, c]) + else: + for r, c in T.Parallel(BT, V): + t = bos + r + if t < eos: + u[t, i_h, c] = T.Cast(in_dtype, u_frag[r, c]) + + return wu_fwd_tl + + +def prepare_wy_repr_fwd( + ag: torch.Tensor, + v: torch.Tensor, + A_ak: torch.Tensor, + A_ab: torch.Tensor, + cu_seqlens: torch.Tensor | None = None, + chunk_size: int = 16, + chunk_layout: ChunkLayout | None = None, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + B, T_, H, K = ag.shape + V = v.shape[-1] + BT = chunk_size + is_varlen = cu_seqlens is not None + + if is_varlen: + assert B == 1 + layout = chunk_layout if chunk_layout is not None else build_varlen_chunk_layout(cu_seqlens, BT, T_) + else: + layout = build_rect_chunk_layout(B, T_, BT, ag.device) + N_tokens = B * T_ + in_dtype = str(ag.dtype).split(".")[-1] + + A_ab_f = A_ab.reshape(N_tokens, H, BT).contiguous() + A_ak_f = A_ak.reshape(N_tokens, H, BT).contiguous() + + if BT == 64: + inv_kernel = _prepare_wy_repr_fwd_kernel_bt64(H, in_dtype, threads=32) + else: + inv_kernel = _prepare_wy_repr_fwd_kernel(H, BT, in_dtype, threads=32) + A_ab_inv_f = torch.empty((N_tokens, H, BT), dtype=torch.float32, device=ag.device) + inv_kernel(A_ab_f, layout.cu_seqlens, layout.chunk_indices, A_ab_inv_f) + + ag_f = ag.reshape(N_tokens, H, K).contiguous() + v_f = v.reshape(N_tokens, H, V).contiguous() + + wu_kernel = _wu_fwd_kernel( + H, K, V, BT, in_dtype, + **_wu_fwd_config(BT, K, ag.device), + ) + w_f = torch.empty((N_tokens, H, K), dtype=ag.dtype, device=ag.device) + u_f = torch.empty((N_tokens, H, V), dtype=v.dtype, device=v.device) + wu_kernel(ag_f, v_f, A_ab_inv_f, A_ak_f, layout.cu_seqlens, layout.chunk_indices, w_f, u_f) + + w = w_f.view(B, T_, H, K) + u = u_f.view(B, T_, H, V) + A_ab_inv = A_ab_inv_f.view(B, T_, H, BT) + return w, u, A_ab_inv diff --git a/fla/ops/generalized_delta_rule/dplr/chunk.py b/fla/ops/generalized_delta_rule/dplr/chunk.py index 6bf1c3ab9d..fd2bf93836 100644 --- a/fla/ops/generalized_delta_rule/dplr/chunk.py +++ b/fla/ops/generalized_delta_rule/dplr/chunk.py @@ -9,6 +9,7 @@ import torch +from fla.ops.backends import dispatch from fla.ops.cp import FLACPContext from fla.ops.cp.chunk_delta_h import ( chunk_gated_delta_rule_bwd_dhu_pre_process, @@ -30,6 +31,20 @@ from fla.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, @@ -149,6 +164,7 @@ def forward( 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, @@ -169,6 +185,10 @@ def forward( 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) @@ -434,10 +454,11 @@ def backward( 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""" @@ -497,6 +519,10 @@ def chunk_dplr_delta_rule( 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 +541,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 +575,7 @@ def chunk_dplr_delta_rule( cu_seqlens, cu_seqlens_cpu, safe_gate, + lower_bound, chunk_size, disable_recompute, cp_context, diff --git a/fla/ops/rwkv7/chunk.py b/fla/ops/rwkv7/chunk.py index df09b76421..ea6158e34a 100644 --- a/fla/ops/rwkv7/chunk.py +++ b/fla/ops/rwkv7/chunk.py @@ -27,6 +27,7 @@ def chunk_rwkv7( chunk_size: int | None = None, disable_recompute: bool = False, cp_context: FLACPContext | None = None, + lower_bound: float | None = None, **kwargs, ): """ @@ -68,6 +69,10 @@ def chunk_rwkv7( 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( @@ -89,4 +94,5 @@ def chunk_rwkv7( chunk_size=chunk_size, disable_recompute=disable_recompute, cp_context=cp_context, + lower_bound=lower_bound, ) diff --git a/fla/utils/__init__.py b/fla/utils/__init__.py index 81daca6083..8272e833b4 100644 --- a/fla/utils/__init__.py +++ b/fla/utils/__init__.py @@ -58,6 +58,8 @@ device_torch_lib, get_all_max_shared_mem, get_available_device, + get_device_capability, + get_device_smem_optin, get_multiprocessor_count, map_triton_backend_to_torch_device, ) diff --git a/fla/utils/_device.py b/fla/utils/_device.py index 5a5711abb9..ebe9dbdcf1 100644 --- a/fla/utils/_device.py +++ b/fla/utils/_device.py @@ -92,6 +92,18 @@ def get_multiprocessor_count(tensor_idx: int = 0) -> int: return 1 +@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: diff --git a/tests/context_parallel/test_cp_dplr.py b/tests/context_parallel/test_cp_dplr.py index 014a0a1b6b..80887f5d5b 100644 --- a/tests/context_parallel/test_cp_dplr.py +++ b/tests/context_parallel/test_cp_dplr.py @@ -44,6 +44,7 @@ from fla.ops.cp import build_cp_context from fla.ops.generalized_delta_rule.dplr import chunk_dplr_delta_rule +from fla.ops.generalized_delta_rule.dplr.backends import DPLRTileLangBackend from fla.ops.generalized_delta_rule.dplr.naive import dplr_recurrence from fla.utils import assert_close @@ -82,6 +83,8 @@ def run_cp_dplr_test_worker( dtype, safe_gate: bool = False, op=chunk_dplr_delta_rule, + op_chunk_size: int = 64, + assert_tilelang: bool = False, ): """ Worker function for CP DPLR test. @@ -242,21 +245,40 @@ def run_cp_dplr_test_worker( f"pre_num_ranks: {context.pre_num_ranks}") dist.barrier() - # CP Forward - o_local, _ = op( - q=q_local, - k=k_local, - v=v_local, - a=a_local, - b=b_local, - gk=gk_local, - cp_context=context, - safe_gate=True, - chunk_size=64, - ) - - # CP Backward - o_local.backward(do_local) + route_spy = None + if assert_tilelang: + from fla.ops.generalized_delta_rule.dplr.backends.tilelang import DPLRTileLangBackend + route_spy = [] + _orig = DPLRTileLangBackend.chunk_dplr_delta_rule + + def _spy(self, *args, **kwargs): + route_spy.append(1) + return _orig(self, *args, **kwargs) + + DPLRTileLangBackend.chunk_dplr_delta_rule = _spy + + try: + # CP Forward + o_local, _ = op( + q=q_local, + k=k_local, + v=v_local, + a=a_local, + b=b_local, + gk=gk_local, + cp_context=context, + safe_gate=True, + chunk_size=op_chunk_size, + ) + + # CP Backward + o_local.backward(do_local) + + if assert_tilelang: + assert route_spy, "TileLang backend route was not taken" + finally: + if assert_tilelang: + DPLRTileLangBackend.chunk_dplr_delta_rule = _orig # Step 4: Result Aggregation and Verification o_gathered = [torch.zeros_like(o_local) for _ in range(world_size)] @@ -338,6 +360,8 @@ def run_cp_test_with_spawn( dtype=torch.bfloat16, safe_gate: bool = False, op=chunk_dplr_delta_rule, + op_chunk_size: int = 64, + assert_tilelang: bool = False, ): """ Run CP test using torch.multiprocessing.spawn. @@ -345,7 +369,7 @@ def run_cp_test_with_spawn( """ mp.start_processes( run_cp_dplr_test_worker, - args=(world_size, test_name, T, H, D, lengths, dtype, safe_gate, op), + args=(world_size, test_name, T, H, D, lengths, dtype, safe_gate, op, op_chunk_size, assert_tilelang), nprocs=world_size, join=True, start_method='spawn', @@ -441,6 +465,32 @@ def test_cp2_safe_gate(): ) +def test_cp2_tilelang_route(): + """CP2 through the TileLang DPLR backend (FLA_TILELANG=1), with route assertion.""" + if torch.cuda.device_count() < 2: + pytest.skip("At least 2 GPUs required") + if not DPLRTileLangBackend.is_available(): + pytest.skip("TileLang backend not available") + + prev = os.environ.get('FLA_TILELANG') + os.environ['FLA_TILELANG'] = '1' + try: + run_cp_test_with_spawn( + world_size=2, + test_name="CP2_TileLang", + T=1024, H=128, D=64, + lengths=[400, 624], + dtype=torch.bfloat16, + op_chunk_size=32, + assert_tilelang=True, + ) + finally: + if prev is None: + del os.environ['FLA_TILELANG'] + else: + os.environ['FLA_TILELANG'] = prev + + # ============================================================ # Main Entry Point (for torchrun) # ============================================================ diff --git a/tests/ops/test_backends.py b/tests/ops/test_backends.py index 02d52f2160..faf4ceb86b 100644 --- a/tests/ops/test_backends.py +++ b/tests/ops/test_backends.py @@ -14,6 +14,7 @@ import torch import fla.ops.common.backends.tilelang as common_tilelang_backend +import fla.ops.generalized_delta_rule.dplr.backends.tilelang as dplr_tilelang_backend import fla.ops.kda.backends.tilelang as kda_tilelang_backend import fla.ops.rwkv6.backends.tilelang as rwkv6_tilelang_backend from fla.utils import _compat @@ -100,10 +101,14 @@ def _backend_cls(backend_module): return backend_module.TileLangBackend if backend_module is rwkv6_tilelang_backend: return backend_module.RWKV6TileLangBackend - return backend_module.KDATileLangBackend + if backend_module is kda_tilelang_backend: + return backend_module.KDATileLangBackend + if backend_module is dplr_tilelang_backend: + return backend_module.DPLRTileLangBackend + raise ValueError(f"unrecognized TileLang backend module: {backend_module}") -@pytest.mark.parametrize("backend_module", [common_tilelang_backend, kda_tilelang_backend, rwkv6_tilelang_backend]) +@pytest.mark.parametrize("backend_module", [common_tilelang_backend, kda_tilelang_backend, rwkv6_tilelang_backend, dplr_tilelang_backend]) def test_tilelang_backend_gated_by_nvcc_probe(monkeypatch, backend_module): monkeypatch.setattr(backend_module, "_TILELANG_AVAILABLE", True) monkeypatch.setattr(backend_module, "has_usable_nvcc", lambda: False) @@ -113,7 +118,7 @@ def test_tilelang_backend_gated_by_nvcc_probe(monkeypatch, backend_module): assert _backend_cls(backend_module).is_available() is True -@pytest.mark.parametrize("backend_module", [common_tilelang_backend, kda_tilelang_backend, rwkv6_tilelang_backend]) +@pytest.mark.parametrize("backend_module", [common_tilelang_backend, kda_tilelang_backend, rwkv6_tilelang_backend, dplr_tilelang_backend]) def test_tilelang_backend_unavailable_without_tilelang(monkeypatch, backend_module): monkeypatch.setattr(backend_module, "_TILELANG_AVAILABLE", False) monkeypatch.setattr(backend_module, "has_usable_nvcc", lambda: True) diff --git a/tests/ops/test_dplr_tilelang.py b/tests/ops/test_dplr_tilelang.py new file mode 100644 index 0000000000..1d385a0d2d --- /dev/null +++ b/tests/ops/test_dplr_tilelang.py @@ -0,0 +1,1011 @@ +# Copyright (c) 2023-2026, Songlin Yang, Yu Zhang, Zhiyuan Li +# +# This source code is licensed under the MIT license found in the +# LICENSE file in 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 +from types import SimpleNamespace + +import pytest +import torch +import torch.nn.functional as F + +import fla.ops.generalized_delta_rule.dplr.backends.tilelang as dplr_tilelang_backend +from fla.ops.generalized_delta_rule.dplr import chunk_dplr_delta_rule +from fla.ops.generalized_delta_rule.dplr.backends.tilelang import DPLRTileLangBackend +from fla.ops.generalized_delta_rule.dplr.backends.tilelang.schedules import chunk64_schedule_or_none +from fla.ops.generalized_delta_rule.dplr.chunk import gate_bound_is_safe +from fla.ops.generalized_delta_rule.dplr.naive import dplr_recurrence +from fla.utils import assert_close, device, get_device_capability, get_device_smem_optin + +_TILELANG_USABLE = DPLRTileLangBackend.is_available() +_DISPATCH_DISABLED = os.environ.get("FLA_DISABLE_BACKEND_DISPATCH") == "1" +_CUDA_AVAILABLE = torch.cuda.is_available() + +requires_cuda = pytest.mark.skipif( + not _CUDA_AVAILABLE, + reason='verifier device queries need CUDA', +) +requires_tilelang_route = pytest.mark.skipif( + _DISPATCH_DISABLED or not _TILELANG_USABLE, + reason='TileLang backend not available or dispatch disabled', +) + + +def _cs64_routed(K: int) -> bool: + if not _CUDA_AVAILABLE: + return False + cc_major, cc_minor = get_device_capability(0) + if cc_major < 12: + # schedulable but measurably slower than Triton on pre-cc12 devices + return False + return chunk64_schedule_or_none( + K=K, V=K, in_dtype='bfloat16', + smem_cap=get_device_smem_optin(0), cc=cc_major * 10 + cc_minor, + ) is not None + + +def _varlen_cs16_routed() -> bool: + if not _CUDA_AVAILABLE: + return False + cc_major, _ = get_device_capability(0) + # schedulable but not measurably faster on pre-cc12 devices + return cc_major >= 12 + + +def _verifier_inputs( + K: int = 64, + V: int | None = None, + dtype: torch.dtype = torch.bfloat16, + gk_dtype: torch.dtype | None = None, + B: int = 4, + H: int = 32, +): + V = K if V is None else V + + def make(d, dt=dtype): + return torch.empty(B, 16, H, d, dtype=dt, device=device) + + return make(K), make(K), make(V), make(K), make(K), make(K, gk_dtype or dtype) + + +@requires_cuda +@pytest.mark.parametrize(('K', 'dtype', 'chunk_size'), [(64, torch.bfloat16, 32), (128, torch.bfloat16, 32), (64, torch.float16, 16)]) +def test_chunk_verifier_accepts(K: int, dtype: torch.dtype, chunk_size: int): + ok, reason = DPLRTileLangBackend().chunk_dplr_delta_rule_verifier( + *_verifier_inputs(K=K, dtype=dtype), safe_gate=True, chunk_size=chunk_size, + ) + assert ok and reason is None + + +@requires_cuda +def test_chunk_verifier_accepts_cp(): + # CP requires a real context: cu_seqlens set, no initial_state, no final state + cp_context = SimpleNamespace( + cu_seqlens=torch.tensor([0, 16, 32, 48, 64], dtype=torch.int32, device=device), + ) + ok, reason = DPLRTileLangBackend().chunk_dplr_delta_rule_verifier( + *_verifier_inputs(), safe_gate=True, chunk_size=32, cp_context=cp_context, + ) + assert ok and reason is None + + +@requires_cuda +def test_chunk_verifier_accepts_fp32_gk(): + # gk may stay fp32 while activations are fp16/bf16 (FLA's own test convention) + ok, reason = DPLRTileLangBackend().chunk_dplr_delta_rule_verifier( + *_verifier_inputs(gk_dtype=torch.float32), safe_gate=True, chunk_size=32, + ) + assert ok and reason is None + + +@requires_cuda +@pytest.mark.parametrize('K', [64, 128]) +def test_chunk_verifier_rejects_chunk64_on_cc90(monkeypatch, K: int): + # every BT=64 schedule launches on sm_90's 232448B optin, but cs64 measured + # 0.68-1.05x vs Triton on H800 (rect and varlen, both head dims), so the + # route is accepted only on cc12x + monkeypatch.setattr(dplr_tilelang_backend, 'get_device_smem_optin', lambda idx: 232448) + monkeypatch.setattr(dplr_tilelang_backend, 'get_device_capability', lambda idx: (9, 0)) + ok, reason = DPLRTileLangBackend().chunk_dplr_delta_rule_verifier( + *_verifier_inputs(K=K), safe_gate=True, lower_bound=-0.61, chunk_size=64, + ) + assert not ok and 'slower than Triton at chunk_size 64' in reason + + +@requires_cuda +def test_chunk_verifier_rejects_safe_gate_chunk64(monkeypatch): + # safe_gate's documented [-5, 0) range overflows the centered scheme at + # chunk_size 64 (half-range 231 log2), so the route is rejected even where + # every K=64 BT=64 stage fits (cc120's 101376B optin) + monkeypatch.setattr(dplr_tilelang_backend, 'get_device_smem_optin', lambda idx: 101376) + monkeypatch.setattr(dplr_tilelang_backend, 'get_device_capability', lambda idx: (12, 0)) + ok, reason = DPLRTileLangBackend().chunk_dplr_delta_rule_verifier( + *_verifier_inputs(K=64), safe_gate=True, chunk_size=64, + ) + assert not ok and 'requires safe_gate or a lower_bound' in reason + + +@requires_cuda +@pytest.mark.parametrize(('cc', 'accepted'), [((9, 0), False), ((12, 0), True)]) +def test_chunk_verifier_varlen_cs16_gate(monkeypatch, cc, accepted): + # varlen cs16 measured 0.83-1.02x vs Triton on sm_90 (no win at any size), + # while cc12x keeps it (1.08-1.23x at h2560/h4096) — rejected below cc12 + monkeypatch.setattr(dplr_tilelang_backend, 'get_device_smem_optin', lambda idx: 232448) + monkeypatch.setattr(dplr_tilelang_backend, 'get_device_capability', lambda idx: cc) + cu = torch.tensor([0, 256, 500, 760, 1000], dtype=torch.int32, device=device) + ok, reason = DPLRTileLangBackend().chunk_dplr_delta_rule_verifier( + *_verifier_inputs(K=64), safe_gate=True, chunk_size=16, cu_seqlens=cu, + ) + if accepted: + assert ok and reason is None + else: + assert not ok and 'chunk_size 16 on variable-length inputs' in reason + + +@requires_cuda +def test_chunk_verifier_accepts_lower_bound(): + # a caller-asserted gate bound licenses the route without safe_gate=True: + # |-1| keeps the cs32 half-range at 23 log2 + ok, reason = DPLRTileLangBackend().chunk_dplr_delta_rule_verifier( + *_verifier_inputs(), safe_gate=False, lower_bound=-1, chunk_size=32, + ) + assert ok and reason is None + + +@requires_cuda +def test_chunk_verifier_accepts_lower_bound_chunk64(monkeypatch): + # RWKV7's w is architecturally clamped to (-0.61, 0), which keeps the cs64 + # half-range at 28 log2 + monkeypatch.setattr(dplr_tilelang_backend, 'get_device_smem_optin', lambda idx: 101376) + monkeypatch.setattr(dplr_tilelang_backend, 'get_device_capability', lambda idx: (12, 0)) + ok, reason = DPLRTileLangBackend().chunk_dplr_delta_rule_verifier( + *_verifier_inputs(K=64), safe_gate=False, lower_bound=-0.61, chunk_size=64, + ) + assert ok and reason is None + + +def test_gate_bound_is_safe_exact_span(): + # the a-side operand spans chunk_size/2 + 1 rows: the exclusive cumsum + # ge[0] = 0 is centered against the inclusive gi[mid]. The helper must + # reject bounds whose true exponent overflows fp32 even where the + # chunk_size/2 estimate still passed (9*10.39*log2(e) = 135 log2) + assert not gate_bound_is_safe(-10.39, 16) + # the documented safe_gate range [-5, 0) stays licensed at cs16/cs32 + # (65 / 123 log2) and rejected at cs64 (238 log2) + assert gate_bound_is_safe(-5, 16) + assert gate_bound_is_safe(-5, 32) + assert not gate_bound_is_safe(-5, 64) + + +@requires_cuda +@pytest.mark.parametrize( + ('case', 'reason'), + [ + ('fp32', 'does not support dtype'), + ('dtype_mismatch', 'dtypes to match'), + ('kv_mismatch', 'K == V'), + ('head_dim', 'head dim'), + ('safe_gate', 'requires safe_gate or a lower_bound'), + ('lower_bound_unfit', 'requires safe_gate or a lower_bound'), + ('lower_bound_positive', 'lower_bound < 0'), + ('chunk64_a100_k128', 'no launchable backward schedule'), + ('chunk64_small_smem_k128', 'no launchable backward schedule'), + ('chunk16_k128', 'slower than Triton'), + ('chunk48', 'chunk_size'), + ('small_grid', 'small grids'), + ('low_small_grid', 'low-smem stream backward'), + ('cp_initial_state', 'initial_state with CP'), + ('cp_final_state', 'output_final_state with CP'), + ('cp_no_cu_seqlens', 'cu_seqlens for CP'), + ], +) +def test_chunk_verifier_rejects(monkeypatch, case: str, reason: str): + kwargs = {'safe_gate': True} + if case == 'fp32': + args = _verifier_inputs(dtype=torch.float32) + elif case == 'dtype_mismatch': + args = list(_verifier_inputs()) + args[1] = torch.empty_like(args[1], dtype=torch.float32) + args = tuple(args) + elif case == 'kv_mismatch': + args = _verifier_inputs(K=64, V=128) + elif case == 'head_dim': + args = _verifier_inputs(K=100, V=100) + elif case == 'safe_gate': + args = _verifier_inputs() + kwargs['safe_gate'] = False + elif case == 'lower_bound_unfit': + # |-5| pushes the cs64 half-range to 231 log2, past the fp32 limit + args = _verifier_inputs() + kwargs['safe_gate'] = False + kwargs['lower_bound'] = -5 + kwargs['chunk_size'] = 64 + elif case == 'lower_bound_positive': + # a non-negative bound is invalid (gk is a log-decay < 0) even with + # safe_gate=True, so both routes surface the same ValueError + args = _verifier_inputs() + kwargs['lower_bound'] = 0.5 + elif case == 'chunk64_a100_k128': + # high=297472B, mid=215552B and low=167936B all exceed A100's 166912B optin + monkeypatch.setattr(dplr_tilelang_backend, 'get_device_smem_optin', lambda idx: 166912) + monkeypatch.setattr(dplr_tilelang_backend, 'get_device_capability', lambda idx: (8, 0)) + args = _verifier_inputs(K=128) + kwargs['lower_bound'] = -0.61 + kwargs['chunk_size'] = 64 + elif case == 'chunk64_small_smem_k128': + # every K=128 stream schedule (mid=215552B, low=167936B) exceeds the + # 99KB cap; the fused A-backward (98432B off cc90) is not the binding + # stage here + monkeypatch.setattr(dplr_tilelang_backend, 'get_device_smem_optin', lambda idx: 101376) + monkeypatch.setattr(dplr_tilelang_backend, 'get_device_capability', lambda idx: (12, 0)) + args = _verifier_inputs(K=128) + kwargs['lower_bound'] = -0.61 + kwargs['chunk_size'] = 64 + elif case == 'chunk16_k128': + args = _verifier_inputs(K=128) + kwargs['chunk_size'] = 16 + elif case == 'small_grid': + args = _verifier_inputs(B=1, H=1) + elif case == 'low_small_grid': + # cc120-class 99KB cap forces the low stream schedule at K=128; + # N*H=64 below half the (pinned) SM count must fall back + monkeypatch.setattr(dplr_tilelang_backend, 'get_device_smem_optin', lambda idx: 101376) + monkeypatch.setattr(dplr_tilelang_backend, 'get_device_capability', lambda idx: (12, 0)) + monkeypatch.setattr(dplr_tilelang_backend, 'get_multiprocessor_count', lambda idx: 188) + args = _verifier_inputs(B=2, H=32, K=128) + elif case == 'cp_initial_state': + args = _verifier_inputs() + kwargs['initial_state'] = torch.empty(4, 32, 64, 64, device=device) + kwargs['cp_context'] = SimpleNamespace( + cu_seqlens=torch.tensor([0, 16, 32, 48, 64], dtype=torch.int32, device=device), + ) + elif case == 'cp_final_state': + args = _verifier_inputs() + kwargs['output_final_state'] = True + kwargs['cp_context'] = SimpleNamespace( + cu_seqlens=torch.tensor([0, 16, 32, 48, 64], dtype=torch.int32, device=device), + ) + elif case == 'cp_no_cu_seqlens': + args = _verifier_inputs() + kwargs['cp_context'] = SimpleNamespace() + else: + args = _verifier_inputs() + kwargs['chunk_size'] = 48 + if 'chunk_size' not in kwargs: + kwargs['chunk_size'] = 32 + ok, why = DPLRTileLangBackend().chunk_dplr_delta_rule_verifier(*args, **kwargs) + assert not ok and reason in why + + +@requires_cuda +def test_chunk_lower_bound_must_be_negative(): + with pytest.raises(ValueError, match='lower_bound'): + chunk_dplr_delta_rule(*_verifier_inputs(), safe_gate=False, lower_bound=0.5) + + +def _spy_on_tilelang_route(monkeypatch): + calls = [] + orig = DPLRTileLangBackend.chunk_dplr_delta_rule + + def spy(self, *args, **kwargs): + calls.append(None) + return orig(self, *args, **kwargs) + + monkeypatch.setattr(DPLRTileLangBackend, 'chunk_dplr_delta_rule', spy) + return calls + + +def _assert_route_parity(monkeypatch, run, names): + monkeypatch.setenv('FLA_TILELANG', '0') + ref = run() + calls = _spy_on_tilelang_route(monkeypatch) + monkeypatch.setenv('FLA_TILELANG', '1') + tri = run() + assert calls, 'TileLang backend route was not taken' + # cross-backend fp32 accumulation order differs slightly; worst observed 0.0088 + for name, r, t in zip(names, ref, tri): + assert_close(name, r, t, 0.01) + + +@requires_tilelang_route +@pytest.mark.parametrize( + ('B', 'T', 'H', 'D', 'dtype', 'chunk_size', 'disable_recompute', 'use_state'), + [ + pytest.param( + *test, + id="B{}-T{}-H{}-D{}-{}-chunk_size{}-disable_recompute{}-use_state{}".format(*test), + marks=pytest.mark.skipif( + test[5] == 64 and not _cs64_routed(test[3]), + reason='chunk_size 64 is not routed on this device', + ), + ) + for test in [ + (8, 512, 32, 64, torch.bfloat16, 32, False, True), + (8, 512, 32, 64, torch.bfloat16, 32, True, True), + (8, 512, 32, 128, torch.bfloat16, 32, False, True), + (8, 512, 32, 64, torch.bfloat16, 16, False, True), + (8, 512, 32, 64, torch.float16, 32, False, True), + (8, 63, 32, 64, torch.float16, 16, False, True), + (8, 512, 32, 64, torch.bfloat16, 64, False, True), + # RWKV7 training branch: no initial state, no final state + (8, 512, 32, 64, torch.bfloat16, 32, False, False), + (8, 512, 32, 64, torch.bfloat16, 64, False, False), + ] + ], +) +def test_chunk_tilelang_route_parity( + monkeypatch, + B: int, + T: int, + H: int, + D: int, + dtype: torch.dtype, + chunk_size: int, + disable_recompute: bool, + use_state: bool, +): + torch.manual_seed(42) + q = torch.randn(B, T, H, D, dtype=dtype) + k = torch.randn(B, T, H, D, dtype=dtype) + v = torch.randn(B, T, H, D, dtype=dtype) + a = F.normalize(torch.rand(B, T, H, D, dtype=dtype), p=2, dim=-1) + b = -a + if chunk_size == 64: + # safe_gate's [-5, 0) bound does not fit chunk_size 64; RWKV7's + # architectural clamp on w does + gk = (-0.61 * torch.sigmoid(5 * torch.randn(B, T, H, D, dtype=torch.float))).to(dtype) + else: + gk = F.logsigmoid(torch.randn(B, T, H, D, dtype=torch.float)).to(dtype).clamp(-5, 0) + h0 = torch.randn(B, H, D, D, dtype=torch.float) + q, k, v, a, b, gk, h0 = (x.to(device) for x in (q, k, v, a, b, gk, h0)) + do = torch.randn_like(v) + dht = torch.randn_like(h0) + + def run(): + q_, k_, v_, a_, b_, gk_ = (x.detach().clone().requires_grad_(True) for x in (q, k, v, a, b, gk)) + h0_ = h0.detach().clone().requires_grad_(True) if use_state else None + o, st = chunk_dplr_delta_rule( + q=q_, k=k_, v=v_, a=a_, b=b_, gk=gk_, + scale=1.0, + initial_state=h0_, + output_final_state=use_state, + safe_gate=True, + lower_bound=-0.61 if chunk_size == 64 else None, + chunk_size=chunk_size, + disable_recompute=disable_recompute, + ) + loss = (o * do).sum() + if use_state: + loss = loss + (st * dht).sum() + loss.backward() + outs = [o, q_.grad, k_.grad, v_.grad, a_.grad, b_.grad, gk_.grad] + if use_state: + outs += [st, h0_.grad] + return outs + + names = ['o', 'dq', 'dk', 'dv', 'da', 'db', 'dgk'] + (['ht', 'dh0'] if use_state else []) + _assert_route_parity(monkeypatch, run, names) + + +@requires_tilelang_route +def test_chunk_tilelang_route_parity_fp32_gk(monkeypatch): + torch.manual_seed(42) + B, T, H, D = 8, 512, 32, 64 + dtype = torch.bfloat16 + q = torch.randn(B, T, H, D, dtype=dtype) + k = torch.randn(B, T, H, D, dtype=dtype) + v = torch.randn(B, T, H, D, dtype=dtype) + a = F.normalize(torch.rand(B, T, H, D, dtype=dtype), p=2, dim=-1) + b = -a + # FLA's own tests keep gk in fp32 with bf16 activations + gk = F.logsigmoid(torch.randn(B, T, H, D, dtype=torch.float)).clamp(-5, 0) + h0 = torch.randn(B, H, D, D, dtype=torch.float) + q, k, v, a, b, gk, h0 = (x.to(device) for x in (q, k, v, a, b, gk, h0)) + do = torch.randn_like(v) + dht = torch.randn_like(h0) + + def run(): + q_, k_, v_, a_, b_, gk_, h0_ = (x.detach().clone().requires_grad_(True) for x in (q, k, v, a, b, gk, h0)) + o, st = chunk_dplr_delta_rule( + q=q_, k=k_, v=v_, a=a_, b=b_, gk=gk_, + scale=1.0, + initial_state=h0_, + output_final_state=True, + safe_gate=True, + chunk_size=32, + ) + ((o * do).sum() + (st * dht).sum()).backward() + return o, st, q_.grad, k_.grad, v_.grad, a_.grad, b_.grad, gk_.grad, h0_.grad + + _assert_route_parity(monkeypatch, run, ('o', 'ht', 'dq', 'dk', 'dv', 'da', 'db', 'dgk', 'dh0')) + + +@requires_tilelang_route +@pytest.mark.parametrize('chunk_size', [16, 32, 64]) +def test_chunk_tilelang_unsafe_gate_stays_on_triton(monkeypatch, chunk_size: int): + # safe_gate=False asserts no gate bound, so dispatch must reject the + # TileLang route and still return the Triton (sub_intra) results. The + # gates here reach magnitudes that would overflow the mid-chunk-centered + # tensor-core scheme both backends share. + torch.manual_seed(42) + B, T, H, D = 8, 512, 32, 64 + dtype = torch.bfloat16 + q = torch.randn(B, T, H, D, dtype=dtype) + k = torch.randn(B, T, H, D, dtype=dtype) + v = torch.randn(B, T, H, D, dtype=dtype) + a = F.normalize(torch.rand(B, T, H, D, dtype=dtype), p=2, dim=-1) + b = -a + gk = (20 * F.logsigmoid(torch.randn(B, T, H, D, dtype=torch.float))).to(dtype) + q, k, v, a, b, gk = (x.to(device) for x in (q, k, v, a, b, gk)) + do = torch.randn_like(v) + + def run(): + q_, k_, v_, a_, b_, gk_ = (x.detach().clone().requires_grad_(True) for x in (q, k, v, a, b, gk)) + o, _ = chunk_dplr_delta_rule( + q=q_, k=k_, v=v_, a=a_, b=b_, gk=gk_, + scale=1.0, + safe_gate=False, + chunk_size=chunk_size, + ) + (o * do).sum().backward() + return o, q_.grad, k_.grad, v_.grad, a_.grad, b_.grad, gk_.grad + + monkeypatch.setenv('FLA_TILELANG', '0') + ref = run() + calls = _spy_on_tilelang_route(monkeypatch) + monkeypatch.setenv('FLA_TILELANG', '1') + tri = run() + assert not calls, 'TileLang route must stay rejected for safe_gate=False' + assert torch.isfinite(tri[0]).all() + for name, r, t in zip(('o', 'dq', 'dk', 'dv', 'da', 'db', 'dgk'), ref, tri): + assert_close(name, r, t, 0.01) + + +@requires_tilelang_route +@pytest.mark.parametrize('chunk_size', [16, 32]) +def test_chunk_tilelang_lower_bound_takes_tilelang_route(monkeypatch, chunk_size: int): + # safe_gate=False with a fitting lower_bound must take the TileLang route; + # the Triton reference canonicalizes to the same centered tensor-core + # scheme, so both sides compute identical math here + torch.manual_seed(42) + B, T, H, D = 8, 512, 32, 64 + dtype = torch.bfloat16 + q = torch.randn(B, T, H, D, dtype=dtype) + k = torch.randn(B, T, H, D, dtype=dtype) + v = torch.randn(B, T, H, D, dtype=dtype) + a = F.normalize(torch.rand(B, T, H, D, dtype=dtype), p=2, dim=-1) + b = -a + gk = F.logsigmoid(torch.randn(B, T, H, D, dtype=torch.float)).to(dtype).clamp(-1, 0) + q, k, v, a, b, gk = (x.to(device) for x in (q, k, v, a, b, gk)) + do = torch.randn_like(v) + + def run(): + q_, k_, v_, a_, b_, gk_ = (x.detach().clone().requires_grad_(True) for x in (q, k, v, a, b, gk)) + o, _ = chunk_dplr_delta_rule( + q=q_, k=k_, v=v_, a=a_, b=b_, gk=gk_, + scale=1.0, + safe_gate=False, + lower_bound=-1, + chunk_size=chunk_size, + ) + (o * do).sum().backward() + return o, q_.grad, k_.grad, v_.grad, a_.grad, b_.grad, gk_.grad + + _assert_route_parity(monkeypatch, run, ('o', 'dq', 'dk', 'dv', 'da', 'db', 'dgk')) + + +@requires_tilelang_route +@pytest.mark.parametrize( + ('B', 'T', 'H', 'D', 'dtype', 'chunk_size', 'gate_style'), + [ + pytest.param( + *test, + id="B{}-T{}-H{}-D{}-{}-chunk_size{}-{}".format(*test), + marks=pytest.mark.skipif( + test[5] == 64 and not _cs64_routed(test[3]), + reason='chunk_size 64 is not routed on this device', + ), + ) + for test in [ + # gates pushed toward the documented -5 bound: gk is natural-log + # decay, so the mid-chunk-centered exponents reach up to + # (BT/2)*5*log2(e) = 115 (BT=32) / 58 (BT=16) log2, the worst + # spread the safe_gate contract licenses at these chunk sizes + (8, 512, 32, 64, torch.bfloat16, 32, 'saturated_lb'), + (8, 512, 32, 64, torch.bfloat16, 16, 'saturated_lb'), + (8, 512, 32, 128, torch.bfloat16, 32, 'saturated_lb'), + # RWKV7's w is architecturally clamped to (-0.61, 0), which keeps + # the BT=64 half-range at 28 log2; saturated to push both ends + (8, 512, 32, 64, torch.bfloat16, 64, 'rwkv7'), + (8, 512, 32, 128, torch.bfloat16, 64, 'rwkv7'), + ] + ], +) +def test_chunk_tilelang_route_parity_gate_stress( + monkeypatch, + B: int, + T: int, + H: int, + D: int, + dtype: torch.dtype, + chunk_size: int, + gate_style: str, +): + torch.manual_seed(42) + q = torch.randn(B, T, H, D, dtype=dtype) + k = torch.randn(B, T, H, D, dtype=dtype) + v = torch.randn(B, T, H, D, dtype=dtype) + a = F.normalize(torch.rand(B, T, H, D, dtype=dtype), p=2, dim=-1) + b = -a + if gate_style == 'saturated_lb': + # KDA's convention (tests/ops/test_kda.py): amplify the logits so the + # clamp pins a fraction of gates at the documented -5 bound. Sustained + # pinning of most gates is excluded on purpose: at cs32 the Triton + # reference's dgk diverges from the true recurrence far beyond parity + # tolerance there (measured err ratio ~0.5 vs TileLang's ~0.17 on + # sm_120), so parity is undefined + gk = (F.logsigmoid(torch.randn(B, T, H, D, dtype=torch.float)) / 0.6).clamp(-5, 0).to(dtype) + else: + # sigmoid saturates to 0/1, pinning gk at the ends of (-0.61, 0) + gk = (-0.61 * torch.sigmoid(5 * torch.randn(B, T, H, D, dtype=torch.float))).to(dtype) + h0 = torch.randn(B, H, D, D, dtype=torch.float) + q, k, v, a, b, gk, h0 = (x.to(device) for x in (q, k, v, a, b, gk, h0)) + do = torch.randn_like(v) + dht = torch.randn_like(h0) + + def run(): + q_, k_, v_, a_, b_, gk_, h0_ = (x.detach().clone().requires_grad_(True) for x in (q, k, v, a, b, gk, h0)) + o, st = chunk_dplr_delta_rule( + q=q_, k=k_, v=v_, a=a_, b=b_, gk=gk_, + scale=1.0, + initial_state=h0_, + output_final_state=True, + safe_gate=True, + lower_bound=-0.61 if chunk_size == 64 else None, + chunk_size=chunk_size, + ) + ((o * do).sum() + (st * dht).sum()).backward() + return o, st, q_.grad, k_.grad, v_.grad, a_.grad, b_.grad, gk_.grad, h0_.grad + + _assert_route_parity(monkeypatch, run, ('o', 'ht', 'dq', 'dk', 'dv', 'da', 'db', 'dgk', 'dh0')) + + +@requires_tilelang_route +def test_chunk_tilelang_safe_gate_chunk64_stays_on_triton(monkeypatch): + # safe_gate's documented [-5, 0) range overflows the centered scheme at + # chunk_size 64, so dispatch must reject the route regardless of the + # (here mild) gates actually supplied + torch.manual_seed(42) + B, T, H, D = 8, 512, 32, 64 + dtype = torch.bfloat16 + q = torch.randn(B, T, H, D, dtype=dtype) + k = torch.randn(B, T, H, D, dtype=dtype) + v = torch.randn(B, T, H, D, dtype=dtype) + a = F.normalize(torch.rand(B, T, H, D, dtype=dtype), p=2, dim=-1) + b = -a + gk = F.logsigmoid(torch.randn(B, T, H, D, dtype=torch.float)).to(dtype).clamp(-5, 0) + q, k, v, a, b, gk = (x.to(device) for x in (q, k, v, a, b, gk)) + do = torch.randn_like(v) + + def run(): + q_, k_, v_, a_, b_, gk_ = (x.detach().clone().requires_grad_(True) for x in (q, k, v, a, b, gk)) + o, _ = chunk_dplr_delta_rule( + q=q_, k=k_, v=v_, a=a_, b=b_, gk=gk_, + scale=1.0, + safe_gate=True, + chunk_size=64, + ) + (o * do).sum().backward() + return o, q_.grad, k_.grad, v_.grad, a_.grad, b_.grad, gk_.grad + + monkeypatch.setenv('FLA_TILELANG', '0') + ref = run() + calls = _spy_on_tilelang_route(monkeypatch) + monkeypatch.setenv('FLA_TILELANG', '1') + tri = run() + assert not calls, 'TileLang route must stay rejected for safe_gate=True with chunk_size=64' + for name, r, t in zip(('o', 'dq', 'dk', 'dv', 'da', 'db', 'dgk'), ref, tri): + assert_close(name, r, t, 0.01) + + +@requires_tilelang_route +@pytest.mark.parametrize( + ('chunk_size', 'pin'), + [ + pytest.param( + chunk_size, + pin, + id=f'chunk_size{chunk_size}-pin{pin}', + marks=pytest.mark.skipif( + chunk_size == 64 and not _cs64_routed(64), + reason='chunk_size 64 is not routed on this device', + ), + ) + for chunk_size, pin in [(16, -5.0), (32, -5.0), (64, -2.0), (64, -0.61)] + ], +) +def test_chunk_tilelang_sustained_gate_pin_stress(monkeypatch, chunk_size: int, pin: float): + # every gate pinned at the bound for the whole sequence: the licensed + # half-range must stay finite and track the fp32 recurrence. dgk is + # ill-conditioned in this regime for any chunked scheme (RMS err ratio vs + # the baseline at the cs32 -5 pin, measured on sm_120: TileLang ~0.17, + # the Triton reference ~0.5), so parity is asserted against the baseline + # with a relaxed dgk tolerance; all other outputs match the usual + # naive-parity bars + torch.manual_seed(42) + B, T, H, D = 8, 512, 16, 64 + dtype = torch.bfloat16 + q = torch.randn(B, T, H, D, dtype=dtype) + k = torch.randn(B, T, H, D, dtype=dtype) + v = torch.randn(B, T, H, D, dtype=dtype) + a = F.normalize(torch.rand(B, T, H, D, dtype=dtype), p=2, dim=-1) + b = -a + gk = torch.full((B, T, H, D), pin, dtype=dtype) + h0 = torch.randn(B, H, D, D, dtype=torch.float) + q, k, v, a, b, gk, h0 = (x.to(device) for x in (q, k, v, a, b, gk, h0)) + do = torch.randn_like(v) + dht = torch.randn_like(h0) + + def run(fn): + q_, k_, v_, a_, b_, gk_, h0_ = (x.detach().clone().requires_grad_(True) for x in (q, k, v, a, b, gk, h0)) + if fn == 'naive': + o, st = _naive_recurrence(q_, k_, v_, a_, b_, gk_, h0_) + else: + o, st = chunk_dplr_delta_rule( + q=q_, k=k_, v=v_, a=a_, b=b_, gk=gk_, + initial_state=h0_, output_final_state=True, + safe_gate=True, lower_bound=pin if chunk_size == 64 else None, + chunk_size=chunk_size, + ) + ((o * do).sum() + (st * dht).sum()).backward() + return o, st, q_.grad, k_.grad, v_.grad, a_.grad, b_.grad, gk_.grad, h0_.grad + + calls = _spy_on_tilelang_route(monkeypatch) + monkeypatch.setenv('FLA_TILELANG', '1') + tri = run('op') + assert calls, 'TileLang backend route was not taken' + for x in tri: + assert torch.isfinite(x).all() + ref = run('naive') + for name, r, t in zip(_NAIVE_PARITY_NAMES, ref, tri): + assert_close(name, r, t, 0.007 if name == 'o' else (0.25 if name == 'dgk' else 0.008)) + + +@requires_tilelang_route +@pytest.mark.parametrize( + ('H', 'D', 'cu_seqlens', 'dtype', 'chunk_size', 'use_state'), + [ + pytest.param( + *test, + id="H{}-D{}-cu_seqlens{}-{}-chunk_size{}-use_state{}".format(*test), + marks=pytest.mark.skipif( + test[4] == 16 and not _varlen_cs16_routed(), + reason='varlen chunk_size 16 is not routed on this device', + ), + ) + for test in [ + (32, 64, [0, 256, 500, 760, 1000], torch.bfloat16, 32, True), + (32, 128, [0, 256, 500, 760, 1000], torch.bfloat16, 32, True), + (32, 64, [0, 256, 500, 760, 1000], torch.bfloat16, 16, True), + (32, 64, [0, 256, 500, 760, 1000], torch.float16, 32, True), + (32, 64, [0, 256, 500, 760, 1000], torch.bfloat16, 32, False), + ] + ], +) +def test_chunk_varlen_tilelang_route_parity( + monkeypatch, + H: int, + D: int, + cu_seqlens: list[int], + dtype: torch.dtype, + chunk_size: int, + use_state: bool, +): + torch.manual_seed(42) + N = len(cu_seqlens) - 1 + T = cu_seqlens[-1] + cu_seqlens = torch.tensor(cu_seqlens, dtype=torch.int32, device=device) + + q = torch.randn(1, T, H, D, dtype=dtype) + k = torch.randn(1, T, H, D, dtype=dtype) + v = torch.randn(1, T, H, D, dtype=dtype) + a = F.normalize(torch.rand(1, T, H, D, dtype=dtype), p=2, dim=-1) + b = -a + gk = F.logsigmoid(torch.randn(1, T, H, D, dtype=torch.float)).to(dtype).clamp(-5, 0) + h0 = torch.randn(N, H, D, D, dtype=torch.float) + q, k, v, a, b, gk, h0 = (x.to(device) for x in (q, k, v, a, b, gk, h0)) + do = torch.randn_like(v) + dht = torch.randn_like(h0) + + def run(): + q_, k_, v_, a_, b_, gk_ = (x.detach().clone().requires_grad_(True) for x in (q, k, v, a, b, gk)) + h0_ = h0.detach().clone().requires_grad_(True) if use_state else None + o, st = chunk_dplr_delta_rule( + q=q_, k=k_, v=v_, a=a_, b=b_, gk=gk_, + initial_state=h0_, + output_final_state=use_state, + cu_seqlens=cu_seqlens, + safe_gate=True, + chunk_size=chunk_size, + ) + loss = (o * do).sum() + if use_state: + loss = loss + (st * dht).sum() + loss.backward() + outs = [o, q_.grad, k_.grad, v_.grad, a_.grad, b_.grad, gk_.grad] + if use_state: + outs += [st, h0_.grad] + return outs + + names = ['o', 'dq', 'dk', 'dv', 'da', 'db', 'dgk'] + (['ht', 'dh0'] if use_state else []) + _assert_route_parity(monkeypatch, run, names) + + +def _naive_recurrence(q, k, v, a, b, gk, h0): + # per-token fp32 PyTorch baseline (no chunk math); dplr_recurrence works + # on [B, H, T, D] and applies the K**-0.5 scale internally + o, st = dplr_recurrence( + q.transpose(1, 2), k.transpose(1, 2), v.transpose(1, 2), + a.transpose(1, 2), b.transpose(1, 2), gk.transpose(1, 2), + initial_state=h0, output_final_state=True, + ) + return o.transpose(1, 2), st + + +def _assert_naive_parity(monkeypatch, inputs, do, dht, names, call): + q, k, v, a, b, gk, h0 = inputs + + def run(fn): + q_, k_, v_, a_, b_, gk_ = (x.detach().clone().requires_grad_(True) for x in (q, k, v, a, b, gk)) + h0_ = h0.detach().clone().requires_grad_(True) + o, st = call(fn, q_, k_, v_, a_, b_, gk_, h0_) + ((o * do).sum() + (st * dht).sum()).backward() + return o, st, q_.grad, k_.grad, v_.grad, a_.grad, b_.grad, gk_.grad, h0_.grad + + calls = _spy_on_tilelang_route(monkeypatch) + monkeypatch.setenv('FLA_TILELANG', '1') + tri = run('op') + assert calls, 'TileLang backend route was not taken' + ref = run('naive') + # same ratios test_dplr_delta uses for the Triton path vs the fp32 ref + for name, r, t in zip(names, ref, tri): + assert_close(name, r, t, 0.007 if name == 'o' else 0.008) + + +_NAIVE_PARITY_NAMES = ('o', 'ht', 'dq', 'dk', 'dv', 'da', 'db', 'dgk', 'dh0') + + +@requires_tilelang_route +@pytest.mark.parametrize( + ('T', 'D', 'dtype', 'chunk_size', 'gate_style'), + [ + pytest.param( + *test, + id="T{}-D{}-{}-chunk_size{}-{}".format(*test), + marks=pytest.mark.skipif( + test[3] == 64 and not _cs64_routed(test[1]), + reason='chunk_size 64 is not routed on this device', + ), + ) + for test in [ + (256, 64, torch.bfloat16, 32, 'standard'), + (256, 64, torch.bfloat16, 16, 'standard'), + (256, 128, torch.bfloat16, 32, 'standard'), + (256, 64, torch.float16, 32, 'standard'), + (256, 64, torch.bfloat16, 64, 'standard'), + (256, 64, torch.bfloat16, 32, 'saturated_lb'), + (256, 128, torch.bfloat16, 64, 'rwkv7'), + ] + ], +) +def test_chunk_tilelang_naive_ref_parity( + monkeypatch, + T: int, + D: int, + dtype: torch.dtype, + chunk_size: int, + gate_style: str, +): + torch.manual_seed(42) + B, H = 8, 32 + q = torch.randn(B, T, H, D, dtype=dtype) + k = torch.randn(B, T, H, D, dtype=dtype) + v = torch.randn(B, T, H, D, dtype=dtype) + a = F.normalize(torch.rand(B, T, H, D, dtype=dtype), p=2, dim=-1) + b = -a + if gate_style == 'saturated_lb': + gk = (F.logsigmoid(torch.randn(B, T, H, D, dtype=torch.float)) / 0.6).clamp(-5, 0).to(dtype) + elif gate_style == 'rwkv7': + gk = (-0.61 * torch.sigmoid(5 * torch.randn(B, T, H, D, dtype=torch.float))).to(dtype) + else: + gk = F.logsigmoid(torch.randn(B, T, H, D, dtype=torch.float)).to(dtype).clamp(-5, 0) + h0 = torch.randn(B, H, D, D, dtype=torch.float) + q, k, v, a, b, gk, h0 = (x.to(device) for x in (q, k, v, a, b, gk, h0)) + do = torch.randn_like(v) + dht = torch.randn_like(h0) + + def call(fn, q_, k_, v_, a_, b_, gk_, h0_): + if fn == 'naive': + return _naive_recurrence(q_, k_, v_, a_, b_, gk_, h0_) + return chunk_dplr_delta_rule( + q=q_, k=k_, v=v_, a=a_, b=b_, gk=gk_, + initial_state=h0_, output_final_state=True, + safe_gate=True, lower_bound=-0.61 if chunk_size == 64 else None, + chunk_size=chunk_size, + ) + + _assert_naive_parity(monkeypatch, (q, k, v, a, b, gk, h0), do, dht, _NAIVE_PARITY_NAMES, call) + + +@requires_tilelang_route +@pytest.mark.parametrize( + ('D', 'cu_seqlens', 'dtype', 'chunk_size'), + [ + pytest.param(*test, id="D{}-cu_seqlens{}-{}-chunk_size{}".format(*test)) + for test in [ + (64, [0, 130, 260, 390, 512], torch.bfloat16, 32), + (128, [0, 130, 260, 390, 512], torch.bfloat16, 32), + ] + ], +) +def test_chunk_varlen_tilelang_naive_ref_parity( + monkeypatch, + D: int, + cu_seqlens: list[int], + dtype: torch.dtype, + chunk_size: int, +): + torch.manual_seed(42) + N = len(cu_seqlens) - 1 + T = cu_seqlens[-1] + H = 32 + cu = torch.tensor(cu_seqlens, dtype=torch.int32, device=device) + + q = torch.randn(1, T, H, D, dtype=dtype) + k = torch.randn(1, T, H, D, dtype=dtype) + v = torch.randn(1, T, H, D, dtype=dtype) + a = F.normalize(torch.rand(1, T, H, D, dtype=dtype), p=2, dim=-1) + b = -a + gk = F.logsigmoid(torch.randn(1, T, H, D, dtype=torch.float)).to(dtype).clamp(-5, 0) + h0 = torch.randn(N, H, D, D, dtype=torch.float) + q, k, v, a, b, gk, h0 = (x.to(device) for x in (q, k, v, a, b, gk, h0)) + do = torch.randn_like(v) + dht = torch.randn_like(h0) + + def call(fn, q_, k_, v_, a_, b_, gk_, h0_): + if fn == 'naive': + # the baseline has no varlen form; run each sequence and repack + outs, sts = [], [] + for i in range(N): + s, e = cu_seqlens[i], cu_seqlens[i + 1] + o_i, st_i = _naive_recurrence( + q_[:, s:e], k_[:, s:e], v_[:, s:e], + a_[:, s:e], b_[:, s:e], gk_[:, s:e], h0_[i: i + 1], + ) + outs.append(o_i) + sts.append(st_i) + return torch.cat(outs, 1), torch.cat(sts, 0) + return chunk_dplr_delta_rule( + q=q_, k=k_, v=v_, a=a_, b=b_, gk=gk_, + initial_state=h0_, output_final_state=True, + cu_seqlens=cu, safe_gate=True, chunk_size=chunk_size, + ) + + _assert_naive_parity(monkeypatch, (q, k, v, a, b, gk, h0), do, dht, _NAIVE_PARITY_NAMES, call) + + +@requires_tilelang_route +def test_chunk_tilelang_fwd_opcheck(): + torch.manual_seed(42) + B, T, H, D = 2, 64, 8, 64 + dtype = torch.bfloat16 + q = torch.randn(B, T, H, D, dtype=dtype, device=device) + k = torch.randn(B, T, H, D, dtype=dtype, device=device) + v = torch.randn(B, T, H, D, dtype=dtype, device=device) + a = F.normalize(torch.rand(B, T, H, D, dtype=dtype, device=device), p=2, dim=-1) + b = -a + gk = F.logsigmoid(torch.randn(B, T, H, D, dtype=torch.float, device=device)).clamp(-5, 0) + h0 = torch.randn(B, H, D, D, dtype=torch.float, device=device) + cu = torch.empty((0,), dtype=torch.int32, device=device) + + torch.library.opcheck( + torch.ops.fla.chunk_dplr_delta_rule_fwd, + (q, k, v, a, b, gk, h0, cu, 1.0, True, True, False, 32), + ) + + +@requires_tilelang_route +def test_chunk_tilelang_torch_compile_fullgraph_smoke(): + from fla.ops.generalized_delta_rule.dplr.backends.tilelang.chunk import chunk_dplr_delta_rule_tilelang + + torch.manual_seed(42) + B, T, H, D = 2, 256, 8, 64 + dtype = torch.bfloat16 + q = torch.randn(B, T, H, D, dtype=dtype, device=device) + k = torch.randn(B, T, H, D, dtype=dtype, device=device) + v = torch.randn(B, T, H, D, dtype=dtype, device=device) + a = F.normalize(torch.rand(B, T, H, D, dtype=dtype, device=device), p=2, dim=-1) + b = -a + gk = F.logsigmoid(torch.randn(B, T, H, D, dtype=torch.float, device=device)).clamp(-5, 0).to(dtype) + h0 = torch.randn(B, H, D, D, dtype=torch.float, device=device) + + def fn(q, k, v, a, b, gk, h0): + return chunk_dplr_delta_rule_tilelang( + q=q, k=k, v=v, a=a, b=b, gk=gk, + scale=1.0, + initial_state=h0, + output_final_state=True, + safe_gate=True, + chunk_size=32, + ) + + ref_o, ref_st = fn(q, k, v, a, b, gk, h0) + compiled = torch.compile(fn, fullgraph=True) + tri_o, tri_st = compiled(q, k, v, a, b, gk, h0) + assert_close('o', ref_o, tri_o, 0.005) + assert_close('ht', ref_st, tri_st, 0.005) + + +@requires_tilelang_route +def test_chunk_tilelang_checkpoint_elision(monkeypatch): + from fla.ops.generalized_delta_rule.dplr.backends.tilelang import chunk as tl_chunk + from fla.ops.generalized_delta_rule.dplr.backends.tilelang.chunk import ( + chunk_dplr_delta_rule_tilelang, + dplr_checkpoint_context_fn, + ) + + torch.manual_seed(42) + B, T, H, D = 2, 256, 8, 64 + dtype = torch.bfloat16 + q = torch.randn(B, T, H, D, dtype=dtype, device=device) + k = torch.randn(B, T, H, D, dtype=dtype, device=device) + v = torch.randn(B, T, H, D, dtype=dtype, device=device) + a = F.normalize(torch.rand(B, T, H, D, dtype=dtype, device=device), p=2, dim=-1) + b = -a + gk = F.logsigmoid(torch.randn(B, T, H, D, dtype=torch.float, device=device)).clamp(-5, 0).to(dtype) + h0 = torch.randn(B, H, D, D, dtype=torch.float, device=device) + do = torch.randn_like(v) + dht = torch.randn_like(h0) + + elided_calls = [] + orig_elided = tl_chunk._chunk_dplr_delta_rule_fwd_ctx_elided_op + + def elided_spy(*args, **kwargs): + elided_calls.append(None) + return orig_elided(*args, **kwargs) + + monkeypatch.setattr(tl_chunk, '_chunk_dplr_delta_rule_fwd_ctx_elided_op', elided_spy) + + def fwd(q, k, v, a, b, gk, h0): + return chunk_dplr_delta_rule_tilelang( + q=q, k=k, v=v, a=a, b=b, gk=gk, + scale=1.0, + initial_state=h0, + output_final_state=True, + safe_gate=True, + chunk_size=32, + disable_recompute=True, + ) + + def run(use_checkpoint): + leaves = (x.detach().clone().requires_grad_(True) for x in (q, k, v, a, b, gk, h0)) + q_, k_, v_, a_, b_, gk_, h0_ = leaves + if use_checkpoint: + o, st = torch.utils.checkpoint.checkpoint( + fwd, q_, k_, v_, a_, b_, gk_, h0_, + use_reentrant=False, + context_fn=dplr_checkpoint_context_fn, + ) + else: + o, st = fwd(q_, k_, v_, a_, b_, gk_, h0_) + ((o * do).sum() + (st * dht).sum()).backward() + return o, st, q_.grad, k_.grad, v_.grad, a_.grad, b_.grad, gk_.grad, h0_.grad + + ref = run(use_checkpoint=False) + tri = run(use_checkpoint=True) + assert elided_calls, 'checkpoint forward did not take the ctx-elided route' + for name, r, t in zip(('o', 'ht', 'dq', 'dk', 'dv', 'da', 'db', 'dgk', 'dh0'), ref, tri): + assert_close(name, r, t, 0.01)