--- 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")