Flash-KMeans는 IO-인식(IO-aware) 배치 K-평균(K-Means) 군집화(clustering)를 위한 구현체로, Triton GPU 커널을 활용하여 최적화된 성능을 제공합니다. 이 프로젝트는 Sparse VideoGen2의 공식 K-Means 구현으로 사용됩니다.

핵심 방법론은 다음과 같습니다:

1. Triton GPU 커널 최적화 (Large-D 및 데이터 타입 지원):

K-Means의 핵심 단계인 할당(assignment) 커널은 두 가지 주요 방식으로 구현되어 입력 데이터의 차원(D)과 데이터 타입에 따라 자동으로 디스패치(dispatch)됩니다.

• Small-D 경로 (D $\le$ 512): 이 경로는 한 번의 프로그램 실행으로 x_tile (BN, D) 데이터를 로드하고, BLOCK_K 단위로 중심점(centroids)을 스트리밍(streaming)하면서 처리합니다. NVIDIA H200, H100, A100, GB10과 같은 특정 아키텍처에 대해서는 flash-kmeans-tune을 통해 도출된 손으로 튜닝된 테이블(hand-tuned table)을 기반으로 최적의 (BLOCK_N, BLOCK_K, num_warps, num_stages) 조합을 선택합니다. 이는 각 GPU 아키텍처의 특성에 맞춰 성능을 극대화합니다.

• Split-D 경로 (D $>$ 512 또는 Small-D 커널이 공유 메모리(shared memory)에 맞지 않을 경우): 이 경로는 외곽 K 루프(outer K loop)와 BLOCK_D 단위로 타일링(tiling)된 내곽 D 루프(inner D loop)를 사용합니다. 교차 누산기(cross accumulator)인 (BN, BK)는 D 루프 전반에 걸쳐 레지스터(registers)에 유지되므로, 대규모 거리 행렬(distance matrix)을 구체화(materialized)하지 않는 K-스트리밍 속성(K-streaming property)이 유지됩니다. 알려지지 않은 GPU의 경우, 이 경로는 보수적인 폴백(fallback) 설정($BLOCK_N=32, BLOCK_K=32, BLOCK_D=32, num_warps=4, num_stages=1$)으로 자동 디스패치됩니다. 이는 튜닝 데이터에 대한 의존성을 피하며, _fit_config_to_smem_split_d 후처리 과정을 통해 공유 메모리 예산에 관계없이 실행이 가능하도록 보장합니다.

2. 멀티-GPU 지원 (대규모 N 워크로드):

kmeans_largeN 워크로드의 경우, Flash-KMeans는 자동 멀티-GPU 스케일링을 지원합니다. $device=None$으로 설정하면 모든 가시적인 GPU를 자동으로 활용합니다.

• 데이터 파티셔닝(Data Partitioning): N개의 데이터 포인트는 연속적인 블록 파티션(block partitions)으로 분할되어 각 GPU에 할당됩니다. 각 GPU는 독립적으로 자신의 파티션에 대해 이중 버퍼링(double-buffered)된 H2D (Host-to-Device) 및 연산 파이프라인(compute pipeline)을 실행합니다. 이 방식은 PCIe 대역폭(bandwidth)이 GPU 수에 비례하여 확장되도록 합니다.

• 경량 AllReduce: 모든 GPU가 로컬 누적(local accumulation)을 완료한 후, 부분적인 중심점 합(약 4MB)과 카운트(약 32KB)는 D2D (Device-to-Device) 복사(NVLink)를 통해 GPU 0으로 수집(gather)됩니다. 이후 합산(reduced)되고, 새로운 중심점은 다시 브로드캐스트(broadcast)됩니다. 데이터 크기가 작기 때문에 NCCL 의존성 없이 수동 복사(manual copies)가 더 빠르고 모든 것을 단일 프로세스 내에서 유지할 수 있습니다.

• H2D / D2D 오버랩(Overlap): H2D 전송(PCIe)과 AllReduce (NVLink)는 서로 다른 하드웨어 경로를 사용하므로 동시에 실행될 수 있습니다. 다음 이터레이션(iteration)의 첫 H2D 블록은 현재 이터레이션의 AllReduce 동안 미리 가져와져(prefetched), AllReduce의 지연 시간(latency)을 전송 시간 뒤에 숨길 수 있습니다.

Flash-KMeans는 fast_pytorch_kmeans, fastkmeans, batched torch kmeans와 같은 기존 구현 대비 현저한 성능 향상을 제공하며, 특히 대규모 입력 데이터 및 대규모 N 워크로드에서 효율성을 보여줍니다.