The One-Time Warmup: Precompiling Triton Kernels to Salvage FLA for GDN Hidden State Extraction

Introduction

In the middle of a high-stakes optimization sprint to build a hidden state extraction pipeline for training a DFlash speculative decoding drafter, the assistant encountered a stubborn performance bottleneck that threatened to turn an overnight job into a multi-day ordeal. The user's terse observation—"Stil mostly cpu in sys" ([msg 7393])—was the catalyst for a critical re-evaluation of the approach. The assistant's response in [msg 7394] represents a turning point: the realization that a previously discarded solution (flash-linear-attention, or FLA) could work if its one-time Triton JIT compilation cost was paid upfront, rather than letting it fragment across four competing processes. This message is a textbook example of how understanding the nature of a bottleneck—not just its symptoms—can unlock a fix that had been prematurely abandoned.

The Scene: A Pipeline Under Pressure

The broader context is the training of a DFlash drafter for Qwen3.6-27B, a large language model with a GDN (Gated Delta Network) hybrid architecture that interleaves standard causal attention with linear attention layers. To train the drafter, the assistant needed to extract hidden states from the target model for a 913,786-sample dataset. The extraction pipeline used HuggingFace Transformers to run forward passes through the 27B-parameter model, capturing intermediate activations at specific layer positions.

The pipeline had already undergone significant optimization. The assistant had moved from writing safetensors files to disk (which caused ~50% SYS CPU usage from kernel overhead) to writing to /dev/shm tmpfs, which eliminated disk I/O overhead. But the SYS CPU remained high—21–53% during GPU-active periods—because of a deeper issue: PyTorch's SDPA (Scaled Dot-Product Attention) implementation for the GDN linear attention layers was falling back to a CPU-mediated path involving CPU-GPU copies.

The False Start with FLA

The assistant had previously attempted to fix this by installing flash-linear-attention (FLA), a library that provides optimized Triton kernels for linear attention operations like the gated delta rule used in GDN layers. When FLA is available, Transformers can dispatch these operations to GPU-native Triton kernels instead of the PyTorch fallback that requires CPU coordination.

However, the first attempt with FLA was a disaster. As shown in <msg id=7382–7383>, GPU utilization dropped to 0%, CPU usage climbed to 8490%, and throughput collapsed to 3.8–6.5 samples/s per GPU—worse than the 8–11/s without FLA. The culprit was Triton JIT (Just-In-Time) compilation: each of the four parallel extraction processes was independently compiling the same Triton kernels on first use, creating a CPU contention storm. The assistant correctly diagnosed this and uninstalled FLA, reverting to the slower but stable PyTorch fallback.

The Insight: Precompilation Changes Everything

The critical reasoning in [msg 7394] is the assistant's recognition that the problem wasn't FLA itself, but when the compilation happened. The JIT cost is one-time per kernel variant; once compiled, Triton caches the binaries in ~/.triton/cache. If the kernels could be precompiled on a single GPU before launching the parallel extractors, all four processes would reuse the same cache without competing for CPU.

This is a subtle but crucial distinction. The assistant had initially concluded "FLA is SLOWER for inference-only" ([msg 7384]), but the real issue was deployment timing, not fundamental performance. The message opens with this re-diagnosis:

The SYS CPU is from PyTorch's SDPA doing the GDN linear attention layers via CPU-GPU copies. The only real fix is FLA, but it needs Triton JIT warmup. Let me try a different approach — precompile FLA's Triton kernels once, then restart extraction. The JIT cost is one-time.

This reasoning demonstrates a sophisticated understanding of the system's performance profile. The assistant correctly attributes the SYS CPU to the PyTorch SDPA fallback for GDN layers—not to file I/O, not to Python overhead, but to the fundamental attention computation path. And it correctly identifies that FLA's Triton kernels would eliminate this overhead, if only the compilation bottleneck could be sidestepped.

The Execution: A Carefully Crafted Warmup

The bash command in [msg 7394] executes the precompilation strategy with careful attention to detail:

  1. Kill existing extractors: Clean slate, no competing GPU memory usage.
  2. Reinstall FLA: The package had been uninstalled after the previous failed attempt.
  3. Warmup with a single GPU: CUDA_VISIBLE_DEVICES=0 ensures only one process compiles kernels.
  4. Multiple sequence lengths: The warmup runs forward passes at 32, 128, and 512 tokens to trigger compilation of different kernel variants. Triton specializes kernels based on tensor shapes, so a range of sizes ensures broader cache coverage.
  5. Batch warmup: 8×200 and 32×200 sequences compile the batched kernel variants needed for the actual extraction pipeline.
  6. Timing output: Each forward pass is timed, providing immediate feedback on whether compilation is happening (slow first call) versus cache reuse (fast subsequent calls). The choice of warmup sizes is not arbitrary. The extraction pipeline processes batches of sequences with varying lengths, and the warmup covers the likely range. The 32×200 batch test at the end serves as a validation: if the first call is slow (JIT compilation) and subsequent calls are fast (cache hit), the strategy is working.

Assumptions and Potential Pitfalls

The assistant makes several assumptions in this approach:

  1. Triton cache is shared across processes: Triton stores compiled kernels in ~/.triton/cache by default. If the extraction processes run as different users or with different $HOME directories, they wouldn't share the cache. The assistant assumes all processes run as root with the same home directory, which is correct for this setup.
  2. CUDA version compatibility: FLA's Triton kernels must be compatible with CUDA 13.0 on the target machine. The assistant had already verified this in [msg 7379] by importing FLA successfully.
  3. The warmup covers all needed kernel variants: If the extraction pipeline uses tensor shapes not covered by the warmup (e.g., very long sequences or unusual batch sizes), those would still trigger JIT compilation. The assistant mitigates this by choosing representative sizes.
  4. FLA actually accelerates inference: The assumption is that FLA's Triton kernels are faster than PyTorch's SDPA fallback for GDN layers. This is generally true for linear attention operations, but the assistant had no benchmark data for this specific model and hardware combination. The warmup output showing 536 tok/s for 32×200 on a single GPU would later validate this assumption ([msg 7395]).
  5. No side effects from preloading on GPU 0: Loading the full 27B model on GPU 0 consumes ~55GB of VRAM. After warmup, the assistant kills this process, but the Triton cache persists on disk. The extraction processes then load their own model instances on their respective GPUs.

The Output Knowledge Created

[msg 7394] produces several forms of knowledge:

  1. A confirmed root cause: The assistant explicitly states that PyTorch's SDPA for GDN linear attention layers is the source of SYS CPU overhead. This is a diagnosis that had been hypothesized but not proven until the FLA experiment showed the alternative path.
  2. A validated strategy: The precompilation approach is a general technique applicable to any Triton-based library (e.g., flash-attn, vLLM kernels) where JIT compilation overhead is a concern in multi-process deployments.
  3. A reusable warmup script: The Python warmup code can be adapted for other models and pipelines, providing a template for kernel precompilation.
  4. Performance baseline data: The warmup output shows the model loading time and forward pass speeds at various sizes, which serves as a reference for future optimization efforts.

The Thinking Process: A Window into Debugging Methodology

The reasoning visible in [msg 7394] reveals a structured debugging methodology:

  1. Symptom observation: High SYS CPU during GPU-active periods.
  2. Hypothesis generation: The SYS CPU comes from PyTorch SDPA's CPU-GPU copies for GDN layers.
  3. Solution identification: FLA would eliminate these copies by providing GPU-native kernels.
  4. Obstacle recognition: Previous FLA attempt failed due to Triton JIT compilation contention.
  5. Reframing: The problem isn't FLA itself but the timing of compilation.
  6. Strategy formulation: Precompile once, then deploy.
  7. Execution: Carefully structured warmup covering multiple kernel variants. This is not a linear process—the assistant had already tried and discarded FLA. The user's prompt ("Stil mostly cpu in sys") forced a re-examination of the discarded solution. The assistant didn't invent new information; it recontextualized existing knowledge (FLA exists, Triton JIT is expensive, cache is shared) into a new strategy.

The Broader Significance

This message exemplifies a pattern that recurs throughout the opencode session: the gap between research-quality code and production deployment. The DFlash training pipeline relies on the speculators library, which was designed for standard Transformer architectures. Qwen3.6-27B's GDN hybrid attention breaks the assumptions baked into both the extraction pipeline and the serving frameworks (vLLM, SGLang). Every optimization requires understanding the intersection of model architecture, framework internals, and hardware characteristics.

The precompilation strategy is also a microcosm of a larger theme: the assistant's willingness to tear down and rebuild when a solution doesn't work. Rather than accepting the 8–11 samples/s throughput as "good enough" (it would have completed in 8–10 hours), the assistant recognized that the SYS CPU indicated a fundamental inefficiency worth fixing. The result—as shown in [msg 7395] where the warmup achieved 536 tok/s on a single GPU—would dramatically accelerate the extraction, potentially reducing completion time from hours to minutes for the warmup validation, and significantly improving the overall pipeline throughput.

Conclusion

[msg 7394] is a masterclass in performance debugging under pressure. It demonstrates that the most valuable optimization insights often come not from discovering new techniques, but from re-examining discarded solutions through a new lens. The assistant's willingness to revisit FLA—a library it had just uninstalled and declared slower—reflects intellectual honesty and a deep understanding of system behavior. The precompilation strategy is elegant in its simplicity: pay the one-time compilation cost upfront, then reap the benefits across all parallel workers. It's a lesson that applies far beyond this specific pipeline, wherever Triton JIT compilation creates a bottleneck in distributed or multi-process GPU workloads.