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

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 2, 02:15 UTC

Enable perf_token_group_quant/_C_stable_libtorch for ROCm (#42758)

Signed-off-by: charlifu <charlifu@amd.com>

Charlie Fu committed Jun 3, 2026 at 01:23 UTC 71df063c494c111ab60f6a33c54aafe7b9ae1d02

12 files changed +146 -110

CMakeLists.txt

+2 -2

@@ -633,6 +633,8 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP")

633 "csrc/libtorch_stable/activation_kernels.cu"

634 "csrc/libtorch_stable/quantization/w8a8/int8/scaled_quant.cu"

635 "csrc/libtorch_stable/quantization/w8a8/fp8/common.cu"

636 + "csrc/libtorch_stable/quantization/w8a8/fp8/per_token_group_quant.cu"

637 + "csrc/libtorch_stable/quantization/w8a8/int8/per_token_group_quant.cu"

638 "csrc/libtorch_stable/quantization/gptq/q_gemm.cu"

639 "csrc/libtorch_stable/quantization/gguf/gguf_kernel.cu"

640 "csrc/libtorch_stable/pos_encoding_kernels.cu"

@@ -657,8 +659,6 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP")

659 "csrc/libtorch_stable/quantization/fp4/nvfp4_quant_entry.cu"

660 "csrc/libtorch_stable/quantization/fp4/nvfp4_scaled_mm_entry.cu"

661 "csrc/libtorch_stable/permute_cols.cu"

660 - "csrc/libtorch_stable/quantization/w8a8/fp8/per_token_group_quant.cu"

661 - "csrc/libtorch_stable/quantization/w8a8/int8/per_token_group_quant.cu"

662 "csrc/libtorch_stable/quantization/awq/gemm_kernels.cu")

663

664 set_gencode_flags_for_srcs(

cmake/hipify.py

+31 -10

@@ -14,7 +14,18 @@ import argparse

14 import os

15 import shutil

17 -from torch.utils.hipify.hipify_python import hipify

17 +from torch.utils.hipify.hipify_python import get_hip_file_path, hipify

18 +

19 +

20 +def _expected_hip_build_path(source_abs: str, output_directory: str) -> str:

21 + """Match torch.utils.hipify.hipify_python.preprocessor fout_path naming."""

22 + rel = os.path.relpath(source_abs, output_directory)

23 + return os.path.abspath(

24 + os.path.join(

25 + output_directory, get_hip_file_path(rel, is_pytorch_extension=True)

26 + )

27 + )

28 +

30 if __name__ == "__main__":

31 parser = argparse.ArgumentParser()

@@ -53,7 +64,11 @@ if __name__ == "__main__":

64 hipify_result = hipify(

65 project_directory=args.project_dir,

66 output_directory=args.output_dir,

56 - header_include_dirs=[],

67 + # Hipify resolves quoted includes next to the including file first; vLLM

68 + # uses paths relative to csrc/ (e.g. "libtorch_stable/torch_utils.h"

69 + # from quantization/w8a8/fp8/*.cu). Without an include root here, those

70 + # headers are never found and are not hipified or rewritten in dependents.

71 + header_include_dirs=["."],

72 includes=includes,

73 extra_files=extra_files,

74 show_detailed=True,

@@ -64,14 +79,20 @@ if __name__ == "__main__":

79 hipified_sources = []

80 for source in args.sources:

81 s_abs = os.path.abspath(source)

67 - hipified_s_abs = (

68 - hipify_result[s_abs].hipified_path

69 - if (

70 - s_abs in hipify_result

71 - and hipify_result[s_abs].hipified_path is not None

72 - )

73 - else s_abs

74 - )

82 + if s_abs in hipify_result and hipify_result[s_abs].hipified_path is not None:

83 + path = hipify_result[s_abs].hipified_path

84 + # PyTorch skips writing when is_pytorch_extension and text unchanged;

85 + # hipified_path then stays *.cu. CMake expects *.hip under output_dir.

86 + if s_abs.endswith(".cu") and path.endswith(".cu"):

87 + dest = _expected_hip_build_path(s_abs, args.output_dir)

88 + if os.path.normpath(path) != os.path.normpath(dest):

89 + os.makedirs(os.path.dirname(dest), exist_ok=True)

90 + shutil.copy2(path, dest)

91 + hipified_s_abs = dest

92 + else:

93 + hipified_s_abs = path

94 + else:

95 + hipified_s_abs = s_abs

96 hipified_sources.append(hipified_s_abs)

98 assert len(hipified_sources) == len(args.sources)

csrc/libtorch_stable/ops.h

+4 -4

@@ -3,10 +3,6 @@

3 #include <torch/csrc/stable/library.h>

4 #include <torch/csrc/stable/tensor.h>

6 -#ifndef USE_ROCM

7 -torch::stable::Tensor permute_cols(torch::stable::Tensor const& A,

8 - torch::stable::Tensor const& perm);

9 -

6 void per_token_group_quant_fp8(const torch::stable::Tensor& input,

7 torch::stable::Tensor& output_q,

8 torch::stable::Tensor& output_s,

@@ -28,6 +24,10 @@ void per_token_group_quant_int8(const torch::stable::Tensor& input,

24 int64_t group_size, double eps, double int8_min,

25 double int8_max);

27 +#ifndef USE_ROCM

28 +torch::stable::Tensor permute_cols(torch::stable::Tensor const& A,

29 + torch::stable::Tensor const& perm);

30 +

31 bool cutlass_scaled_mm_supports_fp8(int64_t cuda_device_capability);

32 bool cutlass_scaled_mm_supports_block_fp8(int64_t cuda_device_capability);

33 bool cutlass_group_gemm_supported(int64_t cuda_device_capability);

csrc/libtorch_stable/quantization/w8a8/fp8/per_token_group_quant.cu

+33 -11

@@ -7,7 +7,11 @@

8 #include <cmath>

10 -#include <cuda_fp8.h>

10 +#ifdef USE_ROCM

11 + #include <hip/hip_fp8.h>

12 +#else

13 + #include <cuda_fp8.h>

14 +#endif

16 #include "libtorch_stable/quantization/vectorization.cuh"

17 #include "libtorch_stable/quantization/vectorization_utils.cuh"

@@ -15,12 +19,23 @@

19 #include "libtorch_stable/torch_utils.h"

21 __device__ __forceinline__ float GroupReduceMax(float val) {

22 +#ifdef USE_ROCM

23 + // 16-thread logical groups may pack up to four per 64-lane wavefront; use a

24 + // 64-bit mask and explicit width so shuffles stay within each group.

25 + const int lane_in_wave = threadIdx.x % warpSize;

26 + const unsigned long long mask = 0xFFFFull << ((lane_in_wave / 16) * 16);

27 + val = fmaxf(val, __shfl_xor_sync(mask, val, 8, 16));

28 + val = fmaxf(val, __shfl_xor_sync(mask, val, 4, 16));

29 + val = fmaxf(val, __shfl_xor_sync(mask, val, 2, 16));

30 + val = fmaxf(val, __shfl_xor_sync(mask, val, 1, 16));

31 +#else

32 unsigned mask = threadIdx.x % 32 >= 16 ? 0xffff0000 : 0x0000ffff;

34 val = fmaxf(val, __shfl_xor_sync(mask, val, 8));

35 val = fmaxf(val, __shfl_xor_sync(mask, val, 4));

36 val = fmaxf(val, __shfl_xor_sync(mask, val, 2));

37 val = fmaxf(val, __shfl_xor_sync(mask, val, 1));

38 +#endif

39 return val;

40 }

@@ -237,10 +252,12 @@ void per_token_group_quant_8bit(const torch::stable::Tensor& input,

252

253 VLLM_STABLE_DISPATCH_FLOATING_TYPES(

254 input.scalar_type(), "per_token_group_quant_8bit", ([&] {

240 - if (dst_type == torch::headeronly::ScalarType::Float8_e4m3fn) {

241 - LAUNCH_KERNEL(scalar_t, __nv_fp8_e4m3);

242 - } else if (dst_type == torch::headeronly::ScalarType::Char) {

255 + if (dst_type == torch::headeronly::ScalarType::Char) {

256 LAUNCH_KERNEL(scalar_t, int8_t);

257 + } else {

258 + VLLM_STABLE_DISPATCH_FP8_TYPES(

259 + dst_type, "per_token_group_quant_8bit_fp8",

260 + ([&] { LAUNCH_KERNEL(scalar_t, fp8_t); }));

261 }

262 }));

263

@@ -317,10 +334,18 @@ __global__ void per_token_group_quant_8bit_packed_register_kernel(

334

335 // 8-lane subgroup shuffle reduce (octet of the warp). The mask selects the

336 // 8 lanes within the warp that share a group.

337 +#ifdef USE_ROCM

338 + const int lane_in_wave = threadIdx.x % warpSize;

339 + const unsigned long long mask = 0xFFull << (lane_in_wave & ~7);

340 + local_absmax = fmaxf(local_absmax, __shfl_xor_sync(mask, local_absmax, 4, 8));

341 + local_absmax = fmaxf(local_absmax, __shfl_xor_sync(mask, local_absmax, 2, 8));

342 + local_absmax = fmaxf(local_absmax, __shfl_xor_sync(mask, local_absmax, 1, 8));

343 +#else

344 unsigned mask = 0xffu << (threadIdx.x & 24u);

345 local_absmax = fmaxf(local_absmax, __shfl_xor_sync(mask, local_absmax, 4));

346 local_absmax = fmaxf(local_absmax, __shfl_xor_sync(mask, local_absmax, 2));

347 local_absmax = fmaxf(local_absmax, __shfl_xor_sync(mask, local_absmax, 1));

348 +#endif

349

350 float y_s = local_absmax / max_8bit;

351 y_s = fmaxf(y_s, 1e-10f);

@@ -503,15 +528,12 @@ void per_token_group_quant_8bit_packed(const torch::stable::Tensor& input,

528

529 VLLM_STABLE_DISPATCH_HALF_TYPES(

530 input.scalar_type(), "per_token_group_quant_8bit_packed_register", ([&] {

506 - if (dst_type == torch::headeronly::ScalarType::Float8_e4m3fn) {

507 - LAUNCH_REG_KERNEL(scalar_t, __nv_fp8_e4m3);

508 - } else if (dst_type == torch::headeronly::ScalarType::Char) {

531 + if (dst_type == torch::headeronly::ScalarType::Char) {

532 LAUNCH_REG_KERNEL(scalar_t, int8_t);

533 } else {

511 - STD_TORCH_CHECK(

512 - false,

513 - "per_token_group_quant_8bit_packed only supports FP8/INT8 "

514 - "outputs.");

534 + VLLM_STABLE_DISPATCH_FP8_TYPES(

535 + dst_type, "per_token_group_quant_8bit_packed_fp8",

536 + ([&] { LAUNCH_REG_KERNEL(scalar_t, fp8_t); }));

537 }

538 }));

539

csrc/libtorch_stable/torch_bindings.cpp

+10 -10

@@ -7,11 +7,6 @@

7 // Note: We register under namespace "_C" so ops are accessible as

8 // torch.ops._C.<op_name> for compatibility with existing code.

9 STABLE_TORCH_LIBRARY_FRAGMENT(_C, ops) {

10 -#ifndef USE_ROCM

11 - ops.def("permute_cols(Tensor A, Tensor perm) -> Tensor");

12 -#endif

13 -

14 -#ifndef USE_ROCM

10 // Compute per-token-group FP8 quantized tensor and scaling factor.

11 // The dummy arguments are here so we can correctly fuse with RMSNorm.

12 ops.def(

@@ -32,6 +27,11 @@ STABLE_TORCH_LIBRARY_FRAGMENT(_C, ops) {

27 "output_s, int group_size, float eps, float int8_min, float int8_max) -> "

28 "()");

30 +#ifndef USE_ROCM

31 + ops.def("permute_cols(Tensor A, Tensor perm) -> Tensor");

32 +#endif

33 +

34 +#ifndef USE_ROCM

35 // CUTLASS w8a8 GEMM, supporting symmetric per-tensor or per-row/column

36 // quantization, as well as bias

37 ops.def(

@@ -508,11 +508,6 @@ STABLE_TORCH_LIBRARY_FRAGMENT(_C, ops) {

508 }

509

510 STABLE_TORCH_LIBRARY_IMPL(_C, CUDA, ops) {

511 -#ifndef USE_ROCM

512 - ops.impl("permute_cols", TORCH_BOX(&permute_cols));

513 -#endif

514 -

515 -#ifndef USE_ROCM

511 // Per-token group quantization

512 ops.impl("per_token_group_fp8_quant", TORCH_BOX(&per_token_group_quant_fp8));

513 ops.impl("per_token_group_fp8_quant_packed",

@@ -520,6 +515,11 @@ STABLE_TORCH_LIBRARY_IMPL(_C, CUDA, ops) {

515 ops.impl("per_token_group_quant_int8",

516 TORCH_BOX(&per_token_group_quant_int8));

517

518 +#ifndef USE_ROCM

519 + ops.impl("permute_cols", TORCH_BOX(&permute_cols));

520 +#endif

521 +

522 +#ifndef USE_ROCM

523 // CUTLASS scaled_mm ops

524 ops.impl("cutlass_scaled_mm", TORCH_BOX(&cutlass_scaled_mm));

525 ops.impl("cutlass_scaled_mm_azp", TORCH_BOX(&cutlass_scaled_mm_azp));

csrc/libtorch_stable/torch_utils.h

+1 -5

@@ -6,11 +6,7 @@

6 #include <torch/csrc/stable/tensor.h>

7 #include <torch/headeronly/util/shim_utils.h>

9 -#ifndef USE_ROCM

10 - #include <cuda_runtime.h>

11 -#else

12 - #include <hip/hip_runtime.h>

13 -#endif

9 +#include <cuda_runtime.h>

10 #include <cublas_v2.h>

12 #include <deque>

tests/kernels/quantization/test_per_token_group_quant.py

+24 -17

@@ -6,6 +6,7 @@ import pytest

6 import torch

8 from vllm.model_executor.layers.quantization.utils import fp8_utils, int8_utils

9 +from vllm.model_executor.layers.quantization.utils.quant_utils import get_fp8_min_max

10 from vllm.platforms import current_platform

@@ -16,7 +17,9 @@ from vllm.platforms import current_platform

17 @pytest.mark.parametrize("tma_aligned", [False, True])

18 @pytest.mark.parametrize("scale_ue8m0", [False, True])

19 @pytest.mark.parametrize("group_size", [64, 128])

19 -@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available")

20 +@pytest.mark.skipif(

21 + not current_platform.is_cuda_alike(), reason="Only test on CUDA/ROCm."

22 +)

23 def test_per_token_group_quant_fp8(

24 shape, column_major: bool, tma_aligned: bool, scale_ue8m0: bool, group_size: int

25 ):

@@ -37,7 +40,7 @@ def test_per_token_group_quant_fp8(

40 )

42 # triton ref

40 - with patch("vllm.platforms.current_platform.is_cuda", return_value=False):

43 + with patch("vllm.platforms.current_platform.is_cuda_alike", return_value=False):

44 ref_q, ref_s = fp8_utils.per_token_group_quant_fp8(

45 x,

46 group_size,

@@ -77,7 +80,8 @@ def test_per_token_group_quant_fp8(

80 )

81 @pytest.mark.parametrize("poisoned_scales", [False, True])

82 @pytest.mark.skipif(

80 - not current_platform.is_cuda(), reason="DeepGEMM not available on this platform"

83 + not current_platform.is_cuda_alike(),

84 + reason="DeepGEMM not available on this platform",

85 )

86 def test_per_token_group_quant_fp8_packed(

87 num_tokens, hidden_dim, group_size, poisoned_scales

@@ -99,8 +103,8 @@ def test_per_token_group_quant_fp8_packed(

103 if poisoned_scales:

104 # Call the kernel with poisoned scale buffer to

105 # ensure padded indices are correctly zeroed.

102 - fp8_dtype = torch.float8_e4m3fn

103 - finfo = torch.finfo(fp8_dtype)

106 + fp8_dtype = current_platform.fp8_dtype()

107 + fp8_min, fp8_max = get_fp8_min_max()

108 out_q = torch.empty_like(x, dtype=fp8_dtype)

109 out_s_packed = torch.empty_strided(

110 (mn, k_num_packed),

@@ -115,8 +119,8 @@ def test_per_token_group_quant_fp8_packed(

119 out_s_packed,

120 group_size,

121 1e-10,

118 - finfo.min,

119 - finfo.max,

122 + fp8_min,

123 + fp8_max,

124 )

125 else:

126 out_q, out_s_packed = fp8_utils.per_token_group_quant_fp8_packed_for_deepgemm(

@@ -126,7 +130,7 @@ def test_per_token_group_quant_fp8_packed(

130 )

131

132 # Triton reference (row-major float32 scales, UE8M0)

129 - with patch("vllm.platforms.current_platform.is_cuda", return_value=False):

133 + with patch("vllm.platforms.current_platform.is_cuda_alike", return_value=False):

134 ref_q, ref_s = fp8_utils.per_token_group_quant_fp8(

135 x,

136 group_size,

@@ -157,7 +161,8 @@ def test_per_token_group_quant_fp8_packed(

161

162

163 @pytest.mark.skipif(

160 - not current_platform.is_cuda(), reason="DeepGEMM not available on this platform"

164 + not current_platform.is_cuda_alike(),

165 + reason="DeepGEMM not available on this platform",

166 )

167 def test_per_token_group_quant_fp8_packed_all_zero():

168 """All-zero input must produce well-defined UE8M0 scale bytes via the eps

@@ -216,7 +221,8 @@ def test_per_token_group_quant_fp8_packed_all_zero():

221

222

223 @pytest.mark.skipif(

219 - not current_platform.is_cuda(), reason="DeepGEMM not available on this platform"

224 + not current_platform.is_cuda_alike(),

225 + reason="DeepGEMM not available on this platform",

226 )

227 def test_per_token_group_quant_fp8_packed_mantissa_rounds_up():

228 """Inputs whose absmax/max_8bit produces a non-power-of-2 force the

@@ -244,7 +250,7 @@ def test_per_token_group_quant_fp8_packed_mantissa_rounds_up():

250 use_ue8m0=True,

251 )

252

247 - with patch("vllm.platforms.current_platform.is_cuda", return_value=False):

253 + with patch("vllm.platforms.current_platform.is_cuda_alike", return_value=False):

254 ref_q, ref_s = fp8_utils.per_token_group_quant_fp8(

255 x,

256 group_size,

@@ -286,7 +292,8 @@ def test_per_token_group_quant_fp8_packed_mantissa_rounds_up():

292 ],

293 )

294 @pytest.mark.skipif(

289 - not current_platform.is_cuda(), reason="DeepGEMM not available on this platform"

295 + not current_platform.is_cuda_alike(),

296 + reason="DeepGEMM not available on this platform",

297 )

298 def test_per_token_group_quant_fp8_packed_zero_fills_padded_output_q(

299 num_tokens, hidden_dim

@@ -305,8 +312,8 @@ def test_per_token_group_quant_fp8_packed_zero_fills_padded_output_q(

312 k_num_packed = (groups_per_row + 3) // 4

313 tma_aligned_mn = ((mn + 3) // 4) * 4

314

308 - fp8_dtype = torch.float8_e4m3fn

309 - finfo = torch.finfo(fp8_dtype)

315 + fp8_dtype = current_platform.fp8_dtype()

316 + fp8_min, fp8_max = get_fp8_min_max()

317 # Allocate output_q with the padded mn extent and pre-fill with 0xFF

318 # so the kernel cannot rely on a clean buffer.

319 out_q = torch.empty((tma_aligned_mn, hidden_dim), device=device, dtype=fp8_dtype)

@@ -320,11 +327,11 @@ def test_per_token_group_quant_fp8_packed_zero_fills_padded_output_q(

327 )

328

329 torch.ops._C.per_token_group_fp8_quant_packed(

323 - x, out_q, out_s_packed, group_size, 1e-10, finfo.min, finfo.max

330 + x, out_q, out_s_packed, group_size, 1e-10, fp8_min, fp8_max

331 )

332

333 # Live rows must match the Triton reference.

327 - with patch("vllm.platforms.current_platform.is_cuda", return_value=False):

334 + with patch("vllm.platforms.current_platform.is_cuda_alike", return_value=False):

335 ref_q, _ = fp8_utils.per_token_group_quant_fp8(x, group_size, use_ue8m0=True)

336 assert torch.equal(out_q[:mn], ref_q), "Live region mismatch"

337

@@ -356,7 +363,7 @@ def test_per_token_group_quant_int8(shape, group_size: int):

363 )

364

365 # triton ref

359 - with patch("vllm.platforms.current_platform.is_cuda", return_value=False):

366 + with patch("vllm.platforms.current_platform.is_cuda_alike", return_value=False):

367 ref_q, ref_s = int8_utils.per_token_group_quant_int8(

368 x,

369 group_size,

vllm/compilation/passes/fusion/matcher_utils.py

+2 -3

@@ -36,14 +36,13 @@ QUANT_OPS: dict[QuantKey, OpOverload] = {

36 kFp8StaticTensorSym: torch.ops._C.static_scaled_fp8_quant.default, # noqa: E501

37 kFp8DynamicTensorSym: torch.ops._C.dynamic_scaled_fp8_quant.default, # noqa: E501

38 kFp8DynamicTokenSym: torch.ops._C.dynamic_per_token_scaled_fp8_quant.default, # noqa: E501

39 + kFp8Dynamic128Sym: torch.ops._C.per_token_group_fp8_quant.default, # noqa: E501

40 + kFp8Dynamic64Sym: torch.ops._C.per_token_group_fp8_quant.default, # noqa: E501

41 }

43 if current_platform.is_cuda() and hasattr(torch.ops._C, "scaled_fp4_quant"):

44 QUANT_OPS[kNvfp4Dynamic] = torch.ops._C.scaled_fp4_quant.out # noqa: E501

44 -if current_platform.is_cuda():

45 - QUANT_OPS[kFp8Dynamic128Sym] = torch.ops._C.per_token_group_fp8_quant.default # noqa: E501

46 - QUANT_OPS[kFp8Dynamic64Sym] = torch.ops._C.per_token_group_fp8_quant.default # noqa: E501

47 SILU_MUL_OP = torch.ops._C.silu_and_mul.default

vllm/compilation/passes/fusion/rms_quant_fusion.py

+29 -34

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

16 GroupShape,

17 QuantKey,

18 ScaleDesc,

19 + get_fp8_min_max,

20 kFp8Dynamic64Sym,

21 kFp8Dynamic128Sym,

22 kFp8DynamicTensorSym,

@@ -83,12 +84,11 @@ QUANT_OPS: dict[QuantKey, OpOverload] = {

84 kFp8StaticTensorSym: torch.ops._C.static_scaled_fp8_quant.default, # noqa: E501

85 kFp8DynamicTensorSym: torch.ops._C.dynamic_scaled_fp8_quant.default, # noqa: E501

86 kFp8DynamicTokenSym: torch.ops._C.dynamic_per_token_scaled_fp8_quant.default, # noqa: E501

87 + kFp8Dynamic128Sym: torch.ops._C.per_token_group_fp8_quant.default, # noqa: E501

88 + kFp8Dynamic64Sym: torch.ops._C.per_token_group_fp8_quant.default, # noqa: E501

89 }

90 if current_platform.is_cuda() and hasattr(torch.ops._C, "scaled_fp4_quant"):

91 QUANT_OPS[kNvfp4Dynamic] = torch.ops._C.scaled_fp4_quant.out

89 -if current_platform.is_cuda():

90 - QUANT_OPS[kFp8Dynamic128Sym] = torch.ops._C.per_token_group_fp8_quant.default # noqa: E501

91 - QUANT_OPS[kFp8Dynamic64Sym] = torch.ops._C.per_token_group_fp8_quant.default # noqa: E501

94 class FusedRMSQuantKey(NamedTuple):

@@ -327,9 +327,7 @@ class FusedAddRMSNormGroupQuantPattern(RMSNormQuantPattern):

327 dtype=self.quant_matcher.quant_key.dtype,

328 )

329 assert scale is not None

330 - finfo = torch.finfo(self.quant_matcher.quant_key.dtype)

331 - fp8_min = finfo.min

332 - fp8_max = finfo.max

330 + fp8_min, fp8_max = get_fp8_min_max()

331

332 _, result, scale = auto_functionalized(

333 self.quant_matcher.QUANT_OP,

@@ -430,9 +428,7 @@ class RMSNormGroupQuantPattern(RMSNormQuantPattern):

428 dtype=self.quant_matcher.quant_key.dtype,

429 )

430 assert scale is not None

433 - finfo = torch.finfo(self.quant_matcher.quant_key.dtype)

434 - fp8_min = finfo.min

435 - fp8_max = finfo.max

431 + fp8_min, fp8_max = get_fp8_min_max()

432

433 _, result, scale = auto_functionalized(

434 self.quant_matcher.QUANT_OP,

@@ -645,31 +641,30 @@ class RMSNormQuantFusionPass(VllmPatternMatcherPass):

641 # Fuse rms_norm + dynamic per-token fp8 quant

642 RMSNormDynamicQuantPattern(epsilon, FP8_DTYPE).register(self.patterns)

643

648 - # Only register group quant patterns on CUDA where the C++ op exists

649 - if current_platform.is_cuda():

650 - for group_shape in [GroupShape(1, 128), GroupShape(1, 64)]:

651 - for has_col_major_scales in [True, False]:

652 - for is_e8m0 in [True, False]:

653 - for is_tma_aligned in [False, True]:

654 - # Fuse fused_add_rms_norm + fp8 group quant

655 - FusedAddRMSNormGroupQuantPattern(

656 - epsilon,

657 - FP8_DTYPE,

658 - group_shape=group_shape,

659 - is_e8m0=is_e8m0,

660 - has_col_major_scales=has_col_major_scales,

661 - is_tma_aligned=is_tma_aligned,

662 - ).register(self.patterns)

663 -

664 - # Fuse rms_norm + fp8 group quant

665 - RMSNormGroupQuantPattern(

666 - epsilon,

667 - FP8_DTYPE,

668 - group_shape=group_shape,

669 - is_e8m0=is_e8m0,

670 - has_col_major_scales=has_col_major_scales,

671 - is_tma_aligned=is_tma_aligned,

672 - ).register(self.patterns)

644 + # Only register group quant patterns on CUDA/ROCm where the C++ op exists

645 + for group_shape in [GroupShape(1, 128), GroupShape(1, 64)]:

646 + for has_col_major_scales in [True, False]:

647 + for is_e8m0 in [True, False]:

648 + for is_tma_aligned in [False, True]:

649 + # Fuse fused_add_rms_norm + fp8 group quant

650 + FusedAddRMSNormGroupQuantPattern(

651 + epsilon,

652 + FP8_DTYPE,

653 + group_shape=group_shape,

654 + is_e8m0=is_e8m0,

655 + has_col_major_scales=has_col_major_scales,

656 + is_tma_aligned=is_tma_aligned,

657 + ).register(self.patterns)

658 +

659 + # Fuse rms_norm + fp8 group quant

660 + RMSNormGroupQuantPattern(

661 + epsilon,

662 + FP8_DTYPE,

663 + group_shape=group_shape,

664 + is_e8m0=is_e8m0,

665 + has_col_major_scales=has_col_major_scales,

666 + is_tma_aligned=is_tma_aligned,

667 + ).register(self.patterns)

668

669 self.dump_patterns(config, self.patterns)

670

vllm/model_executor/layers/quantization/input_quant_fp8.py

-5

@@ -158,11 +158,6 @@ class QuantFP8(CustomOp):

158 if use_aiter_per_token_quant:

159 return rocm_aiter_ops.per_token_quant(x, _FP8_DTYPE, scale)

160

161 - # Fallback to native implementation for group quantization.

162 - if self.is_group_quant:

163 - assert scale is None, "Dynamic group quantization does not use scale"

164 - return self._quantize_group_native(x)

165 -

161 # Fallback to CUDA implementation

162 return self.forward_cuda(x, scale, scale_ub)

163

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

+8 -7

@@ -575,7 +575,9 @@ def per_token_group_quant_fp8(

575

576 # prefer CUDA/XPU kernel if available

577 # TODO(bnell): this causes some fp8 moe test to fail.

578 - if (current_platform.is_cuda() or current_platform.is_xpu()) and x.is_contiguous():

578 + if (

579 + current_platform.is_cuda_alike() or current_platform.is_xpu()

580 + ) and x.is_contiguous():

581 torch.ops._C.per_token_group_fp8_quant(

582 x,

583 x_q,

@@ -664,8 +666,7 @@ def per_token_group_quant_fp8_packed_for_deepgemm(

666 )

667 assert x.stride(-1) == 1, "`x` groups must be contiguous"

668

667 - finfo = torch.finfo(dtype)

668 - fp8_min, fp8_max = finfo.min, finfo.max

669 + fp8_min, fp8_max = get_fp8_min_max()

670

671 # compute DeepGEMM-style packed scale tensor shape.

672 hidden_dim = x.shape[-1]

@@ -681,10 +682,10 @@ def per_token_group_quant_fp8_packed_for_deepgemm(

682 dtype=torch.int32,

683 )

684

684 - # CUDA kernel path only (DeepGEMM + E8M0 is CUDA-specific).

685 - assert current_platform.is_cuda(), (

686 - "per_token_group_quant_fp8_packed_for_deepgemm is only valid on CUDA "

687 - "platforms using DeepGEMM."

685 + # Native kernel (libtorch stable); used with DeepGEMM on CUDA and

686 + # available on ROCm for the same packed UE8M0 scale layout.

687 + assert current_platform.is_cuda_alike(), (

688 + "per_token_group_quant_fp8_packed_for_deepgemm requires a CUDA or ROCm GPU."

689 )

690

691 x_contiguous = x.contiguous()

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

+2 -2

@@ -235,8 +235,8 @@ def per_token_group_quant_int8(

235 device=x.device,

236 dtype=torch.float32,

237 )

238 - # prefer CUDA kernel if available

239 - if current_platform.is_cuda():

238 + # Prefer native stable kernel on CUDA/ROCm when available.

239 + if current_platform.is_cuda_alike():

240 torch.ops._C.per_token_group_quant_int8(

241 x, x_q, x_s, group_size, eps, float(int8_min), float(int8_max)

242 )