Distributed Training and Inference Engineer
Sciforium· Engineering
San Francisco, CA$190k - $250kFull-timePosted 8mo ago
About the role
Sciforium is seeking a Distributed Training and Inference Engineer to build, optimize, and maintain the software stack for large-scale AI training and serving workloads, spanning from CUDA/ROCm runtimes to high-level frameworks like JAX and PyTorch to enable fast, scalable, and reliable distributed training and inference of next-generation models.
Responsibilities
- Software Stack Maintenance: Maintain, update, and optimize critical ML libraries and frameworks including JAX, PyTorch, CUDA, and ROCm across multiple environments and hardware configurations.
- End-to-End Stack Ownership: Build, maintain, and continuously improve the entire ML software stack from ROCm/CUDA drivers to high-level JAX/PyTorch tooling.
- Distributed System Optimization: Ensure all model implementations are efficiently sharded, partitioned, and configured for large-scale distributed training and serving.
- System Integration: Continuously integrate and validate modules for runtime correctness, memory efficiency, and scalability across multi-node GPU/accelerator clusters.
- Profiling & Performance Analysis: Conduct detailed profiling of compilation graphs, training workloads, and runtime execution to optimize performance and eliminate bottlenecks.
- Debugging & Reliability: Troubleshoot complex hardware–software interaction issues, including vLLM compilation failures on ROCm, CUDA memory leaks, distributed runtime failures, and kernel-level inconsistencies.
- Collaborate with research, infrastructure, and kernel engineering teams to improve system throughput, stability, and developer experience.
Qualifications
- 5+ years of industry experience in ML systems, distributed training, or related fields.
- Bachelor’s or Master’s degree in Computer Science, Computer Engineering, Electrical Engineering, or related technical fields.
- Strong programming experience in Python, C++, and familiarity with ML tooling and distributed systems.
- Deep understanding of profiling tools (e.g., Nsight, ROCm Profiler, XLA profiler, TPU tools).
- Deep expertise with partitioning configuration on the modern ML frameworks such as PyTorch and JAX.
- Experience with multi-node distributed training systems and orchestration frameworks (DTensor, GSPMD, etc.).
- Hands-on experience maintaining or building ML training stacks involving CUDA, ROCm, NCCL, XLA, or similar technologies.
Skills
Required:PythonC++ML tooling and distributed systemsProfiling tools
Benefits
- Medical, dental, and vision insurance
- 401k plan
- Daily lunch, snacks, and beverages
- Flexible time off
- Competitive salary and equity
Last updated Oct 2, 2026. Interested? Apply on the company site.
About Sciforium
Byte-Native AI Models & Infrastructure
Sciforium is pioneering the future of AI infrastructure and research. We're developing byte-native multimodal foundation models while delivering serverless LLM serving at a fraction of traditional costs
More open roles at Sciforium
- Office Manager2d agoSan Francisco, CA · Operations
- Research Engineer - Model Evaluation & MLOps1mo agoSan Francisco, CA · Engineering
- Pre-training Research Engineer1mo agoSan Francisco, CA · Engineering
- GPU Cluster Engineer, Networking2mo agoSan Francisco, CA · Engineering
- GPU Cluster Engineer, Systems & Platform2mo agoSan Francisco, CA · Engineering