diff --git a/miles_plugins/models/linear_attn.py b/miles_plugins/models/linear_attn.py index d2539be7d36..cdb2548894d 100644 --- a/miles_plugins/models/linear_attn.py +++ b/miles_plugins/models/linear_attn.py @@ -13,6 +13,9 @@ from __future__ import annotations +import functools +import logging +import os from abc import ABC, abstractmethod from typing import NamedTuple @@ -42,6 +45,8 @@ except ImportError: pass +logger = logging.getLogger(__name__) + WEIGHT_LAYOUT_VERSION = 1 INT32_ELEMENTS = 2**31 - 1 _CHUNK_ELEMENTS = 2**30 @@ -68,11 +73,17 @@ def gdn_kernel(backend: str): raise ValueError(f"Unsupported GDN backend: {backend}") +@functools.cache def kda_kernel(): + """fla's chunk_kda. On Blackwell its backward runs the Triton dqkg kernel (FLA_TILELANG=0), which + beats fla's TileLang one there at every tiling; an explicit FLA_TILELANG wins.""" try: from fla.ops.kda import chunk_kda except ImportError as exc: raise ImportError("KDA requires flash-linear-attention >= 0.5 (fla.ops.kda).") from exc + if torch.cuda.is_available() and torch.cuda.get_device_capability()[0] >= 10: + os.environ.setdefault("FLA_TILELANG", "0") + logger.info(f"KDA backward: FLA_TILELANG={os.environ.get('FLA_TILELANG', 'unset')}") return chunk_kda @@ -124,6 +135,22 @@ def kda_recurrence(q, k, v, beta_logits, decay, A_log, dt_bias, *, gate_lower_bo return out +@functools.cache +def short_conv_backend() -> str: + """fla's "mix" (Triton forward, causal-conv1d CUDA backward) when causal-conv1d's backward imports, + Triton otherwise; ``FLA_CONV_BACKEND`` overrides, as it does for fla's own ShortConvolution.""" + backend = os.environ.get("FLA_CONV_BACKEND") + if backend is None: + try: + from causal_conv1d.cpp_functions import causal_conv1d_bwd_function # noqa: F401 + + backend = "mix" + except ImportError: + backend = "triton" + logger.info(f"Linear-attention short conv backend: {backend}") + return backend + + class _ContiguousGrad(torch.autograd.Function): """Identity whose backward hands on a contiguous gradient. The gradient of one channel chunk of a concatenation is a strided view whose row stride is the full width; fla's conv backward indexes it @@ -178,7 +205,7 @@ def __init__(self, channels: int, kernel_size: int, tp_group, device=None, dtype dtype=dtype, ) self.activation = "silu" - self.backend = "triton" + self.backend = short_conv_backend() self.tp_group = tp_group set_tensor_model_parallel_attributes(self.weight, True, 0, 1) diff --git a/tests/fast-gpu/test_linear_attn_short_conv.py b/tests/fast-gpu/test_linear_attn_short_conv.py index 4180dc74aea..bd453d3eb2e 100644 --- a/tests/fast-gpu/test_linear_attn_short_conv.py +++ b/tests/fast-gpu/test_linear_attn_short_conv.py @@ -1,12 +1,14 @@ """The head-sharded short conv splits inputs past fla's int32 indexing limit into channel chunks. The conv is per channel, so the chunked forward matches the single call bit for bit. fla picks its backward tiling from the channel count, so on Blackwell the gradients differ from the single call by fp32 -accumulation order (~1e-6 relative); a wrong channel mapping is off by ~1.""" +accumulation order (~1e-6 relative); a wrong channel mapping is off by ~1. The conv, on whichever +backend the layer selects, also matches a float64 torch reference to bf16 rounding (< 5e-3).""" import sys import pytest import torch +import torch.nn.functional as F from tests.ci.ci_register import register_cuda_ci register_cuda_ci(est_time=60, suite="stage-b-2-gpu-h200", labels=["precision"], hardware=["hopper", "blackwell"]) @@ -15,14 +17,47 @@ pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") from miles_plugins.models import linear_attn # noqa: E402 -from miles_plugins.models.linear_attn import ShardedShortConv # noqa: E402 + + +def _conv(channels): + return linear_attn.ShardedShortConv(channels, 4, tp_group=None, device="cuda", dtype=torch.bfloat16) + + +def _reference(x, weight, cu_seqlens): + x = x.double().transpose(1, 2) + weight = weight.double() + segments = [ + F.silu(F.conv1d(F.pad(x[..., start:end], (weight.shape[-1] - 1, 0)), weight, groups=weight.shape[0])) + for start, end in zip(cu_seqlens[:-1].tolist(), cu_seqlens[1:].tolist(), strict=True) + ] + return torch.cat(segments, dim=-1).transpose(1, 2) + + +@pytest.mark.parametrize("backend", sorted({"triton", linear_attn.short_conv_backend()})) +def test_conv_matches_float64_reference(backend): + torch.manual_seed(0) + tokens, channels = 2048, 768 + conv = _conv(channels) + conv.backend = backend + x = torch.randn(1, tokens, channels, device="cuda", dtype=torch.bfloat16) + grad = torch.randn_like(x) + cu_seqlens = torch.tensor([0, 700, tokens], device="cuda", dtype=torch.int32) + xi = x.clone().requires_grad_() + conv(xi, cu_seqlens=cu_seqlens).backward(grad) + xr = x.double().requires_grad_() + wr = conv.weight.detach().double().requires_grad_() + out_ref = _reference(xr, wr, cu_seqlens) + out_ref.backward(grad.double()) + out = conv(x, cu_seqlens=cu_seqlens) + for actual, expected in ((out, out_ref), (xi.grad, xr.grad), (conv.weight.grad, wr.grad)): + assert ((actual.double() - expected).norm() / expected.norm()).item() < 5e-3 @pytest.mark.parametrize("channels", [768, 1000]) def test_chunked_conv_matches_single_call(monkeypatch, channels): torch.manual_seed(0) tokens = 2048 - conv = ShardedShortConv(channels, 4, tp_group=None, device="cuda", dtype=torch.bfloat16) + conv = _conv(channels) x = torch.randn(1, tokens, channels, device="cuda", dtype=torch.bfloat16) grad = torch.randn_like(x) cu_seqlens = torch.tensor([0, 700, tokens], device="cuda", dtype=torch.int32) diff --git a/tests/fast/models/test_linear_attn_kernel_selection.py b/tests/fast/models/test_linear_attn_kernel_selection.py index 0f39e3e9076..0c078acd0e3 100644 --- a/tests/fast/models/test_linear_attn_kernel_selection.py +++ b/tests/fast/models/test_linear_attn_kernel_selection.py @@ -3,14 +3,38 @@ The equivalence tests need a Hopper (SM90+) GPU with both `fla` and `flash_qla`; they skip otherwise. """ +import os +import sys +import types + import pytest -from miles_plugins.models.linear_attn import gdn_kernel +from miles_plugins.models import linear_attn def test_unknown_backend_raises_value_error(): with pytest.raises(ValueError, match="Unsupported GDN backend"): - gdn_kernel("nope") + linear_attn.gdn_kernel("nope") + + +@pytest.mark.parametrize( + "capability,user_value,expected", [((10, 0), None, "0"), ((9, 0), None, None), ((10, 0), "1", "1")] +) +def test_kda_uses_the_triton_backward_on_blackwell_unless_set(monkeypatch, capability, user_value, expected): + torch = pytest.importorskip("torch") + monkeypatch.setitem(sys.modules, "fla.ops.kda", types.SimpleNamespace(chunk_kda="chunk_kda")) + monkeypatch.setattr(torch.cuda, "is_available", lambda: True) + monkeypatch.setattr(torch.cuda, "get_device_capability", lambda: capability) + monkeypatch.delenv("FLA_TILELANG", raising=False) + if user_value is not None: + monkeypatch.setenv("FLA_TILELANG", user_value) + assert linear_attn.kda_kernel.__wrapped__() == "chunk_kda" + assert os.environ.get("FLA_TILELANG") == expected + + +def test_short_conv_backend_respects_fla_conv_backend(monkeypatch): + monkeypatch.setenv("FLA_CONV_BACKEND", "cuda") + assert linear_attn.short_conv_backend.__wrapped__() == "cuda" NUM_HEADS = 4 @@ -75,8 +99,8 @@ def test_fla_flashqla_equivalence(dtype_name, atol, rtol): torch = _require_backends() dtype = getattr(torch, dtype_name) - fla_kernel = gdn_kernel("fla") - flashqla_kernel = gdn_kernel("flashqla") + fla_kernel = linear_attn.gdn_kernel("fla") + flashqla_kernel = linear_attn.gdn_kernel("flashqla") query, key, value, g, beta, cu = _make_inputs(torch, dtype, device="cuda")