[flashinfer] 4-stage K/V ring에서 불필요해진 QK(t)→PV(t-1) 대기 제거: SM107 FP8 sparse MLA decode 최대 11% 가속

PR 링크: flashinfer-ai/flashinfer#8828 상태: Merged | 변경: +124 / -159

들어가며

FlashInfer의 trtllm_batch_decode_sparse_mla_dsv4는 DeepSeek-V4의 sparse MLA decode를 담당합니다. 이 커널의 Cake 백엔드는 50개의 program으로 구성되어 있습니다. BF16·FP8 program이 query head 수(8/16/32/64/128)별로 여러 body를 두고 shape에 따라 host route가 하나를 고르며, 여기에 NVFP4-KV decode family가 더해집니다. sm_107a에서는 이 중 FP8 decode-route program 7개만 배포되는데, 그중 하나가 #8818에서 추가된 64-head M64 multi-tile program의 four-stage K/V ring 버전입니다.

이 program은 single-wave launch이면서 K/V tile이 4개 이상인 work item을 처리할 때, tile t의 QK MMA를 tile t-1의 PV MMA가 issue될 때까지 기다리게 했습니다. 이 순서 제약은 2-stage ring에서는 의미가 있었습니다. PV(t-1)가 먼저 나가야 K/V stage가 일찍 비고, 그래야 다음 tile을 로드할 자리가 생기기 때문입니다.

하지만 stage가 4개라면 ring에 이미 여유 슬롯이 있습니다. 이때 같은 제약은 stage를 일찍 비워 주는 이득은 없고 QK → softmax → PV를 직렬화하는 비용만 남깁니다. 이 PR은 four-stage program에서 이 제약을 제거합니다. 인자, host route, 출력 bit는 그대로 유지됩니다.

코드 분석

변경 파일은 csrc/cake_dsv4/sm_107a/cake_dsv4_8afa079cbcc688b8b5b9_kernel.cu 하나입니다. 코드 생성기가 만든 커널이라 diff의 상당 부분은 변수 번호가 바뀐 것(_elected_3 → _elected_2, _tcgen05_cp_desc_4 → _tcgen05_cp_desc_3 등)이고, 의미 있는 변경은 아래 세 군데입니다.

1. single-wave 판정 로직 제거

Before:

int pv_qk_cursor = 0;
int pv_qk_single_wave = 0;
uint32_t _read_special_reg_0;
asm("mov.u32 %0, %%nsmid;" : "=r"(_read_special_reg_0));
if (num_query_tokens <= (int)_read_special_reg_0) {
    pv_qk_single_wave = 1;
}

After: 이 블록 전체가 삭제되었습니다.

%nsmid special register로 SM 수를 읽고, num_query_tokens가 그 이하이면 single-wave로 판정해 pv_qk_single_wave 플래그를 켜던 코드입니다. pv_qk_cursor는 PV가 몇 번째 tile까지 issue되었는지를 추적해, QK(t)가 PV(t-1)를 기다리게 하는 데 쓰였습니다. 이 순서 제약이 사라지면서 판정 로직도 함께 필요 없어졌습니다. diff가 잘려 있어 QK issue 루프 안의 대기 지점까지는 보이지 않지만, 제약을 걸던 상태 변수가 모두 사라졌다는 점이 변경의 핵심을 보여 줍니다.

2. pv_pair_ready barrier의 arrival count 변경

Before:

// pv_pair_ready: 2 barriers, init_count=1
asm volatile("... selp.u32 %0, %3, %0, p; }" : "+r"(_mbarrier_init_count_0_32) : "r"(lane), "n"(15), "r"((uint32_t)(1)));
asm volatile("... selp.u32 %0, %3, %0, p; }" : "+r"(_mbarrier_init_count_0_32) : "r"(lane), "n"(13), "r"((uint32_t)(64)));

After:

// pv_pair_ready: 2 barriers, init_count=64
asm volatile("... selp.u32 %0, %3, %0, p; }" : "+r"(_mbarrier_init_count_0_32) : "r"(lane), "n"(15), "r"((uint32_t)(64)));

이 코드는 lane 번호로 각 lane이 초기화할 mbarrier의 init count를 고르는 방식입니다. setp.lt.u32와 selp를 이어 붙여 lane별 값을 덮어씁니다. 이전에는 lane 1314가 64, 15가 1이었고, pv_pair_ready 두 개 중 하나의 count가 1이었습니다. 이 count 1짜리 barrier는 MMA 쪽 단일 스레드가 "PV issue됨"을 알리는 signal 용도였던 것으로 보입니다. 이제 두 barrier 모두 64(warp pair 64 threads의 arrival)로 통일되어, selp 한 줄로 1315 lane을 함께 처리합니다. 순서 강제용 signal이 빠지면서 barrier 구성도 단순해진 것입니다.

