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 3, 22:45 UTC
[Perf] Batch invariance with Cutlass fp8 support, 28.9% E2E latency improvement (#40408)
Signed-off-by: yewentao256 <zhyanwentao@126.com> Signed-off-by: Wentao Ye <44945378+yewentao256@users.noreply.github.com>
Wentao Ye committed
May 11, 2026 at 12:20 UTC
0d453e23364701f6fe2f20deafbb372b28714a86
11 files changed
+305
-3
csrc/libtorch_stable/quantization/w8a8/cutlass/c3x/scaled_mm_sm100_fp8.cu
+9
@@ -1,5 +1,6 @@
1
#include "scaled_mm_kernels.hpp"
2
#include "scaled_mm_sm100_fp8_dispatch.cuh"
3
+#include "core/batch_invariant.hpp"
4
5
namespace vllm {
6
@@ -13,9 +14,17 @@ void cutlass_scaled_mm_sm100_fp8(
14
STD_TORCH_CHECK(bias->scalar_type() == out.scalar_type(),
15
"currently bias dtype must match output dtype ",
16
out.scalar_type());
17
+ if (vllm_is_batch_invariant()) {
18
+ return cutlass_scaled_mm_sm100_fp8_batch_invariant_epilogue<true>(
19
+ out, a, b, a_scales, b_scales, *bias);
20
+ }
21
return cutlass_scaled_mm_sm100_fp8_epilogue<true>(out, a, b, a_scales,
22
b_scales, *bias);
23
} else {
24
+ if (vllm_is_batch_invariant()) {
25
+ return cutlass_scaled_mm_sm100_fp8_batch_invariant_epilogue<false>(
26
+ out, a, b, a_scales, b_scales);
27
+ }
28
return cutlass_scaled_mm_sm100_fp8_epilogue<false>(out, a, b, a_scales,
29
b_scales);
30
}
csrc/libtorch_stable/quantization/w8a8/cutlass/c3x/scaled_mm_sm100_fp8_dispatch.cuh
+52
@@ -294,6 +294,34 @@ inline void cutlass_gemm_sm100_fp8_dispatch(
294
}
295
}
296
297
+template <typename InType, typename OutType, bool EnableBias,
298
+ typename... EpilogueArgs>
299
+inline void cutlass_gemm_sm100_fp8_batch_invariant_dispatch(
300
+ torch::stable::Tensor& out, torch::stable::Tensor const& a,
301
+ torch::stable::Tensor const& b, torch::stable::Tensor const& a_scales,
302
+ torch::stable::Tensor const& b_scales, EpilogueArgs&&... args) {
303
+ static_assert(std::is_same<InType, cutlass::float_e4m3_t>());
304
+ STD_TORCH_CHECK(a.scalar_type() ==
305
+ torch::headeronly::ScalarType::Float8_e4m3fn);
306
+ STD_TORCH_CHECK(b.scalar_type() ==
307
+ torch::headeronly::ScalarType::Float8_e4m3fn);
308
+
309
+ using Cutlass3xGemmM64SwapAB =
310
+ typename sm100_fp8_config_M64_swap_ab<InType, OutType,
311
+ EnableBias>::Cutlass3xGemm;
312
+ using Cutlass3xGemmM64 =
313
+ typename sm100_fp8_config_M64<InType, OutType, EnableBias>::Cutlass3xGemm;
314
+
315
+ // keep the CUTLASS config independent of M for batch invariance
316
+ uint32_t const k = a.size(1);
317
+ if (k < 4096) {
318
+ return cutlass_gemm_caller_sm100_fp8<Cutlass3xGemmM64>(
319
+ out, a, b, a_scales, b_scales, std::forward<EpilogueArgs>(args)...);
320
+ }
321
+ return cutlass_gemm_caller_sm100_fp8<Cutlass3xGemmM64SwapAB>(
322
+ out, a, b, b_scales, a_scales, std::forward<EpilogueArgs>(args)...);
323
+}
324
+
325
template <bool EnableBias, typename... EpilogueArgs>
326
void cutlass_scaled_mm_sm100_fp8_epilogue(torch::stable::Tensor& out,
327
torch::stable::Tensor const& a,
@@ -320,4 +348,28 @@ void cutlass_scaled_mm_sm100_fp8_epilogue(torch::stable::Tensor& out,
348
}
349
}
350
351
+template <bool EnableBias, typename... EpilogueArgs>
352
+void cutlass_scaled_mm_sm100_fp8_batch_invariant_epilogue(
353
+ torch::stable::Tensor& out, torch::stable::Tensor const& a,
354
+ torch::stable::Tensor const& b, torch::stable::Tensor const& a_scales,
355
+ torch::stable::Tensor const& b_scales, EpilogueArgs&&... epilogue_args) {
356
+ STD_TORCH_CHECK(a.scalar_type() ==
357
+ torch::headeronly::ScalarType::Float8_e4m3fn);
358
+ STD_TORCH_CHECK(b.scalar_type() ==
359
+ torch::headeronly::ScalarType::Float8_e4m3fn);
360
+
361
+ if (out.scalar_type() == torch::headeronly::ScalarType::BFloat16) {
362
+ return cutlass_gemm_sm100_fp8_batch_invariant_dispatch<
363
+ cutlass::float_e4m3_t, cutlass::bfloat16_t, EnableBias>(
364
+ out, a, b, a_scales, b_scales,
365
+ std::forward<EpilogueArgs>(epilogue_args)...);
366
+ } else {
367
+ STD_TORCH_CHECK(out.scalar_type() == torch::headeronly::ScalarType::Half);
368
+ return cutlass_gemm_sm100_fp8_batch_invariant_dispatch<
369
+ cutlass::float_e4m3_t, cutlass::half_t, EnableBias>(
370
+ out, a, b, a_scales, b_scales,
371
+ std::forward<EpilogueArgs>(epilogue_args)...);
372
+ }
373
+}
374
+
375
} // namespace vllm
csrc/libtorch_stable/quantization/w8a8/cutlass/c3x/scaled_mm_sm120_fp8.cu
+9
@@ -1,5 +1,6 @@
1
#include "scaled_mm_kernels.hpp"
2
#include "scaled_mm_sm120_fp8_dispatch.cuh"
3
+#include "core/batch_invariant.hpp"
4
#include "cutlass_extensions/epilogue/scaled_mm_epilogues_c3x.hpp"
5
6
namespace vllm {
@@ -14,9 +15,17 @@ void cutlass_scaled_mm_sm120_fp8(
15
STD_TORCH_CHECK(bias->scalar_type() == out.scalar_type(),
16
"currently bias dtype must match output dtype ",
17
out.scalar_type());
18
+ if (vllm_is_batch_invariant()) {
19
+ return cutlass_scaled_mm_sm120_fp8_batch_invariant_epilogue<
20
+ c3x::ScaledEpilogueBias>(out, a, b, a_scales, b_scales, *bias);
21
+ }
22
return cutlass_scaled_mm_sm120_fp8_epilogue<c3x::ScaledEpilogueBias>(
23
out, a, b, a_scales, b_scales, *bias);
24
} else {
25
+ if (vllm_is_batch_invariant()) {
26
+ return cutlass_scaled_mm_sm120_fp8_batch_invariant_epilogue<
27
+ c3x::ScaledEpilogue>(out, a, b, a_scales, b_scales);
28
+ }
29
return cutlass_scaled_mm_sm120_fp8_epilogue<c3x::ScaledEpilogue>(
30
out, a, b, a_scales, b_scales);
31
}
csrc/libtorch_stable/quantization/w8a8/cutlass/c3x/scaled_mm_sm120_fp8_dispatch.cuh
+42
@@ -179,6 +179,26 @@ inline void cutlass_gemm_sm120_fp8_dispatch(torch::stable::Tensor& out,
179
out, a, b, std::forward<EpilogueArgs>(args)...);
180
}
181
182
+template <typename InType, typename OutType,
183
+ template <typename, typename, typename> typename Epilogue,
184
+ typename... EpilogueArgs>
185
+inline void cutlass_gemm_sm120_fp8_batch_invariant_dispatch(
186
+ torch::stable::Tensor& out, torch::stable::Tensor const& a,
187
+ torch::stable::Tensor const& b, EpilogueArgs&&... args) {
188
+ static_assert(std::is_same<InType, cutlass::float_e4m3_t>());
189
+ STD_TORCH_CHECK(a.scalar_type() ==
190
+ torch::headeronly::ScalarType::Float8_e4m3fn);
191
+ STD_TORCH_CHECK(b.scalar_type() ==
192
+ torch::headeronly::ScalarType::Float8_e4m3fn);
193
+
194
+ using Cutlass3xGemmM64 =
195
+ typename sm120_fp8_config_M64<InType, OutType, Epilogue>::Cutlass3xGemm;
196
+
197
+ // keep the CUTLASS config independent of M for batch invariance
198
+ return cutlass_gemm_caller<Cutlass3xGemmM64>(
199
+ out, a, b, std::forward<EpilogueArgs>(args)...);
200
+}
201
+
202
template <template <typename, typename, typename> typename Epilogue,
203
typename... EpilogueArgs>
204
void cutlass_scaled_mm_sm120_fp8_epilogue(torch::stable::Tensor& out,
@@ -202,4 +222,26 @@ void cutlass_scaled_mm_sm120_fp8_epilogue(torch::stable::Tensor& out,
222
}
223
}
224
225
+template <template <typename, typename, typename> typename Epilogue,
226
+ typename... EpilogueArgs>
227
+void cutlass_scaled_mm_sm120_fp8_batch_invariant_epilogue(
228
+ torch::stable::Tensor& out, torch::stable::Tensor const& a,
229
+ torch::stable::Tensor const& b, EpilogueArgs&&... epilogue_args) {
230
+ STD_TORCH_CHECK(a.scalar_type() ==
231
+ torch::headeronly::ScalarType::Float8_e4m3fn);
232
+ STD_TORCH_CHECK(b.scalar_type() ==
233
+ torch::headeronly::ScalarType::Float8_e4m3fn);
234
+
235
+ if (out.scalar_type() == torch::headeronly::ScalarType::BFloat16) {
236
+ return cutlass_gemm_sm120_fp8_batch_invariant_dispatch<
237
+ cutlass::float_e4m3_t, cutlass::bfloat16_t, Epilogue>(
238
+ out, a, b, std::forward<EpilogueArgs>(epilogue_args)...);
239
+ } else {
240
+ STD_TORCH_CHECK(out.scalar_type() == torch::headeronly::ScalarType::Half);
241
+ return cutlass_gemm_sm120_fp8_batch_invariant_dispatch<
242
+ cutlass::float_e4m3_t, cutlass::half_t, Epilogue>(
243
+ out, a, b, std::forward<EpilogueArgs>(epilogue_args)...);
244
+ }
245
+}
246
+
247
} // namespace vllm
csrc/libtorch_stable/quantization/w8a8/cutlass/c3x/scaled_mm_sm90_fp8.cu
+9
@@ -1,5 +1,6 @@
1
#include "scaled_mm_kernels.hpp"
2
#include "scaled_mm_sm90_fp8_dispatch.cuh"
3
+#include "core/batch_invariant.hpp"
4
5
namespace vllm {
6
@@ -13,9 +14,17 @@ void cutlass_scaled_mm_sm90_fp8(
14
STD_TORCH_CHECK(bias->scalar_type() == out.scalar_type(),
15
"currently bias dtype must match output dtype ",
16
out.scalar_type());
17
+ if (vllm_is_batch_invariant()) {
18
+ return cutlass_scaled_mm_sm90_fp8_batch_invariant_epilogue<true>(
19
+ out, a, b, a_scales, b_scales, *bias);
20
+ }
21
return cutlass_scaled_mm_sm90_fp8_epilogue<true>(out, a, b, a_scales,
22
b_scales, *bias);
23
} else {
24
+ if (vllm_is_batch_invariant()) {
25
+ return cutlass_scaled_mm_sm90_fp8_batch_invariant_epilogue<false>(
26
+ out, a, b, a_scales, b_scales);
27
+ }
28
return cutlass_scaled_mm_sm90_fp8_epilogue<false>(out, a, b, a_scales,
29
b_scales);
30
}
csrc/libtorch_stable/quantization/w8a8/cutlass/c3x/scaled_mm_sm90_fp8_dispatch.cuh
+53
@@ -349,6 +349,35 @@ inline void cutlass_gemm_sm90_fp8_dispatch(
349
}
350
}
351
352
+template <typename InType, typename OutType, bool EnableBias,
353
+ typename... EpilogueArgs>
354
+inline void cutlass_gemm_sm90_fp8_batch_invariant_dispatch(
355
+ torch::stable::Tensor& out, torch::stable::Tensor const& a,
356
+ torch::stable::Tensor const& b, torch::stable::Tensor const& a_scales,
357
+ torch::stable::Tensor const& b_scales, EpilogueArgs&&... args) {
358
+ static_assert(std::is_same<InType, cutlass::float_e4m3_t>());
359
+ STD_TORCH_CHECK(a.scalar_type() ==
360
+ torch::headeronly::ScalarType::Float8_e4m3fn);
361
+ STD_TORCH_CHECK(b.scalar_type() ==
362
+ torch::headeronly::ScalarType::Float8_e4m3fn);
363
+
364
+ using Cutlass3xGemmM64_N1280 =
365
+ typename sm90_fp8_config_M64_N1280<InType, OutType,
366
+ EnableBias>::Cutlass3xGemm;
367
+ using Cutlass3xGemmM64_N8192 =
368
+ typename sm90_fp8_config_M64_N8192<InType, OutType,
369
+ EnableBias>::Cutlass3xGemm;
370
+
371
+ // keep the CUTLASS config independent of M for batch invariance
372
+ uint32_t const n = b.size(1);
373
+ if (n <= 1280) {
374
+ return cutlass_gemm_caller_sm90_fp8<Cutlass3xGemmM64_N1280>(
375
+ out, a, b, b_scales, a_scales, std::forward<EpilogueArgs>(args)...);
376
+ }
377
+ return cutlass_gemm_caller_sm90_fp8<Cutlass3xGemmM64_N8192>(
378
+ out, a, b, b_scales, a_scales, std::forward<EpilogueArgs>(args)...);
379
+}
380
+
381
template <bool EnableBias, typename... EpilogueArgs>
382
void cutlass_scaled_mm_sm90_fp8_epilogue(torch::stable::Tensor& out,
383
torch::stable::Tensor const& a,
@@ -375,4 +404,28 @@ void cutlass_scaled_mm_sm90_fp8_epilogue(torch::stable::Tensor& out,
404
}
405
}
406
407
+template <bool EnableBias, typename... EpilogueArgs>
408
+void cutlass_scaled_mm_sm90_fp8_batch_invariant_epilogue(
409
+ torch::stable::Tensor& out, torch::stable::Tensor const& a,
410
+ torch::stable::Tensor const& b, torch::stable::Tensor const& a_scales,
411
+ torch::stable::Tensor const& b_scales, EpilogueArgs&&... epilogue_args) {
412
+ STD_TORCH_CHECK(a.scalar_type() ==
413
+ torch::headeronly::ScalarType::Float8_e4m3fn);
414
+ STD_TORCH_CHECK(b.scalar_type() ==
415
+ torch::headeronly::ScalarType::Float8_e4m3fn);
416
+
417
+ if (out.scalar_type() == torch::headeronly::ScalarType::BFloat16) {
418
+ return cutlass_gemm_sm90_fp8_batch_invariant_dispatch<
419
+ cutlass::float_e4m3_t, cutlass::bfloat16_t, EnableBias>(
420
+ out, a, b, a_scales, b_scales,
421
+ std::forward<EpilogueArgs>(epilogue_args)...);
422
+ } else {
423
+ STD_TORCH_CHECK(out.scalar_type() == torch::headeronly::ScalarType::Half);
424
+ return cutlass_gemm_sm90_fp8_batch_invariant_dispatch<
425
+ cutlass::float_e4m3_t, cutlass::half_t, EnableBias>(
426
+ out, a, b, a_scales, b_scales,
427
+ std::forward<EpilogueArgs>(epilogue_args)...);
428
+ }
429
+}
430
+
431
} // namespace vllm
csrc/libtorch_stable/quantization/w8a8/cutlass/scaled_mm_c2x.cu
+9
@@ -9,6 +9,7 @@
9
#include "scaled_mm_c2x_sm89_fp8_dispatch.cuh"
10
#include "scaled_mm_c2x_sm89_int8_dispatch.cuh"
11
12
+#include "core/batch_invariant.hpp"
13
#include "libtorch_stable/cutlass_extensions/epilogue/scaled_mm_epilogues_c2x.hpp"
14
15
using namespace vllm;
@@ -191,9 +192,17 @@ void cutlass_scaled_mm_sm89(torch::stable::Tensor& out,
192
STD_TORCH_CHECK(bias->scalar_type() == out.scalar_type(),
193
"currently bias dtype must match output dtype ",
194
out.scalar_type());
195
+ if (vllm_is_batch_invariant()) {
196
+ return cutlass_scaled_mm_sm89_fp8_batch_invariant_epilogue<
197
+ c2x::ScaledEpilogueBias>(out, a, b, a_scales, b_scales, *bias);
198
+ }
199
return cutlass_scaled_mm_sm89_epilogue<c2x::ScaledEpilogueBias>(
200
out, a, b, a_scales, b_scales, *bias);
201
} else {
202
+ if (vllm_is_batch_invariant()) {
203
+ return cutlass_scaled_mm_sm89_fp8_batch_invariant_epilogue<
204
+ c2x::ScaledEpilogue>(out, a, b, a_scales, b_scales);
205
+ }
206
return cutlass_scaled_mm_sm89_epilogue<c2x::ScaledEpilogue>(
207
out, a, b, a_scales, b_scales);
208
}
csrc/libtorch_stable/quantization/w8a8/cutlass/scaled_mm_c2x_sm89_fp8_dispatch.cuh
+39
@@ -381,4 +381,43 @@ inline void cutlass_gemm_sm89_fp8_dispatch(torch::stable::Tensor& out,
381
}
382
}
383
384
+template <typename InType, typename OutType,
385
+ template <typename, typename> typename Epilogue,
386
+ typename... EpilogueArgs>
387
+inline void cutlass_gemm_sm89_fp8_batch_invariant_dispatch(
388
+ torch::stable::Tensor& out, torch::stable::Tensor const& a,
389
+ torch::stable::Tensor const& b, EpilogueArgs&&... args) {
390
+ static_assert(std::is_same<InType, cutlass::float_e4m3_t>());
391
+ STD_TORCH_CHECK(a.scalar_type() ==
392
+ torch::headeronly::ScalarType::Float8_e4m3fn);
393
+ STD_TORCH_CHECK(b.scalar_type() ==
394
+ torch::headeronly::ScalarType::Float8_e4m3fn);
395
+
396
+ // keep the CUTLASS config independent of M for batch invariance
397
+ return sm89_fp8_config_M64::dispatch<InType, OutType, Epilogue>(
398
+ out, a, b, std::forward<EpilogueArgs>(args)...);
399
+}
400
+
401
+template <template <typename, typename> typename Epilogue,
402
+ typename... EpilogueArgs>
403
+void cutlass_scaled_mm_sm89_fp8_batch_invariant_epilogue(
404
+ torch::stable::Tensor& out, torch::stable::Tensor const& a,
405
+ torch::stable::Tensor const& b, EpilogueArgs&&... epilogue_args) {
406
+ STD_TORCH_CHECK(a.scalar_type() ==
407
+ torch::headeronly::ScalarType::Float8_e4m3fn);
408
+ STD_TORCH_CHECK(b.scalar_type() ==
409
+ torch::headeronly::ScalarType::Float8_e4m3fn);
410
+
411
+ if (out.scalar_type() == torch::headeronly::ScalarType::BFloat16) {
412
+ return cutlass_gemm_sm89_fp8_batch_invariant_dispatch<
413
+ cutlass::float_e4m3_t, cutlass::bfloat16_t, Epilogue>(
414
+ out, a, b, std::forward<EpilogueArgs>(epilogue_args)...);
415
+ } else {
416
+ STD_TORCH_CHECK(out.scalar_type() == torch::headeronly::ScalarType::Half);
417
+ return cutlass_gemm_sm89_fp8_batch_invariant_dispatch<
418
+ cutlass::float_e4m3_t, cutlass::half_t, Epilogue>(
419
+ out, a, b, std::forward<EpilogueArgs>(epilogue_args)...);
420
+ }
421
+}
422
+
423
} // namespace vllm
tests/v1/determinism/test_cutlass_batch_invariance.py
new
+68
@@ -0,0 +1,68 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+
4
+import pytest
5
+import torch
6
+
7
+import vllm.envs as envs
8
+from tests.utils import TestFP8Layer, requires_fp8
9
+from vllm.model_executor.kernels.linear.scaled_mm.cutlass import (
10
+ CutlassFP8ScaledMMLinearKernel,
11
+)
12
+from vllm.model_executor.layers.quantization.utils.quant_utils import (
13
+ kFp8DynamicTokenSym,
14
+ kFp8StaticTensorSym,
15
+)
16
+from vllm.platforms import current_platform
17
+
18
+pytest.importorskip("torch.cuda")
19
+
20
+
21
+@pytest.fixture(autouse=True)
22
+def setup_cuda():
23
+ if not current_platform.is_cuda():
24
+ pytest.skip("CUTLASS FP8 kernels require CUDA.")
25
+ torch.set_default_device("cuda")
26
+
27
+
28
+@requires_fp8
29
+@pytest.mark.parametrize("weight_shape", [(1024, 2048), (4608, 4096)])
30
+@pytest.mark.parametrize("batch_size", [1, 16, 17, 32, 64, 65, 256, 257])
31
+@torch.inference_mode()
32
+def test_cutlass_fp8_batch_invariant_fixed_config(
33
+ weight_shape: tuple[int, int],
34
+ batch_size: int,
35
+ default_vllm_config,
36
+ monkeypatch: pytest.MonkeyPatch,
37
+):
38
+ monkeypatch.setenv("VLLM_BATCH_INVARIANT", "1")
39
+ monkeypatch.setattr(envs, "VLLM_BATCH_INVARIANT", True)
40
+
41
+ torch.manual_seed(0)
42
+ layer = TestFP8Layer(
43
+ weight_shape=weight_shape,
44
+ activation_quant_key=kFp8DynamicTokenSym,
45
+ weight_quant_key=kFp8StaticTensorSym,
46
+ input_dtype=torch.bfloat16,
47
+ out_dtype=torch.bfloat16,
48
+ device=torch.device("cuda"),
49
+ force_kernel=CutlassFP8ScaledMMLinearKernel,
50
+ )
51
+ assert isinstance(layer.kernel, CutlassFP8ScaledMMLinearKernel)
52
+
53
+ in_features = weight_shape[1]
54
+ needle = torch.randn((1, in_features), device="cuda", dtype=torch.bfloat16)
55
+ baseline = layer(needle)[0]
56
+
57
+ filler = torch.randn(
58
+ (max(batch_size - 1, 0), in_features), device="cuda", dtype=torch.bfloat16
59
+ )
60
+
61
+ front_batch = torch.cat([needle, filler], dim=0)
62
+ back_batch = torch.cat([filler, needle], dim=0)
63
+
64
+ front_output = layer(front_batch)[0]
65
+ back_output = layer(back_batch)[-1]
66
+
67
+ torch.testing.assert_close(front_output, baseline, rtol=0, atol=0)
68
+ torch.testing.assert_close(back_output, baseline, rtol=0, atol=0)
vllm/model_executor/layers/quantization/fp8.py
+9
-3
@@ -16,7 +16,10 @@ from vllm.logger import init_logger
16
from vllm.model_executor.kernels.linear import (
17
init_fp8_linear_kernel,
18
)
19
-from vllm.model_executor.kernels.linear.scaled_mm import MarlinFP8ScaledMMLinearKernel
19
+from vllm.model_executor.kernels.linear.scaled_mm import (
20
+ CutlassFP8ScaledMMLinearKernel,
21
+ MarlinFP8ScaledMMLinearKernel,
22
+)
23
from vllm.model_executor.layers.attention import Attention
24
from vllm.model_executor.layers.fused_moe import (
25
FusedMoE,
@@ -437,8 +440,8 @@ class Fp8LinearMethod(LinearMethodBase):
440
x: torch.Tensor,
441
bias: torch.Tensor | None = None,
442
) -> torch.Tensor:
440
- # if batch invariant mode is enabled, prefer DeepGEMM FP8 path
441
- # we will use BF16 dequant when DeepGEMM is not supported.
443
+ # if batch invariant mode is enabled, prefer direct FP8 path
444
+ # we will use BF16 dequant when direct FP8 is not supported.
445
if envs.VLLM_BATCH_INVARIANT:
446
if self.block_quant:
447
assert self.weight_block_size is not None
@@ -448,6 +451,9 @@ class Fp8LinearMethod(LinearMethodBase):
451
bias,
452
)
453
else:
454
+ if isinstance(self.fp8_linear, CutlassFP8ScaledMMLinearKernel):
455
+ return self.fp8_linear.apply_weights(layer, x, bias)
456
+
457
# per-tensor/channel: dequant to BF16 and run GEMM
458
weight_fp8 = layer.weight.to(torch.bfloat16)
459
weight_scale = layer.weight_scale.to(torch.bfloat16)
vllm/model_executor/layers/quantization/online/fp8.py
+6
@@ -18,6 +18,9 @@ import vllm.envs as envs
18
from vllm import _custom_ops as ops
19
from vllm.config import get_current_vllm_config
20
from vllm.model_executor.kernels.linear import init_fp8_linear_kernel
21
+from vllm.model_executor.kernels.linear.scaled_mm import (
22
+ CutlassFP8ScaledMMLinearKernel,
23
+)
24
from vllm.model_executor.layers.fused_moe.oracle.fp8 import (
25
select_fp8_moe_backend,
26
)
@@ -161,6 +164,9 @@ class Fp8PerTensorOnlineLinearMethod(_Fp8OnlineLinearBase):
164
) -> torch.Tensor:
165
# if batch invariant mode is enabled, use BF16 dequant
166
if envs.VLLM_BATCH_INVARIANT:
167
+ if isinstance(self.fp8_linear, CutlassFP8ScaledMMLinearKernel):
168
+ return self.fp8_linear.apply_weights(layer, x, bias)
169
+
170
weight_fp8 = layer.weight.to(torch.bfloat16)
171
weight_scale = layer.weight_scale.to(torch.bfloat16)
172
if weight_scale.numel() == 1: