Staff Software Engineer, Machine Learning, JAX Third Party

Company: Google Inc.
Location: Sunnyvale
Posted on: May 2, 2025

Job Description:

Staff Software Engineer, Machine Learning, JAX Third Partycorporate_fare Google place Sunnyvale, CA, USAApply

• Bachelor's degree or equivalent practical experience.
• 8 years of experience in software development.
• 5 years of experience testing, and launching software products.
• 5 years of experience with ML design and ML infrastructure (e.g., model deployment, model evaluation, data processing, debugging, fine-tuning).
• 3 years of experience with software design and architecture.Preferred qualifications:
  • Master's degree or PhD in Engineering, Computer Science, or a related technical field.
  • 8 years of experience with data structures/algorithms.
  • 3 years of experience in a technical leadership role leading project teams and setting technical direction.
  • 3 years of experience working in a complex, matrixed organization involving cross-functional or cross-business projects.About the jobGoogle's software engineers develop next-generation technologies that change how billions of users connect, explore, and interact with information and one another. Our products handle information at massive scale and extend beyond web search. We seek engineers with fresh ideas from areas like information retrieval, distributed computing, large-scale system design, networking, data storage, security, AI, natural language processing, UI design, and mobile development. As a software engineer, you will work on critical projects with opportunities to switch teams and projects as our business evolves. We value versatility, leadership, and enthusiasm for tackling new problems across the full-stack to push technology forward.You will manage project priorities, deadlines, and deliverables, and will design, develop, test, deploy, maintain, and enhance software solutions.Our team's mission is end-to-end training and inference success for Google and the world, with cost-effective and high-performance software stacks across hardware, frameworks, and model types.In this role, you will collaborate with internal and cloud developers to innovate ML-supported products with performant, efficient, and effective ML infrastructure. You will help build the ML infrastructure platform that connects top-level frameworks with target hardware.Google Cloud accelerates digital transformation for organizations worldwide, providing enterprise-grade solutions leveraging cutting-edge technology and tools to foster sustainable development. Customers across more than 200 countries rely on Google Cloud to address critical business challenges.The US base salary range for this full-time position is $197,000-$291,000 plus bonus, equity, and benefits. Salary ranges are role, level, and location-dependent. Specifics will be shared during the hiring process.Note: Listed compensation reflects base salary only, excluding bonus, equity, or benefits. Learn more about benefits at Google.Responsibilities
    • Develop and lead the OSS JAX AI Inference Stack for third-party customers in GCP, supporting GPU and TPU.
    • Implement strategies for optimal inference performance, including framework, runtime, and compiler integration.
    • Support GCP's diverse ML ecosystem needs for JAX and hardware through a cross-stack, co-design approach.
    • Develop stable, standardized APIs and modular components for each layer of the inference stack, ensuring support for large-scale inference.
    • Build multi-cloud migration support for third-party users.Google is committed to equal opportunity employment, fostering a diverse and inclusive workforce. We provide an environment of belonging and do not discriminate based on race, creed, color, religion, gender, sexual orientation, gender identity/expression, national origin, disability, age, genetic information, veteran status, marital status, pregnancy, or related conditions, among others. See Google's EEO Policy and related resources.English proficiency is required for global collaboration unless otherwise specified. Google does not accept unsolicited resumes from recruitment agencies and is not responsible for fees related to such resumes.
      #J-18808-Ljbffr

Keywords: Google Inc., North Highlands , Staff Software Engineer, Machine Learning, JAX Third Party, IT / Software / Systems , Sunnyvale, California