W4A4 CUDA 内核用于语音克隆：从性能分析到比 FP16 快 2.6 倍 | dotieuthien 来源：https://blog.dotieuthien.com/posts/nunchaku-w4-quantization-memory-bound

目录 打开目录

• 动机 (https://blog.dotieuthien.com/posts/nunchaku-w4-quantization-memory-bound#span-stylecolor-1d4ed8motivationspan)
• 1. 性能分析：为什么优化前要先测量 (https://blog.dotieuthien.com/posts/nunchaku-w4-quantization-memory-bound#span-stylecolor-1d4ed81-profiling-why-measure-before-optimizingspan)
  • 1.1 性能分析：先测量 (https://blog.dotieuthien.com/posts/nunchaku-w4-quantization-memory-bound#span-stylecolor-1d4ed811-profiling-measure-firstspan)
  • 1.2 屋顶线分析：确认内存受限 (https://blog.dotieuthien.com/posts/nunchaku-w4-quantization-memory-bound#span-stylecolor-1d4ed812-roofline-analysis-confirming-memory-boundspan)
  • 1.3 理论：内存受限时，减小数据宽度 = 加速 (https://blog.dotieuthien.com/posts/nunchaku-w4-quantization-memory-bound#span-stylecolor-1d4ed813-theory-when-memory-bound-reducing-data-width–speedupspan)
  • 1.4 实践中：选择 W4A4 是因为 Nunchaku 已存在 (https://blog.dotieuthien.com/posts/nunchaku-w4-quantization-memory-bound#span-stylecolor-1d4ed814-in-practice-chose-w4a4-because-nunchaku-already-existsspan)
• 2. 量化流水线：FP16 → INT4 检查点 (https://blog.dotieuthien.com/posts/nunchaku-w4-quantization-memory-bound#span-stylecolor-1d4ed82-quantization-pipeline-fp16–int4-checkpointspan)
  • 2.1 问题：INT4 太粗糙，需要 LoRA 修正 (https://blog.dotieuthien.com/posts/nunchaku-w4-quantization-memory-bound#span-stylecolor-1d4ed821-the-problem-int4-is-too-coarse-needs-lora-correctionspan)
  • 2.2 GPTQ + 迭代 SVD (https://blog.dotieuthien.com/posts/nunchaku-w4-quantization-memory-bound#span-stylecolor-1d4ed822-gptq–iterative-svdspan)
  • 2.3 检查点格式 (https://blog.dotieuthien.com/posts/nunchaku-w4-quantization-memory-bound#span-stylecolor-1d4ed823-checkpoint-formatspan)
• 3. CUDA 内核：INT4 MMA + 融合 LoRA (https://blog.dotieuthien.com/posts/nunchaku-w4-quantization-memory-bound#span-stylecolor-1d4ed83-cuda-kernel-int4-mma–fused-loraspan)
  • 3.1 OmniVoice 模型结构 (https://blog.dotieuthien.com/posts/nunchaku-w4-quantization-memory-bound#span-stylecolor-1d4ed831-omnivoice-model-structurespan)
  • 3.2 CUDA 概念 + 为什么需要自定义内核 (https://blog.dotieuthien.com/posts/nunchaku-w4-quantization-memory-bound#span-stylecolor-1d4ed832-cuda-concepts–why-a-custom-kernel-is-neededspan)
  • 3.3 内核组织 — C++ 绑定 (https://blog.dotieuthien.com/posts/nunchaku-w4-quantization-memory-bound#span-stylecolor-1d4ed833-kernel-organization–c-bindingspan)
  • 3.4 流水线：1 个 Linear 层需要什么？ (https://blog.dotieuthien.com/posts/nunchaku-w4-quantization-memory-bound#span-stylecolor-1d4ed834-pipeline-what-does-1-linear-layer-needspan)
  • 3.5 为什么 tile 大小是 128×128，8 个 warp？ (https://blog.dotieuthien.com/posts/nunchaku-w4-quantization-memory-bound#span-stylecolor-1d4ed835-why-tile-size-128128-8-warpsspan)
  • 3.6 深入：fused_quantize_repack_act_kernel (https://blog.dotieuthien.com/posts/nunchaku-w4-quantization-memory-bound#span-stylecolor-1d4ed836-deep-dive-fused_quantize_repack_act_kernelspan)
  • 3.7 深入：w4a4_gemm_lora_kernel (https://blog.dotieuthien.com/posts/nunchaku-w4-quantization-memory-bound#span-stylecolor-1d4ed837-deep-dive-w4a4_gemm_lora_kernelspan)
  • 3.8 深入：pack_lora_act_kernel (https://blog.dotieuthien.com/posts/nunchaku-w4-quantization-memory-bound#span-stylecolor-1d4ed838-deep-dive-pack_lora_act_kernelspan)
  • 3.9 内核 4-5：RMSNorm + RoPE（注意力路径）(https://blog.dotieuthien.com/posts/nunchaku-w4-quantization-memory-bound#span-stylecolor-1d4ed839-kernels-4-5-rmsnorm–rope-attention-pathspan)
  • 3.10 内核绑定：CUDA → C++ → Python (https://blog.dotieuthien.com/posts/nunchaku-w4-quantization-memory-bound#span-stylecolor-1d4ed8310-kernel-binding-cuda–c–pythonspan)
• 4. CUDA Graph (https://blog.dotieuthien.com/posts/nunchaku-w4-quantization-memory-bound#span-stylecolor-1d4ed84-cuda-graphspan)
  • 4.1 为什么需要 CUDA graphs (https://blog.dotieuthien.com/posts/nunchaku-w4-quantization-memory-bound#span-stylecolor-1d4ed841-why-cuda-graphs-are-neededspan)
  • 4.2 问题：每个输入形状 = 一个新 graph (https://blog.dotieuthien.com/posts/nunchaku-w4-quantization-memory-bound#span-stylecolor-1d4ed842-the-problem-each-input-shape–a-new-graphspan)
  • 4.3 解决方案：flash_attn_varlen + bucket 预捕获 (https://blog.dotieuthien.com/posts/nunchaku-w4-quantization-memory-bound#span-stylecolor-1d4ed843-solution-flash_attn_varlen–bucket-pre-capturespan)
• 5. 基准测试 (https://blog.dotieuthien.com/posts/nunchaku-w4-quantization-memory-bound#span-stylecolor-1d4ed85-benchmarkspan)

动机

这项工作由我和 Triet Le (https://www.linkedin.com/in/triet-le-thong-minh-120217225/) 共同完成，是优化工作中 OmniVoice 推理的一部分。我们之前都没有 CUDA 内核开发经验——我们边做边学，用真实模型作为练习。大部分过程是与 Claude Code 结对编程：我们阅读 Nunchaku 源代码，询问每个概念的解释（MMA 布局、共享内存组冲突、warp shuffle…），然后一起编写内核。Claude Code 编写代码，我们负责审查并决定下一步做什么——选择设计方案、权衡利弊、调试性能分析、验证正确性。下面的结果来自边学习边构建的过程。

OmniVoice——带掩码流匹配的语音克隆模型（Qwen3 骨干网络，28 层，196 个线性层）——在 NVIDIA L4 上：

*RTF = 实时因子。0.056x 表示 1 秒音频在 56ms 内生成。 *W4A4 质量：INT4 激活会丢失信息——与 FP16 相比，某些音频段听起来稍微失真或带有轻微噪声。对于流式/实时使用场景，这是可接受的权衡，但尚未达到高保真语音克隆的生产级质量。

1. 性能分析：为什么优化前要先测量

1.1 性能分析：先测量

使用 torch.profiler 追踪一次完整的生成过程（32 步），导出 JSON 并在 Perfetto UI 中打开。追踪文件约 300MB，约 90 万个事件，约 5.8 万个 CUDA 内核。

Perfetto 概览——32 个解码步骤，每个步骤都是 GPU 流上的一个 llm_forward 块

放大后，可以看到 32 个 llm_forward 块在 GPU 流上重复（黄色行）。CPU 调用栈（下方）显示 patched_generate_iterative → OmniVoice.forward → Qwen3Model → 各层。

放大两个 CUDA 内核之间的间隙，揭示了真正的问题：GPU 在两个内核之间空闲——CPU 花费约 100μs 遍历 Python 包装器堆栈，然后启动下一个内核

GPU 空闲时间 >80%。 GEMM 内核 ampere_fp16_s16816gemm 只运行约 12μs，但之后的间隙拉长到约 100μs——GPU 等待 CPU 遍历 15 个以上的 Python 包装器层（nn.Module._call_impl → linear.forward → aten::matmul → aten::mm → cudaLaunch），然后才启动下一个内核（RMSNorm → aten::pow）。这正是 CUDA graphs 要消除的启动开销（第 4 节）。

Perfetto 时间线清晰地显示了这种模式，但仅通过观察很难判断工作负载是否受内存限制——我们需要数据。让 Claude Code 解析追踪 JSON 得到分解：

CUDA 内核总时间：1857 ms（32 步）

类别                时间(ms)   %
─────────────────────────────────────────────
GEMM (cuBLAS matmul)    674     36%
逐元素操作 (add/mul/cast)  620     33%
注意力 (flash/fmha)     393     21%
内存操作 (cat/copy)      111      6%
规约操作 (norm/mean)     45      2%
其他                    16      1%

两件事引人注目：

1. 逐元素操作占 33%——几乎和 GEMM 一样多。43,918 个小内核，平均每个 14μs。每个内核的实际计算只有约 1-2μs，但启动开销约 7μs/内核 × 44K = 307ms。逐元素操作时间的一半以上是开销，而非计算。
2. GEMM 平均 82-146μs——对于 GPU 来说很小。M=700（典型的 OmniVoice 批处理大小）、hidden=1024 的 cuBLAS matmul 每次调用只产生约 3 GFLOPS——L4 有 120 TFLOPS，所以每个 GEMM 在 82μs 内只使用了峰值计算的 0.002%。

两者都指向同一个结论：模型对于 GPU 来说太小了 → 内存受限。

1.2 屋顶线分析：确认内存受限

性能分析显示了问题，屋顶线分析用数学确认了这一点：

q_proj: Y[700, 2048] = X[700, 1024] × W[1024, 2048]
FLOPs = 2 × M × N × K = 2 × 700 × 2048 × 1024 = 2.94 GFLOP
   ↑ 乘以 2 是因为每个输出元素需要 K 次乘法和 K 次加法（GFLOP = 10^9 次浮点运算）
字节数 = 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/字节
L4 INT4 盈亏平衡点 = 485 TOPS / 300 GB/s = 1617 FLOP/字节
358 << 1617 → 内存受限（GPU 计算仅利用约 22%，其余时间等待数据）

(L4 规格：485 INT8 TOPS，300 GB/s 带宽 — NVIDIA L4 数据表 (https://www.nvidia.com/en-us/data-center/l4/)。INT4 吞吐量 = 2× INT8。)

OmniVoice 的 hidden_size=1024——对于 GPU 来说太小了。与 FLUX (https://huggingface.co/black-forest-labs/FLUX.1-dev)（hidden=3072）对比：

FLUX 线性层 (M=4096, K=3072, N=3072):
FLOPs = 2 × 4096 × 3072 × 3072 = 77.3 GFLOP
字节数 = W(18MB) + X(25MB) + Y(25MB) = 68 MB
AI = 1137 FLOP/字节

OmniVoice AI = 358 → 内存受限（GPU 等待数据）
FLUX AI = 1137 → 接近盈亏平衡点 1617（GPU 几乎饱和计算）

FLUX 足够大，INT4 张量核心能利用约 70% 的峰值 TOPS → Nunchaku 通过计算加速。OmniVoice 的 FLOPs 少约 26 倍 → INT4 通过减少字节数来加速，而不是通过 TOPS。

1.3 理论：内存受限时，减小数据宽度 = 加速

GPU 空闲时间 >60% 等待 HBM。增加 TOPS 没有帮助。有帮助的是减少加载的字节数：

每次 GEMM，带宽 = 300 GB/s：
        权重      激活      总数据     时间
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

理论上：权重带宽是主要瓶颈（当 M 较小时，权重 >> 激活）。W4A8 和 W4A4 的权重带宽相同——唯一的区别是激活（0.35 MB）。所以在速度上，任何 W4 变体在内存受限的工作负载上表现几乎相同。

1.4 实践中：选择 W4A4 是因为 Nunchaku 已存在

我确实尝试通过 TRT-LLM（ModelOpt AWQ 量化）构建 W4A8，但遇到了太多兼容性问题：ModelOpt 不支持 Qwen3 导出，transformers 版本冲突，TRT-LLM 检查点格式不匹配。经过多轮调试依赖地狱后，得出结论：在 TRT-LLM 0.18.2 + Qwen3 上实现 W4A8 目前还不可行。

选择 W4A4 是出于实际原因：Nunchaku (https://github.com/mit-han-lab/nunchaku)（MIT-HAN Lab）已经有一套完整的 W4A4 GEMM 内核实现用于扩散模型——包括 INT4 MMA、融合 LoRA 和权重打包。基于这个设计编写自定义内核，比调试 W4A8 TRT-LLM 流水线快得多。根据上面的理论，在内存受限的工作负载上，W4A4 和 W4A8 差异很小（约 1μs/GEMM）。唯一的权衡是质量：INT4 激活比 FP8 损失更多。但借助 SVDQuant LoRA 修正（第 2 节），质量仍然可以接受。

2. 量化流水线：FP16 → INT4 检查点

2.1 问题：INT4 太粗糙，需要 LoRA 修正

INT4 只有 16 个离散值。直接将 FP16 → INT4 量化：

W_fp16 = [0.127, -0.891, 0.003, ...] (65536 个不同值)
W_int4 = round(W / scale) × scale (只有 16 个值)
每个元素误差：~5%
误差在 28 层中累积 → 输出严重退化

SVDQuant (MIT-HAN Lab) 解决了这个问题：将量化误差分解为低秩矩阵，在运行时用 LoRA 补偿：

量化：E = W - Q(W)  ← 误差矩阵 [N, K]
U, S, V = SVD(E, rank=16)  ← 低秩近似
proj_up = U[:, :16] × S[:16]  ← [N, 16]
proj_down = V[:16, :].T  ← [K, 16]
运行时：Y = INT4_GEMM(X, Q(W)) + X @ proj_down @ proj_up.T ≈ X @ W  ← FP16 精度恢复
使用 SVDQuant LoRA 修正后：质量可接受

2.2 GPTQ + 迭代 SVD

每层的量化流水线：

for svd_iter in range(3):
    # 步骤 1：GPTQ 量化（基于 Hessian 矩阵的舍入）
    Q_W = gptq_quantize(W - LoRA, H=X^T @ X, group_size=64)
    # 步骤 2：对残差误差进行 SVD
    error = W - Q_W
    proj_up, proj_down = SVD(error, rank=16)
    # 步骤 3：检查收敛
    if new_error >= old_error:
        break

GPTQ 使用 Hessian 矩阵 H = X^T @ X（来自校准数据）来决定舍入：激活值大的列获得更谨慎的舍入。迭代 SVD（3 轮）：每轮量化 W - LoRA_previous，然后重新计算 SVD → LoRA 修正变得更精确。

每层输出：qweight（INT4 打包），w_scales（FP16 每组），proj_down [K, 16]，proj_up [N, 16]。

2.3 检查点格式

w4a4_checkpoint.safetensors:
layers.0.self_attn.q_proj.qweight: [2048, 512] uint8 (INT4 打包，每字节 2 个)
layers.0.self_attn.q_proj.w_scales: [2048, 16] fp16 (每组 64 的缩放因子)
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 层

运行时：W4A4Linear 加载检查点 → 将权重重新打包为 MMA 张量核心布局（离线，一次） → forward 使用自定义 CUDA 内核。

3. CUDA 内核：INT4 MMA + 融合 LoRA

3.1 OmniVoice 模型结构

OmniVoice 骨干网络是一个 Qwen3 transformer——28 个相同的层，每层包含：

输入 hidden_states [batch, seq_len, 1024]
│
├─ input_layernorm (RMSNorm)
├─ self_attn:
│   ├─ q_proj [1024 → 2048] ← Linear (16 个头 × 128 维)
│   ├─ k_proj [1024 → 1024] ← Linear (8 KV 头 × 128 维)
│   ├─ v_proj [1024 → 1024] ← Linear
│   ├─ q_norm, k_norm (每个头 RMSNorm)
│   ├─ RoPE (旋转位置编码)
│   ├─ Attention (Q × K^T → softmax → × V)
│   └─ o_proj [2048 → 1024] ← Linear
├─ 残差相加
├─ post_attention_layernorm (RMSNorm)
├─ mlp:
│   ├─ gate_proj [1024 → 3072] ← Linear
│   ├─ up_proj [1024 → 3072] ← Linear
│   ├─ SiLU(gate) × up
│   └─ down_proj [3072 → 1024] ← Linear
└─ 残差相加
│
输出 hidden_states [batch, seq_len, 1024]

每块 7 个 Linear 层 × 28 块 = 196 个 Linear 层——这些是 W4A4 的量化目标。Norm、注意力和 RoPE 保持 FP16。

3.2 CUDA 概念 + 为什么需要自定义内核

CUDA 术语——阅读内核代码前需要了解：

为什么需要自定义内核？ W4A4 需要额外步骤：量化 X，按组反量化，添加 LoRA 修正。如果每个步骤都作为单独的内核运行 → X 从 HBM 读取 2-3 次 → 内存受限，所以比 FP8 还慢（第 1 节 (https://blog.dotieuthien.com/posts/nunchaku-w4-quantization-memory-bound#1-profiling-why-measure-before-optimizing)）。Nunchaku (MIT-HAN Lab) 解决了这个问题。

3.3 内核组织 — C++ 绑定

文件结构：

cuda_kernels/
├── w4a4_gemm_lora.h         # 主内核 (kernel 2)：INT4 GEMM + 融合 LoRA
├── fused_quant.h            # 内核 1：融合量化 + 重新打包
├── pack_lora_act.h          # 内核 3：打包 LoRA 激活用于后续步骤
├── rmsnorm.cu               # 内核 4：NVIDIA 的 fused_rmsnorm (https://github.com/NVIDIA/transformer-libs/blob/main/sm8x/transformer_engine/rmsnorm.cu)
├── rope.cu                  # 内核 5：融合 RoPE (cos/sin 应用)
├── flash_attn.py            # Flash attention (不自定义，使用 Dao-AILab 的)
├── cuda_ext.cpp             # pybind11 绑定和 launch 封装
└── setup.py                 # pip install -e .

3.4 流水线：1 个 Linear 层需要什么？

一个 W4A4 Linear 层的前向传播流水线。目标：在 1 个 CUDA graph 中捕获以下所有内容（第 4 节）：

输入: X [M=700, K=1024] FP16
│
1. fused_quantize_repack_act_kernel (w4a4_gemm_lora 的准备工作)
   ┌────────────────────────────────────┐
   │ 加载 X tile → 量化到 INT4 → 打包   │
   │   → 保存到 SMEM                     │
   │ 加载 W tile (INT4 打包) → MMA 准备   │
   └────────────────────────────────────┘
   │
2. w4a4_gemm_lora_kernel (主计算)
   ┌────────────────────────────────────┐
   │ INT4 MMA: X_int4 × W_int4          │
   │ + 反量化 (每组缩放因子)              │
   │ + LoRA 修正: X_fp16 @ proj_down    │
   │   → 融合：从 SMEM 读取同一个 X tile  │
   │ → 输出 Y [M, N] FP16                │
   └────────────────────────────────────┘
   │
3. pack_lora_act_kernel (后续层的准备)
   ┌────────────────────────────────────┐
   │ 如果下一层也需要 LoRA:              │
   │ 保存 X_fp16 @ proj_down 到激活内存  │
   │ 供下一层的 LoRA 使用                 │
   └────────────────────────────────────┘
   │
输出: Y [M, N] FP16

3.5 为什么 tile 大小是 128×128，8 个 warp？

NVIDIA 张量核心的 INT4 MMA 指令是 m16n8k64：

• M 维度：16（跨 warp 的线程）
• N 维度：8（每个线程持有 8 个 INT4 值 × 32 线程 = 256 个 INT4 = 8 列输出？需要更精确）
• K 维度：64（INT4 值，每组缩放因子应用于 K=64 组）

一个 warp 一次计算 16×8 输出 tile。128×128 tile 意味着：

• 沿 M：128/16 = 8 个 warp → 每个 warp 负责 16 行
• 沿 N：128/8 = 16 个 warp“列”组 → 如果每个 warp 持有 8 列，需要 16 个 warp？
• 但 Nunchaku 每个块使用 8 个 warp，原因：L4 有 48KB SMEM，每个块 128×128 tile 的共享内存预算：
  • X tile：128×128 INT4 = 128×128×0.5B = 8KB
  • W tile：128×64 INT4（加上 4×因子）= ~5KB
  • LoRA 激活：128×16 FP16 = 4KB
  • 总计：约 20KB → 48KB 可以容纳 2 个 tile（乒乓缓冲区），每个 8 个 warp。

为什么是 128×128？ 平衡 M 和 N 维度：M 通常是 700（批处理大小 × 序列长度），N 从 1024（o_proj）到 3072（gate_proj）。128×128 对于所有 M/N 组合都足够，同时保持 SMEM 在限制内存占用。

3.6 深入：fused_quantize_repack_act_kernel

这是所有准备工作的内核。它加载 X 的 FP16 值，将其量化到 INT4，然后以 m16n8k64 MMA 布局重新打包，以便 w4a4_gemm_lora_kernel 可以直接读取它。

__global__ void fused_quantize_repack_act_kernel(
    const half* __restrict__ X,           // 输入激活 [M, K] FP16
    uint8_t* __restrict__ X_int4,         // 输出量化激活 [M, K/2] INT4 打包
    float* __restrict__ X_scale,          // 输出缩放到因子 [M] (每个 token 1 个缩放因子)
    int M, int K
) {
    // 每个块处理 128×128 个 X
    int tile_row = blockIdx.y;    // 沿 M
    int tile_col = blockIdx.x;    // 沿 K
    int tid = threadIdx.x;        // 块内线程 ID
    
    // 加载：FP16 X → 寄存器和 SMEM
    // 对于 128×128 tile，每个 warp 加载 16×128 个元素
    // 使用 vectorized (float4 = 8 FP16) 以利用 128B 内存事务
    float4 x_chunk = *reinterpret_cast<const float4*>(
        &X[(tile_row * 128 + warp_row) * K + tile_col * 128 + col_offset]
    );
    
    // 量化：FP16 → INT4（每个 warp/lane 原子操作）
    half x_vals[8];    // 从 float4 解包
    float absmax = 0.0f;
    #pragma unroll
    for (int i = 0; i < 8; i++) {
        absmax = fmaxf(absmax, fabsf(__half2float(x_vals[i])));
    }
    
    // warp 级规约：找出这 16 行中的最大 abs (每个 warp 处理 16 行)
    float warp_absmax = warp_reduce_max(absmax);  // __shfl_xor_sync
    
    // 共享缩放因子：整个块共享 1 个？不，每个 warp/行？Nunchaku 对 X 做跨 K 的 per-tensor 量化
    // 每组 128 列使用 1 个缩放因子 → 对于 128 列 tile 是 1 个缩放因子
    if (lane_id == 0) {
        atomicMax(&X_scale[tile_row * 128 + warp_row], warp_absmax);
    }
    __syncthreads();
    
    float scale = X_scale[tile_row * 128 + warp_row] / 7.0f;  // INT4 范围为 -8..7
    
    // 量化为 INT4
    uint8_t packed_byte = 0;
    #pragma unroll
    for (int i = 0; i < 4; i++) {  // 每字节打包 2 个 INT4
        int quant_val = __float2int_rn(x_vals[i] / scale);  // 量化
        quant_val = max(-8, min(7, quant_val));              // 截断
        packed_byte |= (quant_val & 0xF) << (i * 4);         // 打包
    }
    
    // 以 MMA 友好布局写入 SMEM
    // 第 5/6 页有布局图：X_int4 以 "列主序" 打包，以便快速加载到 MMA 寄存器
    __shared__ uint8_t X_smem[128 * 128 / 2];  // INT4 打包
    // ... SMEM 解包布局，使每个 warp 可以读取 16×64 块用于 MMA ...
    X_smem[smem_idx] = packed_byte;
    
    // 同步：所有 warp 必须在 SMEM 中都准备好 X 块
    __syncthreads();
    
    // 从 SMEM 以 MMA 布局重新打包到 HBM（供后续内核使用*）
    // ... 将 X_smem 转置写入 X_int4 ...
}

省略了大量细节：

• Warp 级规约：使用 __shfl_xor_sync 在一个 warp 内的 32 个线程中找出最大值——比块级 atomicMax 快得多，因为整个 warp 在同一个锁步中执行（无需同步）。
• SMEM 布局：NVIDIA 的 MMA 指令要求数据以特定布局存储——连续的线程读取连续的 128 位。Nunchaku 以一种布局写入 SMEM，然后以另一种布局读取，这样 MMA 加载就是合并的。
• 乒乓缓冲区：同时处理 2 个 tile——一个正在被 w4a4_gemm_lora_kernel 处理，另一个正在被 fused_quantize_repack_act_kernel 量化。这不在此内核中；它由启动逻辑处理。

3.7 深入：w4a4_gemm_lora_kernel

这是主内核：INT4矩阵乘法、反量化和LoRA修正都融合在一起。

// 简化：实际代码使用共享内存和寄存器分片以实现最佳布局
// 目标设备架构：sm_89 (L4 Ada Lovelace)

__global__ void w4a4_gemm_lora_kernel(
    const uint8_t* __restrict__ X_int4,    // 量化激活 [M, K/2] 以 Nunchaku 布局打包
    const float* __restrict__ X_scale,     // per-tensor 缩放因子 [M]
    const uint8_t* __restrict__ W_int4,    // 量化权重 [N, K/2] 以 MMA 布局打包
    const half* __restrict__ W_scale,      // per-group 缩放因子 [N, K/GROUP]
    const half* __restrict__ proj_down,    // LoRA down [K, rank] FP16
    half* __restrict__ lora_act,           // 输出激活 @ proj_down [M, rank] FP16
    half* __restrict__ Y,                  // 输出 [M, N] FP16
    int M, int N, int K, int rank
) {
    // 每个块：128×128 输出 tile
    // 每个 warp：16×8 输出