[Minor] more code cleanup (#4077)

2025-03-04 21:23:47 -08:00
parent 4725e3f652
commit e074d84e5b
15 changed files with 123 additions and 31 deletions
--- a/python/sglang/srt/models/deepseek_v2.py
+++ b/python/sglang/srt/models/deepseek_v2.py
@@ -818,8 +818,8 @@ def all_gather(
    if world_size == 1:
        return input_tensor

-    all_lens = forward_batch.global_num_tokens
-    max_len = max(forward_batch.global_num_tokens)
+    all_lens = forward_batch.global_num_tokens_cpu
+    max_len = max(forward_batch.global_num_tokens_cpu)

    padded_tensor = torch.nn.functional.pad(
        input_tensor, (0, 0, 0, max_len - input_tensor.shape[0])