[vllm] AMD ROCm GFX11 GPU의 L2 캐시 충돌을 피하기 위한 GEMM 가중치 행 스트라이드 최적화

2026년 10월 15일수정: 2026년 10월 15일

PR 링크: vllm-project/vllm#55090 상태: Merged | 변경: +198 / -13

들어가며

최근 vLLM 프로젝트의 GitHub PR에서는 AMD의 ROCm 환경, 특히 GFX11 아키텍처(RDNA3/3.5)를 사용하는 GPU에서 발생하는 성능 병목 현상을 해결하기 위한 중요한 최적화가 이루어졌습니다. 이 PR의 핵심은 행렬 곱셈(GEMM) 연산 시 가중치(weight) 텐서의 행 스트라이드(row stride)가 특정 조건에서 L2 캐시의 해시 충돌을 일으켜 대역폭 성능을 저하시키는 문제를 해결하는 것입니다. 이 글에서는 해당 PR의 코드 변경 사항을 분석하고, 이 최적화가 왜 효과적이며 어떤 교훈을 주는지 상세히 살펴보겠습니다.

PR 설명에 따르면, GFX11 아키텍처에서 가중치 행 스트라이드가 2048 바이트의 배수일 때 L2/MALL 해시 충돌이 발생하여 상당한 대역폭 손실이 발생합니다. 이 문제를 해결하기 위해 가중치 행 스트라이드를 한 캐시 라인(128 바이트)만큼 오프셋하여 이러한 충돌을 회피합니다. 이 변경은 특히 TTFT(Time To First Token)를 최대 13%, TPOT(Time Per Output Token)를 최대 3%까지 개선하는 효과를 가져왔습니다. 메모리 사용량 증가는 해당 레이어에 한해 미미하며, Qwen3-8B 모델의 경우 가중치 메모리의 약 1.29% (171 MiB) 증가에 그쳤습니다.

이 최적화는 K * element_size가 2048의 배수일 때 적용되며, bf16/fp16 데이터 타입에서는 K가 1024의 배수일 때 대부분의 Llama/Qwen 계열 모델에서 발생하는 조건입니다. 예를 들어 Qwen3-8B 모델의 모든 선형 레이어(qkv/o/gate_up의 K=4096, down의 K=12288)가 이 조건에 해당합니다.

코드 분석

이번 PR의 변경 사항은 주로 vllm/model_executor/layers/utils.py와 vllm/model_executor/layers/linear.py 파일에 집중되어 있으며, 테스트 코드(tests/model_executor/layers/test_rocm_unquantized_gemm.py)에서도 관련 로직 검증이 추가되었습니다.

1. `vllm/model_executor/layers/utils.py`의 핵심 변경 사항

이 파일에서는 L2 캐시 충돌을 회피하기 위한 새로운 로직이 추가되었습니다.

`cache_aliasing_pad_elems` 함수 추가

이 함수는 특정 조건에서 필요한 패딩 크기를 계산합니다. GFX11 아키텍처에서 행 스트라이드가 2048 바이트의 배수일 때 발생하는 캐시 충돌을 피하기 위해, 128 바이트 캐시 라인 크기만큼의 패딩을 요소(element) 단위로 변환하여 반환합니다. 만약 충돌이 발생하지 않으면 0을 반환합니다.

def cache_aliasing_pad_elems(k: int, element_size: int) -> int:
    """Row-stride pad, in elements, that stops an ``[n, k]`` weight aliasing on gfx11,
    or 0 when the stride does not alias.

    Args:
        k: Number of elements per weight row.
        element_size: Size of one element in bytes.

    Returns:
        The number of elements to add to the row stride.

    """
    if (k * element_size) % 2048 != 0:
        return 0
    return 128 // element_size

`maybe_pad_weight_avoid_cache_aliasing` 함수 추가

이 함수는 실제 가중치 텐서에 패딩을 적용하는 로직을 캡슐화합니다. GFX11 ROCm 환경이 아니고, 텐서가 2D가 아니거나 연속적(contiguous)이지 않으면 원본 텐서를 그대로 반환합니다. 해당 조건이 만족되면 cache_aliasing_pad_elems 함수를 호출하여 필요한 패딩 크기를 얻고, torch.nn.functional.pad를 사용하여 행 스트라이드를 조정합니다. 중요한 점은 이 함수가 원본 텐서의 뷰(view)를 반환한다는 것입니다. 즉, 메모리 복사가 최소화되며, 이미 패딩된 가중치에 대해 다시 호출될 경우 원본 포인터를 그대로 반환하여 불필요한 연산을 방지합니다.

def maybe_pad_weight_avoid_cache_aliasing(weight: torch.Tensor) -> torch.Tensor:
    """Offset a dense weight's row stride so it stops aliasing.

    Returns a view with the original shape whose row stride is one cache line
    wider, or ``weight`` unchanged when the pad does not apply.
    """
    from vllm.platforms.rocm import on_gfx11

    if not (current_platform.is_rocm() and on_gfx11()):
        return weight
    if weight.dim() != 2 or not weight.is_contiguous():
        return weight
    pad = cache_aliasing_pad_elems(weight.shape[1], weight.element_size())
    if pad == 0:
        return weight
    return torch.nn.functional.pad(weight, (0, pad))[..., :-pad]

`rocm_unquantized_gemm_impl` 함수의 로직 개선

기존의 rocm_unquantized_gemm_impl 함수는 ROCm 환경에서 다양한 GEMM 커널(wvSplitK, LLMM1 등)을 선택하는 로직을 포함하고 있었습니다. 이번 PR에서는 이 로직이 GFX11 아키텍처의 캐시 충돌 문제를 고려하도록 수정되었습니다. 특히 wvSplitK 커널의 weight_ok_for_wvsplitk 조건이 변경되었습니다.

기존에는 weight.is_contiguous()만 만족하면 wvSplitK를 사용할 수 있었지만, 이제는 GFX11 환경에서 weight.stride(1) == 1인 경우에도 wvSplitK를 사용할 수 있도록 허용됩니다. 이는 maybe_pad_weight_avoid_cache_aliasing 함수가 적용되어 행 스트라이드가 조정된 가중치 텐서도 wvSplitK 커널에서 처리될 수 있음을 의미합니다. 이 변경은 리뷰어(@wjabbour)의 지적에 따라 GFX11에만 적용되도록 명확히 구분되었습니다.

# Before (simplified)
# skinny_gemm_enabled = (
#     envs.VLLM_ROCM_USE_SKINNY_GEMM
#     and (on_gfx9() or on_gfx1x())
#     and x.dtype in [torch.float16, torch.bfloat16]
#     and k % 8 == 0
#     and skinny_operands_compatible
# )
#
# if skinny_gemm_enabled:
#     if (m == 1 or m > 8) and 0 < n <= 5:
#         x_view = x.reshape(-1, x.size(-1)).contiguous()
#         cu_count = num_compute_units()
#         out = ops.wvSplitK(weight, x_view, cu_count, bias)
#         return out.reshape(*x.shape[:-1], weight.shape[0])
#     elif m % 4 == 0 and n == 1 and k <= 8192 and bias is None:
#         x_view = x.reshape(-1, x.size(-1)).contiguous()
#         out = ops.LLMM1(weight, x_view, 4)
#         return out.reshape(*x.shape[:-1], weight.shape[0])

# After (simplified and refactored based on review comments)
skinny_gemm_enabled = (
    envs.VLLM_ROCM_USE_SKINNY_GEMM
    and (on_gfx9() or on_gfx1x())
    and x.dtype in (torch.float16, torch.bfloat16)
    and k % 8 == 0
)

bias_ok = bias is None or bias.is_contiguous()
weight_dense = weight.is_contiguous()

# wvSplitK reads the weight row stride, so a gfx11 cache-aliasing-padded
# (K-contiguous) weight is fine; LLMM1 indexes off a raw pointer and needs dense.
weight_ok_for_wvsplitk = weight_dense or (on_gfx11() and weight.stride(1) == 1)

