본문으로 건너뛰기

[flashinfer] Blackwell NVFP4 양자화 최적화: TMA OOB Zero-fill을 이용한 메모리 복사 오버헤드 제거

PR 링크: flashinfer-ai/flashinfer#4210 상태: Merged | 변경: +79 / -14

들어가며

최근 대규모 언어 모델(LLM)의 추론 성능을 극대화하기 위해 NVIDIA Blackwell 아키텍처에서 도입된 NVFP4(4-bit Floating Point) 양자화 기술이 주목받고 있습니다. NVFP4는 높은 압축률과 함께 하드웨어 가속을 통해 연산 처리량을 비약적으로 높일 수 있는 기술입니다.

하지만 하드웨어의 성능을 100% 끌어내기 위해서는 메모리 레이아웃과 정렬(Alignment) 조건이 까다롭습니다. 특히 FlashInfer와 같은 고성능 라이브러리에서 NVFP4 양자화를 구현할 때, 입력 행렬의 행 크기(M)가 특정 블록 크기(예: 128)에 정렬되지 않으면 이를 맞추기 위한 추가적인 패딩(Padding) 작업이 필요했습니다.

이번 글에서는 FlashInfer PR(#3905 관련)에서 수행된 "NVFP4 TMA 입력 패딩 복사 제거" 최적화를 분석합니다. 이 최적화는 하드웨어의 특성을 활용해 불필요한 메모리 할당과 복사를 제거함으로써, 비정렬 데이터에 대한 성능을 최대 72%까지 끌어올렸습니다.

문제 상황: 128-aligned 'M' Copy Cliff

기존 구현에서는 NVFP4 양자화 시 M이 128의 배수가 아닐 경우, 하드웨어의 TMA(Tensor Memory Accelerator)가 요구하는 레이아웃을 맞추기 위해 다음과 같은 비효율적인 과정을 거쳤습니다.

  1. padded_M 크기의 임시 텐서를 할당 (0으로 초기화).
  2. 원본 입력 데이터를 임시 텐서로 복사.
  3. 패딩된 텐서를 TMA를 통해 커널로 전달.

이 과정은 M이 정렬되어 있을 때는 발생하지 않지만, 단 1개라도 행이 어긋나면 전체 텐서를 복사해야 하는 "Performance Cliff(성능 절벽)" 현상을 야기했습니다. 특히 B200과 같은 고성능 GPU에서 메모리 대역폭을 낭비하는 주범이 되었습니다.

코드 분석: 불필요한 복사 로직의 제거

1. Python 호스트 측의 임시 할당 제거

가장 먼저 눈에 띄는 변화는 nvfp4_quantize_cute_dsl 함수에서 패딩된 텐서를 생성하던 로직이 완전히 사라진 것입니다.

Before:

input_padded = input
if padded_m > m:
    input_padded = torch.zeros(
        padded_m, k, dtype=input.dtype, device=input.device
    )
    input_padded[:m, :] = input

kernel_fn(
    input_padded,
    # ...
)

After:

# input_padded 생성 로직 제거
kernel_fn(
    input,
    # ...
)

기존에는 padded_m > m인 경우 torch.zeros를 호출하고 슬라이싱 복사를 수행했습니다. 이는 GPU 메모리 할당 오버헤드와 데이터 이동 오버헤드를 동시에 발생시켰습니다. 변경 후에는 원본 input을 그대로 커널에 전달합니다.

2. TMA Descriptor의 물리적 크기 반영

TMA가 데이터를 읽어올 때 참조하는 레이아웃 정보도 수정되었습니다. 기존에는 패딩된 크기(padded_M)를 기준으로 텐서 맵을 생성했지만, 이제는 실제 물리적 크기(M)를 사용합니다.

Before (flashinfer/quantization/kernels/nvfp4_quantize.py):

gInput = cute.make_tensor(
    mInput.iterator,
    cute.make_layout(
        (padded_M, self.K // _TMA_COL_TILE, _TMA_COL_TILE),
        stride=(self.K, _TMA_COL_TILE, 1),
    ),
)

After:

gInput = cute.make_tensor(
    mInput.iterator,
    cute.make_layout(
        (M, self.K // _TMA_COL_TILE, _TMA_COL_TILE),
        stride=(self.K, _TMA_COL_TILE, 1),
    ),
)

여기서 핵심은 TMA의 Out-of-bounds(OOB) Zero Fill 기능입니다. TMA는 텐서 맵에 정의된 경계를 벗어나는 영역을 읽으려고 할 때, 하드웨어 수준에서 자동으로 0을 채워주는 기능을 제공합니다. 따라서 굳이 메모리에 0을 써둘 필요 없이, TMA 설정만 실제 크기로 잡아주면 패딩 영역은 하드웨어가 알아서 0으로 처리해 줍니다.

왜 이게 좋은 최적화인가?

1. 극적인 성능 향상 (Performance Data)

PR에서 제공된 B200 벤치마크 결과는 놀랍습니다. M=32283 (비정렬) 케이스를 보면:

  • 기존 (TMA + input copy): 212.446 µs
  • 최적화 후 (This PR): 59.248 µs
  • 개선율: 약 72.1% 감소

정렬된 케이스(M=32384)와 거의 동일한 성능을 보여줌으로써, 정렬 여부에 따른 성능 편차를 완전히 제거했습니다. 또한 기존의 Vectorized 경로보다도 4~5% 더 빠른 성능을 보여줍니다.

2. 하드웨어-소프트웨어 공동 설계 (Co-design)

소프트웨어에서 루프를 돌며 0을 채우거나 복사하는 대신, Blackwell 아키텍처의 TMA 유닛이 가진 특성(OOB Zero-fill)을 정확히 파악하고 이를 활용했습니다. 이는 전형적인 "하드웨어 가속기를 가장 똑똑하게 쓰는 방법"입니다.

3. 메모리 효율성

대규모 배치 사이즈에서 M이 클 경우, 패딩을 위한 임시 텐서 할당은 VRAM 부족(OOM)을 유발하거나 메모리 단편화를 초래할 수 있습니다. 이번 변경으로 추가 메모리 점유 없이 인플레이스(In-place)에 가까운 처리가 가능해졌습니다.

결론

이번 FlashInfer의 PR은 고성능 컴퓨팅에서 "추상화된 데이터 구조와 실제 하드웨어 동작 사이의 간극을 어떻게 메울 것인가"에 대한 훌륭한 답을 제시합니다.

단순히 코드를 깔끔하게 만드는 것을 넘어, 하드웨어의 TMA Descriptor가 OOB를 처리하는 방식을 이해함으로써 대규모 데이터 복사라는 병목 지점을 완전히 제거했습니다. Blackwell과 같은 최신 GPU를 타겟팅하는 개발자라면, CuTe/CUTLASS를 사용할 때 이러한 하드웨어 특화 기능을 적극적으로 탐색해 볼 가치가 있습니다.

참고 자료

참고 자료

⚠️ 알림: 이 분석은 AI가 실제 코드 diff를 기반으로 작성했습니다.

댓글

관련 포스트

PR Analysis 의 다른글