perf: fp16 autocast, parallel Redis KNN, and Docker runtime fixes - #14
Merged
Conversation
Speed up inference throughput on two fronts: - process_vits: run the ViT forward pass under fp16 autocast (CUDA) with model.eval() and torch.inference_mode(), non-blocking host->device copies, and an optional torch.compile path. Embeddings are cast back to float32 to match the FLOAT32 Redis index. Toggle via VSS_AMP / VSS_TORCH_COMPILE. - vector_similarity: add search_vectors() to issue a batch of KNN queries concurrently over the connection pool instead of one serial round-trip per image; predict() now uses it. Concurrency tunable via VSS_KNN_WORKERS. Co-authored-by: Cursor <cursoragent@cursor.com>
torch._dynamo (imported transitively by transformers) calls getpass.getuser() at import time, which raises KeyError when the container runs as an arbitrary uid with no /etc/passwd entry (e.g. compose `user: 12078:20`). Set USER/HOME so getpass short-circuits on the env var, and route torch/HF caches to writable /tmp so the optional torch.compile path works for any uid. Co-authored-by: Cursor <cursoragent@cursor.com>
Switch the runtime image to the CUDA devel base (plus python3.11-dev) so the Inductor/Triton backend used by VSS_TORCH_COMPILE=1 has the gcc/nvcc toolchain and headers it JIT-compiles against, fixing the "Failed to find C compiler" error. Also catch lazy torch.compile failures on the first forward and fall back to eager execution so a compile problem no longer fails every prediction. Co-authored-by: Cursor <cursoragent@cursor.com>
The devel image sets LD_LIBRARY_PATH=/usr/local/cuda/lib64, which shadows the CUDA libraries bundled with the pip torch wheel and caused CUBLAS_STATUS_NOT_INITIALIZED (cublasLtMatmulAlgoGetHeuristic) at runtime. Revert to the lightweight base image and install just gcc/g++ so the optional VSS_TORCH_COMPILE (Triton) path still has a C/C++ compiler without dragging in the conflicting system CUDA libs. Co-authored-by: Cursor <cursoragent@cursor.com>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
VSS_AMP, default on for CUDA),model.eval(), andtorch.inference_mode(); embeddings are cast back to float32 for the Redis index.search_vectors()instead of one serial round-trip per image (VSS_KNN_WORKERS, default 16).VSS_TORCH_COMPILE: uid-agnostic cache env vars, gcc/g++ on the CUDA base image (avoid devel LD_LIBRARY_PATH cuBLAS conflicts), and eager fallback if lazy compile fails on first forward.Test plan
user: 12078:20(nogetpwuid/ cuBLAS init errors)./knnand/embedwith default env and withVSS_TORCH_COMPILE=1.Env knobs
VSS_AMP=1(default) — fp16 autocast on CUDAVSS_KNN_WORKERS=16(default) — concurrent KNN queries per batchVSS_TORCH_COMPILE=0(default) — optionaltorch.compile; requires gcc/g++ in image