[vllm] Rubin 텐서 코어로 LUT-B MoE 가중치를 직접 디코딩하기: MXFP8 활성화 기반 네이티브 커널과 에뮬레이션

PR 링크: vllm-project/vllm#62759 상태: Merged | 변경: +1513 / -0

들어가며

MoE 모델에서 메모리의 대부분은 routed expert 가중치가 차지합니다. GLM-4.7-Flash를 BF16으로 올리면 56.3 GiB가 필요합니다. 가중치를 3-bit 수준으로 줄이면 메모리는 크게 줄지만, 보통은 GEMM 직전에 dequantize 커널을 따로 돌려야 하고 이 비용이 이득을 상쇄합니다.

NVIDIA Rubin(SM 10.7)은 이 문제를 하드웨어로 풉니다. tcgen05.mma.kind::mxf8f6f4.block_scale.decompress::lut::b 명령은 B 피연산자를 lookup table(LUT) 인덱스 형태 그대로 받고, 텐서 코어 안에서 codebook으로 디코딩합니다. 다만 이 명령의 A 피연산자는 FP8/FP6/FP4만 받습니다.

이전 시도(#51880)는 LUT-B 가중치에 BF16 활성화를 조합해 에뮬레이션했기 때문에 Rubin에서 네이티브로 돌릴 수 없었고, fused Triton decode GEMM과 dequantize fallback까지 필요했습니다. 이 PR은 활성화를 MXFP8(E4M3 + 1x32 단위 UE8M0 scale)로 바꿔 하드웨어 제약에 맞추고, 다음 두 경로를 제공합니다.

• Native (LutBSm107Experts + ops.lut_b_moe_gemm): Rubin 전용 grouped GEMM
• Emulation (LutBEmulationTritonExperts): 다른 모든 GPU에서 같은 수치를 재현해 정확도 검증

포맷: 3-bit 인덱스 + E4M3 codebook + UE8M0 exponent

가중치 하나는 8개 엔트리짜리 E4M3 codebook을 가리키는 3-bit 인덱스입니다. codebook은 8 output channel × 64 K 타일 하나가 공유하고, 타일마다 UE8M0 exponent가 하나씩 붙습니다. 논리적 비트수는 약 3.14 bits/weight입니다.

exponent를 추가한 이유가 핵심입니다. 이 값은 명령어가 원래 요구하는 per-32-K B scale 슬롯에 그대로 매핑되므로 런타임 비용이 0입니다. exponent가 없으면 크기가 작은 타일이 E4M3 subnormal 영역에 떨어져 표현력을 잃습니다. Qwen3.6-35B-A3B expert 기준으로 relative weight MSE는 exponent가 없을 때 0.0440.045, 있을 때 0.0380.040이었습니다. 하드웨어가 어차피 읽는 필드를 정확도 개선에 재활용한 셈입니다.

코드 분석

CMakeLists.txt: 아키텍처/툴체인 게이팅

# Rubin LUT-B MoE GEMM; tcgen05 decompress::lut::b needs CUDA 13.4.
if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 13.4)
  cuda_archs_loose_intersection(LUT_B_MOE_ARCHS "10.7f" "${CUDA_ARCHS}")
endif()
if(LUT_B_MOE_ARCHS)
  set(LUT_B_MOE_SRC "csrc/moe/lut_b_moe_sm107.cu")
  ...
endif()

if(LUT_B_MOE_ARCHS)
  target_compile_definitions(_moe_C_stable_libtorch PRIVATE
    VLLM_ENABLE_LUT_B_MOE=1)
endif()

CUDA 13.4 이상이면서 10.7f 아키텍처가 타깃에 포함될 때만 소스를 빌드하고 VLLM_ENABLE_LUT_B_MOE를 정의합니다. 그래서 다른 환경의 빌드에는 전혀 영향을 주지 않습니다. 새 하드웨어 전용 커널을 메인 빌드에 넣을 때 쓸 수 있는 깔끔한 패턴입니다.

csrc/moe/lut_b_moe_sm107.cu: 네이티브 grouped GEMM

타일과 메모리 예산

constexpr int kBlockM = 128;
constexpr int kBlockN = 128;
constexpr int kBlockK = 128;
constexpr int kStages = 3;

constexpr int kABytes = kBlockM * kBlockK;               // E4M3 activations
constexpr int kMetaBytes = 512;                          // one 32x128b copy
constexpr int kBIndexBytes = kBlockN * kBlockK * 3 / 8;  // 3-bit indices
constexpr int kBTileBytes = kBIndexBytes + 2 * kMetaBytes;
...
static_assert(kStages * kStageBytes >= kBlockM * kBlockN * 2,
              "the epilogue reuses the pipeline buffers");

128N×128K B 타일은 인덱스 6144 B, B scale 512 B, codebook 512 B로 약 7 KB입니다. 같은 타일을 BF16으로 두면 32 KB이므로 shared memory와 global 대역폭 부담이 4배 이상 줄어듭니다. static_assert는 epilogue가 파이프라인 버퍼를 재사용한다는 전제를 컴파일 타임에 검증합니다.

가중치는 로드 시점에 smem 이미지로 미리 배치

// B indices, scales and codebooks are already in their smem image.
uint8_t* b_dst = st + kABytes + kMetaBytes;
uint8_t const* b_src = b_tiles + int64_t(k_tile) * kBTileBytes;
for (int i = tid; i < kBTileBytes / 16; i += kThreads)
  cp_async_16(b_dst + i * 16, b_src + i * 16, true);

Python 쪽 lut_b_to_sm107_tiles가 모델 로드 때 가중치를 텐서 코어가 기대하는 shared memory 레이아웃 그대로 재배열해 둡니다. 덕분에 커널은 재배열 없이 16 B cp.async로 그대로 복사만 하면 됩니다. 패딩을 포함한 상주 비용은 3.5 bits/weight입니다. 비용이 큰 레이아웃 변환을 1회성 오프라인 작업으로 옮긴 전형적인 사례입니다.

A는 gather하면서 canonical no-swizzle 레이아웃으로

int32_t const sid = sorted_token_ids[m_block * kBlockM + tid];
bool const row_valid = sid < num_valid_tokens;
...
uint8_t* a_dst = st + (tid / 8) * 1024 + (tid % 8) * 16;
for (int chunk = 0; chunk < kBlockK / 16; ++chunk)
  cp_async_16(a_dst + chunk * 128, a_src + chunk * 16, row_valid);

스레드 하나가 M 타일의 한 행을 맡습니다. moe_align_block_size로 정렬된 sorted_token_ids를 통해 토큰 행을 gather하고, 8행×16 B core matrix 단위의 K-major 레이아웃으로 배치합니다. cp.async의 src-size 인자에 valid ? 16 : 0을 넘겨, 패딩 행은 분기 없이 0으로 채웁니다.

핵심: 두 번의 K64 LUT MMA와 B collector 재사용

// The two K64 halves of a K128 tile share one 48-byte compressed B row per
// output channel; the first fills the B collector and the second reuses it.
"tcgen05.mma.cta_group::1.kind::mxf8f6f4.block_scale.decompress::lut::b."
"block32.collector::b::fill [%0], %1, %2, [%3], %4, [%5], [%6], p;"
...
"block32.collector::b::lastuse [%0], %1, %2, [%3], %4, [%5], [%6], p;"

K128 타일 하나는 두 번의 K64 MMA로 처리됩니다. 첫 번째 MMA는 collector::b::fill로 압축된 B 행(채널당 48 B)을 collector에 적재하고, 두 번째 MMA는 lastuse로 이를 재사용합니다. 두 번째 절반에서 사용할 scale 바이트는 instruction descriptor로 지정합니다.

uint32_t const idesc_hi = idesc | (2u << 29) | (2u << 4);

scale과 codebook은 tcgen05.cp ... 32x128b.warpx4로 TMEM에 복사되어 4개의 32-lane partition에 복제됩니다. MMA 발행은 스레드 하나가 담당하고, 나머지 warp는 3-stage cp.async 로딩과 epilogue를 맡습니다. 각 stage가 끝나면 mbarrier로 완료를 확인하고, 생산된 데이터는 fence.proxy.async.shared::cta로 텐서 코어(async proxy)에 보이게 만듭니다.

Expert parallelism 처리

if (expert < 0) {
  // Tokens routed to an expert that is not local to this rank.
  if (row_valid) { ... dst[i] = make_uint4(0, 0, 0, 0); }
  return;
}

로컬 rank에 없는 expert로 라우팅된 토큰은 출력을 0으로 채우고 바로 종료합니다. 표준 expert_map 처리와 맞물리는 방식이지만, PR 설명에 따르면 실제 하드웨어에서는 아직 검증하지 않았습니다.

Emulation 경로

LutBEmulationTritonExperts는 로드 시 가중치를 한 번 dequantize합니다. 그리고 각 GEMM 전에 활성화를 MXFP8로 quantize→dequantize한 뒤 기존 BF16 TritonExperts를 실행합니다. 메모리나 속도 이점은 없고 정확도 검증 전용입니다. 이 경로 덕분에 Rubin이 없는 CI와 개발 환경에서도 동일한 수치를 재현할 수 있습니다. hidden/intermediate 크기가 128의 배수가 아니거나 bias, LoRA가 있으면 네이티브 경로 대신 자동으로 이 경로를 사용합니다.

왜 이게 좋은가

정확도 (GLM-4.7-Flash, GSM8K 5-shot):

경로	strict
BF16 Triton	0.857
LUT-B emulated	0.844
LUT-B native	0.839

네이티브와 에뮬레이션의 차이는 표준오차(±0.010) 이내입니다. 따라서 에뮬레이션이 네이티브 커널의 정확도를 신뢰성 있게 대변한다고 볼 수 있습니다. BF16 대비 손실은 약 1.5pt입니다.

메모리: GLM-4.7-Flash는 56.3 → 15.9 GiB, Qwen3.6-35B-A3B는 65.5 → 18.6 GiB로 약 3.5배 줄었습니다.

레이어 지연 (E=256, top-8, H=2048, I=512): 16토큰에서 80.3 us 대 105.4 us, 256토큰에서 172.9 us 대 239.7 us로 BF16 Triton보다 빠릅니다. 반면 1토큰(25.2 대 19.3)과 4096토큰(358.8 대 327.6)에서는 느립니다. M을 expert당 128행으로 패딩하기 때문에 decode에서 텐서 코어 작업이 낭비되고, swap-AB, TMA, persistent scheduling도 아직 없습니다. 의도적으로 단순하게 만든 커널인데도 중간 배치 구간에서 이미 이득이 있다는 점이 고무적입니다.

일반적 교훈

1. 하드웨어가 요구하는 필드를 정확도에 재활용하라. 명령어가 어차피 요구하는 B scale 슬롯에 타일 exponent를 넣어 공짜로 MSE를 개선했습니다.
2. 레이아웃 변환은 로드 타임으로. 런타임 커널은 memcpy 수준으로 단순해집니다.
3. 새 하드웨어 커널에는 수치 동등한 에뮬레이션을 짝지어라. 하드웨어 없는 CI에서도 정확도 회귀를 잡을 수 있고, 네이티브 커널의 정합성 기준점이 됩니다.
4. 하드웨어 제약에 맞춰 포맷을 재설계하라. BF16 활성화 대신 MXFP8을 택하면서 별도 decode 커널과 fallback이 통째로 필요 없어졌습니다.

참고 자료

• vLLM Quantization 문서 — vLLM 양자화 기능 개요
• PTX ISA: tcgen05 instructions — tcgen05.mma/cp/alloc 명령 공식 문서
• PTX ISA: cp.async — 비동기 global→shared 복사
• OCP Microscaling Formats (MX) v1.0 Spec — MXFP8/UE8M0 포맷 정의

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