[WIP] Add cudnn fused multi-head attention for JAX by zlsh80826 · Pull Request #105 · NVIDIA/TransformerEngine

zlsh80826 · 2023-03-16T10:33:35Z

Move scale_factor to core attention to align fused multi-head attention implementation [transformer.py, module.py]
Add cudnn-frontend submodule into 3rdparty/cudnn-frontend
Add cudnn-frontend fused multi-head attention with both self attention and cross attention
Fused multi-head attention will be auto enabled if the network satisify the following rules
- not decode
- not transpose_batch_sequence
- fuse_qkv
- dropout_rate = 0 (dropout can be fused into FMHA, but we lack model convergence test. Will add it in the future)
- dtype = bfloat16 or float16
- q_seqlen and kv_seqlen = 128 or 256 or 384 or 512
MHA supports DP and TP sharding

Signed-off-by: Reese Wang <rewang@nvidia.com>

zlsh80826 · 2023-03-16T11:27:00Z

/te-ci

Signed-off-by: Reese Wang <rewang@nvidia.com>

zlsh80826 · 2023-03-16T15:26:29Z

/te-ci

Signed-off-by: Reese Wang <rewang@nvidia.com>

ptrendx · 2023-03-16T17:50:03Z

The proposed API is completely different than the rest of TE APIs and is not acceptable.
@cyanguwa is working on a better abstraction for multiple different fused MHA algorithms, please coordinate with her to plug under that API.

OK, I will work with @cyanguwa to see how to integrate the FMHA

ptrendx · 2023-04-06T23:37:43Z

+jax.config.update('experimental_xmap_spmd_lowering_manual', True)
+
+
+def self_fmha(qkv: jnp.ndarray,


Is this the public API that people should use? If it is, that is not great:
the name is not descriptive - what is "fmha"?
what it implements is not "Multihead attention" but rather "Dot Product Attention" from the "Attention is all you need paper", so should not really use mha in the name
docstrings are missing
it will not extend to the non-zero dropout case since there you need offset too. Even if we do not implement that case right away, we should nto expose the API that will prevent us from doing so in the future.
what is "scaling factor"?
why isn't it a module like the other high-level APIs?

The public API that people should use is MultiHeadAttention(TE API doc), the self_fmha and cross_fmha are custom_calls used internally for JAX-TE instead of public users.
The "fmha" naming is inherited by apex, I can rename them to self(cross)_fused_dot_product_attention if they are preferred in TE.
As described above, the custom_calls are not intented to public users. The doc strings for MultiHeadAttetion are written on MultiHeadAttention
The original cuDNN sample has missed the offset and used a CPU-side seed. But yes, we should use both seed and offset and keep them as the device pointers.
Is scaling_factor ambiguous here? From the "Attention Is All You Need paper", the scaling factor is rsqrt(head_dim) that scaled to dot_product(Q, tranpose(K))
MultiHeadAttention(TE API doc) is implemented as a module. self_fmha is a custom_call wrapper, it is the same level as fp8_dot
I will add some commits for fmha renaming and seed/offset changes. About the naming, which is preferred in TE? self(cross)_fused_dot_product_attention, self(cross)_fused_attention or other?

Oh boy... This means we have a divergence between pyTorch and JAX APIs. pyTorch doesn't expose the MultiHeadAttention API and now this kind of forces us to do that. @ksivaman for visibility. Still, for consistency I would really like for JAX to expose DotProductAttention to be in line with pyTorch (and then of course the MultiHeadAttention API can use that similarly to how it uses the other modules).
If those are not public APIs then most of my comments are not applicable. I would probably settle for self/cross_fused_attention.
About scaling factor - sorry, my bad, I forgot about this scaling and thought about this qk layer scaling that megatron uses. It's fine as is 👍.

I would really like for JAX to expose DotProductAttention to be in line with pyTorch
Reese will submit another PR to solve this issue

Yes, I will modularize DotProductAttention for JAX in the coming days before integrating the fused attentions.

timmoon10 · 2023-04-24T18:35:55Z

PyTorch support for fused attention is added in #155.

zlsh80826 · 2023-04-25T14:38:11Z

Close this PR and I will open a new one for the fp16/bf16 fused attention (max_seq <= 512)

zlsh80826 · 2023-03-16T10:33:35Z

Move scale_factor to core attention to align fused multi-head attention implementation [transformer.py, module.py]
Add cudnn-frontend submodule into 3rdparty/cudnn-frontend
Add cudnn-frontend fused multi-head attention with both self attention and cross attention
Fused multi-head attention will be auto enabled if the network satisify the following rules
- not decode
- not transpose_batch_sequence
- fuse_qkv
- dropout_rate = 0 (dropout can be fused into FMHA, but we lack model convergence test. Will add it in the future)
- dtype = bfloat16 or float16
- q_seqlen and kv_seqlen = 128 or 256 or 384 or 512
MHA supports DP and TP sharding

Signed-off-by: Reese Wang <rewang@nvidia.com>

zlsh80826 · 2023-03-16T11:27:00Z

/te-ci

Signed-off-by: Reese Wang <rewang@nvidia.com>

zlsh80826 · 2023-03-16T15:26:29Z

/te-ci

Signed-off-by: Reese Wang <rewang@nvidia.com>

ptrendx · 2023-03-16T17:50:03Z

The proposed API is completely different than the rest of TE APIs and is not acceptable.
@cyanguwa is working on a better abstraction for multiple different fused MHA algorithms, please coordinate with her to plug under that API.

OK, I will work with @cyanguwa to see how to integrate the FMHA

ptrendx · 2023-04-06T23:37:43Z

+jax.config.update('experimental_xmap_spmd_lowering_manual', True)
+
+
+def self_fmha(qkv: jnp.ndarray,


Is this the public API that people should use? If it is, that is not great:
the name is not descriptive - what is "fmha"?
what it implements is not "Multihead attention" but rather "Dot Product Attention" from the "Attention is all you need paper", so should not really use mha in the name
docstrings are missing
it will not extend to the non-zero dropout case since there you need offset too. Even if we do not implement that case right away, we should nto expose the API that will prevent us from doing so in the future.
what is "scaling factor"?
why isn't it a module like the other high-level APIs?

The public API that people should use is MultiHeadAttention(TE API doc), the self_fmha and cross_fmha are custom_calls used internally for JAX-TE instead of public users.
The "fmha" naming is inherited by apex, I can rename them to self(cross)_fused_dot_product_attention if they are preferred in TE.
As described above, the custom_calls are not intented to public users. The doc strings for MultiHeadAttetion are written on MultiHeadAttention
The original cuDNN sample has missed the offset and used a CPU-side seed. But yes, we should use both seed and offset and keep them as the device pointers.
Is scaling_factor ambiguous here? From the "Attention Is All You Need paper", the scaling factor is rsqrt(head_dim) that scaled to dot_product(Q, tranpose(K))
MultiHeadAttention(TE API doc) is implemented as a module. self_fmha is a custom_call wrapper, it is the same level as fp8_dot
I will add some commits for fmha renaming and seed/offset changes. About the naming, which is preferred in TE? self(cross)_fused_dot_product_attention, self(cross)_fused_attention or other?

Oh boy... This means we have a divergence between pyTorch and JAX APIs. pyTorch doesn't expose the MultiHeadAttention API and now this kind of forces us to do that. @ksivaman for visibility. Still, for consistency I would really like for JAX to expose DotProductAttention to be in line with pyTorch (and then of course the MultiHeadAttention API can use that similarly to how it uses the other modules).
If those are not public APIs then most of my comments are not applicable. I would probably settle for self/cross_fused_attention.
About scaling factor - sorry, my bad, I forgot about this scaling and thought about this qk layer scaling that megatron uses. It's fine as is 👍.

I would really like for JAX to expose DotProductAttention to be in line with pyTorch
Reese will submit another PR to solve this issue

Yes, I will modularize DotProductAttention for JAX in the coming days before integrating the fused attentions.

timmoon10 · 2023-04-24T18:35:55Z

PyTorch support for fused attention is added in #155.

zlsh80826 · 2023-04-25T14:38:11Z

Close this PR and I will open a new one for the fp16/bf16 fused attention (max_seq <= 512)

zlsh80826 · 2023-03-16T10:33:35Z

Move scale_factor to core attention to align fused multi-head attention implementation [transformer.py, module.py]
Add cudnn-frontend submodule into 3rdparty/cudnn-frontend
Add cudnn-frontend fused multi-head attention with both self attention and cross attention
Fused multi-head attention will be auto enabled if the network satisify the following rules
- not decode
- not transpose_batch_sequence
- fuse_qkv
- dropout_rate = 0 (dropout can be fused into FMHA, but we lack model convergence test. Will add it in the future)
- dtype = bfloat16 or float16
- q_seqlen and kv_seqlen = 128 or 256 or 384 or 512
MHA supports DP and TP sharding

Signed-off-by: Reese Wang <rewang@nvidia.com>

zlsh80826 · 2023-03-16T11:27:00Z

/te-ci

Signed-off-by: Reese Wang <rewang@nvidia.com>

zlsh80826 · 2023-03-16T15:26:29Z

/te-ci

Signed-off-by: Reese Wang <rewang@nvidia.com>

ptrendx · 2023-03-16T17:50:03Z

The proposed API is completely different than the rest of TE APIs and is not acceptable.
@cyanguwa is working on a better abstraction for multiple different fused MHA algorithms, please coordinate with her to plug under that API.

OK, I will work with @cyanguwa to see how to integrate the FMHA

ptrendx · 2023-04-06T23:37:43Z

+jax.config.update('experimental_xmap_spmd_lowering_manual', True)
+
+
+def self_fmha(qkv: jnp.ndarray,


Is this the public API that people should use? If it is, that is not great:
the name is not descriptive - what is "fmha"?
what it implements is not "Multihead attention" but rather "Dot Product Attention" from the "Attention is all you need paper", so should not really use mha in the name
docstrings are missing
it will not extend to the non-zero dropout case since there you need offset too. Even if we do not implement that case right away, we should nto expose the API that will prevent us from doing so in the future.
what is "scaling factor"?
why isn't it a module like the other high-level APIs?

The public API that people should use is MultiHeadAttention(TE API doc), the self_fmha and cross_fmha are custom_calls used internally for JAX-TE instead of public users.
The "fmha" naming is inherited by apex, I can rename them to self(cross)_fused_dot_product_attention if they are preferred in TE.
As described above, the custom_calls are not intented to public users. The doc strings for MultiHeadAttetion are written on MultiHeadAttention
The original cuDNN sample has missed the offset and used a CPU-side seed. But yes, we should use both seed and offset and keep them as the device pointers.
Is scaling_factor ambiguous here? From the "Attention Is All You Need paper", the scaling factor is rsqrt(head_dim) that scaled to dot_product(Q, tranpose(K))
MultiHeadAttention(TE API doc) is implemented as a module. self_fmha is a custom_call wrapper, it is the same level as fp8_dot
I will add some commits for fmha renaming and seed/offset changes. About the naming, which is preferred in TE? self(cross)_fused_dot_product_attention, self(cross)_fused_attention or other?

Oh boy... This means we have a divergence between pyTorch and JAX APIs. pyTorch doesn't expose the MultiHeadAttention API and now this kind of forces us to do that. @ksivaman for visibility. Still, for consistency I would really like for JAX to expose DotProductAttention to be in line with pyTorch (and then of course the MultiHeadAttention API can use that similarly to how it uses the other modules).
If those are not public APIs then most of my comments are not applicable. I would probably settle for self/cross_fused_attention.
About scaling factor - sorry, my bad, I forgot about this scaling and thought about this qk layer scaling that megatron uses. It's fine as is 👍.

I would really like for JAX to expose DotProductAttention to be in line with pyTorch
Reese will submit another PR to solve this issue

Yes, I will modularize DotProductAttention for JAX in the coming days before integrating the fused attentions.

timmoon10 · 2023-04-24T18:35:55Z

PyTorch support for fused attention is added in #155.

zlsh80826 · 2023-04-25T14:38:11Z

Close this PR and I will open a new one for the fp16/bf16 fused attention (max_seq <= 512)

zlsh80826 · 2023-03-16T10:33:35Z

Move scale_factor to core attention to align fused multi-head attention implementation [transformer.py, module.py]
Add cudnn-frontend submodule into 3rdparty/cudnn-frontend
Add cudnn-frontend fused multi-head attention with both self attention and cross attention
Fused multi-head attention will be auto enabled if the network satisify the following rules
- not decode
- not transpose_batch_sequence
- fuse_qkv
- dropout_rate = 0 (dropout can be fused into FMHA, but we lack model convergence test. Will add it in the future)
- dtype = bfloat16 or float16
- q_seqlen and kv_seqlen = 128 or 256 or 384 or 512
MHA supports DP and TP sharding

Signed-off-by: Reese Wang <rewang@nvidia.com>

zlsh80826 · 2023-03-16T11:27:00Z

/te-ci

Signed-off-by: Reese Wang <rewang@nvidia.com>

zlsh80826 · 2023-03-16T15:26:29Z

/te-ci

Signed-off-by: Reese Wang <rewang@nvidia.com>

ptrendx · 2023-03-16T17:50:03Z

The proposed API is completely different than the rest of TE APIs and is not acceptable.
@cyanguwa is working on a better abstraction for multiple different fused MHA algorithms, please coordinate with her to plug under that API.

OK, I will work with @cyanguwa to see how to integrate the FMHA

ptrendx · 2023-04-06T23:37:43Z

+jax.config.update('experimental_xmap_spmd_lowering_manual', True)
+
+
+def self_fmha(qkv: jnp.ndarray,


Is this the public API that people should use? If it is, that is not great:
the name is not descriptive - what is "fmha"?
what it implements is not "Multihead attention" but rather "Dot Product Attention" from the "Attention is all you need paper", so should not really use mha in the name
docstrings are missing
it will not extend to the non-zero dropout case since there you need offset too. Even if we do not implement that case right away, we should nto expose the API that will prevent us from doing so in the future.
what is "scaling factor"?
why isn't it a module like the other high-level APIs?

The public API that people should use is MultiHeadAttention(TE API doc), the self_fmha and cross_fmha are custom_calls used internally for JAX-TE instead of public users.
The "fmha" naming is inherited by apex, I can rename them to self(cross)_fused_dot_product_attention if they are preferred in TE.
As described above, the custom_calls are not intented to public users. The doc strings for MultiHeadAttetion are written on MultiHeadAttention
The original cuDNN sample has missed the offset and used a CPU-side seed. But yes, we should use both seed and offset and keep them as the device pointers.
Is scaling_factor ambiguous here? From the "Attention Is All You Need paper", the scaling factor is rsqrt(head_dim) that scaled to dot_product(Q, tranpose(K))
MultiHeadAttention(TE API doc) is implemented as a module. self_fmha is a custom_call wrapper, it is the same level as fp8_dot
I will add some commits for fmha renaming and seed/offset changes. About the naming, which is preferred in TE? self(cross)_fused_dot_product_attention, self(cross)_fused_attention or other?

Oh boy... This means we have a divergence between pyTorch and JAX APIs. pyTorch doesn't expose the MultiHeadAttention API and now this kind of forces us to do that. @ksivaman for visibility. Still, for consistency I would really like for JAX to expose DotProductAttention to be in line with pyTorch (and then of course the MultiHeadAttention API can use that similarly to how it uses the other modules).
If those are not public APIs then most of my comments are not applicable. I would probably settle for self/cross_fused_attention.
About scaling factor - sorry, my bad, I forgot about this scaling and thought about this qk layer scaling that megatron uses. It's fine as is 👍.

I would really like for JAX to expose DotProductAttention to be in line with pyTorch
Reese will submit another PR to solve this issue

Yes, I will modularize DotProductAttention for JAX in the coming days before integrating the fused attentions.

timmoon10 · 2023-04-24T18:35:55Z

PyTorch support for fused attention is added in #155.

zlsh80826 · 2023-04-25T14:38:11Z

Close this PR and I will open a new one for the fp16/bf16 fused attention (max_seq <= 512)

zlsh80826 · 2023-03-16T10:33:35Z

Move scale_factor to core attention to align fused multi-head attention implementation [transformer.py, module.py]
Add cudnn-frontend submodule into 3rdparty/cudnn-frontend
Add cudnn-frontend fused multi-head attention with both self attention and cross attention
Fused multi-head attention will be auto enabled if the network satisify the following rules
- not decode
- not transpose_batch_sequence
- fuse_qkv
- dropout_rate = 0 (dropout can be fused into FMHA, but we lack model convergence test. Will add it in the future)
- dtype = bfloat16 or float16
- q_seqlen and kv_seqlen = 128 or 256 or 384 or 512
MHA supports DP and TP sharding

Signed-off-by: Reese Wang <rewang@nvidia.com>

zlsh80826 · 2023-03-16T11:27:00Z

/te-ci

Signed-off-by: Reese Wang <rewang@nvidia.com>

zlsh80826 · 2023-03-16T15:26:29Z

/te-ci

Signed-off-by: Reese Wang <rewang@nvidia.com>

ptrendx · 2023-03-16T17:50:03Z

The proposed API is completely different than the rest of TE APIs and is not acceptable.
@cyanguwa is working on a better abstraction for multiple different fused MHA algorithms, please coordinate with her to plug under that API.

OK, I will work with @cyanguwa to see how to integrate the FMHA

ptrendx · 2023-04-06T23:37:43Z

+jax.config.update('experimental_xmap_spmd_lowering_manual', True)
+
+
+def self_fmha(qkv: jnp.ndarray,


Is this the public API that people should use? If it is, that is not great:
the name is not descriptive - what is "fmha"?
what it implements is not "Multihead attention" but rather "Dot Product Attention" from the "Attention is all you need paper", so should not really use mha in the name
docstrings are missing
it will not extend to the non-zero dropout case since there you need offset too. Even if we do not implement that case right away, we should nto expose the API that will prevent us from doing so in the future.
what is "scaling factor"?
why isn't it a module like the other high-level APIs?

The public API that people should use is MultiHeadAttention(TE API doc), the self_fmha and cross_fmha are custom_calls used internally for JAX-TE instead of public users.
The "fmha" naming is inherited by apex, I can rename them to self(cross)_fused_dot_product_attention if they are preferred in TE.
As described above, the custom_calls are not intented to public users. The doc strings for MultiHeadAttetion are written on MultiHeadAttention
The original cuDNN sample has missed the offset and used a CPU-side seed. But yes, we should use both seed and offset and keep them as the device pointers.
Is scaling_factor ambiguous here? From the "Attention Is All You Need paper", the scaling factor is rsqrt(head_dim) that scaled to dot_product(Q, tranpose(K))
MultiHeadAttention(TE API doc) is implemented as a module. self_fmha is a custom_call wrapper, it is the same level as fp8_dot
I will add some commits for fmha renaming and seed/offset changes. About the naming, which is preferred in TE? self(cross)_fused_dot_product_attention, self(cross)_fused_attention or other?

Oh boy... This means we have a divergence between pyTorch and JAX APIs. pyTorch doesn't expose the MultiHeadAttention API and now this kind of forces us to do that. @ksivaman for visibility. Still, for consistency I would really like for JAX to expose DotProductAttention to be in line with pyTorch (and then of course the MultiHeadAttention API can use that similarly to how it uses the other modules).
If those are not public APIs then most of my comments are not applicable. I would probably settle for self/cross_fused_attention.
About scaling factor - sorry, my bad, I forgot about this scaling and thought about this qk layer scaling that megatron uses. It's fine as is 👍.

I would really like for JAX to expose DotProductAttention to be in line with pyTorch
Reese will submit another PR to solve this issue

Yes, I will modularize DotProductAttention for JAX in the coming days before integrating the fused attentions.

timmoon10 · 2023-04-24T18:35:55Z

PyTorch support for fused attention is added in #155.

zlsh80826 · 2023-04-25T14:38:11Z

Close this PR and I will open a new one for the fp16/bf16 fused attention (max_seq <= 512)

zlsh80826 · 2023-03-16T10:33:35Z

Move scale_factor to core attention to align fused multi-head attention implementation [transformer.py, module.py]
Add cudnn-frontend submodule into 3rdparty/cudnn-frontend
Add cudnn-frontend fused multi-head attention with both self attention and cross attention
Fused multi-head attention will be auto enabled if the network satisify the following rules
- not decode
- not transpose_batch_sequence
- fuse_qkv
- dropout_rate = 0 (dropout can be fused into FMHA, but we lack model convergence test. Will add it in the future)
- dtype = bfloat16 or float16
- q_seqlen and kv_seqlen = 128 or 256 or 384 or 512
MHA supports DP and TP sharding

Signed-off-by: Reese Wang <rewang@nvidia.com>

zlsh80826 · 2023-03-16T11:27:00Z

/te-ci

Signed-off-by: Reese Wang <rewang@nvidia.com>

zlsh80826 · 2023-03-16T15:26:29Z

/te-ci

Signed-off-by: Reese Wang <rewang@nvidia.com>

ptrendx · 2023-03-16T17:50:03Z

The proposed API is completely different than the rest of TE APIs and is not acceptable.
@cyanguwa is working on a better abstraction for multiple different fused MHA algorithms, please coordinate with her to plug under that API.

OK, I will work with @cyanguwa to see how to integrate the FMHA

ptrendx · 2023-04-06T23:37:43Z

+jax.config.update('experimental_xmap_spmd_lowering_manual', True)
+
+
+def self_fmha(qkv: jnp.ndarray,


Is this the public API that people should use? If it is, that is not great:
the name is not descriptive - what is "fmha"?
what it implements is not "Multihead attention" but rather "Dot Product Attention" from the "Attention is all you need paper", so should not really use mha in the name
docstrings are missing
it will not extend to the non-zero dropout case since there you need offset too. Even if we do not implement that case right away, we should nto expose the API that will prevent us from doing so in the future.
what is "scaling factor"?
why isn't it a module like the other high-level APIs?

The public API that people should use is MultiHeadAttention(TE API doc), the self_fmha and cross_fmha are custom_calls used internally for JAX-TE instead of public users.
The "fmha" naming is inherited by apex, I can rename them to self(cross)_fused_dot_product_attention if they are preferred in TE.
As described above, the custom_calls are not intented to public users. The doc strings for MultiHeadAttetion are written on MultiHeadAttention
The original cuDNN sample has missed the offset and used a CPU-side seed. But yes, we should use both seed and offset and keep them as the device pointers.
Is scaling_factor ambiguous here? From the "Attention Is All You Need paper", the scaling factor is rsqrt(head_dim) that scaled to dot_product(Q, tranpose(K))
MultiHeadAttention(TE API doc) is implemented as a module. self_fmha is a custom_call wrapper, it is the same level as fp8_dot
I will add some commits for fmha renaming and seed/offset changes. About the naming, which is preferred in TE? self(cross)_fused_dot_product_attention, self(cross)_fused_attention or other?

Oh boy... This means we have a divergence between pyTorch and JAX APIs. pyTorch doesn't expose the MultiHeadAttention API and now this kind of forces us to do that. @ksivaman for visibility. Still, for consistency I would really like for JAX to expose DotProductAttention to be in line with pyTorch (and then of course the MultiHeadAttention API can use that similarly to how it uses the other modules).
If those are not public APIs then most of my comments are not applicable. I would probably settle for self/cross_fused_attention.
About scaling factor - sorry, my bad, I forgot about this scaling and thought about this qk layer scaling that megatron uses. It's fine as is 👍.

I would really like for JAX to expose DotProductAttention to be in line with pyTorch
Reese will submit another PR to solve this issue

Yes, I will modularize DotProductAttention for JAX in the coming days before integrating the fused attentions.

timmoon10 · 2023-04-24T18:35:55Z

PyTorch support for fused attention is added in #155.

zlsh80826 · 2023-04-25T14:38:11Z

Close this PR and I will open a new one for the fp16/bf16 fused attention (max_seq <= 512)

zlsh80826 · 2023-03-16T10:33:35Z

Move scale_factor to core attention to align fused multi-head attention implementation [transformer.py, module.py]
Add cudnn-frontend submodule into 3rdparty/cudnn-frontend
Add cudnn-frontend fused multi-head attention with both self attention and cross attention
Fused multi-head attention will be auto enabled if the network satisify the following rules
- not decode
- not transpose_batch_sequence
- fuse_qkv
- dropout_rate = 0 (dropout can be fused into FMHA, but we lack model convergence test. Will add it in the future)
- dtype = bfloat16 or float16
- q_seqlen and kv_seqlen = 128 or 256 or 384 or 512
MHA supports DP and TP sharding

Signed-off-by: Reese Wang <rewang@nvidia.com>

zlsh80826 · 2023-03-16T11:27:00Z

/te-ci

Signed-off-by: Reese Wang <rewang@nvidia.com>

zlsh80826 · 2023-03-16T15:26:29Z

/te-ci

Signed-off-by: Reese Wang <rewang@nvidia.com>

ptrendx · 2023-03-16T17:50:03Z

The proposed API is completely different than the rest of TE APIs and is not acceptable.
@cyanguwa is working on a better abstraction for multiple different fused MHA algorithms, please coordinate with her to plug under that API.

OK, I will work with @cyanguwa to see how to integrate the FMHA

ptrendx · 2023-04-06T23:37:43Z

+jax.config.update('experimental_xmap_spmd_lowering_manual', True)
+
+
+def self_fmha(qkv: jnp.ndarray,


Is this the public API that people should use? If it is, that is not great:
the name is not descriptive - what is "fmha"?
what it implements is not "Multihead attention" but rather "Dot Product Attention" from the "Attention is all you need paper", so should not really use mha in the name
docstrings are missing
it will not extend to the non-zero dropout case since there you need offset too. Even if we do not implement that case right away, we should nto expose the API that will prevent us from doing so in the future.
what is "scaling factor"?
why isn't it a module like the other high-level APIs?

The public API that people should use is MultiHeadAttention(TE API doc), the self_fmha and cross_fmha are custom_calls used internally for JAX-TE instead of public users.
The "fmha" naming is inherited by apex, I can rename them to self(cross)_fused_dot_product_attention if they are preferred in TE.
As described above, the custom_calls are not intented to public users. The doc strings for MultiHeadAttetion are written on MultiHeadAttention
The original cuDNN sample has missed the offset and used a CPU-side seed. But yes, we should use both seed and offset and keep them as the device pointers.
Is scaling_factor ambiguous here? From the "Attention Is All You Need paper", the scaling factor is rsqrt(head_dim) that scaled to dot_product(Q, tranpose(K))
MultiHeadAttention(TE API doc) is implemented as a module. self_fmha is a custom_call wrapper, it is the same level as fp8_dot
I will add some commits for fmha renaming and seed/offset changes. About the naming, which is preferred in TE? self(cross)_fused_dot_product_attention, self(cross)_fused_attention or other?

Oh boy... This means we have a divergence between pyTorch and JAX APIs. pyTorch doesn't expose the MultiHeadAttention API and now this kind of forces us to do that. @ksivaman for visibility. Still, for consistency I would really like for JAX to expose DotProductAttention to be in line with pyTorch (and then of course the MultiHeadAttention API can use that similarly to how it uses the other modules).
If those are not public APIs then most of my comments are not applicable. I would probably settle for self/cross_fused_attention.
About scaling factor - sorry, my bad, I forgot about this scaling and thought about this qk layer scaling that megatron uses. It's fine as is 👍.

I would really like for JAX to expose DotProductAttention to be in line with pyTorch
Reese will submit another PR to solve this issue

Yes, I will modularize DotProductAttention for JAX in the coming days before integrating the fused attentions.

timmoon10 · 2023-04-24T18:35:55Z

PyTorch support for fused attention is added in #155.

zlsh80826 · 2023-04-25T14:38:11Z

Close this PR and I will open a new one for the fp16/bf16 fused attention (max_seq <= 512)

zlsh80826 · 2023-03-16T10:33:35Z

Move scale_factor to core attention to align fused multi-head attention implementation [transformer.py, module.py]
Add cudnn-frontend submodule into 3rdparty/cudnn-frontend
Add cudnn-frontend fused multi-head attention with both self attention and cross attention
Fused multi-head attention will be auto enabled if the network satisify the following rules
- not decode
- not transpose_batch_sequence
- fuse_qkv
- dropout_rate = 0 (dropout can be fused into FMHA, but we lack model convergence test. Will add it in the future)
- dtype = bfloat16 or float16
- q_seqlen and kv_seqlen = 128 or 256 or 384 or 512
MHA supports DP and TP sharding

Signed-off-by: Reese Wang <rewang@nvidia.com>

zlsh80826 · 2023-03-16T11:27:00Z

/te-ci

Signed-off-by: Reese Wang <rewang@nvidia.com>

zlsh80826 · 2023-03-16T15:26:29Z

/te-ci

Signed-off-by: Reese Wang <rewang@nvidia.com>

ptrendx · 2023-03-16T17:50:03Z

The proposed API is completely different than the rest of TE APIs and is not acceptable.
@cyanguwa is working on a better abstraction for multiple different fused MHA algorithms, please coordinate with her to plug under that API.

OK, I will work with @cyanguwa to see how to integrate the FMHA

ptrendx · 2023-04-06T23:37:43Z

+jax.config.update('experimental_xmap_spmd_lowering_manual', True)
+
+
+def self_fmha(qkv: jnp.ndarray,


Is this the public API that people should use? If it is, that is not great:
the name is not descriptive - what is "fmha"?
what it implements is not "Multihead attention" but rather "Dot Product Attention" from the "Attention is all you need paper", so should not really use mha in the name
docstrings are missing
it will not extend to the non-zero dropout case since there you need offset too. Even if we do not implement that case right away, we should nto expose the API that will prevent us from doing so in the future.
what is "scaling factor"?
why isn't it a module like the other high-level APIs?

The public API that people should use is MultiHeadAttention(TE API doc), the self_fmha and cross_fmha are custom_calls used internally for JAX-TE instead of public users.
The "fmha" naming is inherited by apex, I can rename them to self(cross)_fused_dot_product_attention if they are preferred in TE.
As described above, the custom_calls are not intented to public users. The doc strings for MultiHeadAttetion are written on MultiHeadAttention
The original cuDNN sample has missed the offset and used a CPU-side seed. But yes, we should use both seed and offset and keep them as the device pointers.
Is scaling_factor ambiguous here? From the "Attention Is All You Need paper", the scaling factor is rsqrt(head_dim) that scaled to dot_product(Q, tranpose(K))
MultiHeadAttention(TE API doc) is implemented as a module. self_fmha is a custom_call wrapper, it is the same level as fp8_dot
I will add some commits for fmha renaming and seed/offset changes. About the naming, which is preferred in TE? self(cross)_fused_dot_product_attention, self(cross)_fused_attention or other?

Oh boy... This means we have a divergence between pyTorch and JAX APIs. pyTorch doesn't expose the MultiHeadAttention API and now this kind of forces us to do that. @ksivaman for visibility. Still, for consistency I would really like for JAX to expose DotProductAttention to be in line with pyTorch (and then of course the MultiHeadAttention API can use that similarly to how it uses the other modules).
If those are not public APIs then most of my comments are not applicable. I would probably settle for self/cross_fused_attention.
About scaling factor - sorry, my bad, I forgot about this scaling and thought about this qk layer scaling that megatron uses. It's fine as is 👍.

I would really like for JAX to expose DotProductAttention to be in line with pyTorch
Reese will submit another PR to solve this issue

Yes, I will modularize DotProductAttention for JAX in the coming days before integrating the fused attentions.

timmoon10 · 2023-04-24T18:35:55Z

PyTorch support for fused attention is added in #155.

zlsh80826 · 2023-04-25T14:38:11Z

Close this PR and I will open a new one for the fp16/bf16 fused attention (max_seq <= 512)

zlsh80826 added 2 commits March 15, 2023 16:46

Refine test_layer for different transpose_bs settings
aeabddc
Signed-off-by: Reese Wang <rewang@nvidia.com>

Move query scale into core_attention
e1be3a7
Signed-off-by: Reese Wang <rewang@nvidia.com>

zlsh80826 requested review from ksivaman, mingxu1067, ptrendx and timmoon10 March 16, 2023 10:33

zlsh80826 added 3 commits March 16, 2023 14:20

Add cudnn fused multi-head attention
075102e
Signed-off-by: Reese Wang <rewang@nvidia.com>

Add sharding for JAX FMHA
f2397af
Signed-off-by: Reese Wang <rewang@nvidia.com>

Add MHA fall back warning
35f4ae1
Signed-off-by: Reese Wang <rewang@nvidia.com>

zlsh80826 force-pushed the rewang/add-fmha branch from b0d2b26 to 35f4ae1Compare March 16, 2023 14:21

Use CMake FetchContent_Declare instead of a submodule
466c06f
Signed-off-by: Reese Wang <rewang@nvidia.com>

zlsh80826 force-pushed the rewang/add-fmha branch from 9ceaa84 to 466c06fCompare March 16, 2023 15:26

Fix fmha.h license format
fd7406d
Signed-off-by: Reese Wang <rewang@nvidia.com>

ptrendx requested changes Mar 16, 2023
View reviewed changes

ptrendx requested a review from cyanguwa March 16, 2023 17:50

zlsh80826 marked this pull request as draft March 20, 2023 12:39

zlsh80826 changed the title ~~Add cudnn fused multi-head attention for JAX~~[WIP] Add cudnn fused multi-head attention for JAXMar 20, 2023

ptrendx reviewed Apr 6, 2023
View reviewed changes

zlsh80826 closed this Apr 25, 2023

zlsh80826 mentioned this pull request Apr 25, 2023
Add FP16/BF16 fused_attention support with max_seqlen=512 #175
Merged
8 tasks