diff --git a/python/sglang/srt/layers/cp/interleave.py b/python/sglang/srt/layers/cp/interleave.py index 7db762ab9..91769e223 100644 --- a/python/sglang/srt/layers/cp/interleave.py +++ b/python/sglang/srt/layers/cp/interleave.py @@ -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: diff --git a/python/sglang/srt/layers/cp/zigzag.py b/python/sglang/srt/layers/cp/zigzag.py index 3e5af00ba..717962a71 100644 --- a/python/sglang/srt/layers/cp/zigzag.py +++ b/python/sglang/srt/layers/cp/zigzag.py @@ -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: