@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / b4601ad43f

Fork 0

A high-throughput and memory-efficient inference and serving engine for LLMs

Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git

Synced Nov 3, 11:45 UTC

[CPU] Add fused GDN support for AMX CPU platform (#42707)

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

Li, Jiang committed May 18, 2026 at 18:04 UTC b4601ad43ff7ff2b9e2f52379144481e45bcf6c5

8 files changed +407 -14

csrc/cpu/sgl-kernels/conv.cpp

+21 -10

@@ -213,7 +213,7 @@ struct tinygemm_kernel<at::BFloat16, K, BLOCK_N, has_bias, has_silu> {

213 weight + nb_start * width, \

214 out + bs * seqlen * dim + mb_start * dim + nb_start, \

215 has_bias ? bias + nb_start : nullptr, \

216 - has_conv_states ? conv_states + conv_state_index * (K - 1) * dim + nb_start : nullptr, \

216 + has_conv_states ? conv_states + conv_state_index * conv_state_slot_stride + nb_start : nullptr, \

217 has_initial_states_value, \

218 mb_size, \

219 dim, \

@@ -233,7 +233,8 @@ void causal_conv1d_fwd_kernel_impl(

233 int64_t dim,

234 int64_t seqlen,

235 int64_t width,

236 - int64_t num_seq_blocks) {

236 + int64_t num_seq_blocks,

237 + int64_t conv_state_slot_stride) {

238 // handle 32 x 64 per block

239 constexpr int64_t BLOCK_M = block_size_m();

240 constexpr int64_t BLOCK_N = block_size_n() * 2;

@@ -282,7 +283,7 @@ void causal_conv1d_fwd_kernel_impl(

283 at::parallel_for(0, batch, 0, [&](int64_t begin, int64_t end) {

284 for (int64_t bs = begin; bs < end; ++bs) {

285 update_conv_state(

285 - conv_states + bs * (width - 1) * dim, input + bs * seqlen * dim, width, dim, seqlen, has_initial_state[bs]);

286 + conv_states + bs * conv_state_slot_stride, input + bs * seqlen * dim, width, dim, seqlen, has_initial_state[bs]);

287 }

288 });

289 }

@@ -316,7 +317,8 @@ void causal_conv1d_fwd_varlen_kernel_impl(

317 int64_t batch,

318 int64_t dim,

319 int64_t width,

319 - int64_t num_seq_blocks) {

320 + int64_t num_seq_blocks,

321 + int64_t conv_state_slot_stride) {

322 // handle 32 x 64 per block

323 constexpr int64_t BLOCK_M = block_size_m();

324 constexpr int64_t BLOCK_N = block_size_n() * 2;

@@ -366,7 +368,7 @@ void causal_conv1d_fwd_varlen_kernel_impl(

368 int32_t seqlen = query_start_loc[bs + 1] - query_start_loc[bs];

369 int32_t batch_offset = query_start_loc[bs];

370 update_conv_state(

369 - conv_states + conv_state_index * (width - 1) * dim,

371 + conv_states + conv_state_index * conv_state_slot_stride,

372 input + batch_offset * dim,

373 width,

374 dim,

@@ -389,7 +391,8 @@ void causal_conv1d_update_kernel_impl(

391 int64_t batch,

392 int64_t dim,

393 int64_t seqlen,

392 - int64_t width) {

394 + int64_t width,

395 + int64_t conv_state_slot_stride) {

396 // handle 32 x 64 per block

397 constexpr int64_t BLOCK_M = block_size_m();

398 constexpr int64_t BLOCK_N = block_size_n() * 2;

@@ -430,7 +433,7 @@ void causal_conv1d_update_kernel_impl(

433 });

434 });

435

433 -#define CONV_STATE_INDEXR(w) conv_states + conv_state_index*(width - 1) * dim + (w) * dim

436 +#define CONV_STATE_INDEXR(w) conv_states + conv_state_index*conv_state_slot_stride + (w) * dim

437

438 // update conv_states

439 at::parallel_for(0, batch, 0, [&](int64_t begin, int64_t end) {

@@ -592,6 +595,9 @@ at::Tensor causal_conv1d_fwd_cpu(

595 }

596 }

597

598 +// IMPORTANT: To make the kernal compatible with vLLM KV cache layout

599 + int64_t conv_state_slot_stride = conv_states->stride(0);

600 +

601 // block size for sequence blocks, 32

602 constexpr int64_t BLOCK_M = block_size_m();

603

@@ -618,7 +624,8 @@ at::Tensor causal_conv1d_fwd_cpu(

624 batch,

625 dim,

626 width,

621 - num_seq_blocks);

627 + num_seq_blocks,

628 + conv_state_slot_stride);

629 } else {

630 causal_conv1d_fwd_kernel_impl<scalar_t>(

631 out.data_ptr<scalar_t>(),

@@ -633,7 +640,8 @@ at::Tensor causal_conv1d_fwd_cpu(

640 dim,

641 seqlen,

642 width,

636 - num_seq_blocks);

643 + num_seq_blocks,

644 + conv_state_slot_stride);

645 }

646 });

