Spaces:
Running on Zero
Running on Zero
atakan Claude Sonnet 5 commited on
Commit ·
72bc69e
1
Parent(s): 1498459
perf: Use SDPA attention on CUDA instead of the eager default
Browse filesMeaningful 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>
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.
|