[onnxruntime] GGUF 블록 패킹 가중치를 그대로 실행하기: ONNX Runtime의 GatherPacked·MatMulPacked·MoEPacked

PR 링크: microsoft/onnxruntime#33953 상태: Merged | 변경: +6303 / -9

들어가며

llama.cpp 생태계의 GGUF 포맷은 IQ2_XXS, Q4_0, Q6_K 같은 다양한 블록 양자화 형식을 사용합니다. 각 블록 안에 scale과 양자화 값이 interleave되어 저장되는 구조라, ONNX Runtime의 기존 MatMulNBits처럼 "weight 텐서 + 별도 scale 텐서"를 기대하는 연산자에 넣으려면 오프라인 repacking이 필요했습니다. repacking은 변환 파이프라인을 복잡하게 만들고, 포맷마다 정밀도/레이아웃 차이(예: offset-UINT8 Q8 vs. signed byte + interleaved scale의 native Q8_0)를 정확히 맞춰야 하는 부담도 있습니다.

이 PR은 com.microsoft 도메인에 세 가지 contrib op를 추가해 이 문제를 정면으로 해결합니다.

• GatherPacked (CPU/CUDA): 패킹된 embedding table에서 row를 gather하며 즉석 dequantize
• MatMulPacked (CUDA): Y = A @ B^T, B는 2D block-packed blob
• MoEPacked (CUDA): gate/up/down 세 개의 3D packed expert blob + softmax top-k routing + SwiGLU

핵심은 15개 포맷을 native block layout 그대로 소비한다는 점입니다. #33877의 direct external initializer buffer와 결합하면 .gguf shard를 로딩할 때 host 측 중복 복사도 줄일 수 있습니다(단, CUDA 실행 시에는 여전히 device 메모리로 복사됩니다).

코드 분석

1. 스키마: 포맷 ID를 공개 계약으로

docs/ContribOperators.md에 생성된 스키마를 보면 설계 의도가 드러납니다.

GatherPacked
  K : int (required)  — packed format 블록 크기의 양의 배수
  packed_type : int (required) — Stable packed-format ID
  output_dtype : int  (default 16 = BFLOAT16)
  data : uint8 [vocab_size, row_bytes]
  indices : int32/int64, Negative indices and indices >= vocab_size select row 0.

입력이 단순히 uint8 바이트 blob이고, 해석 방법은 packed_type이라는 stable ID로 결정됩니다. 리뷰 과정에서 이 ID와 바이트 레이아웃을 docs/PackedWeightFormats.md라는 별도 계약 문서로 고정했는데, 이는 "내부 디코더 구현은 바뀌어도 모델 파일의 의미는 바뀌지 않는다"는 보장을 주기 위한 것입니다.

MoEPacked는 gate/up/down 각각에 gate_type, up_type, down_type을 따로 받습니다. GGUF 모델은 레이어/텐서별로 다른 양자화(예: down projection만 더 높은 비트)를 섞어 쓰는 경우가 흔하기 때문에 이 분리가 실용적입니다.

MoEPacked
  gate_experts_weights : uint8 [num_experts, inter_size, gate_row_bytes]
  up_experts_weights   : uint8 [num_experts, inter_size, up_row_bytes]
  down_experts_weights : uint8 [num_experts, hidden_size, down_row_bytes]

2. gguf_quant_cuda.cuh: 서브블록 디코더와 Tensor Core 래퍼

CUDA 쪽 공통 헤더에는 DecodeGgufSubBlock8Impl, DecodeGgufSubBlock16Impl, shared memory LUT 버전 DecodeGgufSubBlock8SmemLutTpl, 그리고 레지스터만 쓰는 Iq2xsSignsDev / Iq4NlValueDev 헬퍼가 들어 있습니다. IQ 계열은 codebook grid(iq2xxs_grid 등)를 조회해야 하므로, 자주 쓰이는 grid는 shared memory에 올리고 작은 테이블은 레지스터 연산으로 대체해 global memory 접근을 줄이는 전략입니다.

여기에 mma.sync.aligned.m16n8k16을 감싼 MmaBf16 래퍼가 있는데, 리뷰에서 중요한 버그가 지적되었습니다.

Below SM80, this helper does nothing, so BF16 MoEPacked and large-N BF16 MatMulPacked decode can return all-zero results successfully.

