atakan Claude Sonnet 5 commited on
Commit
72bc69e
·
1 Parent(s): 1498459

perf: Use SDPA attention on CUDA instead of the eager default

Browse files

Meaningful speedup for prefilling the long tool-schema system prompt
on ZeroGPU, at no cost -- SDPA is a drop-in faster attention backend
transformers already ships, no extra dependency needed.

Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>

Files changed (1) hide show
  1. controlai_agent/orchestrator.py +4 -0
controlai_agent/orchestrator.py CHANGED
@@ -221,12 +221,16 @@ class ControlAIAgent:
221
  print(f"Loading PyTorch model: {hf_id} (threads: {num_threads})...")
222
  self.hf_tokenizer = AutoTokenizer.from_pretrained(hf_id, trust_remote_code=True)
223
  dtype = torch.bfloat16 if torch.cuda.is_available() else torch.float32
 
 
 
224
  self.model = AutoModelForCausalLM.from_pretrained(
225
  hf_id,
226
  torch_dtype=dtype,
227
  low_cpu_mem_usage=True,
228
  device_map="auto",
229
  trust_remote_code=True,
 
230
  )
231
  # The fused ControlAI model already has the LoRA weights merged in;
232
  # only apply a separate adapter when loading a plain base model.
 
221
  print(f"Loading PyTorch model: {hf_id} (threads: {num_threads})...")
222
  self.hf_tokenizer = AutoTokenizer.from_pretrained(hf_id, trust_remote_code=True)
223
  dtype = torch.bfloat16 if torch.cuda.is_available() else torch.float32
224
+ # SDPA is a large, free speedup over the "eager" attention
225
+ # default -- meaningful for the long tool-schema prompt prefix.
226
+ attn_impl = "sdpa" if torch.cuda.is_available() else None
227
  self.model = AutoModelForCausalLM.from_pretrained(
228
  hf_id,
229
  torch_dtype=dtype,
230
  low_cpu_mem_usage=True,
231
  device_map="auto",
232
  trust_remote_code=True,
233
+ attn_implementation=attn_impl,
234
  )
235
  # The fused ControlAI model already has the LoRA weights merged in;
236
  # only apply a separate adapter when loading a plain base model.