Skip to content

Qwen3.5-397B-A17B ​

源码核对版本:SGLang commit 95140a7b0c9fc2f87a2a6cf6f6f0df8640a73174。

提交日期:2026-09-14 17:11:51(UTC+8)。

1. Prefill-GDN ​

图中展示 Gated DeltaNet 的 Chunk 计算路径,包含卷积状态、SSM 状态更新,以及后续共享专家与路由专家计算。蓝色节点标记 TP 通信。

流程图

正在绘制图表…

flowchart TB
      X["hidden_states"] --> K1["1<br/>input_add_rmsnorm"]
      R["residual"] --> K1

      K1 --> K2["2<br/>in_proj_qkvz_input_fp8_quant"]
      K2 --> K3["3<br/>in_proj_qkvz_fp8_gemm"]
      K1 --> K4["4<br/>in_proj_ba_bf16_gemm"]

      P["extend_prefix_lens"] --> K5["5<br/>initial_state_mask"]
      K5 --> K6["6<br/>causal_conv1d_silu"]
      CS["conv_state"] --> K6
      K3 -->|qkv| K6
      K6 --> CSO["updated_conv_state"]
      K6 --> K7["7<br/>qkv_split_layout"]

      K4 -->|a, b| K8["8<br/>gdn_gating"]
      K8 -->|g| K11["11<br/>chunk_gate_cumsum"]

      K7 -->|q| K9["9<br/>q_l2norm"]
      K7 -->|k| K10["10<br/>k_l2norm"]

      K12["12<br/>chunk_A_zero_init"] --> K13["13<br/>chunk_kkt_solve"]
      K10 -->|k| K13
      K8 -->|beta| K13
      K11 -->|g_cumsum| K13

      K13 -->|A| K14["14<br/>chunk_recompute_w_u"]
      K10 -->|k| K14
      K7 -->|v| K14
      K8 -->|beta| K14
      K11 -->|g_cumsum| K14

      K14 -->|w, u| K15["15<br/>chunk_state_h"]
      K10 -->|k| K15
      K11 -->|g_cumsum| K15
      SS["ssm_state"] --> K15
      K15 --> SSO["updated_ssm_state"]

      K16["16<br/>chunk_output_zero_init"] --> K17["17<br/>chunk_output_o"]
      K9 -->|q| K17
      K10 -->|k| K17
      K11 -->|g_cumsum| K17
      K17 --> K18["18<br/>output_rmsnorm_silu_gate"]
      K3 -->|z| K18
      K18 --> K19["19<br/>out_proj_input_fp8_quant"]
      K19 --> K20["20<br/>out_proj_fp8_gemm"]
      K20 --> K21["21<br/>attention_tp_allreduce"]

      K21 --> K22["22<br/>post_attention_add_rmsnorm"]
      K1 -->|residual| K22

      K22 --> K23["23<br/>shared_gate_up_input_fp8_quant"]
      K23 --> K24["24<br/>shared_gate_up_fp8_gemm"]
      K24 --> K25["25<br/>shared_silu_mul"]
      K25 --> K26["26<br/>shared_down_input_fp8_quant"]
      K26 --> K27["27<br/>shared_down_fp8_gemm"]

      K22 --> K28["28<br/>router_bf16_gemm"]
      K28 --> K30["30<br/>routing_histogram_scores"]
      K30 --> K31["31<br/>routing_indices"]

      K22 --> K29["29<br/>routed_input_nvfp4_quant"]
      K29 --> K32["32<br/>routed_gate_up_swiglu_fp4_gemm"]
      K31 -->|routing_metadata| K32
      K32 --> K33["33<br/>routed_down_fp4_gemm"]
      K31 -->|routing_metadata| K33

      K33 --> K34["34<br/>routed_finalize"]
      K30 -->|routing_weights| K34
      K31 -->|token_mapping| K34

      K27 -->|shared_output| K35["35<br/>shared_gate_sigmoid_mul_add"]
      K34 -->|routed_output| K35
      K22 -->|hidden_states| K35

      K35 --> K36["36<br/>moe_tp_allreduce"]
      K36 --> N1["next_layer:<br/>1. input_add_rmsnorm"]
      K22 -->|residual| N1

      classDef green fill:#c6efce,stroke:#548235,color:#000;
      classDef pink fill:#edc2ed,stroke:#a64ca6,color:#000;
      classDef orange fill:#f8cbad,stroke:#c65911,color:#000;
      classDef purple fill:#d886df,stroke:#9237a0,color:#000;
      classDef red fill:#ff9999,stroke:#c00000,color:#000;
      classDef blue fill:#bdd7ee,stroke:#4472c4,color:#000;
      classDef gray fill:#d9d9d9,stroke:#808080,color:#000;

      class K21,K36 blue;

2. Prefll-FullAttn ​

箭头表示数据或资源依赖,编号用于定位节点。图中的分叉表示分支之间没有直接数据依赖;实际能否并行,还取决于 CUDA stream、同步和 GPU 资源。蓝色节点标记 TP 通信。

流程图

正在绘制图表…

flowchart TB
      K01["1<br/>input_add_rmsnorm"] --> K02["2<br/>qkv_proj_input_fp8_quant"]
      K02 --> K03["3<br/>qkv_proj_fp8_gemm"]
      K03 --> K04["4<br/>qk_rmsnorm_rope_gate_split"]
      K03 -->|V| K05["5<br/>store_kv_cache"]
      K04 -->|K| K05
      K04 -->|Q| K07["7<br/>trtllm_full_attention"]
      K05 -->|KV cache| K07
      K06["6<br/>attention_workspace_init"] --> K07
      K07 --> K08["8<br/>output_sigmoid_gate"]
      K04 -->|gate| K08
      K08 --> K09["9<br/>o_proj_input_fp8_quant"]
      K09 --> K10["10<br/>o_proj_fp8_gemm"]
      K10 --> K11["11<br/>attention_tp_allreduce"]
      K11 --> K12["12<br/>post_attention_add_rmsnorm"]
      K01 -->|residual| K12

      K12 --> K13["13<br/>shared_gate_up_input_fp8_quant"]
      K13 --> K14["14<br/>shared_gate_up_fp8_gemm"]
      K14 --> K15["15<br/>shared_silu_mul"]
      K15 --> K16["16<br/>shared_down_input_fp8_quant"]
      K16 --> K17["17<br/>shared_down_fp8_gemm"]

      K12 --> K18["18<br/>router_bf16_gemm"]
      K18 --> K20["20<br/>routing_histogram_scores"]
      K20 --> K21["21<br/>routing_indices"]
      K12 --> K19["19<br/>routed_input_nvfp4_quant"]

      K19 --> K22["22<br/>routed_gate_up_swiglu_fp4_gemm"]
      K21 --> K22
      K22 --> K23["23<br/>routed_down_fp4_gemm"]
      K21 -->|routing metadata| K23
      K23 --> K24["24<br/>routed_finalize"]
      K20 -->|routing weights| K24
      K21 -->|token mapping| K24

      K17 --> K25["25<br/>shared_gate_sigmoid_mul_add"]
      K24 --> K25
      K12 -->|hidden states| K25
      K25 --> K26["26<br/>moe_tp_allreduce"]
      K26 --> NEXT["next:<br/>input_add_rmsnorm"]
      K12 -->|residual| NEXT

      classDef blue fill:#cfe2f3,stroke:#4472c4,color:#000
      classDef purple fill:#e8c1e8,stroke:#9955aa,color:#000
      classDef orange fill:#f9dccb,stroke:#cc8855,color:#000
      classDef green fill:#d5edc7,stroke:#70a050,color:#000

      class K11,K26 blue

3. Decode-GDN ​

流程图

正在绘制图表…

flowchart TB
      X["previous_layer.shared_gate_sigmoid_mul_add"] --> K1["1<br/>input_allreduce_add_rmsnorm"]
      R["residual"] --> K1

      K1 -->|bf16| K2["2<br/>in_proj_ba_bf16_splitk_gemm"]
      K2 -->|partial_sums| K4["4<br/>in_proj_ba_splitk_reduce"]

      K1 --> K3["3<br/>in_proj_qkvz_input_fp8_quant"]
      K3 --> K5["5<br/>in_proj_qkvz_fp8_gemm"]

      K5 -->|qkvz| K6["6<br/>qkvzba_unpack_causal_conv1d_silu"]
      K4 -->|ba| K6
      CS["conv_state"] --> K6
      K6 --> CSO["updated_conv_state"]

      K6 -->|qkv, a, b| K7["7<br/>gdn_packed_recurrent_qk_norm_gating_state_update"]
      SS["ssm_state"] --> K7
      K7 --> SSO["updated_ssm_state"]

      K7 -->|attention_output| K8["8<br/>output_rmsnorm_silu_gate"]
      K6 -->|z| K8
      K8 --> K9["9<br/>out_proj_input_fp8_quant"]
      K9 --> K10["10<br/>out_proj_fp8_gemm"]

      K10 --> K11["11<br/>attention_allreduce_add_rmsnorm"]
      K1 -->|residual| K11

      K11 -->|bf16| K12["12<br/>router_bf16_gemm"]
      K12 --> K16["16<br/>routing_block_scores"]
      K16 --> K18["18<br/>routing_cluster_indices"]

      K11 --> K13["13<br/>shared_gate_up_input_fp8_quant"]
      K13 --> K14["14<br/>shared_gate_up_fp8_gemm"]
      K14 --> K17["17<br/>shared_silu_mul"]
      K17 --> K19["19<br/>shared_down_input_fp8_quant"]
      K19 --> K20["20<br/>shared_down_fp8_gemm"]

      K11 --> K15["15<br/>routed_input_nvfp4_quant"]
      K15 --> K21["21<br/>routed_gate_up_swiglu_fp4_gemm"]
      K18 -->|routing_metadata| K21

      K21 --> K22["22<br/>routed_down_fp4_gemm"]
      K18 -->|routing_metadata| K22

      K22 --> K23["23<br/>routed_finalize"]
      K18 -->|routing_weights, token_mapping| K23

      K23 -->|routed_output| K24["24<br/>shared_gate_sigmoid_mul_add"]
      K20 -->|shared_output| K24
      K11 -->|hidden_states| K24

      K24 --> N1["next_layer:<br/>1. input_allreduce_add_rmsnorm"]
      K11 -->|residual| N1

      classDef blue fill:#bdd7ee,stroke:#4472c4,color:#000;
      class K1,K11,N1 blue;

3.1 阅读范围与符号 ​

本节对应 decode_bs54.nsys-rep 中 GPU 0、decode step 1、layer 4 的 24 次 kernel launch,输入上下文长度为 8192,batch size 为 54,TP=2。编号来自 decode_bs54_step1_gpu0_layer4_GDN.csv,按 kernel 开始时间排序;它不是严格的串行执行顺序。这里解释的是原版 SGLang 路径,不能把后续实验融合的行为代入这份 trace。

一个 decode step 中,每个活跃请求处理一个 token。图中的主干可以先理解为:

合并上一层结果并归一化 → GDN 投影、状态更新、输出 → 合并 attention 结果并归一化 → shared/routed 两条 MoE 分支 → 合并两条分支 → 下一层负责跨卡求和。

下面先用实数公式解释计算目标,再标注关键的量化和舍入边界。公式不等于逐 bit 的 CUDA 仿真:归约顺序、BF16 舍入、FP8/FP4 编码和近似数学指令都会影响实际数值。

符号含义
nbatch 中的一条请求;一次 decode 的一行数据
t该请求当前处理的 token 位置,不是 kernel 编号
pTP rank,本场景为 0 或 1
D模型 hidden size,本场景为 4096
H节点 1 输出的 normalized hidden states,形状为 [54,D]
M节点 11 输出的 normalized hidden states,供 MoE 使用
R1,R2节点 1、11 分别输出并保留的 residual
dk,dv一个 GDN head 的 key、value 维度
St一个请求、一个 value head 的历史状态,本文采用 [dv,dk] 布局
Cmathrmbf16(x)将结果舍入为 BF16
odot对应元素相乘,不是矩阵乘法
sigma(x)sigmoid,1/(1+exp⁡(−x))
SiLU(x)xσ(x),输出不局限于 0~1

外层 input/post-attention norm 使用 GemmaRMSNorm。用 checkpoint 中的权重 w 表示时:

GN(x;w)=x1D∑j=1Dxj2+ϵ⊙(1+w).

运行时可能预先构造有效权重 γ=1+w,所以底层 kernel 接收到的 gamma 不一定还是 checkpoint 原始的 w。节点 8 的 head-wise RMSNorm 使用自己的直接乘法权重,不能照搬这里的 1+w。

3.2 逐个 kernel 的目标与公式 ​

1. input_allreduce_add_rmsnorm ​

目标:把上一层 MoE 的两卡局部结果合起来,加回 residual,再生成本层 GDN 的归一化输入。

设上一层节点 24 在 rank p 上输出 Uprev(p),传入的 residual 为 Rin:

Aprev=∑p=01Uprev(p),R~1=Rin+Aprev,R1=Cbf16(R~1),H=Cbf16(GN(R~1;win)).

这是数学上的拆解;实际 fused AR 的归约结果也有数据类型转换,不能据此把整个公式当成无限精度的一次运算。尤其不能认为 norm 必然重新读取已舍入的 R1:融合实现可以直接使用寄存器内的 residual 加法结果。

输出有两份,用途不同: H 交给 2、3;R1 绕过 GDN 主干,交给 11。AllReduce 是求和,不是除以卡数的平均值。上一层 MoE 的通信已经包含在这里,不需要额外虚构一个独立的 moe_tp_allreduce kernel。

2. in_proj_ba_bf16_splitk_gemm ​

目标:计算 BA 投影的局部和,生成控制 GDN 状态更新的参数。

完整投影应为:

BA=H(Wba(p))T=[b,a].

这次 cuBLAS 选择了 Split-K 算法,把求和维度 D 分为若干集合 Ki:

Pi,n,j=∑d∈KiHn,dWba,j,d(p).

节点 2 只产生各段 Pi,还不是完整的 a,b。输入和权重是 BF16,这条 trace 的 Split-K workspace/归约使用 FP32 局部结果。具体切成几段由 GEMM 算法决定,不能从“TP=2”推断为两段。

这里 a,b 都是 每个 token、每个 value head 的标量控制输入。节点 2 还没有做 sigmoid 或 softplus,也没有更新历史状态。

3. in_proj_qkvz_input_fp8_quant ​

目标:把同一个 H 转成 QKVZ GEMM 需要的 FP8 输入。

对一个非零 token 向量 x=Hn,per-token E4M3 量化的主要步骤为:

mn=maxj|xj|,sn=mn448,qn,j=castE4M3(clip(xjsn,−448,448)),xj≈snqn,j.

输出是 FP8 数值和 FP32 scale,二者共同表示原数据;scale 不是可丢弃的辅助信息。零行、特殊值的行为以具体 kernel 的边界处理为准,不能直接在数学公式中除以零。

下文用 Q8(X)=(QX,sX) 表示这类 per-token 量化,用 D8(QX,sX) 表示按 scale 解释后的近似实数值。节点 3 不依赖节点 2 的输出;2 和 3 是读取同一份 H 的两条分支。

4. in_proj_ba_splitk_reduce ​

目标:把节点 2 的 Split-K 局部和合并,得到最终的 a,b。

BA=Cbf16(∑iPi),[b,a]=split(BA).

“split”在这里先表示逻辑上识别两部分,实际整理布局由节点 6 完成。节点 4 的 reduce 是 单张 GPU 内、同一个 GEMM 的局部和归约,不是两张 GPU 的 AllReduce。

因此 2 + 4 才等价于一次完整 BA 投影。nvjet 是节点 2 的 GEMM kernel 名称,不是 shared expert,也不是通信算子。

5. in_proj_qkvz_fp8_gemm ​

目标:从 hidden states 生成 GDN 的内容向量和输出门控向量。

令 H^=D8(Q8(H)),W^qkvz(p) 为 FP8 权重连同其 scale 所表示的有效权重:

C=Cbf16(H^(W^qkvz(p))T)=[qraw,kraw,vraw,z].

这是带输入/权重 scale 的低精度 GEMM,不是直接把 FP8 编码当普通整数相乘。输出为 BF16。这里的拼接表示逻辑分量,实际存储顺序需要节点 6 按模型布局解包。

q,k,v 是后续记忆读写的内容;z 是 从记忆读取结果之后 使用的输出 gate。z 不决定状态遗忘,不参与节点 7 的状态更新。

6. qkvzba_unpack_causal_conv1d_silu ​

目标:解包两路投影、对 QKV 做短因果卷积和 SiLU,并更新卷积缓存。

设卷积窗口长度为 Lc,xt,c 是投影产生的某个 Q/K/V 通道。把卷积系数按滞后量编号,则:

x~t,c=SiLU(∑ℓ=0Lc−1wc,ℓxt−ℓ,c+bc).

没有卷积 bias 的配置取 bc=0。这是逐通道的短卷积:使用当前投影和缓存中的历史投影,不是全序列 attention,也不是所有通道之间的稠密 GEMM。

节点同时执行布局整理:

(C,BA,conv_statet−1)⟶(qt,kt,vt,zt,at,bt,conv_statet).

