[vllm] Blackwell에서 fp8 KV 캐시 기반 QSA sparse attention을 2.6배 빠르게: split-K를 대체하는 single-pass 커널

PR 링크: vllm-project/vllm#66512 상태: Merged | 변경: +673 / -6

들어가며

Qwen4Exp 모델은 QSA(sparse attention)를 사용합니다. 각 query row가 전체 context 중 indexer가 선택한 약 2048개의 토큰만 attend하는 구조입니다. 기존 구현은 Triton 기반 split-K 커널이었습니다. 선택된 토큰을 여러 split으로 나눠 부분 attention을 계산하고, merge 단계에서 합칩니다.

문제는 fp8-e4m3 KV 캐시 + 큰 batch 조합에서 나타났습니다. MTP(Multi-Token Prediction) decode/verify처럼 row 수가 수백 개를 넘어가면 split-K가 크게 비효율적입니다. 원인은 세 가지입니다.

• fp8 코드를 byte 단위로 shared memory에서 읽어 tile당 load 명령이 360개에 달합니다.
• 94개의 register spill이 발생합니다.
• 페이지 테이블 변환(logical → physical slot)이 attention loop 안에 섞여 있어 K/V 주소를 일찍 알 수 없고, 따라서 prefetch가 어렵습니다.

이 PR은 B200/B300(sm_100/sm_103)에서 큰 fp8 batch 전용 single-pass 커널을 새 모듈 qsa_fp8.py에 추가합니다. 결과는 decode/verify shape에서 2.63×, end-to-end median TPOT −4.4%입니다.

코드 분석

1. qsa.py: dispatch gate와 split-K 분리

핵심은 public entry인 qsa_sparse_paged_attention에 들어간 gate입니다.

Before: shape 검증 이후 곧바로 split-K 로직이 이어졌습니다.

    if not q.shape[0]:
        return out

    group_size = q.shape[1] // k_cache.shape[2]
    block_m = triton.next_power_of_2(group_size)
    selection_width = logical_indices.shape[1] - 1  # trailing column is the count

After: 조건을 만족하면 새 커널로 보내고, 기존 로직은 _qsa_sparse_paged_splitk로 그대로 옮겼습니다.

    if qsa_fp8.use_fp8_decode_kernel(q.shape[0] * k_cache.shape[2], k_cache.dtype):
        return qsa_fp8.qsa_sparse_paged_decode(
            q, k_cache, v_cache, logical_indices, block_table,
            token_to_req, out, output_gate_view, softmax_scale, output_scale,
        )
    return _qsa_sparse_paged_splitk(
        q, k_cache, v_cache, logical_indices, block_table,
        token_to_req, use_prefill_config, out, output_gate_view,
        softmax_scale, output_scale,
    )

이 리팩터링의 장점은 fallback 경로가 bitwise-identical하게 유지된다는 점입니다. bf16 캐시, 작은 batch, 다른 GPU에서는 기존 코드가 한 줄도 바뀌지 않은 채 실행됩니다. 새 커널의 위험 범위가 gate 안으로 한정되는 셈입니다.

Gate 조건은 rows × KV heads per rank > 64입니다. PR 설명에 따르면 이 기준 아래에서는 split-K가 더 빠릅니다(새 커널이 0.7–0.9×). 프로그램 수가 적으면 split-K의 K 방향 병렬화가 SM을 채우는 데 유리하기 때문입니다. 기준 위에서는 row × head만으로 이미 충분한 parallelism이 확보되어, split/merge 오버헤드가 없는 single-pass가 이깁니다.

Warmup도 새 커널의 specialization까지 미리 컴파일하도록 확장됐습니다.

    warmed += qsa_fp8.warmup_qsa_fp8_decode(
        kv_cache,
        block_table,
        num_query_heads=num_query_heads,
        selection_width=selection_width,
    )
    return tuple(warmed)

반환 타입도 tuple[tuple[int, int, int], ...]에서 tuple[tuple[int | str, ...], ...]로 넓혔습니다. 서빙 중 JIT 컴파일 stall을 막으려면 새 경로도 반드시 warmup 대상에 포함돼야 합니다.

2. qsa_fp8.py: slot pass + attention kernel

새 모듈은 두 단계로 구성됩니다.

def _is_sm100_family() -> bool:
    """True on sm_100/sm_103 (B200/B300): enables the fp8 decode kernel."""
    return current_platform.is_device_capability_family(100)


