From Architecture Tuning to Data-Centric Improvements: Expanding a DFlash Training Pipeline with 193K Diverse Prompts

Introduction

In the lifecycle of a machine learning project, there comes a critical moment when the team must step back from architecture optimization and confront the data. This chunk of the opencode session captures exactly such a pivot — a strategic shift from tuning the DFlash speculative decoding drafter's architecture to radically expanding its training data. What unfolds is a multi-stage operation spanning SGLang deployment on cutting-edge Blackwell GPUs, the construction of a diverse prompt pipeline drawing from six distinct datasets, large-scale batch inference producing over half a billion tokens, and the integration of this new data with an existing 902K-sample corpus — all culminating in a training restart that encountered memory pressure and required careful diagnosis.

This article traces that journey through the session's key phases, examining how the assistant navigated the technical challenges of deploying SGLang on SM120 hardware, the engineering decisions behind prompt selection, the impressive scale of the generation run, and the OOM diagnosis that followed.

The Strategic Pivot: From Architecture to Data

The session began in the midst of architecture-focused work on the DFlash block-diffusion speculative decoding drafter. The team had been tuning parameters like anchors=1024 and block_size=32, running DDTree optimization experiments, and working to maximize the drafter's acceptance rate. However, a critical document changed the trajectory: DATA_EXPANSION.md, which laid out the rationale for expanding the training data mix.

The existing training data was heavily skewed toward coding tasks — approximately 77% coding, 15% general instruction, 6% tool calling, and only 1% math. The DFlash paper had used a more diverse mix (Nemotron + CodeAlpaca), and the LK losses paper cited Infinity-Instruct-0625 as a key data source. The narrow distribution was a plausible bottleneck on drafter performance: if the drafter had only seen coding examples during training, it would struggle to predict the target model's behavior on general knowledge, mathematical reasoning, or tool-use prompts.

The decision was made to halt the DDTree training run on CT200 and repurpose the 8× RTX PRO 6000 Blackwell GPUs for high-throughput batch inference instead. This was a significant commitment — these GPUs represented the project's primary compute resource, and redirecting them from training to inference meant delaying architecture work in favor of data expansion. The gamble was that a more diverse dataset would yield a better drafter than any amount of architecture tuning on the existing narrow data.

Deploying SGLang on SM120 Blackwell: An Environment Debugging Odyssey

The first technical challenge was getting SGLang running on the SM120 (desktop Blackwell) architecture. The RTX PRO 6000 Blackwell GPUs use compute capability SM 12.0, which at the time of this work was still relatively new and not fully supported by all ML frameworks. The assistant embarked on an extensive environment debugging process that revealed the gritty reality of deploying cutting-edge hardware.

The initial attempt involved installing sglang==0.5.12, but this immediately ran into CUDA version mismatches. The system had CUDA 13.2 installed, but the pip-installed CUDA headers expected a different version layout. The assistant had to manually match the CUDA 13.2 nvcc compiler with the pip-installed CUDA headers — a process that required creating symlinks for libcudart and libcuda stubs to satisfy the build system's expectations.

The next roadblock came from CCCL (CUDA C++ Core Libraries) headers. The flashinfer library, which SGLang uses for attention computations, bundles its own version of libcudacxx. However, the include paths were not resolving correctly, leading to nv/target include errors during compilation. The assistant resolved this by overlaying CCCL headers from flashinfer's bundled libcudacxx onto the system include path, effectively creating a hybrid header environment that satisfied all dependencies.

The attention backend itself proved to be a critical configuration choice. SGLang supports multiple attention backends, including Flash Attention 3 (FA3) and Flash Attention 4 (FA4). However, these newer flash-attention versions did not support SM120 at the time. The assistant had to switch to --attention-backend flashinfer, which uses the flashinfer library's attention implementation instead. This was a necessary compromise — flashinfer's attention kernels were not as optimized as FA3/FA4 for some operations, but they were the only option that worked on the Blackwell architecture.

