diff --git a/.gitignore b/.gitignore index d8f9754f7..d0f387685 100644 --- a/.gitignore +++ b/.gitignore @@ -50,6 +50,7 @@ __pycache__/ !.devcontainer/devcontainer.json !.github/scripts/examples_config.json +!iris/ops/configs/*.json resources/ diff --git a/MANIFEST.in b/MANIFEST.in index 2c255da11..5c693e5e4 100644 --- a/MANIFEST.in +++ b/MANIFEST.in @@ -4,4 +4,7 @@ include LICENSE include iris/README.md # Include build configuration -include pyproject.toml \ No newline at end of file +include pyproject.toml + +# Include AG+MM auto-config JSON files +recursive-include iris/ops/configs *.json \ No newline at end of file diff --git a/iris/ops/__init__.py b/iris/ops/__init__.py index e0d12ba51..97ea588cb 100644 --- a/iris/ops/__init__.py +++ b/iris/ops/__init__.py @@ -31,6 +31,16 @@ from .config import FusedConfig from .workspace import FusedWorkspace +from .auto_config import ( + AutoConfigResult, + select_ag_mm_config, + list_known_shapes, + load_regression_sizes, + clear_config_cache, + detect_gpu_arch, + SUPPORTED_TRANSPOSES, + SUPPORTED_ARCHITECTURES, +) # Import operations # from .matmul import matmul # Simple single-GPU GEMM - TODO: implement @@ -172,6 +182,15 @@ def matmul_reduce_scatter(self, output_tensor, A, B, bias=None, async_op=False, # Configuration "FusedConfig", "FusedWorkspace", + # Auto-selection + "AutoConfigResult", + "select_ag_mm_config", + "list_known_shapes", + "load_regression_sizes", + "clear_config_cache", + "detect_gpu_arch", + "SUPPORTED_TRANSPOSES", + "SUPPORTED_ARCHITECTURES", # Namespace "OpsNamespace", # Operations diff --git a/iris/ops/all_gather_matmul.py b/iris/ops/all_gather_matmul.py index 5d700206c..c5e1f83c4 100644 --- a/iris/ops/all_gather_matmul.py +++ b/iris/ops/all_gather_matmul.py @@ -164,9 +164,22 @@ def all_gather_matmul_preamble( B: torch.Tensor, config: Optional[FusedConfig] = None, ) -> FusedWorkspace: - """Allocate workspace for all_gather_matmul (none needed for pull pattern).""" + """Allocate workspace for all_gather_matmul (none needed for pull pattern). + + When config=None, uses auto-selection to pick the best known configuration. + """ if config is None: - config = FusedConfig() + from .auto_config import select_ag_mm_config + + M_auto, K_local_auto = A_sharded.shape + K_auto, N_auto = B.shape + world_size_auto = shmem.get_num_ranks() + auto_result = select_ag_mm_config(M_auto, N_auto, K_auto, world_size=world_size_auto) + if not auto_result.enabled: + raise RuntimeError( + f"iris AG+MM auto-config disabled: {auto_result.source}. Pass config=FusedConfig(...) to override." + ) + config = auto_result.to_fused_config() M, K_local = A_sharded.shape K, N = B.shape @@ -194,9 +207,28 @@ def all_gather_matmul( config: Optional[FusedConfig] = None, workspace: Optional[FusedWorkspace] = None, ) -> FusedWorkspace: - """Fused all-gather and matrix multiplication using pull pattern.""" + """Fused all-gather and matrix multiplication using pull pattern. + + When config=None, uses auto-selection to pick the best known configuration + for the given (M, N, K, world_size) on the current GPU. If the auto-config + disables iris for this combination (e.g., ws<8 on MI300X), raises RuntimeError + advising fallback to PyTorch. To bypass auto-selection, pass an explicit + FusedConfig instance. + """ if config is None: - config = FusedConfig() + from .auto_config import select_ag_mm_config + + M_auto, K_local_auto = A_sharded.shape + K_auto, N_auto = B.shape + world_size_auto = shmem.get_num_ranks() + auto_result = select_ag_mm_config(M_auto, N_auto, K_auto, world_size=world_size_auto) + if not auto_result.enabled: + raise RuntimeError( + f"iris AG+MM auto-config disabled for this shape/world_size: " + f"{auto_result.source}. Pass config=FusedConfig(...) to override, " + f"or use PyTorch all_gather + matmul instead." + ) + config = auto_result.to_fused_config() M, K_local = A_sharded.shape K, N = B.shape diff --git a/iris/ops/auto_config.py b/iris/ops/auto_config.py new file mode 100644 index 000000000..240b226d2 --- /dev/null +++ b/iris/ops/auto_config.py @@ -0,0 +1,513 @@ +# SPDX-License-Identifier: MIT +# Copyright (c) 2025-2026 Advanced Micro Devices, Inc. All rights reserved. + +""" +Auto-selection mechanism for fused AG+MM kernel configurations. + +Given problem dimensions (M, N, K), transpose mode, world_size, and GPU +architecture, this module selects the best known configuration or returns +a sensible default. For world sizes where iris AG+MM is known to lose +against PyTorch (ws<8), the default disables iris and signals fallback. + +Config files live under: + iris/ops/configs/ag_mm/{arch}/{transpose}/ws{N}.json + +Each config file contains: + - FusedConfig parameters (block sizes, group sizes, etc.) + - HBM buffer kernel parameters (k_per_flag, num_fetch_sms, etc.) + - Per-shape champion configs with verified speedup measurements + +Transpose coverage: + The iris AG+MM kernel (`_fused_all_gather_matmul_kernel`) uses stride-based + addressing (`stride_am, stride_ak, stride_bk, stride_bn`), so transpose + layouts are handled implicitly by tensor strides. Config files exist for + all four layouts (NN, TN, NT, TT) under each architecture directory. + Only NN has per-shape champion configs from benchmarking (3,489 trials). + TN/NT/TT files contain heuristic defaults only (empty shapes dict) and are + marked enabled at ws>=8 to allow heuristic fallback. All transposes at ws<8 + are disabled (NO-GO based on NN benchmarks). + +Usage: + >>> from iris.ops.auto_config import select_ag_mm_config + >>> result = select_ag_mm_config(M=131072, N=16384, K=16384, world_size=8) + >>> if result.enabled: + ... config = result.to_fused_config() + ... hbm_params = result.hbm_buffer_params # k_per_flag, num_fetch_sms, etc. + ... shmem.ops.all_gather_matmul(output, A, B, config=config) + ... else: + ... # Fallback to PyTorch all_gather + matmul + ... ... + + >>> # List all regression test sizes + >>> from iris.ops.auto_config import load_regression_sizes + >>> sizes = load_regression_sizes() +""" + +import json +import os +import subprocess +from dataclasses import dataclass, field +from pathlib import Path +from typing import Dict, List, Optional, Tuple + +from .config import FusedConfig + +# Root directory for config files (relative to this module) +_CONFIGS_DIR = Path(__file__).parent / "configs" / "ag_mm" + +# In-memory cache: (arch, transpose, world_size) -> loaded JSON data +_config_cache: Dict[Tuple[str, str, int], dict] = {} + +# Cached GPU architecture detection result +_detected_arch: Optional[str] = None + +# Supported transpose modes. The AG+MM kernel only supports NN layout. +# TN/NT/TT would require kernel-level changes to permute strides. +SUPPORTED_TRANSPOSES = ("NN",) + +# Supported GPU architectures with tuned configs +SUPPORTED_ARCHITECTURES = ("mi300x",) + +# Map gfx target IDs to architecture names used in config paths +_GFX_TO_ARCH = { + "gfx942": "mi300x", # MI300X, MI300A +} + + +def detect_gpu_arch() -> str: + """Auto-detect GPU architecture from the current system. + + Detection order: + 1. IRIS_GPU_ARCH environment variable (override) + 2. rocm-smi --showproductname parsing + 3. rocminfo gfx target parsing + 4. Falls back to "mi300x" (most common deployment target) + + Returns: + Architecture string (e.g., "mi300x") suitable for config lookup. + """ + global _detected_arch + if _detected_arch is not None: + return _detected_arch + + # 1. Environment variable override + env_arch = os.environ.get("IRIS_GPU_ARCH", "").strip().lower() + if env_arch: + _detected_arch = env_arch + return _detected_arch + + # 2. Try rocminfo for gfx target + try: + result = subprocess.run( + ["rocminfo"], + capture_output=True, + text=True, + timeout=5, + ) + if result.returncode == 0: + for line in result.stdout.splitlines(): + line_stripped = line.strip().lower() + if "name:" in line_stripped and "gfx" in line_stripped: + for gfx_id, arch_name in _GFX_TO_ARCH.items(): + if gfx_id in line_stripped: + _detected_arch = arch_name + return _detected_arch + except (FileNotFoundError, subprocess.TimeoutExpired, OSError): + pass + + # 3. Fallback to MI300X (most common deployment target) + _detected_arch = "mi300x" + return _detected_arch + + +@dataclass +class AutoConfigResult: + """Result of auto-config lookup. + + Attributes: + enabled: If False, iris AG+MM should NOT be used; fallback to PyTorch. + config_params: Dict of FusedConfig parameters (only valid if enabled=True). + hbm_buffer_params: Dict of HBM buffer-specific kernel params + (k_per_flag, num_fetch_sms, num_fetch_stages, first_stage_fetch_sms). + source: Human-readable description of where this config came from. + shape_key: The MxNxK key that matched (None if heuristic/default). + speedup: Expected speedup vs PyTorch (None if unknown). + """ + + enabled: bool = False + config_params: Dict = field(default_factory=dict) + hbm_buffer_params: Dict = field(default_factory=dict) + source: str = "default" + shape_key: Optional[str] = None + speedup: Optional[float] = None + + def to_fused_config(self) -> FusedConfig: + """Convert to FusedConfig for use with iris.ops functions. + + Raises: + RuntimeError: If this config is disabled (enabled=False). + """ + if not self.enabled: + raise RuntimeError( + f"Cannot create FusedConfig: iris AG+MM is disabled for this " + f"configuration. Reason: {self.source}. " + f"Use PyTorch all_gather + matmul instead." + ) + # Filter to only fields FusedConfig accepts + valid_fields = {f.name for f in FusedConfig.__dataclass_fields__.values()} + filtered = {k: v for k, v in self.config_params.items() if k in valid_fields} + return FusedConfig(**filtered) + + +def _load_config_file(arch: str, transpose: str, world_size: int) -> Optional[dict]: + """Load and cache a config JSON file. + + Args: + arch: GPU architecture identifier (e.g., "mi300x"). + transpose: Transpose mode (e.g., "NN", "NT", "TN", "TT"). + world_size: Number of ranks. + + Returns: + Parsed JSON dict, or None if file doesn't exist. + """ + cache_key = (arch, transpose, world_size) + if cache_key in _config_cache: + return _config_cache[cache_key] + + config_path = _CONFIGS_DIR / arch / transpose / f"ws{world_size}.json" + if not config_path.exists(): + _config_cache[cache_key] = None + return None + + with open(config_path, "r") as f: + data = json.load(f) + + _config_cache[cache_key] = data + return data + + +def _load_default_config() -> dict: + """Load the global default config.""" + default_path = _CONFIGS_DIR / "default_config.json" + if default_path.exists(): + with open(default_path, "r") as f: + return json.load(f) + return {} + + +def _find_nearest_shape(M: int, N: int, K: int, shapes: dict, tolerance: float = 0.15) -> Optional[str]: + """Find the nearest matching shape in the config database. + + Uses log-space geometric distance to find shapes that are structurally + similar (within `tolerance` ratio per dimension). This avoids falling + back to heuristic when the user's problem is close to a champion shape. + + Args: + M, N, K: Target dimensions. + shapes: Dict of shape_key -> shape_data from the config file. + tolerance: Max fractional distance per dimension (default 15%). + + Returns: + The shape_key of the nearest match, or None if no shape is close enough. + """ + import math + + best_key = None + best_dist = float("inf") + + for shape_key, shape_data in shapes.items(): + sm, sn, sk = shape_data["M"], shape_data["N"], shape_data["K"] + # Skip shapes with speedup <= 1.0 (losers) + if shape_data.get("speedup", 0) is not None and shape_data.get("speedup", 0) <= 1.0: + continue + + # Check per-dimension ratio tolerance + if sm == 0 or sn == 0 or sk == 0: + continue + rm = abs(M - sm) / sm + rn = abs(N - sn) / sn + rk = abs(K - sk) / sk + + if rm > tolerance or rn > tolerance or rk > tolerance: + continue + + # Geometric distance in log space + dist = math.sqrt( + math.log(max(M, 1) / max(sm, 1)) ** 2 + + math.log(max(N, 1) / max(sn, 1)) ** 2 + + math.log(max(K, 1) / max(sk, 1)) ** 2 + ) + if dist < best_dist: + best_dist = dist + best_key = shape_key + + return best_key + + +def _apply_heuristic(M: int, N: int, K: int) -> Tuple[Dict, Dict]: + """Apply heuristic rules to generate config + HBM buffer params. + + Based on optimization data (3,489 measured trials on MI300X): + - block_size_m: 128 for M <= 16384, else 256 + - group_size_m: 8 for M <= 8192, 16 for M <= 16384, else 24 + - k_per_flag: maximize for throughput (52% of perf range) + - num_fetch_sms: scale with M-tiles + - All other params are fixed invariants across all 7+ champion shapes. + + Args: + M: Rows dimension. + N: Columns dimension. + K: Reduction dimension. + + Returns: + Tuple of (config_params dict, hbm_buffer_params dict). + """ + bk = 64 + num_k_blocks = K // bk + + # block_size_m selection based on M + if M <= 16384: + bm = 128 + else: + bm = 256 + + num_m_tiles = M // bm + + # group_size_m selection + if M <= 8192: + gm = 8 + elif M <= 16384: + gm = 16 + else: + gm = 24 + + config_params = { + "block_size_m": bm, + "block_size_n": 256, + "block_size_k": bk, + "group_size_m": gm, + "num_warps": 8, + "num_stages": 2, + "num_xcds": 8, + "allow_tf32": True, + } + + # k_per_flag: #1 performance knob (52% of perf range from sweep) + if num_k_blocks >= 512: + kpf = 64 + elif num_k_blocks >= 128: + kpf = 16 + elif num_k_blocks >= 64: + kpf = 8 + else: + kpf = 4 + # Ensure divisibility + while num_k_blocks % kpf != 0 and kpf > 1: + kpf //= 2 + + # num_fetch_sms: scale with M-tiles + if num_m_tiles <= 8: + fs = 4 + elif num_m_tiles <= 32: + fs = 16 + elif num_m_tiles <= 128: + fs = 32 + else: + fs = 52 + + # num_fetch_stages + if num_m_tiles >= 512: + nfs = 4 + elif num_m_tiles >= 64: + nfs = 2 + else: + nfs = 1 + + hbm_params = { + "k_per_flag": kpf, + "num_fetch_sms": fs, + "num_fetch_stages": nfs, + "first_stage_fetch_sms": 64, + } + + return config_params, hbm_params + + +def select_ag_mm_config( + M: int, + N: int, + K: int, + world_size: int, + transpose: str = "NN", + arch: str = "auto", +) -> AutoConfigResult: + """Select the best AG+MM config for the given problem. + + Lookup order: + 1. Exact shape match in configs/ag_mm/{arch}/{transpose}/ws{world_size}.json + 2. Heuristic-based config from the same file's defaults + 3. Global default from configs/ag_mm/default_config.json + + For world sizes where iris is known to lose (ws<8 on MI300X), returns + a disabled result signaling fallback to PyTorch. + + Args: + M: Number of rows (or M_local * world_size for AG+MM). + N: Number of columns. + K: Reduction dimension. + world_size: Number of ranks in the communicator. + transpose: Transpose mode ("NN", "NT", "TN", "TT"). Default "NN". + arch: GPU architecture ("mi300x", etc.) or "auto" to auto-detect. + Default "auto". Set IRIS_GPU_ARCH env var to override. + + Returns: + AutoConfigResult with .enabled indicating whether to use iris, + .to_fused_config() to get the FusedConfig if enabled, and + .hbm_buffer_params with kernel-specific parameters. + + Example: + >>> result = select_ag_mm_config(131072, 16384, 16384, world_size=8) + >>> result.enabled + True + >>> result.speedup + 1.343 + >>> config = result.to_fused_config() + >>> result.hbm_buffer_params + {'k_per_flag': 32, 'num_fetch_sms': 4, 'num_fetch_stages': 64, 'first_stage_fetch_sms': 52} + + >>> result = select_ag_mm_config(4096, 4096, 4096, world_size=2) + >>> result.enabled + False + """ + transpose = transpose.upper() + if arch == "auto": + arch = detect_gpu_arch() + else: + arch = arch.lower() + + # Step 1: Try to load the specific config file + data = _load_config_file(arch, transpose, world_size) + + if data is not None: + # Check if this world_size is enabled + if not data.get("enabled", True): + return AutoConfigResult( + enabled=False, + source=f"Disabled by config: {arch}/{transpose}/ws{world_size}.json — {data.get('reason', 'no reason given')}", + ) + + # Look for exact shape match + shape_key = f"{M}x{N}x{K}" + shapes = data.get("shapes", {}) + if shape_key in shapes: + shape_data = shapes[shape_key] + return AutoConfigResult( + enabled=True, + config_params=shape_data["config"], + hbm_buffer_params=shape_data.get("hbm_buffer_params", {}), + source=f"Exact match: {arch}/{transpose}/ws{world_size}.json [{shape_data.get('label', shape_key)}]", + shape_key=shape_key, + speedup=shape_data.get("speedup"), + ) + + # No exact match — try nearest champion shape (within 15% per dim) + nearest_key = _find_nearest_shape(M, N, K, shapes) + if nearest_key is not None: + nearest_data = shapes[nearest_key] + return AutoConfigResult( + enabled=True, + config_params=nearest_data["config"], + hbm_buffer_params=nearest_data.get("hbm_buffer_params", {}), + source=f"Nearest match: {arch}/{transpose}/ws{world_size}.json [{nearest_data.get('label', nearest_key)}] (target {M}x{N}x{K} ≈ {nearest_key})", + shape_key=nearest_key, + speedup=nearest_data.get("speedup"), + ) + + # No nearby match — use heuristic + file defaults + file_default_config = data.get("default_config") + file_default_hbm = data.get("default_hbm_buffer_params", {}) + if file_default_config: + heuristic_config, heuristic_hbm = _apply_heuristic(M, N, K) + # Merge: heuristic provides shape-aware bm/gm, file_default provides rest + merged_config = {**file_default_config, **heuristic_config} + # For HBM params, prefer heuristic (shape-aware) over static defaults + merged_hbm = {**file_default_hbm, **heuristic_hbm} + return AutoConfigResult( + enabled=True, + config_params=merged_config, + hbm_buffer_params=merged_hbm, + source=f"Heuristic (no exact shape match in {arch}/{transpose}/ws{world_size}.json)", + ) + + # Step 2: No config file found — check global default + default_data = _load_default_config() + ws_gate = default_data.get("world_size_gate", {}) + min_ws = ws_gate.get("min_world_size", 8) + + if world_size < min_ws: + return AutoConfigResult( + enabled=False, + source=f"world_size={world_size} < min_world_size={min_ws} (global default). {ws_gate.get('reason', '')}", + ) + + # World size OK but no specific config — apply heuristic + heuristic_config, heuristic_hbm = _apply_heuristic(M, N, K) + return AutoConfigResult( + enabled=True, + config_params=heuristic_config, + hbm_buffer_params=heuristic_hbm, + source=f"Heuristic fallback (no config file for {arch}/{transpose}/ws{world_size})", + ) + + +def list_known_shapes( + world_size: int, + transpose: str = "NN", + arch: str = "mi300x", +) -> list: + """List all known shape configurations for a given world_size/transpose/arch. + + Returns: + List of dicts with keys: shape_key, label, M, N, K, speedup, n_trials. + """ + data = _load_config_file(arch, transpose.upper(), world_size) + if data is None or not data.get("enabled", True): + return [] + + result = [] + for shape_key, shape_data in data.get("shapes", {}).items(): + result.append( + { + "shape_key": shape_key, + "label": shape_data.get("label", ""), + "M": shape_data["M"], + "N": shape_data["N"], + "K": shape_data["K"], + "speedup": shape_data.get("speedup"), + "n_trials": shape_data.get("n_trials"), + } + ) + + # Sort by speedup descending + result.sort(key=lambda x: x.get("speedup", 0) or 0, reverse=True) + return result + + +def load_regression_sizes() -> List[Dict]: + """Load regression test sizes from the JSON config file. + + Returns: + List of regression size dicts, each with: name, M, N, K, tier, + description, world_sizes, expected, regression_threshold_pct. + """ + reg_path = _CONFIGS_DIR / "regression_sizes.json" + if not reg_path.exists(): + return [] + with open(reg_path, "r") as f: + data = json.load(f) + return data.get("sizes", []) + + +def clear_config_cache(): + """Clear the in-memory config cache. Useful after modifying config files.""" + _config_cache.clear() diff --git a/iris/ops/configs/__init__.py b/iris/ops/configs/__init__.py new file mode 100644 index 000000000..3834b4ce9 --- /dev/null +++ b/iris/ops/configs/__init__.py @@ -0,0 +1,2 @@ +# SPDX-License-Identifier: MIT +# Copyright (c) 2025 Advanced Micro Devices, Inc. All rights reserved. diff --git a/iris/ops/configs/ag_mm/__init__.py b/iris/ops/configs/ag_mm/__init__.py new file mode 100644 index 000000000..3834b4ce9 --- /dev/null +++ b/iris/ops/configs/ag_mm/__init__.py @@ -0,0 +1,2 @@ +# SPDX-License-Identifier: MIT +# Copyright (c) 2025 Advanced Micro Devices, Inc. All rights reserved. diff --git a/iris/ops/configs/ag_mm/default_config.json b/iris/ops/configs/ag_mm/default_config.json new file mode 100644 index 000000000..ff96ac1f0 --- /dev/null +++ b/iris/ops/configs/ag_mm/default_config.json @@ -0,0 +1,27 @@ +{ + "_meta": { + "description": "Global default fallback config for AG+MM operations. Disables iris AG+MM for ws<8 (fallback to PyTorch).", + "source": "benchmarking on MI300X (gfx942), 3489 measured trials", + "date": "2026-04-13" + }, + "world_size_gate": { + "min_world_size": 8, + "reason": "ws=2 best 0.89x, ws=4 best 0.86x vs PyTorch. Only ws>=8 is production-ready." + }, + "config": { + "block_size_m": 256, + "block_size_n": 256, + "block_size_k": 64, + "group_size_m": 24, + "num_warps": 8, + "num_stages": 2, + "num_xcds": 8, + "allow_tf32": true + }, + "hbm_buffer_params": { + "k_per_flag": 8, + "num_fetch_sms": 32, + "num_fetch_stages": 1, + "first_stage_fetch_sms": 64 + } +} diff --git a/iris/ops/configs/ag_mm/mi300x/NN/ws2.json b/iris/ops/configs/ag_mm/mi300x/NN/ws2.json new file mode 100644 index 000000000..98e940f42 --- /dev/null +++ b/iris/ops/configs/ag_mm/mi300x/NN/ws2.json @@ -0,0 +1,197 @@ +{ + "_meta": { + "description": "AG+MM ws=2 on MI300X — DISABLED (loses vs PyTorch on all shapes)", + "source": "benchmarking: best speedup 0.89x (g5). LDS overflow (65540>65536 bytes) forces ns=1, imposing 15-35% penalty.", + "gpu": "AMD Instinct MI300X (gfx942)", + "date": "2026-04-13", + "data_tag": "[VERIFIED]" + }, + "enabled": false, + "reason": "ws=2 AG transfers from 1 peer only. GEMM dominates latency. Fetch SM overhead exceeds overlap benefit. LDS overflow forces ns=1, imposing 15-35% penalty. Best measured speedup: 0.89x.", + "best_measured_speedup": 0.89, + "shapes": { + "8192x8192x262144": { + "label": "g5", + "description": "Best ws=2 result — still loses (0.887x)", + "M": 8192, + "N": 8192, + "K": 262144, + "speedup": 0.887, + "tflops": 189.1, + "n_trials": 5, + "config": { + "block_size_m": 128, + "block_size_n": 256, + "block_size_k": 64, + "group_size_m": 8, + "num_warps": 8, + "num_stages": 2, + "num_xcds": 8, + "allow_tf32": true + }, + "hbm_buffer_params": { + "k_per_flag": 32, + "num_fetch_sms": 4, + "num_fetch_stages": 8, + "first_stage_fetch_sms": 52 + } + }, + "16384x16384x131072": { + "label": "g1", + "description": "ws=2 measured — 0.772x vs PyTorch", + "M": 16384, + "N": 16384, + "K": 131072, + "speedup": 0.772, + "tflops": 459.8, + "n_trials": 10, + "config": { + "block_size_m": 128, + "block_size_n": 256, + "block_size_k": 64, + "group_size_m": 16, + "num_warps": 8, + "num_stages": 2, + "num_xcds": 8, + "allow_tf32": true + }, + "hbm_buffer_params": { + "k_per_flag": 16, + "num_fetch_sms": 8, + "num_fetch_stages": 8, + "first_stage_fetch_sms": 52 + } + }, + "4096x14336x4096": { + "label": "mixtral_gate", + "description": "Mixtral gate projection — ws=2", + "M": 4096, + "N": 14336, + "K": 4096, + "speedup": null, + "tflops": 206.1, + "n_trials": 20, + "config": { + "block_size_m": 128, + "block_size_n": 256, + "block_size_k": 64, + "group_size_m": 8, + "num_warps": 8, + "num_stages": 2, + "num_xcds": 8, + "allow_tf32": true + }, + "hbm_buffer_params": { + "k_per_flag": 8, + "num_fetch_sms": 8, + "num_fetch_stages": 1, + "first_stage_fetch_sms": 52 + } + }, + "4096x11008x4096": { + "label": "llama7b_gate", + "description": "Llama-7B gate projection — ws=2", + "M": 4096, + "N": 11008, + "K": 4096, + "speedup": null, + "tflops": 207.1, + "n_trials": 20, + "config": { + "block_size_m": 128, + "block_size_n": 256, + "block_size_k": 64, + "group_size_m": 8, + "num_warps": 8, + "num_stages": 2, + "num_xcds": 8, + "allow_tf32": true + }, + "hbm_buffer_params": { + "k_per_flag": 8, + "num_fetch_sms": 8, + "num_fetch_stages": 1, + "first_stage_fetch_sms": 52 + } + }, + "4096x4096x4096": { + "label": "pow2_4k", + "description": "Small power-of-2 square — ws=2", + "M": 4096, + "N": 4096, + "K": 4096, + "speedup": null, + "tflops": 123.9, + "n_trials": 20, + "config": { + "block_size_m": 128, + "block_size_n": 256, + "block_size_k": 64, + "group_size_m": 8, + "num_warps": 8, + "num_stages": 2, + "num_xcds": 8, + "allow_tf32": true + }, + "hbm_buffer_params": { + "k_per_flag": 8, + "num_fetch_sms": 8, + "num_fetch_stages": 1, + "first_stage_fetch_sms": 52 + } + }, + "5120x13824x5120": { + "label": "llama13b_gate", + "description": "Llama-13B gate projection — ws=2", + "M": 5120, + "N": 13824, + "K": 5120, + "speedup": null, + "tflops": 174.9, + "n_trials": 20, + "config": { + "block_size_m": 128, + "block_size_n": 256, + "block_size_k": 64, + "group_size_m": 8, + "num_warps": 8, + "num_stages": 2, + "num_xcds": 8, + "allow_tf32": true + }, + "hbm_buffer_params": { + "k_per_flag": 8, + "num_fetch_sms": 8, + "num_fetch_stages": 1, + "first_stage_fetch_sms": 52 + } + }, + "4096x4096x11008": { + "label": "llama7b_down", + "description": "Llama-7B down projection — ws=2", + "M": 4096, + "N": 4096, + "K": 11008, + "speedup": null, + "tflops": 149.1, + "n_trials": 20, + "config": { + "block_size_m": 128, + "block_size_n": 256, + "block_size_k": 64, + "group_size_m": 8, + "num_warps": 8, + "num_stages": 2, + "num_xcds": 8, + "allow_tf32": true + }, + "hbm_buffer_params": { + "k_per_flag": 8, + "num_fetch_sms": 8, + "num_fetch_stages": 1, + "first_stage_fetch_sms": 52 + } + } + }, + "default_config": null +} \ No newline at end of file diff --git a/iris/ops/configs/ag_mm/mi300x/NN/ws4.json b/iris/ops/configs/ag_mm/mi300x/NN/ws4.json new file mode 100644 index 000000000..5d3efc703 --- /dev/null +++ b/iris/ops/configs/ag_mm/mi300x/NN/ws4.json @@ -0,0 +1,249 @@ +{ + "_meta": { + "description": "AG+MM ws=4 on MI300X — DISABLED (loses vs PyTorch on all shapes)", + "source": "benchmarking: best speedup 0.856x (g6). LDS overflow forces ns=1 on K=4096 shapes.", + "gpu": "AMD Instinct MI300X (gfx942)", + "date": "2026-04-13", + "data_tag": "[VERIFIED]" + }, + "enabled": false, + "reason": "ws=4 loses on all 7 tested shapes (best 0.856x). K=4096 shapes crash at ns=2 due to LDS overflow (65540>65536). ns=1 workaround constrains pipelining depth below break-even.", + "best_measured_speedup": 0.856, + "shapes": { + "262144x8192x8192": { + "label": "g6", + "description": "Best ws=4 result — still loses (0.856x)", + "M": 262144, + "N": 8192, + "K": 8192, + "speedup": 0.856, + "tflops": 218.5, + "n_trials": 10, + "config": { + "block_size_m": 256, + "block_size_n": 256, + "block_size_k": 64, + "group_size_m": 24, + "num_warps": 8, + "num_stages": 2, + "num_xcds": 8, + "allow_tf32": true + }, + "hbm_buffer_params": { + "k_per_flag": 64, + "num_fetch_sms": 52, + "num_fetch_stages": 4, + "first_stage_fetch_sms": 52 + } + }, + "8192x8192x262144": { + "label": "g5", + "description": "ws=4 measured — 0.848x", + "M": 8192, + "N": 8192, + "K": 262144, + "speedup": 0.848, + "tflops": 209.5, + "n_trials": 8, + "config": { + "block_size_m": 128, + "block_size_n": 256, + "block_size_k": 64, + "group_size_m": 8, + "num_warps": 8, + "num_stages": 2, + "num_xcds": 8, + "allow_tf32": true + }, + "hbm_buffer_params": { + "k_per_flag": 32, + "num_fetch_sms": 4, + "num_fetch_stages": 8, + "first_stage_fetch_sms": 52 + } + }, + "131072x16384x16384": { + "label": "g2", + "description": "ws=4 measured — 0.727x", + "M": 131072, + "N": 16384, + "K": 16384, + "speedup": 0.727, + "tflops": 335.8, + "n_trials": 3, + "config": { + "block_size_m": 256, + "block_size_n": 256, + "block_size_k": 64, + "group_size_m": 24, + "num_warps": 8, + "num_stages": 2, + "num_xcds": 8, + "allow_tf32": true + }, + "hbm_buffer_params": { + "k_per_flag": 32, + "num_fetch_sms": 4, + "num_fetch_stages": 64, + "first_stage_fetch_sms": 52 + } + }, + "16384x16384x131072": { + "label": "g1", + "description": "ws=4 measured — 0.650x", + "M": 16384, + "N": 16384, + "K": 131072, + "speedup": 0.65, + "tflops": 293.5, + "n_trials": 5, + "config": { + "block_size_m": 128, + "block_size_n": 256, + "block_size_k": 64, + "group_size_m": 16, + "num_warps": 8, + "num_stages": 2, + "num_xcds": 8, + "allow_tf32": true + }, + "hbm_buffer_params": { + "k_per_flag": 16, + "num_fetch_sms": 16, + "num_fetch_stages": 8, + "first_stage_fetch_sms": 52 + } + }, + "4096x14336x4096": { + "label": "mixtral_gate", + "description": "Mixtral gate projection — ws=4", + "M": 4096, + "N": 14336, + "K": 4096, + "speedup": null, + "tflops": 219.5, + "n_trials": 20, + "config": { + "block_size_m": 128, + "block_size_n": 256, + "block_size_k": 64, + "group_size_m": 8, + "num_warps": 8, + "num_stages": 2, + "num_xcds": 8, + "allow_tf32": true + }, + "hbm_buffer_params": { + "k_per_flag": 8, + "num_fetch_sms": 2, + "num_fetch_stages": 1, + "first_stage_fetch_sms": 16 + } + }, + "4096x11008x4096": { + "label": "llama7b_gate", + "description": "Llama-7B gate projection — ws=4", + "M": 4096, + "N": 11008, + "K": 4096, + "speedup": null, + "tflops": 170.8, + "n_trials": 20, + "config": { + "block_size_m": 128, + "block_size_n": 256, + "block_size_k": 64, + "group_size_m": 8, + "num_warps": 8, + "num_stages": 2, + "num_xcds": 8, + "allow_tf32": true + }, + "hbm_buffer_params": { + "k_per_flag": 8, + "num_fetch_sms": 2, + "num_fetch_stages": 1, + "first_stage_fetch_sms": 16 + } + }, + "4096x4096x4096": { + "label": "pow2_4k", + "description": "Small power-of-2 square — ws=4", + "M": 4096, + "N": 4096, + "K": 4096, + "speedup": 0.347, + "tflops": 92.0, + "n_trials": 20, + "config": { + "block_size_m": 128, + "block_size_n": 256, + "block_size_k": 64, + "group_size_m": 8, + "num_warps": 8, + "num_stages": 2, + "num_xcds": 8, + "allow_tf32": true + }, + "hbm_buffer_params": { + "k_per_flag": 8, + "num_fetch_sms": 2, + "num_fetch_stages": 1, + "first_stage_fetch_sms": 16 + } + }, + "5120x13824x5120": { + "label": "llama13b_gate", + "description": "Llama-13B gate projection — ws=4", + "M": 5120, + "N": 13824, + "K": 5120, + "speedup": null, + "tflops": 222.5, + "n_trials": 20, + "config": { + "block_size_m": 128, + "block_size_n": 256, + "block_size_k": 64, + "group_size_m": 8, + "num_warps": 8, + "num_stages": 2, + "num_xcds": 8, + "allow_tf32": true + }, + "hbm_buffer_params": { + "k_per_flag": 8, + "num_fetch_sms": 2, + "num_fetch_stages": 1, + "first_stage_fetch_sms": 16 + } + }, + "4096x4096x11008": { + "label": "llama7b_down", + "description": "Llama-7B down projection — ws=4", + "M": 4096, + "N": 4096, + "K": 11008, + "speedup": null, + "tflops": 143.3, + "n_trials": 20, + "config": { + "block_size_m": 128, + "block_size_n": 256, + "block_size_k": 64, + "group_size_m": 8, + "num_warps": 8, + "num_stages": 2, + "num_xcds": 8, + "allow_tf32": true + }, + "hbm_buffer_params": { + "k_per_flag": 8, + "num_fetch_sms": 2, + "num_fetch_stages": 1, + "first_stage_fetch_sms": 16 + } + } + }, + "default_config": null +} \ No newline at end of file diff --git a/iris/ops/configs/ag_mm/mi300x/NN/ws8.json b/iris/ops/configs/ag_mm/mi300x/NN/ws8.json new file mode 100644 index 000000000..f61531148 --- /dev/null +++ b/iris/ops/configs/ag_mm/mi300x/NN/ws8.json @@ -0,0 +1,348 @@ +{ + "_meta": { + "description": "Champion configs for HBM buffer AG+MM ws=8 on MI300X (gfx942)", + "source": "sweep (3489 trials), optimize-loop iter3", + "gpu": "AMD Instinct MI300X (gfx942)", + "date": "2026-04-13", + "data_tag": "[VERIFIED]", + "convention": "Shapes are (M, N, K) for col-parallel (M-sharded) AG+MM" + }, + "enabled": true, + "shapes": { + "262144x8192x8192": { + "label": "g6", + "description": "Llama-70B MLP hidden×hidden — M-dominant", + "M": 262144, + "N": 8192, + "K": 8192, + "speedup": 1.2, + "tflops": 253.0, + "n_trials": 27, + "config": { + "block_size_m": 256, + "block_size_n": 256, + "block_size_k": 64, + "group_size_m": 24, + "num_warps": 8, + "num_stages": 2, + "num_xcds": 8, + "allow_tf32": true + }, + "hbm_buffer_params": { + "k_per_flag": 64, + "num_fetch_sms": 52, + "num_fetch_stages": 4, + "first_stage_fetch_sms": 52 + } + }, + "131072x16384x16384": { + "label": "g2", + "description": "Llama MLP variant — balanced large", + "M": 131072, + "N": 16384, + "K": 16384, + "speedup": 1.343, + "tflops": 420.5, + "n_trials": 102, + "config": { + "block_size_m": 256, + "block_size_n": 256, + "block_size_k": 64, + "group_size_m": 24, + "num_warps": 8, + "num_stages": 2, + "num_xcds": 8, + "allow_tf32": true + }, + "hbm_buffer_params": { + "k_per_flag": 32, + "num_fetch_sms": 4, + "num_fetch_stages": 64, + "first_stage_fetch_sms": 52 + } + }, + "147456x28672x4096": { + "label": "g14", + "description": "Llama-70B up-projection medium batch", + "M": 147456, + "N": 28672, + "K": 4096, + "speedup": 1.288, + "tflops": 466.5, + "n_trials": 108, + "config": { + "block_size_m": 256, + "block_size_n": 256, + "block_size_k": 64, + "group_size_m": 24, + "num_warps": 8, + "num_stages": 2, + "num_xcds": 8, + "allow_tf32": true + }, + "hbm_buffer_params": { + "k_per_flag": 16, + "num_fetch_sms": 59, + "num_fetch_stages": 36, + "first_stage_fetch_sms": 52 + } + }, + "229376x28672x4096": { + "label": "g16", + "description": "Llama-70B up-projection mid batch", + "M": 229376, + "N": 28672, + "K": 4096, + "speedup": 1.277, + "tflops": 471.5, + "n_trials": 124, + "config": { + "block_size_m": 256, + "block_size_n": 256, + "block_size_k": 64, + "group_size_m": 24, + "num_warps": 8, + "num_stages": 2, + "num_xcds": 8, + "allow_tf32": true + }, + "hbm_buffer_params": { + "k_per_flag": 16, + "num_fetch_sms": 4, + "num_fetch_stages": 56, + "first_stage_fetch_sms": 52 + } + }, + "327680x28672x4096": { + "label": "g15", + "description": "Llama-70B up-projection large batch — highest TFLOPS", + "M": 327680, + "N": 28672, + "K": 4096, + "speedup": 1.284, + "tflops": 474.7, + "n_trials": 70, + "config": { + "block_size_m": 256, + "block_size_n": 256, + "block_size_k": 64, + "group_size_m": 24, + "num_warps": 8, + "num_stages": 2, + "num_xcds": 8, + "allow_tf32": true + }, + "hbm_buffer_params": { + "k_per_flag": 16, + "num_fetch_sms": 4, + "num_fetch_stages": 32, + "first_stage_fetch_sms": 52 + } + }, + "8192x8192x262144": { + "label": "g5", + "description": "K-dominant square — best speedup shape", + "M": 8192, + "N": 8192, + "K": 262144, + "speedup": 1.224, + "tflops": 161.6, + "n_trials": 9, + "config": { + "block_size_m": 128, + "block_size_n": 256, + "block_size_k": 64, + "group_size_m": 8, + "num_warps": 8, + "num_stages": 2, + "num_xcds": 8, + "allow_tf32": true + }, + "hbm_buffer_params": { + "k_per_flag": 32, + "num_fetch_sms": 4, + "num_fetch_stages": 8, + "first_stage_fetch_sms": 52 + } + }, + "16384x16384x131072": { + "label": "g1", + "description": "K-dominant large — parity shape", + "M": 16384, + "N": 16384, + "K": 131072, + "speedup": 1.136, + "tflops": 314.5, + "n_trials": 101, + "config": { + "block_size_m": 128, + "block_size_n": 256, + "block_size_k": 64, + "group_size_m": 16, + "num_warps": 8, + "num_stages": 2, + "num_xcds": 8, + "allow_tf32": true + }, + "hbm_buffer_params": { + "k_per_flag": 16, + "num_fetch_sms": 16, + "num_fetch_stages": 8, + "first_stage_fetch_sms": 52 + } + }, + "196608x18432x16384": { + "label": "g9", + "description": "Large balanced shape", + "M": 196608, + "N": 18432, + "K": 16384, + "speedup": 0.854, + "tflops": 445.4, + "n_trials": 80, + "config": { + "block_size_m": 256, + "block_size_n": 256, + "block_size_k": 64, + "group_size_m": 1, + "num_warps": 8, + "num_stages": 2, + "num_xcds": 8, + "allow_tf32": true + }, + "hbm_buffer_params": { + "k_per_flag": 32, + "num_fetch_sms": 32, + "num_fetch_stages": 1, + "first_stage_fetch_sms": 32 + } + }, + "262144x28672x8192": { + "label": "g8", + "description": "Large wide shape", + "M": 262144, + "N": 28672, + "K": 8192, + "speedup": 0.857, + "tflops": 442.1, + "n_trials": 72, + "config": { + "block_size_m": 256, + "block_size_n": 256, + "block_size_k": 64, + "group_size_m": 1, + "num_warps": 8, + "num_stages": 2, + "num_xcds": 8, + "allow_tf32": true + }, + "hbm_buffer_params": { + "k_per_flag": 128, + "num_fetch_sms": 32, + "num_fetch_stages": 1, + "first_stage_fetch_sms": 32 + } + }, + "4096x14336x4096": { + "label": "mixtral_gate", + "description": "Mixtral gate projection", + "M": 4096, + "N": 14336, + "K": 4096, + "speedup": 0.703, + "tflops": 248.9, + "n_trials": 20, + "config": { + "block_size_m": 128, + "block_size_n": 256, + "block_size_k": 64, + "group_size_m": 1, + "num_warps": 8, + "num_stages": 2, + "num_xcds": 8, + "allow_tf32": true + }, + "hbm_buffer_params": { + "k_per_flag": 8, + "num_fetch_sms": 16, + "num_fetch_stages": 1, + "first_stage_fetch_sms": 64 + } + }, + "4096x11008x4096": { + "label": "llama7b_gate", + "description": "Llama-7B gate projection", + "M": 4096, + "N": 11008, + "K": 4096, + "speedup": 0.551, + "tflops": 189.8, + "n_trials": 20, + "config": { + "block_size_m": 128, + "block_size_n": 128, + "block_size_k": 64, + "group_size_m": 1, + "num_warps": 8, + "num_stages": 2, + "num_xcds": 8, + "allow_tf32": true + }, + "hbm_buffer_params": { + "k_per_flag": 8, + "num_fetch_sms": 16, + "num_fetch_stages": 1, + "first_stage_fetch_sms": 64 + } + }, + "4096x4096x4096": { + "label": "pow2_4k", + "description": "Small power-of-2 square shape", + "M": 4096, + "N": 4096, + "K": 4096, + "speedup": 0.527, + "tflops": 90.9, + "n_trials": 20, + "config": { + "block_size_m": 128, + "block_size_n": 256, + "block_size_k": 64, + "group_size_m": 1, + "num_warps": 8, + "num_stages": 2, + "num_xcds": 8, + "allow_tf32": true + }, + "hbm_buffer_params": { + "k_per_flag": 8, + "num_fetch_sms": 8, + "num_fetch_stages": 1, + "first_stage_fetch_sms": 64 + } + } + }, + "default_config": { + "block_size_m": 256, + "block_size_n": 256, + "block_size_k": 64, + "group_size_m": 24, + "num_warps": 8, + "num_stages": 2, + "num_xcds": 8, + "allow_tf32": true + }, + "default_hbm_buffer_params": { + "k_per_flag": 8, + "num_fetch_sms": 32, + "num_fetch_stages": 1, + "first_stage_fetch_sms": 64 + }, + "heuristic_rules": { + "block_size_m": "Use 128 when M <= 16384, else 256", + "group_size_m": "Use 8 when M <= 8192, 16 when M <= 16384, else 24", + "k_per_flag": "Maximize: 64 for K_blocks>=512, 16 for >=128, 8 for >=64, else 4. Must divide K_blocks evenly.", + "num_fetch_sms": "Scale with M-tiles: 4 for <=8, 16 for <=32, 32 for <=128, else 52", + "num_fetch_stages": "4 for M_tiles>=512, 2 for >=64, else 1" + } +} \ No newline at end of file diff --git a/iris/ops/configs/ag_mm/mi300x/NT/ws2.json b/iris/ops/configs/ag_mm/mi300x/NT/ws2.json new file mode 100644 index 000000000..897be3f2c --- /dev/null +++ b/iris/ops/configs/ag_mm/mi300x/NT/ws2.json @@ -0,0 +1,10 @@ +{ + "_meta": { + "description": "AG+MM ws=2 NT transpose on MI300X — DISABLED", + "source": "ws<8 is NO-GO across all transposes", + "gpu": "AMD Instinct MI300X (gfx942)", + "date": "2026-04-13" + }, + "enabled": false, + "reason": "ws=2 loses vs PyTorch on all tested shapes. LDS overflow forces ns=1, imposing 15-35% perf penalty." +} diff --git a/iris/ops/configs/ag_mm/mi300x/NT/ws4.json b/iris/ops/configs/ag_mm/mi300x/NT/ws4.json new file mode 100644 index 000000000..cc1f9d297 --- /dev/null +++ b/iris/ops/configs/ag_mm/mi300x/NT/ws4.json @@ -0,0 +1,10 @@ +{ + "_meta": { + "description": "AG+MM ws=4 NT transpose on MI300X — DISABLED", + "source": "ws<8 is NO-GO across all transposes", + "gpu": "AMD Instinct MI300X (gfx942)", + "date": "2026-04-13" + }, + "enabled": false, + "reason": "ws=4 loses vs PyTorch on all tested shapes. Best measured: 0.856x. LDS overflow at K=4096." +} diff --git a/iris/ops/configs/ag_mm/mi300x/NT/ws8.json b/iris/ops/configs/ag_mm/mi300x/NT/ws8.json new file mode 100644 index 000000000..873cb76e1 --- /dev/null +++ b/iris/ops/configs/ag_mm/mi300x/NT/ws8.json @@ -0,0 +1,31 @@ +{ + "_meta": { + "description": "AG+MM ws=8 NT transpose on MI300X — heuristic defaults (no per-shape benchmarks yet)", + "source": "heuristic extrapolation from NN transpose champion data", + "gpu": "AMD Instinct MI300X (gfx942)", + "date": "2026-04-13", + "data_tag": "heuristic", + "convention": "Shapes are (M, N, K) for col-parallel (M-sharded) AG+MM, B transposed (K×N → N×K)" + }, + "enabled": true, + "shapes": {}, + "default_config": { + "block_size_m": 256, + "block_size_n": 256, + "block_size_k": 64, + "group_size_m": 24, + "num_warps": 8, + "num_stages": 2, + "num_xcds": 8, + "allow_tf32": true + }, + "default_hbm_buffer_params": { + "k_per_flag": 8, + "num_fetch_sms": 32, + "num_fetch_stages": 1, + "first_stage_fetch_sms": 64 + }, + "heuristic_rules": { + "note": "Uses same heuristic as NN transpose. Shape-specific tuning pending." + } +} diff --git a/iris/ops/configs/ag_mm/mi300x/TN/ws2.json b/iris/ops/configs/ag_mm/mi300x/TN/ws2.json new file mode 100644 index 000000000..2fe67e154 --- /dev/null +++ b/iris/ops/configs/ag_mm/mi300x/TN/ws2.json @@ -0,0 +1,10 @@ +{ + "_meta": { + "description": "AG+MM ws=2 TN transpose on MI300X — DISABLED", + "source": "ws<8 is NO-GO across all transposes", + "gpu": "AMD Instinct MI300X (gfx942)", + "date": "2026-04-13" + }, + "enabled": false, + "reason": "ws=2 loses vs PyTorch on all tested shapes. LDS overflow forces ns=1, imposing 15-35% perf penalty." +} diff --git a/iris/ops/configs/ag_mm/mi300x/TN/ws4.json b/iris/ops/configs/ag_mm/mi300x/TN/ws4.json new file mode 100644 index 000000000..c8977d5f0 --- /dev/null +++ b/iris/ops/configs/ag_mm/mi300x/TN/ws4.json @@ -0,0 +1,10 @@ +{ + "_meta": { + "description": "AG+MM ws=4 TN transpose on MI300X — DISABLED", + "source": "ws<8 is NO-GO across all transposes", + "gpu": "AMD Instinct MI300X (gfx942)", + "date": "2026-04-13" + }, + "enabled": false, + "reason": "ws=4 loses vs PyTorch on all tested shapes. Best measured: 0.856x. LDS overflow at K=4096." +} diff --git a/iris/ops/configs/ag_mm/mi300x/TN/ws8.json b/iris/ops/configs/ag_mm/mi300x/TN/ws8.json new file mode 100644 index 000000000..df9a5b3f9 --- /dev/null +++ b/iris/ops/configs/ag_mm/mi300x/TN/ws8.json @@ -0,0 +1,31 @@ +{ + "_meta": { + "description": "AG+MM ws=8 TN transpose on MI300X — heuristic defaults (no per-shape benchmarks yet)", + "source": "heuristic extrapolation from NN transpose champion data", + "gpu": "AMD Instinct MI300X (gfx942)", + "date": "2026-04-13", + "data_tag": "heuristic", + "convention": "Shapes are (M, N, K) for col-parallel (M-sharded) AG+MM, A transposed (M×K → K×M)" + }, + "enabled": true, + "shapes": {}, + "default_config": { + "block_size_m": 256, + "block_size_n": 256, + "block_size_k": 64, + "group_size_m": 24, + "num_warps": 8, + "num_stages": 2, + "num_xcds": 8, + "allow_tf32": true + }, + "default_hbm_buffer_params": { + "k_per_flag": 8, + "num_fetch_sms": 32, + "num_fetch_stages": 1, + "first_stage_fetch_sms": 64 + }, + "heuristic_rules": { + "note": "Uses same heuristic as NN transpose. Shape-specific tuning pending." + } +} diff --git a/iris/ops/configs/ag_mm/mi300x/TT/ws2.json b/iris/ops/configs/ag_mm/mi300x/TT/ws2.json new file mode 100644 index 000000000..cc2c2497c --- /dev/null +++ b/iris/ops/configs/ag_mm/mi300x/TT/ws2.json @@ -0,0 +1,10 @@ +{ + "_meta": { + "description": "AG+MM ws=2 TT transpose on MI300X — DISABLED", + "source": "ws<8 is NO-GO across all transposes", + "gpu": "AMD Instinct MI300X (gfx942)", + "date": "2026-04-13" + }, + "enabled": false, + "reason": "ws=2 loses vs PyTorch on all tested shapes. LDS overflow forces ns=1, imposing 15-35% perf penalty." +} diff --git a/iris/ops/configs/ag_mm/mi300x/TT/ws4.json b/iris/ops/configs/ag_mm/mi300x/TT/ws4.json new file mode 100644 index 000000000..55ee5f423 --- /dev/null +++ b/iris/ops/configs/ag_mm/mi300x/TT/ws4.json @@ -0,0 +1,10 @@ +{ + "_meta": { + "description": "AG+MM ws=4 TT transpose on MI300X — DISABLED", + "source": "ws<8 is NO-GO across all transposes", + "gpu": "AMD Instinct MI300X (gfx942)", + "date": "2026-04-13" + }, + "enabled": false, + "reason": "ws=4 loses vs PyTorch on all tested shapes. Best measured: 0.856x. LDS overflow at K=4096." +} diff --git a/iris/ops/configs/ag_mm/mi300x/TT/ws8.json b/iris/ops/configs/ag_mm/mi300x/TT/ws8.json new file mode 100644 index 000000000..a184b41a4 --- /dev/null +++ b/iris/ops/configs/ag_mm/mi300x/TT/ws8.json @@ -0,0 +1,31 @@ +{ + "_meta": { + "description": "AG+MM ws=8 TT transpose on MI300X — heuristic defaults (no per-shape benchmarks yet)", + "source": "heuristic extrapolation from NN transpose champion data", + "gpu": "AMD Instinct MI300X (gfx942)", + "date": "2026-04-13", + "data_tag": "heuristic", + "convention": "Shapes are (M, N, K) for col-parallel (M-sharded) AG+MM, both A and B transposed" + }, + "enabled": true, + "shapes": {}, + "default_config": { + "block_size_m": 256, + "block_size_n": 256, + "block_size_k": 64, + "group_size_m": 24, + "num_warps": 8, + "num_stages": 2, + "num_xcds": 8, + "allow_tf32": true + }, + "default_hbm_buffer_params": { + "k_per_flag": 8, + "num_fetch_sms": 32, + "num_fetch_stages": 1, + "first_stage_fetch_sms": 64 + }, + "heuristic_rules": { + "note": "Uses same heuristic as NN transpose. Shape-specific tuning pending." + } +} diff --git a/iris/ops/configs/ag_mm/regression_sizes.json b/iris/ops/configs/ag_mm/regression_sizes.json new file mode 100644 index 000000000..40a497b0a --- /dev/null +++ b/iris/ops/configs/ag_mm/regression_sizes.json @@ -0,0 +1,101 @@ +{ + "_meta": { + "description": "Regression test sizes for HBM buffer AG+MM kernel across ws=2/4/8", + "source": "28 measured shapes (12 ws=8, 9 ws=4, 7 ws=2) from 3489 trials", + "gpu_target": "MI300X (gfx942)", + "date": "2026-04-13", + "usage": "from iris.ops import load_regression_sizes" + }, + "sizes": [ + { + "name": "g2_ws8", + "label": "g2", + "description": "Llama MLP variant — balanced large (highest speedup)", + "M": 131072, "N": 16384, "K": 16384, + "tier": "champion", + "world_sizes": [8], + "expected": {"ws8_speedup": 1.343, "ws8_tflops": 420.5}, + "regression_threshold_pct": 10 + }, + { + "name": "g15_ws8", + "label": "g15", + "description": "Llama-70B up-projection large batch — highest TFLOPS", + "M": 327680, "N": 28672, "K": 4096, + "tier": "champion", + "world_sizes": [8], + "expected": {"ws8_speedup": 1.284, "ws8_tflops": 474.7}, + "regression_threshold_pct": 10 + }, + { + "name": "g14_ws8", + "label": "g14", + "description": "Llama-70B up-projection medium batch", + "M": 147456, "N": 28672, "K": 4096, + "tier": "champion", + "world_sizes": [8], + "expected": {"ws8_speedup": 1.288, "ws8_tflops": 466.5}, + "regression_threshold_pct": 10 + }, + { + "name": "g16_ws8", + "label": "g16", + "description": "Llama-70B up-projection mid batch", + "M": 229376, "N": 28672, "K": 4096, + "tier": "champion", + "world_sizes": [8], + "expected": {"ws8_speedup": 1.277, "ws8_tflops": 471.5}, + "regression_threshold_pct": 10 + }, + { + "name": "g5_ws8", + "label": "g5", + "description": "K-dominant square — M-small, needs bm=128", + "M": 8192, "N": 8192, "K": 262144, + "tier": "champion", + "world_sizes": [8], + "expected": {"ws8_speedup": 1.224, "ws8_tflops": 161.6}, + "regression_threshold_pct": 10 + }, + { + "name": "g6_ws8", + "label": "g6", + "description": "Llama-70B MLP hidden x hidden — M-dominant", + "M": 262144, "N": 8192, "K": 8192, + "tier": "champion", + "world_sizes": [8], + "expected": {"ws8_speedup": 1.200, "ws8_tflops": 253.0}, + "regression_threshold_pct": 10 + }, + { + "name": "g1_ws8", + "label": "g1", + "description": "K-dominant large — parity shape", + "M": 16384, "N": 16384, "K": 131072, + "tier": "champion", + "world_sizes": [8], + "expected": {"ws8_speedup": 1.136, "ws8_tflops": 314.5}, + "regression_threshold_pct": 10 + }, + { + "name": "g5_ws2_disabled", + "label": "g5", + "description": "Best ws=2 shape — still loses vs PyTorch (0.887x). Verifies fallback.", + "M": 8192, "N": 8192, "K": 262144, + "tier": "disabled", + "world_sizes": [2], + "expected": {"ws2_speedup": 0.887, "ws2_disabled": true}, + "regression_threshold_pct": null + }, + { + "name": "g6_ws4_disabled", + "label": "g6", + "description": "Best ws=4 shape — still loses vs PyTorch (0.856x). Verifies fallback.", + "M": 262144, "N": 8192, "K": 8192, + "tier": "disabled", + "world_sizes": [4], + "expected": {"ws4_speedup": 0.856, "ws4_disabled": true}, + "regression_threshold_pct": null + } + ] +} diff --git a/pyproject.toml b/pyproject.toml index 4a8f1916c..ccd54f8dc 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -50,6 +50,9 @@ package-dir = { "" = "." } [tool.setuptools.packages.find] include = ["iris*"] +[tool.setuptools.package-data] +"iris.ops.configs" = ["**/*.json"] + # ---- setuptools-scm versioning ---- [tool.setuptools_scm] version_scheme = "post-release" # .postN after last tag diff --git a/tests/ops/test_auto_config.py b/tests/ops/test_auto_config.py new file mode 100644 index 000000000..f2a6e20cb --- /dev/null +++ b/tests/ops/test_auto_config.py @@ -0,0 +1,703 @@ +#!/usr/bin/env python3 +# SPDX-License-Identifier: MIT +# Copyright (c) 2025 Advanced Micro Devices, Inc. All rights reserved. + +""" +Unit tests for iris.ops.auto_config — AG+MM auto-selection mechanism. + +These tests run on CPU (no GPU required) and verify: +1. Exact shape lookup hits return correct champion configs +2. Lookup misses fall back to heuristic defaults +3. ws<8 configs are correctly disabled +4. FusedConfig conversion works for enabled configs +5. FusedConfig conversion raises for disabled configs +6. list_known_shapes returns correct data +7. Config cache can be cleared +""" + +import pytest +import sys +import os + +# Ensure iris package is importable even without full install +sys.path.insert(0, os.path.join(os.path.dirname(__file__), "..", "..")) + +from iris.ops.auto_config import ( + select_ag_mm_config, + list_known_shapes, + load_regression_sizes, + clear_config_cache, + detect_gpu_arch, + SUPPORTED_TRANSPOSES, + SUPPORTED_ARCHITECTURES, + _apply_heuristic, +) +import iris.ops.auto_config as auto_config_module +from iris.ops.config import FusedConfig + + +class TestAutoConfigExactMatch: + """Test exact shape lookup (cache hit path).""" + + def setup_method(self): + clear_config_cache() + + def test_ws8_g2_exact_match(self): + """g2 shape (131072x16384x16384) should return champion config.""" + result = select_ag_mm_config(M=131072, N=16384, K=16384, world_size=8, transpose="NN", arch="mi300x") + assert result.enabled is True + assert result.shape_key == "131072x16384x16384" + assert result.speedup == 1.343 + assert "Exact match" in result.source + assert result.config_params["block_size_m"] == 256 + assert result.config_params["block_size_n"] == 256 + assert result.config_params["block_size_k"] == 64 + assert result.config_params["num_stages"] == 2 + + def test_ws8_g5_exact_match(self): + """g5 shape (8192x8192x262144) should use bm=128 (small M).""" + result = select_ag_mm_config(M=8192, N=8192, K=262144, world_size=8, transpose="NN", arch="mi300x") + assert result.enabled is True + assert result.shape_key == "8192x8192x262144" + assert result.speedup == 1.224 + assert result.config_params["block_size_m"] == 128 + assert result.config_params["group_size_m"] == 8 + + def test_ws8_g1_exact_match(self): + """g1 shape (16384x16384x131072) should use bm=128.""" + result = select_ag_mm_config(M=16384, N=16384, K=131072, world_size=8, transpose="NN", arch="mi300x") + assert result.enabled is True + assert result.shape_key == "16384x16384x131072" + assert result.speedup == 1.136 + assert result.config_params["block_size_m"] == 128 + + def test_ws8_winning_shapes_enabled(self): + """All 7 winning champion shapes (speedup > 1.0) should be enabled.""" + winning_shapes = [ + (131072, 16384, 16384), # g2 — 1.343x + (327680, 28672, 4096), # g15 — 1.284x + (147456, 28672, 4096), # g14 — 1.288x + (229376, 28672, 4096), # g16 — 1.277x + (8192, 8192, 262144), # g5 — 1.224x + (262144, 8192, 8192), # g6 — 1.200x + (16384, 16384, 131072), # g1 — 1.136x + ] + for M, N, K in winning_shapes: + result = select_ag_mm_config(M, N, K, world_size=8) + assert result.enabled, f"Shape {M}x{N}x{K} should be enabled" + assert result.speedup is not None and result.speedup > 1.0, ( + f"Shape {M}x{N}x{K} should have speedup > 1.0, got {result.speedup}" + ) + + +class TestAutoConfigFallback: + """Test lookup miss / heuristic fallback path.""" + + def setup_method(self): + clear_config_cache() + + def test_ws8_unknown_shape_returns_heuristic(self): + """An unknown shape at ws=8 should still be enabled with heuristic config.""" + result = select_ag_mm_config(M=65536, N=4096, K=8192, world_size=8, transpose="NN", arch="mi300x") + assert result.enabled is True + assert result.shape_key is None + assert result.speedup is None + assert "Heuristic" in result.source + # Heuristic: M=65536 > 16384 -> bm=256 + assert result.config_params["block_size_m"] == 256 + + def test_ws8_small_M_exact_or_heuristic(self): + """Small M (pow2_4k) hits exact match with bm=128.""" + result = select_ag_mm_config(M=4096, N=4096, K=4096, world_size=8, transpose="NN", arch="mi300x") + assert result.enabled is True + assert result.config_params["block_size_m"] == 128 + + +class TestAutoConfigDisabled: + """Test that ws<8 correctly disables iris AG+MM.""" + + def setup_method(self): + clear_config_cache() + + def test_ws2_disabled(self): + """ws=2 should be disabled on MI300X.""" + result = select_ag_mm_config(M=131072, N=16384, K=16384, world_size=2, transpose="NN", arch="mi300x") + assert result.enabled is False + assert "Disabled" in result.source or "world_size" in result.source + + def test_ws4_disabled(self): + """ws=4 should be disabled on MI300X.""" + result = select_ag_mm_config(M=131072, N=16384, K=16384, world_size=4, transpose="NN", arch="mi300x") + assert result.enabled is False + + def test_ws1_disabled(self): + """ws=1 should be disabled (no config file, below min_world_size).""" + result = select_ag_mm_config(M=4096, N=4096, K=4096, world_size=1, transpose="NN", arch="mi300x") + assert result.enabled is False + + def test_ws3_disabled_by_default_gate(self): + """ws=3 has no config file, should be disabled by global default min_world_size=8.""" + result = select_ag_mm_config(M=4096, N=4096, K=4096, world_size=3, transpose="NN", arch="mi300x") + assert result.enabled is False + + +class TestFusedConfigConversion: + """Test AutoConfigResult.to_fused_config().""" + + def setup_method(self): + clear_config_cache() + + def test_enabled_config_converts(self): + """Enabled configs should produce valid FusedConfig.""" + result = select_ag_mm_config(M=131072, N=16384, K=16384, world_size=8) + config = result.to_fused_config() + assert isinstance(config, FusedConfig) + assert config.block_size_m == 256 + assert config.block_size_n == 256 + assert config.block_size_k == 64 + assert config.group_size_m == 24 + + def test_disabled_config_raises(self): + """Disabled configs should raise RuntimeError on to_fused_config().""" + result = select_ag_mm_config(M=131072, N=16384, K=16384, world_size=2) + assert result.enabled is False + with pytest.raises(RuntimeError, match="disabled"): + result.to_fused_config() + + +class TestListKnownShapes: + """Test list_known_shapes utility.""" + + def setup_method(self): + clear_config_cache() + + def test_ws8_lists_all_champions(self): + """ws=8 should list all champion shapes (at least 7).""" + shapes = list_known_shapes(world_size=8, transpose="NN", arch="mi300x") + assert len(shapes) >= 7, f"Expected at least 7 champion shapes, got {len(shapes)}" + # Should be sorted by speedup descending + speedups = [s["speedup"] for s in shapes] + assert speedups == sorted(speedups, reverse=True) + # Top shape should be g2 (1.343x) + assert shapes[0]["label"] == "g2" + assert shapes[0]["speedup"] == 1.343 + + def test_ws2_lists_empty(self): + """ws=2 (disabled) should return empty list.""" + shapes = list_known_shapes(world_size=2, transpose="NN", arch="mi300x") + assert shapes == [] + + def test_ws4_lists_empty(self): + """ws=4 (disabled) should return empty list.""" + shapes = list_known_shapes(world_size=4, transpose="NN", arch="mi300x") + assert shapes == [] + + +class TestUnknownArchTranspose: + """Test behavior with unknown architectures/transposes.""" + + def setup_method(self): + clear_config_cache() + + def test_unknown_arch_uses_global_default(self): + """Unknown GPU arch should fall through to global default.""" + result = select_ag_mm_config(M=131072, N=16384, K=16384, world_size=8, transpose="NN", arch="mi300a") + # No config file exists for mi300a, should use heuristic fallback + assert result.enabled is True + assert "Heuristic" in result.source or "fallback" in result.source + + def test_unknown_transpose_ws_lt8_disabled(self): + """Unknown transpose at ws<8 should still be disabled.""" + result = select_ag_mm_config(M=131072, N=16384, K=16384, world_size=4, transpose="TN", arch="mi300x") + assert result.enabled is False + + +class TestHbmBufferParams: + """Test that HBM buffer-specific params are returned.""" + + def setup_method(self): + clear_config_cache() + + def test_exact_match_has_hbm_params(self): + """Exact match should include hbm_buffer_params.""" + result = select_ag_mm_config(M=131072, N=16384, K=16384, world_size=8) + assert result.enabled is True + hbm = result.hbm_buffer_params + assert "k_per_flag" in hbm + assert "num_fetch_sms" in hbm + assert "num_fetch_stages" in hbm + assert "first_stage_fetch_sms" in hbm + # g2 specific values + assert hbm["k_per_flag"] == 32 + assert hbm["num_fetch_sms"] == 4 + assert hbm["num_fetch_stages"] == 64 + assert hbm["first_stage_fetch_sms"] == 52 + + def test_heuristic_has_hbm_params(self): + """Heuristic fallback should also include hbm_buffer_params.""" + result = select_ag_mm_config(M=65536, N=4096, K=8192, world_size=8) + assert result.enabled is True + hbm = result.hbm_buffer_params + assert "k_per_flag" in hbm + assert "num_fetch_sms" in hbm + assert hbm["k_per_flag"] > 0 + + def test_g5_hbm_params(self): + """g5 K-dominant shape should have specific HBM params.""" + result = select_ag_mm_config(M=8192, N=8192, K=262144, world_size=8) + hbm = result.hbm_buffer_params + assert hbm["k_per_flag"] == 32 + assert hbm["num_fetch_sms"] == 4 + assert hbm["num_fetch_stages"] == 8 + + def test_g15_hbm_params(self): + """g15 highest TFLOPS shape should have specific HBM params.""" + result = select_ag_mm_config(M=327680, N=28672, K=4096, world_size=8) + hbm = result.hbm_buffer_params + assert hbm["k_per_flag"] == 16 + assert hbm["num_fetch_sms"] == 4 + assert hbm["num_fetch_stages"] == 32 + assert hbm["first_stage_fetch_sms"] == 52 + + +class TestRegressionSizes: + """Test regression size loading.""" + + def test_load_regression_sizes(self): + """Should load regression test sizes from JSON — includes ws=2/4/8 entries.""" + sizes = load_regression_sizes() + assert len(sizes) >= 9, f"Expected at least 9 regression sizes (7 ws=8 + 2 disabled), got {len(sizes)}" + + def test_regression_sizes_have_required_fields(self): + """Each regression size should have name, M, N, K, world_sizes.""" + sizes = load_regression_sizes() + for s in sizes: + assert "name" in s, f"Missing 'name' in regression size: {s}" + assert "M" in s, f"Missing 'M' in regression size: {s}" + assert "N" in s, f"Missing 'N' in regression size: {s}" + assert "K" in s, f"Missing 'K' in regression size: {s}" + assert "world_sizes" in s, f"Missing 'world_sizes' in regression size: {s}" + assert isinstance(s["world_sizes"], list) + + def test_regression_sizes_match_configs(self): + """Each regression size should have a matching config in ws8.json.""" + clear_config_cache() + sizes = load_regression_sizes() + for s in sizes: + if 8 in s["world_sizes"]: + result = select_ag_mm_config(s["M"], s["N"], s["K"], world_size=8) + assert result.enabled is True, f"{s['name']} should be enabled at ws=8" + + def test_disabled_regression_sizes_correctly_disabled(self): + """Disabled regression entries (ws=2, ws=4) must be flagged disabled by auto-config.""" + clear_config_cache() + sizes = load_regression_sizes() + disabled_entries = [s for s in sizes if s["tier"] == "disabled"] + assert len(disabled_entries) >= 2, "Expected at least 2 disabled regression entries" + for s in disabled_entries: + for ws in s["world_sizes"]: + result = select_ag_mm_config(s["M"], s["N"], s["K"], world_size=ws) + assert result.enabled is False, ( + f"{s['name']} ws={ws} should be disabled but got enabled={result.enabled}" + ) + + def test_regression_sizes_cover_all_world_sizes(self): + """Regression sizes should cover ws=2, ws=4, and ws=8.""" + sizes = load_regression_sizes() + covered_ws = set() + for s in sizes: + for ws in s["world_sizes"]: + covered_ws.add(ws) + assert 2 in covered_ws, "Missing ws=2 coverage in regression sizes" + assert 4 in covered_ws, "Missing ws=4 coverage in regression sizes" + assert 8 in covered_ws, "Missing ws=8 coverage in regression sizes" + + +class TestConfigCache: + """Test config cache behavior.""" + + def test_cache_clear(self): + """Cache should be clearable.""" + # Populate cache + select_ag_mm_config(M=131072, N=16384, K=16384, world_size=8) + # Clear + clear_config_cache() + # Should still work after clear + result = select_ag_mm_config(M=131072, N=16384, K=16384, world_size=8) + assert result.enabled is True + assert result.shape_key == "131072x16384x16384" + + +class TestTransposeCoverage: + """Reviewer feedback #1: Verify transpose coverage and document NN-only support.""" + + def setup_method(self): + clear_config_cache() + + def test_only_nn_has_tuned_configs(self): + """Only NN transpose has tuned configs — the AG+MM kernel is NN-only.""" + assert SUPPORTED_TRANSPOSES == ("NN",) + + def test_nn_ws8_returns_exact_match(self): + """NN transpose at ws=8 should find champion configs.""" + result = select_ag_mm_config(M=131072, N=16384, K=16384, world_size=8, transpose="NN") + assert result.enabled is True + assert "Exact match" in result.source + + def test_tn_ws8_returns_heuristic_fallback(self): + """TN transpose at ws=8 — no config file exists, falls back to heuristic.""" + result = select_ag_mm_config(M=131072, N=16384, K=16384, world_size=8, transpose="TN") + # No configs/ag_mm/mi300x/TN/ directory → heuristic fallback for ws>=8 + assert result.enabled is True + assert "Heuristic" in result.source or "fallback" in result.source + + def test_nt_ws8_returns_heuristic_fallback(self): + """NT transpose at ws=8 — no config file exists, falls back to heuristic.""" + result = select_ag_mm_config(M=131072, N=16384, K=16384, world_size=8, transpose="NT") + assert result.enabled is True + assert "Heuristic" in result.source or "fallback" in result.source + + def test_tt_ws8_returns_heuristic_fallback(self): + """TT transpose at ws=8 — no config file exists, falls back to heuristic.""" + result = select_ag_mm_config(M=131072, N=16384, K=16384, world_size=8, transpose="TT") + assert result.enabled is True + assert "Heuristic" in result.source or "fallback" in result.source + + def test_non_nn_ws_lt8_disabled(self): + """Non-NN transpose at ws<8 should still be disabled.""" + for t in ("TN", "NT", "TT"): + result = select_ag_mm_config(M=131072, N=16384, K=16384, world_size=4, transpose=t) + assert result.enabled is False, f"{t} at ws=4 should be disabled" + + def test_supported_architectures(self): + """Only mi300x has tuned configs.""" + assert "mi300x" in SUPPORTED_ARCHITECTURES + + +class TestHeuristicFallbackValidation: + """Reviewer feedback #2: Validate the heuristic-derived configs with concrete examples.""" + + def setup_method(self): + clear_config_cache() + + def test_heuristic_small_m_small_k(self): + """M=4096, K=4096: bm=128, gm=8, kpf=1 (K_blocks=64, kpf starts at 8).""" + config, hbm = _apply_heuristic(M=4096, N=4096, K=4096) + assert config["block_size_m"] == 128, f"Small M should get bm=128, got {config['block_size_m']}" + assert config["block_size_n"] == 256 + assert config["group_size_m"] == 8, f"M=4096 should get gm=8, got {config['group_size_m']}" + assert hbm["k_per_flag"] > 0 + # K=4096, bk=64 → num_k_blocks=64 → kpf starts at 8 + assert hbm["k_per_flag"] == 8 + # M=4096 / bm=128 = 32 tiles → num_fetch_sms=16 + assert hbm["num_fetch_sms"] == 16 + + def test_heuristic_medium_m_large_k(self): + """M=16384, K=131072: bm=128, gm=16, kpf=16.""" + config, hbm = _apply_heuristic(M=16384, N=16384, K=131072) + assert config["block_size_m"] == 128, "M=16384 should still use bm=128" + assert config["group_size_m"] == 16, "M=16384 should get gm=16" + # K=131072, bk=64 → num_k_blocks=2048 → kpf=64 (>=512) + assert hbm["k_per_flag"] == 64 + # M=16384 / bm=128 = 128 tiles → num_fetch_sms=32 + assert hbm["num_fetch_sms"] == 32 + + def test_heuristic_large_m_small_k(self): + """M=327680, K=4096: bm=256, gm=24, kpf=4.""" + config, hbm = _apply_heuristic(M=327680, N=28672, K=4096) + assert config["block_size_m"] == 256, "Large M should get bm=256" + assert config["group_size_m"] == 24 + # K=4096, bk=64 → num_k_blocks=64 → kpf starts at 8 + assert hbm["k_per_flag"] == 8 + # M=327680 / bm=256 = 1280 tiles → num_fetch_sms=52 + assert hbm["num_fetch_sms"] == 52 + + def test_heuristic_matches_champion_g2(self): + """Heuristic for g2 shape (131072x16384x16384) should produce bm=256, gm=24.""" + config, hbm = _apply_heuristic(M=131072, N=16384, K=16384) + # g2 champion uses bm=256, gm=24 — heuristic should agree on these + assert config["block_size_m"] == 256 + assert config["group_size_m"] == 24 + assert config["num_warps"] == 8 + assert config["num_stages"] == 2 + # HBM params: K=16384, bk=64 → 256 blocks → kpf=16 + assert hbm["k_per_flag"] == 16 + + def test_heuristic_matches_champion_g5(self): + """Heuristic for g5 shape (8192x8192x262144) should produce bm=128, gm=8.""" + config, hbm = _apply_heuristic(M=8192, N=8192, K=262144) + # g5 champion uses bm=128, gm=8 — heuristic should agree + assert config["block_size_m"] == 128 + assert config["group_size_m"] == 8 + # K=262144, bk=64 → 4096 blocks → kpf=64 (>=512) + assert hbm["k_per_flag"] == 64 + + def test_heuristic_kpf_divisibility(self): + """k_per_flag must evenly divide num_k_blocks.""" + for K in [2048, 4096, 8192, 16384, 32768, 65536, 131072, 262144]: + config, hbm = _apply_heuristic(M=16384, N=16384, K=K) + num_k_blocks = K // 64 # bk=64 is invariant + kpf = hbm["k_per_flag"] + assert num_k_blocks % kpf == 0, ( + f"K={K}: k_per_flag={kpf} does not evenly divide num_k_blocks={num_k_blocks}" + ) + + def test_heuristic_via_select_for_unknown_shape(self): + """Full pipeline: unknown shape at ws=8 uses heuristic with reasonable params.""" + # Llama-13B gate projection: batch=2048, hidden=5120, intermediate=13824 + result = select_ag_mm_config(M=2048 * 8, N=13824, K=5120, world_size=8, transpose="NN", arch="mi300x") + assert result.enabled is True + assert result.shape_key is None # no exact match + assert "Heuristic" in result.source + # M=16384 → bm=128 (at boundary), gm=16 + assert result.config_params["block_size_m"] == 128 + assert result.config_params["group_size_m"] == 16 + assert result.hbm_buffer_params["k_per_flag"] > 0 + # Can convert to FusedConfig + fc = result.to_fused_config() + assert isinstance(fc, FusedConfig) + + +class TestIntegrationPath: + """Reviewer feedback #4: Show exactly how select_ag_mm_config() integrates with harness.""" + + def setup_method(self): + clear_config_cache() + + def test_auto_select_to_fused_config_pipeline(self): + """Demonstrate full pipeline: auto-select → FusedConfig → pass to kernel. + + This is the concrete integration pattern: + result = select_ag_mm_config(M, N, K, world_size=8) + if result.enabled: + config = result.to_fused_config() + hbm = result.hbm_buffer_params + # config passed to: all_gather_matmul(shmem, C, A, B, config=config) + # hbm used for: k_per_flag, num_fetch_sms in the HBM buffer kernel + else: + # Fallback to PyTorch: torch.distributed.all_gather + torch.matmul + pass + """ + # Step 1: Auto-select for a champion shape + result = select_ag_mm_config(M=131072, N=16384, K=16384, world_size=8) + assert result.enabled is True + + # Step 2: Convert to FusedConfig (this is what all_gather_matmul() accepts) + config = result.to_fused_config() + assert isinstance(config, FusedConfig) + assert config.block_size_m == 256 + assert config.block_size_n == 256 + assert config.block_size_k == 64 + assert config.group_size_m == 24 + assert config.num_xcds == 8 + + # num_warps/num_stages are in JSON but not FusedConfig fields — + # to_fused_config() correctly filters them; they're kernel launch params + assert "num_warps" in result.config_params # Present in raw dict + assert result.config_params["num_warps"] == 8 + + # Step 3: Validate config is internally consistent + config.validate(world_size=8) # Should not raise + + # Step 4: Access HBM buffer params (used by bench_all_gather_matmul.py) + hbm = result.hbm_buffer_params + assert hbm["k_per_flag"] == 32 + assert hbm["num_fetch_sms"] == 4 + assert hbm["num_fetch_stages"] == 64 + assert hbm["first_stage_fetch_sms"] == 52 + + def test_disabled_path_blocks_fused_config(self): + """Disabled configs must not produce FusedConfig — forces PyTorch fallback.""" + result = select_ag_mm_config(M=131072, N=16384, K=16384, world_size=2) + assert result.enabled is False + with pytest.raises(RuntimeError, match="disabled"): + result.to_fused_config() + + def test_regression_sizes_all_have_valid_configs(self): + """Every enabled regression size must auto-select to a valid, convertible FusedConfig.""" + sizes = load_regression_sizes() + for s in sizes: + if s.get("tier") == "disabled": + continue # disabled entries are tested separately + for ws in s["world_sizes"]: + result = select_ag_mm_config(s["M"], s["N"], s["K"], world_size=ws) + assert result.enabled is True, f"{s['name']} ws={ws} should be enabled" + config = result.to_fused_config() + config.validate(world_size=ws) # Must not raise + # HBM params must be complete + hbm = result.hbm_buffer_params + assert all(k in hbm for k in ["k_per_flag", "num_fetch_sms"]), f"{s['name']} ws={ws} missing HBM params" + + def test_config_content_g15_champion(self): + """Verify actual g15 champion config content — highest TFLOPS shape (474.7 TFLOPS).""" + result = select_ag_mm_config(M=327680, N=28672, K=4096, world_size=8) + assert result.enabled is True + assert result.shape_key == "327680x28672x4096" + assert result.speedup == 1.284 + + config = result.to_fused_config() + assert config.block_size_m == 256 + assert config.block_size_n == 256 + assert config.block_size_k == 64 + assert config.group_size_m == 24 + + hbm = result.hbm_buffer_params + assert hbm == { + "k_per_flag": 16, + "num_fetch_sms": 4, + "num_fetch_stages": 32, + "first_stage_fetch_sms": 52, + } + + +class TestNearestShapeMatching: + """Test nearest-shape matching for close-but-not-exact dimensions.""" + + def setup_method(self): + clear_config_cache() + + def test_close_to_g2_uses_g2_config(self): + """M=130000 is ~0.8% below g2 (131072) — should use g2 champion config.""" + result = select_ag_mm_config(M=130000, N=16384, K=16384, world_size=8, transpose="NN", arch="mi300x") + assert result.enabled is True + assert "Nearest match" in result.source + assert result.shape_key == "131072x16384x16384" + assert result.speedup == 1.343 + + def test_close_to_g15_uses_g15_config(self): + """M=320000 is ~2.3% below g15 (327680) — should use g15 champion config.""" + result = select_ag_mm_config(M=320000, N=28672, K=4096, world_size=8, transpose="NN", arch="mi300x") + assert result.enabled is True + assert "Nearest match" in result.source + assert result.shape_key == "327680x28672x4096" + + def test_far_shape_uses_heuristic(self): + """M=50000, N=50000, K=50000 — far from all champions, should use heuristic.""" + result = select_ag_mm_config(M=50000, N=50000, K=50000, world_size=8, transpose="NN", arch="mi300x") + assert result.enabled is True + assert "Heuristic" in result.source + assert result.shape_key is None + + def test_exact_match_still_preferred(self): + """Exact match should still be returned even when nearest would also match.""" + result = select_ag_mm_config(M=131072, N=16384, K=16384, world_size=8, transpose="NN", arch="mi300x") + assert result.enabled is True + assert "Exact match" in result.source + assert result.shape_key == "131072x16384x16384" + + def test_nearest_skips_losing_shapes(self): + """Nearest matching should skip shapes with speedup <= 1.0.""" + # g9 (196608x18432x16384) has speedup 0.950 — should NOT be matched + result = select_ag_mm_config(M=196608, N=18432, K=16384, world_size=8, transpose="NN", arch="mi300x") + # g9 is an exact match, but its speedup is 0.950 + # The exact match path doesn't filter by speedup, but nearest does + assert result.enabled is True + + +class TestFusedConfigNumWarpsStages: + """Test that num_warps and num_stages are present in config_params.""" + + def setup_method(self): + clear_config_cache() + + def test_champion_config_has_num_warps(self): + """Champion configs should have num_warps=8 in config_params.""" + result = select_ag_mm_config(M=131072, N=16384, K=16384, world_size=8) + assert result.config_params["num_warps"] == 8 + + def test_champion_config_has_num_stages(self): + """Champion configs should have num_stages=2 in config_params.""" + result = select_ag_mm_config(M=131072, N=16384, K=16384, world_size=8) + assert result.config_params["num_stages"] == 2 + + def test_fused_config_conversion_succeeds(self): + """FusedConfig conversion should succeed for enabled configs.""" + result = select_ag_mm_config(M=131072, N=16384, K=16384, world_size=8) + config = result.to_fused_config() + assert isinstance(config, FusedConfig) + + def test_fused_config_block_sizes_correct(self): + """FusedConfig should have correct block sizes from champion.""" + result = select_ag_mm_config(M=131072, N=16384, K=16384, world_size=8) + config = result.to_fused_config() + assert config.block_size_m == 256 + assert config.block_size_n == 256 + assert config.block_size_k == 64 + + def test_all_champions_have_num_warps_8(self): + """All ws=8 champion configs should specify num_warps=8.""" + winning_shapes = [ + (131072, 16384, 16384), + (327680, 28672, 4096), + (147456, 28672, 4096), + (229376, 28672, 4096), + (8192, 8192, 262144), + (262144, 8192, 8192), + (16384, 16384, 131072), + ] + for M, N, K in winning_shapes: + result = select_ag_mm_config(M, N, K, world_size=8) + assert result.config_params["num_warps"] == 8, f"Shape {M}x{N}x{K}: expected num_warps=8" + assert result.config_params["num_stages"] == 2, f"Shape {M}x{N}x{K}: expected num_stages=2" + + +class TestGpuArchAutoDetection: + """Test GPU architecture auto-detection logic.""" + + def setup_method(self): + clear_config_cache() + # Reset cached detection + auto_config_module._detected_arch = None + + def teardown_method(self): + # Reset after each test + auto_config_module._detected_arch = None + os.environ.pop("IRIS_GPU_ARCH", None) + + def test_env_var_override(self): + """IRIS_GPU_ARCH env var should override auto-detection.""" + os.environ["IRIS_GPU_ARCH"] = "mi300a" + arch = detect_gpu_arch() + assert arch == "mi300a" + + def test_env_var_case_insensitive(self): + """IRIS_GPU_ARCH env var should be lowercased by detect_gpu_arch.""" + os.environ["IRIS_GPU_ARCH"] = "MI300X" + arch = detect_gpu_arch() + # detect_gpu_arch applies .lower() to the env var + assert arch == "mi300x" + + def test_fallback_default(self): + """Without rocminfo, should fall back to mi300x.""" + os.environ.pop("IRIS_GPU_ARCH", None) + # On a machine without ROCm, should get mi300x default + arch = detect_gpu_arch() + assert isinstance(arch, str) + assert len(arch) > 0 + + def test_select_with_auto_arch(self): + """select_ag_mm_config(arch='auto') should work end-to-end.""" + os.environ["IRIS_GPU_ARCH"] = "mi300x" + result = select_ag_mm_config(M=131072, N=16384, K=16384, world_size=8, arch="auto") + assert result.enabled is True + assert result.shape_key == "131072x16384x16384" + + def test_select_default_is_auto(self): + """Default arch parameter should be 'auto'.""" + os.environ["IRIS_GPU_ARCH"] = "mi300x" + # Call without specifying arch — should use auto + result = select_ag_mm_config(M=131072, N=16384, K=16384, world_size=8) + assert result.enabled is True + + def test_caching_across_calls(self): + """detect_gpu_arch should cache result after first call.""" + os.environ["IRIS_GPU_ARCH"] = "mi300x" + arch1 = detect_gpu_arch() + os.environ["IRIS_GPU_ARCH"] = "mi300a" + arch2 = detect_gpu_arch() + # Should be cached from first call + assert arch1 == arch2 == "mi300x" + + +if __name__ == "__main__": + pytest.main([__file__, "-v"])