diff --git a/studio/backend/core/inference/diffusion_precision.py b/studio/backend/core/inference/diffusion_precision.py index 68036aae8d..030b6da7f2 100644 --- a/studio/backend/core/inference/diffusion_precision.py +++ b/studio/backend/core/inference/diffusion_precision.py @@ -120,7 +120,10 @@ def _cast_fp8(encoder: Any, target: Any) -> None: # gets cast to fp8 and, sharing one tensor, drags the embedding to fp8 with it. The # embedding then emits fp8 activations that crash the first RMSNorm. Skip the tied # projection so the shared tensor stays dense (lm_head is unused for prompt encoding). - get_out, get_in = getattr(encoder, "get_output_embeddings", None), getattr(encoder, "get_input_embeddings", None) + get_out, get_in = ( + getattr(encoder, "get_output_embeddings", None), + getattr(encoder, "get_input_embeddings", None), + ) out_emb = get_out() if callable(get_out) else None in_emb = get_in() if callable(get_in) else None if out_emb is not None and in_emb is not None and out_emb.weight is in_emb.weight: diff --git a/studio/backend/core/inference/diffusion_transformer_quant.py b/studio/backend/core/inference/diffusion_transformer_quant.py index 6e1341983f..7cb1074425 100644 --- a/studio/backend/core/inference/diffusion_transformer_quant.py +++ b/studio/backend/core/inference/diffusion_transformer_quant.py @@ -257,6 +257,7 @@ def _make_quant_config(scheme: str, fast_accum: Optional[bool] = None) -> Any: return Float8DynamicActivationFloat8WeightConfig() if scheme == TQ_NVFP4: from torchao.prototype.mx_formats import NVFP4DynamicActivationNVFP4WeightConfig + # Select the CUTLASS FP4 path, not the default Triton kernel: torchao defaults # use_triton_kernel=True, which needs MSLK installed. On a Blackwell box with the # CUTLASS FP4 extension but no MSLK, the default would make the smoke probe fail