diff --git a/vllm/v1/attention/ops/triton_merge_attn_states.py b/vllm/v1/attention/ops/triton_merge_attn_states.py
index ca06c2970b..e8f90efce9 100644
--- a/vllm/v1/attention/ops/triton_merge_attn_states.py
+++ b/vllm/v1/attention/ops/triton_merge_attn_states.py
@@ -158,10 +158,10 @@ def merge_attn_states(
         prefix_head_stride,
         output_head_stride,
         output_scale,
+        prefill_tokens_with_context,
         head_size,
         padded_head_size,
         output_lse is not None,
-        prefill_tokens_with_context,
         output_scale is not None,
     )
 
@@ -177,10 +177,10 @@ def merge_attn_states_kernel(
     prefix_head_stride,
     output_head_stride,
     output_scale,  # scale tensor or None
+    prefill_tokens_with_context,
     HEAD_SIZE: tl.constexpr,
     PADDED_HEAD_SIZE: tl.constexpr,
     OUTPUT_LSE: tl.constexpr,
-    prefill_tokens_with_context: tl.constexpr,
     USE_FP8: tl.constexpr,
     FP8_MIN: tl.constexpr = float8_info.min,
     FP8_MAX: tl.constexpr = float8_info.max,