Add Triton kernels for vocab-parallel logprobs, entropy and KL

[jinyouzhi]

Follow-up to #8791 (RFC #8790). The vocab-parallel logprobs / entropy / KL ops ran each row chunk as a dozen separate eager fp32 kernels. On the same local shard they were about 2–2.4× slower than a dense log_softmax, which offsets most of the memory benefit when they are used for OPSD / RL losses (e.g. deepspeedai/DeepSpeedExamples#1019).

Changes

• New deepspeed/ops/triton_ops/vocab_parallel_softmax.py: Triton kernels with one program per row, looping over the local vocabulary in blocks. The module has no communication.
  • Local max and sum-exp are computed in a single read. They are then rescaled to the global max, so TP still needs only one MAX and one SUM all-reduce.
  • Fused forward and backward for entropy and KL, and a fused backward for logprobs (the top-k one-hot terms are added in fp32 before the single cast, as before).
• deepspeed/sequence/vocab_parallel_ops.py dispatches to these kernels on CUDA/XPU when the inputs are contiguous and deepspeed.HAS_TRITON is true. Otherwise it runs the existing eager path, which is unchanged. Collectives and the autograd structure are the same.
• Docs: a note in autotp-training.md.

Performance

Intel Arc Pro B70, bf16, 75968-column shard (Qwen vocab / TP2), forward + backward:

op	N	master	this PR	extra peak memory
reverse KL	8192	340.8 ms	15.5 ms	510 MiB → ~0
forward KL	8192	313.0 ms	15.5 ms	446 MiB → ~0
entropy	8192	205.5 ms	9.1 ms	335 MiB → ~0
logprobs (k=1)	8192	99.5 ms	7.2 ms	192 MiB → ~0
reverse KL	2048	85.1 ms	4.0 ms

• KL is within ~10% of the 608 GB/s memory-bandwidth floor.
• Logprobs with a large top-k is slower (k=256: 35 ms) but still about 3× faster than master.
• For reference, torch.compile of the eager chunk functions reached 26.9 ms for reverse KL (N=8192), so I chose hand-written kernels.

Tests

• New test_triton_kernels_match_eager_path (fp32 / bf16 × logprobs / entropy / forward KL / reverse KL). On an accelerator it compares the Triton outputs and gradients with the CPU eager path. Cases: multiple kernel blocks, partially masked vocabulary, all--inf rows, support mismatch with underflowing probability, large offsets, duplicate / ignored top-k ids, and a fully ignored row.
  • I checked that it fails when the kernel's support-mismatch inf is removed.
• The TP=2 tests (uneven and empty shards) now run through the Triton path.

Executed on 2× Intel Arc Pro B70 (XPU), torch 2.14.0.dev+xpu, Triton 3.8.0:

• tests/unit/v1/sequence_parallelism/test_vocab_parallel_ops.py: 41 passed.
• The existing single-process semantic tests, with XPU as the default device (so they run through Triton): 30 passed.

Update: single-pass entropy/KL forward (by @0z5a)

The tables above were measured on the first commit. @0z5a's e8038bed3 reads each logit once in the entropy/KL forward, using online weighted statistics. Both commits were re-measured with the same script on 1× B70: bf16, a 75968-column shard, forward + backward, median of 20 runs.

op	N	99f054192	head 97d4f66df
reverse KL	8192	18.7 ms	14.6 ms
forward KL	8192	18.7 ms	14.6 ms
entropy	8192	9.2 ms	7.1 ms
logprobs (k=1)	8192	7.4 ms	7.4 ms
reverse KL	2048	4.8 ms	3.9 ms

Extra peak memory, excluding inputs and gradients, is below 1 MiB on both commits.

Numerics vs the first commit. Compared on 1× and 2× B70 (TP=1, TP=2), fp32 and bf16. Cases: 512×32000, 256×151936, and 64×32000 with a +1000 offset and -inf masks. Each output and gradient was also compared with an fp64 CPU reference.

• logprobs: outputs and gradients are bitwise identical.
• entropy / KL: not bitwise identical, because the summation order changed. Forward outputs differ by at most 35 ULP. The forward is closer to fp64 on almost every case: e.g. fp32 entropy relative error drops from 3e-7 to 1.5e-7, and KL from about 6–9e-7 to 3–8e-7. Gradient errors stay at the same level, and bf16 gradient errors are unchanged (about 2e-3).
• There are no inf/NaN placement differences.

Contribution checklist

Related work. This is a performance follow-up to #8791 (RFC #8790). A search for vocab-parallel Triton kernels and entropy/KL found no overlapping PR.

Commands and results. Hardware and software: 2× Intel Arc Pro B70 (XPU, XCCL), Level Zero driver 1.15.39122 (compute runtime 26.27.39122.11), torch 2.14.0.dev20260804+xpu, Triton 3.8.0.

pytest -p no:randomly tests/unit/v1/sequence_parallelism/test_vocab_parallel_ops.py

At head 8f463213f: 49 passed on 2× B70; with DS_ACCELERATOR=cpu, 33 passed and the 16 Triton tests are skipped.

Compatibility matrix.

	Validated	Not validated
Precision	fp32, bf16	fp16
Parallelism	TP=1, TP=2 (uneven and empty shards)	TP>2 on device
Layout	Contiguous rows, plus row_layout strides from e8038bed3
Accelerator	XPU (Triton); CUDA (Triton, 2× H100, validated by @0z5a); CPU (eager fallback, Triton tests skipped)

ZeRO and offload do not interact with these ops.

AI assistance. This PR was developed with AI assistance (GitHub Copilot CLI). I reviewed every changed line and ran the tests above.

[0z5a]

Thanks for the Triton follow-up. I validated it on H100 and tried a few kernel-level changes on top of it; numbers and a patch are below.

#8838 on H100

• tests/unit/v1/sequence_parallelism/test_vocab_parallel_ops.py on 2× H100 (CUDA path, including the TP=2 uneven and empty-shard tests): 41 passed. The single-process semantic tests with CUDA as the default device (so they run through Triton): 39 passed.
• torch 2.12.1+cu130, Triton 3.7.1, driver 595.91.07.
• The kernels are already at the memory-bandwidth floor for their pass count. For reverse KL with a detached teacher, fwd+bwd reads the student and teacher shards three times and writes one gradient (7 full passes). At the measured copy bandwidth that is about 5.9 ms for a bf16 8192 × 75968 shard, and the kernels take 5.83 ms.
• The default launch configuration is already the best on H100. In a sweep of BLOCK 1024–8192 × 4–32 warps, nothing was more than 1% faster than 4096 / 16; some settings were up to 9% slower.

Suggested change: read each logit once in the entropy / KL forward

Since the kernels are bandwidth bound, the remaining lever is the number of passes. The forward currently reads each input twice: once for the local max and sum-exp, then again for the entropy or KL terms. The first pass can also accumulate the probability-weighted sum of logits centred on the running max: Σ e^(z−m)(z−m) for entropy, and Σ e^(a−m_a)((a−m_a)−(b−m_b)) for KL(A‖B). When the max grows, these sums are rescaled the same way as the sum-exp.

After the MAX all-reduce, each shard's sums are re-centred on the global max:

• H = log S − W / S
• KL = W / S_A − log S_A + log S_B

Centring on the max keeps a large common offset from cancelling in fp32; your large-offset test still passes. Support mismatch, meaning B masks a token that A supports, is a flag decided on the logits; it shares the MAX all-reduce. With TP the forward then needs one MAX and one SUM all-reduce instead of one MAX and two SUMs. The backward kernels are unchanged.

Separately, rows are addressed with up to two strides, so a slice such as logits[:, :-1] is read in place. Today _flatten copies such a slice with reshape in both forward and backward. Layouts the kernels cannot address, such as a strided vocabulary dimension, are still copied first, as before. Gradients are returned contiguous. I added one sentence about this to the tutorial note.

H100, fwd+bwd, 8192 × 75968 shard. The #8838 and patch columns are medians of 3 interleaved runs with nothing else on the node; the eager column comes from a single run.

op	eager (#8791)	#8838	#8838 + patch
reverse KL, bf16	133 ms	5.83 ms	4.48 ms
reverse KL, fp32 input	117 ms	11.5 ms	8.39 ms
entropy, bf16	79.6 ms	3.46 ms	2.69 ms
logprobs k=1, bf16	37.7 ms	2.80 ms	2.82 ms
reverse KL, bf16 logits[:, :-1]	141 ms	13.6 ms (+2.3 GiB peak)	4.47 ms (no extra peak)

• The forward alone is 1.7× (bf16) to 1.9× (fp32) faster. fp32 reverse KL is at the 5-pass bandwidth floor.
• The single-pass statistics keep up to six running values per lane, so they use 2048-wide blocks. With 4096 they ran about 1.2× slower on H100.
• logprobs is unchanged: within 1% at 8192 rows, and about 25 µs slower at 1024 rows.
• At an OPSD chunk of 1024 rows, bf16 reverse KL is about the same (0.82 → 0.81 ms), and fp32 input is 1.27× faster.

All outputs and gradients match the eager path within your test tolerances, with NaN/inf in the same places. Validation:

• The full test file passes on 2× H100: 49 passed, including the new test. The new test covers a sequence slice and a strided vocabulary; positions outside the view hold NaN, so reading them would fail.
• The single-process semantic tests pass on CUDA with fp32 default tolerances: 47 passed.
• A CPU run under TRITON_INTERPRET=1, including gloo TP=2 with uneven and empty shards, matches the eager path.

One side effect I noticed, on a 1024 × 75968 KL in both directions: with the patch, bf16 and fp32 inputs give bitwise identical outputs, and the bf16 gradient equals the round-to-nearest-even of the fp32 one. With the current kernels they differ by up to 4e-6. My guess is that the per-thread layout, and so the reduction order, depends on the dtype; I have not verified that. Either way it is harmless, just a note.

I have not run anything on XPU, so the 2048-wide statistics blocks are only tuned on H100. I opened the patch as a single-commit PR against your branch: jinyouzhi#5. Feel free to take all or part of it.

AI assistance was used for the implementation and analysis; I reviewed the changes and the measurements.

[jinyouzhi]

@0z5a Thank you for your optimization patch. I merged it.

[jinyouzhi]

One side effect I noticed, on a 1024 × 75968 KL in both directions: with the patch, bf16 and fp32 inputs give bitwise identical outputs, and the bf16 gradient equals the round-to-nearest-even of the fp32 one. With the current kernels they differ by up to 4e-6. My guess is that the per-thread layout, and so the reduction order, depends on the dtype; I have not verified that. Either way it is harmless, just a note.

It’s worth checking batch invariance when we move on to RL rollouts.