W4A4 CUDA Kernel for Voice Clone: From Profiling to 2.6x Faster Than FP16 | dotieuthien

Source: https://blog.dotieuthien.com/posts/nunchaku-w4-quantization-memory-bound

Table of contents

Open Table of contents- Motivation

• 1. Profiling: Why Measure Before Optimizing- 1.1 Profiling: measure first - 1.2 Roofline analysis: confirming memory-bound - 1.3 Theory: when memory-bound, reducing data width = speedup - 1.4 In practice: chose W4A4 because Nunchaku already exists
• 2. Quantization Pipeline: FP16 → INT4 Checkpoint- 2.1 The problem: INT4 is too coarse, needs LoRA correction - 2.2 GPTQ + Iterative SVD - 2.3 Checkpoint format
• 3. CUDA Kernel: INT4 MMA + Fused LoRA- 3.1 OmniVoice model structure - 3.2 CUDA concepts + why a custom kernel is needed - 3.3 Kernel organization — C++ binding - 3.4 Pipeline: what does 1 Linear layer need? - 3.5 Why tile size 128×128, 8 warps? - 3.6 Deep dive: fused_quantize_repack_act_kernel - 3.7 Deep dive: w4a4_gemm_lora_kernel - 3.8 Deep dive: pack_lora_act_kernel - 3.9 Kernels 4-5: RMSNorm + RoPE (attention path) - 3.10 Kernel binding: CUDA → C++ → Python
• 4. CUDA Graph- 4.1 Why CUDA graphs are needed - 4.2 The problem: each input shape = a new graph - 4.3 Solution: flash_attn_varlen + bucket pre-capture
• 5. Benchmark

Motivation

This work was done by me andTriet Leas part of optimizing OmniVoice inference at work. Neither of us had CUDA kernel experience before — we learned by doing, using the real model as our exercise. Most of the process waspair programming with Claude Code: we read Nunchaku source code, asked for explanations of each concept (MMA layout, shared memory banking, warp shuffle…), then wrote the kernel together. Claude Code wrote code, we reviewed and decided what to do next — choosing the design, evaluating trade-offs, debugging profiling, verifying correctness. The results below come from learning and building at the same time.

OmniVoice— masked flow-matching voice clone model (Qwen3 backbone, 28 layers, 196 linears) — onNVIDIA L4:

EngineMean RTFvs PyTorchQualityPyTorch FP16 (baseline)0.145x1.0xbaselineTRT-LLM FP160.094x1.5xgoodTRT-LLM FP80.071x2.0xgoodW4A4 + CUDA graph0.056x2.6xacceptable*RTF = Real-Time Factor. 0.056x means 1 second of audio is generated in 56ms.

*W4A4 quality: INT4 activations lose information — some audio segments sound slightly garbled or have light noise compared to FP16. An acceptable trade-off for streaming/realtime use cases, but not yet production quality for high-fidelity voice cloning.

---

1. Profiling: Why Measure Before Optimizing

1.1 Profiling: measure first

Usedtorch\.profilerto trace a full generation (32 steps), exported the JSON and opened it in Perfetto UI. The trace file is ~300MB, ~900K events, ~58K CUDA kernels.

Zooming out shows 32llm\_forwardblocks repeating on the GPU stream (yellow row). The CPU call stack (below) showspatched\_generate\_iterative→OmniVoice\.forward→Qwen3Model→ layers.

Zooming in between two CUDA kernels reveals the real problem:

**GPU idle >80% of the time.**The GEMM kernelampere\_fp16\_s16816gemmruns for only ~12µs, but the gap after it stretches ~100µs — the GPU waits for the CPU to traverse 15+ Python wrapper layers (nn\.Module\.\_call\_impl→linear\.forward→aten::matmul→aten::mm→cudaLaunch) before the next kernel (RMSNorm→aten::pow). This is exactly the launch overhead that CUDA graphs eliminate (section 4).

The Perfetto timeline shows the pattern clearly, but it’shard to tell if the workload is memory-bound just by looking— we need numbers. Had Claude Code parse the trace JSON into a breakdown:

Total CUDA kernel time: 1857 ms (32 steps)

Category                    Time(ms)    %
─────────────────────────────────────────
GEMM (cuBLAS matmul)          674     36%
Elementwise (add/mul/cast)    620     33%
Attention (flash/fmha)        393     21%
Memory (cat/copy)             111      6%
Reduce (norm/mean)             45      2%
Other                          16      1%

Two things stand out:

