minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

[CPU] add mamba fla kernels for Qwen3-next (#12324)

Browse Source
This commit is contained in:
blzheng
2025-12-06 14:16:23 +08:00
committed by GitHub
parent 7b0c7ad163
commit d257bf87b9
4 changed files with 577 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files
sgl-kernel/csrc/cpu/common.h
+10
View File
@@ -128,6 +128,16 @@ namespace {
#define CHECK_EQ(a, b) TORCH_CHECK((a) == (b), "CHECK_EQ(" #a ", " #b ") failed. ", a, " vs ", b)
template <bool is_only_lastdim_contiguous>
static inline void CHECK_INPUT_SHAPE_DTYPE(const at::Tensor& tensor, const at::IntArrayRef sizes, at::ScalarType st) {
TORCH_CHECK(tensor.sizes() == sizes, "Input tensor shape mismatch: expected ", sizes, ", got ", tensor.sizes());
TORCH_CHECK(tensor.scalar_type() == st, "Input tensor dtype mismatch");
if constexpr (is_only_lastdim_contiguous) {
CHECK_LAST_DIM_CONTIGUOUS_INPUT(tensor);
} else {
CHECK_INPUT(tensor);
}
}
#define CHECK_GE(a, b) TORCH_CHECK((a) >= (b), "CHECK_GE(" #a ", " #b ") failed. ", a, " vs ", b)
// [NB] Parallel Routines