AsyncT vllm适配、加速笔记（二）

续着前文继续做优化。

1. CUDA 优化扫尾

这里的一些内容严格意义上应该放在上篇内容的末尾，一阶段还有一些比较平凡的 CUDA 优化。

1.1. Cooperative QK

原 CUDA kernel 每个 thread 顺序算 HEAD_SIZE=128 的完整 QK dot。改 THREAD_GROUP_SIZE=8，每个 thread 负责 16 dim，然后 warp shuffle reduce：

constexpr int THREAD_GROUP_SIZE = 8;
const int thread_in_group = tid % THREAD_GROUP_SIZE;

for (int t = group_id; t < seq_len; t += K_TOKENS_PER_ITER) {
    float partial = 0.0f;
    #pragma unroll
    for (int i = 0; i < ELEMS_PER_THREAD; i++) {
        const int d = thread_in_group * ELEMS_PER_THREAD + i;
        partial += q_smem[d] * to_float(k_ptr[d]);
    }
    // Warp shuffle reduce
    for (int mask = THREAD_GROUP_SIZE / 2; mask > 0; mask /= 2) {
        partial += __shfl_xor_sync(0xffffffffu, partial, mask);
    }
    if (thread_in_group == 0) {
        logits[t + FIR_K - 1] = partial;
    }
}

结果写完发现没精度，是因为__shfl_xor_sync(0xffffffffu, ...)用的是 32-lane warp mask，但一个 warp 里有 4 个 group（每个 8 threads）。当 seq_len % K_TOKENS_PER_ITER != 0，不同 group 的外层 loop 迭代数不同，某些 lane 已经退出 loop 而其他 lane 还在 shuffle，all-lanes mask 触发 deadlock。修复：

const int lane_in_warp   = tid & 31;
const int group_in_warp  = lane_in_warp / THREAD_GROUP_SIZE;
const unsigned int group_mask = 0xFFu << (group_in_warp * THREAD_GROUP_SIZE);
for (int mask = THREAD_GROUP_SIZE / 2; mask > 0; mask /= 2) {
    partial += __shfl_xor_sync(group_mask, partial, mask);  // group_mask, NOT 0xffffffff
}

1.2. 假装 prefill 是 decode

做的时候发现 prefill 的占比实际上比想象中大很多，之前的 prefill kernel 实际上用的还是 torch path，效率太差了。但重新写一个 prefill kernel 之后每次优化都需要两边一起测试一起调，难度加倍。在和 codex 的讨论中，它提示：

对每个 Q token，它的有效 causal seq_len 就是 prefix_len + local_q_pos + 1。也就是说，每个 Q token can be treated as its own “decode” with custom seq_len。我们只需要 ”fake” 出 decode-shape 的 metadata：

# 对prefill里每个request，展开其所有Q token变成fake decode
fake_orig_req   = torch.repeat_interleave(p_orig_reqs_t, p_t_qs_t)
expanded_cum_tq = torch.repeat_interleave(cum_tq[:-1], p_t_qs_t)
local_q_pos     = torch.arange(total_tq, ...) - expanded_cum_tq
expanded_prefix = torch.repeat_interleave(p_prefix_t, p_t_qs_t)
fake_seq_lens   = (expanded_prefix + local_q_pos + 1).to(torch.int32)
fake_block_table = block_table[fake_orig_req].to(torch.int32)

# 调用同一个Triton decode kernel
torch.ops._C.hp1_decode_triton(
    output, query[first_q_start:last_q_end],
    key_cache, value_cache, fake_block_table, fake_seq_lens, ...)

就能把 prefill 伪装成 decode，复用已经优化过的 Triton decode kernel，不需要写新的 prefill kernel。

进一步地，上面的 fake prefill metadata 也能避免每层重建，并且略微调整了它的 split k（因为发现 fake decode request 在这里的长度一般不大），从这里一下子能获得～30% 的 gen time 收益，证明检查 workload 形态真的非常重要。

1.3. AsyncI 的简易稀疏优化

回忆 AsyncI 的算法结构，clamp(score - sigma * AvgPooling(ReLU(score)), 0, 1)与 softmax 几乎永远不产生任何 0 不同，里面有 0 的概率相当大，因而在 PV 阶段可以直接跳过对应的 V Load 和 FMA：

for (int t = 0; t < seq_len; t++) {
    const float p = p_buf[t];
    if (p == 0.0f) continue;   // skip V load + FMA when gate is exactly zero
    // ... 正常PV ...
    acc += p * to_float(*v_ptr);
}

做稀疏不可避免地要讨论 warp divergence 问题：128 个 thread 都在执行同一个 if (p == 0.0f)（p_buf[t]是 shared memory，所有 thread 看到同一个值），所以 branch 是在一个 warp 内的行为都是完全一样的，不会有 divergence 问题！

加上去之后只带来了～2% 的提升，主要原因是此处的优化过于 naive。在之后这里还有进一步发掘的空间。

间章：相比于 FA2？FA3 呢？

到目前为止，我们已经实现了非常多（非常常规）的优化，得到的结果：

HP1 wall-time 较 softmax 快 45%（init 开销差异），但生成阶段为 1.286× softmax。考虑到使用的两个模型层数不太一样（AsyncT 35 Layers vs. Qwen3 28 Layers），这个轻微的 Gen Time 还比较容易接受。

然而，Lab 这边能够支持的默认镜像源里面只有 FA2，FA3 需要自己编译，而我们想 Claim 的很多优势都需要 Hopper 架构上的 TMA、Warp Specialization 等方面进行，和 FA2 的对比显然不太公平。因此从这里开始，我们的 baseline 需要改进到 FA3 级别。改进后：

