iHore/vLLM
8
0
Fork 0
Code Issues Pull Requests Actions 1 Packages Projects Releases Wiki Activity

[Bugfix][CPU] Fix several bugs (#54042)

Browse Source 
Signed-off-by: jiang1.li <jiang1.li@intel.com>
This commit is contained in:
Li, Jiang
2026-08-31 15:33:06 +08:00
committed by GitHub
parent 2cf82bcdd1
commit fdbf2ddbd2
11 changed files with 118 additions and 85 deletions
Download Patch File Download Diff File Expand all files Collapse all files
csrc/cpu/cpu_fused_moe.cpp
+33 -26
View File
@@ -136,60 +136,67 @@ void fused_moe_impl(scalar_t* __restrict__ output, scalar_t* __restrict__ input,
}();
// allocate buffers
int32_t common_buffer_offset = 0;
int32_t w13_thread_buffer_offset = 0;
int32_t ws_thread_buffer_offset = 0;
// Sizes/offsets below are byte counts across the whole (token_num *
// topk_num) batch, which for large hidden sizes and batch sizes can
// exceed INT32_MAX. Use int64_t throughout to avoid silently wrapping
// into a negative size that gets sign-extended into a huge value when
// passed to ScratchPadManager::realloc(size_t).
int64_t common_buffer_offset = 0;
int64_t w13_thread_buffer_offset = 0;
int64_t ws_thread_buffer_offset = 0;
// common buffers
const int32_t token_num_per_group_buffer_size =
const int64_t token_num_per_group_buffer_size =
cpu_utils::round_up<64>(expert_num * sizeof(int32_t));
const int32_t token_num_per_group_buffer_offset = common_buffer_offset;
const int64_t token_num_per_group_buffer_offset = common_buffer_offset;
common_buffer_offset += token_num_per_group_buffer_size;
const int32_t cu_token_num_per_group_buffer_size =
const int64_t cu_token_num_per_group_buffer_size =
cpu_utils::round_up<64>((expert_num + 1) * sizeof(int32_t));
const int32_t cu_token_num_per_group_buffer_offset = common_buffer_offset;
const int64_t cu_token_num_per_group_buffer_offset = common_buffer_offset;
common_buffer_offset += cu_token_num_per_group_buffer_size;
const int32_t expand_token_id_buffer_size =
cpu_utils::round_up<64>(token_num * topk_num * sizeof(int32_t));
const int32_t expand_token_id_buffer_offset = common_buffer_offset;
const int64_t expand_token_id_buffer_size = cpu_utils::round_up<64>(
static_cast<int64_t>(token_num) * topk_num * sizeof(int32_t));
const int64_t expand_token_id_buffer_offset = common_buffer_offset;
common_buffer_offset += expand_token_id_buffer_size;
const int32_t expand_token_id_index_buffer_size =
cpu_utils::round_up<64>(token_num * topk_num * sizeof(int32_t));
const int32_t expand_token_id_index_buffer_offset = common_buffer_offset;
const int64_t expand_token_id_index_buffer_size = cpu_utils::round_up<64>(
static_cast<int64_t>(token_num) * topk_num * sizeof(int32_t));
const int64_t expand_token_id_index_buffer_offset = common_buffer_offset;
common_buffer_offset += expand_token_id_index_buffer_size;
const int32_t w13_gemm_output_buffer_size = cpu_utils::round_up<64>(
token_num * topk_num * (output_size_13 / 2) * sizeof(scalar_t));
const int32_t w13_gemm_output_buffer_offset = common_buffer_offset;
const int64_t w13_gemm_output_buffer_size =
cpu_utils::round_up<64>(static_cast<int64_t>(token_num) * topk_num *
(output_size_13 / 2) * sizeof(scalar_t));
const int64_t w13_gemm_output_buffer_offset = common_buffer_offset;
common_buffer_offset += w13_gemm_output_buffer_size;
const int32_t w2_gemm_output_buffer_size = cpu_utils::round_up<64>(
token_num * topk_num * output_size_2 * sizeof(float));
const int32_t w2_gemm_output_buffer_offset = common_buffer_offset;
const int64_t w2_gemm_output_buffer_size =
cpu_utils::round_up<64>(static_cast<int64_t>(token_num) * topk_num *
output_size_2 * sizeof(float));
const int64_t w2_gemm_output_buffer_offset = common_buffer_offset;
common_buffer_offset += w2_gemm_output_buffer_size;
// w13 GEMM thread buffers
const int32_t w13_input_buffer_offset = w13_thread_buffer_offset;
const int64_t w13_input_buffer_offset = w13_thread_buffer_offset;
w13_thread_buffer_offset += w13_input_buffer_size;
const int32_t w13_output_buffer_size = cpu_utils::round_up<64>(
const int64_t w13_output_buffer_size = cpu_utils::round_up<64>(
gemm_m_tile_size * w13_n_tile_size * sizeof(float));
const int32_t w13_output_buffer_offset = w13_thread_buffer_offset;
const int64_t w13_output_buffer_offset = w13_thread_buffer_offset;
w13_thread_buffer_offset += w13_output_buffer_size;
const int32_t w2_input_buffer_offset = w13_thread_buffer_offset;
const int64_t w2_input_buffer_offset = w13_thread_buffer_offset;
w13_thread_buffer_offset += w2_input_buffer_size;
// Weighted sum thread buffer
const int32_t ws_output_buffer_size =
const int64_t ws_output_buffer_size =
cpu_utils::round_up<64>(output_size_2 * sizeof(float));
const int32_t ws_output_buffer_offset = ws_thread_buffer_offset;
const int64_t ws_output_buffer_offset = ws_thread_buffer_offset;
ws_thread_buffer_offset += ws_output_buffer_size;
const int32_t buffer_size =
const int64_t buffer_size =
common_buffer_offset +
std::max(w13_thread_buffer_offset, ws_thread_buffer_offset) * thread_num;
cpu_utils::ScratchPadManager::get_scratchpad_manager()->realloc(buffer_size);
csrc/cpu/sgl-kernels/gemm.cpp
+4 -7
View File
@@ -723,9 +723,10 @@ at::Tensor convert_scale_packed(at::Tensor& scale) {
// bias : [N]
// out : [*, N]
//
at::Tensor
weight_packed_linear(at::Tensor& mat1, at::Tensor& mat2, const std::optional<at::Tensor>& bias, bool is_vnni) {
auto packed_w = is_vnni ? mat2 : convert_weight_packed(mat2);
void
weight_packed_linear(at::Tensor& out, const at::Tensor& mat1, const at::Tensor& mat2, const std::optional<at::Tensor>& bias, bool is_vnni) {
TORCH_CHECK(is_vnni);
auto packed_w = mat2;
bool use_fma_gemm = false;
if (packed_w.scalar_type() == at::kFloat) {
use_fma_gemm = true;
@@ -746,7 +747,6 @@ weight_packed_linear(at::Tensor& mat1, at::Tensor& mat2, const std::optional<at:
}
auto dispatch_type = mat1.scalar_type();
auto out = at::empty({M, N}, mat1.options());
// strides
int64_t out_strideM = out.stride(0);
int64_t mat1_strideM = mat1.stride(-2);
@@ -784,9 +784,6 @@ weight_packed_linear(at::Tensor& mat1, at::Tensor& mat2, const std::optional<at:
out_strideM);
}
});
input_sizes[ndim - 1] = N;
return out.view(input_sizes);
}
// mat1 : [M, K]
csrc/cpu/torch_bindings.cpp
+6 -5
View File
@@ -58,9 +58,9 @@ std::vector<torch::Tensor> shm_recv_tensor_list(int64_t handle, int64_t src);
// SGL CPU kernels
at::Tensor weight_packed_linear(at::Tensor& mat1, at::Tensor& mat2,
const std::optional<at::Tensor>& bias,
bool is_vnni);
void weight_packed_linear(at::Tensor& out, const at::Tensor& mat1,
const at::Tensor& mat2,
const std::optional<at::Tensor>& bias, bool is_vnni);
at::Tensor convert_weight_packed(at::Tensor& weight);
@@ -507,8 +507,9 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
// sgl-kernels
#if defined(__AVX512BF16__) && defined(__AVX512F__) && defined(__AVX512VNNI__)
ops.def(
"weight_packed_linear(Tensor(a0!) mat1, Tensor(a1!) mat2, Tensor(a2!)? "
"bias, bool is_vnni) -> Tensor");
"weight_packed_linear(Tensor(a0!) out, Tensor(a1) mat1, Tensor(a2) mat2, "
"Tensor(a3)? "
"bias, bool is_vnni) -> ()");
ops.impl("weight_packed_linear", torch::kCPU, &weight_packed_linear);
ops.def("convert_weight_packed(Tensor! weight) -> Tensor");
ops.impl("convert_weight_packed", torch::kCPU, &convert_weight_packed);
tests/conftest.py
+9 -1
View File
@@ -934,6 +934,14 @@ def hf_runner():
return HfRunner
def _default_block_size() -> int:
if torch.xpu.is_available():
return 64
if current_platform.is_cpu():
return 128
return 16
class VllmRunner:
"""
The default value of some arguments have been modified from
@@ -962,7 +970,7 @@ class VllmRunner:
dtype: str = "auto",
disable_log_stats: bool = True,
tensor_parallel_size: int = 1,
block_size: int = 16 if not torch.xpu.is_available() else 64,
block_size: int = _default_block_size(),
enable_chunked_prefill: bool | None = False,
enforce_eager: bool | None = False,
# Set this to avoid hanging issue
tests/v1/attention/test_group_head_counts.py
+3 -2
View File
@@ -44,9 +44,10 @@ def _build(layer_num_heads: list[int]) -> CPUAttentionMetadataBuilder:
vllm_config = MagicMock()
vllm_config.model_config.dtype = torch.bfloat16
vllm_config.model_config.get_num_attention_heads.return_value = MODEL_WIDE_NUM_HEADS
vllm_config.cache_config.block_size = 16
vllm_config.cache_config.cache_dtype = "auto"
kv_cache_spec = SimpleNamespace(num_kv_heads=NUM_KV_HEADS, head_size=64)
kv_cache_spec = SimpleNamespace(
num_kv_heads=NUM_KV_HEADS, head_size=64, block_size=16
)
with (
patch(
tests/v1/kv_connector/unit/test_bidirectional_kv_transfer.py
+4 -2
View File
@@ -87,8 +87,10 @@ def _make_connector_with_fake_worker(
hand_shake_latency=0, cycles_before_done=0, do_handshake=True
):
"""Create a NixlConnector with FakeNixlConnectorWorker."""
vllm_config = create_vllm_config(kv_connector_extra_config=BIDIR_KV_EXTRA_CONFIG)
kv_cache_config = make_kv_cache_config(block_size=16, num_blocks=2)
vllm_config = create_vllm_config(
kv_connector_extra_config=BIDIR_KV_EXTRA_CONFIG, block_size=32
)
kv_cache_config = make_kv_cache_config(block_size=32, num_blocks=2)
connector = NixlConnector(vllm_config, KVConnectorRole.WORKER, kv_cache_config)
connector.connector_worker = FakeNixlConnectorWorker(
vllm_config,
vllm/_custom_ops.py
+11 -12
View File
@@ -3328,18 +3328,17 @@ def dsv3_fused_a_gemm(
torch.ops._C.dsv3_fused_a_gemm(output, mat_a, mat_b, enable_pdl)
if hasattr(torch.ops._C, "weight_packed_linear"):
@register_fake("_C::weight_packed_linear")
def weight_packed_linear_fake(
mat1: torch.Tensor,
mat2: torch.Tensor,
bias: torch.Tensor | None,
is_vnni: bool,
) -> torch.Tensor:
return torch.empty(
(mat1.size(0), mat2.size(0)), dtype=mat1.dtype, device=mat2.device
)
def weight_packed_linear_cpu(
x: torch.Tensor,
packed_weight: torch.Tensor,
n: int,
bias: torch.Tensor | None,
) -> torch.Tensor:
output = torch.empty((*x.shape[0:-1], n), dtype=x.dtype)
torch.ops._C.weight_packed_linear(
output.reshape(-1, n), x.reshape(-1, x.size(-1)), packed_weight, bias, True
)
return output
class CPUQuantMethod(IntEnum):
vllm/model_executor/layers/mamba/gdn/qwen_gdn_linear_attn.py
+13 -1
View File
@@ -150,6 +150,15 @@ def _log_gdn_backend_decision(
head_k_dim = getattr(
vllm_config.model_config.hf_text_config, "linear_key_head_dim", None
)
if current_platform.is_cpu():
logger.info_once(
"Using %s GDN prefill kernel (head_k_dim=%s).",
"CPU",
head_k_dim,
)
return
chosen = {
"flashinfer": "FlashInfer",
"cutedsl": "CuteDSL",
@@ -492,7 +501,7 @@ class QwenGatedDeltaNetAttention(GatedDeltaNetAttention):
envs.VLLM_ENABLE_FLA_PACKED_RECURRENT_DECODE
)
self.gdn_decode_kernel = envs.VLLM_GDN_DECODE_KERNEL.strip().lower()
if self.gdn_decode_kernel == "cuda":
if self.gdn_decode_kernel == "cuda" and current_platform.is_cuda_alike():
reason = self._fused_gdn_decode_unsupported_reason(vllm_config)
if reason is not None:
if "VLLM_GDN_DECODE_KERNEL" in os.environ:
@@ -503,6 +512,9 @@ class QwenGatedDeltaNetAttention(GatedDeltaNetAttention):
"Falling back to the Triton GDN decode path: %s", reason
)
self.gdn_decode_kernel = "triton"
elif current_platform.is_cpu():
self.gdn_decode_kernel = "CPU"
self.enable_fused_gdn_decode = self.gdn_decode_kernel == "cuda"
logger.info_once("GDN decode kernel: %s", self.gdn_decode_kernel)
vllm/model_executor/layers/utils.py
+7 -5
View File
@@ -505,13 +505,15 @@ def dispatch_cpu_unquantized_gemm(
# cross over to favoring oneDNN once batch size grows past decode-sized
# M, so they keep using oneDNN below.
if check_cpu_sgl_kernel(N, K, dtype):
# For small size GEMM, packed_weight might be float32
packed_weight = torch.ops._C.convert_weight_packed(layer.weight)
if getattr(layer, "bias", None) is not None:
bias_f32 = layer.bias.to(torch.float32)
else:
bias_f32 = None
layer.cpu_linear = lambda x, weight, bias: torch.ops._C.weight_packed_linear(
x, packed_weight, bias_f32 if bias is not None else None, True
layer.bias.data = layer.bias.to(torch.float32)
layer.cpu_linear = lambda x, weight, bias: ops.weight_packed_linear_cpu(
x,
packed_weight,
N,
bias,
)
if remove_weight:
layer.weight = torch.nn.Parameter(torch.empty(0), requires_grad=False)
vllm/triton_utils/importing.py
+14 -11
View File
@@ -34,6 +34,20 @@ if HAS_TRITON:
if name != "cpu" and x.driver and x.driver.is_active()
]
# Check Triton CPU
if "cpu" in version("vllm"):
if "cpu" in backends:
HAS_TRITON = True
# Suppress following warnings on CPU-only platforms
if len(active_drivers) == 0:
active_drivers.append(backends["cpu"].driver)
else:
logger.warning(
"Triton is installed, but doesn't include CPU backend. "
"Disabling Triton."
)
HAS_TRITON = False
# Check if we're in a distributed environment where CUDA_VISIBLE_DEVICES
# or HIP_VISIBLE_DEVICES might be temporarily empty (e.g., Ray sets it to ""
# during actor init)
@@ -63,17 +77,6 @@ if HAS_TRITON:
len(active_drivers),
)
HAS_TRITON = False
# Check Triton CPU
if "cpu" in version("vllm"):
if "cpu" in backends:
HAS_TRITON = True
else:
logger.warning(
"Triton is installed, but doesn't include CPU backend. "
"Disabling Triton."
)
HAS_TRITON = False
except ImportError:
# This can occur if Triton is partially installed or triton.backends
# is missing.
vllm/v1/attention/backends/cpu_attn.py
+14 -13
View File
@@ -13,7 +13,6 @@ from vllm import _custom_ops as ops
from vllm import envs
from vllm.config import (
VllmConfig,
get_current_vllm_config,
get_layers_from_vllm_config,
)
from vllm.logger import init_logger
@@ -59,7 +58,7 @@ class CPUAttentionBackend(AttentionBackend):
@staticmethod
def get_supported_kernel_block_sizes() -> list[int | MultipleOf]:
return [MultipleOf(16)]
return [MultipleOf(32)]
@classmethod
def get_supported_head_sizes(cls) -> list[int]:
@@ -152,7 +151,7 @@ class CPUAttentionMetadataBuilder(AttentionMetadataBuilder[CPUAttentionMetadata]
self.head_dim = kv_cache_spec.head_size
self.dtype = vllm_config.model_config.dtype
self.window_size = self._group_sliding_window()
self.block_size = vllm_config.cache_config.block_size
self.block_size = kv_cache_spec.block_size
self.kv_cache_dtype = vllm_config.cache_config.cache_dtype
self.isa = _get_attn_isa(
self.dtype,
@@ -160,11 +159,21 @@ class CPUAttentionMetadataBuilder(AttentionMetadataBuilder[CPUAttentionMetadata]
self.head_dim,
self.kv_cache_dtype,
)
self._set_isa_to_layers(layer_names)
self.is_cross_attention = isinstance(kv_cache_spec, CrossAttentionSpec)
self.is_encoder_only_attention = isinstance(
kv_cache_spec, EncoderOnlyAttentionSpec
)
def _set_isa_to_layers(self, layer_names: list[str]) -> None:
attn_layers = get_layers_from_vllm_config(
self.vllm_config,
Attention,
layer_names,
)
for layer in attn_layers.values():
layer.isa = self.isa # type: ignore
def _group_sliding_window(self) -> int:
"""The window shared by every layer in this group, else -1 (no window).
@@ -325,14 +334,6 @@ class CPUAttentionBackendImpl(AttentionImpl):
"heads in the layer"
)
vllm_config = get_current_vllm_config()
self.isa = _get_attn_isa(
vllm_config.model_config.dtype,
vllm_config.cache_config.block_size,
self.head_size,
self.kv_cache_dtype,
)
def forward(
self,
layer: AttentionLayer,
@@ -391,7 +392,7 @@ class CPUAttentionBackendImpl(AttentionImpl):
key_cache,
value_cache,
attn_metadata.slot_mapping,
self.isa,
layer.isa, # type: ignore
k_scale=layer._k_scale_float,
v_scale=layer._v_scale_float,
kv_cache_dtype=self.kv_cache_dtype,
@@ -440,7 +441,7 @@ class CPUAttentionBackendImpl(AttentionImpl):
key_cache,
value_cache,
slot_mapping,
self.isa,
layer.isa, # type: ignore
k_scale=layer._k_scale_float,
v_scale=layer._v_scale_float,
kv_cache_dtype=self.kv_cache_dtype,