@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 6cdf015c3c

Fork 0

A high-throughput and memory-efficient inference and serving engine for LLMs

Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git

Synced Nov 10, 11:30 UTC

[Misc] Fix `Current vLLM config is not set.` warnings, assert to avoid issues in the future (#31747)

Signed-off-by: Lucas Wilkinson <lwilkins@redhat.com> Signed-off-by: Lucas Wilkinson <LucasWilkinson@users.noreply.github.com> Co-authored-by: Luka Govedič <ProExpertProg@users.noreply.github.com>

Lucas Wilkinson committed Jan 8, 2026 at 18:20 UTC 6cdf015c3cd8ee600dadf17ecd72fc7f44099d03

48 files changed +380 -240

tests/compile/distributed/test_async_tp.py

+29 -24

@@ -15,6 +15,7 @@ from vllm.config import (

15 ModelConfig,

16 PassConfig,

17 VllmConfig,

18 + set_current_vllm_config,

19 )

20 from vllm.distributed import (

21 tensor_model_parallel_all_gather,

@@ -340,38 +341,42 @@ def async_tp_pass_on_test_model(

341 )

342

343 async_tp_pass = AsyncTPPass(vllm_config)

343 - backend = TestBackend(async_tp_pass)

344

345 - assert (

346 - async_tp_pass.compilation_config.splitting_ops

347 - == vllm_config.compilation_config.splitting_ops

348 - )

349 - assert (

350 - async_tp_pass.compilation_config.use_inductor_graph_partition

351 - == vllm_config.compilation_config.use_inductor_graph_partition

352 - )

345 + # Set the global vllm_config for TestBackend which calls

346 + # get_current_vllm_config()

347 + with set_current_vllm_config(vllm_config):

348 + backend = TestBackend(async_tp_pass)

349

354 - model = test_model_cls(hidden_size, dtype) # Pass dtype to model constructor

350 + assert (

351 + async_tp_pass.compilation_config.splitting_ops

352 + == vllm_config.compilation_config.splitting_ops

353 + )

354 + assert (

355 + async_tp_pass.compilation_config.use_inductor_graph_partition

356 + == vllm_config.compilation_config.use_inductor_graph_partition

357 + )

358

356 - hidden_states = torch.randn(

357 - (batch_size * seq_len, hidden_size), dtype=dtype, requires_grad=False

358 - )

359 + model = test_model_cls(hidden_size, dtype) # Pass dtype to model constructor

360 +

361 + hidden_states = torch.randn(

362 + (batch_size * seq_len, hidden_size), dtype=dtype, requires_grad=False

363 + )

364

360 - if dynamic:

361 - torch._dynamo.mark_dynamic(hidden_states, 0)

365 + if dynamic:

366 + torch._dynamo.mark_dynamic(hidden_states, 0)

367

363 - compiled_model = torch.compile(model, backend=backend)

364 - compiled_model(hidden_states)

368 + compiled_model = torch.compile(model, backend=backend)

369 + compiled_model(hidden_states)

370

366 - assert async_tp_pass.matched_count == 1

371 + assert async_tp_pass.matched_count == 1

372

368 - # In pre-nodes, all gather or reduce scatter should exist,

369 - # fused_matmul_reduce_scatter or fused_all_gather_matmul should not

370 - backend.check_before_ops(model.ops_in_model_before(), fully_replaced=False)

373 + # In pre-nodes, all gather or reduce scatter should exist,

374 + # fused_matmul_reduce_scatter or fused_all_gather_matmul should not

375 + backend.check_before_ops(model.ops_in_model_before(), fully_replaced=False)

376

372 - # In post-nodes, fused_matmul_reduce_scatter or \

373 - # fused_all_gather_matmul should exist

374 - backend.check_after_ops(model.ops_in_model_after())

377 + # In post-nodes, fused_matmul_reduce_scatter or \

378 + # fused_all_gather_matmul should exist

379 + backend.check_after_ops(model.ops_in_model_after())

380

381

382 @create_new_process_for_each_test()

tests/compile/test_config.py

+1 -1

@@ -430,7 +430,7 @@ def test_cudagraph_sizes_post_init(

430 )

431

432

433 -def test_cached_compilation_config():

433 +def test_cached_compilation_config(default_vllm_config):

434 import torch

435 from torch._inductor.utils import run_and_get_code

436

tests/conftest.py

+11

@@ -189,6 +189,17 @@ def dist_init():

189 cleanup_dist_env_and_memory()

190

191

192 +@pytest.fixture

193 +def default_vllm_config():

194 + """Set a default VllmConfig for tests that directly test CustomOps or pathways

195 + that use get_current_vllm_config() outside of a full engine context.

196 + """

197 + from vllm.config import VllmConfig, set_current_vllm_config

198 +

199 + with set_current_vllm_config(VllmConfig()):

200 + yield

201 +

202 +

203 @pytest.fixture()

204 def should_do_global_cleanup_after_test(request) -> bool:

205 """Allow subdirectories to skip global cleanup by overriding this fixture.

tests/kernels/attention/test_flashinfer_trtllm_attention.py

+1 -1

@@ -458,7 +458,7 @@ def test_flashinfer_trtllm_prefill_with_baseline(

458 )

459

460

461 -def test_trtllm_attention_rejects_num_kv_heads_1() -> None:

461 +def test_trtllm_attention_rejects_num_kv_heads_1(default_vllm_config) -> None:

462 """Test that TRTLLM attention correctly rejects num_kv_heads=1.

463

464 When num_kv_heads=1 (MQA), the KV cache strides become degenerate

tests/kernels/attention/test_mha_attn.py

+3 -1

@@ -36,7 +36,7 @@ if current_platform.is_rocm():

38 @pytest.mark.parametrize("device", devices)

39 -def test_mha_attn_platform(device: str):

39 +def test_mha_attn_platform(default_vllm_config, device: str):

40 """

41 Test the attention selector between different platform and device.

42 """

@@ -116,6 +116,7 @@ CUDA_DEVICES = ["cuda"]

116 @pytest.mark.parametrize("dtype", DTYPES)

117 @pytest.mark.parametrize("device", CUDA_DEVICES)

118 def test_mha_attn_forward(

119 + default_vllm_config,

120 batch_size: int,

121 seq_len: int,

122 num_heads: int,

@@ -162,6 +163,7 @@ def test_mha_attn_forward(

163 @pytest.mark.parametrize("dtype", DTYPES)

164 @pytest.mark.parametrize("device", CUDA_DEVICES)

165 def test_mha_attn_varlen_forward(

166 + default_vllm_config,

167 var_seq_len: list[int],

168 num_heads: int,

169 num_kv_heads: int,

tests/kernels/core/test_activation.py

@@ -45,6 +45,7 @@ CUDA_DEVICES = [f"cuda:{i}" for i in range(1 if torch.cuda.device_count() == 1 e

45 @pytest.mark.parametrize("device", CUDA_DEVICES)

46 @torch.inference_mode()

47 def test_act_and_mul(

48 + default_vllm_config,

49 activation: str,

50 num_tokens: int,

51 d: int,

@@ -122,6 +123,7 @@ def test_act_and_mul(

123 @pytest.mark.parametrize("device", CUDA_DEVICES)

124 @torch.inference_mode()

125 def test_activation(

126 + default_vllm_config,

127 activation: type[torch.nn.Module],

128 num_tokens: int,

129 d: int,

tests/kernels/core/test_fused_qk_norm_rope.py

@@ -57,6 +57,7 @@ def _apply_qk_norm_rope(

57 @pytest.mark.parametrize("rotary_ratio", [1.0, 0.5, 0.25])

58 @torch.inference_mode()

59 def test_fused_qk_norm_rope_matches_reference(

60 + default_vllm_config,

61 device: str,

62 dtype: torch.dtype,

63 is_neox: bool,

tests/kernels/core/test_fused_quant_layernorm.py

@@ -147,6 +147,7 @@ def ops_impl(

147 @pytest.mark.parametrize("device", CUDA_DEVICES)

148 @torch.inference_mode()

149 def test_rms_norm(

150 + default_vllm_config,

151 num_tokens: int,

152 hidden_size: int,

153 add_residual: bool,

tests/kernels/core/test_layernorm.py

@@ -26,6 +26,7 @@ CUDA_DEVICES = [f"cuda:{i}" for i in range(1 if torch.cuda.device_count() == 1 e

26 @pytest.mark.parametrize("strided_input", [False, True])

27 @torch.inference_mode()

28 def test_rms_norm(

29 + default_vllm_config,

30 num_tokens: int,

31 hidden_size: int,

32 add_residual: bool,

tests/kernels/core/test_mrope.py

@@ -90,6 +90,7 @@ num_tokens_list = [11, 8192]

90 @pytest.mark.parametrize("dtype", [torch.bfloat16])

91 @pytest.mark.parametrize("num_tokens", num_tokens_list)

92 def test_mrope(

93 + default_vllm_config,

94 model_name: str,

95 model_info: MRoPETestInfo,

96 tp_size: int,

@@ -159,6 +160,7 @@ def test_mrope(

160 @pytest.mark.parametrize("dtype", [torch.bfloat16])

161 @pytest.mark.parametrize("num_tokens", num_tokens_list)

162 def test_mrope_torch_compile_tracing(

163 + default_vllm_config,

164 model_name: str,

165 model_info: MRoPETestInfo,

166 tp_size: int,

tests/kernels/core/test_pos_encoding.py

+2 -1

@@ -62,6 +62,7 @@ TENSORS_SHAPES_FN = [

62 @pytest.mark.parametrize("use_key", USE_KEY)

63 @torch.inference_mode()

64 def test_rotary_embedding(

65 + default_vllm_config,

66 is_neox_style: bool,

67 tensor_shape_fn: Callable[[int, int, int, int], tuple[int, ...]],

68 batch_size: int,

@@ -123,7 +124,7 @@ def test_rotary_embedding(

124

125

126 @torch.inference_mode()

126 -def test_rope_module_cache():

127 +def test_rope_module_cache(default_vllm_config):

128 MAX_POSITIONS = [123, 1234]

129 ROPE_THETAS = [10000, 1000000]

130 ROPE_PARAMETERS = (

tests/kernels/core/test_rotary_embedding.py

@@ -36,6 +36,7 @@ def rotary_embedding_opcheck(

36 @pytest.mark.parametrize("use_key", [True, False])

37 @pytest.mark.parametrize("head_stride_is_contiguous", [True, False])

38 def test_rotary_embedding_opcheck(

39 + default_vllm_config,

40 dist_init,

41 device,

42 max_position,

tests/kernels/moe/test_cpu_fused_moe.py

+2 -6

@@ -6,7 +6,7 @@ import torch

7 from tests.kernels.allclose_default import get_default_atol, get_default_rtol

8 from vllm._custom_ops import cpu_fused_moe, cpu_prepack_moe_weight

9 -from vllm.model_executor.layers.activation import SiluAndMul, SwigluOAIAndMul

9 +from vllm.model_executor.layers.fused_moe.cpu_fused_moe import _CPU_MOE_ACT

10 from vllm.platforms import current_platform

11 from vllm.utils.torch_utils import set_random_seed

@@ -24,11 +24,6 @@ USE_BIAS = [True, False]

24 ISA = ["amx", "vec"] if torch._C._cpu._is_amx_tile_supported() else ["vec"]

25 DTYPE = [torch.bfloat16]

27 -_CPU_MOE_ACT = {

28 - "silu": SiluAndMul(),

29 - "swigluoai": SwigluOAIAndMul(),

30 -}

31 -

28 def ref_fused_moe(

29 input: torch.Tensor,

@@ -106,6 +101,7 @@ def ref_fused_moe(

101 @pytest.mark.parametrize("act", ACT)

102 @pytest.mark.parametrize("isa", ISA)

103 def test_cpu_fused_moe(

104 + default_vllm_config,

105 batch_size: int,

106 expert_num: int,

107 hidden_size: int,

tests/kernels/moe/test_moe.py

+6 -1

@@ -468,7 +468,12 @@ def test_fused_moe_wn16(

468 )

469 @torch.inference_mode()

470 def test_mixtral_moe(

471 - dist_init, dtype: torch.dtype, padding: bool, use_rocm_aiter: bool, monkeypatch

471 + default_vllm_config,

472 + dist_init,

473 + dtype: torch.dtype,

474 + padding: bool,

475 + use_rocm_aiter: bool,

476 + monkeypatch,

477 ):

478 """Make sure our Mixtral MoE implementation agrees with the one from

479 huggingface."""

tests/kernels/quantization/test_fp8_quant_group.py

+10 -3

@@ -23,7 +23,12 @@ from vllm.utils.torch_utils import set_random_seed

23 @pytest.mark.parametrize("use_ue8m0", [True, False])

24 @torch.inference_mode()

25 def test_quantfp8_group_functionality(

26 - batch_size: int, hidden_dim: int, group_size: int, seed: int, use_ue8m0: bool

26 + default_vllm_config,

27 + batch_size: int,

28 + hidden_dim: int,

29 + group_size: int,

30 + seed: int,

31 + use_ue8m0: bool,

32 ) -> None:

33 """Test QuantFP8 group quantization with various configurations.

@@ -82,7 +87,9 @@ def test_quantfp8_group_functionality(

87 @pytest.mark.parametrize("seed", [42])

88 @pytest.mark.parametrize("use_ue8m0", [True, False])

89 @torch.inference_mode()

85 -def test_quantfp8_group_multidimensional(seed: int, use_ue8m0: bool) -> None:

90 +def test_quantfp8_group_multidimensional(

91 + default_vllm_config, seed: int, use_ue8m0: bool

92 +) -> None:

93 set_random_seed(seed)

95 group_size = 64

@@ -135,7 +142,7 @@ def test_quantfp8_group_multidimensional(seed: int, use_ue8m0: bool) -> None:

142

143 @pytest.mark.parametrize("seed", [42])

144 @torch.inference_mode()

138 -def test_quantfp8_group_edge_cases(seed: int) -> None:

145 +def test_quantfp8_group_edge_cases(default_vllm_config, seed: int) -> None:

146 set_random_seed(seed)

147

148 batch_size = 16

tests/kernels/quantization/test_int8_kernel.py

+1 -1

@@ -102,7 +102,7 @@ SEEDS = [0]

102 itertools.product(M, N, K, E, TOP_KS, DTYPES, SEEDS),

103 )

104 @torch.inference_mode()

105 -def test_w8a8_fp8_fused_moe(M, N, K, E, topk, dtype, seed):

105 +def test_w8a8_fp8_fused_moe(default_vllm_config, M, N, K, E, topk, dtype, seed):

106 torch.manual_seed(seed)

107 # Initialize int8 quantization parameters

108 factor_for_scale = 1e-2

tests/kernels/quantization/test_silu_mul_nvfp4_quant.py

@@ -31,6 +31,7 @@ BLOCK_SIZE = 16

31 @pytest.mark.parametrize("shape", SHAPES)

32 @torch.inference_mode()

33 def test_silu_mul_nvfp4_quant(

34 + default_vllm_config,

35 dtype: torch.dtype,

36 shape: tuple[int, int],

37 ) -> None:

tests/kernels/test_fused_quant_activation.py

@@ -39,6 +39,7 @@ def ops_impl(x: torch.Tensor, scale: torch.Tensor) -> torch.Tensor:

39 @pytest.mark.parametrize("device", CUDA_DEVICES)

40 @torch.inference_mode()

41 def test_silu_and_mul(

42 + default_vllm_config,

43 num_tokens: int,

44 hidden_size: int,

45 dtype: torch.dtype,

tests/lora/conftest.py

+2 -2

@@ -82,7 +82,7 @@ class DummyLoRAModel(nn.Sequential, SupportsLoRA):

84 @pytest.fixture

85 -def dummy_model() -> nn.Module:

85 +def dummy_model(default_vllm_config) -> nn.Module:

86 model = DummyLoRAModel(

87 OrderedDict(

88 [

@@ -115,7 +115,7 @@ def dummy_model() -> nn.Module:

115

116

117 @pytest.fixture

118 -def dummy_model_gate_up() -> nn.Module:

118 +def dummy_model_gate_up(default_vllm_config) -> nn.Module:

119 model = DummyLoRAModel(

120 OrderedDict(

121 [

tests/lora/test_layers.py

+8 -5

@@ -252,7 +252,9 @@ def check_punica_wrapper(punica_wrapper) -> bool:

252 @pytest.mark.parametrize("device", DEVICES)

253 @pytest.mark.parametrize("vocab_size", [512, 32000, 64000, 128000])

254 @pytest.mark.parametrize("stage", STAGES)

255 -def test_embeddings(dist_init, num_loras, device, vocab_size, stage) -> None:

255 +def test_embeddings(

256 + default_vllm_config, dist_init, num_loras, device, vocab_size, stage

257 +) -> None:

258 # For multi-GPU testing of Triton kernel, we must explicitly set the CUDA

259 # device, see: https://github.com/triton-lang/triton/issues/2925

260 # Same below.

@@ -353,7 +355,7 @@ def test_embeddings(dist_init, num_loras, device, vocab_size, stage) -> None:

355 @pytest.mark.parametrize("vocab_size", [512, 32000, 64000, 256512])

356 @pytest.mark.parametrize("stage", STAGES)

357 def test_lm_head_logits_processor(

356 - dist_init, num_loras, device, vocab_size, stage

358 + default_vllm_config, dist_init, num_loras, device, vocab_size, stage

359 ) -> None:

360 if current_platform.is_cuda_alike():

361 torch.cuda.set_device(device)

@@ -470,6 +472,7 @@ def test_lm_head_logits_processor(

472 @pytest.mark.parametrize("device", DEVICES)

473 @pytest.mark.parametrize("stage", STAGES)

474 def test_linear_replicated(

475 + default_vllm_config,

476 dist_init,

477 num_loras,

478 device,

@@ -580,7 +583,7 @@ def test_linear_replicated(

583 @pytest.mark.parametrize("device", DEVICES)

584 @pytest.mark.parametrize("stage", STAGES)

585 def test_linear_parallel(

583 - dist_init, num_loras, orientation, fully_shard, device, stage

586 + default_vllm_config, dist_init, num_loras, orientation, fully_shard, device, stage

587 ) -> None:

588 if current_platform.is_cuda_alike():

589 torch.cuda.set_device(device)

@@ -705,7 +708,7 @@ def test_linear_parallel(

708 @pytest.mark.parametrize("device", DEVICES)

709 @pytest.mark.parametrize("stage", STAGES)

710 def test_column_parallel_packed(

708 - dist_init, num_loras, repeats, fully_shard, device, stage

711 + default_vllm_config, dist_init, num_loras, repeats, fully_shard, device, stage

712 ) -> None:

713 if current_platform.is_cuda_alike():

714 torch.cuda.set_device(device)

@@ -851,7 +854,7 @@ def test_column_parallel_packed(

854 @pytest.mark.parametrize(

855 "seed", list(range(VOCAB_PARALLEL_EMBEDDING_TEST_NUM_RANDOM_SEEDS))

856 )

854 -def test_vocab_parallel_embedding_indices(tp_size, seed):

857 +def test_vocab_parallel_embedding_indices(tp_size, seed, default_vllm_config):

858 random.seed(seed)

859 vocab_size = random.randint(4000, 64000)

860 added_vocab_size = random.randint(0, 1024)

tests/lora/test_lora_manager.py

+13 -7

@@ -111,7 +111,7 @@ def create_packed_lora(

111 return LoRAModel(lora_id, 8, loras)

112

113

114 -def test_replace_submodules(dist_init, dummy_model):

114 +def test_replace_submodules(default_vllm_config, dist_init, dummy_model):

115 model = dummy_model

116 manager = LoRAModelManager(

117 model,

@@ -133,7 +133,7 @@ def test_replace_submodules(dist_init, dummy_model):

133

134

135 @pytest.mark.parametrize("device", DEVICES)

136 -def test_lora_model_manager(dist_init, dummy_model, device):

136 +def test_lora_model_manager(default_vllm_config, dist_init, dummy_model, device):

137 model = dummy_model

138 model_lora1 = create_lora(

139 1, model, ["layer1.dense1", "dense2", "lm_head"], device=device

@@ -199,7 +199,9 @@ def test_lora_model_manager(dist_init, dummy_model, device):

199

200

201 @pytest.mark.parametrize("device", DEVICES)

202 -def test_lora_lru_cache_model_manager(dist_init, dummy_model, device):

202 +def test_lora_lru_cache_model_manager(

203 + default_vllm_config, dist_init, dummy_model, device

204 +):

205 model = dummy_model

206 model_lora1 = create_lora(

207 1, model, ["layer1.dense1", "dense2", "lm_head"], device=device

@@ -289,7 +291,7 @@ def test_lora_lru_cache_model_manager(dist_init, dummy_model, device):

291

292

293 @pytest.mark.parametrize("device", DEVICES)

292 -def test_lru_lora_model_manager(dist_init, dummy_model, device):

294 +def test_lru_lora_model_manager(default_vllm_config, dist_init, dummy_model, device):

295 # This tests just the LRU cache functionality, everything else is

296 # tested in test_lora_model_manager

297 model = dummy_model

@@ -415,7 +417,9 @@ def test_lru_lora_model_manager(dist_init, dummy_model, device):

417

418

419 @pytest.mark.parametrize("device", DEVICES)

418 -def test_lru_cache_worker_adapter_manager(dist_init, dummy_model, device, tmp_path):

420 +def test_lru_cache_worker_adapter_manager(

421 + default_vllm_config, dist_init, dummy_model, device, tmp_path

422 +):

423 lora_config = LoRAConfig(

424 max_lora_rank=8, max_cpu_loras=4, max_loras=4, lora_dtype=DEFAULT_DTYPE

425 )

@@ -529,7 +533,9 @@ def test_lru_cache_worker_adapter_manager(dist_init, dummy_model, device, tmp_pa

533

534

535 @pytest.mark.parametrize("device", DEVICES)

532 -def test_worker_adapter_manager(dist_init, dummy_model_gate_up, device, tmp_path):

536 +def test_worker_adapter_manager(

537 + default_vllm_config, dist_init, dummy_model_gate_up, device, tmp_path

538 +):

539 # Should remove every LoRA not specified in the request.

540 lora_config = LoRAConfig(

541 max_lora_rank=8, max_cpu_loras=4, max_loras=4, lora_dtype=DEFAULT_DTYPE

@@ -636,7 +642,7 @@ def test_worker_adapter_manager(dist_init, dummy_model_gate_up, device, tmp_path

642

643

644 @pytest.mark.parametrize("device", DEVICES)

639 -def test_packed_loras(dist_init, dummy_model_gate_up, device):

645 +def test_packed_loras(default_vllm_config, dist_init, dummy_model_gate_up, device):

646 model = dummy_model_gate_up

647 model_lora = create_packed_lora(

648 1,

tests/model_executor/test_eagle_quantization.py

+1 -1

@@ -55,7 +55,7 @@ def test_get_draft_quant_config_without_draft_model():

56 @torch.inference_mode()

57 @pytest.mark.parametrize("device", DEVICES)

58 -def test_fc_layer_quant_config_usage(dist_init, device) -> None:

58 +def test_fc_layer_quant_config_usage(default_vllm_config, dist_init, device) -> None:

59 import torch

61 from vllm.model_executor.layers.linear import ReplicatedLinear

tests/models/multimodal/pooling/test_intern_vit.py

+3 -1

@@ -73,7 +73,9 @@ def run_intern_vit_test(

73 ],

74 )

75 @pytest.mark.parametrize("dtype", ["half"])

76 -def test_models(dist_init, image_assets, model_id, dtype: str) -> None:

76 +def test_models(

77 + default_vllm_config, dist_init, image_assets, model_id, dtype: str

78 +) -> None:

79 run_intern_vit_test(

80 image_assets,

81 model_id,

tests/models/multimodal/pooling/test_radio.py

+3 -1

@@ -92,7 +92,9 @@ def run_radio_test(

92 ],

93 )

94 @pytest.mark.parametrize("dtype", ["half", "bfloat16"])

95 -def test_radio(dist_init, image_assets, model_id, dtype: str) -> None:

95 +def test_radio(

96 + default_vllm_config, dist_init, image_assets, model_id, dtype: str

97 +) -> None:

98 run_radio_test(

99 image_assets,

100 model_id,

tests/models/multimodal/processing/test_tensor_schema.py

+9 -9

@@ -145,18 +145,18 @@ def initialize_dummy_model(

145 model_config: ModelConfig,

146 ):

147 temp_file = tempfile.mkstemp()[1]

148 - init_distributed_environment(

149 - world_size=1,

150 - rank=0,

151 - distributed_init_method=f"file://{temp_file}",

152 - local_rank=0,

153 - backend="nccl",

154 - )

155 - initialize_model_parallel(tensor_model_parallel_size=1)

156 -

148 current_device = torch.get_default_device()

149 vllm_config = VllmConfig(model_config=model_config)

150 with set_current_vllm_config(vllm_config=vllm_config):

151 + init_distributed_environment(

152 + world_size=1,

153 + rank=0,

154 + distributed_init_method=f"file://{temp_file}",

155 + local_rank=0,

156 + backend="nccl",

157 + )

158 + initialize_model_parallel(tensor_model_parallel_size=1)

159 +

160 with set_default_torch_dtype(model_config.dtype):

161 torch.set_default_device(current_platform.device_type)

162 model = model_cls(vllm_config=vllm_config)

tests/plugins_tests/test_platform_plugins.py

+1 -1

@@ -31,7 +31,7 @@ def test_platform_plugins():

31 )

34 -def test_oot_custom_op(monkeypatch: pytest.MonkeyPatch):

34 +def test_oot_custom_op(default_vllm_config, monkeypatch: pytest.MonkeyPatch):

35 # simulate workload by running an example

36 load_general_plugins()

37 from vllm.model_executor.layers.rotary_embedding import RotaryEmbedding

tests/quantization/test_fp8.py

@@ -277,6 +277,7 @@ def test_scaled_fp8_quant(dtype) -> None:

277 # this is the case for marlin as well as per-tensor Fp8MoEMethod

278 @pytest.mark.parametrize("use_marlin", [False]) # skip True

279 def test_fp8_reloading(

280 + default_vllm_config,

281 method_cls,

282 is_checkpoint_fp8_serialized,

283 weight_block_size,

tests/utils.py

+27 -6

@@ -721,13 +721,34 @@ def init_test_distributed_environment(

721 distributed_init_port: str,

722 local_rank: int = -1,

723 ) -> None:

724 - distributed_init_method = f"tcp://localhost:{distributed_init_port}"

725 - init_distributed_environment(

726 - world_size=pp_size * tp_size,

727 - rank=rank,

728 - distributed_init_method=distributed_init_method,

729 - local_rank=local_rank,

724 + # Note: This function is often called from Ray worker processes, so we

725 + # can't rely on pytest fixtures to set the config. We check if the config

726 + # is already set and only create a default one if needed.

727 + from vllm.config import (

728 + VllmConfig,

729 + get_current_vllm_config_or_none,

730 + set_current_vllm_config,

731 )

732 +

733 + distributed_init_method = f"tcp://localhost:{distributed_init_port}"

734 +

735 + if get_current_vllm_config_or_none() is not None:

736 + # Config already set, use it directly

737 + init_distributed_environment(

738 + world_size=pp_size * tp_size,

739 + rank=rank,

740 + distributed_init_method=distributed_init_method,

741 + local_rank=local_rank,

742 + )

743 + else:

744 + # No config set, create a default one for the test

745 + with set_current_vllm_config(VllmConfig()):

746 + init_distributed_environment(

747 + world_size=pp_size * tp_size,

748 + rank=rank,

749 + distributed_init_method=distributed_init_method,

750 + local_rank=local_rank,

751 + )

752 ensure_model_parallel_initialized(tp_size, pp_size)

753

754

tests/v1/attention/test_attention_backends.py

+1 -1

@@ -556,7 +556,7 @@ def _test_backend_correctness(

556 @pytest.mark.parametrize("model", ["meta-llama/Meta-Llama-3-8B"])

557 @pytest.mark.parametrize("tensor_parallel_size", [1, 2, 4])

558 def test_causal_backend_correctness(

559 - batch_spec_name: str, model: str, tensor_parallel_size: int

559 + default_vllm_config, batch_spec_name: str, model: str, tensor_parallel_size: int

560 ):

561 """Test backend's correctness with causal attention."""

562

tests/v1/attention/test_attention_backends_selection.py

+6 -1

@@ -79,7 +79,12 @@ from vllm.v1.attention.backends.short_conv_attn import ShortConvAttentionBackend

79 ],

80 )

81 def test_mamba_layers_get_attn_backend(

82 - dist_init, layer_class, init_kwargs, expected_backend, expected_mamba_type

82 + default_vllm_config,

83 + dist_init,

84 + layer_class,

85 + init_kwargs,

86 + expected_backend,

87 + expected_mamba_type,

88 ):

89 """Test that Mamba-like layers return the correct attention backend."""

90 layer = layer_class(**init_kwargs)

tests/v1/attention/test_mla_backends.py

+5 -1

@@ -394,7 +394,11 @@ def run_attention_backend(

394 @pytest.mark.parametrize("model", ["deepseek-ai/DeepSeek-R1"])

395 @pytest.mark.parametrize("tensor_parallel_size", [1, 4, 8, 16])

396 def test_backend_correctness(

397 - dist_init, batch_spec_name: str, model: str, tensor_parallel_size: int

397 + default_vllm_config,

398 + dist_init,

399 + batch_spec_name: str,

400 + model: str,

401 + tensor_parallel_size: int,

402 ):

403 """

404 Test that all backends produce similar outputs to a reference implementation

tests/v1/attention/test_sparse_mla_backends.py

+6 -1

@@ -124,7 +124,12 @@ def _quantize_dequantize_fp8_ds_mla(

124 reason="FlashMLASparseBackend requires CUDA 9.0 or higher",

125 )

126 def test_sparse_backend_decode_correctness(

127 - dist_init, batch_name, kv_cache_dtype, tensor_parallel_size, workspace_init

127 + default_vllm_config,

128 + dist_init,

129 + batch_name,

130 + kv_cache_dtype,

131 + tensor_parallel_size,

132 + workspace_init,

133 ):

134 if current_platform.is_rocm():

135 pytest.skip("ROCm does not support fp8_ds_mla data type for kv cache.")

tests/v1/determinism/test_rms_norm_batch_invariant.py

+12 -6

@@ -21,7 +21,11 @@ from vllm.model_executor.layers.layernorm import RMSNorm

21 @pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16])

22 @pytest.mark.parametrize("eps", [1e-6, 1e-5])

23 def test_rms_norm_batch_invariant_vs_standard(

24 - batch_size: int, hidden_size: int, dtype: torch.dtype, eps: float

24 + default_vllm_config,

25 + batch_size: int,

26 + hidden_size: int,

27 + dtype: torch.dtype,

28 + eps: float,

29 ):

30 """

31 Compare batch-invariant Triton RMS norm against standard CUDA implementation.

@@ -68,7 +72,9 @@ def test_rms_norm_batch_invariant_vs_standard(

72 @pytest.mark.parametrize("batch_size", [1, 16, 128])

73 @pytest.mark.parametrize("seq_len", [1, 32, 512])

74 @pytest.mark.parametrize("hidden_size", [2048, 4096])

71 -def test_rms_norm_3d_input(batch_size: int, seq_len: int, hidden_size: int):

75 +def test_rms_norm_3d_input(

76 + default_vllm_config, batch_size: int, seq_len: int, hidden_size: int

77 +):

78 """

79 Test RMS norm with 3D input tensors (batch, seq_len, hidden_size).

@@ -107,7 +113,7 @@ def test_rms_norm_3d_input(batch_size: int, seq_len: int, hidden_size: int):

113

114

115 @skip_unsupported

110 -def test_rms_norm_numerical_stability():

116 +def test_rms_norm_numerical_stability(default_vllm_config):

117 """

118 Test RMS norm numerical stability with extreme values.

119

@@ -167,7 +173,7 @@ def test_rms_norm_numerical_stability():

173

174

175 @skip_unsupported

170 -def test_rms_norm_formula():

176 +def test_rms_norm_formula(default_vllm_config):

177 """

178 Test that RMS norm follows the correct mathematical formula.

179

@@ -201,7 +207,7 @@ def test_rms_norm_formula():

207

208 @skip_unsupported

209 @pytest.mark.parametrize("hidden_size", [128, 1024, 4096, 16384])

204 -def test_rms_norm_different_hidden_sizes(hidden_size: int):

210 +def test_rms_norm_different_hidden_sizes(default_vllm_config, hidden_size: int):

211 """

212 Test RMS norm with various hidden sizes to ensure block size handling.

213

@@ -238,7 +244,7 @@ def test_rms_norm_different_hidden_sizes(hidden_size: int):

244

245

246 @skip_unsupported

241 -def test_rms_norm_determinism():

247 +def test_rms_norm_determinism(default_vllm_config):

248 """

249 Test that batch-invariant RMS norm produces deterministic results.

250

tests/v1/kv_connector/unit/test_nixl_connector.py

+93 -83

@@ -299,6 +299,7 @@ def test_prompt_less_than_block_size():

299 )

300 def test_kv_transfer_handshake(dist_init):

301 """Unit test for basic NixlConnector interface functionality."""

302 + from vllm.config import set_current_vllm_config

303

304 # Test setup, we creates a scheduler that contains a NixlConnector

305 # of role SCHEDULER, and expect it to be serving NixlAgentMetadata from

@@ -308,81 +309,82 @@ def test_kv_transfer_handshake(dist_init):

309 vllm_config.kv_transfer_config.kv_buffer_device = "cpu"

310 scheduler = create_scheduler(vllm_config)

311

311 - # Create two NixlConnector of role WORKER, one is the worker of

312 - # the scheduler (prefill), the other is a worker of decode instance.

312 + with set_current_vllm_config(vllm_config):

313 + # Create two NixlConnector of role WORKER, one is the worker of

314 + # the scheduler (prefill), the other is a worker of decode instance.

315

314 - # Prefill connector will register KV cache to populate proper handshake

315 - # metadata.

316 - prefill_connector = NixlConnector(vllm_config, KVConnectorRole.WORKER)

317 - kv_cache_shape = FlashAttentionBackend.get_kv_cache_shape(

318 - num_blocks=2, block_size=16, num_kv_heads=4, head_size=64

319 - )

320 - shared_tensor = torch.zeros(*kv_cache_shape, dtype=torch.float16)

321 - unique_tensor = torch.zeros(*kv_cache_shape, dtype=torch.float16)

322 - kv_caches = {

323 - "layer0": shared_tensor,

324 - "layer1": unique_tensor,

325 - "layer2": shared_tensor,

326 - }

327 - prefill_connector.register_kv_caches(kv_caches)

328 -

329 - # Simulate EngineCore initialization that would gather connector

330 - # metadata from all workers

331 - metadata = prefill_connector.get_handshake_metadata()

332 -

333 - # metadata is a NixlHandshakePayload, decode it to get NixlAgentMetadata

334 - decoder = msgspec.msgpack.Decoder(NixlAgentMetadata)

335 - expected_agent_metadata = decoder.decode(metadata.agent_metadata_bytes)

336 -

337 - # The scheduler connector expects metadata to be in

338 - # dict[int, KVConnectorHandshakeMetadata], where the first key is

339 - # the dp_rank, the second key is the tp_rank.

340 - scheduler_connector = scheduler.get_kv_connector()

341 - scheduler_connector.set_xfer_handshake_metadata({0: metadata})

342 -

343 - # Simulate a request that finishes prefill, which returns

344 - # corresponding NixlConnectorMetadata for decode instance.

345 - BLOCK_SIZE = vllm_config.cache_config.block_size

346 - NUM_EXTERNAL_FULL_BLOCKS = 2

347 - NUM_TOKENS = int(BLOCK_SIZE * (NUM_EXTERNAL_FULL_BLOCKS + 0.5))

348 -

349 - request = create_request(

350 - request_id=1,

351 - block_size=BLOCK_SIZE,

352 - num_tokens=NUM_TOKENS,

353 - do_remote_decode=True,

354 - )

355 - request.status = RequestStatus.FINISHED_LENGTH_CAPPED

356 - delay, kv_connector_metadata = scheduler.get_kv_connector().request_finished(

357 - request, [0, 1, 2]

358 - )

359 - assert delay

360 -

361 - # Decode connector will be able to create handshake with the prefill connector.

362 - decode_connector = NixlConnector(vllm_config, KVConnectorRole.WORKER)

363 -

364 - # Here we are testing the retrieval of NIXLAgentMetadata.

365 - # Knowing the implementation detail, we override the add_remote_agent

366 - # to validate the metadata received is the same as the one in prefill_connector.

367 - with patch.object(

368 - decode_connector.connector_worker, "add_remote_agent"

369 - ) as mock_add_remote_agent:

370 - mock_add_remote_agent.return_type = "remote_agent"

371 -

372 - decode_connector.connector_worker._nixl_handshake(

373 - kv_connector_metadata["remote_host"],

374 - kv_connector_metadata["remote_port"],

375 - kv_connector_metadata["tp_size"],

376 - kv_connector_metadata["remote_engine_id"],

316 + # Prefill connector will register KV cache to populate proper handshake

317 + # metadata.

318 + prefill_connector = NixlConnector(vllm_config, KVConnectorRole.WORKER)

319 + kv_cache_shape = FlashAttentionBackend.get_kv_cache_shape(

320 + num_blocks=2, block_size=16, num_kv_heads=4, head_size=64

321 )

322 + shared_tensor = torch.zeros(*kv_cache_shape, dtype=torch.float16)

323 + unique_tensor = torch.zeros(*kv_cache_shape, dtype=torch.float16)

324 + kv_caches = {

325 + "layer0": shared_tensor,

326 + "layer1": unique_tensor,

327 + "layer2": shared_tensor,

328 + }

329 + prefill_connector.register_kv_caches(kv_caches)

330 +

331 + # Simulate EngineCore initialization that would gather connector

332 + # metadata from all workers

333 + metadata = prefill_connector.get_handshake_metadata()

334 +

335 + # metadata is a NixlHandshakePayload, decode it to get NixlAgentMetadata

336 + decoder = msgspec.msgpack.Decoder(NixlAgentMetadata)

337 + expected_agent_metadata = decoder.decode(metadata.agent_metadata_bytes)

338 +

339 + # The scheduler connector expects metadata to be in

340 + # dict[int, KVConnectorHandshakeMetadata], where the first key is

341 + # the dp_rank, the second key is the tp_rank.

342 + scheduler_connector = scheduler.get_kv_connector()

343 + scheduler_connector.set_xfer_handshake_metadata({0: metadata})

344 +

345 + # Simulate a request that finishes prefill, which returns

346 + # corresponding NixlConnectorMetadata for decode instance.

347 + BLOCK_SIZE = vllm_config.cache_config.block_size

348 + NUM_EXTERNAL_FULL_BLOCKS = 2

349 + NUM_TOKENS = int(BLOCK_SIZE * (NUM_EXTERNAL_FULL_BLOCKS + 0.5))

350 +

351 + request = create_request(

352 + request_id=1,

353 + block_size=BLOCK_SIZE,

354 + num_tokens=NUM_TOKENS,

355 + do_remote_decode=True,

356 + )

357 + request.status = RequestStatus.FINISHED_LENGTH_CAPPED

358 + delay, kv_connector_metadata = scheduler.get_kv_connector().request_finished(

359 + request, [0, 1, 2]

360 + )

361 + assert delay

362 +

363 + # Decode connector will be able to create handshake with the prefill connector.

364 + decode_connector = NixlConnector(vllm_config, KVConnectorRole.WORKER)

365 +

366 + # Here we are testing the retrieval of NIXLAgentMetadata.

367 + # Knowing the implementation detail, we override the add_remote_agent

368 + # to validate the metadata received is the same as the one in prefill_connector.

369 + with patch.object(

370 + decode_connector.connector_worker, "add_remote_agent"

371 + ) as mock_add_remote_agent:

372 + mock_add_remote_agent.return_type = "remote_agent"

373 +

374 + decode_connector.connector_worker._nixl_handshake(

375 + kv_connector_metadata["remote_host"],

376 + kv_connector_metadata["remote_port"],

377 + kv_connector_metadata["tp_size"],

378 + kv_connector_metadata["remote_engine_id"],

379 + )

380

379 - received_metadata = mock_add_remote_agent.call_args.args

380 - assert received_metadata[0] == expected_agent_metadata

381 - assert received_metadata[1] == 0 # remote_tp_rank

382 - assert received_metadata[2] == 1 # remote_tp_size

381 + received_metadata = mock_add_remote_agent.call_args.args

382 + assert received_metadata[0] == expected_agent_metadata

383 + assert received_metadata[1] == 0 # remote_tp_rank

384 + assert received_metadata[2] == 1 # remote_tp_size

385

384 - # Need to shutdown the background thread to release NIXL side channel port

385 - scheduler_connector.shutdown()

386 + # Need to shutdown the background thread to release NIXL side channel port

387 + scheduler_connector.shutdown()

388

389

390 class FakeNixlConnectorWorker(NixlConnectorWorker):

@@ -458,6 +460,7 @@ class TestNixlHandshake:

460 )

461 def test_multi_xfer_one_engine(

462 self,

463 + default_vllm_config,

464 # dist_init is a fixture that initializes the distributed environment.

465 dist_init,

466 ):

@@ -547,6 +550,7 @@ class TestNixlHandshake:

550 )

551 def test_async_load_kv(

552 self,

553 + default_vllm_config,

554 # Fixture that initializes the distributed environment.

555 dist_init,

556 # Simulate consumer-producer TP sizes.

@@ -605,7 +609,7 @@ class TestNixlHandshake:

609 )

610 @pytest.mark.parametrize("local_tp_size", [1, 2])

611 def test_prefill_tp_size_greater_than_decode_tp_size(

608 - self, local_tp_size: int, dist_init

612 + self, local_tp_size: int, default_vllm_config, dist_init

613 ):

614 """

615 Verify remote TP > local TP handshake succeeds with different

@@ -670,7 +674,7 @@ class TestNixlHandshake:

674 )

675 @pytest.mark.parametrize("local_tp_size", [1, 2])

676 def test_prefill_tp_size_greater_than_decode_tp_size_mla(

673 - self, local_tp_size: int, dist_init

677 + self, local_tp_size: int, default_vllm_config, dist_init

678 ):

679 """

680 Verify remote TP > local TP handshake succeeds with different

@@ -770,6 +774,7 @@ class TestNixlHandshake:

774 )

775 def test_concurrent_load_kv(

776 self,

777 + default_vllm_config,

778 # dist_init is a fixture that initializes the distributed environment.

779 dist_init,

780 ):

@@ -830,7 +835,9 @@ class TestNixlHandshake:

835 "vllm.distributed.kv_transfer.kv_connector.v1.nixl_connector.NixlWrapper",

836 FakeNixlWrapper,

837 )

833 - def test_handshake_fails_on_kv_cache_layout_mismatch(self, dist_init):

838 + def test_handshake_fails_on_kv_cache_layout_mismatch(

839 + self, default_vllm_config, dist_init

840 + ):

841 """

842 Verify that adding a remote agent fails if kv_cache_layout differs.

843 This test is only relevant for heterogeneous TP.

@@ -879,7 +886,7 @@ class TestNixlHandshake:

886 FakeNixlWrapper,

887 )

888 def test_handshake_succeed_on_kv_cache_layout_mismatch_with_experimental(

882 - self, dist_init

889 + self, default_vllm_config, dist_init

890 ):

891 """

892 Verify that adding a remote agent fails if kv_cache_layout differs.

@@ -934,7 +941,7 @@ class TestNixlHandshake:

941 "vllm.distributed.kv_transfer.kv_connector.v1.nixl_connector.NixlWrapper",

942 FakeNixlWrapper,

943 )

937 -def test_kv_connector_stats(dist_init):

944 +def test_kv_connector_stats(default_vllm_config, dist_init):

945 """Test that KV transfer stats are properly recorded and retrieved."""

946 vllm_config = create_vllm_config()

947

@@ -1357,7 +1364,7 @@ def _run_abort_timeout_test(llm: LLM, timeout: int):

1364 "TRITON_ATTN",

1365 ],

1366 )

1360 -def test_register_kv_caches(dist_init, attn_backend):

1367 +def test_register_kv_caches(default_vllm_config, dist_init, attn_backend):

1368 """

1369 Test that register_kv_caches() properly calls nixl_wrapper methods with

1370 correct data.

@@ -1518,7 +1525,9 @@ class FakePlatform(Platform):

1525 ("oot", "VRAM"),

1526 ],

1527 )

1521 -def test_kv_buffer_to_nixl_memory_types(dist_init, kv_buffer_device, nixl_memory_type):

1528 +def test_kv_buffer_to_nixl_memory_types(

1529 + default_vllm_config, dist_init, kv_buffer_device, nixl_memory_type

1530 +):

1531 """

1532 Test that register_kv_caches() passes the correct memory types from the

1533 config to the nixl_wrapper.

@@ -1563,7 +1572,7 @@ def test_kv_buffer_to_nixl_memory_types(dist_init, kv_buffer_device, nixl_memory

1572 "vllm.distributed.kv_transfer.kv_connector.v1.nixl_connector.NixlWrapper",

1573 FakeNixlWrapper,

1574 )

1566 -def test_shutdown_cleans_up_resources(dist_init):

1575 +def test_shutdown_cleans_up_resources(default_vllm_config, dist_init):

1576 """Test that shutdown() properly cleans up all resources."""

1577 vllm_config = create_vllm_config()

1578

@@ -1622,7 +1631,7 @@ def test_shutdown_cleans_up_resources(dist_init):

1631 "vllm.distributed.kv_transfer.kv_connector.v1.nixl_connector.NixlWrapper",

1632 FakeNixlWrapper,

1633 )

1625 -def test_aborted_request_removed_from_worker_in_batch(dist_init):

1634 +def test_aborted_request_removed_from_worker_in_batch(default_vllm_config, dist_init):

1635 """

1636 Create and schedule a request so that P adds it to in-batch tracking via

1637 the real scheduler, then simulate an abort (request not in next scheduler

@@ -1731,7 +1740,7 @@ class FailingNixlWrapper(FakeNixlWrapper):

1740 "vllm.distributed.kv_transfer.kv_connector.v1.nixl_connector.NixlWrapper",

1741 FailingNixlWrapper,

1742 )

1734 -def test_handshake_failure_returns_finished(dist_init):

1743 +def test_handshake_failure_returns_finished(default_vllm_config, dist_init):

1744 """Test that handshake failures mark blocks invalid and return via get_finished."""

1745 vllm_config = create_vllm_config()

1746

@@ -1780,7 +1789,7 @@ def test_handshake_failure_returns_finished(dist_init):

1789 "vllm.distributed.kv_transfer.kv_connector.v1.nixl_connector.NixlWrapper",

1790 FailingNixlWrapper,

1791 )

1783 -def test_transfer_setup_failure_returns_finished(dist_init):

1792 +def test_transfer_setup_failure_returns_finished(default_vllm_config, dist_init):

1793 """Test that transfer setup failures mark blocks invalid

1794 and return via get_finished."""

1795 vllm_config = create_vllm_config()

@@ -1855,6 +1864,7 @@ def test_transfer_setup_failure_returns_finished(dist_init):

1864 FakeNixlWrapper,

1865 )

1866 def test_compatibility_hash_validation(

1867 + default_vllm_config,

1868 dist_init,

1869 mismatch_type,

1870 config_overrides,

@@ -1967,7 +1977,7 @@ def test_compatibility_hash_validation(

1977 "vllm.distributed.kv_transfer.kv_connector.v1.nixl_connector.NixlWrapper",

1978 FakeNixlWrapper,

1979 )

1970 -def test_handshake_decode_errors(dist_init, error_scenario):

1980 +def test_handshake_decode_errors(default_vllm_config, dist_init, error_scenario):

1981 """

1982 Test that msgspec decode errors are properly handled during handshake.

1983

tests/v1/kv_offload/test_cpu_gpu.py

@@ -50,6 +50,7 @@ NUM_MAPPINGS = [3]

50 @pytest.mark.parametrize("device", CUDA_DEVICES)

51 @torch.inference_mode()

52 def test_transfer(

53 + default_vllm_config,

54 gpu_to_cpu: bool,

55 num_mappings: int,

56 head_size: int,

tests/v1/worker/test_gpu_model_runner.py

+16 -15

@@ -112,15 +112,16 @@ def get_vllm_config():

112 @pytest.fixture

113 def model_runner():

114 vllm_config = get_vllm_config()

115 - model_config = vllm_config.model_config

116 - num_heads = model_config.get_num_kv_heads(vllm_config.parallel_config)

117 - head_size = model_config.get_head_size()

118 - vllm_config.compilation_config.static_forward_context["layer.0"] = Attention(

119 - num_heads, head_size, 0.1

120 - )

121 - runner = GPUModelRunner(vllm_config, DEVICE)

122 - initialize_kv_cache(runner)

123 - return runner

115 + with set_current_vllm_config(vllm_config):

116 + model_config = vllm_config.model_config

117 + num_heads = model_config.get_num_kv_heads(vllm_config.parallel_config)

118 + head_size = model_config.get_head_size()

119 + vllm_config.compilation_config.static_forward_context["layer.0"] = Attention(

120 + num_heads, head_size, 0.1

121 + )

122 + runner = GPUModelRunner(vllm_config, DEVICE)

123 + initialize_kv_cache(runner)

124 + yield runner

125

126

127 model_runner_2 = model_runner

@@ -546,7 +547,7 @@ def test_reload_weights_before_load_model(model_runner):

547 model_runner.reload_weights()

548

549

549 -def test_init_kv_cache_with_kv_sharing_invalid_target_layer_order():

550 +def test_init_kv_cache_with_kv_sharing_invalid_target_layer_order(default_vllm_config):

551 torch.set_default_dtype(torch.float16)

552 layer_0 = "model.layers.0.self_attn.attn"

553 layer_1 = "model.layers.1.self_attn.attn"

@@ -573,7 +574,7 @@ def test_init_kv_cache_with_kv_sharing_invalid_target_layer_order():

574 assert fwd_context is not None

575

576

576 -def test_init_kv_cache_with_kv_sharing_target_layer_not_exist():

577 +def test_init_kv_cache_with_kv_sharing_target_layer_not_exist(default_vllm_config):

578 torch.set_default_dtype(torch.float16)

579 layer_0 = "model.layers.0.self_attn.attn"

580 layer_1 = "model.layers.1.self_attn.attn"

@@ -600,7 +601,7 @@ def test_init_kv_cache_with_kv_sharing_target_layer_not_exist():

601 assert fwd_context is not None

602

603

603 -def test_init_kv_cache_with_kv_sharing_target_same_as_current():

604 +def test_init_kv_cache_with_kv_sharing_target_same_as_current(default_vllm_config):

605 torch.set_default_dtype(torch.float16)

606 layer_0 = "model.layers.0.self_attn.attn"

607 layer_1 = "model.layers.1.self_attn.attn"

@@ -627,7 +628,7 @@ def test_init_kv_cache_with_kv_sharing_target_same_as_current():

628 assert fwd_context is not None

629

630

630 -def test_init_kv_cache_without_kv_sharing():

631 +def test_init_kv_cache_without_kv_sharing(default_vllm_config):

632 torch.set_default_dtype(torch.float16)

633 layer_0 = "model.layers.0.self_attn.attn"

634 layer_1 = "model.layers.1.self_attn.attn"

@@ -694,7 +695,7 @@ def test_init_kv_cache_without_kv_sharing():

695 assert kv_cache_config.kv_cache_groups[0].layer_names[1] == layer_1

696

697

697 -def test_init_kv_cache_with_kv_sharing_valid():

698 +def test_init_kv_cache_with_kv_sharing_valid(default_vllm_config):

699 torch.set_default_dtype(torch.float16)

700 layer_0 = "model.layers.0.self_attn.attn"

701 layer_1 = "model.layers.1.self_attn.attn"

@@ -1047,7 +1048,7 @@ def test_input_batch_with_kernel_block_sizes():

1048 assert block_table.block_size == kernel_size

1049

1050

1050 -def test_hybrid_cache_integration(model_runner, dist_init):

1051 +def test_hybrid_cache_integration(default_vllm_config, dist_init):

1052 """Test hybrid cache architecture integration with GPUModelRunner."""

1053 # Create a new model runner with hybrid cache configuration

1054 vllm_config = get_vllm_config()

tests/v1/worker/test_utils.py

+8 -8

@@ -6,14 +6,14 @@ import torch

6 from vllm.v1.worker.utils import bind_kv_cache

9 -def test_bind_kv_cache():

9 +def test_bind_kv_cache(default_vllm_config):

10 from vllm.attention.layer import Attention

12 ctx = {

13 - "layers.0.self_attn": Attention(32, 128, 0.1),

14 - "layers.1.self_attn": Attention(32, 128, 0.1),

15 - "layers.2.self_attn": Attention(32, 128, 0.1),

16 - "layers.3.self_attn": Attention(32, 128, 0.1),

13 + "layers.0.self_attn": Attention(32, 128, 0.1, prefix="layers.0.self_attn"),

14 + "layers.1.self_attn": Attention(32, 128, 0.1, prefix="layers.1.self_attn"),

15 + "layers.2.self_attn": Attention(32, 128, 0.1, prefix="layers.2.self_attn"),

16 + "layers.3.self_attn": Attention(32, 128, 0.1, prefix="layers.3.self_attn"),

17 }

18 kv_cache = {

19 "layers.0.self_attn": torch.zeros((1,)),

@@ -34,13 +34,13 @@ def test_bind_kv_cache():

34 assert runner_kv_caches[3] is kv_cache["layers.3.self_attn"]

37 -def test_bind_kv_cache_non_attention():

37 +def test_bind_kv_cache_non_attention(default_vllm_config):

38 from vllm.attention.layer import Attention

40 # example from Jamba PP=2

41 ctx = {

42 - "model.layers.20.attn": Attention(32, 128, 0.1),

43 - "model.layers.28.attn": Attention(32, 128, 0.1),

42 + "model.layers.20.attn": Attention(32, 128, 0.1, prefix="model.layers.20.attn"),

43 + "model.layers.28.attn": Attention(32, 128, 0.1, prefix="model.layers.28.attn"),

44 }

45 kv_cache = {

46 "model.layers.20.attn": torch.zeros((1,)),

vllm/attention/utils/fa_utils.py

+6 -3

@@ -59,10 +59,13 @@ def get_flash_attn_version(requires_alibi: bool = False) -> int | None:

59 )

61 # 2. override if passed by environment or config

62 - from vllm.config import get_current_vllm_config

62 + from vllm.config import get_current_vllm_config_or_none

64 - vllm_config = get_current_vllm_config()

65 - if vllm_config.attention_config.flash_attn_version is not None:

64 + vllm_config = get_current_vllm_config_or_none()

65 + if (

66 + vllm_config is not None

67 + and vllm_config.attention_config.flash_attn_version is not None

68 + ):

69 fa_version = vllm_config.attention_config.flash_attn_version

71 # 3. fallback for unsupported combinations

vllm/config/__init__.py

@@ -42,6 +42,7 @@ from vllm.config.vllm import (

42 VllmConfig,

43 get_cached_compilation_config,

44 get_current_vllm_config,

45 + get_current_vllm_config_or_none,

46 get_layers_from_vllm_config,

47 set_current_vllm_config,

48 )

@@ -105,6 +106,7 @@ __all__ = [

106 "VllmConfig",

107 "get_cached_compilation_config",

108 "get_current_vllm_config",

109 + "get_current_vllm_config_or_none",

110 "set_current_vllm_config",

111 "get_layers_from_vllm_config",

112 ]

vllm/config/vllm.py

+12 -7

@@ -1441,13 +1441,18 @@ def get_cached_compilation_config():

1441

1442 def get_current_vllm_config() -> VllmConfig:

1443 if _current_vllm_config is None:

1444 - # in ci, usually when we test custom ops/modules directly,

1445 - # we don't set the vllm config. In that case, we set a default

1446 - # config.

1447 - # Use stack level 2 so the log contains the line of the caller,

1448 - # so it's easier to track down the source of the warning.

1449 - logger.warning("Current vLLM config is not set.", stacklevel=2)

1450 - return VllmConfig()

1444 + raise AssertionError(

1445 + "Current vLLM config is not set. This typically means "

1446 + "get_current_vllm_config() was called outside of a "

1447 + "set_current_vllm_config() context, or a CustomOp was instantiated "

1448 + "at module import time or model forward time when config is not set. "

1449 + "For tests that directly test custom ops/modules, use the "

1450 + "'default_vllm_config' pytest fixture from tests/conftest.py."

1451 + )

1452 + return _current_vllm_config

1453 +

1454 +

1455 +def get_current_vllm_config_or_none() -> VllmConfig | None:

1456 return _current_vllm_config

1457

1458

vllm/distributed/device_communicators/base_device_communicator.py

+2 -2

@@ -117,9 +117,9 @@ class DeviceCommunicatorBase:

117

118 use_ep = False

119 all2all_backend = None

120 - from vllm.config import get_current_vllm_config

120 + from vllm.config import get_current_vllm_config_or_none

121

122 - config = get_current_vllm_config()

122 + config = get_current_vllm_config_or_none()

123 if config is not None:

124 # as long as we use data parallel (coupled data parallel

125 # where all data parallel ranks execute forward together),

vllm/distributed/device_communicators/quick_all_reduce.py

+2 -2

@@ -9,7 +9,7 @@ from torch.distributed import ProcessGroup

10 import vllm.envs as envs

11 from vllm import _custom_ops as ops

12 -from vllm.config import get_current_vllm_config

12 +from vllm.config import get_current_vllm_config_or_none

13 from vllm.distributed.parallel_state import in_the_same_node_as

14 from vllm.logger import init_logger

15 from vllm.platforms import current_platform

@@ -184,7 +184,7 @@ class QuickAllReduce:

184 )

185 return

186 self.qr_quant_level = QuickReduceRegime[regime_str]

187 - vllm_config = get_current_vllm_config()

187 + vllm_config = get_current_vllm_config_or_none()

188 if (

189 vllm_config is not None

190 and hasattr(vllm_config, "model_config")

vllm/distributed/parallel_state.py

+6 -6

@@ -1177,9 +1177,9 @@ def init_distributed_environment(

1177 distributed_init_method,

1178 backend,

1179 )

1180 - from vllm.config import get_current_vllm_config

1180 + from vllm.config import get_current_vllm_config_or_none

1181

1182 - config = get_current_vllm_config()

1182 + config = get_current_vllm_config_or_none()

1183 if (

1184 config is not None

1185 and config.parallel_config.distributed_executor_backend != "external_launcher"

@@ -1251,7 +1251,7 @@ def init_distributed_environment(

1251 if _WORLD is None:

1252 ranks = list(range(torch.distributed.get_world_size()))

1253 _WORLD = init_world_group(ranks, local_rank, backend)

1254 - if config.parallel_config.nnodes > 1:

1254 + if config is not None and config.parallel_config.nnodes > 1:

1255 _NODE_COUNT = config.parallel_config.nnodes

1256 else:

1257 _NODE_COUNT = _node_count(_WORLD.cpu_group)

@@ -1260,7 +1260,7 @@ def init_distributed_environment(

1260 assert _WORLD.world_size == torch.distributed.get_world_size(), (

1261 "world group already initialized with a different world size"

1262 )

1263 - if config.parallel_config.nnodes_within_dp > 1:

1263 + if config is not None and config.parallel_config.nnodes_within_dp > 1:

1264 if parallel_config.data_parallel_size > 1:

1265 world_size_inner_dp = parallel_config.world_size

1266 group_ranks = [

@@ -1316,9 +1316,9 @@ def initialize_model_parallel(

1316 backend = backend or torch.distributed.get_backend(get_world_group().device_group)

1317

1318 data_parallel_size = 1

1319 - from vllm.config import get_current_vllm_config

1319 + from vllm.config import get_current_vllm_config_or_none

1320

1321 - config = get_current_vllm_config()

1321 + config = get_current_vllm_config_or_none()

1322 if config is not None:

1323 data_parallel_size = config.parallel_config.data_parallel_size

1324

vllm/model_executor/layers/fused_moe/cpu_fused_moe.py

+23 -5

@@ -13,10 +13,28 @@ from vllm.model_executor.layers.quantization.utils.layer_utils import replace_pa

13 from vllm.utils.torch_utils import direct_register_custom_op

15 _CPU_MOE_LAYER_CACHE = {}

16 -_CPU_MOE_ACT = {

17 - "silu": SiluAndMul(),

18 - "swigluoai": SwigluOAIAndMul(),

19 -}

16 +

17 +

18 +class _LazyActivationDict(dict):

19 + """Lazily instantiate activation functions on first access.

20 +

21 + Avoids triggering CustomOp.__init__() at module import time,

22 + which would call get_current_vllm_config() before config is set.

23 + """

24 +

25 + _factories: dict[str, type[SiluAndMul] | type[SwigluOAIAndMul]] = {

26 + "silu": SiluAndMul,

27 + "swigluoai": SwigluOAIAndMul,

28 + }

29 +

30 + def __missing__(self, key: str) -> SiluAndMul | SwigluOAIAndMul:

31 + if key not in self._factories:

32 + raise KeyError(f"{key} is not a supported activation")

33 + self[key] = self._factories[key]()

34 + return self[key]

35 +

36 +

37 +_CPU_MOE_ACT = _LazyActivationDict()

40 def grouped_topk(

@@ -212,7 +230,7 @@ class CPUFusedMOE:

230 apply_router_weight_on_input: bool = False,

231 activation: str = "silu",

232 ) -> torch.Tensor:

215 - assert activation in _CPU_MOE_ACT, f"{activation} is not supported."

233 + assert activation in _CPU_MOE_ACT._factories, f"{activation} is not supported."

234 assert not apply_router_weight_on_input

235

236 topk_weights, topk_ids = select_experts(

vllm/model_executor/layers/fused_moe/layer.py

+16 -13

@@ -540,6 +540,20 @@ class FusedMoE(CustomOp):

540 self.apply_router_weight_on_input = apply_router_weight_on_input

541 self.activation = activation

542

543 + self._grouped_topk_impl: GroupedTopk | None = None

544 + if self.use_grouped_topk:

545 + assert self.num_expert_group is not None

546 + assert self.topk_group is not None

547 + self._grouped_topk_impl = GroupedTopk(

548 + topk=self.top_k,

549 + renormalize=self.renormalize,

550 + num_expert_group=self.num_expert_group,

551 + topk_group=self.topk_group,

552 + scoring_func=self.scoring_func,

553 + routed_scaling_factor=self.routed_scaling_factor,

554 + num_fused_shared_experts=self.num_fused_shared_experts,

555 + )

556 +

557 if self.scoring_func != "softmax" and not self.use_grouped_topk:

558 raise ValueError(

559 "Only softmax scoring function is supported for non-grouped topk."

@@ -1588,19 +1602,8 @@ class FusedMoE(CustomOp):

1602

1603 # DeepSeekv2 uses grouped_top_k

1604 elif self.use_grouped_topk and valid_grouping():

1591 - assert self.topk_group is not None

1592 - assert self.num_expert_group is not None

1593 - grouped_topk_impl = GroupedTopk(

1594 - topk=self.top_k,

1595 - renormalize=self.renormalize,

1596 - num_expert_group=self.num_expert_group,

1597 - topk_group=self.topk_group,

1598 - scoring_func=self.scoring_func,

1599 - routed_scaling_factor=self.routed_scaling_factor,

1600 - num_fused_shared_experts=self.num_fused_shared_experts,

1601 - )

1602 -

1603 - topk_weights, topk_ids = grouped_topk_impl(

1605 + assert self._grouped_topk_impl is not None

1606 + topk_weights, topk_ids = self._grouped_topk_impl(

1607 hidden_states=hidden_states,

1608 gating_output=router_logits,

1609 e_score_correction_bias=self.e_score_correction_bias,

vllm/model_executor/models/keye.py

+7 -6

@@ -339,15 +339,11 @@ def apply_rotary_pos_emb_flashatt(

339 k: torch.Tensor,

340 cos: torch.Tensor,

341 sin: torch.Tensor,

342 + apply_rotary_emb: ApplyRotaryEmb,

343 ) -> tuple[torch.Tensor, torch.Tensor]:

344 cos = cos.chunk(2, dim=-1)[0].contiguous()

345 sin = sin.chunk(2, dim=-1)[0].contiguous()

346

346 - apply_rotary_emb = ApplyRotaryEmb(

347 - enforce_enable=True,

348 - enable_fp32_compute=True,

349 - )

350 -

347 q_embed = apply_rotary_emb(q, cos, sin)

348 k_embed = apply_rotary_emb(k, cos, sin)

349

@@ -410,6 +406,11 @@ class KeyeSiglipAttention(nn.Module):

406 multimodal_config=multimodal_config,

407 )

408

409 + self.apply_rotary_emb = ApplyRotaryEmb(

410 + enforce_enable=True,

411 + enable_fp32_compute=True,

412 + )

413 +

414 def forward(

415 self,

416 hidden_states: torch.Tensor,

@@ -448,7 +449,7 @@ class KeyeSiglipAttention(nn.Module):

449 self.num_kv_heads,

450 self.head_dim,

451 )

451 - q, k = apply_rotary_pos_emb_flashatt(q, k, cos, sin)

452 + q, k = apply_rotary_pos_emb_flashatt(q, k, cos, sin, self.apply_rotary_emb)

453 v = v.view(

454 *v.shape[:-1],

455 self.num_kv_heads,

vllm/model_executor/models/siglip2navit.py

+8 -6

@@ -152,16 +152,12 @@ def apply_rotary_pos_emb(

152 k: torch.Tensor,

153 cos: torch.Tensor,

154 sin: torch.Tensor,

155 - is_flash_attn_backend: bool = False,

155 + is_flash_attn_backend: bool,

156 + apply_rotary_emb: ApplyRotaryEmb,

157 ) -> tuple[torch.Tensor, torch.Tensor]:

158 cos = cos.chunk(2, dim=-1)[0].contiguous()

159 sin = sin.chunk(2, dim=-1)[0].contiguous()

160

160 - apply_rotary_emb = ApplyRotaryEmb(

161 - enforce_enable=True,

162 - enable_fp32_compute=True,

163 - )

164 -

161 if is_flash_attn_backend and current_platform.is_cuda():

162 apply_rotary_emb_func = apply_rotary_emb.forward_cuda

163 elif is_flash_attn_backend and current_platform.is_rocm():

@@ -235,6 +231,11 @@ class Siglip2Attention(nn.Module):

231 multimodal_config=multimodal_config,

232 )

233

234 + self.apply_rotary_emb = ApplyRotaryEmb(

235 + enforce_enable=True,

236 + enable_fp32_compute=True,

237 + )

238 +

239 def forward(

240 self,

241 hidden_states: torch.Tensor,

@@ -260,6 +261,7 @@ class Siglip2Attention(nn.Module):

261 cos,

262 sin,

263 self.attn.is_flash_attn_backend,

264 + self.apply_rotary_emb,

265 )

266 queries = queries.squeeze(0)

267 keys = keys.squeeze(0)

vllm/v1/worker/gpu_worker.py

+4 -2

@@ -14,7 +14,7 @@ import torch.distributed

14 import torch.nn as nn

16 import vllm.envs as envs

17 -from vllm.config import CUDAGraphMode, VllmConfig

17 +from vllm.config import CUDAGraphMode, VllmConfig, set_current_vllm_config

18 from vllm.config.compilation import CompilationMode

19 from vllm.distributed import (

20 ensure_model_parallel_initialized,

@@ -268,7 +268,9 @@ class Worker(WorkerBase):

268 # to hijack tensor allocation.

269 def load_model(self) -> None:

270 eep_scale_up = os.environ.get("VLLM_ELASTIC_EP_SCALE_UP_LAUNCH") == "1"

271 - with self._maybe_get_memory_pool_context(tag="weights"):

271 + with self._maybe_get_memory_pool_context(

272 + tag="weights"

273 + ) and set_current_vllm_config(self.vllm_config):

274 self.model_runner.load_model(eep_scale_up=eep_scale_up)

275

276 def update_config(self, overrides: dict[str, Any]) -> None: