본문으로 건너뛰기

[vllm] vLLM Qwen3.8-Flash-Next: QSA Indexer 캐시를 위한 FP8 지원으로 메모리 효율 및 성능 최적화

PR 링크: vllm-project/vllm#54890 상태: Merged | 변경: +142 / -24

vLLM Qwen3.8-Flash-Next: QSA Indexer 캐시를 위한 FP8 지원으로 메모리 효율 및 성능 최적화

들어가며

대규모 언어 모델(LLM)의 추론 성능을 최적화하는 것은 항상 중요한 과제입니다. 특히, KV 캐시(Key-Value Cache)는 모델의 컨텍스트 길이에 비례하여 메모리 사용량이 급증하기 때문에, 이를 효율적으로 관리하는 것이 전체 시스템의 확장성과 처리량에 큰 영향을 미칩니다. vLLM 프로젝트는 이러한 LLM 추론 최적화에 앞장서고 있으며, 이번에 분석할 PR은 Qwen3.8-Flash-Next 모델의 QSA(Quantized Sparse Attention) Indexer 캐시에 FP8(8-bit Floating Point) 지원을 추가하여 메모리 효율성을 극대화하고 성능을 개선하는 중요한 변경사항을 담고 있습니다.

기존에는 주로 bfloat16과 같은 16비트 부동소수점 형식을 사용하여 KV 캐시를 저장했지만, FP8은 8비트만을 사용하여 동일한 정보를 저장함으로써 메모리 사용량을 절반으로 줄일 수 있습니다. 이는 특히 긴 컨텍스트 길이에서 막대한 메모리 절감 효과를 가져오며, FP8 연산을 가속화하는 최신 GPU 하드웨어(예: NVIDIA Hopper 및 Blackwell 아키텍처의 Tensor Cores)를 활용하여 연산 속도까지 향상시킬 수 있습니다. 이 PR은 QSA Indexer 캐시의 데이터 타입을 bfloat16에서 torch.float8_e4m3fn으로 확장함으로써 이러한 이점을 실현하고자 합니다.

코드 분석

이번 PR의 핵심은 QSA Indexer 캐시가 FP8 데이터 타입을 사용할 수 있도록 시스템을 확장하고, 이에 따른 정확성 검증 로직을 추가하는 것입니다. 주요 변경사항을 파일별로 살펴보겠습니다.

tests/models/qwen4_exp/test_qsa_pre_indexer.py

이 테스트 파일은 QSA Indexer의 fused 및 unfused 구현 간의 일치성을 검증합니다. FP8 지원을 위해 가장 중요한 변경사항은 다음과 같습니다.

  1. FP8 정밀도 검증 함수 추가 (assert_fp8_within_one_ulp) FP8과 같은 저정밀도 형식에서는 bfloat16처럼 비트 단위의 완벽한 일치(bitwise identical)를 기대하기 어렵습니다. 중간 계산 과정에서의 반올림 오차나 누적 순서의 차이로 인해 결과가 미세하게 달라질 수 있기 때문입니다. 이를 위해 PR은 assert_fp8_within_one_ulp라는 새로운 검증 함수를 도입했습니다. 이 함수는 FP8 값들이 ULP(Unit in the Last Place) 기준으로 1 이내의 차이를 보이거나, 특정 절대 오차 범위(2^-8) 내에 있는지 확인합니다.

    # Before: (No specific FP8 assertion logic)
# torch.testing.assert_close(fused_query, unfused_query, rtol=RTOL, atol=ATOL)

# After:
def assert_fp8_within_one_ulp(actual: torch.Tensor, expected: torch.Tensor) -> None:
    # e4m3 is sign-magnitude, so within a sign the uint8 code order matches the
    # value order and one ulp is one code step. The two paths' intermediates
    # differ in the pooling accumulation order, which at denormal magnitudes
    # (absolute grid step 2^-9) shows up as up to 2 code steps.
    code_diff = (
        actual.view(torch.uint8).int() - expected.view(torch.uint8).int()
    ).abs()
    abs_diff = (actual.float() - expected.float()).abs()
    assert bool(((code_diff <= 1) | (abs_diff <= 2**-8)).all())

  2. indexer_dtype 파라미터화 및 조건부 검증 test_qsa_fused_pre_indexer_matches_unfused 테스트 함수에 indexer_dtype 파라미터를 추가하여 torch.bfloat16과 torch.float8_e4m3fn 두 가지 데이터 타입으로 테스트를 실행할 수 있게 했습니다. 또한, fused_compressed_storage와 fused_query 텐서 생성 시 이 indexer_dtype을 사용하도록 변경했습니다.

    --- a/tests/models/qwen4_exp/test_qsa_pre_indexer.py
