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 15, 11:00 UTC
[Kernel] fp4 marlin kernel (#17687)
Signed-off-by: Jinzhen Lin <linjinzhen@hotmail.com>
Jinzhen Lin committed
May 11, 2025 at 10:58 UTC
d74e5f37bc63d297408460ff6cbf4d74d3c605b5
21 files changed
+1215
-330
csrc/core/scalar_type.hpp
+3
@@ -315,6 +315,8 @@ static inline constexpr auto kS8 = ScalarType::int_(8);
315
static inline constexpr auto kU8 = ScalarType::uint(8);
316
static inline constexpr auto kU8B128 = ScalarType::uint(8, 128);
317
318
+static inline constexpr auto kFE2M1f =
319
+ ScalarType::float_(2, 1, true, ScalarType::NAN_NONE);
320
static inline constexpr auto kFE3M2f =
321
ScalarType::float_(3, 2, true, ScalarType::NAN_NONE);
322
static inline constexpr auto kFE4M3fn =
@@ -332,6 +334,7 @@ static inline constexpr auto kInt8 = kS8;
334
static inline constexpr auto kUint8 = kU8;
335
static inline constexpr auto kUint8b128 = kU8B128;
336
337
+static inline constexpr auto kFloat4_e2m1f = kFE2M1f;
338
static inline constexpr auto kFloat6_e3m2f = kFE3M2f;
339
static inline constexpr auto kFloat8_e4m3fn = kFE4M3fn;
340
static inline constexpr auto kFloat8_e5m2 = kFE5M2;
csrc/moe/marlin_moe_wna16/generate_kernels.py
+11
-2
@@ -31,7 +31,10 @@ TEMPLATE = ("template __global__ void Marlin<"
31
32
# int8 with zero point case (vllm::kU8) is also supported,
33
# we don't add it to reduce wheel size.
34
-SCALAR_TYPES = ["vllm::kU4", "vllm::kU4B8", "vllm::kU8B128", "vllm::kFE4M3fn"]
34
+SCALAR_TYPES = [
35
+ "vllm::kU4", "vllm::kU4B8", "vllm::kU8B128", "vllm::kFE4M3fn",
36
+ "vllm::kFE2M1f"
37
+]
38
THREAD_CONFIGS = [(128, 128, 256), (64, 256, 256), (64, 128, 128)]
39
40
THREAD_M_BLOCKS = [0.5, 1, 2, 3, 4]
@@ -39,7 +42,7 @@ THREAD_M_BLOCKS = [0.5, 1, 2, 3, 4]
42
# = 0 : act order case
43
# = -1 : channelwise quantization
44
# > 0 : group_size=16*group_blocks
42
-GROUP_BLOCKS = [0, -1, 2, 4, 8]
45
+GROUP_BLOCKS = [0, -1, 1, 2, 4, 8]
46
DTYPES = ["fp16", "bf16"]
47
48
@@ -72,6 +75,12 @@ def generate_new_kernels():
75
# for fp8
76
if scalar_type == "vllm::kFE4M3fn" and group_blocks not in [-1, 8]:
77
continue
78
+ # nvfp4 only supports group_size == 16
79
+ if scalar_type == "vllm::kFE2M1f" and group_blocks not in [1, 2]:
80
+ continue
81
+ # other quantization methods don't support group_size = 16
82
+ if scalar_type != "vllm::kFE2M1f" and group_blocks == 1:
83
+ continue
84
85
k_blocks = thread_configs[0] // 16
86
n_blocks = thread_configs[1] // 16
csrc/moe/marlin_moe_wna16/kernel.h
+12
-11
@@ -7,17 +7,18 @@
7
#include "quantization/gptq_marlin/marlin_dtypes.cuh"
8
#include "core/scalar_type.hpp"
9
10
-#define MARLIN_KERNEL_PARAMS \
11
- const int4 *__restrict__ A, const int4 *__restrict__ B, \
12
- int4 *__restrict__ C, int4 *__restrict__ C_tmp, \
13
- const int4 *__restrict__ scales_ptr, const int4 *__restrict__ zp_ptr, \
14
- const int *__restrict__ g_idx, \
15
- const int32_t *__restrict__ sorted_token_ids_ptr, \
16
- const int32_t *__restrict__ expert_ids_ptr, \
17
- const int32_t *__restrict__ num_tokens_past_padded_ptr, \
18
- const float *__restrict__ topk_weights_ptr, int top_k, \
19
- bool mul_topk_weights, bool is_ep, int num_groups, int prob_m, \
20
- int prob_n, int prob_k, int *locks, bool use_atomic_add, \
10
+#define MARLIN_KERNEL_PARAMS \
11
+ const int4 *__restrict__ A, const int4 *__restrict__ B, \
12
+ int4 *__restrict__ C, int4 *__restrict__ C_tmp, \
13
+ const int4 *__restrict__ scales_ptr, \
14
+ const uint16_t *__restrict__ scale2_ptr, \
15
+ const int4 *__restrict__ zp_ptr, const int *__restrict__ g_idx, \
16
+ const int32_t *__restrict__ sorted_token_ids_ptr, \
17
+ const int32_t *__restrict__ expert_ids_ptr, \
18
+ const int32_t *__restrict__ num_tokens_past_padded_ptr, \
19
+ const float *__restrict__ topk_weights_ptr, int top_k, \
20
+ bool mul_topk_weights, bool is_ep, int num_groups, int prob_m, \
21
+ int prob_n, int prob_k, int *locks, bool use_atomic_add, \
22
bool use_fp32_reduce, int max_shared_mem
23
24
namespace MARLIN_NAMESPACE_NAME {
csrc/moe/marlin_moe_wna16/marlin_template.h
+99
-40
@@ -301,9 +301,11 @@ __global__ void Marlin(
301
int4* __restrict__ C_tmp, // fp32 tmp output buffer (for reduce)
302
const int4* __restrict__ scales_ptr, // fp16 quantization scales of shape
303
// (k/groupsize)xn
304
- const int4* __restrict__ zp_ptr, // 4bit packed zero-points of shape
305
- // (k/groupsize)x(n/pack_factor)
306
- const int* __restrict__ g_idx, // int32 group indices of shape k
304
+ const uint16_t* __restrict__ scale2_ptr, // fp16 global scale (for nvfp4
305
+ // only)
306
+ const int4* __restrict__ zp_ptr, // 4bit packed zero-points of shape
307
+ // (k/groupsize)x(n/pack_factor)
308
+ const int* __restrict__ g_idx, // int32 group indices of shape k
309
const int32_t* __restrict__ sorted_token_ids_ptr, // moe sorted_ids
310
const int32_t* __restrict__ expert_ids_ptr, // moe expert ids
311
const int32_t* __restrict__ num_tokens_past_padded_ptr, // moe num tokens
@@ -341,6 +343,16 @@ __global__ void Marlin(
343
extern __shared__ int4 sh[];
344
static constexpr auto w_type = vllm::ScalarType::from_id(w_type_id);
345
constexpr bool has_zp = w_type == vllm::kU4 || w_type == vllm::kU8;
346
+ constexpr bool is_int_type = w_type == vllm::kU4 || w_type == vllm::kU8 ||
347
+ w_type == vllm::kU4B8 || w_type == vllm::kU8B128;
348
+ // see comments of dequant.h for more details
349
+ constexpr bool dequant_skip_flop =
350
+ !is_int_type ||
351
+ has_zp && !is_zp_float && !std::is_same<scalar_t, nv_bfloat16>::value ||
352
+ has_zp && !is_zp_float && !(w_type == vllm::kU8);
353
+
354
+ scalar_t2 global_scale;
355
+
356
constexpr bool has_act_order = group_blocks == 0;
357
358
constexpr int pack_factor = 32 / w_type.size_bits();
@@ -348,7 +360,8 @@ __global__ void Marlin(
360
constexpr int moe_block_size = m_block_size_8 ? 8 : (16 * thread_m_blocks);
361
const int group_size =
362
(!has_act_order && group_blocks == -1) ? prob_k : prob_k / num_groups;
351
- const int scales_expert_stride = prob_n * prob_k / group_size / 8;
363
+ const int scales_expert_stride =
364
+ prob_n * prob_k / group_size / (w_type == vllm::kFE2M1f ? 16 : 8);
365
const int zp_expert_stride =
366
is_zp_float ? prob_n * prob_k / group_size / 8
367
: prob_n * prob_k / group_size / (pack_factor * 4);
@@ -460,9 +473,16 @@ __global__ void Marlin(
473
if (mul_topk_weights) {
474
#pragma unroll
475
for (int i = 0; i < 4; i++) {
463
- sh_block_topk_weights[tid4 * 4 + i] =
464
- Dtype::num2num2(Dtype::float2num(
465
- topk_weights_ptr[sh_block_sorted_ids[tid4 * 4 + i]]));
476
+ if constexpr (w_type == vllm::kFE2M1f) {
477
+ sh_block_topk_weights[tid4 * 4 + i] = __hmul2(
478
+ global_scale,
479
+ Dtype::num2num2(Dtype::float2num(
480
+ topk_weights_ptr[sh_block_sorted_ids[tid4 * 4 + i]])));
481
+ } else {
482
+ sh_block_topk_weights[tid4 * 4 + i] =
483
+ Dtype::num2num2(Dtype::float2num(
484
+ topk_weights_ptr[sh_block_sorted_ids[tid4 * 4 + i]]));
485
+ }
486
}
487
}
488
}
@@ -493,6 +513,11 @@ __global__ void Marlin(
513
expert_id = expert_ids_ptr[block_id];
514
}
515
516
+ if constexpr (w_type == vllm::kFE2M1f) {
517
+ uint16_t val = scale2_ptr[expert_id];
518
+ global_scale = Dtype::num2num2(*reinterpret_cast<scalar_t*>(&val));
519
+ }
520
+
521
B_expert_off = expert_id * prob_n * prob_k / (pack_factor * 4);
522
scales_ptr += (expert_id - old_expert_id) * scales_expert_stride;
523
if constexpr (has_zp) {
@@ -606,7 +631,7 @@ __global__ void Marlin(
631
constexpr int s_sh_stride = 16 * thread_n_blocks / 8;
632
constexpr int s_tb_groups =
633
!has_act_order && group_blocks != -1 && group_blocks < thread_k_blocks
609
- ? thread_k_blocks / group_blocks
634
+ ? thread_k_blocks / group_blocks / (w_type == vllm::kFE2M1f ? 2 : 1)
635
: 1;
636
constexpr int s_sh_stage = s_tb_groups * s_sh_stride;
637
int s_gl_rd_delta = s_gl_stride;
@@ -664,7 +689,8 @@ __global__ void Marlin(
689
if constexpr (group_blocks == -1) {
690
s_gl_rd = s_sh_stride * slice_col + threadIdx.x;
691
} else {
667
- s_gl_rd = s_gl_stride * ((thread_k_blocks * slice_row) / group_blocks) +
692
+ s_gl_rd = s_gl_stride * ((thread_k_blocks * slice_row) / group_blocks) /
693
+ (w_type == vllm::kFE2M1f ? 2 : 1) +
694
s_sh_stride * slice_col + threadIdx.x;
695
}
696
}
@@ -688,10 +714,20 @@ __global__ void Marlin(
714
// we scale a `half2` tile in column-major layout in the former and in
715
// row-major in the latter case.
716
int s_sh_rd;
691
- if constexpr (group_blocks != -1)
717
+ if constexpr (group_blocks != -1 && w_type == vllm::kFE2M1f) {
718
+ auto warp_id = threadIdx.x / 32;
719
+ int n_warps = thread_n_blocks / 4;
720
+ int warp_row = warp_id / n_warps;
721
+
722
s_sh_rd = 8 * ((threadIdx.x / 32) % (thread_n_blocks / 4)) +
723
(threadIdx.x % 32) / 4;
694
- else if constexpr (group_blocks == -1 && (m_block_size_8 || has_zp))
724
+ s_sh_rd = s_sh_rd * 2 + warp_row % 2;
725
+
726
+ } else if constexpr (group_blocks != -1)
727
+ s_sh_rd = 8 * ((threadIdx.x / 32) % (thread_n_blocks / 4)) +
728
+ (threadIdx.x % 32) / 4;
729
+ else if constexpr (group_blocks == -1 &&
730
+ (m_block_size_8 || (has_zp && !dequant_skip_flop)))
731
s_sh_rd = 8 * ((threadIdx.x / 32) % (thread_n_blocks / 4)) +
732
(threadIdx.x % 32) / 8;
733
else
@@ -801,7 +837,7 @@ __global__ void Marlin(
837
sh_first_group_id = first_group_id;
838
sh_num_groups = last_group_id - first_group_id + 1;
839
804
- if (sh_num_groups < act_s_max_num_groups) {
840
+ if (sh_num_groups > act_s_max_num_groups) {
841
sh_num_groups = act_s_max_num_groups;
842
}
843
@@ -1021,12 +1057,19 @@ __global__ void Marlin(
1057
cur_k += k_iter_size * (k % b_sh_wr_iters);
1058
1059
int k_blocks = cur_k / 16;
1024
- int cur_group_id = k_blocks / group_blocks;
1060
+ int cur_group_id =
1061
+ k_blocks / (group_blocks * (w_type == vllm::kFE2M1f ? 2 : 1));
1062
1063
int4* sh_s_stage = sh_s + s_sh_stage * pipe;
1064
1028
- reinterpret_cast<int4*>(&frag_s[k % 2])[0] =
1029
- sh_s_stage[s_sh_rd + cur_group_id * s_sh_stride];
1065
+ if constexpr (w_type_id != vllm::kFE2M1f.id()) {
1066
+ reinterpret_cast<int4*>(&frag_s[k % 2])[0] =
1067
+ sh_s_stage[s_sh_rd + cur_group_id * s_sh_stride];
1068
+ } else {
1069
+ reinterpret_cast<int2*>(&frag_s[k % 2])[0] =
1070
+ reinterpret_cast<int2*>(
1071
+ sh_s_stage)[s_sh_rd + cur_group_id * (2 * s_sh_stride)];
1072
+ }
1073
}
1074
}
1075
@@ -1199,22 +1242,7 @@ __global__ void Marlin(
1242
};
1243
1244
auto dequant_data = [&](int q, scalar_t2* frag_b_ptr) {
1202
- if constexpr (has_zp && is_zp_float || !has_zp) {
1203
- dequant<scalar_t2, w_type_id>(q, frag_b_ptr);
1204
- } else {
1205
- static_assert(has_zp && !is_zp_float);
1206
- static_assert(w_type_id == vllm::kU4.id() || w_type_id == vllm::kU8.id());
1207
- // If (has_zp && !is_zp_float),
1208
- // we use not-zp version `dequant` function
1209
- // to improve numerical accuracy.
1210
- // Since both weight and zero point are dequanted using this logic,
1211
- // the final dequanted weight would be correct.
1212
- if constexpr (w_type_id == vllm::kU4.id()) {
1213
- dequant<scalar_t2, vllm::kU4B8.id()>(q, frag_b_ptr);
1214
- } else if constexpr (w_type_id == vllm::kU8.id()) {
1215
- dequant<scalar_t2, vllm::kU8B128.id()>(q, frag_b_ptr);
1216
- }
1217
- }
1245
+ dequant<scalar_t2, w_type_id, dequant_skip_flop>(q, frag_b_ptr);
1246
};
1247
1248
// Execute the actual tensor core matmul of a sub-tile.
@@ -1244,13 +1272,23 @@ __global__ void Marlin(
1272
dequant_data(zp_quant_1, reinterpret_cast<scalar_t2*>(&frag_zp) + 2);
1273
}
1274
}
1247
- if constexpr (has_zp && is_zp_float) {
1275
+ if constexpr (!dequant_skip_flop && has_zp && is_zp_float) {
1276
if (is_new_zp) {
1277
reinterpret_cast<int4*>(&frag_zp)[0] =
1278
reinterpret_cast<int4*>(&frag_zpf[k2])[0];
1279
}
1280
}
1281
1282
+ if constexpr (w_type == vllm::kFE2M1f) {
1283
+ int s_quant_0 = reinterpret_cast<int*>(frag_s[k2])[0];
1284
+ int s_quant_1 = reinterpret_cast<int*>(frag_s[k2])[1];
1285
+
1286
+ dequant_fp8_scales<scalar_t2>(s_quant_0,
1287
+ reinterpret_cast<scalar_t2*>(&frag_s[k2]));
1288
+ dequant_fp8_scales<scalar_t2>(
1289
+ s_quant_1, reinterpret_cast<scalar_t2*>(&frag_s[k2]) + 2);
1290
+ }
1291
+
1292
// We have the m dimension as the inner loop in order to encourage overlapping
1293
// dequantization and matmul operations.
1294
#pragma unroll
@@ -1259,7 +1297,10 @@ __global__ void Marlin(
1297
FragB frag_b1;
1298
int b_quant_0, b_quant_1;
1299
1262
- if constexpr (w_type.size_bits() == 4) {
1300
+ if constexpr (w_type_id == vllm::kFE2M1f.id()) {
1301
+ b_quant_1 = frag_b_quant[k2][0][j];
1302
+ b_quant_0 = b_quant_1 << 8;
1303
+ } else if constexpr (w_type.size_bits() == 4) {
1304
b_quant_0 = frag_b_quant[k2][0][j];
1305
b_quant_1 = b_quant_0 >> 8;
1306
} else {
@@ -1272,6 +1313,11 @@ __global__ void Marlin(
1313
dequant_data(b_quant_0, reinterpret_cast<scalar_t2*>(&frag_b0));
1314
dequant_data(b_quant_1, reinterpret_cast<scalar_t2*>(&frag_b1));
1315
1316
+ if constexpr (dequant_skip_flop && has_zp && !is_zp_float) {
1317
+ sub_zp<scalar_t>(frag_b0, frag_zp[j], 0);
1318
+ sub_zp<scalar_t>(frag_b1, frag_zp[j], 1);
1319
+ }
1320
+
1321
// Apply scale to frag_b0
1322
if constexpr (has_act_order) {
1323
static_assert(group_blocks != -1);
@@ -1279,7 +1325,8 @@ __global__ void Marlin(
1325
act_frag_s[k2][2][j], act_frag_s[k2][3][j], 0);
1326
scale4<scalar_t>(frag_b1, act_frag_s[k2][0][j], act_frag_s[k2][1][j],
1327
act_frag_s[k2][2][j], act_frag_s[k2][3][j], 1);
1282
- } else if constexpr (has_zp && !is_zp_float && group_blocks == -1) {
1328
+ } else if constexpr (!dequant_skip_flop && has_zp && !is_zp_float &&
1329
+ group_blocks == -1) {
1330
int idx = (threadIdx.x / 4) % 2;
1331
scalar_t2 s2 = Dtype::nums2num2(
1332
reinterpret_cast<scalar_t*>(&frag_s[j / 2][j % 2 * 2 + 0])[idx],
@@ -1287,7 +1334,7 @@ __global__ void Marlin(
1334
if (is_new_zp) frag_zp[j] = __hmul2(frag_zp[j], s2);
1335
scale_and_sub<scalar_t>(frag_b0, s2.x, frag_zp[j].x);
1336
scale_and_sub<scalar_t>(frag_b1, s2.y, frag_zp[j].y);
1290
- } else if constexpr (has_zp && group_blocks != -1) {
1337
+ } else if constexpr (!dequant_skip_flop && has_zp && group_blocks != -1) {
1338
if (is_new_zp)
1339
frag_zp[j] = __hmul2(frag_zp[j],
1340
*reinterpret_cast<scalar_t2*>(&frag_s[k2][j]));
@@ -1554,10 +1601,17 @@ __global__ void Marlin(
1601
// For per-column quantization we finally apply the scale here (only for
1602
// 4-bit)
1603
if constexpr (!has_act_order && group_blocks == -1 &&
1557
- w_type.size_bits() == 4 && !has_zp) {
1604
+ w_type.size_bits() == 4 &&
1605
+ (has_zp && dequant_skip_flop || !has_zp)) {
1606
res = __hmul2(res, s[0]);
1607
}
1608
1609
+ if constexpr (w_type == vllm::kFE2M1f) {
1610
+ if (!mul_topk_weights) {
1611
+ res = __hmul2(res, global_scale);
1612
+ }
1613
+ }
1614
+
1615
if constexpr (m_block_size_8) {
1616
((scalar_t*)sh_red)[idx] = res.x;
1617
((scalar_t*)sh_red)[idx + 8 * c_sh_stride] = res.y;
@@ -1648,7 +1702,9 @@ __global__ void Marlin(
1702
if constexpr (has_zp && !is_zp_float && group_blocks == -1) {
1703
if (i == 0) {
1704
fetch_col_zp_to_shared();
1651
- fetch_col_scale_to_shared();
1705
+ if constexpr (!dequant_skip_flop) {
1706
+ fetch_col_scale_to_shared();
1707
+ }
1708
}
1709
}
1710
fetch_to_shared(i, i, i < slice_iters, i);
@@ -1737,7 +1793,8 @@ __global__ void Marlin(
1793
bool last = slice_idx == slice_count - 1;
1794
// For per-column scales, we only fetch them here in the final step before
1795
// write-out
1740
- if constexpr (!has_act_order && group_blocks == -1 && !has_zp) {
1796
+ if constexpr (!has_act_order && group_blocks == -1 &&
1797
+ (has_zp && dequant_skip_flop || !has_zp)) {
1798
if (w_type.size_bits() == 8 || (last || use_atomic_add)) {
1799
if (s_sh_wr_pred) {
1800
cp_async4(&sh_s[s_sh_wr], &scales_ptr[s_gl_rd]);
@@ -1747,7 +1804,8 @@ __global__ void Marlin(
1804
}
1805
1806
thread_block_reduce();
1750
- if constexpr (!has_act_order && group_blocks == -1 && !has_zp) {
1807
+ if constexpr (!has_act_order && group_blocks == -1 &&
1808
+ (has_zp && dequant_skip_flop || !has_zp)) {
1809
if (w_type.size_bits() == 8 || (last || use_atomic_add)) {
1810
cp_async_wait<0>();
1811
__syncthreads();
@@ -1771,7 +1829,8 @@ __global__ void Marlin(
1829
// that converts the fp32 results to fp16 (so that we avoid possible
1830
// overflow in fp16)
1831
if constexpr (!has_act_order && group_blocks == -1 &&
1774
- w_type.size_bits() == 8 && !has_zp) {
1832
+ w_type.size_bits() == 8 &&
1833
+ (has_zp && dequant_skip_flop || !has_zp)) {
1834
if (threadIdx.x / 32 < thread_n_blocks / 4) {
1835
#pragma unroll
1836
for (int i = 0; i < thread_m_blocks; i++) {
csrc/moe/marlin_moe_wna16/ops.cu
+66
-18
@@ -291,6 +291,7 @@ bool is_valid_config(thread_config_t const& th_config, bool m_block_size_8,
291
// BIGGROUP: cases for big group size (group_blocks in [-1, 8])
292
// FZP: cases for float-zero-point (is_zp_float = true)
293
// ACT: cases for act order case (group_blocks == 0)
294
+ // FP4: cases for nvfp4(e2m1) (group_blocks == 1)
295
#define COMMON_GET_IF_M1(W_TYPE, N_BLOCKS, K_BLOCKS, NUM_THREADS) \
296
_GET_IF(W_TYPE, 1, N_BLOCKS, K_BLOCKS, true, -1, NUM_THREADS, false) \
297
_GET_IF(W_TYPE, 1, N_BLOCKS, K_BLOCKS, true, 2, NUM_THREADS, false) \
@@ -338,6 +339,21 @@ bool is_valid_config(thread_config_t const& th_config, bool m_block_size_8,
339
_GET_IF(W_TYPE, 4, N_BLOCKS, K_BLOCKS, false, -1, NUM_THREADS, false) \
340
_GET_IF(W_TYPE, 4, N_BLOCKS, K_BLOCKS, false, 8, NUM_THREADS, false)
341
342
+ #define FP4_GET_IF_M1(W_TYPE, N_BLOCKS, K_BLOCKS, NUM_THREADS) \
343
+ _GET_IF(W_TYPE, 1, N_BLOCKS, K_BLOCKS, true, 1, NUM_THREADS, false) \
344
+ _GET_IF(W_TYPE, 1, N_BLOCKS, K_BLOCKS, false, 1, NUM_THREADS, false)
345
+
346
+ #define FP4_GET_IF_M234(W_TYPE, N_BLOCKS, K_BLOCKS, NUM_THREADS) \
347
+ _GET_IF(W_TYPE, 2, N_BLOCKS, K_BLOCKS, false, 1, NUM_THREADS, false) \
348
+ _GET_IF(W_TYPE, 3, N_BLOCKS, K_BLOCKS, false, 1, NUM_THREADS, false) \
349
+ _GET_IF(W_TYPE, 4, N_BLOCKS, K_BLOCKS, false, 1, NUM_THREADS, false)
350
+
351
+ #define FP4_GET_IF(W_TYPE) \
352
+ FP4_GET_IF_M1(W_TYPE, 8, 8, 256) \
353
+ FP4_GET_IF_M1(W_TYPE, 8, 4, 128) \
354
+ FP4_GET_IF_M234(W_TYPE, 16, 4, 256) \
355
+ FP4_GET_IF_M234(W_TYPE, 8, 4, 128)
356
+
357
#define BIGGROUP_GET_IF(W_TYPE) \
358
BIGGROUP_GET_IF_M1(W_TYPE, 8, 8, 256) \
359
BIGGROUP_GET_IF_M1(W_TYPE, 8, 4, 128) \
@@ -394,6 +410,8 @@ MarlinFuncPtr get_marlin_kernel(const vllm::ScalarType q_type,
410
411
BIGGROUP_GET_IF(vllm::kFE4M3fn)
412
413
+ FP4_GET_IF(vllm::kFE2M1f)
414
+
415
ACT_GET_IF(vllm::kU4B8)
416
ACT_GET_IF(vllm::kU8B128)
417
@@ -465,7 +483,7 @@ exec_config_t determine_exec_config(const vllm::ScalarType& q_type, int prob_m,
483
484
template <typename scalar_t>
485
void marlin_mm(const void* A, const void* B, void* C, void* C_tmp, void* s,
468
- void* zp, void* g_idx, void* perm, void* a_tmp,
486
+ void* s2, void* zp, void* g_idx, void* perm, void* a_tmp,
487
void* sorted_token_ids, void* expert_ids,
488
void* num_tokens_past_padded, void* topk_weights,
489
int moe_block_size, int top_k, bool mul_topk_weights, bool is_ep,
@@ -479,14 +497,16 @@ void marlin_mm(const void* A, const void* B, void* C, void* C_tmp, void* s,
497
bool m_block_size_8 = moe_block_size == 8;
498
499
if (has_zp) {
482
- TORCH_CHECK(q_type == vllm::kU4,
483
- "q_type must be u4 when has_zp = True. Got = ", q_type.str());
500
+ TORCH_CHECK(
501
+ q_type == vllm::kU4 || q_type == vllm::kU8,
502
+ "q_type must be u4 or u8 when has_zp = True. Got = ", q_type.str());
503
} else {
485
- TORCH_CHECK(q_type == vllm::kU4B8 || q_type == vllm::kU8B128 ||
486
- q_type == vllm::kFE4M3fn,
487
- "q_type must be uint4b8, uint8b128 or fp8e4m3 when has_zp = "
488
- "False. Got = ",
489
- q_type.str());
504
+ TORCH_CHECK(
505
+ q_type == vllm::kU4B8 || q_type == vllm::kU8B128 ||
506
+ q_type == vllm::kFE4M3fn || q_type == vllm::kFE2M1f,
507
+ "q_type must be uint4b8, uint8b128, float8_e4m3fn or float4_e2m1f when "
508
+ "has_zp = False. Got = ",
509
+ q_type.str());
510
}
511
512
TORCH_CHECK(prob_m > 0 && prob_n > 0 && prob_k > 0, "Invalid MNK = [", prob_m,
@@ -519,6 +539,7 @@ void marlin_mm(const void* A, const void* B, void* C, void* C_tmp, void* s,
539
int4* C_ptr = (int4*)C;
540
int4* C_tmp_ptr = (int4*)C_tmp;
541
const int4* s_ptr = (const int4*)s;
542
+ const uint16_t* s2_ptr = (const uint16_t*)s2;
543
const int4* zp_ptr = (const int4*)zp;
544
const int* g_idx_ptr = (const int*)g_idx;
545
const int* perm_ptr = (const int*)perm;
@@ -627,7 +648,7 @@ void marlin_mm(const void* A, const void* B, void* C, void* C_tmp, void* s,
648
// avoid ">>>" being formatted to "> > >"
649
// clang-format off
650
kernel<<<blocks, num_threads, max_shared_mem, stream>>>(
630
- A_ptr, B_ptr, C_ptr, C_tmp_ptr, s_ptr, zp_ptr, g_idx_ptr,
651
+ A_ptr, B_ptr, C_ptr, C_tmp_ptr, s_ptr, s2_ptr, zp_ptr, g_idx_ptr,
652
sorted_token_ids_ptr, expert_ids_ptr, num_tokens_past_padded_ptr,
653
topk_weights_ptr, top_k, mul_topk_weights, is_ep, num_groups, prob_m,
654
prob_n, prob_k, locks, use_atomic_add, use_fp32_reduce, max_shared_mem);
@@ -639,6 +660,7 @@ void marlin_mm(const void* A, const void* B, void* C, void* C_tmp, void* s,
660
torch::Tensor moe_wna16_marlin_gemm(
661
torch::Tensor& a, std::optional<torch::Tensor> const& c_or_none,
662
torch::Tensor& b_q_weight, torch::Tensor& b_scales,
663
+ std::optional<torch::Tensor> const& global_scale_or_none,
664
std::optional<torch::Tensor> const& b_zeros_or_none,
665
std::optional<torch::Tensor> const& g_idx_or_none,
666
std::optional<torch::Tensor> const& perm_or_none, torch::Tensor& workspace,
@@ -790,6 +812,17 @@ torch::Tensor moe_wna16_marlin_gemm(
812
}
813
}
814
815
+ torch::Tensor global_scale;
816
+ if (global_scale_or_none.has_value()) {
817
+ global_scale = global_scale_or_none.value();
818
+ TORCH_CHECK(b_q_type == vllm::kFE2M1f,
819
+ "global_scale can only be used for float4_e2m1f.");
820
+ } else {
821
+ global_scale = torch::empty({0}, options);
822
+ TORCH_CHECK(!(b_q_type == vllm::kFE2M1f),
823
+ "the global_scale parameter must be passed for float4_e2m1f.");
824
+ }
825
+
826
torch::Tensor b_zeros;
827
if (b_zeros_or_none.has_value()) {
828
b_zeros = b_zeros_or_none.value();
@@ -802,13 +835,14 @@ torch::Tensor moe_wna16_marlin_gemm(
835
836
if (has_zp) {
837
TORCH_CHECK(
805
- b_q_type == vllm::kU4,
806
- "b_q_type must be u4 when has_zp = True. Got = ", b_q_type.str());
838
+ b_q_type == vllm::kU4 || b_q_type == vllm::kU8,
839
+ "b_q_type must be u4 or u8 when has_zp = True. Got = ", b_q_type.str());
840
} else {
841
TORCH_CHECK(b_q_type == vllm::kU4B8 || b_q_type == vllm::kU8B128 ||
809
- b_q_type == vllm::kFE4M3fn,
810
- "b_q_type must be uint4b8, uint8b128 or fp8e4m3 when has_zp = "
811
- "False. Got = ",
842
+ b_q_type == vllm::kFE4M3fn || b_q_type == vllm::kFE2M1f,
843
+ "b_q_type must be uint4b8, uint8b128, float8_e4m3fn or "
844
+ "float4_e2m1f when "
845
+ "has_zp = False. Got = ",
846
b_q_type.str());
847
}
848
@@ -854,9 +888,16 @@ torch::Tensor moe_wna16_marlin_gemm(
888
889
int dev = a.get_device();
890
if (a.scalar_type() == at::ScalarType::Half) {
891
+ void* scales_ptr;
892
+ if (b_q_type == vllm::kFE2M1f) {
893
+ scales_ptr = b_scales.data_ptr<at::Float8_e4m3fn>();
894
+ } else {
895
+ scales_ptr = b_scales.data_ptr<at::Half>();
896
+ }
897
+
898
MARLIN_NAMESPACE_NAME::marlin_mm<half>(
899
a.data_ptr<at::Half>(), b_q_weight.data_ptr(), c.data_ptr<at::Half>(),
859
- c_tmp.data_ptr<float>(), b_scales.data_ptr<at::Half>(),
900
+ c_tmp.data_ptr<float>(), scales_ptr, global_scale.data_ptr<at::Half>(),
901
b_zeros.data_ptr(), g_idx.data_ptr(), perm.data_ptr(),
902
a_tmp.data_ptr<at::Half>(), sorted_token_ids.data_ptr(),
903
expert_ids.data_ptr(), num_tokens_past_padded.data_ptr(),
@@ -866,11 +907,18 @@ torch::Tensor moe_wna16_marlin_gemm(
907
at::cuda::getCurrentCUDAStream(dev), thread_k, thread_n, sms,
908
use_atomic_add, use_fp32_reduce, is_zp_float);
909
} else if (a.scalar_type() == at::ScalarType::BFloat16) {
910
+ void* scales_ptr;
911
+ if (b_q_type == vllm::kFE2M1f) {
912
+ scales_ptr = b_scales.data_ptr<at::Float8_e4m3fn>();
913
+ } else {
914
+ scales_ptr = b_scales.data_ptr<at::BFloat16>();
915
+ }
916
+
917
MARLIN_NAMESPACE_NAME::marlin_mm<nv_bfloat16>(
918
a.data_ptr<at::BFloat16>(), b_q_weight.data_ptr(),
871
- c.data_ptr<at::BFloat16>(), c_tmp.data_ptr<float>(),
872
- b_scales.data_ptr<at::BFloat16>(), b_zeros.data_ptr(), g_idx.data_ptr(),
873
- perm.data_ptr(), a_tmp.data_ptr<at::BFloat16>(),
919
+ c.data_ptr<at::BFloat16>(), c_tmp.data_ptr<float>(), scales_ptr,
920
+ global_scale.data_ptr<at::BFloat16>(), b_zeros.data_ptr(),
921
+ g_idx.data_ptr(), perm.data_ptr(), a_tmp.data_ptr<at::BFloat16>(),
922
sorted_token_ids.data_ptr(), expert_ids.data_ptr(),
923
num_tokens_past_padded.data_ptr(), topk_weights.data_ptr(),
924
moe_block_size, top_k, mul_topk_weights, is_ep, size_m, size_n, size_k,
csrc/moe/torch_bindings.cpp
+2
-1
@@ -44,7 +44,8 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, m) {
44
45
m.def(
46
"moe_wna16_marlin_gemm(Tensor! a, Tensor? c_or_none,"
47
- "Tensor! b_q_weight, Tensor! b_scales, Tensor? b_zeros_or_none,"
47
+ "Tensor! b_q_weight, Tensor! b_scales, Tensor? global_scale, Tensor? "
48
+ "b_zeros_or_none,"
49
"Tensor? g_idx_or_none, Tensor? perm_or_none, Tensor! workspace,"
50
"Tensor sorted_token_ids,"
51
"Tensor! expert_ids, Tensor! num_tokens_past_padded,"
csrc/quantization/gptq_marlin/dequant.h
+286
-70
@@ -1,3 +1,67 @@
1
+/*
2
+Fast Dequantization (Converting INT4/INT8/FP4/FP8 to FP16/BF16)
3
+
4
+The process of fast dequantization can be summarized as a combination
5
+of bitwise operations and floating-point computations:
6
+
7
+weight =>(bit_op / bitwise operations)=>
8
+f16_value =>(flop / floating-point computation)=>
9
+dequantized_weight
10
+
11
+Since the dequantized weights typically require subtracting the zero point and
12
+applying a scale factor, the floating-point computation step can be fused with
13
+the zero-point subtraction and scaling operations.
14
+
15
+The following are the parts that need to be modified for the fused operation
16
+of zero-point subtraction and scaling.
17
+
18
+## INT4 => FP16/BF16 or INT8 => FP16
19
+
20
+The floating-point computation is `__hsub2`
21
+
22
+If has zero points:
23
+
24
+ flop(bit_op(weight)) - flop(bit_op(zp))
25
+ = sub(bit_op(weight), bias) - sub(bit_op(zp), bias)
26
+ = bit_op(weight) - bit_op(zp)
27
+
28
+so we don't need additional modification.
29
+
30
+If has float zero points:
31
+
32
+ flop(bit_op(weight)) - fzp
33
+ = sub(bit_op(weight), bias) - fzp
34
+ = bit_op(weight) - (fzp + bias)
35
+
36
+where the `fzp + bias` can be computed at weight loading. But this
37
+may have accuracy issue, so we should not use this in most cases.
38
+
39
+If has not zero points:
40
+
41
+ scale(flop(bit_op(weight)))
42
+ = scale(sub(bit_op(weight), bias))
43
+ = scale(bit_op(weight)) - scale(bias)
44
+ = fma(bit_op(weight), scale_factor, scale(bias))
45
+
46
+where the `scale(bias)` can be cached. But this may have accuracy issue,
47
+so we should not use this in most cases.
48
+
49
+
50
+## INT8 => BF16
51
+
52
+INT8 => BF16 is a special case, it use byte_perm instead of flop.
53
+We cannot fused byte_perm with scaling.
54
+
55
+
56
+## FP4/FP8 => FP16/BF16
57
+
58
+ scale(flop(bit_op(weight)))
59
+ = scale(mul(bit_op(weight), multiplier))
60
+ = mul(bit_op(weight), scale_factor * multiplier)
61
+
62
+where `scale_factor * multiplier` can be computed at weight loading.
63
+
64
+*/
65
66
#include "marlin_dtypes.cuh"
67
@@ -27,7 +91,8 @@ __device__ inline uint32_t prmt(uint32_t a) {
91
return res;
92
}
93
30
-template <typename scalar_t2, vllm::ScalarTypeId w_type_id>
94
+template <typename scalar_t2, vllm::ScalarTypeId w_type_id,
95
+ bool skip_flop = false>
96
__device__ inline void dequant(int q, scalar_t2* frag_b);
97
98
//
@@ -40,7 +105,22 @@ __device__ inline void dequant(int q, scalar_t2* frag_b);
105
// https://github.com/NVIDIA/FasterTransformer/blob/release/v5.3_tag/src/fastertransformer/cutlass_extensions/include/cutlass_extensions/interleaved_numeric_conversion.h#L327-L385
106
//
107
template <>
43
-__device__ inline void dequant<half2, vllm::kU4B8.id()>(int q, half2* frag_b) {
108
+__device__ inline void dequant<half2, vllm::kU4B8.id(), true>(int q,
109
+ half2* frag_b) {
110
+ const int MASK = 0x000f000f;
111
+ const int EX = 0x64006400;
112
+ // Guarantee that the `(a & b) | c` operations are LOP3s.
113
+ int lo = lop3<(0xf0 & 0xcc) | 0xaa>(q, MASK, EX);
114
+ q >>= 4;
115
+ int hi = lop3<(0xf0 & 0xcc) | 0xaa>(q, MASK, EX);
116
+
117
+ frag_b[0] = *reinterpret_cast<half2*>(&lo);
118
+ frag_b[1] = *reinterpret_cast<half2*>(&hi);
119
+}
120
+
121
+template <>
122
+__device__ inline void dequant<half2, vllm::kU4B8.id(), false>(int q,
123
+ half2* frag_b) {
124
const int LO = 0x000f000f;
125
const int HI = 0x00f000f0;
126
const int EX = 0x64006400;
@@ -62,7 +142,14 @@ __device__ inline void dequant<half2, vllm::kU4B8.id()>(int q, half2* frag_b) {
142
}
143
144
template <>
65
-__device__ inline void dequant<half2, vllm::kU4.id()>(int q, half2* frag_b) {
145
+__device__ inline void dequant<half2, vllm::kU4.id(), true>(int q,
146
+ half2* frag_b) {
147
+ dequant<half2, vllm::kU4B8.id(), true>(q, frag_b);
148
+}
149
+
150
+template <>
151
+__device__ inline void dequant<half2, vllm::kU4.id(), false>(int q,
152
+ half2* frag_b) {
153
const int LO = 0x000f000f;
154
const int HI = 0x00f000f0;
155
const int EX = 0x64006400;
@@ -84,7 +171,7 @@ __device__ inline void dequant<half2, vllm::kU4.id()>(int q, half2* frag_b) {
171
}
172
173
template <>
87
-__device__ inline void dequant<nv_bfloat162, vllm::kU4B8.id()>(
174
+__device__ inline void dequant<nv_bfloat162, vllm::kU4B8.id(), true>(
175
int q, nv_bfloat162* frag_b) {
176
static constexpr uint32_t MASK = 0x000f000f;
177
static constexpr uint32_t EX = 0x43004300;
@@ -96,39 +183,36 @@ __device__ inline void dequant<nv_bfloat162, vllm::kU4B8.id()>(
183
int hi = lop3<(0xf0 & 0xcc) | 0xaa>(q, MASK, EX);
184
// clang-format on
185
99
- static constexpr uint32_t MUL = 0x3F803F80;
100
- static constexpr uint32_t ADD = 0xC308C308;
186
+ frag_b[0] = *reinterpret_cast<nv_bfloat162*>(&lo);
187
+ frag_b[1] = *reinterpret_cast<nv_bfloat162*>(&hi);
188
+}
189
+
190
+template <>
191
+__device__ inline void dequant<nv_bfloat162, vllm::kU4B8.id(), false>(
192
+ int q, nv_bfloat162* frag_b) {
193
+ dequant<nv_bfloat162, vllm::kU4B8.id(), true>(q, frag_b);
194
102
- frag_b[0] = __hfma2(*reinterpret_cast<nv_bfloat162*>(&lo),
103
- *reinterpret_cast<const nv_bfloat162*>(&MUL),
104
- *reinterpret_cast<const nv_bfloat162*>(&ADD));
105
- frag_b[1] = __hfma2(*reinterpret_cast<nv_bfloat162*>(&hi),
106
- *reinterpret_cast<const nv_bfloat162*>(&MUL),
107
- *reinterpret_cast<const nv_bfloat162*>(&ADD));
195
+ static constexpr uint32_t SUB = 0x43084308;
196
+
197
+ frag_b[0] = __hsub2(frag_b[0], *reinterpret_cast<const nv_bfloat162*>(&SUB));
198
+ frag_b[1] = __hsub2(frag_b[1], *reinterpret_cast<const nv_bfloat162*>(&SUB));
199
}
200
201
template <>
111
-__device__ inline void dequant<nv_bfloat162, vllm::kU4.id()>(
202
+__device__ inline void dequant<nv_bfloat162, vllm::kU4.id(), true>(
203
int q, nv_bfloat162* frag_b) {
113
- static constexpr uint32_t MASK = 0x000f000f;
114
- static constexpr uint32_t EX = 0x43004300;
204
+ dequant<nv_bfloat162, vllm::kU4B8.id(), true>(q, frag_b);
205
+}
206
116
- // Guarantee that the `(a & b) | c` operations are LOP3s.
117
- // clang-format off
118
- int lo = lop3<(0xf0 & 0xcc) | 0xaa>(q, MASK, EX);
119
- q >>= 4;
120
- int hi = lop3<(0xf0 & 0xcc) | 0xaa>(q, MASK, EX);
121
- // clang-format on
207
+template <>
208
+__device__ inline void dequant<nv_bfloat162, vllm::kU4.id(), false>(
209
+ int q, nv_bfloat162* frag_b) {
210
+ dequant<nv_bfloat162, vllm::kU4.id(), true>(q, frag_b);
211
123
- static constexpr uint32_t MUL = 0x3F803F80;
124
- static constexpr uint32_t ADD = 0xC300C300;
212
+ static constexpr uint32_t SUB = 0x43004300;
213
126
- frag_b[0] = __hfma2(*reinterpret_cast<nv_bfloat162*>(&lo),
127
- *reinterpret_cast<const nv_bfloat162*>(&MUL),
128
- *reinterpret_cast<const nv_bfloat162*>(&ADD));
129
- frag_b[1] = __hfma2(*reinterpret_cast<nv_bfloat162*>(&hi),
130
- *reinterpret_cast<const nv_bfloat162*>(&MUL),
131
- *reinterpret_cast<const nv_bfloat162*>(&ADD));
214
+ frag_b[0] = __hsub2(frag_b[0], *reinterpret_cast<const nv_bfloat162*>(&SUB));
215
+ frag_b[1] = __hsub2(frag_b[1], *reinterpret_cast<const nv_bfloat162*>(&SUB));
216
}
217
218
//
@@ -140,8 +224,8 @@ __device__ inline void dequant<nv_bfloat162, vllm::kU4.id()>(
224
// https://github.com/NVIDIA/FasterTransformer/blob/release/v5.3_tag/src/fastertransformer/cutlass_extensions/include/cutlass_extensions/interleaved_numeric_conversion.h#L125-L175
225
//
226
template <>
143
-__device__ inline void dequant<half2, vllm::kU8B128.id()>(int q,
144
- half2* frag_b) {
227
+__device__ inline void dequant<half2, vllm::kU8B128.id(), true>(int q,
228
+ half2* frag_b) {
229
static constexpr uint32_t mask_for_elt_01 = 0x5250;
230
static constexpr uint32_t mask_for_elt_23 = 0x5351;
231
static constexpr uint32_t start_byte_for_fp16 = 0x64646464;
@@ -149,33 +233,42 @@ __device__ inline void dequant<half2, vllm::kU8B128.id()>(int q,
233
uint32_t lo = prmt<start_byte_for_fp16, mask_for_elt_01>(q);
234
uint32_t hi = prmt<start_byte_for_fp16, mask_for_elt_23>(q);
235
152
- static constexpr uint32_t I8s_TO_F16s_MAGIC_NUM = 0x64806480;
236
+ frag_b[0] = *reinterpret_cast<half2*>(&lo);
237
+ frag_b[1] = *reinterpret_cast<half2*>(&hi);
238
+}
239
154
- frag_b[0] = __hsub2(*reinterpret_cast<half2*>(&lo),
240
+template <>
241
+__device__ inline void dequant<half2, vllm::kU8B128.id(), false>(
242
+ int q, half2* frag_b) {
243
+ dequant<half2, vllm::kU8B128.id(), true>(q, frag_b);
244
+
245
+ static constexpr uint32_t I8s_TO_F16s_MAGIC_NUM = 0x64806480;
246
+ frag_b[0] = __hsub2(frag_b[0],
247
*reinterpret_cast<const half2*>(&I8s_TO_F16s_MAGIC_NUM));
156
- frag_b[1] = __hsub2(*reinterpret_cast<half2*>(&hi),
248
+ frag_b[1] = __hsub2(frag_b[1],
249
*reinterpret_cast<const half2*>(&I8s_TO_F16s_MAGIC_NUM));
250
}
251
252
template <>
161
-__device__ inline void dequant<half2, vllm::kU8.id()>(int q, half2* frag_b) {
162
- static constexpr uint32_t mask_for_elt_01 = 0x5250;
163
- static constexpr uint32_t mask_for_elt_23 = 0x5351;
164
- static constexpr uint32_t start_byte_for_fp16 = 0x64646464;
253
+__device__ inline void dequant<half2, vllm::kU8.id(), true>(int q,
254
+ half2* frag_b) {
255
+ dequant<half2, vllm::kU8B128.id(), true>(q, frag_b);
256
+}
257
166
- uint32_t lo = prmt<start_byte_for_fp16, mask_for_elt_01>(q);
167
- uint32_t hi = prmt<start_byte_for_fp16, mask_for_elt_23>(q);
258
+template <>
259
+__device__ inline void dequant<half2, vllm::kU8.id(), false>(int q,
260
+ half2* frag_b) {
261
+ dequant<half2, vllm::kU8.id(), true>(q, frag_b);
262
263
static constexpr uint32_t I8s_TO_F16s_MAGIC_NUM = 0x64006400;
170
-
171
- frag_b[0] = __hsub2(*reinterpret_cast<half2*>(&lo),
264
+ frag_b[0] = __hsub2(frag_b[0],
265
*reinterpret_cast<const half2*>(&I8s_TO_F16s_MAGIC_NUM));
173
- frag_b[1] = __hsub2(*reinterpret_cast<half2*>(&hi),
266
+ frag_b[1] = __hsub2(frag_b[1],
267
*reinterpret_cast<const half2*>(&I8s_TO_F16s_MAGIC_NUM));
268
}
269
270
template <>
178
-__device__ inline void dequant<nv_bfloat162, vllm::kU8B128.id()>(
271
+__device__ inline void dequant<nv_bfloat162, vllm::kU8B128.id(), false>(
272
int q, nv_bfloat162* frag_b) {
273
float fp32_intermediates[4];
274
uint32_t* fp32_intermediates_casted =
@@ -200,7 +293,7 @@ __device__ inline void dequant<nv_bfloat162, vllm::kU8B128.id()>(
293
}
294
295
template <>
203
-__device__ inline void dequant<nv_bfloat162, vllm::kU8.id()>(
296
+__device__ inline void dequant<nv_bfloat162, vllm::kU8.id(), false>(
297
int q, nv_bfloat162* frag_b) {
298
float fp32_intermediates[4];
299
uint32_t* fp32_intermediates_casted =
@@ -225,22 +318,30 @@ __device__ inline void dequant<nv_bfloat162, vllm::kU8.id()>(
318
}
319
320
template <>
228
-__device__ inline void dequant<half2, vllm::kFE4M3fn.id()>(int q,
229
- half2* frag_b) {
321
+__device__ inline void dequant<half2, vllm::kFE4M3fn.id(), true>(
322
+ int q, half2* frag_b) {
323
// Constants for FP8 (E4M3) and FP16 formats
231
- constexpr int FP8_EXPONENT = 4, FP8_MANTISSA = 3, FP16_EXPONENT = 5;
324
+ constexpr int FP8_EXPONENT = 4, FP16_EXPONENT = 5;
325
constexpr int RIGHT_SHIFT = FP16_EXPONENT - FP8_EXPONENT;
233
-
234
- // Calculate MASK for extracting mantissa and exponent
235
- constexpr int MASK1 = 0x80000000;
236
- constexpr int MASK2 = MASK1 >> (FP8_EXPONENT + FP8_MANTISSA);
237
- constexpr int MASK3 = MASK2 & 0x7fffffff;
238
- constexpr int MASK = MASK3 | (MASK3 >> 16);
239
- // Final MASK value: 0x7F007F00
326
+ constexpr int MASK = 0x7F007F00;
327
328
// Extract and shift FP8 values to FP16 format
329
int Out1 = (q & 0x80008000) | ((q & MASK) >> RIGHT_SHIFT);
243
- int Out2 = ((q << 8) & 0x80008000) | (((q << 8) & MASK) >> RIGHT_SHIFT);
330
+ q <<= 8;
331
+ int Out2 = (q & 0x80008000) | ((q & MASK) >> RIGHT_SHIFT);
332
+
333
+ // Note: reverse indexing is intentional because weights are permuted
334
+ frag_b[1] = *reinterpret_cast<const half2*>(&Out1);
335
+ frag_b[0] = *reinterpret_cast<const half2*>(&Out2);
336
+}
337
+
338
+template <>
339
+__device__ inline void dequant<half2, vllm::kFE4M3fn.id(), false>(
340
+ int q, half2* frag_b) {
341
+ dequant<half2, vllm::kFE4M3fn.id(), true>(q, frag_b);
342
+
343
+ // Constants for FP8 (E4M3) and FP16 formats
344
+ constexpr int FP8_EXPONENT = 4, FP16_EXPONENT = 5;
345
346
// Construct and apply exponent bias
347
constexpr int BIAS_OFFSET =
@@ -248,28 +349,36 @@ __device__ inline void dequant<half2, vllm::kFE4M3fn.id()>(int q,
349
const half2 bias_reg = __float2half2_rn(float(1 << BIAS_OFFSET));
350
351
// Convert to half2 and apply bias
251
- // Note: reverse indexing is intentional because weights are permuted
252
- frag_b[1] = __hmul2(*reinterpret_cast<const half2*>(&Out1), bias_reg);
253
- frag_b[0] = __hmul2(*reinterpret_cast<const half2*>(&Out2), bias_reg);
352
+ frag_b[1] = __hmul2(frag_b[1], bias_reg);
353
+ frag_b[0] = __hmul2(frag_b[0], bias_reg);
354
}
355
356
template <>
257
-__device__ inline void dequant<nv_bfloat162, vllm::kFE4M3fn.id()>(
357
+__device__ inline void dequant<nv_bfloat162, vllm::kFE4M3fn.id(), true>(
358
int q, nv_bfloat162* frag_b) {
359
// Constants for FP8 (E4M3) and BF16 formats
260
- constexpr int FP8_EXPONENT = 4, FP8_MANTISSA = 3, BF16_EXPONENT = 8;
360
+ constexpr int FP8_EXPONENT = 4, BF16_EXPONENT = 8;
361
constexpr int RIGHT_SHIFT = BF16_EXPONENT - FP8_EXPONENT;
362
263
- // Calculate MASK for extracting mantissa and exponent
264
- constexpr int MASK1 = 0x80000000;
265
- constexpr int MASK2 = MASK1 >> (FP8_EXPONENT + FP8_MANTISSA);
266
- constexpr int MASK3 = MASK2 & 0x7fffffff;
267
- constexpr int MASK = MASK3 | (MASK3 >> 16);
268
- // Final MASK value: 0x7F007F00
363
+ constexpr int MASK = 0x7F007F00;
364
365
// Extract and shift FP8 values to BF16 format
366
int Out1 = (q & 0x80008000) | ((q & MASK) >> RIGHT_SHIFT);
272
- int Out2 = ((q << 8) & 0x80008000) | (((q << 8) & MASK) >> RIGHT_SHIFT);
367
+ q <<= 8;
368
+ int Out2 = (q & 0x80008000) | ((q & MASK) >> RIGHT_SHIFT);
369
+
370
+ // Note: reverse indexing is intentional because weights are permuted
371
+ frag_b[1] = *reinterpret_cast<const nv_bfloat162*>(&Out1);
372
+ frag_b[0] = *reinterpret_cast<const nv_bfloat162*>(&Out2);
373
+}
374
+
375
+template <>
376
+__device__ inline void dequant<nv_bfloat162, vllm::kFE4M3fn.id(), false>(
377
+ int q, nv_bfloat162* frag_b) {
378
+ dequant<nv_bfloat162, vllm::kFE4M3fn.id(), true>(q, frag_b);
379
+
380
+ // Constants for FP8 (E4M3) and BF16 formats
381
+ constexpr int FP8_EXPONENT = 4, BF16_EXPONENT = 8;
382
383
// Construct and apply exponent bias
384
constexpr int BIAS_OFFSET =
@@ -281,9 +390,116 @@ __device__ inline void dequant<nv_bfloat162, vllm::kFE4M3fn.id()>(
390
__float2bfloat162_rn(*reinterpret_cast<const float*>(&BIAS));
391
392
// Convert to bfloat162 and apply bias
393
+ frag_b[1] = __hmul2(frag_b[1], bias_reg);
394
+ frag_b[0] = __hmul2(frag_b[0], bias_reg);
395
+}
396
+
397
+template <>
398
+__device__ inline void dequant<half2, vllm::kFE2M1f.id(), true>(int q,
399
+ half2* frag_b) {
400
+ // Constants for FP4 (E2M1) and FP16 formats
401
+ constexpr int FP4_EXPONENT = 2, FP16_EXPONENT = 5;
402
+ constexpr int RIGHT_SHIFT = FP16_EXPONENT - FP4_EXPONENT;
403
+ constexpr int MASK = 0x70007000;
404
+
405
+ // Extract and shift FP4 values to FP16 format
406
+ int Out1 = (q & 0x80008000) | ((q & MASK) >> RIGHT_SHIFT);
407
+ q <<= 4;
408
+ int Out2 = (q & 0x80008000) | ((q & MASK) >> RIGHT_SHIFT);
409
+
410
+ // Note: reverse indexing is intentional because weights are permuted
411
+ frag_b[1] = *reinterpret_cast<const half2*>(&Out1);
412
+ frag_b[0] = *reinterpret_cast<const half2*>(&Out2);
413
+}
414
+
415
+template <>
416
+__device__ inline void dequant<half2, vllm::kFE2M1f.id(), false>(
417
+ int q, half2* frag_b) {
418
+ dequant<half2, vllm::kFE2M1f.id(), true>(q, frag_b);
419
+
420
+ // Constants for FP4 (E2M1) and FP16 formats
421
+ constexpr int FP4_EXPONENT = 2, FP16_EXPONENT = 5;
422
+
423
+ // Construct and apply exponent bias
424
+ constexpr int BIAS_OFFSET =
425
+ (1 << (FP16_EXPONENT - 1)) - (1 << (FP4_EXPONENT - 1));
426
+ const half2 bias_reg = __float2half2_rn(float(1 << BIAS_OFFSET));
427
+
428
+ // Convert to half2 and apply bias
429
+ frag_b[1] = __hmul2(frag_b[1], bias_reg);
430
+ frag_b[0] = __hmul2(frag_b[0], bias_reg);
431
+}
432
+
433
+template <>
434
+__device__ inline void dequant<nv_bfloat162, vllm::kFE2M1f.id(), true>(
435
+ int q, nv_bfloat162* frag_b) {
436
+ // Constants for FP4 (E2M1) and FP16 formats
437
+ constexpr int FP4_EXPONENT = 2, BF16_EXPONENT = 8;
438
+ constexpr int RIGHT_SHIFT = BF16_EXPONENT - FP4_EXPONENT;
439
+ constexpr int MASK = 0x70007000;
440
+
441
+ // Extract and shift FP4 values to FP16 format
442
+ int Out1 = (q & 0x80008000) | ((q & MASK) >> RIGHT_SHIFT);
443
+ q <<= 4;
444
+ int Out2 = (q & 0x80008000) | ((q & MASK) >> RIGHT_SHIFT);
445
+
446
+ // Note: reverse indexing is intentional because weights are permuted
447
+ frag_b[1] = *reinterpret_cast<const nv_bfloat162*>(&Out1);
448
+ frag_b[0] = *reinterpret_cast<const nv_bfloat162*>(&Out2);
449
+}
450
+
451
+template <>
452
+__device__ inline void dequant<nv_bfloat162, vllm::kFE2M1f.id(), false>(
453
+ int q, nv_bfloat162* frag_b) {
454
+ dequant<nv_bfloat162, vllm::kFE2M1f.id(), true>(q, frag_b);
455
+
456
+ // Constants for FP4 (E2M1) and BF16 formats
457
+ constexpr int FP4_EXPONENT = 2, BF16_EXPONENT = 8;
458
+
459
+ // Construct and apply exponent bias
460
+ constexpr int BIAS_OFFSET =
461
+ (1 << (BF16_EXPONENT - 1)) - (1 << (FP4_EXPONENT - 1));
462
+ // Add 127 (float exponent bias) to BIAS_OFFSET and shift to float exponent
463
+ // position
464
+ constexpr uint32_t BIAS = (BIAS_OFFSET + 127) << 23;
465
+ const nv_bfloat162 bias_reg =
466
+ __float2bfloat162_rn(*reinterpret_cast<const float*>(&BIAS));
467
+
468
+ // Convert to half2 and apply bias
469
+ frag_b[1] = __hmul2(frag_b[1], bias_reg);
470
+ frag_b[0] = __hmul2(frag_b[0], bias_reg);
471
+}
472
+
473
+template <typename scalar_t2>
474
+__device__ inline void dequant_fp8_scales(int q, scalar_t2* frag_b);
475
+
476
+template <>
477
+__device__ inline void dequant_fp8_scales<half2>(int q, half2* frag_b) {
478
+ int Out1 = (q & 0xFF00FF00) >> 1;
479
+ ;
480
+ q <<= 8;
481
+ int Out2 = (q & 0xFF00FF00) >> 1;
482
+
483
+ // Note: reverse indexing is intentional because weights are permuted
484
+ frag_b[1] = *reinterpret_cast<const half2*>(&Out1);
485
+ frag_b[0] = *reinterpret_cast<const half2*>(&Out2);
486
+};
487
+
488
+template <>
489
+__device__ inline void dequant_fp8_scales<nv_bfloat162>(int q,
490
+ nv_bfloat162* frag_b) {
491
+ constexpr int FP8_EXPONENT = 4, BF16_EXPONENT = 8;
492
+ constexpr int RIGHT_SHIFT = BF16_EXPONENT - FP8_EXPONENT;
493
+ constexpr int MASK = 0x7F007F00;
494
+
495
+ // Extract and shift FP8 values to BF16 format
496
+ int Out1 = ((q & 0x80008000) >> 1) | ((q & MASK) >> RIGHT_SHIFT);
497
+ q <<= 8;
498
+ int Out2 = ((q & 0x80008000) >> 1) | ((q & MASK) >> RIGHT_SHIFT);
499
+
500
// Note: reverse indexing is intentional because weights are permuted
285
- frag_b[1] = __hmul2(*reinterpret_cast<const nv_bfloat162*>(&Out1), bias_reg);
286
- frag_b[0] = __hmul2(*reinterpret_cast<const nv_bfloat162*>(&Out2), bias_reg);
501
+ frag_b[1] = *reinterpret_cast<const nv_bfloat162*>(&Out1);
502
+ frag_b[0] = *reinterpret_cast<const nv_bfloat162*>(&Out2);
503
}
504
505
#endif
csrc/quantization/gptq_marlin/generate_kernels.py
+11
-2
@@ -31,7 +31,10 @@ TEMPLATE = ("template __global__ void Marlin<"
31
32
# int8 with zero point case (vllm::kU8) is also supported,
33
# we don't add it to reduce wheel size.
34
-SCALAR_TYPES = ["vllm::kU4", "vllm::kU4B8", "vllm::kU8B128", "vllm::kFE4M3fn"]
34
+SCALAR_TYPES = [
35
+ "vllm::kU4", "vllm::kU4B8", "vllm::kU8B128", "vllm::kFE4M3fn",
36
+ "vllm::kFE2M1f"
37
+]
38
THREAD_CONFIGS = [(128, 128, 256), (64, 256, 256), (64, 128, 128),
39
(128, 64, 128)]
40
@@ -40,7 +43,7 @@ THREAD_M_BLOCKS = [0.5, 1, 2, 3, 4]
43
# = 0 : act order case
44
# = -1 : channelwise quantization
45
# > 0 : group_size=16*group_blocks
43
-GROUP_BLOCKS = [0, -1, 2, 4, 8]
46
+GROUP_BLOCKS = [0, 1, -1, 2, 4, 8]
47
DTYPES = ["fp16", "bf16"]
48
49
@@ -73,6 +76,12 @@ def generate_new_kernels():
76
# for fp8
77
if scalar_type == "vllm::kFE4M3fn" and group_blocks not in [-1, 8]:
78
continue
79
+ # nvfp4 only supports group_size == 16
80
+ if scalar_type == "vllm::kFE2M1f" and group_blocks != 1:
81
+ continue
82
+ # other quantization methods don't support group_size = 16
83
+ if scalar_type != "vllm::kFE2M1f" and group_blocks == 1:
84
+ continue
85
86
k_blocks = thread_configs[0] // 16
87
n_blocks = thread_configs[1] // 16
csrc/quantization/gptq_marlin/gptq_marlin.cu
+65
-16
@@ -258,6 +258,7 @@ bool is_valid_config(thread_config_t const& th_config, int thread_m_blocks,
258
// BIGGROUP: cases for big group size (group_blocks in [-1, 8])
259
// FZP: cases for float-zero-point (is_zp_float = true)
260
// ACT: cases for act order case (group_blocks == 0)
261
+ // FP4: cases for nvfp4(e2m1) (group_blocks == 1)
262
#define COMMON_GET_IF_M1(W_TYPE, N_BLOCKS, K_BLOCKS, NUM_THREADS) \
263
_GET_IF(W_TYPE, 1, N_BLOCKS, K_BLOCKS, true, -1, NUM_THREADS, false) \
264
_GET_IF(W_TYPE, 1, N_BLOCKS, K_BLOCKS, true, 2, NUM_THREADS, false) \
@@ -314,6 +315,23 @@ bool is_valid_config(thread_config_t const& th_config, int thread_m_blocks,
315
BIGGROUP_GET_IF_M234(W_TYPE, 8, 4, 128) \
316
BIGGROUP_GET_IF_M234(W_TYPE, 4, 8, 128)
317
318
+ #define FP4_GET_IF_M1(W_TYPE, N_BLOCKS, K_BLOCKS, NUM_THREADS) \
319
+ _GET_IF(W_TYPE, 1, N_BLOCKS, K_BLOCKS, true, 1, NUM_THREADS, false) \
320
+ _GET_IF(W_TYPE, 1, N_BLOCKS, K_BLOCKS, false, 1, NUM_THREADS, false)
321
+
322
+ #define FP4_GET_IF_M234(W_TYPE, N_BLOCKS, K_BLOCKS, NUM_THREADS) \
323
+ _GET_IF(W_TYPE, 2, N_BLOCKS, K_BLOCKS, false, 1, NUM_THREADS, false) \
324
+ _GET_IF(W_TYPE, 3, N_BLOCKS, K_BLOCKS, false, 1, NUM_THREADS, false) \
325
+ _GET_IF(W_TYPE, 4, N_BLOCKS, K_BLOCKS, false, 1, NUM_THREADS, false)
326
+
327
+ #define FP4_GET_IF(W_TYPE) \
328
+ FP4_GET_IF_M1(W_TYPE, 8, 8, 256) \
329
+ FP4_GET_IF_M1(W_TYPE, 8, 4, 128) \
330
+ FP4_GET_IF_M1(W_TYPE, 4, 8, 128) \
331
+ FP4_GET_IF_M234(W_TYPE, 16, 4, 256) \
332
+ FP4_GET_IF_M234(W_TYPE, 8, 4, 128) \
333
+ FP4_GET_IF_M234(W_TYPE, 4, 8, 128)
334
+
335
// We currently have 4-bit models only with group_blocks == 4
336
#define FZP_GET_IF_M1(W_TYPE, N_BLOCKS, K_BLOCKS, NUM_THREADS) \
337
_GET_IF(W_TYPE, 1, N_BLOCKS, K_BLOCKS, true, 4, NUM_THREADS, true) \
@@ -366,6 +384,8 @@ MarlinFuncPtr get_marlin_kernel(const vllm::ScalarType q_type,
384
COMMON_GET_IF(vllm::kU4B8)
385
COMMON_GET_IF(vllm::kU8B128)
386
387
+ FP4_GET_IF(vllm::kFE2M1f)
388
+
389
BIGGROUP_GET_IF(vllm::kFE4M3fn)
390
391
ACT_GET_IF(vllm::kU4B8)
@@ -434,8 +454,8 @@ exec_config_t determine_exec_config(const vllm::ScalarType& q_type, int prob_m,
454
455
template <typename scalar_t>
456
void marlin_mm(const void* A, const void* B, void* C, void* C_tmp, void* s,
437
- void* zp, void* g_idx, void* perm, void* a_tmp, int prob_m,
438
- int prob_n, int prob_k, int lda, void* workspace,
457
+ void* s2, void* zp, void* g_idx, void* perm, void* a_tmp,
458
+ int prob_m, int prob_n, int prob_k, int lda, void* workspace,
459
vllm::ScalarType const& q_type, bool has_act_order,
460
bool is_k_full, bool has_zp, int num_groups, int group_size,
461
int dev, cudaStream_t stream, int thread_k_init,
@@ -446,11 +466,12 @@ void marlin_mm(const void* A, const void* B, void* C, void* C_tmp, void* s,
466
q_type == vllm::kU4 || q_type == vllm::kU8,
467
"q_type must be u4 or u8 when has_zp = True. Got = ", q_type.str());
468
} else {
449
- TORCH_CHECK(q_type == vllm::kU4B8 || q_type == vllm::kU8B128 ||
450
- q_type == vllm::kFE4M3fn,
451
- "q_type must be uint4b8, uint8b128 or float8_e4m3fn when "
452
- "has_zp = False. Got = ",
453
- q_type.str());
469
+ TORCH_CHECK(
470
+ q_type == vllm::kU4B8 || q_type == vllm::kU8B128 ||
471
+ q_type == vllm::kFE4M3fn || q_type == vllm::kFE2M1f,
472
+ "q_type must be uint4b8, uint8b128, float8_e4m3fn or float4_e2m1f when "
473
+ "has_zp = False. Got = ",
474
+ q_type.str());
475
}
476
477
TORCH_CHECK(prob_m > 0 && prob_n > 0 && prob_k > 0, "Invalid MNK = [", prob_m,
@@ -483,6 +504,7 @@ void marlin_mm(const void* A, const void* B, void* C, void* C_tmp, void* s,
504
int4* C_ptr = (int4*)C;
505
int4* C_tmp_ptr = (int4*)C_tmp;
506
const int4* s_ptr = (const int4*)s;
507
+ const uint16_t* s2_ptr = (const uint16_t*)s2;
508
const int4* zp_ptr = (const int4*)zp;
509
const int* g_idx_ptr = (const int*)g_idx;
510
const int* perm_ptr = (const int*)perm;
@@ -601,7 +623,7 @@ void marlin_mm(const void* A, const void* B, void* C, void* C_tmp, void* s,
623
// avoid ">>>" being formatted to "> > >"
624
// clang-format off
625
kernel<<<blocks, num_threads, max_shared_mem_new, stream>>>(
604
- A_ptr, B_ptr, C_ptr, C_tmp_ptr, s_ptr, zp_ptr, g_idx_ptr, num_groups,
626
+ A_ptr, B_ptr, C_ptr, C_tmp_ptr, s_ptr, s2_ptr, zp_ptr, g_idx_ptr, num_groups,
627
prob_m_split, prob_n, prob_k, lda, locks, part_use_atomic_add,
628
use_fp32_reduce, max_shared_mem_new);
629
// clang-format on
@@ -617,6 +639,7 @@ void marlin_mm(const void* A, const void* B, void* C, void* C_tmp, void* s,
639
torch::Tensor gptq_marlin_gemm(
640
torch::Tensor& a, std::optional<torch::Tensor> c_or_none,
641
torch::Tensor& b_q_weight, torch::Tensor& b_scales,
642
+ std::optional<torch::Tensor> const& global_scale_or_none,
643
std::optional<torch::Tensor> const& b_zeros_or_none,
644
std::optional<torch::Tensor> const& g_idx_or_none,
645
std::optional<torch::Tensor> const& perm_or_none, torch::Tensor& workspace,
@@ -759,6 +782,17 @@ torch::Tensor gptq_marlin_gemm(
782
}
783
}
784
785
+ torch::Tensor global_scale;
786
+ if (global_scale_or_none.has_value()) {
787
+ global_scale = global_scale_or_none.value();
788
+ TORCH_CHECK(b_q_type == vllm::kFE2M1f,
789
+ "global_scale can only be used for float4_e2m1f.");
790
+ } else {
791
+ global_scale = torch::empty({0}, options);
792
+ TORCH_CHECK(!(b_q_type == vllm::kFE2M1f),
793
+ "the global_scale parameter must be passed for float4_e2m1f.");
794
+ }
795
+
796
torch::Tensor b_zeros;
797
if (b_zeros_or_none.has_value()) {
798
b_zeros = b_zeros_or_none.value();
@@ -774,8 +808,9 @@ torch::Tensor gptq_marlin_gemm(
808
"b_q_type must be u4 or u8 when has_zp = True. Got = ", b_q_type.str());
809
} else {
810
TORCH_CHECK(b_q_type == vllm::kU4B8 || b_q_type == vllm::kU8B128 ||
777
- b_q_type == vllm::kFE4M3fn,
778
- "b_q_type must be uint4b8, uint8b128 or float8_e4m3fn when "
811
+ b_q_type == vllm::kFE4M3fn || b_q_type == vllm::kFE2M1f,
812
+ "b_q_type must be uint4b8, uint8b128, float8_e4m3fn or "
813
+ "float4_e2m1f when "
814
"has_zp = False. Got = ",
815
b_q_type.str());
816
}
@@ -820,22 +855,36 @@ torch::Tensor gptq_marlin_gemm(
855
856
int dev = a.get_device();
857
if (a.scalar_type() == at::ScalarType::Half) {
858
+ void* scales_ptr;
859
+ if (b_q_type == vllm::kFE2M1f) {
860
+ scales_ptr = b_scales.data_ptr<at::Float8_e4m3fn>();
861
+ } else {
862
+ scales_ptr = b_scales.data_ptr<at::Half>();
863
+ }
864
+
865
marlin::marlin_mm<half>(
866
a.data_ptr<at::Half>(), b_q_weight.data_ptr(), c.data_ptr<at::Half>(),
825
- c_tmp.data_ptr<float>(), b_scales.data_ptr<at::Half>(),
867
+ c_tmp.data_ptr<float>(), scales_ptr, global_scale.data_ptr<at::Half>(),
868
b_zeros.data_ptr(), g_idx.data_ptr(), perm.data_ptr(),
869
a_tmp.data_ptr<at::Half>(), size_m, size_n, size_k, a.stride(0),
870
workspace.data_ptr(), b_q_type, has_act_order, is_k_full, has_zp,
871
num_groups, group_size, dev, at::cuda::getCurrentCUDAStream(dev),
872
thread_k, thread_n, sms, use_atomic_add, use_fp32_reduce, is_zp_float);
873
} else if (a.scalar_type() == at::ScalarType::BFloat16) {
874
+ void* scales_ptr;
875
+ if (b_q_type == vllm::kFE2M1f) {
876
+ scales_ptr = b_scales.data_ptr<at::Float8_e4m3fn>();
877
+ } else {
878
+ scales_ptr = b_scales.data_ptr<at::BFloat16>();
879
+ }
880
+
881
marlin::marlin_mm<nv_bfloat16>(
882
a.data_ptr<at::BFloat16>(), b_q_weight.data_ptr(),
834
- c.data_ptr<at::BFloat16>(), c_tmp.data_ptr<float>(),
835
- b_scales.data_ptr<at::BFloat16>(), b_zeros.data_ptr(), g_idx.data_ptr(),
836
- perm.data_ptr(), a_tmp.data_ptr<at::BFloat16>(), size_m, size_n, size_k,
837
- a.stride(0), workspace.data_ptr(), b_q_type, has_act_order, is_k_full,
838
- has_zp, num_groups, group_size, dev,
883
+ c.data_ptr<at::BFloat16>(), c_tmp.data_ptr<float>(), scales_ptr,
884
+ global_scale.data_ptr<at::BFloat16>(), b_zeros.data_ptr(),
885
+ g_idx.data_ptr(), perm.data_ptr(), a_tmp.data_ptr<at::BFloat16>(),
886
+ size_m, size_n, size_k, a.stride(0), workspace.data_ptr(), b_q_type,
887
+ has_act_order, is_k_full, has_zp, num_groups, group_size, dev,
888
at::cuda::getCurrentCUDAStream(dev), thread_k, thread_n, sms,
889
use_atomic_add, use_fp32_reduce, is_zp_float);
890
} else {
csrc/quantization/gptq_marlin/kernel.h
+8
-7
@@ -7,13 +7,14 @@
7
#include "marlin_dtypes.cuh"
8
#include "core/scalar_type.hpp"
9
10
-#define MARLIN_KERNEL_PARAMS \
11
- const int4 *__restrict__ A, const int4 *__restrict__ B, \
12
- int4 *__restrict__ C, int4 *__restrict__ C_tmp, \
13
- const int4 *__restrict__ scales_ptr, const int4 *__restrict__ zp_ptr, \
14
- const int *__restrict__ g_idx, int num_groups, int prob_m, int prob_n, \
15
- int prob_k, int lda, int *locks, bool use_atomic_add, \
16
- bool use_fp32_reduce, int max_shared_mem
10
+#define MARLIN_KERNEL_PARAMS \
11
+ const int4 *__restrict__ A, const int4 *__restrict__ B, \
12
+ int4 *__restrict__ C, int4 *__restrict__ C_tmp, \
13
+ const int4 *__restrict__ scales_ptr, \
14
+ const uint16_t *__restrict__ scale2_ptr, \
15
+ const int4 *__restrict__ zp_ptr, const int *__restrict__ g_idx, \
16
+ int num_groups, int prob_m, int prob_n, int prob_k, int lda, int *locks, \
17
+ bool use_atomic_add, bool use_fp32_reduce, int max_shared_mem
18
19
namespace MARLIN_NAMESPACE_NAME {
20
template <typename scalar_t, // compute dtype, half or nv_float16
csrc/quantization/gptq_marlin/marlin_template.h
+85
-36
@@ -292,9 +292,11 @@ __global__ void Marlin(
292
int4* __restrict__ C_tmp, // fp32 tmp output buffer (for reduce)
293
const int4* __restrict__ scales_ptr, // fp16 quantization scales of shape
294
// (k/groupsize)xn
295
- const int4* __restrict__ zp_ptr, // 4bit packed zero-points of shape
296
- // (k/groupsize)x(n/pack_factor)
297
- const int* __restrict__ g_idx, // int32 group indices of shape k
295
+ const uint16_t* __restrict__ scale2_ptr, // fp16 global scale (for nvfp4
296
+ // only)
297
+ const int4* __restrict__ zp_ptr, // 4bit packed zero-points of shape
298
+ // (k/groupsize)x(n/pack_factor)
299
+ const int* __restrict__ g_idx, // int32 group indices of shape k
300
int num_groups, // number of scale groups per output channel
301
int prob_m, // batch dimension m
302
int prob_n, // output dimension n
@@ -325,6 +327,21 @@ __global__ void Marlin(
327
328
static constexpr auto w_type = vllm::ScalarType::from_id(w_type_id);
329
constexpr bool has_zp = w_type == vllm::kU4 || w_type == vllm::kU8;
330
+ constexpr bool is_int_type = w_type == vllm::kU4 || w_type == vllm::kU8 ||
331
+ w_type == vllm::kU4B8 || w_type == vllm::kU8B128;
332
+ // see comments of dequant.h for more details
333
+ constexpr bool dequant_skip_flop =
334
+ !is_int_type ||
335
+ has_zp && !is_zp_float && !std::is_same<scalar_t, nv_bfloat16>::value ||
336
+ has_zp && !is_zp_float && !(w_type == vllm::kU8);
337
+
338
+ scalar_t2 global_scale;
339
+
340
+ if constexpr (w_type == vllm::kFE2M1f) {
341
+ uint16_t val = scale2_ptr[0];
342
+ global_scale = Dtype::num2num2(*reinterpret_cast<scalar_t*>(&val));
343
+ }
344
+
345
constexpr bool has_act_order = group_blocks == 0;
346
constexpr int m_block_size = m_block_size_8 ? 8 : (16 * thread_m_blocks);
347
@@ -481,7 +498,7 @@ __global__ void Marlin(
498
constexpr int s_sh_stride = 16 * thread_n_blocks / 8;
499
constexpr int s_tb_groups =
500
!has_act_order && group_blocks != -1 && group_blocks < thread_k_blocks
484
- ? thread_k_blocks / group_blocks
501
+ ? thread_k_blocks / group_blocks / (w_type == vllm::kFE2M1f ? 2 : 1)
502
: 1;
503
constexpr int s_sh_stage = s_tb_groups * s_sh_stride;
504
int s_gl_rd_delta = s_gl_stride;
@@ -540,7 +557,8 @@ __global__ void Marlin(
557
if constexpr (group_blocks == -1) {
558
s_gl_rd = s_sh_stride * slice_col + threadIdx.x;
559
} else {
543
- s_gl_rd = s_gl_stride * ((thread_k_blocks * slice_row) / group_blocks) +
560
+ s_gl_rd = s_gl_stride * ((thread_k_blocks * slice_row) / group_blocks) /
561
+ (w_type == vllm::kFE2M1f ? 2 : 1) +
562
s_sh_stride * slice_col + threadIdx.x;
563
}
564
}
@@ -564,10 +582,20 @@ __global__ void Marlin(
582
// we scale a `half2` tile in column-major layout in the former and in
583
// row-major in the latter case.
584
int s_sh_rd;
567
- if constexpr (group_blocks != -1)
585
+ if constexpr (group_blocks != -1 && w_type == vllm::kFE2M1f) {
586
+ auto warp_id = threadIdx.x / 32;
587
+ int n_warps = thread_n_blocks / 4;
588
+ int warp_row = warp_id / n_warps;
589
+
590
s_sh_rd = 8 * ((threadIdx.x / 32) % (thread_n_blocks / 4)) +
591
(threadIdx.x % 32) / 4;
570
- else if constexpr (group_blocks == -1 && (m_block_size_8 || has_zp))
592
+ s_sh_rd = s_sh_rd * 2 + warp_row % 2;
593
+
594
+ } else if constexpr (group_blocks != -1)
595
+ s_sh_rd = 8 * ((threadIdx.x / 32) % (thread_n_blocks / 4)) +
596
+ (threadIdx.x % 32) / 4;
597
+ else if constexpr (group_blocks == -1 &&
598
+ (m_block_size_8 || (has_zp && !dequant_skip_flop)))
599
s_sh_rd = 8 * ((threadIdx.x / 32) % (thread_n_blocks / 4)) +
600
(threadIdx.x % 32) / 8;
601
else
@@ -681,7 +709,7 @@ __global__ void Marlin(
709
sh_first_group_id = first_group_id;
710
sh_num_groups = last_group_id - first_group_id + 1;
711
684
- if (sh_num_groups < act_s_max_num_groups) {
712
+ if (sh_num_groups > act_s_max_num_groups) {
713
sh_num_groups = act_s_max_num_groups;
714
}
715
@@ -887,12 +915,19 @@ __global__ void Marlin(
915
cur_k += k_iter_size * (k % b_sh_wr_iters);
916
917
int k_blocks = cur_k / 16;
890
- int cur_group_id = k_blocks / group_blocks;
918
+ int cur_group_id =
919
+ k_blocks / (group_blocks * (w_type == vllm::kFE2M1f ? 2 : 1));
920
921
int4* sh_s_stage = sh_s + s_sh_stage * pipe;
922
894
- reinterpret_cast<int4*>(&frag_s[k % 2])[0] =
895
- sh_s_stage[s_sh_rd + cur_group_id * s_sh_stride];
923
+ if constexpr (w_type_id != vllm::kFE2M1f.id()) {
924
+ reinterpret_cast<int4*>(&frag_s[k % 2])[0] =
925
+ sh_s_stage[s_sh_rd + cur_group_id * s_sh_stride];
926
+ } else {
927
+ reinterpret_cast<int2*>(&frag_s[k % 2])[0] =
928
+ reinterpret_cast<int2*>(
929
+ sh_s_stage)[s_sh_rd + cur_group_id * (2 * s_sh_stride)];
930
+ }
931
}
932
}
933
@@ -1065,22 +1100,7 @@ __global__ void Marlin(
1100
};
1101
1102
auto dequant_data = [&](int q, scalar_t2* frag_b_ptr) {
1068
- if constexpr (has_zp && is_zp_float || !has_zp) {
1069
- dequant<scalar_t2, w_type_id>(q, frag_b_ptr);
1070
- } else {
1071
- static_assert(has_zp && !is_zp_float);
1072
- static_assert(w_type_id == vllm::kU4.id() || w_type_id == vllm::kU8.id());
1073
- // If (has_zp && !is_zp_float),
1074
- // we use not-zp version `dequant` function
1075
- // to improve numerical accuracy.
1076
- // Since both weight and zero point are dequanted using this logic,
1077
- // the final dequanted weight would be correct.
1078
- if constexpr (w_type_id == vllm::kU4.id()) {
1079
- dequant<scalar_t2, vllm::kU4B8.id()>(q, frag_b_ptr);
1080
- } else if constexpr (w_type_id == vllm::kU8.id()) {
1081
- dequant<scalar_t2, vllm::kU8B128.id()>(q, frag_b_ptr);
1082
- }
1083
- }
1103
+ dequant<scalar_t2, w_type_id, dequant_skip_flop>(q, frag_b_ptr);
1104
};
1105
1106
// Execute the actual tensor core matmul of a sub-tile.
@@ -1110,13 +1130,23 @@ __global__ void Marlin(
1130
dequant_data(zp_quant_1, reinterpret_cast<scalar_t2*>(&frag_zp) + 2);
1131
}
1132
}
1113
- if constexpr (has_zp && is_zp_float) {
1133
+ if constexpr (!dequant_skip_flop && has_zp && is_zp_float) {
1134
if (is_new_zp) {
1135
reinterpret_cast<int4*>(&frag_zp)[0] =
1136
reinterpret_cast<int4*>(&frag_zpf[k2])[0];
1137
}
1138
}
1139
1140
+ if constexpr (w_type == vllm::kFE2M1f) {
1141
+ int s_quant_0 = reinterpret_cast<int*>(frag_s[k2])[0];
1142
+ int s_quant_1 = reinterpret_cast<int*>(frag_s[k2])[1];
1143
+
1144
+ dequant_fp8_scales<scalar_t2>(s_quant_0,
1145
+ reinterpret_cast<scalar_t2*>(&frag_s[k2]));
1146
+ dequant_fp8_scales<scalar_t2>(
1147
+ s_quant_1, reinterpret_cast<scalar_t2*>(&frag_s[k2]) + 2);
1148
+ }
1149
+
1150
// We have the m dimension as the inner loop in order to encourage overlapping
1151
// dequantization and matmul operations.
1152
#pragma unroll
@@ -1125,7 +1155,10 @@ __global__ void Marlin(
1155
FragB frag_b1;
1156
int b_quant_0, b_quant_1;
1157
1128
- if constexpr (w_type.size_bits() == 4) {
1158
+ if constexpr (w_type_id == vllm::kFE2M1f.id()) {
1159
+ b_quant_1 = frag_b_quant[k2][0][j];
1160
+ b_quant_0 = b_quant_1 << 8;
1161
+ } else if constexpr (w_type.size_bits() == 4) {
1162
b_quant_0 = frag_b_quant[k2][0][j];
1163
b_quant_1 = b_quant_0 >> 8;
1164
} else {
@@ -1138,6 +1171,11 @@ __global__ void Marlin(
1171
dequant_data(b_quant_0, reinterpret_cast<scalar_t2*>(&frag_b0));
1172
dequant_data(b_quant_1, reinterpret_cast<scalar_t2*>(&frag_b1));
1173
1174
+ if constexpr (dequant_skip_flop && has_zp && !is_zp_float) {
1175
+ sub_zp<scalar_t>(frag_b0, frag_zp[j], 0);
1176
+ sub_zp<scalar_t>(frag_b1, frag_zp[j], 1);
1177
+ }
1178
+
1179
// Apply scale to frag_b0
1180
if constexpr (has_act_order) {
1181
static_assert(group_blocks != -1);
@@ -1145,7 +1183,8 @@ __global__ void Marlin(
1183
act_frag_s[k2][2][j], act_frag_s[k2][3][j], 0);
1184
scale4<scalar_t>(frag_b1, act_frag_s[k2][0][j], act_frag_s[k2][1][j],
1185
act_frag_s[k2][2][j], act_frag_s[k2][3][j], 1);
1148
- } else if constexpr (has_zp && !is_zp_float && group_blocks == -1) {
1186
+ } else if constexpr (!dequant_skip_flop && has_zp && !is_zp_float &&
1187
+ group_blocks == -1) {
1188
int idx = (threadIdx.x / 4) % 2;
1189
scalar_t2 s2 = Dtype::nums2num2(
1190
reinterpret_cast<scalar_t*>(&frag_s[j / 2][j % 2 * 2 + 0])[idx],
@@ -1153,7 +1192,7 @@ __global__ void Marlin(
1192
if (is_new_zp) frag_zp[j] = __hmul2(frag_zp[j], s2);
1193
scale_and_sub<scalar_t>(frag_b0, s2.x, frag_zp[j].x);
1194
scale_and_sub<scalar_t>(frag_b1, s2.y, frag_zp[j].y);
1156
- } else if constexpr (has_zp && group_blocks != -1) {
1195
+ } else if constexpr (!dequant_skip_flop && has_zp && group_blocks != -1) {
1196
if (is_new_zp)
1197
frag_zp[j] = __hmul2(frag_zp[j],
1198
*reinterpret_cast<scalar_t2*>(&frag_s[k2][j]));
@@ -1408,10 +1447,15 @@ __global__ void Marlin(
1447
// For per-column quantization we finally apply the scale here (only for
1448
// 4-bit)
1449
if constexpr (!has_act_order && group_blocks == -1 &&
1411
- w_type.size_bits() == 4 && !has_zp) {
1450
+ w_type.size_bits() == 4 &&
1451
+ (has_zp && dequant_skip_flop || !has_zp)) {
1452
res = __hmul2(res, s[0]);
1453
}
1454
1455
+ if constexpr (w_type == vllm::kFE2M1f) {
1456
+ res = __hmul2(res, global_scale);
1457
+ }
1458
+
1459
if constexpr (m_block_size_8) {
1460
((scalar_t*)sh_red)[idx] = res.x;
1461
((scalar_t*)sh_red)[idx + 8 * c_sh_stride] = res.y;
@@ -1488,7 +1532,9 @@ __global__ void Marlin(
1532
if constexpr (has_zp && !is_zp_float && group_blocks == -1) {
1533
if (i == 0) {
1534
fetch_col_zp_to_shared();
1491
- fetch_col_scale_to_shared();
1535
+ if constexpr (!dequant_skip_flop) {
1536
+ fetch_col_scale_to_shared();
1537
+ }
1538
}
1539
}
1540
fetch_to_shared(i, i, i < slice_iters);
@@ -1563,7 +1609,8 @@ __global__ void Marlin(
1609
bool last = slice_idx == slice_count - 1;
1610
// For per-column scales, we only fetch them here in the final step before
1611
// write-out
1566
- if constexpr (!has_act_order && group_blocks == -1 && !has_zp) {
1612
+ if constexpr (!has_act_order && group_blocks == -1 &&
1613
+ (has_zp && dequant_skip_flop || !has_zp)) {
1614
if (w_type.size_bits() == 8 || (last || use_atomic_add)) {
1615
if (s_sh_wr_pred) {
1616
cp_async4(&sh_s[s_sh_wr], &scales_ptr[s_gl_rd]);
@@ -1573,7 +1620,8 @@ __global__ void Marlin(
1620
}
1621
1622
thread_block_reduce();
1576
- if constexpr (!has_act_order && group_blocks == -1 && !has_zp) {
1623
+ if constexpr (!has_act_order && group_blocks == -1 &&
1624
+ (has_zp && dequant_skip_flop || !has_zp)) {
1625
if (w_type.size_bits() == 8 || (last || use_atomic_add)) {
1626
cp_async_wait<0>();
1627
__syncthreads();
@@ -1597,7 +1645,8 @@ __global__ void Marlin(
1645
// that converts the fp32 results to fp16 (so that we avoid possible
1646
// overflow in fp16)
1647
if constexpr (!has_act_order && group_blocks == -1 &&
1600
- w_type.size_bits() == 8 && !has_zp) {
1648
+ w_type.size_bits() == 8 &&
1649
+ (has_zp && dequant_skip_flop || !has_zp)) {
1650
if (threadIdx.x / 32 < thread_n_blocks / 4) {
1651
#pragma unroll
1652
for (int i = 0; i < thread_m_blocks; i++) {
csrc/torch_bindings.cpp
+2
-2
@@ -292,8 +292,8 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
292
// gptq_marlin Optimized Quantized GEMM for GPTQ.
293
ops.def(
294
"gptq_marlin_gemm(Tensor a, Tensor? c_or_none, Tensor b_q_weight, "
295
- "Tensor b_scales, Tensor? b_zeros_or_none, Tensor? g_idx_or_none, "
296
- "Tensor? perm_or_none, Tensor workspace, int b_q_type, "
295
+ "Tensor b_scales, Tensor? global_scale, Tensor? b_zeros_or_none, Tensor? "
296
+ "g_idx_or_none, Tensor? perm_or_none, Tensor workspace, int b_q_type, "
297
"SymInt size_m, SymInt size_n, SymInt size_k, bool is_k_full, "
298
"bool use_atomic_add, bool use_fp32_reduce, bool is_zp_float) -> Tensor",
299
{stride_tag});
tests/kernels/moe/test_moe.py
+103
-32
@@ -16,6 +16,8 @@ from vllm.model_executor.layers.fused_moe import fused_moe
16
from vllm.model_executor.layers.fused_moe.fused_moe import fused_topk
17
from vllm.model_executor.layers.fused_moe.moe_torch_iterative import (
18
fused_moe as iterative_moe)
19
+from vllm.model_executor.layers.quantization.utils.marlin_utils_fp4 import (
20
+ rand_marlin_weight_fp4_like)
21
from vllm.model_executor.layers.quantization.utils.marlin_utils_fp8 import (
22
marlin_quant_fp8_torch)
23
from vllm.model_executor.layers.quantization.utils.marlin_utils_test import (
@@ -286,21 +288,64 @@ def test_mixtral_moe(dtype: torch.dtype, padding: bool, use_rocm_aiter: bool,
288
atol=mixtral_moe_tol[dtype])
289
290
291
+def marlin_moe_generate_valid_test_cases():
292
+ import itertools
293
+ m_list = [1, 123, 666]
294
+ n_list = [128, 1024]
295
+ k_list = [256, 2048]
296
+ e_list = [4, 12]
297
+ topk_list = [2, 3]
298
+ ep_size_list = [1, 4]
299
+ dtype_list = [torch.half, torch.bfloat16]
300
+ group_size_list = [-1, 16, 32, 128]
301
+ act_order_list = [True, False]
302
+ quant_type_list = [
303
+ scalar_types.float4_e2m1f,
304
+ scalar_types.float8_e4m3fn,
305
+ scalar_types.uint4,
306
+ scalar_types.uint4b8,
307
+ scalar_types.uint8b128,
308
+ ]
309
+ is_k_full_list = [True, False]
310
+
311
+ all_combinations = itertools.product(m_list, n_list, k_list, e_list,
312
+ topk_list, ep_size_list, dtype_list,
313
+ group_size_list, act_order_list,
314
+ quant_type_list, is_k_full_list)
315
+
316
+ def is_invalid(m, n, k, e, topk, ep_size, dtype, group_size, act_order,
317
+ quant_type, is_k_full):
318
+
319
+ if quant_type == scalar_types.float8_e4m3fn and \
320
+ group_size not in [-1, 128]:
321
+ return False
322
+ if quant_type == scalar_types.float4_e2m1f and group_size != 16:
323
+ return False
324
+ if quant_type != scalar_types.float4_e2m1f and group_size == 16:
325
+ return False
326
+
327
+ # Filter act_order
328
+ if act_order:
329
+ if group_size in (-1, k, n):
330
+ return False
331
+ if quant_type not in [scalar_types.uint4b8]:
332
+ return False
333
+ elif not is_k_full:
334
+ return False
335
+
336
+ return True
337
+
338
+ cases = []
339
+ for case in all_combinations:
340
+ if is_invalid(*case):
341
+ cases.append(case)
342
+ return cases
343
+
344
+
345
@pytest.mark.flaky(reruns=2)
290
-@pytest.mark.parametrize("m", [1, 123, 666])
291
-@pytest.mark.parametrize("n", [128, 1024])
292
-@pytest.mark.parametrize("k", [256, 2048])
293
-@pytest.mark.parametrize("e", [4, 12])
294
-@pytest.mark.parametrize("topk", [2, 3])
295
-@pytest.mark.parametrize("ep_size", [1, 4])
296
-@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16])
297
-@pytest.mark.parametrize("group_size", [-1, 32, 128])
298
-@pytest.mark.parametrize("act_order", [True, False])
299
-@pytest.mark.parametrize("quant_type", [
300
- scalar_types.uint4, scalar_types.uint8b128, scalar_types.uint4b8,
301
- scalar_types.float8_e4m3fn
302
-])
303
-@pytest.mark.parametrize("is_k_full", [True, False])
346
+@pytest.mark.parametrize(("m, n, k, e, topk, ep_size, dtype, group_size,"
347
+ "act_order, quant_type, is_k_full"),
348
+ marlin_moe_generate_valid_test_cases())
349
@pytest.mark.skipif(current_platform.is_rocm(), reason="Skip for rocm")
350
def test_fused_marlin_moe(
351
m: int,
@@ -338,6 +383,11 @@ def test_fused_marlin_moe(
383
if not is_k_full:
384
return
385
386
+ if quant_type == scalar_types.float4_e2m1f and group_size != 16:
387
+ return
388
+ if quant_type != scalar_types.float4_e2m1f and group_size == 16:
389
+ return
390
+
391
a = torch.randn((m, k), device="cuda", dtype=dtype) / 10
392
w1 = torch.randn((e, 2 * n, k), device="cuda", dtype=dtype) / 20
393
w2 = torch.randn((e, k, n), device="cuda", dtype=dtype) / 20
@@ -355,12 +405,27 @@ def test_fused_marlin_moe(
405
w_ref1_l = []
406
qweight1_l = []
407
scales1_l = []
408
+ global_scale1_l = []
409
zeros1_l = []
410
g_idx1_l = []
411
sort_indices1_l = []
412
413
for i in range(w1.shape[0]):
363
- if has_zp:
414
+ if quant_type == scalar_types.float4_e2m1f:
415
+ w_ref1, qweight1, scales1, global_scale1 = \
416
+ rand_marlin_weight_fp4_like(w1[i], group_size)
417
+
418
+ w_ref1_l.append(w_ref1.T)
419
+ qweight1_l.append(qweight1)
420
+ scales1_l.append(scales1)
421
+ global_scale1_l.append(global_scale1)
422
+ elif quant_type == scalar_types.float8_e4m3fn:
423
+ w_ref1, qweight1, scales1 = marlin_quant_fp8_torch(
424
+ w1[i], group_size)
425
+ w_ref1_l.append(w_ref1.T)
426
+ qweight1_l.append(qweight1)
427
+ scales1_l.append(scales1)
428
+ elif has_zp:
429
w_ref1, qweight1, scales1, zeros1 = awq_marlin_quantize(
430
w1[i].transpose(1, 0), quant_type, group_size)
431
@@ -368,7 +433,7 @@ def test_fused_marlin_moe(
433
qweight1_l.append(qweight1)
434
scales1_l.append(scales1)
435
zeros1_l.append(zeros1)
371
- elif quant_type != scalar_types.float8_e4m3fn:
436
+ else:
437
test_perm = torch.randperm(k)
438
w_ref1, qweight1, scales1, g_idx1, sort_indices1, _ = \
439
marlin_quantize(w1[i].transpose(1, 0), quant_type,
@@ -379,16 +444,11 @@ def test_fused_marlin_moe(
444
scales1_l.append(scales1)
445
g_idx1_l.append(g_idx1)
446
sort_indices1_l.append(sort_indices1)
382
- else:
383
- w_ref1, qweight1, scales1 = marlin_quant_fp8_torch(
384
- w1[i], group_size)
385
- w_ref1_l.append(w_ref1.T)
386
- qweight1_l.append(qweight1)
387
- scales1_l.append(scales1)
447
448
w_ref1 = stack_and_dev(w_ref1_l)
449
qweight1 = stack_and_dev(qweight1_l).contiguous()
450
scales1 = stack_and_dev(scales1_l)
451
+ global_scale1 = stack_and_dev(global_scale1_l) if global_scale1_l else None
452
g_idx1 = stack_and_dev(g_idx1_l) if g_idx1_l else None
453
zeros1 = stack_and_dev(zeros1_l) if zeros1_l else None
454
sort_indices1 = stack_and_dev(sort_indices1_l) if sort_indices1_l else None
@@ -396,12 +456,27 @@ def test_fused_marlin_moe(
456
w_ref2_l = []
457
qweight2_l = []
458
scales2_l = []
459
+ global_scale2_l = []
460
zeros2_l = []
461
g_idx2_l = []
462
sort_indices2_l = []
463
464
for i in range(w2.shape[0]):
404
- if has_zp:
465
+ if quant_type == scalar_types.float4_e2m1f:
466
+ w_ref2, qweight2, scales2, global_scale2 = \
467
+ rand_marlin_weight_fp4_like(w2[i], group_size)
468
+
469
+ w_ref2_l.append(w_ref2.T)
470
+ qweight2_l.append(qweight2)
471
+ scales2_l.append(scales2)
472
+ global_scale2_l.append(global_scale2)
473
+ elif quant_type == scalar_types.float8_e4m3fn:
474
+ w_ref2, qweight2, scales2 = marlin_quant_fp8_torch(
475
+ w2[i], group_size)
476
+ w_ref2_l.append(w_ref2.T)
477
+ qweight2_l.append(qweight2)
478
+ scales2_l.append(scales2)
479
+ elif has_zp:
480
w_ref2, qweight2, scales2, zeros2 = awq_marlin_quantize(
481
w2[i].transpose(1, 0), quant_type, group_size)
482
@@ -409,7 +484,7 @@ def test_fused_marlin_moe(
484
qweight2_l.append(qweight2)
485
scales2_l.append(scales2)
486
zeros2_l.append(zeros2)
412
- elif quant_type != scalar_types.float8_e4m3fn:
487
+ else:
488
test_perm = torch.randperm(n)
489
w_ref2, qweight2, scales2, g_idx2, sort_indices2, _ = \
490
marlin_quantize(w2[i].transpose(1, 0), quant_type,
@@ -420,24 +495,18 @@ def test_fused_marlin_moe(
495
scales2_l.append(scales2)
496
g_idx2_l.append(g_idx2)
497
sort_indices2_l.append(sort_indices2)
423
- else:
424
- w_ref2, qweight2, scales2 = marlin_quant_fp8_torch(
425
- w2[i], group_size)
426
- w_ref2_l.append(w_ref2.T)
427
- qweight2_l.append(qweight2)
428
- scales2_l.append(scales2)
498
499
w_ref2 = stack_and_dev(w_ref2_l)
500
qweight2 = stack_and_dev(qweight2_l).contiguous()
501
scales2 = stack_and_dev(scales2_l)
502
+ global_scale2 = stack_and_dev(global_scale2_l) if global_scale2_l else None
503
g_idx2 = stack_and_dev(g_idx2_l) if g_idx2_l else None
504
zeros2 = stack_and_dev(zeros2_l) if zeros2_l else None
505
sort_indices2 = stack_and_dev(sort_indices2_l) if sort_indices2_l else None
506
507
score = torch.randn((m, e), device="cuda", dtype=dtype)
508
439
- topk_weights, topk_ids, token_expert_indices = fused_topk(
440
- a, score, topk, False)
509
+ topk_weights, topk_ids, _ = fused_topk(a, score, topk, False)
510
511
torch_output = torch_moe(a, w_ref1, w_ref2, score, topk, e_map)
512
@@ -452,6 +521,8 @@ def test_fused_marlin_moe(
521
topk_ids,
522
global_num_experts=e,
523
expert_map=e_map,
524
+ global_scale1=global_scale1,
525
+ global_scale2=global_scale2,
526
g_idx1=g_idx1,
527
g_idx2=g_idx2,
528
sort_indices1=sort_indices1,
tests/kernels/quantization/test_marlin_gemm.py
+41
-73
@@ -20,6 +20,8 @@ from vllm.model_executor.layers.quantization.utils.marlin_utils import (
20
MARLIN_SUPPORTED_GROUP_SIZES, marlin_make_empty_g_idx,
21
marlin_make_workspace_new, marlin_permute_scales,
22
query_marlin_supported_quant_types)
23
+from vllm.model_executor.layers.quantization.utils.marlin_utils_fp4 import (
24
+ FP4_MARLIN_SUPPORTED_GROUP_SIZES, rand_marlin_weight_fp4_like)
25
from vllm.model_executor.layers.quantization.utils.marlin_utils_fp8 import (
26
marlin_quant_fp8_torch)
27
from vllm.model_executor.layers.quantization.utils.marlin_utils_test import (
@@ -190,9 +192,10 @@ def test_awq_marlin_repack(k_chunk, n_chunk, quant_type, group_size,
192
reason="Marlin is not supported on this GPU type.")
193
@pytest.mark.parametrize("k_chunk", MARLIN_K_CHUNKS)
194
@pytest.mark.parametrize("n_chunk", MARLIN_N_CHUNKS)
193
-@pytest.mark.parametrize("quant_type",
194
- query_marlin_supported_quant_types(False))
195
-@pytest.mark.parametrize("group_size", MARLIN_SUPPORTED_GROUP_SIZES)
195
+@pytest.mark.parametrize("quant_type", query_marlin_supported_quant_types())
196
+@pytest.mark.parametrize(
197
+ "group_size",
198
+ set(MARLIN_SUPPORTED_GROUP_SIZES + FP4_MARLIN_SUPPORTED_GROUP_SIZES))
199
@pytest.mark.parametrize("mnk_factors", MNK_FACTORS)
200
@pytest.mark.parametrize("act_order", ACT_ORDER_OPTS)
201
@pytest.mark.parametrize("is_k_full", K_FULL_OPTS)
@@ -210,6 +213,7 @@ def test_gptq_marlin_gemm(
213
use_fp32_reduce,
214
):
215
m_factor, n_factor, k_factor = mnk_factors
216
+ has_zp = quant_type in [scalar_types.uint4, scalar_types.uint8]
217
218
size_m = m_factor
219
size_k = k_chunk * k_factor
@@ -220,6 +224,8 @@ def test_gptq_marlin_gemm(
224
return
225
if group_size == size_k:
226
return
227
+ if has_zp:
228
+ return
229
230
if size_k % group_size != 0:
231
return
@@ -227,7 +233,15 @@ def test_gptq_marlin_gemm(
233
a_input = rand_data((size_m, size_k))
234
b_weight = rand_data((size_k, size_n))
235
230
- if quant_type == scalar_types.float8_e4m3fn:
236
+ if quant_type == scalar_types.float4_e2m1f:
237
+ if group_size != 16 or act_order:
238
+ return
239
+ w_ref, marlin_q_w, marlin_s, marlin_s2 = rand_marlin_weight_fp4_like(
240
+ b_weight.T, group_size)
241
+ g_idx = None
242
+ sort_indices = None
243
+ marlin_zp = None
244
+ elif quant_type == scalar_types.float8_e4m3fn:
245
if group_size not in [-1, 128]:
246
return
247
if act_order:
@@ -236,26 +250,39 @@ def test_gptq_marlin_gemm(
250
b_weight.T, group_size)
251
g_idx = None
252
sort_indices = None
253
+ marlin_zp = None
254
+ marlin_s2 = None
255
+ elif has_zp:
256
+ if group_size == 16:
257
+ return
258
+ w_ref, marlin_q_w, marlin_s, marlin_zp = awq_marlin_quantize(
259
+ b_weight, quant_type, group_size)
260
+ g_idx = None
261
+ sort_indices = None
262
+ marlin_s2 = None
263
else:
264
+ if group_size == 16:
265
+ return
266
w_ref, marlin_q_w, marlin_s, g_idx, sort_indices, _ = marlin_quantize(
267
b_weight, quant_type, group_size, act_order)
242
-
243
- marlin_zp = marlin_make_empty_g_idx(marlin_s.device)
268
+ marlin_zp = None
269
+ marlin_s2 = None
270
271
workspace = marlin_make_workspace_new(w_ref.device)
272
247
- opcheck(
248
- torch.ops._C.gptq_marlin_gemm,
249
- (a_input, None, marlin_q_w, marlin_s, marlin_zp, g_idx, sort_indices,
250
- workspace, quant_type.id, a_input.shape[0], b_weight.shape[1],
251
- a_input.shape[1], is_k_full, use_atomic_add, use_fp32_reduce, False),
252
- test_utils=DEFAULT_OPCHECK_TEST_UTILS)
273
+ opcheck(torch.ops._C.gptq_marlin_gemm,
274
+ (a_input, None, marlin_q_w, marlin_s, marlin_s2, marlin_zp, g_idx,
275
+ sort_indices, workspace, quant_type.id, a_input.shape[0],
276
+ b_weight.shape[1], a_input.shape[1], is_k_full, use_atomic_add,
277
+ use_fp32_reduce, False),
278
+ test_utils=DEFAULT_OPCHECK_TEST_UTILS)
279
280
output = ops.gptq_marlin_gemm(
281
a_input,
282
None,
283
marlin_q_w,
284
marlin_s,
285
+ marlin_s2,
286
marlin_zp,
287
g_idx,
288
sort_indices,
@@ -339,67 +366,6 @@ def test_gptq_marlin_24_gemm(k_chunk, n_chunk, quant_type, group_size,
366
assert max_diff < 0.04
367
368
342
-@pytest.mark.skipif(not is_quant_method_supported("gptq_marlin"),
343
- reason="Marlin is not supported on this GPU type.")
344
-@pytest.mark.parametrize("k_chunk", MARLIN_K_CHUNKS)
345
-@pytest.mark.parametrize("n_chunk", MARLIN_N_CHUNKS)
346
-@pytest.mark.parametrize("quant_type",
347
- query_marlin_supported_quant_types(True))
348
-@pytest.mark.parametrize("group_size", MARLIN_SUPPORTED_GROUP_SIZES)
349
-@pytest.mark.parametrize("mnk_factors", MNK_FACTORS)
350
-@pytest.mark.parametrize("use_fp32_reduce", USE_FP32_REDUCE_OPTS)
351
-def test_awq_marlin_gemm(
352
- k_chunk,
353
- n_chunk,
354
- quant_type,
355
- group_size,
356
- mnk_factors,
357
- use_fp32_reduce,
358
-):
359
- m_factor, n_factor, k_factor = mnk_factors
360
-
361
- size_m = m_factor
362
- size_k = k_chunk * k_factor
363
- size_n = n_chunk * n_factor
364
-
365
- a_input = rand_data((size_m, size_k))
366
- b_weight = rand_data((size_k, size_n))
367
-
368
- w_ref, marlin_q_w, marlin_s, marlin_zp = awq_marlin_quantize(
369
- b_weight, quant_type, group_size)
370
-
371
- g_idx = torch.empty(0, dtype=torch.int, device=marlin_q_w.device)
372
- sort_indices = torch.empty(0, dtype=torch.int, device=marlin_q_w.device)
373
- is_k_full = True
374
-
375
- workspace = marlin_make_workspace_new(a_input.device)
376
-
377
- output = ops.gptq_marlin_gemm(
378
- a_input,
379
- None,
380
- marlin_q_w,
381
- marlin_s,
382
- marlin_zp,
383
- g_idx,
384
- sort_indices,
385
- workspace,
386
- quant_type,
387
- a_input.shape[0],
388
- b_weight.shape[1],
389
- a_input.shape[1],
390
- is_k_full=is_k_full,
391
- use_fp32_reduce=use_fp32_reduce,
392
- is_zp_float=False,
393
- )
394
- output_ref = torch.matmul(a_input, w_ref)
395
-
396
- torch.cuda.synchronize()
397
-
398
- max_diff = compute_max_diff(output, output_ref)
399
-
400
- assert max_diff < 0.04
401
-
402
-
369
@pytest.mark.skipif(not is_quant_method_supported("gptq_marlin"),
370
reason="Marlin is not supported on this GPU type.")
371
@pytest.mark.parametrize("k_chunk", MARLIN_K_CHUNKS)
@@ -452,6 +418,7 @@ def test_hqq_marlin_gemm(
418
None,
419
marlin_w_q,
420
marlin_s,
421
+ None,
422
marlin_zp,
423
g_idx,
424
g_idx_sort_indices,
@@ -564,6 +531,7 @@ def test_marlin_gemm_subset_input():
531
None,
532
marlin_q_w,
533
marlin_s,
534
+ None,
535
marlin_zp,
536
g_idx,
537
sort_indices,
vllm/_custom_ops.py
+12
-8
@@ -333,6 +333,7 @@ if hasattr(torch.ops._C, "gptq_marlin_24_gemm"):
333
c: Optional[torch.Tensor],
334
b_q_weight: torch.Tensor,
335
b_scales: torch.Tensor,
336
+ global_scale: Optional[torch.Tensor],
337
b_zeros: Optional[torch.Tensor],
338
g_idx: Optional[torch.Tensor],
339
perm: Optional[torch.Tensor],
@@ -866,6 +867,7 @@ def gptq_marlin_gemm(a: torch.Tensor,
867
c: Optional[torch.Tensor],
868
b_q_weight: torch.Tensor,
869
b_scales: torch.Tensor,
870
+ global_scale: Optional[torch.Tensor],
871
b_zeros: Optional[torch.Tensor],
872
g_idx: Optional[torch.Tensor],
873
perm: Optional[torch.Tensor],
@@ -878,9 +880,10 @@ def gptq_marlin_gemm(a: torch.Tensor,
880
use_atomic_add: bool = False,
881
use_fp32_reduce: bool = False,
882
is_zp_float: bool = False) -> torch.Tensor:
881
- return torch.ops._C.gptq_marlin_gemm(a, c, b_q_weight, b_scales, b_zeros,
882
- g_idx, perm, workspace, b_q_type.id,
883
- size_m, size_n, size_k, is_k_full,
883
+ return torch.ops._C.gptq_marlin_gemm(a, c, b_q_weight, b_scales,
884
+ global_scale, b_zeros, g_idx, perm,
885
+ workspace, b_q_type.id, size_m,
886
+ size_n, size_k, is_k_full,
887
use_atomic_add, use_fp32_reduce,
888
is_zp_float)
889
@@ -1381,6 +1384,7 @@ def topk_softmax(topk_weights: torch.Tensor, topk_ids: torch.Tensor,
1384
1385
def moe_wna16_marlin_gemm(input: torch.Tensor, output: Optional[torch.Tensor],
1386
b_qweight: torch.Tensor, b_scales: torch.Tensor,
1387
+ global_scale: Optional[torch.Tensor],
1388
b_qzeros: Optional[torch.Tensor],
1389
g_idx: Optional[torch.Tensor],
1390
perm: Optional[torch.Tensor],
@@ -1395,11 +1399,11 @@ def moe_wna16_marlin_gemm(input: torch.Tensor, output: Optional[torch.Tensor],
1399
use_fp32_reduce: bool,
1400
is_zp_float: bool) -> torch.Tensor:
1401
return torch.ops._moe_C.moe_wna16_marlin_gemm(
1398
- input, output, b_qweight, b_scales, b_qzeros, g_idx, perm, workspace,
1399
- sorted_token_ids, expert_ids, num_tokens_past_padded, topk_weights,
1400
- moe_block_size, top_k, mul_topk_weights, is_ep, b_q_type.id, size_m,
1401
- size_n, size_k, is_k_full, use_atomic_add, use_fp32_reduce,
1402
- is_zp_float)
1402
+ input, output, b_qweight, b_scales, global_scale, b_qzeros, g_idx,
1403
+ perm, workspace, sorted_token_ids, expert_ids, num_tokens_past_padded,
1404
+ topk_weights, moe_block_size, top_k, mul_topk_weights, is_ep,
1405
+ b_q_type.id, size_m, size_n, size_k, is_k_full, use_atomic_add,
1406
+ use_fp32_reduce, is_zp_float)
1407
1408
1409
if supports_moe_ops and hasattr(torch.ops._moe_C, "marlin_gemm_moe"):
vllm/model_executor/layers/fused_moe/fused_marlin_moe.py
+11
-3
@@ -25,6 +25,8 @@ def fused_marlin_moe(hidden_states: torch.Tensor,
25
quant_type_id: int,
26
global_num_experts: int = -1,
27
expert_map: Optional[torch.Tensor] = None,
28
+ global_scale1: Optional[torch.Tensor] = None,
29
+ global_scale2: Optional[torch.Tensor] = None,
30
g_idx1: Optional[torch.Tensor] = None,
31
g_idx2: Optional[torch.Tensor] = None,
32
sort_indices1: Optional[torch.Tensor] = None,
@@ -64,11 +66,13 @@ def fused_marlin_moe(hidden_states: torch.Tensor,
66
quant_type = ScalarType.from_id(quant_type_id)
67
assert quant_type in [
68
scalar_types.uint4, scalar_types.uint8b128, scalar_types.uint4b8,
67
- scalar_types.float8_e4m3fn
69
+ scalar_types.float8_e4m3fn, scalar_types.float4_e2m1f
70
]
71
70
- int4_scalar_types = [scalar_types.uint4, scalar_types.uint4b8]
71
- num_bits = 4 if quant_type in int4_scalar_types else 8
72
+ bit4_scalar_types = [
73
+ scalar_types.uint4, scalar_types.uint4b8, scalar_types.float4_e2m1f
74
+ ]
75
+ num_bits = 4 if quant_type in bit4_scalar_types else 8
76
77
# Check constraints.
78
assert hidden_states.shape[0] == gating_output.shape[
@@ -133,6 +137,7 @@ def fused_marlin_moe(hidden_states: torch.Tensor,
137
intermediate_cache1,
138
w1,
139
w1_scale,
140
+ global_scale1,
141
w1_zeros,
142
g_idx1,
143
sort_indices1,
@@ -165,6 +170,7 @@ def fused_marlin_moe(hidden_states: torch.Tensor,
170
intermediate_cache3,
171
w2,
172
w2_scale,
173
+ global_scale2,
174
w2_zeros,
175
g_idx2,
176
sort_indices2,
@@ -202,6 +208,8 @@ def fused_marlin_moe_fake(hidden_states: torch.Tensor,
208
topk_ids: torch.Tensor,
209
quant_type_id: int,
210
global_num_experts: int = -1,
211
+ global_scale1: Optional[torch.Tensor] = None,
212
+ global_scale2: Optional[torch.Tensor] = None,
213
expert_map: Optional[torch.Tensor] = None,
214
g_idx1: Optional[torch.Tensor] = None,
215
g_idx2: Optional[torch.Tensor] = None,
vllm/model_executor/layers/quantization/hqq_marlin.py
+3
-1
@@ -304,8 +304,10 @@ class HQQMarlinMethod(LinearMethodBase):
304
305
marlin_out = ops.gptq_marlin_gemm(
306
x,
307
+ None,
308
layer.marlin_qweight,
309
scales,
310
+ None,
311
zeros,
312
layer.g_idx,
313
layer.g_idx_sort_indices,
@@ -315,7 +317,7 @@ class HQQMarlinMethod(LinearMethodBase):
317
self.output_size_per_partition,
318
self.input_size_per_partition,
319
True, # is_k_full
318
- True, # has_zp
320
+ False, # use atomic add
321
True, # use 32-bit reduce
322
True, # use float zp
323
)
vllm/model_executor/layers/quantization/modelopt.py
+80
-4
@@ -17,6 +17,9 @@ from vllm.model_executor.layers.quantization import QuantizationMethods
17
from vllm.model_executor.layers.quantization.base_config import (
18
QuantizationConfig, QuantizeMethodBase)
19
from vllm.model_executor.layers.quantization.kv_cache import BaseKVCacheMethod
20
+from vllm.model_executor.layers.quantization.utils.marlin_utils_fp4 import (
21
+ apply_fp4_marlin_linear, is_fp4_marlin_supported,
22
+ prepare_fp4_layer_for_marlin, prepare_moe_fp4_layer_for_marlin)
23
from vllm.model_executor.layers.quantization.utils.quant_utils import (
24
is_layer_skipped)
25
from vllm.model_executor.layers.quantization.utils.w8a8_utils import (
@@ -24,6 +27,7 @@ from vllm.model_executor.layers.quantization.utils.w8a8_utils import (
27
from vllm.model_executor.parameter import (ModelWeightParameter,
28
PerTensorScaleParameter)
29
from vllm.platforms import current_platform
30
+from vllm.scalar_type import scalar_types
31
32
logger = init_logger(__name__)
33
@@ -196,7 +200,7 @@ class ModelOptNvFp4Config(QuantizationConfig):
200
201
@classmethod
202
def get_min_capability(cls) -> int:
199
- return 100
203
+ return 80
204
205
@classmethod
206
def get_config_filenames(cls) -> List[str]:
@@ -278,9 +282,15 @@ class ModelOptNvFp4LinearMethod(LinearMethodBase):
282
def __init__(self, quant_config: ModelOptNvFp4Config):
283
self.quant_config = quant_config
284
self.cutlass_nvfp4_supported = cutlass_fp4_supported()
285
+ self.use_marlin = False
286
+
287
if not self.cutlass_nvfp4_supported:
282
- raise ValueError("Current platform does not support NVFP4"
283
- " quantization. Please use Blackwell and above.")
288
+ if is_fp4_marlin_supported():
289
+ self.use_marlin = True
290
+ else:
291
+ raise ValueError("Current platform does not support NVFP4"
292
+ " quantization. Please use Blackwell and"
293
+ " above.")
294
295
def create_weights(
296
self,
@@ -392,12 +402,29 @@ class ModelOptNvFp4LinearMethod(LinearMethodBase):
402
layer.weight_scale_swizzled = Parameter(swizzled_weight_scale,
403
requires_grad=False)
404
405
+ if self.use_marlin:
406
+ prepare_fp4_layer_for_marlin(layer)
407
+ del layer.alpha
408
+ del layer.input_scale
409
+ del layer.weight_scale_swizzled
410
+
411
def apply(
412
self,
413
layer: torch.nn.Module,
414
x: torch.Tensor,
415
bias: Optional[torch.Tensor] = None,
416
) -> torch.Tensor:
417
+ if self.use_marlin:
418
+ return apply_fp4_marlin_linear(
419
+ input=x,
420
+ weight=layer.weight,
421
+ weight_scale=layer.weight_scale,
422
+ weight_scale_2=layer.weight_scale_2,
423
+ workspace=layer.workspace,
424
+ size_n=layer.output_size_per_partition,
425
+ size_k=layer.input_size_per_partition,
426
+ bias=bias)
427
+
428
output_dtype = x.dtype
429
430
# for input only the contracting dimension has a constraint.
@@ -434,6 +461,16 @@ class ModelOptNvFp4FusedMoE(FusedMoEMethodBase):
461
462
def __init__(self, quant_config: ModelOptNvFp4Config):
463
self.quant_config = quant_config
464
+ self.cutlass_nvfp4_supported = cutlass_fp4_supported()
465
+ self.use_marlin = False
466
+
467
+ if not self.cutlass_nvfp4_supported:
468
+ if is_fp4_marlin_supported():
469
+ self.use_marlin = True
470
+ else:
471
+ raise ValueError("Current platform does not support NVFP4"
472
+ " quantization. Please use Blackwell and"
473
+ " above.")
474
475
def create_weights(self, layer: torch.nn.Module, num_experts: int,
476
hidden_size: int, intermediate_size_per_partition: int,
@@ -442,6 +479,8 @@ class ModelOptNvFp4FusedMoE(FusedMoEMethodBase):
479
raise ValueError("NVFP4 quantization was selected, "
480
" dynamic quantization is not supported.")
481
482
+ layer.num_experts = num_experts
483
+ layer.params_dtype = params_dtype
484
layer.quant_config = self.quant_config
485
weight_dtype = torch.uint8
486
weight_scale_dtype = torch.float8_e4m3fn
@@ -594,7 +633,15 @@ class ModelOptNvFp4FusedMoE(FusedMoEMethodBase):
633
634
layer.w2_blockscale_swizzled = Parameter(w2_blockscale_swizzled,
635
requires_grad=False)
597
- return
636
+
637
+ if self.use_marlin:
638
+ prepare_moe_fp4_layer_for_marlin(layer)
639
+ del layer.g1_alphas
640
+ del layer.g2_alphas
641
+ del layer.w13_input_scale_quant
642
+ del layer.w2_input_scale_quant
643
+ del layer.w13_blockscale_swizzled
644
+ del layer.w2_blockscale_swizzled
645
646
def apply(
647
self,
@@ -614,6 +661,35 @@ class ModelOptNvFp4FusedMoE(FusedMoEMethodBase):
661
apply_router_weight_on_input: bool = False,
662
activation: str = "silu",
663
):
664
+ if self.use_marlin:
665
+ topk_weights, topk_ids = FusedMoE.select_experts(
666
+ hidden_states=x,
667
+ router_logits=router_logits,
668
+ use_grouped_topk=use_grouped_topk,
669
+ top_k=top_k,
670
+ renormalize=renormalize,
671
+ topk_group=topk_group,
672
+ num_expert_group=num_expert_group,
673
+ custom_routing_function=custom_routing_function,
674
+ scoring_func=scoring_func,
675
+ e_score_correction_bias=e_score_correction_bias,
676
+ )
677
+
678
+ return torch.ops.vllm.fused_marlin_moe(
679
+ x,
680
+ layer.w13_weight,
681
+ layer.w2_weight,
682
+ layer.w13_weight_scale,
683
+ layer.w2_weight_scale,
684
+ router_logits,
685
+ topk_weights,
686
+ topk_ids,
687
+ global_scale1=layer.w13_weight_scale_2,
688
+ global_scale2=layer.w2_weight_scale_2,
689
+ quant_type_id=scalar_types.float4_e2m1f.id,
690
+ global_num_experts=global_num_experts,
691
+ expert_map=expert_map)
692
+
693
assert activation == "silu", "Only SiLU activation is supported."
694
assert not apply_router_weight_on_input, (
695
"Router weight on input is not "
vllm/model_executor/layers/quantization/utils/marlin_utils.py
+14
-2
@@ -33,7 +33,7 @@ USE_FP32_REDUCE_DEFAULT = True
33
# without runtime zero-point. We support common cases, i.e. AWQ and GPTQ.
34
# TODO: we may want to move this into the C++ so its closer to the actual impl
35
def query_marlin_supported_quant_types(
36
- has_zp: bool,
36
+ has_zp: Optional[bool] = None,
37
include_fp_type: bool = True,
38
device_capability: Optional[int] = None,
39
):
@@ -45,6 +45,16 @@ def query_marlin_supported_quant_types(
45
if device_capability < 80:
46
return []
47
48
+ # - has_zp is True: return quant_types that has zero points
49
+ # - has_zp is False: return quant_types that has not zero points
50
+ # - has_zp is None: both
51
+ if has_zp is None:
52
+ types0 = query_marlin_supported_quant_types(False, include_fp_type,
53
+ device_capability)
54
+ types1 = query_marlin_supported_quant_types(True, include_fp_type,
55
+ device_capability)
56
+ return types0 + types1
57
+
58
if has_zp:
59
# AWQ style, unsigned + runtime zero-point
60
return [scalar_types.uint4]
@@ -52,7 +62,7 @@ def query_marlin_supported_quant_types(
62
# GPTQ style, unsigned + symmetric bias
63
res = [scalar_types.uint4b8, scalar_types.uint8b128]
64
if include_fp_type:
55
- res += [scalar_types.float8_e4m3fn]
65
+ res += [scalar_types.float8_e4m3fn, scalar_types.float4_e2m1f]
66
return res
67
68
@@ -394,6 +404,7 @@ def apply_gptq_marlin_linear(
404
None,
405
weight,
406
weight_scale,
407
+ None,
408
weight_zp,
409
g_idx,
410
g_idx_sort_indices,
@@ -439,6 +450,7 @@ def apply_awq_marlin_linear(
450
None,
451
weight,
452
weight_scale,
453
+ None,
454
weight_zp,
455
g_idx,
456
g_idx_sort_indices,
vllm/model_executor/layers/quantization/utils/marlin_utils_fp4.py
new
+277
@@ -0,0 +1,277 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+
3
+from typing import Optional
4
+
5
+import torch
6
+
7
+import vllm._custom_ops as ops
8
+from vllm.logger import init_logger
9
+from vllm.model_executor.layers.quantization.utils.marlin_utils import (
10
+ USE_FP32_REDUCE_DEFAULT, marlin_make_workspace_new, marlin_permute_scales,
11
+ should_use_atomic_add_reduce)
12
+from vllm.platforms import current_platform
13
+from vllm.scalar_type import scalar_types
14
+
15
+FP4_MARLIN_SUPPORTED_GROUP_SIZES = [16]
16
+
17
+logger = init_logger(__name__)
18
+
19
+
20
+def is_fp4_marlin_supported():
21
+ return current_platform.has_device_capability(80)
22
+
23
+
24
+def fp4_marlin_process_scales(marlin_scales):
25
+ assert (marlin_scales >= 0).all()
26
+
27
+ # convert to half first, we would convert to fp8 later
28
+ marlin_scales = marlin_scales.to(torch.half)
29
+
30
+ # 8 is the number of scale number using by one thread
31
+ marlin_scales = marlin_scales.view(marlin_scales.size(0) // 2, 2, -1, 8)
32
+ marlin_scales = marlin_scales.permute(0, 2, 1, 3).reshape(
33
+ marlin_scales.size(0) * 2, -1)
34
+
35
+ # fit the layout of fp8 dequantization
36
+ marlin_scales = marlin_scales.view(-1, 4)[:, [0, 2, 1, 3]].view(
37
+ marlin_scales.size(0), -1)
38
+
39
+ # We assume that weight_scale (FP8-S1E4M3) is always greater
40
+ # than or equal to 0. So we can convert
41
+ # (weight_scale * (2 ** 7) to a special FP8-S0E5M3 format.
42
+ # After multiplying by 2 ** 7, the top bit of FP8-S0E5M3 would always be 1
43
+ # when weight_scale > 0. This allows us to have an exponent bias
44
+ # closer to zero after dequantization.
45
+
46
+ marlin_scales = (marlin_scales * (2**7)).view(torch.int16) << 1
47
+ marlin_scales = marlin_scales.view(torch.float8_e4m3fn)
48
+ marlin_scales = marlin_scales[:, 1::2].contiguous()
49
+
50
+ return marlin_scales
51
+
52
+
53
+def fp4_marlin_process_global_scale(global_scale):
54
+ assert global_scale.dtype in [torch.half, torch.bfloat16]
55
+ fp4_exponent = 2
56
+ if global_scale.dtype == torch.half:
57
+ target_exponent = 5
58
+ elif global_scale.dtype == torch.bfloat16:
59
+ target_exponent = 8
60
+ # exponent_bias_fp16 = 2 ** 4 - 2 ** 1 = 14
61
+ # exponent_bias_bf16 = 2 ** 7 - 2 ** 1 = 126
62
+ exponent_bias = 2**(target_exponent - 1) - 2**(fp4_exponent - 1)
63
+ return global_scale * (2.0**(exponent_bias - 7))
64
+
65
+
66
+def apply_fp4_marlin_linear(
67
+ input: torch.Tensor,
68
+ weight: torch.Tensor,
69
+ weight_scale: torch.Tensor,
70
+ weight_scale_2: torch.Tensor,
71
+ workspace: torch.Tensor,
72
+ size_n: int,
73
+ size_k: int,
74
+ bias: Optional[torch.Tensor] = None,
75
+ use_fp32_reduce: bool = USE_FP32_REDUCE_DEFAULT) -> torch.Tensor:
76
+ # For GPUs that lack FP4 hardware support, we can leverage the
77
+ # Marlin kernel for fast weight-only FP4 quantization
78
+
79
+ reshaped_x = input.reshape(-1, input.shape[-1])
80
+ out_shape = input.shape[:-1] + (size_n, )
81
+
82
+ use_atomic_add = should_use_atomic_add_reduce(m=reshaped_x.size(0),
83
+ n=size_n,
84
+ k=size_k,
85
+ device=input.device,
86
+ dtype=input.dtype)
87
+
88
+ output = ops.gptq_marlin_gemm(a=reshaped_x,
89
+ c=None,
90
+ b_q_weight=weight,
91
+ b_scales=weight_scale,
92
+ global_scale=weight_scale_2,
93
+ b_zeros=None,
94
+ g_idx=None,
95
+ perm=None,
96
+ workspace=workspace,
97
+ b_q_type=scalar_types.float4_e2m1f,
98
+ size_m=reshaped_x.size(0),
99
+ size_n=size_n,
100
+ size_k=size_k,
101
+ use_atomic_add=use_atomic_add,
102
+ use_fp32_reduce=use_fp32_reduce)
103
+
104
+ if bias is not None:
105
+ output.add_(bias) # In-place add
106
+
107
+ return output.reshape(out_shape)
108
+
109
+
110
+def prepare_fp4_layer_for_marlin(layer: torch.nn.Module) -> None:
111
+ logger.warning_once(
112
+ "Your GPU does not have native support for FP4 computation but "
113
+ "FP4 quantization is being used. Weight-only FP4 compression will "
114
+ "be used leveraging the Marlin kernel. This may degrade "
115
+ "performance for compute-heavy workloads.")
116
+
117
+ part_size_n = layer.output_size_per_partition
118
+ part_size_k = layer.input_size_per_partition
119
+ param_dtype = layer.params_dtype
120
+
121
+ assert layer.weight.shape == (part_size_n, part_size_k // 2)
122
+
123
+ device = layer.weight.device
124
+
125
+ # WORKSPACE
126
+ layer.workspace = marlin_make_workspace_new(device)
127
+
128
+ # WEIGHT
129
+ # Repack weights to marlin format
130
+ perm = torch.empty(0, dtype=torch.int, device=device)
131
+ qweight = layer.weight.view(torch.int32).T.contiguous()
132
+
133
+ marlin_qweight = ops.gptq_marlin_repack(b_q_weight=qweight,
134
+ perm=perm,
135
+ size_k=part_size_k,
136
+ size_n=part_size_n,
137
+ num_bits=4)
138
+ layer.weight = torch.nn.Parameter(marlin_qweight, requires_grad=False)
139
+
140
+ # WEIGHT SCALES
141
+ # Permute scales
142
+ weight_scale = layer.weight_scale.T.to(param_dtype)
143
+ weight_scale = marlin_permute_scales(s=weight_scale,
144
+ size_k=part_size_k,
145
+ size_n=part_size_n,
146
+ group_size=16)
147
+ weight_scale = fp4_marlin_process_scales(weight_scale)
148
+ layer.weight_scale = torch.nn.Parameter(weight_scale, requires_grad=False)
149
+
150
+ weight_scale_2 = layer.weight_scale_2.to(param_dtype)
151
+ weight_scale_2 = fp4_marlin_process_global_scale(weight_scale_2)
152
+ layer.weight_scale_2 = torch.nn.Parameter(weight_scale_2,
153
+ requires_grad=False)
154
+
155
+ return
156
+
157
+
158
+def prepare_moe_fp4_layer_for_marlin(layer: torch.nn.Module) -> None:
159
+ logger.warning_once(
160
+ "Your GPU does not have native support for FP4 computation but "
161
+ "FP4 quantization is being used. Weight-only FP4 compression will "
162
+ "be used leveraging the Marlin kernel. This may degrade "
163
+ "performance for compute-heavy workloads.")
164
+
165
+ e = layer.num_experts
166
+ k = layer.hidden_size
167
+ n = layer.intermediate_size_per_partition
168
+
169
+ # WORKSPACE
170
+ device = layer.w13_weight.device
171
+ param_dtype = layer.params_dtype
172
+ layer.workspace = marlin_make_workspace_new(device, 4)
173
+ perm = torch.empty(0, dtype=torch.int, device=device)
174
+
175
+ # WEIGHT
176
+ # Repack weights to marlin format
177
+ for name in ["w13_weight", "w2_weight"]:
178
+ weight = getattr(layer, name)
179
+ tensor_list = []
180
+ if "w13" in name:
181
+ size_n, size_k = n * 2, k
182
+ else:
183
+ size_n, size_k = k, n
184
+
185
+ assert weight.shape == (e, size_n, size_k // 2)
186
+
187
+ for i in range(e):
188
+ qweight = weight[i].view(torch.int32).T.contiguous()
189
+
190
+ marlin_qweight = ops.gptq_marlin_repack(b_q_weight=qweight,
191
+ perm=perm,
192
+ size_k=size_k,
193
+ size_n=size_n,
194
+ num_bits=4)
195
+ tensor_list.append(marlin_qweight)
196
+
197
+ weight = torch.cat([x.unsqueeze(0) for x in tensor_list], 0)
198
+ weight = torch.nn.Parameter(weight, requires_grad=False)
199
+
200
+ setattr(layer, name, weight)
201
+
202
+ # WEIGHT SCALES
203
+ # Permute scales
204
+ for name in ["w13", "w2"]:
205
+ scales = getattr(layer, name + "_weight_scale").to(param_dtype)
206
+ global_scale = getattr(layer, name + "_weight_scale_2").to(param_dtype)
207
+
208
+ tensor_list = []
209
+ if "w13" in name:
210
+ size_n, size_k = n * 2, k
211
+ else:
212
+ size_n, size_k = k, n
213
+
214
+ for i in range(e):
215
+ marlin_scales = marlin_permute_scales(s=scales[i].T,
216
+ size_k=size_k,
217
+ size_n=size_n,
218
+ group_size=16)
219
+ marlin_scales = fp4_marlin_process_scales(marlin_scales)
220
+ tensor_list.append(marlin_scales)
221
+
222
+ scales = torch.cat([x.unsqueeze(0) for x in tensor_list], 0)
223
+ scales = torch.nn.Parameter(scales, requires_grad=False)
224
+ setattr(layer, name + "_weight_scale", scales)
225
+
226
+ global_scale = fp4_marlin_process_global_scale(global_scale)
227
+ global_scale = torch.nn.Parameter(global_scale, requires_grad=False)
228
+ setattr(layer, name + "_weight_scale_2", global_scale)
229
+
230
+
231
+def rand_marlin_weight_fp4_like(weight, group_size):
232
+ assert group_size > 0
233
+ size_n, size_k = weight.shape
234
+ device = weight.device
235
+
236
+ scales = weight.view(size_n, -1, group_size).abs().max(-1)[0] / 6
237
+ global_scale = scales.max() / 448
238
+ scales = (scales / global_scale).to(torch.float8_e4m3fn)
239
+
240
+ fp4_weight = torch.randint(0,
241
+ 256, (size_n, size_k // 2),
242
+ dtype=torch.uint8,
243
+ device=weight.device)
244
+ fp4_weight_part_1 = ((fp4_weight & 0b10000000) |
245
+ ((fp4_weight & 0b01110000) >> 2))
246
+ fp4_weight_part_1 = fp4_weight_part_1.view(torch.float8_e4m3fn)
247
+ fp4_weight_part_1 = fp4_weight_part_1.to(weight.dtype) * (2**6)
248
+
249
+ fp4_weight2 = fp4_weight << 4
250
+ fp4_weight_part_2 = ((fp4_weight2 & 0b10000000) |
251
+ ((fp4_weight2 & 0b01110000) >> 2))
252
+ fp4_weight_part_2 = fp4_weight_part_2.view(torch.float8_e4m3fn)
253
+ fp4_weight_part_2 = fp4_weight_part_2.to(weight.dtype) * (2**6)
254
+
255
+ weight_ref = torch.cat(
256
+ [fp4_weight_part_2.unsqueeze(2),
257
+ fp4_weight_part_1.unsqueeze(2)], 2).view(size_n, size_k)
258
+ weight_ref = weight_ref * global_scale.to(weight.dtype) * \
259
+ scales.repeat_interleave(group_size, 1).to(weight.dtype)
260
+
261
+ marlin_qweight = ops.gptq_marlin_repack(
262
+ b_q_weight=fp4_weight.view(torch.int32).T.contiguous(),
263
+ perm=torch.empty(0, dtype=torch.int, device=device),
264
+ size_k=size_k,
265
+ size_n=size_n,
266
+ num_bits=4,
267
+ )
268
+
269
+ marlin_scales = marlin_permute_scales(s=scales.T.to(weight.dtype),
270
+ size_k=size_k,
271
+ size_n=size_n,
272
+ group_size=group_size)
273
+ marlin_scales = fp4_marlin_process_scales(marlin_scales)
274
+
275
+ global_scale = fp4_marlin_process_global_scale(global_scale)
276
+
277
+ return weight_ref.T, marlin_qweight, marlin_scales, global_scale
vllm/model_executor/layers/quantization/utils/marlin_utils_fp8.py
+24
-2
@@ -19,6 +19,20 @@ def is_fp8_marlin_supported():
19
return current_platform.has_device_capability(80)
20
21
22
+def fp8_fused_exponent_bias_into_scales(scales):
23
+ fp8_exponent = 4
24
+ if scales.dtype == torch.half:
25
+ target_exponent = 5
26
+ elif scales.dtype == torch.bfloat16:
27
+ target_exponent = 8
28
+ # exponent_bias_fp16 = 2 ** 4 - 2 ** 3 = 8
29
+ # exponent_bias_bf16 = 2 ** 7 - 2 ** 3 = 120
30
+ exponent_bias = 2**(target_exponent - 1) - 2**(fp8_exponent - 1)
31
+ s = torch.ones_like(scales) * 2
32
+ s = s**exponent_bias
33
+ return scales * s
34
+
35
+
36
def apply_fp8_marlin_linear(
37
input: torch.Tensor,
38
weight: torch.Tensor,
@@ -44,6 +58,7 @@ def apply_fp8_marlin_linear(
58
c=None,
59
b_q_weight=weight,
60
b_scales=weight_scale,
61
+ global_scale=None,
62
b_zeros=None,
63
g_idx=None,
64
perm=None,
@@ -132,8 +147,10 @@ def prepare_fp8_layer_for_marlin(layer: torch.nn.Module,
147
# block-wise quantization -> group-wise quantization
148
# (size_k // block_size[1], ceil(size_n / block_size[0]))
149
# =>(repeat)=> (size_k // block_size[1], size_n)
150
+ if not size_k_first:
151
+ scales = scales.T.contiguous()
152
block_n = layer.weight_block_size[0]
136
- scales = scales.T.repeat_interleave(block_n, 1)
153
+ scales = scales.repeat_interleave(block_n, 1)
154
# size_n may not divisible by block_size[0]
155
scales = scales[:, :part_size_n]
156
@@ -141,6 +158,7 @@ def prepare_fp8_layer_for_marlin(layer: torch.nn.Module,
158
size_k=part_size_k,
159
size_n=part_size_n,
160
group_size=group_size)
161
+ marlin_scales = fp8_fused_exponent_bias_into_scales(marlin_scales)
162
layer.weight_scale = torch.nn.Parameter(marlin_scales, requires_grad=False)
163
164
@@ -239,8 +257,10 @@ def prepare_moe_fp8_layer_for_marlin(layer: torch.nn.Module,
257
# block-wise quantization -> group-wise quantization
258
# (e, size_k // block_size[1], ceil(size_n / block_size[0]))
259
# =>(repeat)=> (e, size_k // block_size[1], size_n)
260
+ if not size_k_first:
261
+ scales = scales.permute(0, 2, 1)
262
block_n = layer.weight_block_size[0]
243
- scales = scales.permute(0, 2, 1).repeat_interleave(block_n, 2)
263
+ scales = scales.repeat_interleave(block_n, 2)
264
# size_n may not divisible by block_size[0]
265
scales = scales[..., :size_n].contiguous()
266
@@ -302,4 +322,6 @@ def marlin_quant_fp8_torch(weight, group_size):
322
size_n=size_n,
323
group_size=group_size)
324
325
+ marlin_scales = fp8_fused_exponent_bias_into_scales(marlin_scales)
326
+
327
return weight_ref.T, marlin_qweight, marlin_scales