Cache packed sequence metadata to reduce D2H syncs across layers

[ichbinhandsome]

Added per-forward-pass caching to eliminate redundant D2H copies and cudaStreamSynchronize calls across layers.

When packing (padding-free) is enabled, three functions are called on every layer of the model during the forward pass:

• get_packed_info_from_kwargs: calls lengths.max().item() — triggers D2H copy + sync
• build_sdpa_packed_attention_mask (SDPA backend): calls seq_lengths.sum().item() and seq_lengths.tolist() — triggers 2 D2H copies + syncs
• build_xformers_block_causal_mask (XFormers backend): calls seq_lengths.to("cpu") — triggers D2H copy + sync

For a model with N layers, this results in N unnecessary D2H synchronizations per function, even though the packed sequence metadata (seq_lengths) is identical across all layers within the same forward pass.

Solution

Cache the output of each function using Python object identity (is) comparison on the seq_lengths tensor. Since the same seq_lengths tensor object is passed to all layers within a single forward pass, subsequent layers hit the cache and skip the D2H operations entirely. A new batch produces a new seq_lengths tensor object, which naturally invalidates the cache.

This reduces D2H synchronizations per forward pass:

Function	Before	After
get_packed_info_from_kwargs	N	1
build_sdpa_packed_attention_mask	2N	2
build_xformers_block_causal_mask	N	1

Nsys profiling traces:

• Without this PR:

• With this PR:

Performance

With the caching strategy, CudaStreamSync only appears in the first layer and in the following layers it disappears. We achieve around 43.3% speedup for forward, 5.8% speedup for backward, 14.3% speedup for each batch for Qwen3 14B QLoRA SFT.

[gemini-code-assist]

Summary of Changes

Hello, I'm Gemini Code Assist¹! I'm currently reviewing this pull request and will post my feedback shortly. In the meantime, here's a summary to help you and other reviewers quickly get up to speed!

This pull request introduces a caching mechanism to optimize the performance of models utilizing packed sequences. By intelligently caching metadata related to sequence lengths and attention masks, it drastically reduces the number of costly Device-to-Host synchronizations that previously occurred in every layer of a multi-layer model. This optimization leads to substantial speedups in both forward and backward passes, improving overall training and inference efficiency.

Highlights

• D2H Synchronization Reduction: Implemented per-forward-pass caching for get_packed_info_from_kwargs, build_sdpa_packed_attention_mask, and build_xformers_block_causal_mask to eliminate redundant Device-to-Host (D2H) copies and cudaStreamSynchronize calls across model layers.
• Caching Strategy: Caches are invalidated naturally by new seq_lengths tensor objects for each new batch, ensuring correctness while maximizing cache hits for subsequent layers within the same forward pass by using Python object identity (is) comparison.
• Performance Improvement: Achieved significant speedups, including approximately 43.3% for forward pass, 5.8% for backward pass, and 14.3% per batch for Qwen3 14B QLoRA SFT.

Changelog

• unsloth/utils/packing.py
  • Introduced global cache dictionaries for packed info, SDPA masks, and XFormers block masks.
  • Modified get_packed_info_from_kwargs to check and utilize the _PACKED_INFO_CACHE, and to store results.
  • Updated cu_seqlens initialization from torch.empty to torch.zeros in get_packed_info_from_kwargs.
  • Modified build_xformers_block_causal_mask to incorporate caching logic using _XFORMERS_BLOCK_MASK_CACHE.
  • Modified build_sdpa_packed_attention_mask to integrate caching logic using _SDPA_MASK_CACHE.

Activity

• Pull request created by ichbinhandsome.
• Initial code changes implemented to introduce caching for packed sequence metadata.
• Performance benchmarks provided in the PR description demonstrating significant speedups.

Using Gemini Code Assist

The full guide for Gemini Code Assist can be found on our documentation page, here are some quick tips.

Invoking Gemini

You can request assistance from Gemini at any point by creating a comment using either /gemini <command> or @gemini-code-assist <command>. Below is a summary of the supported commands on the current page.

Feature	Command	Description
Code Review	/gemini review	Performs a code review for the current pull request in its current state.
Pull Request Summary	/gemini summary	Provides a summary of the current pull request in its current state.
Comment	@gemini-code-assist	Responds in comments when explicitly tagged, both in pull request comments and review comments.
Help	/gemini help	Displays a list of available commands.

Customization

To customize Gemini Code Assist for GitHub experience, repository maintainers can create a configuration file and/or provide a custom code review style guide (such as PEP-8 for Python) by creating and adding files to a .gemini/ folder in the base of the repository. Detailed instructions can be found here.

Limitations & Feedback

Gemini Code Assist may make mistakes. Please leave feedback on any instances where its feedback is incorrect or counter productive. You can react with 👍 and 👎 on @gemini-code-assist comments. If you're interested in giving your feedback about your experience with Gemini Code Assist for Github and other Google products, sign up here.

You can also get AI-powered code generation, chat, as well as code reviews directly in the IDE at no cost with the Gemini Code Assist IDE Extension.

Footnotes

1. Review the Privacy Notices, Generative AI Prohibited Use Policy, Terms of Service, and learn how to configure Gemini Code Assist in GitHub here. Gemini can make mistakes, so double check it and use code with caution. ↩

[chatgpt-codex-connector]

💡 Codex Review

Here are some automated review suggestions for this pull request.

Reviewed commit: ab927312f7

ℹ️ About Codex in GitHub

Codex has been enabled to automatically review pull requests in this repo. Reviews are triggered when you

• Open a pull request for review
• Mark a draft as ready
• Comment "@codex review".

If Codex has suggestions, it will comment; otherwise it will react with 👍.

When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".

[chatgpt-codex-connector]

Include device in packed-info cache key

