Skip to content

Cache packed sequence metadata to reduce D2H syncs across layers #4243

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

Merged
danielhanchen merged 5 commits into from
Mar 12, 2026
Merged

Cache packed sequence metadata to reduce D2H syncs across layers #4243

Changes from all commits
Commits
Show all changes
5 commits
Select commit Hold shift + click to select a range
2efe2ef
packing optimziation with cache to reduce D2H copy
ruixiang63 Feb 18, 2026
837b5f8
[pre-commit.ci] auto fixes from pre-commit.com hooks
pre-commit-ci[bot] Mar 6, 2026
67930ba
cache per device to avoid race condition for multi-gpu
ruixiang63 Mar 11, 2026
ecc5106
[pre-commit.ci] auto fixes from pre-commit.com hooks
pre-commit-ci[bot] Mar 11, 2026
60da93d
add cache freeing up func
ruixiang63 Mar 11, 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
64 changes: 60 additions & 4 deletions unsloth/utils/packing.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -36,6 +36,15 @@
_XFORMERS_MASK_CACHE_MAXSIZE = 32
_XFORMERS_MASK_CACHE: OrderedDict[Tuple[Tuple[int, ...], int], Any] = OrderedDict()

# Cache per device for get_packed_info_from_kwargs to avoid repeated D2H sync across layers
_PACKED_INFO_CACHE: dict = {}

# Cache per device for build_sdpa_packed_attention_mask to avoid repeated D2H sync across layers
_SDPA_MASK_CACHE: dict = {}

# Cache per device for build_xformers_block_causal_mask to avoid repeated D2H sync across layers
_XFORMERS_BLOCK_MASK_CACHE: dict = {}
Comment on lines +39 to +46

@gemini-code-assist gemini-code-assist Bot Mar 12, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

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

medium

The type hints for the new cache dictionaries are dict. For better code clarity and type safety, consider using more specific types from the typing module. You'll need to add Dict and Any to your imports from typing.

Using Dict[torch.device, Dict[str, Any]] would be more descriptive. For even better documentation of the cache entry structure, you could consider using TypedDict.

Suggested change
# Cache per device for get_packed_info_from_kwargs to avoid repeated D2H sync across layers
_PACKED_INFO_CACHE: dict = {}
# Cache per device for build_sdpa_packed_attention_mask to avoid repeated D2H sync across layers
_SDPA_MASK_CACHE: dict = {}
# Cache per device for build_xformers_block_causal_mask to avoid repeated D2H sync across layers
_XFORMERS_BLOCK_MASK_CACHE: dict = {}
# Cache per device for get_packed_info_from_kwargs to avoid repeated D2H sync across layers
_PACKED_INFO_CACHE: "Dict[torch.device, Dict[str, Any]]" = {}
# Cache per device for build_sdpa_packed_attention_mask to avoid repeated D2H sync across layers
_SDPA_MASK_CACHE: "Dict[torch.device, Dict[str, Any]]" = {}
# Cache per device for build_xformers_block_causal_mask to avoid repeated D2H sync across layers
_XFORMERS_BLOCK_MASK_CACHE: "Dict[torch.device, Dict[str, Any]]" = {}



def _window_cache_key(sliding_window: Optional[int]) -> int:
if sliding_window is None or sliding_window <= 0:
Expand Down Expand Up @@ -224,13 +233,18 @@ def get_packed_info_from_kwargs(
if seq_lengths is None:
return None

entry = _PACKED_INFO_CACHE.get(device)
if entry is not None and entry["seq_lengths"] is seq_lengths:
return entry["result"]

lengths = seq_lengths.to(device = device, dtype = torch.int32, non_blocking = True)
cu_seqlens = torch.empty(lengths.numel() + 1, dtype = torch.int32, device = device)
cu_seqlens[0] = 0
cu_seqlens = torch.zeros(lengths.numel() + 1, dtype = torch.int32, device = device)
torch.cumsum(lengths, dim = 0, dtype = torch.int32, out = cu_seqlens[1:])

max_seqlen = int(lengths.max().item())
return lengths, cu_seqlens, max_seqlen
result = (lengths, cu_seqlens, max_seqlen)
_PACKED_INFO_CACHE[device] = {"seq_lengths": seq_lengths, "result": result}
return result


def build_xformers_block_causal_mask(
Expand All @@ -243,11 +257,28 @@ def build_xformers_block_causal_mask(
return None
if seq_info is not None:
seq_lengths, _, _ = seq_info
# Cache the mask to avoid repeated D2H sync across layers
device = seq_lengths.device
params = (sliding_window,)
entry = _XFORMERS_BLOCK_MASK_CACHE.get(device)
if (
entry is not None
and entry["seq_lengths"] is seq_lengths
and entry["params"] == params
):
return entry["mask"]

lengths_tensor = seq_lengths.to("cpu", torch.int32)
if lengths_tensor.numel() == 0:
return None
lengths = tuple(int(x) for x in lengths_tensor.tolist())
mask = _get_cached_block_mask(lengths, sliding_window)

_XFORMERS_BLOCK_MASK_CACHE[device] = {
"seq_lengths": seq_lengths,
"params": params,
"mask": mask,
}
else:
mask = base_mask

Expand All @@ -269,6 +300,16 @@ def build_sdpa_packed_attention_mask(
sliding_window: Optional[int] = None,
) -> torch.Tensor:
seq_lengths, _, _ = seq_info

params = (dtype, sliding_window)
entry = _SDPA_MASK_CACHE.get(device)
if (
entry is not None
and entry["seq_lengths"] is seq_lengths
and entry["params"] == params
):
return entry["mask"]

total_tokens = int(seq_lengths.sum().item())
mask = torch.full(
(total_tokens, total_tokens),
Expand Down Expand Up @@ -297,7 +338,14 @@ def build_sdpa_packed_attention_mask(
block = block.masked_fill(window_mask, float("-inf"))
mask[offset : offset + length, offset : offset + length] = block
offset += length
return mask.unsqueeze(0).unsqueeze(0)

result = mask.unsqueeze(0).unsqueeze(0)
_SDPA_MASK_CACHE[device] = {
"seq_lengths": seq_lengths,
"params": params,
"mask": result,
}
return result


def _normalize_packed_lengths(
Expand Down Expand Up @@ -341,6 +389,13 @@ def mask_packed_sequence_boundaries(
return True


def clear_packed_caches():
"""Release cached masks/metadata to free device memory."""
_PACKED_INFO_CACHE.clear()
_SDPA_MASK_CACHE.clear()
_XFORMERS_BLOCK_MASK_CACHE.clear()


__all__ = [
"configure_sample_packing",
"configure_padding_free",
Expand All @@ -351,4 +406,5 @@ def mask_packed_sequence_boundaries(
"build_xformers_block_causal_mask",
"build_sdpa_packed_attention_mask",
"mask_packed_sequence_boundaries",
"clear_packed_caches",
]
Loading