[jit_kernel] Move JIT kernels into namespace sglang

.claude/skills/add-jit-kernel/SKILL.md

@@ -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

@@ -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

@@ -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

@@ -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

@@ -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

@@ -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

@@ -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

@@ -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

@@ -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

@@ -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

@@ -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

@@ -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

@@ -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

@@ -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

@@ -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

@@ -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

@@ -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

@@ -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

@@ -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

@@ -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