feat(speculative): add Ulysses context parallelism for the EAGLE-3 draft (#3356)
* feat(speculative): add Ulysses context parallelism for the EAGLE-3 draft The draft CP added in #2918 is a hand-written causal ring that is single-document causal, so it is gated off with sequence packing. Add a second CP mode, "ulysses", that redistributes block-0 attention with an all-to-all: each cp rank holds the full sequence for a subset of the heads, runs one dense or packed varlen FlashAttention, then all-to-alls the output back. The per-position TTT diagonals stay shard-local and merge into block-0 via the online-softmax identity. Because block-0 sees the whole sequence, cu_seqlens document boundaries carry through, so cp_size>1 composes with packed_sequence_size>0. In this mode the frozen target runs its normal block-causal forward (unsharded) and the recipe slices its hidden states to each draft shard. _CachedUlyssesAttention is a hand-written autograd.Function whose backward runs the flash backward against the merged joint-softmax out/lse (reusing the ring's dense wrappers and the varlen private entry points), so the block-0 q/k gradients are correct with TTT diagonals present. Select with cp_mode=ulysses (default "ring" unchanged). Requires the draft num_attention_heads and seq_length/packed_sequence_size divisible by cp_size, and flash-attn 2.8.x. Signed-off-by: khazic <khazzz1c@gmail.com> * fix(speculative): pass a contiguous lse to the Ulysses block-0 backward _gather_lse returned a transposed view, and flash-attn's _flash_attn_backward forwards softmax_lse to the kernel without the maybe_contiguous it applies to q/k/v/out/dout. The kernel then read the merged log-sum-exp with the wrong strides, so every block-0 query gradient was wrong once the Ulysses degree exceeded 1. At degree 1 the gather is a no-op and the tensor stayed contiguous, which is why only the world-size-2 gradient test caught it. Signed-off-by: khazic <khazzz1c@gmail.com> --------- Signed-off-by: khazic <khazzz1c@gmail.com> Signed-off-by: HuiyingLi <willwin.lee@gmail.com> Co-authored-by: HuiyingLi <willwin.lee@gmail.com>
K
khazzz1c committed
e2a3ee6499f948c9df1ac28bda28aeaf2c15272a
Parent: 116eec7
Committed by GitHub <noreply@github.com>
on 8/23/2026, 8:37:12 AM