"git@developer.sourcefind.cn:change/sglang.git" did not exist on "45360b2fa9337b41576a1434d302663ec7fbc51d"
Unverified Commit 4455b26e authored by DavidChan's avatar DavidChan Committed by GitHub
Browse files

[Bug fixed] fixed the crash when enable the dp-attention on the single card (#3958)

parent c553e160
...@@ -848,12 +848,12 @@ class DeepseekV2AttentionMLA(nn.Module): ...@@ -848,12 +848,12 @@ class DeepseekV2AttentionMLA(nn.Module):
def all_gather( def all_gather(
input_tensor: torch.Tensor, forward_batch: ForwardBatch, rank, world_size, group input_tensor: torch.Tensor, forward_batch: ForwardBatch, rank, world_size, group
): ):
if world_size == 1:
return input_tensor
all_lens = forward_batch.global_num_tokens_cpu all_lens = forward_batch.global_num_tokens_cpu
max_len = max(forward_batch.global_num_tokens_cpu) max_len = max(forward_batch.global_num_tokens_cpu)
if world_size == 1:
return input_tensor, 0, all_lens[0]
padded_tensor = torch.nn.functional.pad( padded_tensor = torch.nn.functional.pad(
input_tensor, (0, 0, 0, max_len - input_tensor.shape[0]) input_tensor, (0, 0, 0, max_len - input_tensor.shape[0])
) )
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment