A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 3, 04:15 UTC
[Minor] Bigger overlap for FI AR (#43103)
Signed-off-by: Jee Jee Li <jeejeelee@inferact.ai>
Jee Jee Li committed
May 21, 2026 at 12:20 UTC
ee05e8137ec48b8e7375228a1142b4c5f2e3360c
2 files changed
+10
-1
vllm/compilation/passes/fusion/allreduce_rms_fusion.py
+5
@@ -44,6 +44,10 @@ from .matcher_utils import MatcherQuantFP8
44
45
FP8_DTYPE = current_platform.fp8_dtype()
46
47
+
48
+# The empirical value for small batch
49
+PDL_ADVANCE_LAUNCH_TOKENS = 16
50
+
51
logger = init_logger(__name__)
52
53
flashinfer_comm: ModuleType | None = None
@@ -204,6 +208,7 @@ if flashinfer_comm is not None:
208
layout_code=layout_code,
209
use_oneshot=use_oneshot,
210
fp32_acc=fp32_acc,
211
+ trigger_completion_at_end=num_tokens > PDL_ADVANCE_LAUNCH_TOKENS,
212
)
213
214
def call_trtllm_fused_allreduce_norm_fake(
vllm/distributed/device_communicators/flashinfer_all_reduce.py
+5
-1
@@ -19,6 +19,8 @@ from vllm.platforms import current_platform
19
20
logger = init_logger(__name__)
21
22
+# The empirical value for small batch
23
+PDL_ADVANCE_LAUNCH_TOKENS = 16
24
25
fi_ar_available = False
26
try:
@@ -312,7 +314,7 @@ class FlashInferAllReduce:
314
return self._ensure_workspace(hidden_dim, input_tensor.dtype)
315
316
def all_reduce(self, input_tensor: torch.Tensor) -> torch.Tensor:
315
- _, hidden_dim = input_tensor.shape
317
+ num_tokens, hidden_dim = input_tensor.shape
318
workspace = get_fi_ar_workspace(
319
world_size=self.world_size,
320
rank=self.rank,
@@ -325,6 +327,8 @@ class FlashInferAllReduce:
327
input=input_tensor,
328
workspace=workspace,
329
pattern=flashinfer_comm.AllReduceFusionPattern.kAllReduce,
330
+ launch_with_pdl=True,
331
+ trigger_completion_at_end=num_tokens > PDL_ADVANCE_LAUNCH_TOKENS,
332
)
333
334
def destroy(self):