[flashinfer] FlashInfer SM120 NVFP4 어텐션 최적화: N64 스코어-슬롯 재사용을 통한 성능 향상

2026년 8월 25일수정: 2026년 8월 25일

PR 링크: flashinfer-ai/flashinfer#4502 상태: Merged | 변경: +1429 / -379

들어가며

최근 LLM 및 Diffusion 모델의 발전과 함께, 모델의 추론 속도 향상은 매우 중요한 과제가 되었습니다. 특히 어텐션 메커니즘은 이러한 모델들의 핵심 구성 요소로, 연산량이 많아 병목 현상의 주된 원인이 됩니다. NVIDIA GPU의 최신 아키텍처인 SM120 (Hopper) 환경에서 NVFP4 (NVIDIA Variable Precision Format 4-bit) 어텐션 연산을 더욱 가속하기 위한 PR이 올라왔습니다. 이 PR은 기존의 N128 어텐션 타일에 대한 최적화를 통해, 특히 Cosmos 추론에서 사용되는 헤드 차원 128 워크로드의 성능을 크게 향상시키는 것을 목표로 합니다.

이번 PR의 핵심은 N128 어텐션 타일 내에서의 대칭적인 N64/N64 스코어-슬롯 재사용입니다. 이를 통해 기존에는 메모리에 유지해야 했던 추가적인 스코어 프래그먼트를 줄이고, QK(Query-Key) 스코어 계산과 PV(Position-Value) MMA(Matrix Multiply-Accumulate) 연산 간의 오버랩을 늘려 전체적인 처리량을 높입니다. 또한, LSE(Log-Sum-Exp) 계산을 선택적으로 비활성화하여 불필요한 연산을 제거하고, 커널의 구조를 개선하는 등 다양한 최적화 기법이 적용되었습니다.

이 글에서는 해당 PR의 코드 변경 사항을 상세히 분석하고, 각 변경이 왜 성능 향상으로 이어지는지, 그리고 어떤 기술적 교훈을 얻을 수 있는지 살펴보겠습니다.

코드 분석

1. `benchmarks/bench_sm120_attention.py` - 벤치마킹 스크립트 수정

벤치마킹 스크립트는 새로운 기능과 설정을 반영하기 위해 수정되었습니다. 가장 눈에 띄는 변경은 헤드 관련 인자 처리 방식입니다. 기존에는 num_heads 하나로 쿼리, 키, 밸류 헤드 수를 모두 지정했지만, 이제는 GQA(Grouped-Query Attention) 및 MQA(Multi-Query Attention)를 지원하기 위해 num_qo_heads (Query/Output Heads)와 num_kv_heads (Key/Value Heads)로 분리되었습니다.

Before:

DEFAULT_CONFIGS = (
    (4, 8, 4096, 128, False),
    (1, 8, 32768, 128, False),
)
# ...
parser.add_argument(
    "--num-heads",
    "--num_heads",
    type=int,
    nargs="+",
    default=None,
    help="Number of attention heads.",
)

After:

DEFAULT_CONFIGS = (
    (4, 8, 8, 4096, 128, False), # B, Hq, Hkv, S, D, Causal
    (1, 8, 8, 32768, 128, False),
)
# ...
parser.add_argument(
    "--num-qo-heads",
    "--num_qo_heads",
    type=int,
    nargs="+",
    default=None,
    help="Number of query/output attention heads.",
)
parser.add_argument(
    "--num-kv-heads",
    "--num_kv_heads",
    type=int,
    nargs="+",
    default=None,
    help="Number of key/value attention heads.",
)
parser.add_argument(
    "--return-lse",
    action="store_true",
    help="Compute log-sum-exp in addition to the attention output.",
)

또한, return_lse 옵션이 추가되어 LSE(Log-Sum-Exp) 값을 반환할지 여부를 제어할 수 있게 되었습니다. 이는 LSE가 필요 없는 일반적인 LLM 추론 시나리오에서 불필요한 연산과 메모리 할당을 줄여 성능을 향상시킵니다.

