[SSE] Add SSE integration - #882
Conversation
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request introduces the SSE (Scaling Linear Attention with Sparse State Expansion) architecture, implementing SSEGLA and SSEGDN layers, optimized Triton kernels, and the SSEForCausalLM model. Feedback identifies a typo in the all export list and a critical indexing bug in the Triton kernels where query strides are incorrectly applied to gate tensors. Additionally, recommendations were made to remove performance-degrading host-device synchronizations during state handling and to synchronize default parameter values between the configuration and layer classes for better consistency.
|
|
||
| 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 |
There was a problem hiding this comment.
The _fused_mask_fwd_kernel uses stride_k_* to access the g tensor. This is incorrect when g has different strides than k, which happens in SSEGDN.sse_linear_attention_mask where v (with head dimension V) is passed as the g argument. Since v has a different head dimension and different strides than k, this leads to incorrect memory indexing. The kernel should accept stride_g_* or the caller should ensure strides match.
| 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 |
There was a problem hiding this comment.
Similar to the forward kernel, _fused_mask_bwd_kernel uses stride_k_* for dg. This will cause incorrect gradient updates when dg has different strides than k (e.g., in SSEGDN).
|
|
||
| recurrent_state_rec = None | ||
| if use_cache: | ||
| state_id = torch.nonzero(state_sizes.flatten(), as_tuple=True)[0].cpu() |
There was a problem hiding this comment.
Calling .cpu() on state_id triggers a host-device synchronization, which can be a significant performance bottleneck during inference. Since recurrent_state2 is on the GPU, you should keep state_id on the GPU as well for indexing.
| state_id = torch.nonzero(state_sizes.flatten(), as_tuple=True)[0].cpu() | |
| state_id = torch.nonzero(state_sizes.flatten(), as_tuple=True)[0] |
|
|
||
| recurrent_state_rec = None | ||
| if use_cache: | ||
| state_id = torch.nonzero(state_sizes.flatten(), as_tuple=True)[0].cpu() |
There was a problem hiding this comment.
Similar to the SSEGLA implementation, avoid moving state_id to the CPU to prevent unnecessary synchronization overhead.
| state_id = torch.nonzero(state_sizes.flatten(), as_tuple=True)[0].cpu() | |
| state_id = torch.nonzero(state_sizes.flatten(), as_tuple=True)[0] |
| num_writer: int = 2, | ||
| num_reader: int = 2, |
There was a problem hiding this comment.
The default values for num_writer and num_reader in SSEConfig (2) do not match the defaults in the SSEGLA and SSEGDN layers (1). It is recommended to synchronize these defaults to avoid confusion.
| num_writer: int = 2, | |
| num_reader: int = 2, | |
| num_writer: int = 1, | |
| num_reader: int = 1, |
Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>
1 participant
Description
This PR adds SSE integration, where SSE refers to Scaling Linear Attention with Sparse State Expansion.
The implementation introduces the SSE model/layer/operator components into the repository and wires them into the existing module structure so that SSE can be used consistently with other supported architectures.
Changes
fla/layers/fla/models/sse/fla/ops/sse/__init__.pyfiles