From 7 to 155 Samples Per Second: The Hidden State Extraction Optimization Odyssey

Introduction

In the span of roughly 40 messages spanning a few hours of wall-clock time, a hidden state extraction pipeline for training a DFlash speculative decoding drafter underwent a transformation so dramatic that it changed the project's trajectory from "overnight batch job" to "coffee break task." The journey from 7–11 samples per second per GPU to 140–155 samples per second—a 17× improvement—was not the result of buying faster hardware or switching to a more efficient model architecture. It was the product of systematic debugging, a series of false starts, one critical architectural insight, and the operational challenges that emerged when a pipeline became too fast for its own good.

This article synthesizes the full arc of that optimization journey, drawing on the 28 message articles that document each step. It traces the path from the initial bottleneck diagnosis through the failed FLA experiment, the Triton kernel prewarming gambit, the GPU-side concatenation breakthrough, the documentation pivot, the tmpfs overflow crisis, the backpressure fix, and the final verification that the pipeline was robust enough to run unattended. Along the way, it examines the assumptions, mistakes, and engineering judgments that shaped each decision.

The Problem: A Pipeline in the Doldrums

The goal was straightforward: extract hidden states from the Qwen3.6-27B model—a 27-billion-parameter language model with a GDN (Gated Delta Network) hybrid attention architecture—for a 913,786-sample training dataset. These hidden states would serve as training targets for a 2B-parameter DFlash speculative decoding drafter, a small model that learns to predict the target model's internal representations, enabling faster inference through speculative decoding [1].

The initial extraction pipeline, built on HuggingFace Transformers and running across four NVIDIA RTX PRO 6000 Blackwell GPUs (96GB each), achieved a meager 7–11 samples per second per GPU. More concerning than the raw throughput was the resource profile: CPU utilization was dominated by system-level overhead, with the "sy" field in top showing approximately 50% of CPU cycles consumed in kernel mode [2]. GPU utilization hovered near zero for extended periods, punctuated by brief spikes to 50–100% during the forward pass [3].

The assistant systematically investigated the root causes. The first hypothesis was filesystem I/O: the pipeline was writing per-sample safetensors files to the container's overlay filesystem, which incurs significant kernel overhead. Moving writes to /dev/shm (tmpfs, a RAM-backed filesystem) eliminated the filesystem overhead but only improved throughput marginally, from 7–11 to 8–11 samples/s per GPU [4]. The kernel-mode CPU time remained high during GPU-active periods, suggesting a deeper issue.

The False Start: Flash Linear Attention and the Triton JIT Disaster

The assistant identified that the GDN hybrid attention layers in Qwen3.6-27B were falling back to PyTorch's SDPA (Scaled Dot-Product Attention) implementation, which performs CPU-GPU synchronization and data transfers for the linear attention components. The solution seemed obvious: install flash-linear-attention (FLA), a library of optimized Triton kernels that could accelerate the GDN layers entirely on GPU [5].

The result was catastrophic. When the assistant installed FLA and launched four parallel extractors, CPU usage spiked to 8,490% across all cores, GPU utilization dropped to 0%, and throughput actually degraded to 3.8–6.5 samples/s per GPU—roughly half the already-poor baseline rate [6]. The root cause was Triton's Just-In-Time (JIT) compilation system: each of the four extractor processes independently encountered the GDN linear attention layers for the first time, triggering simultaneous compilation of the same Triton kernels. 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 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" [7]. This was a reasonable hypothesis but ultimately incorrect. The real culprit was the JIT compilation contention, not the kernel performance itself. The assistant uninstalled FLA and reverted to the PyTorch fallback, settling for the 8–11 samples/s baseline.

The Triton Cache Gambit: Precompiling JIT Kernels

The user's observation—"Stil mostly cpu in sys" [8]—reframed the problem. The high system CPU time wasn't going away; it was intrinsic to the PyTorch SDPA fallback path for GDN layers. The assistant realized 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 [9]. The mistake in the first attempt was launching all four extractors simultaneously—they all tried to compile the same kernels in parallel. 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.

The warmup script ran forward passes at multiple sequence lengths (32, 128, 512 tokens) and batch sizes (8×200, 32×200) to trigger compilation of different kernel variants [10]. The results validated the approach: the first forward pass took 10 seconds (Triton JIT compilation), but subsequent passes completed in 0.12 seconds—an 80× speedup. A batch test achieved 536 tokens per second, confirming that the FLA kernels were dramatically faster than the PyTorch fallback [11].

