diff --git a/tests/jax/test_fused_attn.py b/tests/jax/test_fused_attn.py
index 811fe3ba6..8d8e73721 100644
--- a/tests/jax/test_fused_attn.py
+++ b/tests/jax/test_fused_attn.py
@@ -43,7 +43,7 @@
     CPStrategy,
     ReorderStrategy,
 )
-from transformer_engine.jax.cpp_extensions import FusedAttnHelper
+from transformer_engine.jax.cpp_extensions import FusedAttnHelper, fused_attn_fwd
 from transformer_engine_jax import (
     NVTE_Fused_Attn_Backend,
     get_cudnn_version,
@@ -59,6 +59,12 @@
 # CI test level
 _TEST_LEVEL = get_test_level()
 
+# AITER FP8 fused attention is inference-only and ships only for gfx950.
+rocm_fp8_attention_support = pytest.mark.skipif(
+    not is_hip_extension() or get_device_compute_capability(0) != 95,
+    reason="AITER FP8 ASM fused attention requires gfx950",
+)
+
 
 @pytest.fixture(autouse=True, scope="module")
 def init():
@@ -70,6 +76,191 @@ def init():
     yield
 
 
+@rocm_fp8_attention_support
+@pytest.mark.parametrize("head_dim", [128, 256])
+def test_aiter_fp8_fwd_backend_selection(head_dim):
+    """gfx950 fp8bf16 inference selects CK/AITER.
+
+    Training stays unsupported. NVTE_CK_USES_FWD_V3=0 also selects no backend
+    until a CK tile FP8 path exists.
+    """
+    kwargs = dict(
+        batch_size=1,
+        q_dtype=jnp.float8_e4m3fn,
+        kv_dtype=jnp.float8_e4m3fn,
+        qkv_layout=QKVLayout.BSHD_BSHD_BSHD,
+        attn_bias_type=AttnBiasType.NO_BIAS,
+        attn_mask_type=AttnMaskType.NO_MASK,
+        softmax_type=AttnSoftmaxType.VANILLA_SOFTMAX,
+        dropout_probability=0.0,
+        q_num_heads=8,
+        kv_num_heads=8,
+        q_max_seqlen=512,
+        kv_max_seqlen=512,
+        head_dim_qk=head_dim,
+        head_dim_v=head_dim,
+        window_size=(-1, -1),
+    )
+    inference_backend, _ = FusedAttnHelper(is_training=False, **kwargs).get_fused_attn_backend()
+    training_backend, _ = FusedAttnHelper(is_training=True, **kwargs).get_fused_attn_backend()
+    sliding_window_backend, _ = FusedAttnHelper(
+        is_training=False, **(kwargs | {"window_size": (64, 64)})
+    ).get_fused_attn_backend()
+    assert inference_backend in (
+        NVTE_Fused_Attn_Backend.NVTE_CK,
+        NVTE_Fused_Attn_Backend.NVTE_No_Backend,
+    )
+    assert training_backend == NVTE_Fused_Attn_Backend.NVTE_No_Backend
+    assert sliding_window_backend == NVTE_Fused_Attn_Backend.NVTE_No_Backend
+
+
+@rocm_fp8_attention_support
+@pytest.mark.parametrize("head_dim", [128, 256])
+def test_aiter_fp8_fwd_asm(head_dim):
+    """Small dense smoke: load the fp8bf16 ASM kernel and check vs a float32 JAX reference.
+
+    Broader HD128/HD256 accuracy (masks, THD padding, GQA, vs BF16 ASM) lives in
+    TestFusedAttnFP8ASM.
+    """
+    backend, _ = FusedAttnHelper(
+        is_training=False,
+        batch_size=1,
+        q_dtype=jnp.float8_e4m3fn,
+        kv_dtype=jnp.float8_e4m3fn,
+        qkv_layout=QKVLayout.BSHD_BSHD_BSHD,
+        attn_bias_type=AttnBiasType.NO_BIAS,
+        attn_mask_type=AttnMaskType.NO_MASK,
+        softmax_type=AttnSoftmaxType.VANILLA_SOFTMAX,
+        dropout_probability=0.0,
+        q_num_heads=8,
+        kv_num_heads=8,
+        q_max_seqlen=128,
+        kv_max_seqlen=128,
+        head_dim_qk=head_dim,
+        head_dim_v=head_dim,
+        window_size=(-1, -1),
+    ).get_fused_attn_backend()
+    if backend == NVTE_Fused_Attn_Backend.NVTE_No_Backend:
+        pytest.skip("AITER FP8 ASM fused attention backend is disabled")
+    assert backend == NVTE_Fused_Attn_Backend.NVTE_CK
+
+    shape = (1, 128, 8, head_dim)
+    key = jax.random.PRNGKey(7)
+    # Distinct inverse scales, so a missing or unit descale cannot match the reference.
+    q_scale_inv = jnp.asarray([0.5], dtype=jnp.float32)
+    k_scale_inv = jnp.asarray([1.5], dtype=jnp.float32)
+    v_scale_inv = jnp.asarray([0.75], dtype=jnp.float32)
+    q = (jax.random.normal(key, shape) * 0.2 / q_scale_inv).astype(jnp.float8_e4m3fn)
+    k = (jax.random.normal(jax.random.fold_in(key, 1), shape) * 0.2 / k_scale_inv).astype(
+        jnp.float8_e4m3fn
+    )
+    v = (jax.random.normal(jax.random.fold_in(key, 2), shape) * 0.2 / v_scale_inv).astype(
+        jnp.float8_e4m3fn
+    )
+    q_dq = q.astype(jnp.float32) * q_scale_inv
+    k_dq = k.astype(jnp.float32) * k_scale_inv
+    v_dq = v.astype(jnp.float32) * v_scale_inv
+    sequence_descriptor = SequenceDescriptor.from_seqlens(jnp.array([128], dtype=jnp.int32))
+    scale = head_dim**-0.5
+
+    output, _, _, _ = fused_attn_fwd(
+        (q, k, v),
+        None,
+        None,
+        sequence_descriptor,
+        None,
+        AttnBiasType.NO_BIAS,
+        AttnMaskType.NO_MASK,
+        AttnSoftmaxType.VANILLA_SOFTMAX,
+        QKVLayout.BSHD_BSHD_BSHD,
+        scale,
+        0.0,
+        False,
+        1,
+        q_scale_inv=q_scale_inv,
+        k_scale_inv=k_scale_inv,
+        v_scale_inv=v_scale_inv,
+    )
+    scores = jnp.einsum("bqhd,bkhd->bhqk", q_dq, k_dq) * scale
+    reference = jnp.einsum("bhqk,bkhd->bqhd", jax.nn.softmax(scores, axis=-1), v_dq)
+
+    assert output.dtype == jnp.bfloat16
+    assert float(jnp.max(jnp.abs(output.astype(jnp.float32) - reference))) < 0.01
+
+
+@rocm_fp8_attention_support
+@pytest.mark.parametrize("head_dim", [128, 256])
+def test_aiter_fp8_fwd_asm_jit_scales(head_dim):
+    """Inference jit must accept scale inverses as traced arguments.
+
+    Binding them with partial outside jit hides the custom_vjp nondiff check.
+    """
+    backend, _ = FusedAttnHelper(
+        is_training=False,
+        batch_size=1,
+        q_dtype=jnp.float8_e4m3fn,
+        kv_dtype=jnp.float8_e4m3fn,
+        qkv_layout=QKVLayout.BSHD_BSHD_BSHD,
+        attn_bias_type=AttnBiasType.NO_BIAS,
+        attn_mask_type=AttnMaskType.NO_MASK,
+        softmax_type=AttnSoftmaxType.VANILLA_SOFTMAX,
+        dropout_probability=0.0,
+        q_num_heads=8,
+        kv_num_heads=8,
+        q_max_seqlen=128,
+        kv_max_seqlen=128,
+        head_dim_qk=head_dim,
+        head_dim_v=head_dim,
+        window_size=(-1, -1),
+    ).get_fused_attn_backend()
+    if backend == NVTE_Fused_Attn_Backend.NVTE_No_Backend:
+        pytest.skip("AITER FP8 ASM fused attention backend is disabled")
+    assert backend == NVTE_Fused_Attn_Backend.NVTE_CK
+
+    shape = (1, 128, 8, head_dim)
+    key = jax.random.PRNGKey(7)
+    q_scale_inv = jnp.asarray([0.5], dtype=jnp.float32)
+    k_scale_inv = jnp.asarray([1.5], dtype=jnp.float32)
+    v_scale_inv = jnp.asarray([0.75], dtype=jnp.float32)
+    q = (jax.random.normal(key, shape) * 0.2 / q_scale_inv).astype(jnp.float8_e4m3fn)
+    k = (jax.random.normal(jax.random.fold_in(key, 1), shape) * 0.2 / k_scale_inv).astype(
+        jnp.float8_e4m3fn
+    )
+    v = (jax.random.normal(jax.random.fold_in(key, 2), shape) * 0.2 / v_scale_inv).astype(
+        jnp.float8_e4m3fn
+    )
+    sequence_descriptor = SequenceDescriptor.from_seqlens(jnp.array([128], dtype=jnp.int32))
+    scale = head_dim**-0.5
+
+    @jax.jit
+    def run(q, k, v, q_scale_inv, k_scale_inv, v_scale_inv):
+        return fused_attn(
+            (q, k, v),
+            None,
+            sequence_descriptor,
+            None,
+            attn_bias_type=AttnBiasType.NO_BIAS,
+            attn_mask_type=AttnMaskType.NO_MASK,
+            qkv_layout=QKVLayout.BSHD_BSHD_BSHD,
+            softmax_type=AttnSoftmaxType.VANILLA_SOFTMAX,
+            scaling_factor=scale,
+            dropout_probability=0.0,
+            is_training=False,
+            q_scale_inv=q_scale_inv,
+            k_scale_inv=k_scale_inv,
+            v_scale_inv=v_scale_inv,
+        )
+
+    output = run(q, k, v, q_scale_inv, k_scale_inv, v_scale_inv)
+    q_dq = q.astype(jnp.float32) * q_scale_inv
+    k_dq = k.astype(jnp.float32) * k_scale_inv
+    v_dq = v.astype(jnp.float32) * v_scale_inv
+    scores = jnp.einsum("bqhd,bkhd->bhqk", q_dq, k_dq) * scale
+    reference = jnp.einsum("bhqk,bkhd->bqhd", jax.nn.softmax(scores, axis=-1), v_dq)
+    assert output.dtype == jnp.bfloat16
+    assert float(jnp.max(jnp.abs(output.astype(jnp.float32) - reference))) < 0.01
+
+
 @partial(jax.jit, static_argnums=(6, 7, 8, 9, 11, 12, 13))
 def general_dot_product_attention(
     query: ArrayLike,
@@ -2678,3 +2869,265 @@ def test_smallseq(
         )
         runner.test_forward()
         runner.test_backward()
+
+def _pack_qkv_args(query, key, value, qkv_layout):
+    """Pack Q/K/V the same way as customcall_fused_dpa, without recasting the output dtype."""
+    match qkv_layout:
+        case QKVLayout.BS3HD | QKVLayout.T3HD:
+            query, key, value = map(partial(jnp.expand_dims, axis=-3), [query, key, value])
+            return (jnp.concatenate((query, key, value), axis=-3),)
+        case QKVLayout.BSHD_BS2HD | QKVLayout.THD_T2HD:
+            key, value = map(partial(jnp.expand_dims, axis=-3), [key, value])
+            return (query, jnp.concatenate((key, value), axis=-3))
+        case QKVLayout.BSHD_BSHD_BSHD | QKVLayout.THD_THD_THD:
+            return (query, key, value)
+        case _:
+            raise ValueError(f"Unsupported {qkv_layout=}")
+
+
+def _fused_attn_keep_output_dtype(
+    query, key, value, sequence_descriptor, **kwargs
+):
+    """Call fused_attn without recasting O back to the Q dtype (FP8 in, BF16 out)."""
+    return fused_attn(
+        _pack_qkv_args(query, key, value, kwargs["qkv_layout"]),
+        None,
+        sequence_descriptor,
+        None,
+        softmax_offset=None,
+        **kwargs,
+    )
+
+
+@rocm_fp8_attention_support
+@pytest.mark.parametrize(
+    "attn_mask_type",
+    [
+        pytest.param(AttnMaskType.NO_MASK, id="NO_MASK"),
+        pytest.param(AttnMaskType.CAUSAL_MASK, id="CAUSAL"),
+        pytest.param(AttnMaskType.PADDING_MASK, id="PADDING"),
+        pytest.param(AttnMaskType.PADDING_CAUSAL_MASK, id="PADDING_CAUSAL"),
+    ],
+)
+@pytest.mark.parametrize(
+    "b, s_q, s_kv, h_q, h_kv, d, qkv_layout",
+    [
+        pytest.param(
+            2,
+            512,
+            512,
+            8,
+            8,
+            128,
+            QKVLayout.BSHD_BSHD_BSHD,
+            id="2-512-8-128-BSHD-MHA",
+        ),
+        pytest.param(
+            2,
+            512,
+            512,
+            8,
+            4,
+            128,
+            QKVLayout.BSHD_BSHD_BSHD,
+            id="2-512-8-4-128-BSHD-GQA",
+        ),
+        pytest.param(
+            2,
+            512,
+            512,
+            8,
+            8,
+            256,
+            QKVLayout.BSHD_BSHD_BSHD,
+            id="2-512-8-256-BSHD-MHA",
+        ),
+        pytest.param(
+            2,
+            512,
+            512,
+            8,
+            4,
+            256,
+            QKVLayout.BSHD_BSHD_BSHD,
+            id="2-512-8-4-256-BSHD-GQA",
+        ),
+        pytest.param(
+            2,
+            512,
+            512,
+            8,
+            8,
+            128,
+            QKVLayout.THD_THD_THD,
+            id="2-512-8-128-THD-MHA",
+        ),
+        pytest.param(
+            2,
+            512,
+            512,
+            8,
+            4,
+            128,
+            QKVLayout.THD_THD_THD,
+            id="2-512-8-4-128-THD-GQA",
+        ),
+        pytest.param(
+            2,
+            512,
+            512,
+            8,
+            8,
+            256,
+            QKVLayout.THD_THD_THD,
+            id="2-512-8-256-THD-MHA",
+        ),
+        pytest.param(
+            2,
+            512,
+            512,
+            8,
+            4,
+            256,
+            QKVLayout.THD_THD_THD,
+            id="2-512-8-4-256-THD-GQA",
+        ),
+    ],
+)
+class TestFusedAttnFP8ASM:
+    """gfx950 fp8bf16 ASM fused-attn accuracy vs JAX softmax and vs BF16 ASM.
+
+    Coverage mirrors the BF16 HD128/HD256 runner tests: dense BSHD (no/causal)
+    and varlen THD (padding / padding-causal), MHA and GQA, hd128 and hd256.
+    Inference only — FP8 ASM has no training path.
+    """
+
+    @staticmethod
+    def test_forward(attn_mask_type, b, s_q, s_kv, h_q, h_kv, d, qkv_layout):
+        if qkv_layout.is_thd() and not attn_mask_type.is_padding():
+            pytest.skip("THD format requires padding masks.")
+        if attn_mask_type.is_padding() and not qkv_layout.is_thd():
+            pytest.skip("FP8 ASM requires THD layout for padding masks.")
+
+        runner = FusedAttnRunner(
+            b,
+            s_q,
+            s_kv,
+            h_q,
+            h_kv,
+            d,
+            d,
+            AttnBiasType.NO_BIAS,
+            attn_mask_type,
+            AttnSoftmaxType.VANILLA_SOFTMAX,
+            0.0,
+            jnp.bfloat16,
+            False,
+            qkv_layout,
+            None,
+            None,
+            SeqDescFormat.Seqlens,
+        )
+        runner._setup_inputs()
+
+        # Distinct inverse scales, so a missing or unit descale cannot match the reference.
+        q_scale_inv = jnp.asarray([0.5], dtype=jnp.float32)
+        k_scale_inv = jnp.asarray([1.5], dtype=jnp.float32)
+        v_scale_inv = jnp.asarray([0.75], dtype=jnp.float32)
+        q_fp8 = (runner.q.astype(jnp.float32) / q_scale_inv).astype(jnp.float8_e4m3fn)
+        k_fp8 = (runner.k.astype(jnp.float32) / k_scale_inv).astype(jnp.float8_e4m3fn)
+        v_fp8 = (runner.v.astype(jnp.float32) / v_scale_inv).astype(jnp.float8_e4m3fn)
+        q_dq = q_fp8.astype(jnp.float32) * q_scale_inv
+        k_dq = k_fp8.astype(jnp.float32) * k_scale_inv
+        v_dq = v_fp8.astype(jnp.float32) * v_scale_inv
+        q_bf16 = q_dq.astype(jnp.bfloat16)
+        k_bf16 = k_dq.astype(jnp.bfloat16)
+        v_bf16 = v_dq.astype(jnp.bfloat16)
+
+        fp8_backend, _ = FusedAttnHelper(
+            is_training=False,
+            batch_size=b,
+            q_dtype=jnp.float8_e4m3fn,
+            kv_dtype=jnp.float8_e4m3fn,
+            qkv_layout=qkv_layout,
+            attn_bias_type=AttnBiasType.NO_BIAS,
+            attn_mask_type=attn_mask_type,
+            softmax_type=AttnSoftmaxType.VANILLA_SOFTMAX,
+            dropout_probability=0.0,
+            q_num_heads=h_q,
+            kv_num_heads=h_kv,
+            q_max_seqlen=s_q,
+            kv_max_seqlen=s_kv,
+            head_dim_qk=d,
+            head_dim_v=d,
+            window_size=(-1, -1),
+        ).get_fused_attn_backend()
+        if fp8_backend == NVTE_Fused_Attn_Backend.NVTE_No_Backend:
+            pytest.skip("AITER FP8 ASM fused attention backend is disabled")
+        assert fp8_backend == NVTE_Fused_Attn_Backend.NVTE_CK, (
+            f"FP8 ASM was not selected: backend={fp8_backend}"
+        )
+
+        kwargs = dict(
+            attn_bias_type=AttnBiasType.NO_BIAS,
+            attn_mask_type=attn_mask_type,
+            softmax_type=AttnSoftmaxType.VANILLA_SOFTMAX,
+            scaling_factor=runner.scaling_factor,
+            dropout_probability=0.0,
+            is_training=False,
+            qkv_layout=qkv_layout,
+            max_segments_per_seq=runner._get_max_segments_per_sequence(),
+            window_size=None,
+        )
+
+        fp8_kwargs = dict(
+            kwargs,
+            q_scale_inv=q_scale_inv,
+            k_scale_inv=k_scale_inv,
+            v_scale_inv=v_scale_inv,
+        )
+        fp8_out = jit(
+            partial(_fused_attn_keep_output_dtype, **fp8_kwargs),
+            static_argnames=kwargs.keys(),
+        )(q_fp8, k_fp8, v_fp8, runner.sequence_desciptor)
+        bf16_out = jit(
+            partial(customcall_fused_dpa, **kwargs),
+            static_argnames=kwargs.keys(),
+        )(q_bf16, k_bf16, v_bf16, None, None, runner.sequence_desciptor, None)
+        reference_out = jax_dpa(
+            q_dq, k_dq, v_dq, None, None, runner.mask, None, **kwargs
+        )
+
+        assert fp8_out.dtype == jnp.bfloat16
+        assert bf16_out.dtype == jnp.bfloat16
+
+        fp8_valid, fp8_invalid, ref_valid, _ = _split_valid_and_invalid(
+            fp8_out, reference_out, runner.pad_q
+        )
+        bf16_valid, _, _, _ = _split_valid_and_invalid(
+            bf16_out, reference_out, runner.pad_q
+        )
+
+        assert_allclose(
+            fp8_invalid,
+            jnp.zeros_like(fp8_invalid),
+            dtype=jnp.bfloat16,
+            err_msg="FP8 ASM wrote non-zero values into padded tokens",
+        )
+        # JAX softmax reference on the same dequantized FP8 inputs.
+        assert_allclose(
+            fp8_valid,
+            ref_valid,
+            dtype=jnp.float8_e4m3fn,
+            err_msg="FP8 ASM vs JAX softmax reference",
+        )
+        # Same inputs through the BF16 ASM/CK path; residual should stay in e4m3 range.
+        assert_allclose(
+            fp8_valid,
+            bf16_valid,
+            dtype=jnp.float8_e4m3fn,
+            err_msg="FP8 ASM vs BF16 ASM on identical dequantized QKV",
+        )
+        print_debug_tensor_stats("fp8_asm", fp8_valid)
+        print_debug_tensor_stats("bf16_asm", bf16_valid)
+        print_debug_tensor_stats("jax_ref", ref_valid)
diff --git a/transformer_engine/common/ck_fused_attn/include/ck_fused_attn/ck_fused_attn.hpp b/transformer_engine/common/ck_fused_attn/include/ck_fused_attn/ck_fused_attn.hpp
index 7302b8891..cae8916b3 100644
--- a/transformer_engine/common/ck_fused_attn/include/ck_fused_attn/ck_fused_attn.hpp
+++ b/transformer_engine/common/ck_fused_attn/include/ck_fused_attn/ck_fused_attn.hpp
@@ -16,9 +16,10 @@ namespace ck_fused_attn{
 
 // input qkv dtypes
 enum class DType {
-  kFloat16    = 0,  /*!< 16-bit float (E5M10) */
-  kBFloat16   = 1,  /*!< 16-bit bfloat (E8M7) */
-  kNumTypes         /*!< Number of supported types */
+  kFloat16  = 0,  /*!< 16-bit float (E5M10) */
+  kBFloat16 = 1,  /*!< 16-bit bfloat (E8M7) */
+  kFloat8   = 2,  /*!< 8-bit float (E4M3) input */
+  kNumTypes       /*!< Number of supported types */
 };
 
 // keep sync with mask_enum in mask.hpp
@@ -99,6 +100,13 @@ struct CKAttnFwdArgs : CKAttnCommonArgs {
   // Output (writable)
   void* o_ptr = nullptr;
   void* lse_ptr = nullptr;
+  // For FP8 E4M3 input, the AITER fp8bf16 kernels require BF16.  
+  DType o_dtype = DType::kNumTypes;
+
+  // Per-tensor inverse scales for FP8 Q/K/V: one FP32 value per tensor.
+  const void* q_descale_ptr = nullptr;
+  const void* k_descale_ptr = nullptr;
+  const void* v_descale_ptr = nullptr;
 
   // V3 ASM kernel selection
   bool uses_fwd_v3 = false;
diff --git a/transformer_engine/common/ck_fused_attn/src/ck_fused_attn_fwd.cpp b/transformer_engine/common/ck_fused_attn/src/ck_fused_attn_fwd.cpp
index cac418504..6d21973f9 100644
--- a/transformer_engine/common/ck_fused_attn/src/ck_fused_attn_fwd.cpp
+++ b/transformer_engine/common/ck_fused_attn/src/ck_fused_attn_fwd.cpp
@@ -292,11 +292,14 @@ aiter::mha_fwd_args build_fwd_fmha_args(const CKAttnFwdArgs& args){
   fmha_args.is_group_mode   = args.is_group_mode();
   fmha_args.bias_type       = static_cast<int>(bias_type);
   fmha_args.has_lse         = args.lse_ptr!=nullptr;
-  fmha_args.qscale_type     = static_cast<int>(quant_scale_enum::no_scale);
+  fmha_args.qscale_type     = static_cast<int>(
+      args.dtype == DType::kFloat8
+          ? quant_scale_enum::pertensor
+          : quant_scale_enum::no_scale);
   fmha_args.has_sink        = args.has_sink;
-  fmha_args.q_descale_ptr    = nullptr;
-  fmha_args.k_descale_ptr    = nullptr;
-  fmha_args.v_descale_ptr    = nullptr;
+  fmha_args.q_descale_ptr    = args.q_descale_ptr;
+  fmha_args.k_descale_ptr    = args.k_descale_ptr;
+  fmha_args.v_descale_ptr    = args.v_descale_ptr;
   // sink_size is CK's StreamingLLM sink *prefix width* in key columns, which is a
   // different feature from the learnable softmax offset NVTE asks for here. The
   // offset only needs has_sink + sink_ptr (CK folds it into the softmax
@@ -313,9 +316,10 @@ aiter::mha_fwd_args build_fwd_fmha_args(const CKAttnFwdArgs& args){
   return fmha_args;
 }
 
-//Split-KV and Sink dispatchers do not honor v3_api_check
-//so we cannot rely on the AITER backend to tell us whether v3 is available.
-//Currently this method is not used so just opt it out.
+// Split-KV and sink dispatchers do not honor v3_api_check, so a naive
+// mha_fwd(..., v3_api_check=true) can still launch a real kernel when those
+// paths are compiled in. Unused today; the FP8 probe in ck_attn_fwd clears
+// sink/split-KV fields before using the same dry-run.
 #if !(FA_WITH_SINK || FA_WITH_NATIVE_SPLITKV)
 // Probe whether AITER's v3 (asm) forward path will run for this config, without
 // launching any kernel. Builds the same args as ck_attn_fwd and relies on AITER's
@@ -340,6 +344,29 @@ hipError_t ck_attn_fwd(const CKAttnFwdArgs& args, hipStream_t stream){
 
   aiter::mha_fwd_args fmha_args = build_fwd_fmha_args(args);
 
+  // This PR wires AITER gfx950 FP8 FMHA through the v3 ASM kernels only.
+  // Upstream CK has an FP8 tile path, but it is not enabled here, so a failed
+  // ASM probe must not fall through to an unvalidated CK launch. Enabling CK
+  // tile FP8 is a separate change.
+  //
+  // QoLA builds with FA_WITH_SINK and FA_WITH_NATIVE_SPLITKV, and those
+  // dispatchers run before fmha_fwd_v3 without honoring v3_api_check. Clear
+  // sink/split-KV fields so the probe reaches the v3 dry-run and cannot launch
+  // a real kernel (same reason ck_attn_fwd_uses_v3 is compiled out).
+  if(args.dtype == DType::kFloat8){
+    aiter::mha_fwd_args probe_args = fmha_args;
+    probe_args.v3_api_check = true;
+    probe_args.num_splits = 0;
+    probe_args.splitkv_workspace_ptr = nullptr;
+    probe_args.has_sink = false;
+    probe_args.sink_ptr = nullptr;
+    ck_tile::stream_config probe_config{nullptr, false, false};
+    if(QOLA_NS(mha_fwd)(probe_args, probe_config) != 1){
+      throw std::runtime_error(
+        "ck_fused_attn fwd: no AITER fp8bf16 ASM kernel for this configuration.");
+    }
+  }
+
   if(const char* env_p = std::getenv("NVTE_CK_RUNTIME_MAX_SEQLEN")){
     if(args.is_group_mode() && std::string(env_p) == "1"){
       if(log_file){
diff --git a/transformer_engine/common/ck_fused_attn/src/ck_fused_attn_utils.cpp b/transformer_engine/common/ck_fused_attn/src/ck_fused_attn_utils.cpp
index 936639f27..23d4390b0 100644
--- a/transformer_engine/common/ck_fused_attn/src/ck_fused_attn_utils.cpp
+++ b/transformer_engine/common/ck_fused_attn/src/ck_fused_attn_utils.cpp
@@ -106,14 +106,18 @@ std::ostream* get_ck_log_stream() {
   return log_stream;
 }
 
+// Return AITER's data_type kernel-selection tag. The tag may encode a
+// combination of input and output types rather than the input dtype alone.
 std::string get_data_type_str(DType dtype){
   std::string data_type_str;
   if(dtype==DType::kFloat16){
     data_type_str = "fp16";
   }else if(dtype==DType::kBFloat16){
     data_type_str = "bf16";
+  }else if(dtype==DType::kFloat8){
+    data_type_str = "fp8bf16";
+    // "fp8bf16" selects FP8 input with BF16 output; "fp8" selects FP8 output.
   }else{
-    //TODO: better error out system
     throw std::runtime_error("Invalid dtype in ck_fused_attn.");
   }
   return data_type_str;
diff --git a/transformer_engine/common/fused_attn_rocm/fused_attn.cpp b/transformer_engine/common/fused_attn_rocm/fused_attn.cpp
index b3776343f..9ff25acdc 100644
--- a/transformer_engine/common/fused_attn_rocm/fused_attn.cpp
+++ b/transformer_engine/common/fused_attn_rocm/fused_attn.cpp
@@ -299,6 +299,7 @@ NVTE_Fused_Attn_Backend nvte_get_fused_attn_backend(
 
   // first check whether ck can be used, then check aotriton
   if(nvte_fused_attn_ck && fused_attn_rocm::is_ck_backend_supported(
+        is_training,
         q_dtype,
         kv_dtype,
         qkv_layout,
diff --git a/transformer_engine/common/fused_attn_rocm/fused_attn_ck.cpp b/transformer_engine/common/fused_attn_rocm/fused_attn_ck.cpp
index c90e60fa4..a8864605f 100644
--- a/transformer_engine/common/fused_attn_rocm/fused_attn_ck.cpp
+++ b/transformer_engine/common/fused_attn_rocm/fused_attn_ck.cpp
@@ -21,6 +21,7 @@ namespace fused_attn_rocm {
 // check the fused attn config to see whether it's ck backend supported
 // single filtering followed by joint filtering
 bool is_ck_backend_supported(
+  bool is_training,
   NVTEDType q_dtype,
   NVTEDType kv_dtype,
   NVTE_QKV_Layout qkv_layout,
@@ -49,11 +50,13 @@ bool is_ck_backend_supported(
     return false;
   }
 
-  // filter based on data type
-  // Q and KV must have the same data type, in fp16 or bf16
-  if((q_dtype!=kv_dtype) || !((q_dtype==NVTEDType::kNVTEFloat16) || (q_dtype == NVTEDType::kNVTEBFloat16))){
+  const bool is_fp8 = q_dtype == NVTEDType::kNVTEFloat8E4M3;
+  const bool is_16bit =
+      q_dtype == NVTEDType::kNVTEFloat16 || q_dtype == NVTEDType::kNVTEBFloat16;
+  //  q and kv have to have same dtype: FP16/BF16  or FP8 E4M3 
+  if(q_dtype != kv_dtype || (!is_16bit && !is_fp8)){
     if(nvte_log_ck_config){
-      std::cout<<"q, k, v data type has to be fp16 or bf16"<<std::endl;
+      std::cout<<"q, k, v data type has to be fp16, bf16, or fp8 e4m3"<<std::endl;
     }
     return false;
   }
@@ -136,6 +139,33 @@ bool is_ck_backend_supported(
     }
     return false;
   }
+  if(is_fp8){
+    // FP8 is wired to AITER v3 ASM only. CK tile FP8 is not enabled, so turning
+    // ASM off must report unsupported and let the caller select NVTE_No_Backend.
+    // A later PR that enables CK FP8 can fall back to that path here instead.
+    if(getenv<int>("NVTE_CK_USES_FWD_V3", 1) == 0){
+      if(nvte_log_ck_config){
+        std::cout<<"AITER FP8 ASM fwd is disabled by NVTE_CK_USES_FWD_V3=0"<<std::endl;
+      }
+      return false;
+    }
+    const bool fp8_hdim_supported =
+        (head_dim_qk == 128 && head_dim_v == 128) ||
+        (head_dim_qk == 256 && head_dim_v == 256);
+    const bool fp8_window_supported =
+        is_causal ? (window_size_left == -1 && window_size_right == 0)
+                  : (window_size_left == -1 && window_size_right == -1);
+    if(cuda::sm_arch() != 95 || is_training || !fp8_hdim_supported ||
+       !fp8_window_supported || dropout != 0.f ||
+       bias_type != NVTE_Bias_Type::NVTE_NO_BIAS ||
+       (is_padding && !is_ragged)){
+      if(nvte_log_ck_config){
+        std::cout<<"AITER FP8 ASM fwd requires gfx950 inference, hd128/hd256, "
+                   "no bias/dropout/sliding-window, and THD for padding masks"<<std::endl;
+      }
+      return false;
+    }
+  }
   return true;
 #else
   return false;
@@ -176,6 +206,7 @@ ck_fused_attn::DType nvte_to_ck_dtype(DType t_dtype){
 #define CAST_TYPE(aname, dtname) if (t_dtype == DType::aname) return ck_fused_attn::DType::dtname
   CAST_TYPE(kFloat16, kFloat16);
   CAST_TYPE(kBFloat16, kBFloat16);
+  CAST_TYPE(kFloat8E4M3, kFloat8);
   return ck_fused_attn::DType::kNumTypes;
 #undef CAST_TYPE
 }
@@ -490,7 +521,8 @@ void fused_attn_ck_fwd_impl(
   void* devPtrDropoutSeed, void* devPtrDropoutOffset,
   void* devPtrCuSeqlensQ, void* devPtrCuSeqlensKV,
   void* devPtrSeqOffsetsQ, void* devPtrSeqOffsetsKV,
-  DType dtype,
+  void* devPtrDescaleQ, void* devPtrDescaleK, void* devPtrDescaleV,
+  DType dtype, DType output_dtype,
   void *workspace, 
   size_t *workspace_size,
   cudaStream_t stream){
@@ -520,11 +552,18 @@ void fused_attn_ck_fwd_impl(
   size_t q_storage_bytes = max_tokens_q*h*d_qk*nvte_dtype_size(dtype);
   size_t k_storage_bytes = max_tokens_kv*hg*d_qk*nvte_dtype_size(dtype);
   size_t v_storage_bytes = max_tokens_kv*hg*d_v*nvte_dtype_size(dtype);
-  size_t o_storage_bytes = max_tokens_q*h*d_v*nvte_dtype_size(dtype);
+  size_t o_storage_bytes = max_tokens_q*h*d_v*nvte_dtype_size(output_dtype);
 
   // Reserve workspace chunks. Same allocation sequence runs in sizing mode
   // (planner returns nullptr, accumulates total) and execution mode.
   WorkspacePlanner planner(workspace);
+  // A real FP8 launch must carry the inverse scale with each input. Workspace
+  // sizing passes null tensors, so that pass does not require scales.
+  if(dtype == DType::kFloat8E4M3 && !planner.is_sizing()){
+    NVTE_CHECK(devPtrDescaleQ != nullptr && devPtrDescaleK != nullptr &&
+                   devPtrDescaleV != nullptr,
+               "AITER FP8 ASM fused attention requires one FP32 scale_inv for each of Q, K, and V.");
+  }
 
   if(ck_small_seq_env_enabled) {
     if(cuda::sm_arch() == 94 || cuda::sm_arch() == 95) {
@@ -589,6 +628,10 @@ void fused_attn_ck_fwd_impl(
   // Common fields filled here; mode-specific fields are overwritten below.
   ck_fused_attn::CKAttnFwdArgs ck_args;
   ck_args.dtype = nvte_to_ck_dtype(dtype);
+  ck_args.o_dtype = nvte_to_ck_dtype(output_dtype);
+  ck_args.q_descale_ptr = devPtrDescaleQ;
+  ck_args.k_descale_ptr = devPtrDescaleK;
+  ck_args.v_descale_ptr = devPtrDescaleV;
   ck_args.b = b; ck_args.h = h; ck_args.hg = hg;
   ck_args.s_q = s_q; ck_args.s_kv = s_kv; ck_args.d_qk = d_qk; ck_args.d_v = d_v;
   ck_args.is_training = is_training;
@@ -611,7 +654,10 @@ void fused_attn_ck_fwd_impl(
     ck_args.cu_seqlen_q_ptr = devPtrCuSeqlensQ;
   }
 
-  ck_args.num_splits = (ck_args.uses_fwd_v3 && !has_sink) ? ck_attn_fwd_num_splits(ck_args) : 0;
+  ck_args.num_splits =
+      (ck_args.uses_fwd_v3 && !has_sink && dtype != DType::kFloat8E4M3)
+          ? ck_attn_fwd_num_splits(ck_args)
+          : 0;
   if (ck_args.num_splits > 0)
   {
     size_t splitkv_workspace_bytes = ck_attn_fwd_workspace_size(ck_args);
@@ -1352,11 +1398,41 @@ void fused_attn_ck_fwd(
 
 #ifdef USE_FUSED_ATTN_CK
   const DType QKV_type = input_Q->data.dtype;
+  const DType O_type = output_O->data.dtype;
+  const bool is_fp8 = QKV_type == DType::kFloat8E4M3;
+
+  NVTE_CHECK(input_K->data.dtype == QKV_type && input_V->data.dtype == QKV_type,
+             "ROCm fused attention requires Q, K, and V to have the same dtype.");
+  if(is_fp8){
+    NVTE_CHECK(!is_training, "AITER FP8 ASM fused attention is forward-inference only.");
+    // FP8 E4M3 input is accepted only with BF16 output for AITER fp8bf16 kernels.
+    NVTE_CHECK(O_type == DType::kBFloat16,
+               "AITER fp8bf16 fused attention requires BF16 output.");
+    NVTE_CHECK(input_Q->scaling_mode == NVTE_DELAYED_TENSOR_SCALING &&
+                   input_K->scaling_mode == NVTE_DELAYED_TENSOR_SCALING &&
+                   input_V->scaling_mode == NVTE_DELAYED_TENSOR_SCALING,
+               "AITER FP8 ASM fused attention supports delayed tensor scaling only.");
+    auto valid_descale = [](const Tensor* tensor) {
+      return tensor->scale_inv.dptr != nullptr &&
+             tensor->scale_inv.dtype == DType::kFloat32 &&
+             tensor->scale_inv.numel() == 1;
+    };
+    const bool all_descales_present =
+        valid_descale(input_Q) && valid_descale(input_K) && valid_descale(input_V);
+    // Workspace sizing uses null tensors. A real launch must carry the scales.
+    if(input_Q->data.dptr != nullptr){
+      NVTE_CHECK(all_descales_present,
+                 "AITER FP8 ASM fused attention requires one FP32 scale_inv for each of Q, K, and V.");
+    }
+  }
 
   void *devPtrQ = input_Q->data.dptr;
   void *devPtrK = input_K->data.dptr;
   void *devPtrV = input_V->data.dptr;
   void *devPtrO = output_O->data.dptr;
+  void *devPtrDescaleQ = is_fp8 ? input_Q->scale_inv.dptr : nullptr;
+  void *devPtrDescaleK = is_fp8 ? input_K->scale_inv.dptr : nullptr;
+  void *devPtrDescaleV = is_fp8 ? input_V->scale_inv.dptr : nullptr;
   void *devPtrS = nullptr;
   void *devPtrBias = nullptr;
   size_t bias_b = 0;
@@ -1444,7 +1520,8 @@ void fused_attn_ck_fwd(
     reinterpret_cast<void *>(reinterpret_cast<uint64_t *>(rng_state->data.dptr) + 1),
     devPtrCuSeqlensQ, devPtrCuSeqlensKV,
     devPtrSeqOffsetsQ, devPtrSeqOffsetsKV,
-    QKV_type,
+    devPtrDescaleQ, devPtrDescaleK, devPtrDescaleV,
+    QKV_type, O_type,
     workspace->data.dptr,
     &workspace_size,
     stream);
diff --git a/transformer_engine/common/fused_attn_rocm/fused_attn_ck.h b/transformer_engine/common/fused_attn_rocm/fused_attn_ck.h
index d2b565c82..be5b5ea27 100644
--- a/transformer_engine/common/fused_attn_rocm/fused_attn_ck.h
+++ b/transformer_engine/common/fused_attn_rocm/fused_attn_ck.h
@@ -17,6 +17,7 @@ namespace transformer_engine {
 namespace fused_attn_rocm {
 // check the fused attn config to see whether it's ck backend supported
 bool is_ck_backend_supported(
+  bool is_training,
   NVTEDType q_dtype,
   NVTEDType kv_dtype,
   NVTE_QKV_Layout qkv_layout,
diff --git a/transformer_engine/jax/attention.py b/transformer_engine/jax/attention.py
index 7b2e41129..6553a7801 100644
--- a/transformer_engine/jax/attention.py
+++ b/transformer_engine/jax/attention.py
@@ -1158,6 +1158,9 @@ def _legacy_fused_attn(
         qkv,
         bias,
         softmax_offset,
+        None,
+        None,
+        None,
         SequenceDescriptor.from_seqlens((q_seq_lens, kv_seq_lens)),
         seed,
         attn_bias_type=attn_bias_type,
@@ -1237,6 +1240,9 @@ def fused_attn_thd(
         qkv,
         bias,
         softmax_offset,
+        None,
+        None,
+        None,
         SequenceDescriptor.from_seqlens_and_offsets(
             (q_seq_lens, kv_seq_lens), (q_seq_offsets, kv_seq_offsets)
         ),
@@ -1261,12 +1267,18 @@ def fused_attn_thd(
 
 @partial(
     jax.custom_vjp,
-    nondiff_argnums=(5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19),
+    # Scale inverses sit with the other primal inputs so jit can trace them.
+    # Their cotangents are None: this path does not differentiate the scales.
+    # Indexes are positions in this signature: 0-7 are differentiable, 8-22 are not.
+    nondiff_argnums=(8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22),
 )
 def _fused_attn(
     qkv: Tuple[jnp.ndarray, ...],
     bias: Optional[jnp.ndarray],
     softmax_offset: Optional[jnp.ndarray],
+    q_scale_inv: Optional[jnp.ndarray],
+    k_scale_inv: Optional[jnp.ndarray],
+    v_scale_inv: Optional[jnp.ndarray],
     sequence_descriptor: SequenceDescriptor,
     seed: Optional[jnp.ndarray],
     attn_bias_type: AttnBiasType,
@@ -1289,6 +1301,9 @@ def _fused_attn(
         qkv,
         bias,
         softmax_offset,
+        q_scale_inv,
+        k_scale_inv,
+        v_scale_inv,
         sequence_descriptor,
         seed,
         attn_bias_type,
@@ -1303,9 +1318,9 @@ def _fused_attn(
         context_parallel_strategy,
         context_parallel_causal_load_balanced,
         context_parallel_axis,
-        context_checkpoint_name=context_checkpoint_name,
-        stripe_size=stripe_size,
-        return_max_logit=return_max_logit,
+        context_checkpoint_name,
+        stripe_size,
+        return_max_logit,
     )
     return output
 
@@ -1314,6 +1329,9 @@ def _fused_attn_fwd_rule(
     qkv,
     bias,
     softmax_offset,
+    q_scale_inv,
+    k_scale_inv,
+    v_scale_inv,
     sequence_descriptor,
     seed,
     attn_bias_type,
@@ -1351,6 +1369,9 @@ def _fused_attn_fwd_rule(
         context_parallel_causal_load_balanced=context_parallel_causal_load_balanced,
         context_parallel_axis=context_parallel_axis,
         stripe_size=stripe_size,
+        q_scale_inv=q_scale_inv,
+        k_scale_inv=k_scale_inv,
+        v_scale_inv=v_scale_inv,
         return_max_logit=return_max_logit,
     )
     output = checkpoint_name(output, context_checkpoint_name)
@@ -1427,12 +1448,18 @@ def _fused_attn_bwd_rule(
         grad_bias = None
     if softmax_type != AttnSoftmaxType.LEARNABLE_SOFTMAX:
         grad_softmax_offset = None
+    # Cotangents follow the differentiable inputs, in signature order:
+    # qkv, bias, softmax_offset, q_scale_inv, k_scale_inv, v_scale_inv,
+    # sequence_descriptor, seed.
     return (
         grad_qkv,
         grad_bias,
         grad_softmax_offset,
         None,
         None,
+        None,
+        None,
+        None,
     )
 
 
@@ -1516,6 +1543,9 @@ def fused_attn(
     score_mod_bprop: Optional[Callable] = None,
     score_mod_tensors: Optional[Mapping[str, Any]] = None,
     score_mod_bprop_tensors: Optional[Mapping[str, Any]] = None,
+    q_scale_inv: Optional[jnp.ndarray] = None,
+    k_scale_inv: Optional[jnp.ndarray] = None,
+    v_scale_inv: Optional[jnp.ndarray] = None,
     return_max_logit: bool = False,
 ):
     """
@@ -1705,6 +1735,9 @@ def fused_attn(
         qkv,
         bias,
         softmax_offset,
+        q_scale_inv,
+        k_scale_inv,
+        v_scale_inv,
         sequence_descriptor,
         seed,
         attn_bias_type=attn_bias_type,
diff --git a/transformer_engine/jax/cpp_extensions/attention.py b/transformer_engine/jax/cpp_extensions/attention.py
index 3d4590e77..bc93956a2 100644
--- a/transformer_engine/jax/cpp_extensions/attention.py
+++ b/transformer_engine/jax/cpp_extensions/attention.py
@@ -101,6 +101,18 @@ def setup_logging():
         AttentionLogging._is_logging_setup = True
 
 
+def _scale_inv_operand(scale_inv):
+    """One FP32 inverse scale, or an empty buffer when the caller has no scale."""
+    if scale_inv is None:
+        return jnp.zeros((0,), dtype=jnp.float32)
+    scale_inv = jnp.asarray(scale_inv, dtype=jnp.float32)
+    if scale_inv.size != 1:
+        raise ValueError(
+            f"FP8 fused attention scale_inv must contain one value, got shape {scale_inv.shape}"
+        )
+    return scale_inv.reshape((1,))
+
+
 @partial(
     jax.tree_util.register_dataclass,
     data_fields=[],
@@ -463,7 +475,7 @@ class FusedAttnFwdPrimitive(BasePrimitive):
 
     name = "te_fused_attn_forward_ffi"
     multiple_results = True
-    impl_static_args = (14,)
+    impl_static_args = (17,)
     inner_primitive = None
     outer_primitive = None
 
@@ -483,6 +495,9 @@ def abstract(
         _kv_segment_ids,
         _q_segment_pos,
         _kv_segment_pos,
+        q_scale_inv_aval,
+        k_scale_inv_aval,
+        v_scale_inv_aval,
         *,
         config: _FusedAttnConfig,
     ):
@@ -496,6 +511,20 @@ def abstract(
         assert (
             q_dtype == k_dtype == v_dtype == bias_dtype
         ), f"q_dtype={q_dtype}, k_dtype={k_dtype}, v_dtype={v_dtype}, bias_dtype={bias_dtype}"
+        fp8_input = q_dtype in {
+            dtypes.canonicalize_dtype(jnp.float8_e4m3fn),
+            dtypes.canonicalize_dtype(jnp.float8_e4m3fnuz),
+        }
+        expected_scale_shape = (1,) if fp8_input else (0,)
+        for name, aval in (
+            ("q_scale_inv", q_scale_inv_aval),
+            ("k_scale_inv", k_scale_inv_aval),
+            ("v_scale_inv", v_scale_inv_aval),
+        ):
+            assert aval.dtype == jnp.float32, f"{name} must be float32, got {aval.dtype}"
+            assert aval.shape == expected_scale_shape, (
+                f"{name} shape must be {expected_scale_shape} for this dtype, got {aval.shape}"
+            )
         assert q_seqlen_or_cu_seqlen_aval.dtype == kv_seqlen_or_cu_seqlen_aval.dtype, (
             f"q_seqlen_or_cu_seqlen_aval={q_seqlen_or_cu_seqlen_aval},"
             f" kv_seqlen_or_cu_seqlen_aval={kv_seqlen_or_cu_seqlen_aval}"
@@ -512,7 +541,13 @@ def abstract(
         ) = FusedAttnHelper.parse_qkv_aval(q_aval, k_aval, v_aval, config.qkv_layout)
 
         output_shape = (*batch_shape, q_max_seqlen, attn_heads, v_head_dim)
-        out_aval = q_aval.update(shape=output_shape, dtype=q_dtype)
+        # ROCm FP8 E4M3 fused attention returns BF16.
+        output_dtype = (
+            dtypes.canonicalize_dtype(jnp.bfloat16)
+            if is_hip_extension() and fp8_input
+            else q_dtype
+        )
+        out_aval = q_aval.update(shape=output_shape, dtype=output_dtype)
 
         # backend determines the softmax buffer shape/dtype
         input_batch = reduce(operator.mul, batch_shape)
@@ -682,6 +717,9 @@ def lowering(
         _kv_segment_ids,
         _q_segment_pos,
         _kv_segment_pos,
+        q_scale_inv,
+        k_scale_inv,
+        v_scale_inv,
         *,
         config: _FusedAttnConfig,
     ):
@@ -725,7 +763,10 @@ def lowering(
             _q_segment_ids,
             _kv_segment_ids,
             _q_segment_pos,
-            _kv_segment_pos,  # ffi_lowering needs number of parameters meets primitive.lowering
+            _kv_segment_pos,
+            q_scale_inv,
+            k_scale_inv,
+            v_scale_inv,  # ffi_lowering needs number of parameters meets primitive.lowering
             input_batch=input_batch,
             bias_batch=bias_batch,
             q_max_seqlen=q_max_seqlen,
@@ -766,6 +807,9 @@ def impl(
         _kv_segment_ids,
         _q_segment_pos,
         _kv_segment_pos,
+        q_scale_inv,
+        k_scale_inv,
+        v_scale_inv,
         config: _FusedAttnConfig,
     ):
         assert (
@@ -860,6 +904,9 @@ def convert_to_2d(offsets, batch, max_seqlen):
             _kv_segment_ids,
             _q_segment_pos,
             _kv_segment_pos,
+            q_scale_inv,
+            k_scale_inv,
+            v_scale_inv,
             config=config,
         )
         # Reduce cuDNN's raw Max tensor to TE's public per-head [H] max_logit.
@@ -1028,8 +1075,9 @@ def partition(config, mesh, arg_infos, result_infos):
         )
         arg_shardings = [arg_i.sharding for arg_i in arg_infos]
         arg_shardings[5] = seed_sharding
-        arg_shardings[-1] = arg_shardings[-3]
-        arg_shardings[-2] = arg_shardings[-4]
+        # segment_pos follows segment_ids. scale_inv is appended after those four operands.
+        arg_shardings[13] = arg_shardings[11]
+        arg_shardings[12] = arg_shardings[10]
         arg_shardings = tuple(arg_shardings)
         out_shardings = (out_sharding, softmax_aux_sharding, rng_state_sharding, max_logit_sharding)
         max_logit_reduce_axes = (
@@ -2163,6 +2211,9 @@ def impl(
             _kv_segment_ids,
             _q_segment_pos,
             _kv_segment_pos,
+            q_scale_inv,
+            k_scale_inv,
+            v_scale_inv,
         ):
             cp_size = get_mesh_axis_size(config.cp_axis, mesh)
             cp_rank = get_mesh_axis_rank(config.cp_axis, mesh)
@@ -2210,6 +2261,9 @@ def _cross_attn(idx, q, k, v, bias, softmax_offset, q_seqlen, kv_seqlen, seed):
                         _kv_segment_ids,
                         _q_segment_pos,
                         _kv_segment_pos,
+                        q_scale_inv,
+                        k_scale_inv,
+                        v_scale_inv,
                         config=helper.get_step_config(),
                     )
                     results.append((output, softmax_aux, rng_state, max_logit))
@@ -2470,6 +2524,9 @@ def impl(
             _kv_segment_ids,
             _q_segment_pos,
             _kv_segment_pos,
+            q_scale_inv,
+            k_scale_inv,
+            v_scale_inv,
         ):  # pylint: disable=unused-argument
             cp_size = get_mesh_axis_size(config.cp_axis, mesh)
             cp_rank = get_mesh_axis_rank(config.cp_axis, mesh)
@@ -2532,6 +2589,9 @@ def _cross_attn(
                     jnp.zeros(0),
                     jnp.zeros(0),
                     jnp.zeros(0),
+                    q_scale_inv,
+                    k_scale_inv,
+                    v_scale_inv,
                     config=helper.get_step_config_for_striped(
                         max_seqlen=kv_max_seqlen, cp_size=cp_size
                     ),
@@ -2909,9 +2969,9 @@ def partition(config, mesh, arg_infos, result_infos):
         )
         arg_shardings = [arg_i.sharding for arg_i in arg_infos]
         arg_shardings[5] = seed_sharding
-        # Ensure segment_pos gets same sharding as ID.
-        arg_shardings[-1] = arg_shardings[-3]
-        arg_shardings[-2] = arg_shardings[-4]
+        # segment_pos follows segment_ids. scale_inv is appended after those four operands.
+        arg_shardings[13] = arg_shardings[11]
+        arg_shardings[12] = arg_shardings[10]
         arg_shardings = tuple(arg_shardings)
         out_shardings = (out_sharding, softmax_aux_sharding, rng_state_sharding, max_logit_sharding)
         max_logit_reduce_axes = (
@@ -2935,6 +2995,9 @@ def ring_attn_fwd_impl(
             _kv_segment_ids,
             _q_segment_pos,
             _kv_segment_pos,
+            q_scale_inv,
+            k_scale_inv,
+            v_scale_inv,
         ):
             _not_used = jnp.zeros(0, dtype=v.dtype)
 
@@ -2983,6 +3046,9 @@ def mask_compute(attn_mask_type):
                             _kv_segment_ids,
                             _q_segment_pos,
                             _kv_segment_pos,
+                            q_scale_inv,
+                            k_scale_inv,
+                            v_scale_inv,
                             config=helper.get_step_config(attn_mask_type),
                         )
                     )
@@ -3011,6 +3077,9 @@ def half_kv_no_mask_compute():
                             _kv_segment_ids,
                             _q_segment_pos,
                             _kv_segment_pos,
+                            q_scale_inv,
+                            k_scale_inv,
+                            v_scale_inv,
                             config=helper.get_step_config(AttnMaskType.NO_MASK),
                         )
                     )
@@ -3036,6 +3105,9 @@ def half_q_no_mask_compute():
                             _kv_segment_ids,
                             _q_segment_pos,
                             _kv_segment_pos,
+                            q_scale_inv,
+                            k_scale_inv,
+                            v_scale_inv,
                             config=helper.get_step_config(AttnMaskType.NO_MASK),
                         )
                     )
@@ -3437,9 +3509,9 @@ def partition(config, mesh, arg_infos, result_infos):
         )
         arg_shardings = [arg_i.sharding for arg_i in arg_infos]
         arg_shardings[5] = seed_sharding
-        # Ensure segment_pos gets same sharding as ID.
-        arg_shardings[-1] = arg_shardings[-3]
-        arg_shardings[-2] = arg_shardings[-4]
+        # segment_pos follows segment_ids. scale_inv is appended after those four operands.
+        arg_shardings[13] = arg_shardings[11]
+        arg_shardings[12] = arg_shardings[10]
         arg_shardings = tuple(arg_shardings)
         out_shardings = (out_sharding, softmax_aux_sharding, rng_state_sharding, max_logit_sharding)
         max_logit_reduce_axes = (
@@ -3463,6 +3535,9 @@ def fwd_impl(
             kv_segment_ids,
             q_segment_pos,
             kv_segment_pos,
+            q_scale_inv,
+            k_scale_inv,
+            v_scale_inv,
         ):
             if q_segment_ids.size == 0 or kv_segment_ids.size == 0:
                 raise ValueError("THD + ring attn only supports passing seqment_ids/pos")
@@ -3516,6 +3591,9 @@ def compute(config):
                         kv_segment_ids,
                         q_segment_pos,
                         kv_segment_pos,
+                        q_scale_inv,
+                        k_scale_inv,
+                        v_scale_inv,
                         config=config,
                     )
 
@@ -3778,6 +3856,9 @@ def fused_attn_fwd(
     context_parallel_causal_load_balanced: bool = False,
     context_parallel_axis: str = "",
     stripe_size: int | None = None,
+    q_scale_inv: Optional[jnp.ndarray] = None,
+    k_scale_inv: Optional[jnp.ndarray] = None,
+    v_scale_inv: Optional[jnp.ndarray] = None,
     return_max_logit: bool = False,
 ) -> jnp.ndarray:
     """
@@ -3912,12 +3993,19 @@ def fused_attn_fwd(
                 primitive = FusedRingAttnFwdPrimitive.outer_primitive
 
     seq_desc_flatten, _ = jax.tree.flatten(sequence_descriptor)
+    # FP16/BF16 pass empty buffers. FP8 passes one FP32 inverse scale per tensor.
+    q_scale_inv = _scale_inv_operand(q_scale_inv)
+    k_scale_inv = _scale_inv_operand(k_scale_inv)
+    v_scale_inv = _scale_inv_operand(v_scale_inv)
     output, softmax_aux, rng_state, max_logit = primitive.bind(
         *qkv_for_primitive,
         bias,
         softmax_offset,
         seed,
         *seq_desc_flatten,
+        q_scale_inv,
+        k_scale_inv,
+        v_scale_inv,
         config=fused_config,
     )
     rng_state = with_sharding_constraint(rng_state, PartitionSpec(get_all_mesh_axes(), None))
diff --git a/transformer_engine/jax/csrc/extensions/attention.cpp b/transformer_engine/jax/csrc/extensions/attention.cpp
index 55bade457..cec175a04 100644
--- a/transformer_engine/jax/csrc/extensions/attention.cpp
+++ b/transformer_engine/jax/csrc/extensions/attention.cpp
@@ -250,13 +250,18 @@ pybind11::tuple GetFusedAttnForwardWorkspaceSizes(
   auto v_tensor = TensorWrapper(nullptr, v_shape, dtype);
   auto o_shape = is_ragged ? std::vector<size_t>{input_batch * q_max_seqlen, attn_heads, v_head_dim}
                            : std::vector<size_t>{input_batch, q_max_seqlen, attn_heads, v_head_dim};
+#ifdef USE_ROCM
+  const DType o_dtype = dtype == DType::kFloat8E4M3 ? DType::kBFloat16 : dtype;
+#else
+  const DType o_dtype = dtype;
+#endif
 
   auto bias_shape = std::vector<size_t>{bias_batch, bias_heads, q_max_seqlen, kv_max_seqlen};
   auto bias_tensor = TensorWrapper(nullptr, bias_shape, dtype);
 
   // F16 doesn't use this tensor
   auto s_tensor = TensorWrapper(nullptr, std::vector<size_t>{1}, dtype);
-  auto o_tensor = TensorWrapper(nullptr, o_shape, dtype);
+  auto o_tensor = TensorWrapper(nullptr, o_shape, o_dtype);
 
   auto dummy_rng_state_tensor = TensorWrapper(nullptr, std::vector<size_t>{2}, DType::kInt64);
   auto dummy_page_table_tensor = TensorWrapper(nullptr, std::vector<size_t>{1}, DType::kInt32);
@@ -407,9 +412,16 @@ static void FusedAttnForwardImpl(
     float dropout_probability, NVTE_Bias_Type bias_type, NVTE_Mask_Type mask_type,
     NVTE_Softmax_Type softmax_type, NVTE_QKV_Layout qkv_layout, DType dtype, DType wkspace_dtype,
     bool is_training, bool return_max_logit, bool deterministic, int64_t window_size_left,
-    int64_t window_size_right, bool bottom_right_diagonal) {
+    int64_t window_size_right, bool bottom_right_diagonal, void *q_scale_inv, void *k_scale_inv,
+    void *v_scale_inv) {
   FUSED_ATTN_IMPL_COMMON_BLOCK;
 
+#ifdef USE_ROCM
+  const DType o_dtype = dtype == DType::kFloat8E4M3 ? DType::kBFloat16 : dtype;
+#else
+  const DType o_dtype = dtype;
+#endif
+
   /* Input tensors */
   auto bias_tensor = TensorWrapper(bias, bias_shape, dtype);
   auto softmax_offset_tensor =
@@ -417,7 +429,7 @@ static void FusedAttnForwardImpl(
 
   if (is_ragged) {
     auto output_size = input_batch * q_max_seqlen * attn_heads * v_head_dim;
-    (void)cudaMemsetAsync(output, 0, output_size * typeToSize(dtype), stream);
+    (void)cudaMemsetAsync(output, 0, output_size * typeToSize(o_dtype), stream);
 
     // Memset to 0xF0 for filling large negative numbers
     auto softmax_aux_size = input_batch * q_max_seqlen * attn_heads;
@@ -431,7 +443,7 @@ static void FusedAttnForwardImpl(
   auto s_tensor = TensorWrapper(nullptr, std::vector<size_t>{1}, dtype);  // not used in F16
   auto o_shape = is_ragged ? std::vector<size_t>{input_batch * q_max_seqlen, attn_heads, v_head_dim}
                            : std::vector<size_t>{input_batch, q_max_seqlen, attn_heads, v_head_dim};
-  auto o_tensor = TensorWrapper(output, o_shape, dtype);
+  auto o_tensor = TensorWrapper(output, o_shape, o_dtype);
 
   /* Prepare RNG state */
   auto rng_state_tensor = TensorWrapper(rng_state, std::vector<size_t>{2}, DType::kInt64);
@@ -501,6 +513,14 @@ static void FusedAttnForwardImpl(
   auto q_tensor = TensorWrapper(q_ptr, q_shape, dtype);
   auto k_tensor = TensorWrapper(k_ptr, k_shape, dtype);
   auto v_tensor = TensorWrapper(v_ptr, v_shape, dtype);
+  // FP16/BF16 leave scale_inv null. FP8 carries one FP32 inverse scale per tensor.
+  if (dtype == DType::kFloat8E4M3) {
+    NVTE_CHECK(q_scale_inv != nullptr && k_scale_inv != nullptr && v_scale_inv != nullptr,
+               "FP8 fused attention requires one FP32 scale_inv for each of Q, K, and V.");
+    q_tensor.set_rowwise_scale_inv(q_scale_inv, DType::kFloat32, std::vector<size_t>{1});
+    k_tensor.set_rowwise_scale_inv(k_scale_inv, DType::kFloat32, std::vector<size_t>{1});
+    v_tensor.set_rowwise_scale_inv(v_scale_inv, DType::kFloat32, std::vector<size_t>{1});
+  }
 
   FusedAttnFwdParamsWrapper params;
   params.set_Q(q_tensor.data())
@@ -578,13 +598,36 @@ Error_Type FusedAttnForwardFFI(cudaStream_t stream, Buffer_Type q_buf, Buffer_Ty
                                Buffer_Type softmax_offset_buf, Buffer_Type seed_buf,
                                Buffer_Type q_cu_seqlens_buf, Buffer_Type kv_cu_seqlens_buf,
                                Buffer_Type q_seq_offsets_buf, Buffer_Type k_seq_offsets_buf,
-                               Variadic_Buffer_Type _unused_args, Result_Type output_buf,
-                               Result_Type softmax_aux_buf, Result_Type max_tensor_buf,
-                               Result_Type rng_state_buf, Result_Type workspace_buf,
+                               Buffer_Type /*q_segment_ids*/, Buffer_Type /*kv_segment_ids*/,
+                               Buffer_Type /*q_segment_pos*/, Buffer_Type /*kv_segment_pos*/,
+                               Buffer_Type q_scale_inv_buf, Buffer_Type k_scale_inv_buf,
+                               Buffer_Type v_scale_inv_buf, Variadic_Buffer_Type _unused_args,
+                               Result_Type output_buf, Result_Type softmax_aux_buf,
+                               Result_Type max_tensor_buf, Result_Type rng_state_buf,
+                               Result_Type workspace_buf,
                                Dictionary attrs) {
   FUSED_ATTN_FFI_GET_ATTRS;
   bool return_max_logit = get_attr_value_or_default<bool>(attrs, "return_max_logit", false);
 
+  // Non-FP8 passes an empty buffer. A real FP8 launch passes one FP32 value.
+  auto scale_inv_arg = [](const Buffer_Type &buf) -> void * {
+    if (product(buf.dimensions()) == 0) {
+      return nullptr;
+    }
+    NVTE_CHECK(product(buf.dimensions()) == 1 &&
+                   convert_ffi_datatype_to_te_dtype(buf.element_type()) == DType::kFloat32,
+               "FP8 fused attention scale_inv must be one FP32 value.");
+    return buf.untyped_data();
+  };
+  void *q_scale_inv = nullptr;
+  void *k_scale_inv = nullptr;
+  void *v_scale_inv = nullptr;
+  if (dtype == DType::kFloat8E4M3) {
+    q_scale_inv = scale_inv_arg(q_scale_inv_buf);
+    k_scale_inv = scale_inv_arg(k_scale_inv_buf);
+    v_scale_inv = scale_inv_arg(v_scale_inv_buf);
+  }
+
   FusedAttnForwardImpl(
       stream, q_buf.untyped_data(), k_buf.untyped_data(), v_buf.untyped_data(),
       bias_buf.untyped_data(), softmax_offset_buf.untyped_data(), seed_buf.untyped_data(),
@@ -596,7 +639,8 @@ Error_Type FusedAttnForwardFFI(cudaStream_t stream, Buffer_Type q_buf, Buffer_Ty
       q_max_seqlen, kv_max_seqlen, attn_heads, num_gqa_groups, bias_heads, qk_head_dim, v_head_dim,
       max_segments_per_seq, wkspace_size, scaling_factor, dropout_probability, bias_type, mask_type,
       softmax_type, qkv_layout, dtype, wkspace_dtype, is_training, return_max_logit, deterministic,
-      window_size_left, window_size_right, bottom_right_diagonal);
+      window_size_left, window_size_right, bottom_right_diagonal, q_scale_inv, k_scale_inv,
+      v_scale_inv);
   return ffi_with_cuda_error_check();
 }
 
@@ -613,6 +657,13 @@ XLA_FFI_DEFINE_HANDLER_SYMBOL(FusedAttnForwardHandler, FusedAttnForwardFFI,
                                   .Arg<Buffer_Type>()      // kv_cu_seqlens
                                   .Arg<Buffer_Type>()      // q_seq_offsets
                                   .Arg<Buffer_Type>()      // k_seq_offsets
+                                  .Arg<Buffer_Type>()      // q_segment_ids
+                                  .Arg<Buffer_Type>()      // kv_segment_ids
+                                  .Arg<Buffer_Type>()      // q_segment_pos
+                                  .Arg<Buffer_Type>()      // kv_segment_pos
+                                  .Arg<Buffer_Type>()      // q_scale_inv
+                                  .Arg<Buffer_Type>()      // k_scale_inv
+                                  .Arg<Buffer_Type>()      // v_scale_inv
                                   .RemainingArgs()         // _cp_aux_args unused
                                   .Ret<Buffer_Type>()      // output
                                   .Ret<Buffer_Type>()      // softmax_aux
diff --git a/transformer_engine/pytorch/attention/dot_product_attention/utils.py b/transformer_engine/pytorch/attention/dot_product_attention/utils.py
index d7a5a2593..406b83167 100644
--- a/transformer_engine/pytorch/attention/dot_product_attention/utils.py
+++ b/transformer_engine/pytorch/attention/dot_product_attention/utils.py
@@ -725,6 +725,14 @@ def _disable_all_flash_attention() -> None:
         fp8_recipe = fp8_meta["recipe"]
         if fp8_meta.get("local_recipes", None) is not None:
             fp8_recipe = fp8_meta["local_recipes"][0]
+        if use_fused_attention and IS_HIP_EXTENSION and (
+            is_training or not fp8_recipe.delayed()
+        ):
+            logger.debug(
+                "Disabling ROCm FP8 FusedAttention: AITER FP8 ASM supports "
+                "inference with delayed tensor scaling only"
+            )
+            use_fused_attention = False
         if use_flash_attention_2 and FlashAttentionUtils.is_installed:
             logger.debug("Disabling FlashAttention 2 for FP8 attention")
             use_flash_attention_2 = False
diff --git a/transformer_engine/pytorch/attention/multi_head_attention.py b/transformer_engine/pytorch/attention/multi_head_attention.py
index e07e203ca..3a89565b4 100644
--- a/transformer_engine/pytorch/attention/multi_head_attention.py
+++ b/transformer_engine/pytorch/attention/multi_head_attention.py
@@ -1,3 +1,5 @@
+# This file was modified for portability to AMDGPU
+# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved.
 # Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 #
 # See LICENSE for license information.
@@ -7,6 +9,7 @@
 import collections
 from typing import Any, Callable, Dict, List, Optional, Tuple, Union
 import torch
+from torch.utils.cpp_extension import IS_HIP_EXTENSION
 
 from transformer_engine.pytorch.quantization import FP8GlobalStateManager, QuantizerRole
 from transformer_engine.pytorch.quantized_tensor import QuantizedTensorStorage
@@ -958,10 +961,14 @@ def forward(
         # DPA: produce FP8 output to take advantage of O amax from DPA; Projection Gemm can take FP8 or F16 inputs
         # 1. FP8DS/FP8CS recipe: produce FP8 output
         # 2. MXFP8 recipe: produce F16 output; again, due to quantization dimensions mismatch
+        # gfx950 AITER FP8 attention kernels are fp8bf16: Q/K/V are FP8 but O
+        # is BF16. CUDA fused attention may continue producing FP8 output.
         # For CustomRecipe, fp8_dpa only controls DPA-internal quantization.
         # External MHA boundary tensors become FP8 only when fp8_mha is enabled.
         dpa_fp8_output_enabled = fp8_mha if custom_recipe else (fp8_dpa or fp8_mha)
-        dpa_fp8_output = fp8 and dpa_fp8_output_enabled and not mxfp8_scaling
+        dpa_fp8_output = (
+            fp8 and dpa_fp8_output_enabled and not mxfp8_scaling and not IS_HIP_EXTENSION
+        )
         # Projection Gemm: match DPA output except
         # 1. FP8CS recipe: produce F16 grads; again, due to cuBLAS limitation
         proj_fp8_grad = dpa_fp8_output and not float8_current_scaling