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

benchmark/kernels/minmax-text-01-lightning_attention/benchmark_lightning_attention_decode.py

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

benchmark/kernels/minmax-text-01-lightning_attention/benchmark_lightning_attention_prefill.py

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

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"

sgl-kernel/benchmark/bench_lightning_attention_decode.py

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

sgl-kernel/csrc/attention/lightning_attention_decode_kernel.cu

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

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

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(

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

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,

sgl-kernel/tests/test_lightning_attention_decode.py

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