areno 0.0.0.dev0__tar.gz

areno-0.0.0.dev0/PKG-INFO

@@ -0,0 +1,19 @@
+Metadata-Version: 2.4
+Name: areno
+Version: 0.0.0.dev0
+Summary: Local LLM post-training stack with the areno engine and model plugins.
+Requires-Python: >=3.10
+Requires-Dist: torch>=2.6
+Requires-Dist: flash-attn>=2.7
+Requires-Dist: flash-linear-attention>=0.2
+Requires-Dist: safetensors>=0.4
+Requires-Dist: transformers>=4.56
+Requires-Dist: huggingface-hub>=0.25
+Requires-Dist: datasets>=3.3.0
+Requires-Dist: numpy
+Requires-Dist: tensorboard
+Requires-Dist: einops
+Requires-Dist: fastapi>=0.110
+Requires-Dist: click>=8.1
+Requires-Dist: tqdm>=4.66
+Requires-Dist: uvicorn>=0.27

areno-0.0.0.dev0/areno/__init__.py

@@ -0,0 +1,66 @@
+"""Top-level areno package.
+
+Sets process-wide knobs that must be in place before any CUDA/Triton kernel or
+torch.compile call runs: a single CUDA stream for collectives and a generous
+TorchDynamo cache so the engine can compile many specialized graphs (per shape
+bucket, prefill vs decode, train vs infer) without thrashing.
+
+Exposes the user-facing surface: configuration dataclasses, the rollout
+output container, and the `ArenoEngine` coordinator.
+"""
+
+from __future__ import annotations
+
+import os
+
+# A single CUDA stream connection keeps NCCL collectives ordered with compute,
+# which is what areno's TP/DP all-reduce + all-gather patterns assume.
+os.environ.setdefault("CUDA_DEVICE_MAX_CONNECTIONS", "1")
+
+try:
+    import torch._dynamo as _dynamo
+except ModuleNotFoundError:
+    _dynamo = None
+
+if _dynamo is not None:
+    # Train, prefill, decode, scoring and multiple shape buckets all produce
+    # distinct compiled artifacts; raise the cache limits so recompilation does
+    # not evict graphs that will be replayed across RL steps.
+    _dynamo.config.cache_size_limit = max(_dynamo.config.cache_size_limit, 64)
+    try:
+        _dynamo.config.accumulated_cache_size_limit = max(_dynamo.config.accumulated_cache_size_limit, 256)
+    except AttributeError:
+        pass
+
+from areno.engine.log import configure_default_logging
+
+configure_default_logging()
+
+
+def __getattr__(name: str):
+    """Lazily expose engine symbols without importing kernel-heavy modules."""
+
+    if name == "ArenoEngine":
+        from areno.engine import ArenoEngine
+
+        return ArenoEngine
+    if name in {"EngineConfig", "ModelConfig", "OptimizerConfig", "RuntimeConfig"}:
+        from areno.engine import config
+
+        return getattr(config, name)
+    if name in {"RolloutOutput", "SamplingParams", "TrainStats"}:
+        from areno.engine import data
+
+        return getattr(data, name)
+    raise AttributeError(f"module {__name__!r} has no attribute {name!r}")
+
+__all__ = [
+    "ArenoEngine",
+    "EngineConfig",
+    "ModelConfig",
+    "OptimizerConfig",
+    "RolloutOutput",
+    "RuntimeConfig",
+    "SamplingParams",
+    "TrainStats",
+]

areno-0.0.0.dev0/areno/accel/__init__.py

