Skip to content

Commit b5a810f

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 b5a810f

6 files changed

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

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

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

0 commit comments

Comments
 (0)