monkeypatch.utils

monkeypatch.utils

Shared utils for the monkeypatches

Functions

Name	Description
get_cu_seqlens	generate a cumulative sequence length mask for flash attention using attn mask
get_cu_seqlens_from_pos_ids	generate a cumulative sequence length mask for flash attention using pos ids
get_max_seqlen_in_batch	Token counts for each packed sub-sequence in a multipack attention mask.

get_cu_seqlens

monkeypatch.utils.get_cu_seqlens(attn_mask)

generate a cumulative sequence length mask for flash attention using attn mask

get_cu_seqlens_from_pos_ids

monkeypatch.utils.get_cu_seqlens_from_pos_ids(position_ids)

generate a cumulative sequence length mask for flash attention using pos ids

get_max_seqlen_in_batch

monkeypatch.utils.get_max_seqlen_in_batch(attention_mask)

Token counts for each packed sub-sequence in a multipack attention mask.

Multipack encodes every packed document as a distinct positive integer id (1, 2, 3, ...) with 0 for padding, so a row looks like [1, 1, 1, 2, 2, 3, 3, 0, 0]. This returns the length of each non-pad document, flattened across the batch.

Vectorized on purpose: the previous implementation used int(mask.max().item()) as a Python loop bound, which forced a device->host sync and broke torch.compile. This variant computes segment boundaries with tensor ops only.