A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 5, 21:00 UTC
[Bugfix]fix output Nan/Inf in marlin if dtype=float16 (#33972)
Signed-off-by: IriKa Qiu <qiujie.jq@gmail.com>
IriKa committed
Mar 28, 2026 at 07:36 UTC
148a5c1226f8668cb52c4900b5ff2c80344e78f2
8 files changed
+83
-55
csrc/moe/marlin_moe_wna16/kernel.h
+1
-1
@@ -13,7 +13,7 @@
13
const int4 *__restrict__ b_bias_ptr, \
14
const float *__restrict__ a_scales_ptr, \
15
const int4 *__restrict__ scales_ptr, \
16
- const uint16_t *__restrict__ global_scale_ptr, \
16
+ const float *__restrict__ global_scale_ptr, \
17
const int4 *__restrict__ zp_ptr, const int *__restrict__ g_idx, \
18
const int32_t *__restrict__ sorted_token_ids_ptr, \
19
const int32_t *__restrict__ expert_ids_ptr, \
csrc/moe/marlin_moe_wna16/marlin_template.h
+22
-13
@@ -260,7 +260,7 @@ __global__ void Marlin(
260
// fp16 quantization scales. shape (k/groupsize, n)
261
const int4* __restrict__ scales_ptr,
262
// fp16 global scale (for nvfp4// only)
263
- const uint16_t* __restrict__ global_scale_ptr,
263
+ const float* __restrict__ global_scale_ptr,
264
// 4bit packed zero-points of shape
265
// (k/groupsize, n/pack_factor)
266
const int4* __restrict__ zp_ptr,
@@ -308,7 +308,14 @@ __global__ void Marlin(
308
constexpr int moe_block_size = m_block_size_8 ? 8 : (16 * thread_m_blocks);
309
310
#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ == 750
311
- constexpr bool use_fp16_accum = a_type_id == vllm::kFloat16.id();
311
+ static constexpr auto num_bits =
312
+ vllm::ScalarType::from_id(b_type_id).size_bits();
313
+ // Disable use_fp16_accum for NVFP4 and cases when group_size == -1 &&
314
+ // num_bits == 4
315
+ constexpr bool use_fp16_accum =
316
+ a_type_id == vllm::kFloat16.id() &&
317
+ (!(b_type_id == vllm::kFE2M1f.id() && s_type_id == vllm::kFE4M3fn.id()) &&
318
+ !(group_blocks == -1 && num_bits == 4));
319
#else
320
constexpr bool use_fp16_accum = false;
321
#endif
@@ -357,7 +364,7 @@ __global__ void Marlin(
364
has_zp && !is_zp_float && !std::is_same<scalar_t, nv_bfloat16>::value ||
365
has_zp && !is_zp_float && !(b_type == vllm::kU8);
366
360
- c_scalar_t2 global_scale;
367
+ float global_scale_f32 = 1.0f;
368
369
constexpr bool has_act_order = group_blocks == 0;
370
@@ -507,11 +514,12 @@ __global__ void Marlin(
514
515
if (mul_topk_weights) {
516
idx = idx < prob_m_top_k ? idx : 0;
510
- c_scalar_t2 topk_weight_val =
511
- Cdtype::num2num2(Cdtype::float2num(topk_weights_ptr[idx]));
517
+ float topk_weight_tmp = topk_weights_ptr[idx];
518
if constexpr (b_type == vllm::kFE2M1f && s_type == vllm::kFE4M3fn) {
513
- topk_weight_val = __hmul2(topk_weight_val, global_scale);
519
+ topk_weight_tmp *= global_scale_f32;
520
}
521
+ c_scalar_t2 topk_weight_val =
522
+ Cdtype::num2num2(Cdtype::float2num(topk_weight_tmp));
523
sh_block_topk_weights[threadIdx.x] = topk_weight_val;
524
}
525
}
@@ -532,8 +540,7 @@ __global__ void Marlin(
540
expert_id = expert_ids_ptr[block_id];
541
542
if constexpr (b_type == vllm::kFE2M1f && s_type == vllm::kFE4M3fn) {
535
- uint16_t val = global_scale_ptr[expert_id];
536
- global_scale = Cdtype::num2num2(*reinterpret_cast<c_scalar_t*>(&val));
543
+ global_scale_f32 = global_scale_ptr[expert_id];
544
}
545
546
B_expert_off = expert_id * prob_n * prob_k / (pack_factor * 4);
@@ -1784,6 +1791,13 @@ __global__ void Marlin(
1791
// We first reorder in shared memory to guarantee the most efficient final
1792
// global write patterns
1793
auto write = [&](int idx, float c0, float c1, FragS& s, FragS& b_bias) {
1794
+ if constexpr (b_type == vllm::kFE2M1f && s_type == vllm::kFE4M3fn) {
1795
+ if (!mul_topk_weights) {
1796
+ c0 *= global_scale_f32;
1797
+ c1 *= global_scale_f32;
1798
+ }
1799
+ }
1800
+
1801
c_scalar_t2 res =
1802
Cdtype::nums2num2(Cdtype::float2num(c0), Cdtype::float2num(c1));
1803
@@ -1800,11 +1814,6 @@ __global__ void Marlin(
1814
res = __hmul2(res, tmp_scale);
1815
}
1816
1803
- if constexpr (b_type == vllm::kFE2M1f && s_type == vllm::kFE4M3fn) {
1804
- if (!mul_topk_weights) {
1805
- res = __hmul2(res, global_scale);
1806
- }
1807
- }
1817
if (has_bias && last) {
1818
c_scalar_t2 tmp_bias = b_bias[0];
1819
if constexpr (m_block_size_8) {
csrc/moe/marlin_moe_wna16/ops.cu
+4
-4
@@ -382,7 +382,7 @@ void marlin_mm(const void* A, const void* B, void* C, void* C_tmp, void* b_bias,
382
const int4* bias_ptr = (const int4*)b_bias;
383
const float* a_s_ptr = (const float*)a_s;
384
const int4* b_s_ptr = (const int4*)b_s;
385
- const uint16_t* g_s_ptr = (const uint16_t*)g_s;
385
+ const float* g_s_ptr = (const float*)g_s;
386
const int4* zp_ptr = (const int4*)zp;
387
const int* g_idx_ptr = (const int*)g_idx;
388
const int* perm_ptr = (const int*)perm;
@@ -759,7 +759,7 @@ torch::Tensor moe_wna16_marlin_gemm(
759
TORCH_CHECK(b_type == vllm::kFE2M1f && s_type == vllm::kFE4M3fn,
760
"global_scale can only be used for nvfp4 format.");
761
} else {
762
- global_scale = torch::empty({0}, options);
762
+ global_scale = torch::empty({0}, options_fp32);
763
TORCH_CHECK(!(b_type == vllm::kFE2M1f && s_type == vllm::kFE4M3fn),
764
"the global_scale parameter must be passed for nvfp4 format.");
765
}
@@ -842,8 +842,8 @@ torch::Tensor moe_wna16_marlin_gemm(
842
843
TORCH_CHECK(a_scales.scalar_type() == at::ScalarType::Float,
844
"scalar type of a_scales must be float");
845
- TORCH_CHECK(global_scale.scalar_type() == c.scalar_type(),
846
- "scalar type of global_scale must be the same with c");
845
+ TORCH_CHECK(global_scale.scalar_type() == at::ScalarType::Float,
846
+ "scalar type of global_scale must be float");
847
if (a_type.size_bits() == 16) {
848
TORCH_CHECK(
849
a.scalar_type() == c.scalar_type(),
csrc/quantization/activation_kernels.cu
-3
@@ -189,10 +189,7 @@ __device__ __forceinline__ void cp_async_wait<0>() {
189
}
190
191
__device__ __forceinline__ float clip(float v, float mmin, float mmax) {
192
-#if __CUDACC_VER_MAJOR__ >= 11 && __CUDA_ARCH__ >= 800
192
return fminf(mmax, fmaxf(v, mmin));
194
-#else
195
-#endif
193
}
194
195
__device__ __forceinline__ __nv_bfloat16 clip(__nv_bfloat16 v,
csrc/quantization/marlin/kernel.h
+1
-1
@@ -13,7 +13,7 @@
13
const int4 *__restrict__ b_bias_ptr, \
14
const float *__restrict__ a_scales_ptr, \
15
const int4 *__restrict__ scales_ptr, \
16
- const uint16_t *__restrict__ global_scale_ptr, \
16
+ const float *__restrict__ global_scale_ptr, \
17
const int4 *__restrict__ zp_ptr, const int *__restrict__ g_idx, \
18
int num_groups, int prob_m, int prob_n, int prob_k, int lda, int *locks, \
19
bool has_bias, bool use_atomic_add, bool use_fp32_reduce, \
csrc/quantization/marlin/marlin.cu
+5
-5
@@ -57,7 +57,7 @@ torch::Tensor marlin_gemm(
57
int64_t size_k, bool is_k_full, bool use_atomic_add, bool use_fp32_reduce,
58
bool is_zp_float) {
59
TORCH_CHECK_NOT_IMPLEMENTED(false,
60
- "marlin_gemm(..) requires CUDA_ARCH >= 8.0");
60
+ "marlin_gemm(..) requires CUDA_ARCH >= 7.5");
61
return torch::empty({1, 1});
62
}
63
@@ -356,7 +356,7 @@ void marlin_mm(const void* A, const void* B, void* C, void* C_tmp, void* b_bias,
356
const int4* bias_ptr = (const int4*)b_bias;
357
const float* a_s_ptr = (const float*)a_s;
358
const int4* b_s_ptr = (const int4*)b_s;
359
- const uint16_t* g_s_ptr = (const uint16_t*)g_s;
359
+ const float* g_s_ptr = (const float*)g_s;
360
361
const int4* zp_ptr = (const int4*)zp;
362
const int* g_idx_ptr = (const int*)g_idx;
@@ -751,7 +751,7 @@ torch::Tensor marlin_gemm(
751
TORCH_CHECK(b_type == vllm::kFE2M1f && s_type == vllm::kFE4M3fn,
752
"global_scale can only be used for nvfp4 format.");
753
} else {
754
- global_scale = torch::empty({0}, options);
754
+ global_scale = torch::empty({0}, options_fp32);
755
TORCH_CHECK(!(b_type == vllm::kFE2M1f && s_type == vllm::kFE4M3fn),
756
"the global_scale parameter must be passed for nvfp4 format.");
757
}
@@ -832,8 +832,8 @@ torch::Tensor marlin_gemm(
832
833
TORCH_CHECK(a_scales.scalar_type() == at::ScalarType::Float,
834
"scalar type of a_scales must be float");
835
- TORCH_CHECK(global_scale.scalar_type() == c.scalar_type(),
836
- "scalar type of global_scale must be the same with c");
835
+ TORCH_CHECK(global_scale.scalar_type() == at::ScalarType::Float,
836
+ "scalar type of global_scale must be float");
837
if (a_type.size_bits() == 16) {
838
TORCH_CHECK(
839
a.scalar_type() == c.scalar_type(),
csrc/quantization/marlin/marlin_template.h
+15
-10
@@ -251,8 +251,8 @@ __global__ void Marlin(
251
const float* __restrict__ a_scales_ptr,
252
// fp16 quantization scales. shape (k/groupsize, n)
253
const int4* __restrict__ scales_ptr,
254
- // fp16 global scale (for nvfp4// only)
255
- const uint16_t* __restrict__ global_scale_ptr,
254
+ // float global scale (for nvfp4// only)
255
+ const float* __restrict__ global_scale_ptr,
256
// 4bit packed zero-points of shape
257
// (k/groupsize, n/pack_factor)
258
const int4* __restrict__ zp_ptr,
@@ -292,7 +292,13 @@ __global__ void Marlin(
292
#endif
293
294
#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ == 750
295
- constexpr bool use_fp16_accum = a_type_id == vllm::kFloat16.id();
295
+ constexpr auto num_bits = vllm::ScalarType::from_id(b_type_id).size_bits();
296
+ // Disable use_fp16_accum for NVFP4 and cases when group_size == -1 &&
297
+ // num_bits == 4
298
+ constexpr bool use_fp16_accum =
299
+ a_type_id == vllm::kFloat16.id() &&
300
+ (!(b_type_id == vllm::kFE2M1f.id() && s_type_id == vllm::kFE4M3fn.id()) &&
301
+ !(group_blocks == -1 && num_bits == 4));
302
#else
303
constexpr bool use_fp16_accum = false;
304
#endif
@@ -342,11 +348,10 @@ __global__ void Marlin(
348
has_zp && !is_zp_float && !std::is_same<scalar_t, nv_bfloat16>::value ||
349
has_zp && !is_zp_float && !(b_type == vllm::kU8);
350
345
- c_scalar_t2 global_scale;
351
+ float global_scale_f32 = 1.0f;
352
353
if constexpr (b_type == vllm::kFE2M1f && s_type == vllm::kFE4M3fn) {
348
- uint16_t val = global_scale_ptr[0];
349
- global_scale = Cdtype::num2num2(*reinterpret_cast<c_scalar_t*>(&val));
354
+ global_scale_f32 = global_scale_ptr[0];
355
}
356
357
constexpr bool has_act_order = group_blocks == 0;
@@ -1644,6 +1649,10 @@ __global__ void Marlin(
1649
// We first reorder in shared memory to guarantee the most efficient final
1650
// global write patterns
1651
auto write = [&](int idx, float c0, float c1, FragS& s, FragS& b_bias) {
1652
+ if constexpr (b_type == vllm::kFE2M1f && s_type == vllm::kFE4M3fn) {
1653
+ c0 *= global_scale_f32;
1654
+ c1 *= global_scale_f32;
1655
+ }
1656
c_scalar_t2 res =
1657
Cdtype::nums2num2(Cdtype::float2num(c0), Cdtype::float2num(c1));
1658
@@ -1659,10 +1668,6 @@ __global__ void Marlin(
1668
}
1669
res = __hmul2(res, tmp_scale);
1670
}
1662
-
1663
- if constexpr (b_type == vllm::kFE2M1f && s_type == vllm::kFE4M3fn) {
1664
- res = __hmul2(res, global_scale);
1665
- }
1671
if (has_bias && last) {
1672
c_scalar_t2 tmp_bias = b_bias[0];
1673
if constexpr (m_block_size_8) {
vllm/model_executor/layers/quantization/utils/marlin_utils_fp4.py
+35
-18
@@ -27,10 +27,19 @@ def is_fp4_marlin_supported():
27
return current_platform.has_device_capability(75)
28
29
30
-def _nvfp4_compute_scale_factor(marlin_scales: torch.Tensor) -> float:
30
+def _nvfp4_compute_scale_factor(
31
+ marlin_scales: torch.Tensor,
32
+ a_dtype: torch.dtype | None = None,
33
+) -> float:
34
"""Compute the power-of-2 scale_factor needed so that all non-zero
35
values in marlin_scales * 2^7 are >= 2 after rescaling.
36
Returns a Python float (power of 2, >= 1.0)."""
37
+
38
+ # Since half has a smaller dynamic range compared to bfloat16,
39
+ # no rescaling is applied here if active dtype is half.
40
+ if a_dtype is not None and a_dtype == torch.half:
41
+ return 1.0
42
+
43
ws_float = marlin_scales.float() * (2**7)
44
nonzero_mask = ws_float > 0
45
if nonzero_mask.any():
@@ -44,6 +53,7 @@ def _nvfp4_compute_scale_factor(marlin_scales: torch.Tensor) -> float:
53
def nvfp4_marlin_process_scales(
54
marlin_scales: torch.Tensor,
55
scale_factor: float | None = None,
56
+ a_dtype: torch.dtype | None = None,
57
) -> tuple[torch.Tensor, float]:
58
"""Process NVFP4 weight scales into the special S0E5M3 format for Marlin.
59
@@ -91,7 +101,7 @@ def nvfp4_marlin_process_scales(
101
# to fully utilize the E4M3 dynamic range (e.g., global_scale=1).
102
# The caller must compensate by dividing global_scale by scale_factor.
103
if scale_factor is None:
94
- scale_factor = _nvfp4_compute_scale_factor(marlin_scales)
104
+ scale_factor = _nvfp4_compute_scale_factor(marlin_scales, a_dtype)
105
if scale_factor > 1.0:
106
marlin_scales = (marlin_scales.float() * scale_factor).to(torch.half)
107
@@ -119,12 +129,14 @@ def mxfp4_marlin_process_scales(marlin_scales, input_dtype=None):
129
return marlin_scales
130
131
122
-def nvfp4_marlin_process_global_scale(global_scale):
123
- assert global_scale.dtype in [torch.half, torch.bfloat16]
132
+def nvfp4_marlin_process_global_scale(global_scale, a_dtype: torch.dtype | None = None):
133
+ if a_dtype is None:
134
+ a_dtype = global_scale.dtype
135
+ assert a_dtype in [torch.half, torch.bfloat16]
136
fp4_exponent = 2
125
- if global_scale.dtype == torch.half:
137
+ if a_dtype == torch.half:
138
target_exponent = 5
127
- elif global_scale.dtype == torch.bfloat16:
139
+ elif a_dtype == torch.bfloat16:
140
target_exponent = 8
141
# exponent_bias_fp16 = 2 ** 4 - 2 ** 1 = 14
142
# exponent_bias_bf16 = 2 ** 7 - 2 ** 1 = 126
@@ -244,11 +256,15 @@ def prepare_fp4_layer_for_marlin(
256
)
257
258
if is_nvfp4:
247
- weight_scale, scale_factor = nvfp4_marlin_process_scales(weight_scale)
259
+ weight_scale, scale_factor = nvfp4_marlin_process_scales(
260
+ weight_scale, a_dtype=param_dtype
261
+ )
262
layer.weight_scale = torch.nn.Parameter(weight_scale, requires_grad=False)
263
250
- weight_global_scale = layer.weight_global_scale.to(param_dtype)
251
- weight_global_scale = nvfp4_marlin_process_global_scale(weight_global_scale)
264
+ weight_global_scale = layer.weight_global_scale.to(torch.float32)
265
+ weight_global_scale = nvfp4_marlin_process_global_scale(
266
+ weight_global_scale, param_dtype
267
+ )
268
weight_global_scale = weight_global_scale / scale_factor
269
layer.weight_global_scale = torch.nn.Parameter(
270
weight_global_scale, requires_grad=False
@@ -339,7 +355,6 @@ def prepare_nvfp4_moe_layer_for_marlin(
355
scales: torch.Tensor, g_scales: torch.Tensor, name: str
356
) -> tuple[torch.Tensor, torch.Tensor]:
357
scales = scales.to(param_dtype)
342
- g_scales = g_scales.to(param_dtype)
358
359
tensor_list = []
360
num_shards = 2 if is_act_and_mul else 1
@@ -350,7 +365,7 @@ def prepare_nvfp4_moe_layer_for_marlin(
365
366
# All experts share one global_scale, so compute the max
367
# scale_factor across all experts first, then apply uniformly.
353
- combined_scale_factor = _nvfp4_compute_scale_factor(scales)
368
+ combined_scale_factor = _nvfp4_compute_scale_factor(scales, param_dtype)
369
370
for i in range(E):
371
scale = scales[i].T
@@ -362,12 +377,12 @@ def prepare_nvfp4_moe_layer_for_marlin(
377
is_a_8bit=is_a_8bit,
378
)
379
marlin_scales, _ = nvfp4_marlin_process_scales(
365
- marlin_scales, scale_factor=combined_scale_factor
380
+ marlin_scales, scale_factor=combined_scale_factor, a_dtype=param_dtype
381
)
382
tensor_list.append(marlin_scales)
383
384
scales = torch.cat([x.unsqueeze(0) for x in tensor_list], 0)
370
- g_scales = nvfp4_marlin_process_global_scale(g_scales)
385
+ g_scales = nvfp4_marlin_process_global_scale(g_scales, param_dtype)
386
g_scales = g_scales / combined_scale_factor
387
return scales, g_scales
388
@@ -438,7 +453,7 @@ def prepare_moe_fp4_layer_for_marlin(
453
scales = scales.view(torch.float8_e8m0fnu)
454
scales = scales.to(param_dtype)
455
if is_nvfp4:
441
- global_scale = getattr(layer, name + "_weight_scale_2").to(param_dtype)
456
+ global_scale = getattr(layer, name + "_weight_scale_2")
457
458
tensor_list = []
459
if "w13" in name:
@@ -449,7 +464,7 @@ def prepare_moe_fp4_layer_for_marlin(
464
# For NVFP4: compute unified scale_factor across all experts
465
combined_scale_factor = None
466
if is_nvfp4:
452
- combined_scale_factor = _nvfp4_compute_scale_factor(scales)
467
+ combined_scale_factor = _nvfp4_compute_scale_factor(scales, param_dtype)
468
469
for i in range(e):
470
scale = scales[i].T
@@ -463,7 +478,9 @@ def prepare_moe_fp4_layer_for_marlin(
478
)
479
if is_nvfp4:
480
marlin_scales, _ = nvfp4_marlin_process_scales(
466
- marlin_scales, scale_factor=combined_scale_factor
481
+ marlin_scales,
482
+ scale_factor=combined_scale_factor,
483
+ a_dtype=param_dtype,
484
)
485
else:
486
marlin_scales = mxfp4_marlin_process_scales(
@@ -477,7 +494,7 @@ def prepare_moe_fp4_layer_for_marlin(
494
495
if is_nvfp4:
496
assert combined_scale_factor is not None
480
- global_scale = nvfp4_marlin_process_global_scale(global_scale)
497
+ global_scale = nvfp4_marlin_process_global_scale(global_scale, param_dtype)
498
global_scale = global_scale / combined_scale_factor
499
global_scale = torch.nn.Parameter(global_scale, requires_grad=False)
500
setattr(layer, name + "_weight_scale_2", global_scale)
@@ -665,7 +682,7 @@ def rand_marlin_weight_nvfp4_like(weight, group_size, input_dtype=None):
682
)
683
marlin_scales, scale_factor = nvfp4_marlin_process_scales(marlin_scales)
684
668
- global_scale = nvfp4_marlin_process_global_scale(global_scale)
685
+ global_scale = nvfp4_marlin_process_global_scale(global_scale).to(torch.float32)
686
global_scale = global_scale / scale_factor
687
688
return weight_ref.T, marlin_qweight, marlin_scales, global_scale