[Triton] AMD PartitionedSharedEncodingAttr 도입으로 shared memory 파티셔닝 지원

PR 링크: triton-lang/triton#9374 상태: Merged | 변경: +650 / -65

들어가며

GPU의 shared memory는 물리적으로 bank 단위로 나뉘어 있으며, 동시에 같은 bank에 접근하면 conflict가 발생하여 성능이 저하된다. 이 PR은 텐서를 여러 물리적 shared memory 파티션에 분산 배치할 수 있는 PartitionedSharedEncodingAttr를 도입한다. 이전 패치가 revert된 이유(ProxyFenceInsertion에서 getBufferIds 대신 getAllBufferIdsWithAliases를 사용해야 했던 문제)도 함께 수정했다.

핵심 코드 분석

Before: 단일 buffer 매핑

// Before: Value -> 단일 Buffer
using ValueBufferMapT = llvm::MapVector<Value, BufferT *>;

BufferId getBufferId(Value value) const {
  if (valueBuffer.count(value)) {
    return valueBuffer.lookup(value)->id;
  } else {
    return InvalidBufferId;
  }
}

// alias buffer 포함한 모든 buffer ids
BufferIdSetT getBufferIds(Value value) const {
  BufferIdSetT bufferIds;
  auto allocBufferId = getBufferId(value);
  if (allocBufferId != InvalidBufferId)
    bufferIds.insert(allocBufferId);
  for (auto *buffer : aliasBuffer.lookup(value)) {
    if (buffer->id != InvalidBufferId)
      bufferIds.insert(buffer->id);
  }
  return bufferIds;
}

After: 다중 buffer 매핑 + 명확한 API 분리

// After: Value -> 여러 Buffer (partitioned tensor 지원)
using ValueBufferMapT = llvm::MapVector<Value, SmallVector<BufferT *>>;

// 해당 value의 직접 할당된 buffer ids만 반환
SmallVector<BufferId> getBufferIds(Value value) const {
  SmallVector<BufferId> bufferIds;
  auto it = valueBuffer.find(value);
  if (it == valueBuffer.end())
    return bufferIds;
  for (auto *buffer : it->second) {
    bufferIds.push_back(buffer->id);
  }
  return bufferIds;
}

// alias buffer까지 포함하는 상위집합
BufferIdSetT getAllBufferIdsWithAliases(Value value) const {
  BufferIdSetT bufferIds;
  for (auto bufferId : getBufferIds(value)) {
    bufferIds.insert(bufferId);
  }
  for (auto *buffer : aliasBuffer.lookup(value)) {
    if (buffer->id != InvalidBufferId)
      bufferIds.insert(buffer->id);
  }
  return bufferIds;
}

PartitionedSharedEncoding 정의

def PartitionedSharedEncodingAttr
    : TritonGPU_Attr<"PartitionedSharedEncoding",
                     "partitioned_shared_encoding", [...]> {
  // numPartitions: 물리적 파티션 수 (서로 다른 메모리 슬롯에 배치)
  // numGroups: 그룹 수 (각 그룹은 numPartitions개의 연속 조각 포함)
  // partitionDim: 텐서를 분할하는 차원
  // partitionLayout: 각 조각 내부의 shared memory layout
}

왜 이게 좋은가

1. Bank conflict 감소: 텐서 조각을 서로 다른 물리적 shared memory 슬롯에 배치하여 동시 접근 시 conflict를 줄인다.
2. API 명확성: getBufferIds(직접 buffer만)와 getAllBufferIdsWithAliases(alias 포함)를 분리하여 사용 의도를 명확히 한다.
3. 이웃 관계 추적: BufferT::neighbors를 통해 서로 다른 파티션에 배치되어야 하는 buffer 간 관계를 추적한다.
4. 기존 호환성: partitioned가 아닌 일반 tensor는 기존과 동일하게 단일 buffer로 동작한다.

정리

이 PR은 Triton의 allocation system을 확장하여 partitioned shared memory encoding을 지원한다. 단일 value가 여러 buffer를 가질 수 있도록 ValueBufferMapT를 변경하고, buffer ID 조회 API를 명확히 분리했다. 이는 AMD GPU에서 shared memory partition conflict를 하드웨어 수준에서 해결하기 위한 기반이 된다.

참고 자료

• triton-lang/triton#9374
• triton-lang/triton#9367 (원본 PR)

---

이 글은 AI를 활용하여 PR의 핵심 변경사항을 분석하고 정리한 것입니다. 실제 코드의 맥락은 원본 PR을 참고해 주세요.