[flashinfer] DeepSeek-V4.1 Sparse MLA Decode를 위한 Prims-TS 백엔드: FP4 캐시에서 바로 E4M3 MMA까지

PR 링크: flashinfer-ai/flashinfer#6988 상태: Merged | 변경: +9529 / -463

들어가며

DeepSeek-V4.1(DS4.1)의 sparse MLA decode는 독특한 형태를 가집니다. H64 / Sq1 / D512 attention이 128행의 SWA(sliding window) 윈도우와 최대 512개로 선택된 압축(FP4) 행을 동시에 attend합니다. 기존 FlashInfer의 CuTe DSL(Frost) 백엔드는 다음과 같은 한계가 있었습니다.

• SM100 전용이라 GB300(SM103), Rubin(SM107)에서는 동작하지 않음
• Q/P MMA를 BF16으로 수행 → FP8 MMA peak(BF16의 2배)를 활용하지 못함

이 PR은 task-scheduled CuTe-DSL 커널 프레임워크인 Prims-TS 기반의 실험적 decode 백엔드를 추가합니다. 핵심 아이디어는 "serving cache 포맷(E4M3 SWA, packed FP4 main) 그대로 읽어서 E4M3 QK/PV MMA로 계산"하는 것이며, 런타임 메타데이터 갱신과 CUDA Graph replay까지 지원합니다. 기본값은 여전히 backend="cute-dsl"이므로 기존 호출자에는 영향이 없습니다.

코드 분석

1. benchmarks/bench_deepseek_v41_decode.py — 백엔드별로 다른 입력 계약

가장 먼저 눈에 띄는 변화는 벤치마크 fixture가 "하나의 입력"이 아니라 "백엔드별 arm"을 만든다는 점입니다.

Before:

def fixture(batch, context, seed):
    torch.manual_seed(seed + batch)
    q = torch.randn(batch, 1, 64, 512, device="cuda", dtype=torch.bfloat16)
    w = torch.randn(batch * 128, 512, device="cuda", dtype=torch.bfloat16)
    swa = deepseek_v41_quantize_cache(w, format="swa_mxfp8")

After:

shared = main, wi, ci, sink
arms = {}
if "cute-dsl" in backends or flashmla:
    legacy = (
        q_source,
        deepseek_v41_quantize_cache(w_source, format="swa_mxfp8", kv_layout="NHD"),
        *shared,
    )
    ...
if "prims-ts" in backends:
    arms["prims-ts"] = (
        q_source.to(torch.float8_e4m3fn),
        deepseek_v41_quantize_cache(w_source, format="swa_fp8", kv_layout="NHD"),
        *shared,
    )

주목할 점은 두 가지입니다.

1. 같은 seeded BF16 source에서 파생되지만, Prims-TS arm은 Q를 float8_e4m3fn으로, SWA를 새 format="swa_fp8"로 양자화합니다. 정밀도 계약이 다르다는 사실을 코드 수준에서 명시합니다.
2. 주석의 "Do not allocate a second production-sized global pool or converted cache"처럼, FP4 main pool과 인덱스는 모든 arm이 공유합니다. 대형 context(32768)에서 메모리 중복 할당을 피하면서 동일 데이터로 공정 비교를 보장합니다.

2. FP64 oracle — 전체 캐시를 변환하지 않는 검증

정밀도가 다른 커널을 비교할 때 "어떤 기준으로 맞다고 할 것인가"가 중요합니다. PR은 Prims-TS 전용 독립 oracle을 추가했습니다.