BF16 MMA는 SM80 이상에서만 지원되므로, 이전 아키텍처에서는 에러 없이 0을 반환하는 최악의 silent failure가 발생할 수 있었습니다. 최종 버전은 configured device의 capability를 probe하고 BF16 SIMT fallback을 사용하도록 수정되었으며, PR 설명에 따르면 packed 커널을 compute_75 PTX로만 빌드해 이 fallback 경로를 강제 실행하는 테스트까지 수행했습니다.

3. matmul_packed.cu: decode와 prefill의 이원화

LLM 추론의 두 단계는 성격이 정반대입니다.

• Decode (M <= 8): memory-bound. 가중치를 한 번 읽는 비용이 전부입니다. → warp-K-split Tensor Core 경로와 warp GEMV 경로에서 로드 직후 레지스터에서 디코딩하여 연산에 바로 투입합니다.
• Prefill (M > 8): compute-bound. → bounded tile 단위로 dequantize한 뒤 cuBLAS GEMM에 맡깁니다. scratch 메모리가 tile 크기로 제한되므로 전체 가중치를 풀어놓는 메모리 폭발이 없습니다.

문서에 명시된 정밀도 규칙도 흥미롭습니다.

Decode (M <= 8) uses float32-decoded weights for FP16 and for BF16 with N < 65536.
BF16 decode with N >= 65536 rounds decoded weights to BF16, including the SIMT fallback.
Prefill (M > 8) uses ... BF16-rounded decoded weights for BF16;
FP16 prefill uses float32-decoded weights unless the packed format is F16.
Consequently, BF16 results may differ across the M and N dispatch boundaries.

FP16 prefill에서 F16 포맷이 아니면 FP32 decoded weight를 쓰는 이유는, Q6_K처럼 scale × value가 FP16 range/정밀도를 벗어날 수 있기 때문입니다. 그리고 dispatch 경계에서 결과가 미세하게 달라질 수 있다는 사실을 스키마 문서에 명시한 점은 리뷰어 요청("explicit rounding boundaries")을 반영한 좋은 사례입니다.

리뷰에서 지적된 또 하나의 핵심 이슈는 warp 수렴 조건입니다.

valid_lo[m] varies across lanes for partial output tiles, but mma.sync.aligned requires the entire warp to participate. BF16 decode with N=65537 ... giving undefined behavior. Zero-fill invalid weight fragments and execute both MMA calls unconditionally, retaining the final-store guards.

mma.sync.aligned는 warp 전체가 같은 명령을 실행해야 합니다. 경계 tile에서 일부 lane만 MMA를 호출하면 UB입니다. 해법은 정석적인 패턴: 연산은 무조건 수행하되 범위 밖 fragment는 0으로 채우고, store 단계에서만 guard하는 것입니다.

4. moe_packed.cu: 라우팅부터 finalize까지 fusion

MoEPacked는 다음을 하나의 파이프라인으로 처리합니다.

1. fused top-k softmax routing
2. parallel warp-scan으로 route를 tile 단위로 bucketing (decode 8 routes/tile, prefill 32 routes/tile)
3. decode: warp-K-split Tensor Core FC1 + SwiGLU fusion / prefill: 4-subtile Tensor Core FC1/FC2
4. 128-bit vectorized finalize (가중합)

리뷰를 통해 수정된 부분이 특히 교육적입니다.

• k <= num_experts 검증: num_experts=2, k=4일 때 top-k 루프가 유효 key를 소진해 key 0을 선택하고, 이것이 expert 8191로 디코딩되어 finalize가 쓰이지 않은 FC2 버퍼를 읽는 문제가 있었습니다.
• dynamic shared memory 48 KiB 제한: hidden_size=4096 FP16이면 64 KiB가 필요한데 cudaFuncSetAttribute로 opt-in하지 않으면 launch가 실패합니다. bounded staging 또는 device limit 확인 후 opt-in + fallback이 필요합니다.
• launch clamp: M=17, k=8, E=512, H=2560, I=7168 BF16 워크로드에서 FC1 grid (112,136,1), FC2 grid (20,136,1)로 총 17,952 CTA가 실행되었고, 이전 bound 계산으로는 68,244 CTA였습니다. 리뷰어는 이를 "launch-count evidence, not a measured latency improvement"라고 정확히 구분해 기술했습니다.

