본문으로 건너뛰기

[sglang] Hy-Nomas INT8 DiT의 groupwise dequantization을 Triton 커널 하나로 융합하기

PR 링크: sgl-project/sglang#49270 상태: Merged | 변경: +273 / -4

들어가며

SGLang의 diffusion 런타임은 Hy-Nomas 모델의 DiT를 groupwise INT8 weight(group size 128)로 서빙합니다. 이 방식은 weight를 INT8로 저장하고, 각 row의 128개 원소마다 FP32 scale 하나를 둡니다. 메모리는 아낄 수 있지만 GEMM 직전에 매번 weight를 BF16으로 복원(dequantize)해야 합니다.

기존 PyTorch 경로는 이 작업을 여러 단계로 나눠 처리했습니다.

  1. INT8 weight 전체를 FP32로 변환하면서 FP32 중간 텐서를 만든다 (weight 크기의 4배)
  2. group scale을 broadcast해서 곱한다. 결과로 또 하나의 FP32 텐서가 생긴다
  3. activation dtype(BF16)으로 cast한다

단계마다 별도의 GPU kernel launch가 일어나고, 단계 사이에서 FP32 데이터가 HBM을 오갑니다. DiT는 denoising step마다 모든 linear를 다시 실행하므로 이 비용이 그대로 누적됩니다. 이 PR은 세 단계를 Triton 커널 하나로 합칩니다. 출력은 기존 구현과 bit 단위까지 같습니다.

코드 분석

1. 새 커널: groupwise_int8_dequant.py

핵심 Triton 커널은 다음과 같습니다.

@triton.jit
def _groupwise_int8_dequant_kernel(
    weight, scales, output,
    N: tl.constexpr, WIDTH: tl.constexpr, BLOCK: tl.constexpr,
):
    offsets = tl.program_id(0).to(tl.int64) * BLOCK + tl.arange(0, BLOCK)
    if WIDTH % 128 == 0:
        scale_offsets = offsets // 128
    else:
        # A partial final group restarts at the next row.
        rows = offsets // WIDTH
        groups = (offsets % WIDTH) // 128
        scale_offsets = rows * tl.cdiv(WIDTH, 128) + groups
    values = tl.load(weight + offsets, offsets < N, other=0).to(tl.float32)
    scale = tl.load(scales + scale_offsets, offsets < N, other=0)
    result = (values * scale).to(tl.bfloat16, fp_downcast_rounding="rtne")
    tl.store(output + offsets, result, offsets < N)

눈여겨볼 설계 포인트가 네 가지 있습니다.

  • 2-D 텐서를 1-D로 평탄화해서 처리합니다. weight가 contiguous하면 원소의 선형 offset만으로 위치를 알 수 있습니다. 그래서 block 하나가 1024개 원소를 연속으로 읽고 씁니다. memory-bound 커널에서는 coalesced access가 가장 중요합니다.
  • WIDTH를 tl.constexpr로 받습니다. WIDTH % 128 == 0 분기는 컴파일 타임에 결정됩니다. width가 128의 배수인 일반적인 경우에는 row 경계를 신경 쓸 필요 없이 offsets // 128 한 번으로 scale index가 나옵니다. width가 128의 배수가 아니면 마지막 group이 128보다 짧습니다(partial final group). 이때만 row/group을 따로 계산합니다. shape마다 재컴파일이 일어나는 대신, 런타임 분기 비용은 0입니다.
  • tl.int64 offset을 씁니다. 12288×8192처럼 큰 weight에서 int32 overflow가 나지 않도록 미리 막아 둡니다.
  • FP32에서 곱하고 BF16으로는 한 번만 rounding합니다. fp_downcast_rounding="rtne"(round-to-nearest-even)는 PyTorch의 .to(torch.bfloat16)과 같은 rounding 모드입니다. 덕분에 기존 결과와 bit 단위로 일치합니다. 중간 결과는 레지스터 안에만 있고 HBM에는 쓰이지 않습니다.

2. Checked API와 custom op 등록

@register_custom_op(op_name="groupwise_int8_dequant", mutates_args=["output"])
def _dequant_into(weight, scales, output) -> None:
    ...
    if scales.shape != (rows, triton.cdiv(width, 128)):
        raise ValueError("Scales must have shape [rows, ceil(width / 128)]")
    ...
    with torch.cuda.device(weight.device):
        _launch_dequant(weight, scales, output)
  • register_custom_op로 등록했기 때문에 torch.compile(fullgraph=True)에서도 graph break 없이 opaque op으로 취급됩니다. mutates_args=["output"]로 side effect를 명시해 둔 것도 정확한 처리입니다.
  • with torch.cuda.device(...)는 multi-GPU 환경에서 weight가 cuda:1에 있어도 해당 device에서 launch되게 합니다. 테스트에서 assert torch.cuda.current_device() == current로 current device가 오염되지 않는지도 확인합니다.
  • _launch_dequant는 count가 0이면 launch 자체를 건너뜁니다. 그래서 (0, 128), (3, 0) 같은 빈 텐서도 안전하게 처리됩니다.

PR 설명에 따르면 validation을 넣으면 호출당 약 10–14µs의 host 오버헤드가 생깁니다. 작성자는 이 비용을 감수하고 checked interface를 유지했습니다. 레이어에 "검증 완료" 같은 상태를 캐싱하면 device move나 offload 이후 stale state가 될 위험이 있기 때문입니다. 정확성을 우선한 선택입니다.

3. 디스패치: quanto_int8.py

elif (
    self.group_size == 128
    and x.is_cuda
    and torch.version.hip is None
    and x.dtype == torch.bfloat16
    and layer.weight.dtype == torch.int8
    and layer.weight_scale.dtype == torch.float32
    and layer.weight.is_contiguous()
    and layer.weight_scale.is_contiguous()
):
    weight = groupwise_int8_dequant(layer.weight, layer.weight_scale)
else:
    out, width = layer.weight.shape
    ...  # 기존 PyTorch 경로

조건을 모두 만족할 때만 fused 경로를 탑니다. ROCm, CPU, FP16 activation 등 나머지 경우는 기존 경로 그대로입니다. 적용 범위를 좁게 잡아 회귀 위험을 줄인 방식입니다.

shape 검증은 런타임이 아니라 로딩 시점(process_weights_after_loading)으로 옮겼습니다.

if scales.shape != (rows, (width + self.group_size - 1) // self.group_size):
    raise ValueError("INT8 scales have an incompatible groupwise shape")

잘못된 checkpoint는 첫 forward가 아니라 로딩 단계에서 바로 실패합니다.

4. 테스트

test_fp32_multiply_before_bf16_rounding은 scale을 먼저 BF16으로 반올림하거나 FP32 subnormal을 flush하면 출력 bit가 달라지는 케이스를 직접 겨냥합니다. 경계 테스트는 offset=1로 정렬되지 않은 contiguous view까지 확인합니다. 비교는 torch.equal(actual.view(torch.int16), ...) 형태의 bit 비교입니다. tolerance 비교가 아닙니다.

왜 이게 좋은가

성능 수치 (RTX 4090 D)

Weight shape PyTorch (µs) Fused Triton (µs)
2048×2048 41.65 13.93
8192×2048 310.38 55.97
12288×2048 516.65 83.91

커널 단위로 약 3~6배 빨라졌습니다. 감소폭은 weight가 클수록 커집니다. PyTorch 경로가 원소당 HBM에 오가는 바이트를 대략 계산하면 이유가 보입니다.

  • PyTorch 경로: 1(INT8 read) + 4(FP32 write) + 4(FP32 read) + 4(FP32 write) + 4(read) + 2(BF16 write) ≈ 19 bytes
  • fused 커널: 1 + 2 = 3 bytes

memory-bound 연산에서는 이 트래픽 감소가 거의 그대로 속도 향상으로 이어집니다.

end-to-end 결과는 다음과 같습니다.

  • DiT denoising: 7.181s → 5.791s (19.4% 감소)
  • 전체 요청: 진단 pass 기준 11.1% 감소
  • E2E 표준편차: 2.238s → 0.728s

작성자는 노이즈가 큰 일반 timing pass 결과를 전부 커널 효과로 돌리지 않았습니다. synchronized breakdown을 함께 제시해, 개선이 DiT 구간에 집중되어 있고 VAE 구간은 그대로임을 보여줍니다. 측정을 정직하게 보고한 예입니다. 마이크로벤치마크도 CUDA Graph replay와 input/output rotation으로 hot-L2 재사용 효과를 배제했습니다.

일반적 교훈

  1. Weight-only quantization에서 진짜 병목은 dequant 경로일 수 있습니다. GEMM 자체보다 GEMM에 넣을 weight를 준비하는 elementwise 연산 체인이 HBM 대역폭을 더 많이 쓸 수 있습니다.
  2. Fusion할 때는 수치 의미를 그대로 보존해야 합니다. FP32 곱셈 후 RTNE로 한 번 반올림하는 순서를 지켰기 때문에, bitwise 동일성이라는 강한 보장을 얻었습니다.
  3. 좁은 디스패치 조건과 fallback 유지. 최적화 경로는 확실한 경우에만 쓰고, 나머지는 검증된 기존 경로로 보냅니다.
  4. compile, CUDA Graph, offload 호환성까지 테스트해야 합니다. 커널은 단독으로 빠른 것만으로는 부족합니다. 실제 serving 환경의 device move와 torch.compile 아래에서도 동작하는지 확인해야 합니다.

다음 단계로는 dequant를 GEMM prologue에 합쳐 BF16 weight를 아예 materialize하지 않는 방식을 생각해 볼 수 있습니다. 이번 PR은 GEMM은 그대로 두고 그 앞 단계만 정리해 위험을 낮춘 접근입니다.

참고 자료

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

댓글

관련 포스트

PR Analysis 의 다른글