The Compact Draft Cache: Enforcing Sliding Window Attention in SGLang's DFlash Speculative Decoding

Introduction

In the high-stakes world of large language model inference, every optimization matters. When deploying a 1-trillion-parameter MoE model like Kimi K2.6 with speculative decoding, the difference between a working system and an optimal one often lies in the details of how attention windows are managed, how memory is allocated, and how the drafter model's architectural constraints are faithfully respected at runtime. Message [msg 11695] captures a pivotal moment in this optimization journey: the assistant discovers a critical mismatch between the drafter model's configuration and its runtime behavior, and traces through SGLang's source code to determine how to properly enforce sliding window attention.

This message sits at the intersection of model architecture understanding, inference engine internals, and practical deployment optimization. It is not a flashy breakthrough—no throughput numbers jump, no bugs are fixed—but it represents the kind of careful, systematic investigation that separates a merely working deployment from a truly optimized one. The assistant is asking a seemingly simple question: "Does the drafter actually respect its 2048-token sliding window during inference?" The answer turns out to be far from trivial, and the path to finding it reveals deep insights about how SGLang's speculative decoding pipeline is structured.

The Context: A Chain of Discoveries

To understand why message [msg 11695] was written, we must trace back through the preceding conversation. The user had just asked the assistant to "sweep budgets/topk" and "make sure we have sliding window attention on the drafter, up to 2k context" ([msg 11690]). This request came after an extensive period of deploying and benchmarking the K2.6 model with DFlash speculative decoding on 8× RTX PRO 6000 Blackwell GPUs. The assistant had already fixed three critical bugs—a CUDA graph sizing mismatch, a custom-mask corruption that silently broke tree verification, and temperature sampling support—and had achieved promising results: DDTree with budget=8 and CUDA graphs reached 150 tok/s at single concurrency, a 1.53× speedup over the autoregressive baseline.

But the user's request about sliding window attention revealed a potential blind spot. The drafter model—a smaller "draft" network that predicts tokens for the main model to verify—had been trained with a 2048-token sliding window. This means each attention layer in the drafter can only attend to the last 2048 tokens, not the full context. If the inference engine was ignoring this constraint and allowing full-context attention, the drafter would be operating out of distribution, potentially degrading both quality and performance.

The assistant's first investigation ([msg 11691]) confirmed the drafter's configuration: sliding_window=2048, with 5 out of 6 layers using sliding_attention and only one using full_attention. But the runtime logs showed a troubling discrepancy: draft_window_size=None, compact_cache=False. The draft KV cache was storing the full context rather than clamping to 2048 tokens. This meant the drafter was attending over the entire input, not just the sliding window it was trained for.

A series of follow-up investigations ([msg 11692], [msg 11693], [msg 11694]) traced through the SGLang source code to understand how sliding window attention was implemented in the DFlash draft model. The assistant checked the HuggingFace reference implementation (dflash.py), then SGLang's own dflash.py model file, and finally the dflash_worker.py that orchestrates the draft model's runtime behavior. Each step revealed more about the architecture but also raised new questions.

The Message Itself: Reasoning and Investigation

Message [msg 11695] is the culmination of this investigation chain. The assistant opens with a clear statement of the problem:

I'm checking whether SGLang's DFlash draft model exposes the get_attention_sliding_window_size() method or if model_config.sliding_window_size gets populated from the drafter's configuration—if neither is true, the draft model would run with full attention across the entire context instead of using a sliding window.

This reasoning reveals two possible paths by which SGLang could enforce sliding window attention. The first is a method called get_attention_sliding_window_size() that a model class can implement to declare its sliding window size. The second is the model_config.sliding_window_size field, which might be populated automatically from the model's configuration file. If neither path is active, the draft model defaults to full attention—which is what the runtime logs were showing.

The assistant then articulates the practical solution:

The practical solution is to use --speculative-dflash-draft-window-size 2048, which constrains the draft KV cache to 2048 tokens, naturally implementing the sliding window the drafter was trained with while also reducing memory overhead at longer contexts.

This is a crucial insight. Rather than trying to modify the attention computation itself to enforce the sliding window—which would require changes to SGLang's triton kernels—the assistant recognizes that simply clamping the KV cache to 2048 tokens achieves the same effect. If the draft model only has access to the last 2048 tokens of history, it cannot attend to older tokens, regardless of what the attention implementation does. This is an elegant, low-risk solution that leverages existing infrastructure (the compact draft cache feature) rather than requiring invasive code changes.

The assistant then runs a bash command to verify the plumbing. The command probes three things:

Whether the DFlashDraftModel class or its attention layers expose get_attention_sliding_window_size
The compact cache clamp logic in dflash_worker.py around line 470-480
The _compute_compact_draft_seq_lens method that implements the actual cache clamping The output confirms the draft model class hierarchy (DFlashAttention, DFlashMLP, DFlashDecoderLayer, DFlashDraftModel) and shows the compact cache method signature. Notably, the grep for get_attention_sliding_window_size returns no results—confirming that the draft model does not implement this method, which is why draft_window_size=None in the logs.

Decisions Made and Assumptions Challenged

The most important decision in this message is the choice to use --speculative-dflash-draft-window-size as the mechanism for enforcing sliding window attention. This decision rests on several assumptions:

The compact draft cache is functionally equivalent to architectural sliding window attention. This is a reasonable assumption: if the model can only see 2048 tokens of history, it cannot attend beyond that window, regardless of the attention kernel's capabilities. However, there is a subtle difference: architectural sliding window attention applies per-layer (some layers might be full attention), while the compact cache is a global constraint. The drafter config shows 5 sliding layers and 1 full layer—with the compact cache, even the full layer would be limited to 2048 tokens. This is likely acceptable since the full layer was presumably trained to operate within the 2048-token context anyway.
The compact cache feature is compatible with the DDTree speculative decoding path. The assistant has already fixed several DDTree-specific bugs (CUDA graph sizing, mask corruption), and adding the compact cache introduces another variable. The assistant implicitly assumes this won't break anything, or that any issues will be caught in subsequent benchmarking.
The --speculative-dflash-draft-window-size flag actually works as documented. The assistant hasn't yet tested this—the message ends with the verification still in progress. The assumption is that SGLang's compact cache implementation correctly clamps the KV cache and that the draft model's forward pass handles the truncated cache gracefully. The message also challenges an implicit assumption from earlier in the conversation: that because the model config specified sliding_window=2048, the runtime would automatically respect it. The discovery that draft_window_size=None disproves this—SGLang's DFlash draft model does not automatically read the sliding window from the model config, at least not through the paths the assistant checked.

Input Knowledge Required

To fully understand this message, one needs:

SGLang's speculative decoding architecture: How DFlash works, the relationship between the draft model and the target model, and the role of the draft KV cache. The assistant has been working with this codebase extensively, so terms like "draft runner," "compact cache," and "dflash_worker" are familiar.
The concept of sliding window attention: How it constrains each token to only attend to a fixed-size window of previous tokens, and why this is important for both computational efficiency and distributional correctness. The drafter was trained with a 2048-token window, so using full attention would be out-of-distribution.
The specific model architecture: Kimi K2.6's DFlash drafter has 6 layers, 5 of which use sliding attention and 1 uses full attention. The sliding_window=2048 config parameter controls this.
Python and PyTorch model conventions: The get_attention_sliding_window_size() method is a convention in HuggingFace transformers and SGLang for models to declare their sliding window size. Understanding this convention is essential to interpreting the investigation.
Linux command-line tools: The bash command uses grep, sed, head, and piping to probe source code on a remote machine. The ssh -o ConnectTimeout=5 pattern is used throughout the session for remote execution.

Output Knowledge Created

This message produces several valuable pieces of knowledge:

The DFlash draft model does not implement get_attention_sliding_window_size(). This is confirmed by the grep returning no results. This explains why draft_window_size=None in the runtime logs—SGLang has no way to discover the drafter's sliding window size automatically.
The compact cache mechanism exists and is the correct lever. The _compute_compact_draft_seq_lens method in dflash_worker.py implements the actual cache clamping. The --speculative-dflash-draft-window-size flag controls whether this is active and what window size to use.
The draft model class hierarchy is confirmed. The output shows DFlashAttention, DFlashMLP, DFlashDecoderLayer, and DFlashDraftModel as the key classes. This is useful for understanding where to make changes if deeper modifications are needed.
A clear action plan is established. The assistant will set --speculative-dflash-draft-window-size 2048 and then proceed to benchmark across budget, top-k, context length, and concurrency combinations. This directly addresses the user's request.

The Thinking Process: A Model of Systematic Debugging

The reasoning visible in this message exemplifies a disciplined approach to debugging inference engine behavior. The assistant:

States the hypothesis clearly: "I'm checking whether SGLang's DFlash draft model exposes the get_attention_sliding_window_size() method or if model_config.sliding_window_size gets populated from the drafter's configuration."
Identifies the consequence of the hypothesis being false: "if neither is true, the draft model would run with full attention across the entire context instead of using a sliding window."
Proposes a practical solution before fully confirming the hypothesis: "The practical solution is to use --speculative-dflash-draft-window-size 2048." This is not premature—it's efficient. The assistant recognizes that even if the hypothesis is wrong (i.e., if one of the automatic paths does work), the explicit flag would override it and provide the desired behavior. The flag is a superset of the automatic paths.
Verifies the solution's plumbing: The bash command checks that the compact cache mechanism exists and understands how it works. The assistant doesn't just assume the flag works—it traces through the code to confirm.
Plans the next steps: "I should verify this works correctly and then benchmark across different budget, top-k, context length, and concurrency combinations." This pattern—hypothesize, propose solution, verify plumbing, plan verification—is a hallmark of effective engineering. The assistant avoids both the trap of endless investigation (spending hours tracing code paths without proposing a fix) and the trap of blind action (setting a flag without understanding what it does).

Broader Significance

While this message might seem like a minor technical investigation, it has significant implications for the overall deployment. The sliding window attention constraint is not optional—it is baked into the drafter's training. Ignoring it would mean the drafter operates out-of-distribution on long contexts, potentially producing lower-quality drafts and reducing the acceptance rate. By enforcing the 2048-token window, the assistant ensures that:

Distributional correctness: The drafter operates in the regime it was trained for, producing drafts that the target model is more likely to accept.
Memory efficiency: The draft KV cache is capped at 2048 tokens regardless of context length. For a 32K-token context, this saves ~15/16 of the draft KV cache memory—a substantial reduction that frees GPU memory for larger batches or longer contexts.
Computational efficiency: The attention computation in the drafter is naturally bounded by the window size, reducing FLOPs per draft token. The message also demonstrates the value of understanding the full stack—from model architecture (the sliding window config) through the inference engine (SGLang's draft runner) to the runtime behavior (the compact cache). This holistic view is essential for effective optimization.

Conclusion

Message [msg 11695] is a masterclass in targeted investigation. In a single round, the assistant identifies a critical gap between model configuration and runtime behavior, proposes a practical solution, verifies the solution's plumbing in the source code, and establishes a clear plan for validation and benchmarking. The reasoning is precise, the assumptions are explicit, and the action is grounded in code-level understanding.

This message may not contain dramatic speedups or bug fixes, but it represents the kind of careful, systematic work that makes those speedups possible. The compact draft cache, once enabled, will ensure the drafter operates correctly and efficiently—a foundation upon which all subsequent optimization can build.