jalr/tabbyAPI-ollama
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

Main: Enable cudaMallocAsync backend by default

Browse source 
Works on cuda 12.4 and up. If CUDA doesn't exist, then don't enable
the backend. This is an env var that needs to be set, so it's not really
possible to set it via config.yml.

This used to be experimental, but it's probably fine to keep it enabled
since it only provides a benefit.

Signed-off-by: kingbri <8082010+kingbri1@users.noreply.github.com>
This commit is contained in:
kingbri 2025-07-27 22:29:46 -04:00
parent 0b4ca567f8
commit 113643c0df
2 changed files with 5 additions and 7 deletions
Download patch file Download diff file Expand all files Collapse all files

1
docs/02.-Server-options.md
View file

@ -47,7 +47,6 @@ Note: These are experimental flags that may be removed at any point.
| ------------------------- | -------------- | ----------------------------------------------------------------------------------------------------------------------------------------------- |
| unsafe_launch | Bool (False) | Skips dependency checks on startup. Only recommended for debugging. |
| disable_request_streaming | Bool (False) | Forcefully disables streaming requests |
| cuda_malloc_backend | Bool (False) | Uses pytorch's CUDA malloc backend to load models. Helps save VRAM.<br><br>Safe to enable. |
| realtime_process_priority | Bool (False) | Set the process priority to "Realtime". Administrator/sudo access required, otherwise the priority is set to the highest it can go in userland. |
### Model Options

11
main.py
View file

@ -1,8 +1,12 @@
"""The main tabbyAPI module. Contains the FastAPI server and endpoints."""
# Set this env var for cuda malloc async before torch is initalized
import os
os.environ["PYTORCH_CUDA_ALLOC_CONF"] = "backend:cudaMallocAsync"
import argparse
import asyncio
import os
import pathlib
import platform
import signal
@ -154,11 +158,6 @@ def entrypoint(
raise SystemExit(install_message)
# Enable CUDA malloc backend
if config.developer.cuda_malloc_backend:
os.environ["PYTORCH_CUDA_ALLOC_CONF"] = "backend:cudaMallocAsync"
logger.warning("EXPERIMENTAL: Enabled the pytorch CUDA malloc backend.")
# Set the process priority
if config.developer.realtime_process_priority:
import psutil