From 05bc0b8f673dad69dba4681ed4ccfe37d0134b03 Mon Sep 17 00:00:00 2001 From: Gregory Comer Date: Wed, 8 Jul 2026 16:58:49 -0700 Subject: [PATCH] Allow f16/bf16 custom_sdpa --- optimum/executorch/attentions/custom_kv_cache.py | 2 +- optimum/executorch/attentions/custom_sdpa.py | 4 ---- 2 files changed, 1 insertion(+), 5 deletions(-) diff --git a/optimum/executorch/attentions/custom_kv_cache.py b/optimum/executorch/attentions/custom_kv_cache.py index 64b7322..515d1c7 100644 --- a/optimum/executorch/attentions/custom_kv_cache.py +++ b/optimum/executorch/attentions/custom_kv_cache.py @@ -174,7 +174,7 @@ def from_legacy_cache( dtype = legacy_cache.k_cache.dtype # assert device is None or device == "cpu" - assert dtype is None or dtype == torch.float32 + assert dtype is None or dtype in (torch.float32, torch.bfloat16, torch.float16) # Use the legacy cache's max_seq_len if max_cache_len is not specified if max_cache_len is None and hasattr(legacy_cache, "max_seq_len"): diff --git a/optimum/executorch/attentions/custom_sdpa.py b/optimum/executorch/attentions/custom_sdpa.py index 0f5d0fc..d59e418 100644 --- a/optimum/executorch/attentions/custom_sdpa.py +++ b/optimum/executorch/attentions/custom_sdpa.py @@ -81,11 +81,7 @@ def custom_sdpa_with_start_pos_forward( key = key.transpose(1, 2) value = value.transpose(1, 2) - # Convert the hell out of the inputs to fp32 and back input_dtype = query.dtype - query = query.to(torch.float32) - key = key.to(torch.float32) - value = value.to(torch.float32) # Ignore the causal flag from kwargs but use the one in module kwargs.pop("is_causal", None)