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 14, 05:30 UTC
[BUGFIX] Add cu_tokens_across_sp to DPMetadata (#26457)
Signed-off-by: Sage Moore <sage@neuralmagic.com>
Sage Moore committed
Oct 9, 2025 at 10:13 UTC
0d37450eb7e38ac82df35d4e0f21d4254435049d
1 file changed
+11
vllm/forward_context.py
+11
@@ -161,6 +161,17 @@ class DPMetadata:
161
assert self.local_sizes is not None
162
return self.local_sizes
163
164
+ # Get the cumulative tokens across sequence parallel ranks.
165
+ # In this case the input to the MoEs will be distributed w.r.t both
166
+ # DP and TP rank.
167
+ # When sp_size==1, this is just the cummulative num tokens across DP.
168
+ def cu_tokens_across_sp(self, sp_size: int) -> torch.Tensor:
169
+ num_tokens_across_sp_cpu = (
170
+ self.num_tokens_across_dp_cpu - 1 + sp_size
171
+ ) // sp_size
172
+ num_tokens_across_sp_cpu = num_tokens_across_sp_cpu.repeat_interleave(sp_size)
173
+ return torch.cumsum(num_tokens_across_sp_cpu, dim=0)
174
+
175
176
@dataclass
177
class ForwardContext: