TL;DR
nn.MultiheadAttention은 추론(inference) 모드이면서 기울기 계산(gradient)이 필요 없고 마스크를 전달할 경우, need_weights=False 옵션을 사용하더라도 SDPA (Scaled Dot-Product Attention)를 사용하지 않고 헤드 수 × P × P 크기의 어텐션 행렬을 통째로 생성합니다. Xiaomi-Lens-0에서는 flash-attn이 없는 경우의 대체 경로가 바로 이 방식이었는데, 3 MP 이미지에서 마스크 생성 단계만 추가로 VRAM을 18.02 GB 소모했고, 8 MP에서는 한 번에 30.75 GiB를 확보해야 하여 OOM(Out of Memory)이 발생했습니다.
- 기존에 구현되어 있던 SDPA 경로를 대체용으로 사용하는 단 하나의 PR 파일 수정만으로 VRAM 사용량이 2.28 GB (약 87% 감소)로 줄어들었고, 다음 날 병합되었습니다. 이 현상은 학습 시나 테스트 환경에서는 재현되지 않는 함정(trap)이므로, 옵션 의존성을 제외한 추론 경로는 별도로 벤치마킹해야 합니다.
배경: 왜 파고들었는가
Xiaomi-Lens-0은 Qwen3.5-4B 기반의 멀티모달 모델로, 문장으로 답변하는 동시에 <|mask|> 토큰을 픽셀 단위의 마스크로 변환합니다. 마스크를 만드는 디코더인 PatchSelfAttention은 flash-attn의 FA2, PyTorch의 SDPA, 그리고 nn.MultiheadAttention 세 가지 구현체를 동일한 가중치(weight)로 전환할 수 있도록 설계되어 있습니다.
기본 설정은 flash였으나, flash-attn이 없거나 GPU가 sm80 미만이거나 CPU인 경우 MultiheadAttention으로 폴백(fallback)되는 구조였습니다. Windows 환경에서는 flash-attn 빌드가 번거롭기 때문에, 작성자 본인의 환경은 정확히 그 대체 경로를 거치고 있었습니다. 큰 이미지를 넣었더니 메모리 부족이 발생했다는 것이 이 글의 시작점입니다.
작성자가 가장 의아하게 여긴 부분은 해당 코드에 need_weights=False가 제대로 붙어 있었다는 점이었습니다. 문서 내용을 읽어보면, 이것만으로 SDPA를 사용해야 할 것 같았습니다. 이렇게 정확하게 작성된 코드가 18 GB를 소모한다면, 이야기는 Xiaomi-Lens-0에서 끝나지 않습니다.
환경
| 항목 | GPU 측정 (PR) | 메커니즘 재현 |
|---|
| 하드웨어 | RTX 5090 32 GB | CPU 2 코어 / 메모리 7 GB |
| ... | | |
무엇이 일어나고 있었는가
PR에 올린 수치로, 모델 로딩 후 상태부터 마스크 생성 단계 (infer.py의 run_forward)가 추가적으로 사용한 VRAM의 피크 값입니다. 전체 애플리케이션 사용량이 아닙니다.
| 입력 | 패치 수 P | MultiheadAttention | SDPA |
|---|
| 0.5 MP | 1,944 | 0.86 GB | 0.46 GB |
| ... | | | |
이미지가 6배가 되자 P는 21배 증가했습니다. P가 이미지를 16×16 화소로 나눈 구역 수이므로, 이차적으로 늘어나는 것은 한눈에 알 수 있습니다. 16 헤드 × P × P를 bf16으로 저장하면 3 MP에서 약 4.32 GB가 필요하며, 18.02 GB는 그 약 4장 분량에 해당합니다. 8 MP (3840×2560)에서는 P가 약 3.2만까지 늘어나 한 장만으로도 30.75 GiB의 확보를 요구하여 OOM이 발생했습니다.
측정 구조는 다음과 같습니다.
def mask_pass_extra_gb(model, processor, mask_id, image, text, device):
inputs, _ = build_inputs(processor, image, text, device)
ids, _, n_masks = generate_completion(...)
범인 찾기: need_weights=False는 무엇을 약속했었나
문서의 약속
PyTorch 2.14의 nn.MultiheadAttention에는 다음과 같이 쓰여 있습니다.
Set need_weights=False to use the optimized scaled_dot_product_attention
(번역: 최적화된 scaled_dot_product_attention을 사용하려면 need_weights=False를 설정하시오)
같은 페이지에는 추론 전용의 fast path에 진입하는 조건들도 나열되어 있습니다. 자기 주의(self-attention)여야 하는 것, batch_first의 3차원 입력이어야 하는 것, autograd 비활성화(inference_mode 또는 no_grad)가 되어 있어야 하는 것, .eval()이 호출되었어야 하는 것 등입니다. (infer.py)
model.eval() 이후에 torch.inference_mode()로 추론하기 때문에, 모두 충족한다.
소스의 현실
fast path의 CUDA 구현인 native_multi_head_attention_cuda를 읽어보면, SDPA(Scaled Dot Product Attention)로 진입하는 조건의 시작이 이것이었다.
if (!mask.has_value() && no_seq_len_1_nested &&
마스크가 있으면 SDPA로는 가지 않고, bmm_nt를 사용해 스코어 행렬을 만든 후 masked_softmax를 적용한다. 바로 앞 주석에는 형태가 [B, num_head, T, T]임을 명시하고 있다.
게다가 Xiaomi-Lens-0의 디코더는 패딩용 key_padding_mask를 항상 전달하고 있었다. 즉, need_weights=False라는 약속은 fast path에 진입하는 시점에서 지켜지지 않는다. 필자는 문서의 그 문장은 fast path 외부(Python 측의 F.multi_head_attention_forward)이거나 CUDA에서 마스크를 전달하지 않는 경우에만 성립한다고 해석하는 것이 정확하다고 생각한다.
프로파일러로 검증하기
소스만 읽기에는 불안했기 때문에, CPU에서 동일한 호출을 4가지 조건으로 실행하여 호출된 aten 연산을 포착했다.
import torch
import torch.nn as nn
from torch.profiler import ProfilerActivity, profile
...
eval none ['aten::_native_multi_head_attention', 'aten::_softmax', 'aten::bmm']
eval mask ['aten::_masked_softmax', 'aten::_native_multi_head_attention', 'aten::bmm']
train none ['aten::_scaled_dot_product_flash_attention_for_cpu', 'aten::scaled_dot_product_attention']
...
eval・마스크가 있는 _masked_softmax에는 [1, 16, 512, 512]가 전달되었다. CPU의 fast path는 마스크가 없어도 행렬을 만든다는 점에서 CUDA와 동작이 다르지만, 마스크가 있는 추론에서 행렬을 만드는 점은 동일하다.
여기서 가장 의외였던 것은 train의 두 줄이다. 같은 레이어・같은 인자임에도, 학습 모드일 때는 SDPA를 거친다.
CPU에서 양을 측정하기
경로가 파악되었다고 해도, 그것이 메모리 주범인지는 양을 측정하지 않고는 말할 수 없다. forward를 한 번 실행할 때마다 별도의 프로세스를 생성하고, /proc/self/clear_refs로 피크 상주 메모리(VmHWM)를 리셋한 후 증가분을 측정했다. SDPA 측은 Xiaomi-Lens-0의 _sdpa_self_attention과 동일한 구성 방식(MHA의 투영 가중치를 재사용하고, 패딩을 -inf 덧셈 마스크로 처리)을 사용한다.
peak_mem.py (전체)
import sys
import torch
import torch.nn as nn
...
| P | 16 × P² × 4 B | MultiheadAttention | SDPA |
|---|
| 1,024 | 64 MiB | 172 MiB | 44 MiB |
| ... | | | |
두 번 실행해도 같은 값이 나왔다. MultiheadAttention은 이론값의 약 2배(스코어 행렬과 softmax 출력)에 수렴해 갔고, SDPA는 P 1,024 부근에서 20 MiB의 직선을 그렸다. GPU의 실측 결과와 동일한 형태가 CPU에서도 나왔다.
PR의 4개 지점에 이차식과 일차식을 대입하면, 32 GB급 GPU에서 MultiheadAttention 경로는 약 3 MP까지 실용적이며, SDPA 경로는 최대 8.4 MP까지 추가로 약 6 GB가 필요할 것으로 추정된다 (4 MP 이상은 외삽).
fit_scaling.py
import numpy as np
P = np.array([1944, 3952, 7776, 11616], dtype=float)
mha_gb = np.array([0.86, 2.58, 8.53, 18.02])
...
수정 방법: 새로 작성하지 않기
선택지는 3가지가 있었다.
- MultiheadAttention의 호출 방식을 조정하여 fast path를 우회한다. 기울기(gradient) 유무에 의존하는 비공식적인 방법으로, PyTorch 버전별로 동작이 달라질 수 있어 기각함.
- 마스크를 전달하지 않는다. 패딩을 포함한 입력에서는 의미가 바뀌므로 기각함.
- 이미 존재하던 SDPA 경로를 대체재로 사용한다. 동일한 가중치(weight)를 사용하며, 환경 변수
XIAOMI_LENS_DECODER_ATTN=sdpa
으로 선택할 수 있는 상태였다.
3번을 선택했다. 차이점은 실질적으로 이 경로 대체뿐이다 (docstring 생략, 시그니처는 1줄로 압축함).
def _self_attention(self, normed, key_padding_mask):
- if _DECODER_ATTN == "sdpa":
- return self._sdpa_self_attention(normed, key_padding_mask)
...
XIAOMI_LENS_DECODER_ATTN=mha
로 기존 동작을 명시적으로 선택할 수 있도록 했고, 경고 메시지도 "falling back to SDPA"로 수정했다. 환경 변수는 그것을 아는 사람만이 도움받을 수 있다. 기본 대체 경로를 변경하면 모르는 사람에게도 도달한다.
출력은 달라졌는가
PR에서는 6가지 케이스(고양이, 커피잔, 23개의 동전, 로켓, 깃발, 추론이 포함된 요청)에서 토큰 ID, 박스, 전체 마스크 로그잇(logit)을 비교했다.
| 비교 | 결과 |
|---|
| 새로운 기본값 vs 기존 코드에서 sdpa 명시 | 6/6 비트 단위 일치 |
| ... | |
| 상위 2줄은 배선 확인, 하단 2줄이 경로 차이다. 경계 부분에 약간의 차이가 발생하는 것은 반올림(rounding) 때문이며, 6가지 케이스는 포괄적인 품질 평가가 아니다. CPU에서도 동일한 가중치로 비교했을 때, 최대 절대차는 fp32에서 5.22e-08, bf16에서 4.88e-04였으며 NaN은 없었다. | |
compare_out.py
import torch
import torch.nn as nn
from peak_mem import E, H, sdpa_like_repo
...
직접 MHA에서 SDPA로 이식할 경우, 참/거짓 마스크의 방향에 주의해야 한다. key_padding_mask는 True가 무시(ignore)를 의미하고, SDPA의 attn_mask는 True가 참여(participate)를 의미하여 방향이 반대이다. 그대로 전달하면 결과가 손상된다 (아래 스크립트에서 최대 차 3.46). Xiaomi-Lens-0은 가산 마스크로 변환했기 때문에 문제없었다.
import torch
import torch.nn.functional as F
torch.manual_seed(0)
...
고찰: 대체 경로는 왜 방치되는가
이번 건을 통해 필자가 배운 것은, 옵션에 의존하는 대체 경로는 이중으로 눈에 띄지 않는다는 점이다.
첫째는 개발자의 손안에는 본 경로(flash-attn)가 있기 때문에, 대체 경로를 거치지 않는다. 둘째는 대체 경로를 거친다 하더라도, 학습이나 기울기가 있는 테스트에서는 fast path로 진입하지 않기 때문에 문제가 발생하지 않는다. 대체 경로가 2차 메모리를 소모하는 지점은 "flash-attn이 없을 때" 그리고 "실제 서비스와 동등한 추론(inference)"의 교차점뿐이며, 그곳은 작성자가 가장 잘 건드리지 않는 곳이다.
같은 교차점은 다른 모델에도 존재한다. 표준 nn.TransformerEncoderLayer 역시, src_key_padding_mask를 전달하여 추론을 수행했을 때 비슷한 형태의 행렬을 만들었다. 손안의 모델을 점검하기 위해, 추론과 동일한 inference_mode로 한 번 실행하고 마지막 2차원이 같은 큰 텐서를 가져오는 함수를 작성했다.
import torch
from torch.profiler import ProfilerActivity, profile
WATCH = ("softmax", "bmm", "baddbmm", "matmul")
...
MHA [('aten::_masked_softmax', (1, 16, 1024, 1024)), ('aten::bmm', (16, 1024, 1024))]
Enc [('aten::_masked_softmax', (1, 16, 1024, 1024)), ('aten::bmm', (16, 1024, 1024))]
SDPA []
이 함수를 작성하는 과정에서 필자는 한 번 속았습니다. `.eval()`을 호출하는 것을 잊은 레이어에서 테스트하여 빈 리스트가 반환되자 '문제없다'고 결론 내렸습니다. 검사 자체를 실제 운영 조건(production condition)으로 돌리지 않으면, 이 함정은 테스트를 빠져나갈 수 있습니다.
수정 자체는 몇 줄에 불과하지만, 유지보수자가 짧은 시간 안에 '올바르다'고 판단할 근거가 없으면 병합(merge)되지 않습니다. PR 본문은 '문제 / 변경 / 측정 / 출력 비교'로 나누어 다음 3가지 사항을 명시했습니다.
- 측정한 범위: 전체 VRAM이 아닌 추가적인 피크 값임
- 우회 경로: `mha`를 명시하면 이전 동작으로 되돌릴 수 있음 - 출력 차이의 성질: 배선 일치(비트 단위)와 경로 차이의 크기(IoU)를 나누어 제시함
10월 19일에 제출했고, 20일에 병합되었습니다.
- `need_weights=False`는 추론 시 fast path에 마스크를 전달할 때 SDPA를 보장하지 않습니다. 소스 코드의 분기는 마스크 유무로 결정됩니다 - Xiaomi-Lens-0에서는 대체 경로를 SDPA로 교체하는 것만으로 3 MP 추가 VRAM이 18.02 GB에서 2.28 GB로 줄었습니다.
- 대체 경로 문제는 메인 기능이 들어간 로컬 환경과, 그래디언트(gradient)가 있는 테스트 모두에서 숨겨집니다. 실제 운영 환경과 같습니다.
`eval`과 `inference_mode`를 사용하여 옵션 의존성을 제거하고 한 번 측정하는 것이 가장 저렴한 대책입니다.
- PyTorch 2.14 torch.nn.MultiheadAttention ── fast path의 조건 및 need_weights 설명
- PyTorch 2.14 scaled_dot_product_attention ── 참/거짓 마스크의 방향, dropout, 구현 자동 선택
- pytorch/pytorch attention.cu ── fast path가 마스크가 있을 때 SDPA로 진행하지 않는 분기 (2026-10-24 시점 main)
- SeerRay-Lab/Xiaomi-Lens-0 PR
전체 표의 비유를 사용하여 기호와 식을 하나씩 추적한 입문자용 해설은 Qiita에 작성했습니다.
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
