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 14, 05:30 UTC
[Kernel] [Quantization] Add MXFP4 and bias support for marlin kernel (#22428)
Signed-off-by: rongfu.leng <rongfu.leng@daocloud.io> Signed-off-by: Jinzhen Lin <linjinzhen@hotmail.com> Signed-off-by: Huzaifa Sidhpurwala <huzaifas@redhat.com> Signed-off-by: Varun Sundar Rabindranath <vsundarr@redhat.com> Signed-off-by: Harry Mellor <19981378+hmellor@users.noreply.github.com> Signed-off-by: Jee Jee Li <pandaleefree@gmail.com> Signed-off-by: mgoin <mgoin64@gmail.com> Signed-off-by: Animesh Jain <anijain@umich.edu> Signed-off-by: Rui Qiao <ruisearch42@gmail.com> Signed-off-by: Xiongfei Wei <isaacwxf23@gmail.com> Signed-off-by: Nick Hill <nhill@redhat.com> Signed-off-by: yewentao256 <zhyanwentao@126.com> Signed-off-by: kf <kuanfu.liu@embeddedllm.com> Signed-off-by: vllmellm <vllm.ellm@embeddedllm.com> Signed-off-by: NickLucche <nlucches@redhat.com> Signed-off-by: Dipika Sikka <dipikasikka1@gmail.com> Signed-off-by: Sage Moore <sage@neuralmagic.com> Signed-off-by: tjtanaavllm <tunjian.tan@amd.com> Signed-off-by: Yong Hoon Shin <yhshin@meta.com> Signed-off-by: Chih-Chieh-Yang <7364402+cyang49@users.noreply.github.com> Signed-off-by: Roger Wang <hey@rogerw.me> Signed-off-by: Vadim Gimpelson <vadim.gimpelson@centml.ai> Signed-off-by: Isotr0py <2037008807@qq.com> Signed-off-by: zRzRzRzRzRzRzR <2448370773@qq.com> Signed-off-by: Chih-Chieh Yang <7364402+cyang49@users.noreply.github.com> Signed-off-by: DarkLight1337 <tlleungac@connect.ust.hk> Signed-off-by: Thomas Parnell <tpa@zurich.ibm.com> Signed-off-by: yan <yan.ma@intel.com> Signed-off-by: Yan Ma <yan.ma@intel.com> Signed-off-by: Xiao Liu <xiszishu@gmail.com> Signed-off-by: jiahanc <173873397+jiahanc@users.noreply.github.com> Signed-off-by: Isotr0py <mozf@mail2.sysu.edu.cn> Signed-off-by: Ye (Charlotte) Qi <yeq@meta.com> Signed-off-by: LopezCastroRoberto <roberto.lopez.castro@udc.es> Signed-off-by: Andy Xie <andy.xning@gmail.com> Signed-off-by: Haibin Lin <haibin.lin@bytedance.com> Signed-off-by: David Ben-David <davidb@pliops.com> Signed-off-by: Woosuk Kwon <woosuk.kwon@berkeley.edu> Signed-off-by: jiang1.li <jiang1.li@intel.com> Signed-off-by: Seiji Eicher <seiji@anyscale.com> Signed-off-by: zitian.zhao <zitian.zhao@tencentmusic.com> Signed-off-by: 22quinn <33176974+22quinn@users.noreply.github.com> Signed-off-by: Abirdcfly <fp544037857@gmail.com> Signed-off-by: Giancarlo Delfin <gdelfin@meta.com> Signed-off-by: Tyler Michael Smith <tyler@neuralmagic.com> Signed-off-by: huangweixiao <huangweixiao@msh.team> Signed-off-by: alyosha-swamy <raghav@arcee.ai> Signed-off-by: Eric Hanley <ericehanley@google.com> Signed-off-by: Abatom <abzhonghua@gmail.com> Signed-off-by: CLFutureX <775523362@qq.com> Signed-off-by: Linkun Chen <github@lkchen.net> Signed-off-by: tjtanaa <tunjian.tan@embeddedllm.com> Signed-off-by: Gregory Shtrasberg <Gregory.Shtrasberg@amd.com> Signed-off-by: tlipoca9 <tlipoca9@gmail.com> Signed-off-by: elvischenv <219235043+elvischenv@users.noreply.github.com> Signed-off-by: zitian zhao <zitian.zhao@tencentmusic.com> Signed-off-by: mgoin <michael@neuralmagic.com> Signed-off-by: wang.yuqi <noooop@126.com> Signed-off-by: Benji Beck <benjibeck@meta.com> Signed-off-by: Siyuan Liu <lsiyuan@google.com> Signed-off-by: Benjamin Chislett <benjamin.chislett@centml.ai> Signed-off-by: isotr0py <2037008807@qq.com> Signed-off-by: Chen Zhang <zhangch99@outlook.com> Signed-off-by: simon-mo <xmo@berkeley.edu> Signed-off-by: LucasWilkinson <lwilkinson@neuralmagic.com> Signed-off-by: Zhang Jason <ning.zhang2@amd.com> Signed-off-by: Yongye Zhu <zyy1102000@gmail.com> Signed-off-by: asafg <asafg@ai21.com> Signed-off-by: Siyuan Fu <siyuanf@nvidia.com> Signed-off-by: Lain <fusiyuan2000@hotmail.com> Signed-off-by: Max de Bayser <mbayser@br.ibm.com> Signed-off-by: Lucas Wilkinson <lwilkins@redhat.com> Signed-off-by: Kunshang Ji <kunshang.ji@intel.com> Signed-off-by: Tao He <linzhu.ht@alibaba-inc.com> Signed-off-by: Michael Goin <mgoin64@gmail.com> Signed-off-by: QscQ <qscqesze@gmail.com> Signed-off-by: qingjun <qingjun@minimaxi.com> Signed-off-by: Syed Muhammad Bin Asif <syedmba7@connect.hku.hk> Signed-off-by: Lionel Villard <villard@us.ibm.com> Signed-off-by: ycyaw66 <497410282@qq.com> Signed-off-by: David Chen <530634352@qq.com> Signed-off-by: Linkun <github@lkchen.net> Signed-off-by: Moritz Sanft <58110325+msanft@users.noreply.github.com> Signed-off-by: Ming Yang <minos.future@gmail.com> Signed-off-by: Adrian Garcia <adrian.garcia@inceptionai.ai> Signed-off-by: shaojunqi <shaojunqi.sjq@alibaba-inc.com> Signed-off-by: Ricardo Decal <rdecal@anyscale.com> Signed-off-by: Andrew Chan <andrewkchan.akc@gmail.com> Signed-off-by: Felix Marty <Felix.Marty@amd.com> Signed-off-by: Andrew Sansom <andrew@protopia.ai> Signed-off-by: Zhiyu Cheng <zhiyuc@nvidia.com> Signed-off-by: Shu Wang <shuw@nvidia.com> Signed-off-by: Po-Han Huang <pohanh@nvidia.com> Signed-off-by: Shu Wang. <shuw@nvidia.com> Signed-off-by: XIn Li <xinli@nvidia.com> Signed-off-by: Junhao Li <junhao@ubicloud.com> Signed-off-by: chaunceyjiang <chaunceyjiang@gmail.com> Signed-off-by: iAmir97 <Amir.balwel@embeddedllm.com> Signed-off-by: iAmir97 <71513472+iAmir97@users.noreply.github.com> Signed-off-by: <zyy1102000@gmail.com> Signed-off-by: Guy Stone <guys@spotify.com> Signed-off-by: <yyweiss@gmail.com> Signed-off-by: yyw <yyweiss@gmail.com> Signed-off-by: Russell Bryant <rbryant@redhat.com> Signed-off-by: Pradyun Ramadorai <pradyunr@amazon.com> Signed-off-by: Pradyun92 <142861237+Pradyun92@users.noreply.github.com> Signed-off-by: Jinzhen Lin <jinzhen.ljz@antgroup.com> Co-authored-by: rongfu.leng <rongfu.leng@daocloud.io> Co-authored-by: Huzaifa Sidhpurwala <huzaifas@redhat.com> Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com> Co-authored-by: Russell Bryant <rbryant@redhat.com> Co-authored-by: Varun Sundar Rabindranath <varunsundar08@gmail.com> Co-authored-by: Varun Sundar Rabindranath <vsundarr@redhat.com> Co-authored-by: Harry Mellor <19981378+hmellor@users.noreply.github.com> Co-authored-by: Jee Jee Li <pandaleefree@gmail.com> Co-authored-by: Michael Goin <mgoin64@gmail.com> Co-authored-by: Animesh Jain <jainanimesh2305@yahoo.com> Co-authored-by: Rui Qiao <161574667+ruisearch42@users.noreply.github.com> Co-authored-by: XiongfeiWei <isaacwxf23@gmail.com> Co-authored-by: Nick Hill <nhill@redhat.com> Co-authored-by: Wentao Ye <44945378+yewentao256@users.noreply.github.com> Co-authored-by: JartX <sagformas@gmail.com> Co-authored-by: fhl2000 <63384265+fhl2000@users.noreply.github.com> Co-authored-by: vllmellm <vllm.ellm@embeddedllm.com> Co-authored-by: kf <kuanfu.liu@embeddedllm.com> Co-authored-by: Nicolò Lucchesi <nlucches@redhat.com> Co-authored-by: Dipika Sikka <dipikasikka1@gmail.com> Co-authored-by: Sage Moore <sage@neuralmagic.com> Co-authored-by: tjtanaavllm <tunjian.tan@amd.com> Co-authored-by: Yong Hoon Shin <48474650+sarckk@users.noreply.github.com> Co-authored-by: Chih-Chieh Yang <7364402+cyang49@users.noreply.github.com> Co-authored-by: Roger Wang <hey@rogerw.me> Co-authored-by: Vadim Gimpelson <156319763+vadiklyutiy@users.noreply.github.com> Co-authored-by: Yuxuan Zhang <2448370773@qq.com> Co-authored-by: Isotr0py <2037008807@qq.com> Co-authored-by: Cyrus Leung <tlleungac@connect.ust.hk> Co-authored-by: Thomas Parnell <tpa@zurich.ibm.com> Co-authored-by: Yan Ma <yan.ma@intel.com> Co-authored-by: Xiao <xiszishu@gmail.com> Co-authored-by: jiahanc <173873397+jiahanc@users.noreply.github.com> Co-authored-by: Isotr0py <mozf@mail2.sysu.edu.cn> Co-authored-by: Ye (Charlotte) Qi <yeq@meta.com> Co-authored-by: Roberto L. Castro <38211239+LopezCastroRoberto@users.noreply.github.com> Co-authored-by: Ning Xie <andy.xning@gmail.com> Co-authored-by: H <linhaibin.eric@gmail.com> Co-authored-by: David Ben-David <sdavidbd@gmail.com> Co-authored-by: David Ben-David <davidb@pliops.com> Co-authored-by: Woosuk Kwon <woosuk.kwon@berkeley.edu> Co-authored-by: Li, Jiang <jiang1.li@intel.com> Co-authored-by: TankNee <nee@tanknee.cn> Co-authored-by: Cyrus Leung <cyrus.tl.leung@gmail.com> Co-authored-by: Seiji Eicher <58963096+eicherseiji@users.noreply.github.com> Co-authored-by: ZiTian.Zhao <zitian.zhao@tencentmusic.com> Co-authored-by: 22quinn <33176974+22quinn@users.noreply.github.com> Co-authored-by: Abirdcfly <fp544037857@gmail.com> Co-authored-by: Giancarlo Delfin <32987265+TheEpicDolphin@users.noreply.github.com> Co-authored-by: Chenxi Yang <cxyang@cs.utexas.edu> Co-authored-by: Chenxi Yang <cxyang@meta.com> Co-authored-by: Tyler Michael Smith <tyler@neuralmagic.com> Co-authored-by: Weixiao Huang <hwx.simle@gmail.com> Co-authored-by: Raghav Ravishankar <113712354+alyosha-swamy@users.noreply.github.com> Co-authored-by: ericehanley <ericehanley@google.com> Co-authored-by: Zhonghua Deng <abzhonghua@gmail.com> Co-authored-by: Po-Han Huang (NVIDIA) <53919306+nvpohanh@users.noreply.github.com> Co-authored-by: PiteXChen <44110731+CLFutureX@users.noreply.github.com> Co-authored-by: lkchen <github@lkchen.net> Co-authored-by: TJian <tunjian.tan@embeddedllm.com> Co-authored-by: Gregory Shtrasberg <156009573+gshtras@users.noreply.github.com> Co-authored-by: tlipoca9 <160737620+tlipoca9@users.noreply.github.com> Co-authored-by: elvischenv <219235043+elvischenv@users.noreply.github.com> Co-authored-by: wang.yuqi <noooop@126.com> Co-authored-by: Benji Beck <benjibeck@meta.com> Co-authored-by: youkaichao <youkaichao@gmail.com> Co-authored-by: Siyuan Liu <lsiyuan@google.com> Co-authored-by: Benjamin Chislett <chislett.ben@gmail.com> Co-authored-by: LiuXiaoxuanPKU <lilyliupku@gmail.com> Co-authored-by: simon-mo <xmo@berkeley.edu> Co-authored-by: Chen Zhang <zhangch99@outlook.com> Co-authored-by: Hongxia Yang <62075498+hongxiayang@users.noreply.github.com> Co-authored-by: Minseok Lee <47620120+minseokl@users.noreply.github.com> Co-authored-by: Yongye Zhu <zyy1102000@gmail.com> Co-authored-by: Lucas Wilkinson <LucasWilkinson@users.noreply.github.com> Co-authored-by: Zhang Jason <ning.zhang2@amd.com> Co-authored-by: Asaf Joseph Gardin <39553475+Josephasafg@users.noreply.github.com> Co-authored-by: asafg <asafg@ai21.com> Co-authored-by: Lain <siyuanf@nvidia.com> Co-authored-by: tc-mb <157115220+tc-mb@users.noreply.github.com> Co-authored-by: imning3 <hbning@pku.edu.cn> Co-authored-by: Maximilien de Bayser <mbayser@br.ibm.com> Co-authored-by: Kunshang Ji <kunshang.ji@intel.com> Co-authored-by: Tao He <linzhu.ht@alibaba-inc.com> Co-authored-by: qscqesze <qingjun@minimaxi.com> Co-authored-by: Syed Muhammad Bin Asif <92625830+syedmba@users.noreply.github.com> Co-authored-by: Lionel Villard <villard@us.ibm.com> Co-authored-by: WeiQing Chen <40507679+david6666666@users.noreply.github.com> Co-authored-by: ycyaw66 <497410282@qq.com> Co-authored-by: Moritz Sanft <58110325+msanft@users.noreply.github.com> Co-authored-by: Ming Yang <minos.future@gmail.com> Co-authored-by: Adrián García García <adrigarvk8@gmail.com> Co-authored-by: Michael Goin <mgoin@redhat.com> Co-authored-by: JaceyShao <65159281+JaceyShao@users.noreply.github.com> Co-authored-by: shaojunqi <shaojunqi.sjq@alibaba-inc.com> Co-authored-by: Ricardo Decal <crypdick@users.noreply.github.com> Co-authored-by: Andrew Chan <andrewkchan.akc@gmail.com> Co-authored-by: fxmarty-amd <felmarty@amd.com> Co-authored-by: Andrew Sansom <andrew@protopia.ai> Co-authored-by: Zhiyu <zhiyuc@nvidia.com> Co-authored-by: Shu Wang <shuw@nvidia.com> Co-authored-by: XIn Li <xinli@nvidia.com> Co-authored-by: Junhao Li <streaver91@gmail.com> Co-authored-by: Chauncey <chaunceyjiang@gmail.com> Co-authored-by: iAmir97 <71513472+iAmir97@users.noreply.github.com> Co-authored-by: iAmir97 <Amir.balwel@embeddedllm.com> Co-authored-by: Hong Hanh <hanh.usth@gmail.com> Co-authored-by: Daniel Serebrenik <74646983+pliops-daniels@users.noreply.github.com> Co-authored-by: yewentao256 <zhyanwentao@126.com> Co-authored-by: Guy Stone <guys@spotify.com> Co-authored-by: yyweiss <70619747+yyweiss@users.noreply.github.com> Co-authored-by: Pradyun92 <142861237+Pradyun92@users.noreply.github.com> Co-authored-by: Pradyun Ramadorai <pradyunr@amazon.com> Co-authored-by: Nicolò Lucchesi <nicolo.lucchesi@gmail.com>
Jinzhen Lin committed
Aug 15, 2025 at 02:23 UTC
33c63e95474f9354b362e72d447c225d1c64dfc2
34 files changed
+1126
-322
CMakeLists.txt
+7
@@ -351,6 +351,8 @@ if(VLLM_GPU_LANG STREQUAL "CUDA")
351
set_gencode_flags_for_srcs(
352
SRCS "${MARLIN_TEMPLATE_KERNEL_SRC}"
353
CUDA_ARCHS "${MARLIN_ARCHS}")
354
+ set_source_files_properties(${MARLIN_TEMPLATE_KERNEL_SRC}
355
+ PROPERTIES COMPILE_FLAGS "-static-global-template-stub=false")
356
357
list(APPEND VLLM_EXT_SRC ${MARLIN_TEMPLATE_KERNEL_SRC})
358
@@ -364,7 +366,10 @@ if(VLLM_GPU_LANG STREQUAL "CUDA")
366
set_gencode_flags_for_srcs(
367
SRCS "${MARLIN_SRCS}"
368
CUDA_ARCHS "${MARLIN_ARCHS}")
369
+ set_source_files_properties("csrc/quantization/gptq_marlin/gptq_marlin.cu"
370
+ PROPERTIES COMPILE_FLAGS "-static-global-template-stub=false")
371
list(APPEND VLLM_EXT_SRC "${MARLIN_SRCS}")
372
+
373
message(STATUS "Building Marlin kernels for archs: ${MARLIN_ARCHS}")
374
else()
375
message(STATUS "Not building Marlin kernels as no compatible archs found"
@@ -854,6 +859,8 @@ if(VLLM_GPU_LANG STREQUAL "CUDA")
859
set_gencode_flags_for_srcs(
860
SRCS "${MOE_WNAA16_MARLIN_SRC}"
861
CUDA_ARCHS "${MARLIN_MOE_ARCHS}")
862
+ set_source_files_properties(${MOE_WNAA16_MARLIN_SRC}
863
+ PROPERTIES COMPILE_FLAGS "-static-global-template-stub=false")
864
865
list(APPEND VLLM_MOE_EXT_SRC ${MOE_WNAA16_MARLIN_SRC})
866
benchmarks/kernels/benchmark_machete.py
+1
@@ -236,6 +236,7 @@ def marlin_create_bench_fn(bt: BenchmarkTensors) -> Callable:
236
a=bt.a,
237
c=None,
238
b_q_weight=w_q,
239
+ b_bias=None,
240
b_scales=w_s,
241
global_scale=None,
242
b_zeros=w_zp,
csrc/core/scalar_type.hpp
+2
@@ -321,6 +321,8 @@ static inline constexpr auto kFE3M2f =
321
ScalarType::float_(3, 2, true, ScalarType::NAN_NONE);
322
static inline constexpr auto kFE4M3fn =
323
ScalarType::float_(4, 3, true, ScalarType::NAN_EXTD_RANGE_MAX_MIN);
324
+static inline constexpr auto kFE8M0fnu =
325
+ ScalarType(8, 0, false, 0, true, ScalarType::NAN_EXTD_RANGE_MAX_MIN);
326
static inline constexpr auto kFE5M2 = ScalarType::float_IEEE754(5, 2);
327
static inline constexpr auto kFE8M7 = ScalarType::float_IEEE754(8, 7);
328
static inline constexpr auto kFE5M10 = ScalarType::float_IEEE754(5, 10);
csrc/moe/marlin_moe_wna16/generate_kernels.py
+15
@@ -20,6 +20,7 @@ namespace MARLIN_NAMESPACE_NAME {
20
TEMPLATE = ("template __global__ void Marlin<"
21
"{{scalar_t}}, "
22
"{{w_type_id}}, "
23
+ "{{s_type_id}}, "
24
"{{threads}}, "
25
"{{thread_m_blocks}}, "
26
"{{thread_n_blocks}}, "
@@ -77,6 +78,7 @@ def generate_new_kernels():
78
if scalar_type == "vllm::kFE4M3fn" and group_blocks not in [-1, 8]:
79
continue
80
# nvfp4 only supports group_size == 16
81
+ # mxfp4 only supports group_size == 32
82
if scalar_type == "vllm::kFE2M1f" and group_blocks not in [1, 2]:
83
continue
84
# other quantization methods don't support group_size = 16
@@ -89,9 +91,22 @@ def generate_new_kernels():
91
92
c_dtype = "half" if dtype == "fp16" else "nv_bfloat16"
93
94
+ if scalar_type == "vllm::kFE2M1f" and group_blocks == 1:
95
+ s_type = "vllm::kFE4M3fn"
96
+ elif scalar_type == "vllm::kFE2M1f" and group_blocks == 2:
97
+ s_type = "vllm::kFE8M0fnu"
98
+ if dtype == "fp16":
99
+ # we cannot safely dequantize e8m0 to fp16, so skip this
100
+ continue
101
+ elif dtype == "fp16":
102
+ s_type = "vllm::kFloat16"
103
+ elif dtype == "bf16":
104
+ s_type = "vllm::kBFloat16"
105
+
106
template_str = jinja2.Template(TEMPLATE).render(
107
scalar_t=c_dtype,
108
w_type_id=scalar_type + ".id()",
109
+ s_type_id=s_type + ".id()",
110
threads=threads,
111
thread_m_blocks=max(m_blocks, 1),
112
thread_n_blocks=n_blocks,
csrc/moe/marlin_moe_wna16/kernel.h
+14
-12
@@ -7,23 +7,25 @@
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, \
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, \
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__ b_bias_ptr, \
14
+ const int4 *__restrict__ scales_ptr, \
15
+ const uint16_t *__restrict__ scale2_ptr, \
16
+ const int4 *__restrict__ zp_ptr, const int *__restrict__ g_idx, \
17
+ const int32_t *__restrict__ sorted_token_ids_ptr, \
18
+ const int32_t *__restrict__ expert_ids_ptr, \
19
+ const int32_t *__restrict__ num_tokens_past_padded_ptr, \
20
+ const float *__restrict__ topk_weights_ptr, int top_k, \
21
+ bool mul_topk_weights, bool is_ep, int num_groups, int prob_m, \
22
+ int prob_n, int prob_k, int *locks, bool has_bias, bool use_atomic_add, \
23
bool use_fp32_reduce, int max_shared_mem
24
25
namespace MARLIN_NAMESPACE_NAME {
26
template <typename scalar_t, // compute dtype, half or nv_float16
27
const vllm::ScalarTypeId w_type_id, // weight ScalarType id
28
+ const vllm::ScalarTypeId s_type_id, // weight scale ScalarType id
29
const int threads, // number of threads in a threadblock
30
const int thread_m_blocks, // number of 16x16 blocks in the m
31
// dimension (batchsize) of the
csrc/moe/marlin_moe_wna16/marlin_template.h
+111
-26
@@ -280,6 +280,7 @@ __device__ inline void wait_negative_and_add(int* lock) {
280
281
template <typename scalar_t, // compute dtype, half or nv_float16
282
const vllm::ScalarTypeId w_type_id, // weight ScalarType id
283
+ const vllm::ScalarTypeId s_type_id, // weight scale ScalarType id
284
const int threads, // number of threads in a threadblock
285
const int thread_m_blocks, // number of 16x16 blocks in the m
286
// dimension (batchsize) of the
@@ -299,6 +300,7 @@ __global__ void Marlin(
300
const int4* __restrict__ B, // 4bit quantized weight matrix of shape kxn
301
int4* __restrict__ C, // fp16 output buffer of shape mxn
302
int4* __restrict__ C_tmp, // fp32 tmp output buffer (for reduce)
303
+ const int4* __restrict__ b_bias_ptr,
304
const int4* __restrict__ scales_ptr, // fp16 quantization scales of shape
305
// (k/groupsize)xn
306
const uint16_t* __restrict__ scale2_ptr, // fp16 global scale (for nvfp4
@@ -318,8 +320,9 @@ __global__ void Marlin(
320
int prob_n, // output dimension n
321
int prob_k, // reduction dimension k
322
int* locks, // extra global storage for barrier synchronization
321
- bool use_atomic_add, // whether to use atomic add to reduce
322
- bool use_fp32_reduce, // whether to use fp32 global reduce
323
+ bool has_bias,
324
+ bool use_atomic_add, // whether to use atomic add to reduce
325
+ bool use_fp32_reduce, // whether to use fp32 global reduce
326
int max_shared_mem) {
327
// Each threadblock processes one "stripe" of the B matrix with (roughly) the
328
// same size, which might involve multiple column "slices" (of width 16 *
@@ -342,12 +345,23 @@ __global__ void Marlin(
345
346
extern __shared__ int4 sh[];
347
static constexpr auto w_type = vllm::ScalarType::from_id(w_type_id);
348
+ static constexpr auto s_type = vllm::ScalarType::from_id(s_type_id);
349
+ if constexpr (w_type == vllm::kFE2M1f) {
350
+ static_assert(s_type == vllm::kFE4M3fn && group_blocks == 1 ||
351
+ s_type == vllm::kFE8M0fnu && group_blocks == 2);
352
+ } else if constexpr (std::is_same<scalar_t, nv_bfloat16>::value) {
353
+ static_assert(s_type == vllm::kBFloat16);
354
+ } else if constexpr (std::is_same<scalar_t, half>::value) {
355
+ static_assert(s_type == vllm::kFloat16);
356
+ }
357
+
358
constexpr bool has_zp = w_type == vllm::kU4 || w_type == vllm::kU8;
359
constexpr bool is_int_type = w_type == vllm::kU4 || w_type == vllm::kU8 ||
360
w_type == vllm::kU4B8 || w_type == vllm::kU8B128;
361
// see comments of dequant.h for more details
362
constexpr bool dequant_skip_flop =
350
- !is_int_type ||
363
+ w_type == vllm::kFE4M3fn ||
364
+ w_type == vllm::kFE2M1f && s_type == vllm::kFE4M3fn ||
365
has_zp && !is_zp_float && !std::is_same<scalar_t, nv_bfloat16>::value ||
366
has_zp && !is_zp_float && !(w_type == vllm::kU8);
367
@@ -365,6 +379,7 @@ __global__ void Marlin(
379
const int zp_expert_stride =
380
is_zp_float ? prob_n * prob_k / group_size / 8
381
: prob_n * prob_k / group_size / (pack_factor * 4);
382
+ const int b_bias_expert_stride = prob_n / 8;
383
384
// parallel: num valid moe blocks
385
int num_tokens_past_padded = num_tokens_past_padded_ptr[0];
@@ -475,7 +490,7 @@ __global__ void Marlin(
490
for (int i = 0; i < 4; i++) {
491
int idx = tid4 * 4 + i;
492
idx = idx < block_num_valid_tokens ? idx : 0;
478
- if constexpr (w_type == vllm::kFE2M1f) {
493
+ if constexpr (w_type == vllm::kFE2M1f && s_type == vllm::kFE4M3fn) {
494
sh_block_topk_weights[idx] = __hmul2(
495
global_scale, Dtype::num2num2(Dtype::float2num(
496
topk_weights_ptr[sh_block_sorted_ids[idx]])));
@@ -513,7 +528,7 @@ __global__ void Marlin(
528
expert_id = expert_ids_ptr[block_id];
529
}
530
516
- if constexpr (w_type == vllm::kFE2M1f) {
531
+ if constexpr (w_type == vllm::kFE2M1f && s_type == vllm::kFE4M3fn) {
532
uint16_t val = scale2_ptr[expert_id];
533
global_scale = Dtype::num2num2(*reinterpret_cast<scalar_t*>(&val));
534
}
@@ -526,6 +541,9 @@ __global__ void Marlin(
541
if constexpr (has_act_order) {
542
g_idx += (expert_id - old_expert_id) * prob_k;
543
}
544
+ if (has_bias) {
545
+ b_bias_ptr += (expert_id - old_expert_id) * b_bias_expert_stride;
546
+ }
547
548
read_moe_block_data(block_id);
549
};
@@ -721,7 +739,7 @@ __global__ void Marlin(
739
740
s_sh_rd = 8 * ((threadIdx.x / 32) % (thread_n_blocks / 4)) +
741
(threadIdx.x % 32) / 4;
724
- s_sh_rd = s_sh_rd * 2 + warp_row % 2;
742
+ s_sh_rd = s_sh_rd * 2 + (warp_row / group_blocks) % 2;
743
744
} else if constexpr (group_blocks != -1)
745
s_sh_rd = 8 * ((threadIdx.x / 32) % (thread_n_blocks / 4)) +
@@ -734,6 +752,18 @@ __global__ void Marlin(
752
s_sh_rd = 8 * ((threadIdx.x / 32) % (thread_n_blocks / 4)) +
753
(threadIdx.x % 32) % 4;
754
755
+ int bias_sh_rd;
756
+ if constexpr (m_block_size_8) {
757
+ bias_sh_rd = 8 * ((threadIdx.x / 32) % (thread_n_blocks / 4)) +
758
+ (threadIdx.x % 32) / 8;
759
+ } else {
760
+ bias_sh_rd = 8 * ((threadIdx.x / 32) % (thread_n_blocks / 4)) +
761
+ (threadIdx.x % 32) % 4;
762
+ }
763
+
764
+ int bias_sh_wr = threadIdx.x;
765
+ int bias_gl_rd = (thread_n_blocks * 16 / 8) * slice_col + threadIdx.x;
766
+
767
// Zero-points have the same read layout as the scales
768
// (without column-wise case)
769
constexpr int num_col_threads = 8;
@@ -793,7 +823,19 @@ __global__ void Marlin(
823
constexpr int sh_b_size = stages * b_sh_stage;
824
int4* sh_b = sh_new;
825
int4* sh_red = sh_new;
796
- int4* sh_g_idx = sh_b + (sh_red_size > sh_b_size ? sh_red_size : sh_b_size);
826
+
827
+ constexpr int sh_size_b_red_min =
828
+ (sh_red_size < sh_b_size ? sh_red_size : sh_b_size);
829
+ constexpr int sh_size_b_red_max =
830
+ (sh_red_size > sh_b_size ? sh_red_size : sh_b_size);
831
+ constexpr int sh_bias_size = (thread_n_blocks * 16 / 8);
832
+ constexpr int sh_b_red_bias_size =
833
+ sh_size_b_red_max > (sh_size_b_red_min + sh_bias_size)
834
+ ? sh_size_b_red_max
835
+ : (sh_size_b_red_min + sh_bias_size);
836
+
837
+ int4* sh_bias = sh_new + sh_size_b_red_min;
838
+ int4* sh_g_idx = sh_new + sh_b_red_bias_size;
839
int4* sh_zp = sh_g_idx + (stages * g_idx_stage);
840
constexpr int sh_s_size = has_act_order ? (act_s_max_num_groups * s_sh_stride)
841
: (stages * s_sh_stage);
@@ -803,9 +845,9 @@ __global__ void Marlin(
845
static_assert(thread_m_blocks * 16 * thread_n_blocks * 16 / 8 <=
846
stages * b_sh_stage);
847
int4* sh_a = sh_s + sh_s_size;
806
- constexpr int shm_size_used =
807
- moe_block_size + stages * (g_idx_stage + zp_sh_stage) + sh_s_size +
808
- (sh_red_size > sh_b_size ? sh_red_size : sh_b_size);
848
+ constexpr int shm_size_used = moe_block_size +
849
+ stages * (g_idx_stage + zp_sh_stage) +
850
+ sh_s_size + sh_b_red_bias_size;
851
852
// all remaining shared memory is used to cache A (input)
853
// sh_a_max_row is at least ` stages * 16 * thread_m_blocks `
@@ -816,7 +858,8 @@ __global__ void Marlin(
858
FragA frag_a[2][thread_m_blocks];
859
I4 frag_b_quant[2][b_thread_vecs];
860
FragC frag_c[thread_m_blocks][4][2];
819
- FragS frag_s[2][4]; // No act-order
861
+ FragS frag_s[2][4]; // No act-order
862
+ FragS frag_bias[2][4];
863
FragS act_frag_s[2][4][4]; // For act-order
864
int frag_qzp[2][num_ints_per_thread]; // Zero-points
865
FragZP frag_zp; // Zero-points in fp16
@@ -1065,10 +1108,15 @@ __global__ void Marlin(
1108
if constexpr (w_type_id != vllm::kFE2M1f.id()) {
1109
reinterpret_cast<int4*>(&frag_s[k % 2])[0] =
1110
sh_s_stage[s_sh_rd + cur_group_id * s_sh_stride];
1068
- } else {
1111
+ } else if constexpr (group_blocks == 1 || thread_k_blocks > 4) {
1112
reinterpret_cast<int2*>(&frag_s[k % 2])[0] =
1113
reinterpret_cast<int2*>(
1114
sh_s_stage)[s_sh_rd + cur_group_id * (2 * s_sh_stride)];
1115
+ } else {
1116
+ reinterpret_cast<int2*>(&frag_s[k % 2])[0] =
1117
+ reinterpret_cast<int2*>(
1118
+ sh_s_stage)[s_sh_rd + cur_group_id * (2 * s_sh_stride) +
1119
+ k % 2];
1120
}
1121
}
1122
}
@@ -1281,9 +1329,9 @@ __global__ void Marlin(
1329
int s_quant_0 = reinterpret_cast<int*>(frag_s[k2])[0];
1330
int s_quant_1 = reinterpret_cast<int*>(frag_s[k2])[1];
1331
1284
- dequant_fp8_scales<scalar_t2>(s_quant_0,
1285
- reinterpret_cast<scalar_t2*>(&frag_s[k2]));
1286
- dequant_fp8_scales<scalar_t2>(
1332
+ dequant_fp8_scales<scalar_t2, s_type_id>(
1333
+ s_quant_0, reinterpret_cast<scalar_t2*>(&frag_s[k2]));
1334
+ dequant_fp8_scales<scalar_t2, s_type_id>(
1335
s_quant_1, reinterpret_cast<scalar_t2*>(&frag_s[k2]) + 2);
1336
}
1337
@@ -1566,7 +1614,7 @@ __global__ void Marlin(
1614
// Write out the reduce final result in the correct layout. We only actually
1615
// reshuffle matrix fragments in this step, the reduction above is performed
1616
// in fragment layout.
1569
- auto write_result = [&]() {
1617
+ auto write_result = [&](bool last) {
1618
int c_gl_stride = prob_n / 8;
1619
constexpr int c_sh_stride = 2 * thread_n_blocks + 1;
1620
int c_gl_wr_delta = c_gl_stride * (threads / (2 * thread_n_blocks));
@@ -1592,7 +1640,7 @@ __global__ void Marlin(
1640
1641
// We first reorder in shared memory to guarantee the most efficient final
1642
// global write patterns
1595
- auto write = [&](int idx, float c0, float c1, FragS& s) {
1643
+ auto write = [&](int idx, float c0, float c1, FragS& s, FragS& b_bias) {
1644
scalar_t2 res =
1645
Dtype::nums2num2(Dtype::float2num(c0), Dtype::float2num(c1));
1646
@@ -1601,14 +1649,27 @@ __global__ void Marlin(
1649
if constexpr (!has_act_order && group_blocks == -1 &&
1650
w_type.size_bits() == 4 &&
1651
(has_zp && dequant_skip_flop || !has_zp)) {
1604
- res = __hmul2(res, s[0]);
1652
+ scalar_t2 tmp_scale = s[0];
1653
+ if constexpr (m_block_size_8) {
1654
+ tmp_scale = Dtype::num2num2(
1655
+ reinterpret_cast<scalar_t*>(&s[0])[(threadIdx.x % 8) / 4]);
1656
+ }
1657
+ res = __hmul2(res, tmp_scale);
1658
}
1659
1607
- if constexpr (w_type == vllm::kFE2M1f) {
1660
+ if constexpr (w_type == vllm::kFE2M1f && s_type == vllm::kFE4M3fn) {
1661
if (!mul_topk_weights) {
1662
res = __hmul2(res, global_scale);
1663
}
1664
}
1665
+ if (has_bias && last) {
1666
+ scalar_t2 tmp_bias = b_bias[0];
1667
+ if constexpr (m_block_size_8) {
1668
+ tmp_bias = Dtype::num2num2(
1669
+ reinterpret_cast<scalar_t*>(&b_bias[0])[(threadIdx.x % 8) / 4]);
1670
+ }
1671
+ res = __hadd2(res, tmp_bias);
1672
+ }
1673
1674
if constexpr (m_block_size_8) {
1675
((scalar_t*)sh_red)[idx] = res.x;
@@ -1626,19 +1687,25 @@ __global__ void Marlin(
1687
if constexpr (m_block_size_8) {
1688
int wr = c_sh_wr + 16 * j;
1689
write(wr, frag_c[i][j][0][0], frag_c[i][j][0][1],
1629
- frag_s[j / 2][2 * (j % 2) + 0]);
1690
+ frag_s[j / 2][2 * (j % 2) + 0],
1691
+ frag_bias[j / 2][2 * (j % 2) + 0]);
1692
write(wr + 8, frag_c[i][j][0][2], frag_c[i][j][0][3],
1631
- frag_s[j / 2][2 * (j % 2) + 1]);
1693
+ frag_s[j / 2][2 * (j % 2) + 1],
1694
+ frag_bias[j / 2][2 * (j % 2) + 1]);
1695
} else {
1696
int wr = c_sh_wr + 8 * j;
1697
write(wr + (4 * c_sh_stride) * 0 + 0, frag_c[i][j][0][0],
1635
- frag_c[i][j][0][1], frag_s[j / 2][2 * (j % 2) + 0]);
1698
+ frag_c[i][j][0][1], frag_s[j / 2][2 * (j % 2) + 0],
1699
+ frag_bias[j / 2][2 * (j % 2) + 0]);
1700
write(wr + (4 * c_sh_stride) * 8 + 0, frag_c[i][j][0][2],
1637
- frag_c[i][j][0][3], frag_s[j / 2][2 * (j % 2) + 0]);
1701
+ frag_c[i][j][0][3], frag_s[j / 2][2 * (j % 2) + 0],
1702
+ frag_bias[j / 2][2 * (j % 2) + 0]);
1703
write(wr + (4 * c_sh_stride) * 0 + 4, frag_c[i][j][1][0],
1639
- frag_c[i][j][1][1], frag_s[j / 2][2 * (j % 2) + 1]);
1704
+ frag_c[i][j][1][1], frag_s[j / 2][2 * (j % 2) + 1],
1705
+ frag_bias[j / 2][2 * (j % 2) + 1]);
1706
write(wr + (4 * c_sh_stride) * 8 + 4, frag_c[i][j][1][2],
1641
- frag_c[i][j][1][3], frag_s[j / 2][2 * (j % 2) + 1]);
1707
+ frag_c[i][j][1][3], frag_s[j / 2][2 * (j % 2) + 1],
1708
+ frag_bias[j / 2][2 * (j % 2) + 1]);
1709
}
1710
}
1711
c_sh_wr += 16 * (4 * c_sh_stride);
@@ -1805,6 +1872,14 @@ __global__ void Marlin(
1872
}
1873
1874
thread_block_reduce();
1875
+
1876
+ if (has_bias && last) {
1877
+ __syncthreads();
1878
+ cp_async4_pred(&sh_bias[bias_sh_wr], &b_bias_ptr[bias_gl_rd],
1879
+ threadIdx.x < 16 * thread_n_blocks / 8);
1880
+ cp_async_fence();
1881
+ }
1882
+
1883
if constexpr (!has_act_order && group_blocks == -1 &&
1884
(has_zp && dequant_skip_flop || !has_zp)) {
1885
if (w_type.size_bits() == 8 || (last || use_atomic_add)) {
@@ -1867,11 +1942,20 @@ __global__ void Marlin(
1942
}
1943
barrier_release(&locks[locks_off], last);
1944
}
1945
+
1946
+ if (has_bias && last) {
1947
+ cp_async_wait<0>();
1948
+ __syncthreads();
1949
+ reinterpret_cast<int4*>(&frag_bias)[0] = sh_bias[bias_sh_rd];
1950
+ reinterpret_cast<int4*>(&frag_bias)[1] = sh_bias[bias_sh_rd + 4];
1951
+ __syncthreads();
1952
+ }
1953
+
1954
if (use_atomic_add && slice_count > 1 && slice_idx != 0)
1955
wait_negative_and_add(&locks[locks_off]);
1956
if (last || use_atomic_add)
1957
// only the last block in a slice actually writes the result
1874
- write_result();
1958
+ write_result(last);
1959
int old_slice_row = slice_row;
1960
slice_row = 0;
1961
slice_col_par++;
@@ -1904,6 +1988,7 @@ __global__ void Marlin(
1988
for (int i = 0; i < b_sh_wr_iters; i++) B_ptr[i] -= b_gl_stride;
1989
}
1990
1991
+ bias_gl_rd = (thread_n_blocks * 16 / 8) * slice_col + threadIdx.x;
1992
// Update slice k/n for scales loading
1993
if constexpr (has_act_order) {
1994
slice_k_start = tb_k * slice_row;
csrc/moe/marlin_moe_wna16/ops.cu
+120
-61
@@ -51,8 +51,9 @@ __global__ void permute_cols_kernel(
51
} // namespace marlin
52
53
torch::Tensor moe_wna16_marlin_gemm(
54
- torch::Tensor& a, std::optional<torch::Tensor> const& c_or_none,
55
- torch::Tensor& b_q_weight, torch::Tensor& b_scales,
54
+ torch::Tensor& a, std::optional<torch::Tensor> c_or_none,
55
+ torch::Tensor& b_q_weight,
56
+ std::optional<torch::Tensor> const& b_bias_or_none, torch::Tensor& b_scales,
57
std::optional<torch::Tensor> const& b_zeros_or_none,
58
std::optional<torch::Tensor> const& g_idx_or_none,
59
std::optional<torch::Tensor> const& perm_or_none, torch::Tensor& workspace,
@@ -212,7 +213,7 @@ int get_kernel_cache_size(thread_config_t const& th_config, bool m_block_size_8,
213
// Get B size
214
int tb_k = th_config.thread_k;
215
int tb_n = th_config.thread_n;
215
- int tb_m = thread_m_blocks * (m_block_size_8 ? 8 : 16);
216
+ int tb_m = thread_m_blocks * 16;
217
218
// shm size for block_sorted_ids/rd_block_sorted_ids/block_topk_weights
219
// both of them requires tb_m * 4 bytes (tb_m * int32 or tb_m * float32)
@@ -220,6 +221,11 @@ int get_kernel_cache_size(thread_config_t const& th_config, bool m_block_size_8,
221
int sh_a_size = pipe_stages * (tb_m * tb_k) * 2;
222
int sh_b_size = pipe_stages * (tb_k * tb_n / pack_factor) * 4;
223
int sh_red_size = tb_m * (tb_n + 8) * 2;
224
+ int sh_bias_size = tb_n * 2;
225
+ int tmp_size =
226
+ (sh_b_size > sh_red_size ? sh_red_size : sh_b_size) + sh_bias_size;
227
+ tmp_size = max(max(sh_b_size, sh_red_size), tmp_size);
228
+
229
int sh_s_size =
230
get_scales_cache_size(th_config, prob_m, prob_n, prob_k, num_bits,
231
group_size, has_act_order, is_k_full);
@@ -234,8 +240,8 @@ int get_kernel_cache_size(thread_config_t const& th_config, bool m_block_size_8,
240
sh_zp_size = sh_s_size / 2;
241
}
242
237
- int total_size = max(sh_b_size, sh_red_size) + sh_a_size + sh_s_size +
238
- sh_zp_size + sh_g_idx_size + sh_block_meta_size;
243
+ int total_size = tmp_size + sh_a_size + sh_s_size + sh_zp_size +
244
+ sh_g_idx_size + sh_block_meta_size;
245
246
return total_size;
247
}
@@ -270,20 +276,25 @@ bool is_valid_config(thread_config_t const& th_config, bool m_block_size_8,
276
int cache_size = get_kernel_cache_size(
277
th_config, m_block_size_8, thread_m_blocks, prob_m, prob_n, prob_k,
278
num_bits, group_size, has_act_order, is_k_full, has_zp, is_zp_float);
273
- return cache_size <= max_shared_mem;
279
+ return cache_size + 512 <= max_shared_mem;
280
}
281
276
- #define _GET_IF(W_TYPE, THREAD_M_BLOCKS, THREAD_N_BLOCKS, THREAD_K_BLOCKS, \
277
- M_BLOCK_SIZE_8, GROUP_BLOCKS, NUM_THREADS, IS_ZP_FLOAT) \
278
- else if (q_type == W_TYPE && thread_m_blocks == THREAD_M_BLOCKS && \
279
- thread_n_blocks == THREAD_N_BLOCKS && \
280
- thread_k_blocks == THREAD_K_BLOCKS && \
281
- m_block_size_8 == M_BLOCK_SIZE_8 && \
282
- group_blocks == GROUP_BLOCKS && num_threads == NUM_THREADS && \
283
- is_zp_float == IS_ZP_FLOAT) { \
284
- kernel = Marlin<scalar_t, W_TYPE.id(), NUM_THREADS, THREAD_M_BLOCKS, \
285
- THREAD_N_BLOCKS, THREAD_K_BLOCKS, M_BLOCK_SIZE_8, \
286
- pipe_stages, GROUP_BLOCKS, IS_ZP_FLOAT>; \
282
+ #define _GET_IF(W_TYPE, THREAD_M_BLOCKS, THREAD_N_BLOCKS, THREAD_K_BLOCKS, \
283
+ M_BLOCK_SIZE_8, GROUP_BLOCKS, NUM_THREADS, IS_ZP_FLOAT) \
284
+ else if (q_type == W_TYPE && thread_m_blocks == THREAD_M_BLOCKS && \
285
+ thread_n_blocks == THREAD_N_BLOCKS && \
286
+ thread_k_blocks == THREAD_K_BLOCKS && \
287
+ m_block_size_8 == M_BLOCK_SIZE_8 && \
288
+ group_blocks == GROUP_BLOCKS && num_threads == NUM_THREADS && \
289
+ is_zp_float == IS_ZP_FLOAT) { \
290
+ constexpr auto S_TYPE = \
291
+ W_TYPE == vllm::kFE2M1f \
292
+ ? (GROUP_BLOCKS == 1 ? vllm::kFE4M3fn : vllm::kFE8M0fnu) \
293
+ : (std::is_same<scalar_t, half>::value ? vllm::kFloat16 \
294
+ : vllm::kBFloat16); \
295
+ kernel = Marlin<scalar_t, W_TYPE.id(), S_TYPE.id(), NUM_THREADS, \
296
+ THREAD_M_BLOCKS, THREAD_N_BLOCKS, THREAD_K_BLOCKS, \
297
+ M_BLOCK_SIZE_8, pipe_stages, GROUP_BLOCKS, IS_ZP_FLOAT>; \
298
}
299
300
// COMMON: cases for (group_blocks in [-1, 2, 4, 8] and is_zp_float == false)
@@ -335,30 +346,44 @@ bool is_valid_config(thread_config_t const& th_config, bool m_block_size_8,
346
_GET_IF(W_TYPE, 2, N_BLOCKS, K_BLOCKS, false, 8, NUM_THREADS, false) \
347
_GET_IF(W_TYPE, 3, N_BLOCKS, K_BLOCKS, false, -1, NUM_THREADS, false) \
348
_GET_IF(W_TYPE, 3, N_BLOCKS, K_BLOCKS, false, 8, NUM_THREADS, false) \
338
- \
349
_GET_IF(W_TYPE, 4, N_BLOCKS, K_BLOCKS, false, -1, NUM_THREADS, false) \
350
_GET_IF(W_TYPE, 4, N_BLOCKS, K_BLOCKS, false, 8, NUM_THREADS, false)
351
342
- #define FP4_GET_IF_M1(W_TYPE, N_BLOCKS, K_BLOCKS, NUM_THREADS) \
352
+ #define BIGGROUP_GET_IF(W_TYPE) \
353
+ BIGGROUP_GET_IF_M1(W_TYPE, 8, 8, 256) \
354
+ BIGGROUP_GET_IF_M1(W_TYPE, 8, 4, 128) \
355
+ BIGGROUP_GET_IF_M234(W_TYPE, 16, 4, 256) \
356
+ BIGGROUP_GET_IF_M234(W_TYPE, 8, 4, 128)
357
+
358
+ #define NVFP4_GET_IF_M1(W_TYPE, N_BLOCKS, K_BLOCKS, NUM_THREADS) \
359
_GET_IF(W_TYPE, 1, N_BLOCKS, K_BLOCKS, true, 1, NUM_THREADS, false) \
360
_GET_IF(W_TYPE, 1, N_BLOCKS, K_BLOCKS, false, 1, NUM_THREADS, false)
361
346
- #define FP4_GET_IF_M234(W_TYPE, N_BLOCKS, K_BLOCKS, NUM_THREADS) \
362
+ #define NVFP4_GET_IF_M234(W_TYPE, N_BLOCKS, K_BLOCKS, NUM_THREADS) \
363
_GET_IF(W_TYPE, 2, N_BLOCKS, K_BLOCKS, false, 1, NUM_THREADS, false) \
364
_GET_IF(W_TYPE, 3, N_BLOCKS, K_BLOCKS, false, 1, NUM_THREADS, false) \
365
_GET_IF(W_TYPE, 4, N_BLOCKS, K_BLOCKS, false, 1, NUM_THREADS, false)
366
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)
367
+ #define NVFP4_GET_IF(W_TYPE) \
368
+ NVFP4_GET_IF_M1(W_TYPE, 8, 8, 256) \
369
+ NVFP4_GET_IF_M1(W_TYPE, 8, 4, 128) \
370
+ NVFP4_GET_IF_M234(W_TYPE, 16, 4, 256) \
371
+ NVFP4_GET_IF_M234(W_TYPE, 8, 4, 128)
372
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) \
360
- BIGGROUP_GET_IF_M234(W_TYPE, 16, 4, 256) \
361
- BIGGROUP_GET_IF_M234(W_TYPE, 8, 4, 128)
373
+ #define MXFP4_GET_IF_M1(W_TYPE, N_BLOCKS, K_BLOCKS, NUM_THREADS) \
374
+ _GET_IF(W_TYPE, 1, N_BLOCKS, K_BLOCKS, true, 2, NUM_THREADS, false) \
375
+ _GET_IF(W_TYPE, 1, N_BLOCKS, K_BLOCKS, false, 2, NUM_THREADS, false)
376
+
377
+ #define MXFP4_GET_IF_M234(W_TYPE, N_BLOCKS, K_BLOCKS, NUM_THREADS) \
378
+ _GET_IF(W_TYPE, 2, N_BLOCKS, K_BLOCKS, false, 2, NUM_THREADS, false) \
379
+ _GET_IF(W_TYPE, 3, N_BLOCKS, K_BLOCKS, false, 2, NUM_THREADS, false) \
380
+ _GET_IF(W_TYPE, 4, N_BLOCKS, K_BLOCKS, false, 2, NUM_THREADS, false)
381
+
382
+ #define MXFP4_GET_IF(W_TYPE) \
383
+ MXFP4_GET_IF_M1(W_TYPE, 8, 8, 256) \
384
+ MXFP4_GET_IF_M1(W_TYPE, 8, 4, 128) \
385
+ MXFP4_GET_IF_M234(W_TYPE, 16, 4, 256) \
386
+ MXFP4_GET_IF_M234(W_TYPE, 8, 4, 128)
387
388
// We currently have 4-bit models only with group_blocks == 4
389
#define FZP_GET_IF_M1(W_TYPE, N_BLOCKS, K_BLOCKS, NUM_THREADS) \
@@ -408,12 +433,17 @@ MarlinFuncPtr get_marlin_kernel(const vllm::ScalarType q_type,
433
COMMON_GET_IF(vllm::kU4B8)
434
COMMON_GET_IF(vllm::kU8B128)
435
411
- BIGGROUP_GET_IF(vllm::kFE4M3fn)
436
+ NVFP4_GET_IF(vllm::kFE2M1f)
437
413
- FP4_GET_IF(vllm::kFE2M1f)
438
+ BIGGROUP_GET_IF(vllm::kFE4M3fn)
439
440
ACT_GET_IF(vllm::kU4B8)
441
ACT_GET_IF(vllm::kU8B128)
442
+ if (std::is_same<scalar_t, nv_bfloat16>::value) {
443
+ if (false) {
444
+ }
445
+ MXFP4_GET_IF(vllm::kFE2M1f)
446
+ }
447
448
return kernel;
449
}
@@ -482,16 +512,16 @@ exec_config_t determine_exec_config(const vllm::ScalarType& q_type, int prob_m,
512
}
513
514
template <typename scalar_t>
485
-void marlin_mm(const void* A, const void* B, void* C, void* C_tmp, void* s,
486
- void* s2, void* zp, void* g_idx, void* perm, void* a_tmp,
487
- void* sorted_token_ids, void* expert_ids,
515
+void marlin_mm(const void* A, const void* B, void* C, void* C_tmp, void* b_bias,
516
+ void* s, void* s2, void* zp, void* g_idx, void* perm,
517
+ void* a_tmp, void* sorted_token_ids, void* expert_ids,
518
void* num_tokens_past_padded, void* topk_weights,
519
int moe_block_size, int top_k, bool mul_topk_weights, bool is_ep,
520
int prob_m, int prob_n, int prob_k, void* workspace,
491
- vllm::ScalarType const& q_type, bool has_act_order,
492
- bool is_k_full, bool has_zp, int num_groups, int group_size,
493
- int dev, cudaStream_t stream, int thread_k, int thread_n,
494
- int sms, bool use_atomic_add, bool use_fp32_reduce,
521
+ vllm::ScalarType const& q_type, bool has_bias,
522
+ bool has_act_order, bool is_k_full, bool has_zp, int num_groups,
523
+ int group_size, int dev, cudaStream_t stream, int thread_k,
524
+ int thread_n, int sms, bool use_atomic_add, bool use_fp32_reduce,
525
bool is_zp_float) {
526
int thread_m_blocks = div_ceil(moe_block_size, 16);
527
bool m_block_size_8 = moe_block_size == 8;
@@ -538,6 +568,7 @@ void marlin_mm(const void* A, const void* B, void* C, void* C_tmp, void* s,
568
const int4* B_ptr = (const int4*)B;
569
int4* C_ptr = (int4*)C;
570
int4* C_tmp_ptr = (int4*)C_tmp;
571
+ const int4* bias_ptr = (const int4*)b_bias;
572
const int4* s_ptr = (const int4*)s;
573
const uint16_t* s2_ptr = (const uint16_t*)s2;
574
const int4* zp_ptr = (const int4*)zp;
@@ -648,10 +679,10 @@ void marlin_mm(const void* A, const void* B, void* C, void* C_tmp, void* s,
679
// avoid ">>>" being formatted to "> > >"
680
// clang-format off
681
kernel<<<blocks, num_threads, max_shared_mem, stream>>>(
651
- A_ptr, B_ptr, C_ptr, C_tmp_ptr, s_ptr, s2_ptr, zp_ptr, g_idx_ptr,
682
+ A_ptr, B_ptr, C_ptr, C_tmp_ptr, bias_ptr, s_ptr, s2_ptr, zp_ptr, g_idx_ptr,
683
sorted_token_ids_ptr, expert_ids_ptr, num_tokens_past_padded_ptr,
684
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);
685
+ prob_n, prob_k, locks, has_bias, use_atomic_add, use_fp32_reduce, max_shared_mem);
686
// clang-format on
687
}
688
@@ -659,7 +690,8 @@ void marlin_mm(const void* A, const void* B, void* C, void* C_tmp, void* s,
690
691
torch::Tensor moe_wna16_marlin_gemm(
692
torch::Tensor& a, std::optional<torch::Tensor> const& c_or_none,
662
- torch::Tensor& b_q_weight, torch::Tensor& b_scales,
693
+ torch::Tensor& b_q_weight,
694
+ std::optional<torch::Tensor> const& b_bias_or_none, torch::Tensor& b_scales,
695
std::optional<torch::Tensor> const& global_scale_or_none,
696
std::optional<torch::Tensor> const& b_zeros_or_none,
697
std::optional<torch::Tensor> const& g_idx_or_none,
@@ -766,7 +798,6 @@ torch::Tensor moe_wna16_marlin_gemm(
798
num_groups = b_scales.size(1);
799
800
torch::Tensor g_idx, perm, a_tmp;
769
- ;
801
if (g_idx_or_none.has_value() && perm_or_none.has_value()) {
802
g_idx = g_idx_or_none.value();
803
perm = perm_or_none.value();
@@ -815,12 +846,24 @@ torch::Tensor moe_wna16_marlin_gemm(
846
torch::Tensor global_scale;
847
if (global_scale_or_none.has_value()) {
848
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.");
849
+ TORCH_CHECK(b_q_type == vllm::kFE2M1f && group_size == 16,
850
+ "global_scale can only be used for nvfp4 format.");
851
} else {
852
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.");
853
+ TORCH_CHECK(!(b_q_type == vllm::kFE2M1f && group_size == 16),
854
+ "the global_scale parameter must be passed for nvfp4 format.");
855
+ }
856
+
857
+ bool has_bias = b_bias_or_none.has_value();
858
+ torch::Tensor b_bias;
859
+ if (has_bias) {
860
+ b_bias = b_bias_or_none.value();
861
+ TORCH_CHECK(b_bias.device().is_cuda(), "b_bias is not on GPU");
862
+ TORCH_CHECK(b_bias.is_contiguous(), "b_bias is not contiguous");
863
+ TORCH_CHECK(b_bias.size(1) == size_n, "b_bias.size(0) != size_n");
864
+ TORCH_CHECK(b_bias.stride(1) == 1, "b_bias.stride(1) != 1");
865
+ } else {
866
+ b_bias = torch::empty({0}, options);
867
}
868
869
torch::Tensor b_zeros;
@@ -832,7 +875,6 @@ torch::Tensor moe_wna16_marlin_gemm(
875
b_zeros = torch::empty({0}, options);
876
}
877
bool has_zp = b_zeros.size(-1) > 0;
835
-
878
if (has_zp) {
879
TORCH_CHECK(
880
b_q_type == vllm::kU4 || b_q_type == vllm::kU8,
@@ -890,41 +932,58 @@ torch::Tensor moe_wna16_marlin_gemm(
932
if (a.scalar_type() == at::ScalarType::Half) {
933
void* scales_ptr;
934
if (b_q_type == vllm::kFE2M1f) {
893
- scales_ptr = b_scales.data_ptr<at::Float8_e4m3fn>();
935
+ if (group_size == 16)
936
+ scales_ptr = b_scales.data_ptr<at::Float8_e4m3fn>();
937
+ else if (group_size == 32)
938
+ scales_ptr = b_scales.data_ptr<at::Float8_e8m0fnu>();
939
+ else
940
+ TORCH_CHECK(false,
941
+ "float4_e2m1f only supports group_size == 16 (NVFP4) ",
942
+ "and group_size == 32 (MXFP4)");
943
} else {
944
scales_ptr = b_scales.data_ptr<at::Half>();
945
}
946
947
MARLIN_NAMESPACE_NAME::marlin_mm<half>(
948
a.data_ptr<at::Half>(), b_q_weight.data_ptr(), c.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(),
904
- topk_weights.data_ptr(), moe_block_size, top_k, mul_topk_weights, is_ep,
905
- size_m, size_n, size_k, workspace.data_ptr(), b_q_type, has_act_order,
906
- is_k_full, has_zp, num_groups, group_size, dev,
949
+ c_tmp.data_ptr<float>(), b_bias.data_ptr<at::Half>(), scales_ptr,
950
+ global_scale.data_ptr<at::Half>(), b_zeros.data_ptr(), g_idx.data_ptr(),
951
+ perm.data_ptr(), a_tmp.data_ptr<at::Half>(),
952
+ sorted_token_ids.data_ptr(), expert_ids.data_ptr(),
953
+ num_tokens_past_padded.data_ptr(), topk_weights.data_ptr(),
954
+ moe_block_size, top_k, mul_topk_weights, is_ep, size_m, size_n, size_k,
955
+ workspace.data_ptr(), b_q_type, has_bias, has_act_order, is_k_full,
956
+ has_zp, num_groups, group_size, dev,
957
at::cuda::getCurrentCUDAStream(dev), thread_k, thread_n, sms,
958
use_atomic_add, use_fp32_reduce, is_zp_float);
959
} else if (a.scalar_type() == at::ScalarType::BFloat16) {
960
void* scales_ptr;
961
if (b_q_type == vllm::kFE2M1f) {
912
- scales_ptr = b_scales.data_ptr<at::Float8_e4m3fn>();
962
+ if (group_size == 16)
963
+ scales_ptr = b_scales.data_ptr<at::Float8_e4m3fn>();
964
+ else if (group_size == 32)
965
+ scales_ptr = b_scales.data_ptr<at::Float8_e8m0fnu>();
966
+ else
967
+ TORCH_CHECK(false,
968
+ "float4_e2m1f only supports group_size == 16 (NVFP4) ",
969
+ "and group_size == 32 (MXFP4)");
970
} else {
971
scales_ptr = b_scales.data_ptr<at::BFloat16>();
972
}
973
974
MARLIN_NAMESPACE_NAME::marlin_mm<nv_bfloat16>(
975
a.data_ptr<at::BFloat16>(), b_q_weight.data_ptr(),
919
- c.data_ptr<at::BFloat16>(), c_tmp.data_ptr<float>(), scales_ptr,
976
+ c.data_ptr<at::BFloat16>(), c_tmp.data_ptr<float>(),
977
+ b_bias.data_ptr<at::BFloat16>(), scales_ptr,
978
global_scale.data_ptr<at::BFloat16>(), b_zeros.data_ptr(),
979
g_idx.data_ptr(), perm.data_ptr(), a_tmp.data_ptr<at::BFloat16>(),
980
sorted_token_ids.data_ptr(), expert_ids.data_ptr(),
981
num_tokens_past_padded.data_ptr(), topk_weights.data_ptr(),
982
moe_block_size, top_k, mul_topk_weights, is_ep, size_m, size_n, size_k,
925
- workspace.data_ptr(), b_q_type, has_act_order, is_k_full, has_zp,
926
- num_groups, group_size, dev, at::cuda::getCurrentCUDAStream(dev),
927
- thread_k, thread_n, sms, use_atomic_add, use_fp32_reduce, is_zp_float);
983
+ workspace.data_ptr(), b_q_type, has_bias, has_act_order, is_k_full,
984
+ has_zp, num_groups, group_size, dev,
985
+ at::cuda::getCurrentCUDAStream(dev), thread_k, thread_n, sms,
986
+ use_atomic_add, use_fp32_reduce, is_zp_float);
987
} else {
988
TORCH_CHECK(false,
989
"moe_wna16_marlin_gemm only supports bfloat16 and float16");
csrc/moe/torch_bindings.cpp
+2
-1
@@ -35,7 +35,8 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, m) {
35
36
m.def(
37
"moe_wna16_marlin_gemm(Tensor! a, Tensor? c_or_none,"
38
- "Tensor! b_q_weight, Tensor! b_scales, Tensor? global_scale, Tensor? "
38
+ "Tensor! b_q_weight, Tensor? b_bias_or_none,"
39
+ "Tensor! b_scales, Tensor? global_scale, Tensor? "
40
"b_zeros_or_none,"
41
"Tensor? g_idx_or_none, Tensor? perm_or_none, Tensor! workspace,"
42
"Tensor sorted_token_ids,"
csrc/quantization/gptq_marlin/dequant.h
+19
-4
@@ -470,11 +470,12 @@ __device__ inline void dequant<nv_bfloat162, vllm::kFE2M1f.id(), false>(
470
frag_b[0] = __hmul2(frag_b[0], bias_reg);
471
}
472
473
-template <typename scalar_t2>
473
+template <typename scalar_t2, vllm::ScalarTypeId s_type_id>
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) {
477
+__device__ inline void dequant_fp8_scales<half2, vllm::kFE4M3fn.id()>(
478
+ int q, half2* frag_b) {
479
int Out1 = (q & 0xFF00FF00) >> 1;
480
;
481
q <<= 8;
@@ -486,8 +487,8 @@ __device__ inline void dequant_fp8_scales<half2>(int q, half2* frag_b) {
487
};
488
489
template <>
489
-__device__ inline void dequant_fp8_scales<nv_bfloat162>(int q,
490
- nv_bfloat162* frag_b) {
490
+__device__ inline void dequant_fp8_scales<nv_bfloat162, vllm::kFE4M3fn.id()>(
491
+ int q, nv_bfloat162* frag_b) {
492
constexpr int FP8_EXPONENT = 4, BF16_EXPONENT = 8;
493
constexpr int RIGHT_SHIFT = BF16_EXPONENT - FP8_EXPONENT;
494
constexpr int MASK = 0x7F007F00;
@@ -502,6 +503,20 @@ __device__ inline void dequant_fp8_scales<nv_bfloat162>(int q,
503
frag_b[0] = *reinterpret_cast<const nv_bfloat162*>(&Out2);
504
}
505
506
+template <>
507
+__device__ inline void dequant_fp8_scales<nv_bfloat162, vllm::kFE8M0fnu.id()>(
508
+ int q, nv_bfloat162* frag_b) {
509
+ // In this conversion, 2 ** -127 in FP8E8M0 would become 0 in BF16,
510
+ // but we assume that such a extreme value would not occur in real models.
511
+ int Out1 = (q & 0xFF00FF00) >> 1;
512
+ q <<= 7;
513
+ int Out2 = q & 0x7F807F80;
514
+
515
+ // Note: reverse indexing is intentional because weights are permuted
516
+ frag_b[1] = *reinterpret_cast<const nv_bfloat162*>(&Out1);
517
+ frag_b[0] = *reinterpret_cast<const nv_bfloat162*>(&Out2);
518
+}
519
+
520
#endif
521
522
} // namespace MARLIN_NAMESPACE_NAME
csrc/quantization/gptq_marlin/generate_kernels.py
+16
-1
@@ -20,6 +20,7 @@ namespace MARLIN_NAMESPACE_NAME {
20
TEMPLATE = ("template __global__ void Marlin<"
21
"{{scalar_t}}, "
22
"{{w_type_id}}, "
23
+ "{{s_type_id}}, "
24
"{{threads}}, "
25
"{{thread_m_blocks}}, "
26
"{{thread_n_blocks}}, "
@@ -78,7 +79,8 @@ def generate_new_kernels():
79
if scalar_type == "vllm::kFE4M3fn" and group_blocks not in [-1, 8]:
80
continue
81
# nvfp4 only supports group_size == 16
81
- if scalar_type == "vllm::kFE2M1f" and group_blocks != 1:
82
+ # mxfp4 only supports group_size == 32
83
+ if scalar_type == "vllm::kFE2M1f" and group_blocks not in [1, 2]:
84
continue
85
# other quantization methods don't support group_size = 16
86
if scalar_type != "vllm::kFE2M1f" and group_blocks == 1:
@@ -97,10 +99,23 @@ def generate_new_kernels():
99
# 4bit quantization and fp16
100
is_zp_float_list.append(True)
101
102
+ if scalar_type == "vllm::kFE2M1f" and group_blocks == 1:
103
+ s_type = "vllm::kFE4M3fn"
104
+ elif scalar_type == "vllm::kFE2M1f" and group_blocks == 2:
105
+ s_type = "vllm::kFE8M0fnu"
106
+ if dtype == "fp16":
107
+ # we cannot safely dequantize e8m0 to fp16, so skip this
108
+ continue
109
+ elif dtype == "fp16":
110
+ s_type = "vllm::kFloat16"
111
+ elif dtype == "bf16":
112
+ s_type = "vllm::kBFloat16"
113
+
114
for is_zp_float in is_zp_float_list:
115
template_str = jinja2.Template(TEMPLATE).render(
116
scalar_t=c_dtype,
117
w_type_id=scalar_type + ".id()",
118
+ s_type_id=s_type + ".id()",
119
threads=threads,
120
thread_m_blocks=max(m_blocks, 1),
121
thread_n_blocks=n_blocks,
csrc/quantization/gptq_marlin/gptq_marlin.cu
+113
-49
@@ -48,7 +48,8 @@ __global__ void permute_cols_kernel(int4 const* __restrict__ a_int4_ptr,
48
49
torch::Tensor gptq_marlin_gemm(
50
torch::Tensor& a, std::optional<torch::Tensor> c_or_none,
51
- torch::Tensor& b_q_weight, torch::Tensor& b_scales,
51
+ torch::Tensor& b_q_weight,
52
+ std::optional<torch::Tensor> const& b_bias_or_none, torch::Tensor& b_scales,
53
std::optional<torch::Tensor> const& b_zeros_or_none,
54
std::optional<torch::Tensor> const& g_idx_or_none,
55
std::optional<torch::Tensor> const& perm_or_none, torch::Tensor& workspace,
@@ -187,7 +188,12 @@ int get_kernel_cache_size(thread_config_t const& th_config, int thread_m_blocks,
188
int tb_m = thread_m_blocks * 16;
189
int sh_a_size = pipe_stages * (tb_m * tb_k) * 2;
190
int sh_b_size = pipe_stages * (tb_k * tb_n / pack_factor) * 4;
190
- int sh_red_size = tb_m * (tb_n + 8);
191
+ int sh_red_size = tb_m * (tb_n + 8) * 2;
192
+ int sh_bias_size = tb_n * 2;
193
+ int tmp_size =
194
+ (sh_b_size > sh_red_size ? sh_red_size : sh_b_size) + sh_bias_size;
195
+ tmp_size = max(max(sh_b_size, sh_red_size), tmp_size);
196
+
197
int sh_s_size =
198
get_scales_cache_size(th_config, prob_m, prob_n, prob_k, num_bits,
199
group_size, has_act_order, is_k_full);
@@ -202,8 +208,8 @@ int get_kernel_cache_size(thread_config_t const& th_config, int thread_m_blocks,
208
sh_zp_size = sh_s_size / 2;
209
}
210
205
- int total_size = max(sh_b_size, sh_red_size) + sh_a_size + sh_s_size +
206
- sh_zp_size + sh_g_idx_size;
211
+ int total_size =
212
+ tmp_size + sh_a_size + sh_s_size + sh_zp_size + sh_g_idx_size;
213
214
return total_size;
215
}
@@ -237,20 +243,25 @@ bool is_valid_config(thread_config_t const& th_config, int thread_m_blocks,
243
int cache_size = get_kernel_cache_size(
244
th_config, thread_m_blocks, prob_m, prob_n, prob_k, num_bits, group_size,
245
has_act_order, is_k_full, has_zp, is_zp_float);
240
- return cache_size <= max_shared_mem;
246
+ return cache_size + 512 <= max_shared_mem;
247
}
248
243
- #define _GET_IF(W_TYPE, THREAD_M_BLOCKS, THREAD_N_BLOCKS, THREAD_K_BLOCKS, \
244
- M_BLOCK_SIZE_8, GROUP_BLOCKS, NUM_THREADS, IS_ZP_FLOAT) \
245
- else if (q_type == W_TYPE && thread_m_blocks == THREAD_M_BLOCKS && \
246
- thread_n_blocks == THREAD_N_BLOCKS && \
247
- thread_k_blocks == THREAD_K_BLOCKS && \
248
- m_block_size_8 == M_BLOCK_SIZE_8 && \
249
- group_blocks == GROUP_BLOCKS && num_threads == NUM_THREADS && \
250
- is_zp_float == IS_ZP_FLOAT) { \
251
- kernel = Marlin<scalar_t, W_TYPE.id(), NUM_THREADS, THREAD_M_BLOCKS, \
252
- THREAD_N_BLOCKS, THREAD_K_BLOCKS, M_BLOCK_SIZE_8, \
253
- pipe_stages, GROUP_BLOCKS, IS_ZP_FLOAT>; \
249
+ #define _GET_IF(W_TYPE, THREAD_M_BLOCKS, THREAD_N_BLOCKS, THREAD_K_BLOCKS, \
250
+ M_BLOCK_SIZE_8, GROUP_BLOCKS, NUM_THREADS, IS_ZP_FLOAT) \
251
+ else if (q_type == W_TYPE && thread_m_blocks == THREAD_M_BLOCKS && \
252
+ thread_n_blocks == THREAD_N_BLOCKS && \
253
+ thread_k_blocks == THREAD_K_BLOCKS && \
254
+ m_block_size_8 == M_BLOCK_SIZE_8 && \
255
+ group_blocks == GROUP_BLOCKS && num_threads == NUM_THREADS && \
256
+ is_zp_float == IS_ZP_FLOAT) { \
257
+ constexpr auto S_TYPE = \
258
+ W_TYPE == vllm::kFE2M1f \
259
+ ? (GROUP_BLOCKS == 1 ? vllm::kFE4M3fn : vllm::kFE8M0fnu) \
260
+ : (std::is_same<scalar_t, half>::value ? vllm::kFloat16 \
261
+ : vllm::kBFloat16); \
262
+ kernel = Marlin<scalar_t, W_TYPE.id(), S_TYPE.id(), NUM_THREADS, \
263
+ THREAD_M_BLOCKS, THREAD_N_BLOCKS, THREAD_K_BLOCKS, \
264
+ M_BLOCK_SIZE_8, pipe_stages, GROUP_BLOCKS, IS_ZP_FLOAT>; \
265
}
266
267
// COMMON: cases for (group_blocks in [-1, 2, 4, 8] and is_zp_float == false)
@@ -315,22 +326,39 @@ bool is_valid_config(thread_config_t const& th_config, int thread_m_blocks,
326
BIGGROUP_GET_IF_M234(W_TYPE, 8, 4, 128) \
327
BIGGROUP_GET_IF_M234(W_TYPE, 4, 8, 128)
328
318
- #define FP4_GET_IF_M1(W_TYPE, N_BLOCKS, K_BLOCKS, NUM_THREADS) \
329
+ #define NVFP4_GET_IF_M1(W_TYPE, N_BLOCKS, K_BLOCKS, NUM_THREADS) \
330
_GET_IF(W_TYPE, 1, N_BLOCKS, K_BLOCKS, true, 1, NUM_THREADS, false) \
331
_GET_IF(W_TYPE, 1, N_BLOCKS, K_BLOCKS, false, 1, NUM_THREADS, false)
332
322
- #define FP4_GET_IF_M234(W_TYPE, N_BLOCKS, K_BLOCKS, NUM_THREADS) \
333
+ #define NVFP4_GET_IF_M234(W_TYPE, N_BLOCKS, K_BLOCKS, NUM_THREADS) \
334
_GET_IF(W_TYPE, 2, N_BLOCKS, K_BLOCKS, false, 1, NUM_THREADS, false) \
335
_GET_IF(W_TYPE, 3, N_BLOCKS, K_BLOCKS, false, 1, NUM_THREADS, false) \
336
_GET_IF(W_TYPE, 4, N_BLOCKS, K_BLOCKS, false, 1, NUM_THREADS, false)
337
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)
338
+ #define NVFP4_GET_IF(W_TYPE) \
339
+ NVFP4_GET_IF_M1(W_TYPE, 8, 8, 256) \
340
+ NVFP4_GET_IF_M1(W_TYPE, 8, 4, 128) \
341
+ NVFP4_GET_IF_M1(W_TYPE, 4, 8, 128) \
342
+ NVFP4_GET_IF_M234(W_TYPE, 16, 4, 256) \
343
+ NVFP4_GET_IF_M234(W_TYPE, 8, 4, 128) \
344
+ NVFP4_GET_IF_M234(W_TYPE, 4, 8, 128)
345
+
346
+ #define MXFP4_GET_IF_M1(W_TYPE, N_BLOCKS, K_BLOCKS, NUM_THREADS) \
347
+ _GET_IF(W_TYPE, 1, N_BLOCKS, K_BLOCKS, true, 2, NUM_THREADS, false) \
348
+ _GET_IF(W_TYPE, 1, N_BLOCKS, K_BLOCKS, false, 2, NUM_THREADS, false)
349
+
350
+ #define MXFP4_GET_IF_M234(W_TYPE, N_BLOCKS, K_BLOCKS, NUM_THREADS) \
351
+ _GET_IF(W_TYPE, 2, N_BLOCKS, K_BLOCKS, false, 2, NUM_THREADS, false) \
352
+ _GET_IF(W_TYPE, 3, N_BLOCKS, K_BLOCKS, false, 2, NUM_THREADS, false) \
353
+ _GET_IF(W_TYPE, 4, N_BLOCKS, K_BLOCKS, false, 2, NUM_THREADS, false)
354
+
355
+ #define MXFP4_GET_IF(W_TYPE) \
356
+ MXFP4_GET_IF_M1(W_TYPE, 8, 8, 256) \
357
+ MXFP4_GET_IF_M1(W_TYPE, 8, 4, 128) \
358
+ MXFP4_GET_IF_M1(W_TYPE, 4, 8, 128) \
359
+ MXFP4_GET_IF_M234(W_TYPE, 16, 4, 256) \
360
+ MXFP4_GET_IF_M234(W_TYPE, 8, 4, 128) \
361
+ MXFP4_GET_IF_M234(W_TYPE, 4, 8, 128)
362
363
// We currently have 4-bit models only with group_blocks == 4
364
#define FZP_GET_IF_M1(W_TYPE, N_BLOCKS, K_BLOCKS, NUM_THREADS) \
@@ -384,7 +412,7 @@ MarlinFuncPtr get_marlin_kernel(const vllm::ScalarType q_type,
412
COMMON_GET_IF(vllm::kU4B8)
413
COMMON_GET_IF(vllm::kU8B128)
414
387
- FP4_GET_IF(vllm::kFE2M1f)
415
+ NVFP4_GET_IF(vllm::kFE2M1f)
416
417
BIGGROUP_GET_IF(vllm::kFE4M3fn)
418
@@ -396,6 +424,11 @@ MarlinFuncPtr get_marlin_kernel(const vllm::ScalarType q_type,
424
}
425
FZP_GET_IF(vllm::kU4)
426
}
427
+ if (std::is_same<scalar_t, nv_bfloat16>::value) {
428
+ if (false) {
429
+ }
430
+ MXFP4_GET_IF(vllm::kFE2M1f)
431
+ }
432
433
return kernel;
434
}
@@ -453,12 +486,12 @@ exec_config_t determine_exec_config(const vllm::ScalarType& q_type, int prob_m,
486
}
487
488
template <typename scalar_t>
456
-void marlin_mm(const void* A, const void* B, void* C, void* C_tmp, void* s,
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,
489
+void marlin_mm(const void* A, const void* B, void* C, void* C_tmp, void* b_bias,
490
+ void* s, void* s2, void* zp, void* g_idx, void* perm,
491
+ void* a_tmp, int prob_m, int prob_n, int prob_k, int lda,
492
+ void* workspace, vllm::ScalarType const& q_type, bool has_bias,
493
+ bool has_act_order, bool is_k_full, bool has_zp, int num_groups,
494
+ int group_size, int dev, cudaStream_t stream, int thread_k_init,
495
int thread_n_init, int sms, bool use_atomic_add,
496
bool use_fp32_reduce, bool is_zp_float) {
497
if (has_zp) {
@@ -503,6 +536,7 @@ void marlin_mm(const void* A, const void* B, void* C, void* C_tmp, void* s,
536
const int4* B_ptr = (const int4*)B;
537
int4* C_ptr = (int4*)C;
538
int4* C_tmp_ptr = (int4*)C_tmp;
539
+ const int4* bias_ptr = (const int4*)b_bias;
540
const int4* s_ptr = (const int4*)s;
541
const uint16_t* s2_ptr = (const uint16_t*)s2;
542
const int4* zp_ptr = (const int4*)zp;
@@ -623,8 +657,9 @@ void marlin_mm(const void* A, const void* B, void* C, void* C_tmp, void* s,
657
// avoid ">>>" being formatted to "> > >"
658
// clang-format off
659
kernel<<<blocks, num_threads, max_shared_mem_new, stream>>>(
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,
660
+ A_ptr, B_ptr, C_ptr, C_tmp_ptr, bias_ptr, s_ptr, s2_ptr, zp_ptr,
661
+ g_idx_ptr, num_groups,
662
+ prob_m_split, prob_n, prob_k, lda, locks, has_bias, part_use_atomic_add,
663
use_fp32_reduce, max_shared_mem_new);
664
// clang-format on
665
@@ -638,7 +673,8 @@ void marlin_mm(const void* A, const void* B, void* C, void* C_tmp, void* s,
673
674
torch::Tensor gptq_marlin_gemm(
675
torch::Tensor& a, std::optional<torch::Tensor> c_or_none,
641
- torch::Tensor& b_q_weight, torch::Tensor& b_scales,
676
+ torch::Tensor& b_q_weight,
677
+ std::optional<torch::Tensor> const& b_bias_or_none, torch::Tensor& b_scales,
678
std::optional<torch::Tensor> const& global_scale_or_none,
679
std::optional<torch::Tensor> const& b_zeros_or_none,
680
std::optional<torch::Tensor> const& g_idx_or_none,
@@ -785,12 +821,24 @@ torch::Tensor gptq_marlin_gemm(
821
torch::Tensor global_scale;
822
if (global_scale_or_none.has_value()) {
823
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.");
824
+ TORCH_CHECK(b_q_type == vllm::kFE2M1f && group_size == 16,
825
+ "global_scale can only be used for nvfp4 format.");
826
} else {
827
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.");
828
+ TORCH_CHECK(!(b_q_type == vllm::kFE2M1f && group_size == 16),
829
+ "the global_scale parameter must be passed for nvfp4 format.");
830
+ }
831
+
832
+ bool has_bias = b_bias_or_none.has_value();
833
+ torch::Tensor b_bias;
834
+ if (has_bias) {
835
+ b_bias = b_bias_or_none.value();
836
+ TORCH_CHECK(b_bias.device().is_cuda(), "b_bias is not on GPU");
837
+ TORCH_CHECK(b_bias.is_contiguous(), "b_bias is not contiguous");
838
+ TORCH_CHECK(b_bias.size(0) == size_n, "b_bias.size(0) != size_n");
839
+ TORCH_CHECK(b_bias.stride(0) == 1, "b_bias.stride(0) != 1");
840
+ } else {
841
+ b_bias = torch::empty({0}, options);
842
}
843
844
torch::Tensor b_zeros;
@@ -857,34 +905,50 @@ torch::Tensor gptq_marlin_gemm(
905
if (a.scalar_type() == at::ScalarType::Half) {
906
void* scales_ptr;
907
if (b_q_type == vllm::kFE2M1f) {
860
- scales_ptr = b_scales.data_ptr<at::Float8_e4m3fn>();
908
+ if (group_size == 16)
909
+ scales_ptr = b_scales.data_ptr<at::Float8_e4m3fn>();
910
+ else if (group_size == 32)
911
+ scales_ptr = b_scales.data_ptr<at::Float8_e8m0fnu>();
912
+ else
913
+ TORCH_CHECK(false,
914
+ "float4_e2m1f only supports group_size == 16 (NVFP4) ",
915
+ "and group_size == 32 (MXFP4)");
916
} else {
917
scales_ptr = b_scales.data_ptr<at::Half>();
918
}
919
920
marlin::marlin_mm<half>(
921
a.data_ptr<at::Half>(), b_q_weight.data_ptr(), c.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);
922
+ c_tmp.data_ptr<float>(), b_bias.data_ptr<at::Half>(), scales_ptr,
923
+ global_scale.data_ptr<at::Half>(), b_zeros.data_ptr(), g_idx.data_ptr(),
924
+ perm.data_ptr(), a_tmp.data_ptr<at::Half>(), size_m, size_n, size_k,
925
+ a.stride(0), workspace.data_ptr(), b_q_type, has_bias, has_act_order,
926
+ is_k_full, has_zp, num_groups, group_size, dev,
927
+ at::cuda::getCurrentCUDAStream(dev), thread_k, thread_n, sms,
928
+ use_atomic_add, use_fp32_reduce, is_zp_float);
929
} else if (a.scalar_type() == at::ScalarType::BFloat16) {
930
void* scales_ptr;
931
if (b_q_type == vllm::kFE2M1f) {
876
- scales_ptr = b_scales.data_ptr<at::Float8_e4m3fn>();
932
+ if (group_size == 16)
933
+ scales_ptr = b_scales.data_ptr<at::Float8_e4m3fn>();
934
+ else if (group_size == 32)
935
+ scales_ptr = b_scales.data_ptr<at::Float8_e8m0fnu>();
936
+ else
937
+ TORCH_CHECK(false,
938
+ "float4_e2m1f only supports group_size == 16 (NVFP4) ",
939
+ "and group_size == 32 (MXFP4)");
940
} else {
941
scales_ptr = b_scales.data_ptr<at::BFloat16>();
942
}
943
944
marlin::marlin_mm<nv_bfloat16>(
945
a.data_ptr<at::BFloat16>(), b_q_weight.data_ptr(),
883
- c.data_ptr<at::BFloat16>(), c_tmp.data_ptr<float>(), scales_ptr,
946
+ c.data_ptr<at::BFloat16>(), c_tmp.data_ptr<float>(),
947
+ b_bias.data_ptr<at::BFloat16>(), scales_ptr,
948
global_scale.data_ptr<at::BFloat16>(), b_zeros.data_ptr(),
949
g_idx.data_ptr(), perm.data_ptr(), a_tmp.data_ptr<at::BFloat16>(),
950
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,
951
+ has_bias, has_act_order, is_k_full, has_zp, num_groups, group_size, dev,
952
at::cuda::getCurrentCUDAStream(dev), thread_k, thread_n, sms,
953
use_atomic_add, use_fp32_reduce, is_zp_float);
954
} else {
csrc/quantization/gptq_marlin/kernel.h
+4
-1
@@ -10,15 +10,18 @@
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__ b_bias_ptr, \
14
const int4 *__restrict__ scales_ptr, \
15
const uint16_t *__restrict__ scale2_ptr, \
16
const int4 *__restrict__ zp_ptr, const int *__restrict__ g_idx, \
17
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
+ bool has_bias, bool use_atomic_add, bool use_fp32_reduce, \
19
+ int max_shared_mem
20
21
namespace MARLIN_NAMESPACE_NAME {
22
template <typename scalar_t, // compute dtype, half or nv_float16
23
const vllm::ScalarTypeId w_type_id, // weight ScalarType id
24
+ const vllm::ScalarTypeId s_type_id, // weight ScalarType id
25
const int threads, // number of threads in a threadblock
26
const int thread_m_blocks, // number of 16x16 blocks in the m
27
// dimension (batchsize) of the
csrc/quantization/gptq_marlin/marlin_template.h
+109
-30
@@ -39,6 +39,7 @@ namespace MARLIN_NAMESPACE_NAME {
39
40
template <typename scalar_t, // compute dtype, half or nv_float16
41
const vllm::ScalarTypeId w_type_id, // weight ScalarType id
42
+ const vllm::ScalarTypeId s_type_id, // weight scale ScalarType id
43
const int threads, // number of threads in a threadblock
44
const int thread_m_blocks, // number of 16x16 blocks in the m
45
// dimension (batchsize) of the
@@ -271,6 +272,7 @@ __device__ inline void wait_negative_and_add(int* lock) {
272
273
template <typename scalar_t, // compute dtype, half or nv_float16
274
const vllm::ScalarTypeId w_type_id, // weight ScalarType id
275
+ const vllm::ScalarTypeId s_type_id, // weight scale ScalarType id
276
const int threads, // number of threads in a threadblock
277
const int thread_m_blocks, // number of 16x16 blocks in the m
278
// dimension (batchsize) of the
@@ -290,6 +292,7 @@ __global__ void Marlin(
292
const int4* __restrict__ B, // 4bit quantized weight matrix of shape kxn
293
int4* __restrict__ C, // fp16 output buffer of shape mxn
294
int4* __restrict__ C_tmp, // fp32 tmp output buffer (for reduce)
295
+ const int4* __restrict__ b_bias_ptr,
296
const int4* __restrict__ scales_ptr, // fp16 quantization scales of shape
297
// (k/groupsize)xn
298
const uint16_t* __restrict__ scale2_ptr, // fp16 global scale (for nvfp4
@@ -297,12 +300,13 @@ __global__ void Marlin(
300
const int4* __restrict__ zp_ptr, // 4bit packed zero-points of shape
301
// (k/groupsize)x(n/pack_factor)
302
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
303
- int prob_k, // reduction dimension k
304
- int lda, // A.stride(0), equal to prob_k is A is contiguous
305
- int* locks, // extra global storage for barrier synchronization
303
+ int num_groups, // number of scale groups per output channel
304
+ int prob_m, // batch dimension m
305
+ int prob_n, // output dimension n
306
+ int prob_k, // reduction dimension k
307
+ int lda, // A.stride(0), equal to prob_k is A is contiguous
308
+ int* locks, // extra global storage for barrier synchronization
309
+ bool has_bias,
310
bool use_atomic_add, // whether to use atomic add to reduce
311
bool use_fp32_reduce, // whether to use fp32 global reduce
312
int max_shared_mem) {
@@ -326,18 +330,29 @@ __global__ void Marlin(
330
using FragZP = typename ScalarType<scalar_t>::FragZP;
331
332
static constexpr auto w_type = vllm::ScalarType::from_id(w_type_id);
333
+ static constexpr auto s_type = vllm::ScalarType::from_id(s_type_id);
334
+ if constexpr (w_type == vllm::kFE2M1f) {
335
+ static_assert(s_type == vllm::kFE4M3fn && group_blocks == 1 ||
336
+ s_type == vllm::kFE8M0fnu && group_blocks == 2);
337
+ } else if constexpr (std::is_same<scalar_t, nv_bfloat16>::value) {
338
+ static_assert(s_type == vllm::kBFloat16);
339
+ } else if constexpr (std::is_same<scalar_t, half>::value) {
340
+ static_assert(s_type == vllm::kFloat16);
341
+ }
342
+
343
constexpr bool has_zp = w_type == vllm::kU4 || w_type == vllm::kU8;
344
constexpr bool is_int_type = w_type == vllm::kU4 || w_type == vllm::kU8 ||
345
w_type == vllm::kU4B8 || w_type == vllm::kU8B128;
346
// see comments of dequant.h for more details
347
constexpr bool dequant_skip_flop =
334
- !is_int_type ||
348
+ w_type == vllm::kFE4M3fn ||
349
+ w_type == vllm::kFE2M1f && s_type == vllm::kFE4M3fn ||
350
has_zp && !is_zp_float && !std::is_same<scalar_t, nv_bfloat16>::value ||
351
has_zp && !is_zp_float && !(w_type == vllm::kU8);
352
353
scalar_t2 global_scale;
339
-
340
- if constexpr (w_type == vllm::kFE2M1f) {
354
+ if constexpr (w_type == vllm::kFE2M1f && s_type == vllm::kFE4M3fn) {
355
+ // NVFP4 format requires global scale
356
uint16_t val = scale2_ptr[0];
357
global_scale = Dtype::num2num2(*reinterpret_cast<scalar_t*>(&val));
358
}
@@ -589,7 +604,7 @@ __global__ void Marlin(
604
605
s_sh_rd = 8 * ((threadIdx.x / 32) % (thread_n_blocks / 4)) +
606
(threadIdx.x % 32) / 4;
592
- s_sh_rd = s_sh_rd * 2 + warp_row % 2;
607
+ s_sh_rd = s_sh_rd * 2 + (warp_row / group_blocks) % 2;
608
609
} else if constexpr (group_blocks != -1)
610
s_sh_rd = 8 * ((threadIdx.x / 32) % (thread_n_blocks / 4)) +
@@ -602,6 +617,18 @@ __global__ void Marlin(
617
s_sh_rd = 8 * ((threadIdx.x / 32) % (thread_n_blocks / 4)) +
618
(threadIdx.x % 32) % 4;
619
620
+ int bias_sh_rd;
621
+ if constexpr (m_block_size_8) {
622
+ bias_sh_rd = 8 * ((threadIdx.x / 32) % (thread_n_blocks / 4)) +
623
+ (threadIdx.x % 32) / 8;
624
+ } else {
625
+ bias_sh_rd = 8 * ((threadIdx.x / 32) % (thread_n_blocks / 4)) +
626
+ (threadIdx.x % 32) % 4;
627
+ }
628
+
629
+ int bias_sh_wr = threadIdx.x;
630
+ int bias_gl_rd = (thread_n_blocks * 16 / 8) * slice_col + threadIdx.x;
631
+
632
// Zero-points have the same read layout as the scales
633
// (without column-wise case)
634
constexpr int num_col_threads = 8;
@@ -670,7 +697,19 @@ __global__ void Marlin(
697
constexpr int sh_b_size = stages * b_sh_stage;
698
int4* sh_b = sh;
699
int4* sh_red = sh;
673
- int4* sh_g_idx = sh_b + (sh_red_size > sh_b_size ? sh_red_size : sh_b_size);
700
+
701
+ constexpr int sh_size_b_red_min =
702
+ (sh_red_size < sh_b_size ? sh_red_size : sh_b_size);
703
+ constexpr int sh_size_b_red_max =
704
+ (sh_red_size > sh_b_size ? sh_red_size : sh_b_size);
705
+ constexpr int sh_bias_size = (thread_n_blocks * 16 / 8);
706
+ constexpr int sh_b_red_bias_size =
707
+ sh_size_b_red_max > (sh_size_b_red_min + sh_bias_size)
708
+ ? sh_size_b_red_max
709
+ : (sh_size_b_red_min + sh_bias_size);
710
+
711
+ int4* sh_bias = sh + sh_size_b_red_min;
712
+ int4* sh_g_idx = sh + sh_b_red_bias_size;
713
int4* sh_zp = sh_g_idx + (stages * g_idx_stage);
714
constexpr int sh_s_size = has_act_order ? (act_s_max_num_groups * s_sh_stride)
715
: (stages * s_sh_stage);
@@ -680,15 +719,13 @@ __global__ void Marlin(
719
static_assert(thread_m_blocks * 16 * thread_n_blocks * 16 / 8 <=
720
stages * b_sh_stage);
721
int4* sh_a = sh_s + sh_s_size;
683
- // constexpr int shm_size_used =
684
- // stages * (g_idx_stage + zp_sh_stage) + sh_s_size +
685
- // (sh_red_size > sh_b_size ? sh_red_size : sh_b_size);
722
723
// Register storage for double buffer of shared memory reads.
724
FragA frag_a[2][thread_m_blocks];
725
I4 frag_b_quant[2][b_thread_vecs];
726
FragC frag_c[thread_m_blocks][4][2];
691
- FragS frag_s[2][4]; // No act-order
727
+ FragS frag_s[2][4]; // No act-order
728
+ FragS frag_bias[2][4];
729
FragS act_frag_s[2][4][4]; // For act-order
730
int frag_qzp[2][num_ints_per_thread]; // Zero-points
731
FragZP frag_zp; // Zero-points in fp16
@@ -923,10 +960,15 @@ __global__ void Marlin(
960
if constexpr (w_type_id != vllm::kFE2M1f.id()) {
961
reinterpret_cast<int4*>(&frag_s[k % 2])[0] =
962
sh_s_stage[s_sh_rd + cur_group_id * s_sh_stride];
926
- } else {
963
+ } else if constexpr (group_blocks == 1 || thread_k_blocks > 4) {
964
reinterpret_cast<int2*>(&frag_s[k % 2])[0] =
965
reinterpret_cast<int2*>(
966
sh_s_stage)[s_sh_rd + cur_group_id * (2 * s_sh_stride)];
967
+ } else {
968
+ reinterpret_cast<int2*>(&frag_s[k % 2])[0] =
969
+ reinterpret_cast<int2*>(
970
+ sh_s_stage)[s_sh_rd + cur_group_id * (2 * s_sh_stride) +
971
+ k % 2];
972
}
973
}
974
}
@@ -1139,9 +1181,9 @@ __global__ void Marlin(
1181
int s_quant_0 = reinterpret_cast<int*>(frag_s[k2])[0];
1182
int s_quant_1 = reinterpret_cast<int*>(frag_s[k2])[1];
1183
1142
- dequant_fp8_scales<scalar_t2>(s_quant_0,
1143
- reinterpret_cast<scalar_t2*>(&frag_s[k2]));
1144
- dequant_fp8_scales<scalar_t2>(
1184
+ dequant_fp8_scales<scalar_t2, s_type_id>(
1185
+ s_quant_0, reinterpret_cast<scalar_t2*>(&frag_s[k2]));
1186
+ dequant_fp8_scales<scalar_t2, s_type_id>(
1187
s_quant_1, reinterpret_cast<scalar_t2*>(&frag_s[k2]) + 2);
1188
}
1189
@@ -1411,7 +1453,7 @@ __global__ void Marlin(
1453
// Write out the reduce final result in the correct layout. We only actually
1454
// reshuffle matrix fragments in this step, the reduction above is performed
1455
// in fragment layout.
1414
- auto write_result = [&]() {
1456
+ auto write_result = [&](bool last) {
1457
int c_gl_stride = prob_n / 8;
1458
constexpr int c_sh_stride = 2 * thread_n_blocks + 1;
1459
int c_gl_wr_delta = c_gl_stride * (threads / (2 * thread_n_blocks));
@@ -1438,7 +1480,7 @@ __global__ void Marlin(
1480
int c_gl_wr_end = c_gl_stride * prob_m;
1481
// We first reorder in shared memory to guarantee the most efficient final
1482
// global write patterns
1441
- auto write = [&](int idx, float c0, float c1, FragS& s) {
1483
+ auto write = [&](int idx, float c0, float c1, FragS& s, FragS& b_bias) {
1484
scalar_t2 res =
1485
Dtype::nums2num2(Dtype::float2num(c0), Dtype::float2num(c1));
1486
@@ -1447,12 +1489,25 @@ __global__ void Marlin(
1489
if constexpr (!has_act_order && group_blocks == -1 &&
1490
w_type.size_bits() == 4 &&
1491
(has_zp && dequant_skip_flop || !has_zp)) {
1450
- res = __hmul2(res, s[0]);
1492
+ scalar_t2 tmp_scale = s[0];
1493
+ if constexpr (m_block_size_8) {
1494
+ tmp_scale = Dtype::num2num2(
1495
+ reinterpret_cast<scalar_t*>(&s[0])[(threadIdx.x % 8) / 4]);
1496
+ }
1497
+ res = __hmul2(res, tmp_scale);
1498
}
1499
1453
- if constexpr (w_type == vllm::kFE2M1f) {
1500
+ if constexpr (w_type == vllm::kFE2M1f && s_type == vllm::kFE4M3fn) {
1501
res = __hmul2(res, global_scale);
1502
}
1503
+ if (has_bias && last) {
1504
+ scalar_t2 tmp_bias = b_bias[0];
1505
+ if constexpr (m_block_size_8) {
1506
+ tmp_bias = Dtype::num2num2(
1507
+ reinterpret_cast<scalar_t*>(&b_bias[0])[(threadIdx.x % 8) / 4]);
1508
+ }
1509
+ res = __hadd2(res, tmp_bias);
1510
+ }
1511
1512
if constexpr (m_block_size_8) {
1513
((scalar_t*)sh_red)[idx] = res.x;
@@ -1470,19 +1525,25 @@ __global__ void Marlin(
1525
if constexpr (m_block_size_8) {
1526
int wr = c_sh_wr + 16 * j;
1527
write(wr, frag_c[i][j][0][0], frag_c[i][j][0][1],
1473
- frag_s[j / 2][2 * (j % 2) + 0]);
1528
+ frag_s[j / 2][2 * (j % 2) + 0],
1529
+ frag_bias[j / 2][2 * (j % 2) + 0]);
1530
write(wr + 8, frag_c[i][j][0][2], frag_c[i][j][0][3],
1475
- frag_s[j / 2][2 * (j % 2) + 1]);
1531
+ frag_s[j / 2][2 * (j % 2) + 1],
1532
+ frag_bias[j / 2][2 * (j % 2) + 1]);
1533
} else {
1534
int wr = c_sh_wr + 8 * j;
1535
write(wr + (4 * c_sh_stride) * 0 + 0, frag_c[i][j][0][0],
1479
- frag_c[i][j][0][1], frag_s[j / 2][2 * (j % 2) + 0]);
1536
+ frag_c[i][j][0][1], frag_s[j / 2][2 * (j % 2) + 0],
1537
+ frag_bias[j / 2][2 * (j % 2) + 0]);
1538
write(wr + (4 * c_sh_stride) * 8 + 0, frag_c[i][j][0][2],
1481
- frag_c[i][j][0][3], frag_s[j / 2][2 * (j % 2) + 0]);
1539
+ frag_c[i][j][0][3], frag_s[j / 2][2 * (j % 2) + 0],
1540
+ frag_bias[j / 2][2 * (j % 2) + 0]);
1541
write(wr + (4 * c_sh_stride) * 0 + 4, frag_c[i][j][1][0],
1483
- frag_c[i][j][1][1], frag_s[j / 2][2 * (j % 2) + 1]);
1542
+ frag_c[i][j][1][1], frag_s[j / 2][2 * (j % 2) + 1],
1543
+ frag_bias[j / 2][2 * (j % 2) + 1]);
1544
write(wr + (4 * c_sh_stride) * 8 + 4, frag_c[i][j][1][2],
1485
- frag_c[i][j][1][3], frag_s[j / 2][2 * (j % 2) + 1]);
1545
+ frag_c[i][j][1][3], frag_s[j / 2][2 * (j % 2) + 1],
1546
+ frag_bias[j / 2][2 * (j % 2) + 1]);
1547
}
1548
}
1549
c_sh_wr += 16 * (4 * c_sh_stride);
@@ -1622,6 +1683,14 @@ __global__ void Marlin(
1683
}
1684
1685
thread_block_reduce();
1686
+
1687
+ if (has_bias && last) {
1688
+ __syncthreads();
1689
+ cp_async4_pred(&sh_bias[bias_sh_wr], &b_bias_ptr[bias_gl_rd],
1690
+ threadIdx.x < 16 * thread_n_blocks / 8);
1691
+ cp_async_fence();
1692
+ }
1693
+
1694
if constexpr (!has_act_order && group_blocks == -1 &&
1695
(has_zp && dequant_skip_flop || !has_zp)) {
1696
if (w_type.size_bits() == 8 || (last || use_atomic_add)) {
@@ -1684,11 +1753,20 @@ __global__ void Marlin(
1753
}
1754
barrier_release(&locks[locks_off], last);
1755
}
1756
+
1757
+ if (has_bias && last) {
1758
+ cp_async_wait<0>();
1759
+ __syncthreads();
1760
+ reinterpret_cast<int4*>(&frag_bias)[0] = sh_bias[bias_sh_rd];
1761
+ reinterpret_cast<int4*>(&frag_bias)[1] = sh_bias[bias_sh_rd + 4];
1762
+ __syncthreads();
1763
+ }
1764
+
1765
if (use_atomic_add && slice_count > 1 && slice_idx != 0)
1766
wait_negative_and_add(&locks[locks_off]);
1767
if (last || use_atomic_add)
1768
// only the last block in a slice actually writes the result
1691
- write_result();
1769
+ write_result(last);
1770
slice_row = 0;
1771
slice_col_par++;
1772
slice_col++;
@@ -1706,6 +1784,7 @@ __global__ void Marlin(
1784
for (int i = 0; i < b_sh_wr_iters; i++) B_ptr[i] -= b_gl_stride;
1785
}
1786
1787
+ bias_gl_rd = (thread_n_blocks * 16 / 8) * slice_col + threadIdx.x;
1788
// Update slice k/n for scales loading
1789
if constexpr (has_act_order) {
1790
slice_k_start = tb_k * slice_row;
csrc/torch_bindings.cpp
+1
@@ -326,6 +326,7 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
326
// gptq_marlin Optimized Quantized GEMM for GPTQ.
327
ops.def(
328
"gptq_marlin_gemm(Tensor a, Tensor? c_or_none, Tensor b_q_weight, "
329
+ "Tensor? b_bias_or_none,"
330
"Tensor b_scales, Tensor? global_scale, Tensor? b_zeros_or_none, Tensor? "
331
"g_idx_or_none, Tensor? perm_or_none, Tensor workspace, int b_q_type, "
332
"SymInt size_m, SymInt size_n, SymInt size_k, bool is_k_full, "
tests/kernels/moe/test_moe.py
+141
-34
@@ -24,8 +24,10 @@ from vllm.model_executor.layers.fused_moe.fused_moe import (
24
fused_topk, modular_triton_fused_moe)
25
from vllm.model_executor.layers.fused_moe.moe_torch_iterative import (
26
fused_moe as iterative_moe)
27
+from vllm.model_executor.layers.quantization.utils.marlin_utils import (
28
+ marlin_permute_bias)
29
from vllm.model_executor.layers.quantization.utils.marlin_utils_fp4 import (
28
- rand_marlin_weight_fp4_like)
30
+ rand_marlin_weight_mxfp4_like, rand_marlin_weight_nvfp4_like)
31
from vllm.model_executor.layers.quantization.utils.marlin_utils_fp8 import (
32
marlin_quant_fp8_torch)
33
from vllm.model_executor.layers.quantization.utils.marlin_utils_test import (
@@ -476,8 +478,11 @@ def marlin_moe_generate_valid_test_cases():
478
if quant_type == scalar_types.float8_e4m3fn and \
479
group_size not in [-1, 128]:
480
return False
479
- if quant_type == scalar_types.float4_e2m1f and group_size != 16:
480
- return False
481
+ if quant_type == scalar_types.float4_e2m1f:
482
+ if group_size not in [16, 32]:
483
+ return False
484
+ if dtype == torch.float16 and group_size == 32:
485
+ return False
486
if quant_type != scalar_types.float4_e2m1f and group_size == 16:
487
return False
488
@@ -520,31 +525,6 @@ def test_fused_marlin_moe(
525
torch.cuda.manual_seed(0)
526
has_zp = quant_type in [scalar_types.uint4, scalar_types.uint8]
527
523
- if quant_type == scalar_types.float8_e4m3fn:
524
- if group_size not in [-1, 128]:
525
- return
526
- if act_order:
527
- return
528
-
529
- # Filter act_order
530
- if act_order:
531
- if quant_type == scalar_types.float8_e4m3fn:
532
- return
533
- if group_size == -1:
534
- return
535
- if group_size in (k, n):
536
- return
537
- if has_zp:
538
- return
539
- else:
540
- if not is_k_full:
541
- return
542
-
543
- if quant_type == scalar_types.float4_e2m1f and group_size != 16:
544
- return
545
- if quant_type != scalar_types.float4_e2m1f and group_size == 16:
546
- return
547
-
528
a = torch.randn((m, k), device="cuda", dtype=dtype) / 10
529
w1 = torch.randn((e, 2 * n, k), device="cuda", dtype=dtype) / 20
530
w2 = torch.randn((e, k, n), device="cuda", dtype=dtype) / 20
@@ -569,13 +549,19 @@ def test_fused_marlin_moe(
549
550
for i in range(w1.shape[0]):
551
if quant_type == scalar_types.float4_e2m1f:
572
- w_ref1, qweight1, scales1, global_scale1 = \
573
- rand_marlin_weight_fp4_like(w1[i], group_size)
552
+ if group_size == 16:
553
+ w_ref1, qweight1, scales1, global_scale1 = \
554
+ rand_marlin_weight_nvfp4_like(w1[i], group_size)
555
+ else:
556
+ w_ref1, qweight1, scales1 = \
557
+ rand_marlin_weight_mxfp4_like(w1[i], group_size)
558
+ global_scale1 = None
559
560
w_ref1_l.append(w_ref1.T)
561
qweight1_l.append(qweight1)
562
scales1_l.append(scales1)
578
- global_scale1_l.append(global_scale1)
563
+ if global_scale1 is not None:
564
+ global_scale1_l.append(global_scale1)
565
elif quant_type == scalar_types.float8_e4m3fn:
566
w_ref1, qweight1, scales1 = marlin_quant_fp8_torch(
567
w1[i], group_size)
@@ -620,13 +606,19 @@ def test_fused_marlin_moe(
606
607
for i in range(w2.shape[0]):
608
if quant_type == scalar_types.float4_e2m1f:
623
- w_ref2, qweight2, scales2, global_scale2 = \
624
- rand_marlin_weight_fp4_like(w2[i], group_size)
609
+ if group_size == 16:
610
+ w_ref2, qweight2, scales2, global_scale2 = \
611
+ rand_marlin_weight_nvfp4_like(w2[i], group_size)
612
+ else:
613
+ w_ref2, qweight2, scales2 = \
614
+ rand_marlin_weight_mxfp4_like(w2[i], group_size)
615
+ global_scale2 = None
616
617
w_ref2_l.append(w_ref2.T)
618
qweight2_l.append(qweight2)
619
scales2_l.append(scales2)
629
- global_scale2_l.append(global_scale2)
620
+ if global_scale2 is not None:
621
+ global_scale2_l.append(global_scale2)
622
elif quant_type == scalar_types.float8_e4m3fn:
623
w_ref2, qweight2, scales2 = marlin_quant_fp8_torch(
624
w2[i], group_size)
@@ -677,6 +669,8 @@ def test_fused_marlin_moe(
669
a,
670
qweight1,
671
qweight2,
672
+ None,
673
+ None,
674
scales1,
675
scales2,
676
score,
@@ -698,6 +692,119 @@ def test_fused_marlin_moe(
692
torch.testing.assert_close(marlin_output, torch_output, atol=5e-2, rtol=0)
693
694
695
+@pytest.mark.flaky(reruns=2)
696
+@pytest.mark.skipif(current_platform.is_rocm(), reason="Skip for rocm")
697
+@pytest.mark.parametrize("m", [1, 256])
698
+def test_fused_marlin_moe_with_bias(m):
699
+ torch.cuda.manual_seed(0)
700
+
701
+ e, topk = 32, 4
702
+ n, k = 2048, 2048
703
+ group_size = 128
704
+ act_order = False
705
+ is_k_full = True
706
+ quant_type = scalar_types.uint4b8
707
+ dtype = torch.half
708
+
709
+ a = torch.randn((m, k), device="cuda", dtype=dtype) / 10
710
+ w1 = torch.randn((e, 2 * n, k), device="cuda", dtype=dtype) / 10
711
+ w2 = torch.randn((e, k, n), device="cuda", dtype=dtype) / 10
712
+ b_bias1 = torch.randn((e, 2 * n), device="cuda", dtype=dtype) / 10
713
+ b_bias2 = torch.randn((e, k), device="cuda", dtype=dtype) / 10
714
+
715
+ b_bias1_l = []
716
+ w_ref1_l = []
717
+ qweight1_l = []
718
+ scales1_l = []
719
+ g_idx1_l = []
720
+ sort_indices1_l = []
721
+
722
+ for i in range(w1.shape[0]):
723
+ test_perm = torch.randperm(k)
724
+ w_ref1, qweight1, scales1, g_idx1, sort_indices1, _ = \
725
+ marlin_quantize(w1[i].transpose(1, 0), quant_type,
726
+ group_size, act_order, test_perm)
727
+
728
+ w_ref1_l.append(w_ref1.T)
729
+ qweight1_l.append(qweight1)
730
+ scales1_l.append(scales1)
731
+ g_idx1_l.append(g_idx1)
732
+ sort_indices1_l.append(sort_indices1)
733
+ b_bias1_l.append(marlin_permute_bias(b_bias1[i]))
734
+
735
+ w_ref1 = stack_and_dev(w_ref1_l)
736
+ qweight1 = stack_and_dev(qweight1_l).contiguous()
737
+ scales1 = stack_and_dev(scales1_l)
738
+ global_scale1 = None
739
+ g_idx1 = stack_and_dev(g_idx1_l) if g_idx1_l else None
740
+ zeros1 = None
741
+ sort_indices1 = stack_and_dev(sort_indices1_l) if sort_indices1_l else None
742
+ marlin_bias1 = stack_and_dev(b_bias1_l) if b_bias1_l else None
743
+
744
+ b_bias2_l = []
745
+ w_ref2_l = []
746
+ qweight2_l = []
747
+ scales2_l = []
748
+ g_idx2_l = []
749
+ sort_indices2_l = []
750
+
751
+ for i in range(w2.shape[0]):
752
+ test_perm = torch.randperm(n)
753
+ w_ref2, qweight2, scales2, g_idx2, sort_indices2, _ = \
754
+ marlin_quantize(w2[i].transpose(1, 0), quant_type,
755
+ group_size, act_order, test_perm)
756
+
757
+ w_ref2_l.append(w_ref2.T)
758
+ qweight2_l.append(qweight2)
759
+ scales2_l.append(scales2)
760
+ g_idx2_l.append(g_idx2)
761
+ sort_indices2_l.append(sort_indices2)
762
+ b_bias2_l.append(marlin_permute_bias(b_bias2[i]))
763
+
764
+ w_ref2 = stack_and_dev(w_ref2_l)
765
+ qweight2 = stack_and_dev(qweight2_l).contiguous()
766
+ scales2 = stack_and_dev(scales2_l)
767
+ global_scale2 = None
768
+ g_idx2 = stack_and_dev(g_idx2_l) if g_idx2_l else None
769
+ zeros2 = None
770
+ sort_indices2 = stack_and_dev(sort_indices2_l) if sort_indices2_l else None
771
+ marlin_bias2 = stack_and_dev(b_bias2_l) if b_bias2_l else None
772
+
773
+ score = torch.randn((m, e), device="cuda", dtype=dtype)
774
+
775
+ topk_weights, topk_ids, _ = fused_topk(a, score, topk, False)
776
+
777
+ with set_current_vllm_config(vllm_config):
778
+ torch_output = torch_moe(a, w_ref1, w_ref2, score, topk, b_bias1,
779
+ b_bias2)
780
+
781
+ marlin_output = torch.ops.vllm.fused_marlin_moe(
782
+ a,
783
+ qweight1,
784
+ qweight2,
785
+ marlin_bias1,
786
+ marlin_bias2,
787
+ scales1,
788
+ scales2,
789
+ score,
790
+ topk_weights,
791
+ topk_ids,
792
+ global_num_experts=e,
793
+ expert_map=None,
794
+ global_scale1=global_scale1,
795
+ global_scale2=global_scale2,
796
+ g_idx1=g_idx1,
797
+ g_idx2=g_idx2,
798
+ sort_indices1=sort_indices1,
799
+ sort_indices2=sort_indices2,
800
+ w1_zeros=zeros1,
801
+ w2_zeros=zeros2,
802
+ quant_type_id=quant_type.id,
803
+ is_k_full=is_k_full)
804
+
805
+ torch.testing.assert_close(marlin_output, torch_output, atol=5e-2, rtol=0)
806
+
807
+
808
def test_moe_align_block_size_opcheck():
809
num_experts = 4
810
block_size = 4
tests/kernels/quantization/test_marlin_gemm.py
+74
-21
@@ -19,10 +19,11 @@ from vllm.model_executor.layers.quantization.qqq import (
19
from vllm.model_executor.layers.quantization.utils.marlin_utils import (
20
GPTQ_MARLIN_MAX_PARALLEL, GPTQ_MARLIN_MIN_THREAD_N,
21
MARLIN_SUPPORTED_GROUP_SIZES, marlin_make_empty_g_idx,
22
- marlin_make_workspace_new, marlin_permute_scales,
22
+ marlin_make_workspace_new, marlin_permute_bias, marlin_permute_scales,
23
query_marlin_supported_quant_types)
24
from vllm.model_executor.layers.quantization.utils.marlin_utils_fp4 import (
25
- FP4_MARLIN_SUPPORTED_GROUP_SIZES, rand_marlin_weight_fp4_like)
25
+ FP4_MARLIN_SUPPORTED_GROUP_SIZES, rand_marlin_weight_mxfp4_like,
26
+ rand_marlin_weight_nvfp4_like)
27
from vllm.model_executor.layers.quantization.utils.marlin_utils_fp8 import (
28
marlin_quant_fp8_torch)
29
from vllm.model_executor.layers.quantization.utils.marlin_utils_test import (
@@ -39,7 +40,7 @@ from vllm.scalar_type import scalar_types
40
ACT_ORDER_OPTS = [False, True]
41
K_FULL_OPTS = [False, True]
42
USE_ATOMIC_ADD_OPTS = [False, True]
42
-USE_FP32_REDUCE_OPTS = [False, True]
43
+USE_FP32_REDUCE_OPTS = [True]
44
45
MARLIN_K_CHUNKS = [128]
46
MARLIN_N_CHUNKS = [64, 256]
@@ -202,17 +203,10 @@ def test_awq_marlin_repack(k_chunk, n_chunk, quant_type, group_size,
203
@pytest.mark.parametrize("is_k_full", K_FULL_OPTS)
204
@pytest.mark.parametrize("use_atomic_add", USE_ATOMIC_ADD_OPTS)
205
@pytest.mark.parametrize("use_fp32_reduce", USE_FP32_REDUCE_OPTS)
205
-def test_gptq_marlin_gemm(
206
- k_chunk,
207
- n_chunk,
208
- quant_type,
209
- group_size,
210
- mnk_factors,
211
- act_order,
212
- is_k_full,
213
- use_atomic_add,
214
- use_fp32_reduce,
215
-):
206
+@pytest.mark.parametrize("dtype", DTYPES)
207
+def test_gptq_marlin_gemm(k_chunk, n_chunk, quant_type, group_size,
208
+ mnk_factors, act_order, is_k_full, use_atomic_add,
209
+ use_fp32_reduce, dtype):
210
m_factor, n_factor, k_factor = mnk_factors
211
has_zp = quant_type in [scalar_types.uint4, scalar_types.uint8]
212
@@ -231,14 +225,23 @@ def test_gptq_marlin_gemm(
225
if size_k % group_size != 0:
226
return
227
234
- a_input = rand_data((size_m, size_k))
235
- b_weight = rand_data((size_k, size_n))
228
+ a_input = rand_data((size_m, size_k), dtype)
229
+ b_weight = rand_data((size_k, size_n), dtype)
230
231
if quant_type == scalar_types.float4_e2m1f:
238
- if group_size != 16 or act_order:
232
+ if group_size not in [16, 32] or act_order:
233
return
240
- w_ref, marlin_q_w, marlin_s, marlin_s2 = rand_marlin_weight_fp4_like(
241
- b_weight.T, group_size)
234
+ if group_size == 32 and dtype == torch.float16:
235
+ return
236
+
237
+ if group_size == 16:
238
+ w_ref, marlin_q_w, marlin_s, marlin_s2 = \
239
+ rand_marlin_weight_nvfp4_like(b_weight.T, group_size)
240
+ else:
241
+ w_ref, marlin_q_w, marlin_s = \
242
+ rand_marlin_weight_mxfp4_like(b_weight.T, group_size)
243
+ marlin_s2 = None
244
+
245
g_idx = None
246
sort_indices = None
247
marlin_zp = None
@@ -272,8 +275,8 @@ def test_gptq_marlin_gemm(
275
workspace = marlin_make_workspace_new(w_ref.device)
276
277
opcheck(torch.ops._C.gptq_marlin_gemm,
275
- (a_input, None, marlin_q_w, marlin_s, marlin_s2, marlin_zp, g_idx,
276
- sort_indices, workspace, quant_type.id, a_input.shape[0],
278
+ (a_input, None, marlin_q_w, None, marlin_s, marlin_s2, marlin_zp,
279
+ g_idx, sort_indices, workspace, quant_type.id, a_input.shape[0],
280
b_weight.shape[1], a_input.shape[1], is_k_full, use_atomic_add,
281
use_fp32_reduce, False),
282
test_utils=DEFAULT_OPCHECK_TEST_UTILS)
@@ -282,6 +285,7 @@ def test_gptq_marlin_gemm(
285
a_input,
286
None,
287
marlin_q_w,
288
+ None,
289
marlin_s,
290
marlin_s2,
291
marlin_zp,
@@ -418,6 +422,7 @@ def test_hqq_marlin_gemm(
422
a_input,
423
None,
424
marlin_w_q,
425
+ None,
426
marlin_s,
427
None,
428
marlin_zp,
@@ -531,6 +536,7 @@ def test_marlin_gemm_subset_input():
536
a_input,
537
None,
538
marlin_q_w,
539
+ None,
540
marlin_s,
541
None,
542
marlin_zp,
@@ -555,6 +561,53 @@ def test_marlin_gemm_subset_input():
561
assert max_diff < 0.04
562
563
564
+@pytest.mark.parametrize("size_m", [1, 256])
565
+def test_marlin_gemm_with_bias(size_m):
566
+ quant_type = scalar_types.uint4b8
567
+ group_size = 128
568
+
569
+ size_k, size_n = 1024, 2048
570
+ a_input = rand_data((size_m, size_k))
571
+ b_weight = rand_data((size_k, size_n))
572
+ b_bias = rand_data((size_n, )) * 10
573
+
574
+ marlin_bias = marlin_permute_bias(b_bias)
575
+
576
+ w_ref, marlin_q_w, marlin_s, g_idx, sort_indices, _ = marlin_quantize(
577
+ b_weight, quant_type, group_size, False)
578
+
579
+ marlin_zp = marlin_make_empty_g_idx(marlin_s.device)
580
+ workspace = marlin_make_workspace_new(a_input.device)
581
+
582
+ output = ops.gptq_marlin_gemm(
583
+ a_input,
584
+ None,
585
+ marlin_q_w,
586
+ marlin_bias,
587
+ marlin_s,
588
+ None,
589
+ marlin_zp,
590
+ g_idx,
591
+ sort_indices,
592
+ workspace,
593
+ quant_type,
594
+ a_input.shape[0],
595
+ b_weight.shape[1],
596
+ a_input.shape[1],
597
+ is_k_full=True,
598
+ use_atomic_add=False,
599
+ use_fp32_reduce=True,
600
+ is_zp_float=False,
601
+ )
602
+ output_ref = torch.matmul(a_input, w_ref) + b_bias.view(1, -1)
603
+
604
+ torch.cuda.synchronize()
605
+
606
+ max_diff = compute_max_diff(output, output_ref)
607
+
608
+ assert max_diff < 0.04
609
+
610
+
611
def test_marlin_gemm_opcheck():
612
size_m = 2048
613
size_n = 4096
tests/kernels/utils.py
+20
-1
@@ -1064,6 +1064,8 @@ def torch_experts(
1064
topk_weight: torch.Tensor,
1065
topk_ids: torch.Tensor,
1066
global_num_experts: int = -1,
1067
+ b_bias1: Optional[torch.Tensor] = None,
1068
+ b_bias2: Optional[torch.Tensor] = None,
1069
expert_map: Optional[torch.Tensor] = None,
1070
w1_scale: Optional[torch.Tensor] = None,
1071
w2_scale: Optional[torch.Tensor] = None,
@@ -1108,8 +1110,13 @@ def torch_experts(
1110
if mask.sum():
1111
if quant_dtype is None:
1112
tmp1 = a[mask] @ w1[i].transpose(0, 1)
1113
+ if b_bias1 is not None:
1114
+ tmp1 = tmp1 + b_bias1[i].view(1, -1).to(tmp1.dtype)
1115
tmp2 = SiluAndMul()(tmp1)
1116
out[mask] = tmp2 @ w2[i].transpose(0, 1)
1117
+ if b_bias2 is not None:
1118
+ out[mask] = out[mask] + b_bias2[i].view(1, -1).to(
1119
+ tmp1.dtype)
1120
elif block_shape is not None:
1121
# block quantized
1122
assert (a_scale is not None and w1_scale is not None
@@ -1117,6 +1124,8 @@ def torch_experts(
1124
tmp1 = native_w8a8_block_matmul(a[mask], w1[i], a_scale[mask],
1125
w1_scale[i], block_shape,
1126
out.dtype)
1127
+ if b_bias1 is not None:
1128
+ tmp1 = tmp1 + b_bias1[i].view(1, -1).to(tmp1.dtype)
1129
tmp2 = SiluAndMul()(tmp1)
1130
tmp2, b_scale = moe_kernel_quantize_input(
1131
tmp2, a2_scale, quant_dtype, per_act_token_quant,
@@ -1125,6 +1134,9 @@ def torch_experts(
1134
out[mask] = native_w8a8_block_matmul(tmp2, w2[i], b_scale,
1135
w2_scale[i], block_shape,
1136
out.dtype)
1137
+ if b_bias2 is not None:
1138
+ out[mask] = out[mask] + b_bias2[i].view(1, -1).to(
1139
+ tmp1.dtype)
1140
else:
1141
assert (a_scale is not None and w1_scale is not None
1142
and w2_scale is not None)
@@ -1133,6 +1145,8 @@ def torch_experts(
1145
tmp1 = a[mask].to(f32) * scales
1146
w1_dq = (w1[i].to(f32) * w1_scale[i]).transpose(0, 1)
1147
tmp1 = (tmp1 @ w1_dq).to(out.dtype)
1148
+ if b_bias1 is not None:
1149
+ tmp1 = tmp1 + b_bias1[i].view(1, -1).to(out.dtype)
1150
1151
tmp2 = SiluAndMul()(tmp1).to(out.dtype)
1152
@@ -1144,6 +1158,9 @@ def torch_experts(
1158
tmp2 = tmp2.to(f32) * b_scale
1159
w2_dq = (w2[i].to(f32) * w2_scale[i]).transpose(0, 1)
1160
out[mask] = (tmp2 @ w2_dq).to(out.dtype)
1161
+ if b_bias2 is not None:
1162
+ out[mask] = out[mask] + b_bias2[i].view(1, -1).to(
1163
+ out.dtype)
1164
1165
if apply_router_weights_on_input:
1166
return out
@@ -1157,12 +1174,14 @@ def torch_moe(a: torch.Tensor,
1174
w2: torch.Tensor,
1175
score: torch.Tensor,
1176
topk: int,
1177
+ b_bias1: Optional[torch.Tensor] = None,
1178
+ b_bias2: Optional[torch.Tensor] = None,
1179
global_num_experts: int = -1,
1180
expert_map: Optional[torch.Tensor] = None) -> torch.Tensor:
1181
score = torch.softmax(score, dim=-1, dtype=torch.float32)
1182
topk_weight, topk_ids = torch.topk(score, topk)
1183
return torch_experts(a, w1, w2, topk_weight, topk_ids, global_num_experts,
1165
- expert_map)
1184
+ b_bias1, b_bias2, expert_map)
1185
1186
1187
def torch_moe_single(a, w, score, topk):
vllm/_custom_ops.py
+11
-7
@@ -452,6 +452,7 @@ if hasattr(torch.ops._C, "gptq_marlin_24_gemm"):
452
def _gptq_marlin_gemm_fake(a: torch.Tensor,
453
c: Optional[torch.Tensor],
454
b_q_weight: torch.Tensor,
455
+ b_bias: Optional[torch.Tensor],
456
b_scales: torch.Tensor,
457
global_scale: Optional[torch.Tensor],
458
b_zeros: Optional[torch.Tensor],
@@ -1048,6 +1049,7 @@ def awq_marlin_moe_repack(b_q_weight: torch.Tensor, perm: torch.Tensor,
1049
def gptq_marlin_gemm(a: torch.Tensor,
1050
c: Optional[torch.Tensor],
1051
b_q_weight: torch.Tensor,
1052
+ b_bias: Optional[torch.Tensor],
1053
b_scales: torch.Tensor,
1054
global_scale: Optional[torch.Tensor],
1055
b_zeros: Optional[torch.Tensor],
@@ -1062,7 +1064,7 @@ def gptq_marlin_gemm(a: torch.Tensor,
1064
use_atomic_add: bool = False,
1065
use_fp32_reduce: bool = False,
1066
is_zp_float: bool = False) -> torch.Tensor:
1065
- return torch.ops._C.gptq_marlin_gemm(a, c, b_q_weight, b_scales,
1067
+ return torch.ops._C.gptq_marlin_gemm(a, c, b_q_weight, b_bias, b_scales,
1068
global_scale, b_zeros, g_idx, perm,
1069
workspace, b_q_type.id, size_m,
1070
size_n, size_k, is_k_full,
@@ -1540,7 +1542,9 @@ def topk_softmax(topk_weights: torch.Tensor, topk_ids: torch.Tensor,
1542
1543
1544
def moe_wna16_marlin_gemm(input: torch.Tensor, output: Optional[torch.Tensor],
1543
- b_qweight: torch.Tensor, b_scales: torch.Tensor,
1545
+ b_qweight: torch.Tensor,
1546
+ b_bias: Optional[torch.Tensor],
1547
+ b_scales: torch.Tensor,
1548
global_scale: Optional[torch.Tensor],
1549
b_qzeros: Optional[torch.Tensor],
1550
g_idx: Optional[torch.Tensor],
@@ -1556,11 +1560,11 @@ def moe_wna16_marlin_gemm(input: torch.Tensor, output: Optional[torch.Tensor],
1560
use_fp32_reduce: bool,
1561
is_zp_float: bool) -> torch.Tensor:
1562
return torch.ops._moe_C.moe_wna16_marlin_gemm(
1559
- input, output, b_qweight, b_scales, global_scale, b_qzeros, g_idx,
1560
- perm, workspace, sorted_token_ids, expert_ids, num_tokens_past_padded,
1561
- topk_weights, moe_block_size, top_k, mul_topk_weights, is_ep,
1562
- b_q_type.id, size_m, size_n, size_k, is_k_full, use_atomic_add,
1563
- use_fp32_reduce, is_zp_float)
1563
+ input, output, b_qweight, b_bias, b_scales, global_scale, b_qzeros,
1564
+ g_idx, perm, workspace, sorted_token_ids, expert_ids,
1565
+ num_tokens_past_padded, topk_weights, moe_block_size, top_k,
1566
+ mul_topk_weights, is_ep, b_q_type.id, size_m, size_n, size_k,
1567
+ is_k_full, use_atomic_add, use_fp32_reduce, is_zp_float)
1568
1569
1570
if supports_moe_ops and hasattr(torch.ops._moe_C, "marlin_gemm_moe"):
vllm/envs.py
+11
@@ -122,6 +122,7 @@ if TYPE_CHECKING:
122
VLLM_MOE_DP_CHUNK_SIZE: int = 256
123
VLLM_RANDOMIZE_DP_DUMMY_INPUTS: bool = False
124
VLLM_MARLIN_USE_ATOMIC_ADD: bool = False
125
+ VLLM_MXFP4_USE_MARLIN: Optional[bool] = None
126
VLLM_V0_USE_OUTLINES_CACHE: bool = False
127
VLLM_V1_USE_OUTLINES_CACHE: bool = False
128
VLLM_TPU_BUCKET_PADDING_GAP: int = 0
@@ -182,6 +183,12 @@ def maybe_convert_int(value: Optional[str]) -> Optional[int]:
183
return int(value)
184
185
186
+def maybe_convert_bool(value: Optional[str]) -> Optional[bool]:
187
+ if value is None:
188
+ return None
189
+ return bool(int(value))
190
+
191
+
192
def get_vllm_port() -> Optional[int]:
193
"""Get the port from VLLM_PORT environment variable.
194
@@ -906,6 +913,10 @@ environment_variables: dict[str, Callable[[], Any]] = {
913
"VLLM_MARLIN_USE_ATOMIC_ADD":
914
lambda: os.environ.get("VLLM_MARLIN_USE_ATOMIC_ADD", "0") == "1",
915
916
+ # Whether to use marlin kernel in mxfp4 quantization method
917
+ "VLLM_MXFP4_USE_MARLIN":
918
+ lambda: maybe_convert_bool(os.environ.get("VLLM_MXFP4_USE_MARLIN", None)),
919
+
920
# Whether to turn on the outlines cache for V0
921
# This cache is unbounded and on disk, so it's not safe to use in
922
# an environment with potentially malicious users.
vllm/model_executor/layers/fused_moe/fused_marlin_moe.py
+28
-2
@@ -18,6 +18,8 @@ from vllm.utils import direct_register_custom_op
18
def fused_marlin_moe(hidden_states: torch.Tensor,
19
w1: torch.Tensor,
20
w2: torch.Tensor,
21
+ bias1: Optional[torch.Tensor],
22
+ bias2: Optional[torch.Tensor],
23
w1_scale: torch.Tensor,
24
w2_scale: torch.Tensor,
25
gating_output: torch.Tensor,
@@ -26,6 +28,7 @@ def fused_marlin_moe(hidden_states: torch.Tensor,
28
quant_type_id: int,
29
apply_router_weight_on_input: bool = False,
30
global_num_experts: int = -1,
31
+ activation: Optional[str] = "silu",
32
expert_map: Optional[torch.Tensor] = None,
33
global_scale1: Optional[torch.Tensor] = None,
34
global_scale2: Optional[torch.Tensor] = None,
@@ -88,6 +91,7 @@ def fused_marlin_moe(hidden_states: torch.Tensor,
91
assert w2.is_contiguous(), "Expert weights2 must be contiguous"
92
assert hidden_states.dtype in [torch.float16, torch.bfloat16]
93
assert num_bits in [4, 8]
94
+ assert topk_weights.dtype == torch.float32
95
96
M, K = hidden_states.shape
97
E = w1.shape[0]
@@ -138,6 +142,7 @@ def fused_marlin_moe(hidden_states: torch.Tensor,
142
hidden_states,
143
intermediate_cache1,
144
w1,
145
+ bias1,
146
w1_scale,
147
global_scale1,
148
w1_zeros,
@@ -161,8 +166,28 @@ def fused_marlin_moe(hidden_states: torch.Tensor,
166
use_fp32_reduce=True,
167
is_zp_float=False)
168
164
- torch.ops._C.silu_and_mul(intermediate_cache2,
165
- intermediate_cache1.view(-1, 2 * N))
169
+ if activation == "silu":
170
+ torch.ops._C.silu_and_mul(intermediate_cache2,
171
+ intermediate_cache1.view(-1, 2 * N))
172
+ elif activation == "swiglu_oai":
173
+ # NOTE: in gpt-oss, the gate_proj and up_proj is interleaved
174
+ # - interleaved: gate, up = gate_up[..., ::2], gate_up[..., 1::2]
175
+ # - origin: gate, up = gate_up[..., :N], gate_up[..., N:]
176
+
177
+ @torch.compile(dynamic=True)
178
+ def swiglu_oai(gate_up):
179
+ alpha = 1.702
180
+ limit = 7.0
181
+ gate, up = gate_up[..., ::2], gate_up[..., 1::2]
182
+ gate = gate.clamp(min=None, max=limit)
183
+ up = up.clamp(min=-limit, max=limit)
184
+ glu = gate * torch.sigmoid(gate * alpha)
185
+ return (up + 1) * glu
186
+
187
+ intermediate_cache2 = swiglu_oai(intermediate_cache1)
188
+ else:
189
+ raise ValueError(f"Unsupported activation: {activation}. "
190
+ "Only silu and swiglu_oai activations are supported.")
191
192
if expert_map is not None:
193
intermediate_cache3.zero_()
@@ -171,6 +196,7 @@ def fused_marlin_moe(hidden_states: torch.Tensor,
196
intermediate_cache2,
197
intermediate_cache3,
198
w2,
199
+ bias2,
200
w2_scale,
201
global_scale2,
202
w2_zeros,
vllm/model_executor/layers/fused_moe/layer.py
+6
-14
@@ -36,7 +36,7 @@ from vllm.model_executor.utils import set_weight_attrs
36
from vllm.platforms import current_platform
37
from vllm.platforms.interface import CpuArchEnum
38
from vllm.utils import (direct_register_custom_op, has_deep_ep, has_pplx,
39
- has_triton_kernels, is_torch_equal_or_newer, round_up)
39
+ round_up)
40
from vllm.utils.flashinfer import has_flashinfer
41
42
if current_platform.is_cuda_alike():
@@ -751,19 +751,11 @@ class FusedMoE(CustomOp):
751
self.global_num_experts = num_experts + num_redundant_experts
752
753
# we padding globally so EP buffer allocation works
754
- if quant_config and quant_config.get_name() == "mxfp4":
755
- if not current_platform.is_device_capability(100):
756
- if not is_torch_equal_or_newer("2.8.0"):
757
- raise RuntimeError(
758
- "Mxfp4 on non-blackwell requires torch >= 2.8.0")
759
- if not has_triton_kernels():
760
- raise NotImplementedError(
761
- "triton_kernels must be installed for "
762
- "mxfp4 on non-blackwell")
763
- if (current_platform.is_rocm()
764
- or envs.VLLM_USE_FLASHINFER_MOE_MXFP4_MXFP8
765
- or envs.VLLM_USE_FLASHINFER_MOE_MXFP4_BF16):
766
- hidden_size = round_up(hidden_size, 256)
754
+ if (quant_config and quant_config.get_name() == "mxfp4"
755
+ and (current_platform.is_rocm()
756
+ or envs.VLLM_USE_FLASHINFER_MOE_MXFP4_MXFP8
757
+ or envs.VLLM_USE_FLASHINFER_MOE_MXFP4_BF16)):
758
+ hidden_size = round_up(hidden_size, 256)
759
760
# For smuggling this layer into the fused moe custom op
761
compilation_config = vllm_config.compilation_config
vllm/model_executor/layers/quantization/awq_marlin.py
+12
-1
@@ -25,7 +25,7 @@ from vllm.model_executor.layers.quantization.utils.marlin_utils import (
25
apply_awq_marlin_linear, awq_to_marlin_zero_points, check_marlin_supported,
26
check_marlin_supports_layer, check_moe_marlin_supports_layer,
27
marlin_make_empty_g_idx, marlin_make_workspace_new,
28
- marlin_moe_permute_scales, marlin_permute_scales,
28
+ marlin_moe_permute_scales, marlin_permute_bias, marlin_permute_scales,
29
moe_awq_to_marlin_zero_points, verify_marlin_supported,
30
verify_marlin_supports_shape)
31
from vllm.model_executor.layers.vocab_parallel_embedding import ParallelLMHead
@@ -303,6 +303,9 @@ class AWQMarlinLinearMethod(LinearMethodBase):
303
layer.g_idx = marlin_make_empty_g_idx(device)
304
layer.g_idx_sort_indices = marlin_make_empty_g_idx(device)
305
306
+ if hasattr(layer, "bias") and layer.bias is not None:
307
+ layer.bias.data = marlin_permute_bias(layer.bias)
308
+
309
def apply(
310
self,
311
layer: torch.nn.Module,
@@ -469,6 +472,12 @@ class AWQMoEMethod(FusedMoEMethodBase):
472
num_bits=self.quant_config.weight_bits)
473
replace_parameter(layer, "w2_qzeros", marlin_w2_zp)
474
475
+ if hasattr(layer, "w13_bias") and layer.w13_bias is not None:
476
+ layer.w13_bias.data = marlin_permute_bias(layer.w13_bias)
477
+
478
+ if hasattr(layer, "w2_bias") and layer.w2_bias is not None:
479
+ layer.w2_bias.data = marlin_permute_bias(layer.w2_bias)
480
+
481
def apply(
482
self,
483
layer: torch.nn.Module,
@@ -513,6 +522,8 @@ class AWQMoEMethod(FusedMoEMethodBase):
522
x,
523
layer.w13_qweight,
524
layer.w2_qweight,
525
+ getattr(layer, "w13_bias", None),
526
+ getattr(layer, "w2_bias", None),
527
layer.w13_scales,
528
layer.w2_scales,
529
router_logits,
vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe.py
+6
@@ -324,6 +324,8 @@ class CompressedTensorsW4A4MoeMethod(CompressedTensorsMoEMethod):
324
x,
325
layer.w13_weight,
326
layer.w2_weight,
327
+ None,
328
+ None,
329
layer.w13_weight_scale,
330
layer.w2_weight_scale,
331
router_logits,
@@ -795,6 +797,8 @@ class CompressedTensorsW8A8Fp8MoEMethod(CompressedTensorsMoEMethod):
797
x,
798
layer.w13_weight,
799
layer.w2_weight,
800
+ None,
801
+ None,
802
layer.w13_weight_scale,
803
layer.w2_weight_scale,
804
router_logits,
@@ -1253,6 +1257,8 @@ class CompressedTensorsWNA16MarlinMoEMethod(CompressedTensorsMoEMethod):
1257
x,
1258
layer.w13_weight_packed,
1259
layer.w2_weight_packed,
1260
+ None,
1261
+ None,
1262
layer.w13_weight_scale,
1263
layer.w2_weight_scale,
1264
router_logits,
vllm/model_executor/layers/quantization/fp8.py
+2
@@ -983,6 +983,8 @@ class Fp8MoEMethod(FusedMoEMethodBase):
983
x,
984
layer.w13_weight,
985
layer.w2_weight,
986
+ None,
987
+ None,
988
layer.w13_weight_scale,
989
layer.w2_weight_scale,
990
router_logits,
vllm/model_executor/layers/quantization/gptq_marlin.py
+9
-1
@@ -24,7 +24,7 @@ from vllm.model_executor.layers.quantization.utils.gptq_utils import (
24
get_dynamic_override, get_linear_quant_method, override_config)
25
from vllm.model_executor.layers.quantization.utils.marlin_utils import (
26
check_marlin_supported, check_moe_marlin_supports_layer,
27
- marlin_make_workspace_new, marlin_moe_permute_scales,
27
+ marlin_make_workspace_new, marlin_moe_permute_scales, marlin_permute_bias,
28
marlin_repeat_scales_on_all_ranks, verify_marlin_supported)
29
from vllm.model_executor.parameter import (ChannelQuantScaleParameter,
30
GroupQuantScaleParameter,
@@ -618,6 +618,12 @@ class GPTQMarlinMoEMethod(FusedMoEMethodBase):
618
)
619
replace_parameter(layer, "w2_scales", marlin_w2_scales)
620
621
+ if hasattr(layer, "w13_bias") and layer.w13_bias is not None:
622
+ layer.w13_bias.data = marlin_permute_bias(layer.w13_bias)
623
+
624
+ if hasattr(layer, "w2_bias") and layer.w2_bias is not None:
625
+ layer.w2_bias.data = marlin_permute_bias(layer.w2_bias)
626
+
627
def apply(
628
self,
629
layer: torch.nn.Module,
@@ -662,6 +668,8 @@ class GPTQMarlinMoEMethod(FusedMoEMethodBase):
668
x,
669
layer.w13_qweight,
670
layer.w2_qweight,
671
+ getattr(layer, "w13_bias", None),
672
+ getattr(layer, "w2_bias", None),
673
layer.w13_scales,
674
layer.w2_scales,
675
router_logits,
vllm/model_executor/layers/quantization/hqq_marlin.py
+5
-4
@@ -14,7 +14,7 @@ from vllm.model_executor.layers.quantization.base_config import (
14
QuantizationConfig, QuantizeMethodBase)
15
from vllm.model_executor.layers.quantization.utils.marlin_utils import (
16
GPTQ_MARLIN_MAX_PARALLEL, GPTQ_MARLIN_MIN_THREAD_N,
17
- marlin_make_empty_g_idx, marlin_permute_scales)
17
+ marlin_make_empty_g_idx, marlin_permute_bias, marlin_permute_scales)
18
from vllm.model_executor.layers.quantization.utils.marlin_utils_test import (
19
MarlinWorkspace)
20
from vllm.model_executor.layers.quantization.utils.quant_utils import gptq_pack
@@ -284,6 +284,9 @@ class HQQMarlinMethod(LinearMethodBase):
284
layer.marlin_zeros = marlin_zp
285
layer.marlin_scales = marlin_s
286
287
+ if hasattr(layer, "bias") and layer.bias is not None:
288
+ layer.bias.data = marlin_permute_bias(layer.bias)
289
+
290
def apply(
291
self,
292
layer: torch.nn.Module,
@@ -307,6 +310,7 @@ class HQQMarlinMethod(LinearMethodBase):
310
x,
311
None,
312
layer.marlin_qweight,
313
+ bias,
314
scales,
315
None,
316
zeros,
@@ -326,7 +330,4 @@ class HQQMarlinMethod(LinearMethodBase):
330
if orig_type != torch.float16:
331
marlin_out = marlin_out.to(orig_type)
332
329
- if bias is not None:
330
- marlin_out.add_(bias)
331
-
333
return marlin_out
vllm/model_executor/layers/quantization/kernels/mixed_precision/marlin.py
+6
-2
@@ -9,8 +9,9 @@ from vllm import _custom_ops as ops
9
from vllm.model_executor.layers.quantization.utils.marlin_utils import (
10
MARLIN_SUPPORTED_GROUP_SIZES, apply_gptq_marlin_linear,
11
check_marlin_supports_shape, marlin_is_k_full, marlin_make_empty_g_idx,
12
- marlin_make_workspace_new, marlin_permute_scales, marlin_sort_g_idx,
13
- marlin_zero_points, query_marlin_supported_quant_types, unpack_cols)
12
+ marlin_make_workspace_new, marlin_permute_bias, marlin_permute_scales,
13
+ marlin_sort_g_idx, marlin_zero_points, query_marlin_supported_quant_types,
14
+ unpack_cols)
15
from vllm.model_executor.parameter import (BasevLLMParameter,
16
permute_param_layout_)
17
from vllm.platforms import current_platform
@@ -111,6 +112,9 @@ class MarlinLinearKernel(MPLinearKernel):
112
self._transform_param(layer, self.w_q_name, transform_w_q)
113
self._transform_param(layer, self.w_s_name, transform_w_s)
114
115
+ if hasattr(layer, "bias") and layer.bias is not None:
116
+ layer.bias.data = marlin_permute_bias(layer.bias)
117
+
118
def apply_weights(self,
119
layer: torch.nn.Module,
120
x: torch.Tensor,
vllm/model_executor/layers/quantization/modelopt.py
+2
@@ -1330,6 +1330,8 @@ class ModelOptNvFp4FusedMoE(FusedMoEMethodBase):
1330
x,
1331
layer.w13_weight,
1332
layer.w2_weight,
1333
+ None,
1334
+ None,
1335
layer.w13_weight_scale,
1336
layer.w2_weight_scale,
1337
router_logits,
vllm/model_executor/layers/quantization/mxfp4.py
+81
-10
@@ -15,13 +15,17 @@ from vllm.model_executor.layers.linear import (LinearBase,
15
from vllm.model_executor.layers.quantization import QuantizationMethods
16
from vllm.model_executor.layers.quantization.base_config import (
17
QuantizationConfig, QuantizeMethodBase)
18
+from vllm.model_executor.layers.quantization.utils.marlin_utils_fp4 import (
19
+ prepare_moe_fp4_layer_for_marlin)
20
from vllm.model_executor.layers.quantization.utils.mxfp4_utils import (
21
_can_support_mxfp4, _swizzle_mxfp4)
22
from vllm.model_executor.layers.quantization.utils.quant_utils import (
23
is_layer_skipped)
24
from vllm.model_executor.utils import set_weight_attrs
25
from vllm.platforms import current_platform
24
-from vllm.utils import next_power_of_2, round_up
26
+from vllm.scalar_type import scalar_types
27
+from vllm.utils import (has_triton_kernels, is_torch_equal_or_newer,
28
+ next_power_of_2, round_up)
29
30
if (envs.VLLM_USE_FLASHINFER_MOE_MXFP4_MXFP8
31
or envs.VLLM_USE_FLASHINFER_MOE_MXFP4_BF16):
@@ -81,6 +85,21 @@ class Mxfp4MoEMethod(FusedMoEMethodBase):
85
super().__init__()
86
self.topk_indices_dtype = None
87
self.moe = moe
88
+ self.use_marlin = self._should_use_marlin()
89
+
90
+ def _should_use_marlin(self):
91
+ if envs.VLLM_MXFP4_USE_MARLIN is not None:
92
+ return envs.VLLM_MXFP4_USE_MARLIN
93
+ if current_platform.is_cuda() and \
94
+ not current_platform.has_device_capability(100):
95
+ if not current_platform.is_device_capability(90):
96
+ # marlin kernel has better performance on ampere
97
+ return True
98
+ if not has_triton_kernels():
99
+ return True
100
+ if not is_torch_equal_or_newer("2.8.0"):
101
+ return True
102
+ return False
103
104
def create_weights(self, layer: torch.nn.Module, num_experts: int,
105
hidden_size: int, intermediate_size_per_partition: int,
@@ -101,11 +120,29 @@ class Mxfp4MoEMethod(FusedMoEMethodBase):
120
121
intermediate_size_per_partition_after_pad = \
122
intermediate_size_per_partition
104
- # pad the intermediate size to be a multiple of 2 * mxfp4_block
105
- # for to hold non-uniform sharded tensor as well as swizzling
106
- # other padding to increase performance
107
- if (envs.VLLM_USE_FLASHINFER_MOE_MXFP4_MXFP8
108
- or envs.VLLM_USE_FLASHINFER_MOE_MXFP4_BF16):
123
+ if self.use_marlin:
124
+ # The moe marlin kernel requires that for each linear
125
+ # n % 256 == 0 and k % 128 == 0.
126
+ # In gate_up_proj:
127
+ # n = 2 * intermediate_size_per_partition_after_pad
128
+ # k = hidden_size
129
+ # In down_proj
130
+ # n = hidden_size
131
+ # k = intermediate_size_per_partition_after_pad
132
+ intermediate_size_per_partition_after_pad = round_up(
133
+ intermediate_size_per_partition, 128)
134
+ hidden_size = round_up(hidden_size, 256)
135
+
136
+ layer.params_dtype = params_dtype
137
+ layer.num_experts = num_experts
138
+ layer.hidden_size = hidden_size
139
+ layer.intermediate_size_per_partition = \
140
+ intermediate_size_per_partition_after_pad
141
+ elif (envs.VLLM_USE_FLASHINFER_MOE_MXFP4_MXFP8
142
+ or envs.VLLM_USE_FLASHINFER_MOE_MXFP4_BF16):
143
+ # pad the intermediate size to be a multiple of 2 * mxfp4_block
144
+ # for to hold non-uniform sharded tensor as well as swizzling
145
+ # other padding to increase performance
146
intermediate_size_per_partition_after_pad = round_up(
147
intermediate_size_per_partition, 256)
148
hidden_size = round_up(hidden_size, 256)
@@ -191,8 +228,10 @@ class Mxfp4MoEMethod(FusedMoEMethodBase):
228
set_weight_attrs(w2_bias, extra_weight_attrs)
229
230
def process_weights_after_loading(self, layer):
194
- if (envs.VLLM_USE_FLASHINFER_MOE_MXFP4_MXFP8
195
- or envs.VLLM_USE_FLASHINFER_MOE_MXFP4_BF16):
231
+ if self.use_marlin:
232
+ prepare_moe_fp4_layer_for_marlin(layer)
233
+ elif (envs.VLLM_USE_FLASHINFER_MOE_MXFP4_MXFP8
234
+ or envs.VLLM_USE_FLASHINFER_MOE_MXFP4_BF16):
235
layer.gemm1_alpha = Parameter(torch.tensor(
236
[1.702] * self.num_experts, dtype=torch.float32).cuda(),
237
requires_grad=False)
@@ -399,13 +438,45 @@ class Mxfp4MoEMethod(FusedMoEMethodBase):
438
if enable_eplb:
439
raise NotImplementedError("EPLB is not supported for mxfp4")
440
441
+ if self.use_marlin:
442
+ topk_weights, topk_ids = FusedMoE.select_experts(
443
+ hidden_states=x,
444
+ router_logits=router_logits,
445
+ use_grouped_topk=use_grouped_topk,
446
+ top_k=top_k,
447
+ renormalize=renormalize,
448
+ topk_group=topk_group,
449
+ num_expert_group=num_expert_group,
450
+ custom_routing_function=custom_routing_function,
451
+ scoring_func=scoring_func,
452
+ e_score_correction_bias=e_score_correction_bias)
453
+
454
+ return torch.ops.vllm.fused_marlin_moe(
455
+ x,
456
+ layer.w13_weight,
457
+ layer.w2_weight,
458
+ layer.w13_bias,
459
+ layer.w2_bias,
460
+ layer.w13_weight_scale,
461
+ layer.w2_weight_scale,
462
+ router_logits,
463
+ topk_weights,
464
+ topk_ids,
465
+ global_scale1=None,
466
+ global_scale2=None,
467
+ quant_type_id=scalar_types.float4_e2m1f.id,
468
+ apply_router_weight_on_input=apply_router_weight_on_input,
469
+ global_num_experts=global_num_experts,
470
+ activation=activation,
471
+ expert_map=expert_map)
472
+
473
assert _can_support_mxfp4(
474
use_grouped_topk, topk_group, num_expert_group, expert_map,
475
custom_routing_function, e_score_correction_bias,
476
apply_router_weight_on_input, scoring_func, activation,
477
expert_load_view, logical_to_physical_map,
407
- logical_replica_count), ("MXFP4 are not supported\
408
- with this configuration.")
478
+ logical_replica_count), (
479
+ "MXFP4 are not supported with this configuration.")
480
481
if (envs.VLLM_USE_FLASHINFER_MOE_MXFP4_MXFP8
482
or envs.VLLM_USE_FLASHINFER_MOE_MXFP4_BF16):
vllm/model_executor/layers/quantization/utils/marlin_utils.py
+9
-6
@@ -261,6 +261,13 @@ def marlin_permute_scales(s: torch.Tensor, size_k: int, size_n: int,
261
return s
262
263
264
+def marlin_permute_bias(s: torch.Tensor) -> torch.Tensor:
265
+ origin_shape = s.shape
266
+ _, scale_perm_single = get_scale_perms()
267
+ s = s.reshape((-1, len(scale_perm_single)))[:, scale_perm_single]
268
+ return s.reshape(*origin_shape).contiguous()
269
+
270
+
271
def marlin_moe_permute_scales(
272
s: torch.Tensor,
273
size_k: int,
@@ -410,6 +417,7 @@ def apply_gptq_marlin_linear(
417
output = ops.gptq_marlin_gemm(reshaped_x,
418
None,
419
weight,
420
+ bias,
421
weight_scale,
422
None,
423
weight_zp,
@@ -425,9 +433,6 @@ def apply_gptq_marlin_linear(
433
use_fp32_reduce=use_fp32_reduce,
434
is_zp_float=False)
435
428
- if bias is not None:
429
- output.add_(bias) # In-place add
430
-
436
return output.reshape(out_shape)
437
438
@@ -456,6 +461,7 @@ def apply_awq_marlin_linear(
461
output = ops.gptq_marlin_gemm(reshaped_x,
462
None,
463
weight,
464
+ bias,
465
weight_scale,
466
None,
467
weight_zp,
@@ -470,7 +476,4 @@ def apply_awq_marlin_linear(
476
use_fp32_reduce=use_fp32_reduce,
477
is_zp_float=False)
478
473
- if bias is not None:
474
- output.add_(bias) # In-place add
475
-
479
return output.reshape(out_shape)
vllm/model_executor/layers/quantization/utils/marlin_utils_fp4.py
+141
-28
@@ -8,8 +8,8 @@ import torch
8
import vllm._custom_ops as ops
9
from vllm.logger import init_logger
10
from vllm.model_executor.layers.quantization.utils.marlin_utils import (
11
- USE_FP32_REDUCE_DEFAULT, marlin_make_workspace_new, marlin_permute_scales,
12
- should_use_atomic_add_reduce)
11
+ USE_FP32_REDUCE_DEFAULT, marlin_make_workspace_new, marlin_permute_bias,
12
+ marlin_permute_scales, should_use_atomic_add_reduce)
13
from vllm.platforms import current_platform
14
from vllm.scalar_type import scalar_types
15
@@ -22,7 +22,7 @@ def is_fp4_marlin_supported():
22
return current_platform.has_device_capability(80)
23
24
25
-def fp4_marlin_process_scales(marlin_scales):
25
+def nvfp4_marlin_process_scales(marlin_scales):
26
if not (marlin_scales >= 0).all():
27
logger.warning_once(
28
"NVFP4 Marlin assumes the scales to be >=0, but has encountered "
@@ -56,7 +56,20 @@ def fp4_marlin_process_scales(marlin_scales):
56
return marlin_scales
57
58
59
-def fp4_marlin_process_global_scale(global_scale):
59
+def mxfp4_marlin_process_scales(marlin_scales):
60
+ # 8 is the number of scale number using by one thread
61
+ marlin_scales = marlin_scales.view(marlin_scales.size(0) // 2, 2, -1, 8)
62
+ marlin_scales = marlin_scales.permute(0, 2, 1, 3).reshape(
63
+ marlin_scales.size(0) * 2, -1)
64
+
65
+ # fit the layout of fp8 dequantization
66
+ marlin_scales = marlin_scales.view(-1, 4)[:, [0, 2, 1, 3]].view(
67
+ marlin_scales.size(0), -1)
68
+ marlin_scales = marlin_scales.to(torch.float8_e8m0fnu)
69
+ return marlin_scales
70
+
71
+
72
+def nvfp4_marlin_process_global_scale(global_scale):
73
assert global_scale.dtype in [torch.half, torch.bfloat16]
74
fp4_exponent = 2
75
if global_scale.dtype == torch.half:
@@ -73,7 +86,7 @@ def apply_fp4_marlin_linear(
86
input: torch.Tensor,
87
weight: torch.Tensor,
88
weight_scale: torch.Tensor,
76
- weight_scale_2: torch.Tensor,
89
+ weight_scale_2: Optional[torch.Tensor],
90
workspace: torch.Tensor,
91
size_n: int,
92
size_k: int,
@@ -94,6 +107,7 @@ def apply_fp4_marlin_linear(
107
output = ops.gptq_marlin_gemm(a=reshaped_x,
108
c=None,
109
b_q_weight=weight,
110
+ b_bias=bias,
111
b_scales=weight_scale,
112
global_scale=weight_scale_2,
113
b_zeros=None,
@@ -107,9 +121,6 @@ def apply_fp4_marlin_linear(
121
use_atomic_add=use_atomic_add,
122
use_fp32_reduce=use_fp32_reduce)
123
110
- if bias is not None:
111
- output.add_(bias) # In-place add
112
-
124
return output.reshape(out_shape)
125
126
@@ -120,6 +131,9 @@ def prepare_fp4_layer_for_marlin(layer: torch.nn.Module) -> None:
131
"be used leveraging the Marlin kernel. This may degrade "
132
"performance for compute-heavy workloads.")
133
134
+ is_nvfp4 = hasattr(layer, "weight_scale_2")
135
+ group_size = 16 if is_nvfp4 else 32
136
+
137
part_size_n = layer.output_size_per_partition
138
part_size_k = layer.input_size_per_partition
139
param_dtype = layer.params_dtype
@@ -145,18 +159,35 @@ def prepare_fp4_layer_for_marlin(layer: torch.nn.Module) -> None:
159
160
# WEIGHT SCALES
161
# Permute scales
148
- weight_scale = layer.weight_scale.T.to(param_dtype)
162
+ weight_scale = layer.weight_scale.T.contiguous()
163
+
164
+ if not is_nvfp4:
165
+ weight_scale = weight_scale.view(torch.float8_e8m0fnu)
166
+
167
+ weight_scale = weight_scale.to(param_dtype)
168
weight_scale = marlin_permute_scales(s=weight_scale,
169
size_k=part_size_k,
170
size_n=part_size_n,
152
- group_size=16)
153
- weight_scale = fp4_marlin_process_scales(weight_scale)
154
- layer.weight_scale = torch.nn.Parameter(weight_scale, requires_grad=False)
155
-
156
- weight_scale_2 = layer.weight_scale_2.to(param_dtype)
157
- weight_scale_2 = fp4_marlin_process_global_scale(weight_scale_2)
158
- layer.weight_scale_2 = torch.nn.Parameter(weight_scale_2,
159
- requires_grad=False)
171
+ group_size=group_size)
172
+
173
+ if is_nvfp4:
174
+ weight_scale = nvfp4_marlin_process_scales(weight_scale)
175
+ layer.weight_scale = torch.nn.Parameter(weight_scale,
176
+ requires_grad=False)
177
+
178
+ weight_scale_2 = layer.weight_scale_2.to(param_dtype)
179
+ weight_scale_2 = nvfp4_marlin_process_global_scale(weight_scale_2)
180
+ layer.weight_scale_2 = torch.nn.Parameter(weight_scale_2,
181
+ requires_grad=False)
182
+ else:
183
+ weight_scale = mxfp4_marlin_process_scales(weight_scale)
184
+ layer.weight_scale = torch.nn.Parameter(weight_scale,
185
+ requires_grad=False)
186
+
187
+ if hasattr(layer, "bias") and layer.bias is not None:
188
+ assert layer.bias.shape == (part_size_n, )
189
+ bias = marlin_permute_bias(layer.bias)
190
+ layer.bias = torch.nn.Parameter(bias, requires_grad=False)
191
192
return
193
@@ -168,6 +199,9 @@ def prepare_moe_fp4_layer_for_marlin(layer: torch.nn.Module) -> None:
199
"be used leveraging the Marlin kernel. This may degrade "
200
"performance for compute-heavy workloads.")
201
202
+ is_nvfp4 = hasattr(layer, "w13_weight_scale_2")
203
+ group_size = 16 if is_nvfp4 else 32
204
+
205
e = layer.num_experts
206
k = layer.hidden_size
207
n = layer.intermediate_size_per_partition
@@ -208,8 +242,13 @@ def prepare_moe_fp4_layer_for_marlin(layer: torch.nn.Module) -> None:
242
# WEIGHT SCALES
243
# Permute scales
244
for name in ["w13", "w2"]:
211
- scales = getattr(layer, name + "_weight_scale").to(param_dtype)
212
- global_scale = getattr(layer, name + "_weight_scale_2").to(param_dtype)
245
+ scales = getattr(layer, name + "_weight_scale")
246
+ if not is_nvfp4:
247
+ scales = scales.view(torch.float8_e8m0fnu)
248
+ scales = scales.to(param_dtype)
249
+ if is_nvfp4:
250
+ global_scale = getattr(layer,
251
+ name + "_weight_scale_2").to(param_dtype)
252
253
tensor_list = []
254
if "w13" in name:
@@ -218,23 +257,47 @@ def prepare_moe_fp4_layer_for_marlin(layer: torch.nn.Module) -> None:
257
size_n, size_k = k, n
258
259
for i in range(e):
221
- marlin_scales = marlin_permute_scales(s=scales[i].T,
260
+ scale = scales[i].T
261
+
262
+ marlin_scales = marlin_permute_scales(s=scale,
263
size_k=size_k,
264
size_n=size_n,
224
- group_size=16)
225
- marlin_scales = fp4_marlin_process_scales(marlin_scales)
265
+ group_size=group_size)
266
+ if is_nvfp4:
267
+ marlin_scales = nvfp4_marlin_process_scales(marlin_scales)
268
+ else:
269
+ marlin_scales = mxfp4_marlin_process_scales(marlin_scales)
270
tensor_list.append(marlin_scales)
271
272
scales = torch.cat([x.unsqueeze(0) for x in tensor_list], 0)
273
scales = torch.nn.Parameter(scales, requires_grad=False)
274
setattr(layer, name + "_weight_scale", scales)
275
232
- global_scale = fp4_marlin_process_global_scale(global_scale)
233
- global_scale = torch.nn.Parameter(global_scale, requires_grad=False)
234
- setattr(layer, name + "_weight_scale_2", global_scale)
276
+ if is_nvfp4:
277
+ global_scale = nvfp4_marlin_process_global_scale(global_scale)
278
+ global_scale = torch.nn.Parameter(global_scale,
279
+ requires_grad=False)
280
+ setattr(layer, name + "_weight_scale_2", global_scale)
281
+
282
+ # BIAS
283
+ # Permute bias
284
+ for name in ["w13_bias", "w2_bias"]:
285
+ if not hasattr(layer, name):
286
+ continue
287
+ bias = getattr(layer, name).to(param_dtype)
288
+
289
+ tensor_list = []
290
+ for i in range(e):
291
+ expert_bias = bias[i]
292
+
293
+ tensor_list.append(marlin_permute_bias(expert_bias))
294
+
295
+ bias = torch.cat([x.unsqueeze(0) for x in tensor_list], 0)
296
+ bias = torch.nn.Parameter(bias, requires_grad=False)
297
+ setattr(layer, name, bias)
298
299
237
-def rand_marlin_weight_fp4_like(weight, group_size):
300
+def rand_marlin_weight_nvfp4_like(weight, group_size):
301
assert group_size > 0
302
size_n, size_k = weight.shape
303
device = weight.device
@@ -276,8 +339,58 @@ def rand_marlin_weight_fp4_like(weight, group_size):
339
size_k=size_k,
340
size_n=size_n,
341
group_size=group_size)
279
- marlin_scales = fp4_marlin_process_scales(marlin_scales)
342
+ marlin_scales = nvfp4_marlin_process_scales(marlin_scales)
343
281
- global_scale = fp4_marlin_process_global_scale(global_scale)
344
+ global_scale = nvfp4_marlin_process_global_scale(global_scale)
345
346
return weight_ref.T, marlin_qweight, marlin_scales, global_scale
347
+
348
+
349
+def rand_marlin_weight_mxfp4_like(weight, group_size):
350
+ assert group_size > 0
351
+ size_n, size_k = weight.shape
352
+ device = weight.device
353
+
354
+ scales = torch.randint(100,
355
+ 125, (size_n, size_k // group_size),
356
+ dtype=torch.uint8,
357
+ device=weight.device)
358
+ scales = scales.view(torch.float8_e8m0fnu)
359
+
360
+ fp4_weight = torch.randint(0,
361
+ 256, (size_n, size_k // 2),
362
+ dtype=torch.uint8,
363
+ device=weight.device)
364
+ fp4_weight_part_1 = ((fp4_weight & 0b10000000) |
365
+ ((fp4_weight & 0b01110000) >> 2))
366
+ fp4_weight_part_1 = fp4_weight_part_1.view(torch.float8_e4m3fn)
367
+ fp4_weight_part_1 = fp4_weight_part_1.to(weight.dtype) * (2**6)
368
+
369
+ fp4_weight2 = fp4_weight << 4
370
+ fp4_weight_part_2 = ((fp4_weight2 & 0b10000000) |
371
+ ((fp4_weight2 & 0b01110000) >> 2))
372
+ fp4_weight_part_2 = fp4_weight_part_2.view(torch.float8_e4m3fn)
373
+ fp4_weight_part_2 = fp4_weight_part_2.to(weight.dtype) * (2**6)
374
+
375
+ weight_ref = torch.cat(
376
+ [fp4_weight_part_2.unsqueeze(2),
377
+ fp4_weight_part_1.unsqueeze(2)], 2).view(size_n, size_k)
378
+ weight_ref = weight_ref * \
379
+ scales.repeat_interleave(group_size, 1).to(weight.dtype)
380
+
381
+ marlin_qweight = ops.gptq_marlin_repack(
382
+ b_q_weight=fp4_weight.view(torch.int32).T.contiguous(),
383
+ perm=torch.empty(0, dtype=torch.int, device=device),
384
+ size_k=size_k,
385
+ size_n=size_n,
386
+ num_bits=4,
387
+ )
388
+
389
+ marlin_scales = marlin_permute_scales(s=scales.T.to(weight.dtype),
390
+ size_k=size_k,
391
+ size_n=size_n,
392
+ group_size=group_size)
393
+
394
+ marlin_scales = mxfp4_marlin_process_scales(marlin_scales)
395
+
396
+ return weight_ref.T, marlin_qweight, marlin_scales.to(torch.float8_e8m0fnu)
vllm/model_executor/layers/quantization/utils/marlin_utils_fp8.py
+25
-5
@@ -8,8 +8,8 @@ import torch
8
import vllm._custom_ops as ops
9
from vllm.logger import init_logger
10
from vllm.model_executor.layers.quantization.utils.marlin_utils import (
11
- USE_FP32_REDUCE_DEFAULT, marlin_make_workspace_new, marlin_permute_scales,
12
- should_use_atomic_add_reduce)
11
+ USE_FP32_REDUCE_DEFAULT, marlin_make_workspace_new, marlin_permute_bias,
12
+ marlin_permute_scales, should_use_atomic_add_reduce)
13
from vllm.platforms import current_platform
14
from vllm.scalar_type import scalar_types
15
@@ -58,6 +58,7 @@ def apply_fp8_marlin_linear(
58
output = ops.gptq_marlin_gemm(a=reshaped_x,
59
c=None,
60
b_q_weight=weight,
61
+ b_bias=bias,
62
b_scales=weight_scale,
63
global_scale=None,
64
b_zeros=None,
@@ -71,9 +72,6 @@ def apply_fp8_marlin_linear(
72
use_atomic_add=use_atomic_add,
73
use_fp32_reduce=use_fp32_reduce)
74
74
- if bias is not None:
75
- output.add_(bias) # In-place add
76
-
75
return output.reshape(out_shape)
76
77
@@ -160,6 +158,11 @@ def prepare_fp8_layer_for_marlin(layer: torch.nn.Module,
158
marlin_scales = fp8_fused_exponent_bias_into_scales(marlin_scales)
159
layer.weight_scale = torch.nn.Parameter(marlin_scales, requires_grad=False)
160
161
+ if hasattr(layer, "bias") and layer.bias is not None:
162
+ assert layer.bias.shape == (part_size_n, )
163
+ bias = marlin_permute_bias(layer.bias)
164
+ layer.bias = torch.nn.Parameter(bias, requires_grad=False)
165
+
166
167
def prepare_moe_fp8_layer_for_marlin(layer: torch.nn.Module,
168
size_k_first: bool = True) -> None:
@@ -274,6 +277,23 @@ def prepare_moe_fp8_layer_for_marlin(layer: torch.nn.Module,
277
278
setattr(layer, name + "_weight_scale", scales)
279
280
+ # BIAS
281
+ # Permute bias
282
+ for name in ["w13_bias", "w2_bias"]:
283
+ if not hasattr(layer, name):
284
+ continue
285
+ bias = getattr(layer, name).to(layer.orig_dtype)
286
+
287
+ tensor_list = []
288
+ for i in range(e):
289
+ expert_bias = bias[i]
290
+
291
+ tensor_list.append(marlin_permute_bias(expert_bias))
292
+
293
+ bias = torch.cat([x.unsqueeze(0) for x in tensor_list], 0)
294
+ bias = torch.nn.Parameter(bias, requires_grad=False)
295
+ setattr(layer, name, bias)
296
+
297
298
def pack_fp8_to_int32(fp8_tensor: torch.Tensor,
299
size_k_first: bool = True) -> torch.Tensor:
vllm/model_executor/layers/quantization/utils/mxfp4_utils.py
+1
-1
@@ -61,7 +61,7 @@ def _can_support_mxfp4(use_grouped_topk: bool = False,
61
e_score_correction_bias: Optional[torch.Tensor] = None,
62
apply_router_weight_on_input: bool = False,
63
scoring_func: str = "softmax",
64
- activation: str = "silu",
64
+ activation: str = "swiglu_oai",
65
expert_load_view: Optional[torch.Tensor] = None,
66
logical_to_physical_map: Optional[torch.Tensor] = None,
67
logical_replica_count: Optional[torch.Tensor] = None):
vllm/scalar_type.py
+2
@@ -327,6 +327,8 @@ class scalar_types:
327
uint8 = ScalarType.uint(8, None)
328
float8_e4m3fn = ScalarType.float_(4, 3, True, NanRepr.EXTD_RANGE_MAX_MIN)
329
float8_e5m2 = ScalarType.float_IEEE754(5, 2)
330
+ float8_e8m0fnu = ScalarType(8, 0, False, 0, True,
331
+ NanRepr.EXTD_RANGE_MAX_MIN)
332
float16_e8m7 = ScalarType.float_IEEE754(8, 7)
333
float16_e5m10 = ScalarType.float_IEEE754(5, 10)
334