[flashinfer] FlashInfer, NVFP4를 활용한 DeepSeek-V4 디코딩 성능 최적화

2026년 9월 24일수정: 2026년 9월 24일

PR 링크: flashinfer-ai/flashinfer#5443 상태: Merged | 변경: +14286 / -2

들어가며

최근 대규모 언어 모델(LLM)의 발전 속도는 눈부십니다. 특히 추론 단계에서의 성능 최적화는 모델의 실질적인 활용도를 결정짓는 중요한 요소입니다. NVIDIA의 최신 GPU 아키텍처를 최대한 활용하여 LLM 추론 속도를 극대화하는 라이브러리인 FlashInfer는 지속적으로 새로운 최적화 기법을 도입하고 있습니다. 이번 PR은 DeepSeek-V4 모델의 디코딩 과정에서 NVFP4(NVIDIA Variable Precision Floating Point 4) 양자화 기술을 실험적으로 도입하여 성능을 크게 향상시키는 내용을 담고 있습니다.

이 PR은 특히 #5403 이슈에서 요청된 DeepSeek-V4 모델의 디코딩 성능 개선을 목표로 합니다. 기존의 FP8 기반 디코딩 커널 대비 NVFP4를 사용하여 메모리 대역폭을 절약하고 연산 효율성을 높이는 것이 핵심입니다. 본 글에서는 이 PR의 코드 변경 사항을 분석하고, NVFP4 도입이 왜 성능 향상으로 이어지는지, 그리고 이 최적화가 가지는 일반적인 교훈은 무엇인지 살펴보겠습니다.

코드 분석

이번 PR의 핵심은 flashinfer/experimental/nvfp4_mla_decode/ 디렉토리에 새로운 실험적 API와 백엔드를 추가한 것입니다. 주요 변경 사항은 다음과 같습니다.

1. `benchmarks/bench_cake_nvfp4_mla_decode.py` - 성능 벤치마크 추가

새로운 NVFP4 디코딩 커널의 성능을 측정하기 위한 벤치마크 스크립트가 추가되었습니다. 이 스크립트는 다양한 KV 캐시 길이에 대해 NVFP4 커널의 지연 시간과 처리량을 측정하며, 선택적으로 기존 FlashInfer FP8 CuTe-DSL MLA 디코딩 커널과 비교합니다.

# benchmarks/bench_cake_nvfp4_mla_decode.py

# ... (기존 코드) ...

def bench_cake(
    batch,
    kv_len,
    num_heads,
    q_len,
    device,
):
    # ... (입력 생성 및 NVFP4 양자화) ...
    decode = prepare_nvfp4_batch_decode_with_kv_cache_mla(
        query,
        query_scale,
        kv_cache,
        kv_scale,
        block_tables,
        seq_lens,
        workspace,
        sm_scale=HEAD_DIM**-0.5,
        seq_lens_cpu=seq_lens.cpu(),
    )
    return median_ms(decode), decode.plan

# ... (trtllm_batch_decode_with_kv_cache_mla 비교 함수 등) ...

def main():
    # ... (인자 파싱 및 루프) ...
    for kv_len in args.kv:
        ms, plan = bench_cake(args.batch, kv_len, args.heads, args.q_len, device)
        # ... (성능 지표 계산 및 출력)

이 스크립트는 prepare_nvfp4_batch_decode_with_kv_cache_mla 함수를 호출하여 NVFP4 디코딩을 수행하고, bench_gpu_time_with_cupti를 통해 CUPTI를 이용한 GPU 커널 실행 시간을 측정합니다. 또한, TFLOPS 및 GB/s와 같은 성능 지표를 계산하여 출력합니다.

2. `flashinfer/experimental/nvfp4_mla_decode/README.md` - 실험적 API 문서화

새로운 NVFP4 디코딩 API에 대한 설명 문서가 추가되었습니다. 이 문서는 API의 목적, 지원하는 하드웨어(SM100/SM103), 데이터 형식(NVFP4 E2M1, UE4M3 스케일), 입력 텐서의 형태와 타입, 그리고 준비(prepare) 단계와 실행(run) 단계의 계약 사항을 명확히 설명합니다.

# Experimental NVFP4 DeepSeek-V4 decode attention

Both this API and its Cake backend are experimental and may change without
backward compatibility. Calling the API explicitly opts into the experimental
feature and emits FlashInfer's experimental API warning. There is no automatic
backend selection. Tracking: flashinfer-ai/flashinfer#5403.

`flashinfer.mla.prepare_nvfp4_batch_decode_with_kv_cache_mla(...)` serves the
DeepSeek-V4 main-attention decode geometry: paged MQA over one 512-wide NVFP4
latent row per token that is both K and V, 64-token pages, `q_len` query
tokens per request (derived from the query rows; six for DeepSeek-V4) with a
causal mask inside the block against the request's last tokens, an optional per-head attention sink, BF16 output and natural-log FP32
LSE. Query and KV are E2M1 values packed two per byte (256 bytes per row)
with UE4M3 block-16 scales (32 bytes per row); `quantize_nvfp4` in
`cake_backend.py` produces both from a float tensor.

| Tensor | Shape | dtype |
| --- | --- | --- |
| `query` | `[batch * q_len, num_heads, 256]` | uint8 (packed E2M1) |
| `query_scale` | `[batch * q_len, num_heads, 32]` | uint8 (UE4M3) |
| `kv_cache` | `[num_pages, 64, 256]` | uint8 (packed E2M1) |
| `kv_scale` | `[num_pages, 64, 32]` | uint8 (UE4M3) |
| `block_tables` | `[batch, max_pages]` | int32 |
| `seq_lens` | `[batch]` | int32 |
| `sinks` (optional) | `[num_heads]` | float32 |
| `out` | `[batch * q_len, num_heads, 512]` | bfloat16 |
| `lse` | `[batch * q_len, num_heads]` | float32 (natural log) |

# ... (이하 생략) ...

문서에는 NVFP4 형식의 세부 사항, 입력 텐서의 크기 및 타입, 그리고 workspace_buffer를 활용한 메모리 관리 방식 등이 상세히 기술되어 있습니다.

3. `flashinfer/experimental/nvfp4_mla_decode/cake_backend.py` - NVFP4 디코딩 백엔드 구현

이 파일은 실제 NVFP4 디코딩 커널을 구현하고 관리하는 핵심 부분입니다. DeepSeek-V4의 특정 요구사항(64개 헤드, head_dim 512, 페이지 크기 64)에 맞춰 최적화되었습니다.

주요 특징은 다음과 같습니다:

NVFP4 E2M1 양자화: Q, K, V 데이터는 NVFP4 E2M1 형식으로 저장됩니다. 이는 1바이트당 2개의 FP4 값을 저장하여 메모리 사용량을 크게 줄입니다. UE4M3 블록-16 스케일이 함께 사용됩니다.
BMM1 (QK) 최적화: QK 행렬 곱셈은 NVFP4 블록 스케일 텐서 코어 MMA(Matrix Multiply-Accumulate) 연산으로 수행됩니다. 스케일 값은 텐서 메모리에 게시됩니다.
P (Attention Scores) 처리: 스케일링된 P 행렬은 FP8 E4M3 형식으로 변환됩니다.
BMM2 (PV) 최적화: PV 행렬 곱셈은 FP8 MMA 연산으로 수행됩니다.
Split-KV Combine Kernel: 긴 시퀀스를 처리하기 위해 CTA(Compute Thread Array) 간에 KV 캐시를 분할하고 결합하는 커널이 사용됩니다. 이는 호스트 워크 플랜에 의해 관리됩니다.
하드웨어 지원: Compute Capability 10.0 (B200) 및 10.3 (GB300/B300)을 지원합니다. SM120/SM121은 tcgen05/TMEM 요구사항 때문에 지원되지 않습니다.

# flashinfer/experimental/nvfp4_mla_decode/cake_backend.py

def prepare_nvfp4_batch_decode_with_kv_cache_mla(
    query: torch.Tensor,
    query_scale: torch.Tensor,
    kv_cache: torch.Tensor,
    kv_scale: torch.Tensor,
    block_tables: torch.Tensor,
    seq_lens: torch.Tensor,
    workspace_buffer: torch.Tensor,
    sm_scale: float = HEAD_DIM**-0.5,
    seq_lens_cpu: Optional[torch.Tensor] = None,
    attention_sinks: Optional[torch.Tensor] = None,
    backend: str = "cake",
) -> NVFP4MLADecodeRunner:
    # ... (입력 검증 및 워크 플랜 생성)

    if backend == "cake":
        # ... (Cake 백엔드 초기화 및 Runner 반환)
        return NVFP4MLADecodeRunner(
            plan,
            query,
            query_scale,
            kv_cache,
            kv_scale,
            block_tables,
            seq_lens,
            workspace_buffer,
            attention_sinks,
        )
    else:
        raise ValueError(f"Unknown backend: {backend}")

