diff --git a/csrc/cumem_allocator.cpp b/csrc/cumem_allocator.cpp index 6965d2124a..efb1f879d2 100644 --- a/csrc/cumem_allocator.cpp +++ b/csrc/cumem_allocator.cpp @@ -226,28 +226,6 @@ void create_and_map(unsigned long long device, ssize_t size, CUdeviceptr d_mem, return; } -#ifdef USE_ROCM - // ROCm-only: zero the freshly (re)mapped region. - // - // On some amdgpu driver versions the physical VRAM backing a buffer is not - // scrubbed on release, so the pages handed back here (e.g. on wake_up after a - // sleep that discarded a tag) can still contain stale data from a previously - // freed allocation -- possibly from another process. Callers of the device - // allocator expect newly mapped memory to be well defined; in particular - // hybrid Mamba / gated-delta-net (GDN) models keep a persisted conv/recurrent - // state cache, and a new sequence that reads an unzeroed slot will propagate - // garbage/NaN, collapsing generation to token id 0 (decoded as "!"). - // - // This is scoped to the ROCm path on purpose: the CUDA path is unchanged. - // Use CUDA_CHECK so a failure sets the global error_code; the Python caller - // (python_create_and_map) only inspects error_code, so returning without - // setting it would make a failed memset look successful and still hand back - // unzeroed memory. - CUDA_CHECK((CUresult)hipMemset(reinterpret_cast(d_mem), 0, size)); - if (error_code != 0) { - return; - } -#endif // std::cout << "create_and_map: device=" << device << ", size=" << size << ", // d_mem=" << d_mem << ", p_memHandle=" << p_memHandle << std::endl; } diff --git a/docs/features/sleep_mode.md b/docs/features/sleep_mode.md index a8aa15af7b..0db82de9d3 100644 --- a/docs/features/sleep_mode.md +++ b/docs/features/sleep_mode.md @@ -107,7 +107,7 @@ By default, CUDA graph memory stays on the GPU while asleep. With `sleep_mode_offload_cudagraph=True` (off by default), CUDA graphs are captured into a cuMem pool that sleep backs up to CPU memory at both levels and any wake restores in place, so graphs are replayed, not recaptured. It needs the default -`cumem` backend, CUDA and CUDA graphs; otherwise it has no effect. +`cumem` backend, CUDA or ROCm, and CUDA graphs; otherwise it has no effect. ```python llm = LLM("Qwen/Qwen3-8B", enable_sleep_mode=True, sleep_mode_offload_cudagraph=True) diff --git a/tests/basic_correctness/memory/cumem/test_cumem.py b/tests/basic_correctness/memory/cumem/test_cumem.py index df8c604882..c7ebfd2609 100644 --- a/tests/basic_correctness/memory/cumem/test_cumem.py +++ b/tests/basic_correctness/memory/cumem/test_cumem.py @@ -359,7 +359,7 @@ def test_runtime_state_survives_sleep(level): @create_new_process_for_each_test("fork" if current_platform.is_cuda() else "spawn") @pytest.mark.skipif(current_platform.is_xpu(), reason="Uses the CuMem allocator") def test_level2_discards_ordinary_tensor_with_weights_tag(): - """Reproduce the level-2 variant for an ordinary tensor in weights.""" + """Discarded weights-tag memory is remapped; ROCm zeroes it over stale pages.""" allocator = get_mem_allocator_instance() with allocator.use_memory_pool("weights"): @@ -374,8 +374,9 @@ def test_level2_discards_ordinary_tensor_with_weights_tag(): torch.accelerator.synchronize() assert (fake_weight.data_ptr(), ordinary_tensor.data_ptr()) == pointers - assert torch.all(fake_weight == 0xA5) - assert torch.all(ordinary_tensor == 0xA5) + expected = 0 if current_platform.is_rocm() else 0xA5 + assert torch.all(fake_weight == expected) + assert torch.all(ordinary_tensor == expected) @create_new_process_for_each_test("fork" if current_platform.is_cuda() else "spawn") @@ -427,7 +428,9 @@ def test_cumem_with_cudagraph(): @pytest.mark.parametrize("level", [1, 2], ids=["sleep-1", "sleep-2"]) @create_new_process_for_each_test("fork" if current_platform.is_cuda() else "spawn") -@pytest.mark.skipif(not current_platform.is_cuda(), reason="cuMem CUDA graph pool") +@pytest.mark.skipif( + not current_platform.is_cuda_alike(), reason="cuMem CUDA graph pool" +) def test_cudagraph_pool_sleep(level): """Routing, and the graph pool backed up at both sleep levels and restored in place by any wake.""" diff --git a/tests/basic_correctness/memory/sleep_mode/test_sleep_mode.py b/tests/basic_correctness/memory/sleep_mode/test_sleep_mode.py index 9f6180335a..d897d71001 100644 --- a/tests/basic_correctness/memory/sleep_mode/test_sleep_mode.py +++ b/tests/basic_correctness/memory/sleep_mode/test_sleep_mode.py @@ -209,8 +209,20 @@ def _cudagraph_bytes(worker) -> int: def _custom_ar_active(worker) -> bool: from vllm.distributed.parallel_state import get_tp_group - ca_comm = get_tp_group().device_communicator.ca_comm - return ca_comm is not None and not ca_comm.disabled + comm = get_tp_group().device_communicator + return any( + c is not None and not c.disabled for c in (comm.ca_comm, comm.aiter_ar_comm) + ) + + +def _capture_registers_buffers(worker) -> bool: + from vllm.distributed.parallel_state import get_tp_group + + comm = get_tp_group().device_communicator + ca, aiter = comm.ca_comm, comm.aiter_ar_comm + return (ca is not None and ca._capture_registered) or ( + aiter is not None and aiter.aiter_ca.enable_register_for_capturing + ) @pytest.mark.parametrize( @@ -227,7 +239,9 @@ def _custom_ar_active(worker) -> bool: ids=["full", "piecewise", "breakable", "custom-ar-tp2", "off-by-default"], ) @create_new_process_for_each_test() -@pytest.mark.skipif(not current_platform.is_cuda(), reason="cuMem CUDA graph pool") +@pytest.mark.skipif( + not current_platform.is_cuda_alike(), reason="cuMem CUDA graph pool" +) def test_sleep_cudagraph_pool(monkeypatch, mode, breakable, tp, offload): """Each capture site uses the pool, exact across sleeps; TP=2 needs unregistered custom AR; off changes nothing.""" @@ -255,8 +269,10 @@ def test_sleep_cudagraph_pool(monkeypatch, mode, breakable, tp, offload): assert all(graph_bytes) if offload else not any(graph_bytes) if not offload: return - if tp > 1 and not all(llm.collective_rpc(_custom_ar_active)): - pytest.skip("Custom allreduce is unavailable on these GPUs") + if tp > 1: + if not all(llm.collective_rpc(_custom_ar_active)): + pytest.skip("Custom allreduce is unavailable on these GPUs") + assert not any(llm.collective_rpc(_capture_registers_buffers)) prompt, params = "How are you?", SamplingParams(temperature=0, max_tokens=10) expected = llm.generate(prompt, params)[0].outputs[0].text for level in (1, 2): @@ -268,6 +284,42 @@ def test_sleep_cudagraph_pool(monkeypatch, mode, breakable, tp, offload): assert llm.generate(prompt, params)[0].outputs[0].text == expected +@multi_gpu_test(num_gpus=2) +@create_new_process_for_each_test() +@pytest.mark.skipif(not current_platform.is_rocm(), reason="AITER fused allreduce") +def test_sleep_cudagraph_pool_aiter_fused_ar_group_quant(monkeypatch): + """The fused AITER AR + RMSNorm + per-group FP8 quant captures into the + cuMem pool without IPC-registering its input, across sleeps.""" + for name, value in [ + ("VLLM_USE_V2_MODEL_RUNNER", "1"), + ("VLLM_ROCM_USE_AITER", "1"), + ("VLLM_ALLOW_INSECURE_SERIALIZATION", "1"), + ]: + monkeypatch.setenv(name, value) + monkeypatch.delenv("NCCL_GRAPH_REGISTER", raising=False) + llm = LLM( + "Qwen/Qwen3-0.6B-FP8", + enable_sleep_mode=True, + sleep_mode_offload_cudagraph=True, + tensor_parallel_size=2, + max_model_len=1024, + compilation_config={"pass_config": {"fuse_allreduce_rms": True}}, + ) + if not all(llm.collective_rpc(_custom_ar_active)): + pytest.skip("AITER custom allreduce is unavailable on these GPUs") + assert not any(llm.collective_rpc(_capture_registers_buffers)) + assert all(llm.collective_rpc(_cudagraph_bytes)) + prompt, params = "How are you?", SamplingParams(temperature=0, max_tokens=10) + assert llm.generate(prompt, params)[0].outputs[0].text + for level in (1, 2): + llm.sleep(level=level) + llm.wake_up(tags=["weights"]) + if level == 2: + llm.collective_rpc("reload_weights") + llm.wake_up(tags=["kv_cache"]) + assert llm.generate(prompt, params)[0].outputs[0].text + + @create_new_process_for_each_test() def test_deep_sleep_async(gpu_memory_cleared): async def test(): diff --git a/vllm/_aiter_ops.py b/vllm/_aiter_ops.py index 44b89aeaf1..0ed8490618 100644 --- a/vllm/_aiter_ops.py +++ b/vllm/_aiter_ops.py @@ -1167,7 +1167,8 @@ def _rocm_aiter_fused_allreduce_rmsnorm_quant_per_group_impl( w=weight, eps=epsilon, group_size=group_size, - registered=torch.cuda.is_current_stream_capturing(), + registered=torch.cuda.is_current_stream_capturing() + and ca.enable_register_for_capturing, use_1stage=use_1stage, ) assert result is not None @@ -1219,7 +1220,8 @@ def _rocm_aiter_fused_allreduce_rmsnorm_quant_per_group_with_bf16_norm_impl( w=weight, eps=epsilon, group_size=group_size, - registered=torch.cuda.is_current_stream_capturing(), + registered=torch.cuda.is_current_stream_capturing() + and ca.enable_register_for_capturing, use_1stage=use_1stage, emit_bf16=True, ) diff --git a/vllm/config/vllm.py b/vllm/config/vllm.py index 5d38b67653..bbda00af40 100644 --- a/vllm/config/vllm.py +++ b/vllm/config/vllm.py @@ -702,7 +702,7 @@ class VllmConfig: and model_config.enable_sleep_mode and model_config.sleep_mode_backend == "cumem" and self.compilation_config.cudagraph_mode != CUDAGraphMode.NONE - and current_platform.is_cuda() + and current_platform.is_cuda_alike() ) @property diff --git a/vllm/device_allocator/cumem.py b/vllm/device_allocator/cumem.py index b85be08d28..a2e582f85b 100644 --- a/vllm/device_allocator/cumem.py +++ b/vllm/device_allocator/cumem.py @@ -98,7 +98,7 @@ class CuMemAllocator: offloaded to CPU memory, and the rest of the tensors will be discarded. When we call `wake_up`, all tensors that are previously offloaded will be loaded back to GPU memory, and the rest of the tensors will - have empty memory. + have empty memory (zeroed on ROCm). Why it needs to be a singleton? When allocated tensors are garbage collected, PyTorch will call @@ -331,7 +331,7 @@ class CuMemAllocator: def wake_up(self, tags: list[str] | None = None) -> None: """Wake up the allocator from sleep mode. All data that is previously offloaded will be loaded back to GPU - memory, and the rest of the data will have empty memory. + memory, and the rest will have empty memory (zeroed on ROCm). Args: tags: The deferrable tags (weights, kv_cache) to load back to GPU @@ -358,6 +358,10 @@ class CuMemAllocator: cpu_ptr = cpu_backup_tensor.data_ptr() libcudart.cudaMemcpy(ptr, cpu_ptr, size_in_bytes) data.cpu_backup_tensor = None + elif current_platform.is_rocm(): + # amdgpu <6.14 may return stale VRAM (#44972); drop with its support + libcudart.cudaMemset(ptr, 0, handle[1]) + torch.accelerator.synchronize() @contextmanager def use_memory_pool(self, tag: str | None = None): diff --git a/vllm/distributed/device_communicators/aiter_custom_all_reduce.py b/vllm/distributed/device_communicators/aiter_custom_all_reduce.py index 6fbffcaaf7..9a5970c565 100644 --- a/vllm/distributed/device_communicators/aiter_custom_all_reduce.py +++ b/vllm/distributed/device_communicators/aiter_custom_all_reduce.py @@ -30,6 +30,8 @@ class AiterCustomAllreduce: group: ProcessGroup, device: int | str | torch.device, max_size: int | None = None, + *, + register_graph_buffers: bool = True, ): from aiter.dist.device_communicators.custom_all_reduce import ( CustomAllreduce as _AiterCustomAllreduce, @@ -38,7 +40,10 @@ class AiterCustomAllreduce: if max_size is None: max_size = self.MAX_SIZE - self._impl = _AiterCustomAllreduce(group, device, max_size=max_size) + kwargs = {} + if not register_graph_buffers: + kwargs["enable_register_for_capturing"] = False + self._impl = _AiterCustomAllreduce(group, device, max_size=max_size, **kwargs) @property def aiter_ca(self): diff --git a/vllm/distributed/device_communicators/cuda_communicator.py b/vllm/distributed/device_communicators/cuda_communicator.py index d4d588bf45..992d07e15d 100644 --- a/vllm/distributed/device_communicators/cuda_communicator.py +++ b/vllm/distributed/device_communicators/cuda_communicator.py @@ -151,6 +151,7 @@ class CudaCommunicator(DeviceCommunicatorBase): self.aiter_ar_comm = AiterCustomAllreduce( group=self.cpu_group, device=self.device, + register_graph_buffers=register, ) if use_custom_allreduce and self.aiter_ar_comm is None and self.world_size > 1: @@ -174,6 +175,7 @@ class CudaCommunicator(DeviceCommunicatorBase): self.aiter_ar_comm = AiterCustomAllreduce( group=self.cpu_group, device=self.device, + register_graph_buffers=register, ) if self.aiter_ar_comm.disabled: self.aiter_ar_comm = None