From 6dfd87f59658eba0af0354418ca725c64ef1ed7d Mon Sep 17 00:00:00 2001 From: JohnQinAMD Date: Sun, 20 Sep 2026 19:17:11 -0700 Subject: [PATCH] [ROCm][DSv4.1] Keep the Engram tables in host memory on ROCm (#57491) --- tests/kernels/test_engram.py | 8 ++++++-- tests/test_config.py | 12 +++++++----- vllm/config/engram.py | 2 +- vllm/config/vllm.py | 2 +- vllm/models/deepseek_v41/amd/model.py | 6 +++++- 5 files changed, 20 insertions(+), 10 deletions(-) diff --git a/tests/kernels/test_engram.py b/tests/kernels/test_engram.py index c2a7607cba..45cebcfe57 100644 --- a/tests/kernels/test_engram.py +++ b/tests/kernels/test_engram.py @@ -599,7 +599,9 @@ def test_engram_rejects_empty_head_shards(tp_size, dp_size, n_heads, monkeypatch ParallelEngramEmbedding(n_heads * 17, 64, (17,) * n_heads) -@pytest.mark.skipif(not current_platform.is_cuda(), reason="CUDA required") +@pytest.mark.skipif( + not current_platform.is_cuda_alike(), reason="CUDA or ROCm required" +) @pytest.mark.parametrize("cpu_offload", [False, True]) @pytest.mark.parametrize("tp_size", [1, 2, 4, 8]) 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(): assert cache_sizes == [1] * len(cache_sizes) -@pytest.mark.skipif(not current_platform.is_cuda(), reason="CUDA required") +@pytest.mark.skipif( + not current_platform.is_cuda_alike(), reason="CUDA or ROCm required" +) @pytest.mark.parametrize("cpu_offload", [False, True]) @pytest.mark.parametrize("background", [False, True]) @pytest.mark.parametrize("num_tokens", [1, 7, 256]) diff --git a/tests/test_config.py b/tests/test_config.py index 9da6eced94..f31adfdb8d 100644 --- a/tests/test_config.py +++ b/tests/test_config.py @@ -1299,7 +1299,7 @@ def test_engram_dp_shared_memory_config_validation( monkeypatch, dp_size, load_format, multithread, error ): """Reject invalid shared configs before distributed init; allow threaded loads.""" - monkeypatch.setattr(current_platform, "is_cuda", lambda: True) + monkeypatch.setattr(current_platform, "is_cuda_alike", lambda: True) config = cast( VllmConfig, SimpleNamespace( @@ -1324,7 +1324,7 @@ def test_engram_dp_shared_memory_config_validation( @pytest.mark.parametrize( - "architecture, ple_layers, cuda, supported", + "architecture, ple_layers, accelerator, supported", [ ("DeepseekV41ForCausalLM", [1], True, True), ("DeepseekV41ForCausalLM", [], True, False), @@ -1339,9 +1339,11 @@ def test_engram_dp_shared_memory_config_validation( (None, None, True, False), ], ) -def test_engram_model_support(monkeypatch, architecture, ple_layers, cuda, supported): +def test_engram_model_support( + monkeypatch, architecture, ple_layers, accelerator, supported +): """A similarly named HF field must not enable unsupported implementations.""" - monkeypatch.setattr(current_platform, "is_cuda", lambda: cuda) + monkeypatch.setattr(current_platform, "is_cuda_alike", lambda: accelerator) model = ( cast( ModelConfig, @@ -1430,7 +1432,7 @@ def test_engram_explicit_config_requires_supported_model(): @pytest.mark.parametrize("explicit", [False, True]) def test_engram_draft_config_validates_target(monkeypatch, target_has_ple, explicit): """MTP may inherit cross-DP sharding without having its own PLE layers.""" - monkeypatch.setattr(current_platform, "is_cuda", lambda: True) + monkeypatch.setattr(current_platform, "is_cuda_alike", lambda: True) target = SimpleNamespace( architecture="Qwen4ExpForCausalLM", hf_text_config=SimpleNamespace(ple_layer_ids=[1] if target_has_ple else []), diff --git a/vllm/config/engram.py b/vllm/config/engram.py index 5abb2bd34a..773ca20911 100644 --- a/vllm/config/engram.py +++ b/vllm/config/engram.py @@ -74,7 +74,7 @@ class EngramConfig: if ( model_config is None or field is None - or not current_platform.is_cuda() + or not current_platform.is_cuda_alike() or not getattr(model_config.hf_text_config, field, None) ): raise ValueError( diff --git a/vllm/config/vllm.py b/vllm/config/vllm.py index 78e30cf9d0..6974d1fbb3 100644 --- a/vllm/config/vllm.py +++ b/vllm/config/vllm.py @@ -1260,7 +1260,7 @@ class VllmConfig: "Disable --enable-dbo and set --ubatch-size to 0." ) if self.engram_config is None: - if not current_platform.is_cuda() or not model_has_engram_layers( + if not current_platform.is_cuda_alike() or not model_has_engram_layers( model_config ): return diff --git a/vllm/models/deepseek_v41/amd/model.py b/vllm/models/deepseek_v41/amd/model.py index b3eddc2024..1383bac820 100644 --- a/vllm/models/deepseek_v41/amd/model.py +++ b/vllm/models/deepseek_v41/amd/model.py @@ -61,9 +61,13 @@ from vllm.utils.math_utils import cdiv from vllm.v1.attention.backends.registry import AttentionBackendEnum from vllm.v1.worker.ubatching import dbo_current_ubatch_id -from ..common.engram import Engram, EngramLayout, NgramHashState +from ..common.engram import EngramLayout, NgramHashState from ..common.mm_preprocess import IMAGE_SENTINEL_BASE_ID, image_sentinel_mask +# Engram host offload and its prefetch stream are neither ROCm- nor +# NVIDIA-specific, so they are imported rather than duplicated. +from ..nvidia.engram import Engram + if typing.TYPE_CHECKING: from vllm.v1.attention.backends.mla.sparse_swa import DeepseekSparseSWAMetadata