Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom

[COMMON/JAX] Support sliding window on THD format #1327

Open
wants to merge 3 commits into
base: main
Choose a base branch
from
Open

[COMMON/JAX] Support sliding window on THD format #1327

wants to merge 3 commits into from

Conversation

zlsh80826
Copy link
Collaborator

Description

Support THD format with sliding window

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 refractor

Changes

  • [COMMON] Relax the condition for THD + sliding window
  • [COMMON] Replace THD with is_ragged and BSHD||SBHD with !is_ragged in nvte_get_fused_attn_backend
  • [JAX] Add some util functions for AttnMaskType

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

@zlsh80826
Copy link
Collaborator Author

/te-ci jax L1

@zlsh80826
Copy link
Collaborator Author

/te-ci

@zlsh80826
Add doc string
00c169f 
Signed-off-by: Reese Wang <[email protected]>
mgoldfarb-nvidia
mgoldfarb-nvidia reviewed Nov 11, 2024
View reviewed changes
transformer_engine/jax/attention.py
@@ -46,6 +46,30 @@ class AttnMaskType(Enum):
CAUSAL_BOTTOM_RIGHT_MASK = NVTE_Mask_Type.NVTE_CAUSAL_BOTTOM_RIGHT_MASK
PADDING_CAUSAL_BOTTOM_RIGHT_MASK = NVTE_Mask_Type.NVTE_PADDING_CAUSAL_BOTTOM_RIGHT_MASK

def is_causal(self):
Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Nice pythonic helpers!

cyanguwa
cyanguwa reviewed Nov 12, 2024
View reviewed changes
transformer_engine/common/fused_attn/fused_attn.cpp
// sliding window
((cudnn_runtime_version < 90200 && window_size_left == -1 &&
(window_size_right == -1 || window_size_right == 0)) ||
(cudnn_runtime_version >= 90200 &&
((window_size_left == -1 && (window_size_right == -1 || window_size_right == 0)) ||
((window_size_left >= 0 || window_size_left == -1) && window_size_right == 0 &&
(attn_mask_type == NVTE_Mask_Type::NVTE_CAUSAL_MASK ||
((!is_ragged && attn_mask_type == NVTE_Mask_Type::NVTE_CAUSAL_MASK) ||
(is_ragged && attn_mask_type == NVTE_Mask_Type::NVTE_PADDING_CAUSAL_MASK) ||
Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Could you factor the self-attn bottom_right combination into this condition as well?

(!is_ragged && (attn_mask_type == NVTE_Mask_Type::NVTE_CAUSAL_MASK || (attn_mask_type == NVTE_Mask_Type::NVTE_CAUSAL_BOTTOM_RIGHT_MASK && max_seqlen_q == max_seqlen_kv))) || (is_ragged && (attn_mask_type == NVTE_Mask_Type::NVTE_PADDING_CAUSAL_MASK || (attn_mask_type == NVTE_Mask_Type::NVTE_PADDING_CAUSAL_BOTTOM_RIGHT_MASK && max_seqlen_q == max_seqlen_kv)))

cyanguwa
cyanguwa approved these changes Nov 12, 2024
View reviewed changes
Copy link
Collaborator

@cyanguwa cyanguwa left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Other than that one comment, LGTM, thanks!

huanghua1994
huanghua1994 requested changes Nov 12, 2024
View reviewed changes
Copy link
Collaborator

@huanghua1994 huanghua1994 left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The changes to the code LGTM. The CI builds failed, need to check and re-run.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Reviewers

@huanghua1994 huanghua1994 huanghua1994 requested changes

@mgoldfarb-nvidia mgoldfarb-nvidia mgoldfarb-nvidia approved these changes

@cyanguwa cyanguwa cyanguwa approved these changes

Requested changes must be addressed to merge this pull request.

Assignees
No one assigned
Labels
None yet
Projects
None yet
Milestone
No milestone
Development

Successfully merging this pull request may close these issues.
4 participants
@zlsh80826 @cyanguwa @huanghua1994 @mgoldfarb-nvidia