Commit b2eb84de authored by Isotr0py's avatar Isotr0py Committed by Kevin H. Luu
Browse files

[Bugfix] Remove `tile_size=64` for mm_prefix triton attention (#30973)


Signed-off-by: default avatarIsotr0py <mozf@mail2.sysu.edu.cn>
(cherry picked from commit d2dc5dfc)
parent ac43367c
......@@ -800,7 +800,6 @@ def _get_tile_size(
head_size: int,
sliding_window: int,
element_size: int,
is_mm_prefix: bool,
is_prefill: bool,
) -> int:
"""Select tile size with Gemma3-specific optimization.
......@@ -809,10 +808,6 @@ def _get_tile_size(
the larger head dimension (128/256). For other models, use
the default vLLM behavior.
"""
if is_mm_prefix:
# Multimodal bidirectional attention needs a larger tile size
return 64
if _is_gemma3_attention(head_size, sliding_window):
# Gemma3: use 32 for decode (default is 16)
return 32
......@@ -903,14 +898,12 @@ def unified_attention(
head_size,
sliding_window_val,
q.element_size(),
is_mm_prefix=use_mm_prefix,
is_prefill=True,
)
TILE_SIZE_DECODE = _get_tile_size(
head_size,
sliding_window_val,
q.element_size(),
is_mm_prefix=use_mm_prefix,
is_prefill=False,
)
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment