diff --git a/tests/distributed/test_weight_transfer.py b/tests/distributed/test_weight_transfer.py index 715a8069d7..8a13b24dc5 100644 --- a/tests/distributed/test_weight_transfer.py +++ b/tests/distributed/test_weight_transfer.py @@ -37,23 +37,18 @@ from vllm.platforms import current_platform from vllm.utils.network_utils import get_open_port -def _weight_transfer_ray_env_vars() -> dict[str, str]: - if not current_platform.is_rocm(): - return {} - - return { - "RAY_EXPERIMENTAL_NOSET_CUDA_VISIBLE_DEVICES": "1", - "RAY_EXPERIMENTAL_NOSET_HIP_VISIBLE_DEVICES": "1", - "RAY_EXPERIMENTAL_NOSET_ROCR_VISIBLE_DEVICES": "1", - } - - def _init_ray_for_weight_transfer() -> None: if ray.is_initialized(): return ray.init( ignore_reinit_error=True, - runtime_env={"env_vars": _weight_transfer_ray_env_vars()}, + runtime_env={ + "env_vars": { + "RAY_EXPERIMENTAL_NOSET_CUDA_VISIBLE_DEVICES": "1", + "RAY_EXPERIMENTAL_NOSET_HIP_VISIBLE_DEVICES": "1", + "RAY_EXPERIMENTAL_NOSET_ROCR_VISIBLE_DEVICES": "1", + } + }, ) @@ -66,7 +61,7 @@ def _get_ray_assigned_device() -> torch.device: def _set_ray_assigned_device() -> torch.device: device = _get_ray_assigned_device() - torch.accelerator.set_device_index(device) + current_platform.set_device(device) return device diff --git a/tests/v1/e2e/general/test_mamba_prefix_cache.py b/tests/v1/e2e/general/test_mamba_prefix_cache.py index 4644a6cc7e..dd9efc6696 100644 --- a/tests/v1/e2e/general/test_mamba_prefix_cache.py +++ b/tests/v1/e2e/general/test_mamba_prefix_cache.py @@ -403,6 +403,7 @@ def _run_ref_mamba_state_worker(): GPUModelRunner._sample = fake_sample_fn engine = LLM( model=MODEL, + load_format="dummy", block_size=BLOCK_SIZE, hf_overrides={"num_hidden_layers": NUM_HIDDEN_LAYERS}, seed=42, @@ -750,6 +751,7 @@ def test_mamba_prefix_cache_mrv1(monkeypatch: pytest.MonkeyPatch): engine = LLM( model=MODEL, + load_format="dummy", enable_prefix_caching=True, block_size=BLOCK_SIZE, mamba_cache_mode="align", diff --git a/vllm/model_executor/layers/fused_moe/routed_experts.py b/vllm/model_executor/layers/fused_moe/routed_experts.py index 1b1deb250c..ce87d0b1b3 100644 --- a/vllm/model_executor/layers/fused_moe/routed_experts.py +++ b/vllm/model_executor/layers/fused_moe/routed_experts.py @@ -907,7 +907,7 @@ class RoutedExperts(PluggableLayer): # Unified loading logic for fused and non-fused experts loaded_experts = experts_shard.unbind() for expert_id, loaded_expert in enumerate(loaded_experts, start=start): - success = self.weight_loader( + success = param.weight_loader( param=param, loaded_weight=loaded_expert, weight_name=weight_name, diff --git a/vllm/model_executor/models/deepseek_v2.py b/vllm/model_executor/models/deepseek_v2.py index 211c88129c..9530ac58b6 100644 --- a/vllm/model_executor/models/deepseek_v2.py +++ b/vllm/model_executor/models/deepseek_v2.py @@ -1445,6 +1445,8 @@ class DeepseekV2Model(nn.Module): hidden_states, residual = combined_states.split( [self.hidden_size, self.hidden_size], dim=-1 ) + # fused_add_rms_norm requires a contiguous residual + residual = residual.contiguous() if idx in self.aux_hidden_state_layers: aux_hidden_state = hidden_states + residual if aux_hidden_state.shape[0] != positions.shape[0]: @@ -1469,6 +1471,8 @@ class DeepseekV2Model(nn.Module): hidden_states, residual = combined_states.split( [self.hidden_size, self.hidden_size], dim=-1 ) + # fused_add_rms_norm requires a contiguous residual + residual = residual.contiguous() if self.end_layer in self.aux_hidden_state_layers: aux_hidden_states.append(hidden_states + residual) diff --git a/vllm/model_executor/models/gemma3_mm.py b/vllm/model_executor/models/gemma3_mm.py index 9e58438f4c..0d9f8f1418 100644 --- a/vllm/model_executor/models/gemma3_mm.py +++ b/vllm/model_executor/models/gemma3_mm.py @@ -761,6 +761,7 @@ class Gemma3ForConditionalGeneration( max_frames_per_batch: int, device: torch.device, dtype: torch.dtype, + path: str = "default", ): from vllm.v1.worker.encoder_cudagraph_defs import ( EncoderCudaGraphCaptureInputs, @@ -792,6 +793,7 @@ class Gemma3ForConditionalGeneration( mm_kwargs: dict[str, Any], max_batch_size: int, max_frames_per_batch: int, + path: str = "default", ): from vllm.v1.worker.encoder_cudagraph_defs import ( EncoderCudaGraphReplayBuffers, @@ -804,6 +806,7 @@ class Gemma3ForConditionalGeneration( def encoder_cudagraph_forward( self, values: dict[str, torch.Tensor], + path: str = "default", ) -> torch.Tensor: pixel_values = values["pixel_values"] image_features = self.vision_tower(pixel_values)