Fix for gemma4 sub-func and continous batching

[quic-hemagnih]

This PR fixes the issue - Gemma4 sub-function w/ CB enabled.

Issue reported :
In QEffGemma4TextModel.forward, the attention mask target_length for global layers was determined by:
target_length = int(inputs_embeds.shape[1]) # = seq_len = 312

then overridden by:

target_length = int(layer_keys.shape[-2]) # = dummy cache size = 312

Both paths resolve to 312 during ONNX tracing. The int() call converts the tensor dimension to a Python integer at trace time, which PyTorch's ONNX exporter bakes in as a static Constant node:

Constant(312) → arange → mask shape [1, 1, 256, 312]

Meanwhile, the KV cache input past_key. N is declared with a dynamic symbol ctx_len in the ONNX graph (shape [full_batch_size, 1, ctx_len, 512]). The compiler resolves ctx_len = 2048 from specializations.json. After the cache scatter+gather update, key_states.shape[-2] = ctx_len = 2048,
so:

attn_weights = matmul(query, key.T) → shape [1, 8, 256, 2048]

The Add node (attn_weights + attention_mask) then receives:

[1, 8, 256, 2048] ← attn_weights (ctx_len resolved dynamically to 2048)
[1, 1, 256, 312] ← attention_mask (ctx_len baked statically as 312)

Why Sliding Layers Were Fine

Sliding attention layers use target_length = min(sliding_window, max_position_embeddings) = 512, which is a config constant the same static value the compiler resolves from the specialization. No mismatch.

The Fix

Move mask construction inside QEffGemma4TextAttention.forward, after the KV cache update, using key_states.shape[-2] without int():

key_states is the post-update tensor whose shape[-2] is a live

symbolic dimension in the ONNX graph — not a traced-away constant

kv_target_length = key_states.shape[-2] # dynamic, no int()

attention_mask = _build_additive_attention_mask(
position_ids=position_ids,
target_length=kv_target_length, # emits Shape→Gather→Cast→Range
dtype=query_states.dtype,
sliding_window=self.sliding_window,
)

Without int(), torch.arange(kv_target_length) emits a dynamic Range node in the ONNX graph. The compiler resolves kv_target_length to ctx_len =
2048 (global) or sliding_window = 512 (sliding) from the specialization — always matching attn_weights.

Tests ran:
Ran gemma4_cb.py for Gemma4-E2B-it and Gemma4-31B-it models (both 2 layers and full layers)

--- Response [0] ---
This is a charming, somewhat whimsical image featuring a stylized, anthropomorphic cat character set in a lush, outdoor, possibly pastoral or garden-like environment.

Central Subject:
The main focus is a cat character with light tan or cream-colored fur. The cat has strikingly large, bright blue eyes that are looking slightly upward and to the side with a gentle, curious expression. Its ears are pointed, and its facial features are soft and rounded.

The cat is dressed in a formal

--- Response [1] ---
Based on the image, here are the objects I can identify:

Living Things:

• Rabbit/Hare: The central subject of the image.
• Trees: Visible on the right side and in the background.
• Grass/Vegetation: Covering the fields and foreground.
• Flowers: Various types of wildflowers are visible in the foreground and fields.
• Hills/Mountains: In the distant background.

Man-made Objects/Structures:

--- Response [2] ---
The main subject of the image is a cat.
The cat appears to be a character from an animated work, given its stylized appearance.

--- Response [3] ---
The predominant colors in the image are:

• Earthy tones: Various shades of brown (for the dirt path, the animal's fur, and some of the foliage) and tan/beige (for the clothing and some of the dry grass).
• Greens: Various shades of green are prominent in the lush grass, trees, and distant rolling hills.
• Blues/Teals: A noticeable blue or teal color is present

Execution info:
Average Prefill time a.k.a TTFT is= 11.7 sec
Decode is= 91.29 tokens/sec
Total is= 7.74 tokens/sec
Total (E2E) inference time is= 51.15 sec