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

Add dtype for more operations (#1705)

Browse Source
This commit is contained in:
Lianmin Zheng
2024-10-18 12:18:15 -07:00
committed by GitHub
parent 6d0fa73ece
commit 392f2863c8
3 changed files with 5 additions and 4 deletions
Download Patch File Download Diff File Expand all files Collapse all files

4
python/sglang/srt/managers/schedule_batch.py
View File

@@ -537,8 +537,8 @@ class ScheduleBatch:
# Set fields # Set fields
with out_cache_loc.device: with out_cache_loc.device:
self.input_ids = torch.tensor(sum(input_ids, []), dtype=torch.int32) self.input_ids = torch.tensor(sum(input_ids, []), dtype=torch.int32)
self.req_pool_indices = torch.tensor(req_pool_indices) self.req_pool_indices = torch.tensor(req_pool_indices, dtype=torch.int32)
self.seq_lens = torch.tensor(seq_lens) self.seq_lens = torch.tensor(seq_lens, dtype=torch.int32)
self.extend_num_tokens = extend_num_tokens self.extend_num_tokens = extend_num_tokens
self.out_cache_loc = out_cache_loc self.out_cache_loc = out_cache_loc

3
python/sglang/srt/model_executor/forward_batch_info.py
View File

@@ -145,8 +145,9 @@ class ForwardBatch:
], ],
axis=0, axis=0,
), ),
dtype=torch.int64,
device=device, device=device,
).to(torch.int64) )
ret.image_inputs = batch.image_inputs ret.image_inputs = batch.image_inputs
ret.extend_seq_lens = torch.tensor(batch.extend_seq_lens, device=device) ret.extend_seq_lens = torch.tensor(batch.extend_seq_lens, device=device)

2
python/sglang/srt/sampling/sampling_batch_info.py
View File

@@ -57,7 +57,7 @@ class SamplingBatchInfo:
[r.sampling_params.top_p for r in reqs], dtype=torch.float [r.sampling_params.top_p for r in reqs], dtype=torch.float
) )
top_ks = torch.tensor( top_ks = torch.tensor(
[r.sampling_params.top_k for r in reqs], dtype=torch.int [r.sampling_params.top_k for r in reqs], dtype=torch.int32
) )
min_ps = torch.tensor( min_ps = torch.tensor(
[r.sampling_params.min_p for r in reqs], dtype=torch.float [r.sampling_params.min_p for r in reqs], dtype=torch.float