diff --git a/unsloth/utils/packing.py b/unsloth/utils/packing.py index 63a57c04da1..bd6631d4618 100644 --- a/unsloth/utils/packing.py +++ b/unsloth/utils/packing.py @@ -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 = {} + def _window_cache_key(sliding_window: Optional[int]) -> int: if sliding_window is None or sliding_window <= 0: @@ -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( @@ -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 @@ -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), @@ -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( @@ -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", @@ -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", ]