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 5, 06:00 UTC
Bug/test eagle dp v0 (#38938)
Signed-off-by: Monishver Chandrasekaran <monishverchandrasekaran@gmail.com> Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com> Co-authored-by: Nicolò Lucchesi <nlucches@redhat.com>
Monishver committed
Apr 13, 2026 at 16:50 UTC
8213e8f88023726799ae48e2ebef06ef0c021128
3 files changed
+13
-8
.buildkite/test_areas/distributed.yaml
+1
@@ -196,6 +196,7 @@ steps:
196
- VLLM_ALLOW_INSECURE_SERIALIZATION=1 python3 examples/rl/rlhf_async_new_apis.py
197
- VLLM_USE_DEEP_GEMM=1 VLLM_LOGGING_LEVEL=DEBUG python3 examples/offline_inference/data_parallel.py --model=Qwen/Qwen1.5-MoE-A2.7B -tp=1 -dp=2 --max-model-len=2048 --all2all-backend=deepep_high_throughput
198
- pytest -v -s tests/v1/distributed/test_dbo.py
199
+ - TP_SIZE=1 DP_SIZE=2 pytest -v -s tests/v1/distributed/test_eagle_dp.py
200
201
- label: Distributed Tests (2 GPUs)(B200)
202
device: b200
tests/v1/distributed/test_eagle_dp.py
+6
-8
@@ -20,16 +20,14 @@ if current_platform.is_rocm():
20
else:
21
ATTN_BACKENDS = ["FLASH_ATTN"]
22
23
+# On SM<90 (e.g., L4), batch invariance does not support CUDA graphs.
24
+# See https://github.com/vllm-project/vllm/pull/30018 and
25
+# tests/v1/determinism/utils.py for the documented limitation.
26
+IS_DEVICE_CAPABILITY_BELOW_90 = not current_platform.has_device_capability(90)
27
+
28
29
@pytest.mark.asyncio
30
@pytest.mark.parametrize("attn_backend", ATTN_BACKENDS)
26
-@pytest.mark.xfail(
27
- not current_platform.is_rocm(),
28
- reason="EAGLE + DP > 1 produces wrong outputs when async spec decode "
29
- "correction is active. Root cause under investigation. "
30
- "See: https://github.com/vllm-project/vllm/issues/31913",
31
- strict=False,
32
-)
31
@pytest.mark.xfail(
32
current_platform.is_rocm(),
33
reason="Test may fail on ROCm until batch invariance is enabled. "
@@ -57,7 +55,7 @@ async def test_run_eagle_dp(monkeypatch: pytest.MonkeyPatch, attn_backend: str):
55
engine_args = AsyncEngineArgs(
56
model=target_model,
57
tokenizer_mode="auto",
60
- enforce_eager=False,
58
+ enforce_eager=IS_DEVICE_CAPABILITY_BELOW_90,
59
tensor_parallel_size=int(os.getenv("TP_SIZE", 1)),
60
data_parallel_size=DP_SIZE,
61
data_parallel_backend="mp", # ray takes more time
vllm/model_executor/layers/vocab_parallel_embedding.py
+6
@@ -8,6 +8,7 @@ import torch
8
import torch.nn.functional as F
9
from torch.nn.parameter import Parameter, UninitializedParameter
10
11
+import vllm.envs as envs
12
from vllm.distributed import (
13
divide,
14
get_tensor_model_parallel_rank,
@@ -15,6 +16,9 @@ from vllm.distributed import (
16
tensor_model_parallel_all_reduce,
17
)
18
from vllm.model_executor.custom_op import PluggableLayer
19
+from vllm.model_executor.layers.batch_invariant import (
20
+ linear_batch_invariant,
21
+)
22
from vllm.model_executor.layers.quantization.base_config import (
23
QuantizationConfig,
24
QuantizeMethodBase,
@@ -66,6 +70,8 @@ class UnquantizedEmbeddingMethod(QuantizeMethodBase):
70
x: torch.Tensor,
71
bias: torch.Tensor | None = None,
72
) -> torch.Tensor:
73
+ if envs.VLLM_BATCH_INVARIANT and current_platform.is_cuda_alike():
74
+ return linear_batch_invariant(x, layer.weight, bias)
75
return dispatch_unquantized_gemm()(layer, x, layer.weight, bias)
76
77
def embedding(self, layer: torch.nn.Module, input_: torch.Tensor) -> torch.Tensor: