TheRock/rocm-libraries/projects/hipdnn/flatbuffers_sdk/schemas/sdpa_attributes.fbs

// Copyright © Advanced Micro Devices, Inc., or its affiliates.
// SPDX-License-Identifier: MIT

include "data_types.fbs";

namespace hipdnn_flatbuffers_sdk.data_objects;

enum DiagonalAlignment : byte {
    TOP_LEFT     = 0,
    BOTTOM_RIGHT = 1,
}

enum AttentionImplementation : byte {
    AUTO      = 0,
    COMPOSITE = 1,
    UNIFIED   = 2,
}

table SdpaAttributes {
    // Required input tensor UIDs
    q_tensor_uid:             long;
    k_tensor_uid:             long;
    v_tensor_uid:             long;

    // Required output tensor UID
    o_tensor_uid:             long;

    // Optional input tensor UIDs
    attn_mask_tensor_uid:     long = null;  // Additive attention bias (Bias)
    scale_tensor_uid:         long = null;  // Attention scale tensor (Attn_scale)
    seq_len_q_tensor_uid:     long = null;
    seq_len_kv_tensor_uid:    long = null;
    seed_tensor_uid:          long = null;  // Dropout seed
    offset_tensor_uid:        long = null;  // Dropout offset
    dropout_mask_tensor_uid:  long = null;
    dropout_scale_tensor_uid: long = null;
    page_table_k_tensor_uid:  long = null;
    page_table_v_tensor_uid:  long = null;
    block_mask_tensor_uid:    long = null;
    sink_token_tensor_uid:    long = null;

    // FP8-specific input tensor UIDs
    descale_q_tensor_uid:     long = null;
    descale_k_tensor_uid:     long = null;
    descale_v_tensor_uid:     long = null;
    descale_s_tensor_uid:     long = null;
    scale_s_tensor_uid:       long = null;
    scale_o_tensor_uid:       long = null;

    // Optional output tensor UIDs
    stats_tensor_uid:         long = null;
    max_tensor_uid:           long = null;  // Logit max (Max)
    sum_exp_tensor_uid:       long = null;  // Score sum exp (Sum_exp)
    rng_dump_tensor_uid:      long = null;  // RNG dump for dropout debugging

    // FP8-specific output tensor UIDs
    amax_s_tensor_uid:        long = null;  // Absolute max for attention scores
    amax_o_tensor_uid:        long = null;  // Absolute max for output tensor

    // Boolean flags
    generate_stats:           bool = null;
    alibi_mask:               bool = false;
    padding_mask:             bool = false;
    causal_mask:              bool = false;              // Deprecated
    causal_mask_bottom_right: bool = false;              // Deprecated

    // Scalar attributes
    dropout_probability:      float = null;
    attn_scale_value:         float = null;
    left_bound:               long = null;
    right_bound:              long = null;
    max_seq_len_kv:           int = null;

    // Enum attributes
    diagonal_alignment:       DiagonalAlignment = TOP_LEFT;
    mma_core_mode:            DataType = UNSET;
    implementation:           AttentionImplementation = AUTO;
}