[feat] Support swa for npu#135

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

Open
POI-WX wants to merge 9 commits into
base:main
Choose a base branch
from
Open

[feat] Support swa for npu #135

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
9 commits
Select commit Hold shift + click to select a range
49d1a9d
fix bug of rope for npu and support the interleaved option
Oct 21, 2024
0eebc4b
optimize when interleaved is True
Oct 22, 2024
d69d17b
use fused rope op from mindspeed
Oct 22, 2024
6681c0a
optimize
Oct 23, 2024
ee0baeb
using fused rms norm op with 8.0 RC3 software stack on npu
Jan 8, 2025
5565fbc
optimize
Jan 9, 2025
c39b8c5
modify atol and rtol
Jan 9, 2025
a346f6c
use the global attention mask and support sliding window local attention
Jan 21, 2025
4b133df
optimize code for npu
Jan 21, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
221 changes: 114 additions & 107 deletions deeplink_ext/internevo_ops/_flash_attention_npu.py
View file
Open in desktop
Original file line numberDiff line numberDiff line change
Expand Up@@ -12,6 +12,43 @@
"flash_attn_varlen_kvpacked_func",
]

# construct a global attention mask for npu
_GLOBAL_ATTN_MASK = None


def get_attention_mask(seqlen, causal, window_size):
global _GLOBAL_ATTN_MASK

if _GLOBAL_ATTN_MASK is not None:
return _GLOBAL_ATTN_MASK

# causal attention
if causal:
if seqlen > 2048:
_GLOBAL_ATTN_MASK = torch.triu(
torch.ones([2048, 2048], dtype=bool, device=torch.npu.current_device()),
diagonal=1,
)
else:
_GLOBAL_ATTN_MASK = torch.triu(
torch.ones(
[seqlen, seqlen], dtype=bool, device=torch.npu.current_device()
),
diagonal=1,
)

# sliding window attention
if window_size[0] >= 0 or window_size[1] >= 0:
_GLOBAL_ATTN_MASK = torch.tril(
torch.ones([seqlen, seqlen], dtype=bool, device=torch.npu.current_device()),
diagonal=-((seqlen - 1 if window_size[0] < 0 else window_size[0]) + 1),
) + torch.triu(
torch.ones([seqlen, seqlen], dtype=bool, device=torch.npu.current_device()),
diagonal=(seqlen - 1 if window_size[1] < 0 else window_size[1]) + 1,
)

return _GLOBAL_ATTN_MASK


