build/html/attn_2params_8h_source.html

// Copyright © 2024 Apple Inc.


#pragma once


// Attn param classes


namespace mlx {

namespace steel {


struct AttnParams {

  int B;

  int H;

  int D;


  int qL;

  int kL;


  int gqa_factor;

  float scale;


  int NQ;

  int NK;


  int NQ_aligned;

  int NK_aligned;


  int64_t Q_strides[3];

  int64_t K_strides[3];

  int64_t V_strides[3];

  int64_t O_strides[3];

};


} // namespace steel

} // namespace mlx

mlx::steel
Definition attn.h:19

mlx
Definition allocator.h:7

mlx::steel::AttnParams
Definition params.h:12

mlx::steel::AttnParams::D
int D
Head Dim.
Definition params.h:15

mlx::steel::AttnParams::B
int B
Batch Size.
Definition params.h:13

mlx::steel::AttnParams::gqa_factor
int gqa_factor
Group Query factor.
Definition params.h:20

mlx::steel::AttnParams::H
int H
Heads.
Definition params.h:14

mlx::steel::AttnParams::NQ
int NQ
Number of query blocks.
Definition params.h:23

mlx::steel::AttnParams::kL
int kL
Key Sequence Length.
Definition params.h:18

mlx::steel::AttnParams::NQ_aligned
int NQ_aligned
Number of full query blocks.
Definition params.h:26

mlx::steel::AttnParams::qL
int qL
Query Sequence Length.
Definition params.h:17

mlx::steel::AttnParams::NK
int NK
Number of key/value blocks.
Definition params.h:24

mlx::steel::AttnParams::Q_strides
int64_t Q_strides[3]
Query strides (B, H, L, D = 1)
Definition params.h:29

mlx::steel::AttnParams::NK_aligned
int NK_aligned
Number of full key/value blocks.
Definition params.h:27

mlx::steel::AttnParams::O_strides
int64_t O_strides[3]
Output strides (B, H, L, D = 1)
Definition params.h:32

mlx::steel::AttnParams::V_strides
int64_t V_strides[3]
Value strides (B, H, L, D = 1)
Definition params.h:31

mlx::steel::AttnParams::scale
float scale
Attention scale.
Definition params.h:21

mlx::steel::AttnParams::K_strides
int64_t K_strides[3]
Key strides (B, H, L, D = 1)
Definition params.h:30