Once SGLang was operational, the assistant turned to performance tuning. The initial configuration used the extra_buffer mamba strategy, which allocated additional GPU memory as a buffer for request handling. However, this strategy limited the number of concurrent requests that could be served. By swapping to no_buffer mode, which eliminates the extra buffer and relies on more efficient request scheduling, the assistant doubled the maximum concurrent requests from 37 to 72 per GPU. This change, combined with other optimizations, yielded approximately 1,180 tokens per second per GPU, or about 9,400 tokens per second aggregate across the 8 GPUs.

Building the Diverse Prompt Pipeline

With SGLang deployed and tuned, the assistant turned to the prompt preparation pipeline. The goal was to extract approximately 193,000 diverse prompts from a carefully curated set of data sources, each chosen to address a specific gap in the existing training distribution.

The prompt sources were:

Infinity-Instruct-0625 (~99K prompts): This was the largest single source, drawn from the Infinity-Instruct dataset that the LK losses paper had used. It contains a broad range of general instruction-following examples, covering topics from creative writing to analytical reasoning. This dataset was the primary vehicle for expanding beyond the coding-heavy distribution.
WebInstructSub (~40K prompts): A subset of the WebInstruct dataset, which contains web-derived instruction-response pairs. This source added diversity in domain coverage, including topics like science, history, and everyday knowledge that were underrepresented in the coding-focused original data.
CodeFeedback (~29K prompts): While the original dataset was already coding-heavy, the CodeFeedback dataset provided a different flavor of code examples — specifically, examples where the model receives feedback on its code output and must iteratively improve. This added a conversational coding dimension that was missing from the static code examples in the original data.
MetaMathQA (~24K prompts): This was the primary vehicle for addressing the math gap. The original dataset was only 1% math; MetaMathQA added nearly 24K math reasoning prompts covering arithmetic, algebra, geometry, and more complex mathematical problem-solving. This was critical for improving the drafter's ability to predict the target model's mathematical reasoning traces.
Hermes Function Calling v1 (~1.2K prompts): A smaller but high-value source focused on tool-use scenarios. These prompts included proper tool XML specifications, teaching the model to call functions with correct syntax and parameter formatting. While only 1.2K prompts, they represented a qualitative addition — the original dataset had tool-calling examples, but the Hermes dataset provided more structured and diverse tool-use patterns.
Agent Training (~553 prompts): The smallest source, but strategically important. These prompts focused on multi-step agentic tasks where the model must plan, execute actions, and reason about outcomes. Even a few hundred such examples could help the drafter learn the patterns of agentic reasoning. The prompt preparation pipeline was rewritten from scratch to handle the heterogeneous formats of these datasets. Each source had its own schema — some used the ShareGPT conversation format, others used OpenAI-style message arrays, and some used custom JSON structures. The pipeline normalized all prompts into a consistent format suitable for SGLang batch inference, ensuring that the chat template and tokenization would be applied uniformly.

The Generation Run: Scale and Reliability

With the prompts prepared and SGLang deployed, the assistant launched the batch generation run. The scale was impressive: 193,010 prompts were sent to the Qwen3.6-27B model running on the 8 Blackwell GPUs, with thinking mode enabled to capture the model's internal reasoning traces.

The generation completed 192,995 out of 193,010 prompts — a success rate of 99.992%. Only 15 prompts failed, representing a failure rate of 0.008%. This extraordinary reliability is a testament to both the stability of the SGLang deployment (after the extensive debugging) and the robustness of the prompt preparation pipeline. In large-scale batch inference, failures typically arise from edge cases in prompt formatting, unexpected token sequences that trigger model errors, or transient hardware issues. The near-perfect completion rate suggests that the pipeline handled all these edge cases gracefully.

The total output was 523 million tokens, with an average of approximately 2,712 tokens per completion. This average is notably high — typical instruction-following datasets average 500-1,000 tokens per response. The higher average reflects the use of thinking mode, which causes the model to generate extended reasoning traces before producing its final answer. These thinking traces are valuable for drafter training because they expose the model's internal reasoning process, which the drafter must learn to predict.

