[Spec] Add LiLiCorr: a candidate-lattice reranker for DFlash drafts

[mrusanovsky]

Motivation

DFlash is trained on per-position marginals rather than on the joint block distribution, so its
drafted tokens are individually plausible yet jointly incoherent — slot 3 may be a fine
continuation of the prompt while contradicting the token DFlash itself chose for slot 2.

LiLiCorr keeps the top-k candidates DFlash already produces at each block position and scores
transitions between adjacent candidates with a small two-layer transformer, then commits a path
through the resulting lattice greedily, left to right, instead of taking the per-slot argmax. One
network pass produces every vector, the pairwise scores are a single batched matmul, and only the
greedy walk stays sequential. Verify is untouched, so outputs remain distributionally identical to
the target model's — only acceptance length and throughput change.

• Paper: LiLiCorr: Lightweight Likelihood Correlation of Parallel Drafts for Speculative Decoding (arXiv:2608.20530)
• Blog: https://research.nvidia.com/labs/nemotron/lilicorr/
• Companion PR — training: LiLiCorr training NVIDIA/Model-Optimizer#2342

This PR is the serving half. It ships no checkpoints and no training code; the companion PR
above trains the drafters and exports them in the format this loader reads.

Performance

How these were produced. Six drafters for a Qwen3-8B target, all trained in NVIDIA
ModelOpt on one matched contract — the same corpus, schedule and block geometry for every arm, so
no row carries a training advantage. Training data is NVIDIA's Nemotron Post-Training Dataset v2
with the multilingual split excluded, generated from the target with thinking disabled;
6 epochs; block size 16 (15 drafted slots, 16 verified); DFlash decay objective at gamma 7;
8 nodes × 8 H100, global batch size 64 (one sequence per device, no gradient accumulation).

All six were then exported and served through the code in this PR on a single H100 80GB,
tp_size 1, at concurrency 1, greedy T=0, --attention-backend fa3, draft_length 15, mean of
two replicates, with the whole node held exclusive per benchmark so nothing else shared it.
Speedup is output tokens/s against an autoregressive baseline measured in the same allocation,
because a denominator borrowed from another node carries that node's clock into every ratio.

Three of the five alternatives — DFlash, DSpark, Domino — already live in this repo.

Cells are acceptance length / speedup-vs-AR; ★ fastest, ☆ second fastest:

benchmark	LiLiCorr+conv	LiLiCorr	DSpark	DFlash2	Domino	DFlash
gsm8k	★ 7.715 / 5.26x	☆ 7.557 / 5.22x	7.375 / 4.86x	7.252 / 5.06x	7.225 / 4.87x	6.341 / 4.59x
math500	★ 9.241 / 6.54x	☆ 9.064 / 6.52x	9.012 / 6.15x	8.999 / 6.49x	8.976 / 6.25x	7.909 / 5.88x
aime25	★ 8.285 / 6.03x	☆ 8.156 / 6.03x	8.043 / 5.61x	7.967 / 5.91x	8.066 / 5.77x	7.126 / 5.44x
humaneval	★ 7.393 / 4.01x	7.077 / 3.93x	7.163 / 3.72x	☆ 7.081 / 3.95x	6.864 / 3.73x	6.156 / 3.68x
mbpp_sanitized	★ 5.999 / 4.18x	☆ 5.849 / 4.13x	5.888 / 3.95x	5.685 / 4.05x	5.679 / 3.91x	5.027 / 3.70x
livecodebench	★ 7.975 / 5.40x	☆ 7.754 / 5.33x	7.775 / 5.10x	7.601 / 5.26x	7.553 / 5.04x	6.808 / 4.88x
alpaca_eval	☆ 3.697 / 2.69x	★ 3.656 / 2.70x	3.588 / 2.52x	3.467 / 2.58x	3.627 / 2.59x	3.222 / 2.46x
mtbench	★ 4.014 / 2.94x	☆ 3.939 / 2.93x	3.957 / 2.78x	3.748 / 2.80x	3.948 / 2.84x	3.478 / 2.67x

