The Triton Cache Gambit: Precompiling JIT Kernels to Salvage a Hidden State Extraction Pipeline

Introduction

In the sprawling effort to train a better DFlash speculative decoding drafter for Qwen3.6-27B, the assistant and user had constructed an elaborate offline hidden state extraction pipeline. The pipeline used HuggingFace Transformers to run forward passes through the 27B-parameter GDN hybrid model, capturing intermediate hidden states that would serve as training data for a 2B-parameter drafter. After migrating to a new 4× RTX PRO 6000 Blackwell node and installing dependencies, the extraction was running at a disappointing 8–11 samples per second per GPU—far below the throughput needed to process a 914K-sample dataset in a reasonable timeframe.

The bottleneck was perplexing. GPU utilization hovered near zero while CPU usage was dominated by system time (the "sy" field in top), sometimes exceeding 50% of all CPU cycles. The assistant had already tried multiple fixes: switching from per-sample safetensors writes to batched GPU-side concatenation, moving file I/O from the container overlay filesystem to /dev/shm tmpfs, and integrating async S3 uploads. Each fix helped incrementally, but the fundamental problem remained—the model's GDN (Gated Delta Network) linear attention layers were falling back to a PyTorch implementation that performed CPU-GPU transfers on every forward pass.

The solution seemed obvious: install flash-linear-attention (FLA), a library of optimized Triton kernels that could accelerate the GDN layers on GPU, eliminating the CPU overhead entirely. But when the assistant first tried this, the result was a disaster.

The False Start: FLA Makes Everything Worse

When the assistant initially installed FLA and restarted the four parallel extractors (msg 7382–7384), the results were catastrophic. CPU usage climbed to 8490% across all cores, GPU utilization stayed at 0%, and the throughput actually dropped to 3.8–6.5 samples/s—roughly half the already-poor baseline rate. The extractors were spending all their time in Triton's Just-In-Time (JIT) compilation system, compiling GPU kernels for the first time.

The problem was that Triton compiles GPU kernels lazily: the first time a particular kernel configuration is encountered, it must be compiled from Python to PTX to machine code. With four independent extractor processes each encountering GDN layers for the first time, all four were simultaneously invoking the Triton JIT compiler, overwhelming the CPU cores. The compilation itself is CPU-intensive (involving LLVM-based optimization and code generation), and while compilation is happening, the GPUs sit idle waiting for compiled kernels.

The assistant's initial diagnosis was mistaken in an interesting way. In msg 7384, the assistant speculated that "FLA might be SLOWER for inference-only than the PyTorch fallback in some cases—FLA is optimized for training (recurrent mode) not for the kind of forward-only extraction we're doing." This was a reasonable hypothesis—FLA's fused_recurrent_gated_delta_rule kernel is designed for the recurrent computation used in training, which could theoretically be slower than the parallel PyTorch SDPA fallback for short sequences. But the real culprit was the JIT compilation overhead, not the kernel performance itself. The assistant uninstalled FLA and reverted to the PyTorch fallback, settling for the 8–11 samples/s baseline and planning to let the extraction run overnight.

The User's Observation and the Pivot

The user then observed (msg 7393): "Stil mostly cpu in sys." This simple observation reframed the problem. The high system CPU time wasn't going away—it was intrinsic to the PyTorch SDPA fallback path for GDN layers. Each of the model's 48 GDN layers was performing CPU-side operations and CPU-GPU synchronization, and this overhead was the hard ceiling on throughput.

The assistant's response (msg 7394) shows a critical insight: "The only real fix is FLA, but it needs Triton JIT warmup." The key realization was that the JIT compilation was a one-time cost, not a recurring one. Triton caches compiled kernels in ~/.triton/cache, and subsequent invocations reuse the cached binaries. The mistake in the first attempt was launching all four extractors simultaneously—they all tried to compile the same kernels in parallel, creating a CPU contention storm. The fix was to compile the kernels once on a single GPU before any extractors started, then let all four extractors share the pre-populated cache.

This is a classic systems optimization pattern: amortize a fixed upfront cost across many subsequent operations. The compilation cost is paid once; the benefit (faster kernel execution) is reaped for every forward pass thereafter.

Message 7395: The Warmup Execution

The subject message (msg 7395) documents the execution of this strategy. The assistant runs a warmup script on a single GPU that:

  1. Loads the full 27B Qwen3.6-27B model onto GPU 0
  2. Runs forward passes at multiple sequence lengths (32, 128, 512 tokens) to trigger compilation of different kernel variants
  3. Runs batch forward passes (8×200, 32×200 tokens) to trigger batch kernel variants
  4. Reports the timing of each warmup pass The results are striking. The first forward pass (32 tokens) takes 10 seconds—this is the Triton JIT compilation time. The second pass (128 tokens) takes 0.12 seconds—a >80× speedup, showing that the kernels are now cached and running at full speed. The batch test (32×200 tokens) achieves 536 tokens per second, confirming that the FLA kernels are dramatically faster than the PyTorch fallback. The assistant then restarts all four extractors, this time with the Triton cache already populated. The command is nearly identical to the earlier failed attempt, but with one crucial difference: the LD_LIBRARY_PATH is set to include CUDA libraries, and the Triton cache at ~/.triton/cache already contains all the compiled kernels. The extractors will load the cached kernels instead of trying to compile them from scratch.

Input Knowledge Required

To understand this message, one needs knowledge of several technical domains:

Triton compiler internals: Triton is a domain-specific language and compiler for GPU kernels. It uses a JIT compilation model where Python code defining a kernel is compiled to Triton IR, then optimized and lowered to PTX (NVIDIA's GPU assembly language), then to machine code. The compilation process involves LLVM and can take seconds to minutes for complex kernels. The compiled artifacts are cached on disk in ~/.triton/cache and reused across processes.

Flash Linear Attention (FLA): FLA provides optimized Triton kernels for linear attention mechanisms, including the Gated Delta Network (GDN) used in Qwen3.6-27B. The fused_recurrent_gated_delta_rule kernel implements the recurrent form of GDN attention, which is more efficient than the PyTorch fallback for both training and inference—contrary to the assistant's earlier speculation.

GDN hybrid attention: Qwen3.6-27B uses a hybrid architecture combining standard softmax attention with GDN linear attention layers. The GDN layers are the ones that benefit from FLA; without it, they fall back to a PyTorch implementation that involves CPU-GPU transfers.

Multi-process Triton cache sharing: Triton's kernel cache is stored on disk and is readable by any process on the same machine. This is what makes the warmup strategy viable—one process compiles the kernels, and all other processes can reuse them without recompiling.

CUDA context isolation: Each GPU has its own CUDA context. The warmup runs on GPU 0, but the Triton kernels compiled for GPU 0's architecture (SM120 Blackwell) are also valid for GPUs 1–3 on the same machine, since they share the same microarchitecture.

Output Knowledge Created

This message creates several important pieces of knowledge:

Quantified warmup cost: The first forward pass takes 10 seconds due to Triton JIT compilation. Subsequent passes take 0.12 seconds. This quantifies the one-time overhead that must be paid before FLA becomes beneficial.

Confirmed FLA speedup: The batch throughput of 536 tok/s for 32×200 sequences is dramatically higher than the 8–11 samples/s per GPU achieved with the PyTorch fallback. This confirms that FLA is indeed the right solution for GDN inference, contrary to the earlier speculation that it might be slower.

Validated warmup strategy: The approach of pre-compiling kernels on a single GPU before launching multiple workers is proven to work. The Triton cache is shared across processes, and the JIT overhead is avoided.

Reproducible restart procedure: The exact commands to kill extractors, clean progress state, and relaunch with FLA are documented and can be reused.

Assumptions and Potential Pitfalls

The assistant makes several assumptions in this message:

Triton cache is shared across processes: This is correct—Triton's disk cache at ~/.triton/cache is readable by all processes on the machine. However, if the extractors run in containers with different home directories or if the Triton cache path is overridden by environment variables, this assumption could fail.

All kernel variants are covered by the warmup: The warmup runs sequences of length 32, 128, and 512, and batches of 8 and 32. If the extractor processes encounter sequences of different lengths or batch sizes, they might trigger compilation of additional kernel variants. However, the warmup covers the most common configurations, and any additional compilation would be incremental.

The Triton cache persists across process restarts: The cache is on disk, so it should persist. But if the machine is rebooted or the cache directory is cleaned, the warmup would need to be repeated.

FLA kernels are actually faster than the PyTorch fallback: The warmup results (536 tok/s) suggest dramatic improvement, but this is for a single GPU doing inference without the overhead of hidden state extraction. The actual extractor throughput depends on many factors including data loading, hidden state capture, and S3 upload.

The Thinking Process Visible in the Reasoning

The assistant's reasoning in this message reveals a sophisticated understanding of the system's bottlenecks. The progression from "FLA made things worse" to "FLA needs warmup" to "warmup works and gives 536 tok/s" shows a systematic debugging process:

  1. Observation: High SYS CPU, low GPU utilization
  2. Hypothesis 1: File I/O to container overlay filesystem (fixed with tmpfs, helped marginally)
  3. Hypothesis 2: PyTorch SDPA fallback for GDN layers causes CPU-GPU transfers
  4. Attempted fix: Install FLA to accelerate GDN layers
  5. Failure mode: FLA causes 8490% CPU usage, throughput drops further
  6. Diagnosis: Triton JIT compilation overhead from 4 parallel processes
  7. Hypothesis 3: FLA might be slower for inference than PyTorch fallback (incorrect)
  8. Revert: Uninstall FLA, continue with PyTorch fallback
  9. User observation: "Still mostly cpu in sys" reframes the problem
  10. New insight: The JIT compilation is one-time; pre-compile on one GPU, share cache
  11. Validation: Warmup shows 10s first pass, 0.12s subsequent passes, 536 tok/s batch
  12. Execution: Restart extractors with prewarmed cache This is a textbook example of how systems debugging requires iterating through hypotheses, testing each one, and using observations to refine understanding. The assistant's initial dismissal of FLA as "optimized for training" was a reasonable but incorrect inference—the real issue was compilation overhead, not kernel performance. The user's concise observation ("Stil mostly cpu in sys") provided the nudge needed to revisit the FLA approach with a different strategy.

Conclusion

Message 7395 represents a turning point in the hidden state extraction pipeline. By recognizing that Triton's JIT compilation is a one-time cost that can be amortized across all workers, the assistant transforms FLA from a performance disaster into a dramatic speedup. The 536 tok/s batch throughput suggests that the extraction pipeline could see a 5–10× improvement over the 8–11 samples/s baseline, potentially reducing the 8–10 hour ETA to under an hour.

The deeper lesson is about the nature of JIT compilation in ML systems. Triton's lazy compilation model is designed for interactive use where the first invocation is acceptable slow, but it creates pathological behavior when multiple processes simultaneously trigger compilation. The fix—pre-compiling on a single worker and sharing the cache—is simple in retrospect but required understanding both the Triton caching mechanism and the specific failure mode of parallel JIT contention. This kind of systems knowledge, built through trial and error, is what separates a working pipeline from a production-grade one.