BenchConfig 클래스도 num_qo_heads, num_kv_heads, return_lse 필드를 포함하도록 업데이트되었습니다.

attention_flops 함수에서도 num_heads 대신 num_qo_heads를 사용하여 FLOPs 계산이 정확하게 이루어지도록 수정되었습니다.

bench_config 함수 내부에서는 return_lse 플래그에 따라 lse 텐서를 초기화하는 로직이 추가되었습니다. return_lse=False (기본값)일 경우 lse는 None으로 유지되어 관련 연산이 수행되지 않습니다.

    lse = None
    if return_lse:
        lse = torch.empty(
            config.batch_size,
            config.num_qo_heads,
            config.seq_len,
            dtype=torch.float32,
            device="cuda",
        )

FP8 베이스라인과의 비교를 위해 fp8_baseline_available 조건이 추가되었고, return_lse 옵션이 FP8 경로에도 적용되도록 수정되었습니다.

2. `flashinfer/nvfp4_attention_sm120.py` - API 변경 및 기본값 조정

이 파일은 NVFP4 어텐션 커널의 Python 인터페이스를 정의합니다. 가장 중요한 변경 사항은 return_lse의 기본값이 True에서 False로 변경된 것입니다.

기존 동작 (v0.6.17 이전): flashinfer.nvfp4_attention_sm120_fwd(...) 호출 시 기본적으로 (output, lse) 튜플을 반환했습니다.

새로운 기본값 (PR 이후): flashinfer.nvfp4_attention_sm120_fwd(...) 호출 시 기본적으로 output 텐서만 반환합니다.

리뷰어 코멘트 분석:

[tiffany940107] 이 변경은 이전 버전과의 호환성 문제를 야기할 수 있다는 점을 지적했습니다. return_lse=False가 기본값이 되면, return_lse 인자를 명시적으로 전달하지 않던 기존 호출자들은 예상치 못한 동작 변경을 겪을 수 있습니다.
[saltyminty] 이 우려에 동의하며, 기본값을 True로 유지하는 것이 더 안전할 수 있다고 제안했습니다.
[tiffany940107] (이후 수정) return_lse=False가 기본값이 된 것은 성능 최적화를 위한 의도적인 결정임을 명확히 했습니다. LSE가 필요 없는 일반적인 LLM 추론 시나리오에서는 이 옵션이 더 빠른 경로를 제공합니다. LSE가 필요한 경우 return_lse=True를 명시적으로 전달하면 기존과 동일하게 (out, lse)를 반환받을 수 있습니다. 따라서 이는 성능 지향적인 기본값 설정이며, 호환성 경로는 유지됩니다.

이 변경은 API의 기본 동작을 변경하는 것이므로, 사용자는 자신의 워크로드에 따라 return_lse=True를 명시적으로 설정해야 할 수 있습니다.

3. `include/flashinfer/attention/sm120/nvfp4_attention_sm120/compute/consumer/softmax.cuh` - Softmax 최적화

이 파일은 SM120 아키텍처에서 NVFP4 어텐션의 Softmax 계산 부분을 담당하는 CUDA 커널 코드입니다. 여러 최적화가 이루어졌습니다.

마스킹 없는 경로에서 분기 제거: 비인과적(noncausal) 어텐션에서 마스크가 없는 경우, Softmax 계산 루프에서 불필요한 분기(branch)를 제거했습니다. 이는 GPU의 SIMT(Single Instruction, Multiple Threads) 실행 모델에서 분기 예측 실패로 인한 성능 저하를 방지합니다.
SOFTMAX_FMA_EXP2 분기 처리: SOFTMAX_FMA_EXP2 매크로가 정의되지 않은 현재 PTX 경로에서는 ex2.approx.ftz.f32 (approximate exponential with flush-to-zero)를 사용합니다. 리뷰어는 SOFTMAX_FMA_EXP2가 활성화될 경우 발생할 수 있는 잠재적인 오버플로우/언더플로우 문제를 지적했고, 이에 대한 대응으로 해당 분기 내부에 값의 범위를 제한하는 로직이 추가되었습니다. 하지만 현재 코드에서는 이 분기가 활성화되지 않으므로 실제 PTX 코드에는 영향을 미치지 않습니다.
수학적 정확성: 2^-126은 부동 소수점의 최소 정규값(normal value)이며, 이를 반환하는 것은 수학적으로 유효합니다. 이전에는 이를 0으로 강제하는 근사가 있었으나, 제거되었습니다.

