[sglang] ROCm aiter에서 Ring Attention 지원과 K/V 버퍼 재사용으로 MiniMax-H3 2-GPU 추론 11% 단축

PR 링크: sgl-project/sglang#49078 상태: Merged | 변경: +139 / -22

들어가며

sglang-diffusion에서 MiniMax-H3(T2VA, 비디오+오디오 생성 DiT)를 MI355X 2장으로 돌리면 118.2초가 걸립니다. GPU-seconds로 환산하면 1 GPU보다 36%를 더 씁니다. 같은 구성에서 B200의 오버헤드는 6%입니다.

원인은 interconnect 토폴로지입니다. 2 GPU는 XGMI 링크 하나(~57 GB/s)로만 통신하므로 Ulysses sequence parallelism의 all-to-all이 레이어 연산 대비 큰 비중을 차지합니다. head-group pipeline(#47567)으로 일부를 overlap해도 첫 그룹의 입력 교환과 마지막 그룹의 출력 교환은 가려지지 않습니다. 그 결과 약 46ms짜리 레이어에서 5ms 정도가 통신으로 노출됩니다.

Ring attention(degree 2)은 이 상황에 맞는 방식입니다. K/V를 한 번만 보내고, 그 전송을 로컬 chunk의 attention 계산과 겹치며, 출력 교환이 아예 없습니다. 문제는 --ring-degree가 ROCm에서 가장 빠른 exact backend인 aiter를 거부했다는 점입니다. 이 PR은 두 가지를 바꿉니다.

1. aiter가 ring merge에 필요한 softmax LSE를 반환하도록 합니다.
2. ring의 K/V 버퍼를 호출마다 새로 할당하지 않고 재사용해 OOM을 해결합니다.

코드 분석

1. aiter.py — LSE 반환과 ring 지원 선언

Ring attention은 각 hop에서 받은 KV chunk에 대한 부분 attention 결과를 합칩니다. 이때 out = Σ exp(lse_i - lse) * out_i 형태의 online-softmax 병합을 하려면 kernel이 softmax의 log-sum-exp(LSE)를 함께 돌려줘야 합니다.

Before에서는 LSE를 버렸습니다.

    ) -> torch.Tensor:
        return aiter.flash_attn_func(
            query, key, value,
            ...
            return_lse=False,
            how_v3_bf16_cvt=AITER_BF16_CVT_MODE,
        )

After에서는 옵션으로 LSE를 받습니다.

        return_lse: bool = False,
    ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
        output = aiter.flash_attn_func(
            ...
            return_lse=return_lse,
            how_v3_bf16_cvt=AITER_BF16_CVT_MODE,
        )
        return tuple(output[:2]) if return_lse else output

forward(..., return_softmax_lse=True)이면 FP8 경로를 건너뛰고 항상 BF16 kernel을 씁니다. FP8 prefill 경로는 LSE를 노출하지 않기 때문입니다. 정확도가 필요한 merge 경로에서 정밀도 경로를 하나로 고정한 판단입니다.

        if return_softmax_lse:
            return self._forward_bf16(query, key, value, return_lse=True)

packed 경로용 forward_ring_kv_chunk도 새로 추가됐습니다. [T, H, D] query를 회전되어 온 KV chunk 하나에 attend시키고, [H, Tq] LSE를 반환합니다.

    @torch.compiler.disable
    def forward_ring_kv_chunk(self, query, key, value):
        output, lse = aiter.flash_attn_varlen_func(
            q=query.contiguous(), k=key.contiguous(), v=value.contiguous(),
            cu_seqlens_q=_cu_seqlens((0, query.shape[0]), query.device),
            cu_seqlens_k=_cu_seqlens((0, key.shape[0]), key.device),
            ...
            causal=False,
            return_lse=True,
        )[:2]
        return output, lse

여기서 눈여겨볼 부분은 _cu_seqlens의 lru_cache입니다.

@functools.lru_cache(maxsize=64)
def _cu_seqlens(bounds: tuple[int, ...], device: torch.device) -> torch.Tensor:
    # one upload per bounds and device; callers must not mutate the result
    return torch.tensor(bounds, dtype=torch.int32, device=device)

매 hop마다 torch.tensor(..., device=cuda)를 만들면 그때마다 작은 H2D 복사가 생기고, 경우에 따라 동기화 지점이 됩니다. shape는 레이어와 step 사이에 거의 변하지 않으므로 (bounds, device)를 키로 캐싱하면 업로드가 한 번으로 끝납니다. 주석에 적힌 대로 반환된 텐서를 mutate하면 안 된다는 계약이 붙습니다.

마지막으로 supports_ring_rotation()을 True로 선언했고, server args의 이름 기반 화이트리스트에도 추가했습니다.

-RING_CAPABLE_ATTENTION_BACKENDS = ("fa", "sage_attn")
+RING_CAPABLE_ATTENTION_BACKENDS = ("fa", "sage_attn", "aiter")

2. usp.py — Ring K/V 버퍼 재사용 (OOM 해결)

이 PR에서 가장 교훈적인 변경입니다. backend와 무관하게 적용됩니다.

Before에서는 호출마다 버퍼 쌍을 새로 할당했습니다.

    kv0 = torch.stack((k, v))
    kv_bufs = [kv0, torch.empty_like(kv0)]

After에서는 staging 버퍼를 재사용합니다.

    kv_bufs = [
        _a2a_staging_buffer(f"ring_kv_{i}", (2, *k.shape), k.dtype, k.device)
        for i in range(2)
    ]
    torch.stack((k, v), out=kv_bufs[0])

얼핏 보면 PyTorch caching allocator가 알아서 재활용할 것 같지만, 실제로는 그렇지 않았습니다. P2P send/recv는 별도 stream에서 돌기 때문에, 그 stream에서 쓰이는 텐서는 GPU가 해당 작업에 실제로 도달할 때까지 해제되지 않습니다(record_stream 의미론). eager 모드에서 host는 GPU보다 수백 레이어 앞서 kernel을 enqueue할 수 있고, 그러면 대기 중인 레이어 하나마다 ~1GB짜리 K/V 쌍이 하나씩 묶여 있게 됩니다. 실제로 H3를 ring 2로 돌리면 denoise 4 step 안에 HSA_STATUS_ERROR_OUT_OF_RESOURCES가 났고, standalone 루프로도 재현됐습니다.

이름이 붙은 고정 staging 버퍼를 쓰면 메모리 사용량이 큐 깊이와 상관없이 일정해집니다. torch.stack(..., out=)으로 바로 써 넣으므로 contiguous 레이아웃과 "K/V를 한 번의 P2P로 보낸다"는 기존 설계도 그대로 유지됩니다.

3. minimax_h3.py — 게이트 확장

FA만 허용하던 하드코딩을 튜플 상수로 바꿨습니다.

# packed kernels that return the softmax LSE the ring merge needs
_RING_ATTENTION_BACKENDS = (AttentionBackendEnum.FA, AttentionBackendEnum.AITER)
...
-        if attention._attention_backend_enum is not AttentionBackendEnum.FA:
+        if attention._attention_backend_enum not in _RING_ATTENTION_BACKENDS:

pipeline config의 validate_server_args도 같은 방식으로 aiter를 허용합니다. 허용 조건을 "LSE를 반환하는 packed kernel"이라는 의미 단위로 묶어 두었기 때문에, 이후 backend를 추가할 때 수정할 곳이 한 군데로 모입니다.

왜 이게 좋은가

성능

8x MI350X(gfx950, MI355X와 같은 XGMI mesh)에서 H3 T2VA 1344x768, 124 frames, 49 steps로 측정한 결과입니다.

2 GPUs	GPUs 0-1	GPUs 4-5
main, Ulysses 2	117.70s	115.22s
this PR, ring 2	104.51s (-11.2%)	105.52s (-8.4%)

1 GPU는 204초가 걸리므로, ring 2는 이상적인 절반에서 2.5% 이내까지 줄었습니다.

정확도

두 KV chunk를 ring merge한 결과는 전체 key에 대한 단일 호출과 같습니다(packed/dense 경로 모두 단위 테스트로 확인). E2E 결과는 결정적이고(MP4 md5 일치), 1-GPU 대비 PSNR은 ring 2가 19.5dB, Ulysses 2가 18.7dB입니다. 따라서 ring 2는 Ulysses와 같은 수준의 floating-point 재정렬 오차 범위 안에 있습니다.

일반적 교훈

1. 병렬화 전략은 토폴로지에 따라 달라집니다. 저자는 ring이 항상 이기지 않는다는 점도 솔직하게 보고합니다. 4 GPU에서는 Ulysses 4(59.47s)가 Ulysses 2 x ring 2(66.97s)보다, 8 GPU에서도 Ulysses 8이 더 빠릅니다. 링크가 하나뿐인 2-GPU 구성에서 all-to-all이 비효율적이라는 특정 조건에 맞춘 해법입니다.
2. Side stream에 넘긴 텐서는 생각보다 오래 살아 있습니다. async P2P와 host run-ahead가 결합되면 "호출마다 할당"하는 패턴이 큐 깊이에 비례하는 메모리 누수처럼 동작합니다. 고정 버퍼 재사용이 정석입니다.
3. 작은 device tensor 생성도 캐싱할 가치가 있습니다. cu_seqlens처럼 shape에만 의존하는 메타데이터는 lru_cache로 H2D 업로드를 없앨 수 있습니다.
4. 대안을 측정으로 배제했습니다. pipeline group을 늘리는 방안(E2E 0.25% 개선)과 copy engine P2P(28GB/s, RCCL의 절반)는 효과가 없다는 것을 수치로 확인했고, 덕분에 ring을 택한 근거가 분명합니다.

참고 자료

• torch.stack — out= 인자로 staging 버퍼에 직접 K/V를 쓰는 데 사용
• torch.Tensor.record_stream — side stream에서 쓰이는 텐서의 수명 연장 의미론
• CUDA semantics - memory management — caching allocator와 stream 간 메모리 재사용
• functools.lru_cache — cu_seqlens 텐서 캐싱
• torch.compiler.disable — forward_ring_kv_chunk를 컴파일 대상에서 제외
• ROCm aiter — AMD GPU용 attention kernel 라이브러리

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