Fix quantized text encoder matmul gating and Gemma4 prefill cache positions (#16185)

comfy/ops.py

@@ -1342,11 +1342,8 @@ def mixed_precision_ops(quant_config={}, compute_dtype=torch.bfloat16, full_prec
                         compute_dtype=compute_dtype,
                         want_requant=want_requant,
                     ) as (weight, bias):
-                        if self._full_precision_mm and isinstance(weight, QuantizedTensor):
-                            # cast_bias_weight only dequantizes on a dtype change, which is a
-                            # no-op here when the quantized weight's orig_dtype already equals
-                            # the compute dtype. Force it so the disabled/unsupported-format
-                            # fallback doesn't hand a QuantizedTensor to a plain linear() call.
+                        if isinstance(weight, QuantizedTensor) and (self._full_precision_mm_config or getattr(self, "quant_format", None) in self._disabled_formats or not weight.layout_cls.supports_fast_matmul()):
+                            # explicit per-layer full precision, or a format this device can't run: don't reach the fast quantized matmul
                             weight = weight.dequantize()
                         return self._forward(input, weight, bias)
 

comfy/text_encoders/gemma4.py

@@ -528,17 +528,18 @@ class Gemma4Transformer(nn.Module):
                         and comfy.model_management.is_device_cuda(x.device))
         decode_bias = None
         decode_masks = None
-        if decode:
+        if fixed_kv:
+            # prefill must advance the device-side write position of the global caches too
             prepared = set()
             for kv in past_key_values:
                 if isinstance(kv, FixedKV) and id(kv.position) not in prepared:
                     kv.prepare(seq_len)
                     prepared.add(id(kv.position))
-            if mask is not None:
-                decode_masks = {}
-                for kv in past_key_values:
-                    if isinstance(kv, FixedKV) and id(kv.position) not in decode_masks:
-                        decode_masks[id(kv.position)] = _fixed_kv_decode_mask(mask, kv, min_val)
+        if decode and mask is not None:
+            decode_masks = {}
+            for kv in past_key_values:
+                if isinstance(kv, FixedKV) and id(kv.position) not in decode_masks:
+                    decode_masks[id(kv.position)] = _fixed_kv_decode_mask(mask, kv, min_val)
         if compiled_decode:
             capacities = tuple(sorted({kv.key.shape[2] for kv in past_key_values if isinstance(kv, FixedKV)}))
             valid = past_len + 1