use_wvsplitk = (
    skinny_gemm_enabled and weight_ok_for_wvsplitk and bias_ok
    and m > 8 and 0 < n <= 5
)

use_llmm1 = (
    skinny_gemm_enabled and weight_dense and bias is None
    and m % 4 == 0 and n == 1 and k <= 8192
)

if use_wvsplitk or use_llmm1:
    x_view = x.reshape(-1, x.size(-1)).contiguous()
    if use_wvsplitk:
        out = ops.wvSplitK(weight, x_view, num_compute_units(), bias)
        return out.reshape(*x.shape[:-1], weight.shape[0])
    out = ops.LLMM1(weight, x_view, 4)
    return out.reshape(*x.shape[:-1], weight.shape[0])

2. `vllm/model_executor/layers/linear.py`의 변경 사항

UnquantizedLinearMethod 클래스의 process_weights_after_loading 메서드에서 가중치 로딩 후 처리 로직이 수정되었습니다. 기존에는 가중치가 연속적이지 않을 경우 t().contiguous().t()를 통해 연속성을 확보했지만, 이제 ROCm 환경에서는 maybe_pad_weight_avoid_cache_aliasing 함수를 호출하여 가중치 데이터에 패딩을 적용합니다.

# Before
def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
    # ... (other logic)
    if (
        current_platform.is_rocm()
        and weight.dim() == 2
        and weight.stride(0) != 1
        and weight.stride(1) != 1
        and weight.stride(0) != weight.shape[1]
        and weight.stride(1) != weight.shape[0]
        and weight.stride(0) != weight.shape[1] * weight.element_size()
        and weight.stride(1) != weight.shape[0] * weight.element_size()
        and weight.stride(0) != weight.shape[1] + weight.shape[1] * weight.element_size()
        and weight.stride(1) != weight.shape[0] + weight.shape[0] * weight.element_size()
        and weight.stride(0) != weight.shape[1] + weight.shape[1] * weight.element_size() * 2
        and weight.stride(1) != weight.shape[0] + weight.shape[0] * weight.element_size() * 2
        and weight.stride(0) != weight.shape[1] * 2 + weight.shape[1] * weight.element_size()
        and weight.stride(1) != weight.shape[0] * 2 + weight.shape[0] * weight.element_size()
        and weight.stride(0) != weight.shape[1] * 2 + weight.shape[1] * weight.element_size() * 2
        and weight.stride(1) != weight.shape[0] * 2 + weight.shape[0] * weight.element_size() * 2
    ):
        layer.weight.data = weight.t().contiguous().t()

# After
def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
    # ... (other logic)
    elif current_platform.is_rocm():
        layer.weight.data = maybe_pad_weight_avoid_cache_aliasing(layer.weight.data)

3. 테스트 코드 변경 사항

tests/model_executor/layers/test_rocm_unquantized_gemm.py 파일에는 새로운 테스트 케이스들이 추가되어 변경된 로직을 검증합니다.

test_rocm_unquantized_gemm_wvsplitk_accepts_row_stride_padded_weight: 행 스트라이드가 패딩된 가중치가 wvSplitK 커널에서 올바르게 처리되는지 확인합니다.
test_cache_aliasing_pad_elems: cache_aliasing_pad_elems 함수의 정확성을 다양한 K와 element_size 조합으로 테스트합니다.
test_maybe_pad_weight_avoid_cache_aliasing_skips_non_gfx11: GFX11이 아닌 환경에서는 패딩이 적용되지 않음을 확인합니다.
test_maybe_pad_weight_avoid_cache_aliasing: GFX11 환경에서 패딩이 올바르게 적용되고, 동일한 가중치에 대해 반복 호출 시 원본 포인터를 유지하는지 확인합니다.
test_rocm_unquantized_gemm_padded_weight_real_kernel: 실제 커널 실행 시 패딩된 가중치가 올바르게 동작하는지 검증합니다.

