Skip to content

Add cu_seqlens support to OlmoHybridGatedDeltaNet for packed sequences #44836

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Closed
tyler-romero wants to merge 4 commits into from
Closed

Add cu_seqlens support to OlmoHybridGatedDeltaNet for packed sequences #44836

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
4 commits
Select commit Hold shift + click to select a range
5b7b7f3
Add cu_seqlens support to OlmoHybridGatedDeltaNet for packed sequences
Mar 18, 2026
39fea8f
Simplify conv1d calls by always passing both use_precomputed and cu_s…
Mar 18, 2026
67fb636
Simplify unpad/repad to use boolean mask indexing
Mar 18, 2026
6713d5e
Apply ruff formatting
Mar 18, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
58 changes: 48 additions & 10 deletions src/transformers/models/olmo_hybrid/modeling_olmo_hybrid.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -621,6 +621,24 @@ def torch_recurrent_gated_delta_rule(
)


def _cu_seqlens_from_packed_mask(attention_mask: torch.Tensor) -> torch.Tensor:

@vasqu vasqu Mar 19, 2026

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I have to intervene a bit here. We should not manually unpad and pad - this is something that I explicitly disallowed in the model addition PR

We can still support padding-free training/forward tho by preparing the input properly before passing it the model. We are kinda stuck with FA supporting it with position ids only as well which was a mistake imo.

Tl;dr: Use our data collator

transformers/src/transformers/data/data_collator.py

Line 1364 in 16a5b09

class DataCollatorWithFlattening(DefaultDataCollator):

And pass the prepared input with its metadata (cu seq lens, max seq lens etc)

"""Derive ``cu_seqlens`` from a packed attention mask with unique sequence IDs.

For a mask like ``[1, 1, 1, 2, 2, 0, 0]``, returns ``cu_seqlens = [0, 3, 5]``
(ignoring padding). For a standard ``0/1`` mask, returns ``[0, num_ones]``.
"""
flat = attention_mask.flatten()
non_pad = flat > 0
non_pad_ids = flat[non_pad]
if len(non_pad_ids) == 0:
return torch.tensor([0], dtype=torch.int32, device=attention_mask.device)
boundaries = torch.where(non_pad_ids[1:] != non_pad_ids[:-1])[0] + 1
cu_seqlens = torch.zeros(len(boundaries) + 2, dtype=torch.int32, device=attention_mask.device)
cu_seqlens[1:-1] = boundaries
cu_seqlens[-1] = len(non_pad_ids)
return cu_seqlens


class OlmoHybridGatedDeltaNet(nn.Module):
"""
GatedDeltaNet linear attention for OLMo Hybrid.
Expand Down Expand Up @@ -719,14 +737,26 @@ def forward(
attention_mask: torch.Tensor | None = None,
**kwargs: Unpack[TransformersKwargs],
) -> torch.Tensor:
# Requires LEFT padding to work correctly
hidden_states = apply_mask_to_padding_states(hidden_states, attention_mask)
Comment on lines -722 to -723

@vasqu vasqu Mar 19, 2026

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This will stay then


batch_size, seq_len, _ = hidden_states.shape

use_cache = cache_params is not None
use_precomputed = use_cache and getattr(cache_params, "has_previous_state", False) and seq_len == 1

# For packed sequences (attention_mask with unique sequence IDs > 1), derive
# cu_seqlens and unpad so recurrent state doesn't leak across sequence boundaries.
# Requires the FLA fast path; torch fallbacks don't support cu_seqlens.
cu_seqlens = None
unpad_indices = None
if attention_mask is not None and not use_precomputed and is_fast_path_available and attention_mask.max() > 1:
cu_seqlens = _cu_seqlens_from_packed_mask(attention_mask)
unpad_indices = attention_mask.flatten() > 0
hidden_states = hidden_states[:, unpad_indices, :]
else:
# Requires LEFT padding to work correctly
hidden_states = apply_mask_to_padding_states(hidden_states, attention_mask)

effective_batch, effective_len, _ = hidden_states.shape

conv_state_q = cache_params.conv_states_q[self.layer_idx] if cache_params else None
conv_state_k = cache_params.conv_states_k[self.layer_idx] if cache_params else None
conv_state_v = cache_params.conv_states_v[self.layer_idx] if cache_params else None
Expand All @@ -737,23 +767,23 @@ def forward(
v = self.v_proj(hidden_states)

q, new_conv_state_q = self.q_conv1d(
q, cache=conv_state_q, use_precomputed=use_precomputed, output_final_state=use_cache
q, cache=conv_state_q, use_precomputed=use_precomputed, output_final_state=use_cache, cu_seqlens=cu_seqlens
)
k, new_conv_state_k = self.k_conv1d(
k, cache=conv_state_k, use_precomputed=use_precomputed, output_final_state=use_cache
k, cache=conv_state_k, use_precomputed=use_precomputed, output_final_state=use_cache, cu_seqlens=cu_seqlens
)
v, new_conv_state_v = self.v_conv1d(
v, cache=conv_state_v, use_precomputed=use_precomputed, output_final_state=use_cache
v, cache=conv_state_v, use_precomputed=use_precomputed, output_final_state=use_cache, cu_seqlens=cu_seqlens
Comment on lines +770 to +776

@vasqu vasqu Mar 19, 2026

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Should be hidden behind kwargs if possible (from the signature)

I hope FLA has **kwargs in its signature and it doesnt destroy anything

)

if cache_params is not None:
cache_params.conv_states_q[self.layer_idx] = new_conv_state_q
cache_params.conv_states_k[self.layer_idx] = new_conv_state_k
cache_params.conv_states_v[self.layer_idx] = new_conv_state_v

q = q.view(batch_size, seq_len, -1, self.head_k_dim)
k = k.view(batch_size, seq_len, -1, self.head_k_dim)
v = v.view(batch_size, seq_len, -1, self.head_v_dim)
q = q.view(effective_batch, effective_len, -1, self.head_k_dim)
k = k.view(effective_batch, effective_len, -1, self.head_k_dim)
v = v.view(effective_batch, effective_len, -1, self.head_v_dim)

if self.num_v_heads > self.num_k_heads:
expand_ratio = self.num_v_heads // self.num_k_heads
Expand All @@ -778,6 +808,7 @@ def forward(
use_qk_l2norm_in_kernel=True,
)
else:
chunk_extra_kwargs = {"cu_seqlens": cu_seqlens} if cu_seqlens is not None else {}
output, new_recurrent_state = self.chunk_gated_delta_rule(
q,
k,
Expand All @@ -787,6 +818,7 @@ def forward(
initial_state=recurrent_state,
output_final_state=use_cache,
use_qk_l2norm_in_kernel=True,
**chunk_extra_kwargs,
)

if cache_params is not None:
Expand All @@ -796,10 +828,16 @@ def forward(
output = output.reshape(-1, self.head_v_dim)
gate = gate.reshape(-1, self.head_v_dim)
output = self.o_norm(output, gate)
output = output.reshape(batch_size, seq_len, -1)
output = output.reshape(effective_batch, effective_len, -1)

output = self.o_proj(output)

# Re-pad output to original shape for packed sequences
if unpad_indices is not None:
output_padded = output.new_zeros(batch_size, seq_len, output.shape[-1])
output_padded[:, unpad_indices, :] = output
output = output_padded

return output


Expand Down
58 changes: 48 additions & 10 deletions src/transformers/models/olmo_hybrid/modular_olmo_hybrid.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -391,6 +391,24 @@ def forward(self, x, position_ids):
return cos, sin


def _cu_seqlens_from_packed_mask(attention_mask: torch.Tensor) -> torch.Tensor:
"""Derive ``cu_seqlens`` from a packed attention mask with unique sequence IDs.

For a mask like ``[1, 1, 1, 2, 2, 0, 0]``, returns ``cu_seqlens = [0, 3, 5]``
(ignoring padding). For a standard ``0/1`` mask, returns ``[0, num_ones]``.
"""
flat = attention_mask.flatten()
non_pad = flat > 0
non_pad_ids = flat[non_pad]
if len(non_pad_ids) == 0:
return torch.tensor([0], dtype=torch.int32, device=attention_mask.device)
boundaries = torch.where(non_pad_ids[1:] != non_pad_ids[:-1])[0] + 1
cu_seqlens = torch.zeros(len(boundaries) + 2, dtype=torch.int32, device=attention_mask.device)
cu_seqlens[1:-1] = boundaries
cu_seqlens[-1] = len(non_pad_ids)
return cu_seqlens


class OlmoHybridGatedDeltaNet(nn.Module):
"""
GatedDeltaNet linear attention for OLMo Hybrid.
Expand Down Expand Up @@ -489,14 +507,26 @@ def forward(
attention_mask: torch.Tensor | None = None,
**kwargs: Unpack[TransformersKwargs],
) -> torch.Tensor:
# Requires LEFT padding to work correctly
hidden_states = apply_mask_to_padding_states(hidden_states, attention_mask)

