[vllm] [vLLM] ROCm 환경에서 MXFP4 GEMM 커널 최적화: 환경 변수에서 동적 설정으로

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

PR 링크: vllm-project/vllm#61278 상태: Merged | 변경: +66 / -43

들어가며

LLM 추론 엔진인 vLLM에서 AMD GPU(ROCm) 환경의 성능을 최적화하려는 노력은 계속되고 있습니다. 특히 MXFP4(Microscaling FP4)와 같은 저정밀도 양자화 기술은 메모리 대역폭을 절약하고 연산 속도를 높이는 데 핵심적인 역할을 합니다.

기존 vLLM 코드에서는 VLLM_ROCM_USE_AITER_FP4_ASM_GEMM이라는 환경 변수를 통해 모든 레이어에 대해 일괄적으로 Weight Layout(Preshuffled vs Plain)을 결정했습니다. 하지만 이 방식에는 치명적인 단점이 있었습니다. Llama와 같은 모델에서는 성능 이득이 크지만, DeepSeek(DS)나 Kimi 같은 모델에서는 튜닝된 Preshuffled Triton 커널이 없어 오히려 성능 저하(Regression)가 발생했기 때문입니다.

이번 PR은 이러한 '전부 아니면 전무(All-or-nothing)' 식의 설정을 제거하고, 레이어 로드 시점에 AITER의 튜닝된 설정(Tuned Configs)이 있는지 확인하여 최적의 커널을 동적으로 선택하도록 개선했습니다. 이를 통해 성능 이득은 챙기면서 안정성까지 확보한 과정을 살펴보겠습니다.

코드 분석: 무엇이 어떻게 바뀌었나?

1. 환경 변수 의존성 제거 (`vllm/envs.py` & `vllm/_aiter_ops.py`)

가장 먼저 눈에 띄는 변화는 전역 환경 변수의 삭제입니다. 이제 사용자가 수동으로 커널 사용 여부를 결정할 필요가 없습니다.

Before:

# vllm/envs.py
"VLLM_ROCM_USE_AITER_FP4_ASM_GEMM": lambda: (
    os.getenv("VLLM_ROCM_USE_AITER_FP4_ASM_GEMM", "False").lower() in ("true", "1")
),

After:

# vllm/envs.py
# 해당 환경 변수 정의 및 로직이 완전히 삭제됨

이와 함께 vllm/_aiter_ops.py에서도 단순히 환경 변수 값을 읽어오던 is_asm_fp4_gemm_dynamic_quant_enabled 메서드가 하드웨어 지원 여부만 체크하는 is_asm_fp4_gemm_supported로 변경되었습니다.

2. AITER 튜닝 설정 확인 로직 추가 (`vllm/_aiter_ops.py`)

핵심은 "현재 레이어의 Shape(N, K)에 대해 최적화된 커널이 존재하는가?"를 판단하는 로직입니다. rocm_aiter_ops 클래스에 is_asm_fp4_gemm_tuned 메서드가 추가되었습니다.

After:

@staticmethod
@functools.cache
def is_asm_fp4_gemm_tuned(n: int, k_bytes: int) -> bool:
    if not current_platform.is_rocm():
        return False
    try:
        from aiter.ops.gemm_op_a4w4 import get_GEMM_config

        # M=1024, 8192에 대해 설정이 있는지 확인하는 heuristic 사용
        return all(
            get_GEMM_config(m, n, 2 * k_bytes) is not None for m in (1024, 8192)
        )
    except ImportError:
        return False

여기서 재미있는 점은 m 값을 (1024, 8192) 두 가지만 체크한다는 것입니다. 리뷰어 @simondanielsson의 질문에 대해 작성자 @afriedri는 "표준 M sweep에는 항상 1024와 8192가 포함되므로, 이 두 지점에 설정이 있다면 해당 (N, K) 형상은 전반적으로 튜닝되어 있다고 판단할 수 있는 저렴한(cheap) 확인 방법"이라고 설명했습니다.

3. 레이어별 동적 커널 선택 (`vllm/model_executor/kernels/linear/mxfp4/aiter.py`)

이제 모델이 로드된 후, 각 레이어는 자신의 Shape에 맞춰 최적의 경로를 선택합니다.

Before:

def __init__(self, config: MxFp4LinearLayerConfig) -> None:
    super().__init__(config)
    self.use_asm_gemm = rocm_aiter_ops.is_asm_fp4_gemm_dynamic_quant_enabled()

After:

def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
    # AITER에서 ASM(Prefill용)과 Triton Preshuffled(Decode용) 커널이 모두 튜닝된 경우에만 사용
    if self.asm_gemm_supported:
        n, k_bytes = layer.weight.shape
        self.use_asm_gemm = rocm_aiter_ops.is_asm_fp4_gemm_tuned(
            n, k_bytes
        ) and rocm_aiter_ops.is_triton_gemm_afp4wfp4_presh_ws_tuned(n, k_bytes)

    # (이후 Scale Swizzle 지원 여부 등에 따라 최종 결정)

이 변경을 통해 use_asm_gemm은 단순히 환경 변수를 따르는 것이 아니라, ASM GEMM과 Triton Preshuffled GEMM이 모두 해당 레이어 크기에 최적화되어 있을 때만 활성화됩니다.

왜 이게 좋은 최적화인가?

1. 압도적인 성능 향상 (Llama-3.3-70B 기준)

벤치마크 결과는 놀랍습니다. MI350X 환경에서 Llama-3.3-70B-Instruct-MXFP4 모델의 경우:

TTFT (Time To First Token): 모든 Concurrency와 TP(Tensor Parallel) 설정에서 18% ~ 82% 감소했습니다.
Throughput (tok/s): TP1 환경에서 최대 53.3% 향상되었습니다.

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

PR Analysis 의 다른글

이전글 [vllm] vLLM FlashInfer Allreduce 최적화: 단일 노드 성능 회복과 안정성 확보
현재글 : [vllm] [vLLM] ROCm 환경에서 MXFP4 GEMM 커널 최적화: 환경 변수에서 동적 설정으로
다음글 [flashinfer] FlashInfer의 새로운 실험적 기능: NVFP4 기반 MLA 디코드 최적화

[vllm] [vLLM] ROCm 환경에서 MXFP4 GEMM 커널 최적화: 환경 변수에서 동적 설정으로

들어가며

코드 분석: 무엇이 어떻게 바뀌었나?

1. 환경 변수 의존성 제거 (`vllm/envs.py` & `vllm/_aiter_ops.py`)

2. AITER 튜닝 설정 확인 로직 추가 (`vllm/_aiter_ops.py`)

3. 레이어별 동적 커널 선택 (`vllm/model_executor/kernels/linear/mxfp4/aiter.py`)

왜 이게 좋은 최적화인가?

1. 압도적인 성능 향상 (Llama-3.3-70B 기준)

2. Regression 방지

3. 유지보수성 향상

결론 및 교훈

참고 자료

댓글

관련 포스트

PR Analysis 의 다른글