SIGN IN SIGN UP

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