diff --git a/csrc/cpu/cpu_fused_moe.cpp b/csrc/cpu/cpu_fused_moe.cpp index 6b64c018f8..416f20352e 100644 --- a/csrc/cpu/cpu_fused_moe.cpp +++ b/csrc/cpu/cpu_fused_moe.cpp @@ -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(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(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(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(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); diff --git a/csrc/cpu/sgl-kernels/gemm.cpp b/csrc/cpu/sgl-kernels/gemm.cpp index 212e1db529..e227923d6c 100644 --- a/csrc/cpu/sgl-kernels/gemm.cpp +++ b/csrc/cpu/sgl-kernels/gemm.cpp @@ -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& 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& 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 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& bias, - bool is_vnni); +void weight_packed_linear(at::Tensor& out, const at::Tensor& mat1, + const at::Tensor& mat2, + const std::optional& 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); diff --git a/tests/conftest.py b/tests/conftest.py index f9156a43e9..5fd594a04a 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -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 diff --git a/tests/v1/attention/test_group_head_counts.py b/tests/v1/attention/test_group_head_counts.py index 9691926819..54fe3cfdbc 100644 --- a/tests/v1/attention/test_group_head_counts.py +++ b/tests/v1/attention/test_group_head_counts.py @@ -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( diff --git a/tests/v1/kv_connector/unit/test_bidirectional_kv_transfer.py b/tests/v1/kv_connector/unit/test_bidirectional_kv_transfer.py index 0801e3817e..78a6d930f9 100644 --- a/tests/v1/kv_connector/unit/test_bidirectional_kv_transfer.py +++ b/tests/v1/kv_connector/unit/test_bidirectional_kv_transfer.py @@ -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, diff --git a/vllm/_custom_ops.py b/vllm/_custom_ops.py index dba31ef3e2..badadc162a 100644 --- a/vllm/_custom_ops.py +++ b/vllm/_custom_ops.py @@ -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): diff --git a/vllm/model_executor/layers/mamba/gdn/qwen_gdn_linear_attn.py b/vllm/model_executor/layers/mamba/gdn/qwen_gdn_linear_attn.py index e733cd9a78..108a73223e 100644 --- a/vllm/model_executor/layers/mamba/gdn/qwen_gdn_linear_attn.py +++ b/vllm/model_executor/layers/mamba/gdn/qwen_gdn_linear_attn.py @@ -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) diff --git a/vllm/model_executor/layers/utils.py b/vllm/model_executor/layers/utils.py index 08ce3b28d7..caa52f2851 100644 --- a/vllm/model_executor/layers/utils.py +++ b/vllm/model_executor/layers/utils.py @@ -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) diff --git a/vllm/triton_utils/importing.py b/vllm/triton_utils/importing.py index 2828ed467e..966ab55f32 100644 --- a/vllm/triton_utils/importing.py +++ b/vllm/triton_utils/importing.py @@ -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. diff --git a/vllm/v1/attention/backends/cpu_attn.py b/vllm/v1/attention/backends/cpu_attn.py index 7f06ec1392..dd3f34c916 100644 --- a/vllm/v1/attention/backends/cpu_attn.py +++ b/vllm/v1/attention/backends/cpu_attn.py @@ -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,