리뷰어 코멘트 분석:

[tiffany940107] 는 SOFTMAX_FMA_EXP2 관련 코드가 실제 사용되지 않는 경로임을 확인하고, 해당 분기 내부에만 범위를 제한하는 수정으로 현재 코드에는 영향을 주지 않음을 명확히 했습니다.

4. `include/flashinfer/attention/sm120/nvfp4_attention_sm120/compute/mainloop.cuh` - 메인 루프 및 스코어-슬롯 재사용

이 파일은 어텐션 커널의 메인 계산 루프를 담당하며, 이번 PR의 핵심 성능 향상이 이루어진 부분입니다.

N64 스코어-슬롯 재사용: 가장 중요한 최적화입니다. 기존에는 N128 크기의 QK 스코어 타일을 계산한 후, 전체 N128 컬럼에 대한 최대값 및 온라인 소프트맥스 재스케일링을 수행했습니다. 이후, 두 개의 N64 스코어 슬롯을 순차적으로 소비했습니다. 이 PR에서는 다음과 같은 방식으로 스코어-슬롯을 재사용합니다:
1. 전체 N128 QK 스코어 타일을 계산하고, 모든 128개 컬럼에 대한 행 최대값 및 온라인 소프트맥스 재스케일링을 수행합니다.
2. 첫 번째 N64 스코어 슬롯을 소프트맥스 양자화하고 PV MMA에서 소비합니다.
3. 이 스코어-레지스터 슬롯을 즉시 다음 타일의 QK 스코어 계산에 재사용하는 동안, 두 번째 N64 슬롯이 소비됩니다.
4. 두 번째 슬롯에 대해서도 대칭적으로 반복합니다.
이 방식은 두 개의 전체 N128 스코어 프래그먼트를 메모리에 유지할 필요성을 없애고, QK 스코어 계산과 PV MMA 연산 간의 오버랩을 극대화합니다. 이는 전체 N128 소프트맥스 감소 연산의 정확성을 유지하면서 성능을 향상시킵니다.
메인 루프 분할: 메인 루프를 프롤로그(prologue), 분기 없는 스테디 스테이트(steady state), 최종 드레인(drain)으로 분할했습니다. 이는 커널의 구조를 개선하고 최적화 가능성을 높입니다.
동기화 감소: 불필요한 math-order 배리어를 제거하고 워프 동기화 횟수를 줄였습니다. 이는 GPU 워프 스케줄링의 효율성을 높여 레이턴시를 감소시킵니다.
마스킹 최적화: 비인과적(noncausal) 특수화에서 마스킹을 컴파일 타임에 제거했습니다. 인과적(causal) 순회 시, 첫 번째 타일만 대각선과 겹칠 수 있으며 후속 타일은 완전히 유효하기 때문입니다.
좌표 계산: ID 티(identity tensor)를 유지하는 대신 마스크의 행/열 좌표를 직접 계산합니다.
레지스터 할당 및 CTA 스케줄링 튜닝: 소비자 레지스터 할당, 지속적인 CTA(Cooperative Thread Array) 스케줄링 및 작업 분배를 짧고 긴 인과적/비인과적 워크로드 모두에 맞게 조정했습니다.
LSE 특수화: 출력 전용 경로에서 LSE 할당 및 쓰기를 피하기 위해 컴파일 타임 LSE 특수화를 도입했습니다.

리뷰어 코멘트 분석:

