sglang 0.5.0rc0__py3-none-any.whl → 0.5.0rc1__py3-none-any.whl

sglang/srt/layers/attention/wave_ops/decode_attention.py

@@ -0,0 +1,186 @@
+"""
+Memory-efficient attention for decoding.
+It supports page size = 1.
+"""
+
+import functools
+import logging
+
+from wave_lang.kernel.lang.global_symbols import *
+from wave_lang.kernel.wave.compile import WaveCompileOptions, wave_compile
+from wave_lang.kernel.wave.constraints import GenericDot, MMAOperand, MMAType
+from wave_lang.kernel.wave.templates.paged_decode_attention import (
+    get_paged_decode_attention_kernels,
+    get_paged_decode_intermediate_arrays_shapes,
+    paged_decode_attention_shape,
+)
+from wave_lang.kernel.wave.utils.general_utils import get_default_scheduling_params
+from wave_lang.kernel.wave.utils.run_utils import set_default_run_config
+
+logger = logging.getLogger(__name__)
+import os
+
+dump_generated_mlir = int(os.environ.get("WAVE_DUMP_MLIR", 0))
+
+
+@functools.lru_cache(maxsize=4096)
+def get_wave_kernel(
+    shape: paged_decode_attention_shape,
+    max_kv_splits,
+    input_dtype,
+    output_dtype,
+    logit_cap,
+):
+    mha = (shape.num_query_heads // shape.num_kv_heads) == 1
+
+    # Get the kernels (either compile or load from cache).
+    if mha:
+        mfma_variant = (
+            GenericDot(along_dim=MMAOperand.M, k_vec_size=4, k_mult=1),
+            GenericDot(along_dim=MMAOperand.M, k_vec_size=1, k_mult=64),
+        )
+    else:
+        mfma_variant = (MMAType.F32_16x16x16_F16, MMAType.F32_16x16x16_F16)
+
+    (
+        phase_0,
+        phase_1,
+        hyperparams_0,
+        hyperparams_1,
+        dynamic_symbols_0,
+        dynamic_symbols_1,
+    ) = get_paged_decode_attention_kernels(
+        shape,
+        mfma_variant,
+        max_kv_splits,
+        input_dtype=input_dtype,
+        output_dtype=output_dtype,
+        logit_cap=logit_cap,
+    )
+    hyperparams_0.update(get_default_scheduling_params())
+    hyperparams_1.update(get_default_scheduling_params())
+
+    options = WaveCompileOptions(
+        subs=hyperparams_0,
+        canonicalize=True,
+        run_bench=False,
+        use_buffer_load_ops=True,
+        use_buffer_store_ops=True,
+        waves_per_eu=2,
+        dynamic_symbols=dynamic_symbols_0,
+        wave_runtime=True,
+    )
+    options = set_default_run_config(options)
+    phase_0 = wave_compile(options, phase_0)
+
+    options = WaveCompileOptions(
+        subs=hyperparams_1,
+        canonicalize=True,
+        run_bench=False,
+        use_buffer_load_ops=False,
+        use_buffer_store_ops=False,
+        waves_per_eu=4,
+        dynamic_symbols=dynamic_symbols_1,
+        wave_runtime=True,
+    )
+    options = set_default_run_config(options)
+    phase_1 = wave_compile(options, phase_1)
+
+    return phase_0, phase_1
+
+
+def decode_attention_intermediate_arrays_shapes(
+    num_seqs, head_size_kv, num_query_heads, max_kv_splits
+):
+    # Not all fields are used, but we need to pass them to the function
+    shape = paged_decode_attention_shape(
+        num_query_heads=num_query_heads,
+        num_kv_heads=0,
+        head_size=0,
+        head_size_kv=head_size_kv,
+        block_size=0,
+        num_seqs=num_seqs,
+    )
+    return get_paged_decode_intermediate_arrays_shapes(shape, max_kv_splits)
+
+
+def decode_attention_wave(
+    q,
+    k_buffer,
+    v_buffer,
+    o,
+    b_req_idx,
+    req_to_token,
+    attn_logits,
+    attn_logits_max,
+    num_kv_splits,
+    max_kv_splits,
+    sm_scale,
+    logit_cap,
+):
+    num_seqs, num_query_heads, head_size = q.shape
+    _, num_kv_heads, _ = k_buffer.shape
+    _, _, head_size_kv = v_buffer.shape
+    block_size = 32
+    shape = paged_decode_attention_shape(
+        num_query_heads,
+        num_kv_heads,
+        head_size,
+        head_size_kv,
+        block_size,
+        num_seqs,
+    )
+
+    phase_0, phase_1 = get_wave_kernel(
+        shape, max_kv_splits, q.dtype, o.dtype, logit_cap
+    )
+
+    mb_qk = phase_0(
+        q,
+        k_buffer,
+        v_buffer,
+        b_req_idx,
+        req_to_token,
+        attn_logits,
+        attn_logits_max,
+    )
+    if dump_generated_mlir:
+        filename = f"wave_decode_attention_phase0_{'x'.join(map(str, shape))}.mlir"
+        with open(filename, "w") as f:
+            f.write(mb_qk.module_op.get_asm())
+
+    mb_sv = phase_1(attn_logits, attn_logits_max, b_req_idx, o)
+    if dump_generated_mlir:
+        filename = f"wave_decode_attention_phase1_{'x'.join(map(str, shape))}.mlir"
+        with open(filename, "w") as f:
+            f.write(mb_sv.module_op.get_asm())
+
+
+def decode_attention_fwd(
+    q,
+    k_buffer,
+    v_buffer,
+    o,
+    b_req_idx,
+    req_to_token,
+    attn_logits,
+    attn_logits_max,
+    num_kv_splits,
+    max_kv_splits,
+    sm_scale,
+    logit_cap=0.0,
+):
+    decode_attention_wave(
+        q,
+        k_buffer,
+        v_buffer,
+        o,
+        b_req_idx,
+        req_to_token,
+        attn_logits,
+        attn_logits_max,
+        num_kv_splits,
+        max_kv_splits,
+        sm_scale,
+        logit_cap,
+    )

sglang/srt/layers/attention/wave_ops/extend_attention.py

@@ -0,0 +1,149 @@
+"""
+Memory-efficient attention for prefill.
+It support page size = 1.
+"""
+
+import functools
+import os
+
+import torch
+from wave_lang.kernel.lang.global_symbols import *
+from wave_lang.kernel.wave.compile import WaveCompileOptions, wave_compile
+from wave_lang.kernel.wave.constraints import MMAType
+from wave_lang.kernel.wave.scheduling.schedule import SchedulingType
+from wave_lang.kernel.wave.templates.attention_common import AttentionShape
+from wave_lang.kernel.wave.templates.extend_attention import get_extend_attention_kernel
+from wave_lang.kernel.wave.utils.general_utils import get_default_scheduling_params
+from wave_lang.kernel.wave.utils.run_utils import set_default_run_config
+
+dump_generated_mlir = int(os.environ.get("WAVE_DUMP_MLIR", 0))
+
+
+@functools.lru_cache
+def get_wave_kernel(
+    shape: AttentionShape,
+    q_shape: tuple[int],
+    k_shape: tuple[int],
+    v_shape: tuple[int],
+    k_cache_shape: tuple[int],
+    v_cache_shape: tuple[int],
+    o_shape: tuple[int],
+    input_dtype: torch.dtype,
+    output_dtype: torch.dtype,
+    size_dtype: torch.dtype,
+    is_causal: bool,
+    logit_cap: float,
+    layer_scaling: float,
+):
+    assert shape.num_query_heads % shape.num_kv_heads == 0
+
+    mfma_variant = (MMAType.F32_16x16x32_K8_F16, MMAType.F32_16x16x16_F16)
+    (
+        extend_attention,
+        hyperparams,
+        dynamic_symbols,
+    ) = get_extend_attention_kernel(
+        shape,
+        mfma_variant,
+        q_shape,
+        k_shape,
+        v_shape,
+        k_cache_shape,
+        v_cache_shape,
+        o_shape,
+        input_dtype=input_dtype,
+        output_dtype=output_dtype,
+        size_dtype=size_dtype,
+        is_causal=is_causal,
+        layer_scaling=layer_scaling,
+        logit_cap=logit_cap,
+    )
+
+    hyperparams.update(get_default_scheduling_params())
+    options = WaveCompileOptions(
+        subs=hyperparams,
+        canonicalize=True,
+        run_bench=False,
+        schedule=SchedulingType.NONE,
+        use_scheduling_barriers=False,
+        dynamic_symbols=dynamic_symbols,
+        use_buffer_load_ops=True,
+        use_buffer_store_ops=True,
+        waves_per_eu=2,
+        denorm_fp_math_f32="preserve-sign",
+        gpu_native_math_precision=True,
+        wave_runtime=True,
+    )
+    options = set_default_run_config(options)
+    extend_attention = wave_compile(options, extend_attention)
+
+    return extend_attention
+
+
+def extend_attention_wave(
+    q_extend,
+    k_extend,
+    v_extend,
+    k_buffer,
+    v_buffer,
+    qo_indptr,
+    kv_indptr,
+    kv_indices,
+    custom_mask,
+    mask_indptr,
+    max_seq_len,
+    output,
+    is_causal=True,
+    layer_scaling=None,
+    logit_cap=0,
+):
+    shape = AttentionShape(
+        num_query_heads=q_extend.shape[1],
+        num_kv_heads=k_extend.shape[1],
+        head_size=q_extend.shape[2],
+        head_size_kv=k_extend.shape[2],
+        num_seqs=kv_indptr.shape[0] - 1,
+        max_seq_len=max_seq_len,
+    )
+
+    # Run the wave kernel.
+    extend_attention = get_wave_kernel(
+        shape,
+        q_extend.shape,
+        k_extend.shape,
+        v_extend.shape,
+        k_buffer.shape,
+        v_buffer.shape,
+        output.shape,
+        input_dtype=q_extend.dtype,
+        output_dtype=output.dtype,
+        size_dtype=qo_indptr.dtype,
+        is_causal=is_causal,
+        layer_scaling=layer_scaling,
+        logit_cap=logit_cap,
+    )
+
+    mb = extend_attention(
+        q_extend,
+        k_extend,
+        v_extend,
+        k_buffer,
+        v_buffer,
+        qo_indptr,
+        kv_indptr,
+        kv_indices,
+        max_seq_len,
+        output,
+    )
+
+    if dump_generated_mlir:
+        shape_list = [
+            q_extend.shape[0],
+            q_extend.shape[1],
+            k_extend.shape[1],
+            q_extend.shape[2],
+            k_extend.shape[2],
+        ]
+        filename = f"wave_prefill_attention_{'x'.join(map(str, shape_list))}.mlir"
+        with open(filename, "w") as f:
+            f.write(mb.module_op.get_asm())

sglang/srt/layers/attention/wave_ops/prefill_attention.py

@@ -0,0 +1,79 @@
+"""
+Memory-efficient attention for prefill.
+It support page size = 1.
+"""
+
+import math
+import os
+
+from wave_lang.kernel.lang.global_symbols import *
+from wave_lang.kernel.wave.compile import WaveCompileOptions, wave_compile
+from wave_lang.kernel.wave.constraints import MMAType
+from wave_lang.kernel.wave.templates.attention_common import AttentionShape
+from wave_lang.kernel.wave.templates.prefill_attention import (
+    get_prefill_attention_kernel,
+)
+from wave_lang.kernel.wave.utils.general_utils import get_default_scheduling_params
+from wave_lang.kernel.wave.utils.run_utils import set_default_run_config
+
+dump_generated_mlir = int(os.environ.get("WAVE_DUMP_MLIR", 0))
+
+
+def prefill_attention_wave(
+    q, k, v, o, b_start_loc, b_seq_len, max_seq_len, is_causal=True
+):
+
+    shape = AttentionShape(
+        num_query_heads=q.shape[1],
+        num_kv_heads=k.shape[1],
+        head_size=q.shape[2],
+        head_size_kv=k.shape[2],
+        num_seqs=b_seq_len.shape[0],
+        max_seq_len=max_seq_len,
+        total_seq_len=q.shape[0],
+    )
+
+    assert shape.num_query_heads % shape.num_kv_heads == 0
+
+    output_shape = (shape.total_seq_len, shape.num_query_heads, shape.head_size_kv)
+    # Run the wave kernel.
+    mfma_variant = (MMAType.F32_16x16x16_F16, MMAType.F32_16x16x16_F16)
+    (prefill, hyperparams) = get_prefill_attention_kernel(
+        shape,
+        mfma_variant,
+        q.shape,
+        k.shape,
+        v.shape,
+        output_shape,
+        input_dtype=q.dtype,
+        output_dtype=o.dtype,
+        size_dtype=b_seq_len.dtype,
+    )
+
+    hyperparams.update(get_default_scheduling_params())
+
+    log2e = 1.44269504089
+    dk_sqrt = math.sqrt(1.0 / shape.head_size)
+
+    options = WaveCompileOptions(
+        subs=hyperparams,
+        canonicalize=True,
+        run_bench=False,
+        use_scheduling_barriers=False,
+    )
+    options = set_default_run_config(options)
+    prefill = wave_compile(options, prefill)
+
+    mb = prefill(
+        q * dk_sqrt * log2e,
+        k,
+        v,
+        b_start_loc,
+        b_seq_len,
+        o,
+    )
+    if dump_generated_mlir:
+        shape_list = [q.shape[0], q.shape[1], k.shape[1], q.shape[2], k.shape[2]]
+        filename = f"wave_prefill_attention_{'x'.join(map(str, shape_list))}.mlir"
+        with open(filename, "w") as f:
+            f.write(mb.module_op.get_asm())

sglang/srt/layers/communicator.py

@@ -408,9 +408,9 @@ class CommunicateWithAllReduceAndLayerNormFn:
     ):
         if residual_input_mode == ScatterMode.SCATTERED and context.attn_tp_size > 1:
             residual, local_residual = (
-                forward_batch.gathered_buffer[
-                    : forward_batch.input_ids.shape[0]
-                ].clone(),
+                torch.empty_like(
+                    forward_batch.gathered_buffer[: forward_batch.input_ids.shape[0]]
+                ),
                 residual,
             )
             attn_tp_all_gather_into_tensor(residual, local_residual)
@@ -420,13 +420,11 @@ class CommunicateWithAllReduceAndLayerNormFn:
 
             # Perform layernorm on smaller data before comm. Only valid when attn_tp_size is 1 (tp_size == dp_size)
             use_layer_norm_before_gather = context.attn_tp_size == 1
-            if use_layer_norm_before_gather:
-                residual.copy_(hidden_states)
-                if hidden_states.shape[0] != 0:
-                    hidden_states = layernorm(hidden_states)
-
+            if use_layer_norm_before_gather and hidden_states.shape[0] != 0:
+                residual = hidden_states
+                hidden_states = layernorm(hidden_states)
             hidden_states, local_hidden_states = (
-                forward_batch.gathered_buffer,
+                torch.empty_like(forward_batch.gathered_buffer),
                 hidden_states,
             )
             dp_gather_partial(hidden_states, local_hidden_states, forward_batch)
@@ -443,7 +441,7 @@ class CommunicateWithAllReduceAndLayerNormFn:
                 and _is_flashinfer_available
                 and hasattr(layernorm, "forward_with_allreduce_fusion")
                 and global_server_args_dict["enable_flashinfer_allreduce_fusion"]
-                and hidden_states.shape[0] <= 128
+                and hidden_states.shape[0] <= 2048
             ):
                 hidden_states, residual = layernorm.forward_with_allreduce_fusion(
                     hidden_states, residual

sglang/srt/layers/flashinfer_comm_fusion.py

@@ -1,5 +1,5 @@
 import logging
-from typing import Tuple
+from typing import Optional, Tuple
 
 import torch
 import torch.distributed as dist
@@ -92,7 +92,7 @@ _workspace_manager = FlashInferWorkspaceManager()
 
 
 def ensure_workspace_initialized(
-    max_token_num: int = 128, hidden_dim: int = 4096, use_fp32_lamport: bool = False
+    max_token_num: int = 2048, hidden_dim: int = 4096, use_fp32_lamport: bool = False
 ):
     """Ensure workspace is initialized"""
     if not is_flashinfer_available() or _flashinfer_comm is None:
@@ -124,8 +124,8 @@ def flashinfer_allreduce_residual_rmsnorm(
     residual: torch.Tensor,
     weight: torch.Tensor,
     eps: float = 1e-6,
-    max_token_num: int = 128,
-    use_oneshot: bool = True,
+    max_token_num: int = 2048,
+    use_oneshot: Optional[bool] = None,
     trigger_completion_at_end: bool = False,
     fp32_acc: bool = False,
 ) -> Tuple[torch.Tensor, torch.Tensor]:

sglang/srt/layers/linear.py

@@ -1294,6 +1294,7 @@ class RowParallelLinear(LinearBase):
         with use_symmetric_memory(parallel_state.get_tp_group()) as sm:
             output_parallel = self.quant_method.apply(self, input_parallel, bias=bias_)
             sm.tag(output_parallel)
+
         if self.reduce_results and self.tp_size > 1 and not skip_all_reduce:
             output = tensor_model_parallel_all_reduce(output_parallel)
         else:

sglang/srt/layers/moe/cutlass_moe.py

@@ -9,7 +9,7 @@ from typing import Any, Callable, Dict, List, Optional, Tuple
 import torch
 
 from sglang.srt.layers.moe.cutlass_moe_params import CutlassMoEParams
-from sglang.srt.layers.utils import is_sm100_supported
+from sglang.srt.layers.utils import is_sm90_supported, is_sm100_supported
 from sglang.srt.utils import is_cuda
 
 _is_cuda = is_cuda()
@@ -124,6 +124,7 @@ def cutlass_fused_experts_fp8(
 
     if is_cuda:
         from sglang.srt.layers.quantization.fp8_kernel import (
+            per_group_transpose,
             per_token_group_quant_fp8_hopper_moe_mn_major,
             sglang_per_token_group_quant_fp8,
         )
@@ -152,15 +153,12 @@ def cutlass_fused_experts_fp8(
         k,
     )
 
-    if is_sm100_supported():
-        a_q, a1_scale = sglang_per_token_group_quant_fp8(a, 128)
-        rep_a_q = shuffle_rows(a_q, a_map, (m * topk, k))
-        rep_a1_scales = shuffle_rows(a1_scale, a_map, (m * topk, int(k / 128)))
-    else:
-        rep_a = shuffle_rows(a, a_map, (m * topk, k))
-        rep_a_q, rep_a1_scales = per_token_group_quant_fp8_hopper_moe_mn_major(
-            rep_a, expert_offsets, problem_sizes1, 128
-        )
+    a_q, a1_scale = sglang_per_token_group_quant_fp8(a, 128)
+    rep_a_q = shuffle_rows(a_q, a_map, (m * topk, k))
+    rep_a1_scales = shuffle_rows(a1_scale, a_map, (m * topk, int(k / 128)))
+
+    if not is_sm100_supported():
+        rep_a1_scales = per_group_transpose(rep_a1_scales, expert_offsets)
         w1_scale = w1_scale.contiguous()
 
     c1 = torch.empty((m * topk, n * 2), device=device, dtype=out_dtype)
@@ -193,12 +191,9 @@ def cutlass_fused_experts_fp8(
     intermediate = torch.empty((m * topk, n), device=device, dtype=out_dtype)
     silu_and_mul(c1, intermediate)
 
-    if is_sm100_supported():
-        intemediate_q, a2_scale = sglang_per_token_group_quant_fp8(intermediate, 128)
-    else:
-        intemediate_q, a2_scale = per_token_group_quant_fp8_hopper_moe_mn_major(
-            intermediate, expert_offsets, problem_sizes2, 128
-        )
+    intemediate_q, a2_scale = sglang_per_token_group_quant_fp8(intermediate, 128)
+    if not is_sm100_supported():
+        a2_scale = per_group_transpose(a2_scale, expert_offsets)
         w2_scale = w2_scale.contiguous()
 
     fp8_blockwise_scaled_grouped_mm(

sglang/srt/layers/moe/cutlass_w4a8_moe.py

@@ -11,7 +11,7 @@ from sgl_kernel import (
 )
 
 from sglang.srt.layers.moe.ep_moe.kernels import (
-    post_reorder_triton_kernel,
+    post_reorder_triton_kernel_for_cutlass_moe,
     pre_reorder_triton_kernel_for_cutlass_moe,
     run_cutlass_moe_ep_preproess,
 )
@@ -199,14 +199,13 @@ def cutlass_w4a8_moe(
     )
 
     output = torch.empty_like(a)
-    post_reorder_triton_kernel[(m,)](
+    post_reorder_triton_kernel_for_cutlass_moe[(m,)](
         c2,
         output,
         src2dst,
-        topk_ids_,
+        local_topk_ids,
         topk_weights,
-        start_expert_id,
-        end_expert_id,
+        num_experts,
         topk,
         k,
         0,

sglang/srt/layers/moe/ep_moe/kernels.py

@@ -581,6 +581,49 @@ def post_reorder_triton_kernel(
             )
 
 
+@triton.jit
+def post_reorder_triton_kernel_for_cutlass_moe(
+    down_output_ptr,
+    output_ptr,
+    src2dst_ptr,
+    topk_ids_ptr,
+    topk_weights_ptr,
+    num_experts,
+    topk,
+    hidden_size,
+    dst_start,
+    BLOCK_SIZE: tl.constexpr,
+):
+    InDtype = down_output_ptr.dtype.element_ty
+
+    src_idx_int32 = tl.program_id(0)
+    src_idx = src_idx_int32.to(tl.int64)
+    src2dst_ptr = src2dst_ptr + src_idx * topk
+    topk_ids_ptr = topk_ids_ptr + src_idx * topk
+    topk_weights_ptr = topk_weights_ptr + src_idx * topk
+
+    store_ptr = output_ptr + src_idx * hidden_size
+
+    vec = tl.arange(0, BLOCK_SIZE)
+
+    for start_offset in tl.range(0, hidden_size, BLOCK_SIZE):
+        offset = start_offset + vec
+        mask = offset < hidden_size
+
+        sum_vec = tl.zeros([BLOCK_SIZE], dtype=InDtype)
+        for idx in range(topk):
+            expert_id = tl.load(topk_ids_ptr + idx)
+            if expert_id != num_experts:
+                dst_idx_int32 = tl.load(src2dst_ptr + idx)
+                dst_idx = dst_idx_int32.to(tl.int64)
+                dst_idx = dst_idx - dst_start
+                weigh_scale = tl.load(topk_weights_ptr + idx).to(InDtype)
+                load_ptr = down_output_ptr + dst_idx * hidden_size
+                in_data = tl.load(load_ptr + offset, mask=mask)
+                sum_vec += in_data * weigh_scale
+        tl.store(store_ptr + offset, sum_vec, mask=mask)
+
+
 @triton.jit
 def compute_m_range(
     pid,