# Copyright (c) 2026 LightSeek Foundation # # Permission is hereby granted, free of charge, to any person obtaining a copy # of this software and associated documentation files (the "Software"), to deal # in the Software without restriction, including without limitation the rights # to use, copy, modify, merge, publish, distribute, sublicense, and/or sell # copies of the Software, and to permit persons to whom the Software is # furnished to do so, subject to the following conditions: # # The above copyright notice and this permission notice shall be included in # all copies or substantial portions of the Software. # # THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR # IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, # FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE # AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER # LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, # OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE # SOFTWARE. import importlib import logging import pkgutil from typing import List, Tuple import torch import torch.distributed as dist from tokenspeed_kernel._triton import redirect_triton_to_tokenspeed_triton, tl, triton # iris does plain ``import triton`` at module load time; route those bindings # to the vendored ``tokenspeed_triton`` so iris and tokenspeed-kernel share a # single triton distribution. See # :func:`redirect_triton_to_tokenspeed_triton` for details. with redirect_triton_to_tokenspeed_triton(): import iris # noqa: E402 # Pre-import every iris kernel module that does ``import triton`` at module # load time (the CCL APIs above lazy-import them at call time, when the # redirect is no longer active). import iris.ccl.triton # noqa: E402 from iris.ccl import Config as _IrisConfig # noqa: E402 from iris.ccl.all_gather import all_gather as _iris_all_gather # noqa: E402 from iris.ccl.all_reduce import all_reduce as _iris_all_reduce # noqa: E402 from iris.ccl.reduce_scatter import ( # noqa: E402 reduce_scatter as _iris_reduce_scatter, ) for _info in pkgutil.walk_packages( iris.ccl.triton.__path__, prefix="iris.ccl.triton." ): importlib.import_module(_info.name) from tokenspeed_kernel.platform import current_platform # noqa: E402 logger = logging.getLogger(__file__) _platform = current_platform() __all__ = [ "IrisAllReduce", "IrisRSAG", "IrisAllReduceResidualRMSNorm", "create_iris_state", "iris_all_reduce", "create_iris_rsag_state", "create_iris_ar_rmsnorm_state", "iris_allreduce_residual_rmsnorm", "IRIS_AR_RMSNORM_STATES", ] IRIS_AR_RMSNORM_STATES: dict = {} def _get_available_gpu_memory(gpu_id: int, empty_cache: bool = True) -> float: if torch.cuda.is_available(): with torch.cuda.device(gpu_id): if empty_cache: torch.cuda.empty_cache() free_gpu_memory, _ = torch.cuda.mem_get_info() return free_gpu_memory / (1 << 30) return 0.0 _iris_ctx_singleton = None def _get_or_create_iris_context(heap_size: int): global _iris_ctx_singleton if _iris_ctx_singleton is None: _iris_ctx_singleton = iris.iris(heap_size=heap_size) return _iris_ctx_singleton class IrisRSAG(object): def __init__( self, group: dist.ProcessGroup, rank_in_group: int, max_tokens: int, hidden_size: int, device: torch.device = None, heap_size: int | None = None, ) -> None: assert ( type(group) == dist.ProcessGroup ), f"Expected dist.ProcessGroup, got {type(group)}" assert dist.is_initialized(), ( "torch.distributed must be initialized before constructing " "IrisRSAG; call dist.init_process_group() first." ) assert _platform.is_amd, ( "IrisRSAG currently targets AMD ROCm; " f"got non-AMD platform: {_platform}" ) assert ( group == dist.group.WORLD or group.size() == dist.get_world_size() ), "iris.ccl all_gather/reduce_scatter do not accept a sub-group." self.group = group self.rank_in_group = rank_in_group self.device = device or torch.device(f"cuda:{torch.cuda.current_device()}") self.max_tokens = max_tokens self.hidden_size = hidden_size self.dtype = torch.bfloat16 self.world_size = group.size() # Heap holds in/out flat buffers plus iris bookkeeping; over-provision # similarly to ``IrisAllReduce`` to leave room for ring/spinlock flags. if heap_size is None: buf_bytes = max_tokens * hidden_size * self.dtype.itemsize heap_size = max(1 << 28, 4 * buf_bytes + (16 << 20)) free_gpu_memory_begin = _get_available_gpu_memory(torch.cuda.current_device()) self._ctx = _get_or_create_iris_context(heap_size) self._in_buff = self._ctx.empty((max_tokens, hidden_size), dtype=self.dtype) self._out_buff = self._ctx.empty((max_tokens, hidden_size), dtype=self.dtype) free_gpu_memory_after = _get_available_gpu_memory(torch.cuda.current_device()) logger.info( "Iris RSAG symmetric-heap buffers allocated: %s GB", free_gpu_memory_begin - free_gpu_memory_after, ) assert self._ctx.get_num_ranks() == dist.get_world_size(), ( f"Iris world size {self._ctx.get_num_ranks()} " f"!= torch world size {dist.get_world_size()}" ) assert self.rank_in_group == self._ctx.get_rank(), ( f"rank mismatch: rank_in_group={self.rank_in_group}, " f"iris rank={self._ctx.get_rank()}" ) # -- token-distribution helpers (mirror sibling classes) ---------------- def get_token_dist(self, total_tokens_in_group: int) -> list: token_list_in_group = [] for rank in range(self.world_size): num_tokens_per_rank = total_tokens_in_group // self.world_size + ( 1 if (rank < total_tokens_in_group % self.world_size) else 0 ) token_list_in_group.append(num_tokens_per_rank) return token_list_in_group def get_context(self, token_list_in_group: list) -> Tuple[int, int, int]: total_num_tokens = sum(token_list_in_group) assert ( total_num_tokens <= self.max_tokens ), f"The inner comm buffer is too small: {total_num_tokens=} is not <= {self.max_tokens=}" local_num_tokens = token_list_in_group[self.rank_in_group] local_token_offset = sum(token_list_in_group[: self.rank_in_group]) return total_num_tokens, local_num_tokens, local_token_offset # -- internal helpers --------------------------------------------------- def _assert_uniform(self, token_list_in_group: List[int]) -> int: first = token_list_in_group[0] assert all(t == first for t in token_list_in_group), ( "IrisRSAG requires uniform tokens per rank; got " f"token_list_in_group={token_list_in_group}" ) return first @staticmethod def _pick_block_n(hidden_size: int) -> int: # Pick the largest power-of-two block that divides hidden_size, capped # at 256. This keeps the iris kernel on its no-mask fast path and # still produces enough tiles (world_size * hidden/block_n) to fill # ``comm_sms`` SMs on MI300-class chips. for cand in (256, 128, 64, 32, 16): if hidden_size % cand == 0: return cand return hidden_size def _make_config(self, local_num_tokens: int, hidden_size: int): # ``swizzle_size=1`` keeps tile_id ordering row-major in M, which is # required so that block-distribution (DISTRIBUTION=1) hands rank r # exactly the K tiles spanning rows [r*local, (r+1)*local) in the # reduce-scatter kernel. ``all_gather`` is rank-agnostic on tile order # so the same config is fine. return _IrisConfig( block_size_m=local_num_tokens, block_size_n=self._pick_block_n(hidden_size), swizzle_size=1, all_reduce_distribution=1, ) # -- public collective ops --------------------------------------------- def reduce_scatter( self, hidden_states: torch.Tensor, tp_num_tokens: int = None, token_list_in_group: List[int] = None, safe=True, ) -> torch.Tensor: assert ( tp_num_tokens is not None or token_list_in_group is not None ), "Either tp_num_tokens or token_list_in_group must be provided" if token_list_in_group is None: token_list_in_group = self.get_token_dist(tp_num_tokens) assert ( hidden_states.dtype == self.dtype ), f"Only {self.dtype} is supported, got {hidden_states.dtype}" local_num_tokens = self._assert_uniform(token_list_in_group) total_num_tokens, _, local_token_offset = self.get_context(token_list_in_group) assert (hidden_states.shape[0] == total_num_tokens) and ( hidden_states.shape[-1] == self.hidden_size ), ( f"Mismatched shape, {hidden_states.shape[0]=} != {total_num_tokens=} " f"or {hidden_states.shape[-1]=} != {self.hidden_size=} " f"{hidden_states.shape=}" ) if local_num_tokens == 0: return torch.empty( (0, self.hidden_size), dtype=hidden_states.dtype, device=hidden_states.device, ) in_view = self._in_buff[:total_num_tokens, : self.hidden_size] out_view = self._out_buff[:total_num_tokens, : self.hidden_size] in_view.copy_(hidden_states) self._ctx.device_barrier() config = self._make_config(local_num_tokens, self.hidden_size) _iris_reduce_scatter(out_view, in_view, self._ctx, config=config) output = out_view[local_token_offset : local_token_offset + local_num_tokens, :] return output.clone() if safe else output def all_gather( self, hidden_states: torch.Tensor, tp_num_tokens: int = None, token_list_in_group: List[int] = None, safe=True, ) -> torch.Tensor: assert ( tp_num_tokens is not None or token_list_in_group is not None ), "Either tp_num_tokens or token_list_in_group must be provided" if token_list_in_group is None: token_list_in_group = self.get_token_dist(tp_num_tokens) assert ( hidden_states.dtype == self.dtype ), f"Only {self.dtype} is supported, got {hidden_states.dtype}" local_num_tokens = self._assert_uniform(token_list_in_group) total_num_tokens, _, _ = self.get_context(token_list_in_group) hidden_size = hidden_states.shape[-1] assert (hidden_states.shape[0] == local_num_tokens) and ( hidden_size <= self.hidden_size ), ( f"{hidden_states.shape=}|{local_num_tokens=}|{hidden_states.device=} " "Mismatched shape" ) if local_num_tokens == 0: return torch.empty( (0, hidden_size), dtype=hidden_states.dtype, device=hidden_states.device, ) in_view = self._in_buff[:local_num_tokens, :hidden_size] out_view = self._out_buff[:total_num_tokens, :hidden_size] in_view.copy_(hidden_states) self._ctx.device_barrier() config = self._make_config(local_num_tokens, hidden_size) _iris_all_gather(out_view, in_view, self._ctx, config=config) return out_view.clone() if safe else out_view class IrisAllReduce(object): def __init__( self, group: dist.ProcessGroup, rank_in_group: int, max_numel: int, dtype: torch.dtype = torch.bfloat16, heap_size: int | None = None, device: torch.device = None, config=None, ) -> None: assert ( type(group) == dist.ProcessGroup ), f"Expected dist.ProcessGroup, got {type(group)}" assert dist.is_initialized(), ( "torch.distributed must be initialized before constructing " "IrisAllReduce; call dist.init_process_group() first." ) assert _platform.is_amd, ( "IrisAllReduce currently targets AMD ROCm; " f"got non-AMD platform: {_platform}" ) self.group = group self.rank_in_group = rank_in_group self.max_numel = max_numel self.dtype = dtype self.device = device or torch.device(f"cuda:{torch.cuda.current_device()}") self._config = config or _IrisConfig( block_size_m=32, block_size_n=64, all_reduce_distribution=1 ) # Heap holds two flat buffers of ``max_numel * itemsize`` plus iris # bookkeeping; we leave generous headroom (~16 MiB) for internal # workspaces such as ring/spinlock flags. if heap_size is None: buf_bytes = max_numel * dtype.itemsize heap_size = max(1 << 28, 4 * buf_bytes + (16 << 20)) free_gpu_memory_begin = _get_available_gpu_memory(torch.cuda.current_device()) self._ctx = _get_or_create_iris_context(heap_size) self._input_buf = self._ctx.zeros((max_numel,), dtype=dtype) self._output_buf = self._ctx.zeros((max_numel,), dtype=dtype) free_gpu_memory_after = _get_available_gpu_memory(torch.cuda.current_device()) logger.info( "Iris all-reduce symmetric-heap buffers allocated: %s GB", free_gpu_memory_begin - free_gpu_memory_after, ) self.world_size = group.size() def all_reduce( self, tensor: torch.Tensor, op=None, safe: bool = True, async_op: bool = False, ) -> torch.Tensor: assert tensor.dtype == self.dtype, ( f"Iris all-reduce dtype mismatch: tensor={tensor.dtype}, " f"backend={self.dtype}" ) numel = tensor.numel() assert numel <= self.max_numel, ( f"tensor numel ({numel}) exceeds iris buffer capacity " f"({self.max_numel})" ) if tensor.dim() >= 2: n_dim = tensor.shape[-1] m_dim = numel // n_dim else: m_dim, n_dim = 1, numel in_view = self._input_buf.narrow(0, 0, numel).view(m_dim, n_dim) out_view = self._output_buf.narrow(0, 0, numel).view(m_dim, n_dim) in_view.view(-1).copy_(tensor.view(-1)) self._ctx.device_barrier() ar_group = None if self.group == dist.group.WORLD else self.group _iris_all_reduce( out_view, in_view, self._ctx, op=op, group=ar_group, async_op=async_op, config=self._config, ) result = out_view.view(tensor.shape) return result.clone() if safe else result @triton.jit def iris_allreduce_residual_rmsnorm_kernel( input_sym_ptr, # base of symmetric (M, HIDDEN_SIZE) input buffer residual_ptr, # local (M, HIDDEN_SIZE) weight_ptr, # local (HIDDEN_SIZE,) norm_out_ptr, # local (M, HIDDEN_SIZE) residual_out_ptr, # local (M, HIDDEN_SIZE) M, heap_bases, iris_rank: tl.constexpr, world_size: tl.constexpr, rank_start: tl.constexpr, rank_stride: tl.constexpr, HIDDEN_SIZE: tl.constexpr, BLOCK_SIZE: tl.constexpr, EPS: tl.constexpr, ): row = tl.program_id(0) if row >= M: return offsets = tl.arange(0, BLOCK_SIZE) mask = offsets < HIDDEN_SIZE row_offsets = row * HIDDEN_SIZE + offsets in_row_ptr = input_sym_ptr + row_offsets acc = tl.zeros((BLOCK_SIZE,), dtype=tl.float32) for i in tl.static_range(0, world_size): remote_rank = rank_start + i * rank_stride acc += iris.load( in_row_ptr, iris_rank, remote_rank, heap_bases, mask=mask, other=0.0, ).to(tl.float32) residual = tl.load(residual_ptr + row_offsets, mask=mask, other=0.0).to(tl.float32) residual_out = acc + residual res_out_dtype = residual_out_ptr.type.element_ty tl.store( residual_out_ptr + row_offsets, residual_out.to(res_out_dtype), mask=mask, ) variance = tl.sum(residual_out * residual_out, axis=0) / HIDDEN_SIZE scale = tl.rsqrt(variance + EPS) weight = tl.load(weight_ptr + offsets, mask=mask, other=0.0).to(tl.float32) norm = residual_out * scale * weight norm_dtype = norm_out_ptr.type.element_ty tl.store( norm_out_ptr + row_offsets, norm.to(norm_dtype), mask=mask, ) @triton.jit def iris_allreduce_residual_rmsnorm_kernel_persistent( input_sym_ptr, residual_ptr, weight_ptr, norm_out_ptr, residual_out_ptr, M, heap_bases, iris_rank: tl.constexpr, world_size: tl.constexpr, rank_start: tl.constexpr, rank_stride: tl.constexpr, HIDDEN_SIZE: tl.constexpr, BLOCK_SIZE: tl.constexpr, EPS: tl.constexpr, ): pid = tl.program_id(0) num_programs = tl.num_programs(0) offsets = tl.arange(0, BLOCK_SIZE) mask = offsets < HIDDEN_SIZE weight = tl.load(weight_ptr + offsets, mask=mask, other=0.0).to(tl.float32) res_out_dtype = residual_out_ptr.type.element_ty norm_dtype = norm_out_ptr.type.element_ty for row in range(pid, M, num_programs): row_offsets = row * HIDDEN_SIZE + offsets in_row_ptr = input_sym_ptr + row_offsets acc = tl.zeros((BLOCK_SIZE,), dtype=tl.float32) for i in tl.static_range(0, world_size): remote_rank = rank_start + i * rank_stride acc += iris.load( in_row_ptr, iris_rank, remote_rank, heap_bases, mask=mask, other=0.0, ).to(tl.float32) residual = tl.load(residual_ptr + row_offsets, mask=mask, other=0.0).to( tl.float32 ) residual_out = acc + residual tl.store( residual_out_ptr + row_offsets, residual_out.to(res_out_dtype), mask=mask, ) variance = tl.sum(residual_out * residual_out, axis=0) / HIDDEN_SIZE scale = tl.rsqrt(variance + EPS) norm = residual_out * scale * weight tl.store( norm_out_ptr + row_offsets, norm.to(norm_dtype), mask=mask, ) class IrisAllReduceResidualRMSNorm(object): def __init__( self, group: dist.ProcessGroup, rank_in_group: int, max_token_num: int, hidden_dim: int, dtype: torch.dtype = torch.bfloat16, heap_size: int | None = None, device: torch.device = None, persistent: bool = False, ) -> None: assert ( type(group) == dist.ProcessGroup ), f"Expected dist.ProcessGroup, got {type(group)}" assert dist.is_initialized(), ( "torch.distributed must be initialized before constructing " "IrisAllReduceResidualRMSNorm; call dist.init_process_group() first." ) assert _platform.is_amd, ( "IrisAllReduceResidualRMSNorm currently targets AMD ROCm; " f"got non-AMD platform: {_platform}" ) self.group = group self.rank_in_group = rank_in_group self.world_size = group.size() self.max_token_num = max_token_num self.hidden_dim = hidden_dim self.dtype = dtype self.device = device or torch.device(f"cuda:{torch.cuda.current_device()}") if heap_size is None: buf_bytes = max_token_num * hidden_dim * dtype.itemsize heap_size = max(1 << 28, 4 * buf_bytes + (16 << 20)) free_gpu_memory_begin = _get_available_gpu_memory(torch.cuda.current_device()) self._ctx = _get_or_create_iris_context(heap_size) self._input_buf = self._ctx.zeros((max_token_num, hidden_dim), dtype=dtype) free_gpu_memory_after = _get_available_gpu_memory(torch.cuda.current_device()) logger.info( "Iris AR+RMSNorm symmetric-heap buffer allocated: %s GB", free_gpu_memory_begin - free_gpu_memory_after, ) self._rank_start = 0 self._rank_stride = 1 self._iris_rank = dist.get_rank() self.persistent = persistent self._num_programs = ( torch.cuda.get_device_properties(self.device).multi_processor_count if persistent else 0 ) def fused( self, input_tensor: torch.Tensor, residual: torch.Tensor, weight: torch.Tensor, eps: float, norm_out: torch.Tensor | None = None, residual_out: torch.Tensor | None = None, ) -> Tuple[torch.Tensor, torch.Tensor]: assert input_tensor.dtype == self.dtype, ( f"Iris AR+RMSNorm dtype mismatch: input={input_tensor.dtype}, " f"backend={self.dtype}" ) assert input_tensor.dim() == 2, ( f"input must be 2-D (num_tokens, hidden_dim), got " f"shape={input_tensor.shape}" ) assert ( input_tensor.shape == residual.shape ), f"residual shape {residual.shape} != input shape {input_tensor.shape}" assert input_tensor.shape[1] == self.hidden_dim, ( f"hidden_dim mismatch: input={input_tensor.shape[1]} vs " f"backend={self.hidden_dim}" ) num_tokens = input_tensor.shape[0] assert num_tokens <= self.max_token_num, ( f"num_tokens ({num_tokens}) exceeds max_token_num " f"({self.max_token_num})" ) assert weight.shape == ( self.hidden_dim, ), f"weight shape {weight.shape} != ({self.hidden_dim},)" assert input_tensor.is_contiguous() and residual.is_contiguous() in_view = self._input_buf[:num_tokens, :] in_view.copy_(input_tensor) if norm_out is None: norm_out = torch.empty_like(input_tensor) if residual_out is None: residual_out = torch.empty_like(residual) self._ctx.device_barrier() heap_bases = self._ctx.get_heap_bases() BLOCK_SIZE = triton.next_power_of_2(self.hidden_dim) if self.persistent: kernel = iris_allreduce_residual_rmsnorm_kernel_persistent grid = (min(num_tokens, self._num_programs),) else: kernel = iris_allreduce_residual_rmsnorm_kernel grid = (num_tokens,) kernel[grid]( in_view, residual, weight, norm_out, residual_out, num_tokens, heap_bases, iris_rank=self._iris_rank, world_size=self.world_size, rank_start=self._rank_start, rank_stride=self._rank_stride, HIDDEN_SIZE=self.hidden_dim, BLOCK_SIZE=BLOCK_SIZE, EPS=eps, num_warps=8, ) return norm_out, residual_out def create_iris_state( group: dist.ProcessGroup, rank_in_group: int, max_numel: int, dtype: torch.dtype = torch.bfloat16, heap_size: int | None = None, device: torch.device = None, ) -> "IrisAllReduce": return IrisAllReduce( group=group, rank_in_group=rank_in_group, max_numel=max_numel, dtype=dtype, heap_size=heap_size, device=device, ) def iris_all_reduce( state: "IrisAllReduce", tensor: torch.Tensor, op=None, safe: bool = True, async_op: bool = False, ) -> torch.Tensor: return state.all_reduce(tensor, op=op, safe=safe, async_op=async_op) def create_iris_rsag_state( group: dist.ProcessGroup, rank_in_group: int, max_tokens: int, hidden_size: int, device: torch.device = None, heap_size: int | None = None, ) -> "IrisRSAG": return IrisRSAG( group=group, rank_in_group=rank_in_group, max_tokens=max_tokens, hidden_size=hidden_size, device=device, heap_size=heap_size, ) def create_iris_ar_rmsnorm_state( group: dist.ProcessGroup, rank_in_group: int, max_token_num: int, hidden_dim: int, dtype: torch.dtype = torch.bfloat16, heap_size: int | None = None, device: torch.device = None, persistent: bool = False, ) -> "IrisAllReduceResidualRMSNorm": return IrisAllReduceResidualRMSNorm( group=group, rank_in_group=rank_in_group, max_token_num=max_token_num, hidden_dim=hidden_dim, dtype=dtype, heap_size=heap_size, device=device, persistent=persistent, ) def iris_allreduce_residual_rmsnorm( state: "IrisAllReduceResidualRMSNorm", input_tensor: torch.Tensor, residual: torch.Tensor, weight: torch.Tensor, eps: float = 1e-6, norm_out: torch.Tensor | None = None, residual_out: torch.Tensor | None = None, ) -> Tuple[torch.Tensor, torch.Tensor]: return state.fused( input_tensor=input_tensor, residual=residual, weight=weight, eps=eps, norm_out=norm_out, residual_out=residual_out, )