Against every other approach in the table, LiLiCorr with convolutions is the fastest on all eight
benchmarks. Plain LiLiCorr is the fastest on seven of the eight; the exception is humaneval, a
164-prompt slice, where DFlash2 is ahead by 0.5%.

DFlash is the deliberately head-free control. Against it, the reranked drafter is worth
+9.1% to +14.6% output tokens/s — the number to read if the question is "what does this buy over
what I already run":

benchmark	gsm8k	mbpp_sanitized	math500	livecodebench	mtbench	humaneval
throughput vs head-free DFlash	+14.6%	+12.9%	+11.1%	+10.5%	+9.9%	+9.1%

Every head also clears that control by +7.60% to +21.67% on acceptance, on every arm and every
benchmark, which is the check that a head actually loaded rather than silently falling back.

Acceptance length is bit-reproducible under greedy decoding; throughput is not. The replicate
spread on acceptance was 0.00% on every benchmark; on throughput it is about 0.2% within one
allocation, and about 1% across allocations.

ⓘ Reproducing these needs --attention-backend fa3. Backends that publish seq_lens_cpu every
step pay a device-to-host sync per block, worth about 8% here. That is a property of DFlash decoding
rather than of LiLiCorr, and the new docs subsection explains which backends are affected and why.
Backends also differ in attention numerics, so acceptance measured on different backends does not
belong in one table.

Modifications

No new SpeculativeAlgorithm, no worker subclass, no registration, no CLI flag, and no new
entry in any choices list. Two environment variables gate the opt-in sampled commit described
below; nothing else reads the environment. LiLiCorr rides --speculative-algorithm DFLASH and is
selected by the checkpoint declaring architectures: ["LiLiCorrDraftModel"], exactly as
DFlash2DraftModel selects the candidate selector. If you go looking for the registration, that is
why there isn't one.

10 files, +3,076 / −1. The single deleted line is one if in the draft loop that gained a
branch; everything else is an addition.

Only four files already existed — +164 / −1 between them:

file	added	what
srt/speculative/dflash_worker_v2.py	+95 / −1	five dispatch seams: two __init__ assignments, the fold branch in _maybe_build_draft_sampler, the guarded anchor publication inside _append_target_hidden_to_draft_kv_by_loc, and a fourth arm on the existing three-way draft-loop branch
srt/models/dflash.py	+2	DFlashDraftModel declares lilicorr: Optional[nn.Module] = None beside the existing candidate_selector, and stores the rms_norm_eps it already computes
docs/.../speculative_decoding.mdx	+34	two subsections under "DFlash Decoding"
test/registered/unit/spec/test_dflash_logits.py	+33	one worker-stub test: the sampled arm's device gate

For scale, selector (DFlash2 — the closest precedent, also a head on the same draft model, also
checkpoint-selected) appears on 43 lines of that worker; lilicorr appears on 30.

New, self-contained:

• srt/models/lilicorr.py (632) — the head, LiLiCorrDraftModel(DFlashDraftModel), EntryClass, and the two weight-coverage checks
• kernels/ops/speculative/lilicorr.py (520) — three Triton kernels; see "Why the custom kernels" below
• srt/speculative/lilicorr_utils.py (656) — head geometry, the candidate lattice, the eager draft seam and the CUDA-graph-folded draft sampler, in one module beside dflash_utils.py
• test/registered/unit/spec/test_lilicorr.py (805) — 46 CPU tests
• test/registered/kernels/ops/speculative/test_lilicorr_cuda.py (216) — 18 GPU tests pinning the Triton kernels against their torch references
• test/registered/e2e/speculative/test_dflash_lilicorr.py (84) — the E2E server test, registered disabled= until a drafter is published

The head runs inside the draft CUDA graph, and that is the whole optimization

The head is a long tail of small kernels, so what it costs is host dispatch, not FLOPs. It
therefore rides the seam DFlash already has for exactly this: _maybe_build_draft_sampler runs
before init_cuda_graphs, and the sampler is registered on
draft_model_runner.capture_tail_hooks through make_draft_sampler_capture_hook — the same
mechanism DFlash2's _SelectorDraftSampler and DSpark's DsparkDraftSampler use, and
LiLiCorrDraftSampler has the same shape as the latter: static buffers sized from the capture
buckets, host-side staging of the per-row sampling params before the replay, an in-graph philox
draw, and the drafted tokens written into a static out the worker reads after the replay.

Nothing in the head is capture-hostile: static shapes, no collectives at tp=1, no host syncs, a
fixed decode trip count, and every parameter-derived buffer built by
materialize_inference_buffers before capture.

There is no torch.compile anywhere in this PR. An earlier revision compiled the captured
body, on the theory that the graph removes the launches' host cost while leaving their number and
memory traffic alone. Measured end to end on one H100 — Qwen3-8B, gsm8k, greedy, fa3, tp_size 1,
both arms as sequential cells in one allocation so the node term cancels — it is a null:

requests	concurrency	compiled	no compile	Δ acceptance	Δ output tok/s
1319	1	6.9951 / 731.5	6.9947 / 735.6	−0.005%	+0.56%
1319	1	6.9951 / 739.7	6.9947 / 735.0	−0.005%	−0.62%
1319	32	7.0085 / 10179.5	7.0022 / 10131.9	−0.090%	−0.47%
320	32	7.0650 / 9410.8	7.0576 / 9489.0	−0.105%	+0.83%

Four pairs straddling zero, against a 1.12% spread for the same arm measured in two different
allocations. The fold, by contrast, is load-bearing: in the same allocation and on the same tree,
forcing the head out of the graph with SGLANG_DFLASH_EAGER_DRAFT_SAMPLER=1 costs −8.4%.

So there is no dynamo surface here at all — no compile prewarm, no recompile-limit handling, no
torch._inductor.config mutation. The one operational note that remains is that the head should run
on the folded path: an eager fallback exists for the steps the graph cannot serve and is correct but
costs a large fraction of throughput, and because acceptance is identical either way that shows up
only in tokens/s. build_lilicorr_draft_sampler warns with a reason when it declines to fold, and
the worker warns once if a decode step lands on the eager head.

Two capture gaps remain, both shared with the existing DFlash heads. Prefill and extend run eager —
the eager seam is the third arm of the same elif chain that already carries the selector's and
plain DFlash's. And tp>1 declines to fold, because the candidate top-k needs a packed all-gather
inside the graph.

Convolutional drafters are supported, and the support is a refusal

The leading column of the table above is LiLiCorr composed with the grouped dynamic convolutions
that wrap each DFlash sublayer. It serves through this PR, and that is worth stating precisely,
because it is not a feature this PR adds.

The convolution belongs to the DFlash backbone, not to the reranker: LiLiCorrDraftModel
extends DFlashDraftModel and therefore inherits the __init__ that builds it from
conv_kernel_size and conv_group_size in dflash_config. Both the plain and the convolutional
drafter already load and run through the same path, with no new flag and no new code.

What this PR contributes is the refusal. Both geometry keys default to 0, so a checkpoint whose
config lost them builds no convolution modules at all — and the backbone loader then drops every
convolution tensor it cannot resolve. That is correct behaviour for rotary caches and the worst
possible behaviour here: all 20 tensors of a five-layer drafter disappear in silence, the draft
serves as its convolution-free parent, and the only symptom is an acceptance length that is lower
than it should be and entirely believable. check_conv_weight_coverage refuses any checkpoint whose
convolution tensors and built modules disagree, in either direction. Five of the CPU tests cover it:
the matched case, the convolution-free case, both silent-drop directions, and a partial checkpoint.

Sampling the committed path (opt-in)

At T > 0 the target samples, but the head committed the per-slot argmax and verify was handed a
point mass, so each slot was paid p(argmax) rather than the overlap of the two distributions.
LILICORR_SAMPLING=1 draws each slot from softmax(scores / T) over that slot's candidates, at the
request's own temperature, and publishes the proposal it drew from into the same buffers the DFlash2
selector already fills — so the accept path is untouched and verify stays exact. Greedy requests
sharing the batch take the argmax and report a one-hot proposal. On Qwen3-8B it is worth about 6%
acceptance and 6% end-to-end at T = 1, at under 1% of the per-block wall.

Off by default, and with it unset the captured body is select itself rather than something
equivalent to it, so T = 0 output is unchanged bit for bit.
SGLANG_LILICORR_REQUIRE_SAMPLING=1 refuses to start if the mode was asked for but is not enabled.
Both are read once at import: the sampler is replayed from a captured draft graph, so a per-call
read would be a host-side branch inside a replay — happy to move them to ServerArgs fields
resolved at init if you would rather.

The sampled commit is additionally gated on the same supported-device condition as the selector's,
on both the folded and the eager path, because it publishes into an accept kernel that does not run
everywhere. Where it is unsupported the arm falls back to the argmax commit and the existing verify.

Why the custom kernels

kernels/ops/speculative/lilicorr.py is the single largest new file after the head itself, so it is
worth saying up front what it buys. All three kernels dispatch to a value-identical torch
implementation off CUDA, which is also what makes the head exercisable in a CPU unit test.

lilicorr_topk_lse — exact per-row top-k and the full-vocab log-partition from one pass over
[n, V]. The head scores log-probs normalized over the whole vocabulary, so it needs the partition
as well as the top-k. Composing that from existing ops (torch.topk followed by a logsumexp
epilogue) reads the vocabulary logits three times and materializes two more [n, V] temporaries,
and at a 150k vocabulary that read is the single largest cost in the block. Fusing them makes it one
read and no temporaries. The tile pre-selection is exact, not approximate — the argument is in the
function docstring and worth reading before trusting the kernel: if a top-k element's tile were not
among the k tiles with the largest maxima, k other tiles would each hold an element at least as
large, contradicting its rank.

lilicorr_greedy_path — the whole left-to-right commit in one kernel. The torch form issues
roughly three kernels per slot over [bs, k] tensors, so at 15 slots the decode is dozens of
few-microsecond launches and is pure launch overhead rather than arithmetic. Same recurrence, and
tl.argmax breaks ties toward the lower index as Tensor.argmax does, so the committed path is
identical.

lilicorr_sample_path — that same walk with a sampled commit, emitting the per-slot proposal the
verify needs to accept it by rejection sampling. A row with greedy_mask set takes tl.argmax on
the same fp32 node values as the greedy kernel, so one captured graph serves greedy and sampling
batches and the greedy rows walk a bit-identical path.

Correctness

Verify is unmodified, so this change cannot alter the distribution of emitted tokens. Every
drafted token is still checked against the target model exactly as before; a worse drafter would show
up as lower acceptance, never as different output. That is the primary correctness argument, and it
is why the section above measures acceptance rather than task scores.

Unit tests: 46 CPU tests (test_lilicorr.py), one worker-stub test in
test_dflash_logits.py, and 18 GPU tests (test_lilicorr_cuda.py) pinning all three Triton
kernels against their torch reference implementations on device — vocabulary sizes straddling the
kernel's tile boundary in both directions, bf16 logits as production uses, the narrow-vocabulary
fallback, and the greedy commit at k = 1…16 including tie-breaking. Verified on an H100: 75/75 pass
across the three files. The registry validates — validate_all_suites() is clean over 2,487
registered tests — and every pinned pre-commit hook passes over the PR range.

E2E: test/registered/e2e/speculative/test_dflash_lilicorr.py serves the head through
--speculative-algorithm DFLASH with only the draft path changed, and asserts GSM8K accuracy plus
two acceptance floors. It is registered disabled= because no drafter is published yet, but it has
been run against one on an H100: 1 passed in 177.19s, GSM8K 0.950, accept length 5.0392 against a
head-free DFLASH control reading 4.1694 on the same eval.

Checklist

• Format your code according to Format code with pre-commit — every pinned hook passes over the PR range
• Add unit tests according to Run and add unit tests — 46 CPU tests (register_cpu_ci), 18 GPU kernel-parity tests (register_cuda_ci), and an E2E server test on the DFlash fixture (register_cuda_ci, disabled= pending a published drafter)
• Update documentation according to Write documentations — a subsection under "DFlash Decoding"
• Provide accuracy and speed benchmark results — see Performance
• Follow the SGLang code style guidance

Design notes

The choices most likely to prompt a question, each also commented at the site.

