Quartz 5

Kimi-k3 MoE路径完整分析

4 min read

SM100

核心class为KimiK3MoE，里面核心调用get_moe_impl_class(…)创建expert class

# Routed experts (operate in moe_hidden_size space)
# gate_up_interleaved=False: K3 loads per-expert w1/w3 into non-interleaved layout
self.experts = get_moe_impl_class(moe_quant_config)(
    num_experts=getattr(config, "n_routed_experts", config.num_experts),
    top_k=config.num_experts_per_token,
    hidden_size=self.moe_hidden_size,
    intermediate_size=config.moe_intermediate_size,
    layer_id=self.layer_idx,
    quant_config=moe_quant_config,
    routed_scaling_factor=self.routed_scaling_factor,
    activation=config.hidden_act,
    gemm1_alpha=config.activation_situ_beta,
    gemm1_clamp_limit=config.activation_situ_linear_beta,
    gate_up_interleaved=False,
    # trtllm fused-routing MoE backends (e.g. nvfp4 w4a4) route inside
    # the kernel and require the routing method; K3 uses DSv3-style
    # grouped topk with e_score_correction_bias.
    routing_method_type=getattr(
        config, "routing_method_type", RoutingMethodType.DeepSeekV3
    ),
    prefix=add_prefix("experts", prefix),)

其中，对于普通K3而言，采用量化方法是mxfp4，get_moe_impl_class实际使用mxfp4.py中的Mxfp4MoEMethod class，其中开启MegaMoE后，具体的调用方法有些许区别；

  • 在create_weights(...)中，如果开启MegaMoE，则不对intermediate size进行padding，随后，统一的将gate_up_proj合成第一个GEMM，也即[num_experts, 2*intermediate_size, H/2]，其中，H/2是因为具体的weight是FP4，每个元素只占据1/2个FP8，其对应的scale是32个K group对应同一个scale，具体的layout为[num_experts, 2*intermediate_size, H/32]，对于self.w2，也即fc2，同样获取对应的weight和scale；
  • 在process_weights_after_loading(...)中，如果是MegaMoE的backend，则首先将weight进行pack，因为DeepGEMM的期望权重输入是INT8，调用两次view，获得weight，随后，对w13和w2的两个scale，分别将其转换为fp32，然后调用DeepGEMM的transform_sf_into_required_layout(...)将他们转换成megaMoE需要的layout，具体而言，将4个scale字节打包为一个INT32，将其改为MN-major，和TMA的layout对齐，因为实际的GEMM计算是一个MN的部分K可能一次算完，此时需要做对应的scale，就得要求scale连续，因此，需要做layout调整；调整完scale后，核心调整w13和w2的weight，核心调用transform_weights_for_mega_moe(...)，在其中，作如下操作
    • L1 权重和 scale 做 gate/up 交错：[gate 0..7, up 0..7, gate 8..15, …]，粒度是 8 行。这样 L1 epilogue 从 TMEM 读一块，gate 和 up 就挨在一起，能当场算 SiTU/SwiGLU，不用再读一遍。
    • L1 和 L2 的 scale 都做 UTCCP 转置：每 128 行里按 [4,32] → [32,4] 重排。这是 tcgen05.cp 把 scale 拷进 TMEM 时要求的 4x32 布局。
    • L2 权重本身不动。
    •

原来 (K-major): 内存顺序 [e][row][k_group] 同一行的各组 scale 相邻
现在 (MN-major): 内存顺序 [e][k_pack][row] 同一 k_pack 下各行的 scale 相邻

为什么要 MN-major：kernel 每次处理的是一个 BLOCK_N 行 × 一段 K 的权重块，需要用一条 TMA 把”这一段 K 上、这 BLOCK_N 行”的 scale 一次搬进来。行方向连续才能这样搬。
```

  • 在kimi_k3.py中，获取symm buffer，核心调用_get_mega_moe_symm_buffer(…)，具体而言，symmertric memory中有如下的view
    • buf.x：[T, H] fp8 e4m3
    • buf.x_sf：[T, H/128] int32（每个 int32 装 4 个 UE8M0，所以实际是每 32 个元素一个 scale）
    • buf.topk_idx：[T, k] int64
    • buf.topk_weights：[T, k] fp32
  • 具体forward(…)时候，调用_forward_unfused(...)随后调用run_experts(…)，随后再调用_forward_mega_experts(...)函数，
    • 首先调用mega_moe_pre_dispatch(…)，bf16按照每32个元素一组求absmax，算出UE8M0 scale，量化成fp8 e4m3写入buf.x，scale字节写进buf.x_sf，topk的id和权重直接写入symm buffer；
    • 随后调用_configure_mega_moe_deep_gemm_num_sms(…)，根据SGLANG_OPT_DEEPGEMM_MEGA_MOE_RESERVED_SMS空出对应的SM給Shared expert使用；
    • 调用deep_gemm.fp8_fp4_mega_moe(…)进行实际的计算；

Graph View