minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

[Refactor] Unify logprob results into a single LogprobResult and rename chunk env vars (#31733)

Browse Source
This commit is contained in:
Liangsheng Yin
2026-07-20 12:44:27 -07:00
committed by GitHub
parent e149cdb337
commit ff6c755952
11 changed files with 139 additions and 151 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/environ.py
+7 -3
View File
@@ -834,9 +834,13 @@ class Envs:
# Sparse Embeddings
SGLANG_EMBEDDINGS_SPARSE_HEAD = EnvStr(None)
# Logits processor
SGLANG_ENABLE_LOGITS_PROCESSER_CHUNK = EnvBool(True)
SGLANG_LOGITS_PROCESSER_CHUNK_SIZE = EnvInt(2048)
# Logprob processor
SGLANG_ENABLE_LOGPROB_CHUNK = EnvBoolWithAlias(
True, deprecated_name="SGLANG_ENABLE_LOGITS_PROCESSER_CHUNK"
)
SGLANG_LOGPROB_CHUNK_SIZE = EnvIntWithAlias(
2048, deprecated_name="SGLANG_LOGITS_PROCESSER_CHUNK_SIZE"
)
# Tool-Call behavior
SGLANG_TOOL_STRICT_LEVEL = EnvInt(ToolStrictLevel.OFF)
python/sglang/srt/layers/logits_processor.py
+3 -6
View File
@@ -473,16 +473,13 @@ class LogitsProcessor(nn.Module):
skip_chunking_for_dp_attn=self.do_tensor_parallel_all_gather_dp_attn,
)
return LogitsProcessorOutput(
logits_output = LogitsProcessorOutput(
next_token_logits=sampled_logits,
hidden_states=hidden_states_to_store,
input_token_logprobs=logprobs_result.input_token_logprobs,
input_top_logprobs_val=logprobs_result.input_top_logprobs_val,
input_top_logprobs_idx=logprobs_result.input_top_logprobs_idx,
input_token_ids_logprobs_val=logprobs_result.input_token_ids_logprobs_val,
input_token_ids_logprobs_idx=logprobs_result.input_token_ids_logprobs_idx,
mm_input_embeds=logits_metadata.mm_input_embeds,
)
logprobs_result.write_input_to(logits_output)
return logits_output
def _get_pruned_states(
self,
python/sglang/srt/layers/logprob_processor.py
+105 -113
View File
@@ -12,7 +12,6 @@ from sglang.srt.environ import envs
if TYPE_CHECKING:
from sglang.srt.layers.logits_processor import LogitsMetadata, LogitsProcessorOutput
from sglang.srt.layers.vocab_parallel_embedding import VocabParallelEmbedding
from sglang.srt.sampling.sampling_batch_info import SamplingBatchInfo
logger = logging.getLogger(__name__)
@@ -23,12 +22,43 @@ class LogprobStage(Enum):
@dataclasses.dataclass
class InputLogprobsResult:
input_token_logprobs: torch.Tensor
input_top_logprobs_val: Optional[List] = None
input_top_logprobs_idx: Optional[List] = None
input_token_ids_logprobs_val: Optional[List] = None
input_token_ids_logprobs_idx: Optional[List] = None
class LogprobResult:
"""Logprob fields produced by Input/OutputLogprobProcessor.
Input (prefill) always fills token_logprobs; output (decode / scoring)
fills on demand. write_input_to / write_output_to flush populated fields
onto LogitsProcessorOutput, so the IPC / D2H wire format stays unchanged.
"""
token_logprobs: Optional[torch.Tensor] = None
top_logprobs_val: Optional[List] = None
top_logprobs_idx: Optional[List] = None
token_ids_logprobs_val: Optional[List] = None
token_ids_logprobs_idx: Optional[List] = None
def write_input_to(self, logits_output: LogitsProcessorOutput) -> None:
if self.token_logprobs is not None:
logits_output.input_token_logprobs = self.token_logprobs
if self.top_logprobs_val is not None:
logits_output.input_top_logprobs_val = self.top_logprobs_val
logits_output.input_top_logprobs_idx = self.top_logprobs_idx
if self.token_ids_logprobs_val is not None:
logits_output.input_token_ids_logprobs_val = self.token_ids_logprobs_val
logits_output.input_token_ids_logprobs_idx = self.token_ids_logprobs_idx
def write_output_to(self, logits_output: LogitsProcessorOutput) -> None:
if self.token_logprobs is not None:
logits_output.next_token_logprobs = self.token_logprobs
if self.top_logprobs_val is not None:
logits_output.next_token_top_logprobs_val = self.top_logprobs_val
logits_output.next_token_top_logprobs_idx = self.top_logprobs_idx
if self.token_ids_logprobs_val is not None:
logits_output.next_token_token_ids_logprobs_val = (
self.token_ids_logprobs_val
)
logits_output.next_token_token_ids_logprobs_idx = (
self.token_ids_logprobs_idx
)
def get_top_logprobs_raw(
@@ -135,29 +165,27 @@ def get_token_ids_logprobs(logprobs, token_ids_logprobs, no_copy_to_cpu=False):
def get_top_logprobs_chunk(
logprobs: torch.Tensor,
logits_metadata: LogitsMetadata,
top_k_nums: List[int],
pruned_lens: List[int],
input_top_logprobs_val: List,
input_top_logprobs_idx: List,
top_logprobs_val: List,
top_logprobs_idx: List,
split_pruned_len: int,
) -> int:
"""Get top-k logprobs for each sequence in the chunk.
Args:
logprobs: Log probabilities tensor of shape [seq_len, vocab_size]
logits_metadata: Metadata containing top-k and pruned length info
top_k_nums: List of top-k numbers for each sequence
pruned_lens: List of pruned lengths for each sequence
input_top_logprobs_val: List to store top-k logprob values
input_top_logprobs_idx: List to store top-k token indices
top_logprobs_val: List to store top-k logprob values
top_logprobs_idx: List to store top-k token indices
split_pruned_len: Length of pruned tokens from previous chunk
Returns:
int: Number of remaining tokens to process in next chunk
"""
# Empty chunks still walk the slice to emit placeholder entries.
max_k = max(logits_metadata.top_logprobs_nums)
max_k = max(top_k_nums)
ret = logprobs.topk(max_k, dim=1)
values = ret.values.tolist()
indices = ret.indices.tolist()
@@ -175,8 +203,8 @@ def get_top_logprobs_chunk(
if pruned_len <= 0:
# if pruned length is less than or equal to 0,
# there is no top-k logprobs to process
input_top_logprobs_val.append([])
input_top_logprobs_idx.append([])
top_logprobs_val.append([])
top_logprobs_idx.append([])
continue
# Get the top-k logprobs
@@ -195,11 +223,11 @@ def get_top_logprobs_chunk(
# Split-sequence continuations extend; everyone else owns a fresh
# (possibly empty) entry.
if split_pruned_len > 0:
input_top_logprobs_val[-1].extend(val)
input_top_logprobs_idx[-1].extend(idx)
top_logprobs_val[-1].extend(val)
top_logprobs_idx[-1].extend(idx)
else:
input_top_logprobs_val.append(val)
input_top_logprobs_idx.append(idx)
top_logprobs_val.append(val)
top_logprobs_idx.append(idx)
pt += pruned_len
return next_split_pruned_len
@@ -209,19 +237,18 @@ def get_token_ids_logprobs_chunk(
logprobs: torch.Tensor,
token_ids_logprobs: List[int],
pruned_lens: List[int],
input_token_ids_logprobs_val: List,
input_token_ids_logprobs_idx: List,
token_ids_logprobs_val: List,
token_ids_logprobs_idx: List,
split_pruned_len: int = 0,
):
"""Get token_ids logprobs for each sequence in the chunk.
Args:
logprobs: Log probabilities tensor of shape [seq_len, vocab_size]
logits_metadata: Metadata containing token IDs and pruned length info
token_ids_logprobs: List of token IDs for each sequence
pruned_lens: List of pruned lengths for each sequence
input_token_ids_logprobs_val: List to store token logprob values
input_token_ids_logprobs_idx: List to store token indices
token_ids_logprobs_val: List to store token logprob values
token_ids_logprobs_idx: List to store token indices
split_pruned_len: Length of pruned tokens from previous chunk
Returns:
@@ -245,8 +272,8 @@ def get_token_ids_logprobs_chunk(
if pruned_len <= 0:
# if pruned length is less than or equal to 0,
# there is no token ids logprobs to process
input_token_ids_logprobs_val.append([])
input_token_ids_logprobs_idx.append([])
token_ids_logprobs_val.append([])
token_ids_logprobs_idx.append([])
continue
# Get the token ids logprobs
@@ -264,11 +291,11 @@ def get_token_ids_logprobs_chunk(
# Split-sequence continuations extend; everyone else owns a fresh
# (possibly empty) entry.
if split_pruned_len > 0:
input_token_ids_logprobs_val[-1].extend(val)
input_token_ids_logprobs_idx[-1].extend(idx)
token_ids_logprobs_val[-1].extend(val)
token_ids_logprobs_idx[-1].extend(idx)
else:
input_token_ids_logprobs_val.append(val)
input_token_ids_logprobs_idx.append(idx)
token_ids_logprobs_val.append(val)
token_ids_logprobs_idx.append(idx)
pt += pruned_len
return next_split_pruned_len
@@ -349,9 +376,9 @@ class InputLogprobProcessor:
def __init__(self):
# enable chunked logprobs processing
self.enable_logprobs_chunk = envs.SGLANG_ENABLE_LOGITS_PROCESSER_CHUNK.get()
self.enable_logprobs_chunk = envs.SGLANG_ENABLE_LOGPROB_CHUNK.get()
# chunk size for logprobs processing
self.logprobs_chunk_size = envs.SGLANG_LOGITS_PROCESSER_CHUNK_SIZE.get()
self.logprobs_chunk_size = envs.SGLANG_LOGPROB_CHUNK_SIZE.get()
def forward(
self,
@@ -363,7 +390,7 @@ class InputLogprobProcessor:
get_logits_fn: Callable,
logits_metadata: LogitsMetadata,
skip_chunking_for_dp_attn: bool = False,
) -> Tuple[InputLogprobsResult, torch.Tensor]:
) -> Tuple[LogprobResult, torch.Tensor]:
# Non-chunked = one chunk covering every row. DP-attention must stay
# single-chunk: the collective schedule cannot depend on per-rank rows.
if (
@@ -396,24 +423,24 @@ class InputLogprobProcessor:
get_logits_fn: Callable,
logits_metadata: LogitsMetadata,
chunk_size: int,
) -> Tuple[InputLogprobsResult, torch.Tensor]:
) -> Tuple[LogprobResult, torch.Tensor]:
"""Compute input logprobs chunk by chunk to cap peak memory."""
total_size = pruned_states.shape[0]
num_chunks = (total_size + chunk_size - 1) // chunk_size
input_token_logprobs = []
token_logprobs = []
if logits_metadata.extend_return_top_logprob:
input_top_logprobs_val = []
input_top_logprobs_idx = []
top_logprobs_val = []
top_logprobs_idx = []
else:
input_top_logprobs_val = None
input_top_logprobs_idx = None
top_logprobs_val = None
top_logprobs_idx = None
if logits_metadata.extend_token_ids_logprob:
input_token_ids_logprobs_val = []
input_token_ids_logprobs_idx = []
token_ids_logprobs_val = []
token_ids_logprobs_idx = []
else:
input_token_ids_logprobs_val = None
input_token_ids_logprobs_idx = None
token_ids_logprobs_val = None
token_ids_logprobs_idx = None
# If a single sequence is split into multiple chunks, we need to keep track
# of the pruned length of the sequences in the previous chunks.
@@ -472,11 +499,9 @@ class InputLogprobProcessor:
# Compute the logprobs of the chunk. Free the raw logits before the
# out-of-place log_softmax: keeping all three alive is a 3x peak,
# which OOMs when the single chunk covers a large batch.
chunk_input_logprobs = chunk_logits[chunk_indices]
chunk_logprobs = chunk_logits[chunk_indices]
del chunk_logits
chunk_input_logprobs = torch.nn.functional.log_softmax(
chunk_input_logprobs, dim=-1
)
chunk_logprobs = torch.nn.functional.log_softmax(chunk_logprobs, dim=-1)
# End at the last row inside the chunk; token_to_seq_idx[end_idx]
# belongs to the next chunk and would emit its sequence twice.
@@ -491,12 +516,11 @@ class InputLogprobProcessor:
chunk_slice
]
split_len_topk = get_top_logprobs_chunk(
chunk_input_logprobs,
logits_metadata,
chunk_logprobs,
top_k_nums,
pruned_lens,
input_top_logprobs_val,
input_top_logprobs_idx,
top_logprobs_val,
top_logprobs_idx,
split_len_topk,
)
@@ -507,24 +531,22 @@ class InputLogprobProcessor:
chunk_slice
]
split_len_token_ids = get_token_ids_logprobs_chunk(
chunk_input_logprobs,
chunk_logprobs,
token_ids_logprobs,
pruned_lens,
input_token_ids_logprobs_val,
input_token_ids_logprobs_idx,
token_ids_logprobs_val,
token_ids_logprobs_idx,
split_len_token_ids,
)
# Get the logprob of the requested token ids
chunk_input_token_logprobs = chunk_input_logprobs[
torch.arange(
chunk_input_logprobs.shape[0], device=chunk_input_logprobs.device
),
chunk_token_logprobs = chunk_logprobs[
torch.arange(chunk_logprobs.shape[0], device=chunk_logprobs.device),
logits_metadata.extend_input_logprob_token_ids_gpu[mask_indices],
]
input_token_logprobs.append(chunk_input_token_logprobs)
token_logprobs.append(chunk_token_logprobs)
# Free before the next chunk's logits (bf16 + fp32) materialize.
del chunk_input_logprobs
del chunk_logprobs
# Restore the full-pruned lm_head batch_info after chunk iteration.
if num_chunks > 1 and hasattr(lm_head, "reset_lm_head_pass"):
@@ -534,15 +556,15 @@ class InputLogprobProcessor:
lm_head.reset_lm_head_pass()
# Concatenate the results
input_token_logprobs = torch.cat(input_token_logprobs, dim=0)
token_logprobs = torch.cat(token_logprobs, dim=0)
return (
InputLogprobsResult(
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,
LogprobResult(
token_logprobs=token_logprobs,
top_logprobs_val=top_logprobs_val,
top_logprobs_idx=top_logprobs_idx,
token_ids_logprobs_val=token_ids_logprobs_val,
token_ids_logprobs_idx=token_ids_logprobs_idx,
),
sampled_logits,
)
@@ -636,55 +658,26 @@ def get_token_ids_logprobs_batch_optimized(
return output_token_ids_logprobs_val, output_token_ids_logprobs_idx
@dataclasses.dataclass
class OutputLogprobsResult:
"""Output-side counterpart of InputLogprobsResult.
Built by OutputLogprobProcessor; write_to() flushes the populated fields
onto LogitsProcessorOutput, so the IPC / D2H wire format stays unchanged.
"""
token_logprobs: Optional[torch.Tensor] = None
top_logprobs_val: Optional[List] = None
top_logprobs_idx: Optional[List] = None
token_ids_logprobs_val: Optional[List] = None
token_ids_logprobs_idx: Optional[List] = None
def write_to(self, logits_output: LogitsProcessorOutput) -> None:
if self.token_logprobs is not None:
logits_output.next_token_logprobs = self.token_logprobs
if self.top_logprobs_val is not None:
logits_output.next_token_top_logprobs_val = self.top_logprobs_val
logits_output.next_token_top_logprobs_idx = self.top_logprobs_idx
if self.token_ids_logprobs_val is not None:
logits_output.next_token_token_ids_logprobs_val = (
self.token_ids_logprobs_val
)
logits_output.next_token_token_ids_logprobs_idx = (
self.token_ids_logprobs_idx
)
class OutputLogprobProcessor:
"""Output (decode) logprob processing: logprobs -> topk / token-ids /
sampled-token gather, attached onto LogitsProcessorOutput.
sampled-token gather, returned as a LogprobResult for the caller to
write back onto LogitsProcessorOutput.
Only logits/logprobs are needed here; sampler-side concerns (custom
logit processors, NaN sanitizing) are injected via ``preprocess_fn``.
"""
def attach_logprobs_to_output(
def compute_logprobs(
self,
logits_output: LogitsProcessorOutput,
logprobs: torch.Tensor,
top_logprobs_nums: List[int],
token_ids_logprobs: List[List[int]],
batch_next_token_ids: torch.Tensor,
):
) -> LogprobResult:
# clamp to avoid -inf values
logprobs.clamp_(min=torch.finfo(logprobs.dtype).min)
result = OutputLogprobsResult()
result = LogprobResult()
if any(x > 0 for x in top_logprobs_nums):
(
result.top_logprobs_val,
@@ -703,16 +696,15 @@ class OutputLogprobProcessor:
torch.arange(len(batch_next_token_ids), device=batch_next_token_ids.device),
batch_next_token_ids,
]
result.write_to(logits_output)
return result
def compute_logprobs_only(
self,
logits_output: LogitsProcessorOutput,
sampling_info: SamplingBatchInfo,
next_token_logits: Optional[torch.Tensor],
top_logprobs_nums: List[int],
token_ids_logprobs: List[List[int]],
preprocess_fn: Callable,
) -> None:
preprocess_fn: Callable[[torch.Tensor], torch.Tensor],
) -> Optional[LogprobResult]:
"""
Compute logprobs for requested token IDs without performing sampling.
@@ -720,9 +712,9 @@ class OutputLogprobProcessor:
but don't require next token generation.
"""
if logits_output.next_token_logits is None:
if next_token_logits is None:
logger.warning("No logits available for logprob computation")
return
return None
# Check if any requests actually need logprobs computation
needs_token_ids_logprobs = any(
@@ -732,15 +724,15 @@ class OutputLogprobProcessor:
needs_top_logprobs = any(x > 0 for x in top_logprobs_nums)
if not (needs_token_ids_logprobs or needs_top_logprobs):
return
return None
# Preprocess logits (custom processors and NaN handling)
logits = preprocess_fn(logits_output.next_token_logits, sampling_info)
logits = preprocess_fn(next_token_logits)
# Compute logprobs
logprobs = torch.nn.functional.log_softmax(logits, dim=-1)
result = OutputLogprobsResult()
result = LogprobResult()
# Handle top logprobs if requested
if needs_top_logprobs:
(
@@ -754,4 +746,4 @@ class OutputLogprobProcessor:
result.token_ids_logprobs_val,
result.token_ids_logprobs_idx,
) = get_token_ids_logprobs_batch_optimized(logprobs, token_ids_logprobs)
result.write_to(logits_output)
return result
python/sglang/srt/layers/sampler.py
+8 -8
View File
@@ -1,4 +1,5 @@
import logging
from functools import partial
from typing import Callable, Dict, List, Optional, Tuple
import torch
@@ -210,17 +211,16 @@ class Sampler(nn.Module):
)
del probs
# Attach logprobs to logits_output (in-place modification)
if return_logprob:
if SGLANG_RETURN_ORIGINAL_LOGPROB:
logprobs = original_logprobs
self.output_logprob_processor.attach_logprobs_to_output(
logits_output,
logprob_result = self.output_logprob_processor.compute_logprobs(
logprobs,
top_logprobs_nums,
token_ids_logprobs,
batch_next_token_ids,
)
logprob_result.write_output_to(logits_output)
self._sync_token_ids_across_tp(batch_next_token_ids, sampling_info)
@@ -494,17 +494,17 @@ class Sampler(nn.Module):
self,
logits_output: LogitsProcessorOutput,
sampling_info: SamplingBatchInfo,
return_logprob: bool,
top_logprobs_nums: List[int],
token_ids_logprobs: List[List[int]],
) -> None:
self.output_logprob_processor.compute_logprobs_only(
logits_output=logits_output,
sampling_info=sampling_info,
logprob_result = self.output_logprob_processor.compute_logprobs_only(
next_token_logits=logits_output.next_token_logits,
top_logprobs_nums=top_logprobs_nums,
token_ids_logprobs=token_ids_logprobs,
preprocess_fn=self._preprocess_logits,
preprocess_fn=partial(self._preprocess_logits, sampling_info=sampling_info),
)
if logprob_result is not None:
logprob_result.write_output_to(logits_output)
def register_sampler_backend(backend: str, factory: Callable[[], "Sampler"]) -> None:
python/sglang/srt/lora/backend/lmhead_mixing.py
+2 -2
View File
@@ -32,8 +32,8 @@ class LoRABackendLmHeadMixing:
Returns None if logprobs chunking is disabled or the pruned token
count does not exceed the logprobs chunk size.
"""
logprobs_chunk_size = envs.SGLANG_LOGITS_PROCESSER_CHUNK_SIZE.get()
enable_logprobs_chunk = envs.SGLANG_ENABLE_LOGITS_PROCESSER_CHUNK.get()
logprobs_chunk_size = envs.SGLANG_LOGPROB_CHUNK_SIZE.get()
enable_logprobs_chunk = envs.SGLANG_ENABLE_LOGPROB_CHUNK.get()
pruned_total = sum(pruned_lens)
if not enable_logprobs_chunk or pruned_total <= logprobs_chunk_size:
python/sglang/srt/model_executor/model_runner.py
-1
View File
@@ -1552,7 +1552,6 @@ class ModelRunner:
self.sampler.compute_logprobs_only(
logits_output,
forward_batch.sampling_info,
forward_batch.return_logprob,
forward_batch.top_logprobs_nums,
forward_batch.token_ids_logprobs,
)
test/manual/test_logprobs.py
+3 -3
View File
@@ -242,10 +242,10 @@ class TestLogprobsDense(unittest.TestCase):
chunk_size = kwargs.pop("chunk_size", None)
if chunk_size is not None:
print(f"Setting chunk size to {chunk_size}")
os.environ["SGLANG_ENABLE_LOGITS_PROCESSER_CHUNK"] = "True"
os.environ["SGLANG_LOGITS_PROCESSER_CHUNK_SIZE"] = str(chunk_size)
os.environ["SGLANG_ENABLE_LOGPROB_CHUNK"] = "True"
os.environ["SGLANG_LOGPROB_CHUNK_SIZE"] = str(chunk_size)
else:
os.environ["SGLANG_ENABLE_LOGITS_PROCESSER_CHUNK"] = "False"
os.environ["SGLANG_ENABLE_LOGPROB_CHUNK"] = "False"
# Create engine with merged configuration
engine_config = {**DEFAULT_ENGINE_CONFIG, **kwargs}
test/registered/core/test_srt_endpoint.py
+1 -1
View File
@@ -46,7 +46,7 @@ class TestSRTEndpoint(CustomTestCase):
# The tiny logprob chunk size routes this file's logprob tests
# through the multi-chunk stitching path (requests at or below 64
# rows still cover the non-chunked path).
env={**SERVER_ENV, "SGLANG_LOGITS_PROCESSER_CHUNK_SIZE": "64"},
env={**SERVER_ENV, "SGLANG_LOGPROB_CHUNK_SIZE": "64"},
other_args=(
"--enable-custom-logit-processor",
"--mem-fraction-static",
test/registered/lora/test_lora_hf_sgl_logprob_diff.py
+2 -2
View File
@@ -536,8 +536,8 @@ class TestLoRAHFSGLLogprobDifference(CustomTestCase):
"""
saved = {}
env_overrides = {
"SGLANG_ENABLE_LOGITS_PROCESSER_CHUNK": "true",
"SGLANG_LOGITS_PROCESSER_CHUNK_SIZE": "4",
"SGLANG_ENABLE_LOGPROB_CHUNK": "true",
"SGLANG_LOGPROB_CHUNK_SIZE": "4",
}
for key, val in env_overrides.items():
saved[key] = os.environ.get(key)
test/registered/lora/test_lora_moe_tp_logprob_diff.py
+1 -1
View File
@@ -161,7 +161,7 @@ class TestMoELoRATP2Logprobs(CustomTestCase):
prompts = MOE_LORA_TEST_PROMPTS[:3]
baseline = _run_sglang_moe_lora(tp_size=2, prompts=prompts)
torch.cuda.empty_cache()
with envs.SGLANG_LOGITS_PROCESSER_CHUNK_SIZE.override(16):
with envs.SGLANG_LOGPROB_CHUNK_SIZE.override(16):
chunked = _run_sglang_moe_lora(tp_size=2, prompts=prompts)
for i in range(len(prompts)):
test/registered/unit/layers/test_logprob_chunk_stitching.py
+7 -11
View File
@@ -108,25 +108,21 @@ class TestLogprobChunkStitching(CustomTestCase):
ref, ref_sampled = _run(proc, batch, False, 10**9)
got, got_sampled = _run(proc, batch, True, chunk_size)
label = f"specs={list(combo)} chunk={chunk_size}"
self.assertEqual(
ref.input_top_logprobs_val, got.input_top_logprobs_val, label
)
self.assertEqual(
ref.input_top_logprobs_idx, got.input_top_logprobs_idx, label
)
self.assertEqual(ref.top_logprobs_val, got.top_logprobs_val, label)
self.assertEqual(ref.top_logprobs_idx, got.top_logprobs_idx, label)
if with_token_ids:
self.assertEqual(
ref.input_token_ids_logprobs_val,
got.input_token_ids_logprobs_val,
ref.token_ids_logprobs_val,
got.token_ids_logprobs_val,
label,
)
self.assertEqual(
ref.input_token_ids_logprobs_idx,
got.input_token_ids_logprobs_idx,
ref.token_ids_logprobs_idx,
got.token_ids_logprobs_idx,
label,
)
torch.testing.assert_close(
ref.input_token_logprobs, got.input_token_logprobs, msg=label
ref.token_logprobs, got.token_logprobs, msg=label
)
torch.testing.assert_close(ref_sampled, got_sampled, msg=label)
self.assertGreater(tried, 1000)