[flashinfer] GB300 NVFP4 Context Attention: scale slot 조기 반납과 8단 packed K/V ring으로 최대 1.26x 가속

PR 링크: flashinfer-ai/flashinfer#9177 상태: Merged | 변경: +8169 / -8546

들어가며

FlashInfer의 Cake NVFP4 paged-context attention은 FP8(E4M3) query/output과 NVFP4 K/V page(E4M3 block scale 포함)를 입력으로 받아, shared memory에서 page를 E4M3로 dequantize한 뒤 attention을 수행합니다. head dim 128, bottom-right causal mask, HND page layout을 사용하며 sm_100a(B200)와 sm_103a(GB300)에 각각 11개씩 총 22개의 AOT program이 배포됩니다.

문제는 sm_103a에서의 파이프라인 스케줄이었습니다.

1. softmax stage가 scale slot을 너무 늦게 돌려받음: 각 KV block에서 softmax warp가 scale slot을 재사용하려면, correction(rescale) warp가 output rescale과 다음 K/V tile의 dequantization까지 끝낼 때까지 기다려야 했습니다. 실제로는 scale 값을 읽기만 하면 slot은 비워도 되는데도 말이죠.
2. packed NVFP4 K/V ring이 2단뿐: TMA load가 dequantization보다 최대 한 block만 앞서 갈 수 있었습니다. long-KV 시나리오에서는 메모리 latency를 숨길 여유가 거의 없습니다.

이 PR은 sm_103a의 11개 program만 변경해 두 문제를 동시에 해결합니다. 산술 연산은 바뀌지 않으며, 모든 측정 row와 layout에서 출력 O는 이전 program 및 JIT 모듈과 bitwise 동일합니다.

코드 분석

1. Binding: dynamic shared memory 확대

8단 ring을 담기 위해 CTA당 dynamic shared memory가 늘었습니다.

Before

CheckCuda(cudaFuncSetAttribute(reinterpret_cast<const void*>(kernel_..._8fab5c6a2f2e50fd6e51),
    cudaFuncAttributeMaxDynamicSharedMemorySize, 104960), "cudaFuncSetAttribute");
...
config.dynamicSmemBytes = 104960;

After

CheckCuda(cudaFuncSetAttribute(reinterpret_cast<const void*>(kernel_..._099089a6787f767da51e),
    cudaFuncAttributeMaxDynamicSharedMemorySize, 160256), "cudaFuncSetAttribute");
...
config.dynamicSmemBytes = 160256;

104,960 B → 160,256 B로 약 55 KB 증가했습니다. 48 KB를 넘는 dynamic smem은 cudaFuncSetAttribute로 opt-in이 필요하며, 이 값은 device별로 한 번만 설정되도록 prepared_devices 비트마스크로 캐싱됩니다. __launch_bounds__(512, 1)로 SM당 CTA 1개를 가정하는 커널이므로 occupancy 손실 없이 smem을 파이프라인 깊이에 투자할 수 있는 구조입니다. 커널 이름의 해시 suffix가 바뀐 것은 Cake 생성기가 스케줄 변경을 별도 artifact로 식별하기 때문입니다.

2. Kernel: mbarrier 레이아웃 재배치

pk pipeline의 barrier가 2개씩에서 8개씩으로 늘면서 이후 barrier 주소가 모두 밀렸습니다.

Before

#define pk_full_addr (mbar_base + 64)
#define pk_empty_addr (mbar_base + 80)
#define s_full_addr (mbar_base + 96)
...
#define work_id_empty_addr (mbar_base + 320)

After

#define pk_full_addr (mbar_base + 64)
#define pk_empty_addr (mbar_base + 128)
#define s_full_addr (mbar_base + 192)
...
#define work_id_empty_addr (mbar_base + 416)

pk_full이 64128(8×8 B), pk_empty가 128192(8×8 B)를 차지합니다. 총 barrier 수는 41개(328 B)에서 53개(424 B)가 되었습니다. pk_full은 init_count=1(TMA producer), pk_empty는 init_count=4(4개 consumer warp가 arrive)로, 전형적인 multi-stage producer/consumer ring입니다.

3. Kernel: lane별 mbarrier 초기화

가장 흥미로운 부분은 초기화 방식입니다.

Before

#include <cuda_awbarrier_primitives.h>

__device__ __forceinline__ void mbarrier_init_owner_lane(void* mbar_ptr, uint32_t count) {
    __mbarrier_init(reinterpret_cast<__mbarrier_t*>(mbar_ptr), count);
}
...
if (tid == 0) {
    mbarrier_init_owner_lane(smem_raw + 0, 1);
    mbarrier_init_owner_lane(smem_raw + 8, 1);
    mbarrier_init_owner_lane(smem_raw + 16, 4);
    ... // 41개를 thread 0이 순차 초기화

After

__device__ __forceinline__ void mbarrier_init(int mbar_addr, int count) {
    asm volatile("mbarrier.init.shared::cta.b64 [%0], %1;"
        :: "r"(mbar_addr), "r"(count) : "memory");
}
...
if (warp == 0) {
    // pk_full: 8 barriers, init_count=1
    // pk_empty: 8 barriers, init_count=4
    ... // 각 barrier를 warp 0의 서로 다른 lane이 초기화

기존에는 tid == 0 한 스레드가 모든 barrier를 순차적으로 init했습니다. ring을 8단으로 늘리면 barrier가 12개 늘어나 prologue가 그만큼 길어집니다. 대부분의 시간이 prologue인 one-tile launch(예: 1 x 1 x 16)에서는 이 비용이 그대로 드러납니다. PR은 warp 0의 각 lane이 자기 담당 barrier를 init하도록 바꿔, 53개 init이 사실상 병렬로 처리되게 했습니다. 또한 cuda_awbarrier_primitives.h의 __mbarrier_init 대신 shared-window 32비트 주소를 받는 inline PTX mbarrier.init.shared::cta.b64를 직접 사용해 generic→shared 주소 변환도 생략합니다. PR 설명의 "short one-tile launches do not pay for the deeper ring"이 바로 이 부분입니다.

4. Scale slot 조기 release (스케줄 변경)

diff가 잘려 해당 부분 코드는 보이지 않지만, PR 설명에 따르면 correction warp는 이제 scale을 읽는 즉시 해당 stage의 scale slot을 release합니다. 이전에는 release 시점이 "output rescale + 다음 K/V tile dequant 완료" 이후였기 때문에, softmax warp가 다음 block의 scale을 쓰려면 correction warp의 무거운 작업 전체를 기다려야 했습니다. 이제 이 의존성이 "scale read"라는 짧은 구간으로 축소되어 softmax와 correction이 실질적으로 overlap됩니다.

왜 이게 좋은가

GB300(152 SMs, CUDA 13.4, cold L2, 10개 memory layout median) 기준:

• 이전 Cake 대비 1.009x ~ 1.263x, 28개 row 전부 빨라졌고 느려진 row는 없음
• JIT 모듈 대비 1.457x ~ 3.812x
• 최대 이득은 query가 적고 KV가 긴 경우: 2 x 1 x 4096 56.2 → 44.5 µs(1.263x), 1 x 32 x 32768 403.1 → 329.8 µs(1.222x)
• 최소 이득은 one-tile launch: 1 x 1 x 16 1.010x, 1 x 16 x 64 1.009x — 회귀 없음이 핵심
• trtllm-gen 대비 22개 row에서 최대 5.158x 우세, 6개 row에서는 여전히 0.842~0.982

패턴이 명확합니다. query token이 적으면 CTA당 연산량은 작고 K/V 스트리밍이 지배적입니다. 이때 2단 ring은 load latency를 숨기지 못하고, 8단 ring이 가장 큰 효과를 냅니다. 반대로 tile이 하나뿐인 launch에서는 ring이 깊어도 쓸 일이 없는데, lane-parallel 초기화 덕분에 추가 비용도 거의 없습니다.

일반적 교훈

1. Release는 "마지막 사용" 직후에: 버퍼 slot의 반납 시점을 논리적 작업 단위 끝이 아니라 실제 데이터의 마지막 read 직후로 당기면, 산술 변경 없이 warp-specialized 파이프라인의 직렬화를 제거할 수 있습니다.
2. SM당 CTA 1개라면 smem은 파이프라인 깊이에 쓰라: occupancy를 이미 1로 고정한 persistent 커널에서 남는 shared memory는 stage 수로 바꾸는 것이 가장 직접적인 latency hiding 수단입니다.
3. 깊은 파이프라인의 고정비를 의식하라: stage가 늘면 barrier init 같은 prologue 비용도 늘어납니다. 이를 lane 병렬로 분산해 짧은 workload의 회귀를 막는 것이 균형 잡힌 최적화입니다.
4. bitwise 동일성 검증: 스케줄만 바꾸는 최적화는 출력이 bit 단위로 같아야 한다는 강력한 불변식으로 검증할 수 있고, 이 PR은 모든 row/layout에서 이를 확인했습니다.

리뷰 과정에서는 메인테이너(yyihuang)가 tests/attention/test_cake_fmha.py와 test_cake_fmha_context_programs.py CI 실행을 요청해 기능 정합성을 확인했으며, sm_100a program은 스케줄·이름·바이트가 그대로 유지되어 변경 범위가 sm_103a로 엄격히 제한되었습니다.

참고 자료

• cudaFuncSetAttribute — dynamic shared memory 상한 설정 API (CUDA Runtime Execution Control)
• PTX mbarrier — mbarrier.init 등 PTX barrier 명령어 공식 문서
• CUDA C++ Programming Guide — shared memory, asynchronous barrier 및 TMA 개념

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