# Triton Ops

This folder contains Triton-based kernels and wrappers used by AITER.

## AITER Triton Ops Code Reorganization

Reorganized the flat Triton files into categorized subdirectories. **All existing imports still work** via backward compatibility.

### What Changed

#### Old Structure

```text
aiter/ops/triton/
├── __init__.py
├── gemm_a16w16.py
├── gemm_a8w8.py
├── batched_gemm_a16wfp4.py
├── mha.py
├── fused_gemm_afp4wfp4_a16w16.py
├── ... (other files in a flat structure)
└── _triton_kernels/
    ├── gemm_a16w16.py
    ├── gemm_a8w8.py
    └── ... (matching kernel files)
```

#### New Structure

```text
aiter/ops/triton/
├── __init__.py (with backward compatibility)
├── gemm/
│   ├── basic/          # Basic GEMM
│   ├── batched/        # Batched GEMM
│   ├── feed_forward/   # Feed-forward specific GEMMs
│   └── fused/          # Fused GEMM
├── attention/          # Attention (MHA, MQA, etc.)
├── moe/                # MOE
├── normalization/      # Normalization
├── quant/              # Quantization
├── rope/               # Rope
├── fusions/            # Other fusion
└── utils/              # Utility functions

aiter/ops/triton/_triton_kernels/
└── (uses the same new structure as above)
```

### Backward Compatibility

Old imports **still work** - no changes needed:

```python
# Both work identically:
# Old (still works via backward compatibility (_BACKWARD_COMPAT_MAP) in `aiter/ops/triton/__init__.py`)
    from aiter.ops.triton.gemm_a16w16 import gemm_a16w16
# New
from aiter.ops.triton.gemm.basic.gemm_a16w16 import gemm_a16w16
```

For **Imports**: Use new structure in the wrapper

## GEMM config loading (single utility)

GEMM kernels now use a single utility function instead of managing their own config loading.

- Location: `aiter/ops/triton/utils/gemm_config_utils.py`
- Function: `get_gemm_config(config_name, M, N=None, K=None, bounds=None, specialized_filename=None)`

### Config file formats

#### Deprecated format

```json
{
  "large": { "...": "..." },
  "small": { "...": "..." }
}
```

#### New format (required)

```json
{
  "M_LEQ_64": { "...": "..." },
  "M_GEQ_4096": { "...": "..." },
  "any": { "...": "..." }
}
```

Meaning:

- `M_LEQ_64`: applies when `M <= 64`
- `M_GEQ_4096`: applies when `M >= 4096`
- `any`: fallback for all other `M` values (**must exist if a config exists**)

### How config selection works

Config files are loaded from:

- Default configs: `aiter/ops/triton/configs/gemm/{arch}-{config_name}.json`
- Specialized configs (optional): `{arch}-{config_name}-N={N}-K={K}.json`

Selection rules for the `M` dimension:

1. Searches `M_LEQ_x` keys in ascending order using this default list:
   `[4, 8, 16, 32, 64, 128, 256, 512, 1024, 2048, 4096, 8192]`
2. If no match, searches `M_GEQ_x` keys in descending order.
3. Falls back to `any` if no bounds match.

### Usage patterns

#### Normal case (most common)

```python
def _get_config(M: int, N: int, K: int):
    return get_gemm_config("GEMM-A16W16", M, N, K)
```

#### With Split-K

```python
def _get_config(M: int, N: int, K: int):
    config, is_tuned = get_gemm_config("GEMM-A16W16", M, N, K)
    return compute_splitk_params(config, K), is_tuned
```

#### With custom bounds

```python
def _get_config(M: int, N: int, K: int):
    return get_gemm_config("FF-A16W16-fused", M, N, K, bounds=[4, 8, 64, 4096])
```

#### With specialized filename (mostly for fused kernels)

```python
def _get_config(M: int, N_fp4: int, N_bf16: int, K: int, shuffle: bool = False):
    config_name = "FUSED-GEMM-AFP4WFP4-A16W16"
    specialized_filename = f"N4={N_fp4}-N16={N_bf16}-K={2*K}"
    return get_gemm_config(config_name, M, specialized_filename=specialized_filename)
```

### Config file naming convention

#### Default config

- Pattern: `{arch}-{CONFIG_NAME}.json`
- Example: `gfx950-GEMM-A16W16.json`

#### Specialized config (N, K specific)

- Pattern: `{arch}-{CONFIG_NAME}-N={N}-K={K}.json`
- Example: `gfx950-GEMM-A16W16-N=256-K=7168.json`

#### Specialized config (custom)

- Pattern: `{arch}-{CONFIG_NAME}-{custom_suffix}.json`
- Example: `gfx950-FUSED-GEMM-AFP4WFP4-A16W16-N4=512-N16=256-K=7168.json`

### Config name patterns

- Basic GEMM: `GEMM-A{x}W{y}` (e.g. `GEMM-A16W16`, `GEMM-A8W8`)
- Batched GEMM: `BATCHED_GEMM-A{x}W{y}` (e.g. `BATCHED_GEMM-A16W16`)
- Fused ops: `FUSED-GEMM-{operation}` (e.g. `FUSED-GEMM-A8W8_BLOCKSCALE-A16W16`)
- Feed-forward: `FF-A{x}W{y}-fused` (e.g. `FF-A16W16-fused`)

Variants: append suffix like `_BLOCKSCALE`, `_PRESHUFFLED`, etc.

### Adding a new GEMM kernel (checklist)

- Import and use `get_gemm_config`.
- Add a `_get_config(M, N, K)` wrapper that calls `get_gemm_config(...)`.
- Use config keys `M_LEQ_x`, `M_GEQ_x`, `any` (do **not** use `large`/`small`).
- Name config files per the conventions above: `{arch}-{CONFIG_NAME}.json`.
- Add specialized configs when needed using the proper suffix.
- If the kernel supports split-K, use `compute_splitk_params()`.

## Config-aware kernel names in traces (repr)

Triton kernel names in traces now embed compile-time configuration parameters, making it easier to correlate performance results with the exact tuned config.

- Repr helper: `aiter/ops/triton/utils/_triton/kernel_repr.py`
- Function: `make_kernel_repr(base_name, config_keys)`

### Example usage

```python
kernel_repr = make_kernel_repr(
    "_gemm_a16_w16_kernel",
    ["BLOCK_SIZE_M", "BLOCK_SIZE_N", "..."],
)

@triton.jit(repr=kernel_repr)
def _gemm_a16_w16_kernel(...):
    ...
```

How it works:

- Reads `specialization.constants` and appends `KEY_VALUE` pairs to the kernel name.
- Values are sanitized for readability:
  - `None -> NONE`
  - `bool -> 0/1`
  - strings are uppercased and non-alnum characters become `_`

## Kernel doc comments

Kernels should include a concise doc comment (preferably a docstring near the wrapper):

Include:

- What the kernel computes
- Which configuration parameters are included (and what they control)
- What it returns
- Any special considerations

Example docstring template:

```python
"""
Computes batched 8-bit matrix multiplication Y[i] = X[i] @ W[i]^T with active activation quantization.
X is quantized to INT8 during computation using per-token grouped quantization.
W is pre-quantized INT8 with per-batch-element scaling.

Args:
    X (torch.Tensor): Higher precision input batch with shape (B, M, K) or (M, B, K) if transpose_bm_in=True.
        Quantized to INT8 on-the-fly during GEMM.
    WQ (torch.Tensor): Pre-quantized INT8 weight batch with shape (B, N, K), internally transposed.
    w_scale (torch.Tensor): Per-batch scale for WQ with shape (1,).
    group_size (int): Group size for per-token grouped quantization of X. Must be power of 2.
    bias (Optional[torch.Tensor]): Bias batch with shape (B, 1, N).
    dtype (Optional[torch.dtype]): Output datatype (BF16 or FP16).
    splitK (Optional[int]): Not supported. Must be None.
    YQ (Optional[torch.Tensor]): Pre-allocated output tensor with shape (B, M, N) or (M, B, N) if transpose_bm=True.
    transpose_bm (Optional[bool]): Transpose batch and M dimensions in output.
    transpose_bm_in (Optional[bool]): Transpose batch and M dimensions in input.
    config (Optional[dict]): Kernel tuning parameters (BLOCK_SIZE_M, BLOCK_SIZE_N, GROUP_SIZE_M).

Returns:
    torch.Tensor: Output batch with shape (B, M, N) or (M, B, N) if transpose_bm=True.
"""
```

## Architecture naming convention

Across Triton configuration files and related code, key behavior by GPU architecture identifiers (like `gfx950`) rather than product names.

- Config filenames must use `{arch}` like `gfx950`.
- Conditionals should compare against `DEVICE_ARCH` values like `"gfx950"`.

Preferred:

```python
if impl == "gluon" and DEVICE_ARCH not in ("gfx950",):
    ...
```

Avoid product-name parsing like:

```python
# Don't do this
if impl == "gluon" and int(DEVICE_ARCH.split("MI")[1].replace("X", "")) < mi_chip_number:
    ...
```

## Triton test organization

`op_tests/triton_tests` has been reorganized into categorized subfolders for easier maintenance and targeted test runs.

Categories include: `gemm/`, `moe/`, `attention/`, `quant/`, etc.

### How to run

- All Triton tests:
  - `pytest op_tests/triton_tests/`
- GEMM tests only:
  - `pytest op_tests/triton_tests/gemm/`
- Basic GEMM subset:
  - `pytest op_tests/triton_tests/gemm/basic/`
- MOE tests only:
  - `pytest op_tests/triton_tests/moe/`

### Test selection script

There is a test selection script, that is currently under evaluation, available in `.github/scripts/select_triton_tests.py` file. Its goal is to programmatically select which Triton tests to run, based on the diff content of a given AITER PR, so Triton test CI step can run faster.

The script builds a dependency graph of Triton source files, including kernel config files. The hard part is to relate kernel config files with their respective kernel. The script basically search for the following code patterns:

```python
# Interpolated strings referring to JSON config files:
f"{AITER_TRITON_CONFIGS_PATH}/{dev}-KERNEL.json"
```

```python
# Calls to `get_gemm_config` utility function in which `config_name` is a constant string:
get_gemm_config("AMAZING-GEMM", M, N, K)
```

It is a good practice, when writing kerels, to follow these patterns.

## Quick checklists

### Quick checklist when adding a new kernel

- Config filenames use `{arch}` like `gfx950`.
- Add config-aware trace naming via `make_kernel_repr(...)` and `@triton.jit(repr=...)`.
- Add a kernel docstring describing behavior + config params + returns + notes.
- If possible, adopt patterns recognized by test selection script when loading kernel configs.

### Quick checklist when adding arch-specific exceptions

- Compare `DEVICE_ARCH` to `("gfx950", "gfx942", ...)`.
- Do not parse product names.

### Quick checklist when adding Triton tests

- Place tests under the appropriate category in `op_tests/triton_tests/<category>/`.
