"""Helion-dependency-free runtime launch helpers for the Triton backend.
This module holds the small set of runtime symbols that Helion's *generated*
Triton code depends on at execution time:
* :func:`default_launcher` -- invokes a compiled ``triton.jit`` kernel.
* :func:`get_num_sm` -- persistent-kernel grid size (host statement).
* :func:`set_triton_allocator` -- installs the scratch allocator used by TMA /
tensor-descriptor kernels (device-function prefix statement).
It depends only on ``torch`` and ``triton`` -- no other ``helion`` module -- so
the ahead-of-time precompiler can bulk-export this file verbatim into a
standalone kernel with zero Helion runtime dependency.
Helion-specific behavior that is only meaningful in-process (translating
Triton's opaque shape errors into :class:`helion.exc.ShapeMismatch`, and the
CPU/TPU cases of :func:`get_num_sm`) lives in thin wrappers in
:mod:`helion.runtime`, not here.
"""
from __future__ import annotations
import contextvars
import math
import weakref
import torch
try:
import triton
except ImportError:
triton = None # type: ignore[assignment]
# (Triton driver, device of its scratch allocations): querying the driver's
# target reads device properties, too slow to repeat on every kernel launch.
# Top level, so the standalone runtime shim rewrites its ``global`` to
# ``nonlocal``.
_alloc_device: tuple[object, torch.device] | None = None
if triton is not None:
def _alloc_fn(size: int, alignment: int, stream: int | None) -> torch.Tensor:
global _alloc_device
driver = triton.runtime.driver.active
cached = _alloc_device
if cached is None or cached[0] is not driver:
# Dynamically get device from Triton backend
current_target = driver.get_current_target()
if current_target is None:
raise RuntimeError("No active Triton target available")
# No device index: allocate on the current device, where Triton
# launches.
cached = _alloc_device = (driver, torch.device(current_target.backend))
return torch.empty(size, device=cached[1], dtype=torch.int8)
def set_triton_allocator() -> None:
try:
from triton import set_allocator
from triton.runtime._allocation import NullAllocator
from triton.runtime._allocation import _allocator
except ImportError:
return
if isinstance(_allocator, contextvars.ContextVar):
existing = _allocator.get()
else: # older versions of Triton
existing = _allocator
# if allocator isn't NullAllocator, we assume it is set by the user
if isinstance(existing, NullAllocator):
set_allocator(_alloc_fn)
else:
[docs]
def set_triton_allocator() -> None: # type: ignore[misc]
pass
def get_num_sm(device: torch.device, *, reserved_sms: int = 0) -> int:
"""
Get the number of streaming multiprocessors (SMs) for the specified GPU.
Args:
device: Device to query. Must be a GPU device (``cuda``/``xpu``/``mps``/
``mtia``); CPU/TPU handling lives in :func:`helion.runtime.get_num_sm`.
reserved_sms: Number of SMs to keep free for other work (e.g., communication
kernels). Defaults to 0 meaning all device SMs are available to Helion.
Returns:
Grid size to use for a persistent kernel on the device after accounting
for any reserved SMs. Always at least 1.
"""
available_sms: int
assert device.type in [
"cuda",
"xpu",
"mtia",
"mps",
], "TODO: implement for other devices"
if device.type == "cuda":
available_sms = torch.cuda.get_device_properties(
device.index
).multi_processor_count
# TODO(EikanWang): gpu_subslice_count is an out-of-date term. we change update it to XeCore number.
elif device.type == "xpu":
available_sms = torch.xpu.get_device_properties(device.index).gpu_subslice_count
elif device.type == "mps":
available_sms = torch.backends.mps.get_core_count()
elif device.type == "mtia":
device_props = torch.mtia.get_device_properties(device.index)
if "max_grid_height" in device_props and "max_grid_width" in device_props:
available_sms = (
device_props["max_grid_height"] * device_props["max_grid_width"]
)
else:
raise RuntimeError(
f"Unable to determine SM count for MTIA device. "
f"Available properties: {list(device_props.keys())}"
)
else:
raise NotImplementedError(
f"get_num_sm not implemented for device type: {device.type}"
)
if reserved_sms <= 0:
return available_sms
return max(available_sms - reserved_sms, 1)
# CUs per XCD by base CDNA architecture. Used to derive the live,
# partition-visible XCD count from the observed CU count (see get_num_xcd).
_CUS_PER_XCD: dict[str, int] = {
"gfx942": 38, # CDNA3 (MI300)
"gfx950": 32, # CDNA4 (MI350)
"gfx951": 32, # CDNA4 (MI355)
}
def get_num_xcd(device: torch.device | int | None = None) -> int:
"""Number of XCDs visible for ``device`` on AMD CDNA, else ``1``.
Derived from the live, partition-visible compute-unit count rather than the
architecture name, so MI300A (6 XCDs) and compute-partition modes such as CPX
(which expose a single XCD) are handled correctly. Returns ``1`` -- which
disables xcd_remap -- for unknown architectures or a CU count that does not
look like an integer number of XCDs.
"""
if not torch.cuda.is_available():
return 1
try:
props = torch.cuda.get_device_properties(
device if device is not None else torch.cuda.current_device()
)
except Exception:
return 1
arch = getattr(props, "gcnArchName", None)
if not arch:
return 1
cus_per_xcd = _CUS_PER_XCD.get(arch.split(":")[0])
if cus_per_xcd is None:
return 1
cu_count = props.multi_processor_count
num_xcd = round(cu_count / cus_per_xcd)
# Tolerate harvested parts, but bail out (return 1) if the live CU count does
# not look like an integer number of XCDs.
if num_xcd < 1 or abs(num_xcd * cus_per_xcd - cu_count) > cus_per_xcd // 4:
return 1
return num_xcd
def default_launcher(
triton_kernel: object,
grid: tuple[int, ...],
*args: object,
num_warps: int,
num_stages: int,
_remote_copy_signal_dst: torch.Tensor | None = None,
_remote_copy_signal_slots_per_program: int = 0,
_remote_copy_process_group_name: str | None = None,
_remote_barrier_signal_slots_per_program: int = 0,
_remote_barrier_process_group_name: str | None = None,
_remote_copy_scratch_specs: tuple[tuple[torch.Tensor, int], ...] = (),
_persistent_state_specs: tuple[
tuple[torch.Tensor, int, torch.dtype, bool], ...
] = (),
_persistent_state_process_group_name: str | None = None,
_persistent_state_rank_digest: str | None = None,
_minimum_resident_programs: int = 0,
ptx_options: str | None = None,
launch_cooperative_grid: bool = False,
**kwargs: dict,
) -> object:
"""Default launcher function that executes the kernel immediately."""
if _remote_copy_signal_slots_per_program:
if _remote_copy_signal_dst is None or _remote_copy_process_group_name is None:
raise RuntimeError(
"remote-copy completion storage requires a symmetric destination "
"and process group"
)
signal = _get_remote_copy_signal(
triton_kernel,
_remote_copy_signal_dst,
_remote_copy_process_group_name,
math.prod(grid) * _remote_copy_signal_slots_per_program,
)
# Allocation zeroes new pads and receive waits reset consumed slots.
# Clearing here could erase a completion sent before this rank launches.
args = (*args, signal)
if _remote_barrier_signal_slots_per_program:
if _remote_barrier_process_group_name is None:
raise RuntimeError(
"remote-barrier completion storage requires a process group"
)
signal = _get_remote_barrier_signal(
triton_kernel,
_remote_barrier_process_group_name,
math.prod(grid) * _remote_barrier_signal_slots_per_program,
)
args = (*args, signal)
for slot, (scratch_like, numel_per_program) in enumerate(
_remote_copy_scratch_specs
):
scratch = _get_remote_copy_scratch(
triton_kernel,
scratch_like,
slot,
math.prod(grid) * numel_per_program,
)
args = (*args, scratch)
if _persistent_state_specs:
persistent_state_namespace = (
tuple(grid),
num_warps,
num_stages,
ptx_options,
launch_cooperative_grid,
tuple(sorted((name, repr(value)) for name, value in kwargs.items())),
tuple(spec[1:] for spec in _persistent_state_specs),
# Symmetric state exchanges the namespace, so ranks compare digests.
_persistent_state_rank_digest,
)
for slot, (state_like, numel, dtype, symmetric) in enumerate(
_persistent_state_specs
):
if _persistent_state_process_group_name is None:
state_args = (
_get_persistent_state(
triton_kernel,
state_like,
persistent_state_namespace,
slot,
numel,
dtype,
),
)
else:
state_args = _get_process_group_state(
triton_kernel,
state_like,
persistent_state_namespace,
slot,
numel,
dtype,
symmetric=symmetric,
process_group_name=_persistent_state_process_group_name,
)
args = (*args, *state_args)
# For both CUDA and MTIA, use the same kernel execution.
run_kwargs: dict = {
"grid": grid,
"warmup": False,
"num_warps": num_warps,
"num_stages": num_stages,
"launch_cooperative_grid": launch_cooperative_grid,
**kwargs,
}
if ptx_options is not None:
run_kwargs["ptx_options"] = ptx_options
if _minimum_resident_programs:
# ``triton_kernel`` is a JITFunction. Resource information belongs to
# its exact compiled specialization, so compile (but do not launch)
# that specialization before asking CUDA for its occupancy.
compiled_kernel = triton_kernel.run( # type: ignore[union-attr]
*args,
**{**run_kwargs, "warmup": True},
)
_validate_resident_program_capacity(
compiled_kernel,
args,
num_warps=num_warps,
required_programs=_minimum_resident_programs,
)
return triton_kernel.run( # type: ignore[union-attr]
*args,
**run_kwargs,
)
def compile_only_launch_args(
*args: object,
_remote_copy_signal_slots_per_program: int = 0,
_remote_barrier_signal_slots_per_program: int = 0,
_remote_copy_scratch_specs: tuple[tuple[torch.Tensor, int], ...] = (),
_persistent_state_specs: tuple[
tuple[torch.Tensor, int, torch.dtype, bool], ...
] = (),
**kwargs: object,
) -> tuple[tuple[object, ...], dict[str, object]]:
"""Return ``triton_kernel.run`` arguments for compiling without launching.
The arguments ``default_launcher`` appends become empty tensors of the same
dtypes, which is all compilation reads, and launcher-only options drop.
"""
dtypes = [torch.int64] * (
bool(_remote_copy_signal_slots_per_program)
+ bool(_remote_barrier_signal_slots_per_program)
)
dtypes += [like.dtype for like, _ in _remote_copy_scratch_specs]
for _, _, dtype, symmetric in _persistent_state_specs:
dtypes += [dtype, torch.int64] if symmetric else [dtype]
kwargs = {name: value for name, value in kwargs.items() if not name.startswith("_")}
if not dtypes:
return args, kwargs
device = next(arg.device for arg in args if isinstance(arg, torch.Tensor))
placeholders = [torch.empty(0, dtype=dtype, device=device) for dtype in dtypes]
return (*args, *placeholders), kwargs
def _get_remote_copy_signal(
triton_kernel: object,
dst: torch.Tensor,
process_group_name: str,
required_slots: int,
) -> torch.Tensor:
"""Return compiler-owned completion slots from ``dst``'s signal pad."""
import torch.distributed._symmetric_memory as symm_mem
cache = vars(triton_kernel).setdefault("_helion_remote_copy_signal_cache", {})
key = (id(dst), process_group_name)
entry = cache.get(key)
if entry is not None and entry[0]() is dst:
signal_pad = entry[1]
else:
handle = symm_mem.rendezvous(
dst,
group=process_group_name, # pyrefly: ignore[bad-argument-type]
)
signal_pad = handle.get_signal_pad(handle.rank, dtype=torch.int64)
def remove_from_cache(_ref: object) -> None:
cache.pop(key, None)
cache[key] = (weakref.ref(dst, remove_from_cache), signal_pad)
capacity = signal_pad.numel()
if required_slots > capacity:
raise RuntimeError(
"Helion remote copies require "
f"{required_slots} int64 completion slots, but the symmetric-memory "
f"signal pad has capacity {capacity}. Increase the signal pad size "
"before allocating symmetric tensors."
)
# Reserve from the end so Helion's slots do not overlap PyTorch's standard
# low-offset signal-pad protocols.
return signal_pad.narrow(0, capacity - required_slots, required_slots)
def _get_remote_barrier_signal(
triton_kernel: object,
process_group_name: str,
required_slots: int,
) -> torch.Tensor:
"""Return compiler-owned peer-barrier counters from a group workspace."""
import torch.distributed._symmetric_memory as symm_mem
device = torch.device("cuda", torch.cuda.current_device())
cache = vars(triton_kernel).setdefault("_helion_remote_barrier_signal_cache", {})
key = (device, process_group_name)
entry = cache.get(key)
if entry is None:
workspace = symm_mem.empty(1, dtype=torch.uint8, device=device)
handle = symm_mem.rendezvous(
workspace,
group=process_group_name, # pyrefly: ignore[bad-argument-type]
)
cache[key] = (workspace, handle)
else:
_, handle = entry
signal_pad = handle.get_signal_pad(handle.rank, dtype=torch.int64)
capacity = signal_pad.numel()
if required_slots > capacity:
raise RuntimeError(
"Helion remote barriers require "
f"{required_slots} int64 completion slots, but the symmetric-memory "
f"signal pad has capacity {capacity}. Increase the signal pad size "
"before launching the kernel."
)
return signal_pad.narrow(0, capacity - required_slots, required_slots)
def _get_remote_copy_scratch(
triton_kernel: object,
like: torch.Tensor,
slot: int,
required_numel: int,
) -> torch.Tensor:
"""Return stream-local global scratch for one computed DMA source."""
if like.device.type != "cuda":
raise RuntimeError("NVSHMEM remote-copy scratch requires a CUDA tensor")
stream = torch.cuda.current_stream(like.device)
cache = vars(triton_kernel).setdefault("_helion_remote_copy_scratch_cache", {})
key = (like.device, like.dtype, stream.cuda_stream, slot)
scratch = cache.get(key)
if scratch is None or scratch.numel() < required_numel:
scratch = torch.empty(
required_numel,
dtype=like.dtype,
device=like.device,
)
cache[key] = scratch
return scratch
def _get_persistent_state(
triton_kernel: object,
like: torch.Tensor,
namespace: tuple[object, ...],
slot: int,
required_numel: int,
dtype: torch.dtype,
) -> torch.Tensor:
"""Return stream-local compiler state retained across kernel launches."""
if like.device.type != "cuda":
raise RuntimeError("persistent Triton state requires a CUDA tensor")
stream = torch.cuda.current_stream(like.device)
cache = vars(triton_kernel).setdefault("_helion_persistent_state_cache", {})
key = (like.device, dtype, stream.cuda_stream, namespace, slot)
state = cache.get(key)
if state is None or state.numel() < required_numel:
state = torch.zeros(required_numel, dtype=dtype, device=like.device)
cache[key] = state
return state
def _get_process_group_state(
triton_kernel: object,
like: torch.Tensor,
namespace: tuple[object, ...],
slot: int,
required_numel: int,
dtype: torch.dtype,
*,
symmetric: bool,
process_group_name: str,
) -> tuple[object, ...]:
"""Return the kernel arguments of retained state shared by every stream.
One epoch then covers the dispatch ticket and peer counters, so launches of
one kernel must not overlap across streams. A symmetric state is followed by
its per-rank base pointer table.
"""
if like.device.type != "cuda":
raise RuntimeError("persistent Triton state requires a CUDA tensor")
cache = vars(triton_kernel).setdefault("_helion_persistent_state_cache", {})
# The numel is part of the namespace, so a cached entry always fits.
key = (like.device, dtype, process_group_name, namespace, slot)
entry = cache.get(key)
if entry is None:
if torch.cuda.is_current_stream_capturing():
raise RuntimeError(
"Helion allocates distributed launch state on first use; run the "
"kernel once before CUDA graph capture"
)
if symmetric:
entry = _new_symmetric_state(
like.device,
(
triton_kernel.__name__, # type: ignore[attr-defined]
namespace,
slot,
required_numel,
str(dtype),
),
required_numel,
dtype,
process_group_name,
)
else:
entry = (torch.zeros(required_numel, dtype=dtype, device=like.device),)
cache[key] = entry
return entry[:2]
def _new_symmetric_state(
device: torch.device,
fingerprint: tuple[object, ...],
numel: int,
dtype: torch.dtype,
process_group_name: str,
) -> tuple[torch.Tensor, torch.Tensor, object]:
"""Allocate zeroed symmetric state once every rank agrees on the launch.
Returns the state, its per-rank base pointer table followed by this rank,
and the owning handle.
"""
import torch.distributed as dist
import torch.distributed._symmetric_memory as symm_mem
from torch.distributed.distributed_c10d import _resolve_process_group
group = _resolve_process_group(
process_group_name # pyrefly: ignore[bad-argument-type]
)
with torch.cuda.device(device):
fingerprint = (*fingerprint, torch.cuda.get_device_capability(device))
fingerprints: list[object] = [None] * dist.get_world_size(group)
dist.all_gather_object(fingerprints, fingerprint, group=group)
if any(other != fingerprint for other in fingerprints):
raise RuntimeError(
"Helion distributed kernels require the same launch on every "
f"rank; got {fingerprints!r}"
)
state = symm_mem.empty(numel, dtype=dtype, device=device)
state.zero_()
# Zero before rendezvous, since peers may publish right after it.
torch.cuda.current_stream(device).synchronize()
handle = symm_mem.rendezvous(
state,
group=process_group_name, # pyrefly: ignore[bad-argument-type]
)
ptrs = torch.tensor(
[*handle.buffer_ptrs, handle.rank], dtype=torch.int64, device=device
)
return state, ptrs, handle
def _validate_resident_program_capacity(
compiled_kernel: object,
args: tuple[object, ...],
*,
num_warps: int,
required_programs: int,
) -> None:
"""Reject a polling schedule whose required CTA cohort cannot be resident."""
import importlib
tensor = next((arg for arg in args if isinstance(arg, torch.Tensor)), None)
if tensor is None or tensor.device.type != "cuda":
raise RuntimeError("cross-loop residency checks require a CUDA tensor")
if compiled_kernel is None:
raise RuntimeError("unable to compile cross-loop scheduled kernel")
# Accessing ``run`` initializes Triton's module/function handles without
# launching the kernel. Cache the exact driver result on the compiled
# specialization because this wrapper is also called during graph capture.
_run = compiled_kernel.run # type: ignore[attr-defined]
function = getattr(compiled_kernel, "function", None)
metadata = getattr(compiled_kernel, "metadata", None)
shared = getattr(metadata, "shared", None)
if function is None or not isinstance(shared, int):
raise RuntimeError("unable to query cross-loop kernel occupancy")
device = tensor.device
cache = vars(compiled_kernel).setdefault(
"_helion_resident_program_capacity_cache", {}
)
key = (device, num_warps, shared)
capacity = cache.get(key)
if capacity is None:
cuda_driver = importlib.import_module("cuda.bindings.driver")
with torch.cuda.device(device):
error, blocks_per_sm = (
cuda_driver.cuOccupancyMaxActiveBlocksPerMultiprocessor(
cuda_driver.CUfunction(int(function)),
num_warps * 32,
shared,
)
)
if error != cuda_driver.CUresult.CUDA_SUCCESS:
raise RuntimeError(
f"CUDA occupancy query failed for cross-loop kernel: {error}"
)
properties = torch.cuda.get_device_properties(device)
capacity = int(blocks_per_sm) * int(properties.multi_processor_count)
cache[key] = capacity
if required_programs > capacity:
raise RuntimeError(
"Cross-loop scheduling requires "
f"{required_programs} concurrently resident programs, but this "
f"kernel/device can residently execute only {capacity}. Choose a "
"lower-resource configuration, a smaller ready prefix, or a "
"root barrier."
)
