A Deep Dive into Building a Modern AI Toolchain

08 Mar, 2026

If you're like me, you're probably always trying to squeeze as much performance out of your systems and make each GPU hour count. It's a constant battle against bottlenecks, inefficiencies, and the sheer complexity of modern hardware. For a while, I've been working on a project called PyC, a toolchain of technologies with some experimental methods I'm trying to run and see if they have meaningful performance improvements. I just wanted to detail what I have so far.

This journey culminated in a moment that made all the effort worthwhile. It started with a Realization while reading the flash attention 4 paper:

The math hasn't changed — modern AI isn't limited by compute anymore, it's limited by how efficiently the machine moves data and keeps the pipeline running.

This blog post is the story of that journey. It's about the architectural decisions, the hard-won lessons, and the surprising bottlenecks I discovered while building a system designed to squeeze every last drop of performance out of modern AI hardware. It's a story that takes us from NUMA-aware memory allocators in Rust to policy-driven kernel selection in C, and finally, to running 4-bit quantized diffusion models that experts can't distinguish from full precision.

Modern AI is a systems discipline. This is what we learned.

The Vision: A Unified HPC Toolchain

The goal was always bigger than just one project. Over the years, I had built Nexa_Inference for optimized serving and Nexa_Vortex as a powerful runtime. But the real vision was to merge them into a single, cohesive toolchain for High-Performance Computing (HPC) called PyC. This open-source project, which you can find on my GitHub under the name DarkStarStrix, is an ambitious attempt to create the ultimate toolchain for HPC. The idea was to create a vertically integrated stack, from the Python user interface down to the metal, that could intelligently adapt to the workload and the hardware.

This wasn't about building yet another deep learning framework. It was about building the ultimate toolchain that would sit underneath the frameworks, providing a level of performance and control that is impossible to achieve when you treat the compiler, the runtime, and the kernels as separate, black-box components.

The final architecture is a multi-layered system where each component is designed to communicate with the others:

This stack is composed of five distinct layers:

Building this system required a constant dialogue between layers. The runtime needed to understand the hardware to inform the compiler's decisions. The compiler needed a way to select from a library of specialized kernels. And the whole system needed to be observable through a lightweight telemetry pipeline. Let's break down how each layer works.

The Engine Room: The Rust Vortex Runtime

At the heart of PyC is the Vortex Runtime, a component written entirely in Rust for safety, performance, and fearless concurrency. Its primary job is to act as the engine room of the toolchain, orchestrating the complex dance of data movement and computation. It's designed around an asynchronous "conveyor belt" model, ensuring that the expensive GPU hardware never sits idle.

This is not just a simple task dispatcher. The Vortex runtime is a sophisticated system that includes several key components:

• Hardware Profiler (hw_profile.rs): Before any work begins, the runtime inspects the machine. It detects the number of CPU cores, the memory available, and, most importantly, the NUMA (Non-Uniform Memory Access) topology. On multi-socket server CPUs, accessing memory on a "remote" NUMA node is significantly slower. The profiler identifies which NUMA node is physically closest to the GPU, a critical piece of information for our memory allocator. As others in the field have noted, this awareness of NUMA topology is crucial for high-performance deep learning.

• NUMA-Aware Pinned Memory Allocator (allocator.rs): When you transfer data from the CPU to the GPU, the CUDA driver typically performs a hidden staging copy if the memory is pageable. By using "pinned" (or page-locked) memory, we can eliminate this copy and nearly double the effective host-to-device bandwidth. Our allocator goes a step further: it uses the information from the hardware profiler to ensure that the pinned memory is allocated on the NUMA node closest to the target GPU. This simple trick reduces latency by another ~19% on dual-socket systems. The allocator also pre-warms a pool of these buffers at startup to avoid allocation overhead in the critical path.

• Asynchronous Pipeline (pipeline.rs): This is the orchestrator. It takes a compiled module from the C layer and manages its execution. It uses the tokio runtime to create an asynchronous pipeline where CPU preprocessing, DMA transfers (CPU to GPU), and GPU computation can all happen in parallel. While the GPU is busy with one batch, the CPU is already preparing the next one. This overlap is the key to achieving high utilization.

• Telemetry Sink (telemetry.rs): To understand what's happening inside the runtime without slowing it down, we use a lightweight, non-blocking telemetry system. The hot path emits structured events (like BatchComplete or KernelSelected) into a crossbeam channel. A separate, low-priority thread listens on this channel, serializes the events to JSON, and logs them. This gives us perfect observability with near-zero performance impact.

