[triton] Blackwell Two-CTA MMA에 맞춘 CGA 레이아웃: M-first 분할로 멀티 CTA matmul 최대 9% 향상

PR 링크: triton-lang/triton#10817 상태: Merged | 변경: +146 / -16

들어가며

NVIDIA Blackwell(sm_100 계열)의 5세대 Tensor Core는 two-CTA MMA를 지원합니다. 같은 cluster 안의 인접한 두 CTA가 하나의 큰 MMA 타일을 함께 계산하는 방식입니다. 보통 두 CTA가 M 방향으로 나란히 서고, 각 CTA가 M의 절반씩(예: 128행)을 맡습니다. 이 구조의 장점을 살리려면 컴파일러가 CTA 간 텐서 분할(CGA layout)을 정할 때 다음 두 조건을 지켜야 합니다.

1. M 축을 CTA 사이에 나눠야 합니다.
2. 인접한 CTA ID(0과 1, 2와 3 …)가 M 축에서 짝이 되어야 합니다.

그런데 기존 AssignCGALayouts pass의 getDotCGASplit은 아키텍처와 무관한 휴리스틱을 썼습니다. 그 결과 2-CTA 테스트에서는 N 축만 분할했고(CGALayout = [[0, 1]]), 4-CTA에서는 첫 번째 CTA 비트가 N을 가리켰습니다([[0, 1], [1, 0]]). 하드웨어가 기대하는 "M 페어"와 컴파일러가 만든 레이아웃이 어긋나 있었던 것입니다. 이 PR은 Blackwell에서 M-first CTA 레이아웃을 우선 선택하도록 바꿉니다.

코드 분석

1. preferTwoCTALayout: 언제 적용할지 판단

bool preferTwoCTALayout(triton::DotOp dot) {
  if (ttg::lookupNumCTAs(dot) < 2)
    return false;
  auto module = dot->getParentOfType<ModuleOp>();
  int cc = getNVIDIAComputeCapability(module);
  return cc >= 100 && cc < 120;
}

CTA가 2개 이상이고 compute capability가 100~119일 때만 새 정책을 켭니다. 이는 two-CTA MMA를 지원하는 datacenter Blackwell 계열입니다. sm_120(consumer Blackwell)은 범위에서 빠지므로, 기존 동작이 유지되어 회귀 위험이 그 범위 안에 갇힙니다.

리뷰에서 lezcano가 "이 체크가 tritonnvidiagpu dialect에 꼭 필요한가"라고 물었습니다. 이에 Jokeren은 처음에 있던 CUDA target 확인 코드를 제거했습니다(7c7a349e1e). 근거는 두 가지입니다. 이 pass는 이미 NVIDIA 전용이라 CUDA target을 전제하고, getNVIDIAComputeCapability가 target을 자체 검증합니다. 같은 검사를 중복하지 않고 기존 유틸리티의 보장에 맡긴 것입니다.

2. getDotCGASplit: M 페어 + N 분할

Before:

SmallVector<unsigned> getDotCGASplit(ArrayRef<int64_t> shape,
                                     unsigned numCTAs) {
  ...
  unsigned splitM = 1;
  unsigned splitN = numCTAs;
  // Prefer a larger M chunk, up to 128 elements, by assigning splitM first.

After:

SmallVector<unsigned> getDotCGASplit(ArrayRef<int64_t> shape, unsigned numCTAs,
                                     bool preferTwoCTA) {
  ...
  if (preferTwoCTA && numCTAs >= 2) {
    unsigned splitM = std::max(2u, numCTAs / 2);
    unsigned splitN = numCTAs / splitM;
    if (m / splitM >= kMinChunkSize && m / splitM <= kPreferredChunkSize &&
        n / splitN >= kMinChunkSize && n / splitN <= 256) {
      ctaSplit[rank - 2] = splitM;
      ctaSplit[rank - 1] = splitN;
      return ctaSplit;
    }
  }

CTA 수별 분할은 다음과 같습니다.

• 2 CTAs: splitM=2, splitN=1 → 한 쌍이 M을 나눕니다.
• 4 CTAs: splitM=2, splitN=2 → M 페어 두 개가 N을 반씩 맡습니다.
• 8 CTAs: splitM=4, splitN=2

이 경로는 조건을 만족할 때만 탑니다. CTA당 M 청크가 kMinChunkSize(64) 이상 kPreferredChunkSize 이하이고, N 청크가 64~256 범위여야 합니다. 이는 MMAv5 타일로서 유효하고 유용한 크기입니다. 조건을 벗어나면 기존 휴리스틱으로 fallback하므로 작은 shape나 극단적인 shape에서 안전합니다.

3. ctaOrder swap: 인접 CTA가 M에서 짝이 되도록

  SmallVector<unsigned> ctaOrder;
  for (unsigned dim = dTy.getRank(); dim > 0; --dim)
    ctaOrder.push_back(dim - 1);
  // The first two order entries are the last tensor dimensions: N, then M.
  if (preferTwoCTA)
    std::swap(ctaOrder[0], ctaOrder[1]);

분할 개수만 바꿔서는 부족합니다. CTA ID의 가장 빠르게 변하는 축이 N이면, CTA 0과 1은 N 방향으로 이웃하게 됩니다. order를 swap해 M을 fastest 축으로 만들어야 CTA 0/1, 2/3이 M 방향 페어가 되고, two-CTA MMA의 하드웨어 페어링과 정확히 맞습니다.

4. 테스트 변화로 보는 레이아웃

4-CTA dot_split_mn 테스트입니다.

- CGALayout = [[0, 1], [1, 0]]
+ CGALayout = [[1, 0], [0, 1]]

첫 번째 basis(CTA ID bit 0)가 이제 M([1, 0])을 가리킵니다. 2-CTA 테스트에서는 dot 레이아웃이 [[0, 1]]에서 [[1, 0]](M 분할)로 바뀌었습니다. B operand load 레이아웃은 [[0, 0]]이 되었는데, 이는 B가 두 CTA에 broadcast된다는 의미입니다. M 페어는 같은 N 범위를 공유하므로 자연스러운 결과입니다. threadsPerWarp도 [2, 16]에서 [1, 32]로 바뀌었습니다. CTA 내부 타일이 N 방향으로 넓어진 만큼 warp 내 스레드 배치도 그에 맞게 재계산된 것입니다.

왜 이게 좋은가

GB200, FP16, M=N=8192, SM clock 1.2 GHz 고정, WS(warp specialization) 비활성 조건의 측정 결과(TFLOPS)입니다.

K	2 CTAs	4 CTAs	8 CTAs
512	810.8 → 807.8	468.0 → 493.5	529.5 → 542.3
2048	1189.5 → 1189.6	525.8 → 562.7	716.1 → 774.4
8192	1354.3 → 1335.9	531.4 → 579.3	743.9 → 805.6

4-CTA는 K=8192에서 약 +9%, 8-CTA는 약 +8% 향상됐습니다. 2-CTA는 사실상 변화가 없거나 노이즈 범위입니다. 이 표의 절대값을 보면 4/8-CTA 구성은 여전히 2-CTA보다 느리므로, 멀티 CTA 경로는 앞으로도 최적화 여지가 큽니다.

일반적 교훈

• 레이아웃은 하드웨어 협업 단위를 반영해야 합니다. "몇 개로 나누느냐"만큼 "어떤 순서로 이웃하느냐"(ctaOrder)가 중요합니다. 하드웨어가 인접 ID끼리 협업한다면 그 축을 fastest로 두어야 합니다.
• 아키텍처 특화 휴리스틱은 좁게 gating하고 fallback을 남겨야 합니다. compute capability 범위와 청크 크기 조건을 둘 다 만족할 때만 새 경로를 타게 해서 기존 동작을 보존했습니다.
• 검증은 한 곳에서 합니다. 리뷰 피드백에 따라 중복 target 체크를 제거하고 getNVIDIAComputeCapability의 검증에 맡긴 것은 작지만 좋은 정리입니다.

참고 자료

• CUDA C++ Programming Guide — Thread Block Cluster 개념(멀티 CTA 협업의 기반) 공식 문서
• PTX ISA — Blackwell tcgen05(two-CTA MMA) 명령어가 정의된 PTX 공식 문서
• std::swap — ctaOrder 교환에 사용된 표준 함수

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