997 lines
41 KiB
Python
997 lines
41 KiB
Python
# Copyright 2023-2024 SGLang Team
|
|
# Licensed under the Apache License, Version 2.0 (the "License");
|
|
# you may not use this file except in compliance with the License.
|
|
# You may obtain a copy of the License at
|
|
#
|
|
# http://www.apache.org/licenses/LICENSE-2.0
|
|
#
|
|
# Unless required by applicable law or agreed to in writing, software
|
|
# distributed under the License is distributed on an "AS IS" BASIS,
|
|
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
|
# See the License for the specific language governing permissions and
|
|
# limitations under the License.
|
|
# ==============================================================================
|
|
"""Logits processing."""
|
|
|
|
import dataclasses
|
|
import logging
|
|
from contextlib import contextmanager
|
|
from typing import Any, Dict, List, Optional, Tuple, Union
|
|
|
|
import torch
|
|
from torch import nn
|
|
|
|
from sglang.kernels.ops.activation.softcap import (
|
|
softcap_inplace_logits as fused_softcap,
|
|
)
|
|
from sglang.srt.distributed.device_communicators import triton_symm_mem_ag
|
|
from sglang.srt.layers.aux_hidden_states import (
|
|
AuxHiddenStates,
|
|
pack_aux_hidden_states,
|
|
)
|
|
from sglang.srt.layers.dp_attention import (
|
|
DpPaddingMode,
|
|
attn_tp_all_gather,
|
|
attn_tp_all_gather_into_tensor,
|
|
dp_gather_replicate,
|
|
dp_scatter,
|
|
get_dp_device,
|
|
get_dp_dtype,
|
|
get_dp_hidden_size,
|
|
)
|
|
from sglang.srt.layers.logprob_processor import (
|
|
InputLogprobProcessor,
|
|
LogprobStage,
|
|
get_token_ids_logprobs_raw,
|
|
get_top_logprobs_raw,
|
|
)
|
|
from sglang.srt.layers.vocab_parallel_embedding import VocabParallelEmbedding
|
|
from sglang.srt.model_executor.forward_batch_info import (
|
|
CaptureHiddenMode,
|
|
ForwardBatch,
|
|
ForwardMode,
|
|
)
|
|
from sglang.srt.runtime_context import get_exec, get_parallel
|
|
from sglang.srt.utils.common import (
|
|
is_cpu,
|
|
is_npu,
|
|
is_pin_memory_available,
|
|
use_intel_amx_backend,
|
|
)
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
_is_npu = is_npu()
|
|
_is_cpu = is_cpu()
|
|
|
|
_UNQUANTIZED_LM_HEAD_METHODS = {
|
|
"UnquantizedEmbeddingMethod",
|
|
"UnquantizedLinearMethod",
|
|
"PackWeightMethod",
|
|
}
|
|
|
|
# None outside a FlashInfer autotune pass; inside one, whether that pass runs the
|
|
# LM head. Not-None means the forward's output is discarded -- attention backends
|
|
# read that via get_in_autotune_dummy_run() to skip a cross-node exchange.
|
|
# Skipping the LM head skips its [batch * dp_size, vocab] all-gather, which OOMs
|
|
# under DP attention with a tight mem_fraction_static.
|
|
_autotune_run_lm_head: Optional[bool] = None
|
|
|
|
|
|
def get_in_autotune_dummy_run() -> bool:
|
|
return _autotune_run_lm_head is not None
|
|
|
|
|
|
@contextmanager
|
|
def autotune_dummy_run_mode(*, run_lm_head: bool):
|
|
global _autotune_run_lm_head
|
|
_autotune_run_lm_head = run_lm_head
|
|
try:
|
|
yield
|
|
finally:
|
|
_autotune_run_lm_head = None
|
|
|
|
|
|
@dataclasses.dataclass
|
|
class LogitsProcessorOutput:
|
|
## Part 1: This part will be assigned in python/sglang/srt/layers/logits_processor.py::LogitsProcessor
|
|
# The logits of the next tokens. shape: [#seq, vocab_size]
|
|
# Can be None for certain prefill-only requests (e.g., multi-item scoring) that don't need next token generation
|
|
next_token_logits: Optional[torch.Tensor]
|
|
# Used by speculative decoding (EAGLE)
|
|
# The last hidden layers
|
|
hidden_states: Optional[torch.Tensor] = None
|
|
|
|
## Part 2: This part will be assigned in python/sglang/srt/layers/sampler.py::Sampler
|
|
# he log probs of output tokens, if SGLANG_RETURN_ORIGINAL_LOGPROB = True, will get the log probs before applying temperature. If False, will get the log probs before applying temperature.
|
|
next_token_logprobs: Optional[torch.Tensor] = None
|
|
# The logprobs and ids of the top-k tokens in output positions. shape: [#seq, k]
|
|
next_token_top_logprobs_val: Optional[List] = None
|
|
next_token_top_logprobs_idx: Optional[List] = None
|
|
# The logprobs and ids of the requested token ids in output positions. shape: [#seq, n] (n is the number of requested token ids)
|
|
# Can contain either lists or GPU tensors (for delayed copy optimization in prefill-only requests)
|
|
next_token_token_ids_logprobs_val: Optional[
|
|
List[Union[List[float], torch.Tensor]]
|
|
] = None
|
|
next_token_token_ids_logprobs_idx: Optional[List] = None
|
|
# Sparse top-k/top-p/min-p support ids and selected-token logprob after
|
|
# truncation/renormalization. Only populated when requested.
|
|
next_token_sampling_mask_idx: Optional[List[Optional[List[int]]]] = None
|
|
next_token_sampling_logprobs: Optional[List[Optional[float]]] = None
|
|
|
|
## Part 3: Prefill-only. This part will be assigned in python/sglang/srt/layers/logits_processor.py::LogitsProcessor
|
|
# The logprobs of input tokens. shape: [#token]
|
|
input_token_logprobs: Optional[torch.Tensor] = None
|
|
# The logprobs and ids of the top-k tokens in input positions. shape: [#seq, #token, k]
|
|
input_top_logprobs_val: Optional[List] = None
|
|
input_top_logprobs_idx: Optional[List] = None
|
|
# The logprobs and ids of the requested token ids in input positions. shape: [#seq, n] (n is the number of requested token ids)
|
|
# Can contain either lists or GPU tensors (for delayed GPU-to-CPU transfer optimization)
|
|
input_token_ids_logprobs_val: Optional[List[Union[List[float], torch.Tensor]]] = (
|
|
None
|
|
)
|
|
input_token_ids_logprobs_idx: Optional[List] = None
|
|
|
|
## Part 4: Diffusion LLM only.
|
|
full_logits: Optional[torch.Tensor] = None
|
|
|
|
## Part 5: Customized Info
|
|
customized_info: Optional[Dict[str, List[Any]]] = None
|
|
|
|
## Part 6: Temporary variables
|
|
# FIXME: These fields are not logits-related but are passed through here as a
|
|
# workaround since ForwardBatch is local to forward_batch_generation().
|
|
# They should be moved to GenerationBatchResult to keep this class clean.
|
|
mm_input_embeds: Optional[torch.Tensor] = None
|
|
|
|
|
|
@dataclasses.dataclass
|
|
class LogitsMetadata:
|
|
forward_mode: ForwardMode
|
|
capture_hidden_mode: CaptureHiddenMode = CaptureHiddenMode.NULL
|
|
next_token_logits_buffer: Optional[torch.Tensor] = None
|
|
|
|
extend_return_logprob: bool = False
|
|
extend_return_top_logprob: bool = False
|
|
extend_token_ids_logprob: bool = False
|
|
extend_seq_lens: Optional[torch.Tensor] = None
|
|
extend_seq_lens_cpu: Optional[List[int]] = None
|
|
extend_logprob_start_lens_cpu: Optional[List[int]] = None
|
|
extend_logprob_pruned_lens_cpu: Optional[List[int]] = None
|
|
top_logprobs_nums: Optional[List[int]] = None
|
|
extend_input_logprob_token_ids_gpu: Optional[torch.Tensor] = None
|
|
token_ids_logprobs: Optional[List[List[int]]] = None
|
|
|
|
# logits and logprobs post processing
|
|
temperature: torch.Tensor = None
|
|
top_p: torch.Tensor = None
|
|
|
|
# DP attention metadata. Not needed when DP attention is not used.
|
|
# Number of tokens in the request.
|
|
global_num_tokens_gpu: Optional[torch.Tensor] = None
|
|
# The start position of local hidden states.
|
|
dp_local_start_pos: Optional[torch.Tensor] = None
|
|
dp_local_num_tokens: Optional[torch.Tensor] = None
|
|
global_dp_buffer_len: Optional[int] = None
|
|
# Number of tokens to sample per DP rank
|
|
global_num_tokens_for_logprob_cpu: Optional[torch.Tensor] = None
|
|
global_num_tokens_for_logprob_gpu: Optional[torch.Tensor] = None
|
|
# The gather mode for DP attention
|
|
dp_padding_mode: Optional[DpPaddingMode] = None
|
|
|
|
# Whether this batch is prefill-only (no token generation needed)
|
|
is_prefill_only: bool = False
|
|
|
|
mm_input_embeds: Optional[torch.Tensor] = None
|
|
|
|
# DRAFT_EXTEND_V2: when set, lm_head runs only on these rows (see
|
|
# EagleDraftExtendInput.select_index).
|
|
draft_extend_select_index: Optional[torch.Tensor] = None
|
|
|
|
@classmethod
|
|
def from_forward_batch(cls, forward_batch: ForwardBatch):
|
|
if (
|
|
forward_batch.forward_mode.is_extend()
|
|
and forward_batch.return_logprob
|
|
and not forward_batch.forward_mode.is_target_verify()
|
|
):
|
|
extend_return_top_logprob = any(
|
|
x > 0 for x in forward_batch.top_logprobs_nums
|
|
)
|
|
extend_token_ids_logprob = any(
|
|
x is not None for x in forward_batch.token_ids_logprobs
|
|
)
|
|
extend_return_logprob = False
|
|
extend_logprob_pruned_lens_cpu = []
|
|
for extend_len, start_len in zip(
|
|
forward_batch.extend_seq_lens_cpu,
|
|
forward_batch.extend_logprob_start_lens_cpu,
|
|
):
|
|
if extend_len - start_len > 0:
|
|
extend_return_logprob = True
|
|
extend_logprob_pruned_lens_cpu.append(extend_len - start_len)
|
|
else:
|
|
extend_return_logprob = extend_return_top_logprob = (
|
|
extend_token_ids_logprob
|
|
) = extend_logprob_pruned_lens_cpu = False
|
|
|
|
if forward_batch.forward_mode.is_draft_extend_v2():
|
|
draft_extend_select_index = forward_batch.spec_info.select_index
|
|
else:
|
|
draft_extend_select_index = None
|
|
|
|
return cls(
|
|
forward_mode=forward_batch.forward_mode,
|
|
capture_hidden_mode=forward_batch.capture_hidden_mode,
|
|
next_token_logits_buffer=forward_batch.next_token_logits_buffer,
|
|
extend_return_logprob=extend_return_logprob,
|
|
extend_return_top_logprob=extend_return_top_logprob,
|
|
extend_token_ids_logprob=extend_token_ids_logprob,
|
|
extend_seq_lens=forward_batch.extend_seq_lens,
|
|
extend_seq_lens_cpu=forward_batch.extend_seq_lens_cpu,
|
|
extend_logprob_start_lens_cpu=forward_batch.extend_logprob_start_lens_cpu,
|
|
extend_logprob_pruned_lens_cpu=extend_logprob_pruned_lens_cpu,
|
|
top_logprobs_nums=forward_batch.top_logprobs_nums,
|
|
token_ids_logprobs=forward_batch.token_ids_logprobs,
|
|
extend_input_logprob_token_ids_gpu=forward_batch.extend_input_logprob_token_ids_gpu,
|
|
is_prefill_only=forward_batch.is_prefill_only,
|
|
global_num_tokens_gpu=forward_batch.global_num_tokens_gpu,
|
|
dp_local_start_pos=forward_batch.dp_local_start_pos,
|
|
dp_local_num_tokens=forward_batch.dp_local_num_tokens,
|
|
global_dp_buffer_len=forward_batch.global_dp_buffer_len,
|
|
global_num_tokens_for_logprob_cpu=forward_batch.global_num_tokens_for_logprob_cpu,
|
|
global_num_tokens_for_logprob_gpu=forward_batch.global_num_tokens_for_logprob_gpu,
|
|
dp_padding_mode=DpPaddingMode.SUM_LEN,
|
|
mm_input_embeds=forward_batch.mm_input_embeds,
|
|
draft_extend_select_index=draft_extend_select_index,
|
|
)
|
|
|
|
def compute_dp_attention_metadata(self):
|
|
cumtokens = torch.cumsum(self.global_num_tokens_for_logprob_gpu, dim=0)
|
|
dp_rank = get_parallel().attn_dp_rank
|
|
if dp_rank == 0:
|
|
dp_local_start_pos = torch.zeros_like(
|
|
self.global_num_tokens_for_logprob_gpu[0]
|
|
)
|
|
else:
|
|
dp_local_start_pos = cumtokens[dp_rank - 1]
|
|
|
|
self.dp_local_start_pos = dp_local_start_pos
|
|
self.dp_local_num_tokens = self.global_num_tokens_for_logprob_gpu[dp_rank]
|
|
|
|
hidden_size = get_dp_hidden_size()
|
|
dtype = get_dp_dtype()
|
|
device = get_dp_device()
|
|
|
|
if self.global_num_tokens_for_logprob_cpu is not None:
|
|
# create a smaller buffer to reduce peak memory usage
|
|
self.global_dp_buffer_len = sum(self.global_num_tokens_for_logprob_cpu)
|
|
else:
|
|
self.global_dp_buffer_len = self.global_dp_buffer_len
|
|
|
|
self.gathered_buffer = torch.empty(
|
|
(
|
|
self.global_dp_buffer_len,
|
|
hidden_size,
|
|
),
|
|
dtype=dtype,
|
|
device=device,
|
|
)
|
|
|
|
|
|
class LogitsProcessor(nn.Module):
|
|
def __init__(
|
|
self,
|
|
config,
|
|
skip_all_gather: bool = False,
|
|
logit_scale: Optional[float] = None,
|
|
return_full_logits: bool = False,
|
|
):
|
|
super().__init__()
|
|
self.config = config
|
|
self.vocab_size = config.vocab_size
|
|
self.logit_scale = logit_scale
|
|
self.use_attn_tp_group = get_parallel().enable_dp_lm_head
|
|
self.use_fp32_lm_head = get_exec().features.enable_fp32_lm_head
|
|
if self.use_attn_tp_group:
|
|
self.attn_tp_size = get_parallel().attn_tp_size
|
|
self.do_tensor_parallel_all_gather = (
|
|
not skip_all_gather and self.attn_tp_size > 1
|
|
)
|
|
self.do_tensor_parallel_all_gather_dp_attn = False
|
|
else:
|
|
self.do_tensor_parallel_all_gather = (
|
|
not skip_all_gather and get_parallel().tp_size > 1
|
|
)
|
|
self.do_tensor_parallel_all_gather_dp_attn = (
|
|
self.do_tensor_parallel_all_gather and get_parallel().attn_dp_size != 1
|
|
)
|
|
self.final_logit_softcapping = getattr(
|
|
self.config, "final_logit_softcapping", None
|
|
)
|
|
if (
|
|
self.final_logit_softcapping is not None
|
|
and self.final_logit_softcapping < 0
|
|
):
|
|
self.final_logit_softcapping = None
|
|
|
|
self.return_full_logits = return_full_logits
|
|
self.enable_mis = get_exec().features.enable_mis
|
|
self.rl_on_policy_target = get_exec().deterministic.rl_on_policy_target
|
|
|
|
self._logits_gatherer = triton_symm_mem_ag.MultimemAllGatherer(
|
|
max_tokens=triton_symm_mem_ag.recommended_max_tokens(
|
|
include_prefill=False, floor=128
|
|
),
|
|
enabled=self.do_tensor_parallel_all_gather and not self.use_attn_tp_group,
|
|
skip_entry_sync=True,
|
|
)
|
|
|
|
self.input_logprob_processor = InputLogprobProcessor()
|
|
|
|
def forward(
|
|
self,
|
|
input_ids,
|
|
hidden_states,
|
|
lm_head: VocabParallelEmbedding,
|
|
logits_metadata: Union[LogitsMetadata, ForwardBatch],
|
|
aux_hidden_states: Optional[AuxHiddenStates] = None,
|
|
hidden_states_before_norm: Optional[torch.Tensor] = None,
|
|
) -> LogitsProcessorOutput:
|
|
# Extract MIS indices before ForwardBatch → LogitsMetadata conversion
|
|
multi_item_delimiter_indices = None
|
|
if isinstance(logits_metadata, ForwardBatch):
|
|
multi_item_delimiter_indices = logits_metadata.multi_item_delimiter_indices
|
|
logits_metadata = LogitsMetadata.from_forward_batch(logits_metadata)
|
|
|
|
# Autotune dummy run discards this output. `is False` not `not`: None
|
|
# means no autotune pass, which must not skip. Placed before the MIS /
|
|
# DLLM / common dispatch so all three LM-head paths are skipped.
|
|
if _autotune_run_lm_head is False:
|
|
return LogitsProcessorOutput(next_token_logits=None)
|
|
|
|
# Multi-item scoring only for prefill-only requests with pre-computed indices.
|
|
if multi_item_delimiter_indices is not None and logits_metadata.is_prefill_only:
|
|
return self.compute_logprobs_for_multi_item_scoring(
|
|
input_ids,
|
|
hidden_states,
|
|
lm_head,
|
|
logits_metadata,
|
|
multi_item_delimiter_indices,
|
|
)
|
|
|
|
# Diffusion LLM only.
|
|
if logits_metadata.forward_mode.is_dllm_extend():
|
|
return self._get_dllm_logits(hidden_states, lm_head, logits_metadata)
|
|
|
|
# Get the last hidden states and last logits for the next token prediction
|
|
(
|
|
pruned_states,
|
|
pruned_states_before_norm,
|
|
aux_pruned_states,
|
|
sample_indices,
|
|
input_logprob_indices,
|
|
token_to_seq_idx,
|
|
) = self._get_pruned_states(
|
|
hidden_states,
|
|
hidden_states_before_norm,
|
|
aux_hidden_states,
|
|
logits_metadata,
|
|
)
|
|
|
|
hidden_states_to_store = self._get_hidden_states_to_store(
|
|
hidden_states,
|
|
hidden_states_before_norm,
|
|
aux_hidden_states,
|
|
pruned_states,
|
|
pruned_states_before_norm,
|
|
aux_pruned_states,
|
|
sample_indices,
|
|
logits_metadata,
|
|
)
|
|
del hidden_states
|
|
|
|
if not logits_metadata.extend_return_logprob:
|
|
# Compute logits for both input and sampled tokens.
|
|
logits = self._get_logits(pruned_states, lm_head, logits_metadata)
|
|
sampled_logits = (
|
|
logits[sample_indices] if sample_indices is not None else logits
|
|
)
|
|
|
|
# Decode mode or extend mode without return_logprob.
|
|
return LogitsProcessorOutput(
|
|
next_token_logits=sampled_logits,
|
|
hidden_states=hidden_states_to_store,
|
|
mm_input_embeds=logits_metadata.mm_input_embeds,
|
|
)
|
|
|
|
logprobs_result, sampled_logits = self.input_logprob_processor.forward(
|
|
pruned_states=pruned_states,
|
|
sample_indices=sample_indices,
|
|
input_logprob_indices=input_logprob_indices,
|
|
token_to_seq_idx=token_to_seq_idx,
|
|
lm_head=lm_head,
|
|
get_logits_fn=self._get_logits,
|
|
logits_metadata=logits_metadata,
|
|
skip_chunking_for_dp_attn=self.do_tensor_parallel_all_gather_dp_attn,
|
|
)
|
|
|
|
logits_output = LogitsProcessorOutput(
|
|
next_token_logits=sampled_logits,
|
|
hidden_states=hidden_states_to_store,
|
|
mm_input_embeds=logits_metadata.mm_input_embeds,
|
|
)
|
|
logprobs_result.write_input_to(logits_output)
|
|
return logits_output
|
|
|
|
def _get_pruned_states(
|
|
self,
|
|
hidden_states: torch.Tensor,
|
|
hidden_states_before_norm: Optional[torch.Tensor],
|
|
aux_hidden_states: Optional[AuxHiddenStates],
|
|
logits_metadata: LogitsMetadata,
|
|
):
|
|
pruned_states_before_norm: Optional[torch.Tensor] = None
|
|
aux_pruned_states = None
|
|
token_to_seq_idx = []
|
|
|
|
if (
|
|
logits_metadata.forward_mode.is_decode_or_idle()
|
|
or logits_metadata.forward_mode.is_target_verify()
|
|
or logits_metadata.forward_mode.is_draft_extend_v2()
|
|
):
|
|
if logits_metadata.draft_extend_select_index is not None:
|
|
# Only next_token_logits narrows to [bs, vocab]; the
|
|
# FULL-capture hidden stays unpruned.
|
|
pruned_states = hidden_states[logits_metadata.draft_extend_select_index]
|
|
else:
|
|
pruned_states = hidden_states
|
|
pruned_states_before_norm = hidden_states_before_norm
|
|
if aux_hidden_states is not None:
|
|
aux_pruned_states = (
|
|
aux_hidden_states
|
|
if isinstance(aux_hidden_states, torch.Tensor)
|
|
else [hidden for hidden in aux_hidden_states]
|
|
)
|
|
sample_indices = None
|
|
input_logprob_indices = None
|
|
|
|
elif (
|
|
logits_metadata.forward_mode.is_extend()
|
|
and not logits_metadata.extend_return_logprob
|
|
):
|
|
# Prefill without input logprobs.
|
|
last_index = torch.cumsum(logits_metadata.extend_seq_lens, dim=0) - 1
|
|
pruned_states = hidden_states[last_index]
|
|
if hidden_states_before_norm is not None:
|
|
pruned_states_before_norm = hidden_states_before_norm[last_index]
|
|
if aux_hidden_states is not None:
|
|
aux_pruned_states = (
|
|
aux_hidden_states[last_index]
|
|
if isinstance(aux_hidden_states, torch.Tensor)
|
|
else [hidden[last_index] for hidden in aux_hidden_states]
|
|
)
|
|
sample_indices = None
|
|
input_logprob_indices = None
|
|
else:
|
|
# Prefill with input logprobs.
|
|
# Find 4 different indices.
|
|
# 1. pruned_states: hidden states that we want logprobs from.
|
|
# 2. sample_indices: Indices that have sampled tokens.
|
|
# 3. input_logprob_indices: Indices that have input logprob tokens.
|
|
# 4. token_to_seq_idx: map each token to its sequence index
|
|
#
|
|
# Example
|
|
# -------
|
|
# Suppose a batch (flattened by sequence):
|
|
# [t00, t01, t02, t03, t10, t11, t12, t13, t14, t20, t21, t22, t23, t24, t25]
|
|
# extend_seq_lens_cpu = [4, 5, 6]
|
|
# extend_logprob_start_lens_cpu = [0, 5, 3]
|
|
#
|
|
# Then, the indices are:
|
|
# pruned_states -> [t00, t01, t02, t03, t14, t23, t24, t25]
|
|
# sample_indices -> [3, 4, 7]
|
|
# input_logprob_indices -> [0, 1, 2, 3, 5, 6, 7]
|
|
# token_to_seq_idx -> [0, 0, 0, 0, 1, 2, 2, 2]
|
|
#
|
|
# If chunk is enabled and chunk_size = 3, the chunks will be computed in a chunked manner:
|
|
# [t00, t01, t02], [t03, t14, t23], [t24, t25]
|
|
|
|
sample_index_pt = -1
|
|
sample_indices = []
|
|
input_logprob_indices_pt = 0
|
|
input_logprob_indices = []
|
|
pt, pruned_states_list, pruned_states_before_norm_list = 0, [], []
|
|
is_packed_aux_hidden_states = isinstance(aux_hidden_states, torch.Tensor)
|
|
aux_pruned_states_lists = None
|
|
if aux_hidden_states is not None:
|
|
aux_pruned_states_lists = (
|
|
[]
|
|
if is_packed_aux_hidden_states
|
|
else [[] for _ in aux_hidden_states]
|
|
)
|
|
|
|
for idx, (extend_logprob_start_len, extend_len) in enumerate(
|
|
zip(
|
|
logits_metadata.extend_logprob_start_lens_cpu,
|
|
logits_metadata.extend_seq_lens_cpu,
|
|
)
|
|
):
|
|
# It can happen in chunked prefill. We still need to sample 1 token,
|
|
# But we don't want to include it in input logprob.
|
|
if extend_len == extend_logprob_start_len:
|
|
start_len = extend_logprob_start_len - 1
|
|
else:
|
|
start_len = extend_logprob_start_len
|
|
|
|
# We always need at least 1 token to sample because that's required
|
|
# by a caller.
|
|
assert extend_len > start_len
|
|
pruned_states_list.append(
|
|
hidden_states[pt + start_len : pt + extend_len]
|
|
)
|
|
if hidden_states_before_norm is not None:
|
|
pruned_states_before_norm_list.append(
|
|
hidden_states_before_norm[pt + start_len : pt + extend_len]
|
|
)
|
|
if aux_pruned_states_lists is not None:
|
|
if is_packed_aux_hidden_states:
|
|
aux_pruned_states_lists.append(
|
|
aux_hidden_states[pt + start_len : pt + extend_len]
|
|
)
|
|
else:
|
|
for j, hidden in enumerate(aux_hidden_states):
|
|
aux_pruned_states_lists[j].append(
|
|
hidden[pt + start_len : pt + extend_len]
|
|
)
|
|
# Map each token to its sequence index, for chunked computation
|
|
# of input logprobs
|
|
token_to_seq_idx.extend([idx] * (extend_len - start_len))
|
|
pt += extend_len
|
|
sample_index_pt += extend_len - start_len
|
|
sample_indices.append(sample_index_pt)
|
|
input_logprob_indices.extend(
|
|
[
|
|
input_logprob_indices_pt + i
|
|
for i in range(extend_len - extend_logprob_start_len)
|
|
]
|
|
)
|
|
input_logprob_indices_pt += extend_len - start_len
|
|
|
|
pruned_states = torch.cat(pruned_states_list)
|
|
if hidden_states_before_norm is not None:
|
|
pruned_states_before_norm = torch.cat(pruned_states_before_norm_list)
|
|
if aux_pruned_states_lists is not None:
|
|
aux_pruned_states = (
|
|
torch.cat(aux_pruned_states_lists)
|
|
if is_packed_aux_hidden_states
|
|
else [torch.cat(lst) for lst in aux_pruned_states_lists]
|
|
)
|
|
|
|
# Build the index tensors via pinned host memory + non-blocking H2D
|
|
# so the small copy doesn't drain the stream.
|
|
sample_indices = torch.tensor(
|
|
sample_indices,
|
|
dtype=torch.int64,
|
|
pin_memory=is_pin_memory_available(),
|
|
).to(pruned_states.device, non_blocking=True)
|
|
input_logprob_indices = torch.tensor(
|
|
input_logprob_indices,
|
|
dtype=torch.int64,
|
|
pin_memory=is_pin_memory_available(),
|
|
).to(pruned_states.device, non_blocking=True)
|
|
|
|
return (
|
|
pruned_states,
|
|
pruned_states_before_norm,
|
|
aux_pruned_states,
|
|
sample_indices,
|
|
input_logprob_indices,
|
|
token_to_seq_idx,
|
|
)
|
|
|
|
def _get_hidden_states_to_store(
|
|
self,
|
|
hidden_states: torch.Tensor,
|
|
hidden_states_before_norm: Optional[torch.Tensor],
|
|
aux_hidden_states: Optional[AuxHiddenStates],
|
|
pruned_states: torch.Tensor,
|
|
pruned_states_before_norm: Optional[torch.Tensor],
|
|
aux_pruned_states: Optional[AuxHiddenStates],
|
|
sample_indices: Optional[torch.Tensor],
|
|
logits_metadata: LogitsMetadata,
|
|
) -> Optional[torch.Tensor]:
|
|
hidden_states_to_store: Optional[torch.Tensor] = None
|
|
hidden_states_to_store_before_norm: Optional[torch.Tensor] = None
|
|
if logits_metadata.capture_hidden_mode.need_capture():
|
|
if logits_metadata.capture_hidden_mode.is_full():
|
|
if aux_hidden_states is not None:
|
|
hidden_states_to_store = pack_aux_hidden_states(aux_hidden_states)
|
|
else:
|
|
hidden_states_to_store = hidden_states
|
|
hidden_states_to_store_before_norm = hidden_states_before_norm
|
|
elif logits_metadata.capture_hidden_mode.is_last():
|
|
# Get the last token hidden states. If sample_indices is None,
|
|
# pruned states only contain the last tokens already.
|
|
if aux_hidden_states is not None:
|
|
assert aux_pruned_states is not None
|
|
aux_pruned_states = pack_aux_hidden_states(aux_pruned_states)
|
|
hidden_states_to_store = (
|
|
aux_pruned_states[sample_indices]
|
|
if sample_indices is not None
|
|
else aux_pruned_states
|
|
)
|
|
else:
|
|
hidden_states_to_store = (
|
|
pruned_states[sample_indices]
|
|
if sample_indices is not None
|
|
else pruned_states
|
|
)
|
|
if hidden_states_before_norm is not None:
|
|
hidden_states_to_store_before_norm = (
|
|
pruned_states_before_norm[sample_indices]
|
|
if sample_indices is not None
|
|
else pruned_states_before_norm
|
|
)
|
|
else:
|
|
assert False, "Should never reach"
|
|
|
|
if hidden_states_to_store_before_norm is not None:
|
|
# NOTE: when hidden_states_before_norm is provided, we always
|
|
# prefer to return it.
|
|
hidden_states_to_store = hidden_states_to_store_before_norm
|
|
|
|
return hidden_states_to_store
|
|
|
|
def _get_logits(
|
|
self,
|
|
hidden_states: torch.Tensor,
|
|
lm_head: VocabParallelEmbedding,
|
|
logits_metadata: LogitsMetadata,
|
|
embedding_bias: Optional[torch.Tensor] = None,
|
|
use_logits_buffer: bool = True,
|
|
) -> torch.Tensor:
|
|
"""Get logits from hidden_states.
|
|
|
|
If sampled_logits_only is True, it means hidden_states only contain the
|
|
last position (e.g., extend without input logprobs). The caller should
|
|
guarantee the given hidden_states follow this constraint.
|
|
"""
|
|
hidden_states, local_hidden_states = self._gather_dp_attn_hidden_states(
|
|
hidden_states, logits_metadata
|
|
)
|
|
|
|
logits = self._compute_lm_head(hidden_states, lm_head, embedding_bias)
|
|
|
|
if self.logit_scale is not None:
|
|
logits.mul_(self.logit_scale)
|
|
|
|
if self.do_tensor_parallel_all_gather:
|
|
if self.use_attn_tp_group:
|
|
logits = self._gather_attn_tp_logits(logits)
|
|
else:
|
|
logits = self._logits_gatherer(logits)
|
|
|
|
logits = self._scatter_dp_attn_logits(
|
|
logits, local_hidden_states, logits_metadata
|
|
)
|
|
|
|
logits = self._copy_logits_to_buffer(
|
|
logits, logits_metadata, use_buffer=use_logits_buffer
|
|
)
|
|
|
|
if self.final_logit_softcapping:
|
|
if not (_is_npu or _is_cpu):
|
|
fused_softcap(logits, self.final_logit_softcapping)
|
|
else:
|
|
logits = self.final_logit_softcapping * torch.tanh(
|
|
logits / self.final_logit_softcapping
|
|
)
|
|
|
|
return logits
|
|
|
|
def _compute_lm_head(
|
|
self,
|
|
hidden_states: torch.Tensor,
|
|
lm_head: VocabParallelEmbedding,
|
|
embedding_bias: Optional[torch.Tensor] = None,
|
|
) -> torch.Tensor:
|
|
quant_method = getattr(lm_head, "quant_method", None)
|
|
if hasattr(lm_head, "set_lora") and hasattr(lm_head, "apply_lora"):
|
|
# This is a LoRA-wrapped module, use its forward method
|
|
logits = lm_head(hidden_states)
|
|
elif should_apply_lm_head_quant_method(lm_head, quant_method):
|
|
logits = quant_method.apply(lm_head, hidden_states, embedding_bias)
|
|
elif hasattr(lm_head, "weight"):
|
|
# Normal linear layer
|
|
if self.use_fp32_lm_head:
|
|
logits = torch.matmul(
|
|
hidden_states.to(torch.float32), lm_head.weight.to(torch.float32).T
|
|
)
|
|
elif use_intel_amx_backend(lm_head):
|
|
logits = torch.ops.sgl_kernel.weight_packed_linear(
|
|
hidden_states.to(lm_head.weight.dtype),
|
|
lm_head.weight,
|
|
None, # bias
|
|
True, # is_vnni
|
|
)
|
|
elif self.rl_on_policy_target is not None:
|
|
# Due to tie-weight, we may not be able to change lm_head's weight dtype
|
|
logits = torch.matmul(
|
|
hidden_states.bfloat16(), lm_head.weight.T.bfloat16()
|
|
)
|
|
else:
|
|
logits = torch.matmul(
|
|
hidden_states.to(lm_head.weight.dtype), lm_head.weight.T
|
|
)
|
|
else:
|
|
# GGUF models
|
|
# TODO: use weight_packed_linear for GGUF models
|
|
if self.use_fp32_lm_head:
|
|
with torch.cuda.amp.autocast(enabled=False):
|
|
logits = lm_head.quant_method.apply(
|
|
lm_head, hidden_states.to(torch.float32), embedding_bias
|
|
)
|
|
else:
|
|
logits = lm_head.quant_method.apply(
|
|
lm_head, hidden_states, embedding_bias
|
|
)
|
|
return logits
|
|
|
|
def _gather_dp_attn_hidden_states(
|
|
self, hidden_states: torch.Tensor, logits_metadata: LogitsMetadata
|
|
) -> Tuple[torch.Tensor, torch.Tensor]:
|
|
if self.do_tensor_parallel_all_gather_dp_attn:
|
|
logits_metadata.compute_dp_attention_metadata()
|
|
local_hidden_states = hidden_states
|
|
hidden_states = logits_metadata.gathered_buffer
|
|
dp_gather_replicate(hidden_states, local_hidden_states, logits_metadata)
|
|
return hidden_states, local_hidden_states
|
|
return hidden_states, hidden_states
|
|
|
|
def _gather_attn_tp_logits(self, logits: torch.Tensor) -> torch.Tensor:
|
|
if self.vocab_size % self.attn_tp_size == 0:
|
|
global_logits = torch.empty(
|
|
(
|
|
self.attn_tp_size,
|
|
logits.shape[0],
|
|
self.vocab_size // self.attn_tp_size,
|
|
),
|
|
device=logits.device,
|
|
dtype=logits.dtype,
|
|
)
|
|
attn_tp_all_gather_into_tensor(global_logits, logits)
|
|
global_logits = global_logits.permute(1, 0, 2).reshape(
|
|
logits.shape[0], self.vocab_size
|
|
)
|
|
else:
|
|
global_logits = torch.empty(
|
|
(self.vocab_size, logits.shape[0]),
|
|
device=logits.device,
|
|
dtype=logits.dtype,
|
|
)
|
|
global_logits = global_logits.T
|
|
attn_tp_all_gather(
|
|
list(global_logits.tensor_split(self.attn_tp_size, dim=-1)),
|
|
logits,
|
|
)
|
|
return global_logits
|
|
|
|
def _scatter_dp_attn_logits(
|
|
self,
|
|
logits: torch.Tensor,
|
|
local_hidden_states: torch.Tensor,
|
|
logits_metadata: LogitsMetadata,
|
|
) -> torch.Tensor:
|
|
if self.do_tensor_parallel_all_gather_dp_attn:
|
|
global_logits = logits
|
|
logits = torch.empty(
|
|
(local_hidden_states.shape[0], global_logits.shape[1]),
|
|
device=global_logits.device,
|
|
dtype=global_logits.dtype,
|
|
)
|
|
dp_scatter(logits, global_logits, logits_metadata)
|
|
return logits
|
|
|
|
def _copy_logits_to_buffer(
|
|
self,
|
|
logits: torch.Tensor,
|
|
logits_metadata: LogitsMetadata,
|
|
use_buffer: bool = True,
|
|
) -> torch.Tensor:
|
|
logits_buffer = logits_metadata.next_token_logits_buffer if use_buffer else None
|
|
if logits.shape[-1] > self.vocab_size:
|
|
logits = logits[:, : self.vocab_size]
|
|
logits_width = logits.shape[-1]
|
|
# The shared logits buffer is keyed by vocab width and rows; skip it
|
|
# when this batch has a different logits shape than the graph buffer.
|
|
if logits_buffer is not None and tuple(logits_buffer.shape) == tuple(
|
|
logits.shape
|
|
):
|
|
assert logits_buffer.dtype == torch.float
|
|
logits_buffer.copy_(logits)
|
|
logits = logits_buffer
|
|
else:
|
|
logits = logits.float()
|
|
return logits
|
|
|
|
def _get_dllm_logits(
|
|
self,
|
|
hidden_states: torch.Tensor,
|
|
lm_head: VocabParallelEmbedding,
|
|
logits_metadata: LogitsMetadata,
|
|
) -> LogitsProcessorOutput:
|
|
assert self.return_full_logits
|
|
full_logits = self._get_logits(hidden_states, lm_head, logits_metadata)
|
|
return LogitsProcessorOutput(
|
|
full_logits=full_logits,
|
|
next_token_logits=None,
|
|
)
|
|
|
|
def compute_logprobs_for_multi_item_scoring(
|
|
self,
|
|
input_ids,
|
|
hidden_states,
|
|
lm_head: VocabParallelEmbedding,
|
|
logits_metadata: Union[LogitsMetadata, ForwardBatch],
|
|
multi_item_delimiter_indices: List[torch.Tensor],
|
|
):
|
|
"""
|
|
Compute logprobs for multi-item scoring using pre-computed delimiter indices.
|
|
|
|
Sequence format: Query<delimiter>Item1<delimiter>Item2<delimiter>...
|
|
Scoring positions: Extracts logprobs at positions before each <delimiter>
|
|
|
|
Args:
|
|
input_ids: Input token IDs. Shape: [total_sequence_length].
|
|
hidden_states: Hidden states from the model. Shape: [sequence_length, hidden_dim].
|
|
lm_head: Language model head for computing logits.
|
|
logits_metadata: Metadata containing batch info and logprob specs.
|
|
multi_item_delimiter_indices: Pre-computed delimiter positions per request (CPU tensors).
|
|
"""
|
|
# Compute positions just before each delimiter.
|
|
# Build offset-adjusted indices on CPU, then do a single CPU→GPU transfer.
|
|
device = input_ids.device
|
|
all_tensors = []
|
|
if logits_metadata.extend_seq_lens_cpu is not None:
|
|
offset = 0
|
|
for req_seq_len, indices_tensor in zip(
|
|
logits_metadata.extend_seq_lens_cpu, multi_item_delimiter_indices
|
|
):
|
|
if len(indices_tensor) > 0:
|
|
# Note: if the first delimiter is at position 0 (empty query),
|
|
# indices - 1 wraps to -1. This is harmless — the first
|
|
# delimiter entry is always discarded by
|
|
# _process_multi_item_scoring_results.
|
|
all_tensors.append(indices_tensor + (offset - 1))
|
|
offset += req_seq_len
|
|
else:
|
|
all_tensors.append(multi_item_delimiter_indices[0] - 1)
|
|
multi_item_indices = torch.cat(all_tensors).to(device, non_blocking=True)
|
|
|
|
# Extract hidden states at delimiter positions for multi-item scoring
|
|
sliced_hidden = hidden_states[multi_item_indices]
|
|
|
|
sliced_logits = self._get_logits(sliced_hidden, lm_head, logits_metadata)
|
|
sliced_logprobs = torch.nn.functional.log_softmax(sliced_logits, dim=-1)
|
|
|
|
# Initialize return values
|
|
input_token_ids_logprobs_val = []
|
|
input_token_ids_logprobs_idx = []
|
|
input_top_logprobs_val = None
|
|
input_top_logprobs_idx = None
|
|
|
|
# Recalculate extend_logprob_pruned_lens_cpu to match delimiter counts per request
|
|
if (
|
|
logits_metadata.token_ids_logprobs
|
|
or logits_metadata.extend_return_top_logprob
|
|
):
|
|
logits_metadata.extend_logprob_pruned_lens_cpu = [
|
|
len(t) for t in multi_item_delimiter_indices
|
|
]
|
|
|
|
# Get the logprobs of specified token ids
|
|
if logits_metadata.extend_token_ids_logprob:
|
|
(
|
|
input_token_ids_logprobs_val,
|
|
input_token_ids_logprobs_idx,
|
|
) = get_token_ids_logprobs_raw(
|
|
sliced_logprobs,
|
|
logits_metadata.token_ids_logprobs,
|
|
stage=LogprobStage.PREFILL,
|
|
extend_logprob_pruned_lens_cpu=logits_metadata.extend_logprob_pruned_lens_cpu,
|
|
no_copy_to_cpu=True,
|
|
)
|
|
|
|
# Get the logprob of top-k tokens
|
|
if logits_metadata.extend_return_top_logprob:
|
|
(
|
|
input_top_logprobs_val,
|
|
input_top_logprobs_idx,
|
|
) = get_top_logprobs_raw(
|
|
sliced_logprobs,
|
|
logits_metadata.top_logprobs_nums,
|
|
stage=LogprobStage.PREFILL,
|
|
extend_logprob_pruned_lens_cpu=logits_metadata.extend_logprob_pruned_lens_cpu,
|
|
)
|
|
|
|
# MIS scores come from input_token_ids_logprobs_val (label-token logprobs),
|
|
# not from per-position input_token_logprobs. However, the shared logprob
|
|
# pipeline (add_input_logprob_return_values) asserts input_token_logprobs is
|
|
# non-None, converts it to a tuple, slices it, and validates its length —
|
|
# all before score_request() ever sees the result. We can't set it to None
|
|
# without changing those shared asserts, so we fill with zeros to satisfy
|
|
# the pipeline. score_request() ignores this field entirely.
|
|
input_token_logprobs = torch.zeros(multi_item_indices.shape[0], device=device)
|
|
|
|
return LogitsProcessorOutput(
|
|
next_token_logits=None,
|
|
input_token_logprobs=input_token_logprobs,
|
|
input_top_logprobs_val=input_top_logprobs_val,
|
|
input_top_logprobs_idx=input_top_logprobs_idx,
|
|
input_token_ids_logprobs_val=input_token_ids_logprobs_val,
|
|
input_token_ids_logprobs_idx=input_token_ids_logprobs_idx,
|
|
mm_input_embeds=logits_metadata.mm_input_embeds,
|
|
)
|
|
|
|
|
|
def _has_lm_head_runtime_attrs(lm_head, attr_names: Tuple[str, ...]) -> bool:
|
|
return all(hasattr(lm_head, attr_name) for attr_name in attr_names)
|
|
|
|
|
|
def should_apply_lm_head_quant_method(lm_head, quant_method) -> bool:
|
|
if (
|
|
quant_method is None
|
|
or not hasattr(lm_head, "weight")
|
|
or not callable(getattr(quant_method, "apply", None))
|
|
):
|
|
return False
|
|
|
|
method_name = type(quant_method).__name__
|
|
if method_name in _UNQUANTIZED_LM_HEAD_METHODS:
|
|
return False
|
|
|
|
# Some draft models share an unquantized target lm_head tensor while still
|
|
# carrying the draft model's stale ModelOpt quant_method. Only use the
|
|
# ModelOpt lm_head kernel when the runtime quantization state matches it.
|
|
if method_name == "ModelOptFp4LinearMethod":
|
|
if lm_head.weight.dtype == torch.int32 and _has_lm_head_runtime_attrs(
|
|
lm_head,
|
|
(
|
|
"weight_scale",
|
|
"weight_global_scale",
|
|
"workspace",
|
|
"input_size_per_partition",
|
|
"output_size_per_partition",
|
|
),
|
|
):
|
|
return True
|
|
return lm_head.weight.dtype == torch.uint8 and _has_lm_head_runtime_attrs(
|
|
lm_head,
|
|
(
|
|
"weight_scale_interleaved",
|
|
"alpha",
|
|
"input_scale_inv",
|
|
"input_size_per_partition",
|
|
"output_size_per_partition",
|
|
),
|
|
)
|
|
if method_name == "ModelOptNvFp4A16LinearMethod":
|
|
return lm_head.weight.dtype == torch.int32 and _has_lm_head_runtime_attrs(
|
|
lm_head,
|
|
(
|
|
"weight_scale",
|
|
"weight_global_scale",
|
|
"workspace",
|
|
"input_size_per_partition",
|
|
"output_size_per_partition",
|
|
),
|
|
)
|
|
if method_name == "ModelOptFp8LinearMethod":
|
|
return (
|
|
lm_head.weight.dtype == torch.float8_e4m3fn
|
|
and _has_lm_head_runtime_attrs(lm_head, ("weight_scale", "input_scale"))
|
|
)
|
|
|
|
return True
|