Qwen3 Baseline + FA2:    154s wall, 49s gen
Qwen3 Baseline + FA3:    56s wall, 13s gen   (3.6×!)
AsyncT:                  84s wall, 63s gen

差距还是很大！进一步做 Attention 之外的 breakdown profiling：

看起来差距主要体现在 Attention 上，如果我们能做到 Attention 速度对齐甚至超越，模型整体的速度应该就会有一个比较好的提升。进一步观察：

AsyncT prefill-only:  74,556 tok/s
Softmax FA3 prefill: 297,707 tok/s

原来是之前偷懒用的 ”假装 Prefill 是 Decode” 成为了瓶颈。看来必须得手写一个 Prefill Kernel 了。

1.4. Multi-Q Prefill CUDA Kernel

第一版：

grid = (q_head, prefill_req, q_tile)
BLOCK_Q = 16
THREADS_PER_Q = 8
每个 CTA 处理 16 个 query rows
流式遍历 K
每个 K position 做 QK、FIR/gate、PV

理论上它能复用 K/V load，但实际 gsm8k 回退，因为每个 K position 都有 __syncthreads，小 T_q 时同步成本高。sweep BLOCK_Q：

BLOCK_Q=16 -> gsm8k 112s
BLOCK_Q=8  -> 79s
BLOCK_Q=4  -> 71s, gen 51s, pure prefill 134K tok/s

有点反直觉，但似乎也能解释：小 BLOCK_Q 增加 block 数，提高 occupancy，H200 有 132 SM，不能用太少 CTA 喂它。

再尝试做 K-Tiling，这样遍历 K 的过程中能减少一些 sync 的问题，反而造成了性能的回退。通过 nsight compute 进行观察，发现现在是 CUDA core compute/instruction bound，DRAM Utilization 仅 0.1%。

到此为止，CUDA Core 能做的优化已经基本做完了，还想获取算力上的巨大提升，就得靠 Tensor Core 了，下一步开始走向 MMA Tensor Core。

2. MMA Tensor Core

直接从 inline 开始，让 QK 和 PV 都走 tensor core：

asm volatile(
  "mma.sync.aligned.m16n8k16.row.col.f32.bf16.bf16.f32 "
  "{%0, %1, %2, %3}, {%4, %5, %6, %7}, {%8, %9}, {%10, %11, %12, %13};\n"
  : "=f"(d[0]), "=f"(d[1]), "=f"(d[2]), "=f"(d[3])
  : "r"(a[0]), "r"(a[1]), "r"(a[2]), "r"(a[3]),
    "r"(b[0]), "r"(b[1]),
    "f"(c[0]), "f"(c[1]), "f"(c[2]), "f"(c[3]));

Fragment layout（per PTX docs）：

• A (M=16, K=16)：每个 32 lane 持 4×uint32 (=8 BF16)
• B (K=16, N=8 col-major)：每 thread 持 2×uint32 (=4 BF16)
• C/D (M=16, N=8)：每 thread 持 4×fp32

Kernel 结构：

• QK via mma_m16n8k16(8 mma per K-tile, 覆盖 HEAD_SIZE=128 in d-chunks of 16)
• FIR + clamp gate (sequential per Q row, register tail [FIR_K-1=2])
• PV via mma_m16n8k8(16 mma per K-tile, covers HEAD_SIZE=128 in d-chunks of 8)

Micro Bench 上效果还不错，但 gsm8k 的 wall time 反而有回退。猜想是反复调用的问题？正好刷知乎看到有人讲 persistent kernel，直接用 atomicAdd 加上去试试看，结果效率大幅度降低 3×。手算一下：264 个 CTA atomicAdd 到同一个 counter 地址 = massive contention。

改成每个 warp 拥有自己的 Shared Memory Slice 用来存内容，warp 之间独立 pull work 来解决 Warp Contention Compromise：

// 4 warps × 32 lanes = 128 threads/CTA
const int warp_id = tid / 32;
const int lane = tid & 31;
constexpr int NUM_WARPS = 4;

__shared__ int work_id_smem[NUM_WARPS];
while (true) {
    if (lane == 0) work_id_smem[warp_id] = atomicAdd(work_counter, 1);
    __syncwarp();
    int work_id = work_id_smem[warp_id];
    if (work_id >= total_work) return;
    // Each warp processes independently, uses LANE not TID, no inter-warp sync
}

得到 68s wall / 47s gen 的结果，提升不错。既然已经开始用 mma.sync 了，也许现在计算性能已经不是瓶颈，是不是该考虑尝试优化访存相关的内容？开始尝试：

1. 从 manual pack_bf16 换成 ldmatrix.sync.aligned.m8n8.x4.shared.b16。Short scoreboard stall 从 1.37 → 1.10，但 tensor op active 从 1.47% → 1.53%（好像没啥用）。
2. 用 cp.async，但性能几乎没发生什么变化，Long scoreboard 7.87 stall 是 ldmatrix 的过程，cp.async 解决的主要是 global load，跟这里没什么关系
3. 扫 Block_K 的 size：BK=8: 0.60 ms baseline，BK=16: 0.62 ms (-3%)，BK=32: 1.13 ms (-47%!)，看起来是因为每个 SM 能读取到的块变大之后、L1 Cache 难以通过 Warp 间的合适 pipeline 掩盖了；

这些尝试证明：用了 tensor core” 不等于 “tensor core 是主导，如果每次 MMA 周围有太多 scalar/shared-memory work，tensor pipe active 仍然可能低到 1-2%！

看来 MMA 和 WGMMA 还是差的很远，进入 warpgroup-level async MMA.

3. Hopper WGMMA Prefill

3.1. 先解决正确性问题

从 QK-only 开始：