The new cache in get_packed_info_from_kwargs returns a previously computed (lengths, cu_seqlens, max_seqlen) based only on packed_seq_lengths object identity, but those tensors are device-specific. When a forward spans layers on different devices (for example with a device_map/pipeline-parallel setup), later layers can receive cached cu_seqlens from the wrong device, which causes runtime failures in varlen attention paths that require all inputs on the same device. The cache lookup should also key on device (or store per-device results) before reusing the tuple.

Useful? React with 👍 / 👎.

[gemini-code-assist]

Code Review

This pull request introduces a caching mechanism to optimize performance by reducing D2H synchronizations. However, the current implementation has critical thread-safety issues due to the use of global state without proper synchronization, which can lead to race conditions and data corruption in concurrent environments. Specifically, cache updates are not atomic, and one caching function ignores the target device, potentially causing crashes in multi-GPU or pipeline-parallel configurations. It is crucial to address these issues by ensuring thread safety and including all relevant parameters in the cache keys, possibly by making cache updates atomic through storing entries as tuples.

[gemini-code-assist]

The use of global dictionaries (_PACKED_INFO_CACHE, _SDPA_MASK_CACHE, _XFORMERS_BLOCK_MASK_CACHE) for caching metadata across forward passes is not thread-safe. In concurrent environments, multiple threads will overwrite these global variables without synchronization, leading to race conditions, silent data corruption, or incorrect model behavior. Specifically, updating dictionary keys in separate statements is not atomic. Consider using thread-local storage (threading.local()) or appropriate locking mechanisms, or storing the entire cache entry as a single tuple to ensure atomicity and thread safety.

Suggested change

	_PACKED_INFO_CACHE: dict = {"seq_lengths": None, "result": None}
	# Cache for build_sdpa_packed_attention_mask to avoid repeated D2H sync across layers
	_SDPA_MASK_CACHE: dict = {"seq_lengths": None, "params": None, "mask": None}
	# Cache for build_xformers_block_causal_mask to avoid repeated D2H sync across layers
	_XFORMERS_BLOCK_MASK_CACHE: dict = {"seq_lengths": None, "params": None, "mask": None}
	_PACKED_INFO_CACHE: dict = {"entry": (None, None)}
	# Cache for build_sdpa_packed_attention_mask to avoid repeated D2H sync across layers
	_SDPA_MASK_CACHE: dict = {"entry": (None, None, None)}
	# Cache for build_xformers_block_causal_mask to avoid repeated D2H sync across layers
	_XFORMERS_BLOCK_MASK_CACHE: dict = {"entry": (None, None, None)}

[gemini-code-assist]

To ensure thread-safety as discussed for the cache definitions, please update this function to use the atomic cache update pattern. The cache check and update should be modified to work with a single tuple entry.

Suggested change

	params = (sliding_window,)
	if (
	_XFORMERS_BLOCK_MASK_CACHE["seq_lengths"] is seq_lengths
	and _XFORMERS_BLOCK_MASK_CACHE["params"] == params
	):
	return _XFORMERS_BLOCK_MASK_CACHE["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["seq_lengths"] = seq_lengths
	_XFORMERS_BLOCK_MASK_CACHE["params"] = params
	_XFORMERS_BLOCK_MASK_CACHE["mask"] = mask
	params = (sliding_window,)
	cached_seq_lengths, cached_params, cached_mask = _XFORMERS_BLOCK_MASK_CACHE["entry"]
	if (
	cached_seq_lengths is seq_lengths
	and cached_params == params
	):
	return cached_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["entry"] = (seq_lengths, params, mask)

[gemini-code-assist]

The cache check in get_packed_info_from_kwargs currently ignores the device argument, which can lead to returning cached tensors on the wrong device and cause a PyTorch RuntimeError in multi-GPU or pipeline-parallel setups. The cache key should include the device to ensure correctness across different execution contexts. Additionally, consider the thread-safety implications of updating cache entries, as non-atomic operations can lead to race conditions.

Suggested change

	if _PACKED_INFO_CACHE["seq_lengths"] is seq_lengths:
	return _PACKED_INFO_CACHE["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["seq_lengths"] = seq_lengths
	_PACKED_INFO_CACHE["result"] = result
	if _PACKED_INFO_CACHE.get("seq_lengths") is seq_lengths and _PACKED_INFO_CACHE.get("device") == device:
	return _PACKED_INFO_CACHE["result"]
	lengths = seq_lengths.to(device = device, dtype = torch.int32, non_blocking = True)
	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())
	result = (lengths, cu_seqlens, max_seqlen)
	_PACKED_INFO_CACHE["seq_lengths"] = seq_lengths
	_PACKED_INFO_CACHE["device"] = device
	_PACKED_INFO_CACHE["result"] = result

[chatgpt-codex-connector]

💡 Codex Review

Here are some automated review suggestions for this pull request.

Reviewed commit: 2711dbc81a

ℹ️ About Codex in GitHub

Codex has been enabled to automatically review pull requests in this repo. Reviews are triggered when you

• Open a pull request for review
• Mark a draft as ready
• Comment "@codex review".

If Codex has suggestions, it will comment; otherwise it will react with 👍.

When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".

[chatgpt-codex-connector]

Avoid pinning full SDPA masks in global cache

This stores the full packed SDPA mask tensor in a module-global dict, so the latest (1,1,T,T) mask for a device stays strongly referenced until another packed SDPA call replaces it. With long packed sequences this can be a very large allocation, and if execution later switches to non-packed/xformers/flash paths (so this function is not called again) the memory remains live and can cause avoidable OOMs in later steps. A per-forward or explicitly evictable cache would keep the sync optimization without holding large masks indefinitely.

Useful? React with 👍 / 👎.