[8/8] perf(models): per-arch short-conv and KDA backward kernels for linear attention

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)
 

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)

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")