--- baked-ours/distributed/communication_op.py	2026-09-10 19:56:14.000000000 +0200
+++ baked/distributed/communication_op.py	2026-09-11 10:13:47.000000000 +0200
@@ -9,8 +9,17 @@
 from .parallel_state import get_tp_group
 
 
+_DSV41_PREFILL_TIMING = __import__("os").environ.get("DSV41_PREFILL_TIMING", "0") == "1"
+
+
 def tensor_model_parallel_all_reduce(input_: torch.Tensor) -> torch.Tensor:
     """All-reduce the input tensor across model parallel group."""
+    if _DSV41_PREFILL_TIMING and input_.dim() >= 2 and input_.shape[0] >= 1024:
+        # Stage 7: phase "allreduce" of the prefill budget.
+        from vllm.model_executor.kernels.attention.dsa.candidate_blocks import _PrefillTimer
+
+        with _PrefillTimer("allreduce", int(input_.shape[0])):
+            return get_tp_group().all_reduce(input_)
     return get_tp_group().all_reduce(input_)