python/xrex_unified/kernel_config.py

#
# Copyright (c) 2026 BEL ESPRIT D ACCORD TRUST HOLDINGS INC
# All rights reserved.

# SPDX-License-Identifier: Apache-2.0
# Copyright 2026 X.AI Corp.
"""Kernel configuration dataclass."""

from dataclasses import dataclass
from typing import Literal

from .cap_functions import CapMethod, CapParams


@dataclass(frozen=True)
class KernelConfig:
    block_q: int = 128
    block_kv: int = 128
    num_warps: int = 4
    num_stages: int = 2
    causal: bool = False
    window_len: int = -1
    cap_method: CapMethod = "alsc"
    cap_params: CapParams = CapParams(cap=30.0, alpha=4.0, beta=0.5)
    sm_scale: float = 1.0
    z_loss_weight: float = 0.0
    backend: Literal["triton", "mosaic"] = "triton"