5. CPU GatherPacked와 입력 검증

Copilot 리뷰는 두 가지 검증 누락을 짚었습니다.

• [0, row_bytes] 형태의 빈 vocabulary에 non-empty indices가 들어오면, out-of-range index가 row 0으로 매핑되는 규칙 때문에 빈 버퍼를 읽습니다. empty-indices fast path 이후 빈 vocab을 reject해야 합니다.
• GgmlRowBytes는 Q4_1(packed_type=3)처럼 디코더가 구현되지 않은 포맷도 geometry 계산은 통과시켜, 결과적으로 0 가중치가 조용히 출력될 수 있었습니다. "레이아웃을 아는 것"과 "디코딩할 수 있는 것"을 분리해 검증해야 한다는 교훈입니다.

6. 테스트 강화

test_gguf_packed_ops.py에 대한 리뷰도 실질적이었습니다. CUDA provider를 지정해도 CPU fallback이 허용되면 CPU vs CUDA 비교 테스트가 둘 다 CPU에서 실행되어 의미 없이 통과할 수 있습니다. 또 MoE 테스트가 shape/finiteness만 보면 all-zero 출력도 통과합니다. 최종적으로 15개 포맷 bit-exactness, 독립적인 high-index IQ fixture, 해석적 상수 출력 assertion 등이 추가되어 33개 테스트 / 811 subtests가 통과했습니다.

왜 이게 좋은가

H200에서 Qwen3.8-Flash-Next-GGUF(UD-IQ1_S, UD-IQ2_XS, UD-IQ3_XXS) end-to-end 결과:

항목	결과
Prefill 128 tok TTFT	0.048–0.052 s
Prefill 2K	0.222–0.244 s (8,400–9,226 tok/s)
Prefill 8K	0.851–0.944 s (8,682–9,622 tok/s, baseline SIMT 대비 17.1x–18.5x)
Speculative decode (DFlash2 S=8)	20.0–20.9 ms/cycle, 중앙값 183–269 tok/s
MMLU-Pro (200문항)	87.50% / 85.50% / 86.00%

일반적 교훈

1. 포맷을 바꾸지 말고 커널이 포맷을 이해하게 하라. repacking 대신 native layout을 디코딩하면 변환 파이프라인과 정밀도 불일치 위험이 사라집니다. 단, 그 대가로 byte layout을 public contract로 고정해야 합니다.
2. Decode와 prefill은 다른 문제다. memory-bound 영역은 fused dequant + GEMV/MMA, compute-bound 영역은 bounded dequant + cuBLAS가 정답입니다.
3. Silent failure를 가장 경계하라. SM80 미만의 no-op MMA, 미구현 포맷의 zero output, CPU fallback으로 통과하는 테스트 모두 "에러 없이 틀린 값"이라는 공통점이 있습니다.
4. Warp-synchronous 명령은 predicate 대신 zero-fill로. mma.sync는 전체 warp 참여가 전제이므로 경계 처리는 store 시점에 하세요.
5. 정밀도 경계를 문서화하라. dispatch 경로에 따라 BF16 결과가 달라질 수 있다는 사실을 스키마에 명시한 것은 사용자의 디버깅 비용을 크게 줄여줍니다.

리뷰 논의에서 tianleiwu가 정리했듯, native storage가 반드시 새 operator 이름을 요구하는 것은 아닙니다(#33940은 기존 스키마에 guarded profile을 추가하는 대안을 보여줌). 하지만 공개 계약—바이트 코덱, 차원, 반올림, routing 의미—을 명시적으로 정의해야 한다는 점은 어느 접근에서든 동일하며, 이 PR은 그 계약을 균일한 인터페이스로 노출한 사례입니다.

참고 자료

• ONNX Runtime Contrib Operators — GatherPacked/MatMulPacked/MoEPacked 스키마가 추가된 문서
• ONNX Runtime Operator Kernels — 커널 등록 문서
• PTX ISA: Warp-level matrix multiply (mma) — mma.sync.aligned.m16n8k16 명세
• CUDA C++ Programming Guide — dynamic shared memory 및 cudaFuncSetAttribute opt-in
• cuBLAS Documentation — prefill 경로의 GEMM
• GGUF format (ggml) — 블록 양자화 가중치 포맷

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