[Inkling] silu_and_mul: replace helion kernels with plain Triton (#33903)
Co-authored-by: Claude Fable 5 <noreply@anthropic.com>
This commit is contained in:
Sam Shleifer
co-authored by
Claude Fable 5
Claude Fable 5
parent
d238e36b24
commit
afb4f37ca5
@@ -33,7 +33,6 @@ dependencies = [
|
||||
"flash-attn-4>=4.0.0b18",
|
||||
"flashinfer_python[cu13]==0.6.15.post1", # keep it aligned with jit-cache version in Dockerfile
|
||||
"gguf",
|
||||
"helion==1.4",
|
||||
"humming-kernels[cu13]==0.1.10",
|
||||
"interegular",
|
||||
"IPython",
|
||||
|
||||
@@ -32,7 +32,6 @@ runtime_base = [
|
||||
"einops",
|
||||
"fastapi",
|
||||
"gguf",
|
||||
"helion==1.4",
|
||||
"interegular",
|
||||
"IPython",
|
||||
"llguidance>=1.7.6,<2.0.0",
|
||||
|
||||
@@ -1,142 +1,195 @@
|
||||
from functools import partial
|
||||
|
||||
import helion
|
||||
import helion.language as hl
|
||||
import torch
|
||||
import triton
|
||||
import triton.language as tl
|
||||
|
||||
from sglang.kernels.jit.utils import is_arch_support_pdl
|
||||
from sglang.srt.layers.moe.moe_runner.triton_utils.helion_utils import (
|
||||
get_model_depths,
|
||||
helion_aot_autotune,
|
||||
)
|
||||
from sglang.srt.utils.common import is_sm121
|
||||
|
||||
DEFAULT_BLOCK_SIZE = 4096
|
||||
BLOCK_SIZE_M = 128
|
||||
|
||||
|
||||
def silu_and_mul_key(
|
||||
gateup_output: torch.Tensor,
|
||||
topk_weights: torch.Tensor | None,
|
||||
out_dtype: object | None = None,
|
||||
@triton.jit(do_not_specialize=["M"])
|
||||
def silu_and_mul_interleaved_kernel(
|
||||
gateup_out_ptr, # type: ignore # [M, 2N], rows are [g0, u0, g1, u1, ...]
|
||||
topk_weights_ptr, # type: ignore # [M] (unused when HAS_TOPK_WEIGHTS=False)
|
||||
down_inp_ptr, # type: ignore # [M, N]
|
||||
M, # type: ignore
|
||||
stride_gm, # type: ignore
|
||||
N: tl.constexpr,
|
||||
HAS_TOPK_WEIGHTS: tl.constexpr,
|
||||
INT64_INDEX: tl.constexpr,
|
||||
BLOCK_M: tl.constexpr,
|
||||
BLOCK_N: tl.constexpr,
|
||||
):
|
||||
# Keep this stable across inputs for Helion AOT autotune.
|
||||
del out_dtype
|
||||
return gateup_output.shape[1], (gateup_output.dtype,), (topk_weights is not None,)
|
||||
# Flat grid, N-tile fastest: consecutive CTAs read consecutive spans of
|
||||
# the same rows, which measures faster than row-fastest rasterization.
|
||||
pid = tl.program_id(0)
|
||||
if INT64_INDEX:
|
||||
pid = pid.to(tl.int64)
|
||||
NUM_BLOCKS_N: tl.constexpr = tl.cdiv(N, BLOCK_N)
|
||||
pid_m = pid // NUM_BLOCKS_N
|
||||
pid_n = pid % NUM_BLOCKS_N
|
||||
|
||||
offs_m = pid_m * BLOCK_M + tl.arange(0, BLOCK_M)
|
||||
offs_n = pid_n * BLOCK_N + tl.arange(0, BLOCK_N)
|
||||
mask_n = offs_n < N
|
||||
offs_2n = pid_n * (2 * BLOCK_N) + tl.arange(0, 2 * BLOCK_N)
|
||||
mask_2n = offs_2n < 2 * N
|
||||
if BLOCK_M == 1:
|
||||
# The grid covers M exactly; eliding the row mask measurably helps
|
||||
# at bandwidth-bound shapes.
|
||||
mask_mn = mask_n[None, :]
|
||||
mask_m2n = mask_2n[None, :]
|
||||
else:
|
||||
mask_m = offs_m < M
|
||||
mask_mn = mask_m[:, None] & mask_n[None, :]
|
||||
mask_m2n = mask_m[:, None] & mask_2n[None, :]
|
||||
|
||||
def silu_and_mul_inputs(sizes: list[int]):
|
||||
# Used only for Helion autotune input generation.
|
||||
inputs = []
|
||||
numel = 2**30
|
||||
with torch.device("cuda"):
|
||||
for size in sizes:
|
||||
x = torch.randn(numel // size, 2 * size, dtype=torch.bfloat16)
|
||||
inputs.append((x, None, None))
|
||||
inputs.append((x, torch.randn(numel // size, dtype=torch.bfloat16), None))
|
||||
return inputs
|
||||
|
||||
|
||||
@helion_aot_autotune(
|
||||
"silu_and_mul_interleaved",
|
||||
kernel_key=silu_and_mul_key,
|
||||
primary_inputs=partial(silu_and_mul_inputs, sizes=[512, 2048, 48 * 96, 6144, 8192]),
|
||||
secondary_inputs=partial(
|
||||
silu_and_mul_inputs,
|
||||
sizes=[512]
|
||||
+ [i * 96 for i in get_model_depths()]
|
||||
+ list(range(1024, 8192 + 1, 1024)),
|
||||
),
|
||||
)
|
||||
@helion.kernel(static_shapes=False)
|
||||
def _silu_and_mul_helion_interleaved_kernel(
|
||||
gateup_output,
|
||||
topk_weights: torch.Tensor | None = None,
|
||||
out_dtype: hl.constexpr | None = None,
|
||||
):
|
||||
"""
|
||||
Interleaved version of silu_and_mul using Helion kernel.
|
||||
Input format: [gate[0], up[0], gate[1], up[1], ...]
|
||||
This matches the interleaved w13 weight format.
|
||||
"""
|
||||
batch_size, hidden_size = gateup_output.shape
|
||||
hidden_size = hl.specialize(hidden_size)
|
||||
assert hidden_size % 2 == 0, f"{hidden_size=}"
|
||||
|
||||
half_hidden_size = hidden_size // 2
|
||||
down_input = gateup_output.new_empty(
|
||||
batch_size, half_hidden_size, dtype=out_dtype or gateup_output.dtype
|
||||
# Contiguous load of the [g, u] pairs, deinterleaved in registers via
|
||||
# tl.split; element-strided gather loads halve the effective bandwidth.
|
||||
gateup = tl.load(
|
||||
gateup_out_ptr + offs_m[:, None] * stride_gm + offs_2n[None, :],
|
||||
mask=mask_m2n,
|
||||
other=0.0,
|
||||
)
|
||||
for batch_tile, hidden_tile in hl.tile([batch_size, half_hidden_size]):
|
||||
gate_output = gateup_output[batch_tile, 2 * hidden_tile.index].to(torch.float32)
|
||||
up_output = gateup_output[batch_tile, 2 * hidden_tile.index + 1].to(
|
||||
torch.float32
|
||||
)
|
||||
silu_mul_output = gate_output * torch.sigmoid(gate_output) * up_output
|
||||
if topk_weights is not None:
|
||||
weight_scale = topk_weights[batch_tile, None].to(torch.float32)
|
||||
silu_mul_output = silu_mul_output * weight_scale
|
||||
down_input[batch_tile, hidden_tile] = silu_mul_output
|
||||
return down_input
|
||||
gate, up = tl.split(tl.reshape(gateup, (BLOCK_M, BLOCK_N, 2)))
|
||||
gate = gate.to(tl.float32)
|
||||
up = up.to(tl.float32)
|
||||
|
||||
down_inp = gate * tl.sigmoid(gate) * up
|
||||
if HAS_TOPK_WEIGHTS:
|
||||
if BLOCK_M == 1:
|
||||
weight_scale = tl.load(
|
||||
topk_weights_ptr + offs_m, eviction_policy="evict_last"
|
||||
).to(tl.float32)
|
||||
else:
|
||||
weight_scale = tl.load(
|
||||
topk_weights_ptr + offs_m,
|
||||
mask=offs_m < M,
|
||||
eviction_policy="evict_last",
|
||||
).to(tl.float32)
|
||||
down_inp = down_inp * weight_scale[:, None]
|
||||
|
||||
@helion_aot_autotune(
|
||||
"silu_and_mul",
|
||||
kernel_key=silu_and_mul_key,
|
||||
primary_inputs=partial(silu_and_mul_inputs, sizes=[512, 2048, 48 * 96, 6144, 8192]),
|
||||
secondary_inputs=partial(
|
||||
silu_and_mul_inputs,
|
||||
sizes=[512]
|
||||
+ [i * 96 for i in get_model_depths()]
|
||||
+ list(range(1024, 8192 + 1, 1024)),
|
||||
),
|
||||
)
|
||||
@helion.kernel(static_shapes=False)
|
||||
def _silu_and_mul_helion_non_interleaved_kernel(
|
||||
gateup_output,
|
||||
topk_weights: torch.Tensor | None = None,
|
||||
out_dtype: hl.constexpr | None = None,
|
||||
):
|
||||
"""
|
||||
Non-interleaved version of silu_and_mul using Helion kernel.
|
||||
Input format: [gate[0], gate[1], ..., gate[N-1], up[0], up[1], ..., up[N-1]]
|
||||
"""
|
||||
batch_size, hidden_size = gateup_output.shape
|
||||
hidden_size = hl.specialize(hidden_size)
|
||||
assert hidden_size % 2 == 0, f"{hidden_size=}"
|
||||
|
||||
half_hidden_size = hidden_size // 2
|
||||
down_input = gateup_output.new_empty(
|
||||
batch_size, half_hidden_size, dtype=out_dtype or gateup_output.dtype
|
||||
offs_mn = offs_m[:, None] * N + offs_n[None, :]
|
||||
tl.store(
|
||||
down_inp_ptr + offs_mn,
|
||||
down_inp.to(down_inp_ptr.dtype.element_ty),
|
||||
mask=mask_mn,
|
||||
)
|
||||
for batch_tile, hidden_tile in hl.tile([batch_size, half_hidden_size]):
|
||||
gate_output = gateup_output[batch_tile, hidden_tile.index].to(torch.float32)
|
||||
up_output = gateup_output[batch_tile, hidden_tile.index + half_hidden_size].to(
|
||||
torch.float32
|
||||
)
|
||||
silu_mul_output = gate_output * torch.sigmoid(gate_output) * up_output
|
||||
if topk_weights is not None:
|
||||
weight_scale = topk_weights[batch_tile, None].to(torch.float32)
|
||||
silu_mul_output = silu_mul_output * weight_scale
|
||||
down_input[batch_tile, hidden_tile] = silu_mul_output
|
||||
return down_input
|
||||
|
||||
|
||||
def silu_and_mul_helion(
|
||||
@triton.jit(do_not_specialize=["M"])
|
||||
def silu_and_mul_non_interleaved_kernel(
|
||||
gateup_out_ptr, # type: ignore # [M, 2N], rows are [g0, ..., gN-1, u0, ..., uN-1]
|
||||
topk_weights_ptr, # type: ignore # [M] (unused when HAS_TOPK_WEIGHTS=False)
|
||||
down_inp_ptr, # type: ignore # [M, N]
|
||||
M, # type: ignore
|
||||
stride_gm, # type: ignore
|
||||
N: tl.constexpr,
|
||||
HAS_TOPK_WEIGHTS: tl.constexpr,
|
||||
INT64_INDEX: tl.constexpr,
|
||||
BLOCK_M: tl.constexpr,
|
||||
BLOCK_N: tl.constexpr,
|
||||
):
|
||||
# Flat grid, N-tile fastest: consecutive CTAs read consecutive spans of
|
||||
# the same rows, which measures faster than row-fastest rasterization.
|
||||
pid = tl.program_id(0)
|
||||
if INT64_INDEX:
|
||||
pid = pid.to(tl.int64)
|
||||
NUM_BLOCKS_N: tl.constexpr = tl.cdiv(N, BLOCK_N)
|
||||
pid_m = pid // NUM_BLOCKS_N
|
||||
pid_n = pid % NUM_BLOCKS_N
|
||||
|
||||
offs_m = pid_m * BLOCK_M + tl.arange(0, BLOCK_M)
|
||||
offs_n = pid_n * BLOCK_N + tl.arange(0, BLOCK_N)
|
||||
mask_n = offs_n < N
|
||||
if BLOCK_M == 1:
|
||||
# The grid covers M exactly; eliding the row mask measurably helps
|
||||
# at bandwidth-bound shapes.
|
||||
mask_mn = mask_n[None, :]
|
||||
else:
|
||||
mask_m = offs_m < M
|
||||
mask_mn = mask_m[:, None] & mask_n[None, :]
|
||||
|
||||
# Cache hints match the tuned helion config (numerics-neutral): the gate
|
||||
# stream stays in L2 while the up stream of the same rows loads.
|
||||
offs_gate = offs_m[:, None] * stride_gm + offs_n[None, :]
|
||||
gate = tl.load(
|
||||
gateup_out_ptr + offs_gate,
|
||||
mask=mask_mn,
|
||||
other=0.0,
|
||||
eviction_policy="evict_last",
|
||||
).to(tl.float32)
|
||||
up = tl.load(
|
||||
gateup_out_ptr + offs_gate + N,
|
||||
mask=mask_mn,
|
||||
other=0.0,
|
||||
eviction_policy="evict_first",
|
||||
).to(tl.float32)
|
||||
|
||||
down_inp = gate * tl.sigmoid(gate) * up
|
||||
if HAS_TOPK_WEIGHTS:
|
||||
if BLOCK_M == 1:
|
||||
weight_scale = tl.load(
|
||||
topk_weights_ptr + offs_m, eviction_policy="evict_last"
|
||||
).to(tl.float32)
|
||||
else:
|
||||
weight_scale = tl.load(
|
||||
topk_weights_ptr + offs_m,
|
||||
mask=offs_m < M,
|
||||
eviction_policy="evict_last",
|
||||
).to(tl.float32)
|
||||
down_inp = down_inp * weight_scale[:, None]
|
||||
|
||||
offs_mn = offs_m[:, None] * N + offs_n[None, :]
|
||||
tl.store(
|
||||
down_inp_ptr + offs_mn,
|
||||
down_inp.to(down_inp_ptr.dtype.element_ty),
|
||||
mask=mask_mn,
|
||||
)
|
||||
|
||||
|
||||
def _select_silu_and_mul_config(
|
||||
M: int, N: int, interleaved: bool
|
||||
) -> tuple[int, int, int]:
|
||||
"""(BLOCK_M, BLOCK_N, num_warps) heuristic from a GB300 (sm_100) sweep.
|
||||
|
||||
One row per CTA keeps the grid large enough to fill the GPU at decode row
|
||||
counts. BLOCK_N=512 avoids a quarter-masked tail block for N=1536-style
|
||||
widths; other widths prefer the widest contiguous load even when the tail
|
||||
block is masked. The non-interleaved kernel reads two separate streams per
|
||||
row and prefers (2, 512) tiles at prefill row counts.
|
||||
"""
|
||||
if not interleaved and M > 4096 and N % 512 == 0:
|
||||
return 2, 512, 1
|
||||
BLOCK_M = 1
|
||||
if N % 512 == 0 and N % 1024 != 0:
|
||||
BLOCK_N = 512
|
||||
else:
|
||||
BLOCK_N = min(1024, triton.next_power_of_2(N))
|
||||
if M <= 4096:
|
||||
num_warps = 4
|
||||
else:
|
||||
num_warps = 2 if interleaved else 1
|
||||
return BLOCK_M, BLOCK_N, num_warps
|
||||
|
||||
|
||||
def silu_and_mul(
|
||||
gateup_output: torch.Tensor,
|
||||
topk_weights: torch.Tensor | None = None,
|
||||
out_dtype: torch.dtype | None = None,
|
||||
use_interleaved: bool = True,
|
||||
) -> torch.Tensor:
|
||||
"""
|
||||
Unified silu_and_mul function using Helion kernel.
|
||||
Supports both interleaved and non-interleaved input formats.
|
||||
"""silu(gate) * up over routed rows, optionally scaled by per-row weights.
|
||||
|
||||
Bitwise-identical port of the former Helion kernels (fp32 math,
|
||||
gate * sigmoid(gate) * up, weight scale last, one rounding cast at the
|
||||
store).
|
||||
|
||||
Args:
|
||||
gateup_output: Input tensor of shape (batch_size, hidden_size)
|
||||
topk_weights: Optional topk weights tensor
|
||||
topk_weights: Optional per-row weights tensor of shape (batch_size,)
|
||||
out_dtype: Optional output dtype
|
||||
use_interleaved: If True, expects interleaved format [gate[0], up[0], gate[1], up[1], ...]
|
||||
If False, expects non-interleaved format [gate[0], ..., gate[N-1], up[0], ..., up[N-1]]
|
||||
@@ -144,20 +197,51 @@ def silu_and_mul_helion(
|
||||
Returns:
|
||||
Output tensor of shape (batch_size, hidden_size // 2)
|
||||
"""
|
||||
if use_interleaved:
|
||||
return _silu_and_mul_helion_interleaved_kernel(
|
||||
gateup_output, topk_weights, out_dtype
|
||||
)
|
||||
else:
|
||||
return _silu_and_mul_helion_non_interleaved_kernel(
|
||||
gateup_output, topk_weights, out_dtype
|
||||
)
|
||||
assert gateup_output.ndim == 2, f"{gateup_output.shape=}"
|
||||
assert gateup_output.stride(1) == 1, f"{gateup_output.stride()=}"
|
||||
M, hidden_size = gateup_output.shape
|
||||
assert hidden_size % 2 == 0, f"{hidden_size=}"
|
||||
N = hidden_size // 2
|
||||
if topk_weights is not None:
|
||||
assert topk_weights.is_contiguous(), f"{topk_weights.stride()=}"
|
||||
assert topk_weights.shape == (M,), f"{topk_weights.shape=} {M=}"
|
||||
|
||||
down_input = torch.empty(
|
||||
(M, N), device=gateup_output.device, dtype=out_dtype or gateup_output.dtype
|
||||
)
|
||||
if M == 0:
|
||||
return down_input
|
||||
|
||||
BLOCK_M, BLOCK_N, num_warps = _select_silu_and_mul_config(M, N, use_interleaved)
|
||||
grid = (triton.cdiv(M, BLOCK_M) * triton.cdiv(N, BLOCK_N),)
|
||||
kernel = (
|
||||
silu_and_mul_interleaved_kernel
|
||||
if use_interleaved
|
||||
else silu_and_mul_non_interleaved_kernel
|
||||
)
|
||||
kernel[grid](
|
||||
gateup_out_ptr=gateup_output,
|
||||
topk_weights_ptr=topk_weights,
|
||||
down_inp_ptr=down_input,
|
||||
M=M,
|
||||
stride_gm=gateup_output.stride(0),
|
||||
N=N,
|
||||
HAS_TOPK_WEIGHTS=topk_weights is not None,
|
||||
INT64_INDEX=gateup_output.nbytes >= 2**31 or down_input.nbytes >= 2**31,
|
||||
BLOCK_M=BLOCK_M,
|
||||
BLOCK_N=BLOCK_N,
|
||||
num_warps=num_warps,
|
||||
)
|
||||
return down_input
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Triton silu_and_mul
|
||||
# Used by InklingBatchDenseMLP._swiglu because the helion kernel above produces
|
||||
# NaN for small shared-expert batches in EP+DP configs.
|
||||
# Persistent-grid Triton silu_and_mul
|
||||
# Used by InklingBatchDenseMLP._swiglu. It predates the plain-Triton port of
|
||||
# silu_and_mul above: the helion kernels that silu_and_mul replaced were
|
||||
# reported to produce NaN for small shared-expert batches in EP+DP configs.
|
||||
# Unlike silu_and_mul it can read the row count from a device tensor (M_ptr)
|
||||
# and supports PDL.
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
|
||||
-39
@@ -1,39 +0,0 @@
|
||||
{
|
||||
"useful_configs": {
|
||||
"0": "helion.Config(block_sizes=[2, 512], indexing=['tensor_descriptor', 'tensor_descriptor', 'tensor_descriptor', 'tensor_descriptor'], l2_groupings=[1], load_eviction_policies=['last', '', ''], loop_orders=[[1, 0]], num_stages=3, num_warps=1, pid_type='flat', range_flattens=[None], range_multi_buffers=[None], range_num_stages=[0], range_unroll_factors=[0], range_warp_specializes=[None])",
|
||||
"1": "helion.Config(block_sizes=[1, 1024], indexing=['tensor_descriptor', 'tensor_descriptor', 'pointer', 'pointer'], l2_groupings=[1], load_eviction_policies=['last', 'first', 'last'], loop_orders=[[1, 0]], num_stages=5, num_warps=1, pid_type='flat', range_flattens=[None], range_multi_buffers=[None], range_num_stages=[0], range_unroll_factors=[0], range_warp_specializes=[None])",
|
||||
"4": "helion.Config(block_sizes=[2, 512], indexing=['pointer', 'tensor_descriptor', 'pointer'], l2_groupings=[2], load_eviction_policies=['last', ''], loop_orders=[[0, 1]], num_stages=2, num_warps=1, pid_type='flat', range_flattens=[None], range_multi_buffers=[None], range_num_stages=[0], range_unroll_factors=[0], range_warp_specializes=[None])"
|
||||
},
|
||||
"hash_configs": {
|
||||
"(512, (2,), (False,))": 0,
|
||||
"(512, (2,), (True,))": 1,
|
||||
"(1024, (2,), (False,))": 0,
|
||||
"(1024, (2,), (True,))": 0,
|
||||
"(1536, (2,), (False,))": 1,
|
||||
"(1536, (2,), (True,))": 1,
|
||||
"(2048, (2,), (False,))": 4,
|
||||
"(2048, (2,), (True,))": 1,
|
||||
"(3072, (2,), (False,))": 0,
|
||||
"(3072, (2,), (True,))": 0,
|
||||
"(4096, (2,), (False,))": 4,
|
||||
"(4096, (2,), (True,))": 1,
|
||||
"(4608, (2,), (False,))": 0,
|
||||
"(4608, (2,), (True,))": 0,
|
||||
"(6144, (2,), (False,))": 0,
|
||||
"(6144, (2,), (True,))": 1,
|
||||
"(7680, (2,), (False,))": 1,
|
||||
"(7680, (2,), (True,))": 1,
|
||||
"(8192, (2,), (False,))": 4,
|
||||
"(8192, (2,), (True,))": 1,
|
||||
"(9216, (2,), (False,))": 0,
|
||||
"(9216, (2,), (True,))": 0,
|
||||
"(10240, (2,), (False,))": 0,
|
||||
"(10240, (2,), (True,))": 1,
|
||||
"(12288, (2,), (False,))": 0,
|
||||
"(12288, (2,), (True,))": 1,
|
||||
"(14336, (2,), (False,))": 0,
|
||||
"(14336, (2,), (True,))": 1,
|
||||
"(16384, (2,), (False,))": 4,
|
||||
"(16384, (2,), (True,))": 1
|
||||
}
|
||||
}
|
||||
-39
@@ -1,39 +0,0 @@
|
||||
{
|
||||
"useful_configs": {
|
||||
"0": "helion.Config(block_sizes=[2, 512], indexing=['tensor_descriptor', 'tensor_descriptor', 'tensor_descriptor', 'tensor_descriptor'], l2_groupings=[1], load_eviction_policies=['last', '', ''], loop_orders=[[1, 0]], num_stages=3, num_warps=1, pid_type='flat', range_flattens=[None], range_multi_buffers=[None], range_num_stages=[0], range_unroll_factors=[0], range_warp_specializes=[None])",
|
||||
"1": "helion.Config(block_sizes=[1, 1024], indexing=['tensor_descriptor', 'tensor_descriptor', 'pointer', 'pointer'], l2_groupings=[1], load_eviction_policies=['last', 'first', 'last'], loop_orders=[[1, 0]], num_stages=5, num_warps=1, pid_type='flat', range_flattens=[None], range_multi_buffers=[None], range_num_stages=[0], range_unroll_factors=[0], range_warp_specializes=[None])",
|
||||
"4": "helion.Config(block_sizes=[2, 512], indexing=['pointer', 'tensor_descriptor', 'pointer'], l2_groupings=[2], load_eviction_policies=['last', ''], loop_orders=[[0, 1]], num_stages=2, num_warps=1, pid_type='flat', range_flattens=[None], range_multi_buffers=[None], range_num_stages=[0], range_unroll_factors=[0], range_warp_specializes=[None])"
|
||||
},
|
||||
"hash_configs": {
|
||||
"(512, (2,), (False,))": 0,
|
||||
"(512, (2,), (True,))": 1,
|
||||
"(1024, (2,), (False,))": 0,
|
||||
"(1024, (2,), (True,))": 0,
|
||||
"(1536, (2,), (False,))": 1,
|
||||
"(1536, (2,), (True,))": 1,
|
||||
"(2048, (2,), (False,))": 4,
|
||||
"(2048, (2,), (True,))": 1,
|
||||
"(3072, (2,), (False,))": 0,
|
||||
"(3072, (2,), (True,))": 0,
|
||||
"(4096, (2,), (False,))": 4,
|
||||
"(4096, (2,), (True,))": 1,
|
||||
"(4608, (2,), (False,))": 0,
|
||||
"(4608, (2,), (True,))": 0,
|
||||
"(6144, (2,), (False,))": 0,
|
||||
"(6144, (2,), (True,))": 1,
|
||||
"(7680, (2,), (False,))": 1,
|
||||
"(7680, (2,), (True,))": 1,
|
||||
"(8192, (2,), (False,))": 4,
|
||||
"(8192, (2,), (True,))": 1,
|
||||
"(9216, (2,), (False,))": 0,
|
||||
"(9216, (2,), (True,))": 0,
|
||||
"(10240, (2,), (False,))": 0,
|
||||
"(10240, (2,), (True,))": 1,
|
||||
"(12288, (2,), (False,))": 0,
|
||||
"(12288, (2,), (True,))": 1,
|
||||
"(14336, (2,), (False,))": 0,
|
||||
"(14336, (2,), (True,))": 1,
|
||||
"(16384, (2,), (False,))": 4,
|
||||
"(16384, (2,), (True,))": 1
|
||||
}
|
||||
}
|
||||
-40
@@ -1,40 +0,0 @@
|
||||
{
|
||||
"useful_configs": {
|
||||
"0": "helion.Config(block_sizes=[2, 512], indexing=['pointer', 'pointer', 'tensor_descriptor'], l2_groupings=[2], load_eviction_policies=['', ''], loop_orders=[[0, 1]], num_stages=2, num_warps=1, pid_type='flat', range_flattens=[None], range_multi_buffers=[None], range_num_stages=[0], range_unroll_factors=[0], range_warp_specializes=[])",
|
||||
"1": "helion.Config(block_sizes=[1, 512], indexing=['pointer', 'tensor_descriptor', 'tensor_descriptor', 'pointer'], l2_groupings=[1], load_eviction_policies=['', 'first', 'first'], loop_orders=[[1, 0]], num_stages=7, num_warps=1, pid_type='flat', range_flattens=[None], range_multi_buffers=[None], range_num_stages=[0], range_unroll_factors=[0], range_warp_specializes=[])",
|
||||
"2": "helion.Config(block_sizes=[1, 1024], indexing=['pointer', 'pointer', 'pointer', 'pointer'], l2_groupings=[1], load_eviction_policies=['last', '', 'last'], loop_orders=[[1, 0]], num_stages=3, num_warps=1, pid_type='flat', range_flattens=[None], range_multi_buffers=[None], range_num_stages=[0], range_unroll_factors=[0], range_warp_specializes=[])",
|
||||
"3": "helion.Config(block_sizes=[1, 1024], indexing=['pointer', 'tensor_descriptor', 'tensor_descriptor', 'tensor_descriptor'], l2_groupings=[1], load_eviction_policies=['', 'first', ''], loop_orders=[[1, 0]], num_stages=7, num_warps=2, pid_type='flat', range_flattens=[None], range_multi_buffers=[None], range_num_stages=[0], range_unroll_factors=[0], range_warp_specializes=[])"
|
||||
},
|
||||
"hash_configs": {
|
||||
"(512, (2,), (False,))": 0,
|
||||
"(512, (2,), (True,))": 1,
|
||||
"(1024, (2,), (False,))": 0,
|
||||
"(1024, (2,), (True,))": 1,
|
||||
"(1536, (2,), (False,))": 2,
|
||||
"(1536, (2,), (True,))": 2,
|
||||
"(2048, (2,), (False,))": 3,
|
||||
"(2048, (2,), (True,))": 3,
|
||||
"(3072, (2,), (False,))": 1,
|
||||
"(3072, (2,), (True,))": 1,
|
||||
"(4096, (2,), (False,))": 3,
|
||||
"(4096, (2,), (True,))": 3,
|
||||
"(4608, (2,), (False,))": 1,
|
||||
"(4608, (2,), (True,))": 1,
|
||||
"(6144, (2,), (False,))": 3,
|
||||
"(6144, (2,), (True,))": 3,
|
||||
"(7680, (2,), (False,))": 2,
|
||||
"(7680, (2,), (True,))": 2,
|
||||
"(8192, (2,), (False,))": 3,
|
||||
"(8192, (2,), (True,))": 3,
|
||||
"(9216, (2,), (False,))": 1,
|
||||
"(9216, (2,), (True,))": 1,
|
||||
"(10240, (2,), (False,))": 3,
|
||||
"(10240, (2,), (True,))": 3,
|
||||
"(12288, (2,), (False,))": 3,
|
||||
"(12288, (2,), (True,))": 3,
|
||||
"(14336, (2,), (False,))": 3,
|
||||
"(14336, (2,), (True,))": 3,
|
||||
"(16384, (2,), (False,))": 3,
|
||||
"(16384, (2,), (True,))": 3
|
||||
}
|
||||
}
|
||||
-40
@@ -1,40 +0,0 @@
|
||||
{
|
||||
"useful_configs": {
|
||||
"0": "helion.Config(block_sizes=[2, 512], indexing=['pointer', 'pointer', 'pointer'], l2_groupings=[2], load_eviction_policies=['', ''], loop_orders=[[0, 1]], num_stages=2, num_warps=1, pid_type='flat', range_flattens=[None], range_multi_buffers=[None], range_num_stages=[0], range_unroll_factors=[0], range_warp_specializes=[])",
|
||||
"1": "helion.Config(block_sizes=[1, 512], indexing=['pointer', 'pointer', 'pointer', 'pointer'], l2_groupings=[1], load_eviction_policies=['', 'first', 'first'], loop_orders=[[1, 0]], num_stages=7, num_warps=1, pid_type='flat', range_flattens=[None], range_multi_buffers=[None], range_num_stages=[0], range_unroll_factors=[0], range_warp_specializes=[])",
|
||||
"2": "helion.Config(block_sizes=[1, 1024], indexing=['pointer', 'pointer', 'pointer', 'pointer'], l2_groupings=[1], load_eviction_policies=['last', '', 'last'], loop_orders=[[1, 0]], num_stages=3, num_warps=1, pid_type='flat', range_flattens=[None], range_multi_buffers=[None], range_num_stages=[0], range_unroll_factors=[0], range_warp_specializes=[])",
|
||||
"3": "helion.Config(block_sizes=[1, 1024], indexing=['pointer', 'pointer', 'pointer', 'pointer'], l2_groupings=[1], load_eviction_policies=['', 'first', ''], loop_orders=[[1, 0]], num_stages=7, num_warps=2, pid_type='flat', range_flattens=[None], range_multi_buffers=[None], range_num_stages=[0], range_unroll_factors=[0], range_warp_specializes=[])"
|
||||
},
|
||||
"hash_configs": {
|
||||
"(512, (2,), (False,))": 0,
|
||||
"(512, (2,), (True,))": 1,
|
||||
"(1024, (2,), (False,))": 0,
|
||||
"(1024, (2,), (True,))": 1,
|
||||
"(1536, (2,), (False,))": 2,
|
||||
"(1536, (2,), (True,))": 2,
|
||||
"(2048, (2,), (False,))": 3,
|
||||
"(2048, (2,), (True,))": 3,
|
||||
"(3072, (2,), (False,))": 1,
|
||||
"(3072, (2,), (True,))": 1,
|
||||
"(4096, (2,), (False,))": 3,
|
||||
"(4096, (2,), (True,))": 3,
|
||||
"(4608, (2,), (False,))": 1,
|
||||
"(4608, (2,), (True,))": 1,
|
||||
"(6144, (2,), (False,))": 3,
|
||||
"(6144, (2,), (True,))": 3,
|
||||
"(7680, (2,), (False,))": 2,
|
||||
"(7680, (2,), (True,))": 2,
|
||||
"(8192, (2,), (False,))": 3,
|
||||
"(8192, (2,), (True,))": 3,
|
||||
"(9216, (2,), (False,))": 1,
|
||||
"(9216, (2,), (True,))": 1,
|
||||
"(10240, (2,), (False,))": 3,
|
||||
"(10240, (2,), (True,))": 3,
|
||||
"(12288, (2,), (False,))": 3,
|
||||
"(12288, (2,), (True,))": 3,
|
||||
"(14336, (2,), (False,))": 3,
|
||||
"(14336, (2,), (True,))": 3,
|
||||
"(16384, (2,), (False,))": 3,
|
||||
"(16384, (2,), (True,))": 3
|
||||
}
|
||||
}
|
||||
@@ -1,39 +0,0 @@
|
||||
{
|
||||
"useful_configs": {
|
||||
"0": "helion.Config(block_sizes=[2, 512], indexing=['tensor_descriptor', 'tensor_descriptor', 'tensor_descriptor', 'tensor_descriptor'], l2_groupings=[1], load_eviction_policies=['last', '', ''], loop_orders=[[1, 0]], num_stages=3, num_warps=1, pid_type='flat', range_flattens=[None], range_multi_buffers=[None], range_num_stages=[0], range_unroll_factors=[0], range_warp_specializes=[None])",
|
||||
"1": "helion.Config(block_sizes=[1, 1024], indexing=['tensor_descriptor', 'tensor_descriptor', 'pointer', 'pointer'], l2_groupings=[1], load_eviction_policies=['last', 'first', 'last'], loop_orders=[[1, 0]], num_stages=5, num_warps=1, pid_type='flat', range_flattens=[None], range_multi_buffers=[None], range_num_stages=[0], range_unroll_factors=[0], range_warp_specializes=[None])",
|
||||
"4": "helion.Config(block_sizes=[2, 512], indexing=['pointer', 'tensor_descriptor', 'pointer'], l2_groupings=[2], load_eviction_policies=['last', ''], loop_orders=[[0, 1]], num_stages=2, num_warps=1, pid_type='flat', range_flattens=[None], range_multi_buffers=[None], range_num_stages=[0], range_unroll_factors=[0], range_warp_specializes=[None])"
|
||||
},
|
||||
"hash_configs": {
|
||||
"(512, (2,), (False,))": 0,
|
||||
"(512, (2,), (True,))": 1,
|
||||
"(1024, (2,), (False,))": 0,
|
||||
"(1024, (2,), (True,))": 0,
|
||||
"(1536, (2,), (False,))": 1,
|
||||
"(1536, (2,), (True,))": 1,
|
||||
"(2048, (2,), (False,))": 4,
|
||||
"(2048, (2,), (True,))": 1,
|
||||
"(3072, (2,), (False,))": 0,
|
||||
"(3072, (2,), (True,))": 0,
|
||||
"(4096, (2,), (False,))": 4,
|
||||
"(4096, (2,), (True,))": 1,
|
||||
"(4608, (2,), (False,))": 0,
|
||||
"(4608, (2,), (True,))": 0,
|
||||
"(6144, (2,), (False,))": 0,
|
||||
"(6144, (2,), (True,))": 1,
|
||||
"(7680, (2,), (False,))": 1,
|
||||
"(7680, (2,), (True,))": 1,
|
||||
"(8192, (2,), (False,))": 4,
|
||||
"(8192, (2,), (True,))": 1,
|
||||
"(9216, (2,), (False,))": 0,
|
||||
"(9216, (2,), (True,))": 0,
|
||||
"(10240, (2,), (False,))": 0,
|
||||
"(10240, (2,), (True,))": 1,
|
||||
"(12288, (2,), (False,))": 0,
|
||||
"(12288, (2,), (True,))": 1,
|
||||
"(14336, (2,), (False,))": 0,
|
||||
"(14336, (2,), (True,))": 1,
|
||||
"(16384, (2,), (False,))": 4,
|
||||
"(16384, (2,), (True,))": 1
|
||||
}
|
||||
}
|
||||
@@ -1,40 +0,0 @@
|
||||
{
|
||||
"useful_configs": {
|
||||
"0": "helion.Config(block_sizes=[2, 512], indexing=['pointer', 'pointer', 'tensor_descriptor'], l2_groupings=[2], load_eviction_policies=['', ''], loop_orders=[[0, 1]], num_stages=2, num_warps=1, pid_type='flat', range_flattens=[None], range_multi_buffers=[None], range_num_stages=[0], range_unroll_factors=[0], range_warp_specializes=[])",
|
||||
"1": "helion.Config(block_sizes=[1, 512], indexing=['pointer', 'tensor_descriptor', 'tensor_descriptor', 'pointer'], l2_groupings=[1], load_eviction_policies=['', 'first', 'first'], loop_orders=[[1, 0]], num_stages=7, num_warps=1, pid_type='flat', range_flattens=[None], range_multi_buffers=[None], range_num_stages=[0], range_unroll_factors=[0], range_warp_specializes=[])",
|
||||
"2": "helion.Config(block_sizes=[1, 1024], indexing=['pointer', 'pointer', 'pointer', 'pointer'], l2_groupings=[1], load_eviction_policies=['last', '', 'last'], loop_orders=[[1, 0]], num_stages=3, num_warps=1, pid_type='flat', range_flattens=[None], range_multi_buffers=[None], range_num_stages=[0], range_unroll_factors=[0], range_warp_specializes=[])",
|
||||
"3": "helion.Config(block_sizes=[1, 1024], indexing=['pointer', 'tensor_descriptor', 'tensor_descriptor', 'tensor_descriptor'], l2_groupings=[1], load_eviction_policies=['', 'first', ''], loop_orders=[[1, 0]], num_stages=7, num_warps=2, pid_type='flat', range_flattens=[None], range_multi_buffers=[None], range_num_stages=[0], range_unroll_factors=[0], range_warp_specializes=[])"
|
||||
},
|
||||
"hash_configs": {
|
||||
"(512, (2,), (False,))": 0,
|
||||
"(512, (2,), (True,))": 1,
|
||||
"(1024, (2,), (False,))": 0,
|
||||
"(1024, (2,), (True,))": 1,
|
||||
"(1536, (2,), (False,))": 2,
|
||||
"(1536, (2,), (True,))": 2,
|
||||
"(2048, (2,), (False,))": 3,
|
||||
"(2048, (2,), (True,))": 3,
|
||||
"(3072, (2,), (False,))": 1,
|
||||
"(3072, (2,), (True,))": 1,
|
||||
"(4096, (2,), (False,))": 3,
|
||||
"(4096, (2,), (True,))": 3,
|
||||
"(4608, (2,), (False,))": 1,
|
||||
"(4608, (2,), (True,))": 1,
|
||||
"(6144, (2,), (False,))": 3,
|
||||
"(6144, (2,), (True,))": 3,
|
||||
"(7680, (2,), (False,))": 2,
|
||||
"(7680, (2,), (True,))": 2,
|
||||
"(8192, (2,), (False,))": 3,
|
||||
"(8192, (2,), (True,))": 3,
|
||||
"(9216, (2,), (False,))": 1,
|
||||
"(9216, (2,), (True,))": 1,
|
||||
"(10240, (2,), (False,))": 3,
|
||||
"(10240, (2,), (True,))": 3,
|
||||
"(12288, (2,), (False,))": 3,
|
||||
"(12288, (2,), (True,))": 3,
|
||||
"(14336, (2,), (False,))": 3,
|
||||
"(14336, (2,), (True,))": 3,
|
||||
"(16384, (2,), (False,))": 3,
|
||||
"(16384, (2,), (True,))": 3
|
||||
}
|
||||
}
|
||||
@@ -1,40 +0,0 @@
|
||||
{
|
||||
"useful_configs": {
|
||||
"0": "helion.Config(block_sizes=[2, 512], indexing=['pointer', 'pointer', 'pointer'], l2_groupings=[2], load_eviction_policies=['', ''], loop_orders=[[0, 1]], num_stages=2, num_warps=1, pid_type='flat', range_flattens=[None], range_multi_buffers=[None], range_num_stages=[0], range_unroll_factors=[0], range_warp_specializes=[])",
|
||||
"1": "helion.Config(block_sizes=[1, 512], indexing=['pointer', 'pointer', 'pointer', 'pointer'], l2_groupings=[1], load_eviction_policies=['', 'first', 'first'], loop_orders=[[1, 0]], num_stages=7, num_warps=1, pid_type='flat', range_flattens=[None], range_multi_buffers=[None], range_num_stages=[0], range_unroll_factors=[0], range_warp_specializes=[])",
|
||||
"2": "helion.Config(block_sizes=[1, 1024], indexing=['pointer', 'pointer', 'pointer', 'pointer'], l2_groupings=[1], load_eviction_policies=['last', '', 'last'], loop_orders=[[1, 0]], num_stages=3, num_warps=1, pid_type='flat', range_flattens=[None], range_multi_buffers=[None], range_num_stages=[0], range_unroll_factors=[0], range_warp_specializes=[])",
|
||||
"3": "helion.Config(block_sizes=[1, 1024], indexing=['pointer', 'pointer', 'pointer', 'pointer'], l2_groupings=[1], load_eviction_policies=['', 'first', ''], loop_orders=[[1, 0]], num_stages=7, num_warps=2, pid_type='flat', range_flattens=[None], range_multi_buffers=[None], range_num_stages=[0], range_unroll_factors=[0], range_warp_specializes=[])"
|
||||
},
|
||||
"hash_configs": {
|
||||
"(512, (2,), (False,))": 0,
|
||||
"(512, (2,), (True,))": 1,
|
||||
"(1024, (2,), (False,))": 0,
|
||||
"(1024, (2,), (True,))": 1,
|
||||
"(1536, (2,), (False,))": 2,
|
||||
"(1536, (2,), (True,))": 2,
|
||||
"(2048, (2,), (False,))": 3,
|
||||
"(2048, (2,), (True,))": 3,
|
||||
"(3072, (2,), (False,))": 1,
|
||||
"(3072, (2,), (True,))": 1,
|
||||
"(4096, (2,), (False,))": 3,
|
||||
"(4096, (2,), (True,))": 3,
|
||||
"(4608, (2,), (False,))": 1,
|
||||
"(4608, (2,), (True,))": 1,
|
||||
"(6144, (2,), (False,))": 3,
|
||||
"(6144, (2,), (True,))": 3,
|
||||
"(7680, (2,), (False,))": 2,
|
||||
"(7680, (2,), (True,))": 2,
|
||||
"(8192, (2,), (False,))": 3,
|
||||
"(8192, (2,), (True,))": 3,
|
||||
"(9216, (2,), (False,))": 1,
|
||||
"(9216, (2,), (True,))": 1,
|
||||
"(10240, (2,), (False,))": 3,
|
||||
"(10240, (2,), (True,))": 3,
|
||||
"(12288, (2,), (False,))": 3,
|
||||
"(12288, (2,), (True,))": 3,
|
||||
"(14336, (2,), (False,))": 3,
|
||||
"(14336, (2,), (True,))": 3,
|
||||
"(16384, (2,), (False,))": 3,
|
||||
"(16384, (2,), (True,))": 3
|
||||
}
|
||||
}
|
||||
@@ -1,296 +0,0 @@
|
||||
import copy
|
||||
import dataclasses
|
||||
import functools
|
||||
import json
|
||||
import logging
|
||||
import os
|
||||
from collections import defaultdict
|
||||
from collections.abc import Iterable
|
||||
from enum import Enum
|
||||
from pathlib import Path
|
||||
from typing import Any, Callable
|
||||
|
||||
import helion # noqa: F401
|
||||
import torch
|
||||
from triton.testing import do_bench
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
AutotuneInputFn = Callable[[], Iterable[tuple[Any, ...]]]
|
||||
|
||||
|
||||
def get_model_depths() -> list[int]:
|
||||
return [8, 16, 24, 32, 40, 48, 64]
|
||||
|
||||
|
||||
def get_cuda_device_capability() -> tuple[int, int]:
|
||||
if torch.cuda.is_available():
|
||||
dev = torch.cuda.current_device()
|
||||
cc_major, cc_minor = torch.cuda.get_device_capability(dev)
|
||||
return cc_major, cc_minor
|
||||
return (0, 0)
|
||||
|
||||
|
||||
class AOTAutotuneMode(Enum):
|
||||
NONE = "none"
|
||||
CREATE = "create"
|
||||
RETUNE = "retune"
|
||||
|
||||
@classmethod
|
||||
def from_str(cls, mode: str) -> "AOTAutotuneMode":
|
||||
return cls[mode.upper()]
|
||||
|
||||
|
||||
def load_autotune_data_from_json(
|
||||
path: Path,
|
||||
) -> tuple[dict[int, Any], dict[tuple[Any, Any, Any], int]]:
|
||||
with open(path, "r") as f:
|
||||
json_obj = json.load(f)
|
||||
useful_configs = {
|
||||
int(k): eval(v) for k, v in json_obj["useful_configs"].items()
|
||||
}
|
||||
hash_configs = {tuple(eval(k)): v for k, v in json_obj["hash_configs"].items()}
|
||||
return useful_configs, hash_configs
|
||||
|
||||
|
||||
def bind_and_compile_kernel(
|
||||
kernel: helion.Kernel, args: Any, config: helion.Config
|
||||
) -> Callable:
|
||||
has_int64 = any([i.numel() >= 2**31 for i in args if isinstance(i, torch.Tensor)])
|
||||
if has_int64:
|
||||
new_settings = dataclasses.replace(kernel.settings, index_dtype=torch.int64)
|
||||
kernel = helion.Kernel(
|
||||
kernel.fn,
|
||||
configs=kernel.configs,
|
||||
settings=new_settings,
|
||||
key=kernel._key_fn,
|
||||
)
|
||||
config = copy.deepcopy(config)
|
||||
return kernel.bind(args).compile_config(config, allow_print=False)
|
||||
|
||||
|
||||
def helion_aot_autotune(
|
||||
config_path: str,
|
||||
kernel_key: Callable[..., tuple[Any, Any, Any]],
|
||||
primary_inputs: AutotuneInputFn,
|
||||
secondary_inputs: AutotuneInputFn | None = None,
|
||||
int64_threshold: Callable[..., bool] | None = None,
|
||||
warn_on_hash_miss: bool = False,
|
||||
):
|
||||
"""
|
||||
A decorator that automatically tunes and dispatches a Helion kernel based off of the kernel_key and provided inputs.
|
||||
|
||||
The general flow is this:
|
||||
1. We first run helion_kernel.autotune on all primary_inputs. This will give us a list of configs, one per primary_input.
|
||||
2. We benchmark every config on every primary_input and secondary_input.
|
||||
3. For every primary input/secondary input, we keep the config that is the fastest (NB: We aim to do some deduplication by reusing configs if they're within some threshold of the fastest config).
|
||||
4. We'll save the configs and the dispatch choices to a json file.
|
||||
5. We create a dispatch function that will lookup to see whether each kernel is one we've tuned for in step 2. If so, we'll use that config. Otherwise, we'll use some heuristic to (deterministically) find a reasonable config for the kernel.
|
||||
|
||||
There are 3 modes (set by env variable HELION_AOT_AUTOTUNE):
|
||||
- none: No autotuning will be done. We will skip to step 5. If the config json file doesn't exist, we'll raise an error.
|
||||
- retune: We only benchmark the existing useful configs on the primary/secondary inputs. We'll skip to step 2. This finishes much faster than create (albeit won't fully retune for each shape), but requires create to have been run first. For example, for rmsnorm, retune takes maybe one minute for 10 shapes, but create might take 30 minutes.
|
||||
- create: The kernel will be fully autotuned, starting from step 1.
|
||||
|
||||
You can also minimize the kernels to be autotuned by setting the env variable HELION_AOT_AUTOTUNE_KERNEL to the name of the kernel. For example, if you want to only autotune rmsnorm, you can set HELION_AOT_AUTOTUNE_KERNEL="rms_norm_fwd".
|
||||
|
||||
kernel_key: Callable[..., tuple[Any, Any, Any]]
|
||||
returns:
|
||||
(numeric_key, hash_key, exact_key)
|
||||
The semantics are that if all 3 match a saved key, we'll use that config.
|
||||
Otherwise, we'll use some heuristic to find a reasonable config for the
|
||||
kernel. The heuristic is:
|
||||
- We require exact_key to match the saved config.
|
||||
- We prioritize hash_key that match the saved config.
|
||||
- We then prioritize the highest numeric_key that's <= the current numeric_key.
|
||||
"""
|
||||
# Threshold for how much faster a config has to be for some shape to be considered "useful"
|
||||
threshold = 1.01
|
||||
# How many ms to run the kernel when retuning
|
||||
retune_rep_ms = 1000
|
||||
|
||||
def inner_autotune(kernel: helion.Kernel):
|
||||
helion_dir = Path(__file__).parent / "configs"
|
||||
|
||||
cc_major, cc_minor = get_cuda_device_capability()
|
||||
if cc_minor == 3:
|
||||
cc_minor = 0
|
||||
gpu_arch = f"sm_{cc_major}{cc_minor}"
|
||||
|
||||
path = helion_dir / Path(f"{config_path}_{gpu_arch}.json")
|
||||
|
||||
@functools.wraps(kernel_key)
|
||||
def wrapped_kernel_key(*inps: Any) -> tuple[Any, Any, Any]:
|
||||
"""
|
||||
A wrapper that handles dtype specially (since dtype is not
|
||||
serializable to json).
|
||||
"""
|
||||
|
||||
numeric_key, hash_key, exact_key = kernel_key(*inps)
|
||||
assert isinstance(hash_key, tuple)
|
||||
assert isinstance(exact_key, tuple)
|
||||
hash_key = tuple(
|
||||
i.itemsize if isinstance(i, torch.dtype) else i for i in hash_key
|
||||
)
|
||||
if exact_key is not None:
|
||||
exact_key = tuple(
|
||||
i.itemsize if isinstance(i, torch.dtype) else i for i in exact_key
|
||||
)
|
||||
return numeric_key, hash_key, exact_key
|
||||
|
||||
autotune_mode_str = os.environ.get("HELION_AOT_AUTOTUNE", "none")
|
||||
autotune_mode = AOTAutotuneMode.from_str(autotune_mode_str)
|
||||
autotune_kernel = os.environ.get("HELION_AOT_AUTOTUNE_KERNEL", "all")
|
||||
if autotune_kernel != "all" and autotune_kernel != config_path:
|
||||
autotune_mode = AOTAutotuneMode.NONE
|
||||
|
||||
@functools.cache
|
||||
def get_configs_from_autotuning(autotune_mode: AOTAutotuneMode):
|
||||
if autotune_mode == AOTAutotuneMode.NONE:
|
||||
if not path.exists():
|
||||
raise RuntimeError(
|
||||
f"Helion kernel not tuned yet. Run with HELION_AOT_AUTOTUNE=create HELION_AOT_AUTOTUNE_KERNEL={config_path} to generate the config at {path}"
|
||||
)
|
||||
return load_autotune_data_from_json(path)
|
||||
|
||||
inputs = sorted(
|
||||
list(primary_inputs()), key=lambda x: wrapped_kernel_key(*x)[0]
|
||||
)
|
||||
if autotune_mode == AOTAutotuneMode.CREATE:
|
||||
useful_configs = []
|
||||
for idx, input in enumerate(inputs):
|
||||
print(
|
||||
f"Autotuning for {config_path} with key: ",
|
||||
wrapped_kernel_key(*input),
|
||||
)
|
||||
config = kernel.autotune(input)
|
||||
useful_configs.append(repr(config))
|
||||
elif autotune_mode == AOTAutotuneMode.RETUNE:
|
||||
with open(path, "r") as f:
|
||||
json_obj = json.load(f)
|
||||
useful_configs = list(json_obj["useful_configs"].values())
|
||||
else:
|
||||
raise RuntimeError(f"Unexpected autotune mode: {autotune_mode}")
|
||||
|
||||
logger.info("Candidate useful configs: ")
|
||||
for idx, config in enumerate(useful_configs):
|
||||
logger.info(f"{idx}:, Config: {config}")
|
||||
|
||||
input_timings = []
|
||||
if secondary_inputs is not None:
|
||||
inputs += list(secondary_inputs())
|
||||
inputs = sorted(inputs, key=lambda x: kernel_key(*x)[0])
|
||||
|
||||
for input in inputs:
|
||||
cur_input_key = wrapped_kernel_key(*input)
|
||||
timings = []
|
||||
for idx, config in enumerate(useful_configs):
|
||||
try:
|
||||
cur_kernel = bind_and_compile_kernel(
|
||||
kernel, input, eval(config)
|
||||
)
|
||||
timings.append(
|
||||
do_bench(lambda: cur_kernel(*input), rep=retune_rep_ms)
|
||||
) # noqa: B023
|
||||
except Exception as e:
|
||||
logger.info(f"Error compiling config {config}: {e}")
|
||||
timings.append(float("inf"))
|
||||
input_timings.append((cur_input_key, timings))
|
||||
for idx, (key, timing) in enumerate(input_timings):
|
||||
logger.info(
|
||||
f"Key {key} timings: {' '.join([f'{i:.5f}' for i in timing])}"
|
||||
)
|
||||
|
||||
hash_configs_timings: dict[tuple[Any, Any], tuple[float, int | None]] = (
|
||||
defaultdict(lambda: (float("inf"), None))
|
||||
)
|
||||
for cur_kernel_key, input_timings in input_timings:
|
||||
for config_idx, (config, timing) in enumerate(
|
||||
zip(useful_configs, input_timings)
|
||||
):
|
||||
if timing < hash_configs_timings[cur_kernel_key][0] * threshold:
|
||||
hash_configs_timings[cur_kernel_key] = (timing, config_idx)
|
||||
|
||||
kept_configs = {}
|
||||
hash_configs = {k: v[1] for k, v in hash_configs_timings.items()}
|
||||
for key, config_idx in hash_configs.items():
|
||||
assert config_idx is not None
|
||||
kept_configs[config_idx] = useful_configs[config_idx]
|
||||
|
||||
json_obj = {
|
||||
"useful_configs": kept_configs,
|
||||
"hash_configs": {repr(k): v for k, v in hash_configs.items()},
|
||||
}
|
||||
|
||||
with open(path, "w") as f:
|
||||
json.dump(json_obj, f, indent=2)
|
||||
f.write("\n")
|
||||
return load_autotune_data_from_json(path)
|
||||
|
||||
cached_kernels = {}
|
||||
|
||||
def default_int64_threshold(*args: Any) -> bool:
|
||||
return any(
|
||||
[i.numel() >= 2**31 for i in args if isinstance(i, torch.Tensor)]
|
||||
)
|
||||
|
||||
used_int64_threshold = (
|
||||
int64_threshold if int64_threshold is not None else default_int64_threshold
|
||||
)
|
||||
|
||||
def wrapped_func(*args: Any):
|
||||
nonlocal cached_kernels
|
||||
cur_kernel_key = wrapped_kernel_key(*args)
|
||||
has_int64 = used_int64_threshold(*args)
|
||||
size1 = tuple(
|
||||
tuple(shape == 1 for shape in arg.shape)
|
||||
for arg in args
|
||||
if isinstance(arg, torch.Tensor)
|
||||
)
|
||||
dtypes = tuple(arg.dtype for arg in args if isinstance(arg, torch.Tensor))
|
||||
scalar_args = tuple(a for a in args if not isinstance(a, torch.Tensor))
|
||||
key = (cur_kernel_key, has_int64, size1, dtypes, scalar_args)
|
||||
if key in cached_kernels:
|
||||
out = cached_kernels[key](*args)
|
||||
return out
|
||||
if has_int64:
|
||||
kernel.settings = dataclasses.replace(
|
||||
kernel.settings, index_dtype=torch.int64
|
||||
)
|
||||
useful_configs, hash_configs = get_configs_from_autotuning(autotune_mode)
|
||||
key_to_config = {k: useful_configs[v] for k, v in hash_configs.items()}
|
||||
if key not in cached_kernels and cur_kernel_key in hash_configs:
|
||||
cached_kernels[key] = bind_and_compile_kernel(
|
||||
kernel, args, key_to_config[cur_kernel_key]
|
||||
)
|
||||
else:
|
||||
if warn_on_hash_miss:
|
||||
logger.warning(
|
||||
f"No config found for key {cur_kernel_key} for kernel={config_path}. Finding best match. This is *not* a correctness issue, but means that the performance of this kernel could potentially be improved."
|
||||
)
|
||||
used_config = None
|
||||
|
||||
def config_key_sort(
|
||||
config_key: tuple[Any, Any, Any],
|
||||
) -> tuple[Any, ...]:
|
||||
return (
|
||||
config_key[2] == cur_kernel_key[2],
|
||||
config_key[1] == cur_kernel_key[1],
|
||||
config_key[0] <= cur_kernel_key[0],
|
||||
config_key[0],
|
||||
)
|
||||
|
||||
sorted_keys = sorted(
|
||||
hash_configs.keys(), key=config_key_sort, reverse=True
|
||||
)
|
||||
used_config = key_to_config[sorted_keys[0]]
|
||||
if len(used_config) == 3:
|
||||
assert (
|
||||
used_config[2] == cur_kernel_key[2]
|
||||
), "Exact key not found in configs"
|
||||
cached_kernels[key] = bind_and_compile_kernel(kernel, args, used_config)
|
||||
out = cached_kernels[key](*args)
|
||||
return out
|
||||
|
||||
return wrapped_func
|
||||
|
||||
return inner_autotune
|
||||
@@ -25,7 +25,7 @@ from sglang.kernels.ops.moe.inkling_moe import (
|
||||
post_reorder,
|
||||
pre_reorder,
|
||||
select_grouped_gemm_block_m,
|
||||
silu_and_mul_helion,
|
||||
silu_and_mul,
|
||||
)
|
||||
from sglang.kernels.ops.moe.sigmoid_gate_topk_renorm import (
|
||||
sigmoid_gate_topk_renorm,
|
||||
@@ -572,7 +572,7 @@ def activation(
|
||||
*gateup_output.shape[:-1], gateup_output.shape[-1] // 2, dtype=out_dtype
|
||||
)
|
||||
|
||||
return silu_and_mul_helion(
|
||||
return silu_and_mul(
|
||||
gateup_output, topk_weights, out_dtype, use_interleaved=use_interleaved
|
||||
)
|
||||
raise ValueError(f"Unsupported activation: {activation_type}")
|
||||
|
||||
Reference in New Issue
Block a user