+++ b/tests/models/qwen4_exp/test_qsa_pre_indexer.py
@@ -49,8 +49,21 @@ def _make_block_table(block_counts):
     return table, num_blocks
 
 
+def assert_fp8_within_one_ulp(actual: torch.Tensor, expected: torch.Tensor) -> None:
+    # ... (assert_fp8_within_one_ulp function content) ...
+
+
 @requires_qsa_kernels
 @pytest.mark.usefixtures("default_vllm_config")
+@pytest.mark.parametrize("indexer_dtype", [torch.bfloat16, torch.float8_e4m3fn])
 @pytest.mark.parametrize(
     "mrope,is_2d_positions,cache_rope_positions,state_size,seq_lens,query_lens,history_lens",
     [
@@ -74,6 +87,7 @@ def _make_block_table(block_counts):
     ],
 )
 def test_qsa_fused_pre_indexer_matches_unfused(
+    indexer_dtype,
     mrope,
     is_2d_positions,
     cache_rope_positions,
@@ -215,7 +229,7 @@ def test_qsa_fused_pre_indexer_matches_unfused(
     fused_compressed_storage = torch.zeros(
         num_compressed_blocks,
         compressed_page_elements + 16,
-        dtype=torch.bfloat16,
+        dtype=indexer_dtype,
         device=device,
     )
     fused_compressed = torch.as_strided(
@@ -231,7 +245,7 @@ def test_qsa_fused_pre_indexer_matches_unfused(
     q_weight = torch.randn(D, dtype=torch.bfloat16, device=device) * 0.2
     k_weight = torch.randn(D, dtype=torch.bfloat16, device=device) * 0.2
 
-    fused_query = torch.empty(num_tokens, HQ, D, dtype=torch.bfloat16, device=device)
+    fused_query = torch.empty(num_tokens, HQ, D, dtype=indexer_dtype, device=device)
     qsa_pre_indexer(
         projected_qk[:, : HQ * D],
         projected_qk[:, HQ * D :],
@@ -290,8 +304,19 @@ def test_qsa_fused_pre_indexer_matches_unfused(
     if rope_positions is not None:
         qsa_store_cache_rows(rope_positions, raw_slots, position_rows)
 
-    torch.testing.assert_close(fused_query, unfused_query, rtol=RTOL, atol=ATOL)
+    if indexer_dtype == torch.float8_e4m3fn:
+        # Both paths round the same ~bf16 intermediates to e4m3. Bitwise
+        # equality (the dsv4 indexer test's bar) does not hold here: the 1D
+        # reference rope (forward_cuda) differs from _norm_rope by up to 1
+        # bf16 ulp, and even the MRoPE pair flips a code occasionally.
+        unfused_query = unfused_query.to(indexer_dtype)
+        assert_fp8_within_one_ulp(fused_query, unfused_query)
+    else:
+        torch.testing.assert_close(fused_query, unfused_query, rtol=RTOL, atol=ATOL)
     assert torch.equal(fused_raw.view(torch.int16), unfused_raw.view(torch.int16))
-    torch.testing.assert_close(
-        fused_compressed, unfused_compressed, rtol=RTOL, atol=ATOL
-    )
+    if indexer_dtype == torch.float8_e4m3fn:
+        assert_fp8_within_one_ulp(fused_compressed, unfused_compressed)
+    else:
+        torch.testing.assert_close(
+            fused_compressed, unfused_compressed, rtol=RTOL, atol=ATOL
+        )

    이러한 조건부 로직은 FP8의 특성을 고려하여 테스트의 견고성을 높입니다.

tests/models/qwen4_exp/test_qsa_reference.py

이 파일은 QSA의 디코드(decode) 및 프리필(prefill) 선택(selection) 로직의 정확성을 검증합니다. 여기에서도 dtype 파라미터화를 통해 FP8 지원을 추가했습니다.

  1. dtype 파라미터화 및 텐서 변환 test_qsa_decode_selection_correctness 및 test_qsa_prefill_selection_correctness 함수에 dtype 파라미터를 추가하고, q 및 cache 텐서를 해당 dtype으로 변환하도록 변경했습니다.

    --- a/tests/models/qwen4_exp/test_qsa_reference.py
+++ b/tests/models/qwen4_exp/test_qsa_reference.py
@@ -666,13 +668,16 @@ def make_buffers() -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tens
         (4, 33),
     ],
 )
+@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float8_e4m3fn])
 def test_qsa_decode_selection_correctness(
-    decode_query_len: int, num_requests: int
+    decode_query_len: int, num_requests: int, dtype: torch.dtype
 ) -> None:
     torch.manual_seed(1)
     heads, head_dim = 4, 128
     rows = num_requests * decode_query_len
-    q = torch.randn(rows, heads, head_dim, device="cuda", dtype=torch.bfloat16)
+    q = torch.randn(rows, heads, head_dim, device="cuda", dtype=torch.bfloat16).to(
+        dtype
+    )
     page_size, pages_per_request, max_sequence_length = (
         (16, 40, 2560) if num_requests > 32 else (4, 20, 320)
     )
@@ -684,7 +689,7 @@ def test_qsa_decode_selection_correctness(
         head_dim,
         device="cuda",
         dtype=torch.bfloat16,
-    )
+    ).to(dtype)
     page_table = torch.randperm(num_pages, device="cuda", dtype=torch.int32).reshape(
         num_requests, pages_per_request
     )

  2. FP8 로짓(logits) 검증 로직 FP8 로짓의 경우, bf16보다 top-k 경계에서 동점(tie)이 더 자주 발생할 수 있으므로, 인덱스 일치성보다는 선택된 값들의 멀티셋(multiset)을 비교하는 방식으로 검증 로직을 변경했습니다. 또한, SM90(Hopper) 아키텍처에서는 wgmma가 FP8을 낮은 정밀도로 누적하기 때문에 rtol=atol=1e-3와 같은 허용 오차를 적용하고, SM100(Blackwell)에서는 tcgen05가 정확한 fp32를 사용하므로 오차를 적용하지 않는 등 하드웨어별 특성을 고려했습니다.

    --- a/tests/models/qwen4_exp/test_qsa_reference.py
+++ b/tests/models/qwen4_exp/test_qsa_reference.py
@@ -736,14 +741,39 @@ def test_qsa_decode_selection_correctness(
         compress_ratio,
     )
 
+    if dtype == torch.float8_e4m3fn:
+        # fp8 logits tie at the top-k boundary more often than bf16, so index
+        # identity is not stable; compare the selected value multisets.
+        # SM90 wgmma accumulates fp8 in reduced precision (~3e-4 abs
+        # observed); SM100 tcgen05 is exact fp32.
+        rtol = atol = 1e-3 if current_platform.is_device_capability(90) else None
+        logits = _qsa_mqa_paged_reference(
+            q, cache, page_table, token_to_req, visible_blocks
+        )
+        for row in range(rows):
+            selected = actual[row][actual[row] >= 0]
+            wanted = expected[row][expected[row] >= 0]
+            assert selected.numel() == wanted.numel()
+            torch.testing.assert_close(
+                logits[row, selected.long()].sort().values,
+                logits[row, wanted.long()].sort().values,
+                rtol=rtol,
+                atol=atol,
+            )
+        return
+
     torch.testing.assert_close(actual.sort().values, expected.sort().values)

vllm/models/qwen4_exp/common/qsa_cache.py

이 파일은 QSA 캐시의 백엔드를 정의합니다. 여기서 FP8 지원을 공식적으로 선언합니다.

  1. supported_kv_cache_dtypes 확장 QSAStateBackend 클래스의 supported_kv_cache_dtypes에 `

참고 자료

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

댓글

관련 포스트

PR Analysis 의 다른글