Tokenization and Dataset Integration

Once the completions were generated, the assistant tokenized them using the Qwen3.6 chat template, which includes special tokens for marking thinking boundaries (<|thinking|> and <|think_end|>). The tokenization pipeline also computed loss masks — a critical step that determines which tokens contribute to the training loss. In the DFlash training setup, only assistant tokens (including thinking tokens) are trained on; system prompts and user messages are masked out.

The new completions were then merged with the existing 902K-sample dataset. The combined dataset totaled 1,095,082 samples with 2.411 billion tokens. This represented a 21% increase in sample count and a substantial expansion in domain diversity. The tokenized data was stored in Arrow format across multiple shards, following the same schema as the original dataset: input_ids (the tokenized sequence), loss_mask (indicating which tokens are trainable), and seq_len (the sequence length).

Training Resumption and the OOM Challenge

With the expanded dataset ready, the assistant resumed DFlash training from the step 690 checkpoint. The training configuration preserved the critical architectural parameters: anchors=1024 (the number of anchor positions used for block-diffusion prediction) and block_size=32 (the number of tokens predicted per block). These parameters were considered essential for maintaining the drafter's performance and should not be compromised.

However, during the training ramp-up phase, GPU 6 suffered an out-of-memory (OOM) error. The assistant diagnosed the issue as stemming from the increased memory consumption of the expanded dataset combined with the existing training configuration. The new data had longer average sequences (due to the thinking traces), which increased the memory footprint of each training batch.

The user directed the assistant to restart training from scratch (not from the checkpoint) with all three drafter GPUs engaged. The key constraint was to preserve the anchors=1024 and block_size=32 configuration while tuning non-harmful parameters like batch size to resolve the memory pressure. This was a delicate balancing act: reducing the batch size too aggressively would hurt training throughput and stability, but keeping it too high would trigger OOM errors.

The assistant's response to this challenge — and the subsequent attempts to find a working configuration — would set the stage for the next phase of the session, where the tension between training signal preservation and memory management would drive a series of iterative adjustments and ultimately a torch version rollback.

Reflections on the Data Expansion Strategy

The work in this chunk represents a mature understanding of how training data composition affects model performance. The team recognized that their coding-heavy dataset was a likely bottleneck and took decisive action to diversify it. The selection of data sources was strategic rather than opportunistic: each source addressed a specific gap in the distribution, and the quantities were calibrated to meaningfully shift the overall mix without overwhelming the existing data.

The decision to use thinking-mode completions is particularly noteworthy. Many speculative decoding pipelines train on final answers only, ignoring the target model's reasoning traces. By including thinking tokens in the training data, the DFlash drafter can learn to predict not just what the target model says, but how it thinks — potentially improving its ability to anticipate the target model's next tokens in contexts where reasoning is required.

The SGLang deployment effort also deserves recognition. Getting a relatively new inference engine running on cutting-edge hardware required deep debugging of CUDA toolchains, header paths, and attention backend configurations. The assistant's systematic approach — identifying each error, understanding its root cause, and applying targeted fixes — demonstrates the kind of low-level systems engineering that is often required when working at the frontier of ML hardware support.

Conclusion

This chunk of the opencode session captures a pivotal transition in the DFlash training project: the shift from architecture optimization to data-centric improvement. The assistant successfully deployed SGLang on SM120 Blackwell GPUs after extensive environment debugging, constructed a diverse prompt pipeline drawing from six datasets, executed a near-perfect batch generation run producing 523M tokens, and integrated the new data into a combined 1.095M-sample corpus. The training restart encountered OOM challenges that would drive further iteration, but the foundation was laid for a significantly more diverse and capable training dataset.

The work demonstrates that in modern ML engineering, data pipeline construction is as challenging — and as critical — as model architecture design. The ability to deploy inference servers on new hardware, prepare diverse prompts at scale, execute reliable batch generation, and integrate the results into an existing training pipeline is a complex engineering feat that requires deep knowledge across multiple domains. This chunk shows that feat in action.