[Bugfix] Pin EPLB and MLA host-to-device transfer buffers (#56138)
Signed-off-by: khluu <khluu000@gmail.com> Signed-off-by: Kevin H. Luu <khluu000@gmail.com> Co-authored-by: OpenAI Codex <noreply@openai.com>
Kevin H. Luu committed
Sep 11, 2026 at 01:57 UTC
1cf6555214446cfb98c42ff5d07baab784925c1b
4 files changed
+81
-5
tests/distributed/test_eplb_execute.py
+25
@@ -7,6 +7,8 @@ import pytest
7
import torch
8
import torch.distributed
9
10
+import vllm.distributed.eplb.eplb_communicator as eplb_comm
11
+import vllm.utils.gpu_sync_debug as gsd
12
from vllm.config import VllmConfig, set_current_vllm_config
13
from vllm.distributed.eplb.eplb_communicator import (
14
create_eplb_communicator,
@@ -25,6 +27,29 @@ from vllm.distributed.parallel_state import (
27
from .eplb_utils import distributed_run, set_env_vars_and_device
28
29
30
+@pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA")
31
+def test_gloo_receive_staging_does_not_force_gpu_sync(monkeypatch):
32
+ """Gloo receives must reach the GPU without an implicit pageable copy."""
33
+ monkeypatch.setattr(gsd, "_SYNC_CHECK_MODE", "error")
34
+ monkeypatch.setattr(gsd, "_sync_check_enabled", True)
35
+ gsd._install_copy_checkers()
36
+ monkeypatch.setattr(eplb_comm, "is_local_first_rank", lambda: False)
37
+
38
+ monkeypatch.setattr(eplb_comm, "P2POp", lambda op, tensor, peer, group: tensor)
39
+
40
+ def receive(tensors):
41
+ for tensor in tensors:
42
+ tensor.fill_(7)
43
+ return []
44
+
45
+ monkeypatch.setattr(eplb_comm, "batch_isend_irecv", receive)
46
+ communicator = eplb_comm.TorchDistGlooStagedEplbCommunicator(cpu_group=None)
47
+ dst = torch.empty(32, device="cuda")
48
+ communicator.add_recv([dst], src_rank=1, expert_id=0)
49
+ gsd.with_gpu_sync_check(communicator.execute)()
50
+ torch.testing.assert_close(dst.cpu(), torch.full((32,), 7.0))
51
+
52
+
53
def create_expert_indices_with_redundancy(
54
num_layers: int,
55
num_logical_experts: int,
tests/v1/attention/test_mla_context_chunks.py
+40
@@ -8,14 +8,20 @@ chunk never covers a prefill without context (which is why the partial no
8
longer needs an empty-span masking pass).
9
"""
10
11
+from types import SimpleNamespace
12
+
13
import pytest
14
import torch
15
16
+import vllm.utils.gpu_sync_debug as gsd
17
from vllm.model_executor.layers.attention.mla_attention import (
18
build_mla_chunked_context_metadata,
19
init_mla_context_partial,
20
reorg_kvcache,
21
)
22
+from vllm.model_executor.layers.attention.sparse_mla_attention import (
23
+ SparseMLACommonMetadataBuilder,
24
+)
25
26
BLOCK_SIZE = 16
27
@@ -45,6 +51,40 @@ def build_chunked_context(
51
)
52
53
54
+@pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA")
55
+def test_sparse_context_lengths_do_not_force_gpu_sync(monkeypatch):
56
+ """Sparse MLA must pin computed context lengths before the H2D copy."""
57
+ builder = SimpleNamespace(
58
+ chunked_prefill_workspace=torch.empty((2048, 1)),
59
+ chunked_prefill_workspace_size=1024,
60
+ kv_cache_spec=SimpleNamespace(block_size=BLOCK_SIZE),
61
+ device=torch.device("cuda"),
62
+ dcp_world_size=1,
63
+ dcp_local_block_size=1,
64
+ dcp_virtual_block_size=1,
65
+ dcp_manager=None,
66
+ )
67
+ common_metadata = SimpleNamespace(
68
+ seq_lens_cpu_upper_bound=torch.tensor([17, 2052, 324], dtype=torch.int32),
69
+ query_start_loc_cpu=torch.tensor([0, 1, 5, 9], dtype=torch.int32),
70
+ )
71
+ monkeypatch.setattr(gsd, "_SYNC_CHECK_MODE", "error")
72
+ monkeypatch.setattr(gsd, "_sync_check_enabled", True)
73
+ gsd._install_copy_checkers()
74
+ metadata = gsd.with_gpu_sync_check(
75
+ SparseMLACommonMetadataBuilder._build_chunked_context_fields
76
+ )(
77
+ builder,
78
+ common_metadata,
79
+ num_decodes=1,
80
+ num_prefills=2,
81
+ prefill_query_lens_cpu=torch.tensor([4, 4], dtype=torch.int32),
82
+ )
83
+ assert metadata is not None
84
+ assert metadata.context_lens.device.type == "cuda"
85
+ assert metadata.context_lens.tolist() == [2048, 320]
86
+
87
+
88
@pytest.mark.parametrize(
89
"context_lens,workspace_size",
90
[
vllm/distributed/eplb/eplb_communicator.py
+4
-1
@@ -34,6 +34,7 @@ from vllm.distributed.utils import is_weak_contiguous
34
from vllm.logger import init_logger
35
from vllm.platforms import current_platform
36
from vllm.utils.gpu_sync_debug import gpu_sync_allowed
37
+from vllm.utils.torch_utils import PIN_MEMORY
38
39
logger = init_logger(__name__)
40
@@ -204,7 +205,9 @@ class TorchDistGlooStagedEplbCommunicator(EplbCommunicator):
205
)
206
)
207
continue
207
- cpu_tensor = torch.empty_like(tensor, device="cpu")
208
+ cpu_tensor = torch.empty_like(
209
+ tensor, device="cpu", pin_memory=PIN_MEMORY
210
+ )
211
p2p_ops.append(
212
P2POp(
213
torch.distributed.irecv,
vllm/model_executor/layers/attention/sparse_mla_attention.py
+12
-4
@@ -28,7 +28,11 @@ from vllm.model_executor.layers.attention.mla_attention import (
28
from vllm.platforms import current_platform
29
from vllm.triton_utils import tl, triton
30
from vllm.utils.flashinfer import has_flashinfer
31
-from vllm.utils.torch_utils import is_quantized_kv_cache, np_to_pinned_tensor
31
+from vllm.utils.torch_utils import (
32
+ PIN_MEMORY,
33
+ is_quantized_kv_cache,
34
+ np_to_pinned_tensor,
35
+)
36
from vllm.v1.attention.backend import AttentionMetadata, AttentionMetadataBuilder
37
from vllm.v1.attention.backends.fa_utils import get_flash_attn_version
38
from vllm.v1.attention.backends.utils import split_decodes_and_prefills
@@ -241,9 +245,13 @@ class SparseMLACommonMetadataBuilder(AttentionMetadataBuilder[T]):
245
246
seq_lens_cpu = common_attn_metadata.seq_lens_cpu_upper_bound
247
assert seq_lens_cpu is not None
244
- context_lens_cpu = (
245
- seq_lens_cpu[num_decodes : num_decodes + num_prefills]
246
- - prefill_query_lens_cpu
248
+ context_lens_cpu = torch.empty(
249
+ num_prefills, dtype=seq_lens_cpu.dtype, pin_memory=PIN_MEMORY
250
+ )
251
+ torch.subtract(
252
+ seq_lens_cpu[num_decodes : num_decodes + num_prefills],
253
+ prefill_query_lens_cpu,
254
+ out=context_lens_cpu,
255
)
256
qsl_cpu = common_attn_metadata.query_start_loc_cpu
257
prefill_query_start_loc_cpu = qsl_cpu[num_decodes:] - qsl_cpu[num_decodes]