[Tiny] Fix docstring in CP abstractions (#30579)

This commit is contained in:
Baizhou Zhang
2026-07-08 18:41:40 -07:00
committed by GitHub
parent 0ffed946f2
commit 2b3d9ad375
2 changed files with 8 additions and 8 deletions
+4 -4
View File
@@ -16,10 +16,10 @@
For ``cp_size = 4``, each rank owns every fourth token:
dp_attn_tp0: token0, token4, token8, token12, token16, ...
dp_attn_tp1: token1, token5, token9, token13, token17, ...
dp_attn_tp2: token2, token6, token10, token14, token18, ...
dp_attn_tp3: token3, token7, token11, token15, token19, ...
cp0: token0, token4, token8, token12, token16, ...
cp1: token1, token5, token9, token13, token17, ...
cp2: token2, token6, token10, token14, token18, ...
cp3: token3, token7, token11, token15, token19, ...
After all-gather, tokens are restored to the original order:
+4 -4
View File
@@ -17,10 +17,10 @@
For ``cp_size = 4``, each sequence is split into ``2 * cp_size`` blocks. Each
rank owns one early block and one late block:
dp_attn_tp0: block0, block7
dp_attn_tp1: block1, block6
dp_attn_tp2: block2, block5
dp_attn_tp3: block3, block4
cp0: block0, block7
cp1: block1, block6
cp2: block2, block5
cp3: block3, block4
After all-gather, the blocks are reranged back to their original order: