Hector: An efficient programming and compilation framework for implementing relational graph neural networks in GPU architectures

By Kun Wu, Mert Hidayetoğlu, Xiang Song, Sitao Zhang, Da Zheng, Israt Nisa, Wen-mei Hwu
2024
Download Copy BibTeX
Copy BibTeX
Relational graph neural networks (RGNNs) are graph neural networks with dedicated structures for modeling the different types of nodes and edges in heterogeneous graphs. While RGNNs have been increasingly adopted in many real-world applications due to their versatility and accuracy, they pose performance and system design challenges: inherent memory-intensive computation patterns, the gap between the programming interface and kernel APIs, and heavy programming effort required to optimize kernels caused by their coupling with data layout and heterogeneity. To systematically address these challenges, we propose Hector, a novel two-level intermediate representation and its code genera-tor framework that (a) captures the key properties of RGNN models, and opportunities to reduce memory accesses in inter-operator scheduling and materialization, (b) generates code with flexible data access schemes to eliminate redundant data copies, and (c) decouples model semantics, data layout, and operators-specific optimizations from each other to reduce programming effort. By building on one general matrix multiply (GEMM) template and a node/edge traversal template, Hector achieves up to 9.9× speed-up in inference and 43.7× speed-up in training compared with the state-of-the-art public systems on select models, RGCN, RGAT and HGT, when running heterogeneous graphs provided by Deep Graph Library (DGL) and Open Graph Benchmark (OGB). In addition, Hector does not trigger any out-of-memory (OOM) exception in these tests. We also propose linear operator reordering and compact materialization to further accelerate the system by up to 3.8×. As an indicator of the reduction of programming effort, Hector takes in 51 lines of code ex-pressing the three models and generates a total of 8K lines of CUDA and C++ code. Through profiling, we found that higher memory efficiency allows Hector to accommodate larger input and therefore attain higher throughput in for-ward propagation, while backward propagation is bound by latency introduced by atomic updates and outer products.
Research areas

Latest news

US, NY, New York
The Sponsored Products and Brands team at Amazon Ads is re-imagining the advertising landscape through generative AI technologies, revolutionizing how millions of customers discover products and engage with brands across Amazon.com and beyond. We are at the forefront of re-inventing advertising experiences, bridging human creativity with artificial intelligence to transform every aspect of the advertising lifecycle from ad creation and optimization to performance analysis and customer insights. We are a passionate group of innovators dedicated to developing responsible and intelligent AI technologies that balance the needs of advertisers, enhance the shopping experience, and strengthen the marketplace. If you're energized by solving complex challenges and pushing the boundaries of what's possible with AI, join us in shaping the future of advertising. About the team SPB Agent team's vision is to build a highly personalized and context-aware agentic advertiser guidance system that seamlessly integrates Large Language Models (LLMs) with sophisticated tooling, operating across all experiences. The SPB-Agent is the central agent that interfaces with advertisers across Ads Console, Selling Partner portals (Seller Central, KDP, Vendor Central), and internal Sales systems. We identify high-impact opportunities spanning from strategic product guidance to granular optimization and deliver them through personalized, scalable experiences grounded in state-of-the-art agent architectures, reasoning frameworks, sophisticated tool integration, and model customization approaches including fine-tuning, MCP, and preference optimization. This presents an exceptional opportunity to shape the future of e-commerce advertising through advanced AI technology at unprecedented scale, creating solutions that directly impact millions of advertisers.
DE, BE, Berlin
At Audible, we believe stories have the power to transform lives. It’s why we work with some of the world’s leading creators to produce and share audio storytelling with our millions of global listeners. We are dreamers and inventors who come from a wide range of backgrounds and experiences to empower and inspire each other. Imagine your future with us. ABOUT THIS ROLE As an Applied Scientist, you will solve large complex real-world problems at scale, draw inspiration from the latest science and technology to empower undefined/untapped business use cases, delve into customer requirements, collaborate with tech and product teams on design, and create production-ready models that span various domains, including Machine Learning (ML), Artificial Intelligence (AI) and Generative AI, Natural Language Processing (NLP), Reinforcement Learning (RL), real-time and distributed systems. ABOUT YOU Your work will focus on inventing or adapting scientific approaches, models, and algorithms driven by customer needs at the project level. You will develop components and/or end-to-end solutions that are deployed into production or directly support production systems, delivering consistently high-quality work that meets both scientific and engineering best practices. You will develop reusable science components and services that resolve architecture deficiencies and customers’ pain points, while making technical trade-offs for long-term/short- term. You will work semi-autonomously to deliver solutions, contribute to research papers at peer-reviewed venues when appropriate, and document your work thoroughly to enable others to understand and reproduce it. Your decision-making will consistently incorporate robust, data-driven business and technical judgment. You will collaborate with other scientists to raise the bar of both scientific and engineering complexity for the team and to foster valuable scientific partnership opportunities to help/guide science decisions. We work in a highly collaborative, fast-paced environment where scientists, engineers, and product managers work to test and build scalable foundational capabilities, as well as customer facing experiences. You will have the opportunity to innovate and think big within your projects scope, implement optimization services and algorithms, and influence the experiences of millions of customers. We are looking for a results-oriented Applied Scientist with deep knowledge in ML, NLP, Deep Learning, GenAI, and/or large-scale distributed computation. As an Applied Scientist, you will... - Understand use cases across the business and adopt/extend/design/invent solutions/models that are scalable, efficient, and automated for difficult problems that are not well defined - Work closely with fellow scientists and software engineers (at Audible and Amazon) to build and productionize models, deliver novel and highly impactful features - Review models of peers for the purpose of reducing and managing risk to the business, while improving customer experience - Design, develop, and deploy modeling techniques and solutions for Content Understanding, Recommendations, GenAI-based product features, by employing a wide range of methodologies, working from simple to complex - Contribute to initiatives that employ the most recent advances in ML/AI in a fast-paced, experimental environment - Push the boundary of innovation ABOUT AUDIBLE Audible is the leading producer and provider of audio storytelling. We spark listeners’ imaginations, offering immersive, cinematic experiences full of inspiration and insight to enrich our customers daily lives. We are a global company with an entrepreneurial spirit. We are dreamers and inventors who are passionate about the positive impact Audible can make for our customers and our neighbors. This spirit courses throughout Audible, supporting a culture of creativity and inclusion built on our People Principles and our mission to build more equitable communities in the cities we call home.
IN, KA, Bengaluru
This position is based in Bangalore, India The Last Mile team helps get packages from delivery stations to a customer’s doorstep. To provide new innovations for customers, awe are inventing the next-generation smart delivery operation. We are combining innovative mobile and IoT technologies, data streams (video, vehicle telematics, location, and presence), together with machine learning models and algorithms – all to create solutions that allow us to deliver faster, and with more confidence. Playing a key role in the Last Mile Driver Experience team, as a Applied Scientist you will be responsible for building machine learning models and algorithms in areas including mapping and location, pattern detection in sensor data, and computer vision. Using your research, you will work with your engineering and product management peers to drive designs from ideation through development and into production. You will bring your experience of research for similar products and solutions, preferably in consumer or industrial verticals. This role requires autonomy and an ability to deliver results, often within the ambiguity of building a v1 product. You will need to work efficiently to build the right things with limited guidance, raising the bar to create an amazing experience for our customers.
ES, M, Madrid
Amazon's EU International Technology (EU INTech) organisation is creating new ways for customers to discover products through innovative customer experiences. We are a science-only team within EU INTech, responsible for designing and developing AI/ML science solutions that support business needs across Amazon's global search and discovery experiences. Our mission is to make Amazon navigation easier for customers worldwide. We achieve this through two strategic pillars: making Amazon navigation more visual and improving Amazon navigation with more inspiring discovery tools and narrowing navigation. To support this vision, we build and deploy AI/ML models that surface the most relevant content to hundreds of millions of Amazon customers worldwide. Our team comprises Applied Scientists and we partner with other teams, collaborating with ML Engineers, Software Developers, Product Managers, Technical Product Managers, and UX Designers. We are located in the Madrid Technical Hub. We are looking for Applied Scientists who are passionate about solving highly ambiguous and challenging problems at global scale. This is a hands-on, end-to-end applied science role where you will own the full lifecycle of science solutions — from business problem analysis and science plan design, through development and experimentation, to production deployment. We are looking for AI/ML experts with knowledge on ranking, computer vision, recommendation systems, search, and customer experience design. What makes this role unique: • End-to-end ownership – You will analyse business problems, map them to science plans, and design and develop solutions from ideation to production. We are owners of the full science lifecycle. • Applied science with a research edge – While our focus is on delivering applied science solutions that drive measurable business impact, our team actively pushes the state of the art in areas such as computer vision and Generative AI. • Hands-on execution – We need scientists who thrive in building, experimenting, and shipping. What are we looking for? • A scientist who can independently analyse any business problem and design a rigorous science approach to solve it • Strong hands-on engineering skills — you build and ship, not just theorise • Deep expertise in one or more of: computer vision, generative AI, recommendation systems, ranking, or NLP • Experience taking ML models from research to production at scale • Comfort with ambiguity and the ability to structure complex, undefined problems • A passion for customer-centric innovation and measurable impact • A strong communicator capable to adapt the message from a science audience, to engineering or leadership Key job responsibilities • Analyse complex business problems and translate them into well-defined science plans with clear milestones and success criteria • Design, develop, and deliver ML/AI models end-to-end — from research and prototyping through to production systems at Amazon scale and extending solutions going beyond the state of the art • Work with state-of-the-art models in computer vision, ranking and generative AI to power new customer experiences globally • Own major science challenges for the team, driving solutions from ideation through experimentation to production deployment • Collaborate with a variety of roles and partner teams around the world to deliver integrated solutions • Influence scientific direction and best practices across the team • Maintain high quality standards on team deliverables • Contribute to expanding the state of the art in computer vision, ranking and GenAI through publications and internal knowledge sharing
ES, M, Madrid
Amazon's EU International Technology (EU INTech) organisation is creating new ways for customers to discover products through innovative customer experiences. We are a science-only team within EU INTech, responsible for designing and developing AI/ML science solutions that support business needs across Amazon's global search and discovery experiences. Our mission is to make Amazon navigation easier for customers worldwide. We achieve this through two strategic pillars: making Amazon navigation more visual and improving Amazon navigation with more inspiring discovery tools and narrowing navigation. To support this vision, we build and deploy AI/ML models that surface the most relevant content to hundreds of millions of Amazon customers worldwide. Our team comprises Applied Scientists and we partner with other teams, collaborating with ML Engineers, Software Developers, Product Managers, Technical Product Managers, and UX Designers. We are located in the Madrid Technical Hub. We are looking for Applied Scientists who are passionate about solving highly ambiguous and challenging problems at global scale. This is a hands-on, end-to-end applied science role where you will own the full lifecycle of science solutions — from business problem analysis and science plan design, through development and experimentation, to production deployment. We are looking for AI/ML experts with knowledge on ranking, computer vision, recommendation systems, search, and customer experience design. What makes this role unique: • End-to-end ownership – You will analyse business problems, map them to science plans, and design and develop solutions from ideation to production. We are owners of the full science lifecycle. • Applied science with a research edge – While our focus is on delivering applied science solutions that drive measurable business impact, our team actively pushes the state of the art in areas such as computer vision and Generative AI. • Hands-on execution – We need scientists who thrive in building, experimenting, and shipping. What are we looking for? • A scientist who can independently analyse any business problem and design a rigorous science approach to solve it • Strong hands-on engineering skills — you build and ship, not just theorise • Deep expertise in one or more of: computer vision, generative AI, recommendation systems, ranking, or NLP • Experience taking ML models from research to production at scale • Comfort with ambiguity and the ability to structure complex, undefined problems • A passion for customer-centric innovation and measurable impact • A strong communicator capable to adapt the message from a science audience, to engineering or leadership Key job responsibilities • Analyse complex business problems and translate them into well-defined science plans with clear milestones and success criteria • Design, develop, and deliver ML/AI models end-to-end — from research and prototyping through to production systems at Amazon scale and extending solutions going beyond the state of the art • Work with state-of-the-art models in computer vision, ranking and generative AI to power new customer experiences globally • Own major science challenges for the team, driving solutions from ideation through experimentation to production deployment • Collaborate with a variety of roles and partner teams around the world to deliver integrated solutions • Influence scientific direction and best practices across the team • Maintain high quality standards on team deliverables • Contribute to expanding the state of the art in computer vision, ranking and GenAI through publications and internal knowledge sharing
US, CA, Sunnyvale
The Artificial General Intelligence (AGI) Customization Team is seeking a highly skilled and experienced Applied Scientist to support adoption and enable customization of Amazon Nova. The role focuses on developing state-of-the-art services and tools for model customization, including supervised fine-tuning, reinforcement learning, and knowledge distillation across large language models. As an Applied Scientist, you will play a important role in developing advanced customization capabilities that enable enterprises to build highly performant application-specific models without the need for training models from scratch. Your work will directly impact how companies leverage Amazon Nova models for their specific use cases. Key job responsibilities - Contribute to the development of novel customization techniques including extended post-training, continued pre-training, and advanced knowledge distillation - Collaborate with cross-functional teams to design and implement enterprise-ready tooling for various training techniques on Amazon SageMaker - Design and execute experiments to optimize model accuracy, latency, and cost across different customization approaches (SFT, DPO, PPO) - Develop and enhance preference learning algorithms and training curricula for customer-specific applications - Create robust evaluation frameworks for assessing model performance across different domains and use cases - Contribute to the development of the Responsible AI toolkit, including creating training and evaluation datasets for model alignment - Design and implement secure access mechanisms for early model checkpoints and weights - Communicate technical insights and results to both technical and non-technical stakeholders through presentations and documentation
IN, KA, Bengaluru
Amazon is seeking a passionate and inventive Applied Scientist II with a strong machine learning background to build industry-leading Speech and Language technology. Our mission is to deliver delightful customer experiences by advancing Automatic Speech Recognition (ASR), Natural Language Understanding (NLU), Machine Learning (ML), and Computer Vision (CV). You will work alongside internationally recognized experts to develop novel algorithms and modeling techniques that advance the state-of-the-art in human language technology. Your work will directly impact millions of customers through products and services powered by speech and language technology. You will gain hands-on experience with Amazon's heterogeneous speech, text, and structured data sources, and leverage large-scale computing resources to accelerate advances in spoken language understanding. We are hiring across all areas of human language technology: ASR, Machine Translation (MT), NLU, Text-to-Speech (TTS), Dialog Management, and Computer Vision. We also seek talent experienced in building large-scale, high-performing systems. Key job responsibilities Basic Qualifications PhD or M.Tech in Computer Science, Electrical Engineering, Mathematics, or Physics with specialization in one or more of: speech recognition, natural language processing, machine translation, time series analysis, signal processing, or machine learning 1-2 years of industry or research experience (including internships, co-ops, or post-doctoral work) in applied ML or related areas Proficiency in programming languages such as Python, C/C++, or Java Strong foundation in machine learning fundamentals and statistical modeling Preferred Qualifications Experience building speech recognition, machine translation, or natural language processing systems (e.g., commercial products, government projects, or published research with working prototypes) Hands-on experience with deep learning frameworks (e.g., PyTorch, TensorFlow) Track record of publications in top-tier conferences (e.g., NeurIPS, ICML, ACL, Interspeech, CVPR) Scientific thinking with demonstrated ability to innovate and contribute to advancing the field Solid software development practices and experience shipping production-quality code Strong written and verbal communication skills A day in the life 0
US, CA, San Jose
Are you excited about making business decisions using science and data? Are you interested in supporting consumer device concepts from idea inception to launch? Do you want to work on a Science Product team focused on scaling statistics and econometrics with custom tools? If so, this may be the role for you! Amazon.com strives to be Earth's most customer-centric company. The Amazon Devices and Services team focuses on delighting customer by enabling seamless functionality in supplying, entertaining, and managing the home -- and beyond. We seek and hire the world's brightest minds, offering them a fast-paced, technologically-sophisticated, and friendly work environment, where economic theory meets real-world industry. The Decision Science team in Devices owns demand estimates and pricing recommendations of concept devices before customers know they exist. We support devices and services ranging from Echo Frames to Kindle Paperwhite to Blink Video Camera …all prior to launch. We are a cross-functional Product team working to scale Econometrics through Amazon and beyond by incorporating Science into internal facing tools and making it easier for others to do so as well. In this role, you will have input in decision meetings with Amazon senior leadership, which include go/no-go decisions for brand new devices and services and build volume decisions for manufacture prior to receiving any customer signal. You will have direct input to pricing decisions. You will leverage Science and Tools produced by the Decision Science team such as conjoint demand models to produce these recommendations. You will work with Scientists, Economists, Product Managers, and Software Developers to provide meaningful feedback about stakeholder problems to inform business solutions and increase the velocity, quality, and scope behind our recommendations. You will also have the opportunity to work on special projects to both guide the business and advance your own knowledge and understanding of specific topics. Key job responsibilities Applies expertise to develop econometric/machine learning models to measure the demand of devices and the business; Reviews models and results for other scientists, mentors junior scientists; Generates economic insights for the Devices and Services business and work with stakeholders to run the business for effectively; Describes strategic importance of vision inside and outside of team; and, Identifies business opportunities, defines the problem and how to solve it; Engages with senior scientists, business leadership outside Devices and Services to understand interplay between different business units.
AU, VIC, Melbourne
Are you excited about leveraging and extending state-of-the-art Deep Learning, Information Retrieval, Natural Language Processing, Computer Vision algorithms to solve customer problems at the scale of Amazon? As an Applied Scientist Intern, you will be working in the Melbourne office in a fast-paced, cross-disciplinary team of experienced R&D scientists. You will take on complex problems, work on solutions that leverage existing academic and industrial research, and utilize your own out-of-the-box pragmatic thinking. In addition to coming up with novel solutions and prototypes, you may even deliver these to production in customer facing products. Key job responsibilities - Develop novel solutions and build prototypes - Work on complex problems in Deep Learning and Generative AI - Contribute to research that could significantly impact Amazon operations - Collaborate with a diverse team of experts in a fast-paced environment - Present your research findings to both technical and non-technical audiences - Collaborate with scientists on writing and submitting papers to top ML conferences, e.g. NeurIPS, ICML, ICLR, AISTATS, ACL ICCV, CVPR, KDD. Key Opportunities: - Work in a team of ML scientists to solve applied science problems at the scale of Amazon - Access to Amazon services and hardware - Potentially deliver solutions to production in customer-facing applications - Opportunities to be hired full-time after the internship Join us in shaping the future of AI at Amazon. Apply now and turn your research into real-world solutions!
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 - Lead research and development of speech and audio generation technology and end-to-end speech-to-speech architecture - Develop audio processing solutions for production environments, including source separation, enhancement, and mixing - Define the research roadmap for your area, identify high-impact problems, and communicate technical direction to senior leadership - Publish research, contribute to the broader scientific community, and bring external advances into production systems - Hire, mentor, and develop applied scientists. Grow the team's capabilities to meet evolving customer and business needs 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.