diff --git a/comfy/ops.py b/comfy/ops.py index d9df909bb..d53886de0 100644 --- a/comfy/ops.py +++ b/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) diff --git a/comfy/text_encoders/gemma4.py b/comfy/text_encoders/gemma4.py index 96a547b60..ee035d522 100644 --- a/comfy/text_encoders/gemma4.py +++ b/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