[sglang] Kimi-K3.1 전용 MXFP4 Decode MoE 커널을 Shape 템플릿으로 일반화해 DeepSeek-V4까지 확장하기

PR 링크: sgl-project/sglang#49091 상태: Merged | 변경: +1209 / -671

들어가며

이전 PR #48638은 Kimi-K3.1 TP8 decode 전용의 persistent MXFP4 MoE 커널을 도입했습니다. SM100(Blackwell)의 tcgen05 MMA와 TMA, PDL(Programmatic Dependent Launch)을 활용해 flashinfer의 TRT-LLM-gen MXFP4 MoE보다 빠른 decode 성능을 냈습니다. 하지만 이 커널에는 구조적인 한계가 있었습니다. 모델 shape이 상수로 박혀 있었다는 점입니다.

// Before: k31_decode/moe_mxfp4.cuh
namespace device::k31_moe {
inline constexpr uint32_t kNumExperts = 896;
inline constexpr uint32_t kTopK = 16;
inline constexpr uint32_t kHidden = 3584;
inline constexpr uint32_t kInter = 384;
...
inline constexpr uint32_t kMaxV = 1024;

이렇게 되어 있으면 같은 K3.1이라도 TP4 EP2처럼 rank당 expert 수가 달라지는 배치에서는 쓸 수 없습니다. DeepSeek-V4/V4.1처럼 hidden·intermediate·top-k·activation이 다른 모델에서도 마찬가지입니다. 이 PR은 커널을 shape traits 템플릿으로 일반화하고, activation을 템플릿 인자로 분리하며, EP용 expert offset을 추가합니다. 그 결과 하나의 커널로 7개 shape을 다룹니다.

코드 분석

1. moe_mxfp4.cuh: 상수 → Shape traits

파일 경로도 moe/k31_decode/에서 moe/mxfp4_decode/로, namespace도 k31_moe에서 mxfp4_moe로 바뀌었습니다. 가장 중요한 변경은 모델별 상수를 Shape 템플릿으로 묶은 것입니다.

// After
template <uint32_t kNumExperts_, uint32_t kTopK_, uint32_t kHidden_, uint32_t kInter_, uint32_t kBlockN_>
struct Shape {
  static constexpr uint32_t kNumExperts = kNumExperts_;  // local experts
  static constexpr uint32_t kTopK = kTopK_;
  ...
  static constexpr uint32_t kMaxRows = kMaxM * kTopK;
  // virtual experts: sum over touched experts of ceil(n_e / kBlockN) <= E + T * top-k / kBlockN
  static constexpr uint32_t kMaxV = kNumExperts + kMaxRows / kBlockN;
  static constexpr uint32_t kFc1Tiles = 2 * kInter / 128;
  static constexpr uint32_t kFc2Items = kHidden / 256;
  static constexpr uint32_t kFc1Stages = kHidden / 256;
  static constexpr uint32_t kFc2Stages = kInter / 128;

주목할 부분은 kMaxV입니다. 이전에는 1024가 손으로 계산되어 박혀 있었습니다(896 + 256*16/32). 지금은 E + T*top-k/BlockN이라는 상한 공식으로 도출됩니다. 그래서 shape이 바뀌어도 counter 배열 크기와 claim cursor 위치(kClaim1Word, kClaim2Word)가 자동으로 맞춰집니다. 값은 여전히 constexpr이므로 런타임 분기나 레지스터 비용이 없습니다. PR 설명에 따르면 K3.1 TP8 인스턴스의 SASS는 #48638과 완전히 동일합니다. 일반화 과정에서 기존 경로 성능이 전혀 희생되지 않았다는 강력한 증거입니다.

2. 커널 가정을 static_assert로 명문화

이전 코드에서는 "VInfo는 expert 10비트, row 12비트…" 같은 비트 패킹 가정이 주석에만 있었습니다. 일반화하면 이 가정이 깨질 수 있으므로 컴파일 타임에 검증합니다.

static_assert(kHidden % 512 == 0, "FC1 / FC2 K 256 stages, token scale rows in 16-byte vectors");
static_assert(kInter % 128 == 0, "FC2 K-atoms of 128, intermediate scale atoms of 4 K-blocks");
static_assert(kBlockN == 16 || kBlockN == 32, "MMA N % 16 at M 128; ...");
static_assert(kNumExperts <= kRouteThreads, "routing: one thread per expert");
// VInfo fields: expert 10 bits, first row 12 bits, tokens 6 bits, chunk 4 bits
static_assert(kNumExperts <= 1024 && kMaxRows <= 4096 && kBlockN < 64 && kMaxM / kBlockN <= 16);
static_assert(kMaxV < (1u << 16), "routing scan packs (virtual experts, rows) into 16-bit halves");

kBlockN을 16으로 줄일 수 있게 되면서 chunk 수가 최대 16이 됩니다. 그래서 VInfo 주석도 chunk < 8에서 chunk < 16으로 바뀌었습니다. 이처럼 템플릿화와 함께 불변식을 static_assert로 옮겨 두면, 지원 범위를 벗어난 shape을 인스턴스화하려 할 때 조용한 오동작 대신 컴파일 에러가 납니다.

3. Workspace 레이아웃을 constexpr로 한 곳에서 carve

static constexpr uint32_t kActRows = kMaxRows + kBlockN;  // the last virtual expert's token box may read past
static constexpr uint64_t kVInfoOffset = align_up(kCountersOffset + kCounterWords * 4);
static constexpr uint64_t kTokenBitsOffset = align_up(kVInfoOffset + (kMaxV + 1) * 4);
static constexpr uint64_t kActSfOffset = align_up(kTokenBitsOffset + uint64_t(kMaxV) * kMaxWords * 4);
static constexpr uint64_t kActOffset = align_up(kActSfOffset + uint64_t(kMaxV) * kActSfAtoms * kAtomBytes);
static constexpr uint64_t kWorkspaceBytes = align_up(kActOffset + uint64_t(kActRows) * kInter);

counters, vinfo, token bitmask, intermediate scale, intermediate activation을 하나의 byte tensor에서 256B 정렬로 잘라 씁니다. 덕분에 host 쪽은 kWorkspaceBytes 하나만 할당하면 됩니다. routing과 MoE를 한 번의 host entry로 launch하는 구조와 잘 맞습니다. kActRows에 kBlockN을 더한 것은 마지막 virtual expert의 TMA token box가 경계를 넘어 읽는 경우를 위한 padding입니다. 이런 엣지 케이스가 주석으로 남아 있다는 점도 좋습니다.

4. Activation을 템플릿 함수 인자로

기존에는 SiTU가 상수 beta=4, lb=25로 하드코딩되어 있었습니다.

// Before (removed)
SGL_DEVICE float situ(float gate, float up) {
  const float g = 2.0f * ptx::tanh_approx(gate * 0.25f) * (1.0f + ptx::tanh_approx(gate * 0.5f));
  const float u = 25.0f * ptx::tanh_approx(up * 0.04f);
  return g * u;
}

이제 FC1 epilogue의 activation은 situ<beta, linear_beta> 또는 swiglu<limit> 형태의 템플릿 인자로 주입됩니다. DeepSeek-V4/V4.1이 쓰는 clamped SwiGLU(flashinfer fused FC1 epilogue와 같은 의미)도 같은 커널 본문으로 처리됩니다. 상수가 템플릿 파라미터이므로 gate * 0.25f 같은 곱셈이 그대로 상수 폴딩되고, 함수 포인터나 분기 없이 인라인됩니다.

5. Routing: EP offset과 -1 id 처리

struct RouteParams {
  const int32_t* __restrict__ topk_ids;  // [T, top-k]: global expert ids, ...
  int32_t* __restrict__ expanded_idx;    // [T * top-k] -> permuted row, -1 for an id that is not a local expert
  ...
  uint32_t expert_offset;  // global id of local expert 0
};

EP 환경에서는 다른 rank의 expert id나 -1(마스킹된 슬롯)이 들어옵니다. routing kernel은 [expert_offset, expert_offset + E) 범위 밖 id에 row를 배정하지 않고 expanded_idx = -1을 기록합니다. 그러면 engine의 deferred-finalize 단계가 해당 슬롯을 자연스럽게 건너뜁니다. top-k도 flat하게 읽으므로 16 외의 값도 지원합니다. bitmask를 atomicOr로 만들고 popcount로 rank를 매기는 방식은 그대로 유지했습니다. 따라서 결정론성(determinism)도 보존됩니다.

6. Small-T 스케줄링 개선

• 한 CTA가 동시에 보유하는 FC1 item을 최대 2개로 제한했습니다. 토큰 수가 작아 item이 적을 때 앞쪽 CTA가 queue를 독점하면 tail 불균형이 생기는데, 이를 막습니다.
• routing kernel을 호출에 필요한 bitmask word 수(kWords)로 인스턴스화합니다. 작은 T에서 불필요한 word 순회를 줄입니다.
• FC1/FC2 claim prefetch 깊이도 템플릿 인자가 되었습니다.

그 결과 #48638 대비 K3.1 TP8에서 T ≤ 12일 때 3~5% 더 빨라졌습니다.

7. Python Op / Engine 연결

op은 k31_decode_moe에서 sglang.kernels.ops.moe.mxfp4_decode_moe로 바뀌었습니다. activation, activation_constants, expert_offset, is_supported_shape()를 받습니다. 플래그는 SGLANG_USE_CUDA_MXFP4_DECODE_MOE 하나로 통합되었고 기본값은 off입니다. 기존 SGLANG_K31_DECODE_MOE는 deprecated alias로 남았습니다. shape, activation, MXFP8 activation이 지원 범위에 들고 토큰이 256개 이하인 호출만 이 커널을 씁니다. 나머지는 flashinfer로 fallback합니다. Mxfp4MoEMethod(SiTU, 이제 EP 포함)와 Mxfp4FlashinferTrtllmMoEMethod(DeepSeek-V4 SwiGLU) 두 경로가 run_cuda_mxfp4_decode_moe를 공유합니다.

왜 이게 좋은가

B300, cold L2, CUDA graph 조건에서 flashinfer 0.7.2의 최적 tactic과 비교한 결과입니다(us/layer).

shape	T	this PR	flashinfer best	ratio
K3.1 TP8	6 / 24 / 96 / 252	27.1 / 65.8 / 150.7 / 208.0	41.2 / 89.3 / 187.9 / 257.1	0.66 / 0.74 / 0.80 / 0.81
K3.1 TP4 EP2	6 / 24 / 96 / 252	26.9 / 65.3 / 152.8 / 206.4	36.1 / 81.9 / 167.9 / 227.2	0.75 / 0.80 / 0.91 / 0.91
DeepSeek-V4.1-Flash TP4	6 / 24 / 96 / 252	28.7 / 58.4 / 123.9 / 180.7	31.4 / 70.7 / 139.5 / 206.2	0.91 / 0.83 / 0.89 / 0.88

특히 decode의 핵심 구간인 작은 T에서 K3.1 TP8은 34% 단축됩니다. DeepSeek-V4.1 speculative verify(요청당 6 토큰)에서도 일관되게 9~17% 빠릅니다.

정확성 검증도 탄탄합니다. 7개 shape × T{1,7,48,96,256} × BlockN{16,32}를 fp64 reference와 비교했습니다. DeepSeek-V4 shape에서는 FC2 row가 flashinfer와 bitwise 동일합니다. apply() 플래그 on/off 비교에서도 출력이 bitwise 동일합니다. 리뷰어(DarkSharpness)는 /rerun-test로 커널 단위 테스트뿐 아니라 gpt-oss, Kimi-K3, DeepSeek-V4 Flash(B200/GB300, TP2EP2, MegaMoE 포함) e2e 테스트까지 돌렸습니다. 공유 경로를 건드린 변경인 만큼 영향받는 모든 모델 경로를 회귀 검증한 것입니다.

일반적 교훈

1. 특화 커널은 상수를 템플릿 traits로 올려 일반화하라. constexpr 기반이므로 SASS가 동일하게 유지되어, 일반화 비용이 0입니다.
2. 암묵적 가정은 static_assert로. 비트 패킹·스레드 수·정렬 같은 불변식을 컴파일 타임에 강제하면 새 shape 추가가 안전해집니다.
3. Activation 같은 epilogue는 functor 템플릿으로 주입하라. 같은 메인 루프를 재사용하면서 모델별 차이를 흡수할 수 있습니다.
4. 작은 batch의 persistent kernel에서는 work hoarding을 막아라. CTA당 claim 상한 같은 단순한 정책이 tail latency를 줄입니다.
5. opt-in 플래그와 엄격한 fallback 조건(is_supported_shape(), T ≤ 256)으로 신규 커널의 위험 범위를 좁혀라.

참고 자료

• CUDA C++ Programming Guide — PDL, persistent kernel, atomic 연산 등 커널 구현의 기반 문서
• PTX ISA - tcgen05 — SM100 tcgen05 MMA/TMEM 명령어 정의
• FlashInfer — 비교 대상인 TRT-LLM-gen MXFP4 MoE 구현
• SGLang — PR 대상 레포지토리

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