[Feature] JIT Fused QK norm + qk norm clean up (#15835)

python/sglang/jit_kernel/include/sgl_kernel/runtime.cuh

@@ -0,0 +1,26 @@
+#pragma once
+
+#include <sgl_kernel/utils.cuh>
+
+#include <cstddef>
+#include <cstdint>
+
+namespace host::runtime {
+
+// Return the maximum number of active blocks per SM for the given kernel
+template <typename T>
+inline auto get_blocks_per_sm(T&& kernel, int32_t block_dim, std::size_t dynamic_smem = 0) -> uint32_t {
+  int num_blocks_per_sm = 0;
+  RuntimeDeviceCheck(
+      cudaOccupancyMaxActiveBlocksPerMultiprocessor(&num_blocks_per_sm, kernel, block_dim, dynamic_smem));
+  return static_cast<uint32_t>(num_blocks_per_sm);
+}
+
+// Return the number of SMs for the given device
+inline auto get_sm_count(int device_id) -> uint32_t {
+  cudaDeviceProp device_prop;
+  RuntimeDeviceCheck(cudaGetDeviceProperties(&device_prop, device_id));
+  return static_cast<uint32_t>(device_prop.multiProcessorCount);
+}
+
+}  // namespace host::runtime

python/sglang/jit_kernel/include/sgl_kernel/tensor.h

@@ -153,6 +153,11 @@ inline auto& operator<<(std::ostream& os, PrintAbleSpan<T> span) {
 
 }  // namespace details
 
+template <typename T>
+inline bool is_type(DLDataType dtype) {
+  return dtype == details::dtype_trait<T>::value;
+}
+
 struct SymbolicSize {
  public:
   SymbolicSize(std::string_view annotation = {}) : m_value(details::kNullSize), m_annotation(annotation) {}
@@ -259,6 +264,11 @@ struct SymbolicDType {
     }
   }
 
+  template <typename T>
+  auto is_type() const -> bool {
+    return ::host::is_type<T>(m_value);
+  }
+
  private:
   auto m_check(DLDataType value) const -> bool {
     return stdr::empty(m_options) || (stdr::find(m_options, value) != stdr::end(m_options));

python/sglang/jit_kernel/include/sgl_kernel/utils.cuh

@@ -79,6 +79,24 @@ struct device_vec {
   T data[N];
 };
 
+template <bool kUsePDL>
+__forceinline__ __device__ void PDLWaitPrimary() {
+#ifndef USE_ROCM
+  if constexpr (kUsePDL) {
+    asm volatile("griddepcontrol.wait;");
+  }
+#endif
+}
+
+template <bool kUsePDL>
+__forceinline__ __device__ void PDLTriggerSecondary() {
+#ifndef USE_ROCM
+  if constexpr (kUsePDL) {
+    asm volatile("griddepcontrol.launch_dependents;");
+  }
+#endif
+}
+
 }  // namespace device
 
 namespace host {
@@ -120,6 +138,18 @@ struct LaunchKernel {
     return static_cast<cudaStream_t>(::TVMFFIEnvGetStream(device.device_type, device.device_id));
   }
 
+  auto enable_pdl(bool enabled = true) -> LaunchKernel& {
+    if (enabled) {
+      m_attrs[0].id = cudaLaunchAttributeProgrammaticStreamSerialization;
+      m_attrs[0].val.programmaticStreamSerializationAllowed = true;
+      m_config.numAttrs = 1;
+      m_config.attrs = m_attrs;
+    } else {
+      m_config.numAttrs = 0;
+    }
+    return *this;
+  }
+
   template <typename T, typename... Args>
   auto operator()(T&& kernel, Args&&... args) const -> void {
     RuntimeDeviceCheck(::cudaLaunchKernelEx(&m_config, kernel, std::forward<Args>(args)...), m_location);
@@ -142,7 +172,7 @@ struct LaunchKernel {
 
   cudaLaunchConfig_t m_config;
   const DebugInfo m_location;
-  /// TODO: We can add a queue to store the attributes (e.g. for PDL) if needed in the future.
+  cudaLaunchAttribute m_attrs[1];
 };
 
 }  // namespace host

python/sglang/jit_kernel/include/sgl_kernel/warp.cuh

@@ -0,0 +1,14 @@
+#pragma once
+
+// Some warp primitives
+namespace device::warp {
+
+template <typename T>
+__always_inline __device__ T reduce_sum(T val, uint32_t active_mask = 0xffffffff) {
+#pragma unroll
+  for (int mask = 16; mask > 0; mask >>= 1)
+    val += __shfl_xor_sync(active_mask, val, mask, 32);
+  return val;
+}
+
+}  // namespace device::warp