Support logits soft-capping in the Pallas-Triton attention VJP. - #1121
Open
bzantium wants to merge 1 commit into
Open
Support logits soft-capping in the Pallas-Triton attention VJP.#1121bzantium wants to merge 1 commit into
bzantium wants to merge 1 commit into
Conversation
Support logits soft-capping in the Pallas-Triton attention VJP. The VJP raised `NotImplementedError` for `logits_soft_cap`, so models that use it had a forward pass but no backward pass on this kernel. Both backward kernels now apply the cap where the forward does - after the bias add, before masking - and chain `1 - tanh(u / c)^2` into `ds`. The factor is captured before masking, while the logits are finite; afterwards they are `mask_value` and the square would overflow. `dk` and `dbias` are both gradients wrt the pre-cap logits, so the factor is applied before either is computed. Verified on an A100 against an unfused fp32 reference differentiated by JAX, with soft cap 30 and pre-cap logits reaching 168, so the cap is active. Maximum relative error over dq, dk and dv is 4.3e-3 for causal attention and for two packed variants (`k_start` and a dense boolean mask, which agree to every digit); dbias is 3.4e-3. The uncapped path is unchanged, and the cost of the cap in the backward pass is 1.04x at head dim 128. Signed-off-by: bzantium <ryumin93@gmail.com>
4 tasks
bzantium
added a commit
to bzantium/maxtext
that referenced
this pull request
Jul 30, 2026
…GPU backend AttentionOp.cudnn_jax_flash_attention received decoder_segment_ids but its train branch called the jax cuDNN SDPA with mask_type=CAUSAL only. With packing on, which is the default for any real dataset, tokens attended across packed segment boundaries: no guard, no error, silently wrong loss and gradients. Against the dot_product reference on a packed batch, every token after the boundary diverged (max abs diff 5.4 versus 0.009 bf16 noise before it). On Hopper and newer, segment ids are now lowered to the cuDNN packed layout, per-segment q/kv seqlens plus offsets with PADDING_CAUSAL. cuDNN rejects that layout on older GPUs, so pre-Hopper falls back to a fused flash kernel that masks the boundaries and logs the downgrade, rather than attending across them. Not every packing implementation caps the run count at max_segments_per_seq (grain 'first_fit' and the hf pipeline pass max_sequences_per_bin; 'best_fit', 'concat_then_split' and tfds do not), so ids above the bound fold into the last segment; leaving them out of every segment is what previously handed cuDNN uninitialized memory for those tokens while still training on them. That fallback needs a fused kernel that runs where FA2 cannot, which is what the other two pieces provide: attention=cutlass_flash routes to FlashAttention-2 CUTLASS kernels through the optional flash-attn-jax package, with native GQA/MQA, sliding window, and packing lowered to FA2's varlen kernel with cu_seqlens built inside jit. Run starts are selected per row, so a row holding more segments than the budget can only merge two of its own rather than producing one sequence spanning unrelated rows. It stays opt-in rather than a declared dependency: flash-attn-jax pins jax<0.9 while MaxText needs jax>=0.10, and installing it downgrades jax until the package stops importing. CI therefore never exercises this backend. src/maxtext/kernels/attention/gpu_pallas_flash.py forks jax.experimental.pallas.ops.gpu.attention and adds sliding-window masking, gemma2-style logit soft cap with the matching 1 - tanh^2 chain rule, and chunked head_dim accumulation above 256. It exists because the tokamax kernel that attention=flash uses has no backward for soft caps or for head_dim > 256; fixes for both are in review upstream (openxla/tokamax#1121, AI-Hypercomputer#1122), after which this file should go away. backward_pass_impl='auto' always takes the fused backward: the XLA fallback materializes [b, h, q, kv], needing 12.5GiB against 3.3GiB at 8k tokens on an A100, though it is the more accurate of the two (~4e-3 against ~2e-2 on bf16 gradients) and stays available explicitly. A single dispatch serves attention=cutlass_flash and the pre-Hopper fallback: FA2 where eligible, pallas otherwise, so soft caps, fp32 and head_dim > 256 all land on a kernel that handles them instead of being rejected. attention=flash on GPU is left on the Tokamax path from AI-Hypercomputer#4465 and is not touched here. Context parallelism is rejected for these backends, since the kernels mask from shard-local positions; cudnn_flash_te remains the kernel for that. Packing with dropout_rate > 0 is rejected on pre-Hopper, where the fallback kernels have no dropout and would otherwise train with it on one GPU generation and without it on another. Tests: tests/unit/gpu_attention_packing_test.py compares each kernel's forward output and gradient wrt query against dot_product on an unpacked control, a packed batch, and a packed batch with a tile-unaligned boundary, reporting error either side of the boundary so cross-segment leaks are visible. On A100 all kernels match within 2 percent; the cutlass cases skip without flash-attn-jax. Known gaps: the Hopper-native packed path needs jax-ml/jax#39521 to be correct under GQA/MQA, where kv_offsets are currently scaled by the query head count; attention=cutlass_flash is reviewed but never executed for the dependency reason above; and the TE unfused-fallback crash on sm80 is a separate upstream issue. FIXES: AI-Hypercomputer#4476
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Support logits soft-capping in the Pallas-Triton attention VJP.
The VJP raised
NotImplementedErrorforlogits_soft_cap, so models that use it had a forward pass but no backward pass on this kernel. Both backward kernels now apply the cap where the forward does, after the bias add and before masking, and chain1 - tanh(u / c)^2intods. The factor is captured before masking, while the logits are finite; afterwards they aremask_valueand the square would overflow.dkanddbiasare both gradients wrt the pre-cap logits, so the factor is applied before either is computed.Verified on an A100 against an unfused fp32 reference differentiated by JAX, with soft cap 30 and pre-cap logits reaching 168, so the cap is active. Maximum relative error over dq, dk and dv is 4.3e-3 for causal attention and for two packed variants (
k_startand a dense boolean mask, which agree to every digit); dbias is 3.4e-3. The uncapped path is unchanged. No runtime number for the cap: the only machine available to me is shared, and the same config varied by up to 8x between runs there.PallasTritonFlashAttentionWithPallasTritonVjpTestno longer excludeslogits_soft_capfrom the cases it expects to be supported, so the existing harness now covers it.dropout_maskremains excluded.