diff --git a/csrc/cpu/cpu_attn.cpp b/csrc/cpu/cpu_attn.cpp index 2634e649a71..ec1a2b162de 100644 --- a/csrc/cpu/cpu_attn.cpp +++ b/csrc/cpu/cpu_attn.cpp @@ -11,6 +11,17 @@ static inline cpu_attention::Fp8KVCacheDataType parse_fp8_kv_dtype( return cpu_attention::Fp8KVCacheDataType::kAuto; } +bool cpu_attn_has_isa(const std::string& isa) { + if (isa == "rvv") { +#if defined(__riscv) && defined(__riscv_v_min_vlen) && __riscv_v_min_vlen == 128 + return true; +#else + return false; +#endif + } + return false; +} + torch::Tensor get_scheduler_metadata( const int64_t num_req, const int64_t num_heads_q, const int64_t num_heads_kv, const int64_t head_dim, diff --git a/csrc/cpu/torch_bindings.cpp b/csrc/cpu/torch_bindings.cpp index 2aad5e2387d..0204f266b82 100644 --- a/csrc/cpu/torch_bindings.cpp +++ b/csrc/cpu/torch_bindings.cpp @@ -146,6 +146,8 @@ at::Tensor causal_conv1d_update_cpu( void activation_lut_bf16(torch::Tensor& out, torch::Tensor& input, const std::string& activation); +bool cpu_attn_has_isa(const std::string& isa); + torch::Tensor get_scheduler_metadata( const int64_t num_req, const int64_t num_heads_q, const int64_t num_heads_kv, const int64_t head_dim, @@ -497,6 +499,7 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) { ops.impl("fused_gdn_gating_cpu", torch::kCPU, &fused_gdn_gating_cpu); // CPU attention kernels + ops.def("cpu_attn_has_isa(str isa) -> bool", &cpu_attn_has_isa); ops.def( "get_scheduler_metadata(int num_req, int num_heads_q, int num_heads_kv, " "int head_dim, Tensor seq_lens, ScalarType dtype, Tensor " diff --git a/vllm/v1/attention/backends/cpu_attn.py b/vllm/v1/attention/backends/cpu_attn.py index e0670769adb..b2e186ac3b7 100644 --- a/vllm/v1/attention/backends/cpu_attn.py +++ b/vllm/v1/attention/backends/cpu_attn.py @@ -438,9 +438,22 @@ def _riscv_supports_rvv() -> bool: cpuinfo = f.read() except OSError: return False - return any(f"zvl{n}b" in cpuinfo for n in (128, 256)) and all( - f"zvl{n}b" not in cpuinfo for n in (512, 1024) - ) + # If VLEN >= 512 is detected, the RVV kernel was not compiled. + if any(f"zvl{n}b" in cpuinfo for n in (512, 1024)): + return False + + # zvl128b or zvl256b explicitly advertised -> RVV kernel available. + if any(f"zvl{n}b" in cpuinfo for n in (128, 256)): + return True + + # No zvlb flag at all (e.g. some hardware reports zve* without + # a VLEN hint). Delegate to the C++ compile-time check instead. + try: + import torch + + return torch.ops._C.cpu_attn_has_isa("rvv") + except Exception: + return False def _get_attn_isa(