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

[jit_kernel] Move JIT kernels into namespace sglang (#33400)

Browse Source 
Co-authored-by: Claude Fable 5 <noreply@anthropic.com>
This commit is contained in:
DarkSharpness
2026-08-08 16:10:15 +08:00
committed by GitHub
co-authored by Claude Fable 5
parent 5fdf6cd18f
commit 4ad5bb5d9a
178 changed files with 626 additions and 399 deletions
Download Patch File Download Diff File Expand all files Collapse all files
.claude/skills/add-jit-kernel/SKILL.md
+3 -2
View File
@@ -27,6 +27,7 @@ Add a new operation that scales each element of a tensor by a scalar factor:
These hold for every step below.
- **`namespace sglang` is where JIT code lives.** Open it after the include block and close it at the end of the file, with the device kernels, traits and host wrapper inside. The shared `host::` / `device::` helpers are nested in it too, so they resolve unqualified. `load_jit` emits the `TVM_FFI_DLL_EXPORT_TYPED_FUNC` wrapper inside `namespace sglang` as well, so the `kernel_name` you pass from Python needs no `sglang::` prefix.
- **Check where the check is cheapest: `static_assert` > C++ host check > cached Python > per-call Python.** Anything fixed at compile time is a `static_assert`. Anything about the tensors is a `TensorMatcher` / `CHECK_HOST` in the C++ launcher, free next to a kernel launch. A check Python cannot delegate goes inside the `@cache_once` module factory, where it runs once per specialisation. What remains in the per-call entry point costs interpreter time on *every* forward, so it should be nothing but picking the module and allocating `out`.
- **Fixed-width integer types.** Prefer `int32_t` / `int64_t` / `uint32_t` / `size_t` over `int`, `long`, or `long long`, so an index has the same width on both sides of the FFI boundary. Bare `int` is fine only where the width plainly cannot matter — an unrolled loop counter over a `constexpr` bound, a template `int` parameter. Shapes arrive as `int64_t` (`SymbolicSize::unwrap()`); narrowing to `uint32_t` for in-kernel indexing is a deliberate act, so write the `static_cast` explicitly and only where the range is known.
- **Doxygen comments in C++.** Document exported entities with `///` or `/** ... */` blocks using `\brief`, `\param`, `\tparam`, `\return`, the way `include/sgl_kernel/` does. `python -m sglang.kernels.jit` writes `CommentFormat: Doxygen` into `.clangd` when clangd is 21 or newer, so these render on hover in the editor. Plain `//` remains fine for implementation notes inside a function body.
@@ -251,7 +252,7 @@ The implementation fully uses the project abstractions described above:
#include <dlpack/dlpack.h>
#include <tvm/ffi/container/tensor.h>
namespace {
namespace sglang {
/**
* \brief Element-wise scale using vectorized 128-bit loads/stores.
@@ -357,7 +358,7 @@ void scale(tvm::ffi::TensorView dst, tvm::ffi::TensorView src, float factor) {
n);
}
} // namespace
} // namespace sglang
```
**Key points:**
docs/docs/developer_guide/development_jit_kernel_guide.mdx
+9 -2
View File
@@ -23,6 +23,11 @@ After generating the file, restart the clangd language server. It should now rec
C++ source code is located in `python/sglang/kernels/jit/csrc`.
Reusable functions should be placed in `python/sglang/kernels/jit/include`.
JIT C++ lives in `namespace sglang`: open it after the include block and close it at the
end of the file, with the device kernels and the host wrapper both inside.
The shared `host::` and `device::` helpers are nested in it as well, so they resolve unqualified
and need no `sglang::` prefix.
We use [tvm-ffi](https://github.com/apache/tvm-ffi) for efficient foreign language bindings.
Refer to the [documentation](https://tvm.apache.org/ffi/) for advanced usage, such as exporting C++ objects.
Typically, `tvm::ffi::TensorView` is sufficient for passing PyTorch Tensors from Python.
@@ -33,6 +38,8 @@ Python interfaces are defined in `python/sglang/kernels/jit`.
The `load_jit` utility function in `python/sglang/kernels/jit/utils/compile.py` loads and returns the compiled module.
To export a C++ function (e.g., `cpp_func`), pass `cuda_wrappers=[("func", "cpp_func")]` to `load_jit`.
The function can then be called in Python as `module.func`.
`load_jit` emits the export wrapper inside `namespace sglang`, so write `cpp_func` without a
`sglang::` prefix.
For caching compiled modules, prefer `sglang.kernels.jit.utils.cache_once` over `functools.lru_cache`.
`functools.lru_cache` is not compatible with `torch.compile`.
@@ -174,7 +181,7 @@ Write your CUDA kernel in [kernels/jit/csrc/add_constant.cuh](https://github.com
#include <cstddef>
#include <cstdint>
namespace {
namespace sglang {
template <int32_t kConstant>
__global__ void add_constant_kernel(int32_t* dst, const int32_t* src, size_t length) {
@@ -217,7 +224,7 @@ void add_constant(tvm::ffi::TensorView dst, tvm::ffi::TensorView src) {
num_elements);
}
} // namespace
} // namespace sglang
```
python/sglang/kernels/jit/csrc/add_constant.cuh
+2 -2
View File
@@ -10,7 +10,7 @@
#include <cstddef>
#include <cstdint>
namespace {
namespace sglang {
constexpr size_t kBlockSize = 256;
constexpr size_t kVectorizedMinElements = 1 << 20;
@@ -98,4 +98,4 @@ void add_constant(tvm::ffi::TensorView dst, tvm::ffi::TensorView src) {
}
}
} // namespace
} // namespace sglang
python/sglang/kernels/jit/csrc/attention/fixup_zero_kv.cuh
+2 -2
View File
@@ -11,7 +11,7 @@
#include <cstdint>
namespace {
namespace sglang {
constexpr int kFixupBlockSize = 256;
@@ -136,4 +136,4 @@ void fixup_zero_kv_rows(
nh);
}
} // namespace
} // namespace sglang
python/sglang/kernels/jit/csrc/attention/fused_fp8_qkv_kv_cache.cuh
+2 -2
View File
@@ -9,7 +9,7 @@
#include <cstdint>
namespace {
namespace sglang {
struct FusedQkvParams {
const void* __restrict__ q;
@@ -199,4 +199,4 @@ struct FusedFp8QkvKvCache {
}
};
} // namespace
} // namespace sglang
python/sglang/kernels/jit/csrc/attention/kda_fused_decode.cuh
+3 -3
View File
@@ -47,6 +47,8 @@
// Local PTX primitives (cp.async / mbarrier / async-proxy fence)
namespace sglang {
namespace ptx {
// Generic ptr -> 32-bit `.shared` address: inline-PTX `.shared` instructions
@@ -133,8 +135,6 @@ static SGL_DEVICE void fence_async_smem() {
} // namespace ptx
namespace {
constexpr int kDimK = 128;
constexpr int kDimV = 128;
constexpr int kKernelWidth = 4;
@@ -1061,4 +1061,4 @@ struct KdaFusedDecodeKernel {
}
};
} // namespace
} // namespace sglang
python/sglang/kernels/jit/csrc/attention/kda_packed_decode.cuh
+2 -2
View File
@@ -26,7 +26,7 @@
#include <cstdint>
namespace {
namespace sglang {
struct KdaPackedDecodeParams {
const bf16_t* __restrict__ mixed_qkv; // [B, 2*H*K + HV*V]
@@ -237,4 +237,4 @@ struct KdaPackedDecodeKernel {
}
};
} // namespace
} // namespace sglang
python/sglang/kernels/jit/csrc/deepseek_v32/indexer_k.cuh
+2 -2
View File
@@ -24,7 +24,7 @@
#include <bit>
#include <cstdint>
namespace {
namespace sglang {
using deepseek_v4::fp8::pack_fp8;
@@ -426,4 +426,4 @@ struct FusedKIndexerNormRopeStoreKernel {
}
};
} // namespace
} // namespace sglang
python/sglang/kernels/jit/csrc/deepseek_v4/c128.cuh
+2 -2
View File
@@ -16,7 +16,7 @@
#include <cstdint>
namespace {
namespace sglang {
using Plan128 = device::compress::PrefillPlan;
using IndiceT = int32_t;
@@ -519,4 +519,4 @@ struct FlashCompress128Kernel {
}
};
} // namespace
} // namespace sglang
python/sglang/kernels/jit/csrc/deepseek_v4/c128_online.cuh
+3 -7
View File
@@ -19,6 +19,8 @@
#include <cfloat>
#include <cstdint>
namespace sglang {
namespace device::compress {
/// \brief Plan entry for online compress 128 prefill.
@@ -68,8 +70,6 @@ static_assert(sizeof(OnlinePrefillPlan) == kOnlinePrefillPlanDim * sizeof(Online
} // namespace host::compress
namespace {
using OnlinePlan = device::compress::OnlinePrefillPlan;
using IndiceT = int32_t;
@@ -594,8 +594,6 @@ struct FlashCompress128OnlineKernel {
}
};
} // namespace
namespace host::compress {
using OnlinePlanResult = tvm::ffi::Tuple<uint32_t, uint32_t>;
@@ -718,9 +716,7 @@ inline OnlinePlanResult plan_online_prefill(
} // namespace host::compress
namespace {
[[maybe_unused]]
constexpr auto& plan_compress_online_prefill = host::compress::plan_online_prefill;
} // namespace
} // namespace sglang
python/sglang/kernels/jit/csrc/deepseek_v4/c128_online_v2.cuh
+2 -6
View File
@@ -22,7 +22,7 @@
#include <cstring>
#include <type_traits>
namespace {
namespace sglang {
using PlanD = device::compress::DecodePlan;
using PlanC = device::compress::CompressPlan;
@@ -561,8 +561,6 @@ struct FlashCompress128OnlineKernel {
}
};
} // namespace
// ===========================================================================
// Plan builders. Mirrors the offline v2 pattern (`c_plan.cuh`):
// - Decode: a single GPU kernel reads seq_lens / req_to_token /
@@ -925,9 +923,7 @@ inline OnlinePrefillPlan plan_online_prefill(
} // namespace host::compress
namespace {
[[maybe_unused]] constexpr auto& plan_compress_128_online_decode = host::compress::plan_online_decode;
[[maybe_unused]] constexpr auto& plan_compress_128_online_prefill = host::compress::plan_online_prefill;
} // namespace
} // namespace sglang
python/sglang/kernels/jit/csrc/deepseek_v4/c128_v2.cuh
+2 -2
View File
@@ -31,7 +31,7 @@
#include <cstdint>
#include <type_traits>
namespace {
namespace sglang {
using PlanD = device::compress::DecodePlan;
using PlanC = device::compress::CompressPlan;
@@ -509,4 +509,4 @@ struct FlashCompress128Kernel {
}
};
} // namespace
} // namespace sglang
python/sglang/kernels/jit/csrc/deepseek_v4/c4.cuh
+2 -2
View File
@@ -15,7 +15,7 @@
#include <cstdint>
namespace {
namespace sglang {
using Plan4 = device::compress::PrefillPlan;
using IndiceT = int32_t;
@@ -546,4 +546,4 @@ struct FlashCompress4Kernel {
}
};
} // namespace
} // namespace sglang
python/sglang/kernels/jit/csrc/deepseek_v4/c4_v2.cuh
+2 -2
View File
@@ -30,7 +30,7 @@
#include <cstdint>
#include <type_traits>
namespace {
namespace sglang {
using PlanD = device::compress::DecodePlan;
using PlanC = device::compress::CompressPlan;
@@ -488,4 +488,4 @@ struct FlashCompress4Kernel {
}
};
} // namespace
} // namespace sglang
python/sglang/kernels/jit/csrc/deepseek_v4/c_plan.cuh
+4
View File
@@ -13,6 +13,8 @@
#include <cstdint>
#include <limits>
namespace sglang {
namespace host::compress {
constexpr auto kDLUInt8 = DLDataType{.code = kDLUInt, .bits = 8, .lanes = 1};
@@ -840,3 +842,5 @@ inline tvm::ffi::Tensor plan_compress_decode_legacy(
} // namespace host::compress
using namespace host::compress; // expose binding
} // namespace sglang
python/sglang/kernels/jit/csrc/deepseek_v4/common.cuh
+3 -3
View File
@@ -5,6 +5,8 @@
#include <dlpack/dlpack.h>
namespace sglang {
namespace host::compress {
using PlanResult = tvm::ffi::Tuple<uint32_t, uint32_t>;
@@ -200,9 +202,7 @@ inline PlanResult plan_prefill(
} // namespace host::compress
namespace {
[[maybe_unused]]
constexpr auto& plan_compress_prefill = host::compress::plan_prefill;
} // namespace
} // namespace sglang
python/sglang/kernels/jit/csrc/deepseek_v4/fp8_wo_a_group_major_quant.cuh
+2 -2
View File
@@ -22,7 +22,7 @@
#include <cstdint>
#include <cuda_fp8.h>
namespace {
namespace sglang {
using deepseek_v4::fp8::cast_to_ue8m0;
using deepseek_v4::fp8::inv_scale_ue8m0;
@@ -166,4 +166,4 @@ struct FP8WoAGroupMajorQuantUE8M0Kernel {
}
};
} // namespace
} // namespace sglang
python/sglang/kernels/jit/csrc/deepseek_v4/fused_norm_rope.cuh
+2 -2
View File
@@ -14,7 +14,7 @@
#include <cstdint>
#include <type_traits>
namespace {
namespace sglang {
using Plan = device::compress::PrefillPlan;
@@ -251,4 +251,4 @@ struct FusedNormRopeKernel {
}
};
} // namespace
} // namespace sglang
python/sglang/kernels/jit/csrc/deepseek_v4/fused_norm_rope_v2.cuh
+2 -2
View File
@@ -14,7 +14,7 @@
#include <cstdint>
namespace {
namespace sglang {
using PlanC = device::compress::CompressPlan;
using PlanD = device::compress::DecodePlan;
@@ -679,4 +679,4 @@ struct FusedNormRopeKernel {
}
};
} // namespace
} // namespace sglang
python/sglang/kernels/jit/csrc/deepseek_v4/hash_topk.cuh
+2 -2
View File
@@ -10,7 +10,7 @@
#include <cmath>
#include <cstdint>
namespace {
namespace sglang {
[[maybe_unused]]
SGL_DEVICE float act_sqrt_softplus(float x) {
@@ -211,4 +211,4 @@ struct MaskKernel {
}
};
} // namespace
} // namespace sglang
python/sglang/kernels/jit/csrc/deepseek_v4/main_norm_rope.cuh
+2 -2
View File
@@ -15,7 +15,7 @@
#include <bit>
#include <cstdint>
namespace {
namespace sglang {
using deepseek_v4::fp8::cast_to_ue8m0;
using deepseek_v4::fp8::inv_scale_ue8m0;
@@ -879,4 +879,4 @@ struct FusedQIndexerRopeHadamardFp4QuantKernel {
}
};
} // namespace
} // namespace sglang
python/sglang/kernels/jit/csrc/deepseek_v4/mega_moe_pre_dispatch.cuh
+2 -2
View File
@@ -12,7 +12,7 @@
#include <cstdint>
#include <cuda_fp8.h>
namespace {
namespace sglang {
using deepseek_v4::fp8::cast_to_ue8m0;
using deepseek_v4::fp8::pack_fp8;
@@ -218,4 +218,4 @@ struct MegaMoEPreDispatchKernel {
}
};
} // namespace
} // namespace sglang
python/sglang/kernels/jit/csrc/deepseek_v4/online_c128_mtp.cuh
+2 -2
View File
@@ -12,7 +12,7 @@
#include <cstdint>
#include <type_traits>
namespace {
namespace sglang {
SGL_DEVICE int64_t clamp_accept_len(int64_t delta, int64_t max_accept) {
if (delta < 0) return 0;
@@ -402,4 +402,4 @@ struct OnlineC128MTPCommitPendingKernel {
}
};
} // namespace
} // namespace sglang
python/sglang/kernels/jit/csrc/deepseek_v4/paged_mqa_metadata.cuh
+2 -2
View File
@@ -7,7 +7,7 @@
#include <dlpack/dlpack.h>
#include <tvm/ffi/container/tensor.h>
namespace {
namespace sglang {
constexpr uint32_t kBlockSize = 1024;
constexpr uint32_t kSplitKV = 256; // const for both SM90 and SM100
@@ -116,4 +116,4 @@ struct IndexerMetadataKernel {
}
};
} // namespace
} // namespace sglang
python/sglang/kernels/jit/csrc/deepseek_v4/rope.cuh
+2 -2
View File
@@ -9,7 +9,7 @@
#include <cstdint>
namespace {
namespace sglang {
using DType = bf16_t;
constexpr int64_t kRopeDim = 64;
@@ -166,4 +166,4 @@ struct FusedQKRopeKernel {
}
};
} // namespace
} // namespace sglang
python/sglang/kernels/jit/csrc/deepseek_v4/silu_and_mul_masked_post_quant.cuh
+2 -2
View File
@@ -14,7 +14,7 @@
#include <cuda_fp8.h>
#include <type_traits>
namespace {
namespace sglang {
using deepseek_v4::fp8::cast_to_ue8m0;
using deepseek_v4::fp8::pack_fp8;
@@ -537,4 +537,4 @@ struct SiluAndMulContigPostQuantKernel {
}
};
} // namespace
} // namespace sglang
python/sglang/kernels/jit/csrc/deepseek_v4/store.cuh
+2 -2
View File
@@ -16,7 +16,7 @@
#include <cstdint>
#include <cuda_fp8.h>
namespace {
namespace sglang {
using deepseek_v4::fp8::cast_to_ue8m0;
using deepseek_v4::fp8::inv_scale_ue8m0;
@@ -202,4 +202,4 @@ struct FusedStoreCacheIndexerKernel {
}
};
} // namespace
} // namespace sglang
python/sglang/kernels/jit/csrc/deepseek_v4/topk_v1.cuh
+2 -2
View File
@@ -9,7 +9,7 @@
#include <bit>
#include <cstdint>
namespace {
namespace sglang {
#ifndef SGL_TOPK
#define SGL_TOPK 512
@@ -337,4 +337,4 @@ struct TopKKernel {
}
};
} // namespace
} // namespace sglang
python/sglang/kernels/jit/csrc/deepseek_v4/topk_v2.cuh
+2 -2
View File
@@ -23,7 +23,7 @@
#include <cstdint>
#include <iterator>
namespace {
namespace sglang {
namespace impl = device::topk;
using impl::TopKProblem;
@@ -460,4 +460,4 @@ struct TopKKernel {
}
};
} // namespace
} // namespace sglang
python/sglang/kernels/jit/csrc/diffusion/causal_conv3d_cat_pad.cuh
+5 -5
View File
@@ -19,9 +19,9 @@
#include <cstdint>
namespace sglang_causal_conv3d_cat_pad {
namespace sglang {
namespace {
namespace causal_conv3d_cat_pad {
constexpr int kBlockSize = 256;
@@ -151,8 +151,6 @@ void launch_cat_pad_flat(
pad_w_left);
}
} // namespace
template <typename T>
struct CausalConv3dCatPadKernel {
static void
@@ -250,4 +248,6 @@ struct CausalConv3dCatPadKernel {
}
};
} // namespace sglang_causal_conv3d_cat_pad
} // namespace causal_conv3d_cat_pad
} // namespace sglang
python/sglang/kernels/jit/csrc/diffusion/ltx2_qknorm_split_rope.cuh
+5 -5
View File
@@ -16,9 +16,9 @@
#include <cstdint>
#include <cuda_bf16.h>
namespace sglang_ltx2_qknorm_split_rope {
namespace sglang {
namespace {
namespace ltx2_qknorm_split_rope {
constexpr int kThreads = 128;
@@ -172,8 +172,6 @@ inline void launch_one(
stride_sin_t);
}
} // namespace
struct LTX2QKNormSplitRopeKernel {
static void
run(tvm::ffi::TensorView q_out,
@@ -273,4 +271,6 @@ struct LTX2QKNormSplitRopeKernel {
}
};
} // namespace sglang_ltx2_qknorm_split_rope
} // namespace ltx2_qknorm_split_rope
} // namespace sglang
python/sglang/kernels/jit/csrc/diffusion/norm_scale_shift.cuh
+5 -5
View File
@@ -23,9 +23,9 @@
#include <cstdint>
namespace sglang_norm_scale_shift {
namespace sglang {
namespace {
namespace norm_scale_shift {
constexpr int kHidden = 3072;
constexpr int kVecElems = 16; // 32B/thread for bf16 on Blackwell.
@@ -142,8 +142,6 @@ inline uint32_t verify_qwen_geometry(host::SymbolicSize& num_rows) {
return static_cast<uint32_t>(num_rows.unwrap());
}
} // namespace
struct QwenImageNormScaleShiftKernel {
static void
run(tvm::ffi::TensorView y,
@@ -213,4 +211,6 @@ struct QwenImageScaleResidualNormScaleShiftKernel {
}
};
} // namespace sglang_norm_scale_shift
} // namespace norm_scale_shift
} // namespace sglang
python/sglang/kernels/jit/csrc/diffusion/qknorm_rope.cuh
+2 -2
View File
@@ -13,7 +13,7 @@
#include <cstdint>
#include <type_traits>
namespace {
namespace sglang {
struct QKNormRopeParams {
void* __restrict__ q_ptr;
@@ -313,4 +313,4 @@ struct QKNormRopeKernel {
}
};
} // namespace
} // namespace sglang
python/sglang/kernels/jit/csrc/diffusion/residual_gate_add.cuh
+5 -5
View File
@@ -21,9 +21,9 @@
#include <cstdint>
namespace sglang_residual_gate_add {
namespace sglang {
namespace {
namespace residual_gate_add {
constexpr int kBlockSize = 256;
constexpr int kBcastRowsPerBlock = 4;
@@ -303,8 +303,6 @@ inline GateMode validate_residual_gate_add(
return GateMode::kBcastRow;
}
} // namespace
template <typename T>
struct ResidualGateAddKernel {
static void
@@ -314,4 +312,6 @@ struct ResidualGateAddKernel {
}
};
} // namespace sglang_residual_gate_add
} // namespace residual_gate_add
} // namespace sglang
python/sglang/kernels/jit/csrc/diffusion/timestep_embedding.cuh
+5 -5
View File
@@ -17,9 +17,9 @@
#include <cuda_runtime.h>
#include <type_traits>
namespace sglang_timestep_embedding {
namespace sglang {
namespace {
namespace timestep_embedding {
constexpr int kVec = 4; // 16B float vector store
@@ -120,8 +120,6 @@ inline void launch_timestep_embedding(
}
}
} // namespace
template <typename TIn>
void timestep_embedding(
tvm::ffi::TensorView input,
@@ -151,4 +149,6 @@ void timestep_embedding(
launch_timestep_embedding<TIn>(input, output, dim, flip_sin_to_cos, downscale_freq_shift, scale, max_period);
}
} // namespace sglang_timestep_embedding
} // namespace timestep_embedding
} // namespace sglang
python/sglang/kernels/jit/csrc/diffusion/usp_relayout.cuh
+5 -5
View File
@@ -22,9 +22,9 @@
#include <cstdint>
namespace sglang_usp_relayout {
namespace sglang {
namespace {
namespace usp_relayout {
constexpr int kBlockSize = 256;
constexpr int64_t kMaxGrid = 65535;
@@ -135,8 +135,6 @@ __global__ void usp_merge_heads_scalar_kernel(
}
}
} // namespace
template <typename T>
struct UspMergeHeadsKernel {
static void run(tvm::ffi::TensorView out, tvm::ffi::TensorView x) {
@@ -179,4 +177,6 @@ struct UspMergeHeadsKernel {
}
};
} // namespace sglang_usp_relayout
} // namespace usp_relayout
} // namespace sglang
python/sglang/kernels/jit/csrc/distributed/communicator.cuh
+4
View File
@@ -13,6 +13,8 @@
#include <string>
#include <vector>
namespace sglang {
namespace host::distributed {
inline CommunicatorObj::CommunicatorObj(
@@ -111,3 +113,5 @@ inline void register_communicator() {
.def_ro("rank", &Class::rank)
.def("_config", &Class::config);
}
} // namespace sglang
python/sglang/kernels/jit/csrc/distributed/custom_all_reduce.cuh
+2 -2
View File
@@ -30,7 +30,7 @@
#include <string>
#include <variant>
namespace {
namespace sglang {
using device::distributed::Counter, device::distributed::Semaphore;
using host::distributed::CommunicatorRef;
@@ -649,4 +649,4 @@ tvm::ffi::Tensor custom_all_reduce(
return AllReduceKernel<T, kWorldSize, kUsePDL>::run(comm, input, algo, pull_arg);
}
} // namespace
} // namespace sglang
python/sglang/kernels/jit/csrc/distributed/ipc.cuh
+4
View File
@@ -19,6 +19,8 @@
#include <unordered_map>
#include <unordered_set>
namespace sglang {
namespace host::distributed {
struct AllocationRange {
@@ -192,3 +194,5 @@ inline void register_ipc_manager() {
.def("batch_open_handles", &Class::batch_open_handles)
.def("destroy", &Class::destroy);
}
} // namespace sglang
python/sglang/kernels/jit/csrc/distributed/tp_qknorm.cuh
+2 -2
View File
@@ -16,7 +16,7 @@
#include <cstdint>
#include <cstring>
namespace {
namespace sglang {
using device::distributed::Counter;
using host::distributed::CommunicatorObj, host::distributed::CommunicatorRef;
@@ -324,4 +324,4 @@ struct FusedParallelQKNormAcrossHead {
}
};
} // namespace
} // namespace sglang
python/sglang/kernels/jit/csrc/dsa/fused_store_index_cache.cuh
+2 -2
View File
@@ -14,7 +14,7 @@
#include <cstdint>
#include <cuda_fp8.h>
namespace {
namespace sglang {
struct FusedStoreCacheParam {
const void* __restrict__ input;
@@ -121,4 +121,4 @@ struct FusedStoreCacheIndexerKernel {
}
};
} // namespace
} // namespace sglang
python/sglang/kernels/jit/csrc/elementwise/activation.cuh
+2 -2
View File
@@ -13,7 +13,7 @@
#include <limits>
#include <string>
namespace {
namespace sglang {
enum class ActivationKind : uint32_t {
kSiLU,
@@ -284,4 +284,4 @@ struct ActivationKernel {
}
};
} // namespace
} // namespace sglang
python/sglang/kernels/jit/csrc/elementwise/clamp_position.cuh
+2 -2
View File
@@ -9,7 +9,7 @@
#include <cstddef>
#include <cstdint>
namespace {
namespace sglang {
template <typename T>
__global__ void clamp_position_kernel(T* __restrict__ dst, const T* __restrict__ seq_lens, size_t n) {
@@ -51,4 +51,4 @@ struct ClampPosition {
}
};
} // namespace
} // namespace sglang
python/sglang/kernels/jit/csrc/elementwise/concat_mla.cuh
+2 -2
View File
@@ -8,7 +8,7 @@
#include <cuda_bf16.h>
#include <cuda_runtime.h>
namespace {
namespace sglang {
// ======================= Memory Utilities =======================
// Adapted from DeepEP: https://github.com/deepseek-ai/DeepEP/blob/main/csrc/kernels/utils.cuh
@@ -329,4 +329,4 @@ struct ConcatMlaAbsorbQKernel {
}
};
} // namespace
} // namespace sglang
python/sglang/kernels/jit/csrc/elementwise/fused_add_rmsnorm.cuh
+2 -2
View File
@@ -12,7 +12,7 @@
#include <cooperative_groups.h>
#include <type_traits>
namespace {
namespace sglang {
template <typename T, int VEC_SIZE_IN_BYTE>
struct VecTypeTrait;
@@ -194,4 +194,4 @@ struct FusedAddRMSNormKernel {
}
};
} // namespace
} // namespace sglang
python/sglang/kernels/jit/csrc/elementwise/fused_eh_norm.cuh
+2 -2
View File
@@ -10,7 +10,7 @@
#include <tvm/ffi/container/tensor.h>
namespace {
namespace sglang {
struct FusedEHNormParams {
const void* __restrict__ embeds;
@@ -110,4 +110,4 @@ struct FusedEHNormKernel {
}
};
} // namespace
} // namespace sglang
python/sglang/kernels/jit/csrc/elementwise/fused_metadata_copy.cuh
+3 -3
View File
@@ -41,6 +41,8 @@
#include <hip/hip_runtime.h>
#endif
namespace sglang {
// Forward mode enum (must match Python ForwardMode in sglang/srt/layers/attention/dsa_backend.py)
enum ForwardModeEnum { DECODE = 0, TARGET_VERIFY = 1, DRAFT_EXTEND = 2 };
@@ -372,8 +374,6 @@ __global__ void fused_metadata_copy_multi_kernel(const FusedMetadataCopyMultiPar
// Host-side launcher wrappers for JIT compilation
// ============================================================================
namespace {
// Launch configuration constants
constexpr int THREADS_PER_BLOCK = 256;
constexpr int MAX_GRID_SIZE = 1024; // Limit to prevent excessive resource usage
@@ -723,4 +723,4 @@ struct FusedMetadataCopyMultiKernel {
}
};
} // namespace
} // namespace sglang
python/sglang/kernels/jit/csrc/elementwise/fused_qknorm_rope.cuh
+2 -2
View File
@@ -30,7 +30,7 @@
#include <cuda_bf16.h>
#include <cuda_runtime.h>
namespace {
namespace sglang {
// ---------------------------------------------------------------------------
// YaRN-aware frequency computation
@@ -343,4 +343,4 @@ void fused_qk_norm_rope(
rotary_dim);
}
} // namespace
} // namespace sglang
python/sglang/kernels/jit/csrc/elementwise/kvcache.cuh
+2 -2
View File
@@ -11,7 +11,7 @@
#include <cassert>
#include <cstdint>
namespace {
namespace sglang {
struct StoreKVCacheParams {
const void* __restrict__ k;
@@ -318,4 +318,4 @@ struct StoreKVCacheKernel {
}
};
} // namespace
} // namespace sglang
python/sglang/kernels/jit/csrc/elementwise/pos_enc.cuh
+2 -2
View File
@@ -11,7 +11,7 @@
#include <cuda_fp16.h>
#include <cuda_runtime.h>
namespace {
namespace sglang {
template <typename scalar_t, bool IS_NEOX>
inline __device__ void apply_token_rotary_embedding(
@@ -310,4 +310,4 @@ struct RotaryEmbeddingKernel {
}
};
} // namespace
} // namespace sglang
python/sglang/kernels/jit/csrc/elementwise/qknorm.cuh
+2 -2
View File
@@ -15,7 +15,7 @@
#include <cuda_fp16.h>
#include <type_traits>
namespace {
namespace sglang {
struct QKNormParams {
void* __restrict__ q;
@@ -254,4 +254,4 @@ using QKNormKernel = std::conditional_t<
QKNormKernelCTA<kHeadDim, kUsePDL, DType>,
QKNormKernelWarp<kHeadDim, kUsePDL, DType>>;
} // namespace
} // namespace sglang
python/sglang/kernels/jit/csrc/elementwise/qknorm_across_heads.cuh
+2 -2
View File
@@ -12,7 +12,7 @@
#include <cooperative_groups.h>
#include <type_traits>
namespace {
namespace sglang {
template <typename T, int VEC_SIZE_IN_BYTE>
struct VecTypeTrait;
@@ -176,4 +176,4 @@ struct QKNormAcrossHeadsKernel {
}
};
} // namespace
} // namespace sglang
python/sglang/kernels/jit/csrc/elementwise/rmsnorm.cuh
+2 -2
View File
@@ -10,7 +10,7 @@
#include <tvm/ffi/container/tensor.h>
namespace {
namespace sglang {
struct RMSNormParams {
const void* input;
@@ -368,4 +368,4 @@ struct RMSNormHalfKernel {
}
};
} // namespace
} // namespace sglang
python/sglang/kernels/jit/csrc/elementwise/rmsnorm_hf.cuh
+2 -2
View File
@@ -22,7 +22,7 @@
#include <tvm/ffi/container/tensor.h>
namespace {
namespace sglang {
struct RMSNormHFParams {
const void* input;
@@ -250,4 +250,4 @@ struct HFRMSNormKernel {
}
};
} // namespace
} // namespace sglang
python/sglang/kernels/jit/csrc/elementwise/rope.cuh
+2 -2
View File
@@ -11,7 +11,7 @@
#include <numeric>
namespace {
namespace sglang {
struct FusedRopeParams {
void* __restrict__ q_ptr;
@@ -466,4 +466,4 @@ struct FusedRopeKernel {
}
};
} // namespace
} // namespace sglang
python/sglang/kernels/jit/csrc/elementwise/set_mla_kv_buffer.cuh
+2 -2
View File
@@ -35,7 +35,7 @@
#include <cstdint>
namespace {
namespace sglang {
struct SetMlaKVBufferParams {
const void* __restrict__ k_nope;
@@ -213,4 +213,4 @@ struct SetMlaKVBufferKernel {
}
};
} // namespace
} // namespace sglang
python/sglang/kernels/jit/csrc/elementwise/set_mla_kv_concat_q.cuh
+2 -2
View File
@@ -17,7 +17,7 @@
#include <cstdint>
#include <cuda_fp8.h>
namespace {
namespace sglang {
struct SetMlaKVConcatQParams {
// KV scatter side (byte-typed: dtype-agnostic row copies).
@@ -604,4 +604,4 @@ struct SetMlaKVConcatQFp8Kernel {
}
};
} // namespace
} // namespace sglang
python/sglang/kernels/jit/csrc/fast-hadamard-transform/hadamard_jit.cuh
+2 -4
View File
@@ -19,10 +19,8 @@
#include <cstdint>
#include <cstring>
namespace {
namespace sglang {
using ::bf16_t;
using ::fp16_t;
using ::HadamardParamsBase;
constexpr inline int ceil_log2(int val) {
@@ -479,4 +477,4 @@ struct Hadamard40NKernel {
}
};
} // namespace
} // namespace sglang
python/sglang/kernels/jit/csrc/gemm/awq_dequantize.cuh
+4
View File
@@ -6,6 +6,8 @@
#include <sgl_kernel/utils.cuh>
namespace sglang {
namespace device::awq {
template <int lut>
@@ -225,3 +227,5 @@ void awq_dequantize(
static_cast<int>(qweight_cols),
static_cast<int>(qweight_rows));
}
} // namespace sglang
python/sglang/kernels/jit/csrc/gemm/dsv3_fused_a_gemm.cuh
+3 -3
View File
@@ -29,7 +29,7 @@
#include <cuda_bf16.h>
#include <utility>
namespace {
namespace sglang {
using bf16_t = __nv_bfloat16;
@@ -419,7 +419,7 @@ struct MmaComputer {
}
}
}
::arrive_barrier(smem_barrier + 1 + stage_idx * 2);
arrive_barrier(smem_barrier + 1 + stage_idx * 2);
stage_idx += 1;
phase_bit = stage_idx == stage_cnt ? phase_bit ^ 1 : phase_bit;
stage_idx = stage_idx == stage_cnt ? 0 : stage_idx;
@@ -647,4 +647,4 @@ struct DSV3FusedAGemmKernel {
}
};
} // namespace
} // namespace sglang
python/sglang/kernels/jit/csrc/gemm/dsv3_router_gemm.cuh
+2 -2
View File
@@ -31,7 +31,7 @@
#include <type_traits>
namespace {
namespace sglang {
using namespace device;
@@ -181,4 +181,4 @@ struct DSV3RouterGemmKernel {
}
};
} // namespace
} // namespace sglang
python/sglang/kernels/jit/csrc/gemm/fp8_blockwise/fp8_blockwise_scaled_mm_entry.cuh
+4
View File
@@ -15,6 +15,8 @@ limitations under the License.
#include "fp8_blockwise_scaled_mm_sm120.cuh"
namespace sglang {
void fp8_blockwise_scaled_mm(
tvm::ffi::TensorView out,
tvm::ffi::TensorView mat_a,
@@ -23,3 +25,5 @@ void fp8_blockwise_scaled_mm(
tvm::ffi::TensorView scales_b) {
fp8_blockwise_scaled_mm_sm120(out, mat_a, mat_b, scales_a, scales_b);
}
} // namespace sglang
python/sglang/kernels/jit/csrc/gemm/fp8_blockwise/fp8_blockwise_scaled_mm_sm120.cuh
+6 -2
View File
@@ -25,8 +25,6 @@ limitations under the License.
#include <cstdint>
#include <cuda_runtime.h>
using namespace host;
// clang-format off
#include "cutlass/cutlass.h"
#include "cutlass/detail/blockwise_scale_layout.hpp"
@@ -38,6 +36,10 @@ using namespace host;
#include "cutlass/util/packed_stride.hpp"
// clang-format on
namespace sglang {
using namespace host;
#define CUTLASS_CHECK(status) \
{ \
cutlass::Status error = status; \
@@ -500,3 +502,5 @@ inline void fp8_blockwise_scaled_mm_sm120(
}
#endif // defined(CUTLASS_ARCH_MMA_SM120_SUPPORTED) || defined(CUTLASS_ARCH_MMA_SM121_SUPPORTED)
} // namespace sglang
python/sglang/kernels/jit/csrc/gemm/marlin/awq_marlin_repack.cuh
+4
View File
@@ -6,6 +6,8 @@
#include "marlin.cuh"
namespace sglang {
namespace device::marlin {
#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ < 800
@@ -249,3 +251,5 @@ void awq_marlin_repack(
RuntimeCheck(false, "Unsupported repack config: num_bits = ", num_bits);
}
}
} // namespace sglang
python/sglang/kernels/jit/csrc/gemm/marlin/dequant.h
+4
View File
@@ -65,6 +65,8 @@ where `scale_factor * multiplier` can be computed at weight loading.
#include "marlin_dtypes.cuh"
namespace sglang {
namespace device::marlin {
#if !defined(__CUDA_ARCH__) || __CUDA_ARCH__ >= 800
@@ -502,3 +504,5 @@ __device__ inline void dequant_fp8_scales<nv_bfloat162, host::kFE8M0fnu.id()>(in
#endif
} // namespace device::marlin
} // namespace sglang
python/sglang/kernels/jit/csrc/gemm/marlin/gptq_marlin.cuh
+4
View File
@@ -28,6 +28,8 @@
#include "kernel.h"
#include "marlin_template.h"
namespace sglang {
namespace device::marlin {
__global__ void MarlinDefault(MARLIN_KERNEL_PARAMS){};
@@ -999,3 +1001,5 @@ void gptq_marlin_gemm(
use_fp32_reduce,
is_zp_float);
}
} // namespace sglang
python/sglang/kernels/jit/csrc/gemm/marlin/gptq_marlin_repack.cuh
+4
View File
@@ -27,6 +27,8 @@
#include "marlin.cuh"
namespace sglang {
namespace device::marlin {
#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ < 800
@@ -359,4 +361,6 @@ void gptq_marlin_repack(
}
}
} // namespace sglang
#undef CALL_IF_REPACK
python/sglang/kernels/jit/csrc/gemm/marlin/kernel.h
+4
View File
@@ -10,6 +10,8 @@
const int *__restrict__ g_idx, int num_groups, int prob_m, int prob_n, int prob_k, int lda, int *locks, \
bool use_atomic_add, bool use_fp32_reduce, int max_shared_mem
namespace sglang {
namespace device::marlin {
template <
typename scalar_t, // compute dtype, half or nv_float16
@@ -31,3 +33,5 @@ template <
__global__ void Marlin(MARLIN_KERNEL_PARAMS);
} // namespace device::marlin
} // namespace sglang
python/sglang/kernels/jit/csrc/gemm/marlin/marlin.cuh
+4
View File
@@ -4,6 +4,8 @@
#include <iostream>
namespace sglang {
namespace device::marlin {
// Marlin params
@@ -81,3 +83,5 @@ __device__ inline void cp_async_wait() {
#endif
} // namespace device::marlin
} // namespace sglang
python/sglang/kernels/jit/csrc/gemm/marlin/marlin_dtypes.cuh
+4
View File
@@ -4,6 +4,8 @@
#include "marlin.cuh"
namespace sglang {
namespace device::marlin {
template <typename scalar_t>
@@ -74,4 +76,6 @@ class ScalarType<bf16_t> {
} // namespace device::marlin
} // namespace sglang
#endif
python/sglang/kernels/jit/csrc/gemm/marlin/marlin_template.h
+4
View File
@@ -29,6 +29,8 @@
std::is_same<scalar_t, half>::value || std::is_same<scalar_t, nv_bfloat16>::value, \
"only float16 and bfloat16 is supported");
namespace sglang {
namespace device::marlin {
#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ < 800
@@ -1618,3 +1620,5 @@ __global__ void Marlin(
} // namespace device::marlin
#endif
} // namespace sglang
python/sglang/kernels/jit/csrc/gemm/marlin_moe/kernel.h
+4
View File
@@ -13,6 +13,8 @@
bool mul_topk_weights, bool is_ep, int num_groups, int prob_m, int prob_n, int prob_k, int *locks, \
bool has_bias, bool use_atomic_add, bool use_fp32_reduce, int max_shared_mem
namespace sglang {
namespace device::marlin_moe {
template <
typename scalar_t, // compute dtype, half or nv_float16
@@ -37,3 +39,5 @@ template <
__global__ void Marlin(MARLIN_KERNEL_PARAMS);
} // namespace device::marlin_moe
} // namespace sglang
python/sglang/kernels/jit/csrc/gemm/marlin_moe/marlin_template.h
+4
View File
@@ -31,6 +31,8 @@
std::is_same<scalar_t, half>::value || std::is_same<scalar_t, nv_bfloat16>::value, \
"only float16 and bfloat16 is supported");
namespace sglang {
namespace device::marlin_moe {
using namespace device::marlin;
@@ -1906,3 +1908,5 @@ __global__ void Marlin(
} // namespace device::marlin_moe
#endif
} // namespace sglang
python/sglang/kernels/jit/csrc/gemm/marlin_moe/moe_wna16_marlin.cuh
+4
View File
@@ -28,6 +28,8 @@
#include "kernel.h"
#include "marlin_template.h"
namespace sglang {
namespace device::marlin_moe {
__global__ void MarlinDefault(MARLIN_KERNEL_PARAMS){};
@@ -1114,3 +1116,5 @@ void moe_wna16_marlin_gemm(
use_fp32_reduce,
is_zp_float);
}
} // namespace sglang
python/sglang/kernels/jit/csrc/gemm/per_tensor_quant_fp8.cuh
+2 -2
View File
@@ -12,7 +12,7 @@
#include <cstddef>
#include <cstdint>
namespace {
namespace sglang {
constexpr size_t kBlockSize = 256;
@@ -159,4 +159,4 @@ void per_tensor_absmax_fp8(tvm::ffi::TensorView input, tvm::ffi::TensorView outp
DType>(input, output_s, output_s);
}
} // namespace
} // namespace sglang
python/sglang/kernels/jit/csrc/gemm/per_token_group_quant.cuh
+15 -15
View File
@@ -16,9 +16,9 @@
#include <cstdint>
#include <cuda_fp8.h>
namespace {
namespace sglang {
namespace details {
namespace detail {
SGL_DEVICE float silu(const float val) {
// silu(x) = x * sigmoid(x)
@@ -190,12 +190,12 @@ struct TensorArgs {
}
};
} // namespace details
} // namespace detail
struct QuantKernelParams {
details::TensorArgs input;
details::TensorArgs output;
details::ScaleStoreArgs scale;
detail::TensorArgs input;
detail::TensorArgs output;
detail::ScaleStoreArgs scale;
uint32_t num_tokens; // tokens_pad for the masked kernel
uint32_t hidden_size; // = num_groups * kGroupSize
};
@@ -248,9 +248,9 @@ struct QuantTrait {
using T = InputType;
using T2 = packed_t<T>;
using Q = QuantType;
using WTrait = details::WeightTrait<Q>;
using WTrait = detail::WeightTrait<Q>;
using Q2 = typename WTrait::packed2_t;
using in_vec_t = details::Vec32B<T2>;
using in_vec_t = detail::Vec32B<T2>;
using out_vec_t = AlignedVector<Q2, kVecSize / 2>;
constexpr float kMaxValue = WTrait::kMaxValue;
constexpr float kMaxValueInv = 1.f / kMaxValue;
@@ -268,7 +268,7 @@ struct QuantTrait {
#pragma unroll
for (uint32_t i = 0; i < kVecSize / 2; ++i) {
const auto gate = cast<float2>(in[i]);
const auto act = cast<T2>(float2{details::silu(gate.x), details::silu(gate.y)});
const auto act = cast<T2>(float2{detail::silu(gate.x), detail::silu(gate.y)});
in[i] = __hmul2(act, up[i]);
}
}
@@ -284,7 +284,7 @@ struct QuantTrait {
const float raw_scale = amax * kMaxValueInv; // the dequant scale the GEMM consumes
out_vec_t out;
details::scale_t<kUe8m0> scale_inv;
detail::scale_t<kUe8m0> scale_inv;
if constexpr (kUe8m0) {
// ue8m0 scale: pow-2 quant multiplier is exact in float16/bfloat16 type
static_assert(std::is_same_v<Q, fp8_e4m3_t>, "ue8m0 scales imply fp8 quantization");
@@ -309,7 +309,7 @@ struct QuantTrait {
const float2 quant_scale2 = {quant_scale, quant_scale};
#pragma unroll
for (uint32_t i = 0; i < kVecSize / 2; ++i) {
out[i] = WTrait::quant(details::mul2(cast<float2>(in[i]), quant_scale2));
out[i] = WTrait::quant(detail::mul2(cast<float2>(in[i]), quant_scale2));
}
}
@@ -440,7 +440,7 @@ QuantHostContext<Trait> build_quant_context( //
if constexpr (Trait::kUe8m0) {
CHECK_HOST(Trait::kAligned == (num_groups % 4 == 0));
}
auto scale_args = details::ScaleStoreArgs{
auto scale_args = detail::ScaleStoreArgs{
.base = output_s.data_ptr(),
.expert_stride = static_cast<uint32_t>(kMasked ? output_s.stride(0) : 0),
.token_stride = static_cast<uint32_t>(output_s.stride(-2)),
@@ -465,12 +465,12 @@ QuantHostContext<Trait> build_quant_context( //
}
// The scale store indexes with uint32 strides; guard against overflow.
scale_args.check_overflow(num_experts, num_tokens);
const auto input_args = details::TensorArgs{
const auto input_args = detail::TensorArgs{
.ptr = input.data_ptr(),
.expert_stride = kMasked ? input.stride(0) : 0,
.token_stride = input.stride(-2),
};
const auto output_args = details::TensorArgs{
const auto output_args = detail::TensorArgs{
.ptr = output_q.data_ptr(),
.expert_stride = kMasked ? output_q.stride(0) : 0,
.token_stride = output_q.stride(-2),
@@ -566,4 +566,4 @@ struct PerTokenGroupQuantMaskedKernel {
}
};
} // namespace
} // namespace sglang
python/sglang/kernels/jit/csrc/gemm/per_token_group_quant_8bit_v2.cuh
+2 -2
View File
@@ -17,7 +17,7 @@
#include <cuda_fp8.h>
#include <type_traits>
namespace {
namespace sglang {
constexpr float LOCAL_ABSMAX_ABS = 1e-10f;
constexpr uint32_t INPUT_PRIMARY_VEC_NUM_BYTES = 32;
@@ -536,4 +536,4 @@ struct PerTokenGroupQuant8bitV2Kernel {
}
};
} // namespace
} // namespace sglang
python/sglang/kernels/jit/csrc/gemm/tiny_gemm.cuh
+2 -4
View File
@@ -143,10 +143,6 @@ __global__ __launch_bounds__(N_SPLIT* K / kTinyKGemmVecSize, 1) // control the
PDLTriggerSecondary<kUsePDL>();
}
} // namespace sglang
using namespace sglang;
template <uint32_t N, uint32_t K, uint32_t kMaxM, uint32_t N_SPLIT, typename OutT, bool kUsePDL>
struct TinyNGemmKernel {
static constexpr uint32_t kBlockSize = K / kTinyNGemmVecSize;
@@ -229,3 +225,5 @@ struct TinyKGemmKernel {
x_stride);
}
};
} // namespace sglang
python/sglang/kernels/jit/csrc/hisparse.cuh
+2 -2
View File
@@ -12,7 +12,7 @@
#include <stdint.h>
#include <string>
namespace {
namespace sglang {
#ifdef USE_ROCM
constexpr int WARP_SIZE = 64;
@@ -670,4 +670,4 @@ void load_cache_to_device_buffer(
}
}
} // namespace
} // namespace sglang
python/sglang/kernels/jit/csrc/inkling/causal_conv1d.cuh
+2 -2
View File
@@ -26,7 +26,7 @@
#include <cstdint>
#include <cuda_bf16.h>
namespace {
namespace sglang {
struct CausalConv1dParams {
const void* __restrict__ x; // [T, D]
@@ -214,4 +214,4 @@ struct CausalConv1dKernel {
}
};
} // namespace
} // namespace sglang
python/sglang/kernels/jit/csrc/inkling/draft_extend_sconv.cuh
+2 -2
View File
@@ -21,7 +21,7 @@
#include <cstdint>
#include <cuda_bf16.h>
namespace {
namespace sglang {
struct DraftExtendParams {
const void* __restrict__ hidden; // [B*T, D], channel-contiguous
@@ -144,4 +144,4 @@ struct DraftExtendSconvKernel {
}
};
} // namespace
} // namespace sglang
python/sglang/kernels/jit/csrc/inkling/fused_decode_update.cuh
+2 -2
View File
@@ -23,7 +23,7 @@
#include <cstdint>
#include <cuda_bf16.h>
namespace {
namespace sglang {
struct DecodeUpdateParams {
const void* __restrict__ x; // [T, D], channel-contiguous
@@ -184,4 +184,4 @@ struct FusedDecodeUpdateKernel {
}
};
} // namespace
} // namespace sglang
python/sglang/kernels/jit/csrc/inkling/gather_scatter_sconv.cuh
+2 -2
View File
@@ -15,7 +15,7 @@
#include <cstdint>
#include <cuda_bf16.h>
namespace {
namespace sglang {
struct GatherScatterParams {
const void* __restrict__ hidden; // [T, D], channel-contiguous
@@ -106,4 +106,4 @@ struct GatherScatterSconvKernel {
}
};
} // namespace
} // namespace sglang
python/sglang/kernels/jit/csrc/inkling/inkling_all_reduce.cuh
+2 -2
View File
@@ -44,7 +44,7 @@
#include <type_traits>
#include <unordered_map>
namespace {
namespace sglang {
template <typename DType, uint32_t kNumGPU>
struct InklingAllReduceTrait {
@@ -669,4 +669,4 @@ void inkling_multimem_full_oneshot(
n);
}
} // namespace
} // namespace sglang
python/sglang/kernels/jit/csrc/inkling/inkling_ar_barrier.cuh
+4
View File
@@ -16,6 +16,8 @@
#include <cstdint>
namespace sglang {
namespace inkling_ar {
constexpr uint32_t kLeaderStateWords = 8;
@@ -191,3 +193,5 @@ block_system_barrier(uint32_t* __restrict__ st, void* const* __restrict__ flag_p
}
} // namespace inkling_ar
} // namespace sglang
python/sglang/kernels/jit/csrc/inkling/inkling_ar_fused_decode.cuh
+2 -2
View File
@@ -49,7 +49,7 @@
#include <cuda_bf16.h>
#include <type_traits>
namespace {
namespace sglang {
constexpr int kPadSlot = -1;
constexpr uint32_t kVecElems = 8; // bf16x8 = 16 B
@@ -827,4 +827,4 @@ struct ArSconvNormVerifyKernel {
}
};
} // namespace
} // namespace sglang
python/sglang/kernels/jit/csrc/inkling/inkling_ar_scattered_sconv.cuh
+2 -2
View File
@@ -55,7 +55,7 @@
#include <type_traits>
#include <unordered_map>
namespace {
namespace sglang {
constexpr uint32_t kSsVecElems = 8; // bf16x8 = 16 B
constexpr int kSsPadSlot = -1;
@@ -2031,4 +2031,4 @@ struct SsconvNormDecodeKernel {
}
};
} // namespace
} // namespace sglang
python/sglang/kernels/jit/csrc/inkling/inkling_attn_prologue_fused.cuh
+2 -2
View File
@@ -33,7 +33,7 @@
#include <cuda_fp8.h>
#include <type_traits>
namespace {
namespace sglang {
constexpr int kPadSlot = -1;
constexpr uint32_t kVecElems = 8;
@@ -1438,4 +1438,4 @@ struct AttnPrologueExtendKernel {
}
};
} // namespace
} // namespace sglang
python/sglang/kernels/jit/csrc/inkling/inkling_rel_proj.cuh
+2 -2
View File
@@ -23,7 +23,7 @@
#include <cstdint>
namespace {
namespace sglang {
constexpr uint32_t kRpVec = 8; // bf16x8 = 16 B
constexpr uint32_t kRpBlock = 256;
@@ -142,4 +142,4 @@ void rel_proj_small_t(
}
}
} // namespace
} // namespace sglang
python/sglang/kernels/jit/csrc/inkling/inkling_row_scale.cuh
+2 -2
View File
@@ -19,7 +19,7 @@
#include <cstdint>
namespace {
namespace sglang {
constexpr uint32_t kRsVec = 8; // bf16x8 = 16 B
constexpr uint32_t kRsBlock = 256;
@@ -113,4 +113,4 @@ void row_compact(tvm::ffi::TensorView x, tvm::ffi::TensorView out) {
row_scale_launch<kUsePDL, false>(x, nullptr, out, R, N, dev);
}
} // namespace
} // namespace sglang
python/sglang/kernels/jit/csrc/inkling/update_sconv_cache.cuh
+2 -2
View File
@@ -23,7 +23,7 @@
#include <cstdint>
#include <cuda_bf16.h>
namespace {
namespace sglang {
struct UpdateSconvParams {
const void* __restrict__ x; // [T, D], channel-contiguous
@@ -135,4 +135,4 @@ struct UpdateSconvCacheKernel {
}
};
} // namespace
} // namespace sglang
python/sglang/kernels/jit/csrc/kimi_k3/attn_res/fused_tma.cuh
+26 -30
View File
@@ -20,6 +20,8 @@
// Local PTX primitives (mbarrier / bulk TMA / tcgen05 / warp-group sync)
namespace sglang {
namespace ptx {
// ---- bulk 1D TMA (PTX ISA §9.7.9.25) ---------------------------------------
@@ -160,8 +162,6 @@ static SGL_DEVICE void tcgen05_wait_st() {
} // namespace ptx
namespace sglang {
struct AttnResTMAParams {
const bf16_t* __restrict__ prefix_sum; // [T, H]
const bf16_t* __restrict__ bank; // [T, NB_total, H]
@@ -275,19 +275,19 @@ SGL_DEVICE void KimiK3AttnResTrait<kDim_, kNumBankRows_, kChunkRows_, kConsumerR
const auto lane_id = tx % kWarpThreads;
if (warp_id == 0 && lane_id < kNumStages) {
::ptx::mbar_init(&smem->bar_full[lane_id], 1);
::ptx::mbar_init(&smem->bar_free[lane_id], kNumConsumerWarps * kWarpThreads);
::ptx::fence_mbarrier_init();
ptx::mbar_init(&smem->bar_full[lane_id], 1);
ptx::mbar_init(&smem->bar_free[lane_id], kNumConsumerWarps * kWarpThreads);
ptx::fence_mbarrier_init();
} else if (warp_id == 1) {
::ptx::tcgen05_alloc(::ptx::to_shared(&smem->tmem_base), kTmemCols);
::ptx::tcgen05_relinquish();
ptx::tcgen05_alloc(ptx::to_shared(&smem->tmem_base), kTmemCols);
ptx::tcgen05_relinquish();
}
__syncthreads();
if (warp_id >= kNumConsumerWarps) { // producer warp (group); first warp works
if constexpr (kConsumerRegs > 0) ::ptx::setmaxnreg_dec<kProducerRegs>();
if constexpr (kConsumerRegs > 0) ptx::setmaxnreg_dec<kProducerRegs>();
// TODO: reduce the register usage
if (warp_id == kNumConsumerWarps && ::ptx::elect_one()) {
if (warp_id == kNumConsumerWarps && ptx::elect_one()) {
uint32_t global_chunks = 0;
constexpr uint32_t kRowBytes = kDim * sizeof(bf16_t);
for (auto token = blockIdx.x; token < params.num_tokens; token += gridDim.x) {
@@ -298,10 +298,10 @@ SGL_DEVICE void KimiK3AttnResTrait<kDim_, kNumBankRows_, kChunkRows_, kConsumerR
const auto slot = global_chunks % kNumStages;
const auto phase = (global_chunks / kNumStages) & 1;
if (global_chunks >= kNumStages) {
::ptx::mbar_wait_parity(&smem->bar_free[slot], phase ^ 1);
ptx::mbar_wait_parity(&smem->bar_free[slot], phase ^ 1);
}
// One barrier per chunk; each row still gets its own bulk copy.
::ptx::mbar_arrive_expect_tx(&smem->bar_full[slot], an * kRowBytes);
ptx::mbar_arrive_expect_tx(&smem->bar_full[slot], an * kRowBytes);
#pragma unroll
for (uint32_t r = 0; r < an; ++r) {
const auto row = base_row + r;
@@ -310,14 +310,14 @@ SGL_DEVICE void KimiK3AttnResTrait<kDim_, kNumBankRows_, kChunkRows_, kConsumerR
// Only prefix_sum is written by the immediately-preceding kernel;
// one wait before the first token's prefix load covers the rest.
if (token == blockIdx.x && row == kNumRows) PDLWaitPrimary<true>();
::ptx::cp_async_bulk_1d_load(&smem->buf[slot][r], src, kRowBytes, &smem->bar_full[slot]);
ptx::cp_async_bulk_1d_load(&smem->buf[slot][r], src, kRowBytes, &smem->bar_full[slot]);
}
}
}
PDLTriggerSecondary<true>();
}
} else { // 2 consumer warp groups; one chunk per rendezvous
if constexpr (kConsumerRegs > 0) ::ptx::setmaxnreg_inc<kConsumerRegs>();
if constexpr (kConsumerRegs > 0) ptx::setmaxnreg_inc<kConsumerRegs>();
const auto group = warp_id / (kNumConsumerWarps / kNumGroups);
const auto tid_in_group = tx % kGroupThreads;
const auto tmem_cw = smem->tmem_base + group * kTmemColsPerGroup;
@@ -338,8 +338,7 @@ SGL_DEVICE void KimiK3AttnResTrait<kDim_, kNumBankRows_, kChunkRows_, kConsumerR
}
#pragma unroll
for (uint32_t si = 0; si < kSlicesPerGroup; ++si) {
::ptx::tcgen05_st_32x32b_x8(
tmem_cw + si * kVecElems, reinterpret_cast<const uint32_t*>(&staged[si * kVecElems]));
ptx::tcgen05_st_32x32b_x8(tmem_cw + si * kVecElems, reinterpret_cast<const uint32_t*>(&staged[si * kVecElems]));
}
#pragma unroll
for (uint32_t si = 0; si < kSlicesPerGroup; ++si) {
@@ -353,10 +352,9 @@ SGL_DEVICE void KimiK3AttnResTrait<kDim_, kNumBankRows_, kChunkRows_, kConsumerR
}
#pragma unroll
for (uint32_t si = 0; si < kSlicesPerGroup; ++si) {
::ptx::tcgen05_st_32x32b_x8(
tmem_ow + si * kVecElems, reinterpret_cast<const uint32_t*>(&staged[si * kVecElems]));
ptx::tcgen05_st_32x32b_x8(tmem_ow + si * kVecElems, reinterpret_cast<const uint32_t*>(&staged[si * kVecElems]));
}
::ptx::tcgen05_wait_st();
ptx::tcgen05_wait_st();
}
uint32_t global_chunks = 0; // mirrors the producer's chunk counter
@@ -372,7 +370,7 @@ SGL_DEVICE void KimiK3AttnResTrait<kDim_, kNumBankRows_, kChunkRows_, kConsumerR
const uint32_t an = (kNumRows + 1 - base_row) < kChunkRows ? (kNumRows + 1 - base_row) : kChunkRows;
const auto slot = global_chunks % kNumStages;
const auto phase = (global_chunks / kNumStages) & 1;
::ptx::mbar_wait_parity(&smem->bar_full[slot], phase);
ptx::mbar_wait_parity(&smem->bar_full[slot], phase);
// Score pass: the cw slice is loaded once and reused across the
// chunk's rows; each row's 16B slices land in registers. rms/dot
@@ -386,7 +384,7 @@ SGL_DEVICE void KimiK3AttnResTrait<kDim_, kNumBankRows_, kChunkRows_, kConsumerR
const auto tile = si * kNumGroups + group;
if (tile >= kNumTiles) continue;
float q[kVecElems];
::ptx::tcgen05_ld_32x32b_x8(tmem_cw + si * kVecElems, reinterpret_cast<uint32_t*>(q));
ptx::tcgen05_ld_32x32b_x8(tmem_cw + si * kVecElems, reinterpret_cast<uint32_t*>(q));
const auto* q2 = reinterpret_cast<const float2*>(q);
const auto offset = tile * kTile + tid_in_group * kVecElems;
#pragma unroll
@@ -403,7 +401,7 @@ SGL_DEVICE void KimiK3AttnResTrait<kDim_, kNumBankRows_, kChunkRows_, kConsumerR
}
}
}
::ptx::mbar_arrive(&smem->bar_free[slot]);
ptx::mbar_arrive(&smem->bar_free[slot]);
// Fused bank write: the prefix row (last row of the last chunk) is
// already in registers; snapshot it to bank row nvb with plain
@@ -440,7 +438,7 @@ SGL_DEVICE void KimiK3AttnResTrait<kDim_, kNumBankRows_, kChunkRows_, kConsumerR
smem->warp_dot[warp_id][r] = acc_dot[r];
}
}
::ptx::named_barrier_sync(kConsumerBarId, kNumConsumerThreads);
ptx::named_barrier_sync(kConsumerBarId, kNumConsumerThreads);
// Lane r totals row r, then broadcasts: an*16 smem loads per warp
// instead of per thread.
float lane_logit = 0.f;
@@ -519,7 +517,7 @@ SGL_DEVICE void KimiK3AttnResTrait<kDim_, kNumBankRows_, kChunkRows_, kConsumerR
}
float acc_sq = warp::reduce_sum(acc_sq2.x + acc_sq2.y);
if (lane_id == 0) smem->warp_ssq[warp_id] = acc_sq;
::ptx::named_barrier_sync(kConsumerBarId, kNumConsumerThreads);
ptx::named_barrier_sync(kConsumerBarId, kNumConsumerThreads);
float total_sq = 0.f;
#pragma unroll
for (uint32_t w = 0; w < kNumConsumerWarps; ++w) {
@@ -534,7 +532,7 @@ SGL_DEVICE void KimiK3AttnResTrait<kDim_, kNumBankRows_, kChunkRows_, kConsumerR
const auto tile = si * kNumGroups + group;
if (tile >= kNumTiles) continue;
float q[kVecElems];
::ptx::tcgen05_ld_32x32b_x8(tmem_ow + si * kVecElems, reinterpret_cast<uint32_t*>(q));
ptx::tcgen05_ld_32x32b_x8(tmem_ow + si * kVecElems, reinterpret_cast<uint32_t*>(q));
const auto* q2 = reinterpret_cast<const float2*>(q);
row_vec_t out_vec;
#pragma unroll
@@ -552,9 +550,9 @@ SGL_DEVICE void KimiK3AttnResTrait<kDim_, kNumBankRows_, kChunkRows_, kConsumerR
}
}
}
::ptx::named_barrier_sync(kConsumerBarId, kNumConsumerThreads);
ptx::named_barrier_sync(kConsumerBarId, kNumConsumerThreads);
if (warp_id == 1) {
::ptx::tcgen05_dealloc(smem->tmem_base, kTmemCols);
ptx::tcgen05_dealloc(smem->tmem_base, kTmemCols);
}
}
}
@@ -677,10 +675,6 @@ __global__ void __launch_bounds__(Trait::kNumThreads, kOccupancy)
;
}
}
} // namespace sglang
using namespace sglang;
using host::distributed::CommunicatorRef;
// Host launcher: constexpr kernel table over nvb.
@@ -944,3 +938,5 @@ struct AttnResFusedTmaKernel {
LaunchKernel(grid, kNumThreads, device.unwrap(), kSmemBytes).enable_pdl(true)(kAgTable[nvb], params);
}
};
} // namespace sglang
python/sglang/kernels/jit/csrc/kimi_k3/comm/ar_fusion.cuh
+2 -4
View File
@@ -613,10 +613,6 @@ __launch_bounds__(kNormRowVecs, 1) void all_reduce_pull_norm_kernel(const __grid
pull_barrier_exit<kUsePDL>(params, barrier_window);
}
} // namespace sglang
using namespace sglang;
// Host entry points
template <uint32_t kWorldSize, bool kUsePDL>
@@ -906,3 +902,5 @@ struct AllReduceFusionKernel {
launch_pull_norm(params, num_blocks, unroll, input.device());
}
};
} // namespace sglang
python/sglang/kernels/jit/csrc/kimi_k3/comm/gemm_ag.cuh
+2 -3
View File
@@ -198,9 +198,6 @@ __global__ void spin_add3_kernel(const __grid_constant__ ConsumerParams params)
}
} // namespace gemm_ag
} // namespace sglang
using namespace sglang;
using host::distributed::CommunicatorRef;
// Host entry point (tiny_gemm style: one GEMV instantiation per M in
@@ -301,3 +298,5 @@ struct GEMMAGKernel {
.enable_pdl(kUsePDL)(kernel, consumer_params);
}
};
} // namespace sglang
python/sglang/kernels/jit/csrc/kimi_k3/comm/gemm_ar.cuh
+8
View File
@@ -32,6 +32,8 @@
// exactly one consumer. Anything cute already provides goes through cute
// (`set_block_rank` below, the tensor-map driver wrapper in `w_maps`).
namespace sglang {
namespace ptx {
// ---- generic → shared address conversion (PTX ISA §10.4) --------------------
@@ -1356,6 +1358,8 @@ struct Launcher {
} // namespace oproj_ar
} // namespace sglang
// ================= sglang tvm-ffi adapter =================
#include <sgl_kernel/tensor.h> // For TensorMatcher, SymbolicSize, SymbolicDevice
@@ -1369,6 +1373,8 @@ struct Launcher {
#include <mutex>
#include <unordered_map>
namespace sglang {
namespace oproj_ar_ffi {
using namespace oproj_ar;
@@ -1567,3 +1573,5 @@ struct GemmArKernel {
} // namespace oproj_ar_ffi
using oproj_ar_ffi::GemmArKernel;
} // namespace sglang
python/sglang/kernels/jit/csrc/kimi_k3/comm/ptx_sys.cuh
+4
View File
@@ -14,6 +14,8 @@
#include <cstdint>
namespace sglang {
namespace device::distributed {
// Peer-visible flag increment. `.sys` scope, relaxed: ordering is established by
@@ -66,3 +68,5 @@ SGL_DEVICE void multimem_red_add_release(uint32_t* mc_flag) {
}
} // namespace device::distributed
} // namespace sglang
python/sglang/kernels/jit/csrc/kimi_k3/comm/sp_collective.cuh
+6 -4
View File
@@ -20,7 +20,9 @@
#include "../../distributed/custom_all_reduce.cuh"
namespace sglang::sp_collective {
namespace sglang {
namespace sp_collective {
using device::distributed::Counter;
using device::distributed::Semaphore;
@@ -292,9 +294,7 @@ __global__ void reduce_scatter_pull_kernel(const __grid_constant__ Params params
}
}
} // namespace sglang::sp_collective
using namespace sglang;
} // namespace sp_collective
using host::distributed::CommunicatorRef;
template <uint32_t kWorldSize, bool kUsePDL>
@@ -437,3 +437,5 @@ struct SPCollectiveKernel {
host::LaunchKernel(num_blocks, block_size, input.device()).enable_pdl(kUsePDL)(kernel, params);
}
};
} // namespace sglang
python/sglang/kernels/jit/csrc/kimi_k3/mla_output_gate.cuh
+2 -2
View File
@@ -14,7 +14,7 @@
#include <cstdint>
namespace {
namespace sglang {
struct MlaOutputGateParams {
const bf16_t* __restrict__ x; // [N] contiguous (flattened [T, H])
@@ -81,4 +81,4 @@ struct MlaOutputGateKernel {
}
};
} // namespace
} // namespace sglang
python/sglang/kernels/jit/csrc/kimi_k3/situ_and_mul.cuh
+4 -9
View File
@@ -44,10 +44,6 @@ SGL_DEVICE float situ_activate(float g, float u, float beta, float inv_beta, flo
} // namespace kimi_k3
} // namespace sglang
namespace {
// SiTU (SoftCap-GLU) activation:
// gate_out = beta * tanh(gate / beta) * sigmoid(gate)
// up_out = linear_beta * tanh(up / linear_beta)
@@ -104,8 +100,7 @@ __global__ void situ_and_mul_kernel(const __grid_constant__ SituAndMulParams par
const float g = cast<fp32_t>(gate[i]);
const float u = cast<fp32_t>(up[i]);
out[i] =
cast<T>(sglang::kimi_k3::situ_activate<kHasLinearBeta>(g, u, beta, inv_beta, linear_beta, inv_linear_beta));
out[i] = cast<T>(kimi_k3::situ_activate<kHasLinearBeta>(g, u, beta, inv_beta, linear_beta, inv_linear_beta));
}
store_as<vec_t>(params.out, out, output_offset);
@@ -218,8 +213,8 @@ situ_and_mul(DType2 gate, DType2 up, float beta, float inv_beta, float linear_be
const auto [g0, g1] = cast<fp32x2_t>(gate);
const auto [u0, u1] = cast<fp32x2_t>(up);
// kHasLinearBeta=true: this path always softcaps the up operand, as before.
const float val0 = sglang::kimi_k3::situ_activate<true>(g0, u0, beta, inv_beta, linear_beta, inv_linear_beta);
const float val1 = sglang::kimi_k3::situ_activate<true>(g1, u1, beta, inv_beta, linear_beta, inv_linear_beta);
const float val0 = kimi_k3::situ_activate<true>(g0, u0, beta, inv_beta, linear_beta, inv_linear_beta);
const float val1 = kimi_k3::situ_activate<true>(g1, u1, beta, inv_beta, linear_beta, inv_linear_beta);
if constexpr (kPrecise) {
return {val0, val1};
} else {
@@ -449,4 +444,4 @@ struct SituAndMulMaskedPostQuantKernel {
}
};
} // namespace
} // namespace sglang
python/sglang/kernels/jit/csrc/kv_canary/canary_common.cuh
+4
View File
@@ -5,6 +5,8 @@
#include "consts.cuh"
#include <cstdint>
namespace sglang {
namespace canary {
// Device-side handle for one real-KV source.
@@ -134,3 +136,5 @@ SGL_DEVICE uint64_t compute_slot_hash(const uint8_t* canary_buf, int64_t slot_st
}
} // namespace canary
} // namespace sglang

Some files were not shown because too many files have changed in this diff Show More