[Perf] Improve BF16 Qwen3-VL GEMMs, mixed attention and cache writes on Blackwell

python/sglang/benchmark/one_batch.py

@@ -81,6 +81,7 @@
 from sglang.srt.layers.moe import initialize_moe_config
 from sglang.srt.layers.quantization.fp4_utils import initialize_fp4_gemm_config
 from sglang.srt.layers.quantization.fp8_utils import initialize_fp8_gemm_config
+from sglang.srt.layers.quantization.unquant import initialize_bf16_gemm_config
 from sglang.srt.managers.schedule_batch import Req, ScheduleBatch
 from sglang.srt.managers.scheduler_components.dp_attn import prepare_mlp_sync_batch_raw
 from sglang.srt.mem_cache.base_prefix_cache import EvictParams
@@ -123,6 +124,7 @@ def _init_process_global_configs() -> None:
     """
     initialize_moe_config()
     initialize_fp8_gemm_config()
+    initialize_bf16_gemm_config()
     initialize_fp4_gemm_config()
 
 

python/sglang/kernels/jit/csrc/elementwise/activation.cuh

@@ -63,11 +63,11 @@ template <
     bool kFilterExpert,
     bool kRoundActivation = false,
     bool kReuseInput = false,
-    bool kClamp = false>
+    bool kClamp = false,
+    uint32_t kVecSize = device::kMaxVecBytes / sizeof(T)>
 __global__ void act_and_mul_kernel(const __grid_constant__ ActivationParams params) {
   using namespace device;
-  constexpr auto kVecSize = kMaxVecBytes / sizeof(T);
-  using vec_t = AlignedVector<T, kMaxVecBytes / sizeof(T)>;
+  using vec_t = AlignedVector<T, kVecSize>;
   const auto num_vecs = params.hidden_dim / kVecSize;  // per token
   const auto tid = blockIdx.x * blockDim.x + threadIdx.x;
   const auto token_id = tid / num_vecs;
@@ -114,11 +114,10 @@ struct UnaryActivationParams {
   uint32_t num_vecs;
 };
 
-template <typename T, ActivationKind kAct, bool kUsePDL>
+template <typename T, ActivationKind kAct, bool kUsePDL, uint32_t kVecSize = device::kMaxVecBytes / sizeof(T)>
 __global__ void act_kernel(const __grid_constant__ UnaryActivationParams params) {
   using namespace device;
-  constexpr auto kVecSize = kMaxVecBytes / sizeof(T);
-  using vec_t = AlignedVector<T, kMaxVecBytes / sizeof(T)>;
+  using vec_t = AlignedVector<T, kVecSize>;
   const auto vec_id = blockIdx.x * blockDim.x + threadIdx.x;
   if (vec_id >= params.num_vecs) return;
   PDLWaitPrimary<kUsePDL>();
@@ -132,17 +131,14 @@ __global__ void act_kernel(const __grid_constant__ UnaryActivationParams params)
   PDLTriggerSecondary<kUsePDL>();
 }
 
-template <typename T, bool kUsePDL>
+template <typename T, bool kUsePDL, uint32_t kVecSize = device::kMaxVecBytes / sizeof(T), uint32_t kBlockSize = 256>
 struct ActivationKernel {
-  static constexpr auto kVecSize = device::kMaxVecBytes / sizeof(T);
-  static constexpr auto kBlockSize = 256u;
-
   using kernel_fn_t = decltype(&act_and_mul_kernel<T, ActivationKind::kSiLU, kUsePDL, false>);
   using unary_kernel_fn_t = decltype(&act_kernel<T, ActivationKind::kReLU2, kUsePDL>);
 
   template <ActivationKind kAct, bool kFilterExpert, bool kRoundActivation = false, bool kReuseInput = false>
   static constexpr kernel_fn_t activation_kernel =
-      act_and_mul_kernel<T, kAct, kUsePDL, kFilterExpert, kRoundActivation, kReuseInput>;
+      act_and_mul_kernel<T, kAct, kUsePDL, kFilterExpert, kRoundActivation, kReuseInput, false, kVecSize>;
 
   static_assert(device::kMaxVecBytes % sizeof(T) == 0, "unsupported data type");
 
@@ -212,7 +208,7 @@ struct ActivationKernel {
     if constexpr (kClamp) {
       RuntimeCheck(type == "silu" && expert_ids == nullptr, "clamping requires unfiltered SiLU");
       const auto kernel =
-          act_and_mul_kernel<T, ActivationKind::kSiLU, kUsePDL, false, kRoundActivation, kReuseInput, true>;
+          act_and_mul_kernel<T, ActivationKind::kSiLU, kUsePDL, false, kRoundActivation, kReuseInput, true, kVecSize>;
       LaunchKernel(num_blocks, kBlockSize, device).enable_pdl(kUsePDL)(kernel, params);
     } else if (expert_ids != nullptr) {
       RuntimeCheck(expert_step > 0, "expert_step must be positive");
@@ -255,7 +251,7 @@ struct ActivationKernel {
   }
 
   template <ActivationKind kAct>
-  static constexpr auto unary_kernel = act_kernel<T, kAct, kUsePDL>;
+  static constexpr auto unary_kernel = act_kernel<T, kAct, kUsePDL, kVecSize>;
 
   // Use the explicit non-const function-pointer type (mirrors select_kernel's
   // kernel_fn_t) rather than a trailing `decltype(unary_kernel<...>)` return,

python/sglang/kernels/jit/csrc/elementwise/fused_qknorm_rope.cuh

@@ -26,6 +26,7 @@
 
 #include <tvm/ffi/container/tensor.h>
 
+#include "qknorm.cuh"
 #include <cmath>
 #include <cuda_bf16.h>
 #include <cuda_runtime.h>
@@ -343,4 +344,246 @@ void fused_qk_norm_rope(
           rotary_dim);
 }
 
+struct QKNormMRoPEParams {
+  void* q;
+  void* k;
+  const void* q_weight;
+  const void* k_weight;
+  const bf16_t* cache;
+  const int64_t* positions;
+  const int64_t* axis_map;
+  int64_t q_stride;
+  int64_t k_stride;
+  int64_t position_stride;
+  uint32_t q_heads;
+  uint32_t k_heads;
+  uint32_t tokens;
+  float eps;
+  const void* v;
+  void* key_cache;
+  void* value_cache;
+  const int64_t* slots;
+  int64_t v_stride;
+  int64_t block_stride;
+  int64_t token_stride;
+  int64_t head_stride;
+  int64_t page_size;
+};
+
+template <bool kWriteCache>
+struct QKNormMRoPEEpilogue {
+  struct Params {
+    QKNormParams norm;
+    QKNormMRoPEParams rope;
+  };
+  static __device__ const QKNormParams& norm_params(const Params& p) {
+    return p.norm;
+  }
+  template <typename Storage>
+  static __device__ Storage apply(Storage out, const Params& params, int64_t token, int64_t head) {
+    using namespace device;
+    const auto& p = params.rope;
+    const auto lane = get_lane_id();
+    const auto is_q = head < p.q_heads;
+#pragma unroll
+    for (int j = 0; j < 2; ++j) {
+      const auto f = cast<fp32x2_t>(out[j]);
+      const auto partner = make_float2(__shfl_xor_sync(0xffffffff, f.x, 16), __shfl_xor_sync(0xffffffff, f.y, 16));
+      bf16_t values[2] = {__float2bfloat16(f.x), __float2bfloat16(f.y)};
+      const bf16_t partners[2] = {__float2bfloat16(partner.x), __float2bfloat16(partner.y)};
+#pragma unroll
+      for (int v = 0; v < 2; ++v) {
+        const auto d = (lane * 4 + 2 * j + v) % 64;
+        const auto position = p.positions[p.axis_map[d] * p.position_stride + token];
+        const auto c = p.cache[position * 128 + d];
+        const auto s = p.cache[position * 128 + d + 64];
+        const auto a = lane < 16 ? values[v] : partners[v];
+        const auto b = lane < 16 ? partners[v] : values[v];
+        // Preserve the Triton MRoPE kernel's BF16 FMA contraction order.
+        values[v] = lane < 16 ? __hfma(a, c, __hneg(__hmul(b, s))) : __hfma(a, s, __hmul(b, c));
+      }
+      out[j] = __halves2bfloat162(values[0], values[1]);
+    }
+    return out;
+  }
+  template <typename Storage>
+  static __device__ void after_store(Storage out, const Params& params, int64_t token, int64_t head) {
+    using namespace device;
+    const auto& p = params.rope;
+    const auto is_q = head < p.q_heads;
+    const auto gmem = tile::Memory<Storage>::warp();
+    if constexpr (kWriteCache) {
+      if (!is_q) {
+        const auto slot = p.slots[token];
+        if (slot >= 0) {
+          const auto kv_head = head - p.q_heads;
+          const auto offset =
+              (slot / p.page_size) * p.block_stride + (slot % p.page_size) * p.token_stride + kv_head * p.head_stride;
+          gmem.store(pointer::offset(p.key_cache, 2 * offset), out);
+          const auto value = gmem.load(pointer::offset(p.v, 2 * (token * p.v_stride + kv_head * 128)));
+          gmem.store(pointer::offset(p.value_cache, 2 * offset), value);
+        }
+      }
+    }
+  }
+};
+
+/// Validated entry point for BF16 Q/K with 128-element heads.
+template <bool kUsePDL>
+struct FusedQKNormMRoPE {
+  static void
+  run(tvm::ffi::TensorView q,
+      tvm::ffi::TensorView k,
+      tvm::ffi::TensorView q_weight,
+      tvm::ffi::TensorView k_weight,
+      tvm::ffi::TensorView cache,
+      tvm::ffi::TensorView positions,
+      tvm::ffi::TensorView axis_map,
+      float eps) {
+    using namespace host;
+    auto N = SymbolicSize{"tokens"};
+    auto Q = SymbolicSize{"q_width"};
+    auto K = SymbolicSize{"k_width"};
+    auto Sq = SymbolicSize{"q_stride"};
+    auto Sk = SymbolicSize{"k_stride"};
+    auto Sp = SymbolicSize{"position_stride"};
+    auto C = SymbolicSize{"cache_length"};
+    auto device = SymbolicDevice{};
+    device.set_options<kDLCUDA>();
+    TensorMatcher({N, Q}).with_strides({Sq, 1}).with_dtype<bf16_t>().with_device(device).ensure_alignment(8).verify(q);
+    TensorMatcher({N, K}).with_strides({Sk, 1}).with_dtype<bf16_t>().with_device(device).ensure_alignment(8).verify(k);
+    TensorMatcher({128}).with_dtype<bf16_t>().with_device(device).ensure_alignment(8).verify(q_weight).verify(k_weight);
+    TensorMatcher({C, 128}).with_dtype<bf16_t>().with_device(device).verify(cache);
+    TensorMatcher({3, N}).with_strides({Sp, 1}).with_dtype<int64_t>().with_device(device).verify(positions);
+    TensorMatcher({64}).with_dtype<int64_t>().with_device(device).verify(axis_map);
+    CHECK_HOST(Q.unwrap() > 0 && K.unwrap() > 0);
+    CHECK_HOST(Q.unwrap() % 128 == 0 && K.unwrap() % 128 == 0);
+    if (N.unwrap() == 0) return;
+    const auto p = QKNormMRoPEParams{
+        q.data_ptr(),
+        k.data_ptr(),
+        q_weight.data_ptr(),
+        k_weight.data_ptr(),
+        static_cast<const bf16_t*>(cache.data_ptr()),
+        static_cast<const int64_t*>(positions.data_ptr()),
+        static_cast<const int64_t*>(axis_map.data_ptr()),
+        Sq.unwrap(),
+        Sk.unwrap(),
+        Sp.unwrap(),
+        static_cast<uint32_t>(Q.unwrap() / 128),
+        static_cast<uint32_t>(K.unwrap() / 128),
+        static_cast<uint32_t>(N.unwrap()),
+        eps};
+    using Epilogue = QKNormMRoPEEpilogue<false>;
+    const typename Epilogue::Params params{
+        QKNormParams{
+            p.q,
+            pointer::offset(p.k, -2 * static_cast<int64_t>(p.q_heads) * 128),
+            p.q_stride,
+            p.k_stride,
+            p.q_heads,
+            p.k_heads,
+            p.eps,
+            p.q_weight,
+            p.k_weight,
+            p.tokens},
+        p};
+    LaunchKernel(div_ceil((p.q_heads + p.k_heads) * p.tokens, 4u), 128, device.unwrap())
+        .enable_pdl(kUsePDL)(fused_qknorm_warp<128, kUsePDL, bf16_t, Epilogue, false>, params);
+  }
+
+  /// Normalize/rotate Q/K and write BF16 K/V to matching paged cache views.
+  static void run_with_cache(
+      tvm::ffi::TensorView q,
+      tvm::ffi::TensorView k,
+      tvm::ffi::TensorView v,
+      tvm::ffi::TensorView q_weight,
+      tvm::ffi::TensorView k_weight,
+      tvm::ffi::TensorView cache,
+      tvm::ffi::TensorView positions,
+      tvm::ffi::TensorView axis_map,
+      tvm::ffi::TensorView key_cache,
+      tvm::ffi::TensorView value_cache,
+      tvm::ffi::TensorView slots,
+      float eps) {
+    using namespace host;
+    auto N = SymbolicSize{"tokens"};
+    auto Q = SymbolicSize{"q_width"};
+    auto K = SymbolicSize{"k_width"};
+    auto Sq = SymbolicSize{"q_stride"};
+    auto Sk = SymbolicSize{"k_stride"};
+    auto Sp = SymbolicSize{"position_stride"};
+    auto C = SymbolicSize{"cache_length"};
+    auto device = SymbolicDevice{};
+    device.set_options<kDLCUDA>();
+    TensorMatcher({N, Q}).with_strides({Sq, 1}).with_dtype<bf16_t>().with_device(device).ensure_alignment(8).verify(q);
+    TensorMatcher({N, K}).with_strides({Sk, 1}).with_dtype<bf16_t>().with_device(device).ensure_alignment(8).verify(k);
+    TensorMatcher({128}).with_dtype<bf16_t>().with_device(device).ensure_alignment(8).verify(q_weight).verify(k_weight);
+    TensorMatcher({C, 128}).with_dtype<bf16_t>().with_device(device).verify(cache);
+    TensorMatcher({3, N}).with_strides({Sp, 1}).with_dtype<int64_t>().with_device(device).verify(positions);
+    TensorMatcher({64}).with_dtype<int64_t>().with_device(device).verify(axis_map);
+    CHECK_HOST(Q.unwrap() > 0 && K.unwrap() > 0);
+    CHECK_HOST(Q.unwrap() % 128 == 0 && K.unwrap() % 128 == 0);
+    auto Sv = SymbolicSize{"v_stride"};
+    auto Pages = SymbolicSize{"pages"};
+    auto PageSize = SymbolicSize{"page_size"};
+    auto Heads = SymbolicSize{"cache_heads"};
+    Heads.set_value(K.unwrap() / 128);
+    auto Sb = SymbolicSize{"block_stride"};
+    auto St = SymbolicSize{"token_stride"};
+    auto Sh = SymbolicSize{"head_stride"};
+    TensorMatcher({N, K}).with_strides({Sv, 1}).with_dtype<bf16_t>().with_device(device).ensure_alignment(8).verify(v);
+    TensorMatcher({Pages, PageSize, Heads, 128})
+        .with_strides({Sb, St, Sh, 1})
+        .with_dtype<bf16_t>()
+        .with_device(device)
+        .ensure_alignment(8)
+        .verify(key_cache)
+        .verify(value_cache);
+    TensorMatcher({N}).with_dtype<int64_t>().with_device(device).verify(slots);
+    CHECK_HOST(PageSize.unwrap() > 0);
+    if (N.unwrap() == 0) return;
+    const auto p = QKNormMRoPEParams{
+        q.data_ptr(),
+        k.data_ptr(),
+        q_weight.data_ptr(),
+        k_weight.data_ptr(),
+        static_cast<const bf16_t*>(cache.data_ptr()),
+        static_cast<const int64_t*>(positions.data_ptr()),
+        static_cast<const int64_t*>(axis_map.data_ptr()),
+        Sq.unwrap(),
+        Sk.unwrap(),
+        Sp.unwrap(),
+        static_cast<uint32_t>(Q.unwrap() / 128),
+        static_cast<uint32_t>(K.unwrap() / 128),
+        static_cast<uint32_t>(N.unwrap()),
+        eps,
+        v.data_ptr(),
+        key_cache.data_ptr(),
+        value_cache.data_ptr(),
+        static_cast<const int64_t*>(slots.data_ptr()),
+        Sv.unwrap(),
+        Sb.unwrap(),
+        St.unwrap(),
+        Sh.unwrap(),
+        PageSize.unwrap()};
+    using Epilogue = QKNormMRoPEEpilogue<true>;
+    const typename Epilogue::Params params{
+        QKNormParams{
+            p.q,
+            pointer::offset(p.k, -2 * static_cast<int64_t>(p.q_heads) * 128),
+            p.q_stride,
+            p.k_stride,
+            p.q_heads,
+            p.k_heads,
+            p.eps,
+            p.q_weight,
+            p.k_weight,
+            p.tokens},
+        p};
+    LaunchKernel(div_ceil((p.q_heads + p.k_heads) * p.tokens, 4u), 128, device.unwrap())
+        .enable_pdl(kUsePDL)(fused_qknorm_warp<128, kUsePDL, bf16_t, Epilogue, false>, params);
+  }
+};
+
 }  // namespace sglang