1. Elementwise takes 33%— almost as much as GEMM. 43,918 tiny kernels, averaging 14µs each. The actual compute per kernel is only ~1-2µs, butlaunch overhead of ~7µs/kernel × 44K = 307ms. More than half of elementwise time is overhead, not compute.
2. GEMM averages 82-146µs— tiny for a GPU. cuBLAS matmul with M=700 (typical OmniVoice batch), hidden=1024 produces only ~3 GFLOP/call — L4 has 120 TFLOPS, so each GEMM uses only**0.002%**of peak compute in 82µs.

Both point to the same conclusion: the model is too small for the GPU →memory-bound.

1.2 Roofline analysis: confirming memory-bound

Profiling showed the problem, roofline analysis confirms it with math:

q_proj: Y[700, 2048] = X[700, 1024] × W[1024, 2048]

  FLOPs = 2 × M × N × K = 2 × 700 × 2048 × 1024 = 2.94 GFLOP
          ↑
          factor of 2 because each output element needs K multiplies + K adds
          (GFLOP = 10⁹ floating-point operations)

  Bytes = W(N×K×2B = 4.0MB) + X(M×K×2B = 1.4MB) + Y(M×N×2B = 2.8MB) = 8.2 MB
  AI    = 2.94 GFLOP / 8.2 MB = 358 FLOP/byte

  L4 INT4 break-even = 485 TOPS / 300 GB/s = 1617 FLOP/byte

  358 << 1617 → MEMORY-BOUND (GPU compute only ~22% utilized, rest waiting for data)

(L4 specs: 485 INT8 TOPS, 300 GB/s bandwidth —NVIDIA L4 Datasheet. INT4 throughput = 2× INT8.)

OperationAI (FLOP/byte)% of break-evenBoundq_proj (700,1024,2048)35822%Memorygate_proj (700,1024,3072)55334%MemoryAttention Q×K^T1308%MemoryRMSNorm0.250.02%MemoryOmniVoice hidden_size=1024 — too small for the GPU. Compare withFLUX(hidden=3072):

FLUX linear (M=4096, K=3072, N=3072):
  FLOPs = 2 × 4096 × 3072 × 3072 = 77.3 GFLOP
  Bytes = W(18MB) + X(25MB) + Y(25MB) = 68 MB
  AI    = 1137 FLOP/byte

  OmniVoice AI = 358   → memory-bound (GPU waiting for data)
  FLUX AI      = 1137  → near break-even 1617 (GPU nearly saturating compute)

FLUX is large enough for INT4 tensor cores to utilize ~70% peak TOPS → Nunchaku speeds things up throughcompute. OmniVoice has ~26× fewer FLOPs → INT4 speeds things up byreducing bytes, not through TOPS.

1.3 Theory: when memory-bound, reducing data width = speedup

GPU is idle >60% waiting for HBM. Increasing TOPS doesn’t help. What helps isreducing bytes loaded:

Per GEMM, bandwidth = 300 GB/s:

              Weight    Act       Total data   Time
FP16:        4.0 MB    1.4 MB    8.2 MB       27 µs
FP8 (W8A8):  2.0 MB    0.7 MB    5.5 MB       18 µs
W4A8:        1.0 MB    0.7 MB    4.5 MB       15 µs
W4A4:        1.0 MB    0.35 MB   4.15 MB      14 µs

In theory: weight bandwidth is the main bottleneck (weight >> activation when M is small). W4A8 and W4A4 have the same weight bandwidth — the only difference is activation (0.35 MB). So speed-wise, any W4 variant performs nearly the same on memory-bound workloads.

1.4 In practice: chose W4A4 because Nunchaku already exists

I did try building W4A8 via TRT-LLM (ModelOpt AWQ quantization) but ran into too many compatibility issues: ModelOpt doesn’t support Qwen3 export, transformers version conflicts, TRT-LLM checkpoint format mismatches. After many rounds of debugging dependency hell, concluded that W4A8 on TRT-LLM 0.18.2 + Qwen3 isn’t feasible yet.

Chose W4A4 for a practical reason:Nunchaku(MIT-HAN Lab) already has a complete W4A4 GEMM kernel implementation for diffusion models — with INT4 MMA, fused LoRA, and weight packing ready to go. Writing a custom kernel based on this design was much faster than debugging the W4A8 TRT-LLM pipeline.

And per the theory above, on memory-bound workloads W4A4 and W4A8 differ very little (~1µs/GEMM). The only trade-off is quality: INT4 activations are lossier than FP8. But with SVDQuant LoRA correction (section 2), quality remains acceptable.

---

2. Quantization Pipeline: FP16 → INT4 Checkpoint

2.1 The problem: INT4 is too coarse, needs LoRA correction

INT4 has only 16 discrete values. Quantizing FP16 → INT4 directly:

W_fp16 = [0.127, -0.891, 0.003, ...]   (65536 distinct values)
W_int4 = round(W / scale) × scale       (only 16 values)

Error per element: ~5%
Error accumulates across 28 layers → severely degraded output

SVDQuant(MIT-HAN Lab) solves this: decompose the quantization error into a low-rank matrix, compensate at runtime with LoRA:

Quantize:
  E = W - Q(W)                          ← error matrix [N, K]
  U, S, V = SVD(E, rank=16)             ← low-rank approximation
  proj_up = U[:, :16] × S[:16]          ← [N, 16]
  proj_down = V[:16, :].T               ← [K, 16]

Runtime:
  Y = INT4_GEMM(X, Q(W)) + X @ proj_down @ proj_up.T
    ≈ X @ W                              ← FP16 accuracy recovered

With SVDQuant LoRA correction: quality is acceptable

2.2 GPTQ + Iterative SVD

Quantization pipeline per layer:

for svd_iter in range(3):
    # Step 1: GPTQ quantize (Hessian-weighted rounding)
    Q_W = gptq_quantize(W - LoRA, H=X^T @ X, group_size=64)

    # Step 2: SVD on residual error
    error = W - Q_W
    proj_up, proj_down = SVD(error, rank=16)

    # Step 3: Check convergence
    if new_error >= old_error: break

GPTQuses the Hessian matrixH = X^T @ X(from calibration data) to decide rounding: columns with large activations get more careful rounding.

Iterative SVD(3 rounds): each round quantizesW \- LoRA\_previousthen re-computes SVD → LoRA correction becomes more precise.

Output per layer:qweight(INT4 packed),w\_scales(FP16 per-group),proj\_down[K, 16],proj\_up[N, 16].

2.3 Checkpoint format

w4a4_checkpoint.safetensors:
  layers.0.self_attn.q_proj.qweight:    [2048, 512] uint8  (INT4 packed, 2 per byte)
  layers.0.self_attn.q_proj.w_scales:   [2048, 16]  fp16   (per-group-64 scales)
  layers.0.self_attn.q_proj.proj_down:  [1024, 16]  fp16   (LoRA down)
  layers.0.self_attn.q_proj.proj_up:    [2048, 16]  fp16   (LoRA up)
  ... × 196 layers

At runtime:W4A4Linearloads the checkpoint → repacks weights into MMA tensor core layout (offline, once) → forward uses the custom CUDA kernel.

---

3. CUDA Kernel: INT4 MMA + Fused LoRA

3.1 OmniVoice model structure

The OmniVoice backbone is aQwen3 transformer— 28 identical layers, each containing:

Input hidden_states [batch, seq_len, 1024]
  │
  ├─ input_layernorm (RMSNorm)
  ├─ self_attn:
  │    ├─ q_proj  [1024 → 2048]  ← Linear (16 heads × 128 dim)
  │    ├─ k_proj  [1024 → 1024]  ← Linear (8 KV heads × 128 dim)
  │    ├─ v_proj  [1024 → 1024]  ← Linear
  │    ├─ q_norm, k_norm (RMSNorm per head)
  │    ├─ RoPE (rotary position encoding)
  │    ├─ Attention (Q × K^T → softmax → × V)
  │    └─ o_proj  [2048 → 1024]  ← Linear
  ├─ residual add
  ├─ post_attention_layernorm (RMSNorm)
  ├─ mlp:
  │    ├─ gate_proj [1024 → 3072]  ← Linear
  │    ├─ up_proj   [1024 → 3072]  ← Linear
  │    ├─ SiLU(gate) × up
  │    └─ down_proj  [3072 → 1024] ← Linear
  └─ residual add
  │
Output hidden_states [batch, seq_len, 1024]

7 Linear layersper block × 28 blocks =196 Linears— these are the quantization targets for W4A4. Norms, attention, and RoPE stay in FP16.

3.2 CUDA concepts + why a custom kernel is needed

CUDA terminology— need to know before reading kernel code:

TermWhat it isExampleGEMMGeneral Matrix Multiply:Y = X × W>90% of transformer computeKernelA function running in parallel on the GPU. Each launch has ~5-10µs overheadw4a4\_gemm\_lora\_kernelWarp(hardware)32 threads executing the same instruction in lockstep. Smallest GPU scheduling unitMMA instruction = 1 warpSMEMShared memory: ~128KB, shared within a block. ~13x faster than HBMHolds X tile for both quantize + LoRAMMAMatrix Multiply-Accumulate on tensor core. Multiplies an entire matrix block at oncemma\.m16n8k64: 8192 multiplies/~8 cyclesFuseMerge operations into 1 kernel → intermediate data lives in registers/SMEM, no HBM round-tripQuantize + LoRA Down = 1 kernelWhy a custom kernel?W4A4 needs extra steps: quantize X, dequant per-group, add LoRA correction. If each step runs as a separate kernel → X gets read from HBM 2-3 times → memory-bound soslower than FP8(section 1). Nunchaku (MIT-HAN Lab) solves this by fusing: share input X in SMEM, share output in registers.

Key numbers in the kernel— you’ll see 16, 64, 128 everywhere in the code. Here’s why:

• MMA tile size— fixed by hardware. The INT4 MMA instructionm16n8k64computes one small tile of the GEMMY = X × Win ~8 cycles: X tile [16, 64] × W tile [64, 8] → Y tile [16, 8] (16 tokens, (64 input dims, (16 tokens, 64 input dims) 8 output dims) 8 output dims) The numbers 16, 8, 64 are set by NVIDIA’s tensor core design — code must align to them.
• 16 rows— this isfixed by hardware, not a choice. The MMA instruction defines M=16: NVIDIA designed the tensor core to process 16 rows per instruction (32 threads split across 16 rows = 2 threads/row, fitting the MMA register layout). Code must align: quantize tile [16, 64], LoRA tile [16, 16], each warp handles 16 rows.
• Group size 64— adesign choice, but aligned to hardware. INT4 MMA processes K=64 columns per instruction. Choosing group_size=64 for quantization means each tile [16, 64] can be quantized and fed straight into MMA in the same iteration — no cross-group buffering needed.
• SVD rank 16— aquality vs speed trade-off. Higher rank (32, 64) corrects errors better but costs proportionally more compute. Rank 16 captures ~95% of quantization error energy, and fits exactly one MMA K-dimension (LoRA MMAm16n8k16→ 1 instruction per rank-tile). Rank 64 would need 4× MMA instructions + 4× memory traffic — adding ~8% overhead instead of ~2%.

3.3 Kernel organization — C++ binding

PyTorch calls cuBLAS fornn\.Linear. W4A4 replaces it with custom CUDA kernels through 3 layers:

Python (PyTorch)
  │  W4A4Linear.forward(x)     ← nn.Module, drop-in replacement for nn.Linear
  │  calls w4a4ops.w4a4_gemm(...)
  ▼
C++ binding (pybind11)
  │  w4a4_gemm_cuda(torch::Tensor act, torch::Tensor wgt, ...)
  │  extract raw CUDA pointers, launch kernel
  ▼
CUDA kernel
  │  w4a4_gemm_kernel<<<grid, block>>>(act_ptr, wgt_ptr, ...)
  │  runs on GPU

Why C++ binding?torch::Tensorin C++ shares CUDA memory with Python —no data copy. Python calls C++, C++ grabs the pointer directly, launches the kernel. Data never leaves the GPU.

Build:torch\.utils\.cpp\_extension→ nvcc compiles\.cu→\.so→import w4a4ops.

3.4 Pipeline: what does 1 Linear layer need?

PyTorchnn\.Linearmakes a single cuBLAS GEMM call. A W4A4 Linear needs more: quantize input, compute LoRA correction, dequant per-group. Everything is fused into3 kernel launches:

W4A4Linear.forward(x):    ← replaces nn.Linear

  x [700, 1024] FP16      (input: 700 tokens, hidden=1024)
  │
  ▼
  Kernel 1: Quantize + LoRA Down
  │  Reads x from HBM once, uses it for both operations:
  │  ├─ Quantize x → INT4 (pack 2 values/byte, arrange in MMA layout)
  │  └─ LoRA Down: x × proj_down[1024,16] → lora_act[700,16]
  │
  ▼
  Kernel 2: Repack LoRA activation
  │  lora_act[700,16] row-major → MMA register layout
  │  (very small: 44KB, runs in <0.1ms)
  │
  ▼
  Kernel 3: INT4 GEMM + LoRA Up
  │  Main: x_int4 × w_int4 → fpsum (INT4 MMA, dequant per-group)
  │  Epilogue: fpsum += lora_act × proj_up[2048,16]^T (FP16 MMA)
  │  fpsum lives in registers throughout → writes to HBM once at the end
  │
  ▼
  y [700, 2048] FP16       (output)

Beyond these 3 kernels (running 196×/step for 196 Linear layers), there are 2 helper kernels for attention (running 28×/step):

• Kernel 4: Fused per-head RMSNorm + RoPE for Q/K
• Kernel 5: Computes1/sqrt\(mean\(x²\)\+ε\)per row (needs to scan all K=1024 columns, separate because kernel 1 only sees 64-column tiles)

Attention SDPA usesflash\_attn\_varlen\_func(not custom).

**What is “MMA layout”?**The tensor core MMA instruction requires data arranged in a special pattern — each thread in the warp holds specific elements at non-contiguous positions. For example, thread 0 holds\[row 0 col 0\-1, row 8 col 0\-1, \.\.\.\]. Weights are repacked once at model load (offline). Activations are repacked each forward (kernels 1-2).

3.5 Why tile size 128×128, 8 warps?

First, how GPU work is organized. There are two sides:hardware(fixed) andsoftware(programmer-defined):

Software (code defines)          Hardware (fixed on GPU)
─────────────────────           ──────────────────────
Grid (all blocks)            →   GPU chip
  └─ Thread block            →   SM (Streaming Multiprocessor)
      └─ Warp (32 threads)   →   Warp scheduler + CUDA cores + Tensor cores
          └─ Thread          →   1 CUDA core (or 1 lane in tensor core)

Thread(software) — smallest unit, 1 thread runs 1 line of code. Programmer creates as many as needed.

Warp(hardware) — GPUalwaysschedules 32 threads together (lockstep). This is a fixed hardware unit — can’t run 16 or 64. MMA instructions run on 1 warp — 32 threads cooperate.

Thread block(software) — programmer groups multiple warps into a block (e.g. 8 warps = 256 threads). Threads in a block shareshared memory(hardware: ~128KB per SM). The GPU scheduler assigns each block to anSM(hardware).

SM(hardware) — Streaming Multiprocessor. L4 has58 SMs. Each SM has its own CUDA cores, tensor cores, shared memory, and register file. Multiple blocks can run on the same SM if resources allow.

Grid(software) — the collection of all blocks. GPU scheduler distributes blocks across SMs.

The GEMMY\[M, N\] = X\[M, K\] × W\[K, N\]divides the output matrix Y into tiles, each tile computed by one thread block:

Output Y [M=700, N=2048]:

  ┌────────┬────────┬────────┬─── ... ───┬────────┐
  │ Block  │ Block  │ Block  │           │ Block  │
  │ (0,0)  │ (0,1)  │ (0,2)  │    ...    │ (0,15) │  ← 2048/128 = 16 blocks along N
  │ 128×128│ 128×128│ 128×128│           │ 128×128│
  ├────────┼────────┼────────┼─── ... ───┼────────┤
  │  ...   │  ...   │  ...   │           │  ...   │
  ├────────┼────────┼────────┼─── ... ───┼────────┤
  │ Block  │ Block  │        │           │        │
  │ (5,0)  │ (5,1)  │  ...   │    ...    │ (5,15) │  ← 700/128 ≈ 6 blocks along M
  └────────┴────────┴────────┴─── ... ───┴────────┘

  Grid: 6 × 16 = 96 blocks total
  L4 has 58 SMs — each SM runs 1 block at a time (simplified).
  96 blocks > 58 SMs → every SM has work, none sit idle.

Inside each block (128×128 output),8 warpssplit the work — because MMA processes 16 rows/warp, and the block has 128 rows → 128/16 = 8 warps:

Block (i, j) — output tile 128×128:

  Warp 0: rows   0-15,  cols 0-127  (16×128 = 2048 elements)
  Warp 1: rows  16-31,  cols 0-127
  ...
  Warp 7: rows 112-127, cols 0-127

  8 warps × 32 threads = 256 threads per block

**Why 128×128?**Balancing 3 factors:

Too small(16×16): each block has little work, thousands of blocks needed. Shared memory overhead per block doesn’t amortize. Launch overhead dominates.

Too large(256×256): each block has too much work, too few blocks. Grid = (700/256, 2048/256) = (3, 8) = 24 blocks < 58 SMs → 58% of GPU sits idle.

128×128 sweet spot: 96 blocks > 58 SMs (GPU saturated), each block has enough work for 8 warps, register usage fits 1 block per SM.

Why does the GEMM need a loop?Each block computes Y[128, 128] = X[128,1024] × W[1024, 128]. Shared memory is only ~128KB — can’t hold all K=1024 columns at once. So itloops over K, loading 1 small chunk (64 columns) into shared memory each step, computing a partial sum, then loading the next chunk:

K = 1024, group_size = 64 → 16 chunks

Step 0:  load X[:,0:64] + W[0:64,:] → SMEM → MMA → partial_sum
Step 1:  load X[:,64:128] + W[64:128,:] → SMEM → MMA → partial_sum += ...
...
Step 15: load X[:,960:1024] + W[960:1024,:] → SMEM → MMA → partial_sum += ...