@triton.jit(do_not_specialize=["num_requests"])
def _qsa_physical_slots_kernel(
    indices_ptr,
    block_table_ptr,
    token_to_req_ptr,
    slots_ptr,
    ...
    TOPK: tl.constexpr,
    PAGE_SIZE: tl.constexpr,

(1) Slot pass. 각 row의 logical index를 page table로 한 번만 physical slot으로 변환해 int32 workspace에 기록합니다. 이 결과는 모든 KV head가 공유합니다. 기존 방식은 attention loop 안에서 block_table lookup을 하는 의존 load 체인이었습니다. 이를 별도 pass로 빼내면, attention 커널은 slot 배열만 읽으면 곧바로 K/V 주소를 알 수 있습니다. 덕분에 cp.async 기반 다단계 prefetch가 가능해집니다. do_not_specialize=["num_requests"]는 batch 크기가 바뀔 때마다 재컴파일되는 것을 막습니다.

(2) Attention kernel. e4m3 코드 두 개를 int16 한 쌍으로 묶어 gather하고, MMA 입력으로 쓸 fp16 half로 정확히 변환합니다. 효과는 다음과 같습니다.

• Byte 단위 shared-memory load가 tile당 360개에서 40개로 줄어듭니다.
• Split-K에서 발생하던 94개의 register spill이 사라집니다.
• Register 여유가 생겨 4 warps + 5-stage cp.async pipeline 구성이 가능해졌습니다.

3. 수치 정확성 설계

QK는 bf16, PV는 fp16으로 유지합니다. e4m3 → bf16/fp16 변환은 exact합니다. e4m3의 지수·가수 범위가 두 포맷 안에 완전히 포함되기 때문입니다. 따라서 개별 QK/PV 곱은 split-K와 완전히 같고, 다른 것은 fp32 누적의 summation grouping뿐입니다. 최적화를 하면서도 수치 차이의 원인을 명확히 한정한 좋은 설계입니다.

4. 테스트

@pytest.mark.parametrize(
    "num_query_heads,num_kv_heads", [(24, 2), (12, 1), (6, 1), (3, 1)]
)
@pytest.mark.parametrize("num_requests", [20, 100])
def test_qsa_fp8_kernel_correctness(...):
    ...
    token_to_req[-1] = -1  # CUDA-graph padding
    ...
    assert qsa_fp8_ops.use_fp8_decode_kernel(num_rows * num_kv_heads, k_cache.dtype)
    ...
    torch.testing.assert_close(actual[:-1], expected[:-1], rtol=2e-2, atol=2e-2)
    assert not actual[-1].any()

주목할 부분은 두 가지입니다.

• Gate가 실제로 열렸는지 assert합니다. 이 assert가 없으면, shape가 gate를 통과하지 못해 split-K로 빠지더라도 테스트가 무의미하게 통과합니다. 주석의 표현대로 "vacuous"한 테스트를 방지하는 장치입니다.
• CUDA-graph padding row(token_to_req = -1)의 출력이 0으로 남는지 확인합니다. 실서빙에서 graph capture 시 생기는 padding row를 정확히 반영한 검증입니다.

CI 측면에서는 B200 step key를 qwen4-exp-hc-kernels-b200에서 qwen4-exp-kernels-b200으로 일반화했습니다. 또 qsa.py/qsa_fp8.py 변경 시 테스트가 트리거되도록 의존성을 추가했습니다. 리뷰에서는 메인테이너(gau-nernst)가 /ci run을 여러 차례 실행하며 B200 CI 검증을 진행했습니다.

왜 이게 좋은가

커널 성능 (GB200, 2048 selected tokens/row, 24 query / 2 KV heads):

Workload	Rows	Split-K	New	Speedup
Decode/verify	400	396 µs	151 µs	2.63×
Decode/verify	1600	1153 µs	525 µs	2.20×
Prefill	4096–8192			1.37–1.42×

E2E (4×GB200, TP4, MTP3, concurrency 512):

Metric	main	this PR	Delta
Median TPOT	54.64 ms	52.22 ms	−4.4%
Output throughput	7695 tok/s	7975 tok/s	+3.6%

정확도는 GSM8K 96.78% → 96.70%로 사실상 동일합니다. AIME25 차이(90.8% → 85.8%)는 30문제·temperature 0.7 조건이므로 sampling noise 범위로 설명됩니다.

일반적 교훈

1. 간접 주소 변환은 hot loop 밖으로. Page table lookup처럼 의존적인 load를 별도 pass로 분리하면 주소를 미리 알 수 있게 되고, 그래야 async prefetch pipeline이 제대로 작동합니다.
2. Load 폭을 넓혀라. 8-bit 데이터를 byte 단위로 읽지 말고 int16/int32로 packing해서 읽으면 명령 수와 register 압력이 함께 줄어듭니다.
3. Register spill은 occupancy와 pipeline depth를 함께 죽인다. Spill을 없애야 warp 수와 stage 수를 늘릴 여유가 생깁니다.
4. 기존 커널을 고치기보다 분리하라. 작성자는 같은 아이디어를 split-K에 접목하면 2.46×까지 나온다는 것도 측정했습니다. 그럼에도 별도 커널(2.63×)을 택해 fallback의 bitwise 동일성을 지켰습니다.
5. Gate는 측정으로 정하라. Parallelism이 부족한 영역에서는 split-K가 여전히 유리합니다. 그래서 rows × KV heads > 64라는 경험적 임계값을 두었습니다.

한계도 명시돼 있습니다. 새 경로는 split-K와 bitwise-identical하지 않고, sm_103은 검증되지 않았으며, call마다 slot pass launch 1회와 rows × selection width 크기의 int32 workspace가 추가됩니다. 이런 trade-off를 숨기지 않고 정리한 점도 좋은 PR의 모범입니다.

참고 자료

• triton.jit — slot pass/attention 커널 정의에 사용된 Triton JIT 데코레이터
• torch.testing.assert_close — fp8 커널 정확성 테스트에 사용된 비교 함수
• torch.float8_e4m3fn / Tensor attributes — fp8 KV 캐시 dtype

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