@@ -0,0 +1,42 @@
+"""Public entry point for the ARENO acceleration shims.
+
+Re-exports thin Python wrappers around the ``areno.accel._areno_accel`` C++/CUDA
+extension. Each submodule defines a ``torch.autograd.Function`` (where backward
+is needed) and a ``@torch._dynamo.disable``-decorated user-facing function that
+performs argument validation, dispatches to the fused CUDA kernel and exposes a
+PyTorch-friendly signature. The kernels themselves live in ``csrc/`` and are
+built by ``setup.py``; if the extension is not built the first call into the
+shim raises ``ModuleNotFoundError`` at import time.
+"""
+from areno.accel.activations import areno_gelu_tanh_and_mul, areno_sigmoid, areno_silu, areno_silu_and_mul, areno_softplus
+from areno.accel.conv import areno_depthwise_causal_conv1d_silu, areno_depthwise_causal_conv1d_silu_decode, areno_packed_depthwise_causal_conv1d_silu
+from areno.accel.embedding import areno_vocab_embedding
+from areno.accel.linear import areno_grouped_linear, areno_linear
+from areno.accel.moe import areno_moe_permute, areno_moe_topk_permute, areno_moe_unpermute
+from areno.accel.normalization import areno_optional_scale_rmsnorm, areno_rmsnorm, areno_rmsnorm_silu_gate
+from areno.accel.router import areno_grouped_topk_router
+from areno.accel.routing import areno_moe_align
+from areno.accel.topk import areno_topk_softmax
+
+__all__ = [
+    "areno_depthwise_causal_conv1d_silu",
+    "areno_depthwise_causal_conv1d_silu_decode",
+    "areno_packed_depthwise_causal_conv1d_silu",
+    "areno_gelu_tanh_and_mul",
+    "areno_grouped_topk_router",
+    "areno_grouped_linear",
+    "areno_linear",
+    "areno_moe_align",
+    "areno_moe_permute",
+    "areno_moe_topk_permute",
+    "areno_moe_unpermute",
+    "areno_optional_scale_rmsnorm",
+    "areno_rmsnorm",
+    "areno_rmsnorm_silu_gate",
+    "areno_sigmoid",
+    "areno_silu",
+    "areno_silu_and_mul",
+    "areno_softplus",
+    "areno_topk_softmax",
+    "areno_vocab_embedding",
+]

areno-0.0.0.dev0/areno/accel/_extension.py

@@ -0,0 +1,24 @@
+"""Lazy loader for the compiled ``areno.accel._areno_accel`` C++/CUDA extension.
+
+The extension module is imported on first use rather than at package import
+time so that ``import areno.accel`` succeeds in environments where only the
+Python shims are needed (e.g. for type checking). Each shim calls
+``extension()`` to obtain the compiled module and dispatch into the fused
+kernel. There is no pure-Python fallback: if the extension was not built the
+``importlib.import_module`` call below raises ``ModuleNotFoundError``.
+"""
+from __future__ import annotations
+
+import importlib
+from types import ModuleType
+
+# Cached reference to the compiled extension; populated on first call.
+_EXT: ModuleType | None = None
+
+
+def extension() -> ModuleType:
+    """Return the compiled C++/CUDA extension module, importing it lazily."""
+    global _EXT
+    if _EXT is None:
+        _EXT = importlib.import_module("areno.accel._areno_accel")
+    return _EXT

areno-0.0.0.dev0/areno/accel/activations.py

@@ -0,0 +1,172 @@
+"""Fused activation kernels (SiLU, GELU-tanh, sigmoid, softplus, gated variants).
+
+Each wrapper dispatches to the ARENO CUDA kernel via the compiled extension and
+preserves autograd by routing through ``torch.autograd.Function`` whenever the
+input requires gradients. The ``*_and_mul`` variants implement the common
+"gated MLP" pattern where the last input dimension is split in two and the
+first half is activated then element-wise multiplied with the second half,
+producing an output with half the last-dimension size.
+"""
+import torch
+
+from areno.accel._extension import extension as _extension
+
+
+def _activation_out(x: torch.Tensor, out: torch.Tensor | None) -> torch.Tensor:
+    """Allocate or validate the half-width output tensor for ``*_and_mul`` ops."""
+    if x.shape[-1] % 2 != 0:
+        raise ValueError(f"activation input last dimension must be even, got {x.shape[-1]}")
+    hidden = x.shape[-1] // 2
+    expected_shape = (*x.shape[:-1], hidden)
+    if out is None:
+        return torch.empty(expected_shape, device=x.device, dtype=x.dtype)
+    if tuple(out.shape) != expected_shape:
+        raise ValueError(f"activation output shape must be {expected_shape}, got {tuple(out.shape)}")
+    return out
+
+
+def _can_use_cuda_extension(x: torch.Tensor) -> bool:
+    """Guard that the input lives on CUDA; the kernels have no CPU path."""
+    if not x.is_cuda:
+        raise RuntimeError("ARENO activation kernels require CUDA tensors")
+    return True
+
+
+class _SiluMul(torch.autograd.Function):
+    """Autograd glue for fused SiLU(x[..., :H]) * x[..., H:]."""
+
+    @staticmethod
+    def forward(ctx, x: torch.Tensor) -> torch.Tensor:
+        out = _activation_out(x, None)
+        _extension().areno_silu_and_mul(out, x.contiguous())
+        ctx.save_for_backward(x)
+        return out
+
+    @staticmethod
+    def backward(ctx, grad_output: torch.Tensor) -> tuple[torch.Tensor]:
+        (x,) = ctx.saved_tensors
+        grad_input = torch.empty_like(x)
+        _extension().areno_d_silu_and_mul(grad_input, grad_output.contiguous(), x.contiguous())
+        return (grad_input,)
+
+
+class _Silu(torch.autograd.Function):
+    """Autograd glue for the element-wise SiLU kernel."""
+
+    @staticmethod
+    def forward(ctx, x: torch.Tensor) -> torch.Tensor:
+        out = _extension().areno_silu(x.contiguous())
+        ctx.save_for_backward(x)
+        return out
+
+    @staticmethod
+    def backward(ctx, grad_output: torch.Tensor) -> tuple[torch.Tensor]:
+        (x,) = ctx.saved_tensors
+        return (_extension().areno_d_silu(grad_output.contiguous(), x.contiguous()),)
+
+
+class _Sigmoid(torch.autograd.Function):
+    """Autograd glue for the element-wise sigmoid kernel."""
+
+    @staticmethod
+    def forward(ctx, x: torch.Tensor) -> torch.Tensor:
+        out = _extension().areno_sigmoid(x.contiguous())
+        ctx.save_for_backward(out)
+        return out
+
+    @staticmethod
+    def backward(ctx, grad_output: torch.Tensor) -> tuple[torch.Tensor]:
+        (out,) = ctx.saved_tensors
+        return (_extension().areno_d_sigmoid(grad_output.contiguous(), out.contiguous()),)
+
+
+class _Softplus(torch.autograd.Function):
+    """Autograd glue for the element-wise softplus kernel."""
+
+    @staticmethod
+    def forward(ctx, x: torch.Tensor) -> torch.Tensor:
+        out = _extension().areno_softplus(x.contiguous())
+        ctx.save_for_backward(x)
+        return out
+
+    @staticmethod
+    def backward(ctx, grad_output: torch.Tensor) -> tuple[torch.Tensor]:
+        (x,) = ctx.saved_tensors
+        return (_extension().areno_d_softplus(grad_output.contiguous(), x.contiguous()),)
+
+
+class _GeluTanhMul(torch.autograd.Function):
+    """Autograd glue for fused tanh-approx GELU(x[..., :H]) * x[..., H:]."""
+
+    @staticmethod
+    def forward(ctx, x: torch.Tensor) -> torch.Tensor:
+        out = _activation_out(x, None)
+        _extension().areno_gelu_tanh_and_mul(out, x.contiguous())
+        ctx.save_for_backward(x)
+        return out
+
+    @staticmethod
+    def backward(ctx, grad_output: torch.Tensor) -> tuple[torch.Tensor]:
+        (x,) = ctx.saved_tensors
+        grad_input = torch.empty_like(x)
+        _extension().areno_d_gelu_tanh_and_mul(grad_input, grad_output.contiguous(), x.contiguous())
+        return (grad_input,)
+
+
+@torch._dynamo.disable
+def areno_silu_and_mul(x: torch.Tensor, out: torch.Tensor | None = None) -> torch.Tensor:
+    """Apply SiLU to the first half of the last dimension and multiply by the second half.
+
+    Input shape (..., 2H) -> output shape (..., H). When an ``out`` tensor is
+    supplied the autograd path is skipped and the kernel writes in-place.
+    """
+    result = _activation_out(x, out)
+    _can_use_cuda_extension(x)
+    # Only enter the autograd Function when we actually need gradients; the
+    # plain kernel path is used for inference / when an out buffer is given.
+    if out is None and torch.is_grad_enabled() and x.requires_grad:
+        return _SiluMul.apply(x)
+    _extension().areno_silu_and_mul(result, x.contiguous())
+    return result
+
+
+@torch._dynamo.disable
+def areno_silu(x: torch.Tensor) -> torch.Tensor:
+    """Apply SiLU with an ARENO CUDA kernel."""
+    _can_use_cuda_extension(x)
+    if torch.is_grad_enabled() and x.requires_grad:
+        return _Silu.apply(x)
+    return _extension().areno_silu(x.contiguous())
+
+
+@torch._dynamo.disable
+def areno_sigmoid(x: torch.Tensor) -> torch.Tensor:
+    """Apply sigmoid with an ARENO CUDA kernel."""
+    _can_use_cuda_extension(x)
+    if torch.is_grad_enabled() and x.requires_grad:
+        return _Sigmoid.apply(x)
+    return _extension().areno_sigmoid(x.contiguous())
+
+
+@torch._dynamo.disable
+def areno_softplus(x: torch.Tensor) -> torch.Tensor:
+    """Apply softplus(beta=1, threshold=20) with an ARENO CUDA kernel."""
+    _can_use_cuda_extension(x)
+    if torch.is_grad_enabled() and x.requires_grad:
+        return _Softplus.apply(x)
+    return _extension().areno_softplus(x.contiguous())
+
+
+@torch._dynamo.disable
+def areno_gelu_tanh_and_mul(x: torch.Tensor, out: torch.Tensor | None = None) -> torch.Tensor:
+    """Apply tanh-approximate GELU to the first half and multiply by the second half.
+
+    Input shape (..., 2H) -> output shape (..., H). Matches the GeGLU
+    formulation used by Gemma-family MLPs.
+    """
+    result = _activation_out(x, out)
+    _can_use_cuda_extension(x)
+    if out is None and torch.is_grad_enabled() and x.requires_grad:
+        return _GeluTanhMul.apply(x)
+    _extension().areno_gelu_tanh_and_mul(result, x.contiguous())
+    return result

