본문으로 건너뛰기

[sglang] H200 GPU에서 GLM-5.2 MoE를 위한 W4A8 GEMM 커널 최적화 분석

PR 링크: sgl-project/sglang#38220 상태: Merged | 변경: +89 / -0

들어가며

최근 대규모 언어 모델(LLM)의 효율적인 추론은 AI 서비스의 핵심 경쟁력으로 부상하고 있습니다. 특히 Mixture-of-Experts(MoE) 모델은 희소 활성화(sparse activation)를 통해 모델 파라미터 수를 크게 늘리면서도 계산 비용을 효율적으로 유지할 수 있어 주목받고 있습니다. SGLang은 이러한 LLM 추론을 최적화하는 데 중점을 둔 프레임워크이며, GPU 커널 최적화는 SGLang의 성능에 지대한 영향을 미칩니다.

이번에 분석할 PR(sgl-project/sglang#1161)은 H200 GPU에서 GLM-5.2 MoE 모델의 W4A8(Weight 4-bit, Activation 8-bit) 양자화된 Grouped GEMM(General Matrix Multiply) 연산의 성능을 획기적으로 개선한 사례입니다. 기존 SGLang의 dispatch_w4a8_moe_mm_sm90 커널은 주로 DeepSeek 모델의 특정 shape에 맞춰 CUTLASS 설정을 튜닝했기 때문에, GLM-5.2 모델의 고유한 GEMM shape에서는 최적화되지 않은 일반(generic) fallback 커널을 사용해야 했습니다. 이는 GLM-5.2 모델의 추론 성능에 병목 현상을 일으켰습니다.

이 PR의 목표는 GLM-5.2 모델의 실제 GEMM shape(TP-8 및 EP-8 레이아웃)에 맞춰 H200 GPU에 특화된 CUTLASS 설정을 추가하여, MoE 레이어의 성능을 극대화하는 것입니다. 이를 통해 GLM-5.2 모델의 추론 속도를 크게 향상시킬 수 있었습니다.

코드 분석

핵심 변경사항은 python/sglang/kernels/aot/csrc/moe/cutlass_moe/w4a8/w4a8_grouped_mm_c3x.cu 파일에 집중되어 있습니다. 이 파일은 W4A8 MoE Grouped GEMM 연산을 위한 CUTLASS 커널 디스패치 로직을 포함하고 있습니다.

1. H200 장치 감지 및 전용 디스패치 로직 추가

cutlass_w4a8_moe_mm_sm90 함수는 이제 현재 실행 중인 GPU 장치가 H200인지 감지하는 로직을 포함합니다. H200 장치로 확인되면, GLM-5.2에 특화된 새로운 try_dispatch_w4a8_moe_mm_h200 함수를 먼저 호출하여 최적화된 커널을 시도합니다. 만약 H200이 아니거나 try_dispatch_w4a8_moe_mm_h200에서 해당 shape에 맞는 튜닝된 커널을 찾지 못하면, 기존의 일반 dispatch_w4a8_moe_mm_sm90 함수로 fallback합니다.

Before:

void cutlass_w4a8_moe_mm_sm90(
    torch::Tensor& d_tensors,
    torch::Tensor const& a_tensors,
    torch::Tensor const& b_tensors,
    torch::Tensor const& a_scales,
    torch::Tensor const& b_scales,
    torch::Tensor const& expert_offsets,
    torch::Tensor const& problem_sizes,
    torch::Tensor const& a_strides,
    torch::Tensor const& b_strides,
    torch::Tensor const& d_strides,
    torch::Tensor const& s_strides,
    int64_t chunk_size,
    int64_t topk) {
  dispatch_w4a8_moe_mm_sm90(
      d_tensors,
      a_tensors,
      b_tensors,
      a_scales,
      b_scales,
      expert_offsets,
      problem_sizes,
      a_strides,
      b_strides,
      d_strides,
      s_strides,
      chunk_size,
      topk);
}

After:

void cutlass_w4a8_moe_mm_sm90(
    torch::Tensor& d_tensors,
    torch::Tensor const& a_tensors,
    torch::Tensor const& b_tensors,
    torch::Tensor const& a_scales,
    torch::Tensor const& b_scales,
    torch::Tensor const& expert_offsets,
    torch::Tensor const& problem_sizes,
    torch::Tensor const& a_strides,
    torch::Tensor const& b_strides,
    torch::Tensor const& d_strides,
    torch::Tensor const& s_strides,
    int64_t chunk_size,
    int64_t topk) {
  const c10::cuda::CUDAGuard device_guard(a_tensors.device());
  // Detect H200 devices, including names such as "NVIDIA H200 SXM".
  const std::string_view device_name(at::cuda::getCurrentDeviceProperties()->name);
  const auto model_pos = device_name.find("H200");
  const bool is_h200 = model_pos != std::string_view::npos && (model_pos == 0 || device_name[model_pos - 1] == ' ') &&
                       (model_pos + 4 == device_name.size() || device_name[model_pos + 4] == ' ');
  if (is_h200 && try_dispatch_w4a8_moe_mm_h200(
                     d_tensors,
                     a_tensors,
                     b_tensors,
                     a_scales,
                     b_scales,
                     expert_offsets,
                     problem_sizes,
                     a_strides,
                     b_strides,
                     d_strides,
                     s_strides,
                     chunk_size,
                     topk)) {
    return;
  }

  dispatch_w4a8_moe_mm_sm90(
      d_tensors,
      a_tensors,
      b_tensors,
      a_scales,
      b_scales,
      expert_offsets,
      problem_sizes,
      a_strides,
      b_strides,
      d_strides,
      s_strides,
      chunk_size,
      topk);
}

이 변경은 at::cuda::getCurrentDeviceProperties()를 사용하여 현재 GPU의 이름을 가져와

참고 자료

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

댓글

관련 포스트

PR Analysis 의 다른글