647 return out;

@@ -690,6 +698,8 @@ at::Tensor causal_conv1d_update_cpu(

698 conv_states.copy_(conv_states_copy);

699 }

700

701 + // IMPORTANT: To make the kernal compatible with vLLM KV cache layout

702 + int64_t conv_state_slot_stride = conv_states.stride(0);

703 at::Tensor out = at::empty_like(x);

704 AT_DISPATCH_REDUCED_FLOATING_TYPES(scalar_type, "causal_conv1d_update_kernel_impl", [&] {

705 causal_conv1d_update_kernel_impl<scalar_t>(

@@ -703,7 +713,8 @@ at::Tensor causal_conv1d_update_cpu(

713 batch,

714 dim,

715 seqlen,

706 - width);

716 + width,

717 + conv_state_slot_stride);

718 });

719 return out;

720 }

csrc/cpu/sgl-kernels/fla.cpp

+5 -1

@@ -847,6 +847,7 @@ void fused_sigmoid_gating_delta_rule_update_kernel_impl(

847 int64_t v_strideB,

848 int64_t v_strideS,

849 int64_t v_strideH,

850 + int64_t state_slot_stride,

851 bool use_qk_l2norm_in_kernel,

852 double softplus_threshold) {

853 using bVec = at::vec::Vectorized<scalar_t>;

@@ -907,7 +908,7 @@ void fused_sigmoid_gating_delta_rule_update_kernel_impl(

908 data_index_init(begin, bi, batch_size, si, seq_len, ni, v_num_heads);

909 for (int64_t i = begin; i < end; ++i) {

910 int64_t cache_index = indices_ptr[bi];

910 - int64_t state_offset = (cache_index * v_num_heads + ni) * head_dim * v_head_dim;

911 + int64_t state_offset = cache_index * state_slot_stride + ni * head_dim * v_head_dim;

912 float g_val = -std::exp(float(A_log_ptr[ni])) *

913 softplus(float(a_ptr[bi * v_num_heads + ni]) + float(dt_bias_ptr[ni]), softplus_threshold);

914 float g_val_exp = std::exp(g_val);

@@ -1321,6 +1322,8 @@ at::Tensor fused_sigmoid_gating_delta_rule_update_cpu(

1322 int64_t v_strideB = v.stride(1);

1323 int64_t v_strideS = v.stride(0);

1324 int64_t v_strideH = v.stride(2);

1325 + // IMPORTANT: To make the kernal compatible with vLLM KV cache layout

1326 + int64_t state_slot_stride = initial_state_source.stride(0);

1327 at::Tensor core_attn_out = at::empty({batch_size, seq_len, v_num_heads, v_head_dim}, q.options());

1328 at::Tensor qk_scale_buf = at::empty({2 * batch_size, seq_len, num_heads}, at::kFloat);

1329

@@ -1353,6 +1356,7 @@ at::Tensor fused_sigmoid_gating_delta_rule_update_cpu(

1356 v_strideB,

1357 v_strideS,

1358 v_strideH,

1359 + state_slot_stride,

1360 use_qk_l2norm_in_kernel,

1361 softplus_threshold);

1362 });

csrc/cpu/torch_bindings.cpp

+80

@@ -98,6 +98,45 @@ at::Tensor int4_scaled_mm_cpu(at::Tensor& x, at::Tensor& w, at::Tensor& w_zeros,

98 at::Tensor& w_scales,

99 std::optional<at::Tensor> bias);

100

101 +// Adapted from sglang: GDN

102 +std::tuple<at::Tensor, at::Tensor> chunk_gated_delta_rule_cpu(

103 + const at::Tensor& query, const at::Tensor& key, const at::Tensor& value,

104 + const at::Tensor& g, const at::Tensor& beta,

105 + const at::Tensor& initial_state, bool output_final_state,

106 + const at::Tensor& cu_seqlens, bool head_first, bool use_qk_l2norm_in_kernel,

107 + double eps = 1e-5);

108 +

109 +at::Tensor fused_sigmoid_gating_delta_rule_update_cpu(

110 + const at::Tensor& A_log, const at::Tensor& dt_bias, const at::Tensor& q,

111 + const at::Tensor& k, const at::Tensor& v, const at::Tensor& a,

112 + const at::Tensor& b, at::Tensor& initial_state_source,

113 + const at::Tensor& initial_state_indices, const at::Tensor& cu_seqlens,

114 + bool use_qk_l2norm_in_kernel, double softplus_beta = 1.0,

115 + double softplus_threshold = 20.0);

116 +

117 +std::tuple<at::Tensor, at::Tensor> fused_gdn_gating_cpu(

118 + const at::Tensor& A_log, const at::Tensor& a, const at::Tensor& b,

119 + const at::Tensor& dt_bias);

120 +

121 +// Adapted from sglang: casual_conv1d kernels

122 +at::Tensor causal_conv1d_weight_pack(const at::Tensor& weight);

123 +

124 +at::Tensor causal_conv1d_fwd_cpu(

125 + const at::Tensor& x, const at::Tensor& weight,

126 + const std::optional<at::Tensor>& bias,

127 + const std::optional<at::Tensor>& conv_states,

128 + const std::optional<at::Tensor>& query_start_loc,

129 + const std::optional<at::Tensor>& cache_indices,

130 + const std::optional<at::Tensor>& has_initial_state, bool silu_activation,

131 + int64_t pad_slot_id, bool is_vnni);

132 +

133 +at::Tensor causal_conv1d_update_cpu(

134 + const at::Tensor& x, const at::Tensor& conv_states,

135 + const at::Tensor& weight, const std::optional<at::Tensor>& bias,

136 + bool silu_activation, const std::optional<at::Tensor>& cache_seqlens,

137 + const std::optional<at::Tensor>& conv_state_indices, int64_t pad_slot_id,

138 + bool is_vnni);

139 +

140 void activation_lut_bf16(torch::Tensor& out, torch::Tensor& input,

141 const std::string& activation);

142

@@ -397,6 +436,47 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {

436 "scales2, SymInt[] block_size, Tensor? bias, ScalarType out_dtype, "

437 "bool is_vnni) -> Tensor");

438 ops.impl("fp8_scaled_mm_cpu", torch::kCPU, &fp8_scaled_mm_cpu);

439 +

440 + // Adapted from sglang: GDN kernels

441 + ops.def(

442 + "chunk_gated_delta_rule_cpu(Tensor query, Tensor key, Tensor value, "

443 + "Tensor g, Tensor beta, "

444 + "Tensor initial_state, bool output_final_state, Tensor cu_seqlens, bool "

445 + "head_first, "

446 + "bool use_qk_l2norm_in_kernel, float eps=1e-5) -> (Tensor, Tensor)");

447 + ops.impl("chunk_gated_delta_rule_cpu", torch::kCPU,

448 + &chunk_gated_delta_rule_cpu);

449 + ops.def(

450 + "fused_sigmoid_gating_delta_rule_update_cpu(Tensor A_log, Tensor "

451 + "dt_bias, Tensor q, Tensor k, Tensor v, Tensor "

452 + "a, Tensor b, Tensor(a!) initial_state_source, Tensor "

453 + "initial_state_indices, Tensor cu_seqlens, bool "

454 + "use_qk_l2norm_in_kernel, float softplus_beta=1.0, float "

455 + "softplus_threshold=20.0) -> Tensor");

456 + ops.impl("fused_sigmoid_gating_delta_rule_update_cpu", torch::kCPU,

457 + &fused_sigmoid_gating_delta_rule_update_cpu);

458 + ops.def(

459 + "fused_gdn_gating_cpu(Tensor A_log, Tensor a, Tensor b, Tensor dt_bias) "

460 + "-> (Tensor, Tensor)");

461 + ops.impl("fused_gdn_gating_cpu", torch::kCPU, &fused_gdn_gating_cpu);

462 +

463 + // Adapted from sglang: casual_conv1d kernels

464 + ops.def("causal_conv1d_weight_pack(Tensor weight) -> Tensor");

465 + ops.impl("causal_conv1d_weight_pack", torch::kCPU,

466 + &causal_conv1d_weight_pack);

467 + ops.def(

468 + "causal_conv1d_fwd_cpu(Tensor x, Tensor weight, Tensor? bias, Tensor? "

469 + "conv_states, Tensor? query_start_loc,"

470 + "Tensor? cache_indices, Tensor? has_initial_state, bool silu_activation, "

471 + "int pad_slot_id, bool is_vnni) -> "

472 + "Tensor");

473 + ops.impl("causal_conv1d_fwd_cpu", torch::kCPU, &causal_conv1d_fwd_cpu);

474 + ops.def(

475 + "causal_conv1d_update_cpu(Tensor x, Tensor(a!) conv_states, Tensor "

476 + "weight, Tensor? bias, bool silu_activation,"

477 + "Tensor? cache_seqlens, Tensor? conv_state_indices, int pad_slot_id, "

478 + "bool is_vnni) -> Tensor");

479 + ops.impl("causal_conv1d_update_cpu", torch::kCPU, &causal_conv1d_update_cpu);

480 #endif

481

482 // CPU attention kernels

vllm/_custom_ops.py

+129

@@ -3370,6 +3370,135 @@ def fp8_scaled_mm_cpu(

3370 )

3371

3372

3373 +def chunk_gated_delta_rule_cpu(

3374 + query: torch.Tensor,

3375 + key: torch.Tensor,

3376 + value: torch.Tensor,

3377 + g: torch.Tensor,

3378 + beta: torch.Tensor,

3379 + initial_state: torch.Tensor,

3380 + output_final_state: bool,

3381 + cu_seqlens: torch.Tensor,

3382 + head_first: bool,

3383 + use_qk_l2norm_in_kernel: bool,

3384 + eps: float = 1e-5,

3385 +) -> tuple[torch.Tensor, torch.Tensor]:

3386 + return torch.ops._C.chunk_gated_delta_rule_cpu(

3387 + query,

3388 + key,

3389 + value,

3390 + g,

3391 + beta,

3392 + initial_state,

3393 + output_final_state,

3394 + cu_seqlens,

3395 + head_first,

3396 + use_qk_l2norm_in_kernel,

3397 + eps,

3398 + )

3399 +

3400 +

3401 +def fused_sigmoid_gating_delta_rule_update_cpu(

3402 + A_log: torch.Tensor,

3403 + dt_bias: torch.Tensor,

3404 + q: torch.Tensor,

3405 + k: torch.Tensor,

3406 + v: torch.Tensor,

3407 + a: torch.Tensor,

3408 + b: torch.Tensor,

3409 + initial_state_source: torch.Tensor,

3410 + initial_state_indices: torch.Tensor,

3411 + cu_seqlens: torch.Tensor,

3412 + use_qk_l2norm_in_kernel: bool,

3413 + softplus_beta: float = 1.0,

3414 + softplus_threshold: float = 20.0,

3415 +) -> torch.Tensor:

3416 + return torch.ops._C.fused_sigmoid_gating_delta_rule_update_cpu(

3417 + A_log,

3418 + dt_bias,

3419 + q,

3420 + k,

3421 + v,

3422 + a,

3423 + b,

3424 + initial_state_source,

3425 + initial_state_indices,

3426 + cu_seqlens,

3427 + use_qk_l2norm_in_kernel,

3428 + softplus_beta,

3429 + softplus_threshold,

3430 + )

3431 +

3432 +

3433 +def fused_gdn_gating_cpu(

3434 + A_log: torch.Tensor,

3435 + a: torch.Tensor,

3436 + b: torch.Tensor,

3437 + dt_bias: torch.Tensor,

3438 +) -> tuple[torch.Tensor, torch.Tensor]:

3439 + return torch.ops._C.fused_gdn_gating_cpu(

3440 + A_log,

3441 + a,

3442 + b,

3443 + dt_bias,

3444 + )

3445 +

3446 +

3447 +def causal_conv1d_weight_pack(

3448 + weight: torch.Tensor,

3449 +) -> torch.Tensor:

3450 + return torch.ops._C.causal_conv1d_weight_pack(

3451 + weight,

3452 + )

3453 +

3454 +

3455 +def causal_conv1d_fwd_cpu(

3456 + x: torch.Tensor,

3457 + weight: torch.Tensor,

3458 + bias: torch.Tensor | None,

3459 + conv_states: torch.Tensor | None,

3460 + query_start_loc: torch.Tensor | None,

3461 + cache_indices: torch.Tensor | None,

3462 + has_initial_state: torch.Tensor | None,

3463 + silu_activation: bool,

3464 + is_vnni: bool,

3465 +) -> torch.Tensor:

3466 + return torch.ops._C.causal_conv1d_fwd_cpu(

3467 + x,

3468 + weight,

3469 + bias,

3470 + conv_states,

3471 + query_start_loc,

3472 + cache_indices,

3473 + has_initial_state,

3474 + silu_activation,

3475 + -1,

3476 + is_vnni,

3477 + )

3478 +

3479 +

3480 +def causal_conv1d_update_cpu(

3481 + x: torch.Tensor,

3482 + conv_states: torch.Tensor,

3483 + weight: torch.Tensor,

3484 + bias: torch.Tensor | None,

3485 + silu_activation: bool,

3486 + conv_state_indices: torch.Tensor | None,

3487 + is_vnni: bool,

3488 +) -> torch.Tensor:

3489 + return torch.ops._C.causal_conv1d_update_cpu(

3490 + x,

3491 + conv_states,

3492 + weight,

3493 + bias,

3494 + silu_activation,

3495 + None,

3496 + conv_state_indices,

3497 + -1,

3498 + is_vnni,

3499 + )

3500 +

3501 +

3502 class CPUDNNLGEMMHandler:

3503 def __init__(self) -> None:

3504 self.handler_tensor: torch.Tensor | None = None

vllm/model_executor/layers/linear.py

-3

@@ -213,9 +213,6 @@ class UnquantizedLinearMethod(LinearMethodBase):

213

214 def process_weights_after_loading(self, layer: torch.nn.Module) -> None:

215 if current_platform.is_cpu():

216 - if layer.weight.ndim != 2:

217 - # this is not a linear layer

218 - return

216 from vllm.model_executor.layers.utils import dispatch_cpu_unquantized_gemm

217

218 dispatch_cpu_unquantized_gemm(layer, remove_weight=True)

vllm/model_executor/layers/mamba/ops/cpu/gdn_attention.py

+140

@@ -5,6 +5,7 @@ from __future__ import annotations

6 import torch

8 +import vllm._custom_ops as ops

9 from vllm.forward_context import ForwardContext, get_forward_context

10 from vllm.model_executor.layers.mamba.mamba_utils import is_conv_state_dim_first

11 from vllm.model_executor.layers.mamba.ops.cpu.causal_conv1d import (

@@ -55,6 +56,16 @@ def cpu_gdn_attention_core(

56 and attn_metadata_i.num_accepted_tokens is None

57 ), "speculative decode not supported in CPU GDN attention."

59 + if torch.cpu._is_amx_tile_supported():

60 + return cpu_gdn_attention_core_amx(

61 + mixed_qkv,

62 + b,

63 + a,

64 + core_attn_out,

65 + attn_metadata_i,

66 + layer,

67 + )

68 +

69 state_indices_tensor = attn_metadata_i.non_spec_state_indices_tensor

70 query_start_loc = attn_metadata_i.non_spec_query_start_loc

71 assert state_indices_tensor is not None

@@ -194,6 +205,135 @@ def cpu_gdn_attention_core_fake(

205 return

206

207

208 +def cpu_gdn_attention_core_amx(

209 + mixed_qkv: torch.Tensor,

210 + b: torch.Tensor,

211 + a: torch.Tensor,

212 + core_attn_out: torch.Tensor,

213 + attn_metadata_i: GDNAttentionMetadata,

214 + layer: torch.nn.Module,

215 +):

216 + state_indices_tensor = attn_metadata_i.non_spec_state_indices_tensor

217 + query_start_loc = attn_metadata_i.non_spec_query_start_loc

218 + assert state_indices_tensor is not None

219 + assert query_start_loc is not None

220 +

221 + # [num_allocated_slots, kernel - 1, conv_dim]

222 + conv_state = layer.kv_cache[0]

223 + if is_conv_state_dim_first():

224 + raise RuntimeError("AMX GDN attention requires `SD` conv_state layout.")

225 + # reshape to [num_allocated_slots, conv_dim, kernel - 1]

226 + conv_state_t = conv_state.transpose(1, 2)

227 +

228 + # [num_allocated_slots, num_v_heads / tp_size, v_dim, k_dim]

229 + ssm_state = layer.kv_cache[1]

230 + # rehape to [num_allocated_slots, num_v_heads / tp_size, k_dim, v_dim]

231 + num_allocated_slots, head_num, v_dim, k_dim = ssm_state.size()

232 + ssm_state = ssm_state.view(

233 + num_allocated_slots,

234 + head_num,

235 + k_dim,

236 + v_dim,

237 + )

238 +

239 + mixed_qkv = mixed_qkv.contiguous()

240 + a = a.contiguous()

241 + b = b.contiguous()

242 +

243 + num_decodes = attn_metadata_i.num_decodes

244 + num_decode_tokens = attn_metadata_i.num_decode_tokens

245 + num_prefills = attn_metadata_i.num_prefills

246 + num_prefill_tokens = attn_metadata_i.num_prefill_tokens

247 +

248 + if num_decodes > 0:

249 + decode_mixed_qkv = mixed_qkv[:num_decode_tokens]

250 + decode_b = b[:num_decode_tokens]

251 + decode_a = a[:num_decode_tokens]

252 + decode_state_indices = state_indices_tensor[:num_decodes]

253 +

254 + decode_mixed_qkv = ops.causal_conv1d_update_cpu(

255 + x=decode_mixed_qkv,

256 + conv_states=conv_state_t,

257 + weight=layer.conv1d.weight,

258 + bias=layer.conv1d.bias,

259 + silu_activation=layer.activation == "silu",

260 + conv_state_indices=decode_state_indices,

261 + is_vnni=True,

262 + )

263 +

264 + query, key, value = layer.rearrange_mixed_qkv(decode_mixed_qkv)

265 + attn_out = ops.fused_sigmoid_gating_delta_rule_update_cpu(

266 + A_log=layer.A_log,

267 + dt_bias=layer.dt_bias,

268 + q=query,

269 + k=key,

270 + v=value,

271 + a=decode_a,

272 + b=decode_b,

273 + initial_state_source=ssm_state,

274 + initial_state_indices=decode_state_indices,

275 + cu_seqlens=query_start_loc[: num_decodes + 1],

276 + use_qk_l2norm_in_kernel=True,

277 + )

278 + core_attn_out[:num_decode_tokens] = attn_out.squeeze(1)

279 +

280 + if num_prefills > 0:

281 + has_initial_state = attn_metadata_i.has_initial_state

282 + assert has_initial_state is not None

283 +

284 + prefill_token_start = num_decode_tokens

285 + prefill_token_end = prefill_token_start + num_prefill_tokens

286 + prefill_mixed_qkv = mixed_qkv[prefill_token_start:prefill_token_end]

287 + prefill_b = b[prefill_token_start:prefill_token_end]

288 + prefill_a = a[prefill_token_start:prefill_token_end]

289 + prefill_state_indices = state_indices_tensor[

290 + num_decodes : num_decodes + num_prefills

291 + ]

292 + prefill_query_start_loc = (

293 + query_start_loc[num_decodes : num_decodes + num_prefills + 1]

294 + - num_decode_tokens

295 + )

296 + prefill_has_initial_state = has_initial_state[

297 + num_decodes : num_decodes + num_prefills

298 + ]

299 +

300 + prefill_mixed_qkv = ops.causal_conv1d_fwd_cpu(

301 + x=prefill_mixed_qkv.transpose(0, 1),

302 + weight=layer.conv1d.weight,

303 + bias=layer.conv1d.bias,

304 + conv_states=conv_state_t,

305 + query_start_loc=prefill_query_start_loc,

306 + cache_indices=prefill_state_indices,

307 + has_initial_state=prefill_has_initial_state,

308 + silu_activation=layer.activation == "silu",

309 + is_vnni=True,

310 + ).transpose(0, 1)

311 +

312 + query, key, value = layer.rearrange_mixed_qkv(prefill_mixed_qkv)

313 + g, beta = ops.fused_gdn_gating_cpu(

314 + A_log=layer.A_log, a=prefill_a, b=prefill_b, dt_bias=layer.dt_bias

315 + )

316 +

317 + initial_state = ssm_state[prefill_state_indices]

318 + initial_state[~prefill_has_initial_state, ...] = 0

319 + attn_out, last_recurrent_state = ops.chunk_gated_delta_rule_cpu(

320 + query=query,

321 + key=key,

322 + value=value,

323 + g=g,

324 + beta=beta,

325 + initial_state=initial_state,

326 + output_final_state=True,

327 + cu_seqlens=prefill_query_start_loc,

328 + head_first=False,

329 + use_qk_l2norm_in_kernel=True,

330 + )

331 + ssm_state[prefill_state_indices] = last_recurrent_state.to(

332 + ssm_state.dtype, copy=False

333 + )

334 + core_attn_out[prefill_token_start:prefill_token_end] = attn_out.squeeze(0)

335 +

336 +

337 def register_cpu_gdn_attention_ops() -> None:

338 global _CPU_GDN_ATTENTION_OPS_REGISTERED

339 if _CPU_GDN_ATTENTION_OPS_REGISTERED:

vllm/model_executor/layers/utils.py

+13

@@ -233,6 +233,19 @@ def dispatch_cpu_unquantized_gemm(

233 layer.cpu_linear = torch.nn.functional.linear

234 return

235

236 + if layer.weight.ndim != 2:

237 + # this is not a linear layer

238 + # For now it should be a causal_conv1d op

239 + if torch.cpu._is_amx_tile_supported():

240 + # prepack conv weight

241 + layer.weight.data = ops.causal_conv1d_weight_pack(

242 + layer.weight.view(

243 + layer.weight.size(0),

244 + layer.weight.size(2),

245 + )

246 + )

247 + return

248 +

249 N, K = layer.weight.size()

250 dtype = layer.weight.dtype

251

vllm/platforms/cpu.py

+19

@@ -139,6 +139,7 @@ class CpuPlatform(Platform):

139 # supported as it is not possible to set the OMP environment correctly

140 if parallel_config.distributed_executor_backend == "uni":

141 parallel_config.distributed_executor_backend = "mp"

142 +

143 if parallel_config.worker_cls == "auto":

144 parallel_config.worker_cls = "vllm.v1.worker.cpu_worker.CPUWorker"

145 # Disable DBO

@@ -146,6 +147,20 @@ class CpuPlatform(Platform):

147 logger.warning("Dual-Batch Overlap is not supported on CPU, disabled.")

148 parallel_config.enable_dbo = False

149

150 + if torch.cpu._is_amx_tile_supported() and (

151 + model_config is not None

152 + and model_config.get_num_layers_by_block_type(

153 + parallel_config, "linear_attention"

154 + )

155 + > 0

156 + ):

157 + cache_config.enable_prefix_caching = False

158 + scheduler_config.enable_chunked_prefill = False

159 + logger.warning(

160 + "Disabled unsupported prefix caching and chunked prefill "

161 + "for linear attention on AMX CPU platforms."

162 + )

163 +

164 # Note: workaround for v1 gpu_model_runner

165 from vllm.config import CompilationMode

166

@@ -211,6 +226,10 @@ class CpuPlatform(Platform):

226 # Avoid inductor generates num_thread() and breaks the thread binding

227 os.environ["TORCHINDUCTOR_CPP_DYNAMIC_THREADS"] = "1"

228

229 + # For efficient conv state memory access

230 + if torch.cpu._is_amx_tile_supported():

231 + os.environ["VLLM_SSM_CONV_STATE_LAYOUT"] = "SD"

232 +

233 ld_preload_str = os.getenv("LD_PRELOAD", "")

234 cpu_architecture = Platform.get_cpu_architecture()

235