@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / a7fdfeef72
Fork 0

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

[CPU] Support Gemma Diffusion (#45690)

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

Li, Jiang committed Jun 16, 2026 at 14:39 UTC a7fdfeef72323eb3db6f0620e4ea200290d0ca5a
7 files changed +213 -100
csrc/cpu/cpu_attn.cpp
+14 -25
@@ -15,9 +15,10 @@ torch::Tensor get_scheduler_metadata(
15 const int64_t num_req, const int64_t num_heads_q,
16 const int64_t num_heads_kv, const int64_t head_dim,
17 const torch::Tensor& seq_lens, at::ScalarType dtype,
18 - const torch::Tensor& query_start_loc, const bool casual,
18 + const torch::Tensor& query_start_loc, const bool causal,
19 const int64_t window_size, const std::string& isa_hint,
20 - const bool enable_kv_split) {
20 + const bool enable_kv_split,
21 + const std::optional<torch::Tensor>& dynamic_causal) {
22 cpu_attention::ISA isa;
23 if (isa_hint == "amx") {
24 isa = cpu_attention::ISA::AMX;
@@ -44,24 +45,13 @@ torch::Tensor get_scheduler_metadata(
45 input.head_dim = head_dim;
46 input.query_start_loc = query_start_loc.data_ptr<int32_t>();
47 input.seq_lens = seq_lens.data_ptr<int32_t>();
47 - if (window_size != -1) {
48 - input.left_sliding_window_size = window_size - 1;
49 - if (casual) {
50 - input.right_sliding_window_size = 0;
51 - } else {
52 - input.right_sliding_window_size = window_size - 1;
53 - }
54 - } else {
55 - input.left_sliding_window_size = -1;
56 - if (casual) {
57 - input.right_sliding_window_size = 0;
58 - } else {
59 - input.right_sliding_window_size = -1;
60 - }
61 - }
62 - input.casual = casual;
48 +
49 + input.sliding_window_size = window_size;
50 + input.causal = causal;
51 input.isa = isa;
52 input.enable_kv_split = enable_kv_split;
53 + input.dynamic_causal =
54 + dynamic_causal.has_value() ? dynamic_causal->data_ptr<bool>() : nullptr;
55
56 VLLM_DISPATCH_FLOATING_TYPES(dtype, "get_scheduler_metadata", [&]() {
57 CPU_ATTN_DISPATCH(head_dim, isa, 0, [&]() {
@@ -175,10 +165,11 @@ void cpu_attention_with_kv_cache(
165 const torch::Tensor& seq_lens, // [num_tokens]
166 const double scale, const bool causal,
167 const std::optional<torch::Tensor>& alibi_slopes, // [num_heads]
178 - const int64_t sliding_window_left, const int64_t sliding_window_right,
168 + const int64_t sliding_window,
169 const torch::Tensor& block_table, // [num_tokens, max_block_num]
170 const double softcap, const torch::Tensor& scheduler_metadata,
181 - const std::optional<torch::Tensor>& s_aux, // [num_heads]
171 + const std::optional<torch::Tensor>& s_aux, // [num_heads]
172 + const std::optional<torch::Tensor>& dynamic_causal, // [num_reqs]
173 const double k_scale = 1.0, const double v_scale = 1.0,
174 const std::string& kv_cache_dtype = "auto") {
175 TORCH_CHECK_EQ(query.dim(), 3);
@@ -220,13 +211,11 @@ void cpu_attention_with_kv_cache(
211 input.alibi_slopes =
212 alibi_slopes.has_value() ? alibi_slopes->data_ptr<float>() : nullptr;
213 input.s_aux = s_aux.has_value() ? s_aux->data_ptr<c10::BFloat16>() : nullptr;
214 + input.dynamic_causal =
215 + dynamic_causal.has_value() ? dynamic_causal->data_ptr<bool>() : nullptr;
216 input.scale = scale;
217 input.causal = causal;
225 - input.sliding_window_left = sliding_window_left;
226 - input.sliding_window_right = sliding_window_right;
227 - if (input.causal) {
228 - input.sliding_window_right = 0;
229 - }
218 + input.sliding_window_size = sliding_window;
219 input.softcap = static_cast<float>(softcap);
220
221 if (is_fp8) {
csrc/cpu/cpu_attn_impl.hpp
+62 -27
@@ -388,13 +388,13 @@ class AttentionScheduler {
388 int32_t head_dim;
389 int32_t* query_start_loc;
390 int32_t* seq_lens;
391 - int32_t left_sliding_window_size;
392 - int32_t right_sliding_window_size;
393 - bool casual;
391 + int32_t sliding_window_size;
392 + bool causal;
393 cpu_attention::ISA isa;
394 int32_t max_num_q_per_iter; // max Q head num can be hold in registers
395 int32_t kv_block_alignment; // context length alignment requirement
396 bool enable_kv_split;
397 + bool* dynamic_causal;
398 };
399
400 static constexpr int32_t MaxQTileIterNum = 128;
@@ -403,7 +403,8 @@ class AttentionScheduler {
403 : available_cache_size_(cpu_utils::get_available_l2_size()) {}
404
405 torch::Tensor schedule(const ScheduleInput& input) const {
406 - const bool casual = input.casual;
406 + const bool causal = input.causal;
407 + const bool is_dynamic_causal = input.dynamic_causal != nullptr;
408 const int32_t thread_num = omp_get_max_threads();
409 const int64_t cache_size = cpu_utils::get_available_l2_size();
410 const int32_t max_num_q_per_iter = input.max_num_q_per_iter;
@@ -434,8 +435,7 @@ class AttentionScheduler {
435 const int32_t default_tile_token_num = default_tile_size / q_head_per_kv;
436 const int32_t split_kv_q_token_num_threshold =
437 input.enable_kv_split ? 1 : 0;
437 - const int32_t left_sliding_window_size = input.left_sliding_window_size;
438 - const int32_t right_sliding_window_size = input.right_sliding_window_size;
438 + const int32_t sliding_window_size = input.sliding_window_size;
439 TORCH_CHECK_LE(split_kv_q_token_num_threshold * q_head_per_kv, 16);
440
441 // get total kv len
@@ -444,7 +444,9 @@ class AttentionScheduler {
444 const int32_t seq_len = input.seq_lens[req_id];
445 const int32_t q_token_num =
446 input.query_start_loc[req_id + 1] - input.query_start_loc[req_id];
447 - const int32_t q_start_pos = (casual ? (seq_len - q_token_num) : 0);
447 + const bool req_causal =
448 + is_dynamic_causal ? input.dynamic_causal[req_id] : causal;
449 + const int32_t q_start_pos = seq_len - q_token_num;
450 const int32_t kv_start_pos = 0;
451 const int32_t kv_end_pos = seq_len;
452
@@ -456,7 +458,7 @@ class AttentionScheduler {
458 const int32_t q_tile_pos_right = q_tile_pos_left + q_tile_token_num;
459 const auto [kv_tile_pos_left, kv_tile_pos_right] = calcu_kv_tile_pos(
460 kv_start_pos, kv_end_pos, q_tile_pos_left, q_tile_pos_right,
459 - left_sliding_window_size, right_sliding_window_size);
461 + sliding_window_size, req_causal);
462 const auto [aligned_kv_tile_pos_left, aligned_kv_tile_pos_right] =
463 align_kv_tile_pos(kv_tile_pos_left, kv_tile_pos_right,
464 kv_len_alignment);
@@ -484,7 +486,9 @@ class AttentionScheduler {
486 const int32_t seq_len = input.seq_lens[req_id];
487 const int32_t q_token_num =
488 input.query_start_loc[req_id + 1] - input.query_start_loc[req_id];
487 - const int32_t q_start_pos = (casual ? (seq_len - q_token_num) : 0);
489 + const bool req_causal =
490 + is_dynamic_causal ? input.dynamic_causal[req_id] : causal;
491 + const int32_t q_start_pos = seq_len - q_token_num;
492 const int32_t kv_start_pos = 0;
493 const int32_t kv_end_pos = seq_len;
494 int32_t local_split_id = 0;
@@ -498,7 +502,7 @@ class AttentionScheduler {
502 const int32_t q_tile_pos_right = q_tile_pos_left + q_tile_token_num;
503 const auto [kv_tile_pos_left, kv_tile_pos_right] = calcu_kv_tile_pos(
504 kv_start_pos, kv_end_pos, q_tile_pos_left, q_tile_pos_right,
501 - left_sliding_window_size, right_sliding_window_size);
505 + sliding_window_size, req_causal);
506 const auto [aligned_kv_tile_pos_left, aligned_kv_tile_pos_right] =
507 align_kv_tile_pos(kv_tile_pos_left, kv_tile_pos_right,
508 kv_len_alignment);
@@ -708,15 +712,41 @@ class AttentionScheduler {
712 return metadata_tensor;
713 }
714
715 + FORCE_INLINE static std::pair<int32_t, int32_t> calcu_sliding_window_size(
716 + int32_t window_size, bool causal) {
717 + int32_t left_sliding_window_size, right_sliding_window_size;
718 + if (window_size != -1) {
719 + left_sliding_window_size = window_size - 1;
720 + if (causal) {
721 + right_sliding_window_size = 0;
722 + } else {
723 + right_sliding_window_size = window_size - 1;
724 + }
725 + } else {
726 + left_sliding_window_size = -1;
727 + if (causal) {
728 + right_sliding_window_size = 0;
729 + } else {
730 + right_sliding_window_size = -1;
731 + }
732 + }
733 +
734 + return {left_sliding_window_size, right_sliding_window_size};
735 + }
736 +
737 FORCE_INLINE static std::pair<int32_t, int32_t> calcu_kv_tile_pos(
738 int32_t kv_left_pos, int32_t kv_right_pos, int32_t q_left_pos,
713 - int32_t q_right_pos, int32_t sliding_window_left,
714 - int32_t sliding_window_right) {
715 - if (sliding_window_left != -1) {
716 - kv_left_pos = std::max(kv_left_pos, q_left_pos - sliding_window_left);
739 + int32_t q_right_pos, int32_t window_size, bool causal) {
740 + auto [left_sliding_window_size, right_sliding_window_size] =
741 + calcu_sliding_window_size(window_size, causal);
742 +
743 + if (left_sliding_window_size != -1) {
744 + kv_left_pos =
745 + std::max(kv_left_pos, q_left_pos - left_sliding_window_size);
746 }
718 - if (sliding_window_right != -1) {
719 - kv_right_pos = std::min(kv_right_pos, q_right_pos + sliding_window_right);
747 + if (right_sliding_window_size != -1) {
748 + kv_right_pos =
749 + std::min(kv_right_pos, q_right_pos + right_sliding_window_size);
750 }
751 return {kv_left_pos, kv_right_pos};
752 }
@@ -805,10 +835,10 @@ struct AttentionInput {
835 int32_t* block_table;
836 float* alibi_slopes;
837 c10::BFloat16* s_aux;
838 + bool* dynamic_causal;
839 float scale;
840 bool causal;
810 - int32_t sliding_window_left;
811 - int32_t sliding_window_right;
841 + int32_t sliding_window_size;
842 float softcap;
843 // FP8 KV cache scales (used by FP8 attention implementations)
844 float k_scale_fp8 = 1.0f;
@@ -1442,15 +1472,16 @@ class AttentionMainLoop {
1472 const int64_t q_head_num_stride = input->query_num_heads_stride;
1473 const int64_t kv_cache_head_num_stride = input->cache_num_kv_heads_stride;
1474 const int64_t kv_cache_block_num_stride = input->cache_num_blocks_stride;
1445 - const int32_t sliding_window_left = input->sliding_window_left;
1446 - const int32_t sliding_window_right = input->sliding_window_right;
1475 + const int32_t sliding_window_size = input->sliding_window_size;
1476 const int32_t block_size = input->block_size;
1477 const float scale = input->scale;
1478 const float softcap_scale = input->softcap;
1479 const float* alibi_slopes = input->alibi_slopes;
1480 const c10::BFloat16* s_aux = input->s_aux;
1481 + const bool* dynamic_causal = input->dynamic_causal;
1482 + const bool is_dynamic_causal = dynamic_causal != nullptr;
1483
1453 - const bool casual = input->causal;
1484 + const bool causal = input->causal;
1485 int32_t* const block_table = input->block_table;
1486 const int64_t block_table_stride = input->blt_num_tokens_stride;
1487
@@ -1533,6 +1564,11 @@ class AttentionMainLoop {
1564 &curr_workitem_groups[workitem_group_idx];
1565
1566 const int32_t current_group_idx = current_workitem_group->req_id;
1567 + const int32_t current_group_causal =
1568 + is_dynamic_causal ? dynamic_causal[current_group_idx] : causal;
1569 + auto [sliding_window_left, sliding_window_right] =
1570 + AttentionScheduler::calcu_sliding_window_size(
1571 + sliding_window_size, current_group_causal);
1572 const int32_t kv_start_pos =
1573 current_workitem_group->kv_split_pos_start;
1574 const int32_t kv_end_pos = current_workitem_group->kv_split_pos_end;
@@ -1560,8 +1596,7 @@ class AttentionMainLoop {
1596 const int32_t q_end = input->query_start_loc[current_group_idx + 1];
1597 const int32_t q_start = input->query_start_loc[current_group_idx];
1598 const int32_t seq_len = input->seq_lens[current_group_idx];
1563 - const int32_t q_start_pos =
1564 - (casual ? seq_len - (q_end - q_start) : 0);
1599 + const int32_t q_start_pos = seq_len - (q_end - q_start);
1600 const int32_t block_num = (seq_len + block_size - 1) / block_size;
1601 // Only apply sink for the first KV split
1602 bool use_sink = (s_aux != nullptr &&
@@ -1611,8 +1646,8 @@ class AttentionMainLoop {
1646 const auto [kv_tile_start_pos, kv_tile_end_pos] =
1647 AttentionScheduler::calcu_kv_tile_pos(
1648 kv_start_pos, kv_end_pos, q_tile_start_pos,
1614 - q_tile_end_pos, sliding_window_left,
1615 - sliding_window_right);
1649 + q_tile_end_pos, sliding_window_size,
1650 + current_group_causal);
1651 const auto [rounded_kv_tile_start_pos, rounded_kv_tile_end_pos] =
1652 AttentionScheduler::align_kv_tile_pos(
1653 kv_tile_start_pos, kv_tile_end_pos, blocksize_alignment);
@@ -1725,8 +1760,8 @@ class AttentionMainLoop {
1760 actual_kv_tile_pos_right] =
1761 AttentionScheduler::calcu_kv_tile_pos(
1762 kv_tile_pos_left, kv_tile_pos_right, q_tile_pos_left,
1728 - q_tile_pos_right, sliding_window_left,
1729 - sliding_window_right);
1763 + q_tile_pos_right, sliding_window_size,
1764 + current_group_causal);
1765 const int32_t q_iter_idx =
1766 q_head_tile_token_offset / curr_max_q_token_num_per_iter;
1767
csrc/cpu/cpu_fused_moe.cpp
+5 -7
@@ -1,3 +1,5 @@
1 +#include <sleef.h>
2 +
3 #include "cpu/cpu_types.hpp"
4 #include "cpu/utils.hpp"
5 #include "cpu/micro_gemm/cpu_micro_gemm_vec.hpp"
@@ -163,7 +165,6 @@ void gelu_tanh_and_mul(float* __restrict__ input, scalar_t* __restrict__ output,
165 vec_op::FP32Vec16 w1_vec(0.7978845608028654);
166 vec_op::FP32Vec16 w2_vec(0.5);
167 vec_op::FP32Vec16 w3_vec(0.044715);
166 - alignas(64) float temp[16];
168
169 for (int32_t m = 0; m < m_size; ++m) {
170 for (int32_t n = 0; n < dim; n += 16) {
@@ -171,12 +172,9 @@ void gelu_tanh_and_mul(float* __restrict__ input, scalar_t* __restrict__ output,
172 vec_op::FP32Vec16 up_vec(up + n);
173 auto gate_pow3_vec = gate_vec * gate_vec * gate_vec;
174 auto inner_vec = w1_vec * (gate_vec + w3_vec * gate_pow3_vec);
174 -
175 - inner_vec.save(temp);
176 - for (int32_t i = 0; i < 16; ++i) {
177 - temp[i] = std::tanh(temp[i]);
178 - }
179 - vec_op::FP32Vec16 tanh_vec(temp);
175 + // Note: can't use fast_exp form because diffusiongemma will generate
176 + // wrong results
177 + vec_op::FP32Vec16 tanh_vec(Sleef_tanhf16_u10(inner_vec.reg));
178 auto gelu_tanh = gate_vec * w2_vec * (one_vec + tanh_vec);
179 auto gated_output_fp32 = up_vec * gelu_tanh;
180 scalar_vec_t gated_output = scalar_vec_t(gated_output_fp32);
csrc/cpu/torch_bindings.cpp
+10 -8
@@ -152,7 +152,8 @@ torch::Tensor get_scheduler_metadata(
152 const torch::Tensor& seq_lens, at::ScalarType dtype,
153 const torch::Tensor& query_start_loc, const bool casual,
154 const int64_t window_size, const std::string& isa_hint,
155 - const bool enable_kv_split);
155 + const bool enable_kv_split,
156 + const std::optional<torch::Tensor>& dynamic_causal);
157
158 void cpu_attn_reshape_and_cache(const torch::Tensor& key,
159 const torch::Tensor& value,
@@ -169,10 +170,10 @@ void cpu_attention_with_kv_cache(
170 const torch::Tensor& query_start_loc, const torch::Tensor& seq_lens,
171 const double scale, const bool causal,
172 const std::optional<torch::Tensor>& alibi_slopes,
172 - const int64_t sliding_window_left, const int64_t sliding_window_right,
173 - const torch::Tensor& block_table, const double softcap,
174 - const torch::Tensor& scheduler_metadata,
175 - const std::optional<torch::Tensor>& s_aux, const double k_scale,
173 + const int64_t sliding_window_left, const torch::Tensor& block_table,
174 + const double softcap, const torch::Tensor& scheduler_metadata,
175 + const std::optional<torch::Tensor>& s_aux,
176 + const std::optional<torch::Tensor>& dynamic_causal, const double k_scale,
177 const double v_scale, const std::string& kv_cache_dtype);
178
179 // Note: just for avoiding importing errors
@@ -500,7 +501,7 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
501 "get_scheduler_metadata(int num_req, int num_heads_q, int num_heads_kv, "
502 "int head_dim, Tensor seq_lens, ScalarType dtype, Tensor "
503 "query_start_loc, bool casual, int window_size, str isa_hint, bool "
503 - "enable_kv_split) -> Tensor",
504 + "enable_kv_split, Tensor? dynamic_causal) -> Tensor",
505 &get_scheduler_metadata);
506 ops.def(
507 "cpu_attn_reshape_and_cache(Tensor key, Tensor value, Tensor(a2!) "
@@ -512,8 +513,9 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
513 "cpu_attention_with_kv_cache(Tensor query, Tensor key_cache, Tensor "
514 "value_cache, Tensor(a3!) output, Tensor query_start_loc, Tensor "
515 "seq_lens, float scale, bool causal, Tensor? alibi_slopes, SymInt "
515 - "sliding_window_left, SymInt sliding_window_right, Tensor block_table, "
516 - "float softcap, Tensor scheduler_metadata, Tensor? s_aux, "
516 + "sliding_window_size, Tensor block_table, "
517 + "float softcap, Tensor scheduler_metadata, Tensor? s_aux, Tensor? "
518 + "dynamic_causal, "
519 "float k_scale=1.0, float v_scale=1.0, str kv_cache_dtype=\"auto\") -> "
520 "()",
521 &cpu_attention_with_kv_cache);
tests/kernels/attention/test_cpu_attn.py
+100 -25
@@ -107,6 +107,7 @@ def ref_paged_attn(
107 soft_cap: float | None = None,
108 alibi_slopes: torch.Tensor | None = None,
109 s_aux: torch.Tensor | None = None,
110 + dynamic_causal: list[bool] | None = None,
111 ) -> torch.Tensor:
112 num_seqs = len(query_lens)
113 block_tables = block_tables.cpu().numpy()
@@ -142,17 +143,30 @@ def ref_paged_attn(
143 v = torch.repeat_interleave(v, q.shape[1] // v.shape[1], dim=1)
144 attn = torch.einsum("qhd,khd->hqk", q, k).float()
145 empty_mask = torch.ones(query_len, kv_len)
145 - mask = torch.triu(empty_mask, diagonal=kv_len - query_len + 1).bool()
146
147 - if sliding_window is not None:
148 - sliding_window_mask = (
149 - torch.triu(
150 - empty_mask, diagonal=kv_len - (query_len + sliding_window) + 1
147 + if dynamic_causal is None or dynamic_causal[i]:
148 + mask = torch.triu(empty_mask, diagonal=kv_len - query_len + 1).bool()
149 + if sliding_window is not None:
150 + sliding_window_mask = (
151 + torch.triu(
152 + empty_mask, diagonal=kv_len - (query_len + sliding_window) + 1
153 + )
154 + .bool()
155 + .logical_not()
156 )
152 - .bool()
153 - .logical_not()
154 - )
155 - mask |= sliding_window_mask
157 + mask |= sliding_window_mask
158 + else:
159 + if sliding_window is not None:
160 + mask = (
161 + torch.triu(
162 + empty_mask, diagonal=1 - sliding_window + kv_len - query_len
163 + ).bool()
164 + ^ torch.triu(
165 + empty_mask, diagonal=sliding_window + kv_len - query_len
166 + ).bool()
167 + ).logical_not()
168 + else:
169 + mask = empty_mask.logical_not()
170
171 if soft_cap is not None:
172 attn = soft_cap * torch.tanh(attn / soft_cap)
@@ -243,11 +257,6 @@ def varlen_encoder_attention(
257 num_query_heads = num_heads[0]
258 num_kv_heads = num_heads[1]
259 assert num_query_heads % num_kv_heads == 0
246 - window_size = (
247 - (sliding_window - 1, sliding_window - 1)
248 - if sliding_window is not None
249 - else (-1, -1)
250 - )
260 scale = head_size**-0.5
261 token_num = sum(seq_lens)
262
@@ -343,7 +352,7 @@ def varlen_encoder_attention(
352 scale=scale,
353 causal=False,
354 alibi_slopes=None,
346 - sliding_window=window_size,
355 + sliding_window=sliding_window if sliding_window is not None else -1,
356 block_table=encoder_block_table,
357 softcap=0,
358 scheduler_metadata=metadata,
@@ -375,7 +384,7 @@ def varlen_encoder_attention(
384 scale=scale,
385 causal=False,
386 alibi_slopes=None,
378 - sliding_window=window_size,
387 + sliding_window=sliding_window if sliding_window is not None else -1,
388 block_table=encoder_block_table,
389 softcap=0,
390 scheduler_metadata=metadata,
@@ -418,6 +427,7 @@ def varlen_with_paged_kv(
427 kv_cache_dtype: str = "auto",
428 k_scale: float = 1.0,
429 v_scale: float = 1.0,
430 + dynamic_causal: list[bool] | None = None,
431 ) -> None:
432 set_random_seed(0)
433 num_seqs = len(seq_lens)
@@ -427,9 +437,13 @@ def varlen_with_paged_kv(
437 num_kv_heads = num_heads[1]
438 assert num_query_heads % num_kv_heads == 0
439 max_kv_len = max(kv_lens)
430 - window_size = (sliding_window - 1, 0) if sliding_window is not None else (-1, -1)
440 scale = head_size**-0.5
441 token_num = sum(query_lens)
442 + dynamic_causal_tensor = (
443 + torch.tensor(dynamic_causal, dtype=torch.bool)
444 + if dynamic_causal is not None
445 + else None
446 + )
447
448 # for n heads the set of slopes is the geometric sequence that starts
449 # 2^(-8/n)
@@ -515,10 +529,11 @@ def varlen_with_paged_kv(
529 seq_lens=kv_lens_tensor,
530 dtype=dtype,
531 query_start_loc=cu_query_lens,
518 - causal=True,
532 + causal=dynamic_causal is None,
533 sliding_window_size=sliding_window if sliding_window is not None else -1,
534 isa=isa,
535 enable_kv_split=False,
536 + dynamic_causal=dynamic_causal_tensor,
537 )
538
539 out_without_split = torch.empty_like(query)
@@ -530,13 +545,14 @@ def varlen_with_paged_kv(
545 query_start_loc=cu_query_lens,
546 seq_lens=kv_lens_tensor,
547 scale=scale,
533 - causal=True,
548 + causal=dynamic_causal is None,
549 alibi_slopes=alibi_slopes,
535 - sliding_window=window_size,
550 + sliding_window=sliding_window if sliding_window is not None else -1,
551 block_table=block_tables,
552 softcap=soft_cap if soft_cap is not None else 0,
553 scheduler_metadata=metadata,
554 s_aux=s_aux,
555 + dynamic_causal=dynamic_causal_tensor,
556 **fp8_kwargs,
557 )
558
@@ -548,10 +564,11 @@ def varlen_with_paged_kv(
564 seq_lens=kv_lens_tensor,
565 dtype=dtype,
566 query_start_loc=cu_query_lens,
551 - causal=True,
567 + causal=dynamic_causal is None,
568 sliding_window_size=sliding_window if sliding_window is not None else -1,
569 isa=isa,
570 enable_kv_split=True,
571 + dynamic_causal=dynamic_causal_tensor,
572 )
573
574 out_with_split = torch.empty_like(query)
@@ -563,13 +580,14 @@ def varlen_with_paged_kv(
580 query_start_loc=cu_query_lens,
581 seq_lens=kv_lens_tensor,
582 scale=scale,
566 - causal=True,
583 + causal=dynamic_causal is None,
584 alibi_slopes=alibi_slopes,
568 - sliding_window=window_size,
585 + sliding_window=sliding_window if sliding_window is not None else -1,
586 block_table=block_tables,
587 softcap=soft_cap if soft_cap is not None else 0,
588 scheduler_metadata=metadata,
589 s_aux=s_aux,
590 + dynamic_causal=dynamic_causal_tensor,
591 **fp8_kwargs,
592 )
593
@@ -597,13 +615,14 @@ def varlen_with_paged_kv(
615 query_start_loc=cu_query_lens,
616 seq_lens=kv_lens_tensor,
617 scale=scale,
600 - causal=True,
618 + causal=dynamic_causal is None,
619 alibi_slopes=alibi_slopes,
602 - sliding_window=window_size,
620 + sliding_window=sliding_window if sliding_window is not None else -1,
621 block_table=block_tables,
622 softcap=soft_cap if soft_cap is not None else 0,
623 scheduler_metadata=metadata,
624 s_aux=s_aux,
625 + dynamic_causal=dynamic_causal_tensor,
626 )
627 atol = _FP8_ATOL[kv_cache_dtype]
628 rtol = _FP8_RTOL
@@ -620,6 +639,7 @@ def varlen_with_paged_kv(
639 soft_cap=soft_cap,
640 alibi_slopes=alibi_slopes,
641 s_aux=s_aux,
642 + dynamic_causal=dynamic_causal,
643 )
644 atol, rtol = 1.5e-2, 1e-2
645
@@ -1035,3 +1055,58 @@ def test_varlen_with_paged_kv_sink(
1055 isa=isa,
1056 kv_cache_dtype=kv_cache_dtype,
1057 )
1058 +
1059 +
1060 +@pytest.mark.parametrize(
1061 + "kv_cache_dtype",
1062 + [
1063 + "auto",
1064 + ],
1065 +)
1066 +@pytest.mark.parametrize("seq_lens", SEQ_LENS)
1067 +@pytest.mark.parametrize("num_heads", NUM_HEADS)
1068 +@pytest.mark.parametrize(
1069 + "head_size",
1070 + [
1071 + 128,
1072 + ],
1073 +)
1074 +@pytest.mark.parametrize("block_size", [96, 128])
1075 +@pytest.mark.parametrize("sliding_window", SLIDING_WINDOWS)
1076 +@pytest.mark.parametrize("dtype", [torch.bfloat16])
1077 +@pytest.mark.parametrize("soft_cap", [None])
1078 +@pytest.mark.parametrize("num_blocks", NUM_BLOCKS)
1079 +@pytest.mark.parametrize("use_alibi", [False])
1080 +@pytest.mark.parametrize("use_sink", [False])
1081 +@pytest.mark.parametrize("isa", ["amx"])
1082 +@pytest.mark.skipif(not torch.cpu._is_amx_tile_supported(), reason="no AMX support.")
1083 +def test_varlen_with_paged_kv_dynamic_causal(
1084 + seq_lens: list[tuple[int, int]],
1085 + num_heads: tuple[int, int],
1086 + head_size: int,
1087 + sliding_window: int | None,
1088 + dtype: torch.dtype,
1089 + block_size: int,
1090 + soft_cap: float | None,
1091 + num_blocks: int,
1092 + use_alibi: bool,
1093 + use_sink: bool,
1094 + isa: str,
1095 + kv_cache_dtype: str,
1096 +) -> None:
1097 + dynamic_causal = [bool(i % 2) for i in range(len(seq_lens))]
1098 + varlen_with_paged_kv(
1099 + seq_lens=seq_lens,
1100 + num_heads=num_heads,
1101 + head_size=head_size,
1102 + sliding_window=sliding_window,
1103 + dtype=dtype,
1104 + block_size=block_size,
1105 + soft_cap=soft_cap,
1106 + num_blocks=num_blocks,
1107 + use_alibi=use_alibi,
1108 + use_sink=use_sink,
1109 + isa=isa,
1110 + kv_cache_dtype=kv_cache_dtype,
1111 + dynamic_causal=dynamic_causal,
1112 + )
vllm/_custom_ops.py
+6 -3
@@ -3619,6 +3619,7 @@ def cpu_attn_get_scheduler_metadata(
3619 sliding_window_size: int,
3620 isa: str,
3621 enable_kv_split: bool,
3622 + dynamic_causal: torch.Tensor | None = None,
3623 ) -> torch.Tensor:
3624 scheduler_metadata = torch.ops._C.get_scheduler_metadata(
3625 num_reqs,
@@ -3632,6 +3633,7 @@ def cpu_attn_get_scheduler_metadata(
3633 sliding_window_size,
3634 isa,
3635 enable_kv_split,
3636 + dynamic_causal,
3637 )
3638 return scheduler_metadata
3639
@@ -3670,11 +3672,12 @@ def cpu_attention_with_kv_cache(
3672 scale: float,
3673 causal: bool,
3674 alibi_slopes: torch.Tensor | None,
3673 - sliding_window: tuple[int, int],
3675 + sliding_window: int,
3676 block_table: torch.Tensor,
3677 softcap: float,
3678 scheduler_metadata: torch.Tensor,
3679 s_aux: torch.Tensor | None,
3680 + dynamic_causal: torch.Tensor | None = None,
3681 k_scale: float = 1.0,
3682 v_scale: float = 1.0,
3683 kv_cache_dtype: str = "auto",
@@ -3689,12 +3692,12 @@ def cpu_attention_with_kv_cache(
3692 scale,
3693 causal,
3694 alibi_slopes,
3692 - sliding_window[0],
3693 - sliding_window[1],
3695 + sliding_window,
3696 block_table,
3697 softcap,
3698 scheduler_metadata,
3699 s_aux,
3700 + dynamic_causal,
3701 k_scale,
3702 v_scale,
3703 kv_cache_dtype,
vllm/v1/attention/backends/cpu_attn.py
+16 -5
@@ -112,6 +112,7 @@ class CPUAttentionMetadata:
112 slot_mapping: torch.Tensor
113 scheduler_metadata: torch.Tensor | None
114 causal: bool = True
115 + dynamic_causal: torch.Tensor | None = None
116
117 # can be removed after deprecate sdpa
118 use_sdpa_prefill: bool = False
@@ -172,7 +173,16 @@ class CPUAttentionMetadataBuilder(AttentionMetadataBuilder[CPUAttentionMetadata]
173 seq_lens = common_attn_metadata.seq_lens
174 block_table_tensor = common_attn_metadata.block_table_tensor
175 slot_mapping = common_attn_metadata.slot_mapping
175 - causal = False if self.is_cross_attention else common_attn_metadata.causal
176 + is_dynamic_casual = isinstance(common_attn_metadata.causal, torch.Tensor)
177 + dynamic_casual = None
178 + if is_dynamic_casual:
179 + dynamic_casual = common_attn_metadata.causal
180 +
181 + causal = (
182 + False
183 + if self.is_cross_attention or is_dynamic_casual
184 + else common_attn_metadata.causal
185 + )
186
187 encoder_cache_tensor = None
188 if self.is_encoder_only_attention:
@@ -215,6 +225,7 @@ class CPUAttentionMetadataBuilder(AttentionMetadataBuilder[CPUAttentionMetadata]
225 sliding_window_size=self.window_size,
226 isa=self.isa,
227 enable_kv_split=envs.VLLM_CPU_ATTN_SPLIT_KV,
228 + dynamic_causal=dynamic_casual,
229 )
230
231 attn_metadata = CPUAttentionMetadata(
@@ -228,6 +239,7 @@ class CPUAttentionMetadataBuilder(AttentionMetadataBuilder[CPUAttentionMetadata]
239 scheduler_metadata=scheduler_metadata,
240 causal=causal,
241 encoder_cache=encoder_cache_tensor,
242 + dynamic_causal=dynamic_casual,
243 )
244
245 return attn_metadata
@@ -269,11 +281,9 @@ class CPUAttentionBackendImpl(AttentionImpl):
281 alibi_slopes = torch.tensor(alibi_slopes, dtype=torch.float32)
282 self.alibi_slopes = alibi_slopes
283 if sliding_window is None:
272 - self.sliding_window = (-1, -1)
273 - elif attn_type == AttentionType.ENCODER_ONLY:
274 - self.sliding_window = (sliding_window - 1, sliding_window - 1)
284 + self.sliding_window = -1
285 else:
276 - self.sliding_window = (sliding_window - 1, 0)
286 + self.sliding_window = sliding_window
287 self.kv_cache_dtype = kv_cache_dtype
288 self.num_queries_per_kv = self.num_heads // self.num_kv_heads
289
@@ -378,6 +388,7 @@ class CPUAttentionBackendImpl(AttentionImpl):
388 softcap=self.logits_soft_cap,
389 scheduler_metadata=attn_metadata.scheduler_metadata,
390 s_aux=self.sinks,
391 + dynamic_causal=attn_metadata.dynamic_causal,
392 k_scale=layer._k_scale_float,
393 v_scale=layer._v_scale_float,
394 kv_cache_dtype=self.kv_cache_dtype,