[JAX] Remove unnecessary SWA calculation in _segment_ids_pos_to_seqlens_offsets()

[KshitijLakhani]

Description

I noticed that in _segment_ids_pos_to_seqlens_offsets(), a temp / intermediate mask is created, the purpose of which is only to help with the calculation of seqlens and offsets which is then passed onto the cuDNN backend.
This intermediate attn_mask should only need to use the segment ids and pos to create the attn_mask followed by a layer of masking based on the application of the type of mask (causal or brcm) as that would influence the final seqlens and offsets generated.
However, SW mask calculation and application to attn_mask in no way should affect the final seqlens and offsets generated hence this is an unnecessary operation. I tested this out as well and it seems to check out.

I had put in a TODO for this observation and am addressing the same in this PR

Type of change

• Documentation change (change only to the documentation, either a fix or a new content)
• Bug fix (non-breaking change which fixes an issue)
• New feature (non-breaking change which adds functionality)
• Breaking change (fix or feature that would cause existing functionality to not work as expected)
• Infra/Build change
• Code refactoring

Checklist:

• I have read and followed the contributing guidelines
• The functionality is complete
• I have commented my code, particularly in hard-to-understand areas
• I have made corresponding changes to the documentation
• My changes generate no new warnings
• I have added tests that prove my fix is effective or that my feature works
• New and existing unit tests pass locally with my changes

[greptile-apps]

Greptile Overview

Greptile Summary

Removed unnecessary sliding window attention (SWA) mask calculation from the _segment_ids_pos_to_seqlens_offsets function in JAX attention module.

Key Changes:

• Removed 15 lines of SWA mask computation that was previously applied during seqlens/offsets calculation
• Added clarifying comment explaining that sliding window logic is handled by the cuDNN kernel, not during mask-to-seqlens conversion
• The window_size parameter is still passed through to tex.fused_attn_fwd where the actual SWA logic is applied

Analysis:
The change is correct because _segment_ids_pos_to_seqlens_offsets is only responsible for computing sequence lengths and offsets from segment IDs and positions. The sliding window attention logic is properly handled downstream by the cuDNN kernel in tex.fused_attn_fwd, which receives the window_size parameter directly. Applying the SWA mask during seqlens/offsets calculation was redundant and potentially incorrect, as it would constrain the metadata computation rather than the actual attention computation.

Confidence Score: 4/5

• This PR is safe to merge with minimal risk. The change removes redundant computation that was incorrectly applied during metadata calculation.
• Score of 4 reflects a well-reasoned optimization that removes unnecessary code. The change is architecturally correct - sliding window logic belongs in the attention kernel, not in the seqlens/offsets computation. However, the PR lacks a complete description, test evidence, and the checklist is unfilled, which prevents a score of 5.
• No files require special attention. The single-file change is straightforward and removes redundant code.

Important Files Changed

File Analysis

Filename	Score	Overview
transformer_engine/jax/attention.py	4/5	Removed unnecessary SWA mask calculation from _segment_ids_pos_to_seqlens_offsets function. The sliding window logic is handled by the cuDNN kernel, not in seqlens/offsets computation.

Sequence Diagram

sequenceDiagram
    participant User
    participant fused_attn
    participant SequenceDescriptor
    participant _segment_ids_pos_to_seqlens_offsets
    participant _mask_to_seqlens_offset
    participant tex.fused_attn_fwd
    
    User->>fused_attn: Call with QKV, mask_type, window_size
    fused_attn->>SequenceDescriptor: get_seqlens_and_offsets()
    SequenceDescriptor->>_segment_ids_pos_to_seqlens_offsets: segment_ids, segment_pos, attn_mask_type, window_size
    
    Note over _segment_ids_pos_to_seqlens_offsets: Check if fast path applies<br/>(causal without window_size)
    
    alt Fast path: causal without window
        _segment_ids_pos_to_seqlens_offsets->>_segment_ids_pos_to_seqlens_offsets: Use fast causal path
    else Slow path: BRCM or with window
        _segment_ids_pos_to_seqlens_offsets->>_segment_ids_pos_to_seqlens_offsets: Create segment_mask
        Note over _segment_ids_pos_to_seqlens_offsets: REMOVED: SWA mask calculation<br/>Window is handled by cuDNN kernel
        alt Bottom-right causal
            _segment_ids_pos_to_seqlens_offsets->>_segment_ids_pos_to_seqlens_offsets: Apply BRCM mask
        else Regular causal
            _segment_ids_pos_to_seqlens_offsets->>_segment_ids_pos_to_seqlens_offsets: Apply causal mask
        end
        _segment_ids_pos_to_seqlens_offsets->>_mask_to_seqlens_offset: attn_mask_with_id
        _mask_to_seqlens_offset-->>_segment_ids_pos_to_seqlens_offsets: q_seqlen, kv_seqlen, q_offset, kv_offset
    end
    
    _segment_ids_pos_to_seqlens_offsets-->>SequenceDescriptor: seqlens and offsets
    SequenceDescriptor-->>fused_attn: seqlens and offsets
    fused_attn->>tex.fused_attn_fwd: qkv, seqlens, offsets, window_size
    Note over tex.fused_attn_fwd: cuDNN kernel applies<br/>sliding window attention<br/>using window_size parameter
    tex.fused_attn_fwd-->>fused_attn: output
    fused_attn-->>User: result

Loading

[greptile-apps]

_{1 file reviewed, no comments}

_{Edit Code Review Agent Settings | Greptile}

[KshitijLakhani]

/te-ci jax

[mgoldfarb-nvidia]

LGTM

[huanghua1994]

LGTM

[KshitijLakhani]

CI pipeline passes: 38909687

[mingxu1067]

LGTM