Skip to content

Commit bbd6c36

Browse files
kuoihaokuoihao
committed
[Ops] Add TLE and FusedInfer backends for KDA inference
- Add TLE backend with BT=16 inference kernels for chunk_kda - Supports fp16/bf16, K in {64,128,192,256}, arbitrary V - Fuses q/k L2 norm, beta sigmoid, and KDA safe gate - Priority 2 (highest) - Add FusedInfer backend (Triton) for chunk_kda inference - Priority 4 - Register both backends in kda backend registry - Update test_chunk_fwd_inference with 10 comprehensive test cases - 5 fixed-length cases covering various B/T/H/D/dtype - 5 varlen cases covering different sequence distributions - All test cases compatible with TLE requirements
1 parent e8e791f commit bbd6c36

6 files changed

Lines changed: 1851 additions & 0 deletions

File tree

‎fla/ops/kda/backends/__init__.py‎

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -9,11 +9,15 @@
99

1010
from fla.ops.backends import BackendRegistry, dispatch
1111
from fla.ops.kda.backends.flashkda import FlashKDABackend
12+
from fla.ops.kda.backends.fused_infer import KDAFusedInferBackend
1213
from fla.ops.kda.backends.tilelang import KDATileLangBackend
14+
from fla.ops.kda.backends.tle import KDATLEBackend
1315

1416
kda_registry = BackendRegistry("kda")
1517
kda_registry.register(FlashKDABackend())
1618
kda_registry.register(KDATileLangBackend())
19+
kda_registry.register(KDAFusedInferBackend())
20+
kda_registry.register(KDATLEBackend())
1721

1822

1923
__all__ = ['dispatch', 'kda_registry']

‎fla/ops/kda/backends/fused_infer.py‎

Lines changed: 115 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,115 @@
1+
# Copyright (c) 2023-2026, Songlin Yang, Yu Zhang, Zhiyuan Li
2+
#
3+
# This source code is licensed under the MIT license found in the
4+
# LICENSE file in the root directory of this source tree.
5+
# For a list of all contributors, visit:
6+
# https://github.com/fla-org/flash-linear-attention/graphs/contributors
7+
#
8+
# Copyright (c) 2026 FlagOS Contributors
9+
10+
"""Fused Triton inference backend for KDA operations."""
11+
12+
from __future__ import annotations
13+
14+
from typing import TYPE_CHECKING
15+
16+
import torch
17+
18+
from fla.ops.backends import BaseBackend
19+
20+
if TYPE_CHECKING:
21+
from fla.ops.cp import FLACPContext
22+
23+
24+
class KDAFusedInferBackend(BaseBackend):
25+
26+
backend_type = "fused_infer"
27+
package_name = "triton"
28+
env_var = "FLA_FUSED_INFER"
29+
default_enable = True
30+
priority = 4
31+
32+
@classmethod
33+
def is_available(cls) -> bool:
34+
try:
35+
import triton # noqa: F401
36+
return True
37+
except ImportError:
38+
return False
39+
40+
def chunk_kda_verifier(
41+
self,
42+
q: torch.Tensor,
43+
k: torch.Tensor,
44+
v: torch.Tensor,
45+
g: torch.Tensor,
46+
beta: torch.Tensor,
47+
scale: float | None = None,
48+
initial_state: torch.Tensor | None = None,
49+
output_final_state: bool = False,
50+
use_qk_l2norm_in_kernel: bool = False,
51+
use_gate_in_kernel: bool = False,
52+
use_beta_sigmoid_in_kernel: bool = False,
53+
state_v_first: bool = False,
54+
cu_seqlens: torch.LongTensor | None = None,
55+
cu_seqlens_cpu: torch.LongTensor | None = None,
56+
safe_gate: bool = False,
57+
lower_bound: float | None = None,
58+
disable_recompute: bool = False,
59+
return_intermediate_states: bool = False,
60+
cp_context: FLACPContext | None = None,
61+
**kwargs,
62+
) -> tuple[bool, str | None]:
63+
if torch.is_grad_enabled():
64+
return False, "Fused infer backend only supports inference mode"
65+
if return_intermediate_states:
66+
return False, "Fused infer backend does not support return_intermediate_states"
67+
chunk_size = kwargs.get("chunk_size", 16)
68+
if chunk_size != 16:
69+
return False, f"Fused infer backend requires chunk_size=16, got {chunk_size}"
70+
return True, None
71+
72+
def chunk_kda(
73+
self,
74+
q: torch.Tensor,
75+
k: torch.Tensor,
76+
v: torch.Tensor,
77+
g: torch.Tensor,
78+
beta: torch.Tensor,
79+
scale: float | None = None,
80+
initial_state: torch.Tensor | None = None,
81+
output_final_state: bool = False,
82+
use_qk_l2norm_in_kernel: bool = False,
83+
use_gate_in_kernel: bool = False,
84+
use_beta_sigmoid_in_kernel: bool = False,
85+
state_v_first: bool = False,
86+
cu_seqlens: torch.LongTensor | None = None,
87+
cu_seqlens_cpu: torch.LongTensor | None = None,
88+
safe_gate: bool = False,
89+
lower_bound: float | None = None,
90+
disable_recompute: bool = False,
91+
cp_context: FLACPContext | None = None,
92+
**kwargs,
93+
) -> tuple[torch.Tensor, torch.Tensor | None]:
94+
from fla.ops.kda.chunk_fwd_infer import chunk_kda_fwd_infer
95+
return chunk_kda_fwd_infer(
96+
q=q,
97+
k=k,
98+
v=v,
99+
g=g,
100+
beta=beta,
101+
scale=scale,
102+
initial_state=initial_state,
103+
output_final_state=output_final_state,
104+
use_qk_l2norm_in_kernel=use_qk_l2norm_in_kernel,
105+
use_gate_in_kernel=use_gate_in_kernel,
106+
use_beta_sigmoid_in_kernel=use_beta_sigmoid_in_kernel,
107+
state_v_first=state_v_first,
108+
cu_seqlens=cu_seqlens,
109+
cu_seqlens_cpu=cu_seqlens_cpu,
110+
safe_gate=safe_gate,
111+
lower_bound=lower_bound,
112+
disable_recompute=disable_recompute,
113+
cp_context=cp_context,
114+
**kwargs,
115+
)

‎fla/ops/kda/backends/tle/__init__.py‎

Lines changed: 169 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,169 @@
1+
# Copyright (c) 2023-2026, Songlin Yang, Yu Zhang, Zhiyuan Li
2+
#
3+
# This source code is licensed under the MIT license found in the
4+
# LICENSE file in the root directory of this source tree.
5+
# For a list of all contributors, visit:
6+
# https://github.com/fla-org/flash-linear-attention/graphs/contributors
7+
#
8+
# Copyright (c) 2026 FlagOS Contributors
9+
10+
"""TLE inference backend for KDA."""
11+
12+
from __future__ import annotations
13+
14+
from typing import TYPE_CHECKING
15+
16+
import torch
17+
18+
from fla.ops.backends import BaseBackend
19+
20+
if TYPE_CHECKING:
21+
from fla.ops.cp import FLACPContext
22+
23+
24+
class KDATLEBackend(BaseBackend):
25+
"""TLE forward backend for inference-only chunk_kda.
26+
27+
The kernel fuses q/k L2 norm, beta sigmoid, and the KDA safe gate. Compared
28+
with the FlashKDA CUDA backend, this path also supports fp16, GVA, flexible
29+
state layout, K in {64, 128, 192, 256}, and arbitrary positive V.
30+
"""
31+
32+
backend_type = "tle"
33+
package_name = "triton"
34+
env_var = "FLA_TLE"
35+
default_enable = True
36+
# Higher priority than FlashKDA (3); the verifier falls back automatically
37+
# when a call does not match TLE's inference-only contract.
38+
priority = 2
39+
40+
@classmethod
41+
def is_available(cls) -> bool:
42+
try:
43+
import triton.experimental.tle.language # noqa: F401
44+
return True
45+
except ImportError:
46+
return False
47+
48+
def chunk_kda_verifier(
49+
self,
50+
q: torch.Tensor,
51+
k: torch.Tensor,
52+
v: torch.Tensor,
53+
g: torch.Tensor,
54+
beta: torch.Tensor,
55+
scale: float | None = None,
56+
initial_state: torch.Tensor | None = None,
57+
output_final_state: bool = False,
58+
use_qk_l2norm_in_kernel: bool = False,
59+
use_gate_in_kernel: bool = False,
60+
use_beta_sigmoid_in_kernel: bool = False,
61+
allow_neg_eigval: bool = False,
62+
state_v_first: bool = False,
63+
cu_seqlens: torch.LongTensor | None = None,
64+
cu_seqlens_cpu: torch.LongTensor | None = None,
65+
safe_gate: bool = False,
66+
lower_bound: float | None = None,
67+
disable_recompute: bool = False,
68+
return_intermediate_states: bool = False,
69+
cp_context: FLACPContext | None = None,
70+
**kwargs,
71+
) -> tuple[bool, str | None]:
72+
if torch.is_grad_enabled():
73+
return False, "TLE backend only supports inference mode"
74+
if return_intermediate_states:
75+
return False, "TLE backend does not support return_intermediate_states"
76+
chunk_size = kwargs.get("chunk_size", 16)
77+
if chunk_size != 16:
78+
return False, f"TLE backend requires chunk_size=16, got {chunk_size}"
79+
supported_dtypes = (torch.bfloat16, torch.float16)
80+
for name, tensor in (("q", q), ("k", k), ("v", v), ("g", g), ("beta", beta)):
81+
if tensor.dtype not in supported_dtypes:
82+
return False, f"TLE backend requires {name} dtype to be bf16 or fp16, got {tensor.dtype}"
83+
K = q.shape[-1]
84+
H = q.shape[2]
85+
HV = v.shape[2]
86+
V = v.shape[-1]
87+
if K not in (64, 128, 192, 256):
88+
return False, f"TLE backend requires K in {{64, 128, 192, 256}}, got {K}"
89+
if H <= 0:
90+
return False, f"TLE backend requires H > 0, got {H}"
91+
if V <= 0:
92+
return False, f"TLE backend requires V > 0, got {V}"
93+
if k.shape[2] != H:
94+
return False, f"TLE backend requires k heads to match q heads, got H={H}, HK={k.shape[2]}"
95+
if k.shape[-1] != K or g.shape[-1] != K:
96+
return False, f"TLE backend requires q/k/g to share K={K}, got HK={k.shape[-1]}, KG={g.shape[-1]}"
97+
if g.shape[2] != HV or beta.shape[2] != HV:
98+
return False, f"TLE backend requires v/g/beta to share HV={HV}, got HG={g.shape[2]}, HB={beta.shape[2]}"
99+
if HV < H or HV % H != 0:
100+
return False, f"TLE backend requires HV % H == 0 for GVA, got H={H}, HV={HV}"
101+
if initial_state is not None:
102+
N = len(cu_seqlens) - 1 if cu_seqlens is not None else q.shape[0]
103+
expected_shape = (N, HV, V, K) if state_v_first else (N, HV, K, V)
104+
if tuple(initial_state.shape) != expected_shape:
105+
return False, f"TLE backend requires initial_state shape {expected_shape}, got {tuple(initial_state.shape)}"
106+
if cp_context is not None:
107+
return False, "TLE backend does not support context parallel"
108+
if not use_qk_l2norm_in_kernel:
109+
return False, "TLE backend requires use_qk_l2norm_in_kernel=True"
110+
if not use_gate_in_kernel:
111+
return False, "TLE backend requires use_gate_in_kernel=True"
112+
A_log = kwargs.get("A_log")
113+
dt_bias = kwargs.get("dt_bias")
114+
if A_log is None:
115+
return False, "TLE backend requires A_log"
116+
if A_log.numel() != HV:
117+
return False, f"TLE backend requires A_log.numel() == HV={HV}, got {A_log.numel()}"
118+
if dt_bias is None:
119+
return False, "TLE backend requires dt_bias"
120+
if dt_bias.numel() != HV * K:
121+
return False, f"TLE backend requires dt_bias.numel() == HV*K={HV * K}, got {dt_bias.numel()}"
122+
if not use_beta_sigmoid_in_kernel:
123+
return False, "TLE backend requires use_beta_sigmoid_in_kernel=True"
124+
if allow_neg_eigval:
125+
return False, "TLE backend does not support allow_neg_eigval=True"
126+
if lower_bound is None:
127+
return False, "TLE backend requires lower_bound"
128+
return True, None
129+
130+
def chunk_kda(
131+
self,
132+
q: torch.Tensor,
133+
k: torch.Tensor,
134+
v: torch.Tensor,
135+
g: torch.Tensor,
136+
beta: torch.Tensor,
137+
scale: float | None = None,
138+
initial_state: torch.Tensor | None = None,
139+
output_final_state: bool = False,
140+
use_qk_l2norm_in_kernel: bool = False,
141+
use_gate_in_kernel: bool = False,
142+
use_beta_sigmoid_in_kernel: bool = False,
143+
allow_neg_eigval: bool = False,
144+
state_v_first: bool = False,
145+
cu_seqlens: torch.LongTensor | None = None,
146+
cu_seqlens_cpu: torch.LongTensor | None = None,
147+
safe_gate: bool = False,
148+
lower_bound: float | None = None,
149+
disable_recompute: bool = False,
150+
cp_context: FLACPContext | None = None,
151+
**kwargs,
152+
) -> tuple[torch.Tensor, torch.Tensor | None]:
153+
from fla.ops.kda.backends.tle.chunk_fwd import chunk_kda_fwd_infer
154+
return chunk_kda_fwd_infer(
155+
q=q,
156+
k=k,
157+
v=v,
158+
g=g,
159+
beta=beta,
160+
scale=scale,
161+
initial_state=initial_state,
162+
output_final_state=output_final_state,
163+
state_v_first=state_v_first,
164+
cu_seqlens=cu_seqlens,
165+
safe_gate=safe_gate,
166+
lower_bound=lower_bound,
167+
A_log=kwargs.get("A_log"),
168+
dt_bias=kwargs.get("dt_bias"),
169+
)

0 commit comments

Comments
 (0)