def selected_main_fp8(pool, safe, valid):
    """Gather/mask selected bytes before decoding; never convert the full cache."""
    pages = pool.view(pool.shape[0], -1)
    page, row = safe // 64, safe % 64
    column = torch.arange(512, device=safe.device)
    payload = pages[page[..., None], row[..., None] * 256 + column // 2]
    payload = torch.where(valid[..., None], payload, 0)
    code = (payload >> (4 * (column % 2))) & 15
    lut = torch.tensor(
        [0, 0.5, 1, 1.5, 2, 3, 4, 6, -0.0, -0.5, -1, -1.5, -2, -3, -4, -6], ...)
    data = lut[code.long()]
    scale = pages[page[..., None], 64 * 256 + row[..., None] * 32 + column // 16]
    ...
    decoded = data * scale  # E2M1 * finite E4M3 is exact in float32.
    return common_unit_main(decoded)

이 코드는 packed FP4 page 레이아웃을 그대로 드러냅니다. 한 행은 256 byte(512개 E2M1 nibble) payload이고, page 뒤쪽에 group16 단위 E4M3 scale(행당 32 byte)이 붙습니다. 즉 행당 288 byte로, dense E4M3(512 B/token)를 요구하는 TRTLLM-GEN 대비 약 44% 작은 캐시 폭입니다.

common_unit_main은 amax나 request descale 없이 clamp(-448, 448) 후 E4M3로 반올림합니다. 커널이 "FP4 → E4M3 직접 변환"을 수행하는 것과 동일한 계약을 production 헬퍼 재사용 없이 독립적으로 재현한 것입니다.

오차 판정도 단순 rtol 대신 componentwise bound를 씁니다.

sensitivity = torch.einsum("bhk,bkd->bhd", weights, kv.abs()) * abs(bmm2_scale)
...
bound = (2**-4 + 3 * 2**-8 + 1e-5) * sensitivity + (2**-10 / 448) * underflow + 1e-6

E4M3 P의 반올림 오차(2^-4)와 입력 반올림을 |weights|·|kv| 감도에 곱해 원소별 허용 오차를 산출하고, LSE는 rtol=0, atol=1e-5로 엄격하게 검사합니다. 저정밀 커널 검증에서 "loose한 allclose로 대충 통과"를 막는 좋은 패턴입니다.

3. 커널 측 핵심 변경 (PR 설명 및 리뷰 기반)

diff가 잘려 커널 소스 전체를 인용할 수는 없지만, PR 작성자가 리뷰 스레드에 남긴 변경 요약이 최적화의 핵심입니다.

• FP4→E4M3 packed 변환: SM100/SM103에서 PTX 9.4의 mul.e4m3x4.e2m1p4x4.e4m3x4.satfinite를 사용해 8개 값당 12개 → 7개 명령어로 줄였고 결과는 bitwise 동일합니다. SM107은 ptxas에 packed 형태가 없어 F16 경로를 사용합니다.
• 소배치 KV split: 8 * B <= SMs이고 K >= 192이면 BK128 tile 하나를 split 단위로 CTA에 분산하고, PDL-chained reducer로 병합합니다. split 0이 SWA tile을 담당합니다.
• K=0 최적화: 완전히 마스킹된 main tile을 attend하지 않고 SWA tile 하나만 처리합니다.
• PDL(Programmatic Dependent Launch): 첫 입력 read 직전에 wait, 작은 grid에서는 early release로 kernel 간 launch 지연을 숨깁니다.
• Loader warp 4 → 8: FP4 변환과 SWA cp.async gather의 처리량을 늘렸습니다.

4. Workspace API — CUDA Graph 친화적 설계

from flashinfer.deepseek_v41 import (
    deepseek_v41_decode,
    deepseek_v41_decode_workspace_size,
    deepseek_v41_quantize_cache,
)

deepseek_v41_decode_workspace_size(B, K, backend="prims-ts")는 순수 산술 함수이며 B, K에 대해 단조 증가합니다. 따라서 최대 shape로 한 번 할당한 버퍼를 모든 작은 호출에서 재사용할 수 있고, CUDA Graph capture 시 caller-owned out/workspace만 넘기면 됩니다. sparse_topk_lens/seq_lens는 매 replay마다 갱신 가능해 graph 재캡처가 필요 없습니다.

왜 이게 좋은가

성능 수치

warm CUDA Graph replay, ABBA 순서, 각 arm FP64 gate 통과 후 측정 결과입니다.

비교	GPU	결과
vs FlashMLA 395cb3a	B200	1.51-2.71x 빠름
vs FlashMLA (sm_107a 빌드)	Rubin R100	1.30-3.04x 빠름
vs TRTLLM-GEN DSv4	B200	0.93-1.37x

예컨대 B200에서 B=1/K=512는 FlashMLA 21.33us → Prims-TS 7.88us입니다. 리뷰 스레드의 중간 비교를 보면 이전 head(5a584f72b)의 같은 shape가 14.38us였는데, 위의 packed 변환·KV split·PDL 도입 후 7.76us로 거의 절반이 되었습니다. 소배치에서 SM을 채우지 못하던 문제를 KV split이 해결한 효과가 가장 큽니다.

작성자는 공정성 측면도 명확히 밝힙니다. FP8 MMA peak가 BF16의 2배이므로 FlashMLA 대비 이득의 일부는 정밀도 계약 차이에서 옵니다. 또한 TRTLLM-GEN 대비로는 K=128, B<=32 등에서 최대 8% 느리며, split 비용 테이블이 B200(148 SM)에 맞춰져 R100(216 SM)의 B=32/K=512가 B=64보다 느린 등 미튜닝 구간도 솔직하게 공개했습니다.

일반적 교훈

1. 저장 포맷에서 바로 계산하라. 전체 캐시를 BF16으로 dequantize하지 않고 선택된 tile만 staging하며 E4M3로 변환하면 메모리 트래픽과 연산 정밀도 비용을 동시에 줄일 수 있습니다.
2. 하드웨어 신규 명령어를 적극 활용하되 fallback을 유지하라. PTX 9.4 packed multiply로 명령어 수를 40% 줄이면서도, 지원하지 않는 SM107에는 F16 경로를 두었습니다.
3. 소배치 decode는 parallelism 문제다. batch가 작으면 KV 축으로 split하고 PDL로 reducer 지연을 숨기는 것이 유효합니다.
4. 정밀도가 다르면 oracle도 달라야 한다. 독립적인 FP64 represented-input oracle과 componentwise bound는 저정밀 커널의 정확성을 정량적으로 보증합니다.
5. 실험적 기능은 명시적 opt-in으로. backend="prims-ts" 없이 Prims-TS 전용 키워드를 넘기면 에러를 내 기존 시그니처로 새지 않도록 막은 점도 좋은 API 위생입니다.

참고 자료

• CUDA Graphs - CUDA C++ Programming Guide — 벤치마크·API가 전제하는 CUDA Graph capture/replay
• PTX ISA — FP4→E4M3 변환에 사용한 packed mul 및 cp.async 명령어 문서
• torch.cuda.CUDAGraph — 벤치마크의 graph replay 타이밍
• FlashMLA — 비교 대상 upstream sparse MLA decode 구현
• CUTLASS CuTe DSL — Prims-TS 커널이 작성된 CuTe-DSL

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