[tiffany940107] 는 group_id가 N32 스코어 그룹을 선택하며, _0 좌표는 K-타일 좌표이지 헤드 차원만을 나타내는 것이 아니라고 설명했습니다. 또한, D128의 멀티-타일 정확도 테스트가 통과함을 확인하며 K 데이터 재사용 관련 이슈가 없음을 밝혔습니다.

5. `include/flashinfer/attention/sm120/nvfp4_attention_sm120/compute/consumer/pv_mma.cuh` - PV MMA 최적화

이 파일은 PV(Position-Value) MMA 연산을 담당합니다. 리뷰어 코멘트에 따르면, 이 부분은 기존 코드에서 큰 변경 없이 유지되었으며, 특히 N64 스코어-슬롯 재사용 전략이 PV MMA 연산과 잘 통합되도록 설계되었습니다.

6. `tests/attention/test_nvfp4_attention_sm120.py` - 테스트 케이스 추가

새로운 기능과 최적화를 검증하기 위해 테스트 케이스가 추가되었습니다.

return_lse=False (출력 전용) 및 return_lse=True (출력 및 LSE) 두 경우 모두에 대한 테스트가 추가되었습니다.
출력 전용 및 LSE 활성화 시의 어텐션 출력 결과가 rtol=0, atol=5e-4 조건 하에서 일치함을 확인합니다.
LSE 활성화 시의 참조 검사도 통과함을 확인합니다.

리뷰어 코멘트 분석:

[tiffany940107] 는 LSE 관련 테스트가 누락된 점을 지적했고, 이후 출력 전용 및 LSE 활성화 예제를 추가하고 JSON 픽스처를 검증하겠다고 언급했습니다.

왜 이게 좋은가?

성능 향상

PR 설명에 포함된 성능 측정 결과는 매우 인상적입니다. RTX PRO 6000 Blackwell GPU (SM120)에서 BF16 입력/출력, D=128, return_lse=False 조건 하에서 측정된 결과입니다. (QKV 양자화 시간은 제외)

Shape	Mode	#3640 baseline (reported)	CuTe DSL reference (reported)	This PR (measured)	Latency vs #3640	Latency vs CuTe DSL
B4, H8, S4096, D128	Noncausal	0.299 ms / 920.5 TFLOP/s	0.262 ms / 1050.2 TFLOP/s	0.262 ms / 1050.1 TFLOP/s	-12.4%	same at reported precision
B1, H8, S32768, D128	Noncausal	4.970 ms / 884.9 TFLOP/s	4.398 ms / 1000.0 TFLOP/s	4.079 ms / 1078.2 TFLOP/s	-17.9%	-7.3%
B4, H8, S4096, D128	Causal	0.223 ms / 616.6 TFLOP/s	0.180 ms / 764.1 TFLOP/s	0.165 ms / 830.6 TFLOP/s	-26.0%	-8.3%
B1, H8, S32768, D128	Causal	2.958 ms / 743.3 TFLOP/s	2.516 ms / 874.0 TFLOP/s	2.207 ms / 996.5 TFLOP/s	-25.4%	-12.3%

특히 인과적(Causal) 어텐션에서 최대 -26.0%의 레이턴시 감소를 보여주며, 비인과적(Noncausal) 어텐션에서도 상당한 성능 향상이 있었습니다. TFLOPs 측면에서도 기존 대비 향상되거나 동등한 수준을 유지하며 효율성을 높였습니다.

기술적 교훈