1. The head geometry is parsed in srt/speculative/lilicorr_utils.py, not as a field on
   DFlashDraftConfig. The shared DFLASH config therefore carries no knowledge of this head, at the
   cost of one extra dflash_config read at model build, and dflash_utils.py stays at zero
   changes.
2. The lattice attention bias is materialized [batch*heads, S, S] rather than broadcast. The
   broadcast form is mathematically identical and saves a copy, but a stride-0 mask measured
   −1.75pp. Please measure before changing it back.
3. F.rms_norm rather than sglang.srt.layers.layernorm.RMSNorm, and a hand-rolled attention
   module carrying nn.MultiheadAttention's parameter layout. Same math and the same state_dict
   keys in both cases. The custom-op norm is not capturable here, and the module's Python control flow
   and need_weights bookkeeping are what make it uncapturable.
4. The weight loader raises in both directions — missing head tensors and surplus ones — and
   candidate_topk must be a power of two, refused at config parse. The surplus direction is live,
   not hypothetical: one projection is an Identity when the head is as wide as the draft, so a
   config omitting the head width would otherwise silently drop it and serve a different architecture
   than the one trained. The power-of-two constraint comes from the tiled top-k holding its selected
   tiles in a single Triton lane group; supporting other widths would mean serving the head on the
   slow reference path, so it is refused at load instead of silently deoptimized.
5. The convolution coverage check is separate from the head coverage check, because the two look
   at different namespaces: the convolution tensors are backbone parameters named
   layers.*.{attention,mlp}_conv.*, none of which is under lilicorr., so the head check cannot
   see them. See "Convolutional drafters" above for why a missing check is dangerous rather than
   merely untidy.
6. The anchor is worker-owned rather than carried on DFlashDraftInputV2. Carrying it there so
   that filter and merge keep it aligned with the batch is the tidier design, but it measured
   −4.67% acceptance at concurrency 32, and the drift it prevents is not detectable: reordering
   cannot happen at concurrency 1 and happens constantly at 32, where acceptance reads 7.5573 and
   7.5669. The measurement is recorded at the site in set_anchor.

Known gaps

• Checkpoints are not yet published. The docs example carries a placeholder path. The companion
  ModelOpt PR linked at the top adds the training and checkpoint export, including the optional
  convolutions.
• tp>1 serves the head eagerly. Folding it needs the candidate top-k's packed all-gather
  inside the draft graph.

---

CI States

Latest PR Test (Base): ❌ Run #36355573025
Latest PR Test (Extra): ❌ Run #36355572926
Latest PR Test (AMD ROCm 10): ❌ Run #36355573123

[kpham-sgl]

Thank you for the contribution! Some very high level comments before I review further

1. Please point your agents to these rules we have in .claude/rules/. Some to point out

• Avoid extensive AI comments and docstrings
• Avoid using getattr/hasattr
• Avoid writing unnecessary unit tests

2. Can you add an E2E test with a small target for this new algorithm? Make sure to use existing spec decoding test kit?

[kpham-sgl]

Can you also resolve conflicts @mrusanovsky? Thanks!

[nvpohanh]

[by Codex] Five inline review findings are attached.

[mrusanovsky]

Thank you for the contribution! Some very high level comments before I review further

1. Please point your agents to these rules we have in .claude/rules/. Some to point out

• Avoid extensive AI comments and docstrings
• Avoid using getattr/hasattr
• Avoid writing unnecessary unit tests

2. Can you add an E2E test with a small target for this new algorithm? Make sure to use existing spec decoding test kit?

Reading the rest of .claude/rules/ turned up several issues that were addressed.

getattr/hasattr : all removed:

• worker.lilicorr probed the draft model. Now reads DFlashDraftModel.lilicorr, declared None
  next to candidate_selector.
• LiLiCorrHead re-derived rms_norm_eps with a default. The base __init__ already computes it.
• resolve_vocab_shard probed for shard_indices. Now narrows on VocabParallelEmbedding.
• torch._dynamo / torch._inductor knobs assigned directly.

Remaining hasattr is in tearDownClass, required by write-sglang-test.

Comments and docstrings : docstring lines in non-test source down from 476 to 270. Removed
multi-line rationale on private helpers. Kept module docstrings, tensor shape contracts, and
provenance of measured constants.

Tests : 48 CPU to 37, 16 GPU to 12. Back to 40 and 13 after @nvpohanh's findings needed guards.
Deleted a tautology against if topk & (topk-1): raise, a case asserting a property of the test's
own fake, and one made unfalsifiable by an API change.

Also fixed : LiLiCorrConfig was a frozen @dataclass, now msgspec.Struct.
candidate_token_ids renamed to candidate_tokens per speculative-naming Rule 5.
build_lilicorr_draft_sampler and target_input_embeddings no longer take the whole worker.

E2E test : test/registered/spec/dflash/test_dflash_lilicorr.py, same fixture shape as
test_dflash.py with GSM8KMixin. Launch args are plain DFLASH with only the draft path changed.

Run on an H100: 1 passed in 177.19s, GSM8K 0.950, accept length 5.0392. The head-free control
reads 4.1694 on the same eval, so the floor is 4.6. est_time=300 measured at 177 s.

That accept length is the mixin's figure over 200 short 5-shot completions, not the block-weighted
acceptance in the PR body. Different quantities, noted at the site.

Registered disabled= since no drafter is published, same as test_spec_eagle_cpu.py. Draft path
is the literal <unpublished>, so enabling it without an id fails immediately.

SpecDecodingMixin is mixed in too, with both floors measured against the head-free control on the
same test: acceptance 8.192 vs 7.314, speed 883.8 vs 788.0 tok/s on an H100. Floors are 7.7 and 600.
The speed one is deliberately loose since throughput isn't reproducible.

One gap left: no small-target drafter exists, all of ours target Qwen3-8B, so it runs on
1-gpu-large like test_uno.py.

In addition, I resolved conflicts. Rebased on main. Only dflash_worker_v2.py conflicted, against the NPU DFlash2 change on the same __init__ lines. Both kept.

[nvpohanh]

@mrusanovsky could you fix lint issue?
https://github.com/sgl-project/sglang/actions/runs/34270114821/job/102209278752?pr=37462

[nvpohanh]

[by Codex] One inline review finding is attached.

[nvpohanh]

[by Codex] Severity: functional | Confidence: High

Issue

This return treats LiLiCorr sampling as supported on NPU. The selector path explicitly disables sampling there, but LiLiCorr later sets _selector_sample and _selector_sampling_accept calls accept_sampling, which unconditionally launches chain_speculative_sampling_triton. NPU tensors cannot run that kernel, so LILICORR_SAMPLING=1 with a non-greedy request fails during verification instead of taking the existing fallback.

Fix

Gate LiLiCorr sampling on the same supported-device condition as the selector, or add a supported NPU accept implementation. Add a focused NPU-path test for the fallback.

[mrusanovsky]

Fixed. The selector gates both publish sites on _selector_sampling_enabled; LiLiCorr's didn't, and the accept branch keys on _selector_sample is not None alone, so a non-greedy request on NPU reached chain_speculative_sampling_triton.

Gating only the publish would be worse than the crash: the draft would still sample while verify treated the drawn token as a point mass, making acceptance p(x) instead of min(1, p(x)/q(x)) with a residual from the wrong q : not distribution-preserving. So sampling_enabled is threaded into the sampler as done for the selector, deciding which body compiles. Unsupported devices take the argmax commit and the existing verify, with the warning that path already emits.

Fallback test added in test_dflash_logits.py, stub-based so it runs on CPU; the dispatch test now also asserts the flag reaches the builder. On CUDA the effective flag is unchanged : re-ran on an H100 and τ is identical to before: 6.8158 sampled, 6.3944 argmax, T=0 gate 7.557260716552241 on 51999 blocks.

[kpham-sgl]

Reuse DFlash2's compute_candidates (radix top-k + TP gather) and add only the lse, instead of a new top-k kernel and TP combine. It also keeps quantized lm_head support, which the eager path here breaks.

[mrusanovsky]

Done for the structure: compute_candidates body is now candidate_topk() in models/dflash.py and LiLiCorr calls it, so there is one projection (through quant_method), one TP gather, and our TP combine is gone.

