EngineX-Hygon/sglang
5
0
Fork 0
Code Issues Pull Requests Actions 7 Projects Releases Wiki Activity

Move sampler into CUDA graph (#1201)

Browse Source 
Co-authored-by: Yineng Zhang <me@zhyncs.com>
This commit is contained in:
Liangsheng Yin
2024-08-26 07:02:50 -07:00
committed by GitHub
parent 97589a60a2
commit 75ce37f401
28 changed files with 336 additions and 110 deletions
Download Patch File Download Diff File Expand all files Collapse all files

4
python/sglang/srt/models/llama_classification.py
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 LogitProcessorOutput
from sglang.srt.layers.logits_processor import LogitsProcessorOutput
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 LogitProcessorOutput(
return LogitsProcessorOutput(
next_token_logits=scores,
next_token_logprobs=scores,
normalized_prompt_logprobs=scores,