chore: import upstream snapshot with attribution

python/sglang/srt/layers/attention/cute_utils/__init__.py

@@ -0,0 +1,135 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# Adapted from https://github.com/vllm-project/vllm/blob/4868b542c9dfd166662eecc4bb8be3a36a3feaa2/vllm/cute_utils/__init__.py
+from cutlass import BFloat16, Float32, Int64, Uint32, cute
+from cutlass._mlir import ir
+from cutlass._mlir.dialects import llvm, vector
+from cutlass.cute.nvgpu import cpasync
+from cutlass.cutlass_dsl import T, dsl_user_op
+
+# https://github.com/NVIDIA/cutlass/blob/v4.3.2/include/cute/arch/copy_sm90_desc.hpp#L193-L197
+EVICT_NORMAL = Int64(0x1000000000000000)
+EVICT_FIRST = Int64(0x12F0000000000000)
+EVICT_LAST = Int64(0x14F0000000000000)
+
+
+@dsl_user_op
+def recast_val(x, dtype, *, loc=None, ip=None):
+    return dtype(llvm.bitcast(dtype.mlir_type, x.ir_value(loc=loc, ip=ip)))
+
+
+def simple_tma_copy(atom, src, dst, mbar=None, cache_policy=None):
+    """A simple helper that wraps group_modes() and tma_partition()
+    NOTE: this should be called WITHOUT cute.elect_one()
+    """
+    if isinstance(atom.op, cpasync.CopyBulkTensorTileG2SOp):
+        gmem = src
+        smem = dst
+    elif isinstance(atom.op, cpasync.CopyBulkTensorTileS2GOp):
+        smem = src
+        gmem = dst
+    else:
+        raise ValueError
+
+    s_part, g_part = cpasync.tma_partition(
+        atom,
+        0,
+        cute.make_layout(1),
+        cute.group_modes(smem, 0),
+        cute.group_modes(gmem, 0),
+    )
+
+    if isinstance(atom.op, cpasync.CopyBulkTensorTileG2SOp):
+        cute.copy(atom, g_part, s_part, tma_bar_ptr=mbar, cache_policy=cache_policy)
+    elif isinstance(atom.op, cpasync.CopyBulkTensorTileS2GOp):
+        cute.copy(atom, s_part, g_part, cache_policy=cache_policy)
+    else:
+        raise ValueError
+
+
+# can't find the equivalent in nvvm
+@dsl_user_op
+def fence_before_tma_store(*, loc=None, ip=None):
+    llvm.inline_asm(
+        T.i32(),
+        [],
+        "mov.u32 $0, 0;\n\t"
+        "fence.proxy.async::generic.release.sync_restrict::shared::cta.cluster;",
+        "=r",
+        has_side_effects=True,
+        is_align_stack=False,
+        loc=loc,
+        ip=ip,
+    )
+
+
+@dsl_user_op
+def mma_bf16(
+    a: cute.TensorSSA, b: cute.TensorSSA, c: cute.TensorSSA, *, loc=None, ip=None
+):
+    if a.element_type == BFloat16:
+        a = cute.recast_tensor(a, Uint32)
+    if b.element_type == BFloat16:
+        b = cute.recast_tensor(b, Uint32)
+
+    mlir_ty = Float32.mlir_type
+    out = llvm.inline_asm(
+        llvm.StructType.get_literal([mlir_ty] * 4),
+        [a[i].ir_value(loc=loc, ip=ip) for i in range(4)]
+        + [b[i].ir_value(loc=loc, ip=ip) for i in range(2)]
+        + [c[i].ir_value(loc=loc, ip=ip) for i in range(4)],
+        "mma.sync.aligned.m16n8k16.row.col.f32.bf16.bf16.f32 "
+        "{$0, $1, $2, $3}, {$4, $5, $6, $7}, {$8, $9}, "
+        "{$10, $11, $12, $13};",
+        "=f,=f,=f,=f,r,r,r,r,r,r,f,f,f,f",
+        has_side_effects=False,
+        is_align_stack=False,
+        loc=loc,
+        ip=ip,
+    )
+    vec = vector.from_elements(
+        ir.VectorType.get([4], mlir_ty, loc=loc),
+        [llvm.extractvalue(mlir_ty, out, [i], loc=loc, ip=ip) for i in range(4)],
+        loc=loc,
+        ip=ip,
+    )
+    return cute.TensorSSA(vec, 4, Float32)
+
+
+@dsl_user_op
+def _bf16x2_abs(a: Uint32, *, loc=None, ip=None) -> Uint32:
+    out = llvm.inline_asm(
+        T.i32(),
+        [a.ir_value(loc=loc, ip=ip)],
+        "abs.bf16x2 $0, $1;",
+        "=r,r",
+        has_side_effects=False,
+        is_align_stack=False,
+    )
+    return Uint32(out)
+
+
+@dsl_user_op
+def _bf16x2_max(a: Uint32, b: Uint32, *, loc=None, ip=None) -> Uint32:
+    out = llvm.inline_asm(
+        T.i32(),
+        [a.ir_value(loc=loc, ip=ip), b.ir_value(loc=loc, ip=ip)],
+        "max.bf16x2 $0, $1, $2;",
+        "=r,r,r",
+        has_side_effects=False,
+        is_align_stack=False,
+    )
+    return Uint32(out)
+
+
+@dsl_user_op
+def _bf16x2_mul(a: Uint32, b: Uint32, *, loc=None, ip=None) -> Uint32:
+    out = llvm.inline_asm(
+        T.i32(),
+        [a.ir_value(loc=loc, ip=ip), b.ir_value(loc=loc, ip=ip)],
+        "mul.rn.bf16x2 $0, $1, $2;",
+        "=r,r,r",
+        has_side_effects=False,
+        is_align_stack=False,
+    )
+    return Uint32(out)

python/sglang/srt/layers/attention/cute_utils/_tcgen05.py

@@ -0,0 +1,220 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# Adapted from https://github.com/vllm-project/vllm/blob/4868b542c9dfd166662eecc4bb8be3a36a3feaa2/vllm/cute_utils/_tcgen05.py
+# this module is named _tcgen05 to avoid name collision with cute.nvgpu.tcgen05
+
+import cutlass
+from cutlass import Boolean, Float32, Int32, Uint32, Uint64, cute
+from cutlass._mlir import ir
+from cutlass._mlir.dialects import llvm, nvvm, vector
+from cutlass.cutlass_dsl import dsl_user_op
+
+NVVM_CTA_GROUP_MAP = [
+    None,
+    nvvm.Tcgen05GroupKind.CTA_1,
+    nvvm.Tcgen05GroupKind.CTA_2,
+]
+LDST_MAP = {
+    "32x32b": (nvvm.Tcgen05LdStShape.SHAPE_32X32B, 1),
+    "16x128b": (nvvm.Tcgen05LdStShape.SHAPE_16X128B, 2),
+    "16x256b": (nvvm.Tcgen05LdStShape.SHAPE_16X256B, 4),
+}
+
+
+def _make_tmem_llvm_ptr(addr, *, loc=None, ip=None):
+    ptr_ty = llvm.PointerType.get(cute.AddressSpace.tmem.value)
+    val = Int32(addr).ir_value(loc=loc, ip=ip)
+    return llvm.inttoptr(ptr_ty, val, loc=loc, ip=ip)
+
+
+@dsl_user_op
+def alloc(
+    taddr: cute.Pointer,
+    cta_group: int = 1,
+    *,
+    loc=None,
+    ip=None,
+) -> None:
+    nvvm.tcgen05_alloc(
+        taddr.to_llvm_ptr(loc=loc, ip=ip),
+        Uint32(512).ir_value(loc=loc, ip=ip),
+        group=NVVM_CTA_GROUP_MAP[cta_group],
+        loc=loc,
+        ip=ip,
+    )
+
+
+@dsl_user_op
+def dealloc(cta_group: int = 1, *, loc=None, ip=None) -> None:
+    nvvm.tcgen05_dealloc(
+        _make_tmem_llvm_ptr(0, loc=loc, ip=ip),
+        Int32(512).ir_value(loc=loc, ip=ip),
+        group=NVVM_CTA_GROUP_MAP[cta_group],
+        loc=loc,
+        ip=ip,
+    )
+
+
+def make_bf16_idesc(
+    MMA_M: int,
+    MMA_N: int,
+    *,
+    negate_A: bool = False,
+    negate_B: bool = False,
+    transpose_A: bool = False,
+    transpose_B: bool = False,
+):
+    idesc = Uint32(
+        (1 << 4) | (1 << 7) | (1 << 10) | ((MMA_N >> 3) << 17) | ((MMA_M >> 4) << 24)
+    )
+    idesc |= Uint32(negate_A) << 13
+    idesc |= Uint32(negate_B) << 14
+    idesc |= Uint32(transpose_A) << 15
+    idesc |= Uint32(transpose_B) << 16
+    return idesc
+
+
+def make_sdesc_128B_swizzle(LBO: int):
+    SBO = 8 * 128
+    return Uint64((LBO >> 4 << 16) | (SBO >> 4 << 32) | (1 << 46) | (2 << 61))
+
+
+@dsl_user_op
+def mma_f16(
+    d_tmem,
+    a_desc,
+    b_desc,
+    idesc,
+    enable_input_d,
+    cta_group: int = 1,
+    *,
+    loc=None,
+    ip=None,
+) -> None:
+    nvvm.tcgen05_mma(
+        nvvm.Tcgen05MMAKind.F16,
+        NVVM_CTA_GROUP_MAP[cta_group],
+        _make_tmem_llvm_ptr(d_tmem, loc=loc, ip=ip),
+        Uint64(a_desc).ir_value(loc=loc, ip=ip),
+        Uint64(b_desc).ir_value(loc=loc, ip=ip),
+        Int32(idesc).ir_value(loc=loc, ip=ip),
+        Boolean(enable_input_d).ir_value(loc=loc, ip=ip),
+        loc=loc,
+        ip=ip,
+    )
+
+
+@dsl_user_op
+def mma_ts_f16(
+    d_tmem,
+    a_tmem,
+    b_desc,
+    idesc,
+    enable_input_d,
+    cta_group: int = 1,
+    *,
+    loc=None,
+    ip=None,
+) -> None:
+    nvvm.tcgen05_mma(
+        nvvm.Tcgen05MMAKind.F16,
+        NVVM_CTA_GROUP_MAP[cta_group],
+        _make_tmem_llvm_ptr(d_tmem, loc=loc, ip=ip),
+        _make_tmem_llvm_ptr(a_tmem, loc=loc, ip=ip),
+        Uint64(b_desc).ir_value(loc=loc, ip=ip),
+        Int32(idesc).ir_value(loc=loc, ip=ip),
+        Boolean(enable_input_d).ir_value(loc=loc, ip=ip),
+        loc=loc,
+        ip=ip,
+    )
+
+
+@dsl_user_op
+def commit(mbar, cta_mask=None, cta_group: int = 1, *, loc=None, ip=None):
+    mbar_llvm = mbar.to_llvm_ptr(loc=loc, ip=ip)
+    group = NVVM_CTA_GROUP_MAP[cta_group]
+    if cutlass.const_expr(cta_mask is not None):
+        nvvm.tcgen05_commit_arrive(
+            mbar_llvm,
+            multicast_mask=cta_mask.ir_value(loc=loc, ip=ip),
+            group=group,
+            loc=loc,
+            ip=ip,
+        )
+    else:
+        nvvm.tcgen05_commit_arrive(mbar_llvm, group=group, loc=loc, ip=ip)
+
+
+@dsl_user_op
+def ld(row, col, shape: str, num: int, *, loc=None, ip=None):
+    nvvm_shape, regs_per_num = LDST_MAP[shape]
+    num_regs = regs_per_num * num
+    tmem = (Int32(row) << Int32(16)) | Int32(col)
+    tmem_ptr = _make_tmem_llvm_ptr(tmem, loc=loc, ip=ip)
+
+    if num_regs == 1:
+        reg = nvvm.tcgen05_ld(Int32.mlir_type, nvvm_shape, tmem_ptr, loc=loc, ip=ip)
+        reg_f32 = llvm.bitcast(Float32.mlir_type, reg, loc=loc, ip=ip)
+        return Float32(reg_f32)
+
+    else:
+        vec_i32_ty = ir.VectorType.get([num_regs], Int32.mlir_type, loc=loc)
+        vec_f32_ty = ir.VectorType.get([num_regs], Float32.mlir_type, loc=loc)
+        regs = nvvm.tcgen05_ld(vec_i32_ty, nvvm_shape, tmem_ptr, loc=loc, ip=ip)
+        regs_f32 = llvm.bitcast(vec_f32_ty, regs, loc=loc, ip=ip)
+        return cute.TensorSSA(regs_f32, (num_regs,), Float32)
+
+
+@dsl_user_op
+def st(row, col, shape: str, num: int, vals, *, loc=None, ip=None) -> None:
+    # if input is TensorSSA, convert to Tensor so we can bitcast
+    if isinstance(vals, cute.TensorSSA):
+        vals_ = cute.make_rmem_tensor_like(vals)
+        vals_.store(vals)
+        vals = vals_
+
+    # bitcast to Int32
+    vals = cute.recast_tensor(vals, Int32)
+
+    nvvm_shape, regs_per_num = LDST_MAP[shape]
+    num_regs = regs_per_num * num
+    tmem = (Int32(row) << Int32(16)) | Int32(col)
+    tmem_ptr = _make_tmem_llvm_ptr(tmem, loc=loc, ip=ip)
+
+    if num_regs == 1:
+        nvvm.tcgen05_st(
+            nvvm_shape,
+            tmem_ptr,
+            vals[0].ir_value(loc=loc, ip=ip),
+            loc=loc,
+            ip=ip,
+        )
+    else:
+        vec_i32_ty = ir.VectorType.get([num_regs], Int32.mlir_type, loc=loc)
+        val_vec = vector.from_elements(
+            vec_i32_ty,
+            [vals[i].ir_value(loc=loc, ip=ip) for i in range(num_regs)],
+            loc=loc,
+            ip=ip,
+        )
+        nvvm.tcgen05_st(nvvm_shape, tmem_ptr, val_vec, loc=loc, ip=ip)
+
+
+@dsl_user_op
+def fence_after_thread_sync(*, loc=None, ip=None):
+    nvvm.tcgen05_fence(nvvm.Tcgen05FenceKind.AFTER_THREAD_SYNC, loc=loc, ip=ip)
+
+
+@dsl_user_op
+def fence_before_thread_sync(*, loc=None, ip=None):
+    nvvm.tcgen05_fence(nvvm.Tcgen05FenceKind.BEFORE_THREAD_SYNC, loc=loc, ip=ip)
+
+
+@dsl_user_op
+def wait_ld(*, loc=None, ip=None):
+    nvvm.tcgen05_wait(nvvm.Tcgen05WaitKind.LOAD, loc=loc, ip=ip)
+
+
+@dsl_user_op
+def wait_st(*, loc=None, ip=None):
+    nvvm.tcgen05_wait(nvvm.Tcgen05WaitKind.STORE, loc=loc, ip=ip)

python/sglang/srt/layers/attention/cute_utils/cvt.py

@@ -0,0 +1,146 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# Adapted from https://github.com/vllm-project/vllm/blob/4868b542c9dfd166662eecc4bb8be3a36a3feaa2/vllm/cute_utils/cvt.py
+from cutlass import Constexpr, Float32, Uint32, cute
+from cutlass._mlir import ir
+from cutlass._mlir.dialects import llvm, vector
+from cutlass.cutlass_dsl import T, dsl_user_op
+
+
+@dsl_user_op
+def fp32x2_to_bf16x2(a: Float32, b: Float32, *, loc=None, ip=None) -> Uint32:
+    out = llvm.inline_asm(
+        T.i32(),
+        [a.ir_value(loc=loc, ip=ip), b.ir_value(loc=loc, ip=ip)],
+        "cvt.rn.bf16x2.f32 $0, $2, $1;",
+        "=r,f,f",
+        has_side_effects=False,
+        is_align_stack=False,
+    )
+    return Uint32(out)
+
+
+@dsl_user_op
+def bf16x2_to_fp32x2(data, *, loc=None, ip=None) -> tuple[Float32, Float32]:
+    if isinstance(data, Uint32):
+        out = llvm.inline_asm(
+            llvm.StructType.get_literal([T.f32(), T.f32()]),
+            [data.ir_value(loc=loc, ip=ip)],
+            "shl.b32 $0, $2, 16;\n\tand.b32 $1, $2, 0xFFFF0000;",
+            "=f,=f,r",
+            has_side_effects=False,
+            is_align_stack=False,
+            loc=loc,
+            ip=ip,
+        )
+        return (
+            Float32(llvm.extractvalue(T.f32(), out, [0], loc=loc, ip=ip)),
+            Float32(llvm.extractvalue(T.f32(), out, [1], loc=loc, ip=ip)),
+        )
+
+    elif isinstance(data, (cute.Tensor, cute.TensorSSA)):
+        # NOTE: the output is always 1D
+        size = cute.size(data.shape)
+        out = cute.make_rmem_tensor(size * 2, Float32)
+        for i in range(size):
+            out[i * 2], out[i * 2 + 1] = bf16x2_to_fp32x2(data[i])
+        return out
+
+    else:
+        raise ValueError(f"Unsupported type {type(data)}")
+
+
+@dsl_user_op
+def fp8x4_to_bf16x4(x: Uint32, *, loc=None, ip=None) -> cute.TensorSSA:
+    # there is only fp8->fp16 conversion, hence we need to go
+    # round trip through fp16.
+    out = llvm.inline_asm(
+        llvm.StructType.get_literal([T.i32()] * 2),
+        [x.ir_value(loc=loc, ip=ip)],
+        "{\n\t"
+        ".reg .b16 x0, x1;\n\t"
+        ".reg .b16 t00, t01, t10, t11;\n\t"
+        "mov.b32 {x0, x1}, $2;\n\t"
+        "cvt.rn.f16x2.e4m3x2 $0, x0;\n\t"
+        "cvt.rn.f16x2.e4m3x2 $1, x1;\n\t"
+        "mov.b32 {t00, t01}, $0;\n\t"
+        "mov.b32 {t10, t11}, $1;\n\t"
+        "cvt.rn.bf16.f16 t00, t00;\n\t"
+        "cvt.rn.bf16.f16 t01, t01;\n\t"
+        "cvt.rn.bf16.f16 t10, t10;\n\t"
+        "cvt.rn.bf16.f16 t11, t11;\n\t"
+        "mov.b32 $0, {t00, t01};\n\t"
+        "mov.b32 $1, {t10, t11};\n\t"
+        "}\n",
+        "=r,=r,r",
+        has_side_effects=False,
+        is_align_stack=False,
+    )
+    vec = vector.from_elements(
+        ir.VectorType.get([2], T.i32(), loc=loc),
+        [llvm.extractvalue(T.i32(), out, [i], loc=loc, ip=ip) for i in range(2)],
+        loc=loc,
+        ip=ip,
+    )
+    return cute.TensorSSA(vec, 2, Uint32)
+
+
+@dsl_user_op
+def fp32x4_to_fp8x4(
+    a0: Float32,
+    a1: Float32,
+    a2: Float32,
+    a3: Float32,
+    *,
+    loc=None,
+    ip=None,
+) -> Uint32:
+    # Pack four FP32 values into one b32 of four e4m3 bytes, byte order
+    # {a0, a1, a2, a3} from low to high address.
+    out = llvm.inline_asm(
+        T.i32(),
+        [
+            a0.ir_value(loc=loc, ip=ip),
+            a1.ir_value(loc=loc, ip=ip),
+            a2.ir_value(loc=loc, ip=ip),
+            a3.ir_value(loc=loc, ip=ip),
+        ],
+        "{\n\t"
+        ".reg .b16 t0, t1;\n\t"
+        "cvt.rn.satfinite.e4m3x2.f32 t0, $2, $1;\n\t"
+        "cvt.rn.satfinite.e4m3x2.f32 t1, $4, $3;\n\t"
+        "mov.b32 $0, {t0, t1};\n\t"
+        "}\n",
+        "=r,f,f,f,f",
+        has_side_effects=False,
+        is_align_stack=False,
+    )
+    return Uint32(out)
+
+
+@dsl_user_op
+def fp32x8_to_fp4x8(
+    vals: cute.Tensor,
+    offset: Constexpr[int],
+    *,
+    loc=None,
+    ip=None,
+) -> Uint32:
+    # Pack eight scaled FP32 values into four E2M1x2 bytes, returned as one b32.
+    assert vals.element_type is Float32
+    out = llvm.inline_asm(
+        T.i32(),
+        [vals[offset + i].ir_value(loc=loc, ip=ip) for i in range(8)],
+        "{\n\t"
+        ".reg .b8 x0, x1, x2, x3;\n\t"
+        "cvt.rn.satfinite.e2m1x2.f32 x0, $2, $1;\n\t"
+        "cvt.rn.satfinite.e2m1x2.f32 x1, $4, $3;\n\t"
+        "cvt.rn.satfinite.e2m1x2.f32 x2, $6, $5;\n\t"
+        "cvt.rn.satfinite.e2m1x2.f32 x3, $8, $7;\n\t"
+        "mov.b32 $0, {x0, x1, x2, x3};\n\t"
+        "}\n",
+        "=r,f,f,f,f,f,f,f,f",
+        has_side_effects=False,
+        is_align_stack=False,
+    )
+    return Uint32(out)