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 4, 20:00 UTC
[Bugfix][CPU][RISC-V] Clamp exp() input to prevent NaN (#40428)
Signed-off-by: liuyudong <liuyudong@iscas.ac.cn>
lyd1992 committed
Apr 22, 2026 at 17:38 UTC
04eac6ba24d22e0e280dff53695544facfaf5ca0
1 file changed
+26
-12
csrc/cpu/cpu_types_riscv_impl.hpp
+26
-12
@@ -15,16 +15,12 @@
15
#include <torch/all.h>
16
namespace vec_op {
17
18
-#ifdef RISCV_BF16_SUPPORT
19
- #define VLLM_DISPATCH_CASE_FLOATING_TYPES(...) \
20
- AT_DISPATCH_CASE(at::ScalarType::Float, __VA_ARGS__) \
21
- AT_DISPATCH_CASE(at::ScalarType::Half, __VA_ARGS__) \
22
- AT_DISPATCH_CASE(at::ScalarType::BFloat16, __VA_ARGS__)
23
-#else
24
- #define VLLM_DISPATCH_CASE_FLOATING_TYPES(...) \
25
- AT_DISPATCH_CASE(at::ScalarType::Float, __VA_ARGS__) \
26
- AT_DISPATCH_CASE(at::ScalarType::Half, __VA_ARGS__)
27
-#endif
18
+// BFloat16 is always supported on RISC-V: natively when RISCV_BF16_SUPPORT
19
+// is defined, otherwise via the FP32-simulation fallback path.
20
+#define VLLM_DISPATCH_CASE_FLOATING_TYPES(...) \
21
+ AT_DISPATCH_CASE(at::ScalarType::Float, __VA_ARGS__) \
22
+ AT_DISPATCH_CASE(at::ScalarType::Half, __VA_ARGS__) \
23
+ AT_DISPATCH_CASE(at::ScalarType::BFloat16, __VA_ARGS__)
24
25
#define VLLM_DISPATCH_FLOATING_TYPES(TYPE, NAME, ...) \
26
AT_DISPATCH_SWITCH(TYPE, NAME, VLLM_DISPATCH_CASE_FLOATING_TYPES(__VA_ARGS__))
@@ -486,9 +482,18 @@ struct FP32Vec8 : public Vec<FP32Vec8> {
482
}
483
484
FP32Vec8 exp() const {
485
+ // Clamp input to prevent NaN: exp(-inf) must return 0, not NaN.
486
+ // Without clamping, -inf * 0.0 = NaN in the final poly * scale step.
487
+ // Matches the clamping strategy used by x86 AVX-512 and ARM NEON.
488
+ constexpr float exp_lo = -87.3365447505f; // ln(FLT_MIN)
489
+ constexpr float exp_hi = 88.7228391117f; // ln(FLT_MAX)
490
+ fixed_fp32x8_t x = RVVI(__riscv_vfmin_vf_f32, LMUL_256)(
491
+ RVVI(__riscv_vfmax_vf_f32, LMUL_256)(reg, exp_lo, VEC_ELEM_NUM), exp_hi,
492
+ VEC_ELEM_NUM);
493
+
494
const float inv_ln2 = 1.44269504088896341f;
495
fixed_fp32x8_t x_scaled =
491
- RVVI(__riscv_vfmul_vf_f32, LMUL_256)(reg, inv_ln2, VEC_ELEM_NUM);
496
+ RVVI(__riscv_vfmul_vf_f32, LMUL_256)(x, inv_ln2, VEC_ELEM_NUM);
497
fixed_i32x8_t n_int =
498
RVVI(__riscv_vfcvt_x_f_v_i32, LMUL_256)(x_scaled, VEC_ELEM_NUM);
499
fixed_fp32x8_t n_float =
@@ -706,9 +711,18 @@ struct FP32Vec16 : public Vec<FP32Vec16> {
711
}
712
713
FP32Vec16 exp() const {
714
+ // Clamp input to prevent NaN: exp(-inf) must return 0, not NaN.
715
+ // Without clamping, -inf * 0.0 = NaN in the final poly * scale step.
716
+ // Matches the clamping strategy used by x86 AVX-512 and ARM NEON.
717
+ constexpr float exp_lo = -87.3365447505f; // ln(FLT_MIN)
718
+ constexpr float exp_hi = 88.7228391117f; // ln(FLT_MAX)
719
+ fixed_fp32x16_t x = RVVI(__riscv_vfmin_vf_f32, LMUL_512)(
720
+ RVVI(__riscv_vfmax_vf_f32, LMUL_512)(reg, exp_lo, VEC_ELEM_NUM), exp_hi,
721
+ VEC_ELEM_NUM);
722
+
723
const float inv_ln2 = 1.44269504088896341f;
724
fixed_fp32x16_t x_scaled =
711
- RVVI(__riscv_vfmul_vf_f32, LMUL_512)(reg, inv_ln2, VEC_ELEM_NUM);
725
+ RVVI(__riscv_vfmul_vf_f32, LMUL_512)(x, inv_ln2, VEC_ELEM_NUM);
726
fixed_i32x16_t n_int =
727
RVVI(__riscv_vfcvt_x_f_v_i32, LMUL_512)(x_scaled, VEC_ELEM_NUM);
728
fixed_fp32x16_t n_float =