The 8 vs 33 Mismatch: Diagnosing a CUDA Graph Dimension Bug in DDTree Speculative Decoding

Introduction

In the high-stakes world of deploying large language models with speculative decoding, every tensor dimension must align perfectly—or the entire service crashes. Message [msg 11605] captures one such moment of crisis and insight: the assistant, having just launched a Kimi K2.6 model with DDTree (Draft-Driven Tree) speculative decoding on 8× RTX PRO 6000 Blackwell GPUs, watches the service fail with a cryptic tensor size mismatch. The error message—"The size of tensor a (8) must match the size of tensor b (33) at non-singleton dimension 0"—is a puzzle that the assistant must solve before any further progress can be made.

This message is a turning point in a longer debugging arc. The assistant had previously been chasing a "NoneType buffer" crash in the CUDA graph path ([msg 11592]), had added diagnostic instrumentation, and had now triggered a different failure mode. The shift from one error to another is itself informative: it means the earlier diagnostic patches had changed the execution path, exposing a deeper architectural issue.

The Context: DDTree on Blackwell GPUs

To understand why this error matters, we need to understand what DDTree is and why CUDA graphs are involved. DDTree (Draft-Driven Tree) is a speculative decoding algorithm that improves upon linear DFlash by constructing a tree of candidate continuations rather than a single chain. At each step, the draft model produces block_size=8 hidden states; from these, the system selects up to budget=32 candidate tokens organized into a tree structure. The key insight is that all budget+1=33 nodes (32 candidates plus the root) can be verified in a single forward pass through the target model, using a custom tree-attention mask.

CUDA graphs are a performance optimization that captures a GPU kernel launch sequence and replays it with minimal CPU overhead. However, CUDA graphs have a critical limitation: all tensor shapes must be fixed at capture time. If the graph was captured expecting 8 tokens per batch but the actual forward pass sends 33 tokens, the replay will crash with a dimension mismatch—exactly the error seen here.

The assistant had deployed the service with --speculative-dflash-block-size 8 and --speculative-ddtree-budget 32 ([msg 11592]), then waited 390 seconds for the 590 GB model to load and the CUDA graphs to be captured ([msg 11604]). The first user request triggered the crash.

The Message: Reasoning and Action

Message [msg 11605] is deceptively brief. It contains two parts: a reasoning block and a bash command. But the reasoning reveals a sophisticated diagnostic process compressed into a single paragraph:

The crash reveals a tensor size mismatch—the draft block size is 8 but the DDTree budget is 33 (32+1), so I need to trace through where these dimensions are getting out of sync.

The assistant immediately recognizes the numbers: 8 is block_size, 33 is budget + 1. This is not obvious to someone unfamiliar with the DDTree implementation. The assistant knows that:

The Bash Command: Tracing the Stack

The assistant then executes a bash command to extract the full traceback from the systemd journal:

ssh -o ConnectTimeout=5 root@10.1.2.200 "journalctl -u sglang-k26-ddtree.service --no-pager -n 60 | grep -E 'File|Error|line|RuntimeError|size of tensor|DFLASH_DIAG' | tail -30" 2>&1

The command filters for specific keywords: File and line to capture Python traceback entries, RuntimeError and size of tensor for the error itself, and DFLASH_DIAG for any custom diagnostic output the assistant had previously injected. The tail -30 limits output to the most relevant lines.

The result shows the crash originates in model_runner.py line 3223 (the _forward_raw method) and passes through cuda_graph_runner.py line 1306 (the replay method). This confirms that the error occurs during CUDA graph replay—the graph was captured with one set of dimensions but is being replayed with another.

The Thinking Process: What the Reasoning Reveals

The reasoning in [msg 11605] is a model of efficient debugging. The assistant:

  1. Recognizes the error pattern: The numbers 8 and 33 map directly to algorithm parameters (block_size and budget + 1). This is pattern-matching against deep architectural knowledge.
  2. Formulates a hypothesis: The CUDA graph was captured with num_tokens_per_bs=8 (the draft block size), but the DDTree verify forward sends budget+1=33 tokens. The graph replay fails because the tensor buffers are sized for 8 tokens but receive 33.
  3. Plans the next step: "Let me get the full traceback" — the assistant needs to confirm the hypothesis by tracing through the exact call stack, identifying which specific buffer allocation is mismatched.
  4. Signals the shift in debugging focus: "Different error now" acknowledges that the diagnostic work has evolved. The NoneType buffer issue may have been a symptom of the same root cause, or it may have been a separate bug that was fixed by earlier patches. This thinking process is notable for what it doesn't include: there is no speculation about unrelated causes, no detour into checking hardware or configuration files, no panic about the service being down. The assistant has developed a mental model of how DDTree and CUDA graphs interact, and the error message fits cleanly into that model.

Assumptions and Potential Mistakes

The assistant makes several assumptions in this message:

  1. That the CUDA graph capture uses speculative_num_draft_tokens for both draft and verify workers: This turns out to be correct—the subsequent investigation in [msg 11606] through [msg 11610] confirms that num_tokens_per_bs is uniformly set to speculative_num_draft_tokens=8 regardless of whether the worker is doing draft or verify work.
  2. That the DDTree verify forward always processes exactly budget+1 tokens: This is also correct—the DDTree builder pads the tree to exactly budget+1 nodes, making it a fixed-size batch suitable for CUDA graph capture.
  3. That the fix involves making the target verify graph capture with budget+1 instead of block_size: This is the correct architectural fix, but the assistant may be underestimating the complexity. The cuda_graph_runner.py code needs to distinguish between draft and target workers, and the DDTree verify worker needs a different num_tokens_per_bs than the draft worker. The subsequent messages show the assistant tracing through exactly this logic. One potential blind spot: the assistant assumes the CUDA graph capture happens after the DDTree configuration is fully loaded. If the graph capture occurs during model loading (before the DDTree parameters are fully initialized), the budget value might not be available yet. The assistant's investigation in later messages confirms this isn't the case—the graph capture happens during the warmup phase, after all parameters are set.

Input Knowledge Required

To understand this message fully, a reader needs:

  1. DDTree algorithm mechanics: Understanding that DDTree builds a tree of budget candidate tokens from block_size draft hidden states, and verifies all budget+1 nodes in a single forward pass.
  2. CUDA graph constraints: Knowing that CUDA graphs require fixed tensor shapes at capture time, and that replaying with different shapes causes runtime errors.
  3. SGLang speculative decoding architecture: Awareness that SGLang has separate draft and target workers, and that the CUDA graph runner is shared between them with a single num_tokens_per_bs parameter.
  4. The debugging history: Understanding that the assistant had been chasing a NoneType buffer error, added diagnostic patches, and is now seeing a different failure mode.
  5. Systemd journal inspection: The use of journalctl to extract Python tracebacks from a crashed systemd service.

Output Knowledge Created

This message creates several pieces of actionable knowledge:

  1. Root cause hypothesis: The CUDA graph for the target verify forward is captured with the wrong token count (block_size instead of budget+1).
  2. Debugging direction: The assistant needs to trace how num_tokens_per_bs is set in cuda_graph_runner.py, specifically whether it distinguishes between draft and verify workers for DDTree.
  3. Confirmation of the error path: The traceback confirms the crash occurs in cuda_graph_runner.replay(), validating the hypothesis that this is a CUDA graph dimension issue rather than a model loading or memory allocation problem.
  4. A shift in diagnostic strategy: The assistant can now stop investigating the NoneType buffer issue and focus on the dimension mismatch, which is a more fundamental architectural problem.

The Broader Significance

This message represents a classic debugging inflection point. The assistant had been working on one theory (a missing buffer in the CUDA graph capture) and had added diagnostic code to identify which buffer was None. But the act of adding that diagnostic code changed the execution path, causing a different bug to surface first. This is the Heisenbug phenomenon in reverse: instead of a bug disappearing when you try to observe it, a different bug appears because the observation changed the system state.

The 8 vs 33 mismatch is also a cautionary tale about the tension between optimization and correctness. CUDA graphs are a powerful performance optimization—they can 3–4× throughput by eliminating Python overhead—but they require fixed tensor shapes. DDTree's variable-sized tree (which could theoretically have any number of nodes up to budget) is fundamentally at odds with CUDA graph's fixed-size requirement. The assistant's insight that DDTree always pads to exactly budget+1 nodes is the key that makes the optimization viable: the tree size is effectively fixed, even though it's logically variable.

The subsequent investigation in [msg 11606] through [msg 11610] traces through the code to confirm the hypothesis, ultimately finding that speculative_num_draft_tokens is set to block_size (8) and used uniformly for both draft and verify workers. The fix requires adding a conditional: when DDTree is active and the worker is the target verify worker, use budget + 1 instead of block_size. This is a small code change with large performance implications—getting it wrong means either crashes (if the graph is too small) or wasted GPU memory (if the graph is too large).

Conclusion

Message [msg 11605] is a masterclass in efficient debugging. In two sentences of reasoning and one bash command, the assistant identifies a tensor dimension mismatch, maps it to the correct algorithm parameters, formulates a root cause hypothesis, and plans the next investigation step. The message captures the moment when a confusing error message becomes a clear architectural insight—the moment when "8 vs 33" transforms from a cryptic runtime error into a well-understood design flaw in the CUDA graph capture logic.

This message also illustrates the iterative nature of debugging complex ML systems. The assistant didn't find the bug on the first try; it found a different bug, which led to a deeper understanding, which eventually led to the real fix. The 8 vs 33 mismatch was not the final bug—the assistant would go on to discover that even with the correct token count, the B300 SXM6 GPUs had their own CUDA graph instability issues with budgets larger than 8 ([chunk 64.2]). But this message was where the assistant first understood the fundamental tension between DDTree's tree-structured verification and CUDA graph's fixed-size constraints—an understanding that would prove essential for the entire deployment effort.