Technical Summary: P-MOSS

Problem Statement

Modern multi-core servers utilizing NUMA (Non-Uniform Memory Access) and Chiplet architectures exhibit significant hardware heterogeneity. This complexity creates a vast design space where the performance of main-memory indexes (e.g., B+-Trees) can vary by up to an order of magnitude depending on two factors:

1. Data Partitioning: The location of index partitions across NUMA nodes (memory controllers).
2. Core Scheduling: The assignment of logical cores to execute queries targeting specific index partitions.

This combination is termed spatial scheduling. Traditional approaches often rely on static heuristics (e.g., round-robin, interleaving) or OS-level defaults, which fail to adapt to the diverse topologies of modern hardware (Intel, AMD, NVIDIA, ARM) and varying workload patterns (read-heavy, write-heavy, scans). Existing works either focus on temporal scheduling (when to run a query) or address data placement for relational tables, often ignoring the critical spatial aspect of core-to-data affinity in main-memory indexes. Furthermore, heuristic-based solutions struggle to generalize across different CPU vendors and architectures due to the combinatorial nature of the scheduling problem.

Methodology: P-MOSS

The paper introduces P-MOSS (Performance MOnitoring Unit-driven Spatial Query Scheduling), a learned framework that leverages principles from Large Language Models (LLMs) to optimize spatial scheduling.

Core Paradigm: Next Token Prediction (NTP)

P-MOSS reframes the spatial scheduling problem as a Next Token Prediction task. Instead of predicting the next word in a sentence, the system predicts the next Core ID in a sequence of scheduling decisions.

• Sequence: The sequence represents the scheduling policy for an index, where each "token" corresponds to the assignment of an index slice (a logical partition based on key range) to a specific CPU core.
• Context: Unlike natural language tokens, Core IDs lack inherent syntactic regularity. P-MOSS addresses this by enriching the "tokens" with context derived from hardware performance statistics.

Architecture and Training Pipeline

P-MOSS utilizes a Decision Transformer (DT) architecture, trained via Offline Reinforcement Learning (RL). The system operates in three distinct stages:

1. Pre-training (Generative Pre-training):

   • Data Source: An offline dataset is curated from diverse hardware (Intel, AMD, NVIDIA, IBM), various NUMA configurations (1NPS, 4NPS, Chiplet), and multiple workloads (YCSB variants).
   • Process: The system probes hardware Performance Monitoring Units (PMUs) during query execution to generate "Hardware Snapshots." These snapshots capture low-level statistics (instruction counts, cache misses, memory traffic) without software-level bookkeeping.
   • Tokenization: The hardware snapshot, current scheduling policy, and performance metrics are tokenized into State, Action, Return-to-Go (RTG), and Meta tokens.
   • Goal: The DT learns generalizable patterns of efficient scheduling across diverse environments.

2. Inference (Online Execution):

   • The trained "base" model predicts scheduling policies for the target hardware and workload.
   • P-MOSS bootstraps with a random or previous policy, monitors PMU statistics in real-time, and feeds them to the DT to infer the optimal mapping of index slices to cores.
   • The system enforces these decisions by migrating index slices to the corresponding NUMA nodes and binding query execution to specific cores.

3. Post-training (Fine-tuning):

   • As the system operates on the target hardware, it collects a "fine-tuning dataset" specific to that environment.
   • The base model is incrementally updated (offline) to create an "assistant" model, adapting to the unique characteristics of the specific server and workload.

Key Technical Components

• Hardware Profiling (PMU): P-MOSS relies exclusively on low-level hardware statistics collected via the PMU (e.g., LLC misses, memory controller traffic). This avoids the overhead and inaccuracy of software-level bookkeeping.
• Tokenization Strategy:
  • State Token: Represents the current hardware state (occupied cores, cache/memory stats) via View, Position, and Machine tokens.
  • Action Token: The predicted Core ID.
  • RTG Token: The desired future reward (target throughput) minus current rewards, guiding the model toward high-performance outcomes.
  • Meta Token: System-wide context (socket count, NUMA topology, CPU vendor).
• Offline RL: By decoupling the learning process from the critical path of query execution, P-MOSS ensures the DBMS does not suffer performance degradation during the learning phase.

Key Contributions

1. P-MOSS Framework: The first learned framework for spatial query scheduling in main-memory indexes on modern NUMA and Chiplet servers, prioritizing both data placement and core scheduling.
2. PMU-Driven Optimization: A novel approach to DBMS kernel optimization that relies solely on low-level hardware statistics (PMU) without software-level bookkeeping, ensuring adaptability across diverse hardware vendors.
3. NTP Formulation: The formulation of query scheduling as a Next Token Prediction task, enabling the use of Offline RL and Decision Transformers to learn scheduling policies without online interaction with the DBMS.
4. Empirical Validation: Extensive evaluation across a wide range of CPUs (Intel, AMD, ARM, IBM), NUMA architectures (1-4 sockets, 2-8 nodes), and workloads.

Experimental Results

The evaluation was conducted on a main-memory B+-Tree with 30M to 1B records across five distinct server configurations (AMD Milan, Intel Skylake X, Intel Sandy Bridge, NVIDIA Grace Hopper, IBM Power).

• Throughput Improvement: P-MOSS achieves up to 6× improvement in query throughput compared to traditional scheduling policies.
• Comparison with Baselines:
  • Against OS defaults (OS:D, OS:I) and NUMA-aware strategies (SE:N, SN:N), P-MOSS shows an average speedup of 1.92× to 3.66×.
  • Against other learned or heuristic-based "Shared Nothing" strategies (Grouped, Mixed, Spread, Random), P-MOSS consistently outperforms them, with improvements ranging from 1.27% to 3.95% depending on the specific hardware and workload.
• Generalization: P-MOSS demonstrates strong generalization capabilities, performing well on unseen hardware (IBM Power) and unseen data distributions (OSM dataset) without retraining from scratch, leveraging its pre-trained base model.
• Overhead: The overhead of collecting hardware statistics via PMU is minimal (~0.73% performance degradation), and the use of SIMD instructions for processing these statistics improves system performance by 60.37% during the profiling phase.

Significance and Claims

The paper claims that P-MOSS demonstrates the viability of applying foundational LLM concepts, specifically Next Token Prediction, to database optimization tasks. By treating scheduling as a sequence modeling problem, the framework can navigate the combinatorial complexity of modern hardware design spaces more effectively than static heuristics.

The authors emphasize that:

• Spatial scheduling is critical for maximizing performance on modern heterogeneous hardware.
• Hardware performance statistics (PMU) provide a robust, low-overhead feature representation for ML-driven database optimization.
• Offline RL allows for the decoupling of learning from execution, enabling safe and scalable adaptation to diverse environments.
• The approach is not limited to spatial scheduling but offers a blueprint for generalizing LLM-inspired techniques to broader DBMS optimization challenges.