Here is a snippet from the pipeline that shows how these pieces come together:

// runtime/vortex_core/src/pipeline.rs

/// Configuration for the pipeline, derived from hardware topology.
#[derive(Debug, Clone)]
pub struct PipelineConfig {
    pub cpu_workers: usize,
    pub queue_depth: usize,
    pub policy_mode: pyc_objective_mode,
    pub memory_budget_bytes: usize,
    pub numa_node: Option<usize>,
}

impl PipelineConfig {
    /// Derive sensible defaults from the detected hardware topology.
    pub fn from_hardware(hw: &HardwareProfile) -> Self {
        PipelineConfig {
            cpu_workers: (hw.cpu_cores / 2).max(2),
            queue_depth: hw.gpu_count.max(1) * 4,
            policy_mode: pyc_objective_mode::PYC_MODE_UTILIZATION_FIRST,
            memory_budget_bytes: 0,
            numa_node: hw.gpu_numa_node, // Use the NUMA node closest to the GPU!
        }
    }
}

This tight integration between hardware awareness and runtime configuration is a core design principle of PyC. The system doesn't just run code; it adapts its execution strategy to the physical layout of the machine.

The Brains of the Operation: The C Compiler Layer

If the Rust runtime is the engine, the C compiler layer is the brain. This is where the high-level user intent is translated into a concrete execution plan. It's responsible for analyzing the incoming computation, applying optimizations, and, most critically, selecting the right GPU kernel for the job. This layer is written in C for portability and tight control over memory layout and FFI (Foreign Function Interface) boundaries.

Key components include:

• IR (Intermediate Representation): A simple, stable IR defines the operations, tensors, and data types. This provides a contract between the compiler and the runtime.

• Pass Manager: A standard pass manager architecture allows us to apply a series of transformations to the IR, such as operator fusion or layout optimization.

• AI Bridge (ai_bridge.c): This is a fascinating component that translates high-level policies into concrete compiler options. For example, a user might specify a policy of PYC_MODE_UTILIZATION_FIRST or PYC_MODE_MEMORY_FIRST. The AI bridge maps this abstract goal to low-level knobs that control kernel selection and memory budgeting.

• Kernel Registry (kernel_registry.c): This is the centerpiece of the compiler. It's a dynamic registry that holds a list of all available GPU kernels for a given operation (like matmul). Each kernel is registered with a rich set of metadata.

When the compiler needs to select a kernel, it doesn't just pick the fastest one. It runs a scoring function that balances all of these factors against the current policy and the state of the system.

For example, if the PYC_MODE_MEMORY_FIRST policy is active and the runtime has reported high memory pressure, the scoring function will apply a penalty to kernels that use a lot of shared memory or have high register pressure. This might cause it to select a slower but more memory-efficient kernel, preventing an out-of-memory error and keeping the pipeline flowing. This dynamic, policy-driven selection is what allows PyC to adapt to different hardware and workloads.

Here is the heart of the scoring logic:

// compiler/runtime/kernel_registry.c

static double kernel_pressure_penalty(
    const pyc_kernel_desc* desc, 
    pyc_objective_mode mode, 
    double pressure_score) {
    if (pressure_score <= 0.0 || mode == PYC_MODE_UTILIZATION_FIRST) {
        return 0.0;
    }
    return pressure_score * (double)(desc->shared_mem_bytes / 1024U + 
                                  (size_t)(desc->reg_pressure_class * 8));
}

static double kernel_score(
    const kernel_slot* slot, 
    pyc_objective_mode mode, 
    double pressure_score, 
    double* out_penalty) {
    double base = (double)slot->desc.priority * 100.0;
    double occ_weight = mode == PYC_MODE_UTILIZATION_FIRST ? 12.0: 6.0;
    double util = occ_weight * slot->desc.estimated_occupancy;
    double tensor_core_bonus = slot->desc.tensor_core_eligible ? 25.0: 0.0;
    double time = kernel_time_component(slot->bench.best_time_ms);
    double penalty = kernel_pressure_penalty(&slot->desc, mode, pressure_score);
    
    if (out_penalty) {
        *out_penalty = penalty;
    }
    return base + util + tensor_core_bonus + time - penalty;
}

This tight feedback loop between the runtime (which measures pressure) and the compiler (which selects kernels) is what makes PyC so powerful.

The Muscle: Hand-Tuned CUTLASS Kernels

At the bottom of the stack, we have the muscle: a library of highly optimized GPU kernels built using CUTLASS. CUTLASS is a C++ template library from NVIDIA that provides a framework for building high-performance matrix multiplication (GEMM) and convolution operations. It gives you fine-grained control over every aspect of the kernel, from threadblock shapes to the software pipeline stages.

A surprising realization when working with modern accelerators like H100 or B200 GPUs is that raw arithmetic throughput is rarely the limiting factor. The real constraint tends to be memory movement. This is why techniques like FlashAttention focus on reorganizing algorithms to minimize memory traffic.

We don't just use one generic kernel. We implement a hierarchy of them, each tuned for a specific purpose. For a standard matmul operation, our registry contains at least three variants:

1. cutlass_gemm_tensorcore_f16 (Priority 100): This is our champion kernel. It uses 16-bit floating-point (FP16) precision and is designed to run on the GPU's Tensor Cores. It offers the highest theoretical throughput and is ideal for inference workloads where peak performance is critical. It has a measured occupancy of 87% on an A100 GPU.

2. cutlass_gemm_tensorcore_bf16 (Priority 90): This kernel also uses Tensor Cores but with BFloat16 (BF16) precision. BF16 has a lower precision than FP16 but a much wider dynamic range, which makes it more resilient to overflow and underflow during training. It's the preferred choice for many training workloads.

3. cutlass_gemm_simt_f32 (Priority 10): This is our universal fallback. It uses standard 32-bit floating-point (FP32) and runs on the GPU's SIMT (Single Instruction, Multiple Thread) cores, not the Tensor Cores. It's slower but compatible with a wider range of hardware and provides a crucial baseline for correctness and a safe option when memory pressure is extremely high.

These kernels are registered with the C compiler layer at library load time. The cutlass_gemm.cu file contains a registration function that populates the kernel registry with the metadata needed for the policy-driven selection logic.

// compiler/cutlass_kernels/cutlass_gemm.cu

extern "C" void pyc_cutlass_register_gemm_kernels(void) {
    pyc_kernel_desc desc;

    /* --- FP16 Tensor Core GEMM --- */
    memset(&desc, 0, sizeof(desc));
    strncpy(desc.op_key,  "matmul",                       PYC_KERNEL_OP_KEY_MAX - 1);
    strncpy(desc.symbol,  "cutlass_gemm_tensorcore_f16",  PYC_KERNEL_SYMBOL_MAX - 1);
    desc.backend              = PYC_BACKEND_CUDA;
    desc.priority             = 100;          /* highest priority */
    desc.estimated_occupancy  = 0.87;         /* measured on A100 */
    desc.tensor_core_eligible = 1;
    pyc_kernel_register(&desc);

    /* --- BF16 Tensor Core GEMM --- */
    // ... (similar registration for BF16 and FP32 kernels)
}

This explicit registration of kernel capabilities is what allows the higher-level compiler to make intelligent, hardware-aware decisions.

Scaling Out: The Distributed Layer

Modern AI models are often too large to fit on a single GPU, and training them requires the coordinated effort of a whole cluster. This is where the distributed layer comes in. PyC's distributed runtime is designed to be pluggable, allowing it to use different communication backends depending on the hardware and environment.

We primarily use NCCL (NVIDIA Collective Communications Library), which is the industry standard for high-performance communication on NVIDIA GPUs. It provides highly optimized implementations of collective operations like AllReduce, AllGather, and ReduceScatter. These are the building blocks of distributed training algorithms like Fully Sharded Data Parallel (FSDP), a technique detailed in foundational work like Microsoft's ZeRO paper and further explored in recent research on communication optimization for distributed training.

Our NCCL backend (comm_backend_nccl.c) is a thin wrapper around the NCCL library. It dynamically loads the libnccl.so shared library at runtime, which means PyC doesn't have a hard dependency on NCCL being installed. If the library isn't found, the distributed capabilities are simply disabled. This makes the toolchain more portable.

// compiler/runtime/comm_backend_nccl.c

static pyc_comm_status nccl_all_reduce(
    void* backend_ctx,
    pyc_comm_handle_t comm,
    const void* send_buf,
    void* recv_buf,
    size_t count,
    pyc_dtype dtype,
    pyc_reduce_op op,
    void* stream) {
    
    pyc_nccl_backend_ctx* ctx = (pyc_nccl_backend_ctx*)backend_ctx;
    nccl_data_type_t nccl_dtype;
    nccl_red_op_t nccl_op;

    // ... validation and type mapping ...

    return map_nccl_result(ctx->all_reduce(send_buf, recv_buf, count, 
                                          nccl_dtype, nccl_op, 
                                          (nccl_comm_t)comm, stream));
}

This pluggable design also allows us to support other backends like MPI (Message Passing Interface) or AMD's RCCL for different hardware environments, making PyC a truly cross-platform HPC toolchain.

The Payoff: The 4-Bit Diffusion Model

This brings us to the 4-bit quantized Qwen Image Model. This project was the ultimate test of the PyC toolchain. Could our vertically integrated stack, with all its layers of abstraction and policy-driven decisions, actually deliver on the promise of extreme performance without sacrificing quality?

The challenge was immense. Diffusion models are notoriously sensitive to quantization. A naive post-training quantization (PTQ) approach, where you simply quantize the weights after training, results in catastrophic image degradation. The subtle velocities that guide the denoising process are easily corrupted.

This is where Quantization-Aware Distillation (QAD) comes in. Instead of just quantizing a trained model, we use a full-precision "teacher" model to guide the training of a 4-bit "student" model. The student has the same architecture as the teacher, but all its linear layers are quantized to 4-bits. During training, both models receive the same input (a noisy latent, a timestep, and a text embedding). The loss is calculated as the difference between the teacher's velocity prediction and the student's velocity prediction. The gradients are then backpropagated only through the student model.

This process forces the student model to learn how to mimic the behavior of the full-precision teacher, even with its limited 4-bit weights. It learns to compensate for the quantization errors.

This is where PyC shines. The QAD training process is incredibly demanding. It requires running two large models simultaneously, performing distributed training across multiple GPUs, and executing custom quantized kernels. Our full-stack toolchain was able to handle this complexity with ease.

The result? A 4-bit quantized diffusion model that produces images virtually indistinguishable from the full-precision original, while running significantly faster and using a fraction of the memory. The fact that only one person on a team of seven experts could spot the difference is a testament to the power of this approach.

Economic Realities and the Need for Observability

When running workloads on rented GPU clusters, inefficiency is not just a technical problem—it is a financial one. High-end GPUs such as H100 nodes often cost tens of dollars per hour, meaning that idle compute time directly translates into wasted money. A poorly optimized pipeline can burn through significant budgets without producing meaningful progress.

This reality changes the way engineers think about system performance. High utilization, therefore, becomes a key metric in production environments. Even small inefficiencies can accumulate across thousands of training steps and become expensive over time. But you can't optimize what you can't see. This is where a comprehensive, real-time telemetry dashboard becomes an indispensable tool.

This dashboard, generated directly from the telemetry data captured by PyC's Vortex runtime, gives us a holistic, multi-faceted view of the system's health and performance. We can see at a glance:

• Compute Utilization: Are all GPUs firing on all cylinders, or are some lagging behind?
• HBM Memory Usage: Are we approaching the memory capacity of any GPU, risking an out-of-memory error?
• Power Draw and Temperature: Are the GPUs operating within their thermal design power (TDP) and temperature limits, or are they throttling?
• Utilization Heatmap: How does utilization vary over time and across different GPUs? Are there periodic dips that suggest a data loading bottleneck?

This level of observability is not a luxury; it is a necessity for performance engineering. It allows us to move from guesswork to data-driven optimization, identifying the true bottlenecks in the system and measuring the impact of our changes.

Conclusion: The Future is Vertical

The journey of building PyC has been a powerful lesson in modern AI engineering. The biggest gains no longer come from purely algorithmic innovations. They come from a deep, holistic understanding of the entire system, from the user's Python script down to the silicon.

Vertical integration is the future. When the compiler, the runtime, the kernels, and the distributed layer are all designed to work together, they can achieve a level of performance and adaptability that is simply out of reach for a collection of disparate, black-box tools. The ability to make policy-driven decisions that propagate through the entire stack—from selecting a NUMA node for memory allocation to choosing a specific CUTLASS kernel based on memory pressure—is a game-changer.

The 4-bit diffusion model is just one example of what's possible. As models become larger and more complex, the need for this kind of deeply integrated, hardware-aware toolchain will only grow. The future of AI belongs to those who are willing to build the whole machine.

Thanks for reading! If you found this blog insightful, share it with your colleagues and others

Link to full repo: https://github.com/DarkStarStrix/PyC

We quantized every linear layer in the Qwen Image Model to 4 bits. Every. Single. One.
Only one out of the seven people on my team correctly guessed which image was quantized.
ViDiT-Q (ICLR 2025) called it a training bottleneck. PTQ4DiT (NeurIPS 2024) called it non-trivial. Q-DiT (CVPR 2025) called it performance degradation.
We called it Thursday.