# ... (quantize_nvfp4, nvfp4_mla_decode_workspace_size 등 헬퍼 함수)

prepare_nvfp4_batch_decode_with_kv_cache_mla 함수는 입력 데이터를 검증하고, 호스트 워크 플랜을 생성하며, workspace_buffer에서 필요한 메모리를 할당합니다. 이후 NVFP4MLADecodeRunner 객체를 반환하며, 이 객체의 run 메소드를 호출하면 실제 디코딩 커널이 실행됩니다.

4. `.pre-commit-config.yaml` - CI 설정 업데이트

새로운 실험적 기능이 추가됨에 따라, CI 설정 파일인 .pre-commit-config.yaml에 새로운 파일 경로가 exclude 목록에 추가되었습니다. 이는 해당 파일들이 코드 포맷팅이나 린팅 대상에서 제외됨을 의미하며, 실험적인 코드임을 나타냅니다.

--- a/.pre-commit-config.yaml
+++ b/.pre-commit-config.yaml
@@ -33,7 +33,7 @@
 # benefit and obscure its history.
 # GDN CP generated CUDA and its source registry are also receipt-bound exports.
 # Preserve their bytes; the handwritten GDN adapters and tests stay checked.
-exclude: ^(?:flashinfer/msa_ops/cute_dsl/sparse_decode_nvfp4_sm100\.py|csrc/(?:cake_dsv4/|cake_sampling/generated/|cake_minimax_h3_mxfp8_pre_attention/|cake_minimax_h3_nvfp4_pre_attention/|kda/flashkda_generated_.*|blackwell_msa/|concat_mla/|cake_all_gather_matmul/|cake_moe_finalize_allreduce_fusion/|cake_nvfp4_svdquant_gemm/|cake_selective_state_update/(?:cuda|host)/|cake_mamba_ssd_combined/generated/|cake_trtllm_moe_allreduce_fusion/|fused_moe/warp_decode/generated/|kda/cake_flashkda_blackwell_evolution_.*\.cu)|flashinfer/(?:experimental/(?:cake_mxfp8_megamoe_ep16|sm110_gqa_decode|nvfp4_attention|balanced_gqa_decode|minimax_h3_varlen_attention)/csrc/|moe_ep/kernel_src/(?:sm100/cutedsl_megamoe|sm107/next_cutedsl_megamoe|sm90/pull_style_cutedsl_megakernel|sm120/swapab_cutedsl_megakernel)/src/)|csrc/gdn/gdn_cp/sm_103a/|csrc/generated/dcp_alltoall/|flashinfer/jit/cake_gdn_cp_generated\.py$)
+exclude: ^(?:flashinfer/msa_ops/cute_dsl/sparse_decode_nvfp4_sm100\.py|csrc/(?:cake_dsv4/|cake_sampling/generated/|cake_minimax_h3_mxfp8_pre_attention/|cake_minimax_h3_nvfp4_pre_attention/|kda/flashkda_generated_.*|blackwell_msa/|concat_mla/|cake_all_gather_matmul/|cake_moe_finalize_allreduce_fusion/|cake_nvfp4_svdquant_gemm/|cake_selective_state_update/(?:cuda|host)/|cake_mamba_ssd_combined/generated/|cake_trtllm_moe_allreduce_fusion/|fused_moe/warp_decode/generated/|kda/cake_flashkda_blackwell_evolution_.*\.cu)|flashinfer/(?:experimental/(?:cake_mxfp8_megamoe_ep16|sm110_gqa_decode|nvfp4_attention|nvfp4_mla_decode|balanced_gqa_decode|minimax_h3_varlen_attention)/csrc/|moe_ep/kernel_src/(?:sm100/cutedsl_megamoe|sm107/next_cutedsl_megamoe|sm90/pull_style_cutedsl_megakernel|sm120/swapab_cutedsl_megakernel)/src/)|csrc/gdn/gdn_cp/sm_103a/|csrc/generated/dcp_alltoall/|flashinfer/jit/cake_gdn_cp_generated\.py$)

repos:
  # Standard hooks

nvfp4_mla_decode가 추가된 것을 확인할 수 있습니다.

왜 이게 좋은가?

이 PR은 다음과 같은 이유로 좋은 최적화 및 개선이라고 할 수 있습니다.

메모리 대역폭 절약 및 연산 효율성 증대: NVFP4는 FP8보다 더 낮은 정밀도를 가지므로, 동일한 데이터를 저장하는 데 더 적은 메모리를 사용합니다. 이는 메모리 대역폭이 병목 현상이 되는 LLM 추론에서 매우 중요합니다. 또한, 더 적은 비트를 처리하므로 연산 자체의 효율성도 높아질 수 있습니다. PR 설명에 따르면, NVFP4 E2M1 형식은 1바이트당 2개의 FP4 값을 저장하여 기존 FP8보다 2배의 압축률을 제공합니다.
최신 하드웨어 기능 활용: 이 커널은 NVIDIA의 최신 GPU 아키텍처(SM100/SM103)에서 제공하는 텐서 코어 기능을 활용하도록 설계되었습니다. 특히 NVFP4 및 FP8 MMA 연산을 지원하는 하드웨어 기능을 통해 높은 성능을 달성합니다.
성능 향상: 벤치마크 결과에 따르면, NVFP4 기반 커널은 기존 FlashInfer FP8 CuTe-DSL MLA 디코딩 커널 대비 상당한 속도 향상을 보여줍니다. 예를 들어, GB300 (SM103) GPU에서 KV 길이가 128K일 때 1.240배의 속도 향상을 기록했습니다. 이는 더 적은 메모리 사용량과 효율적인 연산 덕분입니다.

KV length GB300 this kernel ms GB300 FI FP8 CuTe-DSL ms GB300 speedup B200 this kernel ms B200 FI FP8 CuTe-DSL ms B200 speedup

128K 1.1641 1.4437 1.240x 1.6053 1.7817 1.110x
정확도 유지: PR 설명에 따르면, NVFP4 양자화로 인한 오차는 FP32 참조 모델 대비 허용 가능한 수준으로 유지됩니다. BF16 출력의 경우 atol/rtol 0.1, per-(row, head) relative L2 <= 0.05, LSE atol 0.05 수준이며, 잔여 오차는 주로 V 재양자화 과정에서 발생한다고 명시되어 있습니다 (relative L2 0.01-0.04).
실험적 기능 도입: 이 PR은 experimental 디렉토리에 코드를 추가함으로써, 아직 안정화되지 않았지만 잠재력이 큰 기술을 선제적으로 도입하고 검증하는 FlashInfer의 개발 철학을 잘 보여줍니다. 이는 향후 LLM 추론 기술 발전에 기여할 수 있습니다.

KV length	GB300 this kernel ms	GB300 FI FP8 CuTe-DSL ms	GB300 speedup	B200 this kernel ms	B200 FI FP8 CuTe-DSL ms	B200 speedup
128K	1.1641	1.4437	1.240x	1.6053	1.7817	1.110x

일반적인 교훈:

양자화의 중요성: LLM 추론 성능 향상에 있어 양자화는 필수적인 기법입니다. 특히 FP8을 넘어 NVFP4와 같은 더 낮은 정밀도의 양자화를 탐색하는 것은 메모리 및 연산 효율성을 극대화하는 데 중요합니다.
하드웨어 특화 최적화: 최신 GPU 아키텍처의 특정 기능을 활용하는 커널 개발은 성능 향상의 핵심입니다. NVFP4, FP8 MMA 등 최신 기술을 적극적으로 도입해야 합니다.
실험적 기능의 가치: 새로운 기술을 실험적으로 도입하고 검증하는 과정은 라이브러리의 경쟁력을 유지하고 기술 트렌드를 선도하는 데 중요합니다.

리뷰 댓글 분석

리뷰 댓글은 주로 CI 파이프라인 실행 및 결과에 대한 내용이었습니다. 여러 차례 run tests/experimental/test_cake_nvfp4_mla_decode.py 명령이 실행되었으나, 초기에는 파이프라인 실패 및 타임아웃 문제가 있었던 것으로 보입니다. 이는 실험적인 코드나 새로운 테스트 환경 설정 시 흔히 발생하는 문제이며, 지속적인 테스트 실행과 디버깅을 통해 해결되었습니다.

yyihuang님의 "Status: this revision is not the final one for review. A further optimization round on the decode kernel is in progress; the branch will be re-exported and this PR updated when it lands. The current head (a64cf168) is complete and validated as described in the body." 라는 댓글은 이 PR이 최종 검토 단계 이전에 추가적인 최적화가 진행 중임을 시사합니다. 이는 개발 과정에서 성능 개선을 위한 지속적인 노력이 이루어지고 있음을 보여줍니다.