Commit 827fad6

leisuzzJ石页a-r-r-o-w

authored

Improve performance of NPU FA (#12260)

Co-authored-by: J石页 <jiangshuo9@h-partners.com> Co-authored-by: Aryan <aryan@huggingface.co>

1 parent 9b721db commit 827fad6Copy full SHA for 827fad6

File tree

-3

lines changed

-3

lines changed

Lines changed: 6 additions & 3 deletions

Original file line number	Diff line number	Diff line change
`@@ -955,12 +955,13 @@ def _native_npu_attention(`
`955`	`955`	`dropout_p: float = 0.0,`
`956`	`956`	`scale: Optional[float] = None,`
`957`	`957`	`) -> torch.Tensor:`
`958`		`- return npu_fusion_attention(`
	`958`	`+ query, key, value = (x.transpose(1, 2).contiguous() for x in (query, key, value))`
	`959`	`+ out = npu_fusion_attention(`
`959`	`960`	`query,`
`960`	`961`	`key,`
`961`	`962`	`value,`
`962`		`- query.size(2), # num_heads`
`963`		`- input_layout="BSND",`
	`963`	`+ query.size(1), # num_heads`
	`964`	`+ input_layout="BNSD",`
`964`	`965`	`pse=None,`
`965`	`966`	`scale=1.0 / math.sqrt(query.shape[-1]) if scale is None else scale,`
`966`	`967`	`pre_tockens=65536,`
`@@ -969,6 +970,8 @@ def _native_npu_attention(`
`969`	`970`	`sync=False,`
`970`	`971`	`inner_precise=0,`
`971`	`972`	`)[0]`
	`973`	`+ out = out.transpose(1, 2).contiguous()`
	`974`	`+ return out`
`972`	`975`
`973`	`976`
`974`	`977`	`# Reference: https://github.com/pytorch/xla/blob/06c5533de6588f6b90aa1655d9850bcf733b90b4/torch_xla/experimental/custom_kernel.py#L853`

Comments

(0)