[flashinfer] packed qkv에서 q/k를 복사 없이 바로 읽기: Cake GDN 결정적 학습 커널의 TMA stride 활용 최적화

PR 링크: flashinfer-ai/flashinfer#9103 상태: Merged | 변경: +2727 / -2758

들어가며

FlashInfer의 실험적 패키지 flashinfer.cake_gdn_det_train은 Gated DeltaNet(GDN)의 chunked training을 결정적(deterministic)으로 수행하는 single-GPU 커널 패키지입니다. Cake GDN 학습 패밀리는 deterministic/non-deterministic 두 종류가 있고, 각각 single-GPU 형태와 context-parallel 형태(_kcp)를 가집니다. 이 PR은 그중 deterministic single-GPU 패키지만 다룹니다.

실제 모델 코드에서 q, k, v는 보통 하나의 projection 결과를 쪼개서 만듭니다. 즉 [T, (H + H + HV) * D] 같은 packed qkv row buffer 하나가 있고, q와 k는 그 버퍼 위의 head view입니다. 이런 view는 마지막 차원은 연속이지만 row stride가 packed row 전체 폭이므로 contiguous하지 않습니다.

기존 구현에서는 TMA prep 단계에 들어가기 전에 q와 k를 contiguous tensor로 복사했습니다. 긴 시퀀스(T≈262k)에서는 이 복사가 global memory를 한 번 더 읽고 쓰는 비용이 되고, forward 시간의 상당 부분을 차지했습니다. 이 PR은 q와 k가 packed qkv의 head view일 때 복사 없이 view의 stride를 그대로 담은 TMA tensor map으로 읽도록 바꿉니다.

코드 분석

1. 생성된 커널 소스의 변경 (build.json)

이 패키지는 Cake라는 상위 코드 생성기에서 렌더링되는 AOT 패키지입니다. 그래서 diff 대부분은 사람이 직접 쓴 코드가 아니라 생성기 revision과 산출물 해시의 갱신으로 나타납니다.

   "generator": {
-    "cake_revision": "ffade710ba0f2de1f73a1658a155c9564b6394b5",
+    "cake_revision": "92886786a6509200be1f0c67eae82423f35c8787",
     "cuda": "13.4.92",
-    "rendered_at": "2026-11-05T05:32:54Z"
+    "rendered_at": "2026-11-07T13:40:55Z"
   },

주목할 점은 바뀐 파일의 종류입니다. 각 variant마다 _binding.cu의 해시는 그대로이고 _kernel.cu의 해시만 바뀌었습니다.

       "sha256": {
         "csrc/cake_gdn_det_train_0378c622e8993662f832_binding.cu": "b2ceb2a8bfff...",
-        "csrc/cake_gdn_det_train_0378c622e8993662f832_kernel.cu": "b7848a458764..."
+        "csrc/cake_gdn_det_train_0378c622e8993662f832_kernel.cu": "1f8c3a24fefd..."
       },

_wait.cu가 있는 variant도 마찬가지입니다.

-        "csrc/cake_gdn_det_train_0b4d86ff7e9093ca0ca9_kernel.cu": "45ab8884fa44...",
+        "csrc/cake_gdn_det_train_0b4d86ff7e9093ca0ca9_kernel.cu": "dad4cc41782a...",
         "csrc/cake_gdn_det_train_0b4d86ff7e9093ca0ca9_wait.cu": "29fd232fe88b..."

즉 Python ↔ C++ 경계(binding)의 ABI는 그대로 두고, 커널 내부 구현만 바뀌었다는 뜻입니다. 사용자 API나 호출 시그니처에는 영향이 없으므로 기존 호출 코드가 그대로 이득을 봅니다. sm_100a / sm_103a / sm_107a 세 아키텍처 각각의 closure_sha256도 일괄 갱신되어, 모든 Blackwell 계열 타깃이 같은 Cake revision에서 재생성되었음을 보장합니다.

2. 커널 측 변경의 핵심: stride-aware TMA tensor map

(커널 소스 diff는 잘려 있어 원문 인용이 불가능하므로, PR 설명을 바탕으로 개념만 정리합니다.)

• Before: q, k가 어떤 layout이든 TMA prep 전에 contiguous 사본을 만들고, 사본 위에 tensor map을 생성.
• After: q, k가 동일한 packed qkv row buffer의 head view이면, view의 base pointer와 실제 row stride(packed row 폭)를 그대로 tensor map에 인코딩해 원본에서 직접 TMA load. 그 외 layout은 기존처럼 복사 경로 유지.

