gemma4_e2b.yaml

# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: Apache-2.0

# Gemma 4 E2B-it — text-only AD export path.
# Uses triton paged attention backend: supports head_dim=512 (global_head_dim),
# paged KV cache, CUDA-graph-compatible, FlashDecoding for decode.
model_factory: Gemma4ForConditionalGeneration
# Use the instruction-tuned tokenizer/chat template for end-to-end prompting.
tokenizer: google/gemma-4-E2B-it
attn_backend: triton
compile_backend: torch-cudagraph
cuda_graph_config:
  batch_sizes: [1, 2, 4, 8, 16, 32, 64, 128, 256, 512]
max_num_tokens: 8192
max_batch_size: 512
max_seq_len: 8192
enable_chunked_prefill: true
kv_cache_config:
  enable_block_reuse: false
  free_gpu_memory_fraction: 0.8
transforms:
  compile_model:
    # Multi-pool VSWA capture is supported (warmup replicates cache metadata per pool).
    piecewise_enabled: true
  mlir_elementwise_fusion:
    # MLIR elementwise kernels currently corrupt piecewise CUDA graph replay for Gemma4 E2B.
    enabled: false
  gather_logits_before_lm_head:
    enabled: true
  fuse_gemms:
    enabled: true