Jax bug fixes for the dot product attention

tests/jax/test_fused_attn.py

@@ -113,7 +113,7 @@ def customcall_cross_fused_attn(q, kv, q_token, kv_token, dropout_rng, **kwargs)
  reason="Fused attention kernel is not supported.")
 class TestSelfFusedAttnMax512():
 
- def set_input(self, b, s, h, d, dtype, attn_mask_type, pad_ratio):
+ def set_input(self, b, s, h, d, dtype, attn_mask_type, pad_ratio, with_bias):
  key = jax.random.PRNGKey(0)
  subkeys = jax.random.split(key, 2)
 
@@ -125,16 +125,17 @@ def set_input(self, b, s, h, d, dtype, attn_mask_type, pad_ratio):
 
  min_val, max_val = -1, 1
  self.qkv = jax.random.uniform(subkeys[0], qkv_shape, dtype, min_val, max_val)
- self.bias = jax.random.uniform(subkeys[1], bias_shape, dtype, min_val, max_val)
+ self.bias = jax.random.uniform(subkeys[1], bias_shape, dtype, min_val,
+ max_val) if with_bias else None
 
  self.q_token = jnp.concatenate((jnp.ones((b, self.valid_len)), jnp.zeros((b, pad_len))),
  axis=-1)
  self.kv_token = self.q_token
 
  self.scaling_factor = 1. / math.sqrt(d)
  self.dropout_probability = 0.
- self.dropout_rng = jax.random.PRNGKey(0)
- self.attn_bias_type = AttnBiasType.POST_SCALE_BIAS
+ self.dropout_rng = jax.random.PRNGKey(0) if self.dropout_probability > 0 else None
+ self.attn_bias_type = AttnBiasType.NO_BIAS if self.bias is None else AttnBiasType.POST_SCALE_BIAS
  # deterministic = not is_training
  self.deterministic = False
 
@@ -143,9 +144,17 @@ def set_input(self, b, s, h, d, dtype, attn_mask_type, pad_ratio):
  @pytest.mark.parametrize('attn_mask_type',
  [AttnMaskType.PADDING_MASK, AttnMaskType.CAUSAL_MASK])
  @pytest.mark.parametrize('pad_ratio', PAD_RATIO)
- def test_forward(self, b, s, h, d, dtype, attn_mask_type, pad_ratio):
+ @pytest.mark.parametrize('with_bias', [True, False])
+ def test_forward(self, b, s, h, d, dtype, attn_mask_type, pad_ratio, with_bias):
 
- self.set_input(b, s, h, d, dtype=dtype, attn_mask_type=attn_mask_type, pad_ratio=pad_ratio)
+ self.set_input(b,
+ s,
+ h,
+ d,
+ dtype=dtype,
+ attn_mask_type=attn_mask_type,
+ pad_ratio=pad_ratio,
+ with_bias=with_bias)
 
  primitive_out = customcall_self_fused_attn(self.qkv,
  self.bias,
@@ -183,8 +192,16 @@ def test_forward(self, b, s, h, d, dtype, attn_mask_type, pad_ratio):
  [AttnMaskType.PADDING_MASK, AttnMaskType.CAUSAL_MASK])
  @pytest.mark.parametrize('dtype', DTYPES)
  @pytest.mark.parametrize('pad_ratio', PAD_RATIO)
- def test_forward_backward(self, b, s, h, d, dtype, attn_mask_type, pad_ratio):
- self.set_input(b, s, h, d, dtype=dtype, attn_mask_type=attn_mask_type, pad_ratio=pad_ratio)
+ @pytest.mark.parametrize('with_bias', [True, False])
+ def test_forward_backward(self, b, s, h, d, dtype, attn_mask_type, pad_ratio, with_bias):
+ self.set_input(b,
+ s,
+ h,
+ d,
+ dtype=dtype,
+ attn_mask_type=attn_mask_type,
+ pad_ratio=pad_ratio,
+ with_bias=with_bias)
 
  def grad_func(fused_attn_max_512_func, *args, **kwargs):
  # Gradient is small, use a gradient multiplier to amplify the graident
@@ -221,11 +238,11 @@ def grad_func(fused_attn_max_512_func, *args, **kwargs):
  (0, 1)))
 
  primitive_out, (primitive_dqkv,
-primitive_dbeta) = jitted_primitive(self.qkv, self.bias, self.q_token,
+primitive_dbias) = jitted_primitive(self.qkv, self.bias, self.q_token,
  self.kv_token, self.dropout_rng)
 
  reference_out, (reference_dqkv,
-reference_dbeta) = jitted_reference(self.qkv, self.bias, self.q_token,
+reference_dbias) = jitted_reference(self.qkv, self.bias, self.q_token,
  self.kv_token, self.dropout_rng)
 
  np.testing.assert_allclose(jnp.asarray(primitive_out, np.float32),
@@ -261,20 +278,22 @@ def grad_func(fused_attn_max_512_func, *args, **kwargs):
  # Padded part should be 0s
  assert jnp.allclose(invalid_primitive_dqkv, jnp.zeros_like(invalid_primitive_dqkv))
 
- # dbeta valid part
- np.testing.assert_allclose(
- jnp.asarray(primitive_dbeta[:, :, :self.valid_len, :self.valid_len], np.float32),
- jnp.asarray(reference_dbeta[:, :, :self.valid_len, :self.valid_len], np.float32),
- rtol=1e-4,
- atol=3e-5)
-
- # dbeta padded part
- np.testing.assert_allclose(
- jnp.asarray(primitive_dbeta[:, :, self.valid_len:, self.valid_len:], np.float32),
- jnp.asarray(reference_dbeta[:, :, self.valid_len:, self.valid_len:], np.float32))
-
- assert jnp.allclose(primitive_dbeta[:, :, self.valid_len:, self.valid_len:],
- jnp.zeros_like(primitive_dbeta[:, :, self.valid_len:, self.valid_len:]))
+ if self.attn_bias_type != AttnBiasType.NO_BIAS:
+ # dbias valid part
+ np.testing.assert_allclose(
+ jnp.asarray(primitive_dbias[:, :, :self.valid_len, :self.valid_len], np.float32),
+ jnp.asarray(reference_dbias[:, :, :self.valid_len, :self.valid_len], np.float32),
+ rtol=1e-4,
+ atol=3e-5)
+
+ # dbias padded part
+ np.testing.assert_allclose(
+ jnp.asarray(primitive_dbias[:, :, self.valid_len:, self.valid_len:], np.float32),
+ jnp.asarray(reference_dbias[:, :, self.valid_len:, self.valid_len:], np.float32))
+
+ assert jnp.allclose(
+ primitive_dbias[:, :, self.valid_len:, self.valid_len:],
+ jnp.zeros_like(primitive_dbias[:, :, self.valid_len:, self.valid_len:]))
 
 
 @pytest.mark.skipif(not is_fused_attn_kernel_available(),

tests/jax/test_layer.py

@@ -102,6 +102,12 @@ def compare_frozen_dict(ref_fd, test_fd, rtol=1e-05, atol=1e-08):
  _KEY_OF_DROPOUT_RATE: 0.0,
  _KEY_OF_MLP_ACTIVATIONS: (('gelu', 'linear')),
  _KEY_OF_FUSE_MLP_WI: True
+}, {
+ _KEY_OF_SCALE_ATTN_LOGITS: True,
+ _KEY_OF_LAYERNORM_TYPE: 'rmsnorm',
+ _KEY_OF_DROPOUT_RATE: 0.8,
+ _KEY_OF_MLP_ACTIVATIONS: (('gelu', 'linear')),
+ _KEY_OF_FUSE_MLP_WI: True
 }, {
  _KEY_OF_TRANSPOSE_BS: False,
  _KEY_OF_SCALE_ATTN_LOGITS: True,

transformer_engine/common/fused_attn/fused_attn_fp16_bf16_max_seqlen_512.cu

@@ -327,7 +327,6 @@ static cudnn_frontend::Tensor createSoftmaxForward(
  // NOLINTNEXTLINE(runtime/references)
  std::vector<cudnn_frontend::Operation> &ops,
  cudnn_frontend::Tensor const &prevBlockOutputTensor) {
-
  int64_t afterBMM1_dim[4] = {b, h, s_q, s_kv};
  int64_t afterBMM1_stride[4] = {h * s_q * s_kv, s_q * s_kv, s_kv, 1};
 
@@ -645,7 +644,7 @@ void fused_attn_max_512_fwd_impl(int64_t b, int64_t h, int64_t s_q, int64_t s_kv
  mask_type, tensorType};
 
  using CacheType = std::map<FADescriptor, cudnn_frontend::ExecutionPlan>;
- static CacheType fmha_fprop_cache;
+ static thread_local CacheType fmha_fprop_cache;
 
  bool enable_dropout = (dropout_probability != 0.0f);
 
@@ -668,7 +667,8 @@ void fused_attn_max_512_fwd_impl(int64_t b, int64_t h, int64_t s_q, int64_t s_kv
  createScale(b, h, s_q, s_kv, d, layout, tensorType, ops);
 
  // if bias, we need to memset the S buffer to correctly computate dbias
- auto zero_s = (bias_type != NVTE_Bias_Type::NVTE_NO_BIAS);
+ auto zero_s = (bias_type != NVTE_Bias_Type::NVTE_NO_BIAS) ||
+ (mask_type == NVTE_Mask_Type::NVTE_CAUSAL_MASK);
  auto bmm1_output = createBMM1(b, h, s_q, s_kv, d, layout, tensorType, zero_s, ops);
 
  NVTE_CHECK(bias_type != NVTE_Bias_Type::NVTE_PRE_SCALE_BIAS,
@@ -814,7 +814,7 @@ void fused_attn_max_512_bwd_impl(int64_t b, int64_t h, int64_t s_q, int64_t s_kv
  layout, bias_type, mask_type, tensorType};
 
  using CacheType = std::map<FADescriptor, cudnn_frontend::ExecutionPlan>;
- static CacheType fmha_bprop_cache;
+ static thread_local CacheType fmha_bprop_cache;
 
  auto get_plan = [&](CacheType &cache, const FADescriptor &descriptor) {
  auto it = cache.find(descriptor);

transformer_engine/common/fused_attn/fused_attn_fp8.cu

@@ -1016,7 +1016,7 @@ void fa_fwd_fp8(int64_t b, int64_t s_q, int64_t s_kv, int64_t h, int64_t d,
  NVTE_Bias_Type::NVTE_NO_BIAS, NVTE_Mask_Type::NVTE_PADDING_MASK, tensorType};
 
  using CacheType = std::map<FADescriptor, cudnn_frontend::ExecutionPlan>;
- static CacheType fa_fprop_cache;
+ static thread_local CacheType fa_fprop_cache;
 
  // Get plan from cache if cache is available, otherwise create one
  auto get_plan = [&](CacheType &cache, const FADescriptor &descriptor) {
@@ -1332,7 +1332,7 @@ void fa_bwd_fp8(int64_t b, int64_t s_q, int64_t s_kv, int64_t h, int64_t d,
  NVTE_Bias_Type::NVTE_NO_BIAS, NVTE_Mask_Type::NVTE_PADDING_MASK, tensorType};
 
  using CacheType = std::map<FADescriptor, cudnn_frontend::ExecutionPlan>;
- static CacheType fa_bprop_cache;
+ static thread_local CacheType fa_bprop_cache;
 
  // Get plan from cache if cache is available, otherwise create one
  auto get_plan = [&](CacheType &cache, const FADescriptor &descriptor) {

transformer_engine/jax/flax/transformer.py

@@ -7,6 +7,7 @@
 import functools
 from enum import Enum
 from math import sqrt
+import os
 from typing import Any, Callable, Optional, Sequence, Tuple, Union
 import warnings
 
@@ -165,8 +166,17 @@ def core_attention(query: Array,
  else:
  attn_weights = jnp.einsum('bqhd,bkhd->bhqk', query, key)
 
+ # When a bias is present, the computation is performed as Softmax(attn_weights * scale + bias).
+ # In this case, the scale can not fused into the Softmax module.
+ if bias is not None:
+ attn_weights = attn_weights * scale_factor
+ fused_scale_factor = 1.
+ else:
+ # If no bias, the scale can be fused into Softmax module
+ fused_scale_factor = scale_factor
+
  attn_weights = Softmax(softmax_type=softmax_type,
- scale_factor=scale_factor,
+ scale_factor=fused_scale_factor,
  sharding_type=softmax_sharding_type)(attn_weights, mask, bias)
 
  if not deterministic and dropout_rate > 0.:
@@ -360,12 +370,13 @@ def kv_init(key, shape, dtype):
  q_seqlen = inputs_q.shape[0] if self.transpose_batch_sequence else inputs_q.shape[1]
  kv_seqlen = inputs_kv.shape[0] if self.transpose_batch_sequence else inputs_kv.shape[1]
  fused_attn_supported_seqlen = [128, 256, 384, 512]
+ enable_fused_attn = int(os.getenv("NVTE_FUSED_ATTN", "0"))
  use_fused_attn = not decode and not self.transpose_batch_sequence and self.fuse_qkv and \
  self.dropout_rate == 0 and canonicalize_dtype in [jnp.bfloat16, jnp.float16] and \
  q_seqlen in fused_attn_supported_seqlen and kv_seqlen in fused_attn_supported_seqlen \
- and is_fused_attn_kernel_available()
+ and is_fused_attn_kernel_available() and enable_fused_attn
 
- if not use_fused_attn:
+ if enable_fused_attn and not use_fused_attn:
  reason = ""
  if decode:
  reason += f"decode=False is required but got {decode}, "

transformer_engine/jax/sharding.py

@@ -386,7 +386,7 @@ def _get_dptp_sharding_meta(input_shapes: Tuple[Tuple[int, ...]],
 
  for input_shape, dp_dim, tp_dim in zip(input_shapes, input_dp_dims, input_tp_dims):
  in_axis = {}
- if dp_dim is not None:
+ if dp_dim is not None and input_shape is not None:
  in_axis[dp_dim] = dp_axis_name
  assert input_shape[dp_dim] % dp_size == 0, \
  f"The dimension of batch in input_shape should be a multiple of " \
@@ -398,7 +398,7 @@ def _get_dptp_sharding_meta(input_shapes: Tuple[Tuple[int, ...]],
  if tp_dim is not None and tp_dim >= dp_dim:
  tp_dim = tp_dim + 1
 
- if tp_dim is not None:
+ if tp_dim is not None and input_shape is not None:
  in_axis[tp_dim] = tp_axis_name
  assert input_shape[tp_dim] % tp_size == 0, \
  f"The dimension of tensor parallel in input_shape should be a multiple of " \

tests/jax/test_fused_attn.py

@@ -113,7 +113,7 @@ def customcall_cross_fused_attn(q, kv, q_token, kv_token, dropout_rng, **kwargs)
  reason="Fused attention kernel is not supported.")
 class TestSelfFusedAttnMax512():
 
- def set_input(self, b, s, h, d, dtype, attn_mask_type, pad_ratio):
+ def set_input(self, b, s, h, d, dtype, attn_mask_type, pad_ratio, with_bias):
  key = jax.random.PRNGKey(0)
  subkeys = jax.random.split(key, 2)
 
@@ -125,16 +125,17 @@ def set_input(self, b, s, h, d, dtype, attn_mask_type, pad_ratio):
 
  min_val, max_val = -1, 1
  self.qkv = jax.random.uniform(subkeys[0], qkv_shape, dtype, min_val, max_val)
- self.bias = jax.random.uniform(subkeys[1], bias_shape, dtype, min_val, max_val)
+ self.bias = jax.random.uniform(subkeys[1], bias_shape, dtype, min_val,
+ max_val) if with_bias else None
 
  self.q_token = jnp.concatenate((jnp.ones((b, self.valid_len)), jnp.zeros((b, pad_len))),
  axis=-1)
  self.kv_token = self.q_token
 
  self.scaling_factor = 1. / math.sqrt(d)
  self.dropout_probability = 0.
- self.dropout_rng = jax.random.PRNGKey(0)
- self.attn_bias_type = AttnBiasType.POST_SCALE_BIAS
+ self.dropout_rng = jax.random.PRNGKey(0) if self.dropout_probability > 0 else None
+ self.attn_bias_type = AttnBiasType.NO_BIAS if self.bias is None else AttnBiasType.POST_SCALE_BIAS
  # deterministic = not is_training
  self.deterministic = False
 
@@ -143,9 +144,17 @@ def set_input(self, b, s, h, d, dtype, attn_mask_type, pad_ratio):
  @pytest.mark.parametrize('attn_mask_type',
  [AttnMaskType.PADDING_MASK, AttnMaskType.CAUSAL_MASK])
  @pytest.mark.parametrize('pad_ratio', PAD_RATIO)
- def test_forward(self, b, s, h, d, dtype, attn_mask_type, pad_ratio):
+ @pytest.mark.parametrize('with_bias', [True, False])
+ def test_forward(self, b, s, h, d, dtype, attn_mask_type, pad_ratio, with_bias):
 
- self.set_input(b, s, h, d, dtype=dtype, attn_mask_type=attn_mask_type, pad_ratio=pad_ratio)
+ self.set_input(b,
+ s,
+ h,
+ d,
+ dtype=dtype,
+ attn_mask_type=attn_mask_type,
+ pad_ratio=pad_ratio,
+ with_bias=with_bias)
 
  primitive_out = customcall_self_fused_attn(self.qkv,
  self.bias,
@@ -183,8 +192,16 @@ def test_forward(self, b, s, h, d, dtype, attn_mask_type, pad_ratio):
  [AttnMaskType.PADDING_MASK, AttnMaskType.CAUSAL_MASK])
  @pytest.mark.parametrize('dtype', DTYPES)
  @pytest.mark.parametrize('pad_ratio', PAD_RATIO)
- def test_forward_backward(self, b, s, h, d, dtype, attn_mask_type, pad_ratio):
- self.set_input(b, s, h, d, dtype=dtype, attn_mask_type=attn_mask_type, pad_ratio=pad_ratio)
+ @pytest.mark.parametrize('with_bias', [True, False])
+ def test_forward_backward(self, b, s, h, d, dtype, attn_mask_type, pad_ratio, with_bias):
+ self.set_input(b,
+ s,
+ h,
+ d,
+ dtype=dtype,
+ attn_mask_type=attn_mask_type,
+ pad_ratio=pad_ratio,
+ with_bias=with_bias)
 
  def grad_func(fused_attn_max_512_func, *args, **kwargs):
  # Gradient is small, use a gradient multiplier to amplify the graident
@@ -221,11 +238,11 @@ def grad_func(fused_attn_max_512_func, *args, **kwargs):
  (0, 1)))
 
  primitive_out, (primitive_dqkv,
-primitive_dbeta) = jitted_primitive(self.qkv, self.bias, self.q_token,
+primitive_dbias) = jitted_primitive(self.qkv, self.bias, self.q_token,
  self.kv_token, self.dropout_rng)
 
  reference_out, (reference_dqkv,
-reference_dbeta) = jitted_reference(self.qkv, self.bias, self.q_token,
+reference_dbias) = jitted_reference(self.qkv, self.bias, self.q_token,
  self.kv_token, self.dropout_rng)
 
  np.testing.assert_allclose(jnp.asarray(primitive_out, np.float32),
@@ -261,20 +278,22 @@ def grad_func(fused_attn_max_512_func, *args, **kwargs):
  # Padded part should be 0s
  assert jnp.allclose(invalid_primitive_dqkv, jnp.zeros_like(invalid_primitive_dqkv))
 
- # dbeta valid part
- np.testing.assert_allclose(
- jnp.asarray(primitive_dbeta[:, :, :self.valid_len, :self.valid_len], np.float32),
- jnp.asarray(reference_dbeta[:, :, :self.valid_len, :self.valid_len], np.float32),
- rtol=1e-4,
- atol=3e-5)
-
- # dbeta padded part
- np.testing.assert_allclose(
- jnp.asarray(primitive_dbeta[:, :, self.valid_len:, self.valid_len:], np.float32),
- jnp.asarray(reference_dbeta[:, :, self.valid_len:, self.valid_len:], np.float32))
-
- assert jnp.allclose(primitive_dbeta[:, :, self.valid_len:, self.valid_len:],
- jnp.zeros_like(primitive_dbeta[:, :, self.valid_len:, self.valid_len:]))
+ if self.attn_bias_type != AttnBiasType.NO_BIAS:
+ # dbias valid part
+ np.testing.assert_allclose(
+ jnp.asarray(primitive_dbias[:, :, :self.valid_len, :self.valid_len], np.float32),
+ jnp.asarray(reference_dbias[:, :, :self.valid_len, :self.valid_len], np.float32),
+ rtol=1e-4,
+ atol=3e-5)
+
+ # dbias padded part
+ np.testing.assert_allclose(
+ jnp.asarray(primitive_dbias[:, :, self.valid_len:, self.valid_len:], np.float32),
+ jnp.asarray(reference_dbias[:, :, self.valid_len:, self.valid_len:], np.float32))
+
+ assert jnp.allclose(
+ primitive_dbias[:, :, self.valid_len:, self.valid_len:],
+ jnp.zeros_like(primitive_dbias[:, :, self.valid_len:, self.valid_len:]))
 
 
 @pytest.mark.skipif(not is_fused_attn_kernel_available(),

tests/jax/test_layer.py

@@ -102,6 +102,12 @@ def compare_frozen_dict(ref_fd, test_fd, rtol=1e-05, atol=1e-08):
  _KEY_OF_DROPOUT_RATE: 0.0,
  _KEY_OF_MLP_ACTIVATIONS: (('gelu', 'linear')),
  _KEY_OF_FUSE_MLP_WI: True
+}, {
+ _KEY_OF_SCALE_ATTN_LOGITS: True,
+ _KEY_OF_LAYERNORM_TYPE: 'rmsnorm',
+ _KEY_OF_DROPOUT_RATE: 0.8,
+ _KEY_OF_MLP_ACTIVATIONS: (('gelu', 'linear')),
+ _KEY_OF_FUSE_MLP_WI: True
 }, {
  _KEY_OF_TRANSPOSE_BS: False,
  _KEY_OF_SCALE_ATTN_LOGITS: True,

transformer_engine/common/fused_attn/fused_attn_fp16_bf16_max_seqlen_512.cu

@@ -327,7 +327,6 @@ static cudnn_frontend::Tensor createSoftmaxForward(
  // NOLINTNEXTLINE(runtime/references)
  std::vector<cudnn_frontend::Operation> &ops,
  cudnn_frontend::Tensor const &prevBlockOutputTensor) {
-
  int64_t afterBMM1_dim[4] = {b, h, s_q, s_kv};
  int64_t afterBMM1_stride[4] = {h * s_q * s_kv, s_q * s_kv, s_kv, 1};
 
@@ -645,7 +644,7 @@ void fused_attn_max_512_fwd_impl(int64_t b, int64_t h, int64_t s_q, int64_t s_kv
  mask_type, tensorType};
 
  using CacheType = std::map<FADescriptor, cudnn_frontend::ExecutionPlan>;
- static CacheType fmha_fprop_cache;
+ static thread_local CacheType fmha_fprop_cache;
 
  bool enable_dropout = (dropout_probability != 0.0f);
 
@@ -668,7 +667,8 @@ void fused_attn_max_512_fwd_impl(int64_t b, int64_t h, int64_t s_q, int64_t s_kv
  createScale(b, h, s_q, s_kv, d, layout, tensorType, ops);
 
  // if bias, we need to memset the S buffer to correctly computate dbias
- auto zero_s = (bias_type != NVTE_Bias_Type::NVTE_NO_BIAS);
+ auto zero_s = (bias_type != NVTE_Bias_Type::NVTE_NO_BIAS) ||
+ (mask_type == NVTE_Mask_Type::NVTE_CAUSAL_MASK);
  auto bmm1_output = createBMM1(b, h, s_q, s_kv, d, layout, tensorType, zero_s, ops);
 
  NVTE_CHECK(bias_type != NVTE_Bias_Type::NVTE_PRE_SCALE_BIAS,
@@ -814,7 +814,7 @@ void fused_attn_max_512_bwd_impl(int64_t b, int64_t h, int64_t s_q, int64_t s_kv
  layout, bias_type, mask_type, tensorType};
 
  using CacheType = std::map<FADescriptor, cudnn_frontend::ExecutionPlan>;
- static CacheType fmha_bprop_cache;
+ static thread_local CacheType fmha_bprop_cache;
 
  auto get_plan = [&](CacheType &cache, const FADescriptor &descriptor) {
  auto it = cache.find(descriptor);

transformer_engine/common/fused_attn/fused_attn_fp8.cu

@@ -1016,7 +1016,7 @@ void fa_fwd_fp8(int64_t b, int64_t s_q, int64_t s_kv, int64_t h, int64_t d,
  NVTE_Bias_Type::NVTE_NO_BIAS, NVTE_Mask_Type::NVTE_PADDING_MASK, tensorType};
 
  using CacheType = std::map<FADescriptor, cudnn_frontend::ExecutionPlan>;
- static CacheType fa_fprop_cache;
+ static thread_local CacheType fa_fprop_cache;
 
  // Get plan from cache if cache is available, otherwise create one
  auto get_plan = [&](CacheType &cache, const FADescriptor &descriptor) {
@@ -1332,7 +1332,7 @@ void fa_bwd_fp8(int64_t b, int64_t s_q, int64_t s_kv, int64_t h, int64_t d,
  NVTE_Bias_Type::NVTE_NO_BIAS, NVTE_Mask_Type::NVTE_PADDING_MASK, tensorType};
 
  using CacheType = std::map<FADescriptor, cudnn_frontend::ExecutionPlan>;
- static CacheType fa_bprop_cache;
+ static thread_local CacheType fa_bprop_cache;
 
  // Get plan from cache if cache is available, otherwise create one
  auto get_plan = [&](CacheType &cache, const FADescriptor &descriptor) {

transformer_engine/jax/flax/transformer.py

@@ -7,6 +7,7 @@
 import functools
 from enum import Enum
 from math import sqrt
+import os
 from typing import Any, Callable, Optional, Sequence, Tuple, Union
 import warnings
 
@@ -165,8 +166,17 @@ def core_attention(query: Array,
  else:
  attn_weights = jnp.einsum('bqhd,bkhd->bhqk', query, key)
 
+ # When a bias is present, the computation is performed as Softmax(attn_weights * scale + bias).
+ # In this case, the scale can not fused into the Softmax module.
+ if bias is not None:
+ attn_weights = attn_weights * scale_factor
+ fused_scale_factor = 1.
+ else:
+ # If no bias, the scale can be fused into Softmax module
+ fused_scale_factor = scale_factor
+
  attn_weights = Softmax(softmax_type=softmax_type,
- scale_factor=scale_factor,
+ scale_factor=fused_scale_factor,
  sharding_type=softmax_sharding_type)(attn_weights, mask, bias)
 
  if not deterministic and dropout_rate > 0.:
@@ -360,12 +370,13 @@ def kv_init(key, shape, dtype):
  q_seqlen = inputs_q.shape[0] if self.transpose_batch_sequence else inputs_q.shape[1]
  kv_seqlen = inputs_kv.shape[0] if self.transpose_batch_sequence else inputs_kv.shape[1]
  fused_attn_supported_seqlen = [128, 256, 384, 512]
+ enable_fused_attn = int(os.getenv("NVTE_FUSED_ATTN", "0"))
  use_fused_attn = not decode and not self.transpose_batch_sequence and self.fuse_qkv and \
  self.dropout_rate == 0 and canonicalize_dtype in [jnp.bfloat16, jnp.float16] and \
  q_seqlen in fused_attn_supported_seqlen and kv_seqlen in fused_attn_supported_seqlen \
- and is_fused_attn_kernel_available()
+ and is_fused_attn_kernel_available() and enable_fused_attn
 
- if not use_fused_attn:
+ if enable_fused_attn and not use_fused_attn:
  reason = ""
  if decode:
  reason += f"decode=False is required but got {decode}, "

transformer_engine/jax/sharding.py

@@ -386,7 +386,7 @@ def _get_dptp_sharding_meta(input_shapes: Tuple[Tuple[int, ...]],
 
  for input_shape, dp_dim, tp_dim in zip(input_shapes, input_dp_dims, input_tp_dims):
  in_axis = {}
- if dp_dim is not None:
+ if dp_dim is not None and input_shape is not None:
  in_axis[dp_dim] = dp_axis_name
  assert input_shape[dp_dim] % dp_size == 0, \
  f"The dimension of batch in input_shape should be a multiple of " \
@@ -398,7 +398,7 @@ def _get_dptp_sharding_meta(input_shapes: Tuple[Tuple[int, ...]],
  if tp_dim is not None and tp_dim >= dp_dim:
  tp_dim = tp_dim + 1
 
- if tp_dim is not None:
+ if tp_dim is not None and input_shape is not None:
  in_axis[tp_dim] = tp_axis_name
  assert input_shape[tp_dim] % tp_size == 0, \
  f"The dimension of tensor parallel in input_shape should be a multiple of " \

tests/jax/test_fused_attn.py

@@ -113,7 +113,7 @@ def customcall_cross_fused_attn(q, kv, q_token, kv_token, dropout_rng, **kwargs)
  reason="Fused attention kernel is not supported.")
 class TestSelfFusedAttnMax512():
 
- def set_input(self, b, s, h, d, dtype, attn_mask_type, pad_ratio):
+ def set_input(self, b, s, h, d, dtype, attn_mask_type, pad_ratio, with_bias):
  key = jax.random.PRNGKey(0)
  subkeys = jax.random.split(key, 2)
 
@@ -125,16 +125,17 @@ def set_input(self, b, s, h, d, dtype, attn_mask_type, pad_ratio):
 
  min_val, max_val = -1, 1
  self.qkv = jax.random.uniform(subkeys[0], qkv_shape, dtype, min_val, max_val)
- self.bias = jax.random.uniform(subkeys[1], bias_shape, dtype, min_val, max_val)
+ self.bias = jax.random.uniform(subkeys[1], bias_shape, dtype, min_val,
+ max_val) if with_bias else None
 
  self.q_token = jnp.concatenate((jnp.ones((b, self.valid_len)), jnp.zeros((b, pad_len))),
  axis=-1)
  self.kv_token = self.q_token
 
  self.scaling_factor = 1. / math.sqrt(d)
  self.dropout_probability = 0.
- self.dropout_rng = jax.random.PRNGKey(0)
- self.attn_bias_type = AttnBiasType.POST_SCALE_BIAS
+ self.dropout_rng = jax.random.PRNGKey(0) if self.dropout_probability > 0 else None
+ self.attn_bias_type = AttnBiasType.NO_BIAS if self.bias is None else AttnBiasType.POST_SCALE_BIAS
  # deterministic = not is_training
  self.deterministic = False
 
@@ -143,9 +144,17 @@ def set_input(self, b, s, h, d, dtype, attn_mask_type, pad_ratio):
  @pytest.mark.parametrize('attn_mask_type',
  [AttnMaskType.PADDING_MASK, AttnMaskType.CAUSAL_MASK])
  @pytest.mark.parametrize('pad_ratio', PAD_RATIO)
- def test_forward(self, b, s, h, d, dtype, attn_mask_type, pad_ratio):
+ @pytest.mark.parametrize('with_bias', [True, False])
+ def test_forward(self, b, s, h, d, dtype, attn_mask_type, pad_ratio, with_bias):
 
- self.set_input(b, s, h, d, dtype=dtype, attn_mask_type=attn_mask_type, pad_ratio=pad_ratio)
+ self.set_input(b,
+ s,
+ h,
+ d,
+ dtype=dtype,
+ attn_mask_type=attn_mask_type,
+ pad_ratio=pad_ratio,
+ with_bias=with_bias)
 
  primitive_out = customcall_self_fused_attn(self.qkv,
  self.bias,
@@ -183,8 +192,16 @@ def test_forward(self, b, s, h, d, dtype, attn_mask_type, pad_ratio):
  [AttnMaskType.PADDING_MASK, AttnMaskType.CAUSAL_MASK])
  @pytest.mark.parametrize('dtype', DTYPES)
  @pytest.mark.parametrize('pad_ratio', PAD_RATIO)
- def test_forward_backward(self, b, s, h, d, dtype, attn_mask_type, pad_ratio):
- self.set_input(b, s, h, d, dtype=dtype, attn_mask_type=attn_mask_type, pad_ratio=pad_ratio)
+ @pytest.mark.parametrize('with_bias', [True, False])
+ def test_forward_backward(self, b, s, h, d, dtype, attn_mask_type, pad_ratio, with_bias):
+ self.set_input(b,
+ s,
+ h,
+ d,
+ dtype=dtype,
+ attn_mask_type=attn_mask_type,
+ pad_ratio=pad_ratio,
+ with_bias=with_bias)
 
  def grad_func(fused_attn_max_512_func, *args, **kwargs):
  # Gradient is small, use a gradient multiplier to amplify the graident
@@ -221,11 +238,11 @@ def grad_func(fused_attn_max_512_func, *args, **kwargs):
  (0, 1)))
 
  primitive_out, (primitive_dqkv,
-primitive_dbeta) = jitted_primitive(self.qkv, self.bias, self.q_token,
+primitive_dbias) = jitted_primitive(self.qkv, self.bias, self.q_token,
  self.kv_token, self.dropout_rng)
 
  reference_out, (reference_dqkv,
-reference_dbeta) = jitted_reference(self.qkv, self.bias, self.q_token,
+reference_dbias) = jitted_reference(self.qkv, self.bias, self.q_token,
  self.kv_token, self.dropout_rng)
 
  np.testing.assert_allclose(jnp.asarray(primitive_out, np.float32),
@@ -261,20 +278,22 @@ def grad_func(fused_attn_max_512_func, *args, **kwargs):
  # Padded part should be 0s
  assert jnp.allclose(invalid_primitive_dqkv, jnp.zeros_like(invalid_primitive_dqkv))
 
- # dbeta valid part
- np.testing.assert_allclose(
- jnp.asarray(primitive_dbeta[:, :, :self.valid_len, :self.valid_len], np.float32),
- jnp.asarray(reference_dbeta[:, :, :self.valid_len, :self.valid_len], np.float32),
- rtol=1e-4,
- atol=3e-5)
-
- # dbeta padded part
- np.testing.assert_allclose(
- jnp.asarray(primitive_dbeta[:, :, self.valid_len:, self.valid_len:], np.float32),
- jnp.asarray(reference_dbeta[:, :, self.valid_len:, self.valid_len:], np.float32))
-
- assert jnp.allclose(primitive_dbeta[:, :, self.valid_len:, self.valid_len:],
- jnp.zeros_like(primitive_dbeta[:, :, self.valid_len:, self.valid_len:]))
+ if self.attn_bias_type != AttnBiasType.NO_BIAS:
+ # dbias valid part
+ np.testing.assert_allclose(
+ jnp.asarray(primitive_dbias[:, :, :self.valid_len, :self.valid_len], np.float32),
+ jnp.asarray(reference_dbias[:, :, :self.valid_len, :self.valid_len], np.float32),
+ rtol=1e-4,
+ atol=3e-5)
+
+ # dbias padded part
+ np.testing.assert_allclose(
+ jnp.asarray(primitive_dbias[:, :, self.valid_len:, self.valid_len:], np.float32),
+ jnp.asarray(reference_dbias[:, :, self.valid_len:, self.valid_len:], np.float32))
+
+ assert jnp.allclose(
+ primitive_dbias[:, :, self.valid_len:, self.valid_len:],
+ jnp.zeros_like(primitive_dbias[:, :, self.valid_len:, self.valid_len:]))
 
 
 @pytest.mark.skipif(not is_fused_attn_kernel_available(),

tests/jax/test_layer.py

@@ -102,6 +102,12 @@ def compare_frozen_dict(ref_fd, test_fd, rtol=1e-05, atol=1e-08):
  _KEY_OF_DROPOUT_RATE: 0.0,
  _KEY_OF_MLP_ACTIVATIONS: (('gelu', 'linear')),
  _KEY_OF_FUSE_MLP_WI: True
+}, {
+ _KEY_OF_SCALE_ATTN_LOGITS: True,
+ _KEY_OF_LAYERNORM_TYPE: 'rmsnorm',
+ _KEY_OF_DROPOUT_RATE: 0.8,
+ _KEY_OF_MLP_ACTIVATIONS: (('gelu', 'linear')),
+ _KEY_OF_FUSE_MLP_WI: True
 }, {
  _KEY_OF_TRANSPOSE_BS: False,
  _KEY_OF_SCALE_ATTN_LOGITS: True,

transformer_engine/common/fused_attn/fused_attn_fp16_bf16_max_seqlen_512.cu

@@ -327,7 +327,6 @@ static cudnn_frontend::Tensor createSoftmaxForward(
  // NOLINTNEXTLINE(runtime/references)
  std::vector<cudnn_frontend::Operation> &ops,
  cudnn_frontend::Tensor const &prevBlockOutputTensor) {
-
  int64_t afterBMM1_dim[4] = {b, h, s_q, s_kv};
  int64_t afterBMM1_stride[4] = {h * s_q * s_kv, s_q * s_kv, s_kv, 1};
 
@@ -645,7 +644,7 @@ void fused_attn_max_512_fwd_impl(int64_t b, int64_t h, int64_t s_q, int64_t s_kv
  mask_type, tensorType};
 
  using CacheType = std::map<FADescriptor, cudnn_frontend::ExecutionPlan>;
- static CacheType fmha_fprop_cache;
+ static thread_local CacheType fmha_fprop_cache;
 
  bool enable_dropout = (dropout_probability != 0.0f);
 
@@ -668,7 +667,8 @@ void fused_attn_max_512_fwd_impl(int64_t b, int64_t h, int64_t s_q, int64_t s_kv
  createScale(b, h, s_q, s_kv, d, layout, tensorType, ops);
 
  // if bias, we need to memset the S buffer to correctly computate dbias
- auto zero_s = (bias_type != NVTE_Bias_Type::NVTE_NO_BIAS);
+ auto zero_s = (bias_type != NVTE_Bias_Type::NVTE_NO_BIAS) ||
+ (mask_type == NVTE_Mask_Type::NVTE_CAUSAL_MASK);
  auto bmm1_output = createBMM1(b, h, s_q, s_kv, d, layout, tensorType, zero_s, ops);
 
  NVTE_CHECK(bias_type != NVTE_Bias_Type::NVTE_PRE_SCALE_BIAS,
@@ -814,7 +814,7 @@ void fused_attn_max_512_bwd_impl(int64_t b, int64_t h, int64_t s_q, int64_t s_kv
  layout, bias_type, mask_type, tensorType};
 
  using CacheType = std::map<FADescriptor, cudnn_frontend::ExecutionPlan>;
- static CacheType fmha_bprop_cache;
+ static thread_local CacheType fmha_bprop_cache;
 
  auto get_plan = [&](CacheType &cache, const FADescriptor &descriptor) {
  auto it = cache.find(descriptor);

transformer_engine/common/fused_attn/fused_attn_fp8.cu

@@ -1016,7 +1016,7 @@ void fa_fwd_fp8(int64_t b, int64_t s_q, int64_t s_kv, int64_t h, int64_t d,
  NVTE_Bias_Type::NVTE_NO_BIAS, NVTE_Mask_Type::NVTE_PADDING_MASK, tensorType};
 
  using CacheType = std::map<FADescriptor, cudnn_frontend::ExecutionPlan>;
- static CacheType fa_fprop_cache;
+ static thread_local CacheType fa_fprop_cache;
 
  // Get plan from cache if cache is available, otherwise create one
  auto get_plan = [&](CacheType &cache, const FADescriptor &descriptor) {
@@ -1332,7 +1332,7 @@ void fa_bwd_fp8(int64_t b, int64_t s_q, int64_t s_kv, int64_t h, int64_t d,
  NVTE_Bias_Type::NVTE_NO_BIAS, NVTE_Mask_Type::NVTE_PADDING_MASK, tensorType};
 
  using CacheType = std::map<FADescriptor, cudnn_frontend::ExecutionPlan>;
- static CacheType fa_bprop_cache;
+ static thread_local CacheType fa_bprop_cache;
 
  // Get plan from cache if cache is available, otherwise create one
  auto get_plan = [&](CacheType &cache, const FADescriptor &descriptor) {

transformer_engine/jax/flax/transformer.py

@@ -7,6 +7,7 @@
 import functools
 from enum import Enum
 from math import sqrt
+import os
 from typing import Any, Callable, Optional, Sequence, Tuple, Union
 import warnings
 
@@ -165,8 +166,17 @@ def core_attention(query: Array,
  else:
  attn_weights = jnp.einsum('bqhd,bkhd->bhqk', query, key)
 
+ # When a bias is present, the computation is performed as Softmax(attn_weights * scale + bias).
+ # In this case, the scale can not fused into the Softmax module.
+ if bias is not None:
+ attn_weights = attn_weights * scale_factor
+ fused_scale_factor = 1.
+ else:
+ # If no bias, the scale can be fused into Softmax module
+ fused_scale_factor = scale_factor
+
  attn_weights = Softmax(softmax_type=softmax_type,
- scale_factor=scale_factor,
+ scale_factor=fused_scale_factor,
  sharding_type=softmax_sharding_type)(attn_weights, mask, bias)
 
  if not deterministic and dropout_rate > 0.:
@@ -360,12 +370,13 @@ def kv_init(key, shape, dtype):
  q_seqlen = inputs_q.shape[0] if self.transpose_batch_sequence else inputs_q.shape[1]
  kv_seqlen = inputs_kv.shape[0] if self.transpose_batch_sequence else inputs_kv.shape[1]
  fused_attn_supported_seqlen = [128, 256, 384, 512]
+ enable_fused_attn = int(os.getenv("NVTE_FUSED_ATTN", "0"))
  use_fused_attn = not decode and not self.transpose_batch_sequence and self.fuse_qkv and \
  self.dropout_rate == 0 and canonicalize_dtype in [jnp.bfloat16, jnp.float16] and \
  q_seqlen in fused_attn_supported_seqlen and kv_seqlen in fused_attn_supported_seqlen \
- and is_fused_attn_kernel_available()
+ and is_fused_attn_kernel_available() and enable_fused_attn
 
- if not use_fused_attn:
+ if enable_fused_attn and not use_fused_attn:
  reason = ""
  if decode:
  reason += f"decode=False is required but got {decode}, "

transformer_engine/jax/sharding.py

@@ -386,7 +386,7 @@ def _get_dptp_sharding_meta(input_shapes: Tuple[Tuple[int, ...]],
 
  for input_shape, dp_dim, tp_dim in zip(input_shapes, input_dp_dims, input_tp_dims):
  in_axis = {}
- if dp_dim is not None:
+ if dp_dim is not None and input_shape is not None:
  in_axis[dp_dim] = dp_axis_name
  assert input_shape[dp_dim] % dp_size == 0, \
  f"The dimension of batch in input_shape should be a multiple of " \
@@ -398,7 +398,7 @@ def _get_dptp_sharding_meta(input_shapes: Tuple[Tuple[int, ...]],
  if tp_dim is not None and tp_dim >= dp_dim:
  tp_dim = tp_dim + 1
 
- if tp_dim is not None:
+ if tp_dim is not None and input_shape is not None:
  in_axis[tp_dim] = tp_axis_name
  assert input_shape[tp_dim] % tp_size == 0, \
  f"The dimension of tensor parallel in input_shape should be a multiple of " \

tests/jax/test_fused_attn.py

@@ -113,7 +113,7 @@ def customcall_cross_fused_attn(q, kv, q_token, kv_token, dropout_rng, **kwargs)
  reason="Fused attention kernel is not supported.")
 class TestSelfFusedAttnMax512():
 
- def set_input(self, b, s, h, d, dtype, attn_mask_type, pad_ratio):
+ def set_input(self, b, s, h, d, dtype, attn_mask_type, pad_ratio, with_bias):
  key = jax.random.PRNGKey(0)
  subkeys = jax.random.split(key, 2)
 
@@ -125,16 +125,17 @@ def set_input(self, b, s, h, d, dtype, attn_mask_type, pad_ratio):
 
  min_val, max_val = -1, 1
  self.qkv = jax.random.uniform(subkeys[0], qkv_shape, dtype, min_val, max_val)
- self.bias = jax.random.uniform(subkeys[1], bias_shape, dtype, min_val, max_val)
+ self.bias = jax.random.uniform(subkeys[1], bias_shape, dtype, min_val,
+ max_val) if with_bias else None
 
  self.q_token = jnp.concatenate((jnp.ones((b, self.valid_len)), jnp.zeros((b, pad_len))),
  axis=-1)
  self.kv_token = self.q_token
 
  self.scaling_factor = 1. / math.sqrt(d)
  self.dropout_probability = 0.
- self.dropout_rng = jax.random.PRNGKey(0)
- self.attn_bias_type = AttnBiasType.POST_SCALE_BIAS
+ self.dropout_rng = jax.random.PRNGKey(0) if self.dropout_probability > 0 else None
+ self.attn_bias_type = AttnBiasType.NO_BIAS if self.bias is None else AttnBiasType.POST_SCALE_BIAS
  # deterministic = not is_training
  self.deterministic = False
 
@@ -143,9 +144,17 @@ def set_input(self, b, s, h, d, dtype, attn_mask_type, pad_ratio):
  @pytest.mark.parametrize('attn_mask_type',
  [AttnMaskType.PADDING_MASK, AttnMaskType.CAUSAL_MASK])
  @pytest.mark.parametrize('pad_ratio', PAD_RATIO)
- def test_forward(self, b, s, h, d, dtype, attn_mask_type, pad_ratio):
+ @pytest.mark.parametrize('with_bias', [True, False])
+ def test_forward(self, b, s, h, d, dtype, attn_mask_type, pad_ratio, with_bias):
 
- self.set_input(b, s, h, d, dtype=dtype, attn_mask_type=attn_mask_type, pad_ratio=pad_ratio)
+ self.set_input(b,
+ s,
+ h,
+ d,
+ dtype=dtype,
+ attn_mask_type=attn_mask_type,
+ pad_ratio=pad_ratio,
+ with_bias=with_bias)
 
  primitive_out = customcall_self_fused_attn(self.qkv,
  self.bias,
@@ -183,8 +192,16 @@ def test_forward(self, b, s, h, d, dtype, attn_mask_type, pad_ratio):
  [AttnMaskType.PADDING_MASK, AttnMaskType.CAUSAL_MASK])
  @pytest.mark.parametrize('dtype', DTYPES)
  @pytest.mark.parametrize('pad_ratio', PAD_RATIO)
- def test_forward_backward(self, b, s, h, d, dtype, attn_mask_type, pad_ratio):
- self.set_input(b, s, h, d, dtype=dtype, attn_mask_type=attn_mask_type, pad_ratio=pad_ratio)
+ @pytest.mark.parametrize('with_bias', [True, False])
+ def test_forward_backward(self, b, s, h, d, dtype, attn_mask_type, pad_ratio, with_bias):
+ self.set_input(b,
+ s,
+ h,
+ d,
+ dtype=dtype,
+ attn_mask_type=attn_mask_type,
+ pad_ratio=pad_ratio,
+ with_bias=with_bias)
 
  def grad_func(fused_attn_max_512_func, *args, **kwargs):
  # Gradient is small, use a gradient multiplier to amplify the graident
@@ -221,11 +238,11 @@ def grad_func(fused_attn_max_512_func, *args, **kwargs):
  (0, 1)))
 
  primitive_out, (primitive_dqkv,
-primitive_dbeta) = jitted_primitive(self.qkv, self.bias, self.q_token,
+primitive_dbias) = jitted_primitive(self.qkv, self.bias, self.q_token,
  self.kv_token, self.dropout_rng)
 
  reference_out, (reference_dqkv,
-reference_dbeta) = jitted_reference(self.qkv, self.bias, self.q_token,
+reference_dbias) = jitted_reference(self.qkv, self.bias, self.q_token,
  self.kv_token, self.dropout_rng)
 
  np.testing.assert_allclose(jnp.asarray(primitive_out, np.float32),
@@ -261,20 +278,22 @@ def grad_func(fused_attn_max_512_func, *args, **kwargs):
  # Padded part should be 0s
  assert jnp.allclose(invalid_primitive_dqkv, jnp.zeros_like(invalid_primitive_dqkv))
 
- # dbeta valid part
- np.testing.assert_allclose(
- jnp.asarray(primitive_dbeta[:, :, :self.valid_len, :self.valid_len], np.float32),
- jnp.asarray(reference_dbeta[:, :, :self.valid_len, :self.valid_len], np.float32),
- rtol=1e-4,
- atol=3e-5)
-
- # dbeta padded part
- np.testing.assert_allclose(
- jnp.asarray(primitive_dbeta[:, :, self.valid_len:, self.valid_len:], np.float32),
- jnp.asarray(reference_dbeta[:, :, self.valid_len:, self.valid_len:], np.float32))
-
- assert jnp.allclose(primitive_dbeta[:, :, self.valid_len:, self.valid_len:],
- jnp.zeros_like(primitive_dbeta[:, :, self.valid_len:, self.valid_len:]))
+ if self.attn_bias_type != AttnBiasType.NO_BIAS:
+ # dbias valid part
+ np.testing.assert_allclose(
+ jnp.asarray(primitive_dbias[:, :, :self.valid_len, :self.valid_len], np.float32),
+ jnp.asarray(reference_dbias[:, :, :self.valid_len, :self.valid_len], np.float32),
+ rtol=1e-4,
+ atol=3e-5)
+
+ # dbias padded part
+ np.testing.assert_allclose(
+ jnp.asarray(primitive_dbias[:, :, self.valid_len:, self.valid_len:], np.float32),
+ jnp.asarray(reference_dbias[:, :, self.valid_len:, self.valid_len:], np.float32))
+
+ assert jnp.allclose(
+ primitive_dbias[:, :, self.valid_len:, self.valid_len:],
+ jnp.zeros_like(primitive_dbias[:, :, self.valid_len:, self.valid_len:]))
 
 
 @pytest.mark.skipif(not is_fused_attn_kernel_available(),

tests/jax/test_layer.py

@@ -102,6 +102,12 @@ def compare_frozen_dict(ref_fd, test_fd, rtol=1e-05, atol=1e-08):
  _KEY_OF_DROPOUT_RATE: 0.0,
  _KEY_OF_MLP_ACTIVATIONS: (('gelu', 'linear')),
  _KEY_OF_FUSE_MLP_WI: True
+}, {
+ _KEY_OF_SCALE_ATTN_LOGITS: True,
+ _KEY_OF_LAYERNORM_TYPE: 'rmsnorm',
+ _KEY_OF_DROPOUT_RATE: 0.8,
+ _KEY_OF_MLP_ACTIVATIONS: (('gelu', 'linear')),
+ _KEY_OF_FUSE_MLP_WI: True
 }, {
  _KEY_OF_TRANSPOSE_BS: False,
  _KEY_OF_SCALE_ATTN_LOGITS: True,

transformer_engine/common/fused_attn/fused_attn_fp16_bf16_max_seqlen_512.cu

@@ -327,7 +327,6 @@ static cudnn_frontend::Tensor createSoftmaxForward(
  // NOLINTNEXTLINE(runtime/references)
  std::vector<cudnn_frontend::Operation> &ops,
  cudnn_frontend::Tensor const &prevBlockOutputTensor) {
-
  int64_t afterBMM1_dim[4] = {b, h, s_q, s_kv};
  int64_t afterBMM1_stride[4] = {h * s_q * s_kv, s_q * s_kv, s_kv, 1};
 
@@ -645,7 +644,7 @@ void fused_attn_max_512_fwd_impl(int64_t b, int64_t h, int64_t s_q, int64_t s_kv
  mask_type, tensorType};
 
  using CacheType = std::map<FADescriptor, cudnn_frontend::ExecutionPlan>;
- static CacheType fmha_fprop_cache;
+ static thread_local CacheType fmha_fprop_cache;
 
  bool enable_dropout = (dropout_probability != 0.0f);
 
@@ -668,7 +667,8 @@ void fused_attn_max_512_fwd_impl(int64_t b, int64_t h, int64_t s_q, int64_t s_kv
  createScale(b, h, s_q, s_kv, d, layout, tensorType, ops);
 
  // if bias, we need to memset the S buffer to correctly computate dbias
- auto zero_s = (bias_type != NVTE_Bias_Type::NVTE_NO_BIAS);
+ auto zero_s = (bias_type != NVTE_Bias_Type::NVTE_NO_BIAS) ||
+ (mask_type == NVTE_Mask_Type::NVTE_CAUSAL_MASK);
  auto bmm1_output = createBMM1(b, h, s_q, s_kv, d, layout, tensorType, zero_s, ops);
 
  NVTE_CHECK(bias_type != NVTE_Bias_Type::NVTE_PRE_SCALE_BIAS,
@@ -814,7 +814,7 @@ void fused_attn_max_512_bwd_impl(int64_t b, int64_t h, int64_t s_q, int64_t s_kv
  layout, bias_type, mask_type, tensorType};
 
  using CacheType = std::map<FADescriptor, cudnn_frontend::ExecutionPlan>;
- static CacheType fmha_bprop_cache;
+ static thread_local CacheType fmha_bprop_cache;
 
  auto get_plan = [&](CacheType &cache, const FADescriptor &descriptor) {
  auto it = cache.find(descriptor);

transformer_engine/common/fused_attn/fused_attn_fp8.cu

@@ -1016,7 +1016,7 @@ void fa_fwd_fp8(int64_t b, int64_t s_q, int64_t s_kv, int64_t h, int64_t d,
  NVTE_Bias_Type::NVTE_NO_BIAS, NVTE_Mask_Type::NVTE_PADDING_MASK, tensorType};
 
  using CacheType = std::map<FADescriptor, cudnn_frontend::ExecutionPlan>;
- static CacheType fa_fprop_cache;
+ static thread_local CacheType fa_fprop_cache;
 
  // Get plan from cache if cache is available, otherwise create one
  auto get_plan = [&](CacheType &cache, const FADescriptor &descriptor) {
@@ -1332,7 +1332,7 @@ void fa_bwd_fp8(int64_t b, int64_t s_q, int64_t s_kv, int64_t h, int64_t d,
  NVTE_Bias_Type::NVTE_NO_BIAS, NVTE_Mask_Type::NVTE_PADDING_MASK, tensorType};
 
  using CacheType = std::map<FADescriptor, cudnn_frontend::ExecutionPlan>;
- static CacheType fa_bprop_cache;
+ static thread_local CacheType fa_bprop_cache;
 
  // Get plan from cache if cache is available, otherwise create one
  auto get_plan = [&](CacheType &cache, const FADescriptor &descriptor) {

transformer_engine/jax/flax/transformer.py

@@ -7,6 +7,7 @@
 import functools
 from enum import Enum
 from math import sqrt
+import os
 from typing import Any, Callable, Optional, Sequence, Tuple, Union
 import warnings
 
@@ -165,8 +166,17 @@ def core_attention(query: Array,
  else:
  attn_weights = jnp.einsum('bqhd,bkhd->bhqk', query, key)
 
+ # When a bias is present, the computation is performed as Softmax(attn_weights * scale + bias).
+ # In this case, the scale can not fused into the Softmax module.
+ if bias is not None:
+ attn_weights = attn_weights * scale_factor
+ fused_scale_factor = 1.
+ else:
+ # If no bias, the scale can be fused into Softmax module
+ fused_scale_factor = scale_factor
+
  attn_weights = Softmax(softmax_type=softmax_type,
- scale_factor=scale_factor,
+ scale_factor=fused_scale_factor,
  sharding_type=softmax_sharding_type)(attn_weights, mask, bias)
 
  if not deterministic and dropout_rate > 0.:
@@ -360,12 +370,13 @@ def kv_init(key, shape, dtype):
  q_seqlen = inputs_q.shape[0] if self.transpose_batch_sequence else inputs_q.shape[1]
  kv_seqlen = inputs_kv.shape[0] if self.transpose_batch_sequence else inputs_kv.shape[1]
  fused_attn_supported_seqlen = [128, 256, 384, 512]
+ enable_fused_attn = int(os.getenv("NVTE_FUSED_ATTN", "0"))
  use_fused_attn = not decode and not self.transpose_batch_sequence and self.fuse_qkv and \
  self.dropout_rate == 0 and canonicalize_dtype in [jnp.bfloat16, jnp.float16] and \
  q_seqlen in fused_attn_supported_seqlen and kv_seqlen in fused_attn_supported_seqlen \
- and is_fused_attn_kernel_available()
+ and is_fused_attn_kernel_available() and enable_fused_attn
 
- if not use_fused_attn:
+ if enable_fused_attn and not use_fused_attn:
  reason = ""
  if decode:
  reason += f"decode=False is required but got {decode}, "

transformer_engine/jax/sharding.py

@@ -386,7 +386,7 @@ def _get_dptp_sharding_meta(input_shapes: Tuple[Tuple[int, ...]],
 
  for input_shape, dp_dim, tp_dim in zip(input_shapes, input_dp_dims, input_tp_dims):
  in_axis = {}
- if dp_dim is not None:
+ if dp_dim is not None and input_shape is not None:
  in_axis[dp_dim] = dp_axis_name
  assert input_shape[dp_dim] % dp_size == 0, \
  f"The dimension of batch in input_shape should be a multiple of " \
@@ -398,7 +398,7 @@ def _get_dptp_sharding_meta(input_shapes: Tuple[Tuple[int, ...]],
  if tp_dim is not None and tp_dim >= dp_dim:
  tp_dim = tp_dim + 1
 
- if tp_dim is not None:
+ if tp_dim is not None and input_shape is not None:
  in_axis[tp_dim] = tp_axis_name
  assert input_shape[tp_dim] % tp_size == 0, \
  f"The dimension of tensor parallel in input_shape should be a multiple of " \

tests/jax/test_fused_attn.py

@@ -113,7 +113,7 @@ def customcall_cross_fused_attn(q, kv, q_token, kv_token, dropout_rng, **kwargs)
  reason="Fused attention kernel is not supported.")
 class TestSelfFusedAttnMax512():
 
- def set_input(self, b, s, h, d, dtype, attn_mask_type, pad_ratio):
+ def set_input(self, b, s, h, d, dtype, attn_mask_type, pad_ratio, with_bias):
  key = jax.random.PRNGKey(0)
  subkeys = jax.random.split(key, 2)
 
@@ -125,16 +125,17 @@ def set_input(self, b, s, h, d, dtype, attn_mask_type, pad_ratio):
 
  min_val, max_val = -1, 1
  self.qkv = jax.random.uniform(subkeys[0], qkv_shape, dtype, min_val, max_val)
- self.bias = jax.random.uniform(subkeys[1], bias_shape, dtype, min_val, max_val)
+ self.bias = jax.random.uniform(subkeys[1], bias_shape, dtype, min_val,
+ max_val) if with_bias else None
 
  self.q_token = jnp.concatenate((jnp.ones((b, self.valid_len)), jnp.zeros((b, pad_len))),
  axis=-1)
  self.kv_token = self.q_token
 
  self.scaling_factor = 1. / math.sqrt(d)
  self.dropout_probability = 0.
- self.dropout_rng = jax.random.PRNGKey(0)
- self.attn_bias_type = AttnBiasType.POST_SCALE_BIAS
+ self.dropout_rng = jax.random.PRNGKey(0) if self.dropout_probability > 0 else None
+ self.attn_bias_type = AttnBiasType.NO_BIAS if self.bias is None else AttnBiasType.POST_SCALE_BIAS
  # deterministic = not is_training
  self.deterministic = False
 
@@ -143,9 +144,17 @@ def set_input(self, b, s, h, d, dtype, attn_mask_type, pad_ratio):
  @pytest.mark.parametrize('attn_mask_type',
  [AttnMaskType.PADDING_MASK, AttnMaskType.CAUSAL_MASK])
  @pytest.mark.parametrize('pad_ratio', PAD_RATIO)
- def test_forward(self, b, s, h, d, dtype, attn_mask_type, pad_ratio):
+ @pytest.mark.parametrize('with_bias', [True, False])
+ def test_forward(self, b, s, h, d, dtype, attn_mask_type, pad_ratio, with_bias):
 
- self.set_input(b, s, h, d, dtype=dtype, attn_mask_type=attn_mask_type, pad_ratio=pad_ratio)
+ self.set_input(b,
+ s,
+ h,
+ d,
+ dtype=dtype,
+ attn_mask_type=attn_mask_type,
+ pad_ratio=pad_ratio,
+ with_bias=with_bias)
 
  primitive_out = customcall_self_fused_attn(self.qkv,
  self.bias,
@@ -183,8 +192,16 @@ def test_forward(self, b, s, h, d, dtype, attn_mask_type, pad_ratio):
  [AttnMaskType.PADDING_MASK, AttnMaskType.CAUSAL_MASK])
  @pytest.mark.parametrize('dtype', DTYPES)
  @pytest.mark.parametrize('pad_ratio', PAD_RATIO)
- def test_forward_backward(self, b, s, h, d, dtype, attn_mask_type, pad_ratio):
- self.set_input(b, s, h, d, dtype=dtype, attn_mask_type=attn_mask_type, pad_ratio=pad_ratio)
+ @pytest.mark.parametrize('with_bias', [True, False])
+ def test_forward_backward(self, b, s, h, d, dtype, attn_mask_type, pad_ratio, with_bias):
+ self.set_input(b,
+ s,
+ h,
+ d,
+ dtype=dtype,
+ attn_mask_type=attn_mask_type,
+ pad_ratio=pad_ratio,
+ with_bias=with_bias)
 
  def grad_func(fused_attn_max_512_func, *args, **kwargs):
  # Gradient is small, use a gradient multiplier to amplify the graident
@@ -221,11 +238,11 @@ def grad_func(fused_attn_max_512_func, *args, **kwargs):
  (0, 1)))
 
  primitive_out, (primitive_dqkv,
-primitive_dbeta) = jitted_primitive(self.qkv, self.bias, self.q_token,
+primitive_dbias) = jitted_primitive(self.qkv, self.bias, self.q_token,
  self.kv_token, self.dropout_rng)
 
  reference_out, (reference_dqkv,
-reference_dbeta) = jitted_reference(self.qkv, self.bias, self.q_token,
+reference_dbias) = jitted_reference(self.qkv, self.bias, self.q_token,
  self.kv_token, self.dropout_rng)
 
  np.testing.assert_allclose(jnp.asarray(primitive_out, np.float32),
@@ -261,20 +278,22 @@ def grad_func(fused_attn_max_512_func, *args, **kwargs):
  # Padded part should be 0s
  assert jnp.allclose(invalid_primitive_dqkv, jnp.zeros_like(invalid_primitive_dqkv))
 
- # dbeta valid part
- np.testing.assert_allclose(
- jnp.asarray(primitive_dbeta[:, :, :self.valid_len, :self.valid_len], np.float32),
- jnp.asarray(reference_dbeta[:, :, :self.valid_len, :self.valid_len], np.float32),
- rtol=1e-4,
- atol=3e-5)
-
- # dbeta padded part
- np.testing.assert_allclose(
- jnp.asarray(primitive_dbeta[:, :, self.valid_len:, self.valid_len:], np.float32),
- jnp.asarray(reference_dbeta[:, :, self.valid_len:, self.valid_len:], np.float32))
-
- assert jnp.allclose(primitive_dbeta[:, :, self.valid_len:, self.valid_len:],
- jnp.zeros_like(primitive_dbeta[:, :, self.valid_len:, self.valid_len:]))
+ if self.attn_bias_type != AttnBiasType.NO_BIAS:
+ # dbias valid part
+ np.testing.assert_allclose(
+ jnp.asarray(primitive_dbias[:, :, :self.valid_len, :self.valid_len], np.float32),
+ jnp.asarray(reference_dbias[:, :, :self.valid_len, :self.valid_len], np.float32),
+ rtol=1e-4,
+ atol=3e-5)
+
+ # dbias padded part
+ np.testing.assert_allclose(
+ jnp.asarray(primitive_dbias[:, :, self.valid_len:, self.valid_len:], np.float32),
+ jnp.asarray(reference_dbias[:, :, self.valid_len:, self.valid_len:], np.float32))
+
+ assert jnp.allclose(
+ primitive_dbias[:, :, self.valid_len:, self.valid_len:],
+ jnp.zeros_like(primitive_dbias[:, :, self.valid_len:, self.valid_len:]))
 
 
 @pytest.mark.skipif(not is_fused_attn_kernel_available(),

tests/jax/test_layer.py

@@ -102,6 +102,12 @@ def compare_frozen_dict(ref_fd, test_fd, rtol=1e-05, atol=1e-08):
  _KEY_OF_DROPOUT_RATE: 0.0,
  _KEY_OF_MLP_ACTIVATIONS: (('gelu', 'linear')),
  _KEY_OF_FUSE_MLP_WI: True
+}, {
+ _KEY_OF_SCALE_ATTN_LOGITS: True,
+ _KEY_OF_LAYERNORM_TYPE: 'rmsnorm',
+ _KEY_OF_DROPOUT_RATE: 0.8,
+ _KEY_OF_MLP_ACTIVATIONS: (('gelu', 'linear')),
+ _KEY_OF_FUSE_MLP_WI: True
 }, {
  _KEY_OF_TRANSPOSE_BS: False,
  _KEY_OF_SCALE_ATTN_LOGITS: True,

transformer_engine/common/fused_attn/fused_attn_fp16_bf16_max_seqlen_512.cu

@@ -327,7 +327,6 @@ static cudnn_frontend::Tensor createSoftmaxForward(
  // NOLINTNEXTLINE(runtime/references)
  std::vector<cudnn_frontend::Operation> &ops,
  cudnn_frontend::Tensor const &prevBlockOutputTensor) {
-
  int64_t afterBMM1_dim[4] = {b, h, s_q, s_kv};
  int64_t afterBMM1_stride[4] = {h * s_q * s_kv, s_q * s_kv, s_kv, 1};
 
@@ -645,7 +644,7 @@ void fused_attn_max_512_fwd_impl(int64_t b, int64_t h, int64_t s_q, int64_t s_kv
  mask_type, tensorType};
 
  using CacheType = std::map<FADescriptor, cudnn_frontend::ExecutionPlan>;
- static CacheType fmha_fprop_cache;
+ static thread_local CacheType fmha_fprop_cache;
 
  bool enable_dropout = (dropout_probability != 0.0f);
 
@@ -668,7 +667,8 @@ void fused_attn_max_512_fwd_impl(int64_t b, int64_t h, int64_t s_q, int64_t s_kv
  createScale(b, h, s_q, s_kv, d, layout, tensorType, ops);
 
  // if bias, we need to memset the S buffer to correctly computate dbias
- auto zero_s = (bias_type != NVTE_Bias_Type::NVTE_NO_BIAS);
+ auto zero_s = (bias_type != NVTE_Bias_Type::NVTE_NO_BIAS) ||
+ (mask_type == NVTE_Mask_Type::NVTE_CAUSAL_MASK);
  auto bmm1_output = createBMM1(b, h, s_q, s_kv, d, layout, tensorType, zero_s, ops);
 
  NVTE_CHECK(bias_type != NVTE_Bias_Type::NVTE_PRE_SCALE_BIAS,
@@ -814,7 +814,7 @@ void fused_attn_max_512_bwd_impl(int64_t b, int64_t h, int64_t s_q, int64_t s_kv
  layout, bias_type, mask_type, tensorType};
 
  using CacheType = std::map<FADescriptor, cudnn_frontend::ExecutionPlan>;
- static CacheType fmha_bprop_cache;
+ static thread_local CacheType fmha_bprop_cache;
 
  auto get_plan = [&](CacheType &cache, const FADescriptor &descriptor) {
  auto it = cache.find(descriptor);

transformer_engine/common/fused_attn/fused_attn_fp8.cu

@@ -1016,7 +1016,7 @@ void fa_fwd_fp8(int64_t b, int64_t s_q, int64_t s_kv, int64_t h, int64_t d,
  NVTE_Bias_Type::NVTE_NO_BIAS, NVTE_Mask_Type::NVTE_PADDING_MASK, tensorType};
 
  using CacheType = std::map<FADescriptor, cudnn_frontend::ExecutionPlan>;
- static CacheType fa_fprop_cache;
+ static thread_local CacheType fa_fprop_cache;
 
  // Get plan from cache if cache is available, otherwise create one
  auto get_plan = [&](CacheType &cache, const FADescriptor &descriptor) {
@@ -1332,7 +1332,7 @@ void fa_bwd_fp8(int64_t b, int64_t s_q, int64_t s_kv, int64_t h, int64_t d,
  NVTE_Bias_Type::NVTE_NO_BIAS, NVTE_Mask_Type::NVTE_PADDING_MASK, tensorType};
 
  using CacheType = std::map<FADescriptor, cudnn_frontend::ExecutionPlan>;
- static CacheType fa_bprop_cache;
+ static thread_local CacheType fa_bprop_cache;
 
  // Get plan from cache if cache is available, otherwise create one
  auto get_plan = [&](CacheType &cache, const FADescriptor &descriptor) {

transformer_engine/jax/flax/transformer.py

@@ -7,6 +7,7 @@
 import functools
 from enum import Enum
 from math import sqrt
+import os
 from typing import Any, Callable, Optional, Sequence, Tuple, Union
 import warnings
 
@@ -165,8 +166,17 @@ def core_attention(query: Array,
  else:
  attn_weights = jnp.einsum('bqhd,bkhd->bhqk', query, key)
 
+ # When a bias is present, the computation is performed as Softmax(attn_weights * scale + bias).
+ # In this case, the scale can not fused into the Softmax module.
+ if bias is not None:
+ attn_weights = attn_weights * scale_factor
+ fused_scale_factor = 1.
+ else:
+ # If no bias, the scale can be fused into Softmax module
+ fused_scale_factor = scale_factor
+
  attn_weights = Softmax(softmax_type=softmax_type,
- scale_factor=scale_factor,
+ scale_factor=fused_scale_factor,
  sharding_type=softmax_sharding_type)(attn_weights, mask, bias)
 
  if not deterministic and dropout_rate > 0.:
@@ -360,12 +370,13 @@ def kv_init(key, shape, dtype):
  q_seqlen = inputs_q.shape[0] if self.transpose_batch_sequence else inputs_q.shape[1]
  kv_seqlen = inputs_kv.shape[0] if self.transpose_batch_sequence else inputs_kv.shape[1]
  fused_attn_supported_seqlen = [128, 256, 384, 512]
+ enable_fused_attn = int(os.getenv("NVTE_FUSED_ATTN", "0"))
  use_fused_attn = not decode and not self.transpose_batch_sequence and self.fuse_qkv and \
  self.dropout_rate == 0 and canonicalize_dtype in [jnp.bfloat16, jnp.float16] and \
  q_seqlen in fused_attn_supported_seqlen and kv_seqlen in fused_attn_supported_seqlen \
- and is_fused_attn_kernel_available()
+ and is_fused_attn_kernel_available() and enable_fused_attn
 
- if not use_fused_attn:
+ if enable_fused_attn and not use_fused_attn:
  reason = ""
  if decode:
  reason += f"decode=False is required but got {decode}, "

transformer_engine/jax/sharding.py

@@ -386,7 +386,7 @@ def _get_dptp_sharding_meta(input_shapes: Tuple[Tuple[int, ...]],
 
  for input_shape, dp_dim, tp_dim in zip(input_shapes, input_dp_dims, input_tp_dims):
  in_axis = {}
- if dp_dim is not None:
+ if dp_dim is not None and input_shape is not None:
  in_axis[dp_dim] = dp_axis_name
  assert input_shape[dp_dim] % dp_size == 0, \
  f"The dimension of batch in input_shape should be a multiple of " \
@@ -398,7 +398,7 @@ def _get_dptp_sharding_meta(input_shapes: Tuple[Tuple[int, ...]],
  if tp_dim is not None and tp_dim >= dp_dim:
  tp_dim = tp_dim + 1
 
- if tp_dim is not None:
+ if tp_dim is not None and input_shape is not None:
  in_axis[tp_dim] = tp_axis_name
  assert input_shape[tp_dim] % tp_size == 0, \
  f"The dimension of tensor parallel in input_shape should be a multiple of " \

tests/jax/test_fused_attn.py

@@ -113,7 +113,7 @@ def customcall_cross_fused_attn(q, kv, q_token, kv_token, dropout_rng, **kwargs)
  reason="Fused attention kernel is not supported.")
 class TestSelfFusedAttnMax512():
 
- def set_input(self, b, s, h, d, dtype, attn_mask_type, pad_ratio):
+ def set_input(self, b, s, h, d, dtype, attn_mask_type, pad_ratio, with_bias):
  key = jax.random.PRNGKey(0)
  subkeys = jax.random.split(key, 2)
 
@@ -125,16 +125,17 @@ def set_input(self, b, s, h, d, dtype, attn_mask_type, pad_ratio):
 
  min_val, max_val = -1, 1
  self.qkv = jax.random.uniform(subkeys[0], qkv_shape, dtype, min_val, max_val)
- self.bias = jax.random.uniform(subkeys[1], bias_shape, dtype, min_val, max_val)
+ self.bias = jax.random.uniform(subkeys[1], bias_shape, dtype, min_val,
+ max_val) if with_bias else None
 
  self.q_token = jnp.concatenate((jnp.ones((b, self.valid_len)), jnp.zeros((b, pad_len))),
  axis=-1)
  self.kv_token = self.q_token
 
  self.scaling_factor = 1. / math.sqrt(d)
  self.dropout_probability = 0.
- self.dropout_rng = jax.random.PRNGKey(0)
- self.attn_bias_type = AttnBiasType.POST_SCALE_BIAS
+ self.dropout_rng = jax.random.PRNGKey(0) if self.dropout_probability > 0 else None
+ self.attn_bias_type = AttnBiasType.NO_BIAS if self.bias is None else AttnBiasType.POST_SCALE_BIAS
  # deterministic = not is_training
  self.deterministic = False
 
@@ -143,9 +144,17 @@ def set_input(self, b, s, h, d, dtype, attn_mask_type, pad_ratio):
  @pytest.mark.parametrize('attn_mask_type',
  [AttnMaskType.PADDING_MASK, AttnMaskType.CAUSAL_MASK])
  @pytest.mark.parametrize('pad_ratio', PAD_RATIO)
- def test_forward(self, b, s, h, d, dtype, attn_mask_type, pad_ratio):
+ @pytest.mark.parametrize('with_bias', [True, False])
+ def test_forward(self, b, s, h, d, dtype, attn_mask_type, pad_ratio, with_bias):
 
- self.set_input(b, s, h, d, dtype=dtype, attn_mask_type=attn_mask_type, pad_ratio=pad_ratio)
+ self.set_input(b,
+ s,
+ h,
+ d,
+ dtype=dtype,
+ attn_mask_type=attn_mask_type,
+ pad_ratio=pad_ratio,
+ with_bias=with_bias)
 
  primitive_out = customcall_self_fused_attn(self.qkv,
  self.bias,
@@ -183,8 +192,16 @@ def test_forward(self, b, s, h, d, dtype, attn_mask_type, pad_ratio):
  [AttnMaskType.PADDING_MASK, AttnMaskType.CAUSAL_MASK])
  @pytest.mark.parametrize('dtype', DTYPES)
  @pytest.mark.parametrize('pad_ratio', PAD_RATIO)
- def test_forward_backward(self, b, s, h, d, dtype, attn_mask_type, pad_ratio):
- self.set_input(b, s, h, d, dtype=dtype, attn_mask_type=attn_mask_type, pad_ratio=pad_ratio)
+ @pytest.mark.parametrize('with_bias', [True, False])
+ def test_forward_backward(self, b, s, h, d, dtype, attn_mask_type, pad_ratio, with_bias):
+ self.set_input(b,
+ s,
+ h,
+ d,
+ dtype=dtype,
+ attn_mask_type=attn_mask_type,
+ pad_ratio=pad_ratio,
+ with_bias=with_bias)
 
  def grad_func(fused_attn_max_512_func, *args, **kwargs):
  # Gradient is small, use a gradient multiplier to amplify the graident
@@ -221,11 +238,11 @@ def grad_func(fused_attn_max_512_func, *args, **kwargs):
  (0, 1)))
 
  primitive_out, (primitive_dqkv,
-primitive_dbeta) = jitted_primitive(self.qkv, self.bias, self.q_token,
+primitive_dbias) = jitted_primitive(self.qkv, self.bias, self.q_token,
  self.kv_token, self.dropout_rng)
 
  reference_out, (reference_dqkv,
-reference_dbeta) = jitted_reference(self.qkv, self.bias, self.q_token,
+reference_dbias) = jitted_reference(self.qkv, self.bias, self.q_token,
  self.kv_token, self.dropout_rng)
 
  np.testing.assert_allclose(jnp.asarray(primitive_out, np.float32),
@@ -261,20 +278,22 @@ def grad_func(fused_attn_max_512_func, *args, **kwargs):
  # Padded part should be 0s
  assert jnp.allclose(invalid_primitive_dqkv, jnp.zeros_like(invalid_primitive_dqkv))
 
- # dbeta valid part
- np.testing.assert_allclose(
- jnp.asarray(primitive_dbeta[:, :, :self.valid_len, :self.valid_len], np.float32),
- jnp.asarray(reference_dbeta[:, :, :self.valid_len, :self.valid_len], np.float32),
- rtol=1e-4,
- atol=3e-5)
-
- # dbeta padded part
- np.testing.assert_allclose(
- jnp.asarray(primitive_dbeta[:, :, self.valid_len:, self.valid_len:], np.float32),
- jnp.asarray(reference_dbeta[:, :, self.valid_len:, self.valid_len:], np.float32))
-
- assert jnp.allclose(primitive_dbeta[:, :, self.valid_len:, self.valid_len:],
- jnp.zeros_like(primitive_dbeta[:, :, self.valid_len:, self.valid_len:]))
+ if self.attn_bias_type != AttnBiasType.NO_BIAS:
+ # dbias valid part
+ np.testing.assert_allclose(
+ jnp.asarray(primitive_dbias[:, :, :self.valid_len, :self.valid_len], np.float32),
+ jnp.asarray(reference_dbias[:, :, :self.valid_len, :self.valid_len], np.float32),
+ rtol=1e-4,
+ atol=3e-5)
+
+ # dbias padded part
+ np.testing.assert_allclose(
+ jnp.asarray(primitive_dbias[:, :, self.valid_len:, self.valid_len:], np.float32),
+ jnp.asarray(reference_dbias[:, :, self.valid_len:, self.valid_len:], np.float32))
+
+ assert jnp.allclose(
+ primitive_dbias[:, :, self.valid_len:, self.valid_len:],
+ jnp.zeros_like(primitive_dbias[:, :, self.valid_len:, self.valid_len:]))
 
 
 @pytest.mark.skipif(not is_fused_attn_kernel_available(),

tests/jax/test_layer.py

@@ -102,6 +102,12 @@ def compare_frozen_dict(ref_fd, test_fd, rtol=1e-05, atol=1e-08):
  _KEY_OF_DROPOUT_RATE: 0.0,
  _KEY_OF_MLP_ACTIVATIONS: (('gelu', 'linear')),
  _KEY_OF_FUSE_MLP_WI: True
+}, {
+ _KEY_OF_SCALE_ATTN_LOGITS: True,
+ _KEY_OF_LAYERNORM_TYPE: 'rmsnorm',
+ _KEY_OF_DROPOUT_RATE: 0.8,
+ _KEY_OF_MLP_ACTIVATIONS: (('gelu', 'linear')),
+ _KEY_OF_FUSE_MLP_WI: True
 }, {
  _KEY_OF_TRANSPOSE_BS: False,
  _KEY_OF_SCALE_ATTN_LOGITS: True,

transformer_engine/common/fused_attn/fused_attn_fp16_bf16_max_seqlen_512.cu

@@ -327,7 +327,6 @@ static cudnn_frontend::Tensor createSoftmaxForward(
  // NOLINTNEXTLINE(runtime/references)
  std::vector<cudnn_frontend::Operation> &ops,
  cudnn_frontend::Tensor const &prevBlockOutputTensor) {
-
  int64_t afterBMM1_dim[4] = {b, h, s_q, s_kv};
  int64_t afterBMM1_stride[4] = {h * s_q * s_kv, s_q * s_kv, s_kv, 1};
 
@@ -645,7 +644,7 @@ void fused_attn_max_512_fwd_impl(int64_t b, int64_t h, int64_t s_q, int64_t s_kv
  mask_type, tensorType};
 
  using CacheType = std::map<FADescriptor, cudnn_frontend::ExecutionPlan>;
- static CacheType fmha_fprop_cache;
+ static thread_local CacheType fmha_fprop_cache;
 
  bool enable_dropout = (dropout_probability != 0.0f);
 
@@ -668,7 +667,8 @@ void fused_attn_max_512_fwd_impl(int64_t b, int64_t h, int64_t s_q, int64_t s_kv
  createScale(b, h, s_q, s_kv, d, layout, tensorType, ops);
 
  // if bias, we need to memset the S buffer to correctly computate dbias
- auto zero_s = (bias_type != NVTE_Bias_Type::NVTE_NO_BIAS);
+ auto zero_s = (bias_type != NVTE_Bias_Type::NVTE_NO_BIAS) ||
+ (mask_type == NVTE_Mask_Type::NVTE_CAUSAL_MASK);
  auto bmm1_output = createBMM1(b, h, s_q, s_kv, d, layout, tensorType, zero_s, ops);
 
  NVTE_CHECK(bias_type != NVTE_Bias_Type::NVTE_PRE_SCALE_BIAS,
@@ -814,7 +814,7 @@ void fused_attn_max_512_bwd_impl(int64_t b, int64_t h, int64_t s_q, int64_t s_kv
  layout, bias_type, mask_type, tensorType};
 
  using CacheType = std::map<FADescriptor, cudnn_frontend::ExecutionPlan>;
- static CacheType fmha_bprop_cache;
+ static thread_local CacheType fmha_bprop_cache;
 
  auto get_plan = [&](CacheType &cache, const FADescriptor &descriptor) {
  auto it = cache.find(descriptor);

transformer_engine/common/fused_attn/fused_attn_fp8.cu

@@ -1016,7 +1016,7 @@ void fa_fwd_fp8(int64_t b, int64_t s_q, int64_t s_kv, int64_t h, int64_t d,
  NVTE_Bias_Type::NVTE_NO_BIAS, NVTE_Mask_Type::NVTE_PADDING_MASK, tensorType};
 
  using CacheType = std::map<FADescriptor, cudnn_frontend::ExecutionPlan>;
- static CacheType fa_fprop_cache;
+ static thread_local CacheType fa_fprop_cache;
 
  // Get plan from cache if cache is available, otherwise create one
  auto get_plan = [&](CacheType &cache, const FADescriptor &descriptor) {
@@ -1332,7 +1332,7 @@ void fa_bwd_fp8(int64_t b, int64_t s_q, int64_t s_kv, int64_t h, int64_t d,
  NVTE_Bias_Type::NVTE_NO_BIAS, NVTE_Mask_Type::NVTE_PADDING_MASK, tensorType};
 
  using CacheType = std::map<FADescriptor, cudnn_frontend::ExecutionPlan>;
- static CacheType fa_bprop_cache;
+ static thread_local CacheType fa_bprop_cache;
 
  // Get plan from cache if cache is available, otherwise create one
  auto get_plan = [&](CacheType &cache, const FADescriptor &descriptor) {

transformer_engine/jax/flax/transformer.py

@@ -7,6 +7,7 @@
 import functools
 from enum import Enum
 from math import sqrt
+import os
 from typing import Any, Callable, Optional, Sequence, Tuple, Union
 import warnings
 
@@ -165,8 +166,17 @@ def core_attention(query: Array,
  else:
  attn_weights = jnp.einsum('bqhd,bkhd->bhqk', query, key)
 
+ # When a bias is present, the computation is performed as Softmax(attn_weights * scale + bias).
+ # In this case, the scale can not fused into the Softmax module.
+ if bias is not None:
+ attn_weights = attn_weights * scale_factor
+ fused_scale_factor = 1.
+ else:
+ # If no bias, the scale can be fused into Softmax module
+ fused_scale_factor = scale_factor
+
  attn_weights = Softmax(softmax_type=softmax_type,
- scale_factor=scale_factor,
+ scale_factor=fused_scale_factor,
  sharding_type=softmax_sharding_type)(attn_weights, mask, bias)
 
  if not deterministic and dropout_rate > 0.:
@@ -360,12 +370,13 @@ def kv_init(key, shape, dtype):
  q_seqlen = inputs_q.shape[0] if self.transpose_batch_sequence else inputs_q.shape[1]
  kv_seqlen = inputs_kv.shape[0] if self.transpose_batch_sequence else inputs_kv.shape[1]
  fused_attn_supported_seqlen = [128, 256, 384, 512]
+ enable_fused_attn = int(os.getenv("NVTE_FUSED_ATTN", "0"))
  use_fused_attn = not decode and not self.transpose_batch_sequence and self.fuse_qkv and \
  self.dropout_rate == 0 and canonicalize_dtype in [jnp.bfloat16, jnp.float16] and \
  q_seqlen in fused_attn_supported_seqlen and kv_seqlen in fused_attn_supported_seqlen \
- and is_fused_attn_kernel_available()
+ and is_fused_attn_kernel_available() and enable_fused_attn
 
- if not use_fused_attn:
+ if enable_fused_attn and not use_fused_attn:
  reason = ""
  if decode:
  reason += f"decode=False is required but got {decode}, "

transformer_engine/jax/sharding.py

@@ -386,7 +386,7 @@ def _get_dptp_sharding_meta(input_shapes: Tuple[Tuple[int, ...]],
 
  for input_shape, dp_dim, tp_dim in zip(input_shapes, input_dp_dims, input_tp_dims):
  in_axis = {}
- if dp_dim is not None:
+ if dp_dim is not None and input_shape is not None:
  in_axis[dp_dim] = dp_axis_name
  assert input_shape[dp_dim] % dp_size == 0, \
  f"The dimension of batch in input_shape should be a multiple of " \
@@ -398,7 +398,7 @@ def _get_dptp_sharding_meta(input_shapes: Tuple[Tuple[int, ...]],
  if tp_dim is not None and tp_dim >= dp_dim:
  tp_dim = tp_dim + 1
 
- if tp_dim is not None:
+ if tp_dim is not None and input_shape is not None:
  in_axis[tp_dim] = tp_axis_name
  assert input_shape[tp_dim] % tp_size == 0, \
  f"The dimension of tensor parallel in input_shape should be a multiple of " \

tests/jax/test_fused_attn.py

@@ -113,7 +113,7 @@ def customcall_cross_fused_attn(q, kv, q_token, kv_token, dropout_rng, **kwargs)
  reason="Fused attention kernel is not supported.")
 class TestSelfFusedAttnMax512():
 
- def set_input(self, b, s, h, d, dtype, attn_mask_type, pad_ratio):
+ def set_input(self, b, s, h, d, dtype, attn_mask_type, pad_ratio, with_bias):
  key = jax.random.PRNGKey(0)
  subkeys = jax.random.split(key, 2)
 
@@ -125,16 +125,17 @@ def set_input(self, b, s, h, d, dtype, attn_mask_type, pad_ratio):
 
  min_val, max_val = -1, 1
  self.qkv = jax.random.uniform(subkeys[0], qkv_shape, dtype, min_val, max_val)
- self.bias = jax.random.uniform(subkeys[1], bias_shape, dtype, min_val, max_val)
+ self.bias = jax.random.uniform(subkeys[1], bias_shape, dtype, min_val,
+ max_val) if with_bias else None
 
  self.q_token = jnp.concatenate((jnp.ones((b, self.valid_len)), jnp.zeros((b, pad_len))),
  axis=-1)
  self.kv_token = self.q_token
 
  self.scaling_factor = 1. / math.sqrt(d)
  self.dropout_probability = 0.
- self.dropout_rng = jax.random.PRNGKey(0)
- self.attn_bias_type = AttnBiasType.POST_SCALE_BIAS
+ self.dropout_rng = jax.random.PRNGKey(0) if self.dropout_probability > 0 else None
+ self.attn_bias_type = AttnBiasType.NO_BIAS if self.bias is None else AttnBiasType.POST_SCALE_BIAS
  # deterministic = not is_training
  self.deterministic = False
 
@@ -143,9 +144,17 @@ def set_input(self, b, s, h, d, dtype, attn_mask_type, pad_ratio):
  @pytest.mark.parametrize('attn_mask_type',
  [AttnMaskType.PADDING_MASK, AttnMaskType.CAUSAL_MASK])
  @pytest.mark.parametrize('pad_ratio', PAD_RATIO)
- def test_forward(self, b, s, h, d, dtype, attn_mask_type, pad_ratio):
+ @pytest.mark.parametrize('with_bias', [True, False])
+ def test_forward(self, b, s, h, d, dtype, attn_mask_type, pad_ratio, with_bias):
 
- self.set_input(b, s, h, d, dtype=dtype, attn_mask_type=attn_mask_type, pad_ratio=pad_ratio)
+ self.set_input(b,
+ s,
+ h,
+ d,
+ dtype=dtype,
+ attn_mask_type=attn_mask_type,
+ pad_ratio=pad_ratio,
+ with_bias=with_bias)
 
  primitive_out = customcall_self_fused_attn(self.qkv,
  self.bias,
@@ -183,8 +192,16 @@ def test_forward(self, b, s, h, d, dtype, attn_mask_type, pad_ratio):
  [AttnMaskType.PADDING_MASK, AttnMaskType.CAUSAL_MASK])
  @pytest.mark.parametrize('dtype', DTYPES)
  @pytest.mark.parametrize('pad_ratio', PAD_RATIO)
- def test_forward_backward(self, b, s, h, d, dtype, attn_mask_type, pad_ratio):
- self.set_input(b, s, h, d, dtype=dtype, attn_mask_type=attn_mask_type, pad_ratio=pad_ratio)
+ @pytest.mark.parametrize('with_bias', [True, False])
+ def test_forward_backward(self, b, s, h, d, dtype, attn_mask_type, pad_ratio, with_bias):
+ self.set_input(b,
+ s,
+ h,
+ d,
+ dtype=dtype,
+ attn_mask_type=attn_mask_type,
+ pad_ratio=pad_ratio,
+ with_bias=with_bias)
 
  def grad_func(fused_attn_max_512_func, *args, **kwargs):
  # Gradient is small, use a gradient multiplier to amplify the graident
@@ -221,11 +238,11 @@ def grad_func(fused_attn_max_512_func, *args, **kwargs):
  (0, 1)))
 
  primitive_out, (primitive_dqkv,
-primitive_dbeta) = jitted_primitive(self.qkv, self.bias, self.q_token,
+primitive_dbias) = jitted_primitive(self.qkv, self.bias, self.q_token,
  self.kv_token, self.dropout_rng)
 
  reference_out, (reference_dqkv,
-reference_dbeta) = jitted_reference(self.qkv, self.bias, self.q_token,
+reference_dbias) = jitted_reference(self.qkv, self.bias, self.q_token,
  self.kv_token, self.dropout_rng)
 
  np.testing.assert_allclose(jnp.asarray(primitive_out, np.float32),
@@ -261,20 +278,22 @@ def grad_func(fused_attn_max_512_func, *args, **kwargs):
  # Padded part should be 0s
  assert jnp.allclose(invalid_primitive_dqkv, jnp.zeros_like(invalid_primitive_dqkv))
 
- # dbeta valid part
- np.testing.assert_allclose(
- jnp.asarray(primitive_dbeta[:, :, :self.valid_len, :self.valid_len], np.float32),
- jnp.asarray(reference_dbeta[:, :, :self.valid_len, :self.valid_len], np.float32),
- rtol=1e-4,
- atol=3e-5)
-
- # dbeta padded part
- np.testing.assert_allclose(
- jnp.asarray(primitive_dbeta[:, :, self.valid_len:, self.valid_len:], np.float32),
- jnp.asarray(reference_dbeta[:, :, self.valid_len:, self.valid_len:], np.float32))
-
- assert jnp.allclose(primitive_dbeta[:, :, self.valid_len:, self.valid_len:],
- jnp.zeros_like(primitive_dbeta[:, :, self.valid_len:, self.valid_len:]))
+ if self.attn_bias_type != AttnBiasType.NO_BIAS:
+ # dbias valid part
+ np.testing.assert_allclose(
+ jnp.asarray(primitive_dbias[:, :, :self.valid_len, :self.valid_len], np.float32),
+ jnp.asarray(reference_dbias[:, :, :self.valid_len, :self.valid_len], np.float32),
+ rtol=1e-4,
+ atol=3e-5)
+
+ # dbias padded part
+ np.testing.assert_allclose(
+ jnp.asarray(primitive_dbias[:, :, self.valid_len:, self.valid_len:], np.float32),
+ jnp.asarray(reference_dbias[:, :, self.valid_len:, self.valid_len:], np.float32))
+
+ assert jnp.allclose(
+ primitive_dbias[:, :, self.valid_len:, self.valid_len:],
+ jnp.zeros_like(primitive_dbias[:, :, self.valid_len:, self.valid_len:]))
 
 
 @pytest.mark.skipif(not is_fused_attn_kernel_available(),

tests/jax/test_layer.py

@@ -102,6 +102,12 @@ def compare_frozen_dict(ref_fd, test_fd, rtol=1e-05, atol=1e-08):
  _KEY_OF_DROPOUT_RATE: 0.0,
  _KEY_OF_MLP_ACTIVATIONS: (('gelu', 'linear')),
  _KEY_OF_FUSE_MLP_WI: True
+}, {
+ _KEY_OF_SCALE_ATTN_LOGITS: True,
+ _KEY_OF_LAYERNORM_TYPE: 'rmsnorm',
+ _KEY_OF_DROPOUT_RATE: 0.8,
+ _KEY_OF_MLP_ACTIVATIONS: (('gelu', 'linear')),
+ _KEY_OF_FUSE_MLP_WI: True
 }, {
  _KEY_OF_TRANSPOSE_BS: False,
  _KEY_OF_SCALE_ATTN_LOGITS: True,

transformer_engine/common/fused_attn/fused_attn_fp16_bf16_max_seqlen_512.cu

@@ -327,7 +327,6 @@ static cudnn_frontend::Tensor createSoftmaxForward(
  // NOLINTNEXTLINE(runtime/references)
  std::vector<cudnn_frontend::Operation> &ops,
  cudnn_frontend::Tensor const &prevBlockOutputTensor) {
-
  int64_t afterBMM1_dim[4] = {b, h, s_q, s_kv};
  int64_t afterBMM1_stride[4] = {h * s_q * s_kv, s_q * s_kv, s_kv, 1};
 
@@ -645,7 +644,7 @@ void fused_attn_max_512_fwd_impl(int64_t b, int64_t h, int64_t s_q, int64_t s_kv
  mask_type, tensorType};
 
  using CacheType = std::map<FADescriptor, cudnn_frontend::ExecutionPlan>;
- static CacheType fmha_fprop_cache;
+ static thread_local CacheType fmha_fprop_cache;
 
  bool enable_dropout = (dropout_probability != 0.0f);
 
@@ -668,7 +667,8 @@ void fused_attn_max_512_fwd_impl(int64_t b, int64_t h, int64_t s_q, int64_t s_kv
  createScale(b, h, s_q, s_kv, d, layout, tensorType, ops);
 
  // if bias, we need to memset the S buffer to correctly computate dbias
- auto zero_s = (bias_type != NVTE_Bias_Type::NVTE_NO_BIAS);
+ auto zero_s = (bias_type != NVTE_Bias_Type::NVTE_NO_BIAS) ||
+ (mask_type == NVTE_Mask_Type::NVTE_CAUSAL_MASK);
  auto bmm1_output = createBMM1(b, h, s_q, s_kv, d, layout, tensorType, zero_s, ops);
 
  NVTE_CHECK(bias_type != NVTE_Bias_Type::NVTE_PRE_SCALE_BIAS,
@@ -814,7 +814,7 @@ void fused_attn_max_512_bwd_impl(int64_t b, int64_t h, int64_t s_q, int64_t s_kv
  layout, bias_type, mask_type, tensorType};
 
  using CacheType = std::map<FADescriptor, cudnn_frontend::ExecutionPlan>;
- static CacheType fmha_bprop_cache;
+ static thread_local CacheType fmha_bprop_cache;
 
  auto get_plan = [&](CacheType &cache, const FADescriptor &descriptor) {
  auto it = cache.find(descriptor);

transformer_engine/common/fused_attn/fused_attn_fp8.cu

@@ -1016,7 +1016,7 @@ void fa_fwd_fp8(int64_t b, int64_t s_q, int64_t s_kv, int64_t h, int64_t d,
  NVTE_Bias_Type::NVTE_NO_BIAS, NVTE_Mask_Type::NVTE_PADDING_MASK, tensorType};
 
  using CacheType = std::map<FADescriptor, cudnn_frontend::ExecutionPlan>;
- static CacheType fa_fprop_cache;
+ static thread_local CacheType fa_fprop_cache;
 
  // Get plan from cache if cache is available, otherwise create one
  auto get_plan = [&](CacheType &cache, const FADescriptor &descriptor) {
@@ -1332,7 +1332,7 @@ void fa_bwd_fp8(int64_t b, int64_t s_q, int64_t s_kv, int64_t h, int64_t d,
  NVTE_Bias_Type::NVTE_NO_BIAS, NVTE_Mask_Type::NVTE_PADDING_MASK, tensorType};
 
  using CacheType = std::map<FADescriptor, cudnn_frontend::ExecutionPlan>;
- static CacheType fa_bprop_cache;
+ static thread_local CacheType fa_bprop_cache;
 
  // Get plan from cache if cache is available, otherwise create one
  auto get_plan = [&](CacheType &cache, const FADescriptor &descriptor) {

transformer_engine/jax/flax/transformer.py

@@ -7,6 +7,7 @@
 import functools
 from enum import Enum
 from math import sqrt
+import os
 from typing import Any, Callable, Optional, Sequence, Tuple, Union
 import warnings
 
@@ -165,8 +166,17 @@ def core_attention(query: Array,
  else:
  attn_weights = jnp.einsum('bqhd,bkhd->bhqk', query, key)
 
+ # When a bias is present, the computation is performed as Softmax(attn_weights * scale + bias).
+ # In this case, the scale can not fused into the Softmax module.
+ if bias is not None:
+ attn_weights = attn_weights * scale_factor
+ fused_scale_factor = 1.
+ else:
+ # If no bias, the scale can be fused into Softmax module
+ fused_scale_factor = scale_factor
+
  attn_weights = Softmax(softmax_type=softmax_type,
- scale_factor=scale_factor,
+ scale_factor=fused_scale_factor,
  sharding_type=softmax_sharding_type)(attn_weights, mask, bias)
 
  if not deterministic and dropout_rate > 0.:
@@ -360,12 +370,13 @@ def kv_init(key, shape, dtype):
  q_seqlen = inputs_q.shape[0] if self.transpose_batch_sequence else inputs_q.shape[1]
  kv_seqlen = inputs_kv.shape[0] if self.transpose_batch_sequence else inputs_kv.shape[1]
  fused_attn_supported_seqlen = [128, 256, 384, 512]
+ enable_fused_attn = int(os.getenv("NVTE_FUSED_ATTN", "0"))
  use_fused_attn = not decode and not self.transpose_batch_sequence and self.fuse_qkv and \
  self.dropout_rate == 0 and canonicalize_dtype in [jnp.bfloat16, jnp.float16] and \
  q_seqlen in fused_attn_supported_seqlen and kv_seqlen in fused_attn_supported_seqlen \
- and is_fused_attn_kernel_available()
+ and is_fused_attn_kernel_available() and enable_fused_attn
 
- if not use_fused_attn:
+ if enable_fused_attn and not use_fused_attn:
  reason = ""
  if decode:
  reason += f"decode=False is required but got {decode}, "

transformer_engine/jax/sharding.py

@@ -386,7 +386,7 @@ def _get_dptp_sharding_meta(input_shapes: Tuple[Tuple[int, ...]],
 
  for input_shape, dp_dim, tp_dim in zip(input_shapes, input_dp_dims, input_tp_dims):
  in_axis = {}
- if dp_dim is not None:
+ if dp_dim is not None and input_shape is not None:
  in_axis[dp_dim] = dp_axis_name
  assert input_shape[dp_dim] % dp_size == 0, \
  f"The dimension of batch in input_shape should be a multiple of " \
@@ -398,7 +398,7 @@ def _get_dptp_sharding_meta(input_shapes: Tuple[Tuple[int, ...]],
  if tp_dim is not None and tp_dim >= dp_dim:
  tp_dim = tp_dim + 1
 
- if tp_dim is not None:
+ if tp_dim is not None and input_shape is not None:
  in_axis[tp_dim] = tp_axis_name
  assert input_shape[tp_dim] % tp_size == 0, \
  f"The dimension of tensor parallel in input_shape should be a multiple of " \

tests/jax/test_fused_attn.py

@@ -113,7 +113,7 @@ def customcall_cross_fused_attn(q, kv, q_token, kv_token, dropout_rng, **kwargs)
  reason="Fused attention kernel is not supported.")
 class TestSelfFusedAttnMax512():
 
- def set_input(self, b, s, h, d, dtype, attn_mask_type, pad_ratio):
+ def set_input(self, b, s, h, d, dtype, attn_mask_type, pad_ratio, with_bias):
  key = jax.random.PRNGKey(0)
  subkeys = jax.random.split(key, 2)
 
@@ -125,16 +125,17 @@ def set_input(self, b, s, h, d, dtype, attn_mask_type, pad_ratio):
 
  min_val, max_val = -1, 1
  self.qkv = jax.random.uniform(subkeys[0], qkv_shape, dtype, min_val, max_val)
- self.bias = jax.random.uniform(subkeys[1], bias_shape, dtype, min_val, max_val)
+ self.bias = jax.random.uniform(subkeys[1], bias_shape, dtype, min_val,
+ max_val) if with_bias else None
 
  self.q_token = jnp.concatenate((jnp.ones((b, self.valid_len)), jnp.zeros((b, pad_len))),
  axis=-1)
  self.kv_token = self.q_token
 
  self.scaling_factor = 1. / math.sqrt(d)
  self.dropout_probability = 0.
- self.dropout_rng = jax.random.PRNGKey(0)
- self.attn_bias_type = AttnBiasType.POST_SCALE_BIAS
+ self.dropout_rng = jax.random.PRNGKey(0) if self.dropout_probability > 0 else None
+ self.attn_bias_type = AttnBiasType.NO_BIAS if self.bias is None else AttnBiasType.POST_SCALE_BIAS
  # deterministic = not is_training
  self.deterministic = False
 
@@ -143,9 +144,17 @@ def set_input(self, b, s, h, d, dtype, attn_mask_type, pad_ratio):
  @pytest.mark.parametrize('attn_mask_type',
  [AttnMaskType.PADDING_MASK, AttnMaskType.CAUSAL_MASK])
  @pytest.mark.parametrize('pad_ratio', PAD_RATIO)
- def test_forward(self, b, s, h, d, dtype, attn_mask_type, pad_ratio):
+ @pytest.mark.parametrize('with_bias', [True, False])
+ def test_forward(self, b, s, h, d, dtype, attn_mask_type, pad_ratio, with_bias):
 
- self.set_input(b, s, h, d, dtype=dtype, attn_mask_type=attn_mask_type, pad_ratio=pad_ratio)
+ self.set_input(b,
+ s,
+ h,
+ d,
+ dtype=dtype,
+ attn_mask_type=attn_mask_type,
+ pad_ratio=pad_ratio,
+ with_bias=with_bias)
 
  primitive_out = customcall_self_fused_attn(self.qkv,
  self.bias,
@@ -183,8 +192,16 @@ def test_forward(self, b, s, h, d, dtype, attn_mask_type, pad_ratio):
  [AttnMaskType.PADDING_MASK, AttnMaskType.CAUSAL_MASK])
  @pytest.mark.parametrize('dtype', DTYPES)
  @pytest.mark.parametrize('pad_ratio', PAD_RATIO)
- def test_forward_backward(self, b, s, h, d, dtype, attn_mask_type, pad_ratio):
- self.set_input(b, s, h, d, dtype=dtype, attn_mask_type=attn_mask_type, pad_ratio=pad_ratio)
+ @pytest.mark.parametrize('with_bias', [True, False])
+ def test_forward_backward(self, b, s, h, d, dtype, attn_mask_type, pad_ratio, with_bias):
+ self.set_input(b,
+ s,
+ h,
+ d,
+ dtype=dtype,
+ attn_mask_type=attn_mask_type,
+ pad_ratio=pad_ratio,
+ with_bias=with_bias)
 
  def grad_func(fused_attn_max_512_func, *args, **kwargs):
  # Gradient is small, use a gradient multiplier to amplify the graident
@@ -221,11 +238,11 @@ def grad_func(fused_attn_max_512_func, *args, **kwargs):
  (0, 1)))
 
  primitive_out, (primitive_dqkv,
-primitive_dbeta) = jitted_primitive(self.qkv, self.bias, self.q_token,
+primitive_dbias) = jitted_primitive(self.qkv, self.bias, self.q_token,
  self.kv_token, self.dropout_rng)
 
  reference_out, (reference_dqkv,
-reference_dbeta) = jitted_reference(self.qkv, self.bias, self.q_token,
+reference_dbias) = jitted_reference(self.qkv, self.bias, self.q_token,
  self.kv_token, self.dropout_rng)
 
  np.testing.assert_allclose(jnp.asarray(primitive_out, np.float32),
@@ -261,20 +278,22 @@ def grad_func(fused_attn_max_512_func, *args, **kwargs):
  # Padded part should be 0s
  assert jnp.allclose(invalid_primitive_dqkv, jnp.zeros_like(invalid_primitive_dqkv))
 
- # dbeta valid part
- np.testing.assert_allclose(
- jnp.asarray(primitive_dbeta[:, :, :self.valid_len, :self.valid_len], np.float32),
- jnp.asarray(reference_dbeta[:, :, :self.valid_len, :self.valid_len], np.float32),
- rtol=1e-4,
- atol=3e-5)
-
- # dbeta padded part
- np.testing.assert_allclose(
- jnp.asarray(primitive_dbeta[:, :, self.valid_len:, self.valid_len:], np.float32),
- jnp.asarray(reference_dbeta[:, :, self.valid_len:, self.valid_len:], np.float32))
-
- assert jnp.allclose(primitive_dbeta[:, :, self.valid_len:, self.valid_len:],
- jnp.zeros_like(primitive_dbeta[:, :, self.valid_len:, self.valid_len:]))
+ if self.attn_bias_type != AttnBiasType.NO_BIAS:
+ # dbias valid part
+ np.testing.assert_allclose(
+ jnp.asarray(primitive_dbias[:, :, :self.valid_len, :self.valid_len], np.float32),
+ jnp.asarray(reference_dbias[:, :, :self.valid_len, :self.valid_len], np.float32),
+ rtol=1e-4,
+ atol=3e-5)
+
+ # dbias padded part
+ np.testing.assert_allclose(
+ jnp.asarray(primitive_dbias[:, :, self.valid_len:, self.valid_len:], np.float32),
+ jnp.asarray(reference_dbias[:, :, self.valid_len:, self.valid_len:], np.float32))
+
+ assert jnp.allclose(
+ primitive_dbias[:, :, self.valid_len:, self.valid_len:],
+ jnp.zeros_like(primitive_dbias[:, :, self.valid_len:, self.valid_len:]))
 
 
 @pytest.mark.skipif(not is_fused_attn_kernel_available(),

tests/jax/test_layer.py

@@ -102,6 +102,12 @@ def compare_frozen_dict(ref_fd, test_fd, rtol=1e-05, atol=1e-08):
  _KEY_OF_DROPOUT_RATE: 0.0,
  _KEY_OF_MLP_ACTIVATIONS: (('gelu', 'linear')),
  _KEY_OF_FUSE_MLP_WI: True
+}, {
+ _KEY_OF_SCALE_ATTN_LOGITS: True,
+ _KEY_OF_LAYERNORM_TYPE: 'rmsnorm',
+ _KEY_OF_DROPOUT_RATE: 0.8,
+ _KEY_OF_MLP_ACTIVATIONS: (('gelu', 'linear')),
+ _KEY_OF_FUSE_MLP_WI: True
 }, {
  _KEY_OF_TRANSPOSE_BS: False,
  _KEY_OF_SCALE_ATTN_LOGITS: True,

transformer_engine/common/fused_attn/fused_attn_fp16_bf16_max_seqlen_512.cu

@@ -327,7 +327,6 @@ static cudnn_frontend::Tensor createSoftmaxForward(
  // NOLINTNEXTLINE(runtime/references)
  std::vector<cudnn_frontend::Operation> &ops,
  cudnn_frontend::Tensor const &prevBlockOutputTensor) {
-
  int64_t afterBMM1_dim[4] = {b, h, s_q, s_kv};
  int64_t afterBMM1_stride[4] = {h * s_q * s_kv, s_q * s_kv, s_kv, 1};
 
@@ -645,7 +644,7 @@ void fused_attn_max_512_fwd_impl(int64_t b, int64_t h, int64_t s_q, int64_t s_kv
  mask_type, tensorType};
 
  using CacheType = std::map<FADescriptor, cudnn_frontend::ExecutionPlan>;
- static CacheType fmha_fprop_cache;
+ static thread_local CacheType fmha_fprop_cache;
 
  bool enable_dropout = (dropout_probability != 0.0f);
 
@@ -668,7 +667,8 @@ void fused_attn_max_512_fwd_impl(int64_t b, int64_t h, int64_t s_q, int64_t s_kv
  createScale(b, h, s_q, s_kv, d, layout, tensorType, ops);
 
  // if bias, we need to memset the S buffer to correctly computate dbias
- auto zero_s = (bias_type != NVTE_Bias_Type::NVTE_NO_BIAS);
+ auto zero_s = (bias_type != NVTE_Bias_Type::NVTE_NO_BIAS) ||
+ (mask_type == NVTE_Mask_Type::NVTE_CAUSAL_MASK);
  auto bmm1_output = createBMM1(b, h, s_q, s_kv, d, layout, tensorType, zero_s, ops);
 
  NVTE_CHECK(bias_type != NVTE_Bias_Type::NVTE_PRE_SCALE_BIAS,
@@ -814,7 +814,7 @@ void fused_attn_max_512_bwd_impl(int64_t b, int64_t h, int64_t s_q, int64_t s_kv
  layout, bias_type, mask_type, tensorType};
 
  using CacheType = std::map<FADescriptor, cudnn_frontend::ExecutionPlan>;
- static CacheType fmha_bprop_cache;
+ static thread_local CacheType fmha_bprop_cache;
 
  auto get_plan = [&](CacheType &cache, const FADescriptor &descriptor) {
  auto it = cache.find(descriptor);

transformer_engine/common/fused_attn/fused_attn_fp8.cu

@@ -1016,7 +1016,7 @@ void fa_fwd_fp8(int64_t b, int64_t s_q, int64_t s_kv, int64_t h, int64_t d,
  NVTE_Bias_Type::NVTE_NO_BIAS, NVTE_Mask_Type::NVTE_PADDING_MASK, tensorType};
 
  using CacheType = std::map<FADescriptor, cudnn_frontend::ExecutionPlan>;
- static CacheType fa_fprop_cache;
+ static thread_local CacheType fa_fprop_cache;
 
  // Get plan from cache if cache is available, otherwise create one
  auto get_plan = [&](CacheType &cache, const FADescriptor &descriptor) {
@@ -1332,7 +1332,7 @@ void fa_bwd_fp8(int64_t b, int64_t s_q, int64_t s_kv, int64_t h, int64_t d,
  NVTE_Bias_Type::NVTE_NO_BIAS, NVTE_Mask_Type::NVTE_PADDING_MASK, tensorType};
 
  using CacheType = std::map<FADescriptor, cudnn_frontend::ExecutionPlan>;
- static CacheType fa_bprop_cache;
+ static thread_local CacheType fa_bprop_cache;
 
  // Get plan from cache if cache is available, otherwise create one
  auto get_plan = [&](CacheType &cache, const FADescriptor &descriptor) {

transformer_engine/jax/flax/transformer.py

@@ -7,6 +7,7 @@
 import functools
 from enum import Enum
 from math import sqrt
+import os
 from typing import Any, Callable, Optional, Sequence, Tuple, Union
 import warnings
 
@@ -165,8 +166,17 @@ def core_attention(query: Array,
  else:
  attn_weights = jnp.einsum('bqhd,bkhd->bhqk', query, key)
 
+ # When a bias is present, the computation is performed as Softmax(attn_weights * scale + bias).
+ # In this case, the scale can not fused into the Softmax module.
+ if bias is not None:
+ attn_weights = attn_weights * scale_factor
+ fused_scale_factor = 1.
+ else:
+ # If no bias, the scale can be fused into Softmax module
+ fused_scale_factor = scale_factor
+
  attn_weights = Softmax(softmax_type=softmax_type,
- scale_factor=scale_factor,
+ scale_factor=fused_scale_factor,
  sharding_type=softmax_sharding_type)(attn_weights, mask, bias)
 
  if not deterministic and dropout_rate > 0.:
@@ -360,12 +370,13 @@ def kv_init(key, shape, dtype):
  q_seqlen = inputs_q.shape[0] if self.transpose_batch_sequence else inputs_q.shape[1]
  kv_seqlen = inputs_kv.shape[0] if self.transpose_batch_sequence else inputs_kv.shape[1]
  fused_attn_supported_seqlen = [128, 256, 384, 512]
+ enable_fused_attn = int(os.getenv("NVTE_FUSED_ATTN", "0"))
  use_fused_attn = not decode and not self.transpose_batch_sequence and self.fuse_qkv and \
  self.dropout_rate == 0 and canonicalize_dtype in [jnp.bfloat16, jnp.float16] and \
  q_seqlen in fused_attn_supported_seqlen and kv_seqlen in fused_attn_supported_seqlen \
- and is_fused_attn_kernel_available()
+ and is_fused_attn_kernel_available() and enable_fused_attn
 
- if not use_fused_attn:
+ if enable_fused_attn and not use_fused_attn:
  reason = ""
  if decode:
  reason += f"decode=False is required but got {decode}, "

transformer_engine/jax/sharding.py

@@ -386,7 +386,7 @@ def _get_dptp_sharding_meta(input_shapes: Tuple[Tuple[int, ...]],
 
  for input_shape, dp_dim, tp_dim in zip(input_shapes, input_dp_dims, input_tp_dims):
  in_axis = {}
- if dp_dim is not None:
+ if dp_dim is not None and input_shape is not None:
  in_axis[dp_dim] = dp_axis_name
  assert input_shape[dp_dim] % dp_size == 0, \
  f"The dimension of batch in input_shape should be a multiple of " \
@@ -398,7 +398,7 @@ def _get_dptp_sharding_meta(input_shapes: Tuple[Tuple[int, ...]],
  if tp_dim is not None and tp_dim >= dp_dim:
  tp_dim = tp_dim + 1
 
- if tp_dim is not None:
+ if tp_dim is not None and input_shape is not None:
  in_axis[tp_dim] = tp_axis_name
  assert input_shape[tp_dim] % tp_size == 0, \
  f"The dimension of tensor parallel in input_shape should be a multiple of " \