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

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

Browse Source
This commit is contained in:
Khoa Pham
2026-08-31 18:53:46 -07:00
committed by GitHub
parent 9a85473a89
commit 97744189b8
5 changed files with 59 additions and 32 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/kernels/jit/csrc/deepseek_v4/silu_and_mul_masked_post_quant.cuh
+34 -12
View File
@@ -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 tile = tile::Memory<Vec>::cta();
const auto row = blockIdx.x / params.blocks_per_row;
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);
const auto up = tile.load(params.input, bid * 2 + 1);
Vec out;
if (vec_id < params.out_vecs) {
const auto input = static_cast<const Vec*>(params.input);
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) {
out[i] = cast<DType2>(silu_and_mul<true>(cast<bf16x2_t>(gate[i]), cast<bf16x2_t>(up[i]), limit));
}
for (uint32_t i = 0; i < kVecSize / 2; ++i) {
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;
RuntimeCheck(num_threads <= 1024, "out_dim too large for single-block-per-row launch");
const auto out_vecs = out_dim / kVecSize;
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
+1 -2
View File
@@ -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
+2 -3
View File
@@ -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_n, i_hv = i_nh // HV, i_nh % HV
i_v, i_n, i_hv = tl.program_id(0), tl.program_id(1), tl.program_id(2)
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
+21 -13
View File
@@ -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[
int
] = None, # kept for API compat; stride is derived from ``intermediate_states_buffer.shape[1]``
cache_steps: Optional[int] = None,
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
# (they can differ under --speculative-adaptive).
cache_stride_steps = (
intermediate_states_buffer.shape[1]
if intermediate_states_buffer is not None
else 0
)
# Adaptive spec changes the runtime draft count without changing the
# allocated per-request pitch, which is preserved in stride(0).
if intermediate_states_buffer is not None:
cache_stride_steps = intermediate_states_buffer.stride(0) // (HV * K * V)
elif cache_steps is not None and cache_steps > 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
+1 -2
View File
@@ -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")
)