batch_size, seq_len, _ = hidden_states.shape

use_cache = cache_params is not None
use_precomputed = use_cache and getattr(cache_params, "has_previous_state", False) and seq_len == 1

# For packed sequences (attention_mask with unique sequence IDs > 1), derive
# cu_seqlens and unpad so recurrent state doesn't leak across sequence boundaries.
# Requires the FLA fast path; torch fallbacks don't support cu_seqlens.
cu_seqlens = None
unpad_indices = None
if attention_mask is not None and not use_precomputed and is_fast_path_available and attention_mask.max() > 1:
cu_seqlens = _cu_seqlens_from_packed_mask(attention_mask)
unpad_indices = attention_mask.flatten() > 0
hidden_states = hidden_states[:, unpad_indices, :]
else:
# Requires LEFT padding to work correctly
hidden_states = apply_mask_to_padding_states(hidden_states, attention_mask)

effective_batch, effective_len, _ = hidden_states.shape

conv_state_q = cache_params.conv_states_q[self.layer_idx] if cache_params else None
conv_state_k = cache_params.conv_states_k[self.layer_idx] if cache_params else None
conv_state_v = cache_params.conv_states_v[self.layer_idx] if cache_params else None
Expand All @@ -507,23 +537,23 @@ def forward(
v = self.v_proj(hidden_states)

q, new_conv_state_q = self.q_conv1d(
q, cache=conv_state_q, use_precomputed=use_precomputed, output_final_state=use_cache
q, cache=conv_state_q, use_precomputed=use_precomputed, output_final_state=use_cache, cu_seqlens=cu_seqlens
)
k, new_conv_state_k = self.k_conv1d(
k, cache=conv_state_k, use_precomputed=use_precomputed, output_final_state=use_cache
k, cache=conv_state_k, use_precomputed=use_precomputed, output_final_state=use_cache, cu_seqlens=cu_seqlens
)
v, new_conv_state_v = self.v_conv1d(
v, cache=conv_state_v, use_precomputed=use_precomputed, output_final_state=use_cache
v, cache=conv_state_v, use_precomputed=use_precomputed, output_final_state=use_cache, cu_seqlens=cu_seqlens
)

if cache_params is not None:
cache_params.conv_states_q[self.layer_idx] = new_conv_state_q
cache_params.conv_states_k[self.layer_idx] = new_conv_state_k
cache_params.conv_states_v[self.layer_idx] = new_conv_state_v

q = q.view(batch_size, seq_len, -1, self.head_k_dim)
k = k.view(batch_size, seq_len, -1, self.head_k_dim)
v = v.view(batch_size, seq_len, -1, self.head_v_dim)
q = q.view(effective_batch, effective_len, -1, self.head_k_dim)
k = k.view(effective_batch, effective_len, -1, self.head_k_dim)
v = v.view(effective_batch, effective_len, -1, self.head_v_dim)

if self.num_v_heads > self.num_k_heads:
expand_ratio = self.num_v_heads // self.num_k_heads
Expand All @@ -548,6 +578,7 @@ def forward(
use_qk_l2norm_in_kernel=True,
)
else:
chunk_extra_kwargs = {"cu_seqlens": cu_seqlens} if cu_seqlens is not None else {}
output, new_recurrent_state = self.chunk_gated_delta_rule(
q,
k,
Expand All @@ -557,6 +588,7 @@ def forward(
initial_state=recurrent_state,
output_final_state=use_cache,
use_qk_l2norm_in_kernel=True,
**chunk_extra_kwargs,
)

if cache_params is not None:
Expand All @@ -566,10 +598,16 @@ def forward(
output = output.reshape(-1, self.head_v_dim)
gate = gate.reshape(-1, self.head_v_dim)
output = self.o_norm(output, gate)
output = output.reshape(batch_size, seq_len, -1)
output = output.reshape(effective_batch, effective_len, -1)

output = self.o_proj(output)

# Re-pad output to original shape for packed sequences
if unpad_indices is not None:
output_padded = output.new_zeros(batch_size, seq_len, output.shape[-1])
output_padded[:, unpad_indices, :] = output
output = output_padded

return output


Expand Down
Loading