[Bugfix][DiffusionGemma] Read the causal mask as int32 and compile the sample step (#59992)

Signed-off-by: ojeda-e <eojeda@league.com>

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<bool>() : nullptr;
+  input.dynamic_causal = dynamic_causal.has_value()
+                             ? dynamic_causal->data_ptr<int32_t>()
+                             : nullptr;
 
   const int64_t kv_cache_idx =
       static_cast<int64_t>(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<bool>() : nullptr;
+  input.dynamic_causal = dynamic_causal.has_value()
+                             ? dynamic_causal->data_ptr<int32_t>()
+                             : nullptr;
   input.scale = scale;
   input.causal = causal;
   input.sliding_window_size = sliding_window;

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);

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
     )

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

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