One difference: with with_partition=True the top-k and lse come from a single fused pass instead of radix top-k plus a separate logsumexp. The second pass re-reads the [N, V] logits, which dominates at batch: 143 vs 680 us per step at c=32 (56 vs 69 us at c=1; H100, bf16, V=151,936, CUDA graph), and about -1.5% output throughput end to end at c=32. The selector keeps _radix_topk.

[kpham-sgl]

The -1.75pp was measured on the compiled body (see 56a1572); without compile you measured +0.73% for broadcast. Compile is gone, so pass self._attn_bias.unsqueeze(0) straight to SDPA and drop the per-step expand/reshape copy, or re-measure without compile if you want to keep it.

[mrusanovsky]

Done, passing self._attn_bias.unsqueeze(0) to SDPA. Re-measured without compile, it's neutral: -0.004% / -0.04% tok/s at c=1, +0.46% / -0.19% at c=32 (two A/B runs on one node, the second with the arms swapped), acceptance identical.

[mrusanovsky]

Addressed the latest review, replies inline.

One small addition, 17 lines in models/lilicorr.py only: LiLiCorr now serves --speculative-num-draft-tokens below the trained block, like base DFlash / DFlash2 / Domino already do. LiLiCorrDraftModel.set_block_size used to refuse a mismatch; it now has the head use a prefix of its trained slots (every relative offset a shorter lattice needs was trained). Nothing else changes, it is bit-identical at the trained size (same greedy acceptance fingerprint on gsm8k), and it serves correctly at 8 and 4 tokens.

[kpham-sgl]

/rerun-tests test/registered/core/test_basic_sanity_dflash.py test/registered/core/test_basic_sanity_dspark.py test/registered/spec/dflash/test_dflash.py test/registered/spec/test_spec_mixed_chunk.py test/registered/spec/test_gemma4_dflash_31b_extra.py test/registered/spec/dflash/test_muse_glimmer_dflash_assistant_gsm8k.py test/registered/e2e/speculative/test_dflash_domino.py test/registered/e2e/models/test_nvidia_nemotron_3_nano.py test/registered/e2e/models/test_glm53_flash_b200.py test/registered/cuda_graph/piecewise/test_pcg_with_speculative_decoding_dflash.py test/registered/kernels/ops/speculative/test_lilicorr_cuda.py test/registered/kernels/ops/speculative/test_dflash_domino.py test/registered/unit/spec/test_dflash_logits.py test/registered/unit/spec/test_dflash_overlap_hostsync.py test/registered/unit/spec/test_dflash_extra_buffer_lazy.py test/registered/unit/spec/test_dflash_domino.py test/registered/unit/spec/test_oot_dflash_hooks.py test/registered/unit/spec/test_draft_construction_isolation.py test/registered/unit/spec/test_dspark_target_hidden_projection.py test/registered/unit/models/test_glm5_next_dflash_capture.py

[github-actions]

⚠️ Rebase Required Before Re-run

A major update has landed on main. Your PR is diverged relative to required base commit 3fc7a669bfe3.

Re-run was not dispatched. What to do:

• Rebase your branch onto the latest main and push again
• Follow issue #21065 for context
• CI-fix PRs may request the bypass-maintenance label to skip this check

[kpham-sgl]

/rerun-tests test/registered/core/test_basic_sanity_dflash.py test/registered/core/test_basic_sanity_dspark.py test/registered/spec/dflash/test_dflash.py test/registered/spec/test_spec_mixed_chunk.py test/registered/spec/test_gemma4_dflash_31b_extra.py test/registered/spec/dflash/test_muse_glimmer_dflash_assistant_gsm8k.py test/registered/e2e/speculative/test_dflash_domino.py test/registered/e2e/models/test_nvidia_nemotron_3_nano.py test/registered/e2e/models/test_glm53_flash_b200.py test/registered/cuda_graph/piecewise/test_pcg_with_speculative_decoding_dflash.py test/registered/kernels/ops/speculative/test_lilicorr_cuda.py test/registered/kernels/ops/speculative/test_dflash_domino.py test/registered/unit/spec/test_dflash_logits.py test/registered/unit/spec/test_dflash_overlap_hostsync.py test/registered/unit/spec/test_dflash_extra_buffer_lazy.py test/registered/unit/spec/test_dflash_domino.py test/registered/unit/spec/test_oot_dflash_hooks.py test/registered/unit/spec/test_draft_construction_isolation.py test/registered/unit/spec/test_dspark_target_hidden_projection.py test/registered/unit/models/test_glm5_next_dflash_capture.py

[github-actions]

Results for /rerun-tests test/registered/core/test_basic_sanity_dflash.py test/registered/core/test_basic_sanity_dspark.py test/registered/spec/dflash/test_dflash.py test/registered/spec/test_spec_mixed_chunk.py test/registered/spec/test_gemma4_dflash_31b_extra.py test/registered/spec/dflash/test_muse_glimmer_dflash_assistant_gsm8k.py test/registered/e2e/speculative/test_dflash_domino.py test/registered/e2e/models/test_nvidia_nemotron_3_nano.py test/registered/e2e/models/test_glm53_flash_b200.py test/registered/cuda_graph/piecewise/test_pcg_with_speculative_decoding_dflash.py test/registered/kernels/ops/speculative/test_lilicorr_cuda.py test/registered/kernels/ops/speculative/test_dflash_domino.py test/registered/unit/spec/test_dflash_logits.py test/registered/unit/spec/test_dflash_overlap_hostsync.py test/registered/unit/spec/test_dflash_extra_buffer_lazy.py test/registered/unit/spec/test_dflash_domino.py test/registered/unit/spec/test_oot_dflash_hooks.py test/registered/unit/spec/test_draft_construction_isolation.py test/registered/unit/spec/test_dspark_target_hidden_projection.py test/registered/unit/models/test_glm5_next_dflash_capture.py:

🚀 1-gpu-5090 (4 tests): ✅ View workflow run

cd test/ && python3 registered/core/test_basic_sanity_dflash.py
cd test/ && python3 registered/spec/dflash/test_dflash.py
cd test/ && python3 registered/unit/spec/test_dflash_overlap_hostsync.py
cd test/ && python3 registered/unit/spec/test_dflash_extra_buffer_lazy.py

🚀 1-gpu-h100 (6 tests): ✅ View workflow run

cd test/ && python3 registered/core/test_basic_sanity_dspark.py
cd test/ && python3 registered/spec/test_spec_mixed_chunk.py
cd test/ && python3 registered/spec/dflash/test_muse_glimmer_dflash_assistant_gsm8k.py
cd test/ && python3 registered/cuda_graph/piecewise/test_pcg_with_speculative_decoding_dflash.py
cd test/ && python3 registered/kernels/ops/speculative/test_lilicorr_cuda.py
cd test/ && python3 registered/kernels/ops/speculative/test_dflash_domino.py

🚀 2-gpu-h100 (2 tests): ✅ View workflow run

cd test/ && python3 registered/spec/test_gemma4_dflash_31b_extra.py
cd test/ && python3 registered/e2e/speculative/test_dflash_domino.py

🚀 4-gpu-b200 (2 tests): ✅ View workflow run

cd test/ && python3 registered/e2e/models/test_nvidia_nemotron_3_nano.py
cd test/ && python3 registered/e2e/models/test_glm53_flash_b200.py

🚀 ubuntu-latest (6 tests): ❌ View workflow run

cd test/ && python3 registered/unit/spec/test_dflash_logits.py
cd test/ && python3 registered/unit/spec/test_dflash_domino.py
cd test/ && python3 registered/unit/spec/test_oot_dflash_hooks.py
cd test/ && python3 registered/unit/spec/test_draft_construction_isolation.py
cd test/ && python3 registered/unit/spec/test_dspark_target_hidden_projection.py
cd test/ && python3 registered/unit/models/test_glm5_next_dflash_capture.py

[kpham-sgl]

/rerun-tests test/registered/unit/spec/test_dflash_logits.py

[github-actions]

Results for /rerun-tests test/registered/unit/spec/test_dflash_logits.py:

🚀 ubuntu-latest (1 test): ✅ View workflow run

cd test/ && python3 registered/unit/spec/test_dflash_logits.py