[Bugfix][CPU][RISC-V] Fix VLEN detection for RVV attention path (#47532)
Signed-off-by: liutong <liutong@iscas.ac.cn> Co-authored-by: Claude <noreply@anthropic.com>
This commit is contained in:
@@ -178,7 +178,10 @@ elseif (CMAKE_SYSTEM_PROCESSOR MATCHES "riscv64")
|
||||
# Override with -DVLLM_RVV_VLEN=128 or -DVLLM_RVV_VLEN=256 for RVV.
|
||||
if(NOT DEFINED VLLM_RVV_VLEN)
|
||||
# Auto-detect: find the largest zvl<N>b in /proc/cpuinfo isa line.
|
||||
if(EXISTS /proc/cpuinfo)
|
||||
# Skip when cross-compiling — /proc/cpuinfo describes the build host.
|
||||
if(CMAKE_CROSSCOMPILING)
|
||||
message(STATUS "Cross-compiling: skipping VLEN auto-detection from /proc/cpuinfo")
|
||||
elseif(EXISTS /proc/cpuinfo)
|
||||
file(READ /proc/cpuinfo _cpuinfo)
|
||||
set(_best 0)
|
||||
foreach(_n IN ITEMS 128 256 512 1024)
|
||||
@@ -186,6 +189,13 @@ elseif (CMAKE_SYSTEM_PROCESSOR MATCHES "riscv64")
|
||||
set(_best ${_n})
|
||||
endif()
|
||||
endforeach()
|
||||
# Only VLEN=128 and VLEN=256 are supported by the RVV kernels.
|
||||
if(_best GREATER 256)
|
||||
message(WARNING
|
||||
"Detected VLEN=${_best} but only 128/256 are supported; "
|
||||
"clamping to 256")
|
||||
set(_best 256)
|
||||
endif()
|
||||
if(_best GREATER 0)
|
||||
set(VLLM_RVV_VLEN ${_best})
|
||||
endif()
|
||||
|
||||
@@ -13,7 +13,8 @@ static inline cpu_attention::Fp8KVCacheDataType parse_fp8_kv_dtype(
|
||||
|
||||
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
|
||||
#if defined(__riscv) && defined(__riscv_v_min_vlen) && \
|
||||
(__riscv_v_min_vlen == 128 || __riscv_v_min_vlen == 256)
|
||||
return true;
|
||||
#else
|
||||
return false;
|
||||
|
||||
@@ -437,27 +437,24 @@ def _riscv_supports_rvv() -> bool:
|
||||
The RVV path is compiled whenever __riscv_v_min_vlen is defined, so
|
||||
we check that at least one supported zvl<N>b is advertised.
|
||||
"""
|
||||
# The C++ compile-time check is the ground truth: it knows which
|
||||
# VLEN the binary was actually compiled for. The cpuinfo check
|
||||
# below is only a fast-path shortcut.
|
||||
try:
|
||||
import torch
|
||||
|
||||
if torch.ops._C.cpu_attn_has_isa("rvv"):
|
||||
return True
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
# Fallback: check /proc/cpuinfo for zvl128b/zvl256b.
|
||||
try:
|
||||
with open("/proc/cpuinfo") as f:
|
||||
cpuinfo = f.read()
|
||||
except OSError:
|
||||
return False
|
||||
# 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 zvl<N>b 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
|
||||
return any(f"zvl{n}b" in cpuinfo for n in (128, 256))
|
||||
|
||||
|
||||
def _get_attn_isa(
|
||||
|
||||
Reference in New Issue
Block a user