[ROCm][DSv4.1] Keep the Engram tables in host memory on ROCm (#57491)
JohnQinAMD committed
Sep 20, 2026 at 19:17 UTC
6dfd87f59658eba0af0354418ca725c64ef1ed7d
5 files changed
+20
-10
tests/kernels/test_engram.py
+6
-2
@@ -599,7 +599,9 @@ def test_engram_rejects_empty_head_shards(tp_size, dp_size, n_heads, monkeypatch
599
ParallelEngramEmbedding(n_heads * 17, 64, (17,) * n_heads)
600
601
602
-@pytest.mark.skipif(not current_platform.is_cuda(), reason="CUDA required")
602
+@pytest.mark.skipif(
603
+ not current_platform.is_cuda_alike(), reason="CUDA or ROCm required"
604
+)
605
@pytest.mark.parametrize("cpu_offload", [False, True])
606
@pytest.mark.parametrize("tp_size", [1, 2, 4, 8])
607
def test_engram_head_shards_reconstruct_checkpoint(cpu_offload, tp_size, monkeypatch):
@@ -701,7 +703,9 @@ def test_engram_lookup_reuses_jit_across_token_shapes():
703
assert cache_sizes == [1] * len(cache_sizes)
704
705
704
-@pytest.mark.skipif(not current_platform.is_cuda(), reason="CUDA required")
706
+@pytest.mark.skipif(
707
+ not current_platform.is_cuda_alike(), reason="CUDA or ROCm required"
708
+)
709
@pytest.mark.parametrize("cpu_offload", [False, True])
710
@pytest.mark.parametrize("background", [False, True])
711
@pytest.mark.parametrize("num_tokens", [1, 7, 256])
tests/test_config.py
+7
-5
@@ -1299,7 +1299,7 @@ def test_engram_dp_shared_memory_config_validation(
1299
monkeypatch, dp_size, load_format, multithread, error
1300
):
1301
"""Reject invalid shared configs before distributed init; allow threaded loads."""
1302
- monkeypatch.setattr(current_platform, "is_cuda", lambda: True)
1302
+ monkeypatch.setattr(current_platform, "is_cuda_alike", lambda: True)
1303
config = cast(
1304
VllmConfig,
1305
SimpleNamespace(
@@ -1324,7 +1324,7 @@ def test_engram_dp_shared_memory_config_validation(
1324
1325
1326
@pytest.mark.parametrize(
1327
- "architecture, ple_layers, cuda, supported",
1327
+ "architecture, ple_layers, accelerator, supported",
1328
[
1329
("DeepseekV41ForCausalLM", [1], True, True),
1330
("DeepseekV41ForCausalLM", [], True, False),
@@ -1339,9 +1339,11 @@ def test_engram_dp_shared_memory_config_validation(
1339
(None, None, True, False),
1340
],
1341
)
1342
-def test_engram_model_support(monkeypatch, architecture, ple_layers, cuda, supported):
1342
+def test_engram_model_support(
1343
+ monkeypatch, architecture, ple_layers, accelerator, supported
1344
+):
1345
"""A similarly named HF field must not enable unsupported implementations."""
1344
- monkeypatch.setattr(current_platform, "is_cuda", lambda: cuda)
1346
+ monkeypatch.setattr(current_platform, "is_cuda_alike", lambda: accelerator)
1347
model = (
1348
cast(
1349
ModelConfig,
@@ -1430,7 +1432,7 @@ def test_engram_explicit_config_requires_supported_model():
1432
@pytest.mark.parametrize("explicit", [False, True])
1433
def test_engram_draft_config_validates_target(monkeypatch, target_has_ple, explicit):
1434
"""MTP may inherit cross-DP sharding without having its own PLE layers."""
1433
- monkeypatch.setattr(current_platform, "is_cuda", lambda: True)
1435
+ monkeypatch.setattr(current_platform, "is_cuda_alike", lambda: True)
1436
target = SimpleNamespace(
1437
architecture="Qwen4ExpForCausalLM",
1438
hf_text_config=SimpleNamespace(ple_layer_ids=[1] if target_has_ple else []),
vllm/config/engram.py
+1
-1
@@ -74,7 +74,7 @@ class EngramConfig:
74
if (
75
model_config is None
76
or field is None
77
- or not current_platform.is_cuda()
77
+ or not current_platform.is_cuda_alike()
78
or not getattr(model_config.hf_text_config, field, None)
79
):
80
raise ValueError(
vllm/config/vllm.py
+1
-1
@@ -1260,7 +1260,7 @@ class VllmConfig:
1260
"Disable --enable-dbo and set --ubatch-size to 0."
1261
)
1262
if self.engram_config is None:
1263
- if not current_platform.is_cuda() or not model_has_engram_layers(
1263
+ if not current_platform.is_cuda_alike() or not model_has_engram_layers(
1264
model_config
1265
):
1266
return
vllm/models/deepseek_v41/amd/model.py
+5
-1
@@ -61,9 +61,13 @@ from vllm.utils.math_utils import cdiv
61
from vllm.v1.attention.backends.registry import AttentionBackendEnum
62
from vllm.v1.worker.ubatching import dbo_current_ubatch_id
63
64
-from ..common.engram import Engram, EngramLayout, NgramHashState
64
+from ..common.engram import EngramLayout, NgramHashState
65
from ..common.mm_preprocess import IMAGE_SENTINEL_BASE_ID, image_sentinel_mask
66
67
+# Engram host offload and its prefetch stream are neither ROCm- nor
68
+# NVIDIA-specific, so they are imported rather than duplicated.
69
+from ..nvidia.engram import Engram
70
+
71
if typing.TYPE_CHECKING:
72
from vllm.v1.attention.backends.mla.sparse_swa import DeepseekSparseSWAMetadata
73