리뷰 과정에서 @wjabbour가 지적한 것처럼, 일부 테스트는 GFX11 환경에서만 실행되도록 @pytest.mark.skipif(_on_gfx11(), reason=...)와 같이 명시적으로 제한되었습니다.

왜 이게 좋은가?

이 PR의 최적화는 다음과 같은 이유로 매우 효과적입니다:

직접적인 성능 향상: GFX11 GPU의 L2 캐시 충돌 문제를 해결함으로써, GEMM 연산의 메모리 대역폭을 최대한 활용할 수 있게 되었습니다. 벤치마크 결과에 따르면, Qwen3-8B 모델에서 TTFT가 최대 13.20%, TPOT가 최대 3.39%까지 개선되었습니다. 이는 특히 대규모 언어 모델의 추론 속도에 직접적인 영향을 미칩니다.
정밀한 조건 기반 적용: 캐시 충돌은 특정 하드웨어 아키텍처(GFX11)와 특정 연산 조건(K * element_size가 2048의 배수)에서만 발생합니다. 이 PR은 maybe_pad_weight_avoid_cache_aliasing 함수를 통해 이러한 조건을 정확히 검사하고, 필요한 경우에만 최소한의 패딩을 적용합니다. 이는 불필요한 연산이나 메모리 오버헤드를 방지합니다.
메모리 효율성: 패딩은 필요한 경우에만 적용되며, 그 영향은 미미합니다. Qwen3-8B 모델의 경우 가중치 메모리의 약 1.29% 증가에 그쳤습니다. 또한, torch.nn.functional.pad를 사용하여 원본 텐서의 뷰를 생성하므로, 대부분의 경우 추가적인 메모리 복사가 발생하지 않습니다.
코드의 명확성 및 유지보수성: 리뷰어(@wjabbour)의 피드백을 반영하여, 커널 선택 로직이 더 명확하게 분리되고 조건부로 적용되도록 개선되었습니다. 이는 코드의 가독성을 높이고 향후 유지보수를 용이하게 합니다.
일반적인 교훈: 이 최적화는 하드웨어 아키텍처의 미묘한 특성(캐시 동작, 메모리 로딩 패턴)이 고성능 컴퓨팅에서 얼마나 중요한지를 다시 한번 보여줍니다. 특정 하드웨어에서 발생하는 성능 병목 현상을 해결하기 위해, 데이터 구조의 스트라이드를 조정하는 것은 매우 효과적인 기법이 될 수 있습니다. 또한, 최적화 로직을 특정 하드웨어 아키텍처에만 국한시키는 것은 중요하며, 이를 위해 플랫폼별 조건 검사를 명확히 하는 것이 필수적입니다.

결론

이번 vLLM의 PR은 AMD ROCm 환경의 GFX11 GPU에서 GEMM 연산의 L2 캐시 충돌 문제를 해결하여 상당한 성능 향상을 이끌어냈습니다. 가중치 행 스트라이드를 미세하게 조정하는 이 최적화는 복잡한 하드웨어 특성을 이해하고 이를 코드에 반영하는 엔지니어링의 중요성을 잘 보여줍니다. 앞으로도 이러한 세밀한 최적화를 통해 LLM 추론 성능을 지속적으로 개선해 나갈 것으로 기대됩니다.

References

AMD ROCm Documentation
torch.nn.functional.pad
vLLM GitHub Repository
GFX11 Architecture Details (Note: Specific L2 cache behavior details might require deeper ROCm/driver documentation)

참고 자료

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

PR Analysis 의 다른글

이전글 [vllm] vLLM의 Mamba 블록 테이블 가더(Gather) 최적화: 7개 커널을 1개로 통합하여 6배 속도 향상
현재글 : [vllm] AMD ROCm GFX11 GPU의 L2 캐시 충돌을 피하기 위한 GEMM 가중치 행 스트라이드 최적화
다음글 [flashinfer] Kimi K3를 위한 초고속 통신 최적화: FlashInfer의 Blackwell MNNVL CuTe DSL 프로파일링 분석