Skip to content

[Easy] replace c10::optional with std::optional #58

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 1 commit into from
Mar 27, 2025
Merged

[Easy] replace c10::optional with std::optional #58

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
44 changes: 22 additions & 22 deletions csrc/common/pytorch_shim.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -15,13 +15,13 @@
* that the following types are not support by PyTorch libary bindings:
* - `int`
* - `float`
* - `c10::optional<T> &`
* - `c10::optional<const at::Tensor> &`
* - `std::optional<T> &`
* - `std::optional<const at::Tensor> &`
* So we convert them to (respectively):
* - `int64_t`
* - `double`
* - `const c10::optional<T>&`
* - `const c10::optional<at::Tensor>&`
* - `const std::optional<T>&`
* - `const std::optional<at::Tensor>&`
*/

template<typename T>
Expand All @@ -38,36 +38,36 @@ template<typename T>
T convert_from_pytorch_compatible_type(pytorch_library_compatible_type_t<T> arg)
{ return pytorch_library_compatible_type<T>::convert_from_type(arg); }

// Map `c10::optional<T> &` -> `const c10::optional<T>&`
// Map `std::optional<T> &` -> `const std::optional<T>&`
// (NOTE: this is bit unsafe but non of the ops in flash_attn mutate
// the optional container)
template<typename T>
struct pytorch_library_compatible_type<c10::optional<T> &> {
using type = const c10::optional<T>&;
static c10::optional<T>& convert_from_type(const c10::optional<T> &arg) {
return const_cast<c10::optional<T>&>(arg);
struct pytorch_library_compatible_type<std::optional<T> &> {
using type = const std::optional<T>&;
static std::optional<T>& convert_from_type(const std::optional<T> &arg) {
return const_cast<std::optional<T>&>(arg);
}
};

// Map `c10::optional<T>` ->
// `c10::optional<pytorch_library_compatible_type_t<T>>`
// (NOTE: tested for `c10::optional<int>` -> `c10::optional<int64_t>`)
// Map `std::optional<T>` ->
// `std::optional<pytorch_library_compatible_type_t<T>>`
// (NOTE: tested for `std::optional<int>` -> `std::optional<int64_t>`)
template<typename T>
struct pytorch_library_compatible_type<c10::optional<T>> {
using type = c10::optional<pytorch_library_compatible_type_t<T>>;
static c10::optional<pytorch_library_compatible_type_t<T>> convert_from_type(c10::optional<T> arg) {
struct pytorch_library_compatible_type<std::optional<T>> {
using type = std::optional<pytorch_library_compatible_type_t<T>>;
static std::optional<pytorch_library_compatible_type_t<T>> convert_from_type(std::optional<T> arg) {
return arg;
}
};

// Map `c10::optional<const at::Tensor>&` -> `const c10::optional<at::Tensor>&`
// Map `std::optional<const at::Tensor>&` -> `const std::optional<at::Tensor>&`
template<>
struct pytorch_library_compatible_type<c10::optional<const at::Tensor> &> {
using type = const c10::optional<at::Tensor>&;
static c10::optional<const at::Tensor>& convert_from_type(
const c10::optional<at::Tensor> &arg) {
return const_cast<c10::optional<const at::Tensor>&>(
reinterpret_cast<const c10::optional<const at::Tensor>&>(arg));
struct pytorch_library_compatible_type<std::optional<const at::Tensor> &> {
using type = const std::optional<at::Tensor>&;
static std::optional<const at::Tensor>& convert_from_type(
const std::optional<at::Tensor> &arg) {
return const_cast<std::optional<const at::Tensor>&>(
reinterpret_cast<const std::optional<const at::Tensor>&>(arg));
}
};

Expand Down
8 changes: 4 additions & 4 deletions csrc/flash_attn/flash_api_sparse.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -318,11 +318,11 @@ mha_varlen_fwd_sparse(at::Tensor &q, // total_q x num_heads x head_size, total_
const at::Tensor &block_offset,
const at::Tensor &column_count,
const at::Tensor &column_index,
const c10::optional<at::Tensor> &out_, // total_q x num_heads x head_size, total_k := \sum_{i=0}^{b} s_i
const std::optional<at::Tensor> &out_, // total_q x num_heads x head_size, total_k := \sum_{i=0}^{b} s_i
const at::Tensor &cu_seqlens_q, // b+1
const at::Tensor &cu_seqlens_k, // b+1
const c10::optional<at::Tensor> &seqused_k, // b. If given, only this many elements of each batch element's keys are used.
const c10::optional<at::Tensor> &alibi_slopes_, // num_heads or b x num_heads
const std::optional<at::Tensor> &seqused_k, // b. If given, only this many elements of each batch element's keys are used.
const std::optional<at::Tensor> &alibi_slopes_, // num_heads or b x num_heads
int64_t max_seqlen_q,
const int64_t max_seqlen_k,
const double p_dropout,
Expand All @@ -331,7 +331,7 @@ mha_varlen_fwd_sparse(at::Tensor &q, // total_q x num_heads x head_size, total_
bool is_causal,
const double softcap,
const bool return_softmax,
c10::optional<at::Generator> gen_) {
std::optional<at::Generator> gen_) {

auto [cc_major, cc_minor] = get_compute_capability(get_current_device());
bool is_sm8x = cc_major == 8 && cc_minor >= 0;
Expand Down
8 changes: 4 additions & 4 deletions csrc/flash_attn/flash_api_torch_lib.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -86,11 +86,11 @@ mha_varlen_fwd_sparse(at::Tensor &q, // total_q x num_heads x head_size, total_
const at::Tensor &block_offset,
const at::Tensor &column_count,
const at::Tensor &column_index,
const c10::optional<at::Tensor> &out_, // total_q x num_heads x head_size, total_k := \sum_{i=0}^{b} s_i
const std::optional<at::Tensor> &out_, // total_q x num_heads x head_size, total_k := \sum_{i=0}^{b} s_i
const at::Tensor &cu_seqlens_q, // b+1
const at::Tensor &cu_seqlens_k, // b+1
const c10::optional<at::Tensor> &seqused_k, // b. If given, only this many elements of each batch element's keys are used.
const c10::optional<at::Tensor> &alibi_slopes_, // num_heads or b x num_heads
const std::optional<at::Tensor> &seqused_k, // b. If given, only this many elements of each batch element's keys are used.
const std::optional<at::Tensor> &alibi_slopes_, // num_heads or b x num_heads
int64_t max_seqlen_q,
const int64_t max_seqlen_k,
const double p_dropout,
Expand All @@ -99,7 +99,7 @@ mha_varlen_fwd_sparse(at::Tensor &q, // total_q x num_heads x head_size, total_
bool is_causal,
const double softcap,
const bool return_softmax,
c10::optional<at::Generator> gen_);
std::optional<at::Generator> gen_);

/**
* Torch Library Registration
Expand Down