续着前文继续做优化。

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。

做完之后gsm8k 123s wall, 99s gen，比softmax的FA2 152s wal还快¹。

进一步地，上面的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]

于是可以用 __shfl_up_sync在 quad 内拿前驱，再加上每行两个 tail register，直接从 QK fragment 生成 P tile，不再需要 scores_smem：²

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%

• GEMM 48.5%
• HP1 WGMMA 12.6%
• norm 11.1%
• sampler ArgMax 5.3% ³

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。

不知道为什么K这边不是一条大的vectorized的指令。要解决也很简单了。对于K，把两个bf16配成一个FP32⁴：

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打平！