hotfix: revert sampler CUDA Graph (#1242)

This commit is contained in:
Yineng Zhang
2024-08-28 21:16:47 +10:00
committed by GitHub
parent 184ae1c683
commit f25f4dfde5
33 changed files with 119 additions and 348 deletions

View File

@@ -24,7 +24,7 @@ from vllm.distributed import get_tensor_model_parallel_rank
from vllm.model_executor.layers.quantization.base_config import QuantizationConfig
from vllm.model_executor.model_loader.weight_utils import default_weight_loader
from sglang.srt.layers.logits_processor import LogitsProcessorOutput
from sglang.srt.layers.logits_processor import LogitProcessorOutput
from sglang.srt.model_executor.forward_batch_info import InputMetadata
from sglang.srt.models.llama2 import LlamaModel
@@ -65,7 +65,7 @@ class LlamaForClassification(nn.Module):
(input_metadata.batch_size, self.config.classification_out_size)
).to(input_ids.device)
return LogitsProcessorOutput(
return LogitProcessorOutput(
next_token_logits=scores,
next_token_logprobs=scores,
normalized_prompt_logprobs=scores,