[ROCm] Stage large pageable H2D copies instead of registering them (#56343)
Signed-off-by: JohnQinAMD <yanyuan.qin@amd.com> Co-authored-by: Claude Opus 5 (1M context) <noreply@anthropic.com> Co-authored-by: Shanshan Shen <467638484@qq.com>
JohnQinAMD committed
Sep 17, 2026 at 00:14 UTC
528fa835fd37d0805a824d516e22f88df61a771a
1 file changed
+14
vllm/platforms/rocm.py
+14
@@ -21,6 +21,7 @@ from .interface import DeviceCapability, Platform, PlatformEnum, in_wsl
21
if TYPE_CHECKING:
22
from vllm.config import VllmConfig
23
from vllm.config.kernel import IrOpPriorityConfig
24
+ from vllm.utils.argparse_utils import FlexibleArgumentParser
25
from vllm.v1.attention.selector import AttentionSelectorConfig
26
27
logger = init_logger(__name__)
@@ -493,6 +494,9 @@ class RocmPlatform(Platform):
494
dist_backend: str = "nccl"
495
# rocm shares the same device control env var as CUDA
496
device_control_env_var: str = "CUDA_VISIBLE_DEVICES"
497
+ # Set in pre_register_and_update, so it exists only on the driver; Ray
498
+ # workers are separate processes and copy env vars by allowlist.
499
+ additional_env_vars: list[str] = ["GPU_PINNED_MIN_XFER_SIZE"]
500
ray_noset_device_env_vars: list[str] = [
501
"RAY_EXPERIMENTAL_NOSET_HIP_VISIBLE_DEVICES",
502
"RAY_EXPERIMENTAL_NOSET_CUDA_VISIBLE_DEVICES",
@@ -862,6 +866,16 @@ class RocmPlatform(Platform):
866
)
867
return torch.cuda.get_device_properties(device_id).total_memory
868
869
+ @classmethod
870
+ def pre_register_and_update(
871
+ cls, parser: "FlexibleArgumentParser | None" = None
872
+ ) -> None:
873
+ # Keep mmap'd weight pages on the HIP staging path: above this
874
+ # threshold the runtime registers the pageable source instead, and each
875
+ # registration's MMU notifier makes KFD suspend our queues. In KB, so
876
+ # 4 GiB.
877
+ os.environ.setdefault("GPU_PINNED_MIN_XFER_SIZE", str(4 * 1024 * 1024))
878
+
879
@classmethod
880
def apply_config_platform_defaults(cls, vllm_config: "VllmConfig") -> None:
881
from vllm._aiter_ops import rocm_aiter_ops