From f3d467e2d246139422d9de5f73bffdc64ce3ca54 Mon Sep 17 00:00:00 2001
From: rjzhb <rjzhb222@163.com>
Date: Fri, 22 May 2026 20:50:52 +0000
Subject: [PATCH 01/18] perf(eagle3): skip dead-position compute in draft
 catch-up step

Signed-off-by: rjzhb <rjzhb222@163.com>
---
 .../tokenspeed/runtime/execution/context.py   |  5 ++
 .../runtime/execution/drafter/eagle.py        | 17 +++++++
 .../runtime/layers/logits_processor.py        |  6 ++-
 .../tokenspeed/runtime/models/llama_eagle3.py | 47 +++++++++++++++----
 4 files changed, 66 insertions(+), 9 deletions(-)

diff --git a/python/tokenspeed/runtime/execution/context.py b/python/tokenspeed/runtime/execution/context.py
index e5cb59f39..1c3af8dac 100644
--- a/python/tokenspeed/runtime/execution/context.py
+++ b/python/tokenspeed/runtime/execution/context.py
@@ -58,3 +58,8 @@ class ForwardContext:
 
     # --- logits processor ---
     gather_ids: torch.Tensor | None = None
+
+    # When True, the draft head's first step prunes attn/MLP/post-norms to the
+    # one live position per request using ctx.gather_ids. LogitsProcessor
+    # bypasses its own slicing (gather_ids set to None on LogitsMetadata).
+    draft_reduce_to_last: bool = False
diff --git a/python/tokenspeed/runtime/execution/drafter/eagle.py b/python/tokenspeed/runtime/execution/drafter/eagle.py
index 1e15cb581..dcb31975c 100644
--- a/python/tokenspeed/runtime/execution/drafter/eagle.py
+++ b/python/tokenspeed/runtime/execution/drafter/eagle.py
@@ -20,6 +20,7 @@
 
 from __future__ import annotations
 
+import os
 from dataclasses import dataclass
 from typing import TYPE_CHECKING
 
@@ -41,6 +42,13 @@
 
 logger = get_colorful_logger(__name__)
 
+# Env-var gate for the EAGLE draft first-step reduce optimization. Temporary —
+# set to 0 to disable the optimization and run the baseline catch-up path for
+# A/B perf comparison. Will be removed once stabilized.
+_DRAFT_REDUCE_FIRST_STEP_ENABLED = (
+    os.environ.get("TOKENSPEED_EAGLE_DRAFT_REDUCE", "1") == "1"
+)
+
 if TYPE_CHECKING:
     from tokenspeed.runtime.execution.input_buffer import InputBuffers
     from tokenspeed.runtime.execution.model_runner import ModelRunner
@@ -212,6 +220,14 @@ def _run_first_step(
             draft_input, bs, draft_input.input_num_tokens
         )
 
+        # Catch-up step with multi-token decode input: every position except
+        # the live one per request is purely a KV-cache write. The midlayer
+        # slices Q after KV write via ctx.gather_ids and runs attn/MLP/post-
+        # norms on just the [bs] live positions.
+        draft_reduce_to_last = (
+            forward_mode.is_decode() and _DRAFT_REDUCE_FIRST_STEP_ENABLED
+        )
+
         ctx = ForwardContext(
             attn_backend=self.attn_backend,
             token_to_kv_pool=self.token_to_kv_pool,
@@ -225,6 +241,7 @@ def _run_first_step(
             global_num_tokens=draft_input.global_num_tokens,
             global_bs=draft_input.global_bs,
             all_decode_or_idle=draft_input.all_decode_or_idle,
+            draft_reduce_to_last=draft_reduce_to_last,
         )
 
         return self.draft_model_runner.forward(
diff --git a/python/tokenspeed/runtime/layers/logits_processor.py b/python/tokenspeed/runtime/layers/logits_processor.py
index 0b7a26865..cfe26ed97 100755
--- a/python/tokenspeed/runtime/layers/logits_processor.py
+++ b/python/tokenspeed/runtime/layers/logits_processor.py
@@ -111,10 +111,14 @@ def from_forward_context(
         ctx: ForwardContext,
         input_lengths: torch.Tensor,
     ):
+        # When the midlayer already pruned to one row per request (EAGLE draft
+        # first-step reduce), drop gather_ids so LogitsProcessor passes through
+        # instead of re-indexing.
+        gather_ids = None if ctx.draft_reduce_to_last else ctx.gather_ids
         return cls(
             forward_mode=ctx.forward_mode,
             capture_hidden_mode=ctx.capture_hidden_mode,
-            gather_ids=ctx.gather_ids,
+            gather_ids=gather_ids,
             extend_seq_lens=input_lengths,
         )
 
diff --git a/python/tokenspeed/runtime/models/llama_eagle3.py b/python/tokenspeed/runtime/models/llama_eagle3.py
index 8e1e9ab13..ab5f22f25 100644
--- a/python/tokenspeed/runtime/models/llama_eagle3.py
+++ b/python/tokenspeed/runtime/models/llama_eagle3.py
@@ -35,6 +35,7 @@
 from tokenspeed.runtime.configs.utils import get_rope_theta
 from tokenspeed.runtime.distributed.mapping import Mapping
 from tokenspeed.runtime.execution.context import ForwardContext
+from tokenspeed.runtime.execution.forward_batch_info import ForwardMode
 from tokenspeed.runtime.layers.activation import SiluAndMul
 from tokenspeed.runtime.layers.common import concat
 from tokenspeed.runtime.layers.layernorm import RMSNorm
@@ -185,14 +186,34 @@ def forward(
                 output_q_rope=q_rope,
                 enable_pdl=pdl_enabled(),
             )
-            attn_output = self.attn(
-                q_rope,
-                None,
-                None,
-                save_kv_cache=False,
-                ctx=ctx,
-                out_cache_loc=out_cache_loc,
-            )
+            if ctx.draft_reduce_to_last:
+                # KV cache for all bs*spec_num_tokens positions already written
+                # by fused_set_kv_buffer_arg above. Reduce Q to one query per
+                # request (live position) and route to decode kernel: bs queries
+                # against the full cache (which now includes the just-written
+                # spec_num_tokens new tokens per request). DRAFT_EXTEND init
+                # populates forward_decode_metadata precisely for this case.
+                q_rope = q_rope.index_select(0, ctx.gather_ids)
+                attn_output = ctx.attn_backend.forward(
+                    q_rope,
+                    None,
+                    None,
+                    self.attn,
+                    out_cache_loc,
+                    ctx.token_to_kv_pool,
+                    ForwardMode.DECODE,
+                    ctx.bs,
+                    save_kv_cache=False,
+                )
+            else:
+                attn_output = self.attn(
+                    q_rope,
+                    None,
+                    None,
+                    save_kv_cache=False,
+                    ctx=ctx,
+                    out_cache_loc=out_cache_loc,
+                )
         else:
             q, k = self.rotary_emb(positions, q, k)
             attn_output = self.attn(q, k, v, ctx=ctx, out_cache_loc=out_cache_loc)
@@ -361,6 +382,11 @@ def forward_low_latency(
             ctx=ctx,
             out_cache_loc=out_cache_loc,
         )
+        if ctx.draft_reduce_to_last:
+            # self_attn returned [bs, H]; gather residual to match before the
+            # fused allreduce+norm so shapes line up. Everything downstream
+            # (post-norm, MLP, final norm) now runs on [bs, H].
+            residual = residual.index_select(0, ctx.gather_ids)
 
         # Fused post-attn allreduce + norm (uses attn tp group)
         block_scale = None
@@ -426,6 +452,11 @@ def forward(
             ctx=ctx,
             out_cache_loc=out_cache_loc,
         )
+        if ctx.draft_reduce_to_last:
+            # self_attn returned [bs, H]; gather residual to match before the
+            # post-attn norm+comm so shapes line up. Everything downstream
+            # (post-norm, MLP, final norm) now runs on [bs, H].
+            residual = residual.index_select(0, ctx.gather_ids)
         hidden_states, residual = self.comm_manager.post_attn_comm(
             hidden_states, residual, ctx
         )

From e911f18479b00c6c732f842711e94d49b83faf7d Mon Sep 17 00:00:00 2001
From: rjzhb <rjzhb222@163.com>
Date: Mon, 25 May 2026 22:51:37 +0000
Subject: [PATCH 02/18] fix(eagle3): slice attn_output on non-prewrite attn
 path

Signed-off-by: rjzhb <rjzhb222@163.com>
---
 python/tokenspeed/runtime/models/llama_eagle3.py | 5 +++++
 1 file changed, 5 insertions(+)

diff --git a/python/tokenspeed/runtime/models/llama_eagle3.py b/python/tokenspeed/runtime/models/llama_eagle3.py
index ab5f22f25..bcfcb609b 100644
--- a/python/tokenspeed/runtime/models/llama_eagle3.py
+++ b/python/tokenspeed/runtime/models/llama_eagle3.py
@@ -217,6 +217,11 @@ def forward(
         else:
             q, k = self.rotary_emb(positions, q, k)
             attn_output = self.attn(q, k, v, ctx=ctx, out_cache_loc=out_cache_loc)
+            if ctx.draft_reduce_to_last:
+                # Non-prewrite backend: KV was written by self.attn above;
+                # slice attn_output to one row per request so o_proj and the
+                # rest of the layer only see the live positions.
+                attn_output = attn_output.index_select(0, ctx.gather_ids)
 
         output, _ = self.o_proj(attn_output)
         return output

From a0df4e37ede039cc7f86cc91b049b291f14526c0 Mon Sep 17 00:00:00 2001
From: rjzhb <rjzhb222@163.com>
Date: Mon, 25 May 2026 22:51:37 +0000
Subject: [PATCH 03/18] fix(eagle3): restrict draft-reduce to llama eagle3 head
 class

Signed-off-by: rjzhb <rjzhb222@163.com>
---
 .../runtime/execution/drafter/eagle.py          | 17 ++++++-----------
 1 file changed, 6 insertions(+), 11 deletions(-)

diff --git a/python/tokenspeed/runtime/execution/drafter/eagle.py b/python/tokenspeed/runtime/execution/drafter/eagle.py
index dcb31975c..2a291c27e 100644
--- a/python/tokenspeed/runtime/execution/drafter/eagle.py
+++ b/python/tokenspeed/runtime/execution/drafter/eagle.py
@@ -20,7 +20,6 @@
 
 from __future__ import annotations
 
-import os
 from dataclasses import dataclass
 from typing import TYPE_CHECKING
 
@@ -37,18 +36,12 @@
     CaptureHiddenMode,
     ForwardMode,
 )
+from tokenspeed.runtime.models.llama_eagle3 import LlamaForCausalLMEagle3
 from tokenspeed.runtime.utils import get_colorful_logger
 from tokenspeed.runtime.utils.nvtx import nvtx_range
 
 logger = get_colorful_logger(__name__)
 
-# Env-var gate for the EAGLE draft first-step reduce optimization. Temporary —
-# set to 0 to disable the optimization and run the baseline catch-up path for
-# A/B perf comparison. Will be removed once stabilized.
-_DRAFT_REDUCE_FIRST_STEP_ENABLED = (
-    os.environ.get("TOKENSPEED_EAGLE_DRAFT_REDUCE", "1") == "1"
-)
-
 if TYPE_CHECKING:
     from tokenspeed.runtime.execution.input_buffer import InputBuffers
     from tokenspeed.runtime.execution.model_runner import ModelRunner
@@ -223,9 +216,11 @@ def _run_first_step(
         # Catch-up step with multi-token decode input: every position except
         # the live one per request is purely a KV-cache write. The midlayer
         # slices Q after KV write via ctx.gather_ids and runs attn/MLP/post-
-        # norms on just the [bs] live positions.
-        draft_reduce_to_last = (
-            forward_mode.is_decode() and _DRAFT_REDUCE_FIRST_STEP_ENABLED
+        # norms on just the [bs] live positions. Only LlamaForCausalLMEagle3
+        # implements the midlayer slice today; MLA / NextN drafts fall back
+        # to the full path until a follow-up PR adds matching slice code.
+        draft_reduce_to_last = forward_mode.is_decode() and isinstance(
+            self.draft_model_runner.model, LlamaForCausalLMEagle3
         )
 
         ctx = ForwardContext(

From fae72911fd35384c5845d9c7760c6ff575cb501b Mon Sep 17 00:00:00 2001
From: rjzhb <rjzhb222@163.com>
Date: Mon, 25 May 2026 22:57:28 +0000
Subject: [PATCH 04/18] chore(eagle3): trim comments on draft-reduce path

Signed-off-by: rjzhb <rjzhb222@163.com>
---
 .../runtime/execution/drafter/eagle.py        |  8 ++------
 .../tokenspeed/runtime/models/llama_eagle3.py | 19 ++++++-------------
 2 files changed, 8 insertions(+), 19 deletions(-)

diff --git a/python/tokenspeed/runtime/execution/drafter/eagle.py b/python/tokenspeed/runtime/execution/drafter/eagle.py
index 2a291c27e..ecc4a1fbc 100644
--- a/python/tokenspeed/runtime/execution/drafter/eagle.py
+++ b/python/tokenspeed/runtime/execution/drafter/eagle.py
@@ -213,12 +213,8 @@ def _run_first_step(
             draft_input, bs, draft_input.input_num_tokens
         )
 
-        # Catch-up step with multi-token decode input: every position except
-        # the live one per request is purely a KV-cache write. The midlayer
-        # slices Q after KV write via ctx.gather_ids and runs attn/MLP/post-
-        # norms on just the [bs] live positions. Only LlamaForCausalLMEagle3
-        # implements the midlayer slice today; MLA / NextN drafts fall back
-        # to the full path until a follow-up PR adds matching slice code.
+        # Only LlamaForCausalLMEagle3 implements the midlayer slice today; MLA
+        # / NextN draft heads fall back to the standard path.
         draft_reduce_to_last = forward_mode.is_decode() and isinstance(
             self.draft_model_runner.model, LlamaForCausalLMEagle3
         )
diff --git a/python/tokenspeed/runtime/models/llama_eagle3.py b/python/tokenspeed/runtime/models/llama_eagle3.py
index bcfcb609b..87f435031 100644
--- a/python/tokenspeed/runtime/models/llama_eagle3.py
+++ b/python/tokenspeed/runtime/models/llama_eagle3.py
@@ -187,12 +187,8 @@ def forward(
                 enable_pdl=pdl_enabled(),
             )
             if ctx.draft_reduce_to_last:
-                # KV cache for all bs*spec_num_tokens positions already written
-                # by fused_set_kv_buffer_arg above. Reduce Q to one query per
-                # request (live position) and route to decode kernel: bs queries
-                # against the full cache (which now includes the just-written
-                # spec_num_tokens new tokens per request). DRAFT_EXTEND init
-                # populates forward_decode_metadata precisely for this case.
+                # KV already written via fused_set_kv_buffer_arg above; slice Q
+                # to one query per request and route attn as decode.
                 q_rope = q_rope.index_select(0, ctx.gather_ids)
                 attn_output = ctx.attn_backend.forward(
                     q_rope,
@@ -218,9 +214,8 @@ def forward(
             q, k = self.rotary_emb(positions, q, k)
             attn_output = self.attn(q, k, v, ctx=ctx, out_cache_loc=out_cache_loc)
             if ctx.draft_reduce_to_last:
-                # Non-prewrite backend: KV was written by self.attn above;
-                # slice attn_output to one row per request so o_proj and the
-                # rest of the layer only see the live positions.
+                # KV written by self.attn above; slice attn_output so o_proj
+                # and the rest of the layer only run on the live rows.
                 attn_output = attn_output.index_select(0, ctx.gather_ids)
 
         output, _ = self.o_proj(attn_output)
@@ -389,8 +384,7 @@ def forward_low_latency(
         )
         if ctx.draft_reduce_to_last:
             # self_attn returned [bs, H]; gather residual to match before the
-            # fused allreduce+norm so shapes line up. Everything downstream
-            # (post-norm, MLP, final norm) now runs on [bs, H].
+            # fused allreduce+norm.
             residual = residual.index_select(0, ctx.gather_ids)
 
         # Fused post-attn allreduce + norm (uses attn tp group)
@@ -459,8 +453,7 @@ def forward(
         )
         if ctx.draft_reduce_to_last:
             # self_attn returned [bs, H]; gather residual to match before the
-            # post-attn norm+comm so shapes line up. Everything downstream
-            # (post-norm, MLP, final norm) now runs on [bs, H].
+            # post-attn norm+comm.
             residual = residual.index_select(0, ctx.gather_ids)
         hidden_states, residual = self.comm_manager.post_attn_comm(
             hidden_states, residual, ctx

From c642f45df6c4b78b9a9613e836ee4d1d5215c9cd Mon Sep 17 00:00:00 2001
From: rjzhb <rjzhb222@163.com>
Date: Tue, 26 May 2026 00:52:48 +0000
Subject: [PATCH 05/18] fix(eagle3): trim draft first-step seq_lens to accept
 length

Signed-off-by: rjzhb <rjzhb222@163.com>
---
 python/tokenspeed/runtime/execution/drafter/eagle.py | 11 +++++++++++
 1 file changed, 11 insertions(+)

diff --git a/python/tokenspeed/runtime/execution/drafter/eagle.py b/python/tokenspeed/runtime/execution/drafter/eagle.py
index ecc4a1fbc..708488c09 100644
--- a/python/tokenspeed/runtime/execution/drafter/eagle.py
+++ b/python/tokenspeed/runtime/execution/drafter/eagle.py
@@ -219,6 +219,17 @@ def _run_first_step(
             self.draft_model_runner.model, LlamaForCausalLMEagle3
         )
 
+        if draft_reduce_to_last:
+            # draft_seq_lens_buf aliases the backend's cache_seqlens_int32 at
+            # valid_cache_len + spec_num_tokens. After Q is sliced to one row
+            # per request, the decode kernel takes seq_lens as the attention
+            # upper bound, so trim by the rejected-draft count to keep the live
+            # query from reading them.
+            correction = (self.spec_num_tokens - draft_input.accept_lengths).to(
+                self.draft_seq_lens_buf.dtype
+            )
+            self.draft_seq_lens_buf[:bs].sub_(correction)
+
         ctx = ForwardContext(
             attn_backend=self.attn_backend,
             token_to_kv_pool=self.token_to_kv_pool,

From b64b835a584a2bf29d4d6ee9ac2cd47667688a53 Mon Sep 17 00:00:00 2001
From: rjzhb <rjzhb222@163.com>
Date: Tue, 26 May 2026 00:52:53 +0000
Subject: [PATCH 06/18] fix(comm): scatter on bs / global_bs under draft
 first-step reduce

Signed-off-by: rjzhb <rjzhb222@163.com>
---
 .../tokenspeed/runtime/distributed/comm_manager.py | 14 +++++++++-----
 1 file changed, 9 insertions(+), 5 deletions(-)

diff --git a/python/tokenspeed/runtime/distributed/comm_manager.py b/python/tokenspeed/runtime/distributed/comm_manager.py
index 134b63a27..3854357cb 100755
--- a/python/tokenspeed/runtime/distributed/comm_manager.py
+++ b/python/tokenspeed/runtime/distributed/comm_manager.py
@@ -60,17 +60,21 @@ def get_num_tokens(self, ctx: ForwardContext):
         return sum(scattered), max(scattered)
 
     def scattered_num_tokens(self, ctx: ForwardContext) -> list[int]:
-        if ctx.global_num_tokens is not None:
+        # Under draft first-step reduce, comm operates on bs / global_bs since
+        # the midlayer pruned activations to one row per request.
+        global_counts = (
+            ctx.global_bs if ctx.draft_reduce_to_last else ctx.global_num_tokens
+        )
+        if global_counts is not None:
             scattered = []
             for attn_dp_rank in range(self.mapping.attn.dp_size):
-                num_tokens = ctx.global_num_tokens[
-                    attn_dp_rank * self.mapping.attn.tp_size
-                ]
+                num_tokens = global_counts[attn_dp_rank * self.mapping.attn.tp_size]
                 scattered.extend(
                     self._scatter_count(num_tokens, self.mapping.attn.tp_size)
                 )
             return scattered
-        return self._scatter_count(ctx.input_num_tokens, self.mapping.attn.tp_size)
+        num_tokens = ctx.bs if ctx.draft_reduce_to_last else ctx.input_num_tokens
+        return self._scatter_count(num_tokens, self.mapping.attn.tp_size)
 
     def attn_tp_group_scattered_num_tokens(self, ctx: ForwardContext) -> list[int]:
         start = self.mapping.attn.tp_size * self.mapping.attn.dp_rank

From d151cfb829d3ba30c7cc8909add8607032ed7954 Mon Sep 17 00:00:00 2001
From: rjzhb <rjzhb222@163.com>
Date: Tue, 26 May 2026 01:35:44 +0000
Subject: [PATCH 07/18] fix(eagle3): gate draft seq_lens correction to prewrite
 backends

Signed-off-by: rjzhb <rjzhb222@163.com>
---
 .../tokenspeed/runtime/execution/drafter/eagle.py   | 13 +++++++------
 1 file changed, 7 insertions(+), 6 deletions(-)

diff --git a/python/tokenspeed/runtime/execution/drafter/eagle.py b/python/tokenspeed/runtime/execution/drafter/eagle.py
index 708488c09..a09442a92 100644
--- a/python/tokenspeed/runtime/execution/drafter/eagle.py
+++ b/python/tokenspeed/runtime/execution/drafter/eagle.py
@@ -219,12 +219,13 @@ def _run_first_step(
             self.draft_model_runner.model, LlamaForCausalLMEagle3
         )
 
-        if draft_reduce_to_last:
-            # draft_seq_lens_buf aliases the backend's cache_seqlens_int32 at
-            # valid_cache_len + spec_num_tokens. After Q is sliced to one row
-            # per request, the decode kernel takes seq_lens as the attention
-            # upper bound, so trim by the rejected-draft count to keep the live
-            # query from reading them.
+        if draft_reduce_to_last and self.attn_backend.support_kv_cache_prewrite:
+            # On the prewrite path, Q is sliced to one row per request before
+            # the decode kernel, which takes seq_lens as the attention upper
+            # bound. Trim by the rejected-draft count so the live query does
+            # not read the dead positions. The non-prewrite path runs full
+            # multi-token attention with per-position causal masking, so
+            # seq_lens must stay at valid_cache_len + spec_num_tokens there.
             correction = (self.spec_num_tokens - draft_input.accept_lengths).to(
                 self.draft_seq_lens_buf.dtype
             )

From e5ae5fa4afc0f5cd30cc41f1157a61e4ece0824b Mon Sep 17 00:00:00 2001
From: rjzhb <rjzhb222@163.com>
Date: Tue, 26 May 2026 05:11:06 +0000
Subject: [PATCH 08/18] refactor(eagle3): duck-type draft first-step reduce
 capability

Signed-off-by: rjzhb <rjzhb222@163.com>
---
 .../runtime/distributed/comm_manager.py       |  4 ++--
 .../tokenspeed/runtime/execution/context.py   |  7 ++----
 .../runtime/execution/drafter/eagle.py        | 24 +++++++++----------
 .../runtime/layers/logits_processor.py        | 10 ++++----
 .../tokenspeed/runtime/models/llama_eagle3.py |  9 +++----
 5 files changed, 24 insertions(+), 30 deletions(-)

diff --git a/python/tokenspeed/runtime/distributed/comm_manager.py b/python/tokenspeed/runtime/distributed/comm_manager.py
index 3854357cb..a4e52a839 100755
--- a/python/tokenspeed/runtime/distributed/comm_manager.py
+++ b/python/tokenspeed/runtime/distributed/comm_manager.py
@@ -63,7 +63,7 @@ def scattered_num_tokens(self, ctx: ForwardContext) -> list[int]:
         # Under draft first-step reduce, comm operates on bs / global_bs since
         # the midlayer pruned activations to one row per request.
         global_counts = (
-            ctx.global_bs if ctx.draft_reduce_to_last else ctx.global_num_tokens
+            ctx.global_bs if ctx.draft_first_step_reduce else ctx.global_num_tokens
         )
         if global_counts is not None:
             scattered = []
@@ -73,7 +73,7 @@ def scattered_num_tokens(self, ctx: ForwardContext) -> list[int]:
                     self._scatter_count(num_tokens, self.mapping.attn.tp_size)
                 )
             return scattered
-        num_tokens = ctx.bs if ctx.draft_reduce_to_last else ctx.input_num_tokens
+        num_tokens = ctx.bs if ctx.draft_first_step_reduce else ctx.input_num_tokens
         return self._scatter_count(num_tokens, self.mapping.attn.tp_size)
 
     def attn_tp_group_scattered_num_tokens(self, ctx: ForwardContext) -> list[int]:
diff --git a/python/tokenspeed/runtime/execution/context.py b/python/tokenspeed/runtime/execution/context.py
index 1c3af8dac..80609de8c 100644
--- a/python/tokenspeed/runtime/execution/context.py
+++ b/python/tokenspeed/runtime/execution/context.py
@@ -50,6 +50,8 @@ class ForwardContext:
     forward_mode: ForwardMode | None
     req_to_page: torch.Tensor | None = None
     capture_hidden_mode: CaptureHiddenMode | None = CaptureHiddenMode.NULL
+    # EAGLE draft head's first step prunes to one live row per request.
+    draft_first_step_reduce: bool = False
 
     # --- dp attention ---
     global_num_tokens: list[int] | None = None
@@ -58,8 +60,3 @@ class ForwardContext:
 
     # --- logits processor ---
     gather_ids: torch.Tensor | None = None
-
-    # When True, the draft head's first step prunes attn/MLP/post-norms to the
-    # one live position per request using ctx.gather_ids. LogitsProcessor
-    # bypasses its own slicing (gather_ids set to None on LogitsMetadata).
-    draft_reduce_to_last: bool = False
diff --git a/python/tokenspeed/runtime/execution/drafter/eagle.py b/python/tokenspeed/runtime/execution/drafter/eagle.py
index a09442a92..b6ac40eb5 100644
--- a/python/tokenspeed/runtime/execution/drafter/eagle.py
+++ b/python/tokenspeed/runtime/execution/drafter/eagle.py
@@ -36,7 +36,6 @@
     CaptureHiddenMode,
     ForwardMode,
 )
-from tokenspeed.runtime.models.llama_eagle3 import LlamaForCausalLMEagle3
 from tokenspeed.runtime.utils import get_colorful_logger
 from tokenspeed.runtime.utils.nvtx import nvtx_range
 
@@ -213,19 +212,18 @@ def _run_first_step(
             draft_input, bs, draft_input.input_num_tokens
         )
 
-        # Only LlamaForCausalLMEagle3 implements the midlayer slice today; MLA
-        # / NextN draft heads fall back to the standard path.
-        draft_reduce_to_last = forward_mode.is_decode() and isinstance(
-            self.draft_model_runner.model, LlamaForCausalLMEagle3
+        # Models declare ``supports_draft_first_step_reduce`` to opt in; MLA /
+        # NextN heads that don't declare it fall back to the standard path.
+        draft_first_step_reduce = forward_mode.is_decode() and getattr(
+            self.draft_model_runner.model, "supports_draft_first_step_reduce", False
         )
 
-        if draft_reduce_to_last and self.attn_backend.support_kv_cache_prewrite:
-            # On the prewrite path, Q is sliced to one row per request before
-            # the decode kernel, which takes seq_lens as the attention upper
-            # bound. Trim by the rejected-draft count so the live query does
-            # not read the dead positions. The non-prewrite path runs full
-            # multi-token attention with per-position causal masking, so
-            # seq_lens must stay at valid_cache_len + spec_num_tokens there.
+        if draft_first_step_reduce and self.attn_backend.support_kv_cache_prewrite:
+            # Prewrite path slices Q to one row per request and uses the decode
+            # kernel, which reads seq_lens as the attention upper bound. Trim
+            # by the rejected-draft count so the live query does not read dead
+            # positions. Non-prewrite runs full multi-token attention with
+            # per-position causal masking, so seq_lens must stay unchanged.
             correction = (self.spec_num_tokens - draft_input.accept_lengths).to(
                 self.draft_seq_lens_buf.dtype
             )
@@ -244,7 +242,7 @@ def _run_first_step(
             global_num_tokens=draft_input.global_num_tokens,
             global_bs=draft_input.global_bs,
             all_decode_or_idle=draft_input.all_decode_or_idle,
-            draft_reduce_to_last=draft_reduce_to_last,
+            draft_first_step_reduce=draft_first_step_reduce,
         )
 
         return self.draft_model_runner.forward(
diff --git a/python/tokenspeed/runtime/layers/logits_processor.py b/python/tokenspeed/runtime/layers/logits_processor.py
index ed3ae9b07..161e0b987 100755
--- a/python/tokenspeed/runtime/layers/logits_processor.py
+++ b/python/tokenspeed/runtime/layers/logits_processor.py
@@ -111,14 +111,10 @@ def from_forward_context(
         ctx: ForwardContext,
         input_lengths: torch.Tensor,
     ):
-        # When the midlayer already pruned to one row per request (EAGLE draft
-        # first-step reduce), drop gather_ids so LogitsProcessor passes through
-        # instead of re-indexing.
-        gather_ids = None if ctx.draft_reduce_to_last else ctx.gather_ids
         return cls(
             forward_mode=ctx.forward_mode,
             capture_hidden_mode=ctx.capture_hidden_mode,
-            gather_ids=gather_ids,
+            gather_ids=ctx.gather_ids,
             extend_seq_lens=input_lengths,
         )
 
@@ -214,7 +210,9 @@ def forward(
         # Get the last hidden states and last logits for the next token prediction
         if not logits_metadata.extend_return_logprob:
             gather_ids = logits_metadata.gather_ids
-            if gather_ids is None:
+            # Shapes align iff midlayer already pruned to one row per request
+            # (EAGLE draft first-step reduce). Other paths emit [N, H] with N > bs.
+            if gather_ids is None or gather_ids.shape[0] == hidden_states.shape[0]:
                 pruned_states = hidden_states
                 if aux_hidden_states is not None:
                     aux_pruned_states = list(aux_hidden_states)
diff --git a/python/tokenspeed/runtime/models/llama_eagle3.py b/python/tokenspeed/runtime/models/llama_eagle3.py
index 87f435031..c5a976982 100644
--- a/python/tokenspeed/runtime/models/llama_eagle3.py
+++ b/python/tokenspeed/runtime/models/llama_eagle3.py
@@ -186,7 +186,7 @@ def forward(
                 output_q_rope=q_rope,
                 enable_pdl=pdl_enabled(),
             )
-            if ctx.draft_reduce_to_last:
+            if ctx.draft_first_step_reduce:
                 # KV already written via fused_set_kv_buffer_arg above; slice Q
                 # to one query per request and route attn as decode.
                 q_rope = q_rope.index_select(0, ctx.gather_ids)
@@ -213,7 +213,7 @@ def forward(
         else:
             q, k = self.rotary_emb(positions, q, k)
             attn_output = self.attn(q, k, v, ctx=ctx, out_cache_loc=out_cache_loc)
-            if ctx.draft_reduce_to_last:
+            if ctx.draft_first_step_reduce:
                 # KV written by self.attn above; slice attn_output so o_proj
                 # and the rest of the layer only run on the live rows.
                 attn_output = attn_output.index_select(0, ctx.gather_ids)
@@ -382,7 +382,7 @@ def forward_low_latency(
             ctx=ctx,
             out_cache_loc=out_cache_loc,
         )
-        if ctx.draft_reduce_to_last:
+        if ctx.draft_first_step_reduce:
             # self_attn returned [bs, H]; gather residual to match before the
             # fused allreduce+norm.
             residual = residual.index_select(0, ctx.gather_ids)
@@ -451,7 +451,7 @@ def forward(
             ctx=ctx,
             out_cache_loc=out_cache_loc,
         )
-        if ctx.draft_reduce_to_last:
+        if ctx.draft_first_step_reduce:
             # self_attn returned [bs, H]; gather residual to match before the
             # post-attn norm+comm.
             residual = residual.index_select(0, ctx.gather_ids)
@@ -572,6 +572,7 @@ def forward(
 class LlamaForCausalLMEagle3(BaseCausalLM):
 
     model_cls = Eagle3LlamaModel
+    supports_draft_first_step_reduce = True
 
     def __init__(
         self,

From e242d39a684ac0d43128eac14c997e737645c3e7 Mon Sep 17 00:00:00 2001
From: rjzhb <rjzhb222@163.com>
Date: Tue, 26 May 2026 05:11:11 +0000
Subject: [PATCH 09/18] feat(eagle3): support draft first-step reduce for MLA
 head

Signed-off-by: rjzhb <rjzhb222@163.com>
---
 python/tokenspeed/runtime/models/deepseek_v3.py | 10 ++++++++++
 1 file changed, 10 insertions(+)

diff --git a/python/tokenspeed/runtime/models/deepseek_v3.py b/python/tokenspeed/runtime/models/deepseek_v3.py
index 9c8b5e634..1ff6c742c 100644
--- a/python/tokenspeed/runtime/models/deepseek_v3.py
+++ b/python/tokenspeed/runtime/models/deepseek_v3.py
@@ -675,6 +675,10 @@ def forward(
                 attn_output[num_prefill_tokens:],
             )
 
+        if ctx.draft_first_step_reduce:
+            # KV already written; drop dead-position rows so o_proj / MLP /
+            # post-norms only run on one live row per request.
+            attn_output = attn_output.index_select(0, ctx.gather_ids)
         output, _ = self.o_proj(attn_output)
         return output
 
@@ -1697,6 +1701,10 @@ def forward(
                 comm_manager=self.comm_manager,
             )
 
+            if ctx.draft_first_step_reduce:
+                # self_attn returned [bs, H]; gather residual to match before
+                # the fused reduce+norm.
+                residual = residual.index_select(0, ctx.gather_ids)
             hidden_states, residual = self.comm_manager.post_attn_reduce_norm(
                 hidden_states, residual, ctx
             )
@@ -1824,6 +1832,8 @@ class Eagle3DeepseekV2ForCausalLM(DeepseekV3ForCausalLM):
     concatenated [embeds || hidden_states] as input.
     """
 
+    supports_draft_first_step_reduce = True
+
     def __init__(
         self,
         config: PretrainedConfig,

From eecfc84c4011b595e928a428fe6af7a4d90f9e7d Mon Sep 17 00:00:00 2001
From: rjzhb <rjzhb222@163.com>
Date: Wed, 27 May 2026 01:21:58 +0000
Subject: [PATCH 10/18] fix(eagle3): mirror draft_first_step_reduce on idle DP
 ranks

Signed-off-by: rjzhb <rjzhb222@163.com>
---
 .../runtime/execution/model_executor.py       | 34 ++++++++++++-------
 1 file changed, 22 insertions(+), 12 deletions(-)

diff --git a/python/tokenspeed/runtime/execution/model_executor.py b/python/tokenspeed/runtime/execution/model_executor.py
index e583e8712..4cff78d5d 100644
--- a/python/tokenspeed/runtime/execution/model_executor.py
+++ b/python/tokenspeed/runtime/execution/model_executor.py
@@ -882,19 +882,29 @@ def execute_idle_forward(
         # NCCL collectives. Idle ranks must match those collectives:
         # 1 first-step forward + (spec_num_steps - 1) multi-step decode forwards.
         if self.drafter is not None:
-            draft_ctx = ForwardContext(
-                attn_backend=self.drafter.attn_backend,
-                token_to_kv_pool=self.drafter.token_to_kv_pool,
-                req_to_page=self.drafter.req_to_page,
-                bs=0,
-                num_extends=0,
-                input_num_tokens=0,
-                forward_mode=ForwardMode.IDLE,
-                global_num_tokens=global_num_tokens,
-                global_bs=global_bs,
-                all_decode_or_idle=all_decode_or_idle,
+            # Mirror active ranks: when all non-idle ranks are decoding and the
+            # draft model opts in, the catch-up step sizes collectives from
+            # bs/global_bs instead of bs*spec_num_tokens. Idle ranks must use
+            # the same basis or RSAG/AR shapes diverge.
+            supports_reduce = all_decode_or_idle and getattr(
+                self.drafter.draft_model_runner.model,
+                "supports_draft_first_step_reduce",
+                False,
             )
-            for _ in range(self.drafter.spec_num_steps):
+            for step_idx in range(self.drafter.spec_num_steps):
+                draft_ctx = ForwardContext(
+                    attn_backend=self.drafter.attn_backend,
+                    token_to_kv_pool=self.drafter.token_to_kv_pool,
+                    req_to_page=self.drafter.req_to_page,
+                    bs=0,
+                    num_extends=0,
+                    input_num_tokens=0,
+                    forward_mode=ForwardMode.IDLE,
+                    global_num_tokens=global_num_tokens,
+                    global_bs=global_bs,
+                    all_decode_or_idle=all_decode_or_idle,
+                    draft_first_step_reduce=(step_idx == 0 and supports_reduce),
+                )
                 self.drafter.draft_model_runner.forward(
                     draft_ctx,
                     input_ids=empty,

From 9e3c15a50c64b782b0fd14c58f5788cd6a01608d Mon Sep 17 00:00:00 2001
From: rjzhb <rjzhb222@163.com>
Date: Wed, 27 May 2026 05:16:38 +0000
Subject: [PATCH 11/18] fix(eagle3): guard llama Eagle3 reduce slice against
 idle DP ranks

Signed-off-by: rjzhb <rjzhb222@163.com>
---
 python/tokenspeed/runtime/models/llama_eagle3.py | 10 ++++------
 1 file changed, 4 insertions(+), 6 deletions(-)

diff --git a/python/tokenspeed/runtime/models/llama_eagle3.py b/python/tokenspeed/runtime/models/llama_eagle3.py
index c5a976982..b5534d83d 100644
--- a/python/tokenspeed/runtime/models/llama_eagle3.py
+++ b/python/tokenspeed/runtime/models/llama_eagle3.py
@@ -382,9 +382,8 @@ def forward_low_latency(
             ctx=ctx,
             out_cache_loc=out_cache_loc,
         )
-        if ctx.draft_first_step_reduce:
-            # self_attn returned [bs, H]; gather residual to match before the
-            # fused allreduce+norm.
+        if ctx.draft_first_step_reduce and not ctx.forward_mode.is_idle():
+            # Gather residual to self_attn's [bs, H]; idle has no gather_ids.
             residual = residual.index_select(0, ctx.gather_ids)
 
         # Fused post-attn allreduce + norm (uses attn tp group)
@@ -451,9 +450,8 @@ def forward(
             ctx=ctx,
             out_cache_loc=out_cache_loc,
         )
-        if ctx.draft_first_step_reduce:
-            # self_attn returned [bs, H]; gather residual to match before the
-            # post-attn norm+comm.
+        if ctx.draft_first_step_reduce and not ctx.forward_mode.is_idle():
+            # Gather residual to self_attn's [bs, H]; idle has no gather_ids.
             residual = residual.index_select(0, ctx.gather_ids)
         hidden_states, residual = self.comm_manager.post_attn_comm(
             hidden_states, residual, ctx

From 9973dbecc6160fc7f1ee7b2b6f5566eb01796e5d Mon Sep 17 00:00:00 2001
From: rjzhb <rjzhb222@163.com>
Date: Wed, 27 May 2026 05:17:31 +0000
Subject: [PATCH 12/18] feat(eagle3): support draft first-step reduce for
 DeepSeek NextN

Signed-off-by: rjzhb <rjzhb222@163.com>
---
 python/tokenspeed/runtime/models/deepseek_v3.py | 3 +++
 1 file changed, 3 insertions(+)

diff --git a/python/tokenspeed/runtime/models/deepseek_v3.py b/python/tokenspeed/runtime/models/deepseek_v3.py
index 1ff6c742c..24af5d376 100644
--- a/python/tokenspeed/runtime/models/deepseek_v3.py
+++ b/python/tokenspeed/runtime/models/deepseek_v3.py
@@ -1157,6 +1157,9 @@ def forward(
                 out_cache_loc=out_cache_loc,
                 comm_manager=self.comm_manager,
             )
+            if ctx.draft_first_step_reduce:
+                # Gather residual to self_attn's [bs, H].
+                residual = residual.index_select(0, ctx.gather_ids)
             hidden_states, residual = self.comm_manager.post_attn_reduce_norm(
                 hidden_states, residual, ctx
             )

From 4a2a2b7d1be5f26dd5af099c23206d4f9ddf1aca Mon Sep 17 00:00:00 2001
From: rjzhb <rjzhb222@163.com>
Date: Wed, 27 May 2026 05:18:21 +0000
Subject: [PATCH 13/18] feat(eagle3): support draft first-step reduce for
 Qwen3.5 MTP

Signed-off-by: rjzhb <rjzhb222@163.com>
---
 python/tokenspeed/runtime/models/qwen3_5.py | 7 +++++++
 1 file changed, 7 insertions(+)

diff --git a/python/tokenspeed/runtime/models/qwen3_5.py b/python/tokenspeed/runtime/models/qwen3_5.py
index dd94964d0..e3e333538 100644
--- a/python/tokenspeed/runtime/models/qwen3_5.py
+++ b/python/tokenspeed/runtime/models/qwen3_5.py
@@ -747,6 +747,10 @@ def self_attention(
         if self.attn_output_gate:
             sigmoid_mul(attn_output, gate)
 
+        if ctx.draft_first_step_reduce:
+            # Slice attn_output to [bs, H] so o_proj runs on live rows only.
+            attn_output = attn_output.index_select(0, ctx.gather_ids)
+
         output, _ = self.o_proj(attn_output)
         return output
 
@@ -774,6 +778,9 @@ def forward(
                 ctx=ctx,
                 out_cache_loc=out_cache_loc,
             )
+            if ctx.draft_first_step_reduce:
+                # Gather residual to self_attention's [bs, H].
+                residual = residual.index_select(0, ctx.gather_ids)
             hidden_states, residual = self.comm_manager.post_attn_reduce_norm(
                 hidden_states, residual, ctx
             )

From 161a324eeb541757b0fccd805a163905b0aa3820 Mon Sep 17 00:00:00 2001
From: rjzhb <rjzhb222@163.com>
Date: Wed, 27 May 2026 05:20:56 +0000
Subject: [PATCH 14/18] refactor(eagle3): drop per-model
 supports_draft_first_step_reduce flag

Signed-off-by: rjzhb <rjzhb222@163.com>
---
 .../tokenspeed/runtime/execution/drafter/eagle.py   |  6 +-----
 .../tokenspeed/runtime/execution/model_executor.py  | 13 +++----------
 python/tokenspeed/runtime/models/deepseek_v3.py     |  2 --
 python/tokenspeed/runtime/models/llama_eagle3.py    |  1 -
 4 files changed, 4 insertions(+), 18 deletions(-)

diff --git a/python/tokenspeed/runtime/execution/drafter/eagle.py b/python/tokenspeed/runtime/execution/drafter/eagle.py
index b6ac40eb5..ee2722582 100644
--- a/python/tokenspeed/runtime/execution/drafter/eagle.py
+++ b/python/tokenspeed/runtime/execution/drafter/eagle.py
@@ -212,11 +212,7 @@ def _run_first_step(
             draft_input, bs, draft_input.input_num_tokens
         )
 
-        # Models declare ``supports_draft_first_step_reduce`` to opt in; MLA /
-        # NextN heads that don't declare it fall back to the standard path.
-        draft_first_step_reduce = forward_mode.is_decode() and getattr(
-            self.draft_model_runner.model, "supports_draft_first_step_reduce", False
-        )
+        draft_first_step_reduce = forward_mode.is_decode()
 
         if draft_first_step_reduce and self.attn_backend.support_kv_cache_prewrite:
             # Prewrite path slices Q to one row per request and uses the decode
diff --git a/python/tokenspeed/runtime/execution/model_executor.py b/python/tokenspeed/runtime/execution/model_executor.py
index 4cff78d5d..4e7134c02 100644
--- a/python/tokenspeed/runtime/execution/model_executor.py
+++ b/python/tokenspeed/runtime/execution/model_executor.py
@@ -882,16 +882,9 @@ def execute_idle_forward(
         # NCCL collectives. Idle ranks must match those collectives:
         # 1 first-step forward + (spec_num_steps - 1) multi-step decode forwards.
         if self.drafter is not None:
-            # Mirror active ranks: when all non-idle ranks are decoding and the
-            # draft model opts in, the catch-up step sizes collectives from
-            # bs/global_bs instead of bs*spec_num_tokens. Idle ranks must use
-            # the same basis or RSAG/AR shapes diverge.
-            supports_reduce = all_decode_or_idle and getattr(
-                self.drafter.draft_model_runner.model,
-                "supports_draft_first_step_reduce",
-                False,
-            )
             for step_idx in range(self.drafter.spec_num_steps):
+                # Mirror active rank's catch-up step: when all non-idle ranks
+                # are decoding, step 0 sizes collectives from bs/global_bs.
                 draft_ctx = ForwardContext(
                     attn_backend=self.drafter.attn_backend,
                     token_to_kv_pool=self.drafter.token_to_kv_pool,
@@ -903,7 +896,7 @@ def execute_idle_forward(
                     global_num_tokens=global_num_tokens,
                     global_bs=global_bs,
                     all_decode_or_idle=all_decode_or_idle,
-                    draft_first_step_reduce=(step_idx == 0 and supports_reduce),
+                    draft_first_step_reduce=(step_idx == 0 and all_decode_or_idle),
                 )
                 self.drafter.draft_model_runner.forward(
                     draft_ctx,
diff --git a/python/tokenspeed/runtime/models/deepseek_v3.py b/python/tokenspeed/runtime/models/deepseek_v3.py
index 24af5d376..bf50a6c94 100644
--- a/python/tokenspeed/runtime/models/deepseek_v3.py
+++ b/python/tokenspeed/runtime/models/deepseek_v3.py
@@ -1835,8 +1835,6 @@ class Eagle3DeepseekV2ForCausalLM(DeepseekV3ForCausalLM):
     concatenated [embeds || hidden_states] as input.
     """
 
-    supports_draft_first_step_reduce = True
-
     def __init__(
         self,
         config: PretrainedConfig,
diff --git a/python/tokenspeed/runtime/models/llama_eagle3.py b/python/tokenspeed/runtime/models/llama_eagle3.py
index b5534d83d..503d9a48f 100644
--- a/python/tokenspeed/runtime/models/llama_eagle3.py
+++ b/python/tokenspeed/runtime/models/llama_eagle3.py
@@ -570,7 +570,6 @@ def forward(
 class LlamaForCausalLMEagle3(BaseCausalLM):
 
     model_cls = Eagle3LlamaModel
-    supports_draft_first_step_reduce = True
 
     def __init__(
         self,

From e0f2e685c53bb5738d31262149c865039a842632 Mon Sep 17 00:00:00 2001
From: rjzhb <rjzhb222@163.com>
Date: Wed, 27 May 2026 05:24:15 +0000
Subject: [PATCH 15/18] chore(eagle3): trim verbose comments on reduce path

Signed-off-by: rjzhb <rjzhb222@163.com>
---
 python/tokenspeed/runtime/execution/drafter/eagle.py | 7 ++-----
 python/tokenspeed/runtime/models/deepseek_v3.py      | 3 +--
 2 files changed, 3 insertions(+), 7 deletions(-)

diff --git a/python/tokenspeed/runtime/execution/drafter/eagle.py b/python/tokenspeed/runtime/execution/drafter/eagle.py
index ee2722582..ced059c22 100644
--- a/python/tokenspeed/runtime/execution/drafter/eagle.py
+++ b/python/tokenspeed/runtime/execution/drafter/eagle.py
@@ -215,11 +215,8 @@ def _run_first_step(
         draft_first_step_reduce = forward_mode.is_decode()
 
         if draft_first_step_reduce and self.attn_backend.support_kv_cache_prewrite:
-            # Prewrite path slices Q to one row per request and uses the decode
-            # kernel, which reads seq_lens as the attention upper bound. Trim
-            # by the rejected-draft count so the live query does not read dead
-            # positions. Non-prewrite runs full multi-token attention with
-            # per-position causal masking, so seq_lens must stay unchanged.
+            # Trim seq_lens by rejected-draft count so the sliced decode
+            # query does not attend to dead positions.
             correction = (self.spec_num_tokens - draft_input.accept_lengths).to(
                 self.draft_seq_lens_buf.dtype
             )
diff --git a/python/tokenspeed/runtime/models/deepseek_v3.py b/python/tokenspeed/runtime/models/deepseek_v3.py
index bf50a6c94..78beecd58 100644
--- a/python/tokenspeed/runtime/models/deepseek_v3.py
+++ b/python/tokenspeed/runtime/models/deepseek_v3.py
@@ -1705,8 +1705,7 @@ def forward(
             )
 
             if ctx.draft_first_step_reduce:
-                # self_attn returned [bs, H]; gather residual to match before
-                # the fused reduce+norm.
+                # Gather residual to self_attn's [bs, H].
                 residual = residual.index_select(0, ctx.gather_ids)
             hidden_states, residual = self.comm_manager.post_attn_reduce_norm(
                 hidden_states, residual, ctx

From 4500a4fce0f3ee6543cbc597d9964f604159744c Mon Sep 17 00:00:00 2001
From: rjzhb <rjzhb222@163.com>
Date: Wed, 27 May 2026 05:29:09 +0000
Subject: [PATCH 16/18] chore(eagle3): replace EAGLE-only wording on
 draft_first_step_reduce comments

Signed-off-by: rjzhb <rjzhb222@163.com>
---
 python/tokenspeed/runtime/execution/context.py       | 2 +-
 python/tokenspeed/runtime/layers/logits_processor.py | 2 +-
 2 files changed, 2 insertions(+), 2 deletions(-)

diff --git a/python/tokenspeed/runtime/execution/context.py b/python/tokenspeed/runtime/execution/context.py
index 80609de8c..b4e80dae5 100644
--- a/python/tokenspeed/runtime/execution/context.py
+++ b/python/tokenspeed/runtime/execution/context.py
@@ -50,7 +50,7 @@ class ForwardContext:
     forward_mode: ForwardMode | None
     req_to_page: torch.Tensor | None = None
     capture_hidden_mode: CaptureHiddenMode | None = CaptureHiddenMode.NULL
-    # EAGLE draft head's first step prunes to one live row per request.
+    # Spec decode draft head's first step prunes to one live row per request.
     draft_first_step_reduce: bool = False
 
     # --- dp attention ---
diff --git a/python/tokenspeed/runtime/layers/logits_processor.py b/python/tokenspeed/runtime/layers/logits_processor.py
index 161e0b987..8df400e8c 100755
--- a/python/tokenspeed/runtime/layers/logits_processor.py
+++ b/python/tokenspeed/runtime/layers/logits_processor.py
@@ -211,7 +211,7 @@ def forward(
         if not logits_metadata.extend_return_logprob:
             gather_ids = logits_metadata.gather_ids
             # Shapes align iff midlayer already pruned to one row per request
-            # (EAGLE draft first-step reduce). Other paths emit [N, H] with N > bs.
+            # (draft first-step reduce). Other paths emit [N, H] with N > bs.
             if gather_ids is None or gather_ids.shape[0] == hidden_states.shape[0]:
                 pruned_states = hidden_states
                 if aux_hidden_states is not None:

From db82f61222bdbc27e06d2dccf9f16d983008147c Mon Sep 17 00:00:00 2001
From: rjzhb <rjzhb222@163.com>
Date: Wed, 27 May 2026 05:48:12 +0000
Subject: [PATCH 17/18] fix(eagle3): align MoE collectives and final-norm
 decision with draft reduce

Signed-off-by: rjzhb <rjzhb222@163.com>
---
 .../runtime/distributed/comm_manager.py          | 16 +++++++++-------
 python/tokenspeed/runtime/models/llama_eagle3.py |  4 +++-
 2 files changed, 12 insertions(+), 8 deletions(-)

diff --git a/python/tokenspeed/runtime/distributed/comm_manager.py b/python/tokenspeed/runtime/distributed/comm_manager.py
index a4e52a839..1660d493e 100755
--- a/python/tokenspeed/runtime/distributed/comm_manager.py
+++ b/python/tokenspeed/runtime/distributed/comm_manager.py
@@ -88,15 +88,17 @@ def dense_tp_group_scattered_num_tokens(self, ctx: ForwardContext) -> list[int]:
 
     def moe_tp_ep_group_scattered_num_tokens(self, ctx: ForwardContext) -> list[int]:
         tp_ep_size = self.mapping.moe.tp_ep_size
-        if ctx.global_num_tokens is not None:
-            # DP path: global_num_tokens has one entry per world rank.
-            # MoE group = tp_ep_size contiguous ranks starting at dp_rank * tp_ep_size.
+        # Under draft first-step reduce, the midlayer pruned activations to bs
+        # rows before pre_moe_comm; MoE collectives must size accordingly.
+        global_counts = (
+            ctx.global_bs if ctx.draft_first_step_reduce else ctx.global_num_tokens
+        )
+        if global_counts is not None:
             start = self.mapping.moe.dp_rank * tp_ep_size
-            return list(ctx.global_num_tokens[start : start + tp_ep_size])
-        # No global_num_tokens (e.g., drafter context or non-DP).
-        # This rank has all input_num_tokens; other ranks in the EP group have 0.
+            return list(global_counts[start : start + tp_ep_size])
+        num_tokens = ctx.bs if ctx.draft_first_step_reduce else ctx.input_num_tokens
         result = [0] * tp_ep_size
-        result[self.mapping.moe.tp_ep_rank] = ctx.input_num_tokens
+        result[self.mapping.moe.tp_ep_rank] = num_tokens
         return result
 
     # ---- Communication patterns ----
diff --git a/python/tokenspeed/runtime/models/llama_eagle3.py b/python/tokenspeed/runtime/models/llama_eagle3.py
index 503d9a48f..3ac3745b0 100644
--- a/python/tokenspeed/runtime/models/llama_eagle3.py
+++ b/python/tokenspeed/runtime/models/llama_eagle3.py
@@ -551,7 +551,9 @@ def forward(
             fuse_embed_reduce=fuse_embed_reduce,
         )
 
-        if midlayer.comm_manager.should_fuse(hidden_states.shape[0]):
+        # Decide on pre-slice token count so this matches the path midlayer
+        # actually took; under draft reduce, hidden_states.shape[0] shrinks.
+        if midlayer.comm_manager.should_fuse(input_ids.shape[0]):
             hidden_states_to_logits, hidden_states_to_aux = hidden_states, residual
         else:
             hidden_states_to_logits, hidden_states_to_aux = self.norm(

From 8905345fc1ffc7f40d64f0994e578af9449787bd Mon Sep 17 00:00:00 2001
From: rjzhb <rjzhb222@163.com>
Date: Wed, 27 May 2026 18:50:33 +0000
Subject: [PATCH 18/18] fix(deepseek-nextn): delegate final norm to
 comm_manager for fused allreduce

Signed-off-by: rjzhb <rjzhb222@163.com>
---
 python/tokenspeed/runtime/models/deepseek_nextn.py | 7 ++++---
 1 file changed, 4 insertions(+), 3 deletions(-)

diff --git a/python/tokenspeed/runtime/models/deepseek_nextn.py b/python/tokenspeed/runtime/models/deepseek_nextn.py
index cfa0ec6e8..fee450073 100755
--- a/python/tokenspeed/runtime/models/deepseek_nextn.py
+++ b/python/tokenspeed/runtime/models/deepseek_nextn.py
@@ -141,11 +141,12 @@ def forward(
         )
 
         if not ctx.forward_mode.is_idle():
-            hidden_states, _ = self.shared_head.norm(hidden_states, residual)
             if not ENABLE_CP:
-                hidden_states, _ = self.decoder.comm_manager.post_final_norm_comm(
-                    hidden_states, residual, ctx
+                hidden_states = self.decoder.comm_manager.final_norm(
+                    hidden_states, residual, ctx, self.shared_head.norm
                 )
+            else:
+                hidden_states, _ = self.shared_head.norm(hidden_states, residual)
         if CP_METADATA:
             hidden_states = cp_all_gather_rerange_output(
                 hidden_states,