[flashinfer] Cake 결정적 MoE Grouped GEMM: FP32 wgrad 누적을 커널에 fuse해 scratch를 없애고 최대 1.29x 가속

PR 링크: flashinfer-ai/flashinfer#9132 상태: Merged | 변경: +33562 / -1007

들어가며

flashinfer.cake_moe_grouped_gemm_det_train은 학습용 ragged BF16 MoE grouped GEMM으로, 실행할 때마다 비트 단위로 같은 결과(run-to-run deterministic)를 보장합니다. Megatron 스타일 학습에서는 weight gradient를 w.grad에 따로 만들지 않고 FP32 main_grad에 바로 더합니다(gradient-accumulation fusion). 문제는 이 누적 경로가 지금까지 composed 방식이었다는 점입니다.

1. 결정적 overwrite 커널이 dW를 out과 같은 dtype의 연속 scratch에 씁니다.
2. 같은 stream에서 out.add_(scratch)를 실행합니다.

S = E * N * K라고 하면 FP32 scratch는 4 * S 바이트를 차지합니다. gate/up(N 4096, K 6144)에서 E=32이면 3 GiB입니다. 여기에 scratch를 쓰고 다시 읽는 8 * S 바이트의 추가 메모리 트래픽이 붙습니다. 큰 E일수록 GEMM보다 이 부수 비용이 커집니다.

이 PR은 세 가지를 바꿉니다.

1. FP32 wgrad accumulate를 커널 안으로 fuse합니다. 결과는 out.add_(dW)와 비트 단위로 같습니다.
2. wide forward와 dgrad 프로그램에 L2 evict-first 변형을 추가합니다.
3. sm_100a/sm_103a에서 512-wide wgrad k tile을 쓰기 시작하는 기준을 그룹당 4096행에서 3072행으로 낮춥니다.

코드 분석

참고: 제공된 diff는 커널 소스 부분이 잘려 있어, 아래 인용은 API 문서와 docstring 변경 위주입니다. 동작 계약이 어떻게 바뀌었는지는 이 부분에 명확히 드러납니다.

1. docs/api/cake_moe_grouped_gemm_det_train.rst: accumulate 경로 이원화

Before: 경로는 composed 하나뿐이었습니다.

The launch composes (``plan["accumulate"] ==
"composed"``): the overwrite of the same plan (route, k tile, stage, tail split and its
ordered tail reduce) into a contiguous scratch of ``out``'s dtype, then ``out.add_`` on the
same stream, outside autograd.

After: fused와 composed 두 경로가 생겼습니다.

``"fused"``: for an fp32 ``out`` on the 512-wide k
tile's half-release forms, where the accumulate builds of the stage and, for a split tail,
of the tail reduce ship and can load, and ``out`` is 32-byte aligned, the kernels add into
``out`` in place (no scratch).  ``"composed"`` otherwise (``plan["accumulate_reason"]``):
...  Both routes give the same bits.

핵심은 fused 경로를 어떤 조건에서만 쓰는지입니다. 조건은 네 가지가 모두 맞아야 합니다.

• out이 FP32일 것
• 512-wide k tile의 half-release 형태일 것
• main stage와 (split tail이면) ordered tail reduce의 accumulate 빌드가 존재하고 로드 가능할 것
• out이 32-byte 정렬일 것

하나라도 어긋나면 composed로 돌아가고, 그 이유를 plan["accumulate_reason"]에 남깁니다. 성능 경로를 공격적으로 열면서도 정확성은 fallback으로 보장하는 설계입니다.

비용 문서도 이에 맞게 바뀌었습니다.

- Cost (``S = E * N * K``): the fused route takes no scratch.  The composed route's
  scratch takes ``4 * S`` bytes for fp32 ...

2. cake_moe_grouped_gemm_det_train.py: prepare 계약과 결정성 유지

prepare_wgrad docstring의 Before는 다음과 같습니다.

``plan["accumulate"]`` is ``"composed"``
(``"off"`` for an overwrite) and every other plan key is the overwrite's.

After는 다음과 같습니다.

``plan["accumulate"]`` is ``"fused"`` (the accumulate builds of the main
stage and, for a split tail, of the tail reduce add into ``out`` in place:
``stages`` is the overwrite's with ``ACC_SUFFIX`` on each) or
``"composed"`` (...)

주목할 점은 fused 플랜의 stages가 overwrite 플랜과 같은 route, k tile, stage, tail split을 쓰고 이름에 ACC_SUFFIX만 붙는다는 것입니다. reduction 순서가 overwrite와 완전히 같으니, 마지막 store 단계에서 out = out + dW_fp32를 한 번 수행하면 out.add_(dW)와 같은 비트가 나옵니다. 결정성의 원천인 "고정된 reduction 순서"는 그대로 두고 epilogue만 바꾼 것입니다.

빈 그룹 처리도 더 명확해졌습니다. -0.0 may become +0.0이 -0.0 becomes +0.0으로 바뀌어, 동작을 확정적으로 문서화합니다.

3. forward: L2 evict-first 변형

# Before
512-column tile when ``N % 512 == 0`` on the architectures of
``wide_fwd_archs``, its 3 / 3 head / tail twin where ``fwd_wide_ts33`` holds).
# After
... its L2 evict-first twin where ``fwd_lwh_ef`` or
``fwd_evict_first`` holds for ``K``, its 3 / 3 head / tail twin ...

한 번만 읽고 다시 쓰지 않는 스트리밍 operand에 evict-first 힌트를 주면, 재사용되는 타일이 L2에 더 오래 남습니다. 그래서 fwd gu와 dgrad gu에서 1~3% 수준의 개선이 나옵니다.

왜 이게 좋은가

수치

• GB200: 210개 shape 중 FI main 대비 98개 빨라졌고, 112개는 noise 범위, 느려진 것은 0개입니다. 최대 wgrad_acc32 dn E32 bal에서 1.207x(2.934 → 2.423 ms)입니다.
• GB300: 같은 shape에서 1.236x(2.566 → 2.075 ms)입니다.
• R200: 1.294x(1.519 → 1.173 ms)입니다.
• wgrad_acc32 gu E32 bal은 GB200에서 5.882 → 4.889 ms(1.202x)입니다.
• end-to-end step_acc32 gu E32도 3~6% 줄었습니다.

개선 폭이 E가 클수록 커지는 것은 예측과 맞습니다. scratch 트래픽 8 * S가 E에 비례해 늘기 때문입니다. 반면 wgrad_acc16(BF16 main_grad)은 fused 대상이 아니어서 그대로입니다. 이것은 변경 범위가 정확히 FP32 경로에 한정됐다는 증거이기도 합니다.

측정 방법론도 눈여겨볼 만합니다. 같은 프로세스에서 구현들을 interleave해 실행하고, cold L2와 CUPTI 시간을 씁니다. 같은 커널 두 벌을 비교한 A/A noise를 최소 검출 효과로 삼고, non-Cake 비교에는 95% CI를 적용했습니다. 그래서 "within noise" 판정을 믿을 수 있습니다.

리뷰에서 드러난 엣지 케이스

병합 뒤 CodeRabbit 지적을 반영한 follow-up(#9173)에서 다음이 고쳐졌습니다.

• JIT 툴킷 판정: nvcc가 없거나 FLASHINFER_DISABLE_JIT가 설정된 경우 torch.version.cuda 값만 보고 툴킷이 있다고 판단하던 문제를 고쳤습니다. 이제 accumulate 빌드를 불러올 수 없으면 composed로 안전하게 떨어집니다.
• 32-bit 주소 범위 검사: num_groups * stride(0) 대신 실제로 주소를 매기는 전체 범위 1 + Σ(size-1)*stride를 검사합니다.
• stale prebuilt .so: 인자 수가 19개에서 21개로 늘어 생기는 불일치는 공유 로더 템플릿에서 closure digest로 다루기로 했습니다.

fused 경로가 "조건부 최적화 + 투명한 fallback" 구조인 만큼, fallback 판정 로직의 견고함이 곧 정확성을 좌우한다는 점을 보여주는 사례입니다.

일반적 교훈

1. Epilogue fusion은 메모리 바운드 후처리를 없애는 가장 값싼 방법입니다. GEMM 출력을 다시 읽어 더하는 elementwise 연산은 거의 전부가 트래픽 비용입니다.
2. 결정성을 지키며 fuse하려면 reduction 순서를 그대로 두고 마지막 연산만 바꿔야 합니다. 그래야 "bitwise equal"을 주장할 수 있습니다.
3. 조건부 fast path에는 이유가 남는 fallback을 두세요. accumulate_reason 같은 진단 필드는 사용자와 테스트 모두에 도움이 됩니다.

참고 자료

• torch.Tensor.add_ — composed 경로에서 scratch를 main_grad에 더하던 in-place add
• torch.use_deterministic_algorithms — scratch NaN fill 동작과 관련된 결정성 설정
• CUDA Graphs in PyTorch — prepared launch의 graph capture/replay 동작
• CUDA C++ Programming Guide — L2 cache access policy(evict-first 등) 관련 공식 문서

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