perf(m3): enable prefill CUDA graph + EAGLE3 draft-loop kernel cleanup

docs/recipes/models.md

@@ -89,7 +89,6 @@ tokenspeed serve nvidia/MiniMax-M3-NVFP4 \
     --speculative-eagle-topk 1 \
     --speculative-num-draft-tokens 4 \
     --drafter-attention-backend fa4 \
-    --disable-prefill-graph \
     --disable-kvstore \
     --block-size 128 \
     --trust-remote-code \

python/tokenspeed/runtime/execution/prefill_graph.py

@@ -49,6 +49,7 @@
 from typing import TYPE_CHECKING, NamedTuple
 
 import torch
+import tqdm
 
 from tokenspeed.runtime.execution.breakable_cuda_graph import (
     BreakableCapture,
@@ -61,7 +62,10 @@
 )
 from tokenspeed.runtime.layers.logits_processor import LogitsMetadata
 from tokenspeed.runtime.utils import get_colorful_logger
-from tokenspeed.runtime.utils.common import maybe_inference_mode
+from tokenspeed.runtime.utils.common import (
+    get_available_gpu_memory,
+    maybe_inference_mode,
+)
 
 logger = get_colorful_logger(__name__)
 
@@ -341,7 +345,17 @@ def capture(self, decode_wrapper: CudaGraphWrapper | None = None) -> None:
             self.disable = True
 
     def _capture_all_buckets(self, decode_wrapper: CudaGraphWrapper | None) -> None:
-        for bucket in sorted(self.capture_buckets, reverse=True):
+        rank = self.config.global_rank
+        buckets = sorted(self.capture_buckets, reverse=True)
+        capture_range = tqdm.tqdm(buckets) if rank == 0 else buckets
+        for bucket in capture_range:
+            if rank == 0:
+                avail_mem = get_available_gpu_memory(
+                    self.config.device, self.config.gpu_id, empty_cache=False
+                )
+                capture_range.set_description(
+                    f"Capturing prefill buckets ({bucket=} {avail_mem=:.2f} GB)"
+                )
             self._ctx = self._make_dummy_batch(bucket, decode_wrapper)
             self._land_input_embeds(
                 self._embed_tokens(self.input_buffers.input_ids_buf[:bucket]), bucket

python/tokenspeed/runtime/layers/layernorm.py

@@ -106,7 +106,7 @@ def forward(
         self,
         x: torch.Tensor,
         residual: torch.Tensor | None = None,
-        inplace: bool = False,
+        out: torch.Tensor | None = None,
     ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
         # There might be no tokens here
         if x.shape[0] == 0:
@@ -117,10 +117,8 @@ def forward(
 
         if _is_amd:
             if residual is not None:
-                if inplace:
-                    raise ValueError(
-                        "fused add rmsnorm does not support inplace operation"
-                    )
+                if out is not None:
+                    raise ValueError("fused add rmsnorm does not support out")
                 return triton_rmsnorm(
                     x,
                     self.weight.data,
@@ -131,14 +129,12 @@ def forward(
                 x,
                 self.weight.data,
                 self.variance_epsilon,
-                out=x if inplace else None,
+                out=out,
             )
         else:
             if residual is not None:
-                if inplace:
-                    raise ValueError(
-                        "fused_add_rmsnorm does not support inplace operation"
-                    )
+                if out is not None:
+                    raise ValueError("fused_add_rmsnorm does not support out")
                 fused_add_rmsnorm(
                     x,
                     residual,
@@ -151,7 +147,7 @@ def forward(
                 x,
                 self.weight.data,
                 self.variance_epsilon,
-                out=x if inplace else None,
+                out=out,
                 enable_pdl=pdl_enabled(),
             )
             return out

python/tokenspeed/runtime/models/base/transformer_model.py

@@ -178,7 +178,7 @@ def forward(
                 )
             elif isinstance(first_layer, BaseDecoderLayer):
                 fuse_embed_reduce = (
-                    self.mapping.attn.tp_size > 1
+                    self.embed_tokens.tp_size > 1
                     and first_layer.comm_manager.should_fuse(input_ids.shape[0])
                 )
             else:

python/tokenspeed/runtime/models/deepseek_v3.py

@@ -684,7 +684,7 @@ def _project_q_latent(
             q = self.q_proj(hidden_states)[0]
             latent_cache = self.kv_a_proj_with_mqa(hidden_states)[0]
             kv_a = latent_cache[..., : self.kv_lora_rank]
-            self.kv_a_layernorm(kv_a, inplace=True)
+            self.kv_a_layernorm(kv_a, out=kv_a)
         return q, latent_cache
 
     @break_point
@@ -2040,6 +2040,16 @@ def __init__(
             if getattr(config, "fc_norm", False)
             else None
         )
+        self.fused_fc_norms = (
+            nn.ModuleList(
+                [
+                    FusedRMSNorm(self.fc_norm[i], self.fc_norm[i + 1])
+                    for i in range(0, self.num_fc_input_dim - 1, 2)
+                ]
+            )
+            if self.fc_norm is not None
+            else None
+        )
         self.norm_output = getattr(config, "norm_output", False)
 
     def forward(
@@ -2060,15 +2070,23 @@ def forward(
 
         hidden_states = captured_hidden_states
         if hidden_states.size(-1) != embeds.size(-1):
-            if self.fc_norm is not None:
+            if self.fc_norm is not None and hidden_states.shape[0] > 0:
                 chunks = hidden_states.chunk(self.num_fc_input_dim, dim=-1)
-                hidden_states = torch.cat(
-                    [
-                        norm(chunk)
-                        for norm, chunk in zip(self.fc_norm, chunks, strict=True)
-                    ],
-                    dim=-1,
-                )
+                normed = torch.empty_like(hidden_states)
+                out_chunks = normed.chunk(self.num_fc_input_dim, dim=-1)
+                i = 0
+                for fused in self.fused_fc_norms:
+                    fused(
+                        input_q_a=chunks[i],
+                        input_kv_a=chunks[i + 1],
+                        output_q_a=out_chunks[i],
+                        output_kv_a=out_chunks[i + 1],
+                    )
+                    i += 2
+                if i < self.num_fc_input_dim:
+                    # Odd count: single norm into the last slice.
+                    self.fc_norm[i](chunks[i], out=out_chunks[i])
+                hidden_states = normed
             hidden_states, _ = self.fc(hidden_states)
 
         residual = None

python/tokenspeed/runtime/models/llama_eagle3.py

@@ -40,15 +40,18 @@
 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
+from tokenspeed.runtime.layers.layernorm import FusedRMSNorm, RMSNorm
 from tokenspeed.runtime.layers.linear import (
     ColumnParallelLinear,
     MergedColumnParallelLinear,
     RowParallelLinear,
 )
 from tokenspeed.runtime.layers.logits_processor import LogitsProcessor
 from tokenspeed.runtime.layers.quantization.base_config import QuantizationConfig
-from tokenspeed.runtime.layers.vocab_parallel_embedding import ParallelLMHead
+from tokenspeed.runtime.layers.vocab_parallel_embedding import (
+    ParallelLMHead,
+    VocabParallelEmbedding,
+)
 from tokenspeed.runtime.model_loader.weight_utils import default_weight_loader
 from tokenspeed.runtime.models.base import (
     BaseCausalLM,
@@ -230,6 +233,10 @@ def __init__(
         )
 
         self.hidden_norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+        self.fused_input_hidden_norm = FusedRMSNorm(
+            self.input_layernorm,
+            self.hidden_norm,
+        )
 
     def resolve_attn(self, prefix: str) -> nn.Module:
 
@@ -295,11 +302,27 @@ def forward_low_latency(
                 embeds,
                 torch.zeros_like(embeds),
             )
+            hidden_states = self.hidden_norm(hidden_states)
+            hidden_states = concat(embeds, hidden_states)
         else:
-            embeds = self.input_layernorm(embeds)
-
-        hidden_states = self.hidden_norm(hidden_states)
-        hidden_states = concat(embeds, hidden_states)
+            h = embeds.size(-1)
+            fused_norm_out = torch.empty(
+                embeds.size(0),
+                h + hidden_states.size(-1),
+                dtype=embeds.dtype,
+                device=embeds.device,
+            )
+            # FusedRMSNorm's q_a/kv_a kwargs are MLA-specific names. Here
+            # embeds and hidden_states correspond to q_a and kv_a, separately.
+            # Skip the launch on zero-token (idle) forwards.
+            if embeds.shape[0] > 0:
+                self.fused_input_hidden_norm(
+                    input_q_a=embeds,
+                    input_kv_a=hidden_states,
+                    output_q_a=fused_norm_out[..., :h],
+                    output_kv_a=fused_norm_out[..., h:],
+                )
+            hidden_states = fused_norm_out
 
         # Attention
         hidden_states = self.comm_manager.pre_attn_comm(hidden_states, ctx)
@@ -363,9 +386,24 @@ def forward(
         # Non-fused path: fuse_embed_reduce is always False here because
         # the model only sets it when should_fuse() is True.
         residual = hidden_states
-        embeds = self.input_layernorm(embeds)
-        hidden_states = self.hidden_norm(hidden_states)
-        hidden_states = torch.cat([embeds, hidden_states], dim=-1)
+        h = embeds.size(-1)
+        fused_norm_out = torch.empty(
+            embeds.size(0),
+            h + hidden_states.size(-1),
+            dtype=embeds.dtype,
+            device=embeds.device,
+        )
+        # FusedRMSNorm's q_a/kv_a kwargs are MLA-specific names. Here
+        # embeds and hidden_states correspond to q_a and kv_a, separately.
+        # Skip the launch on zero-token (idle) forwards.
+        if embeds.shape[0] > 0:
+            self.fused_input_hidden_norm(
+                input_q_a=embeds,
+                input_kv_a=hidden_states,
+                output_q_a=fused_norm_out[..., :h],
+                output_kv_a=fused_norm_out[..., h:],
+            )
+        hidden_states = fused_norm_out
 
         # Attention
         hidden_states = self.comm_manager.pre_attn_comm(hidden_states, ctx)
@@ -458,6 +496,16 @@ def __init__(
             if getattr(config, "fc_norm", False)
             else None
         )
+        self.fused_fc_norms = (
+            nn.ModuleList(
+                [
+                    FusedRMSNorm(self.fc_norm[i], self.fc_norm[i + 1])
+                    for i in range(0, self.num_fc_input_dim - 1, 2)
+                ]
+            )
+            if self.fc_norm is not None
+            else None
+        )
         self.norm_output = getattr(config, "norm_output", False)
 
     def forward(
@@ -471,13 +519,13 @@ def forward(
     ) -> torch.Tensor:
 
         if input_embeds is None:
-            # When TP > 1 and fused allreduce+norm is available, skip the
-            # NCCL allreduce in the embedding and let the midlayer fuse it
-            # with the input_layernorm.
+            # When the embedding is vocab-parallel (tp_size > 1), skip its
+            # NCCL allreduce and let the midlayer fuse it with input_layernorm.
+            # Non-sharded embeddings (tp_size == 1) need no allreduce at all.
             midlayer = self.midlayer
             num_tokens = input_ids.shape[0]
             fuse_embed_reduce = (
-                self.mapping.attn.tp_size > 1
+                self.embed_tokens.tp_size > 1
                 and midlayer.comm_manager.should_fuse(num_tokens)
             )
             embeds = self.embed_tokens(input_ids, reduce_results=not fuse_embed_reduce)
@@ -491,15 +539,23 @@ def forward(
         if hidden_states.size(-1) != embeds.size(-1):
             if self.input_norm is not None:
                 hidden_states = self.input_norm(hidden_states)
-            if self.fc_norm is not None:
+            if self.fc_norm is not None and hidden_states.shape[0] > 0:
                 chunks = hidden_states.chunk(self.num_fc_input_dim, dim=-1)
-                hidden_states = torch.cat(
-                    [
-                        norm(chunk)
-                        for norm, chunk in zip(self.fc_norm, chunks, strict=True)
-                    ],
-                    dim=-1,
-                )
+                normed = torch.empty_like(hidden_states)
+                out_chunks = normed.chunk(self.num_fc_input_dim, dim=-1)
+                i = 0
+                for fused in self.fused_fc_norms:
+                    fused(
+                        input_q_a=chunks[i],
+                        input_kv_a=chunks[i + 1],
+                        output_q_a=out_chunks[i],
+                        output_kv_a=out_chunks[i + 1],
+                    )
+                    i += 2
+                if i < self.num_fc_input_dim:
+                    # Odd count: single norm into the last slice.
+                    self.fc_norm[i](chunks[i], out=out_chunks[i])
+                hidden_states = normed
             hidden_states, _ = self.fc(hidden_states)
 
         residual = None
@@ -673,6 +729,22 @@ def set_embed_and_head(self, embed, head):
             and self.config.target_hidden_size != self.config.hidden_size
         ):
             return
+        if embed.shape != self.model.embed_tokens.weight.shape:
+            # Target embedding layout differs (e.g. replicated instead of
+            # vocab-parallel); rebuild the draft embedding to match it.
+            with torch.device("meta"):
+                replicated = VocabParallelEmbedding(
+                    self.config.vocab_size,
+                    self.config.hidden_size,
+                    prefix="model.embed_tokens",
+                )
+            if embed.shape != replicated.weight.shape:
+                raise ValueError(
+                    f"Cannot share target embed of shape {tuple(embed.shape)}: "
+                    f"draft expects {tuple(self.model.embed_tokens.weight.shape)} "
+                    f"(sharded) or {tuple(replicated.weight.shape)} (replicated)."
+                )
+            self.model.embed_tokens = replicated
         del self.model.embed_tokens.weight
         self.model.embed_tokens.weight = embed
         if head is not None and self.load_lm_head_from_target:

python/tokenspeed/runtime/models/minimax_m3.py

@@ -72,6 +72,7 @@
 )
 from tokenspeed.runtime.layers.quantization.base_config import QuantizationConfig
 from tokenspeed.runtime.layers.rotary_embedding import get_rope
+from tokenspeed.runtime.layers.vocab_parallel_embedding import VocabParallelEmbedding
 from tokenspeed.runtime.model_loader.weight_utils import default_weight_loader
 from tokenspeed.runtime.models.base import BaseCausalLM, BaseTransformerModel
 from tokenspeed.runtime.models.base.comm_ops import FinalNormOp
@@ -701,8 +702,9 @@ def qknorm_rope(
             self.q_norm.gemma_weight,
             self.k_norm.gemma_weight,
             self.q_norm.variance_epsilon,
+            enable_pdl=pdl_enabled(),
         )
-        q, k = self.rotary_emb(positions, q, k)
+        q, k = self.rotary_emb(positions, q, k, enable_pdl=pdl_enabled())
         q = q.view(-1, self.num_heads, self.head_dim)
         k = k.view(-1, self.num_kv_heads, self.head_dim)
         v = v.view(-1, self.num_kv_heads, self.head_dim)
@@ -861,6 +863,13 @@ class MiniMaxM3Model(BaseTransformerModel):
 
     layer_cls = MiniMaxM3DecoderLayer
 
+    def resolve_embed(self, config: MiniMaxM3VLTextConfig, prefix: str) -> nn.Module:
+        return VocabParallelEmbedding(
+            config.vocab_size,
+            config.hidden_size,
+            prefix=add_prefix("embed_tokens", prefix),
+        )
+
     def resolve_layers(
         self,
         config: MiniMaxM3VLTextConfig,

python/tokenspeed/runtime/models/qwen3_5.py

@@ -1008,7 +1008,9 @@ def forward(
         if input_embeds is None:
             # Only skip embedding allreduce when the first layer's fused
             # allreduce+residual+norm will handle it
-            if self.layers[0].comm_manager.should_fuse(input_ids.shape[0]):
+            if self.embed_tokens.tp_size > 1 and self.layers[
+                0
+            ].comm_manager.should_fuse(input_ids.shape[0]):
                 hidden_states = self.embed_tokens(input_ids, reduce_results=False)
                 residual = torch.zeros_like(hidden_states)
             else:

test/agentic_benchmark/minimax_m3/tokenspeed/configs/attn_tp4_moe_ep4.sh

@@ -22,7 +22,6 @@ exec ts serve \
     --speculative-eagle-topk 1 \
     --speculative-num-draft-tokens 4 \
     --drafter-attention-backend fa4 \
-    --disable-prefill-graph \
     --disable-kvstore \
     --block-size 128 \
     --enable-cache-report \