AWQ and Gemma3 compatibility problem

---

Title: Gemma3 + AWQ produces NaN/overflow due to float16 limitations with large RMSNorm weights

Description

Gemma3 models fail to produce meaningful output when using AWQ quantization in CTranslate2. The model outputs only <pad> tokens with a score of exactly -65504.0 (the negative maximum value of float16), indicating numerical overflow.

This issue is specific to AWQ quantization - the same Gemma3 model works correctly with int8 quantization in CTranslate2.

Root Cause

Gemma3 models have unusually large RMSNorm weights compared to other architectures. For example, in gemma-3-1b-it:

LayerNorm	Min Weight	Max Weight
input_layernorm	2.27	27.88
post_attention_layernorm	-1.41	51.25
pre_feedforward_layernorm	-284.00	0.25
post_feedforward_layernorm	-1.00	66.50

These extreme values (especially -284) cause overflow when computed in float16, which has a maximum representable value of ~65504.

The AWQ path in CTranslate2 enforces float16 for all operations:

// From ops/awq/gemm.cc
if (a.dtype() != DataType::FLOAT16 && b.dtype() != DataType::INT32)
    throw std::invalid_argument("Awq gemm is only supported for float16 input and int32 weight");

While the RMSNorm CUDA kernel does compute internally in float32, the output is cast back to float16, and subsequent operations (particularly the FFN's silu(gate) * up multiplication) overflow.

Evidence

1. AWQ model produces garbage output:

Output tokens: ['<pad>', '<pad>', '<pad>', '<pad>', ...]
Score: -65504.0

2. The same AWQ model works in PyTorch with float32 patching:

When we wrap the MLP to compute silu(gate) * up in float32, the model works correctly:

class Gemma3MLPWrapper(nn.Module):
    def __init__(self, original_mlp):
        super().__init__()
        self.gate_proj = original_mlp.gate_proj
        self.up_proj = original_mlp.up_proj
        self.down_proj = original_mlp.down_proj
        self.act_fn = original_mlp.act_fn
def forward(self, x):
    gate_output = self.gate_proj(x)
    up_output = self.up_proj(x)
    activated = self.act_fn(gate_output)
    # CRITICAL: Compute in float32 to avoid overflow
    intermediate = (activated.float() * up_output.float()).to(x.dtype)
    return self.down_proj(intermediate)

With this patch applied:

AWQ logits range: [-29.33, 13.83]
AWQ PAD token rank: #110,051 (correct - not predicted)
Top prediction: '█' with logit 13.83

3. Int8 quantization works fine:

The same Gemma3 model converted with int8 quantization produces correct output, suggesting the issue is specific to the AWQ float16 enforcement.

Affected Files

1. ops/awq/gemm.cc - Enforces float16 input type
2. ops/awq/gemv.cc - Same float16 enforcement
3. FFN/MLP computation (location unknown) - The silu(gate) * up multiplication needs float32 intermediate computation

Proposed Solution

The fix requires computing critical operations in float32 before casting back to float16:

Option A: Model-specific flag for Gemma3

Add a flag in the Gemma3 spec to enable float32 intermediate computations for models with large layernorm weights.

Option B: Automatic detection

Detect when layernorm weights exceed a threshold (e.g., abs > 100) and automatically use float32 for intermediate computations.

Option C: Modify AWQ kernels

Update the AWQ GEMM/GEMV kernels to support float32 accumulation and output, similar to how the RMSNorm kernel already computes internally in float32:

// Current RMSNorm kernel already uses float32 internally:
float sum_squares = 0;
for (cuda::index_t i = threadIdx.x; i < depth; i += blockDim.x)
    sum_squares += float(input[i]) * float(input[i]);  // float32

The same pattern should be applied to the FFN multiplication.

Environment

• CTranslate2 version: 4.6.2
• PyTorch version: 2.9.0+cu128
• CUDA: 12.8
• GPU: NVIDIA RTX 4090
• Model: gemma-3-1b-it quantized to AWQ (4-bit, group_size=128, version=gemm)

---

Start by reproducing the Gemma3 AWQ failure described in the issue and compare it with the working int8 case. Read ops/awq/gemm.cc and ops/awq/gemv.cc, then locate the FFN/MLP computation where silu(gate) * up is performed. Done means Gemma3 AWQ no longer emits pad-only output or -65504.0 scores while preserving the existing AWQ path.