3. 이름 재번호 (의미 변화 없음)

- int _max_1_4 = ((sparse_topk_lens[mapped_query_idx_4] + sparse_topk_lens_offset) > (0) ? ...
+ int _max_0_4 = ((sparse_topk_lens[mapped_query_idx_4] + sparse_topk_lens_offset) > (0) ? ...

active_topk을 [1, sparse_topk]로 clamp한 뒤 128 단위 tile 수(all_num_kv_tiles)를 계산하는 로직은 그대로입니다. 생성기가 임시 변수를 하나 덜 쓰게 되면서 번호만 밀린 것입니다. 이런 diff를 리뷰할 때는 재번호와 의미 변경을 분리해서 보는 것이 중요합니다.

왜 이게 좋은가

성능 수치

측정 환경은 NVIDIA R200(208 SMs), sm_107a, CUDA 13.4이며 보드 2장에서 측정했습니다. 한 CUDA graph 안에서 PDL로 8회 연속 호출하고, graph launch마다 L2를 한 번 flush했습니다. 비교 대상은 #8818의 four-stage program입니다.

shape	rows	speedup
top-k 4x (640), 64 heads, q_len 4	120	1.1045 / 1.1141
top-k 4x (640), 64 heads, q_len 4	132	1.1071 / 1.1029
top-k 4x (640), 64 heads, q_len 4	192	1.0917 / 1.0907
top-k 4x (1152), 64 heads, q_len 1	192	1.0488 / 1.0486

가장 큰 이득은 rows가 SM 수(208)보다 적은 single-wave shape에서 나왔습니다. 예를 들어 10.705us → 9.692us입니다. 이 경우 SM마다 work item이 하나뿐이라 inter-CTA overlap이 없고, 성능이 CTA 내부 파이프라인의 직렬화에 그대로 노출됩니다. 제거된 제약이 정확히 single-wave에서만 걸려 있었다는 점과도 들어맞습니다.

rows 256 이상의 multi-wave shape에서는 원래 제약이 적용되지 않았으므로 1.001.03x 수준입니다. 2048-row prefill(32/16 heads)은 0.99550.9979x로 약간 느려졌지만 노이즈 수준이고, 이 PR이 건드리지 않은 SWA-only program도 0.998~1.000x로 측정 변동 범위를 보여 줍니다.

일반적 교훈

1. 파이프라인 제약은 버퍼 깊이에 종속된다. "QK(t)가 PV(t-1)를 기다린다"는 2-stage ring에서는 deadlock과 stall을 피하는 장치였지만, 4-stage에서는 슬롯 여유가 충분해 순수한 직렬화 비용이 됩니다. stage 수를 늘린 variant를 만들 때는 기존 동기화가 여전히 필요한지 다시 따져 봐야 합니다.
2. single-wave는 latency-bound다. wave가 하나뿐이면 다른 CTA가 stall을 가려 주지 못합니다. CTA 내부의 tensor core(QK/PV)와 CUDA core(softmax)를 겹치는 것이 곧 성능입니다.
3. 동기화 제거는 barrier 단순화로 이어진다. 순서용 signal(count 1)이 사라지면 mbarrier 초기화와 상태 변수(pv_qk_cursor, %nsmid 조회)도 함께 줄어들어, 레지스터와 명령어 오버헤드가 감소합니다.
4. 출력 bit를 유지하는 변경이라 검증 부담이 낮다. 리뷰어(yyihuang)는 test_cake_dsv4.py, test_cake_dsv4_out_placement.py, test_cake_dsv4_hardening.py를 bot으로 실행하도록 요청했습니다. 수치 결과가 바뀌지 않는 스케줄링 변경이므로, 정확성과 배치(out placement) 회귀 테스트만으로 안전성을 확인할 수 있는 구조입니다.

참고 자료

• PTX ISA: mbarrier — 커널이 사용하는 mbarrier init/arrive/wait 명령
• PTX ISA: Special Registers %nsmid — single-wave 판정에 쓰이다 제거된 special register
• PTX ISA: tcgen05 instructions — Blackwell tcgen05.cp / MMA 명령
• CUDA Programming Guide: Programmatic Dependent Launch — 벤치마크에 사용된 PDL

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