def flash_attn_func(
q,
Expand All@@ -32,22 +69,15 @@ def flash_attn_func(
seqlen_k = k.shape[1]
head_num = q.shape[-2]

if seqlen_q == seqlen_k and seqlen_q < 2048 and seqlen_k < 2048:
sparse_mode = 0
else:
sparse_mode = 2
assert seqlen_q == seqlen_k, "Npu currently only supports seqlen_q = seqlen_k."
attention_mask = get_attention_mask(seqlen_q, causal, window_size)
sparse_mode = 0 if attention_mask is None or seqlen_q <= 2048 else 4

seqlen_q = min(seqlen_q, 2048)
seqlen_k = min(seqlen_k, 2048)

attention_mask = (
torch.triu(
torch.ones([seqlen_q, seqlen_k], dtype=torch.bool, device=q.device),
diagonal=1,
)
if causal
else None
)
pre_tokens = seqlen_q - 1
next_tokens = 0
if window_size[0] >= 0 or window_size[1] >= 0:
pre_tokens = seqlen_q - 1 if window_size[0] < 0 else window_size[0]
next_tokens = seqlen_q - 1 if window_size[1] < 0 else window_size[1]

out = torch_npu.npu_fusion_attention(
q,
Expand All@@ -58,8 +88,8 @@ def flash_attn_func(
atten_mask=attention_mask,
scale=softmax_scale,
keep_prob=1 - dropout_p,
pre_tockens=seqlen_q,
next_tockens=0,
pre_tockens=pre_tokens,
next_tockens=next_tokens,
sparse_mode=sparse_mode,
)[0]

Expand DownExpand Up@@ -89,22 +119,18 @@ def flash_attn_varlen_func(

cu_seqlens_q = cu_seqlens_q[1:].tolist()
cu_seqlens_k = cu_seqlens_k[1:].tolist()
seqlen_q = min(max_seqlen_q, 2048)
seqlen_k = min(max_seqlen_k, 2048)

if max_seqlen_q < 2048:
sparse_mode = 0
else:
sparse_mode = 2

attention_mask = (
torch.triu(
torch.ones([seqlen_q, seqlen_k], dtype=torch.bool, device=q.device),
diagonal=1,
)
if causal
else None
)

assert (
max_seqlen_q == max_seqlen_k
), "Npu currently only supports max_seqlen_q = max_seqlen_k."
attention_mask = get_attention_mask(max_seqlen_q, causal, window_size)
sparse_mode = 0 if attention_mask is None or max_seqlen_q <= 2048 else 4

pre_tokens = max_seqlen_q - 1
next_tokens = 0
if window_size[0] >= 0 or window_size[1] >= 0:
pre_tokens = max_seqlen_q - 1 if window_size[0] < 0 else window_size[0]
next_tokens = max_seqlen_q - 1 if window_size[1] < 0 else window_size[1]

out = torch_npu.npu_fusion_attention(
q,
Expand All@@ -114,8 +140,8 @@ def flash_attn_varlen_func(
"TND",
atten_mask=attention_mask,
scale=softmax_scale,
pre_tockens=q.shape[0], # seq_len
next_tockens=0, # 0
pre_tockens=pre_tokens,
next_tockens=next_tokens,
keep_prob=1 - dropout_p,
sparse_mode=sparse_mode,
actual_seq_qlen=cu_seqlens_q,
Expand DownExpand Up@@ -143,21 +169,14 @@ def flash_attn_qkvpacked_func(
seqlen_qkv = qkv.shape[1]
head_num = q.shape[-2]

if seqlen_qkv < 2048:
sparse_mode = 0
else:
sparse_mode = 2

seqlen_qkv = min(qkv.shape[1], 2048)
attention_mask = get_attention_mask(seqlen_qkv, causal, window_size)
sparse_mode = 0 if attention_mask is None or seqlen_qkv <= 2048 else 4

attention_mask = (
torch.triu(
torch.ones([seqlen_qkv, seqlen_qkv], dtype=torch.bool, device=q.device),
diagonal=1,
)
if causal
else None
)
pre_tokens = seqlen_qkv - 1
next_tokens = 0
if window_size[0] >= 0 or window_size[1] >= 0:
pre_tokens = seqlen_qkv - 1 if window_size[0] < 0 else window_size[0]
next_tokens = seqlen_qkv - 1 if window_size[1] < 0 else window_size[1]

out = torch_npu.npu_fusion_attention(
q,
Expand All@@ -168,8 +187,8 @@ def flash_attn_qkvpacked_func(
atten_mask=attention_mask,
scale=softmax_scale,
keep_prob=1 - dropout_p,
pre_tockens=seqlen_qkv,
next_tockens=0,
pre_tockens=pre_tokens,
next_tockens=next_tokens,
sparse_mode=sparse_mode,
)[0]

Expand All@@ -192,26 +211,19 @@ def flash_attn_kvpacked_func(
k = kv[:, :, 0]
v = kv[:, :, 1]

s0 = q.shape[1]
s1 = kv.shape[1]
seqlen_q = q.shape[1]
seqlen_kv = kv.shape[1]
head_num = q.shape[-2]

if s0 == s1 and s0 < 2048 and s1 < 2048:
sparse_mode = 0
else:
sparse_mode = 2

seqlen_q = min(s0, 2048)
seqlen_k = min(s1, 2048)
assert seqlen_q == seqlen_kv, "Npu currently only supports seqlen_q = seqlen_kv."
attention_mask = get_attention_mask(seqlen_q, causal, window_size)
sparse_mode = 0 if attention_mask is None or seqlen_q <= 2048 else 4

attention_mask = (
torch.triu(
torch.ones([seqlen_q, seqlen_k], dtype=torch.bool, device=q.device),
diagonal=1,
)
if causal
else None
)
pre_tokens = seqlen_q - 1
next_tokens = 0
if window_size[0] >= 0 or window_size[1] >= 0:
pre_tokens = seqlen_q - 1 if window_size[0] < 0 else window_size[0]
next_tokens = seqlen_q - 1 if window_size[1] < 0 else window_size[1]

out = torch_npu.npu_fusion_attention(
q,
Expand All@@ -222,8 +234,8 @@ def flash_attn_kvpacked_func(
atten_mask=attention_mask,
scale=softmax_scale,
keep_prob=1 - dropout_p,
pre_tockens=seqlen_k,
next_tockens=0,
pre_tockens=pre_tokens,
next_tockens=next_tokens,
sparse_mode=sparse_mode,
)[0]

Expand All@@ -247,32 +259,30 @@ def flash_attn_varlen_qkvpacked_func(
q = qkv[:, 0]
k = qkv[:, 1]
v = qkv[:, 2]
n = q.shape[1]
if max_seqlen > 2048:
sparse_mode = 2
else:
sparse_mode = 0
head_num = q.shape[1]

cu_seqlens_q = cu_seqlens[1:].tolist()
cu_seqlens_k = cu_seqlens[1:].tolist()
seqlen = min(max_seqlen, 2048)
attention_mask = (
torch.triu(
torch.ones([seqlen, seqlen], dtype=torch.bool, device=q.device),
diagonal=1,
)
if causal
else None
)

attention_mask = get_attention_mask(max_seqlen, causal, window_size)
sparse_mode = 0 if attention_mask is None or max_seqlen <= 2048 else 4

pre_tokens = max_seqlen - 1
next_tokens = 0
if window_size[0] >= 0 or window_size[1] >= 0:
pre_tokens = max_seqlen - 1 if window_size[0] < 0 else window_size[0]
next_tokens = max_seqlen - 1 if window_size[1] < 0 else window_size[1]

out = torch_npu.npu_fusion_attention(
q,
k,
v,
n,
head_num,
"TND",
atten_mask=attention_mask,
scale=softmax_scale,
pre_tockens=q.shape[0], # seq_len
next_tockens=0, # 0
pre_tockens=pre_tokens,
next_tockens=next_tokens,
keep_prob=1 - dropout_p,
sparse_mode=sparse_mode,
actual_seq_qlen=cu_seqlens_q,
Expand DownExpand Up@@ -300,35 +310,32 @@ def flash_attn_varlen_kvpacked_func(
softmax_scale = q.shape[-1] ** (-0.5)
k = kv[:, 0]
v = kv[:, 1]
n = q.shape[1]
head_num = q.shape[1]
cu_seqlens_q = cu_seqlens_q[1:].tolist()
cu_seqlens_k = cu_seqlens_k[1:].tolist()
seqlen_q = min(max_seqlen_q, 2048)
seqlen_k = min(max_seqlen_k, 2048)

if max_seqlen_q > 2048:
sparse_mode = 2
else:
sparse_mode = 0

attention_mask = (
torch.triu(
torch.ones([seqlen_q, seqlen_k], dtype=torch.bool, device=q.device),
diagonal=1,
)
if causal
else None
)

assert (
max_seqlen_q == max_seqlen_k
), "Npu currently only supports max_seqlen_q = max_seqlen_k."
attention_mask = get_attention_mask(max_seqlen_q, causal, window_size)
sparse_mode = 0 if attention_mask is None or max_seqlen_q <= 2048 else 4

pre_tokens = max_seqlen_q - 1
next_tokens = 0
if window_size[0] >= 0 or window_size[1] >= 0:
pre_tokens = max_seqlen_q - 1 if window_size[0] < 0 else window_size[0]
next_tokens = max_seqlen_k - 1 if window_size[1] < 0 else window_size[1]

out = torch_npu.npu_fusion_attention(
q,
k,
v,
n,
head_num,
"TND",
atten_mask=attention_mask,
scale=softmax_scale,
pre_tockens=q.shape[0], # seq_len
next_tockens=0, # 0
pre_tockens=pre_tokens,
next_tockens=next_tokens,
keep_prob=1 - dropout_p,
sparse_mode=sparse_mode,
actual_seq_qlen=cu_seqlens_q,
Expand Down
Loading
, 'i'); if (__m === '*' || __re.test(location.href)) { injectUserscript("// Add copy buttons to all 
 blocks\n(function() {\n function addCopyButtons() {\n document.querySelectorAll('pre code').forEach(function(codeBlock) {\n if (codeBlock.parentElement.hasAttribute('data-copy-added')) return;\n codeBlock.parentElement.setAttribute('data-copy-added', 'true');\n \n var btn = document.createElement('button');\n btn.textContent = 'Copy';\n btn.style.cssText = 'position:absolute;top:4px;right:4px;padding:2px 8px;font-size:11px;background:#4ecdc4;border:none;border-radius:4px;color:#1a1a2e;cursor:pointer;opacity:0.7;transition:opacity 0.2s;';\n btn.onmouseover = function() { this.style.opacity = '1'; };\n btn.onmouseout = function() { this.style.opacity = '0.7'; };\n btn.onclick = function() {\n navigator.clipboard.writeText(codeBlock.textContent).then(function() {\n btn.textContent = 'Copied!';\n setTimeout(function() { btn.textContent = 'Copy'; }, 1500);\n });\n };\n codeBlock.parentElement.style.position = 'relative';\n codeBlock.parentElement.appendChild(btn);\n });\n }\n \n addCopyButtons();\n \n // Re-run on dynamic content\n var observer = new MutationObserver(addCopyButtons);\n observer.observe(document.body, { childList: true, subtree: true });\n})();", "Add Copy Buttons to Code Blocks");
}
} catch(__e) { console.warn('[Userscript:Add Copy Buttons to Code Blocks]', __e); }
})();
(function(){
try {
var __m = "github.com";
var __re = new RegExp('^' + "github\\.com" + '

Skip to content
[feat] Support swa for npu#135

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

 Open

POI-WX
wants to merge
9
commits into

base:main
Choose a base branch

from

 Open


[feat] Support swa for npu
#135
Show file tree
Hide file tree
 Changes from all commits
Commits

Show all changes


9 commits


Select commit
Hold shift + click to select a range
49d1a9d

fix bug of rope for npu and support the interleaved option

Oct 21, 2024
0eebc4b

optimize when interleaved is True

Oct 22, 2024
d69d17b

use fused rope op from mindspeed

Oct 22, 2024
6681c0a

optimize

Oct 23, 2024
ee0baeb

using fused rms norm op with 8.0 RC3 software stack on npu

Jan 8, 2025
5565fbc

optimize

Jan 9, 2025
c39b8c5

modify atol and rtol

Jan 9, 2025
a346f6c

use the global attention mask and support sliding window local attention

Jan 21, 2025
4b133df

optimize code for npu

Jan 21, 2025

File filter

Filter by extension
Filter by extension
Conversations

Failed to load comments. 
Loading
Jump to

Jump to file


Failed to load files. 
Loading
 Diff view

 Diff view


221 changes: 114 additions & 107 deletions
deeplink_ext/internevo_ops/_flash_attention_npu.py

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -12,6 +12,43 @@
 "flash_attn_varlen_kvpacked_func",
]

# construct a global attention mask for npu
_GLOBAL_ATTN_MASK = None


def get_attention_mask(seqlen, causal, window_size):
 global _GLOBAL_ATTN_MASK

 if _GLOBAL_ATTN_MASK is not None:
 return _GLOBAL_ATTN_MASK

 # causal attention
 if causal:
 if seqlen > 2048:
 _GLOBAL_ATTN_MASK = torch.triu(
 torch.ones([2048, 2048], dtype=bool, device=torch.npu.current_device()),
 diagonal=1,
 )
 else:
 _GLOBAL_ATTN_MASK = torch.triu(
 torch.ones(
 [seqlen, seqlen], dtype=bool, device=torch.npu.current_device()
 ),
 diagonal=1,
 )

 # sliding window attention
 if window_size[0] >= 0 or window_size[1] >= 0:
 _GLOBAL_ATTN_MASK = torch.tril(
 torch.ones([seqlen, seqlen], dtype=bool, device=torch.npu.current_device()),
 diagonal=-((seqlen - 1 if window_size[0] < 0 else window_size[0]) + 1),
 ) + torch.triu(
 torch.ones([seqlen, seqlen], dtype=bool, device=torch.npu.current_device()),
 diagonal=(seqlen - 1 if window_size[1] < 0 else window_size[1]) + 1,
 )

 return _GLOBAL_ATTN_MASK


def flash_attn_func(
 q,
Expand All@@ -32,22 +69,15 @@ def flash_attn_func(
 seqlen_k = k.shape[1]
 head_num = q.shape[-2]

 if seqlen_q == seqlen_k and seqlen_q < 2048 and seqlen_k < 2048:
 sparse_mode = 0
 else:
 sparse_mode = 2
 assert seqlen_q == seqlen_k, "Npu currently only supports seqlen_q = seqlen_k."
 attention_mask = get_attention_mask(seqlen_q, causal, window_size)
 sparse_mode = 0 if attention_mask is None or seqlen_q <= 2048 else 4

 seqlen_q = min(seqlen_q, 2048)
 seqlen_k = min(seqlen_k, 2048)

 attention_mask = (
 torch.triu(
 torch.ones([seqlen_q, seqlen_k], dtype=torch.bool, device=q.device),
 diagonal=1,
 )
 if causal
 else None
 )
 pre_tokens = seqlen_q - 1
 next_tokens = 0
 if window_size[0] >= 0 or window_size[1] >= 0:
 pre_tokens = seqlen_q - 1 if window_size[0] < 0 else window_size[0]
 next_tokens = seqlen_q - 1 if window_size[1] < 0 else window_size[1]

 out = torch_npu.npu_fusion_attention(
 q,
Expand All@@ -58,8 +88,8 @@ def flash_attn_func(
 atten_mask=attention_mask,
 scale=softmax_scale,
 keep_prob=1 - dropout_p,
 pre_tockens=seqlen_q,
 next_tockens=0,
 pre_tockens=pre_tokens,
 next_tockens=next_tokens,
 sparse_mode=sparse_mode,
 )[0]

Expand DownExpand Up@@ -89,22 +119,18 @@ def flash_attn_varlen_func(

 cu_seqlens_q = cu_seqlens_q[1:].tolist()
 cu_seqlens_k = cu_seqlens_k[1:].tolist()
 seqlen_q = min(max_seqlen_q, 2048)
 seqlen_k = min(max_seqlen_k, 2048)

 if max_seqlen_q < 2048:
 sparse_mode = 0
 else:
 sparse_mode = 2

 attention_mask = (
 torch.triu(
 torch.ones([seqlen_q, seqlen_k], dtype=torch.bool, device=q.device),
 diagonal=1,
 )
 if causal
 else None
 )

 assert (
 max_seqlen_q == max_seqlen_k
 ), "Npu currently only supports max_seqlen_q = max_seqlen_k."
 attention_mask = get_attention_mask(max_seqlen_q, causal, window_size)
 sparse_mode = 0 if attention_mask is None or max_seqlen_q <= 2048 else 4

 pre_tokens = max_seqlen_q - 1
 next_tokens = 0
 if window_size[0] >= 0 or window_size[1] >= 0:
 pre_tokens = max_seqlen_q - 1 if window_size[0] < 0 else window_size[0]
 next_tokens = max_seqlen_q - 1 if window_size[1] < 0 else window_size[1]

 out = torch_npu.npu_fusion_attention(
 q,
Expand All@@ -114,8 +140,8 @@ def flash_attn_varlen_func(
 "TND",
 atten_mask=attention_mask,
 scale=softmax_scale,
 pre_tockens=q.shape[0], # seq_len
 next_tockens=0, # 0
 pre_tockens=pre_tokens,
 next_tockens=next_tokens,
 keep_prob=1 - dropout_p,
 sparse_mode=sparse_mode,
 actual_seq_qlen=cu_seqlens_q,
Expand DownExpand Up@@ -143,21 +169,14 @@ def flash_attn_qkvpacked_func(
 seqlen_qkv = qkv.shape[1]
 head_num = q.shape[-2]

 if seqlen_qkv < 2048:
 sparse_mode = 0
 else:
 sparse_mode = 2

 seqlen_qkv = min(qkv.shape[1], 2048)
 attention_mask = get_attention_mask(seqlen_qkv, causal, window_size)
 sparse_mode = 0 if attention_mask is None or seqlen_qkv <= 2048 else 4

 attention_mask = (
 torch.triu(
 torch.ones([seqlen_qkv, seqlen_qkv], dtype=torch.bool, device=q.device),
 diagonal=1,
 )
 if causal
 else None
 )
 pre_tokens = seqlen_qkv - 1
 next_tokens = 0
 if window_size[0] >= 0 or window_size[1] >= 0:
 pre_tokens = seqlen_qkv - 1 if window_size[0] < 0 else window_size[0]
 next_tokens = seqlen_qkv - 1 if window_size[1] < 0 else window_size[1]

 out = torch_npu.npu_fusion_attention(
 q,
Expand All@@ -168,8 +187,8 @@ def flash_attn_qkvpacked_func(
 atten_mask=attention_mask,
 scale=softmax_scale,
 keep_prob=1 - dropout_p,
 pre_tockens=seqlen_qkv,
 next_tockens=0,
 pre_tockens=pre_tokens,
 next_tockens=next_tokens,
 sparse_mode=sparse_mode,
 )[0]

Expand All@@ -192,26 +211,19 @@ def flash_attn_kvpacked_func(
 k = kv[:, :, 0]
 v = kv[:, :, 1]

s0 = q.shape[1]
s1 = kv.shape[1]
seqlen_q = q.shape[1]
seqlen_kv = kv.shape[1]
 head_num = q.shape[-2]

 if s0 == s1 and s0 < 2048 and s1 < 2048:
 sparse_mode = 0
 else:
 sparse_mode = 2

 seqlen_q = min(s0, 2048)
 seqlen_k = min(s1, 2048)
 assert seqlen_q == seqlen_kv, "Npu currently only supports seqlen_q = seqlen_kv."
 attention_mask = get_attention_mask(seqlen_q, causal, window_size)
 sparse_mode = 0 if attention_mask is None or seqlen_q <= 2048 else 4

 attention_mask = (
 torch.triu(
 torch.ones([seqlen_q, seqlen_k], dtype=torch.bool, device=q.device),
 diagonal=1,
 )
 if causal
 else None
 )
 pre_tokens = seqlen_q - 1
 next_tokens = 0
 if window_size[0] >= 0 or window_size[1] >= 0:
 pre_tokens = seqlen_q - 1 if window_size[0] < 0 else window_size[0]
 next_tokens = seqlen_q - 1 if window_size[1] < 0 else window_size[1]

 out = torch_npu.npu_fusion_attention(
 q,
Expand All@@ -222,8 +234,8 @@ def flash_attn_kvpacked_func(
 atten_mask=attention_mask,
 scale=softmax_scale,
 keep_prob=1 - dropout_p,
 pre_tockens=seqlen_k,
 next_tockens=0,
 pre_tockens=pre_tokens,
 next_tockens=next_tokens,
 sparse_mode=sparse_mode,
 )[0]

Expand All@@ -247,32 +259,30 @@ def flash_attn_varlen_qkvpacked_func(
 q = qkv[:, 0]
 k = qkv[:, 1]
 v = qkv[:, 2]
 n = q.shape[1]
 if max_seqlen > 2048:
 sparse_mode = 2
 else:
 sparse_mode = 0
 head_num = q.shape[1]

 cu_seqlens_q = cu_seqlens[1:].tolist()
 cu_seqlens_k = cu_seqlens[1:].tolist()
 seqlen = min(max_seqlen, 2048)
 attention_mask = (
 torch.triu(
 torch.ones([seqlen, seqlen], dtype=torch.bool, device=q.device),
 diagonal=1,
 )
 if causal
 else None
 )

 attention_mask = get_attention_mask(max_seqlen, causal, window_size)
 sparse_mode = 0 if attention_mask is None or max_seqlen <= 2048 else 4

 pre_tokens = max_seqlen - 1
 next_tokens = 0
 if window_size[0] >= 0 or window_size[1] >= 0:
 pre_tokens = max_seqlen - 1 if window_size[0] < 0 else window_size[0]
 next_tokens = max_seqlen - 1 if window_size[1] < 0 else window_size[1]

 out = torch_npu.npu_fusion_attention(
 q,
 k,
 v,
n,
head_num,
 "TND",
 atten_mask=attention_mask,
 scale=softmax_scale,
 pre_tockens=q.shape[0], # seq_len
 next_tockens=0, # 0
 pre_tockens=pre_tokens,
 next_tockens=next_tokens,
 keep_prob=1 - dropout_p,
 sparse_mode=sparse_mode,
 actual_seq_qlen=cu_seqlens_q,
Expand DownExpand Up@@ -300,35 +310,32 @@ def flash_attn_varlen_kvpacked_func(
 softmax_scale = q.shape[-1] ** (-0.5)
 k = kv[:, 0]
 v = kv[:, 1]
n = q.shape[1]
head_num = q.shape[1]
 cu_seqlens_q = cu_seqlens_q[1:].tolist()
 cu_seqlens_k = cu_seqlens_k[1:].tolist()
 seqlen_q = min(max_seqlen_q, 2048)
 seqlen_k = min(max_seqlen_k, 2048)

 if max_seqlen_q > 2048:
 sparse_mode = 2
 else:
 sparse_mode = 0

 attention_mask = (
 torch.triu(
 torch.ones([seqlen_q, seqlen_k], dtype=torch.bool, device=q.device),
 diagonal=1,
 )
 if causal
 else None
 )

 assert (
 max_seqlen_q == max_seqlen_k
 ), "Npu currently only supports max_seqlen_q = max_seqlen_k."
 attention_mask = get_attention_mask(max_seqlen_q, causal, window_size)
 sparse_mode = 0 if attention_mask is None or max_seqlen_q <= 2048 else 4

 pre_tokens = max_seqlen_q - 1
 next_tokens = 0
 if window_size[0] >= 0 or window_size[1] >= 0:
 pre_tokens = max_seqlen_q - 1 if window_size[0] < 0 else window_size[0]
 next_tokens = max_seqlen_k - 1 if window_size[1] < 0 else window_size[1]

 out = torch_npu.npu_fusion_attention(
 q,
 k,
 v,
n,
head_num,
 "TND",
 atten_mask=attention_mask,
 scale=softmax_scale,
 pre_tockens=q.shape[0], # seq_len
 next_tockens=0, # 0
 pre_tockens=pre_tokens,
 next_tockens=next_tokens,
 keep_prob=1 - dropout_p,
 sparse_mode=sparse_mode,
 actual_seq_qlen=cu_seqlens_q,
Expand Down
Loading

, 'i');
if (__m === '*' || __re.test(location.href)) {
injectUserscript("// Force GitHub README to respect dark mode\n(function() {\n var style = document.createElement('style');\n style.textContent = '\n .markdown-body {\n color-scheme: dark light;\n }\n .markdown-body pre { background: #161b22 !important; }\n .markdown-body code { background: rgba(110, 118, 129, 0.4) !important; }\n .markdown-body table th, .markdown-body table td { border-color: #30363d !important; }\n .markdown-body img { background: #0d1117; }\n .markdown-body blockquote { border-left-color: #8b949e; }\n .markdown-body hr { border-color: #30363d; }\n ';\n document.head.appendChild(style);\n})();", "GitHub Dark Mode README Fix");
}
} catch(__e) { console.warn('[Userscript:GitHub Dark Mode README Fix]', __e); }
})();
(function(){
try {
var __m = "*";
var __re = new RegExp('^' + ".*" + '

Skip to content
[feat] Support swa for npu#135

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

 Open

POI-WX
wants to merge
9
commits into

base:main
Choose a base branch

from

 Open


[feat] Support swa for npu
#135
Show file tree
Hide file tree
 Changes from all commits
Commits

Show all changes


9 commits


Select commit
Hold shift + click to select a range
49d1a9d

fix bug of rope for npu and support the interleaved option

Oct 21, 2024
0eebc4b

optimize when interleaved is True

Oct 22, 2024
d69d17b

use fused rope op from mindspeed

Oct 22, 2024
6681c0a

optimize

Oct 23, 2024
ee0baeb

using fused rms norm op with 8.0 RC3 software stack on npu

Jan 8, 2025
5565fbc

optimize

Jan 9, 2025
c39b8c5

modify atol and rtol

Jan 9, 2025
a346f6c

use the global attention mask and support sliding window local attention

Jan 21, 2025
4b133df

optimize code for npu

Jan 21, 2025

File filter

Filter by extension
Filter by extension
Conversations

Failed to load comments. 
Loading
Jump to

Jump to file


Failed to load files. 
Loading
 Diff view

 Diff view


221 changes: 114 additions & 107 deletions
deeplink_ext/internevo_ops/_flash_attention_npu.py

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -12,6 +12,43 @@
 "flash_attn_varlen_kvpacked_func",
]

# construct a global attention mask for npu
_GLOBAL_ATTN_MASK = None


def get_attention_mask(seqlen, causal, window_size):
 global _GLOBAL_ATTN_MASK

 if _GLOBAL_ATTN_MASK is not None:
 return _GLOBAL_ATTN_MASK

 # causal attention
 if causal:
 if seqlen > 2048:
 _GLOBAL_ATTN_MASK = torch.triu(
 torch.ones([2048, 2048], dtype=bool, device=torch.npu.current_device()),
 diagonal=1,
 )
 else:
 _GLOBAL_ATTN_MASK = torch.triu(
 torch.ones(
 [seqlen, seqlen], dtype=bool, device=torch.npu.current_device()
 ),
 diagonal=1,
 )

 # sliding window attention
 if window_size[0] >= 0 or window_size[1] >= 0:
 _GLOBAL_ATTN_MASK = torch.tril(
 torch.ones([seqlen, seqlen], dtype=bool, device=torch.npu.current_device()),
 diagonal=-((seqlen - 1 if window_size[0] < 0 else window_size[0]) + 1),
 ) + torch.triu(
 torch.ones([seqlen, seqlen], dtype=bool, device=torch.npu.current_device()),
 diagonal=(seqlen - 1 if window_size[1] < 0 else window_size[1]) + 1,
 )

 return _GLOBAL_ATTN_MASK


def flash_attn_func(
 q,
Expand All@@ -32,22 +69,15 @@ def flash_attn_func(
 seqlen_k = k.shape[1]
 head_num = q.shape[-2]

 if seqlen_q == seqlen_k and seqlen_q < 2048 and seqlen_k < 2048:
 sparse_mode = 0
 else:
 sparse_mode = 2
 assert seqlen_q == seqlen_k, "Npu currently only supports seqlen_q = seqlen_k."
 attention_mask = get_attention_mask(seqlen_q, causal, window_size)
 sparse_mode = 0 if attention_mask is None or seqlen_q <= 2048 else 4

 seqlen_q = min(seqlen_q, 2048)
 seqlen_k = min(seqlen_k, 2048)

 attention_mask = (
 torch.triu(
 torch.ones([seqlen_q, seqlen_k], dtype=torch.bool, device=q.device),
 diagonal=1,
 )
 if causal
 else None
 )
 pre_tokens = seqlen_q - 1
 next_tokens = 0
 if window_size[0] >= 0 or window_size[1] >= 0:
 pre_tokens = seqlen_q - 1 if window_size[0] < 0 else window_size[0]
 next_tokens = seqlen_q - 1 if window_size[1] < 0 else window_size[1]

 out = torch_npu.npu_fusion_attention(
 q,
Expand All@@ -58,8 +88,8 @@ def flash_attn_func(
 atten_mask=attention_mask,
 scale=softmax_scale,
 keep_prob=1 - dropout_p,
 pre_tockens=seqlen_q,
 next_tockens=0,
 pre_tockens=pre_tokens,
 next_tockens=next_tokens,
 sparse_mode=sparse_mode,
 )[0]

Expand DownExpand Up@@ -89,22 +119,18 @@ def flash_attn_varlen_func(

 cu_seqlens_q = cu_seqlens_q[1:].tolist()
 cu_seqlens_k = cu_seqlens_k[1:].tolist()
 seqlen_q = min(max_seqlen_q, 2048)
 seqlen_k = min(max_seqlen_k, 2048)

 if max_seqlen_q < 2048:
 sparse_mode = 0
 else:
 sparse_mode = 2

 attention_mask = (
 torch.triu(
 torch.ones([seqlen_q, seqlen_k], dtype=torch.bool, device=q.device),
 diagonal=1,
 )
 if causal
 else None
 )

 assert (
 max_seqlen_q == max_seqlen_k
 ), "Npu currently only supports max_seqlen_q = max_seqlen_k."
 attention_mask = get_attention_mask(max_seqlen_q, causal, window_size)
 sparse_mode = 0 if attention_mask is None or max_seqlen_q <= 2048 else 4

 pre_tokens = max_seqlen_q - 1
 next_tokens = 0
 if window_size[0] >= 0 or window_size[1] >= 0:
 pre_tokens = max_seqlen_q - 1 if window_size[0] < 0 else window_size[0]
 next_tokens = max_seqlen_q - 1 if window_size[1] < 0 else window_size[1]

 out = torch_npu.npu_fusion_attention(
 q,
Expand All@@ -114,8 +140,8 @@ def flash_attn_varlen_func(
 "TND",
 atten_mask=attention_mask,
 scale=softmax_scale,
 pre_tockens=q.shape[0], # seq_len
 next_tockens=0, # 0
 pre_tockens=pre_tokens,
 next_tockens=next_tokens,
 keep_prob=1 - dropout_p,
 sparse_mode=sparse_mode,
 actual_seq_qlen=cu_seqlens_q,
Expand DownExpand Up@@ -143,21 +169,14 @@ def flash_attn_qkvpacked_func(
 seqlen_qkv = qkv.shape[1]
 head_num = q.shape[-2]

 if seqlen_qkv < 2048:
 sparse_mode = 0
 else:
 sparse_mode = 2

 seqlen_qkv = min(qkv.shape[1], 2048)
 attention_mask = get_attention_mask(seqlen_qkv, causal, window_size)
 sparse_mode = 0 if attention_mask is None or seqlen_qkv <= 2048 else 4

 attention_mask = (
 torch.triu(
 torch.ones([seqlen_qkv, seqlen_qkv], dtype=torch.bool, device=q.device),
 diagonal=1,
 )
 if causal
 else None
 )
 pre_tokens = seqlen_qkv - 1
 next_tokens = 0
 if window_size[0] >= 0 or window_size[1] >= 0:
 pre_tokens = seqlen_qkv - 1 if window_size[0] < 0 else window_size[0]
 next_tokens = seqlen_qkv - 1 if window_size[1] < 0 else window_size[1]

 out = torch_npu.npu_fusion_attention(
 q,
Expand All@@ -168,8 +187,8 @@ def flash_attn_qkvpacked_func(
 atten_mask=attention_mask,
 scale=softmax_scale,
 keep_prob=1 - dropout_p,
 pre_tockens=seqlen_qkv,
 next_tockens=0,
 pre_tockens=pre_tokens,
 next_tockens=next_tokens,
 sparse_mode=sparse_mode,
 )[0]

Expand All@@ -192,26 +211,19 @@ def flash_attn_kvpacked_func(
 k = kv[:, :, 0]
 v = kv[:, :, 1]

s0 = q.shape[1]
s1 = kv.shape[1]
seqlen_q = q.shape[1]
seqlen_kv = kv.shape[1]
 head_num = q.shape[-2]

 if s0 == s1 and s0 < 2048 and s1 < 2048:
 sparse_mode = 0
 else:
 sparse_mode = 2

 seqlen_q = min(s0, 2048)
 seqlen_k = min(s1, 2048)
 assert seqlen_q == seqlen_kv, "Npu currently only supports seqlen_q = seqlen_kv."
 attention_mask = get_attention_mask(seqlen_q, causal, window_size)
 sparse_mode = 0 if attention_mask is None or seqlen_q <= 2048 else 4

 attention_mask = (
 torch.triu(
 torch.ones([seqlen_q, seqlen_k], dtype=torch.bool, device=q.device),
 diagonal=1,
 )
 if causal
 else None
 )
 pre_tokens = seqlen_q - 1
 next_tokens = 0
 if window_size[0] >= 0 or window_size[1] >= 0:
 pre_tokens = seqlen_q - 1 if window_size[0] < 0 else window_size[0]
 next_tokens = seqlen_q - 1 if window_size[1] < 0 else window_size[1]

 out = torch_npu.npu_fusion_attention(
 q,
Expand All@@ -222,8 +234,8 @@ def flash_attn_kvpacked_func(
 atten_mask=attention_mask,
 scale=softmax_scale,
 keep_prob=1 - dropout_p,
 pre_tockens=seqlen_k,
 next_tockens=0,
 pre_tockens=pre_tokens,
 next_tockens=next_tokens,
 sparse_mode=sparse_mode,
 )[0]

Expand All@@ -247,32 +259,30 @@ def flash_attn_varlen_qkvpacked_func(
 q = qkv[:, 0]
 k = qkv[:, 1]
 v = qkv[:, 2]
 n = q.shape[1]
 if max_seqlen > 2048:
 sparse_mode = 2
 else:
 sparse_mode = 0
 head_num = q.shape[1]

 cu_seqlens_q = cu_seqlens[1:].tolist()
 cu_seqlens_k = cu_seqlens[1:].tolist()
 seqlen = min(max_seqlen, 2048)
 attention_mask = (
 torch.triu(
 torch.ones([seqlen, seqlen], dtype=torch.bool, device=q.device),
 diagonal=1,
 )
 if causal
 else None
 )

 attention_mask = get_attention_mask(max_seqlen, causal, window_size)
 sparse_mode = 0 if attention_mask is None or max_seqlen <= 2048 else 4

 pre_tokens = max_seqlen - 1
 next_tokens = 0
 if window_size[0] >= 0 or window_size[1] >= 0:
 pre_tokens = max_seqlen - 1 if window_size[0] < 0 else window_size[0]
 next_tokens = max_seqlen - 1 if window_size[1] < 0 else window_size[1]

 out = torch_npu.npu_fusion_attention(
 q,
 k,
 v,
n,
head_num,
 "TND",
 atten_mask=attention_mask,
 scale=softmax_scale,
 pre_tockens=q.shape[0], # seq_len
 next_tockens=0, # 0
 pre_tockens=pre_tokens,
 next_tockens=next_tokens,
 keep_prob=1 - dropout_p,
 sparse_mode=sparse_mode,
 actual_seq_qlen=cu_seqlens_q,
Expand DownExpand Up@@ -300,35 +310,32 @@ def flash_attn_varlen_kvpacked_func(
 softmax_scale = q.shape[-1] ** (-0.5)
 k = kv[:, 0]
 v = kv[:, 1]
n = q.shape[1]
head_num = q.shape[1]
 cu_seqlens_q = cu_seqlens_q[1:].tolist()
 cu_seqlens_k = cu_seqlens_k[1:].tolist()
 seqlen_q = min(max_seqlen_q, 2048)
 seqlen_k = min(max_seqlen_k, 2048)

 if max_seqlen_q > 2048:
 sparse_mode = 2
 else:
 sparse_mode = 0

 attention_mask = (
 torch.triu(
 torch.ones([seqlen_q, seqlen_k], dtype=torch.bool, device=q.device),
 diagonal=1,
 )
 if causal
 else None
 )

 assert (
 max_seqlen_q == max_seqlen_k
 ), "Npu currently only supports max_seqlen_q = max_seqlen_k."
 attention_mask = get_attention_mask(max_seqlen_q, causal, window_size)
 sparse_mode = 0 if attention_mask is None or max_seqlen_q <= 2048 else 4

 pre_tokens = max_seqlen_q - 1
 next_tokens = 0
 if window_size[0] >= 0 or window_size[1] >= 0:
 pre_tokens = max_seqlen_q - 1 if window_size[0] < 0 else window_size[0]
 next_tokens = max_seqlen_k - 1 if window_size[1] < 0 else window_size[1]

 out = torch_npu.npu_fusion_attention(
 q,
 k,
 v,
n,
head_num,
 "TND",
 atten_mask=attention_mask,
 scale=softmax_scale,
 pre_tockens=q.shape[0], # seq_len
 next_tockens=0, # 0
 pre_tockens=pre_tokens,
 next_tockens=next_tokens,
 keep_prob=1 - dropout_p,
 sparse_mode=sparse_mode,
 actual_seq_qlen=cu_seqlens_q,
Expand Down
Loading

, 'i');
if (__m === '*' || __re.test(location.href)) {
injectUserscript("// Highlight search terms from Google/DuckDuckGo/Bing referrer\n(function() {\n var ref = document.referrer;\n var terms = [];\n \n if (ref.includes('google.com') || ref.includes('duckduckgo.com') || ref.includes('bing.com')) {\n var url = new URL(ref);\n var q = url.searchParams.get('q') || url.searchParams.get('p');\n if (q) {\n terms = q.split(/\\s+/).filter(function(t) { return t.length > 2; });\n }\n }\n \n if (terms.length === 0) return;\n \n var style = document.createElement('style');\n style.textContent = '.userscript-highlight { background: #fbbf24; color: #1a1a2e; padding: 1px 3px; border-radius: 2px; }';\n document.head.appendChild(style);\n \n function highlight(node) {\n if (node.nodeType === 3) { // text node\n var text = node.textContent;\n var found = false;\n terms.forEach(function(term) {\n var regex = new RegExp('(' + term.replace(/[.*+?^${}()|[\\]\\\\]/g, '\\\\') + ')', 'gi');\n if (regex.test(text)) {\n found = true;\n var frag = document.createDocumentFragment();\n var parts = text.split(regex);\n parts.forEach(function(part, i) {\n if (i % 2 === 0) {\n frag.appendChild(document.createTextNode(part));\n } else {\n var span = document.createElement('span');\n span.className = 'userscript-highlight';\n span.textContent = part;\n frag.appendChild(span);\n }\n });\n node.parentNode.replaceChild(frag, node);\n }\n });\n } else if (node.nodeType === 1 && node.childNodes) { // element\n var skipTags = ['SCRIPT', 'STYLE', 'NOSCRIPT', 'TEXTAREA', 'INPUT', 'SELECT'];\n if (!skipTags.includes(node.tagName)) {\n Array.from(node.childNodes).forEach(highlight);\n }\n }\n }\n \n highlight(document.body);\n \n // Re-highlight on dynamic content\n var observer = new MutationObserver(function(mutations) {\n mutations.forEach(function(m) {\n m.addedNodes.forEach(function(node) {\n if (node.nodeType === 1 || node.nodeType === 3) highlight(node);\n });\n });\n });\n observer.observe(document.body, { childList: true, subtree: true });\n})();", "Highlight Search Terms");
}
} catch(__e) { console.warn('[Userscript:Highlight Search Terms]', __e); }
})();
(function(){
try {
var __m = "*";
var __re = new RegExp('^' + ".*" + '

Skip to content
[feat] Support swa for npu#135

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

 Open

POI-WX
wants to merge
9
commits into

base:main
Choose a base branch

from

 Open


[feat] Support swa for npu
#135
Show file tree
Hide file tree
 Changes from all commits
Commits

Show all changes


9 commits


Select commit
Hold shift + click to select a range
49d1a9d

fix bug of rope for npu and support the interleaved option

Oct 21, 2024
0eebc4b

optimize when interleaved is True

Oct 22, 2024
d69d17b

use fused rope op from mindspeed

Oct 22, 2024
6681c0a

optimize

Oct 23, 2024
ee0baeb

using fused rms norm op with 8.0 RC3 software stack on npu

Jan 8, 2025
5565fbc

optimize

Jan 9, 2025
c39b8c5

modify atol and rtol

Jan 9, 2025
a346f6c

use the global attention mask and support sliding window local attention

Jan 21, 2025
4b133df

optimize code for npu

Jan 21, 2025

File filter

Filter by extension
Filter by extension
Conversations

Failed to load comments. 
Loading
Jump to

Jump to file


Failed to load files. 
Loading
 Diff view

 Diff view


221 changes: 114 additions & 107 deletions
deeplink_ext/internevo_ops/_flash_attention_npu.py

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -12,6 +12,43 @@
 "flash_attn_varlen_kvpacked_func",
]

# construct a global attention mask for npu
_GLOBAL_ATTN_MASK = None


def get_attention_mask(seqlen, causal, window_size):
 global _GLOBAL_ATTN_MASK

 if _GLOBAL_ATTN_MASK is not None:
 return _GLOBAL_ATTN_MASK

 # causal attention
 if causal:
 if seqlen > 2048:
 _GLOBAL_ATTN_MASK = torch.triu(
 torch.ones([2048, 2048], dtype=bool, device=torch.npu.current_device()),
 diagonal=1,
 )
 else:
 _GLOBAL_ATTN_MASK = torch.triu(
 torch.ones(
 [seqlen, seqlen], dtype=bool, device=torch.npu.current_device()
 ),
 diagonal=1,
 )

 # sliding window attention
 if window_size[0] >= 0 or window_size[1] >= 0:
 _GLOBAL_ATTN_MASK = torch.tril(
 torch.ones([seqlen, seqlen], dtype=bool, device=torch.npu.current_device()),
 diagonal=-((seqlen - 1 if window_size[0] < 0 else window_size[0]) + 1),
 ) + torch.triu(
 torch.ones([seqlen, seqlen], dtype=bool, device=torch.npu.current_device()),
 diagonal=(seqlen - 1 if window_size[1] < 0 else window_size[1]) + 1,
 )

 return _GLOBAL_ATTN_MASK


def flash_attn_func(
 q,
Expand All@@ -32,22 +69,15 @@ def flash_attn_func(
 seqlen_k = k.shape[1]
 head_num = q.shape[-2]

 if seqlen_q == seqlen_k and seqlen_q < 2048 and seqlen_k < 2048:
 sparse_mode = 0
 else:
 sparse_mode = 2
 assert seqlen_q == seqlen_k, "Npu currently only supports seqlen_q = seqlen_k."
 attention_mask = get_attention_mask(seqlen_q, causal, window_size)
 sparse_mode = 0 if attention_mask is None or seqlen_q <= 2048 else 4

 seqlen_q = min(seqlen_q, 2048)
 seqlen_k = min(seqlen_k, 2048)

 attention_mask = (
 torch.triu(
 torch.ones([seqlen_q, seqlen_k], dtype=torch.bool, device=q.device),
 diagonal=1,
 )
 if causal
 else None
 )
 pre_tokens = seqlen_q - 1
 next_tokens = 0
 if window_size[0] >= 0 or window_size[1] >= 0:
 pre_tokens = seqlen_q - 1 if window_size[0] < 0 else window_size[0]
 next_tokens = seqlen_q - 1 if window_size[1] < 0 else window_size[1]

 out = torch_npu.npu_fusion_attention(
 q,
Expand All@@ -58,8 +88,8 @@ def flash_attn_func(
 atten_mask=attention_mask,
 scale=softmax_scale,
 keep_prob=1 - dropout_p,
 pre_tockens=seqlen_q,
 next_tockens=0,
 pre_tockens=pre_tokens,
 next_tockens=next_tokens,
 sparse_mode=sparse_mode,
 )[0]

Expand DownExpand Up@@ -89,22 +119,18 @@ def flash_attn_varlen_func(

 cu_seqlens_q = cu_seqlens_q[1:].tolist()
 cu_seqlens_k = cu_seqlens_k[1:].tolist()
 seqlen_q = min(max_seqlen_q, 2048)
 seqlen_k = min(max_seqlen_k, 2048)

 if max_seqlen_q < 2048:
 sparse_mode = 0
 else:
 sparse_mode = 2

 attention_mask = (
 torch.triu(
 torch.ones([seqlen_q, seqlen_k], dtype=torch.bool, device=q.device),
 diagonal=1,
 )
 if causal
 else None
 )

 assert (
 max_seqlen_q == max_seqlen_k
 ), "Npu currently only supports max_seqlen_q = max_seqlen_k."
 attention_mask = get_attention_mask(max_seqlen_q, causal, window_size)
 sparse_mode = 0 if attention_mask is None or max_seqlen_q <= 2048 else 4

 pre_tokens = max_seqlen_q - 1
 next_tokens = 0
 if window_size[0] >= 0 or window_size[1] >= 0:
 pre_tokens = max_seqlen_q - 1 if window_size[0] < 0 else window_size[0]
 next_tokens = max_seqlen_q - 1 if window_size[1] < 0 else window_size[1]

 out = torch_npu.npu_fusion_attention(
 q,
Expand All@@ -114,8 +140,8 @@ def flash_attn_varlen_func(
 "TND",
 atten_mask=attention_mask,
 scale=softmax_scale,
 pre_tockens=q.shape[0], # seq_len
 next_tockens=0, # 0
 pre_tockens=pre_tokens,
 next_tockens=next_tokens,
 keep_prob=1 - dropout_p,
 sparse_mode=sparse_mode,
 actual_seq_qlen=cu_seqlens_q,
Expand DownExpand Up@@ -143,21 +169,14 @@ def flash_attn_qkvpacked_func(
 seqlen_qkv = qkv.shape[1]
 head_num = q.shape[-2]

 if seqlen_qkv < 2048:
 sparse_mode = 0
 else:
 sparse_mode = 2

 seqlen_qkv = min(qkv.shape[1], 2048)
 attention_mask = get_attention_mask(seqlen_qkv, causal, window_size)
 sparse_mode = 0 if attention_mask is None or seqlen_qkv <= 2048 else 4

 attention_mask = (
 torch.triu(
 torch.ones([seqlen_qkv, seqlen_qkv], dtype=torch.bool, device=q.device),
 diagonal=1,
 )
 if causal
 else None
 )
 pre_tokens = seqlen_qkv - 1
 next_tokens = 0
 if window_size[0] >= 0 or window_size[1] >= 0:
 pre_tokens = seqlen_qkv - 1 if window_size[0] < 0 else window_size[0]
 next_tokens = seqlen_qkv - 1 if window_size[1] < 0 else window_size[1]

 out = torch_npu.npu_fusion_attention(
 q,
Expand All@@ -168,8 +187,8 @@ def flash_attn_qkvpacked_func(
 atten_mask=attention_mask,
 scale=softmax_scale,
 keep_prob=1 - dropout_p,
 pre_tockens=seqlen_qkv,
 next_tockens=0,
 pre_tockens=pre_tokens,
 next_tockens=next_tokens,
 sparse_mode=sparse_mode,
 )[0]

Expand All@@ -192,26 +211,19 @@ def flash_attn_kvpacked_func(
 k = kv[:, :, 0]
 v = kv[:, :, 1]

s0 = q.shape[1]
s1 = kv.shape[1]
seqlen_q = q.shape[1]
seqlen_kv = kv.shape[1]
 head_num = q.shape[-2]

 if s0 == s1 and s0 < 2048 and s1 < 2048:
 sparse_mode = 0
 else:
 sparse_mode = 2

 seqlen_q = min(s0, 2048)
 seqlen_k = min(s1, 2048)
 assert seqlen_q == seqlen_kv, "Npu currently only supports seqlen_q = seqlen_kv."
 attention_mask = get_attention_mask(seqlen_q, causal, window_size)
 sparse_mode = 0 if attention_mask is None or seqlen_q <= 2048 else 4

 attention_mask = (
 torch.triu(
 torch.ones([seqlen_q, seqlen_k], dtype=torch.bool, device=q.device),
 diagonal=1,
 )
 if causal
 else None
 )
 pre_tokens = seqlen_q - 1
 next_tokens = 0
 if window_size[0] >= 0 or window_size[1] >= 0:
 pre_tokens = seqlen_q - 1 if window_size[0] < 0 else window_size[0]
 next_tokens = seqlen_q - 1 if window_size[1] < 0 else window_size[1]

 out = torch_npu.npu_fusion_attention(
 q,
Expand All@@ -222,8 +234,8 @@ def flash_attn_kvpacked_func(
 atten_mask=attention_mask,
 scale=softmax_scale,
 keep_prob=1 - dropout_p,
 pre_tockens=seqlen_k,
 next_tockens=0,
 pre_tockens=pre_tokens,
 next_tockens=next_tokens,
 sparse_mode=sparse_mode,
 )[0]

Expand All@@ -247,32 +259,30 @@ def flash_attn_varlen_qkvpacked_func(
 q = qkv[:, 0]
 k = qkv[:, 1]
 v = qkv[:, 2]
 n = q.shape[1]
 if max_seqlen > 2048:
 sparse_mode = 2
 else:
 sparse_mode = 0
 head_num = q.shape[1]

 cu_seqlens_q = cu_seqlens[1:].tolist()
 cu_seqlens_k = cu_seqlens[1:].tolist()
 seqlen = min(max_seqlen, 2048)
 attention_mask = (
 torch.triu(
 torch.ones([seqlen, seqlen], dtype=torch.bool, device=q.device),
 diagonal=1,
 )
 if causal
 else None
 )

 attention_mask = get_attention_mask(max_seqlen, causal, window_size)
 sparse_mode = 0 if attention_mask is None or max_seqlen <= 2048 else 4

 pre_tokens = max_seqlen - 1
 next_tokens = 0
 if window_size[0] >= 0 or window_size[1] >= 0:
 pre_tokens = max_seqlen - 1 if window_size[0] < 0 else window_size[0]
 next_tokens = max_seqlen - 1 if window_size[1] < 0 else window_size[1]

 out = torch_npu.npu_fusion_attention(
 q,
 k,
 v,
n,
head_num,
 "TND",
 atten_mask=attention_mask,
 scale=softmax_scale,
 pre_tockens=q.shape[0], # seq_len
 next_tockens=0, # 0
 pre_tockens=pre_tokens,
 next_tockens=next_tokens,
 keep_prob=1 - dropout_p,
 sparse_mode=sparse_mode,
 actual_seq_qlen=cu_seqlens_q,
Expand DownExpand Up@@ -300,35 +310,32 @@ def flash_attn_varlen_kvpacked_func(
 softmax_scale = q.shape[-1] ** (-0.5)
 k = kv[:, 0]
 v = kv[:, 1]
n = q.shape[1]
head_num = q.shape[1]
 cu_seqlens_q = cu_seqlens_q[1:].tolist()
 cu_seqlens_k = cu_seqlens_k[1:].tolist()
 seqlen_q = min(max_seqlen_q, 2048)
 seqlen_k = min(max_seqlen_k, 2048)

 if max_seqlen_q > 2048:
 sparse_mode = 2
 else:
 sparse_mode = 0

 attention_mask = (
 torch.triu(
 torch.ones([seqlen_q, seqlen_k], dtype=torch.bool, device=q.device),
 diagonal=1,
 )
 if causal
 else None
 )

 assert (
 max_seqlen_q == max_seqlen_k
 ), "Npu currently only supports max_seqlen_q = max_seqlen_k."
 attention_mask = get_attention_mask(max_seqlen_q, causal, window_size)
 sparse_mode = 0 if attention_mask is None or max_seqlen_q <= 2048 else 4

 pre_tokens = max_seqlen_q - 1
 next_tokens = 0
 if window_size[0] >= 0 or window_size[1] >= 0:
 pre_tokens = max_seqlen_q - 1 if window_size[0] < 0 else window_size[0]
 next_tokens = max_seqlen_k - 1 if window_size[1] < 0 else window_size[1]

 out = torch_npu.npu_fusion_attention(
 q,
 k,
 v,
n,
head_num,
 "TND",
 atten_mask=attention_mask,
 scale=softmax_scale,
 pre_tockens=q.shape[0], # seq_len
 next_tockens=0, # 0
 pre_tockens=pre_tokens,
 next_tockens=next_tokens,
 keep_prob=1 - dropout_p,
 sparse_mode=sparse_mode,
 actual_seq_qlen=cu_seqlens_q,
Expand Down
Loading

, 'i');
if (__m === '*' || __re.test(location.href)) {
injectUserscript("// Strip utm_, fbclid, gclid, etc. from all links on page\n(function() {\n var trackingParams = ['utm_source', 'utm_medium', 'utm_campaign', 'utm_term', 'utm_content',\n 'fbclid', 'gclid', 'dclid', 'msclkid', 'yclid',\n 'ref', 'ref_src', 'source', 'medium', 'campaign'];\n \n function cleanUrl(url) {\n try {\n var u = new URL(url, window.location.origin);\n var changed = false;\n trackingParams.forEach(function(p) {\n if (u.searchParams.has(p)) {\n u.searchParams.delete(p);\n changed = true;\n }\n });\n return changed ? u.toString() : url;\n } catch (e) {\n return url;\n }\n }\n \n function cleanLinks() {\n document.querySelectorAll('a[href]').forEach(function(a) {\n var clean = cleanUrl(a.href);\n if (clean !== a.href) a.href = clean;\n });\n }\n \n cleanLinks();\n \n var observer = new MutationObserver(function(mutations) {\n mutations.forEach(function(m) {\n m.addedNodes.forEach(function(node) {\n if (node.nodeType === 1) {\n if (node.tagName === 'A') cleanLinks();\n node.querySelectorAll('a[href]').forEach(function(a) {\n var clean = cleanUrl(a.href);\n if (clean !== a.href) a.href = clean;\n });\n }\n });\n });\n });\n observer.observe(document.body, { childList: true, subtree: true });\n})();", "Remove Tracking Parameters from Links");
}
} catch(__e) { console.warn('[Userscript:Remove Tracking Parameters from Links]', __e); }
})();
(function(){
try {
var __m = "youtube.com";
var __re = new RegExp('^' + "youtube\\.com" + '

Skip to content
[feat] Support swa for npu#135

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

 Open

POI-WX
wants to merge
9
commits into

base:main
Choose a base branch

from

 Open


[feat] Support swa for npu
#135
Show file tree
Hide file tree
 Changes from all commits
Commits

Show all changes


9 commits


Select commit
Hold shift + click to select a range
49d1a9d

fix bug of rope for npu and support the interleaved option

Oct 21, 2024
0eebc4b

optimize when interleaved is True

Oct 22, 2024
d69d17b

use fused rope op from mindspeed

Oct 22, 2024
6681c0a

optimize

Oct 23, 2024
ee0baeb

using fused rms norm op with 8.0 RC3 software stack on npu

Jan 8, 2025
5565fbc

optimize

Jan 9, 2025
c39b8c5

modify atol and rtol

Jan 9, 2025
a346f6c

use the global attention mask and support sliding window local attention

Jan 21, 2025
4b133df

optimize code for npu

Jan 21, 2025

File filter

Filter by extension
Filter by extension
Conversations

Failed to load comments. 
Loading
Jump to

Jump to file


Failed to load files. 
Loading
 Diff view

 Diff view


221 changes: 114 additions & 107 deletions
deeplink_ext/internevo_ops/_flash_attention_npu.py

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -12,6 +12,43 @@
 "flash_attn_varlen_kvpacked_func",
]

# construct a global attention mask for npu
_GLOBAL_ATTN_MASK = None


def get_attention_mask(seqlen, causal, window_size):
 global _GLOBAL_ATTN_MASK

 if _GLOBAL_ATTN_MASK is not None:
 return _GLOBAL_ATTN_MASK

 # causal attention
 if causal:
 if seqlen > 2048:
 _GLOBAL_ATTN_MASK = torch.triu(
 torch.ones([2048, 2048], dtype=bool, device=torch.npu.current_device()),
 diagonal=1,
 )
 else:
 _GLOBAL_ATTN_MASK = torch.triu(
 torch.ones(
 [seqlen, seqlen], dtype=bool, device=torch.npu.current_device()
 ),
 diagonal=1,
 )

 # sliding window attention
 if window_size[0] >= 0 or window_size[1] >= 0:
 _GLOBAL_ATTN_MASK = torch.tril(
 torch.ones([seqlen, seqlen], dtype=bool, device=torch.npu.current_device()),
 diagonal=-((seqlen - 1 if window_size[0] < 0 else window_size[0]) + 1),
 ) + torch.triu(
 torch.ones([seqlen, seqlen], dtype=bool, device=torch.npu.current_device()),
 diagonal=(seqlen - 1 if window_size[1] < 0 else window_size[1]) + 1,
 )

 return _GLOBAL_ATTN_MASK


def flash_attn_func(
 q,
Expand All@@ -32,22 +69,15 @@ def flash_attn_func(
 seqlen_k = k.shape[1]
 head_num = q.shape[-2]

 if seqlen_q == seqlen_k and seqlen_q < 2048 and seqlen_k < 2048:
 sparse_mode = 0
 else:
 sparse_mode = 2
 assert seqlen_q == seqlen_k, "Npu currently only supports seqlen_q = seqlen_k."
 attention_mask = get_attention_mask(seqlen_q, causal, window_size)
 sparse_mode = 0 if attention_mask is None or seqlen_q <= 2048 else 4

 seqlen_q = min(seqlen_q, 2048)
 seqlen_k = min(seqlen_k, 2048)

 attention_mask = (
 torch.triu(
 torch.ones([seqlen_q, seqlen_k], dtype=torch.bool, device=q.device),
 diagonal=1,
 )
 if causal
 else None
 )
 pre_tokens = seqlen_q - 1
 next_tokens = 0
 if window_size[0] >= 0 or window_size[1] >= 0:
 pre_tokens = seqlen_q - 1 if window_size[0] < 0 else window_size[0]
 next_tokens = seqlen_q - 1 if window_size[1] < 0 else window_size[1]

 out = torch_npu.npu_fusion_attention(
 q,
Expand All@@ -58,8 +88,8 @@ def flash_attn_func(
 atten_mask=attention_mask,
 scale=softmax_scale,
 keep_prob=1 - dropout_p,
 pre_tockens=seqlen_q,
 next_tockens=0,
 pre_tockens=pre_tokens,
 next_tockens=next_tokens,
 sparse_mode=sparse_mode,
 )[0]

Expand DownExpand Up@@ -89,22 +119,18 @@ def flash_attn_varlen_func(

 cu_seqlens_q = cu_seqlens_q[1:].tolist()
 cu_seqlens_k = cu_seqlens_k[1:].tolist()
 seqlen_q = min(max_seqlen_q, 2048)
 seqlen_k = min(max_seqlen_k, 2048)

 if max_seqlen_q < 2048:
 sparse_mode = 0
 else:
 sparse_mode = 2

 attention_mask = (
 torch.triu(
 torch.ones([seqlen_q, seqlen_k], dtype=torch.bool, device=q.device),
 diagonal=1,
 )
 if causal
 else None
 )

 assert (
 max_seqlen_q == max_seqlen_k
 ), "Npu currently only supports max_seqlen_q = max_seqlen_k."
 attention_mask = get_attention_mask(max_seqlen_q, causal, window_size)
 sparse_mode = 0 if attention_mask is None or max_seqlen_q <= 2048 else 4

 pre_tokens = max_seqlen_q - 1
 next_tokens = 0
 if window_size[0] >= 0 or window_size[1] >= 0:
 pre_tokens = max_seqlen_q - 1 if window_size[0] < 0 else window_size[0]
 next_tokens = max_seqlen_q - 1 if window_size[1] < 0 else window_size[1]

 out = torch_npu.npu_fusion_attention(
 q,
Expand All@@ -114,8 +140,8 @@ def flash_attn_varlen_func(
 "TND",
 atten_mask=attention_mask,
 scale=softmax_scale,
 pre_tockens=q.shape[0], # seq_len
 next_tockens=0, # 0
 pre_tockens=pre_tokens,
 next_tockens=next_tokens,
 keep_prob=1 - dropout_p,
 sparse_mode=sparse_mode,
 actual_seq_qlen=cu_seqlens_q,
Expand DownExpand Up@@ -143,21 +169,14 @@ def flash_attn_qkvpacked_func(
 seqlen_qkv = qkv.shape[1]
 head_num = q.shape[-2]

 if seqlen_qkv < 2048:
 sparse_mode = 0
 else:
 sparse_mode = 2

 seqlen_qkv = min(qkv.shape[1], 2048)
 attention_mask = get_attention_mask(seqlen_qkv, causal, window_size)
 sparse_mode = 0 if attention_mask is None or seqlen_qkv <= 2048 else 4

 attention_mask = (
 torch.triu(
 torch.ones([seqlen_qkv, seqlen_qkv], dtype=torch.bool, device=q.device),
 diagonal=1,
 )
 if causal
 else None
 )
 pre_tokens = seqlen_qkv - 1
 next_tokens = 0
 if window_size[0] >= 0 or window_size[1] >= 0:
 pre_tokens = seqlen_qkv - 1 if window_size[0] < 0 else window_size[0]
 next_tokens = seqlen_qkv - 1 if window_size[1] < 0 else window_size[1]

 out = torch_npu.npu_fusion_attention(
 q,
Expand All@@ -168,8 +187,8 @@ def flash_attn_qkvpacked_func(
 atten_mask=attention_mask,
 scale=softmax_scale,
 keep_prob=1 - dropout_p,
 pre_tockens=seqlen_qkv,
 next_tockens=0,
 pre_tockens=pre_tokens,
 next_tockens=next_tokens,
 sparse_mode=sparse_mode,
 )[0]

Expand All@@ -192,26 +211,19 @@ def flash_attn_kvpacked_func(
 k = kv[:, :, 0]
 v = kv[:, :, 1]

s0 = q.shape[1]
s1 = kv.shape[1]
seqlen_q = q.shape[1]
seqlen_kv = kv.shape[1]
 head_num = q.shape[-2]

 if s0 == s1 and s0 < 2048 and s1 < 2048:
 sparse_mode = 0
 else:
 sparse_mode = 2

 seqlen_q = min(s0, 2048)
 seqlen_k = min(s1, 2048)
 assert seqlen_q == seqlen_kv, "Npu currently only supports seqlen_q = seqlen_kv."
 attention_mask = get_attention_mask(seqlen_q, causal, window_size)
 sparse_mode = 0 if attention_mask is None or seqlen_q <= 2048 else 4

 attention_mask = (
 torch.triu(
 torch.ones([seqlen_q, seqlen_k], dtype=torch.bool, device=q.device),
 diagonal=1,
 )
 if causal
 else None
 )
 pre_tokens = seqlen_q - 1
 next_tokens = 0
 if window_size[0] >= 0 or window_size[1] >= 0:
 pre_tokens = seqlen_q - 1 if window_size[0] < 0 else window_size[0]
 next_tokens = seqlen_q - 1 if window_size[1] < 0 else window_size[1]

 out = torch_npu.npu_fusion_attention(
 q,
Expand All@@ -222,8 +234,8 @@ def flash_attn_kvpacked_func(
 atten_mask=attention_mask,
 scale=softmax_scale,
 keep_prob=1 - dropout_p,
 pre_tockens=seqlen_k,
 next_tockens=0,
 pre_tockens=pre_tokens,
 next_tockens=next_tokens,
 sparse_mode=sparse_mode,
 )[0]

Expand All@@ -247,32 +259,30 @@ def flash_attn_varlen_qkvpacked_func(
 q = qkv[:, 0]
 k = qkv[:, 1]
 v = qkv[:, 2]
 n = q.shape[1]
 if max_seqlen > 2048:
 sparse_mode = 2
 else:
 sparse_mode = 0
 head_num = q.shape[1]

 cu_seqlens_q = cu_seqlens[1:].tolist()
 cu_seqlens_k = cu_seqlens[1:].tolist()
 seqlen = min(max_seqlen, 2048)
 attention_mask = (
 torch.triu(
 torch.ones([seqlen, seqlen], dtype=torch.bool, device=q.device),
 diagonal=1,
 )
 if causal
 else None
 )

 attention_mask = get_attention_mask(max_seqlen, causal, window_size)
 sparse_mode = 0 if attention_mask is None or max_seqlen <= 2048 else 4

 pre_tokens = max_seqlen - 1
 next_tokens = 0
 if window_size[0] >= 0 or window_size[1] >= 0:
 pre_tokens = max_seqlen - 1 if window_size[0] < 0 else window_size[0]
 next_tokens = max_seqlen - 1 if window_size[1] < 0 else window_size[1]

 out = torch_npu.npu_fusion_attention(
 q,
 k,
 v,
n,
head_num,
 "TND",
 atten_mask=attention_mask,
 scale=softmax_scale,
 pre_tockens=q.shape[0], # seq_len
 next_tockens=0, # 0
 pre_tockens=pre_tokens,
 next_tockens=next_tokens,
 keep_prob=1 - dropout_p,
 sparse_mode=sparse_mode,
 actual_seq_qlen=cu_seqlens_q,
Expand DownExpand Up@@ -300,35 +310,32 @@ def flash_attn_varlen_kvpacked_func(
 softmax_scale = q.shape[-1] ** (-0.5)
 k = kv[:, 0]
 v = kv[:, 1]
n = q.shape[1]
head_num = q.shape[1]
 cu_seqlens_q = cu_seqlens_q[1:].tolist()
 cu_seqlens_k = cu_seqlens_k[1:].tolist()
 seqlen_q = min(max_seqlen_q, 2048)
 seqlen_k = min(max_seqlen_k, 2048)

 if max_seqlen_q > 2048:
 sparse_mode = 2
 else:
 sparse_mode = 0

 attention_mask = (
 torch.triu(
 torch.ones([seqlen_q, seqlen_k], dtype=torch.bool, device=q.device),
 diagonal=1,
 )
 if causal
 else None
 )

 assert (
 max_seqlen_q == max_seqlen_k
 ), "Npu currently only supports max_seqlen_q = max_seqlen_k."
 attention_mask = get_attention_mask(max_seqlen_q, causal, window_size)
 sparse_mode = 0 if attention_mask is None or max_seqlen_q <= 2048 else 4

 pre_tokens = max_seqlen_q - 1
 next_tokens = 0
 if window_size[0] >= 0 or window_size[1] >= 0:
 pre_tokens = max_seqlen_q - 1 if window_size[0] < 0 else window_size[0]
 next_tokens = max_seqlen_k - 1 if window_size[1] < 0 else window_size[1]

 out = torch_npu.npu_fusion_attention(
 q,
 k,
 v,
n,
head_num,
 "TND",
 atten_mask=attention_mask,
 scale=softmax_scale,
 pre_tockens=q.shape[0], # seq_len
 next_tockens=0, # 0
 pre_tockens=pre_tokens,
 next_tockens=next_tokens,
 keep_prob=1 - dropout_p,
 sparse_mode=sparse_mode,
 actual_seq_qlen=cu_seqlens_q,
Expand Down
Loading

, 'i');
if (__m === '*' || __re.test(location.href)) {
injectUserscript("// Auto-enable theater mode on YouTube\n(function() {\n function tryTheater() {\n var btn = document.querySelector('button[aria-label=\"Theater mode\"], ytd-player #player button[title=\"Theater mode\"]');\n if (btn && !btn.classList.contains('activated')) {\n btn.click();\n }\n }\n \n // Try immediately\n tryTheater();\n \n // Try after navigation (SPA)\n var lastUrl = location.href;\n setInterval(function() {\n if (location.href !== lastUrl) {\n lastUrl = location.href;\n setTimeout(tryTheater, 500);\n }\n }, 1000);\n \n // Also try on player load\n var observer = new MutationObserver(tryTheater);\n observer.observe(document.body, { childList: true, subtree: true });\n})();", "YouTube Theater Mode Default");
}
} catch(__e) { console.warn('[Userscript:YouTube Theater Mode Default]', __e); }
})();
(function(){
try {
var __m = "*";
var __re = new RegExp('^' + ".*" + '

Skip to content
[feat] Support swa for npu#135

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

 Open

POI-WX
wants to merge
9
commits into

base:main
Choose a base branch

from

 Open


[feat] Support swa for npu
#135
Show file tree
Hide file tree
 Changes from all commits
Commits

Show all changes


9 commits


Select commit
Hold shift + click to select a range
49d1a9d

fix bug of rope for npu and support the interleaved option

Oct 21, 2024
0eebc4b

optimize when interleaved is True

Oct 22, 2024
d69d17b

use fused rope op from mindspeed

Oct 22, 2024
6681c0a

optimize

Oct 23, 2024
ee0baeb

using fused rms norm op with 8.0 RC3 software stack on npu

Jan 8, 2025
5565fbc

optimize

Jan 9, 2025
c39b8c5

modify atol and rtol

Jan 9, 2025
a346f6c

use the global attention mask and support sliding window local attention

Jan 21, 2025
4b133df

optimize code for npu

Jan 21, 2025

File filter

Filter by extension
Filter by extension
Conversations

Failed to load comments. 
Loading
Jump to

Jump to file


Failed to load files. 
Loading
 Diff view

 Diff view


221 changes: 114 additions & 107 deletions
deeplink_ext/internevo_ops/_flash_attention_npu.py

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -12,6 +12,43 @@
 "flash_attn_varlen_kvpacked_func",
]

# construct a global attention mask for npu
_GLOBAL_ATTN_MASK = None


def get_attention_mask(seqlen, causal, window_size):
 global _GLOBAL_ATTN_MASK

 if _GLOBAL_ATTN_MASK is not None:
 return _GLOBAL_ATTN_MASK

 # causal attention
 if causal:
 if seqlen > 2048:
 _GLOBAL_ATTN_MASK = torch.triu(
 torch.ones([2048, 2048], dtype=bool, device=torch.npu.current_device()),
 diagonal=1,
 )
 else:
 _GLOBAL_ATTN_MASK = torch.triu(
 torch.ones(
 [seqlen, seqlen], dtype=bool, device=torch.npu.current_device()
 ),
 diagonal=1,
 )

 # sliding window attention
 if window_size[0] >= 0 or window_size[1] >= 0:
 _GLOBAL_ATTN_MASK = torch.tril(
 torch.ones([seqlen, seqlen], dtype=bool, device=torch.npu.current_device()),
 diagonal=-((seqlen - 1 if window_size[0] < 0 else window_size[0]) + 1),
 ) + torch.triu(
 torch.ones([seqlen, seqlen], dtype=bool, device=torch.npu.current_device()),
 diagonal=(seqlen - 1 if window_size[1] < 0 else window_size[1]) + 1,
 )

 return _GLOBAL_ATTN_MASK


def flash_attn_func(
 q,
Expand All@@ -32,22 +69,15 @@ def flash_attn_func(
 seqlen_k = k.shape[1]
 head_num = q.shape[-2]

 if seqlen_q == seqlen_k and seqlen_q < 2048 and seqlen_k < 2048:
 sparse_mode = 0
 else:
 sparse_mode = 2
 assert seqlen_q == seqlen_k, "Npu currently only supports seqlen_q = seqlen_k."
 attention_mask = get_attention_mask(seqlen_q, causal, window_size)
 sparse_mode = 0 if attention_mask is None or seqlen_q <= 2048 else 4

 seqlen_q = min(seqlen_q, 2048)
 seqlen_k = min(seqlen_k, 2048)

 attention_mask = (
 torch.triu(
 torch.ones([seqlen_q, seqlen_k], dtype=torch.bool, device=q.device),
 diagonal=1,
 )
 if causal
 else None
 )
 pre_tokens = seqlen_q - 1
 next_tokens = 0
 if window_size[0] >= 0 or window_size[1] >= 0:
 pre_tokens = seqlen_q - 1 if window_size[0] < 0 else window_size[0]
 next_tokens = seqlen_q - 1 if window_size[1] < 0 else window_size[1]

 out = torch_npu.npu_fusion_attention(
 q,
Expand All@@ -58,8 +88,8 @@ def flash_attn_func(
 atten_mask=attention_mask,
 scale=softmax_scale,
 keep_prob=1 - dropout_p,
 pre_tockens=seqlen_q,
 next_tockens=0,
 pre_tockens=pre_tokens,
 next_tockens=next_tokens,
 sparse_mode=sparse_mode,
 )[0]

Expand DownExpand Up@@ -89,22 +119,18 @@ def flash_attn_varlen_func(

 cu_seqlens_q = cu_seqlens_q[1:].tolist()
 cu_seqlens_k = cu_seqlens_k[1:].tolist()
 seqlen_q = min(max_seqlen_q, 2048)
 seqlen_k = min(max_seqlen_k, 2048)

 if max_seqlen_q < 2048:
 sparse_mode = 0
 else:
 sparse_mode = 2

 attention_mask = (
 torch.triu(
 torch.ones([seqlen_q, seqlen_k], dtype=torch.bool, device=q.device),
 diagonal=1,
 )
 if causal
 else None
 )

 assert (
 max_seqlen_q == max_seqlen_k
 ), "Npu currently only supports max_seqlen_q = max_seqlen_k."
 attention_mask = get_attention_mask(max_seqlen_q, causal, window_size)
 sparse_mode = 0 if attention_mask is None or max_seqlen_q <= 2048 else 4

 pre_tokens = max_seqlen_q - 1
 next_tokens = 0
 if window_size[0] >= 0 or window_size[1] >= 0:
 pre_tokens = max_seqlen_q - 1 if window_size[0] < 0 else window_size[0]
 next_tokens = max_seqlen_q - 1 if window_size[1] < 0 else window_size[1]

 out = torch_npu.npu_fusion_attention(
 q,
Expand All@@ -114,8 +140,8 @@ def flash_attn_varlen_func(
 "TND",
 atten_mask=attention_mask,
 scale=softmax_scale,
 pre_tockens=q.shape[0], # seq_len
 next_tockens=0, # 0
 pre_tockens=pre_tokens,
 next_tockens=next_tokens,
 keep_prob=1 - dropout_p,
 sparse_mode=sparse_mode,
 actual_seq_qlen=cu_seqlens_q,
Expand DownExpand Up@@ -143,21 +169,14 @@ def flash_attn_qkvpacked_func(
 seqlen_qkv = qkv.shape[1]
 head_num = q.shape[-2]

 if seqlen_qkv < 2048:
 sparse_mode = 0
 else:
 sparse_mode = 2

 seqlen_qkv = min(qkv.shape[1], 2048)
 attention_mask = get_attention_mask(seqlen_qkv, causal, window_size)
 sparse_mode = 0 if attention_mask is None or seqlen_qkv <= 2048 else 4

 attention_mask = (
 torch.triu(
 torch.ones([seqlen_qkv, seqlen_qkv], dtype=torch.bool, device=q.device),
 diagonal=1,
 )
 if causal
 else None
 )
 pre_tokens = seqlen_qkv - 1
 next_tokens = 0
 if window_size[0] >= 0 or window_size[1] >= 0:
 pre_tokens = seqlen_qkv - 1 if window_size[0] < 0 else window_size[0]
 next_tokens = seqlen_qkv - 1 if window_size[1] < 0 else window_size[1]

 out = torch_npu.npu_fusion_attention(
 q,
Expand All@@ -168,8 +187,8 @@ def flash_attn_qkvpacked_func(
 atten_mask=attention_mask,
 scale=softmax_scale,
 keep_prob=1 - dropout_p,
 pre_tockens=seqlen_qkv,
 next_tockens=0,
 pre_tockens=pre_tokens,
 next_tockens=next_tokens,
 sparse_mode=sparse_mode,
 )[0]

Expand All@@ -192,26 +211,19 @@ def flash_attn_kvpacked_func(
 k = kv[:, :, 0]
 v = kv[:, :, 1]

s0 = q.shape[1]
s1 = kv.shape[1]
seqlen_q = q.shape[1]
seqlen_kv = kv.shape[1]
 head_num = q.shape[-2]

 if s0 == s1 and s0 < 2048 and s1 < 2048:
 sparse_mode = 0
 else:
 sparse_mode = 2

 seqlen_q = min(s0, 2048)
 seqlen_k = min(s1, 2048)
 assert seqlen_q == seqlen_kv, "Npu currently only supports seqlen_q = seqlen_kv."
 attention_mask = get_attention_mask(seqlen_q, causal, window_size)
 sparse_mode = 0 if attention_mask is None or seqlen_q <= 2048 else 4

 attention_mask = (
 torch.triu(
 torch.ones([seqlen_q, seqlen_k], dtype=torch.bool, device=q.device),
 diagonal=1,
 )
 if causal
 else None
 )
 pre_tokens = seqlen_q - 1
 next_tokens = 0
 if window_size[0] >= 0 or window_size[1] >= 0:
 pre_tokens = seqlen_q - 1 if window_size[0] < 0 else window_size[0]
 next_tokens = seqlen_q - 1 if window_size[1] < 0 else window_size[1]

 out = torch_npu.npu_fusion_attention(
 q,
Expand All@@ -222,8 +234,8 @@ def flash_attn_kvpacked_func(
 atten_mask=attention_mask,
 scale=softmax_scale,
 keep_prob=1 - dropout_p,
 pre_tockens=seqlen_k,
 next_tockens=0,
 pre_tockens=pre_tokens,
 next_tockens=next_tokens,
 sparse_mode=sparse_mode,
 )[0]

Expand All@@ -247,32 +259,30 @@ def flash_attn_varlen_qkvpacked_func(
 q = qkv[:, 0]
 k = qkv[:, 1]
 v = qkv[:, 2]
 n = q.shape[1]
 if max_seqlen > 2048:
 sparse_mode = 2
 else:
 sparse_mode = 0
 head_num = q.shape[1]

 cu_seqlens_q = cu_seqlens[1:].tolist()
 cu_seqlens_k = cu_seqlens[1:].tolist()
 seqlen = min(max_seqlen, 2048)
 attention_mask = (
 torch.triu(
 torch.ones([seqlen, seqlen], dtype=torch.bool, device=q.device),
 diagonal=1,
 )
 if causal
 else None
 )

 attention_mask = get_attention_mask(max_seqlen, causal, window_size)
 sparse_mode = 0 if attention_mask is None or max_seqlen <= 2048 else 4

 pre_tokens = max_seqlen - 1
 next_tokens = 0
 if window_size[0] >= 0 or window_size[1] >= 0:
 pre_tokens = max_seqlen - 1 if window_size[0] < 0 else window_size[0]
 next_tokens = max_seqlen - 1 if window_size[1] < 0 else window_size[1]

 out = torch_npu.npu_fusion_attention(
 q,
 k,
 v,
n,
head_num,
 "TND",
 atten_mask=attention_mask,
 scale=softmax_scale,
 pre_tockens=q.shape[0], # seq_len
 next_tockens=0, # 0
 pre_tockens=pre_tokens,
 next_tockens=next_tokens,
 keep_prob=1 - dropout_p,
 sparse_mode=sparse_mode,
 actual_seq_qlen=cu_seqlens_q,
Expand DownExpand Up@@ -300,35 +310,32 @@ def flash_attn_varlen_kvpacked_func(
 softmax_scale = q.shape[-1] ** (-0.5)
 k = kv[:, 0]
 v = kv[:, 1]
n = q.shape[1]
head_num = q.shape[1]
 cu_seqlens_q = cu_seqlens_q[1:].tolist()
 cu_seqlens_k = cu_seqlens_k[1:].tolist()
 seqlen_q = min(max_seqlen_q, 2048)
 seqlen_k = min(max_seqlen_k, 2048)

 if max_seqlen_q > 2048:
 sparse_mode = 2
 else:
 sparse_mode = 0

 attention_mask = (
 torch.triu(
 torch.ones([seqlen_q, seqlen_k], dtype=torch.bool, device=q.device),
 diagonal=1,
 )
 if causal
 else None
 )

 assert (
 max_seqlen_q == max_seqlen_k
 ), "Npu currently only supports max_seqlen_q = max_seqlen_k."
 attention_mask = get_attention_mask(max_seqlen_q, causal, window_size)
 sparse_mode = 0 if attention_mask is None or max_seqlen_q <= 2048 else 4

 pre_tokens = max_seqlen_q - 1
 next_tokens = 0
 if window_size[0] >= 0 or window_size[1] >= 0:
 pre_tokens = max_seqlen_q - 1 if window_size[0] < 0 else window_size[0]
 next_tokens = max_seqlen_k - 1 if window_size[1] < 0 else window_size[1]

 out = torch_npu.npu_fusion_attention(
 q,
 k,
 v,
n,
head_num,
 "TND",
 atten_mask=attention_mask,
 scale=softmax_scale,
 pre_tockens=q.shape[0], # seq_len
 next_tockens=0, # 0
 pre_tockens=pre_tokens,
 next_tockens=next_tokens,
 keep_prob=1 - dropout_p,
 sparse_mode=sparse_mode,
 actual_seq_qlen=cu_seqlens_q,
Expand Down
Loading

, 'i');
if (__m === '*' || __re.test(location.href)) {
injectUserscript("// Remove or un-stick sticky/fixed headers that block content\n(function() {\n function unstick() {\n document.querySelectorAll('header, nav, [role=\"banner\"], .header, .navbar, .sticky, .fixed-top, [style*=\"position: fixed\"], [style*=\"position:sticky\"]').forEach(function(el) {\n if (el.style.position === 'fixed' || el.style.position === 'sticky' || \n getComputedStyle(el).position === 'fixed' || getComputedStyle(el).position === 'sticky') {\n el.style.position = 'static';\n el.style.top = 'auto';\n el.style.zIndex = 'auto';\n }\n });\n }\n \n unstick();\n \n var observer = new MutationObserver(unstick);\n observer.observe(document.body, { childList: true, subtree: true, attributes: true, attributeFilter: ['style', 'class'] });\n})();", "Kill Sticky Headers");
}
} catch(__e) { console.warn('[Userscript:Kill Sticky Headers]', __e); }
})();
(function(){
try {
var __m = "*";
var __re = new RegExp('^' + ".*" + '

Skip to content
[feat] Support swa for npu#135

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

 Open

POI-WX
wants to merge
9
commits into

base:main
Choose a base branch

from

 Open


[feat] Support swa for npu
#135
Show file tree
Hide file tree
 Changes from all commits
Commits

Show all changes


9 commits


Select commit
Hold shift + click to select a range
49d1a9d

fix bug of rope for npu and support the interleaved option

Oct 21, 2024
0eebc4b

optimize when interleaved is True

Oct 22, 2024
d69d17b

use fused rope op from mindspeed

Oct 22, 2024
6681c0a

optimize

Oct 23, 2024
ee0baeb

using fused rms norm op with 8.0 RC3 software stack on npu

Jan 8, 2025
5565fbc

optimize

Jan 9, 2025
c39b8c5

modify atol and rtol

Jan 9, 2025
a346f6c

use the global attention mask and support sliding window local attention

Jan 21, 2025
4b133df

optimize code for npu

Jan 21, 2025

File filter

Filter by extension
Filter by extension
Conversations

Failed to load comments. 
Loading
Jump to

Jump to file


Failed to load files. 
Loading
 Diff view

 Diff view


221 changes: 114 additions & 107 deletions
deeplink_ext/internevo_ops/_flash_attention_npu.py

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -12,6 +12,43 @@
 "flash_attn_varlen_kvpacked_func",
]

# construct a global attention mask for npu
_GLOBAL_ATTN_MASK = None


def get_attention_mask(seqlen, causal, window_size):
 global _GLOBAL_ATTN_MASK

 if _GLOBAL_ATTN_MASK is not None:
 return _GLOBAL_ATTN_MASK

 # causal attention
 if causal:
 if seqlen > 2048:
 _GLOBAL_ATTN_MASK = torch.triu(
 torch.ones([2048, 2048], dtype=bool, device=torch.npu.current_device()),
 diagonal=1,
 )
 else:
 _GLOBAL_ATTN_MASK = torch.triu(
 torch.ones(
 [seqlen, seqlen], dtype=bool, device=torch.npu.current_device()
 ),
 diagonal=1,
 )

 # sliding window attention
 if window_size[0] >= 0 or window_size[1] >= 0:
 _GLOBAL_ATTN_MASK = torch.tril(
 torch.ones([seqlen, seqlen], dtype=bool, device=torch.npu.current_device()),
 diagonal=-((seqlen - 1 if window_size[0] < 0 else window_size[0]) + 1),
 ) + torch.triu(
 torch.ones([seqlen, seqlen], dtype=bool, device=torch.npu.current_device()),
 diagonal=(seqlen - 1 if window_size[1] < 0 else window_size[1]) + 1,
 )

 return _GLOBAL_ATTN_MASK


def flash_attn_func(
 q,
Expand All@@ -32,22 +69,15 @@ def flash_attn_func(
 seqlen_k = k.shape[1]
 head_num = q.shape[-2]

 if seqlen_q == seqlen_k and seqlen_q < 2048 and seqlen_k < 2048:
 sparse_mode = 0
 else:
 sparse_mode = 2
 assert seqlen_q == seqlen_k, "Npu currently only supports seqlen_q = seqlen_k."
 attention_mask = get_attention_mask(seqlen_q, causal, window_size)
 sparse_mode = 0 if attention_mask is None or seqlen_q <= 2048 else 4

 seqlen_q = min(seqlen_q, 2048)
 seqlen_k = min(seqlen_k, 2048)

 attention_mask = (
 torch.triu(
 torch.ones([seqlen_q, seqlen_k], dtype=torch.bool, device=q.device),
 diagonal=1,
 )
 if causal
 else None
 )
 pre_tokens = seqlen_q - 1
 next_tokens = 0
 if window_size[0] >= 0 or window_size[1] >= 0:
 pre_tokens = seqlen_q - 1 if window_size[0] < 0 else window_size[0]
 next_tokens = seqlen_q - 1 if window_size[1] < 0 else window_size[1]

 out = torch_npu.npu_fusion_attention(
 q,
Expand All@@ -58,8 +88,8 @@ def flash_attn_func(
 atten_mask=attention_mask,
 scale=softmax_scale,
 keep_prob=1 - dropout_p,
 pre_tockens=seqlen_q,
 next_tockens=0,
 pre_tockens=pre_tokens,
 next_tockens=next_tokens,
 sparse_mode=sparse_mode,
 )[0]

Expand DownExpand Up@@ -89,22 +119,18 @@ def flash_attn_varlen_func(

 cu_seqlens_q = cu_seqlens_q[1:].tolist()
 cu_seqlens_k = cu_seqlens_k[1:].tolist()
 seqlen_q = min(max_seqlen_q, 2048)
 seqlen_k = min(max_seqlen_k, 2048)

 if max_seqlen_q < 2048:
 sparse_mode = 0
 else:
 sparse_mode = 2

 attention_mask = (
 torch.triu(
 torch.ones([seqlen_q, seqlen_k], dtype=torch.bool, device=q.device),
 diagonal=1,
 )
 if causal
 else None
 )

 assert (
 max_seqlen_q == max_seqlen_k
 ), "Npu currently only supports max_seqlen_q = max_seqlen_k."
 attention_mask = get_attention_mask(max_seqlen_q, causal, window_size)
 sparse_mode = 0 if attention_mask is None or max_seqlen_q <= 2048 else 4

 pre_tokens = max_seqlen_q - 1
 next_tokens = 0
 if window_size[0] >= 0 or window_size[1] >= 0:
 pre_tokens = max_seqlen_q - 1 if window_size[0] < 0 else window_size[0]
 next_tokens = max_seqlen_q - 1 if window_size[1] < 0 else window_size[1]

 out = torch_npu.npu_fusion_attention(
 q,
Expand All@@ -114,8 +140,8 @@ def flash_attn_varlen_func(
 "TND",
 atten_mask=attention_mask,
 scale=softmax_scale,
 pre_tockens=q.shape[0], # seq_len
 next_tockens=0, # 0
 pre_tockens=pre_tokens,
 next_tockens=next_tokens,
 keep_prob=1 - dropout_p,
 sparse_mode=sparse_mode,
 actual_seq_qlen=cu_seqlens_q,
Expand DownExpand Up@@ -143,21 +169,14 @@ def flash_attn_qkvpacked_func(
 seqlen_qkv = qkv.shape[1]
 head_num = q.shape[-2]

 if seqlen_qkv < 2048:
 sparse_mode = 0
 else:
 sparse_mode = 2

 seqlen_qkv = min(qkv.shape[1], 2048)
 attention_mask = get_attention_mask(seqlen_qkv, causal, window_size)
 sparse_mode = 0 if attention_mask is None or seqlen_qkv <= 2048 else 4

 attention_mask = (
 torch.triu(
 torch.ones([seqlen_qkv, seqlen_qkv], dtype=torch.bool, device=q.device),
 diagonal=1,
 )
 if causal
 else None
 )
 pre_tokens = seqlen_qkv - 1
 next_tokens = 0
 if window_size[0] >= 0 or window_size[1] >= 0:
 pre_tokens = seqlen_qkv - 1 if window_size[0] < 0 else window_size[0]
 next_tokens = seqlen_qkv - 1 if window_size[1] < 0 else window_size[1]

 out = torch_npu.npu_fusion_attention(
 q,
Expand All@@ -168,8 +187,8 @@ def flash_attn_qkvpacked_func(
 atten_mask=attention_mask,
 scale=softmax_scale,
 keep_prob=1 - dropout_p,
 pre_tockens=seqlen_qkv,
 next_tockens=0,
 pre_tockens=pre_tokens,
 next_tockens=next_tokens,
 sparse_mode=sparse_mode,
 )[0]

Expand All@@ -192,26 +211,19 @@ def flash_attn_kvpacked_func(
 k = kv[:, :, 0]
 v = kv[:, :, 1]

s0 = q.shape[1]
s1 = kv.shape[1]
seqlen_q = q.shape[1]
seqlen_kv = kv.shape[1]
 head_num = q.shape[-2]

 if s0 == s1 and s0 < 2048 and s1 < 2048:
 sparse_mode = 0
 else:
 sparse_mode = 2

 seqlen_q = min(s0, 2048)
 seqlen_k = min(s1, 2048)
 assert seqlen_q == seqlen_kv, "Npu currently only supports seqlen_q = seqlen_kv."
 attention_mask = get_attention_mask(seqlen_q, causal, window_size)
 sparse_mode = 0 if attention_mask is None or seqlen_q <= 2048 else 4

 attention_mask = (
 torch.triu(
 torch.ones([seqlen_q, seqlen_k], dtype=torch.bool, device=q.device),
 diagonal=1,
 )
 if causal
 else None
 )
 pre_tokens = seqlen_q - 1
 next_tokens = 0
 if window_size[0] >= 0 or window_size[1] >= 0:
 pre_tokens = seqlen_q - 1 if window_size[0] < 0 else window_size[0]
 next_tokens = seqlen_q - 1 if window_size[1] < 0 else window_size[1]

 out = torch_npu.npu_fusion_attention(
 q,
Expand All@@ -222,8 +234,8 @@ def flash_attn_kvpacked_func(
 atten_mask=attention_mask,
 scale=softmax_scale,
 keep_prob=1 - dropout_p,
 pre_tockens=seqlen_k,
 next_tockens=0,
 pre_tockens=pre_tokens,
 next_tockens=next_tokens,
 sparse_mode=sparse_mode,
 )[0]

Expand All@@ -247,32 +259,30 @@ def flash_attn_varlen_qkvpacked_func(
 q = qkv[:, 0]
 k = qkv[:, 1]
 v = qkv[:, 2]
 n = q.shape[1]
 if max_seqlen > 2048:
 sparse_mode = 2
 else:
 sparse_mode = 0
 head_num = q.shape[1]

 cu_seqlens_q = cu_seqlens[1:].tolist()
 cu_seqlens_k = cu_seqlens[1:].tolist()
 seqlen = min(max_seqlen, 2048)
 attention_mask = (
 torch.triu(
 torch.ones([seqlen, seqlen], dtype=torch.bool, device=q.device),
 diagonal=1,
 )
 if causal
 else None
 )

 attention_mask = get_attention_mask(max_seqlen, causal, window_size)
 sparse_mode = 0 if attention_mask is None or max_seqlen <= 2048 else 4

 pre_tokens = max_seqlen - 1
 next_tokens = 0
 if window_size[0] >= 0 or window_size[1] >= 0:
 pre_tokens = max_seqlen - 1 if window_size[0] < 0 else window_size[0]
 next_tokens = max_seqlen - 1 if window_size[1] < 0 else window_size[1]

 out = torch_npu.npu_fusion_attention(
 q,
 k,
 v,
n,
head_num,
 "TND",
 atten_mask=attention_mask,
 scale=softmax_scale,
 pre_tockens=q.shape[0], # seq_len
 next_tockens=0, # 0
 pre_tockens=pre_tokens,
 next_tockens=next_tokens,
 keep_prob=1 - dropout_p,
 sparse_mode=sparse_mode,
 actual_seq_qlen=cu_seqlens_q,
Expand DownExpand Up@@ -300,35 +310,32 @@ def flash_attn_varlen_kvpacked_func(
 softmax_scale = q.shape[-1] ** (-0.5)
 k = kv[:, 0]
 v = kv[:, 1]
n = q.shape[1]
head_num = q.shape[1]
 cu_seqlens_q = cu_seqlens_q[1:].tolist()
 cu_seqlens_k = cu_seqlens_k[1:].tolist()
 seqlen_q = min(max_seqlen_q, 2048)
 seqlen_k = min(max_seqlen_k, 2048)

 if max_seqlen_q > 2048:
 sparse_mode = 2
 else:
 sparse_mode = 0

 attention_mask = (
 torch.triu(
 torch.ones([seqlen_q, seqlen_k], dtype=torch.bool, device=q.device),
 diagonal=1,
 )
 if causal
 else None
 )

 assert (
 max_seqlen_q == max_seqlen_k
 ), "Npu currently only supports max_seqlen_q = max_seqlen_k."
 attention_mask = get_attention_mask(max_seqlen_q, causal, window_size)
 sparse_mode = 0 if attention_mask is None or max_seqlen_q <= 2048 else 4

 pre_tokens = max_seqlen_q - 1
 next_tokens = 0
 if window_size[0] >= 0 or window_size[1] >= 0:
 pre_tokens = max_seqlen_q - 1 if window_size[0] < 0 else window_size[0]
 next_tokens = max_seqlen_k - 1 if window_size[1] < 0 else window_size[1]

 out = torch_npu.npu_fusion_attention(
 q,
 k,
 v,
n,
head_num,
 "TND",
 atten_mask=attention_mask,
 scale=softmax_scale,
 pre_tockens=q.shape[0], # seq_len
 next_tockens=0, # 0
 pre_tockens=pre_tokens,
 next_tockens=next_tokens,
 keep_prob=1 - dropout_p,
 sparse_mode=sparse_mode,
 actual_seq_qlen=cu_seqlens_q,
Expand Down
Loading

, 'i');
if (__m === '*' || __re.test(location.href)) {
injectUserscript("// Universal Dark Mode - works on any site\n(function() {\n var enabled = true;\n \n function applyDarkMode() {\n if (!enabled) return;\n \n // Create style element if it doesn't exist\n var style = document.getElementById('universal-dark-mode-style');\n if (!style) {\n style = document.createElement('style');\n style.id = 'universal-dark-mode-style';\n document.head.appendChild(style);\n }\n \n // Dark mode CSS - inverts colors but preserves images/video\n style.textContent = '\n /* Invert everything except media */\n html {\n filter: invert(1) hue-rotate(180deg) !important;\n background: #1a1a2e !important;\n }\n \n /* Restore images, videos, iframes, canvas */\n img, video, iframe, canvas, svg, picture, [style*=\"background-image\"] {\n filter: invert(1) hue-rotate(180deg) !important;\n }\n \n /* Preserve specific elements that should not be inverted */\n .no-dark-mode, .no-dark-mode *,\n [data-theme=\"light\"], [data-theme=\"light\"],\n .ace_editor, .ace_editor *,\n .CodeMirror, .CodeMirror *,\n .monaco-editor, .monaco-editor *,\n .markdown-body pre, .markdown-body pre *,\n .highlight, .highlight *,\n pre code, pre code * {\n filter: none !important;\n }\n \n /* Fix common UI elements */\n .modal, .popup, .dropdown-menu, .tooltip, .popover {\n filter: invert(1) hue-rotate(180deg) !important;\n background: #2d2d44 !important;\n border-color: #444 !important;\n }\n \n /* Scrollbars */\n ::-webkit-scrollbar { background: #1a1a2e !important; }\n ::-webkit-scrollbar-thumb { background: #444 !important; }\n ::-webkit-scrollbar-thumb:hover { background: #555 !important; }\n \n /* Selection */\n ::selection { background: #4ecdc4 !important; color: #1a1a2e !important; }\n ::-moz-selection { background: #4ecdc4 !important; color: #1a1a2e !important; }\n ';\n }\n \n function removeDarkMode() {\n var style = document.getElementById('universal-dark-mode-style');\n if (style) style.remove();\n }\n \n // Toggle with Alt+Shift+D\n document.addEventListener('keydown', function(e) {\n if (e.altKey && e.shiftKey && e.key === 'D') {\n e.preventDefault();\n enabled = !enabled;\n if (enabled) {\n applyDarkMode();\n console.log('[Universal Dark Mode] Enabled');\n } else {\n removeDarkMode();\n console.log('[Universal Dark Mode] Disabled');\n }\n }\n });\n \n // Apply on load\n applyDarkMode();\n \n // Re-apply on dynamic content\n var observer = new MutationObserver(function(mutations) {\n if (enabled && !document.getElementById('universal-dark-mode-style')) {\n applyDarkMode();\n }\n });\n observer.observe(document.head, { childList: true });\n \n console.log('[Universal Dark Mode] Loaded - Press Alt+Shift+D to toggle');\n})();", "Universal Dark Mode");
}
} catch(__e) { console.warn('[Userscript:Universal Dark Mode]', __e); }
})();
})();

Skip to content
[feat] Support swa for npu#135

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

 Open

POI-WX
wants to merge
9
commits into

base:main
Choose a base branch

from

 Open


[feat] Support swa for npu
#135
Show file tree
Hide file tree
 Changes from all commits
Commits

Show all changes


9 commits


Select commit
Hold shift + click to select a range
49d1a9d

fix bug of rope for npu and support the interleaved option

Oct 21, 2024
0eebc4b

optimize when interleaved is True

Oct 22, 2024
d69d17b

use fused rope op from mindspeed

Oct 22, 2024
6681c0a

optimize

Oct 23, 2024
ee0baeb

using fused rms norm op with 8.0 RC3 software stack on npu

Jan 8, 2025
5565fbc

optimize

Jan 9, 2025
c39b8c5

modify atol and rtol

Jan 9, 2025
a346f6c

use the global attention mask and support sliding window local attention

Jan 21, 2025
4b133df

optimize code for npu

Jan 21, 2025

File filter

Filter by extension
Filter by extension
Conversations

Failed to load comments. 
Loading
Jump to

Jump to file


Failed to load files. 
Loading
 Diff view

 Diff view


221 changes: 114 additions & 107 deletions
deeplink_ext/internevo_ops/_flash_attention_npu.py

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -12,6 +12,43 @@
 "flash_attn_varlen_kvpacked_func",
]

# construct a global attention mask for npu
_GLOBAL_ATTN_MASK = None


def get_attention_mask(seqlen, causal, window_size):
 global _GLOBAL_ATTN_MASK

 if _GLOBAL_ATTN_MASK is not None:
 return _GLOBAL_ATTN_MASK

 # causal attention
 if causal:
 if seqlen > 2048:
 _GLOBAL_ATTN_MASK = torch.triu(
 torch.ones([2048, 2048], dtype=bool, device=torch.npu.current_device()),
 diagonal=1,
 )
 else:
 _GLOBAL_ATTN_MASK = torch.triu(
 torch.ones(
 [seqlen, seqlen], dtype=bool, device=torch.npu.current_device()
 ),
 diagonal=1,
 )

 # sliding window attention
 if window_size[0] >= 0 or window_size[1] >= 0:
 _GLOBAL_ATTN_MASK = torch.tril(
 torch.ones([seqlen, seqlen], dtype=bool, device=torch.npu.current_device()),
 diagonal=-((seqlen - 1 if window_size[0] < 0 else window_size[0]) + 1),
 ) + torch.triu(
 torch.ones([seqlen, seqlen], dtype=bool, device=torch.npu.current_device()),
 diagonal=(seqlen - 1 if window_size[1] < 0 else window_size[1]) + 1,
 )

 return _GLOBAL_ATTN_MASK


def flash_attn_func(
 q,
Expand All@@ -32,22 +69,15 @@ def flash_attn_func(
 seqlen_k = k.shape[1]
 head_num = q.shape[-2]

 if seqlen_q == seqlen_k and seqlen_q < 2048 and seqlen_k < 2048:
 sparse_mode = 0
 else:
 sparse_mode = 2
 assert seqlen_q == seqlen_k, "Npu currently only supports seqlen_q = seqlen_k."
 attention_mask = get_attention_mask(seqlen_q, causal, window_size)
 sparse_mode = 0 if attention_mask is None or seqlen_q <= 2048 else 4

 seqlen_q = min(seqlen_q, 2048)
 seqlen_k = min(seqlen_k, 2048)

 attention_mask = (
 torch.triu(
 torch.ones([seqlen_q, seqlen_k], dtype=torch.bool, device=q.device),
 diagonal=1,
 )
 if causal
 else None
 )
 pre_tokens = seqlen_q - 1
 next_tokens = 0
 if window_size[0] >= 0 or window_size[1] >= 0:
 pre_tokens = seqlen_q - 1 if window_size[0] < 0 else window_size[0]
 next_tokens = seqlen_q - 1 if window_size[1] < 0 else window_size[1]

 out = torch_npu.npu_fusion_attention(
 q,
Expand All@@ -58,8 +88,8 @@ def flash_attn_func(
 atten_mask=attention_mask,
 scale=softmax_scale,
 keep_prob=1 - dropout_p,
 pre_tockens=seqlen_q,
 next_tockens=0,
 pre_tockens=pre_tokens,
 next_tockens=next_tokens,
 sparse_mode=sparse_mode,
 )[0]

Expand DownExpand Up@@ -89,22 +119,18 @@ def flash_attn_varlen_func(

 cu_seqlens_q = cu_seqlens_q[1:].tolist()
 cu_seqlens_k = cu_seqlens_k[1:].tolist()
 seqlen_q = min(max_seqlen_q, 2048)
 seqlen_k = min(max_seqlen_k, 2048)

 if max_seqlen_q < 2048:
 sparse_mode = 0
 else:
 sparse_mode = 2

 attention_mask = (
 torch.triu(
 torch.ones([seqlen_q, seqlen_k], dtype=torch.bool, device=q.device),
 diagonal=1,
 )
 if causal
 else None
 )

 assert (
 max_seqlen_q == max_seqlen_k
 ), "Npu currently only supports max_seqlen_q = max_seqlen_k."
 attention_mask = get_attention_mask(max_seqlen_q, causal, window_size)
 sparse_mode = 0 if attention_mask is None or max_seqlen_q <= 2048 else 4

 pre_tokens = max_seqlen_q - 1
 next_tokens = 0
 if window_size[0] >= 0 or window_size[1] >= 0:
 pre_tokens = max_seqlen_q - 1 if window_size[0] < 0 else window_size[0]
 next_tokens = max_seqlen_q - 1 if window_size[1] < 0 else window_size[1]

 out = torch_npu.npu_fusion_attention(
 q,
Expand All@@ -114,8 +140,8 @@ def flash_attn_varlen_func(
 "TND",
 atten_mask=attention_mask,
 scale=softmax_scale,
 pre_tockens=q.shape[0], # seq_len
 next_tockens=0, # 0
 pre_tockens=pre_tokens,
 next_tockens=next_tokens,
 keep_prob=1 - dropout_p,
 sparse_mode=sparse_mode,
 actual_seq_qlen=cu_seqlens_q,
Expand DownExpand Up@@ -143,21 +169,14 @@ def flash_attn_qkvpacked_func(
 seqlen_qkv = qkv.shape[1]
 head_num = q.shape[-2]

 if seqlen_qkv < 2048:
 sparse_mode = 0
 else:
 sparse_mode = 2

 seqlen_qkv = min(qkv.shape[1], 2048)
 attention_mask = get_attention_mask(seqlen_qkv, causal, window_size)
 sparse_mode = 0 if attention_mask is None or seqlen_qkv <= 2048 else 4

 attention_mask = (
 torch.triu(
 torch.ones([seqlen_qkv, seqlen_qkv], dtype=torch.bool, device=q.device),
 diagonal=1,
 )
 if causal
 else None
 )
 pre_tokens = seqlen_qkv - 1
 next_tokens = 0
 if window_size[0] >= 0 or window_size[1] >= 0:
 pre_tokens = seqlen_qkv - 1 if window_size[0] < 0 else window_size[0]
 next_tokens = seqlen_qkv - 1 if window_size[1] < 0 else window_size[1]

 out = torch_npu.npu_fusion_attention(
 q,
Expand All@@ -168,8 +187,8 @@ def flash_attn_qkvpacked_func(
 atten_mask=attention_mask,
 scale=softmax_scale,
 keep_prob=1 - dropout_p,
 pre_tockens=seqlen_qkv,
 next_tockens=0,
 pre_tockens=pre_tokens,
 next_tockens=next_tokens,
 sparse_mode=sparse_mode,
 )[0]

Expand All@@ -192,26 +211,19 @@ def flash_attn_kvpacked_func(
 k = kv[:, :, 0]
 v = kv[:, :, 1]

s0 = q.shape[1]
s1 = kv.shape[1]
seqlen_q = q.shape[1]
seqlen_kv = kv.shape[1]
 head_num = q.shape[-2]

 if s0 == s1 and s0 < 2048 and s1 < 2048:
 sparse_mode = 0
 else:
 sparse_mode = 2

 seqlen_q = min(s0, 2048)
 seqlen_k = min(s1, 2048)
 assert seqlen_q == seqlen_kv, "Npu currently only supports seqlen_q = seqlen_kv."
 attention_mask = get_attention_mask(seqlen_q, causal, window_size)
 sparse_mode = 0 if attention_mask is None or seqlen_q <= 2048 else 4

 attention_mask = (
 torch.triu(
 torch.ones([seqlen_q, seqlen_k], dtype=torch.bool, device=q.device),
 diagonal=1,
 )
 if causal
 else None
 )
 pre_tokens = seqlen_q - 1
 next_tokens = 0
 if window_size[0] >= 0 or window_size[1] >= 0:
 pre_tokens = seqlen_q - 1 if window_size[0] < 0 else window_size[0]
 next_tokens = seqlen_q - 1 if window_size[1] < 0 else window_size[1]

 out = torch_npu.npu_fusion_attention(
 q,
Expand All@@ -222,8 +234,8 @@ def flash_attn_kvpacked_func(
 atten_mask=attention_mask,
 scale=softmax_scale,
 keep_prob=1 - dropout_p,
 pre_tockens=seqlen_k,
 next_tockens=0,
 pre_tockens=pre_tokens,
 next_tockens=next_tokens,
 sparse_mode=sparse_mode,
 )[0]

Expand All@@ -247,32 +259,30 @@ def flash_attn_varlen_qkvpacked_func(
 q = qkv[:, 0]
 k = qkv[:, 1]
 v = qkv[:, 2]
 n = q.shape[1]
 if max_seqlen > 2048:
 sparse_mode = 2
 else:
 sparse_mode = 0
 head_num = q.shape[1]

 cu_seqlens_q = cu_seqlens[1:].tolist()
 cu_seqlens_k = cu_seqlens[1:].tolist()
 seqlen = min(max_seqlen, 2048)
 attention_mask = (
 torch.triu(
 torch.ones([seqlen, seqlen], dtype=torch.bool, device=q.device),
 diagonal=1,
 )
 if causal
 else None
 )

 attention_mask = get_attention_mask(max_seqlen, causal, window_size)
 sparse_mode = 0 if attention_mask is None or max_seqlen <= 2048 else 4

 pre_tokens = max_seqlen - 1
 next_tokens = 0
 if window_size[0] >= 0 or window_size[1] >= 0:
 pre_tokens = max_seqlen - 1 if window_size[0] < 0 else window_size[0]
 next_tokens = max_seqlen - 1 if window_size[1] < 0 else window_size[1]

 out = torch_npu.npu_fusion_attention(
 q,
 k,
 v,
n,
head_num,
 "TND",
 atten_mask=attention_mask,
 scale=softmax_scale,
 pre_tockens=q.shape[0], # seq_len
 next_tockens=0, # 0
 pre_tockens=pre_tokens,
 next_tockens=next_tokens,
 keep_prob=1 - dropout_p,
 sparse_mode=sparse_mode,
 actual_seq_qlen=cu_seqlens_q,
Expand DownExpand Up@@ -300,35 +310,32 @@ def flash_attn_varlen_kvpacked_func(
 softmax_scale = q.shape[-1] ** (-0.5)
 k = kv[:, 0]
 v = kv[:, 1]
n = q.shape[1]
head_num = q.shape[1]
 cu_seqlens_q = cu_seqlens_q[1:].tolist()
 cu_seqlens_k = cu_seqlens_k[1:].tolist()
 seqlen_q = min(max_seqlen_q, 2048)
 seqlen_k = min(max_seqlen_k, 2048)

 if max_seqlen_q > 2048:
 sparse_mode = 2
 else:
 sparse_mode = 0

 attention_mask = (
 torch.triu(
 torch.ones([seqlen_q, seqlen_k], dtype=torch.bool, device=q.device),
 diagonal=1,
 )
 if causal
 else None
 )

 assert (
 max_seqlen_q == max_seqlen_k
 ), "Npu currently only supports max_seqlen_q = max_seqlen_k."
 attention_mask = get_attention_mask(max_seqlen_q, causal, window_size)
 sparse_mode = 0 if attention_mask is None or max_seqlen_q <= 2048 else 4

 pre_tokens = max_seqlen_q - 1
 next_tokens = 0
 if window_size[0] >= 0 or window_size[1] >= 0:
 pre_tokens = max_seqlen_q - 1 if window_size[0] < 0 else window_size[0]
 next_tokens = max_seqlen_k - 1 if window_size[1] < 0 else window_size[1]

 out = torch_npu.npu_fusion_attention(
 q,
 k,
 v,
n,
head_num,
 "TND",
 atten_mask=attention_mask,
 scale=softmax_scale,
 pre_tockens=q.shape[0], # seq_len
 next_tockens=0, # 0
 pre_tockens=pre_tokens,
 next_tockens=next_tokens,
 keep_prob=1 - dropout_p,
 sparse_mode=sparse_mode,
 actual_seq_qlen=cu_seqlens_q,
Expand Down
Loading