[Bugfix][CI] fix typos (#34934)
Signed-off-by: 1195343015 <1195343015@qq.com> Signed-off-by: Jiayi Yan <66017932+1195343015@users.noreply.github.com> Co-authored-by: Harry Mellor <19981378+hmellor@users.noreply.github.com>
This commit is contained in:
@@ -2540,7 +2540,7 @@ class MLACommonImpl(MLAAttentionImpl[M], Generic[M]):
|
||||
)
|
||||
# workspace
|
||||
# |------- N tokens --------|--------- N*dcp_size tokens ----------|
|
||||
# |<- use for loca_gather ->|<--------- use for allgather -------->|
|
||||
# |<- use for local_gather ->|<--------- use for allgather -------->|
|
||||
allgather_offset = workspace.shape[0] // (dcp_world_size + 1)
|
||||
assert allgather_offset * (dcp_world_size + 1) == workspace.shape[0]
|
||||
assert toks <= allgather_offset
|
||||
|
||||
Reference in New Issue
Block a user