Optimize LenVM guided sampling latency

[namezhenzhang]

Summary

Follow-up to the LenVM timing analysis in PR #2. This PR optimizes the in-process LenVM guided sampling hot path and adds a written latency summary in docs/lenvm-guided-sampling-optimization.md.

Main changes:

• Skip in-process LenVM initialization for batches without active value guidance, and skip neutral guidance settings such as centered_exp scale 0.
• Keep compacted candidate ids/probs/masks on GPU for expectation-style guidance, avoiding per-token CPU candidate materialization in the common path.
• Apply guidance in place and only return changed rows, preserving sampling filters for unmodified rows without cloning the full [batch, vocab] probability tensor.
• Add a fused LenVM prefix-extend + candidate-scoring path for tiny decode deltas.
• Avoid .tolist() GPU syncs in Qwen2/Qwen3/Qwen2.5-VL LenVM value slicing by carrying prefix/candidate metadata in the tree-value spec.
• Cache request EOS ids and add opt-in timing logs through SGLANG_LVM_TIMING.
• Add focused tests for the fast path, no-op skips, baseline precheck, mixed GPU guidance modes, and tree-value mask/prefill args.

Benchmark Summary

Reference config: 1x H100 SXM, Qwen/Qwen2.5-7B-Instruct + namezz/lvm-math-0402-a-qwen2.5-7b-instruct-b-qwen2.5-1.5b-instruct, GSM8K 50 questions x 16 samples, max_tokens=6000, T=1.0, top-p=1.0, min-p=0.01. Baseline uses top_k=-1; LenVM uses top_k=5, value_mode=centered_exp, value_scale=0.001, gamma=0.997.

run	wall clock	avg completion tokens / choice	acc_first	acc_any
baseline	19.84 s	295.96	0.88	1.00
optimized LenVM	80.02 s	298.36	0.90	0.96
ratio	4.03x slower	+0.81%	-	-

Compared with the PR #2 reference result, 19.22 s -> 87.44 s (4.55x slower), this reduces the guided run wall clock by 8.5%, the slowdown ratio by 11.4%, and the incremental LenVM overhead by 11.8%.

Model memory from the same run: base 7B bf16 weights 14.30 GB; LenVM 1.5B bf16 weights 3.03 GB.

Profiling (SGLANG_LVM_TIMING=1, q10/n4/topk5/scale=0.001) shows baseline Sampler.forward at about 0.70-0.77 ms. With LenVM active, Sampler.forward is about 13-15 ms after warmup and LvmGuidedSampler.apply accounts for about 97% of it. The remaining bottleneck is still rows that fall back to the two-forward LenVM path: prefix extend plus candidate launch.

Artifacts from the Slurm runs:

• results/timing/pr7b_q50n16_s001_84065_20260523_211347/
• results/timing/pr7b_prof_s001_84067_20260523_212103/

Tests

• .venv-infer/bin/python -m compileall -q ... on modified SGLang files and focused tests
• git diff --check
• Direct focused test invocation through .venv-infer/bin/python because this workspace does not currently have pytest installed:
  • test_lvm_guided_sampling_fast_path.py: 7 tests passed
  • test_tree_value_spec.py: 2 tests passed

Follow-up

The fused path is still often all-or-nothing at the batch level. The next high-value optimization is to split a batch into fusible and fallback rows, then run fused extend+candidate scoring for eligible rows and keep the two-phase path only for the rest.