Context Parallel Causal Split

Problem

Balance causal attention work across context-parallel ranks: split the sequence into 2 * num_gpus chunks of ceil(seq_len / (2 * num_gpus)) tokens (the final chunk may be partial) and give rank r chunk r plus chunk 2 * num_gpus - 1 - r, pairing a light early chunk with a heavy late one.

The signature is context_parallel_chunks(seq_len, num_gpus). Return the token count owned by each rank.

Hint

Causal attention in early chunks is cheaper, so each rank takes one early and one late chunk.

Test cases

case 1

in: [16,2]

expected: [8,8]

case 2

in: [32,4]

expected: [8,8,8,8]

case 3

in: [24,3]

expected: [8,8,8]

case 4

in: [10,2]

expected: [4,6]

case 5

in: [64,4]

expected: [16,16,16,16]

Starter codePython

def context_parallel_chunks(seq_len, num_gpus):
    # Your code here
    pass

5 test cases · run your code against them in the editor