其中 q,k,v 经过卷积和 SiLU;z,a,b 只是按需要拆出和排布,不经过这次卷积。卷积 state 保存后续 token 所需的短窗口,与节点 7 的矩阵记忆 S 是两种不同缓存。节点 6 同时依赖 4 和 5,不能只等 QKVZ 投影结束就开始消费所有输入。

7. gdn_packed_recurrent_qk_norm_gating_state_update ​

目标:一个 kernel 内完成 Q/K 归一化、a/b 门控变换、旧状态衰减、delta 更新和记忆读取。

以下公式针对一个 token、一个 value head。多个 value head 可以按模型分组规则使用对应的 Q/K head。

第一步,Q/K L2 norm,并给 query 乘缩放系数:

k¯t=kt∑jkt,j2+10−6,q¯t=1dkqt∑jqt,j2+10−6.

这里是平方和,不是除以维数后的均值,因此它不是 RMSNorm。1/dk 是另一项 query scale,不能漏掉或重复乘。

第二步,把 BA 投影转换为实际的控制系数:

gt=−exp⁡(Alog)softplus(at+dtbias),αt=exp⁡(gt),βt=float32(Cbf16(σ(bt))).

Alog 和 dtbias 是学习得到的每-head 参数。softplus 在大输入区域使用线性分支。上式特意保留 β 的 BF16 舍入:当前 packed decode 源码先转回 b.dtype,再提升到 FP32 参与运算;把它随意改成始终 FP32,不保证与原版精度一致。

第三步,先衰减旧状态,再计算需要写入的“修正量”:

St−=αtSt−1,v^t=St−k¯t,Δvt=βt(vt−v^t).

v^t 是旧记忆在当前 key 上预测出来的 value。GDN 写入的是新 value 与这个预测的差,而不是无条件把 vt 再累加一遍。

第四步,外积更新状态,再用 query 读取:

St=St−+Δvtk¯tT,ot=Stq¯t.

合并更新式得到:

St=αtSt−1(I−βtk¯tk¯tT)+βtvtk¯tT.

这就是 Gated Delta Rule。采用本文 S∈Rdv×dk 的布局,读取必须写成 Stq¯t;如果采用 ST 的存储约定,需要把整套公式一起转置,不能只在读取公式里添加转置。

kernel 将 state 加载到 FP32 后运算,输出 ot 转回输出 dtype,并将更新后的 St 转回 state cache dtype 写回缓存。缓存类型与输出类型应分别看待,不能因为 hidden states 是 BF16,就断言 SSM state 也一定是 BF16。

它把 prefill 中分散的 gating、Q/K norm 和 chunk 状态计算,换成了适合单步 decode 的 recurrent 更新;这里不存在额外的 chunk_gate_cumsum 或 chunk_kkt_solve launch。

8. output_rmsnorm_silu_gate ​

目标:对记忆读出的结果做 head-wise RMSNorm,再用 z 控制输出。

yt,h=Cbf16[(ot,h1dv∑jot,h,j2+ϵ⊙γgdn)⊙SiLU(zt,h)].

归一化范围是 一个 value head 的 dv 个通道,不是模型全部 4096 个通道。代码设置 norm_before_gate=True:先 norm,后乘 gate;不能交换成先乘 gate 再 norm。

γgdn 是这里直接相乘的 norm 权重,与外层 GemmaRMSNorm 的 checkpoint 参数化不同。z 每个输出通道都有分量;它不在 GDN 的状态更新公式里,而是在 ot 之后出现。SiLU 可以输出负数和大于 1 的值,不能把它理解成严格的 0~1 开关。

9. out_proj_input_fp8_quant ​

目标:量化节点 8 的结果,供 attention 输出投影使用。

Y(p)=concath∈local headsyh,(QY,sY)=Q8(Y(p)).

仍使用节点 3 的 per-token FP8 公式,但每行的归约范围变成该 rank 拼接后的 GDN 输出宽度,不能把节点 3 的 scale 复用到这里。它依赖 8,因此 8 + 9 是直接相邻的融合候选。

10. out_proj_fp8_gemm ​

目标:把本卡 GDN heads 的输出投影回模型 hidden space。

A(p)=Cbf16[D8(QY,sY)(W^o(p))T].

这是 row-parallel 输出投影:每张卡只计算输入维度的一个分片,但都产生形状为 [54,D] 的局部结果。两卡输出的形状相同,数值却是需要相加的局部贡献,还不是完整 attention 输出。

11. attention_allreduce_add_rmsnorm ​

目标:合并两卡 attention 贡献,加回节点 1 保留的 residual,再生成 MoE 输入。

A=∑p=01A(p),R~2=R1+A,R2=Cbf16(R~2),M=Cbf16(GN(R~2;wpost)).

与节点 1 一样,公式表达算法结构,实际 AR 和 norm 保留各自的舍入边界。M 分别供 router、shared expert、routed 输入量化和最终 shared gate 使用;R2 则一直保留到下一层的节点 1。后面的专家分支不应提前把 R2 加进来。

12. router_bf16_gemm ​

目标:为每个 token 计算它对各个 routed expert 的打分。

L=Cbf16(MWrouterT),L∈R54×E.

本模型有 E=512 个 routed experts。Ln,e 是第 n 个 token 对专家 e 的 logit,尚不是概率。这里的 GEMM 使用 BF16 输入,因此不能仅向它提供节点 13 的 FP8 数据。

它不是共享专家的 gate;shared gate 在节点 24 中计算。每张 TP 卡都有相应的 router 计算,使后续本卡专家分片使用一致的路由决策。

13. shared_gate_up_input_fp8_quant ​

目标:把 M 转成共享专家第一层 GEMM 所需的 FP8。

(QM(8),sM(8))=Q8(M).

其 scale 按每个 token 的 D 维向量计算。它与 router 12、NVFP4 量化 15 都读取 M,彼此没有输出到输入的数据依赖。不同 stream 的排队次序可能让它们在 trace 上先后出现,但不改变这条分叉关系。

14. shared_gate_up_fp8_gemm ​

