mirror of
https://github.com/tile-ai/tilelang.git
synced 2026-10-23 07:00:29 +08:00
Wenhao Xie
37051417b7
* upd sparse attn * lint * rename * update test file * update benchmark * lint * update benchmark
17 lines
577 B
C++
17 lines
577 B
C++
// Copyright (c) Microsoft Corporation.
|
|
// Licensed under the MIT license.
|
|
|
|
#include "torch/extension.h"
|
|
#include <vector>
|
|
|
|
std::vector<at::Tensor> convert_vertical_slash_indexes(
|
|
torch::Tensor seqlens, // [BATCH, ]
|
|
torch::Tensor vertical_indexes, // [BATCH, N_HEADS, NNZ_V]
|
|
torch::Tensor slash_indexes, // [BATCH, N_HEADS, NNZ_S]
|
|
int context_size, int block_size_M, int block_size_N);
|
|
|
|
PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
|
|
m.def("convert_vertical_slash_indexes", &convert_vertical_slash_indexes,
|
|
"dynamic sparse index function");
|
|
}
|