ec68d53b2b · tduncan/Nvidia-Optimization-Stack-VLLM

benchmarks/kernels/benchmark_moe_align_block_size.py

+2 -1

@@ -9,6 +9,7 @@ from vllm.model_executor.layers.fused_moe.moe_align_block_size import (

9 moe_align_block_size,

10 )

11 from vllm.triton_utils import triton

12 +from vllm.utils.torch_utils import set_random_seed

13

14

15 def get_topk_ids(num_tokens: int, num_experts: int, topk: int) -> torch.Tensor:

@@ -44,7 +45,7 @@ configs = list(

45 def benchmark(num_tokens, num_experts, topk, ep_size, provider):

46 """Benchmark function for Triton."""

47 block_size = 256

47 - torch.cuda.manual_seed_all(0)

48 + set_random_seed(0)

49 topk_ids = get_topk_ids(num_tokens, num_experts, topk)

50

51 e_map = None

benchmarks/kernels/deepgemm/benchmark_fp8_block_dense_gemm.py

+2 -3

@@ -16,6 +16,7 @@ from vllm.utils.deep_gemm import (

16 fp8_gemm_nt,

17 per_block_cast_to_fp8,

18 )

19 +from vllm.utils.torch_utils import set_random_seed

20

21

22 def benchmark_shape(

@@ -235,9 +236,7 @@ def run_benchmarks(verbose: bool = False):

236 torch.backends.cudnn.allow_tf32 = True

237

238 # Set seeds for reproducibility

238 - torch.manual_seed(42)

239 - torch.cuda.manual_seed(42)

240 -

239 + set_random_seed(42)

240 # Define benchmark shapes (m, n, k)

241 shapes = [

242 (8, 4096, 7168),

tests/kernels/attention/test_lightning_attn.py

-6

@@ -122,8 +122,6 @@ def test_linear_decode_forward_triton(

122 dtype: torch.dtype,

123 ):

124 torch.set_default_device("cuda")

125 - torch.manual_seed(42)

126 - torch.cuda.manual_seed_all(42)

125 set_random_seed(42)

126 base = 0.01

127 q = base * torch.randn(batch_size, num_heads, 1, head_size, dtype=dtype)

@@ -165,8 +163,6 @@ def test_linear_decode_forward_triton_with_padding(

163 dtype: torch.dtype,

164 ):

165 torch.set_default_device("cuda")

168 - torch.manual_seed(42)

169 - torch.cuda.manual_seed_all(42)

166 set_random_seed(42)

167

168 batch_size = 4

@@ -229,8 +225,6 @@ def test_lightning_attention_reference(

225 dtype: torch.dtype,

226 ):

227 torch.set_default_device("cuda")

232 - torch.manual_seed(42)

233 - torch.cuda.manual_seed_all(42)

228 set_random_seed(42)

229

230 base = 0.01

tests/kernels/core/test_fused_quant_layernorm.py

+2 -3

@@ -17,6 +17,7 @@ from vllm.model_executor.layers.quantization.utils.int8_utils import (

17 per_token_group_quant_int8,

18 )

19 from vllm.platforms import current_platform

20 +from vllm.utils.torch_utils import set_random_seed

21

22 DTYPES = [torch.bfloat16, torch.float]

23 QUANT_DTYPES = [torch.int8, current_platform.fp8_dtype()]

@@ -180,9 +181,7 @@ def test_rms_norm(

181 device: str,

182 strided_input: bool,

183 ) -> None:

183 - torch.random.manual_seed(seed)

184 - if torch.cuda.is_available():

185 - torch.cuda.manual_seed(seed)

184 + set_random_seed(seed)

185 torch.set_default_device(device)

186 torch.accelerator.set_device_index(device)

187

tests/kernels/moe/test_marlin_vs_trtllm_mxint4.py

+3 -2

@@ -16,6 +16,7 @@ from vllm.model_executor.layers.quantization.utils.flashinfer_mxint4_moe import

16 )

17 from vllm.platforms import current_platform

18 from vllm.scalar_type import scalar_types

19 +from vllm.utils.torch_utils import set_random_seed

20

21

22 def mxint4_quantize(

@@ -134,7 +135,7 @@ def test_marlin_vs_trtllm_mxint4_moe_kimik2(monkeypatch, m, n, k, e, topk, group

135 pytest.importorskip("flashinfer")

136 monkeypatch.setenv("VLLM_USE_FLASHINFER_MOE_INT4", "1")

137

137 - torch.cuda.manual_seed(0)

138 + set_random_seed(0)

139

140 dtype = torch.bfloat16

141

@@ -289,7 +290,7 @@ def test_flashinfer_trtllm_mxint4_moe_wrapper(m, n, k, e, topk):

290 flashinfer_trtllm_mxint4_moe,

291 )

292

292 - torch.cuda.manual_seed(0)

293 + set_random_seed(0)

294 dtype = torch.bfloat16

295

296 a = torch.randn((m, k), device="cuda", dtype=dtype) * 0.5

tests/kernels/moe/test_moe.py

+5 -5

@@ -1031,7 +1031,7 @@ def test_fused_marlin_moe(

1031 act_order: bool,

1032 is_k_full: bool,

1033 ):

1034 - torch.cuda.manual_seed(1)

1034 + set_random_seed(1)

1035 group_size = group_blocks if group_blocks <= 0 else group_blocks * 16

1036

1037 if c_type == scalar_types.float16:

@@ -1131,7 +1131,7 @@ def test_fused_marlin_moe(

1131 @pytest.mark.skipif(current_platform.is_rocm(), reason="Skip for rocm")

1132 @pytest.mark.parametrize("m", [1, 256])

1133 def test_fused_marlin_moe_with_bias(m):

1134 - torch.cuda.manual_seed(0)

1134 + set_random_seed(0)

1135

1136 e, topk = 32, 4

1137 n, k = 2048, 2048

@@ -1213,7 +1213,7 @@ def test_fused_marlin_moe_non_gated(

1213 Non-gated activations like relu2 don't have the gate-up projection pattern,

1214 so w1 has shape (e, n, k) instead of (e, 2*n, k).

1215 """

1216 - torch.cuda.manual_seed(42)

1216 + set_random_seed(42)

1217

1218 group_size = 16 # NVFP4 group size

1219 is_k_full = True

@@ -1397,7 +1397,7 @@ def test_cpu_fused_moe_basic(

1397 from vllm.model_executor.layers.fused_moe.cpu_fused_moe import CPUFusedMOE

1398

1399 device = "cpu"

1400 - torch.manual_seed(7)

1400 + set_random_seed(7)

1401

1402 a = torch.randn((m, k), device=device, dtype=dtype) / 10

1403 w13 = torch.randn((e, 2 * n, k), device=device, dtype=dtype) / 10

@@ -1469,7 +1469,7 @@ def test_batched_fused_marlin_moe(

1469 f"topk={topk}, "

1470 f"max_tokens_per_batch={max_tokens_per_batch}"

1471 )

1472 - torch.cuda.manual_seed(0)

1472 + set_random_seed(0)

1473

1474 dtype = torch.bfloat16

1475 quant_dtype = scalar_types.float4_e2m1f

tests/kernels/moe/test_shared_fused_moe_routed_transform.py

+2 -3

@@ -15,7 +15,7 @@ from vllm.config import VllmConfig, set_current_vllm_config

15 from vllm.forward_context import set_forward_context

16 from vllm.model_executor.layers.fused_moe.shared_fused_moe import SharedFusedMoE

17 from vllm.platforms import current_platform

18 -from vllm.utils.torch_utils import is_torch_equal_or_newer

18 +from vllm.utils.torch_utils import is_torch_equal_or_newer, set_random_seed

19

20

21 class SimpleLinear(nn.Module):

@@ -144,8 +144,7 @@ def test_routed_input_transform_inside_vs_outside(

144

145 rocm_aiter_ops.refresh_env_variables()

146

147 - torch.manual_seed(42)

148 - torch.cuda.manual_seed(42)

147 + set_random_seed(42)

148

149 num_experts = 8

150 top_k = 2

tests/kernels/test_fused_quant_activation.py

+2 -3

@@ -7,6 +7,7 @@ import vllm._custom_ops as ops

7 from tests.kernels.utils import opcheck

8 from vllm.model_executor.layers.activation import SiluAndMul

9 from vllm.platforms import current_platform

10 +from vllm.utils.torch_utils import set_random_seed

11

12 DTYPES = [torch.bfloat16, torch.float16]

13 QUANT_DTYPES = [current_platform.fp8_dtype()]

@@ -49,9 +50,7 @@ def test_silu_and_mul(

50 seed: int,

51 device: str,

52 ) -> None:

52 - torch.random.manual_seed(seed)

53 - if torch.cuda.is_available():

54 - torch.cuda.manual_seed(seed)

53 + set_random_seed(seed)

54 torch.set_default_device(device)

55

56 layer = SiluAndMul()

tests/models/language/pooling/test_token_classification.py

+2 -7

@@ -1,25 +1,20 @@

1 # SPDX-License-Identifier: Apache-2.0

2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project

3 -import random

3

5 -import numpy as np

4 import pytest

5 import torch

6 from transformers import AutoModelForTokenClassification

7

8 from tests.models.utils import softmax

9 from vllm.platforms import current_platform

10 +from vllm.utils.torch_utils import set_random_seed

11

12

13 @pytest.fixture(autouse=True)

14 def seed_everything():

15 """Seed all random number generators for reproducibility."""

16 seed = 0

18 - random.seed(seed)

19 - np.random.seed(seed)

20 - torch.manual_seed(seed)

21 - if torch.cuda.is_available():

22 - torch.cuda.manual_seed_all(seed)

17 + set_random_seed(seed)

18 torch.backends.cudnn.deterministic = True

19 torch.backends.cudnn.benchmark = False

20 yield

tests/v1/e2e/spec_decode/test_lora_with_spec_decode.py

+2 -7

@@ -5,9 +5,6 @@ This script contains:

5 1. test lora with speculative decoding for batch inference

6 """

7

8 -import random

9 -

10 -import numpy as np

8 import pytest

9 import torch

10

@@ -15,6 +12,7 @@ from vllm import LLM, SamplingParams

12 from vllm.distributed import cleanup_dist_env_and_memory

13 from vllm.lora.request import LoRARequest

14 from vllm.platforms import current_platform

15 +from vllm.utils.torch_utils import set_random_seed

16

17 LORA_TEST_PROMPT_MAP: dict[str, str] = {}

18

@@ -63,10 +61,7 @@ def test_batch_inference_correctness(

61 with monkeypatch.context() as m:

62 # Disable randomness

63 m.setenv("CUBLAS_WORKSPACE_CONFIG", ":4096:8")

66 - torch.manual_seed(SEED)

67 - np.random.seed(SEED)

68 - random.seed(SEED)

69 - torch.cuda.manual_seed_all(SEED)

64 + set_random_seed(SEED)

65 torch.backends.cudnn.benchmark = False

66 torch.backends.cudnn.deterministic = True

67

tests/v1/spec_decode/test_tree_attention.py

+2 -2

@@ -14,6 +14,7 @@ from tests.v1.attention.utils import (

14 )

15 from vllm.config import ParallelConfig, SpeculativeConfig

16 from vllm.platforms import current_platform

17 +from vllm.utils.torch_utils import set_random_seed

18 from vllm.v1.attention.backend import CommonAttentionMetadata

19 from vllm.v1.attention.backends.fa_utils import is_flash_attn_varlen_func_available

20 from vllm.v1.attention.backends.registry import AttentionBackendEnum

@@ -323,8 +324,7 @@ def forward_attention(

324 def test_tree_attn_correctness(

325 reference_backend: AttentionBackendEnum,

326 ) -> None:

326 - torch.manual_seed(42)

327 - torch.cuda.manual_seed_all(42)

327 + set_random_seed(42)

328

329 device = "cuda"

330 tree_attn_masks = {

tools/pre_commit/check_torch_cuda.py

+9

@@ -9,6 +9,7 @@ import regex as re

9 # --------------------------------------------------------------------------- #

10 _TORCH_CUDA_PATTERNS = [

12 + r"\btorch\.cuda\.(manual_seed|manual_seed_all)\b",

13 r"\bwith\storch\.cuda\.device\b",

14 # Calls torch.cuda.{_is_compiled/_device_count_amdsmi/_device_count_nvml} internally

15 r"\bcuda_device_count_stateless\(\)\b",

@@ -24,6 +25,14 @@ def scan_file(path: str) -> int:

25 for match in re.finditer(pattern, content, re.MULTILINE):

26 # Calculate line number from match position

27 line_num = content[: match.start() + 1].count("\n") + 1

28 + matched_text = match.group(0)

29 + if "manual_seed" in matched_text:

30 + print(

31 + f"{path}:{line_num}: "

32 + "\033[91merror:\033[0m "

33 + f"Found {matched_text} API call. Use set_random_seed instead."

34 + )

35 + return 1

36 print(

37 f"{path}:{line_num}: "

38 "\033[91merror:\033[0m " # red color

vllm/platforms/cpu.py

+4

@@ -154,6 +154,10 @@ class CpuPlatform(Platform):

154 """

155 torch.cpu.set_device(device)

156

157 + @classmethod

158 + def manual_seed_all(cls, seed: int) -> None:

159 + pass

160 +

161 @classmethod

162 def inference_mode(cls):

163 return torch.no_grad()

vllm/platforms/cuda.py

+4

@@ -188,6 +188,10 @@ class CudaPlatformBase(Platform):

188 # for why and when it is needed

189 _ = torch.zeros(1, device=device)

190

191 + @classmethod

192 + def manual_seed_all(cls, seed: int) -> None:

193 + torch.cuda.manual_seed_all(seed)

194 +

195 @classmethod

196 def get_device_capability(cls, device_id: int = 0) -> DeviceCapability | None:

197 raise NotImplementedError

vllm/platforms/interface.py

+5

@@ -391,6 +391,11 @@ class Platform:

391 """

392 raise NotImplementedError

393

394 + @classmethod

395 + def manual_seed_all(cls, seed: int) -> None:

396 + """Set RNG seed across all devices for the current platform."""

397 + raise NotImplementedError

398 +

399 @classmethod

400 def pre_register_and_update(

401 cls, parser: FlexibleArgumentParser | None = None

vllm/platforms/rocm.py

+4

@@ -605,6 +605,10 @@ class RocmPlatform(Platform):

605 """

606 torch.cuda.set_device(device)

607

608 + @classmethod

609 + def manual_seed_all(cls, seed: int) -> None:

610 + torch.cuda.manual_seed_all(seed)

611 +

612 @classmethod

613 @lru_cache(maxsize=8)

614 def get_device_capability(cls, device_id: int = 0) -> DeviceCapability | None:

vllm/platforms/xpu.py

+4

@@ -125,6 +125,10 @@ class XPUPlatform(Platform):

125 """

126 torch.xpu.set_device(device)

127

128 + @classmethod

129 + def manual_seed_all(cls, seed: int) -> None:

130 + torch.xpu.manual_seed_all(seed)

131 +

132 @classmethod

133 def get_device_capability(

134 cls,

vllm/utils/torch_utils.py

+3 -2

@@ -365,8 +365,9 @@ def set_random_seed(seed: int | None) -> None:

365 random.seed(seed)

366 np.random.seed(seed)

367 torch.manual_seed(seed)

368 - if torch.cuda.is_available():

369 - torch.cuda.manual_seed_all(seed)

368 + from vllm.platforms import current_platform

369 +

370 + current_platform.manual_seed_all(seed)

371

372

373 def create_kv_caches_with_random_flash(

Add platform manual_seed_all API (#38468)