87 lines
1.9 KiB
Python
87 lines
1.9 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
|
|
|
|
"""Numerical primitives shared by exact LTX-2 Triton kernels."""
|
|
|
|
from vllm.triton_utils import tl, triton
|
|
|
|
|
|
@triton.jit
|
|
def round_bf16_to_fp32(value):
|
|
"""RNE-round FP32 to BF16 precision while retaining an FP32 register."""
|
|
|
|
bits = value.to(tl.int32, bitcast=True)
|
|
rounding_bias = 0x7FFF + ((bits >> 16) & 1)
|
|
rounded_bits = (bits + rounding_bias) & -65536
|
|
return rounded_bits.to(tl.float32, bitcast=True)
|
|
|
|
|
|
@triton.jit
|
|
def add_rn_f32(x, y):
|
|
return tl.inline_asm_elementwise(
|
|
asm="add.rn.f32 $0, $1, $2;",
|
|
constraints="=f,f,f",
|
|
args=[x, y],
|
|
dtype=tl.float32,
|
|
is_pure=True,
|
|
pack=1,
|
|
)
|
|
|
|
|
|
@triton.jit
|
|
def mul_rn_f32(x, y):
|
|
return tl.inline_asm_elementwise(
|
|
asm="mul.rn.f32 $0, $1, $2;",
|
|
constraints="=f,f,f",
|
|
args=[x, y],
|
|
dtype=tl.float32,
|
|
is_pure=True,
|
|
pack=1,
|
|
)
|
|
|
|
|
|
@triton.jit
|
|
def fma_rn_f32(x, y, accumulator):
|
|
return tl.inline_asm_elementwise(
|
|
asm="fma.rn.f32 $0, $1, $2, $3;",
|
|
constraints="=f,f,f,f",
|
|
args=[x, y, accumulator],
|
|
dtype=tl.float32,
|
|
is_pure=True,
|
|
pack=1,
|
|
)
|
|
|
|
|
|
@triton.jit
|
|
def rsqrt_approx_f32(x):
|
|
return tl.inline_asm_elementwise(
|
|
asm="rsqrt.approx.f32 $0, $1;",
|
|
constraints="=f,f",
|
|
args=[x],
|
|
dtype=tl.float32,
|
|
is_pure=True,
|
|
pack=1,
|
|
)
|
|
|
|
|
|
@triton.jit
|
|
def shfl_down_f32(value, delta: tl.constexpr):
|
|
return tl.inline_asm_elementwise(
|
|
asm="shfl.sync.down.b32 $0, $1, $2, 0x1f, 0xffffffff;",
|
|
constraints="=f,f,n",
|
|
args=[value, delta],
|
|
dtype=tl.float32,
|
|
is_pure=True,
|
|
pack=1,
|
|
)
|
|
|
|
|
|
__all__ = [
|
|
"add_rn_f32",
|
|
"fma_rn_f32",
|
|
"mul_rn_f32",
|
|
"round_bf16_to_fp32",
|
|
"rsqrt_approx_f32",
|
|
"shfl_down_f32",
|
|
]
|