__device__ __forceinline__
void wgmma_m64n8k16_f32_bf16(
    float& d0, float& d1, float& d2, float& d3,
    uint64_t desc_a, uint64_t desc_b, int scale_d) {
  asm volatile(
      "{\n"
      " .reg .pred p;\n"
      " setp.ne.b32 p, %6, 0;\n"
      " wgmma.mma_async.sync.aligned.m64n8k16.f32.bf16.bf16 "
      " {%0, %1, %2, %3}, %4, %5, p, 1, 1, 0, 0;\n"
      "}\n"
      : "+f"(d0), "+f"(d1), "+f"(d2), "+f"(d3)
      : "l"(desc_a), "l"(desc_b), "r"(scale_d));
}

最开始把 K-Major 理解为普通的 row-major[row][d]，结果只有每个 N-Block 的第一列正确。查阅资料才知道，WGMMA 需要 8-row interleaved bricks：

int wgmma_kmajor_idx(int row, int d, int row_stride_elems) {
  return (row >> 3) * (8 * row_stride_elems)
       + (d >> 4) * (8 * 16)
       + ((d >> 3) & 1) * (8 * 8)
       + (row & 7) * 8
       + (d & 7);
}

并把 descriptor LBO 从 16 bytes 改为 128 bytes。修完后 QK scores 和 mma reference 对齐。

正确性问题解决，开始把全流程适配：WGMMA QK → FIR Gate → WGMMA PV → output，结果小 batch 正确、大 batch 结果出错、compute-sanitizer 又不报 OOB。最后发现是 Hopper shared-memory proxy ordering 问题。普通 shared store 到 Q/V/P 后，WGMMA async proxy 读取 shared memory，不能只靠 __syncthreads()；需要在 Q/Tail 初始化之后、V Page Load 之后、P Tile 形成后 - PV WGMMA 读取前添加

__device__ __forceinline__ void fence_proxy_async_shared_cta() {
  asm volatile("fence.proxy.async.shared::cta;" ::: "memory");
}

之前每个 dc 一个 wgmma_fence/commit/wait_group<0>，serialize 了 16 个 async group per page。改成一次 fence、16 个 mma issue、一次 commit/wait：

wgmma_fence();
for (int dc = 0; dc < N_D_PV_W; dc++) {
  wgmma_m64n8k16_f32_bf16(...);
}
wgmma_commit();
wgmma_wait_group<0>();

修复后 full WGMMA path 正确，B32/T256 达到 0.219ms vs MMA 0.584ms，2.67x。

3.2. Parallel FIR Gate

之前的做法中仍然是 QK scores 写到 smem 然后 64-thread 顺序遍历每个 Q row 的 16 K position。观察到，HP1 的 FIR 只依赖当前 score 和前两个 score，而 WGMMA fragment 的每行 8 列已经分布在一个 4-lane quad 中：

lane q:     score[k=2q], score[k=2q+1]
prev lane:  score[k=2q-2], score[k=2q-1]

const float prev1_lane = __shfl_up_sync(0xffffffff, r1, 1);
const float prev2_lane = __shfl_up_sync(0xffffffff, r0, 1);
const float prev1 = (lane_quad == 0) ? tail0 : prev1_lane;
const float prev2 = (lane_quad == 0) ? tail1 : prev2_lane;

0.219 → 0.172 ms, 3.39×. Tensor pipe active 3.1 → 3.8-3.9%。scores_smem 变成 dead code，从 shared memory carveout 里移除了 4KB，registers/thread 125 → 122。ncu 中 SM throughput 从约 29% 到 39%，long scoreboard 从约 3.09 到 2.28–2.30。

这个改进比较有趣的地方在于，是通过观察 AsyncT 的数学结构得来的，而非平凡的 GPU 上的性能优化。这应该是第一个可以声称是 AsyncT 特异的加速收益。

进一步把 m64n8k16 改成 m64n16k16，一次覆盖完整 16-token page，PV 输出维度也一次 16 dims。结果 0.172ms 到 0.165ms，提升有限。FA3 在相同的 size (B=32, T=256, H=8/HKV=4, D=128) 是 0.052 ms. WGMMA AsyncT 差不多慢 3.2×，比之前的 4.85× 好一点。

ncu 显示 tensor pipe active 只到约 4.0–4.2%，active warps 约 21%。这说明瓶颈不只是 WGMMA instruction count，大概率是因为

• 寄存器占用太高→一个 WarpGroup CTA 只有四个 warp？
• page-level barrier？
• K/V 的 Staging？
• pipeline 设计还有问题？

4. 重构 Benchmark

做到这里，从 gsm8k 和几个端到端 lm-eval 中抽取 benchmark 模拟各种场景实在是太复杂了，开始构造新的 Benchmarks：

prompt 由固定 filler 文本 tokenizer 后裁到指定 token 数，通过 prompt_token_ids 直接喂给 vLLM；decode 用 min_tokens=max_tokens, ignore_eos=True, temperature=0，所以每次输出长度是精确固定的。之后性能的描述主要通过这些 Benchmark 描述，而 lm-eval 只用于大的修改后评估总体 serving 的能力（以及正确性）。

将上面的 kernel 接入 vllm 时，由于很多没有测试的情况，给 WGMMA Dispatch 加了一些限制：

VLLM_HP1_USE_CUDA=1
VLLM_HP1_PREFILL_WGMMA=1
H * D == 1024
block_size == 16
dtype == bf16
fir_k == 3
max_t_q >= VLLM_HP1_WGMMA_MIN_TQ

测试 WGMMA hit/fallback counter：

prefill_heavy_b16: 0.3418s -> 0.2146s, 1.59x
balanced/mixed:    只有约 1.02-1.04x
decode-heavy:      基本不变

lm-eval 中 WGMMA hit token ratio 只有 21.18%，大部分真实 limited-task dispatch 低于 MIN_TQ=64。

详细做 nsys breakdown，发现：

• hp1_decode_attention34.4%
• gemm_matmul27.6%
• norm13.7%
• sampler7.1%
• hp1_prefill_wgmma2.4%

Decode 反而才是现在的瓶颈！

5. Decode 加速

当前的 Decode Kernel:

grid = (q_head, request)
每个 CTA 处理一个 request 的一个 q_head
shared q_smem[128]
QK:
  THREAD_GROUP_SIZE=8
  每 8 个线程合作一个 K token dot
  写 logits
FIR/gate:
  shared logits -> p_buf
PV:
  tid < 128 的线程累积输出维度
  if p == 0 skip V load

在短 context 下还不错，但长 context 时 grid 太小、每 CTA 工作太长，SM 利用率不够。

当前每个(request, q_head)一个 CTA 的设计完全填不满 GPU！

5.1. GQA-pair decode & streaming decode & Lane-interleaved QK Load

一个很自然的优化是：GQA 下两个 q_heads 共享一个 kv_head，能否一个 CTA 同时算两个 q_heads？这在数学上肯定没什么问题，因为它们确实共享 K/V，但实际跑下来效果并不好，因为 K/V 的 Load 省了，但是 shared state 翻倍、CTA 数量减半了，occupancy 反而更差、也更不好 hide 各种 memory latency。

另一个尝试是 streaming decode：用 AsyncT 的 local FIR 结构把 shared memory 从 MAX_SEQ_LEN 的 full logits/p_buf 缩到一个小 tile，每个 tile 走一次 QK→gate→PV。Shared memory 确实少了，但每个 tile barrier 都是新的 cost，导致所有 tile size 都比 full-buffer 慢。Sync overhead 吃掉了节省下来的 shared mem。

又尝试用 d = i * THREAD_GROUP_SIZE + thread_in_group 代替 thread_in_group * ELEMS_PER_THREAD + i，纸面上讲上应该改善 coalescing，结果反而更差。原 contiguous-per-lane loop 胜在更简单的 scalar stream 和更好的 local cache scheduling。Compiler/L1 对 ”每个 lane 读一段 contiguous” 这种模式优化得比 ”interleaved” 更好。

5.2. Split-K decode

那既然问题是 grid too small，那就通过减少每个 CTA 的大小，大幅度增加 CTA 数量。把每个 request 的 K-range 拆分成多个 split、每个 split 都是独立的 CTA：

stage1 grid:  (q_head, req, kv_split) → fp32 partial out
stage2 grid:  (q_head, req)           → sum partials and cast to output dtype

而且这个做法也完全受益于 Async T 的结构：因为没有 softmax 的 rescale merge。每个 split 只要多读左侧 FIR_K-1 个 score 作为 FIR overlap，就能正确计算本 split 内的 P，然后 partial P @ V 最后直接 sum。Softmax decode 想做 split-K 的话需要保留 row max + row sum 的 LSE state，stage2 要做 scale matching。

结果 split256 在小 batch 上最大直接得到 3-4× 提升，但在 split512 上效果不佳，预期是每个 split 的工作量超出 CTA 的 latency-hiding 能力。

vLLM CUDA Graph 问题

最开始的时候用了一个很简单的判断：

if max_seq_len >= VLLM_HP1_DECODE_SPLIT_MIN_LEN:
    split_size = 256

在 eager mode 下没有问题，但启用 CUDA Graph 后此优化完全失效。原因是：CUDA graph capture 时看到的是 dummy short decode lengths（vLLM warmup 用的是短 sequence），所以 captured graph 走的是 no-split branch。Later replay 的时候即使 max_seq_len=1026 是长 context，replay 还是跑那个 no-split branch——graph 节点的 dispatch path 是 capture 时定死的，runtime 读 max_seq_len 没用。

修复方法是把判断条件从 runtime tensor 改成 static capacity：

split_trigger_len = (
    block_table_blocks * block_size           # 静态capacity
    if VLLM_HP1_DECODE_SPLIT_STATIC=1
    else max_seq_len                          # 动态max_seq_len
)
VLLM_HP1_DECODE_SPLIT_AUTO=1
VLLM_HP1_DECODE_SPLIT_MIN_LEN=1024
VLLM_HP1_DECODE_SPLIT_SIZE=256
VLLM_HP1_MMA_MIN_TQ=64

vLLM graph 捕获看到的是 capacity/bucket，而不是每一步真实长度。只要 bucket capacity 足够大，就应该捕获长上下文 decode 的 split graph。后续 replay 里真实 max_seq_len 在这个 bucket 内变化，但 replay 的 kernel body 是稳定的。

MMA_MIN_TQ=64 是另一个容易忽略的修正。长 decode 的 serving 场景中会出现一些 16-token 左右的 chunked prefill tail。之前这些 tiny prefill tail 会掉进旧的 MMA full prefill path，结果一个很小的 prefill tail 反而占掉大量 kernel time。profile 显示，调高 MMA_MIN_TQ 后，hp1_prefill_mmabucket 从 top attribution 中消失；tiny prefill tail 重新交给 Triton synthesis，而长 prefill 继续走 WGMMA。

修复之后遇到了一个很微妙的情况：

这个结论比较微妙，最终决定做成一个之后可以随时插拔的 flag，暂时不添加到主线上。

6. profiling 改进和 norm/WGMMA 的其他小问题

做到这里，decode 优化一口气也做了不少，整体 decode kernel 在 long-context 上已经有了像样的工具，但还有几个比较烦的现象一直没解决：

