Bitblas — Mixed-Precision GEMM with the Ladder Compiler

강의의 출발점 — 모던 LLM inference 는 weight 만 양자화하는 W4A16 패턴이 dominant. weight 는 INT4 (메모리 절감), activation 은 FP16 (정확도 보존). mma 명령은 같은 dtype input 을 요구하는데, 어떻게 둘을 곱하는가?

강의가 답하려는 두 줄 —

1. INT4 weight 를 FP16 mma 의 input 으로 어떻게 박는가 — dequant 를 어디서 어떻게 할지의 디자인.
2. 그 dequant + matmul 을 cuBLAS 보다 빠르게 어떻게 만드는가 — vendor 라이브러리가 cover 하지 않는 자리에서.

Wang Lei 의 출발점은 명시적이다. “같은 GEMM 이 — Tensor Core 가 요구하는 정확한 layout 을 못 맞추면 — RTX 3090 위에서 cuBLAS 의 28% perf 만 낸다”. layout 은 단순한 메모리 배치가 아니라 perf 의 결정자. Bitblas 의 모든 디자인이 이 한 사실에서 나온다.

강의의 작은 역사 — 2018 년대 양자화는 “모델을 더 작은 칩에 넣기” 위함. 2024 년대는 다르다. 모델이 너무 커서 — “GPU 의 cache 와 HBM 안에 들어가게” 하는 의 도구가 양자화. inference 의 throughput 도 메모리 bandwidth 가 dominant.

Bitblas 는 라이브러리지만, 그 안에서 도는 건 Ladder 컴파일러. TVM 의 후예로 — “스케줄과 컴퓨트의 분리” + “low-bit 형식의 자동 layout 변환” 이 더해진 endtoend 컴파일러.

# Bitblas 사용 — 사용자 시점
import bitblas
from bitblas import Matmul, MatmulConfig

# 1. 추상적 정의 — input/output dtype 만
config = MatmulConfig(
    M=4096, N=4096, K=4096,
    A_dtype="float16",         # activation
    W_dtype="int4",            # weight
    accum_dtype="float16",
    out_dtype="float16",
    layout="nt",
    with_scaling=True,
    group_size=128,
)

# 2. Bitblas 가 자동 — layout 변환 + tuning + 코드 생성
matmul = Matmul(config=config)
matmul.tune()                    # 처음 호출 시 ~수십 초

# 3. 호출 — cuBLAS 같은 인터페이스
out = matmul(activation, weight_int4_packed, scale)

Tensor Core 의 mma 명령은 — input fragment 가 정확한 layout 으로 들어와야. m16n8k16, m16n8k32 같은 명령마다 register 안 데이터 배치가 박혀 있음. 잘못된 layout 이면 shared memory 를 거쳐 layout conversion 이 일어나고 — 그게 perf 의 절반 이상을 먹는다.

강의 후반부의 surprise — Wang Lei 가 만든 새 DSL TileLang (TI Language). Triton 의 사용자 친화적 syntax 와 — TVM 의 schedule 분리, CUTLASS 의 Tensor Core 직접 통제 — 를 합친 시도.

# TileLang 의 GEMM — Triton 과 닮은 모습
import tilelang.language as T

@T.prim_func
def gemm_kernel(
    A: T.Buffer((M, K), "float16"),
    B: T.Buffer((K, N), "int4"),
    C: T.Buffer((M, N), "float16"),
):
    with T.Kernel(M//128, N//128,
                  threads=128) as (bx, by):
        A_s = T.alloc_shared((128, 32), "float16")
        B_s = T.alloc_shared((32, 128), "int4")
        C_l = T.alloc_local((128, 128), "float16")

        T.clear(C_l)
        for k in T.Pipelined(K//32, num_stages=3):
            T.copy(A[bx*128, k*32], A_s)
            T.copy(B[k*32, by*128], B_s)
            T.gemm(A_s, B_s, C_l, transpose_B=True)
        T.copy(C_l, C[bx*128, by*128])

손에 새기기 — 실습 시퀀스

1. Bitblas 설치 + 첫 GEMM — pip install bitblas. MatmulConfig 으로 W4A16 정의. tune() 후 cuBLAS 와 perf 비교.
2. 같은 GEMM 의 다른 group_size — 32, 64, 128, 256 으로 tune. 어떤 size 가 자기 GPU 에서 최적인지 직접 측정.
3. ~/.bitblas 의 cache 보기 — Bitblas 도 specialization cache. 각 (M,N,K,dtype,group_size) 의 생성된 CUDA 코드를 직접 열어본다.
4. TileLang 첫 커널 — 같은 GEMM 을 TileLang 으로. T.Pipelined, T.gemm 의 사용. Triton 으로 짜는 것과 비교.
5. Llama 7B 양자화 모델 inference — AutoGPTQ 또는 AWQ 로 양자화 + Bitblas backend. cuBLAS 기반 baseline 과 토큰/초 비교.
6. 새 dtype 조합 시도 — W2A16 (INT2 weight). 정확도 손실과 perf win 측정.
7. FP4 시도 (Blackwell 사용 가능 시) — Blackwell hardware 위 FP4 mma 의 perf 확인.
8. vLLM 안 Bitblas backend — vLLM 으로 LLM serving + Bitblas. throughput 측정.

• Hopper / Blackwell 위 perf — 강의 시점은 Ampere 측정. Hopper 의 wgmma + TMA, Blackwell 의 FP4 mma 위에서 Bitblas 가 같은 비율의 win 을 내는가.
• TileLang 의 채택 곡선 — Triton 이 dominant 한 자리에서 TileLang 이 어디서 의미 있는 mindshare 를 얻는가. 학습 곡선이 길지 않은가.
• Bitblas vs Marlin 의 정확한 갭 — 같은 W4A16 GEMM 에서 두 도구의 perf 차이. 강의에서 “동급” 라고만.
• 비-NVIDIA backend — Ladder/Bitblas 는 TVM 의 후예라 backend 추가가 가능. AMD MI300 / Intel XPU 위에서의 perf.
• batch > 1 의 자리 — decode 는 batch 1 dominant. prefill 같은 큰 batch 자리에서 Bitblas 가 cuBLAS FP16 의 batched GEMM 와 어떻게 비교되는가.
• tuning time 의 비용 — config 별 tune 이 수십 초. production 에서 cold start 가 문제 안 되는지. cache 의 실용성.
• BitNet 의 INT1 perf — 1-bit weight 의 actual hardware perf. mma 가 INT1 을 받지 않으니 dequant 비용이 dominant 가능.