A Small Attention Fallback Fix Cut Peak Mask-Pass VRAM by 87%

A merged fix in Xiaomi-Lens-0, and why fallback paths deserve their own benchmarks.

When an optional acceleration library is unavailable, what happens next?

In Xiaomi-Lens-0, the pixel mask decoder fell back from flash-attn to nn.MultiheadAttention. In the configuration I tested, that fallback made large images consume enough GPU memory to fail.

An efficient alternative was already implemented. The fix was to make it the fallback.

At a 3-megapixel input size, the additional peak GPU memory used by the mask pass dropped from 18.02 GB to 2.28 GB—about 87%.

The fix has been merged upstream.

The expensive path was the fallback

Xiaomi-Lens-0 can locate and segment objects in images. Its pixel mask decoder includes a PatchSelfAttention module with multiple attention implementations.

Before this change, when the flash-attn path was unavailable, the decoder used its MultiheadAttention implementation.

In the tested execution path, the model’s padding mask caused attention weights to be materialized across all 16 heads. With (P) image patches, the attention matrices scale with (P^2).

That gets expensive quickly as image resolution increases.

The decoder already had a scaled_dot_product_attention implementation, selectable through:

XIAOMI_LENS_DECODER_ATTN=sdpa

It used the same model weights and avoided that memory bottleneck in my test environment.

The opportunity was to make that existing path work automatically for users without flash-attn.

The change

The routing previously looked like this:

if configured_backend == "sdpa":
    return sdpa_attention(...)

if configured_backend == "flash":
    output = try_flash_attention(...)
    if output is not None:
        return output

return mha_attention(...)

After the change:

if configured_backend == "mha":
    return mha_attention(...)

if configured_backend == "flash":
    output = try_flash_attention(...)
    if output is not None:
        return output

return sdpa_attention(...)

These snippets simplify the actual implementation, but capture the change.

Explicit MHA selection remains available. The warning message now reports the correct fallback.

What I measured

The benchmark environment was:

• NVIDIA RTX 5090, 32 GB VRAM
• Windows 11
• PyTorch 2.14.1+cu130
• Transformers 5.4.0
• No flash-attn installed

I used scikit-image’s coffee photograph, resized to different input sizes, and the README prompt targeting “the coffee cup.”

The benchmark followed the inference stages used by infer.py. The table measures peak GPU memory added by run_forward, the mask pass, above the loaded-model baseline.

It does not measure total application VRAM.

Input size	Patches	Previous MHA fallback	SDPA fallback
0.5 MP	1,944	0.86 GB	0.46 GB
1 MP	3,952	2.58 GB	0.83 GB
2 MP	7,776	8.53 GB	1.55 GB
3 MP	11,616	18.02 GB	2.28 GB

At 3 MP, that is 15.74 GB less additional peak memory for this stage.

In a separate large-image test through a ComfyUI integration, the MHA path ran out of memory while attempting a single 30.75 GiB allocation. The SDPA path completed, with a whole-request peak of 15.05 GB and a runtime of approximately six seconds.

That separate test used a byte-identical copy of PatchSelfAttention, rather than running infer.py itself. Its whole-request measurement should not be directly compared with the mask-pass numbers above.

Did the outputs change?

I compared six image-and-prompt cases, including object segmentation and a reasoning request, checking token IDs, boxes, and every mask logit.

Two checks verified that the routing change behaved as intended:

• The new default matched the previous explicit SDPA path bit for bit in all six cases.
• Explicitly selecting MHA still matched the previous default bit for bit in all six cases.

Comparing SDPA against MHA:

• Generated token IDs and boxes were unchanged.
• Masks thresholded at 0.5 had an IoU of 0.9976–0.9998.

There were small differences at mask boundaries. This was a limited six-case comparison, not a comprehensive quality evaluation, but it made the numerical differences visible alongside the memory improvement.

This is not a universal MHA-to-SDPA rule

The result depends on the execution path.

PyTorch’s MultiheadAttention can itself use optimized SDPA, and its documentation recommends need_weights=False to enable that path. SDPA also selects implementations based on the inputs and available backends; calling the API does not guarantee the same memory savings on every device.

Migrating implementations requires care: boolean mask semantics differ between MHA’s key padding mask and SDPA’s attention mask, and SDPA’s dropout must be configured explicitly.

The relevant documentation is here:

• PyTorch MultiheadAttention
• PyTorch scaled_dot_product_attention

The finding in this PR is narrower: this decoder already had an effective SDPA path, and selecting it as the fallback substantially reduced memory use in the measured environment.

Why a small default change matters

An environment variable helps users who know it exists and understand when to set it.

A better fallback reaches users automatically.

This fix changed one file. The work behind it was identifying the expensive execution path, measuring its scaling, checking output differences, and giving maintainers evidence they could review.

For someone whose image previously ran out of memory, that small routing change can determine whether the model is usable at all.

My next step is to investigate a few other models and see whether similar fallback problems recur. If the diagnosis and verification work proves reusable, it could become a shared toolkit.

For now, the practical lesson is simple: benchmark the path your model takes when an optional dependency is missing. That path may deserve more attention than it gets.