Diagnosing a Tensor Size Mismatch: DDTree CUDA Graph Verification Bug

Introduction

In the high-stakes world of deploying large language models with speculative decoding, every optimization comes with hidden assumptions. This article examines a single message from an opencode coding session where an AI assistant diagnosed a critical crash during the deployment of the Kimi K2.6 model with DDTree (Draft-Tree) speculative decoding on an 8× NVIDIA RTX PRO 6000 Blackwell system. The crash manifested as a tensor size mismatch—RuntimeError: The size of tensor a (8) must match the size of tensor b (33) at non-singleton dimension 0—and the assistant's reasoning in this message reveals a deep understanding of the interaction between CUDA graph optimization and the DDTree verification algorithm.

The subject message, <msg id=11606>, captures the moment when the assistant traces the crash through the SGLang inference stack, identifies the root cause, and plans the fix. It is a masterclass in systematic debugging of complex ML inference systems.

The Context: Deploying DDTree on Blackwell

The assistant had been working for hours to deploy Kimi K2.6—a 1-trillion-parameter Mixture-of-Experts model quantized to INT4—with speculative decoding on 8× RTX PRO 6000 Blackwell GPUs connected via PCIe. The speculative decoding approach had evolved through several iterations: first linear DFlash (a single-path draft-verify loop), then DDTree (a tree-structured draft that evaluates multiple candidate paths in a single verify forward pass).

The DDTree configuration used --speculative-ddtree-budget 32 and --speculative-ddtree-topk-cap 8, meaning the tree builder would expand up to 32 candidate nodes from the top-8 candidates at each depth position. The draft model had block_size=8, meaning each draft forward produced 8 hidden states. The key architectural insight—and the source of the bug—is that DDTree's verify forward processes budget + 1 = 33 tokens in a single batch (the root plus 32 tree nodes), not the block_size = 8 tokens that linear DFlash would process.

The service had started successfully, loading the 590 GB model across 8 GPUs, but crashed on the very first inference request. The assistant needed to diagnose why.

The Message: Tracing Through the Stack

The subject message begins with the assistant's reasoning, which reveals a systematic trace through the crash stack:

I'm tracing through the stack to pinpoint where the issue originates—it starts at the forward batch generation call in dflash_worker.py and cascades through the model runner's forward pass, eventually hitting the CUDA graph replay logic where the actual problem surfaces.

This is a classic debugging approach: follow the stack trace from the error back to its origin. The assistant identifies three layers in the call chain:

1. dflash_worker.py — the speculative decoding worker that orchestrates draft generation and verification
2. model_runner.py — the generic model forward pass infrastructure
3. cuda_graph_runner.py — the CUDA graph optimization that pre-records GPU operations for faster execution The crash surfaces in the CUDA graph replay, but the root cause lies in how the graph was captured.

The Root Cause: A Fundamental Dimension Mismatch

The assistant's reasoning crystallizes the problem with remarkable clarity:

The error reveals a fundamental mismatch: the target model's CUDA graph was pre-captured assuming 8 tokens per batch (the speculative block size), but DDTree's verify forward is actually passing 33 tokens per batch (budget + 1). This causes the tensor size mismatch when the graph tries to replay with buffers sized for the wrong token count.

This is the crux of the issue. CUDA graphs are a performance optimization that records a sequence of GPU kernel launches and replays them with minimal CPU overhead. However, the graph capture is done with specific tensor shapes—the buffers are allocated once and reused. If the graph was captured with buffers sized for 8 tokens (the linear DFlash block size), but the actual inference pass sends 33 tokens (the DDTree budget + 1), the tensor operations will fail because the buffers are too small.

The assistant correctly identifies that this is specifically the target verify CUDA graph (referencing line 1521 in the codebase), not the draft graph. The target verify graph is responsible for running the large target model (K2.6) on the candidate tokens to produce verification logits. When DDTree is active, this graph needs to handle budget + 1 tokens simultaneously, not just block_size tokens.## The Investigation: Tracing the Configuration Parameter

After identifying the root cause, the assistant immediately moves to trace how the critical parameter—num_tokens_per_bs—is set in the CUDA graph runner. The message includes a bash command that greps through cuda_graph_runner.py to find all references to this parameter and related variables:

[bash] ssh -o ConnectTimeout=5 root@10.1.2.200 "grep -n 'num_tokens_per_bs\|speculative_num_draft_tokens\|draft_token_num\|ddtree_budget\|is_ddtree' /root/venv_sglang211/lib/python3.12/site-packages/sglang/srt/model_executor/cuda_graph_runner.py | head -30" 2>&1
178:        num_tokens_per_bs: int,
199:                (max_bs * seq_len_fill_value + max_num_token) * num_tokens_per_bs,
290:        num_tokens_per_bs: int,
350:            self.global_num_tokens_gpu.fill_(bs * num_tokens_per_bs)
351:            self.global_num_tokens_for_logprob_gpu.fill_(bs * num_tokens_per_bs)
355:                num_tokens_per_dp = bs * num_tokens_per_bs
502:def get_batch_sizes_to_capture(model_runner: ModelRunner, num_tokens_per_bs=1):
510:        num_tokens_per_bs = 1  #...

This grep output is the evidence the assistant needs to confirm the hypothesis. The key observations:

1. Line 178: num_tokens_per_bs is a constructor parameter, meaning it's set once when the CUDA graph runner is initialized.
2. Lines 350-355: This parameter directly controls buffer sizes: self.global_num_tokens_gpu.fill_(bs * num_tokens_per_bs). If num_tokens_per_bs=8, the buffer holds 8 * bs tokens. If the actual forward pass sends 33 * bs tokens, the buffer is undersized by a factor of ~4.
3. Line 502: get_batch_sizes_to_capture has a default of num_tokens_per_bs=1, suggesting that for non-speculative decoding, each batch slot processes one token at a time.
4. Line 510: There's a comment or override that forces num_tokens_per_bs = 1, which may be a fallback for non-speculative modes. The assistant is looking for where num_tokens_per_bs gets set to 8 (the block size) versus where it should be set to 33 (budget + 1) when DDTree is active. The grep also searches for is_ddtree and ddtree_budget to see if the code already has awareness of DDTree as a distinct mode—but the absence of these in the output (only the searched terms appear) suggests the CUDA graph runner does not distinguish between linear DFlash and DDTree.

Assumptions Made by the Assistant

The assistant makes several assumptions in this message that are worth examining:

1. The crash is reproducible and deterministic: The assistant assumes that restarting the service and sending another request will produce the same crash. This is a reasonable assumption for a CUDA graph capture issue—the graph is captured once at startup, and every replay will hit the same buffer size problem.
2. The bug is in the CUDA graph capture, not in the DDTree verify logic: The assistant immediately focuses on the graph capture dimension rather than suspecting a logic error in the verify forward itself. This is informed by the specific error message—"size of tensor a (8) must match size of tensor b (33)"—which clearly indicates a dimension mismatch rather than a NaN, assertion failure, or illegal memory access.
3. The fix is to change num_tokens_per_bs based on the speculative algorithm: The assistant's planned solution—"ensure the target verify CUDA graph captures with the correct num_tokens_per_bs value when DDTree is active"—assumes that the graph runner already has the machinery to handle variable num_tokens_per_bs and just needs the right value passed. This is a reasonable assumption given that the parameter already exists in the constructor.
4. The draft graph is not affected: The assistant focuses specifically on the target verify graph (line 1521) and does not investigate whether the draft model's CUDA graph also needs adjustment. This is likely correct because the draft model always produces block_size=8 tokens regardless of the tree structure—the tree expansion happens on top of those 8 hidden states.

Input Knowledge Required

To fully understand this message, a reader needs knowledge of several interconnected concepts:

1. Speculative decoding: The technique of using a small "draft" model to propose candidate tokens that a large "target" model verifies in parallel. The assistant had deployed both linear DFlash (single-path) and DDTree (tree-structured) variants.
2. CUDA graphs: A CUDA API feature that captures a sequence of GPU kernel launches into a graph object that can be replayed with near-zero CPU overhead. The graphs are captured with fixed tensor shapes—they cannot dynamically resize buffers between captures.
3. DDTree algorithm: The tree-structured speculative decoding approach where the draft model proposes multiple candidate paths (up to budget nodes), and the target model verifies all paths simultaneously in a single forward pass using a custom tree-attention mask. This is the key difference from linear DFlash: DDTree's verify batch size is budget + 1, not block_size.
4. SGLang inference stack: The serving framework being used, with its layered architecture of speculative workers, model runners, and CUDA graph optimization. The assistant references specific files (dflash_worker.py, model_runner.py, cuda_graph_runner.py) and line numbers from the deployed codebase.
5. The specific model configuration: Kimi K2.6 (1T MoE, INT4 quantized), the DFlash draft model (SubSir/Kimi-K2.6-DFlash-tmp-long with block_size=8), and the DDTree hyperparameters (budget=32, topk_cap=8).

Output Knowledge Created

This message creates several valuable outputs:

1. A confirmed root cause: The tensor size mismatch is definitively traced to the CUDA graph capture dimension, not to a logic error in the verify algorithm. This narrows the fix to a single parameter.
2. A clear fix direction: The assistant identifies that num_tokens_per_bs must be set to budget + 1 (33) when DDTree is active, rather than block_size (8). The fix location is in the CUDA graph runner initialization, where the speculative algorithm type should be checked.
3. A reproducible diagnostic procedure: The assistant demonstrates a systematic approach: trigger the crash, capture the full traceback, identify the failing operation (tensor size check in graph replay), trace back through the call chain to find the parameter origin, and grep the codebase for all relevant parameter references.
4. A codebase map: The grep output provides a map of how num_tokens_per_bs flows through the CUDA graph runner—from constructor parameter (line 178) to buffer allocation (line 199) to tensor fill operations (lines 350-355). This is invaluable for anyone unfamiliar with this code.

The Thinking Process: A Window into Debugging

The assistant's reasoning in this message reveals a structured debugging thought process:

Step 1 — Observe the symptom: The crash is a tensor size mismatch: a (8) vs b (33). The numbers 8 and 33 are not arbitrary—8 is the DFlash block size, and 33 is the DDTree budget (32) plus one (the root node).

Step 2 — Trace the stack: The error surfaces in cuda_graph_runner.py at the replay method (line 1306), called from model_runner.py (line 3223), which is called from dflash_worker.py. The assistant doesn't need to read the full stack—it already knows the call chain from previous work.

Step 3 — Formulate a hypothesis: The CUDA graph was captured with buffers sized for 8 tokens per batch slot, but DDTree's verify forward sends 33 tokens. The mismatch is in the graph capture, not in the forward logic.

Step 4 — Gather evidence: The assistant greps the codebase for the key parameter (num_tokens_per_bs) and related variables to confirm that the parameter controls buffer sizing and to understand how it's currently set.

Step 5 — Plan the fix: The fix is to make num_tokens_per_bs conditional on the speculative algorithm—use budget + 1 for DDTree, block_size for linear DFlash, and 1 for no speculative decoding.

This is a textbook example of debugging a complex system: start from the error message, use the numbers to infer the semantic meaning, trace the call chain, identify the parameter that controls the dimension, and verify with code search.

Broader Implications

The bug diagnosed in this message has implications beyond this specific deployment. It highlights a fundamental tension in CUDA graph optimization: graphs are captured with fixed tensor shapes for maximum performance, but speculative decoding algorithms have variable batch sizes depending on the algorithm variant. Linear DFlash sends block_size tokens per verify step, DDTree sends budget + 1 tokens, and future algorithms might send different numbers.

The cleanest fix—and the one the assistant is implicitly pointing toward—is to make the CUDA graph capture aware of the speculative algorithm configuration. This means either:

• Capturing separate graphs for each algorithm mode (expensive in memory and capture time)
• Capturing graphs for the maximum possible batch size and padding smaller batches (wastes GPU memory and compute)
• Making the verify path fall back to eager mode (no CUDA graph) when the batch size doesn't match the captured graph (loses the graph speedup) The assistant's approach—modifying the num_tokens_per_bs parameter based on the algorithm—is the most elegant solution because it requires only a single graph capture with the correct size, and the graph runner already supports the parameterization.

Conclusion

The subject message <msg id=11606> captures a pivotal moment in a complex ML inference deployment. The assistant's systematic debugging—tracing a tensor size mismatch from the error message through the call stack to the root cause in CUDA graph capture parameters—demonstrates the deep systems knowledge required to deploy large language models with cutting-edge speculative decoding techniques. The fix is conceptually simple (change one parameter), but finding it required understanding the interaction between CUDA graph optimization, the DDTree verification algorithm, and the SGLang inference stack. This message is a microcosm of the entire session: a relentless pursuit of performance through debugging, optimization, and deep architectural understanding.