[flashinfer] Rubin(sm_107a)에서 DRAM 피크 95%까지: UMMA ReplaySSM u-cache flush 커널의 uGPU 로컬라이제이션과 6가지 튜닝

PR 링크: flashinfer-ai/flashinfer#6791 상태: Merged | 변경: +8837 / -102

들어가며

FlashInfer의 Gated Delta Net(GDN) 디코딩에는 MTP(Multi-Token Prediction) 검증과 u-cache flush를 한 번에 처리하는 UMMA(tcgen05 / TMEM) 기반 ReplaySSM 커널 gated_delta_rule_mtp_ucache_flush_umma가 있습니다(#6274). B200(sm_100a)에 맞춰 튜닝된 커널인데, NVIDIA의 차세대 GPU인 Rubin(VR200, sm_107a)에서 그대로 돌리면 DRAM 피크의 52~57% 밖에 쓰지 못했습니다. 이 커널은 state를 읽고 쓰는 memory-bound 커널이라, 대역폭을 못 쓰면 그만큼 시간이 늘어납니다.

이 PR은 Rubin 전용 빌드(gdn_replay_mtp_umma_rubin.py)를 추가하고, Rubin의 두 die가 각자 자기 메모리만 읽게 하는 uGPU localization을 도입합니다. 결과는 B = 2048에서 B200 최고 기록 대비 2.45× / 2.55× / 2.41× (T = 1 / 4 / 8), DRAM 피크의 91 / 95 / 90% 입니다. Blackwell 커널과 결과는 bitwise로 같고, stochastic rounding도 같은 noise를 뽑습니다.

왜 Blackwell 커널이 Rubin에서 느렸나

PR 설명에 정리된 원인 6가지가 이 PR의 핵심입니다. 각각이 "하드웨어가 바뀌면 기존 가정이 깨진다"는 좋은 사례입니다.

1. SM당 CTA 2개 제한 — Rubin은 SM당 named barrier가 16개(B200은 64개)인데, 커널이 id를 5까지 사용했습니다. id를 1~3으로 줄여 SM당 4 CTA를 확보.
2. Instruction cache miss — 같은 SM에서 flush CTA와 verify CTA가 서로 다른 코드를 돌리는데, 합치면 I-cache에 안 들어갔습니다(flush 행이 있는 step이 10~17% 느림). 루프를 roll해서 코드 크기를 줄임.
3. CTA 간 유휴 시간 — CTA마다 TMEM 할당과 ~60개 barrier 초기화를 반복하느라 CTA 슬롯이 13% 비었습니다. T = 28에서 persistent kernel로 한 번만 set-up.
4. die 간 트래픽 — 일반 할당은 메모리를 두 die에 흩뿌려서 바이트의 절반이 die-to-die 링크를 건넙니다. uGPU로 각 die가 자기 메모리만 읽게 함.
5. 중간 batch에서 key head 분할 — B200의 heads-per-CTA 규칙은 key head 하나의 value head 4개를 CTA 2개로 나눴는데, 각 CTA가 prologue를 반복합니다. Rubin은 SM당 DRAM 바이트가 B200의 1.6배라 이 비용을 덜 숨깁니다.
6. 작은 batch는 latency-bound — row-order indirection, launch bound, uGPU fork/join 같은 고정비용이 지배.

코드 분석

benchmarks/bench_gdn_ucache_flush.py — Rubin arm 추가

기존 UMMA arm을 그대로 재사용하면서 Rubin 빌드를 고를 수 있게 했습니다.

Before

if arm in _UMMA_SR:
    from flashinfer.gdn_kernels.gdn_replay_mtp_umma import (
        gated_delta_rule_mtp_ucache_flush_umma,
    )

    fn = gated_delta_rule_mtp_ucache_flush_umma

After

if arm in _UMMA_SR:
    if arm.startswith("rubin"):
        from flashinfer.gdn_kernels.gdn_replay_mtp_umma_rubin import (
            gated_delta_rule_mtp_ucache_flush_umma_rubin as fn,
        )
    else:
        from flashinfer.gdn_kernels.gdn_replay_mtp_umma import (
            gated_delta_rule_mtp_ucache_flush_umma as fn,
        )

stochastic rounding 설정도 한 줄로 복제합니다. umma_ph5 → rubin_ph5처럼 접두사만 바꿔 같은 설정을 공유하므로 두 커널의 비교가 공정해집니다.

_UMMA_SR.update({"rubin" + k[4:]: v for k, v in list(_UMMA_SR.items())})

die-localized 할당: DieSplit (벤치마크의 --localized)

uGPU의 원리를 가장 잘 보여주는 코드는 벤치마크의 DieSplit 클래스입니다. 두 가지를 합니다.

(1) die별 SM으로 green context 만들기. cuDevSmResourceSplit에 LOCALITY_DOMAIN_ID 플래그를 주어 SM 자원을 die 단위로 나누고, 각 그룹으로 green context와 stream을 만듭니다.

p = drv.CU_DEV_SM_RESOURCE_GROUP_PARAMS()
p.smCount = nsm // self.ndom
p.coscheduledSmCount = 2
p.flags = int(flags.CU_DEV_SM_RESOURCE_GROUP_LOCALITY_DOMAIN_ID) | int(
    flags.CU_DEV_SM_RESOURCE_GROUP_BACKFILL
)
p.localityDomainId = i
...
ctx = self._ck(drv.cuGreenCtxCreate(desc, self.dev,
    drv.CUgreenCtxCreate_flags.CU_GREEN_CTX_DEFAULT_STREAM))

BACKFILL은 SM 수가 딱 나눠떨어지지 않아도 모든 SM을 쓰게 해줍니다.

(2) 특정 die 메모리에 할당. CUDA VMM API(cuMemCreate / cuMemAddressReserve / cuMemMap / cuMemSetAccess)로 location type을 DEVICE_LOCALITY_DOMAIN으로 지정합니다.

prop.location.type = (
    drv.CUmemLocationType.CU_MEM_LOCATION_TYPE_DEVICE_LOCALITY_DOMAIN
)
prop.location.localized.deviceId = int(self.dev)
prop.location.localized.localityDomainId = dom

흥미로운 디테일은 copy 직후 검증 루프입니다. 주석에 따르면 새 매핑이 copy-engine fill 직후 SM에 stale line을 줄 수 있어, 읽어서 일치할 때까지 최대 4번 재복사합니다. 신규 하드웨어/드라이버 API를 다룰 때의 방어적 코딩입니다.

커널의 uGPU 버퍼: ugpu_ucache_buffers

실제 API에서는 사용자가 batch를 나눌 필요 없이, ugpu_ucache_buffers로 state pool과 ring을 할당하기만 하면 됩니다. 레이아웃은 [pool, 2, heads / 2, ...]이고, [s, t, j, ...]는 기존 버퍼의 [s, t * heads/2 + j, ...]에 해당합니다. 벤치마크에서도 이 관계로 복사합니다.

bufs = ugpu_ucache_buffers(B, H, HV, u_dtype=uc.dtype, device=DEV)
for dst, src in zip(bufs, (pool, kc, uc, gc), strict=True):
    src5 = src.view(src.shape[0], 2, src.shape[1] // 2, *src.shape[2:])
    dst.copy_(src5)

즉 batch가 아니라 key head를 반으로 나눠 die 0이 앞 절반, die 1이 뒤 절반을 처리합니다. 리뷰어 노트에 이유가 있습니다. #6604처럼 V 행을 나누면 tensor-core 타일이 128 V 행 전체를 덮으므로 M = 64 타일이 필요하지만, key head는 서로 독립이라 head-slice view와 stochastic rounding 카운터의 head offset만 주면 되고, 결과도 bitwise로 같습니다.

dispatcher: gdn_decode_bf16_wy_ucache_flush.py

backend="auto"가 SM107이거나 5-D uGPU 버퍼가 들어오면 Rubin 커널로 보냅니다. 여기서 리뷰가 중요한 버그를 잡았습니다.

• 리뷰어 kahyunnam: uGPU 버퍼의 early return이 hist_len / cache_base 검증을 건너뛰어, hist_len이 20이나 −3이어도 조용히 실행되고 commit됨. → 검증을 hand-off 이전으로 옮기고 test_dispatcher_ugpu_cursor_checks 추가.
• _umma_unavailable_reason이 major == 10만 확인해 sm_103에서 5-D 버퍼가 엉뚱한 NotImplementedError로 떨어짐. → sm_100 / sm_107 외에는 문서화된 ValueError.

fp32 gate 파라미터 캐시와 CUDA Graph

가장 날카로운 리뷰는 CUDA graph 안전성입니다. bf16 A_log / dt_bias의 버전이 바뀌면(copy_, load_state_dict) 다음 eager 호출이 새 fp32 사본을 할당하고 옛것을 해제하는데, 이전에 capture된 graph는 해제된 주소를 계속 읽습니다. 그 블록이 재사용되면 replay 결과가 최대 0.18 차이 났습니다. 수정은 "주소를 유지하고 in-place로 갱신"입니다: 버전 변경 시 device sync → copy → 새 event → 다른 stream이 다시 wait, shape가 바뀔 때만 새 버퍼. 리뷰어는 추가로 이 갱신이 "다음 eager 호출 이후"에만 일어나므로, 가중치를 리로드하고 graph replay만 하는 배포에선 옛 값을 읽는다는 점을 문서에 명시하라고 지적했습니다. graph capture가 일반화된 서빙 환경에서 반드시 기억할 포인트입니다.

왜 이게 좋은가

성능

B	T = 1	T = 4	T = 8
32	7.3 µs · 56% · 1.72×	8.3 µs · 51% · 1.70×	9.1 µs · 46% · 1.72×
512	77.8 µs · 86% · 2.34×	81.1 µs · 88% · 2.42×	90.0 µs · 82% · 2.27×
2,048	292.2 µs · 91% · 2.45×	302.2 µs · 95% · 2.55×	331.5 µs · 90% · 2.41×
8,192	1,149 µs · 93% · 2.48×	1,188 µs · 97% · 2.60×	1,295 µs · 92% · 2.74×

Rubin의 DRAM 대역폭이 B200의 2.5배이므로, 큰 batch에서 2.42.7배는 하드웨어 대역폭 증가를 거의 그대로 속도로 바꿨다는 뜻입니다. heads-per-CTA 규칙(fix 5)만으로 B = 64192에서 718%, small-batch dispatch(fix 6)로 B = 1에서 4.67 → 4.00 µs를 얻었습니다.

일반적 교훈

1. 아키텍처 상수는 커널의 숨은 가정이다. named barrier가 64 → 16개로 줄자 occupancy가 반토막났습니다. barrier id, smem, register 같은 per-SM 자원이 바뀌면 occupancy를 다시 계산해야 합니다.
2. multi-die GPU에선 메모리 배치가 성능이다. NUMA처럼 데이터와 연산을 같은 die에 두는 것이 대역폭 포화의 전제입니다. 배치는 할당 시점에 정해지므로 API 설계상 버퍼 할당 함수를 노출한 것도 합리적입니다.
3. I-cache도 자원이다. 서로 다른 코드 경로가 한 SM에서 섞이면 unroll이 오히려 손해일 수 있습니다.
4. 고정비용은 persistent kernel로 상각한다. 무거운 set-up(TMEM, barrier)을 가진 커널일수록 효과가 큽니다.
5. 튜닝 로그는 코드가 아니라 PR에. 리뷰어 지적대로 ~95줄의 측정 로그가 이미 서로 모순되고 있었고, 상수별 "이유 한 줄"만 남기고 측정치는 PR 본문으로 옮겼습니다.

덧붙여, die 간 launch 지연(2~7 µs)이 공유 hardware work queue의 parity 충돌 때문이라는 분석과, 해결책이 undocumented driver 동작에 의존하므로 PR에서 제외한 판단도 신중한 엔지니어링의 좋은 예입니다.

참고 자료

• CUDA Driver API: Green Contexts — cuGreenCtxCreate, cuDevSmResourceSplit 등 die별 SM 분할에 사용한 API
• CUDA Driver API: Virtual Memory Management — cuMemCreate / cuMemMap / cuMemSetAccess로 locality domain 메모리 할당
• torch.cuda.ExternalStream — green context stream을 PyTorch stream으로 감싸는 데 사용
• torch.cuda.CUDAGraph — 벤치마크의 graph replay 및 gate 캐시 안정성 이슈 관련
• CUDA C++ Programming Guide — named barrier, occupancy 등 커널 자원 개념

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