[Score API] Setwise scoring: CausalLM support (batched + --enable-mis)

[sundar24295s]

Setwise Scoring: CausalLM Support

Follow-up to [Score API] Setwise Scoring Support (#38965) (merged). That PR
added setwise scoring for SequenceClassification models; this PR extends the same
token_indices_to_pool primitive to generative (CausalLM) models.

Summary

Extends setwise scoring to CausalLM models. A generative reranker can now
read label-token logprobs at every occurrence of a caller-chosen extraction
token in each query + item sequence, instead of only at the last token — the
generative analogue of the SequenceClassification setwise path. A candidate set is
packed into one item with one extraction token per candidate; the LM head is read at
each of those positions in a single prefill, so scores comes back as one
[N_i x num_labels] matrix per item (nested per item), where num_labels = len(label_token_ids).

The unifying signal is unchanged: token_indices_to_pool still names the readout
positions. For SequenceClassification the pooler pools the classification head there;
for CausalLM the LM head is applied there and the requested label_token_ids
logprobs are gathered — one row per anchor.

Modes: both batched (default) and fused --enable-mis are supported,
matching the SequenceClassification setwise paths.

Motivation

Many rerankers are decoder-only models that express a candidate's score as the
probability the model assigns to specific label tokens (e.g. yes/no, a rating
digit) right after each candidate. Without setwise support such a caller must issue
one request per candidate, paying the shared query-prefix cost N times. This PR packs
the whole candidate block into one prompt with one anchor per candidate and returns
all N per-candidate label distributions from a single prefill.

API

Engine (in-process)

result = engine.score(
    query="",                        # optional shared prefix
    items=[prompt_with_anchors],     # one item = one candidate set
    label_token_ids=[yes_id, no_id], # required for CausalLM
    score_extraction_token_id=anchor_id,
    apply_softmax=False,
)
# result.scores -> [num_items][N_i x num_labels]   (nested per item)
#   each row = per-anchor label-token scores: softmax over labels when
#   apply_softmax=True, else exp(logprob) probabilities (the existing pointwise
#   CausalLM conversion, applied per anchor).

HTTP (POST /v1/score)

{
  "query": "",
  "items": ["Rank the candidates. C0. C1. C2. Scores:<|object_ref_start|><|object_ref_start|><|object_ref_start|>"],
  "label_token_ids": [9454, 2753],
  "score_extraction_token": "<|object_ref_start|>",
  "apply_softmax": false
}

No new request/response fields: this reuses score_extraction_token /
score_extraction_token_id and the existing nested ScoringResponse.scores union.
label_token_ids is required for CausalLM (as for pointwise CausalLM scoring).

Semantics & validation

• Both execution modes. Batched scores each item as an independent query + item
  sequence; fused --enable-mis packs the items into one sequence with a
  block-diagonal mask that isolates each set, and splits the flat result back per
  item. Both return nested [num_items][Nᵢ x num_labels].
• The SequenceClassification is_score_and_pool_model allow-list is unchanged and
  still governs the non-generation path; generation skips that check.
• All other setwise constraints are shared: every item must contain at least one
  extraction token, the shared query prefix must not contain it, and radix cache /
  chunked prefill must be off with no --allow-auto-truncate (pooling positions are
  full-prompt coordinates). All raise a clean 400 before inference.
• return_pooled_hidden_states stays rejected for CausalLM (no task head).

How it works

token_indices_to_pool already flows to every ForwardBatch. For SequenceClassification
the pooler consumes it; for CausalLM there is no task head, so the LM head is applied
at those positions instead — mirroring the existing multi-item (--enable-mis)
delimiter path, but pooling AT the anchor (no delimiter - 1 shift, no discarded
row). The forward dispatch prioritizes the anchor readout over the delimiter path, so
it serves both modes; the block-diagonal MIS mask is applied independently by the
attention backend (gated on enable_mis), so a batched request never triggers it
while a fused one does.

Implementation

Area	File	Change
LM-head readout	layers/logits_processor.py	compute_logprobs_at_positions() — reads the LM head AT token_indices_to_pool and gathers label_token_ids logprobs (one row per anchor); a forward() branch dispatches to it for prefill-only requests carrying the field, taking precedence over the MIS delimiter path so it serves both modes.
Logprob bookkeeping	managers/scheduler_components/logprob_result_processor.py	Generalized "position-based scoring" (_scoring_positions): setwise anchors take precedence over MIS delimiters (a fused setwise request carries both), so the expected input-logprob count is len(token_indices_to_pool).
Request plumbing	managers/io_struct.py	token_indices_to_pool on GenerateReqInput (per-item split in __getitem__) and TokenizedGenerateReqInput.
	managers/tokenizer_manager.py, managers/scheduler.py	Thread the field into the tokenized generate request and the Req.
Score engine	managers/tokenizer_manager_score_mixin.py	Validation allows CausalLM setwise (batched and --enable-mis); score_request sets token_indices_to_pool + logprob_start_len=0; the batched result branch reads per-request input_token_ids_logprobs, and the fused _process_multi_item_extraction_results gains a generation branch that reads the flat input_token_ids_logprobs and splits per item via per_item_anchor_counts. Both reuse _convert_logprobs_to_scores (per-item labels + temperature).
Docs	entrypoints/openai/protocol.py, entrypoints/engine_score_mixin.py	score_extraction_token(_id) docs note CausalLM support.
Tests	test/registered/unit/managers/test_setwise_score_mixin.py, test/registered/e2e/scoring/test_setwise_scoring.py	CPU result-grouping + validation tests and a CausalLM HF-parity / MIS set-isolation E2E suite.

Tests

# CPU unit tests (no GPU / model needed)
python test/registered/unit/managers/test_setwise_score_mixin.py -v
python test/registered/unit/layers/test_pooler_score_and_pool.py -v

# E2E (GPU); CausalLM setwise uses a decoder-only checkpoint
export TEST_CAUSAL_LM_MODEL=<path-to>/Qwen3-0.6B
python test/registered/e2e/scoring/test_setwise_scoring.py -v

• Unit (test_setwise_score_mixin.py): a _GenHarness (is_generation=True)
  plus a TestGenerationSetwiseResults suite covering nested grouping (batched and
  fused/MIS) from input_token_ids_logprobs, softmax on/off, missing-label → 0.0,
  empty-logprobs / row-count-mismatch errors, and a pointwise-unchanged regression;
  validation tests assert CausalLM setwise is allowed.
• E2E (test_setwise_scoring.py): TestGenerationSetwiseScoringHFParity asserts
  numerical parity against a HuggingFace reference (gather log_softmax(LM-head logits) at each anchor, exponentiate for apply_softmax=False), single- and
  multi-item. TestGenerationSetwiseMISScoring asserts block-diagonal set
  isolation — the property only a real fused forward exercises. Shape/grouping/error
  cases live in the CPU unit tests, not duplicated in the GPU suite.

Validation

Environment: H100 80GB, flashinfer 0.6.18, sglang-kernel 0.4.7, Qwen3-0.6B
(CausalLM base, anchor <|object_ref_start|>), engine float16 on flashinfer.

Suite	Cases	Type	Result
test_setwise_score_mixin.py	50	CPU unit	PASS
test_pooler_score_and_pool.py	12	CPU unit	PASS
test_token_scoring.py	10	CPU unit	PASS
test_setwise_scoring.py::TestGenerationSetwiseScoringHFParity	2	E2E (HF parity, single + multi-item)	PASS
test_setwise_scoring.py::TestGenerationSetwiseMISScoring	1	E2E (fused set isolation)	PASS
lint: ruff / ruff-format / isort / codespell / registered-tests	—	lint	PASS

HTTP smoke test (launch_server + curl /v1/score)

• Setwise, 1 item × 3 anchors, apply_softmax=false → nested [1][3 x 2] probability
  matrix.
• Setwise, 1 item × 2 anchors, apply_softmax=true → nested [1][2 x 2], each row
  sums to 1.0.
• Pointwise (no score_extraction_token), 2 items → flat [2 x 2] — backward
  compatibility intact.

Backward compatibility

Purely additive. Pointwise CausalLM scoring is unchanged (score_extraction_token_id
unset → last-token output_token_ids_logprobs path, flat [num_rows x num_labels]).
SequenceClassification setwise behavior is untouched — the only shared-path change is
the logprob-count bookkeeping, which now recognizes token_indices_to_pool in
addition to MIS delimiters.
}

---

CI States

Latest PR Test (Base): ❌ Run #36188053776
Latest PR Test (Extra): ⚠️ Not enabled -- add run-ci-extra label to opt in.
Latest PR Test (AMD ROCm 10): ⏳ Run #36188053694

[sundar24295s]

/rerun-failed-ci

[sundar24295s]

/rerun-failed-ci

[Qiaolin-Yu]

failed cis are unrelated