[sgl-kernel Code Clean] Remove useless lightning_attention kernel (#13819)
This commit is contained in:
Xiaoyu Zhang
@@ -1,577 +0,0 @@
|
||||
import itertools
|
||||
import math
|
||||
from typing import Optional, Tuple
|
||||
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
import torch.nn.functional as F
|
||||
import triton
|
||||
import triton.language as tl
|
||||
from einops import rearrange
|
||||
from sgl_kernel import lightning_attention_decode as sgl_lightning_attention_decode
|
||||
|
||||
|
||||
@triton.jit
|
||||
def _decode_kernel(
|
||||
Q,
|
||||
K,
|
||||
V,
|
||||
KV,
|
||||
Out,
|
||||
S,
|
||||
b: tl.constexpr,
|
||||
h: tl.constexpr,
|
||||
n: tl.constexpr,
|
||||
d: tl.constexpr,
|
||||
d_original: tl.constexpr,
|
||||
e: tl.constexpr,
|
||||
e_original: tl.constexpr,
|
||||
):
|
||||
off_bh = tl.program_id(0)
|
||||
off_h = off_bh % h
|
||||
|
||||
qk_offset = off_bh * n * d
|
||||
v_offset = off_bh * n * e
|
||||
o_offset = off_bh * n * e
|
||||
kv_offset = off_bh * d * e
|
||||
|
||||
s = tl.load(S + off_h)
|
||||
ratio = tl.exp(-s)
|
||||
|
||||
d_idx = tl.arange(0, d)
|
||||
e_idx = tl.arange(0, e)
|
||||
|
||||
# Create masks for original dimensions
|
||||
d_mask = d_idx < d_original
|
||||
e_mask = e_idx < e_original
|
||||
|
||||
# Load with masking
|
||||
q = tl.load(Q + qk_offset + d_idx, mask=d_mask, other=0.0)
|
||||
k = tl.load(K + qk_offset + d_idx, mask=d_mask, other=0.0)
|
||||
v = tl.load(V + v_offset + e_idx, mask=e_mask, other=0.0)
|
||||
|
||||
# Load KV with 2D masking
|
||||
kv = tl.load(
|
||||
KV + kv_offset + d_idx[:, None] * e + e_idx[None, :],
|
||||
mask=(d_mask[:, None] & e_mask[None, :]),
|
||||
other=0.0,
|
||||
)
|
||||
|
||||
# Compute outer product using element-wise operations
|
||||
k_v_prod = k[:, None] * v[None, :]
|
||||
kv = ratio * kv + k_v_prod
|
||||
|
||||
# Store KV with 2D masking
|
||||
tl.store(
|
||||
KV + kv_offset + d_idx[:, None] * e + e_idx[None, :],
|
||||
kv.to(KV.dtype.element_ty),
|
||||
mask=(d_mask[:, None] & e_mask[None, :]),
|
||||
)
|
||||
|
||||
# Compute matrix-vector multiplication using element-wise operations and reduction
|
||||
o = tl.sum(q[:, None] * kv, axis=0)
|
||||
|
||||
# Store output with masking
|
||||
tl.store(Out + o_offset + e_idx, o.to(Out.dtype.element_ty), mask=e_mask)
|
||||
|
||||
|
||||
def lightning_attn_decode(q, k, v, kv, s):
|
||||
"""Triton implementation of Lightning Attention decode operation"""
|
||||
b, h, n, d = q.shape
|
||||
e = v.shape[-1]
|
||||
assert n == 1, "Sequence length must be 1 in decode mode"
|
||||
|
||||
# Get padded dimensions (power of 2)
|
||||
d_padded = next_power_of_2(d)
|
||||
e_padded = next_power_of_2(e)
|
||||
|
||||
# Create output tensor (padded)
|
||||
o_padded = torch.empty(b, h, n, e_padded, dtype=v.dtype, device=v.device)
|
||||
|
||||
# Create padded tensors without actually padding the data
|
||||
q_padded = torch.empty(b, h, n, d_padded, dtype=q.dtype, device=q.device)
|
||||
k_padded = torch.empty(b, h, n, d_padded, dtype=k.dtype, device=k.device)
|
||||
v_padded = torch.empty(b, h, n, e_padded, dtype=v.dtype, device=v.device)
|
||||
kv_padded = torch.empty(
|
||||
b, h, d_padded, e_padded, dtype=torch.float32, device=kv.device
|
||||
)
|
||||
|
||||
# Copy data to padded tensors
|
||||
q_padded[..., :d] = q
|
||||
k_padded[..., :d] = k
|
||||
v_padded[..., :e] = v
|
||||
kv_padded[..., :d, :e] = kv
|
||||
|
||||
# Launch kernel
|
||||
grid = (b * h, 1)
|
||||
_decode_kernel[grid](
|
||||
q_padded,
|
||||
k_padded,
|
||||
v_padded,
|
||||
kv_padded,
|
||||
o_padded,
|
||||
s,
|
||||
b=b,
|
||||
h=h,
|
||||
n=n,
|
||||
d=d_padded,
|
||||
d_original=d,
|
||||
e=e_padded,
|
||||
e_original=e,
|
||||
)
|
||||
|
||||
# Get unpadded outputs
|
||||
o = o_padded[..., :e]
|
||||
kv_out = kv_padded[..., :d, :e]
|
||||
|
||||
return o, kv_out
|
||||
|
||||
|
||||
def next_power_of_2(n):
|
||||
return 2 ** (int(math.ceil(math.log(n, 2))))
|
||||
|
||||
|
||||
class MiniMaxText01LightningAttention(nn.Module):
|
||||
def __init__(self, config=None, layer_idx: Optional[int] = None, **kwargs):
|
||||
super().__init__()
|
||||
if config is None:
|
||||
config = type("Config", (), kwargs)
|
||||
|
||||
bias = False
|
||||
self.hidden_size = config.hidden_size
|
||||
self.num_heads = config.num_attention_heads
|
||||
self.head_dim = getattr(config, "head_dim", self.hidden_size // self.num_heads)
|
||||
|
||||
self.out_proj = nn.Linear(
|
||||
self.head_dim * self.num_heads, self.hidden_size, bias=bias
|
||||
)
|
||||
self.act = get_activation_fn(config.hidden_act)
|
||||
self.norm = MiniMaxText01RMSNorm(self.head_dim * self.num_heads)
|
||||
|
||||
self.qkv_proj = nn.Linear(
|
||||
self.hidden_size, 3 * self.head_dim * self.num_heads, bias=bias
|
||||
)
|
||||
self.output_gate = nn.Linear(
|
||||
self.hidden_size, self.head_dim * self.num_heads, bias=bias
|
||||
)
|
||||
|
||||
# for inference only
|
||||
self.offset = 0
|
||||
self.layer_idx = layer_idx
|
||||
|
||||
def forward(
|
||||
self,
|
||||
hidden_states,
|
||||
attn_mask: Optional[torch.Tensor] = None, # (b, h, n, m)
|
||||
output_attentions: bool = False,
|
||||
past_key_value: Optional[Tuple[torch.Tensor]] = None,
|
||||
use_cache: bool = False,
|
||||
slope_rate: Optional[torch.Tensor] = None,
|
||||
do_eval: bool = False,
|
||||
**kwargs,
|
||||
):
|
||||
if (not self.training) and (not do_eval):
|
||||
return self.inference(
|
||||
hidden_states,
|
||||
attn_mask,
|
||||
output_attentions,
|
||||
past_key_value,
|
||||
use_cache,
|
||||
slope_rate,
|
||||
)
|
||||
|
||||
def inference(
|
||||
self,
|
||||
x,
|
||||
attn_mask: Optional[torch.Tensor] = None, # (b, n)
|
||||
output_attentions: bool = False,
|
||||
past_key_value: Optional[Tuple[torch.Tensor]] = None,
|
||||
use_cache: bool = False,
|
||||
slope_rate: Optional[torch.Tensor] = None, # (h, 1, 1)
|
||||
):
|
||||
# x: b n d
|
||||
b, n, d = x.shape
|
||||
# linear map
|
||||
qkv = self.act(self.qkv_proj(x))
|
||||
new_shape = qkv.size()[:-1] + (self.num_heads, -1)
|
||||
qkv = qkv.view(*new_shape)
|
||||
q, k, v = torch.split(qkv, [self.head_dim] * 3, dim=3)
|
||||
q = q.transpose(1, 2) # [b, n, h, d] -> [b, h, n, d]
|
||||
k = k.transpose(1, 2) # [b, n, h, d] -> [b, h, n, d]
|
||||
v = v.transpose(1, 2) # [b, n, h, d] -> [b, h, n, e]
|
||||
|
||||
self.offset += 1
|
||||
ratio = torch.exp(-slope_rate) # [h, 1, 1]
|
||||
|
||||
# decode mode
|
||||
kv = past_key_value # [b, h, d, e]
|
||||
output = []
|
||||
for i in range(n):
|
||||
# kv: [b, h, d, e]
|
||||
# ratio: [h, 1, 1]
|
||||
# k: [b, h, n, d]
|
||||
# v: [b, h, n, e]
|
||||
# k[:, :, i : i + 1]: [b, h, 1, d]
|
||||
# v[:, :, i : i + 1]: [b, h, 1, e]
|
||||
# ratio * kv: [b, h, d, e]
|
||||
# torch.einsum(
|
||||
# "... n d, ... n e -> ... d e",
|
||||
# k[:, :, i : i + 1],
|
||||
# v[:, :, i : i + 1],
|
||||
# )
|
||||
# [b, h, d, e] + [b, h, d, e] -> [b, h, d, e]
|
||||
kv = ratio * kv + torch.einsum(
|
||||
"... n d, ... n e -> ... d e",
|
||||
k[:, :, i : i + 1],
|
||||
v[:, :, i : i + 1],
|
||||
)
|
||||
# q[:, :, i : i + 1]: [b, h, 1, d]
|
||||
# kv.to(q.dtype): [b, h, d, e]
|
||||
# torch.einsum(
|
||||
# "... n e, ... e d -> ... n d", q[:, :, i : i + 1], kv.to(q.dtype)
|
||||
# )
|
||||
# [b, h, 1, d] * [b, h, d, e] -> [b, h, 1, e]
|
||||
qkv = torch.einsum(
|
||||
"... n e, ... e d -> ... n d", q[:, :, i : i + 1], kv.to(q.dtype)
|
||||
)
|
||||
output.append(qkv)
|
||||
output = torch.cat(output, dim=-2)
|
||||
|
||||
# reshape
|
||||
output = rearrange(output, "b h n d -> b n (h d)")
|
||||
# normalize
|
||||
output = self.norm(output)
|
||||
# gate
|
||||
output = F.sigmoid(self.output_gate(x)) * output
|
||||
# outproj
|
||||
output = self.out_proj(output)
|
||||
|
||||
attn_weights = None
|
||||
|
||||
return output, attn_weights, kv
|
||||
|
||||
|
||||
def get_activation_fn(activation):
|
||||
if activation == "gelu":
|
||||
return F.gelu
|
||||
elif activation == "relu":
|
||||
return F.relu
|
||||
elif activation == "elu":
|
||||
return F.elu
|
||||
elif activation == "sigmoid":
|
||||
return F.sigmoid
|
||||
elif activation == "exp":
|
||||
|
||||
def f(x):
|
||||
with torch.no_grad():
|
||||
x_max = torch.max(x, dim=-1, keepdims=True).values
|
||||
y = torch.exp(x - x_max)
|
||||
return y
|
||||
|
||||
return f
|
||||
elif activation == "leak":
|
||||
return F.leaky_relu
|
||||
elif activation == "1+elu":
|
||||
|
||||
def f(x):
|
||||
return 1 + F.elu(x)
|
||||
|
||||
return f
|
||||
elif activation == "2+elu":
|
||||
|
||||
def f(x):
|
||||
return 2 + F.elu(x)
|
||||
|
||||
return f
|
||||
elif activation == "silu" or activation == "swish":
|
||||
return F.silu
|
||||
elif activation == "sine":
|
||||
return torch.sin
|
||||
else:
|
||||
return lambda x: x
|
||||
|
||||
|
||||
class MiniMaxText01RMSNorm(nn.Module):
|
||||
def __init__(self, hidden_size, eps=1e-6):
|
||||
"""
|
||||
MiniMaxText01RMSNorm is equivalent to T5LayerNorm
|
||||
"""
|
||||
super().__init__()
|
||||
self.weight = nn.Parameter(torch.ones(hidden_size))
|
||||
self.variance_epsilon = eps
|
||||
|
||||
def forward(self, hidden_states):
|
||||
input_dtype = hidden_states.dtype
|
||||
hidden_states = hidden_states.to(torch.float32)
|
||||
variance = hidden_states.pow(2).mean(-1, keepdim=True)
|
||||
hidden_states = hidden_states * torch.rsqrt(variance + self.variance_epsilon)
|
||||
return self.weight * hidden_states.to(input_dtype)
|
||||
|
||||
|
||||
def test_lightning_attention_implementations(model_params):
|
||||
torch.manual_seed(42)
|
||||
|
||||
batch_size = 64
|
||||
seq_len = 1
|
||||
dtype = torch.bfloat16
|
||||
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
|
||||
|
||||
hidden_states = torch.randn(
|
||||
batch_size, seq_len, model_params["hidden_size"], dtype=dtype, device=device
|
||||
)
|
||||
|
||||
attention_mask = torch.ones(batch_size, seq_len, dtype=dtype, device=device)
|
||||
|
||||
slope_rate = _build_slope_tensor(model_params["num_attention_heads"]).to(device)
|
||||
|
||||
model_attn = MiniMaxText01LightningAttention(**model_params).to(dtype).to(device)
|
||||
model_attn.eval()
|
||||
|
||||
d = model_params["head_dim"]
|
||||
past_kv = torch.randn(
|
||||
batch_size,
|
||||
model_params["num_attention_heads"],
|
||||
d,
|
||||
d,
|
||||
device=device,
|
||||
)
|
||||
with torch.no_grad():
|
||||
model_output, _, new_kv = model_attn.inference(
|
||||
hidden_states,
|
||||
attn_mask=attention_mask,
|
||||
slope_rate=slope_rate,
|
||||
past_key_value=past_kv,
|
||||
)
|
||||
|
||||
qkv = model_attn.act(model_attn.qkv_proj(hidden_states))
|
||||
new_shape = qkv.size()[:-1] + (model_attn.num_heads, -1)
|
||||
qkv = qkv.view(*new_shape)
|
||||
q, k, v = torch.split(qkv, [model_attn.head_dim] * 3, dim=-1)
|
||||
q = q.transpose(1, 2)
|
||||
k = k.transpose(1, 2)
|
||||
v = v.transpose(1, 2)
|
||||
q = q.contiguous()
|
||||
k = k.contiguous()
|
||||
v = v.contiguous()
|
||||
past_kv = past_kv.contiguous()
|
||||
slope_rate = slope_rate.contiguous()
|
||||
|
||||
# Test Triton implementation
|
||||
triton_output, triton_new_kv = lightning_attn_decode(q, k, v, past_kv, slope_rate)
|
||||
triton_output = triton_output.transpose(1, 2).contiguous()
|
||||
triton_output = triton_output.view(batch_size, seq_len, -1)
|
||||
triton_output = model_attn.norm(triton_output)
|
||||
triton_output = torch.sigmoid(model_attn.output_gate(hidden_states)) * triton_output
|
||||
triton_output = model_attn.out_proj(triton_output)
|
||||
|
||||
# Test SGL implementation
|
||||
sgl_output = torch.empty_like(v)
|
||||
sgl_new_kv = torch.empty_like(past_kv)
|
||||
sgl_lightning_attention_decode(q, k, v, past_kv, slope_rate, sgl_output, sgl_new_kv)
|
||||
|
||||
sgl_output = sgl_output.transpose(1, 2).contiguous()
|
||||
sgl_output = sgl_output.view(batch_size, seq_len, -1)
|
||||
sgl_output = model_attn.norm(sgl_output)
|
||||
sgl_output = torch.sigmoid(model_attn.output_gate(hidden_states)) * sgl_output
|
||||
sgl_output = model_attn.out_proj(sgl_output)
|
||||
|
||||
# Verify Triton implementation results
|
||||
torch.testing.assert_close(
|
||||
model_output,
|
||||
triton_output,
|
||||
rtol=1e-3,
|
||||
atol=1e-2,
|
||||
msg="Triton lightning attention implementation produces different output results",
|
||||
)
|
||||
torch.testing.assert_close(
|
||||
new_kv,
|
||||
triton_new_kv,
|
||||
rtol=1e-3,
|
||||
atol=1e-2,
|
||||
msg="Triton lightning attention implementation produces different kv results",
|
||||
)
|
||||
|
||||
# Verify SGL implementation results
|
||||
torch.testing.assert_close(
|
||||
model_output,
|
||||
sgl_output,
|
||||
rtol=1e-3,
|
||||
atol=1e-2,
|
||||
msg="SGL lightning attention implementation produces different output results",
|
||||
)
|
||||
torch.testing.assert_close(
|
||||
new_kv,
|
||||
sgl_new_kv,
|
||||
rtol=1e-3,
|
||||
atol=1e-2,
|
||||
msg="SGL lightning attention implementation produces different kv results",
|
||||
)
|
||||
|
||||
print("✅ All implementations match")
|
||||
|
||||
|
||||
def _build_slope_tensor(n_attention_heads: int):
|
||||
def get_slopes(n):
|
||||
def get_slopes_power_of_2(n):
|
||||
start = 2 ** (-(2 ** -(math.log2(n) - 3)))
|
||||
ratio = start
|
||||
return [start * ratio**i for i in range(n)]
|
||||
|
||||
if math.log2(n).is_integer():
|
||||
return get_slopes_power_of_2(n)
|
||||
else:
|
||||
closest_power_of_2 = 2 ** math.floor(math.log2(n))
|
||||
return (
|
||||
get_slopes_power_of_2(closest_power_of_2)
|
||||
+ get_slopes(2 * closest_power_of_2)[0::2][: n - closest_power_of_2]
|
||||
)
|
||||
|
||||
slopes = torch.tensor(get_slopes(n_attention_heads)).reshape(
|
||||
n_attention_heads, 1, 1
|
||||
)
|
||||
return slopes
|
||||
|
||||
|
||||
def get_benchmark():
|
||||
batch_size_range = [i for i in range(1, 33)] # max 32
|
||||
seq_length_range = [1] # decode mode sequence length is fixed to 1
|
||||
configs = list(itertools.product(batch_size_range, seq_length_range))
|
||||
|
||||
@triton.testing.perf_report(
|
||||
triton.testing.Benchmark(
|
||||
x_names=["batch_size", "seq_len"],
|
||||
x_vals=[list(_) for _ in configs],
|
||||
line_arg="provider",
|
||||
line_vals=["Original", "Triton", "SGL"],
|
||||
line_names=[
|
||||
"Original PyTorch Implementation",
|
||||
"Triton Implementation",
|
||||
"SGL Implementation",
|
||||
],
|
||||
styles=[("blue", "-"), ("green", "-"), ("red", "-")],
|
||||
ylabel="us",
|
||||
plot_name="lightning-attention-decode-performance",
|
||||
args={},
|
||||
)
|
||||
)
|
||||
def benchmark(batch_size, seq_len, provider):
|
||||
dtype = torch.bfloat16
|
||||
device = torch.device("cuda")
|
||||
|
||||
params = {
|
||||
"hidden_size": 6144,
|
||||
"num_attention_heads": 64,
|
||||
"head_dim": 96,
|
||||
"hidden_act": "gelu",
|
||||
}
|
||||
|
||||
hidden_states = torch.randn(
|
||||
batch_size, seq_len, params["hidden_size"], dtype=dtype, device=device
|
||||
)
|
||||
|
||||
attention_mask = torch.ones(batch_size, seq_len, dtype=dtype, device=device)
|
||||
|
||||
slope_rate = _build_slope_tensor(params["num_attention_heads"]).to(device)
|
||||
model_attn = MiniMaxText01LightningAttention(**params).to(dtype).to(device)
|
||||
model_attn.eval()
|
||||
|
||||
d = params["head_dim"]
|
||||
past_kv = torch.randn(
|
||||
batch_size,
|
||||
params["num_attention_heads"],
|
||||
d,
|
||||
d,
|
||||
device=device,
|
||||
)
|
||||
|
||||
quantiles = [0.5, 0.2, 0.8]
|
||||
if provider == "Original":
|
||||
ms, min_ms, max_ms = triton.testing.do_bench(
|
||||
lambda: model_attn.inference(
|
||||
hidden_states,
|
||||
attn_mask=attention_mask,
|
||||
slope_rate=slope_rate,
|
||||
past_key_value=past_kv,
|
||||
),
|
||||
quantiles=quantiles,
|
||||
)
|
||||
elif provider == "Triton":
|
||||
|
||||
def run_triton():
|
||||
qkv = model_attn.act(model_attn.qkv_proj(hidden_states))
|
||||
new_shape = qkv.size()[:-1] + (model_attn.num_heads, -1)
|
||||
qkv = qkv.view(*new_shape)
|
||||
q, k, v = torch.split(qkv, [model_attn.head_dim] * 3, dim=-1)
|
||||
q = q.transpose(1, 2)
|
||||
k = k.transpose(1, 2)
|
||||
v = v.transpose(1, 2)
|
||||
|
||||
output, new_kv = lightning_attn_decode(q, k, v, past_kv, slope_rate)
|
||||
output = output.transpose(1, 2).contiguous()
|
||||
output = output.view(batch_size, seq_len, -1)
|
||||
output = model_attn.norm(output)
|
||||
output = torch.sigmoid(model_attn.output_gate(hidden_states)) * output
|
||||
return model_attn.out_proj(output)
|
||||
|
||||
ms, min_ms, max_ms = triton.testing.do_bench(
|
||||
run_triton,
|
||||
quantiles=quantiles,
|
||||
)
|
||||
else: # SGL
|
||||
|
||||
def run_sgl():
|
||||
qkv = model_attn.act(model_attn.qkv_proj(hidden_states))
|
||||
new_shape = qkv.size()[:-1] + (model_attn.num_heads, -1)
|
||||
qkv = qkv.view(*new_shape)
|
||||
q, k, v = torch.split(qkv, [model_attn.head_dim] * 3, dim=-1)
|
||||
q = q.transpose(1, 2).contiguous()
|
||||
k = k.transpose(1, 2).contiguous()
|
||||
v = v.transpose(1, 2).contiguous()
|
||||
|
||||
output = torch.empty_like(v)
|
||||
new_kv = torch.empty_like(past_kv)
|
||||
sgl_lightning_attention_decode(
|
||||
q, k, v, past_kv, slope_rate, output, new_kv
|
||||
)
|
||||
|
||||
output = output.transpose(1, 2).contiguous()
|
||||
output = output.view(batch_size, seq_len, -1)
|
||||
output = model_attn.norm(output)
|
||||
output = torch.sigmoid(model_attn.output_gate(hidden_states)) * output
|
||||
return model_attn.out_proj(output)
|
||||
|
||||
ms, min_ms, max_ms = triton.testing.do_bench(
|
||||
run_sgl,
|
||||
quantiles=quantiles,
|
||||
)
|
||||
|
||||
return 1000 * ms, 1000 * max_ms, 1000 * min_ms
|
||||
|
||||
return benchmark
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
import argparse
|
||||
|
||||
parser = argparse.ArgumentParser()
|
||||
parser.add_argument(
|
||||
"--save_path",
|
||||
type=str,
|
||||
default="./configs/benchmark_ops/lightning_attention_decode/",
|
||||
help="Path to save lightning attention decode benchmark results",
|
||||
)
|
||||
args = parser.parse_args()
|
||||
|
||||
params = {
|
||||
"hidden_size": 6144,
|
||||
"num_attention_heads": 64,
|
||||
"head_dim": 96,
|
||||
"hidden_act": "silu",
|
||||
}
|
||||
# Run correctness test first
|
||||
# Adapted from https://huggingface.co/MiniMaxAI/MiniMax-Text-01/blob/main/config.json
|
||||
test_lightning_attention_implementations(params)
|
||||
|
||||
# Run performance benchmark
|
||||
benchmark = get_benchmark()
|
||||
benchmark.run(print_data=True, save_path=args.save_path)
|
||||
@@ -1,607 +0,0 @@
|
||||
import itertools
|
||||
import logging
|
||||
import math
|
||||
import os
|
||||
from typing import Optional, Tuple
|
||||
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
import torch.nn.functional as F
|
||||
import triton
|
||||
import triton.language as tl
|
||||
from einops import rearrange
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
# Adapted from https://github.com/OpenNLPLab/lightning-attention/blob/main/lightning_attn/ops/triton/lightning_attn2.py
|
||||
@triton.jit
|
||||
def _fwd_kernel(
|
||||
Q,
|
||||
K,
|
||||
V,
|
||||
Out,
|
||||
S, # log lambda
|
||||
b: tl.constexpr,
|
||||
h: tl.constexpr,
|
||||
n: tl.constexpr,
|
||||
d: tl.constexpr,
|
||||
e: tl.constexpr,
|
||||
BLOCK: tl.constexpr,
|
||||
NUM_BLOCK: tl.constexpr,
|
||||
BLOCK_MODEL: tl.constexpr,
|
||||
):
|
||||
##### get offset
|
||||
off_bh = tl.program_id(0)
|
||||
off_h = off_bh % h
|
||||
off_e = tl.program_id(1)
|
||||
qk_offset = off_bh * n * d
|
||||
v_offset = off_bh * n * e
|
||||
o_offset = off_bh * n * e
|
||||
# channel offset
|
||||
e_offset = off_e * BLOCK_MODEL
|
||||
|
||||
##### get block ptr
|
||||
Q_block_ptr = Q + qk_offset + tl.arange(0, d)[None, :]
|
||||
K_trans_block_ptr = K + qk_offset + tl.arange(0, d)[:, None]
|
||||
V_block_ptr = V + v_offset + e_offset + tl.arange(0, BLOCK_MODEL)[None, :]
|
||||
O_block_ptr = Out + o_offset + e_offset + tl.arange(0, BLOCK_MODEL)[None, :]
|
||||
S_block_ptr = S + off_h
|
||||
|
||||
##### init diag decay(Lambda); q, k decay; kv
|
||||
s = tl.load(S_block_ptr)
|
||||
# q, k decay
|
||||
off_block = tl.arange(
|
||||
0, BLOCK
|
||||
) # Not bug, this is a bit different from algorithm 1, but is mathematically equivalent
|
||||
q_decay = tl.exp(-s.to(tl.float32) * off_block[:, None])
|
||||
k_trans_decay = tl.exp(-s.to(tl.float32) * (BLOCK - off_block[None, :]))
|
||||
block_decay = tl.exp(-s.to(tl.float32) * BLOCK)
|
||||
# diag decay
|
||||
index = off_block[:, None] - off_block[None, :]
|
||||
s_index = s * index
|
||||
s_index = tl.where(index >= 0, -s_index, float("-inf"))
|
||||
diag_decay = tl.exp(s_index)
|
||||
kv = tl.zeros([d, BLOCK_MODEL], dtype=tl.float32)
|
||||
|
||||
##### compute
|
||||
for i in range(NUM_BLOCK):
|
||||
# load
|
||||
q = tl.load(
|
||||
Q_block_ptr + off_block[:, None] * d, mask=off_block[:, None] < n, other=0.0
|
||||
).to(tl.float32)
|
||||
k_trans = tl.load(
|
||||
K_trans_block_ptr + off_block[None, :] * d,
|
||||
mask=off_block[None, :] < n,
|
||||
other=0.0,
|
||||
).to(tl.float32)
|
||||
v = tl.load(
|
||||
V_block_ptr + off_block[:, None] * e, mask=off_block[:, None] < n, other=0.0
|
||||
).to(tl.float32)
|
||||
|
||||
# compute
|
||||
qk = tl.dot(q, k_trans) * diag_decay
|
||||
o_intra = tl.dot(qk, v)
|
||||
o_inter = tl.dot(q, kv) * q_decay
|
||||
o = o_intra + o_inter
|
||||
|
||||
# save and update
|
||||
tl.store(
|
||||
O_block_ptr + off_block[:, None] * e,
|
||||
o.to(O_block_ptr.dtype.element_ty),
|
||||
mask=off_block[:, None] < n,
|
||||
)
|
||||
kv = block_decay * kv + tl.dot(k_trans * k_trans_decay, v)
|
||||
off_block += BLOCK
|
||||
|
||||
|
||||
def lightning_attn2(q, k, v, s):
|
||||
q = q.contiguous()
|
||||
k = k.contiguous()
|
||||
v = v.contiguous()
|
||||
s = s.contiguous()
|
||||
|
||||
b, h, n, d = q.shape
|
||||
e = v.shape[-1]
|
||||
|
||||
# Pad d to next power of 2
|
||||
d_padded = next_power_of_2(d)
|
||||
if d_padded != d:
|
||||
q_padded = F.pad(q, (0, d_padded - d))
|
||||
k_padded = F.pad(k, (0, d_padded - d))
|
||||
else:
|
||||
q_padded = q
|
||||
k_padded = k
|
||||
|
||||
# Pad e to next power of 2
|
||||
e_padded = next_power_of_2(e)
|
||||
if e_padded != e:
|
||||
v_padded = F.pad(v, (0, e_padded - e))
|
||||
else:
|
||||
v_padded = v
|
||||
|
||||
o_padded = torch.empty((b, h, n, e_padded), dtype=q.dtype, device=q.device)
|
||||
|
||||
BLOCK = 64
|
||||
NUM_BLOCK = triton.cdiv(q.shape[2], BLOCK)
|
||||
# parallel over channel
|
||||
BLOCK_MODEL = min(triton.next_power_of_2(e_padded), 32)
|
||||
grid = (b * h, triton.cdiv(e_padded, BLOCK_MODEL))
|
||||
|
||||
_fwd_kernel[grid](
|
||||
q_padded,
|
||||
k_padded,
|
||||
v_padded,
|
||||
o_padded,
|
||||
s,
|
||||
b,
|
||||
h,
|
||||
n,
|
||||
d_padded,
|
||||
e_padded,
|
||||
BLOCK=BLOCK,
|
||||
NUM_BLOCK=NUM_BLOCK,
|
||||
BLOCK_MODEL=BLOCK_MODEL,
|
||||
)
|
||||
|
||||
# Remove padding from output
|
||||
if e_padded != e:
|
||||
o = o_padded[..., :e]
|
||||
else:
|
||||
o = o_padded
|
||||
|
||||
return o
|
||||
|
||||
|
||||
def is_support(dim):
|
||||
return 16 % dim
|
||||
|
||||
|
||||
def next_power_of_2(n):
|
||||
return 2 ** (int(math.ceil(math.log(n, 2))))
|
||||
|
||||
|
||||
def lightning_attn_func(q, k, v, s):
|
||||
b, h, n, d = q.shape
|
||||
e = v.shape[-1]
|
||||
assert is_support(d) and is_support(e)
|
||||
|
||||
# pad v's feature dim to power of 2
|
||||
e_pad = next_power_of_2(e)
|
||||
need_pad = e_pad != e
|
||||
if need_pad:
|
||||
v = F.pad(v, (0, e_pad - e))
|
||||
|
||||
if d > 128:
|
||||
# split over head
|
||||
if 64 % d:
|
||||
m = 64
|
||||
elif 32 % d:
|
||||
m = 32
|
||||
elif 16 % d:
|
||||
m = 16
|
||||
arr = [m * i for i in range(d // m + 1)]
|
||||
if arr[-1] != d:
|
||||
arr.append(d)
|
||||
n = len(arr)
|
||||
o = 0
|
||||
for i in range(n - 1):
|
||||
start = arr[i]
|
||||
end = arr[i + 1]
|
||||
q1 = q[..., start:end]
|
||||
k1 = k[..., start:end]
|
||||
o += lightning_attn2(q1, k1, v, s)
|
||||
else:
|
||||
o = lightning_attn2(q, k, v, s)
|
||||
|
||||
if need_pad:
|
||||
o = o[:, :, :, :e]
|
||||
|
||||
return o
|
||||
|
||||
|
||||
debug = eval(os.environ.get("debug", default="False"))
|
||||
|
||||
BLOCK = 256
|
||||
|
||||
|
||||
# Copied from transformers.models.llama.modeling_llama.LlamaRMSNorm with Llama->MiniMaxText01
|
||||
class MiniMaxText01RMSNorm(nn.Module):
|
||||
def __init__(self, hidden_size, eps=1e-6):
|
||||
"""
|
||||
MiniMaxText01RMSNorm is equivalent to T5LayerNorm
|
||||
"""
|
||||
super().__init__()
|
||||
self.weight = nn.Parameter(torch.ones(hidden_size))
|
||||
self.variance_epsilon = eps
|
||||
|
||||
def forward(self, hidden_states):
|
||||
input_dtype = hidden_states.dtype
|
||||
hidden_states = hidden_states.to(torch.float32)
|
||||
variance = hidden_states.pow(2).mean(-1, keepdim=True)
|
||||
hidden_states = hidden_states * torch.rsqrt(variance + self.variance_epsilon)
|
||||
return self.weight * hidden_states.to(input_dtype)
|
||||
|
||||
|
||||
# Copied from https://huggingface.co/MiniMaxAI/MiniMax-Text-01/blob/main/modeling_minimax_text_01.py
|
||||
def get_activation_fn(activation):
|
||||
if debug:
|
||||
logger.info(f"activation: {activation}")
|
||||
if activation == "gelu":
|
||||
return F.gelu
|
||||
elif activation == "relu":
|
||||
return F.relu
|
||||
elif activation == "elu":
|
||||
return F.elu
|
||||
elif activation == "sigmoid":
|
||||
return F.sigmoid
|
||||
elif activation == "exp":
|
||||
|
||||
def f(x):
|
||||
with torch.no_grad():
|
||||
x_max = torch.max(x, dim=-1, keepdims=True).values
|
||||
y = torch.exp(x - x_max)
|
||||
|
||||
return y
|
||||
|
||||
return f
|
||||
elif activation == "leak":
|
||||
return F.leaky_relu
|
||||
elif activation == "1+elu":
|
||||
|
||||
def f(x):
|
||||
return 1 + F.elu(x)
|
||||
|
||||
return f
|
||||
elif activation == "2+elu":
|
||||
|
||||
def f(x):
|
||||
return 2 + F.elu(x)
|
||||
|
||||
return f
|
||||
elif activation == "silu" or activation == "swish":
|
||||
return F.silu
|
||||
elif activation == "sine":
|
||||
return torch.sin
|
||||
else:
|
||||
logger.info(f"activation: does not support {activation}, use Identity!!!")
|
||||
return lambda x: x
|
||||
|
||||
|
||||
# Copied from https://huggingface.co/MiniMaxAI/MiniMax-Text-01/blob/main/modeling_minimax_text_01.py
|
||||
class MiniMaxText01LightningAttention(nn.Module):
|
||||
def __init__(self, config=None, layer_idx: Optional[int] = None, **kwargs):
|
||||
super().__init__()
|
||||
if config is None:
|
||||
config = type("Config", (), kwargs)
|
||||
|
||||
bias = False
|
||||
self.hidden_size = config.hidden_size
|
||||
self.num_heads = config.num_attention_heads
|
||||
self.head_dim = getattr(config, "head_dim", self.hidden_size // self.num_heads)
|
||||
|
||||
self.out_proj = nn.Linear(
|
||||
self.head_dim * self.num_heads, self.hidden_size, bias=bias
|
||||
)
|
||||
self.act = get_activation_fn(config.hidden_act)
|
||||
self.norm = MiniMaxText01RMSNorm(self.head_dim * self.num_heads)
|
||||
|
||||
self.qkv_proj = nn.Linear(
|
||||
self.hidden_size, 3 * self.head_dim * self.num_heads, bias=bias
|
||||
)
|
||||
self.output_gate = nn.Linear(
|
||||
self.hidden_size, self.head_dim * self.num_heads, bias=bias
|
||||
)
|
||||
|
||||
# for inference only
|
||||
self.offset = 0
|
||||
self.layer_idx = layer_idx
|
||||
|
||||
def forward(
|
||||
self,
|
||||
hidden_states,
|
||||
attn_mask: Optional[torch.Tensor] = None, # (b, h, n, m)
|
||||
output_attentions: bool = False,
|
||||
past_key_value: Optional[Tuple[torch.Tensor]] = None,
|
||||
use_cache: bool = False,
|
||||
slope_rate: Optional[torch.Tensor] = None,
|
||||
do_eval: bool = False,
|
||||
**kwargs,
|
||||
):
|
||||
if (not self.training) and (not do_eval):
|
||||
return self.inference(
|
||||
hidden_states,
|
||||
attn_mask,
|
||||
output_attentions,
|
||||
past_key_value,
|
||||
use_cache,
|
||||
slope_rate,
|
||||
)
|
||||
|
||||
def inference(
|
||||
self,
|
||||
x,
|
||||
attn_mask: Optional[torch.Tensor] = None, # (b, n)
|
||||
output_attentions: bool = False,
|
||||
past_key_value: Optional[Tuple[torch.Tensor]] = None,
|
||||
use_cache: bool = False,
|
||||
slope_rate: Optional[torch.Tensor] = None, # (h, 1, 1)
|
||||
):
|
||||
# x: b n d
|
||||
b, n, d = x.shape
|
||||
# linear map
|
||||
qkv = self.act(self.qkv_proj(x))
|
||||
new_shape = qkv.size()[:-1] + (self.num_heads, -1)
|
||||
qkv = qkv.view(*new_shape)
|
||||
q, k, v = torch.split(qkv, [self.head_dim] * 3, dim=3)
|
||||
q = q.transpose(1, 2)
|
||||
k = k.transpose(1, 2)
|
||||
v = v.transpose(1, 2)
|
||||
|
||||
if past_key_value is None:
|
||||
self.offset = q.shape[-2]
|
||||
else:
|
||||
self.offset += 1
|
||||
|
||||
# for align with metaseq
|
||||
ratio = torch.exp(-slope_rate)
|
||||
|
||||
# only use for the first time
|
||||
if past_key_value is None:
|
||||
slope_rate = slope_rate.to(torch.float32)
|
||||
if attn_mask is not None:
|
||||
v = v.masked_fill(
|
||||
(1 - attn_mask).unsqueeze(1).unsqueeze(-1).to(torch.bool), 0
|
||||
)
|
||||
NUM_BLOCK = (n + BLOCK - 1) // BLOCK
|
||||
b, h, n, d = q.shape
|
||||
e = v.shape[-1]
|
||||
# other
|
||||
array = torch.arange(BLOCK).to(q) + 1
|
||||
q_decay = torch.exp(-slope_rate * array.reshape(-1, 1))
|
||||
k_decay = torch.exp(-slope_rate * (BLOCK - array.reshape(-1, 1)))
|
||||
index = array[:, None] - array[None, :]
|
||||
s_index = (
|
||||
slope_rate
|
||||
* index[
|
||||
None,
|
||||
None,
|
||||
]
|
||||
)
|
||||
s_index = torch.where(index >= 0, -s_index, float("-inf"))
|
||||
diag_decay = torch.exp(s_index)
|
||||
|
||||
kv = torch.zeros(b, h, d, e).to(torch.float32).to(q.device)
|
||||
output = torch.empty((b, h, n, e), dtype=q.dtype, device=q.device)
|
||||
for i in range(NUM_BLOCK):
|
||||
si = i * BLOCK
|
||||
ei = min(si + BLOCK, n)
|
||||
m = ei - si
|
||||
qi = q[:, :, si:ei].contiguous()
|
||||
ki = k[:, :, si:ei].contiguous()
|
||||
vi = v[:, :, si:ei].contiguous()
|
||||
qkv_none_diag = torch.matmul(qi * q_decay[:, :m], kv).to(torch.float32)
|
||||
|
||||
# diag
|
||||
qk = (
|
||||
torch.matmul(qi, ki.transpose(-1, -2)).to(torch.float32)
|
||||
* diag_decay[:, :, :m, :m]
|
||||
)
|
||||
qkv_diag = torch.matmul(qk, vi.to(torch.float32))
|
||||
block_decay = torch.exp(-slope_rate * m)
|
||||
output[:, :, si:ei] = qkv_none_diag + qkv_diag
|
||||
kv = block_decay * kv + torch.matmul(
|
||||
(ki * k_decay[:, -m:]).transpose(-1, -2).to(vi.dtype), vi
|
||||
)
|
||||
|
||||
else:
|
||||
kv = past_key_value
|
||||
output = []
|
||||
for i in range(n):
|
||||
kv = ratio * kv + torch.einsum(
|
||||
"... n d, ... n e -> ... d e",
|
||||
k[:, :, i : i + 1],
|
||||
v[:, :, i : i + 1],
|
||||
)
|
||||
qkv = torch.einsum(
|
||||
"... n e, ... e d -> ... n d", q[:, :, i : i + 1], kv.to(q.dtype)
|
||||
)
|
||||
output.append(qkv)
|
||||
output = torch.cat(output, dim=-2)
|
||||
# reshape
|
||||
output = rearrange(output, "b h n d -> b n (h d)")
|
||||
# normalize
|
||||
output = self.norm(output)
|
||||
# gate
|
||||
output = F.sigmoid(self.output_gate(x)) * output
|
||||
# outproj
|
||||
output = self.out_proj(output)
|
||||
|
||||
attn_weights = None
|
||||
|
||||
return output, attn_weights, kv
|
||||
|
||||
|
||||
def _build_slope_tensor(n_attention_heads: int):
|
||||
def get_slopes(n):
|
||||
def get_slopes_power_of_2(n):
|
||||
start = 2 ** (-(2 ** -(math.log2(n) - 3)))
|
||||
ratio = start
|
||||
return [start * ratio**i for i in range(n)]
|
||||
|
||||
if math.log2(n).is_integer():
|
||||
return get_slopes_power_of_2(
|
||||
n
|
||||
) # In the paper, we only train models that have 2^a heads for some a. This function has
|
||||
else: # some good properties that only occur when the input is a power of 2. To maintain that even
|
||||
closest_power_of_2 = 2 ** math.floor(
|
||||
math.log2(n)
|
||||
) # when the number of heads is not a power of 2, we use this workaround.
|
||||
return (
|
||||
get_slopes_power_of_2(closest_power_of_2)
|
||||
+ get_slopes(2 * closest_power_of_2)[0::2][: n - closest_power_of_2]
|
||||
)
|
||||
|
||||
# h, 1, 1
|
||||
slopes = torch.tensor(get_slopes(n_attention_heads)).reshape(
|
||||
n_attention_heads, 1, 1
|
||||
)
|
||||
|
||||
return slopes
|
||||
|
||||
|
||||
def test_lightning_attention_implementations(model_params):
|
||||
torch.manual_seed(42)
|
||||
|
||||
batch_size = 2
|
||||
seq_len = 1024
|
||||
dtype = torch.bfloat16
|
||||
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
|
||||
|
||||
hidden_states = torch.randn(
|
||||
batch_size, seq_len, model_params["hidden_size"], dtype=dtype, device=device
|
||||
)
|
||||
|
||||
attention_mask = torch.ones(batch_size, seq_len, dtype=dtype, device=device)
|
||||
|
||||
slope_rate = _build_slope_tensor(model_params["num_attention_heads"]).to(device)
|
||||
|
||||
model_attn = MiniMaxText01LightningAttention(**model_params).to(dtype).to(device)
|
||||
model_attn.eval()
|
||||
|
||||
with torch.no_grad():
|
||||
model_output, _, _ = model_attn.inference(
|
||||
hidden_states, attn_mask=attention_mask, slope_rate=slope_rate
|
||||
)
|
||||
|
||||
qkv = model_attn.act(model_attn.qkv_proj(hidden_states))
|
||||
new_shape = qkv.size()[:-1] + (model_attn.num_heads, -1)
|
||||
qkv = qkv.view(*new_shape)
|
||||
q, k, v = torch.split(qkv, [model_attn.head_dim] * 3, dim=-1)
|
||||
q = q.transpose(1, 2)
|
||||
k = k.transpose(1, 2)
|
||||
v = v.transpose(1, 2)
|
||||
|
||||
lib_output = lightning_attn_func(q, k, v, slope_rate)
|
||||
lib_output = lib_output.transpose(1, 2).contiguous()
|
||||
lib_output = lib_output.view(batch_size, seq_len, -1)
|
||||
lib_output = model_attn.norm(lib_output)
|
||||
lib_output = torch.sigmoid(model_attn.output_gate(hidden_states)) * lib_output
|
||||
lib_output = model_attn.out_proj(lib_output)
|
||||
|
||||
torch.testing.assert_close(
|
||||
model_output,
|
||||
lib_output,
|
||||
rtol=1e-3,
|
||||
atol=1e-2,
|
||||
msg="Lightning attention implementations produce different results",
|
||||
)
|
||||
|
||||
print("✅ Two implementations match")
|
||||
|
||||
|
||||
def get_benchmark():
|
||||
batch_size_range = [2**i for i in range(0, 7)] # max 64
|
||||
seq_length_range = [256, 512, 1024, 2048, 4096] # max 4096
|
||||
configs = list(itertools.product(batch_size_range, seq_length_range))
|
||||
|
||||
@triton.testing.perf_report(
|
||||
triton.testing.Benchmark(
|
||||
x_names=["batch_size", "seq_len"],
|
||||
x_vals=[list(_) for _ in configs],
|
||||
line_arg="provider",
|
||||
line_vals=["MiniMax-Text-01", "OpenNLPLab"],
|
||||
line_names=[
|
||||
"MiniMax-Text-01 Model Implementation",
|
||||
"OpenNLPLab Library Implementation",
|
||||
],
|
||||
styles=[("blue", "-"), ("green", "-")],
|
||||
ylabel="us",
|
||||
plot_name="lightning-attention-prefill-performance",
|
||||
args={},
|
||||
)
|
||||
)
|
||||
def benchmark(batch_size, seq_len, provider):
|
||||
dtype = torch.bfloat16
|
||||
device = torch.device("cuda")
|
||||
|
||||
params = {
|
||||
"hidden_size": 6144,
|
||||
"num_attention_heads": 64,
|
||||
"head_dim": 96,
|
||||
"hidden_act": "gelu",
|
||||
}
|
||||
|
||||
hidden_states = torch.randn(
|
||||
batch_size, seq_len, params["hidden_size"], dtype=dtype, device=device
|
||||
)
|
||||
|
||||
attention_mask = torch.ones(batch_size, seq_len, dtype=dtype, device=device)
|
||||
|
||||
slope_rate = _build_slope_tensor(params["num_attention_heads"]).to(device)
|
||||
model_attn = MiniMaxText01LightningAttention(**params).to(dtype).to(device)
|
||||
model_attn.eval()
|
||||
|
||||
quantiles = [0.5, 0.2, 0.8]
|
||||
if provider == "MiniMax-Text-01":
|
||||
ms, min_ms, max_ms = triton.testing.do_bench(
|
||||
lambda: model_attn.inference(
|
||||
hidden_states, attn_mask=attention_mask, slope_rate=slope_rate
|
||||
),
|
||||
quantiles=quantiles,
|
||||
)
|
||||
else:
|
||||
|
||||
def run_lib():
|
||||
qkv = model_attn.act(model_attn.qkv_proj(hidden_states))
|
||||
new_shape = qkv.size()[:-1] + (model_attn.num_heads, -1)
|
||||
qkv = qkv.view(*new_shape)
|
||||
q, k, v = torch.split(qkv, [model_attn.head_dim] * 3, dim=-1)
|
||||
q = q.transpose(1, 2)
|
||||
k = k.transpose(1, 2)
|
||||
v = v.transpose(1, 2)
|
||||
|
||||
lib_output = lightning_attn_func(q, k, v, slope_rate)
|
||||
lib_output = lib_output.transpose(1, 2).contiguous()
|
||||
lib_output = lib_output.view(batch_size, seq_len, -1)
|
||||
lib_output = model_attn.norm(lib_output)
|
||||
lib_output = (
|
||||
torch.sigmoid(model_attn.output_gate(hidden_states)) * lib_output
|
||||
)
|
||||
return model_attn.out_proj(lib_output)
|
||||
|
||||
ms, min_ms, max_ms = triton.testing.do_bench(
|
||||
run_lib,
|
||||
quantiles=quantiles,
|
||||
)
|
||||
|
||||
return 1000 * ms, 1000 * max_ms, 1000 * min_ms
|
||||
|
||||
return benchmark
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
import argparse
|
||||
|
||||
parser = argparse.ArgumentParser()
|
||||
parser.add_argument(
|
||||
"--save_path",
|
||||
type=str,
|
||||
default="./configs/benchmark_ops/lightning_attention_prefill/",
|
||||
help="Path to save lightning attention prefill benchmark results",
|
||||
)
|
||||
args = parser.parse_args()
|
||||
|
||||
# Run correctness test first
|
||||
# Adapted from https://huggingface.co/MiniMaxAI/MiniMax-Text-01/blob/main/config.json
|
||||
params = {
|
||||
"hidden_size": 6144,
|
||||
"num_attention_heads": 64,
|
||||
"head_dim": 96,
|
||||
"hidden_act": "silu",
|
||||
}
|
||||
test_lightning_attention_implementations(params)
|
||||
|
||||
# Run performance benchmark
|
||||
benchmark = get_benchmark()
|
||||
benchmark.run(print_data=True, save_path=args.save_path)
|
||||
@@ -269,7 +269,6 @@ set(SOURCES
|
||||
"csrc/allreduce/mscclpp_allreduce.cu"
|
||||
"csrc/attention/cascade.cu"
|
||||
"csrc/attention/cutlass_mla_kernel.cu"
|
||||
"csrc/attention/lightning_attention_decode_kernel.cu"
|
||||
"csrc/attention/merge_attn_states.cu"
|
||||
"csrc/attention/vertical_slash_index.cu"
|
||||
"csrc/common_extension.cc"
|
||||
|
||||
@@ -1,312 +0,0 @@
|
||||
import itertools
|
||||
import math
|
||||
import os
|
||||
|
||||
import torch
|
||||
import triton
|
||||
import triton.language as tl
|
||||
from sgl_kernel import lightning_attention_decode
|
||||
|
||||
# CI environment detection
|
||||
IS_CI = (
|
||||
os.getenv("CI", "false").lower() == "true"
|
||||
or os.getenv("GITHUB_ACTIONS", "false").lower() == "true"
|
||||
)
|
||||
|
||||
|
||||
def next_power_of_2(n):
|
||||
return 2 ** (int(math.ceil(math.log(n, 2))))
|
||||
|
||||
|
||||
@triton.jit
|
||||
def _decode_kernel(
|
||||
Q,
|
||||
K,
|
||||
V,
|
||||
KV,
|
||||
Out,
|
||||
S,
|
||||
b: tl.constexpr,
|
||||
h: tl.constexpr,
|
||||
n: tl.constexpr,
|
||||
d: tl.constexpr,
|
||||
d_original: tl.constexpr,
|
||||
e: tl.constexpr,
|
||||
e_original: tl.constexpr,
|
||||
):
|
||||
off_bh = tl.program_id(0)
|
||||
off_h = off_bh % h
|
||||
|
||||
qk_offset = off_bh * n * d
|
||||
v_offset = off_bh * n * e
|
||||
o_offset = off_bh * n * e
|
||||
kv_offset = off_bh * d * e
|
||||
|
||||
s = tl.load(S + off_h)
|
||||
ratio = tl.exp(-s)
|
||||
|
||||
d_idx = tl.arange(0, d)
|
||||
e_idx = tl.arange(0, e)
|
||||
|
||||
# Create masks for original dimensions
|
||||
d_mask = d_idx < d_original
|
||||
e_mask = e_idx < e_original
|
||||
|
||||
# Load with masking
|
||||
q = tl.load(Q + qk_offset + d_idx, mask=d_mask, other=0.0)
|
||||
k = tl.load(K + qk_offset + d_idx, mask=d_mask, other=0.0)
|
||||
v = tl.load(V + v_offset + e_idx, mask=e_mask, other=0.0)
|
||||
|
||||
# Load KV with 2D masking
|
||||
kv = tl.load(
|
||||
KV + kv_offset + d_idx[:, None] * e + e_idx[None, :],
|
||||
mask=(d_mask[:, None] & e_mask[None, :]),
|
||||
other=0.0,
|
||||
)
|
||||
|
||||
# Compute outer product using element-wise operations
|
||||
k_v_prod = k[:, None] * v[None, :]
|
||||
kv = ratio * kv + k_v_prod
|
||||
|
||||
# Store KV with 2D masking
|
||||
tl.store(
|
||||
KV + kv_offset + d_idx[:, None] * e + e_idx[None, :],
|
||||
kv.to(KV.dtype.element_ty),
|
||||
mask=(d_mask[:, None] & e_mask[None, :]),
|
||||
)
|
||||
|
||||
# Compute matrix-vector multiplication using element-wise operations and reduction
|
||||
o = tl.sum(q[:, None] * kv, axis=0)
|
||||
|
||||
# Store output with masking
|
||||
tl.store(Out + o_offset + e_idx, o.to(Out.dtype.element_ty), mask=e_mask)
|
||||
|
||||
|
||||
def triton_lightning_attn_decode(q, k, v, kv, s):
|
||||
"""Triton implementation of Lightning Attention decode operation"""
|
||||
b, h, n, d = q.shape
|
||||
e = v.shape[-1]
|
||||
assert n == 1, "Sequence length must be 1 in decode mode"
|
||||
|
||||
# Get padded dimensions (power of 2)
|
||||
d_padded = next_power_of_2(d)
|
||||
e_padded = next_power_of_2(e)
|
||||
|
||||
# Create output tensor (padded)
|
||||
o_padded = torch.empty(b, h, n, e_padded, dtype=v.dtype, device=v.device)
|
||||
|
||||
# Create padded tensors without actually padding the data
|
||||
q_padded = torch.empty(b, h, n, d_padded, dtype=q.dtype, device=q.device)
|
||||
k_padded = torch.empty(b, h, n, d_padded, dtype=k.dtype, device=k.device)
|
||||
v_padded = torch.empty(b, h, n, e_padded, dtype=v.dtype, device=v.device)
|
||||
kv_padded = torch.empty(
|
||||
b, h, d_padded, e_padded, dtype=torch.float32, device=kv.device
|
||||
)
|
||||
|
||||
# Copy data to padded tensors
|
||||
q_padded[..., :d] = q
|
||||
k_padded[..., :d] = k
|
||||
v_padded[..., :e] = v
|
||||
kv_padded[..., :d, :e] = kv
|
||||
|
||||
# Launch kernel
|
||||
grid = (b * h, 1)
|
||||
_decode_kernel[grid](
|
||||
q_padded,
|
||||
k_padded,
|
||||
v_padded,
|
||||
kv_padded,
|
||||
o_padded,
|
||||
s,
|
||||
b=b,
|
||||
h=h,
|
||||
n=n,
|
||||
d=d_padded,
|
||||
d_original=d,
|
||||
e=e_padded,
|
||||
e_original=e,
|
||||
)
|
||||
|
||||
# Get unpadded outputs
|
||||
o = o_padded[..., :e]
|
||||
kv_out = kv_padded[..., :d, :e]
|
||||
|
||||
return o, kv_out
|
||||
|
||||
|
||||
def lightning_attention_decode_naive(q, k, v, past_kv, slope):
|
||||
"""Naive implementation of lightning attention decode"""
|
||||
original_dtype = q.dtype
|
||||
ratio = torch.exp(-slope) # [h, 1, 1]
|
||||
|
||||
kv = past_kv
|
||||
b, h, n, d = q.shape
|
||||
|
||||
output = []
|
||||
for i in range(n):
|
||||
kv = ratio * kv.to(torch.float32) + torch.einsum(
|
||||
"... n d, ... n e -> ... d e",
|
||||
k[:, :, i : i + 1],
|
||||
v[:, :, i : i + 1],
|
||||
)
|
||||
qkv = torch.einsum(
|
||||
"... n e, ... e d -> ... n d",
|
||||
q[:, :, i : i + 1].to(torch.float32),
|
||||
kv.to(torch.float32),
|
||||
)
|
||||
output.append(qkv)
|
||||
output = torch.cat(output, dim=-2)
|
||||
|
||||
return output.to(original_dtype), kv
|
||||
|
||||
|
||||
def lightning_attention_decode_kernel(q, k, v, past_kv, slope, output, new_kv):
|
||||
return lightning_attention_decode(q, k, v, past_kv, slope, output, new_kv)
|
||||
|
||||
|
||||
def calculate_diff(batch_size):
|
||||
dtype = torch.bfloat16
|
||||
device = torch.device("cuda")
|
||||
num_heads = 64
|
||||
head_dim = 96
|
||||
seq_len = 1
|
||||
|
||||
q = torch.randn(
|
||||
batch_size, num_heads, seq_len, head_dim, device=device, dtype=dtype
|
||||
)
|
||||
k = torch.randn(
|
||||
batch_size, num_heads, seq_len, head_dim, device=device, dtype=dtype
|
||||
)
|
||||
v = torch.randn(
|
||||
batch_size, num_heads, seq_len, head_dim, device=device, dtype=dtype
|
||||
)
|
||||
past_kv = torch.randn(batch_size, num_heads, head_dim, head_dim, device=device)
|
||||
slope = torch.randn(num_heads, 1, 1, device=device)
|
||||
|
||||
output_naive, new_kv_naive = lightning_attention_decode_naive(
|
||||
q.clone(), k.clone(), v.clone(), past_kv.clone(), slope.clone()
|
||||
)
|
||||
|
||||
output_kernel = torch.empty_like(output_naive)
|
||||
new_kv_kernel = torch.empty_like(new_kv_naive)
|
||||
lightning_attention_decode_kernel(
|
||||
q.clone(),
|
||||
k.clone(),
|
||||
v.clone(),
|
||||
past_kv.clone(),
|
||||
slope.clone(),
|
||||
output_kernel,
|
||||
new_kv_kernel,
|
||||
)
|
||||
|
||||
output_triton, new_kv_triton = triton_lightning_attn_decode(
|
||||
q.clone(), k.clone(), v.clone(), past_kv.clone(), slope.clone()
|
||||
)
|
||||
|
||||
if (
|
||||
torch.allclose(output_naive, output_kernel, atol=1e-2, rtol=1e-2)
|
||||
and torch.allclose(output_naive, output_triton, atol=1e-2, rtol=1e-2)
|
||||
and torch.allclose(new_kv_naive, new_kv_kernel, atol=1e-2, rtol=1e-2)
|
||||
and torch.allclose(new_kv_naive, new_kv_triton, atol=1e-2, rtol=1e-2)
|
||||
):
|
||||
print("✅ All implementations match")
|
||||
else:
|
||||
print("❌ Implementations differ")
|
||||
|
||||
|
||||
# Simplified for CI environment
|
||||
if IS_CI:
|
||||
batch_size_range = [1] # Single batch size for CI
|
||||
else:
|
||||
batch_size_range = [i for i in range(1, 65)] # 1 to 64
|
||||
|
||||
configs = [(bs,) for bs in batch_size_range]
|
||||
|
||||
|
||||
@triton.testing.perf_report(
|
||||
triton.testing.Benchmark(
|
||||
x_names=["batch_size"],
|
||||
x_vals=[list(_) for _ in configs],
|
||||
line_arg="provider",
|
||||
line_vals=["naive", "kernel", "triton"],
|
||||
line_names=["PyTorch Naive", "SGL Kernel", "Triton"],
|
||||
styles=[("blue", "-"), ("red", "-"), ("green", "-")],
|
||||
ylabel="us",
|
||||
plot_name="lightning-attention-decode-performance",
|
||||
args={},
|
||||
)
|
||||
)
|
||||
def benchmark(batch_size, provider):
|
||||
dtype = torch.bfloat16
|
||||
device = torch.device("cuda")
|
||||
num_heads = 64
|
||||
head_dim = 96
|
||||
seq_len = 1
|
||||
|
||||
q = torch.randn(
|
||||
batch_size, num_heads, seq_len, head_dim, device=device, dtype=dtype
|
||||
)
|
||||
k = torch.randn(
|
||||
batch_size, num_heads, seq_len, head_dim, device=device, dtype=dtype
|
||||
)
|
||||
v = torch.randn(
|
||||
batch_size, num_heads, seq_len, head_dim, device=device, dtype=dtype
|
||||
)
|
||||
past_kv = torch.randn(batch_size, num_heads, head_dim, head_dim, device=device)
|
||||
slope = torch.randn(num_heads, 1, 1, device=device)
|
||||
|
||||
quantiles = [0.5, 0.2, 0.8]
|
||||
|
||||
if provider == "naive":
|
||||
ms, min_ms, max_ms = triton.testing.do_bench_cudagraph(
|
||||
lambda: lightning_attention_decode_naive(
|
||||
q.clone(), k.clone(), v.clone(), past_kv.clone(), slope.clone()
|
||||
),
|
||||
quantiles=quantiles,
|
||||
)
|
||||
elif provider == "kernel":
|
||||
output = torch.empty(
|
||||
batch_size, num_heads, seq_len, head_dim, device=device, dtype=dtype
|
||||
)
|
||||
new_kv = torch.empty(batch_size, num_heads, head_dim, head_dim, device=device)
|
||||
ms, min_ms, max_ms = triton.testing.do_bench_cudagraph(
|
||||
lambda: lightning_attention_decode_kernel(
|
||||
q.clone(),
|
||||
k.clone(),
|
||||
v.clone(),
|
||||
past_kv.clone(),
|
||||
slope.clone(),
|
||||
output,
|
||||
new_kv,
|
||||
),
|
||||
quantiles=quantiles,
|
||||
)
|
||||
elif provider == "triton":
|
||||
ms, min_ms, max_ms = triton.testing.do_bench_cudagraph(
|
||||
lambda: triton_lightning_attn_decode(
|
||||
q.clone(), k.clone(), v.clone(), past_kv.clone(), slope.clone()
|
||||
),
|
||||
quantiles=quantiles,
|
||||
)
|
||||
|
||||
return 1000 * ms, 1000 * max_ms, 1000 * min_ms
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
import argparse
|
||||
|
||||
parser = argparse.ArgumentParser()
|
||||
parser.add_argument(
|
||||
"--save_path",
|
||||
type=str,
|
||||
default="./configs/benchmark_ops/lightning_attention_decode_sgl/",
|
||||
help="Path to save lightning attention decode benchmark results",
|
||||
)
|
||||
args = parser.parse_args()
|
||||
|
||||
# Run correctness test - simplified for CI
|
||||
test_batch_size = 1 if IS_CI else 4
|
||||
calculate_diff(batch_size=test_batch_size)
|
||||
|
||||
# Run performance benchmark
|
||||
benchmark.run(print_data=True)
|
||||
@@ -1,154 +0,0 @@
|
||||
/* Copyright 2025 SGLang Team. All Rights Reserved.
|
||||
|
||||
Licensed under the Apache License, Version 2.0 (the "License");
|
||||
you may not use this file except in compliance with the License.
|
||||
You may obtain a copy of the License at
|
||||
|
||||
http://www.apache.org/licenses/LICENSE-2.0
|
||||
|
||||
Unless required by applicable law or agreed to in writing, software
|
||||
distributed under the License is distributed on an "AS IS" BASIS,
|
||||
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
See the License for the specific language governing permissions and
|
||||
limitations under the License.
|
||||
==============================================================================*/
|
||||
|
||||
#include <ATen/ATen.h>
|
||||
#include <ATen/cuda/CUDAContext.h>
|
||||
#include <c10/cuda/CUDAGuard.h>
|
||||
#include <cuda_fp16.h>
|
||||
#include <cuda_runtime.h>
|
||||
#include <torch/all.h>
|
||||
|
||||
#define THREADS_PER_BLOCK 128
|
||||
|
||||
template <typename T>
|
||||
__global__ void lightning_attention_decode_kernel(
|
||||
const T* __restrict__ q, // [b, h, 1, d]
|
||||
const T* __restrict__ k, // [b, h, 1, d]
|
||||
const T* __restrict__ v, // [b, h, 1, e]
|
||||
const float* __restrict__ past_kv, // [b, h, d, e]
|
||||
const float* __restrict__ slope, // [h, 1, 1]
|
||||
T* __restrict__ output, // [b, h, 1, e]
|
||||
float* __restrict__ new_kv, // [b, h, d, e]
|
||||
const int batch_size,
|
||||
const int num_heads,
|
||||
const int qk_dim,
|
||||
const int v_dim) {
|
||||
extern __shared__ char smem[];
|
||||
T* __restrict__ q_shared = reinterpret_cast<T*>(smem);
|
||||
T* __restrict__ k_shared = reinterpret_cast<T*>(smem + qk_dim * sizeof(T));
|
||||
T* __restrict__ v_shared = reinterpret_cast<T*>(smem + 2 * qk_dim * sizeof(T));
|
||||
float* __restrict__ new_kv_shared = reinterpret_cast<float*>(smem + (2 * qk_dim + v_dim) * sizeof(T));
|
||||
T* __restrict__ output_shared =
|
||||
reinterpret_cast<T*>(smem + (2 * qk_dim + v_dim) * sizeof(T) + qk_dim * (v_dim + 1) * sizeof(float));
|
||||
|
||||
const int32_t tid = threadIdx.x;
|
||||
const int32_t current_head = blockIdx.x;
|
||||
const int32_t b = current_head / num_heads;
|
||||
const int32_t h = current_head % num_heads;
|
||||
|
||||
if (b >= batch_size) return;
|
||||
|
||||
const int32_t qk_offset = b * num_heads * qk_dim + h * qk_dim;
|
||||
const int32_t v_offset = b * num_heads * v_dim + h * v_dim;
|
||||
const int32_t kv_offset = b * num_heads * qk_dim * v_dim + h * qk_dim * v_dim;
|
||||
|
||||
// Load q, k, v into shared memory
|
||||
for (int d = tid; d < qk_dim; d += blockDim.x) {
|
||||
q_shared[d] = q[qk_offset + d];
|
||||
k_shared[d] = k[qk_offset + d];
|
||||
}
|
||||
for (int e = tid; e < v_dim; e += blockDim.x) {
|
||||
v_shared[e] = v[v_offset + e];
|
||||
}
|
||||
|
||||
__syncthreads();
|
||||
|
||||
const float ratio = expf(-1.0f * slope[h]);
|
||||
|
||||
// Compute new_kv
|
||||
for (int d = tid; d < qk_dim; d += blockDim.x) {
|
||||
const T k_val = k_shared[d];
|
||||
for (int e = 0; e < v_dim; ++e) {
|
||||
const int past_kv_idx = kv_offset + d * v_dim + e;
|
||||
const T v_val = v_shared[e];
|
||||
const float new_val = ratio * past_kv[past_kv_idx] + k_val * v_val;
|
||||
const int shared_idx = d * (v_dim + 1) + e;
|
||||
new_kv_shared[shared_idx] = new_val;
|
||||
}
|
||||
}
|
||||
|
||||
__syncthreads();
|
||||
|
||||
// Store new_kv to global memory
|
||||
for (int idx = tid; idx < qk_dim * v_dim; idx += blockDim.x) {
|
||||
const int d = idx / v_dim;
|
||||
const int e = idx % v_dim;
|
||||
const int shared_idx = d * (v_dim + 1) + e;
|
||||
const int global_idx = kv_offset + idx;
|
||||
new_kv[global_idx] = new_kv_shared[shared_idx];
|
||||
}
|
||||
|
||||
__syncthreads();
|
||||
|
||||
// Compute output
|
||||
for (int e = tid; e < v_dim; e += blockDim.x) {
|
||||
float sum = 0.0f;
|
||||
for (int d = 0; d < qk_dim; ++d) {
|
||||
const int shared_idx = d * (v_dim + 1) + e;
|
||||
sum += q_shared[d] * new_kv_shared[shared_idx];
|
||||
}
|
||||
output_shared[e] = static_cast<T>(sum);
|
||||
}
|
||||
|
||||
__syncthreads();
|
||||
|
||||
// Store output to global memory
|
||||
if (tid == 0) {
|
||||
for (int e = 0; e < v_dim; ++e) {
|
||||
output[v_offset + e] = output_shared[e];
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
void lightning_attention_decode(
|
||||
const torch::Tensor& q,
|
||||
const torch::Tensor& k,
|
||||
const torch::Tensor& v,
|
||||
const torch::Tensor& past_kv,
|
||||
const torch::Tensor& slope,
|
||||
torch::Tensor output,
|
||||
torch::Tensor new_kv) {
|
||||
TORCH_CHECK(q.is_contiguous(), "q must be contiguous");
|
||||
TORCH_CHECK(k.is_contiguous(), "k must be contiguous");
|
||||
TORCH_CHECK(v.is_contiguous(), "v must be contiguous");
|
||||
TORCH_CHECK(past_kv.is_contiguous(), "past_kv must be contiguous");
|
||||
|
||||
auto batch_size = q.size(0);
|
||||
auto num_heads = q.size(1);
|
||||
auto qk_dim = q.size(3);
|
||||
auto v_dim = v.size(3);
|
||||
|
||||
dim3 block(THREADS_PER_BLOCK);
|
||||
dim3 grid(batch_size * num_heads);
|
||||
|
||||
const cudaStream_t stream = at::cuda::getCurrentCUDAStream();
|
||||
|
||||
AT_DISPATCH_FLOATING_TYPES_AND2(
|
||||
at::ScalarType::Half, at::ScalarType::BFloat16, q.scalar_type(), "lightning_attention_decode_kernel", ([&] {
|
||||
size_t smem_size = (2 * qk_dim + 2 * v_dim) * sizeof(scalar_t) + qk_dim * (v_dim + 1) * sizeof(float);
|
||||
lightning_attention_decode_kernel<scalar_t><<<grid, block, smem_size, stream>>>(
|
||||
q.data_ptr<scalar_t>(),
|
||||
k.data_ptr<scalar_t>(),
|
||||
v.data_ptr<scalar_t>(),
|
||||
past_kv.data_ptr<float>(),
|
||||
slope.data_ptr<float>(),
|
||||
output.data_ptr<scalar_t>(),
|
||||
new_kv.data_ptr<float>(),
|
||||
batch_size,
|
||||
num_heads,
|
||||
qk_dim,
|
||||
v_dim);
|
||||
}));
|
||||
}
|
||||
@@ -50,10 +50,6 @@ TORCH_LIBRARY_FRAGMENT(sgl_kernel, m) {
|
||||
/*
|
||||
* From csrc/attention
|
||||
*/
|
||||
m.def(
|
||||
"lightning_attention_decode(Tensor q, Tensor k, Tensor v, Tensor past_kv, Tensor slope, Tensor! output, Tensor! "
|
||||
"new_kv) -> ()");
|
||||
m.impl("lightning_attention_decode", torch::kCUDA, &lightning_attention_decode);
|
||||
m.def("merge_state(Tensor v_a, Tensor s_a, Tensor v_b, Tensor s_b, Tensor! v_merged, Tensor! s_merged) -> ()");
|
||||
m.impl("merge_state", torch::kCUDA, &merge_state);
|
||||
m.def("merge_state_v2(Tensor v_a, Tensor s_a, Tensor v_b, Tensor s_b, Tensor! v_merged, Tensor! s_merged) -> ()");
|
||||
|
||||
@@ -103,14 +103,6 @@ void mscclpp_allreduce(fptr_t _context, torch::Tensor& inp, torch::Tensor& out,
|
||||
/*
|
||||
* From csrc/attention
|
||||
*/
|
||||
void lightning_attention_decode(
|
||||
const torch::Tensor& q,
|
||||
const torch::Tensor& k,
|
||||
const torch::Tensor& v,
|
||||
const torch::Tensor& past_kv,
|
||||
const torch::Tensor& slope,
|
||||
torch::Tensor output,
|
||||
torch::Tensor new_kv);
|
||||
void merge_state(
|
||||
at::Tensor v_a, at::Tensor s_a, at::Tensor v_b, at::Tensor s_b, at::Tensor v_merged, at::Tensor s_merged);
|
||||
void merge_state_v2(
|
||||
|
||||
@@ -13,7 +13,6 @@ from sgl_kernel.allreduce import *
|
||||
from sgl_kernel.attention import (
|
||||
cutlass_mla_decode,
|
||||
cutlass_mla_get_workspace_size,
|
||||
lightning_attention_decode,
|
||||
merge_state,
|
||||
merge_state_v2,
|
||||
)
|
||||
|
||||
@@ -3,12 +3,6 @@ from typing import Optional, Tuple
|
||||
import torch
|
||||
|
||||
|
||||
def lightning_attention_decode(q, k, v, past_kv, slope, output, new_kv):
|
||||
torch.ops.sgl_kernel.lightning_attention_decode.default(
|
||||
q, k, v, past_kv, slope, output, new_kv
|
||||
)
|
||||
|
||||
|
||||
def merge_state(
|
||||
v_a: torch.Tensor,
|
||||
s_a: torch.Tensor,
|
||||
|
||||
@@ -1,84 +0,0 @@
|
||||
import pytest
|
||||
import torch
|
||||
from sgl_kernel import lightning_attention_decode
|
||||
|
||||
|
||||
def naive_lightning_attention_decode(q, k, v, past_kv, slope):
|
||||
"""Naive implementation of lightning attention decode"""
|
||||
original_dtype = q.dtype
|
||||
ratio = torch.exp(-slope) # [h, 1, 1]
|
||||
|
||||
kv = past_kv
|
||||
b, h, n, d = q.shape
|
||||
|
||||
output = []
|
||||
for i in range(n):
|
||||
kv = ratio * kv.to(torch.float32) + torch.einsum(
|
||||
"... n d, ... n e -> ... d e",
|
||||
k[:, :, i : i + 1],
|
||||
v[:, :, i : i + 1],
|
||||
)
|
||||
qkv = torch.einsum(
|
||||
"... n e, ... e d -> ... n d",
|
||||
q[:, :, i : i + 1].to(torch.float32),
|
||||
kv.to(torch.float32),
|
||||
)
|
||||
output.append(qkv)
|
||||
output = torch.cat(output, dim=-2)
|
||||
|
||||
return output.to(original_dtype), kv
|
||||
|
||||
|
||||
configs = [
|
||||
# (batch_size, num_heads, dim, embed_dim)
|
||||
(1, 8, 64, 64),
|
||||
(2, 8, 64, 64),
|
||||
(1, 32, 32, 64),
|
||||
(2, 32, 32, 64),
|
||||
(4, 32, 64, 64),
|
||||
(4, 32, 64, 64),
|
||||
(16, 64, 96, 96),
|
||||
(64, 64, 96, 96),
|
||||
]
|
||||
|
||||
dtypes = [torch.float32, torch.float16, torch.bfloat16]
|
||||
|
||||
|
||||
@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available")
|
||||
@pytest.mark.parametrize("dtype", dtypes)
|
||||
@pytest.mark.parametrize("batch_size,num_heads,dim,embed_dim", configs)
|
||||
def test_lightning_attention_decode(dtype, batch_size, num_heads, dim, embed_dim):
|
||||
device = torch.device("cuda")
|
||||
|
||||
q = torch.randn(batch_size, num_heads, 1, dim, device=device, dtype=dtype)
|
||||
k = torch.randn(batch_size, num_heads, 1, dim, device=device, dtype=dtype)
|
||||
v = torch.randn(batch_size, num_heads, 1, embed_dim, device=device, dtype=dtype)
|
||||
past_kv = torch.randn(batch_size, num_heads, dim, embed_dim, device=device)
|
||||
slope = torch.randn(num_heads, 1, 1, device=device)
|
||||
|
||||
ref_output, ref_new_kv = naive_lightning_attention_decode(q, k, v, past_kv, slope)
|
||||
|
||||
output = torch.empty_like(ref_output)
|
||||
new_kv = torch.empty_like(ref_new_kv)
|
||||
lightning_attention_decode(q, k, v, past_kv, slope, output, new_kv)
|
||||
|
||||
rtol = 1e-2
|
||||
atol = 1e-2
|
||||
|
||||
torch.testing.assert_close(
|
||||
output,
|
||||
ref_output,
|
||||
rtol=rtol,
|
||||
atol=atol,
|
||||
)
|
||||
|
||||
torch.testing.assert_close(
|
||||
new_kv,
|
||||
ref_new_kv,
|
||||
rtol=rtol,
|
||||
atol=atol,
|
||||
)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
pytest.main([__file__])
|
||||
Reference in New Issue
Block a user