JetSpec tree drafting on DFlash — lossless; width=1 clean drop-in#8
Open
vincentzed wants to merge 41 commits into
Open
JetSpec tree drafting on DFlash — lossless; width=1 clean drop-in#8vincentzed wants to merge 41 commits into
vincentzed wants to merge 41 commits into
Conversation
Dense (Qwen3-8B) tree verify is TOKEN-EXACT vs the fresh linear oracle:
10/10 prompts byte-identical across tree_width {4,7} x budget {64,128}
(runs jetspec/runs/8b_tree_*_deterministic_default_expanded_nograph_*).
Real target tree verify — the branch-reverify band-aid is removed.
MoE (Qwen3.6-35B-A3B GDN) status:
- decode/linear path: lossless.
- tree verify: WIP. Root cause found + fix landed — the retrieve links are
split so tree topology (all siblings) feeds the GDN conv/SSM state kernels
while duplicate-token overwrite stays only in acceptance. Residual: one
prompt (prompt 3) still diverges in the recurrent-state commit, tangled
with a CUDA illegal-memory-access under debug.
Linear-attn backend: cutedsl/triton (never FlashInfer GDN, which asserts on
tree verify sizes; verify falls back to triton).
CUDA graph for the tree path: not yet (runs eager); to be aligned with the
EAGLE3 topk>1 cuda-graph strategy.
Unit tests: 18 passing. Checkpoint before a fresh debugging pass.
Continues 8f9e732. MoE (Qwen3.6 GDN) tree verify reduced from 10/10 mismatches to a single holdout (prompt 4) on the fixed 10-prompt set. Landed: retrieve-link topology/accept split, an env-selectable full-attention verify backend for the hybrid model (SGLANG_DFLASH_TREE_VERIFY_FULL_ATTN_BACKEND), and direct-snapshot KV/state commit work. KNOWN RISK: a control-flow path in the mamba-state commit branch (accept_index_for_mamba_commit) may be referenced outside its init branch after the direct-snapshot path — must be verified/cleaned. Files compile; 18 unit tests pass. Dense path proven lossless at 8f9e732 (re-verify after these edits).
Attempted dropping the dense accepted-path reverify (EAGLE-style direct KV commit). FAILED the fresh flushed gate (6/10 mismatch), so reverify stays. Localization (per-layer diagnostic, debug-gated): for the accepted path, layer-0 K/V projections are BIT-EQUAL to a clean causal forward (kv_max_abs=0.0) and positions/RoPE are ruled out, but the layer-0 ATTENTION OUTPUT diverges (first_layer_hidden_delta~0.008, growing to hidden_max_abs~8). => the accepted node attends to a different KV SET in the compact FlashInfer one-query-row verify than a multi-token causal prefill would, OR it is a FlashInfer compact-verify kernel numeric difference. Next lever: audit the per-node compact_kv_indices / custom mask = exact ancestor set; if correct, test a reference (non-FlashInfer) attention to decide logic-bug vs kernel-numeric. Canonical baseline unchanged: compact tree w7/b64 + reverify = 5.08 accept / 283.9 tok/s vs linear 4.11 / 774.9 (0.37x). MoE untouched.
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request extends SGLang's DFlash speculative decoding with parallel tree drafting (JetSpec), introducing a top2gap tree construction mode, a paged-tree verification kernel for Blackwell (SM100+), and CUDA graph support for compact and paged tree verification. The review comments are highly constructive and should be addressed to improve numerical stability and robustness. Specifically, you should clamp the input to math.exp to prevent overflow, use torch.log_softmax for entropy calculation, add bounds clamping to prevent out-of-bounds indexing in the ancestor matrix builder, and restrict speculative config key validation to when the active algorithm is DFlash.
Important
The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.
Comment on lines
+198
to
+203
| arg = -float(beta) * (float(gap) - float(g_0)) | ||
| if arg >= 0: | ||
| sigmoid = 1.0 / (1.0 + math.exp(-arg)) | ||
| else: | ||
| exp_arg = math.exp(arg) | ||
| sigmoid = exp_arg / (1.0 + exp_arg) |
There was a problem hiding this comment.
To prevent potential OverflowError in math.exp when arg is extremely large (which can happen with unnormalized logits or large beta values), clamp arg to a safe range (e.g., [-700.0, 700.0]) before evaluating the sigmoid.
Suggested change
| arg = -float(beta) * (float(gap) - float(g_0)) | |
| if arg >= 0: | |
| sigmoid = 1.0 / (1.0 + math.exp(-arg)) | |
| else: | |
| exp_arg = math.exp(arg) | |
| sigmoid = exp_arg / (1.0 + exp_arg) | |
| arg = -float(beta) * (float(gap) - float(g_0)) | |
| arg = max(-700.0, min(700.0, arg)) | |
| if arg >= 0: | |
| sigmoid = 1.0 / (1.0 + math.exp(-arg)) | |
| else: | |
| exp_arg = math.exp(arg) | |
| sigmoid = exp_arg / (1.0 + exp_arg) |
Comment on lines
+11
to
+14
| def compute_per_depth_entropy(logits: torch.Tensor) -> list[float]: | ||
| probs = torch.softmax(logits, dim=-1) | ||
| ent = -(probs * torch.log(probs + 1e-10)).sum(dim=-1) | ||
| return ent.tolist() |
There was a problem hiding this comment.
Using torch.log_softmax is numerically more stable and efficient than adding a small epsilon like 1e-10 to the raw probabilities before taking torch.log.
Suggested change
| def compute_per_depth_entropy(logits: torch.Tensor) -> list[float]: | |
| probs = torch.softmax(logits, dim=-1) | |
| ent = -(probs * torch.log(probs + 1e-10)).sum(dim=-1) | |
| return ent.tolist() | |
| def compute_per_depth_entropy(logits: torch.Tensor) -> list[float]: | |
| probs = torch.softmax(logits, dim=-1) | |
| ent = -(probs * torch.log_softmax(logits, dim=-1)).sum(dim=-1) | |
| return ent.tolist() |
| mask[:, 0] = True | ||
| if num_nodes == 1: | ||
| return mask | ||
|
|
There was a problem hiding this comment.
Clamp the upper bound of current to num_nodes - 1 to prevent potential out-of-bounds indexing on parent_tensor if any invalid parent indices are passed.
Suggested change
| current = torch.where(valid, parent_tensor[current.clamp(min=0, max=num_nodes - 1)], current) |
Comment on lines
+435
to
+443
| for key, value in config.items(): | ||
| dest = _DFLASH_SPECULATIVE_CONFIG_KEYS.get(key) | ||
| if dest is None: | ||
| allowed = ", ".join(sorted(_DFLASH_SPECULATIVE_CONFIG_KEYS)) | ||
| raise ValueError( | ||
| f"Unsupported --speculative-config key {key!r}. " | ||
| f"Supported DFLASH keys: {allowed}." | ||
| ) | ||
| kwargs[dest] = value |
There was a problem hiding this comment.
To avoid breaking --speculative-config validation for other speculative decoding algorithms (e.g., EAGLE) that might be configured or added in the future, restrict the strict DFlash key validation to when the active speculative algorithm is indeed DFLASH.
Suggested change
| for key, value in config.items(): | |
| dest = _DFLASH_SPECULATIVE_CONFIG_KEYS.get(key) | |
| if dest is None: | |
| allowed = ", ".join(sorted(_DFLASH_SPECULATIVE_CONFIG_KEYS)) | |
| raise ValueError( | |
| f"Unsupported --speculative-config key {key!r}. " | |
| f"Supported DFLASH keys: {allowed}." | |
| ) | |
| kwargs[dest] = value | |
| algo = str(kwargs.get("speculative_algorithm", "")).upper() | |
| for key, value in config.items(): | |
| dest = _DFLASH_SPECULATIVE_CONFIG_KEYS.get(key) | |
| if dest is None: | |
| if algo == "DFLASH": | |
| allowed = ", ".join(sorted(_DFLASH_SPECULATIVE_CONFIG_KEYS)) | |
| raise ValueError( | |
| f"Unsupported --speculative-config key {key!r}. " | |
| f"Supported DFLASH keys: {allowed}." | |
| ) | |
| continue | |
| kwargs[dest] = value |
![cubic-dev-ai[bot]](https://avatars.githubusercontent.com/in/1082092?s=60&v=4){kind=small}
There was a problem hiding this comment.
25 issues found across 47 files
Prompt for AI agents (unresolved issues)
Check if these issues are valid — if so, understand the root cause of each and fix them. If appropriate, use sub-agents to investigate and fix each issue separately.
<file name="python/sglang/srt/layers/attention/triton_backend.py">
<violation number="1" location="python/sglang/srt/layers/attention/triton_backend.py:800">
P1: Compact verify metadata must derive custom-mask offsets from compact_kv_indptr/compact_qo_indptr, not forward_batch.seq_lens. With ragged compact spans, mask_indptr drifts and applies another request’s tree mask, corrupting verification logits.</violation>
</file>
<file name="jetspec/run_top2gap_sweep.py">
<violation number="1" location="jetspec/run_top2gap_sweep.py:115">
P2: Do not hard-code GPU 7 in the sweep runner. Respect CUDA_VISIBLE_DEVICES from the caller so the documented local repro works across machines.</violation>
<violation number="2" location="jetspec/run_top2gap_sweep.py:118">
P2: Do not overwrite an existing Hugging Face token with an empty HF_TOKEN fallback. This can make private/gated model downloads fail even when HUGGING_FACE_HUB_TOKEN is correctly set.</violation>
</file>
<file name="python/sglang/srt/models/dflash.py">
<violation number="1" location="python/sglang/srt/models/dflash.py:52">
P1: `--speculative-dflash-head-type causal` does not affect the draft model attention type, so tree mode can run with bidirectional `ENCODER_ONLY` layers. Thread the resolved head type into the model config or reject the override when the checkpoint lacks `causal_head=True`.</violation>
</file>
<file name="jetspec/bench_paper_sglang.py">
<violation number="1" location="jetspec/bench_paper_sglang.py:338">
P2: Warmup reuses measured prompts without clearing the server cache by default. This can inflate benchmark throughput and make DFlash/JetSpec comparisons depend on whether the first samples were pre-cached.</violation>
</file>
<file name="python/sglang/srt/model_executor/runner/eager_runner.py">
<violation number="1" location="python/sglang/srt/model_executor/runner/eager_runner.py:117">
P1: Do not cast speculative_dflash_tree_budget without handling None; DFlash tree mode supports omitting this option and computing the budget. This makes width>1 fail at startup unless users pass an explicit tree budget.</violation>
</file>
<file name="jetspec/_oldattempt/NOTES.md">
<violation number="1" location="jetspec/_oldattempt/NOTES.md:8">
P2: Claimed `realtree_wip.patch` file does not exist in the repository. The note says WIP edits are backed up there, but the file is absent, making this reference misleading.</violation>
</file>
<file name="python/sglang/srt/layers/attention/triton_ops/extend_attention.py">
<violation number="1" location="python/sglang/srt/layers/attention/triton_ops/extend_attention.py:405">
P2: Hard-coding `allow_tf32=False` in the shared extend-attention kernels disables Triton's default NVIDIA TF32 path for all fp32 callers. Gate exact matmul behind the deterministic/JetSpec path or leave generic DFlash precision unchanged.</violation>
</file>
<file name="jetspec/notes/crossproduct_tree_spec.md">
<violation number="1" location="jetspec/notes/crossproduct_tree_spec.md:1">
P2: Missing verify overhead factors: compact tree ~1.29× per-step and paged tree ~1.44× per-step must be documented per README requirement.</violation>
<violation number="2" location="jetspec/notes/crossproduct_tree_spec.md:1">
P2: Missing break‑even condition: the spec must document `accept_ratio > step_cost_ratio` as required by README.</violation>
</file>
<file name="jetspec/verify/CANONICAL_COMMAND.md">
<violation number="1" location="jetspec/verify/CANONICAL_COMMAND.md:21">
P2: Inline shell comment after backslash continuation breaks copy-paste — spaces between `\` and `#` prevent line continuation. Move the comment to its own line or remove `\` on the commented line.</violation>
</file>
<file name="jetspec/hf_configs/JetSpec--jetspec-Qwen3.6-35B-A3B/config.json">
<violation number="1" location="jetspec/hf_configs/JetSpec--jetspec-Qwen3.6-35B-A3B/config.json:11">
P2: bos_token_id 151643 doesn't match the target model's bos_token_id (248044). Likely copy-pasted from the Qwen3-8B JetSpec config. The draft model should share the same bos_token as the target for correct tokenization.</violation>
<violation number="2" location="jetspec/hf_configs/JetSpec--jetspec-Qwen3.6-35B-A3B/config.json:24">
P2: eos_token_id 248046 doesn't match the target model's eos_token_id (248044). In spec-decode the draft shares the target's tokenizer, so EOS must match for correct stop-token detection.</violation>
</file>
<file name="python/sglang/srt/layers/attention/torch_native_backend.py">
<violation number="1" location="python/sglang/srt/layers/attention/torch_native_backend.py:396">
P1: Custom-mask target verify ignores sliding-window attention and can attend to the entire prefix. Pass the layer's sliding window into this path and intersect it with the tree mask before SDPA.</violation>
</file>
<file name="python/sglang/srt/layers/attention/flashinfer_backend.py">
<violation number="1" location="python/sglang/srt/layers/attention/flashinfer_backend.py:1828">
P1: This overstates compact-tree KV lengths during fast_prefill_plan replay. Compact FA4 CUDA graph can plan past the provided compact_kv_indices, producing invalid attention or OOB reads.</violation>
</file>
<file name="jetspec/run_dflash_gate_bench.sh">
<violation number="1" location="jetspec/run_dflash_gate_bench.sh:208">
P2: lean-b16-pair loses the first configuration's summary because summarize overwrites summary.ndjson on every call. Append rows instead so both beta/g0 pairs remain in the CI artifact.</violation>
<violation number="2" location="jetspec/run_dflash_gate_bench.sh:282">
P2: profile mode depends on a hardcoded /root/.claude absolute path, making the documented script non-portable outside the author's machine. Make the analyzer path configurable or use a checked-in/repo-relative tool.</violation>
</file>
<file name="python/sglang/srt/speculative/dflash_info.py">
<violation number="1" location="python/sglang/srt/speculative/dflash_info.py:46">
P1: New verify metadata tensors are not included in overlap record_stream handling. Compact/paged tree verify can then use freed/reused CUDA memory if the producing stream drops these tensors before the forward stream finishes.</violation>
</file>
<file name="python/sglang/srt/speculative/spec_utils.py">
<violation number="1" location="python/sglang/srt/speculative/spec_utils.py:597">
P1: Overlap-safe KV compaction does not move FP4 KV scale buffers. FP4 target KV cache rows copied by DFlash tree verify will dequantize with stale per-token scales after acceptance.</violation>
</file>
<file name="python/sglang/srt/layers/attention/flashattention_backend.py">
<violation number="1" location="python/sglang/srt/layers/attention/flashattention_backend.py:785">
P1: Compact tree verify passes the prefix length as `max_seq_len_k`, so varlen attention can under-size rows that include draft/ancestor KV entries. Compute the max from `compact_kv_indptr` (or `compact_kv_max_seq_len`) instead.</violation>
<violation number="2" location="python/sglang/srt/layers/attention/flashattention_backend.py:1196">
P1: DFlash tree verify is silently routed through the normal MLA attention path when `self.use_mla` is true. Add an explicit unsupported guard or implement the compact/paged tree handling for MLA before enabling these metadata modes.</violation>
</file>
<file name="python/sglang/srt/model_executor/model_runner_kv_cache_mixin.py">
<violation number="1" location="python/sglang/srt/model_executor/model_runner_kv_cache_mixin.py:342">
P2: DFLASH tree extra request rows are not applied to decode-disaggregation pools. Use `req_to_token_pool_size` for the non-Mamba `DecodeReqToTokenPool` too, otherwise width>1 expanded tree verify can fail from insufficient temporary request rows.</violation>
</file>
<file name="python/sglang/srt/arg_groups/speculative_hook.py">
<violation number="1" location="python/sglang/srt/arg_groups/speculative_hook.py:296">
P2: This check does not actually guarantee the root-inclusive linear DFLASH path is present in tree mode. With width>1 and budget=block_size, the breadth-first builder can spend nodes on siblings and shorten the preserved top-1 path, undermining the advertised linear-path preservation for compact trees.</violation>
</file>
<file name="python/sglang/jit_kernel/dflash_paged_tree_verify.py">
<violation number="1" location="python/sglang/jit_kernel/dflash_paged_tree_verify.py:152">
P1: `flash_attn_varlen_func` does not expose a `mask_mod` keyword, so the new paged-tree verify path will fail at runtime instead of applying the tree mask. Route this through a supported FA4 mask API or add proper wrapper support before calling it here.</violation>
</file>
<file name="python/sglang/srt/model_executor/runner/decode_cuda_graph_runner.py">
<violation number="1" location="python/sglang/srt/model_executor/runner/decode_cuda_graph_runner.py:893">
P2: Compact FA4 graph eligibility check is too permissive when compact metadata is missing. This can select CUDA-graph replay with a variant key that was never captured.</violation>
</file>
Tip: cubic can generate docs of your entire codebase and keep them up to date. Try it here.
Re-trigger cubic
| max_extend_len = self.num_draft_tokens | ||
| if custom_mask is not None: | ||
| seq_mask_len = num_draft_tokens * ( | ||
| forward_batch.seq_lens[:bs] + num_draft_tokens |
There was a problem hiding this comment.
P1: Compact verify metadata must derive custom-mask offsets from compact_kv_indptr/compact_qo_indptr, not forward_batch.seq_lens. With ragged compact spans, mask_indptr drifts and applies another request’s tree mask, corrupting verification logits.
Prompt for AI agents
Check if this issue is valid — if so, understand the root cause and fix it. At python/sglang/srt/layers/attention/triton_backend.py, line 800:
<comment>Compact verify metadata must derive custom-mask offsets from compact_kv_indptr/compact_qo_indptr, not forward_batch.seq_lens. With ragged compact spans, mask_indptr drifts and applies another request’s tree mask, corrupting verification logits.</comment>
<file context>
@@ -783,13 +795,18 @@ def init_forward_metadata(self, forward_batch: ForwardBatch):
- max_extend_len = self.num_draft_tokens
+ if custom_mask is not None:
+ seq_mask_len = num_draft_tokens * (
+ forward_batch.seq_lens[:bs] + num_draft_tokens
+ )
+ mask_indptr = torch.empty(
</file context>
| return bool(getattr(dflash_config, "causal_head")) | ||
| if hasattr(config, "causal_head"): | ||
| return bool(getattr(config, "causal_head")) | ||
| return False |
There was a problem hiding this comment.
P1: --speculative-dflash-head-type causal does not affect the draft model attention type, so tree mode can run with bidirectional ENCODER_ONLY layers. Thread the resolved head type into the model config or reject the override when the checkpoint lacks causal_head=True.
Prompt for AI agents
Check if this issue is valid — if so, understand the root cause and fix it. At python/sglang/srt/models/dflash.py, line 52:
<comment>`--speculative-dflash-head-type causal` does not affect the draft model attention type, so tree mode can run with bidirectional `ENCODER_ONLY` layers. Thread the resolved head type into the model config or reject the override when the checkpoint lacks `causal_head=True`.</comment>
<file context>
@@ -41,12 +41,25 @@
+ return bool(getattr(dflash_config, "causal_head"))
+ if hasattr(config, "causal_head"):
+ return bool(getattr(config, "causal_head"))
+ return False
+
+
</file context>
| and mr.spec_algorithm.is_dflash() | ||
| and sa.speculative_dflash_tree_width > 1 | ||
| ): | ||
| max_bs *= int(sa.speculative_dflash_tree_budget) |
There was a problem hiding this comment.
P1: Do not cast speculative_dflash_tree_budget without handling None; DFlash tree mode supports omitting this option and computing the budget. This makes width>1 fail at startup unless users pass an explicit tree budget.
Prompt for AI agents
Check if this issue is valid — if so, understand the root cause and fix it. At python/sglang/srt/model_executor/runner/eager_runner.py, line 117:
<comment>Do not cast speculative_dflash_tree_budget without handling None; DFlash tree mode supports omitting this option and computing the budget. This makes width>1 fail at startup unless users pass an explicit tree budget.</comment>
<file context>
@@ -102,6 +109,12 @@ def __init__(self, model_runner: ModelRunner) -> None:
+ and mr.spec_algorithm.is_dflash()
+ and sa.speculative_dflash_tree_width > 1
+ ):
+ max_bs *= int(sa.speculative_dflash_tree_budget)
# Mirror prepare_mlp_sync_batch padding so the registry holds what load_batch copies.
if require_mlp_sync(sa):
</file context>
| and spec_info is not None | ||
| and getattr(spec_info, "custom_mask", None) is not None | ||
| ): | ||
| self._run_sdpa_forward_target_verify( |
There was a problem hiding this comment.
P1: Custom-mask target verify ignores sliding-window attention and can attend to the entire prefix. Pass the layer's sliding window into this path and intersect it with the tree mask before SDPA.
Prompt for AI agents
Check if this issue is valid — if so, understand the root cause and fix it. At python/sglang/srt/layers/attention/torch_native_backend.py, line 396:
<comment>Custom-mask target verify ignores sliding-window attention and can attend to the entire prefix. Pass the layer's sliding window into this path and intersect it with the tree mask before SDPA.</comment>
<file context>
@@ -309,6 +387,27 @@ def forward_extend(
+ and spec_info is not None
+ and getattr(spec_info, "custom_mask", None) is not None
+ ):
+ self._run_sdpa_forward_target_verify(
+ q_,
+ o_,
</file context>
| isinstance(spec_info, SpecInput) | ||
| and spec_info.spec_input_type == SpecInputType.DFLASH_VERIFY | ||
| ): | ||
| kv_lens_cpu_i32 = kv_lens_cpu_i32 + int(num_tokens_per_req) |
There was a problem hiding this comment.
P1: This overstates compact-tree KV lengths during fast_prefill_plan replay. Compact FA4 CUDA graph can plan past the provided compact_kv_indices, producing invalid attention or OOB reads.
Prompt for AI agents
Check if this issue is valid — if so, understand the root cause and fix it. At python/sglang/srt/layers/attention/flashinfer_backend.py, line 1828:
<comment>This overstates compact-tree KV lengths during fast_prefill_plan replay. Compact FA4 CUDA graph can plan past the provided compact_kv_indices, producing invalid attention or OOB reads.</comment>
<file context>
@@ -1792,6 +1820,12 @@ def call_begin_forward(
+ isinstance(spec_info, SpecInput)
+ and spec_info.spec_input_type == SpecInputType.DFLASH_VERIFY
+ ):
+ kv_lens_cpu_i32 = kv_lens_cpu_i32 + int(num_tokens_per_req)
qo_indptr_host = torch.arange(
0,
</file context>
Suggested change
| kv_lens_cpu_i32 = kv_lens_cpu_i32 + int(num_tokens_per_req) | |
| if getattr(spec_info, "compact_kv_indices", None) is None: | |
| kv_lens_cpu_i32 = kv_lens_cpu_i32 + int(num_tokens_per_req) |
| return | ||
| fi | ||
| launch_tree "$port" "$LOG_DIR/profile_top2gap_w8_b16_${port}.log" 8 16 1.0 1.0 "$PROFILE_MEM_FRACTION_STATIC" | ||
| python3 /root/.claude/skills/llm-torch-profiler-analysis/scripts/analyze_llm_torch_profile.py \ |
There was a problem hiding this comment.
P2: profile mode depends on a hardcoded /root/.claude absolute path, making the documented script non-portable outside the author's machine. Make the analyzer path configurable or use a checked-in/repo-relative tool.
Prompt for AI agents
Check if this issue is valid — if so, understand the root cause and fix it. At jetspec/run_dflash_gate_bench.sh, line 282:
<comment>profile mode depends on a hardcoded /root/.claude absolute path, making the documented script non-portable outside the author's machine. Make the analyzer path configurable or use a checked-in/repo-relative tool.</comment>
<file context>
@@ -0,0 +1,313 @@
+ return
+ fi
+ launch_tree "$port" "$LOG_DIR/profile_top2gap_w8_b16_${port}.log" 8 16 1.0 1.0 "$PROFILE_MEM_FRACTION_STATIC"
+ python3 /root/.claude/skills/llm-torch-profiler-analysis/scripts/analyze_llm_torch_profile.py \
+ --framework sglang \
+ --url "http://127.0.0.1:${port}" \
</file context>
Comment on lines
+208
to
+210
| summary_path.write_text( | ||
| "".join(json.dumps(row, sort_keys=True) + "\n" for row in rows) | ||
| ) |
There was a problem hiding this comment.
P2: lean-b16-pair loses the first configuration's summary because summarize overwrites summary.ndjson on every call. Append rows instead so both beta/g0 pairs remain in the CI artifact.
Prompt for AI agents
Check if this issue is valid — if so, understand the root cause and fix it. At jetspec/run_dflash_gate_bench.sh, line 208:
<comment>lean-b16-pair loses the first configuration's summary because summarize overwrites summary.ndjson on every call. Append rows instead so both beta/g0 pairs remain in the CI artifact.</comment>
<file context>
@@ -0,0 +1,313 @@
+ "tok_s": summary["throughput_wall_tok_s"],
+ }
+ )
+summary_path.write_text(
+ "".join(json.dumps(row, sort_keys=True) + "\n" for row in rows)
+)
</file context>
Suggested change
| summary_path.write_text( | |
| "".join(json.dumps(row, sort_keys=True) + "\n" for row in rows) | |
| ) | |
| with summary_path.open("a") as f: | |
| for row in rows: | |
| f.write(json.dumps(row, sort_keys=True) + "\n") |
| def _init_pools(self: ModelRunner): | ||
| """Initialize the memory pools.""" | ||
| max_num_reqs = self.max_running_requests | ||
| req_to_token_pool_size = max_num_reqs |
There was a problem hiding this comment.
P2: DFLASH tree extra request rows are not applied to decode-disaggregation pools. Use req_to_token_pool_size for the non-Mamba DecodeReqToTokenPool too, otherwise width>1 expanded tree verify can fail from insufficient temporary request rows.
Prompt for AI agents
Check if this issue is valid — if so, understand the root cause and fix it. At python/sglang/srt/model_executor/model_runner_kv_cache_mixin.py, line 342:
<comment>DFLASH tree extra request rows are not applied to decode-disaggregation pools. Use `req_to_token_pool_size` for the non-Mamba `DecodeReqToTokenPool` too, otherwise width>1 expanded tree verify can fail from insufficient temporary request rows.</comment>
<file context>
@@ -335,6 +339,16 @@ def _validate_prefill_only_disable_kv_cache_pool_family(
def _init_pools(self: ModelRunner):
"""Initialize the memory pools."""
max_num_reqs = self.max_running_requests
+ req_to_token_pool_size = max_num_reqs
+ if (
+ self.server_args.speculative_algorithm == "DFLASH"
</file context>
| tree_budget, | ||
| ) | ||
| tree_budget = block_size | ||
| elif tree_budget < block_size: |
There was a problem hiding this comment.
P2: This check does not actually guarantee the root-inclusive linear DFLASH path is present in tree mode. With width>1 and budget=block_size, the breadth-first builder can spend nodes on siblings and shorten the preserved top-1 path, undermining the advertised linear-path preservation for compact trees.
Prompt for AI agents
Check if this issue is valid — if so, understand the root cause and fix it. At python/sglang/srt/arg_groups/speculative_hook.py, line 296:
<comment>This check does not actually guarantee the root-inclusive linear DFLASH path is present in tree mode. With width>1 and budget=block_size, the breadth-first builder can spend nodes on siblings and shorten the preserved top-1 path, undermining the advertised linear-path preservation for compact trees.</comment>
<file context>
@@ -220,6 +221,94 @@ def _handle_dflash(server_args: ServerArgs) -> None:
+ tree_budget,
+ )
+ tree_budget = block_size
+ elif tree_budget < block_size:
+ raise ValueError(
+ "--speculative-dflash-tree-budget must be >= the DFLASH block size "
</file context>
Comment on lines
+893
to
+906
| is_dflash_compact_supported = True | ||
| if ( | ||
| self.dflash_tree_compact_fa4_graph | ||
| and forward_batch.forward_mode.is_target_verify() | ||
| and getattr(forward_batch.spec_info, "compact_kv_indices", None) | ||
| is not None | ||
| ): | ||
| is_dflash_compact_supported = ( | ||
| forward_batch.input_ids.numel() == forward_batch.batch_size | ||
| and forward_batch.batch_size | ||
| <= self.dflash_tree_compact_fa4_graph_max_rows | ||
| and self._dflash_compact_kv_bucket_for_forward_batch(forward_batch) | ||
| is not None | ||
| ) |
There was a problem hiding this comment.
P2: Compact FA4 graph eligibility check is too permissive when compact metadata is missing. This can select CUDA-graph replay with a variant key that was never captured.
Prompt for AI agents
Check if this issue is valid — if so, understand the root cause and fix it. At python/sglang/srt/model_executor/runner/decode_cuda_graph_runner.py, line 893:
<comment>Compact FA4 graph eligibility check is too permissive when compact metadata is missing. This can select CUDA-graph replay with a variant key that was never captured.</comment>
<file context>
@@ -458,12 +890,43 @@ def can_run_graph(self, forward_batch: ForwardBatch):
else True
)
+ is_dflash_compact_supported = True
+ if (
+ self.dflash_tree_compact_fa4_graph
</file context>
Suggested change
| is_dflash_compact_supported = True | |
| if ( | |
| self.dflash_tree_compact_fa4_graph | |
| and forward_batch.forward_mode.is_target_verify() | |
| and getattr(forward_batch.spec_info, "compact_kv_indices", None) | |
| is not None | |
| ): | |
| is_dflash_compact_supported = ( | |
| forward_batch.input_ids.numel() == forward_batch.batch_size | |
| and forward_batch.batch_size | |
| <= self.dflash_tree_compact_fa4_graph_max_rows | |
| and self._dflash_compact_kv_bucket_for_forward_batch(forward_batch) | |
| is not None | |
| ) | |
| is_dflash_compact_supported = True | |
| if self.dflash_tree_compact_fa4_graph and forward_batch.forward_mode.is_target_verify(): | |
| compact_indices = getattr(forward_batch.spec_info, "compact_kv_indices", None) | |
| is_dflash_compact_supported = ( | |
| compact_indices is not None | |
| and forward_batch.input_ids.numel() == forward_batch.batch_size | |
| and forward_batch.batch_size | |
| <= self.dflash_tree_compact_fa4_graph_max_rows | |
| and self._dflash_compact_kv_bucket_for_forward_batch(forward_batch) | |
| is not None | |
| ) |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
JetSpec tree drafting on SGLang DFlash
Extends SGLang's linear DFlash speculative decoding with parallel tree drafting (JetSpec): the DFlash draft head's per-depth block logits feed a token tree (top2gap construction), verified under a tree-causal mask.
tree_width=1is the unchanged linear path;width>1is the new tree path.What this delivers
top2gapconstruction — per-depth fanout from the top-1/top-2 logprob gap (branch where the drafter is uncertain, chain where confident). Reproduces the paper's accept length at a fraction of the budget.width=1is a byte-identical, near-zero-overhead drop-in for linear DFlash.Full speedup matrix (Qwen3-8B, B300, FA4 + page_size 16, greedy)
Why the tree could not improve further (the honest result)
accept_ratio > step_cost_ratio. The tree verify costs ~1.29× (compact) / 1.44× (paged) per step at equal node count, so it needs that much more acceptance to win.The good news: width=1 is a clean drop-in
width>1is available for workloads where verify is cheaper or the draft head concentrates better (e.g. weaker/AR baselines, or future paged-engine work).Reproduce
Detailed run-by-run notes:
jetspec/notes/bench_results.md; status:jetspec/STATUS.md.