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

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 8, 03:45 UTC

[Bugfix] Fix quant RMS norm fusion for quantization with TMA-aligned scales (#33255)

Signed-off-by: ElizaWszola <ewszola@redhat.com> Co-authored-by: Luka Govedič <ProExpertProg@users.noreply.github.com>

ElizaWszola committed Feb 18, 2026 at 08:35 UTC a88b3be7c4c494c896f4a88ca8e6bfc1083625e0

12 files changed +235 -76

csrc/ops.h

+3 -1

@@ -315,7 +315,9 @@ void silu_and_mul_scaled_fp4_experts_quant(

315 void per_token_group_quant_fp8(const torch::Tensor& input,

316 torch::Tensor& output_q, torch::Tensor& output_s,

317 int64_t group_size, double eps, double fp8_min,

318 - double fp8_max, bool scale_ue8m0);

318 + double fp8_max, bool scale_ue8m0,

319 + bool dummy_is_scale_transposed,

320 + bool dummy_is_tma_aligned);

321

322 void per_token_group_quant_int8(const torch::Tensor& input,

323 torch::Tensor& output_q,

csrc/quantization/fused_kernels/fused_layernorm_dynamic_per_token_quant.cu

+12 -4

@@ -97,7 +97,7 @@ __global__ void rms_norm_per_block_quant_kernel(

97 scalar_t const* __restrict__ input, // [..., hidden_size]

98 scalar_t const* __restrict__ weight, // [hidden_size]

99 float const* scale_ub, float const var_epsilon, int32_t const hidden_size,

100 - scalar_t* __restrict__ residual = nullptr) {

100 + scalar_t* __restrict__ residual = nullptr, int64_t outer_scale_stride = 1) {

101 float rms;

102 // Compute RMS

103 // Always able to vectorize due to constraints on hidden_size

@@ -108,7 +108,8 @@ __global__ void rms_norm_per_block_quant_kernel(

108 // Always able to vectorize due to constraints on hidden_size and group_size

109 vllm::vectorized::compute_dynamic_per_token_scales<

110 scalar_t, scalar_out_t, has_residual, is_scale_transposed, group_size>(

111 - nullptr, scales, input, weight, rms, scale_ub, hidden_size, residual);

111 + nullptr, scales, input, weight, rms, scale_ub, hidden_size, residual,

112 + outer_scale_stride);

113

114 // RMS Norm + Quant

115 // Always able to vectorize due to constraints on hidden_size

@@ -119,7 +120,8 @@ __global__ void rms_norm_per_block_quant_kernel(

120 vllm::vectorized::norm_and_quant<

121 scalar_t, scalar_out_t, std::is_same_v<scalar_out_t, int8_t>,

122 has_residual, is_scale_transposed, group_size>(

122 - out, input, weight, rms, scales, hidden_size, residual);

123 + out, input, weight, rms, scales, hidden_size, residual,

124 + outer_scale_stride);

125 }

126

127 } // namespace vllm

@@ -225,7 +227,8 @@ void rms_norm_per_block_quant_dispatch(

227 : nullptr,

228 var_epsilon, hidden_size,

229 has_residual ? residual->data_ptr<scalar_in_t>()

228 - : nullptr);

230 + : nullptr,

231 + scales.stride(1));

232 });

233 });

234 });

