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 5, 16:00 UTC
[Kernel] Add MXFP8 to Marlin GEMM/MoE and refactor Mxfp8LinearOp (#34664)
Signed-off-by: mgoin <mgoin64@gmail.com>
Michael Goin committed
Apr 1, 2026 at 18:41 UTC
db5d0719e1057b876ed673ae6387d940962691bb
15 files changed
+481
-129
csrc/moe/marlin_moe_wna16/generate_kernels.py
+9
@@ -108,6 +108,15 @@ QUANT_CONFIGS = [
108
"thread_m_blocks": THREAD_M_BLOCKS,
109
"group_blocks": [2],
110
},
111
+ # MXFP8
112
+ {
113
+ "a_type": ["kBFloat16"],
114
+ "b_type": "kFE4M3fn",
115
+ "s_type": "kFE8M0fnu",
116
+ "thread_configs": THREAD_CONFIGS,
117
+ "thread_m_blocks": THREAD_M_BLOCKS,
118
+ "group_blocks": [2],
119
+ },
120
# AWQ-INT4 with INT8 activation
121
{
122
"a_type": ["kS8"],
csrc/moe/marlin_moe_wna16/marlin_template.h
+10
-8
@@ -343,6 +343,8 @@ __global__ void Marlin(
343
if constexpr (b_type == vllm::kFE2M1f) {
344
static_assert(s_type == vllm::kFE4M3fn && group_blocks == 1 ||
345
s_type == vllm::kFE8M0fnu && group_blocks == 2);
346
+ } else if constexpr (b_type == vllm::kFE4M3fn && s_type == vllm::kFE8M0fnu) {
347
+ static_assert(group_blocks == 2);
348
} else if constexpr (std::is_same<scalar_t, nv_bfloat16>::value) {
349
static_assert(s_type == vllm::kBFloat16);
350
} else if constexpr (std::is_same<scalar_t, half>::value) {
@@ -357,9 +359,10 @@ __global__ void Marlin(
359
constexpr bool is_int_type = b_type == vllm::kU4 || b_type == vllm::kU8 ||
360
b_type == vllm::kS4 || b_type == vllm::kS8 ||
361
b_type == vllm::kU4B8 || b_type == vllm::kU8B128;
362
+ constexpr bool is_8bit_scale = s_type.size_bits() == 8;
363
// see comments of dequant.h for more details
364
constexpr bool dequant_skip_flop =
362
- is_a_8bit || b_type == vllm::kFE4M3fn ||
365
+ is_a_8bit || (b_type == vllm::kFE4M3fn && !(s_type == vllm::kFE8M0fnu)) ||
366
b_type == vllm::kFE2M1f && s_type == vllm::kFE4M3fn ||
367
has_zp && !is_zp_float && !std::is_same<scalar_t, nv_bfloat16>::value ||
368
has_zp && !is_zp_float && !(b_type == vllm::kU8);
@@ -373,7 +376,7 @@ __global__ void Marlin(
376
const int group_size =
377
(!has_act_order && group_blocks == -1) ? prob_k : prob_k / num_groups;
378
const int scales_expert_stride =
376
- prob_n * prob_k / group_size / (b_type == vllm::kFE2M1f ? 16 : 8);
379
+ prob_n * prob_k / group_size / (is_8bit_scale ? 16 : 8);
380
const int zp_expert_stride =
381
is_zp_float ? prob_n * prob_k / group_size / 8
382
: prob_n * prob_k / group_size / (pack_factor * 4);
@@ -692,9 +695,8 @@ __global__ void Marlin(
695
constexpr int b_sh_wr_iters = b_sh_stage / b_sh_wr_delta;
696
697
// Scale sizes/strides without act_order
695
- int s_gl_stride = prob_n / (b_type == vllm::kFE2M1f ? 16 : 8);
696
- constexpr int s_sh_stride =
697
- 16 * thread_n_blocks / (b_type == vllm::kFE2M1f ? 16 : 8);
698
+ int s_gl_stride = prob_n / (is_8bit_scale ? 16 : 8);
699
+ constexpr int s_sh_stride = 16 * thread_n_blocks / (is_8bit_scale ? 16 : 8);
700
constexpr int s_tb_groups =
701
!has_act_order && group_blocks != -1 && group_blocks < thread_k_blocks
702
? thread_k_blocks / group_blocks
@@ -1131,7 +1133,7 @@ __global__ void Marlin(
1133
1134
int4* sh_s_stage = sh_s + s_sh_stage * pipe;
1135
1134
- if constexpr (b_type_id != vllm::kFE2M1f.id()) {
1136
+ if constexpr (!is_8bit_scale) {
1137
reinterpret_cast<int4*>(&frag_s[k % 2])[0] =
1138
sh_s_stage[s_sh_rd + cur_group_id * s_sh_stride];
1139
} else {
@@ -1140,7 +1142,7 @@ __global__ void Marlin(
1142
sh_s_stage)[s_sh_rd + cur_group_id * (2 * s_sh_stride)];
1143
}
1144
} else if (group_blocks >= b_sh_wr_iters) {
1143
- if constexpr (b_type_id != vllm::kFE2M1f.id()) {
1145
+ if constexpr (!is_8bit_scale) {
1146
reinterpret_cast<int4*>(&frag_s[1])[0] =
1147
reinterpret_cast<int4*>(&frag_s[0])[0];
1148
} else {
@@ -1341,7 +1343,7 @@ __global__ void Marlin(
1343
}
1344
}
1345
1344
- if constexpr (b_type == vllm::kFE2M1f) {
1346
+ if constexpr (s_type == vllm::kFE4M3fn || s_type == vllm::kFE8M0fnu) {
1347
int s_quant_0 = reinterpret_cast<int*>(frag_s[k2])[0];
1348
int s_quant_1 = reinterpret_cast<int*>(frag_s[k2])[1];
1349
csrc/moe/marlin_moe_wna16/ops.cu
+3
@@ -599,6 +599,9 @@ torch::Tensor moe_wna16_marlin_gemm(
599
"When b_type = float4_e2m1f, b_scale scalar type must be",
600
"float8_e4m3fn (for NVFP4) or float8_e8m0fnu (for MXFP4).");
601
}
602
+ } else if (b_type_id == vllm::kFE4M3fn.id() &&
603
+ b_scales.scalar_type() == at::ScalarType::Float8_e8m0fnu) {
604
+ s_type_id = vllm::kFE8M0fnu.id();
605
}
606
607
vllm::ScalarType a_type = vllm::ScalarType::from_id(a_type_id);
csrc/quantization/marlin/generate_kernels.py
+9
@@ -108,6 +108,15 @@ QUANT_CONFIGS = [
108
"thread_m_blocks": THREAD_M_BLOCKS,
109
"group_blocks": [2],
110
},
111
+ # MXFP8
112
+ {
113
+ "a_type": ["kBFloat16"],
114
+ "b_type": "kFE4M3fn",
115
+ "s_type": "kFE8M0fnu",
116
+ "thread_configs": THREAD_CONFIGS,
117
+ "thread_m_blocks": THREAD_M_BLOCKS,
118
+ "group_blocks": [2],
119
+ },
120
# AWQ-INT4 with INT8 activation
121
{
122
"a_type": ["kS8"],
csrc/quantization/marlin/marlin.cu
+3
@@ -591,6 +591,9 @@ torch::Tensor marlin_gemm(
591
"When b_type = float4_e2m1f, b_scale scalar type must be",
592
"float8_e4m3fn (for NVFP4) or float8_e8m0fnu (for MXFP4).");
593
}
594
+ } else if (b_type_id == vllm::kFE4M3fn.id() &&
595
+ b_scales.scalar_type() == at::ScalarType::Float8_e8m0fnu) {
596
+ s_type_id = vllm::kFE8M0fnu.id();
597
}
598
599
vllm::ScalarType a_type = vllm::ScalarType::from_id(a_type_id);
csrc/quantization/marlin/marlin_template.h
+10
-7
@@ -327,6 +327,9 @@ __global__ void Marlin(
327
if constexpr (b_type == vllm::kFE2M1f) {
328
static_assert(s_type == vllm::kFE4M3fn && group_blocks == 1 ||
329
s_type == vllm::kFE8M0fnu && group_blocks == 2);
330
+ } else if constexpr (s_type == vllm::kFE8M0fnu) {
331
+ // MXFP8: FP8 weights with e8m0 microscaling block scales
332
+ static_assert(b_type == vllm::kFE4M3fn && group_blocks == 2);
333
} else if constexpr (std::is_same<scalar_t, nv_bfloat16>::value) {
334
static_assert(s_type == vllm::kBFloat16);
335
} else if constexpr (std::is_same<scalar_t, half>::value) {
@@ -334,6 +337,7 @@ __global__ void Marlin(
337
}
338
339
constexpr bool is_a_8bit = a_type.size_bits() == 8;
340
+ constexpr bool is_8bit_scale = s_type.size_bits() == 8;
341
if constexpr (!is_a_8bit) {
342
static_assert(std::is_same<scalar_t, c_scalar_t>::value);
343
}
@@ -343,7 +347,7 @@ __global__ void Marlin(
347
b_type == vllm::kU4B8 || b_type == vllm::kU8B128;
348
// see comments of dequant.h for more details
349
constexpr bool dequant_skip_flop =
346
- is_a_8bit || b_type == vllm::kFE4M3fn ||
350
+ is_a_8bit || (b_type == vllm::kFE4M3fn && !(s_type == vllm::kFE8M0fnu)) ||
351
b_type == vllm::kFE2M1f && s_type == vllm::kFE4M3fn ||
352
has_zp && !is_zp_float && !std::is_same<scalar_t, nv_bfloat16>::value ||
353
has_zp && !is_zp_float && !(b_type == vllm::kU8);
@@ -555,9 +559,8 @@ __global__ void Marlin(
559
constexpr int b_sh_wr_iters = b_sh_stage / b_sh_wr_delta;
560
561
// Scale sizes/strides without act_order
558
- int s_gl_stride = prob_n / (b_type == vllm::kFE2M1f ? 16 : 8);
559
- constexpr int s_sh_stride =
560
- 16 * thread_n_blocks / (b_type == vllm::kFE2M1f ? 16 : 8);
562
+ int s_gl_stride = prob_n / (is_8bit_scale ? 16 : 8);
563
+ constexpr int s_sh_stride = 16 * thread_n_blocks / (is_8bit_scale ? 16 : 8);
564
constexpr int s_tb_groups =
565
!has_act_order && group_blocks != -1 && group_blocks < thread_k_blocks
566
? thread_k_blocks / group_blocks
@@ -997,7 +1000,7 @@ __global__ void Marlin(
1000
1001
int4* sh_s_stage = sh_s + s_sh_stage * pipe;
1002
1000
- if constexpr (b_type_id != vllm::kFE2M1f.id()) {
1003
+ if constexpr (!is_8bit_scale) {
1004
reinterpret_cast<int4*>(&frag_s[k % 2])[0] =
1005
sh_s_stage[s_sh_rd + cur_group_id * s_sh_stride];
1006
} else {
@@ -1006,7 +1009,7 @@ __global__ void Marlin(
1009
sh_s_stage)[s_sh_rd + cur_group_id * (2 * s_sh_stride)];
1010
}
1011
} else if (group_blocks >= b_sh_wr_iters) {
1009
- if constexpr (b_type_id != vllm::kFE2M1f.id()) {
1012
+ if constexpr (!is_8bit_scale) {
1013
reinterpret_cast<int4*>(&frag_s[1])[0] =
1014
reinterpret_cast<int4*>(&frag_s[0])[0];
1015
} else {
@@ -1207,7 +1210,7 @@ __global__ void Marlin(
1210
}
1211
}
1212
1210
- if constexpr (b_type == vllm::kFE2M1f) {
1213
+ if constexpr (s_type == vllm::kFE4M3fn || s_type == vllm::kFE8M0fnu) {
1214
int s_quant_0 = reinterpret_cast<int*>(frag_s[k2])[0];
1215
int s_quant_1 = reinterpret_cast<int*>(frag_s[k2])[1];
1216
tests/kernels/moe/test_moe.py
+6
@@ -151,6 +151,12 @@ MOE_MARLIN_QUANT_TEST_CONFIGS = [
151
"b_type": scalar_types.float4_e2m1f,
152
"group_blocks": [2],
153
},
154
+ # MXFP8
155
+ {
156
+ "a_type": [scalar_types.bfloat16],
157
+ "b_type": scalar_types.float8_e4m3fn,
158
+ "group_blocks": [2],
159
+ },
160
# AWQ-INT4 with INT8 activation
161
{
162
"a_type": [scalar_types.int8],
tests/models/quantization/test_mxfp8.py
+1
-1
@@ -23,7 +23,7 @@ from tests.quantization.utils import is_quant_method_supported
23
from ..utils import check_logprobs_close
24
25
# A small MoE model that fits on a single GPU and has both linear + MoE layers.
26
-MOE_MODEL = "Qwen/Qwen3-30B-A3B"
26
+MOE_MODEL = "allenai/OLMoE-1B-7B-0125-Instruct"
27
# A small dense model (no MoE) to validate the linear-only path.
28
DENSE_MODEL = "Qwen/Qwen3-0.6B"
29
vllm/model_executor/layers/fused_moe/fused_marlin_moe.py
+2
-1
@@ -41,6 +41,7 @@ from vllm.model_executor.layers.quantization.utils.quant_utils import (
41
kFp8StaticChannelSym,
42
kFp8StaticTensorSym,
43
kMxfp4Static,
44
+ kMxfp8Static,
45
kNvfp4Static,
46
)
47
from vllm.platforms import current_platform
@@ -582,6 +583,7 @@ class MarlinExpertsBase(mk.FusedMoEExpertsModular):
583
kFp8StaticChannelSym,
584
kFp8StaticTensorSym,
585
kMxfp4Static,
586
+ kMxfp8Static,
587
kNvfp4Static,
588
]
589
return weight_key in SUPPORTED_W
@@ -609,7 +611,6 @@ class MarlinExpertsBase(mk.FusedMoEExpertsModular):
611
612
@property
613
def quant_type_id(self) -> int:
612
- # uint4b8 will be set for int4 weight and float4_e2m1f will be used for mxfp4
614
if self.quant_config.use_int4_w4a16:
615
return scalar_types.uint4b8.id
616
elif self.quant_config.use_mxfp4_w4a16 or self.quant_config.use_nvfp4_w4a16:
vllm/model_executor/layers/fused_moe/oracle/fp8.py
+21
-7
@@ -436,13 +436,27 @@ def convert_to_fp8_moe_kernel_format(
436
elif fp8_backend == Fp8MoeBackend.AITER:
437
w13, w2 = rocm_aiter_ops.shuffle_weights(w13, w2)
438
elif fp8_backend == Fp8MoeBackend.MARLIN:
439
- w13, w2, w13_scale, w2_scale = prepare_fp8_moe_layer_for_marlin(
440
- layer,
441
- w13,
442
- w2,
443
- w13_scale,
444
- w2_scale,
445
- )
439
+ weight_block_size = getattr(layer, "weight_block_size", None)
440
+ if weight_block_size == [1, 32]:
441
+ from vllm.model_executor.layers.quantization.utils.marlin_utils_fp8 import (
442
+ prepare_mxfp8_moe_layer_for_marlin,
443
+ )
444
+
445
+ w13, w2, w13_scale, w2_scale = prepare_mxfp8_moe_layer_for_marlin(
446
+ layer,
447
+ w13,
448
+ w2,
449
+ w13_scale,
450
+ w2_scale,
451
+ )
452
+ else:
453
+ w13, w2, w13_scale, w2_scale = prepare_fp8_moe_layer_for_marlin(
454
+ layer,
455
+ w13,
456
+ w2,
457
+ w13_scale,
458
+ w2_scale,
459
+ )
460
elif fp8_backend in [
461
Fp8MoeBackend.FLASHINFER_CUTLASS,
462
Fp8MoeBackend.FLASHINFER_TRTLLM,
vllm/model_executor/layers/fused_moe/oracle/mxfp8.py
+9
-5
@@ -15,14 +15,14 @@ from vllm.model_executor.layers.quantization.utils.quant_utils import (
15
16
logger = init_logger(__name__)
17
18
-_SUPPORTED_BACKENDS: frozenset[Fp8MoeBackend] = frozenset(
19
- {
20
- Fp8MoeBackend.FLASHINFER_TRTLLM,
21
- }
18
+_SUPPORTED_BACKENDS = (
19
+ Fp8MoeBackend.FLASHINFER_TRTLLM,
20
+ Fp8MoeBackend.MARLIN,
21
)
22
23
_BACKEND_NAME_MAP: dict[str, Fp8MoeBackend] = {
24
"flashinfer_trtllm": Fp8MoeBackend.FLASHINFER_TRTLLM,
25
+ "marlin": Fp8MoeBackend.MARLIN,
26
}
27
28
@@ -81,7 +81,11 @@ def select_mxfp8_moe_backend(
81
82
# Auto-select: pick the first supported backend.
83
for backend in _SUPPORTED_BACKENDS:
84
+ try:
85
+ experts_cls = _select_kernel_cls(backend, config)
86
+ except ValueError:
87
+ continue
88
logger.info_once("Using '%s' MxFp8 MoE backend.", backend.value)
85
- return backend, _select_kernel_cls(backend, config)
89
+ return backend, experts_cls
90
91
raise ValueError("No MXFP8 MoE backends available.")
vllm/model_executor/layers/quantization/modelopt.py
+7
-55
@@ -67,10 +67,8 @@ from vllm.model_executor.layers.quantization.utils.mxfp8_utils import (
67
MXFP8_BLOCK_SIZE,
68
MXFP8_SCALE_DTYPE,
69
MXFP8_VALUE_DTYPE,
70
- Mxfp8LinearBackend,
70
Mxfp8LinearOp,
71
mxfp8_e4m3_quantize,
73
- swizzle_mxfp8_scale,
72
)
73
from vllm.model_executor.layers.quantization.utils.nvfp4_utils import (
74
apply_nvfp4_linear,
@@ -1499,8 +1497,8 @@ class ModelOptMxFp8Config(ModelOptQuantConfigBase):
1497
1498
@classmethod
1499
def get_min_capability(cls) -> int:
1502
- # MXFP8 hardware acceleration requires Blackwell (SM100) or newer
1503
- return 100
1500
+ # Marlin kernel supports MXFP8 on SM80+
1501
+ return 80
1502
1503
@classmethod
1504
def override_quantization_method(
@@ -1555,9 +1553,7 @@ class ModelOptMxFp8LinearMethod(LinearMethodBase):
1553
"Dynamic quantization is not supported."
1554
)
1555
1558
- self.backend: Mxfp8LinearBackend = Mxfp8LinearBackend.FLASHINFER_CUTLASS
1559
- self.mxfp8_linear_op = Mxfp8LinearOp(backend=self.backend)
1560
- logger.info_once("Using %s backend for MXFP8 GEMM", self.backend.value)
1556
+ self.mxfp8_linear_op = Mxfp8LinearOp()
1557
1558
def create_weights(
1559
self,
@@ -1615,36 +1611,6 @@ class ModelOptMxFp8LinearMethod(LinearMethodBase):
1611
)
1612
layer.register_parameter("weight_scale", weight_scale)
1613
1618
- def _process_weights_after_loading_scale_2d(self, layer: torch.nn.Module) -> None:
1619
- """Not swizzled - MXFP8 GEMM emulation"""
1620
- weight = layer.weight.data # [N, K]
1621
- N, K = weight.shape
1622
- scale_k = K // MXFP8_BLOCK_SIZE
1623
-
1624
- # Slice weight_scale to match weight dimensions (handles padding)
1625
- weight_scale = layer.weight_scale.data[:N, :scale_k].contiguous()
1626
-
1627
- layer.weight = Parameter(weight.contiguous(), requires_grad=False)
1628
- layer.weight_scale = Parameter(weight_scale, requires_grad=False)
1629
-
1630
- def _process_weights_after_loading_scale_1d(self, layer: torch.nn.Module) -> None:
1631
- """Swizzled - MXFP8 GEMM Flashinfer CUTLASS"""
1632
- weight = layer.weight.data # [N, K]
1633
- N, K = weight.shape
1634
-
1635
- # 2D weight scale
1636
- weight_scale = layer.weight_scale.data
1637
-
1638
- # Swizzle the weight scales
1639
- scale_k = K // MXFP8_BLOCK_SIZE
1640
- weight_scale_2d = weight_scale[:N, :scale_k].contiguous()
1641
- weight_scale_swizzled = swizzle_mxfp8_scale(weight_scale_2d, M=N, K=K)
1642
-
1643
- layer.weight = Parameter(weight.contiguous(), requires_grad=False)
1644
- layer.weight_scale = Parameter(
1645
- weight_scale_swizzled.contiguous(), requires_grad=False
1646
- )
1647
-
1614
def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
1615
# Validate weight tensor
1616
if layer.weight.ndim != 2:
@@ -1669,14 +1635,7 @@ class ModelOptMxFp8LinearMethod(LinearMethodBase):
1635
f" got {layer.weight_scale.dtype}"
1636
)
1637
1672
- if self.backend == Mxfp8LinearBackend.EMULATION:
1673
- # Swizzled layout is not used
1674
- self._process_weights_after_loading_scale_2d(layer)
1675
- return
1676
-
1677
- assert self.backend == Mxfp8LinearBackend.FLASHINFER_CUTLASS
1678
- # Swizzled layout is required for Flashinfer CUTLASS
1679
- self._process_weights_after_loading_scale_1d(layer)
1638
+ self.mxfp8_linear_op.process_weights(layer)
1639
1640
def apply(
1641
self,
@@ -1684,22 +1643,15 @@ class ModelOptMxFp8LinearMethod(LinearMethodBase):
1643
x: torch.Tensor,
1644
bias: torch.Tensor | None = None,
1645
) -> torch.Tensor:
1687
- if layer.weight.dtype != MXFP8_VALUE_DTYPE:
1688
- raise ValueError(
1689
- f"Weight dtype {layer.weight.dtype} != expected {MXFP8_VALUE_DTYPE}"
1690
- )
1691
- if layer.weight_scale.dtype != MXFP8_SCALE_DTYPE:
1692
- raise ValueError(
1693
- f"Weight scale dtype {layer.weight_scale.dtype} != "
1694
- f"expected {MXFP8_SCALE_DTYPE}"
1695
- )
1696
-
1646
return self.mxfp8_linear_op.apply(
1647
input=x,
1648
weight=layer.weight,
1649
weight_scale=layer.weight_scale,
1650
out_dtype=x.dtype,
1651
bias=bias,
1652
+ workspace=getattr(layer, "workspace", None),
1653
+ size_n=layer.output_size_per_partition,
1654
+ size_k=layer.input_size_per_partition,
1655
)
1656
1657
vllm/model_executor/layers/quantization/mxfp8.py
+23
-33
@@ -34,10 +34,8 @@ from vllm.model_executor.layers.quantization.fp8 import (
34
)
35
from vllm.model_executor.layers.quantization.utils.mxfp8_utils import (
36
MXFP8_BLOCK_SIZE,
37
- Mxfp8LinearBackend,
37
Mxfp8LinearOp,
38
mxfp8_e4m3_quantize,
40
- swizzle_mxfp8_scale,
39
)
40
from vllm.model_executor.layers.quantization.utils.quant_utils import (
41
is_layer_skipped,
@@ -71,7 +69,8 @@ class Mxfp8Config(Fp8Config):
69
70
@classmethod
71
def get_min_capability(cls) -> int:
74
- return 100
72
+ # Marlin kernel supports MXFP8 on SM80+
73
+ return 80
74
75
@classmethod
76
def from_config(cls, config: dict[str, Any]) -> "Mxfp8Config":
@@ -128,24 +127,7 @@ class Mxfp8OnlineLinearMethod(Fp8OnlineLinearMethod):
127
def __init__(self, quant_config: "Mxfp8Config"):
128
self.quant_config = quant_config
129
self.out_dtype = torch.get_default_dtype()
131
- self.mxfp8_linear = Mxfp8LinearOp(self._select_backend())
132
- logger.info_once(
133
- "Using %s backend for MXFP8 GEMM", self.mxfp8_linear.backend.value
134
- )
135
-
136
- @staticmethod
137
- def _select_backend() -> Mxfp8LinearBackend:
138
- try:
139
- from vllm.utils import flashinfer as fi
140
-
141
- _ = fi.mm_mxfp8
142
- return Mxfp8LinearBackend.FLASHINFER_CUTLASS
143
- except Exception:
144
- logger.warning(
145
- "FlashInfer mm_mxfp8 not available, "
146
- "falling back to MXFP8 emulation backend."
147
- )
148
- return Mxfp8LinearBackend.EMULATION
130
+ self.mxfp8_linear = Mxfp8LinearOp()
131
132
def create_weights(
133
self,
@@ -180,14 +162,12 @@ class Mxfp8OnlineLinearMethod(Fp8OnlineLinearMethod):
162
163
weight_fp8, weight_scale = mxfp8_e4m3_quantize(layer.weight.contiguous())
164
183
- if self.mxfp8_linear.backend == Mxfp8LinearBackend.FLASHINFER_CUTLASS:
184
- N, K = layer.weight.shape[0], layer.weight.shape[1]
185
- weight_scale = swizzle_mxfp8_scale(weight_scale, N, K)
186
-
165
layer.input_scale = None
166
replace_parameter(layer, "weight", weight_fp8.data)
167
replace_parameter(layer, "weight_scale", weight_scale.data)
168
169
+ self.mxfp8_linear.process_weights(layer)
170
+
171
layer._already_called_process_weights_after_loading = True
172
173
def apply(
@@ -202,6 +182,9 @@ class Mxfp8OnlineLinearMethod(Fp8OnlineLinearMethod):
182
weight_scale=layer.weight_scale,
183
out_dtype=self.out_dtype,
184
bias=bias,
185
+ workspace=getattr(layer, "workspace", None),
186
+ size_n=layer.output_size_per_partition,
187
+ size_k=layer.input_size_per_partition,
188
)
189
190
@@ -255,17 +238,24 @@ class Mxfp8OnlineMoEMethod(Fp8OnlineMoEMethod):
238
self, weight: torch.Tensor
239
) -> tuple[torch.Tensor, torch.Tensor]:
240
"""Batch quantization: bf16/fp16 weights -> MXFP8 (fp8 + uint8 scales)."""
258
- num_batches = weight.size(0)
259
- w_quant = []
260
- w_scales = []
261
- for i in range(num_batches):
262
- mx_fp8_quant, mx_fp8_scale = mxfp8_e4m3_quantize(
241
+ E = weight.size(0)
242
+ first_q, first_s = mxfp8_e4m3_quantize(weight[0], is_sf_swizzled_layout=False)
243
+ # Pre-allocate the output tensors rather than stacking.
244
+ # This is important for consistent memory layout.
245
+ w_quant = torch.empty(
246
+ (E, *first_q.shape), dtype=first_q.dtype, device=weight.device
247
+ )
248
+ w_scales = torch.empty(
249
+ (E, *first_s.shape), dtype=first_s.dtype, device=weight.device
250
+ )
251
+ w_quant[0] = first_q
252
+ w_scales[0] = first_s
253
+ for i in range(1, E):
254
+ w_quant[i], w_scales[i] = mxfp8_e4m3_quantize(
255
weight[i], is_sf_swizzled_layout=False
256
)
265
- w_quant.append(mx_fp8_quant)
266
- w_scales.append(mx_fp8_scale)
257
268
- return torch.stack(w_quant), torch.stack(w_scales)
258
+ return w_quant, w_scales
259
260
def process_weights_after_loading(self, layer: Module) -> None:
261
if getattr(layer, "_already_called_process_weights_after_loading", False):
vllm/model_executor/layers/quantization/utils/marlin_utils_fp8.py
+196
@@ -336,6 +336,202 @@ def pack_fp8_to_int32(
336
return int32_tensor.T.contiguous() if size_k_first else int32_tensor
337
338
339
+def mxfp8_marlin_process_scales(marlin_scales: torch.Tensor) -> torch.Tensor:
340
+ """Reorder scales for e8m0 kernel layout and convert to float8_e8m0fnu."""
341
+ # fit the layout of fp8 dequantization
342
+ marlin_scales = marlin_scales.view(-1, 4)[:, [0, 2, 1, 3]].view(
343
+ marlin_scales.size(0), -1
344
+ )
345
+ marlin_scales = marlin_scales.to(torch.float8_e8m0fnu)
346
+ return marlin_scales
347
+
348
+
349
+def apply_mxfp8_marlin_linear(
350
+ input: torch.Tensor,
351
+ weight: torch.Tensor,
352
+ weight_scale: torch.Tensor,
353
+ workspace: torch.Tensor,
354
+ size_n: int,
355
+ size_k: int,
356
+ bias: torch.Tensor | None = None,
357
+ use_fp32_reduce: bool = USE_FP32_REDUCE_DEFAULT,
358
+) -> torch.Tensor:
359
+ reshaped_x = input.reshape(-1, input.shape[-1])
360
+ out_shape = input.shape[:-1] + (size_n,)
361
+
362
+ use_atomic_add = should_use_atomic_add_reduce(
363
+ m=reshaped_x.size(0),
364
+ n=size_n,
365
+ k=size_k,
366
+ device=input.device,
367
+ dtype=input.dtype,
368
+ )
369
+
370
+ output = ops.marlin_gemm(
371
+ a=reshaped_x,
372
+ c=None,
373
+ b_q_weight=weight,
374
+ b_bias=bias,
375
+ b_scales=weight_scale,
376
+ a_scales=None,
377
+ global_scale=None,
378
+ b_zeros=None,
379
+ g_idx=None,
380
+ perm=None,
381
+ workspace=workspace,
382
+ b_q_type=scalar_types.float8_e4m3fn,
383
+ size_m=reshaped_x.size(0),
384
+ size_n=size_n,
385
+ size_k=size_k,
386
+ use_atomic_add=use_atomic_add,
387
+ use_fp32_reduce=use_fp32_reduce,
388
+ )
389
+
390
+ return output.reshape(out_shape)
391
+
392
+
393
+def prepare_mxfp8_layer_for_marlin(layer: torch.nn.Module) -> None:
394
+ """Repack MXFP8 weights and scales into Marlin kernel format.
395
+
396
+ Expects the layer to have:
397
+ - weight: [N, K] float8_e4m3fn
398
+ - weight_scale: [N, K//32] uint8 (e8m0 encoded)
399
+ - input_size_per_partition / output_size_per_partition
400
+ """
401
+ part_size_n = layer.output_size_per_partition
402
+ part_size_k = layer.input_size_per_partition
403
+ group_size = 32 # MX standard block size
404
+
405
+ device = layer.weight.device
406
+
407
+ # WORKSPACE
408
+ layer.workspace = marlin_make_workspace_new(device)
409
+
410
+ # WEIGHT - repack FP8 weights to Marlin format
411
+ perm = torch.empty(0, dtype=torch.int, device=device)
412
+ qweight = pack_fp8_to_int32(layer.weight, size_k_first=False)
413
+ qweight = qweight.T.contiguous()
414
+
415
+ marlin_qweight = ops.gptq_marlin_repack(
416
+ b_q_weight=qweight,
417
+ perm=perm,
418
+ size_k=part_size_k,
419
+ size_n=part_size_n,
420
+ num_bits=8,
421
+ )
422
+ replace_parameter(layer, "weight", marlin_qweight)
423
+
424
+ # WEIGHT SCALES
425
+ # Convert uint8 scales -> e8m0fnu -> param_dtype for permutation
426
+ # Scales are [N, K//32], need [K//32, N] for marlin_permute_scales
427
+ param_dtype = torch.get_default_dtype()
428
+ scales = layer.weight_scale.data[:part_size_n, : part_size_k // group_size]
429
+ scales = scales.contiguous()
430
+ scales = scales.view(torch.float8_e8m0fnu).to(param_dtype)
431
+ scales = scales.T.contiguous()
432
+
433
+ # Permute scales to Marlin layout
434
+ marlin_scales = marlin_permute_scales(
435
+ s=scales,
436
+ size_k=part_size_k,
437
+ size_n=part_size_n,
438
+ group_size=group_size,
439
+ )
440
+
441
+ # Reorder for e8m0 kernel layout and convert back to e8m0fnu
442
+ marlin_scales = mxfp8_marlin_process_scales(marlin_scales)
443
+ replace_parameter(layer, "weight_scale", marlin_scales)
444
+
445
+ # BIAS
446
+ if hasattr(layer, "bias") and layer.bias is not None:
447
+ assert layer.bias.shape == (part_size_n,)
448
+ bias = marlin_permute_bias(layer.bias)
449
+ replace_parameter(layer, "bias", bias)
450
+
451
+
452
+def prepare_mxfp8_moe_layer_for_marlin(
453
+ layer: torch.nn.Module,
454
+ w13: torch.Tensor,
455
+ w2: torch.Tensor,
456
+ w13_scale: torch.Tensor,
457
+ w2_scale: torch.Tensor,
458
+) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
459
+ """Repack MXFP8 MoE weights and scales into Marlin kernel format.
460
+
461
+ Args:
462
+ layer: MoE layer (used to read params_dtype and attach workspace).
463
+ w13: [E, 2*N, K] float8_e4m3fn weights.
464
+ w2: [E, K, N] float8_e4m3fn weights.
465
+ w13_scale: [E, 2*N, K//32] uint8 e8m0 scales.
466
+ w2_scale: [E, K, N//32] uint8 e8m0 scales.
467
+
468
+ Returns:
469
+ (w13, w2, w13_scale, w2_scale) in Marlin format.
470
+ """
471
+ group_size = 32
472
+ e = w13.shape[0]
473
+ w13_n = w13.shape[1]
474
+ k = w13.shape[2]
475
+ n = w2.shape[2]
476
+
477
+ device = w13.device
478
+ param_dtype = torch.get_default_dtype()
479
+ perm = torch.empty(0, dtype=torch.int, device=device)
480
+
481
+ layer.workspace = marlin_make_workspace_new(device, 4)
482
+
483
+ def repack_weight(weight: torch.Tensor, name: str) -> torch.Tensor:
484
+ if "w13" in name:
485
+ size_n, size_k = w13_n, k
486
+ else:
487
+ size_n, size_k = k, n
488
+
489
+ assert weight.shape == (e, size_n, size_k)
490
+
491
+ tensor_list = []
492
+ for i in range(e):
493
+ qweight = pack_fp8_to_int32(weight[i], size_k_first=False)
494
+ qweight = qweight.T.contiguous()
495
+ marlin_qweight = ops.gptq_marlin_repack(
496
+ b_q_weight=qweight,
497
+ perm=perm,
498
+ size_k=size_k,
499
+ size_n=size_n,
500
+ num_bits=8,
501
+ )
502
+ tensor_list.append(marlin_qweight)
503
+ return torch.cat([x.unsqueeze(0) for x in tensor_list], 0)
504
+
505
+ w13 = repack_weight(w13, "w13")
506
+ w2 = repack_weight(w2, "w2")
507
+
508
+ def permute_scales(scales: torch.Tensor, name: str) -> torch.Tensor:
509
+ if "w13" in name:
510
+ size_n, size_k = w13_n, k
511
+ else:
512
+ size_n, size_k = k, n
513
+
514
+ tensor_list = []
515
+ for i in range(e):
516
+ s = scales[i][:size_n, : size_k // group_size].contiguous()
517
+ s = s.view(torch.float8_e8m0fnu).to(param_dtype)
518
+ s = s.T.contiguous()
519
+ marlin_s = marlin_permute_scales(
520
+ s=s,
521
+ size_k=size_k,
522
+ size_n=size_n,
523
+ group_size=group_size,
524
+ )
525
+ marlin_s = mxfp8_marlin_process_scales(marlin_s)
526
+ tensor_list.append(marlin_s)
527
+ return torch.cat([x.unsqueeze(0) for x in tensor_list], 0)
528
+
529
+ w13_scale = permute_scales(w13_scale, "w13")
530
+ w2_scale = permute_scales(w2_scale, "w2")
531
+
532
+ return w13, w2, w13_scale, w2_scale
533
+
534
+
535
def marlin_quant_fp8_torch(weight, group_size, input_dtype=None):
536
is_a_8bit = input_dtype is not None and input_dtype.itemsize == 1
537
if is_a_8bit:
vllm/model_executor/layers/quantization/utils/mxfp8_utils.py
+172
-12
@@ -4,6 +4,7 @@
4
from enum import Enum
5
6
import torch
7
+from torch.nn.parameter import Parameter
8
9
from vllm.logger import init_logger
10
from vllm.utils import flashinfer as vllm_flashinfer
@@ -15,6 +16,7 @@ logger = init_logger(__name__)
16
class Mxfp8LinearBackend(Enum):
17
EMULATION = "emulation"
18
FLASHINFER_CUTLASS = "flashinfer-cutlass"
19
+ MARLIN = "marlin"
20
21
22
# MXFP8 constants
@@ -23,6 +25,28 @@ MXFP8_SCALE_DTYPE = torch.uint8
25
MXFP8_BLOCK_SIZE = 32
26
27
28
+def select_mxfp8_linear_backend() -> Mxfp8LinearBackend:
29
+ """Select the best MXFP8 linear backend for the current device.
30
+
31
+ - SM100+ (Blackwell): FLASHINFER_CUTLASS (native MXFP8 W8A8 GEMM)
32
+ - SM80+ (Ampere/Ada): MARLIN (MXFP8 W8A16 GEMM)
33
+ - Otherwise: EMULATION (dequant to BF16 fallback)
34
+ """
35
+ from vllm.platforms import current_platform
36
+
37
+ if current_platform.has_device_capability(100):
38
+ return Mxfp8LinearBackend.FLASHINFER_CUTLASS
39
+
40
+ from vllm.model_executor.layers.quantization.utils.marlin_utils_fp8 import (
41
+ is_fp8_marlin_supported,
42
+ )
43
+
44
+ if is_fp8_marlin_supported():
45
+ return Mxfp8LinearBackend.MARLIN
46
+
47
+ return Mxfp8LinearBackend.EMULATION
48
+
49
+
50
def swizzle_mxfp8_scale(sf: torch.Tensor, M: int, K: int) -> torch.Tensor:
51
"""Swizzle MXFP8 scales from row-major 2D to F8_128x4 layout."""
52
scaling_vector_size = MXFP8_BLOCK_SIZE # 32 for MXFP8
@@ -47,17 +71,71 @@ def swizzle_mxfp8_scale(sf: torch.Tensor, M: int, K: int) -> torch.Tensor:
71
return sf_swizzled.contiguous().view(-1)
72
73
74
+def _mxfp8_e4m3_quantize_torch(
75
+ x: torch.Tensor,
76
+ is_sf_swizzled_layout: bool = False,
77
+) -> tuple[torch.Tensor, torch.Tensor]:
78
+ """Naive MXFP8 quantization.
79
+ For each block of 32 elements along the last dimension, compute a
80
+ shared e8m0 scale (the biased exponent of the block-wise amax)
81
+ and quantize each element to float8_e4m3fn.
82
+
83
+ Returns (quantized_values [same shape, fp8], scales uint8).
84
+ Scale shape depends on is_sf_swizzled_layout:
85
+ False -> [..., K//32] (row-major 2D)
86
+ True -> [flat swizzled 1D]
87
+ """
88
+ assert x.shape[-1] % MXFP8_BLOCK_SIZE == 0
89
+ orig_shape = x.shape
90
+ num_blocks = x.shape[-1] // MXFP8_BLOCK_SIZE
91
+
92
+ x_fp32 = x.to(torch.float32)
93
+ x_blocked = x_fp32.view(*orig_shape[:-1], num_blocks, MXFP8_BLOCK_SIZE)
94
+
95
+ amax = x_blocked.abs().amax(dim=-1)
96
+ amax = amax.clamp(min=torch.finfo(torch.float32).tiny)
97
+ scale_biased = torch.floor(torch.log2(amax)) + 127.0
98
+ scale_biased = scale_biased.clamp(0, 254)
99
+ scales_uint8 = scale_biased.to(torch.uint8)
100
+
101
+ descale = torch.exp2(scale_biased - 127.0)
102
+ x_scaled = x_blocked / descale.unsqueeze(-1)
103
+
104
+ x_fp8 = x_scaled.view(orig_shape).to(MXFP8_VALUE_DTYPE)
105
+
106
+ if x.ndim == 2:
107
+ M, K = x.shape
108
+ scales_uint8 = scales_uint8.view(M, -1)
109
+ if is_sf_swizzled_layout:
110
+ scales_uint8 = swizzle_mxfp8_scale(scales_uint8, M=M, K=K)
111
+ elif x.ndim == 3:
112
+ B, M, K = x.shape
113
+ scales_uint8 = scales_uint8.view(B, M, -1)
114
+ if is_sf_swizzled_layout:
115
+ swizzled = []
116
+ for i in range(B):
117
+ swizzled.append(swizzle_mxfp8_scale(scales_uint8[i], M=M, K=K))
118
+ scales_uint8 = torch.cat(swizzled)
119
+
120
+ return x_fp8, scales_uint8
121
+
122
+
123
def _mxfp8_e4m3_quantize_impl(
124
x: torch.Tensor, is_sf_swizzled_layout: bool = False
125
) -> tuple[torch.Tensor, torch.Tensor]:
53
- from flashinfer import mxfp8_quantize as flashinfer_mxfp8_quantize
126
+ from vllm.platforms import current_platform
127
55
- x_q, x_scales = flashinfer_mxfp8_quantize(
56
- x, is_sf_swizzled_layout=is_sf_swizzled_layout
57
- )
58
- if x_scales.ndim == 1 and x.ndim == 2 and not is_sf_swizzled_layout:
59
- x_scales = x_scales.view(x.size(0), -1)
60
- return x_q, x_scales
128
+ if current_platform.has_device_capability(100):
129
+ from flashinfer import mxfp8_quantize as flashinfer_mxfp8_quantize
130
+
131
+ x_q, x_scales = flashinfer_mxfp8_quantize(
132
+ x, is_sf_swizzled_layout=is_sf_swizzled_layout
133
+ )
134
+ if x_scales.ndim == 1 and x.ndim == 2 and not is_sf_swizzled_layout:
135
+ x_scales = x_scales.view(x.size(0), -1)
136
+ return x_q, x_scales
137
+
138
+ return _mxfp8_e4m3_quantize_torch(x, is_sf_swizzled_layout)
139
140
141
def mxfp8_e4m3_quantize(
@@ -128,11 +206,51 @@ direct_register_custom_op(
206
207
208
class Mxfp8LinearOp:
131
- def __init__(self, backend: Mxfp8LinearBackend):
132
- if backend not in Mxfp8LinearBackend:
133
- raise ValueError(f"Unsupported backend: {backend}")
209
+ def __init__(self):
210
+ self.backend = select_mxfp8_linear_backend()
211
+ logger.info_once("Using %s backend for MXFP8 GEMM", self.backend)
212
+
213
+ def process_weights(self, layer: torch.nn.Module) -> None:
214
+ """Process MXFP8 weights after loading into backend-specific format."""
215
+ if self.backend == Mxfp8LinearBackend.MARLIN:
216
+ self._process_weights_marlin(layer)
217
+ elif self.backend == Mxfp8LinearBackend.FLASHINFER_CUTLASS:
218
+ self._process_weights_flashinfer_cutlass(layer)
219
+ else:
220
+ self._process_weights_emulation(layer)
221
+
222
+ def _process_weights_emulation(self, layer: torch.nn.Module) -> None:
223
+ """Keep scales as 2D uint8 for dequant-to-BF16 emulation."""
224
+ weight = layer.weight.data # [N, K]
225
+ N, K = weight.shape
226
+ scale_k = K // MXFP8_BLOCK_SIZE
227
+
228
+ weight_scale = layer.weight_scale.data[:N, :scale_k].contiguous()
229
+
230
+ layer.weight = Parameter(weight.contiguous(), requires_grad=False)
231
+ layer.weight_scale = Parameter(weight_scale, requires_grad=False)
232
+
233
+ def _process_weights_flashinfer_cutlass(self, layer: torch.nn.Module) -> None:
234
+ """Swizzle scales to F8_128x4 layout for flashinfer CUTLASS."""
235
+ weight = layer.weight.data # [N, K]
236
+ N, K = weight.shape
237
135
- self.backend = backend
238
+ scale_k = K // MXFP8_BLOCK_SIZE
239
+ weight_scale_2d = layer.weight_scale.data[:N, :scale_k].contiguous()
240
+ weight_scale_swizzled = swizzle_mxfp8_scale(weight_scale_2d, M=N, K=K)
241
+
242
+ layer.weight = Parameter(weight.contiguous(), requires_grad=False)
243
+ layer.weight_scale = Parameter(
244
+ weight_scale_swizzled.contiguous(), requires_grad=False
245
+ )
246
+
247
+ def _process_weights_marlin(self, layer: torch.nn.Module) -> None:
248
+ """Repack MXFP8 weights and scales into Marlin kernel format."""
249
+ from vllm.model_executor.layers.quantization.utils.marlin_utils_fp8 import (
250
+ prepare_mxfp8_layer_for_marlin,
251
+ )
252
+
253
+ prepare_mxfp8_layer_for_marlin(layer)
254
255
def _apply_emulation(
256
self,
@@ -142,7 +260,6 @@ class Mxfp8LinearOp:
260
out_dtype: torch.dtype,
261
bias: torch.Tensor | None = None,
262
) -> torch.Tensor:
145
- # Validate weight_scale dtype and shape (must be 2D for TORCH backend)
263
if weight_scale.dtype != MXFP8_SCALE_DTYPE:
264
raise ValueError(
265
f"TORCH backend requires {MXFP8_SCALE_DTYPE} weight_scale dtype, "
@@ -219,6 +336,32 @@ class Mxfp8LinearOp:
336
output_shape = (*input_shape[:-1], N)
337
return output.view(output_shape)
338
339
+ def _apply_marlin(
340
+ self,
341
+ input: torch.Tensor,
342
+ weight: torch.Tensor,
343
+ weight_scale: torch.Tensor,
344
+ out_dtype: torch.dtype,
345
+ bias: torch.Tensor | None = None,
346
+ *,
347
+ workspace: torch.Tensor,
348
+ size_n: int,
349
+ size_k: int,
350
+ ) -> torch.Tensor:
351
+ from vllm.model_executor.layers.quantization.utils.marlin_utils_fp8 import (
352
+ apply_mxfp8_marlin_linear,
353
+ )
354
+
355
+ return apply_mxfp8_marlin_linear(
356
+ input=input,
357
+ weight=weight,
358
+ weight_scale=weight_scale,
359
+ workspace=workspace,
360
+ size_n=size_n,
361
+ size_k=size_k,
362
+ bias=bias,
363
+ )
364
+
365
def apply(
366
self,
367
input: torch.Tensor,
@@ -226,10 +369,27 @@ class Mxfp8LinearOp:
369
weight_scale: torch.Tensor,
370
out_dtype: torch.dtype,
371
bias: torch.Tensor | None = None,
372
+ *,
373
+ workspace: torch.Tensor | None = None,
374
+ size_n: int = 0,
375
+ size_k: int = 0,
376
) -> torch.Tensor:
377
if self.backend == Mxfp8LinearBackend.EMULATION:
378
return self._apply_emulation(input, weight, weight_scale, out_dtype, bias)
379
380
+ if self.backend == Mxfp8LinearBackend.MARLIN:
381
+ assert workspace is not None
382
+ return self._apply_marlin(
383
+ input,
384
+ weight,
385
+ weight_scale,
386
+ out_dtype,
387
+ bias,
388
+ workspace=workspace,
389
+ size_n=size_n,
390
+ size_k=size_k,
391
+ )
392
+
393
assert self.backend == Mxfp8LinearBackend.FLASHINFER_CUTLASS
394
return self._apply_flashinfer_cutlass(
395
input, weight, weight_scale, out_dtype, bias