Fusion torch.compile 패스

Fusion torch.compile 패스

vLLM은 컴파일 타임에(커스텀 torch.compile Inductor 패스를 통해) 커널/연산자 퓨전 세트를 적용해 최적화를 모델 정의와 분리하고 모델 코드의 레이어 추상화를 깨지 않게 합니다. 이 퓨전들은 PassConfig의 필드로 제어되며, 적절한 최적화 레벨에서 자동으로 활성화됩니다.

출처: 문서

본문

빠른 참조 (Quick Reference)

아래 표는 각 퓨전을 제어 플래그/구성 노브, 퓨전하는 연산, 기본 활성화 레벨, 대략적인 속도 향상에 매핑합니다. Fullgraph 열은 퓨전이 전체 모델 그래프를 봐야 하는지(Inductor partition 또는 splitting_ops=[]를 통해) 나타내고, 마지막 열은 퓨전이 모든 num_tokens에 대해 활성화되는지 아니면 낮은/높은 끝에서만 활성화되는지를 나타냅니다.

정보: 속도 향상은 정확한 모델, 배치 크기, 하드웨어에 크게 의존합니다. 수동으로 성능을 튜닝한다면 항상 퓨전 있고 없고로 정확한 사용 사례를 벤치마크해서 영향을 검증하세요.

퓨전 PassConfig 플래그 퓨전되는 연산 기본 레벨 E2E 속도 향상 Fullgraph num_tokens
AllReduce + RMSNorm fuse_allreduce_rms All-reduce → RMSNorm (+residual_add) (→ quant) O2 (Hopper/Blackwell + TP > 1) 5-20% No 낮음
Attention + Quant fuse_attn_quant Attention 출력 → FP8/NVFP4 quant 기본 off 3-7% Yes 항상
MLA Attention + Quant fuse_attn_quant MLA Attention 출력 → FP8/NVFP4 quant 기본 off TBD Yes 항상
RoPE + KV-Cache Update fuse_rope_kvcache Rotary embedding → KV cache write O2 (ROCm/AITER 전용) 2-4% No 낮음
QK Norm + RoPE enable_qk_norm_rope_fusion Q/K RMSNorm → rotary embedding 기본 off 2-3% No 낮음
Sequence Parallelism enable_sp AllReduce → ReduceScatter + AllGather 기본 off AsyncTP의 사전조건 Yes 높음
AsyncTP GEMM + collective fuse_gemm_comms GEMM → reduce-scatter / all-gather → GEMM 기본 off 7-10% Yes 높음
RMSNorm + Quant fuse_norm_quant RMSNorm (+residual add) → FP8/FP4 quant O1 (조건부) 1-4% No 항상
SiLU+Mul + Quant fuse_act_quant SiLU+Mul activation → FP8/FP4 quant O1 (조건부) 1-4% No 항상
RMSNorm + Padding fuse_act_padding Residual add + RMSNorm → padding O1 (ROCm/AITER 전용) TBD No 항상
MLA Dual RMSNorm fuse_mla_dual_rms_norm Paired Q + KV RMSNorm (+ FP8 quant) → 1 kernel O1 (ROCm/AITER 전용) 1-2% No 항상

지원 매트릭스 (Support Matrix)

아래 표는 각 퓨전이 각 플랫폼에서 지원하는 양자화 방식을 나열합니다. —는 그 플랫폼에서 퓨전을 사용할 수 없음을 뜻합니다. 최신·진행 중 작업은 추적 이슈 #36066에 있습니다.

퓨전 SM100 (Blackwell) SM90 (Hopper) SM89 (Ada) SM80 (Ampere) ROCm
fuse_allreduce_rms FP16/BF16, FP8 static, NVFP4 FP16/BF16, FP8 static — — —
fuse_attn_quant * FP8 static*, NVFP4* FP8 static* FP8 static* — FP8 static*
fuse_attn_quant (MLA)* FP8 static*, FP8 per-group*, NVFP4* FP8 static*, FP8 per-group* FP8 static*, FP8 per-group* — FP8 static* (미테스트)
fuse_rope_kvcache — — — — FP16/BF16
enable_qk_norm_rope_fusion FP16/BF16 FP16/BF16 FP16/BF16† FP16/BF16† —
enable_sp FP16/BF16, FP8 static† FP16/BF16, FP8 static FP16/BF16† FP16/BF16† —
fuse_gemm_comms FP16/BF16, FP8 static† FP16/BF16, FP8 static FP16/BF16† FP16/BF16† —
fuse_norm_quant FP8 static, FP8 per-token, FP8 per-group FP8 static, FP8 per-token, FP8 per-group FP8 static, FP8 per-token, FP8 per-group — FP8 static, FP8 per-token, FP8 per-group
fuse_act_quant FP8 static, NVFP4 FP8 static, FP8 per-group (128/64) FP8 static, FP8 per-group (128/64) — FP8 per-group
fuse_act_padding — — — — FP16/BF16
fuse_mla_dual_rms_norm — — — — BF16
  • fuse_attn_quant 지원은 사용 중인 attention 백엔드에 달려 있습니다. 모든 백엔드가 퓨전 양자화 출력을 지원하진 않습니다. 백엔드별 세부 사항은 fuse_attn_quant 절에서 확인하세요.

† enable_sp와 fuse_gemm_comms는 오늘날 SM90에서만 자동 구성됩니다. 다른 아키텍처는 PassConfig.sp_min_token_num을 명시적으로 설정해야 합니다. SM100 지원에는 VLLM_DISABLED_KERNELS=FlashInferFP8ScaledMMLinearKernel 설정도 필요합니다.

퓨전 활성화/비활성화 (Enabling / Disabling Fusions)

퓨전은 CompilationConfig 안에 중첩된 PassConfig를 통해 노출됩니다:

from vllm import LLM
from vllm.config import CompilationConfig, PassConfig

llm = LLM(
    model="...",
    optimization_level=2, # Default optimization level
    compilation_config=CompilationConfig(
        pass_config=PassConfig(
            fuse_norm_quant=True,
            fuse_act_quant=True,
            fuse_allreduce_rms=False,  # disable a specific fusion
        )
    ),
)

퓨전은 모든 vllm ... 명령에서 명령줄 플래그로도 활성화할 수 있습니다:

# Enable O2 defaults, but turn off allreduce fusion
vllm serve meta-llama/Llama-3.1-8B-Instruct -O2 -cc.pass_config.fuse_allreduce_rms=False

# The above is equivalent to the more verbose:
vllm serve meta-llama/Llama-3.1-8B-Instruct -O2 --compilation-config '{"pass_config": {"fuse_allreduce_rms": false}}'

# Same syntax in other commands, e.g. vllm bench:
vllm bench latency --model=meta-llama/Llama-3.1-8B-Instruct -O2 -cc.pass_config.fuse_allreduce_rms=False

사용자가 명시적으로 설정한 필드는 항상 최적화 레벨 기본값보다 우선합니다.

퓨전 세부 사항 (Fusion Details)

AllReduce + RMSNorm (fuse_allreduce_rms)

경고: TP+DP 및 TP+PP 조합은 현재 손상돼 있습니다(#34458와 #35426). FlashInfer가 설치된 NVIDIA Hopper(SM90)·Blackwell(SM100)에서만 지원됩니다.

무엇을 퓨전하나: 텐서 병렬 all-reduce collective와 그 뒤 residual add, RMSNorm, 선택적으로 양자화 단계를 단일 FlashInfer / TRT-LLM 통신 커널로 퓨전합니다. 이 퓨전은 작은 num_tokens에서만 이득이 있으므로 더 낮은 컴파일 범위에서만 수행됩니다.

다루는 패턴:

  • AllReduce → RMSNorm(+residual_add): FlashInfer를 쓴 CUDA sm90+
  • AllReduce → RMSNorm(+residual_add) → FP8 static quant: FlashInfer를 쓴 CUDA sm90+
  • AllReduce → RMSNorm(+residual_add) → NVFP4 dynamic quant: FlashInfer를 쓴 CUDA sm100+

퓨전 커널을 사용하는 최대 텐서 크기는 하드웨어 의존적이며(SM90/SM100에서 TP=2면 64MB) PassConfig.fi_allreduce_fusion_max_size_mb로 구성할 수 있습니다.

코드 위치:

Attention + Quantization (fuse_attn_quant)

정보: fuse_attn_quant는 현재 어떤 최적화 레벨에서도 기본 활성화되지 않으며 명시적으로 설정해야 합니다. 전체 모델 그래프가 보여야 합니다(Inductor partition 또는 splitting_ops=[]).

무엇을 퓨전하나: attention 출력 양자화를 attention 계산 직후에 퓨전해 attention 출력의 full-precision 메모리 왕복을 제거합니다. 이 퓨전은 표준 Attention과 (DeepSeek-V2/V3/R1 모델이 쓰는) MLAAttention을 모두 지원합니다. 다루는 패턴:

Attention → FP8 static quant:

  • TRITON_ATTN: CUDA, ROCm
  • FLASHINFER: FlashInfer가 설치된 CUDA sm100+
  • ROCM_ATTN: ROCm
  • ROCM_AITER_UNIFIED_ATTN: AITER를 쓴 ROCm

Attention → NVFP4 dynamic quant:

  • FLASHINFER: FlashInfer가 설치된 CUDA sm100+

MLAAttention → FP8 static, FP8 per-group, NVFP4 dynamic quant

MLA 퓨전은 그래프 레벨에서 unified_mla_attention_with_output 연산에서 동작하며 모든 MLA decode·prefill 백엔드 조합과 함께 동작합니다. 표준 Attention 백엔드(커널이 FP8 출력을 직접 씀)와 달리, 현재 어떤 MLA prefill·decode 백엔드도 직접 FP8/FP4 출력을 지원하지 않습니다. 퓨전은 중간 버퍼에 쓰고 별도 단계에서 양자화하므로 아직 메모리 왕복 제거는 없습니다.

정보: MLA attention 퓨전은 아직 측정 가능한 속도 향상을 기대하지 않습니다. MLA prefill/decode 커널이 직접 FP8/FP4 출력을 지원하면 개선됩니다.

다른 attention 백엔드는 아직 퓨전 출력 양자화를 지원하지 않습니다.

코드 위치:

RoPE + KV-Cache Update (fuse_rope_kvcache)

정보: ROCm/AITER 전용. NVIDIA CUDA나 CPU에선 사용할 수 없습니다. AITER 융합 커널 성능 문제로 퓨전은 기본적으로 num_tokens ≤ 256에서만 활성화됩니다. 이 임계값은 PassConfig.rope_kvcache_fusion_max_token_num으로 구성할 수 있습니다.

무엇을 퓨전하나: rotary positional embedding 커널을 KV-cache scatter/write와 단일 커널로 퓨전해 key·value 텐서의 별도 읽기·쓰기를 피합니다.

요구 사항: AITER로 활성화된 AMD ROCm, 활성 rotary_embedding 커스텀 op(자동), 그리고 kv_cache update op가 그래프에 보여야 합니다 (Inductor graph partition을 쓰거나 splitting_ops에서 제거). 이 조건이 설정되면 최적화 레벨 O1 이상에서 퓨전이 자동 활성화됩니다.

코드 위치:

Sequence Parallelism (enable_sp)

무엇을 퓨전하나: All-reduce collective를 reduce-scatter + 로컬 RMSNorm + all-gather로 교체해 시퀀스 차원을 TP 랭크에 분할합니다. 이는 그래프를 재구성해 이후 AsyncTP 패스가 주변 GEMM과 reduce-scatter / all-gather를 퓨전할 수 있게 합니다.

Sequence Parallelism 자체는 성능을 직접 개선하지 않습니다. 이는 AsyncTP 패스(fuse_gemm_comms)의 사전 조건입니다. SP는 기기 성능과 모델 hidden_size를 기반으로 자동 구성되는 최소 토큰 임계값 위에서만 적용됩니다. 현재 hidden_size >= 8192인 모델의 H100/SM90에서만 활성화됩니다. 임계값은 PassConfig.sp_min_token_num으로 구성할 수 있습니다.

일반적인 변환:

Input → AllReduce → RMSNorm → Output
becomes:
Input → ReduceScatter → local RMSNorm → AllGather → Output

다루는 패턴:

  • 첫 블록: AllReduce → RMSNorm → ReduceScatter → RMSNorm → AllGather
  • 중간 블록: AllReduce → fused_add_RMSNorm → ReduceScatter → fused_add_RMSNorm → AllGather
  • 둘 다 선택적 → FP8 static quant 접미사 포함

요구 사항: use_inductor_graph_partition=True 또는 크기가 tensor_parallel_size로 나누어떨어지는 정적 크기의 piecewise 컴파일.

지원 하드웨어: NVIDIA CUDA에서만 테스트됨, ROCm에서 동작할 가능성 있음. FP8 all-gather는 sm90+가 필요합니다.

코드 위치:

AsyncTP GEMM + Collective Overlap (fuse_gemm_comms)

정보: enable_sp=True가 필요합니다(자동 활성화). Sequence Parallelism이 적용되지 않았다면 이 패스는 no-op입니다.

무엇을 퓨전하나: Sequence Parallelism이 그래프를 변환한 뒤, torch.ops.symm_mem 대칭 메모리 프리미티브로 주변 reduce-scatter(출력 투영)와 all-gather(입력 투영)를 GEMM 커널과 퓨전해 통신과 계산을 겹칩니다. 이 오버랩은 큰 num_tokens에서만 이득이 있으므로 퓨전(및 앞선 SP)은 PassConfig.sp_min_token_num 위의 더 높은 컴파일 범위에서만 수행됩니다.

다루는 패턴:

  • GEMM → reduce-scatter → fused_matmul_reduce_scatter
  • all-gather → GEMM → all_gather_matmul
  • 두 패턴의 FP8 scaled 변형

지원 하드웨어: 대칭 메모리(torch.distributed._symmetric_memory) 지원을 갖춘 NVIDIA CUDA.

B200에서 fp8 FlashInfer scaled MM 패턴 매칭은 지원되지 않으므로 비활성화해야 합니다(#27893):

VLLM_DISABLED_KERNELS=FlashInferFP8ScaledMMLinearKernel ...

코드 위치:

QK Norm + RoPE (enable_qk_norm_rope_fusion)

정보: rotary positional embedding 전에 Q·K에 per-head RMSNorm을 적용하는 모델(예: Qwen)에만 해당합니다. H100에서의 성능 문제로 어떤 최적화 레벨에서도 기본 활성화되지 않습니다: #34391

무엇을 퓨전하나: 분할 QKV → reshape → Q/K RMSNorm → reshape → rotary embedding 시퀀스를 단일 fused_qk_norm_rope CUDA 커널로 퓨전합니다.

# Unfused:
q, k, v = split(qkv)
q_norm = rms_norm(q.view(heads))
k_norm = rms_norm(k.view(kv_heads))
q_rope, k_rope = rotary_embedding(q_norm, k_norm, ...)

# Fused:
fused_qk_norm_rope(qkv, ...)

지원 하드웨어: CUDA(sm80+) 전용, sm90·sm100에서만 테스트.

코드 위치:

RMSNorm + Quantization (fuse_norm_quant)

경고: NVIDIA에서 Inductor는 실제로 우리 커스텀 CUDA 커널보다 더 빠른 융합 커널을 생성합니다. 따라서 이 퓨전은 rms_norm이나 quant_fp8 중 하나가 커스텀 커널을 사용할 때만 활성화됩니다.

무엇을 퓨전하나: 커스텀 rms_norm / fused_add_rms_norm 연산과 이후 양자화를 단일 융합 커널로 결합해 full-precision 활성화 텐서의 중간 읽기·쓰기를 제거합니다. 두 변형이 퓨전됩니다:

  • 일반 RMSNorm + quant: rms_norm(x) → quant_fp8(y)
  • Fused-add RMSNorm + quant: fused_add_rms_norm(x, residual) → quant_fp8(y) — residual도 in-place로 갱신합니다.

AITER 퓨전은 현재 vllm.compilation.passes.fusion.rocm_aiter_fusion의 별도 패스에 있습니다.

지원 양자화 방식/하드웨어 조합:

  • FP8 static per-tensor: CUDA & HIP 커널
  • FP8 dynamic per-token: CUDA & HIP 커널, AITER
  • FP8 dynamic per-token-group (128/64): CUDA & HIP 커널, AITER

코드 위치:

SiLU+Mul + Quantization (fuse_act_quant)

경고: fuse_norm_quant와 동일 — NVIDIA에서 Inductor는 커스텀 op보다 빠른 융합 커널을 생성합니다. 이 퓨전은 silu_and_mul이나 quant_fp8 중 하나가 커스텀 커널을 사용할 때, 또는 NVFP4 양자화 모델(FP4 quant가 항상 커스텀 op인 경우)에서만 활성화됩니다.

무엇을 퓨전하나: silu_and_mul gate-up 프로젝션 활성화와 이후 양자화를 단일 커널로 퓨전해 full-precision post-activation 텐서의 실체화를 피합니다.

AITER 퓨전은 vllm.compilation.passes.fusion.rocm_aiter_fusion의 별도 패스에 있습니다.

지원 양자화 방식/하드웨어 조합:

  • FP8 static per-tensor: CUDA & HIP 커널
  • FP8 dynamic per-group (128/64): CUDA 커널 (sm89+, sm100+에서 DeepGemm 사용 시 비활성)
  • NVFP4 dynamic: FlashInfer를 쓴 CUDA sm100+ 전용
  • FP8 per-token-group (128): ROCm AITER 전용

코드 위치:

RMSNorm + Padding (fuse_act_padding)

정보: ROCm/AITER 전용. GPT-OSS 모델 대상.

무엇을 퓨전하나: residual add + RMSNorm과 이후 hidden 차원을 다운스트림 AITER Triton GEMM 커널이 요구하는 배수로 패딩하는 패딩 연산을 퓨전합니다.

요구 사항: AITER RMSNorm으로 활성화된 AMD ROCm. hidden size가 2880이고 AITER Triton GEMM이 활성화되지 않았을 때 최적화 레벨 O1 이상에서 기본 활성화됩니다.

코드 위치:

MLA Dual RMSNorm (fuse_mla_dual_rms_norm)

정보: ROCm/AITER 전용. DeepSeek-V3 / Kimi-K2 MLA attention 대상.

참고: rms_norm의 네이티브 구현을 사용할 때(현재 CUDA·ROCm 기본값) Inductor의 내장 퓨전이 이 norm 병합을 자동으로 처리합니다. 이 명시적 패스는 Inductor가 스스로 퓨전할 수 없는 AITER의 커스텀 rms_norm op가 활성화된 경우를 대상으로 합니다.

무엇을 퓨전하나: MLA attention의 쌍을 이루는 q_a_layernorm와 kv_a_layernorm RMS norm 연산을 AITER를 통한 단일 fused_qk_rmsnorm HIP 커널 호출로 퓨전해 MLA 레이어당 커널 실행 오버헤드를 2회에서 1회로 줄입니다.

# Unfused:
q_c, kv_lora = split(projected, [q_dim, kv_dim])
kv_c, k_pe   = split(kv_lora,  [kv_c_dim, k_pe_dim])
q_c  = rms_norm(q_c,  q_weight,  eps)
kv_c = rms_norm(kv_c, kv_weight, eps)

# Fused:
q_c, kv_lora = split(projected, [q_dim, kv_dim])
kv_c, k_pe   = split(kv_lora,  [kv_c_dim, k_pe_dim])
q_normed, kv_normed = fused_mla_dual_rms_norm(
    q_c, q_weight, kv_c, kv_weight, eps1, eps2)

요구 사항: AITER로 활성화된 AMD ROCm. AITER 사용 가능 시 최적화 레벨 O1 이상에서 기본 활성화됩니다.

FP8 attention 변형(per-token quant): per-token FP8 q_b_proj를 쓰면 q 잠재 변수만 FP8 양자화되고 kv는 bf16으로 남습니다. RocmAiterRMSNormQuantFusionPass가 먼저 q 쪽을 rocm_aiter_rmsnorm_fused_dynamic_quant로 접어 kv를 일반 rms_norm으로 남깁니다 — 위의 대칭 패턴을 깨뜨립니다. 그런 다음 같은 패스가 이 비대칭 쌍을 매칭해 fused_mla_dual_rms_norm_per_token_quant로 낮춥니다.

# Unfused (q norm+quant fused; kv still plain rms_norm):
q_c, kv_lora = split(projected, [q_dim, kv_dim])
kv_c, k_pe   = split(kv_lora,  [kv_c_dim, k_pe_dim])
q_fp8, q_scale = rocm_aiter_rmsnorm_fused_dynamic_quant(q_c, q_weight, eps, fp8)
kv_normed      = rms_norm(kv_c, kv_weight, eps)          # bf16

# Fused:
q_c, kv_lora = split(projected, [q_dim, kv_dim])
kv_c, k_pe   = split(kv_lora,  [kv_c_dim, k_pe_dim])
q_fp8, q_scale, kv_normed = fused_mla_dual_rms_norm_per_token_quant(
    q_c, q_weight, kv_c, kv_weight, eps1, eps2)

코드 위치:

함께 보기 (See Also)

더 알아보기 (Learn more)