chenchenghao/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 11 Packages Projects Releases Wiki Activity

[JIT Kernel][Feature] Support JIT custom all reduce (rewrite as v2) (#19880)

Browse Source 
Co-authored-by: Xiaoyu Zhang <35585791+BBuf@users.noreply.github.com>
This commit is contained in:
DarkSharpness
2026-03-20 18:24:07 +08:00
committed by GitHub
parent 2099943a49
commit 2dd9196079
16 changed files with 2154 additions and 96 deletions
Download Patch File Download Diff File Expand all files Collapse all files

8
python/sglang/jit_kernel/include/sgl_kernel/vec.cuh
View File

@@ -80,15 +80,11 @@ struct AlignedVector {
public:
/// \brief Vectorized load from `ptr` at the given element `offset`.
template <typename U>
SGL_DEVICE void load(const U* ptr, std::size_t offset = 0) {
static_assert(std::is_same_v<U, T> || std::is_same_v<U, void>);
SGL_DEVICE void load(const void* ptr, int64_t offset = 0) {
m_storage = reinterpret_cast<const storage_t*>(ptr)[offset];
}
/// \brief Vectorized store to `ptr` at the given element `offset`.
template <typename U>
SGL_DEVICE void store(U* ptr, std::size_t offset = 0) const {
static_assert(std::is_same_v<U, T> || std::is_same_v<U, void>);
SGL_DEVICE void store(void* ptr, int64_t offset = 0) const {
reinterpret_cast<storage_t*>(ptr)[offset] = m_storage;
}
/// \brief Fill all N elements with the same `value`.