diff --git a/csrc/cpu/cpu_attn.cpp b/csrc/cpu/cpu_attn.cpp index 13e430fb41..afdc0e0eca 100644 --- a/csrc/cpu/cpu_attn.cpp +++ b/csrc/cpu/cpu_attn.cpp @@ -78,8 +78,9 @@ torch::Tensor get_scheduler_metadata( input.causal = causal; input.isa = isa; input.enable_kv_split = enable_kv_split; - input.dynamic_causal = - dynamic_causal.has_value() ? dynamic_causal->data_ptr() : nullptr; + input.dynamic_causal = dynamic_causal.has_value() + ? dynamic_causal->data_ptr() + : nullptr; const int64_t kv_cache_idx = static_cast(parse_fp8_kv_dtype(kv_cache_dtype)); @@ -256,8 +257,9 @@ void cpu_attention_with_kv_cache( input.s_aux = nullptr; input.s_aux_is_bf16 = false; } - input.dynamic_causal = - dynamic_causal.has_value() ? dynamic_causal->data_ptr() : nullptr; + input.dynamic_causal = dynamic_causal.has_value() + ? dynamic_causal->data_ptr() + : nullptr; input.scale = scale; input.causal = causal; input.sliding_window_size = sliding_window; diff --git a/csrc/cpu/cpu_attn_impl.hpp b/csrc/cpu/cpu_attn_impl.hpp index d914e15063..0666ee413f 100644 --- a/csrc/cpu/cpu_attn_impl.hpp +++ b/csrc/cpu/cpu_attn_impl.hpp @@ -397,7 +397,7 @@ class AttentionScheduler { int32_t max_num_q_per_iter; // max Q head num can be hold in registers int32_t kv_block_alignment; // context length alignment requirement bool enable_kv_split; - bool* dynamic_causal; + int32_t* dynamic_causal; }; static constexpr int32_t MaxQTileIterNum = 128; @@ -450,7 +450,7 @@ class AttentionScheduler { const int32_t q_token_num = input.query_start_loc[req_id + 1] - input.query_start_loc[req_id]; const bool req_causal = - is_dynamic_causal ? input.dynamic_causal[req_id] : causal; + is_dynamic_causal ? input.dynamic_causal[req_id] != 0 : causal; const int32_t q_start_pos = seq_len - q_token_num; const int32_t kv_start_pos = 0; const int32_t kv_end_pos = seq_len; @@ -492,7 +492,7 @@ class AttentionScheduler { const int32_t q_token_num = input.query_start_loc[req_id + 1] - input.query_start_loc[req_id]; const bool req_causal = - is_dynamic_causal ? input.dynamic_causal[req_id] : causal; + is_dynamic_causal ? input.dynamic_causal[req_id] != 0 : causal; const int32_t q_start_pos = seq_len - q_token_num; const int32_t kv_start_pos = 0; const int32_t kv_end_pos = seq_len; @@ -884,7 +884,7 @@ struct AttentionInput { // provided as fp32 and executed in full float precision. const void* s_aux; bool s_aux_is_bf16; - bool* dynamic_causal; + int32_t* dynamic_causal; float scale; bool causal; int32_t sliding_window_size; @@ -1571,7 +1571,7 @@ class AttentionMainLoop { const float* alibi_slopes = input->alibi_slopes; const void* s_aux = input->s_aux; const bool s_aux_is_bf16 = input->s_aux_is_bf16; - const bool* dynamic_causal = input->dynamic_causal; + const int32_t* dynamic_causal = input->dynamic_causal; const bool is_dynamic_causal = dynamic_causal != nullptr; const bool causal = input->causal; @@ -1658,7 +1658,8 @@ class AttentionMainLoop { const int32_t current_group_idx = current_workitem_group->req_id; const int32_t current_group_causal = - is_dynamic_causal ? dynamic_causal[current_group_idx] : causal; + is_dynamic_causal ? dynamic_causal[current_group_idx] != 0 + : causal; auto [sliding_window_left, sliding_window_right] = AttentionScheduler::calcu_sliding_window_size( sliding_window_size, current_group_causal); diff --git a/tests/kernels/attention/test_cpu_attn.py b/tests/kernels/attention/test_cpu_attn.py index f79b1b15d2..7f88f76d0b 100644 --- a/tests/kernels/attention/test_cpu_attn.py +++ b/tests/kernels/attention/test_cpu_attn.py @@ -472,7 +472,7 @@ def varlen_with_paged_kv( scale = head_size**-0.5 token_num = sum(query_lens) dynamic_causal_tensor = ( - torch.tensor(dynamic_causal, dtype=torch.bool) + torch.tensor(dynamic_causal, dtype=torch.int32) if dynamic_causal is not None else None ) diff --git a/tests/v1/sample/test_diffusion_gemma_reads.py b/tests/v1/sample/test_diffusion_gemma_reads.py index 7bd04c0379..4abadc254d 100644 --- a/tests/v1/sample/test_diffusion_gemma_reads.py +++ b/tests/v1/sample/test_diffusion_gemma_reads.py @@ -227,6 +227,10 @@ def test_single_step_tile_skips_self_conditioning(): assert not states.self_conditioning_embeds[0].any() +def test_sample_step_is_compiled(): + assert hasattr(_compiled_sample_step, "_torchdynamo_orig_callable") + + @pytest.mark.parametrize("stance", ["default", "force_eager"]) def test_self_conditioning_stores_a_bf16_model_in_the_fp32_buffer(stance): # The model's embeddings are bf16 while the buffer is fp32. Compiled code diff --git a/vllm/model_executor/models/diffusion_gemma.py b/vllm/model_executor/models/diffusion_gemma.py index 404c154b54..1b27df68ab 100644 --- a/vllm/model_executor/models/diffusion_gemma.py +++ b/vllm/model_executor/models/diffusion_gemma.py @@ -445,11 +445,6 @@ def _mask_rows_to_allowed( return logits if out is None else out -# Tiles specialize the graph on batch size 1, on a tile as wide as the canvas, -# on compute_sc and on sizes that coincide with the state buffers. That set -# is small but passes Dynamo's default of 8, after which every step would run -# eager: about twice as slow for a self-conditioned step. -@torch._dynamo.config.patch(recompile_limit=64) @torch.compile(dynamic=True, backend=current_platform.simple_compile_backend) def _denoise_temperature( step_tensor: torch.Tensor, @@ -464,6 +459,12 @@ def _denoise_temperature( return t_min + (t_max - t_min) * (remaining / max_denoising_steps) +# Tiles specialize the graph on batch size 1, on a tile as wide as the canvas, +# on compute_sc and on sizes that coincide with the state buffers. That set +# is small but passes Dynamo's default of 8, after which every step would run +# eager: about twice as slow for a self-conditioned step. +@torch._dynamo.config.patch(recompile_limit=64) +@torch.compile(dynamic=True, backend=current_platform.simple_compile_backend) def _compiled_sample_step( # Per-position statistics of the temperature-scaled logits, from # sample_row_stats: [num_decode, CL] each, and the softmax