From 4fc9d3e65874dc86019b2c026675514f428365c5 Mon Sep 17 00:00:00 2001 From: Pan-Yuqi Date: Sat, 9 May 2026 18:00:00 +0800 Subject: [PATCH 1/2] [SSE] Add SSE integration --- fla/__init__.py | 8 + fla/layers/__init__.py | 3 + fla/layers/sse.py | 983 ++++++++++++++++++++++++++++ fla/models/__init__.py | 4 + fla/models/sse/__init__.py | 17 + fla/models/sse/configuration_sse.py | 119 ++++ fla/models/sse/modeling_sse.py | 443 +++++++++++++ fla/ops/sse/__init__.py | 14 + fla/ops/sse/index.py | 36 + fla/ops/sse/mask.py | 359 ++++++++++ 10 files changed, 1986 insertions(+) create mode 100644 fla/layers/sse.py create mode 100644 fla/models/sse/__init__.py create mode 100644 fla/models/sse/configuration_sse.py create mode 100644 fla/models/sse/modeling_sse.py create mode 100644 fla/ops/sse/__init__.py create mode 100644 fla/ops/sse/index.py create mode 100644 fla/ops/sse/mask.py diff --git a/fla/__init__.py b/fla/__init__.py index 92d9467e3a..0138793492 100644 --- a/fla/__init__.py +++ b/fla/__init__.py @@ -34,6 +34,8 @@ RodimusAttention, RWKV6Attention, RWKV7Attention, + SSEGLA, + SSEGDN, ) from fla.models import ( ABCForCausalLM, @@ -86,6 +88,8 @@ RWKV6Model, RWKV7ForCausalLM, RWKV7Model, + SSEForCausalLM, + SSEModel, TransformerForCausalLM, TransformerModel, ) @@ -169,6 +173,10 @@ "RodimusAttention", "RodimusForCausalLM", "RodimusModel", + "SSEGGLA", + "SSEGDN", + "SSEForCausalLM", + "SSEModel", "TransformerForCausalLM", "TransformerModel", ] diff --git a/fla/layers/__init__.py b/fla/layers/__init__.py index b04de7052a..9a8b54326f 100644 --- a/fla/layers/__init__.py +++ b/fla/layers/__init__.py @@ -37,6 +37,7 @@ from .rodimus import RodimusAttention, SlidingWindowSharedKeyAttention from .rwkv6 import RWKV6Attention from .rwkv7 import RWKV7Attention +from .sse import SSEGLA, SSEGDN __all__ = [ 'ABCAttention', @@ -72,4 +73,6 @@ 'ReBasedLinearAttention', 'RodimusAttention', 'SlidingWindowSharedKeyAttention', + 'SSEGLA', + 'SSEGDN', ] diff --git a/fla/layers/sse.py b/fla/layers/sse.py new file mode 100644 index 0000000000..209c2884ca --- /dev/null +++ b/fla/layers/sse.py @@ -0,0 +1,983 @@ +# Copyright (c) 2023-2026, Songlin Yang, Yu Zhang, Zhiyuan Li +# +# This source code is licensed under the MIT license found in the +# LICENSE file in the root directory of this source tree. +# For a list of all contributors, visit: +# https://github.com/fla-org/flash-linear-attention/graphs/contributors + +from __future__ import annotations + +import math +import warnings +from typing import TYPE_CHECKING + +import torch +import torch.nn as nn +from einops import rearrange, repeat +from torch.nn import functional as F + +from fla.layers.utils import get_unpad_data, index_first_axis, pad_input +from fla.modules import FusedRMSNormGated, RMSNorm, ShortConvolution +from fla.ops.gla import chunk_gla, fused_recurrent_gla +from fla.ops.gated_delta_rule import chunk_gated_delta_rule, fused_recurrent_gated_delta_rule +from fla.ops.sse import prepare_sample_relpos_global_index_flat, softmax_and_mask + + +if TYPE_CHECKING: + from transformers.processing_utils import Unpack + + from fla.models.utils import Cache + + +def sort_along_l(q, k, v, gk, beta, e, cu_seqlens, K, emulq, emulk): + _, L, H, D = q.shape + N = e.size(-1) + S = len(cu_seqlens) - 1 + + e = F.softmax(e, dim=-1, dtype=torch.float) + topk_value, topk_expert = torch.topk(e, k=K, dim=2) # [1, L, K] + topk_value = topk_value.to(q.dtype) + mask_w = torch.zeros_like(e, dtype=torch.bool).scatter_(dim=-1, index=topk_expert, src=torch.ones_like(topk_expert, dtype=torch.bool)) + experts_flat = topk_expert.reshape(L * K) # [L*K] + values_flat = topk_value.reshape(L * K) # [L*K] + + sample_idx_flat, relpos_flat, global_idx_flat, lengths = prepare_sample_relpos_global_index_flat(cu_seqlens, K) # ([L*K] * 3, S) + assert sample_idx_flat.dtype == torch.long and relpos_flat.dtype == torch.long and global_idx_flat.dtype == torch.long + + bits_pos = int(lengths.max().item()).bit_length() + bits_exp = int((N - 1)).bit_length() + shift_exp = bits_pos + shift_samp = bits_pos + bits_exp + + ## sort by (sample_idx <- expert_idx <- relpos_in_sample) + key = (sample_idx_flat << shift_samp) | (experts_flat << shift_exp) | relpos_flat + order = torch.argsort(key, stable=False) + experts_sorted = experts_flat.take(order) + sample_sorted = sample_idx_flat.take(order) + global_sorted = global_idx_flat.take(order) # gather index + values_sorted = values_flat.take(order) # sorted eta + # pos_sorted = relpos_flat.take(order) + + ## x: [1, L, H, D] -> y: [1, L*K, H, D] + index4gather = global_sorted[None, :, None, None].expand(1, L * K, H, D) + if beta is None: + q, k, v, gk = [torch.gather(x, dim=1, index=index4gather) for x in (q, k, v, gk)] # GLA + else: + q, k, v = [torch.gather(x, dim=1, index=index4gather) for x in (q, k, v)] # GDN + gk, beta = [torch.gather(x, dim=1, index=index4gather[..., 0]) for x in (gk, beta)] + if emulq: + q = q * values_sorted[None, :, None, None] + if emulk: + k = k * values_sorted[None, :, None, None] + + ## calculate offsets (new cu_seqlens) + pair_id = sample_sorted * N + experts_sorted # [L*K] + counts = torch.bincount(pair_id, minlength=S * N) # [S*N] + state_sizes = counts.view(S, N) + offsets = torch.zeros(1 + S * N, dtype=torch.long, device=q.device) + offsets[1:] = counts.cumsum(dim=0) + offsets = torch.unique(offsets) + + return q, k, v, gk, beta, e, mask_w, offsets, state_sizes, global_sorted + + +class SSEGLA(nn.Module): + """ + The layer implementaion for [SSE: Scaling Linear Attention with Sparse State Expansion](https://arxiv.org/pdf/2507.16577). + + Args: + hidden_size (int, Optional): + The hidden size of the input. Default: 2048. + expand_v (float, Optional): + The expansion ratio for the value dim. Default: 2.0. + head_dim (int, Optional): + The dimension of each head. Default: 256. + num_heads (int, Optional): + The number of heads. Default: 4. + num_v_heads (int, Optional): + The number of heads for the value projection, equal to `num_heads` if `None`. + GVA is applied if `num_v_heads` > `num_heads`. Default: `None`. + mode (str, Optional): + Which GLA kernel to use. + Currently available: `chunk` and `fused_recurrent`. + Default: `chunk`. + use_output_gate (bool, Optional): + Whether to use output gate. Default: `True`. + use_short_conv (bool, Optional): + Whether to use short convolutions. Default: `False`. + conv_size (int, Optional): + The kernel size of the short convolution, only used when `use_short_conv` is `True`. Default: 4. + conv_bias (bool, Optional): + Whether to use bias in the short convolution, only used when `use_short_conv` is `True`. Default: `False`. + num_sparse_partition (int, optional): + Number of state partitions. Default: 4. + num_writer (int, optional): + Top-k write size (number of writers). Default: 1. + num_reader (int, optional): + Top-k read size (number of readers). Default: 1. + sse_implementation (str, optional): + SSE implementation to use. One of `"varlen"` or `"mask"`. Default: `"varlen"`. + use_q_softmax (bool, optional): + Whether to apply softmax to the query. Default: `False`. + use_k_softmax (bool, optional): + Whether to apply softmax to the key. Default: `True`. + emulq (bool, optional): + Whether to use a read gate operating on the state output (Q). Default: `True`. + emulk (bool, optional): + Whether to use a write gate operating on the state input (KV). Default: `True`. + gate_logit_normalizer (int, Optional): + The normalizer for the gate logits, appied after `logsigmoid`. Default: 16. + gate_low_rank_dim (int, Optional): + The low rank dim for the gate projection. Default: 16. + layer_idx (int, Optional): + The index of the layer. Default: None. + norm_eps (float, Optional): + The epsilon value for the normalization layer. Default: 1e-5. + """ + + def __init__( + self, + hidden_size: int = 2048, + expand_v: float = 1., + head_dim: int = 256, + num_heads: int = 6, + num_v_heads: int = None, + mode: str = 'chunk', + use_output_gate: bool = True, + use_short_conv: bool = False, + conv_size: int = 4, + conv_bias: bool = False, + num_sparse_partition: int = 4, + num_writer: int = 1, + num_reader: int = 1, + sse_implementation: str = "varlen", + use_q_softmax: bool = False, + use_k_softmax: bool = True, + emulq: bool = True, + emulk: bool = True, + gate_logit_normalizer: int = 16, + gate_low_rank_dim: int = 16, + layer_idx: int = None, + norm_eps: float = 1e-5, + **kwargs, + ) -> SSEGLA: + super().__init__() + + self.mode = mode + self.hidden_size = hidden_size + self.expand_v = expand_v + + assert num_reader < num_sparse_partition and num_writer < num_sparse_partition, \ + "num_reader and num_writer must be less than num_sparse_partition." + assert sse_implementation in ["mask", "varlen"], \ + f"Unknown SSE implementation {sse_implementation}" + + self.num_sparse_partition = num_sparse_partition + self.num_writer = num_writer + self.num_reader = num_reader + self.sse_implementation = { + "mask": self.sse_linear_attention_mask, + "varlen": self.sse_linear_attention_varlen, + }[sse_implementation] + + self.use_output_gate = use_output_gate + self.use_short_conv = use_short_conv + self.conv_size = conv_size + self.conv_bias = conv_bias + + self.use_q_softmax = use_q_softmax + self.use_k_softmax = use_k_softmax + self.emulq = emulq + self.emulk = emulk + + self.head_dim = head_dim + self.num_heads = num_heads + self.num_v_heads = num_v_heads if num_v_heads is not None else num_heads + + self.head_k_dim = head_dim + self.head_v_dim = int(self.head_dim * self.expand_v) + self.key_dim = int(self.num_heads * self.head_k_dim) + self.value_dim = int(self.num_v_heads * self.head_v_dim) + self.layer_idx = layer_idx + + # Consistency check: Ensure expand_v produces integer values + if not math.isclose(self.num_v_heads * self.head_dim * expand_v, self.value_dim, rel_tol=1e-5): + raise ValueError( + f"expand_v={expand_v} does not produce an integer value when multiplied by key_dim={self.key_dim}. " + f"Resulting value_dim would be {self.num_v_heads * self.head_dim * expand_v}, which is invalid for nn.Linear.", + ) + if self.num_v_heads > self.num_heads and self.num_v_heads % self.num_heads != 0: + raise ValueError( + f"num_v_heads={self.num_v_heads} must be divisible by num_heads={self.num_heads}.", + ) + + if not math.isclose(head_dim * expand_v, self.head_v_dim, rel_tol=1e-5): + raise ValueError( + f"expand_v={expand_v} does not produce an integer value when multiplied by head_dim={head_dim}. " + f"Resulting head_v_dim would be {head_dim * expand_v}, which is invalid for FusedRMSNormGated.", + ) + assert mode in ['chunk', 'fused_recurrent'], f"Not supported mode `{mode}`." + + self.q_proj = nn.Linear(hidden_size, self.key_dim, bias=False) + self.k_proj = nn.Linear(hidden_size, self.key_dim, bias=False) + self.v_proj = nn.Linear(hidden_size, self.value_dim, bias=False) + self.lora_q_proj = nn.Sequential(nn.Linear(hidden_size, self.head_v_dim, bias=False), + nn.Linear(self.head_v_dim, self.key_dim, bias=False)) + self.lora_k_proj = nn.Sequential(nn.Linear(hidden_size, self.head_v_dim, bias=False), + nn.Linear(self.head_v_dim, self.key_dim, bias=False)) + + self.gate_logit_normalizer = gate_logit_normalizer + self.gk_proj = nn.ModuleList([nn.Sequential(nn.Linear(hidden_size, gate_low_rank_dim, bias=False), + nn.Linear(gate_low_rank_dim, self.key_dim, bias=True)) + for _ in range(2)]) + + self.e_proj = nn.Linear(hidden_size, self.num_sparse_partition, bias=False) + + if use_short_conv: + self.conv_size = conv_size + self.q_conv1d_shared = ShortConvolution( + hidden_size=self.key_dim, + kernel_size=conv_size, + bias=conv_bias, + activation=None, + ) + self.k_conv1d_shared = ShortConvolution( + hidden_size=self.key_dim, + kernel_size=conv_size, + bias=conv_bias, + activation=None, + ) + + if use_output_gate: + self.g_proj = nn.Sequential(nn.Linear(hidden_size, self.head_v_dim, bias=False), + nn.Linear(self.head_v_dim, self.value_dim, bias=False)) + self.o_norm = FusedRMSNormGated(self.head_v_dim, eps=norm_eps) + else: + self.o_norm = RMSNorm(self.head_v_dim, eps=norm_eps) + + self.o_proj = nn.Linear(self.value_dim, hidden_size, bias=False) + + def sse_linear_attention_varlen(self, q1, q2, k1, k2, v, gk1, gk2, eta, recurrent_state=None, use_cache=False, cu_seqlens=None): + """ + q1: [bsz, qlen, nhead, head_dim] + q2: [bsz, qlen, nhead, head_dim] + k1: [bsz, klen, nhead, head_dim] + k2: [bsz, klen, nhead, head_dim] + v: [bsz, klen, nhead, head_dim] + gk1: [bsz, klen, nhead, head_dim] + gk2: [bsz, klen, nhead, head_dim] + eta: [bsz, klen, num_sparse_partition] + """ + assert self.num_writer == self.num_reader, "varlen only support num_writer == num_reader" + bsz, q_len, nhead, _ = q1.shape + # change to inference mode. + mode = 'fused_recurrent' if q_len <= 64 else self.mode + if self.training: + assert mode == 'chunk', "Only chunk mode is supported in training." + + v1 = v + v2 = v + if cu_seqlens is None: + cu_seqlens = torch.arange(0, (bsz + 1) * q_len, q_len, dtype=torch.int32, device=q1.device) + q1, k1, gk1, v1 = [rearrange(src, 'b l h d -> 1 (b l) h d').contiguous() for src in [q1, k1, gk1, v]] + q2, k2, gk2, v2 = [rearrange(src, 'b l h d -> 1 (b l) h d').contiguous() for src in [q2, k2, gk2, v]] + S = len(cu_seqlens) - 1 + + if use_cache: + recurrent_state1 = recurrent_state[:S] if recurrent_state is not None else \ + torch.zeros(S, self.num_heads, self.head_dim, self.head_dim).to(torch.float32).to(v.device) + recurrent_state2 = recurrent_state[S:] if recurrent_state is not None else \ + torch.zeros(S*self.num_sparse_partition, self.num_heads, self.head_dim, self.head_dim).to(torch.float32).to(v.device) + + q2, k2, v2, gk2, _, eta, mask, offsets, state_sizes, global_sorted = sort_along_l(q2, k2, v2, gk2, None, eta, cu_seqlens, self.num_writer, self.emulq, self.emulk) + + aux_loss = torch.zeros(()).to(eta) + if self.training: + p = torch.mean(eta.float(), dim=(0, 1)) + f = torch.mean(mask.float(), dim=(0, 1)) + aux_loss = torch.sum(p * f) * self.num_sparse_partition / self.num_writer + + q, k, gk, v = [torch.cat(pair, dim=1) for pair in zip((q1, k1, gk1, v1), (q2, k2, gk2, v2))] + offsets = torch.cat([cu_seqlens.to(offsets), offsets[1:] + cu_seqlens[-1]]) + + recurrent_state_rec = None + if use_cache: + state_id = torch.nonzero(state_sizes.flatten(), as_tuple=True)[0].cpu() + recurrent_state_rec = torch.cat((recurrent_state1, recurrent_state2[state_id]), dim=0) + + if mode == 'fused_recurrent': + o, recurrent_state_rec = fused_recurrent_gla( + q=q, + k=k, + v=v, + gk=gk, + initial_state=recurrent_state_rec, + output_final_state=use_cache, + cu_seqlens=offsets, + ) + elif mode == 'chunk': + o, recurrent_state_rec = chunk_gla( + q=q, + k=k, + v=v, + g=gk, + initial_state=recurrent_state_rec, + output_final_state=use_cache, + cu_seqlens=offsets, + ) + else: + raise NotImplementedError(f"Not supported mode `{mode}`.") + + if recurrent_state_rec is not None: + recurrent_state1 = recurrent_state_rec[:S] + recurrent_state2[state_id] = recurrent_state_rec[S:] + recurrent_state = torch.cat((recurrent_state1, recurrent_state2), dim=0) + else: + recurrent_state = None + + o1, o2 = o[:, :cu_seqlens[-1]], o[:, cu_seqlens[-1]:] + o2_reduce = torch.zeros_like(o1) + o2_reduce.index_add_(dim=1, index=global_sorted, source=o2) + o = o1 + o2_reduce + if bsz > 1: + o = rearrange(o, "1 (b l) h d -> b l h d", b=bsz).contiguous() + + return o, recurrent_state, aux_loss + + def sse_linear_attention_mask(self, q1, q2, k1, k2, v, gk1, gk2, eta, recurrent_state=None, use_cache=False, cu_seqlens=None): + """ + q1: [bsz, qlen, nhead, head_dim] + q2: [bsz, qlen, nhead, head_dim] + k1: [bsz, klen, nhead, head_dim] + k2: [bsz, klen, nhead, head_dim] + v: [bsz, klen, nhead, head_dim] + gk1: [bsz, klen, nhead, head_dim] + gk2: [bsz, klen, nhead, head_dim] + eta: [bsz, klen, num_sparse_partition] + """ + bsz, q_len, nhead, _ = q1.shape + # change to inference mode. + mode = 'fused_recurrent' if q_len <= 64 else self.mode + if self.training: + assert mode == 'chunk', "Only chunk mode is supported in training." + + q2, k2, v2, gk2, eta, mask_w, mask_r = softmax_and_mask(q2, k2, v, gk2, eta, self.num_writer, self.num_reader) + + # writer-only auxloss + aux_loss = torch.zeros(()).to(eta) + if self.training: + p = torch.mean(eta.float(), dim=(0, 1)) + f = torch.mean(mask_w.float(), dim=(0, 1)) + aux_loss = torch.sum(p * f) * self.num_sparse_partition / self.num_writer + + q, k, gk, v = [torch.cat(pair, dim=-2) for pair in zip((q1, k1, gk1, v), (q2, k2, gk2, v2))] + + if mode == 'fused_recurrent': + o, recurrent_state = fused_recurrent_gla( + q=q, + k=k, + v=v, + gk=gk, + initial_state=recurrent_state, + output_final_state=use_cache, + cu_seqlens=cu_seqlens, + ) + elif mode == 'chunk': + o, recurrent_state = chunk_gla( + q=q, + k=k, + v=v, + g=gk, + initial_state=recurrent_state, + output_final_state=use_cache, + cu_seqlens=cu_seqlens, + ) + else: + raise NotImplementedError(f"Not supported mode `{mode}`.") + + o = rearrange(o, "b l (n h) d -> b l n h d", n=self.num_sparse_partition+1) + o = o.sum(2) + + return o, recurrent_state, aux_loss + + def forward( + self, + hidden_states: torch.Tensor, + attention_mask: torch.Tensor | None = None, + past_key_values: Cache | None = None, + use_cache: bool | None = False, + output_attentions: bool | None = False, + **kwargs: Unpack[dict], + ) -> tuple[torch.Tensor, torch.Tensor | None, Cache | None]: + if attention_mask is not None: + assert len(attention_mask.shape) == 2, ( + "Expected attention_mask as a 0-1 matrix with shape [batch_size, seq_len] " + "for padding purposes (0 indicating padding). " + "Arbitrary attention masks of shape [batch_size, seq_len, seq_len] are not allowed." + ) + + batch_size, q_len, _ = hidden_states.shape + + last_state = None + if past_key_values is not None and len(past_key_values) > self.layer_idx: + last_state = past_key_values[self.layer_idx] + + cu_seqlens = kwargs.get('cu_seqlens') + if attention_mask is not None: + indices, cu_seqlens, _ = get_unpad_data(attention_mask[:, -q_len:]) + hidden_states = index_first_axis(rearrange(hidden_states, "b s ... -> (b s) ..."), indices).unsqueeze(0) + + q1 = self.q_proj(hidden_states) + k1 = self.k_proj(hidden_states) + q2 = q1 + self.lora_q_proj(hidden_states) + k2 = k1 + self.lora_k_proj(hidden_states) + v = self.v_proj(hidden_states) + + gk1 = self.gk_proj[0](hidden_states) + gk2 = self.gk_proj[1](hidden_states) + + eta = self.e_proj(hidden_states) + + if self.use_short_conv: + conv_state_q, conv_state_k = None, None + if last_state is not None: + conv_state_q, conv_state_k = last_state['conv_state'] + q1, conv_state_q = self.q_conv1d_shared( + x=q1, + cache=conv_state_q, + output_final_state=use_cache, + cu_seqlens=cu_seqlens, + ) + k1, conv_state_k = self.k_conv1d_shared( + x=k1, + cache=conv_state_k, + output_final_state=use_cache, + cu_seqlens=cu_seqlens, + ) + + q1, q2, k1, k2, gk1, gk2 = map(lambda x: rearrange(x, '... (h d) -> ... h d', d=self.head_k_dim), (q1, q2, k1, k2, gk1, gk2)) + v = rearrange(v, '... (h d) -> ... h d', d=self.head_v_dim) + + if self.use_q_softmax: + q1 = F.softmax(q1.float(), dim=-1).to(v) + q2 = F.softmax(q2.float(), dim=-1).to(v) + else: + q1 = F.silu(q1) + q2 = F.silu(q2) + if self.use_k_softmax: + k1 = F.softmax(k1.float(), dim=-1).to(v) + k2 = F.softmax(k2.float(), dim=-1).to(v) + else: + k1 = F.silu(k1) + k2 = F.silu(k2) + v = F.silu(v) + + gk1 = F.logsigmoid(gk1) / self.gate_logit_normalizer + gk2 = F.logsigmoid(gk2) / self.gate_logit_normalizer + + if self.num_v_heads > self.num_heads: + q1, q2, k1, k2, gk1, gk2 = map(lambda x: repeat(x, '... h d -> ... (h g) d', g=self.num_v_heads // self.num_heads), (q1, q2, k1, k2, gk1, gk2)) + + recurrent_state = last_state['recurrent_state'] if last_state is not None else None + o, recurrent_state, aux_loss = self.sse_implementation( + q1, + q2, + k1, + k2, + v, + gk1, + gk2, + eta, + recurrent_state=recurrent_state, + use_cache=use_cache, + cu_seqlens=cu_seqlens, + ) + + if past_key_values is not None: + past_key_values.update( + recurrent_state=recurrent_state, + conv_state=(conv_state_q, conv_state_k) if self.use_short_conv else None, + layer_idx=self.layer_idx, + offset=q_len, + ) + + if self.use_output_gate: + g = rearrange(self.g_proj(hidden_states), '... (h d) -> ... h d', d=self.head_v_dim) + o = self.o_norm(o, g) + else: + o = self.o_norm(o) + o = rearrange(o, 'b t h d -> b t (h d)') + o = self.o_proj(o) + if attention_mask is not None: + o = pad_input(o.squeeze(0), indices, batch_size, q_len) + + return o, (None, aux_loss), past_key_values + + +class SSEGDN(nn.Module): + """ + The layer implementaion for [SSE: Scaling Linear Attention with Sparse State Expansion](https://arxiv.org/pdf/2507.16577). + + Args: + hidden_size (int, Optional): + The hidden size of the input. Default: 2048. + expand_v (float, Optional): + The expansion ratio for the value dim. Default: 2.0. + head_dim (int, Optional): + The dimension of each head. Default: 256. + num_heads (int, Optional): + The number of heads. Default: 4. + num_v_heads (int, Optional): + The number of heads for the value projection, equal to `num_heads` if `None`. + GVA is applied if `num_v_heads` > `num_heads`. Default: `None`. + mode (str, Optional): + Which Gated DeltaNet kernel to use. + Currently available: `chunk` and `fused_recurrent`. + Default: `chunk`. + use_output_gate (bool, Optional): + Whether to use output gate. Default: `True`. + use_short_conv (bool, Optional): + Whether to use short convolutions. Default: `False`. + allow_neg_eigval (bool, Optional): + Allow negative eigenvalues. Default: `False`. If set to `True`, the beta will be multiplied by 2. + See reference: [Unlocking State-Tracking in Linear RNNs Through Negative Eigenvalues](https://arxiv.org/abs/2411.12537) + conv_size (int, Optional): + The kernel size of the short convolution, only used when `use_short_conv` is `True`. Default: 4. + conv_bias (bool, Optional): + Whether to use bias in the short convolution, only used when `use_short_conv` is `True`. Default: `False`. + num_sparse_partition (int, optional): + Number of state partitions. Default: 4. + num_writer (int, optional): + Top-k write size (number of writers). Default: 1. + num_reader (int, optional): + Top-k read size (number of readers). Default: 1. + sse_implementation (str, optional): + SSE implementation to use. One of `"varlen"` or `"mask"`. Default: `"varlen"`. + use_q_softmax (bool, optional): + Whether to apply softmax to the query. Default: `False`. + use_k_softmax (bool, optional): + Whether to apply softmax to the key. Default: `True`. + emulq (bool, optional): + Whether to use a read gate operating on the state output (Q). Default: `True`. + emulk (bool, optional): + Whether to use a write gate operating on the state input (KV). Default: `True`. + layer_idx (int, Optional): + The index of the layer. Default: None. + norm_eps (float, Optional): + The epsilon value for the normalization layer. Default: 1e-5. + """ + + def __init__( + self, + hidden_size: int = 2048, + expand_v: float = 2, + head_dim: int = 256, + num_heads: int = 6, + num_v_heads: int = None, + mode: str = 'chunk', + use_output_gate: bool = True, + use_short_conv: bool = False, + allow_neg_eigval: bool = False, + conv_size: int = 4, + conv_bias: bool = False, + num_sparse_partition: int = 4, + num_writer: int = 1, + num_reader: int = 1, + sse_implementation: str = "varlen", + use_q_softmax: bool = False, + use_k_softmax: bool = False, + emulq: bool = True, + emulk: bool = True, + layer_idx: int = None, + norm_eps: float = 1e-5, + **kwargs, + ) -> SSEGDN: + super().__init__() + + self.mode = mode + self.allow_neg_eigval = allow_neg_eigval + self.hidden_size = hidden_size + self.expand_v = expand_v + + assert num_reader < num_sparse_partition and num_writer < num_sparse_partition, \ + "num_reader and num_writer must be less than num_sparse_partition." + assert sse_implementation in ["mask", "varlen"], \ + f"Unknown SSE implementation {sse_implementation}" + + self.num_sparse_partition = num_sparse_partition + self.num_writer = num_writer + self.num_reader = num_reader + self.sse_implementation = { + "mask": self.sse_linear_attention_mask, + "varlen": self.sse_linear_attention_varlen, + }[sse_implementation] + + self.use_output_gate = use_output_gate + self.use_short_conv = use_short_conv + self.conv_size = conv_size + self.conv_bias = conv_bias + + self.use_q_softmax = use_q_softmax + self.use_k_softmax = use_k_softmax + self.emulq = emulq + self.emulk = emulk + + self.head_dim = head_dim + self.num_heads = num_heads + self.num_v_heads = num_v_heads if num_v_heads is not None else num_heads + + self.head_k_dim = head_dim + self.head_v_dim = int(self.head_dim * self.expand_v) + self.key_dim = int(self.num_heads * self.head_k_dim) + self.value_dim = int(self.num_v_heads * self.head_v_dim) + self.layer_idx = layer_idx + + # Consistency check: Ensure expand_v produces integer values + if not math.isclose(self.num_v_heads * self.head_dim * expand_v, self.value_dim, rel_tol=1e-5): + raise ValueError( + f"expand_v={expand_v} does not produce an integer value when multiplied by key_dim={self.key_dim}. " + f"Resulting value_dim would be {self.num_v_heads * self.head_dim * expand_v}, which is invalid for nn.Linear.", + ) + if self.num_v_heads > self.num_heads and self.num_v_heads % self.num_heads != 0: + raise ValueError( + f"num_v_heads={self.num_v_heads} must be divisible by num_heads={self.num_heads}.", + ) + + if not math.isclose(head_dim * expand_v, self.head_v_dim, rel_tol=1e-5): + raise ValueError( + f"expand_v={expand_v} does not produce an integer value when multiplied by head_dim={head_dim}. " + f"Resulting head_v_dim would be {head_dim * expand_v}, which is invalid for FusedRMSNormGated.", + ) + assert mode in ['chunk', 'fused_recurrent'], f"Not supported mode `{mode}`." + + self.q_proj = nn.Linear(hidden_size, self.key_dim, bias=False) + self.k_proj = nn.Linear(hidden_size, self.key_dim, bias=False) + self.v_proj = nn.Linear(hidden_size, self.value_dim, bias=False) + self.lora_q_proj = nn.Sequential(nn.Linear(hidden_size, self.head_v_dim, bias=False), + nn.Linear(self.head_v_dim, self.key_dim, bias=False)) + self.lora_k_proj = nn.Sequential(nn.Linear(hidden_size, self.head_v_dim, bias=False), + nn.Linear(self.head_v_dim, self.key_dim, bias=False)) + + self.a_proj = nn.Linear(hidden_size, self.num_v_heads*2, bias=False) + self.b_proj = nn.Linear(hidden_size, self.num_v_heads*2, bias=False) + + A = torch.empty(self.num_v_heads*2, dtype=torch.float32).uniform_(0, 16) + self.A_log = nn.Parameter(torch.log(A)) + self.A_log._no_weight_decay = True + # hard coded for now + dt_min = 0.001 + dt_max = 0.1 + dt_init_floor = 1e-4 + dt = torch.exp( + torch.rand(self.num_v_heads*2) * (math.log(dt_max) - math.log(dt_min)) + + math.log(dt_min), + ) + dt = torch.clamp(dt, min=dt_init_floor) + # Inverse of softplus: https://github.com/pytorch/pytorch/issues/72759 + inv_dt = dt + torch.log(-torch.expm1(-dt)) + self.dt_bias = nn.Parameter(inv_dt) + # Just to be explicit. Without this we already don't put wd on dt_bias because of the check + # name.endswith("bias") in param_grouping.py + self.dt_bias._no_weight_decay = True + + self.e_proj = nn.Linear(hidden_size, self.num_sparse_partition, bias=False) + + if use_short_conv: + self.conv_size = conv_size + self.q_conv1d_shared = ShortConvolution( + hidden_size=self.key_dim, + kernel_size=conv_size, + bias=conv_bias, + activation=None, + ) + self.k_conv1d_shared = ShortConvolution( + hidden_size=self.key_dim, + kernel_size=conv_size, + bias=conv_bias, + activation=None, + ) + + if use_output_gate: + self.g_proj = nn.Sequential(nn.Linear(hidden_size, self.head_v_dim, bias=False), + nn.Linear(self.head_v_dim, self.value_dim, bias=False)) + self.o_norm = FusedRMSNormGated(self.head_v_dim, eps=norm_eps) + else: + self.o_norm = RMSNorm(self.head_v_dim, eps=norm_eps) + + self.o_proj = nn.Linear(self.value_dim, hidden_size, bias=False) + + def sse_linear_attention_varlen(self, q1, q2, k1, k2, v, g1, g2, b1, b2, eta, recurrent_state=None, use_cache=False, cu_seqlens=None): + """ + q1: [bsz, qlen, nhead, head_dim] + q2: [bsz, qlen, nhead, head_dim] + k1: [bsz, klen, nhead, head_dim] + k2: [bsz, klen, nhead, head_dim] + v: [bsz, klen, nhead, head_dim] + g1: [bsz, klen, nhead] + g2: [bsz, klen, nhead] + b1: [bsz, klen, nhead] + b2: [bsz, klen, nhead] + eta: [bsz, klen, num_sparse_partition] + """ + assert self.num_writer == self.num_reader, "varlen only support num_writer == num_reader" + bsz, q_len, nhead, _ = q1.shape + # change to inference mode. + mode = 'fused_recurrent' if q_len // self.num_sparse_partition <= 64 else self.mode + if self.training: + assert mode == 'chunk', "Only chunk mode is supported in training." + + v1 = v + v2 = v + if cu_seqlens is None: + cu_seqlens = torch.arange(0, (bsz + 1) * q_len, q_len, dtype=torch.int32, device=q1.device) + q1, k1, v1 = [rearrange(src, 'b l h d -> 1 (b l) h d').contiguous() for src in [q1, k1, v]] + q2, k2, v2 = [rearrange(src, 'b l h d -> 1 (b l) h d').contiguous() for src in [q2, k2, v]] + g1, g2, b1, b2 = [rearrange(src, 'b l h -> 1 (b l) h').contiguous() for src in [g1, g2, b1, b2]] + S = len(cu_seqlens) - 1 + + if use_cache: + recurrent_state1 = recurrent_state[:S] if recurrent_state is not None else \ + torch.zeros(S, self.num_heads, self.head_dim, self.head_dim).to(torch.float32).to(v.device) + recurrent_state2 = recurrent_state[S:] if recurrent_state is not None else \ + torch.zeros(S*self.num_sparse_partition, self.num_heads, self.head_dim, self.head_dim).to(torch.float32).to(v.device) + + q2, k2, v2, g2, b2, eta, mask, offsets, state_sizes, global_sorted = sort_along_l(q2, k2, v2, g2, b2, eta, cu_seqlens, self.num_writer, self.emulq, self.emulk) + + aux_loss = torch.zeros(()).to(eta) + if self.training: + p = torch.mean(eta.float(), dim=(0, 1)) + f = torch.mean(mask.float(), dim=(0, 1)) + aux_loss = torch.sum(p * f) * self.num_sparse_partition / self.num_writer + + q, k, g, b, v = [torch.cat(pair, dim=1) for pair in zip((q1, k1, g1, b1, v1), (q2, k2, g2, b2, v2))] + offsets = torch.cat([cu_seqlens.to(offsets), offsets[1:] + cu_seqlens[-1]]) + + recurrent_state_rec = None + if use_cache: + state_id = torch.nonzero(state_sizes.flatten(), as_tuple=True)[0].cpu() + recurrent_state_rec = torch.cat((recurrent_state1, recurrent_state2[state_id]), dim=0) + + if mode == 'fused_recurrent': + o, recurrent_state_rec = fused_recurrent_gated_delta_rule( + q=q, + k=k, + v=v, + g=g, + beta=b, + initial_state=recurrent_state_rec, + output_final_state=use_cache, + cu_seqlens=offsets, + use_qk_l2norm_in_kernel=True, + ) + elif mode == 'chunk': + o, recurrent_state_rec = chunk_gated_delta_rule( + q=q, + k=k, + v=v, + g=g, + beta=b, + initial_state=recurrent_state_rec, + output_final_state=use_cache, + cu_seqlens=offsets, + use_qk_l2norm_in_kernel=True, + ) + else: + raise NotImplementedError(f"Not supported mode `{mode}`.") + + if recurrent_state_rec is not None: + recurrent_state1 = recurrent_state_rec[:S] + recurrent_state2[state_id] = recurrent_state_rec[S:] + recurrent_state = torch.cat((recurrent_state1, recurrent_state2), dim=0) + else: + recurrent_state = None + + o1, o2 = o[:, :cu_seqlens[-1]], o[:, cu_seqlens[-1]:] + o2_reduce = torch.zeros_like(o1) + o2_reduce.index_add_(dim=1, index=global_sorted, source=o2) + o = o1 + o2_reduce + if bsz > 1: + o = rearrange(o, "1 (b l) h d -> b l h d", b=bsz).contiguous() + + return o, recurrent_state, aux_loss + + def sse_linear_attention_mask(self, q1, q2, k1, k2, v, g1, g2, b1, b2, eta, recurrent_state=None, use_cache=False, cu_seqlens=None): + """ + q1: [bsz, qlen, nhead, head_dim] + q2: [bsz, qlen, nhead, head_dim] + k1: [bsz, klen, nhead, head_dim] + k2: [bsz, klen, nhead, head_dim] + v: [bsz, klen, nhead, head_dim] + g1: [bsz, klen, nhead] + g2: [bsz, klen, nhead] + b1: [bsz, klen, nhead] + b2: [bsz, klen, nhead] + eta: [bsz, klen, num_sparse_partition] + """ + bsz, q_len, nhead, _ = q1.shape + # change to inference mode. + mode = 'fused_recurrent' if q_len // self.num_sparse_partition <= 64 else self.mode + if self.training: + assert mode == 'chunk', "Only chunk mode is supported in training." + + q2, k2, v2, _, eta, mask_w, mask_r = softmax_and_mask(q2, k2, v, v, eta, self.num_writer, self.num_reader) + g2, b2 = [repeat(x, "b l h -> b l n h", n=self.num_sparse_partition) for x in (g2, b2)] + mask_r = mask_r[..., None] + g2, b2 = g2 * mask_r, b2 * mask_r + g2, b2 = [rearrange(x, "b l n h -> b l (n h)") for x in (g2, b2)] + + # writer-only auxloss + aux_loss = torch.zeros(()).to(eta) + if self.training: + p = torch.mean(eta.float(), dim=(0, 1)) + f = torch.mean(mask_w.float(), dim=(0, 1)) + aux_loss = torch.sum(p * f) * self.num_sparse_partition / self.num_writer + + q, k, g, b, v = [torch.cat(pair, dim=2) for pair in zip((q1, k1, g1, b1, v), (q2, k2, g2, b2, v2))] + + if mode == 'chunk': + o, recurrent_state = chunk_gated_delta_rule( + q=q, + k=k, + v=v, + g=g, + beta=b, + initial_state=recurrent_state, + output_final_state=use_cache, + cu_seqlens=cu_seqlens, + use_qk_l2norm_in_kernel=True, + ) + elif mode == 'fused_recurrent': + o, recurrent_state = fused_recurrent_gated_delta_rule( + q=q, + k=k, + v=v, + g=g, + beta=b, + initial_state=recurrent_state, + output_final_state=use_cache, + cu_seqlens=cu_seqlens, + use_qk_l2norm_in_kernel=True, + ) + else: + raise NotImplementedError(f"Not supported mode `{mode}`.") + + o = rearrange(o, "b l (n h) d -> b l n h d", n=self.num_sparse_partition+1) + o = o.sum(2) + + return o, recurrent_state, aux_loss + + def forward( + self, + hidden_states: torch.Tensor, + attention_mask: torch.Tensor | None = None, + past_key_values: Cache | None = None, + use_cache: bool | None = False, + output_attentions: bool | None = False, + **kwargs: Unpack[dict], + ) -> tuple[torch.Tensor, torch.Tensor | None, Cache | None]: + if attention_mask is not None: + assert len(attention_mask.shape) == 2, ( + "Expected attention_mask as a 0-1 matrix with shape [batch_size, seq_len] " + "for padding purposes (0 indicating padding). " + "Arbitrary attention masks of shape [batch_size, seq_len, seq_len] are not allowed." + ) + + batch_size, q_len, _ = hidden_states.shape + + last_state = None + if past_key_values is not None and len(past_key_values) > self.layer_idx: + last_state = past_key_values[self.layer_idx] + + cu_seqlens = kwargs.get('cu_seqlens') + if attention_mask is not None: + indices, cu_seqlens, _ = get_unpad_data(attention_mask[:, -q_len:]) + hidden_states = index_first_axis(rearrange(hidden_states, "b s ... -> (b s) ..."), indices).unsqueeze(0) + + q1 = self.q_proj(hidden_states) + k1 = self.k_proj(hidden_states) + q2 = q1 + self.lora_q_proj(hidden_states) + k2 = k1 + self.lora_k_proj(hidden_states) + v = self.v_proj(hidden_states) + + b = self.b_proj(hidden_states).sigmoid() + if self.allow_neg_eigval: + b = b * 2. + g = -self.A_log.float().exp() * F.softplus(self.a_proj(hidden_states).float() + self.dt_bias) + b1, b2 = torch.chunk(b, 2, dim=-1) + g1, g2 = torch.chunk(g, 2, dim=-1) + + eta = self.e_proj(hidden_states) + + if self.use_short_conv: + conv_state_q, conv_state_k = None, None + if last_state is not None: + conv_state_q, conv_state_k = last_state['conv_state'] + q1, conv_state_q = self.q_conv1d_shared( + x=q1, + cache=conv_state_q, + output_final_state=use_cache, + cu_seqlens=cu_seqlens, + ) + k1, conv_state_k = self.k_conv1d_shared( + x=k1, + cache=conv_state_k, + output_final_state=use_cache, + cu_seqlens=cu_seqlens, + ) + + q1, q2, k1, k2 = map(lambda x: rearrange(x, '... (h d) -> ... h d', d=self.head_k_dim), (q1, q2, k1, k2)) + v = rearrange(v, '... (h d) -> ... h d', d=self.head_v_dim) + + if self.use_q_softmax: + q1 = F.softmax(q1.float(), dim=-1).to(v) + q2 = F.softmax(q2.float(), dim=-1).to(v) + else: + q1 = F.silu(q1) + q2 = F.silu(q2) + if self.use_k_softmax: + k1 = F.softmax(k1.float(), dim=-1).to(v) + k2 = F.softmax(k2.float(), dim=-1).to(v) + else: + k1 = F.silu(k1) + k2 = F.silu(k2) + v = F.silu(v) + + if self.num_v_heads > self.num_heads: + q1, q2, k1, k2 = map(lambda x: repeat(x, '... h d -> ... (h g) d', g=self.num_v_heads // self.num_heads), (q1, q2, k1, k2)) + + recurrent_state = last_state['recurrent_state'] if last_state is not None else None + o, recurrent_state, aux_loss = self.sse_implementation( + q1, + q2, + k1, + k2, + v, + g1, + g2, + b1, + b2, + eta, + recurrent_state=recurrent_state, + use_cache=use_cache, + cu_seqlens=cu_seqlens, + ) + + if past_key_values is not None: + past_key_values.update( + recurrent_state=recurrent_state, + conv_state=(conv_state_q, conv_state_k) if self.use_short_conv else None, + layer_idx=self.layer_idx, + offset=q_len, + ) + + if self.use_output_gate: + g = rearrange(self.g_proj(hidden_states), '... (h d) -> ... h d', d=self.head_v_dim) + o = self.o_norm(o, g) + else: + o = self.o_norm(o) + o = rearrange(o, 'b t h d -> b t (h d)') + o = self.o_proj(o) + if attention_mask is not None: + o = pad_input(o.squeeze(0), indices, batch_size, q_len) + + return o, (None, aux_loss), past_key_values diff --git a/fla/models/__init__.py b/fla/models/__init__.py index 1b2971e002..09aa3bcdea 100644 --- a/fla/models/__init__.py +++ b/fla/models/__init__.py @@ -39,6 +39,7 @@ from fla.models.rwkv6 import RWKV6Config, RWKV6ForCausalLM, RWKV6Model from fla.models.rwkv7 import RWKV7Config, RWKV7ForCausalLM, RWKV7Model from fla.models.samba import SambaConfig, SambaForCausalLM, SambaModel +from fla.models.sse import SSEConfig, SSEForCausalLM, SSEModel from fla.models.transformer import TransformerConfig, TransformerForCausalLM, TransformerModel __all__ = [ @@ -132,6 +133,9 @@ 'SambaConfig', 'SambaForCausalLM', 'SambaModel', + 'SSEConfig', + 'SSEForCausalLM', + 'SSEModel', 'TransformerConfig', 'TransformerForCausalLM', 'TransformerModel', diff --git a/fla/models/sse/__init__.py b/fla/models/sse/__init__.py new file mode 100644 index 0000000000..e0ea821267 --- /dev/null +++ b/fla/models/sse/__init__.py @@ -0,0 +1,17 @@ +# Copyright (c) 2023-2026, Songlin Yang, Yu Zhang, Zhiyuan Li +# +# This source code is licensed under the MIT license found in the +# LICENSE file in the root directory of this source tree. +# For a list of all contributors, visit: +# https://github.com/fla-org/flash-linear-attention/graphs/contributors + +from transformers import AutoConfig, AutoModel, AutoModelForCausalLM + +from fla.models.sse.configuration_sse import SSEConfig +from fla.models.sse.modeling_sse import SSEForCausalLM, SSEModel + +AutoConfig.register(SSEConfig.model_type, SSEConfig, exist_ok=True) +AutoModel.register(SSEConfig, SSEModel, exist_ok=True) +AutoModelForCausalLM.register(SSEConfig, SSEForCausalLM, exist_ok=True) + +__all__ = ['SSEConfig', 'SSEForCausalLM', 'SSEModel'] diff --git a/fla/models/sse/configuration_sse.py b/fla/models/sse/configuration_sse.py new file mode 100644 index 0000000000..7de15261d9 --- /dev/null +++ b/fla/models/sse/configuration_sse.py @@ -0,0 +1,119 @@ +# Copyright (c) 2023-2026, Songlin Yang, Yu Zhang, Zhiyuan Li +# +# This source code is licensed under the MIT license found in the +# LICENSE file in the root directory of this source tree. +# For a list of all contributors, visit: +# https://github.com/fla-org/flash-linear-attention/graphs/contributors + +import warnings + +from transformers.configuration_utils import PretrainedConfig + + +class SSEConfig(PretrainedConfig): + model_type = 'sse' + keys_to_ignore_at_inference = ['past_key_values'] + + def __init__( + self, + attn_mode: str = "chunk", + hidden_size: int = 2048, + expand_v: float = 1.0, + use_output_gate: bool = True, + use_short_conv: bool = False, + allow_neg_eigval: bool = False, + conv_size: int = 4, + head_dim: int = 256, + num_heads: int = 6, + num_v_heads: int | None = None, + num_sparse_partition: int = 4, + num_writer: int = 2, + num_reader: int = 2, + linear_attn_type: str = "gla", + sse_implementation: str = "varlen", + aux_loss_coef: float = 0.01, + max_position_embeddings: int = 2048, + hidden_ratio: int | None = 4, + intermediate_size: int | None = None, + hidden_act: str = "swish", + num_hidden_layers: int = 24, + norm_eps: float = 1e-6, + attn: dict | None = None, + use_cache: bool = True, + pad_token_id: int | None = None, + bos_token_id: int = 1, + eos_token_id: int = 2, + tie_word_embeddings: bool = False, + initializer_range: float = 0.02, + fuse_norm: bool = True, + fuse_swiglu: bool = True, + fuse_cross_entropy: bool = True, + fuse_linear_cross_entropy: bool = False, + use_l2warp: bool = False, + vocab_size: int = 32000, + **kwargs, + ): + self.attn_mode = attn_mode + self.hidden_size = hidden_size + self.expand_v = expand_v + self.use_output_gate = use_output_gate + self.use_short_conv = use_short_conv + self.conv_size = conv_size + self.head_dim = head_dim + self.num_heads = num_heads + self.num_v_heads = num_v_heads + self.num_sparse_partition = num_sparse_partition + self.num_writer = num_writer + self.num_reader = num_reader + self.linear_attn_type = linear_attn_type + self.sse_implementation = sse_implementation + self.aux_loss_coef = aux_loss_coef + self.max_position_embeddings = max_position_embeddings + + self.hidden_ratio = hidden_ratio + self.intermediate_size = intermediate_size + self.hidden_act = hidden_act + self.num_hidden_layers = num_hidden_layers + self.norm_eps = norm_eps + self.attn = attn + self.use_cache = use_cache + self.initializer_range = initializer_range + + self.fuse_norm = fuse_norm + self.fuse_swiglu = fuse_swiglu + self.fuse_cross_entropy = fuse_cross_entropy + self.fuse_linear_cross_entropy = fuse_linear_cross_entropy + self.use_l2warp = use_l2warp + self.vocab_size = vocab_size + self.allow_neg_eigval = allow_neg_eigval + + if fuse_cross_entropy and fuse_linear_cross_entropy: + raise ValueError( + "`fuse_cross_entropy` and `fuse_linear_cross_entropy` cannot be True at the same time.", + ) + if fuse_linear_cross_entropy: + warnings.warn( + "`fuse_linear_cross_entropy` is enabled, which can improves memory efficiency " + "at the potential cost of reduced precision. " + "If you observe issues like loss divergence, consider disabling this setting.", + ) + + if attn is not None: + if not isinstance(attn, dict): + raise ValueError("attn must be a dictionary") + if 'layers' not in attn: + raise ValueError("Layer indices must be provided to initialize hybrid attention layers") + if 'num_heads' not in attn: + raise ValueError("Number of heads must be provided to initialize hybrid attention layers") + attn['num_kv_heads'] = attn.get('num_kv_heads', attn['num_heads']) + attn['qkv_bias'] = attn.get('qkv_bias', False) + attn['window_size'] = attn.get('window_size', None) + attn['rope_theta'] = attn.get('rope_theta', 10000.) + + super().__init__( + pad_token_id=pad_token_id, + bos_token_id=bos_token_id, + eos_token_id=eos_token_id, + tie_word_embeddings=tie_word_embeddings, + **kwargs, + ) diff --git a/fla/models/sse/modeling_sse.py b/fla/models/sse/modeling_sse.py new file mode 100644 index 0000000000..92e6bfd390 --- /dev/null +++ b/fla/models/sse/modeling_sse.py @@ -0,0 +1,443 @@ +# Copyright (c) 2023-2026, Songlin Yang, Yu Zhang, Zhiyuan Li +# +# This source code is licensed under the MIT license found in the +# LICENSE file in the root directory of this source tree. +# For a list of all contributors, visit: +# https://github.com/fla-org/flash-linear-attention/graphs/contributors + +from __future__ import annotations + +import math +import warnings +from dataclasses import dataclass +from typing import TYPE_CHECKING, Optional, Tuple + +import torch +import torch.nn as nn +from transformers.modeling_outputs import BaseModelOutputWithPast, MoeCausalLMOutputWithPast +from transformers.modeling_utils import PreTrainedModel +from transformers.utils import logging +from transformers.utils.deprecation import deprecate_kwarg + +from fla.layers.attn import Attention +from fla.layers.sse import SSEGLA, SSEGDN +from fla.models.sse.configuration_sse import SSEConfig +from fla.models.utils import Cache, FLAGenerationMixin +from fla.modules import FusedCrossEntropyLoss, FusedLinearCrossEntropyLoss, RMSNorm +from fla.modules import GatedMLP as SSEMLP +from fla.modules.l2warp import l2_warp + +if TYPE_CHECKING: + from transformers.processing_utils import Unpack + + +try: + from transformers.modeling_layers import GradientCheckpointingLayer +except ImportError: + from fla.models.modeling_layers import GradientCheckpointingLayer + +logger = logging.get_logger(__name__) + + +class SSEBlock(GradientCheckpointingLayer): + + def __init__(self, config: SSEConfig, layer_idx: int): + super().__init__() + + self.config = config + self.layer_idx = layer_idx + + self.attn_norm = (RMSNorm if config.fuse_norm else nn.RMSNorm)(config.hidden_size, eps=config.norm_eps, dtype=torch.float32) + if config.attn is not None and layer_idx in config.attn['layers']: + self.attn = Attention( + hidden_size=config.hidden_size, + num_heads=config.attn['num_heads'], + num_kv_heads=config.attn['num_kv_heads'], + qkv_bias=config.attn['qkv_bias'], + window_size=config.attn['window_size'], + rope_theta=config.attn['rope_theta'], + max_position_embeddings=config.max_position_embeddings, + layer_idx=layer_idx, + ) + elif config.linear_attn_type == "gla": + self.attn = SSEGLA( + mode=config.attn_mode, + hidden_size=config.hidden_size, + expand_v=config.expand_v, + head_dim=config.head_dim, + num_heads=config.num_heads, + num_v_heads=config.num_v_heads, + use_output_gate=config.use_output_gate, + use_short_conv=config.use_short_conv, + conv_size=config.conv_size, + num_sparse_partition=config.num_sparse_partition, + num_writer=config.num_writer, + num_reader=config.num_reader, + sse_implementation=config.sse_implementation, + norm_eps=config.norm_eps, + layer_idx=layer_idx, + ) + elif config.linear_attn_type == "gdn": + self.attn = SSEGDN( + mode=config.attn_mode, + hidden_size=config.hidden_size, + expand_v=config.expand_v, + head_dim=config.head_dim, + num_heads=config.num_heads, + num_v_heads=config.num_v_heads, + use_output_gate=config.use_output_gate, + use_short_conv=config.use_short_conv, + allow_neg_eigval=config.allow_neg_eigval, + conv_size=config.conv_size, + num_sparse_partition=config.num_sparse_partition, + num_writer=config.num_writer, + num_reader=config.num_reader, + sse_implementation=config.sse_implementation, + norm_eps=config.norm_eps, + layer_idx=layer_idx, + ) + else: + raise ValueError(f"Unknown linear attention type: {config.linear_attn_type}") + self.mlp_norm = (RMSNorm if config.fuse_norm else nn.RMSNorm)(config.hidden_size, eps=config.norm_eps, dtype=torch.float32) + self.mlp = SSEMLP( + hidden_size=config.hidden_size, + hidden_ratio=config.hidden_ratio, + intermediate_size=config.intermediate_size, + hidden_act=config.hidden_act, + fuse_swiglu=config.fuse_swiglu, + ) + + def forward( + self, + hidden_states: torch.Tensor, + attention_mask: torch.Tensor | None = None, + past_key_values: Cache | list[torch.FloatTensor] | None = None, + use_cache: bool | None = False, + output_attentions: bool | None = False, + **kwargs: Unpack[dict], + ) -> tuple[torch.FloatTensor, tuple[torch.FloatTensor, torch.FloatTensor] | None]: + residual = hidden_states + hidden_states = self.attn_norm(hidden_states) + hidden_states, attentions, past_key_values = self.attn( + hidden_states=hidden_states, + attention_mask=attention_mask, + past_key_values=past_key_values, + use_cache=use_cache, + output_attentions=output_attentions, + **kwargs, + ) + if self.config.fuse_norm: + hidden_states, residual = self.mlp_norm(hidden_states, residual, True) + else: + hidden_states = residual + hidden_states + residual = hidden_states + hidden_states = self.mlp_norm(hidden_states) + hidden_states = self.mlp(hidden_states, **kwargs) + hidden_states = residual + hidden_states + + aux_loss = torch.zeros(()).to(hidden_states) + # Compatible with Attention output + if isinstance(attentions, tuple): + attentions, aux_loss = attentions + + outputs = (hidden_states, attentions, past_key_values, aux_loss) + + return outputs + + +class SSEPreTrainedModel(PreTrainedModel): + + config_class = SSEConfig + base_model_prefix = 'model' + supports_gradient_checkpointing = True + _no_split_modules = ['SSEBlock'] + _supports_cache_class = True + + def __init__(self, *inputs, **kwargs): + super().__init__(*inputs, **kwargs) + + def _init_weights( + self, + module: nn.Module, + prenorm_residual_strategy: str | None = None, + num_residuals_per_layer: int = 2, + ): + if isinstance(module, SSEGDN) and next(module.parameters()).device.type != 'meta': + with torch.no_grad(): + module.A_log.copy_(nn.init.uniform_(module.A_log, a=0, b=16).log()) + module.A_log._no_weight_decay = True + dt = torch.exp( + nn.init.uniform_(module.dt_bias) * (math.log(0.1) - math.log(0.001)) + math.log(0.001), + ).clamp(min=1e-4) + # Inverse of softplus: https://github.com/pytorch/pytorch/issues/72759 + inv_dt = dt + torch.log(-torch.expm1(-dt)) + module.dt_bias.copy_(inv_dt) + module.dt_bias._no_weight_decay = True + + elif isinstance(module, (nn.Linear, nn.Conv1d)): + # Slightly different from the TF version which uses truncated_normal for initialization + # cf https://github.com/pytorch/pytorch/pull/5617 + nn.init.normal_(module.weight, mean=0.0, std=self.config.initializer_range) + if module.bias is not None: + nn.init.zeros_(module.bias) + elif isinstance(module, nn.Embedding): + nn.init.normal_(module.weight, mean=0.0, std=self.config.initializer_range) + elif hasattr(module, 'reset_parameters'): + module.reset_parameters() + + if prenorm_residual_strategy is not None: + # Reinitialize selected weights subject to the OpenAI GPT-2 Paper Scheme: + # > A modified initialization which accounts for the accumulation on the residual path with model depth. Scale + # > the weights of residual layers at initialization by a factor of 1/√N where N is the # of residual layers. + # > -- GPT-2 :: https://openai.com/blog/better-language-models/ + # + # Reference (Megatron-LM): https://github.com/NVIDIA/Megatron-LM/blob/main/megatron/model/gpt_model.py + p = None + if hasattr(module, 'o_proj'): + p = module.o_proj.weight + elif hasattr(module, 'down_proj'): + p = module.down_proj.weight + if p is not None: + # Special Scaled Initialization --> There are 2 Layer Norms per Transformer Block + # Following Pytorch init, except scale by 1/sqrt(2 * n_layer) + # We need to reinit p since this code could be called multiple times + # Having just p *= scale would repeatedly scale it down + if prenorm_residual_strategy == 'rescale': + nn.init.kaiming_uniform_(p, a=math.sqrt(5)) + with torch.no_grad(): + p /= math.sqrt(num_residuals_per_layer * self.config.num_hidden_layers) + elif prenorm_residual_strategy == 'zero': + nn.init.zeros_(p) + else: + raise ValueError(f"Invalid prenorm_residual_strategy: {prenorm_residual_strategy}") + + +@dataclass +class MoeModelOutputWithPastAndAuxLosses(BaseModelOutputWithPast): + """ + Base class for model's outputs, with potential hidden states and attentions. + + Args: + aux_losses (`Optional[Tuple[torch.FloatTensor]]`, *optional*, returned when `labels` is provided): + aux_losses for the sparse modules. + """ + + aux_losses: Optional[Tuple[torch.FloatTensor]] = None + + +class SSEModel(SSEPreTrainedModel): + + def __init__(self, config: SSEConfig): + super().__init__(config) + self.padding_idx = config.pad_token_id + self.vocab_size = config.vocab_size + + self.embeddings = nn.Embedding(config.vocab_size, config.hidden_size, self.padding_idx) + self.layers = nn.ModuleList([SSEBlock(config, layer_idx) for layer_idx in range(config.num_hidden_layers)]) + self.norm = (RMSNorm if config.fuse_norm else nn.RMSNorm)(config.hidden_size, eps=config.norm_eps, dtype=torch.float32) + + self.gradient_checkpointing = False + + self.post_init() + + def get_input_embeddings(self): + return self.embeddings + + def set_input_embeddings(self, value): + self.embeddings = value + + def forward( + self, + input_ids: torch.LongTensor | None = None, + attention_mask: Optional[torch.Tensor] = None, # noqa + inputs_embeds: torch.FloatTensor | None = None, + past_key_values: Cache | list[torch.FloatTensor] | None = None, + use_cache: bool | None = None, + output_attentions: bool | None = None, + output_hidden_states: bool | None = None, + return_dict: bool | None = None, + **kwargs: Unpack[dict], + ) -> tuple | BaseModelOutputWithPast: + if output_attentions: + warnings.warn("`SSEModel` does not `output_attentions` now, setting it to `False`.") + output_attentions = False + output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions + output_hidden_states = output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states + output_aux_losses = True + use_cache = use_cache if use_cache is not None else (self.config.use_cache if not self.training else False) + return_dict = return_dict if return_dict is not None else self.config.use_return_dict + + # retrieve input_ids and inputs_embeds + if input_ids is not None and inputs_embeds is not None: + raise ValueError("You cannot specify both input_ids and inputs_embeds at the same time") + if input_ids is None and inputs_embeds is None: + raise ValueError("You have to specify either input_ids or inputs_embeds") + + if inputs_embeds is None: + inputs_embeds = self.embeddings(input_ids) + hidden_states = inputs_embeds + + if use_cache and not isinstance(past_key_values, Cache): + past_key_values = Cache.from_legacy_cache(past_key_values) + + all_hidden_states = () if output_hidden_states else None + all_attns = () if output_attentions else None + all_aux_losses = () if output_aux_losses else None + for layer in self.layers: + if output_hidden_states: + all_hidden_states += (hidden_states,) + + hidden_states, attentions, past_key_values, aux_loss = layer( + hidden_states, + attention_mask=attention_mask, + past_key_values=past_key_values, + use_cache=use_cache, + output_attentions=output_attentions, + **kwargs, + ) + + if output_attentions: + all_attns += (attentions,) + + if output_aux_losses: + all_aux_losses += (aux_loss,) + + hidden_states = self.norm(hidden_states) + + # add hidden states from the last decoder layer + if output_hidden_states: + all_hidden_states += (hidden_states,) + + if not return_dict: + return tuple(i for i in [hidden_states, past_key_values, all_hidden_states, all_attns, all_aux_losses] if i is not None) + return MoeModelOutputWithPastAndAuxLosses( + last_hidden_state=hidden_states, + past_key_values=past_key_values, + hidden_states=all_hidden_states, + attentions=all_attns, + aux_losses=all_aux_losses, + ) + + +class SSEForCausalLM(SSEPreTrainedModel, FLAGenerationMixin): + + _tied_weights_keys = ["lm_head.weight"] + + def __init__(self, config): + super().__init__(config) + self.model = SSEModel(config) + self.vocab_size = config.vocab_size + self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False) + self.criterion = None + self.aux_loss_coef = config.aux_loss_coef + + # Initialize weights and apply final processing + self.post_init() + + def get_input_embeddings(self): + return self.model.embeddings + + def set_input_embeddings(self, value): + self.model.embeddings = value + + def get_output_embeddings(self): + return self.lm_head + + def set_output_embeddings(self, new_embeddings): + self.lm_head = new_embeddings + + def set_decoder(self, decoder): + self.model = decoder + + def get_decoder(self): + return self.model + + def generate(self, *args, **kwargs): + try: + return super().generate(*args, **kwargs) + except AttributeError as exception: + if 'past_key_values' in str(exception): + raise AttributeError( + f"You tried to call `generate` with a decoding strategy that manipulates `past_key_values`, " + f"which is not supported for {self.__class__.__name__}. " + f"Try another generation strategy instead. " + f"For the available generation strategies, check this doc: " + f"https://huggingface.co/docs/transformers/en/generation_strategies#decoding-strategies", + ) + else: + raise exception + + @deprecate_kwarg("num_logits_to_keep", version="4.50", new_name="logits_to_keep") + def forward( + self, + input_ids: torch.LongTensor = None, + attention_mask: torch.Tensor | None = None, + inputs_embeds: torch.Tensor | None = None, + past_key_values: Cache | list[torch.FloatTensor] | None = None, + labels: torch.LongTensor | None = None, + use_cache: bool | None = None, + output_attentions: bool | None = None, + output_hidden_states: bool | None = None, + return_dict: bool | None = None, + logits_to_keep: int | None = 0, + **kwargs: Unpack[dict], + ) -> tuple | MoeCausalLMOutputWithPast: + output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions + output_hidden_states = ( + output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states + ) + return_dict = return_dict if return_dict is not None else self.config.use_return_dict + + outputs = self.model( + input_ids=input_ids, + attention_mask=attention_mask, + inputs_embeds=inputs_embeds, + past_key_values=past_key_values, + use_cache=use_cache, + output_attentions=output_attentions, + output_hidden_states=output_hidden_states, + return_dict=return_dict, + **kwargs, + ) + + hidden_states = outputs[0] + + loss, aux_loss, logits = None, None, None + if not self.config.fuse_linear_cross_entropy or labels is None: + logits = self.lm_head(hidden_states if logits_to_keep is None else hidden_states[:, -logits_to_keep:]) + if labels is not None: + if getattr(self, 'criterion', None) is None: + if self.config.fuse_linear_cross_entropy: + criterion = FusedLinearCrossEntropyLoss(use_l2warp=self.config.use_l2warp) + elif self.config.fuse_cross_entropy: + criterion = FusedCrossEntropyLoss(inplace_backward=True) + else: + criterion = nn.CrossEntropyLoss() + else: + criterion = self.criterion + labels = labels.to(hidden_states.device) + labels = torch.cat((labels[..., 1:], torch.full_like(labels[:, :1], criterion.ignore_index)), 1) + if self.config.fuse_linear_cross_entropy: + loss = criterion(hidden_states, labels, self.lm_head.weight, self.lm_head.bias) + else: + loss = criterion(logits.view(labels.numel(), -1), labels.view(-1)) + loss = l2_warp(loss, logits) if self.config.use_l2warp else loss + + aux_losses = outputs.aux_losses + compute_device = aux_losses[0].device + aux_loss = sum(layer_aux_loss.to(compute_device) for layer_aux_loss in aux_losses) + + loss += self.aux_loss_coef * aux_loss.to(loss.device) + + if not return_dict: + output = (logits,) + outputs[1:] + return (loss,) + output if loss is not None else output + + return MoeCausalLMOutputWithPast( + loss=loss, + aux_loss=aux_loss, + logits=logits, + past_key_values=outputs.past_key_values, + hidden_states=outputs.hidden_states, + attentions=outputs.attentions, + ) diff --git a/fla/ops/sse/__init__.py b/fla/ops/sse/__init__.py new file mode 100644 index 0000000000..13d93c94d6 --- /dev/null +++ b/fla/ops/sse/__init__.py @@ -0,0 +1,14 @@ +# Copyright (c) 2023-2026, Songlin Yang, Yu Zhang, Zhiyuan Li +# +# This source code is licensed under the MIT license found in the +# LICENSE file in the root directory of this source tree. +# For a list of all contributors, visit: +# https://github.com/fla-org/flash-linear-attention/graphs/contributors + +from .index import prepare_sample_relpos_global_index_flat +from .mask import softmax_and_mask + +__all__ = [ + "prepare_sample_relpos_global_index_flat", + "softmax_and_mask", +] diff --git a/fla/ops/sse/index.py b/fla/ops/sse/index.py new file mode 100644 index 0000000000..171c33b381 --- /dev/null +++ b/fla/ops/sse/index.py @@ -0,0 +1,36 @@ +# Copyright (c) 2023-2026, Songlin Yang, Yu Zhang, Zhiyuan Li +# +# This source code is licensed under the MIT license found in the +# LICENSE file in the root directory of this source tree. +# For a list of all contributors, visit: +# https://github.com/fla-org/flash-linear-attention/graphs/contributors + +import torch +from typing import Tuple + +from fla.utils import tensor_cache + + +@tensor_cache +def prepare_sample_relpos_global_index( + offsets: torch.Tensor +) -> Tuple[torch.LongTensor, torch.LongTensor, torch.LongTensor]: + lengths = offsets[1:] - offsets[:-1] + S = lengths.numel() + sample_idx_per_token = torch.repeat_interleave(torch.arange(S, device=offsets.device), lengths) # [L] + token_global_idx = torch.arange(offsets[-1], device=offsets.device) # [L] + token_start_idx = offsets[:-1].index_select(0, sample_idx_per_token) # [L] + relpos_in_sample = token_global_idx - token_start_idx + return sample_idx_per_token, relpos_in_sample, token_global_idx, lengths + + +@tensor_cache +def prepare_sample_relpos_global_index_flat( + offsets: torch.Tensor, + K: int +) -> Tuple[torch.LongTensor, torch.LongTensor, torch.LongTensor]: + sample_idx_per_token, relpos_in_sample, token_global_idx, lengths = prepare_sample_relpos_global_index(offsets) + sample_idx_flat = sample_idx_per_token[:, None].expand(-1, K).reshape(-1) # [L*K] + relpos_flat = relpos_in_sample[:, None].expand(-1, K).reshape(-1) # [L*K] + global_idx_flat = token_global_idx[:, None].expand(-1, K).reshape(-1) # [L*K] + return sample_idx_flat, relpos_flat, global_idx_flat, lengths diff --git a/fla/ops/sse/mask.py b/fla/ops/sse/mask.py new file mode 100644 index 0000000000..ddb1f5d5be --- /dev/null +++ b/fla/ops/sse/mask.py @@ -0,0 +1,359 @@ +# Copyright (c) 2023-2026, Songlin Yang, Yu Zhang, Zhiyuan Li +# +# This source code is licensed under the MIT license found in the +# LICENSE file in the root directory of this source tree. +# For a list of all contributors, visit: +# https://github.com/fla-org/flash-linear-attention/graphs/contributors + +import torch +import triton +import triton.language as tl +import torch.nn.functional as F + +from fla.utils import input_guard +from fla.ops.utils.softmax import softmax_bwd + + +@triton.autotune( + configs=[ + triton.Config({}, num_warps=num_warps, num_stages=num_stages) + for num_warps in [2, 4, 8] + for num_stages in [2, 3] + ], + key=['BN'], +) +@triton.jit +def _fused_softmax_topk_fwd_kernel( + e, + e_o, + mw, + mr, + stride_e_b, + stride_e_l, + B, + T, + N, + NUM_WRITER: tl.constexpr, + NUM_READER: tl.constexpr, + BN: tl.constexpr, +): + i_b, i_t = tl.program_id(0), tl.program_id(1) + + offsets_n = tl.arange(0, BN) + mask_n = offsets_n < N + p_e = e + i_b * stride_e_b + i_t * stride_e_l + offsets_n + p_e_o = e_o + i_b * stride_e_b + i_t * stride_e_l + offsets_n + p_mw = mw + i_b * stride_e_b + i_t * stride_e_l + offsets_n + p_mr = mr + i_b * stride_e_b + i_t * stride_e_l + offsets_n + + ### stable softmax and topk ### + b_e = tl.load(p_e, mask=mask_n, other=-float('inf')).to(tl.float32) + b_m = tl.max(b_e, axis=0) + b_e = tl.exp(b_e - b_m) + b_p = b_e / tl.sum(b_e, axis=0) + b_p = tl.where(mask_n, b_p.to(p_e.dtype.element_ty), -float('inf')) + b_ps = tl.sort(b_p, descending=True) + tl.store(p_e_o, b_p.to(p_e_o.dtype.element_ty), mask=mask_n) + + mask_w = tl.full((BN,), 1, dtype=b_p.dtype) + if NUM_WRITER < N: + threshold_w = tl.sum(b_ps * (offsets_n == NUM_WRITER - 1)) + mask_w_gr = b_p > threshold_w + need = NUM_WRITER - tl.sum(mask_w_gr.to(tl.int32)) + mask_w_eq = b_p == threshold_w + mask_w_eq_need = mask_w_eq & (tl.cumsum(mask_w_eq.to(tl.int32), axis=0) <= need) + mask_w = mask_w_gr | mask_w_eq_need + mask_w = mask_w.to(b_p.dtype) + tl.store(p_mw, mask_w.to(p_mw.dtype.element_ty), mask=mask_n) + + mask_r = tl.full((BN,), 1, dtype=b_p.dtype) + if NUM_READER < N: + threshold_r = tl.sum(b_ps * (offsets_n == NUM_READER - 1)) + mask_r_gr = b_p > threshold_r + need = NUM_READER - tl.sum(mask_r_gr.to(tl.int32)) + mask_r_eq = b_p == threshold_r + mask_r_eq_need = mask_r_eq & (tl.cumsum(mask_r_eq.to(tl.int32), axis=0) <= need) + mask_r = mask_r_gr | mask_r_eq_need + mask_r = mask_r.to(b_p.dtype) + tl.store(p_mr, mask_r.to(p_mr.dtype.element_ty), mask=mask_n) + + +@triton.autotune( + configs=[ + triton.Config({}, num_warps=num_warps, num_stages=num_stages) + for num_warps in [2, 4, 8] + for num_stages in [2, 3] + ], + key=['BN', 'BK', 'BV'], +) +@triton.jit +def _fused_mask_fwd_kernel( + q, k, v, g, e, mw, mr, + q_o, k_o, v_o, g_o, + stride_k_b, stride_k_l, stride_k_h, + stride_v_b, stride_v_l, stride_v_h, + stride_e_b, stride_e_l, + B, T, N, H, K, V, + BN: tl.constexpr, BK: tl.constexpr, BV: tl.constexpr, +): + i_b, i_t, i_h = tl.program_id(0), tl.program_id(1), tl.program_id(2) + + offsets_n = tl.arange(0, BN) + offsets_k = tl.arange(0, BK) + offsets_v = tl.arange(0, BV) + mask_n = offsets_n < N + mask_k = offsets_k < K + mask_v = offsets_v < V + + p_e = e + i_b * stride_e_b + i_t * stride_e_l + offsets_n + p_mw = mw + i_b * stride_e_b + i_t * stride_e_l + offsets_n + p_mr = mr + i_b * stride_e_b + i_t * stride_e_l + offsets_n + + p_q = q + i_b * stride_k_b + i_t * stride_k_l + i_h * stride_k_h + offsets_k + p_k = k + i_b * stride_k_b + i_t * stride_k_l + i_h * stride_k_h + offsets_k + p_g = g + i_b * stride_k_b + i_t * stride_k_l + i_h * stride_k_h + offsets_k + p_v = v + i_b * stride_v_b + i_t * stride_v_l + i_h * stride_v_h + offsets_v + p_q_o = q_o + i_b * stride_k_b * N + i_t * stride_k_l * N + i_h * stride_k_h \ + + offsets_n[:, None] * H * K + offsets_k[None, :] + p_k_o = k_o + i_b * stride_k_b * N + i_t * stride_k_l * N + i_h * stride_k_h \ + + offsets_n[:, None] * H * K + offsets_k[None, :] + p_g_o = g_o + i_b * stride_k_b * N + i_t * stride_k_l * N + i_h * stride_k_h \ + + offsets_n[:, None] * H * K + offsets_k[None, :] + p_v_o = v_o + i_b * stride_v_b * N + i_t * stride_v_l * N + i_h * stride_v_h \ + + offsets_n[:, None] * H * V + offsets_v[None, :] + + b_e = tl.load(p_e, mask=mask_n, other=0.) + mask_w = tl.load(p_mw, mask=mask_n, other=0.).to(b_e.dtype) + mask_r = tl.load(p_mr, mask=mask_n, other=0.).to(b_e.dtype) + b_e_topk_w = b_e * mask_w + b_e_topk_r = b_e * mask_r + + ### mask qkvg ### + b_q = tl.load(p_q, mask=mask_k, other=0.) + b_q = b_q[None, :] * b_e_topk_r[:, None] + + b_k = tl.load(p_k, mask=mask_k, other=0.) + b_k = b_k[None, :] * b_e_topk_w[:, None] + + b_g = tl.load(p_g, mask=mask_k, other=0.) + b_g = b_g[None, :] * mask_w[:, None] + + b_v = tl.load(p_v, mask=mask_v, other=0.) + b_v = b_v[None, :] * mask_w[:, None] + + mask_nk = mask_n[:, None] & mask_k[None, :] + mask_nv = mask_n[:, None] & mask_v[None, :] + tl.store(p_q_o, b_q.to(p_q_o.dtype.element_ty), mask=mask_nk) + tl.store(p_k_o, b_k.to(p_k_o.dtype.element_ty), mask=mask_nk) + tl.store(p_g_o, b_g.to(p_g_o.dtype.element_ty), mask=mask_nk) + tl.store(p_v_o, b_v.to(p_v_o.dtype.element_ty), mask=mask_nv) + + +@triton.autotune( + configs=[ + triton.Config({}, num_warps=num_warps, num_stages=num_stages) + for num_warps in [2, 4, 8] + for num_stages in [2, 3] + ], + key=['BN', 'BK', 'BV'], +) +@triton.jit +def _fused_mask_bwd_kernel( + q, k, e, mw, mr, + dq_o, dk_o, dv_o, dg_o, + dq, dk, dv, dg, de, + stride_k_b, stride_k_l, stride_k_h, + stride_v_b, stride_v_l, stride_v_h, + stride_e_b, stride_e_l, + B, T, N, H, K, V, + BN: tl.constexpr, BK: tl.constexpr, BV: tl.constexpr, +): + i_b, i_t, i_h = tl.program_id(0), tl.program_id(1), tl.program_id(2) + + offsets_n = tl.arange(0, BN) + offsets_k = tl.arange(0, BK) + offsets_v = tl.arange(0, BV) + mask_n = offsets_n < N + mask_k = offsets_k < K + mask_v = offsets_v < V + + p_e = e + i_b * stride_e_b + i_t * stride_e_l + offsets_n + p_de = de + (i_b * stride_e_b + i_t * stride_e_l + offsets_n) * H + i_h + p_mw = mw + i_b * stride_e_b + i_t * stride_e_l + offsets_n + p_mr = mr + i_b * stride_e_b + i_t * stride_e_l + offsets_n + + p_q = q + i_b * stride_k_b + i_t * stride_k_l + i_h * stride_k_h + offsets_k + p_k = k + i_b * stride_k_b + i_t * stride_k_l + i_h * stride_k_h + offsets_k + p_dq = dq + i_b * stride_k_b + i_t * stride_k_l + i_h * stride_k_h + offsets_k + p_dk = dk + i_b * stride_k_b + i_t * stride_k_l + i_h * stride_k_h + offsets_k + p_dg = dg + i_b * stride_k_b + i_t * stride_k_l + i_h * stride_k_h + offsets_k + p_dv = dv + i_b * stride_v_b + i_t * stride_v_l + i_h * stride_v_h + offsets_v + p_dq_o = dq_o + i_b * stride_k_b * N + i_t * stride_k_l * N + i_h * stride_k_h \ + + offsets_n[:, None] * H * K + offsets_k[None, :] + p_dk_o = dk_o + i_b * stride_k_b * N + i_t * stride_k_l * N + i_h * stride_k_h \ + + offsets_n[:, None] * H * K + offsets_k[None, :] + p_dg_o = dg_o + i_b * stride_k_b * N + i_t * stride_k_l * N + i_h * stride_k_h \ + + offsets_n[:, None] * H * K + offsets_k[None, :] + p_dv_o = dv_o + i_b * stride_v_b * N + i_t * stride_v_l * N + i_h * stride_v_h \ + + offsets_n[:, None] * H * V + offsets_v[None, :] + + b_e = tl.load(p_e, mask=mask_n, other=0.) + mask_w = tl.load(p_mw, mask=mask_n, other=0.).to(b_e.dtype) + mask_r = tl.load(p_mr, mask=mask_n, other=0.).to(b_e.dtype) + b_e_topk_w = b_e * mask_w + b_e_topk_r = b_e * mask_r + + mask_nk = mask_n[:, None] & mask_k[None, :] + mask_nv = mask_n[:, None] & mask_v[None, :] + b_dq_o = tl.load(p_dq_o, mask=mask_nk, other=0.) + b_dk_o = tl.load(p_dk_o, mask=mask_nk, other=0.) + b_dg_o = tl.load(p_dg_o, mask=mask_nk, other=0.) + b_dv_o = tl.load(p_dv_o, mask=mask_nv, other=0.) + b_dq = tl.sum((b_dq_o * b_e_topk_r[:, None]).to(tl.float32), axis=0).to(b_dq_o.dtype) + b_dk = tl.sum((b_dk_o * b_e_topk_w[:, None]).to(tl.float32), axis=0).to(b_dk_o.dtype) + b_dg = tl.sum((b_dg_o * mask_w[:, None]).to(tl.float32), axis=0).to(b_dg_o.dtype) + b_dv = tl.sum((b_dv_o * mask_w[:, None]).to(tl.float32), axis=0).to(b_dv_o.dtype) + + b_q = tl.load(p_q, mask=mask_k, other=0.) + b_k = tl.load(p_k, mask=mask_k, other=0.) + b_de = b_dq_o * b_q[None, :] * mask_r[:, None] + b_dk_o * b_k[None, :] * mask_w[:, None] + b_de = tl.sum(b_de.to(tl.float32), axis=1).to(b_de.dtype) + + tl.store(p_de, b_de.to(p_de.dtype.element_ty), mask=mask_n) + tl.store(p_dq, b_dq.to(p_dq.dtype.element_ty), mask=mask_k) + tl.store(p_dk, b_dk.to(p_dk.dtype.element_ty), mask=mask_k) + tl.store(p_dg, b_dg.to(p_dg.dtype.element_ty), mask=mask_k) + tl.store(p_dv, b_dv.to(p_dv.dtype.element_ty), mask=mask_v) + + +class SoftmaxAndMask(torch.autograd.Function): + r""" + Applies softmax to router weights, repeats and masks inputs, + scales queries and keys with the router weights, and generates reader/writer masks. + + Notation: + B: batch size + T: sequence length + H: number of attention heads + K: key/query head dimension + V: value head dimension + N: number of state partitions + + Args: + q (torch.Tensor): + Queries of shape `(B, T, H, K)`. + k (torch.Tensor): + Keys of shape `(B, T, H, K)`. + v (torch.Tensor): + Values of shape `(B, T, H, V)`. + g (torch.Tensor): + Gates of shape `(B, T, H, V)`. + e (torch.Tensor): + Router weights before softmax of shape `(B, T, N)`. + num_writer (int): + Number of state partitions to write. + num_reader (int): + Number of state partitions to read. + + Returns: + q_out (torch.Tensor): + Repeated and masked queries of shape `(B, T, N * H, K)`. + k_out (torch.Tensor): + Repeated and masked keys of shape `(B, T, N * H, K)`. + v_out (torch.Tensor): + Repeated and masked values of shape `(B, T, N * H, V)`. + g_out (torch.Tensor): + Repeated and masked gates of shape `(B, T, N * H, V)`. + e_out (torch.Tensor): + Router weights after softmax of shape `(B, T, N)`. + mask_w (torch.Tensor): + Writer mask of shape `(B, T, N)`. + mask_r (torch.Tensor): + Reader mask of shape `(B, T, N)`. + """ + + @staticmethod + @input_guard + def forward(ctx, q, k, v, g, e, num_writer, num_reader): + B, T, H, K, V, N = *k.shape, v.shape[-1], e.shape[-1] + BN = triton.next_power_of_2(N) + BK = triton.next_power_of_2(K) + BV = triton.next_power_of_2(V) + + q_out = q.new_empty(B, T, N * H, K) + k_out = k.new_empty(B, T, N * H, K) + v_out = v.new_empty(B, T, N * H, V) + g_out = g.new_empty(B, T, N * H, K) + e_out = torch.empty_like(e) + mask_w = torch.empty_like(e, dtype=torch.int32) + mask_r = torch.empty_like(e, dtype=torch.int32) + + _fused_softmax_topk_fwd_kernel[(B, T)]( + e, + e_out, + mask_w, + mask_r, + e.stride(0), + e.stride(1), + B, + T, + N, + NUM_WRITER=num_writer, + NUM_READER=num_reader, + BN=BN, + ) + + _fused_mask_fwd_kernel[(B, T, H)]( + q, k, v, g, e_out, mask_w, mask_r, + q_out, k_out, v_out, g_out, + k.stride(0), k.stride(1), k.stride(2), + v.stride(0), v.stride(1), v.stride(2), + e.stride(0), e.stride(1), + B, T, N, H, K, V, + BN=BN, + BK=BK, + BV=BV, + ) + + ctx.save_for_backward(q, k, v, g, e_out, mask_w, mask_r) + ctx.num_writer = num_writer + ctx.num_reader = num_reader + return q_out, k_out, v_out, g_out, e_out, mask_w, mask_r + + @staticmethod + @input_guard + def backward(ctx, dq_out, dk_out, dv_out, dg_out, de_out, dmask_w, dmask_r): + q, k, v, g, e_out, mask_w, mask_r = ctx.saved_tensors + + B, T, H, K, V, N = *k.shape, v.shape[-1], e_out.shape[-1] + BN = triton.next_power_of_2(N) + BK = triton.next_power_of_2(K) + BV = triton.next_power_of_2(V) + + dq = torch.empty_like(q) + dk = torch.empty_like(k) + dv = torch.empty_like(v) + dg = torch.empty_like(g) + de = g.new_empty(B, T, N, H) + + grid = (B, T, H) + + _fused_mask_bwd_kernel[grid]( + q, k, e_out, mask_w, mask_r, + dq_out, dk_out, dv_out, dg_out, + dq, dk, dv, dg, de, + k.stride(0), k.stride(1), k.stride(2), + v.stride(0), v.stride(1), v.stride(2), + e_out.stride(0), e_out.stride(1), + B, T, N, H, K, V, + BN=BN, + BK=BK, + BV=BV, + ) + + de = de.sum(dim=-1).add_(de_out) + de = softmax_bwd(e_out, de, dtype=de.dtype) + + return dq.to(q), dk.to(k), dv.to(v), dg.to(g), de.to(e_out), None, None + +softmax_and_mask = SoftmaxAndMask.apply From 741fbbf279a102944c39facdfe1fdf323aa85f5f Mon Sep 17 00:00:00 2001 From: Pan-Yuqi <116128711+Pan-Yuqi@users.noreply.github.com> Date: Sat, 9 May 2026 18:13:09 +0800 Subject: [PATCH 2/2] Update fla/__init__.py Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com> --- fla/__init__.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/fla/__init__.py b/fla/__init__.py index 0138793492..cab97fb151 100644 --- a/fla/__init__.py +++ b/fla/__init__.py @@ -173,7 +173,7 @@ "RodimusAttention", "RodimusForCausalLM", "RodimusModel", - "SSEGGLA", + "SSEGLA", "SSEGDN", "SSEForCausalLM", "SSEModel",