[spec decoding] support kimi-k2.6-eagle3.1-mla draft (#26506)

Co-authored-by: Cursor <cursoragent@cursor.com>

python/sglang/srt/models/kimi_k25_eagle3.py

@@ -1,10 +1,18 @@
-"""EAGLE3 draft model with MLA attention for Kimi-K2.5.
+"""EAGLE3 / EAGLE3.1 draft model with MLA attention for Kimi-K2.x.
 
 The ``kimi-k2.5-eagle3-mla`` checkpoint pairs an EAGLE3 layout
 (concatenated [embed_norm, hidden_norm] pre-attention input, fc projection
 over the concatenated multi-layer aux hidden states, single decoder layer,
 dense MLP) with DeepSeek-V2 multi-latent attention. Sharing the MLA layout
-with the Kimi-K2.5 target keeps the draft KV cache small.
+with the Kimi-K2.x target keeps the draft KV cache small.
+
+The eagle3.1 variant (e.g. ``kimi-k2.6-eagle3.1-mla``) adds two optional
+config flags on top of the same layout:
+
+* ``fc_norm``: per-chunk RMSNorm applied to each auxiliary hidden state
+  before the fc projection.
+* ``norm_output``: emit post-norm (rather than pre-norm) hidden states as
+  the auxiliary output consumed by the next draft step.
 """
 
 import copy
@@ -196,13 +204,28 @@ class Eagle3MLAModel(nn.Module):
         target_hidden_size = (
             getattr(config, "target_hidden_size", None) or config.hidden_size
         )
-        num_fc_input = _get_eagle_aux_layer_count(config)
+        self.num_aux_hidden_states = _get_eagle_aux_layer_count(config)
         self.fc = nn.Linear(
-            target_hidden_size * num_fc_input,
+            target_hidden_size * self.num_aux_hidden_states,
             config.hidden_size,
             bias=getattr(config, "bias", False),
         )
 
+        # Per-aux RMSNorm before fc; enabled via `fc_norm` or legacy
+        # `use_aux_norm` flag. Matches the eagle3.1 layout.
+        use_fc_norm = getattr(config, "fc_norm", None) or getattr(
+            config, "use_aux_norm", False
+        )
+        if use_fc_norm:
+            self.fc_norm = nn.ModuleList(
+                [
+                    RMSNorm(target_hidden_size, eps=config.rms_norm_eps)
+                    for _ in range(self.num_aux_hidden_states)
+                ]
+            )
+        else:
+            self.fc_norm = None
+
         if config.num_hidden_layers != 1:
             raise ValueError("EAGLE3 currently only supports 1 layer")
         self.midlayer = Eagle3MLADecoderLayer(
@@ -213,6 +236,9 @@ class Eagle3MLAModel(nn.Module):
         )
 
         self.norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+        # Draft decode captures pre-norm hidden by default; eagle3.1 opts for
+        # post-norm via `norm_output: true`.
+        self.norm_output = getattr(config, "norm_output", False)
 
     def forward(
         self,
@@ -243,6 +269,12 @@ class Eagle3MLAModel(nn.Module):
 
         hidden_states = forward_batch.spec_info.hidden_states
         if hidden_states.shape[-1] != embeds.shape[-1]:
+            if self.fc_norm is not None:
+                chunks = hidden_states.chunk(self.num_aux_hidden_states, dim=-1)
+                hidden_states = torch.cat(
+                    [norm(chunk) for norm, chunk in zip(self.fc_norm, chunks)],
+                    dim=-1,
+                )
             hidden_states = self.fc(hidden_states)
 
         if hidden_states.shape[0] == 0:
@@ -260,7 +292,8 @@ class Eagle3MLAModel(nn.Module):
         hidden_states_to_logits, hidden_states_to_aux = self.norm(
             hidden_states, residual
         )
-        return hidden_states_to_logits, [hidden_states_to_aux]
+        aux = hidden_states_to_logits if self.norm_output else hidden_states_to_aux
+        return hidden_states_to_logits, [aux]
 
 
 class Eagle3DeepseekV2ForCausalLM(nn.Module):