@@ -257,6 +260,11 @@ void rms_norm_per_block_quant(torch::Tensor& out, torch::Tensor const& input,

260 TORCH_CHECK(group_size == 128 || group_size == 64,

261 "Unsupported group size: ", group_size);

262

263 + if (scales.stride(1) > 1) {

264 + TORCH_CHECK(is_scale_transposed,

265 + "Outer scale stride must be 1 when scales are not transposed");

266 + }

267 +

268 rms_norm_per_block_quant_dispatch(out, input, weight, scales, group_size,

269 var_epsilon, scale_ub, residual,

270 is_scale_transposed);

csrc/quantization/fused_kernels/layernorm_utils.cuh

+22 -15

@@ -74,7 +74,7 @@ __device__ void compute_dynamic_per_token_scales(

74 scalar_t const* __restrict__ input, scalar_t const* __restrict__ weight,

75 float const rms, float const* __restrict__ scale_ub,

76 int32_t const hidden_size, scalar_t const* __restrict__ residual = nullptr,

77 - int32_t const group_size = 0) {

77 + int32_t const group_size = 0, int64_t outer_scale_stride = 1) {

78 float block_absmax_val_maybe = 0.0f;

79 constexpr scalar_out_t qmax{quant_type_max_v<scalar_out_t>};

80 __syncthreads();

@@ -133,7 +133,9 @@ __device__ void compute_dynamic_per_token_scales(

133 scale = max(scale / qmax, min_scaling_factor<scalar_out_t>::val());

134 // Global output store

135 if constexpr (is_scale_transposed) {

136 - all_token_scales[(threadIdx.x / threads_per_group) * gridDim.x +

136 + int64_t const scale_rows = (gridDim.x + outer_scale_stride - 1) /

137 + outer_scale_stride * outer_scale_stride;

138 + all_token_scales[(threadIdx.x / threads_per_group) * scale_rows +

139 blockIdx.x] = scale;

140 } else {

141 all_token_scales[blockIdx.x * num_groups +

@@ -180,13 +182,11 @@ __device__ void compute_dynamic_per_token_scales(

182

183 template <typename scalar_t, typename scalar_out_t, bool is_scale_inverted,

184 bool has_residual = false, bool is_scale_transposed = false>

183 -__device__ void norm_and_quant(scalar_out_t* __restrict__ output,

184 - scalar_t const* __restrict__ input,

185 - scalar_t const* __restrict__ weight,

186 - float const rms, float* const scale,

187 - int32_t const hidden_size,

188 - scalar_t* __restrict__ residual = nullptr,

189 - int32_t const group_size = 0) {

185 +__device__ void norm_and_quant(

186 + scalar_out_t* __restrict__ output, scalar_t const* __restrict__ input,

187 + scalar_t const* __restrict__ weight, float const rms, float* const scale,

188 + int32_t const hidden_size, scalar_t* __restrict__ residual = nullptr,

189 + int32_t const group_size = 0, int64_t outer_scale_stride = 1) {

190 int64_t const token_offset = blockIdx.x * static_cast<int64_t>(hidden_size);

191

192 for (auto i = threadIdx.x; i < hidden_size; i += blockDim.x) {

@@ -202,7 +202,9 @@ __device__ void norm_and_quant(scalar_out_t* __restrict__ output,

202 int64_t scale_idx = 0;

203 if (group_size > 0) {

204 if constexpr (is_scale_transposed) {

205 - scale_idx = (i / group_size) * gridDim.x + blockIdx.x;

205 + int64_t const scale_rows = (gridDim.x + outer_scale_stride - 1) /

206 + outer_scale_stride * outer_scale_stride;

207 + scale_idx = (i / group_size) * scale_rows + blockIdx.x;

208 } else {

209 scale_idx = blockIdx.x * (hidden_size / group_size) + i / group_size;

210 }

@@ -286,8 +288,8 @@ __device__ void compute_dynamic_per_token_scales(

288 float* __restrict__ token_scale, float* __restrict__ all_token_scales,

289 scalar_t const* __restrict__ input, scalar_t const* __restrict__ weight,

290 float const rms, float const* __restrict__ scale_ub,

289 - int32_t const hidden_size,

290 - scalar_t const* __restrict__ residual = nullptr) {

291 + int32_t const hidden_size, scalar_t const* __restrict__ residual = nullptr,

292 + int64_t outer_scale_stride = 1) {

293 constexpr scalar_out_t qmax{quant_type_max_v<scalar_out_t>};

294

295 const int VEC_SIZE = 4;

@@ -382,7 +384,9 @@ __device__ void compute_dynamic_per_token_scales(

384 scale = max(scale / qmax, min_scaling_factor<scalar_out_t>::val());

385 // Global output store

386 if constexpr (is_scale_transposed) {

385 - all_token_scales[(threadIdx.x / threads_per_group) * gridDim.x +

387 + int64_t const scale_rows = (gridDim.x + outer_scale_stride - 1) /

388 + outer_scale_stride * outer_scale_stride;

389 + all_token_scales[(threadIdx.x / threads_per_group) * scale_rows +

390 blockIdx.x] = scale;

391 } else {

392 all_token_scales[blockIdx.x * num_groups +

@@ -463,7 +467,8 @@ __device__ void norm_and_quant(scalar_out_t* __restrict__ output,

467 scalar_t const* __restrict__ weight,

468 float const rms, float* const scale,

469 int32_t const hidden_size,

466 - scalar_t* __restrict__ residual = nullptr) {

470 + scalar_t* __restrict__ residual = nullptr,

471 + int64_t outer_scale_stride = 1) {

472 int64_t const token_offset = blockIdx.x * static_cast<int64_t>(hidden_size);

473

474 // Vectorized input/output/weight/residual to better utilize memory bandwidth.

@@ -516,7 +521,9 @@ __device__ void norm_and_quant(scalar_out_t* __restrict__ output,

521 int64_t const num_groups = hidden_size / group_size;

522 int64_t scale_idx = 0;

523 if constexpr (is_scale_transposed) {

519 - scale_idx = (i * VEC_SIZE / group_size) * gridDim.x + blockIdx.x;

524 + int64_t const scale_rows = (gridDim.x + outer_scale_stride - 1) /

525 + outer_scale_stride * outer_scale_stride;

526 + scale_idx = (i * VEC_SIZE / group_size) * scale_rows + blockIdx.x;

527 } else {

528 scale_idx = blockIdx.x * num_groups + i * VEC_SIZE / group_size;

529 }

csrc/quantization/w8a8/fp8/per_token_group_quant.cu

+3 -1

@@ -379,7 +379,9 @@ void per_token_group_quant_8bit_packed(const torch::Tensor& input,

379 void per_token_group_quant_fp8(const torch::Tensor& input,

380 torch::Tensor& output_q, torch::Tensor& output_s,

381 int64_t group_size, double eps, double fp8_min,

382 - double fp8_max, bool scale_ue8m0) {

382 + double fp8_max, bool scale_ue8m0,

383 + bool dummy_is_scale_transposed = false,

384 + bool dummy_is_tma_aligned = false) {

385 per_token_group_quant_8bit(input, output_q, output_s, group_size, eps,

386 fp8_min, fp8_max, scale_ue8m0);

387 }

\ No newline at end of file

csrc/torch_bindings.cpp

+3 -1

@@ -643,11 +643,13 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {

643

644 #ifndef USE_ROCM

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

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

647 ops.def(

648 "per_token_group_fp8_quant(Tensor input, Tensor! output_q, Tensor! "

649 "output_s, "

650 "int group_size, float eps, float fp8_min, float fp8_max, bool "

650 - "scale_ue8m0) -> ()");

651 + "scale_ue8m0, bool dummy_is_scale_transposed, bool dummy_is_tma_aligned "

652 + ") -> ()");

653 ops.impl("per_token_group_fp8_quant", torch::kCUDA,

654 &per_token_group_quant_fp8);

655

tests/compile/fusions_e2e/test_tp1_quant.py

+2 -4

@@ -50,10 +50,9 @@ def test_tp1_fp8_fusions(

50 run_e2e_fusion_test,

51 monkeypatch,

52 ):

53 - if use_deepgemm:

54 - # TODO(luka/eliza) DeepGEMM uses different quants, matching not supported

53 + if use_deepgemm and is_blackwell():

54 + # TODO(luka) DeepGEMM uses different quants, matching not supported

55 # - on Blackwell, uses a special quant fp8, currently not supported

56 - # - on Hopper, tma-aligned scales inhibit matching (fix WIP)

56 pytest.skip("DeepGEMM & quant matching not currently supported")

58 matches = matches_fn(n_layers)

@@ -66,7 +65,6 @@ def test_tp1_fp8_fusions(

65 model_kwargs["hf_overrides"] = hf_overrides(n_layers)

66 model_kwargs["load_format"] = "dummy"

67 model_kwargs["max_model_len"] = 1024

69 -

68 compilation_config = dict(

69 use_inductor_graph_partition=inductor_graph_partition,

70 custom_ops=custom_ops.split(","),

tests/kernels/core/test_fused_quant_layernorm.py

+36 -5

@@ -2,6 +2,8 @@

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

5 +import itertools

6 +

7 import pytest

8 import torch

@@ -21,7 +23,7 @@ QUANT_DTYPES = [torch.int8, current_platform.fp8_dtype()]

23 VEC_HIDDEN_SIZES = [1024, 1025, 1027, 1029]

24 # Avoid combinatorial explosion with full Cartesian product

25 NUM_TOKENS_HIDDEN_SIZES = [

24 - *[(1, i) for i in [1, 64, *VEC_HIDDEN_SIZES, 5120, 5137]],

26 + *[(1, i) for i in [1, 64, 128, *VEC_HIDDEN_SIZES, 5120, 5137]],

27 *[(2048, i) for i in [1, 64, *VEC_HIDDEN_SIZES, 5137]],

28 *[(4096, i) for i in [1, 64, 5137]],

29 ]

@@ -29,6 +31,7 @@ NUM_TOKENS_HIDDEN_SIZES = [

31 ADD_RESIDUAL = [False, True]

32 SCALE_UBS = [True, False]

33 GROUP_SIZES = [None, [1, 64], [1, 128]]

34 +TMA_ALIGNMENTS = [0, 4]

35 SEEDS = [0]

36 CUDA_DEVICES = [f"cuda:{i}" for i in range(1 if torch.cuda.device_count() == 1 else 2)]

@@ -110,12 +113,21 @@ def ops_dynamic_per_token_or_block_quant(

113 residual: torch.Tensor | None,

114 scale_ub: torch.Tensor | None,

115 group_size: list[int] | None,

116 + tma_alignment: int,

117 ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor | None]:

118 if residual is not None:

119 residual = residual.clone()

120 if group_size is not None:

121 out, scales = ops.rms_norm_per_block_quant(

118 - x, weight, EPS, quant_dtype, group_size, scale_ub, residual, True

122 + x,

123 + weight,

124 + EPS,

125 + quant_dtype,

126 + group_size,

127 + scale_ub,

128 + residual,

129 + True,

130 + tma_alignment,

131 )

132 scales = scales.contiguous()

133 else:

@@ -132,9 +144,10 @@ def ops_impl(

144 residual: torch.Tensor | None,

145 scale_ub: torch.Tensor | None,

146 group_size: list[int] | None,

147 + tma_alignment: int,

148 ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor | None]:

149 return ops_dynamic_per_token_or_block_quant(

137 - weight, x, quant_dtype, residual, scale_ub, group_size

150 + weight, x, quant_dtype, residual, scale_ub, group_size, tma_alignment

151 )

152

153

@@ -143,7 +156,10 @@ def ops_impl(

156 @pytest.mark.parametrize("has_scale_ub", SCALE_UBS)

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

158 @pytest.mark.parametrize("quant_dtype", QUANT_DTYPES)

146 -@pytest.mark.parametrize("group_size", GROUP_SIZES)

159 +@pytest.mark.parametrize(

160 + "group_size, tma_alignment",

161 + [(None, 0), *itertools.product(GROUP_SIZES, TMA_ALIGNMENTS)],

162 +)

163 @pytest.mark.parametrize("seed", SEEDS)

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

165 @torch.inference_mode()

@@ -156,6 +172,7 @@ def test_rms_norm(

172 dtype: torch.dtype,

173 quant_dtype: torch.dtype,

174 group_size: list[int] | None,

175 + tma_alignment: int,

176 seed: int,

177 device: str,

178 ) -> None:

@@ -173,6 +190,20 @@ def test_rms_norm(

190 # blockwise baseline doesn't support scale_ub

191 return

192

193 + if (

194 + group_size is None or quant_dtype != current_platform.fp8_dtype()

195 + ) and tma_alignment != 0:

196 + # TMA alignment is only supported for groupwise fp8 kernels

197 + return

198 +

199 + if (

200 + group_size is not None

201 + and tma_alignment != 0

202 + and hidden_size // group_size[1] % tma_alignment == 0

203 + ):

204 + # Skip tests where TMA alignment doesn't create extra padding to save time

205 + return

206 +

207 if has_scale_ub and quant_dtype != current_platform.fp8_dtype():

208 # skip

209 return

@@ -196,7 +227,7 @@ def test_rms_norm(

227 layer, x, quant_dtype, residual, scale_ub, group_size

228 )

229 ops_out, ops_scales, ops_residual = ops_impl(

199 - layer.weight, x, quant_dtype, residual, scale_ub, group_size

230 + layer.weight, x, quant_dtype, residual, scale_ub, group_size, tma_alignment

231 )

232

233 assert ref_out.dtype == quant_dtype

vllm/_custom_ops.py

+24 -5

@@ -450,15 +450,30 @@ def rms_norm_per_block_quant(

450 scale_ub: torch.Tensor | None = None,

451 residual: torch.Tensor | None = None,

452 is_scale_transposed: bool = False,

453 + tma_alignment: int = 0,

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

455 assert len(group_size) == 2

456 output = torch.empty_like(input, dtype=quant_dtype)

457 if is_scale_transposed:

457 - scales = torch.empty(

458 - (input.shape[-1] // group_size[1], input.numel() // input.shape[-1]),

459 - device=input.device,

460 - dtype=torch.float32,

461 - ).transpose(0, 1)

458 + if tma_alignment == 0:

459 + scales = torch.empty(

460 + (input.shape[-1] // group_size[1], input.numel() // input.shape[-1]),

461 + device=input.device,

462 + dtype=torch.float32,

463 + ).transpose(0, 1)

464 + else:

465 + m = input.shape[-2]

466 + sf_k = input.shape[-1] // group_size[1]

467 + tma_aligned_m = (m + tma_alignment - 1) // tma_alignment * tma_alignment

468 + shape = input.shape[:-2] + (m, sf_k)

469 + stride = (

470 + (1, tma_aligned_m)

471 + if input.dim() == 2

472 + else (tma_aligned_m * sf_k, 1, tma_aligned_m)

473 + )

474 + scales = torch.empty_strided(

475 + shape, stride, device=input.device, dtype=torch.float32

476 + )

477 else:

478 scales = torch.empty(

479 (input.numel() // input.shape[-1], input.shape[-1] // group_size[1]),

@@ -466,6 +481,10 @@ def rms_norm_per_block_quant(

481 dtype=torch.float32,

482 )

483

484 + assert tma_alignment in [0, 4], "Expected TMA alignment 0 or 4, but got " + str(

485 + tma_alignment

486 + )

487 +

488 torch.ops._C.rms_norm_per_block_quant(

489 output,

490 input,

vllm/compilation/passes/fusion/matcher_utils.py

+10 -2

@@ -292,6 +292,7 @@ class MatcherQuantFP8(MatcherCustomOp):

292 has_col_major_scales: bool = False,

293 is_e8m0: bool = False,

294 match_rocm_aiter: bool = False,

295 + is_tma_aligned: bool = False,

296 ) -> None:

297 if enabled is None:

298 enabled = QuantFP8.enabled()

@@ -301,6 +302,7 @@ class MatcherQuantFP8(MatcherCustomOp):

302 self.has_col_major_scales = has_col_major_scales

303 self.is_e8m0 = is_e8m0

304 self.match_rocm_aiter = match_rocm_aiter

305 + self.is_tma_aligned = is_tma_aligned

306

307 if match_rocm_aiter:

308 assert not quant_key.scale.group_shape.is_per_tensor(), (

@@ -336,6 +338,7 @@ class MatcherQuantFP8(MatcherCustomOp):

338 quant_key.scale.group_shape,

339 column_major_scales=has_col_major_scales,

340 use_ue8m0=is_e8m0,

341 + tma_aligned_scales=self.is_tma_aligned,

342 compile_native=False,

343 )

344

@@ -367,8 +370,11 @@ class MatcherQuantFP8(MatcherCustomOp):

370 )

371

372 if self.quant_key.scale.group_shape.is_per_group():

370 - assert scale is None

371 - scale = self.make_scale(input, transposed=self.has_col_major_scales)

373 + # for tma_aligned, the scale must be passed to forward_custom

374 + # tma_aligned fusion then matches by custom op arguments

375 + if not self.is_tma_aligned:

376 + assert scale is None

377 + scale = self.make_scale(input, transposed=self.has_col_major_scales)

378

379 finfo = torch.finfo(self.quant_key.dtype)

380 fp8_min = finfo.min

@@ -384,6 +390,8 @@ class MatcherQuantFP8(MatcherCustomOp):

390 fp8_min=fp8_min,

391 fp8_max=fp8_max,

392 scale_ue8m0=self.is_e8m0,

393 + dummy_is_scale_transposed=self.has_col_major_scales,

394 + dummy_is_tma_aligned=self.is_tma_aligned,

395 )

396 return result, scale

397

vllm/compilation/passes/fusion/rms_quant_fusion.py

+109 -36

@@ -121,6 +121,7 @@ class RMSNormQuantPattern:

121 key: FusedRMSQuantKey,

122 has_col_major_scales: bool = False,

123 is_e8m0: bool = False,

124 + is_tma_aligned: bool = False,

125 ) -> None:

126 self.epsilon = epsilon

127 self.quant_dtype = key.quant.dtype

@@ -136,7 +137,10 @@ class RMSNormQuantPattern:

137 else MatcherFusedAddRMSNorm(epsilon)

138 )

139 self.quant_matcher = MatcherQuantFP8(

139 - key.quant, has_col_major_scales=has_col_major_scales, is_e8m0=is_e8m0

140 + key.quant,

141 + has_col_major_scales=has_col_major_scales,

142 + is_e8m0=is_e8m0,

143 + is_tma_aligned=is_tma_aligned,

144 )

145

146

@@ -262,8 +266,9 @@ class FusedAddRMSNormGroupQuantPattern(RMSNormQuantPattern):

266 quant_dtype: torch.dtype,

267 group_shape: GroupShape,

268 symmetric: bool = True,

265 - has_col_major_scales: bool = False,

269 is_e8m0: bool = False,

270 + has_col_major_scales: bool = True,

271 + is_tma_aligned: bool = True,

272 ) -> None:

273 scale = ScaleDesc(torch.float32, False, group_shape)

274 key = FusedRMSQuantKey(

@@ -271,29 +276,63 @@ class FusedAddRMSNormGroupQuantPattern(RMSNormQuantPattern):

276 quant=QuantKey(dtype=quant_dtype, scale=scale, symmetric=symmetric),

277 )

278 self.group_shape = group_shape

274 - self.has_col_major_scales = has_col_major_scales

279 self.is_e8m0 = is_e8m0

280 + self.has_col_major_scales = has_col_major_scales

281 + self.is_tma_aligned = is_tma_aligned

282 super().__init__(

277 - epsilon, key, has_col_major_scales=has_col_major_scales, is_e8m0=is_e8m0

283 + epsilon,

284 + key,

285 + has_col_major_scales=has_col_major_scales,

286 + is_e8m0=is_e8m0,

287 + is_tma_aligned=is_tma_aligned,

288 )

289

290 def register(self, pm_pass: PatternMatcherPass) -> None:

291 def pattern(

282 - input: torch.Tensor, weight: torch.Tensor, residual: torch.Tensor

292 + input: torch.Tensor,

293 + weight: torch.Tensor,

294 + residual: torch.Tensor,

295 + scale: torch.Tensor,

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

297 result_rms, residual = self.rmsnorm_matcher(input, weight, residual)

285 - result, scale = self.quant_matcher(result_rms)

298 + result = torch.empty(

299 + result_rms.shape,

300 + device=result_rms.device,

301 + dtype=self.quant_matcher.quant_key.dtype,

302 + )

303 + assert scale is not None

304 + finfo = torch.finfo(self.quant_matcher.quant_key.dtype)

305 + fp8_min = finfo.min

306 + fp8_max = finfo.max

307 +

308 + _, result, scale = auto_functionalized(

309 + self.quant_matcher.QUANT_OP,

310 + input=result_rms,

311 + output_q=result,

312 + output_s=scale,

313 + group_size=self.quant_matcher.quant_key.scale.group_shape[1],

314 + eps=1e-10,

315 + fp8_min=fp8_min,

316 + fp8_max=fp8_max,

317 + scale_ue8m0=self.quant_matcher.is_e8m0,

318 + dummy_is_scale_transposed=self.has_col_major_scales,

319 + dummy_is_tma_aligned=self.is_tma_aligned,

320 + )

321 +

322 return result, residual, scale

323

324 def replacement(

289 - input: torch.Tensor, weight: torch.Tensor, residual: torch.Tensor

325 + input: torch.Tensor,

326 + weight: torch.Tensor,

327 + residual: torch.Tensor,

328 + scale: torch.Tensor,

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

330 # In case we're matching native rms-norm, conversions might be

331 # optimized out. We convert here just to be safe.

332 input = input.to(dtype=self.model_dtype)

333

334 result = torch.empty_like(input, dtype=self.quant_dtype)

296 - scale = self.quant_matcher.make_scale(input, self.has_col_major_scales)

335 +

336 at = auto_functionalized(

337 self.FUSED_OP,

338 result=result,

@@ -310,10 +349,12 @@ class FusedAddRMSNormGroupQuantPattern(RMSNormQuantPattern):

349 # result, residual, scale

350 return at[1], at[3], at[2]

351

352 + scale = self.quant_matcher.empty_f32(1, 1)

353 +

354 pm.register_replacement(

355 pattern,

356 replacement,

316 - self.rmsnorm_matcher.inputs(),

357 + self.rmsnorm_matcher.inputs() + [scale],

358 pm.fwd_only,

359 pm_pass,

360 )

@@ -326,8 +367,9 @@ class RMSNormGroupQuantPattern(RMSNormQuantPattern):

367 quant_dtype: torch.dtype,

368 group_shape: GroupShape,

369 symmetric: bool = True,

329 - has_col_major_scales: bool = False,

370 is_e8m0: bool = False,

371 + has_col_major_scales: bool = True,

372 + is_tma_aligned: bool = True,

373 ) -> None:

374 scale = ScaleDesc(torch.float32, False, group_shape)

375 key = FusedRMSQuantKey(

@@ -335,29 +377,55 @@ class RMSNormGroupQuantPattern(RMSNormQuantPattern):

377 quant=QuantKey(dtype=quant_dtype, scale=scale, symmetric=symmetric),

378 )

379 self.group_shape = group_shape

380 + self.has_col_major_scales = has_col_major_scales

381 + self.is_tma_aligned = is_tma_aligned

382 super().__init__(

339 - epsilon, key, has_col_major_scales=has_col_major_scales, is_e8m0=is_e8m0

383 + epsilon,

384 + key,

385 + has_col_major_scales=self.has_col_major_scales,

386 + is_e8m0=is_e8m0,

387 + is_tma_aligned=is_tma_aligned,

388 )

389

390 def register(self, pm_pass: PatternMatcherPass) -> None:

391 def pattern(

344 - input: torch.Tensor, weight: torch.Tensor

392 + input: torch.Tensor, weight: torch.Tensor, scale: torch.Tensor

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

394 result_rms = self.rmsnorm_matcher(input, weight)

347 - result, scale = self.quant_matcher(result_rms)

395 + result = torch.empty(

396 + result_rms.shape,

397 + device=result_rms.device,

398 + dtype=self.quant_matcher.quant_key.dtype,

399 + )

400 + assert scale is not None

401 + finfo = torch.finfo(self.quant_matcher.quant_key.dtype)

402 + fp8_min = finfo.min

403 + fp8_max = finfo.max

404 +

405 + _, result, scale = auto_functionalized(

406 + self.quant_matcher.QUANT_OP,

407 + input=result_rms,

408 + output_q=result,

409 + output_s=scale,

410 + group_size=self.quant_matcher.quant_key.scale.group_shape[1],

411 + eps=1e-10,

412 + fp8_min=fp8_min,

413 + fp8_max=fp8_max,

414 + scale_ue8m0=self.quant_matcher.is_e8m0,

415 + dummy_is_scale_transposed=self.has_col_major_scales,

416 + dummy_is_tma_aligned=self.is_tma_aligned,

417 + )

418 +

419 return result, scale

420

421 def replacement(

351 - input: torch.Tensor, weight: torch.Tensor

422 + input: torch.Tensor, weight: torch.Tensor, scale: torch.Tensor

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

424 # In case we're matching native rms-norm, conversions might be

425 # optimized out. We convert here just to be safe.

426 input = input.to(dtype=self.model_dtype)

427

428 result = torch.empty_like(input, dtype=self.quant_dtype)

358 - scale = self.quant_matcher.make_scale(

359 - input, transposed=self.quant_matcher.has_col_major_scales

360 - )

429 at = auto_functionalized(

430 self.FUSED_OP,

431 result=result,

@@ -368,16 +436,18 @@ class RMSNormGroupQuantPattern(RMSNormQuantPattern):

436 scale_ub=None,

437 residual=None,

438 group_size=self.group_shape[1],

371 - is_scale_transposed=self.quant_matcher.has_col_major_scales,

439 + is_scale_transposed=self.has_col_major_scales,

440 )

441

442 # result, scale

443 return at[1], at[2]

444

445 + scale = self.quant_matcher.empty_f32(1, 1)

446 +

447 pm.register_replacement(

448 pattern,

449 replacement,

380 - self.rmsnorm_matcher.inputs(),

450 + self.rmsnorm_matcher.inputs() + [scale],

451 pm.fwd_only,

452 pm_pass,

453 )

@@ -532,23 +602,26 @@ class RMSNormQuantFusionPass(VllmPatternMatcherPass):

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

603 for has_col_major_scales in [True, False]:

604 for is_e8m0 in [True, False]:

535 - # Fuse fused_add_rms_norm + fp8 group quant

536 - FusedAddRMSNormGroupQuantPattern(

537 - epsilon,

538 - FP8_DTYPE,

539 - group_shape=group_shape,

540 - has_col_major_scales=has_col_major_scales,

541 - is_e8m0=is_e8m0,

542 - ).register(self.patterns)

543 -

544 - # Fuse rms_norm + fp8 group quant

545 - RMSNormGroupQuantPattern(

546 - epsilon,

547 - FP8_DTYPE,

548 - group_shape=group_shape,

549 - has_col_major_scales=has_col_major_scales,

550 - is_e8m0=is_e8m0,

551 - ).register(self.patterns)

605 + for is_tma_aligned in [False, True]:

606 + # Fuse fused_add_rms_norm + fp8 group quant

607 + FusedAddRMSNormGroupQuantPattern(

608 + epsilon,

609 + FP8_DTYPE,

610 + group_shape=group_shape,

611 + is_e8m0=is_e8m0,

612 + has_col_major_scales=has_col_major_scales,

613 + is_tma_aligned=is_tma_aligned,

614 + ).register(self.patterns)

615 +

616 + # Fuse rms_norm + fp8 group quant

617 + RMSNormGroupQuantPattern(

618 + epsilon,

619 + FP8_DTYPE,

620 + group_shape=group_shape,

621 + is_e8m0=is_e8m0,

622 + has_col_major_scales=has_col_major_scales,

623 + is_tma_aligned=is_tma_aligned,

624 + ).register(self.patterns)

625

626 self.dump_patterns(config, self.patterns)

627

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

+10 -1

@@ -924,7 +924,16 @@ def per_token_group_quant_fp8(

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

925 if current_platform.is_cuda() and x.is_contiguous():

926 torch.ops._C.per_token_group_fp8_quant(

927 - x, x_q, x_s, group_size, eps, fp8_min, fp8_max, use_ue8m0

927 + x,

928 + x_q,

929 + x_s,

930 + group_size,

931 + eps,

932 + fp8_min,

933 + fp8_max,

934 + use_ue8m0,

935 + column_major_scales,

936 + tma_aligned_scales,

937 )

938 return x_q, x_s

939

vllm/utils/deep_gemm.py

+1 -1

@@ -349,7 +349,7 @@ def _align(x: int, y: int) -> int:

349

350

351 # Taken from https://github.com/deepseek-ai/DeepGEMM/blob/v2.1.1/csrc/utils/math.hpp#L19

352 -def get_tma_aligned_size(x: int, element_size: int):

352 +def get_tma_aligned_size(x: int, element_size: int) -> int:

353 return _align(x, 16 // element_size)

354

355