[Kernel] Raise shape limits in shared FLA and MoE kernels (ported from #36507) (#37317)

python/sglang/kernels/jit/csrc/deepseek_v4/silu_and_mul_masked_post_quant.cuh

@@ -10,6 +10,7 @@
 
 #include <sgl_kernel/deepseek_v4/fp8_utils.cuh>
 
+#include <algorithm>
 #include <cstdint>
 #include <cuda_fp8.h>
 #include <type_traits>
@@ -208,8 +209,19 @@ struct SiluAndMulClampParams {
   const void* __restrict__ input;
   void* __restrict__ output;
   float swiglu_limit;
+  uint32_t out_vecs;
+  uint32_t blocks_per_row;
 };
 
+template <typename DType2>
+SGL_DEVICE bf16x2_t to_bf16x2(DType2 value) {
+  if constexpr (std::is_same_v<DType2, bf16x2_t>) {
+    return value;
+  } else {
+    return device::cast<bf16x2_t>(device::cast<fp32x2_t>(value));
+  }
+}
+
 template <typename DType, bool kUsePDL>
 __global__ __launch_bounds__(1024, 2) void  // maximize occupancy
     silu_mul_clamp_kernel(const SiluAndMulClampParams __grid_constant__ params) {
@@ -219,21 +231,27 @@ __global__ __launch_bounds__(1024, 2) void  // maximize occupancy
   constexpr auto kVecSize = 16 / sizeof(DType);
   static_assert(kVecSize % 2 == 0 && kVecSize > 0);
   using Vec = AlignedVector<DType2, kVecSize / 2>;
-  const auto bid = blockIdx.x;
+  const auto row = blockIdx.x / params.blocks_per_row;
-  const auto tile = tile::Memory<Vec>::cta();
+  const auto block_in_row = blockIdx.x % params.blocks_per_row;
+  const auto vec_id = block_in_row * blockDim.x + threadIdx.x;
   const float limit = params.swiglu_limit;
 
   PDLWaitPrimary<kUsePDL>();
-  const auto gate = tile.load(params.input, bid * 2 + 0);
+  if (vec_id < params.out_vecs) {
-  const auto up = tile.load(params.input, bid * 2 + 1);
+    const auto input = static_cast<const Vec*>(params.input);
-  Vec out;
+    auto output = static_cast<Vec*>(params.output);
+    const auto input_row = row * 2 * params.out_vecs;
+    const auto gate = input[input_row + vec_id];
+    const auto up = input[input_row + params.out_vecs + vec_id];
+    Vec out;
 
 #pragma unroll
-  for (uint32_t i = 0; i < kVecSize / 2; ++i) {
+    for (uint32_t i = 0; i < kVecSize / 2; ++i) {
-    out[i] = cast<DType2>(silu_and_mul<true>(cast<bf16x2_t>(gate[i]), cast<bf16x2_t>(up[i]), limit));
+      out[i] = cast<DType2>(silu_and_mul<true>(to_bf16x2(gate[i]), to_bf16x2(up[i]), limit));
-  }
+    }
 
-  tile.store(params.output, out, bid);
+    output[row * params.out_vecs + vec_id] = out;
+  }
   PDLTriggerSecondary<kUsePDL>();
 }
 
@@ -349,16 +367,20 @@ struct SiluAndMulClampKernel {
     constexpr uint32_t kVecSize = 16 / sizeof(DType);
     const auto out_dim = static_cast<uint32_t>(H.unwrap());
     const auto num_tokens = static_cast<uint32_t>(M.unwrap());
+    RuntimeCheck(out_dim > 0, "out_dim must be positive");
     RuntimeCheck(out_dim % kVecSize == 0, "out_dim must be divisible by vector size");
-    const auto num_threads = out_dim / kVecSize;
+    const auto out_vecs = out_dim / kVecSize;
-    RuntimeCheck(num_threads <= 1024, "out_dim too large for single-block-per-row launch");
+    const auto num_threads = std::min(out_vecs, 1024u);
+    const auto blocks_per_row = host::div_ceil(out_vecs, num_threads);
 
     const auto params = SiluAndMulClampParams{
         .input = input.data_ptr(),
         .output = output.data_ptr(),
         .swiglu_limit = static_cast<float>(swiglu_limit),
+        .out_vecs = out_vecs,
+        .blocks_per_row = blocks_per_row,
     };
-    LaunchKernel(num_tokens, num_threads, device.unwrap())  //
+    LaunchKernel(num_tokens * blocks_per_row, num_threads, device.unwrap())  //
         .enable_pdl(kUsePDL)(kernel, params);
   }
 };

python/sglang/kernels/ops/attention/fla/cumsum.py

@@ -70,10 +70,9 @@ def chunk_local_cumsum_scalar_kernel(
 
 @triton.autotune(
     configs=[
-        triton.Config({"BS": BS}, num_warps=num_warps, num_stages=num_stages)
+        triton.Config({"BS": BS}, num_warps=num_warps)
         for BS in BS_LIST
         for num_warps in [2, 4, 8]
-        for num_stages in [2, 3, 4]
     ],
     key=["B", "H", "S", "BT", "IS_VARLEN", "REVERSE", "HAS_SCALE"],
 )

python/sglang/kernels/ops/attention/fla/fused_recurrent.py

@@ -434,8 +434,7 @@ def fused_recurrent_kda_packed_decode_kernel(
     """KDA packed decode: same shape as the GDN packed decode kernel, but
     with a per-K gate (``a`` is ``[B, HV*K]`` and ``dt_bias`` is ``[HV*K]``),
     so the state decay is a per-K vector ``exp(g)`` rather than a scalar."""
-    i_v, i_nh = tl.program_id(0), tl.program_id(1)
+    i_v, i_n, i_hv = tl.program_id(0), tl.program_id(1), tl.program_id(2)
-    i_n, i_hv = i_nh // HV, i_nh % HV
     i_h = i_hv // (HV // H)
 
     o_k = tl.arange(0, BK)
@@ -674,7 +673,7 @@ def fused_recurrent_kda_packed_decode(
     stride_indices_seq = ssm_state_indices.stride(0)
 
     NV = triton.cdiv(V, BV)
-    grid = (NV, B * HV)
+    grid = (NV, B, HV)
     fused_recurrent_kda_packed_decode_kernel[grid](
         mixed_qkv=mixed_qkv,
         a=a,

python/sglang/kernels/ops/attention/fla/fused_sigmoid_gating_recurrent.py

@@ -69,11 +69,20 @@ def fused_sigmoid_gating_delta_rule_update_kernel(
     stride_beta_slot: tl.constexpr = 0,
     MAX_CACHE_LEN: tl.constexpr = 0,
     CACHE_RING: tl.constexpr = False,
+    SPLIT_N_HV_GRID: tl.constexpr = False,
     USE_GDC: tl.constexpr = False,
 ):
     """
     Fused kernel that combines sigmoid gating computation with recurrent delta rule update.
     """
+    if SPLIT_N_HV_GRID:
+        i_v, i_n, i_hv = tl.program_id(0), tl.program_id(1), tl.program_id(2)
+        # The GPU wrapper asserts NK == 1. Keep N and HV on independent grid
+        # axes so large GLM5 decode batches do not exceed CUDA's grid limit.
+        i_k = 0
+    else:
+        i_k, i_v, i_nh = tl.program_id(0), tl.program_id(1), tl.program_id(2)
+        i_n, i_hv = i_nh // HV, i_nh % HV
     # PDL: overlap this kernel's prologue with the producer (the KDA/GDN
     # conv1d_update). All global loads below happen after the wait, so
     # numerics are unchanged. The immediate trigger releases the LAUNCH of
@@ -83,8 +92,6 @@ def fused_sigmoid_gating_delta_rule_update_kernel(
         tl.extra.cuda.gdc_wait()
         tl.extra.cuda.gdc_launch_dependents()
 
-    i_k, i_v, i_nh = tl.program_id(0), tl.program_id(1), tl.program_id(2)
-    i_n, i_hv = i_nh // HV, i_nh % HV
     i_h = i_hv // (HV // H)
 
     if IS_VARLEN:
@@ -360,9 +367,7 @@ def fused_sigmoid_gating_delta_rule_update(
     disable_state_update: bool = False,
     intermediate_states_buffer: Optional[torch.Tensor] = None,
     intermediate_state_indices: Optional[torch.Tensor] = None,
-    cache_steps: Optional[
+    cache_steps: Optional[int] = None,
-        int
-    ] = None,  # kept for API compat; stride is derived from ``intermediate_states_buffer.shape[1]``
     retrieve_parent_token: Optional[torch.Tensor] = None,
     # fused ReplaySSM ring-write (spec verify). When cache_ring, each draft step
     # stores pre-norm k / raw v / gate / beta into these per-slot rings,
@@ -419,15 +424,17 @@ def fused_sigmoid_gating_delta_rule_update(
 
     NP2_T = triton.next_power_of_2(T)
 
-    grid = (NK, NV, N * HV)
+    split_n_hv_grid = q.device.type == "cuda"
+    grid = (NV, N, HV) if split_n_hv_grid else (NK, NV, N * HV)
 
-    # Per-req stride must match the buffer's allocated dim, not runtime steps
+    # Adaptive spec changes the runtime draft count without changing the
-    # (they can differ under --speculative-adaptive).
+    # allocated per-request pitch, which is preserved in stride(0).
-    cache_stride_steps = (
+    if intermediate_states_buffer is not None:
-        intermediate_states_buffer.shape[1]
+        cache_stride_steps = intermediate_states_buffer.stride(0) // (HV * K * V)
-        if intermediate_states_buffer is not None
+    elif cache_steps is not None and cache_steps > 0:
-        else 0
+        cache_stride_steps = cache_steps
-    )
+    else:
+        cache_stride_steps = 0
 
     # ring strides (per-slot rings are contiguous [num_slots, heads, L, dim];
     # the kernel offsets within a slot with MAX_CACHE_LEN and the dim extents).
@@ -516,6 +523,7 @@ def fused_sigmoid_gating_delta_rule_update(
         stride_beta_slot=stride_beta_slot,
         MAX_CACHE_LEN=max_cache_len,
         CACHE_RING=cache_ring,
+        SPLIT_N_HV_GRID=split_n_hv_grid,
         num_warps=num_warps,
         num_stages=num_stages,
         **pdl_kwargs,

python/sglang/srt/layers/attention/vision.py

@@ -499,8 +499,7 @@ class VisionTritonAttention(nn.Module):
             seq_lens = kwargs.get("sequence_lengths")
             if seq_lens is None:
                 seq_lens = cu_seqlens_gpu[1:] - cu_seqlens_gpu[:-1]
-            else:
+            seq_lens = seq_lens.to(device=q.device, dtype=torch.int32)
-                seq_lens = seq_lens.to(device=q.device, dtype=torch.int32)
             max_seqlen = resolve_precomputed_max_seqlen(
                 cu_seqlens_gpu, kwargs.get("max_seqlen")
             )