@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 4b3803d180
Fork 0

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 9, 01:15 UTC

[BugFix] DPMetadata raises assert error for dense model (#32739)

Co-authored-by: Dezhan Tu <dztu@meta.com>

Dezhan committed Feb 2, 2026 at 16:56 UTC 4b3803d18044499962f331db0a1614b726110e2a
1 file changed +5 -3
vllm/forward_context.py
+5 -3
@@ -203,7 +203,7 @@ class ForwardContext:
203 attn_metadata: dict[str, AttentionMetadata] | list[dict[str, AttentionMetadata]]
204 slot_mapping: dict[str, torch.Tensor] | list[dict[str, torch.Tensor]]
205 """
206 - Type Dict[str, AttentionMetadata] for v1, map from layer_name of each
206 + Type Dict[str, AttentionMetadata] for v1, map from layer_name of each
207 attention layer to its attention metadata
208 Type List[Dict[str, AttentionMetadata]] for DBO. List of size two, one
209 for each microbatch.
@@ -339,8 +339,10 @@ def set_forward_context(
339 forward_start_time = time.perf_counter()
340
341 dp_metadata: DPMetadata | None = None
342 - if vllm_config.parallel_config.data_parallel_size > 1 and (
343 - attn_metadata is not None or num_tokens is not None
342 + if (
343 + vllm_config.parallel_config.data_parallel_size > 1
344 + and vllm_config.parallel_config.is_moe_model is not False
345 + and (attn_metadata is not None or num_tokens is not None)
346 ):
347 # If num_tokens_across_dp hasn't already been initialized, then
348 # initialize it here. Both DP padding and Microbatching will be