wehub-resource-sync
59a0a3844c
PR Test AMD / cancel-on-close (push) Has been skipped
PR Test NVIDIA ARM / scan (push) Has been skipped
PR Test NVIDIA / cancel-on-close (push) Has been skipped
PR Test AMD / scan (push) Has been skipped
PR Test NVIDIA ARM / cancel-on-close (push) Has been skipped
PR Test NVIDIA / scan (push) Has been skipped
Release Docker Images / build (cu129-torch-2.11.0) (push) Has been skipped
Release Docker Images / build (cu130-torch-2.11.0) (push) Has been skipped
Release PyPI / publish (push) Has been skipped
Scheduler Python Test / test (push) Successful in 27m19s
Docs / build (push) Successful in 28m8s
Scheduler C++ Test / test (push) Successful in 28m19s
Scheduler C++ Test / test-flat (push) Successful in 28m18s
Docs / deploy (push) Has been cancelled
PR Test AMD / finish (push) Has been cancelled
PR Test NVIDIA / finish (push) Has been cancelled
PR Test NVIDIA ARM / finish (push) Has been cancelled
PR Test NVIDIA ARM / ${{ matrix.name }} (${{ matrix.runner }}) (push) Has been cancelled
PR Test AMD / ${{ matrix.name }} (${{ matrix.runner }}) (push) Has been cancelled
PR Test NVIDIA / ${{ matrix.name }} (${{ matrix.runner }}) (push) Has been cancelled
1185 lines
38 KiB
Python
1185 lines
38 KiB
Python
"""Local copies of small triton-kernels utilities used by AMD MoE kernels."""
|
|
|
|
# fmt: off
|
|
# isort: off
|
|
from __future__ import annotations
|
|
|
|
from dataclasses import dataclass
|
|
from typing import Optional, TypeAlias, Union
|
|
|
|
import torch
|
|
from tokenspeed_kernel_amd._triton import tl, triton
|
|
|
|
|
|
# activation metadata
|
|
# ---------------------------------------------------------------------------- #
|
|
@dataclass(frozen=True)
|
|
class FnSpecs:
|
|
name: str
|
|
fn: object
|
|
fn_arg_names: tuple[str, ...]
|
|
fn_arg_do_not_specialize: tuple[str, ...] = tuple()
|
|
reduction_n: int = 1
|
|
|
|
@staticmethod
|
|
def default():
|
|
return FnSpecs("dflt", None, tuple())
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class FusedActivation:
|
|
specs: FnSpecs = FnSpecs.default()
|
|
fn_args: tuple[object, ...] = tuple()
|
|
|
|
|
|
# swiglu
|
|
# ---------------------------------------------------------------------------- #
|
|
@triton.jit
|
|
def _swiglu_clip(x, limit, clip_lower: tl.constexpr):
|
|
res = tl.minimum(x, limit)
|
|
if clip_lower:
|
|
res = tl.maximum(-limit, res)
|
|
return res
|
|
|
|
|
|
@triton.jit
|
|
def _compute_swiglu(gelu, linear, scale, alpha, limit):
|
|
gelu = gelu.to(tl.float32) * scale
|
|
if limit is not None:
|
|
gelu = _swiglu_clip(gelu, limit, clip_lower=False)
|
|
linear = linear.to(tl.float32) * scale
|
|
if limit is not None:
|
|
linear = _swiglu_clip(linear, limit, clip_lower=True)
|
|
s = gelu / (1 + tl.exp(-alpha * gelu))
|
|
return tl.fma(s, linear, s)
|
|
|
|
|
|
@triton.jit(repr=lambda _: "_swiglu")
|
|
def swiglu_fn(input, alpha, limit):
|
|
gelu, linear = tl.split(tl.reshape(input, (input.shape[0], input.shape[1] // 2, 2)))
|
|
return _compute_swiglu(gelu, linear, 1.0, alpha, limit)
|
|
|
|
|
|
# data types
|
|
# ---------------------------------------------------------------------------- #
|
|
@dataclass(frozen=True)
|
|
class IntegerType:
|
|
bitwidth: int
|
|
is_signed: bool
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class FloatType:
|
|
bitwidth_exponent: int
|
|
bitwidth_mantissa: int
|
|
is_signed: bool
|
|
unsigned_zero: bool = False
|
|
|
|
@property
|
|
def bitwidth(self):
|
|
return int(self.is_signed) + self.bitwidth_exponent + self.bitwidth_mantissa
|
|
|
|
|
|
BIT = IntegerType(1, is_signed=False)
|
|
UINT8 = IntegerType(8, is_signed=False)
|
|
FP4 = FloatType(bitwidth_exponent=2, bitwidth_mantissa=1, is_signed=True)
|
|
FP8_E4M3FN = FloatType(bitwidth_exponent=4, bitwidth_mantissa=3, is_signed=True)
|
|
FP8_E4M3FNUZ = FloatType(
|
|
bitwidth_exponent=4, bitwidth_mantissa=3, is_signed=True, unsigned_zero=True
|
|
)
|
|
FP8_E5M2 = FloatType(bitwidth_exponent=5, bitwidth_mantissa=2, is_signed=True)
|
|
BF16 = FloatType(bitwidth_exponent=8, bitwidth_mantissa=7, is_signed=True)
|
|
FP16 = FloatType(bitwidth_exponent=5, bitwidth_mantissa=10, is_signed=True)
|
|
FP32 = FloatType(bitwidth_exponent=8, bitwidth_mantissa=23, is_signed=True)
|
|
FP64 = FloatType(bitwidth_exponent=11, bitwidth_mantissa=52, is_signed=True)
|
|
INT16 = IntegerType(16, is_signed=True)
|
|
INT32 = IntegerType(32, is_signed=True)
|
|
INT64 = IntegerType(64, is_signed=True)
|
|
|
|
DataType: TypeAlias = IntegerType | FloatType
|
|
|
|
|
|
# layout utilities
|
|
# ---------------------------------------------------------------------------- #
|
|
@dataclass(frozen=True)
|
|
class StridedLayout:
|
|
major_dim: int = -1
|
|
|
|
def __post_init__(self):
|
|
if not isinstance(self.major_dim, int):
|
|
raise TypeError(
|
|
f"StridedLayout(major_dim=...) must be an int, got {type(self.major_dim)}"
|
|
)
|
|
|
|
@property
|
|
def name(self):
|
|
return "STRIDED"
|
|
|
|
def swizzle_block_shape(self, block_shape):
|
|
return block_shape
|
|
|
|
def order(self, rank: int) -> list[int]:
|
|
"""
|
|
Returns the minor->major dimension order for a given tensor rank.
|
|
|
|
`self.major_dim` supports negative indexing (like Python).
|
|
"""
|
|
if rank <= 0:
|
|
return []
|
|
if not (-rank <= self.major_dim < rank):
|
|
raise ValueError(
|
|
f"Invalid StridedLayout.major_dim={self.major_dim} for rank={rank}"
|
|
)
|
|
major_dim = self.major_dim if self.major_dim >= 0 else self.major_dim + rank
|
|
base = list(reversed(range(rank)))
|
|
idx = base.index(major_dim)
|
|
base[0], base[idx] = base[idx], base[0]
|
|
return base
|
|
|
|
|
|
# storage
|
|
# ---------------------------------------------------------------------------- #
|
|
@dataclass
|
|
class Storage:
|
|
data: torch.Tensor
|
|
layout: StridedLayout
|
|
|
|
@property
|
|
def device(self):
|
|
return self.data.device
|
|
|
|
|
|
# main tensor class
|
|
# ---------------------------------------------------------------------------- #
|
|
@dataclass
|
|
class Tensor:
|
|
storage: Storage
|
|
dtype: IntegerType | FloatType
|
|
shape: list[int] | None = None
|
|
shape_max: list[int] | None = None
|
|
|
|
def __post_init__(self):
|
|
assert isinstance(self.storage, Storage)
|
|
# initialize dtype
|
|
if self.dtype.bitwidth < 8 and self.shape is None:
|
|
raise ValueError("shape must be provided for sub-byte types")
|
|
# initialize shape
|
|
if self.shape is None:
|
|
self.shape = list(self.storage.data.shape)
|
|
self.shape = list(self.shape)
|
|
# validate shape: all elements must be `int` or numel-1 `torch.Tensor`
|
|
is_int = lambda s: isinstance(s, int)
|
|
is_item = lambda s: hasattr(s, "numel") and s.numel() == 1
|
|
assert all(map(lambda s: is_int(s) or is_item(s), self.shape))
|
|
# initialize shape_max
|
|
if self.shape_max is None:
|
|
self.shape_max = [None] * len(self.shape)
|
|
for i, (s, smax) in enumerate(zip(self.shape, self.shape_max)):
|
|
if smax is not None and not is_int(smax):
|
|
raise ValueError(
|
|
f"shape_max[{i}] must be `int` or `None`; got {type(smax)}"
|
|
)
|
|
if smax is None:
|
|
self.shape_max[i] = s
|
|
# validate shape_max: all elements must be `int`
|
|
assert all(map(is_int, self.shape_max))
|
|
|
|
# torch compatibility layer
|
|
@property
|
|
def ndim(self):
|
|
return len(self.shape)
|
|
|
|
@property
|
|
def device(self):
|
|
return self.storage.device
|
|
|
|
def stride(self, i=None):
|
|
return self.storage.data.stride() if i is None else self.storage.data.stride(i)
|
|
|
|
def data_ptr(self):
|
|
return self.storage.data.data_ptr()
|
|
|
|
def numel(self):
|
|
return self.storage.data.numel()
|
|
|
|
def element_size(self):
|
|
return self.dtype.bitwidth // 8
|
|
|
|
@property
|
|
def data(self):
|
|
t = self.storage
|
|
return t.data if isinstance(t, Storage) else t
|
|
|
|
def dim(self):
|
|
return self.ndim
|
|
|
|
def size(self, i=None):
|
|
if i is None:
|
|
return self.shape
|
|
return self.shape[i]
|
|
|
|
|
|
def dtype_to_torch_dtype(dtype: DataType) -> torch.dtype:
|
|
if dtype is None:
|
|
return None
|
|
if not isinstance(dtype, DataType):
|
|
return dtype
|
|
return {
|
|
FP4: torch.uint8,
|
|
UINT8: torch.uint8,
|
|
FP8_E4M3FN: torch.float8_e4m3fn,
|
|
FP8_E4M3FNUZ: torch.float8_e4m3fnuz,
|
|
FP8_E5M2: torch.float8_e5m2,
|
|
BF16: torch.bfloat16,
|
|
FP32: torch.float32,
|
|
FP16: torch.float16,
|
|
FP64: torch.float64,
|
|
INT16: torch.int16,
|
|
INT32: torch.int32,
|
|
INT64: torch.int64,
|
|
}[dtype]
|
|
|
|
|
|
def torch_dtype_to_dtype(dtype: torch.dtype) -> DataType:
|
|
if isinstance(dtype, DataType):
|
|
return dtype
|
|
id = str(dtype).split(".")[-1]
|
|
vals = {
|
|
"uint8": UINT8,
|
|
"float8_e4m3fn": FP8_E4M3FN,
|
|
"float8_e4m3fnuz": FP8_E4M3FNUZ,
|
|
"float8_e5m2": FP8_E5M2,
|
|
"float16": FP16,
|
|
"bfloat16": BF16,
|
|
"float32": FP32,
|
|
"float64": FP64,
|
|
"int16": INT16,
|
|
"int32": INT32,
|
|
"int64": INT64,
|
|
}
|
|
if id in vals:
|
|
return vals[id]
|
|
if "float8" in id:
|
|
return FP8_E4M3FN
|
|
assert False, f"Unknown dtype: {id}"
|
|
|
|
|
|
def wrap_torch_tensor(
|
|
torch_tensor, dtype=None, shape=None, shape_max=None, layout=None
|
|
):
|
|
if dtype is None:
|
|
dtype = torch_tensor.dtype
|
|
dtype = torch_dtype_to_dtype(dtype)
|
|
if shape is None:
|
|
shape = list(torch_tensor.shape)
|
|
if dtype == FP4:
|
|
shape[torch_tensor.stride().index(1)] *= (
|
|
8 * torch_tensor.dtype.itemsize
|
|
) // dtype.bitwidth
|
|
if shape_max is None:
|
|
shape_max = list(shape)
|
|
if layout is None:
|
|
# For a strided (dense) tensor we only track which dimension has unit stride.
|
|
# This is consistent with how we expand `shape` for packed sub-byte dtypes.
|
|
major_dim = torch_tensor.stride().index(1) if 1 in torch_tensor.stride() else -1
|
|
layout = StridedLayout(major_dim=major_dim - torch_tensor.ndim)
|
|
return Tensor(
|
|
Storage(torch_tensor, layout), dtype=dtype, shape=shape, shape_max=shape_max
|
|
)
|
|
|
|
|
|
# sum bitmatrix rows
|
|
# ---------------------------------------------------------------------------- #
|
|
@triton.jit
|
|
def vpopc(x):
|
|
"""
|
|
Vertical popcount
|
|
Input x : uint32[..., N]
|
|
Output y : uint32[..., 32]
|
|
semantics : y[..., i] = sum_j((x[..., j] >> i) & 1)
|
|
credits: @apgoucher
|
|
"""
|
|
|
|
tl.static_assert(
|
|
x.dtype == tl.uint32, "x should consist of 32-bit unsigned integers"
|
|
)
|
|
|
|
BLOCK_N: tl.constexpr = x.shape[-1] # summation axis
|
|
BATCHES: tl.constexpr = x.numel // BLOCK_N # number of batches
|
|
if BLOCK_N >= 8:
|
|
sa1: tl.constexpr = 8
|
|
else:
|
|
sa1: tl.constexpr = BLOCK_N
|
|
# create 8-way sums in 4-bit fields:
|
|
y = tl.reshape(x, [BATCHES, BLOCK_N // sa1, sa1, 1])
|
|
y = (y >> tl.arange(0, 4)[None, None, None, :]) & 0x11111111
|
|
y = tl.sum(y, 2) # [BATCHES, BLOCK_N // sa1, 4]
|
|
if BLOCK_N >= 128:
|
|
sa2: tl.constexpr = 16
|
|
else:
|
|
sa2: tl.constexpr = BLOCK_N // sa1
|
|
# create 128-way sums in 8-bit fields:
|
|
y = tl.reshape(y, [BATCHES, BLOCK_N // (sa1 * sa2), sa2, 1, 4])
|
|
y = (y >> (4 * tl.arange(0, 2))[None, None, None, :, None]) & 0x0F0F0F0F
|
|
y = tl.sum(y, 2) # [BATCHES, BLOCK_N // (sa1 * sa2), 2, 4]
|
|
sa3: tl.constexpr = BLOCK_N // (sa1 * sa2)
|
|
# create N-way sums in 32-bit fields:
|
|
y = tl.reshape(y, [BATCHES, 1, sa3, 8])
|
|
y = (y >> (8 * tl.arange(0, 4))[None, :, None, None]) & 0x000000FF
|
|
y = tl.sum(y, 2) # [BATCHES, 4, 8]
|
|
y = tl.reshape(y, x.shape[:-1] + [32])
|
|
return y
|
|
|
|
|
|
@triton.jit
|
|
def _sum_bitmatrix_rows(
|
|
B,
|
|
shape_bm,
|
|
stride_bm: tl.constexpr,
|
|
stride_bn: tl.constexpr, # input bitmatrix
|
|
Out,
|
|
OutPartials,
|
|
stride_pm: tl.constexpr,
|
|
stride_pn,
|
|
shape_pn, # outputs
|
|
BLOCK_MM: tl.constexpr,
|
|
BLOCK_M: tl.constexpr,
|
|
):
|
|
tl.static_assert(BLOCK_MM % BLOCK_M == 0)
|
|
TILE_SIZE: tl.constexpr = BLOCK_MM // BLOCK_M
|
|
if isinstance(shape_bm, tl.tensor) and shape_bm.dtype.is_ptr():
|
|
shape_bm = tl.load(shape_bm)
|
|
# load input bits
|
|
pid_m = tl.program_id(0)
|
|
pid_n = tl.program_id(1)
|
|
offs_bm = pid_m * BLOCK_MM + tl.arange(0, BLOCK_MM)
|
|
bits = tl.load(
|
|
B + pid_n * stride_bn + offs_bm * stride_bm, mask=offs_bm < shape_bm, other=0
|
|
)
|
|
bits = tl.reshape(bits, [TILE_SIZE, BLOCK_M])
|
|
# partial row sum
|
|
partial_row_sum = vpopc(bits) # [TILE_SIZE, 32]
|
|
# write-back partial row sum
|
|
offs_pm = pid_m * TILE_SIZE + tl.arange(0, TILE_SIZE)
|
|
offs_n = pid_n * 32 + tl.arange(0, 32)
|
|
tl.store(
|
|
OutPartials + offs_pm[:, None] * stride_pm + offs_n[None, :] * stride_pn,
|
|
partial_row_sum,
|
|
)
|
|
# update final row sum
|
|
tl.atomic_add(Out + offs_n, tl.sum(partial_row_sum, 0), sem="relaxed")
|
|
|
|
|
|
def cdiv(x, y):
|
|
return (x + y - 1) // y
|
|
|
|
|
|
def sum_bitmatrix_rows(x, partials_block_size=None):
|
|
assert partials_block_size is not None
|
|
PARTIALS_BLOCK_M = partials_block_size
|
|
n_rows, n_cols = x.shape
|
|
n_rows_max = x.shape_max[0]
|
|
|
|
TILE_SIZE = max(1, 128 // PARTIALS_BLOCK_M)
|
|
BLOCK_MM = PARTIALS_BLOCK_M * TILE_SIZE
|
|
|
|
grid_m = cdiv(n_rows_max, BLOCK_MM)
|
|
grid_n = cdiv(n_cols, 32)
|
|
out = torch.zeros((cdiv(n_cols, 128) * 128,), device=x.device, dtype=torch.int32)[
|
|
:n_cols
|
|
]
|
|
out_partials = torch.empty(
|
|
(grid_n * 32, grid_m * TILE_SIZE), device=x.device, dtype=torch.int32
|
|
)
|
|
out_partials = torch.transpose(out_partials, 0, 1)
|
|
# output tensors
|
|
_sum_bitmatrix_rows[(grid_m, grid_n)](
|
|
x.storage.data,
|
|
n_rows,
|
|
x.stride(0),
|
|
x.stride(1), # input
|
|
out, # output [final reduction]
|
|
out_partials,
|
|
out_partials.stride(0),
|
|
out_partials.stride(1),
|
|
out_partials.shape[1], # output [partial reductions]
|
|
BLOCK_M=PARTIALS_BLOCK_M,
|
|
BLOCK_MM=BLOCK_MM, # constants
|
|
num_warps=8,
|
|
)
|
|
out_partials = out_partials[: cdiv(n_rows_max, PARTIALS_BLOCK_M), :]
|
|
return out, out_partials
|
|
|
|
|
|
# bitmatrix metadata
|
|
# ---------------------------------------------------------------------------- #
|
|
@dataclass
|
|
class BitmatrixMetadata:
|
|
"""
|
|
Example:
|
|
`bitmatrix` = [0 0 1 0 1 1 0
|
|
0 1 0 0 0 1 0
|
|
1 1 1 0 0 0 1
|
|
0 0 1 0 1 0 0]
|
|
`col_sum` = [1 2 3 0 2 2 1]
|
|
`col_sorted_indx` = cat([5], [3 6], [0 7], [], [9 1 10], [2 4], [8])
|
|
`row_sorted_indx` = cat([3 6 8], [1 9], [0 2 4 10], [5 7])
|
|
"""
|
|
|
|
# the number of entries equal to 1 in each column
|
|
col_sum: torch.Tensor
|
|
# indices of nonzero values numbered row-major, grouped by cols, concatenated
|
|
col_sorted_indx: torch.Tensor
|
|
# indices of nonzero values numbered col-major, grouped by rows, concatenated
|
|
row_sorted_indx: torch.Tensor
|
|
|
|
|
|
@triton.jit
|
|
def _keyed_add(x, y):
|
|
# we keep the key in the upper 16 bits of a uint32:
|
|
key_mask: tl.constexpr = 0xFFFF0000
|
|
|
|
kx = x & key_mask
|
|
ky = y & key_mask
|
|
z = tl.where(kx == ky, x + y - kx, y)
|
|
return z
|
|
|
|
|
|
@triton.jit
|
|
def _bitmatrix_metadata_compute_stage2(
|
|
ColSortedIndx,
|
|
RowSortedIndx,
|
|
NonzeroIndx,
|
|
n_tokens,
|
|
ColPartialSum,
|
|
stride_pm,
|
|
stride_pn,
|
|
ColOffs,
|
|
TOKS_PER_ROW: tl.constexpr,
|
|
BLOCK_PER_TOK: tl.constexpr,
|
|
):
|
|
BLOCK_SIZE: tl.constexpr = BLOCK_PER_TOK * TOKS_PER_ROW
|
|
tl.static_assert(BLOCK_SIZE <= 32768)
|
|
if isinstance(n_tokens, tl.tensor) and n_tokens.dtype.is_ptr():
|
|
n_tokens = tl.load(n_tokens)
|
|
nonzero_indx_size = n_tokens * TOKS_PER_ROW
|
|
pid_m = tl.program_id(0)
|
|
# load column indices
|
|
offs_local = tl.arange(0, BLOCK_SIZE)
|
|
offs_global = pid_m * BLOCK_SIZE + offs_local
|
|
mask = offs_global < nonzero_indx_size
|
|
col_indx = tl.load(NonzeroIndx + offs_global, mask=mask, other=-1).to(tl.uint32)
|
|
# stable-sort by columns index
|
|
kv_pairs = ((col_indx << 16) | offs_local).to(tl.uint32)
|
|
kv_pairs = tl.sort(kv_pairs, 0)
|
|
col_indx = kv_pairs >> 16
|
|
offs_global = pid_m * BLOCK_SIZE + (kv_pairs & 0xFFFF)
|
|
mask = col_indx != 0xFFFF
|
|
# compute run lengths in column-sorted order:
|
|
x = kv_pairs & 0xFFFF0000 | 0x00000001
|
|
cols_and_inclusive_run_lengths = tl.associative_scan(x, 0, _keyed_add)
|
|
exclusive_run_lengths = (cols_and_inclusive_run_lengths - 1) & 0xFFFF
|
|
# compute output
|
|
row_sorted_indx = tl.load(
|
|
ColPartialSum + pid_m * stride_pm + col_indx * stride_pn, mask=mask
|
|
)
|
|
row_sorted_indx += tl.load(ColOffs + col_indx, mask=mask)
|
|
row_sorted_indx += exclusive_run_lengths
|
|
# write back output
|
|
tl.store(RowSortedIndx + offs_global, row_sorted_indx, mask=mask)
|
|
tl.store(ColSortedIndx + row_sorted_indx, offs_global, mask=mask)
|
|
|
|
|
|
@triton.jit
|
|
def _bitmatrix_metadata_compute_stage1(
|
|
CombinedIndx,
|
|
n_combined_indx,
|
|
sentinel,
|
|
BLOCK: tl.constexpr,
|
|
ColSum,
|
|
ColOffs,
|
|
n_cols,
|
|
PartialColSum,
|
|
shape_pm,
|
|
stride_pm,
|
|
stride_pn,
|
|
BLOCK_M: tl.constexpr,
|
|
BLOCK_N: tl.constexpr,
|
|
):
|
|
pid = tl.program_id(0)
|
|
# compute col_partial_sums
|
|
if pid < n_cols:
|
|
PartialColSum += pid * stride_pn
|
|
curr_sum = 0
|
|
for start in range(0, shape_pm, BLOCK_M):
|
|
offs = start + tl.arange(0, BLOCK_M) * stride_pm
|
|
partial_col_sum = tl.load(PartialColSum + offs, mask=offs < shape_pm)
|
|
out = tl.cumsum(partial_col_sum, 0) - partial_col_sum + curr_sum
|
|
curr_sum += tl.sum(partial_col_sum, 0)
|
|
tl.store(PartialColSum + offs, out, mask=offs < shape_pm)
|
|
# compute col_offs
|
|
elif pid == n_cols:
|
|
curr_sum = 0
|
|
for start in range(0, n_cols, BLOCK_N):
|
|
offs = start + tl.arange(0, BLOCK_N)
|
|
col_sum = tl.load(ColSum + offs, mask=offs < n_cols)
|
|
col_offs = tl.cumsum(col_sum, 0) - col_sum + curr_sum
|
|
curr_sum += tl.sum(col_sum, 0)
|
|
tl.store(ColOffs + offs, col_offs, mask=offs < n_cols)
|
|
# memset `combined_indx` to `sentinel`
|
|
else:
|
|
offs = (pid - n_cols - 1) * BLOCK + tl.arange(0, BLOCK)
|
|
tl.store(CombinedIndx + offs, sentinel, mask=offs < n_combined_indx)
|
|
|
|
|
|
def make_bitmatrix_metadata(nonzero_indx, bitmatrix):
|
|
assert nonzero_indx.ndim == 2
|
|
PARTIAL_BLOCK_M = 32
|
|
col_sum, col_partial_sum = sum_bitmatrix_rows(
|
|
bitmatrix, partials_block_size=PARTIAL_BLOCK_M
|
|
)
|
|
# allocate memory
|
|
device = bitmatrix.device
|
|
n_indx = nonzero_indx.numel()
|
|
n_cols = bitmatrix.shape[1]
|
|
col_offs = torch.empty(n_cols, dtype=torch.int32, device=device)
|
|
combined_indx = torch.empty(n_indx * 2, dtype=torch.int32, device=device)
|
|
col_sorted_indx = combined_indx[:n_indx]
|
|
row_sorted_indx = combined_indx[n_indx:]
|
|
# this kernel:
|
|
# - initializes `{row,col}_sorted_indx` to `sentinel`
|
|
# - computes col_offs; necessary for computing `{row,col}_sorted_indx`
|
|
# - computes col_partial_sums; necessary for computing `{row,col}_sorted_indx`
|
|
MEMSET_BLOCK = 1024
|
|
memset_grid = (cdiv(n_indx * 2, MEMSET_BLOCK) + n_cols + 1,)
|
|
_bitmatrix_metadata_compute_stage1[memset_grid](
|
|
combined_indx,
|
|
n_indx * 2,
|
|
-1,
|
|
MEMSET_BLOCK,
|
|
col_sum, #
|
|
col_offs,
|
|
col_sum.shape[0],
|
|
col_partial_sum, # inputs
|
|
col_partial_sum.shape[0],
|
|
col_partial_sum.stride(0),
|
|
col_partial_sum.stride(1), # outputs
|
|
BLOCK_M=512,
|
|
BLOCK_N=512, # tunable parameters
|
|
)
|
|
# this kernel computes valid entries of `{row,col}_sorted_indx`
|
|
# using `col_offs` and `col_partial_sums`
|
|
n_indx = nonzero_indx.numel()
|
|
toks_per_row = nonzero_indx.shape[-1]
|
|
compute_grid = (cdiv(bitmatrix.shape_max[0], PARTIAL_BLOCK_M),)
|
|
_bitmatrix_metadata_compute_stage2[compute_grid](
|
|
col_sorted_indx,
|
|
row_sorted_indx, # outputs
|
|
nonzero_indx,
|
|
bitmatrix.shape[0],
|
|
col_partial_sum,
|
|
col_partial_sum.stride(0),
|
|
col_partial_sum.stride(1), # inputs
|
|
col_offs, #
|
|
TOKS_PER_ROW=toks_per_row,
|
|
BLOCK_PER_TOK=PARTIAL_BLOCK_M, #
|
|
)
|
|
return BitmatrixMetadata(
|
|
col_sum=col_sum,
|
|
col_sorted_indx=col_sorted_indx,
|
|
row_sorted_indx=row_sorted_indx,
|
|
)
|
|
|
|
|
|
# sparse matrix
|
|
# ---------------------------------------------------------------------------- #
|
|
@dataclass
|
|
class SparseMatrix:
|
|
indx: torch.Tensor
|
|
vals: torch.Tensor
|
|
mask: Tensor
|
|
|
|
def __post_init__(self):
|
|
self.mask_metadata = make_bitmatrix_metadata(self.indx, self.mask)
|
|
|
|
|
|
# topk forward kernels
|
|
# ---------------------------------------------------------------------------- #
|
|
@triton.jit
|
|
def get_topmask_and_fullmask(x):
|
|
tl.static_assert(
|
|
x.dtype.is_int_unsigned(), "floating-point value must be passed as bits"
|
|
)
|
|
tm: tl.constexpr = 1 << (-1 + x.dtype.primitive_bitwidth)
|
|
fm: tl.constexpr = (1 << x.dtype.primitive_bitwidth) - 1
|
|
tm_arr = tl.full(x.shape, tm, dtype=x.dtype)
|
|
fm_arr = tl.full(x.shape, fm, dtype=x.dtype)
|
|
return tm_arr, fm_arr
|
|
|
|
|
|
@triton.jit
|
|
def fpval_to_key(x):
|
|
tm, fm = get_topmask_and_fullmask(x)
|
|
return x ^ tl.where((x & tm) != 0, fm, tm)
|
|
|
|
|
|
@triton.jit
|
|
def key_to_fpval(x):
|
|
tm, fm = get_topmask_and_fullmask(x)
|
|
return x ^ tl.where((x & tm) == 0, fm, tm)
|
|
|
|
|
|
# stable top-k tie-breaks to value with smaller index
|
|
@triton.jit
|
|
def indx_to_key(indx, N_EXPTS_PAD: tl.constexpr):
|
|
return N_EXPTS_PAD - indx
|
|
|
|
|
|
@triton.jit
|
|
def key_to_indx(indx, N_EXPTS_PAD: tl.constexpr):
|
|
return N_EXPTS_PAD - indx
|
|
|
|
|
|
@triton.jit
|
|
def streaming_topk(
|
|
X,
|
|
stride_xm,
|
|
n_expts_tot,
|
|
offs_m,
|
|
mask_m,
|
|
N_EXPTS_PAD: tl.constexpr,
|
|
N_EXPTS_ACT: tl.constexpr,
|
|
BLOCK_N: tl.constexpr,
|
|
):
|
|
x_nbits: tl.constexpr = X.dtype.element_ty.primitive_bitwidth
|
|
x_utype: tl.constexpr = tl.dtype(f"uint{x_nbits}")
|
|
if x_nbits < 16:
|
|
# this ensures that we leave at least 16 bits for expert index
|
|
# even if the input dtype is smaller than 16 bits:
|
|
y_nbits: tl.constexpr = 32
|
|
else:
|
|
y_nbits: tl.constexpr = x_nbits * 2
|
|
x_ultype: tl.constexpr = tl.dtype(f"uint{y_nbits}")
|
|
x_dtype: tl.constexpr = X.dtype.element_ty
|
|
|
|
# subtract 1 from loop iterations because we peel the first (masked) iteration:
|
|
loop_iterations: tl.constexpr = N_EXPTS_PAD // BLOCK_N - 1
|
|
offs_x_n = loop_iterations * BLOCK_N + tl.arange(0, BLOCK_N)
|
|
mask_n = offs_x_n[None, :] < n_expts_tot
|
|
|
|
# first iteration:
|
|
X_ptrs = X + offs_m[:, None] * stride_xm + offs_x_n[None, :]
|
|
x = tl.load(X_ptrs, mask=(mask_m & mask_n), other=float("-inf"))
|
|
x = fpval_to_key(x.to(x_utype, bitcast=True))
|
|
x = (x.to(x_ultype) << 16) | indx_to_key(offs_x_n, N_EXPTS_PAD)[None, :]
|
|
acc = tl.topk(x, N_EXPTS_ACT, dim=1)
|
|
|
|
# subsequent iterations:
|
|
for _i in (tl.static_range if loop_iterations <= 4 else range)(loop_iterations):
|
|
acc = tl.bitonic_merge(acc) # ensure sorted ascending for the merge
|
|
X_ptrs -= BLOCK_N
|
|
offs_x_n -= BLOCK_N
|
|
x = tl.load(X_ptrs, mask=mask_m, other=float("-inf"))
|
|
x = fpval_to_key(x.to(x_utype, bitcast=True))
|
|
x = (x.to(x_ultype) << 16) | indx_to_key(offs_x_n, N_EXPTS_PAD)[None, :]
|
|
acc = tl.maximum(acc, tl.topk(x, N_EXPTS_ACT, dim=1))
|
|
|
|
# sort packed (value_key, index_key) descending:
|
|
# this keeps outputs ordered by gate value and uses smaller expert index for ties
|
|
acc = tl.sort(acc, dim=1, descending=True)
|
|
# 0000vvvvvvvviiii --> 0000iiii:
|
|
y_indices_raw = (acc & 0xFFFF).to(tl.uint32)
|
|
y_indices = key_to_indx(y_indices_raw, N_EXPTS_PAD)
|
|
# 0000vvvvvvvviiii --> vvvvvvvv:
|
|
y_values_raw = (acc >> 16).to(x_utype)
|
|
y_values = key_to_fpval(y_values_raw).to(x_dtype, bitcast=True)
|
|
|
|
return y_values, y_indices
|
|
|
|
|
|
@triton.jit
|
|
def _topk_forward(
|
|
X,
|
|
stride_xm, # inputs
|
|
PeerYvs,
|
|
PeerYis,
|
|
stride_ym, # topk values/indices
|
|
USE_PROVIDED_INDX: tl.constexpr,
|
|
PeerBits,
|
|
stride_rm: tl.constexpr,
|
|
stride_rn: tl.constexpr, # bitmatrix
|
|
n_rows,
|
|
n_expts_tot, # shape
|
|
dst_offs_m,
|
|
APPLY_SOFTMAX: tl.constexpr, # constant
|
|
BLOCK_M: tl.constexpr,
|
|
N_EXPTS_PAD: tl.constexpr,
|
|
N_EXPTS_ACT: tl.constexpr,
|
|
BLOCK_N: tl.constexpr,
|
|
):
|
|
N_PEERS: tl.constexpr = len(PeerYvs)
|
|
|
|
pid = tl.program_id(0)
|
|
if isinstance(n_rows, tl.tensor) and n_rows.dtype.is_ptr():
|
|
n_rows = tl.load(n_rows)
|
|
|
|
if pid * BLOCK_M >= n_rows:
|
|
# early exit:
|
|
return
|
|
|
|
tl.static_assert(BLOCK_N % 32 == 0)
|
|
tl.static_assert(N_EXPTS_PAD % BLOCK_N == 0)
|
|
x_dtype: tl.constexpr = X.dtype.element_ty
|
|
|
|
# load logits
|
|
offs_m = pid * BLOCK_M + tl.arange(0, BLOCK_M)
|
|
offs_y_n = tl.arange(0, N_EXPTS_ACT)
|
|
mask_m = offs_m[:, None] < n_rows
|
|
if USE_PROVIDED_INDX:
|
|
tl.static_assert(len(PeerYis) == 1)
|
|
Yi_ptrs = (
|
|
PeerYis[0] + (dst_offs_m + offs_m[:, None]) * stride_ym + offs_y_n[None, :]
|
|
)
|
|
y_indices = tl.load(Yi_ptrs, mask=mask_m)
|
|
Xv_ptrs = X + offs_m[:, None] * stride_xm + y_indices
|
|
y_values = tl.load(Xv_ptrs, mask=mask_m)
|
|
else:
|
|
y_values, y_indices = streaming_topk(
|
|
X,
|
|
stride_xm,
|
|
n_expts_tot,
|
|
offs_m,
|
|
mask_m,
|
|
N_EXPTS_PAD,
|
|
N_EXPTS_ACT,
|
|
BLOCK_N,
|
|
)
|
|
|
|
# normalize selected values
|
|
if APPLY_SOFTMAX:
|
|
y_values = tl.softmax(y_values.to(tl.float32), dim=1, keep_dims=True).to(
|
|
x_dtype
|
|
)
|
|
|
|
# write back
|
|
for rank in tl.static_range(N_PEERS):
|
|
Yv_ptrs = (
|
|
PeerYvs[rank]
|
|
+ (dst_offs_m + offs_m[:, None]) * stride_ym
|
|
+ offs_y_n[None, :]
|
|
)
|
|
tl.store(Yv_ptrs, y_values, mask=mask_m)
|
|
if not USE_PROVIDED_INDX:
|
|
for rank in tl.static_range(N_PEERS):
|
|
Yi_ptrs = (
|
|
PeerYis[rank]
|
|
+ (dst_offs_m + offs_m[:, None]) * stride_ym
|
|
+ offs_y_n[None, :]
|
|
)
|
|
tl.store(Yi_ptrs, y_indices, mask=mask_m)
|
|
|
|
# pack into bitmatrix
|
|
y_div = y_indices // 32
|
|
y_rem = y_indices % 32
|
|
loop_iterations = N_EXPTS_PAD // BLOCK_N
|
|
for i in range(loop_iterations):
|
|
offs_r_n = tl.arange(0, BLOCK_N // 32) + i * (BLOCK_N // 32)
|
|
y2 = tl.where(
|
|
y_div[:, :, None] == offs_r_n[None, None, :], (1 << y_rem)[:, :, None], 0
|
|
)
|
|
r = tl.reduce_or(y2, axis=1)
|
|
for rank in tl.static_range(N_PEERS):
|
|
BitsPtrs = (
|
|
PeerBits[rank]
|
|
+ (dst_offs_m + offs_m[:, None]) * stride_rm
|
|
+ offs_r_n[None, :] * stride_rn
|
|
)
|
|
tl.store(BitsPtrs, r, mask=mask_m)
|
|
|
|
|
|
def make_empty(offset, shape, dtype, device, all_gather, symm_mem_pool):
|
|
dtype = dtype_to_torch_dtype(dtype)
|
|
if all_gather:
|
|
rank_id = symm_mem_pool.mesh.local_rank
|
|
ret_bufs = symm_mem_pool.make_empty(
|
|
shape=shape, dtype=dtype, region="topk", region_offset=offset
|
|
)
|
|
ret = ret_bufs[rank_id]
|
|
offset = symm_mem_pool.align_up(
|
|
offset + ret.numel() * ret.element_size(),
|
|
symm_mem_pool.regions["topk"].alignment,
|
|
)
|
|
return ret_bufs, ret, offset
|
|
ret = torch.empty(shape, dtype=dtype, device=device)
|
|
return (ret,), ret, 0
|
|
|
|
|
|
def topk_forward(
|
|
x,
|
|
k,
|
|
apply_softmax=True,
|
|
dim=1,
|
|
y_indx=None,
|
|
n_rows=None,
|
|
all_gather=False,
|
|
symm_mem_pool=None,
|
|
):
|
|
if not isinstance(x, Tensor):
|
|
x_shape = [x.shape[0] if n_rows is None else n_rows, x.shape[1]]
|
|
x_shape_max = [x.shape[0], x.shape[1]]
|
|
x = wrap_torch_tensor(x, shape=x_shape, shape_max=x_shape_max)
|
|
BLOCK_M = 32
|
|
BLOCK_N = 32
|
|
use_provided_indx = y_indx is not None
|
|
assert symm_mem_pool is not None or not all_gather
|
|
assert len(x.shape) == 2
|
|
assert x.shape_max[-1] < 32768
|
|
assert dim == 1
|
|
n_rows, n_cols = x.shape
|
|
n_rows_max, _ = x.shape_max
|
|
dev = x.device
|
|
n_rows_out_max = (
|
|
n_rows_max * symm_mem_pool.mesh.world_size if all_gather else n_rows_max
|
|
)
|
|
# scratchpad tensors
|
|
# NOTE: these are not returned
|
|
y_vals_bufs, y_vals, offset = make_empty(
|
|
0,
|
|
(n_rows_out_max, k),
|
|
x.dtype,
|
|
dev,
|
|
all_gather=all_gather,
|
|
symm_mem_pool=symm_mem_pool,
|
|
)
|
|
if y_indx is None:
|
|
y_indx_bufs, y_indx, offset = make_empty(
|
|
offset,
|
|
(n_rows_out_max, k),
|
|
torch.int16,
|
|
dev,
|
|
all_gather=all_gather,
|
|
symm_mem_pool=symm_mem_pool,
|
|
)
|
|
else:
|
|
y_indx_bufs = (y_indx,)
|
|
# create bitmatrix in transposed memory layout:
|
|
n_cols_pad = cdiv(n_cols, BLOCK_N) * BLOCK_N
|
|
n_cols_words = n_cols_pad // 32
|
|
bitmatrix_bufs, bitmatrix_data, offset = make_empty(
|
|
offset,
|
|
(n_cols_words, cdiv(n_rows_out_max, 32) * 32),
|
|
torch.uint32,
|
|
dev,
|
|
all_gather=all_gather,
|
|
symm_mem_pool=symm_mem_pool,
|
|
)
|
|
bitmatrix_data = torch.transpose(bitmatrix_data, 0, 1)[:n_rows_max]
|
|
pids = cdiv(n_rows_max, BLOCK_M)
|
|
_topk_forward[(pids,)](
|
|
x.storage.data,
|
|
x.stride(0), # inputs
|
|
y_vals_bufs,
|
|
y_indx_bufs,
|
|
y_vals.stride(0),
|
|
use_provided_indx, # output [topk]
|
|
bitmatrix_bufs,
|
|
bitmatrix_data.stride(0),
|
|
bitmatrix_data.stride(1), # output [bitmatrix]
|
|
n_rows,
|
|
n_cols, # shapes
|
|
symm_mem_pool.mesh.local_rank * n_rows_max if all_gather else 0,
|
|
BLOCK_M=BLOCK_M,
|
|
BLOCK_N=BLOCK_N, # tunable parameter
|
|
APPLY_SOFTMAX=apply_softmax,
|
|
N_EXPTS_PAD=n_cols_pad,
|
|
N_EXPTS_ACT=k, # constants
|
|
)
|
|
if all_gather:
|
|
symm_mem_pool.hdl.barrier(channel=0)
|
|
bitmatrix_shape = [
|
|
n_rows * symm_mem_pool.mesh.world_size if all_gather else n_rows,
|
|
n_cols,
|
|
]
|
|
bitmatrix_shape_max = [n_rows_out_max, None]
|
|
bitmatrix = wrap_torch_tensor(
|
|
bitmatrix_data, dtype=BIT, shape=bitmatrix_shape, shape_max=bitmatrix_shape_max
|
|
)
|
|
return y_vals, y_indx, bitmatrix
|
|
|
|
|
|
def topk(
|
|
x: Union[Tensor, torch.Tensor],
|
|
k: int,
|
|
apply_softmax: bool = True,
|
|
dim: int = 1,
|
|
y_indx: Optional[torch.Tensor] = None,
|
|
n_rows: Optional[int] = None,
|
|
all_gather: bool = False,
|
|
symm_mem_pool: object | None = None,
|
|
):
|
|
"""
|
|
Computes the top-k values and indices along a specified dimension of a tensor.
|
|
Note that the input can be either a `Tensor` or a `torch.Tensor`, but the output will always be a `torch.Tensor`.
|
|
|
|
Parameters
|
|
----------
|
|
x : Union[Tensor, torch.Tensor]
|
|
Input tensor of shape (n_tokens, n_expts).
|
|
k : int
|
|
Number of top elements to retrieve.
|
|
apply_softmax : bool, default True
|
|
Whether to apply softmax to the input tensor before computing top-k.
|
|
dim : int, default 1
|
|
Dimension along which to compute top-k.
|
|
y_indx : torch.Tensor, optional
|
|
Pre-allocated tensor for storing indices of top-k elements with shape (n_tokens, k).
|
|
If provided, we skip the computation of top-k indices and use this tensor instead.
|
|
n_rows : int, optional
|
|
Number of rows to apply top-k on. If None, we consider all rows in `x`.
|
|
|
|
Returns
|
|
-------
|
|
SparseMatrix: sparse matrix equal to `x` with non-selected entries set to 0
|
|
"""
|
|
y_vals, y_indx, bitmatrix = topk_forward(
|
|
x,
|
|
k,
|
|
apply_softmax,
|
|
dim,
|
|
y_indx,
|
|
n_rows,
|
|
all_gather,
|
|
symm_mem_pool,
|
|
)
|
|
return SparseMatrix(vals=y_vals, indx=y_indx, mask=bitmatrix)
|
|
|
|
|
|
# ragged tensor metadata
|
|
# ---------------------------------------------------------------------------- #
|
|
@dataclass
|
|
class RaggedTensorMetadata:
|
|
"""
|
|
Example:
|
|
`slice_sizes`= [15 17 0 127]
|
|
`slice_offs`= [0 15 32 32 332]
|
|
`block_offs_data` = {
|
|
16: [0 1 3 3 11]
|
|
32: [0 1 2 2 6]
|
|
64: [0 1 2 2 4]
|
|
128: [0 1 2 2 3]
|
|
}
|
|
`block_schedule_data` = {
|
|
16: [(0, 0) (0, 1) (0, 3) (1, 3) (2, 3) ... (7, 3) -1 ... -1]
|
|
32: [(0, 0) (0, 1) (0, 3) (1, 3) (2, 3) (3, 3) -1 ... -1]
|
|
64: [(0, 0) (0, 1) (0, 3) (1, 3) (2, 3) -1 ... -1]
|
|
128: [(0, 0) (0, 1) (0, 3) (1, 3) -1 ... -1]
|
|
}
|
|
"""
|
|
|
|
# slice_sizes[i] is the number of elements in slice i along the ragged dimension
|
|
slice_sizes: torch.Tensor
|
|
# slice_offs = [0] + cumsum(slice_sizes)
|
|
# i.e., slice_offs[i] is the offset of the first element in slice `i`
|
|
slice_offs: torch.Tensor
|
|
# block_offs_data[k] = [0] + cumsum(ceil_div(slice_sizes, 16 * k))
|
|
# i.e., `block_offs_data[k][i]` is the offset of the first block of
|
|
# `16*k`` token for batch `i` in a `bath_sizes`-shaped ragged tensor
|
|
block_offs_data: torch.Tensor
|
|
# let `num_blocks[k] = block_offs_data[k, 1:] - block_offs_data[k, :-1]
|
|
# block_schedule_data[k] = cat(*[[(batch, blk) for blk in range(blks)] for batch, blks in enumerate(num_blocks)])
|
|
# i.e., if the schedule of batch `i` is [(i, 0), (i, 1), ..., (i, num_blocks[k][i] - 1)]
|
|
# then `block_schedule_data[k]` is the concatenation of the schedules for all batches
|
|
# NOTE 1: `block_schedule_data[k][j]` is a packed 32-bit integer
|
|
# NOTE 2: because the size of `block_schedule_data[k]` is data-dependent, we pad it with -1s
|
|
# up to an user-provided upper bound
|
|
block_schedule_data: torch.Tensor
|
|
# expected slice size (for heuristics)
|
|
expected_slice_size: int | None = None
|
|
# divisibility hint for values in `slice_sizes`
|
|
slice_sizes_divisibility: int = None
|
|
|
|
def __post_init__(self):
|
|
assert self.block_offs_data.shape[0] == len(RaggedTensorMetadata.block_sizes())
|
|
assert self.block_schedule_data.shape[0] == len(
|
|
RaggedTensorMetadata.block_sizes()
|
|
)
|
|
assert self.block_offs_data.dtype == torch.int32
|
|
assert self.block_schedule_data.dtype == torch.int32
|
|
if self.slice_sizes is not None:
|
|
assert self.slice_sizes.dtype == torch.int32
|
|
if self.slice_offs is not None:
|
|
assert self.slice_offs.dtype == torch.int32
|
|
|
|
@property
|
|
def n_slices(self):
|
|
return self.slice_sizes.shape[0]
|
|
|
|
def block_offs(self, block_size):
|
|
return self.block_offs_data[
|
|
RaggedTensorMetadata.block_sizes().index(block_size)
|
|
]
|
|
|
|
def block_schedule(self, block_size):
|
|
return self.block_schedule_data[
|
|
RaggedTensorMetadata.block_sizes().index(block_size)
|
|
]
|
|
|
|
@staticmethod
|
|
def n_blocks(n_slices, n_total_rows, block_size):
|
|
if n_total_rows <= n_slices:
|
|
return n_total_rows
|
|
return n_slices - 1 - ((n_slices - n_total_rows - 1) // block_size)
|
|
|
|
@staticmethod
|
|
def max_n_blocks(n_slices, n_total_rows):
|
|
return RaggedTensorMetadata.n_blocks(
|
|
n_slices, n_total_rows, min(RaggedTensorMetadata.block_sizes())
|
|
)
|
|
|
|
@staticmethod
|
|
def block_sizes_log2():
|
|
return range(4, 9) if torch.version.hip is not None else range(4, 8)
|
|
|
|
@staticmethod
|
|
def block_sizes():
|
|
return [2**x for x in RaggedTensorMetadata.block_sizes_log2()]
|
|
|
|
|
|
def exact_div(x, y):
|
|
assert x % y == 0
|
|
return x // y
|
|
|
|
|
|
def empty_aligned(shape, dtype, device, pad_size):
|
|
pad = lambda x: cdiv(x, pad_size) * pad_size
|
|
ret = torch.empty((*shape[:-1], pad(shape[-1])), dtype=dtype, device=device)
|
|
ret_slices = (*[slice(None)] * (len(shape) - 1), slice(0, shape[-1]))
|
|
return ret[ret_slices], ret.numel()
|
|
|
|
|
|
@triton.jit
|
|
def _cdiv_pow2(n, log2_k):
|
|
# ceil_div(n, 2**log2_k)
|
|
return (n + ((1 << log2_k) - 1)) >> log2_k
|
|
|
|
|
|
@triton.jit
|
|
def _ragged_tensor_metadata_memset(
|
|
SliceSizes,
|
|
n_slices,
|
|
BlockOffs,
|
|
slice_offs_stride_m,
|
|
BlockSchedule,
|
|
first_block_size_log2,
|
|
SIZES: tl.constexpr,
|
|
BLOCK: tl.constexpr,
|
|
):
|
|
pid = tl.program_id(0)
|
|
if pid <= SIZES:
|
|
BlockOffs += pid * slice_offs_stride_m
|
|
BlockOffsPtrs = BlockOffs + tl.arange(0, BLOCK)
|
|
block_size_log2 = tl.where(pid == 0, 0, pid + first_block_size_log2 - 1)
|
|
# total number of blocks in slice processed as the loop iterates
|
|
n_blocks_tot = tl.zeros([BLOCK], dtype=BlockOffs.dtype.element_ty)
|
|
for i in range(0, n_slices + 1, BLOCK):
|
|
# load slice sizes
|
|
offs = tl.arange(0, BLOCK) + i
|
|
mask = offs < n_slices
|
|
slice_sizes = tl.load(SliceSizes + offs, mask=mask, other=0)
|
|
# number of blocks in the slices loaded
|
|
n_blocks = _cdiv_pow2(slice_sizes, block_size_log2)
|
|
# start index of the blocks for the slices loaded
|
|
block_starts = tl.cumsum(n_blocks, 0) + n_blocks_tot
|
|
n_blocks_tot += tl.sum(n_blocks, 0)
|
|
tl.store(BlockOffsPtrs, block_starts - n_blocks)
|
|
BlockOffsPtrs += BLOCK
|
|
else:
|
|
# initialize block schedule to -1
|
|
pid -= SIZES + 1
|
|
offs = pid * BLOCK + tl.arange(0, BLOCK)
|
|
tl.store(BlockSchedule + offs, 0xFFFFFFFF)
|
|
|
|
|
|
@triton.jit
|
|
def _ragged_tensor_metadata_compute(
|
|
SliceSizes, #
|
|
BlockOffs,
|
|
block_offs_stride_m, #
|
|
BlockSchedule,
|
|
block_schedule_stride_m, #
|
|
first_block_size_log2, #
|
|
SIZES: tl.constexpr,
|
|
BLOCK: tl.constexpr,
|
|
):
|
|
pid = tl.program_id(0)
|
|
slice_id = pid // SIZES
|
|
block_size_id = pid % SIZES
|
|
# offset pointers
|
|
BlockOffs += block_size_id * block_offs_stride_m
|
|
BlockSchedule += block_size_id * block_schedule_stride_m
|
|
# load slice sizes
|
|
slice_sizes = tl.load(SliceSizes + slice_id)
|
|
# number of blocks in the slices loaded
|
|
block_size_log2 = first_block_size_log2 + block_size_id
|
|
n_blocks = _cdiv_pow2(slice_sizes, block_size_log2)
|
|
# compute block schedule
|
|
block_off = tl.load(BlockOffs + slice_id)
|
|
BlockSchedule += block_off
|
|
for block_off in range(0, n_blocks, BLOCK):
|
|
block_offs = block_off + tl.arange(0, BLOCK)
|
|
data = (block_offs << 16) + slice_id
|
|
tl.store(BlockSchedule + block_offs, data, mask=block_offs < n_blocks)
|
|
|
|
|
|
def make_ragged_tensor_metadata(slice_sizes, n_total_rows):
|
|
assert slice_sizes.ndim == 1
|
|
n_slices = slice_sizes.shape[0]
|
|
block_sizes_log2 = RaggedTensorMetadata.block_sizes_log2()
|
|
block_size_num = len(block_sizes_log2)
|
|
MEMSET_BLOCK = 512
|
|
dtype = torch.int32
|
|
device = slice_sizes.device
|
|
max_n_blocks = RaggedTensorMetadata.max_n_blocks(n_slices, n_total_rows)
|
|
slice_offs_combined, _ = empty_aligned(
|
|
(block_size_num + 1, n_slices + 1), dtype, device, MEMSET_BLOCK
|
|
)
|
|
block_schedule_data, n_memset_elts = empty_aligned(
|
|
(block_size_num, max_n_blocks), dtype, device, MEMSET_BLOCK
|
|
)
|
|
slice_offs, block_offs_data = slice_offs_combined[0], slice_offs_combined[1:]
|
|
n_memset_blocks = exact_div(n_memset_elts, MEMSET_BLOCK)
|
|
|
|
_ragged_tensor_metadata_memset[(slice_offs_combined.shape[0] + n_memset_blocks,)](
|
|
slice_sizes,
|
|
n_slices, #
|
|
slice_offs_combined,
|
|
slice_offs_combined.stride(0), #
|
|
block_schedule_data, #
|
|
block_sizes_log2[0],
|
|
SIZES=len(block_sizes_log2),
|
|
BLOCK=MEMSET_BLOCK, # optimization parameters
|
|
num_warps=4,
|
|
)
|
|
|
|
_ragged_tensor_metadata_compute[(block_size_num * n_slices,)](
|
|
slice_sizes,
|
|
block_offs_data,
|
|
block_offs_data.stride(0),
|
|
block_schedule_data,
|
|
block_schedule_data.stride(0), # outputs
|
|
block_sizes_log2[0],
|
|
SIZES=len(block_sizes_log2),
|
|
BLOCK=512, # optimization parameters
|
|
num_warps=4,
|
|
)
|
|
|
|
return RaggedTensorMetadata(
|
|
slice_sizes, slice_offs, block_offs_data, block_schedule_data
|
|
)
|
|
|
|
|
|
__all__ = [
|
|
"topk",
|
|
"make_ragged_tensor_metadata",
|
|
]
|