chore: import upstream snapshot with attribution

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

@@ -0,0 +1,184 @@
+"""
+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_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_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,
+    )

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

@@ -0,0 +1,147 @@
+"""
+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_ops=True,
+        waves_per_eu=2,
+        denorm_fp_math_f32="preserve-sign",
+        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())

python/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())