params.h

Go to the documentation of this file.

// 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; 
 
  int qL_rem; 
  int kL_rem; 
  int qL_off; 
 
  int64_t Q_strides[3]; 
  int64_t K_strides[3]; 
  int64_t V_strides[3]; 
  int64_t O_strides[3]; 
};
 
struct AttnMaskParams {
  int64_t M_strides[3]; 
};
 
} // namespace steel
} // namespace mlx