目标:一次 GEMM 同时计算共享专家的 gate projection 和 up projection。

[Gs(p),Us(p)]=Cbf16[D8(QM(8),sM(8))(W^s,g(p)W^s,u(p))T].

“gate_up”表示将两个投影打包到一次矩阵乘法中,这一步还没有做 SiLU 或 gate 与 up 的乘法。这里的 gate 是 MLP 中间特征,与 GDN 的 a,b,z 以及节点 24 的 shared-expert gate 都不是同一个量。

15. routed_input_nvfp4_quant ​

目标:为 routed experts 生成 NVFP4 输入,包含 FP4 数据和块级 scale。

这份 trace 使用 block size 16 的 NVFP4 量化。用 G>0 表示传入量化器的全局编码缩放系数,B 表示一个 token 中连续的 16 个元素:

an,B=maxj∈B|Mn,j|,sn,B=castE4M3(Gan,B6),qn,j=castE2M1(Mn,jGsn,B),Mn,j≈qn,jsn,BG.

上式针对非零、正常 scale 的块表达主要计算关系;实现还包含零值处理、饱和转换、scale 的 FP8 舍入和可能的近似倒数。这里把 G 定义为编码乘数,避免与某些 API 中表示反量化因子的 global_scale 同名但互为倒数的约定混淆。

E2M1 的最大有限幅值是 6,两个 4-bit 元素打包在一个字节中。与 FP8 每行一个 scale 不同,NVFP4 每 16 个元素有块级 scale,再结合全局缩放因子解释数据。

13 和 15 是对同一个 BF16 输入的两种独立量化,不是 BF16 → FP8 → FP4 的连续量化。 一份供共享专家 GEMM,另一份供路由专家 GEMM。下文以 Q4、D4 简记这套量化和反量化解释。

16. routing_block_scores ​

目标:选出 token 的 top-K 专家,并计算这些专家的混合权重。

令 Kact 为配置中的 num_experts_per_tok:

En=TopKIndices(Ln,Kact),πn,e=exp⁡(Ln,e−mn)∑j∈Enexp⁡(Ln,j−mn),e∈En,mn=maxj∈EnLn,j.

该 trace 的策略名为 TopKExpertSelect<NoOpPreprocess, SoftmaxPostprocess>:对选中的 logits 做 softmax,得到归一化的 top-K 权重。它在实数上等价于对所有 logits 做 softmax、选 top-K,再对选中概率重新归一化;实际舍入顺序可能不同。

输出是打包的 (expert_id, weight)。kernel 名称模板中的 16 是该实例支持的 top-K 容量参数,不能仅凭它认定实际每个 token 激活 16 个专家;实际值取运行配置。

17. shared_silu_mul ​

目标:完成共享专家的 SwiGLU 激活。

Fs(p)=Cbf16[SiLU(Gs(p))⊙Us(p)].

节点 14 已经完成两个矩阵投影,这里只做激活和逐元素乘法。输出 BF16 中间特征交给节点 19。这个乘法不涉及 token 之间的交互,也不涉及跨卡通信。

18. routing_cluster_indices ​

目标:把“每个 token 选了哪些专家”变成 grouped GEMM 可使用的专家分组索引。

设 en,r 为 token n 的第 r 个选中专家,先把 token-expert 对展开:

u=nKact+r,Ne=∑n,r1[en,r=e].

再根据专家分组、padding 和调度 tile 布局,生成映射:

P(u)=expert_offset(en,r)+position_within_expert(n,r).

输出包括专家计数、分组后的 token 索引、expanded/permuted 索引映射、调度所需信息以及后续使用的 top-K 权重。真实 offset 还受 padding/布局规则影响,上式表达逻辑关系,不指定某一种物理排列。

在这条两阶段 routing 路径里,16 已算出 top-K;18 读取其 packed 结果做分组与映射,不是再做一次 router GEMM。cluster 指 GPU thread-block cluster,不是跨机器集群。建立索引也不意味着此 kernel 已把全部 hidden states 复制到新的排列中,后续 GEMM 可依据索引读取输入。

19. shared_down_input_fp8_quant ​

目标:把共享专家的激活结果量化,供 down projection 使用。

(QFs,sFs)=Q8(Fs(p)).

它对本 rank 的 shared intermediate 向量重新求最大绝对值并计算 scale。这个 scale 与节点 13 的输入 scale 无关。17 + 19 可以围绕同一个中间向量考虑融合,但若要求保持原版结果,量化须尊重节点 17 的 BF16 舍入边界。

20. shared_down_fp8_gemm ​

目标:把共享专家中间特征映射回模型 hidden space。

Ys(p)=Cbf16[D8(QFs,sFs)(W^s,d(p))T].

输出是 [54,D] 的本卡局部贡献,交给 24。这里不单独 AllReduce,因为可以先把本卡 shared 和 routed 输出加在一起,再统一跨卡求和。

21. routed_gate_up_swiglu_fp4_gemm ​

目标:按路由结果计算各个选中专家的 gate/up 投影,并在融合路径中完成 SwiGLU 和供下一层使用的低精度输出。

对每个选中专家 e∈En,用 M^=D4(Q4(M)) 表示节点 15 的有效输入,用 W^e 表示包含量化 scale 的专家权重:

Gn,e(p)=M^n(W^e,g(p))T,Un,e(p)=M^n(W^e,u(p))T,Fn,e(p)=SiLU(Gn,e(p))⊙Un,e(p),(QFe,sFe)=Q4(Fn,e(p)).

这描述标准 SwiGLU 的数学目标;低精度 GEMM 的 accumulator、epilogue 和量化在同一个实现路径中衔接,不能据此假定中间一定落地成 BF16 张量。实际选中实现的 clamp 等 epilogue 细节也要按其参数保留,不能仅靠简化公式重写一个“精度等价”的 kernel。