With the Triton cache populated, the assistant relaunched all four extractors. But the fundamental bottleneck remained: GPU utilization was still low, and CPU system time was still high during forward passes. The Triton prewarming had eliminated the JIT compilation overhead, but the underlying data movement pattern was still inefficient [12].

The Breakthrough: GPU-Side Concatenation

The user shared a screenshot showing the telltale pattern: GPUs spiking to 50–100% utilization for a few seconds during the forward pass, then dropping to 0% for extended periods while CPU usage remained high, with hundreds of S3 upload subprocesses consuming CPU cycles [13]. The assistant's analysis was precise:

"The issue: the save_file call + s3_queue.put() happens synchronously after each batch, and the save_file to tmpfs is actually fast BUT the torch tensor .cpu() and .to(torch.bfloat16) and concatenation in the per-sample loop after the forward pass is what's eating all that CPU time. With 545 samples per batch, we're doing 545 × 5 layer captures × cpu copy × concat × bfloat16 conversion." [14]

This was the critical insight. The extraction code was performing 2,725 individual GPU→CPU memory transfers per batch (545 samples × 5 captured layers). Each transfer involved a .cpu() call that synchronizes GPU and CPU, a .to(torch.bfloat16) conversion, and a concatenation operation—all on the CPU side. The cumulative overhead was devastating: the GPU would finish its forward pass in seconds, then sit idle while the CPU spent tens of seconds processing these 2,725 tiny tensor copies one by one [15].

The fix was elegantly simple: concatenate all hidden states on the GPU using torch.cat, then perform a single .cpu() transfer for the entire batch. This reduced 2,725 individual copies to just 5 (one per layer, for the entire batch) [16]. The dtype conversion could also happen on GPU before the transfer, leveraging the GPU's massively parallel compute units rather than the CPU.

The result, verified in a monitoring loop [17], was transformative:

The Documentation Pivot: Capturing Knowledge at the Moment of Breakthrough

Immediately after witnessing this breakthrough, the user issued a simple but critical instruction: "save progress so far and detailed plan in the /data/.. dir" [19]. This request reflected a mature engineering instinct: when a pipeline is running well—especially after a period of intense debugging and optimization—it is the right moment to capture what was learned, what was built, and what remains to be done.

The assistant responded by writing three substantial documents. PROGRESS.md (7,493 bytes) captured the optimization journey: the migration from kpro6 to kpro5, the MTP speculation deployment achieving 73.5 tok/s, the DFlash investigation revealing three separate vLLM bugs (layer-ID offset, SWA layer handling, and eagle cache drops), the pivot to offline extraction, and the performance optimization culminating in the GPU-side concat fix [20]. PLAN.md (5,855 bytes) laid out the roadmap ahead: completing hidden state extraction, setting up the training environment, launching the DFlash drafter training run, evaluating the trained drafter, and potentially iterating on DDTree integration [21]. STATUS.md (709 bytes) provided a real-time snapshot of the extraction pipeline's progress [22].

This documentation pivot was not merely a housekeeping exercise. It created a persistent shared artifact that both the user and the assistant could reference, ensuring that the hard-won knowledge about what worked—GPU-side concatenation, tmpfs writes, async S3 uploads, Triton kernel prewarming—would survive conversation truncation, session restarts, and the inevitable context decay that occurs as new messages push older ones out of the assistant's attention window [23].

The New Bottleneck: When Optimization Breaks the Buffer

But success created a new vulnerability. The pipeline wrote extracted hidden states as safetensors files to /dev/shm (a 251GB tmpfs partition), then queued them for asynchronous upload to S3. The extraction was now so fast that it filled tmpfs before the S3 uploader could drain it [24]. The user reported the symptom:

"GPU2: 6%|▌ | 412/7364 [21:12<6:02:37, 3.13s/it] Err: Error while serializing: I/O error: No space left on device (os error 28)" [25]

The df output confirmed the diagnosis: shm 251G 251G 0 100% /dev/shm. The user added the prescient observation: "Probably want some backpressure around 90%.. S3 uploads don't entirely keep up."

The assistant's response was immediate: "Need backpressure: pause extraction when tmpfs is >80% full" [26]. The fix was a simple check before each batch: if /dev/shm usage exceeds 80%, sleep for a configurable interval and retry. This is a classic producer-consumer flow control pattern, applied to an ML pipeline context [27].

The Recovery: Marker-Based Resume and the Pragmatic Restart

The assistant deployed the backpressure fix and attempted to restart the pipeline, but the SSH command timed out after 15 seconds, leaving the system in an uncertain state [28]. The user then asked a critical question: "Do we resume from S3 state and cleanup?" [29]

The assistant began investigating S3 reconciliation—querying the bucket to count uploaded files, comparing against local marker files, and determining exactly which batches needed reprocessing [30]. But before the investigation could complete, the user cut in: "No, it's cheap enough to just resume" [31].

This seven-word message was a masterclass in pragmatic engineering judgment. The user performed a real-time cost-benefit analysis: the cost of re-extracting a few thousand samples at ~600 samples/s aggregate was minutes of GPU time, while the cost of building a correct S3 reconciliation mechanism was engineering time and complexity. The simple restart was the right call [32].

The assistant executed a clean restart: killing all processes, cleaning /dev/shm, preserving marker files (878 batches already completed), and relaunching all four extractors plus the monitoring WebUI [33]. The marker-based resume mechanism—where .done_* files track which batches have been successfully uploaded to S3—ensured that already-processed samples were skipped. The assistant confirmed: "Markers: 878" and "SHM clean" [34].

The Verification and the Empty Message

The assistant launched a verification command to confirm the pipeline was running and skipping properly, but the user aborted it before output could return [35]. This abort was itself a signal: the user had enough confidence in the pipeline's state—perhaps via the monitoring WebUI—that the verification was unnecessary overhead.

Then came an empty user message—message 7421, containing only empty &lt;conversation_data&gt; tags [36]. This seemingly trivial event revealed deep truths about the conversational dynamics of human-AI collaboration. The empty message was a signal to proceed, a learned shorthand that said "acknowledged, move forward." The assistant correctly interpreted it by producing a comprehensive status summary documenting the entire project state: deployment status, DFlash investigation findings, DDTree research, training data preparation, extraction pipeline status, blocked items, key decisions, next steps, and critical technical context [37].

The Throughput Degradation and Final Status

As the extraction progressed through the dataset, the throughput naturally degraded from the peak of 140–155 samples/s to 36–44 samples/s [38]. This was not a regression but a predictable consequence of sequence length distribution: short sequences (which dominated the early batches) process quickly, while longer sequences take proportionally more time. The assistant correctly diagnosed this and calculated the remaining ETA: approximately 50–60 minutes, with ~413K of 914K samples completed [39].

The final status check confirmed the pipeline was healthy: all four shards running at 72–80 samples/s (for the remaining shorter sequences), with consistent skip rates across shards indicating a well-calibrated filtering mechanism [40]. The extraction that had once been an 8-hour ordeal was now finishing in under an hour.

Lessons and Engineering Philosophy

Several themes emerge from this optimization odyssey:

The bottleneck always shifts. Every optimization revealed a new constraint. Fixing the filesystem overhead exposed the Triton JIT compilation problem. Fixing the JIT compilation exposed the per-sample GPU→CPU copy bottleneck. Fixing the copies exposed the S3 upload bandwidth limitation. The backpressure fix addressed the symptom but not the root cause of the upload pipeline being slower than extraction [41].

The simplest fix is often the best. The GPU-side concatenation fix was conceptually simple—batch before transfer—but required understanding the interaction between PyTorch's tensor operations, CUDA memory management, and PCIe transfer characteristics. The backpressure fix was a five-line change to a Python script. The clean restart was a bash command. None of these fixes required new libraries, new hardware, or architectural redesign [42].

Documentation at the moment of insight is invaluable. The user's request to save progress and plan at the peak of the breakthrough ensured that the hard-won knowledge was captured before the next crisis consumed attention. The PROGRESS.md, PLAN.md, and STATUS.md files became the project's institutional memory [43].

Trust in the recovery mechanism enables velocity. The marker-based resume system, designed for failure from the start, allowed the pipeline to survive crashes, restarts, and node migrations without losing progress. The user's confidence in this mechanism enabled the "just resume" decision that avoided hours of S3 reconciliation work [44].

Conclusion

The hidden state extraction pipeline for the DFlash drafter training project began as a painfully slow 7–11 samples/s per GPU and ended as a robust 140–155 samples/s system capable of processing 914K samples in under an hour. The journey between these two points was not a single breakthrough but a cascade of insights: recognizing that Triton JIT compilation could be prewarmed, that 2,725 individual GPU→CPU copies could be replaced by a single bulk transfer, that the S3 upload pipeline needed backpressure, and that a clean restart with marker-based resume was more reliable than complex state reconciliation.

Each of these insights required understanding the system at a different level of abstraction—from Triton compiler internals to PyTorch tensor semantics to producer-consumer flow control patterns. The assistant's willingness to question each layer of the stack, to abandon failed approaches without attachment, and to document insights at the moment of discovery transformed what could have been a multi-day slog into a focused, efficient optimization sprint. The result was not just a faster pipeline, but a deeper understanding of how to build robust, high-throughput ML infrastructure—knowledge that would prove invaluable in the next phase: actually training the DFlash drafter model.## References

[1] "The One-Time Warmup: Precompiling Triton Kernels to Salvage FLA for GDN Hidden State Extraction" — Article on msg 7394's Triton kernel prewarming strategy.

[2] "The Triton Cache Gambit: Precompiling JIT Kernels to Salvage a Hidden State Extraction Pipeline" — Article on msg 7395's warmup execution and validation.

[3] "The Interrupted Checkpoint: A Window into ML Engineering Debugging" — Article on msg 7396's aborted verification command.

[4] "The Bottleneck Shifts: A User's Diagnostic Eye in the Hidden State Extraction Pipeline" — Article on msg 7397's screenshot revealing the GPU idle pattern.

[5] "The Critical Insight: Batching GPU→CPU Transfers in Hidden State Extraction" — Article on msg 7398's diagnosis of the per-sample copy bottleneck.

[6] "The 2,725-to-1 Optimization: Eliminating GPU→CPU Copy Overhead in Hidden State Extraction" — Article on msg 7399's deployment of the GPU-side concat fix.

[7] "The Breakthrough: How a Single GPU-Side Concatenation Transformed Hidden State Extraction from 11 to 155 Samples Per Second" — Article on msg 7400's monitoring verification.

[8] "The 17× Breakthrough: How Eliminating 2,725 GPU→CPU Copies Per Batch Transformed a Hidden State Extraction Pipeline" — Article on msg 7401's triumphant summary.

[9] "The Pivot to Documentation: Capturing Progress at the Moment of Breakthrough" — Article on msg 7402's user request to save progress.

[10] "The Silent Timeout: A Status Check That Reveals Everything" — Article on msg 7403's SSH timeout.

[11] "The Moment of Truth: Verifying a 17× Throughput Improvement in Hidden State Extraction" — Article on msg 7404's progress data.

[12] "The Checkpoint Moment: Preserving a 17× Breakthrough in Hidden State Extraction" — Article on msg 7405's script archiving.

[13] "The Empty Response: When an AI Assistant Produces Nothing" — Article on msg 7406's empty message.

[14] "The Weight of a Single Word: Analyzing 'continue' in a High-Stakes AI Collaboration" — Article on msg 7407's user prompt.

[15] "The Checkpoint Moment: Documenting Progress After a 17× Breakthrough in Hidden State Extraction" — Article on msg 7408's PROGRESS.md.

[16] "The Moment of Documentation: Writing PLAN.md After a 17× Breakthrough" — Article on msg 7409's PLAN.md.

[17] "The Status Check That Tells a Thousand Stories: Monitoring Hidden State Extraction at Scale" — Article on msg 7410's status check.

[18] "The Checkpoint in the Pipeline: Monitoring Throughput Degradation in DFlash Drafter Training" — Article on msg 7411's STATUS.md.

[19] "The Verification: A Quiet Moment in a 17× Performance Breakthrough" — Article on msg 7412's ls -la verification.

[20] "When Optimization Breaks the Buffer: The Tmpfs Overflow in a High-Throughput Hidden State Extraction Pipeline" — Article on msg 7413's tmpfs crisis.

[21] "The Backpressure Fix: When Optimization Creates New Bottlenecks" — Article on msg 7414's backpressure implementation.

[22] "The Backpressure Problem: When Optimization Creates New Bottlenecks" — Article on msg 7415's deployment of the fix.

[23] "The Question That Reveals Engineering Philosophy: 'Do we resume from S3 state and cleanup?'" — Article on msg 7416's user question.

[24] "The Moment of Reckoning: Reconciling State After a Pipeline Disaster" — Article on msg 7417's S3 reconciliation attempt.

[25] "The Pragmatic Pivot: Why 'It's Cheap Enough to Just Resume' Was the Right Call" — Article on msg 7418's user decision.

[26] "The Art of the Clean Restart: Resilience Through Marker-Based Resume in Distributed ML Pipelines" — Article on msg 7419's clean restart.

[27] "The Verification That Almost Wasn't: A Post-Mortem Check-in on a High-Throughput ML Pipeline" — Article on msg 7420's aborted verification.

[28] "The Empty Message: Conversational Dynamics in AI-Assisted Coding Sessions" — Article on msg 7421's empty message.