diff --git a/python/sglang/kernels/jit/csrc/deepseek_v4/silu_and_mul_masked_post_quant.cuh b/python/sglang/kernels/jit/csrc/deepseek_v4/silu_and_mul_masked_post_quant.cuh index f09f9adee..fa8550782 100644 --- a/python/sglang/kernels/jit/csrc/deepseek_v4/silu_and_mul_masked_post_quant.cuh +++ b/python/sglang/kernels/jit/csrc/deepseek_v4/silu_and_mul_masked_post_quant.cuh @@ -10,6 +10,7 @@ #include +#include #include #include #include @@ -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 +SGL_DEVICE bf16x2_t to_bf16x2(DType2 value) { + if constexpr (std::is_same_v) { + return value; + } else { + return device::cast(device::cast(value)); + } +} + template __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; - const auto bid = blockIdx.x; - const auto tile = tile::Memory::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(); - 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(params.input); + auto output = static_cast(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(silu_and_mul(cast(gate[i]), cast(up[i]), limit)); - } + for (uint32_t i = 0; i < kVecSize / 2; ++i) { + out[i] = cast(silu_and_mul(to_bf16x2(gate[i]), to_bf16x2(up[i]), limit)); + } - tile.store(params.output, out, bid); + output[row * params.out_vecs + vec_id] = out; + } PDLTriggerSecondary(); } @@ -349,16 +367,20 @@ struct SiluAndMulClampKernel { constexpr uint32_t kVecSize = 16 / sizeof(DType); const auto out_dim = static_cast(H.unwrap()); const auto num_tokens = static_cast(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(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); } }; diff --git a/python/sglang/kernels/ops/attention/fla/cumsum.py b/python/sglang/kernels/ops/attention/fla/cumsum.py index 775a8f800..182211e55 100644 --- a/python/sglang/kernels/ops/attention/fla/cumsum.py +++ b/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"], ) diff --git a/python/sglang/kernels/ops/attention/fla/fused_recurrent.py b/python/sglang/kernels/ops/attention/fla/fused_recurrent.py index d3ea8b821..457535586 100644 --- a/python/sglang/kernels/ops/attention/fla/fused_recurrent.py +++ b/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_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, diff --git a/python/sglang/kernels/ops/attention/fla/fused_sigmoid_gating_recurrent.py b/python/sglang/kernels/ops/attention/fla/fused_sigmoid_gating_recurrent.py index 38dcd162b..8f7c34200 100644 --- a/python/sglang/kernels/ops/attention/fla/fused_sigmoid_gating_recurrent.py +++ b/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[ - 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, diff --git a/python/sglang/srt/layers/attention/vision.py b/python/sglang/srt/layers/attention/vision.py index 4661fbd89..591ea5c99 100644 --- a/python/sglang/srt/layers/attention/vision.py +++ b/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") )