Does the lightning_attn interface really trainable?

[Arthur-wza]

Dear authors:
when I use lightning_attn_func to replace LlamaAttention in Llama model, I find the convergence of model not really good.
What's more, I don't find a model in this github project, could you plz tell me your experience of training with lightning_attn_func?
Thank you very much!

here are my calling code, if there is any mistake plz let me know.
from lightning_attn.ops import lightning_attn_func
from lightning_attn.utils import _build_slope_tensor
class LightningAttention(nn.Module):
def init(self, config: LlamaConfig, layer_idx: Optional[int] = None):
super().init()
...........
def forward(
) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]:

   ..............
    key_states = repeat_kv(key_states, self.num_key_value_groups)
    value_states = repeat_kv(value_states, self.num_key_value_groups)

    s = _build_slope_tensor(self.config.num_attention_heads).to(query_states.device).to(torch.float32)

    attn_output = lightning_attn_func(
        query_states,
        key_states,
        value_states,
        s
    )

    attn_output = attn_output.transpose(1, 2).contiguous()

    attn_output = attn_output.reshape(bsz, q_len, -1)

    if self.config.pretraining_tp > 1:
        attn_output = attn_output.split(self.hidden_size // self.config.pretraining_tp, dim=2)
        o_proj_slices = self.o_proj.weight.split(self.hidden_size // self.config.pretraining_tp, dim=1)
        attn_output = sum([F.linear(attn_output[i], o_proj_slices[i]) for i in range(self.config.pretraining_tp)])
    else:
        attn_output = self.o_proj(attn_output)

    return attn_output, past_key_value

Thanks again!

[Doraemonzzz]

Hello, this kernel is trainable. Please note that the input shape is (batch, num_heads, seqlen, head_dim). I will update the complete network architecture in the coming days.

[Doraemonzzz]

You can find the full usage example at: https://github.com/Doraemonzzz/xmixers/blob/main/xmixers/modules/token_mixers/linear_attention/tnl_attention.py. This implementation uses the xops library to implement Lightning Attention. Please note that the input shape is (batch, seqlen, num_heads, head_dim). You can also benchmark this kernel, as you mentioned in another issue.

[Arthur-wza]

thanks a lot！I will try on your reply.

[Arthur-wza]

dear authors, when I want to follow your implementation athttps://github.com/Doraemonzzz/xmixers/blob/main/xmixers/modules/token_mixers/linear_attention/tnl_attention.py, I find that your xopes repo do not have installation instruction, could you plz tell me how to install your xopes package?（or use it as a local package）
By the way, do you have wechat ? If you send your wx number to my email Arthurwzaa@gmail.com, it will be my honor to communicate deeper with your lightning_attn real-time.

[Arthur-wza]

Furthermore, my understanding is that Lightning Attention employs a dual-mode approach during inference. It utilizes its parallelized tiling algorithm for the initial prompt processing(prefilling) . Subsequently, for the token-by-token autoregressive generation, it reverts to a standard, non-tiled linear attention mechanism.
Am I understanding this correctly?Thanks a lot.