[vllm] vLLM, H200 GPU를 위한 FP8 MoE 다운 프로젝션 최적화: TMA 활용 및 커널 퓨전으로 7.6% 처리량 향상

2026년 11월 4일수정: 2026년 11월 4일

PR 링크: vllm-project/vllm#64760 상태: Merged | 변경: +282 / -21

들어가며

최근 대규모 언어 모델(LLM)의 발전은 모델의 크기와 복잡성을 증가시키고 있으며, 특히 Mixture-of-Experts (MoE) 아키텍처는 이러한 추세 속에서 효율적인 성능 확장을 위한 핵심 기술로 주목받고 있습니다. MoE 모델은 수많은 전문가(expert) 네트워크를 활용하여 입력에 따라 동적으로 활성화되는 전문가를 선택함으로써, 파라미터 수를 크게 늘리면서도 연산량을 효율적으로 관리할 수 있습니다.

vLLM은 LLM 추론을 위한 고성능 라이브러리로, 지속적으로 모델 아키텍처와 하드웨어 특성에 최적화된 기술을 도입하고 있습니다. 이번 PR은 NVIDIA H200 GPU 환경에서 FP8 정밀도를 사용하는 MoE 모델의 다운 프로젝션(down-projection) 단계 성능을 개선하는 데 초점을 맞추고 있습니다. 구체적으로, MoE 중간 텐서를 전문가 정렬 순서(expert-sorted order)로 유지하여 Tensor Memory Accelerator (TMA)를 활용하고, SiLU 활성화 함수, 행렬 곱셈, 그리고 FP8 양자화를 단일 Triton 커널로 융합(fuse)함으로써 종단 간(end-to-end) 처리량을 6.0%에서 7.6%까지 향상시키는 것을 목표로 합니다.

이 글에서는 해당 PR의 코드 변경 사항을 상세히 분석하고, 이러한 최적화가 왜 효과적인지, 그리고 이를 통해 얻을 수 있는 일반적인 교훈은 무엇인지 살펴보겠습니다.

코드 분석

이번 PR의 핵심 변경 사항은 vllm/model_executor/layers/fused_moe/experts/triton_moe.py 파일과 관련 설정 파일들에 집중되어 있습니다. 주요 변경 내용은 다음과 같습니다.

1. `triton_moe.py`: TMA 활용 및 커널 퓨전 도입

가장 중요한 변경은 TritonExperts 클래스의 apply 메서드에서 이루어졌습니다. 이 메서드는 MoE 레이어의 핵심 연산을 처리하며, 이번 PR에서는 H200 GPU의 TMA 기능을 활용하고 연산들을 융합하기 위한 로직이 추가되었습니다.

1.1. 전문가 정렬 순서 유지 및 TMA 활용 (Sorted-A TMA Path)

기존에는 MoE의 다운 프로젝션 단계에서 중간 활성화 텐서의 순서가 전문가별로 정렬되지 않아 TMA를 효율적으로 사용하기 어려웠습니다. 이번 PR에서는 중간 텐서를 전문가 정렬 순서로 유지하도록 하여, 활성화 텐서와 가중치 모두에 대해 TMA를 활용할 수 있게 되었습니다.

_enable_sorted_a_tma 플래그를 도입하여 이 기능을 제어하며, test_sorted_a_tma_moe_matches_route_major 테스트 케이스에서는 H200 GPU에서 이 경로가 기존 경로와 동일한 결과를 생성하는지 검증합니다. patch.object를 사용하여 invoke_fused_moe_triton_kernel 함수의 호출을 감시하고, A_is_sorted 인자가 True로 설정되었는지 확인합니다. 이는 TMA가 정렬된 입력에 대해 최적화되었음을 나타냅니다.

--- a/tests/kernels/moe/test_block_fp8.py
+++ b/tests/kernels/moe/test_block_fp8.py
@@ -181,6 +181,55 @@
     torch.set_default_device("cuda")
 
 
+@pytest.mark.skipif(
+    "H200" not in torch.cuda.get_device_name(),
+    reason="The sorted-A TMA path is tuned for H200",
+) # H200 GPU에서만 실행
+def test_sorted_a_tma_moe_matches_route_major(workspace_init):
+    """The H200 sorted path preserves routed FP8 block outputs exactly."""
+    torch.manual_seed(123)
+    m, experts, hidden, intermediate, topk = 1024, 256, 4096, 256, 8
+    x = (torch.randn((m, hidden)) / 10).to(torch.bfloat16)
+    w1, w2, quant_config = make_test_quant_config(
+        experts, intermediate, hidden, torch.bfloat16, torch.float8_e4m3fn,
+        block_shape=[128, 128],
+    )
+    moe_config = make_dummy_moe_config(
+        num_experts=experts, experts_per_token=topk, hidden_dim=hidden,
+        intermediate_size=intermediate, max_num_tokens=m,
+    )
+    kernel = modular_triton_fused_moe(moe_config, quant_config)
+    triton_experts = kernel.fused_experts
+    assert isinstance(triton_experts, triton_moe.TritonExperts)
+    topk_ids = torch.randint(experts, (m, topk))
+    topk_weights = torch.softmax(torch.randn((m, topk)), dim=-1)
+    kwargs = dict(
+        activation=MoEActivation.SILU,
+        global_num_experts=experts, expert_map=None,
+        apply_router_weight_on_input=False,
+    )
+    with set_current_vllm_config(vllm_config):
+        triton_experts._enable_sorted_a_tma = False
+        expected = kernel.apply(x, w1, w2, topk_weights, topk_ids, **kwargs)
+        triton_experts._enable_sorted_a_tma = True
+        with patch.object(
+            triton_moe, "invoke_fused_moe_triton_kernel", wraps=triton_moe.invoke_fused_moe_triton_kernel
+        ) as calls:
+            actual = kernel.apply(x, w1, w2, topk_weights, topk_ids, **kwargs)
+    assert any(call.kwargs.get("A_is_sorted") for call in calls.call_args_list)
+    assert torch.equal(expected, actual)
+```

또한, `apply` 메서드 내에서 `use_sorted_a` 조건을 평가하는 로직이 추가되었습니다. 이 조건은 TMA 사용 가능 여부, 활성화 함수 종류 (SILU), FP8 양자화 설정, H200 GPU인지 여부 등 다양한 요소를 고려하여 TMA 최적화 경로를 활성화할지 결정합니다.

```diff
--- a/vllm/model_executor/layers/fused_moe/experts/triton_moe.py
+++ b/vllm/model_executor/layers/fused_moe/experts/triton_moe.py
@@ -394,6 +431,56 @@
             )
         )
 
+        from vllm.model_executor.layers.fused_moe import get_config
+
+        down_configs = (
+            get_moe_configs(
+                E, cache2_dim, self.quant_config.config_name(hidden_states.dtype),
+                128, 128, down_moe=True,
+            )
+            if self.quant_config.use_fp8_w8a8
+            and self.block_shape == [128, 128]
+            and not get_config()
+            else None
+        )
+        tma_down_config = (
+            dict(down_configs[min(down_configs, key=lambda m: abs(m - num_tokens))])
+            if down_configs
+            else None
+        )
+
+        use_sorted_a = bool(
+            tma_down_config
+            and tma_down_config.pop("USE_TMA", False)
+            and self._enable_sorted_a_tma
+            and activation == MoEActivation.SILU
+            and self.activation_config.clamp_limit is None
+            and self._lora_context is None
+            and not hooked
+            and not self.quantization_emulation
+            and is_deep_gemm_e8m0_used()
+            and a2_scale is None
+            and sorted_token_ids is not None
+            and expert_map is None
+            and workspace2.dtype == torch.bfloat16
+            and current_platform.is_cuda()
+            and current_platform.is_device_capability(90)
+            and config["BLOCK_SIZE_M"] % 16 == 0
+            and w2.is_contiguous()
+            and w2.size(1) >= 4096
+            and w2.size(1) % tma_down_config["BLOCK_SIZE_N"] == 0
+            and cache2_dim % 128 == 0
+        )
+        if use_sorted_a:
+            assert sorted_token_ids is not None
+            num_sorted_rows = sorted_token_ids.numel()
+            # GEMM fills padding inside active blocks; quantization skips the tail.
+            intermediate_cache1 = workspace2.new_empty((num_sorted_rows, 1, N))
+
         # LoRA w13: applied to intermediate_cache1 before activation. When
         # the LoRA layer requested a dual-stream schedule, we run base w13
         # GEMM on the default stream and the LoRA fast-path on aux_stream;
@@ -522,7 +679,27 @@
         # The fused kernel has no clamp parameter, so a configured
         # SwiGLU clamp limit falls through to the unfused path.
-        if (
+        if use_sorted_a:
+            qintermediate_cache2 = intermediate_cache1.new_empty(
+                (num_sorted_rows, cache2_dim), dtype=current_platform.fp8_dtype()
+            )
+            a2q_scale = qintermediate_cache2.new_empty(
+                (num_sorted_rows, cache2_dim // 128), dtype=torch.float32
+            )
+            _silu_mul_quant_sorted_kernel[ # 새로운 커널 호출
+                (triton.cdiv(num_sorted_rows, 16), cache2_dim // 128)
+            ](
+                intermediate_cache1, qintermediate_cache2, a2q_scale,
+                num_tokens_post_padded, cache2_dim, BLOCK_M=16, GROUP_K=128,
+            )
+        elif (
             activation == MoEActivation.SILU
             and self.activation_config.clamp_limit is None
             and self.quant_config.use_fp8_w8a8
@@ -555,6 +732,14 @@
         # the w13 pair: base GEMM on default stream, LoRA delta on aux,
         # join via .add_() into intermediate_cache3.
         def _base_w2_fn():
+            if use_sorted_a:
+                assert tma_down_config is not None
+                down_config = {
+                    **tma_down_config,
+                    "BLOCK_SIZE_M": config["BLOCK_SIZE_M"],
+                }
+            else:
+                down_config = config
             invoke_fused_moe_triton_kernel(
                 qintermediate_cache2,
                 w2,
@@ -568,7 +753,7 @@
                 not hooked
                 and self._apply_router_weight_in_gemm2(apply_router_weight_on_input),
                 1,
-                config,
+                down_config,
                 compute_type=compute_type,
                 use_fp8_w8a8=self.quant_config.use_fp8_w8a8,
                 use_int8_w8a8=self.quant_config.use_int8_w8a8,

1.2. SiLU, 곱셈, FP8 양자화 커널 퓨전

기존에는 SiLU 활성화, 행렬 곱셈, 그리고 FP8 양자화가 별도의 단계 또는 커널에서 처리될 수 있었습니다. 이번 PR에서는 이 세 가지 연산을 _silu_mul_quant_sorted_kernel이라는 단일 Triton 커널로 융합했습니다. 이는 연산 간의 오버헤드를 줄이고 데이터 재사용성을 높여 성능을 향상시킵니다.

--- a/vllm/model_executor/layers/fused_moe/experts/triton_moe.py
+++ b/vllm/model_executor/layers/fused_moe/experts/triton_moe.py
@@ -21,6 +21,7 @@
 from vllm.model_executor.layers.fused_moe.fused_moe import (
     _prepare_expert_assignment,
     _weights_host_resident,
+    get_moe_configs,
     invoke_fused_moe_triton_kernel,
     invoke_fused_moe_wna16_triton_kernel,
     try_get_optimal_moe_config,
@@ -60,10 +61,45 @@
     kInt8StaticTensorSym,
 )
 from vllm.platforms import current_platform
-from vllm.triton_utils import tl
+from vllm.triton_utils import tl, triton
 from vllm.utils.multi_stream_utils import maybe_execute_in_parallel
 
 
+@triton.jit
def _silu_mul_quant_sorted_kernel(
+    input_ptr, output_ptr, scale_ptr, num_tokens_post_padded_ptr,
+    intermediate: tl.constexpr, BLOCK_M: tl.constexpr, GROUP_K: tl.constexpr,
+):
+    row_start = tl.program_id(0) * BLOCK_M
+    if row_start >= tl.load(num_tokens_post_padded_ptr):
+        return
+    rows = row_start + tl.arange(0, BLOCK_M)
+    cols = tl.program_id(1) * GROUP_K + tl.arange(0, GROUP_K)
+    gate = tl.load(input_ptr + rows[:, None] * 2 * intermediate + cols[None, :])
+    up = tl.load(
+        input_ptr + rows[:, None] * 2 * intermediate + intermediate + cols[None, :]
+    )
+    gate = gate.to(tl.float32)
+    up = up.to(tl.float32)
+    # match silu_and_mul: round SiLU before multiplying the other half.
+    silu = tl.div_rn(gate, 1.0 + tl.extra.cuda.libdevice.exp(-gate))
+    silu = silu.to(input_ptr.dtype.element_ty).to(tl.float32)
+    result = (silu * up).to(input_ptr.dtype.element_ty).to(tl.float32)
+    scale_raw = tl.maximum(tl.max(tl.abs(result), 1), 1e-10) * (1.0 / 448.0)
+    scale = tl.math.exp2(tl.ceil(tl.log2(scale_raw)))
+    quantized = tl.minimum(tl.maximum(tl.div_rn(result, scale[:, None]), -448.0), 448.0)
+    tl.store(
+        output_ptr + rows[:, None] * intermediate + cols[None, :],
+        quantized.to(tl.float8e4nv),
+    )
+    tl.store(scale_ptr + rows * (intermediate // GROUP_K) + tl.program_id(1), scale)
+```

이 새로운 커널 `_silu_mul_quant_sorted_kernel`은 입력 텐서(`input_ptr`)를 받아 SiLU 활성화를 적용하고, 다른 절반과 곱한 후, FP8 (E4M3FN 형식)으로 양자화하여 출력 텐서(`output_ptr`)에 저장합니다. 또한, 양자화에 사용된 스케일 값(`scale_ptr`)도 함께 저장합니다. 이는 이전에는 여러 단계로 나뉘어 수행되던 작업을 단일 커널로 통합하여 GPU 커널 실행 오버헤드를 크게 줄입니다.

### 2. 설정 파일: TMA 최적화 파라미터 추가

H200 GPU에서 TMA를 활용하기 위한 최적화된 커널 파라미터들이 새로운 설정 파일로 추가되었습니다. 특히 `block_shape=[128, 128]`의 경우, `E=256, N=256` 및 `E=256, N=512`에 대해 `USE_TMA: true` 설정이 포함되어 있습니다.

```diff
--- a/vllm/model_executor/layers/fused_moe/configs/E=256,N=256,device_name=NVIDIA_H200,dtype=fp8_w8a8,block_shape=[128,128]_down.json
+++ b/vllm/model_executor/layers/fused_moe/configs/E=256,N=256,device_name=NVIDIA_H200,dtype=fp8_w8a8,block_shape=[128,128]_down.json
@@ -1,22 +1,22 @@
 {
     "512": {"USE_TMA": false},
     "1024": {
-        "BLOCK_SIZE_M": 64,
-        "BLOCK_SIZE_N": 128,
-        "BLOCK_SIZE_K": 128,
-        "GROUP_SIZE_M": 64,
-        "num_warps": 4,
-        "num_stages": 3,
-        "USE_TMA": true
+        "BLOCK_SIZE_M": 64, "BLOCK_SIZE_N": 128, "BLOCK_SIZE_K": 128,
+        "GROUP_SIZE_M": 64, "num_warps": 4, "num_stages": 3,
+        "USE_TMA": true
     },
     "3072": {
-        "BLOCK_SIZE_M": 64,
-        "BLOCK_SIZE_N": 128,
-        "BLOCK_SIZE_K": 128,
-        "GROUP_SIZE_M": 64,
-        "num_warps": 4,
-        "num_stages": 3,
-        "USE_TMA": true
+        "BLOCK_SIZE_M": 64, "BLOCK_SIZE_N": 128, "BLOCK_SIZE_K": 128,
+        "GROUP_SIZE_M": 64, "num_warps": 4, "num_stages": 3,
+        "USE_TMA": true
     },
     "4096": {"USE_TMA": false}
 }
+---
+--- /dev/null
++++ b/vllm/model_executor/layers/fused_moe/configs/E=256,N=512,device_name=NVIDIA_H200,dtype=fp8_w8a8,block_shape=[128,128]_down.json
@@ -0,0 +1,23 @@
+{
+    "512": {"USE_TMA": false},
+    "1023": {"USE_TMA": false},
+    "1024": {
+        "BLOCK_SIZE_M": 64, "BLOCK_SIZE_N": 128, "BLOCK_SIZE_K": 128,
+        "GROUP_SIZE_M": 64, "num_warps": 4, "num_stages": 3,
+        "USE_TMA": true
+    },
+    "3072": {
+        "BLOCK_SIZE_M": 64, "BLOCK_SIZE_N": 128, "BLOCK_SIZE_K": 128,
+        "GROUP_SIZE_M": 64, "num_warps": 4, "num_stages": 3,
+        "USE_TMA": true
+    },
+    "4096": {"USE_TMA": false}
+}

이 설정들은 특정 hidden_dim (N) 및 intermediate_dim (E) 조합에 대해 USE_TMA를 활성화하여, H200 GPU의 TMA 하드웨어 기능을 최대한 활용하도록 Triton 커널의 동작 방식을 조정합니다.

왜 이게 좋은가?

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

하드웨어 기능 활용 극대화 (TMA): NVIDIA H200 GPU는 TMA라는 강력한 하드웨어 기능을 제공합니다. TMA는 GPU 코어와 독립적으로 데이터를 메모리에서 로드하고 저장하는 작업을 병렬로 처리하여, 데이터 이동 병목 현상을 완화합니다. MoE 모델의 다운 프로젝션 단계는 대규모 행렬 곱셈을 포함하며, 이는 데이터 이동에 민감합니다. 중간 텐서를 전문가 정렬 순서로 유지함으로써 TMA를 효과적으로 활용할 수 있게 되었고, 이는 메모리 대역폭 활용률을 크게 높여줍니다.
커널 퓨전 (Kernel Fusion): 여러 개의 작은 커널을 하나의 큰 커널로 합치는 것은 GPU 컴퓨팅에서 성능 향상을 위한 고전적이면서도 강력한 기법입니다. 각 커널 호출에는 오버헤드가 발생하며, 커널 간 데이터 전송 시 캐시 효율성 저하가 발생할 수 있습니다. SiLU, 곱셈, FP8 양자화를 단일 Triton 커널로 융합함으로써 이러한 오버헤드를 제거하고, 중간 결과를 GPU 레지스터 내에서 직접 처리하여 데이터 재사용성을 극대화했습니다. 이는 연산 속도를 직접적으로 향상시킵니다.
FP8 정밀도 활용: FP8 정밀도는 FP16이나 BF16에 비해 메모리 사용량을 절반으로 줄이고 연산 속도를 높일 수 있습니다. 이번 최적화는 FP8 정밀도를 사용하는 MoE 모델에 초점을 맞춰, 최신 하드웨어의 FP8 연산 능력을 최대한 활용하도록 설계되었습니다.

성능 향상:

PR 설명에 따르면, 이러한 최적화를 통해 H200 GPU에서 GLM 5.3 모델의 종단 간(end-to-end) 처리량이 6.0%에서 7.6%까지 향상되었습니다. 제공된 벤치마크 결과에서도 이를 확인할 수 있습니다:

기존 (main): Request throughput: 24.61 req/s, Total token throughput: 50424.10 tok/s
개선 후 (Now): Request throughput: 26.99 req/s, Total token throughput: 55305.37 tok/s

이는 약 9.6%의 요청 처리량 증가와 9.7%의 토큰 처리량 증가를 보여주며, PR 설명의 목표치 범위 내에 있습니다.

일반적인 교훈:

하드웨어 특화 최적화: 최신 GPU 아키텍처(예: H200의 TMA)의 기능을 이해하고 이를 적극적으로 활용하는 것이 성능 향상의 핵심입니다.
데이터 레이아웃의 중요성: 데이터의 메모리 레이아웃(예: 전문가 정렬 순서)은 하드웨어 가속기(TMA 등)의 효율성에 지대한 영향을 미칩니다. 때로는 데이터 순서를 유지하기 위한 추가 비용이 있더라도, 전체적인 성능 향상을 위해 이를 감수해야 할 수 있습니다.
커널 퓨전의 지속적인 가치: 연산 그래프를 분석하여 융합 가능한 연산들을 찾아 단일 커널로 만드는 것은 여전히 GPU 성능 최적화의 중요한 전략입니다.
정밀도 활용: FP8과 같은 저정밀도 연산은 메모리 및 연산 효율성을 크게 향상시킬 수 있으며, 이를 지원하는 하드웨어와 라이브러리(vLLM)를 활용하는 것이 중요합니다.

결론

이번 vLLM PR은 H200 GPU 환경에서 FP8 MoE 모델의 다운 프로젝션 성능을 획기적으로 개선했습니다. TMA를 활용하기 위한 데이터 정렬 및 새로운 Triton 커널의 도입은 MoE 모델 추론의 병목 현상을 효과적으로 해결했으며, 이는 LLM 서비스의 전반적인 효율성을 높이는 데 크게 기여할 것입니다. 이러한 지속적인 최적화 노력은 vLLM이 최첨단 LLM을 위한 최고의 추론 엔진 중 하나로 자리매김하는 데 중요한 역할을 하고 있습니다.

참고 자료

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

PR Analysis 의 다른글

이전글 [vllm] vLLM 스케줄러 최적화: Cache-aware Admission으로 KV 캐시 효율 극대화하기
현재글 : [vllm] vLLM, H200 GPU를 위한 FP8 MoE 다운 프로젝션 최적화: TMA 활용 및 커널 퓨전으로 7.6% 처리량 향상
다음글 [onnxruntime] ONNX Runtime WebGPU: NVIDIA GPU를 위한 FlashAttention 루프 언롤링 최적화

[vllm] vLLM, H200 GPU를 위한 FP8 MoE 다운 프로젝션 최적화: TMA 활용 및 커널 퓨전으로 7.6% 처리량 향상

들어가며

코드 분석

1. `triton_moe.py`: TMA 활용 및 커널 퓨전 도입

왜 이게 좋은가?

결론

참고 자료

댓글

관련 포스트

PR Analysis 의 다른글

들어가며

코드 분석

1. triton_moe.py: TMA 활용 및 커널 퓨전 도입

왜 이게 좋은가?

결론

참고 자료

댓글

관련 포스트

PR Analysis 의 다른글

1. `triton_moe.py`: TMA 활용 및 커널 퓨전 도입