[Kimi-K3] O(1) expert weight lookup in load_weights (#38805)

Signed-off-by: JinYan Su <751080330@qq.com>
Co-authored-by: Claude Fable 5.1 <noreply@anthropic.com>
Co-authored-by: Liangsheng Yin <hnyls2002@gmail.com>
Co-authored-by: hnyls2002 <lsyincs@gmail.com>

python/sglang/srt/models/kimi_k3.py

@@ -8,6 +8,7 @@
 
 import logging
 import os
+import re
 from collections.abc import Iterable
 from functools import cached_property
 from types import SimpleNamespace
@@ -134,6 +135,10 @@ from sglang.srt.utils.common import (
 )
 
 logger = logging.getLogger(__name__)
+
+# `experts.<expert_id>.<w1|w2|w3>.` fragment of a checkpoint tensor name, the
+# key FusedMoE.make_expert_params_mapping entries match on.
+_EXPERT_WEIGHT_NAME = re.compile(r"experts\.\d+\.w[123]\.")
 _is_hip = is_hip()
 _is_npu = is_npu()
 _aiter_k3_opt = get_bool_env_var("SGLANG_AITER_K3_OPT")
@@ -214,10 +219,10 @@ def _k3_bf16_gemm(
     return torch.mm(x, weight.t(), out=out)
 
 
-# Fully fused KDA decode step (conv1d + delta rule + gated RMSNorm in one
+# Fully fused KDA decode step (kernels/ops/attention/kda_fused_decode). The
-# kernel, kernels/ops/attention/kda_fused_decode). The model hands the output-norm gate
+# model hands the output-norm gate to the KDA backend via an attempt-and-verify
-# to the KDA backend via an attempt-and-verify stash on the attention layer;
+# stash on the attention layer; unconsumed stashes fall back to the unfused
-# unconsumed stashes fall back to the unfused chain + o_norm here.
+# chain + o_norm here.
 
 
 def _merge_weights_as_views(
@@ -246,10 +251,8 @@ def _merge_weights_as_views(
 # input_layernorm / post_attention_layernorm, which the communicator expects
 # to own. Instead the MLP/MoE modules gather/scatter around their own body:
 # attention and the attn-res buffers stay in local (per-DP-rank) token space,
-# the MLP/MoE runs on the DP-gathered global batch with plain full-TP
+# the MLP/MoE runs on the DP-gathered global batch, and the delayed prefix_sum
-# semantics (its internal all-reduces are unchanged and required — the latent
+# add stays local, applied after the scatter back.
-# reduce must happen in latent space before the norm), and the delayed
-# prefix_sum add stays local, applied after the scatter back.
 def _dp_local_buffer_group():
     """Symmetric-memory group for the local DP buffer (mirrors
     CommunicateSummableTensorPairFn._scatter_hidden_states)."""
@@ -518,13 +521,12 @@ class KimiK3MoE(nn.Module):
             ),
         )
 
-        # MegaMoE (deep_gemm fused a2a+GEMM over the EP symm buffer): a drop-in
+        # MegaMoE (deep_gemm fused a2a+GEMM over the EP symm buffer) replaces
-        # replacement for the routed experts call below. K3 routes ALL batches
+        # the routed experts call below. K3 routes ALL batches through it when
-        # through it when enabled — the megamoe backend's non-mega fallback is
+        # enabled -- its non-mega fallback is a StandardDispatcher without a2a,
-        # a StandardDispatcher without a2a, which is wrong for scattered
+        # wrong for scattered tokens -- so
-        # tokens — so SGLANG_OPT_DEEPGEMM_MEGA_MOE_NUM_MAX_TOKENS_PER_RANK must
+        # SGLANG_OPT_DEEPGEMM_MEGA_MOE_NUM_MAX_TOKENS_PER_RANK must cover the
-        # cover the per-rank prefill chunk. The DeepGEMM MegaMoE SiTU kernel
+        # per-rank prefill chunk.
-        # bakes in the K3 activation constants.
         self._use_mega_moe = get_moe_a2a_backend().is_megamoe()
         self._mega_intermediate_size = moe_intermediate_size
         self._mega_top_k = config.num_experts_per_token
@@ -538,12 +540,10 @@ class KimiK3MoE(nn.Module):
                 "got a checkpoint with different constants"
             )
 
-        # EP a2a backends (megamoe / DeepEP / Mooncake / Ascend-FuseEP / MoRI)
+        # EP a2a backends move each row to its experts directly, so the MoE
-        # move each row to its experts directly, so the MoE region can consume
+        # region consumes whatever rows this rank holds (SP-MoE shard or
-        # whatever rows this rank holds — an SP-MoE token shard (attn_tp > 1) or
+        # DP-local batch) with every global token dispatched exactly once.
-        # the DP-local batch (DP attention) — with every global token dispatched
+        # No DP gather and no TP reduce anywhere in the region.
-        # exactly once. No DP gather and no TP reduce is needed anywhere in the
-        # region.
         _a2a_backend = get_moe_a2a_backend()
         self._ep_a2a = (
             _a2a_backend.is_megamoe()
@@ -553,13 +553,10 @@ class KimiK3MoE(nn.Module):
             or _a2a_backend.is_mori()
         )
 
-        # Defer the trtllm-gen finalize (top-k weighted unpermute) out of the
+        # Defer the trtllm-gen finalize (top-k weighted unpermute) into the
-        # MoE op and fuse it into the push all-reduce's staging pass
+        # push all-reduce's staging pass (k3_ar_fusion.finalize_all_reduce_push_norm)
-        # (k3_ar_fusion.finalize_all_reduce_push_norm): the rank-local latent
+        # so the rank-local latent never materializes. Sizes beyond the push
-        # never materializes. Only the situ trtllm-gen path serves the
+        # window fall back to the in-op finalize at runtime.
-        # deferral, on either routing form (packed ids on a fused route+quant
-        # hit, unpacked fp32 weights otherwise); sizes beyond the push window
-        # fall back to the in-op finalize at runtime (finalize_push_fits).
         self._defer_moe_finalize = (
             get_moe_runner_backend().is_flashinfer_mxfp4()
             and config.hidden_act == "situ"
@@ -615,19 +612,11 @@ class KimiK3MoE(nn.Module):
         else:
             self.shared_experts = None
 
-        # SBO (single batch overlap): the shared experts read a fixed slab of
+        # SBO (single batch overlap): shared experts are bf16 + tp1-replicated
-        # weights the routed path never touches (bf16 — the checkpoint leaves
+        # (~264 MB/layer/rank), the routed path is a2a-latency bound in decode
-        # shared_experts unquantized — and tp1-replicated under EP a2a, so
+        # with HBM idle; issue shared on the side stream, join before the tail
-        # ~264 MB per layer per rank), while the routed path is a2a-latency
+        # add. NPU shared-expert TP can instead overlap the shared collectives
-        # bound in decode with HBM mostly idle. Issue the shared experts on the
+        # with SGLANG_NPU_FINE_GRAINED_MOE_DUAL_STREAM.
-        # side stream so the two run concurrently instead of back to back; the
-        # join happens right before the tail add. Measured on 2x4 GB300
-        # (TP8/EP8 MegaMoE + SP-MoE): +4~5% output tok/s and −5% ITL over
-        # bs 1–32, GSM8K unchanged — so it is on whenever the shape allows,
-        # no flag.
-        # NPU shared-expert TP can also overlap the shared
-        # collectives using SGLANG_NPU_FINE_GRAINED_MOE_DUAL_STREAM. Otherwise
-        # the collectives stay on the current stream.
         self._sbo_shared_overlap = (
             self._ep_a2a
             and self.shared_experts is not None
@@ -668,24 +657,19 @@ class KimiK3MoE(nn.Module):
             self.routed_expert_norm = None
             self.routed_expert_up_proj = None
 
-        # Static eligibility for fusing the fused-front latent all-reduce with
+        # Static eligibility for fusing the latent all-reduce with the RMSNorm
-        # the RMSNorm epilogue (SGLANG_K3_AR_FUSION). The kernel views the flat
+        # epilogue (SGLANG_K3_AR_FUSION): kernel needs latent == NORM_DIM and
-        # [latent | shared] buffer as [3N, NORM_DIM] rows and norms the first N,
+        # shared == 2*NORM_DIM (3584 / 7168). Decided once so the hot path
-        # so it requires latent width == NORM_DIM and shared width == 2*NORM_DIM
+        # reads a bool.
-        # (K3: 3584 / 7168). Decided once here so the hot path only reads a bool
-        # and never re-validates dims per forward.
         self.fuse_ar_norm = (
             self.routed_expert_norm is not None
             and self.moe_hidden_size == k3_ar_fusion.NORM_DIM
             and hidden_size == 2 * k3_ar_fusion.NORM_DIM
         )
         # Static eligibility for the column-parallel up_proj tail (gemm_ag):
-        # per-rank 1/8-column GEMV -> multicast all-gather staged in the v2
+        # 1/8-column GEMV + multicast all-gather + spin-add3 replaces the
-        # push workspace -> spin-add3 with shared_output (+ prefix_sum),
+        # replicated [3584, 7168] GEMM (~1.5-2x at decode). Dims fixed to
-        # replacing the replicated [3584, 7168] GEMM + _add3 (~1.5-2x at
+        # fuse_ar_norm's over TP8; per-batch checks in k3_ar_fusion.
-        # decode sizes, 1/8 of the weight bytes read per rank). Kernel dims
-        # are fixed to fuse_ar_norm's (3584 -> 7168) over TP8; per-batch
-        # capacity checks live in k3_ar_fusion.gemm_ag_up_fits.
         self._gemm_ag_up_eligible = (
             self.fuse_ar_norm
             and self.tp_size == 8
@@ -724,12 +708,9 @@ class KimiK3MoE(nn.Module):
                 self.routed_expert_down_proj,
             ]
         elif envs.SGLANG_K3_FUSED_FRONT.get():
-            # EP a2a: the shared experts are tp1-replicated and run on the side
+            # Merge the router gate into the latent down-proj so one GEMM reads
-            # stream, so they stay out of the merge -- but the router gate and the
+            # hidden_states once: the 896-row gate alone is too few to use the
-            # latent down-proj still read the same hidden_states, and merging just
+            # machine well; folded into the 3584-row down-proj it is near-free.
-            # those two is what lets one GEMM read the activations once. The gate
-            # GEMM alone is only 896 rows, which is too few to use the machine
-            # well; folded into the 3584-row down-proj it comes almost free.
             mods = [self.gate, self.routed_expert_down_proj]
         else:
             return
@@ -1078,17 +1059,13 @@ class KimiK3MoE(nn.Module):
     ) -> torch.Tensor:
         """Front section with three separate GEMMs, each reading
         hidden_states: shared-expert MLP, router gate, latent down-proj."""
-        # Shared experts on original hidden_states. Under SBO they go to the
+        # Shared experts on original hidden_states; under SBO they go to the
-        # side stream and are joined at the tail (see _sbo_shared_overlap).
+        # side stream, joined at the tail. CUDA issues this after the front so
-        #
+        # they overlap the routed a2a rather than the front GEMMs; NPU starts
-        # CUDA issues this after the front so the shared experts overlap the
+        # before the front. Fine-grained NPU overlap splits at the dispatch
-        # routed a2a rather than the front GEMMs. NPU starts the shared branch
+        # boundaries:
-        # before the front. Fine-grained NPU overlap splits it at the complete
-        # dispatch boundaries:
         #   current: front ---------- dispatch ---------- routed GEMMs -- tail
         #   alt:     all-gather ----- shared MLP -------- reduce-scatter
-        # Shared and routed GEMMs wait for each other at phase boundaries;
-        # each can run beside the other branch's communication.
         fine_grained_overlap = self._can_overlap_shared_experts_npu(hidden_states)
         shared_input = None
         shared_output = None
@@ -1186,10 +1163,8 @@ class KimiK3MoE(nn.Module):
         if _is_npu and self._sbo_shared_overlap:
             issue_shared()
 
-        # Front: gate + TopK (+ latent down-proj when the merged front covers it).
+        # Front: gate + TopK (+ latent down-proj when merged). Strategy table:
-        # The gate and the latent down-proj read the same hidden_states, so the
+        # kernels/ops/moe/moe_front.py.
-        # merged-weight strategies compute both in one GEMM; see
-        # kernels/ops/moe/moe_front.py for the strategy table.
         routed_input = self._ep_front(hidden_states)
         if routed_input is None and not fine_grained_overlap:
             routed_input = self._ep_front_overlap(hidden_states)
@@ -1414,9 +1389,8 @@ class KimiK3MoE(nn.Module):
                 # routed GEMMs it overlaps (K3 dims are fixed; tuned here)
                 k3_ar_fusion.all_reduce_low_sm(shared_output, num_blocks=4, unroll=8)
             current_stream.wait_stream(self.alt_stream)
-            # NOTE: the latent AR must stay serialized after the shared AR
+            # The latent AR must stay serialized after the shared AR (both
-            # (both reuse the v2 pull semaphores; concurrent calls would
+            # reuse the v2 pull semaphores); the join above does it.
-            # corrupt each other's barrier windows) — the join above does it.
             if defer_finalize:
                 # finalize folded into the push AR's staging pass; the norm
                 # covers every latent row
@@ -1471,10 +1445,8 @@ class KimiK3MoE(nn.Module):
             latent = self._latent_norm(latent)
         out, _ = self.routed_expert_up_proj(latent)
 
-        # prefetch_bc: b (shared_output) was produced by the all-reduce and
+        # prefetch_bc: b and c complete before the norm / up_proj chain
-        # c (prefix_sum) even earlier; the AR is a plain launch (full
+        # starts; only `a`'s producer can still be in flight at PDL entry.
-        # barrier), so both are complete once the norm / up_proj GEMM chain
-        # starts — only `a`'s producer can still be in flight at PDL entry.
         return _add3(out, shared_output, prefix_sum, prefetch_bc=True)
 
     def forward(
@@ -1537,9 +1509,8 @@ class KimiK3DeltaAttention(nn.Module):
         super().__init__()
         self.all_reduce_fusion = all_reduce_fusion
         # Side stream for the [f_a|b] + f_b tiny GEMVs: they read only
-        # hidden_states, so they can run concurrently with the wide fused
+        # hidden_states, so they overlap the wide [q,k,v,g] GEMM on the main
-        # [q,k,v,g] GEMM on the main stream (graphed decode/verify only).
+        # stream (graphed decode/verify only; SM bound as the MLA gate stream).
-        # Same SM bound rationale as the MLA gate stream.
         self._bfa_alt_stream = bfa_alt_stream
         self._bfa_bs_limit = (
             (128 if get_platform().is_blackwell else 64)
@@ -1547,10 +1518,8 @@ class KimiK3DeltaAttention(nn.Module):
             else 0
         )
         self.tp_size = get_parallel().tp_size
-        # KDA is an attention layer: all head-sharded params must follow the
+        # KDA is an attention layer: head-sharded params follow the
-        # attention-TP group (= tp under plain TP, = 1 under DP attention),
+        # attention-TP group, matching the mamba state cache sizing.
-        # matching the mamba state cache sizing (KimiLinearCacheParams uses
-        # get_attention_tp_size). Mirrors GLM5-next's head_shard_size pattern.
         self.attn_tp_size = get_parallel().attn_tp_size
         self.attn_tp_rank = get_parallel().attn_tp_rank
         self.hidden_size = hidden_size
@@ -1584,12 +1553,11 @@ class KimiK3DeltaAttention(nn.Module):
         self.do_fuse_qkvbfg = quant_config is None and self.attn_tp_size == self.tp_size
 
         if self.use_full_rank_gate:
-            # Fuse only the alignment-friendly wide projections [q, k, v, g]
+            # Fuse only the wide projections [q, k, v, g]: folding b (12/rank)
-            # (6144/rank at TP8). Folding b (12/rank) and f_a (128, replicated)
+            # and f_a (128, replicated) in skews the output dim and degrades
-            # in as well skews the output dim to 6284 and measurably degrades
+            # GEMM kernel selection; they stay as separate tiny GEMVs. ROCm
-            # the GEMM kernel selection; they stay as separate tiny GEMVs.
+            # reverses this below the token threshold
-            # (ROCm reverses this below the token threshold -- see
+            # (_merge_kda_inproj_weights_hip).
-            # _merge_kda_inproj_weights_hip.)
             self.fused_qkvg_proj = MergedColumnParallelLinear(
                 self.hidden_size,
                 [
@@ -1762,12 +1730,9 @@ class KimiK3DeltaAttention(nn.Module):
         )
         self.qkv_conv1d.weight.data = self.qkv_conv1d.weight.data.unsqueeze(1)
 
-        # K3 checkpoint stores A_log as [head_dim] (128), but the FLA kernel
+        # Checkpoint stores A_log as [head_dim]; the FLA kernel expects
-        # expects exactly local_num_heads elements.  We define the param as
+        # local_num_heads. The loader handles both the old 4-D and the K3 1-D
-        # [1, 1, local_num_heads, 1] (matching the kimi_linear.py convention)
+        # formats by narrowing to the first num_heads then TP-sharding.
-        # and attach a custom weight_loader that handles both the old 4-D
-        # format and the K3 1-D [head_dim] format by narrowing to the first
-        # num_heads elements then TP-sharding.
         self.A_log = nn.Parameter(
             torch.empty(1, 1, self.local_num_heads, 1, dtype=torch.float32)
         )
@@ -1806,16 +1771,12 @@ class KimiK3DeltaAttention(nn.Module):
             quant_config=quant_config,
             tp_rank=self.attn_tp_rank,
             tp_size=self.attn_tp_size,
-            # Reduce within the attn-TP group: the default reduce path uses
+            # Reduce within the attn-TP group: the default full-TP collective
-            # the full-TP collective, which at attn_tp>1 is both the wrong
+            # is the wrong group at attn_tp>1 (sums across DP groups) and
-            # group (sums across DP groups) and asymmetric vs idle DP ranks
+            # deadlocks idle DP ranks. Off under all_reduce_fusion: the fused
-            # (deadlocks the per-layer DP gather). Off under all_reduce_fusion:
+            # AR reduces in place out of a symmetric buffer, and this flag
-            # the fused AR does the reduce itself (reduce_results=False) and the
+            # would wrap the GEMM in its own symm_ctx, defeating the
-            # forward hands o_proj a slice of a persistent symmetric region —
+            # caller-owned buffer.
-            # leaving this True would wrap the GEMM in
-            # use_symmetric_memory(attn_tp), which allocates its own output and
-            # so defeats the caller-owned buffer. At the fusion config
-            # attn_tp==tp so the fused full-TP reduce is the same group anyway.
             use_dp_attention_reduce=not self.all_reduce_fusion,
             prefix=f"{prefix}.o_proj",
         )
@@ -2061,10 +2022,8 @@ class KimiK3DeltaAttention(nn.Module):
                     and 0 < hidden_states.shape[0] <= self._qkvgbfa_bs_limit
                 ):
                     # ROCm only. One GEMM for the whole in-proj: the [f_a|b]
-                    # tail rides along in the wide projection's bandwidth
+                    # tail rides the wide projection's bandwidth (~30% of the
-                    # instead of paying its own launch. Worth ~30% of the
+                    # in-proj at decode on gfx950, SGLANG_ROCM_K3_FUSE_KDA_INPROJ).
-                    # in-proj at decode on gfx950; see
-                    # SGLANG_ROCM_K3_FUSE_KDA_INPROJ.
                     fused_states = self.fused_qkvg_proj.quant_method.apply(
                         self._qkvgbfa_layer, hidden_states, None
                     )
@@ -2080,9 +2039,8 @@ class KimiK3DeltaAttention(nn.Module):
                     and 0 < hidden_states.shape[0] <= self._bfa_bs_limit
                 ):
                     # Issue the tiny [f_a|b] + f_b GEMVs on the side stream,
-                    # then the wide [q,k,v,g] GEMM on the main stream; both
+                    # then the wide [q,k,v,g] GEMM on the main stream (both
-                    # read only hidden_states. Join before the split's
+                    # read only hidden_states); join before the consumers.
-                    # consumers touch beta/forget_gate.
                     alt = self._bfa_alt_stream
                     cur = torch.cuda.current_stream()
                     alt.wait_stream(cur)
@@ -2282,20 +2240,11 @@ class KimiK3MLAAttention(DeepseekV2AttentionMLA):
             self.o_proj.use_dp_attention_reduce = True
         k3_gemm_ar.maybe_wrap_o_proj(self.o_proj)
         if self.all_reduce_fusion:
-            # reduce_results=False was passed through super().__init__ above;
+            # Hand the GEMM a slice of the persistent symmetric buffer
-            # the fused all-reduce does the reduce itself and reduces the o_proj
+            # (k3_ar_fusion.symm_buffer); the fused AR reduces it in place.
-            # output in place, so hand the GEMM a slice of the persistent
+            # The captured name must differ from the gate block's
-            # symmetric buffer (k3_ar_fusion.symm_buffer) and do NOT set
+            # `_orig_o_proj_forward` — closures capture the __init__ local by
-            # use_dp_attention_reduce — its inner attn_tp symm_ctx allocates its
+            # reference, and reusing it would rebind to this wrapper (recursion).
-            # own output and would defeat the caller-owned buffer. At the fusion
-            # config (attn_tp==tp) the fused full-TP reduce is the same group as
-            # the attn_tp reduce.
-            # The wrap is installed before the output-gate wrap below so the
-            # gate multiply stays outside it and only the o_proj GEMM writes the
-            # region slice. NOTE: the captured name must differ from the
-            # gate block's `_orig_o_proj_forward` — closures capture the
-            # __init__ local by reference, and reusing the name would rebind it
-            # to this wrapper (infinite recursion + nested pool enter).
             _symm_inner_o_proj_forward = self.o_proj.forward
             _symm_o_proj = self.o_proj
 
@@ -2309,11 +2258,9 @@ class KimiK3MLAAttention(DeepseekV2AttentionMLA):
 
             self.o_proj.forward = _symm_o_proj_forward
         else:
-            # K3 has no LayerCommunicator, so o_proj (reduce_results=True by
+            # K3 has no LayerCommunicator, so o_proj must reduce within the
-            # default here, unlike deepseek's communicator flow) must reduce
+            # attn-TP group itself: the default full-TP collective is the wrong
-            # within the attn-TP group itself — the default full-TP collective
+            # group at attn_tp>1 and deadlocks idle DP ranks.
-            # is the wrong group at attn_tp>1 and deadlocks against idle DP
-            # ranks.
             self.o_proj.use_dp_attention_reduce = True
         if self.use_output_gate:
             projection_size = config.num_attention_heads * config.v_head_dim
@@ -2328,11 +2275,10 @@ class KimiK3MLAAttention(DeepseekV2AttentionMLA):
                 tp_size=get_parallel().attn_tp_size,
                 prefix=f"{prefix}.g_proj",
             )
-            # Output gate must multiply the TP-local attention output right
+            # Output gate multiplies the TP-local attention output right
-            # before o_proj (vLLM: attn_out * sigmoid(g_proj(hidden_states))).
+            # before o_proj; o_proj is invoked deep inside
-            # o_proj is invoked deep inside DeepseekV2AttentionMLA forward
+            # DeepseekV2AttentionMLA forward cores, so wrap its forward at
-            # cores, so wrap its forward at the instance level; the module
+            # the instance level (weights, reduce_results, loading untouched).
-            # itself (weights, reduce_results, loading path) is untouched.
             self._gate_hidden_states = None
             # (gate, producer stream) issued on the alt stream by forward();
             # None when the lazy path computes the gate here instead.
@@ -2460,24 +2406,13 @@ class KimiK3DecoderLayer(nn.Module):
             and layer_idx >= config.first_k_dense_replace
             and layer_idx % config.moe_layer_freq == 0
         )
-        # SP-MoE (EP a2a backend — megamoe, DeepEP, Mooncake, Ascend-FuseEP or
+        # SP-MoE (EP a2a backend): o_proj defers its attention-TP reduction;
-        # MoRI): o_proj defers its attention-TP reduction; this layer completes
+        # this layer completes it as a reduce-scatter so the whole MoE region
-        # it as a reduce-scatter
+        # runs on 1/attn_tp of the rows, then all-gathers back after the tail
-        # so the whole MoE region (agg2, norms, gate, latent projs, tp1
+        # add. RS+AG moves the same bytes the AR did, the shared-expert AR
-        # shared experts, EP a2a dispatch) runs on 1/attn_tp of the rows,
-        # then all-gathers rows back after the MoE tail add. RS+AG moves the
-        # same bytes the o_proj all-reduce did, the shared-expert all-reduce
         # disappears via tp1 weights, and each rank dispatches only its shard
-        # through the a2a (kills the attn_tp-fold dispatch redundancy) —
+        # through the a2a. Same under DP attention. Dense layers excluded: no
-        # strictly less communication + MoE-front compute /attn_tp. Works the
+        # per-token decomposition survives a token shard.
-        # same under DP attention: the attn_tp group is then the
-        # within-replica subgroup, rows are the DP-local batch, and
-        # KimiK3MoE skips the DP gather under EP a2a so the shard flows
-        # straight into the a2a. With attn_tp == 1 (full DP attention) there
-        # is no attention reduce to convert — the MoE-side gather skip alone
-        # removes the replication. Dense layers are excluded: their
-        # column-parallel MLP has no per-token decomposition that survives a
-        # token shard.
         _a2a_backend = get_moe_a2a_backend()
         self._sp_moe = (
             (
@@ -2491,14 +2426,9 @@ class KimiK3DecoderLayer(nn.Module):
             and get_parallel().attn_tp_group.world_size > 1
         )
 
-        # The fused all-reduce only serves the attn-res path (attn_res is
-        # config-static), so the standard path stays byte-for-byte untouched
-        # and always sees a reduced attention output.
         # Mutually exclusive with SP-MoE: both complete o_proj's deferred
         # reduction, but SP-MoE reduce-scatters to a shard whereas the fusion
-        # produces the full batch in a symm buffer — an SP-MoE layer builds
+        # produces the full batch in a symm buffer.
-        # o_proj in the plain deferred-reduce config (reduce_results forced off
-        # below) and reduce-scatters instead.
         attn_tp_size = get_parallel().attn_tp_size
         self.all_reduce_fusion = (
             not self._sp_moe
@@ -2658,15 +2588,12 @@ class KimiK3DecoderLayer(nn.Module):
         if forward_batch.forward_mode.is_idle():
             return hidden_states
 
-        # mlp-sync (DP attention OR MoE a2a/EP — require_mlp_sync) pads
+        # mlp-sync pads extend batches to a multiple of attn_tp_size, but the
-        # extend batches to a multiple of attn_tp_size
+        # attention metadata covers only the real tokens: flashinfer ragged
-        # (prepare_mlp_sync_batch ceil_align), but the attention metadata
+        # prefill rejects the row mismatch, and silent paths would write the
-        # (qo_indptr / query_start_loc) covers only the real tokens — the
+        # padded rows' garbage KV through zero-padded out_cache_loc (clobbering
-        # flashinfer ragged prefill rejects the row mismatch, and silent
+        # pool slot 0 -> cross-request corruption). Run attention on the real
-        # paths write the padded rows' garbage KV through the zero-padded
+        # rows and zero-pad the output back.
-        # out_cache_loc entries (clobbering pool slot 0 → cross-request
-        # corruption). Run attention on the real rows and zero-pad the
-        # output back; padded rows are discarded downstream.
         num_padded = hidden_states.shape[0]
         num_real = num_padded
         if self._trim_padded_attn and forward_batch.forward_mode.is_extend():
@@ -2940,13 +2867,10 @@ class KimiK3LinearModel(nn.Module):
         else:
             self.embed_tokens = PPMissingLayer()
 
-        # Multi-stream pool (deepseek_v4 pattern): every alt stream is
+        # Alt streams threaded down to the layers. Slots:
-        # constructed here and threaded down to the layers. Slots:
         #   [0] MoE dual-stream shared-expert tail
         #   [1] DeepseekV2AttentionMLA base internals (forwarded; unused by K3)
         #   [2] MLA output-gate GEMM, overlaps the attention core
-        # (The attn-res bank write no longer needs a stream: it is fused
-        # into the agg1 fast kernel, see AttnResidual.forward(write=True).)
         # Disable on HIP code path.
         self.alt_streams = None if _is_hip else [torch.cuda.Stream() for _ in range(3)]
 
@@ -3341,6 +3265,11 @@ class KimiK3LinearForCausalLM(nn.Module):
 
         params_dict = dict(self.named_parameters())
         loaded_params: set[str] = set()
+        # Keyed by the `experts.<id>.<proj>.` fragment (see _EXPERT_WEIGHT_NAME).
+        expert_params_lookup = {entry[1]: entry for entry in expert_params_mapping}
+        assert all(
+            _EXPERT_WEIGHT_NAME.fullmatch(key) for key in expert_params_lookup
+        ), "ckpt expert names diverged from _EXPERT_WEIGHT_NAME; no expert would load"
 
         num_hidden_layers = self.config.num_hidden_layers
         for args in weights:
@@ -3432,27 +3361,28 @@ class KimiK3LinearForCausalLM(nn.Module):
                 weight_loader(param, loaded_weight, shard_id)
                 break
             else:
-                for idx, (param_name, weight_name, expert_id, shard_id) in enumerate(
+                expert_match = _EXPERT_WEIGHT_NAME.search(name)
-                    expert_params_mapping
+                expert_entry = (
-                ):
+                    expert_params_lookup.get(expert_match.group(0))
-                    if weight_name not in name:
+                    if expert_match
-                        continue
+                    else None
+                )
+                if expert_entry is not None:
+                    param_name, weight_name, expert_id, shard_id = expert_entry
                     name = name.replace(weight_name, param_name)
                     # Skip experts of layers outside a truncated config (e.g.
                     # num_hidden_layers override), mirroring the non-expert
                     # `name not in params_dict` guard below.
-                    if name not in params_dict:
+                    if name in params_dict:
-                        break
+                        param = params_dict[name]
-                    param = params_dict[name]
+                        weight_loader = param.weight_loader
-                    weight_loader = param.weight_loader
+                        weight_loader(
-                    weight_loader(
+                            param,
-                        param,
+                            loaded_weight,
-                        loaded_weight,
+                            name,
-                        name,
+                            expert_id=expert_id,
-                        expert_id=expert_id,
+                            shard_id=shard_id,
-                        shard_id=shard_id,
+                        )
-                    )
-                    break
                 else:
                     if (
                         name.endswith(".bias")
@@ -3483,8 +3413,7 @@ class KimiK3LinearForCausalLM(nn.Module):
     def post_load_weights(self):
         # Also invoked by loader post-load hooks (DummyModelLoader,
         # ShardedStateLoader, remote-instance flows -- none of which call
-        # load_weights), so e.g. dummy-weight benchmarks get w_kc/w_vc and
+        # load_weights), so e.g. dummy-weight benchmarks get the fused buffers.
-        # the fused buffers too. Same pattern as deepseek_v4.
         # Post-load: absorb kv_b_proj into w_kc and w_vc for MLA layers
         for layer_id in self.config.full_attention_layer_ids:
             if layer_id >= len(self.model.layers):
@@ -3511,8 +3440,8 @@ class KimiK3LinearForCausalLM(nn.Module):
 
         # Post-load: precompute the attn-res combined score weights BEFORE
         # cuda graph capture (a lazy first call inside get_cw would bake the
-        # multiply into every captured graph replay otherwise). Warm both
+        # multiply into every replay). Warm both dtypes (bf16 fast kernel,
-        # dtypes: the fast kernel consumes bf16, the triton fallback fp32.
+        # fp32 triton fallback).
         def _warm_cw(proj, norm):
             get_cw(proj, norm, dtype=torch.bfloat16)
             get_cw(proj, norm)
@@ -3526,19 +3455,16 @@ class KimiK3LinearForCausalLM(nn.Module):
         if hasattr(self.model, "output_attn_res_proj"):
             _warm_cw(self.model.output_attn_res_proj, self.model.output_attn_res_norm)
 
-        # Post-load: merge the horizontally-fused decode weights. Module
+        # Post-load: merge the horizontally-fused decode weights (views of the
-        # weights are re-pointed to views of the merged buffers (net extra
+        # merged buffers, ~0 extra memory); must run before cuda graph capture.
-        # memory ~0), so this must run after all weights are loaded and
-        # before cuda graph capture.
         for layer in self.model.layers:
             if isinstance(layer, PPMissingLayer):
                 continue
             if isinstance(layer.mlp, KimiK3MoE):
                 layer.mlp._merge_front_weights()
-                # The router consumes the correction bias in fp32; convert the
+                # Convert the correction bias to fp32 once so the per-call
-                # bf16 checkpoint values once (exact) so the per-call
+                # .to(float32) in topk is a no-op, not one upcast kernel per
-                # .to(float32) in topk becomes a no-op instead of one upcast
+                # MoE layer per step.
-                # kernel per MoE layer per step.
                 bias = layer.mlp.gate.e_score_correction_bias
                 if bias.dtype != torch.float32:
                     bias.data = bias.data.to(torch.float32)