Custom dtype - bf16/fp16 support

QEfficient/base/modeling_qeff.py

@@ -1191,8 +1191,10 @@ def _compile(
             command.append("-sub-functions")
 
         model_in_bfloat16 = hasattr(self, "config") and (self.config.torch_dtype == torch.bfloat16)
+        io_name_prefix = ("past_", "pixel_values", "conv_", "recurrent_")
         pkv_in_bfloat16 = (custom_io is not None) and any(
-            ("past_" in key or "pixel_values" in key) and "bfloat16" in value for key, value in custom_io.items()
+            any(bfloat16_io_name in key for bfloat16_io_name in io_name_prefix) and "bfloat16" in value
+            for key, value in custom_io.items()
         )
         custom_io_for_compiler = custom_io if not (model_in_bfloat16 and pkv_in_bfloat16) else None
 

QEfficient/transformers/models/gemma4/modeling_gemma4.py

@@ -323,7 +323,9 @@ def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
             weight = getattr(self, "_qeff_unit_weight", None)
             if weight is None:
                 weight = hidden_states.new_ones(hidden_states.shape[-1])
-        return CustomRMSNormFunc.apply(hidden_states, weight, self.eps)
+        # Cast weight to match hidden_states dtype so the AIC compiler sees
+        # matching Input/Scale dtypes in the CustomRMSNorm op (e.g. bfloat16).
+        return CustomRMSNormFunc.apply(hidden_states, weight.to(hidden_states.dtype), self.eps)
 
 
 class QEffGemma4TextMoeBlock(QEffMoEBlockMixin, nn.Module):
@@ -476,6 +478,7 @@ def forward(
     ) -> tuple[torch.Tensor, torch.Tensor | None]:
         input_shape = hidden_states.shape[:-1]
         hidden_shape = (*input_shape, -1, self.head_dim)
+        target_dtype = hidden_states.dtype
         cache_kwargs = {"position_ids": position_ids, "batch_index": batch_index}
         token_key_states = None
         token_value_states = None
@@ -506,8 +509,8 @@ def forward(
 
         if self.is_kv_shared_layer and past_key_values is not None:
             key_states, value_states = past_key_values.shared_layers[self.kv_shared_layer_index]
-            key_states = key_states.to(query_states.device)
-            value_states = value_states.to(query_states.device)
+            key_states = key_states.to(query_states.device, dtype=target_dtype)
+            value_states = value_states.to(query_states.device, dtype=target_dtype)
             if hasattr(past_key_values, "shared_layers_token"):
                 token_states = past_key_values.shared_layers_token.get(self.kv_shared_layer_index)
                 if token_states is not None:
@@ -519,9 +522,13 @@ def forward(
             key_states = self.k_norm(key_states)
             key_states = qeff_apply_rotary_pos_emb(key_states, cos, sin)
             key_states = key_states.transpose(1, 2)
+            if key_states.dtype != target_dtype:
+                key_states = key_states.to(target_dtype)
 
             value_states = self.v_norm(value_states)
             value_states = value_states.transpose(1, 2)
+            if value_states.dtype != target_dtype:
+                value_states = value_states.to(target_dtype)
             token_key_states, token_value_states = key_states, value_states
 
         if use_blocking:
@@ -719,6 +726,10 @@ def forward(
         if input_ids is not None:
             inputs_embeds = self.embed_tokens(input_ids)
 
+        target_dtype = self.embed_tokens.weight.dtype
+        if inputs_embeds is not None and inputs_embeds.dtype != target_dtype:
+            inputs_embeds = inputs_embeds.to(target_dtype)
+
         if self.hidden_size_per_layer_input:
             if per_layer_inputs is None:
                 per_layer_inputs = self.get_per_layer_inputs(input_ids, inputs_embeds)
@@ -1072,8 +1083,8 @@ def get_dummy_pkv_cache(self, config, batch_size, seq_len):
             cache_shape = [batch_size, n_heads, layer_seq_len, d_head]
             past_key_values.append(
                 (
-                    torch.zeros(cache_shape, dtype=torch.float32),
-                    torch.zeros(cache_shape, dtype=torch.float32),
+                    torch.zeros(cache_shape, dtype=config.dtype),
+                    torch.zeros(cache_shape, dtype=config.dtype),
                 )
             )
         return past_key_values
@@ -1156,6 +1167,9 @@ def forward(
         llm_input_ids = input_ids.clone()
         llm_input_ids[special_image_mask] = self.config.text_config.pad_token_id
         inputs_embeds = self.model.get_input_embeddings()(llm_input_ids)
+        target_dtype = self.language_model.embed_tokens.weight.dtype
+        if inputs_embeds.dtype != target_dtype:
+            inputs_embeds = inputs_embeds.to(target_dtype)
 
         next_image_idx = image_idx
         if input_ids.shape[1] != 1 and special_image_mask.any() and vision_embeds is None:
@@ -1174,6 +1188,8 @@ def forward(
             indices0 = torch.arange(special_image_mask.shape[0], device=special_image_mask.device).view(-1, 1)
             safe_indices1 = torch.where(indices1 < 0, torch.zeros_like(indices1), indices1)
             gathered_vision_embeds = vision_embeds[indices0, safe_indices1]
+            if gathered_vision_embeds.dtype != target_dtype:
+                gathered_vision_embeds = gathered_vision_embeds.to(target_dtype)
             inputs_embeds = torch.where(special_image_mask.unsqueeze(-1), gathered_vision_embeds, inputs_embeds)
             next_image_idx = (indices1.max() + 1).reshape(1, 1)
 
@@ -1453,8 +1469,8 @@ def get_dummy_pkv_cache(self, config, batch_size, seq_len):
             cache_shape = [batch_size, n_heads, layer_seq_len, d_head]
             past_key_values.append(
                 (
-                    torch.zeros(cache_shape, dtype=torch.float32),
-                    torch.zeros(cache_shape, dtype=torch.float32),
+                    torch.zeros(cache_shape, dtype=config.dtype),
+                    torch.zeros(cache_shape, dtype=config.dtype),
                 )
             )
         return past_key_values
@@ -1496,12 +1512,14 @@ def get_dummy_inputs(
         mm_token_type_ids[:, image_start:image_end] = 1
 
         vision_inputs = {
-            "pixel_values": torch.zeros((bs, max_patches, patch_dim), dtype=torch.float32),
+            "pixel_values": torch.zeros((bs, max_patches, patch_dim), dtype=self.config.dtype),
             "image_position_ids": image_position_ids,
         }
         lang_inputs = {
             "input_ids": input_ids,
-            "vision_embeds": torch.zeros((bs, mm_tokens_per_image, self.model.language_model.config.hidden_size)),
+            "vision_embeds": torch.zeros(
+                (bs, mm_tokens_per_image, self.model.language_model.config.hidden_size), dtype=self.config.dtype
+            ),
             "position_ids": torch.arange(seq_len, dtype=torch.int64).view(1, seq_len).repeat(bs, 1),
             "image_idx": torch.zeros((1, 1), dtype=torch.int64),
             "mm_token_type_ids": mm_token_type_ids,

QEfficient/transformers/models/qwen3_5/modeling_qwen3_5.py

@@ -65,9 +65,11 @@ class QEffQwen3_5GatedDeltaNetCustomRMSNormAIC(nn.Module):
     def forward(self, hidden_states, gate):
         return (
             CustomRMSNormFunc.apply(
-                hidden_states, self.weight, self.variance_epsilon if hasattr(self, "variance_epsilon") else self.eps
+                hidden_states,
+                self.weight.to(hidden_states.dtype),
+                self.variance_epsilon if hasattr(self, "variance_epsilon") else self.eps,
             )
-        ) * F.silu(gate.to(torch.float32))
+        ) * F.silu(gate.to(hidden_states.dtype))
 
 
 class QEffQwen3_5DynamicCache(Cache):
@@ -355,10 +357,10 @@ def eager_attention_forward(
     attn_weights = torch.matmul(query, key_states.transpose(2, 3)) * scaling
     if attention_mask is not None:
         attn_weights = torch.where(
-            attention_mask, torch.tensor(MIN_MASKED_ATTENTION_VALUE, dtype=torch.float32), attn_weights
+            attention_mask, torch.tensor(MIN_MASKED_ATTENTION_VALUE, dtype=module.config.torch_dtype), attn_weights
         )
 
-    attn_weights = nn.functional.softmax(attn_weights, dim=-1, dtype=torch.float32).to(query.dtype)
+    attn_weights = nn.functional.softmax(attn_weights, dim=-1, dtype=module.config.torch_dtype).to(query.dtype)
     attn_output = torch.matmul(attn_weights, value_states)
     attn_output = attn_output.transpose(1, 2).contiguous()
     return attn_output, attn_weights
@@ -492,6 +494,7 @@ class QEffQwen3_5GatedDeltaNet(Qwen3_5GatedDeltaNet):
 
     def __qeff_init__(self):
         self.chunk_gated_delta_rule = self.torch_chunk_gated_delta_rule_qeff
+        self.torch_dtype = self.out_proj.weight.dtype
         chunk_size = 64  # must match what's used in the function
 
         # Precompute all constant masks — no triu/tril with diagonal args at runtime
@@ -552,7 +555,7 @@ def torch_chunk_gated_delta_rule_qeff(
             query = query * torch.rsqrt(torch.einsum("bthd,bthd->bth", query, query).unsqueeze(-1) + 1e-6)
             key = key * torch.rsqrt(torch.einsum("bthd,bthd->bth", key, key).unsqueeze(-1) + 1e-6)
         query, key, value, beta, g = [
-            x.transpose(1, 2).contiguous().to(torch.float32) for x in (query, key, value, beta, g)
+            x.transpose(1, 2).contiguous().to(self.torch_dtype) for x in (query, key, value, beta, g)
         ]
 
         mask = (position_ids[0] != -1).unsqueeze(1)
@@ -570,22 +573,21 @@ def torch_chunk_gated_delta_rule_qeff(
         batch_size, num_heads, sequence_length, k_head_dim = key.shape
         v_head_dim = value.shape[-1]
         pad_size = (chunk_size - sequence_length % chunk_size) % chunk_size
-        # query = F.pad(query, (0, 0, 0, pad_size))
-        # key = F.pad(key, (0, 0, 0, pad_size))
-        # value = F.pad(value, (0, 0, 0, pad_size))
-        # beta = F.pad(beta, (0, pad_size))
-
-        # # ck = g.clone()
-        # g = F.pad(g, (0, pad_size))
-        query = F.pad(query, (0, 0, 0, pad_size), mode="constant", value=0.0)
-        key = F.pad(key, (0, 0, 0, pad_size), mode="constant", value=0.0)
-        value = F.pad(value, (0, 0, 0, pad_size), mode="constant", value=0.0)
-        beta = F.pad(beta, (0, pad_size), mode="constant", value=0.0)
-
-        # ck = g.clone()
-        g = F.pad(g, (0, pad_size), mode="constant", value=0.0)
+        # QAIC's LoadPad kernel only supports float32/int32/int64 inputs, so bf16/fp16
+        # tensors are padded via torch.cat with a zero tensor instead of F.pad.
+        query = torch.cat(
+            [query, torch.zeros(*query.shape[:2], pad_size, query.shape[3], dtype=query.dtype)],
+            dim=2,
+        )
+        key = torch.cat([key, torch.zeros(*key.shape[:2], pad_size, key.shape[3], dtype=key.dtype)], dim=2)
+        value = torch.cat(
+            [value, torch.zeros(*value.shape[:2], pad_size, value.shape[3], dtype=value.dtype)],
+            dim=2,
+        )
+        beta = torch.cat([beta, torch.zeros(*beta.shape[:2], pad_size, dtype=beta.dtype)], dim=2)
+        g = torch.cat([g, torch.zeros(*g.shape[:2], pad_size, dtype=g.dtype)], dim=2)
         total_sequence_length = sequence_length + pad_size
-        scale = 1 / (query.shape[-1] ** 0.5)
+        scale = 1 / (self.head_k_dim**0.5)
         query = query * scale
 
         v_beta = value * beta.unsqueeze(-1)
@@ -614,7 +616,7 @@ def torch_chunk_gated_delta_rule_qeff(
         diff = g.unsqueeze(-1) - g.unsqueeze(-2)  # (B, H, num_chunks, C, C)
         diff = diff * (~mask_strict).float()  # zero upper triangle (strict)
         decay_mask = diff.exp().float()
-        decay_mask = decay_mask * (~mask_strict).float()  # ensure upper is zero
+        decay_mask = (decay_mask * (~mask_strict).float()).to(self.torch_dtype)  # ensure upper is zero
 
         attn = -((k_beta @ key.transpose(-1, -2)) * decay_mask).masked_fill(mask, 0)
         for i in range(1, chunk_size):
@@ -694,6 +696,8 @@ def torch_chunk_gated_delta_rule_qeff(
         )
         core_attn_out = core_attn_out[:, :, :sequence_length]
         core_attn_out = core_attn_out.transpose(1, 2).contiguous().to(initial_dtype)
+        if last_recurrent_state is not None:
+            last_recurrent_state = last_recurrent_state.to(initial_dtype)
         return core_attn_out, last_recurrent_state
 
     def _recurrent_step_batched(self, query, key, value, g, beta, recurrent_state):
@@ -713,7 +717,7 @@ def _recurrent_step_batched(self, query, key, value, g, beta, recurrent_state):
         k = k * torch.rsqrt(torch.einsum("bthd,bthd->bth", k, k).unsqueeze(-1) + 1e-6)
         v = value.float()
 
-        scale = 1.0 / (q.shape[-1] ** 0.5)
+        scale = 1.0 / (self.head_k_dim**0.5)
         q = q * scale  # (B, T, H, d_k)
 
         # For T=1 decode, this is a single step
@@ -1092,8 +1096,8 @@ def get_onnx_retained_state_specs(
             if layer_type == "full_attention":
                 layer_names = [f"past_key.{layer_idx}", f"past_value.{layer_idx}"]
                 layer_tensors = [
-                    torch.zeros(tuple(kv_cache_shape), dtype=torch.float32),
-                    torch.zeros(tuple(kv_cache_shape), dtype=torch.float32),
+                    torch.zeros(tuple(kv_cache_shape), dtype=self.config.torch_dtype),
+                    torch.zeros(tuple(kv_cache_shape), dtype=self.config.torch_dtype),
                 ]
                 layer_axes = [
                     {0: batch_axis_name, 2: "ctx_len"},
@@ -1105,8 +1109,8 @@ def get_onnx_retained_state_specs(
                 recurrent_shape = (batch_size, layer.num_v_heads, layer.head_k_dim, layer.head_v_dim)
                 layer_names = [f"conv_state.{layer_idx}", f"recurrent_state.{layer_idx}"]
                 layer_tensors = [
-                    torch.zeros(conv_shape, dtype=torch.float32),
-                    torch.zeros(recurrent_shape, dtype=torch.float32),
+                    torch.zeros(conv_shape, dtype=self.config.torch_dtype),
+                    torch.zeros(recurrent_shape, dtype=self.config.torch_dtype),
                 ]
                 layer_axes = [{0: batch_axis_name}, {0: batch_axis_name}]
 
@@ -1395,7 +1399,7 @@ def forward(
         q, k = apply_rotary_pos_emb_vision(q, k, cos, sin)
 
         attention_mask = torch.full(
-            [1, seq_length, seq_length], torch.finfo(q.dtype).min, device=q.device, dtype=q.dtype
+            [1, seq_length, seq_length], MIN_MASKED_ATTENTION_VALUE, device=q.device, dtype=q.dtype
         )
         seq_len = attention_mask.shape[-1]
         rows = torch.arange(seq_len).view(1, -1)
@@ -1407,17 +1411,17 @@ def forward(
         col_mask = (cols >= start) & (cols < end)
         block_mask = row_mask & col_mask
 
-        final_mask = torch.ones((seq_len, seq_len), dtype=torch.float32)
+        final_mask = torch.ones((seq_len, seq_len), dtype=self.config.torch_dtype)
         final_mask[block_mask.any(dim=0)] = 0
-        final_mask = torch.where(final_mask == 1.0, torch.finfo(q.dtype).min, final_mask)
+        final_mask = torch.where(final_mask == 1.0, MIN_MASKED_ATTENTION_VALUE, final_mask)
         attention_mask[0] = final_mask
 
         q = q.transpose(0, 1)
         k = k.transpose(0, 1)
         v = v.transpose(0, 1)
         attn_weights = torch.matmul(q, k.transpose(1, 2)) / math.sqrt(self.head_dim)
         attn_weights = attn_weights + attention_mask
-        attn_weights = nn.functional.softmax(attn_weights, dim=-1, dtype=torch.float32).to(q.dtype)
+        attn_weights = nn.functional.softmax(attn_weights, dim=-1, dtype=self.config.torch_dtype).to(q.dtype)
         attn_output = torch.matmul(attn_weights, v)
         attn_output = attn_output.transpose(0, 1)
         attn_output = attn_output.reshape(seq_length, -1)
@@ -1806,10 +1810,10 @@ def get_dummy_inputs(
 
         vision_inputs = {}
         lang_inputs = {}
-        vision_inputs["pixel_values"] = torch.zeros((inputs_shapes["pixel_values"]), dtype=torch.float32)
+        vision_inputs["pixel_values"] = torch.zeros((inputs_shapes["pixel_values"]), dtype=self.config.torch_dtype)
         vision_inputs["image_grid_thw"] = torch.zeros((inputs_shapes["image_grid_thw"]), dtype=torch.int64)
         lang_inputs["input_ids"] = torch.zeros((inputs_shapes["input_ids"]), dtype=torch.int64)
-        lang_inputs["vision_embeds"] = torch.zeros((inputs_shapes["vision_embeds"]), dtype=torch.float32)
+        lang_inputs["vision_embeds"] = torch.zeros((inputs_shapes["vision_embeds"]), dtype=self.config.torch_dtype)
         lang_inputs["position_ids"] = (
             (
                 torch.arange(dummy_seq_len, dtype=torch.int64)
@@ -1836,13 +1840,13 @@ def get_dummy_inputs(
         for i in range(self.model.config.text_config.num_hidden_layers):
             if self.model.config.text_config.layer_types[i] == "full_attention":
                 for kv in ["key", "value"]:
-                    lang_inputs["past_key_values"][i].append(torch.zeros(kv_cache_shape, dtype=torch.float32))
+                    lang_inputs["past_key_values"][i].append(torch.zeros(kv_cache_shape, dtype=self.config.torch_dtype))
             else:
                 layer = self.model.language_model.layers[i].linear_attn
                 conv_shape = (linear_batch_size, layer.conv_dim, layer.conv_kernel_size)
                 recurrent_shape = (linear_batch_size, layer.num_v_heads, layer.head_k_dim, layer.head_v_dim)
-                lang_inputs["past_key_values"][i].append(torch.zeros(conv_shape, dtype=torch.float32))
-                lang_inputs["past_key_values"][i].append(torch.zeros(recurrent_shape, dtype=torch.float32))
+                lang_inputs["past_key_values"][i].append(torch.zeros(conv_shape, dtype=self.config.torch_dtype))
+                lang_inputs["past_key_values"][i].append(torch.zeros(recurrent_shape, dtype=self.config.torch_dtype))
 
         #
         if continuous_batching:
@@ -1885,7 +1889,7 @@ def get_inputs_info(self):
         return [
             IOInfo(name="input_ids", datatype=torch.int64, shape=("batch_size", "seq_len")),
             IOInfo(name="attention_mask", datatype=torch.int64, shape=("batch_size", "seq_len")),
-            # IOInfo(name="pixel_values", datatype=torch.float32, shape=("batch_size", 3, "image_size", "image_size")),
+            # IOInfo(name="pixel_values", datatype=self.config.torch_dtype, shape=("batch_size", 3, "image_size", "image_size")),
         ]
 
     def prepare_inputs_for_generation(self, inputs, prefill_seq_len=32, batch_size=1):