perf(glm5next): dispatch mHC by token batch size

tests/models/test_glm5next_model.py

@@ -442,6 +442,34 @@ def test_glm5next_dflash_maps_target_layers_to_completed_outputs() -> None:
     assert supports_eagle3(Glm5NextForConditionalGeneration)
 
 
+@pytest.mark.parametrize(
+    ("num_tokens", "expected"),
+    [
+        (1, True),
+        (8, True),
+        (9, False),
+    ],
+)
+def test_glm5next_b12x_mhc_dispatches_decode_sized_batches(
+    num_tokens: int, expected: bool
+) -> None:
+    layer = Glm5NextDecoderLayer.__new__(Glm5NextDecoderLayer)
+    torch.nn.Module.__init__(layer)
+    layer._b12x_mhc = object()
+    layer._b12x_mhc_max_tokens = 8
+
+    assert layer._use_b12x_mhc(torch.empty(num_tokens, 4)) is expected
+
+
+def test_glm5next_b12x_mhc_dispatch_requires_available_backend() -> None:
+    layer = Glm5NextDecoderLayer.__new__(Glm5NextDecoderLayer)
+    torch.nn.Module.__init__(layer)
+    layer._b12x_mhc = None
+    layer._b12x_mhc_max_tokens = 8
+
+    assert not layer._use_b12x_mhc(torch.empty(1, 4))
+
+
 def test_glm5next_conditional_post_load_finalizes_language_model() -> None:
     class FakeLanguageModel(torch.nn.Module):
         def __init__(self) -> None:

vllm/models/glm5next/nvidia/model.py

@@ -437,6 +437,17 @@ def __init__(
             self.mhc_post_op = MHCPostOp()
             self.mhc_fused_post_pre_op = MHCFusedPostPreOp()
             self._b12x_mhc = None
+            max_decode_tokens = int(vllm_config.scheduler_config.max_num_seqs) * (
+                1 + int(vllm_config.num_speculative_tokens)
+            )
+            max_cudagraph_tokens = int(
+                vllm_config.compilation_config.max_cudagraph_capture_size or 0
+            )
+            max_b12x_mhc_tokens = max(max_decode_tokens, max_cudagraph_tokens)
+            if self.is_sequence_parallel:
+                tp_size = int(parallel_config.tensor_parallel_size)
+                max_b12x_mhc_tokens = (max_b12x_mhc_tokens + tp_size - 1) // tp_size
+            self._b12x_mhc_max_tokens = max_b12x_mhc_tokens
             if (
                 current_platform.is_cuda()
                 and current_platform.is_device_capability_family(120)
@@ -493,9 +504,11 @@ def forward(
         # hc_post with this layer's attn hc_pre into one kernel (inter-layer
         # fusion). Layer 0 has no incoming state -> standalone hc_pre.
         x = hidden_states
+        use_b12x_mhc = self._use_b12x_mhc(x)
         if post is None:
-            if self._b12x_mhc is not None:
+            if use_b12x_mhc:
                 assert self.hc_attn_fn_broadcast is not None
+                assert self._b12x_mhc is not None
                 residual, post, comb, x = self._b12x_mhc.run_pre(
                     x,
                     self.hc_attn_fn_broadcast,
@@ -527,6 +540,7 @@ def forward(
                 self.hc_attn_base,
                 norm_weight=self.input_layernorm.weight,
                 norm_eps=self.input_layernorm.variance_epsilon,
+                use_b12x_mhc=use_b12x_mhc,
             )
 
         # Attention needs the full token sequence; mHC above ran on the SP
@@ -553,6 +567,7 @@ def forward(
             self.hc_ffn_base,
             norm_weight=self.post_attention_layernorm.weight,
             norm_eps=self.post_attention_layernorm.variance_epsilon,
+            use_b12x_mhc=use_b12x_mhc,
         )
 
         # Fully Connected
@@ -565,7 +580,7 @@ def forward(
         # to fuse with) then contracts; every other layer defers its hc_post to
         # the next layer's fused pre, returning the state.
         if self.layer_idx == self.num_hidden_layers - 1:
-            x = self.hc_post(x, residual, post, comb)
+            x = self.hc_post(x, residual, post, comb, use_b12x_mhc=use_b12x_mhc)
             x = hc_contract(x, self.n)
             return x, None, None, None
 
@@ -595,14 +610,31 @@ def hc_pre(
         )
         return post_mix, res_mix, layer_input
 
+    def _use_b12x_mhc(self, x: torch.Tensor) -> bool:
+        """Select B12X for decode-sized batches, including graph padding.
+
+        TileLang handles larger prefill batches. The threshold is expressed in
+        rank-local tokens so sequence-parallel layers make the same selection
+        before and after attention.
+        """
+        return self._b12x_mhc is not None and x.shape[0] <= self._b12x_mhc_max_tokens
+
     def hc_post(
         self,
         x: torch.Tensor,
         residual: torch.Tensor,
         post: torch.Tensor,
         comb: torch.Tensor,
+        *,
+        use_b12x_mhc: bool | None = None,
     ):
-        if self._b12x_mhc is not None:
+        # Auxiliary-state capture completes deferred mHC state outside the
+        # decoder-layer forward, so it must derive the backend from the same
+        # rank-local token count used by the layer dispatch.
+        if use_b12x_mhc is None:
+            use_b12x_mhc = self._use_b12x_mhc(x)
+        if use_b12x_mhc:
+            assert self._b12x_mhc is not None
             return self._b12x_mhc.run_post(x, residual, post, comb)
         return self.mhc_post_op(x, residual, post, comb)
 
@@ -617,8 +649,11 @@ def hc_fused_post_pre(
         hc_base: torch.Tensor,
         norm_weight: torch.Tensor | None = None,
         norm_eps: float = 0.0,
+        *,
+        use_b12x_mhc: bool,
     ):
-        if self._b12x_mhc is not None:
+        if use_b12x_mhc:
+            assert self._b12x_mhc is not None
             return self._b12x_mhc.run_post_pre(
                 x,
                 residual,