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

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,