[AUTOTUNER] Add (XGB+GA) FlagTune support for Triton v3.6.x

[HenryRenYz]

Summary

This PR adds the initial FlagTune integration for FlagTree/Triton v3.6.x.

FlagTune introduces a model-assisted autotuning path that can propose high-quality kernel configurations before the regular benchmarking loop. The implementation is opt-in through environment variables and preserves the default Triton autotuner behavior when FlagTune is disabled, unavailable, or unable to produce valid candidates.

Motivation

The existing autotuning flow can require benchmarking a large configuration space for every new shape. FlagTune reduces this search cost by using a trained XGBoost ranking model to predict top candidate configurations, with an optional genetic-algorithm refinement stage when benchmark feedback is available.

This is especially useful for matmul/GEMM workloads on Hopper-class backends, where the valid and performant configuration space is large and shape-dependent.

What Changed

• Added the triton.flagtune Python package.
• Added a Flagtuner wrapper and @flagtune decorator as a drop-in autotuner path.
• Added the generic ConfigProposer API:
  • make_config_proposer(...)
  • predict_configs(...)
  • make_early_config_prune(...) for backward compatibility
• Added core abstractions for:
  • operator input spaces
  • parameter spaces
  • feature pipelines
  • benchmark functions
  • config proposers
• Added XGBoost-based ranking and prediction support.
• Added GA-based refinement seeded from XGBoost top-k predictions.
• Added model discovery, validation, caching, and optional remote download support.
• Added an operator registry with lazy adapter discovery.
• Added an initial matmul adapter for mm_general_tma / flagtree/gemm.
• Added shape/config validation for Hopper TMA shared-memory constraints.
• Added benchmark scripts for comparing default autotune vs FlagTune on matmul shapes.

User-Facing Behavior

FlagTune is disabled by default.

To enable it:

TRITON_USE_FLAGTUNE=1

Note: has no relation with USE_FLAGTUNE in FlagGems, that is a legacy variable.

Optional environment variables:

TRITON_FLAGTUNE_MODEL_DIR=/path/to/model
TRITON_FLAGTUNE_TOP_K=10
FLAGTUNE_MODEL_CACHE=~/.flagtree/flagtune_models
FLAGTUNE_MODEL_URLS=/path/to/model_urls.json
FLAGTUNE_DISABLE_REMOTE=1
FLAGTUNE_DISABLE_OPS=flagtree/gemm

Example usage:

from triton.flagtune import flagtune

@flagtune(
    configs=configs,
    key=["M", "N", "K"],
    flagtune_op_id="flagtree/gemm",
)
@triton.jit
def kernel(...):
    ...

When FlagTune is disabled or initialization fails, the code falls back to the normal autotuner path.

Or, if a low-level API is desired:

from triton.flagtune import make_config_proposer

proposer = make_config_proposer({"op_id": "flagtree/gemm"})

shape = {
    "M": 4096,
    "N": 4096,
    "K": 4096,
    "stride_am": 4096,
    "stride_bk": 4096,
}

initial_configs = [
    {
        "BLOCK_M": 128,
        "BLOCK_N": 128,
        "BLOCK_K": 64,
        "GROUP_M": 8,
        "SPLIT_K": 1,
        "num_warps": 4,
        "num_stages": 3,
    }
]

configs = proposer(
    None,  # no benchmark function: return model-ranked candidates
    shape,
    initial_configs,
    {"op_id": "flagtree/gemm"},
)

The proposer is basically a function that maps (benchmark fn, input shape, initial configs, operator metas) into a new set of candidate configs that provides to caller for further selection/tuning.

Compatibility Notes

• The feature is opt-in and should not affect existing users unless TRITON_USE_FLAGTUNE=1 is set.
• Model loading supports user-provided model directories, local cache, bundled fallback models, and remote downloads.
• The matmul adapter includes conservative validation to avoid returning structurally invalid Hopper TMA configurations.
• The existing autotuner pruning logic is still applied after FlagTune proposes candidates.

Testing

Added benchmark/test coverage for matmul FlagTune behavior:

python/test/flagtune/test_bench_matmul.py
python/triton/flagtune/tests/test_bench_matmul.py

Recommended validation commands:

TRITON_USE_FLAGTUNE=0 python python/test/flagtune/test_bench_matmul.py --mode default

TRITON_USE_FLAGTUNE=1 \
TRITON_FLAGTUNE_MODEL_DIR=/path/to/model \
python python/test/flagtune/test_bench_matmul.py --mode flagtune

TRITON_USE_FLAGTUNE=1 \
TRITON_FLAGTUNE_MODEL_DIR=/path/to/model \
python python/test/flagtune/test_bench_matmul.py --compare

Risk

The main risk is model/config quality for newly supported shapes. This is mitigated by:

• keeping FlagTune disabled by default,
• validating proposed configs before use,
• preserving fallback behavior to the default autotuner,
• allowing per-op disablement through FLAGTUNE_DISABLE_OPS.

Files Added

• python/triton/flagtune/
• python/test/flagtune/test_bench_matmul.py
• python/triton/flagtune/tests/test_bench_matmul.py

[CLAassistant]

Thank you for your submission! We really appreciate it. Like many open source projects, we ask that you sign our Contributor License Agreement before we can accept your contribution.
_{You have signed the CLA already but the status is still pending? Let us recheck it.}