[triton] AutoWS에서 TMA와 non-TMA 로드 혼합 시 self-latency 및 MMA 처리 수정

PR 링크: triton-lang/triton#9111 상태: Merged | 변경: +252 / -196

들어가며

Triton의 Auto Warp Specialization(AutoWS)은 로드와 MMA를 서로 다른 warp 그룹에 배치하여 메모리 접근과 연산을 오버랩합니다. 그러나 TMA(Tensor Memory Accelerator)와 non-TMA 로드가 혼합된 경우, MMA의 self-latency 결정과 lowerMMA 처리에 문제가 있었습니다.

핵심 코드 분석

AssignLatencies에서 self-latency 결정:

auto cantWarpSpec = [](Operation *op) { return isa<LoadOp>(op); };
auto warpSpecHelper = ttng::MMAv5PipelineableOperandsHelper(
    mma, forOp, [&](Operation *op) {
      return isLoadToBePipelined(op) && !cantWarpSpec(op);
    });
if (warpSpecHelper.isPipelineable ||
    (warpSpecHelper.isOperandsStateDetermined &&
     llvm::none_of(warpSpecHelper.unpipelineableOperandDefs, cantWarpSpec)))
  mmaSelfLatency[mma] = 0;

모든 MMA 피연산자가 warp specialized될 수 있으면 self-latency를 0으로 설정합니다. 일반 LoadOp은 warp specialize할 수 없으므로, non-TMA 로드가 있으면 self-latency를 유지합니다. lowerMMA도 warp-specialized MMA를 처리하도록 확장되었습니다.

왜 이게 좋은가

Hopper numStages <= 1일 때 autoWS를 끄는 가드도 추가하여, 하류에서 발생하는 다양한 문제를 사전에 방지합니다. TMA와 non-TMA 혼합은 실제 프로덕션 커널에서 흔한 패턴이므로, 이 수정은 AutoWS의 적용 범위를 크게 넓힙니다.

정리

• TMA와 non-TMA 로드 혼합 시 MMA self-latency 올바르게 설정
• lowerMMA에서 warp-specialized MMA 처리 지원
• numStages <= 1일 때 Hopper autoWS 비활성화
• ConSan 통합 테스트 추가로 커버리지 확대

참고 자료

• triton-lang/triton#9111

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