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

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 13, 08:15 UTC

[ROCM] Fix ROCm warnings, environment flag access, and GEMM kernel naming for consistency in `_aiter_ops.py` (#28464)

Signed-off-by: vllmellm <vllm.ellm@embeddedllm.com>

vllmellm committed Nov 13, 2025 at 05:46 UTC d8140b98337a253f73dedd58c7d1e9b8832a9e3f

5 files changed +33 -29

vllm/_aiter_ops.py

+27 -22

@@ -32,13 +32,13 @@ def if_aiter_supported(func: Callable) -> Callable:

32 def wrapper(*args, **kwargs):

33 # checks the platform, device arch and aiter library existance.

35 - from vllm.platforms.rocm import on_gfx9

35 + if current_platform.is_rocm() and IS_AITER_FOUND:

36 + from vllm.platforms.rocm import on_gfx9

37 - if current_platform.is_rocm() and on_gfx9() and IS_AITER_FOUND:

38 - return func(*args, **kwargs)

39 - else:

40 - # Return None or do nothing if not supported

41 - return None

38 + if on_gfx9():

39 + return func(*args, **kwargs)

40 +

41 + return None

43 return wrapper

@@ -296,7 +296,7 @@ def _rocm_aiter_mla_decode_fwd_fake(

296 pass

297

298

299 -def _rocm_aiter_gemm_w8a8_impl(

299 +def _rocm_aiter_gemm_a8w8_impl(

300 A: torch.Tensor,

301 B: torch.Tensor,

302 As: torch.Tensor,

@@ -313,7 +313,7 @@ def _rocm_aiter_gemm_w8a8_impl(

313 return gemm_a8w8_CK(A, B, As, Bs, bias, output_dtype)

314

315

316 -def _rocm_aiter_gemm_w8a8_fake(

316 +def _rocm_aiter_gemm_a8w8_fake(

317 A: torch.Tensor,

318 B: torch.Tensor,

319 As: torch.Tensor,

@@ -327,7 +327,7 @@ def _rocm_aiter_gemm_w8a8_fake(

327 return Y

328

329

330 -def _rocm_aiter_gemm_w8a8_blockscale_impl(

330 +def _rocm_aiter_gemm_a8w8_blockscale_impl(

331 A: torch.Tensor,

332 B: torch.Tensor,

333 As: torch.Tensor,

@@ -339,7 +339,7 @@ def _rocm_aiter_gemm_w8a8_blockscale_impl(

339 return gemm_a8w8_blockscale(A, B, As, Bs, dtype=output_dtype)

340

341

342 -def _rocm_aiter_gemm_w8a8_blockscale_fake(

342 +def _rocm_aiter_gemm_a8w8_blockscale_fake(

343 A: torch.Tensor,

344 B: torch.Tensor,

345 As: torch.Tensor,

@@ -419,6 +419,7 @@ class rocm_aiter_ops:

419 _FP4_GEMM_DYNAMIC_QUANT_ASM = envs.VLLM_ROCM_USE_AITER_FP4_ASM_GEMM

420 _TRITON_ROTARY_EMBED = envs.VLLM_ROCM_USE_AITER_TRITON_ROPE

421 _MOE_SHARED_EXPERTS_ENABLED = envs.VLLM_ROCM_USE_AITER_FUSION_SHARED_EXPERTS

422 + _TRITON_UNQUANT_GEMM = envs.VLLM_ROCM_USE_AITER_TRITON_GEMM

423

424 @classmethod

425 @if_aiter_supported

@@ -494,6 +495,11 @@ class rocm_aiter_ops:

495 def is_triton_rotary_embed_enabled(cls) -> bool:

496 return cls._AITER_ENABLED and cls._TRITON_ROTARY_EMBED

497

498 + @classmethod

499 + @if_aiter_supported

500 + def is_triton_gemm_enabled(cls) -> bool:

501 + return cls._AITER_ENABLED and cls._TRITON_UNQUANT_GEMM

502 +

503 @staticmethod

504 @if_aiter_supported

505 def register_ops_once() -> None:

@@ -555,18 +561,18 @@ class rocm_aiter_ops:

561 )

562

563 direct_register_custom_op(

558 - op_name="rocm_aiter_gemm_w8a8",

559 - op_func=_rocm_aiter_gemm_w8a8_impl,

564 + op_name="rocm_aiter_gemm_a8w8",

565 + op_func=_rocm_aiter_gemm_a8w8_impl,

566 mutates_args=[],

561 - fake_impl=_rocm_aiter_gemm_w8a8_fake,

567 + fake_impl=_rocm_aiter_gemm_a8w8_fake,

568 dispatch_key=current_platform.dispatch_key,

569 )

570

571 direct_register_custom_op(

566 - op_name="rocm_aiter_gemm_w8a8_blockscale",

567 - op_func=_rocm_aiter_gemm_w8a8_blockscale_impl,

572 + op_name="rocm_aiter_gemm_a8w8_blockscale",

573 + op_func=_rocm_aiter_gemm_a8w8_blockscale_impl,

574 mutates_args=[],

569 - fake_impl=_rocm_aiter_gemm_w8a8_blockscale_fake,

575 + fake_impl=_rocm_aiter_gemm_a8w8_blockscale_fake,

576 dispatch_key=current_platform.dispatch_key,

577 )

578

@@ -606,7 +612,7 @@ class rocm_aiter_ops:

612 return torch.ops.vllm.rocm_aiter_rms_norm(x, weight, variance_epsilon)

613

614 @staticmethod

609 - def gemm_w8a8(

615 + def gemm_a8w8(

616 A: torch.Tensor,

617 B: torch.Tensor,

618 As: torch.Tensor,

@@ -614,10 +620,10 @@ class rocm_aiter_ops:

620 bias: torch.Tensor | None = None,

621 output_dtype: torch.dtype = torch.float16,

622 ) -> torch.Tensor:

617 - return torch.ops.vllm.rocm_aiter_gemm_w8a8(A, B, As, Bs, bias, output_dtype)

623 + return torch.ops.vllm.rocm_aiter_gemm_a8w8(A, B, As, Bs, bias, output_dtype)

624

625 @staticmethod

620 - def gemm_w8a8_blockscale(

626 + def gemm_a8w8_blockscale(

627 A: torch.Tensor,

628 B: torch.Tensor,

629 As: torch.Tensor,

@@ -625,7 +631,7 @@ class rocm_aiter_ops:

631 block_size: list[int],

632 output_dtype: torch.dtype = torch.float16,

633 ) -> torch.Tensor:

628 - return torch.ops.vllm.rocm_aiter_gemm_w8a8_blockscale(

634 + return torch.ops.vllm.rocm_aiter_gemm_a8w8_blockscale(

635 A, B, As, Bs, output_dtype

636 )

637

@@ -938,5 +944,4 @@ class rocm_aiter_ops:

944 return tuple(shuffle_weight(tensor, layout=layout) for tensor in tensors)

945

946

941 -if IS_AITER_FOUND:

942 - rocm_aiter_ops.register_ops_once()

947 +rocm_aiter_ops.register_ops_once()

vllm/model_executor/layers/quantization/kernels/scaled_mm/aiter.py

+1 -1

@@ -117,4 +117,4 @@ class AiterScaledMMLinearKernel(CutlassScaledMMLinearKernel):

117 # a to be [M, K]

118 # b to be [N, K]

119 # CutlassScaledMMLinearKernel prepare weight `w_q` in [K, N] format

120 - return rocm_aiter_ops.gemm_w8a8(x_q, w_q.t(), x_s, w_s, bias, out_dtype)

120 + return rocm_aiter_ops.gemm_a8w8(x_q, w_q.t(), x_s, w_s, bias, out_dtype)

vllm/model_executor/layers/quantization/utils/fp8_utils.py

+1 -1

@@ -328,7 +328,7 @@ class W8A8BlockFp8LinearOp:

328 if use_triton:

329 gemm_a8w8_blockscale_op = rocm_aiter_ops.triton_gemm_a8w8_blockscale

330 else:

331 - gemm_a8w8_blockscale_op = rocm_aiter_ops.gemm_w8a8_blockscale

331 + gemm_a8w8_blockscale_op = rocm_aiter_ops.gemm_a8w8_blockscale

332

333 if input_scale is not None:

334 q_input = input_2d

vllm/model_executor/layers/utils.py

+2 -2

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

9 from vllm import _custom_ops as ops

10 from vllm import envs

11 +from vllm._aiter_ops import rocm_aiter_ops

12 from vllm.logger import init_logger

13 from vllm.platforms import CpuArchEnum, current_platform

14 from vllm.utils.torch_utils import direct_register_custom_op

@@ -105,8 +106,7 @@ def default_unquantized_gemm(

106

107 def use_aiter_triton_gemm(n, m, k, dtype):

108 if (

108 - envs.VLLM_ROCM_USE_AITER == 0

109 - or envs.VLLM_ROCM_USE_AITER_TRITON_GEMM == 0

109 + not rocm_aiter_ops.is_triton_gemm_enabled()

110 # MI300's - fp8nuz=True

111 or current_platform.is_fp8_fnuz()

112 or dtype not in [torch.float16, torch.bfloat16]

vllm/platforms/rocm.py

+2 -3

@@ -325,6 +325,7 @@ class RocmPlatform(Platform):

325

326 @classmethod

327 def check_and_update_config(cls, vllm_config: "VllmConfig") -> None:

328 + from vllm._aiter_ops import rocm_aiter_ops

329 from vllm.config.compilation import CUDAGraphMode

330

331 cache_config = vllm_config.cache_config

@@ -332,9 +333,7 @@ class RocmPlatform(Platform):

333 parallel_config = vllm_config.parallel_config

334 is_eager_execution = compilation_config == CUDAGraphMode.NONE

335

335 - use_aiter_rms_norm = (

336 - envs.VLLM_ROCM_USE_AITER and envs.VLLM_ROCM_USE_AITER_RMSNORM

337 - )

336 + use_aiter_rms_norm = rocm_aiter_ops.is_rmsnorm_enabled()

337

338 if cache_config and cache_config.block_size is None:

339 cache_config.block_size = 16