Search
English

Using a mask in PyTorch attention inference spiked GPU memory; the author reduced it by changing the fallback path(zenn.dev)

startupxo Discuss

Writing language: KoreanRead in the original language

Summary / Read source ↗

- The model Xiaomi-Lens-0, which creates masks marking image regions, switches to the MultiheadAttention path when flash-attn is not installed. When a mask was passed, that path created a matrix sized as 16 heads × the square of the patch count. - According to the author's measurements, for a 3-megapixel image the extra VRAM peak in the mask step was 18.02GB on the original path and 2.28GB on the SDPA path. The test environment was RTX 5090 32GB, Windows 11, with flash-attn not installed. - The PyTorch documentation says an optimized path is used when need_weights=False. However, when the author read the CUDA source, the path was not entered when a mask was present. In CPU profiling, the same layer in training mode used SDPA, so the issue reproduced only in eval inference. - The author reported that across a comparison of 6 samples, tokens and boxes matched, and mask IoU ranged from 0.9976 to 0.9998. This is a limited comparison, not a full quality evaluation.
Found on

Zenn ↗ / Published by source

Have you run into something like this? Tell us in a line or two what you did. You do not need to know the product.

Sign in to comment

1 comment

Editorial opinionstartupxo

O: An engineer who runs an in-house inference server could measure the GPU memory peak once with a padding mask, under eval and inference_mode, and compare it with the SDPA path. X: The figure is the extra peak added by the mask step in the author's environment, not total VRAM. First, it is worth checking that results can differ if the GPU, resolution, or model structure is different.
Writing language: Korean

Keyboard shortcuts

Choose a post with J and K or the arrow keys, then press Enter.

K ↑ / J ↓
Previous post / next post
Enter C
Open summary and comments for the selected post
O
Open the original link of the selected post
/
Search
?
Show this help
Tab
Move to the submit or comment button, then press Enter

Type normally in text fields. Tab and Enter are always available.