74 lines
1.9 KiB
Python
74 lines
1.9 KiB
Python
import torch
|
|
|
|
from sglang.srt.layers.attention.linear.kernels.kernel_backend import (
|
|
LinearAttnKernelBase,
|
|
)
|
|
from sglang.srt.utils import is_cpu
|
|
|
|
if not is_cpu():
|
|
from sglang.srt.layers.attention.fla.fused_sigmoid_gating_recurrent import (
|
|
fused_sigmoid_gating_delta_rule_update,
|
|
)
|
|
from sglang.srt.layers.attention.fla.kda import chunk_kda
|
|
|
|
|
|
class TritonKDAKernel(LinearAttnKernelBase):
|
|
"""Triton-based kernel for KDA (Kimi Delta Attention) linear attention."""
|
|
|
|
def decode(
|
|
self,
|
|
q: torch.Tensor,
|
|
k: torch.Tensor,
|
|
v: torch.Tensor,
|
|
a: torch.Tensor,
|
|
b: torch.Tensor,
|
|
*,
|
|
A_log: torch.Tensor,
|
|
dt_bias: torch.Tensor,
|
|
ssm_states: torch.Tensor,
|
|
cache_indices: torch.Tensor,
|
|
query_start_loc: torch.Tensor,
|
|
**kwargs,
|
|
) -> torch.Tensor:
|
|
return fused_sigmoid_gating_delta_rule_update(
|
|
A_log=A_log,
|
|
dt_bias=dt_bias,
|
|
q=q,
|
|
k=k,
|
|
v=v,
|
|
a=a,
|
|
b=b,
|
|
initial_state_source=ssm_states,
|
|
initial_state_indices=cache_indices,
|
|
cu_seqlens=query_start_loc,
|
|
use_qk_l2norm_in_kernel=True,
|
|
softplus_beta=1.0,
|
|
softplus_threshold=20.0,
|
|
is_kda=True,
|
|
)
|
|
|
|
def extend(
|
|
self,
|
|
q: torch.Tensor,
|
|
k: torch.Tensor,
|
|
v: torch.Tensor,
|
|
g: torch.Tensor,
|
|
beta: torch.Tensor,
|
|
*,
|
|
ssm_states: torch.Tensor,
|
|
cache_indices: torch.Tensor,
|
|
query_start_loc: torch.Tensor,
|
|
**kwargs,
|
|
) -> torch.Tensor:
|
|
return chunk_kda(
|
|
q=q,
|
|
k=k,
|
|
v=v,
|
|
g=g,
|
|
beta=beta,
|
|
initial_state=ssm_states,
|
|
initial_state_indices=cache_indices,
|
|
use_qk_l2norm_in_kernel=True,
|
|
cu_seqlens=query_start_loc,
|
|
)
|