Optimizing irregular dense operators of heterogeneous GNN models on GPU

By Israt Nisa, Minjie Wang, Da Zheng, Qiang Fu, Umit Catalyurek, George Karypis
2023
Download Copy BibTeX
Copy BibTeX
GNN models on heterogeneous graphs have achieved state-of-the-art (SOTA) performance in various graph tasks such as link prediction and node classification. Despite their success in providing SOTA results, popular GNN libraries, such as PyG and DGL, fail to provide fast and efficient solutions for heterogeneous GNN models. One common key bottlenecks of models like RGAT, RGCN, and HGT is relation-specific linear projection. In this paper, we propose two high-performing tensor operators: gather-mm and segment-mm to address the issue. We demonstrate the effectiveness of the proposed operators in training two popular heterogeneous GNN models – RGCN and HGT. Our proposed approaches outperform the full-batch training time of RGCN by up to 3× and mini-batch by up to 2×.
Research areas

Latest news

US, NY, New York
We are seeking a Robotics/AI Motor Control Scientist to develop cutting-edge machine learning algorithms for motor control systems in robots. In this role, you will focus on creating and optimizing intelligent motor control strategies to enable robots to perform complex, whole-body tasks. Your contributions will be essential in advancing robotics by enabling fluid, reliable, and safe interactions between robots and their environments. Key job responsibilities - Develop controllers that leverage reinforcement learning, imitation learning, or other advanced AI techniques to achieve natural, robust, and adaptive motor behaviors - Collaborate with multi-disciplinary teams to integrate motor control systems with robotic hardware, ensuring alignment with real-world constraints such as actuator dynamics and energy efficiency - Use simulation and real-world testing to refine and validate control algorithms - Stay updated on advancements in robotics, AI, and control systems to apply advanced techniques to robotic motion challenges - Lead technical projects from conception through production deployment - Mentor junior scientists and engineers - Bridge research initiatives with practical engineering implementation About the team Fauna Robotics, an Amazon company, is building capable, safe, and genuinely delightful robots for everyday life. Our goal is simple: make robots people actually want to live and interact with in everyday human spaces. We believe that future won’t arrive until building for robotics becomes far more accessible. Today, too much effort is spent reinventing the fundamentals. We’re changing that by developing tightly integrated hardware and software systems that make it faster, safer, and more intuitive to create real-world robotic products. Our work spans the full stack: mechanical design, control systems, dynamic modeling, and intelligent software. The focus is not just functionality, but experience. We’re building robots that feel responsive, expressive, and genuinely useful. At Fauna, you’ll work at the frontier of this space, helping define how robots move, manipulate, and interact with people in natural environments. It’s an opportunity to solve hard problems across hardware and software with a team focused on making robotics accessible and joyful to build. If you care about making robotics real for everyone and building systems that are as delightful as they are capable, we’re interested in hearing from you. an opportunity to solve hard problems across hardware and software with a team focused on making robotics accessible and joyful to build. If you care about making robotics real for everyone and building systems that are as delightful as they are capable, we’re interested in hearing from you.
US, WA, Seattle
Prime Video is a first-stop entertainment destination offering customers a vast collection of premium programming in one app available across thousands of devices. Prime members can customize their viewing experience and find their favorite movies, series, documentaries, and live sports – including Amazon MGM Studios-produced series and movies; licensed fan favorites; and programming from Prime Video subscriptions such as Apple TV+, HBO Max, Peacock, Crunchyroll and MGM+. All customers, regardless of whether they have a Prime membership or not, can rent or buy titles via the Prime Video Store, and can enjoy even more content for free with ads. Are you interested in shaping the future of entertainment? Prime Video's technology teams are creating best-in-class digital video experience. As a Prime Video team member, you’ll have end-to-end ownership of the product, user experience, design, and technology required to deliver state-of-the-art experiences for our customers. You’ll get to work on projects that are fast-paced, challenging, and varied. You’ll also be able to experiment with new possibilities, take risks, and collaborate with remarkable people. We’ll look for you to bring your diverse perspectives, ideas, and skill-sets to make Prime Video even better for our customers. With global opportunities for talented technologists, you can decide where a career Prime Video Tech takes you! Key job responsibilities As a highly experienced and seasoned science leader, you will apply state of the art natural language processing and computer vision research to video centric digital media, while also responsible for creating and maintaining the best environment for applied science in order to recruit, retain and develop top talent. You will lead the research direction for a team of deeply talented applied scientists, creating the roadmaps for forward-looking research and communicate them effectively to senior leadership. You will also hire and develop applied scientists - growing the team to meet the evolving needs of our customers. About the team This team's mission is to deeply understand all content and empower all customers with relevant language options, innovative accessibility assists, and rich title-information across all their content-experiences on Prime Video. We create and publish content on-time that's meaningful, accurate, and accessible to every customer globally. We delight our customers by pushing the boundaries of content understanding and enrichment. Through inclusion and innovation, we do the most fulfilling work of our career.
GB, MLN, Edinburgh
Do you want to make a real difference to real people's lives? Want to design and build fair and explainable systems which automate recruitment processes across Amazon? Come and be part of a team that develops new machine learning (ML) technologies, which help Amazon scale for its customers by recruiting diverse teams. Join our Recommendations team within Intelligent Talent Acquisition (ITA) where you’ll build machine learning products that transform how job seekers find opportunities and recruiters discover talent. You’ll develop sophisticated recommendation systems powering both Amazon Jobs and internal hiring platforms, operating at global scale to match the right people with the right positions. Using techniques including representation learning, reinforcement learning, and probabilistic modeling, your work will directly improve efficiency for recruiters and help candidates find their ideal roles. This position offers the chance to solve complex problems with significant impact by creating systems that make Amazon’s entire hiring ecosystem more effective while collaborating with scientists across the organization. Key job responsibilities - Design and implement machine learning models that power recommendation systems for job seekers and recruiters, ensuring high performance, scalability, and reliability at global scale. Our ideal candidate has a strong scientific foundation and experience of statistical analysis and model building and has a passion for fairness and explainability in ML systems. - Collaborate with engineers, scientists, and product managers to define requirements, create solutions, and deliver products that improve the hiring experience. - Participate in the full software development lifecycle including scoping, design, coding, testing, documentation, deployment, and maintenance of recommendation systems and ML models. - Solve complex ML problems using optimal data structures and algorithms, making thoughtful trade-offs between efficiency and maintainability. - Stay current with scientific literature and develop novel approaches that address business challenges in talent acquisition. You will have the opportunity to provide feedback on scientific work across the organization helping the entire Intelligent Talent Acquisition organization improve. A day in the life You might spend the morning reviewing a colleague’s code for a new recommendation algorithm feature, then collaborate with product managers to refine requirements for an upcoming enhancement. After lunch, you’ll dive into model development, analyzing performance metrics from recent A/B tests and implementing improvements to the job-seeker recommendation pipeline. Throughout the day, you’ll participate in scientific discussions with peers across the organization, providing valuable feedback while continuing to refine your expertise. About the team The Recommendations team is a hybrid group of software engineers and applied scientists located in Edinburgh. We build tools that match people to jobs and jobs to people, optimizing experiences for both recruiters and candidates. Our work directly impacts Amazon’s ability to find and hire exceptional talent globally. The team maintains a collaborative environment with regular knowledge sharing and mentorship opportunities. We work closely with our product teams to understand business needs and develop innovative scientific solutions that improve hiring outcomes across both industry and student requisitions worldwide.
US, CA, Pasadena
We are seeking an Applied Scientist to join the SAF Lab. In this role, you will lead the effort in safe reinforcement learning (RL) including the development of legged locomotion algorithms that internalize safety and are deployable on physical hardware—enabling highly dynamic robots to walk, run, avoid collisions and recover from disturbances with agility and robustness. You will develop RL architectures that interface with physics-based models (for dynamic retargeting and reward shaping), internalize safety constraints in training, sim-to-real transfer and interface with safety filters at run-time. Therefore, your work will sit at the intersection of safety-critical control and learning, and you will collaborate with others in the SAF Lab and Amazon working on perception, planning, whole-body and safety-critical control. This is an opportunity to shape the foundations of safe learning on emerging platforms that will remove bottlenecks to deployment and enable these robots to safely operate around humans. Key job responsibilities • Collaborate with product teams and science leaders to set a science roadmap (with eventual impact on real robots). • Design, train, and deploy reinforcement learning (RL) policies for dynamic legged locomotion including walking, running, stair climbing, and fall recovery on physical robots • Develop sim-to-real transfer pipelines that produce policies robust to the reality gap, including domain randomization, system identification, and adaptive strategies • Integrate control-based methods with RL, as inputs to the RL (dynamic retargeting and control-guided rewards), in training (internalizing safety constraints in training), and as the RL feeds into safety layers and whole-body control • Develop and maintain large-scale training infrastructure for locomotion policy learning, including physics simulation environments, domain randomization and GPU parallelization • Investigate the distillation of locomotion policies, integration with whole-body control, foundation models, VLAs, world models, perception and full-stack autonomy • Evaluate policy performance rigorously through simulation benchmarks, hardware experiments, and failure-mode analysis • Publish research at top-tier robotics and ML venues and contribute to Amazon's scientific reputation in advanced robotics • Collaborate with perception and planning teams to enable terrain-aware and goal-conditioned locomotion behaviors A day in the life Amazon offers a full range of benefits that support you and eligible family members, including domestic partners and their children. Benefits can vary by location, the number of regularly scheduled hours you work, length of employment, and job status such as seasonal or temporary employment. The benefits that generally apply to regular, full-time employees include: 1. Medical, Dental, and Vision Coverage 2. Maternity and Parental Leave Options 3. Paid Time Off (PTO) 4. 401(k) Plan If you are not sure that every qualification on the list above describes you exactly, we'd still love to hear from you! At Amazon, we value people with unique backgrounds, experiences, and skillsets. If you’re passionate about this role and want to make an impact on a global scale, please apply! About the team Work with the inventor of control barrier functions in the Safe Autonomy Frontiers (SAF) Lab. The first industry research lab in safe autonomy, developing a universal safety layer for the next generation of robotic systems: mobile robots, manipulators, mobile manipulators, and future platforms with dynamic stability. You will push the frontiers of performant safety for highly dynamic robots: CBF theory integrated with perception and learning, evaluated on next-generation robots. Your work will underpin robots operating alongside people at Amazon's unprecedented scale.
US, WA, Redmond
We are searching for a talented candidate with expertise in orbital mechanics and spaceflight navigation, including LEO Satellite Orbit Determination. This position requires experience in simulation and analysis of spacecraft orbital mechanics and sequential orbit determination methods, including Extended Kalman Filters (EKF) and/or Unscented Kalman Filter (UKF). Strong analysis skills are required to develop engineering studies of complex large-scale dynamical systems. This position requires demonstrated expertise in computational analysis automation and tool development. Key job responsibilities - Perform spacecraft maneuver or navigation analysis in support of multi-disciplinary trades within the Amazon Leo team. - Contribute to prototype software development of flight algorithms. - Test and assess navigation software for integration into flight systems. - Assess and trouble-shoot the performance of Leo on-board GNSS hardware and software systems. - Work closely with GNC engineers to manage on-orbit performance and develop flight dynamics operations processes. Export Control Requirement: Due to applicable export control laws and regulations, candidates must be a U.S. citizen or national, U.S. permanent resident (i.e., current Green Card holder), or lawfully admitted into the U.S. as a refugee or granted asylum. A day in the life - Interacting with GNC teams to evaluate and troubleshoot satellite issues. - Working within the Flight Dynamics Research team to prioritize tasks. - Performing analysis, simulation, testing and documentation to address assigned tasks.
BR, SP, Sao Paulo
Are you passionate about helping customers achieve business transformation through AI? Do you want to lead forward-deployed teams that embed directly into the enterprise and unlock real business outcomes? And are you ready to operate as a general manager across engineering, science, and commercial strategy in the fastest-moving space in AI and infrastructure? The AWS Generative AI Innovation Center (GenAIIC) is on a mission to accelerate enterprise AI transformation across global customers going from beyond isolated use cases to holistic, C-suite-sponsored initiatives that reshape how organizations operate. We combine deep AI expertise across science, strategy, and business transformation. We start with the customer's most critical operational challenges and work backwards, and deploy multidisciplinary teams that embed with the customer, prove impact in 45-day sprints, and expand across the enterprise. We are a fast-moving, entrepreneurial team that values leaders who can operate across technical depth and commercial breadth. You will lead a team of ML engineers, AI scientists, and AI strategists who work alongside customers to architect and deliver AI solutions that move and stay in production, realizing value. You will regularly engage with CFOs, CIOs, and C-suite executives. You must bring equal fluency in engineering, data science, go-to-market, and customer delivery. You are ready to roll up your sleeves alongside the team, whether that means scoping an agentic AI architecture, presenting to a board, or operationalizing a repeatable delivery motion. You will partner with customers, AWS Sales, AWS service teams, AWS industry teams and AWS Professional Services delivery teams to meet the specific needs of the customer, and extend that use to other customers. The successful candidate will possess both technical and customer-facing skills that will allow you to be the technical “face” of AWS within our solution providers’ ecosystem/environment as well as directly to end customers. You will be able to drive discussions with senior technical and management personnel within customers and partners, as well as the technical background that enables them to interact with and give guidance to data/research/applied scientists and software developers. The ideal candidate will also have a demonstrated ability to think strategically about business, product, and technical issues. Finally, and of critical importance, the candidate will be an excellent technical team manager, someone who knows how to hire, develop, and retain high quality technical talent. About the team Diverse Experiences AWS values diverse experiences. Even if you do not meet all of the qualifications and skills listed in the job description, we encourage candidates to apply. If your career is just starting, hasn’t followed a traditional path, or includes alternative experiences, don’t let it stop you from applying. Why AWS? Amazon Web Services (AWS) is the world’s most comprehensive and broadly adopted cloud platform. We pioneered cloud computing and never stopped innovating — that’s why customers from the most successful startups to Global 500 companies trust our robust suite of products and services to power their businesses. Inclusive Team Culture Here at AWS, it’s in our nature to learn and be curious. Our employee-led affinity groups foster a culture of inclusion that empower us to be proud of our differences. Mentorship & Career Growth We’re continuously raising our performance bar as we strive to become Earth’s Best Employer. That’s why you’ll find endless knowledge-sharing, mentorship and other career-advancing resources here to help you develop into a better-rounded professional. Work/Life Balance We value work-life harmony. Achieving success at work should never come at the expense of sacrifices at home, which is why flexible work hours and arrangements are part of our culture. When we feel supported in the workplace and at home, there’s nothing we can’t achieve in the cloud.
US, WA, Seattle
Join the AWS Perimeter Protection team as a Senior Applied Scientist, where you will bring your deep ML engineering expertise to design, build, and scale AI-driven security solutions that protect AWS customers worldwide. This role is ideal for someone who has already built and shipped production ML systems at industry scale and is looking to apply that experience to high-impact security challenges. You will own the full ML lifecycle — from research and prototyping to production deployment and optimization — powering services including Web Application Firewall, DDoS Protection, Bot Management, and Infrastructure Protection. With services spanning all AWS regions and handling trillions of requests per week, you will solve complex engineering and science problems where model performance, system reliability, and low-latency inference are critical. Key job responsibilities - Design, build, and deploy production-grade ML models and systems for real-time threat detection, mitigation, and protection against evolving cyber threats at cloud scale. - Own the full ML lifecycle end-to-end — from problem formulation, data engineering, and model development through to production deployment, monitoring, and continuous improvement. - Architect and optimize ML pipelines, training infrastructure, and serving systems to meet strict latency, throughput, and reliability requirements at AWS scale. - Bridge the gap between research and production by translating novel ML approaches into robust, scalable, and maintainable systems that operate in real-time security environments. - Design and implement feature engineering workflows and large-scale data processing pipelines to support rapid experimentation and reliable model iteration. - Collaborate closely with software engineering teams to integrate ML models into distributed, low-latency security services, driving engineering decisions around model serving, infrastructure, and system design. - Analyze large-scale production data to identify patterns, anomalies, and emerging threat vectors, and translate findings into measurable improvements to detection and mitigation capabilities. - Establish and improve best practices for ML system design, model evaluation, A/B testing, and production monitoring across the team. - Mentor junior scientists and engineers, raising the bar on both scientific rigor and engineering quality.
US, WA, Seattle
The Annapurna ML team is looking for a Senior Applied Scientist to work on the intersection of Artificial Intelligence and program analysis to raise the code quality bar in our state-of-the-art deep learning compiler stack. This stack is designed to optimize application models across diverse domains, including Large Language and Vision, originating from leading frameworks such as PyTorch, TensorFlow, and JAX. Your role will involve working closely with our custom-built Machine Learning accelerators, Inferentia and Trainium, which represent the forefront of Annapurna innovation for advanced ML capabilities, and is the underpinning of Generative AI. As a Senior Applied Scientist, you'll be instrumental in designing, developing, and deploying analyzers for ML compiler stages and compiler IRs. You will architect and implement business-critical tooling, publish research, and mentor a brilliant team of experienced scientists and engineers. You will need to be technically capable, credible, and curious in your own right as a trusted scientist, innovating on behalf of our customers. Your responsibilities will involve tackling crucial challenges alongside a talented engineering team, contributing to leading-edge design and research in compiler technology and deep-learning systems software. Strong experience in programming languages, compilers, program analyzers, and program synthesis engines will be a benefit in this role. A background in machine learning and AI accelerators is preferred but not required. A day in the life Diverse Experiences Amazon values diverse experiences. Even if you do not meet all of the qualifications and skills listed in the job description, we encourage candidates to apply. If your career is just starting, hasn’t followed a traditional path, or includes alternative experiences, don’t let it stop you from applying. Why Amazon? We pioneered cloud computing and never stopped innovating — that’s why customers from the most successful startups to Global 500 companies trust our robust suite of products and services to power their businesses. Inclusive Team Culture Here at Amazon, it’s in our nature to learn and be curious. Our employee-led affinity groups foster a culture of inclusion that empower us to be proud of our differences. Ongoing events and learning experiences, including our Conversations on Race and Ethnicity (CORE) and AmazeCon (diversity) conferences, inspire us to never stop embracing our uniqueness. Mentorship & Career Growth We’re continuously raising our performance bar as we strive to become Earth’s Best Employer. That’s why you’ll find endless knowledge-sharing, mentorship and other career-advancing resources here to help you develop into a better-rounded professional. Work/Life Balance We value work-life harmony. Achieving success at work should never come at the expense of sacrifices at home, which is why we strive for flexibility as part of our working culture. When we feel supported in the workplace and at home, there’s nothing we can’t achieve in the cloud.
US, WA, Seattle
The Automated Reasoning Group in the Amazon Neuron team is looking for an Applied Scientist to work on the intersection of Artificial Intelligence and program analysis to raise the code quality bar in our state-of-the-art deep learning compiler stack. This stack is designed to optimize application models across diverse domains, including Large Language and Vision, originating from leading frameworks such as PyTorch and JAX. Your role will involve working closely with our custom-built Machine Learning accelerator, Trainium, which represents the forefront of innovation for advanced ML capabilities, and is the underpinning of Generative AI. In this role as an Applied Scientist, you'll be instrumental in designing, developing, and deploying analyzers for ML compiler stages and compiler IRs. You will architect and implement business-critical tooling, publish research, and mentor a brilliant team of experienced scientists and engineers. You will need to be technically capable, credible, and curious in your own right as a trusted AWS Neuron engineer, innovating on behalf of our customers. Your responsibilities will involve tackling crucial challenges alongside a talented engineering team, contributing to leading-edge design and research in compiler technology and deep-learning systems software. Strong experience in programming languages, compilers, program analyzers, theorem provers, and program synthesis engines will be a benefit in this role. A background in machine learning and AI accelerators is preferred but not required.
US, WA, Seattle
The Shopping Convo Foundations Team - Pre-purchases Science is looking for an Senior Applied Scientist with expertise in Artificial Intelligence and Machine Learning to drive scientific innovation that expands Amazon's product catalogue. Our goal is to leverage AI/ML solutions to enhance catalogue coverage with high precision. In this role, you will lead the research and development of novel machine learning approaches to solve complex catalogue expansion and product attribute challenges. You lead the design and develop state-of-the-art ML models, conduct rigorous experimentation, translate scientific breakthroughs into production-ready solutions, and guide a set of junior scientists. You will work closely with ML Engineers and Software Development Engineers to optimize model performance, ensure scalability, and deploy low-latency solutions at Amazon scale. About the team Our team is a horizontal applied science group that works across the full lifecycle of brand and product data extraction and quality. We span seven workstreams — from products sourcing and Brand entitlement, to designing and evaluating extraction strategies for ASINs, offers, and brand attributes at scale, as well as relevance modeling and search. We drive root cause analysis through human-in-the-loop evaluation, improve how catalog data surfaces in search, optimize business metrics tied to data quality, and build brand intelligence capabilities. This cross-cutting scope positions the team as a connective layer across product, engineering, and science — ensuring that improvements in one area compound across the system rather than remain siloed.