areno-0.0.0.dev0/areno/accel/conv.py

@@ -0,0 +1,141 @@
+"""Fused depthwise causal Conv1d + SiLU kernels for short-range token mixing.
+
+Used by the linear-attention / Mamba-style layers in areno to mix recent
+tokens within each channel. The kernel performs the convolution with explicit
+causal padding (no future leak) and applies SiLU in the same pass; weights are
+always coerced to ``float32`` so the CUDA path matches the reference layout.
+The ``*_decode`` entry point handles the single-token autoregressive case
+using a separately maintained history cache.
+"""
+import torch
+
+from areno.accel._extension import extension as _extension
+
+
+def _check_weight_shape(weight: torch.Tensor) -> None:
+    """Validate depthwise weight layout (channels, 1, kernel_size)."""
+    if weight.dim() != 3 or weight.shape[1] != 1:
+        raise ValueError(f"weight must have shape (channels, 1, kernel), got {tuple(weight.shape)}")
+
+
+def _kernel_weight(weight: torch.Tensor) -> torch.Tensor:
+    """Cast convolution weights to float32 as required by the CUDA kernel."""
+    _check_weight_shape(weight)
+    return weight if weight.dtype == torch.float32 else weight.float()
+
+
+class _DepthwiseCausalConv1dSilu(torch.autograd.Function):
+    """Autograd glue for fused depthwise causal conv1d + SiLU.
+
+    Forward returns the activated output; ``preact`` is saved so the backward
+    pass can recompute the SiLU derivative cheaply.
+    """
+
+    @staticmethod
+    def forward(ctx, x: torch.Tensor, weight: torch.Tensor) -> torch.Tensor:
+        out, preact = _extension().areno_depthwise_causal_conv1d_silu_forward(x.contiguous(), weight.contiguous())
+        ctx.save_for_backward(x, weight, preact)
+        return out
+
+    @staticmethod
+    def backward(ctx, grad_output: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
+        x, weight, preact = ctx.saved_tensors
+        grad_input, grad_weight = _extension().areno_depthwise_causal_conv1d_silu_backward(
+            grad_output.contiguous(),
+            x.contiguous(),
+            weight.contiguous(),
+            preact,
+        )
+        return grad_input, grad_weight
+
+
+class _PackedDepthwiseCausalConv1dSilu(torch.autograd.Function):
+    """Autograd glue for packed varlen depthwise causal conv1d + SiLU."""
+
+    @staticmethod
+    def forward(ctx, x: torch.Tensor, weight: torch.Tensor, cu_seqlens: torch.Tensor) -> torch.Tensor:
+        cu_seqlens = cu_seqlens.to(device=x.device, dtype=torch.int32).contiguous()
+        out, preact = _extension().areno_packed_depthwise_causal_conv1d_silu_forward(
+            x.contiguous(),
+            weight.contiguous(),
+            cu_seqlens,
+        )
+        ctx.save_for_backward(x, weight, cu_seqlens, preact)
+        return out
+
+    @staticmethod
+    def backward(ctx, grad_output: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor, None]:
+        x, weight, cu_seqlens, preact = ctx.saved_tensors
+        grad_input, grad_weight = _extension().areno_packed_depthwise_causal_conv1d_silu_backward(
+            grad_output.contiguous(),
+            x.contiguous(),
+            weight.contiguous(),
+            cu_seqlens,
+            preact,
+        )
+        return grad_input, grad_weight, None
+
+
+@torch._dynamo.disable
+def areno_depthwise_causal_conv1d_silu(x: torch.Tensor, weight: torch.Tensor) -> torch.Tensor:
+    """Apply depthwise causal conv1d followed by SiLU for (batch, seqlen, channels) tensors.
+
+    The kernel handles arbitrary ``seqlen`` and uses left-padding to keep the
+    convolution causal. Falls through to the inference path when autograd is
+    disabled to avoid stashing the pre-activation tensor.
+    """
+    if not x.is_cuda or not weight.is_cuda:
+        raise RuntimeError("areno_depthwise_causal_conv1d_silu requires CUDA input and weight")
+    if x.dim() != 3:
+        raise ValueError(f"input must have shape (batch, seqlen, channels), got {tuple(x.shape)}")
+    weight = _kernel_weight(weight)
+    if x.shape[-1] != weight.shape[0]:
+        raise ValueError(f"channel mismatch: input={x.shape[-1]} weight={weight.shape[0]}")
+    if torch.is_grad_enabled() and (x.requires_grad or weight.requires_grad):
+        return _DepthwiseCausalConv1dSilu.apply(x, weight)
+    out, _ = _extension().areno_depthwise_causal_conv1d_silu_forward(x.contiguous(), weight.contiguous())
+    return out
+
+
+@torch._dynamo.disable
+def areno_packed_depthwise_causal_conv1d_silu(x: torch.Tensor, weight: torch.Tensor, cu_seqlens: torch.Tensor) -> torch.Tensor:
+    """Apply depthwise causal conv1d followed by SiLU to packed (1, tokens, channels) tensors."""
+    if not x.is_cuda or not weight.is_cuda or not cu_seqlens.is_cuda:
+        raise RuntimeError("areno_packed_depthwise_causal_conv1d_silu requires CUDA tensors")
+    if x.dim() != 3 or x.shape[0] != 1:
+        raise ValueError(f"input must have shape (1, tokens, channels), got {tuple(x.shape)}")
+    weight = _kernel_weight(weight)
+    if x.shape[-1] != weight.shape[0]:
+        raise ValueError(f"channel mismatch: input={x.shape[-1]} weight={weight.shape[0]}")
+    if torch.is_grad_enabled() and (x.requires_grad or weight.requires_grad):
+        return _PackedDepthwiseCausalConv1dSilu.apply(x, weight, cu_seqlens)
+    out, _ = _extension().areno_packed_depthwise_causal_conv1d_silu_forward(
+        x.contiguous(),
+        weight.contiguous(),
+        cu_seqlens.to(device=x.device, dtype=torch.int32).contiguous(),
+    )
+    return out
+
+
+@torch._dynamo.disable
+def areno_depthwise_causal_conv1d_silu_decode(current: torch.Tensor, history: torch.Tensor, weight: torch.Tensor) -> torch.Tensor:
+    """Apply one-token depthwise causal conv1d followed by SiLU for decode.
+
+    ``current`` is the new token activations with shape ``(rows, channels)``
+    and ``history`` provides the previous ``kernel_size - 1`` tokens shaped
+    ``(rows, channels, kernel - 1)``. Returns the activated single-step output;
+    callers are responsible for shifting ``history``.
+    """
+    if not current.is_cuda or not history.is_cuda or not weight.is_cuda:
+        raise RuntimeError("areno_depthwise_causal_conv1d_silu_decode requires CUDA tensors")
+    if current.dim() != 2:
+        raise ValueError(f"current must have shape (rows, channels), got {tuple(current.shape)}")
+    if history.dim() != 3:
+        raise ValueError(f"history must have shape (rows, channels, kernel - 1), got {tuple(history.shape)}")
+    weight = _kernel_weight(weight)
+    out, _ = _extension().areno_depthwise_causal_conv1d_silu_decode(
+        current.contiguous(),
+        history.contiguous(),
+        weight.contiguous(),
+    )
+    return out

areno-0.0.0.dev0/areno/accel/embedding.py

@@ -0,0 +1,49 @@
+"""Tensor-parallel vocab embedding gather backed by an ARENO CUDA kernel.
+
+In TP-sharded embedding tables each rank only holds a contiguous slice of the
+vocabulary ``[vocab_start, vocab_end)``. The kernel gathers rows for ids that
+fall inside that range and writes zeros for out-of-range ids so the subsequent
+all-reduce reconstructs the full embedding.
+"""
+import torch
+
+from areno.accel._extension import extension as _extension
+
+
+class _VocabEmbedding(torch.autograd.Function):
+    """Autograd glue for the vocab-parallel embedding gather/scatter."""
+
+    @staticmethod
+    def forward(ctx, input_ids: torch.Tensor, weight: torch.Tensor, vocab_start: int, vocab_end: int) -> torch.Tensor:
+        out = _extension().areno_vocab_embedding_forward(input_ids.contiguous(), weight.contiguous(), int(vocab_start), int(vocab_end))
+        ctx.save_for_backward(input_ids, weight)
+        ctx.vocab_start = int(vocab_start)
+        ctx.vocab_end = int(vocab_end)
+        return out
+
+    @staticmethod
+    def backward(ctx, grad_output: torch.Tensor) -> tuple[None, torch.Tensor, None, None]:
+        input_ids, weight = ctx.saved_tensors
+        grad_weight = _extension().areno_vocab_embedding_backward(
+            grad_output.contiguous(),
+            input_ids.contiguous(),
+            weight,
+            ctx.vocab_start,
+            ctx.vocab_end,
+        )
+        return None, grad_weight, None, None
+
+
+@torch._dynamo.disable
+def areno_vocab_embedding(input_ids: torch.Tensor, weight: torch.Tensor, vocab_start: int, vocab_end: int) -> torch.Tensor:
+    """Gather rank-local vocab embeddings and zero out non-local token ids.
+
+    ``input_ids`` must be int64 on CUDA. ``weight`` is the local shard with
+    shape ``(vocab_end - vocab_start, hidden)``. Returns embeddings with shape
+    ``(*input_ids.shape, hidden)`` ready for tensor-parallel reduction.
+    """
+    if not input_ids.is_cuda or not weight.is_cuda:
+        raise RuntimeError("areno_vocab_embedding requires CUDA input_ids and weight")
+    if input_ids.dtype != torch.long:
+        raise TypeError("areno_vocab_embedding input_ids must be int64")
+    return _VocabEmbedding.apply(input_ids, weight, int(vocab_start), int(vocab_end))

areno-0.0.0.dev0/areno/accel/kernels/__init__.py

@@ -0,0 +1 @@
+"""Triton kernels that belong to the areno.accel acceleration surface."""