@tduncan / Nvidia-Optimization-Stack-VLLM / commits / fdbf2ddbd2

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

Signed-off-by: jiang1.li <jiang1.li@intel.com>

Li, Jiang committed Aug 31, 2026 at 15:33 UTC fdbf2ddbd22a82f8f59b97ba30e5436f0509d2c8

11 files changed +118 -85

csrc/cpu/cpu_fused_moe.cpp

+33 -26

@@ -136,60 +136,67 @@ void fused_moe_impl(scalar_t* __restrict__ output, scalar_t* __restrict__ input,

136 }();

137

138 // allocate buffers

139 - int32_t common_buffer_offset = 0;

140 - int32_t w13_thread_buffer_offset = 0;

141 - int32_t ws_thread_buffer_offset = 0;

139 + // Sizes/offsets below are byte counts across the whole (token_num *

140 + // topk_num) batch, which for large hidden sizes and batch sizes can

141 + // exceed INT32_MAX. Use int64_t throughout to avoid silently wrapping

142 + // into a negative size that gets sign-extended into a huge value when

143 + // passed to ScratchPadManager::realloc(size_t).

144 + int64_t common_buffer_offset = 0;

145 + int64_t w13_thread_buffer_offset = 0;

146 + int64_t ws_thread_buffer_offset = 0;

147

148 // common buffers

144 - const int32_t token_num_per_group_buffer_size =

149 + const int64_t token_num_per_group_buffer_size =

150 cpu_utils::round_up<64>(expert_num * sizeof(int32_t));

146 - const int32_t token_num_per_group_buffer_offset = common_buffer_offset;

151 + const int64_t token_num_per_group_buffer_offset = common_buffer_offset;

152 common_buffer_offset += token_num_per_group_buffer_size;

153

149 - const int32_t cu_token_num_per_group_buffer_size =

154 + const int64_t cu_token_num_per_group_buffer_size =

155 cpu_utils::round_up<64>((expert_num + 1) * sizeof(int32_t));

151 - const int32_t cu_token_num_per_group_buffer_offset = common_buffer_offset;

156 + const int64_t cu_token_num_per_group_buffer_offset = common_buffer_offset;

157 common_buffer_offset += cu_token_num_per_group_buffer_size;

158

154 - const int32_t expand_token_id_buffer_size =

155 - cpu_utils::round_up<64>(token_num * topk_num * sizeof(int32_t));

156 - const int32_t expand_token_id_buffer_offset = common_buffer_offset;

159 + const int64_t expand_token_id_buffer_size = cpu_utils::round_up<64>(

160 + static_cast<int64_t>(token_num) * topk_num * sizeof(int32_t));

161 + const int64_t expand_token_id_buffer_offset = common_buffer_offset;

162 common_buffer_offset += expand_token_id_buffer_size;

163

159 - const int32_t expand_token_id_index_buffer_size =

160 - cpu_utils::round_up<64>(token_num * topk_num * sizeof(int32_t));

161 - const int32_t expand_token_id_index_buffer_offset = common_buffer_offset;

164 + const int64_t expand_token_id_index_buffer_size = cpu_utils::round_up<64>(

165 + static_cast<int64_t>(token_num) * topk_num * sizeof(int32_t));

166 + const int64_t expand_token_id_index_buffer_offset = common_buffer_offset;

167 common_buffer_offset += expand_token_id_index_buffer_size;

168

164 - const int32_t w13_gemm_output_buffer_size = cpu_utils::round_up<64>(

165 - token_num * topk_num * (output_size_13 / 2) * sizeof(scalar_t));

166 - const int32_t w13_gemm_output_buffer_offset = common_buffer_offset;

169 + const int64_t w13_gemm_output_buffer_size =

170 + cpu_utils::round_up<64>(static_cast<int64_t>(token_num) * topk_num *

171 + (output_size_13 / 2) * sizeof(scalar_t));

172 + const int64_t w13_gemm_output_buffer_offset = common_buffer_offset;

173 common_buffer_offset += w13_gemm_output_buffer_size;

174

169 - const int32_t w2_gemm_output_buffer_size = cpu_utils::round_up<64>(

170 - token_num * topk_num * output_size_2 * sizeof(float));

171 - const int32_t w2_gemm_output_buffer_offset = common_buffer_offset;

175 + const int64_t w2_gemm_output_buffer_size =

176 + cpu_utils::round_up<64>(static_cast<int64_t>(token_num) * topk_num *

177 + output_size_2 * sizeof(float));

178 + const int64_t w2_gemm_output_buffer_offset = common_buffer_offset;

179 common_buffer_offset += w2_gemm_output_buffer_size;

180

181 // w13 GEMM thread buffers

175 - const int32_t w13_input_buffer_offset = w13_thread_buffer_offset;

182 + const int64_t w13_input_buffer_offset = w13_thread_buffer_offset;

183 w13_thread_buffer_offset += w13_input_buffer_size;

184

178 - const int32_t w13_output_buffer_size = cpu_utils::round_up<64>(

185 + const int64_t w13_output_buffer_size = cpu_utils::round_up<64>(

186 gemm_m_tile_size * w13_n_tile_size * sizeof(float));

180 - const int32_t w13_output_buffer_offset = w13_thread_buffer_offset;

187 + const int64_t w13_output_buffer_offset = w13_thread_buffer_offset;

188 w13_thread_buffer_offset += w13_output_buffer_size;

189

183 - const int32_t w2_input_buffer_offset = w13_thread_buffer_offset;

190 + const int64_t w2_input_buffer_offset = w13_thread_buffer_offset;

191 w13_thread_buffer_offset += w2_input_buffer_size;

192

193 // Weighted sum thread buffer

187 - const int32_t ws_output_buffer_size =

194 + const int64_t ws_output_buffer_size =

195 cpu_utils::round_up<64>(output_size_2 * sizeof(float));

189 - const int32_t ws_output_buffer_offset = ws_thread_buffer_offset;

196 + const int64_t ws_output_buffer_offset = ws_thread_buffer_offset;

197 ws_thread_buffer_offset += ws_output_buffer_size;

198

192 - const int32_t buffer_size =

199 + const int64_t buffer_size =

200 common_buffer_offset +

201 std::max(w13_thread_buffer_offset, ws_thread_buffer_offset) * thread_num;

202 cpu_utils::ScratchPadManager::get_scratchpad_manager()->realloc(buffer_size);

csrc/cpu/sgl-kernels/gemm.cpp

+4 -7

@@ -723,9 +723,10 @@ at::Tensor convert_scale_packed(at::Tensor& scale) {

723 // bias : [N]

724 // out : [*, N]

725 //

726 -at::Tensor

727 -weight_packed_linear(at::Tensor& mat1, at::Tensor& mat2, const std::optional<at::Tensor>& bias, bool is_vnni) {

728 - auto packed_w = is_vnni ? mat2 : convert_weight_packed(mat2);

726 +void

727 +weight_packed_linear(at::Tensor& out, const at::Tensor& mat1, const at::Tensor& mat2, const std::optional<at::Tensor>& bias, bool is_vnni) {

728 + TORCH_CHECK(is_vnni);

729 + auto packed_w = mat2;

730 bool use_fma_gemm = false;

731 if (packed_w.scalar_type() == at::kFloat) {

732 use_fma_gemm = true;

@@ -746,7 +747,6 @@ weight_packed_linear(at::Tensor& mat1, at::Tensor& mat2, const std::optional<at:

747 }

748

749 auto dispatch_type = mat1.scalar_type();

749 - auto out = at::empty({M, N}, mat1.options());

750 // strides

751 int64_t out_strideM = out.stride(0);

752 int64_t mat1_strideM = mat1.stride(-2);

@@ -784,9 +784,6 @@ weight_packed_linear(at::Tensor& mat1, at::Tensor& mat2, const std::optional<at:

784 out_strideM);

785 }

786 });

787 -

788 - input_sizes[ndim - 1] = N;

789 - return out.view(input_sizes);

787 }

788

789 // mat1 : [M, K]

csrc/cpu/torch_bindings.cpp

+6 -5

@@ -58,9 +58,9 @@ std::vector<torch::Tensor> shm_recv_tensor_list(int64_t handle, int64_t src);

59 // SGL CPU kernels

61 -at::Tensor weight_packed_linear(at::Tensor& mat1, at::Tensor& mat2,

62 - const std::optional<at::Tensor>& bias,

63 - bool is_vnni);

61 +void weight_packed_linear(at::Tensor& out, const at::Tensor& mat1,

62 + const at::Tensor& mat2,

63 + const std::optional<at::Tensor>& bias, bool is_vnni);

65 at::Tensor convert_weight_packed(at::Tensor& weight);

@@ -507,8 +507,9 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {

507 // sgl-kernels

508 #if defined(__AVX512BF16__) && defined(__AVX512F__) && defined(__AVX512VNNI__)

509 ops.def(

510 - "weight_packed_linear(Tensor(a0!) mat1, Tensor(a1!) mat2, Tensor(a2!)? "

511 - "bias, bool is_vnni) -> Tensor");

510 + "weight_packed_linear(Tensor(a0!) out, Tensor(a1) mat1, Tensor(a2) mat2, "

511 + "Tensor(a3)? "

512 + "bias, bool is_vnni) -> ()");

513 ops.impl("weight_packed_linear", torch::kCPU, &weight_packed_linear);

514 ops.def("convert_weight_packed(Tensor! weight) -> Tensor");

515 ops.impl("convert_weight_packed", torch::kCPU, &convert_weight_packed);

tests/conftest.py

+9 -1

@@ -934,6 +934,14 @@ def hf_runner():

934 return HfRunner

935

936

937 +def _default_block_size() -> int:

938 + if torch.xpu.is_available():

939 + return 64

940 + if current_platform.is_cpu():

941 + return 128

942 + return 16

943 +

944 +

945 class VllmRunner:

946 """

947 The default value of some arguments have been modified from

@@ -962,7 +970,7 @@ class VllmRunner:

970 dtype: str = "auto",

971 disable_log_stats: bool = True,

972 tensor_parallel_size: int = 1,

965 - block_size: int = 16 if not torch.xpu.is_available() else 64,

973 + block_size: int = _default_block_size(),

974 enable_chunked_prefill: bool | None = False,

975 enforce_eager: bool | None = False,

976 # Set this to avoid hanging issue

tests/v1/attention/test_group_head_counts.py

+3 -2

@@ -44,9 +44,10 @@ def _build(layer_num_heads: list[int]) -> CPUAttentionMetadataBuilder:

44 vllm_config = MagicMock()

45 vllm_config.model_config.dtype = torch.bfloat16

46 vllm_config.model_config.get_num_attention_heads.return_value = MODEL_WIDE_NUM_HEADS

47 - vllm_config.cache_config.block_size = 16

47 vllm_config.cache_config.cache_dtype = "auto"

49 - kv_cache_spec = SimpleNamespace(num_kv_heads=NUM_KV_HEADS, head_size=64)

48 + kv_cache_spec = SimpleNamespace(

49 + num_kv_heads=NUM_KV_HEADS, head_size=64, block_size=16

50 + )

52 with (

53 patch(

tests/v1/kv_connector/unit/test_bidirectional_kv_transfer.py

+4 -2

@@ -87,8 +87,10 @@ def _make_connector_with_fake_worker(

87 hand_shake_latency=0, cycles_before_done=0, do_handshake=True

88 ):

89 """Create a NixlConnector with FakeNixlConnectorWorker."""

90 - vllm_config = create_vllm_config(kv_connector_extra_config=BIDIR_KV_EXTRA_CONFIG)

91 - kv_cache_config = make_kv_cache_config(block_size=16, num_blocks=2)

90 + vllm_config = create_vllm_config(

91 + kv_connector_extra_config=BIDIR_KV_EXTRA_CONFIG, block_size=32

92 + )

93 + kv_cache_config = make_kv_cache_config(block_size=32, num_blocks=2)

94 connector = NixlConnector(vllm_config, KVConnectorRole.WORKER, kv_cache_config)

95 connector.connector_worker = FakeNixlConnectorWorker(

96 vllm_config,

vllm/_custom_ops.py

+11 -12

@@ -3328,18 +3328,17 @@ def dsv3_fused_a_gemm(

3328 torch.ops._C.dsv3_fused_a_gemm(output, mat_a, mat_b, enable_pdl)

3329

3330

3331 -if hasattr(torch.ops._C, "weight_packed_linear"):

3332 -

3333 - @register_fake("_C::weight_packed_linear")

3334 - def weight_packed_linear_fake(

3335 - mat1: torch.Tensor,

3336 - mat2: torch.Tensor,

3337 - bias: torch.Tensor | None,

3338 - is_vnni: bool,

3339 - ) -> torch.Tensor:

3340 - return torch.empty(

3341 - (mat1.size(0), mat2.size(0)), dtype=mat1.dtype, device=mat2.device

3342 - )

3331 +def weight_packed_linear_cpu(

3332 + x: torch.Tensor,

3333 + packed_weight: torch.Tensor,

3334 + n: int,

3335 + bias: torch.Tensor | None,

3336 +) -> torch.Tensor:

3337 + output = torch.empty((*x.shape[0:-1], n), dtype=x.dtype)

3338 + torch.ops._C.weight_packed_linear(

3339 + output.reshape(-1, n), x.reshape(-1, x.size(-1)), packed_weight, bias, True

3340 + )

3341 + return output

3342

3343

3344 class CPUQuantMethod(IntEnum):

vllm/model_executor/layers/mamba/gdn/qwen_gdn_linear_attn.py

+13 -1

@@ -150,6 +150,15 @@ def _log_gdn_backend_decision(

150 head_k_dim = getattr(

151 vllm_config.model_config.hf_text_config, "linear_key_head_dim", None

152 )

153 +

154 + if current_platform.is_cpu():

155 + logger.info_once(

156 + "Using %s GDN prefill kernel (head_k_dim=%s).",

157 + "CPU",

158 + head_k_dim,

159 + )

160 + return

161 +

162 chosen = {

163 "flashinfer": "FlashInfer",

164 "cutedsl": "CuteDSL",

@@ -492,7 +501,7 @@ class QwenGatedDeltaNetAttention(GatedDeltaNetAttention):

501 envs.VLLM_ENABLE_FLA_PACKED_RECURRENT_DECODE

502 )

503 self.gdn_decode_kernel = envs.VLLM_GDN_DECODE_KERNEL.strip().lower()

495 - if self.gdn_decode_kernel == "cuda":

504 + if self.gdn_decode_kernel == "cuda" and current_platform.is_cuda_alike():

505 reason = self._fused_gdn_decode_unsupported_reason(vllm_config)

506 if reason is not None:

507 if "VLLM_GDN_DECODE_KERNEL" in os.environ:

@@ -503,6 +512,9 @@ class QwenGatedDeltaNetAttention(GatedDeltaNetAttention):

512 "Falling back to the Triton GDN decode path: %s", reason

513 )

514 self.gdn_decode_kernel = "triton"

515 + elif current_platform.is_cpu():

516 + self.gdn_decode_kernel = "CPU"

517 +

518 self.enable_fused_gdn_decode = self.gdn_decode_kernel == "cuda"

519 logger.info_once("GDN decode kernel: %s", self.gdn_decode_kernel)

520

vllm/model_executor/layers/utils.py

+7 -5

@@ -505,13 +505,15 @@ def dispatch_cpu_unquantized_gemm(

505 # cross over to favoring oneDNN once batch size grows past decode-sized

506 # M, so they keep using oneDNN below.

507 if check_cpu_sgl_kernel(N, K, dtype):

508 + # For small size GEMM, packed_weight might be float32

509 packed_weight = torch.ops._C.convert_weight_packed(layer.weight)

510 if getattr(layer, "bias", None) is not None:

510 - bias_f32 = layer.bias.to(torch.float32)

511 - else:

512 - bias_f32 = None

513 - layer.cpu_linear = lambda x, weight, bias: torch.ops._C.weight_packed_linear(

514 - x, packed_weight, bias_f32 if bias is not None else None, True

511 + layer.bias.data = layer.bias.to(torch.float32)

512 + layer.cpu_linear = lambda x, weight, bias: ops.weight_packed_linear_cpu(

513 + x,

514 + packed_weight,

515 + N,

516 + bias,

517 )

518 if remove_weight:

519 layer.weight = torch.nn.Parameter(torch.empty(0), requires_grad=False)

vllm/triton_utils/importing.py

+14 -11

@@ -34,6 +34,20 @@ if HAS_TRITON:

34 if name != "cpu" and x.driver and x.driver.is_active()

35 ]

37 + # Check Triton CPU

38 + if "cpu" in version("vllm"):

39 + if "cpu" in backends:

40 + HAS_TRITON = True

41 + # Suppress following warnings on CPU-only platforms

42 + if len(active_drivers) == 0:

43 + active_drivers.append(backends["cpu"].driver)

44 + else:

45 + logger.warning(

46 + "Triton is installed, but doesn't include CPU backend. "

47 + "Disabling Triton."

48 + )

49 + HAS_TRITON = False

50 +

51 # Check if we're in a distributed environment where CUDA_VISIBLE_DEVICES

52 # or HIP_VISIBLE_DEVICES might be temporarily empty (e.g., Ray sets it to ""

53 # during actor init)

@@ -63,17 +77,6 @@ if HAS_TRITON:

77 len(active_drivers),

78 )

79 HAS_TRITON = False

66 -

67 - # Check Triton CPU

68 - if "cpu" in version("vllm"):

69 - if "cpu" in backends:

70 - HAS_TRITON = True

71 - else:

72 - logger.warning(

73 - "Triton is installed, but doesn't include CPU backend. "

74 - "Disabling Triton."

75 - )

76 - HAS_TRITON = False

80 except ImportError:

81 # This can occur if Triton is partially installed or triton.backends

82 # is missing.

vllm/v1/attention/backends/cpu_attn.py

+14 -13

@@ -13,7 +13,6 @@ from vllm import _custom_ops as ops

13 from vllm import envs

14 from vllm.config import (

15 VllmConfig,

16 - get_current_vllm_config,

16 get_layers_from_vllm_config,

17 )

18 from vllm.logger import init_logger

@@ -59,7 +58,7 @@ class CPUAttentionBackend(AttentionBackend):

59 @staticmethod

60 def get_supported_kernel_block_sizes() -> list[int | MultipleOf]:

62 - return [MultipleOf(16)]

61 + return [MultipleOf(32)]

63 @classmethod

64 def get_supported_head_sizes(cls) -> list[int]:

@@ -152,7 +151,7 @@ class CPUAttentionMetadataBuilder(AttentionMetadataBuilder[CPUAttentionMetadata]

151 self.head_dim = kv_cache_spec.head_size

152 self.dtype = vllm_config.model_config.dtype

153 self.window_size = self._group_sliding_window()

155 - self.block_size = vllm_config.cache_config.block_size

154 + self.block_size = kv_cache_spec.block_size

155 self.kv_cache_dtype = vllm_config.cache_config.cache_dtype

156 self.isa = _get_attn_isa(

157 self.dtype,

@@ -160,11 +159,21 @@ class CPUAttentionMetadataBuilder(AttentionMetadataBuilder[CPUAttentionMetadata]

159 self.head_dim,

160 self.kv_cache_dtype,

161 )

162 + self._set_isa_to_layers(layer_names)

163 self.is_cross_attention = isinstance(kv_cache_spec, CrossAttentionSpec)

164 self.is_encoder_only_attention = isinstance(

165 kv_cache_spec, EncoderOnlyAttentionSpec

166 )

167

168 + def _set_isa_to_layers(self, layer_names: list[str]) -> None:

169 + attn_layers = get_layers_from_vllm_config(

170 + self.vllm_config,

171 + Attention,

172 + layer_names,

173 + )

174 + for layer in attn_layers.values():

175 + layer.isa = self.isa # type: ignore

176 +

177 def _group_sliding_window(self) -> int:

178 """The window shared by every layer in this group, else -1 (no window).

179

@@ -325,14 +334,6 @@ class CPUAttentionBackendImpl(AttentionImpl):

334 "heads in the layer"

335 )

336

328 - vllm_config = get_current_vllm_config()

329 - self.isa = _get_attn_isa(

330 - vllm_config.model_config.dtype,

331 - vllm_config.cache_config.block_size,

332 - self.head_size,

333 - self.kv_cache_dtype,

334 - )

335 -

337 def forward(

338 self,

339 layer: AttentionLayer,

@@ -391,7 +392,7 @@ class CPUAttentionBackendImpl(AttentionImpl):

392 key_cache,

393 value_cache,

394 attn_metadata.slot_mapping,

394 - self.isa,

395 + layer.isa, # type: ignore

396 k_scale=layer._k_scale_float,

397 v_scale=layer._v_scale_float,

398 kv_cache_dtype=self.kv_cache_dtype,

@@ -440,7 +441,7 @@ class CPUAttentionBackendImpl(AttentionImpl):

441 key_cache,

442 value_cache,

443 slot_mapping,

443 - self.isa,

444 + layer.isa, # type: ignore

445 k_scale=layer._k_scale_float,

446 v_scale=layer._v_scale_float,

447 kv_cache_dtype=self.kv_cache_dtype,