TMA tensor map(cuTensorMapEncodeTiled)은 원래 각 차원의 global stride를 바이트 단위로 받도록 설계되어 있습니다. 차원 0만 연속이고 나머지 stride는 16바이트 정렬만 맞으면 되므로, head view처럼 row 사이에 다른 head 데이터가 끼어 있는 layout도 하드웨어가 그대로 타일 단위로 가져올 수 있습니다. 굳이 소프트웨어로 gather/copy할 이유가 없었던 셈입니다.

CodeRabbit 요약에 언급된 "large tensor dimension을 위한 address calculation 개선"도 같은 맥락으로 볼 수 있습니다. packed buffer의 row stride는 contiguous q보다 수 배 크기 때문에, T가 수십만인 경우 오프셋 계산이 32-bit 범위를 넘을 위험이 커집니다. stride를 원본 그대로 쓰려면 주소 계산을 더 넓은 정수형으로 처리해야 합니다.

3. 부수 변경: cluster scope barrier

재생성된 소스에서는 backward tail의 receive barrier가 cluster scope로 arm됩니다. PR 설명대로 출력과 타이밍에는 변화가 없고, 메모리 모델상 의미를 정확히 맞춘 정합성 수정입니다.

왜 이게 좋은가

성능 수치 (GB300, sm_103a)

shape	forward (prev → PR)	fwd+bwd	ckpt step
H=2, HV=6, T=262,144 (64 seqs)	1.701 → 1.418 ms (1.20x)	1.08x	1.11x
H=16, HV=48, T=32,768	1.731 → 1.443 ms (1.20x)	1.08x	1.12x
H=1, HV=3, T=524,288	2.125 → 1.842 ms (1.15x)	1.06x	1.09x
H=4, HV=12, T=32,768	0.550 → 0.549 ms (1.00x)	1.00x	1.00x

• 12개 shape 중 9개에서 forward가 1.15x~1.20x, 긴 시퀀스에서 호출당 약 0.28 ms 절감.
• Backward는 0.998x~1.001x로 변화 없음. 복사가 forward 쪽 prep에만 있었다는 방증입니다.
• Checkpointed step(2x fwd + bwd)은 forward가 두 번 돌기 때문에 이득이 증폭되어 1.09x~1.12x.
• TMA prep이 선택되지 않는 shape(H=4, H=2 @ T=32,768)은 정확히 1.00x로, 변경이 다른 경로에 회귀를 일으키지 않음.
• 출력과 gradient는 bit-identical. 결정적 학습 패키지에서 이 점은 필수 조건입니다.

솔직하게 짚어둘 부분도 있습니다. cuDNN frontend FROST GDN과 비교하면 forward는 여전히 모든 shape에서 느리고(0.29x~0.49x), backward에서만 대부분 앞섭니다. 이번 PR은 경쟁 구현을 따라잡는 것이 아니라 자기 자신의 불필요한 비용을 제거한 것입니다.

일반적 교훈

1. .contiguous()는 공짜가 아니다. packed projection을 split한 view에 습관적으로 contiguous 사본을 만들면 대역폭 바운드 커널에서 수십 %의 오버헤드가 생깁니다.
2. 하드웨어가 stride를 지원하면 맡겨라. TMA, cp.async 등은 strided layout을 기본 지원합니다. layout 정규화 대신 descriptor에 stride를 정확히 기술하는 편이 낫습니다.
3. fast path + 안전한 fallback. 조건(같은 packed buffer의 head view)을 만족할 때만 in-place 경로를 쓰고 나머지는 기존 복사를 유지해 정확성과 커버리지를 함께 지켰습니다.
4. ABI를 건드리지 않는 최적화. binding 해시가 그대로인 것에서 보듯, 내부 구현만 바꿔 사용자 코드를 수정하지 않고 이득을 전달했습니다.

리뷰 과정에서는 메인테이너(yyihuang)가 test_cake_gdn_det_train.py, _api.py, _names.py 세 테스트 스위트를 CI 봇으로 반복 실행해, 생성 산출물 교체 후 정확성과 API/이름 규약이 유지되는지 검증했습니다.

참고 자료

• cuTensorMapEncodeTiled (CUDA Driver API, Tensor Map Object Management) — TMA tensor map에 global stride를 인코딩하는 API
• torch.Tensor.contiguous — 기존 경로에서 사용되던 contiguous 사본 생성
• torch.Tensor.stride — head view의 stride 확인에 사용되는 API

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