Update on "Use unfused SDPA for short sequences (q_len <= 128 or kv_len <= 128)"

ATT

Differential Revision: [D96044308](https://our.internmc.facebook.com/intern/diff/D96044308/)

[ghstack-poisoned]

Author: kimishpatel

.ci/scripts/test_lora.sh

@@ -133,14 +133,14 @@ else
 fi
 
 ### QUANTIZATION & PROGRAM DATA SEPARATION ###
-EXPECTED_QUANT_PREFIX="<|im_start|>user Calculate 15% of 80?<|im_end|><|im_start|>assistant:
+EXPECTED_QUANT_PREFIX="<|im_start|>user Calculate 15% of 80?<|im_end|><|im_start|>assistant: me
 <think>
 <think>
 Okay, so I need to calculate 15% of 80."
 EXPECTED_QUANT_LORA_PREFIX="
 <|im_start|>user Calculate 15% of 80?<|im_end|><|im_start|>assistant
 To calculate 15% of 80, we can multiply 80 by 15/100.
-80 * 0.15 = 12.
+80 * 15/100 = 12.
 So, 15% of 80 is 12.
 #### 12
 The answer is: 12<|im_end|>"

.ci/scripts/test_lora_multimethod.sh

@@ -67,7 +67,7 @@ EXPECTED_LORA_PREFIX="
 <|im_start|>user Calculate 15% of 80?<|im_end|><|im_start|>assistant
 To calculate 15% of 80"
 
-EXPECTED_BASE_PREFIX="<|im_start|>user Calculate 15% of 80?<|im_end|><|im_start|>assistant:
+EXPECTED_BASE_PREFIX="<|im_start|>user Calculate 15% of 80?<|im_end|><|im_start|>assistant: me
 <think>
 <think>
 Okay, so I need to calculate 15% of 80."

extension/llm/custom_ops/op_sdpa.cpp

@@ -416,7 +416,8 @@ Tensor& custom_sdpa_out_impl(
   // correctly when seq_dim=ONE and seq_len > 1, so keep the conservative
   // condition for quantized inputs.
   bool is_quantized = q.scalar_type() == ScalarType::Char;
-  bool use_unfused_sdpa = (!is_quantized) && (seq_len <= 128 || num_keys_for_causal_attention <= 128);
+  bool use_unfused_sdpa = (!is_quantized) &&
+      (seq_len <= 128 || num_keys_for_causal_attention <= 128);
   if (use_unfused_sdpa) {
     ET_SWITCH_FLOAT_TYPES(output.scalar_type(), ctx, "sdpa", CTYPE, [&] {
       sdpa::impl::cpu_sdpa<CTYPE>(