→ partial_sum = sum of 16 chunks = complete GEMM result

KCHUNK=2: Each step, after loading data into shared memory, all warps must wait for each other (\_\_syncthreads\(\), ~20 cycles) before reading — ensuring data is ready. Normally: 1 chunk (64 cols) per step →16 syncs. KCHUNK=2: load 2 chunks (128 cols) into shared memory at once, sync once then compute both → only16/2 = 8 syncs. Trade-off: uses 2× shared memory for weight staging.

3.6 Deep dive: fused_quantize_repack_act_kernel

Idea: This kernel fuses2 operations that share the same input X: (1) quantize X → INT4 and (2) LoRA DownX @ proj\_down. Both read X — running them separately means X gets read from HBM twice. Fused: load X into shared memory once, both operations read from there.

Concrete examplewithq\_proj(K=1024):

Input:  X[700, 1024] FP16       ← 700 tokens, hidden_size 1024

Grid:   (700/16, 1024/64) = (44, 16) = 704 blocks
        Each block: 1 warp (32 threads), processes tile [16 rows, 64 cols]

Output: act_packed[700, 1024] INT4    ← X quantized, MMA layout
        lora_act[700, 16] FP32        ← X @ proj_down[1024, 16]

704 blocks for 58 SMs → each SM runs ~12 blocks sequentially. Each block is small (32 threads, little shared memory) → many blocks run concurrently on 1 SM.

Inside 1 block— processing tile X[16, 64]:

Load X → shared memory (once) → two parallel branches:

X tile [16, 64] in shared memory
                    ┌────────────────────────────────────┐
                    │                                    │
          ┌─────────┴──────────┐              ┌──────────┴──────────┐
          │  Quantize branch   │              │  LoRA Down branch   │
          │  (reads registers) │              │  (reads SMEM)       │
          │                    │              │                     │
          │  absmax → scale    │              │  x_smem × proj_down │
          │  round → INT4      │              │  (FP16 MMA, 4 tiles)│
          │  pack 8×INT4→uint32│              │  → partial [16, 16] │
          │                    │              │                     │
          ▼                    │              ▼                     │
    act_packed → HBM           │      atomicAdd → lora_act → HBM    │
                               └────────────────────────────────────┘

Quantize path— find absmax per group using warp shuffle, compute scale, round each FP16 value to INT4 [-8, 7]:

// Warp shuffle reduction — 32 threads find shared max in ~5 cycles
for (int mask = 4; mask > 0; mask /= 2)
    maxval = fmaxf(maxval, __shfl_xor_sync(0xFFFFFFFF, maxval, mask));

float scale = fmaxf(maxval / 7.0f, 1e-10f);
int qval = max(-8, min(7, __float2int_rn(val / scale)));  // quantize
qpack |= (uint32_t)(qval & 0xF) << (j * 4);               // pack 4 bits

LoRA Down path— reads X from shared memory (already staged above), multiplies withproj\_down\[64, 16\]using FP16 MMA. 64 cols split into 4 sub-tiles × 16 cols → 4 MMA instructions:

// X is already in shared memory — 0 extra HBM reads!
xf.data[0] = *reinterpret_cast<const half2*>(&x_smem[row][col]);
pd = proj_down_packed[pd_idx];          // pre-packed weights
lp = mma_f16xf16_f32(xf, pd, lp);      // FP16 MMA → FP32 accumulate

atomicAdd(&lora_act_out[row * 16 + col], lp.data[0]);  // 16 blocks accumulate

Result: X is read from HBMoncefor both operations.

3.7 Deep dive: w4a4_gemm_lora_kernel

Idea: This kernel fuses2 operations that share the same output accumulator: (1) INT4 GEMMX\_int4 × W\_int4and (2) LoRA Uplora\_act × proj\_up^T. The GEMM output (fpsum) lives inregisters— LoRA Up adds its correction directly there before writing to HBM. Saves 1 round-trip for the output.

Concrete examplewithq\_proj(K=1024, N=2048):

Input:  act_packed[700, 1024] INT4   ← from kernel 1
        wgt_packed[2048, 1024] INT4  ← pre-packed offline
        lora_packed[700, 16] FP32    ← from kernel 2
        proj_up_packed[2048, 16] FP16 ← pre-packed offline

Grid:   (700/128, 2048/128) = (6, 16) = 96 blocks
        Each block: 8 warps (256 threads), computes output tile [128, 128]

Output: Y[700, 2048] FP16

96 blocks > 58 SMs → GPU saturated.

Inside 1 block— 3 phases:

