feat: init gemma3

Author: YzXiao101

python/minisgl/attention/base.py

@@ -18,7 +18,15 @@ def get_last_indices(self, bs: int) -> torch.Tensor: ...
 class BaseAttnBackend(ABC):
     @abstractmethod
     def forward(
-        self, q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, layer_id: int, batch: Batch
+        self,
+        q: torch.Tensor,
+        k: torch.Tensor,
+        v: torch.Tensor,
+        layer_id: int,
+        batch: Batch,
+        *,
+        window_size: tuple[int, int] = (-1, -1),
+        softmax_scale: float | None = None,
     ) -> torch.Tensor: ...
 
     @abstractmethod
@@ -44,10 +52,26 @@ def __init__(
         self.decode_backend = decode_backend
 
     def forward(
-        self, q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, layer_id: int, batch: Batch
+        self,
+        q: torch.Tensor,
+        k: torch.Tensor,
+        v: torch.Tensor,
+        layer_id: int,
+        batch: Batch,
+        *,
+        window_size: tuple[int, int] = (-1, -1),
+        softmax_scale: float | None = None,
     ) -> torch.Tensor:
         backend = self.prefill_backend if batch.is_prefill else self.decode_backend
-        return backend.forward(q, k, v, layer_id, batch)
+        return backend.forward(
+            q,
+            k,
+            v,
+            layer_id,
+            batch,
+            window_size=window_size,
+            softmax_scale=softmax_scale,
+        )
 
     def prepare_metadata(self, batch: Batch) -> None:
         backend = self.prefill_backend if batch.is_prefill else self.decode_backend

python/minisgl/attention/fa.py

@@ -46,7 +46,15 @@ def __init__(self, config: ModelConfig):
         self.version = 4 if is_sm100_supported() else 3
 
     def forward(
-        self, q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, layer_id: int, batch: Batch
+        self,
+        q: torch.Tensor,
+        k: torch.Tensor,
+        v: torch.Tensor,
+        layer_id: int,
+        batch: Batch,
+        *,
+        window_size: tuple[int, int] = (-1, -1),
+        softmax_scale: float | None = None,
     ) -> torch.Tensor:
         metadata = batch.attn_metadata
         assert isinstance(metadata, FAMetadata)
@@ -60,8 +68,9 @@ def forward(
             cu_seqlens_q=metadata.cu_seqlens_q,
             cu_seqlens_k=metadata.cu_seqlens_k,
             max_seqlen_q=metadata.max_seqlen_q,
-            softmax_scale=self.scale,
+            softmax_scale=self.scale if softmax_scale is None else softmax_scale,
             version=self.version,
+            window_size=window_size,
         )
 
     def prepare_metadata(self, batch: Batch) -> None:

python/minisgl/attention/fi.py

@@ -174,11 +174,22 @@ def _get_ones_cpu(self, bs: int) -> torch.Tensor:
         return self.cached_ones_cpu[:bs]
 
     def forward(
-        self, q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, layer_id: int, batch: Batch
+        self,
+        q: torch.Tensor,
+        k: torch.Tensor,
+        v: torch.Tensor,
+        layer_id: int,
+        batch: Batch,
+        *,
+        window_size: tuple[int, int] = (-1, -1),
+        softmax_scale: float | None = None,
     ) -> torch.Tensor:
         def _flatten_cache(cache: torch.Tensor) -> torch.Tensor:  # treat page = 1
             return cache.view(-1, 1, cache.shape[2], cache.shape[3])
 
+        if window_size != (-1, -1) or softmax_scale is not None:
+            raise NotImplementedError
+
         metadata = batch.attn_metadata
         assert isinstance(metadata, FIMetadata)
         self._initialize_metadata_once(metadata)

python/minisgl/attention/trtllm.py

@@ -47,11 +47,22 @@ def __init__(self, config: ModelConfig):
         )
 
     def forward(
-        self, q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, layer_id: int, batch: Batch
+        self,
+        q: torch.Tensor,
+        k: torch.Tensor,
+        v: torch.Tensor,
+        layer_id: int,
+        batch: Batch,
+        *,
+        window_size: tuple[int, int] = (-1, -1),
+        softmax_scale: float | None = None,
     ) -> torch.Tensor:
         from flashinfer.decode import trtllm_batch_decode_with_kv_cache
         from flashinfer.prefill import trtllm_batch_context_with_kv_cache
 
+        if window_size != (-1, -1) or softmax_scale is not None:
+            raise NotImplementedError
+
         metadata = batch.attn_metadata
         assert isinstance(metadata, TRTLLMMetadata)
         self.kvcache.store_kv(k, v, batch.out_loc, layer_id)

python/minisgl/layers/__init__.py

@@ -1,4 +1,4 @@
-from .activation import gelu_and_mul, silu_and_mul
+from .activation import gelu_and_mul, gelu_tanh_and_mul, silu_and_mul
 from .attention import AttentionLayer
 from .base import BaseOP, OPList, StateLessOP
 from .embedding import ParallelLMHead, VocabParallelEmbedding
@@ -10,12 +10,13 @@
     LinearRowParallel,
 )
 from .moe import MoELayer
-from .norm import RMSNorm, RMSNormFused
+from .norm import Gemma3RMSNorm, RMSNorm, RMSNormFused
 from .rotary import get_rope, set_rope_device
 
 __all__ = [
     "silu_and_mul",
     "gelu_and_mul",
+    "gelu_tanh_and_mul",
     "AttentionLayer",
     "BaseOP",
     "StateLessOP",
@@ -26,6 +27,7 @@
     "LinearRowParallel",
     "LinearOProj",
     "LinearQKVMerged",
+    "Gemma3RMSNorm",
     "RMSNorm",
     "RMSNormFused",
     "get_rope",

python/minisgl/layers/activation.py

@@ -18,4 +18,10 @@ def gelu_and_mul(x: torch.Tensor, out: torch.Tensor | None = None):
     return gelu_and_mul(x, out=out)
 
 
-__all__ = ["silu_and_mul", "gelu_and_mul"]
+def gelu_tanh_and_mul(x: torch.Tensor, out: torch.Tensor | None = None):
+    from flashinfer import gelu_tanh_and_mul
+
+    return gelu_tanh_and_mul(x, out=out)
+
+
+__all__ = ["silu_and_mul", "gelu_and_mul", "gelu_tanh_and_mul"]

python/minisgl/layers/attention.py

@@ -25,6 +25,8 @@ def __init__(
         rotary_config: RotaryConfig,
         q_norm: RMSNorm | None = None,
         k_norm: RMSNorm | None = None,
+        sliding_window_size: int | None = None,
+        softmax_scale: float | None = None,
     ):
         assert num_qo_heads % num_kv_heads == 0
         self.layer_id = layer_id
@@ -43,6 +45,11 @@ def __init__(
         )
         self.q_norm = q_norm
         self.k_norm = k_norm
+        # sliding_window_size: HF-convention (inclusive). Converted to FA (left, right) here.
+        self._window_size = (
+            (sliding_window_size - 1, 0) if sliding_window_size is not None else (-1, -1)
+        )
+        self._softmax_scale = softmax_scale
 
     def forward(self, qkv: torch.Tensor) -> torch.Tensor:
         ctx = get_global_ctx()
@@ -53,5 +60,13 @@ def forward(self, qkv: torch.Tensor) -> torch.Tensor:
             self.k_norm.forward_inplace(k.view(-1, self.num_kv_heads, self.head_dim))
         q, k = self.rotary.forward(ctx.batch.positions, q, k)
         q = q.view(-1, self.num_qo_heads, self.head_dim)
-        o = ctx.attn_backend.forward(q, k, v, self.layer_id, ctx.batch)
+        o = ctx.attn_backend.forward(
+            q,
+            k,
+            v,
+            self.layer_id,
+            ctx.batch,
+            window_size=self._window_size,
+            softmax_scale=self._softmax_scale,
+        )
         return o.view(-1, self.qo_attn_dim)

python/minisgl/layers/linear.py

@@ -100,9 +100,11 @@ def __init__(self, input_size: int, output_size: int, has_bias: bool):
         super().__init__(full_isize, full_osize, local_isize, local_osize, has_bias)
 
     def forward(self, x: torch.Tensor) -> torch.Tensor:
-        y = F.linear(x, self.weight, self.bias)
+        y = F.linear(x, self.weight, None)
         if self._tp_size > 1:
             y = self._comm.all_reduce(y)
+        if self.bias is not None:
+            y = y + self.bias
         return y
 
 
@@ -121,7 +123,9 @@ def __init__(
         super().__init__(input_size, output_size, local_input_size, local_output_size, has_bias)
 
     def forward(self, x: torch.Tensor) -> torch.Tensor:
-        y = F.linear(x, self.weight, self.bias)
+        y = F.linear(x, self.weight, None)
         if self._tp_size > 1:
             y = self._comm.all_reduce(y)
+        if self.bias is not None:
+            y = y + self.bias
         return y

python/minisgl/layers/norm.py

@@ -1,3 +1,5 @@
+from __future__ import annotations
+
 from typing import Tuple
 
 import torch
@@ -36,3 +38,20 @@ def forward(
             return self.rmsnorm(x, self.weight, self.eps), x
         self.fused_add_rmsnorm(x, residual, self.weight, self.eps)
         return x, residual
+
+
+class Gemma3RMSNorm(BaseOP):
+
+    def __init__(self, size: int, eps: float) -> None:
+        from flashinfer import gemma_rmsnorm
+
+        self.eps = eps
+        self.weight = torch.zeros(size)
+        self.gemma_rmsnorm = gemma_rmsnorm
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return self.gemma_rmsnorm(x, self.weight, self.eps)
+
+    def forward_inplace(self, x: torch.Tensor) -> None:
+        shape = x.shape  # [t, h, d]
+        x.copy_(self.gemma_rmsnorm(x.view(-1, shape[-1]), self.weight, self.eps).view(shape))

python/minisgl/layers/rotary.py

@@ -20,7 +20,6 @@ def __init__(
     ) -> None:
         super().__init__()
         self.head_size = head_size
-        assert rotary_dim == head_size
         inv_freq = 1.0 / (base ** (torch.arange(0, rotary_dim, 2, dtype=torch.float) / rotary_dim))
         if post_process is not None:
             inv_freq = post_process(inv_freq)
@@ -30,8 +29,8 @@ def __init__(
         sin = freqs.sin()
         # buffer, so don't load/save
         self._cos_sin_cache = torch.cat((cos, sin), dim=-1)
-        assert self.head_size in [64, 128, 256, 512]
 
+        assert self.head_size in [64, 128, 256, 512]
         from flashinfer import apply_rope_with_cos_sin_cache_inplace
 
         self.apply_rope_with_cos_sin_cache_inplace = apply_rope_with_cos_sin_cache_inplace
@@ -97,15 +96,20 @@ def post_process(inv_freq: torch.Tensor) -> torch.Tensor:
             orig_max_pos: int = rope_scaling["original_max_position_embeddings"]
 
             def _find_correction_dim(num_rotations: float) -> float:
-                return rotary_dim * math.log(orig_max_pos / (num_rotations * 2 * math.pi)) / (2 * math.log(base))
+                return (
+                    rotary_dim
+                    * math.log(orig_max_pos / (num_rotations * 2 * math.pi))
+                    / (2 * math.log(base))
+                )
 
             low = max(math.floor(_find_correction_dim(beta_fast)), 0)
             high = min(math.ceil(_find_correction_dim(beta_slow)), rotary_dim // 2 - 1)
 
             def post_process(inv_freq: torch.Tensor) -> torch.Tensor:
                 ramp = torch.clamp(
                     (torch.arange(rotary_dim // 2, dtype=torch.float32) - low) / max(high - low, 1),
-                    0, 1,
+                    0,
+                    1,
                 )
                 return (inv_freq / factor) * ramp + inv_freq * (1 - ramp)
 
@@ -143,4 +147,4 @@ def get_rope(
     return _get_rope(head_dim, rotary_dim, max_position, base, rope_map)
 
 
-__all__ = ["get_rope", "RotaryEmbedding", "set_rope_device"]
+__all__ = ["get_rope", "RotaryEmbedding", "set_rope_device"]