add cuda gemv (#2400)

This commit is contained in:
Awni Hannun
2025-07-22 08:24:13 -07:00
committed by GitHub
parent 1e496ddb82
commit d107d8d495
12 changed files with 198 additions and 21 deletions

24
mlx/backend/cuda/gemv.h Normal file
View File

@@ -0,0 +1,24 @@
// Copyright © 2025 Apple Inc.
#pragma once
#include "mlx/backend/cuda/device.h"
namespace mlx::core::cu {
bool can_use_gemv(int M, int N, int K, bool a_transposed, bool b_transposed);
void gemv(
const array& a,
const array& b,
array& out,
int M,
int N,
int K,
uint32_t batch_count,
const mlx::core::Shape& batch_shape,
const mlx::core::Strides& a_batch_strides,
const mlx::core::Strides& b_batch_strides,
CommandEncoder& encoder);
} // namespace mlx::core::cu