Phase 1: GEMM main loop (16 iterations over K=1024)
  ┌─────────────────────────────────────────────────┐
  │ for k = 0..15:                                  │
  │   Load weight chunk [128, 64] from HBM → SMEM   │ ← pipeline overlap
  │   8 warps × 8 MMA tiles = 64 INT4 MMA/iteration │
  │   Dequant: INT32 → FP16, × scales, accumulate   │
  │   fpsum lives in REGISTERS across all 16 iters  │
  └─────────────────────────────────────────────────┘

Phase 2: LoRA Up epilogue (in registers, 0 HBM traffic)
  ┌─────────────────────────────────────────────────┐
  │ fpsum += lora_act × proj_up^T  (FP16 MMA)       │
  │ fpsum STAYS in registers — never written to HBM! │
  └─────────────────────────────────────────────────┘

Phase 3: Epilogue — write output
  ┌─────────────────────────────────────────────────┐
  │ registers → SMEM (transpose) → HBM (coalesced)  │
  └─────────────────────────────────────────────────┘

Weight pipeline— As explained in section 3.5, shared memory can’t hold all K=1024 so we loop 16 steps, loading 1 weight chunk from HBM each step. Problem: loading from HBM takes ~500 cycles, MMA compute takes only ~8 cycles — loading then computing wastes 98% of time. The 3-stage pipeline solves this by overlapping the next load with the current compute:

Stage 0: [Load W₀]
Stage 1:            [Load W₁][Compute W₀]    ← overlap!
Stage 2:                     [Load W₂][Compute W₁]
Stage 0:                              [Load W₃][Compute W₂]

// Non-blocking copy HBM → SMEM
asm volatile("cp.async.ca.shared.global [%0], [%1], 16;\n"
             : : "r"(smem_addr), "l"(global_addr));

Per-group dequant— INT4 MMA outputs INT32, multiply by scales to get FP16. Uses\_\_hfma2(fused multiply-add on 2 FP16 values at once, FP16 accumulator saves 50% registers):

fpsum[j].data[0] = __hfma2(
    int2half2(psum.data[0], psum.data[1]),  // INT32 → half2
    __hmul2(act_scale, wgt_scale),           // scales
    fpsum[j].data[0]);                       // accumulate FP16

LoRA Up— fpsum is in registers, add correction via FP16 MMA directly, no HBM write:

// Load lora_act (small: 16 floats per warp = 64 bytes)
lora_act_reg.data[j] = ptr[j * WARP_SIZE];

// FP16 MMA: add to fpsum (still in registers)
psum = mma_f16xf16_f32(lora_act_f16, proj_up_tile, psum);

3.8 Deep dive: pack_lora_act_kernel

Idea: Same repacking as in section 3.6 (rearranging data from row-major to MMA layout), but forlora_actinstead of INT4 activations. Kernel 1 outputslora\_act\[M, 16\]in row-major, kernel 3 (GEMM) needs MMA layout → kernel 2 bridges the gap by shuffling the data.

Why not repack inside kernel 1? Becauselora\_actis computed viaatomicAddfrom many blocks — we need to wait for all blocks to finish before the final values are ready to repack. So it’s a separate kernel.

Data is very small (700×16 = 44KB), kernel runs in <0.1ms.

3.9 Kernels 4-5: RMSNorm + RoPE (attention path)

Beyond the 3 kernels for Linear, there are 2 helper kernels for attention (running 28×/step):

• Kernel 4 (qknorm_rope): Fuses per-head RMSNorm + RoPE for Q/K. Same idea as kernel 1 — merge 2 ops that share the same data to avoid HBM round-trips. head_dim=128 divides evenly across 32 threads (1 warp), so it uses warp shuffle (fast, ~5 cycles) instead of shared memory.
• Kernel 5 (compute_inv_rms): Computes1/sqrt\(mean\(x²\)\+ε\)per row. Separate because kernel 1 (quantize) only sees tile [16, 64] — needs the mean across all K=1024 columns. This kernel scans the full row, kernel 1 just reads the result (1 scalar per row).

3.10 Kernel binding: CUDA → C++ → Python

3 layers connecting CUDA kernels to Python:

Layer 1 — CUDA kernel(w4a4\_gemm\.cu): 32\_\_global\_\_and\_\_device\_\_functions. C wrapper functions taketorch::Tensor, extract raw pointers, launch kernels with grid/block config.

Layer 2 — C++ binding(bindings\.cpp): uses pybind11 to expose C functions as a Python module.PYBIND11\_MODULEdefines the module name,m\.def\("python\_name", &cpp\_name\)maps each function.

Layer 3 — Python wrapper(linear\.py):W4A4Linear\(nn\.Module\)calls 3 kernels per forward:

