[Sampling] Add selected/support sampling logprob modes

[nanjiangwill]

Summary

• Add optional sampling_logprobs_mode: Literal["selected", "support"].
• Preserve the existing default: output_token_sampling_logprobs[t] is the selected output token behavior logprob.
• In "support" mode, return output_token_sampling_logprobs[t] as list[float], aligned elementwise with output_token_sampling_mask[t].
• Preserve the existing singular response key and length key.
• Carry the request-selected representation through raw generation, chat completions, sessions, pipeline parallelism, streaming, multi-tokenizer serving, and P/D disaggregation.

API contract

client.generate(
    ...,
    return_sampling_mask=True,
    sampling_logprobs_mode="support",
)

For output position t, let w_t be the exact post-filter weights captured from the sampler, S_t = {a | w_t(a) > 0}, and

q_t(a) = w_t(a) / sum_{v in S_t} w_t(v).

With return_sampling_mask=false, sampling_logprobs_mode must be omitted.

With return_sampling_mask=true:

• Omitting sampling_logprobs_mode or setting it to "selected" returns log q_t(output_ids[t]) as a float. This keeps the existing default response shape.
• sampling_logprobs_mode="support" returns one log q_t(a) per token ID in output_token_sampling_mask[t] as a list. The token-ID and logprob rows have identical lengths and elementwise alignment.
• Both modes describe the same normalized behavior distribution. The mode changes how much of that distribution is returned, not how it is normalized.
• Greedy sampling returns a singleton support; selected mode returns 0.0, while support mode returns [0.0].
• Any explicit sampling_logprobs_mode requires return_sampling_mask=true; when masks are enabled, omission resolves to selected mode at the internal request boundary.

The support row order is not an API guarantee; elementwise alignment is. Overflow and invalid captured support continue to fail closed, so no partial distribution is returned.

Implementation

• Reuse the exact post-filter weights already captured by the sampling-mask path; there is no second top-k/top-p reconstruction.
• Preserve the existing producer-captured selected-token path. If token synchronization or grammar changes the selected token, keep the existing fallback that recomputes its logprob from the captured weights.
• Build the full support-logprob tensor only when at least one request selects support mode.
• In mixed batches, index-select only support-mode rows. Selected-mode neighbors do not allocate or copy full support-logprob rows.
• Append the mode as a defaulted tail field in the Python scheduler wire, preserving the existing Rust positional prefix and compatibility with shorter senders.
• Widen the opt-in P/D sampling-logprob metadata row to the configured sampling-mask capacity so support-mode handoff tokens remain aligned.

Performance boundary

When sampling-mask output is disabled, this adds no sampler tensor work or payload. With masks enabled:

• Selected mode retains the existing compact result shape: approximately 4RC + 12R bytes for R mask rows and packed capacity C.
• Support mode adds approximately 4R_sC bytes, where R_s is only the number of support-mode rows.
• P/D sampling-mask metadata remains behind SGLANG_ENABLE_DISAGG_SAMPLING_MASK; support mode necessarily makes the aligned float row capacity-sized, matching the token-ID row.

A packed/base64 ragged response format is intentionally not part of this change. It is an independent transport optimization and can be added without changing these selected/support semantics.

Validation

• 314 focused CPU tests passed, 2 skipped, with 142 subtests after the review changes.
• Coverage includes the public option validation matrix, default selected behavior, explicit support behavior, mixed selected/support batches, support-only row packing, pipeline-parallel transport, multi-tokenizer splitting, P/D handoff, overflow, invalid sampled support, streaming, and custom hard exclusions.
• All pre-commit hooks passed for the changed files.
• The registered CUDA sampling-mask suite passed on 2× H100 80 GB GPUs (CUDA 13.0, PyTorch 2.13.0): 34 tests and 28 subtests in 437.53 seconds. Modal run

---

CI States

Latest PR Test (Base): ⏳ Run #35951430674
Latest PR Test (Extra): ❌ Run #35951430010
Latest PR Test (AMD ROCm 10): ❌ Run #35951430457

[ByronHsu]

/tag-and-rerun-ci

[ByronHsu]

/rerun-test test/registered/npu/basic_function/dllm/test_npu_llada2_mini.py

[github-actions]

Results for /rerun-test test/registered/npu/basic_function/dllm/test_npu_llada2_mini.py:

⛔ test/registered/npu/basic_function/dllm/test_npu_llada2_mini.py: test/registered/npu/basic_function/dllm/test_npu_llada2_mini.py is registered for NPU (suite base-b-test-4-npu-a3), not for CUDA or CPU; rerun-test.yml has no NPU job. Rerun it with /rerun-failed-ci, or dispatch the NPU workflow manually.

[ByronHsu]

/rerun-failed-ci