diff --git a/benchmark/kernels/minmax-text-01-lightning_attention/benchmark_lightning_attention_decode.py b/benchmark/kernels/minmax-text-01-lightning_attention/benchmark_lightning_attention_decode.py deleted file mode 100644 index 734a0314b..000000000 --- a/benchmark/kernels/minmax-text-01-lightning_attention/benchmark_lightning_attention_decode.py +++ /dev/null @@ -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) diff --git a/benchmark/kernels/minmax-text-01-lightning_attention/benchmark_lightning_attention_prefill.py b/benchmark/kernels/minmax-text-01-lightning_attention/benchmark_lightning_attention_prefill.py deleted file mode 100644 index 9f11ac904..000000000 --- a/benchmark/kernels/minmax-text-01-lightning_attention/benchmark_lightning_attention_prefill.py +++ /dev/null @@ -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) diff --git a/sgl-kernel/CMakeLists.txt b/sgl-kernel/CMakeLists.txt index fcb158924..168b91a2e 100644 --- a/sgl-kernel/CMakeLists.txt +++ b/sgl-kernel/CMakeLists.txt @@ -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" diff --git a/sgl-kernel/benchmark/bench_lightning_attention_decode.py b/sgl-kernel/benchmark/bench_lightning_attention_decode.py deleted file mode 100644 index db0ef05bd..000000000 --- a/sgl-kernel/benchmark/bench_lightning_attention_decode.py +++ /dev/null @@ -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) diff --git a/sgl-kernel/csrc/attention/lightning_attention_decode_kernel.cu b/sgl-kernel/csrc/attention/lightning_attention_decode_kernel.cu deleted file mode 100644 index 01bd4797c..000000000 --- a/sgl-kernel/csrc/attention/lightning_attention_decode_kernel.cu +++ /dev/null @@ -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 -#include -#include -#include -#include -#include - -#define THREADS_PER_BLOCK 128 - -template -__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(smem); - T* __restrict__ k_shared = reinterpret_cast(smem + qk_dim * sizeof(T)); - T* __restrict__ v_shared = reinterpret_cast(smem + 2 * qk_dim * sizeof(T)); - float* __restrict__ new_kv_shared = reinterpret_cast(smem + (2 * qk_dim + v_dim) * sizeof(T)); - T* __restrict__ output_shared = - reinterpret_cast(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(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<<>>( - q.data_ptr(), - k.data_ptr(), - v.data_ptr(), - past_kv.data_ptr(), - slope.data_ptr(), - output.data_ptr(), - new_kv.data_ptr(), - batch_size, - num_heads, - qk_dim, - v_dim); - })); -} diff --git a/sgl-kernel/csrc/common_extension.cc b/sgl-kernel/csrc/common_extension.cc index 6ca57579d..55ce98697 100644 --- a/sgl-kernel/csrc/common_extension.cc +++ b/sgl-kernel/csrc/common_extension.cc @@ -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) -> ()"); diff --git a/sgl-kernel/include/sgl_kernel_ops.h b/sgl-kernel/include/sgl_kernel_ops.h index 953682902..c80536894 100644 --- a/sgl-kernel/include/sgl_kernel_ops.h +++ b/sgl-kernel/include/sgl_kernel_ops.h @@ -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( diff --git a/sgl-kernel/python/sgl_kernel/__init__.py b/sgl-kernel/python/sgl_kernel/__init__.py index 6627a25b5..5dc22ee62 100644 --- a/sgl-kernel/python/sgl_kernel/__init__.py +++ b/sgl-kernel/python/sgl_kernel/__init__.py @@ -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, ) diff --git a/sgl-kernel/python/sgl_kernel/attention.py b/sgl-kernel/python/sgl_kernel/attention.py index f15b4fa24..44dd6111a 100644 --- a/sgl-kernel/python/sgl_kernel/attention.py +++ b/sgl-kernel/python/sgl_kernel/attention.py @@ -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, diff --git a/sgl-kernel/tests/test_lightning_attention_decode.py b/sgl-kernel/tests/test_lightning_attention_decode.py deleted file mode 100644 index f2f0ba258..000000000 --- a/sgl-kernel/tests/test_lightning_attention_decode.py +++ /dev/null @@ -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__])