def forward(self, x):
    # Kernel 1: Quantize X + LoRA Down (fused, 1 launch)
    act_packed, scales, lora_act = w4a4ops.fused_quantize_repack_lora_down(...)

    # Kernel 2: Pack LoRA activation → MMA layout (1 launch)
    lora_packed = w4a4ops.pack_lora_act(lora_act, M_padded, R)

    # Kernel 3: INT4 GEMM + LoRA Up (fused, 1 launch)
    return w4a4ops.w4a4_gemm_lora(act_packed, wgt_packed, scales, ...)

Build:python setup\.py build\_ext \-\-inplace→ nvcc compiles\.cu→\.soshared library →import w4a4opsin Python as usual.

Weight repacking(offline, once at model load): the checkpoint stores INT4 weights in row-major format. MMA instructions need data in a special layout (each thread holds specific elements).repack\_for\_mma\(\)converts the layout — runs once at load, every forward after that uses pre-packed weights.

---

4. CUDA Graph

4.1 Why CUDA graphs are needed

OmniVoice runs 32 iterative decode steps. Each step launches ~500 CUDA kernels. Kernel launch overhead is ~5-10µs each:

Without graph: 500 launches × 7µs = 3.5ms overhead/step × 32 steps = 112ms
With graph:    1 graph replay = ~10µs/step × 32 steps = 0.3ms

RTF: 0.086x → 0.056x (35% faster)

4.2 The problem: each input shape = a new graph

CUDA graphs bakealltensor shapes and kernel launch configs at capture time. The attention code uses a Python loop:

for b, L in enumerate(lens):        # Python loop — runs at capture time
    SDPA(q[:,:,:L], k[:,:,:L], ...)  # slice size L baked into graph

Different text → different sequence length → differentlens→ different graph → 2-5s cold capture. 10 different requests = 10 graphs. Not acceptable for production.

4.3 Solution: flash_attn_varlen + bucket pre-capture

flash_attn_varlen_funcreplaces the Python loop with a single kernel. Instead of slicing Q/K/V per-segment, pass acu\_seqlenstensor telling the kernel the segment boundaries.cu\_seqlensis a dynamic input — update viacopy\_\(\)before each replay, no graph rebuild needed.

Remaining problem: packed sequence length P differs per request → different tensor shape → different graph. Solution:pad P to a fixed bucket size, pre-capture a graph for each bucket at startup:

_PACKED_BUCKETS = [128, 192, 256, 384, 512, 768, 1024, 1536, 2048]

# At startup: pre-capture 9 graphs with dummy tensors
for bucket_P in _PACKED_BUCKETS:
    graph = capture(dummy_input(bucket_P))

# At inference: pad to nearest bucket, copy data, replay
# Request P=230 → bucket 256, copy data + cu_seqlens, replay graph
# Padding tokens don't affect output (flash_attn_varlen only processes tokens within cu_seqlens range)

Result: 9 graphs cover all input sizes. Zero cold captures at inference time.

---

5. Benchmark

NVIDIA L4, 10 voice-clone samples — 1 English (en_01) + 9 Vietnamese (vi_01–vi_09) with different speakers (male/female, various regional accents).

EngineMean RTFMean Infer(ms)vs PyTorchQualityPyTorch FP160.145x1654ms1.0xbaselineTRT-LLM FP160.094x1106ms1.5xgoodTRT-LLM FP80.071x817ms2.0xgoodW4A4 + CUDA graph0.056x561ms****2.6xacceptable*Per-sample RTF comparison:

SamplePyTorch FP16TRT-LLM FP16TRT-LLM FP8W4A4 + graphen_010.093x0.066x0.052x0.041xvi_010.149x0.110x0.077x0.055xvi_020.323x0.152x0.096x0.077xvi_030.210x0.132x0.091x0.062xvi_040.189x0.128x0.074x0.060xvi_050.150x0.111x0.081x0.068xvi_060.162x0.096x0.075x0.058xvi_070.141x0.084x0.070x0.053xvi_080.139x0.088x0.070x0.063xvi_090.144x0.097x0.072x0.059xMean0.145x0.094x0.071x****0.056xW4A4 audio samples (listen to quality comparison):

en_01— English voice clone:

PyTorch FP16:

W4A4:

vi_03— Vietnamese voice clone:

PyTorch FP16:

W4A4:

vi_06— Vietnamese (different speaker):

PyTorch FP16:

W4A4:

vi_07— Vietnamese (different speaker):

PyTorch FP16:

W4A4:

---

Nunchaku:MIT-HAN Lab. SVDQuant:arxiv.org/abs/2410.02355.