[XPU] Enable Gemma 4 E2B / E4B / 31B/ 26B-A4B on Intel XPU (#23280)
Co-authored-by: Claude Opus 4.6 (1M context) <noreply@anthropic.com> Co-authored-by: jmunetong <jmunetong@users.noreply.github.com> Co-authored-by: Meng, Hengyu <hengyu.meng@intel.com> Co-authored-by: ckvermaAI <ckverma@habana.ai>
This commit is contained in:
co-authored by
Claude Opus 4.6
jmunetong
Meng, Hengyu
ckvermaAI
parent
bcf89928b4
commit
2c8357f794
@@ -13,6 +13,7 @@ from sglang.srt.layers.attention.flashattention_backend import (
|
|||||||
prepare_swa_spec_page_table_triton,
|
prepare_swa_spec_page_table_triton,
|
||||||
)
|
)
|
||||||
from sglang.srt.managers.schedule_batch import get_global_server_args
|
from sglang.srt.managers.schedule_batch import get_global_server_args
|
||||||
|
from sglang.srt.mem_cache.swa_memory_pool import SWAKVPool
|
||||||
from sglang.srt.model_executor.forward_batch_info import ForwardBatch, ForwardMode
|
from sglang.srt.model_executor.forward_batch_info import ForwardBatch, ForwardMode
|
||||||
|
|
||||||
if TYPE_CHECKING:
|
if TYPE_CHECKING:
|
||||||
@@ -72,6 +73,12 @@ class XPUAttentionBackend(AttentionBackend):
|
|||||||
self.use_mla = model_runner.model_config.attention_arch == AttentionArch.MLA
|
self.use_mla = model_runner.model_config.attention_arch == AttentionArch.MLA
|
||||||
self.skip_prefill = skip_prefill
|
self.skip_prefill = skip_prefill
|
||||||
self.is_hybrid_swa = model_runner.is_hybrid_swa
|
self.is_hybrid_swa = model_runner.is_hybrid_swa
|
||||||
|
self.use_sliding_window_kv_pool = (
|
||||||
|
isinstance(model_runner.token_to_kv_pool, SWAKVPool)
|
||||||
|
and model_runner.token_to_kv_pool.swa_layer_nums > 0
|
||||||
|
)
|
||||||
|
if self.use_sliding_window_kv_pool:
|
||||||
|
self.token_to_kv_pool = model_runner.token_to_kv_pool
|
||||||
if self.is_hybrid_swa:
|
if self.is_hybrid_swa:
|
||||||
self.full_to_swa_index_mapping = (
|
self.full_to_swa_index_mapping = (
|
||||||
model_runner.token_to_kv_pool.full_to_swa_index_mapping
|
model_runner.token_to_kv_pool.full_to_swa_index_mapping
|
||||||
@@ -193,6 +200,7 @@ class XPUAttentionBackend(AttentionBackend):
|
|||||||
metadata.page_table = self.req_to_token_pool.req_to_token[
|
metadata.page_table = self.req_to_token_pool.req_to_token[
|
||||||
forward_batch.req_pool_indices, : metadata.max_seq_len_k
|
forward_batch.req_pool_indices, : metadata.max_seq_len_k
|
||||||
]
|
]
|
||||||
|
|
||||||
# TODO: we need to test this part for llama 4 eagle case
|
# TODO: we need to test this part for llama 4 eagle case
|
||||||
self._init_local_attn_metadata(forward_batch, metadata, device)
|
self._init_local_attn_metadata(forward_batch, metadata, device)
|
||||||
elif forward_batch.forward_mode.is_target_verify():
|
elif forward_batch.forward_mode.is_target_verify():
|
||||||
@@ -373,6 +381,16 @@ class XPUAttentionBackend(AttentionBackend):
|
|||||||
),
|
),
|
||||||
]
|
]
|
||||||
|
|
||||||
|
# Translate full-pool indices to SWA-pool indices for hybrid models
|
||||||
|
if self.use_sliding_window_kv_pool:
|
||||||
|
# flash_attn_with_kvcache requires int32 page tables; the SWA index
|
||||||
|
# mapping is int64, so cast (matches flashattention_backend.py).
|
||||||
|
metadata.swa_page_table = (
|
||||||
|
self.token_to_kv_pool.translate_loc_from_full_to_swa(
|
||||||
|
metadata.page_table
|
||||||
|
).to(torch.int32)
|
||||||
|
)
|
||||||
|
|
||||||
if self.use_mla:
|
if self.use_mla:
|
||||||
workspace_size = flash_mla_get_workspace_size(
|
workspace_size = flash_mla_get_workspace_size(
|
||||||
max_seq_len=self.max_context_len,
|
max_seq_len=self.max_context_len,
|
||||||
@@ -389,11 +407,27 @@ class XPUAttentionBackend(AttentionBackend):
|
|||||||
workspace_size, device=self.device, dtype=torch.uint8
|
workspace_size, device=self.device, dtype=torch.uint8
|
||||||
)
|
)
|
||||||
|
|
||||||
|
# Translate full-pool indices to SWA-pool indices for hybrid models
|
||||||
|
if self.use_sliding_window_kv_pool:
|
||||||
|
# flash_attn_with_kvcache requires int32 page tables; the SWA index
|
||||||
|
# mapping is int64, so cast (matches flashattention_backend.py).
|
||||||
|
metadata.swa_page_table = (
|
||||||
|
self.token_to_kv_pool.translate_loc_from_full_to_swa(
|
||||||
|
metadata.page_table
|
||||||
|
).to(torch.int32)
|
||||||
|
)
|
||||||
|
|
||||||
# Convert the page table to a strided format which is needed by FA3 API
|
# Convert the page table to a strided format which is needed by FA3 API
|
||||||
if self.page_size > 1:
|
if self.page_size > 1:
|
||||||
self.strided_indices = torch.arange(
|
self.strided_indices = torch.arange(
|
||||||
0, metadata.page_table.shape[1], self.page_size, device=self.device
|
0, metadata.page_table.shape[1], self.page_size, device=self.device
|
||||||
)
|
)
|
||||||
|
|
||||||
|
if self.use_sliding_window_kv_pool and metadata.swa_page_table is not None:
|
||||||
|
metadata.swa_page_table = (
|
||||||
|
metadata.swa_page_table[:, self.strided_indices] // self.page_size
|
||||||
|
)
|
||||||
|
|
||||||
metadata.page_table = (
|
metadata.page_table = (
|
||||||
metadata.page_table[:, self.strided_indices] // self.page_size
|
metadata.page_table[:, self.strided_indices] // self.page_size
|
||||||
)
|
)
|
||||||
@@ -413,8 +447,17 @@ class XPUAttentionBackend(AttentionBackend):
|
|||||||
k_rope: Optional[torch.Tensor] = None,
|
k_rope: Optional[torch.Tensor] = None,
|
||||||
sinks: Optional[torch.Tensor] = None,
|
sinks: Optional[torch.Tensor] = None,
|
||||||
):
|
):
|
||||||
if k is not None:
|
if k is None and v is None:
|
||||||
assert v is not None
|
# Cross-layer KV sharing (Gemma 4): the layer reuses another
|
||||||
|
# layer's KV cache. The paged kernel reads K/V directly via
|
||||||
|
# page_table, and pool.get_kv_buffer(layer.layer_id) routes
|
||||||
|
# to the correct sub-pool because RadixAttention is initialized
|
||||||
|
# with layer_id=kv_shared_layer_index for shared layers. No
|
||||||
|
# materialization needed; just skip the write path.
|
||||||
|
pass
|
||||||
|
elif k is None or v is None:
|
||||||
|
raise ValueError("Both k and v should be None or not None")
|
||||||
|
else:
|
||||||
if save_kv_cache:
|
if save_kv_cache:
|
||||||
cache_loc = (
|
cache_loc = (
|
||||||
forward_batch.out_cache_loc
|
forward_batch.out_cache_loc
|
||||||
@@ -497,6 +540,13 @@ class XPUAttentionBackend(AttentionBackend):
|
|||||||
cu_seqlens_k = swa_spec_metadata.cu_seqlens_k
|
cu_seqlens_k = swa_spec_metadata.cu_seqlens_k
|
||||||
else:
|
else:
|
||||||
page_table = metadata.page_table
|
page_table = metadata.page_table
|
||||||
|
if is_hybrid_swa and self.use_sliding_window_kv_pool:
|
||||||
|
if metadata.swa_page_table is not None:
|
||||||
|
page_table = metadata.swa_page_table
|
||||||
|
else:
|
||||||
|
page_table = self.token_to_kv_pool.translate_loc_from_full_to_swa(
|
||||||
|
metadata.page_table
|
||||||
|
).to(torch.int32)
|
||||||
cu_seqlens_q = metadata.cu_seqlens_q
|
cu_seqlens_q = metadata.cu_seqlens_q
|
||||||
cache_seqlens = metadata.cache_seqlens_int32
|
cache_seqlens = metadata.cache_seqlens_int32
|
||||||
max_seqlen_q = metadata.max_seq_len_q
|
max_seqlen_q = metadata.max_seq_len_q
|
||||||
@@ -525,7 +575,7 @@ class XPUAttentionBackend(AttentionBackend):
|
|||||||
page_table=page_table,
|
page_table=page_table,
|
||||||
cache_seqlens=cache_seqlens,
|
cache_seqlens=cache_seqlens,
|
||||||
cu_seqlens_q=cu_seqlens_q,
|
cu_seqlens_q=cu_seqlens_q,
|
||||||
cu_seqlens_k_new=cu_seqlens_k if not use_local_attn else None,
|
cu_seqlens_k_new=None,
|
||||||
max_seqlen_q=max_seqlen_q,
|
max_seqlen_q=max_seqlen_q,
|
||||||
softmax_scale=layer.scaling,
|
softmax_scale=layer.scaling,
|
||||||
causal=False if use_cascade_attn else causal,
|
causal=False if use_cascade_attn else causal,
|
||||||
@@ -546,7 +596,7 @@ class XPUAttentionBackend(AttentionBackend):
|
|||||||
page_table=self.forward_metadata_spec_decode_expand.page_table,
|
page_table=self.forward_metadata_spec_decode_expand.page_table,
|
||||||
cache_seqlens=self.forward_metadata_spec_decode_expand.cache_seqlens_int32,
|
cache_seqlens=self.forward_metadata_spec_decode_expand.cache_seqlens_int32,
|
||||||
cu_seqlens_q=self.forward_metadata_spec_decode_expand.cu_seqlens_q,
|
cu_seqlens_q=self.forward_metadata_spec_decode_expand.cu_seqlens_q,
|
||||||
cu_seqlens_k_new=self.forward_metadata_spec_decode_expand.cu_seqlens_k,
|
cu_seqlens_k_new=None,
|
||||||
max_seqlen_q=self.forward_metadata_spec_decode_expand.max_seq_len_q,
|
max_seqlen_q=self.forward_metadata_spec_decode_expand.max_seq_len_q,
|
||||||
softmax_scale=layer.scaling,
|
softmax_scale=layer.scaling,
|
||||||
causal=False,
|
causal=False,
|
||||||
@@ -648,7 +698,7 @@ class XPUAttentionBackend(AttentionBackend):
|
|||||||
page_table=page_table,
|
page_table=page_table,
|
||||||
cache_seqlens=cache_seqlens,
|
cache_seqlens=cache_seqlens,
|
||||||
cu_seqlens_q=cu_seqlens_q,
|
cu_seqlens_q=cu_seqlens_q,
|
||||||
cu_seqlens_k_new=cu_seqlens_k if not use_local_attn else None,
|
cu_seqlens_k_new=None,
|
||||||
max_seqlen_q=max_seqlen_q,
|
max_seqlen_q=max_seqlen_q,
|
||||||
softmax_scale=layer.scaling,
|
softmax_scale=layer.scaling,
|
||||||
causal=False if use_cascade_attn else causal,
|
causal=False if use_cascade_attn else causal,
|
||||||
@@ -668,7 +718,7 @@ class XPUAttentionBackend(AttentionBackend):
|
|||||||
page_table=self.forward_metadata_spec_decode_expand.page_table,
|
page_table=self.forward_metadata_spec_decode_expand.page_table,
|
||||||
cache_seqlens=self.forward_metadata_spec_decode_expand.cache_seqlens_int32,
|
cache_seqlens=self.forward_metadata_spec_decode_expand.cache_seqlens_int32,
|
||||||
cu_seqlens_q=self.forward_metadata_spec_decode_expand.cu_seqlens_q,
|
cu_seqlens_q=self.forward_metadata_spec_decode_expand.cu_seqlens_q,
|
||||||
cu_seqlens_k_new=self.forward_metadata_spec_decode_expand.cu_seqlens_k,
|
cu_seqlens_k_new=None,
|
||||||
max_seqlen_q=self.forward_metadata_spec_decode_expand.max_seq_len_q,
|
max_seqlen_q=self.forward_metadata_spec_decode_expand.max_seq_len_q,
|
||||||
softmax_scale=layer.scaling,
|
softmax_scale=layer.scaling,
|
||||||
causal=False,
|
causal=False,
|
||||||
@@ -688,7 +738,8 @@ class XPUAttentionBackend(AttentionBackend):
|
|||||||
else:
|
else:
|
||||||
o = result
|
o = result
|
||||||
|
|
||||||
return o.view(-1, layer.tp_q_head_num * layer.v_head_dim)
|
out = o.view(-1, layer.tp_q_head_num * layer.v_head_dim)
|
||||||
|
return out
|
||||||
|
|
||||||
def forward_decode(
|
def forward_decode(
|
||||||
self,
|
self,
|
||||||
@@ -703,8 +754,12 @@ class XPUAttentionBackend(AttentionBackend):
|
|||||||
k_rope: Optional[torch.Tensor] = None,
|
k_rope: Optional[torch.Tensor] = None,
|
||||||
sinks: Optional[torch.Tensor] = None,
|
sinks: Optional[torch.Tensor] = None,
|
||||||
) -> torch.Tensor:
|
) -> torch.Tensor:
|
||||||
if k is not None:
|
if k is None and v is None:
|
||||||
assert v is not None
|
# Cross-layer KV sharing (Gemma 4): see forward_extend for details.
|
||||||
|
pass
|
||||||
|
elif k is None or v is None:
|
||||||
|
raise ValueError("Both k and v should be None or not None")
|
||||||
|
else:
|
||||||
if save_kv_cache:
|
if save_kv_cache:
|
||||||
cache_loc = (
|
cache_loc = (
|
||||||
forward_batch.out_cache_loc
|
forward_batch.out_cache_loc
|
||||||
@@ -787,7 +842,7 @@ class XPUAttentionBackend(AttentionBackend):
|
|||||||
page_table=metadata.encoder_page_table,
|
page_table=metadata.encoder_page_table,
|
||||||
cache_seqlens=metadata.encoder_lens_int32,
|
cache_seqlens=metadata.encoder_lens_int32,
|
||||||
cu_seqlens_q=metadata.cu_seqlens_q,
|
cu_seqlens_q=metadata.cu_seqlens_q,
|
||||||
cu_seqlens_k_new=metadata.encoder_cu_seqlens_k,
|
cu_seqlens_k_new=None,
|
||||||
max_seqlen_q=1,
|
max_seqlen_q=1,
|
||||||
softmax_scale=layer.scaling,
|
softmax_scale=layer.scaling,
|
||||||
causal=False,
|
causal=False,
|
||||||
@@ -817,7 +872,24 @@ class XPUAttentionBackend(AttentionBackend):
|
|||||||
**kwargs,
|
**kwargs,
|
||||||
)
|
)
|
||||||
else:
|
else:
|
||||||
|
is_swa_layer = (
|
||||||
|
layer.sliding_window_size is not None
|
||||||
|
and layer.sliding_window_size > -1
|
||||||
|
)
|
||||||
|
|
||||||
page_table = metadata.page_table
|
page_table = metadata.page_table
|
||||||
|
# For SWA layers on hybrid models, use the translated
|
||||||
|
# SWA-pool page table so KV reads hit the correct pool.
|
||||||
|
if is_swa_layer and self.use_sliding_window_kv_pool:
|
||||||
|
if metadata.swa_page_table is not None:
|
||||||
|
page_table = metadata.swa_page_table
|
||||||
|
else:
|
||||||
|
page_table = (
|
||||||
|
self.token_to_kv_pool.translate_loc_from_full_to_swa(
|
||||||
|
metadata.page_table
|
||||||
|
).to(torch.int32)
|
||||||
|
)
|
||||||
|
|
||||||
cache_seqlens = metadata.cache_seqlens_int32
|
cache_seqlens = metadata.cache_seqlens_int32
|
||||||
cu_seqlens_k = metadata.cu_seqlens_k
|
cu_seqlens_k = metadata.cu_seqlens_k
|
||||||
max_seqlen_q = metadata.max_seq_len_q
|
max_seqlen_q = metadata.max_seq_len_q
|
||||||
@@ -833,7 +905,7 @@ class XPUAttentionBackend(AttentionBackend):
|
|||||||
page_table=page_table,
|
page_table=page_table,
|
||||||
cache_seqlens=cache_seqlens,
|
cache_seqlens=cache_seqlens,
|
||||||
cu_seqlens_q=metadata.cu_seqlens_q,
|
cu_seqlens_q=metadata.cu_seqlens_q,
|
||||||
cu_seqlens_k_new=cu_seqlens_k,
|
cu_seqlens_k_new=None,
|
||||||
max_seqlen_q=max_seqlen_q,
|
max_seqlen_q=max_seqlen_q,
|
||||||
softmax_scale=layer.scaling,
|
softmax_scale=layer.scaling,
|
||||||
causal=False if use_cascade_attn else causal,
|
causal=False if use_cascade_attn else causal,
|
||||||
@@ -854,7 +926,7 @@ class XPUAttentionBackend(AttentionBackend):
|
|||||||
page_table=self.forward_metadata_spec_decode_expand.page_table,
|
page_table=self.forward_metadata_spec_decode_expand.page_table,
|
||||||
cache_seqlens=self.forward_metadata_spec_decode_expand.cache_seqlens_int32,
|
cache_seqlens=self.forward_metadata_spec_decode_expand.cache_seqlens_int32,
|
||||||
cu_seqlens_q=self.forward_metadata_spec_decode_expand.cu_seqlens_q,
|
cu_seqlens_q=self.forward_metadata_spec_decode_expand.cu_seqlens_q,
|
||||||
cu_seqlens_k_new=self.forward_metadata_spec_decode_expand.cu_seqlens_k,
|
cu_seqlens_k_new=None,
|
||||||
max_seqlen_q=self.forward_metadata_spec_decode_expand.max_seq_len_q,
|
max_seqlen_q=self.forward_metadata_spec_decode_expand.max_seq_len_q,
|
||||||
softmax_scale=layer.scaling,
|
softmax_scale=layer.scaling,
|
||||||
causal=False,
|
causal=False,
|
||||||
@@ -899,7 +971,8 @@ class XPUAttentionBackend(AttentionBackend):
|
|||||||
layer.scaling,
|
layer.scaling,
|
||||||
)
|
)
|
||||||
|
|
||||||
return o.view(-1, layer.tp_q_head_num * layer.v_head_dim)
|
out = o.view(-1, layer.tp_q_head_num * layer.v_head_dim)
|
||||||
|
return out
|
||||||
|
|
||||||
def get_cuda_graph_seq_len_fill_value(self):
|
def get_cuda_graph_seq_len_fill_value(self):
|
||||||
"""Get the fill value for sequence length in CUDA graph."""
|
"""Get the fill value for sequence length in CUDA graph."""
|
||||||
|
|||||||
@@ -215,7 +215,7 @@ def gemma_qkv_rmsnorm(
|
|||||||
|
|
||||||
If k and v are both None (KV-shared layer), only Q is normalized.
|
If k and v are both None (KV-shared layer), only Q is normalized.
|
||||||
"""
|
"""
|
||||||
assert q.is_cuda
|
assert q.is_cuda or q.is_xpu
|
||||||
assert q.stride(-1) == 1, "Q's last dim must be contiguous"
|
assert q.stride(-1) == 1, "Q's last dim must be contiguous"
|
||||||
assert q_weight.shape[-1] == head_dim
|
assert q_weight.shape[-1] == head_dim
|
||||||
M = q.shape[0] if q.dim() >= 2 else 1
|
M = q.shape[0] if q.dim() >= 2 else 1
|
||||||
@@ -223,7 +223,7 @@ def gemma_qkv_rmsnorm(
|
|||||||
|
|
||||||
has_kv = k is not None and v is not None
|
has_kv = k is not None and v is not None
|
||||||
if has_kv:
|
if has_kv:
|
||||||
assert k.is_cuda and v.is_cuda
|
assert (k.is_cuda and v.is_cuda) or (k.is_xpu and v.is_xpu)
|
||||||
assert k.stride(-1) == 1 and v.stride(-1) == 1
|
assert k.stride(-1) == 1 and v.stride(-1) == 1
|
||||||
assert k_weight is not None and k_weight.shape[-1] == head_dim
|
assert k_weight is not None and k_weight.shape[-1] == head_dim
|
||||||
|
|
||||||
@@ -245,6 +245,75 @@ def gemma_qkv_rmsnorm(
|
|||||||
)
|
)
|
||||||
|
|
||||||
|
|
||||||
|
@triton.jit
|
||||||
|
def _gemma_routing_post_topk_kernel(
|
||||||
|
Logits_ptr,
|
||||||
|
Ids_ptr,
|
||||||
|
Scale_ptr,
|
||||||
|
Out_weights_ptr,
|
||||||
|
Out_ids_ptr,
|
||||||
|
stride_l,
|
||||||
|
stride_ow,
|
||||||
|
stride_oi,
|
||||||
|
K: tl.constexpr,
|
||||||
|
BLOCK_K: tl.constexpr,
|
||||||
|
):
|
||||||
|
"""Fused: softmax(topk_logits) * per_expert_scale[topk_ids] → float32 weights, int32 ids.
|
||||||
|
|
||||||
|
One program per token. K is the number of top-k experts (e.g. 8).
|
||||||
|
"""
|
||||||
|
row = tl.program_id(0)
|
||||||
|
cols = tl.arange(0, BLOCK_K)
|
||||||
|
mask = cols < K
|
||||||
|
|
||||||
|
logits = tl.load(
|
||||||
|
Logits_ptr + row * stride_l + cols, mask=mask, other=float("-inf")
|
||||||
|
).to(tl.float32)
|
||||||
|
ids_i64 = tl.load(Ids_ptr + row * stride_l + cols, mask=mask, other=0)
|
||||||
|
|
||||||
|
# Stable softmax
|
||||||
|
max_val = tl.max(logits, axis=0)
|
||||||
|
exp_val = tl.exp(logits - max_val)
|
||||||
|
sum_exp = tl.sum(exp_val, axis=0)
|
||||||
|
weights = exp_val / sum_exp
|
||||||
|
|
||||||
|
# Gather per_expert_scale and multiply
|
||||||
|
scale = tl.load(Scale_ptr + ids_i64, mask=mask, other=1.0).to(tl.float32)
|
||||||
|
weights = weights * scale
|
||||||
|
|
||||||
|
tl.store(Out_weights_ptr + row * stride_ow + cols, weights, mask=mask)
|
||||||
|
tl.store(Out_ids_ptr + row * stride_oi + cols, ids_i64.to(tl.int32), mask=mask)
|
||||||
|
|
||||||
|
|
||||||
|
def gemma_routing_post_topk(
|
||||||
|
topk_logits: torch.Tensor,
|
||||||
|
topk_ids: torch.Tensor,
|
||||||
|
per_expert_scale: torch.Tensor,
|
||||||
|
) -> tuple[torch.Tensor, torch.Tensor]:
|
||||||
|
"""Fused softmax + scale-gather + casts for Gemma4 routing.
|
||||||
|
|
||||||
|
Replaces: softmax(topk_logits) * per_expert_scale[topk_ids] → (f32, i32).
|
||||||
|
"""
|
||||||
|
B, K = topk_logits.shape
|
||||||
|
BLOCK_K = triton.next_power_of_2(K)
|
||||||
|
out_weights = torch.empty((B, K), dtype=torch.float32, device=topk_logits.device)
|
||||||
|
out_ids = torch.empty((B, K), dtype=torch.int32, device=topk_logits.device)
|
||||||
|
|
||||||
|
_gemma_routing_post_topk_kernel[(B,)](
|
||||||
|
topk_logits,
|
||||||
|
topk_ids,
|
||||||
|
per_expert_scale,
|
||||||
|
out_weights,
|
||||||
|
out_ids,
|
||||||
|
topk_logits.stride(0),
|
||||||
|
out_weights.stride(0),
|
||||||
|
out_ids.stride(0),
|
||||||
|
K=K,
|
||||||
|
BLOCK_K=BLOCK_K,
|
||||||
|
)
|
||||||
|
return out_weights, out_ids
|
||||||
|
|
||||||
|
|
||||||
def gemma_dual_rmsnorm_residual_scalar(
|
def gemma_dual_rmsnorm_residual_scalar(
|
||||||
x1: torch.Tensor,
|
x1: torch.Tensor,
|
||||||
weight1: torch.Tensor,
|
weight1: torch.Tensor,
|
||||||
|
|||||||
@@ -878,6 +878,15 @@ class Gemma4RMSNorm(MultiPlatformOp):
|
|||||||
out = out.reshape(original_shape)
|
out = out.reshape(original_shape)
|
||||||
return out
|
return out
|
||||||
|
|
||||||
|
def forward_xpu(self, x: torch.Tensor) -> torch.Tensor:
|
||||||
|
if x.numel() == 0:
|
||||||
|
return x
|
||||||
|
if self.with_scale and self.scale_shift == 1.0:
|
||||||
|
out = gemma_rmsnorm(x, self.weight.data, self.eps)
|
||||||
|
else:
|
||||||
|
out = rmsnorm(x, self.weight.data, self.eps)
|
||||||
|
return out
|
||||||
|
|
||||||
def forward_hip(self, x: torch.Tensor) -> torch.Tensor:
|
def forward_hip(self, x: torch.Tensor) -> torch.Tensor:
|
||||||
# sgl_kernel's gemma_rmsnorm is not available on ROCm;
|
# sgl_kernel's gemma_rmsnorm is not available on ROCm;
|
||||||
# delegate to the pure-PyTorch implementation.
|
# delegate to the pure-PyTorch implementation.
|
||||||
|
|||||||
@@ -34,6 +34,7 @@ from sglang.srt.layers.gemma4_fused_ops import (
|
|||||||
gemma_dual_rmsnorm_residual_scalar,
|
gemma_dual_rmsnorm_residual_scalar,
|
||||||
gemma_qkv_rmsnorm,
|
gemma_qkv_rmsnorm,
|
||||||
gemma_rmsnorm_residual_scalar,
|
gemma_rmsnorm_residual_scalar,
|
||||||
|
gemma_routing_post_topk,
|
||||||
)
|
)
|
||||||
from sglang.srt.layers.layernorm import Gemma4RMSNorm, RMSNorm
|
from sglang.srt.layers.layernorm import Gemma4RMSNorm, RMSNorm
|
||||||
from sglang.srt.layers.linear import (
|
from sglang.srt.layers.linear import (
|
||||||
@@ -56,6 +57,9 @@ from sglang.srt.model_loader.weight_utils import (
|
|||||||
maybe_remap_kv_scale_name,
|
maybe_remap_kv_scale_name,
|
||||||
)
|
)
|
||||||
from sglang.srt.models.gemma3_causal import Gemma3MLP, Gemma3TextScaledWordEmbedding
|
from sglang.srt.models.gemma3_causal import Gemma3MLP, Gemma3TextScaledWordEmbedding
|
||||||
|
from sglang.srt.models.utils import (
|
||||||
|
create_fused_set_kv_buffer_arg,
|
||||||
|
)
|
||||||
from sglang.srt.server_args import get_global_server_args
|
from sglang.srt.server_args import get_global_server_args
|
||||||
from sglang.srt.utils import add_prefix, make_layers
|
from sglang.srt.utils import add_prefix, make_layers
|
||||||
|
|
||||||
@@ -145,7 +149,8 @@ class Gemma4Router(nn.Module):
|
|||||||
super().__init__()
|
super().__init__()
|
||||||
self.hidden_size = config.hidden_size
|
self.hidden_size = config.hidden_size
|
||||||
|
|
||||||
# RMSNorm without learned weight — pure normalization only
|
# RMSNorm without learned weight — scale is folded into norm weight
|
||||||
|
# after loading so forward is a single fused norm kernel.
|
||||||
self.norm = Gemma4RMSNorm(
|
self.norm = Gemma4RMSNorm(
|
||||||
self.hidden_size, eps=config.rms_norm_eps, with_scale=False
|
self.hidden_size, eps=config.rms_norm_eps, with_scale=False
|
||||||
)
|
)
|
||||||
@@ -165,18 +170,19 @@ class Gemma4Router(nn.Module):
|
|||||||
quant_config=None,
|
quant_config=None,
|
||||||
prefix=add_prefix("proj", prefix),
|
prefix=add_prefix("proj", prefix),
|
||||||
)
|
)
|
||||||
self._fused_scale: Optional[torch.Tensor] = None
|
self._scale_fused = False
|
||||||
|
|
||||||
def fuse_scale(self):
|
def fuse_scale(self):
|
||||||
"""Pre-compute scale * root_size. Call after weights are loaded."""
|
"""Fold scale * root_size into norm.weight so forward needs no extra mul."""
|
||||||
self._fused_scale = (self.scale * self.root_size).to(self.scale.dtype)
|
fused = (self.scale * self.root_size).to(self.norm.weight.dtype)
|
||||||
|
self.norm.weight.data.copy_(fused)
|
||||||
|
self._scale_fused = True
|
||||||
|
|
||||||
def forward(self, x: torch.Tensor) -> torch.Tensor:
|
def forward(self, x: torch.Tensor) -> torch.Tensor:
|
||||||
"""Returns raw router logits [T, E]."""
|
"""Returns raw router logits [T, E]."""
|
||||||
x = self.norm(x)
|
if not self._scale_fused:
|
||||||
if self._fused_scale is None:
|
|
||||||
self.fuse_scale()
|
self.fuse_scale()
|
||||||
x = x * self._fused_scale.to(x.dtype)
|
x = self.norm(x)
|
||||||
router_logits, _ = self.proj(x)
|
router_logits, _ = self.proj(x)
|
||||||
return router_logits
|
return router_logits
|
||||||
|
|
||||||
@@ -230,13 +236,15 @@ class Gemma4MoE(nn.Module):
|
|||||||
return gemma4_fused_routing(gating_output, per_expert_scale, topk)
|
return gemma4_fused_routing(gating_output, per_expert_scale, topk)
|
||||||
|
|
||||||
topk_logits, topk_ids = torch.topk(gating_output, k=topk, dim=-1)
|
topk_logits, topk_ids = torch.topk(gating_output, k=topk, dim=-1)
|
||||||
topk_weights = torch.nn.functional.softmax(topk_logits, dim=-1)
|
|
||||||
|
|
||||||
# Fold per_expert_scale into routing weights
|
# Fused: softmax + per_expert_scale gather + mul + casts in one kernel
|
||||||
|
if topk_logits.is_cuda or topk_logits.is_xpu:
|
||||||
|
return gemma_routing_post_topk(topk_logits, topk_ids, per_expert_scale)
|
||||||
|
|
||||||
|
topk_weights = torch.nn.functional.softmax(topk_logits, dim=-1)
|
||||||
topk_weights = topk_weights * per_expert_scale[topk_ids].to(
|
topk_weights = topk_weights * per_expert_scale[topk_ids].to(
|
||||||
topk_weights.dtype
|
topk_weights.dtype
|
||||||
)
|
)
|
||||||
|
|
||||||
return topk_weights.to(torch.float32), topk_ids.to(torch.int32)
|
return topk_weights.to(torch.float32), topk_ids.to(torch.int32)
|
||||||
|
|
||||||
self.topk = TopK(
|
self.topk = TopK(
|
||||||
@@ -407,14 +415,15 @@ class Gemma4Attention(nn.Module):
|
|||||||
q, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1)
|
q, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1)
|
||||||
|
|
||||||
# Fused Q/K/V RMSNorm: replaces three separate norm kernels with one.
|
# Fused Q/K/V RMSNorm: replaces three separate norm kernels with one.
|
||||||
# Preconditions for the fused path: tensors on CUDA, q_norm/k_norm use
|
# Preconditions for the fused path: tensors on CUDA or XPU (the kernel
|
||||||
# the standard norm*weight (scale_shift==0) and v_norm has weight=ones
|
# is pure Triton and lowers to both backends), q_norm/k_norm use the
|
||||||
|
# standard norm*weight (scale_shift==0) and v_norm has weight=ones
|
||||||
# (with_scale=False) — the canonical Gemma4 attention configuration.
|
# (with_scale=False) — the canonical Gemma4 attention configuration.
|
||||||
is_kv_shared = (
|
is_kv_shared = (
|
||||||
self.is_kv_shared_layer and self.kv_shared_layer_index is not None
|
self.is_kv_shared_layer and self.kv_shared_layer_index is not None
|
||||||
)
|
)
|
||||||
can_fuse_qkv_norm = (
|
can_fuse_qkv_norm = (
|
||||||
q.is_cuda
|
(q.is_cuda or q.is_xpu)
|
||||||
and self.q_norm.scale_shift == 0.0
|
and self.q_norm.scale_shift == 0.0
|
||||||
and self.k_norm.scale_shift == 0.0
|
and self.k_norm.scale_shift == 0.0
|
||||||
and not self.v_norm.with_scale
|
and not self.v_norm.with_scale
|
||||||
@@ -466,9 +475,22 @@ class Gemma4Attention(nn.Module):
|
|||||||
v = self.v_norm(v)
|
v = self.v_norm(v)
|
||||||
|
|
||||||
# Apply rotary embedding
|
# Apply rotary embedding
|
||||||
|
use_fused_kv = False
|
||||||
if k is not None:
|
if k is not None:
|
||||||
k = k.flatten(-2, -1)
|
k = k.flatten(-2, -1)
|
||||||
q, k = self.rotary_emb(positions, q, k)
|
# Fuse RoPE + KV-cache write for non-SWA layers with bf16 cache
|
||||||
|
# DISABLED: causes accuracy regression in launch_server path
|
||||||
|
can_fuse = False
|
||||||
|
if can_fuse:
|
||||||
|
fused_arg = create_fused_set_kv_buffer_arg(
|
||||||
|
value=v.flatten(-2, -1) if v.dim() == 3 else v,
|
||||||
|
layer=self.attn,
|
||||||
|
forward_batch=forward_batch,
|
||||||
|
)
|
||||||
|
use_fused_kv = True
|
||||||
|
else:
|
||||||
|
fused_arg = None
|
||||||
|
q, k = self.rotary_emb(positions, q, k, fused_set_kv_buffer_arg=fused_arg)
|
||||||
k = k.unflatten(-1, (self.num_kv_heads, self.head_dim))
|
k = k.unflatten(-1, (self.num_kv_heads, self.head_dim))
|
||||||
else:
|
else:
|
||||||
# Rotary embedding requires a key input; use zeros since KV is shared from another layer
|
# Rotary embedding requires a key input; use zeros since KV is shared from another layer
|
||||||
@@ -481,7 +503,7 @@ class Gemma4Attention(nn.Module):
|
|||||||
k,
|
k,
|
||||||
v,
|
v,
|
||||||
forward_batch=forward_batch,
|
forward_batch=forward_batch,
|
||||||
save_kv_cache=not self.is_kv_shared_layer,
|
save_kv_cache=not self.is_kv_shared_layer and not use_fused_kv,
|
||||||
)
|
)
|
||||||
if attn_output.dim() == 3:
|
if attn_output.dim() == 3:
|
||||||
attn_output = attn_output.flatten(-2, -1)
|
attn_output = attn_output.flatten(-2, -1)
|
||||||
@@ -666,7 +688,7 @@ class Gemma4DecoderLayer(nn.Module):
|
|||||||
# Fused: (rmsnorm(rmsnorm(h1,w1) + rmsnorm(h2,w2), w3) + residual) * scalar
|
# Fused: (rmsnorm(rmsnorm(h1,w1) + rmsnorm(h2,w2), w3) + residual) * scalar
|
||||||
if (
|
if (
|
||||||
not self.has_ple
|
not self.has_ple
|
||||||
and hidden_states_1.is_cuda
|
and (hidden_states_1.is_cuda or hidden_states_1.is_xpu)
|
||||||
and hidden_states_1.dim() == 2
|
and hidden_states_1.dim() == 2
|
||||||
):
|
):
|
||||||
norm1 = self.post_feedforward_layernorm_1
|
norm1 = self.post_feedforward_layernorm_1
|
||||||
@@ -698,7 +720,12 @@ class Gemma4DecoderLayer(nn.Module):
|
|||||||
)
|
)
|
||||||
hidden_states = self.mlp(hidden_states)
|
hidden_states = self.mlp(hidden_states)
|
||||||
|
|
||||||
if not self.has_ple and hidden_states.is_cuda and hidden_states.dim() == 2:
|
if (
|
||||||
|
not self.has_ple
|
||||||
|
and self.moe is None
|
||||||
|
and (hidden_states.is_cuda or hidden_states.is_xpu)
|
||||||
|
and hidden_states.dim() == 2
|
||||||
|
):
|
||||||
# Fused: (post_ff_norm(h) + residual) * layer_scalar in one kernel
|
# Fused: (post_ff_norm(h) + residual) * layer_scalar in one kernel
|
||||||
norm = self.post_feedforward_layernorm
|
norm = self.post_feedforward_layernorm
|
||||||
hidden_states = gemma_rmsnorm_residual_scalar(
|
hidden_states = gemma_rmsnorm_residual_scalar(
|
||||||
|
|||||||
@@ -2358,12 +2358,12 @@ class ServerArgs:
|
|||||||
self.attention_backend = default_attention_backend
|
self.attention_backend = default_attention_backend
|
||||||
|
|
||||||
prefill_backend, decode_backend = self.get_attention_backends()
|
prefill_backend, decode_backend = self.get_attention_backends()
|
||||||
accepted_backends = ("trtllm_mha", "triton")
|
accepted_backends = ("trtllm_mha", "triton", "intel_xpu")
|
||||||
assert (
|
assert (
|
||||||
prefill_backend in accepted_backends
|
prefill_backend in accepted_backends
|
||||||
and decode_backend in accepted_backends
|
and decode_backend in accepted_backends
|
||||||
), (
|
), (
|
||||||
"Gemma4 only supports trtllm_mha or triton attention backend, "
|
"Gemma4 only supports trtllm_mha, triton, or intel_xpu attention backend, "
|
||||||
f"got prefill={prefill_backend}, decode={decode_backend}"
|
f"got prefill={prefill_backend}, decode={decode_backend}"
|
||||||
)
|
)
|
||||||
|
|
||||||
|
|||||||
@@ -0,0 +1,136 @@
|
|||||||
|
"""
|
||||||
|
Gemma 4 E2B-it: simple text Q&A on XPU (OpenAI /v1), same shape as
|
||||||
|
``test_deepseek_coder_v2_lite_instruct.py``.
|
||||||
|
|
||||||
|
Model card: https://huggingface.co/google/gemma-4-E2B-it
|
||||||
|
|
||||||
|
- XPU test runs when Intel XPU is available.
|
||||||
|
|
||||||
|
Run from test/srt::
|
||||||
|
|
||||||
|
python3 -m unittest xpu.test_gemma_4_e2b.TestGemma4E2BXPU.test_simple_qa
|
||||||
|
|
||||||
|
A single end-to-end test (``test_simple_qa``) verifies the model boots
|
||||||
|
on XPU and returns a coherent reply. On failure the assertion message
|
||||||
|
includes the model's actual output.
|
||||||
|
|
||||||
|
Server is started with ``sglang serve`` (``--model-impl sglang``).
|
||||||
|
"""
|
||||||
|
|
||||||
|
from __future__ import annotations
|
||||||
|
|
||||||
|
import os
|
||||||
|
import unittest
|
||||||
|
|
||||||
|
import openai
|
||||||
|
import torch
|
||||||
|
|
||||||
|
from sglang.srt.utils.common import is_xpu
|
||||||
|
from sglang.test.test_utils import CustomTestCase
|
||||||
|
from sglang.test.vlm_utils import (
|
||||||
|
DEFAULT_URL_FOR_TEST,
|
||||||
|
kill_process_tree,
|
||||||
|
popen_launch_server,
|
||||||
|
)
|
||||||
|
|
||||||
|
MODEL = "google/gemma-4-E2B-it"
|
||||||
|
|
||||||
|
LAUNCH_TIMEOUT = 900
|
||||||
|
|
||||||
|
# The -it model ships its own chat_template.jinja, so no --chat-template needed.
|
||||||
|
# E2B model: single-rank for small model on XPU.
|
||||||
|
XPU_SERVER_ARGS = [
|
||||||
|
"--device",
|
||||||
|
"xpu",
|
||||||
|
"--tp=1",
|
||||||
|
"--trust-remote-code",
|
||||||
|
"--disable-overlap-schedule",
|
||||||
|
"--page-size",
|
||||||
|
"64",
|
||||||
|
"--mem-fraction-static",
|
||||||
|
"0.70",
|
||||||
|
"--attention-backend",
|
||||||
|
"intel_xpu",
|
||||||
|
"--model-impl",
|
||||||
|
"sglang",
|
||||||
|
]
|
||||||
|
|
||||||
|
# Standard sglang e2e Q&A prompt (see test_openai_server.py::run_chat_completion).
|
||||||
|
_SIMPLE_QA_PROMPT = "What is the capital of France? Answer in a few words."
|
||||||
|
|
||||||
|
|
||||||
|
def _simple_text_messages():
|
||||||
|
return [
|
||||||
|
{
|
||||||
|
"role": "user",
|
||||||
|
"content": [
|
||||||
|
{"type": "text", "text": _SIMPLE_QA_PROMPT},
|
||||||
|
],
|
||||||
|
}
|
||||||
|
]
|
||||||
|
|
||||||
|
|
||||||
|
def _empty_xpu_cache() -> None:
|
||||||
|
"""Release cached XPU allocations so back-to-back tests start clean."""
|
||||||
|
if torch.xpu.is_available():
|
||||||
|
torch.xpu.empty_cache()
|
||||||
|
|
||||||
|
|
||||||
|
@unittest.skipUnless(is_xpu(), "Intel XPU not available")
|
||||||
|
class TestGemma4E2BXPU(CustomTestCase):
|
||||||
|
@classmethod
|
||||||
|
def setUpClass(cls):
|
||||||
|
cls.model = MODEL
|
||||||
|
cls.base_url = DEFAULT_URL_FOR_TEST
|
||||||
|
cls.api_key = "sk-123456"
|
||||||
|
os.environ["SGLANG_USE_SGL_XPU"] = "1"
|
||||||
|
|
||||||
|
_empty_xpu_cache()
|
||||||
|
cls.process = popen_launch_server(
|
||||||
|
cls.model,
|
||||||
|
cls.base_url,
|
||||||
|
timeout=LAUNCH_TIMEOUT,
|
||||||
|
api_key=cls.api_key,
|
||||||
|
other_args=list(XPU_SERVER_ARGS),
|
||||||
|
device="xpu",
|
||||||
|
)
|
||||||
|
cls.base_url += "/v1"
|
||||||
|
|
||||||
|
@classmethod
|
||||||
|
def tearDownClass(cls):
|
||||||
|
kill_process_tree(cls.process.pid)
|
||||||
|
_empty_xpu_cache()
|
||||||
|
|
||||||
|
def test_simple_qa(self):
|
||||||
|
client = openai.Client(api_key=self.api_key, base_url=self.base_url)
|
||||||
|
response = client.chat.completions.create(
|
||||||
|
model="default",
|
||||||
|
messages=_simple_text_messages(),
|
||||||
|
temperature=0,
|
||||||
|
max_tokens=96,
|
||||||
|
)
|
||||||
|
msg = response.choices[0].message
|
||||||
|
text = msg.content or ""
|
||||||
|
reasoning = getattr(msg, "reasoning_content", None) or ""
|
||||||
|
combined = f"{text} {reasoning}".strip()
|
||||||
|
|
||||||
|
self.assertEqual(msg.role, "assistant", f"unexpected role; got: {combined!r}")
|
||||||
|
self.assertGreater(len(combined), 0, "empty reply from model")
|
||||||
|
self.assertIn(
|
||||||
|
"paris", combined.lower(), f"expected `Paris` in reply, got: {combined!r}"
|
||||||
|
)
|
||||||
|
self.assertIsNotNone(response.usage)
|
||||||
|
self.assertGreater(
|
||||||
|
response.usage.completion_tokens,
|
||||||
|
0,
|
||||||
|
f"no tokens generated; got: {combined!r}",
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
|
from sglang.test.ci.ci_register import register_xpu_ci
|
||||||
|
|
||||||
|
# Single e2e test: boot + a short Q&A.
|
||||||
|
register_xpu_ci(est_time=240, suite="stage-b-test-1-gpu-xpu")
|
||||||
|
|
||||||
|
if __name__ == "__main__":
|
||||||
|
unittest.main()
|
||||||
Reference in New Issue
Block a user