[Bugfix][Gemma4] Fix FA4 mm_prefix mask: add sliding window and absolute q_idx (#47332)
Signed-off-by: Luciano Martins <lucianommartins@users.noreply.github.com> Co-authored-by: Luciano Martins <lucianommartins@users.noreply.github.com>
Luciano Martins committed
Jul 4, 2026 at 21:46 UTC
34b560b725b401f33477464163244b4744d1c40f
1 file changed
+90
-50
vllm/v1/attention/backends/flash_attn.py
+90
-50
@@ -892,21 +892,28 @@ class FlashAttentionImpl(AttentionImpl):
892
and self.vllm_flash_attn_version == 4
893
):
894
max_ranges = mm_prefix_ranges.shape[1]
895
- # Gemma4: clamp the bidirectional block to the sliding
896
- # window (HF (causal OR blockwise) AND sliding_window on
897
- # local layers). Other mm-prefix models pass sliding_window=0
898
- # -> unclamped, behavior unchanged. sliding_window_size[0]+1
899
- # is the window (self.sliding_window = (sw-1, 0) on sliding
900
- # layers); causal mm_prefix never hits the symmetric path.
895
+ # Sliding window value in Triton convention
896
+ # (1 + window_size[0]). Global-attention layers
897
+ # store (-1, -1) → sw stays None / 0.
898
+ sw_val = (
899
+ 1 + sliding_window_size[0]
900
+ if sliding_window_size is not None
901
+ and sliding_window_size[0] >= 0
902
+ else None
903
+ )
904
+ # Gemma4: also clamp the bidirectional block to the
905
+ # sliding window when the layer opts in
906
+ # (mm_prefix_clamp_sliding_window flag from PR #47217).
907
mm_clamp_sw = 0
908
if (
909
getattr(layer, "mm_prefix_clamp_sliding_window", False)
904
- and sliding_window_size is not None
905
- and sliding_window_size[0] >= 0
910
+ and sw_val is not None
911
):
907
- mm_clamp_sw = sliding_window_size[0] + 1
912
+ mm_clamp_sw = sw_val
913
mm_mask_mod = _make_mm_prefix_mask_mod(
909
- max_ranges, sliding_window=mm_clamp_sw
914
+ max_ranges,
915
+ sliding_window=mm_clamp_sw,
916
+ sliding_window_left=sw_val,
917
)
918
mm_aux = [mm_prefix_ranges]
919
@@ -1204,23 +1211,26 @@ class FlashAttentionImpl(AttentionImpl):
1211
return output
1212
1213
1207
-def _make_mm_prefix_mask_mod(max_ranges: int, sliding_window: int = 0):
1208
- """Build a CuTE-DSL mask_mod implementing (causal OR mm_prefix).
1209
-
1210
- Returns a @cute.jit callable that evaluates:
1211
- keep = (kv_idx <= q_idx) OR
1212
- (q_idx in [r_start,r_end] AND kv_idx in [r_start,r_end]
1213
- [AND (q_idx - kv_idx) < sliding_window])
1214
- for each mm_prefix range stored in aux_tensors[0].
1215
-
1216
- When sliding_window > 0 (Gemma4 sliding layers), the bidirectional block is
1217
- clamped to the sliding window, matching HF's
1218
- (causal OR blockwise) AND sliding_window (== sliding_window_overlay
1219
- kv > q - sw; future kv passes since q - kv < 0). sliding_window <= 0
1220
- (default) leaves the block unclamped, so other mm-prefix models are
1221
- unchanged. q_idx/kv_idx are local offsets, but their difference is
1222
- frame-invariant and image bidirectional attention only matters during
1223
- prefill (where local == absolute).
1214
+def _make_mm_prefix_mask_mod(
1215
+ max_ranges: int,
1216
+ sliding_window: int = 0,
1217
+ sliding_window_left: int | None = None,
1218
+):
1219
+ """Build a CuTE-DSL mask_mod implementing
1220
+ ``(causal AND sliding_window) OR mm_prefix``.
1221
+
1222
+ The FA4 kernel passes *local* ``q_idx`` (0-based within the current
1223
+ prefill chunk) while ``kv_idx`` is absolute (0-based over the full
1224
+ KV cache). We recover the absolute Q position via
1225
+ ``q_abs = q_idx + seqlen_k - seqlen_q`` (the context-length offset)
1226
+ so that causal, sliding-window, and mm_prefix range comparisons all
1227
+ use consistent absolute positions. This matches the Triton
1228
+ reference path (``compute_kv_seq_mask``).
1229
+
1230
+ ``sliding_window_left`` enforces the sliding window on the causal
1231
+ term (None = full causal, no window). ``sliding_window`` clamps the
1232
+ bidirectional block to the window (0 = unclamped; >0 = Gemma4 local
1233
+ layers via ``mm_prefix_clamp_sliding_window``).
1234
"""
1235
import cutlass
1236
import cutlass.cute as cute
@@ -1230,29 +1240,59 @@ def _make_mm_prefix_mask_mod(max_ranges: int, sliding_window: int = 0):
1240
scalar_to_ssa,
1241
)
1242
1233
- @cute.jit
1234
- def mm_prefix_mask_mod(
1235
- batch_idx: cute.TensorSSA,
1236
- head_idx: cute.TensorSSA,
1237
- q_idx: cute.TensorSSA,
1238
- kv_idx: cute.TensorSSA,
1239
- seqlen_info,
1240
- aux_tensors,
1241
- ):
1242
- keep = kv_idx <= q_idx
1243
- ranges = aux_tensors[0]
1244
- b = batch_idx[0]
1245
- for i in cutlass.range_constexpr(max_ranges): # type: ignore[attr-defined]
1246
- r_start = scalar_to_ssa(ranges[b, i, 0], Int32)
1247
- r_end = scalar_to_ssa(ranges[b, i, 1], Int32)
1248
- valid = r_start < r_end
1249
- q_in = (q_idx >= r_start) & (q_idx <= r_end) & valid
1250
- k_in = (kv_idx >= r_start) & (kv_idx <= r_end) & valid
1251
- mm = q_in & k_in
1252
- if sliding_window > 0:
1253
- mm = mm & ((q_idx - kv_idx) < sliding_window)
1254
- keep = keep | mm
1255
- return keep
1243
+ if sliding_window_left is not None:
1244
+
1245
+ @cute.jit
1246
+ def mm_prefix_mask_mod(
1247
+ batch_idx: cute.TensorSSA,
1248
+ head_idx: cute.TensorSSA,
1249
+ q_idx: cute.TensorSSA,
1250
+ kv_idx: cute.TensorSSA,
1251
+ seqlen_info,
1252
+ aux_tensors,
1253
+ ):
1254
+ ctx_off = scalar_to_ssa(seqlen_info.seqlen_k - seqlen_info.seqlen_q, Int32)
1255
+ q_abs = q_idx + ctx_off
1256
+ sw = scalar_to_ssa(Int32(sliding_window_left), Int32)
1257
+ keep = (kv_idx <= q_abs) & ((q_abs - kv_idx) < sw)
1258
+ ranges = aux_tensors[0]
1259
+ b = batch_idx[0]
1260
+ for i in cutlass.range_constexpr(max_ranges): # type: ignore[attr-defined]
1261
+ r_start = scalar_to_ssa(ranges[b, i, 0], Int32)
1262
+ r_end = scalar_to_ssa(ranges[b, i, 1], Int32)
1263
+ valid = r_start < r_end
1264
+ q_in = (q_abs >= r_start) & (q_abs <= r_end) & valid
1265
+ k_in = (kv_idx >= r_start) & (kv_idx <= r_end) & valid
1266
+ mm = q_in & k_in
1267
+ if sliding_window > 0:
1268
+ mm = mm & ((q_abs - kv_idx) < sw)
1269
+ keep = keep | mm
1270
+ return keep
1271
+
1272
+ else:
1273
+
1274
+ @cute.jit
1275
+ def mm_prefix_mask_mod(
1276
+ batch_idx: cute.TensorSSA,
1277
+ head_idx: cute.TensorSSA,
1278
+ q_idx: cute.TensorSSA,
1279
+ kv_idx: cute.TensorSSA,
1280
+ seqlen_info,
1281
+ aux_tensors,
1282
+ ):
1283
+ ctx_off = scalar_to_ssa(seqlen_info.seqlen_k - seqlen_info.seqlen_q, Int32)
1284
+ q_abs = q_idx + ctx_off
1285
+ keep = kv_idx <= q_abs
1286
+ ranges = aux_tensors[0]
1287
+ b = batch_idx[0]
1288
+ for i in cutlass.range_constexpr(max_ranges): # type: ignore[attr-defined]
1289
+ r_start = scalar_to_ssa(ranges[b, i, 0], Int32)
1290
+ r_end = scalar_to_ssa(ranges[b, i, 1], Int32)
1291
+ valid = r_start < r_end
1292
+ q_in = (q_abs >= r_start) & (q_abs <= r_end) & valid
1293
+ k_in = (kv_idx >= r_start) & (kv_idx <= r_end) & valid
1294
+ keep = keep | (q_in & k_in)
1295
+ return keep
1296
1297
mm_prefix_mask_mod.use_fast_sampling = True
1298
return mm_prefix_mask_mod