这个节点同时需要 15 的量化输入和 18 的路由元数据。kernel 名称中的 swiGlu 说明它已经包含激活路径,因此这里没有像 shared 分支那样独立的 shared_silu_mul 和下一次输入量化 launch。

22. routed_down_fp4_gemm ​

目标:对每个 token-expert 对执行第二层专家投影。

Zn,e(p)=Cbf16[D4(QFe,sFe)(W^e,d(p))T].

这里仍是 grouped GEMM,输出按专家计算布局存储。一个 token 选中多个专家,就对应多个专家输出;节点 22 尚未把它们汇总成单个 token 的结果,也尚未完成 TP 跨卡求和。

23. routed_finalize ​

目标:把专家输出映射回 token 顺序,按路由权重加权求和。

Yr,n(p)=Cbf16[∑r=0Kact−1πn,en,rZP(nKact+r)(p)].

节点 18 提供索引映射和权重,22 提供专家输出。该 kernel 在一个 token 内归约多个 expert 的贡献,输出恢复为 [54,D]。这是 本卡上的专家汇总,不是 TP AllReduce;它没有把 shared expert、residual 或下一层 norm 加进去。

这里将权重位置写在 finalize 阶段,是针对本 trace 的独立 finalize 路径。其他 MoE backend 可能把权重缩放放在不同阶段,不能仅凭一个模型公式就认为所有 backend 的 kernel 划分相同。

24. shared_gate_sigmoid_mul_add ​

目标:计算 shared-expert gate,缩放共享专家输出,再与 routed 输出相加。

先从 MoE 输入 Mn 计算每个 token 一个标量 gate:

ηn=σ(∑d=1DMn,dwshared gate,d).

再得到本卡完整 MoE 局部贡献:

Un(p)=Cbf16[Yr,n(p)+ηnYs,n(p)].

源码把 gate 的向量点积、sigmoid、乘 shared output、加 routed output 融在一起,内部加载后用 FP32 做这些运算,再写回输出 dtype。因此图中必须有 11 → 24:只给 24 两路专家结果是不够的,它还需要 M 来计算 gate。

这个 η 是每个 token 的共享专家门控,与节点 8 中逐通道的 SiLU(z) 不同。此处的 add 是 shared + routed,不是加 residual。

3.3 跨层闭环:最后的通信和 residual 在哪里 ​

节点 24 之后,下一层节点 1 完成:

Un=∑p=01Un(p),R~next,n=R2,n+Un,Rnext=Cbf16(R~next),Hnext=Cbf16[GN(R~next;wnext,in)].

在该 TP 路径中,各卡使用相同的 ηn,所以可以利用:

∑p(Yr,n(p)+ηnYs,n(p))=∑pYr,n(p)+ηn∑pYs,n(p).

这解释了为什么不需要分别对 shared、routed 输出各做一次 AllReduce:先在本卡合并,再做一次通信即可。浮点实现仍须保持约定的舍入位置,代数上等价不代表任意改写都逐 bit 相同。

本文按 trace 的记账方式,把层末 AR 归到下一层入口。因此 24 个 kernel 的列表没有单独的末尾 AR,但算法没有遗漏通信。模型最后一层的边界由末尾 norm/通信逻辑处理,不能机械地假定还存在一个普通 decoder layer。

3.4 把几类容易混淆的“gate”和“reduce”区分开 ​

名称所在节点实际含义
a→α2、4 产生,7 使用GDN 历史状态的衰减
b→β2、4 产生,7 使用当前 key 方向上的记忆修正力度
z→SiLU(z)5、6 产生,8 使用GDN 读出结果的逐通道门控
MLP gate/up14、17 或 21专家内部 SwiGLU 的两组中间特征
router logits / π12、16、18、23选择哪些 routed experts,以及如何混合它们
shared gate / η24对共享专家输出的逐 token 标量门控
Split-K reduce4同卡内合并一个 GEMM 的不同 K 分段
expert finalize23同卡内合并一个 token 的多个专家输出
TP AllReduce1、11、下一层 1合并不同 GPU 的局部输出

依赖图表示“读什么数据”,不等于 kernel 时间区间必须完全前后分离。 例如原始 trace 中 BA GEMM 2 的开始时间早于节点 1 的结束时间;PDL 允许消费者提前启动部分工作,但安全读取依赖数据仍需要正确的同步。不能仅凭时间线重叠就断言 Tensor Core 已经消费了尚未完成的数据。

3.5 源码与 trace 核对入口 ​

SGLang 以下链接固定到本文开头的 commit,避免后续主干变化造成解释漂移:

FlashInfer routing/NVFP4 的辅助核对使用本地 checkout 3a5f57d749737cc1b7c2f024ff037b8c78bb24fd 中的 trtllm_fused_moe_routing_custom.cuh、RoutingCustomPolicy.cuh、quantization/kernels/nvfp4_quantize.py。这用于解释与 trace 对应的算法和接口,不将该源码 SHA 冒充为采集环境中已核实的 wheel 构建 SHA。cuBLAS 内部实现不在 SGLang 源码中,节点 2、4 的拆分依据实际 kernel 名称和输入输出类型解释。

3.6 Decode-GDN 融合空间:按依赖和关键路径筛选 ​

本节是对上述原版 trace 与本地源码的静态分析,不是新增 B300 实测结果。遵守当前“不修改 GEMM 主体”的边界;纯计算优化归 KernelForge,涉及 AR 协议与通算融合的实现归 CommWeave。

3.6.1 先看哪些分支真的在等待 ​

以节点 1 的开始为相对时间 0,这个 layer 样本的关键时间点如下。时间来自 nsys 导出的 GPU 0 CSV,不能代表所有层、所有 step 和另一个 rank。

路径或节点相对开始 / 结束(μs)对融合决策的含义
BA:2 → 44.352 / 11.040比 QKVZ 早完成 10.752 μs,当前不是节点 6 的等待来源
QKVZ:3 → 56.688 / 21.792节点 6 需要等待这条较慢分支
6 → 722.240 / 66.335GDN 主链;7 自身耗时 38.143 μs
8 → 966.336 / 71.552输出 norm/gate 和 FP8 quant 位于主链上
router 1285.793 / 94.977已与节点 11 尾部重叠,不应在新实现中丢掉这段重叠
15 → 16 → 1895.713 / 110.241量化与 routing 在同一 stream 排队,直接影响 routed GEMM 启动
shared down 20106.497 / 110.081shared 分支已经完成
routed finalize 23203.042 / 208.162routed 分支比 shared 晚完成约 98.081 μs
shared/routed merge 24208.770 / 211.202两路汇总后才能进入下一层

所有 kernel duration 的和为 235.680 μs,而本列表从首个开始到最后结束的时间跨度为 211.202 μs;差异来自重叠等因素。duration 相加不是层延迟,也不是融合收益。 下一层入口 AR 未包含在这个跨度中,完整跨层比较必须纳入它。

3.6.2 候选清单 ​

下表“现有耗时和”仅用于描述参与改造的工作量,不是预期可节省时间。不同候选互相覆盖,不能把潜在收益相加。

候选现有 kernel 耗时和(μs)判断主要约束
8 + 94.736优先做纯计算原型head-wise norm 与 token-wise FP8 scale 的归约范围不同
16 + 1811.328优先 A/B 已有 fused routing 路径上游选择拆分有布局/资源方面的原因,融合不保证更快
23 + 247.552主链上的直接候选保留 finalize 的 BF16 舍入和前序 GEMM 的依赖同步
13 + 154.096值得试的双量化分支合并保留两种独立量化,改变排队次序时重新建立依赖
15 + 165.312路由入口的备选方案两者输入不同;NVFP4 的逐块量化与 top-K 的线程布局要兼容
17 + 194.960实现简单,但当前端到端优先级低shared 分支有约 98 μs 的完成时间余量
1 + 38.576通算融合候选保留 BF16 分支、FP8 分支、residual 和现有 BA 提前启动
11 + 13 + 1513.025通算融合候选,晚于独立双量化评估BF16 router、FP8 shared、FP4 routed 三路消费者均需正确调度
23 + 24 + next.1 + next.3跨层,当前 CSV 不足以完整计时优先复用已有 CommWeave 经验用当前版本、input8192/bs54 重测,不能套用旧实验百分比
6 + 743.455有空间,但需要重构状态更新分工同一卷积历史可能被多个 recurrent block 读取,原地更新存在竞争风险
7 + 840.799中长期布局研究recurrent 按 V 切块,norm 需要一个 head 的完整 V
6 + 7 + 8 + 948.191不建议作为第一轮同时跨越卷积、状态分块、head norm、token quant 四种布局
2 + 4、3 + 5、9 + 10、14 + 17、19 + 20 等—当前边界外涉及 GEMM 算法或 prologue/epilogue 修改

3.6.3 为什么先看 8 + 9,而不是直接把 7、8、9 全合起来 ​

8 先对每个 head 的 128 维输出做 norm/gate;9 再对当前 rank 的 32 个 head 拼接后的 4096 维向量求 FP8 scale。它们可以由一个 token 级 kernel 完成:

yn,h=Cbf16[RMSNorm(on,h)⊙SiLU(zn,h)],sn=maxh,j|yn,h,j|448.

量化必须基于已按原版舍入到 BF16 的 y,不能直接拿未舍入的 FP32 结果求 scale。可以省去独立的 quant launch,以及在无其他消费者时省去 BF16 中间张量的写回和读取。

但“一个 token 一个 CTA”的朴素方案在 bs54 下只有 54 个 CTA,可能不足以有效利用整张 B300。KernelForge 已有面向 8192/20480 token 的 norm/gate/FP8 实现,只能复用数值逻辑和验证方法;不能直接删掉 shape guard 就宣称 decode 得到收益。应比较整 token CTA、协作分组等布局,必要时保留小型第二阶段,接受“两个更快的 kernel”优于“一个更慢的 kernel”。

节点 7 当前设置 BV=min(next_power_of_2(V),32)。对 V=128,一个 value head 分给 4 个 Triton program,各自产生 32 个输出分量。节点 8 的 norm 需要全部 128 个分量,节点 9 的 scale 又需要全部本地 heads:

32-value tile→128-value head→4096-value token.

因此 7 + 8 需要跨 program 协作或重新扩大 program 的负责范围;7 + 8 + 9 还要解决跨 head 归约。扩大单个 program 可能提高寄存器压力、减少并行度。不能认为只要 o 在寄存器里,就能免费接上所有后处理。

3.6.4 Routing:先比较已有路径,再决定是否写新 kernel ​

本地 FlashInfer 3a5f57d... 的 routingIndicesClusterKernelBody 同时存在两种工作方式:

  • 传入原始 scores:在 cluster 内完成 top-K、softmax 后处理和 permutation。
  • 传入预先生成的 packed top-K:只做后续 permutation,对应当前 trace 的 16 → 18 分拆路径。

所以 16 + 18 不是从零发明一个新算法,先比较上游已有路径更合适。该 checkout 提供 FLASHINFER_ROUTING_FORCE_BLOCK_PER_TOKEN=off 作为关闭强制 split 选择的调试入口;在符合 small-batch/cluster dispatch 的条件下可走合并路径。这不是已确认适用于所有安装版本的稳定接口:实际 B300 wheel 必须核对该开关、编译 tier 和实际选择,并用 trace 确认是否真的由两个 launch 变成一个。环境变量会影响其他层的 routing,需分清单算子 A/B 和整模型效果。

另一个无需先融合的机会在节点 15:它只依赖 11 的 M,不依赖 router logits,但当前排在 stream 277 的 router 12 后面。可以先研究把量化安排到合适的另一条 stream,或将它与 13 合并,提前准备 routed GEMM 的输入。此时要显式保证 21 等到 FP4 数据和 routing metadata 都完成;跨 stream 读取不能只依靠 PDL 启动许可。

15 + 16 则是在 router 之后,同时读取 M 与 logits,一次完成 NVFP4 和 top-K。它主要省 launch 和部分排队开销,不是复用同一份输入张量的读取,因为 quant 读 hidden states,top-K 读 router scores。13 + 15、15 + 16、16 + 18 是可比较的组合选择,不宜直接全部叠在一起。

3.6.5 GDN 核心还有一个比“大融合”更温和的方向 ​

节点 7 的 4 个 value tiles 都会重新计算同一 value head 的 a/b 变换;同一 Q/K head 还可能被多个 value head 共享。这带来重复的 Q/K norm、query scale 和 gating 运算。

可研究把适量的前处理放入节点 6,让 7 更集中地处理 SSM state。这个方案可能仍然是两个 kernel,但减少重复工作,未必需要引入危险的跨 block 同步。必须保留:

  1. 卷积输出原有的 BF16 舍入;
  2. 原节点 7 中 Q/K norm 的 FP32 中间结果精度,不能迁移后随意存成 BF16;
  3. beta 的 BF16 舍入与 alpha 的计算顺序;
  4. padding slot、无效 state index、CUDA Graph 重放时的缓存更新语义。

至于完整 6 + 7,尤其要处理卷积 state 的读取与原地更新:仅指定一个 block 写回还不够,如果它先写回、其他 block 后读取历史,后者就可能读到新状态。需要保证历史读取完成后再更新,或使用独立快照、双缓冲或明确的协作方式。是否值得付出这些代价,必须先用 NCU 检查节点 7 的瓶颈;38.143 μs 的时长本身无法证明它被指数运算、带宽或寄存器中的哪一项限制。

3.6.6 验证顺序与验收口径 ​

建议先做 8 + 9 原型、16 + 18 上游路径 A/B,接着评估 13 + 15/量化提前调度、23 + 24。17 + 19 可作为精度与融合工具链的简单验证项,但不作为主要端到端收益目标。通信方案复用 CommWeave,6/7 的状态布局重构放在后面。

  • 精度:比较相应 BF16 中间值、FP8/FP4 编码和 scale;routing 比较选中专家及权重,permutation 用映射还原后的语义结果检查,不能只要求未规定顺序的工作队列字节相同;GDN 同时比较输出、conv state 和 SSM state,覆盖连续多步、padding 和状态槽重用。
  • 局部性能:使用实际 shape、充分预热和 CUDA Graph。8/9 比较消费它们的 out projection 启动与完成;routing 比较 routed GEMM 何时可用;23/24 比较到下一层投影的跨层子链。
  • 整体性能:同版本、同 GPU 对、相同 input8192/bs54 和相同 decode token 数,多轮交替 A/B,比较完整 decode step。不能把单个 layer 的 kernel 耗时差直接乘层数。
  • 重叠:保留已有 BA/router 提前启动及 routed GEMM/Finalize 的依赖同步。少一个 launch 但破坏原有重叠的版本,不算有效优化。

历史 CommWeave 的 bs54 跨层融合结果使用不同的 SGLang 基线和 input1/output128;AR/norm 双量化实验则使用 8192 行输入。这两类结果分别证明已有实现经验和局部可行性,都不是本文 input8192/bs54 decode 的新实测收益。

4. Decode-FullAttn ​

流程图

正在绘制图表…

flowchart TB
      X["previous_layer.shared_gate_sigmoid_mul_add"] --> K1["1<br/>input_allreduce_add_rmsnorm"]
      R["residual"] --> K1

      K1 --> K2["2<br/>qkv_proj_input_fp8_quant"]
      K2 --> K3["3<br/>qkv_proj_fp8_gemm"]
      K3 -->|q_gate, k| K4["4<br/>qk_rmsnorm_rope_gate_split"]
      K3 -->|v| K5["5<br/>store_kv_cache"]
      K4 -->|k| K5

      K4 -->|q| K6["6<br/>trtllm_full_attention_decode"]
      K5 -->|current_kv_cache| K6
      KV["historical_kv_cache"] --> K6

      K6 -->|attention_output| K7["7<br/>output_sigmoid_gate"]
      K4 -->|gate| K7
      K7 --> K8["8<br/>o_proj_input_fp8_quant"]
      K8 --> K9["9<br/>o_proj_fp8_gemm"]

      K9 --> K10["10<br/>attention_allreduce_add_rmsnorm"]
      K1 -->|residual| K10

      K10 -->|bf16| K11["11<br/>router_bf16_gemm"]
      K11 --> K15["15<br/>routing_block_scores"]
      K15 --> K17["17<br/>routing_cluster_indices"]

      K10 --> K12["12<br/>shared_gate_up_input_fp8_quant"]
      K12 --> K13["13<br/>shared_gate_up_fp8_gemm"]
      K13 --> K16["16<br/>shared_silu_mul"]
      K16 --> K18["18<br/>shared_down_input_fp8_quant"]
      K18 --> K19["19<br/>shared_down_fp8_gemm"]

      K10 --> K14["14<br/>routed_input_nvfp4_quant"]
      K14 --> K20["20<br/>routed_gate_up_swiglu_fp4_gemm"]
      K17 -->|routing_metadata| K20

      K20 --> K21["21<br/>routed_down_fp4_gemm"]
      K17 -->|routing_metadata| K21

      K21 --> K22["22<br/>routed_finalize"]
      K17 -->|routing_weights, token_mapping| K22

      K22 -->|routed_output| K23["23<br/>shared_gate_sigmoid_mul_add"]
      K19 -->|shared_output| K23
      K10 -->|hidden_states| K23

      K23 --> N1["next_layer:<br/>1. input_allreduce_add_rmsnorm"]
      K10 -->|residual| N1

      classDef blue fill:#bdd7ee,stroke:#4472c4,color:#000;
      class K1,K10,N1 blue;