[None][perf] GLM-5.3-Flash: share KDA verify offsets, fuse mHC in every phase, BF16 router GEMM

tensorrt_llm/_torch/models/modeling_glm5_next.py

@@ -867,6 +867,7 @@ def __init__(
         dtype: torch.dtype = torch.bfloat16,
         mapping: Mapping | None = None,
         allreduce_strategy: AllReduceStrategy = AllReduceStrategy.ONESHOT,
+        aux_stream: torch.cuda.Stream | None = None,
     ) -> None:
         del dtype  # the shared mixer is bf16 (fp32 gate parameters and pools)
         # The checkpoint's ``linear_attn_config`` does not name the gate rank:
@@ -879,7 +880,13 @@ def __init__(
         # The mixer's own row-parallel o_proj AllReduce (fixed strategy; the
         # size-aware Glm5NextAllReduce used elsewhere is a perf option, not a
         # correctness requirement, and AUTO's autotuner raced at TP4 decode).
-        super().__init__(config, layer_idx, mapping=mapping, allreduce_strategy=allreduce_strategy)
+        super().__init__(
+            config,
+            layer_idx,
+            mapping=mapping,
+            allreduce_strategy=allreduce_strategy,
+            aux_stream=aux_stream,
+        )
         # Preserve checkpoint FP32 gate parameters. MetaInitMode rejects detach
         # on meta tensors, so create empty FP32 parameters on that branch.
         for name in ("A_log", "dt_bias"):
@@ -1501,10 +1508,12 @@ def forward_verify(
 
 
 class Glm5NextGate(DeepseekV3Gate):
-    """DeepSeek noaux_tc routing with FP32 weights, logits and correction bias.
+    """DeepSeek noaux_tc routing with FP32 logits and correction bias.
 
-    Preserve small inter-expert score differences when adding the correction bias;
-    rounding it to BF16 can change expert selection.
+    The checkpoint router weight is BF16, so the base BF16 GEMM with FP32 output
+    forms the same products as FP32 math. Preserve small inter-expert score
+    differences when adding the correction bias; rounding it to BF16 can change
+    expert selection.
     """
 
     def __init__(self, config: PretrainedConfig, moe_backend: str = "CUTLASS") -> None:
@@ -1517,23 +1526,18 @@ def __init__(self, config: PretrainedConfig, moe_backend: str = "CUTLASS") -> No
             n_group=int(getattr(config, "n_group", 1) or 1),
             topk_group=int(getattr(config, "topk_group", 1) or 1),
             routed_scaling_factor=float(config.routed_scaling_factor),
-            dtype=torch.float32,
+            dtype=torch.bfloat16,
             fuse_routing_kernel=True,
             apply_routing=False,
             moe_backend=moe_backend,
         )
-        self.hidden_size = int(config.hidden_size)
         self.norm_topk_prob = bool(config.norm_topk_prob)
         # FP32 regardless of the MoE backend (the base class picks bf16 for
         # TRTLLM); the fused routing kernels consume the FP32 tensor directly.
         self.e_score_correction_bias = nn.Parameter(
             torch.empty(int(config.n_routed_experts), dtype=torch.float32), requires_grad=False
         )
 
-    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
-        flat = hidden_states.reshape(-1, self.hidden_size)
-        return torch.nn.functional.linear(flat.float(), self.weight)
-
 
 class Glm5NextMoE(nn.Module):
     """Fused routed experts with a shared GatedMLP expert.
@@ -1740,6 +1744,7 @@ def __init__(
         model_config: ModelConfig,
         dtype: torch.dtype = torch.bfloat16,
         aux_stream: torch.cuda.Stream | None = None,
+        attention_aux_stream: torch.cuda.Stream | None = None,
     ) -> None:
         super().__init__()
         self.layer_idx = layer_idx
@@ -1755,6 +1760,7 @@ def __init__(
                 dtype=dtype,
                 mapping=mapping,
                 allreduce_strategy=model_config.allreduce_strategy,
+                aux_stream=attention_aux_stream,
             )
         else:
             self.self_attn = Glm5NextSparseAttention(
@@ -1855,8 +1861,11 @@ def __init__(self, model_config: ModelConfig[PretrainedConfig]) -> None:
 
         # Decoder and draft layers execute serially and share this side stream.
         self.aux_stream_dict = {}
-        if not model_config.mapping.enable_attention_dp and torch.cuda.is_available():
-            self.aux_stream_dict[AuxStreamType.MoeShared] = torch.cuda.Stream()
+        if torch.cuda.is_available():
+            # KDA overlaps its small gate projections with the QKV projection.
+            self.aux_stream_dict[AuxStreamType.Attention] = torch.cuda.Stream()
+            if not model_config.mapping.enable_attention_dp:
+                self.aux_stream_dict[AuxStreamType.MoeShared] = torch.cuda.Stream()
 
         self.embed_tokens = Embedding(int(config.vocab_size), int(config.hidden_size), dtype=dtype)
         self.layers = nn.ModuleList(
@@ -1868,6 +1877,7 @@ def __init__(self, model_config: ModelConfig[PretrainedConfig]) -> None:
                     model_config,
                     dtype=dtype,
                     aux_stream=self.aux_stream_dict.get(AuxStreamType.MoeShared),
+                    attention_aux_stream=self.aux_stream_dict.get(AuxStreamType.Attention),
                 )
                 for i in range(schedule.num_layers)
             ]
@@ -1896,24 +1906,20 @@ def forward(
             raise ValueError("specify exactly one of input_ids or inputs_embeds")
         if inputs_embeds is None:
             inputs_embeds = self.embed_tokens(input_ids)
-        streams = self.expand_streams(inputs_embeds)
         if runtime_ctx is None:
             runtime_ctx = build_glm5_next_runtime_context(attn_metadata)
         # Only the decoder stack: under one-model MTP the speculative base
         # class appends the draft layer(s) to ``self.layers`` (so the
         # checkpoint's ``model.layers.45.*`` names resolve), and those run in
         # the speculative worker's draft loop, not here.
-        # Generation-only batches take the fused hyper-connection loop. Prefill
-        # (and mixed) batches keep the per-layer path so the engine's warmup /
-        # KV-cache sizing pass measures the same activation peak that mixed
-        # context+generation iterations reach in production (a lower
-        # pure-prefill peak over-allocates the KV cache and OOMs later).
-        if (
-            runtime_ctx.num_contexts == 0
-            and runtime_ctx.num_generations > 0
-            and self._fused_hc_ok()
-        ):
-            return self.collapse_streams(self._forward_single_phase_fused_hc(streams, runtime_ctx))
+        # Every batch phase takes the fused hyper-connection loop, so the
+        # engine's warmup / KV-cache sizing pass measures the same activation
+        # peak that mixed context+generation iterations reach in production.
+        if self._fused_hc_ok():
+            return self.collapse_streams(
+                self._forward_single_phase_fused_hc(inputs_embeds, runtime_ctx)
+            )
+        streams = self.expand_streams(inputs_embeds)
         for layer_idx in range(self.schedule.num_layers):
             streams = self.layers[layer_idx](
                 position_ids=position_ids,
@@ -1935,7 +1941,7 @@ def _fused_hc_ok(self) -> bool:
         return cached
 
     def _forward_single_phase_fused_hc(
-        self, streams: torch.Tensor, runtime_ctx: Glm5NextRuntimeContext
+        self, embeds: torch.Tensor, runtime_ctx: Glm5NextRuntimeContext
     ) -> torch.Tensor:
         """Single-phase forward with fused hyper-connection boundaries.
 
@@ -1944,13 +1950,17 @@ def _forward_single_phase_fused_hc(
         RMSNorm; ``mHC.fused_hc`` (the in-tree DeepSeek-V4 boundary op) runs
         the three as one kernel, so a layer costs 2 fused boundaries instead
         of 4 mappings + 2 norms. The math matches the per-layer forward; the
-        first pre-mapping and last post-mapping stay unfused. Only generation-only batches use it
-        (see :meth:`forward`).
+        first pre-mapping and last post-mapping stay unfused. Pipeline-parallel
+        ranks keep the per-layer path (see :meth:`_fused_hc_ok`).
+
+        Only ``residual`` holds the streams and each sublayer output is dropped
+        once its boundary consumed it, so a long prefill keeps one copy of the
+        [tokens, hc_mult, hidden] streams alive, as the per-layer loop does.
         """
         num_layers = self.schedule.num_layers
         first = self.layers[0]
-        residual = streams
-        post, comb, x = first.hc_attn.pre_mapping(streams)
+        residual = self.expand_streams(embeds)
+        post, comb, x = first.hc_attn.pre_mapping(residual)
         x = first.input_layernorm(x)
         for layer_idx in range(num_layers):
             layer = self.layers[layer_idx]
@@ -1966,6 +1976,7 @@ def _forward_single_phase_fused_hc(
                 norm_weight=layer.post_attention_layernorm.weight,
                 norm_eps=layer.post_attention_layernorm.variance_epsilon,
             )
+            del attn_out
             mlp_out = layer.run_mlp(x, runtime_ctx.metadata.all_rank_num_tokens)
             if layer_idx + 1 < num_layers:
                 nxt = self.layers[layer_idx + 1]
@@ -1979,6 +1990,7 @@ def _forward_single_phase_fused_hc(
                 )
             else:
                 residual = layer.hc_ffn.post_mapping(mlp_out, residual, post, comb)
+            del mlp_out
         return residual
 
     def expand_streams(self, embeds: torch.Tensor) -> torch.Tensor: