| 1234567891011121314151617181920 |
- --- 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_)
-
-
|
