chenchenghao/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

[sgl-kernel Code Clean] Remove useless lightning_attention kernel (#13819)

Browse Source
This commit is contained in:
Xiaoyu Zhang
2025-11-24 18:26:25 +08:00
committed by GitHub
parent aeac622058
commit ecefc7904f
10 changed files with 0 additions and 1754 deletions
Download Patch File Download Diff File Expand all files Collapse all files

1
sgl-kernel/CMakeLists.txt
View File

@@ -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"

312
sgl-kernel/benchmark/bench_lightning_attention_decode.py
View File

@@ -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)

154
sgl-kernel/csrc/attention/lightning_attention_decode_kernel.cu
View File

@@ -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);
}));
}

4
sgl-kernel/csrc/common_extension.cc
View File

@@ -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) -> ()");

8
sgl-kernel/include/sgl_kernel_ops.h
View File

@@ -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(

1
sgl-kernel/python/sgl_kernel/__init__.py
View File

@@ -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,
)

6
sgl-kernel/python/sgl_kernel/attention.py
View File

@@ -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,

84
sgl-kernel/tests/test_lightning_attention_decode.py
View File

@@ -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__])