스코어-슬롯 재사용의 중요성: LLM의 어텐션 연산에서 QK 스코어 계산과 PV MMA 연산 간의 데이터 종속성 및 오버랩 기회를 파악하는 것이 중요합니다. N128 타일 내에서 N64 단위로 스코어 슬롯을 재사용함으로써, GPU의 레지스터 파일과 연산 유닛을 더 효율적으로 활용할 수 있습니다. 이는 메모리 대역폭 병목을 줄이고 연산 유닛의 유휴 시간을 최소화하는 데 기여합니다.
API 설계와 기본값: return_lse의 기본값을 False로 변경한 것은 성능 최적화를 위한 과감한 결정입니다. 대부분의 LLM 추론 시나리오에서는 LSE 값이 필요 없으므로, 이를 기본적으로 비활성화하여 성능을 높이는 것은 합리적인 선택입니다. 다만, API 변경 시에는 기존 사용자를 위한 명확한 문서화와 호환성 옵션 제공이 필수적입니다. 리뷰어들의 피드백을 통해 이 점이 명확히 논의되었고, return_lse=True 옵션을 통해 기존 기능을 완벽하게 지원합니다.
컴파일 타임 최적화: 마스킹 제거, LSE 경로 분리 등 컴파일 타임에 특정 조건을 처리하는 것은 런타임 오버헤드를 줄이는 효과적인 방법입니다. GPU 커널에서는 분기 예측 실패나 불필요한 연산이 성능에 큰 영향을 미치므로, 가능한 한 컴파일 타임에 최적화하는 것이 좋습니다.
정확성과 성능의 균형: NVFP4와 같은 저정밀도 연산에서는 정확도 손실이 발생할 수 있습니다. 이 PR은 N64 스코어-슬롯 재사용 전략을 사용하면서도 전체 N128 소프트맥스 감소의 정확성을 유지하는 방법을 찾아냈습니다. 또한, 2^-126과 같은 부동 소수점의 특성을 고려하여 수학적 정확성을 높였습니다. 이는 성능 향상을 추구하되, 모델의 정확도를 해치지 않는 선에서 최적화를 진행하는 것의 중요성을 보여줍니다.
벤치마킹의 중요성: 다양한 시나리오(인과적/비인과적, 짧은/긴 시퀀스)에 대한 상세한 성능 측정 결과는 최적화의 효과를 명확히 보여줍니다. 또한, FP8과 같은 다른 양자화 방식과의 비교는 NVFP4의 이점을 부각시킵니다.

결론

이번 PR은 FlashInfer 라이브러리의 SM120 NVFP4 어텐션 커널 성능을 크게 향상시키는 중요한 개선 사항을 포함하고 있습니다. N64 스코어-슬롯 재사용이라는 핵심 최적화 기법과 더불어, API 기본값 조정, 컴파일 타임 최적화 등 다각적인 접근을 통해 실제 LLM 추론 환경에서 상당한 속도 향상을 기대할 수 있게 되었습니다. 이러한 최적화는 최신 GPU 아키텍처의 기능을 최대한 활용하고, LLM 모델의 효율성을 높이는 데 크게 기여할 것입니다.

참고 자료

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

PR Analysis 의 다른글

이전글 [flashinfer] FlashInfer: Blackwell 아키텍처를 위한 Recurrent-KDA Prefill 최적화
현재글 : [flashinfer] FlashInfer SM120 NVFP4 어텐션 최적화: N64 스코어-슬롯 재사용을 통한 성능 향상
다음글 [vllm] vLLM MoE 성능 최적화: flashinfer_nvlink_one_sided 도입

[flashinfer] FlashInfer SM120 NVFP4 어텐션 최적화: N64 스코어-슬롯 재사용을 통한 성능 향상

들어가며

코드 분석

1. `benchmarks/bench_sm120_attention.py` - 벤치마킹 스크립트 수정

2. `flashinfer/nvfp4_attention_sm120.py` - API 변경 및 기본값 조정

3. `include/flashinfer/attention/sm120/nvfp4_attention_sm120/compute/consumer/softmax.cuh` - Softmax 최적화

4. `include/flashinfer/attention/sm120/nvfp4_attention_sm120/compute/mainloop.cuh` - 메인 루프 및 스코어-슬롯 재사용

5. `include/flashinfer/attention/sm120/nvfp4_attention_sm120/compute/consumer/pv_mma.cuh` - PV MMA 최적화

6. `tests/attention/test_nvfp4_attention_sm120.py` - 테스트 케이스 추가

왜 이게 좋은가?

성능 향상

기술적 교훈

결론

참고 자료

댓글

관련 포스트

PR Analysis 의 다른글