fix: Investigate Gemma3 1B decoder output discrepancy

tensorrt_llm/_torch/attention_backend/vanilla.py

@@ -1,3 +1,4 @@
+import math
 from typing import Optional
 
 import torch
@@ -39,6 +40,23 @@ def generate_causal_mask(batch_size: int, target_length: int,
     return causal_mask
 
 
+def generate_sliding_window_mask(batch_size: int, target_length: int,
+                                 cache_position: torch.Tensor,
+                                 device: torch.device,
+                                 attention_window_size: int):
+    # TRTLLM's sliding window attention is inclusive.
+    effective_window_size = attention_window_size + 1
+    attention_mask_1 = torch.arange(
+        target_length,
+        device=device).unsqueeze(0) <= cache_position.unsqueeze(-1)
+    attention_mask_2 = torch.arange(target_length, device=device).unsqueeze(
+        0) > cache_position.unsqueeze(-1) - effective_window_size
+    attention_mask = attention_mask_1 & attention_mask_2
+    attention_mask = attention_mask[None,
+                                    None, :, :].expand(batch_size, 1, -1, -1)
+    return attention_mask
+
+
 class VanillaAttentionMetadata(AttentionMetadata):
 
     def prepare(self) -> None:
@@ -66,11 +84,17 @@ def __init__(
         head_dim: int,
         num_kv_heads: Optional[int] = None,
         quant_config: Optional[QuantConfig] = None,
+        q_scaling: Optional[float] = None,
         **kwargs,
     ):
-        super().__init__(layer_idx, num_heads, head_dim, num_kv_heads,
-                         quant_config, **kwargs)
+        super().__init__(layer_idx,
+                         num_heads,
+                         head_dim,
+                         num_kv_heads=num_kv_heads,
+                         quant_config=quant_config,
+                         **kwargs)
         self.num_key_value_groups = self.num_heads // self.num_kv_heads
+        self.q_scaling = q_scaling
 
     def _single_request_update_kv_cache(self, k, v, kv_cache_tensor, seq_len,
                                         cache_idx, cache_position):
@@ -86,8 +110,15 @@ def _single_request_update_kv_cache(self, k, v, kv_cache_tensor, seq_len,
 
         return k_out[:, :seq_len, :, :], v_out[:, :seq_len, :, :]
 
-    def _single_request_forward(self, q, k, v, attention_mask: AttentionMask,
-                                kv_cache_tensor, past_seen_token, cache_idx):
+    def _single_request_forward(self,
+                                q,
+                                k,
+                                v,
+                                attention_mask: AttentionMask,
+                                kv_cache_tensor,
+                                past_seen_token,
+                                cache_idx,
+                                attention_window_size: Optional[int] = None):
 
         bsz = 1
         q_len = q.size(0)
@@ -129,7 +160,12 @@ def _single_request_forward(self, q, k, v, attention_mask: AttentionMask,
         is_causal = False
         attn_mask = None
         if attention_mask == PredefinedAttentionMask.CAUSAL:
-            if past_seen_token == 0:
+            # Create custom sliding window mask as sdpa doesn't natively support it.
+            if attention_window_size is not None:
+                attn_mask = generate_sliding_window_mask(
+                    bsz, target_seq_len, cache_position, q.device,
+                    attention_window_size)
+            elif past_seen_token == 0:
                 is_causal = True
             elif q_len != 1:
                 # attn_mask: 4-D tensor (batch_size, 1, query_seq_len, seq_len)
@@ -140,12 +176,17 @@ def _single_request_forward(self, q, k, v, attention_mask: AttentionMask,
         else:
             raise ValueError("Unexpected attention mask type")
 
+        qk_scale = None
+        if self.q_scaling is not None:
+            qk_scale = 1 / (math.sqrt(self.head_dim) * self.q_scaling)
+
         attn_output = torch.nn.functional.scaled_dot_product_attention(
             q,
             key_states,
             value_states,
             is_causal=is_causal,
             attn_mask=attn_mask,
+            scale=qk_scale,
         )
 
         attn_output = attn_output.squeeze(0)
@@ -229,6 +270,7 @@ def forward(self,
                 metadata: VanillaAttentionMetadata,
                 *,
                 attention_mask: AttentionMask = PredefinedAttentionMask.CAUSAL,
+                attention_window_size: Optional[int] = None,
                 **kwargs) -> torch.Tensor:
         if metadata.kv_cache_manager is None:
             # NOTE: WAR for no kv cache attn e.g. BERT,
@@ -270,11 +312,9 @@ def forward(self,
             past_seen_token = past_seen_tokens[i]
             cache_idx = cache_indices[i]
 
-            attn_output = self._single_request_forward(single_q, single_k,
-                                                       single_v, attention_mask,
-                                                       kv_cache_tensor,
-                                                       past_seen_token,
-                                                       cache_idx)
+            attn_output = self._single_request_forward(
+                single_q, single_k, single_v, attention_mask, kv_cache_tensor,
+                past_seen_token, cache_idx, attention_window_size)
             attn_outputs.append(attn_output)
 
             offset += seq_len

tensorrt_llm/_torch/models/modeling_gemma3.py

@@ -65,7 +65,7 @@ def __init__(
         self.attention_window_size = None
         if is_sliding:
             rope_params.theta = 10000
-            self.attention_window_size = config.sliding_window
+            self.attention_window_size = config.sliding_window - 1  # Gemma3 sliding window isn't inclusive.
         pos_embd_params = PositionalEmbeddingParams(
             type=PositionEmbeddingType.rope_gpt_neox,
             rope=rope_params,
@@ -107,15 +107,14 @@ def forward(
         **kwargs,
     ) -> torch.Tensor:
 
-        attention_window_size = self.attention_window_size or attn_metadata.max_seq_len
         return super().forward(position_ids=position_ids,
                                hidden_states=hidden_states,
                                attn_metadata=attn_metadata,
                                attention_mask=attention_mask,
                                mrope_config=mrope_config,
                                all_reduce_params=all_reduce_params,
                                lora_params=lora_params,
-                               attention_window_size=attention_window_size,
+                               attention_window_size=self.attention_window_size,
                                **kwargs)
 
     def apply_qk_norm(self, q, k):

tests/integration/defs/accuracy/test_llm_api_pytorch.py

@@ -437,7 +437,12 @@ class TestGemma3_1BInstruct(LlmapiAccuracyTestHarness):
     MODEL_PATH = f"{llm_models_root()}/gemma/gemma-3-1b-it/"
 
     def test_auto_dtype(self):
-        with LLM(self.MODEL_PATH) as llm:
+        # Disabling kv cache reuse as a WAR to deal with gaps in kernel support for Gemma3's non-inclusive sliding window size.
+        kv_cache_config = KvCacheConfig(
+            enable_block_reuse=False,
+            enable_partial_reuse=False,
+        )
+        with LLM(self.MODEL_PATH, kv_cache_config=kv_cache_config) as llm:
             task = CnnDailymail(self.MODEL_NAME)
             task.evaluate(llm)