• 每次启用 CUDA Graph 的时候 wall-time 反而劣化；
• nsys 里反复看到 _topk_topp_kernel、FillFunctor 等 ”奇怪” 的 kernel；
• norm 在所有 scenario 里都稳定占 11-13% 的 GPU kernel time，但它好像并不在 ”主线” 上。

同时在做上面的 decode 优化流程中，也 come up 了一些新的点子，遂暂时叉出一条支线，来做一些其他的优化。

6.1. Fused QK Norm

UClip 不像 RMSNorm 那样需要 Reduction，所以 Q 和 K 其实可以在一个 Kernel 内处理，减少 launch & graph node。结果：

nsys 上 balanced 的 norm 类别从 73.07 ms 到 71.42 ms，只省 1.65 ms，但 GEMM 类别从 146.30 ms 涨到了 164.38 ms。GEMM 不应该跟这次修改的内容有任何关系，那这只可能是 graph capture 形状/timing 改变之后 cuBLAS 选了不同的 algorithm，或者其他因素？完全不理解为什么如此。

6.2. Timed-only graph profiling

之前的 nsys profile 总是把 model warmup、CUDA Graph capture、sampler warmup 和真正测量的 generation loop 全部混在一起。最明显的副作用是 _topk_topp_kernel 在 synthetic greedy 跑里看起来非常大。给 bench_hp1_suite.py 加一个 timed-only 包装，包住真正的 generation loop，再用 nsys profile --capture-range=cudaProfilerApi --capture-range-end=stop。

另外有一个操作细节：vLLM v1 的多进程模式在 parent 进程启动 profiling 之后才 spawn EngineCore child，timed-only capture 必须配合 VLLM_ENABLE_V1_MULTIPROCESSING=0，否则子进程捕获不到。

修完之后：

FA3 的：

_topk_topp_kernel 完全消失了——synthetic greedy 在稳态实际上走的是 ArgMaxOps，之前 profile 里看到的几次大 _topk_topp_kernel 其实是 vLLM 初始化阶段的 warmup artifact！

CUDA API time 透露了一件更值得注意的事：

HP1 在 cudaEventSynchronize 上多花了 446 ms。kernel 时间总和接近，但 wall-time 上 HP1 还是慢——也就是说两边的 kernel 计算量差不多，差距在 graph replay 或者 launch span 上?

6.3. Fused Peri-LN DyHT

AsyncT Decoder 内部有：

attn_out = post_attn_layernorm(attn_out)
hidden_states, residual = pre_mlp_layernorm(attn_out, residual)

pre_mlp_layernorm 自己已经是 fused add+UClipNorm 了，但 post_attn_layernorm 还是独立 op。两次 norm 中间有一个 bf16/fp16 的点（为了保证 residual 上的数值精度问题）。直觉上可以把它合成一个 kernel：先对 attention output 做 post-attn UClip，得到 post_q，再和 residual 相加，得到新的 residual，同时输出 pre-MLP UClip 的 hidden state。但此处不知道数值精度等是否会产生影响，于是：

hp1_post_attn_pre_mlp_uclip       // v1: 复用 residual buffer
hp1_post_attn_pre_mlp_uclip_out   // v2: 单独 residual_out

direct microbench（N x 1024bf16）：

graph A/B：

想象中更好的 v1 在 balanced 上效果居然不好！同一份算术、相同的 fused kernel，唯一区别就是 in-place vs out-of-place 改变了 residual buffer 的 alias 关系。看起来，CUDA Graph 对这种 alias 改变非常敏感。大概是因为 capture 时 buffer dependency 不同，replay 时 graph 内部的 scheduling 完全不同。

6.4. WGMMA FIR register 残余

之前 3.2 中已经把那个 scores_smem 干掉了，但跨 page 的 FIR 历史（两个值，对应 FIR_K=3 的 score[k-2]、score[k-1]）还存在 shared memory。这个 state 在同一个 CTA 的整个生命周期内只被自己用，所以挪到 register 完全安全。但直接修改为一个 thread 保留一堆 tail 不对：m64n8fragment 在 PTX 的 layout 下，每个 thread 同时持有 row0 和 row1 两个逻辑行的元素。两行的 FIR 历史是各自独立的，所以每个 thread 要持有两对 tail register：

float tail0_row0_reg = 0.0f;
float tail1_row0_reg = 0.0f;
float tail0_row1_reg = 0.0f;
float tail1_row1_reg = 0.0f;

B32/T256 micro 从 0.165 ms 到 0.163 ms（很小但稳定），主要的作用是把 tail_smem 区域从 shared layout 里删了，registers/thread 不变。

6.5. Prefix-zero contiguous prefill

做的过程中的猜测：纯 prefill（prefix_len 全部为 0）时，刚算出的 K/V 还在 post-RoPE 的 contiguous tensor 里，立刻又写进 paged cache、然后 kernel 又从 paged cache 通过 block_table 把同样的 token 读回来。这一次回读看起来很浪费。

用

key[q_start + k_abs, kv_head, d]
value[q_start + k_abs, kv_head, d]

代替

block_table[req, page] -> cache[physical_block, block_offset, kv_head, d]

dispatch 条件是 all(prefix_len == 0) && Q/K/V contiguous。kernel 复用整个 full WGMMA 实现，只换 K/V 的指针。

结果 0 提升，并且在真实 vLLM graph 里 contig 路径根本没触发：post-projection 的 K/V tensor 是 fused QKV projection 的 narrow view，is_contiguous()返回 False。要触发就得 .contiguous()额外一次 copy。

看起来 paged cache 的 layout 已经接近 WGMMA kernel 想要的 access pattern，回读本身不是大开销。

6.6. Gate Specialization

看 WGMMA full kernel 的 FIR/gate 内部，里面有这样逻辑的一段：

const float r0 = hp_relu_pre ? fmaxf(s0, 0.0f) : s0;
// ...
p0 = fmaxf(z0, 0.0f);
p0 = fminf(p0, clip_max);
p0 = fmaxf(p0, clip_min);

hp_relu_pre、clip_min、clip_max 这三个东西在生产配置上其实都是固定的：hp_relu_pre=true、clip_min=0、clip_max=1。但是 kernel 里它们是 runtime 才知道的。在最内层 K-tile loop 里，每次 gate 都要走 runtime branch + 两次 fmin/fmax。改成 compile-time template：

template <bool CONTIG_KV, bool HP_RELU_PRE, bool CLIP_0_1>
__global__ void hp1_prefill_wgmma_full_kernel(...);

const float r0 = HP_RELU_PRE ? fmaxf(s0, 0.0f) : s0;
// ...
if constexpr (CLIP_0_1) {
  p0 = fminf(p0, 1.0f);
} else {
  p0 = fminf(p0, clip_max);
  p0 = fmaxf(p0, clip_min);
}

这样 hp_relu_pre=true 和 clamp [0,1]都变成 compile-time 常量，gate 热路径可以少掉 runtime branch 和双边 clamp。

ncu 上 executed instructions 从 51.8M 到 48.7M（少了 3M 条指令），SM throughput 42% → 42%（几乎不动），但 duration 从 175 us 到 158 us。graph 上 balanced 提了约 5%，其它 scenario 变化不大。这是合理的，因为 micro 是纯 kernel 时间，graph 还有大量 GEMM 和 graph overhead 摊薄。

6.7. Shared Store Conflict?

direct WGMMA 改善确实转化成 timed kernel time，比如 long_decode_b32 的 hp1_prefill_wgmma 从 134.992 ms 降到 122.206 ms。但 balanced serving 中 WGMMA 已经只有 11.6%，GEMM 和 norm 也同样大。

做 ncu line-info 走 source counters 看 shared store 冲突，结果是：

• LDGSTS cp.asyncstores：32-way conflict rows，约 1M wavefronts
• V scalar shared stores：4-way conflict rows
• Q scalar shared stores：4-way conflict rows
• P pair stores：很小，看起来不成问题

做两个尝试：

1. V staging 换成 WGMMA-layout-coalesced kk-pair stores。让相邻的 K 位置在 shared memory 里成对，shared store 冲突会显著减少，但是 global load 就不是 coalesced 了——本来 V 的 global load 是一个 thread 读一个 contiguous chunk，改完之后会按 KV head stride 跳着读。结果慢 10%！global load coalescing 比 shared store conflict 影响大得多。
2. 保留 V，只把 Q 改成 bf16x2 pair store。这次 global side 没有副作用，shared store 上少几条指令，时间上一点没变。

这两个负向结果一起说明：ncu shared-store counter 高的地方不一定是有问题的（或者说修了不一定有好处），必须连带看 global memory side 的代价。WGMMA prefill 当前已经接近 instruction-bound + L1-bound 的混合状态，单一维度的 shared layout 改动很难推动整体。

6.8 再改 fused QK norm & profiling

前面的 fused QK Norm 没有收益的问题还是非常疑惑，怀疑是某种 graph-side 的 aliasing 问题，但还没确认。

尝试加 --disable-async-scheduling，看是不是 async scheduling 引入了某种 reorder。结果 HP1 0.72s → 0.91s、softmax 0.48s → 0.72s，两边都变慢。async scheduling 不是 HP1 wall gap 的解释。又尝试直接直接嵌入之前写的 fused QK Norm 作为默认再测一次，又有~ 3% 的提升，确实很奇怪！

做 timed-only nsys 的时候发现 AsyncT 和 softmax visible GPU kernel total 很接近，但 wall time 差很多，怀疑 host scheduler idle？但这里其实有一个 measurement 上更基础的问题——之前用的是 cuda_gpu_kern_sum，它不包括 CUDA graph replay 的 span。也就是说，所谓的 ”idle” 其实是 graph 在 replay 但被排除在 kernel-sum 之外。

加上-full-shape-warmup, --cudagraph-mode, --cudagraph-copy-inputs。其中 --full-shape-warmup 解决一个隐藏的精度问题：默认 warmup 用的是 batch=4, tokens=4，跟测量的 prompt shape 完全不一样，所以测量 loop 里可能还在 JIT Triton fallback kernel（比如 _hp1_decode_stage1_kernel）。开启 full-shape warmup 之后，timed loop 里就只剩 _compute_slot_mapping_kernel 这种小尾巴。

加上之后：

比例从 1.51x 到 1.45x，HP1 自己也确实快了一些（不再有 inflight JIT）。但 ratio 没有质变，证明 JIT warmup 不是主要 gap。

差的 580 ms 跟 wall gap 吻合了。所谓 idle gap 实际上是 HP1 graph replay 自己的 span 更长。AsyncT 有 35 层，softmax 有 28 层，per-layer graph replay 算下来：

HP1 graph replay per layer:     2.326 ms / 35 = 66.5 us/layer
Softmax replay per layer:       1.567 ms / 28 = 56.0 us/layer

每层多 10 us。这里既有 7 层的层数差，也有 AsyncT 每层 graph body 重量的差（多几个 UClipNorm 节点等）。普通 kernel sum 把 graph replay span 漏掉了，才看起来像 “GPU kernel time 差不多但 wall 差很多”。

所以接下来再优化，还需要减少 AsyncT graph nodes、融合 AsyncT-specific 小节点、缩短 per-layer graph replay span。除非 kernel 改动能真正减少 graph replay mean，否则 direct-op 小规模的优化不一定有意义。

6.9 重做 Specialized fused QK norm

更细的 nsys --cuda-graph-trace=node 可以把 kernel 按 graph node 归类，加上之前 fork-tracing 之前的一些问题，最终确定 vLLM 多进程下要用 --trace-fork-before-exec=true --wait=all，并且要让 worker 进入 profile capture，否则 SQLite 里压根没有 kernel/graph-node 表。

修好之后，AsyncT decode_heavy_b32graph-node attribution：

FA3：

GEMM 在两边都是大头但接近，softmax_attention 153 ms ≈ AsyncT decode_attention 252 ms 多了约 100 ms（35 vs 28 层 + 每层贵）；AsyncT 的 norm 95.7 ms 也明显多于 softmax 的 27.4 ms。看来还是得先把周边的算子改进做掉。

Q/K UClipNorm 是 alpha -> per-channel, gamma -> scalar。之前的 hp1_qk_uclip_norm_kernel 把这些都做成 runtime flag，每个 element 都走 branch。新加一个生产模式 specialized kernel 把这两个写死：

graph A/B：

graph-node 验证：

干掉了 35 个 norm 节点，norm category 减少 13.7 ms。

6.10. 分离 KV-cache 更新

FlashAttention 设 forward_includes_kv_cache_update=False，由 vLLM 的统一 update op 写 KV cache。我们的一直在 attention forward 里直接调 reshape_and_cache_flash。所以尝试模仿一下，结果

反而变慢了！很奇怪，但暂时先放着。

6.11. Decode Thread-group size 自适应

到这里 Decode 之外的 Kernel 基本都走了一遍了，有想法的只剩下改 GEMM 里加稀疏了，所以还是回归 Decode Kernel 小改。

decode kernel 用的 THREAD_GROUP_SIZE=8 是默认值——一个 K token 由 8 个 thread 协作做 QK dot product。但这个值好不好其实跟 batch / context length 有关。重新做了一次 TGS=4/8/16 sweep，顺手修了一个 group mask 的问题：

constexpr unsigned int THREAD_GROUP_MASK =
    (1u << THREAD_GROUP_SIZE) - 1u;
const unsigned int group_mask =
    THREAD_GROUP_MASK << (group_in_warp * THREAD_GROUP_SIZE);

结果：

TGS=4 一片下降；TGS=16 是大 batch 专用的——B≥64 时优势明显，否则不如 TGS=8。

加 auto policy：

unset VLLM_HP1_DECODE_TGS:
  num_decode_reqs >= 64 -> TGS=16
  otherwise             -> TGS=8

VLLM_HP1_DECODE_TGS_AUTO_MIN_BATCH=64

最终在几个 benchmark 上有～6% 提升。

7. Decode 细节优化

继续回 Decode 优化。

7.1. Paired K loads + Register Q Fragment

前面 profile 的过程中，ncu 一直提到有 L1/shared instruction pressure。问了一下 ChatGPT 和 Claude，它们都推荐我详细看看 decode Kernel 里面编译出来的 SASS。遂看 source-counter ncu 上 decode QK 的 SASS，最显眼的两个 pattern：

1. K 侧是一长串 LDG.E.U16.CONSTANT，每 bf16/fp16 K 元素是独立的 scalar load；
2. 每个 K token 的 QK loop 里反复 LDS.128 从 q_smem 读取 Q tile。

template <>
__device__ inline void load2_to_float<__nv_bfloat16>(
    const __nv_bfloat16* ptr, float& x0, float& x1) {
  const uint32_t packed = *reinterpret_cast<const uint32_t*>(ptr);
  x0 = __uint_as_float((packed & 0xffffu) << 16);
  x1 = __uint_as_float(packed & 0xffff0000u);
}

然后把 Q 全部在 CTA load 后 hoist 到 register q_frag[]，内层 K loop 不再反复读 shared Q：

float q_frag[ELEMS_PER_THREAD];

#pragma unroll
for (int i = 0; i < ELEMS_PER_THREAD; i++) {
  const int d = thread_in_group * ELEMS_PER_THREAD + i;
  q_frag[i] = q_smem[d];
}

for (int i = 0; i < ELEMS_PER_THREAD; i += 2) {
  const int d = thread_in_group * ELEMS_PER_THREAD + i;
  float k0, k1;
  load2_to_float(k_ptr + d, k0, k1);
  partial_qk += q_frag[i] * k0;
  partial_qk += q_frag[i + 1] * k1;
}

direct-op 上

ncu 具体 profile：

L1/TEX 直接从 39% 降到 22%！本来 ”L1 instruction pressure” 是真问题，配对 + 寄存器化把它解掉了一大半。效果显著！

再看之前不太好的 decode_heavy_b32：

hp1_decode_attention 占比从之前的约 42% 降到 29.5%，证明 paired K + Q register 确实有效！

7.2. 推广到 split-K

split-K decode 用的是独立的 stage1 kernel。它的内层 QK loop 还是老样子：从 shared memory 读 Q、标量 bf16 load K。把 R64 的 load2_to_float+ q_frag 模式照搬过来：

float q_frag[ELEMS_PER_THREAD];
#pragma unroll
for (int i = 0; i < ELEMS_PER_THREAD; i++) {
  const int d = thread_in_group * ELEMS_PER_THREAD + i;
  q_frag[i] = q_smem[d];
}
...
for (int i = 0; i < ELEMS_PER_THREAD; i += 2) {
  const int d = thread_in_group * ELEMS_PER_THREAD + i;
  float k0, k1;
  load2_to_float(k_ptr + d, k0, k1);
  partial_qk += q_frag[i] * k0;
  partial_qk += q_frag[i + 1] * k1;
}

graph 上 long_decode_b32 从 1.452 s 降到 1.008 s，1.44x。AsyncT/FA3 在 long decode 上的 gap 从 ~2x 降到 1.37x。

做了这个修改之后 Split-K 的效果比之前好很多了，但之前的 Split-K 触发条件还是保守的 split_min_len=1024，这轮重新 sweep 一下：

看起来 split-256 没什么问题。再看看什么时候要 split：

观察上面这张表格不难总结出：split-K 是中 - 小 batch 长 context 的工具，B128 这种大 batch 因为 base kernel 已经能填满 SM，split 的额外 reduce/write 开销反而拖累。

改进后在 mixed 场景上得到～10% 提升。

7.3. 一系列碰壁

做完之后做了一系列比较小的优化尝试，但效果都不好：

1. 个 K/V token 都要通过 block_table[block_idx]找 physical block，是否可以把 block_table 先 preload 到 shared memory？

__shared__ int32_t block_nums[MAX_BLOCKS];

for (int b = tid; b < num_seq_blocks; b += NUM_THREADS) {
  block_nums[b] = block_table[b];
}

physical_block = block_nums[block_idx];

结果性能完全不行，猜测是 block_table 很小、很可能已经 cache；额外 preload loop、shared array、sync 反而增加压力。

1. 会不会 if (p == 0)分支比加载 V 再乘 0 更贵？临时删掉，micro bench 上效果不错，但是真实模型里面真的是很稀疏的，Attention 稀疏也有约 70%！ 做完在真实模型上测试效果明显变差。

2. 把 Peri-LN UClip 像 QK Norm 一样给一个编译期写死的参数形状，短序列上轻微掉速长序列上～3% 收益，效果不好但保留（因为起码不算很负面、理论上有帮助）；

3. 把 PV 改成 outer loop over 16-token KV page，hoist block_idx = t / 16、block_offset = t % 16、physical_block = block_table[block_idx]出 inner loop，这样不仅在数学上等效而且理论上 inner loop 里面没有取模运算了。结果 micro bench 上掉速 50%！GPT 的评价：“t / 16、t % 16 这种 BLOCK_SIZE=16 的整数除模在编译器视角下基本免费（移位 + and），同时新的 per-page 结构引入了内层的 bounds check 和 outer loop overhead。整数除模看着重，但被强制 hoist 出来不一定净赚。”

4. 顺手看了下 FFN 里面的稀疏度，～71% 导致基本做不了 block level（～2.81%）的稀疏。不过按 Sparsing Law 里面提到的，应该会随着训练稀疏度逐渐增大？等之后有 Main Result 的结果再来做。

5. 再看了一下 SASS，最重的几条还是 PV loop 里

   • shared p_buf 的标量读；
   • predicated V global load；
   • bf16→fp32 转换 + 整数位操作；
   • paged V 的地址生成。

   于是尝试改成让一个 thread 累积两个输出维度：

if (tid < HEAD_SIZE / 2) {
  const int d = tid * 2;
  float acc0 = 0.0f, acc1 = 0.0f;
  ...
  float v0, v1;
  load2_to_float(v_ptr, v0, v1);
  acc0 += p * v0;
  acc1 += p * v1;
  ...
  out_ptr[0] = from_float<scalar_t>(acc0);
  out_ptr[1] = from_float<scalar_t>(acc1);
}

减少了 V load/conversion 指令数，但 active lanes 从 128 减到 64，latency hiding 更差。

7.4. Sparse P 压缩

前面最早的 PV zero-skip 是最简单的 sparse 利用：

for t in seq_len:
    if p[t] == 0: continue
    acc += p[t] * V[t]

问题是每个 output dimension thread 都还要遍历完整 seq_len，每次读 p_buf[t]，判断是否为 0。这里的想法是：既然 zero 这么多，就先在 gate pass 中构造 nonzero index list，然后 PV 只遍历 nonzero positions。

__shared__ float logits[MAX_SEQ_LEN + FIR_K - 1];
__shared__ float p_buf[MAX_SEQ_LEN];
__shared__ int nz_count;
__shared__ int nz_idx[MAX_SEQ_LEN];
...

// gate 阶段
if (tid == 0) nz_count = 0;
__syncthreads();
...
const float p_out = gamma_v * p_val;
p_buf[t] = p_out;
if (p_out != 0.0f) {
  const int slot = atomicAdd(&nz_count, 1);
  nz_idx[slot] = t;
}

// PV 阶段：只遍历非零位置
const int local_nz = nz_count;
for (int j = 0; j < local_nz; j++) {
  const int t = nz_idx[j];
  const float p = p_buf[t];
  // ... load V, acc += p * V ...
}

不过这里其实引入了一个小的精度问题：atomicAdd(&nz_count, 1)产生的 nz_idx 顺序并不严格升序——多个 thread 在同一时刻竞争 slot。但这只改变了 PV 中非零项的累加顺序，不改变 V 的内容。decode 的 sentinel 仍然以 PyTorch reference 为基准；可观察的最大误差是 1.5259e-05（bf16），问题不大。这样改完算子里就有～30% 的提升！

同样的，split-K 内部的也可以做同样的修改：

__shared__ int nz_count;
__shared__ int nz_idx[SPLIT_SIZE];

p_buf[rel] = p_out;
if (p_out != 0.0f) {
  int slot = atomicAdd(&nz_count, 1);
  nz_idx[slot] = rel;
}

for (int j = 0; j < nz_count; j++) {
  int rel = nz_idx[j];
  acc += p_buf[rel] * V[rel];
}

在 split-K 上的提升更大，部分 shape 上有～80% 的 op 提速。并且重扫参数后，split-K 算法在更小的 seq_len 上也有优势了！

现在的情况：

至此第一次与 FA3 打平！