chenchenghao/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 11 Packages Projects Releases Wiki Activity

feature: support bidirectional attention for Gemma-3 (#10707)

Browse Source
This commit is contained in:
brimon
2026-02-09 23:17:45 +08:00
committed by GitHub
parent 4f7da5ad0f
commit ddbcfbaaab
5 changed files with 115 additions and 120 deletions
Download Patch File Download Diff File Expand all files Collapse all files

20
docs/supported_models/text_generation/multimodal_language_models.md
View File

@@ -114,3 +114,23 @@ Use this flag when you have sufficient GPU memory and want to minimize latency f
- **Use `--mm-process-config '{"image":{"max_pixels":1048576},"video":{"fps":3,"max_pixels":602112,"max_frames":60}}'`**: To set `image`, `video`, and `audio` input limits.
This can reduce GPU memory usage, improve inference speed, and help to avoid OOM, but may impact model performance, thus set a proper value based on your specific use case. Currently, only `qwen_vl` supports this config. Please refer to [qwen_vl processor](https://github.com/sgl-project/sglang/blob/main/python/sglang/srt/multimodal/processors/qwen_vl.py) for understanding the meaning of each parameter.
### Bidirectional Attention in Multimodal Model Serving
**Note for serving the Gemma-3 multimodal model**:
As mentioned in [Welcome Gemma 3: Google's all new multimodal, multilingual, long context open LLM
](https://huggingface.co/blog/gemma3#multimodality), Gemma-3 employs bidirectional attention between image tokens during the prefill phase. Currently, SGLang only supports bidirectional attention when using the Triton Attention Backend. Note, however, that SGLang's current bidirectional attention implementation is incompatible with both CUDA Graph and Chunked Prefill.
To enable bidirectional attention, you can use the `TritonAttnBackend` while disabling CUDA Graph and Chunked Prefill. Example launch command:
```shell
python -m sglang.launch_server \
--model-path google/gemma-3-4b-it \
--host 0.0.0.0 --port 30000 \
--enable-multimodal \
--dtype bfloat16 --triton-attention-reduce-in-fp32 \
--attention-backend triton \ # Use Triton attention backend
--disable-cuda-graph \ # Disable Cuda Graph
--chunked-prefill-size -1 # Disable Chunked Prefill
```
If higher serving performance is required and a certain degree of accuracy loss is acceptable, you may choose to use other attention backends, and you can also enable features like CUDA Graph and Chunked Prefill for better performance, but note that the model will fall back to using causal attention instead of bidirectional attention.

39
python/sglang/srt/layers/attention/triton_backend.py
View File

@@ -110,6 +110,11 @@ class TritonAttnBackend(AttentionBackend):
)
self.max_kv_splits = model_runner.server_args.triton_attention_num_kv_splits
self.allow_bidirectional_attention_in_extend = (
model_runner.server_args.disable_cuda_graph
and (model_runner.server_args.chunked_prefill_size == -1)
)
# Decide whether enable deterministic inference with batch-invariant operations
self.enable_deterministic = (
model_runner.server_args.enable_deterministic_inference
@@ -389,17 +394,20 @@ class TritonAttnBackend(AttentionBackend):
)
# Sliding window
if self.sliding_window_size is not None and self.sliding_window_size > 0:
window_kv_indptr, window_kv_indices, _, _ = (
update_sliding_window_buffer(
self.window_kv_indptr,
self.req_to_token,
self.sliding_window_size,
forward_batch.extend_prefix_lens,
forward_batch.req_pool_indices,
bs,
self.device,
self.token_to_kv_pool_allocator,
)
(
window_kv_indptr,
window_kv_indices,
window_kv_lens,
window_kv_offsets,
) = update_sliding_window_buffer(
self.window_kv_indptr,
self.req_to_token,
self.sliding_window_size,
forward_batch.extend_prefix_lens,
forward_batch.req_pool_indices,
bs,
self.device,
self.token_to_kv_pool_allocator,
)
qo_indptr = self.qo_indptr
@@ -812,7 +820,14 @@ class TritonAttnBackend(AttentionBackend):
logits_soft_cap = logit_capping_mod(layer.logit_capping_method, layer.logit_cap)
causal = True
if layer.is_cross_attention or layer.attn_type == AttentionType.ENCODER_ONLY:
if (
layer.is_cross_attention
or layer.attn_type == AttentionType.ENCODER_ONLY
or (
layer.attn_type == AttentionType.DECODER_BIDIRECTIONAL
and self.allow_bidirectional_attention_in_extend
)
):
causal = False
# Deterministic mode: use unified 1-stage kernel

2
python/sglang/srt/layers/radix_attention.py
View File

@@ -37,6 +37,8 @@ class AttentionType(Enum):
# Decoder attention between previous layer Q/K/V
DECODER = "decoder"
# Decoder bidirectional attention between image tokens
DECODER_BIDIRECTIONAL = "decoder_bidirectional"
# Encoder attention between previous layer Q/K/V
ENCODER_ONLY = "encoder_only"

50
python/sglang/srt/models/gemma3_causal.py
View File

@@ -16,7 +16,6 @@ from typing import Iterable, Optional, Set, Tuple
import einops
import torch
import torch.nn.functional as F
from torch import nn
from transformers import (
ROPE_INIT_FUNCTIONS,
@@ -35,7 +34,7 @@ from sglang.srt.layers.linear import (
)
from sglang.srt.layers.logits_processor import LogitsProcessor
from sglang.srt.layers.quantization.base_config import QuantizationConfig
from sglang.srt.layers.radix_attention import RadixAttention
from sglang.srt.layers.radix_attention import AttentionType, RadixAttention
from sglang.srt.layers.rotary_embedding import apply_rotary_pos_emb
from sglang.srt.layers.vocab_parallel_embedding import ParallelLMHead
from sglang.srt.model_executor.forward_batch_info import ForwardBatch
@@ -193,58 +192,13 @@ class Gemma3Attention(nn.Module):
sliding_window_size=self.sliding_window,
quant_config=quant_config,
prefix=add_prefix("attn", prefix),
attn_type=AttentionType.DECODER_BIDIRECTIONAL,
)
# Gemma3 adds normalization for q and k
self.q_norm = Gemma3RMSNorm(dim=config.head_dim, eps=config.rms_norm_eps)
self.k_norm = Gemma3RMSNorm(dim=config.head_dim, eps=config.rms_norm_eps)
def naive_attn_with_masks(
self,
q: torch.Tensor,
k: torch.Tensor,
v: torch.Tensor,
out: torch.Tensor,
**kwargs,
) -> torch.Tensor:
q = q.view(-1, self.num_heads, self.head_dim)
# Expand the key and value to handle GQA.
num_queries_per_kv = self.num_heads // self.num_kv_heads
k = k.view(-1, self.num_kv_heads, self.head_dim)
k = k.repeat_interleave(num_queries_per_kv, dim=-2)
v = v.view(-1, self.num_kv_heads, self.head_dim)
v = v.repeat_interleave(num_queries_per_kv, dim=-2)
if self.is_sliding:
attn_masks = kwargs["local_attn_masks"]
else:
attn_masks = kwargs["global_attn_masks"]
seq_lens = kwargs["seq_lens"]
start_idx = 0
for seq_len, attn_mask in zip(seq_lens, attn_masks):
end_idx = start_idx + seq_len
query = q[start_idx:end_idx].unsqueeze(0)
key = k[start_idx:end_idx].unsqueeze(0)
value = v[start_idx:end_idx].unsqueeze(0)
# Transpose.
query = query.transpose(1, 2)
key = key.transpose(1, 2)
value = value.transpose(1, 2)
output = F.scaled_dot_product_attention(
query,
key,
value,
attn_mask,
self.scaling,
)
output = output.transpose(1, 2).flatten(-2, -1)
out[start_idx:end_idx] = output
start_idx = end_idx
return out
def forward(
self,
hidden_states: torch.Tensor,

124
python/sglang/srt/models/gemma3_mm.py
View File

@@ -18,12 +18,13 @@
import logging
import re
from functools import lru_cache
from typing import Dict, Iterable, List, Optional, Set, Tuple, TypedDict
from typing import Iterable, List, Optional, Set, Tuple, TypedDict
import torch
from torch import nn
from transformers import Gemma3Config, PreTrainedModel
from sglang.srt.layers.attention.triton_backend import TritonAttnBackend
from sglang.srt.layers.layernorm import Gemma3RMSNorm
from sglang.srt.layers.logits_processor import LogitsProcessor
from sglang.srt.layers.quantization.base_config import QuantizationConfig
@@ -36,7 +37,7 @@ from sglang.srt.managers.schedule_batch import (
MultimodalInputs,
flatten_nested_list,
)
from sglang.srt.model_executor.forward_batch_info import ForwardBatch
from sglang.srt.model_executor.forward_batch_info import ForwardBatch, ForwardMode
from sglang.srt.model_loader.weight_utils import (
default_weight_loader,
maybe_remap_kv_scale_name,
@@ -212,71 +213,62 @@ class Gemma3ForConditionalGeneration(PreTrainedModel):
def prepare_attn_masks(
self,
forward_batch: ForwardBatch,
input_ids: torch.Tensor,
positions: torch.Tensor,
mask_dtype: torch.dtype,
**kwargs,
) -> Dict:
):
"""Prepare attention masks for multimodal inputs."""
kwargs["has_images"] = True
# Distinguish sequences by position id 0
start_indices = (positions == 0).cpu().nonzero()
num_seqs = len(start_indices)
seq_lens = []
for i in range(num_seqs):
start_idx = start_indices[i].item()
if i < num_seqs - 1:
end_idx = start_indices[i + 1].item()
else:
end_idx = len(input_ids)
seq_lens.append(end_idx - start_idx)
kwargs["seq_lens"] = seq_lens
# Create attention masks
global_attn_masks = []
local_attn_masks = []
sliding_window = self.config.text_config.interleaved_sliding_window
start_idx = 0
for seq_len in seq_lens:
end_idx = start_idx + seq_len
input_token_ids = input_ids[start_idx:end_idx]
start_idx = end_idx
# Create global causal mask
global_attn_mask = torch.empty(
1,
1,
seq_len,
seq_len,
dtype=mask_dtype,
device=input_ids.device,
if isinstance(forward_batch.attn_backend, TritonAttnBackend):
assert forward_batch.forward_mode == ForwardMode.EXTEND
bidirectional_attn_masks_list = []
bidirectional_attn_mask_indptr = torch.zeros(
forward_batch.batch_size + 1, dtype=torch.int32, device=input_ids.device
)
global_attn_mask.fill_(float("-inf"))
global_attn_mask = global_attn_mask.triu(diagonal=1)
# Consider bidirectional attention between image tokens
img_mask = torch.zeros_like(global_attn_mask)
img_pos = input_token_ids == self.config.image_token_index
img_mask[:, :, :, img_pos] += 1
img_mask[:, :, img_pos, :] += 1
global_attn_mask = torch.where(img_mask == 2, 0, global_attn_mask)
global_attn_masks.append(global_attn_mask)
for i in range(forward_batch.batch_size):
bidirectional_attn_mask = torch.empty(
forward_batch.extend_seq_lens[i],
forward_batch.extend_seq_lens[i]
+ forward_batch.extend_prefix_lens[i],
dtype=mask_dtype,
device=input_ids.device,
)
bidirectional_attn_mask.fill_(1)
bidirectional_attn_mask = bidirectional_attn_mask.tril(
diagonal=forward_batch.extend_prefix_lens[i]
)
# Create local causal mask with sliding window
local_attn_mask = torch.ones_like(global_attn_mask)
local_attn_mask = torch.tril(local_attn_mask, diagonal=-sliding_window)
local_attn_mask = torch.where(
local_attn_mask == 0, global_attn_mask, float("-inf")
)
local_attn_masks.append(local_attn_mask)
# Consider bidirectional attention between image tokens
mm_inputs = forward_batch.mm_inputs[i]
for mm_item in mm_inputs.mm_items:
if mm_item.is_image():
for im_begin, im_end in mm_item.offsets:
if (
im_begin >= forward_batch.extend_prefix_lens[i]
): # compatible with radix cache
bidirectional_attn_mask[
im_begin
- forward_batch.extend_prefix_lens[i] : im_end
+ 1
- forward_batch.extend_prefix_lens[i],
im_begin : im_end + 1,
] = 1
bidirectional_attn_masks_list.append(bidirectional_attn_mask.flatten())
bidirectional_attn_mask_indptr[i + 1] = (
bidirectional_attn_mask_indptr[i]
+ bidirectional_attn_mask.nelement()
)
kwargs["global_attn_masks"] = global_attn_masks
kwargs["local_attn_masks"] = local_attn_masks
return kwargs
if bidirectional_attn_masks_list:
bidirectional_attn_masks = torch.cat(
bidirectional_attn_masks_list, dim=0
)
forward_batch.attn_backend.forward_metadata.mask_indptr = (
bidirectional_attn_mask_indptr
)
forward_batch.attn_backend.forward_metadata.custom_mask = (
bidirectional_attn_masks
)
def get_input_embeddings(self) -> nn.Embedding:
return self.language_model.get_input_embeddings()
@@ -402,6 +394,18 @@ class Gemma3ForConditionalGeneration(PreTrainedModel):
else:
llm_input_ids = input_ids
# NOTE: As described in https://huggingface.co/blog/gemma3#multimodality, in the prefill stage of Gemma-3, image tokens use bidirectional attention. Currently, only the TritonAttnBackend supports bidirectional attention; other backends have not yet implemented this. Bidirectional attention is incompatible with CUDA Graph and chunked prefill.
if (
forward_batch.forward_mode
== ForwardMode.EXTEND # only Extend mode is supported for now
and forward_batch.contains_image_inputs() # Gemma-3 only supports image as mm inputs
):
self.prepare_attn_masks(
forward_batch,
llm_input_ids,
mask_dtype=torch.bool,
)
hs = general_mm_embed_routine(
input_ids=llm_input_ids,
forward_batch=forward_batch,