[JAX] Support Ring Attention (Context Parallelism) #1059
Conversation
mingxu1067
force-pushed
the
mingh/ring_attn_primitive
branch
2 times, most recently
from
a42f6a2
to
7a147c8
Compare
mgoldfarb-nvidia
force-pushed
the
mingh/ring_attn_primitive
branch
7 times, most recently
from
f50f6f0
to
9fb3dc3
Compare
mgoldfarb-nvidia
force-pushed
the
mingh/ring_attn_primitive
branch
from
9fb3dc3
to
c67cc11
Compare
|
/te-ci jax L1 |
|
|
||
| @staticmethod | ||
| @cache | ||
| def use_scanloop(): |
There was a problem hiding this comment.
This is a bit ugly but a necessity of using scan loop which is preferred as it gives XLA more control over unrolling. We won't prevent not using scan but it currently warns the user to update their XLA flags. NVTE_FUSED_RING_ATTENTION_USE_SCAN if undefined will default to 1.
Hopefully --xla_experimental_ignore_channel_id becomes default in XLA:GPU at some point.
|
|
||
| # Combine KV tensors if separate for better permute scheduling and performance. | ||
| # Eventually XLA should perform this automatically. | ||
| kv = helper.stack_kv(k, v) |
There was a problem hiding this comment.
XLA:GPU does a poor job scheduling 2+ permutes with overlap so we manually combine K and V into our combined KV format. There is a plan to add functionality into XLA to do this fusion automatically.
mgoldfarb-nvidia
force-pushed
the
mingh/ring_attn_primitive
branch
5 times, most recently
from
e35fd0d
to
51922dc
Compare
|
/te-ci jax L1 |
Signed-off-by: Michael Goldfarb <mgoldfarb@nvidia.com> Signed-off-by: Ming Huang <mingh@nvidia.com>
mgoldfarb-nvidia
force-pushed
the
mingh/ring_attn_primitive
branch
from
d04ef7f
to
eba6d4e
Compare
|
/te-ci jax L1 |
for more information, see https://pre-commit.ci
| return lax.cond((idx <= cp_rank), no_mask_compute, skip_compute) | ||
|
|
||
| output_per_step, softmax_aux_per_step = lax.cond( | ||
| idx == 0, causal_mask_compute, jax_cond_wrap |
There was a problem hiding this comment.
Could you clarify why causal computation is applied to idx == 0? I initially thought that the causal mask would be needed for the trailing block instead.
There was a problem hiding this comment.
Here is a picture of what happens in both unbalanced and load balanced case:
Each GPU_i starts with Q_i and KV_i at the start. These are the causal masked parts along the diagonal. All subsequence iterations are either skipped or non-masked partial computations. In the load balanced case this is where the "half KV" and "half Q" cases come from.
| for idx in range(cp_size): | ||
| output = output + output_per_steps[idx].astype(jnp.float32) * jnp.exp( | ||
| softmax_aux_per_steps[idx] - softmax_aux | ||
| ).transpose(0, 2, 1, 3) |
There was a problem hiding this comment.
Would it be feasible to transpose softmax_aux inside scan_block to allow for pipelining the transpose?
There was a problem hiding this comment.
If it doesn't affect the performance lots we can keep the transpose here
Description
Add ring attention as an additional context parallel strategy to Jax fused attention API.
Type of change
Changes
Adds a new context parallel strategy and unit tests.
Checklist: