launcher/dcp2/diffs/03-model_executor_layers_fused_moe_b12x.py.diff

--- baked-ours/model_executor/layers/fused_moe/b12x.py	2026-09-10 19:56:14.000000000 +0200
+++ baked/model_executor/layers/fused_moe/b12x.py	2026-09-11 10:13:47.000000000 +0200
@@ -755,16 +755,20 @@
         )
         scratch = _workspace_as_b12x_scratch(workspace2, plan)
 
-        _run_b12x_moe_plan(
-            plan=plan,
-            scratch=scratch,
-            hidden_states=hidden_states,
-            prepared=prepared,
-            topk_weights=topk_weights,
-            topk_ids=topk_ids,
-            output=output,
-            unit_scale_contract=self._quant_mode == "w4a16",
-        )
+        # Stage 7: phase "moe" of the prefill budget (inert unless DSV41_PREFILL_TIMING=1).
+        from vllm.model_executor.kernels.attention.dsa.candidate_blocks import _PrefillTimer
+
+        with _PrefillTimer("moe", int(hidden_states.shape[0])):
+            _run_b12x_moe_plan(
+                plan=plan,
+                scratch=scratch,
+                hidden_states=hidden_states,
+                prepared=prepared,
+                topk_weights=topk_weights,
+                topk_ids=topk_ids,
+                output=output,
+                unit_scale_contract=self._quant_mode == "w4a16",
+            )
 
     def moe_sum(self, input: torch.Tensor, output: torch.Tensor) -> None:
         raise NotImplementedError("LoRA is not supported for B12xExperts")