[Bugfix][Engram] Keep THP tables private when resolving shared memory (#59068)
Signed-off-by: Ren Yuzhou <54501155+yuzhouo7@users.noreply.github.com> Co-authored-by: Claude <noreply@anthropic.com>
Ren Yuzhou committed
Sep 29, 2026 at 13:19 UTC
ec5e0c352f079c2cb8f46752fd0317a2745050bf
2 files changed
+27
-5
tests/test_config.py
+26
-5
@@ -1510,21 +1510,42 @@ def test_engram_dp_shared_memory_requires_cpu_offload():
1510
EngramConfig(cpu_offload=False, dp_shared_memory=True)
1511
1512
1513
+@pytest.mark.skip_global_cleanup
1514
@pytest.mark.parametrize(
1514
- "cpu_offload,dp_size,elastic_ep,expected",
1515
- [(True, 2, False, True), (False, 2, False, False), (True, 1, False, False)],
1515
+ "cpu_offload,use_thp,dp_shared_memory,dp_size,elastic_ep,expected",
1516
+ [
1517
+ (True, False, None, 2, False, True),
1518
+ (False, False, None, 2, False, False),
1519
+ (True, False, None, 1, False, False),
1520
+ (True, False, None, 2, True, False),
1521
+ # use_thp backs private tables; sharing would silently ignore it.
1522
+ (True, True, None, 2, False, False),
1523
+ (True, True, False, 2, False, False),
1524
+ ],
1525
)
1526
def test_engram_dp_shared_memory_defaults_when_supported(
1518
- cpu_offload, dp_size, elastic_ep, expected
1527
+ cpu_offload, use_thp, dp_shared_memory, dp_size, elastic_ep, expected
1528
):
1520
- """Unset dp_shared_memory enables sharing only for offloaded, non-elastic DP."""
1529
+ """Resolve sharing only where supported, preserving an explicit false."""
1530
parallel = ParallelConfig(data_parallel_size=dp_size)
1531
parallel.enable_elastic_ep = elastic_ep
1523
- config = EngramConfig(cpu_offload=cpu_offload)
1532
+ config = EngramConfig(
1533
+ cpu_offload=cpu_offload,
1534
+ use_thp=use_thp,
1535
+ dp_shared_memory=dp_shared_memory,
1536
+ )
1537
config.resolve_dp_shared_memory(parallel)
1538
assert config.dp_shared_memory is expected
1539
1540
1541
+@pytest.mark.skip_global_cleanup
1542
+def test_engram_thp_rejects_explicit_shared_memory():
1543
+ with pytest.raises(
1544
+ ValueError, match="use_thp requires cpu_offload=True and dp_shared_memory=False"
1545
+ ):
1546
+ EngramConfig(use_thp=True, dp_shared_memory=True)
1547
+
1548
+
1549
@pytest.mark.parametrize(
1550
"dp_size,load_format,multithread,error",
1551
[
vllm/config/engram.py
+1
@@ -91,6 +91,7 @@ class EngramConfig:
91
if self.dp_shared_memory is None:
92
self.dp_shared_memory = (
93
self.cpu_offload
94
+ and not self.use_thp
95
and parallel_config.data_parallel_size > 1
96
and not parallel_config.enable_elastic_ep
97
)