[akhauriyash]

Motivation

Add native SGLang serving support for UNO, a diffusion-augmented LLM. UNO uses two model pathways:

• The diffusion pathway combines the autoregressive (AR) weights with diffusion weights implemented as LoRA adapters.
• The AR pathway uses only the base AR weights.

The diffusion pathway drafts a block of tokens in parallel, and the AR pathway verifies those proposals against the base model's autoregressive distribution.

UNO is lossless like speculative decoding; however, it does not employ a separate draft model.

Each decoding cycle performs two target-model forward passes:

1. A diffusion-pathway forward pass that produces the first output token and a block of draft proposals.
2. An AR-pathway forward pass that verifies the proposals.

The integration supports two UNO sampling modes. The diffusion pathway constructs marginal distributions for a block of B tokens:

1. Linear sampler: Samples each token independently from its marginal distribution. This mode is designed to maximize system throughput and uses B as its primary hyperparameter.
2. Tree sampler: Considers V candidate sequences by selecting the top K candidates from each marginal distribution. Its hyperparameters are (B, K, V).

Modifications

• Register UNO as a native speculative-decoding algorithm with
  --uno-lora-path.
• Validate UNO's CUDA, FA3, parallelism, LoRA, and tree configuration
  constraints at startup.
• Reuse SGLang's existing multi-LoRA manager, adapter representation, layer
  wrappers, and stacked A/B weight buffers. Add an optimized single-adapter
  cuBLAS execution backend that overlaps LoRA-A computation with the base
  GEMMs.
• Apply the UNO LoRA to all but the first row of each draft forward. Use the
  base model alone for target verification.
• Implement sparse and dense proposal and rejection sampling for linear UNO.
• Implement tree UNO sampling and reuse SGLang's existing EAGLE
  tree-verification path.
• Use sparse target sampling for tree verification when eligible; otherwise
  use EAGLE's dense sampling path.
• Add an UNO-only fused kernel for tree-suffix attention and LSE merge, with
  SGLang's existing FlashAttention and state-merge path as fallback.
• Reserve and page-align KV space for UNO's two-pass decode, and extend
  scheduler bookkeeping for its B + 1 output layout.
• Integrate UNO's carried decode state with SGLang's overlap scheduler.
• Add CUDA-graph support:
  • Linear mode captures separate graph variants for the B-wide LoRA draft
    and base-only verification passes.
  • Tree mode uses separate graph runners for the B-wide LoRA draft pass and
    V-node EAGLE cascade-attention verification pass.
• Add registered CPU dispatch and sampling tests, a CUDA kernel and CUDA-graph
  correctness test for the fused suffix merge, and an end-to-end CUDA-graph
  smoke test using a trained UNO adapter.

Changes to shared SGLang infrastructure

UNO is implemented primarily in UNO-specific modules. The following extensions
touch shared runtime interfaces but are opt-in or preserve existing defaults:

• LoRA interfaces gain optional support for overlapped LoRA-A execution,
  explicit token-row routing, and skipping inactive internal adapters. Only
  the UNO backend enables these capabilities.
• The shared FlashAttention backend uses the fused suffix-attention merge only
  for eligible UNO tree verification. All other cases retain the existing
  cascade-attention path.
• EAGLE verification accepts optional UNO sampling metadata to select sparse
  target sampling; calls without it retain the existing sampling behavior.
• GenerationBatchResult gains optional output-stride and non-draft-token-count
  fields. Scheduler bookkeeping uses them to slice UNO's B + 1 output rows
  and count accepted proposals; existing algorithms retain their previous
  defaults.
• Speculative allocation sizing uses an UNO-specific branch to reserve the
  additional KV position required by its two-pass decode.
• Model-runner and speculative-registry hooks select the UNO worker and
  CUDA-graph runner only when --speculative-algorithm UNO is active.

Validation

H200 unit and CUDA smoke validation

The CPU regression tests and the registered H200 CUDA smoke test passed. The
smoke test captured and replayed both linear and tree CUDA-graph paths using
three fixed GSM8K prompts with 128 generated tokens per prompt.

Mode	B	K	V	TPF
Linear	8	1	8	2.462
Tree	8	16	8	2.783

TPF is useful output tokens divided by both target-model forwards per UNO
cycle: tokens / (2 * spec_verify_ct). Tree TPF exceeded linear TPF.

H200 math evaluation

The runs used Qwen/Qwen3-8B in BF16 on one NVIDIA H200 with FA3, a 40,960-token context, and the published UNO Qwen3-8B adapter for UNO modes. All modes used temperature 1, top-k 50, top-p 0.95, random seed 42, and up to 32,768 output tokens. GSM8K and MATH500 used one sample per problem; AIME25 used 10 samples per problem (300 completions). Each dataset was run separately so it has an independent throughput measurement.

Mode	Dataset	Requests	Accuracy	TPF	tok/s	tok/s/request
AR	GSM8K	64	96.10%	1.000	6095.61	95.24
Linear 8/1/8	GSM8K	64	96.00%	2.590	7197.80	112.47
AR	MATH500	64	96.00%	1.000	3148.58	49.20
Linear 8/1/8	MATH500	64	95.60%	2.720	5534.55	86.48
AR	AIME25	64	69.70%	1.000	1894.82	29.61
Linear 8/1/8	AIME25	64	68.00%	2.730	3841.41	60.02
AR	GSM8K	1	95.90%	1.000	190.09	190.09
Tree 16/32/32	GSM8K	1	96.00%	3.130	470.47	470.47
AR	MATH500	1	96.00%	1.000	183.12	183.12
Tree 16/32/32	MATH500	1	96.80%	3.450	508.56	508.56
AR	AIME25	1	68.00%	1.000	178.32	178.32
Tree 16/32/32	AIME25	1	67.30%	3.430	488.67	488.67

For UNO modes, the three numbers are B/K/V: draft width, top-k candidates kept per position,
and verification width. Thus, Linear 8/1/8 uses an
8-token draft width, one candidate per position, and an 8-token verification
width. tok/s/request is aggregate throughput divided by the number of
concurrent requests.

Reproduction

The table was generated with benchmark/uno/run_math_eval.py, which
instantiates sglang.Engine in-process and excludes engine startup from the
timed interval. Exact commands for reproducing every row are provided in
benchmark/uno/README.md.

The equivalent server launch commands for the two UNO configurations are:

export MODEL_PATH=Qwen/Qwen3-8B
export UNO_LORA_PATH=s-sahoo/uno-qwen3-8B

# Linear UNO: B/K/V = 8/1/8, 64 concurrent requests
python3 -m sglang.launch_server \
  --model-path "$MODEL_PATH" \
  --dtype bfloat16 \
  --attention-backend fa3 \
  --context-length 40960 \
  --max-running-requests 64 \
  --random-seed 42 \
  --speculative-algorithm UNO \
  --uno-lora-path "$UNO_LORA_PATH" \
  --speculative-num-steps 1 \
  --speculative-eagle-topk 1 \
  --speculative-num-draft-tokens 8

# Tree UNO: B/K/V = 16/32/32, one concurrent request
python3 -m sglang.launch_server \
  --model-path "$MODEL_PATH" \
  --dtype bfloat16 \
  --attention-backend fa3 \
  --context-length 40960 \
  --max-running-requests 1 \
  --random-seed 42 \
  --speculative-algorithm UNO \
  --uno-lora-path "$UNO_LORA_PATH" \
  --speculative-num-steps 15 \
  --speculative-eagle-topk 32 \
  --speculative-num-draft-tokens 32

Test coverage

Added regression coverage for:

• Invalid tree dimensions exceeding the current EAGLE parent-list
  representation.
• Clearing stale internal LoRA metadata before base-only verification.
• Sparse tree target-sampling dispatch and compact-support sampling.
• UNO-only fused suffix-attention dispatch with fallback for unsupported
  attention features and layouts.
• Numerical correctness and CUDA-graph replay of the fused suffix-attention
  merge kernel.
• Linear and tree CUDA-graph capture, replay, generation, and nontrivial
  acceptance.

Current test files:

• test/registered/unit/spec/test_uno_tree_config.py
• test/registered/unit/lora/test_uno_inactive_lora_batch.py
• test/registered/unit/spec/test_uno_tree_sparse_sampling.py
• test/registered/unit/spec/test_suffix_attention_merge_dispatch.py
• test/registered/kernels/ops/attention/test_suffix_attention_merge.py
• test/registered/spec/uno/test_uno.py

Checklist

• Follow the SGLang code-style guidance.
• Add unit tests for the new dispatch, sampling, kernel, and end-to-end
  paths.
• Use the published s-sahoo/uno-qwen3-8B adapter.

---

CI States

Latest PR Test (Base): ⏳ Run #33738458871
Latest PR Test (Extra): ⏳ Run #33738458816
Latest PR Test (AMD ROCm 7.2): ⏳ Run #33738458864

[BBuf]

This makes the UNO reserve 2 * (V + 1), but with page_size == 1 get_req_to_token_extra_context_len() still gives each row only V + 4 extra slots. Near the context limit, UnoDraftInput.prepare_for_decode() can hit its max(next_lengths) > row_width error (for example with V=8). Could we size the UNO row headroom from get_alloc_reserve_per_decode() too, as the page-size > 1 path already does?

[drproduck]

I restored reserve-sized headroom for UNO at page size 1 while leaving other speculative algorithms unchanged. For V=8, the row now provides 18 slots, matching 2 * (V + 1). I also added a CPU regression test for this invariant.

This is the same page-size-1 row-sizing issue tracked in #33579; the general DFLASH/DSPARK fix is pending in #33581.

[BBuf]

This renormalizes over top-k before applying top-p. The regular SGLang sampler uses joint filtering (filter_apply_order="joint"), where top-p is computed from the full softmax distribution. The dense UNO path is top-k-first as well, so with both top_k and top_p (including the 50/0.95 benchmark setup), UNO's target distribution differs from normal AR despite the lossless claim. Could we align both UNO target paths with the regular sampler and add a small distribution/support parity test?

[drproduck]

I wanted to note that both dflash and eagle3 do this:

dflash
eagle3

Since applying top-k first avoids materializing full-vocabulary softmax probabilities, I intentionally followed the existing speculative-decoding convention. Changing UNO alone would change the computational cost relative to competing specdecs, affecting apples-to-apples comparisons.

Would you prefer that UNO retain the existing convention in this PR, or align with AR while dflash and eagle remain different?

[BBuf]

num_correct_drafts = accept_len - 2 is right for UNO proposal acceptance, but the downstream token counters still add only batch_size + num_correct_drafts (update_spec_metrics, report_decode_stats, and _record_step_counters). That counts accept_len - 1, so every UNO verify under-reports one generated token per request. Could we carry the non-draft-token count through those counters too?

[drproduck]

I centralized the generated-token count as num_correct_drafts + batch_size * num_non_draft_tokens_per_req.

It is now used consistently for aggregate acceptance length, logged throughput, realtime/MFU metrics, and scheduler decode moments. Correct drafts are tracked separately so UNO's second non-draft token is not incorrectly included in the accept-rate numerator. Existing speculative algorithms retain their default of one non-draft token per request.

I added CPU coverage for this.

[BBuf]

Since this is the only adapter-backed linear/tree CUDA-graph test, can we enable it before the PR leaves draft? Right now the checkpoint is not public and the assertions only cover non-empty output/TPF, so public CI cannot catch a lossless-correctness regression. A fixed greedy UNO-vs-AR prefix check for both modes would cover that.

[drproduck]

The AR-vs-UNO parity test was pushed right after your comment.

The CUDA test now verifies that the first 32 greedily generated tokens from both linear and tree UNO exactly match AR.

The same test exercises CUDA graphs in both modes and reports linear TPF 2.526 and tree TPF 3.048. The test remains disabled in CI until we publish the adapter; we will enable it before the PR leaves draft.

[HF-001]

@akhauriyash @drproduck hi, Uno is a great feature. How does it perform on large-sized models? Do you have experimental data? Will you continue to optimize uno on the large-sized and high TP models in the future?