diff --git a/.gitignore b/.gitignore index 021f8bc9..a743dc66 100755 --- a/.gitignore +++ b/.gitignore @@ -6,6 +6,7 @@ build/ .idea .claude workspace_* +/experiments/ # Auto-cloned agent dependencies agents/geak_optimagentv2/GEAK-agent/ diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index 54fc77cb..5e135369 100644 --- a/.pre-commit-config.yaml +++ b/.pre-commit-config.yaml @@ -6,4 +6,4 @@ repos: entry: make check-perf-helpers language: system pass_filenames: false - files: ^(src/tools/perf/|src/tools/sync_perf_helpers\.py|src/tools/materialize_perf_helpers\.py|src/perf_helper_materialization\.py|tasks/.*/rocmbench/.*/performance_utils_pytest\.py|tasks/triton2triton/vllm/.*/scripts/task_runner\.py)$ + files: ^(src/tools/perf/|src/tools/sync_perf_helpers\.py|src/tools/materialize_perf_helpers\.py|src/perf_helper_materialization\.py|tasks/.*/rocmbench/.*/performance_utils_pytest\.py|tasks/.*/scripts/task_runner\.py)$ diff --git a/Makefile b/Makefile index 6e6e2392..d1b2c9f0 100755 --- a/Makefile +++ b/Makefile @@ -39,7 +39,7 @@ help: @echo "make check-perf-helpers - Verify task perf-helper stubs and markers are valid" @echo "make materialize-perf-workspace WORKSPACE=workspace_x - Inject canonical perf helpers into workspace(s)" @echo "make materialize-perf-task TASK=tasks/... OUT=/tmp/aka-task - Copy task(s) and inject canonical perf helpers" - @echo "make cleanup-works - Remove workspace_* directories and logs" + @echo "make cleanup-works - Remove experiments/workspace_* directories and logs" @echo "make install-cursor-agent - Install the Cursor Agent CLI on the host" DOCKER_RUNNER := src/scripts/docker_benchmark.sh @@ -123,7 +123,7 @@ materialize-perf-task: cleanup-works: @echo "Removing workspace directories and logs..." - @rm -rf workspace_* + @rm -rf experiments/workspace_* @rm -rf logs @echo "✓ Workspace directories and logs removed" diff --git a/README.md b/README.md index b15d264b..7d49b695 100755 --- a/README.md +++ b/README.md @@ -40,8 +40,8 @@ Compare the generated reports directly: ```bash python3 src/tools/compare_runs.py \ - workspace_MI300_claude_code/run__baseline \ - workspace_MI300_claude_code/run__treatment + experiments/workspace_MI300_claude_code/run__baseline \ + experiments/workspace_MI300_claude_code/run__treatment ``` For visual comparison, build the local dashboard as described in the @@ -57,6 +57,7 @@ the observed deltas together with run-to-run variance. AgentKernelArena/ ├── main.py # Run orchestration, resume, and parallel queue ├── example_configs/ # Quickstart and curated benchmark run configs +├── experiments/ # Ignored generated workspaces and run artifacts ├── src/ │ ├── module_registration.py # Agent registration and handler selection │ ├── preprocessing.py # Workspace and repository setup diff --git a/agents/geak_v3/README.md b/agents/geak_v3/README.md index a4a5cfec..4b55fa5e 100644 --- a/agents/geak_v3/README.md +++ b/agents/geak_v3/README.md @@ -95,10 +95,10 @@ with more GPUs via `GEAK_GPU_IDS` / `--num-parallel`. - **AKA run log:** `logs/*.log` (path from `log_directory` in `config.yaml`). - **Per-task result (authoritative):** - `workspace__geak_v3/run_/_/task_result.yaml` + `experiments/workspace__geak_v3/run_/_/task_result.yaml` (plus `baseline_perf.yaml`, `optimized_perf.yaml`, and `build/performance_report.json`). - **GEAK internals:** - `workspace__geak_v3/run_/__logs/` + `experiments/workspace__geak_v3/run_/__logs/` (`final_report.json`, `geak_agent.log`, the winning `.diff`). -- **Aggregate:** `workspace__geak_v3/run_/reports/overall_summary.csv`. +- **Aggregate:** `experiments/workspace__geak_v3/run_/reports/overall_summary.csv`. diff --git a/agents/geak_v3_triton/README.md b/agents/geak_v3_triton/README.md index 6985c412..4daa8a10 100644 --- a/agents/geak_v3_triton/README.md +++ b/agents/geak_v3_triton/README.md @@ -118,7 +118,7 @@ Override per run via env: `GEAK_GPU_IDS`, `GEAK_NUM_PARALLEL`, `GEAK_MODEL`, docker exec geak-aka tail -f /workspace/logs/*.log # Per-kernel AKA results (authoritative) -for f in workspace_*/run_*/*/task_result.yaml; do +for f in experiments/workspace_*/run_*/*/task_result.yaml; do [ -f "$f" ] && echo "$(basename $(dirname "$f")): $(grep speedup_ratio "$f")" done ``` diff --git a/docs/examples/examples.md b/docs/examples/examples.md index ea3b93e6..9d6e6c28 100644 --- a/docs/examples/examples.md +++ b/docs/examples/examples.md @@ -34,7 +34,7 @@ Run the single-task Claude Code quickstart for the physical GPU. 3. Inspect the result: ```text - workspace_MI300_claude_code/ + experiments/workspace_MI300_claude_code/ └── run_/ └── hip2hip_gpumode_GELU_/ └── task_result.yaml @@ -89,7 +89,7 @@ make docker-run CONFIG=config_triton.yaml RUN_ARGS="--run-suffix baseline" make docker-run CONFIG=config_triton.yaml RUN_ARGS="--run-suffix with_capability" ``` -Both runs land in `workspace_MI300_/` with distinct run names. Build the +Both runs land in `experiments/workspace_MI300_/` with distinct run names. Build the dashboard and compare them side-by-side: ```bash @@ -153,7 +153,7 @@ make docker-parallel-run \ The run directory contains normal per-task workspaces plus a scheduler queue: ```text -workspace_MI355X_claude_code/ +experiments/workspace_MI355X_claude_code/ └── run__claude_parallel8/ ├── .parallel/ │ ├── pending/ diff --git a/docs/how-to/add-task.md b/docs/how-to/add-task.md index 4e7f7d5b..fa7ac7e8 100644 --- a/docs/how-to/add-task.md +++ b/docs/how-to/add-task.md @@ -156,8 +156,9 @@ The shared performance timing helpers are generated from `src/tools/perf/` into run workspace. In committed task sources: - `tasks/*/rocmbench/**/performance_utils_pytest.py` is intentionally a stub. -- The `AKA-GENERATED` region in `triton2triton/vllm/*/scripts/task_runner.py` is - intentionally a stub block. +- The `AKA-GENERATED` region in any `tasks/**/scripts/task_runner.py` that opts + into shared timing is intentionally a stub block. The vLLM runners remain + mandatory users of this mechanism. Do not hand-edit those stubs. If a task needs shared timing behavior, add the stub/marker and run `make sync-perf-helpers`. If you need to change timing logic, diff --git a/docs/how-to/agents.md b/docs/how-to/agents.md index 0527c164..26f1d677 100644 --- a/docs/how-to/agents.md +++ b/docs/how-to/agents.md @@ -95,8 +95,8 @@ You can also generate a text comparison directly: ```bash python3 src/tools/compare_runs.py \ - workspace_MI300_claude_code/run__baseline \ - workspace_MI300_claude_code/run__with_capability + experiments/workspace_MI300_claude_code/run__baseline \ + experiments/workspace_MI300_claude_code/run__with_capability ``` The resulting `task_result.yaml` files expose compilation, correctness, timing, diff --git a/docs/how-to/held-out-evaluation.md b/docs/how-to/held-out-evaluation.md index 0760ce06..23544019 100644 --- a/docs/how-to/held-out-evaluation.md +++ b/docs/how-to/held-out-evaluation.md @@ -57,7 +57,7 @@ Pass a completed run directory and the generated held-out configurations: ```bash python3 -m src.held_out.run_heldout_eval \ - --run-dir workspace_MI300_claude_code/run_ \ + --run-dir experiments/workspace_MI300_claude_code/run_ \ --heldout-dir held_out_tests/ \ --tasks-dir tasks/ ``` diff --git a/docs/how-to/parallel-run.md b/docs/how-to/parallel-run.md index 50fd4d1d..5c4a8730 100644 --- a/docs/how-to/parallel-run.md +++ b/docs/how-to/parallel-run.md @@ -52,7 +52,7 @@ make docker-parallel-run \ The runner creates a shared queue in the run directory: ```text -workspace__/ +experiments/workspace__/ └── run_[_suffix]/ └── .parallel/ ├── pending/ diff --git a/docs/how-to/run-evaluation.md b/docs/how-to/run-evaluation.md index ddf0f36f..4f2b0d6a 100644 --- a/docs/how-to/run-evaluation.md +++ b/docs/how-to/run-evaluation.md @@ -95,7 +95,7 @@ Add a suffix to label a run directory (useful for A/B testing): ```bash make docker-run CONFIG="$CONFIG_PATH" RUN_ARGS="--run-suffix cursor_with_mcp" -# → workspace_MI300_cursor/run_20260617_101500_cursor_with_mcp +# → experiments/workspace_MI300_cursor/run_20260617_101500_cursor_with_mcp ``` For debugging, enter the same Docker runtime used by the experiment: @@ -193,10 +193,10 @@ make docker-parallel-run CONFIG="$CONFIG_PATH" GPU_IDS=0,1,2,3 RUN_ARGS="--resum ## Read the results -A run produces this layout under the workspace directory: +A run produces this ignored layout under `experiments/`: ```text -workspace__/ +experiments/workspace__/ └── run_/ ├── .parallel/ # present for docker-parallel-run │ ├── pending/ diff --git a/docs/how-to/task-validator.md b/docs/how-to/task-validator.md index e6ae433d..4d1ea703 100644 --- a/docs/how-to/task-validator.md +++ b/docs/how-to/task-validator.md @@ -37,7 +37,7 @@ make docker-run CONFIG=config_validator.yaml Each task workspace receives a `validation_report.yaml` with per-check results, and a `validation_summary.yaml` with aggregated statistics is written to the -workspace root. Tasks skipped by `platform_support.status: skip` or by a +workspace root under `experiments/`. Tasks skipped by `platform_support.status: skip` or by a non-matching `platform_support.required_arch` are filtered before workspace creation, so they do not produce a validation report or appear in the summary counts. diff --git a/docs/how-to/visualization.md b/docs/how-to/visualization.md index 284c8d49..31d1b18e 100644 --- a/docs/how-to/visualization.md +++ b/docs/how-to/visualization.md @@ -27,13 +27,13 @@ By default, it scans only visualization-specific report bundles: ``` Workspace-run reports, which are usually located at -`workspace__/run_/reports/`, can also be scanned, but this +`experiments/workspace__/run_/reports/`, can also be scanned, but this is opt-in. ## Build the dashboard data and serve it After a normal AgentKernelArena run, reports land in -`workspace__/run_/reports/`. Pass +`experiments/workspace__/run_/reports/`. Pass `--include-workspace-runs` so the module picks them up. Run the following from the AgentKernelArena repository root: diff --git a/docs/install/install.md b/docs/install/install.md index 92748162..44047d3c 100644 --- a/docs/install/install.md +++ b/docs/install/install.md @@ -180,7 +180,7 @@ provider settings. The quickstart run above confirms that the framework, GPU, and agent CLI work together. It creates a timestamped workspace directory -(`workspace__/run_/`), logs to `logs/`, and writes a +(`experiments/workspace__/run_/`), logs to `logs/`, and writes a `task_result.yaml` for the task. To run across multiple GPUs, list host GPU IDs or omit `GPU_IDS` to discover diff --git a/docs/reference/api-reference.md b/docs/reference/api-reference.md index 7b430567..ccc1e668 100644 --- a/docs/reference/api-reference.md +++ b/docs/reference/api-reference.md @@ -21,7 +21,7 @@ A run configuration defines a single experiment. Start from a file under | `tasks` | list of strings | Task selectors relative to `tasks/`. Use `all` for every task, a category prefix for a group, or a full path for a single task. | | `target_gpu_model` | string | Target GPU model, for example `MI300` or `MI355X`. Used to select the Docker image architecture, set `PYTORCH_ROCM_ARCH`, and name the workspace. | | `log_directory` | string | Directory for run logs. | -| `workspace_directory_prefix` | string | Prefix for the workspace directory. The full name is `__`. | +| `workspace_directory_prefix` | string | Prefix for the workspace directory. Runs are stored under the ignored path `experiments/__`. | Specialized GEAK and mini-swe integrations also accept some optional top-level fields: diff --git a/main.py b/main.py index fdff7abf..835bce79 100755 --- a/main.py +++ b/main.py @@ -340,7 +340,9 @@ def _build_context( workspace_directory_name = ( f"{config['workspace_directory_prefix']}_{config['target_gpu_model']}_{agent.value}" ) - workspace_directory = (project_root / workspace_directory_name).resolve() + workspace_directory = ( + project_root / "experiments" / workspace_directory_name + ).resolve() resolved_run = _resolve_run(args, workspace_directory) if resolved_run is None: return None diff --git a/src/held_out/README.md b/src/held_out/README.md index 8d0899be..33b06cf5 100644 --- a/src/held_out/README.md +++ b/src/held_out/README.md @@ -54,7 +54,7 @@ host-CLI integrations). ```bash python3 -m src.held_out.run_heldout_eval \ - --run-dir workspace_MI300_cursor/run_20260417_142419 \ + --run-dir experiments/workspace_MI300_cursor/run_20260417_142419 \ --heldout-dir held_out_tests/ \ --tasks-dir tasks/ ``` diff --git a/src/held_out/run_heldout_eval.py b/src/held_out/run_heldout_eval.py index ac3fb2a8..d13371cb 100644 --- a/src/held_out/run_heldout_eval.py +++ b/src/held_out/run_heldout_eval.py @@ -9,7 +9,7 @@ Usage: python3 -m src.held_out.run_heldout_eval \ - --run-dir workspace_MI300_cursor/run_20260417_142419 \ + --run-dir experiments/workspace_MI300_cursor/run_20260417_142419 \ --heldout-dir held_out_tests/ \ --tasks-dir tasks/ \ [--output-suffix _heldout] @@ -507,7 +507,7 @@ def main(): parser = argparse.ArgumentParser(description="Held-out evaluation for AgentKernelArena") parser.add_argument( "--run-dir", required=True, - help="Path to the completed run directory (e.g. workspace_MI300_cursor/run_20260417_142419)", + help="Path to the completed run directory (e.g. experiments/workspace_MI300_cursor/run_20260417_142419)", ) parser.add_argument( "--heldout-dir", required=True, diff --git a/src/perf_helper_materialization.py b/src/perf_helper_materialization.py index 55fa324b..ce30ca26 100644 --- a/src/perf_helper_materialization.py +++ b/src/perf_helper_materialization.py @@ -92,7 +92,7 @@ def _marker_filtered_targets(root: Path, pattern: str) -> list[Path]: bespoke timer are left untouched. """ targets = [] - for p in sorted(glob.glob(str(root / pattern))): + for p in sorted(glob.glob(str(root / pattern), recursive=True)): text = Path(p).read_text() if any(marker in text for marker in MARK_STARTS) or MARK_END in text: targets.append(Path(p)) @@ -104,6 +104,17 @@ def image_kernel_targets(root: Path = ROOT) -> list[Path]: return _marker_filtered_targets(root, "tasks/image_kernel/*/scripts/task_runner.py") +def marked_task_runner_targets(root: Path = ROOT) -> list[Path]: + """Return any task runner that opts into the shared inline timing helper. + + vLLM runners remain mandatory targets through :func:`vllm_targets`. Other + task suites opt in by carrying the AKA-GENERATED marker block, which lets + new suites use the canonical helper without adding another path-specific + scanner here. + """ + return _marker_filtered_targets(root, "tasks/**/scripts/task_runner.py") + + def canonical_rocmbench_helper(root: Path = ROOT) -> str: return (root / "src" / "tools" / "perf" / "performance_utils_pytest.py").read_text() diff --git a/src/postprocessing.py b/src/postprocessing.py index 7d0026a7..81e603b0 100755 --- a/src/postprocessing.py +++ b/src/postprocessing.py @@ -125,9 +125,9 @@ def _get_run_directory(workspace_paths: List[str]) -> Path: Extract run directory from workspace paths. Workspace paths are task directories like: - workspace_MI300_cursor/run_20250115_143022/task_hip2hip_silu_20250115_143022/ + experiments/workspace_MI300_cursor/run_20250115_143022/task_hip2hip_silu_20250115_143022/ - Returns the run directory: workspace_MI300_cursor/run_20250115_143022/ + Returns the run directory: experiments/workspace_MI300_cursor/run_20250115_143022/ """ if not workspace_paths: raise ValueError("Cannot determine run directory: empty workspace_paths") @@ -360,7 +360,7 @@ def _collect_all_tasks_from_run(run_directory: Path) -> List[str]: Collect all task directories from a run directory that have task_result.yaml. Args: - run_directory: Run-level directory (e.g., workspace_MI300_cursor/run_20250115_143022/) + run_directory: Run-level directory (e.g., experiments/workspace_MI300_cursor/run_20250115_143022/) Returns: List of task directory paths (as strings) that have task_result.yaml @@ -665,5 +665,5 @@ def export_task_results_csv( if __name__ == "__main__": # manually generate report - workspace_path = "workspace_MI300_claude_code" + workspace_path = "experiments/workspace_MI300_claude_code" general_post_processing(workspace_path, logger = None) diff --git a/src/preprocessing.py b/src/preprocessing.py index 3674be78..a58c6573 100755 --- a/src/preprocessing.py +++ b/src/preprocessing.py @@ -448,7 +448,7 @@ def is_task_complete( Check if a task is already completed. Args: - run_directory: Run-level directory (e.g., workspace_MI300_cursor/run_20250115_143022/) + run_directory: Run-level directory (e.g., experiments/workspace_MI300_cursor/run_20250115_143022/) task_name: Full task name (e.g., "hip2hip/gpumode/SiLU") timestamp: Timestamp string used in task directory name agent_name: Agent name. task_validator uses validation_report.yaml; @@ -478,7 +478,7 @@ def setup_workspace(task_config_dir: str, run_directory: Path, timestamp: str, l Args: task_config_dir: Path to task's config.yaml - run_directory: Run-level directory (e.g., workspace_MI300_cursor/run_20250115_143022/) + run_directory: Run-level directory (e.g., experiments/workspace_MI300_cursor/run_20250115_143022/) timestamp: Timestamp string for unique workspace naming logger: Logger instance task_name: Full task name (e.g., "hip2hip/gpumode/SiLU") for unique directory naming diff --git a/src/scripts/docker_benchmark.sh b/src/scripts/docker_benchmark.sh index 00ec90f7..c510cc77 100755 --- a/src/scripts/docker_benchmark.sh +++ b/src/scripts/docker_benchmark.sh @@ -922,7 +922,7 @@ resolve_workspace_dir_for_config() { [[ -n "$prefix" ]] || die "workspace_directory_prefix not found in $config" [[ -n "$model" ]] || die "target_gpu_model not found in $config" [[ -n "$agent" ]] || die "agent.template not found in $config" - printf '%s/%s_%s_%s\n' "$HOST_ROOT" "$prefix" "$model" "$agent" + printf '%s/experiments/%s_%s_%s\n' "$HOST_ROOT" "$prefix" "$model" "$agent" } resolve_latest_run_name() { diff --git a/src/tools/compare_runs.py b/src/tools/compare_runs.py index d85ab20b..f31b5633 100644 --- a/src/tools/compare_runs.py +++ b/src/tools/compare_runs.py @@ -4,7 +4,7 @@ Usage: python3 src/tools/compare_runs.py run1_path run2_path - python3 src/tools/compare_runs.py workspace_MI300_cursor/run_20260714_120000_baseline workspace_MI300_cursor/run_20260714_140000_treatment + python3 src/tools/compare_runs.py experiments/workspace_MI300_cursor/run_20260714_120000_baseline experiments/workspace_MI300_cursor/run_20260714_140000_treatment """ import json @@ -19,7 +19,7 @@ def load_run_data(run_path: Path) -> Dict[str, Any]: Load task_type_breakdown.json from a run directory. Args: - run_path: Path to run directory (e.g., workspace_MI300_cursor/run_20260714_120000_baseline) + run_path: Path to run directory (e.g., experiments/workspace_MI300_cursor/run_20260714_120000_baseline) Returns: Dictionary containing run data from JSON file @@ -342,7 +342,7 @@ def main(): epilog=""" Examples: # Compare two runs - python3 src/tools/compare_runs.py workspace_MI300_cursor/run_20260714_120000_baseline workspace_MI300_cursor/run_20260714_140000_treatment + python3 src/tools/compare_runs.py experiments/workspace_MI300_cursor/run_20260714_120000_baseline experiments/workspace_MI300_cursor/run_20260714_140000_treatment # Compare and save to file python3 src/tools/compare_runs.py run1 run2 --output comparison_report.txt @@ -352,13 +352,13 @@ def main(): parser.add_argument( 'run1', type=str, - help='Path to baseline/first run directory (e.g., workspace_MI300_cursor/run_20260714_120000_baseline)' + help='Path to baseline/first run directory (e.g., experiments/workspace_MI300_cursor/run_20260714_120000_baseline)' ) parser.add_argument( 'run2', type=str, - help='Path to treatment/second run directory (e.g., workspace_MI300_cursor/run_20260714_140000_treatment)' + help='Path to treatment/second run directory (e.g., experiments/workspace_MI300_cursor/run_20260714_140000_treatment)' ) parser.add_argument( diff --git a/src/tools/sync_perf_helpers.py b/src/tools/sync_perf_helpers.py index 1971abd0..a7e9e81c 100644 --- a/src/tools/sync_perf_helpers.py +++ b/src/tools/sync_perf_helpers.py @@ -14,9 +14,9 @@ Two helper families: 1. Every */rocmbench/**/performance_utils_pytest.py should be the committed stub from src.perf_helper_materialization. setup_workspace() replaces it. - 2. Every triton2triton/vllm/*/scripts/task_runner.py should contain the - committed stub block between AKA-GENERATED markers. setup_workspace() - replaces that block. + 2. Every triton2triton/vllm/*/scripts/task_runner.py, plus any other task + runner that opts in with AKA-GENERATED markers, should contain the + committed stub block. setup_workspace() replaces that block. """ import argparse import pathlib @@ -28,7 +28,7 @@ from src.perf_helper_materialization import ( # noqa: E402 ROCMBENCH_HELPER_STUB, VLLM_HELPER_STUB_BLOCK, - image_kernel_targets, + marked_task_runner_targets, replace_marked_region, rocmbench_targets, vllm_targets, @@ -50,7 +50,9 @@ def main() -> int: p.write_text(ROCMBENCH_HELPER_STUB) wrote += 1 - inline_targets = list(vllm_targets(ROOT)) + list(image_kernel_targets(ROOT)) + inline_targets = sorted( + set(vllm_targets(ROOT)) | set(marked_task_runner_targets(ROOT)) + ) for p in inline_targets: cur = p.read_text() new = replace_marked_region(cur, VLLM_HELPER_STUB_BLOCK) @@ -73,8 +75,8 @@ def main() -> int: return 0 print(f"synced {wrote} file(s) " - f"({len(rocmbench_targets(ROOT))} rocmbench + {len(vllm_targets(ROOT))} vllm " - f"+ {len(image_kernel_targets(ROOT))} image_kernel checked)") + f"({len(rocmbench_targets(ROOT))} rocmbench + " + f"{len(inline_targets)} inline task runners checked)") return 0 diff --git a/src/visualization/README.md b/src/visualization/README.md index 302accc7..e4c88e7d 100644 --- a/src/visualization/README.md +++ b/src/visualization/README.md @@ -13,7 +13,7 @@ The module keeps three kinds of files separate: Normal run reports remain under: ```text -workspace__/run_/reports/ +experiments/workspace__/run_/reports/ ``` ## Commands diff --git a/src/visualization/__main__.py b/src/visualization/__main__.py index 8df68b1d..3833d80b 100644 --- a/src/visualization/__main__.py +++ b/src/visualization/__main__.py @@ -19,7 +19,7 @@ def create_parser() -> argparse.ArgumentParser: build_parser.add_argument( "--include-workspace-runs", action="store_true", - help="Include workspace_*/run_*/reports directories.", + help="Include experiments/workspace_*/run_*/reports directories.", ) serve_parser = subparsers.add_parser("serve", help="Serve the dashboard.") @@ -31,7 +31,7 @@ def create_parser() -> argparse.ArgumentParser: run_parser.add_argument( "--include-workspace-runs", action="store_true", - help="Include workspace_*/run_*/reports directories.", + help="Include experiments/workspace_*/run_*/reports directories.", ) _add_server_arguments(run_parser) return parser diff --git a/src/visualization/build_data.py b/src/visualization/build_data.py index 9c5e1285..de2d9ec4 100644 --- a/src/visualization/build_data.py +++ b/src/visualization/build_data.py @@ -12,7 +12,12 @@ from pathlib import Path from typing import Any -from src.visualization.paths import DATA_ROOT, PROJECT_ROOT, REPORTS_ROOT +from src.visualization.paths import ( + DATA_ROOT, + EXPERIMENTS_ROOT, + PROJECT_ROOT, + REPORTS_ROOT, +) OUTPUT_JSON = DATA_ROOT / "data.json" @@ -80,7 +85,7 @@ def parse_args() -> argparse.Namespace: parser.add_argument( "--include-workspace-runs", action="store_true", - help="Also scan workspace_*/run_*/reports outside .visualization/reports.", + help="Also scan experiments/workspace_*/run_*/reports.", ) return parser.parse_args() @@ -101,10 +106,11 @@ def is_workspace_run_report_directory(report_dir: Path) -> bool: return False return ( - len(relative.parts) == 3 - and relative.parts[0].startswith("workspace_") - and relative.parts[1].startswith("run_") - and relative.parts[2] == "reports" + len(relative.parts) == 4 + and relative.parts[0] == "experiments" + and relative.parts[1].startswith("workspace_") + and relative.parts[2].startswith("run_") + and relative.parts[3] == "reports" ) @@ -129,9 +135,11 @@ def discover_report_directories(include_workspace_runs: bool = False) -> list[Pa seen.add(report_dir) report_dirs.append(report_dir) - if include_workspace_runs: + if include_workspace_runs and EXPERIMENTS_ROOT.is_dir(): for workspace_dir in sorted( - p for p in PROJECT_ROOT.iterdir() if p.is_dir() and p.name.startswith("workspace_") + p + for p in EXPERIMENTS_ROOT.iterdir() + if p.is_dir() and p.name.startswith("workspace_") ): for run_dir in sorted( p for p in workspace_dir.iterdir() if p.is_dir() and p.name.startswith("run_") diff --git a/src/visualization/frontend/index.html b/src/visualization/frontend/index.html index f07225d9..0c34ed0c 100644 --- a/src/visualization/frontend/index.html +++ b/src/visualization/frontend/index.html @@ -20,7 +20,7 @@

Controlled agent runs, compared outcome by outcome.

Compare baseline and treatment reports discovered across AgentKernelArena, - typically under `workspace_*/run_*/reports`. Inspect score, correctness, + typically under `experiments/workspace_*/run_*/reports`. Inspect score, correctness, speedup distributions, and per-task regressions or improvements.

diff --git a/src/visualization/paths.py b/src/visualization/paths.py index c7a4db32..61e20362 100644 --- a/src/visualization/paths.py +++ b/src/visualization/paths.py @@ -15,6 +15,7 @@ def _configured_path(environment_name: str, default: Path) -> Path: PROJECT_ROOT = _configured_path("AKA_PROJECT_ROOT", MODULE_ROOT.parents[1]) +EXPERIMENTS_ROOT = PROJECT_ROOT / "experiments" RUNTIME_ROOT = _configured_path( "AKA_VISUALIZATION_RUNTIME_ROOT", PROJECT_ROOT / ".visualization" ) diff --git a/tasks/hip2hip/campaign20/paged_attention_large/Makefile b/tasks/hip2hip/campaign20/paged_attention_large/Makefile new file mode 100644 index 00000000..0f8b4f36 --- /dev/null +++ b/tasks/hip2hip/campaign20/paged_attention_large/Makefile @@ -0,0 +1,6 @@ +# Auto-generated. +all: + python3 scripts/task_runner.py compile + +clean: + rm -rf build diff --git a/tasks/hip2hip/campaign20/paged_attention_large/config.yaml b/tasks/hip2hip/campaign20/paged_attention_large/config.yaml new file mode 100644 index 00000000..c62e53b9 --- /dev/null +++ b/tasks/hip2hip/campaign20/paged_attention_large/config.yaml @@ -0,0 +1,38 @@ +source_file_path: +- src/rocm/attention.cu +target_kernel_functions: +- paged_attention_ll4mi_QKV_mfma4_kernel +- paged_attention_ll4mi_QKV_mfma16_kernel +- paged_attention_ll4mi_reduce_kernel +compile_command: +- python3 scripts/task_runner.py compile +correctness_command: +- python3 scripts/task_runner.py correctness +performance_command: +- python3 scripts/task_runner.py performance +roofline_probe_command: +- env GEAK_PERF_LARGE=1 python3 scripts/task_runner.py performance +task_type: hip2hip +platform_support: + required_arch: gfx950 + status: active + skip_reason: null +task_result_template: null +prompt: + source_code: null + instructions: "Optimize the paged_attention kernel extracted from vllm. It is registered\ + \ into the `extracted_` torch namespace via src/bindings.cpp; do not depend\ + \ on the host engine. Tests generate inputs from test_cases.json on the fly \u2014\ + \ no pre-saved golden tensors. This 'large' variant removes the legacy seq_len=32\ + \ clamp (captured cases now run at their real recorded max_seq_len with\ + \ self-consistent block tables) and the default performance run additionally\ + \ times a sweep of constructed memory-bound decode cases (disjoint KV blocks ->\ + \ real HBM streaming). Regenerate the large cases with scripts/gen_perf_cases.py." + cheatsheet: null +extracted_from: + engine: vllm + kernel_id: torch.ops._rocm_C.paged_attention + op: paged_attention + local_namespace: extracted_paged_attention + source: engine_csrc + test_cases: 11 diff --git a/tasks/hip2hip/campaign20/paged_attention_large/provenance.json b/tasks/hip2hip/campaign20/paged_attention_large/provenance.json new file mode 100644 index 00000000..a687f22f --- /dev/null +++ b/tasks/hip2hip/campaign20/paged_attention_large/provenance.json @@ -0,0 +1,51 @@ +{ + "rank": 5, + "backend": "hip", + "kernel_name": "paged_attention", + "op_category": "attention", + "max_time_pct": 11.73, + "n_source_models": 2, + "n_instances": 4, + "canonical": { + "combo": "Qwen_Qwen3.5-35B-A3B__vllm_default", + "task_dir": "vllm_paged_attention" + }, + "instances": [ + { + "model": "Qwen_Qwen3.5-35B-A3B", + "engine": "vllm", + "mode": "default", + "combo": "Qwen_Qwen3.5-35B-A3B__vllm_default", + "time_pct": 11.73, + "source_dir": "results/Qwen_Qwen3.5-35B-A3B__vllm_default/vllm_paged_attention", + "is_canonical": true + }, + { + "model": "Qwen_Qwen3.5-35B-A3B", + "engine": "vllm", + "mode": "hip", + "combo": "Qwen_Qwen3.5-35B-A3B__vllm_hip", + "time_pct": 6.62, + "source_dir": "results/Qwen_Qwen3.5-35B-A3B__vllm_hip/vllm_paged_attention", + "is_canonical": false + }, + { + "model": "Qwen_Qwen3-235B-A22B-FP8", + "engine": "vllm", + "mode": "default", + "combo": "Qwen_Qwen3-235B-A22B-FP8__vllm_default", + "time_pct": 3.41, + "source_dir": "results/Qwen_Qwen3-235B-A22B-FP8__vllm_default/vllm_paged_attention", + "is_canonical": false + }, + { + "model": "Qwen_Qwen3-235B-A22B-FP8", + "engine": "vllm", + "mode": "hip", + "combo": "Qwen_Qwen3-235B-A22B-FP8__vllm_hip", + "time_pct": 0.52, + "source_dir": "results/Qwen_Qwen3-235B-A22B-FP8__vllm_hip/vllm_paged_attention", + "is_canonical": false + } + ] +} \ No newline at end of file diff --git a/tasks/hip2hip/campaign20/paged_attention_large/scripts/_runtime.py b/tasks/hip2hip/campaign20/paged_attention_large/scripts/_runtime.py new file mode 100644 index 00000000..88c449c5 --- /dev/null +++ b/tasks/hip2hip/campaign20/paged_attention_large/scripts/_runtime.py @@ -0,0 +1,896 @@ +"""Runtime helpers shared by all auto-generated task_runner.py scripts. + +Embedded into each task at build time as ``scripts/_runtime.py`` so the runner +has zero engine dependency at execution time. Two responsibilities: + + 1. ``build_inputs(test_case, seed)``: deterministic random tensor generation + from a captured launch signature (shape + dtype + non-tensor scalars). + 2. ``compare(got, expected, dtype)``: dtype-aware allclose with sensible + tolerances for fp16 / bf16 / fp8. + +A small *reference table* lets known kernels (``rms_norm``, ``silu_and_mul``, +``rotary_embedding``, …) compute an analytic expected value purely in PyTorch. +For unknown kernels the runner falls back to a determinism check (run twice +with the same seed → byte-identical output). +""" +from __future__ import annotations + +import math +from typing import Any, Callable, Dict, List, Optional, Tuple + +import torch + + +# ---------------------------------------------------------------- dtype lookup +_DTYPE_MAP = { + "float32": torch.float32, "float": torch.float32, + "float16": torch.float16, "half": torch.float16, + "bfloat16": torch.bfloat16, + "float8_e4m3fn": torch.float8_e4m3fn, + "float8_e4m3fnuz": getattr(torch, "float8_e4m3fnuz", torch.float8_e4m3fn), + "float8_e5m2": torch.float8_e5m2, + "float8_e5m2fnuz": getattr(torch, "float8_e5m2fnuz", torch.float8_e5m2), + "int64": torch.int64, "long": torch.int64, + "int32": torch.int32, "int": torch.int32, + "int16": torch.int16, "int8": torch.int8, + "uint8": torch.uint8, "uint16": torch.uint16, + "uint32": torch.uint32, "uint64": torch.uint64, + "bool": torch.bool, +} + + +def _dtype(name: str) -> torch.dtype: + s = name.replace("torch.", "").strip() + if s in _DTYPE_MAP: + return _DTYPE_MAP[s] + return torch.float32 + + +# Argument names that need *structured* values for kernels with semantic +# constraints (typically attention / paged-cache kernels). Random ints would +# crash the kernel because, e.g., ``cu_seqlens_q`` must be a non-decreasing +# prefix-sum starting at 0. +_STRUCTURED_KEYS = { + # cumulative-seqlen prefix sums for varlen attention + "cu_seqlens_q", "cu_seqlens_k", "cu_seqlens", "qo_indptr", "kv_indptr", + "reduce_indptr", "num_kv_splits_indptr", "work_indptr", + # block / page bookkeeping + "block_table", "block_tables", "kv_indices", "kv_last_page_lens", + "kv_last_page_len", "page_indices", "page_indptr", + # MoE routing + "topk_ids", "sorted_token_ids", "sorted_expert_ids", "num_valid_ids", + "expert_ids", "topk_indices", "expert_indptr", + # general index/slot tensors + "slot_mapping", "positions", "seq_lens", "context_lens", "query_start_loc", + "query_lens", "cache_indices", "row_starts", "lengths", +} + + +def _make_structured(name: str, sig: dict, gen: torch.Generator, + device: str = "cuda", + ctx: Optional[Dict[str, Any]] = None) -> Optional[torch.Tensor]: + """Generate semantically-valid values for known structured argument names. + + ``ctx`` carries already-resolved scalar args from the same launch (eg. + ``max_seqlen_q`` from kwargs) so cu_seqlens / block_table sizes can be + sized against the kernel's own constraints. + + Returns None if we don't have a recipe for ``name``. + """ + ctx = ctx or {} + shape = tuple(sig.get("shape", [])) + dtype = _dtype(sig.get("dtype", "int32")) + nm = name.lower() + if not shape: + return None + n = shape[0] if len(shape) == 1 else shape[-1] + # query_start_loc for the *decode* paged_attention kernel is a per-sequence + # cumulative query-token index where every sequence contributes EXACTLY ONE + # query token. The kernel enforces this (it skips any sequence whose + # query_start_loc[i+1]-query_start_loc[i] != 1) and no-ops otherwise. The + # generic cu_seqlens prefix-sum below uses a chunk>1 step, which violates + # that precondition and makes the whole kernel a silent no-op. Generate a + # unit-step prefix sum [0, 1, 2, ...] so the decode kernel actually runs. + if "query_start_loc" in nm: + if n < 1: + return torch.zeros(shape, dtype=dtype, device=device) + return torch.arange(0, n, dtype=dtype, device=device).reshape(shape) + # Prefix-sum / indptr style: monotonic non-decreasing starting at 0, + # values bounded by max_seqlen_* if recorded in the same launch. + if any(k in nm for k in ("cu_seqlens", "indptr", "query_start_loc", "row_starts")): + if n < 1: + return torch.zeros(shape, dtype=dtype, device=device) + # Pick the increment so that each "sequence" has a bounded length and + # the total fits whatever batched-token tensor was captured for this + # launch. Defaults: try max_seqlen_q / max_seqlen_k from kwargs first. + max_len = None + for k, v in ctx.items(): + kk = k.lower() + if kk in ("max_seqlen_q", "max_seqlen_k", "max_extend_len", + "max_seqlen", "max_context_len") and isinstance(v, int): + if "_q" in nm and "_q" in kk: + max_len = v; break + if "_k" in nm and "_k" in kk: + max_len = v; break + if max_len is None: + max_len = v + chunk = max(1, min(int(max_len), 64) if max_len else 32) + # n includes the leading 0 entry, so we have n-1 sequences. + out = torch.arange(0, n, dtype=dtype, device=device) * chunk + return out.reshape(shape) + # seqlen / context_len: the *real* recorded context length, NOT the old + # magic-32 clamp that kept the whole KV working set resident in L2 and made + # paged_attention profile as latency-bound. We size to the launch's recorded + # ``max_seq_len`` / ``max_seqlen_*`` so decode actually streams KV from HBM. + # (block_table width / num_blocks consistency is enforced afterwards by + # _sanitize_paged_attention, which is what guards against OOB at this length.) + if any(k in nm for k in ("seq_lens", "context_lens", "query_lens", "kv_last_page_len", "lengths")): + max_len = None + for k, v in ctx.items(): + kk = k.lower() + if (kk.startswith("max_seqlen") or kk in ("max_seq_len", "max_context_len")) and isinstance(v, int): + max_len = v; break + v = int(max_len) if max_len else 1024 + return torch.full(shape, max(1, v), dtype=dtype, device=device) + # slot_mapping / page_indices / cache_indices: write targets for scattered + # cache writes. They MUST be unique — duplicate slots make the kernel + # non-deterministic across runs (concurrent threads racing the same slot). + # We sample a unique permutation. + if any(k in nm for k in ("slot_mapping", "page_indices", "cache_indices", + "kv_indices")): + total = 1 + for d in shape: + total *= d + # Pool size: at least 4× the number of slots so the permutation has + # room. Cap at 32k to keep allocation cheap. + pool = max(total * 4, 256) + pool = min(pool, 32768) + perm = torch.randperm(pool, generator=gen, dtype=dtype, device=device)[:total] + return perm.reshape(shape) + if any(k in nm for k in ("block_table", "positions")): + high = 1024 + return torch.randint(0, high, shape, dtype=dtype, device=device, generator=gen) + # MoE routing — topk_ids/sorted_expert_ids must be valid expert indices. + # We don't know num_experts here; default to 8 which is common. + if "expert" in nm or nm == "topk_ids" or nm == "topk_indices": + return torch.randint(0, 8, shape, dtype=dtype, device=device, generator=gen) + if nm == "num_valid_ids": + return torch.tensor([min(shape[0] if shape else 1, 64)] * (shape[0] if shape else 1), + dtype=dtype, device=device) + if "sorted_token_ids" in nm: + return torch.zeros(shape, dtype=dtype, device=device) + return None + + +# ---------------------------------------------------------------- tensor gen +def _tensor_from_data(sig: dict, device: str = "cuda") -> torch.Tensor: + """Materialize a tensor from captured/curated literal values. + + ``sig["data"]`` is a (possibly nested) Python list — or scalar — holding the + real values recorded at capture time (or curated for a perf workload). We + reshape to ``sig["shape"]`` and cast to ``sig["dtype"]`` so the replayed + tensor is byte-for-byte the intended one. A numel/shape mismatch raises + loudly rather than silently fabricating data. + """ + shape = tuple(sig.get("shape", [])) + dtype = _dtype(sig.get("dtype", "float32")) + t = torch.as_tensor(sig["data"], device=device) + t = t.reshape(shape) + return t.to(dtype) + + +def _make_tensor(sig: dict, gen: torch.Generator, device: str = "cuda", + name: str = "", + ctx: Optional[Dict[str, Any]] = None) -> torch.Tensor: + # Replay real values verbatim when the signature carries them. This takes + # precedence over the structured-name synthesis and random fill below, so + # persisted index tensors (block_table / seq_lens / query_start_loc, …) + # reproduce a self-consistent workload instead of being re-invented (which + # is what clamps synthetic seq_lens to ~32 and hides memory-bound behavior). + if sig.get("data") is not None: + return _tensor_from_data(sig, device) + if name and name.lower() in _STRUCTURED_KEYS: + s = _make_structured(name, sig, gen, device, ctx=ctx) + if s is not None: + return s + shape = tuple(sig.get("shape", [])) + dtype = _dtype(sig.get("dtype", "float32")) + if dtype.is_floating_point: + if dtype in ( + torch.float8_e4m3fn, + torch.float8_e5m2, + getattr(torch, "float8_e4m3fnuz", torch.float8_e4m3fn), + getattr(torch, "float8_e5m2fnuz", torch.float8_e5m2), + ): + base = (torch.randn(shape, dtype=torch.float32, device=device, generator=gen) * 0.1) + return base.clamp(-1.0, 1.0).to(dtype) + return torch.randn(shape, dtype=torch.float32, device=device, generator=gen).to(dtype) + if dtype == torch.bool: + return (torch.randint(0, 2, shape, device=device, generator=gen, dtype=torch.int32) > 0) + # integer tensor — keep values modest so they're plausible indices etc. + if not shape: + return torch.tensor(0, dtype=dtype, device=device) + high = max(2, min(1024, shape[-1] if shape[-1] > 1 else 32)) + # Clamp to dtype range so eg. uint8 doesn't overflow torch.randint's bound + # check (``to - 1 is out of bounds for unsigned char``). + if dtype == torch.uint8: + high = min(high, 256) + elif dtype == torch.int8: + high = min(high, 128) + elif dtype == torch.int16: + high = min(high, 32768) + elif dtype == torch.uint16: + high = min(high, 65536) + return torch.randint(0, high, shape, dtype=dtype, device=device, generator=gen) + + +def _decode_opaque(sig: dict) -> Any: + py_type = (sig.get("py_type") or "").lower() + rep = sig.get("repr", "") + if py_type in ("str", "int", "float", "bool", "nonetype"): + try: + import ast as _ast + return _ast.literal_eval(rep) + except Exception: + return None + if py_type == "dtype" and rep.startswith("torch."): + return _dtype(rep[6:]) + # Captured aiter enums look like ````. Extracted tasks + # call pybind functions directly, and those bindings accept the underlying + # integer rather than the Python enum object from an installed aiter package. + if "." in rep and rep.startswith("<") and ":" in rep: + try: + import re as _re + m = _re.match(r"<([\w.]+)\.(\w+):\s*(-?\d+)>", rep) + if m: + return int(m.group(3)) + except Exception: + pass + return None + + +def _arg_from_sig(sig: dict, gen: torch.Generator, name: str = "", + ctx: Optional[Dict[str, Any]] = None) -> Any: + kind = sig.get("kind", "scalar") + if kind == "tensor": + return _make_tensor(sig, gen, name=name, ctx=ctx) + if kind == "scalar": + return sig.get("value") + if kind == "seq": + return [_arg_from_sig(s, gen, ctx=ctx) for s in sig.get("items", [])] + if kind == "map": + return {k: _arg_from_sig(v, gen, name=k, ctx=ctx) for k, v in sig.get("items", {}).items()} + if kind == "opaque": + return _decode_opaque(sig) + return None + + +def _resolve_scalar_ctx(test_case: dict) -> Dict[str, Any]: + """First-pass scan: pull any scalar / opaque-scalar values out of the + launch's args+kwargs so structured tensor generators (cu_seqlens, etc.) + can size themselves against ``max_seqlen_q`` and friends. + + Both kwargs *and* named positional args are scanned. paged_attention records + ``max_seq_len`` positionally (index 13), so without the positional pass the + seq_lens generator never sees the real context length and the old code fell + back to the magic-32 clamp. + """ + ctx: Dict[str, Any] = {} + + def _stash(k: Any, v: Any) -> None: + if not isinstance(k, str) or not isinstance(v, dict): + return + if v.get("kind") == "scalar": + val = v.get("value") + if isinstance(val, (int, float, bool, str)): + ctx[k] = val + elif v.get("kind") == "opaque": + dec = _decode_opaque(v) + if isinstance(dec, (int, float, bool, str)): + ctx[k] = dec + + args_sig = test_case.get("args_sig", []) or [] + args_names = test_case.get("args_names") or [] + for i, sig in enumerate(args_sig): + name = args_names[i] if i < len(args_names) else "" + _stash(name, sig) + for k, v in (test_case.get("kwargs_sig") or {}).items(): + _stash(k, v) + return ctx + + +def build_inputs(test_case: dict, seed: int = 0xC0FFEE) -> Tuple[List[Any], Dict[str, Any]]: + """Materialize positional + keyword args from a captured launch signature. + + A first pass extracts scalar kwargs (``max_seqlen_q``, etc.) into a context + dict that the tensor builder consults — so eg. ``cu_seqlens_k`` is sized + so that its max value matches the captured ``max_seqlen_k``. + + If the test_case carries ``args_names`` (parsed from the op schema), they + are forwarded to the per-position tensor builder so structured generators + (slot_mapping → unique perm, cu_seqlens_q → prefix sum) fire even when + the kernel is called positionally. + """ + if not torch.cuda.is_available(): + raise RuntimeError("CUDA / ROCm device required to build inputs") + gen = torch.Generator(device="cuda").manual_seed(seed) + ctx = _resolve_scalar_ctx(test_case) + args_sig = test_case.get("args_sig", []) + args_names = test_case.get("args_names") or [""] * len(args_sig) + if len(args_names) < len(args_sig): + args_names = list(args_names) + [""] * (len(args_sig) - len(args_names)) + args = [ + _arg_from_sig(s, gen, name=args_names[i], ctx=ctx) + for i, s in enumerate(args_sig) + ] + kwargs: Dict[str, Any] = {} + for k, v in (test_case.get("kwargs_sig") or {}).items(): + kwargs[k] = _arg_from_sig(v, gen, name=k, ctx=ctx) + # Make the shape-only paged_attention cases self-consistent at the real + # (non-32) seq_len. perf_only cases already carry baked, consistent ``data`` + # and must NOT be touched. + if not test_case.get("perf_only"): + names_l = [n.lower() for n in args_names] + if {"block_tables", "seq_lens", "query_start_loc"}.issubset(set(names_l)): + _sanitize_paged_attention(args, names_l) + return args, kwargs + + +def _sanitize_paged_attention(args: List[Any], names_l: List[str]) -> None: + """Make a shape-only paged_attention launch internally consistent at the + real context length, so removing the magic-32 seq_len clamp can't drive the + kernel/reference out of bounds. + + The capture records tensor *shapes* but no index *values*; the generic + builder then fabricates ``seq_lens`` / ``block_tables`` / ``query_start_loc`` + independently, which only stayed in-bounds because seq_len was clamped to 32. + With the real seq_len we must couple them: + + * seq_len := min(max_seq_len, block_table_cols * block_size) (so every + referenced KV block has a column in block_table), further + capped so the disjoint block assignment fits ``num_blocks``. + * block_tables := disjoint runs seq i -> [i*nblk, (i+1)*nblk) (real + HBM-streaming layout); falls back to modulo into num_blocks + if the cache is too small to give every seq its own run. + * query_start_loc := unit-step prefix sum [0,1,2,…] (decode: 1 query + token per sequence — the kernel no-ops on any other step). + + Deterministic (no RNG), so build_inputs(seed=42) is reproducible across the + twin kernel/reference builds the correctness check relies on. + """ + idx = {n: i for i, n in enumerate(names_l)} + try: + key_cache = args[idx["key_cache"]] + block_tables = args[idx["block_tables"]] + seq_lens = args[idx["seq_lens"]] + qsl = args[idx["query_start_loc"]] + block_size = int(args[idx["block_size"]]) + max_seq_len = int(args[idx["max_seq_len"]]) + except (KeyError, IndexError, TypeError, ValueError): + return + if not isinstance(block_tables, torch.Tensor) or block_tables.dim() != 2: + return + num_blocks = int(key_cache.shape[0]) + num_seqs, cols = int(block_tables.shape[0]), int(block_tables.shape[1]) + if num_seqs < 1 or cols < 1 or block_size < 1 or num_blocks < 1: + return + dev = block_tables.device + + # seq_len bounded so every referenced block has a block_table column. + L = min(max(1, max_seq_len), cols * block_size) + nblk = math.ceil(L / block_size) # blocks read per seq + disjoint = num_seqs * nblk <= num_blocks + if not disjoint: + # Cache can't give every seq its own run; shrink L so it at least fits + # one shared pool, then assign modulo (overlap is fine for correctness). + nblk = max(1, num_blocks // max(1, num_seqs)) if num_seqs <= num_blocks else 1 + L = min(L, nblk * block_size) + + seq_lens.fill_(L) + + base = torch.arange(num_seqs, device=dev, dtype=block_tables.dtype).reshape(num_seqs, 1) * nblk + col = torch.arange(cols, device=dev, dtype=block_tables.dtype).reshape(1, cols) + bt = base + col + if not disjoint: + bt = bt % num_blocks + bt = bt.clamp_(0, num_blocks - 1) + block_tables.copy_(bt) + + # decode: exactly one query token per sequence -> unit-step prefix sum. + if isinstance(qsl, torch.Tensor) and qsl.numel() >= 1: + qsl.copy_(torch.arange(qsl.numel(), device=qsl.device, dtype=qsl.dtype)) + + +# ---------------------------------------------------------------- aiter wrapper -> pybind normalization +def _dtype_to_aiter_string(value: Any) -> Any: + if value is None: + return None + if value is torch.float16: + return "fp16" + if value is torch.bfloat16: + return "bf16" + if value is torch.float32: + return "fp32" + if isinstance(value, str): + return value + return None + + +def _sanitize_moe_routing(args: List[Any], kwargs: Dict[str, Any]) -> None: + if len(args) < 8: + return + hidden, w1, w2 = args[0], args[1], args[2] + sorted_token_ids, sorted_expert_ids, num_valid_ids = args[3], args[4], args[5] + if not all(isinstance(x, torch.Tensor) for x in (hidden, w1, sorted_token_ids, sorted_expert_ids, num_valid_ids)): + return + try: + topk = int(args[7]) + except Exception: + topk = int(kwargs.get("topk", 1) or 1) + tokens = int(hidden.shape[0]) + experts = int(w1.shape[0]) + valid = max(1, tokens * max(1, topk)) + with torch.no_grad(): + ids = torch.arange(sorted_token_ids.numel(), device=sorted_token_ids.device, + dtype=sorted_token_ids.dtype) % valid + sorted_token_ids.copy_(ids.reshape_as(sorted_token_ids)) + eids = torch.arange(sorted_expert_ids.numel(), device=sorted_expert_ids.device, + dtype=sorted_expert_ids.dtype) % max(1, experts) + sorted_expert_ids.copy_(eids.reshape_as(sorted_expert_ids)) + num_valid_ids.fill_(min(sorted_token_ids.numel(), valid)) + # Per instruction_WIP.md kernel_moe_gemm: scale inputs by 1/sqrt(K) so + # bf16 doesn't overflow with random non-pre-shuffled weights. The + # underlying GEMM is K-reduction; without scaling the accumulator hits + # ~3e38 and atomic-add ordering produces non-deterministic overflow. + if hidden.dtype.is_floating_point: + K = int(hidden.shape[-1]) + scale = 1.0 / max(1.0, math.sqrt(float(K))) + hidden.mul_(scale) + if isinstance(w1, torch.Tensor) and w1.dtype.is_floating_point: + w1.mul_(scale) + if isinstance(w2, torch.Tensor) and w2.dtype.is_floating_point: + w2.mul_(scale) + + +def normalize_aiter_call(py_fn_name: str, fc_name: str, + args: List[Any], kwargs: Dict[str, Any]) -> Tuple[List[Any], Dict[str, Any]]: + """Translate captured Python-wrapper arguments to the direct pybind ABI. + + The recorder observes public aiter wrappers such as ``flash_attn_varlen_func`` + and ``ck_moe_stage1_fwd``. Extracted tasks intentionally call the local + pybind module directly, so a few wrapper-only keyword names and defaults need + to be normalized before invocation. + """ + op = py_fn_name or fc_name + args = list(args) + kwargs = dict(kwargs) + + if op in ("ck_moe_stage1_fwd", "ck_moe_stage2_fwd", "ck_moe_stage1", "ck_moe_stage2"): + if "use_non_temporal_load" in kwargs and "non_temporal_load" not in kwargs: + kwargs["non_temporal_load"] = kwargs.pop("use_non_temporal_load") + elif "use_non_temporal_load" in kwargs: + kwargs.pop("use_non_temporal_load", None) + # Wrapper-only kwargs absent from the underlying pybind ABI. + if "dtype" in kwargs and "dst_type" not in kwargs: + kwargs["dst_type"] = _dtype_to_aiter_string(kwargs.pop("dtype")) + else: + kwargs.pop("dtype", None) + if len(args) > 8: + kwargs.pop("kernelName", None) + elif kwargs.get("kernelName") is None: + kwargs["kernelName"] = "" + if "dst_type" in kwargs: + kwargs["dst_type"] = _dtype_to_aiter_string(kwargs["dst_type"]) + # Positional captures include dst_type just before is_shuffled. + if len(args) >= 18: + args[17] = _dtype_to_aiter_string(args[17]) + _sanitize_moe_routing(args, kwargs) + return args, kwargs + + if op == "flash_attn_varlen_func" or fc_name in ("fmha_v3_varlen_fwd", "mha_varlen_fwd"): + # The captured wrapper sometimes passes (q, k, v) positionally and the + # rest as a sparse kwargs dict; the underlying pybind requires every + # non-trailing parameter. Promote positional q/k/v back into kwargs + # then reissue everything by name with explicit defaults. + if args: + kwargs.setdefault("q", args[0] if len(args) > 0 else None) + kwargs.setdefault("k", args[1] if len(args) > 1 else None) + kwargs.setdefault("v", args[2] if len(args) > 2 else None) + args = [] + window = kwargs.pop("window_size", [-1, -1]) + if window is None: + window = [-1, -1] + window_left = int(window[0]) if len(window) > 0 else -1 + window_right = int(window[1]) if len(window) > 1 else -1 + # The captured max_seqlen_q/k may be the model's max-sequence cap + # (e.g. 65536) rather than the actual maximum present in the + # synthesized cu_seqlens_q. Recompute from cu_seqlens since the + # synthesized inputs may be much smaller — passing an oversized + # max_seqlen makes the CK kernel index past the actual tensor. + cu_q = kwargs.get("cu_seqlens_q") + cu_k = kwargs.get("cu_seqlens_k") + captured_max_q = int(kwargs.get("max_seqlen_q", 0) or 0) + captured_max_k = int(kwargs.get("max_seqlen_k", 0) or 0) + actual_max_q = captured_max_q + actual_max_k = captured_max_k + try: + if cu_q is not None and hasattr(cu_q, "diff"): + actual_max_q = int(cu_q.diff().max().item()) + if cu_k is not None and hasattr(cu_k, "diff"): + actual_max_k = int(cu_k.diff().max().item()) + except Exception: + pass + max_q = max(1, min(captured_max_q, actual_max_q) if captured_max_q else actual_max_q) + max_k = max(1, min(captured_max_k, actual_max_k) if captured_max_k else actual_max_k) + kwargs = { + "q": kwargs.get("q"), + "k": kwargs.get("k"), + "v": kwargs.get("v"), + "cu_seqlens_q": cu_q, + "cu_seqlens_k": cu_k, + "max_seqlen_q": max_q, + "max_seqlen_k": max_k, + "min_seqlen_q": int(kwargs.get("min_seqlen_q", 0)), + "dropout_p": float(kwargs.get("dropout_p", 0.0)), + "softmax_scale": float(kwargs.get("softmax_scale", 1.0)), + "logits_soft_cap": float(kwargs.get("logits_soft_cap", 0.0)), + "zero_tensors": bool(kwargs.get("zero_tensors", False)), + "is_causal": bool(kwargs.get("is_causal", kwargs.get("causal", False))), + "window_size_left": window_left, + "window_size_right": window_right, + "return_softmax_lse": bool(kwargs.get("return_softmax_lse", + kwargs.get("return_lse", False))), + "return_dropout_randval": bool(kwargs.get("return_dropout_randval", False)), + "how_v3_bf16_cvt": int(kwargs.get("how_v3_bf16_cvt", 1)), + "out": kwargs.get("out"), + "block_table": kwargs.get("block_table"), + "bias": kwargs.get("bias"), + "alibi_slopes": kwargs.get("alibi_slopes"), + "gen": kwargs.get("gen"), + "cu_seqlens_q_padded": kwargs.get("cu_seqlens_q_padded"), + "cu_seqlens_k_padded": kwargs.get("cu_seqlens_k_padded"), + } + return args, kwargs + + return args, kwargs + + +def normalize_aiter_output(py_fn_name: str, value: Any) -> Any: + if py_fn_name == "flash_attn_varlen_func" and isinstance(value, (list, tuple)): + # AITER returns (out, softmax_lse, dropout_mask, rng_state). With + # dropout disabled the rng_state buffer is not semantically meaningful + # and may contain run-to-run garbage; compare the observable outputs. + return tuple(value[:2]) + return value + + +# ---------------------------------------------------------------- comparison +def _tol_for(dtype: torch.dtype) -> Tuple[float, float]: + if dtype in (torch.float16, torch.bfloat16): + return 5e-2, 5e-2 + if dtype in (torch.float8_e4m3fn, torch.float8_e5m2, + getattr(torch, "float8_e4m3fnuz", torch.float8_e4m3fn), + getattr(torch, "float8_e5m2fnuz", torch.float8_e5m2)): + return 0.5, 0.5 + if dtype.is_floating_point: + return 1e-3, 1e-3 + return 0, 0 + + +def compare(got: Any, expected: Any) -> Optional[str]: + """Return None if equal-within-tolerance, else a short diff message.""" + if isinstance(got, torch.Tensor) and isinstance(expected, torch.Tensor): + if got.shape != expected.shape: + return f"shape mismatch: got {tuple(got.shape)} vs {tuple(expected.shape)}" + atol, rtol = _tol_for(got.dtype) + a = got.detach().to(torch.float32).cpu() + b = expected.detach().to(torch.float32).cpu() + if not torch.isfinite(a).all(): + count = int((~torch.isfinite(a)).sum().item()) + return f"kernel output contains {count} non-finite value(s)" + if not torch.isfinite(b).all(): + count = int((~torch.isfinite(b)).sum().item()) + return f"reference output contains {count} non-finite value(s)" + if not torch.allclose(a, b, atol=atol, rtol=rtol): + diff = (a - b).abs().max().item() + return f"max abs diff {diff:.4g} > atol={atol}" + return None + if isinstance(got, (list, tuple)) and isinstance(expected, (list, tuple)): + if len(got) != len(expected): + return f"len mismatch: {len(got)} vs {len(expected)}" + for i, (g, e) in enumerate(zip(got, expected)): + err = compare(g, e) + if err: + return f"[{i}]: {err}" + return None + if got == expected: + return None + return f"value mismatch" + + +# ---------------------------------------------------------------- references +# Each reference takes the SAME (args, kwargs) as the kernel and returns the +# expected output. For in-place kernels the reference returns a tensor that +# the runner will compare against the (now-mutated) input. The runner picks +# the entry by op_name; if missing, falls back to determinism check. +ReferenceFn = Callable[[List[Any], Dict[str, Any]], Any] + + +def _ref_rms_norm(args: list, kwargs: dict): + # signature: (out, input, weight, epsilon) — vLLM `_C.rms_norm` writes to out + out, inp, weight, eps = args[0], args[1], args[2], args[3] + var = inp.to(torch.float32).pow(2).mean(dim=-1, keepdim=True) + return (inp.to(torch.float32) * torch.rsqrt(var + eps)).to(inp.dtype) * weight.to(inp.dtype) + + +def _ref_fused_add_rms_norm(args: list, kwargs: dict): + # (input, residual, weight, epsilon) — in-place: input = rmsnorm(input+residual) + inp, res, weight, eps = args[0], args[1], args[2], args[3] + s = (inp + res).to(torch.float32) + var = s.pow(2).mean(dim=-1, keepdim=True) + return (s * torch.rsqrt(var + eps)).to(inp.dtype) * weight.to(inp.dtype) + + +def _ref_silu_and_mul(args: list, kwargs: dict): + # (out, input) — out = silu(input[..., :H/2]) * input[..., H/2:] + out, inp = args[0], args[1] + a, b = inp.chunk(2, dim=-1) + return torch.nn.functional.silu(a.to(torch.float32)).to(inp.dtype) * b + + +def _ref_gelu_and_mul(args: list, kwargs: dict): + out, inp = args[0], args[1] + a, b = inp.chunk(2, dim=-1) + return torch.nn.functional.gelu(a.to(torch.float32)).to(inp.dtype) * b + + +def _ref_gelu_tanh_and_mul(args: list, kwargs: dict): + out, inp = args[0], args[1] + a, b = inp.chunk(2, dim=-1) + return torch.nn.functional.gelu(a.to(torch.float32), approximate="tanh").to(inp.dtype) * b + + +def _ref_gelu_quick(args: list, kwargs: dict): + out, inp = args[0], args[1] + f = inp.to(torch.float32) + return (f * torch.sigmoid(1.702 * f)).to(inp.dtype) + + +def _ref_gelu_new(args: list, kwargs: dict): + out, inp = args[0], args[1] + f = inp.to(torch.float32) + return (0.5 * f * (1 + torch.tanh(math.sqrt(2 / math.pi) * (f + 0.044715 * f.pow(3))))).to(inp.dtype) + + +def _ref_gelu_fast(args: list, kwargs: dict): + return _ref_gelu_new(args, kwargs) + + +def _ref_mul_and_silu(args: list, kwargs: dict): + out, inp = args[0], args[1] + a, b = inp.chunk(2, dim=-1) + return (a.to(torch.float32) * torch.nn.functional.silu(b.to(torch.float32))).to(inp.dtype) + + +def _ref_fatrelu_and_mul(args: list, kwargs: dict): + out, inp = args[0], args[1] + threshold = float(kwargs.get("threshold", args[2] if len(args) > 2 else 0.0)) + a, b = inp.chunk(2, dim=-1) + mask = (a.to(torch.float32) > threshold).to(inp.dtype) + return (a * mask) * b + + +# References are keyed by ``":"`` because vLLM and AITER +# expose ops with the same short name but different positional layouts (e.g. +# vLLM ``fused_add_rms_norm(input, residual, weight, eps)`` vs. AITER +# ``rmsnorm2d_fwd_with_add(out, input, residual, out_residual, weight, eps)``). +# The runner passes its source prefix; unknown keys fall back to a determinism +# check, which is a safe no-op rather than producing NaN garbage. +# [BugA-fix] verified numerical references (added by benchmark maintainers) +def _ref_moe_sum(args, kwargs): + inp, out = args[0], args[1] + return {1: inp.to(torch.float32).sum(dim=1).to(out.dtype)} + + +def _ref__apply_rotary(x, cos, sin, rot_dim, is_neox): + xf = x.clone() + embed = rot_dim // 2 + c = cos[:, None, :] + s = sin[:, None, :] + if is_neox: + x1 = xf[..., :embed]; x2 = xf[..., embed:rot_dim] + xf[..., :embed] = x1 * c - x2 * s + xf[..., embed:rot_dim] = x2 * c + x1 * s + else: + x1 = xf[..., 0:rot_dim:2]; x2 = xf[..., 1:rot_dim:2] + xf[..., 0:rot_dim:2] = x1 * c - x2 * s + xf[..., 1:rot_dim:2] = x2 * c + x1 * s + return xf + + +def _ref_rotary_embedding(args, kwargs): + positions, query, key, head_size, cos_sin_cache, is_neox = args[0], args[1], args[2], args[3], args[4], args[5] + rot_dim = cos_sin_cache.shape[-1] + embed = rot_dim // 2 + pos = positions.reshape(-1).long() + cs = cos_sin_cache.index_select(0, pos).to(torch.float32) + cos, sin = cs[:, :embed], cs[:, embed:rot_dim] + T = pos.shape[0] + out = {1: _ref__apply_rotary(query.to(torch.float32).view(T, -1, head_size), cos, sin, rot_dim, is_neox).view_as(query).to(query.dtype)} + if isinstance(key, torch.Tensor): + out[2] = _ref__apply_rotary(key.to(torch.float32).view(T, -1, head_size), cos, sin, rot_dim, is_neox).view_as(key).to(key.dtype) + return out + + +def _ref_reshape_and_cache(args, kwargs): + key, value, key_cache, value_cache, slot_mapping = args[0], args[1], args[2], args[3], args[4] + kc, vc = key_cache.clone(), value_cache.clone() + T, nh, hs = key.shape + nb, nh2, hsx, bs, x = key_cache.shape + slot = slot_mapping.reshape(-1).long() + for t in range(T): + s = int(slot[t]) + if s < 0: + continue + b, off = s // bs, s % bs + kc[b, :, :, off, :] = key[t].to(kc.dtype).view(nh, hsx, x) + vc[b, :, :, off] = value[t].to(vc.dtype) + return {2: kc, 3: vc} + + +def _ref_reshape_and_cache_flash(args, kwargs): + key, value, key_cache, value_cache, slot_mapping = args[0], args[1], args[2], args[3], args[4] + kc, vc = key_cache.clone(), value_cache.clone() + T = key.shape[0] + nb, bs, nh2, hs2 = key_cache.shape + slot = slot_mapping.reshape(-1).long() + for t in range(T): + s = int(slot[t]) + if s < 0: + continue + b, off = s // bs, s % bs + kc[b, off, :, :] = key[t].to(kc.dtype) + vc[b, off, :, :] = value[t].to(vc.dtype) + return {2: kc, 3: vc} + + +def _ref_concat_and_cache_mla(args, kwargs): + kv_c, k_pe, kv_cache, slot_mapping = args[0], args[1], args[2], args[3] + c = kv_cache.clone() + T, lora = kv_c.shape + pe = k_pe.shape[1] + nb, bs, entry = kv_cache.shape + slot = slot_mapping.reshape(-1).long() + for t in range(T): + s = int(slot[t]) + if s < 0: + continue + b, off = s // bs, s % bs + c[b, off, :lora] = kv_c[t].to(c.dtype) + c[b, off, lora:lora + pe] = k_pe[t].to(c.dtype) + return {2: c} + + +def _ref_topk_softmax(args, kwargs): + topk_weights, gating = args[0], args[3] + renorm = bool(args[4]) if len(args) > 4 else False + topk = topk_weights.shape[1] + probs = torch.softmax(gating.to(torch.float32), dim=-1) + vals, _ = probs.topk(topk, dim=-1) + if renorm: + vals = vals / vals.sum(dim=-1, keepdim=True) + # Only the top-k weights are numerically meaningful; expert-index ordering / + # tie-breaking is implementation-defined and intentionally not compared. + return {0: vals.to(topk_weights.dtype)} + + +def _ref_wvSplitK(args, kwargs): + in_a, in_b = args[0], args[1] + bias = args[2] if len(args) > 2 else None + out = in_b.to(torch.float32) @ in_a.to(torch.float32).T + if isinstance(bias, torch.Tensor): + out = out + bias.to(torch.float32) + return out.to(in_b.dtype) + + +def _ref_paged_attention(args, kwargs): + out, query, key_cache, value_cache = args[0], args[4], args[5], args[6] + num_kv_heads, scale = int(args[7]), float(args[8]) + block_tables, seq_lens, query_start_loc, block_size = args[9], args[10], args[11], int(args[12]) + num_seqs = block_tables.shape[0] + num_heads, head_size = query.shape[1], query.shape[2] + gqa = num_heads // num_kv_heads + nb, nkv, hsx, bs, x = key_cache.shape + kc, vc, q = key_cache.float(), value_cache.float(), query.float() + bt, sl = block_tables.long(), seq_lens.long() + qsl = query_start_loc.long() if isinstance(query_start_loc, torch.Tensor) else None + out_ref = out.float().clone() + kv_of_head = torch.arange(num_heads, device=q.device) // gqa + for si in range(num_seqs): + row = int(qsl[si]) if qsl is not None else si + L = int(sl[si]) + if L <= 0: + continue + pos = torch.arange(L, device=q.device) + blk = bt[si, pos // block_size] + off = pos % block_size + ar = torch.arange(L, device=q.device) + K = kc[blk][ar, :, :, off, :].reshape(L, nkv, hsx * x) + V = vc[blk][ar, :, :, off] + Kh, Vh = K[:, kv_of_head, :], V[:, kv_of_head, :] + scores = scale * torch.einsum('hd,phd->hp', q[row], Kh) + attn = torch.softmax(scores, dim=-1) + out_ref[row] = torch.einsum('hp,phd->hd', attn, Vh) + return {0: out_ref.to(out.dtype)} + + +REFERENCES: Dict[str, ReferenceFn] = { + # [BugA-fix] verified numerical references + "vllm:moe_sum": _ref_moe_sum, + "vllm:rotary_embedding": _ref_rotary_embedding, + "vllm:reshape_and_cache": _ref_reshape_and_cache, + "vllm:reshape_and_cache_flash": _ref_reshape_and_cache_flash, + "vllm:concat_and_cache_mla": _ref_concat_and_cache_mla, + "vllm:topk_softmax": _ref_topk_softmax, + "vllm:wvSplitK": _ref_wvSplitK, + "vllm:paged_attention": _ref_paged_attention, + + # vLLM _C ops — runner passes ``vllm:`` + "vllm:rms_norm": _ref_rms_norm, + "vllm:fused_add_rms_norm": _ref_fused_add_rms_norm, + "vllm:silu_and_mul": _ref_silu_and_mul, + "vllm:gelu_and_mul": _ref_gelu_and_mul, + "vllm:gelu_tanh_and_mul": _ref_gelu_tanh_and_mul, + "vllm:gelu_quick": _ref_gelu_quick, + "vllm:gelu_new": _ref_gelu_new, + "vllm:gelu_fast": _ref_gelu_fast, + "vllm:mul_and_silu": _ref_mul_and_silu, + "vllm:fatrelu_and_mul": _ref_fatrelu_and_mul, + # SGLang sgl_kernel ops mostly mirror vLLM's signatures. + "sglang:rms_norm": _ref_rms_norm, + "sglang:fused_add_rms_norm": _ref_fused_add_rms_norm, + "sglang:silu_and_mul": _ref_silu_and_mul, + "sglang:gelu_and_mul": _ref_gelu_and_mul, + "sglang:gelu_tanh_and_mul": _ref_gelu_tanh_and_mul, + "sglang:gelu_quick": _ref_gelu_quick, +} + + +def reference_for(op_name: str, source: str = "") -> Optional[ReferenceFn]: + """Pick a reference for ``op_name``. If ``source`` is given (``vllm`` / + ``aiter`` / ``triton`` / ``sglang``) it is used as a namespace prefix to + disambiguate same-named ops with different signatures.""" + if source: + fn = REFERENCES.get(f"{source}:{op_name}") + if fn is not None: + return fn + return REFERENCES.get(op_name) + + +# ---------------------------------------------------------------- output capture +# Many vLLM/AITER kernels mutate their first ``out`` arg rather than returning +# a value. The runner inspects the first positional tensor argument's bytes +# before/after the call to detect this and treat it as the output. + +def snapshot(args: list) -> list: + return [a.detach().clone() if isinstance(a, torch.Tensor) else None for a in args] + + +def detect_output(pre: list, post: list, ret: Any) -> Any: + """Pick the most plausible output for comparison. + + 1. If the kernel returned a tensor (or tuple), use that. + 2. Otherwise look for the first positional tensor that changed in-place. + 3. Else None. + """ + if isinstance(ret, torch.Tensor) or isinstance(ret, (list, tuple)) and ret and isinstance(ret[0], torch.Tensor): + return ret + for i, (b, a) in enumerate(zip(pre, post)): + if a is None or b is None: + continue + try: + if not torch.equal(b.to(torch.float32).cpu(), a.to(torch.float32).cpu()): + return a + except Exception: + continue + return None diff --git a/tasks/hip2hip/campaign20/paged_attention_large/scripts/gen_perf_cases.py b/tasks/hip2hip/campaign20/paged_attention_large/scripts/gen_perf_cases.py new file mode 100644 index 00000000..34df0da5 --- /dev/null +++ b/tasks/hip2hip/campaign20/paged_attention_large/scripts/gen_perf_cases.py @@ -0,0 +1,154 @@ +#!/usr/bin/env python3 +"""Inject memory-bound performance cases into ``test_cases.json``. + +The captured cases record real tensor *shapes* but no *values*; the runtime then +synthesizes structured index tensors (block_table / seq_lens / query_start_loc) +and clamps seq_lens to ~32 for crash-safety. At that length the whole KV cache +stays resident in MI300X's last-level cache, so paged_attention profiles as +latency-bound and a roofline reads <2% HBM — unrepresentative of real decode. + +This script appends a ``perf_only`` case carrying *real, self-consistent* values +for the structured index tensors (each sequence owns a disjoint run of KV blocks, +seq_lens = L), sized so the KV working set far exceeds the cache and the kernel +must stream from HBM. The big float tensors stay shape-only — their values don't +affect bandwidth/timing, only the index structure does. + + ``seq_lens`` : every sequence has context length L + ``block_table`` : seq i owns blocks [i*nblk, (i+1)*nblk) -> disjoint HBM + ``query_start_loc`` : unit-step prefix sum [0,1,2,…] (1 query token/seq = decode) + +Replay is handled by ``_runtime._make_tensor`` honoring the ``"data"`` field. + + python3 scripts/gen_perf_cases.py # default: L=4096, S=2048, GQA 16:1 + +Idempotent: existing ``perf_only`` cases are removed before re-appending. +""" +import argparse +import copy +import json +import math +import os + +HERE = os.path.dirname(os.path.abspath(__file__)) +TEST_CASES = os.path.join(os.path.dirname(HERE), "test_cases.json") + +# Kernel layout constants (match the captured cases + the compiled template: +# paged_attention_ll4mi_QKV_mfma16_kernel<…, D, PART, …, X, …>). +D = 128 # head size +X = 8 # key_cache packing factor (D // X = 16) +BLK = 16 # block_size (KV cache page size) +PART = 256 # PARTITION_SIZE -> exp_sums/max_logits/tmp_out 3rd dim = ceil(L/PART) + + +def _t(dtype, shape, data=None): + sig = {"kind": "tensor", "dtype": dtype, "shape": list(shape), "device": "cuda:0"} + if data is not None: + sig["data"] = data + return sig + + +def build_case(S: int, L: int, H: int, KVH: int, template: dict) -> dict: + """Construct one perf_only case from a captured case template. + + Scalar / opaque args (kv_cache_dtype, mfma_type, alibi_slopes, …) are copied + verbatim from the template so their exact recorded encoding is preserved; + only tensor shapes, the coupled scalars, and the structured ``data`` change. + """ + nblk = math.ceil(L / BLK) # blocks per sequence + num_blocks = S * nblk # disjoint -> forces real HBM traffic + P = math.ceil(L / PART) # partitions per sequence + + block_table = [[i * nblk + j for j in range(nblk)] for i in range(S)] + seq_lens = [L] * S + query_start_loc = list(range(S + 1)) # 1 query token per sequence (decode) + + tc = copy.deepcopy(template) + a = tc["args_sig"] + a[0] = _t("bfloat16", [S, H, D]) # out + a[1] = _t("float32", [S, H, P]) # exp_sums + a[2] = _t("float32", [S, H, P]) # max_logits + a[3] = _t("bfloat16", [S, H, P, D]) # tmp_out + a[4] = _t("bfloat16", [S, H, D]) # query + a[5] = _t("bfloat16", [num_blocks, KVH, D // X, BLK, X]) # key_cache + a[6] = _t("bfloat16", [num_blocks, KVH, D, BLK]) # value_cache + a[7]["value"] = KVH # num_kv_heads + a[8]["value"] = 1.0 / math.sqrt(D) # scale + a[9] = _t("int32", [S, nblk], data=block_table) # block_tables + a[10] = _t("int32", [S], data=seq_lens) # seq_lens + a[11] = _t("int32", [S + 1], data=query_start_loc) # query_start_loc + a[12]["value"] = BLK # block_size + a[13]["value"] = L # max_seq_len + a[16] = _t("float32", [], data=1.0) # k_scale + a[17] = _t("float32", [], data=1.0) # v_scale + # a[14] alibi_slopes(None), a[15] kv_cache_dtype, a[18] fp8_out_scale(None), + # a[19] mfma_type are left exactly as captured. + + kv_gb = 2 * num_blocks * KVH * (D // X) * BLK * X * 2 / 1e9 # K+V, bf16 + tc["test_case_id"] = f"perf_L{L}_S{S}_h{H}kv{KVH}" + tc["perf_only"] = True + tc["count"] = 1 + tc["params_repr"] = { + "S_seqs": S, "ctx_len": L, "out_len": 1, "heads": H, "kv_heads": KVH, + "head_size": D, "block_size": BLK, "partition": PART, + "kv_alloc_gb": round(kv_gb, 1), "note": "decode; disjoint KV blocks -> HBM-streaming", + } + return tc + + +# Default memory-bound sweep tested by DEFAULT (no env var needed). This is the +# roofline_probe.py (S, L) sweep plus the original single (S=2048, L=4096) case, +# all GQA 16:1 (H=16, KVH=1). Every entry carries baked, self-consistent index +# ``data`` (disjoint KV blocks) so the kernel actually streams KV from HBM. +DEFAULT_CONFIGS = [ + # (S sequences, L context length) + (1024, 1024), (1024, 4096), + (4096, 2048), (8192, 2048), + (2048, 8192), (8192, 8192), + (16384, 4096), + (2048, 4096), # the original gen_perf_cases default +] +DEFAULT_H = 16 +DEFAULT_KVH = 1 + + +def main(): + ap = argparse.ArgumentParser() + ap.add_argument("-S", type=int, default=None, help="single override: concurrent sequences (batch)") + ap.add_argument("-L", type=int, default=None, help="single override: context / KV length per sequence") + ap.add_argument("-H", type=int, default=DEFAULT_H, help="query heads") + ap.add_argument("--kvh", type=int, default=DEFAULT_KVH, help="kv heads (GQA group = H/kvh)") + args = ap.parse_args() + + # -S/-L given -> single case; otherwise emit the full default sweep. + if args.S is not None or args.L is not None: + configs = [(args.S or 2048, args.L or 4096)] + else: + seen, configs = set(), [] # de-dup on (S, L), preserve order + for S, L in DEFAULT_CONFIGS: + if (S, L) not in seen: + seen.add((S, L)); configs.append((S, L)) + + with open(TEST_CASES) as f: + cases = json.load(f) + captured = [c for c in cases if not c.get("perf_only")] + + # Pick a captured template whose GQA layout matches the request so scalar / + # opaque arg encodings carry over cleanly. + def heads_of(c): + q = c["args_sig"][4] + kv = c["args_sig"][7] + return (q.get("shape", [0, 0])[1], kv.get("value")) + tmpl = next((c for c in captured if heads_of(c) == (args.H, args.kvh)), captured[0]) + + new_cases = [build_case(S, L, args.H, args.kvh, tmpl) for S, L in configs] + out = captured + new_cases + with open(TEST_CASES, "w") as f: + json.dump(out, f, indent=2) + print(f"[gen_perf_cases] wrote {len(captured)} captured + {len(new_cases)} perf_only:") + for c in new_cases: + print(f" {c['test_case_id']:28s} KV~{c['params_repr']['kv_alloc_gb']}GB") + print(f" -> {TEST_CASES}") + + +if __name__ == "__main__": + main() diff --git a/tasks/hip2hip/campaign20/paged_attention_large/scripts/gen_regime_cases.py b/tasks/hip2hip/campaign20/paged_attention_large/scripts/gen_regime_cases.py new file mode 100644 index 00000000..ad17deee --- /dev/null +++ b/tasks/hip2hip/campaign20/paged_attention_large/scripts/gen_regime_cases.py @@ -0,0 +1,47 @@ +#!/usr/bin/env python3 +"""Set test_cases.json perf_only cases to the arena WORKLOAD REGIME. + +Regime: decode-style paged_attention. ctx_len (KV length) L = 1024, one query +token per sequence (q_len=1). Concurrency B in {2,32,64} maps to S = num_seqs +(batch). Model dims kept from the captured base case: head_size=128, block_size=16, +H=16 query heads, KVH=1 (GQA 16:1), X=8, PARTITION_SIZE=256. + +Emits exactly 3 perf_only cases with ids c2,c32,c64. Captured (correctness) +cases are preserved untouched. +""" +import json, os, sys +HERE = os.path.dirname(os.path.abspath(__file__)) +sys.path.insert(0, HERE) +import gen_perf_cases as g + +TEST_CASES = os.path.join(os.path.dirname(HERE), "test_cases.json") + +L = 1024 # ctx / KV length per sequence (= arena seqlen) +H = g.DEFAULT_H # 16 query heads (captured) +KVH = g.DEFAULT_KVH # 1 kv head (captured GQA 16:1) +REGIME = [2, 32, 64] + +def main(): + with open(TEST_CASES) as f: + cases = json.load(f) + captured = [c for c in cases if not c.get("perf_only")] + # template with matching GQA layout + def heads_of(c): + q = c["args_sig"][4]; kv = c["args_sig"][7] + return (q.get("shape", [0, 0])[1], kv.get("value")) + tmpl = next((c for c in captured if heads_of(c) == (H, KVH)), captured[0]) + new = [] + for B in REGIME: + tc = g.build_case(B, L, H, KVH, tmpl) + tc["test_case_id"] = f"c{B}" + tc["params_repr"]["concurrency_B"] = B + new.append(tc) + out = captured + new + with open(TEST_CASES, "w") as f: + json.dump(out, f, indent=2) + print(f"[gen_regime_cases] {len(captured)} captured + {len(new)} perf_only:") + for c in new: + print(f" {c['test_case_id']:6s} S={c['params_repr']['S_seqs']} L={c['params_repr']['ctx_len']} KV~{c['params_repr']['kv_alloc_gb']}GB") + +if __name__ == "__main__": + main() diff --git a/tasks/hip2hip/campaign20/paged_attention_large/scripts/roofline_probe.py b/tasks/hip2hip/campaign20/paged_attention_large/scripts/roofline_probe.py new file mode 100644 index 00000000..19a936ac --- /dev/null +++ b/tasks/hip2hip/campaign20/paged_attention_large/scripts/roofline_probe.py @@ -0,0 +1,128 @@ +#!/usr/bin/env python3 +"""Probe how high paged_attention can push HBM bandwidth / GFLOP/s. + +Builds *consistent* large decode inputs directly (bypassing the synthetic +build_inputs, which caps seq_lens to ~32), giving each sequence its own KV +blocks so the kernel must actually stream KV from HBM. Times with CUDA events +and reports achieved bandwidth and compute vs the empirical MI300 peaks. + + python3 scripts/roofline_probe.py +""" +import math +import os +import sys + +sys.path.insert(0, os.path.dirname(os.path.abspath(__file__))) +import task_runner as T # noqa: E402 (module-level: chdir + path setup) + +import torch # noqa: E402 + +# Empirical MI300 peaks (from rocprof-compute roofline.csv on this box). +PEAK_HBM_GBPS = 4170.53 +PEAK_BF16_MFMA_GFLOPS = 510391.03 + +H = 16 # query heads +KVH = 1 # kv heads (GQA ratio 16) +D = 128 # head size +X = 8 # kv-cache packing factor (head_size/x = 16) +BLK = 16 # block size +PART = 256 # PARTITION_SIZE +DTYPE = torch.bfloat16 + + +def build(S: int, L: int): + dev = "cuda" + nblk_per_seq = math.ceil(L / BLK) + num_blocks = S * nblk_per_seq # disjoint blocks -> real HBM traffic + P = math.ceil(L / PART) + + q = torch.randn(S, H, D, dtype=DTYPE, device=dev) + key_cache = torch.randn(num_blocks, KVH, D // X, BLK, X, dtype=DTYPE, device=dev) + value_cache = torch.randn(num_blocks, KVH, D, BLK, dtype=DTYPE, device=dev) + out = torch.empty(S, H, D, dtype=DTYPE, device=dev) + exp_sums = torch.empty(S, H, P, dtype=torch.float32, device=dev) + max_logits = torch.empty(S, H, P, dtype=torch.float32, device=dev) + tmp_out = torch.empty(S, H, P, D, dtype=DTYPE, device=dev) + # each seq i owns blocks [i*nblk_per_seq, (i+1)*nblk_per_seq) + block_tables = ( + torch.arange(num_blocks, dtype=torch.int32, device=dev).reshape(S, nblk_per_seq) + ) + seq_lens = torch.full((S,), L, dtype=torch.int32, device=dev) + query_start_loc = torch.arange(S + 1, dtype=torch.int32, device=dev) # 1 token/seq + k_scale = torch.ones(1, dtype=torch.float32, device=dev) + v_scale = torch.ones(1, dtype=torch.float32, device=dev) + + args = (out, exp_sums, max_logits, tmp_out, q, key_cache, value_cache, + KVH, 1.0 / math.sqrt(D), block_tables, seq_lens, query_start_loc, + BLK, L, None, "auto", k_scale, v_scale, None, "f16") + # bytes that MUST come from HBM (KV is >L2 by construction): K + V, bf16 + kv_bytes = S * L * KVH * D * 2 * 2 + # attention flops: QK^T (2*L*D) + softmax*V (2*L*D), per (seq, query head) + flops = S * H * (2 * L * D + 2 * L * D) + alloc_gb = (key_cache.nbytes + value_cache.nbytes) / 1e9 + return args, kv_bytes, flops, alloc_gb, num_blocks + + +def bench(op, args, iters=30, warmup=10): + for _ in range(warmup): + op(*args) + torch.cuda.synchronize() + starts = [torch.cuda.Event(enable_timing=True) for _ in range(iters)] + ends = [torch.cuda.Event(enable_timing=True) for _ in range(iters)] + for i in range(iters): + starts[i].record(); op(*args); ends[i].record() + torch.cuda.synchronize() + return sum(s.elapsed_time(e) for s, e in zip(starts, ends)) / iters # ms + + +def profile_one(op, S: int, L: int, iters: int): + """Run one config in a steady loop (no timing) for an external profiler.""" + args, kv_bytes, flops, alloc_gb, num_blocks = build(S, L) + print(f"[profile] S={S} L={L} KValloc={alloc_gb:.1f}GB num_blocks={num_blocks} iters={iters}", + flush=True) + for _ in range(iters): + op(*args) + torch.cuda.synchronize() + + +def main(): + op = T._load_op() + + # Single-config profiling mode (for rocprof-compute): set PROBE_S / PROBE_L. + if os.environ.get("PROBE_S"): + S = int(os.environ["PROBE_S"]) + L = int(os.environ.get("PROBE_L", "2048")) + iters = int(os.environ.get("PROBE_ITERS", "20")) + profile_one(op, S, L, iters) + return + + configs = [ + (1024, 1024), (1024, 4096), + (4096, 2048), (8192, 2048), + (2048, 8192), (8192, 8192), + (16384, 4096), + ] + print(f"{'S(seqs)':>8} {'ctx_L':>6} {'KValloc_GB':>10} {'ms':>8} " + f"{'HBM_GB/s':>9} {'HBM%':>6} {'GFLOP/s':>9} {'BF16%':>6}") + for S, L in configs: + try: + args, kv_bytes, flops, alloc_gb, _ = build(S, L) + if alloc_gb > 80: + print(f"{S:>8} {L:>6} {alloc_gb:>10.1f} (skip: alloc too large)") + continue + ms = bench(op, args) + t = ms / 1e3 + bw = kv_bytes / t / 1e9 + gf = flops / t / 1e9 + print(f"{S:>8} {L:>6} {alloc_gb:>10.1f} {ms:>8.3f} " + f"{bw:>9.1f} {100*bw/PEAK_HBM_GBPS:>5.1f}% " + f"{gf:>9.1f} {100*gf/PEAK_BF16_MFMA_GFLOPS:>5.2f}%") + del args + torch.cuda.empty_cache() + except Exception as e: + print(f"{S:>8} {L:>6} ERROR: {str(e)[:80]}") + torch.cuda.empty_cache() + + +if __name__ == "__main__": + main() diff --git a/tasks/hip2hip/campaign20/paged_attention_large/scripts/task_runner.py b/tasks/hip2hip/campaign20/paged_attention_large/scripts/task_runner.py new file mode 100644 index 00000000..522a24ab --- /dev/null +++ b/tasks/hip2hip/campaign20/paged_attention_large/scripts/task_runner.py @@ -0,0 +1,338 @@ +#!/usr/bin/env python3 +"""Auto-generated task runner for vllm_paged_attention (HIP). + +Inputs are generated each run from the shape/dtype signatures in +test_cases.json — no .pt files are loaded. +""" +import sys, os, json, argparse, glob +sys.path.insert(0, os.path.dirname(os.path.abspath(__file__))) +import _runtime as rt + +# >>> AKA-GENERATED: shared CUDA-graph benchmark helpers - edit src/tools/perf/vllm_cuda_graph_block.py then run `make sync-perf-helpers` >>> +def _measure_cuda_event_fallback(*args, **kwargs): + raise RuntimeError( + "CUDA-graph benchmark helpers were not materialized. " + "Run this task through AgentKernelArena so setup_workspace() can inject " + "src/tools/perf/vllm_cuda_graph_block.py into the workspace." + ) + + +def _benchmark_cuda_graph_or_events(*args, **kwargs): + raise RuntimeError( + "CUDA-graph benchmark helpers were not materialized. " + "Run this task through AgentKernelArena so setup_workspace() can inject " + "src/tools/perf/vllm_cuda_graph_block.py into the workspace." + ) +# <<< AKA-GENERATED <<< + + +TASK_DIR = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) +os.chdir(TASK_DIR) +TASK_NAME = "hip2hip/campaign20/paged_attention_large" +NAMESPACE = "extracted_paged_attention" # extracted_ +OP_NAME = "paged_attention" +REF_SOURCE = "vllm" # "vllm" | "sglang" — used to pick the correct reference + +SRC_DIR = os.path.join(TASK_DIR, "src") +BUILD_DIR = os.path.join(TASK_DIR, "build") +TEST_CASES = os.path.join(TASK_DIR, "test_cases.json") + +# Lock the offload arch to the runtime device's gfx so torch.utils.cpp_extension +# does not try to compile for every ROCm target (RDNA gfx1100 etc. break on +# vllm's cub/bf16 templates). Override by setting PYTORCH_ROCM_ARCH externally. +def _detect_gfx(): + try: + import torch + if torch.cuda.is_available(): + arch = torch.cuda.get_device_properties(0).gcnArchName + return arch.split(":")[0] # e.g. "gfx942:sramecc+:xnack-" -> "gfx942" + except Exception: + pass + return "gfx942" + + +def _build(): + import torch + from torch.utils.cpp_extension import load + # vLLM/SGLang csrc IS written for portable CUDA-style code that depends + # on torch.cpp_extension's hipify pass to rewrite ``cudaStream_t`` → + # ``hipStream_t`` etc. Leave hipify enabled here. (AITER tasks use a + # different runner that disables hipify because its CK template trees + # don't survive the rewrite.) + # Lock the offload-arch list to the runtime device. Torch's default + # populates PYTORCH_ROCM_ARCH with EVERY arch its build knows + # (gfx90a;gfx942;...;gfx1100;gfx1101;...), and the RDNA targets fail on + # vllm's cub bf16 templates. Override with the live device's gfx unless + # the user has set a single non-default arch explicitly. + cur = os.environ.get("PYTORCH_ROCM_ARCH", "") + detected = _detect_gfx() + if not cur or ";" in cur or " " in cur or any( + a in cur for a in ("gfx10", "gfx11", "gfx12") + ): + os.environ["PYTORCH_ROCM_ARCH"] = detected + os.makedirs(BUILD_DIR, exist_ok=True) + # Remove stale build locks left behind by a previously *killed* build. + # torch.cpp_extension.load() guards the build dir with a FileBaton + # (``/lock``) and ninja with ``/.ninja_lock``. Neither is + # released if the holding process is SIGKILL/SIGTERM'd mid-compile (which + # the GEAK harness does on profiling/timeout budget caps). A leftover + # ``lock`` makes every later load() spin forever in baton.wait() + # (time.sleep poll loop -> 0% CPU, idle GPU, "hung" build). Each GEAK slot + # owns its own build dir (single writer), so any lock present at the start + # of a fresh _build() is by definition stale -> safe to clear. + for _lk in ("lock", ".ninja_lock"): + _p = os.path.join(BUILD_DIR, _lk) + if os.path.exists(_p): + try: + os.unlink(_p) + print(f"[task_runner] removed stale build lock: {_p}", flush=True) + except Exception: + pass + # Strip any in-place hipify residue from prior builds. torch.cpp_extension + # generates ``foo.hip`` and ``foo_hip.cuh`` next to the original ``foo.cu`` + # / ``foo.cuh``. If both linger between builds the linker gets duplicate + # symbols. The .hip / *_hip.* files are regenerable from the .cu source. + import shutil as _sh + for stale in glob.glob(os.path.join(SRC_DIR, "**", "*.hip"), recursive=True): + try: os.unlink(stale) + except Exception: pass + for stale in glob.glob(os.path.join(SRC_DIR, "**", "*_hip.*"), recursive=True): + try: os.unlink(stale) + except Exception: pass + sources = [] + for ext in ("cu", "cpp"): + sources.extend(sorted(glob.glob(os.path.join(SRC_DIR, "**", f"*.{ext}"), recursive=True))) + if not sources: + raise RuntimeError("no sources under src/") + # is_python_module=False because bindings.cpp registers via TORCH_LIBRARY + # rather than defining a PyInit_* symbol — the op is reached via + # ``torch.ops..`` after load() returns. + # ``torch.utils.cpp_extension`` defaults to passing + # ``-D__HIP_NO_HALF_OPERATORS__=1 -D__HIP_NO_HALF_CONVERSIONS__=1`` to the + # HIP compiler. PyTorch needs those to keep its own ``c10::Half`` happy, + # but vLLM/SGLang csrc uses raw ``__half2 += __half2`` (in + # ``type_convert.cuh`` and similar) which is precisely the operator that + # flag disables. Their official setup.py build doesn't set the flag; we + # have to undefine it here to make the same source compile through + # ``cpp_extension.load``. + cflags = [ + "-U__HIP_NO_HALF_OPERATORS__", + "-U__HIP_NO_HALF_CONVERSIONS__", + # vLLM/SGLang csrc gates large blocks of fp8 / bf16 helpers behind + # ``ENABLE_FP8`` / ``ENABLE_BF16`` (see e.g. + # quantization/w8a8/fp8/amd/quant_utils.cuh's ``namespace fp8``). + # Their CMake/setup.py define these for the ROCm build; cpp_extension + # doesn't, so the symbols vanish and the dependent .cu files fail with + # "no member named 'scaled_vec_conversion' in namespace 'vllm::fp8'". + "-DENABLE_FP8", + "-DENABLE_BF16", + # AITER fp8 sources also expect this to choose the FNUZ vs E4M3 layout. + "-DHIP_FP8_TYPE_FNUZ", + ] + # vLLM's csrc references ``TORCH_HIP_VERSION`` (a macro vLLM's setup.py + # would normally define from ``HIP_VERSION_MAJOR/MINOR``). Recreate it + # here so the same source compiles under cpp_extension.load. + try: + import torch + hv = getattr(torch.version, "hip", None) + if hv: + major, _, rest = hv.partition(".") + minor = rest.split(".", 1)[0] if rest else "0" + # Torch's TORCH_HIP_VERSION convention: HIP_VERSION_MAJOR*100 + + # HIP_VERSION_MINOR (e.g. 702 for ROCm 7.2). NOT *10000 — that + # extra factor broke ``#if TORCH_HIP_VERSION >= 12090`` style + # gates inherited from CUDA-versioned source (after hipify + # rewrites CUDA_VERSION → TORCH_HIP_VERSION) by making 70200 + # satisfy a CUDA 12.9+ comparison. + torch_hip_version = int(major) * 100 + int(minor) + cflags.append("-DTORCH_HIP_VERSION=" + str(torch_hip_version)) + except Exception: + cflags.append("-DTORCH_HIP_VERSION=702") + return load( + name=NAMESPACE, + sources=sources, + extra_include_paths=[ + SRC_DIR, + os.path.join(SRC_DIR, "core"), + os.path.join(SRC_DIR, "include"), + ], + extra_cflags=cflags, + extra_cuda_cflags=cflags, + verbose=False, + with_cuda=True, + is_python_module=False, + build_directory=BUILD_DIR, + ) + + +def _load_op(): + import torch + _build() + return getattr(getattr(torch.ops, NAMESPACE), OP_NAME) + + +def _test_cases(): + if not os.path.isfile(TEST_CASES): + return [] + with open(TEST_CASES) as f: + return json.load(f) + + +def run_compile(): + try: + _build() + return True, None + except Exception as e: + return False, str(e) + + +def run_correctness(): + import torch + try: + fn = _load_op() + except Exception as e: + return False, f"build failed: {e}" + cases = _test_cases() + if not cases: + return True, "no recorded launch signatures (compile-only check)" + # Drop empty signatures (kernel was registered but never called with args + # in the captured run — eg. `_C` ops shadowed by AITER's wrappers). Also drop + # perf_only cases: those are large HBM-streaming workloads for roofline, and + # the pure-PyTorch reference loops per-sequence in Python — unusably slow at + # S=2048/L=4096. Correctness stays on the captured cases. + cases = [c for c in cases + if (c.get("args_sig") or c.get("kwargs_sig")) and not c.get("perf_only")] + if not cases: + return True, "all recorded launches had empty signatures (compile-only check)" + ref = rt.reference_for(OP_NAME, REF_SOURCE) + for tc in cases: + try: + args1, kwargs1 = rt.build_inputs(tc, seed=42) + pre = rt.snapshot(args1) + ret1 = fn(*args1, **kwargs1) + if ref is not None: + args_r, kwargs_r = rt.build_inputs(tc, seed=42) + expected = ref(args_r, kwargs_r) + if isinstance(expected, dict): + # [BugA-fix] reference returns {arg_index: expected_tensor} + # for in-place / multi-output kernels; compare each mutated arg. + err = None + for _idx, _exp in expected.items(): + err = rt.compare(args1[_idx], _exp) + if err: + err = f"arg{_idx}: {err}" + break + else: + out1 = rt.detect_output(pre, args1, ret1) + if out1 is None: + return False, f"{tc['test_case_id']}: kernel produced no observable output" + err = rt.compare(out1, expected) + if err: + return False, f"{tc['test_case_id']}: vs reference: {err}" + else: + # determinism check: same seed, same input, byte-identical out + out1 = rt.detect_output(pre, args1, ret1) + if out1 is None: + return False, f"{tc['test_case_id']}: kernel produced no observable output" + args2, kwargs2 = rt.build_inputs(tc, seed=42) + pre2 = rt.snapshot(args2) + ret2 = fn(*args2, **kwargs2) + out2 = rt.detect_output(pre2, args2, ret2) + if out2 is None: + return False, f"{tc['test_case_id']}: repeat produced no observable output" + err = rt.compare(out1, out2) + if err: + return False, f"{tc['test_case_id']}: non-deterministic: {err}" + except Exception as e: + return False, f"{tc['test_case_id']}: kernel raised {e}" + return True, None + + +def run_performance(): + import signal + try: + fn = _load_op() + except Exception as exc: + return [{"test_case_id": "load", "execution_time_ms": -1.0, + "params": {"error": str(exc)[:160]}}] + allcases = [c for c in _test_cases() if c.get("args_sig") or c.get("kwargs_sig")] + # Default runs ALL cases (captured + perf_only), so the large memory-bound + # perf case is timed by default. GEAK_PERF_LARGE=1 narrows to ONLY perf_only + # cases (clean single-case roofline via roofline_probe_command). + _v = os.environ.get("GEAK_PERF_LARGE", "") + only_perf = _v not in ("", "0", "false", "False", "no") + cases = [c for c in allcases if c.get("perf_only")] if only_perf else allcases + # Iteration counts are env-tunable. GEAK passes GEAK_BENCHMARK_ITERATIONS + # (the harness has no --iterations flag), so honour it. The defaults follow + # the arena's accepted 10-warmup/100-measured-sample methodology. + def _env_int(name, default): + try: return max(1, int(os.environ.get(name, "") or default)) + except Exception: return default + n_iter = _env_int("GEAK_BENCHMARK_ITERATIONS", 100) + n_warmup = _env_int("GEAK_BENCHMARK_WARMUP", 10) + # Hard per-case wall-clock guard so one pathological case (huge inputs / + # runaway kernel) can't hang the whole benchmark. Best-effort: SIGALRM + # fires in the main thread and lands when control returns to Python. + case_timeout = _env_int("GEAK_PERF_CASE_TIMEOUT", 120) + + class _CaseTimeout(Exception): + pass + + def _on_alarm(signum, frame): + raise _CaseTimeout() + + have_alarm = hasattr(signal, "SIGALRM") + if have_alarm: + signal.signal(signal.SIGALRM, _on_alarm) + + out = [] + for tc in cases: + if have_alarm: + signal.alarm(case_timeout) + try: + args, kwargs = rt.build_inputs(tc, seed=42) + avg, metadata = _benchmark_cuda_graph_or_events( + lambda: fn(*args, **kwargs), + warmup=n_warmup, + repetition=n_iter, + ) + out.append({"test_case_id": tc["test_case_id"], + "execution_time_ms": avg, **metadata, + "params": tc.get("params_repr", {})}) + except _CaseTimeout: + out.append({"test_case_id": tc["test_case_id"], "execution_time_ms": -1.0, "params": {"error": f"case exceeded {case_timeout}s timeout"}}) + except Exception as e: + out.append({"test_case_id": tc["test_case_id"], "execution_time_ms": -1.0, "params": {"error": str(e)[:120]}}) + finally: + if have_alarm: + signal.alarm(0) + return out + + +def main(): + ap = argparse.ArgumentParser() + ap.add_argument("mode", choices=["compile", "correctness", "performance"]) + args = ap.parse_args() + os.makedirs(BUILD_DIR, exist_ok=True) + if args.mode == "compile": + ok, err = run_compile() + json.dump({"status": "ok" if ok else "fail", "error": err}, open(os.path.join(BUILD_DIR, "compile_report.json"), "w")) + print(f"Compilation: {'PASS' if ok else 'FAIL'}") + if err: print("Error:", err) + sys.exit(0 if ok else 1) + if args.mode == "correctness": + ok, err = run_correctness() + json.dump({"status": "ok" if ok else "fail", "error": err}, open(os.path.join(BUILD_DIR, "correctness_report.json"), "w")) + print(f"Correctness: {'PASS' if ok else 'FAIL'}") + if err: print("Error:", err) + sys.exit(0 if ok else 1) + cases = run_performance() + json.dump({"test_cases": cases}, open(os.path.join(BUILD_DIR, "performance_report.json"), "w"), indent=2) + for c in cases: + print(f"Performance: {c['execution_time_ms']:.4f} ms ({c['test_case_id']})") + valid = bool(cases) and all(c.get("execution_time_ms", -1) > 0 for c in cases) + sys.exit(0 if valid else 1) + + +if __name__ == "__main__": + main() diff --git a/tasks/hip2hip/campaign20/paged_attention_large/src/attention/attention_dtypes.h b/tasks/hip2hip/campaign20/paged_attention_large/src/attention/attention_dtypes.h new file mode 100644 index 00000000..64f86381 --- /dev/null +++ b/tasks/hip2hip/campaign20/paged_attention_large/src/attention/attention_dtypes.h @@ -0,0 +1,7 @@ +#pragma once + +#include "attention_generic.cuh" +#include "dtype_float16.cuh" +#include "dtype_float32.cuh" +#include "dtype_bfloat16.cuh" +#include "dtype_fp8.cuh" diff --git a/tasks/hip2hip/campaign20/paged_attention_large/src/attention/attention_generic.cuh b/tasks/hip2hip/campaign20/paged_attention_large/src/attention/attention_generic.cuh new file mode 100644 index 00000000..62409c0c --- /dev/null +++ b/tasks/hip2hip/campaign20/paged_attention_large/src/attention/attention_generic.cuh @@ -0,0 +1,65 @@ +/* + * Adapted from + * https://github.com/NVIDIA/FasterTransformer/blob/release/v5.3_tag/src/fastertransformer/kernels/decoder_masked_multihead_attention_utils.h + * Copyright (c) 2023, The vLLM team. + * Copyright (c) 2020-2023, NVIDIA CORPORATION. All rights reserved. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +#pragma once + +#include + +namespace vllm { + +// A vector type to store Q, K, V elements. +template +struct Vec {}; + +// A vector type to store FP32 accumulators. +template +struct FloatVec {}; + +// Template vector operations. +template +inline __device__ Acc mul(A a, B b); + +template +inline __device__ float sum(T v); + +template +inline __device__ float dot(T a, T b) { + return sum(mul(a, b)); +} + +template +inline __device__ float dot(T a, T b) { + return sum(mul(a, b)); +} + +template +inline __device__ void zero(T& dst) { + constexpr int WORDS = sizeof(T) / 4; + union { + T raw; + uint32_t words[WORDS]; + } tmp; + +#pragma unroll + for (int ii = 0; ii < WORDS; ++ii) { + tmp.words[ii] = 0u; + } + dst = tmp.raw; +} + +} // namespace vllm diff --git a/tasks/hip2hip/campaign20/paged_attention_large/src/attention/dtype_bfloat16.cuh b/tasks/hip2hip/campaign20/paged_attention_large/src/attention/dtype_bfloat16.cuh new file mode 100644 index 00000000..97a25baa --- /dev/null +++ b/tasks/hip2hip/campaign20/paged_attention_large/src/attention/dtype_bfloat16.cuh @@ -0,0 +1,463 @@ +/* + * Adapted from + * https://github.com/NVIDIA/FasterTransformer/blob/release/v5.3_tag/src/fastertransformer/kernels/decoder_masked_multihead_attention/decoder_masked_multihead_attention_template.hpp + * and + * https://github.com/NVIDIA/FasterTransformer/blob/release/v5.3_tag/src/fastertransformer/kernels/decoder_masked_multihead_attention_utils.h + * Copyright (c) 2023, The vLLM team. + * Copyright (c) 2020-2023, NVIDIA CORPORATION. All rights reserved. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +#pragma once + +#include "attention_generic.cuh" +#include "dtype_float32.cuh" + +#ifndef USE_ROCM + #include + #include +#else + #include + #include + +typedef __hip_bfloat162 __nv_bfloat162; +typedef __hip_bfloat16 __nv_bfloat16; +#endif + +#include + +namespace vllm { + +// Define custom BF16 vector data types. +struct bf16_4_t { + __nv_bfloat162 x; + __nv_bfloat162 y; +}; + +struct bf16_8_t { + __nv_bfloat162 x; + __nv_bfloat162 y; + __nv_bfloat162 z; + __nv_bfloat162 w; +}; + +// BF16 vector types for Q, K, V. +template <> +struct Vec<__nv_bfloat16, 1> { + using Type = __nv_bfloat16; +}; +template <> +struct Vec<__nv_bfloat16, 2> { + using Type = __nv_bfloat162; +}; +template <> +struct Vec<__nv_bfloat16, 4> { + using Type = bf16_4_t; +}; +template <> +struct Vec<__nv_bfloat16, 8> { + using Type = bf16_8_t; +}; + +// FP32 accumulator vector types corresponding to Vec. +template <> +struct FloatVec<__nv_bfloat16> { + using Type = float; +}; +template <> +struct FloatVec<__nv_bfloat162> { + using Type = float2; +}; +template <> +struct FloatVec { + using Type = Float4_; +}; +template <> +struct FloatVec { + using Type = Float8_; +}; + +// Utility functions for type conversions. +inline __device__ float2 bf1622float2(const __nv_bfloat162 val) { +#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ < 800 + assert(false); +#else + return __bfloat1622float2(val); +#endif + __builtin_unreachable(); // Suppress missing return statement warning +} + +inline __device__ __nv_bfloat162 bf162bf162(const __nv_bfloat16 val) { +#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ < 800 + assert(false); +#else + return __bfloat162bfloat162(val); +#endif + __builtin_unreachable(); // Suppress missing return statement warning +} + +// Vector addition. +inline __device__ __nv_bfloat16 add(__nv_bfloat16 a, __nv_bfloat16 b) { +#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ < 800 + assert(false); +#else + #ifndef USE_ROCM + return a + b; + #else + return __hadd(a, b); + #endif +#endif + __builtin_unreachable(); // Suppress missing return statement warning +} + +inline __device__ __nv_bfloat162 add(__nv_bfloat162 a, __nv_bfloat162 b) { +#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ < 800 + assert(false); +#else + return __hadd2(a, b); +#endif + __builtin_unreachable(); // Suppress missing return statement warning +} + +inline __device__ bf16_4_t add(bf16_4_t a, bf16_4_t b) { + bf16_4_t c; + c.x = add(a.x, b.x); + c.y = add(a.y, b.y); + return c; +} + +inline __device__ bf16_8_t add(bf16_8_t a, bf16_8_t b) { + bf16_8_t c; + c.x = add(a.x, b.x); + c.y = add(a.y, b.y); + c.z = add(a.z, b.z); + c.w = add(a.w, b.w); + return c; +} + +inline __device__ float2 add(__nv_bfloat162 a, float2 fb) { + float2 fa = bf1622float2(a); + return add(fa, fb); +} + +inline __device__ Float4_ add(bf16_4_t a, Float4_ fb) { + Float4_ fc; + fc.x = add(a.x, fb.x); + fc.y = add(a.y, fb.y); + return fc; +} + +inline __device__ Float8_ add(bf16_8_t a, Float8_ fb) { + Float8_ fc; + fc.x = add(a.x, fb.x); + fc.y = add(a.y, fb.y); + fc.z = add(a.z, fb.z); + fc.w = add(a.w, fb.w); + return fc; +} + +// Vector multiplication. +template <> +inline __device__ __nv_bfloat16 mul(__nv_bfloat16 a, __nv_bfloat16 b) { +#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ < 800 + assert(false); +#else + return __hmul(a, b); +#endif + __builtin_unreachable(); // Suppress missing return statement warning +} + +template <> +inline __device__ __nv_bfloat162 mul(__nv_bfloat162 a, __nv_bfloat162 b) { +#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ < 800 + assert(false); +#else + return __hmul2(a, b); +#endif + __builtin_unreachable(); // Suppress missing return statement warning +} + +template <> +inline __device__ __nv_bfloat162 mul(__nv_bfloat16 a, __nv_bfloat162 b) { + return mul<__nv_bfloat162, __nv_bfloat162, __nv_bfloat162>(bf162bf162(a), b); +} + +template <> +inline __device__ bf16_4_t mul(bf16_4_t a, bf16_4_t b) { + bf16_4_t c; + c.x = mul<__nv_bfloat162, __nv_bfloat162, __nv_bfloat162>(a.x, b.x); + c.y = mul<__nv_bfloat162, __nv_bfloat162, __nv_bfloat162>(a.y, b.y); + return c; +} + +template <> +inline __device__ bf16_4_t mul(__nv_bfloat16 a, bf16_4_t b) { + __nv_bfloat162 s = bf162bf162(a); + bf16_4_t c; + c.x = mul<__nv_bfloat162, __nv_bfloat162, __nv_bfloat162>(s, b.x); + c.y = mul<__nv_bfloat162, __nv_bfloat162, __nv_bfloat162>(s, b.y); + return c; +} + +template <> +inline __device__ bf16_8_t mul(bf16_8_t a, bf16_8_t b) { + bf16_8_t c; + c.x = mul<__nv_bfloat162, __nv_bfloat162, __nv_bfloat162>(a.x, b.x); + c.y = mul<__nv_bfloat162, __nv_bfloat162, __nv_bfloat162>(a.y, b.y); + c.z = mul<__nv_bfloat162, __nv_bfloat162, __nv_bfloat162>(a.z, b.z); + c.w = mul<__nv_bfloat162, __nv_bfloat162, __nv_bfloat162>(a.w, b.w); + return c; +} + +template <> +inline __device__ bf16_8_t mul(__nv_bfloat16 a, bf16_8_t b) { + __nv_bfloat162 s = bf162bf162(a); + bf16_8_t c; + c.x = mul<__nv_bfloat162, __nv_bfloat162, __nv_bfloat162>(s, b.x); + c.y = mul<__nv_bfloat162, __nv_bfloat162, __nv_bfloat162>(s, b.y); + c.z = mul<__nv_bfloat162, __nv_bfloat162, __nv_bfloat162>(s, b.z); + c.w = mul<__nv_bfloat162, __nv_bfloat162, __nv_bfloat162>(s, b.w); + return c; +} + +template <> +inline __device__ float mul(__nv_bfloat16 a, __nv_bfloat16 b) { + float fa = __bfloat162float(a); + float fb = __bfloat162float(b); + return fa * fb; +} + +template <> +inline __device__ float2 mul(__nv_bfloat162 a, __nv_bfloat162 b) { + float2 fa = bf1622float2(a); + float2 fb = bf1622float2(b); + return mul(fa, fb); +} + +template <> +inline __device__ float2 mul(__nv_bfloat16 a, __nv_bfloat162 b) { + return mul(bf162bf162(a), b); +} + +template <> +inline __device__ Float4_ mul(bf16_4_t a, bf16_4_t b) { + Float4_ fc; + fc.x = mul(a.x, b.x); + fc.y = mul(a.y, b.y); + return fc; +} + +template <> +inline __device__ Float4_ mul(__nv_bfloat16 a, bf16_4_t b) { + __nv_bfloat162 s = bf162bf162(a); + Float4_ fc; + fc.x = mul(s, b.x); + fc.y = mul(s, b.y); + return fc; +} + +template <> +inline __device__ Float8_ mul(bf16_8_t a, bf16_8_t b) { + Float8_ fc; + fc.x = mul(a.x, b.x); + fc.y = mul(a.y, b.y); + fc.z = mul(a.z, b.z); + fc.w = mul(a.w, b.w); + return fc; +} + +template <> +inline __device__ Float8_ mul(__nv_bfloat16 a, bf16_8_t b) { + __nv_bfloat162 s = bf162bf162(a); + Float8_ fc; + fc.x = mul(s, b.x); + fc.y = mul(s, b.y); + fc.z = mul(s, b.z); + fc.w = mul(s, b.w); + return fc; +} + +// Vector fused multiply-add. +inline __device__ __nv_bfloat162 fma(__nv_bfloat162 a, __nv_bfloat162 b, + __nv_bfloat162 c) { +#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ < 800 + assert(false); +#else + return __hfma2(a, b, c); +#endif + __builtin_unreachable(); // Suppress missing return statement warning +} + +inline __device__ __nv_bfloat162 fma(__nv_bfloat16 a, __nv_bfloat162 b, + __nv_bfloat162 c) { +#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ < 800 + assert(false); +#else + return __hfma2(bf162bf162(a), b, c); +#endif + __builtin_unreachable(); // Suppress missing return statement warning +} + +inline __device__ bf16_4_t fma(bf16_4_t a, bf16_4_t b, bf16_4_t c) { + bf16_4_t d; + d.x = fma(a.x, b.x, c.x); + d.y = fma(a.y, b.y, c.y); + return d; +} + +inline __device__ bf16_4_t fma(__nv_bfloat16 a, bf16_4_t b, bf16_4_t c) { + __nv_bfloat162 s = bf162bf162(a); + bf16_4_t d; + d.x = fma(s, b.x, c.x); + d.y = fma(s, b.y, c.y); + return d; +} + +inline __device__ bf16_8_t fma(bf16_8_t a, bf16_8_t b, bf16_8_t c) { + bf16_8_t d; + d.x = fma(a.x, b.x, c.x); + d.y = fma(a.y, b.y, c.y); + d.z = fma(a.z, b.z, c.z); + d.w = fma(a.w, b.w, c.w); + return d; +} + +inline __device__ bf16_8_t fma(__nv_bfloat16 a, bf16_8_t b, bf16_8_t c) { + __nv_bfloat162 s = bf162bf162(a); + bf16_8_t d; + d.x = fma(s, b.x, c.x); + d.y = fma(s, b.y, c.y); + d.z = fma(s, b.z, c.z); + d.w = fma(s, b.w, c.w); + return d; +} + +inline __device__ float fma(__nv_bfloat16 a, __nv_bfloat16 b, float fc) { + return __bfloat162float(a) * __bfloat162float(b) + fc; +} + +inline __device__ float2 fma(__nv_bfloat162 a, __nv_bfloat162 b, float2 fc) { + float2 fa = bf1622float2(a); + float2 fb = bf1622float2(b); + return fma(fa, fb, fc); +} + +inline __device__ float2 fma(__nv_bfloat16 a, __nv_bfloat162 b, float2 fc) { + return fma(bf162bf162(a), b, fc); +} + +inline __device__ Float4_ fma(bf16_4_t a, bf16_4_t b, Float4_ fc) { + Float4_ fd; + fd.x = fma(a.x, b.x, fc.x); + fd.y = fma(a.y, b.y, fc.y); + return fd; +} + +inline __device__ Float4_ fma(__nv_bfloat16 a, bf16_4_t b, Float4_ fc) { + __nv_bfloat162 s = bf162bf162(a); + Float4_ fd; + fd.x = fma(s, b.x, fc.x); + fd.y = fma(s, b.y, fc.y); + return fd; +} + +inline __device__ Float8_ fma(bf16_8_t a, bf16_8_t b, Float8_ fc) { + Float8_ fd; + fd.x = fma(a.x, b.x, fc.x); + fd.y = fma(a.y, b.y, fc.y); + fd.z = fma(a.z, b.z, fc.z); + fd.w = fma(a.w, b.w, fc.w); + return fd; +} + +inline __device__ Float8_ fma(__nv_bfloat16 a, bf16_8_t b, Float8_ fc) { + __nv_bfloat162 s = bf162bf162(a); + Float8_ fd; + fd.x = fma(s, b.x, fc.x); + fd.y = fma(s, b.y, fc.y); + fd.z = fma(s, b.z, fc.z); + fd.w = fma(s, b.w, fc.w); + return fd; +} + +// Vector sum. +template <> +inline __device__ float sum(__nv_bfloat16 v) { + return __bfloat162float(v); +} + +template <> +inline __device__ float sum(__nv_bfloat162 v) { + float2 vf = bf1622float2(v); + return vf.x + vf.y; +} + +template <> +inline __device__ float sum(bf16_4_t v) { + return sum(v.x) + sum(v.y); +} + +template <> +inline __device__ float sum(bf16_8_t v) { + return sum(v.x) + sum(v.y) + sum(v.z) + sum(v.w); +} + +// From float32 to bfloat16. +inline __device__ void from_float(__nv_bfloat16& dst, float src) { + dst = __float2bfloat16(src); +} + +inline __device__ void from_float(__nv_bfloat162& dst, float2 src) { +#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ < 800 + assert(false); +#else + dst = __float22bfloat162_rn(src); +#endif +} + +inline __device__ void from_float(bf16_4_t& dst, Float4_ src) { +#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ < 800 + assert(false); +#else + dst.x = __float22bfloat162_rn(src.x); + dst.y = __float22bfloat162_rn(src.y); +#endif +} + +inline __device__ void from_float(bf16_8_t& dst, Float8_ src) { +#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ < 800 + assert(false); +#else + dst.x = __float22bfloat162_rn(src.x); + dst.y = __float22bfloat162_rn(src.y); + dst.z = __float22bfloat162_rn(src.z); + dst.w = __float22bfloat162_rn(src.w); +#endif +} + +// From bfloat16 to float32. +inline __device__ float to_float(__nv_bfloat16 u) { + return __bfloat162float(u); +} + +// Zero-out a variable. +inline __device__ void zero(__nv_bfloat16& dst) { +#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ < 800 + assert(false); +#else + // Same as CUDART_ZERO_BF16 introduced in CUDA 12.2. + dst = __ushort_as_bfloat16((unsigned short)0x0000U); +#endif +} + +} // namespace vllm diff --git a/tasks/hip2hip/campaign20/paged_attention_large/src/attention/dtype_float16.cuh b/tasks/hip2hip/campaign20/paged_attention_large/src/attention/dtype_float16.cuh new file mode 100644 index 00000000..3a1815f0 --- /dev/null +++ b/tasks/hip2hip/campaign20/paged_attention_large/src/attention/dtype_float16.cuh @@ -0,0 +1,504 @@ +/* + * Adapted from + * https://github.com/NVIDIA/FasterTransformer/blob/release/v5.3_tag/src/fastertransformer/kernels/decoder_masked_multihead_attention/decoder_masked_multihead_attention_template.hpp + * and + * https://github.com/NVIDIA/FasterTransformer/blob/release/v5.3_tag/src/fastertransformer/kernels/decoder_masked_multihead_attention_utils.h + * Copyright (c) 2023, The vLLM team. + * Copyright (c) 2020-2023, NVIDIA CORPORATION. All rights reserved. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +#pragma once + +#include "attention_generic.cuh" +#include "dtype_float32.cuh" + +#ifdef USE_ROCM + #include +#endif + +#include + +namespace vllm { + +// FP16 vector types for Q, K, V. +template <> +struct Vec { + using Type = uint16_t; +}; +template <> +struct Vec { + using Type = uint32_t; +}; +template <> +struct Vec { + using Type = uint2; +}; +template <> +struct Vec { + using Type = uint4; +}; + +// FP32 accumulator vector types corresponding to Vec. +template <> +struct FloatVec { + using Type = float; +}; +template <> +struct FloatVec { + using Type = float2; +}; +template <> +struct FloatVec { + using Type = Float4_; +}; +template <> +struct FloatVec { + using Type = Float8_; +}; + +// Utility functions for type conversions. +inline __device__ uint32_t h0_h0(uint16_t a) { +#ifndef USE_ROCM + uint32_t b; + asm volatile("mov.b32 %0, {%1, %1};" : "=r"(b) : "h"(a)); + return b; +#else + union { + uint32_t u32; + uint16_t u16[2]; + } tmp; + tmp.u16[0] = a; + tmp.u16[1] = a; + return tmp.u32; +#endif +} + +inline __device__ float half_to_float(uint16_t h) { + float f; +#ifndef USE_ROCM + asm volatile("cvt.f32.f16 %0, %1;\n" : "=f"(f) : "h"(h)); +#else + asm volatile("v_cvt_f32_f16 %0, %1;" : "=v"(f) : "v"(h)); +#endif + return f; +} + +inline __device__ float2 half2_to_float2(uint32_t v) { +#ifndef USE_ROCM + uint16_t lo, hi; + asm volatile("mov.b32 {%0, %1}, %2;\n" : "=h"(lo), "=h"(hi) : "r"(v)); + return make_float2(half_to_float(lo), half_to_float(hi)); +#else + union { + uint32_t u32; + uint16_t u16[2]; + } tmp; + tmp.u32 = v; + float2 ret; + ret.x = half_to_float(tmp.u16[0]); + ret.y = half_to_float(tmp.u16[1]); + return ret; +#endif +} + +inline __device__ uint16_t float_to_half(float f) { + union { + uint32_t u32; + uint16_t u16[2]; + } tmp; +#ifndef USE_ROCM + asm volatile("cvt.rn.f16.f32 %0, %1;\n" : "=h"(tmp.u16[0]) : "f"(f)); +#else + asm volatile("v_cvt_f16_f32 %0, %1;\n" : "=v"(tmp.u32) : "v"(f)); +#endif + return tmp.u16[0]; +} + +inline __device__ uint32_t float2_to_half2(float2 f) { + union { + uint32_t u32; + uint16_t u16[2]; + } tmp; +#ifndef USE_ROCM + #if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 800 + asm volatile("cvt.rn.f16x2.f32 %0, %1, %2;\n" + : "=r"(tmp.u32) + : "f"(f.y), "f"(f.x)); + #else + asm volatile("cvt.rn.f16.f32 %0, %1;\n" : "=h"(tmp.u16[0]) : "f"(f.x)); + asm volatile("cvt.rn.f16.f32 %0, %1;\n" : "=h"(tmp.u16[1]) : "f"(f.y)); + #endif +#else + tmp.u16[0] = float_to_half(f.x); + tmp.u16[1] = float_to_half(f.y); +#endif + return tmp.u32; +} + +// Vector addition. +inline __device__ uint16_t add(uint16_t a, uint16_t b) { + uint16_t c; +#ifndef USE_ROCM + asm volatile("add.f16 %0, %1, %2;\n" : "=h"(c) : "h"(a), "h"(b)); +#else + asm volatile("v_add_f16 %0, %1, %2;\n" : "=v"(c) : "v"(a), "v"(b)); +#endif + return c; +} + +inline __device__ uint32_t add(uint32_t a, uint32_t b) { + uint32_t c; +#ifndef USE_ROCM + asm volatile("add.f16x2 %0, %1, %2;\n" : "=r"(c) : "r"(a), "r"(b)); +#else + asm volatile("v_pk_add_f16 %0, %1, %2;\n" : "=v"(c) : "v"(a), "v"(b)); +#endif + return c; +} + +inline __device__ uint2 add(uint2 a, uint2 b) { + uint2 c; + c.x = add(a.x, b.x); + c.y = add(a.y, b.y); + return c; +} + +inline __device__ uint4 add(uint4 a, uint4 b) { + uint4 c; + c.x = add(a.x, b.x); + c.y = add(a.y, b.y); + c.z = add(a.z, b.z); + c.w = add(a.w, b.w); + return c; +} + +inline __device__ float2 add(uint32_t a, float2 fb) { + float2 fa = half2_to_float2(a); + return add(fa, fb); +} + +inline __device__ Float4_ add(uint2 a, Float4_ fb) { + Float4_ fc; + fc.x = add(a.x, fb.x); + fc.y = add(a.y, fb.y); + return fc; +} + +inline __device__ Float8_ add(uint4 a, Float8_ fb) { + Float8_ fc; + fc.x = add(a.x, fb.x); + fc.y = add(a.y, fb.y); + fc.z = add(a.z, fb.z); + fc.w = add(a.w, fb.w); + return fc; +} + +// Vector multiplication. +template <> +inline __device__ uint16_t mul(uint16_t a, uint16_t b) { + uint16_t c; +#ifndef USE_ROCM + asm volatile("mul.f16 %0, %1, %2;\n" : "=h"(c) : "h"(a), "h"(b)); +#else + asm volatile("v_mul_f16 %0, %1, %2;\n" : "=v"(c) : "v"(a), "v"(b)); +#endif + return c; +} + +template <> +inline __device__ uint32_t mul(uint32_t a, uint32_t b) { + uint32_t c; +#ifndef USE_ROCM + asm volatile("mul.f16x2 %0, %1, %2;\n" : "=r"(c) : "r"(a), "r"(b)); +#else + asm volatile("v_pk_mul_f16 %0, %1, %2;\n" : "=v"(c) : "v"(a), "v"(b)); +#endif + return c; +} + +template <> +inline __device__ uint32_t mul(uint16_t a, uint32_t b) { + return mul(h0_h0(a), b); +} + +template <> +inline __device__ uint2 mul(uint2 a, uint2 b) { + uint2 c; + c.x = mul(a.x, b.x); + c.y = mul(a.y, b.y); + return c; +} + +template <> +inline __device__ uint2 mul(uint16_t a, uint2 b) { + uint32_t s = h0_h0(a); + uint2 c; + c.x = mul(s, b.x); + c.y = mul(s, b.y); + return c; +} + +template <> +inline __device__ uint4 mul(uint4 a, uint4 b) { + uint4 c; + c.x = mul(a.x, b.x); + c.y = mul(a.y, b.y); + c.z = mul(a.z, b.z); + c.w = mul(a.w, b.w); + return c; +} + +template <> +inline __device__ uint4 mul(uint16_t a, uint4 b) { + uint32_t s = h0_h0(a); + uint4 c; + c.x = mul(s, b.x); + c.y = mul(s, b.y); + c.z = mul(s, b.z); + c.w = mul(s, b.w); + return c; +} + +template <> +inline __device__ float mul(uint16_t a, uint16_t b) { + float fa = half_to_float(a); + float fb = half_to_float(b); + return fa * fb; +} + +template <> +inline __device__ float2 mul(uint32_t a, uint32_t b) { + float2 fa = half2_to_float2(a); + float2 fb = half2_to_float2(b); + return mul(fa, fb); +} + +template <> +inline __device__ float2 mul(uint16_t a, uint32_t b) { + return mul(h0_h0(a), b); +} + +template <> +inline __device__ Float4_ mul(uint2 a, uint2 b) { + Float4_ fc; + fc.x = mul(a.x, b.x); + fc.y = mul(a.y, b.y); + return fc; +} + +template <> +inline __device__ Float4_ mul(uint16_t a, uint2 b) { + uint32_t s = h0_h0(a); + Float4_ fc; + fc.x = mul(s, b.x); + fc.y = mul(s, b.y); + return fc; +} + +template <> +inline __device__ Float8_ mul(uint4 a, uint4 b) { + Float8_ fc; + fc.x = mul(a.x, b.x); + fc.y = mul(a.y, b.y); + fc.z = mul(a.z, b.z); + fc.w = mul(a.w, b.w); + return fc; +} + +template <> +inline __device__ Float8_ mul(uint16_t a, uint4 b) { + uint32_t s = h0_h0(a); + Float8_ fc; + fc.x = mul(s, b.x); + fc.y = mul(s, b.y); + fc.z = mul(s, b.z); + fc.w = mul(s, b.w); + return fc; +} + +// Vector fused multiply-add. +inline __device__ uint32_t fma(uint32_t a, uint32_t b, uint32_t c) { + uint32_t d; +#ifndef USE_ROCM + asm volatile("fma.rn.f16x2 %0, %1, %2, %3;\n" + : "=r"(d) + : "r"(a), "r"(b), "r"(c)); +#else + asm volatile("v_pk_fma_f16 %0, %1, %2, %3;\n" + : "=v"(d) + : "v"(a), "v"(b), "v"(c)); +#endif + return d; +} + +inline __device__ uint32_t fma(uint16_t a, uint32_t b, uint32_t c) { + return fma(h0_h0(a), b, c); +} + +inline __device__ uint2 fma(uint2 a, uint2 b, uint2 c) { + uint2 d; + d.x = fma(a.x, b.x, c.x); + d.y = fma(a.y, b.y, c.y); + return d; +} + +inline __device__ uint2 fma(uint16_t a, uint2 b, uint2 c) { + uint32_t s = h0_h0(a); + uint2 d; + d.x = fma(s, b.x, c.x); + d.y = fma(s, b.y, c.y); + return d; +} + +inline __device__ uint4 fma(uint4 a, uint4 b, uint4 c) { + uint4 d; + d.x = fma(a.x, b.x, c.x); + d.y = fma(a.y, b.y, c.y); + d.z = fma(a.z, b.z, c.z); + d.w = fma(a.w, b.w, c.w); + return d; +} + +inline __device__ uint4 fma(uint16_t a, uint4 b, uint4 c) { + uint32_t s = h0_h0(a); + uint4 d; + d.x = fma(s, b.x, c.x); + d.y = fma(s, b.y, c.y); + d.z = fma(s, b.z, c.z); + d.w = fma(s, b.w, c.w); + return d; +} + +inline __device__ float fma(uint16_t a, uint16_t b, float fc) { + float fa = half_to_float(a); + float fb = half_to_float(b); + return fa * fb + fc; +} + +inline __device__ float2 fma(uint32_t a, uint32_t b, float2 fc) { + float2 fa = half2_to_float2(a); + float2 fb = half2_to_float2(b); + return fma(fa, fb, fc); +} + +inline __device__ float2 fma(uint16_t a, uint32_t b, float2 fc) { + return fma(h0_h0(a), b, fc); +} + +inline __device__ Float4_ fma(uint2 a, uint2 b, Float4_ fc) { + Float4_ fd; + fd.x = fma(a.x, b.x, fc.x); + fd.y = fma(a.y, b.y, fc.y); + return fd; +} + +inline __device__ Float4_ fma(uint16_t a, uint2 b, Float4_ fc) { + uint32_t s = h0_h0(a); + Float4_ fd; + fd.x = fma(s, b.x, fc.x); + fd.y = fma(s, b.y, fc.y); + return fd; +} + +inline __device__ Float8_ fma(uint4 a, uint4 b, Float8_ fc) { + Float8_ fd; + fd.x = fma(a.x, b.x, fc.x); + fd.y = fma(a.y, b.y, fc.y); + fd.z = fma(a.z, b.z, fc.z); + fd.w = fma(a.w, b.w, fc.w); + return fd; +} + +inline __device__ Float8_ fma(uint16_t a, uint4 b, Float8_ fc) { + uint32_t s = h0_h0(a); + Float8_ fd; + fd.x = fma(s, b.x, fc.x); + fd.y = fma(s, b.y, fc.y); + fd.z = fma(s, b.z, fc.z); + fd.w = fma(s, b.w, fc.w); + return fd; +} + +// Vector sum. +template <> +inline __device__ float sum(uint16_t v) { + return half_to_float(v); +} + +template <> +inline __device__ float sum(uint32_t v) { + float2 tmp = half2_to_float2(v); + return tmp.x + tmp.y; +} + +template <> +inline __device__ float sum(uint2 v) { + uint32_t c = add(v.x, v.y); + return sum(c); +} + +template <> +inline __device__ float sum(uint4 v) { + uint32_t c = add(v.x, v.y); + c = add(c, v.z); + c = add(c, v.w); + return sum(c); +} + +// From float32 to float16. +inline __device__ void from_float(uint16_t& dst, float src) { + dst = float_to_half(src); +} + +inline __device__ void from_float(uint32_t& dst, float2 src) { + dst = float2_to_half2(src); +} + +inline __device__ void from_float(uint2& dst, Float4_ src) { + dst.x = float2_to_half2(src.x); + dst.y = float2_to_half2(src.y); +} + +inline __device__ void from_float(uint4& dst, Float8_ src) { + dst.x = float2_to_half2(src.x); + dst.y = float2_to_half2(src.y); + dst.z = float2_to_half2(src.z); + dst.w = float2_to_half2(src.w); +} + +// From float16 to float32. +inline __device__ float to_float(uint16_t u) { return half_to_float(u); } + +inline __device__ float2 to_float(uint32_t u) { return half2_to_float2(u); } + +inline __device__ Float4_ to_float(uint2 u) { + Float4_ tmp; + tmp.x = half2_to_float2(u.x); + tmp.y = half2_to_float2(u.y); + return tmp; +} + +inline __device__ Float8_ to_float(uint4 u) { + Float8_ tmp; + tmp.x = half2_to_float2(u.x); + tmp.y = half2_to_float2(u.y); + tmp.z = half2_to_float2(u.z); + tmp.w = half2_to_float2(u.w); + return tmp; +} + +// Zero-out a variable. +inline __device__ void zero(uint16_t& dst) { dst = uint16_t(0); } + +} // namespace vllm diff --git a/tasks/hip2hip/campaign20/paged_attention_large/src/attention/dtype_float32.cuh b/tasks/hip2hip/campaign20/paged_attention_large/src/attention/dtype_float32.cuh new file mode 100644 index 00000000..7c6a686d --- /dev/null +++ b/tasks/hip2hip/campaign20/paged_attention_large/src/attention/dtype_float32.cuh @@ -0,0 +1,251 @@ +/* + * Adapted from + * https://github.com/NVIDIA/FasterTransformer/blob/release/v5.3_tag/src/fastertransformer/kernels/decoder_masked_multihead_attention/decoder_masked_multihead_attention_template.hpp + * and + * https://github.com/NVIDIA/FasterTransformer/blob/release/v5.3_tag/src/fastertransformer/kernels/decoder_masked_multihead_attention_utils.h + * Copyright (c) 2023, The vLLM team. + * Copyright (c) 2020-2023, NVIDIA CORPORATION. All rights reserved. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +#pragma once + +#include "attention_generic.cuh" + +#include + +namespace vllm { + +// Define custom FP32 vector data types. +struct Float4_ { + float2 x; + float2 y; +}; + +struct Float8_ { + float2 x; + float2 y; + float2 z; + float2 w; +}; + +// FP32 vector types for Q, K, V. +template <> +struct Vec { + using Type = float; +}; +template <> +struct Vec { + using Type = float2; +}; +template <> +struct Vec { + using Type = float4; +}; + +// FP32 accumulator vector types corresponding to Vec. +template <> +struct FloatVec { + using Type = float; +}; +template <> +struct FloatVec { + using Type = float2; +}; +template <> +struct FloatVec { + using Type = float4; +}; + +// Vector addition. +inline __device__ float add(float a, float b) { return a + b; } + +inline __device__ float2 add(float2 a, float2 b) { + float2 c; + c.x = add(a.x, b.x); + c.y = add(a.y, b.y); + return c; +} + +inline __device__ float4 add(float4 a, float4 b) { + float4 c; + c.x = add(a.x, b.x); + c.y = add(a.y, b.y); + c.z = add(a.z, b.z); + c.w = add(a.w, b.w); + return c; +} + +// Vector multiplication. +template <> +inline __device__ float mul(float a, float b) { + return a * b; +} + +template <> +inline __device__ float2 mul(float2 a, float2 b) { + float2 c; + c.x = a.x * b.x; + c.y = a.y * b.y; + return c; +} + +template <> +inline __device__ float2 mul(float a, float2 b) { + float2 c; + c.x = a * b.x; + c.y = a * b.y; + return c; +} + +template <> +inline __device__ float4 mul(float4 a, float4 b) { + float4 c; + c.x = a.x * b.x; + c.y = a.y * b.y; + c.z = a.z * b.z; + c.w = a.w * b.w; + return c; +} + +template <> +inline __device__ float4 mul(float a, float4 b) { + float4 c; + c.x = a * b.x; + c.y = a * b.y; + c.z = a * b.z; + c.w = a * b.w; + return c; +} + +// Vector fused multiply-add. +inline __device__ float fma(float a, float b, float c) { return a * b + c; } + +inline __device__ float2 fma(float2 a, float2 b, float2 c) { + float2 d; + d.x = fma(a.x, b.x, c.x); + d.y = fma(a.y, b.y, c.y); + return d; +} + +inline __device__ float2 fma(float a, float2 b, float2 c) { + float2 d; + d.x = fma(a, b.x, c.x); + d.y = fma(a, b.y, c.y); + return d; +} + +inline __device__ float4 fma(float4 a, float4 b, float4 c) { + float4 d; + d.x = fma(a.x, b.x, c.x); + d.y = fma(a.y, b.y, c.y); + d.z = fma(a.z, b.z, c.z); + d.w = fma(a.w, b.w, c.w); + return d; +} + +inline __device__ float4 fma(float a, float4 b, float4 c) { + float4 d; + d.x = fma(a, b.x, c.x); + d.y = fma(a, b.y, c.y); + d.z = fma(a, b.z, c.z); + d.w = fma(a, b.w, c.w); + return d; +} + +inline __device__ Float4_ fma(float a, Float4_ b, Float4_ c) { + Float4_ d; + d.x = fma(a, b.x, c.x); + d.y = fma(a, b.y, c.y); + return d; +} + +inline __device__ Float8_ fma(float a, Float8_ b, Float8_ c) { + Float8_ d; + d.x = fma(a, b.x, c.x); + d.y = fma(a, b.y, c.y); + d.z = fma(a, b.z, c.z); + d.w = fma(a, b.w, c.w); + return d; +} + +// Vector sum. +template <> +inline __device__ float sum(float v) { + return v; +} + +template <> +inline __device__ float sum(float2 v) { + return v.x + v.y; +} + +template <> +inline __device__ float sum(float4 v) { + return v.x + v.y + v.z + v.w; +} + +template <> +inline __device__ float sum(Float4_ v) { + return v.x.x + v.x.y + v.y.x + v.y.y; +} + +template <> +inline __device__ float sum(Float8_ v) { + return v.x.x + v.x.y + v.y.x + v.y.y + v.z.x + v.z.y + v.w.x + v.w.y; +} + +// Vector dot product. +inline __device__ float dot(float a, float b) { return a * b; } + +inline __device__ float dot(float2 a, float2 b) { + float2 c = mul(a, b); + return c.x + c.y; +} + +inline __device__ float dot(Float4_ a, Float4_ b) { + float2 acc = mul(a.x, b.x); + acc = fma(a.y, b.y, acc); + return acc.x + acc.y; +} + +inline __device__ float dot(Float8_ a, Float8_ b) { + float2 acc = mul(a.x, b.x); + acc = fma(a.y, b.y, acc); + acc = fma(a.z, b.z, acc); + acc = fma(a.w, b.w, acc); + return acc.x + acc.y; +} + +// From float to float. +inline __device__ void from_float(float& dst, float src) { dst = src; } + +inline __device__ void from_float(float2& dst, float2 src) { dst = src; } + +inline __device__ void from_float(float4& dst, float4 src) { dst = src; } + +// From float to float. +inline __device__ float to_float(float u) { return u; } + +inline __device__ float2 to_float(float2 u) { return u; } + +inline __device__ float4 to_float(float4 u) { return u; } + +inline __device__ Float4_ to_float(Float4_ u) { return u; } + +inline __device__ Float8_ to_float(Float8_ u) { return u; } + +// Zero-out a variable. +inline __device__ void zero(float& dst) { dst = 0.f; } + +} // namespace vllm diff --git a/tasks/hip2hip/campaign20/paged_attention_large/src/attention/dtype_fp8.cuh b/tasks/hip2hip/campaign20/paged_attention_large/src/attention/dtype_fp8.cuh new file mode 100644 index 00000000..e714e321 --- /dev/null +++ b/tasks/hip2hip/campaign20/paged_attention_large/src/attention/dtype_fp8.cuh @@ -0,0 +1,41 @@ +#pragma once + +#include "attention_generic.cuh" + +#include +#ifdef ENABLE_FP8 + #ifndef USE_ROCM + #include + #endif // USE_ROCM +#endif // ENABLE_FP8 + +namespace vllm { + +enum class Fp8KVCacheDataType { + kAuto = 0, + kFp8E4M3 = 1, + kFp8E5M2 = 2, +}; + +// fp8 vector types for quantization of kv cache +template <> +struct Vec { + using Type = uint8_t; +}; + +template <> +struct Vec { + using Type = uint16_t; +}; + +template <> +struct Vec { + using Type = uint32_t; +}; + +template <> +struct Vec { + using Type = uint2; +}; + +} // namespace vllm diff --git a/tasks/hip2hip/campaign20/paged_attention_large/src/bindings.cpp b/tasks/hip2hip/campaign20/paged_attention_large/src/bindings.cpp new file mode 100644 index 00000000..2952d6f9 --- /dev/null +++ b/tasks/hip2hip/campaign20/paged_attention_large/src/bindings.cpp @@ -0,0 +1,15 @@ +// Auto-generated by extractor.source.hip_extractor. +// Re-registers the extracted op into the `extracted_paged_attention` namespace so the task +// can call torch.ops.extracted_paged_attention.paged_attention without touching the host engine. +#include +#include + +void paged_attention( torch::Tensor& out, torch::Tensor& exp_sums, torch::Tensor& max_logits, torch::Tensor& tmp_out, torch::Tensor& query, torch::Tensor& key_cache, torch::Tensor& value_cache, int64_t num_kv_heads, double scale, torch::Tensor& block_tables, torch::Tensor& seq_lens, const std::optional& query_start_loc, int64_t block_size, int64_t max_seq_len, const std::optional& alibi_slopes, const std::string& kv_cache_dtype, torch::Tensor& k_scale, torch::Tensor& v_scale, const std::optional& fp8_out_scale, const std::string& mfma_type); + +TORCH_LIBRARY(extracted_paged_attention, m) { + m.def("paged_attention(Tensor! out, Tensor exp_sums, Tensor max_logits, Tensor tmp_out, Tensor query, Tensor key_cache, Tensor value_cache, int num_kv_heads, float scale, Tensor block_tables, Tensor seq_lens, Tensor? query_start_loc, int block_size, int max_seq_len, Tensor? alibi_slopes, str kv_cache_dtype, Tensor k_scale, Tensor v_scale, Tensor? fp8_out_scale, str mfma_type) -> ()"); +} + +TORCH_LIBRARY_IMPL(extracted_paged_attention, CUDA, m) { + m.impl("paged_attention", &paged_attention); +} diff --git a/tasks/hip2hip/campaign20/paged_attention_large/src/cuda_compat.h b/tasks/hip2hip/campaign20/paged_attention_large/src/cuda_compat.h new file mode 100644 index 00000000..d7d589db --- /dev/null +++ b/tasks/hip2hip/campaign20/paged_attention_large/src/cuda_compat.h @@ -0,0 +1,76 @@ +#pragma once + +#ifdef USE_ROCM + #include +#endif + +#ifdef USE_ROCM +struct Utils { + static __host__ int get_warp_size() { + static bool is_cached = false; + static int result; + + if (!is_cached) { + int device_id; + cudaDeviceProp deviceProp; + cudaGetDevice(&device_id); + cudaGetDeviceProperties(&deviceProp, device_id); + + result = deviceProp.warpSize; + is_cached = true; + } + + return result; + } + + static __device__ constexpr int get_warp_size() { + #ifdef __GFX9__ + return 64; + #else + return 32; + #endif + } +}; + + #define WARP_SIZE Utils::get_warp_size() +#else + #define WARP_SIZE 32 +#endif + +#ifndef USE_ROCM + #define VLLM_LDG(arg) __ldg(arg) +#else + #define VLLM_LDG(arg) *(arg) +#endif + +#ifndef USE_ROCM + #define VLLM_SHFL_XOR_SYNC(var, lane_mask) \ + __shfl_xor_sync(uint32_t(-1), var, lane_mask) + #define VLLM_SHFL_XOR_SYNC_WIDTH(var, lane_mask, width) \ + __shfl_xor_sync(uint32_t(-1), var, lane_mask, width) +#else + #define VLLM_SHFL_XOR_SYNC(var, lane_mask) __shfl_xor(var, lane_mask) + #define VLLM_SHFL_XOR_SYNC_WIDTH(var, lane_mask, width) \ + __shfl_xor(var, lane_mask, width) +#endif + +#ifndef USE_ROCM + #define VLLM_SHFL_SYNC(var, src_lane) __shfl_sync(uint32_t(-1), var, src_lane) +#else + #define VLLM_SHFL_SYNC(var, src_lane) __shfl(var, src_lane) +#endif + +#ifndef USE_ROCM + #define VLLM_SHFL_DOWN_SYNC(var, lane_delta) \ + __shfl_down_sync(uint32_t(-1), var, lane_delta) +#else + #define VLLM_SHFL_DOWN_SYNC(var, lane_delta) __shfl_down(var, lane_delta) +#endif + +#ifndef USE_ROCM + #define VLLM_DevFuncAttribute_SET_MaxDynamicSharedMemorySize(FUNC, VAL) \ + cudaFuncSetAttribute(FUNC, cudaFuncAttributeMaxDynamicSharedMemorySize, VAL) +#else + #define VLLM_DevFuncAttribute_SET_MaxDynamicSharedMemorySize(FUNC, VAL) \ + hipFuncSetAttribute(FUNC, hipFuncAttributeMaxDynamicSharedMemorySize, VAL) +#endif diff --git a/tasks/hip2hip/campaign20/paged_attention_large/src/hip_compat.h b/tasks/hip2hip/campaign20/paged_attention_large/src/hip_compat.h new file mode 100644 index 00000000..645f0e17 --- /dev/null +++ b/tasks/hip2hip/campaign20/paged_attention_large/src/hip_compat.h @@ -0,0 +1,77 @@ +// !!! This is a file automatically generated by hipify!!! +#pragma once + +#ifdef USE_ROCM + #include +#endif + +#ifdef USE_ROCM +struct Utils { + static __host__ int get_warp_size() { + static bool is_cached = false; + static int result; + + if (!is_cached) { + int device_id; + hipDeviceProp_t deviceProp; + hipGetDevice(&device_id); + hipGetDeviceProperties(&deviceProp, device_id); + + result = deviceProp.warpSize; + is_cached = true; + } + + return result; + } + + static __device__ constexpr int get_warp_size() { + #ifdef __GFX9__ + return 64; + #else + return 32; + #endif + } +}; + + #define WARP_SIZE Utils::get_warp_size() +#else + #define WARP_SIZE 32 +#endif + +#ifndef USE_ROCM + #define VLLM_LDG(arg) __ldg(arg) +#else + #define VLLM_LDG(arg) *(arg) +#endif + +#ifndef USE_ROCM + #define VLLM_SHFL_XOR_SYNC(var, lane_mask) \ + __shfl_xor_sync(uint32_t(-1), var, lane_mask) + #define VLLM_SHFL_XOR_SYNC_WIDTH(var, lane_mask, width) \ + __shfl_xor_sync(uint32_t(-1), var, lane_mask, width) +#else + #define VLLM_SHFL_XOR_SYNC(var, lane_mask) __shfl_xor(var, lane_mask) + #define VLLM_SHFL_XOR_SYNC_WIDTH(var, lane_mask, width) \ + __shfl_xor(var, lane_mask, width) +#endif + +#ifndef USE_ROCM + #define VLLM_SHFL_SYNC(var, src_lane) __shfl_sync(uint32_t(-1), var, src_lane) +#else + #define VLLM_SHFL_SYNC(var, src_lane) __shfl(var, src_lane) +#endif + +#ifndef USE_ROCM + #define VLLM_SHFL_DOWN_SYNC(var, lane_delta) \ + __shfl_down_sync(uint32_t(-1), var, lane_delta) +#else + #define VLLM_SHFL_DOWN_SYNC(var, lane_delta) __shfl_down(var, lane_delta) +#endif + +#ifndef USE_ROCM + #define VLLM_DevFuncAttribute_SET_MaxDynamicSharedMemorySize(FUNC, VAL) \ + hipFuncSetAttribute(FUNC, hipFuncAttributeMaxDynamicSharedMemorySize, VAL) +#else + #define VLLM_DevFuncAttribute_SET_MaxDynamicSharedMemorySize(FUNC, VAL) \ + hipFuncSetAttribute(FUNC, hipFuncAttributeMaxDynamicSharedMemorySize, VAL) +#endif diff --git a/tasks/hip2hip/campaign20/paged_attention_large/src/quantization/w8a8/fp8/amd/quant_utils.cuh b/tasks/hip2hip/campaign20/paged_attention_large/src/quantization/w8a8/fp8/amd/quant_utils.cuh new file mode 100644 index 00000000..81f5cb83 --- /dev/null +++ b/tasks/hip2hip/campaign20/paged_attention_large/src/quantization/w8a8/fp8/amd/quant_utils.cuh @@ -0,0 +1,671 @@ +#pragma once +#include + +#include +#include +#include + +#include "../../../../attention/attention_dtypes.h" + +namespace vllm { +#ifdef USE_ROCM + +namespace fp8 { + #ifdef ENABLE_FP8 + +// Use hardware cvt instruction for fp8 on rocm +template +__device__ __forceinline__ fp8_type cvt_c10(float const r) { + return {}; +} + +// __hip_fp8_e4m3 only exists starting in ROCm 6.3. The macro +// HIP_FP8_TYPE_OCP comes from the hip_fp8.h header and also makes +// its first appearance in ROCm 6.3. Since VLLM_DISPATCH_FP8_TYPES +// on ROCm instantiates both OCP and FNUZ kernels, we need to replace +// the new HW cvt with something reasonable that doesn't rely on the +// ROCm 6.3 feature. This allows compiling on ROCm 6.2 or newer. +template <> +__device__ __forceinline__ c10::Float8_e4m3fn cvt_c10(float const r) { + #if HIP_FP8_TYPE_OCP + return c10::Float8_e4m3fn( + __hip_cvt_float_to_fp8(r, __hip_fp8_e4m3::__default_saturation, + __hip_fp8_e4m3::__default_interpret), + c10::Float8_e4m3fn::from_bits()); + #else + // Cast implemented by pytorch. Uses bit manipulation instead of HW cvt. + // HW cvt above is faster when it is available (ROCm 6.3 or newer). + return static_cast(r); + #endif +} + +template <> +__device__ __forceinline__ c10::Float8_e4m3fnuz cvt_c10(float const r) { + return c10::Float8_e4m3fnuz( + __hip_cvt_float_to_fp8(r, __hip_fp8_e4m3_fnuz::__default_saturation, + __hip_fp8_e4m3_fnuz::__default_interpret), + c10::Float8_e4m3fnuz::from_bits()); +} + +template +__inline__ __device__ Tout vec_conversion(const Tin& x) { + return x; +} + +template +__inline__ __device__ Tout scaled_vec_conversion(const Tin& x, + const float scale) { + return x; +} + + #if HIP_FP8_TYPE_OCP +using fp8_type = __hip_fp8_e4m3; +using fp8x2_type = __hip_fp8x2_e4m3; + #else +using fp8_type = __hip_fp8_e4m3_fnuz; +using fp8x2_type = __hip_fp8x2_e4m3_fnuz; + #endif + +// fp8 -> half +template <> +__inline__ __device__ uint16_t +vec_conversion(const uint8_t& a) { + return __hip_cvt_fp8_to_halfraw(a, fp8_type::__default_interpret).x; +} + +// fp8x2 -> half2 +template <> +__inline__ __device__ uint32_t +vec_conversion(const uint16_t& a) { + union { + __half2_raw h2r; + uint32_t ui32; + } tmp; + tmp.h2r = __hip_cvt_fp8x2_to_halfraw2(a, fp8_type::__default_interpret); + return tmp.ui32; +} + +// fp8x4 -> half2x2 +template <> +__inline__ __device__ uint2 vec_conversion(const uint32_t& a) { + union { + uint2 u32x2; + uint32_t u32[2]; + } tmp; + tmp.u32[0] = vec_conversion((uint16_t)a); + tmp.u32[1] = vec_conversion((uint16_t)(a >> 16U)); + return tmp.u32x2; +} + +// fp8x8 -> half2x4 +template <> +__inline__ __device__ uint4 vec_conversion(const uint2& a) { + union { + uint4 u64x2; + uint2 u64[2]; + } tmp; + tmp.u64[0] = vec_conversion(a.x); + tmp.u64[1] = vec_conversion(a.y); + return tmp.u64x2; +} + +using __nv_bfloat16 = __hip_bfloat16; + +// fp8 -> __nv_bfloat16 +template <> +__inline__ __device__ __nv_bfloat16 +vec_conversion<__nv_bfloat16, uint8_t>(const uint8_t& a) { + fp8_type f8; + f8.__x = a; + return __float2bfloat16(static_cast(f8)); +} + +using __nv_bfloat162 = __hip_bfloat162; + +// fp8x2 -> __nv_bfloat162 +template <> +__inline__ __device__ __nv_bfloat162 +vec_conversion<__nv_bfloat162, uint16_t>(const uint16_t& a) { + __nv_bfloat162 res; + res.x = vec_conversion<__nv_bfloat16, uint8_t>((uint8_t)a); + res.y = vec_conversion<__nv_bfloat16, uint8_t>((uint8_t)(a >> 8U)); + return res; +} + +// fp8x4 -> bf16_4_t +template <> +__inline__ __device__ bf16_4_t +vec_conversion(const uint32_t& a) { + bf16_4_t res; + res.x = vec_conversion<__nv_bfloat162, uint16_t>((uint16_t)a); + res.y = vec_conversion<__nv_bfloat162, uint16_t>((uint16_t)(a >> 16U)); + return res; +} + +// fp8x8 -> bf16_8_t +template <> +__inline__ __device__ bf16_8_t vec_conversion(const uint2& a) { + bf16_4_t tmp1, tmp2; + tmp1 = vec_conversion(a.x); + tmp2 = vec_conversion(a.y); + bf16_8_t res; + res.x = tmp1.x; + res.y = tmp1.y; + res.z = tmp2.x; + res.w = tmp2.y; + return res; +} + +// fp8 -> float +template <> +__inline__ __device__ float vec_conversion(const uint8_t& a) { + fp8_type f8; + f8.__x = a; + return static_cast(f8); +} + +// fp8x2 -> float2 +template <> +__inline__ __device__ float2 +vec_conversion(const uint16_t& a) { + fp8x2_type f8x2; + f8x2.__x = a; + return static_cast(f8x2); +} + +// fp8x4 -> float4 +template <> +__inline__ __device__ Float4_ +vec_conversion(const uint32_t& a) { + Float4_ res; + res.x = vec_conversion((uint16_t)a); + res.y = vec_conversion((uint16_t)(a >> 16U)); + return res; +} + +// fp8x4 -> float4 +template <> +__inline__ __device__ float4 +vec_conversion(const uint32_t& a) { + Float4_ tmp = vec_conversion(a); + float4 res = make_float4(tmp.x.x, tmp.x.y, tmp.y.x, tmp.y.y); + return res; +} + +// fp8x8 -> float8 +template <> +__inline__ __device__ Float8_ vec_conversion(const uint2& a) { + Float4_ tmp1, tmp2; + tmp1 = vec_conversion(a.x); + tmp2 = vec_conversion(a.y); + Float8_ res; + res.x = tmp1.x; + res.y = tmp1.y; + res.z = tmp2.x; + res.w = tmp2.y; + return res; +} + +// half -> fp8 +template <> +__inline__ __device__ uint8_t +vec_conversion(const uint16_t& a) { + __half_raw tmp; + tmp.x = a; + return __hip_cvt_halfraw_to_fp8(tmp, fp8_type::__default_saturation, + fp8_type::__default_interpret); +} + +template <> +__inline__ __device__ uint16_t +vec_conversion(const uint32_t& a) { + union { + uint32_t ui32; + __half2_raw h2r; + } tmp; + tmp.ui32 = a; + return __hip_cvt_halfraw2_to_fp8x2(tmp.h2r, fp8_type::__default_saturation, + fp8_type::__default_interpret); +} + +// bf16 -> fp8 +template <> +__inline__ __device__ uint8_t +vec_conversion(const __nv_bfloat16& a) { + return __hip_cvt_float_to_fp8(__bfloat162float(a), + fp8_type::__default_saturation, + fp8_type::__default_interpret); +} + +// float -> fp8 +template <> +__inline__ __device__ uint8_t vec_conversion(const float& a) { + return __hip_cvt_float_to_fp8(a, fp8_type::__default_saturation, + fp8_type::__default_interpret); +} + +// float2 -> half2 +template <> +__inline__ __device__ uint32_t +vec_conversion(const float2& a) { + union { + half2 float16; + uint32_t uint32; + }; + + float16 = __float22half2_rn(a); + return uint32; +} + +// Float4 -> half2x2 +template <> +__inline__ __device__ uint2 vec_conversion(const Float4_& a) { + uint2 b; + float2 val; + val.x = a.x.x; + val.y = a.x.y; + b.x = vec_conversion(val); + + val.x = a.y.x; + val.y = a.y.y; + b.y = vec_conversion(val); + return b; +} + +// Float4 -> float4 +template <> +__inline__ __device__ float4 vec_conversion(const Float4_& a) { + float4 b; + b.x = a.x.x; + b.y = a.x.y; + b.z = a.y.x; + b.w = a.y.y; + return b; +} + +// Float8 -> half2x4 +template <> +__inline__ __device__ uint4 vec_conversion(const Float8_& a) { + uint4 b; + b.x = vec_conversion(a.x); + b.y = vec_conversion(a.y); + b.z = vec_conversion(a.z); + b.w = vec_conversion(a.w); + return b; +} + +// float2 -> bfloat162 +template <> +__inline__ __device__ __nv_bfloat162 +vec_conversion<__nv_bfloat162, float2>(const float2& a) { + __nv_bfloat162 b = __float22bfloat162_rn(a); + return b; +} + +// Float4 -> bfloat162x2 +template <> +__inline__ __device__ bf16_4_t +vec_conversion(const Float4_& a) { + bf16_4_t b; + b.x = __float22bfloat162_rn(a.x); + b.y = __float22bfloat162_rn(a.y); + return b; +} + +// Float8 -> bfloat162x4 +template <> +__inline__ __device__ bf16_8_t +vec_conversion(const Float8_& a) { + bf16_8_t b; + b.x = __float22bfloat162_rn(a.x); + b.y = __float22bfloat162_rn(a.y); + b.z = __float22bfloat162_rn(a.z); + b.w = __float22bfloat162_rn(a.w); + return b; +} + +/* Scaled and vectorized conversions, for data exchange between high and low + precision domains + + Convention of the scale in API, e.g: FP8_data = Quantization( + High_Precision_data / scale ) s.t. Quantize(HP / scale) => FP8 Dequant(FP8) * + scale => HP + + */ + +using __nv_bfloat16 = __hip_bfloat16; + +// fp8 -> __nv_bfloat16 +template <> +__inline__ __device__ __nv_bfloat16 +scaled_vec_conversion<__nv_bfloat16, uint8_t>(const uint8_t& a, float scale) { + fp8_type f8; + f8.__x = a; + return __float2bfloat16(static_cast(f8) * scale); +} + +// fp8x2 -> __nv_bfloat162 +template <> +__inline__ __device__ __nv_bfloat162 +scaled_vec_conversion<__nv_bfloat162, uint16_t>(const uint16_t& a, + float scale) { + __nv_bfloat162 res; + res.x = scaled_vec_conversion<__nv_bfloat16, uint8_t>((uint8_t)a, scale); + res.y = + scaled_vec_conversion<__nv_bfloat16, uint8_t>((uint8_t)(a >> 8U), scale); + return res; +} + +// fp8x4 -> bf16_4_t +template <> +__inline__ __device__ bf16_4_t +scaled_vec_conversion(const uint32_t& a, float scale) { + bf16_4_t res; + res.x = scaled_vec_conversion<__nv_bfloat162, uint16_t>((uint16_t)a, scale); + res.y = scaled_vec_conversion<__nv_bfloat162, uint16_t>((uint16_t)(a >> 16U), + scale); + return res; +} + +// fp8x8 -> bf16_8_t +template <> +__inline__ __device__ bf16_8_t +scaled_vec_conversion(const uint2& a, float scale) { + bf16_4_t tmp1, tmp2; + tmp1 = scaled_vec_conversion(a.x, scale); + tmp2 = scaled_vec_conversion(a.y, scale); + bf16_8_t res; + res.x = tmp1.x; + res.y = tmp1.y; + res.z = tmp2.x; + res.w = tmp2.y; + return res; +} + +// fp8 -> float +template <> +__inline__ __device__ float scaled_vec_conversion( + const uint8_t& a, float scale) { + fp8_type f8; + f8.__x = a; + return static_cast(f8) * scale; +} + +// fp8x2 -> float2 +template <> +__inline__ __device__ float2 +scaled_vec_conversion(const uint16_t& a, float scale) { + fp8x2_type f8x2; + f8x2.__x = a; + return static_cast(f8x2) * scale; +} + +// fp8x4 -> float4 +template <> +__inline__ __device__ Float4_ +scaled_vec_conversion(const uint32_t& a, const float scale) { + Float4_ res; + res.x = scaled_vec_conversion((uint16_t)a, scale); + res.y = scaled_vec_conversion((uint16_t)(a >> 16U), scale); + return res; +} + +// fp8x4 -> float4 +template <> +__inline__ __device__ float4 +scaled_vec_conversion(const uint32_t& a, float scale) { + Float4_ res = scaled_vec_conversion(a, scale); + return {res.x.x, res.x.y, res.y.x, res.y.y}; +} + +// fp8x8 -> float8 +template <> +__inline__ __device__ Float8_ +scaled_vec_conversion(const uint2& a, float scale) { + Float4_ tmp1, tmp2; + tmp1 = scaled_vec_conversion(a.x, scale); + tmp2 = scaled_vec_conversion(a.y, scale); + Float8_ res; + res.x = tmp1.x; + res.y = tmp1.y; + res.z = tmp2.x; + res.w = tmp2.y; + return res; +} + +// fp8 -> half +template <> +__inline__ __device__ uint16_t +scaled_vec_conversion(const uint8_t& a, float scale) { + __half_raw res; + res.data = scaled_vec_conversion(a, scale); + return res.x; +} + +// fp8x2 -> half2 +template <> +__inline__ __device__ uint32_t +scaled_vec_conversion(const uint16_t& a, float scale) { + union { + __half2_raw h2r; + uint32_t ui32; + } tmp; + tmp.h2r = __hip_cvt_fp8x2_to_halfraw2(a, fp8_type::__default_interpret); + tmp.h2r.x.data *= scale; + tmp.h2r.y.data *= scale; + return tmp.ui32; +} + +// fp8x4 -> half2x2 +template <> +__inline__ __device__ uint2 +scaled_vec_conversion(const uint32_t& a, float scale) { + union { + uint2 u32x2; + uint32_t u32[2]; + } tmp; + tmp.u32[0] = scaled_vec_conversion((uint16_t)a, scale); + tmp.u32[1] = + scaled_vec_conversion((uint16_t)(a >> 16U), scale); + return tmp.u32x2; +} + +// fp8x8 -> half2x4 +template <> +__inline__ __device__ uint4 scaled_vec_conversion(const uint2& a, + float scale) { + union { + uint4 u64x2; + uint2 u64[2]; + } tmp; + tmp.u64[0] = scaled_vec_conversion(a.x, scale); + tmp.u64[1] = scaled_vec_conversion(a.y, scale); + return tmp.u64x2; +} + +// half -> fp8 +template <> +__inline__ __device__ uint8_t +scaled_vec_conversion(const uint16_t& a, float scale) { + __half_raw tmp; + tmp.x = a; + tmp.data /= scale; + return __hip_cvt_halfraw_to_fp8(tmp, fp8_type::__default_saturation, + fp8_type::__default_interpret); +} + +// halfx2 -> fp8x2 +template <> +__inline__ __device__ uint16_t +scaled_vec_conversion(const uint32_t& a, float scale) { + union { + uint32_t ui32; + __half2_raw h2r; + } tmp; + tmp.ui32 = a; + tmp.h2r.x.data /= scale; + tmp.h2r.y.data /= scale; + return __hip_cvt_halfraw2_to_fp8x2(tmp.h2r, fp8_type::__default_saturation, + fp8_type::__default_interpret); +} + +// half2x2 -> fp8x4 +template <> +__inline__ __device__ uint32_t +scaled_vec_conversion(const uint2& a, float scale) { + union { + uint16_t ui16[2]; + uint32_t ui32; + } tmp; + tmp.ui16[0] = scaled_vec_conversion(a.x, scale); + tmp.ui16[1] = scaled_vec_conversion(a.y, scale); + return tmp.ui32; +} + +// half2x4 -> fp8x8 +template <> +__inline__ __device__ uint2 scaled_vec_conversion(const uint4& a, + float scale) { + union { + uint2 ui2[2]; + uint4 ui4; + } tmp; + tmp.ui4 = a; + uint2 res; + res.x = scaled_vec_conversion(tmp.ui2[0], scale); + res.y = scaled_vec_conversion(tmp.ui2[1], scale); + return res; +} + +// bf16 -> fp8 +template <> +__inline__ __device__ uint8_t scaled_vec_conversion( + const __nv_bfloat16& a, float scale) { + return __hip_cvt_float_to_fp8(__bfloat162float(a) / scale, + fp8_type::__default_saturation, + fp8_type::__default_interpret); +} + +// bf16x2 -> fp8x2 +template <> +__inline__ __device__ uint16_t scaled_vec_conversion( + const __nv_bfloat162& a, float scale) { + union { + uint8_t ui8[2]; + uint16_t ui16; + } tmp; + tmp.ui8[0] = scaled_vec_conversion(a.x, scale); + tmp.ui8[1] = scaled_vec_conversion(a.y, scale); + return tmp.ui16; +} + +// bf16x4 -> fp8x4 +template <> +__inline__ __device__ uint32_t +scaled_vec_conversion(const bf16_4_t& a, float scale) { + union { + uint16_t ui16[2]; + uint32_t ui32; + } tmp; + tmp.ui16[0] = scaled_vec_conversion(a.x, scale); + tmp.ui16[1] = scaled_vec_conversion(a.y, scale); + return tmp.ui32; +} + +// bf16x8 -> fp8x8 +template <> +__inline__ __device__ uint2 +scaled_vec_conversion(const bf16_8_t& a, float scale) { + uint2 res; + res.x = scaled_vec_conversion({a.x, a.y}, scale); + res.y = scaled_vec_conversion({a.z, a.w}, scale); + return res; +} + +// float -> fp8 +template <> +__inline__ __device__ uint8_t +scaled_vec_conversion(const float& a, float scale) { + return __hip_cvt_float_to_fp8(a / scale, fp8_type::__default_saturation, + fp8_type::__default_interpret); +} + +// floatx2 -> fp8x2 +template <> +__inline__ __device__ uint16_t +scaled_vec_conversion(const float2& a, float scale) { + return __hip_cvt_float2_to_fp8x2(a / scale, fp8_type::__default_saturation, + fp8_type::__default_interpret); +} + +// floatx4 -> fp8x4 +template <> +__inline__ __device__ uint32_t +scaled_vec_conversion(const float4& a, float scale) { + union { + uint16_t ui16[2]; + uint32_t ui32; + } tmp; + tmp.ui16[0] = scaled_vec_conversion({a.x, a.y}, scale); + tmp.ui16[1] = scaled_vec_conversion({a.z, a.w}, scale); + return tmp.ui32; +} + #endif // ENABLE_FP8 + +template +__inline__ __device__ Tout convert(const Tin& x) { + #ifdef ENABLE_FP8 + if constexpr (kv_dt == Fp8KVCacheDataType::kFp8E4M3) { + return vec_conversion(x); + } + #endif + assert(false); + return {}; // Squash missing return statement warning +} + +template +__inline__ __device__ Tout scaled_convert(const Tin& x, const float scale) { + #ifdef ENABLE_FP8 + if constexpr (kv_dt == Fp8KVCacheDataType::kFp8E4M3) { + return scaled_vec_conversion(x, scale); + } + #endif + assert(false); + return {}; // Squash missing return statement warning +} + + // The following macro is used to dispatch the conversion function based on + // the data type of the key and value cache. The FN is a macro that calls a + // function with template. + #define DISPATCH_BY_KV_CACHE_DTYPE(SRC_DTYPE, KV_DTYPE, FN) \ + if (KV_DTYPE == "auto") { \ + if (SRC_DTYPE == at::ScalarType::Float) { \ + FN(float, float, vllm::Fp8KVCacheDataType::kAuto); \ + } else if (SRC_DTYPE == at::ScalarType::Half) { \ + FN(uint16_t, uint16_t, vllm::Fp8KVCacheDataType::kAuto); \ + } else if (SRC_DTYPE == at::ScalarType::BFloat16) { \ + FN(__nv_bfloat16, __nv_bfloat16, vllm::Fp8KVCacheDataType::kAuto); \ + } else { \ + TORCH_CHECK(false, "Unsupported input type of kv cache: ", SRC_DTYPE); \ + } \ + } else { \ + if (KV_DTYPE == "fp8" || KV_DTYPE == "fp8_e4m3") { \ + if (SRC_DTYPE == at::ScalarType::Float) { \ + FN(float, uint8_t, vllm::Fp8KVCacheDataType::kFp8E4M3); \ + } else if (SRC_DTYPE == at::ScalarType::Half) { \ + FN(uint16_t, uint8_t, vllm::Fp8KVCacheDataType::kFp8E4M3); \ + } else if (SRC_DTYPE == at::ScalarType::BFloat16) { \ + FN(__nv_bfloat16, uint8_t, vllm::Fp8KVCacheDataType::kFp8E4M3); \ + } else { \ + TORCH_CHECK(false, \ + "Unsupported input type of kv cache: ", SRC_DTYPE); \ + } \ + } else { \ + TORCH_CHECK(false, "Unsupported data type of kv cache: ", KV_DTYPE); \ + } \ + } + +} // namespace fp8 +#endif // USE_ROCM +} // namespace vllm diff --git a/tasks/hip2hip/campaign20/paged_attention_large/src/rocm/attention.cu b/tasks/hip2hip/campaign20/paged_attention_large/src/rocm/attention.cu new file mode 100644 index 00000000..906bd0df --- /dev/null +++ b/tasks/hip2hip/campaign20/paged_attention_large/src/rocm/attention.cu @@ -0,0 +1,3715 @@ +/* + * Copyright (c) 2024, The vLLM team. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#include +#include +#include +#include +#include +#include "../cuda_compat.h" + +#include +#include "../attention/dtype_fp8.cuh" +#include "../quantization/w8a8/fp8/amd/quant_utils.cuh" + +// ROCm 6.2 compatibility: map OCP fp8 types to FNUZ variants if OCP is absent +#if !defined(HIP_FP8_TYPE_OCP) +using __hip_fp8_e4m3 = __hip_fp8_e4m3_fnuz; +using __hip_fp8_e5m2 = __hip_fp8_e5m2_fnuz; +#endif + +#if defined(__HIPCC__) && \ + (defined(__gfx90a__) || defined(__gfx942__) || defined(__gfx950__)) + #define __HIP__GFX9__ +#endif + +#if defined(__HIPCC__) && (defined(__gfx942__) || defined(__gfx950__)) + #define __HIP__FP8MFMA__ +#endif + +#if defined(__HIPCC__) && (defined(__gfx1100__) || defined(__gfx1101__) || \ + defined(__gfx1150__) || defined(__gfx1151__)) + #define __HIP__GFX11__ +#endif + +#if defined(__HIPCC__) && (defined(__gfx1200__) || defined(__gfx1201__)) + #define __HIP__GFX12__ +#endif + +#if defined(NDEBUG) + #undef NDEBUG + #include + #define UNREACHABLE_CODE assert(false); + #define NDEBUG +#else + #define UNREACHABLE_CODE assert(false); +#endif + +#define MAX(a, b) ((a) > (b) ? (a) : (b)) +#define MIN(a, b) ((a) < (b) ? (a) : (b)) +#define DIVIDE_ROUND_UP(a, b) (((a) + (b) - 1) / (b)) + +enum class MFMAType { + F16 = 0, + Fp8 = 1, + Fp4 = 2, +}; + +#if defined(__HIP__GFX9__) + + #define GCN_MFMA_INSTR1 __builtin_amdgcn_mfma_f32_16x16x4f32 + #define GCN_MFMA_INSTR __builtin_amdgcn_mfma_f32_4x4x4f16 + +using floatx4 = __attribute__((__vector_size__(4 * sizeof(float)))) float; +using float16x4 = + __attribute__((__vector_size__(4 * sizeof(_Float16)))) _Float16; +typedef float16x4 _Half4; +using float16x2 = + __attribute__((__vector_size__(2 * sizeof(_Float16)))) _Float16; +typedef float16x2 _Half2; +typedef struct _Half8 { + _Half4 xy[2]; +} _Half8; + +using bit16_t = uint16_t; +using bit16x4 = __attribute__((__vector_size__(4 * sizeof(uint16_t)))) uint16_t; +typedef bit16x4 _B16x4; +typedef struct _B16x8 { + _B16x4 xy[2]; +} _B16x8; + +using _B8x8 = uint2; +using _B8x4 = int32_t; // used in builtins +using bit8_t = uint8_t; + +typedef struct _B8x16 { + _B8x8 xy[2]; +} _B8x16; + +template +__device__ __forceinline__ floatx4 gcn_mfma4x4x4_instr(const _B16x4& inpA, + const _B16x4& inpB, + const floatx4& inpC) { + if constexpr (std::is_same::value) { + return __builtin_amdgcn_mfma_f32_4x4x4f16(inpA, inpB, inpC, absz, cbid, + blgp); + } else if constexpr (std::is_same::value) { + return __builtin_amdgcn_mfma_f32_4x4x4bf16_1k(inpA, inpB, inpC, absz, cbid, + blgp); + } else { + static_assert(false, "unsupported 16b dtype"); + } +} + +template +__device__ __forceinline__ floatx4 gcn_mfma16x16x16_instr(const _B16x4& inpA, + const _B16x4& inpB, + const floatx4& inpC) { + if constexpr (std::is_same::value) { + return __builtin_amdgcn_mfma_f32_16x16x16f16(inpA, inpB, inpC, absz, cbid, + blgp); + } else if constexpr (std::is_same::value) { + return __builtin_amdgcn_mfma_f32_16x16x16bf16_1k(inpA, inpB, inpC, absz, + cbid, blgp); + } else { + static_assert(false, "unsupported 16b dtype"); + } +} + +template +__device__ __forceinline__ floatx4 gcn_mfma16x16x32_instr(const long& inpA, + const long& inpB, + const floatx4& inpC) { + if constexpr (std::is_same::value) { + return __builtin_amdgcn_mfma_f32_16x16x32_fp8_fp8(inpA, inpB, inpC, absz, + cbid, blgp); + } else if constexpr (std::is_same::value) { + return __builtin_amdgcn_mfma_f32_16x16x32_bf8_bf8(inpA, inpB, inpC, absz, + cbid, blgp); + } else { + static_assert(false, "unsupported 8b dtype"); + } +} + +template +__device__ __forceinline__ float to_float(const T& inp) { + if constexpr (std::is_same::value) { + return (float)inp; + } else if constexpr (std::is_same::value) { + return __bfloat162float(inp); + } else { + static_assert(false, "unsupported 16b dtype"); + } +} + +template +__device__ __forceinline__ T from_float(const float& inp) { + if constexpr (std::is_same::value) { + return (_Float16)inp; + } else if constexpr (std::is_same::value) { + return __float2bfloat16(inp); + } else { + static_assert(false, "unsupported 16b dtype"); + } +} + +template +__device__ __forceinline__ _B16x4 from_floatx4(const floatx4& inp) { + _B16x4 ret; + if constexpr (std::is_same::value) { + union h2cvt { + __half2 h2[2]; + _B16x4 b16x4; + } u; + u.h2[0] = __float22half2_rn(make_float2(inp[0], inp[1])); + u.h2[1] = __float22half2_rn(make_float2(inp[2], inp[3])); + return u.b16x4; + } else if constexpr (std::is_same::value) { + for (int i = 0; i < 4; i++) { + union fcvt { + uint32_t u32; + float f32; + } u; + u.f32 = inp[i]; + u.u32 += 0x7fff + ((u.u32 >> 16) & 1); // BF16 RNE with no nan/inf check + ret[i] = uint16_t(u.u32 >> 16); + } + return ret; + } else { + static_assert(false, "unsupported 16b dtype"); + } +} + +template +__device__ __forceinline__ _B16x4 addx4(const _B16x4& inp1, + const _B16x4& inp2) { + _B16x4 ret; + if constexpr (std::is_same::value) { + union h2cvt { + _B16x4 b16x4; + __half2 h2[2]; + } u1, u2, s; + u1.b16x4 = inp1; + u2.b16x4 = inp2; + s.h2[0] = u1.h2[0] + u2.h2[0]; + s.h2[1] = u1.h2[1] + u2.h2[1]; + return s.b16x4; + } else if constexpr (std::is_same::value) { + for (int i = 0; i < 4; i++) { + union fcvt { + float f32; + uint32_t i32; + } u1, u2, s; + u1.i32 = uint32_t(inp1[i]) << 16; + u2.i32 = uint32_t(inp2[i]) << 16; + s.f32 = u1.f32 + u2.f32; + ret[i] = uint16_t(s.i32 >> 16); + } + return ret; + } else { + static_assert(false, "unsupported 16b dtype"); + } +} + +__device__ __forceinline__ floatx4 to_float_fp8x4(const _B8x4& inp) { + // From MI300+ platforms, we have v_cvt_pk_f32_fp8 instruction + // to convert 2 packed fp8 to 2 packed fp32 values. + // However, in MI200 platforms, we only have v_cvt_f32_fp8 + // to convert fp8 values individually. So we added + // #else case for fewer instructions (# inst=2) in MI300+, + // and fallback to + // #if case for other platforms (# inst=4). + #if defined(__gfx90a__) + float4 f32x4 = vllm::fp8::vec_conversion( + *reinterpret_cast(&inp)); + return *reinterpret_cast(&f32x4); + #else // MI3xx+ optimized builtins + const auto f0 = __builtin_amdgcn_cvt_pk_f32_fp8(inp, false); + const auto f1 = __builtin_amdgcn_cvt_pk_f32_fp8(inp, true); + floatx4 ret; + ret[0] = f0[0]; + ret[1] = f0[1]; + ret[2] = f1[0]; + ret[3] = f1[1]; + return ret; + #endif +} + +template +__device__ __forceinline__ _B16x4 from_floatx4_rtz(const floatx4& inp) { + _B16x4 ret; + if constexpr (std::is_same::value) { + union h2cvt { + _Half2 h2[2]; + _B16x4 b16x4; + } u; + u.h2[0] = __builtin_amdgcn_cvt_pkrtz(inp[0], inp[1]); + u.h2[1] = __builtin_amdgcn_cvt_pkrtz(inp[2], inp[3]); + return u.b16x4; + } else if constexpr (std::is_same::value) { + for (int i = 0; i < 4; i++) { + union fcvt { + uint32_t i32; + float f32; + } u; + u.f32 = inp[i]; + ret[i] = uint16_t(u.i32 >> 16); + } + return ret; + } else { + static_assert(false, "unsupported 16b dtype"); + } +} + +template +__device__ __forceinline__ _B16x8 convert_b8x8_custom(const _B8x8 input) { + union { + _B8x8 b8x8; + _B8x4 b8x4[2]; + } tmp; + tmp.b8x8 = input; + _B16x8 ret; + for (int i = 0; i < 2; i++) { + ret.xy[i] = from_floatx4_rtz(to_float_fp8x4(tmp.b8x4[i])); + } + return ret; +} + +typedef union u64_cvt { + half f16x4[4]; + int16_t b16x4[4]; + _B8x8 b8x8; + _B16x4 b64; + int64_t i64; +} _T8x8; + +__device__ __forceinline__ _B8x8 convert_b16x8(const _B16x8& input, + _T8x8& Mtemp) { + _T8x8 Qtmp8x8; + + for (int i = 0; i < 2; i++) { + floatx4 q_out = {0, 0, 0, 0}; + q_out = gcn_mfma16x16x16_instr<_Float16, 0, 0, 0>(Mtemp.b64, input.xy[i], + q_out); + Qtmp8x8.b16x4[i * 2] = + __builtin_amdgcn_cvt_pk_fp8_f32(q_out[0], q_out[1], 0, false); + Qtmp8x8.b16x4[i * 2 + 1] = + __builtin_amdgcn_cvt_pk_fp8_f32(q_out[2], q_out[3], 0, false); + } + return Qtmp8x8.b8x8; +} + +__device__ float warpReduceMax(float val) { + for (int offset = warpSize / 2; offset > 0; offset /= 2) { + val = max( + val, __shfl_down(val, offset, WARP_SIZE)); // Using max() for reduction + } + return val; +} + +// grid (num_seqs, num_partitions,num_kv_heads) +// block (256) +// clang-format off +template +__global__ +__launch_bounds__(NUM_THREADS, 5) void paged_attention_ll4mi_QKV_mfma16_kernel( + const scalar_t* __restrict__ q, // [num_seqs, num_heads, head_size] + const cache_t* __restrict__ k_cache, // [num_blocks, num_kv_heads, head_size/x, block_size, x] + const cache_t* __restrict__ v_cache, // [num_blocks, num_kv_heads, head_size, block_size] + const int num_kv_heads, + const float scale, + const int* __restrict__ block_tables, // [num_seqs, max_num_blocks_per_seq] + const int* __restrict__ seq_lens, // [num_seqs] + const int* __restrict__ query_start_loc_ptr, // [num_seqs] + const int max_num_blocks_per_seq, + const float* __restrict__ alibi_slopes, // [num_heads] + const int q_stride, + const int kv_block_stride, + const int kv_head_stride, + float* __restrict__ exp_sums, // [num_seqs, num_heads, max_num_partitions] + float* __restrict__ max_logits, // [num_seqs, num_heads, max_num_partitions] + scalar_t* __restrict__ out, // [num_seqs, num_heads, max_num_partitions, head_size] + OUTT* __restrict__ final_out, // [num_seqs, num_heads, head_size] + int max_ctx_blocks, const float* k_scale, const float* v_scale) { + // clang-format on + constexpr int NWARPS = NUM_THREADS / WARP_SIZE; + const auto warpid = threadIdx.x / WARP_SIZE; + const auto laneid = threadIdx.x % WARP_SIZE; + const int lane4id = laneid % 4; + const int lane16id = laneid % 16; + const int rowid = laneid / 16; + + const auto seq_idx = blockIdx.x; + // NOTE queries with sequence len > 1 are prefills and taken care by another + // kernel. + if (query_start_loc_ptr != nullptr && + (query_start_loc_ptr[seq_idx + 1] - query_start_loc_ptr[seq_idx]) != 1) { + return; + } + + const auto partition_idx = blockIdx.y; + + constexpr int T_PAR_SIZE = 256; // token partition size set to 256 + + const auto max_num_partitions = gridDim.y; + + const int seq_len = seq_lens[seq_idx]; + + const int partition_start_token_idx = + partition_idx * T_PAR_SIZE; // partition_size; + // exit if partition is out of context for seq + if (partition_start_token_idx >= seq_len) { + return; + } + + constexpr int GQA_RATIO4 = DIVIDE_ROUND_UP(GQA_RATIO, 4); + + // shared_logits is used for multiple purposes + __shared__ _B16x4 shared_logits[NWARPS][4][16][4]; + + // for QK mfma16x16, layout is QHead/Tokenx16 across every 16 lanes, 16 Bytes + // HeadElements in each lane, 4x16B HeadElements across 4 rows of warp + constexpr int ROWS_PER_WARP = + WARP_SIZE / 16; // rows refers to 16 lanes; refer DDP (Data Parallel + // Processing) terminology + constexpr int CONTIGUOUS_KV_ELEMS_16B_LOAD = + 16 / sizeof(cache_t); // 8 for 16 bit cache type, 16 for 8 bit types + constexpr int QKHE_PER_FETCH = + CONTIGUOUS_KV_ELEMS_16B_LOAD * + ROWS_PER_WARP; // each fetch across a warp fetches these many elements + constexpr int QK_SIZE_RATIO = + sizeof(scalar_t) / + sizeof(cache_t); // 1 for 16bit types, 2 for 8bit types + constexpr int QKHELOOP = HEAD_SIZE / QKHE_PER_FETCH; // 4xQKHE_16B across + // warp + + _B16x8 Qlocal[QKHELOOP] + [QK_SIZE_RATIO]; // note that 16 contiguous elements of Q should + // be fetched per lane for 8 bit cache types : + // QK_SIZE_RATIO changes for this + + constexpr int CONTIGUOUS_SCALAR_ELEMS_16B = 16 / sizeof(scalar_t); + + constexpr int TOKENS_PER_WARP = + T_PAR_SIZE / + NWARPS; // sub partition of tokens per warp for qk calculation + constexpr int TLOOP = + TOKENS_PER_WARP / + 16; // each mfma16x16x16 instruction processes 16 tokens + + // can be interpreted as B8x16 for 8 bit types + _B16x8 Klocal[TLOOP][QKHELOOP]; + + const auto wg_start_head_idx = blockIdx.z * GQA_RATIO; + const auto wg_start_kv_head_idx = blockIdx.z; + const auto total_num_heads = gridDim.z * GQA_RATIO; + + // for QK mfma, tokens in multiples of TOKENS_PER_WARP are spread across warps + // each mfma takes QH16xT16x16HE across warp + // repeat mfmas across QKHELOOP dimension + // output layout from QKmfma : QH16xT4x4 16 qheads across 16 lanes, 16 tokens + // across 4 rows x 4 tokens per lane + + const int num_seq_blocks = DIVIDE_ROUND_UP(seq_len, BLOCK_SIZE); + const int last_seq_block = num_seq_blocks - 1; + + const int* block_table_seq = block_tables + seq_idx * max_num_blocks_per_seq; + + int kphysical_block_number[TLOOP]; + #if defined(__HIP__FP8MFMA__) + float q_max = 0; + float q_scale = 1.0; + #endif + + // fetch k physical block numbers + for (int token_depth = 0; token_depth < TLOOP; token_depth++) { + const int klocal_token_idx = + TOKENS_PER_WARP * warpid + token_depth * 16 + lane16id; + const int kglobal_token_idx = partition_start_token_idx + klocal_token_idx; + const int kblock_idx = (kglobal_token_idx < seq_len) + ? kglobal_token_idx / BLOCK_SIZE + : last_seq_block; + kphysical_block_number[token_depth] = block_table_seq[kblock_idx]; + } + + // fetch Q in shared across warps and then write to registers + const int local_qhead_idx = 4 * warpid + rowid; + const int global_qhead_idx = wg_start_head_idx + local_qhead_idx; + const int64_t query_start_off = static_cast( + query_start_loc_ptr ? query_start_loc_ptr[seq_idx] : seq_idx); + const scalar_t* q_ptr = + q + query_start_off * q_stride + global_qhead_idx * HEAD_SIZE; + + const int qhead_element = lane16id * CONTIGUOUS_SCALAR_ELEMS_16B; + if ((local_qhead_idx < GQA_RATIO) && (qhead_element < HEAD_SIZE)) { + const scalar_t* q_fetch_ptr = q_ptr + qhead_element; + const _B16x8* q_fetch_ptr_16B = + reinterpret_cast(q_fetch_ptr); + _B16x8 tmp = *q_fetch_ptr_16B; + if constexpr (KV_DTYPE == vllm::Fp8KVCacheDataType::kAuto) { + const int offset1 = + lane16id / + 4; // 16 contiguous chunks of head elems are spread across 4x4lanes + shared_logits[offset1][lane4id][local_qhead_idx][0] = tmp.xy[0]; + shared_logits[offset1][lane4id][local_qhead_idx][1] = tmp.xy[1]; + } else { + for (int i = 0; i < 2; i++) { + const int head_elem = lane16id * 2 + i; // element id in _B16x4 terms + const int offset3 = head_elem % 4; + const int offset2 = (head_elem / 4) % 4; + const int offset1 = head_elem / 4 / 4; + shared_logits[offset1][offset2][local_qhead_idx][offset3] = tmp.xy[i]; + } + } + } + __syncthreads(); + for (int qkhe_depth = 0; qkhe_depth < QKHELOOP; qkhe_depth++) { + for (int qkratio = 0; qkratio < QK_SIZE_RATIO; qkratio++) { + for (int i = 0; i < 2; i++) { + Qlocal[qkhe_depth][qkratio].xy[i] = + shared_logits[qkhe_depth][rowid][lane16id % GQA_RATIO] + [2 * qkratio + i]; + #if defined(__HIP__FP8MFMA__) + if constexpr (KV_DTYPE != vllm::Fp8KVCacheDataType::kAuto && + MFMA_TYPE == MFMAType::Fp8) { + scalar_t* qptr = + reinterpret_cast(&Qlocal[qkhe_depth][qkratio].xy[i]); + for (int k = 0; k < 4; k++) + q_max = fmax(fabs(to_float(qptr[k])), q_max); + } + #endif + } + } + } + + constexpr int KX = + 16 / sizeof(cache_t); // vLLM defines x as 16 Bytes of kv cache elements + const cache_t* k_ptr = k_cache + wg_start_kv_head_idx * kv_head_stride; + + const int row_head_elem = rowid * CONTIGUOUS_KV_ELEMS_16B_LOAD; + // fetch K values + for (int token_depth = 0; token_depth < TLOOP; token_depth++) { + const int64_t kblock_number = + static_cast(kphysical_block_number[token_depth]); + const cache_t* k_ptr2 = k_ptr + kblock_number * kv_block_stride; + const int klocal_token_idx = + TOKENS_PER_WARP * warpid + token_depth * 16 + lane16id; + const int kphysical_block_offset = klocal_token_idx % BLOCK_SIZE; + const cache_t* k_ptr3 = k_ptr2 + kphysical_block_offset * KX; + + for (int qkhe_depth = 0; qkhe_depth < QKHELOOP; qkhe_depth++) { + const int head_elem = row_head_elem + qkhe_depth * QKHE_PER_FETCH; + const int offset1 = head_elem / KX; + const int offset2 = head_elem % KX; + const cache_t* k_fetch_ptr = k_ptr3 + offset1 * BLOCK_SIZE * KX + offset2; + const _B16x8* k_fetch_ptr_16B = + reinterpret_cast(k_fetch_ptr); + Klocal[token_depth][qkhe_depth] = *k_fetch_ptr_16B; + } + } + + float alibi_slope; + if constexpr (ALIBI_ENABLED) { + const int alibi_head_idx = wg_start_head_idx + lane16id; + alibi_slope = (lane16id < GQA_RATIO) ? alibi_slopes[alibi_head_idx] : 0.f; + } + + constexpr int VTOKENS_PER_LANE = + TOKENS_PER_WARP / ROWS_PER_WARP; // 64/4 = 16 contiguous vtokens per lane + constexpr int VBLOCKS_PER_LANE = + 1; // assumes block size >=16, each lane can correspond to 1 block only + constexpr int VTLOOP = NWARPS; // corresponds to tokens across warps + constexpr int VTLANELOOP = DIVIDE_ROUND_UP( + VTOKENS_PER_LANE, + CONTIGUOUS_KV_ELEMS_16B_LOAD); // optimized for 16B fetches; assumes + // minimum block size is 16 + constexpr int VHELOOP = HEAD_SIZE / 16 / NWARPS; + + int vphysical_block_number[VTLOOP][VBLOCKS_PER_LANE]; + + // fetch v physical block numbers + for (int vtoken_depth = 0; vtoken_depth < VTLOOP; vtoken_depth++) { + for (int vblock_depth = 0; vblock_depth < VBLOCKS_PER_LANE; + vblock_depth++) { + const int vlocal_token_idx = + vtoken_depth * VTOKENS_PER_LANE * ROWS_PER_WARP + + rowid * VTOKENS_PER_LANE + vblock_depth * BLOCK_SIZE; + // Safe to use an int32_t here assuming we are working with < 2 billion + // tokens + const int vglobal_token_idx = + partition_start_token_idx + vlocal_token_idx; + const int vblock_idx = (vglobal_token_idx < seq_len) + ? vglobal_token_idx / BLOCK_SIZE + : last_seq_block; + vphysical_block_number[vtoken_depth][vblock_depth] = + block_table_seq[vblock_idx]; + } + } + + _B16x8 Vlocal[VTLOOP][VHELOOP][VTLANELOOP]; // this could be B8x16 too + + const cache_t* v_ptr = v_cache + wg_start_kv_head_idx * kv_head_stride + + ((rowid * VTOKENS_PER_LANE) % BLOCK_SIZE); + + // v fetches are 16head elems across lanes x 16 tokens per lane + for (int vhe_depth = 0; vhe_depth < VHELOOP; vhe_depth++) { + const int vhead_elem = vhe_depth * NWARPS * 16 + warpid * 16 + lane16id; + const cache_t* v_ptr2 = v_ptr + vhead_elem * BLOCK_SIZE; + + for (int vtoken_depth = 0; vtoken_depth < VTLOOP; vtoken_depth++) { + for (int vfetch_depth = 0; vfetch_depth < VTLANELOOP; vfetch_depth++) { + const int vblock_depth = 0; + const int64_t vblock_number = static_cast( + vphysical_block_number[vtoken_depth][vblock_depth]); + const cache_t* v_ptr3 = v_ptr2 + (vblock_number * kv_block_stride); + + const cache_t* v_fetch_ptr = + v_ptr3 + vfetch_depth * CONTIGUOUS_KV_ELEMS_16B_LOAD; + const _B16x8* v_fetch_ptr_16B = + reinterpret_cast(v_fetch_ptr); + Vlocal[vtoken_depth][vhe_depth][vfetch_depth] = *v_fetch_ptr_16B; + } + } + } + + // calculate post qk mfma scale + float scale2 = scale; + if constexpr (KV_DTYPE != vllm::Fp8KVCacheDataType::kAuto) { + // multiply by k_scale if fp8 kv cache + scale2 *= *k_scale; + #if defined(__HIP__FP8MFMA__) + q_max = warpReduceMax(q_max); + constexpr float FP8_E4M3_SCALE_TARGET = 224.0f; + if constexpr (MFMA_TYPE == MFMAType::Fp8) { + q_scale = q_max > 0 ? FP8_E4M3_SCALE_TARGET / q_max : 1.0f; + scale2 /= q_scale; + } + #endif + } + + floatx4 d_out[TLOOP]; + // qk mfma + for (int token_depth = 0; token_depth < TLOOP; token_depth++) { + d_out[token_depth] = {0}; + for (int qkhe_depth = 0; qkhe_depth < QKHELOOP; qkhe_depth++) { + if constexpr (KV_DTYPE == vllm::Fp8KVCacheDataType::kAuto) { + for (int qkratio = 0; qkratio < QK_SIZE_RATIO; qkratio++) { + for (int i = 0; i < 2; i++) { + d_out[token_depth] = gcn_mfma16x16x16_instr( + Klocal[token_depth][qkhe_depth].xy[i], + Qlocal[qkhe_depth][qkratio].xy[i], d_out[token_depth]); + } + } + } else { // kv cache dtype fp8 + auto Ktmp = Klocal[token_depth][qkhe_depth]; + _B8x16 Ktmp8x16 = *reinterpret_cast<_B8x16*>(&Ktmp); + for (int qkratio = 0; qkratio < QK_SIZE_RATIO; qkratio++) { + if constexpr (MFMA_TYPE == MFMAType::F16) { + _B8x8 Ktmp8x8 = Ktmp8x16.xy[qkratio]; + _B16x8 Klocaltmp = convert_b8x8_custom(Ktmp8x8); + for (int i = 0; i < 2; i++) { + d_out[token_depth] = gcn_mfma16x16x16_instr( + Klocaltmp.xy[i], Qlocal[qkhe_depth][qkratio].xy[i], + d_out[token_depth]); + } + } else { + #if defined(__HIP__FP8MFMA__) + _T8x8 Ktmp8x8, Qtmp8x8; + Ktmp8x8.b8x8 = Ktmp8x16.xy[qkratio]; + + for (int n = 0; n < 2; n++) { + scalar_t* qptr = reinterpret_cast( + &Qlocal[qkhe_depth][qkratio].xy[n]); + + Qtmp8x8.b16x4[n * 2] = + vllm::fp8::scaled_vec_conversion( + make_float2(to_float(qptr[0]), + to_float(qptr[1])), + q_scale); + Qtmp8x8.b16x4[n * 2 + 1] = + vllm::fp8::scaled_vec_conversion( + make_float2(to_float(qptr[2]), + to_float(qptr[3])), + q_scale); + } + + d_out[token_depth] = + gcn_mfma16x16x32_instr<__hip_fp8_e4m3, 0, 0, 0>( + Ktmp8x8.i64, Qtmp8x8.i64, d_out[token_depth]); + #else + UNREACHABLE_CODE + #endif + } + } + } + } + d_out[token_depth] *= scale2; + } + + const int qkout_token_idx = + partition_start_token_idx + TOKENS_PER_WARP * warpid + rowid * 4; + + // apply alibi + if constexpr (ALIBI_ENABLED) { + for (int token_depth = 0; token_depth < TLOOP; token_depth++) { + const int local_token_idx = qkout_token_idx + token_depth * 16; + const int alibi_offset = local_token_idx - seq_len + 1; + for (int i = 0; i < 4; i++) { + d_out[token_depth][i] += alibi_slope * (alibi_offset + i); + } + } + } + + // calculate qk_max and exp_sum per warp and write to shared memory + float qk_max = -FLT_MAX; + float exp_sum = 0.0f; + + for (int token_depth = 0; token_depth < TLOOP; token_depth++) { + const int local_token_idx = qkout_token_idx + token_depth * 16; + for (int i = 0; i < 4; i++) { + const float tmp = + (local_token_idx + i < seq_len) ? d_out[token_depth][i] : -FLT_MAX; + qk_max = fmaxf(qk_max, tmp); + } + } + + for (int mask = WARP_SIZE / 2; mask >= 16; mask /= 2) { + qk_max = fmaxf(qk_max, __shfl_xor(qk_max, mask)); + } + + for (int token_depth = 0; token_depth < TLOOP; token_depth++) { + const int local_token_idx = qkout_token_idx + token_depth * 16; + for (int i = 0; i < 4; i++) { + const float tmp = (local_token_idx + i < seq_len) + ? __expf(d_out[token_depth][i] - qk_max) + : 0.0f; + d_out[token_depth][i] = tmp; + exp_sum += tmp; + } + } + + for (int mask = WARP_SIZE / 2; mask >= 16; mask /= 2) { + exp_sum += __shfl_xor(exp_sum, mask); + } + + __syncthreads(); // sync before writing to shared mem + + float* shared_mem = reinterpret_cast(shared_logits); + if (laneid < 16) { + const int qk_max_offset = warpid * 16 + lane16id; + shared_mem[qk_max_offset] = qk_max; + const int exp_sum_offset = NWARPS * 16 + qk_max_offset; + shared_mem[exp_sum_offset] = exp_sum; + } + + __syncthreads(); + + // calculate partition qk_max and exp_sum + float partition_qk_max = -FLT_MAX; + float warp_qk_max_exp[NWARPS]; + float partition_exp_sum = 0.0f; + + for (int w = 0; w < NWARPS; w++) { + warp_qk_max_exp[w] = shared_mem[w * 16 + lane16id]; + partition_qk_max = fmaxf(partition_qk_max, warp_qk_max_exp[w]); + } + + for (int w = 0; w < NWARPS; w++) { + warp_qk_max_exp[w] = __expf(warp_qk_max_exp[w] - partition_qk_max); + partition_exp_sum += + shared_mem[NWARPS * 16 + w * 16 + lane16id] * warp_qk_max_exp[w]; + } + + const float inv_sum_scale = + __fdividef(1.f, partition_exp_sum + 1e-6f) * warp_qk_max_exp[warpid]; + + __syncthreads(); + + // disable rtz conversion due to its impact on accuracy. + constexpr bool LOGITS_RTZ_CONVERSION = false; + + #if defined(__HIP__FP8MFMA__) + int rowid_8x8 = rowid / 2; + int offset = rowid % 2; + #endif + + // write logits to shared mem + for (int token_depth = 0; token_depth < TLOOP; token_depth++) { + d_out[token_depth] *= inv_sum_scale; + if constexpr (MFMA_TYPE != MFMAType::Fp8) { + if constexpr (LOGITS_RTZ_CONVERSION) { + // use rtz conversion for better performance, with negligible impact on + // accuracy + shared_logits[warpid][token_depth][lane16id][rowid] = + from_floatx4_rtz(d_out[token_depth]); + } else { + shared_logits[warpid][token_depth][lane16id][rowid] = + from_floatx4(d_out[token_depth]); + } + } else { + #if defined(__HIP__FP8MFMA__) + // cast _B16x4* to _B8x8* + _T8x8& logits_8x8 = *reinterpret_cast<_T8x8*>( + &shared_logits[warpid][token_depth][lane16id][rowid_8x8]); + logits_8x8.b16x4[offset * 2] = __builtin_amdgcn_cvt_pk_fp8_f32( + d_out[token_depth][0], d_out[token_depth][1], 0, false); + logits_8x8.b16x4[offset * 2 + 1] = __builtin_amdgcn_cvt_pk_fp8_f32( + d_out[token_depth][2], d_out[token_depth][3], 0, false); + #else + UNREACHABLE_CODE + #endif + } + } + + // write out partition max_logits and exp_sum + if (threadIdx.x < GQA_RATIO) { + const int qhead_idx = lane16id; + const int64_t offset = static_cast(seq_idx) * + static_cast(total_num_heads) * + static_cast(max_num_partitions) + + (static_cast(wg_start_head_idx) + + static_cast(qhead_idx)) * + static_cast(max_num_partitions) + + static_cast(partition_idx); + max_logits[offset] = partition_qk_max; + exp_sums[offset] = partition_exp_sum; + } + + __syncthreads(); + + constexpr int ELEMS8_ELEMS4_RATIO = 8 / 4; + constexpr int ELEMS16_ELEMS8_RATIO = 16 / 8; + + _B16x4 outelems[VHELOOP]; + // Softmax V mfma + // v layout: 16he across lanes x 16 tokens per lane + for (int vhe_depth = 0; vhe_depth < VHELOOP; vhe_depth++) { + floatx4 tmp_out = {0}; + + for (int vtoken_depth = 0; vtoken_depth < VTLOOP; vtoken_depth++) { + if constexpr (KV_DTYPE == vllm::Fp8KVCacheDataType::kAuto) { + for (int vfetch_depth = 0; vfetch_depth < VTLANELOOP; vfetch_depth++) { + for (int i = 0; i < ELEMS8_ELEMS4_RATIO; i++) { + const int offset = rowid * VTLANELOOP * ELEMS8_ELEMS4_RATIO + + vfetch_depth * ELEMS8_ELEMS4_RATIO + i; + const int offset1 = offset % ROWS_PER_WARP; + const int offset2 = offset / ROWS_PER_WARP; + // output format is 16 qheads across 16 lanes, 16 head elems spread + // across 4 rows + tmp_out = gcn_mfma16x16x16_instr( + Vlocal[vtoken_depth][vhe_depth][vfetch_depth].xy[i], + shared_logits[vtoken_depth][offset2][lane16id][offset1], + tmp_out); + } + } + // KV cache fp8 + } else { + for (int vfetch_depth = 0; vfetch_depth < VTLANELOOP; vfetch_depth++) { + _B16x8 Vtmp = Vlocal[vtoken_depth][vhe_depth][vfetch_depth]; + // reinterpret V format as 16 elements of 8bits + _B8x16 Vtmp8x16 = *reinterpret_cast<_B8x16*>(&Vtmp); + for (int j = 0; j < ELEMS16_ELEMS8_RATIO; j++) { + _B8x8 Vtmp8x8 = Vtmp8x16.xy[j]; + if constexpr (MFMA_TYPE == MFMAType::F16) { + _B16x8 Vlocaltmp = convert_b8x8_custom(Vtmp8x8); + for (int i = 0; i < ELEMS8_ELEMS4_RATIO; i++) { + const int offset = + rowid * ELEMS16_ELEMS8_RATIO * ELEMS8_ELEMS4_RATIO + + j * ELEMS8_ELEMS4_RATIO + i; + const int offset1 = offset % ROWS_PER_WARP; + const int offset2 = offset / ROWS_PER_WARP; + // output format is 16 qheads across 16 lanes, 16 head elems + // spread across 4 rows + tmp_out = gcn_mfma16x16x16_instr( + Vlocaltmp.xy[i], + shared_logits[vtoken_depth][offset2][lane16id][offset1], + tmp_out); + } + } else { + #if defined(__HIP__FP8MFMA__) + for (int i = 0; i < ELEMS8_ELEMS4_RATIO / 2; i++) { + const int offset = + rowid * ELEMS16_ELEMS8_RATIO * ELEMS8_ELEMS4_RATIO + + j * ELEMS8_ELEMS4_RATIO + i; + const int offset1 = (offset % ROWS_PER_WARP) / 2; + const int offset2 = offset / ROWS_PER_WARP; + // output format is 16 qheads across 16 lanes, 16 head elems + // spread across 4 rows + tmp_out = gcn_mfma16x16x32_instr<__hip_fp8_e4m3, 0, 0, 0>( + reinterpret_cast<_T8x8*>(&Vtmp8x8)->i64, + reinterpret_cast<_T8x8*>( + &shared_logits[vtoken_depth][offset2][lane16id] + [offset1]) + ->i64, + tmp_out); + } + #else + UNREACHABLE_CODE + #endif + } + } + } + } + } + // apply post Softmax V mfma v_scale + if constexpr (KV_DTYPE != vllm::Fp8KVCacheDataType::kAuto) { + tmp_out *= *v_scale; + } + outelems[vhe_depth] = from_floatx4(tmp_out); + } + + __syncthreads(); + + // store Softmax-V mfma output to shared mem + for (int vhe_depth = 0; vhe_depth < VHELOOP; vhe_depth++) { + // lane16 id head dimension; rowid head element dimension + shared_logits[warpid][vhe_depth][lane16id][rowid] = outelems[vhe_depth]; + } + + __syncthreads(); + + // write to tmp_out with coalesced writes after reading from shared mem + if (warpid == 0) { + _B16x8 vout[GQA_RATIO4]; + // each lane writes out 16Bytes of tmp_out along head elem dimension + const int head_elem_idx = lane16id * 8; + if (head_elem_idx < HEAD_SIZE) { + for (int h = 0; h < GQA_RATIO4; h++) { + const int local_head_idx = 4 * h + rowid; + const int offset1 = (head_elem_idx / 16) % 4; + const int offset2 = head_elem_idx / 16 / NWARPS; + const int offset3 = (head_elem_idx / 4) % 4; + for (int i = 0; i < 2; i++) { + vout[h].xy[i] = + shared_logits[offset1][offset2][local_head_idx][offset3 + i]; + } + } + + const int64_t hsz_maxp_mult = + static_cast(HEAD_SIZE * max_num_partitions); + scalar_t* out_ptr = out + seq_idx * total_num_heads * hsz_maxp_mult + + partition_idx * HEAD_SIZE; + for (int h = 0; h < GQA_RATIO4; h++) { + const int local_head_idx = 4 * h + rowid; + if (local_head_idx < GQA_RATIO) { + const int64_t out_head_idx = + static_cast(wg_start_head_idx + local_head_idx); + scalar_t* out_ptr2 = out_ptr + out_head_idx * hsz_maxp_mult; + scalar_t* out_ptr3 = out_ptr2 + head_elem_idx; + _B16x8* out_ptr_B16x8 = reinterpret_cast<_B16x8*>(out_ptr3); + *out_ptr_B16x8 = vout[h]; + } + } + } + } +} + +// grid (num_seqs, num_partitions, num_kv_heads) +// block (256 : partition size) +// each WG handles 1 partition per sequence +// clang-format off +template +__global__ +__launch_bounds__(NUM_THREADS) void paged_attention_ll4mi_QKV_mfma4_kernel( + const scalar_t* __restrict__ q, // [num_seqs, num_heads, head_size] + const cache_t* __restrict__ k_cache, // [num_blocks, num_kv_heads, head_size/x, block_size, x] + const cache_t* __restrict__ v_cache, // [num_blocks, num_kv_heads, head_size, block_size] + const int num_kv_heads, + const float scale, + const int* __restrict__ block_tables, // [num_seqs, max_num_blocks_per_seq] + const int* __restrict__ seq_lens, // [num_seqs] + const int* __restrict__ query_start_loc_ptr, // [num_seqs] + const int max_num_blocks_per_seq, + const float* __restrict__ alibi_slopes, // [num_heads] + const int q_stride, + const int kv_block_stride, + const int kv_head_stride, + float* __restrict__ exp_sums, // [num_seqs, num_heads, max_num_partitions] + float* __restrict__ max_logits, // [num_seqs, num_heads, max_num_partitions] + scalar_t* __restrict__ out, // [num_seqs, num_heads, max_num_partitions, head_size] + OUTT* __restrict__ final_out, // [num_seqs, num_heads, head_size] + int max_ctx_blocks, const float* k_scale, const float* v_scale) { + // clang-format on + constexpr int NWARPS = NUM_THREADS / WARP_SIZE; + const auto warpid = threadIdx.x / WARP_SIZE; + const auto laneid = threadIdx.x % WARP_SIZE; + const int lane4id = laneid % 4; + + const auto seq_idx = blockIdx.x; + // NOTE queries with sequence len > 1 are prefills and taken care by another + // kernel. + if (query_start_loc_ptr != nullptr && + (query_start_loc_ptr[seq_idx + 1] - query_start_loc_ptr[seq_idx] != 1)) { + return; + } + const auto partition_idx = blockIdx.y; + const auto partition_size = blockDim.x; + const auto max_num_partitions = gridDim.y; + + const int seq_len = seq_lens[seq_idx]; + const int partition_start_token_idx = partition_idx * partition_size; + // exit if partition is out of context for seq + if (partition_start_token_idx >= seq_len) { + return; + } + // every 4 lanes fetch 4 different qheads + // qhloop = num loops over qhead dimension + constexpr int QHLOOP = DIVIDE_ROUND_UP(GQA_RATIO, 4); + constexpr int GQA_RATIO4 = 4 * QHLOOP; + __shared__ float shared_qk_max[NWARPS][GQA_RATIO4 + 1]; + __shared__ float shared_exp_sum[NWARPS][GQA_RATIO4 + 1]; + _B16x8 Qlocal[QHLOOP]; + constexpr int x = 16 / sizeof(scalar_t); + // kheloop = num loops over head_size for 16Bytes of Q/dequantized K elements + constexpr int KHELOOP = HEAD_SIZE / x; + _B16x8 Klocal[KHELOOP]; + _B8x8 Klocalb8[KHELOOP]; + // for SoftMax-V Gemm, V head_size dimension is distributed across warp + // vheloop = num loops to cover v head size dimension + constexpr int VHELOOP = HEAD_SIZE / WARP_SIZE; + // softmax out has warp_size tokens across warp + // vtloop = num loops to cover warp_size(64) tokens with 16Bytes of + // dequantized V elements + constexpr int VTLOOP = WARP_SIZE / 8; + // num vblocks to cover warp_size(64) v elements + constexpr int VBLOCKS = 8 * VTLOOP / BLOCK_SIZE; + int vphysical_blocks[VBLOCKS]; + _B16x8 Vlocal[VHELOOP][VTLOOP]; + _B8x8 Vlocalb8[VHELOOP][VTLOOP]; + floatx4 d_out[QHLOOP]; + float qk_max[QHLOOP]; + + __shared__ _B16x4 vout_shared[QHLOOP][VHELOOP][WARP_SIZE][NWARPS + 1]; + + for (int h = 0; h < QHLOOP; h++) { + d_out[h] = {0}; + qk_max[h] = -FLT_MAX; + } + + const auto wg_start_head_idx = blockIdx.z * GQA_RATIO; + const auto wg_start_kv_head_idx = blockIdx.z; + + const int warp_start_token_idx = + partition_start_token_idx + warpid * WARP_SIZE; + + if (warp_start_token_idx >= seq_len) { // warp out of context + #pragma unroll + for (int h = 0; h < GQA_RATIO4; h++) { + shared_qk_max[warpid][h] = -FLT_MAX; + shared_exp_sum[warpid][h] = 0.0f; + } + } else { // warp within context + + const int num_seq_blocks = DIVIDE_ROUND_UP(seq_len, BLOCK_SIZE); + const int last_seq_block = num_seq_blocks - 1; + + const int* block_table = block_tables + seq_idx * max_num_blocks_per_seq; + // token id within partition + const auto local_token_idx = threadIdx.x; + // token id within sequence + const int global_token_idx = partition_start_token_idx + local_token_idx; + + // fetch block number for k + const int block_idx = (global_token_idx < seq_len) + ? global_token_idx / BLOCK_SIZE + : last_seq_block; + + // fetch k physical block number + // int32 physical_block_number leads to overflow when multiplied with + // kv_block_stride + const int64_t physical_block_number = + static_cast(block_table[block_idx]); + + // fetch vphysical block numbers up front + const int warp_start_block_idx = warp_start_token_idx / BLOCK_SIZE; + for (int b = 0; b < VBLOCKS; b++) { + const int vblock_idx = warp_start_block_idx + b; + const int vblock_idx_ctx = + (vblock_idx <= last_seq_block) ? vblock_idx : last_seq_block; + vphysical_blocks[b] = block_table[vblock_idx_ctx]; + } + + // fetch q elements + // every 4 lanes fetch 8 elems, so warp fetches 8*16 = 128 elemsc + const int64_t query_start_off = static_cast( + query_start_loc_ptr ? query_start_loc_ptr[seq_idx] : seq_idx); + const scalar_t* q_ptr = + q + query_start_off * q_stride + wg_start_head_idx * HEAD_SIZE; + const _B16x8* q_ptrh8 = reinterpret_cast(q_ptr); + const int qhead_elemh8 = laneid / 4; + + for (int h = 0; h < QHLOOP - 1; h++) { + const int qhead_idx = h * 4 + lane4id; + Qlocal[h] = q_ptrh8[qhead_idx * HEAD_SIZE / 8 + qhead_elemh8]; + } + const int final_qhead_idx = 4 * (QHLOOP - 1) + lane4id; + if (final_qhead_idx < GQA_RATIO) { + Qlocal[QHLOOP - 1] = + q_ptrh8[final_qhead_idx * HEAD_SIZE / 8 + qhead_elemh8]; + } else { + Qlocal[QHLOOP - 1].xy[0] = {0}; + Qlocal[QHLOOP - 1].xy[1] = {0}; + } + + // fetch k elements + const cache_t* k_ptr = k_cache + physical_block_number * kv_block_stride + + wg_start_kv_head_idx * kv_head_stride; + + // physical_block_offset is already cast in terms of _B16x8 + const int physical_block_offset = local_token_idx % BLOCK_SIZE; + + // each K fetch is for 8 elements of cache_t which are later dequantized to + // scalar_t for fp8 + if constexpr (KV_DTYPE == vllm::Fp8KVCacheDataType::kAuto) { + const _B16x8* k_ptrh8 = reinterpret_cast(k_ptr); + for (int d = 0; d < KHELOOP; d++) { + Klocal[d] = k_ptrh8[d * BLOCK_SIZE + physical_block_offset]; + } + } else { + // vllm defines X as 16 Bytes of elements of cache_t + constexpr int X = 16 / sizeof(cache_t); + const cache_t* k_ptr2 = k_ptr + physical_block_offset * X; + for (int d = 0; d < KHELOOP; d++) { + const int head_elem = d * 8; + const int offset1 = head_elem / X; + const int offset2 = head_elem % X; + const cache_t* k_ptr3 = k_ptr2 + offset1 * BLOCK_SIZE * X + offset2; + Klocalb8[d] = *reinterpret_cast(k_ptr3); + } + } + + // optional alibi fetch + float alibi_slope[QHLOOP]; + if constexpr (ALIBI_ENABLED) { + for (int h = 0; h < QHLOOP; h++) { + const int qhead_idx = h * 4 + lane4id; + alibi_slope[h] = (qhead_idx < GQA_RATIO) + ? alibi_slopes[wg_start_head_idx + qhead_idx] + : 0.f; + } + } + + const cache_t* v_ptr = v_cache + wg_start_kv_head_idx * kv_head_stride; + // fetch vcache in kv cache auto case + if constexpr (KV_DTYPE == vllm::Fp8KVCacheDataType::kAuto) { + const _B16x8* v_ptrh8 = reinterpret_cast(v_ptr); + // iterate over each v block + for (int b = 0; b < VBLOCKS; b++) { + // int32 physical_block_number leads to overflow when multiplied with + // kv_block_stride + const int64_t vphysical_block_number = + static_cast(vphysical_blocks[b]); + const _B16x8* v_ptrh8b = + v_ptrh8 + (vphysical_block_number * kv_block_stride) / 8; + // iterate over each head elem (within head_size) + for (int h = 0; h < VHELOOP; h++) { + const int head_size_elem = h * WARP_SIZE + laneid; + const _B16x8* v_ptrh8be = v_ptrh8b + head_size_elem * BLOCK_SIZE / 8; + // iterate over all velems within block + for (int d = 0; d < BLOCK_SIZE / 8; d++) { + Vlocal[h][b * BLOCK_SIZE / 8 + d] = v_ptrh8be[d]; + } + } + } + } // if constexpr (KV_DTYPE == vllm::Fp8KVCacheDataType::kAuto) + // fetch vcache in fp8 case + else { // if constexpr (KV_DTYPE != vllm::Fp8KVCacheDataType::kAuto) + const _B8x8* v_ptrh8 = reinterpret_cast(v_ptr); + // iterate over each v block + for (int b = 0; b < VBLOCKS; b++) { + // int32 physical_block_number leads to overflow when multiplied with + // kv_block_stride + const int64_t vphysical_block_number = + static_cast(vphysical_blocks[b]); + const _B8x8* v_ptrh8b = + v_ptrh8 + (vphysical_block_number * kv_block_stride) / 8; + // iterate over each head elem (within head_size) + for (int h = 0; h < VHELOOP; h++) { + const int head_size_elem = h * WARP_SIZE + laneid; + const _B8x8* v_ptrh8be = v_ptrh8b + head_size_elem * BLOCK_SIZE / 8; + // iterate over all velems within block + for (int d = 0; d < BLOCK_SIZE / 8; d++) { + Vlocalb8[h][b * BLOCK_SIZE / 8 + d] = v_ptrh8be[d]; + } + } + } + } + + #define QK_mfma(x) \ + if constexpr (KV_DTYPE != vllm::Fp8KVCacheDataType::kAuto) { \ + Klocal[x] = convert_b8x8_custom(Klocalb8[x]); \ + } \ + for (int h = 0; h < QHLOOP; h++) { \ + d_out[h] = gcn_mfma4x4x4_instr( \ + Qlocal[h].xy[0], Klocal[x].xy[0], d_out[h]); \ + d_out[h] = gcn_mfma4x4x4_instr( \ + Qlocal[h].xy[1], Klocal[x].xy[1], d_out[h]); \ + } + // QK mfma with Q mfma block broadcast + // Q values across head_size dimension stored across lanes + // K values across head_size dimension are stored depthwise within lane + // Q broadcast with absz, cbid of mfma instruction + QK_mfma(0); + QK_mfma(1); + QK_mfma(2); + QK_mfma(3); + QK_mfma(4); + QK_mfma(5); + QK_mfma(6); + QK_mfma(7); + // below only needed for head size 128 + if constexpr (KHELOOP > 8) { + QK_mfma(8); + QK_mfma(9); + QK_mfma(10); + QK_mfma(11); + QK_mfma(12); + QK_mfma(13); + QK_mfma(14); + QK_mfma(15); + } + #undef QK_mfma + + float scale2 = scale; + if constexpr (KV_DTYPE != vllm::Fp8KVCacheDataType::kAuto) { + // post mfma scaling for fp8 + scale2 *= *k_scale; + } + + for (int h = 0; h < QHLOOP; h++) { + d_out[h] *= scale2; + } + + // transpose d_out so that 4 token ids are in each lane, and 4 heads are + // across 4 lanes + for (int h = 0; h < QHLOOP; h++) { + floatx4 tmp = {0}; + for (int i = 0; i < 4; i++) { + const float B = (lane4id == i) ? 1.0f : 0.0f; + tmp = __builtin_amdgcn_mfma_f32_4x4x1f32(d_out[h][i], B, tmp, 0, 0, 0); + } + d_out[h] = tmp; + } + + const int lane4_token_idx = 4 * (global_token_idx >> 2); + + if constexpr (ALIBI_ENABLED) { + const int alibi_offset = lane4_token_idx - seq_len + 1; + for (int h = 0; h < QHLOOP; h++) { + for (int i = 0; i < 4; i++) { + d_out[h][i] += alibi_slope[h] * (alibi_offset + i); + } + } + } + + const int bpermute_mask = 4 * (16 * ((laneid >> 2) % 4) + lane4id); + + for (int h = 0; h < QHLOOP; h++) { + qk_max[h] = -FLT_MAX; + for (int i = 0; i < 4; i++) { + qk_max[h] = (lane4_token_idx + i < seq_len) + ? fmaxf(qk_max[h], d_out[h][i]) + : qk_max[h]; + } + + // for (int mask = WARP_SIZE / 2; mask >= 4; mask /= 2) { + // qk_max[h] = fmaxf(qk_max[h], __shfl_xor(qk_max[h], mask)); + // } + // faster version of above code with dpp + asm("v_nop\n v_nop\n v_max_f32_dpp %0, %1, %2 row_ror:4" + : "=v"(qk_max[h]) + : "v"(qk_max[h]), "v"(qk_max[h])); + asm("v_nop\n v_nop\n v_max_f32_dpp %0, %1, %2 row_ror:8" + : "=v"(qk_max[h]) + : "v"(qk_max[h]), "v"(qk_max[h])); + + auto tmp = __builtin_amdgcn_ds_bpermute( + bpermute_mask, *reinterpret_cast(&qk_max[h])); + qk_max[h] = *reinterpret_cast(&tmp); + asm("v_nop\n v_nop\n v_max_f32_dpp %0, %1, %2 row_ror:4" + : "=v"(qk_max[h]) + : "v"(qk_max[h]), "v"(qk_max[h])); + asm("v_nop\n v_nop\n v_max_f32_dpp %0, %1, %2 row_ror:8" + : "=v"(qk_max[h]) + : "v"(qk_max[h]), "v"(qk_max[h])); + } + + float exp_sum[QHLOOP]; + for (int h = 0; h < QHLOOP; h++) { + exp_sum[h] = 0.0f; + for (int i = 0; i < 4; i++) { + d_out[h][i] = (lane4_token_idx + i < seq_len) + ? __expf(d_out[h][i] - qk_max[h]) + : 0.0f; + exp_sum[h] += d_out[h][i]; + } + // for (int mask = WARP_SIZE / 2; mask >= 4; mask /= 2) { + // exp_sum[h] += __shfl_xor(exp_sum[h], mask); + // } + // faster version of above code with dpp + asm("v_nop\n v_nop\n v_add_f32_dpp %0, %1, %2 row_ror:4" + : "=v"(exp_sum[h]) + : "v"(exp_sum[h]), "v"(exp_sum[h])); + asm("v_nop\n v_nop\n v_add_f32_dpp %0, %1, %2 row_ror:8" + : "=v"(exp_sum[h]) + : "v"(exp_sum[h]), "v"(exp_sum[h])); + + auto tmp = __builtin_amdgcn_ds_bpermute( + bpermute_mask, *reinterpret_cast(&exp_sum[h])); + exp_sum[h] = *reinterpret_cast(&tmp); + asm("v_nop\n v_nop\n v_add_f32_dpp %0, %1, %2 row_ror:4" + : "=v"(exp_sum[h]) + : "v"(exp_sum[h]), "v"(exp_sum[h])); + asm("v_nop\n v_nop\n v_add_f32_dpp %0, %1, %2 row_ror:8" + : "=v"(exp_sum[h]) + : "v"(exp_sum[h]), "v"(exp_sum[h])); + } + + if (laneid < 4) { + for (int h = 0; h < QHLOOP; h++) { + const int head_idx = 4 * h + lane4id; + shared_qk_max[warpid][head_idx] = qk_max[h]; + shared_exp_sum[warpid][head_idx] = exp_sum[h]; + } + } + } // warp within context + + __syncthreads(); + + const auto num_heads = gridDim.z * GQA_RATIO; + float* max_logits_ptr = + max_logits + seq_idx * num_heads * max_num_partitions + partition_idx; + float* exp_sums_ptr = + exp_sums + seq_idx * num_heads * max_num_partitions + partition_idx; + // calculate qk_max and exp_sums for partition + for (int h = 0; h < QHLOOP; h++) { + float global_qk_max = -FLT_MAX; + float warp_qk_max[NWARPS]; + const int head_idx = 4 * h + lane4id; + for (int w = 0; w < NWARPS; w++) { + warp_qk_max[w] = shared_qk_max[w][head_idx]; + global_qk_max = fmaxf(global_qk_max, warp_qk_max[w]); + } + float global_exp_sum = 0.0f; + for (int w = 0; w < NWARPS; w++) { + global_exp_sum += + shared_exp_sum[w][head_idx] * __expf(warp_qk_max[w] - global_qk_max); + } + if (head_idx < GQA_RATIO) { + max_logits_ptr[(wg_start_head_idx + head_idx) * max_num_partitions] = + global_qk_max; + exp_sums_ptr[(wg_start_head_idx + head_idx) * max_num_partitions] = + global_exp_sum; + } + const float global_inv_sum_scale = __fdividef(1.f, global_exp_sum + 1e-6f) * + __expf(qk_max[h] - global_qk_max); + d_out[h] *= global_inv_sum_scale; + } + constexpr bool LOGITS_RTZ_CONVERSION = false; + // logits[h] -> every 4 lanes hold 4 heads, each lane holds 4 tokens, there + // are 4x16 tokens across warp + _B16x4 logits[QHLOOP]; + for (int h = 0; h < QHLOOP; h++) { + if constexpr (LOGITS_RTZ_CONVERSION) { + // use rtz for faster performance with no perceivable accuracy loss + logits[h] = from_floatx4_rtz(d_out[h]); + } else { + logits[h] = from_floatx4(d_out[h]); + } + } + + if (warp_start_token_idx >= seq_len) { // warp out of context + for (int qh = 0; qh < QHLOOP; qh++) { + for (int vh = 0; vh < VHELOOP; vh++) { + vout_shared[qh][vh][laneid][warpid] = {0}; + } + } + } else { // warp in context + #define SV_mfma(x) \ + if constexpr (KV_DTYPE != vllm::Fp8KVCacheDataType::kAuto) { \ + Vlocal[vh][x] = convert_b8x8_custom(Vlocalb8[vh][x]); \ + } \ + for (int qh = 0; qh < QHLOOP; qh++) { \ + acc[qh] = gcn_mfma4x4x4_instr( \ + logits[qh], Vlocal[vh][x].xy[0], acc[qh]); \ + acc[qh] = gcn_mfma4x4x4_instr( \ + logits[qh], Vlocal[vh][x].xy[1], acc[qh]); \ + } + + for (int vh = 0; vh < VHELOOP; vh++) { + floatx4 acc[QHLOOP]; + for (int qh = 0; qh < QHLOOP; qh++) { + acc[qh] = {0}; + } + // SoftMax-V calculation + // logits -> token dimension is distributed across lanes + // Vlocal -> token dimension is depthwise within lane + // uses mfma instruction block broadcast for logits + SV_mfma(0); + SV_mfma(1); + SV_mfma(2); + SV_mfma(3); + SV_mfma(4); + SV_mfma(5); + SV_mfma(6); + SV_mfma(7); + + for (int qh = 0; qh < QHLOOP; qh++) { + if constexpr (KV_DTYPE != vllm::Fp8KVCacheDataType::kAuto) { + // post mfma v scale for fp8 + acc[qh] *= *v_scale; + } + vout_shared[qh][vh][laneid][warpid] = from_floatx4(acc[qh]); + } + } + + #undef SV_mfma + } // warp in context + + __syncthreads(); + + // final write to tmp_out after vout accumulation + if (warpid == 0) { + _B16x4 vout[QHLOOP][VHELOOP]; + // iterate across heads + for (int qh = 0; qh < QHLOOP; qh++) { + // iterate over each v head elem (within head_size) + for (int vh = 0; vh < VHELOOP; vh++) { + vout[qh][vh] = {0}; + for (int w = 0; w < NWARPS; w++) { + vout[qh][vh] = + addx4(vout[qh][vh], vout_shared[qh][vh][laneid][w]); + } + } + } + + scalar_t* out_ptr = out + + seq_idx * num_heads * max_num_partitions * HEAD_SIZE + + partition_idx * HEAD_SIZE; + const int out_num_partitions = max_num_partitions; + bit16_t* out_ptr_b16 = reinterpret_cast(out_ptr); + for (int qh = 0; qh < QHLOOP; qh++) { + for (int vh = 0; vh < VHELOOP; vh++) { + const int head_size_elem = vh * WARP_SIZE + laneid; + for (int i = 0; i < 4; i++) { + const int head_idx = 4 * qh + i; + if (head_idx < GQA_RATIO) { + out_ptr_b16[(wg_start_head_idx + head_idx) * out_num_partitions * + HEAD_SIZE + + head_size_elem] = vout[qh][vh][i]; + } + } + } + } + } // warpid == 0 +} + +// Grid: (num_heads, num_seqs). +template +__global__ +__launch_bounds__(NUM_THREADS) void paged_attention_ll4mi_reduce_kernel( + OUTT* __restrict__ out, // [num_seqs, num_heads, head_size] + const float* __restrict__ exp_sums, // [num_seqs, num_heads, + // max_num_partitions] + const float* __restrict__ max_logits, // [num_seqs, num_heads, + // max_num_partitions] + const scalar_t* __restrict__ tmp_out, // [num_seqs, num_heads, + // max_num_partitions, head_size] + const int* __restrict__ seq_lens, // [num_seqs] + const int* __restrict__ query_start_loc_ptr, // [num_seqs] + const int max_num_partitions, const float* __restrict__ fp8_out_scale_ptr) { + const auto num_heads = gridDim.x; + const auto head_idx = blockIdx.x; + const auto seq_idx = blockIdx.y; + + // NOTE queries with sequence len > 1 are prefills and taken care by another + // kernel. + if (query_start_loc_ptr != nullptr && + (query_start_loc_ptr[seq_idx + 1] - query_start_loc_ptr[seq_idx] != 1)) { + return; + } + + const int seq_len = seq_lens[seq_idx]; + const int num_partitions = DIVIDE_ROUND_UP(seq_len, PARTITION_SIZE); + const auto warpid = threadIdx.x / WARP_SIZE; + + __shared__ float shared_global_exp_sum; + // max num partitions supported is warp_size * NPAR_LOOPS + __shared__ float shared_exp_sums[NPAR_LOOPS * WARP_SIZE]; + + if (warpid == 0) { + const float* max_logits_ptr = max_logits + + seq_idx * num_heads * max_num_partitions + + head_idx * max_num_partitions; + + // valid partition is the last valid partition in case threadid > num + // partitions + int valid_partition[NPAR_LOOPS]; + float reg_max_logit[NPAR_LOOPS]; + const int last_valid_partition = num_partitions - 1; + + #pragma unroll + for (int i = 0; i < NPAR_LOOPS; i++) { + const auto partition_no = i * WARP_SIZE + threadIdx.x; + valid_partition[i] = + (partition_no < num_partitions) ? partition_no : last_valid_partition; + } + #pragma unroll + for (int i = 0; i < NPAR_LOOPS; i++) { + reg_max_logit[i] = max_logits_ptr[valid_partition[i]]; + } + float max_logit = reg_max_logit[0]; + #pragma unroll + for (int i = 1; i < NPAR_LOOPS; i++) { + max_logit = fmaxf(max_logit, reg_max_logit[i]); + } + + #pragma unroll + for (int mask = WARP_SIZE / 2; mask >= 1; mask /= 2) { + max_logit = fmaxf(max_logit, __shfl_xor(max_logit, mask)); + } + + const float* exp_sums_ptr = exp_sums + + seq_idx * num_heads * max_num_partitions + + head_idx * max_num_partitions; + + float rescaled_exp_sum[NPAR_LOOPS]; + #pragma unroll + for (int i = 0; i < NPAR_LOOPS; i++) { + rescaled_exp_sum[i] = exp_sums_ptr[valid_partition[i]]; + } + #pragma unroll + for (int i = 0; i < NPAR_LOOPS; i++) { + const auto partition_no = i * WARP_SIZE + threadIdx.x; + rescaled_exp_sum[i] *= (partition_no < num_partitions) + ? expf(reg_max_logit[i] - max_logit) + : 0.0f; + } + float global_exp_sum = rescaled_exp_sum[0]; + #pragma unroll + for (int i = 1; i < NPAR_LOOPS; i++) { + global_exp_sum += rescaled_exp_sum[i]; + } + #pragma unroll + for (int i = 0; i < NPAR_LOOPS; i++) { + const auto partition_no = i * WARP_SIZE + threadIdx.x; + shared_exp_sums[partition_no] = rescaled_exp_sum[i]; + } + + #pragma unroll + for (int mask = WARP_SIZE / 2; mask >= 1; mask /= 2) { + global_exp_sum += __shfl_xor(global_exp_sum, mask); + } + if (threadIdx.x == 0) { + shared_global_exp_sum = global_exp_sum; + } + } // warpid == 0 + const scalar_t* tmp_out_ptr = + tmp_out + seq_idx * num_heads * max_num_partitions * HEAD_SIZE + + head_idx * max_num_partitions * HEAD_SIZE + threadIdx.x; + constexpr int MAX_NPAR = 64; + scalar_t tmps[MAX_NPAR]; + const float dzero = 0.0f; + #pragma unroll + for (int j = 0; j < MAX_NPAR; j++) { + tmps[j] = from_float(dzero); + } + const int last_partition_offset = (num_partitions - 1) * HEAD_SIZE; + const int num_partition_offset = (num_partitions)*HEAD_SIZE; + int idx = 0; + + constexpr int JCHUNK = 16; + + #pragma unroll + for (int j = 0; j < JCHUNK * HEAD_SIZE; j += HEAD_SIZE) { + // lastj is last valid partition + const int lastj_offset = + (j < num_partition_offset) ? j : last_partition_offset; + tmps[idx] = tmp_out_ptr[lastj_offset]; + idx++; + } + __syncthreads(); + + if (num_partitions > JCHUNK) { + #pragma unroll + for (int j = JCHUNK * HEAD_SIZE; j < 2 * JCHUNK * HEAD_SIZE; + j += HEAD_SIZE) { + const int lastj_offset = + (j < num_partition_offset) ? j : last_partition_offset; + tmps[idx] = tmp_out_ptr[lastj_offset]; + idx++; + } + + if (num_partitions > 2 * JCHUNK) { + #pragma unroll + for (int j = 2 * JCHUNK * HEAD_SIZE; j < MAX_NPAR * HEAD_SIZE; + j += HEAD_SIZE) { + const int lastj_offset = + (j < num_partition_offset) ? j : last_partition_offset; + tmps[idx] = tmp_out_ptr[lastj_offset]; + idx++; + } + } + } // num_partitions > JCHUNK + + // Aggregate tmp_out to out. + float acc = 0.0f; + #pragma unroll + for (int j = 0; j < JCHUNK; j++) { + acc += to_float(tmps[j]) * shared_exp_sums[j]; + } + if (num_partitions > JCHUNK) { + #pragma unroll + for (int j = JCHUNK; j < 2 * JCHUNK; j++) { + acc += to_float(tmps[j]) * shared_exp_sums[j]; + } + if (num_partitions > 2 * JCHUNK) { + #pragma unroll + for (int j = 2 * JCHUNK; j < MAX_NPAR; j++) { + acc += to_float(tmps[j]) * shared_exp_sums[j]; + } + } + } + + for (int p = 1; p < NPAR_LOOPS; p++) { + if (num_partitions > p * MAX_NPAR) { + idx = 0; + #pragma unroll + for (int j = p * MAX_NPAR * HEAD_SIZE; j < (p + 1) * MAX_NPAR * HEAD_SIZE; + j += HEAD_SIZE) { + // lastj is last valid partition + const int lastj_offset = + (j < num_partition_offset) ? j : last_partition_offset; + tmps[idx] = tmp_out_ptr[lastj_offset]; + idx++; + } + + #pragma unroll + for (int j = 0; j < MAX_NPAR; j++) { + acc += to_float(tmps[j]) * shared_exp_sums[j + p * MAX_NPAR]; + } + } + } + + const float inv_global_exp_sum = + __fdividef(1.0f, shared_global_exp_sum + 1e-6f); + const float out_scale = + (fp8_out_scale_ptr != nullptr) ? 1.0f / (*fp8_out_scale_ptr) : 1.0f; + acc *= inv_global_exp_sum; + acc *= out_scale; + const int64_t query_start_off = static_cast( + query_start_loc_ptr ? query_start_loc_ptr[seq_idx] : seq_idx); + OUTT* out_ptr = out + query_start_off * num_heads * HEAD_SIZE + + static_cast(head_idx) * HEAD_SIZE; + if constexpr (std::is_same::value) { + out_ptr[threadIdx.x] = + __hip_cvt_float_to_fp8(acc, vllm::fp8::fp8_type::__default_saturation, + vllm::fp8::fp8_type::__default_interpret); + } else { + out_ptr[threadIdx.x] = from_float(acc); + } +} + +#elif defined(__HIP__GFX11__) + +using floatx8 = __attribute__((__vector_size__(8 * sizeof(float)))) float; + +using bit16_t = uint16_t; +using bit16x4 = __attribute__((__vector_size__(4 * sizeof(uint16_t)))) uint16_t; +typedef bit16x4 _B16x4; + +using bit16x8 = __attribute__((__vector_size__(8 * sizeof(uint16_t)))) uint16_t; +union b16x8_u { + bit16x8 u16x8; + _B16x4 xy[2]; +}; +typedef b16x8_u _B16x8; + +using bit16x16 = + __attribute__((__vector_size__(16 * sizeof(uint16_t)))) uint16_t; +union b16x16_u { + bit16x16 u16x16; + _B16x8 xy[2]; +}; +typedef b16x16_u _B16x16; + +using _B8x8 = uint2; +using bit8_t = uint8_t; + +typedef struct _B8x16 { + _B8x8 xy[2]; +} _B8x16; + +template +__device__ __forceinline__ floatx8 gcn_wmma16x16x16_instr(const bit16x16& inpA, + const bit16x16& inpB, + const floatx8& inpC) { + if constexpr (std::is_same::value) { + return __builtin_amdgcn_wmma_f32_16x16x16_f16_w32(inpA, inpB, inpC); + } else if constexpr (std::is_same::value) { + return __builtin_amdgcn_wmma_f32_16x16x16_bf16_w32(inpA, inpB, inpC); + } else { + static_assert(false, "unsupported 16b dtype"); + } +} + +template +__device__ __forceinline__ float to_float(const T& inp) { + if constexpr (std::is_same::value) { + return (float)inp; + } else if constexpr (std::is_same::value) { + return __bfloat162float(inp); + } else { + static_assert(false, "unsupported 16b dtype"); + } +} + +template +__device__ __forceinline__ T from_float(const float& inp) { + if constexpr (std::is_same::value) { + return (_Float16)inp; + } else if constexpr (std::is_same::value) { + return __float2bfloat16(inp); + } else { + static_assert(false, "unsupported 16b dtype"); + } +} + +template +__device__ __forceinline__ _B16x8 from_floatx8(const floatx8& inp) { + if constexpr (std::is_same::value) { + union h2cvt { + __half2 h2[4]; + _B16x8 b16x8; + } u; + u.h2[0] = __float22half2_rn(make_float2(inp[0], inp[1])); + u.h2[1] = __float22half2_rn(make_float2(inp[2], inp[3])); + u.h2[2] = __float22half2_rn(make_float2(inp[4], inp[5])); + u.h2[3] = __float22half2_rn(make_float2(inp[6], inp[7])); + return u.b16x8; + } else if constexpr (std::is_same::value) { + union b2cvt { + __hip_bfloat162 b2[4]; + _B16x8 b16x8; + } u; + + u.b2[0] = __float22bfloat162_rn(make_float2(inp[0], inp[1])); + u.b2[1] = __float22bfloat162_rn(make_float2(inp[2], inp[3])); + u.b2[2] = __float22bfloat162_rn(make_float2(inp[4], inp[5])); + u.b2[3] = __float22bfloat162_rn(make_float2(inp[6], inp[7])); + + return u.b16x8; + } else { + static_assert(false, "unsupported 16b dtype"); + } +} + +// clang-format off +template +__global__ +__launch_bounds__(NUM_THREADS, 3) void paged_attention_ll4mi_QKV_mfma16_kernel( + const scalar_t* __restrict__ q, // [num_seqs, num_heads, head_size] + const cache_t* __restrict__ k_cache, // [num_blocks, num_kv_heads, + // head_size/x, block_size, x] + const cache_t* __restrict__ v_cache, // [num_blocks, num_kv_heads, + // head_size, block_size] + const int num_kv_heads, const float scale, + const int* __restrict__ block_tables, // [num_seqs, max_num_blocks_per_seq] + const int* __restrict__ seq_lens, // [num_seqs] + const int* __restrict__ query_start_loc_ptr, // [num_seqs] + const int max_num_blocks_per_seq, + const float* __restrict__ alibi_slopes, // [num_heads] + const int q_stride, const int kv_block_stride, const int kv_head_stride, + float* __restrict__ exp_sums, // [num_seqs, num_heads, max_num_partitions] + float* __restrict__ max_logits, // [num_seqs, num_heads, + // max_num_partitions] + scalar_t* __restrict__ out, // [num_seqs, num_heads, max_num_partitions, + // head_size] + OUTT* __restrict__ final_out, // [num_seqs, num_heads, head_size] + int max_ctx_blocks, const float* k_scale, const float* v_scale) { + // clang-format on + constexpr int NWARPS = NUM_THREADS / WARP_SIZE; // 8 warps on gfx11 + const int warpid = threadIdx.x / WARP_SIZE; + const int laneid = threadIdx.x % WARP_SIZE; + const int lane2id = laneid % 2; + const int lane16id = laneid % 16; + const int rowid = laneid / 16; + + const int seq_idx = blockIdx.x; + // NOTE queries with sequence len > 1 are prefills and taken care by another + // kernel. + if (query_start_loc_ptr != nullptr && + (query_start_loc_ptr[seq_idx + 1] - query_start_loc_ptr[seq_idx]) != 1) { + return; + } + + const int partition_idx = blockIdx.y; + + constexpr int T_PAR_SIZE = 256; // token partition size set to 256 + + const int max_num_partitions = gridDim.y; + + const int seq_len = seq_lens[seq_idx]; // length of a seq + + const int partition_start_token_idx = partition_idx * T_PAR_SIZE; + // exit if partition is out of context for seq + if (partition_start_token_idx >= seq_len) { + return; + } + + constexpr int GQA_RATIO2 = DIVIDE_ROUND_UP(GQA_RATIO, 2); + + __shared__ float shared_qk_max[NWARPS][16 + 1]; + __shared__ float shared_exp_sum[NWARPS][16 + 1]; + // shared_logits is used for multiple purposes + __shared__ _B16x16 shared_logits[NWARPS][2][16][2]; + + // for QK wmma16x16, layout is QHead/Tokenx16 across every 16 lanes, + // 32 Bytes HeadElements in each lane, 2x16B HeadElements across a row of warp + constexpr int ROWS_PER_WARP = + WARP_SIZE / 16 / 2; // rows refers to 16 lanes; refer dpp terminology + constexpr int CONTIGUOUS_KV_ELEMS_16B_LOAD = + 16 / sizeof(cache_t); // 8 for 16 bit cache type, 16 for 8 bit types + constexpr int QKHE_PER_FETCH = + CONTIGUOUS_KV_ELEMS_16B_LOAD * + ROWS_PER_WARP; // each fetch across a warp fetches these many elements + constexpr int QKHELOOP = HEAD_SIZE / QKHE_PER_FETCH; // 2xQKHE_16B across + // warp + + _B16x16 Qlocal[QKHELOOP / 2]; // note that 16 contiguous elements of Q should + // be fetched per lane for 16 bit cache types + + constexpr int CONTIGUOUS_SCALAR_ELEMS_16B = 16 / sizeof(scalar_t); + + constexpr int TOKENS_PER_WARP = + T_PAR_SIZE / + NWARPS; // sub partition of tokens per warp for qk calculation + constexpr int TLOOP = + TOKENS_PER_WARP / + 16; // each wmma16x16x16 instruction processes 16 tokens + + _B16x16 Klocal[TLOOP] + [QKHELOOP / 2]; // can be interpreted as B8x16 for 8 bit types + + const int wg_start_head_idx = blockIdx.z * GQA_RATIO; + const int wg_start_kv_head_idx = blockIdx.z; + const int total_num_heads = gridDim.z * GQA_RATIO; + + // for QK wmma, tokens in multiples of TOKENS_PER_WARP are spread across warps + // each wmma takes QH16xT16x16HE across warp + // repeat wmma across QKHELOOP dimension + // output layout from QKwmma : QH16xT8x2 16 qheads across 16 lanes, 16 tokens + // across 2 rows x 8 tokens per lane + + const int64_t query_start_off = static_cast( + query_start_loc_ptr ? query_start_loc_ptr[seq_idx] : seq_idx); + + if (GQA_RATIO == 1) { + const int local_qhead_idx = lane16id % GQA_RATIO; + const int global_qhead_idx = wg_start_head_idx + local_qhead_idx; + const scalar_t* q_ptr = + q + query_start_off * q_stride + global_qhead_idx * HEAD_SIZE; + if (lane16id < GQA_RATIO) { + #pragma unroll + for (int qkhe_depth = 0; qkhe_depth < QKHELOOP / 2; qkhe_depth++) { + const scalar_t* q_fetch_ptr = q_ptr + qkhe_depth * QKHE_PER_FETCH * 2; + const _B16x16* q_fetch_ptr_32B = + reinterpret_cast(q_fetch_ptr); + Qlocal[qkhe_depth] = *q_fetch_ptr_32B; + } + } + } else { + // fetch Q in shared across warps and then write to registers + const int local_qhead_idx = 2 * warpid + rowid; + const int global_qhead_idx = wg_start_head_idx + local_qhead_idx; + const scalar_t* q_ptr = + q + query_start_off * q_stride + global_qhead_idx * HEAD_SIZE; + + const int qhead_element = lane16id * CONTIGUOUS_SCALAR_ELEMS_16B; + if ((local_qhead_idx < GQA_RATIO) && (qhead_element < HEAD_SIZE)) { + const scalar_t* q_fetch_ptr = q_ptr + qhead_element; + const _B16x8* q_fetch_ptr_16B = + reinterpret_cast(q_fetch_ptr); + _B16x8 tmp = *q_fetch_ptr_16B; + + const int offset1 = + lane16id / + 2; // 16 contiguous chunks of head elems are spread across 8x2lanes + shared_logits[offset1][lane2id][local_qhead_idx][0].xy[0] = tmp; + } + + __syncthreads(); + + #pragma unroll + for (int qkhe_depth = 0; qkhe_depth < QKHELOOP / 2; qkhe_depth++) { + Qlocal[qkhe_depth].xy[0] = + shared_logits[qkhe_depth][0][lane16id % GQA_RATIO][0].xy[0]; + Qlocal[qkhe_depth].xy[1] = + shared_logits[qkhe_depth][1][lane16id % GQA_RATIO][0].xy[0]; + } + } + + const int num_seq_blocks = DIVIDE_ROUND_UP(seq_len, BLOCK_SIZE); + const int last_seq_block = num_seq_blocks - 1; + + const int* block_table_seq = block_tables + seq_idx * max_num_blocks_per_seq; + + int kphysical_block_number[TLOOP]; + + // fetch k physical block numbers + for (int token_depth = 0; token_depth < TLOOP; token_depth++) { + const int klocal_token_idx = + TOKENS_PER_WARP * warpid + token_depth * 16 + lane16id; + const int kglobal_token_idx = partition_start_token_idx + klocal_token_idx; + const int kblock_idx = (kglobal_token_idx < seq_len) + ? kglobal_token_idx / BLOCK_SIZE + : last_seq_block; + kphysical_block_number[token_depth] = block_table_seq[kblock_idx]; + } + + constexpr int KX = 16 / sizeof(cache_t); + const cache_t* k_ptr = k_cache + wg_start_kv_head_idx * kv_head_stride; + + const int row_head_elem = 0; + + for (int token_depth = 0; token_depth < TLOOP; token_depth++) { + const int64_t kblock_number = + static_cast(kphysical_block_number[token_depth]); + const cache_t* k_ptr2 = k_ptr + kblock_number * kv_block_stride; + const int klocal_token_idx = + TOKENS_PER_WARP * warpid + token_depth * 16 + lane16id; + const int kphysical_block_offset = klocal_token_idx % BLOCK_SIZE; + const cache_t* k_ptr3 = k_ptr2 + kphysical_block_offset * KX; + + for (int qkhe_depth = 0; qkhe_depth < QKHELOOP; qkhe_depth++) { + const int head_elem = row_head_elem + qkhe_depth * QKHE_PER_FETCH; + const int offset1 = head_elem / KX; + const int offset2 = head_elem % KX; + const cache_t* k_fetch_ptr = k_ptr3 + offset1 * BLOCK_SIZE * KX + offset2; + const _B16x8* k_fetch_ptr_16B = + reinterpret_cast(k_fetch_ptr); + Klocal[token_depth][qkhe_depth / 2].xy[qkhe_depth % 2] = *k_fetch_ptr_16B; + } + } + + constexpr int VTOKENS_PER_LANE = + TOKENS_PER_WARP / ROWS_PER_WARP; // 32/1 = 32 vtokens per lane + constexpr int VBLOCKS_PER_LANE = 2; // assumes block size >=16 + constexpr int VTLOOP = NWARPS; // corresponds to tokens across warps + constexpr int VTLANELOOP = DIVIDE_ROUND_UP( + VTOKENS_PER_LANE, + CONTIGUOUS_KV_ELEMS_16B_LOAD); // optimized for 16B fetches; assumes + // minimum block size is 16 + constexpr int VHELOOP = DIVIDE_ROUND_UP( + (HEAD_SIZE / 16), NWARPS); // head_size distributed across warps; each + // wmma instr works on 16 head elements + + int vphysical_block_number[VTLOOP][VBLOCKS_PER_LANE]; + + // fetch v physical block numbers + for (int vtoken_depth = 0; vtoken_depth < VTLOOP; vtoken_depth++) { + for (int vblock_depth = 0; vblock_depth < VBLOCKS_PER_LANE; + vblock_depth++) { + const int vlocal_token_idx = + vtoken_depth * VTOKENS_PER_LANE * ROWS_PER_WARP + + vblock_depth * BLOCK_SIZE; + const int vglobal_token_idx = + partition_start_token_idx + vlocal_token_idx; + const int vblock_idx = (vglobal_token_idx < seq_len) + ? vglobal_token_idx / BLOCK_SIZE + : last_seq_block; + vphysical_block_number[vtoken_depth][vblock_depth] = + block_table_seq[vblock_idx]; + } + } + + _B16x16 Vlocal[VTLOOP][VHELOOP] + [VTLANELOOP / 2]; // this can be interpreted as B8x16 too + + const cache_t* v_ptr = v_cache + wg_start_kv_head_idx * kv_head_stride; + // v fetches are 16head elems across lanes x (16x2) tokens per lane + for (int vhe_depth = 0; vhe_depth < VHELOOP; vhe_depth++) { + const int vhead_elem = vhe_depth * NWARPS * 16 + warpid * 16 + lane16id; + const cache_t* v_ptr2 = v_ptr + vhead_elem * BLOCK_SIZE; + + for (int vtoken_depth = 0; vtoken_depth < VTLOOP; vtoken_depth++) { + for (int vfetch_depth = 0; vfetch_depth < VTLANELOOP; vfetch_depth++) { + const int64_t vblock_number = static_cast( + vphysical_block_number[vtoken_depth] + [vfetch_depth / VBLOCKS_PER_LANE]); + const cache_t* v_ptr3 = v_ptr2 + (vblock_number * kv_block_stride); + + const cache_t* v_fetch_ptr = + v_ptr3 + + (vfetch_depth % VBLOCKS_PER_LANE) * CONTIGUOUS_KV_ELEMS_16B_LOAD; + const _B16x8* v_fetch_ptr_16B = + reinterpret_cast(v_fetch_ptr); + Vlocal[vtoken_depth][vhe_depth][vfetch_depth / 2].xy[vfetch_depth % 2] = + *v_fetch_ptr_16B; + } + } + } + + floatx8 dout[TLOOP]; + // qk wmma + for (int token_depth = 0; token_depth < TLOOP; token_depth++) { + dout[token_depth] = {0}; + for (int qkhe_depth = 0; qkhe_depth < QKHELOOP / 2; qkhe_depth++) { + dout[token_depth] = gcn_wmma16x16x16_instr( + Klocal[token_depth][qkhe_depth].u16x16, Qlocal[qkhe_depth].u16x16, + dout[token_depth]); + } + dout[token_depth] *= scale; + } + + // calculate qk_max and exp_sum per warp and write to shared memory + float qk_max = -FLT_MAX; + float exp_sum = 0.0f; + const int qkout_token_idx = + partition_start_token_idx + TOKENS_PER_WARP * warpid + rowid; + for (int token_depth = 0; token_depth < TLOOP; token_depth++) { + const int local_token_idx = qkout_token_idx + token_depth * 16; + for (int i = 0; i < 8; i++) { + const float tmp = + (local_token_idx + 2 * i < seq_len) ? dout[token_depth][i] : -FLT_MAX; + qk_max = fmaxf(qk_max, tmp); + } + } + + qk_max = fmaxf(qk_max, __shfl_xor(qk_max, 16)); + + for (int token_depth = 0; token_depth < TLOOP; token_depth++) { + const int local_token_idx = qkout_token_idx + token_depth * 16; + for (int i = 0; i < 8; i++) { + const float tmp = (local_token_idx + 2 * i < seq_len) + ? __expf(dout[token_depth][i] - qk_max) + : 0.0f; + dout[token_depth][i] = tmp; + exp_sum += tmp; + } + } + + exp_sum += __shfl_xor(exp_sum, 16); + + __syncthreads(); + + if (laneid < 16) { + shared_qk_max[warpid][lane16id] = qk_max; + shared_exp_sum[warpid][lane16id] = exp_sum; + } + + __syncthreads(); + + // calculate partition qk_max and exp_sum + float partition_qk_max = -FLT_MAX; + float warp_qk_max_exp[NWARPS]; + float partition_exp_sum = 0.0f; + + #pragma unroll + for (int w = 0; w < NWARPS; w++) { + warp_qk_max_exp[w] = shared_qk_max[w][lane16id]; + partition_qk_max = fmaxf(partition_qk_max, warp_qk_max_exp[w]); + } + + for (int w = 0; w < NWARPS; w++) { + warp_qk_max_exp[w] = __expf(warp_qk_max_exp[w] - partition_qk_max); + partition_exp_sum += shared_exp_sum[w][lane16id] * warp_qk_max_exp[w]; + } + + const float inv_sum_scale = + __fdividef(1.f, partition_exp_sum + 1e-6f) * warp_qk_max_exp[warpid]; + + __syncthreads(); + + // write logits to shared mem + #pragma unroll + for (int token_depth = 0; token_depth < TLOOP; token_depth++) { + dout[token_depth] *= inv_sum_scale; + shared_logits[warpid][token_depth][lane16id][0].xy[rowid] = + from_floatx8(dout[token_depth]); + } + __syncthreads(); + + _B16x8 swp_buf[TLOOP][2]; + #pragma unroll + for (int token_depth = 0; token_depth < TLOOP; token_depth++) { + swp_buf[token_depth][0] = + shared_logits[warpid][token_depth][lane16id][0].xy[0]; + swp_buf[token_depth][1] = + shared_logits[warpid][token_depth][lane16id][0].xy[1]; + } + + #pragma unroll + for (int token_depth = 0; token_depth < TLOOP; token_depth++) { + #pragma unroll + for (int i = 0; i < 8; i++) { + shared_logits[warpid][token_depth][lane16id][0].xy[rowid].u16x8[i] = + swp_buf[token_depth][i % 2].u16x8[4 * rowid + (i / 2)]; + } + } + + // write out partition max_logits and exp_sum + if (threadIdx.x < GQA_RATIO) { + const int qhead_idx = lane16id; + const int offset = seq_idx * total_num_heads * max_num_partitions + + (wg_start_head_idx + qhead_idx) * max_num_partitions + + partition_idx; + max_logits[offset] = partition_qk_max; + exp_sums[offset] = partition_exp_sum; + } + + __syncthreads(); + + _B16x8 outelems[VHELOOP]; + // Softmax V wmma + // v layout: 16he across lanes x (16x2) tokens per lane + for (int vhe_depth = 0; vhe_depth < VHELOOP; vhe_depth++) { + floatx8 tmp_out = {0}; + for (int vtoken_depth = 0; vtoken_depth < VTLOOP; vtoken_depth++) { + for (int vfetch_depth = 0; vfetch_depth < VTLANELOOP / 2; + vfetch_depth++) { + const int offset = vfetch_depth; + // if output format is 16 qheads across 16 lanes, 16 head elems spread + // across rows + tmp_out = gcn_wmma16x16x16_instr( + Vlocal[vtoken_depth][vhe_depth][vfetch_depth].u16x16, + shared_logits[vtoken_depth][offset][lane16id][0].u16x16, tmp_out); + } + } + outelems[vhe_depth] = from_floatx8(tmp_out); + } + + __syncthreads(); + + #pragma unroll + for (int vhe_depth = 0; vhe_depth < VHELOOP; vhe_depth++) { + shared_logits[warpid][vhe_depth][lane16id][0].xy[rowid] = + outelems[vhe_depth]; // lane16 id head dimension; rowid head element + // dimension + } + + __syncthreads(); + + #pragma unroll + for (int vhe_depth = 0; vhe_depth < VHELOOP; vhe_depth++) { + swp_buf[vhe_depth][0] = shared_logits[warpid][vhe_depth][lane16id][0].xy[0]; + swp_buf[vhe_depth][1] = shared_logits[warpid][vhe_depth][lane16id][0].xy[1]; + } + + #pragma unroll + for (int vhe_depth = 0; vhe_depth < VHELOOP; vhe_depth++) { + #pragma unroll + for (int i = 0; i < 8; i++) { + shared_logits[warpid][vhe_depth][lane16id][0].xy[rowid].u16x8[i] = + swp_buf[vhe_depth][i % 2].u16x8[4 * rowid + (i / 2)]; + } + } + + __syncthreads(); + + // write to tmp_out with coalesced writes after reading from shared mem + if (warpid == 0) { + _B16x8 vout[GQA_RATIO2]; + // each lane writes out 16Bytes of tmp_out along head elem dimension + const int head_elem_idx = lane16id * 8; + if (head_elem_idx < HEAD_SIZE) { + for (int h = 0; h < GQA_RATIO2; h++) { + const int local_head_idx = 2 * h + rowid; + const int offset1 = (head_elem_idx / 16) % NWARPS; + const int offset2 = head_elem_idx / 16 / NWARPS; + const int offset3 = (head_elem_idx / 8) % 2; // num_he % num_row + vout[h] = + shared_logits[offset1][offset2][local_head_idx][0].xy[offset3]; + } + + const int hsz_maxp_mult = HEAD_SIZE * max_num_partitions; + scalar_t* out_ptr = out + seq_idx * total_num_heads * hsz_maxp_mult + + partition_idx * HEAD_SIZE; + for (int h = 0; h < GQA_RATIO2; h++) { + const int local_head_idx = 2 * h + rowid; + if (local_head_idx < GQA_RATIO) { + const int out_head_idx = wg_start_head_idx + local_head_idx; + scalar_t* out_ptr2 = out_ptr + out_head_idx * hsz_maxp_mult; + scalar_t* out_ptr3 = out_ptr2 + head_elem_idx; + _B16x8* out_ptr_B16x8 = reinterpret_cast<_B16x8*>(out_ptr3); + *out_ptr_B16x8 = vout[h]; + } + } + } + } +} + +template +__global__ +__launch_bounds__(NUM_THREADS) void paged_attention_ll4mi_QKV_mfma4_kernel( + const scalar_t* __restrict__ q, // [num_seqs, num_heads, head_size] + const cache_t* __restrict__ k_cache, // [num_blocks, num_kv_heads, + // head_size/x, block_size, x] + const cache_t* __restrict__ v_cache, // [num_blocks, num_kv_heads, + // head_size, block_size] + const int num_kv_heads, const float scale, + const int* __restrict__ block_tables, // [num_seqs, max_num_blocks_per_seq] + const int* __restrict__ seq_lens, // [num_seqs] + const int* __restrict__ query_start_loc_ptr, // [num_seqs] + const int max_num_blocks_per_seq, + const float* __restrict__ alibi_slopes, // [num_heads] + const int q_stride, const int kv_block_stride, const int kv_head_stride, + float* __restrict__ exp_sums, // [num_seqs, num_heads, max_num_partitions] + float* __restrict__ max_logits, // [num_seqs, num_heads, + // max_num_partitions] + scalar_t* __restrict__ out, // [num_seqs, num_heads, max_num_partitions, + // head_size] + OUTT* __restrict__ final_out, // [num_seqs, num_heads, head_size] + int max_ctx_blocks, const float* k_scale, const float* v_scale) { + UNREACHABLE_CODE +} + +// Grid: (num_heads, num_seqs). +template +__global__ +__launch_bounds__(NUM_THREADS) void paged_attention_ll4mi_reduce_kernel( + OUTT* __restrict__ out, // [num_seqs, num_heads, head_size] + const float* __restrict__ exp_sums, // [num_seqs, num_heads, + // max_num_partitions] + const float* __restrict__ max_logits, // [num_seqs, num_heads, + // max_num_partitions] + const scalar_t* __restrict__ tmp_out, // [num_seqs, num_heads, + // max_num_partitions, head_size] + const int* __restrict__ seq_lens, // [num_seqs] + const int* __restrict__ query_start_loc_ptr, // [num_seqs] + const int max_num_partitions, const float* __restrict__ fp8_out_scale_ptr) { + const auto num_heads = gridDim.x; + const auto head_idx = blockIdx.x; + const auto seq_idx = blockIdx.y; + + // NOTE queries with sequence len > 1 are prefills and taken care by another + // kernel. + if (query_start_loc_ptr != nullptr && + (query_start_loc_ptr[seq_idx + 1] - query_start_loc_ptr[seq_idx] != 1)) { + return; + } + + const int seq_len = seq_lens[seq_idx]; + const int num_partitions = DIVIDE_ROUND_UP(seq_len, PARTITION_SIZE); + const int warpid = threadIdx.x / WARP_SIZE; + + __shared__ float shared_global_exp_sum; + // max num partitions supported is warp_size * NPAR_LOOPS + __shared__ float shared_exp_sums[NPAR_LOOPS * WARP_SIZE]; + + if (warpid == 0) { + const float* max_logits_ptr = max_logits + + seq_idx * num_heads * max_num_partitions + + head_idx * max_num_partitions; + + // valid partition is the last valid partition in case threadid > num + // partitions + int valid_partition[NPAR_LOOPS]; + float reg_max_logit[NPAR_LOOPS]; + const int last_valid_partition = num_partitions - 1; + + #pragma unroll + for (int i = 0; i < NPAR_LOOPS; i++) { + const int partition_no = i * WARP_SIZE + threadIdx.x; + valid_partition[i] = + (partition_no < num_partitions) ? partition_no : last_valid_partition; + } + #pragma unroll + for (int i = 0; i < NPAR_LOOPS; i++) { + reg_max_logit[i] = max_logits_ptr[valid_partition[i]]; + } + float max_logit = reg_max_logit[0]; + #pragma unroll + for (int i = 1; i < NPAR_LOOPS; i++) { + max_logit = fmaxf(max_logit, reg_max_logit[i]); + } + + #pragma unroll + for (int mask = WARP_SIZE / 2; mask >= 1; mask /= 2) { + max_logit = fmaxf(max_logit, __shfl_xor(max_logit, mask)); + } + + const float* exp_sums_ptr = exp_sums + + seq_idx * num_heads * max_num_partitions + + head_idx * max_num_partitions; + + float rescaled_exp_sum[NPAR_LOOPS]; + #pragma unroll + for (int i = 0; i < NPAR_LOOPS; i++) { + rescaled_exp_sum[i] = exp_sums_ptr[valid_partition[i]]; + } + #pragma unroll + for (int i = 0; i < NPAR_LOOPS; i++) { + const int partition_no = i * WARP_SIZE + threadIdx.x; + rescaled_exp_sum[i] *= (partition_no < num_partitions) + ? expf(reg_max_logit[i] - max_logit) + : 0.0f; + } + float global_exp_sum = rescaled_exp_sum[0]; + #pragma unroll + for (int i = 1; i < NPAR_LOOPS; i++) { + global_exp_sum += rescaled_exp_sum[i]; + } + #pragma unroll + for (int i = 0; i < NPAR_LOOPS; i++) { + const int partition_no = i * WARP_SIZE + threadIdx.x; + shared_exp_sums[partition_no] = rescaled_exp_sum[i]; + } + + #pragma unroll + for (int mask = WARP_SIZE / 2; mask >= 1; mask /= 2) { + global_exp_sum += __shfl_xor(global_exp_sum, mask); + } + if (threadIdx.x == 0) { + shared_global_exp_sum = global_exp_sum; + } + } // warpid == 0 + const scalar_t* tmp_out_ptr = + tmp_out + seq_idx * num_heads * max_num_partitions * HEAD_SIZE + + head_idx * max_num_partitions * HEAD_SIZE + threadIdx.x; + constexpr int MAX_NPAR = 32; + scalar_t tmps[MAX_NPAR]; + const float dzero = 0.0f; + #pragma unroll + for (int j = 0; j < MAX_NPAR; j++) { + tmps[j] = from_float(dzero); + } + const int last_partition_offset = (num_partitions - 1) * HEAD_SIZE; + const int num_partition_offset = (num_partitions)*HEAD_SIZE; + int idx = 0; + + constexpr int JCHUNK = 16; + + #pragma unroll + for (int j = 0; j < JCHUNK * HEAD_SIZE; j += HEAD_SIZE) { + // lastj is last valid partition + const int lastj_offset = + (j < num_partition_offset) ? j : last_partition_offset; + tmps[idx] = tmp_out_ptr[lastj_offset]; + idx++; + } + __syncthreads(); + + if (num_partitions > JCHUNK) { + #pragma unroll + for (int j = JCHUNK * HEAD_SIZE; j < 2 * JCHUNK * HEAD_SIZE; + j += HEAD_SIZE) { + const int lastj_offset = + (j < num_partition_offset) ? j : last_partition_offset; + tmps[idx] = tmp_out_ptr[lastj_offset]; + idx++; + } + + if (num_partitions > 2 * JCHUNK) { + #pragma unroll + for (int j = 2 * JCHUNK * HEAD_SIZE; j < MAX_NPAR * HEAD_SIZE; + j += HEAD_SIZE) { + const int lastj_offset = + (j < num_partition_offset) ? j : last_partition_offset; + tmps[idx] = tmp_out_ptr[lastj_offset]; + idx++; + } + } + } // num_partitions > JCHUNK + + // Aggregate tmp_out to out. + float acc = 0.0f; + #pragma unroll + for (int j = 0; j < JCHUNK; j++) { + acc += to_float(tmps[j]) * shared_exp_sums[j]; + } + if (num_partitions > JCHUNK) { + #pragma unroll + for (int j = JCHUNK; j < 2 * JCHUNK; j++) { + acc += to_float(tmps[j]) * shared_exp_sums[j]; + } + if (num_partitions > 2 * JCHUNK) { + #pragma unroll + for (int j = 2 * JCHUNK; j < MAX_NPAR; j++) { + acc += to_float(tmps[j]) * shared_exp_sums[j]; + } + } + } + + for (int p = 1; p < NPAR_LOOPS; p++) { + if (num_partitions > p * MAX_NPAR) { + idx = 0; + #pragma unroll + for (int j = p * MAX_NPAR * HEAD_SIZE; j < (p + 1) * MAX_NPAR * HEAD_SIZE; + j += HEAD_SIZE) { + // lastj is last valid partition + const int lastj_offset = + (j < num_partition_offset) ? j : last_partition_offset; + tmps[idx] = tmp_out_ptr[lastj_offset]; + idx++; + } + + #pragma unroll + for (int j = 0; j < MAX_NPAR; j++) { + acc += to_float(tmps[j]) * shared_exp_sums[j + p * MAX_NPAR]; + } + } + } + + const float inv_global_exp_sum = + __fdividef(1.0f, shared_global_exp_sum + 1e-6f); + acc *= inv_global_exp_sum; + + const int64_t query_start_off = static_cast( + query_start_loc_ptr ? query_start_loc_ptr[seq_idx] : seq_idx); + OUTT* out_ptr = out + query_start_off * num_heads * HEAD_SIZE + + static_cast(head_idx) * HEAD_SIZE; + out_ptr[threadIdx.x] = from_float(acc); +} + +#elif defined(__HIP__GFX12__) + +using floatx8 = __attribute__((__vector_size__(8 * sizeof(float)))) float; + +using bit16_t = uint16_t; +using bit16x4 = __attribute__((__vector_size__(4 * sizeof(uint16_t)))) uint16_t; +typedef bit16x4 _B16x4; + +using bit16x8 = __attribute__((__vector_size__(8 * sizeof(uint16_t)))) uint16_t; +union b16x8_u { + bit16x8 u16x8; + _B16x4 xy[2]; +}; +typedef b16x8_u _B16x8; + +using _B8x8 = uint2; +using bit8_t = uint8_t; + +typedef struct _B8x16 { + _B8x8 xy[2]; +} _B8x16; + +template +__device__ __forceinline__ floatx8 gcn_wmma16x16x16_instr(const bit16x8& inpA, + const bit16x8& inpB, + const floatx8& inpC) { + if constexpr (std::is_same::value) { + return __builtin_amdgcn_wmma_f32_16x16x16_f16_w32_gfx12(inpA, inpB, inpC); + } else if constexpr (std::is_same::value) { + return __builtin_amdgcn_wmma_f32_16x16x16_bf16_w32_gfx12(inpA, inpB, inpC); + } else { + static_assert(false, "unsupported 16b dtype"); + } +} + +template +__device__ __forceinline__ float to_float(const T& inp) { + if constexpr (std::is_same::value) { + return (float)inp; + } else if constexpr (std::is_same::value) { + return __bfloat162float(inp); + } else { + static_assert(false, "unsupported 16b dtype"); + } +} + +template +__device__ __forceinline__ float to_float_b16(const bit16_t& inp) { + union tmpcvt { + bit16_t u; + _Float16 f; + __hip_bfloat16 b; + } t16; + t16.u = inp; + if constexpr (std::is_same::value) { + return (float)t16.f; + } else if constexpr (std::is_same::value) { + return __bfloat162float(t16.b); + } else { + static_assert(false, "unsupported 16b dtype"); + } +} + +template +__device__ __forceinline__ T from_float(const float& inp) { + if constexpr (std::is_same::value) { + return (_Float16)inp; + } else if constexpr (std::is_same::value) { + return __float2bfloat16(inp); + } else { + static_assert(false, "unsupported 16b dtype"); + } +} + +template +__device__ __forceinline__ _B16x8 from_floatx8(const floatx8& inp) { + if constexpr (std::is_same::value) { + union h2cvt { + __half2 h2[4]; + _B16x8 b16x8; + } u; + u.h2[0] = __float22half2_rn(make_float2(inp[0], inp[1])); + u.h2[1] = __float22half2_rn(make_float2(inp[2], inp[3])); + u.h2[2] = __float22half2_rn(make_float2(inp[4], inp[5])); + u.h2[3] = __float22half2_rn(make_float2(inp[6], inp[7])); + return u.b16x8; + } else if constexpr (std::is_same::value) { + union b2cvt { + __hip_bfloat162 b2[4]; + _B16x8 b16x8; + } u; + + u.b2[0] = __float22bfloat162_rn(make_float2(inp[0], inp[1])); + u.b2[1] = __float22bfloat162_rn(make_float2(inp[2], inp[3])); + u.b2[2] = __float22bfloat162_rn(make_float2(inp[4], inp[5])); + u.b2[3] = __float22bfloat162_rn(make_float2(inp[6], inp[7])); + + return u.b16x8; + } else { + static_assert(false, "unsupported 16b dtype"); + } +} + +// clang-format off +template +__global__ +__launch_bounds__(NUM_THREADS, 3) void paged_attention_ll4mi_QKV_mfma16_kernel( + const scalar_t* __restrict__ q, // [num_seqs, num_heads, head_size] + const cache_t* __restrict__ k_cache, // [num_blocks, num_kv_heads, + // head_size/x, block_size, x] + const cache_t* __restrict__ v_cache, // [num_blocks, num_kv_heads, + // head_size, block_size] + const int num_kv_heads, const float scale, + const int* __restrict__ block_tables, // [num_seqs, max_num_blocks_per_seq] + const int* __restrict__ seq_lens, // [num_seqs] + const int* __restrict__ query_start_loc_ptr, // [num_seqs] + const int max_num_blocks_per_seq, + const float* __restrict__ alibi_slopes, // [num_heads] + const int q_stride, const int kv_block_stride, const int kv_head_stride, + float* __restrict__ exp_sums, // [num_seqs, num_heads, max_num_partitions] + float* __restrict__ max_logits, // [num_seqs, num_heads, + // max_num_partitions] + scalar_t* __restrict__ out, // [num_seqs, num_heads, max_num_partitions, + // head_size] + OUTT* __restrict__ final_out, // [num_seqs, num_heads, head_size] + int max_ctx_blocks, const float* k_scale, const float* v_scale) { + // clang-format on + constexpr int NWARPS = NUM_THREADS / WARP_SIZE; // 8 warps on gfx11 + const int warpid = threadIdx.x / WARP_SIZE; + const int laneid = threadIdx.x % WARP_SIZE; + const int lane2id = laneid % 2; + const int lane16id = laneid % 16; + const int rowid = laneid / 16; + + const int seq_idx = blockIdx.x; + // NOTE queries with sequence len > 1 are prefills and taken care by another + // kernel. + if (query_start_loc_ptr != nullptr && + (query_start_loc_ptr[seq_idx + 1] - query_start_loc_ptr[seq_idx] != 1)) { + return; + } + const int partition_idx = blockIdx.y; + + constexpr int T_PAR_SIZE = 256; // token partition size set to 256 + + const int max_num_partitions = gridDim.y; + + const int seq_len = seq_lens[seq_idx]; // length of a seq + + const int partition_start_token_idx = partition_idx * T_PAR_SIZE; + // exit if partition is out of context for seq + if (partition_start_token_idx >= seq_len) { + return; + } + + constexpr int GQA_RATIO2 = DIVIDE_ROUND_UP(GQA_RATIO, 2); + + __shared__ float shared_qk_max[NWARPS][16 + 1]; + __shared__ float shared_exp_sum[NWARPS][16 + 1]; + // shared_logits is used for multiple purposes + __shared__ _B16x8 shared_logits[NWARPS][2][16][2]; + + // for QK wmma16x16_gfx12, layout is QHead/Tokenx16 across every 16 lanes, + // 16 Bytes HeadElements in each lane, 2x16B HeadElements across 2 rows of + // warp + constexpr int ROWS_PER_WARP = + WARP_SIZE / 16; // rows refers to 16 lanes; refer dpp terminology + constexpr int CONTIGUOUS_KV_ELEMS_16B_LOAD = + 16 / sizeof(cache_t); // 8 for 16 bit cache type, 16 for 8 bit types + constexpr int QKHE_PER_FETCH = + CONTIGUOUS_KV_ELEMS_16B_LOAD * + ROWS_PER_WARP; // each fetch across a warp fetches these many elements + constexpr int QKHELOOP = HEAD_SIZE / QKHE_PER_FETCH; // 2xQKHE_16B across + // warp + + _B16x8 Qlocal[QKHELOOP]; // note that 16 contiguous elements of Q should + // be fetched per lane for 16 bit cache types + + constexpr int CONTIGUOUS_SCALAR_ELEMS_16B = 16 / sizeof(scalar_t); + + constexpr int TOKENS_PER_WARP = + T_PAR_SIZE / + NWARPS; // sub partition of tokens per warp for qk calculation + constexpr int TLOOP = + TOKENS_PER_WARP / + 16; // each wmma16x16x16 instruction processes 16 tokens + + _B16x8 Klocal[TLOOP] + [QKHELOOP]; // can be interpreted as B8x16 for 8 bit types + + const int wg_start_head_idx = blockIdx.z * GQA_RATIO; + const int wg_start_kv_head_idx = blockIdx.z; + const int total_num_heads = gridDim.z * GQA_RATIO; + + // for QK wmma, tokens in multiples of TOKENS_PER_WARP are spread across warps + // each wmma takes QH16xT16x16HE across warp + // repeat wmma across QKHELOOP dimension + // output layout from QKwmma : QH16xT8x2 16 qheads across 16 lanes, 16 tokens + // across 2 rows x 8 tokens per lane + + const int64_t query_start_off = static_cast( + query_start_loc_ptr ? query_start_loc_ptr[seq_idx] : seq_idx); + + if (GQA_RATIO == 1) { + const int local_qhead_idx = lane16id % GQA_RATIO; + const int global_qhead_idx = wg_start_head_idx + local_qhead_idx; + const scalar_t* q_ptr = q + query_start_off * q_stride + + global_qhead_idx * HEAD_SIZE + + rowid * CONTIGUOUS_KV_ELEMS_16B_LOAD; + if (lane16id < GQA_RATIO) { + #pragma unroll + for (int qkhe_depth = 0; qkhe_depth < QKHELOOP; qkhe_depth++) { + const scalar_t* q_fetch_ptr = q_ptr + qkhe_depth * QKHE_PER_FETCH; + const _B16x8* q_fetch_ptr_16B = + reinterpret_cast(q_fetch_ptr); + Qlocal[qkhe_depth] = *q_fetch_ptr_16B; + } + } + } else { + // fetch Q in shared across warps and then write to registers + const int local_qhead_idx = 2 * warpid + rowid; + const int global_qhead_idx = wg_start_head_idx + local_qhead_idx; + const scalar_t* q_ptr = + q + query_start_off * q_stride + global_qhead_idx * HEAD_SIZE; + + const int qhead_element = lane16id * CONTIGUOUS_SCALAR_ELEMS_16B; + if ((local_qhead_idx < GQA_RATIO) && (qhead_element < HEAD_SIZE)) { + const scalar_t* q_fetch_ptr = q_ptr + qhead_element; + const _B16x8* q_fetch_ptr_16B = + reinterpret_cast(q_fetch_ptr); + _B16x8 tmp = *q_fetch_ptr_16B; + + const int offset1 = + lane16id / + 2; // 16 contiguous chunks of head elems are spread across 8x2lanes + shared_logits[offset1][lane2id][local_qhead_idx][0] = tmp; + } + + __syncthreads(); + + #pragma unroll + for (int qkhe_depth = 0; qkhe_depth < QKHELOOP; qkhe_depth++) { + Qlocal[qkhe_depth] = + shared_logits[qkhe_depth][rowid][lane16id % GQA_RATIO][0]; + } + } + + const int num_seq_blocks = DIVIDE_ROUND_UP(seq_len, BLOCK_SIZE); + const int last_seq_block = num_seq_blocks - 1; + + const int* block_table_seq = block_tables + seq_idx * max_num_blocks_per_seq; + + int kphysical_block_number[TLOOP]; + + // fetch k physical block numbers + for (int token_depth = 0; token_depth < TLOOP; token_depth++) { + const int klocal_token_idx = + TOKENS_PER_WARP * warpid + token_depth * 16 + lane16id; + const int kglobal_token_idx = partition_start_token_idx + klocal_token_idx; + const int kblock_idx = (kglobal_token_idx < seq_len) + ? kglobal_token_idx / BLOCK_SIZE + : last_seq_block; + kphysical_block_number[token_depth] = block_table_seq[kblock_idx]; + } + + constexpr int KX = 16 / sizeof(cache_t); + const cache_t* k_ptr = k_cache + wg_start_kv_head_idx * kv_head_stride; + + const int row_head_elem = rowid * CONTIGUOUS_KV_ELEMS_16B_LOAD; + + for (int token_depth = 0; token_depth < TLOOP; token_depth++) { + const int64_t kblock_number = + static_cast(kphysical_block_number[token_depth]); + const cache_t* k_ptr2 = k_ptr + kblock_number * kv_block_stride; + const int klocal_token_idx = + TOKENS_PER_WARP * warpid + token_depth * 16 + lane16id; + const int kphysical_block_offset = klocal_token_idx % BLOCK_SIZE; + const cache_t* k_ptr3 = k_ptr2 + kphysical_block_offset * KX; + + for (int qkhe_depth = 0; qkhe_depth < QKHELOOP; qkhe_depth++) { + const int head_elem = row_head_elem + qkhe_depth * QKHE_PER_FETCH; + const int offset1 = head_elem / KX; + const int offset2 = head_elem % KX; + const cache_t* k_fetch_ptr = k_ptr3 + offset1 * BLOCK_SIZE * KX + offset2; + const _B16x8* k_fetch_ptr_16B = + reinterpret_cast(k_fetch_ptr); + Klocal[token_depth][qkhe_depth] = *k_fetch_ptr_16B; + } + } + + constexpr int VTOKENS_PER_LANE = + TOKENS_PER_WARP / ROWS_PER_WARP; // 32/2 = 16 vtokens per lane + constexpr int VBLOCKS_PER_LANE = 1; // assumes block size >=16 + constexpr int VTLOOP = NWARPS; // corresponds to tokens across warps + constexpr int VTLANELOOP = DIVIDE_ROUND_UP( + VTOKENS_PER_LANE, + CONTIGUOUS_KV_ELEMS_16B_LOAD); // optimized for 16B fetches; assumes + // minimum block size is 16 + constexpr int VHELOOP = DIVIDE_ROUND_UP( + (HEAD_SIZE / 16), NWARPS); // head_size distributed across warps; each + // wmma instr works on 16 head elements + + int vphysical_block_number[VTLOOP][VBLOCKS_PER_LANE]; + + // fetch v physical block numbers + for (int vtoken_depth = 0; vtoken_depth < VTLOOP; vtoken_depth++) { + for (int vblock_depth = 0; vblock_depth < VBLOCKS_PER_LANE; + vblock_depth++) { + const int vlocal_token_idx = + vtoken_depth * VTOKENS_PER_LANE * ROWS_PER_WARP + + rowid * VTOKENS_PER_LANE + vblock_depth * BLOCK_SIZE; + const int vglobal_token_idx = + partition_start_token_idx + vlocal_token_idx; + const int vblock_idx = (vglobal_token_idx < seq_len) + ? vglobal_token_idx / BLOCK_SIZE + : last_seq_block; + vphysical_block_number[vtoken_depth][vblock_depth] = + block_table_seq[vblock_idx]; + } + } + + _B16x8 Vlocal[VTLOOP][VHELOOP] + [VTLANELOOP]; // this can be interpreted as B8x16 too + + const cache_t* v_ptr = v_cache + wg_start_kv_head_idx * kv_head_stride + + ((rowid * VTOKENS_PER_LANE) % BLOCK_SIZE); + + // v fetches are 16head elems across lanes x 16 tokens per lane + for (int vhe_depth = 0; vhe_depth < VHELOOP; vhe_depth++) { + const int vhead_elem = vhe_depth * NWARPS * 16 + warpid * 16 + lane16id; + const cache_t* v_ptr2 = v_ptr + vhead_elem * BLOCK_SIZE; + + for (int vtoken_depth = 0; vtoken_depth < VTLOOP; vtoken_depth++) { + for (int vfetch_depth = 0; vfetch_depth < VTLANELOOP; vfetch_depth++) { + const int vblock_depth = 0; + const int64_t vblock_number = static_cast( + vphysical_block_number[vtoken_depth][vblock_depth]); + const cache_t* v_ptr3 = v_ptr2 + (vblock_number * kv_block_stride); + + const cache_t* v_fetch_ptr = + v_ptr3 + vfetch_depth * CONTIGUOUS_KV_ELEMS_16B_LOAD; + const _B16x8* v_fetch_ptr_16B = + reinterpret_cast(v_fetch_ptr); + Vlocal[vtoken_depth][vhe_depth][vfetch_depth] = *v_fetch_ptr_16B; + } + } + } + + floatx8 dout[TLOOP]; + // qk wmma + for (int token_depth = 0; token_depth < TLOOP; token_depth++) { + dout[token_depth] = {0}; + for (int qkhe_depth = 0; qkhe_depth < QKHELOOP; qkhe_depth++) { + dout[token_depth] = gcn_wmma16x16x16_instr( + Klocal[token_depth][qkhe_depth].u16x8, Qlocal[qkhe_depth].u16x8, + dout[token_depth]); + } + dout[token_depth] *= scale; + } + + // calculate qk_max and exp_sum per warp and write to shared memory + float qk_max = -FLT_MAX; + float exp_sum = 0.0f; + const int qkout_token_idx = + partition_start_token_idx + TOKENS_PER_WARP * warpid + rowid * 8; + for (int token_depth = 0; token_depth < TLOOP; token_depth++) { + const int local_token_idx = qkout_token_idx + token_depth * 16; + for (int i = 0; i < 8; i++) { + const float tmp = + (local_token_idx + i < seq_len) ? dout[token_depth][i] : -FLT_MAX; + qk_max = fmaxf(qk_max, tmp); + } + } + + qk_max = fmaxf(qk_max, __shfl_xor(qk_max, 16)); + + for (int token_depth = 0; token_depth < TLOOP; token_depth++) { + const int local_token_idx = qkout_token_idx + token_depth * 16; + for (int i = 0; i < 8; i++) { + const float tmp = (local_token_idx + i < seq_len) + ? __expf(dout[token_depth][i] - qk_max) + : 0.0f; + dout[token_depth][i] = tmp; + exp_sum += tmp; + } + } + + exp_sum += __shfl_xor(exp_sum, 16); + + __syncthreads(); + + if (laneid < 16) { + shared_qk_max[warpid][lane16id] = qk_max; + shared_exp_sum[warpid][lane16id] = exp_sum; + } + + __syncthreads(); + + // calculate partition qk_max and exp_sum + float partition_qk_max = -FLT_MAX; + float warp_qk_max_exp[NWARPS]; + float partition_exp_sum = 0.0f; + + #pragma unroll + for (int w = 0; w < NWARPS; w++) { + warp_qk_max_exp[w] = shared_qk_max[w][lane16id]; + partition_qk_max = fmaxf(partition_qk_max, warp_qk_max_exp[w]); + } + + for (int w = 0; w < NWARPS; w++) { + warp_qk_max_exp[w] = __expf(warp_qk_max_exp[w] - partition_qk_max); + partition_exp_sum += shared_exp_sum[w][lane16id] * warp_qk_max_exp[w]; + } + + const float inv_sum_scale = + __fdividef(1.f, partition_exp_sum + 1e-6f) * warp_qk_max_exp[warpid]; + + __syncthreads(); + + // write logits to shared mem + #pragma unroll + for (int token_depth = 0; token_depth < TLOOP; token_depth++) { + dout[token_depth] *= inv_sum_scale; + shared_logits[warpid][token_depth][lane16id][rowid] = + from_floatx8(dout[token_depth]); + } + + // write out partition max_logits and exp_sum + if (threadIdx.x < GQA_RATIO) { + const int qhead_idx = lane16id; + const int offset = seq_idx * total_num_heads * max_num_partitions + + (wg_start_head_idx + qhead_idx) * max_num_partitions + + partition_idx; + max_logits[offset] = partition_qk_max; + exp_sums[offset] = partition_exp_sum; + } + + __syncthreads(); + + _B16x8 outelems[VHELOOP]; + // Softmax V wmma + // v layout: 16he across lanes x 16 tokens per lane + for (int vhe_depth = 0; vhe_depth < VHELOOP; vhe_depth++) { + floatx8 tmp_out = {0}; + + for (int vtoken_depth = 0; vtoken_depth < VTLOOP; vtoken_depth++) { + for (int vfetch_depth = 0; vfetch_depth < VTLANELOOP; vfetch_depth++) { + const int offset = rowid * VTLANELOOP + vfetch_depth; + const int offset1 = offset % ROWS_PER_WARP; + const int offset2 = offset / ROWS_PER_WARP; + // if output format is 16 qheads across 16 lanes, 16 head elems spread + // across rows + tmp_out = gcn_wmma16x16x16_instr( + Vlocal[vtoken_depth][vhe_depth][vfetch_depth].u16x8, + shared_logits[vtoken_depth][offset2][lane16id][offset1].u16x8, + tmp_out); + } + } + outelems[vhe_depth] = from_floatx8(tmp_out); + } + + __syncthreads(); + + #pragma unroll + for (int vhe_depth = 0; vhe_depth < VHELOOP; vhe_depth++) { + shared_logits[warpid][vhe_depth][lane16id][rowid] = + outelems[vhe_depth]; // lane16 id head dimension; rowid head element + // dimension + } + + __syncthreads(); + + // write to tmp_out with coalesced writes after reading from shared mem + if (warpid == 0) { + _B16x8 vout[GQA_RATIO2]; + // each lane writes out 16Bytes of tmp_out along head elem dimension + const int head_elem_idx = lane16id * 8; + if (head_elem_idx < HEAD_SIZE) { + for (int h = 0; h < GQA_RATIO2; h++) { + const int local_head_idx = 2 * h + rowid; + const int offset1 = (head_elem_idx / 16) % NWARPS; + const int offset2 = head_elem_idx / 16 / NWARPS; + const int offset3 = (head_elem_idx / 8) % 2; // num_he % num_row + vout[h] = shared_logits[offset1][offset2][local_head_idx][offset3]; + } + + const int hsz_maxp_mult = HEAD_SIZE * max_num_partitions; + scalar_t* out_ptr = out + seq_idx * total_num_heads * hsz_maxp_mult + + partition_idx * HEAD_SIZE; + for (int h = 0; h < GQA_RATIO2; h++) { + const int local_head_idx = 2 * h + rowid; + if (local_head_idx < GQA_RATIO) { + const int out_head_idx = wg_start_head_idx + local_head_idx; + scalar_t* out_ptr2 = out_ptr + out_head_idx * hsz_maxp_mult; + scalar_t* out_ptr3 = out_ptr2 + head_elem_idx; + _B16x8* out_ptr_B16x8 = reinterpret_cast<_B16x8*>(out_ptr3); + *out_ptr_B16x8 = vout[h]; + } + } + } + } +} + +template +__global__ +__launch_bounds__(NUM_THREADS) void paged_attention_ll4mi_QKV_mfma4_kernel( + const scalar_t* __restrict__ q, // [num_seqs, num_heads, head_size] + const cache_t* __restrict__ k_cache, // [num_blocks, num_kv_heads, + // head_size/x, block_size, x] + const cache_t* __restrict__ v_cache, // [num_blocks, num_kv_heads, + // head_size, block_size] + const int num_kv_heads, const float scale, + const int* __restrict__ block_tables, // [num_seqs, max_num_blocks_per_seq] + const int* __restrict__ seq_lens, // [num_seqs] + const int* __restrict__ query_start_loc_ptr, // [num_seqs] + const int max_num_blocks_per_seq, + const float* __restrict__ alibi_slopes, // [num_heads] + const int q_stride, const int kv_block_stride, const int kv_head_stride, + float* __restrict__ exp_sums, // [num_seqs, num_heads, max_num_partitions] + float* __restrict__ max_logits, // [num_seqs, num_heads, + // max_num_partitions] + scalar_t* __restrict__ out, // [num_seqs, num_heads, max_num_partitions, + // head_size] + OUTT* __restrict__ final_out, // [num_seqs, num_heads, head_size] + int max_ctx_blocks, const float* k_scale, const float* v_scale) { + UNREACHABLE_CODE +} + +// Grid: (num_heads, num_seqs). +template +__global__ +__launch_bounds__(NUM_THREADS) void paged_attention_ll4mi_reduce_kernel( + OUTT* __restrict__ out, // [num_seqs, num_heads, head_size] + const float* __restrict__ exp_sums, // [num_seqs, num_heads, + // max_num_partitions] + const float* __restrict__ max_logits, // [num_seqs, num_heads, + // max_num_partitions] + const scalar_t* __restrict__ tmp_out, // [num_seqs, num_heads, + // max_num_partitions, head_size] + const int* __restrict__ seq_lens, // [num_seqs] + const int* __restrict__ query_start_loc_ptr, // [num_seqs] + const int max_num_partitions, const float* __restrict__ fp8_out_scale_ptr) { + const auto num_heads = gridDim.x; + const auto head_idx = blockIdx.x; + const auto seq_idx = blockIdx.y; + + // NOTE queries with sequence len > 1 are prefills and taken care by another + // kernel. + if (query_start_loc_ptr != nullptr && + (query_start_loc_ptr[seq_idx + 1] - query_start_loc_ptr[seq_idx] != 1)) { + return; + } + + const int seq_len = seq_lens[seq_idx]; + const int num_partitions = DIVIDE_ROUND_UP(seq_len, PARTITION_SIZE); + const int warpid = threadIdx.x / WARP_SIZE; + + __shared__ float shared_global_exp_sum; + // max num partitions supported is warp_size * NPAR_LOOPS + __shared__ float shared_exp_sums[NPAR_LOOPS * WARP_SIZE]; + + if (warpid == 0) { + const float* max_logits_ptr = max_logits + + seq_idx * num_heads * max_num_partitions + + head_idx * max_num_partitions; + + // valid partition is the last valid partition in case threadid > num + // partitions + int valid_partition[NPAR_LOOPS]; + float reg_max_logit[NPAR_LOOPS]; + const int last_valid_partition = num_partitions - 1; + + #pragma unroll + for (int i = 0; i < NPAR_LOOPS; i++) { + const int partition_no = i * WARP_SIZE + threadIdx.x; + valid_partition[i] = + (partition_no < num_partitions) ? partition_no : last_valid_partition; + } + #pragma unroll + for (int i = 0; i < NPAR_LOOPS; i++) { + reg_max_logit[i] = max_logits_ptr[valid_partition[i]]; + } + float max_logit = reg_max_logit[0]; + #pragma unroll + for (int i = 1; i < NPAR_LOOPS; i++) { + max_logit = fmaxf(max_logit, reg_max_logit[i]); + } + + #pragma unroll + for (int mask = WARP_SIZE / 2; mask >= 1; mask /= 2) { + max_logit = fmaxf(max_logit, __shfl_xor(max_logit, mask)); + } + + const float* exp_sums_ptr = exp_sums + + seq_idx * num_heads * max_num_partitions + + head_idx * max_num_partitions; + + float rescaled_exp_sum[NPAR_LOOPS]; + #pragma unroll + for (int i = 0; i < NPAR_LOOPS; i++) { + rescaled_exp_sum[i] = exp_sums_ptr[valid_partition[i]]; + } + #pragma unroll + for (int i = 0; i < NPAR_LOOPS; i++) { + const int partition_no = i * WARP_SIZE + threadIdx.x; + rescaled_exp_sum[i] *= (partition_no < num_partitions) + ? expf(reg_max_logit[i] - max_logit) + : 0.0f; + } + float global_exp_sum = rescaled_exp_sum[0]; + #pragma unroll + for (int i = 1; i < NPAR_LOOPS; i++) { + global_exp_sum += rescaled_exp_sum[i]; + } + #pragma unroll + for (int i = 0; i < NPAR_LOOPS; i++) { + const int partition_no = i * WARP_SIZE + threadIdx.x; + shared_exp_sums[partition_no] = rescaled_exp_sum[i]; + } + + #pragma unroll + for (int mask = WARP_SIZE / 2; mask >= 1; mask /= 2) { + global_exp_sum += __shfl_xor(global_exp_sum, mask); + } + if (threadIdx.x == 0) { + shared_global_exp_sum = global_exp_sum; + } + } // warpid == 0 + const scalar_t* tmp_out_ptr = + tmp_out + seq_idx * num_heads * max_num_partitions * HEAD_SIZE + + head_idx * max_num_partitions * HEAD_SIZE + threadIdx.x; + constexpr int MAX_NPAR = 32; + scalar_t tmps[MAX_NPAR]; + const float dzero = 0.0f; + #pragma unroll + for (int j = 0; j < MAX_NPAR; j++) { + tmps[j] = from_float(dzero); + } + const int last_partition_offset = (num_partitions - 1) * HEAD_SIZE; + const int num_partition_offset = (num_partitions)*HEAD_SIZE; + int idx = 0; + + constexpr int JCHUNK = 16; + + #pragma unroll + for (int j = 0; j < JCHUNK * HEAD_SIZE; j += HEAD_SIZE) { + // lastj is last valid partition + const int lastj_offset = + (j < num_partition_offset) ? j : last_partition_offset; + tmps[idx] = tmp_out_ptr[lastj_offset]; + idx++; + } + __syncthreads(); + + if (num_partitions > JCHUNK) { + #pragma unroll + for (int j = JCHUNK * HEAD_SIZE; j < 2 * JCHUNK * HEAD_SIZE; + j += HEAD_SIZE) { + const int lastj_offset = + (j < num_partition_offset) ? j : last_partition_offset; + tmps[idx] = tmp_out_ptr[lastj_offset]; + idx++; + } + + if (num_partitions > 2 * JCHUNK) { + #pragma unroll + for (int j = 2 * JCHUNK * HEAD_SIZE; j < MAX_NPAR * HEAD_SIZE; + j += HEAD_SIZE) { + const int lastj_offset = + (j < num_partition_offset) ? j : last_partition_offset; + tmps[idx] = tmp_out_ptr[lastj_offset]; + idx++; + } + } + } // num_partitions > JCHUNK + + // Aggregate tmp_out to out. + float acc = 0.0f; + #pragma unroll + for (int j = 0; j < JCHUNK; j++) { + acc += to_float(tmps[j]) * shared_exp_sums[j]; + } + if (num_partitions > JCHUNK) { + #pragma unroll + for (int j = JCHUNK; j < 2 * JCHUNK; j++) { + acc += to_float(tmps[j]) * shared_exp_sums[j]; + } + if (num_partitions > 2 * JCHUNK) { + #pragma unroll + for (int j = 2 * JCHUNK; j < MAX_NPAR; j++) { + acc += to_float(tmps[j]) * shared_exp_sums[j]; + } + } + } + + for (int p = 1; p < NPAR_LOOPS; p++) { + if (num_partitions > p * MAX_NPAR) { + idx = 0; + #pragma unroll + for (int j = p * MAX_NPAR * HEAD_SIZE; j < (p + 1) * MAX_NPAR * HEAD_SIZE; + j += HEAD_SIZE) { + // lastj is last valid partition + const int lastj_offset = + (j < num_partition_offset) ? j : last_partition_offset; + tmps[idx] = tmp_out_ptr[lastj_offset]; + idx++; + } + + #pragma unroll + for (int j = 0; j < MAX_NPAR; j++) { + acc += to_float(tmps[j]) * shared_exp_sums[j + p * MAX_NPAR]; + } + } + } + + const float inv_global_exp_sum = + __fdividef(1.0f, shared_global_exp_sum + 1e-6f); + acc *= inv_global_exp_sum; + + const int64_t query_start_off = static_cast( + query_start_loc_ptr ? query_start_loc_ptr[seq_idx] : seq_idx); + OUTT* out_ptr = out + query_start_off * num_heads * HEAD_SIZE + + static_cast(head_idx) * HEAD_SIZE; + out_ptr[threadIdx.x] = from_float(acc); +} + +#else + +// clang-format off +template +__global__ +__launch_bounds__(NUM_THREADS) void paged_attention_ll4mi_QKV_mfma16_kernel( + const scalar_t* __restrict__ q, // [num_seqs, num_heads, head_size] + const cache_t* __restrict__ k_cache, // [num_blocks, num_kv_heads, head_size/x, block_size, x] + const cache_t* __restrict__ v_cache, // [num_blocks, num_kv_heads, head_size, block_size] + const int num_kv_heads, + const float scale, + const int* __restrict__ block_tables, // [num_seqs, max_num_blocks_per_seq] + const int* __restrict__ seq_lens, // [num_seqs] + const int* __restrict__ query_start_loc_ptr, // [num_seqs] + const int max_num_blocks_per_seq, + const float* __restrict__ alibi_slopes, // [num_heads] + const int q_stride, + const int kv_block_stride, + const int kv_head_stride, + float* __restrict__ exp_sums, // [num_seqs, num_heads, max_num_partitions] + float* __restrict__ max_logits, // [num_seqs, num_heads, max_num_partitions] + scalar_t* __restrict__ out, // [num_seqs, num_heads, max_num_partitions, head_size] + OUTT* __restrict__ final_out, // [num_seqs, num_heads, head_size] + int max_ctx_blocks, const float* k_scale, const float* v_scale) { + UNREACHABLE_CODE +} + +template +__global__ +__launch_bounds__(NUM_THREADS) void paged_attention_ll4mi_QKV_mfma4_kernel( + const scalar_t* __restrict__ q, // [num_seqs, num_heads, head_size] + const cache_t* __restrict__ k_cache, // [num_blocks, num_kv_heads, head_size/x, block_size, x] + const cache_t* __restrict__ v_cache, // [num_blocks, num_kv_heads, head_size, block_size] + const int num_kv_heads, + const float scale, + const int* __restrict__ block_tables, // [num_seqs, max_num_blocks_per_seq] + const int* __restrict__ seq_lens, // [num_seqs] + const int* __restrict__ query_start_loc_ptr, // [num_seqs] + const int max_num_blocks_per_seq, + const float* __restrict__ alibi_slopes, // [num_heads] + const int q_stride, + const int kv_block_stride, + const int kv_head_stride, + float* __restrict__ exp_sums, // [num_seqs, num_heads, max_num_partitions] + float* __restrict__ max_logits, // [num_seqs, num_heads, max_num_partitions] + scalar_t* __restrict__ out, // [num_seqs, num_heads, max_num_partitions, head_size] + OUTT* __restrict__ final_out, // [num_seqs, num_heads, head_size] + int max_ctx_blocks, const float* k_scale, const float* v_scale) { + UNREACHABLE_CODE +} + +// Grid: (num_heads, num_seqs). +template +__global__ +__launch_bounds__(NUM_THREADS) void paged_attention_ll4mi_reduce_kernel( + OUTT* __restrict__ out, // [num_seqs, num_heads, head_size] + const float* __restrict__ exp_sums, // [num_seqs, num_heads, max_num_partitions] + const float* __restrict__ max_logits, // [num_seqs, num_heads, max_num_partitions] + const scalar_t* __restrict__ tmp_out, // [num_seqs, num_heads, max_num_partitions, head_size] + const int* __restrict__ seq_lens, // [num_seqs] + const int* __restrict__ query_start_loc_ptr, // [num_seqs] + const int max_num_partitions, const float* __restrict__ fp8_out_scale_ptr) { + UNREACHABLE_CODE +} +// clang-format on + +#endif + +#define LAUNCH_CUSTOM_ATTENTION_MFMA16(GQA_RATIO) \ + paged_attention_ll4mi_QKV_mfma16_kernel \ + <<>>( \ + query_ptr, key_cache_ptr, value_cache_ptr, num_kv_heads, scale, \ + block_tables_ptr, seq_lens_ptr, query_start_loc_ptr, \ + max_num_blocks_per_seq, alibi_slopes_ptr, q_stride, kv_block_stride, \ + kv_head_stride, exp_sums_ptr, max_logits_ptr, tmp_out_ptr, out_ptr, \ + max_ctx_blocks, k_scale_ptr, v_scale_ptr); + +#define LAUNCH_CUSTOM_ATTENTION_MFMA4(GQA_RATIO) \ + paged_attention_ll4mi_QKV_mfma4_kernel \ + <<>>( \ + query_ptr, key_cache_ptr, value_cache_ptr, num_kv_heads, scale, \ + block_tables_ptr, seq_lens_ptr, query_start_loc_ptr, \ + max_num_blocks_per_seq, alibi_slopes_ptr, q_stride, kv_block_stride, \ + kv_head_stride, exp_sums_ptr, max_logits_ptr, tmp_out_ptr, out_ptr, \ + max_ctx_blocks, k_scale_ptr, v_scale_ptr); + +#define LAUNCH_CUSTOM_REDUCTION(NPAR_LOOPS) \ + paged_attention_ll4mi_reduce_kernel \ + <<>>( \ + out_ptr, exp_sums_ptr, max_logits_ptr, tmp_out_ptr, seq_lens_ptr, \ + query_start_loc_ptr, max_num_partitions, fp8_out_scale_ptr); + +template +void paged_attention_custom_launcher( + torch::Tensor& out, torch::Tensor& exp_sums, torch::Tensor& max_logits, + torch::Tensor& tmp_out, torch::Tensor& query, torch::Tensor& key_cache, + torch::Tensor& value_cache, const int num_kv_heads, float scale, + torch::Tensor& block_tables, torch::Tensor& seq_lens, + const std::optional& query_start_loc, int max_seq_len, + const std::optional& alibi_slopes, torch::Tensor& k_scale, + torch::Tensor& v_scale, const std::optional& fp8_out_scale) { + int num_seqs = block_tables.size(0); + int num_heads = query.size(1); + int head_size = query.size(2); + int max_num_blocks_per_seq = block_tables.size(1); + int q_stride = query.stride(0); + int kv_block_stride = key_cache.stride(0); + int kv_head_stride = key_cache.stride(1); + + // NOTE: query start location is optional for V0 decode should not be used. + // If batch contains mix of prefills and decode, prefills should be skipped. + const int* query_start_loc_ptr = + query_start_loc + ? reinterpret_cast(query_start_loc.value().data_ptr()) + : nullptr; + + // NOTE: alibi_slopes is optional. + const float* alibi_slopes_ptr = + alibi_slopes + ? reinterpret_cast(alibi_slopes.value().data_ptr()) + : nullptr; + + float* exp_sums_ptr = reinterpret_cast(exp_sums.data_ptr()); + float* max_logits_ptr = reinterpret_cast(max_logits.data_ptr()); + T* tmp_out_ptr = reinterpret_cast(tmp_out.data_ptr()); + T* query_ptr = reinterpret_cast(query.data_ptr()); + KVT* key_cache_ptr = reinterpret_cast(key_cache.data_ptr()); + KVT* value_cache_ptr = reinterpret_cast(value_cache.data_ptr()); + int* block_tables_ptr = block_tables.data_ptr(); + int* seq_lens_ptr = seq_lens.data_ptr(); + const float* k_scale_ptr = reinterpret_cast(k_scale.data_ptr()); + const float* v_scale_ptr = reinterpret_cast(v_scale.data_ptr()); + // NOTE: fp8_out_scale is optional. + const auto fp8_out_scale_ptr = + fp8_out_scale + ? static_cast(fp8_out_scale.value().data_ptr()) + : nullptr; + OUTT* out_ptr = reinterpret_cast(out.data_ptr()); + + const int max_ctx_blocks = DIVIDE_ROUND_UP(max_seq_len, BLOCK_SIZE); + + // partition size is fixed at 256 since both mfma4 and mfma16 kernels support + // it mfma4 kernel also supports partition size 512 + constexpr int PARTITION_SIZE = 256; + const int max_num_partitions = DIVIDE_ROUND_UP(max_seq_len, PARTITION_SIZE); + const int gqa_ratio = num_heads / num_kv_heads; + assert(num_heads % num_kv_heads == 0); + assert(head_size == HEAD_SIZE); + + constexpr int NTHR = 256; + dim3 grid(num_seqs, max_num_partitions, num_kv_heads); + dim3 block(NTHR); + const at::cuda::OptionalCUDAGuard device_guard(device_of(query)); + const cudaStream_t stream = at::cuda::getCurrentCUDAStream(); + + // mfma4 kernel is faster than mfma16 for gqa_ratio <= 4 + switch (gqa_ratio) { + case 1: + LAUNCH_CUSTOM_ATTENTION_MFMA4(1); + break; + case 2: + LAUNCH_CUSTOM_ATTENTION_MFMA4(2); + break; + case 3: + LAUNCH_CUSTOM_ATTENTION_MFMA4(3); + break; + case 4: + LAUNCH_CUSTOM_ATTENTION_MFMA4(4); + break; + case 5: + LAUNCH_CUSTOM_ATTENTION_MFMA16(5); + break; + case 6: + LAUNCH_CUSTOM_ATTENTION_MFMA16(6); + break; + case 7: + LAUNCH_CUSTOM_ATTENTION_MFMA16(7); + break; + case 8: + LAUNCH_CUSTOM_ATTENTION_MFMA16(8); + break; + case 9: + LAUNCH_CUSTOM_ATTENTION_MFMA16(9); + break; + case 10: + LAUNCH_CUSTOM_ATTENTION_MFMA16(10); + break; + case 11: + LAUNCH_CUSTOM_ATTENTION_MFMA16(11); + break; + case 12: + LAUNCH_CUSTOM_ATTENTION_MFMA16(12); + break; + case 13: + LAUNCH_CUSTOM_ATTENTION_MFMA16(13); + break; + case 14: + LAUNCH_CUSTOM_ATTENTION_MFMA16(14); + break; + case 15: + LAUNCH_CUSTOM_ATTENTION_MFMA16(15); + break; + case 16: + LAUNCH_CUSTOM_ATTENTION_MFMA16(16); + break; + default: + TORCH_CHECK(false, "Unsupported gqa ratio: ", gqa_ratio); + break; + } + + dim3 reduce_grid(num_heads, num_seqs); + dim3 reduce_block(head_size); + const int npar_loops = DIVIDE_ROUND_UP(max_num_partitions, WARP_SIZE); + // reduction kernel supports upto 8 NPAR_loops * 64 (warp_size) * 256 + // (partition size) = 128K context length + switch (npar_loops) { + case 1: + LAUNCH_CUSTOM_REDUCTION(1); + break; + case 2: + LAUNCH_CUSTOM_REDUCTION(2); + break; + case 3: + LAUNCH_CUSTOM_REDUCTION(3); + break; + case 4: + LAUNCH_CUSTOM_REDUCTION(4); + break; + case 5: + LAUNCH_CUSTOM_REDUCTION(5); + break; + case 6: + LAUNCH_CUSTOM_REDUCTION(6); + break; + case 7: + LAUNCH_CUSTOM_REDUCTION(7); + break; + case 8: + LAUNCH_CUSTOM_REDUCTION(8); + break; + default: + TORCH_CHECK(false, "Unsupported npar_loops: ", npar_loops); + break; + } +} + +template +void paged_attention_custom_launcher_navi( + torch::Tensor& out, torch::Tensor& exp_sums, torch::Tensor& max_logits, + torch::Tensor& tmp_out, torch::Tensor& query, torch::Tensor& key_cache, + torch::Tensor& value_cache, const int num_kv_heads, float scale, + torch::Tensor& block_tables, torch::Tensor& seq_lens, + const std::optional& query_start_loc, int max_seq_len, + const std::optional& alibi_slopes, torch::Tensor& k_scale, + torch::Tensor& v_scale) { + int num_seqs = block_tables.size(0); + int num_heads = query.size(1); + int head_size = query.size(2); + int max_num_blocks_per_seq = block_tables.size(1); + int q_stride = query.stride(0); + int kv_block_stride = key_cache.stride(0); + int kv_head_stride = key_cache.stride(1); + + // NOTE: query start location is optional for V0 decode should not be used. + // If batch contains mix of prefills and decode, prefills should be skipped. + const int* query_start_loc_ptr = + query_start_loc + ? reinterpret_cast(query_start_loc.value().data_ptr()) + : nullptr; + + // NOTE: Navi does not support alibi_slopes. + const float* alibi_slopes_ptr = nullptr; + + float* exp_sums_ptr = reinterpret_cast(exp_sums.data_ptr()); + float* max_logits_ptr = reinterpret_cast(max_logits.data_ptr()); + T* tmp_out_ptr = reinterpret_cast(tmp_out.data_ptr()); + T* query_ptr = reinterpret_cast(query.data_ptr()); + KVT* key_cache_ptr = reinterpret_cast(key_cache.data_ptr()); + KVT* value_cache_ptr = reinterpret_cast(value_cache.data_ptr()); + int* block_tables_ptr = block_tables.data_ptr(); + int* seq_lens_ptr = seq_lens.data_ptr(); + + const float* k_scale_ptr = reinterpret_cast(k_scale.data_ptr()); + const float* v_scale_ptr = reinterpret_cast(v_scale.data_ptr()); + // NOTE: Navi does not support fp8. + const auto fp8_out_scale_ptr = nullptr; + OUTT* out_ptr = reinterpret_cast(out.data_ptr()); + + const int max_ctx_blocks = DIVIDE_ROUND_UP(max_seq_len, BLOCK_SIZE); + + constexpr int PARTITION_SIZE = 256; + const int max_num_partitions = DIVIDE_ROUND_UP(max_seq_len, PARTITION_SIZE); + const int gqa_ratio = num_heads / num_kv_heads; + assert(num_heads % num_kv_heads == 0); + assert(head_size == HEAD_SIZE); + + constexpr int NTHR = 256; + dim3 grid(num_seqs, max_num_partitions, num_kv_heads); + dim3 block(NTHR); + const at::cuda::OptionalCUDAGuard device_guard(device_of(query)); + const cudaStream_t stream = at::cuda::getCurrentCUDAStream(); + + switch (gqa_ratio) { + case 1: + LAUNCH_CUSTOM_ATTENTION_MFMA16(1); + break; + case 2: + LAUNCH_CUSTOM_ATTENTION_MFMA16(2); + break; + case 3: + LAUNCH_CUSTOM_ATTENTION_MFMA16(3); + break; + case 4: + LAUNCH_CUSTOM_ATTENTION_MFMA16(4); + break; + case 5: + LAUNCH_CUSTOM_ATTENTION_MFMA16(5); + break; + case 6: + LAUNCH_CUSTOM_ATTENTION_MFMA16(6); + break; + case 7: + LAUNCH_CUSTOM_ATTENTION_MFMA16(7); + break; + case 8: + LAUNCH_CUSTOM_ATTENTION_MFMA16(8); + break; + case 9: + LAUNCH_CUSTOM_ATTENTION_MFMA16(9); + break; + case 10: + LAUNCH_CUSTOM_ATTENTION_MFMA16(10); + break; + case 11: + LAUNCH_CUSTOM_ATTENTION_MFMA16(11); + break; + case 12: + LAUNCH_CUSTOM_ATTENTION_MFMA16(12); + break; + case 13: + LAUNCH_CUSTOM_ATTENTION_MFMA16(13); + break; + case 14: + LAUNCH_CUSTOM_ATTENTION_MFMA16(14); + break; + case 15: + LAUNCH_CUSTOM_ATTENTION_MFMA16(15); + break; + case 16: + LAUNCH_CUSTOM_ATTENTION_MFMA16(16); + break; + default: + TORCH_CHECK(false, "Unsupported gqa ratio: ", gqa_ratio); + break; + } + + dim3 reduce_grid(num_heads, num_seqs); + dim3 reduce_block(head_size); + const int warp_size = 32; + const int npar_loops = DIVIDE_ROUND_UP(max_num_partitions, warp_size); + // reduction kernel supports upto 16 NPAR_loops * 32 (warp_size) * 256 + // (partition size) = 128K context length + switch (npar_loops) { + case 1: + LAUNCH_CUSTOM_REDUCTION(1); + break; + case 2: + LAUNCH_CUSTOM_REDUCTION(2); + break; + case 3: + LAUNCH_CUSTOM_REDUCTION(3); + break; + case 4: + LAUNCH_CUSTOM_REDUCTION(4); + break; + case 5: + LAUNCH_CUSTOM_REDUCTION(5); + break; + case 6: + LAUNCH_CUSTOM_REDUCTION(6); + break; + case 7: + LAUNCH_CUSTOM_REDUCTION(7); + break; + case 8: + LAUNCH_CUSTOM_REDUCTION(8); + break; + case 9: + LAUNCH_CUSTOM_REDUCTION(9); + break; + case 10: + LAUNCH_CUSTOM_REDUCTION(10); + break; + case 11: + LAUNCH_CUSTOM_REDUCTION(11); + break; + case 12: + LAUNCH_CUSTOM_REDUCTION(12); + break; + case 13: + LAUNCH_CUSTOM_REDUCTION(13); + break; + case 14: + LAUNCH_CUSTOM_REDUCTION(14); + break; + case 15: + LAUNCH_CUSTOM_REDUCTION(15); + break; + case 16: + LAUNCH_CUSTOM_REDUCTION(16); + break; + default: + TORCH_CHECK(false, "Unsupported npar_loops: ", npar_loops); + break; + } +} + +#define CALL_CUSTOM_LAUNCHER(T, KVT, KV_DTYPE, BLK_SIZE, HEAD_SIZE, OUTT, \ + PSIZE, ALIBI_ENABLED, MFMA_TYPE) \ + if (!is_navi) { \ + paged_attention_custom_launcher( \ + out, exp_sums, max_logits, tmp_out, query, key_cache, value_cache, \ + num_kv_heads, scale, block_tables, seq_lens, query_start_loc, \ + max_seq_len, alibi_slopes, k_scale, v_scale, fp8_out_scale); \ + } else { \ + paged_attention_custom_launcher_navi( \ + out, exp_sums, max_logits, tmp_out, query, key_cache, value_cache, \ + num_kv_heads, scale, block_tables, seq_lens, query_start_loc, \ + max_seq_len, alibi_slopes, k_scale, v_scale); \ + } + +#define CALL_CUSTOM_LAUNCHER_ALIBI(T, KVT, KV_DTYPE, BLK_SIZE, HEAD_SIZE, \ + OUTT, PSIZE, MFMA_TYPE) \ + if (alibi_slopes) { \ + CALL_CUSTOM_LAUNCHER(T, KVT, KV_DTYPE, BLK_SIZE, HEAD_SIZE, OUTT, PSIZE, \ + true, MFMA_TYPE); \ + } else { \ + CALL_CUSTOM_LAUNCHER(T, KVT, KV_DTYPE, BLK_SIZE, HEAD_SIZE, OUTT, PSIZE, \ + false, MFMA_TYPE); \ + } + +#if defined(__HIPCC__) && defined(__gfx90a__) + #define CALL_CUSTOM_LAUNCHER_OUT(T, KVT, KV_DTYPE, BLK_SIZE, HEAD_SIZE, \ + MFMA_TYPE) \ + if (fp8_out_scale) { \ + TORCH_CHECK(false, "fp8 out scale unsupported for gfx90a"); \ + } else { \ + CALL_CUSTOM_LAUNCHER_ALIBI(T, KVT, KV_DTYPE, BLK_SIZE, HEAD_SIZE, T, \ + 256, MFMA_TYPE); \ + } +#else + #define CALL_CUSTOM_LAUNCHER_OUT(T, KVT, KV_DTYPE, BLK_SIZE, HEAD_SIZE, \ + MFMA_TYPE) \ + if (fp8_out_scale) { \ + CALL_CUSTOM_LAUNCHER_ALIBI(T, KVT, KV_DTYPE, BLK_SIZE, HEAD_SIZE, \ + uint8_t, 256, MFMA_TYPE); \ + } else { \ + CALL_CUSTOM_LAUNCHER_ALIBI(T, KVT, KV_DTYPE, BLK_SIZE, HEAD_SIZE, T, \ + 256, MFMA_TYPE); \ + } +#endif + +#define CALL_CUSTOM_LAUNCHER_BLK(T, KVT, KV_DTYPE, HEAD_SIZE, MFMA_TYPE) \ + switch (block_size) { \ + case 16: \ + CALL_CUSTOM_LAUNCHER_OUT(T, KVT, KV_DTYPE, 16, HEAD_SIZE, MFMA_TYPE); \ + break; \ + case 32: \ + CALL_CUSTOM_LAUNCHER_OUT(T, KVT, KV_DTYPE, 32, HEAD_SIZE, MFMA_TYPE); \ + break; \ + default: \ + TORCH_CHECK(false, "Unsupported block size: ", block_size); \ + break; \ + } + +#define CALL_CUSTOM_LAUNCHER_BLK_HEAD(T, KVT, KV_DTYPE, MFMA_TYPE) \ + switch (head_size) { \ + case 64: \ + CALL_CUSTOM_LAUNCHER_BLK(T, KVT, KV_DTYPE, 64, MFMA_TYPE); \ + break; \ + case 128: \ + CALL_CUSTOM_LAUNCHER_BLK(T, KVT, KV_DTYPE, 128, MFMA_TYPE); \ + break; \ + default: \ + TORCH_CHECK(false, "Unsupported head size: ", head_size); \ + break; \ + } + +bool is_navi_gpu() { + static bool is_cached = false; + static bool result; + + if (!is_cached) { + int device_id; + hipDeviceProp_t deviceProp; + hipGetDevice(&device_id); + hipGetDeviceProperties(&deviceProp, device_id); + + std::string arch = deviceProp.gcnArchName; + result = arch.find("gfx11") == 0 || arch.find("gfx12") == 0; + is_cached = true; + } + + return result; +} + +// clang-format off +void paged_attention( + torch::Tensor& out, // [num_seqs, num_heads, head_size] + torch::Tensor& exp_sums, // [num_seqs, num_heads, max_num_partitions] + torch::Tensor& max_logits, // [num_seqs, num_heads, max_num_partitions] + torch::Tensor& tmp_out, // [num_seqs, num_heads, max_num_partitions, head_size] + torch::Tensor& query, // [num_seqs, num_heads, head_size] + torch::Tensor& key_cache, // [num_blocks, num_heads, head_size/x, block_size, x] + torch::Tensor& value_cache, // [num_blocks, num_heads, head_size, block_size] + int64_t num_kv_heads, + double scale, + torch::Tensor& block_tables, // [num_seqs, max_num_blocks_per_seq] + torch::Tensor& seq_lens, // [num_seqs] + const std::optional& query_start_loc, // [num_seqs] + int64_t block_size, int64_t max_seq_len, + const std::optional& alibi_slopes, + const std::string& kv_cache_dtype, torch::Tensor& k_scale, + torch::Tensor& v_scale, + const std::optional& fp8_out_scale, + const std::string& mfma_type) { + // clang-format on + bool is_navi = is_navi_gpu(); + const int head_size = query.size(2); + if (kv_cache_dtype == "auto") { + if (query.dtype() == at::ScalarType::Half) { + CALL_CUSTOM_LAUNCHER_BLK_HEAD( + _Float16, _Float16, vllm::Fp8KVCacheDataType::kAuto, MFMAType::F16); + } else if (query.dtype() == at::ScalarType::BFloat16) { + CALL_CUSTOM_LAUNCHER_BLK_HEAD(__hip_bfloat16, __hip_bfloat16, + vllm::Fp8KVCacheDataType::kAuto, + MFMAType::F16); + } else { + TORCH_CHECK(false, "Unsupported data type: ", query.dtype()); + } + } else if (kv_cache_dtype == "fp8" || kv_cache_dtype == "fp8_e4m3") { + if (query.dtype() == at::ScalarType::Half) { + if (mfma_type == "fp8") { + CALL_CUSTOM_LAUNCHER_BLK_HEAD(_Float16, uint8_t, + vllm::Fp8KVCacheDataType::kFp8E4M3, + MFMAType::Fp8); + } else { + CALL_CUSTOM_LAUNCHER_BLK_HEAD(_Float16, uint8_t, + vllm::Fp8KVCacheDataType::kFp8E4M3, + MFMAType::F16); + } + } else if (query.dtype() == at::ScalarType::BFloat16) { + if (mfma_type == "fp8") { + CALL_CUSTOM_LAUNCHER_BLK_HEAD(__hip_bfloat16, uint8_t, + vllm::Fp8KVCacheDataType::kFp8E4M3, + MFMAType::Fp8); + } else { + CALL_CUSTOM_LAUNCHER_BLK_HEAD(__hip_bfloat16, uint8_t, + vllm::Fp8KVCacheDataType::kFp8E4M3, + MFMAType::F16); + } + } else { + TORCH_CHECK(false, "Unsupported data type: ", query.dtype()); + } + } else { + TORCH_CHECK(false, "Unsupported KV cache dtype: ", kv_cache_dtype); + } +} + +#undef WARP_SIZE +#undef MAX +#undef MIN +#undef DIVIDE_ROUND_UP diff --git a/tasks/hip2hip/campaign20/paged_attention_large/test_cases.json b/tasks/hip2hip/campaign20/paged_attention_large/test_cases.json new file mode 100644 index 00000000..d97a727e --- /dev/null +++ b/tasks/hip2hip/campaign20/paged_attention_large/test_cases.json @@ -0,0 +1,9230 @@ +[ + { + "test_case_id": "sig_7b6340559547", + "count": 376, + "args_sig": [ + { + "kind": "tensor", + "dtype": "bfloat16", + "shape": [ + 1024, + 16, + 128 + ], + "stride": [ + 2048, + 128, + 1 + ], + "device": "cuda:0" + }, + { + "kind": "tensor", + "dtype": "float32", + "shape": [ + 1, + 16, + 4 + ], + "stride": [ + 64, + 4, + 1 + ], + "device": "cuda:0" + }, + { + "kind": "tensor", + "dtype": "float32", + "shape": [ + 1, + 16, + 4 + ], + "stride": [ + 64, + 4, + 1 + ], + "device": "cuda:0" + }, + { + "kind": "tensor", + "dtype": "bfloat16", + "shape": [ + 1, + 16, + 4, + 128 + ], + "stride": [ + 8192, + 512, + 128, + 1 + ], + "device": "cuda:0" + }, + { + "kind": "tensor", + "dtype": "bfloat16", + "shape": [ + 1024, + 16, + 128 + ], + "stride": [ + 2048, + 128, + 1 + ], + "device": "cuda:0" + }, + { + "kind": "tensor", + "dtype": "bfloat16", + "shape": [ + 57718, + 1, + 16, + 16, + 8 + ], + "stride": [ + 2048, + 2048, + 128, + 8, + 1 + ], + "device": "cuda:0" + }, + { + "kind": "tensor", + "dtype": "bfloat16", + "shape": [ + 57718, + 1, + 128, + 16 + ], + "stride": [ + 2048, + 2048, + 16, + 1 + ], + "device": "cuda:0" + }, + { + "kind": "scalar", + "value": 1, + "py_type": "int" + }, + { + "kind": "scalar", + "value": 0.08838834764831845, + "py_type": "float" + }, + { + "kind": "tensor", + "dtype": "int32", + "shape": [ + 1, + 128 + ], + "stride": [ + 128, + 1 + ], + "device": "cuda:0" + }, + { + "kind": "tensor", + "dtype": "int32", + "shape": [ + 1 + ], + "stride": [ + 1 + ], + "device": "cuda:0" + }, + { + "kind": "tensor", + "dtype": "int32", + "shape": [ + 2 + ], + "stride": [ + 1 + ], + "device": "cuda:0" + }, + { + "kind": "scalar", + "value": 16, + "py_type": "int" + }, + { + "kind": "scalar", + "value": 1024, + "py_type": "int" + }, + { + "kind": "scalar", + "value": null, + "py_type": "NoneType" + }, + { + "kind": "opaque", + "py_type": "str", + "repr": "'auto'" + }, + { + "kind": "tensor", + "dtype": "float32", + "shape": [], + "stride": [], + "device": "cuda:0" + }, + { + "kind": "tensor", + "dtype": "float32", + "shape": [], + "stride": [], + "device": "cuda:0" + }, + { + "kind": "scalar", + "value": null, + "py_type": "NoneType" + }, + { + "kind": "opaque", + "py_type": "str", + "repr": "'f16'" + } + ], + "kwargs_sig": {}, + "params_repr": {}, + "args_names": [ + "out", + "exp_sums", + "max_logits", + "tmp_out", + "query", + "key_cache", + "value_cache", + "num_kv_heads", + "scale", + "block_tables", + "seq_lens", + "query_start_loc", + "block_size", + "max_seq_len", + "alibi_slopes", + "kv_cache_dtype", + "k_scale", + "v_scale", + "fp8_out_scale", + "mfma_type" + ], + "sources": [ + { + "model": "Qwen_Qwen3-235B-A22B-FP8", + "engine": "vllm", + "mode": "default", + "combo": "Qwen_Qwen3-235B-A22B-FP8__vllm_default", + "count": 376 + } + ] + }, + { + "test_case_id": "sig_b167bd1c728f", + "count": 376, + "args_sig": [ + { + "kind": "tensor", + "dtype": "bfloat16", + "shape": [ + 1, + 16, + 128 + ], + "stride": [ + 2048, + 128, + 1 + ], + "device": "cuda:0" + }, + { + "kind": "tensor", + "dtype": "float32", + "shape": [ + 1, + 16, + 5 + ], + "stride": [ + 80, + 5, + 1 + ], + "device": "cuda:0" + }, + { + "kind": "tensor", + "dtype": "float32", + "shape": [ + 1, + 16, + 5 + ], + "stride": [ + 80, + 5, + 1 + ], + "device": "cuda:0" + }, + { + "kind": "tensor", + "dtype": "bfloat16", + "shape": [ + 1, + 16, + 5, + 128 + ], + "stride": [ + 10240, + 640, + 128, + 1 + ], + "device": "cuda:0" + }, + { + "kind": "tensor", + "dtype": "bfloat16", + "shape": [ + 1, + 16, + 128 + ], + "stride": [ + 2048, + 128, + 1 + ], + "device": "cuda:0" + }, + { + "kind": "tensor", + "dtype": "bfloat16", + "shape": [ + 57718, + 1, + 16, + 16, + 8 + ], + "stride": [ + 2048, + 2048, + 128, + 8, + 1 + ], + "device": "cuda:0" + }, + { + "kind": "tensor", + "dtype": "bfloat16", + "shape": [ + 57718, + 1, + 128, + 16 + ], + "stride": [ + 2048, + 2048, + 16, + 1 + ], + "device": "cuda:0" + }, + { + "kind": "scalar", + "value": 1, + "py_type": "int" + }, + { + "kind": "scalar", + "value": 0.08838834764831845, + "py_type": "float" + }, + { + "kind": "tensor", + "dtype": "int32", + "shape": [ + 1, + 128 + ], + "stride": [ + 128, + 1 + ], + "device": "cuda:0" + }, + { + "kind": "tensor", + "dtype": "int32", + "shape": [ + 1 + ], + "stride": [ + 1 + ], + "device": "cuda:0" + }, + { + "kind": "tensor", + "dtype": "int32", + "shape": [ + 2 + ], + "stride": [ + 1 + ], + "device": "cuda:0" + }, + { + "kind": "scalar", + "value": 16, + "py_type": "int" + }, + { + "kind": "scalar", + "value": 1025, + "py_type": "int" + }, + { + "kind": "scalar", + "value": null, + "py_type": "NoneType" + }, + { + "kind": "opaque", + "py_type": "str", + "repr": "'auto'" + }, + { + "kind": "tensor", + "dtype": "float32", + "shape": [], + "stride": [], + "device": "cuda:0" + }, + { + "kind": "tensor", + "dtype": "float32", + "shape": [], + "stride": [], + "device": "cuda:0" + }, + { + "kind": "scalar", + "value": null, + "py_type": "NoneType" + }, + { + "kind": "opaque", + "py_type": "str", + "repr": "'f16'" + } + ], + "kwargs_sig": {}, + "params_repr": {}, + "args_names": [ + "out", + "exp_sums", + "max_logits", + "tmp_out", + "query", + "key_cache", + "value_cache", + "num_kv_heads", + "scale", + "block_tables", + "seq_lens", + "query_start_loc", + "block_size", + "max_seq_len", + "alibi_slopes", + "kv_cache_dtype", + "k_scale", + "v_scale", + "fp8_out_scale", + "mfma_type" + ], + "sources": [ + { + "model": "Qwen_Qwen3-235B-A22B-FP8", + "engine": "vllm", + "mode": "default", + "combo": "Qwen_Qwen3-235B-A22B-FP8__vllm_default", + "count": 376 + } + ] + }, + { + "test_case_id": "sig_e211e5aa7aeb", + "count": 376, + "args_sig": [ + { + "kind": "tensor", + "dtype": "bfloat16", + "shape": [ + 1520, + 16, + 128 + ], + "stride": [ + 2048, + 128, + 1 + ], + "device": "cuda:0" + }, + { + "kind": "tensor", + "dtype": "float32", + "shape": [ + 32, + 16, + 4 + ], + "stride": [ + 64, + 4, + 1 + ], + "device": "cuda:0" + }, + { + "kind": "tensor", + "dtype": "float32", + "shape": [ + 32, + 16, + 4 + ], + "stride": [ + 64, + 4, + 1 + ], + "device": "cuda:0" + }, + { + "kind": "tensor", + "dtype": "bfloat16", + "shape": [ + 32, + 16, + 4, + 128 + ], + "stride": [ + 8192, + 512, + 128, + 1 + ], + "device": "cuda:0" + }, + { + "kind": "tensor", + "dtype": "bfloat16", + "shape": [ + 1520, + 16, + 128 + ], + "stride": [ + 2048, + 128, + 1 + ], + "device": "cuda:0" + }, + { + "kind": "tensor", + "dtype": "bfloat16", + "shape": [ + 46209, + 1, + 16, + 16, + 8 + ], + "stride": [ + 2048, + 2048, + 128, + 8, + 1 + ], + "device": "cuda:0" + }, + { + "kind": "tensor", + "dtype": "bfloat16", + "shape": [ + 46209, + 1, + 128, + 16 + ], + "stride": [ + 2048, + 2048, + 16, + 1 + ], + "device": "cuda:0" + }, + { + "kind": "scalar", + "value": 1, + "py_type": "int" + }, + { + "kind": "scalar", + "value": 0.08838834764831845, + "py_type": "float" + }, + { + "kind": "tensor", + "dtype": "int32", + "shape": [ + 32, + 128 + ], + "stride": [ + 128, + 1 + ], + "device": "cuda:0" + }, + { + "kind": "tensor", + "dtype": "int32", + "shape": [ + 32 + ], + "stride": [ + 1 + ], + "device": "cuda:0" + }, + { + "kind": "tensor", + "dtype": "int32", + "shape": [ + 33 + ], + "stride": [ + 1 + ], + "device": "cuda:0" + }, + { + "kind": "scalar", + "value": 16, + "py_type": "int" + }, + { + "kind": "scalar", + "value": 1024, + "py_type": "int" + }, + { + "kind": "scalar", + "value": null, + "py_type": "NoneType" + }, + { + "kind": "opaque", + "py_type": "str", + "repr": "'auto'" + }, + { + "kind": "tensor", + "dtype": "float32", + "shape": [], + "stride": [], + "device": "cuda:0" + }, + { + "kind": "tensor", + "dtype": "float32", + "shape": [], + "stride": [], + "device": "cuda:0" + }, + { + "kind": "scalar", + "value": null, + "py_type": "NoneType" + }, + { + "kind": "opaque", + "py_type": "str", + "repr": "'f16'" + } + ], + "kwargs_sig": {}, + "params_repr": {}, + "args_names": [ + "out", + "exp_sums", + "max_logits", + "tmp_out", + "query", + "key_cache", + "value_cache", + "num_kv_heads", + "scale", + "block_tables", + "seq_lens", + "query_start_loc", + "block_size", + "max_seq_len", + "alibi_slopes", + "kv_cache_dtype", + "k_scale", + "v_scale", + "fp8_out_scale", + "mfma_type" + ], + "sources": [ + { + "model": "Qwen_Qwen3-235B-A22B-FP8", + "engine": "vllm", + "mode": "hip", + "combo": "Qwen_Qwen3-235B-A22B-FP8__vllm_hip", + "count": 376 + } + ] + }, + { + "test_case_id": "sig_1227d7656405", + "count": 376, + "args_sig": [ + { + "kind": "tensor", + "dtype": "bfloat16", + "shape": [ + 32, + 16, + 128 + ], + "stride": [ + 2048, + 128, + 1 + ], + "device": "cuda:0" + }, + { + "kind": "tensor", + "dtype": "float32", + "shape": [ + 32, + 16, + 5 + ], + "stride": [ + 80, + 5, + 1 + ], + "device": "cuda:0" + }, + { + "kind": "tensor", + "dtype": "float32", + "shape": [ + 32, + 16, + 5 + ], + "stride": [ + 80, + 5, + 1 + ], + "device": "cuda:0" + }, + { + "kind": "tensor", + "dtype": "bfloat16", + "shape": [ + 32, + 16, + 5, + 128 + ], + "stride": [ + 10240, + 640, + 128, + 1 + ], + "device": "cuda:0" + }, + { + "kind": "tensor", + "dtype": "bfloat16", + "shape": [ + 32, + 16, + 128 + ], + "stride": [ + 2048, + 128, + 1 + ], + "device": "cuda:0" + }, + { + "kind": "tensor", + "dtype": "bfloat16", + "shape": [ + 46209, + 1, + 16, + 16, + 8 + ], + "stride": [ + 2048, + 2048, + 128, + 8, + 1 + ], + "device": "cuda:0" + }, + { + "kind": "tensor", + "dtype": "bfloat16", + "shape": [ + 46209, + 1, + 128, + 16 + ], + "stride": [ + 2048, + 2048, + 16, + 1 + ], + "device": "cuda:0" + }, + { + "kind": "scalar", + "value": 1, + "py_type": "int" + }, + { + "kind": "scalar", + "value": 0.08838834764831845, + "py_type": "float" + }, + { + "kind": "tensor", + "dtype": "int32", + "shape": [ + 32, + 128 + ], + "stride": [ + 128, + 1 + ], + "device": "cuda:0" + }, + { + "kind": "tensor", + "dtype": "int32", + "shape": [ + 32 + ], + "stride": [ + 1 + ], + "device": "cuda:0" + }, + { + "kind": "tensor", + "dtype": "int32", + "shape": [ + 33 + ], + "stride": [ + 1 + ], + "device": "cuda:0" + }, + { + "kind": "scalar", + "value": 16, + "py_type": "int" + }, + { + "kind": "scalar", + "value": 1025, + "py_type": "int" + }, + { + "kind": "scalar", + "value": null, + "py_type": "NoneType" + }, + { + "kind": "opaque", + "py_type": "str", + "repr": "'auto'" + }, + { + "kind": "tensor", + "dtype": "float32", + "shape": [], + "stride": [], + "device": "cuda:0" + }, + { + "kind": "tensor", + "dtype": "float32", + "shape": [], + "stride": [], + "device": "cuda:0" + }, + { + "kind": "scalar", + "value": null, + "py_type": "NoneType" + }, + { + "kind": "opaque", + "py_type": "str", + "repr": "'f16'" + } + ], + "kwargs_sig": {}, + "params_repr": {}, + "args_names": [ + "out", + "exp_sums", + "max_logits", + "tmp_out", + "query", + "key_cache", + "value_cache", + "num_kv_heads", + "scale", + "block_tables", + "seq_lens", + "query_start_loc", + "block_size", + "max_seq_len", + "alibi_slopes", + "kv_cache_dtype", + "k_scale", + "v_scale", + "fp8_out_scale", + "mfma_type" + ], + "sources": [ + { + "model": "Qwen_Qwen3-235B-A22B-FP8", + "engine": "vllm", + "mode": "hip", + "combo": "Qwen_Qwen3-235B-A22B-FP8__vllm_hip", + "count": 376 + } + ] + }, + { + "test_case_id": "sig_35e859d9b24a", + "count": 48, + "args_sig": [ + { + "kind": "tensor", + "dtype": "bfloat16", + "shape": [ + 1024, + 32, + 128 + ], + "stride": [ + 4096, + 128, + 1 + ], + "device": "cuda:0" + }, + { + "kind": "tensor", + "dtype": "float32", + "shape": [ + 1, + 32, + 4 + ], + "stride": [ + 128, + 4, + 1 + ], + "device": "cuda:0" + }, + { + "kind": "tensor", + "dtype": "float32", + "shape": [ + 1, + 32, + 4 + ], + "stride": [ + 128, + 4, + 1 + ], + "device": "cuda:0" + }, + { + "kind": "tensor", + "dtype": "bfloat16", + "shape": [ + 1, + 32, + 4, + 128 + ], + "stride": [ + 16384, + 512, + 128, + 1 + ], + "device": "cuda:0" + }, + { + "kind": "tensor", + "dtype": "bfloat16", + "shape": [ + 1024, + 32, + 128 + ], + "stride": [ + 4096, + 128, + 1 + ], + "device": "cuda:0" + }, + { + "kind": "tensor", + "dtype": "bfloat16", + "shape": [ + 47679, + 4, + 16, + 16, + 8 + ], + "stride": [ + 8192, + 2048, + 128, + 8, + 1 + ], + "device": "cuda:0" + }, + { + "kind": "tensor", + "dtype": "bfloat16", + "shape": [ + 47679, + 4, + 128, + 16 + ], + "stride": [ + 8192, + 2048, + 16, + 1 + ], + "device": "cuda:0" + }, + { + "kind": "scalar", + "value": 4, + "py_type": "int" + }, + { + "kind": "scalar", + "value": 0.08838834764831845, + "py_type": "float" + }, + { + "kind": "tensor", + "dtype": "int32", + "shape": [ + 1, + 128 + ], + "stride": [ + 128, + 1 + ], + "device": "cuda:0" + }, + { + "kind": "tensor", + "dtype": "int32", + "shape": [ + 1 + ], + "stride": [ + 1 + ], + "device": "cuda:0" + }, + { + "kind": "tensor", + "dtype": "int32", + "shape": [ + 2 + ], + "stride": [ + 1 + ], + "device": "cuda:0" + }, + { + "kind": "scalar", + "value": 16, + "py_type": "int" + }, + { + "kind": "scalar", + "value": 1024, + "py_type": "int" + }, + { + "kind": "scalar", + "value": null, + "py_type": "NoneType" + }, + { + "kind": "opaque", + "py_type": "str", + "repr": "'auto'" + }, + { + "kind": "tensor", + "dtype": "float32", + "shape": [], + "stride": [], + "device": "cuda:0" + }, + { + "kind": "tensor", + "dtype": "float32", + "shape": [], + "stride": [], + "device": "cuda:0" + }, + { + "kind": "scalar", + "value": null, + "py_type": "NoneType" + }, + { + "kind": "opaque", + "py_type": "str", + "repr": "'f16'" + } + ], + "kwargs_sig": {}, + "params_repr": {}, + "args_names": [ + "out", + "exp_sums", + "max_logits", + "tmp_out", + "query", + "key_cache", + "value_cache", + "num_kv_heads", + "scale", + "block_tables", + "seq_lens", + "query_start_loc", + "block_size", + "max_seq_len", + "alibi_slopes", + "kv_cache_dtype", + "k_scale", + "v_scale", + "fp8_out_scale", + "mfma_type" + ], + "sources": [ + { + "model": "Qwen_Qwen3.5-35B-A3B", + "engine": "vllm", + "mode": "default", + "combo": "Qwen_Qwen3.5-35B-A3B__vllm_default", + "count": 48 + } + ] + }, + { + "test_case_id": "sig_c18c27c7fbd9", + "count": 48, + "args_sig": [ + { + "kind": "tensor", + "dtype": "bfloat16", + "shape": [ + 1, + 32, + 128 + ], + "stride": [ + 4096, + 128, + 1 + ], + "device": "cuda:0" + }, + { + "kind": "tensor", + "dtype": "float32", + "shape": [ + 1, + 32, + 5 + ], + "stride": [ + 160, + 5, + 1 + ], + "device": "cuda:0" + }, + { + "kind": "tensor", + "dtype": "float32", + "shape": [ + 1, + 32, + 5 + ], + "stride": [ + 160, + 5, + 1 + ], + "device": "cuda:0" + }, + { + "kind": "tensor", + "dtype": "bfloat16", + "shape": [ + 1, + 32, + 5, + 128 + ], + "stride": [ + 20480, + 640, + 128, + 1 + ], + "device": "cuda:0" + }, + { + "kind": "tensor", + "dtype": "bfloat16", + "shape": [ + 1, + 32, + 128 + ], + "stride": [ + 4096, + 128, + 1 + ], + "device": "cuda:0" + }, + { + "kind": "tensor", + "dtype": "bfloat16", + "shape": [ + 47679, + 4, + 16, + 16, + 8 + ], + "stride": [ + 8192, + 2048, + 128, + 8, + 1 + ], + "device": "cuda:0" + }, + { + "kind": "tensor", + "dtype": "bfloat16", + "shape": [ + 47679, + 4, + 128, + 16 + ], + "stride": [ + 8192, + 2048, + 16, + 1 + ], + "device": "cuda:0" + }, + { + "kind": "scalar", + "value": 4, + "py_type": "int" + }, + { + "kind": "scalar", + "value": 0.08838834764831845, + "py_type": "float" + }, + { + "kind": "tensor", + "dtype": "int32", + "shape": [ + 1, + 128 + ], + "stride": [ + 128, + 1 + ], + "device": "cuda:0" + }, + { + "kind": "tensor", + "dtype": "int32", + "shape": [ + 1 + ], + "stride": [ + 1 + ], + "device": "cuda:0" + }, + { + "kind": "tensor", + "dtype": "int32", + "shape": [ + 2 + ], + "stride": [ + 1 + ], + "device": "cuda:0" + }, + { + "kind": "scalar", + "value": 16, + "py_type": "int" + }, + { + "kind": "scalar", + "value": 1025, + "py_type": "int" + }, + { + "kind": "scalar", + "value": null, + "py_type": "NoneType" + }, + { + "kind": "opaque", + "py_type": "str", + "repr": "'auto'" + }, + { + "kind": "tensor", + "dtype": "float32", + "shape": [], + "stride": [], + "device": "cuda:0" + }, + { + "kind": "tensor", + "dtype": "float32", + "shape": [], + "stride": [], + "device": "cuda:0" + }, + { + "kind": "scalar", + "value": null, + "py_type": "NoneType" + }, + { + "kind": "opaque", + "py_type": "str", + "repr": "'f16'" + } + ], + "kwargs_sig": {}, + "params_repr": {}, + "args_names": [ + "out", + "exp_sums", + "max_logits", + "tmp_out", + "query", + "key_cache", + "value_cache", + "num_kv_heads", + "scale", + "block_tables", + "seq_lens", + "query_start_loc", + "block_size", + "max_seq_len", + "alibi_slopes", + "kv_cache_dtype", + "k_scale", + "v_scale", + "fp8_out_scale", + "mfma_type" + ], + "sources": [ + { + "model": "Qwen_Qwen3.5-35B-A3B", + "engine": "vllm", + "mode": "default", + "combo": "Qwen_Qwen3.5-35B-A3B__vllm_default", + "count": 48 + } + ] + }, + { + "test_case_id": "sig_7ffb6dca9258", + "count": 48, + "args_sig": [ + { + "kind": "tensor", + "dtype": "bfloat16", + "shape": [ + 1024, + 32, + 128 + ], + "stride": [ + 4096, + 128, + 1 + ], + "device": "cuda:0" + }, + { + "kind": "tensor", + "dtype": "float32", + "shape": [ + 1, + 32, + 4 + ], + "stride": [ + 128, + 4, + 1 + ], + "device": "cuda:0" + }, + { + "kind": "tensor", + "dtype": "float32", + "shape": [ + 1, + 32, + 4 + ], + "stride": [ + 128, + 4, + 1 + ], + "device": "cuda:0" + }, + { + "kind": "tensor", + "dtype": "bfloat16", + "shape": [ + 1, + 32, + 4, + 128 + ], + "stride": [ + 16384, + 512, + 128, + 1 + ], + "device": "cuda:0" + }, + { + "kind": "tensor", + "dtype": "bfloat16", + "shape": [ + 1024, + 32, + 128 + ], + "stride": [ + 4096, + 128, + 1 + ], + "device": "cuda:0" + }, + { + "kind": "tensor", + "dtype": "bfloat16", + "shape": [ + 45466, + 4, + 16, + 16, + 8 + ], + "stride": [ + 8192, + 2048, + 128, + 8, + 1 + ], + "device": "cuda:0" + }, + { + "kind": "tensor", + "dtype": "bfloat16", + "shape": [ + 45466, + 4, + 128, + 16 + ], + "stride": [ + 8192, + 2048, + 16, + 1 + ], + "device": "cuda:0" + }, + { + "kind": "scalar", + "value": 4, + "py_type": "int" + }, + { + "kind": "scalar", + "value": 0.08838834764831845, + "py_type": "float" + }, + { + "kind": "tensor", + "dtype": "int32", + "shape": [ + 1, + 128 + ], + "stride": [ + 128, + 1 + ], + "device": "cuda:0" + }, + { + "kind": "tensor", + "dtype": "int32", + "shape": [ + 1 + ], + "stride": [ + 1 + ], + "device": "cuda:0" + }, + { + "kind": "tensor", + "dtype": "int32", + "shape": [ + 2 + ], + "stride": [ + 1 + ], + "device": "cuda:0" + }, + { + "kind": "scalar", + "value": 16, + "py_type": "int" + }, + { + "kind": "scalar", + "value": 1024, + "py_type": "int" + }, + { + "kind": "scalar", + "value": null, + "py_type": "NoneType" + }, + { + "kind": "opaque", + "py_type": "str", + "repr": "'auto'" + }, + { + "kind": "tensor", + "dtype": "float32", + "shape": [], + "stride": [], + "device": "cuda:0" + }, + { + "kind": "tensor", + "dtype": "float32", + "shape": [], + "stride": [], + "device": "cuda:0" + }, + { + "kind": "scalar", + "value": null, + "py_type": "NoneType" + }, + { + "kind": "opaque", + "py_type": "str", + "repr": "'f16'" + } + ], + "kwargs_sig": {}, + "params_repr": {}, + "args_names": [ + "out", + "exp_sums", + "max_logits", + "tmp_out", + "query", + "key_cache", + "value_cache", + "num_kv_heads", + "scale", + "block_tables", + "seq_lens", + "query_start_loc", + "block_size", + "max_seq_len", + "alibi_slopes", + "kv_cache_dtype", + "k_scale", + "v_scale", + "fp8_out_scale", + "mfma_type" + ], + "sources": [ + { + "model": "Qwen_Qwen3.5-35B-A3B", + "engine": "vllm", + "mode": "hip", + "combo": "Qwen_Qwen3.5-35B-A3B__vllm_hip", + "count": 48 + } + ] + }, + { + "test_case_id": "sig_d3abec4575ad", + "count": 48, + "args_sig": [ + { + "kind": "tensor", + "dtype": "bfloat16", + "shape": [ + 1, + 32, + 128 + ], + "stride": [ + 4096, + 128, + 1 + ], + "device": "cuda:0" + }, + { + "kind": "tensor", + "dtype": "float32", + "shape": [ + 1, + 32, + 5 + ], + "stride": [ + 160, + 5, + 1 + ], + "device": "cuda:0" + }, + { + "kind": "tensor", + "dtype": "float32", + "shape": [ + 1, + 32, + 5 + ], + "stride": [ + 160, + 5, + 1 + ], + "device": "cuda:0" + }, + { + "kind": "tensor", + "dtype": "bfloat16", + "shape": [ + 1, + 32, + 5, + 128 + ], + "stride": [ + 20480, + 640, + 128, + 1 + ], + "device": "cuda:0" + }, + { + "kind": "tensor", + "dtype": "bfloat16", + "shape": [ + 1, + 32, + 128 + ], + "stride": [ + 4096, + 128, + 1 + ], + "device": "cuda:0" + }, + { + "kind": "tensor", + "dtype": "bfloat16", + "shape": [ + 45466, + 4, + 16, + 16, + 8 + ], + "stride": [ + 8192, + 2048, + 128, + 8, + 1 + ], + "device": "cuda:0" + }, + { + "kind": "tensor", + "dtype": "bfloat16", + "shape": [ + 45466, + 4, + 128, + 16 + ], + "stride": [ + 8192, + 2048, + 16, + 1 + ], + "device": "cuda:0" + }, + { + "kind": "scalar", + "value": 4, + "py_type": "int" + }, + { + "kind": "scalar", + "value": 0.08838834764831845, + "py_type": "float" + }, + { + "kind": "tensor", + "dtype": "int32", + "shape": [ + 1, + 128 + ], + "stride": [ + 128, + 1 + ], + "device": "cuda:0" + }, + { + "kind": "tensor", + "dtype": "int32", + "shape": [ + 1 + ], + "stride": [ + 1 + ], + "device": "cuda:0" + }, + { + "kind": "tensor", + "dtype": "int32", + "shape": [ + 2 + ], + "stride": [ + 1 + ], + "device": "cuda:0" + }, + { + "kind": "scalar", + "value": 16, + "py_type": "int" + }, + { + "kind": "scalar", + "value": 1025, + "py_type": "int" + }, + { + "kind": "scalar", + "value": null, + "py_type": "NoneType" + }, + { + "kind": "opaque", + "py_type": "str", + "repr": "'auto'" + }, + { + "kind": "tensor", + "dtype": "float32", + "shape": [], + "stride": [], + "device": "cuda:0" + }, + { + "kind": "tensor", + "dtype": "float32", + "shape": [], + "stride": [], + "device": "cuda:0" + }, + { + "kind": "scalar", + "value": null, + "py_type": "NoneType" + }, + { + "kind": "opaque", + "py_type": "str", + "repr": "'f16'" + } + ], + "kwargs_sig": {}, + "params_repr": {}, + "args_names": [ + "out", + "exp_sums", + "max_logits", + "tmp_out", + "query", + "key_cache", + "value_cache", + "num_kv_heads", + "scale", + "block_tables", + "seq_lens", + "query_start_loc", + "block_size", + "max_seq_len", + "alibi_slopes", + "kv_cache_dtype", + "k_scale", + "v_scale", + "fp8_out_scale", + "mfma_type" + ], + "sources": [ + { + "model": "Qwen_Qwen3.5-35B-A3B", + "engine": "vllm", + "mode": "hip", + "combo": "Qwen_Qwen3.5-35B-A3B__vllm_hip", + "count": 48 + } + ] + }, + { + "test_case_id": "c2", + "count": 1, + "args_sig": [ + { + "kind": "tensor", + "dtype": "bfloat16", + "shape": [ + 2, + 16, + 128 + ], + "device": "cuda:0" + }, + { + "kind": "tensor", + "dtype": "float32", + "shape": [ + 2, + 16, + 4 + ], + "device": "cuda:0" + }, + { + "kind": "tensor", + "dtype": "float32", + "shape": [ + 2, + 16, + 4 + ], + "device": "cuda:0" + }, + { + "kind": "tensor", + "dtype": "bfloat16", + "shape": [ + 2, + 16, + 4, + 128 + ], + "device": "cuda:0" + }, + { + "kind": "tensor", + "dtype": "bfloat16", + "shape": [ + 2, + 16, + 128 + ], + "device": "cuda:0" + }, + { + "kind": "tensor", + "dtype": "bfloat16", + "shape": [ + 128, + 1, + 16, + 16, + 8 + ], + "device": "cuda:0" + }, + { + "kind": "tensor", + "dtype": "bfloat16", + "shape": [ + 128, + 1, + 128, + 16 + ], + "device": "cuda:0" + }, + { + "kind": "scalar", + "value": 1, + "py_type": "int" + }, + { + "kind": "scalar", + "value": 0.08838834764831843, + "py_type": "float" + }, + { + "kind": "tensor", + "dtype": "int32", + "shape": [ + 2, + 64 + ], + "device": "cuda:0", + "data": [ + [ + 0, + 1, + 2, + 3, + 4, + 5, + 6, + 7, + 8, + 9, + 10, + 11, + 12, + 13, + 14, + 15, + 16, + 17, + 18, + 19, + 20, + 21, + 22, + 23, + 24, + 25, + 26, + 27, + 28, + 29, + 30, + 31, + 32, + 33, + 34, + 35, + 36, + 37, + 38, + 39, + 40, + 41, + 42, + 43, + 44, + 45, + 46, + 47, + 48, + 49, + 50, + 51, + 52, + 53, + 54, + 55, + 56, + 57, + 58, + 59, + 60, + 61, + 62, + 63 + ], + [ + 64, + 65, + 66, + 67, + 68, + 69, + 70, + 71, + 72, + 73, + 74, + 75, + 76, + 77, + 78, + 79, + 80, + 81, + 82, + 83, + 84, + 85, + 86, + 87, + 88, + 89, + 90, + 91, + 92, + 93, + 94, + 95, + 96, + 97, + 98, + 99, + 100, + 101, + 102, + 103, + 104, + 105, + 106, + 107, + 108, + 109, + 110, + 111, + 112, + 113, + 114, + 115, + 116, + 117, + 118, + 119, + 120, + 121, + 122, + 123, + 124, + 125, + 126, + 127 + ] + ] + }, + { + "kind": "tensor", + "dtype": "int32", + "shape": [ + 2 + ], + "device": "cuda:0", + "data": [ + 1024, + 1024 + ] + }, + { + "kind": "tensor", + "dtype": "int32", + "shape": [ + 3 + ], + "device": "cuda:0", + "data": [ + 0, + 1, + 2 + ] + }, + { + "kind": "scalar", + "value": 16, + "py_type": "int" + }, + { + "kind": "scalar", + "value": 1024, + "py_type": "int" + }, + { + "kind": "scalar", + "value": null, + "py_type": "NoneType" + }, + { + "kind": "opaque", + "py_type": "str", + "repr": "'auto'" + }, + { + "kind": "tensor", + "dtype": "float32", + "shape": [], + "device": "cuda:0", + "data": 1.0 + }, + { + "kind": "tensor", + "dtype": "float32", + "shape": [], + "device": "cuda:0", + "data": 1.0 + }, + { + "kind": "scalar", + "value": null, + "py_type": "NoneType" + }, + { + "kind": "opaque", + "py_type": "str", + "repr": "'f16'" + } + ], + "kwargs_sig": {}, + "params_repr": { + "S_seqs": 2, + "ctx_len": 1024, + "out_len": 1, + "heads": 16, + "kv_heads": 1, + "head_size": 128, + "block_size": 16, + "partition": 256, + "kv_alloc_gb": 0.0, + "note": "decode; disjoint KV blocks -> HBM-streaming", + "concurrency_B": 2 + }, + "args_names": [ + "out", + "exp_sums", + "max_logits", + "tmp_out", + "query", + "key_cache", + "value_cache", + "num_kv_heads", + "scale", + "block_tables", + "seq_lens", + "query_start_loc", + "block_size", + "max_seq_len", + "alibi_slopes", + "kv_cache_dtype", + "k_scale", + "v_scale", + "fp8_out_scale", + "mfma_type" + ], + "sources": [ + { + "model": "Qwen_Qwen3-235B-A22B-FP8", + "engine": "vllm", + "mode": "default", + "combo": "Qwen_Qwen3-235B-A22B-FP8__vllm_default", + "count": 376 + } + ], + "perf_only": true + }, + { + "test_case_id": "c32", + "count": 1, + "args_sig": [ + { + "kind": "tensor", + "dtype": "bfloat16", + "shape": [ + 32, + 16, + 128 + ], + "device": "cuda:0" + }, + { + "kind": "tensor", + "dtype": "float32", + "shape": [ + 32, + 16, + 4 + ], + "device": "cuda:0" + }, + { + "kind": "tensor", + "dtype": "float32", + "shape": [ + 32, + 16, + 4 + ], + "device": "cuda:0" + }, + { + "kind": "tensor", + "dtype": "bfloat16", + "shape": [ + 32, + 16, + 4, + 128 + ], + "device": "cuda:0" + }, + { + "kind": "tensor", + "dtype": "bfloat16", + "shape": [ + 32, + 16, + 128 + ], + "device": "cuda:0" + }, + { + "kind": "tensor", + "dtype": "bfloat16", + "shape": [ + 2048, + 1, + 16, + 16, + 8 + ], + "device": "cuda:0" + }, + { + "kind": "tensor", + "dtype": "bfloat16", + "shape": [ + 2048, + 1, + 128, + 16 + ], + "device": "cuda:0" + }, + { + "kind": "scalar", + "value": 1, + "py_type": "int" + }, + { + "kind": "scalar", + "value": 0.08838834764831843, + "py_type": "float" + }, + { + "kind": "tensor", + "dtype": "int32", + "shape": [ + 32, + 64 + ], + "device": "cuda:0", + "data": [ + [ + 0, + 1, + 2, + 3, + 4, + 5, + 6, + 7, + 8, + 9, + 10, + 11, + 12, + 13, + 14, + 15, + 16, + 17, + 18, + 19, + 20, + 21, + 22, + 23, + 24, + 25, + 26, + 27, + 28, + 29, + 30, + 31, + 32, + 33, + 34, + 35, + 36, + 37, + 38, + 39, + 40, + 41, + 42, + 43, + 44, + 45, + 46, + 47, + 48, + 49, + 50, + 51, + 52, + 53, + 54, + 55, + 56, + 57, + 58, + 59, + 60, + 61, + 62, + 63 + ], + [ + 64, + 65, + 66, + 67, + 68, + 69, + 70, + 71, + 72, + 73, + 74, + 75, + 76, + 77, + 78, + 79, + 80, + 81, + 82, + 83, + 84, + 85, + 86, + 87, + 88, + 89, + 90, + 91, + 92, + 93, + 94, + 95, + 96, + 97, + 98, + 99, + 100, + 101, + 102, + 103, + 104, + 105, + 106, + 107, + 108, + 109, + 110, + 111, + 112, + 113, + 114, + 115, + 116, + 117, + 118, + 119, + 120, + 121, + 122, + 123, + 124, + 125, + 126, + 127 + ], + [ + 128, + 129, + 130, + 131, + 132, + 133, + 134, + 135, + 136, + 137, + 138, + 139, + 140, + 141, + 142, + 143, + 144, + 145, + 146, + 147, + 148, + 149, + 150, + 151, + 152, + 153, + 154, + 155, + 156, + 157, + 158, + 159, + 160, + 161, + 162, + 163, + 164, + 165, + 166, + 167, + 168, + 169, + 170, + 171, + 172, + 173, + 174, + 175, + 176, + 177, + 178, + 179, + 180, + 181, + 182, + 183, + 184, + 185, + 186, + 187, + 188, + 189, + 190, + 191 + ], + [ + 192, + 193, + 194, + 195, + 196, + 197, + 198, + 199, + 200, + 201, + 202, + 203, + 204, + 205, + 206, + 207, + 208, + 209, + 210, + 211, + 212, + 213, + 214, + 215, + 216, + 217, + 218, + 219, + 220, + 221, + 222, + 223, + 224, + 225, + 226, + 227, + 228, + 229, + 230, + 231, + 232, + 233, + 234, + 235, + 236, + 237, + 238, + 239, + 240, + 241, + 242, + 243, + 244, + 245, + 246, + 247, + 248, + 249, + 250, + 251, + 252, + 253, + 254, + 255 + ], + [ + 256, + 257, + 258, + 259, + 260, + 261, + 262, + 263, + 264, + 265, + 266, + 267, + 268, + 269, + 270, + 271, + 272, + 273, + 274, + 275, + 276, + 277, + 278, + 279, + 280, + 281, + 282, + 283, + 284, + 285, + 286, + 287, + 288, + 289, + 290, + 291, + 292, + 293, + 294, + 295, + 296, + 297, + 298, + 299, + 300, + 301, + 302, + 303, + 304, + 305, + 306, + 307, + 308, + 309, + 310, + 311, + 312, + 313, + 314, + 315, + 316, + 317, + 318, + 319 + ], + [ + 320, + 321, + 322, + 323, + 324, + 325, + 326, + 327, + 328, + 329, + 330, + 331, + 332, + 333, + 334, + 335, + 336, + 337, + 338, + 339, + 340, + 341, + 342, + 343, + 344, + 345, + 346, + 347, + 348, + 349, + 350, + 351, + 352, + 353, + 354, + 355, + 356, + 357, + 358, + 359, + 360, + 361, + 362, + 363, + 364, + 365, + 366, + 367, + 368, + 369, + 370, + 371, + 372, + 373, + 374, + 375, + 376, + 377, + 378, + 379, + 380, + 381, + 382, + 383 + ], + [ + 384, + 385, + 386, + 387, + 388, + 389, + 390, + 391, + 392, + 393, + 394, + 395, + 396, + 397, + 398, + 399, + 400, + 401, + 402, + 403, + 404, + 405, + 406, + 407, + 408, + 409, + 410, + 411, + 412, + 413, + 414, + 415, + 416, + 417, + 418, + 419, + 420, + 421, + 422, + 423, + 424, + 425, + 426, + 427, + 428, + 429, + 430, + 431, + 432, + 433, + 434, + 435, + 436, + 437, + 438, + 439, + 440, + 441, + 442, + 443, + 444, + 445, + 446, + 447 + ], + [ + 448, + 449, + 450, + 451, + 452, + 453, + 454, + 455, + 456, + 457, + 458, + 459, + 460, + 461, + 462, + 463, + 464, + 465, + 466, + 467, + 468, + 469, + 470, + 471, + 472, + 473, + 474, + 475, + 476, + 477, + 478, + 479, + 480, + 481, + 482, + 483, + 484, + 485, + 486, + 487, + 488, + 489, + 490, + 491, + 492, + 493, + 494, + 495, + 496, + 497, + 498, + 499, + 500, + 501, + 502, + 503, + 504, + 505, + 506, + 507, + 508, + 509, + 510, + 511 + ], + [ + 512, + 513, + 514, + 515, + 516, + 517, + 518, + 519, + 520, + 521, + 522, + 523, + 524, + 525, + 526, + 527, + 528, + 529, + 530, + 531, + 532, + 533, + 534, + 535, + 536, + 537, + 538, + 539, + 540, + 541, + 542, + 543, + 544, + 545, + 546, + 547, + 548, + 549, + 550, + 551, + 552, + 553, + 554, + 555, + 556, + 557, + 558, + 559, + 560, + 561, + 562, + 563, + 564, + 565, + 566, + 567, + 568, + 569, + 570, + 571, + 572, + 573, + 574, + 575 + ], + [ + 576, + 577, + 578, + 579, + 580, + 581, + 582, + 583, + 584, + 585, + 586, + 587, + 588, + 589, + 590, + 591, + 592, + 593, + 594, + 595, + 596, + 597, + 598, + 599, + 600, + 601, + 602, + 603, + 604, + 605, + 606, + 607, + 608, + 609, + 610, + 611, + 612, + 613, + 614, + 615, + 616, + 617, + 618, + 619, + 620, + 621, + 622, + 623, + 624, + 625, + 626, + 627, + 628, + 629, + 630, + 631, + 632, + 633, + 634, + 635, + 636, + 637, + 638, + 639 + ], + [ + 640, + 641, + 642, + 643, + 644, + 645, + 646, + 647, + 648, + 649, + 650, + 651, + 652, + 653, + 654, + 655, + 656, + 657, + 658, + 659, + 660, + 661, + 662, + 663, + 664, + 665, + 666, + 667, + 668, + 669, + 670, + 671, + 672, + 673, + 674, + 675, + 676, + 677, + 678, + 679, + 680, + 681, + 682, + 683, + 684, + 685, + 686, + 687, + 688, + 689, + 690, + 691, + 692, + 693, + 694, + 695, + 696, + 697, + 698, + 699, + 700, + 701, + 702, + 703 + ], + [ + 704, + 705, + 706, + 707, + 708, + 709, + 710, + 711, + 712, + 713, + 714, + 715, + 716, + 717, + 718, + 719, + 720, + 721, + 722, + 723, + 724, + 725, + 726, + 727, + 728, + 729, + 730, + 731, + 732, + 733, + 734, + 735, + 736, + 737, + 738, + 739, + 740, + 741, + 742, + 743, + 744, + 745, + 746, + 747, + 748, + 749, + 750, + 751, + 752, + 753, + 754, + 755, + 756, + 757, + 758, + 759, + 760, + 761, + 762, + 763, + 764, + 765, + 766, + 767 + ], + [ + 768, + 769, + 770, + 771, + 772, + 773, + 774, + 775, + 776, + 777, + 778, + 779, + 780, + 781, + 782, + 783, + 784, + 785, + 786, + 787, + 788, + 789, + 790, + 791, + 792, + 793, + 794, + 795, + 796, + 797, + 798, + 799, + 800, + 801, + 802, + 803, + 804, + 805, + 806, + 807, + 808, + 809, + 810, + 811, + 812, + 813, + 814, + 815, + 816, + 817, + 818, + 819, + 820, + 821, + 822, + 823, + 824, + 825, + 826, + 827, + 828, + 829, + 830, + 831 + ], + [ + 832, + 833, + 834, + 835, + 836, + 837, + 838, + 839, + 840, + 841, + 842, + 843, + 844, + 845, + 846, + 847, + 848, + 849, + 850, + 851, + 852, + 853, + 854, + 855, + 856, + 857, + 858, + 859, + 860, + 861, + 862, + 863, + 864, + 865, + 866, + 867, + 868, + 869, + 870, + 871, + 872, + 873, + 874, + 875, + 876, + 877, + 878, + 879, + 880, + 881, + 882, + 883, + 884, + 885, + 886, + 887, + 888, + 889, + 890, + 891, + 892, + 893, + 894, + 895 + ], + [ + 896, + 897, + 898, + 899, + 900, + 901, + 902, + 903, + 904, + 905, + 906, + 907, + 908, + 909, + 910, + 911, + 912, + 913, + 914, + 915, + 916, + 917, + 918, + 919, + 920, + 921, + 922, + 923, + 924, + 925, + 926, + 927, + 928, + 929, + 930, + 931, + 932, + 933, + 934, + 935, + 936, + 937, + 938, + 939, + 940, + 941, + 942, + 943, + 944, + 945, + 946, + 947, + 948, + 949, + 950, + 951, + 952, + 953, + 954, + 955, + 956, + 957, + 958, + 959 + ], + [ + 960, + 961, + 962, + 963, + 964, + 965, + 966, + 967, + 968, + 969, + 970, + 971, + 972, + 973, + 974, + 975, + 976, + 977, + 978, + 979, + 980, + 981, + 982, + 983, + 984, + 985, + 986, + 987, + 988, + 989, + 990, + 991, + 992, + 993, + 994, + 995, + 996, + 997, + 998, + 999, + 1000, + 1001, + 1002, + 1003, + 1004, + 1005, + 1006, + 1007, + 1008, + 1009, + 1010, + 1011, + 1012, + 1013, + 1014, + 1015, + 1016, + 1017, + 1018, + 1019, + 1020, + 1021, + 1022, + 1023 + ], + [ + 1024, + 1025, + 1026, + 1027, + 1028, + 1029, + 1030, + 1031, + 1032, + 1033, + 1034, + 1035, + 1036, + 1037, + 1038, + 1039, + 1040, + 1041, + 1042, + 1043, + 1044, + 1045, + 1046, + 1047, + 1048, + 1049, + 1050, + 1051, + 1052, + 1053, + 1054, + 1055, + 1056, + 1057, + 1058, + 1059, + 1060, + 1061, + 1062, + 1063, + 1064, + 1065, + 1066, + 1067, + 1068, + 1069, + 1070, + 1071, + 1072, + 1073, + 1074, + 1075, + 1076, + 1077, + 1078, + 1079, + 1080, + 1081, + 1082, + 1083, + 1084, + 1085, + 1086, + 1087 + ], + [ + 1088, + 1089, + 1090, + 1091, + 1092, + 1093, + 1094, + 1095, + 1096, + 1097, + 1098, + 1099, + 1100, + 1101, + 1102, + 1103, + 1104, + 1105, + 1106, + 1107, + 1108, + 1109, + 1110, + 1111, + 1112, + 1113, + 1114, + 1115, + 1116, + 1117, + 1118, + 1119, + 1120, + 1121, + 1122, + 1123, + 1124, + 1125, + 1126, + 1127, + 1128, + 1129, + 1130, + 1131, + 1132, + 1133, + 1134, + 1135, + 1136, + 1137, + 1138, + 1139, + 1140, + 1141, + 1142, + 1143, + 1144, + 1145, + 1146, + 1147, + 1148, + 1149, + 1150, + 1151 + ], + [ + 1152, + 1153, + 1154, + 1155, + 1156, + 1157, + 1158, + 1159, + 1160, + 1161, + 1162, + 1163, + 1164, + 1165, + 1166, + 1167, + 1168, + 1169, + 1170, + 1171, + 1172, + 1173, + 1174, + 1175, + 1176, + 1177, + 1178, + 1179, + 1180, + 1181, + 1182, + 1183, + 1184, + 1185, + 1186, + 1187, + 1188, + 1189, + 1190, + 1191, + 1192, + 1193, + 1194, + 1195, + 1196, + 1197, + 1198, + 1199, + 1200, + 1201, + 1202, + 1203, + 1204, + 1205, + 1206, + 1207, + 1208, + 1209, + 1210, + 1211, + 1212, + 1213, + 1214, + 1215 + ], + [ + 1216, + 1217, + 1218, + 1219, + 1220, + 1221, + 1222, + 1223, + 1224, + 1225, + 1226, + 1227, + 1228, + 1229, + 1230, + 1231, + 1232, + 1233, + 1234, + 1235, + 1236, + 1237, + 1238, + 1239, + 1240, + 1241, + 1242, + 1243, + 1244, + 1245, + 1246, + 1247, + 1248, + 1249, + 1250, + 1251, + 1252, + 1253, + 1254, + 1255, + 1256, + 1257, + 1258, + 1259, + 1260, + 1261, + 1262, + 1263, + 1264, + 1265, + 1266, + 1267, + 1268, + 1269, + 1270, + 1271, + 1272, + 1273, + 1274, + 1275, + 1276, + 1277, + 1278, + 1279 + ], + [ + 1280, + 1281, + 1282, + 1283, + 1284, + 1285, + 1286, + 1287, + 1288, + 1289, + 1290, + 1291, + 1292, + 1293, + 1294, + 1295, + 1296, + 1297, + 1298, + 1299, + 1300, + 1301, + 1302, + 1303, + 1304, + 1305, + 1306, + 1307, + 1308, + 1309, + 1310, + 1311, + 1312, + 1313, + 1314, + 1315, + 1316, + 1317, + 1318, + 1319, + 1320, + 1321, + 1322, + 1323, + 1324, + 1325, + 1326, + 1327, + 1328, + 1329, + 1330, + 1331, + 1332, + 1333, + 1334, + 1335, + 1336, + 1337, + 1338, + 1339, + 1340, + 1341, + 1342, + 1343 + ], + [ + 1344, + 1345, + 1346, + 1347, + 1348, + 1349, + 1350, + 1351, + 1352, + 1353, + 1354, + 1355, + 1356, + 1357, + 1358, + 1359, + 1360, + 1361, + 1362, + 1363, + 1364, + 1365, + 1366, + 1367, + 1368, + 1369, + 1370, + 1371, + 1372, + 1373, + 1374, + 1375, + 1376, + 1377, + 1378, + 1379, + 1380, + 1381, + 1382, + 1383, + 1384, + 1385, + 1386, + 1387, + 1388, + 1389, + 1390, + 1391, + 1392, + 1393, + 1394, + 1395, + 1396, + 1397, + 1398, + 1399, + 1400, + 1401, + 1402, + 1403, + 1404, + 1405, + 1406, + 1407 + ], + [ + 1408, + 1409, + 1410, + 1411, + 1412, + 1413, + 1414, + 1415, + 1416, + 1417, + 1418, + 1419, + 1420, + 1421, + 1422, + 1423, + 1424, + 1425, + 1426, + 1427, + 1428, + 1429, + 1430, + 1431, + 1432, + 1433, + 1434, + 1435, + 1436, + 1437, + 1438, + 1439, + 1440, + 1441, + 1442, + 1443, + 1444, + 1445, + 1446, + 1447, + 1448, + 1449, + 1450, + 1451, + 1452, + 1453, + 1454, + 1455, + 1456, + 1457, + 1458, + 1459, + 1460, + 1461, + 1462, + 1463, + 1464, + 1465, + 1466, + 1467, + 1468, + 1469, + 1470, + 1471 + ], + [ + 1472, + 1473, + 1474, + 1475, + 1476, + 1477, + 1478, + 1479, + 1480, + 1481, + 1482, + 1483, + 1484, + 1485, + 1486, + 1487, + 1488, + 1489, + 1490, + 1491, + 1492, + 1493, + 1494, + 1495, + 1496, + 1497, + 1498, + 1499, + 1500, + 1501, + 1502, + 1503, + 1504, + 1505, + 1506, + 1507, + 1508, + 1509, + 1510, + 1511, + 1512, + 1513, + 1514, + 1515, + 1516, + 1517, + 1518, + 1519, + 1520, + 1521, + 1522, + 1523, + 1524, + 1525, + 1526, + 1527, + 1528, + 1529, + 1530, + 1531, + 1532, + 1533, + 1534, + 1535 + ], + [ + 1536, + 1537, + 1538, + 1539, + 1540, + 1541, + 1542, + 1543, + 1544, + 1545, + 1546, + 1547, + 1548, + 1549, + 1550, + 1551, + 1552, + 1553, + 1554, + 1555, + 1556, + 1557, + 1558, + 1559, + 1560, + 1561, + 1562, + 1563, + 1564, + 1565, + 1566, + 1567, + 1568, + 1569, + 1570, + 1571, + 1572, + 1573, + 1574, + 1575, + 1576, + 1577, + 1578, + 1579, + 1580, + 1581, + 1582, + 1583, + 1584, + 1585, + 1586, + 1587, + 1588, + 1589, + 1590, + 1591, + 1592, + 1593, + 1594, + 1595, + 1596, + 1597, + 1598, + 1599 + ], + [ + 1600, + 1601, + 1602, + 1603, + 1604, + 1605, + 1606, + 1607, + 1608, + 1609, + 1610, + 1611, + 1612, + 1613, + 1614, + 1615, + 1616, + 1617, + 1618, + 1619, + 1620, + 1621, + 1622, + 1623, + 1624, + 1625, + 1626, + 1627, + 1628, + 1629, + 1630, + 1631, + 1632, + 1633, + 1634, + 1635, + 1636, + 1637, + 1638, + 1639, + 1640, + 1641, + 1642, + 1643, + 1644, + 1645, + 1646, + 1647, + 1648, + 1649, + 1650, + 1651, + 1652, + 1653, + 1654, + 1655, + 1656, + 1657, + 1658, + 1659, + 1660, + 1661, + 1662, + 1663 + ], + [ + 1664, + 1665, + 1666, + 1667, + 1668, + 1669, + 1670, + 1671, + 1672, + 1673, + 1674, + 1675, + 1676, + 1677, + 1678, + 1679, + 1680, + 1681, + 1682, + 1683, + 1684, + 1685, + 1686, + 1687, + 1688, + 1689, + 1690, + 1691, + 1692, + 1693, + 1694, + 1695, + 1696, + 1697, + 1698, + 1699, + 1700, + 1701, + 1702, + 1703, + 1704, + 1705, + 1706, + 1707, + 1708, + 1709, + 1710, + 1711, + 1712, + 1713, + 1714, + 1715, + 1716, + 1717, + 1718, + 1719, + 1720, + 1721, + 1722, + 1723, + 1724, + 1725, + 1726, + 1727 + ], + [ + 1728, + 1729, + 1730, + 1731, + 1732, + 1733, + 1734, + 1735, + 1736, + 1737, + 1738, + 1739, + 1740, + 1741, + 1742, + 1743, + 1744, + 1745, + 1746, + 1747, + 1748, + 1749, + 1750, + 1751, + 1752, + 1753, + 1754, + 1755, + 1756, + 1757, + 1758, + 1759, + 1760, + 1761, + 1762, + 1763, + 1764, + 1765, + 1766, + 1767, + 1768, + 1769, + 1770, + 1771, + 1772, + 1773, + 1774, + 1775, + 1776, + 1777, + 1778, + 1779, + 1780, + 1781, + 1782, + 1783, + 1784, + 1785, + 1786, + 1787, + 1788, + 1789, + 1790, + 1791 + ], + [ + 1792, + 1793, + 1794, + 1795, + 1796, + 1797, + 1798, + 1799, + 1800, + 1801, + 1802, + 1803, + 1804, + 1805, + 1806, + 1807, + 1808, + 1809, + 1810, + 1811, + 1812, + 1813, + 1814, + 1815, + 1816, + 1817, + 1818, + 1819, + 1820, + 1821, + 1822, + 1823, + 1824, + 1825, + 1826, + 1827, + 1828, + 1829, + 1830, + 1831, + 1832, + 1833, + 1834, + 1835, + 1836, + 1837, + 1838, + 1839, + 1840, + 1841, + 1842, + 1843, + 1844, + 1845, + 1846, + 1847, + 1848, + 1849, + 1850, + 1851, + 1852, + 1853, + 1854, + 1855 + ], + [ + 1856, + 1857, + 1858, + 1859, + 1860, + 1861, + 1862, + 1863, + 1864, + 1865, + 1866, + 1867, + 1868, + 1869, + 1870, + 1871, + 1872, + 1873, + 1874, + 1875, + 1876, + 1877, + 1878, + 1879, + 1880, + 1881, + 1882, + 1883, + 1884, + 1885, + 1886, + 1887, + 1888, + 1889, + 1890, + 1891, + 1892, + 1893, + 1894, + 1895, + 1896, + 1897, + 1898, + 1899, + 1900, + 1901, + 1902, + 1903, + 1904, + 1905, + 1906, + 1907, + 1908, + 1909, + 1910, + 1911, + 1912, + 1913, + 1914, + 1915, + 1916, + 1917, + 1918, + 1919 + ], + [ + 1920, + 1921, + 1922, + 1923, + 1924, + 1925, + 1926, + 1927, + 1928, + 1929, + 1930, + 1931, + 1932, + 1933, + 1934, + 1935, + 1936, + 1937, + 1938, + 1939, + 1940, + 1941, + 1942, + 1943, + 1944, + 1945, + 1946, + 1947, + 1948, + 1949, + 1950, + 1951, + 1952, + 1953, + 1954, + 1955, + 1956, + 1957, + 1958, + 1959, + 1960, + 1961, + 1962, + 1963, + 1964, + 1965, + 1966, + 1967, + 1968, + 1969, + 1970, + 1971, + 1972, + 1973, + 1974, + 1975, + 1976, + 1977, + 1978, + 1979, + 1980, + 1981, + 1982, + 1983 + ], + [ + 1984, + 1985, + 1986, + 1987, + 1988, + 1989, + 1990, + 1991, + 1992, + 1993, + 1994, + 1995, + 1996, + 1997, + 1998, + 1999, + 2000, + 2001, + 2002, + 2003, + 2004, + 2005, + 2006, + 2007, + 2008, + 2009, + 2010, + 2011, + 2012, + 2013, + 2014, + 2015, + 2016, + 2017, + 2018, + 2019, + 2020, + 2021, + 2022, + 2023, + 2024, + 2025, + 2026, + 2027, + 2028, + 2029, + 2030, + 2031, + 2032, + 2033, + 2034, + 2035, + 2036, + 2037, + 2038, + 2039, + 2040, + 2041, + 2042, + 2043, + 2044, + 2045, + 2046, + 2047 + ] + ] + }, + { + "kind": "tensor", + "dtype": "int32", + "shape": [ + 32 + ], + "device": "cuda:0", + "data": [ + 1024, + 1024, + 1024, + 1024, + 1024, + 1024, + 1024, + 1024, + 1024, + 1024, + 1024, + 1024, + 1024, + 1024, + 1024, + 1024, + 1024, + 1024, + 1024, + 1024, + 1024, + 1024, + 1024, + 1024, + 1024, + 1024, + 1024, + 1024, + 1024, + 1024, + 1024, + 1024 + ] + }, + { + "kind": "tensor", + "dtype": "int32", + "shape": [ + 33 + ], + "device": "cuda:0", + "data": [ + 0, + 1, + 2, + 3, + 4, + 5, + 6, + 7, + 8, + 9, + 10, + 11, + 12, + 13, + 14, + 15, + 16, + 17, + 18, + 19, + 20, + 21, + 22, + 23, + 24, + 25, + 26, + 27, + 28, + 29, + 30, + 31, + 32 + ] + }, + { + "kind": "scalar", + "value": 16, + "py_type": "int" + }, + { + "kind": "scalar", + "value": 1024, + "py_type": "int" + }, + { + "kind": "scalar", + "value": null, + "py_type": "NoneType" + }, + { + "kind": "opaque", + "py_type": "str", + "repr": "'auto'" + }, + { + "kind": "tensor", + "dtype": "float32", + "shape": [], + "device": "cuda:0", + "data": 1.0 + }, + { + "kind": "tensor", + "dtype": "float32", + "shape": [], + "device": "cuda:0", + "data": 1.0 + }, + { + "kind": "scalar", + "value": null, + "py_type": "NoneType" + }, + { + "kind": "opaque", + "py_type": "str", + "repr": "'f16'" + } + ], + "kwargs_sig": {}, + "params_repr": { + "S_seqs": 32, + "ctx_len": 1024, + "out_len": 1, + "heads": 16, + "kv_heads": 1, + "head_size": 128, + "block_size": 16, + "partition": 256, + "kv_alloc_gb": 0.0, + "note": "decode; disjoint KV blocks -> HBM-streaming", + "concurrency_B": 32 + }, + "args_names": [ + "out", + "exp_sums", + "max_logits", + "tmp_out", + "query", + "key_cache", + "value_cache", + "num_kv_heads", + "scale", + "block_tables", + "seq_lens", + "query_start_loc", + "block_size", + "max_seq_len", + "alibi_slopes", + "kv_cache_dtype", + "k_scale", + "v_scale", + "fp8_out_scale", + "mfma_type" + ], + "sources": [ + { + "model": "Qwen_Qwen3-235B-A22B-FP8", + "engine": "vllm", + "mode": "default", + "combo": "Qwen_Qwen3-235B-A22B-FP8__vllm_default", + "count": 376 + } + ], + "perf_only": true + }, + { + "test_case_id": "c64", + "count": 1, + "args_sig": [ + { + "kind": "tensor", + "dtype": "bfloat16", + "shape": [ + 64, + 16, + 128 + ], + "device": "cuda:0" + }, + { + "kind": "tensor", + "dtype": "float32", + "shape": [ + 64, + 16, + 4 + ], + "device": "cuda:0" + }, + { + "kind": "tensor", + "dtype": "float32", + "shape": [ + 64, + 16, + 4 + ], + "device": "cuda:0" + }, + { + "kind": "tensor", + "dtype": "bfloat16", + "shape": [ + 64, + 16, + 4, + 128 + ], + "device": "cuda:0" + }, + { + "kind": "tensor", + "dtype": "bfloat16", + "shape": [ + 64, + 16, + 128 + ], + "device": "cuda:0" + }, + { + "kind": "tensor", + "dtype": "bfloat16", + "shape": [ + 4096, + 1, + 16, + 16, + 8 + ], + "device": "cuda:0" + }, + { + "kind": "tensor", + "dtype": "bfloat16", + "shape": [ + 4096, + 1, + 128, + 16 + ], + "device": "cuda:0" + }, + { + "kind": "scalar", + "value": 1, + "py_type": "int" + }, + { + "kind": "scalar", + "value": 0.08838834764831843, + "py_type": "float" + }, + { + "kind": "tensor", + "dtype": "int32", + "shape": [ + 64, + 64 + ], + "device": "cuda:0", + "data": [ + [ + 0, + 1, + 2, + 3, + 4, + 5, + 6, + 7, + 8, + 9, + 10, + 11, + 12, + 13, + 14, + 15, + 16, + 17, + 18, + 19, + 20, + 21, + 22, + 23, + 24, + 25, + 26, + 27, + 28, + 29, + 30, + 31, + 32, + 33, + 34, + 35, + 36, + 37, + 38, + 39, + 40, + 41, + 42, + 43, + 44, + 45, + 46, + 47, + 48, + 49, + 50, + 51, + 52, + 53, + 54, + 55, + 56, + 57, + 58, + 59, + 60, + 61, + 62, + 63 + ], + [ + 64, + 65, + 66, + 67, + 68, + 69, + 70, + 71, + 72, + 73, + 74, + 75, + 76, + 77, + 78, + 79, + 80, + 81, + 82, + 83, + 84, + 85, + 86, + 87, + 88, + 89, + 90, + 91, + 92, + 93, + 94, + 95, + 96, + 97, + 98, + 99, + 100, + 101, + 102, + 103, + 104, + 105, + 106, + 107, + 108, + 109, + 110, + 111, + 112, + 113, + 114, + 115, + 116, + 117, + 118, + 119, + 120, + 121, + 122, + 123, + 124, + 125, + 126, + 127 + ], + [ + 128, + 129, + 130, + 131, + 132, + 133, + 134, + 135, + 136, + 137, + 138, + 139, + 140, + 141, + 142, + 143, + 144, + 145, + 146, + 147, + 148, + 149, + 150, + 151, + 152, + 153, + 154, + 155, + 156, + 157, + 158, + 159, + 160, + 161, + 162, + 163, + 164, + 165, + 166, + 167, + 168, + 169, + 170, + 171, + 172, + 173, + 174, + 175, + 176, + 177, + 178, + 179, + 180, + 181, + 182, + 183, + 184, + 185, + 186, + 187, + 188, + 189, + 190, + 191 + ], + [ + 192, + 193, + 194, + 195, + 196, + 197, + 198, + 199, + 200, + 201, + 202, + 203, + 204, + 205, + 206, + 207, + 208, + 209, + 210, + 211, + 212, + 213, + 214, + 215, + 216, + 217, + 218, + 219, + 220, + 221, + 222, + 223, + 224, + 225, + 226, + 227, + 228, + 229, + 230, + 231, + 232, + 233, + 234, + 235, + 236, + 237, + 238, + 239, + 240, + 241, + 242, + 243, + 244, + 245, + 246, + 247, + 248, + 249, + 250, + 251, + 252, + 253, + 254, + 255 + ], + [ + 256, + 257, + 258, + 259, + 260, + 261, + 262, + 263, + 264, + 265, + 266, + 267, + 268, + 269, + 270, + 271, + 272, + 273, + 274, + 275, + 276, + 277, + 278, + 279, + 280, + 281, + 282, + 283, + 284, + 285, + 286, + 287, + 288, + 289, + 290, + 291, + 292, + 293, + 294, + 295, + 296, + 297, + 298, + 299, + 300, + 301, + 302, + 303, + 304, + 305, + 306, + 307, + 308, + 309, + 310, + 311, + 312, + 313, + 314, + 315, + 316, + 317, + 318, + 319 + ], + [ + 320, + 321, + 322, + 323, + 324, + 325, + 326, + 327, + 328, + 329, + 330, + 331, + 332, + 333, + 334, + 335, + 336, + 337, + 338, + 339, + 340, + 341, + 342, + 343, + 344, + 345, + 346, + 347, + 348, + 349, + 350, + 351, + 352, + 353, + 354, + 355, + 356, + 357, + 358, + 359, + 360, + 361, + 362, + 363, + 364, + 365, + 366, + 367, + 368, + 369, + 370, + 371, + 372, + 373, + 374, + 375, + 376, + 377, + 378, + 379, + 380, + 381, + 382, + 383 + ], + [ + 384, + 385, + 386, + 387, + 388, + 389, + 390, + 391, + 392, + 393, + 394, + 395, + 396, + 397, + 398, + 399, + 400, + 401, + 402, + 403, + 404, + 405, + 406, + 407, + 408, + 409, + 410, + 411, + 412, + 413, + 414, + 415, + 416, + 417, + 418, + 419, + 420, + 421, + 422, + 423, + 424, + 425, + 426, + 427, + 428, + 429, + 430, + 431, + 432, + 433, + 434, + 435, + 436, + 437, + 438, + 439, + 440, + 441, + 442, + 443, + 444, + 445, + 446, + 447 + ], + [ + 448, + 449, + 450, + 451, + 452, + 453, + 454, + 455, + 456, + 457, + 458, + 459, + 460, + 461, + 462, + 463, + 464, + 465, + 466, + 467, + 468, + 469, + 470, + 471, + 472, + 473, + 474, + 475, + 476, + 477, + 478, + 479, + 480, + 481, + 482, + 483, + 484, + 485, + 486, + 487, + 488, + 489, + 490, + 491, + 492, + 493, + 494, + 495, + 496, + 497, + 498, + 499, + 500, + 501, + 502, + 503, + 504, + 505, + 506, + 507, + 508, + 509, + 510, + 511 + ], + [ + 512, + 513, + 514, + 515, + 516, + 517, + 518, + 519, + 520, + 521, + 522, + 523, + 524, + 525, + 526, + 527, + 528, + 529, + 530, + 531, + 532, + 533, + 534, + 535, + 536, + 537, + 538, + 539, + 540, + 541, + 542, + 543, + 544, + 545, + 546, + 547, + 548, + 549, + 550, + 551, + 552, + 553, + 554, + 555, + 556, + 557, + 558, + 559, + 560, + 561, + 562, + 563, + 564, + 565, + 566, + 567, + 568, + 569, + 570, + 571, + 572, + 573, + 574, + 575 + ], + [ + 576, + 577, + 578, + 579, + 580, + 581, + 582, + 583, + 584, + 585, + 586, + 587, + 588, + 589, + 590, + 591, + 592, + 593, + 594, + 595, + 596, + 597, + 598, + 599, + 600, + 601, + 602, + 603, + 604, + 605, + 606, + 607, + 608, + 609, + 610, + 611, + 612, + 613, + 614, + 615, + 616, + 617, + 618, + 619, + 620, + 621, + 622, + 623, + 624, + 625, + 626, + 627, + 628, + 629, + 630, + 631, + 632, + 633, + 634, + 635, + 636, + 637, + 638, + 639 + ], + [ + 640, + 641, + 642, + 643, + 644, + 645, + 646, + 647, + 648, + 649, + 650, + 651, + 652, + 653, + 654, + 655, + 656, + 657, + 658, + 659, + 660, + 661, + 662, + 663, + 664, + 665, + 666, + 667, + 668, + 669, + 670, + 671, + 672, + 673, + 674, + 675, + 676, + 677, + 678, + 679, + 680, + 681, + 682, + 683, + 684, + 685, + 686, + 687, + 688, + 689, + 690, + 691, + 692, + 693, + 694, + 695, + 696, + 697, + 698, + 699, + 700, + 701, + 702, + 703 + ], + [ + 704, + 705, + 706, + 707, + 708, + 709, + 710, + 711, + 712, + 713, + 714, + 715, + 716, + 717, + 718, + 719, + 720, + 721, + 722, + 723, + 724, + 725, + 726, + 727, + 728, + 729, + 730, + 731, + 732, + 733, + 734, + 735, + 736, + 737, + 738, + 739, + 740, + 741, + 742, + 743, + 744, + 745, + 746, + 747, + 748, + 749, + 750, + 751, + 752, + 753, + 754, + 755, + 756, + 757, + 758, + 759, + 760, + 761, + 762, + 763, + 764, + 765, + 766, + 767 + ], + [ + 768, + 769, + 770, + 771, + 772, + 773, + 774, + 775, + 776, + 777, + 778, + 779, + 780, + 781, + 782, + 783, + 784, + 785, + 786, + 787, + 788, + 789, + 790, + 791, + 792, + 793, + 794, + 795, + 796, + 797, + 798, + 799, + 800, + 801, + 802, + 803, + 804, + 805, + 806, + 807, + 808, + 809, + 810, + 811, + 812, + 813, + 814, + 815, + 816, + 817, + 818, + 819, + 820, + 821, + 822, + 823, + 824, + 825, + 826, + 827, + 828, + 829, + 830, + 831 + ], + [ + 832, + 833, + 834, + 835, + 836, + 837, + 838, + 839, + 840, + 841, + 842, + 843, + 844, + 845, + 846, + 847, + 848, + 849, + 850, + 851, + 852, + 853, + 854, + 855, + 856, + 857, + 858, + 859, + 860, + 861, + 862, + 863, + 864, + 865, + 866, + 867, + 868, + 869, + 870, + 871, + 872, + 873, + 874, + 875, + 876, + 877, + 878, + 879, + 880, + 881, + 882, + 883, + 884, + 885, + 886, + 887, + 888, + 889, + 890, + 891, + 892, + 893, + 894, + 895 + ], + [ + 896, + 897, + 898, + 899, + 900, + 901, + 902, + 903, + 904, + 905, + 906, + 907, + 908, + 909, + 910, + 911, + 912, + 913, + 914, + 915, + 916, + 917, + 918, + 919, + 920, + 921, + 922, + 923, + 924, + 925, + 926, + 927, + 928, + 929, + 930, + 931, + 932, + 933, + 934, + 935, + 936, + 937, + 938, + 939, + 940, + 941, + 942, + 943, + 944, + 945, + 946, + 947, + 948, + 949, + 950, + 951, + 952, + 953, + 954, + 955, + 956, + 957, + 958, + 959 + ], + [ + 960, + 961, + 962, + 963, + 964, + 965, + 966, + 967, + 968, + 969, + 970, + 971, + 972, + 973, + 974, + 975, + 976, + 977, + 978, + 979, + 980, + 981, + 982, + 983, + 984, + 985, + 986, + 987, + 988, + 989, + 990, + 991, + 992, + 993, + 994, + 995, + 996, + 997, + 998, + 999, + 1000, + 1001, + 1002, + 1003, + 1004, + 1005, + 1006, + 1007, + 1008, + 1009, + 1010, + 1011, + 1012, + 1013, + 1014, + 1015, + 1016, + 1017, + 1018, + 1019, + 1020, + 1021, + 1022, + 1023 + ], + [ + 1024, + 1025, + 1026, + 1027, + 1028, + 1029, + 1030, + 1031, + 1032, + 1033, + 1034, + 1035, + 1036, + 1037, + 1038, + 1039, + 1040, + 1041, + 1042, + 1043, + 1044, + 1045, + 1046, + 1047, + 1048, + 1049, + 1050, + 1051, + 1052, + 1053, + 1054, + 1055, + 1056, + 1057, + 1058, + 1059, + 1060, + 1061, + 1062, + 1063, + 1064, + 1065, + 1066, + 1067, + 1068, + 1069, + 1070, + 1071, + 1072, + 1073, + 1074, + 1075, + 1076, + 1077, + 1078, + 1079, + 1080, + 1081, + 1082, + 1083, + 1084, + 1085, + 1086, + 1087 + ], + [ + 1088, + 1089, + 1090, + 1091, + 1092, + 1093, + 1094, + 1095, + 1096, + 1097, + 1098, + 1099, + 1100, + 1101, + 1102, + 1103, + 1104, + 1105, + 1106, + 1107, + 1108, + 1109, + 1110, + 1111, + 1112, + 1113, + 1114, + 1115, + 1116, + 1117, + 1118, + 1119, + 1120, + 1121, + 1122, + 1123, + 1124, + 1125, + 1126, + 1127, + 1128, + 1129, + 1130, + 1131, + 1132, + 1133, + 1134, + 1135, + 1136, + 1137, + 1138, + 1139, + 1140, + 1141, + 1142, + 1143, + 1144, + 1145, + 1146, + 1147, + 1148, + 1149, + 1150, + 1151 + ], + [ + 1152, + 1153, + 1154, + 1155, + 1156, + 1157, + 1158, + 1159, + 1160, + 1161, + 1162, + 1163, + 1164, + 1165, + 1166, + 1167, + 1168, + 1169, + 1170, + 1171, + 1172, + 1173, + 1174, + 1175, + 1176, + 1177, + 1178, + 1179, + 1180, + 1181, + 1182, + 1183, + 1184, + 1185, + 1186, + 1187, + 1188, + 1189, + 1190, + 1191, + 1192, + 1193, + 1194, + 1195, + 1196, + 1197, + 1198, + 1199, + 1200, + 1201, + 1202, + 1203, + 1204, + 1205, + 1206, + 1207, + 1208, + 1209, + 1210, + 1211, + 1212, + 1213, + 1214, + 1215 + ], + [ + 1216, + 1217, + 1218, + 1219, + 1220, + 1221, + 1222, + 1223, + 1224, + 1225, + 1226, + 1227, + 1228, + 1229, + 1230, + 1231, + 1232, + 1233, + 1234, + 1235, + 1236, + 1237, + 1238, + 1239, + 1240, + 1241, + 1242, + 1243, + 1244, + 1245, + 1246, + 1247, + 1248, + 1249, + 1250, + 1251, + 1252, + 1253, + 1254, + 1255, + 1256, + 1257, + 1258, + 1259, + 1260, + 1261, + 1262, + 1263, + 1264, + 1265, + 1266, + 1267, + 1268, + 1269, + 1270, + 1271, + 1272, + 1273, + 1274, + 1275, + 1276, + 1277, + 1278, + 1279 + ], + [ + 1280, + 1281, + 1282, + 1283, + 1284, + 1285, + 1286, + 1287, + 1288, + 1289, + 1290, + 1291, + 1292, + 1293, + 1294, + 1295, + 1296, + 1297, + 1298, + 1299, + 1300, + 1301, + 1302, + 1303, + 1304, + 1305, + 1306, + 1307, + 1308, + 1309, + 1310, + 1311, + 1312, + 1313, + 1314, + 1315, + 1316, + 1317, + 1318, + 1319, + 1320, + 1321, + 1322, + 1323, + 1324, + 1325, + 1326, + 1327, + 1328, + 1329, + 1330, + 1331, + 1332, + 1333, + 1334, + 1335, + 1336, + 1337, + 1338, + 1339, + 1340, + 1341, + 1342, + 1343 + ], + [ + 1344, + 1345, + 1346, + 1347, + 1348, + 1349, + 1350, + 1351, + 1352, + 1353, + 1354, + 1355, + 1356, + 1357, + 1358, + 1359, + 1360, + 1361, + 1362, + 1363, + 1364, + 1365, + 1366, + 1367, + 1368, + 1369, + 1370, + 1371, + 1372, + 1373, + 1374, + 1375, + 1376, + 1377, + 1378, + 1379, + 1380, + 1381, + 1382, + 1383, + 1384, + 1385, + 1386, + 1387, + 1388, + 1389, + 1390, + 1391, + 1392, + 1393, + 1394, + 1395, + 1396, + 1397, + 1398, + 1399, + 1400, + 1401, + 1402, + 1403, + 1404, + 1405, + 1406, + 1407 + ], + [ + 1408, + 1409, + 1410, + 1411, + 1412, + 1413, + 1414, + 1415, + 1416, + 1417, + 1418, + 1419, + 1420, + 1421, + 1422, + 1423, + 1424, + 1425, + 1426, + 1427, + 1428, + 1429, + 1430, + 1431, + 1432, + 1433, + 1434, + 1435, + 1436, + 1437, + 1438, + 1439, + 1440, + 1441, + 1442, + 1443, + 1444, + 1445, + 1446, + 1447, + 1448, + 1449, + 1450, + 1451, + 1452, + 1453, + 1454, + 1455, + 1456, + 1457, + 1458, + 1459, + 1460, + 1461, + 1462, + 1463, + 1464, + 1465, + 1466, + 1467, + 1468, + 1469, + 1470, + 1471 + ], + [ + 1472, + 1473, + 1474, + 1475, + 1476, + 1477, + 1478, + 1479, + 1480, + 1481, + 1482, + 1483, + 1484, + 1485, + 1486, + 1487, + 1488, + 1489, + 1490, + 1491, + 1492, + 1493, + 1494, + 1495, + 1496, + 1497, + 1498, + 1499, + 1500, + 1501, + 1502, + 1503, + 1504, + 1505, + 1506, + 1507, + 1508, + 1509, + 1510, + 1511, + 1512, + 1513, + 1514, + 1515, + 1516, + 1517, + 1518, + 1519, + 1520, + 1521, + 1522, + 1523, + 1524, + 1525, + 1526, + 1527, + 1528, + 1529, + 1530, + 1531, + 1532, + 1533, + 1534, + 1535 + ], + [ + 1536, + 1537, + 1538, + 1539, + 1540, + 1541, + 1542, + 1543, + 1544, + 1545, + 1546, + 1547, + 1548, + 1549, + 1550, + 1551, + 1552, + 1553, + 1554, + 1555, + 1556, + 1557, + 1558, + 1559, + 1560, + 1561, + 1562, + 1563, + 1564, + 1565, + 1566, + 1567, + 1568, + 1569, + 1570, + 1571, + 1572, + 1573, + 1574, + 1575, + 1576, + 1577, + 1578, + 1579, + 1580, + 1581, + 1582, + 1583, + 1584, + 1585, + 1586, + 1587, + 1588, + 1589, + 1590, + 1591, + 1592, + 1593, + 1594, + 1595, + 1596, + 1597, + 1598, + 1599 + ], + [ + 1600, + 1601, + 1602, + 1603, + 1604, + 1605, + 1606, + 1607, + 1608, + 1609, + 1610, + 1611, + 1612, + 1613, + 1614, + 1615, + 1616, + 1617, + 1618, + 1619, + 1620, + 1621, + 1622, + 1623, + 1624, + 1625, + 1626, + 1627, + 1628, + 1629, + 1630, + 1631, + 1632, + 1633, + 1634, + 1635, + 1636, + 1637, + 1638, + 1639, + 1640, + 1641, + 1642, + 1643, + 1644, + 1645, + 1646, + 1647, + 1648, + 1649, + 1650, + 1651, + 1652, + 1653, + 1654, + 1655, + 1656, + 1657, + 1658, + 1659, + 1660, + 1661, + 1662, + 1663 + ], + [ + 1664, + 1665, + 1666, + 1667, + 1668, + 1669, + 1670, + 1671, + 1672, + 1673, + 1674, + 1675, + 1676, + 1677, + 1678, + 1679, + 1680, + 1681, + 1682, + 1683, + 1684, + 1685, + 1686, + 1687, + 1688, + 1689, + 1690, + 1691, + 1692, + 1693, + 1694, + 1695, + 1696, + 1697, + 1698, + 1699, + 1700, + 1701, + 1702, + 1703, + 1704, + 1705, + 1706, + 1707, + 1708, + 1709, + 1710, + 1711, + 1712, + 1713, + 1714, + 1715, + 1716, + 1717, + 1718, + 1719, + 1720, + 1721, + 1722, + 1723, + 1724, + 1725, + 1726, + 1727 + ], + [ + 1728, + 1729, + 1730, + 1731, + 1732, + 1733, + 1734, + 1735, + 1736, + 1737, + 1738, + 1739, + 1740, + 1741, + 1742, + 1743, + 1744, + 1745, + 1746, + 1747, + 1748, + 1749, + 1750, + 1751, + 1752, + 1753, + 1754, + 1755, + 1756, + 1757, + 1758, + 1759, + 1760, + 1761, + 1762, + 1763, + 1764, + 1765, + 1766, + 1767, + 1768, + 1769, + 1770, + 1771, + 1772, + 1773, + 1774, + 1775, + 1776, + 1777, + 1778, + 1779, + 1780, + 1781, + 1782, + 1783, + 1784, + 1785, + 1786, + 1787, + 1788, + 1789, + 1790, + 1791 + ], + [ + 1792, + 1793, + 1794, + 1795, + 1796, + 1797, + 1798, + 1799, + 1800, + 1801, + 1802, + 1803, + 1804, + 1805, + 1806, + 1807, + 1808, + 1809, + 1810, + 1811, + 1812, + 1813, + 1814, + 1815, + 1816, + 1817, + 1818, + 1819, + 1820, + 1821, + 1822, + 1823, + 1824, + 1825, + 1826, + 1827, + 1828, + 1829, + 1830, + 1831, + 1832, + 1833, + 1834, + 1835, + 1836, + 1837, + 1838, + 1839, + 1840, + 1841, + 1842, + 1843, + 1844, + 1845, + 1846, + 1847, + 1848, + 1849, + 1850, + 1851, + 1852, + 1853, + 1854, + 1855 + ], + [ + 1856, + 1857, + 1858, + 1859, + 1860, + 1861, + 1862, + 1863, + 1864, + 1865, + 1866, + 1867, + 1868, + 1869, + 1870, + 1871, + 1872, + 1873, + 1874, + 1875, + 1876, + 1877, + 1878, + 1879, + 1880, + 1881, + 1882, + 1883, + 1884, + 1885, + 1886, + 1887, + 1888, + 1889, + 1890, + 1891, + 1892, + 1893, + 1894, + 1895, + 1896, + 1897, + 1898, + 1899, + 1900, + 1901, + 1902, + 1903, + 1904, + 1905, + 1906, + 1907, + 1908, + 1909, + 1910, + 1911, + 1912, + 1913, + 1914, + 1915, + 1916, + 1917, + 1918, + 1919 + ], + [ + 1920, + 1921, + 1922, + 1923, + 1924, + 1925, + 1926, + 1927, + 1928, + 1929, + 1930, + 1931, + 1932, + 1933, + 1934, + 1935, + 1936, + 1937, + 1938, + 1939, + 1940, + 1941, + 1942, + 1943, + 1944, + 1945, + 1946, + 1947, + 1948, + 1949, + 1950, + 1951, + 1952, + 1953, + 1954, + 1955, + 1956, + 1957, + 1958, + 1959, + 1960, + 1961, + 1962, + 1963, + 1964, + 1965, + 1966, + 1967, + 1968, + 1969, + 1970, + 1971, + 1972, + 1973, + 1974, + 1975, + 1976, + 1977, + 1978, + 1979, + 1980, + 1981, + 1982, + 1983 + ], + [ + 1984, + 1985, + 1986, + 1987, + 1988, + 1989, + 1990, + 1991, + 1992, + 1993, + 1994, + 1995, + 1996, + 1997, + 1998, + 1999, + 2000, + 2001, + 2002, + 2003, + 2004, + 2005, + 2006, + 2007, + 2008, + 2009, + 2010, + 2011, + 2012, + 2013, + 2014, + 2015, + 2016, + 2017, + 2018, + 2019, + 2020, + 2021, + 2022, + 2023, + 2024, + 2025, + 2026, + 2027, + 2028, + 2029, + 2030, + 2031, + 2032, + 2033, + 2034, + 2035, + 2036, + 2037, + 2038, + 2039, + 2040, + 2041, + 2042, + 2043, + 2044, + 2045, + 2046, + 2047 + ], + [ + 2048, + 2049, + 2050, + 2051, + 2052, + 2053, + 2054, + 2055, + 2056, + 2057, + 2058, + 2059, + 2060, + 2061, + 2062, + 2063, + 2064, + 2065, + 2066, + 2067, + 2068, + 2069, + 2070, + 2071, + 2072, + 2073, + 2074, + 2075, + 2076, + 2077, + 2078, + 2079, + 2080, + 2081, + 2082, + 2083, + 2084, + 2085, + 2086, + 2087, + 2088, + 2089, + 2090, + 2091, + 2092, + 2093, + 2094, + 2095, + 2096, + 2097, + 2098, + 2099, + 2100, + 2101, + 2102, + 2103, + 2104, + 2105, + 2106, + 2107, + 2108, + 2109, + 2110, + 2111 + ], + [ + 2112, + 2113, + 2114, + 2115, + 2116, + 2117, + 2118, + 2119, + 2120, + 2121, + 2122, + 2123, + 2124, + 2125, + 2126, + 2127, + 2128, + 2129, + 2130, + 2131, + 2132, + 2133, + 2134, + 2135, + 2136, + 2137, + 2138, + 2139, + 2140, + 2141, + 2142, + 2143, + 2144, + 2145, + 2146, + 2147, + 2148, + 2149, + 2150, + 2151, + 2152, + 2153, + 2154, + 2155, + 2156, + 2157, + 2158, + 2159, + 2160, + 2161, + 2162, + 2163, + 2164, + 2165, + 2166, + 2167, + 2168, + 2169, + 2170, + 2171, + 2172, + 2173, + 2174, + 2175 + ], + [ + 2176, + 2177, + 2178, + 2179, + 2180, + 2181, + 2182, + 2183, + 2184, + 2185, + 2186, + 2187, + 2188, + 2189, + 2190, + 2191, + 2192, + 2193, + 2194, + 2195, + 2196, + 2197, + 2198, + 2199, + 2200, + 2201, + 2202, + 2203, + 2204, + 2205, + 2206, + 2207, + 2208, + 2209, + 2210, + 2211, + 2212, + 2213, + 2214, + 2215, + 2216, + 2217, + 2218, + 2219, + 2220, + 2221, + 2222, + 2223, + 2224, + 2225, + 2226, + 2227, + 2228, + 2229, + 2230, + 2231, + 2232, + 2233, + 2234, + 2235, + 2236, + 2237, + 2238, + 2239 + ], + [ + 2240, + 2241, + 2242, + 2243, + 2244, + 2245, + 2246, + 2247, + 2248, + 2249, + 2250, + 2251, + 2252, + 2253, + 2254, + 2255, + 2256, + 2257, + 2258, + 2259, + 2260, + 2261, + 2262, + 2263, + 2264, + 2265, + 2266, + 2267, + 2268, + 2269, + 2270, + 2271, + 2272, + 2273, + 2274, + 2275, + 2276, + 2277, + 2278, + 2279, + 2280, + 2281, + 2282, + 2283, + 2284, + 2285, + 2286, + 2287, + 2288, + 2289, + 2290, + 2291, + 2292, + 2293, + 2294, + 2295, + 2296, + 2297, + 2298, + 2299, + 2300, + 2301, + 2302, + 2303 + ], + [ + 2304, + 2305, + 2306, + 2307, + 2308, + 2309, + 2310, + 2311, + 2312, + 2313, + 2314, + 2315, + 2316, + 2317, + 2318, + 2319, + 2320, + 2321, + 2322, + 2323, + 2324, + 2325, + 2326, + 2327, + 2328, + 2329, + 2330, + 2331, + 2332, + 2333, + 2334, + 2335, + 2336, + 2337, + 2338, + 2339, + 2340, + 2341, + 2342, + 2343, + 2344, + 2345, + 2346, + 2347, + 2348, + 2349, + 2350, + 2351, + 2352, + 2353, + 2354, + 2355, + 2356, + 2357, + 2358, + 2359, + 2360, + 2361, + 2362, + 2363, + 2364, + 2365, + 2366, + 2367 + ], + [ + 2368, + 2369, + 2370, + 2371, + 2372, + 2373, + 2374, + 2375, + 2376, + 2377, + 2378, + 2379, + 2380, + 2381, + 2382, + 2383, + 2384, + 2385, + 2386, + 2387, + 2388, + 2389, + 2390, + 2391, + 2392, + 2393, + 2394, + 2395, + 2396, + 2397, + 2398, + 2399, + 2400, + 2401, + 2402, + 2403, + 2404, + 2405, + 2406, + 2407, + 2408, + 2409, + 2410, + 2411, + 2412, + 2413, + 2414, + 2415, + 2416, + 2417, + 2418, + 2419, + 2420, + 2421, + 2422, + 2423, + 2424, + 2425, + 2426, + 2427, + 2428, + 2429, + 2430, + 2431 + ], + [ + 2432, + 2433, + 2434, + 2435, + 2436, + 2437, + 2438, + 2439, + 2440, + 2441, + 2442, + 2443, + 2444, + 2445, + 2446, + 2447, + 2448, + 2449, + 2450, + 2451, + 2452, + 2453, + 2454, + 2455, + 2456, + 2457, + 2458, + 2459, + 2460, + 2461, + 2462, + 2463, + 2464, + 2465, + 2466, + 2467, + 2468, + 2469, + 2470, + 2471, + 2472, + 2473, + 2474, + 2475, + 2476, + 2477, + 2478, + 2479, + 2480, + 2481, + 2482, + 2483, + 2484, + 2485, + 2486, + 2487, + 2488, + 2489, + 2490, + 2491, + 2492, + 2493, + 2494, + 2495 + ], + [ + 2496, + 2497, + 2498, + 2499, + 2500, + 2501, + 2502, + 2503, + 2504, + 2505, + 2506, + 2507, + 2508, + 2509, + 2510, + 2511, + 2512, + 2513, + 2514, + 2515, + 2516, + 2517, + 2518, + 2519, + 2520, + 2521, + 2522, + 2523, + 2524, + 2525, + 2526, + 2527, + 2528, + 2529, + 2530, + 2531, + 2532, + 2533, + 2534, + 2535, + 2536, + 2537, + 2538, + 2539, + 2540, + 2541, + 2542, + 2543, + 2544, + 2545, + 2546, + 2547, + 2548, + 2549, + 2550, + 2551, + 2552, + 2553, + 2554, + 2555, + 2556, + 2557, + 2558, + 2559 + ], + [ + 2560, + 2561, + 2562, + 2563, + 2564, + 2565, + 2566, + 2567, + 2568, + 2569, + 2570, + 2571, + 2572, + 2573, + 2574, + 2575, + 2576, + 2577, + 2578, + 2579, + 2580, + 2581, + 2582, + 2583, + 2584, + 2585, + 2586, + 2587, + 2588, + 2589, + 2590, + 2591, + 2592, + 2593, + 2594, + 2595, + 2596, + 2597, + 2598, + 2599, + 2600, + 2601, + 2602, + 2603, + 2604, + 2605, + 2606, + 2607, + 2608, + 2609, + 2610, + 2611, + 2612, + 2613, + 2614, + 2615, + 2616, + 2617, + 2618, + 2619, + 2620, + 2621, + 2622, + 2623 + ], + [ + 2624, + 2625, + 2626, + 2627, + 2628, + 2629, + 2630, + 2631, + 2632, + 2633, + 2634, + 2635, + 2636, + 2637, + 2638, + 2639, + 2640, + 2641, + 2642, + 2643, + 2644, + 2645, + 2646, + 2647, + 2648, + 2649, + 2650, + 2651, + 2652, + 2653, + 2654, + 2655, + 2656, + 2657, + 2658, + 2659, + 2660, + 2661, + 2662, + 2663, + 2664, + 2665, + 2666, + 2667, + 2668, + 2669, + 2670, + 2671, + 2672, + 2673, + 2674, + 2675, + 2676, + 2677, + 2678, + 2679, + 2680, + 2681, + 2682, + 2683, + 2684, + 2685, + 2686, + 2687 + ], + [ + 2688, + 2689, + 2690, + 2691, + 2692, + 2693, + 2694, + 2695, + 2696, + 2697, + 2698, + 2699, + 2700, + 2701, + 2702, + 2703, + 2704, + 2705, + 2706, + 2707, + 2708, + 2709, + 2710, + 2711, + 2712, + 2713, + 2714, + 2715, + 2716, + 2717, + 2718, + 2719, + 2720, + 2721, + 2722, + 2723, + 2724, + 2725, + 2726, + 2727, + 2728, + 2729, + 2730, + 2731, + 2732, + 2733, + 2734, + 2735, + 2736, + 2737, + 2738, + 2739, + 2740, + 2741, + 2742, + 2743, + 2744, + 2745, + 2746, + 2747, + 2748, + 2749, + 2750, + 2751 + ], + [ + 2752, + 2753, + 2754, + 2755, + 2756, + 2757, + 2758, + 2759, + 2760, + 2761, + 2762, + 2763, + 2764, + 2765, + 2766, + 2767, + 2768, + 2769, + 2770, + 2771, + 2772, + 2773, + 2774, + 2775, + 2776, + 2777, + 2778, + 2779, + 2780, + 2781, + 2782, + 2783, + 2784, + 2785, + 2786, + 2787, + 2788, + 2789, + 2790, + 2791, + 2792, + 2793, + 2794, + 2795, + 2796, + 2797, + 2798, + 2799, + 2800, + 2801, + 2802, + 2803, + 2804, + 2805, + 2806, + 2807, + 2808, + 2809, + 2810, + 2811, + 2812, + 2813, + 2814, + 2815 + ], + [ + 2816, + 2817, + 2818, + 2819, + 2820, + 2821, + 2822, + 2823, + 2824, + 2825, + 2826, + 2827, + 2828, + 2829, + 2830, + 2831, + 2832, + 2833, + 2834, + 2835, + 2836, + 2837, + 2838, + 2839, + 2840, + 2841, + 2842, + 2843, + 2844, + 2845, + 2846, + 2847, + 2848, + 2849, + 2850, + 2851, + 2852, + 2853, + 2854, + 2855, + 2856, + 2857, + 2858, + 2859, + 2860, + 2861, + 2862, + 2863, + 2864, + 2865, + 2866, + 2867, + 2868, + 2869, + 2870, + 2871, + 2872, + 2873, + 2874, + 2875, + 2876, + 2877, + 2878, + 2879 + ], + [ + 2880, + 2881, + 2882, + 2883, + 2884, + 2885, + 2886, + 2887, + 2888, + 2889, + 2890, + 2891, + 2892, + 2893, + 2894, + 2895, + 2896, + 2897, + 2898, + 2899, + 2900, + 2901, + 2902, + 2903, + 2904, + 2905, + 2906, + 2907, + 2908, + 2909, + 2910, + 2911, + 2912, + 2913, + 2914, + 2915, + 2916, + 2917, + 2918, + 2919, + 2920, + 2921, + 2922, + 2923, + 2924, + 2925, + 2926, + 2927, + 2928, + 2929, + 2930, + 2931, + 2932, + 2933, + 2934, + 2935, + 2936, + 2937, + 2938, + 2939, + 2940, + 2941, + 2942, + 2943 + ], + [ + 2944, + 2945, + 2946, + 2947, + 2948, + 2949, + 2950, + 2951, + 2952, + 2953, + 2954, + 2955, + 2956, + 2957, + 2958, + 2959, + 2960, + 2961, + 2962, + 2963, + 2964, + 2965, + 2966, + 2967, + 2968, + 2969, + 2970, + 2971, + 2972, + 2973, + 2974, + 2975, + 2976, + 2977, + 2978, + 2979, + 2980, + 2981, + 2982, + 2983, + 2984, + 2985, + 2986, + 2987, + 2988, + 2989, + 2990, + 2991, + 2992, + 2993, + 2994, + 2995, + 2996, + 2997, + 2998, + 2999, + 3000, + 3001, + 3002, + 3003, + 3004, + 3005, + 3006, + 3007 + ], + [ + 3008, + 3009, + 3010, + 3011, + 3012, + 3013, + 3014, + 3015, + 3016, + 3017, + 3018, + 3019, + 3020, + 3021, + 3022, + 3023, + 3024, + 3025, + 3026, + 3027, + 3028, + 3029, + 3030, + 3031, + 3032, + 3033, + 3034, + 3035, + 3036, + 3037, + 3038, + 3039, + 3040, + 3041, + 3042, + 3043, + 3044, + 3045, + 3046, + 3047, + 3048, + 3049, + 3050, + 3051, + 3052, + 3053, + 3054, + 3055, + 3056, + 3057, + 3058, + 3059, + 3060, + 3061, + 3062, + 3063, + 3064, + 3065, + 3066, + 3067, + 3068, + 3069, + 3070, + 3071 + ], + [ + 3072, + 3073, + 3074, + 3075, + 3076, + 3077, + 3078, + 3079, + 3080, + 3081, + 3082, + 3083, + 3084, + 3085, + 3086, + 3087, + 3088, + 3089, + 3090, + 3091, + 3092, + 3093, + 3094, + 3095, + 3096, + 3097, + 3098, + 3099, + 3100, + 3101, + 3102, + 3103, + 3104, + 3105, + 3106, + 3107, + 3108, + 3109, + 3110, + 3111, + 3112, + 3113, + 3114, + 3115, + 3116, + 3117, + 3118, + 3119, + 3120, + 3121, + 3122, + 3123, + 3124, + 3125, + 3126, + 3127, + 3128, + 3129, + 3130, + 3131, + 3132, + 3133, + 3134, + 3135 + ], + [ + 3136, + 3137, + 3138, + 3139, + 3140, + 3141, + 3142, + 3143, + 3144, + 3145, + 3146, + 3147, + 3148, + 3149, + 3150, + 3151, + 3152, + 3153, + 3154, + 3155, + 3156, + 3157, + 3158, + 3159, + 3160, + 3161, + 3162, + 3163, + 3164, + 3165, + 3166, + 3167, + 3168, + 3169, + 3170, + 3171, + 3172, + 3173, + 3174, + 3175, + 3176, + 3177, + 3178, + 3179, + 3180, + 3181, + 3182, + 3183, + 3184, + 3185, + 3186, + 3187, + 3188, + 3189, + 3190, + 3191, + 3192, + 3193, + 3194, + 3195, + 3196, + 3197, + 3198, + 3199 + ], + [ + 3200, + 3201, + 3202, + 3203, + 3204, + 3205, + 3206, + 3207, + 3208, + 3209, + 3210, + 3211, + 3212, + 3213, + 3214, + 3215, + 3216, + 3217, + 3218, + 3219, + 3220, + 3221, + 3222, + 3223, + 3224, + 3225, + 3226, + 3227, + 3228, + 3229, + 3230, + 3231, + 3232, + 3233, + 3234, + 3235, + 3236, + 3237, + 3238, + 3239, + 3240, + 3241, + 3242, + 3243, + 3244, + 3245, + 3246, + 3247, + 3248, + 3249, + 3250, + 3251, + 3252, + 3253, + 3254, + 3255, + 3256, + 3257, + 3258, + 3259, + 3260, + 3261, + 3262, + 3263 + ], + [ + 3264, + 3265, + 3266, + 3267, + 3268, + 3269, + 3270, + 3271, + 3272, + 3273, + 3274, + 3275, + 3276, + 3277, + 3278, + 3279, + 3280, + 3281, + 3282, + 3283, + 3284, + 3285, + 3286, + 3287, + 3288, + 3289, + 3290, + 3291, + 3292, + 3293, + 3294, + 3295, + 3296, + 3297, + 3298, + 3299, + 3300, + 3301, + 3302, + 3303, + 3304, + 3305, + 3306, + 3307, + 3308, + 3309, + 3310, + 3311, + 3312, + 3313, + 3314, + 3315, + 3316, + 3317, + 3318, + 3319, + 3320, + 3321, + 3322, + 3323, + 3324, + 3325, + 3326, + 3327 + ], + [ + 3328, + 3329, + 3330, + 3331, + 3332, + 3333, + 3334, + 3335, + 3336, + 3337, + 3338, + 3339, + 3340, + 3341, + 3342, + 3343, + 3344, + 3345, + 3346, + 3347, + 3348, + 3349, + 3350, + 3351, + 3352, + 3353, + 3354, + 3355, + 3356, + 3357, + 3358, + 3359, + 3360, + 3361, + 3362, + 3363, + 3364, + 3365, + 3366, + 3367, + 3368, + 3369, + 3370, + 3371, + 3372, + 3373, + 3374, + 3375, + 3376, + 3377, + 3378, + 3379, + 3380, + 3381, + 3382, + 3383, + 3384, + 3385, + 3386, + 3387, + 3388, + 3389, + 3390, + 3391 + ], + [ + 3392, + 3393, + 3394, + 3395, + 3396, + 3397, + 3398, + 3399, + 3400, + 3401, + 3402, + 3403, + 3404, + 3405, + 3406, + 3407, + 3408, + 3409, + 3410, + 3411, + 3412, + 3413, + 3414, + 3415, + 3416, + 3417, + 3418, + 3419, + 3420, + 3421, + 3422, + 3423, + 3424, + 3425, + 3426, + 3427, + 3428, + 3429, + 3430, + 3431, + 3432, + 3433, + 3434, + 3435, + 3436, + 3437, + 3438, + 3439, + 3440, + 3441, + 3442, + 3443, + 3444, + 3445, + 3446, + 3447, + 3448, + 3449, + 3450, + 3451, + 3452, + 3453, + 3454, + 3455 + ], + [ + 3456, + 3457, + 3458, + 3459, + 3460, + 3461, + 3462, + 3463, + 3464, + 3465, + 3466, + 3467, + 3468, + 3469, + 3470, + 3471, + 3472, + 3473, + 3474, + 3475, + 3476, + 3477, + 3478, + 3479, + 3480, + 3481, + 3482, + 3483, + 3484, + 3485, + 3486, + 3487, + 3488, + 3489, + 3490, + 3491, + 3492, + 3493, + 3494, + 3495, + 3496, + 3497, + 3498, + 3499, + 3500, + 3501, + 3502, + 3503, + 3504, + 3505, + 3506, + 3507, + 3508, + 3509, + 3510, + 3511, + 3512, + 3513, + 3514, + 3515, + 3516, + 3517, + 3518, + 3519 + ], + [ + 3520, + 3521, + 3522, + 3523, + 3524, + 3525, + 3526, + 3527, + 3528, + 3529, + 3530, + 3531, + 3532, + 3533, + 3534, + 3535, + 3536, + 3537, + 3538, + 3539, + 3540, + 3541, + 3542, + 3543, + 3544, + 3545, + 3546, + 3547, + 3548, + 3549, + 3550, + 3551, + 3552, + 3553, + 3554, + 3555, + 3556, + 3557, + 3558, + 3559, + 3560, + 3561, + 3562, + 3563, + 3564, + 3565, + 3566, + 3567, + 3568, + 3569, + 3570, + 3571, + 3572, + 3573, + 3574, + 3575, + 3576, + 3577, + 3578, + 3579, + 3580, + 3581, + 3582, + 3583 + ], + [ + 3584, + 3585, + 3586, + 3587, + 3588, + 3589, + 3590, + 3591, + 3592, + 3593, + 3594, + 3595, + 3596, + 3597, + 3598, + 3599, + 3600, + 3601, + 3602, + 3603, + 3604, + 3605, + 3606, + 3607, + 3608, + 3609, + 3610, + 3611, + 3612, + 3613, + 3614, + 3615, + 3616, + 3617, + 3618, + 3619, + 3620, + 3621, + 3622, + 3623, + 3624, + 3625, + 3626, + 3627, + 3628, + 3629, + 3630, + 3631, + 3632, + 3633, + 3634, + 3635, + 3636, + 3637, + 3638, + 3639, + 3640, + 3641, + 3642, + 3643, + 3644, + 3645, + 3646, + 3647 + ], + [ + 3648, + 3649, + 3650, + 3651, + 3652, + 3653, + 3654, + 3655, + 3656, + 3657, + 3658, + 3659, + 3660, + 3661, + 3662, + 3663, + 3664, + 3665, + 3666, + 3667, + 3668, + 3669, + 3670, + 3671, + 3672, + 3673, + 3674, + 3675, + 3676, + 3677, + 3678, + 3679, + 3680, + 3681, + 3682, + 3683, + 3684, + 3685, + 3686, + 3687, + 3688, + 3689, + 3690, + 3691, + 3692, + 3693, + 3694, + 3695, + 3696, + 3697, + 3698, + 3699, + 3700, + 3701, + 3702, + 3703, + 3704, + 3705, + 3706, + 3707, + 3708, + 3709, + 3710, + 3711 + ], + [ + 3712, + 3713, + 3714, + 3715, + 3716, + 3717, + 3718, + 3719, + 3720, + 3721, + 3722, + 3723, + 3724, + 3725, + 3726, + 3727, + 3728, + 3729, + 3730, + 3731, + 3732, + 3733, + 3734, + 3735, + 3736, + 3737, + 3738, + 3739, + 3740, + 3741, + 3742, + 3743, + 3744, + 3745, + 3746, + 3747, + 3748, + 3749, + 3750, + 3751, + 3752, + 3753, + 3754, + 3755, + 3756, + 3757, + 3758, + 3759, + 3760, + 3761, + 3762, + 3763, + 3764, + 3765, + 3766, + 3767, + 3768, + 3769, + 3770, + 3771, + 3772, + 3773, + 3774, + 3775 + ], + [ + 3776, + 3777, + 3778, + 3779, + 3780, + 3781, + 3782, + 3783, + 3784, + 3785, + 3786, + 3787, + 3788, + 3789, + 3790, + 3791, + 3792, + 3793, + 3794, + 3795, + 3796, + 3797, + 3798, + 3799, + 3800, + 3801, + 3802, + 3803, + 3804, + 3805, + 3806, + 3807, + 3808, + 3809, + 3810, + 3811, + 3812, + 3813, + 3814, + 3815, + 3816, + 3817, + 3818, + 3819, + 3820, + 3821, + 3822, + 3823, + 3824, + 3825, + 3826, + 3827, + 3828, + 3829, + 3830, + 3831, + 3832, + 3833, + 3834, + 3835, + 3836, + 3837, + 3838, + 3839 + ], + [ + 3840, + 3841, + 3842, + 3843, + 3844, + 3845, + 3846, + 3847, + 3848, + 3849, + 3850, + 3851, + 3852, + 3853, + 3854, + 3855, + 3856, + 3857, + 3858, + 3859, + 3860, + 3861, + 3862, + 3863, + 3864, + 3865, + 3866, + 3867, + 3868, + 3869, + 3870, + 3871, + 3872, + 3873, + 3874, + 3875, + 3876, + 3877, + 3878, + 3879, + 3880, + 3881, + 3882, + 3883, + 3884, + 3885, + 3886, + 3887, + 3888, + 3889, + 3890, + 3891, + 3892, + 3893, + 3894, + 3895, + 3896, + 3897, + 3898, + 3899, + 3900, + 3901, + 3902, + 3903 + ], + [ + 3904, + 3905, + 3906, + 3907, + 3908, + 3909, + 3910, + 3911, + 3912, + 3913, + 3914, + 3915, + 3916, + 3917, + 3918, + 3919, + 3920, + 3921, + 3922, + 3923, + 3924, + 3925, + 3926, + 3927, + 3928, + 3929, + 3930, + 3931, + 3932, + 3933, + 3934, + 3935, + 3936, + 3937, + 3938, + 3939, + 3940, + 3941, + 3942, + 3943, + 3944, + 3945, + 3946, + 3947, + 3948, + 3949, + 3950, + 3951, + 3952, + 3953, + 3954, + 3955, + 3956, + 3957, + 3958, + 3959, + 3960, + 3961, + 3962, + 3963, + 3964, + 3965, + 3966, + 3967 + ], + [ + 3968, + 3969, + 3970, + 3971, + 3972, + 3973, + 3974, + 3975, + 3976, + 3977, + 3978, + 3979, + 3980, + 3981, + 3982, + 3983, + 3984, + 3985, + 3986, + 3987, + 3988, + 3989, + 3990, + 3991, + 3992, + 3993, + 3994, + 3995, + 3996, + 3997, + 3998, + 3999, + 4000, + 4001, + 4002, + 4003, + 4004, + 4005, + 4006, + 4007, + 4008, + 4009, + 4010, + 4011, + 4012, + 4013, + 4014, + 4015, + 4016, + 4017, + 4018, + 4019, + 4020, + 4021, + 4022, + 4023, + 4024, + 4025, + 4026, + 4027, + 4028, + 4029, + 4030, + 4031 + ], + [ + 4032, + 4033, + 4034, + 4035, + 4036, + 4037, + 4038, + 4039, + 4040, + 4041, + 4042, + 4043, + 4044, + 4045, + 4046, + 4047, + 4048, + 4049, + 4050, + 4051, + 4052, + 4053, + 4054, + 4055, + 4056, + 4057, + 4058, + 4059, + 4060, + 4061, + 4062, + 4063, + 4064, + 4065, + 4066, + 4067, + 4068, + 4069, + 4070, + 4071, + 4072, + 4073, + 4074, + 4075, + 4076, + 4077, + 4078, + 4079, + 4080, + 4081, + 4082, + 4083, + 4084, + 4085, + 4086, + 4087, + 4088, + 4089, + 4090, + 4091, + 4092, + 4093, + 4094, + 4095 + ] + ] + }, + { + "kind": "tensor", + "dtype": "int32", + "shape": [ + 64 + ], + "device": "cuda:0", + "data": [ + 1024, + 1024, + 1024, + 1024, + 1024, + 1024, + 1024, + 1024, + 1024, + 1024, + 1024, + 1024, + 1024, + 1024, + 1024, + 1024, + 1024, + 1024, + 1024, + 1024, + 1024, + 1024, + 1024, + 1024, + 1024, + 1024, + 1024, + 1024, + 1024, + 1024, + 1024, + 1024, + 1024, + 1024, + 1024, + 1024, + 1024, + 1024, + 1024, + 1024, + 1024, + 1024, + 1024, + 1024, + 1024, + 1024, + 1024, + 1024, + 1024, + 1024, + 1024, + 1024, + 1024, + 1024, + 1024, + 1024, + 1024, + 1024, + 1024, + 1024, + 1024, + 1024, + 1024, + 1024 + ] + }, + { + "kind": "tensor", + "dtype": "int32", + "shape": [ + 65 + ], + "device": "cuda:0", + "data": [ + 0, + 1, + 2, + 3, + 4, + 5, + 6, + 7, + 8, + 9, + 10, + 11, + 12, + 13, + 14, + 15, + 16, + 17, + 18, + 19, + 20, + 21, + 22, + 23, + 24, + 25, + 26, + 27, + 28, + 29, + 30, + 31, + 32, + 33, + 34, + 35, + 36, + 37, + 38, + 39, + 40, + 41, + 42, + 43, + 44, + 45, + 46, + 47, + 48, + 49, + 50, + 51, + 52, + 53, + 54, + 55, + 56, + 57, + 58, + 59, + 60, + 61, + 62, + 63, + 64 + ] + }, + { + "kind": "scalar", + "value": 16, + "py_type": "int" + }, + { + "kind": "scalar", + "value": 1024, + "py_type": "int" + }, + { + "kind": "scalar", + "value": null, + "py_type": "NoneType" + }, + { + "kind": "opaque", + "py_type": "str", + "repr": "'auto'" + }, + { + "kind": "tensor", + "dtype": "float32", + "shape": [], + "device": "cuda:0", + "data": 1.0 + }, + { + "kind": "tensor", + "dtype": "float32", + "shape": [], + "device": "cuda:0", + "data": 1.0 + }, + { + "kind": "scalar", + "value": null, + "py_type": "NoneType" + }, + { + "kind": "opaque", + "py_type": "str", + "repr": "'f16'" + } + ], + "kwargs_sig": {}, + "params_repr": { + "S_seqs": 64, + "ctx_len": 1024, + "out_len": 1, + "heads": 16, + "kv_heads": 1, + "head_size": 128, + "block_size": 16, + "partition": 256, + "kv_alloc_gb": 0.0, + "note": "decode; disjoint KV blocks -> HBM-streaming", + "concurrency_B": 64 + }, + "args_names": [ + "out", + "exp_sums", + "max_logits", + "tmp_out", + "query", + "key_cache", + "value_cache", + "num_kv_heads", + "scale", + "block_tables", + "seq_lens", + "query_start_loc", + "block_size", + "max_seq_len", + "alibi_slopes", + "kv_cache_dtype", + "k_scale", + "v_scale", + "fp8_out_scale", + "mfma_type" + ], + "sources": [ + { + "model": "Qwen_Qwen3-235B-A22B-FP8", + "engine": "vllm", + "mode": "default", + "combo": "Qwen_Qwen3-235B-A22B-FP8__vllm_default", + "count": 376 + } + ], + "perf_only": true + } +] \ No newline at end of file diff --git a/tasks/hip2hip/campaign20/wvsplitk/Makefile b/tasks/hip2hip/campaign20/wvsplitk/Makefile new file mode 100644 index 00000000..0f8b4f36 --- /dev/null +++ b/tasks/hip2hip/campaign20/wvsplitk/Makefile @@ -0,0 +1,6 @@ +# Auto-generated. +all: + python3 scripts/task_runner.py compile + +clean: + rm -rf build diff --git a/tasks/hip2hip/campaign20/wvsplitk/canonical/regime_test_cases.json b/tasks/hip2hip/campaign20/wvsplitk/canonical/regime_test_cases.json new file mode 100644 index 00000000..3f0ee6a5 --- /dev/null +++ b/tasks/hip2hip/campaign20/wvsplitk/canonical/regime_test_cases.json @@ -0,0 +1,167 @@ +[ + { + "test_case_id": "c2", + "args_sig": [ + { + "kind": "tensor", + "dtype": "bfloat16", + "shape": [ + 256, + 7168 + ], + "stride": [ + 7168, + 1 + ], + "device": "cuda:0" + }, + { + "kind": "tensor", + "dtype": "bfloat16", + "shape": [ + 2, + 7168 + ], + "stride": [ + 7168, + 1 + ], + "device": "cuda:0" + }, + { + "kind": "scalar", + "value": null, + "py_type": "NoneType" + }, + { + "kind": "scalar", + "value": 304, + "py_type": "int" + } + ], + "kwargs_sig": {}, + "params_repr": { + "concurrency_B": 2, + "tokens": 2, + "K": 7168, + "weight_N": 256, + "note": "decode skinny GEMM; in_b rows=tokens capped at 4 by kernel" + }, + "args_names": [ + "in_a", + "in_b", + "in_bias", + "CuCount" + ] + }, + { + "test_case_id": "c32", + "args_sig": [ + { + "kind": "tensor", + "dtype": "bfloat16", + "shape": [ + 256, + 7168 + ], + "stride": [ + 7168, + 1 + ], + "device": "cuda:0" + }, + { + "kind": "tensor", + "dtype": "bfloat16", + "shape": [ + 4, + 7168 + ], + "stride": [ + 7168, + 1 + ], + "device": "cuda:0" + }, + { + "kind": "scalar", + "value": null, + "py_type": "NoneType" + }, + { + "kind": "scalar", + "value": 304, + "py_type": "int" + } + ], + "kwargs_sig": {}, + "params_repr": { + "concurrency_B": 32, + "tokens": 4, + "K": 7168, + "weight_N": 256, + "note": "decode skinny GEMM; in_b rows=tokens capped at 4 by kernel" + }, + "args_names": [ + "in_a", + "in_b", + "in_bias", + "CuCount" + ] + }, + { + "test_case_id": "c64", + "args_sig": [ + { + "kind": "tensor", + "dtype": "bfloat16", + "shape": [ + 256, + 7168 + ], + "stride": [ + 7168, + 1 + ], + "device": "cuda:0" + }, + { + "kind": "tensor", + "dtype": "bfloat16", + "shape": [ + 4, + 7168 + ], + "stride": [ + 7168, + 1 + ], + "device": "cuda:0" + }, + { + "kind": "scalar", + "value": null, + "py_type": "NoneType" + }, + { + "kind": "scalar", + "value": 304, + "py_type": "int" + } + ], + "kwargs_sig": {}, + "params_repr": { + "concurrency_B": 64, + "tokens": 4, + "K": 7168, + "weight_N": 256, + "note": "decode skinny GEMM; in_b rows=tokens capped at 4 by kernel" + }, + "args_names": [ + "in_a", + "in_b", + "in_bias", + "CuCount" + ] + } +] \ No newline at end of file diff --git a/tasks/hip2hip/campaign20/wvsplitk/config.yaml b/tasks/hip2hip/campaign20/wvsplitk/config.yaml new file mode 100644 index 00000000..ae8f327f --- /dev/null +++ b/tasks/hip2hip/campaign20/wvsplitk/config.yaml @@ -0,0 +1,35 @@ +source_file_path: +- src/rocm/skinny_gemms.cu +target_kernel_functions: +- wvSplitK_hf_ +- wvSplitK_hf_sml_ +- wvSplitK_hf_big_ +- wvSplitKQ_hf_ +- wvSplitKQ_hf_sml_ +- wvSplitKrc_ +compile_command: +- python3 scripts/task_runner.py compile +correctness_command: +- python3 scripts/task_runner.py correctness +performance_command: +- python3 scripts/task_runner.py performance +task_type: hip2hip +platform_support: + required_arch: gfx950 + status: active + skip_reason: null +task_result_template: null +prompt: + source_code: null + instructions: "Optimize the wvsplitk kernel extracted from vllm. It is registered\ + \ into the `extracted_` torch namespace via src/bindings.cpp; do not depend\ + \ on the host engine. Tests generate inputs from test_cases.json on the fly \u2014\ + \ no pre-saved golden tensors." + cheatsheet: null +extracted_from: + engine: vllm + kernel_id: torch.ops._rocm_C.wvSplitK + op: wvsplitk + local_namespace: extracted_wvsplitk + source: engine_csrc + test_cases: 3 diff --git a/tasks/hip2hip/campaign20/wvsplitk/provenance.json b/tasks/hip2hip/campaign20/wvsplitk/provenance.json new file mode 100644 index 00000000..4716aa02 --- /dev/null +++ b/tasks/hip2hip/campaign20/wvsplitk/provenance.json @@ -0,0 +1,132 @@ +{ + "rank": 26, + "backend": "hip", + "kernel_name": "wvSplitK", + "op_category": "gemm", + "max_time_pct": 2.83, + "n_source_models": 6, + "n_instances": 13, + "canonical": { + "combo": "openai_gpt-oss-120b__vllm_default", + "task_dir": "vllm_wvsplitk" + }, + "instances": [ + { + "model": "openai_gpt-oss-120b", + "engine": "vllm", + "mode": "default", + "combo": "openai_gpt-oss-120b__vllm_default", + "time_pct": 2.83, + "source_dir": "results/openai_gpt-oss-120b__vllm_default/vllm_wvsplitk", + "is_canonical": true + }, + { + "model": "Qwen_Qwen3.5-35B-A3B", + "engine": "vllm", + "mode": "default", + "combo": "Qwen_Qwen3.5-35B-A3B__vllm_default", + "time_pct": 1.49, + "source_dir": "results/Qwen_Qwen3.5-35B-A3B__vllm_default/vllm_wvsplitk", + "is_canonical": false + }, + { + "model": "Qwen_Qwen3-Next-80B-A3B-Instruct-FP8", + "engine": "vllm", + "mode": "default", + "combo": "Qwen_Qwen3-Next-80B-A3B-Instruct-FP8__vllm_default", + "time_pct": 1.31, + "source_dir": "results/Qwen_Qwen3-Next-80B-A3B-Instruct-FP8__vllm_default/vllm_wvsplitk", + "is_canonical": false + }, + { + "model": "deepseek-ai_DeepSeek-R1", + "engine": "vllm", + "mode": "triton_mla", + "combo": "deepseek-ai_DeepSeek-R1__vllm_triton_mla", + "time_pct": 1.18, + "source_dir": "results/deepseek-ai_DeepSeek-R1__vllm_triton_mla/vllm_wvsplitk", + "is_canonical": false + }, + { + "model": "openai_gpt-oss-120b", + "engine": "vllm", + "mode": "hip", + "combo": "openai_gpt-oss-120b__vllm_hip", + "time_pct": 1.09, + "source_dir": "results/openai_gpt-oss-120b__vllm_hip/vllm_wvsplitk", + "is_canonical": false + }, + { + "model": "Qwen_Qwen3-235B-A22B-FP8", + "engine": "vllm", + "mode": "default", + "combo": "Qwen_Qwen3-235B-A22B-FP8__vllm_default", + "time_pct": 1.08, + "source_dir": "results/Qwen_Qwen3-235B-A22B-FP8__vllm_default/vllm_wvsplitk", + "is_canonical": false + }, + { + "model": "Qwen_Qwen3.5-35B-A3B", + "engine": "vllm", + "mode": "hip", + "combo": "Qwen_Qwen3.5-35B-A3B__vllm_hip", + "time_pct": 0.54, + "source_dir": "results/Qwen_Qwen3.5-35B-A3B__vllm_hip/vllm_wvsplitk", + "is_canonical": false + }, + { + "model": "deepseek-ai_DeepSeek-R1", + "engine": "vllm", + "mode": "default", + "combo": "deepseek-ai_DeepSeek-R1__vllm_default", + "time_pct": 0.39, + "source_dir": "results/deepseek-ai_DeepSeek-R1__vllm_default/vllm_wvsplitk", + "is_canonical": false + }, + { + "model": "deepseek-ai_DeepSeek-V3.2", + "engine": "vllm", + "mode": "triton_mla", + "combo": "deepseek-ai_DeepSeek-V3.2__vllm_triton_mla", + "time_pct": 0.23, + "source_dir": "results/deepseek-ai_DeepSeek-V3.2__vllm_triton_mla/vllm_wvsplitk", + "is_canonical": false + }, + { + "model": "deepseek-ai_DeepSeek-V3.2", + "engine": "vllm", + "mode": "hip", + "combo": "deepseek-ai_DeepSeek-V3.2__vllm_hip", + "time_pct": 0.22, + "source_dir": "results/deepseek-ai_DeepSeek-V3.2__vllm_hip/vllm_wvsplitk", + "is_canonical": false + }, + { + "model": "deepseek-ai_DeepSeek-V3.2", + "engine": "vllm", + "mode": "default", + "combo": "deepseek-ai_DeepSeek-V3.2__vllm_default", + "time_pct": 0.19, + "source_dir": "results/deepseek-ai_DeepSeek-V3.2__vllm_default/vllm_wvsplitk", + "is_canonical": false + }, + { + "model": "Qwen_Qwen3-Next-80B-A3B-Instruct-FP8", + "engine": "vllm", + "mode": "hip", + "combo": "Qwen_Qwen3-Next-80B-A3B-Instruct-FP8__vllm_hip", + "time_pct": 0.18, + "source_dir": "results/Qwen_Qwen3-Next-80B-A3B-Instruct-FP8__vllm_hip/vllm_wvsplitk", + "is_canonical": false + }, + { + "model": "deepseek-ai_DeepSeek-R1", + "engine": "vllm", + "mode": "hip", + "combo": "deepseek-ai_DeepSeek-R1__vllm_hip", + "time_pct": 0.04, + "source_dir": "results/deepseek-ai_DeepSeek-R1__vllm_hip/vllm_wvsplitk", + "is_canonical": false + } + ] +} \ No newline at end of file diff --git a/tasks/hip2hip/campaign20/wvsplitk/scripts/_runtime.py b/tasks/hip2hip/campaign20/wvsplitk/scripts/_runtime.py new file mode 100644 index 00000000..79c8a40b --- /dev/null +++ b/tasks/hip2hip/campaign20/wvsplitk/scripts/_runtime.py @@ -0,0 +1,778 @@ +"""Runtime helpers shared by all auto-generated task_runner.py scripts. + +Embedded into each task at build time as ``scripts/_runtime.py`` so the runner +has zero engine dependency at execution time. Two responsibilities: + + 1. ``build_inputs(test_case, seed)``: deterministic random tensor generation + from a captured launch signature (shape + dtype + non-tensor scalars). + 2. ``compare(got, expected, dtype)``: dtype-aware allclose with sensible + tolerances for fp16 / bf16 / fp8. + +A small *reference table* lets known kernels (``rms_norm``, ``silu_and_mul``, +``rotary_embedding``, …) compute an analytic expected value purely in PyTorch. +For unknown kernels the runner falls back to a determinism check (run twice +with the same seed → byte-identical output). +""" +from __future__ import annotations + +import math +from typing import Any, Callable, Dict, List, Optional, Tuple + +import torch + + +# ---------------------------------------------------------------- dtype lookup +_DTYPE_MAP = { + "float32": torch.float32, "float": torch.float32, + "float16": torch.float16, "half": torch.float16, + "bfloat16": torch.bfloat16, + "float8_e4m3fn": torch.float8_e4m3fn, + "float8_e4m3fnuz": getattr(torch, "float8_e4m3fnuz", torch.float8_e4m3fn), + "float8_e5m2": torch.float8_e5m2, + "float8_e5m2fnuz": getattr(torch, "float8_e5m2fnuz", torch.float8_e5m2), + "int64": torch.int64, "long": torch.int64, + "int32": torch.int32, "int": torch.int32, + "int16": torch.int16, "int8": torch.int8, + "uint8": torch.uint8, "uint16": torch.uint16, + "uint32": torch.uint32, "uint64": torch.uint64, + "bool": torch.bool, +} + + +def _dtype(name: str) -> torch.dtype: + s = name.replace("torch.", "").strip() + if s in _DTYPE_MAP: + return _DTYPE_MAP[s] + return torch.float32 + + +# Argument names that need *structured* values for kernels with semantic +# constraints (typically attention / paged-cache kernels). Random ints would +# crash the kernel because, e.g., ``cu_seqlens_q`` must be a non-decreasing +# prefix-sum starting at 0. +_STRUCTURED_KEYS = { + # cumulative-seqlen prefix sums for varlen attention + "cu_seqlens_q", "cu_seqlens_k", "cu_seqlens", "qo_indptr", "kv_indptr", + "reduce_indptr", "num_kv_splits_indptr", "work_indptr", + # block / page bookkeeping + "block_table", "block_tables", "kv_indices", "kv_last_page_lens", + "kv_last_page_len", "page_indices", "page_indptr", + # MoE routing + "topk_ids", "sorted_token_ids", "sorted_expert_ids", "num_valid_ids", + "expert_ids", "topk_indices", "expert_indptr", + # general index/slot tensors + "slot_mapping", "positions", "seq_lens", "context_lens", "query_start_loc", + "query_lens", "cache_indices", "row_starts", "lengths", +} + + +def _make_structured(name: str, sig: dict, gen: torch.Generator, + device: str = "cuda", + ctx: Optional[Dict[str, Any]] = None) -> Optional[torch.Tensor]: + """Generate semantically-valid values for known structured argument names. + + ``ctx`` carries already-resolved scalar args from the same launch (eg. + ``max_seqlen_q`` from kwargs) so cu_seqlens / block_table sizes can be + sized against the kernel's own constraints. + + Returns None if we don't have a recipe for ``name``. + """ + ctx = ctx or {} + shape = tuple(sig.get("shape", [])) + dtype = _dtype(sig.get("dtype", "int32")) + nm = name.lower() + if not shape: + return None + n = shape[0] if len(shape) == 1 else shape[-1] + # Prefix-sum / indptr style: monotonic non-decreasing starting at 0, + # values bounded by max_seqlen_* if recorded in the same launch. + if any(k in nm for k in ("cu_seqlens", "indptr", "query_start_loc", "row_starts")): + if n < 1: + return torch.zeros(shape, dtype=dtype, device=device) + # Pick the increment so that each "sequence" has a bounded length and + # the total fits whatever batched-token tensor was captured for this + # launch. Defaults: try max_seqlen_q / max_seqlen_k from kwargs first. + max_len = None + for k, v in ctx.items(): + kk = k.lower() + if kk in ("max_seqlen_q", "max_seqlen_k", "max_extend_len", + "max_seqlen", "max_context_len") and isinstance(v, int): + if "_q" in nm and "_q" in kk: + max_len = v; break + if "_k" in nm and "_k" in kk: + max_len = v; break + if max_len is None: + max_len = v + chunk = max(1, min(int(max_len), 64) if max_len else 32) + # n includes the leading 0 entry, so we have n-1 sequences. + out = torch.arange(0, n, dtype=dtype, device=device) * chunk + return out.reshape(shape) + # seqlen / context_len: small positive ints, bounded by max_seqlen if known. + if any(k in nm for k in ("seq_lens", "context_lens", "query_lens", "kv_last_page_len", "lengths")): + max_len = None + for k, v in ctx.items(): + kk = k.lower() + if kk.startswith("max_seqlen") and isinstance(v, int): + max_len = v; break + v = min(int(max_len), 64) if max_len else 32 + return torch.full(shape, max(1, v), dtype=dtype, device=device) + # slot_mapping / page_indices / cache_indices: write targets for scattered + # cache writes. They MUST be unique — duplicate slots make the kernel + # non-deterministic across runs (concurrent threads racing the same slot). + # We sample a unique permutation. + if any(k in nm for k in ("slot_mapping", "page_indices", "cache_indices", + "kv_indices")): + total = 1 + for d in shape: + total *= d + # Pool size: at least 4× the number of slots so the permutation has + # room. Cap at 32k to keep allocation cheap. + pool = max(total * 4, 256) + pool = min(pool, 32768) + perm = torch.randperm(pool, generator=gen, dtype=dtype, device=device)[:total] + return perm.reshape(shape) + if any(k in nm for k in ("block_table", "positions")): + high = 1024 + return torch.randint(0, high, shape, dtype=dtype, device=device, generator=gen) + # MoE routing — topk_ids/sorted_expert_ids must be valid expert indices. + # We don't know num_experts here; default to 8 which is common. + if "expert" in nm or nm == "topk_ids" or nm == "topk_indices": + return torch.randint(0, 8, shape, dtype=dtype, device=device, generator=gen) + if nm == "num_valid_ids": + return torch.tensor([min(shape[0] if shape else 1, 64)] * (shape[0] if shape else 1), + dtype=dtype, device=device) + if "sorted_token_ids" in nm: + return torch.zeros(shape, dtype=dtype, device=device) + return None + + +# ---------------------------------------------------------------- tensor gen +def _make_tensor(sig: dict, gen: torch.Generator, device: str = "cuda", + name: str = "", + ctx: Optional[Dict[str, Any]] = None) -> torch.Tensor: + if name and name.lower() in _STRUCTURED_KEYS: + s = _make_structured(name, sig, gen, device, ctx=ctx) + if s is not None: + return s + shape = tuple(sig.get("shape", [])) + dtype = _dtype(sig.get("dtype", "float32")) + if dtype.is_floating_point: + if dtype in ( + torch.float8_e4m3fn, + torch.float8_e5m2, + getattr(torch, "float8_e4m3fnuz", torch.float8_e4m3fn), + getattr(torch, "float8_e5m2fnuz", torch.float8_e5m2), + ): + base = (torch.randn(shape, dtype=torch.float32, device=device, generator=gen) * 0.1) + return base.clamp(-1.0, 1.0).to(dtype) + return torch.randn(shape, dtype=torch.float32, device=device, generator=gen).to(dtype) + if dtype == torch.bool: + return (torch.randint(0, 2, shape, device=device, generator=gen, dtype=torch.int32) > 0) + # integer tensor — keep values modest so they're plausible indices etc. + if not shape: + return torch.tensor(0, dtype=dtype, device=device) + high = max(2, min(1024, shape[-1] if shape[-1] > 1 else 32)) + # Clamp to dtype range so eg. uint8 doesn't overflow torch.randint's bound + # check (``to - 1 is out of bounds for unsigned char``). + if dtype == torch.uint8: + high = min(high, 256) + elif dtype == torch.int8: + high = min(high, 128) + elif dtype == torch.int16: + high = min(high, 32768) + elif dtype == torch.uint16: + high = min(high, 65536) + return torch.randint(0, high, shape, dtype=dtype, device=device, generator=gen) + + +def _decode_opaque(sig: dict) -> Any: + py_type = (sig.get("py_type") or "").lower() + rep = sig.get("repr", "") + if py_type in ("str", "int", "float", "bool", "nonetype"): + try: + import ast as _ast + return _ast.literal_eval(rep) + except Exception: + return None + if py_type == "dtype" and rep.startswith("torch."): + return _dtype(rep[6:]) + # Captured aiter enums look like ````. Extracted tasks + # call pybind functions directly, and those bindings accept the underlying + # integer rather than the Python enum object from an installed aiter package. + if "." in rep and rep.startswith("<") and ":" in rep: + try: + import re as _re + m = _re.match(r"<([\w.]+)\.(\w+):\s*(-?\d+)>", rep) + if m: + return int(m.group(3)) + except Exception: + pass + return None + + +def _arg_from_sig(sig: dict, gen: torch.Generator, name: str = "", + ctx: Optional[Dict[str, Any]] = None) -> Any: + kind = sig.get("kind", "scalar") + if kind == "tensor": + return _make_tensor(sig, gen, name=name, ctx=ctx) + if kind == "scalar": + return sig.get("value") + if kind == "seq": + return [_arg_from_sig(s, gen, ctx=ctx) for s in sig.get("items", [])] + if kind == "map": + return {k: _arg_from_sig(v, gen, name=k, ctx=ctx) for k, v in sig.get("items", {}).items()} + if kind == "opaque": + return _decode_opaque(sig) + return None + + +def _resolve_scalar_ctx(test_case: dict) -> Dict[str, Any]: + """First-pass scan: pull any scalar / opaque-scalar values out of the + launch's args+kwargs so structured tensor generators (cu_seqlens, etc.) + can size themselves against ``max_seqlen_q`` and friends. + """ + ctx: Dict[str, Any] = {} + for entry in (test_case.get("kwargs_sig") or {}).items(): + k, v = entry + if not isinstance(v, dict): + continue + if v.get("kind") == "scalar": + val = v.get("value") + if isinstance(val, (int, float, bool, str)): + ctx[k] = val + elif v.get("kind") == "opaque": + dec = _decode_opaque(v) + if isinstance(dec, (int, float, bool, str)): + ctx[k] = dec + return ctx + + +def build_inputs(test_case: dict, seed: int = 0xC0FFEE) -> Tuple[List[Any], Dict[str, Any]]: + """Materialize positional + keyword args from a captured launch signature. + + A first pass extracts scalar kwargs (``max_seqlen_q``, etc.) into a context + dict that the tensor builder consults — so eg. ``cu_seqlens_k`` is sized + so that its max value matches the captured ``max_seqlen_k``. + + If the test_case carries ``args_names`` (parsed from the op schema), they + are forwarded to the per-position tensor builder so structured generators + (slot_mapping → unique perm, cu_seqlens_q → prefix sum) fire even when + the kernel is called positionally. + """ + if not torch.cuda.is_available(): + raise RuntimeError("CUDA / ROCm device required to build inputs") + gen = torch.Generator(device="cuda").manual_seed(seed) + ctx = _resolve_scalar_ctx(test_case) + args_sig = test_case.get("args_sig", []) + args_names = test_case.get("args_names") or [""] * len(args_sig) + if len(args_names) < len(args_sig): + args_names = list(args_names) + [""] * (len(args_sig) - len(args_names)) + args = [ + _arg_from_sig(s, gen, name=args_names[i], ctx=ctx) + for i, s in enumerate(args_sig) + ] + # CuCount was captured on MI300X (304 CUs). Use the live device value so + # the same self-contained task launches the intended grid on MI355X and + # other supported AMD GPUs. + for i, name in enumerate(args_names): + if name.lower() == "cucount": + args[i] = torch.cuda.get_device_properties(0).multi_processor_count + kwargs: Dict[str, Any] = {} + for k, v in (test_case.get("kwargs_sig") or {}).items(): + kwargs[k] = _arg_from_sig(v, gen, name=k, ctx=ctx) + return args, kwargs + + +# ---------------------------------------------------------------- aiter wrapper -> pybind normalization +def _dtype_to_aiter_string(value: Any) -> Any: + if value is None: + return None + if value is torch.float16: + return "fp16" + if value is torch.bfloat16: + return "bf16" + if value is torch.float32: + return "fp32" + if isinstance(value, str): + return value + return None + + +def _sanitize_moe_routing(args: List[Any], kwargs: Dict[str, Any]) -> None: + if len(args) < 8: + return + hidden, w1, w2 = args[0], args[1], args[2] + sorted_token_ids, sorted_expert_ids, num_valid_ids = args[3], args[4], args[5] + if not all(isinstance(x, torch.Tensor) for x in (hidden, w1, sorted_token_ids, sorted_expert_ids, num_valid_ids)): + return + try: + topk = int(args[7]) + except Exception: + topk = int(kwargs.get("topk", 1) or 1) + tokens = int(hidden.shape[0]) + experts = int(w1.shape[0]) + valid = max(1, tokens * max(1, topk)) + with torch.no_grad(): + ids = torch.arange(sorted_token_ids.numel(), device=sorted_token_ids.device, + dtype=sorted_token_ids.dtype) % valid + sorted_token_ids.copy_(ids.reshape_as(sorted_token_ids)) + eids = torch.arange(sorted_expert_ids.numel(), device=sorted_expert_ids.device, + dtype=sorted_expert_ids.dtype) % max(1, experts) + sorted_expert_ids.copy_(eids.reshape_as(sorted_expert_ids)) + num_valid_ids.fill_(min(sorted_token_ids.numel(), valid)) + # Per instruction_WIP.md kernel_moe_gemm: scale inputs by 1/sqrt(K) so + # bf16 doesn't overflow with random non-pre-shuffled weights. The + # underlying GEMM is K-reduction; without scaling the accumulator hits + # ~3e38 and atomic-add ordering produces non-deterministic overflow. + if hidden.dtype.is_floating_point: + K = int(hidden.shape[-1]) + scale = 1.0 / max(1.0, math.sqrt(float(K))) + hidden.mul_(scale) + if isinstance(w1, torch.Tensor) and w1.dtype.is_floating_point: + w1.mul_(scale) + if isinstance(w2, torch.Tensor) and w2.dtype.is_floating_point: + w2.mul_(scale) + + +def normalize_aiter_call(py_fn_name: str, fc_name: str, + args: List[Any], kwargs: Dict[str, Any]) -> Tuple[List[Any], Dict[str, Any]]: + """Translate captured Python-wrapper arguments to the direct pybind ABI. + + The recorder observes public aiter wrappers such as ``flash_attn_varlen_func`` + and ``ck_moe_stage1_fwd``. Extracted tasks intentionally call the local + pybind module directly, so a few wrapper-only keyword names and defaults need + to be normalized before invocation. + """ + op = py_fn_name or fc_name + args = list(args) + kwargs = dict(kwargs) + + if op in ("ck_moe_stage1_fwd", "ck_moe_stage2_fwd", "ck_moe_stage1", "ck_moe_stage2"): + if "use_non_temporal_load" in kwargs and "non_temporal_load" not in kwargs: + kwargs["non_temporal_load"] = kwargs.pop("use_non_temporal_load") + elif "use_non_temporal_load" in kwargs: + kwargs.pop("use_non_temporal_load", None) + # Wrapper-only kwargs absent from the underlying pybind ABI. + if "dtype" in kwargs and "dst_type" not in kwargs: + kwargs["dst_type"] = _dtype_to_aiter_string(kwargs.pop("dtype")) + else: + kwargs.pop("dtype", None) + if len(args) > 8: + kwargs.pop("kernelName", None) + elif kwargs.get("kernelName") is None: + kwargs["kernelName"] = "" + if "dst_type" in kwargs: + kwargs["dst_type"] = _dtype_to_aiter_string(kwargs["dst_type"]) + # Positional captures include dst_type just before is_shuffled. + if len(args) >= 18: + args[17] = _dtype_to_aiter_string(args[17]) + _sanitize_moe_routing(args, kwargs) + return args, kwargs + + if op == "flash_attn_varlen_func" or fc_name in ("fmha_v3_varlen_fwd", "mha_varlen_fwd"): + # The captured wrapper sometimes passes (q, k, v) positionally and the + # rest as a sparse kwargs dict; the underlying pybind requires every + # non-trailing parameter. Promote positional q/k/v back into kwargs + # then reissue everything by name with explicit defaults. + if args: + kwargs.setdefault("q", args[0] if len(args) > 0 else None) + kwargs.setdefault("k", args[1] if len(args) > 1 else None) + kwargs.setdefault("v", args[2] if len(args) > 2 else None) + args = [] + window = kwargs.pop("window_size", [-1, -1]) + if window is None: + window = [-1, -1] + window_left = int(window[0]) if len(window) > 0 else -1 + window_right = int(window[1]) if len(window) > 1 else -1 + # The captured max_seqlen_q/k may be the model's max-sequence cap + # (e.g. 65536) rather than the actual maximum present in the + # synthesized cu_seqlens_q. Recompute from cu_seqlens since the + # synthesized inputs may be much smaller — passing an oversized + # max_seqlen makes the CK kernel index past the actual tensor. + cu_q = kwargs.get("cu_seqlens_q") + cu_k = kwargs.get("cu_seqlens_k") + captured_max_q = int(kwargs.get("max_seqlen_q", 0) or 0) + captured_max_k = int(kwargs.get("max_seqlen_k", 0) or 0) + actual_max_q = captured_max_q + actual_max_k = captured_max_k + try: + if cu_q is not None and hasattr(cu_q, "diff"): + actual_max_q = int(cu_q.diff().max().item()) + if cu_k is not None and hasattr(cu_k, "diff"): + actual_max_k = int(cu_k.diff().max().item()) + except Exception: + pass + max_q = max(1, min(captured_max_q, actual_max_q) if captured_max_q else actual_max_q) + max_k = max(1, min(captured_max_k, actual_max_k) if captured_max_k else actual_max_k) + kwargs = { + "q": kwargs.get("q"), + "k": kwargs.get("k"), + "v": kwargs.get("v"), + "cu_seqlens_q": cu_q, + "cu_seqlens_k": cu_k, + "max_seqlen_q": max_q, + "max_seqlen_k": max_k, + "min_seqlen_q": int(kwargs.get("min_seqlen_q", 0)), + "dropout_p": float(kwargs.get("dropout_p", 0.0)), + "softmax_scale": float(kwargs.get("softmax_scale", 1.0)), + "logits_soft_cap": float(kwargs.get("logits_soft_cap", 0.0)), + "zero_tensors": bool(kwargs.get("zero_tensors", False)), + "is_causal": bool(kwargs.get("is_causal", kwargs.get("causal", False))), + "window_size_left": window_left, + "window_size_right": window_right, + "return_softmax_lse": bool(kwargs.get("return_softmax_lse", + kwargs.get("return_lse", False))), + "return_dropout_randval": bool(kwargs.get("return_dropout_randval", False)), + "how_v3_bf16_cvt": int(kwargs.get("how_v3_bf16_cvt", 1)), + "out": kwargs.get("out"), + "block_table": kwargs.get("block_table"), + "bias": kwargs.get("bias"), + "alibi_slopes": kwargs.get("alibi_slopes"), + "gen": kwargs.get("gen"), + "cu_seqlens_q_padded": kwargs.get("cu_seqlens_q_padded"), + "cu_seqlens_k_padded": kwargs.get("cu_seqlens_k_padded"), + } + return args, kwargs + + return args, kwargs + + +def normalize_aiter_output(py_fn_name: str, value: Any) -> Any: + if py_fn_name == "flash_attn_varlen_func" and isinstance(value, (list, tuple)): + # AITER returns (out, softmax_lse, dropout_mask, rng_state). With + # dropout disabled the rng_state buffer is not semantically meaningful + # and may contain run-to-run garbage; compare the observable outputs. + return tuple(value[:2]) + return value + + +# ---------------------------------------------------------------- comparison +def _tol_for(dtype: torch.dtype) -> Tuple[float, float]: + if dtype in (torch.float16, torch.bfloat16): + return 5e-2, 5e-2 + if dtype in (torch.float8_e4m3fn, torch.float8_e5m2, + getattr(torch, "float8_e4m3fnuz", torch.float8_e4m3fn), + getattr(torch, "float8_e5m2fnuz", torch.float8_e5m2)): + return 0.5, 0.5 + if dtype.is_floating_point: + return 1e-3, 1e-3 + return 0, 0 + + +def compare(got: Any, expected: Any) -> Optional[str]: + """Return None if equal-within-tolerance, else a short diff message.""" + if isinstance(got, torch.Tensor) and isinstance(expected, torch.Tensor): + if got.shape != expected.shape: + return f"shape mismatch: got {tuple(got.shape)} vs {tuple(expected.shape)}" + atol, rtol = _tol_for(got.dtype) + a = got.detach().to(torch.float32).cpu() + b = expected.detach().to(torch.float32).cpu() + if not torch.isfinite(a).all(): + count = int((~torch.isfinite(a)).sum().item()) + return f"kernel output contains {count} non-finite value(s)" + if not torch.isfinite(b).all(): + count = int((~torch.isfinite(b)).sum().item()) + return f"reference output contains {count} non-finite value(s)" + if not torch.allclose(a, b, atol=atol, rtol=rtol): + diff = (a - b).abs().max().item() + return f"max abs diff {diff:.4g} > atol={atol}" + return None + if isinstance(got, (list, tuple)) and isinstance(expected, (list, tuple)): + if len(got) != len(expected): + return f"len mismatch: {len(got)} vs {len(expected)}" + for i, (g, e) in enumerate(zip(got, expected)): + err = compare(g, e) + if err: + return f"[{i}]: {err}" + return None + if got == expected: + return None + return f"value mismatch" + + +# ---------------------------------------------------------------- references +# Each reference takes the SAME (args, kwargs) as the kernel and returns the +# expected output. For in-place kernels the reference returns a tensor that +# the runner will compare against the (now-mutated) input. The runner picks +# the entry by op_name; if missing, falls back to determinism check. +ReferenceFn = Callable[[List[Any], Dict[str, Any]], Any] + + +def _ref_rms_norm(args: list, kwargs: dict): + # signature: (out, input, weight, epsilon) — vLLM `_C.rms_norm` writes to out + out, inp, weight, eps = args[0], args[1], args[2], args[3] + var = inp.to(torch.float32).pow(2).mean(dim=-1, keepdim=True) + return (inp.to(torch.float32) * torch.rsqrt(var + eps)).to(inp.dtype) * weight.to(inp.dtype) + + +def _ref_fused_add_rms_norm(args: list, kwargs: dict): + # (input, residual, weight, epsilon) — in-place: input = rmsnorm(input+residual) + inp, res, weight, eps = args[0], args[1], args[2], args[3] + s = (inp + res).to(torch.float32) + var = s.pow(2).mean(dim=-1, keepdim=True) + return (s * torch.rsqrt(var + eps)).to(inp.dtype) * weight.to(inp.dtype) + + +def _ref_silu_and_mul(args: list, kwargs: dict): + # (out, input) — out = silu(input[..., :H/2]) * input[..., H/2:] + out, inp = args[0], args[1] + a, b = inp.chunk(2, dim=-1) + return torch.nn.functional.silu(a.to(torch.float32)).to(inp.dtype) * b + + +def _ref_gelu_and_mul(args: list, kwargs: dict): + out, inp = args[0], args[1] + a, b = inp.chunk(2, dim=-1) + return torch.nn.functional.gelu(a.to(torch.float32)).to(inp.dtype) * b + + +def _ref_gelu_tanh_and_mul(args: list, kwargs: dict): + out, inp = args[0], args[1] + a, b = inp.chunk(2, dim=-1) + return torch.nn.functional.gelu(a.to(torch.float32), approximate="tanh").to(inp.dtype) * b + + +def _ref_gelu_quick(args: list, kwargs: dict): + out, inp = args[0], args[1] + f = inp.to(torch.float32) + return (f * torch.sigmoid(1.702 * f)).to(inp.dtype) + + +def _ref_gelu_new(args: list, kwargs: dict): + out, inp = args[0], args[1] + f = inp.to(torch.float32) + return (0.5 * f * (1 + torch.tanh(math.sqrt(2 / math.pi) * (f + 0.044715 * f.pow(3))))).to(inp.dtype) + + +def _ref_gelu_fast(args: list, kwargs: dict): + return _ref_gelu_new(args, kwargs) + + +def _ref_mul_and_silu(args: list, kwargs: dict): + out, inp = args[0], args[1] + a, b = inp.chunk(2, dim=-1) + return (a.to(torch.float32) * torch.nn.functional.silu(b.to(torch.float32))).to(inp.dtype) + + +def _ref_fatrelu_and_mul(args: list, kwargs: dict): + out, inp = args[0], args[1] + threshold = float(kwargs.get("threshold", args[2] if len(args) > 2 else 0.0)) + a, b = inp.chunk(2, dim=-1) + mask = (a.to(torch.float32) > threshold).to(inp.dtype) + return (a * mask) * b + + +# References are keyed by ``":"`` because vLLM and AITER +# expose ops with the same short name but different positional layouts (e.g. +# vLLM ``fused_add_rms_norm(input, residual, weight, eps)`` vs. AITER +# ``rmsnorm2d_fwd_with_add(out, input, residual, out_residual, weight, eps)``). +# The runner passes its source prefix; unknown keys fall back to a determinism +# check, which is a safe no-op rather than producing NaN garbage. +# [BugA-fix] verified numerical references (added by benchmark maintainers) +def _ref_moe_sum(args, kwargs): + inp, out = args[0], args[1] + return {1: inp.to(torch.float32).sum(dim=1).to(out.dtype)} + + +def _ref__apply_rotary(x, cos, sin, rot_dim, is_neox): + xf = x.clone() + embed = rot_dim // 2 + c = cos[:, None, :] + s = sin[:, None, :] + if is_neox: + x1 = xf[..., :embed]; x2 = xf[..., embed:rot_dim] + xf[..., :embed] = x1 * c - x2 * s + xf[..., embed:rot_dim] = x2 * c + x1 * s + else: + x1 = xf[..., 0:rot_dim:2]; x2 = xf[..., 1:rot_dim:2] + xf[..., 0:rot_dim:2] = x1 * c - x2 * s + xf[..., 1:rot_dim:2] = x2 * c + x1 * s + return xf + + +def _ref_rotary_embedding(args, kwargs): + positions, query, key, head_size, cos_sin_cache, is_neox = args[0], args[1], args[2], args[3], args[4], args[5] + rot_dim = cos_sin_cache.shape[-1] + embed = rot_dim // 2 + pos = positions.reshape(-1).long() + cs = cos_sin_cache.index_select(0, pos).to(torch.float32) + cos, sin = cs[:, :embed], cs[:, embed:rot_dim] + T = pos.shape[0] + out = {1: _ref__apply_rotary(query.to(torch.float32).view(T, -1, head_size), cos, sin, rot_dim, is_neox).view_as(query).to(query.dtype)} + if isinstance(key, torch.Tensor): + out[2] = _ref__apply_rotary(key.to(torch.float32).view(T, -1, head_size), cos, sin, rot_dim, is_neox).view_as(key).to(key.dtype) + return out + + +def _ref_reshape_and_cache(args, kwargs): + key, value, key_cache, value_cache, slot_mapping = args[0], args[1], args[2], args[3], args[4] + kc, vc = key_cache.clone(), value_cache.clone() + T, nh, hs = key.shape + nb, nh2, hsx, bs, x = key_cache.shape + slot = slot_mapping.reshape(-1).long() + for t in range(T): + s = int(slot[t]) + if s < 0: + continue + b, off = s // bs, s % bs + kc[b, :, :, off, :] = key[t].to(kc.dtype).view(nh, hsx, x) + vc[b, :, :, off] = value[t].to(vc.dtype) + return {2: kc, 3: vc} + + +def _ref_reshape_and_cache_flash(args, kwargs): + key, value, key_cache, value_cache, slot_mapping = args[0], args[1], args[2], args[3], args[4] + kc, vc = key_cache.clone(), value_cache.clone() + T = key.shape[0] + nb, bs, nh2, hs2 = key_cache.shape + slot = slot_mapping.reshape(-1).long() + for t in range(T): + s = int(slot[t]) + if s < 0: + continue + b, off = s // bs, s % bs + kc[b, off, :, :] = key[t].to(kc.dtype) + vc[b, off, :, :] = value[t].to(vc.dtype) + return {2: kc, 3: vc} + + +def _ref_concat_and_cache_mla(args, kwargs): + kv_c, k_pe, kv_cache, slot_mapping = args[0], args[1], args[2], args[3] + c = kv_cache.clone() + T, lora = kv_c.shape + pe = k_pe.shape[1] + nb, bs, entry = kv_cache.shape + slot = slot_mapping.reshape(-1).long() + for t in range(T): + s = int(slot[t]) + if s < 0: + continue + b, off = s // bs, s % bs + c[b, off, :lora] = kv_c[t].to(c.dtype) + c[b, off, lora:lora + pe] = k_pe[t].to(c.dtype) + return {2: c} + + +def _ref_topk_softmax(args, kwargs): + topk_weights, gating = args[0], args[3] + renorm = bool(args[4]) if len(args) > 4 else False + topk = topk_weights.shape[1] + probs = torch.softmax(gating.to(torch.float32), dim=-1) + vals, _ = probs.topk(topk, dim=-1) + if renorm: + vals = vals / vals.sum(dim=-1, keepdim=True) + # Only the top-k weights are numerically meaningful; expert-index ordering / + # tie-breaking is implementation-defined and intentionally not compared. + return {0: vals.to(topk_weights.dtype)} + + +def _ref_wvSplitK(args, kwargs): + in_a, in_b = args[0], args[1] + bias = args[2] if len(args) > 2 else None + out = in_b.to(torch.float32) @ in_a.to(torch.float32).T + if isinstance(bias, torch.Tensor): + out = out + bias.to(torch.float32) + return out.to(in_b.dtype) + + +def _ref_paged_attention(args, kwargs): + out, query, key_cache, value_cache = args[0], args[4], args[5], args[6] + num_kv_heads, scale = int(args[7]), float(args[8]) + block_tables, seq_lens, query_start_loc, block_size = args[9], args[10], args[11], int(args[12]) + num_seqs = block_tables.shape[0] + num_heads, head_size = query.shape[1], query.shape[2] + gqa = num_heads // num_kv_heads + nb, nkv, hsx, bs, x = key_cache.shape + kc, vc, q = key_cache.float(), value_cache.float(), query.float() + bt, sl = block_tables.long(), seq_lens.long() + qsl = query_start_loc.long() if isinstance(query_start_loc, torch.Tensor) else None + out_ref = out.float().clone() + kv_of_head = torch.arange(num_heads, device=q.device) // gqa + for si in range(num_seqs): + row = int(qsl[si]) if qsl is not None else si + L = int(sl[si]) + if L <= 0: + continue + pos = torch.arange(L, device=q.device) + blk = bt[si, pos // block_size] + off = pos % block_size + ar = torch.arange(L, device=q.device) + K = kc[blk][ar, :, :, off, :].reshape(L, nkv, hsx * x) + V = vc[blk][ar, :, :, off] + Kh, Vh = K[:, kv_of_head, :], V[:, kv_of_head, :] + scores = scale * torch.einsum('hd,phd->hp', q[row], Kh) + attn = torch.softmax(scores, dim=-1) + out_ref[row] = torch.einsum('hp,phd->hd', attn, Vh) + return {0: out_ref.to(out.dtype)} + + +REFERENCES: Dict[str, ReferenceFn] = { + # [BugA-fix] verified numerical references + "vllm:moe_sum": _ref_moe_sum, + "vllm:rotary_embedding": _ref_rotary_embedding, + "vllm:reshape_and_cache": _ref_reshape_and_cache, + "vllm:reshape_and_cache_flash": _ref_reshape_and_cache_flash, + "vllm:concat_and_cache_mla": _ref_concat_and_cache_mla, + "vllm:topk_softmax": _ref_topk_softmax, + "vllm:wvSplitK": _ref_wvSplitK, + "vllm:paged_attention": _ref_paged_attention, + + # vLLM _C ops — runner passes ``vllm:`` + "vllm:rms_norm": _ref_rms_norm, + "vllm:fused_add_rms_norm": _ref_fused_add_rms_norm, + "vllm:silu_and_mul": _ref_silu_and_mul, + "vllm:gelu_and_mul": _ref_gelu_and_mul, + "vllm:gelu_tanh_and_mul": _ref_gelu_tanh_and_mul, + "vllm:gelu_quick": _ref_gelu_quick, + "vllm:gelu_new": _ref_gelu_new, + "vllm:gelu_fast": _ref_gelu_fast, + "vllm:mul_and_silu": _ref_mul_and_silu, + "vllm:fatrelu_and_mul": _ref_fatrelu_and_mul, + # SGLang sgl_kernel ops mostly mirror vLLM's signatures. + "sglang:rms_norm": _ref_rms_norm, + "sglang:fused_add_rms_norm": _ref_fused_add_rms_norm, + "sglang:silu_and_mul": _ref_silu_and_mul, + "sglang:gelu_and_mul": _ref_gelu_and_mul, + "sglang:gelu_tanh_and_mul": _ref_gelu_tanh_and_mul, + "sglang:gelu_quick": _ref_gelu_quick, +} + + +def reference_for(op_name: str, source: str = "") -> Optional[ReferenceFn]: + """Pick a reference for ``op_name``. If ``source`` is given (``vllm`` / + ``aiter`` / ``triton`` / ``sglang``) it is used as a namespace prefix to + disambiguate same-named ops with different signatures.""" + if source: + fn = REFERENCES.get(f"{source}:{op_name}") + if fn is not None: + return fn + return REFERENCES.get(op_name) + + +# ---------------------------------------------------------------- output capture +# Many vLLM/AITER kernels mutate their first ``out`` arg rather than returning +# a value. The runner inspects the first positional tensor argument's bytes +# before/after the call to detect this and treat it as the output. + +def snapshot(args: list) -> list: + return [a.detach().clone() if isinstance(a, torch.Tensor) else None for a in args] + + +def detect_output(pre: list, post: list, ret: Any) -> Any: + """Pick the most plausible output for comparison. + + 1. If the kernel returned a tensor (or tuple), use that. + 2. Otherwise look for the first positional tensor that changed in-place. + 3. Else None. + """ + if isinstance(ret, torch.Tensor) or isinstance(ret, (list, tuple)) and ret and isinstance(ret[0], torch.Tensor): + return ret + for i, (b, a) in enumerate(zip(pre, post)): + if a is None or b is None: + continue + try: + if not torch.equal(b.to(torch.float32).cpu(), a.to(torch.float32).cpu()): + return a + except Exception: + continue + return None diff --git a/tasks/hip2hip/campaign20/wvsplitk/scripts/task_runner.py b/tasks/hip2hip/campaign20/wvsplitk/scripts/task_runner.py new file mode 100644 index 00000000..f0aff42e --- /dev/null +++ b/tasks/hip2hip/campaign20/wvsplitk/scripts/task_runner.py @@ -0,0 +1,287 @@ +#!/usr/bin/env python3 +"""Auto-generated task runner for vllm_wvsplitk (HIP). + +Inputs are generated each run from the shape/dtype signatures in +test_cases.json — no .pt files are loaded. +""" +import sys, os, json, argparse, glob +sys.path.insert(0, os.path.dirname(os.path.abspath(__file__))) +import _runtime as rt + +# >>> AKA-GENERATED: shared CUDA-graph benchmark helpers - edit src/tools/perf/vllm_cuda_graph_block.py then run `make sync-perf-helpers` >>> +def _measure_cuda_event_fallback(*args, **kwargs): + raise RuntimeError( + "CUDA-graph benchmark helpers were not materialized. " + "Run this task through AgentKernelArena so setup_workspace() can inject " + "src/tools/perf/vllm_cuda_graph_block.py into the workspace." + ) + + +def _benchmark_cuda_graph_or_events(*args, **kwargs): + raise RuntimeError( + "CUDA-graph benchmark helpers were not materialized. " + "Run this task through AgentKernelArena so setup_workspace() can inject " + "src/tools/perf/vllm_cuda_graph_block.py into the workspace." + ) +# <<< AKA-GENERATED <<< + + +TASK_DIR = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) +os.chdir(TASK_DIR) +TASK_NAME = "hip2hip/campaign20/wvsplitk" +NAMESPACE = "extracted_wvsplitk" # extracted_ +OP_NAME = "wvSplitK" +REF_SOURCE = "vllm" # "vllm" | "sglang" — used to pick the correct reference + +SRC_DIR = os.path.join(TASK_DIR, "src") +BUILD_DIR = os.path.join(TASK_DIR, "build") +# Prefer the workload-regime cases (seqlen=1024, concurrency B in {2,32,64}) +# when present; otherwise fall back to the originally captured launch sigs. +# wvSplitK is a *decode* skinny GEMM: the token dimension (in_b.size(0)) is +# hard-capped at 4 by the kernel's N_in switch, so concurrency B is mapped to +# token count = min(B, 4) -> c2:2 tokens, c32/c64:4 tokens. Model dims (K, the +# weight output rows) are kept from the captured DeepSeek base case. +_REGIME_CASES = os.path.join(TASK_DIR, "canonical", "regime_test_cases.json") +TEST_CASES = _REGIME_CASES if os.path.isfile(_REGIME_CASES) else os.path.join(TASK_DIR, "test_cases.json") + +# Lock the offload arch to the runtime device's gfx so torch.utils.cpp_extension +# does not try to compile for every ROCm target (RDNA gfx1100 etc. break on +# vllm's cub/bf16 templates). Override by setting PYTORCH_ROCM_ARCH externally. +def _detect_gfx(): + try: + import torch + if torch.cuda.is_available(): + arch = torch.cuda.get_device_properties(0).gcnArchName + return arch.split(":")[0] # e.g. "gfx942:sramecc+:xnack-" -> "gfx942" + except Exception: + pass + return "gfx942" + + +def _build(): + import torch + from torch.utils.cpp_extension import load + # vLLM/SGLang csrc IS written for portable CUDA-style code that depends + # on torch.cpp_extension's hipify pass to rewrite ``cudaStream_t`` → + # ``hipStream_t`` etc. Leave hipify enabled here. (AITER tasks use a + # different runner that disables hipify because its CK template trees + # don't survive the rewrite.) + # Lock the offload-arch list to the runtime device. Torch's default + # populates PYTORCH_ROCM_ARCH with EVERY arch its build knows + # (gfx90a;gfx942;...;gfx1100;gfx1101;...), and the RDNA targets fail on + # vllm's cub bf16 templates. Override with the live device's gfx unless + # the user has set a single non-default arch explicitly. + cur = os.environ.get("PYTORCH_ROCM_ARCH", "") + detected = _detect_gfx() + if not cur or ";" in cur or " " in cur or any( + a in cur for a in ("gfx10", "gfx11", "gfx12") + ): + os.environ["PYTORCH_ROCM_ARCH"] = detected + os.makedirs(BUILD_DIR, exist_ok=True) + # Strip any in-place hipify residue from prior builds. torch.cpp_extension + # generates ``foo.hip`` and ``foo_hip.cuh`` next to the original ``foo.cu`` + # / ``foo.cuh``. If both linger between builds the linker gets duplicate + # symbols. The .hip / *_hip.* files are regenerable from the .cu source. + import shutil as _sh + for stale in glob.glob(os.path.join(SRC_DIR, "**", "*.hip"), recursive=True): + try: os.unlink(stale) + except Exception: pass + for stale in glob.glob(os.path.join(SRC_DIR, "**", "*_hip.*"), recursive=True): + try: os.unlink(stale) + except Exception: pass + sources = [] + for ext in ("cu", "cpp"): + sources.extend(sorted(glob.glob(os.path.join(SRC_DIR, "**", f"*.{ext}"), recursive=True))) + if not sources: + raise RuntimeError("no sources under src/") + # is_python_module=False because bindings.cpp registers via TORCH_LIBRARY + # rather than defining a PyInit_* symbol — the op is reached via + # ``torch.ops..`` after load() returns. + # ``torch.utils.cpp_extension`` defaults to passing + # ``-D__HIP_NO_HALF_OPERATORS__=1 -D__HIP_NO_HALF_CONVERSIONS__=1`` to the + # HIP compiler. PyTorch needs those to keep its own ``c10::Half`` happy, + # but vLLM/SGLang csrc uses raw ``__half2 += __half2`` (in + # ``type_convert.cuh`` and similar) which is precisely the operator that + # flag disables. Their official setup.py build doesn't set the flag; we + # have to undefine it here to make the same source compile through + # ``cpp_extension.load``. + cflags = [ + "-U__HIP_NO_HALF_OPERATORS__", + "-U__HIP_NO_HALF_CONVERSIONS__", + # vLLM/SGLang csrc gates large blocks of fp8 / bf16 helpers behind + # ``ENABLE_FP8`` / ``ENABLE_BF16`` (see e.g. + # quantization/w8a8/fp8/amd/quant_utils.cuh's ``namespace fp8``). + # Their CMake/setup.py define these for the ROCm build; cpp_extension + # doesn't, so the symbols vanish and the dependent .cu files fail with + # "no member named 'scaled_vec_conversion' in namespace 'vllm::fp8'". + "-DENABLE_FP8", + "-DENABLE_BF16", + # AITER fp8 sources also expect this to choose the FNUZ vs E4M3 layout. + "-DHIP_FP8_TYPE_FNUZ", + ] + # vLLM's csrc references ``TORCH_HIP_VERSION`` (a macro vLLM's setup.py + # would normally define from ``HIP_VERSION_MAJOR/MINOR``). Recreate it + # here so the same source compiles under cpp_extension.load. + try: + import torch + hv = getattr(torch.version, "hip", None) + if hv: + major, _, rest = hv.partition(".") + minor = rest.split(".", 1)[0] if rest else "0" + # Torch's TORCH_HIP_VERSION convention: HIP_VERSION_MAJOR*100 + + # HIP_VERSION_MINOR (e.g. 702 for ROCm 7.2). NOT *10000 — that + # extra factor broke ``#if TORCH_HIP_VERSION >= 12090`` style + # gates inherited from CUDA-versioned source (after hipify + # rewrites CUDA_VERSION → TORCH_HIP_VERSION) by making 70200 + # satisfy a CUDA 12.9+ comparison. + torch_hip_version = int(major) * 100 + int(minor) + cflags.append("-DTORCH_HIP_VERSION=" + str(torch_hip_version)) + except Exception: + cflags.append("-DTORCH_HIP_VERSION=702") + return load( + name=NAMESPACE, + sources=sources, + extra_include_paths=[ + SRC_DIR, + os.path.join(SRC_DIR, "core"), + os.path.join(SRC_DIR, "include"), + ], + extra_cflags=cflags, + extra_cuda_cflags=cflags, + verbose=False, + with_cuda=True, + is_python_module=False, + build_directory=BUILD_DIR, + ) + + +def _load_op(): + import torch + _build() + return getattr(getattr(torch.ops, NAMESPACE), OP_NAME) + + +def _test_cases(): + if not os.path.isfile(TEST_CASES): + return [] + with open(TEST_CASES) as f: + return json.load(f) + + +def run_compile(): + try: + _build() + return True, None + except Exception as e: + return False, str(e) + + +def run_correctness(): + import torch + try: + fn = _load_op() + except Exception as e: + return False, f"build failed: {e}" + cases = _test_cases() + if not cases: + return True, "no recorded launch signatures (compile-only check)" + # Drop empty signatures (kernel was registered but never called with args + # in the captured run — eg. `_C` ops shadowed by AITER's wrappers). + cases = [c for c in cases if c.get("args_sig") or c.get("kwargs_sig")] + if not cases: + return True, "all recorded launches had empty signatures (compile-only check)" + ref = rt.reference_for(OP_NAME, REF_SOURCE) + for tc in cases: + try: + args1, kwargs1 = rt.build_inputs(tc, seed=42) + pre = rt.snapshot(args1) + ret1 = fn(*args1, **kwargs1) + if ref is not None: + args_r, kwargs_r = rt.build_inputs(tc, seed=42) + expected = ref(args_r, kwargs_r) + if isinstance(expected, dict): + # [BugA-fix] reference returns {arg_index: expected_tensor} + # for in-place / multi-output kernels; compare each mutated arg. + err = None + for _idx, _exp in expected.items(): + err = rt.compare(args1[_idx], _exp) + if err: + err = f"arg{_idx}: {err}" + break + else: + out1 = rt.detect_output(pre, args1, ret1) + if out1 is None: + return False, f"{tc['test_case_id']}: kernel produced no observable output" + err = rt.compare(out1, expected) + if err: + return False, f"{tc['test_case_id']}: vs reference: {err}" + else: + # determinism check: same seed, same input, byte-identical out + out1 = rt.detect_output(pre, args1, ret1) + if out1 is None: + return False, f"{tc['test_case_id']}: kernel produced no observable output" + args2, kwargs2 = rt.build_inputs(tc, seed=42) + pre2 = rt.snapshot(args2) + ret2 = fn(*args2, **kwargs2) + out2 = rt.detect_output(pre2, args2, ret2) + if out2 is None: + return False, f"{tc['test_case_id']}: repeat produced no observable output" + err = rt.compare(out1, out2) + if err: + return False, f"{tc['test_case_id']}: non-deterministic: {err}" + except Exception as e: + return False, f"{tc['test_case_id']}: kernel raised {e}" + return True, None + + +def run_performance(): + try: + fn = _load_op() + except Exception as exc: + return [{"test_case_id": "load", "execution_time_ms": -1.0, + "params": {"error": str(exc)[:160]}}] + cases = [c for c in _test_cases() if c.get("args_sig") or c.get("kwargs_sig")] + out = [] + for tc in cases: + try: + args, kwargs = rt.build_inputs(tc, seed=42) + avg, metadata = _benchmark_cuda_graph_or_events( + lambda: fn(*args, **kwargs), + warmup=10, + repetition=100, + ) + out.append({"test_case_id": tc["test_case_id"], + "execution_time_ms": avg, **metadata, + "params": tc.get("params_repr", {})}) + except Exception as e: + out.append({"test_case_id": tc["test_case_id"], "execution_time_ms": -1.0, "params": {"error": str(e)[:120]}}) + return out + + +def main(): + ap = argparse.ArgumentParser() + ap.add_argument("mode", choices=["compile", "correctness", "performance"]) + args = ap.parse_args() + os.makedirs(BUILD_DIR, exist_ok=True) + if args.mode == "compile": + ok, err = run_compile() + json.dump({"status": "ok" if ok else "fail", "error": err}, open(os.path.join(BUILD_DIR, "compile_report.json"), "w")) + print(f"Compilation: {'PASS' if ok else 'FAIL'}") + if err: print("Error:", err) + sys.exit(0 if ok else 1) + if args.mode == "correctness": + ok, err = run_correctness() + json.dump({"status": "ok" if ok else "fail", "error": err}, open(os.path.join(BUILD_DIR, "correctness_report.json"), "w")) + print(f"Correctness: {'PASS' if ok else 'FAIL'}") + if err: print("Error:", err) + sys.exit(0 if ok else 1) + cases = run_performance() + json.dump({"test_cases": cases}, open(os.path.join(BUILD_DIR, "performance_report.json"), "w"), indent=2) + for c in cases: + print(f"Performance: {c['execution_time_ms']:.4f} ms ({c['test_case_id']})") + valid = bool(cases) and all(c.get("execution_time_ms", -1) > 0 for c in cases) + sys.exit(0 if valid else 1) + + +if __name__ == "__main__": + main() diff --git a/tasks/hip2hip/campaign20/wvsplitk/src/attention/attention_dtypes.h b/tasks/hip2hip/campaign20/wvsplitk/src/attention/attention_dtypes.h new file mode 100644 index 00000000..64f86381 --- /dev/null +++ b/tasks/hip2hip/campaign20/wvsplitk/src/attention/attention_dtypes.h @@ -0,0 +1,7 @@ +#pragma once + +#include "attention_generic.cuh" +#include "dtype_float16.cuh" +#include "dtype_float32.cuh" +#include "dtype_bfloat16.cuh" +#include "dtype_fp8.cuh" diff --git a/tasks/hip2hip/campaign20/wvsplitk/src/attention/attention_generic.cuh b/tasks/hip2hip/campaign20/wvsplitk/src/attention/attention_generic.cuh new file mode 100644 index 00000000..62409c0c --- /dev/null +++ b/tasks/hip2hip/campaign20/wvsplitk/src/attention/attention_generic.cuh @@ -0,0 +1,65 @@ +/* + * Adapted from + * https://github.com/NVIDIA/FasterTransformer/blob/release/v5.3_tag/src/fastertransformer/kernels/decoder_masked_multihead_attention_utils.h + * Copyright (c) 2023, The vLLM team. + * Copyright (c) 2020-2023, NVIDIA CORPORATION. All rights reserved. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +#pragma once + +#include + +namespace vllm { + +// A vector type to store Q, K, V elements. +template +struct Vec {}; + +// A vector type to store FP32 accumulators. +template +struct FloatVec {}; + +// Template vector operations. +template +inline __device__ Acc mul(A a, B b); + +template +inline __device__ float sum(T v); + +template +inline __device__ float dot(T a, T b) { + return sum(mul(a, b)); +} + +template +inline __device__ float dot(T a, T b) { + return sum(mul(a, b)); +} + +template +inline __device__ void zero(T& dst) { + constexpr int WORDS = sizeof(T) / 4; + union { + T raw; + uint32_t words[WORDS]; + } tmp; + +#pragma unroll + for (int ii = 0; ii < WORDS; ++ii) { + tmp.words[ii] = 0u; + } + dst = tmp.raw; +} + +} // namespace vllm diff --git a/tasks/hip2hip/campaign20/wvsplitk/src/attention/dtype_bfloat16.cuh b/tasks/hip2hip/campaign20/wvsplitk/src/attention/dtype_bfloat16.cuh new file mode 100644 index 00000000..97a25baa --- /dev/null +++ b/tasks/hip2hip/campaign20/wvsplitk/src/attention/dtype_bfloat16.cuh @@ -0,0 +1,463 @@ +/* + * Adapted from + * https://github.com/NVIDIA/FasterTransformer/blob/release/v5.3_tag/src/fastertransformer/kernels/decoder_masked_multihead_attention/decoder_masked_multihead_attention_template.hpp + * and + * https://github.com/NVIDIA/FasterTransformer/blob/release/v5.3_tag/src/fastertransformer/kernels/decoder_masked_multihead_attention_utils.h + * Copyright (c) 2023, The vLLM team. + * Copyright (c) 2020-2023, NVIDIA CORPORATION. All rights reserved. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +#pragma once + +#include "attention_generic.cuh" +#include "dtype_float32.cuh" + +#ifndef USE_ROCM + #include + #include +#else + #include + #include + +typedef __hip_bfloat162 __nv_bfloat162; +typedef __hip_bfloat16 __nv_bfloat16; +#endif + +#include + +namespace vllm { + +// Define custom BF16 vector data types. +struct bf16_4_t { + __nv_bfloat162 x; + __nv_bfloat162 y; +}; + +struct bf16_8_t { + __nv_bfloat162 x; + __nv_bfloat162 y; + __nv_bfloat162 z; + __nv_bfloat162 w; +}; + +// BF16 vector types for Q, K, V. +template <> +struct Vec<__nv_bfloat16, 1> { + using Type = __nv_bfloat16; +}; +template <> +struct Vec<__nv_bfloat16, 2> { + using Type = __nv_bfloat162; +}; +template <> +struct Vec<__nv_bfloat16, 4> { + using Type = bf16_4_t; +}; +template <> +struct Vec<__nv_bfloat16, 8> { + using Type = bf16_8_t; +}; + +// FP32 accumulator vector types corresponding to Vec. +template <> +struct FloatVec<__nv_bfloat16> { + using Type = float; +}; +template <> +struct FloatVec<__nv_bfloat162> { + using Type = float2; +}; +template <> +struct FloatVec { + using Type = Float4_; +}; +template <> +struct FloatVec { + using Type = Float8_; +}; + +// Utility functions for type conversions. +inline __device__ float2 bf1622float2(const __nv_bfloat162 val) { +#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ < 800 + assert(false); +#else + return __bfloat1622float2(val); +#endif + __builtin_unreachable(); // Suppress missing return statement warning +} + +inline __device__ __nv_bfloat162 bf162bf162(const __nv_bfloat16 val) { +#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ < 800 + assert(false); +#else + return __bfloat162bfloat162(val); +#endif + __builtin_unreachable(); // Suppress missing return statement warning +} + +// Vector addition. +inline __device__ __nv_bfloat16 add(__nv_bfloat16 a, __nv_bfloat16 b) { +#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ < 800 + assert(false); +#else + #ifndef USE_ROCM + return a + b; + #else + return __hadd(a, b); + #endif +#endif + __builtin_unreachable(); // Suppress missing return statement warning +} + +inline __device__ __nv_bfloat162 add(__nv_bfloat162 a, __nv_bfloat162 b) { +#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ < 800 + assert(false); +#else + return __hadd2(a, b); +#endif + __builtin_unreachable(); // Suppress missing return statement warning +} + +inline __device__ bf16_4_t add(bf16_4_t a, bf16_4_t b) { + bf16_4_t c; + c.x = add(a.x, b.x); + c.y = add(a.y, b.y); + return c; +} + +inline __device__ bf16_8_t add(bf16_8_t a, bf16_8_t b) { + bf16_8_t c; + c.x = add(a.x, b.x); + c.y = add(a.y, b.y); + c.z = add(a.z, b.z); + c.w = add(a.w, b.w); + return c; +} + +inline __device__ float2 add(__nv_bfloat162 a, float2 fb) { + float2 fa = bf1622float2(a); + return add(fa, fb); +} + +inline __device__ Float4_ add(bf16_4_t a, Float4_ fb) { + Float4_ fc; + fc.x = add(a.x, fb.x); + fc.y = add(a.y, fb.y); + return fc; +} + +inline __device__ Float8_ add(bf16_8_t a, Float8_ fb) { + Float8_ fc; + fc.x = add(a.x, fb.x); + fc.y = add(a.y, fb.y); + fc.z = add(a.z, fb.z); + fc.w = add(a.w, fb.w); + return fc; +} + +// Vector multiplication. +template <> +inline __device__ __nv_bfloat16 mul(__nv_bfloat16 a, __nv_bfloat16 b) { +#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ < 800 + assert(false); +#else + return __hmul(a, b); +#endif + __builtin_unreachable(); // Suppress missing return statement warning +} + +template <> +inline __device__ __nv_bfloat162 mul(__nv_bfloat162 a, __nv_bfloat162 b) { +#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ < 800 + assert(false); +#else + return __hmul2(a, b); +#endif + __builtin_unreachable(); // Suppress missing return statement warning +} + +template <> +inline __device__ __nv_bfloat162 mul(__nv_bfloat16 a, __nv_bfloat162 b) { + return mul<__nv_bfloat162, __nv_bfloat162, __nv_bfloat162>(bf162bf162(a), b); +} + +template <> +inline __device__ bf16_4_t mul(bf16_4_t a, bf16_4_t b) { + bf16_4_t c; + c.x = mul<__nv_bfloat162, __nv_bfloat162, __nv_bfloat162>(a.x, b.x); + c.y = mul<__nv_bfloat162, __nv_bfloat162, __nv_bfloat162>(a.y, b.y); + return c; +} + +template <> +inline __device__ bf16_4_t mul(__nv_bfloat16 a, bf16_4_t b) { + __nv_bfloat162 s = bf162bf162(a); + bf16_4_t c; + c.x = mul<__nv_bfloat162, __nv_bfloat162, __nv_bfloat162>(s, b.x); + c.y = mul<__nv_bfloat162, __nv_bfloat162, __nv_bfloat162>(s, b.y); + return c; +} + +template <> +inline __device__ bf16_8_t mul(bf16_8_t a, bf16_8_t b) { + bf16_8_t c; + c.x = mul<__nv_bfloat162, __nv_bfloat162, __nv_bfloat162>(a.x, b.x); + c.y = mul<__nv_bfloat162, __nv_bfloat162, __nv_bfloat162>(a.y, b.y); + c.z = mul<__nv_bfloat162, __nv_bfloat162, __nv_bfloat162>(a.z, b.z); + c.w = mul<__nv_bfloat162, __nv_bfloat162, __nv_bfloat162>(a.w, b.w); + return c; +} + +template <> +inline __device__ bf16_8_t mul(__nv_bfloat16 a, bf16_8_t b) { + __nv_bfloat162 s = bf162bf162(a); + bf16_8_t c; + c.x = mul<__nv_bfloat162, __nv_bfloat162, __nv_bfloat162>(s, b.x); + c.y = mul<__nv_bfloat162, __nv_bfloat162, __nv_bfloat162>(s, b.y); + c.z = mul<__nv_bfloat162, __nv_bfloat162, __nv_bfloat162>(s, b.z); + c.w = mul<__nv_bfloat162, __nv_bfloat162, __nv_bfloat162>(s, b.w); + return c; +} + +template <> +inline __device__ float mul(__nv_bfloat16 a, __nv_bfloat16 b) { + float fa = __bfloat162float(a); + float fb = __bfloat162float(b); + return fa * fb; +} + +template <> +inline __device__ float2 mul(__nv_bfloat162 a, __nv_bfloat162 b) { + float2 fa = bf1622float2(a); + float2 fb = bf1622float2(b); + return mul(fa, fb); +} + +template <> +inline __device__ float2 mul(__nv_bfloat16 a, __nv_bfloat162 b) { + return mul(bf162bf162(a), b); +} + +template <> +inline __device__ Float4_ mul(bf16_4_t a, bf16_4_t b) { + Float4_ fc; + fc.x = mul(a.x, b.x); + fc.y = mul(a.y, b.y); + return fc; +} + +template <> +inline __device__ Float4_ mul(__nv_bfloat16 a, bf16_4_t b) { + __nv_bfloat162 s = bf162bf162(a); + Float4_ fc; + fc.x = mul(s, b.x); + fc.y = mul(s, b.y); + return fc; +} + +template <> +inline __device__ Float8_ mul(bf16_8_t a, bf16_8_t b) { + Float8_ fc; + fc.x = mul(a.x, b.x); + fc.y = mul(a.y, b.y); + fc.z = mul(a.z, b.z); + fc.w = mul(a.w, b.w); + return fc; +} + +template <> +inline __device__ Float8_ mul(__nv_bfloat16 a, bf16_8_t b) { + __nv_bfloat162 s = bf162bf162(a); + Float8_ fc; + fc.x = mul(s, b.x); + fc.y = mul(s, b.y); + fc.z = mul(s, b.z); + fc.w = mul(s, b.w); + return fc; +} + +// Vector fused multiply-add. +inline __device__ __nv_bfloat162 fma(__nv_bfloat162 a, __nv_bfloat162 b, + __nv_bfloat162 c) { +#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ < 800 + assert(false); +#else + return __hfma2(a, b, c); +#endif + __builtin_unreachable(); // Suppress missing return statement warning +} + +inline __device__ __nv_bfloat162 fma(__nv_bfloat16 a, __nv_bfloat162 b, + __nv_bfloat162 c) { +#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ < 800 + assert(false); +#else + return __hfma2(bf162bf162(a), b, c); +#endif + __builtin_unreachable(); // Suppress missing return statement warning +} + +inline __device__ bf16_4_t fma(bf16_4_t a, bf16_4_t b, bf16_4_t c) { + bf16_4_t d; + d.x = fma(a.x, b.x, c.x); + d.y = fma(a.y, b.y, c.y); + return d; +} + +inline __device__ bf16_4_t fma(__nv_bfloat16 a, bf16_4_t b, bf16_4_t c) { + __nv_bfloat162 s = bf162bf162(a); + bf16_4_t d; + d.x = fma(s, b.x, c.x); + d.y = fma(s, b.y, c.y); + return d; +} + +inline __device__ bf16_8_t fma(bf16_8_t a, bf16_8_t b, bf16_8_t c) { + bf16_8_t d; + d.x = fma(a.x, b.x, c.x); + d.y = fma(a.y, b.y, c.y); + d.z = fma(a.z, b.z, c.z); + d.w = fma(a.w, b.w, c.w); + return d; +} + +inline __device__ bf16_8_t fma(__nv_bfloat16 a, bf16_8_t b, bf16_8_t c) { + __nv_bfloat162 s = bf162bf162(a); + bf16_8_t d; + d.x = fma(s, b.x, c.x); + d.y = fma(s, b.y, c.y); + d.z = fma(s, b.z, c.z); + d.w = fma(s, b.w, c.w); + return d; +} + +inline __device__ float fma(__nv_bfloat16 a, __nv_bfloat16 b, float fc) { + return __bfloat162float(a) * __bfloat162float(b) + fc; +} + +inline __device__ float2 fma(__nv_bfloat162 a, __nv_bfloat162 b, float2 fc) { + float2 fa = bf1622float2(a); + float2 fb = bf1622float2(b); + return fma(fa, fb, fc); +} + +inline __device__ float2 fma(__nv_bfloat16 a, __nv_bfloat162 b, float2 fc) { + return fma(bf162bf162(a), b, fc); +} + +inline __device__ Float4_ fma(bf16_4_t a, bf16_4_t b, Float4_ fc) { + Float4_ fd; + fd.x = fma(a.x, b.x, fc.x); + fd.y = fma(a.y, b.y, fc.y); + return fd; +} + +inline __device__ Float4_ fma(__nv_bfloat16 a, bf16_4_t b, Float4_ fc) { + __nv_bfloat162 s = bf162bf162(a); + Float4_ fd; + fd.x = fma(s, b.x, fc.x); + fd.y = fma(s, b.y, fc.y); + return fd; +} + +inline __device__ Float8_ fma(bf16_8_t a, bf16_8_t b, Float8_ fc) { + Float8_ fd; + fd.x = fma(a.x, b.x, fc.x); + fd.y = fma(a.y, b.y, fc.y); + fd.z = fma(a.z, b.z, fc.z); + fd.w = fma(a.w, b.w, fc.w); + return fd; +} + +inline __device__ Float8_ fma(__nv_bfloat16 a, bf16_8_t b, Float8_ fc) { + __nv_bfloat162 s = bf162bf162(a); + Float8_ fd; + fd.x = fma(s, b.x, fc.x); + fd.y = fma(s, b.y, fc.y); + fd.z = fma(s, b.z, fc.z); + fd.w = fma(s, b.w, fc.w); + return fd; +} + +// Vector sum. +template <> +inline __device__ float sum(__nv_bfloat16 v) { + return __bfloat162float(v); +} + +template <> +inline __device__ float sum(__nv_bfloat162 v) { + float2 vf = bf1622float2(v); + return vf.x + vf.y; +} + +template <> +inline __device__ float sum(bf16_4_t v) { + return sum(v.x) + sum(v.y); +} + +template <> +inline __device__ float sum(bf16_8_t v) { + return sum(v.x) + sum(v.y) + sum(v.z) + sum(v.w); +} + +// From float32 to bfloat16. +inline __device__ void from_float(__nv_bfloat16& dst, float src) { + dst = __float2bfloat16(src); +} + +inline __device__ void from_float(__nv_bfloat162& dst, float2 src) { +#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ < 800 + assert(false); +#else + dst = __float22bfloat162_rn(src); +#endif +} + +inline __device__ void from_float(bf16_4_t& dst, Float4_ src) { +#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ < 800 + assert(false); +#else + dst.x = __float22bfloat162_rn(src.x); + dst.y = __float22bfloat162_rn(src.y); +#endif +} + +inline __device__ void from_float(bf16_8_t& dst, Float8_ src) { +#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ < 800 + assert(false); +#else + dst.x = __float22bfloat162_rn(src.x); + dst.y = __float22bfloat162_rn(src.y); + dst.z = __float22bfloat162_rn(src.z); + dst.w = __float22bfloat162_rn(src.w); +#endif +} + +// From bfloat16 to float32. +inline __device__ float to_float(__nv_bfloat16 u) { + return __bfloat162float(u); +} + +// Zero-out a variable. +inline __device__ void zero(__nv_bfloat16& dst) { +#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ < 800 + assert(false); +#else + // Same as CUDART_ZERO_BF16 introduced in CUDA 12.2. + dst = __ushort_as_bfloat16((unsigned short)0x0000U); +#endif +} + +} // namespace vllm diff --git a/tasks/hip2hip/campaign20/wvsplitk/src/attention/dtype_float16.cuh b/tasks/hip2hip/campaign20/wvsplitk/src/attention/dtype_float16.cuh new file mode 100644 index 00000000..3a1815f0 --- /dev/null +++ b/tasks/hip2hip/campaign20/wvsplitk/src/attention/dtype_float16.cuh @@ -0,0 +1,504 @@ +/* + * Adapted from + * https://github.com/NVIDIA/FasterTransformer/blob/release/v5.3_tag/src/fastertransformer/kernels/decoder_masked_multihead_attention/decoder_masked_multihead_attention_template.hpp + * and + * https://github.com/NVIDIA/FasterTransformer/blob/release/v5.3_tag/src/fastertransformer/kernels/decoder_masked_multihead_attention_utils.h + * Copyright (c) 2023, The vLLM team. + * Copyright (c) 2020-2023, NVIDIA CORPORATION. All rights reserved. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +#pragma once + +#include "attention_generic.cuh" +#include "dtype_float32.cuh" + +#ifdef USE_ROCM + #include +#endif + +#include + +namespace vllm { + +// FP16 vector types for Q, K, V. +template <> +struct Vec { + using Type = uint16_t; +}; +template <> +struct Vec { + using Type = uint32_t; +}; +template <> +struct Vec { + using Type = uint2; +}; +template <> +struct Vec { + using Type = uint4; +}; + +// FP32 accumulator vector types corresponding to Vec. +template <> +struct FloatVec { + using Type = float; +}; +template <> +struct FloatVec { + using Type = float2; +}; +template <> +struct FloatVec { + using Type = Float4_; +}; +template <> +struct FloatVec { + using Type = Float8_; +}; + +// Utility functions for type conversions. +inline __device__ uint32_t h0_h0(uint16_t a) { +#ifndef USE_ROCM + uint32_t b; + asm volatile("mov.b32 %0, {%1, %1};" : "=r"(b) : "h"(a)); + return b; +#else + union { + uint32_t u32; + uint16_t u16[2]; + } tmp; + tmp.u16[0] = a; + tmp.u16[1] = a; + return tmp.u32; +#endif +} + +inline __device__ float half_to_float(uint16_t h) { + float f; +#ifndef USE_ROCM + asm volatile("cvt.f32.f16 %0, %1;\n" : "=f"(f) : "h"(h)); +#else + asm volatile("v_cvt_f32_f16 %0, %1;" : "=v"(f) : "v"(h)); +#endif + return f; +} + +inline __device__ float2 half2_to_float2(uint32_t v) { +#ifndef USE_ROCM + uint16_t lo, hi; + asm volatile("mov.b32 {%0, %1}, %2;\n" : "=h"(lo), "=h"(hi) : "r"(v)); + return make_float2(half_to_float(lo), half_to_float(hi)); +#else + union { + uint32_t u32; + uint16_t u16[2]; + } tmp; + tmp.u32 = v; + float2 ret; + ret.x = half_to_float(tmp.u16[0]); + ret.y = half_to_float(tmp.u16[1]); + return ret; +#endif +} + +inline __device__ uint16_t float_to_half(float f) { + union { + uint32_t u32; + uint16_t u16[2]; + } tmp; +#ifndef USE_ROCM + asm volatile("cvt.rn.f16.f32 %0, %1;\n" : "=h"(tmp.u16[0]) : "f"(f)); +#else + asm volatile("v_cvt_f16_f32 %0, %1;\n" : "=v"(tmp.u32) : "v"(f)); +#endif + return tmp.u16[0]; +} + +inline __device__ uint32_t float2_to_half2(float2 f) { + union { + uint32_t u32; + uint16_t u16[2]; + } tmp; +#ifndef USE_ROCM + #if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 800 + asm volatile("cvt.rn.f16x2.f32 %0, %1, %2;\n" + : "=r"(tmp.u32) + : "f"(f.y), "f"(f.x)); + #else + asm volatile("cvt.rn.f16.f32 %0, %1;\n" : "=h"(tmp.u16[0]) : "f"(f.x)); + asm volatile("cvt.rn.f16.f32 %0, %1;\n" : "=h"(tmp.u16[1]) : "f"(f.y)); + #endif +#else + tmp.u16[0] = float_to_half(f.x); + tmp.u16[1] = float_to_half(f.y); +#endif + return tmp.u32; +} + +// Vector addition. +inline __device__ uint16_t add(uint16_t a, uint16_t b) { + uint16_t c; +#ifndef USE_ROCM + asm volatile("add.f16 %0, %1, %2;\n" : "=h"(c) : "h"(a), "h"(b)); +#else + asm volatile("v_add_f16 %0, %1, %2;\n" : "=v"(c) : "v"(a), "v"(b)); +#endif + return c; +} + +inline __device__ uint32_t add(uint32_t a, uint32_t b) { + uint32_t c; +#ifndef USE_ROCM + asm volatile("add.f16x2 %0, %1, %2;\n" : "=r"(c) : "r"(a), "r"(b)); +#else + asm volatile("v_pk_add_f16 %0, %1, %2;\n" : "=v"(c) : "v"(a), "v"(b)); +#endif + return c; +} + +inline __device__ uint2 add(uint2 a, uint2 b) { + uint2 c; + c.x = add(a.x, b.x); + c.y = add(a.y, b.y); + return c; +} + +inline __device__ uint4 add(uint4 a, uint4 b) { + uint4 c; + c.x = add(a.x, b.x); + c.y = add(a.y, b.y); + c.z = add(a.z, b.z); + c.w = add(a.w, b.w); + return c; +} + +inline __device__ float2 add(uint32_t a, float2 fb) { + float2 fa = half2_to_float2(a); + return add(fa, fb); +} + +inline __device__ Float4_ add(uint2 a, Float4_ fb) { + Float4_ fc; + fc.x = add(a.x, fb.x); + fc.y = add(a.y, fb.y); + return fc; +} + +inline __device__ Float8_ add(uint4 a, Float8_ fb) { + Float8_ fc; + fc.x = add(a.x, fb.x); + fc.y = add(a.y, fb.y); + fc.z = add(a.z, fb.z); + fc.w = add(a.w, fb.w); + return fc; +} + +// Vector multiplication. +template <> +inline __device__ uint16_t mul(uint16_t a, uint16_t b) { + uint16_t c; +#ifndef USE_ROCM + asm volatile("mul.f16 %0, %1, %2;\n" : "=h"(c) : "h"(a), "h"(b)); +#else + asm volatile("v_mul_f16 %0, %1, %2;\n" : "=v"(c) : "v"(a), "v"(b)); +#endif + return c; +} + +template <> +inline __device__ uint32_t mul(uint32_t a, uint32_t b) { + uint32_t c; +#ifndef USE_ROCM + asm volatile("mul.f16x2 %0, %1, %2;\n" : "=r"(c) : "r"(a), "r"(b)); +#else + asm volatile("v_pk_mul_f16 %0, %1, %2;\n" : "=v"(c) : "v"(a), "v"(b)); +#endif + return c; +} + +template <> +inline __device__ uint32_t mul(uint16_t a, uint32_t b) { + return mul(h0_h0(a), b); +} + +template <> +inline __device__ uint2 mul(uint2 a, uint2 b) { + uint2 c; + c.x = mul(a.x, b.x); + c.y = mul(a.y, b.y); + return c; +} + +template <> +inline __device__ uint2 mul(uint16_t a, uint2 b) { + uint32_t s = h0_h0(a); + uint2 c; + c.x = mul(s, b.x); + c.y = mul(s, b.y); + return c; +} + +template <> +inline __device__ uint4 mul(uint4 a, uint4 b) { + uint4 c; + c.x = mul(a.x, b.x); + c.y = mul(a.y, b.y); + c.z = mul(a.z, b.z); + c.w = mul(a.w, b.w); + return c; +} + +template <> +inline __device__ uint4 mul(uint16_t a, uint4 b) { + uint32_t s = h0_h0(a); + uint4 c; + c.x = mul(s, b.x); + c.y = mul(s, b.y); + c.z = mul(s, b.z); + c.w = mul(s, b.w); + return c; +} + +template <> +inline __device__ float mul(uint16_t a, uint16_t b) { + float fa = half_to_float(a); + float fb = half_to_float(b); + return fa * fb; +} + +template <> +inline __device__ float2 mul(uint32_t a, uint32_t b) { + float2 fa = half2_to_float2(a); + float2 fb = half2_to_float2(b); + return mul(fa, fb); +} + +template <> +inline __device__ float2 mul(uint16_t a, uint32_t b) { + return mul(h0_h0(a), b); +} + +template <> +inline __device__ Float4_ mul(uint2 a, uint2 b) { + Float4_ fc; + fc.x = mul(a.x, b.x); + fc.y = mul(a.y, b.y); + return fc; +} + +template <> +inline __device__ Float4_ mul(uint16_t a, uint2 b) { + uint32_t s = h0_h0(a); + Float4_ fc; + fc.x = mul(s, b.x); + fc.y = mul(s, b.y); + return fc; +} + +template <> +inline __device__ Float8_ mul(uint4 a, uint4 b) { + Float8_ fc; + fc.x = mul(a.x, b.x); + fc.y = mul(a.y, b.y); + fc.z = mul(a.z, b.z); + fc.w = mul(a.w, b.w); + return fc; +} + +template <> +inline __device__ Float8_ mul(uint16_t a, uint4 b) { + uint32_t s = h0_h0(a); + Float8_ fc; + fc.x = mul(s, b.x); + fc.y = mul(s, b.y); + fc.z = mul(s, b.z); + fc.w = mul(s, b.w); + return fc; +} + +// Vector fused multiply-add. +inline __device__ uint32_t fma(uint32_t a, uint32_t b, uint32_t c) { + uint32_t d; +#ifndef USE_ROCM + asm volatile("fma.rn.f16x2 %0, %1, %2, %3;\n" + : "=r"(d) + : "r"(a), "r"(b), "r"(c)); +#else + asm volatile("v_pk_fma_f16 %0, %1, %2, %3;\n" + : "=v"(d) + : "v"(a), "v"(b), "v"(c)); +#endif + return d; +} + +inline __device__ uint32_t fma(uint16_t a, uint32_t b, uint32_t c) { + return fma(h0_h0(a), b, c); +} + +inline __device__ uint2 fma(uint2 a, uint2 b, uint2 c) { + uint2 d; + d.x = fma(a.x, b.x, c.x); + d.y = fma(a.y, b.y, c.y); + return d; +} + +inline __device__ uint2 fma(uint16_t a, uint2 b, uint2 c) { + uint32_t s = h0_h0(a); + uint2 d; + d.x = fma(s, b.x, c.x); + d.y = fma(s, b.y, c.y); + return d; +} + +inline __device__ uint4 fma(uint4 a, uint4 b, uint4 c) { + uint4 d; + d.x = fma(a.x, b.x, c.x); + d.y = fma(a.y, b.y, c.y); + d.z = fma(a.z, b.z, c.z); + d.w = fma(a.w, b.w, c.w); + return d; +} + +inline __device__ uint4 fma(uint16_t a, uint4 b, uint4 c) { + uint32_t s = h0_h0(a); + uint4 d; + d.x = fma(s, b.x, c.x); + d.y = fma(s, b.y, c.y); + d.z = fma(s, b.z, c.z); + d.w = fma(s, b.w, c.w); + return d; +} + +inline __device__ float fma(uint16_t a, uint16_t b, float fc) { + float fa = half_to_float(a); + float fb = half_to_float(b); + return fa * fb + fc; +} + +inline __device__ float2 fma(uint32_t a, uint32_t b, float2 fc) { + float2 fa = half2_to_float2(a); + float2 fb = half2_to_float2(b); + return fma(fa, fb, fc); +} + +inline __device__ float2 fma(uint16_t a, uint32_t b, float2 fc) { + return fma(h0_h0(a), b, fc); +} + +inline __device__ Float4_ fma(uint2 a, uint2 b, Float4_ fc) { + Float4_ fd; + fd.x = fma(a.x, b.x, fc.x); + fd.y = fma(a.y, b.y, fc.y); + return fd; +} + +inline __device__ Float4_ fma(uint16_t a, uint2 b, Float4_ fc) { + uint32_t s = h0_h0(a); + Float4_ fd; + fd.x = fma(s, b.x, fc.x); + fd.y = fma(s, b.y, fc.y); + return fd; +} + +inline __device__ Float8_ fma(uint4 a, uint4 b, Float8_ fc) { + Float8_ fd; + fd.x = fma(a.x, b.x, fc.x); + fd.y = fma(a.y, b.y, fc.y); + fd.z = fma(a.z, b.z, fc.z); + fd.w = fma(a.w, b.w, fc.w); + return fd; +} + +inline __device__ Float8_ fma(uint16_t a, uint4 b, Float8_ fc) { + uint32_t s = h0_h0(a); + Float8_ fd; + fd.x = fma(s, b.x, fc.x); + fd.y = fma(s, b.y, fc.y); + fd.z = fma(s, b.z, fc.z); + fd.w = fma(s, b.w, fc.w); + return fd; +} + +// Vector sum. +template <> +inline __device__ float sum(uint16_t v) { + return half_to_float(v); +} + +template <> +inline __device__ float sum(uint32_t v) { + float2 tmp = half2_to_float2(v); + return tmp.x + tmp.y; +} + +template <> +inline __device__ float sum(uint2 v) { + uint32_t c = add(v.x, v.y); + return sum(c); +} + +template <> +inline __device__ float sum(uint4 v) { + uint32_t c = add(v.x, v.y); + c = add(c, v.z); + c = add(c, v.w); + return sum(c); +} + +// From float32 to float16. +inline __device__ void from_float(uint16_t& dst, float src) { + dst = float_to_half(src); +} + +inline __device__ void from_float(uint32_t& dst, float2 src) { + dst = float2_to_half2(src); +} + +inline __device__ void from_float(uint2& dst, Float4_ src) { + dst.x = float2_to_half2(src.x); + dst.y = float2_to_half2(src.y); +} + +inline __device__ void from_float(uint4& dst, Float8_ src) { + dst.x = float2_to_half2(src.x); + dst.y = float2_to_half2(src.y); + dst.z = float2_to_half2(src.z); + dst.w = float2_to_half2(src.w); +} + +// From float16 to float32. +inline __device__ float to_float(uint16_t u) { return half_to_float(u); } + +inline __device__ float2 to_float(uint32_t u) { return half2_to_float2(u); } + +inline __device__ Float4_ to_float(uint2 u) { + Float4_ tmp; + tmp.x = half2_to_float2(u.x); + tmp.y = half2_to_float2(u.y); + return tmp; +} + +inline __device__ Float8_ to_float(uint4 u) { + Float8_ tmp; + tmp.x = half2_to_float2(u.x); + tmp.y = half2_to_float2(u.y); + tmp.z = half2_to_float2(u.z); + tmp.w = half2_to_float2(u.w); + return tmp; +} + +// Zero-out a variable. +inline __device__ void zero(uint16_t& dst) { dst = uint16_t(0); } + +} // namespace vllm diff --git a/tasks/hip2hip/campaign20/wvsplitk/src/attention/dtype_float32.cuh b/tasks/hip2hip/campaign20/wvsplitk/src/attention/dtype_float32.cuh new file mode 100644 index 00000000..7c6a686d --- /dev/null +++ b/tasks/hip2hip/campaign20/wvsplitk/src/attention/dtype_float32.cuh @@ -0,0 +1,251 @@ +/* + * Adapted from + * https://github.com/NVIDIA/FasterTransformer/blob/release/v5.3_tag/src/fastertransformer/kernels/decoder_masked_multihead_attention/decoder_masked_multihead_attention_template.hpp + * and + * https://github.com/NVIDIA/FasterTransformer/blob/release/v5.3_tag/src/fastertransformer/kernels/decoder_masked_multihead_attention_utils.h + * Copyright (c) 2023, The vLLM team. + * Copyright (c) 2020-2023, NVIDIA CORPORATION. All rights reserved. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +#pragma once + +#include "attention_generic.cuh" + +#include + +namespace vllm { + +// Define custom FP32 vector data types. +struct Float4_ { + float2 x; + float2 y; +}; + +struct Float8_ { + float2 x; + float2 y; + float2 z; + float2 w; +}; + +// FP32 vector types for Q, K, V. +template <> +struct Vec { + using Type = float; +}; +template <> +struct Vec { + using Type = float2; +}; +template <> +struct Vec { + using Type = float4; +}; + +// FP32 accumulator vector types corresponding to Vec. +template <> +struct FloatVec { + using Type = float; +}; +template <> +struct FloatVec { + using Type = float2; +}; +template <> +struct FloatVec { + using Type = float4; +}; + +// Vector addition. +inline __device__ float add(float a, float b) { return a + b; } + +inline __device__ float2 add(float2 a, float2 b) { + float2 c; + c.x = add(a.x, b.x); + c.y = add(a.y, b.y); + return c; +} + +inline __device__ float4 add(float4 a, float4 b) { + float4 c; + c.x = add(a.x, b.x); + c.y = add(a.y, b.y); + c.z = add(a.z, b.z); + c.w = add(a.w, b.w); + return c; +} + +// Vector multiplication. +template <> +inline __device__ float mul(float a, float b) { + return a * b; +} + +template <> +inline __device__ float2 mul(float2 a, float2 b) { + float2 c; + c.x = a.x * b.x; + c.y = a.y * b.y; + return c; +} + +template <> +inline __device__ float2 mul(float a, float2 b) { + float2 c; + c.x = a * b.x; + c.y = a * b.y; + return c; +} + +template <> +inline __device__ float4 mul(float4 a, float4 b) { + float4 c; + c.x = a.x * b.x; + c.y = a.y * b.y; + c.z = a.z * b.z; + c.w = a.w * b.w; + return c; +} + +template <> +inline __device__ float4 mul(float a, float4 b) { + float4 c; + c.x = a * b.x; + c.y = a * b.y; + c.z = a * b.z; + c.w = a * b.w; + return c; +} + +// Vector fused multiply-add. +inline __device__ float fma(float a, float b, float c) { return a * b + c; } + +inline __device__ float2 fma(float2 a, float2 b, float2 c) { + float2 d; + d.x = fma(a.x, b.x, c.x); + d.y = fma(a.y, b.y, c.y); + return d; +} + +inline __device__ float2 fma(float a, float2 b, float2 c) { + float2 d; + d.x = fma(a, b.x, c.x); + d.y = fma(a, b.y, c.y); + return d; +} + +inline __device__ float4 fma(float4 a, float4 b, float4 c) { + float4 d; + d.x = fma(a.x, b.x, c.x); + d.y = fma(a.y, b.y, c.y); + d.z = fma(a.z, b.z, c.z); + d.w = fma(a.w, b.w, c.w); + return d; +} + +inline __device__ float4 fma(float a, float4 b, float4 c) { + float4 d; + d.x = fma(a, b.x, c.x); + d.y = fma(a, b.y, c.y); + d.z = fma(a, b.z, c.z); + d.w = fma(a, b.w, c.w); + return d; +} + +inline __device__ Float4_ fma(float a, Float4_ b, Float4_ c) { + Float4_ d; + d.x = fma(a, b.x, c.x); + d.y = fma(a, b.y, c.y); + return d; +} + +inline __device__ Float8_ fma(float a, Float8_ b, Float8_ c) { + Float8_ d; + d.x = fma(a, b.x, c.x); + d.y = fma(a, b.y, c.y); + d.z = fma(a, b.z, c.z); + d.w = fma(a, b.w, c.w); + return d; +} + +// Vector sum. +template <> +inline __device__ float sum(float v) { + return v; +} + +template <> +inline __device__ float sum(float2 v) { + return v.x + v.y; +} + +template <> +inline __device__ float sum(float4 v) { + return v.x + v.y + v.z + v.w; +} + +template <> +inline __device__ float sum(Float4_ v) { + return v.x.x + v.x.y + v.y.x + v.y.y; +} + +template <> +inline __device__ float sum(Float8_ v) { + return v.x.x + v.x.y + v.y.x + v.y.y + v.z.x + v.z.y + v.w.x + v.w.y; +} + +// Vector dot product. +inline __device__ float dot(float a, float b) { return a * b; } + +inline __device__ float dot(float2 a, float2 b) { + float2 c = mul(a, b); + return c.x + c.y; +} + +inline __device__ float dot(Float4_ a, Float4_ b) { + float2 acc = mul(a.x, b.x); + acc = fma(a.y, b.y, acc); + return acc.x + acc.y; +} + +inline __device__ float dot(Float8_ a, Float8_ b) { + float2 acc = mul(a.x, b.x); + acc = fma(a.y, b.y, acc); + acc = fma(a.z, b.z, acc); + acc = fma(a.w, b.w, acc); + return acc.x + acc.y; +} + +// From float to float. +inline __device__ void from_float(float& dst, float src) { dst = src; } + +inline __device__ void from_float(float2& dst, float2 src) { dst = src; } + +inline __device__ void from_float(float4& dst, float4 src) { dst = src; } + +// From float to float. +inline __device__ float to_float(float u) { return u; } + +inline __device__ float2 to_float(float2 u) { return u; } + +inline __device__ float4 to_float(float4 u) { return u; } + +inline __device__ Float4_ to_float(Float4_ u) { return u; } + +inline __device__ Float8_ to_float(Float8_ u) { return u; } + +// Zero-out a variable. +inline __device__ void zero(float& dst) { dst = 0.f; } + +} // namespace vllm diff --git a/tasks/hip2hip/campaign20/wvsplitk/src/attention/dtype_fp8.cuh b/tasks/hip2hip/campaign20/wvsplitk/src/attention/dtype_fp8.cuh new file mode 100644 index 00000000..e714e321 --- /dev/null +++ b/tasks/hip2hip/campaign20/wvsplitk/src/attention/dtype_fp8.cuh @@ -0,0 +1,41 @@ +#pragma once + +#include "attention_generic.cuh" + +#include +#ifdef ENABLE_FP8 + #ifndef USE_ROCM + #include + #endif // USE_ROCM +#endif // ENABLE_FP8 + +namespace vllm { + +enum class Fp8KVCacheDataType { + kAuto = 0, + kFp8E4M3 = 1, + kFp8E5M2 = 2, +}; + +// fp8 vector types for quantization of kv cache +template <> +struct Vec { + using Type = uint8_t; +}; + +template <> +struct Vec { + using Type = uint16_t; +}; + +template <> +struct Vec { + using Type = uint32_t; +}; + +template <> +struct Vec { + using Type = uint2; +}; + +} // namespace vllm diff --git a/tasks/hip2hip/campaign20/wvsplitk/src/bindings.cpp b/tasks/hip2hip/campaign20/wvsplitk/src/bindings.cpp new file mode 100644 index 00000000..d6cbec9f --- /dev/null +++ b/tasks/hip2hip/campaign20/wvsplitk/src/bindings.cpp @@ -0,0 +1,15 @@ +// Auto-generated by extractor.source.hip_extractor. +// Re-registers the extracted op into the `extracted_wvsplitk` namespace so the task +// can call torch.ops.extracted_wvsplitk.wvSplitK without touching the host engine. +#include +#include + +torch::Tensor wvSplitK(const at::Tensor& in_a, const at::Tensor& in_b, const std::optional& in_bias, const int64_t CuCount); + +TORCH_LIBRARY(extracted_wvsplitk, m) { + m.def("wvSplitK(Tensor in_a, Tensor in_b, Tensor? in_bias, int CuCount) -> Tensor"); +} + +TORCH_LIBRARY_IMPL(extracted_wvsplitk, CUDA, m) { + m.impl("wvSplitK", &wvSplitK); +} diff --git a/tasks/hip2hip/campaign20/wvsplitk/src/core/batch_invariant.hpp b/tasks/hip2hip/campaign20/wvsplitk/src/core/batch_invariant.hpp new file mode 100644 index 00000000..fffe96b8 --- /dev/null +++ b/tasks/hip2hip/campaign20/wvsplitk/src/core/batch_invariant.hpp @@ -0,0 +1,19 @@ +#pragma once +#include +#include +#include + +namespace vllm { + +// vllm_is_batch_invariant(); returns true +// if env VLLM_BATCH_INVARIANT=1 +inline bool vllm_is_batch_invariant() { + static bool cached = []() { + std::string env_key = "VLLM_BATCH_INVARIANT"; + const char* val = std::getenv(env_key.c_str()); + return (val && std::atoi(val) != 0) ? 1 : 0; + }(); + return cached; +} + +} // namespace vllm diff --git a/tasks/hip2hip/campaign20/wvsplitk/src/cuda_compat.h b/tasks/hip2hip/campaign20/wvsplitk/src/cuda_compat.h new file mode 100644 index 00000000..d7d589db --- /dev/null +++ b/tasks/hip2hip/campaign20/wvsplitk/src/cuda_compat.h @@ -0,0 +1,76 @@ +#pragma once + +#ifdef USE_ROCM + #include +#endif + +#ifdef USE_ROCM +struct Utils { + static __host__ int get_warp_size() { + static bool is_cached = false; + static int result; + + if (!is_cached) { + int device_id; + cudaDeviceProp deviceProp; + cudaGetDevice(&device_id); + cudaGetDeviceProperties(&deviceProp, device_id); + + result = deviceProp.warpSize; + is_cached = true; + } + + return result; + } + + static __device__ constexpr int get_warp_size() { + #ifdef __GFX9__ + return 64; + #else + return 32; + #endif + } +}; + + #define WARP_SIZE Utils::get_warp_size() +#else + #define WARP_SIZE 32 +#endif + +#ifndef USE_ROCM + #define VLLM_LDG(arg) __ldg(arg) +#else + #define VLLM_LDG(arg) *(arg) +#endif + +#ifndef USE_ROCM + #define VLLM_SHFL_XOR_SYNC(var, lane_mask) \ + __shfl_xor_sync(uint32_t(-1), var, lane_mask) + #define VLLM_SHFL_XOR_SYNC_WIDTH(var, lane_mask, width) \ + __shfl_xor_sync(uint32_t(-1), var, lane_mask, width) +#else + #define VLLM_SHFL_XOR_SYNC(var, lane_mask) __shfl_xor(var, lane_mask) + #define VLLM_SHFL_XOR_SYNC_WIDTH(var, lane_mask, width) \ + __shfl_xor(var, lane_mask, width) +#endif + +#ifndef USE_ROCM + #define VLLM_SHFL_SYNC(var, src_lane) __shfl_sync(uint32_t(-1), var, src_lane) +#else + #define VLLM_SHFL_SYNC(var, src_lane) __shfl(var, src_lane) +#endif + +#ifndef USE_ROCM + #define VLLM_SHFL_DOWN_SYNC(var, lane_delta) \ + __shfl_down_sync(uint32_t(-1), var, lane_delta) +#else + #define VLLM_SHFL_DOWN_SYNC(var, lane_delta) __shfl_down(var, lane_delta) +#endif + +#ifndef USE_ROCM + #define VLLM_DevFuncAttribute_SET_MaxDynamicSharedMemorySize(FUNC, VAL) \ + cudaFuncSetAttribute(FUNC, cudaFuncAttributeMaxDynamicSharedMemorySize, VAL) +#else + #define VLLM_DevFuncAttribute_SET_MaxDynamicSharedMemorySize(FUNC, VAL) \ + hipFuncSetAttribute(FUNC, hipFuncAttributeMaxDynamicSharedMemorySize, VAL) +#endif diff --git a/tasks/hip2hip/campaign20/wvsplitk/src/dispatch_utils.h b/tasks/hip2hip/campaign20/wvsplitk/src/dispatch_utils.h new file mode 100644 index 00000000..de0c505b --- /dev/null +++ b/tasks/hip2hip/campaign20/wvsplitk/src/dispatch_utils.h @@ -0,0 +1,158 @@ +/* + * Adapted from + * https://github.com/pytorch/pytorch/blob/v2.0.1/aten/src/ATen/Dispatch.h + */ +#pragma once + +#include + +// Need a special dispatch case macro since we will nest the FP8 dispatch. +// Instead of the usual 'scalar_t', this names the dispatched type 'fp8_t'. +#define AT_DISPATCH_FP8_CASE(enum_type, ...) \ + AT_PRIVATE_CASE_TYPE_USING_HINT(enum_type, fp8_t, __VA_ARGS__) + +#define VLLM_DISPATCH_CASE_FLOATING_TYPES(...) \ + AT_DISPATCH_CASE(at::ScalarType::Float, __VA_ARGS__) \ + AT_DISPATCH_CASE(at::ScalarType::Half, __VA_ARGS__) \ + AT_DISPATCH_CASE(at::ScalarType::BFloat16, __VA_ARGS__) + +#define VLLM_DISPATCH_FLOATING_TYPES(TYPE, NAME, ...) \ + AT_DISPATCH_SWITCH(TYPE, NAME, VLLM_DISPATCH_CASE_FLOATING_TYPES(__VA_ARGS__)) + +#define VLLM_DISPATCH_CASE_HALF_TYPES(...) \ + AT_DISPATCH_CASE(at::ScalarType::Half, __VA_ARGS__) \ + AT_DISPATCH_CASE(at::ScalarType::BFloat16, __VA_ARGS__) + +#define VLLM_DISPATCH_HALF_TYPES(TYPE, NAME, ...) \ + AT_DISPATCH_SWITCH(TYPE, NAME, VLLM_DISPATCH_CASE_HALF_TYPES(__VA_ARGS__)) + +// ROCm devices might use either fn or fnuz, so set up dispatch table for both. +// A host-based check at runtime will create a preferred FP8 type for ROCm +// such that the correct kernel is dispatched. +#ifdef USE_ROCM + #define VLLM_DISPATCH_CASE_FP8_TYPES(...) \ + AT_DISPATCH_FP8_CASE(at::ScalarType::Float8_e4m3fn, __VA_ARGS__) \ + AT_DISPATCH_FP8_CASE(at::ScalarType::Float8_e4m3fnuz, __VA_ARGS__) + + #define VLLM_DISPATCH_CASE_QUANT_TYPES(...) \ + AT_DISPATCH_CASE(at::ScalarType::Float8_e4m3fn, __VA_ARGS__) \ + AT_DISPATCH_CASE(at::ScalarType::Float8_e4m3fnuz, __VA_ARGS__) \ + AT_DISPATCH_CASE(at::ScalarType::Char, __VA_ARGS__) +#else + #define VLLM_DISPATCH_CASE_FP8_TYPES(...) \ + AT_DISPATCH_FP8_CASE(at::ScalarType::Float8_e4m3fn, __VA_ARGS__) + + #define VLLM_DISPATCH_CASE_QUANT_TYPES(...) \ + AT_DISPATCH_CASE(at::ScalarType::Float8_e4m3fn, __VA_ARGS__) \ + AT_DISPATCH_CASE(at::ScalarType::Char, __VA_ARGS__) +#endif + +// When using this dispatch macro, the type is 'fp8_t' not 'scalar_t'. +// See AT_DISPATCH_FP8_CASE above. +#define VLLM_DISPATCH_FP8_TYPES(TYPE, NAME, ...) \ + AT_DISPATCH_SWITCH(TYPE, NAME, VLLM_DISPATCH_CASE_FP8_TYPES(__VA_ARGS__)) + +#define VLLM_DISPATCH_QUANT_TYPES(TYPE, NAME, ...) \ + AT_DISPATCH_SWITCH(TYPE, NAME, VLLM_DISPATCH_CASE_QUANT_TYPES(__VA_ARGS__)) + +#define VLLM_DISPATCH_CASE_FLOATING_AND_BYTE_TYPES(...) \ + AT_DISPATCH_CASE(at::ScalarType::Float, __VA_ARGS__) \ + AT_DISPATCH_CASE(at::ScalarType::Half, __VA_ARGS__) \ + AT_DISPATCH_CASE(at::ScalarType::BFloat16, __VA_ARGS__) \ + AT_DISPATCH_CASE(at::ScalarType::Byte, __VA_ARGS__) + +#define VLLM_DISPATCH_FLOATING_AND_BYTE_TYPES(TYPE, NAME, ...) \ + AT_DISPATCH_SWITCH(TYPE, NAME, \ + VLLM_DISPATCH_CASE_FLOATING_AND_BYTE_TYPES(__VA_ARGS__)) + +#define VLLM_DISPATCH_CASE_INTEGRAL_TYPES(...) \ + AT_DISPATCH_CASE(at::ScalarType::Byte, __VA_ARGS__) \ + AT_DISPATCH_CASE(at::ScalarType::Char, __VA_ARGS__) \ + AT_DISPATCH_CASE(at::ScalarType::Short, __VA_ARGS__) \ + AT_DISPATCH_CASE(at::ScalarType::Int, __VA_ARGS__) \ + AT_DISPATCH_CASE(at::ScalarType::Long, __VA_ARGS__) + +#define VLLM_DISPATCH_CASE_INTEGRAL_AND_UNSIGNED_TYPES(...) \ + AT_DISPATCH_CASE(at::ScalarType::Byte, __VA_ARGS__) \ + AT_DISPATCH_CASE(at::ScalarType::Char, __VA_ARGS__) \ + AT_DISPATCH_CASE(at::ScalarType::Short, __VA_ARGS__) \ + AT_DISPATCH_CASE(at::ScalarType::Int, __VA_ARGS__) \ + AT_DISPATCH_CASE(at::ScalarType::Long, __VA_ARGS__) \ + AT_DISPATCH_CASE(at::ScalarType::UInt16, __VA_ARGS__) \ + AT_DISPATCH_CASE(at::ScalarType::UInt32, __VA_ARGS__) \ + AT_DISPATCH_CASE(at::ScalarType::UInt64, __VA_ARGS__) + +#define VLLM_DISPATCH_INTEGRAL_TYPES(TYPE, NAME, ...) \ + AT_DISPATCH_SWITCH(TYPE, NAME, VLLM_DISPATCH_CASE_INTEGRAL_TYPES(__VA_ARGS__)) + +#define VLLM_DISPATCH_INTEGRAL_AND_UNSIGNED_TYPES(TYPE, NAME, ...) \ + AT_DISPATCH_SWITCH( \ + TYPE, NAME, VLLM_DISPATCH_CASE_INTEGRAL_AND_UNSIGNED_TYPES(__VA_ARGS__)) + +#define VLLM_DISPATCH_VEC_SIZE(VEC_SIZE, ...) \ + switch (VEC_SIZE) { \ + case 16: { \ + constexpr int vec_size = 16; \ + __VA_ARGS__(); \ + break; \ + } \ + case 8: { \ + constexpr int vec_size = 8; \ + __VA_ARGS__(); \ + break; \ + } \ + case 4: { \ + constexpr int vec_size = 4; \ + __VA_ARGS__(); \ + break; \ + } \ + case 2: { \ + constexpr int vec_size = 2; \ + __VA_ARGS__(); \ + break; \ + } \ + default: { \ + constexpr int vec_size = 1; \ + __VA_ARGS__(); \ + break; \ + } \ + } + +#define VLLM_DISPATCH_BOOL(expr, const_expr, ...) \ + if (expr) { \ + constexpr bool const_expr = true; \ + __VA_ARGS__(); \ + } else { \ + constexpr bool const_expr = false; \ + __VA_ARGS__(); \ + } + +#define VLLM_DISPATCH_GROUP_SIZE(group_size, const_group_size, ...) \ + if (group_size == 128) { \ + constexpr int const_group_size = 128; \ + __VA_ARGS__(); \ + } else if (group_size == 64) { \ + constexpr int const_group_size = 64; \ + __VA_ARGS__(); \ + } + +#define VLLM_DISPATCH_RANK234(NUM_DIMS, ...) \ + switch (NUM_DIMS) { \ + case 2: { \ + constexpr int tensor_rank = 2; \ + __VA_ARGS__(); \ + break; \ + } \ + case 3: { \ + constexpr int tensor_rank = 3; \ + __VA_ARGS__(); \ + break; \ + } \ + case 4: { \ + constexpr int tensor_rank = 4; \ + __VA_ARGS__(); \ + break; \ + } \ + default: \ + TORCH_CHECK(false, "Expects rank 2, 3 or 4 tensors but got ", NUM_DIMS); \ + } diff --git a/tasks/hip2hip/campaign20/wvsplitk/src/hip_compat.h b/tasks/hip2hip/campaign20/wvsplitk/src/hip_compat.h new file mode 100644 index 00000000..645f0e17 --- /dev/null +++ b/tasks/hip2hip/campaign20/wvsplitk/src/hip_compat.h @@ -0,0 +1,77 @@ +// !!! This is a file automatically generated by hipify!!! +#pragma once + +#ifdef USE_ROCM + #include +#endif + +#ifdef USE_ROCM +struct Utils { + static __host__ int get_warp_size() { + static bool is_cached = false; + static int result; + + if (!is_cached) { + int device_id; + hipDeviceProp_t deviceProp; + hipGetDevice(&device_id); + hipGetDeviceProperties(&deviceProp, device_id); + + result = deviceProp.warpSize; + is_cached = true; + } + + return result; + } + + static __device__ constexpr int get_warp_size() { + #ifdef __GFX9__ + return 64; + #else + return 32; + #endif + } +}; + + #define WARP_SIZE Utils::get_warp_size() +#else + #define WARP_SIZE 32 +#endif + +#ifndef USE_ROCM + #define VLLM_LDG(arg) __ldg(arg) +#else + #define VLLM_LDG(arg) *(arg) +#endif + +#ifndef USE_ROCM + #define VLLM_SHFL_XOR_SYNC(var, lane_mask) \ + __shfl_xor_sync(uint32_t(-1), var, lane_mask) + #define VLLM_SHFL_XOR_SYNC_WIDTH(var, lane_mask, width) \ + __shfl_xor_sync(uint32_t(-1), var, lane_mask, width) +#else + #define VLLM_SHFL_XOR_SYNC(var, lane_mask) __shfl_xor(var, lane_mask) + #define VLLM_SHFL_XOR_SYNC_WIDTH(var, lane_mask, width) \ + __shfl_xor(var, lane_mask, width) +#endif + +#ifndef USE_ROCM + #define VLLM_SHFL_SYNC(var, src_lane) __shfl_sync(uint32_t(-1), var, src_lane) +#else + #define VLLM_SHFL_SYNC(var, src_lane) __shfl(var, src_lane) +#endif + +#ifndef USE_ROCM + #define VLLM_SHFL_DOWN_SYNC(var, lane_delta) \ + __shfl_down_sync(uint32_t(-1), var, lane_delta) +#else + #define VLLM_SHFL_DOWN_SYNC(var, lane_delta) __shfl_down(var, lane_delta) +#endif + +#ifndef USE_ROCM + #define VLLM_DevFuncAttribute_SET_MaxDynamicSharedMemorySize(FUNC, VAL) \ + hipFuncSetAttribute(FUNC, hipFuncAttributeMaxDynamicSharedMemorySize, VAL) +#else + #define VLLM_DevFuncAttribute_SET_MaxDynamicSharedMemorySize(FUNC, VAL) \ + hipFuncSetAttribute(FUNC, hipFuncAttributeMaxDynamicSharedMemorySize, VAL) +#endif diff --git a/tasks/hip2hip/campaign20/wvsplitk/src/libtorch_stable/quantization/vectorization.cuh b/tasks/hip2hip/campaign20/wvsplitk/src/libtorch_stable/quantization/vectorization.cuh new file mode 100644 index 00000000..9d5eea00 --- /dev/null +++ b/tasks/hip2hip/campaign20/wvsplitk/src/libtorch_stable/quantization/vectorization.cuh @@ -0,0 +1,31 @@ +#pragma once +/** + * __device__ datatypes vectorized by 4 + */ + +// Include both AMD and NVIDIA fp8 types to avoid circular import +#include +#include + +namespace vllm { + +// Vectorization containers +template +struct __align__(vec_size * sizeof(scalar_t)) vec_n_t { + scalar_t val[vec_size]; +}; + +template +struct __align__(vec_size * sizeof(quant_type_t)) q8_n_t { + static_assert(std::is_same_v || + std::is_same_v || + std::is_same_v); + quant_type_t val[vec_size]; +}; + +template +using vec4_t = vec_n_t; +template +using q8x4_t = q8_n_t; + +} // namespace vllm diff --git a/tasks/hip2hip/campaign20/wvsplitk/src/quantization/utils.cuh b/tasks/hip2hip/campaign20/wvsplitk/src/quantization/utils.cuh new file mode 100644 index 00000000..73055a15 --- /dev/null +++ b/tasks/hip2hip/campaign20/wvsplitk/src/quantization/utils.cuh @@ -0,0 +1,59 @@ +#pragma once + +/** + * Quantization utilities including: + * Adjusted maximum values for qtypes. + * Minimum scaling factors for qtypes. + */ + +#include +#include + +#ifndef USE_ROCM + #include + #define MAYBE_HOST_DEVICE C10_HOST_DEVICE +#else + #include + #include + #include + // ROCm doesn't seem to need C10_HOST_DEVICE for static constexpr + #define MAYBE_HOST_DEVICE +#endif + +template || + std::is_same_v || + std::is_same_v>> +struct quant_type_max { + static constexpr T val() { return std::numeric_limits::max(); } +}; + +// Using the default max value from pytorch (240.0 0x7F) will cause accuracy +// issues when running dynamic quantization. Here use 224.0 0x7E for rocm. +template <> +struct quant_type_max { + static constexpr c10::Float8_e4m3fnuz val() { + return c10::Float8_e4m3fnuz(0x7E, c10::Float8_e4m3fnuz::from_bits()); + } +}; + +template +MAYBE_HOST_DEVICE static constexpr T quant_type_max_v = + quant_type_max::val(); + +template || + std::is_same_v || + std::is_same_v>> +struct min_scaling_factor { + C10_DEVICE C10_ALWAYS_INLINE static float val() { + return 1.0f / (quant_type_max_v * 512.0f); + } +}; + +template <> +struct min_scaling_factor { + C10_DEVICE C10_ALWAYS_INLINE static float val() { + return std::numeric_limits::epsilon(); + } +}; \ No newline at end of file diff --git a/tasks/hip2hip/campaign20/wvsplitk/src/quantization/w8a8/fp8/amd/quant_utils.cuh b/tasks/hip2hip/campaign20/wvsplitk/src/quantization/w8a8/fp8/amd/quant_utils.cuh new file mode 100644 index 00000000..81f5cb83 --- /dev/null +++ b/tasks/hip2hip/campaign20/wvsplitk/src/quantization/w8a8/fp8/amd/quant_utils.cuh @@ -0,0 +1,671 @@ +#pragma once +#include + +#include +#include +#include + +#include "../../../../attention/attention_dtypes.h" + +namespace vllm { +#ifdef USE_ROCM + +namespace fp8 { + #ifdef ENABLE_FP8 + +// Use hardware cvt instruction for fp8 on rocm +template +__device__ __forceinline__ fp8_type cvt_c10(float const r) { + return {}; +} + +// __hip_fp8_e4m3 only exists starting in ROCm 6.3. The macro +// HIP_FP8_TYPE_OCP comes from the hip_fp8.h header and also makes +// its first appearance in ROCm 6.3. Since VLLM_DISPATCH_FP8_TYPES +// on ROCm instantiates both OCP and FNUZ kernels, we need to replace +// the new HW cvt with something reasonable that doesn't rely on the +// ROCm 6.3 feature. This allows compiling on ROCm 6.2 or newer. +template <> +__device__ __forceinline__ c10::Float8_e4m3fn cvt_c10(float const r) { + #if HIP_FP8_TYPE_OCP + return c10::Float8_e4m3fn( + __hip_cvt_float_to_fp8(r, __hip_fp8_e4m3::__default_saturation, + __hip_fp8_e4m3::__default_interpret), + c10::Float8_e4m3fn::from_bits()); + #else + // Cast implemented by pytorch. Uses bit manipulation instead of HW cvt. + // HW cvt above is faster when it is available (ROCm 6.3 or newer). + return static_cast(r); + #endif +} + +template <> +__device__ __forceinline__ c10::Float8_e4m3fnuz cvt_c10(float const r) { + return c10::Float8_e4m3fnuz( + __hip_cvt_float_to_fp8(r, __hip_fp8_e4m3_fnuz::__default_saturation, + __hip_fp8_e4m3_fnuz::__default_interpret), + c10::Float8_e4m3fnuz::from_bits()); +} + +template +__inline__ __device__ Tout vec_conversion(const Tin& x) { + return x; +} + +template +__inline__ __device__ Tout scaled_vec_conversion(const Tin& x, + const float scale) { + return x; +} + + #if HIP_FP8_TYPE_OCP +using fp8_type = __hip_fp8_e4m3; +using fp8x2_type = __hip_fp8x2_e4m3; + #else +using fp8_type = __hip_fp8_e4m3_fnuz; +using fp8x2_type = __hip_fp8x2_e4m3_fnuz; + #endif + +// fp8 -> half +template <> +__inline__ __device__ uint16_t +vec_conversion(const uint8_t& a) { + return __hip_cvt_fp8_to_halfraw(a, fp8_type::__default_interpret).x; +} + +// fp8x2 -> half2 +template <> +__inline__ __device__ uint32_t +vec_conversion(const uint16_t& a) { + union { + __half2_raw h2r; + uint32_t ui32; + } tmp; + tmp.h2r = __hip_cvt_fp8x2_to_halfraw2(a, fp8_type::__default_interpret); + return tmp.ui32; +} + +// fp8x4 -> half2x2 +template <> +__inline__ __device__ uint2 vec_conversion(const uint32_t& a) { + union { + uint2 u32x2; + uint32_t u32[2]; + } tmp; + tmp.u32[0] = vec_conversion((uint16_t)a); + tmp.u32[1] = vec_conversion((uint16_t)(a >> 16U)); + return tmp.u32x2; +} + +// fp8x8 -> half2x4 +template <> +__inline__ __device__ uint4 vec_conversion(const uint2& a) { + union { + uint4 u64x2; + uint2 u64[2]; + } tmp; + tmp.u64[0] = vec_conversion(a.x); + tmp.u64[1] = vec_conversion(a.y); + return tmp.u64x2; +} + +using __nv_bfloat16 = __hip_bfloat16; + +// fp8 -> __nv_bfloat16 +template <> +__inline__ __device__ __nv_bfloat16 +vec_conversion<__nv_bfloat16, uint8_t>(const uint8_t& a) { + fp8_type f8; + f8.__x = a; + return __float2bfloat16(static_cast(f8)); +} + +using __nv_bfloat162 = __hip_bfloat162; + +// fp8x2 -> __nv_bfloat162 +template <> +__inline__ __device__ __nv_bfloat162 +vec_conversion<__nv_bfloat162, uint16_t>(const uint16_t& a) { + __nv_bfloat162 res; + res.x = vec_conversion<__nv_bfloat16, uint8_t>((uint8_t)a); + res.y = vec_conversion<__nv_bfloat16, uint8_t>((uint8_t)(a >> 8U)); + return res; +} + +// fp8x4 -> bf16_4_t +template <> +__inline__ __device__ bf16_4_t +vec_conversion(const uint32_t& a) { + bf16_4_t res; + res.x = vec_conversion<__nv_bfloat162, uint16_t>((uint16_t)a); + res.y = vec_conversion<__nv_bfloat162, uint16_t>((uint16_t)(a >> 16U)); + return res; +} + +// fp8x8 -> bf16_8_t +template <> +__inline__ __device__ bf16_8_t vec_conversion(const uint2& a) { + bf16_4_t tmp1, tmp2; + tmp1 = vec_conversion(a.x); + tmp2 = vec_conversion(a.y); + bf16_8_t res; + res.x = tmp1.x; + res.y = tmp1.y; + res.z = tmp2.x; + res.w = tmp2.y; + return res; +} + +// fp8 -> float +template <> +__inline__ __device__ float vec_conversion(const uint8_t& a) { + fp8_type f8; + f8.__x = a; + return static_cast(f8); +} + +// fp8x2 -> float2 +template <> +__inline__ __device__ float2 +vec_conversion(const uint16_t& a) { + fp8x2_type f8x2; + f8x2.__x = a; + return static_cast(f8x2); +} + +// fp8x4 -> float4 +template <> +__inline__ __device__ Float4_ +vec_conversion(const uint32_t& a) { + Float4_ res; + res.x = vec_conversion((uint16_t)a); + res.y = vec_conversion((uint16_t)(a >> 16U)); + return res; +} + +// fp8x4 -> float4 +template <> +__inline__ __device__ float4 +vec_conversion(const uint32_t& a) { + Float4_ tmp = vec_conversion(a); + float4 res = make_float4(tmp.x.x, tmp.x.y, tmp.y.x, tmp.y.y); + return res; +} + +// fp8x8 -> float8 +template <> +__inline__ __device__ Float8_ vec_conversion(const uint2& a) { + Float4_ tmp1, tmp2; + tmp1 = vec_conversion(a.x); + tmp2 = vec_conversion(a.y); + Float8_ res; + res.x = tmp1.x; + res.y = tmp1.y; + res.z = tmp2.x; + res.w = tmp2.y; + return res; +} + +// half -> fp8 +template <> +__inline__ __device__ uint8_t +vec_conversion(const uint16_t& a) { + __half_raw tmp; + tmp.x = a; + return __hip_cvt_halfraw_to_fp8(tmp, fp8_type::__default_saturation, + fp8_type::__default_interpret); +} + +template <> +__inline__ __device__ uint16_t +vec_conversion(const uint32_t& a) { + union { + uint32_t ui32; + __half2_raw h2r; + } tmp; + tmp.ui32 = a; + return __hip_cvt_halfraw2_to_fp8x2(tmp.h2r, fp8_type::__default_saturation, + fp8_type::__default_interpret); +} + +// bf16 -> fp8 +template <> +__inline__ __device__ uint8_t +vec_conversion(const __nv_bfloat16& a) { + return __hip_cvt_float_to_fp8(__bfloat162float(a), + fp8_type::__default_saturation, + fp8_type::__default_interpret); +} + +// float -> fp8 +template <> +__inline__ __device__ uint8_t vec_conversion(const float& a) { + return __hip_cvt_float_to_fp8(a, fp8_type::__default_saturation, + fp8_type::__default_interpret); +} + +// float2 -> half2 +template <> +__inline__ __device__ uint32_t +vec_conversion(const float2& a) { + union { + half2 float16; + uint32_t uint32; + }; + + float16 = __float22half2_rn(a); + return uint32; +} + +// Float4 -> half2x2 +template <> +__inline__ __device__ uint2 vec_conversion(const Float4_& a) { + uint2 b; + float2 val; + val.x = a.x.x; + val.y = a.x.y; + b.x = vec_conversion(val); + + val.x = a.y.x; + val.y = a.y.y; + b.y = vec_conversion(val); + return b; +} + +// Float4 -> float4 +template <> +__inline__ __device__ float4 vec_conversion(const Float4_& a) { + float4 b; + b.x = a.x.x; + b.y = a.x.y; + b.z = a.y.x; + b.w = a.y.y; + return b; +} + +// Float8 -> half2x4 +template <> +__inline__ __device__ uint4 vec_conversion(const Float8_& a) { + uint4 b; + b.x = vec_conversion(a.x); + b.y = vec_conversion(a.y); + b.z = vec_conversion(a.z); + b.w = vec_conversion(a.w); + return b; +} + +// float2 -> bfloat162 +template <> +__inline__ __device__ __nv_bfloat162 +vec_conversion<__nv_bfloat162, float2>(const float2& a) { + __nv_bfloat162 b = __float22bfloat162_rn(a); + return b; +} + +// Float4 -> bfloat162x2 +template <> +__inline__ __device__ bf16_4_t +vec_conversion(const Float4_& a) { + bf16_4_t b; + b.x = __float22bfloat162_rn(a.x); + b.y = __float22bfloat162_rn(a.y); + return b; +} + +// Float8 -> bfloat162x4 +template <> +__inline__ __device__ bf16_8_t +vec_conversion(const Float8_& a) { + bf16_8_t b; + b.x = __float22bfloat162_rn(a.x); + b.y = __float22bfloat162_rn(a.y); + b.z = __float22bfloat162_rn(a.z); + b.w = __float22bfloat162_rn(a.w); + return b; +} + +/* Scaled and vectorized conversions, for data exchange between high and low + precision domains + + Convention of the scale in API, e.g: FP8_data = Quantization( + High_Precision_data / scale ) s.t. Quantize(HP / scale) => FP8 Dequant(FP8) * + scale => HP + + */ + +using __nv_bfloat16 = __hip_bfloat16; + +// fp8 -> __nv_bfloat16 +template <> +__inline__ __device__ __nv_bfloat16 +scaled_vec_conversion<__nv_bfloat16, uint8_t>(const uint8_t& a, float scale) { + fp8_type f8; + f8.__x = a; + return __float2bfloat16(static_cast(f8) * scale); +} + +// fp8x2 -> __nv_bfloat162 +template <> +__inline__ __device__ __nv_bfloat162 +scaled_vec_conversion<__nv_bfloat162, uint16_t>(const uint16_t& a, + float scale) { + __nv_bfloat162 res; + res.x = scaled_vec_conversion<__nv_bfloat16, uint8_t>((uint8_t)a, scale); + res.y = + scaled_vec_conversion<__nv_bfloat16, uint8_t>((uint8_t)(a >> 8U), scale); + return res; +} + +// fp8x4 -> bf16_4_t +template <> +__inline__ __device__ bf16_4_t +scaled_vec_conversion(const uint32_t& a, float scale) { + bf16_4_t res; + res.x = scaled_vec_conversion<__nv_bfloat162, uint16_t>((uint16_t)a, scale); + res.y = scaled_vec_conversion<__nv_bfloat162, uint16_t>((uint16_t)(a >> 16U), + scale); + return res; +} + +// fp8x8 -> bf16_8_t +template <> +__inline__ __device__ bf16_8_t +scaled_vec_conversion(const uint2& a, float scale) { + bf16_4_t tmp1, tmp2; + tmp1 = scaled_vec_conversion(a.x, scale); + tmp2 = scaled_vec_conversion(a.y, scale); + bf16_8_t res; + res.x = tmp1.x; + res.y = tmp1.y; + res.z = tmp2.x; + res.w = tmp2.y; + return res; +} + +// fp8 -> float +template <> +__inline__ __device__ float scaled_vec_conversion( + const uint8_t& a, float scale) { + fp8_type f8; + f8.__x = a; + return static_cast(f8) * scale; +} + +// fp8x2 -> float2 +template <> +__inline__ __device__ float2 +scaled_vec_conversion(const uint16_t& a, float scale) { + fp8x2_type f8x2; + f8x2.__x = a; + return static_cast(f8x2) * scale; +} + +// fp8x4 -> float4 +template <> +__inline__ __device__ Float4_ +scaled_vec_conversion(const uint32_t& a, const float scale) { + Float4_ res; + res.x = scaled_vec_conversion((uint16_t)a, scale); + res.y = scaled_vec_conversion((uint16_t)(a >> 16U), scale); + return res; +} + +// fp8x4 -> float4 +template <> +__inline__ __device__ float4 +scaled_vec_conversion(const uint32_t& a, float scale) { + Float4_ res = scaled_vec_conversion(a, scale); + return {res.x.x, res.x.y, res.y.x, res.y.y}; +} + +// fp8x8 -> float8 +template <> +__inline__ __device__ Float8_ +scaled_vec_conversion(const uint2& a, float scale) { + Float4_ tmp1, tmp2; + tmp1 = scaled_vec_conversion(a.x, scale); + tmp2 = scaled_vec_conversion(a.y, scale); + Float8_ res; + res.x = tmp1.x; + res.y = tmp1.y; + res.z = tmp2.x; + res.w = tmp2.y; + return res; +} + +// fp8 -> half +template <> +__inline__ __device__ uint16_t +scaled_vec_conversion(const uint8_t& a, float scale) { + __half_raw res; + res.data = scaled_vec_conversion(a, scale); + return res.x; +} + +// fp8x2 -> half2 +template <> +__inline__ __device__ uint32_t +scaled_vec_conversion(const uint16_t& a, float scale) { + union { + __half2_raw h2r; + uint32_t ui32; + } tmp; + tmp.h2r = __hip_cvt_fp8x2_to_halfraw2(a, fp8_type::__default_interpret); + tmp.h2r.x.data *= scale; + tmp.h2r.y.data *= scale; + return tmp.ui32; +} + +// fp8x4 -> half2x2 +template <> +__inline__ __device__ uint2 +scaled_vec_conversion(const uint32_t& a, float scale) { + union { + uint2 u32x2; + uint32_t u32[2]; + } tmp; + tmp.u32[0] = scaled_vec_conversion((uint16_t)a, scale); + tmp.u32[1] = + scaled_vec_conversion((uint16_t)(a >> 16U), scale); + return tmp.u32x2; +} + +// fp8x8 -> half2x4 +template <> +__inline__ __device__ uint4 scaled_vec_conversion(const uint2& a, + float scale) { + union { + uint4 u64x2; + uint2 u64[2]; + } tmp; + tmp.u64[0] = scaled_vec_conversion(a.x, scale); + tmp.u64[1] = scaled_vec_conversion(a.y, scale); + return tmp.u64x2; +} + +// half -> fp8 +template <> +__inline__ __device__ uint8_t +scaled_vec_conversion(const uint16_t& a, float scale) { + __half_raw tmp; + tmp.x = a; + tmp.data /= scale; + return __hip_cvt_halfraw_to_fp8(tmp, fp8_type::__default_saturation, + fp8_type::__default_interpret); +} + +// halfx2 -> fp8x2 +template <> +__inline__ __device__ uint16_t +scaled_vec_conversion(const uint32_t& a, float scale) { + union { + uint32_t ui32; + __half2_raw h2r; + } tmp; + tmp.ui32 = a; + tmp.h2r.x.data /= scale; + tmp.h2r.y.data /= scale; + return __hip_cvt_halfraw2_to_fp8x2(tmp.h2r, fp8_type::__default_saturation, + fp8_type::__default_interpret); +} + +// half2x2 -> fp8x4 +template <> +__inline__ __device__ uint32_t +scaled_vec_conversion(const uint2& a, float scale) { + union { + uint16_t ui16[2]; + uint32_t ui32; + } tmp; + tmp.ui16[0] = scaled_vec_conversion(a.x, scale); + tmp.ui16[1] = scaled_vec_conversion(a.y, scale); + return tmp.ui32; +} + +// half2x4 -> fp8x8 +template <> +__inline__ __device__ uint2 scaled_vec_conversion(const uint4& a, + float scale) { + union { + uint2 ui2[2]; + uint4 ui4; + } tmp; + tmp.ui4 = a; + uint2 res; + res.x = scaled_vec_conversion(tmp.ui2[0], scale); + res.y = scaled_vec_conversion(tmp.ui2[1], scale); + return res; +} + +// bf16 -> fp8 +template <> +__inline__ __device__ uint8_t scaled_vec_conversion( + const __nv_bfloat16& a, float scale) { + return __hip_cvt_float_to_fp8(__bfloat162float(a) / scale, + fp8_type::__default_saturation, + fp8_type::__default_interpret); +} + +// bf16x2 -> fp8x2 +template <> +__inline__ __device__ uint16_t scaled_vec_conversion( + const __nv_bfloat162& a, float scale) { + union { + uint8_t ui8[2]; + uint16_t ui16; + } tmp; + tmp.ui8[0] = scaled_vec_conversion(a.x, scale); + tmp.ui8[1] = scaled_vec_conversion(a.y, scale); + return tmp.ui16; +} + +// bf16x4 -> fp8x4 +template <> +__inline__ __device__ uint32_t +scaled_vec_conversion(const bf16_4_t& a, float scale) { + union { + uint16_t ui16[2]; + uint32_t ui32; + } tmp; + tmp.ui16[0] = scaled_vec_conversion(a.x, scale); + tmp.ui16[1] = scaled_vec_conversion(a.y, scale); + return tmp.ui32; +} + +// bf16x8 -> fp8x8 +template <> +__inline__ __device__ uint2 +scaled_vec_conversion(const bf16_8_t& a, float scale) { + uint2 res; + res.x = scaled_vec_conversion({a.x, a.y}, scale); + res.y = scaled_vec_conversion({a.z, a.w}, scale); + return res; +} + +// float -> fp8 +template <> +__inline__ __device__ uint8_t +scaled_vec_conversion(const float& a, float scale) { + return __hip_cvt_float_to_fp8(a / scale, fp8_type::__default_saturation, + fp8_type::__default_interpret); +} + +// floatx2 -> fp8x2 +template <> +__inline__ __device__ uint16_t +scaled_vec_conversion(const float2& a, float scale) { + return __hip_cvt_float2_to_fp8x2(a / scale, fp8_type::__default_saturation, + fp8_type::__default_interpret); +} + +// floatx4 -> fp8x4 +template <> +__inline__ __device__ uint32_t +scaled_vec_conversion(const float4& a, float scale) { + union { + uint16_t ui16[2]; + uint32_t ui32; + } tmp; + tmp.ui16[0] = scaled_vec_conversion({a.x, a.y}, scale); + tmp.ui16[1] = scaled_vec_conversion({a.z, a.w}, scale); + return tmp.ui32; +} + #endif // ENABLE_FP8 + +template +__inline__ __device__ Tout convert(const Tin& x) { + #ifdef ENABLE_FP8 + if constexpr (kv_dt == Fp8KVCacheDataType::kFp8E4M3) { + return vec_conversion(x); + } + #endif + assert(false); + return {}; // Squash missing return statement warning +} + +template +__inline__ __device__ Tout scaled_convert(const Tin& x, const float scale) { + #ifdef ENABLE_FP8 + if constexpr (kv_dt == Fp8KVCacheDataType::kFp8E4M3) { + return scaled_vec_conversion(x, scale); + } + #endif + assert(false); + return {}; // Squash missing return statement warning +} + + // The following macro is used to dispatch the conversion function based on + // the data type of the key and value cache. The FN is a macro that calls a + // function with template. + #define DISPATCH_BY_KV_CACHE_DTYPE(SRC_DTYPE, KV_DTYPE, FN) \ + if (KV_DTYPE == "auto") { \ + if (SRC_DTYPE == at::ScalarType::Float) { \ + FN(float, float, vllm::Fp8KVCacheDataType::kAuto); \ + } else if (SRC_DTYPE == at::ScalarType::Half) { \ + FN(uint16_t, uint16_t, vllm::Fp8KVCacheDataType::kAuto); \ + } else if (SRC_DTYPE == at::ScalarType::BFloat16) { \ + FN(__nv_bfloat16, __nv_bfloat16, vllm::Fp8KVCacheDataType::kAuto); \ + } else { \ + TORCH_CHECK(false, "Unsupported input type of kv cache: ", SRC_DTYPE); \ + } \ + } else { \ + if (KV_DTYPE == "fp8" || KV_DTYPE == "fp8_e4m3") { \ + if (SRC_DTYPE == at::ScalarType::Float) { \ + FN(float, uint8_t, vllm::Fp8KVCacheDataType::kFp8E4M3); \ + } else if (SRC_DTYPE == at::ScalarType::Half) { \ + FN(uint16_t, uint8_t, vllm::Fp8KVCacheDataType::kFp8E4M3); \ + } else if (SRC_DTYPE == at::ScalarType::BFloat16) { \ + FN(__nv_bfloat16, uint8_t, vllm::Fp8KVCacheDataType::kFp8E4M3); \ + } else { \ + TORCH_CHECK(false, \ + "Unsupported input type of kv cache: ", SRC_DTYPE); \ + } \ + } else { \ + TORCH_CHECK(false, "Unsupported data type of kv cache: ", KV_DTYPE); \ + } \ + } + +} // namespace fp8 +#endif // USE_ROCM +} // namespace vllm diff --git a/tasks/hip2hip/campaign20/wvsplitk/src/quantization/w8a8/fp8/common.cuh b/tasks/hip2hip/campaign20/wvsplitk/src/quantization/w8a8/fp8/common.cuh new file mode 100644 index 00000000..7a385f51 --- /dev/null +++ b/tasks/hip2hip/campaign20/wvsplitk/src/quantization/w8a8/fp8/common.cuh @@ -0,0 +1,62 @@ +#pragma once + +#include "libtorch_stable/quantization/vectorization.cuh" +#include "quantization/utils.cuh" + +#include + +#ifndef USE_ROCM + #include "nvidia/quant_utils.cuh" +#else + #include "amd/quant_utils.cuh" +#endif + +// Determines the preferred FP8 type for the current platform. +// Note that for CUDA this just returns true, +// but on ROCm it will check device props. +static bool is_fp8_ocp() { +#ifndef USE_ROCM + return true; +#else + auto dprops = at::cuda::getCurrentDeviceProperties(); + std::string device_arch = dprops->gcnArchName; + size_t substring = device_arch.find("gfx94"); + return substring == std::string::npos; +#endif +} + +namespace vllm { + +__device__ __forceinline__ float atomicMaxFloat(float* addr, float value) { + float old; + old = (value >= 0) + ? __int_as_float(atomicMax((int*)addr, __float_as_int(value))) + : __uint_as_float( + atomicMin((unsigned int*)addr, __float_as_uint(value))); + + return old; +} + +template +__device__ __forceinline__ fp8_type scaled_fp8_conversion(float const val, + float const scale) { + float x = 0.0f; + if constexpr (is_scale_inverted) { + x = val * scale; + } else { + x = val / scale; + } + + float r = + fmaxf(-quant_type_max_v, fminf(x, quant_type_max_v)); +#ifndef USE_ROCM + // Use hardware cvt instruction for fp8 on nvidia + // Currently only support fp8_type = c10::Float8_e4m3fn + return fp8::vec_conversion(r); +#else + // Use hardware cvt instruction for fp8 on rocm + return fp8::cvt_c10(r); +#endif +} + +} // namespace vllm diff --git a/tasks/hip2hip/campaign20/wvsplitk/src/quantization/w8a8/fp8/nvidia/quant_utils.cuh b/tasks/hip2hip/campaign20/wvsplitk/src/quantization/w8a8/fp8/nvidia/quant_utils.cuh new file mode 100644 index 00000000..421e8092 --- /dev/null +++ b/tasks/hip2hip/campaign20/wvsplitk/src/quantization/w8a8/fp8/nvidia/quant_utils.cuh @@ -0,0 +1,597 @@ +#pragma once + +#include "../../../../attention/attention_dtypes.h" +#include +#include +#include +#include + +namespace vllm { +#ifndef USE_ROCM + +namespace fp8 { + #ifdef ENABLE_FP8 + +template +__inline__ __device__ Tout vec_conversion( + const Tin& x, const __nv_fp8_interpretation_t fp8_type = __NV_E4M3) { + return x; +} + +// float -> c10::Float8_e4m3fn +template <> +__inline__ __device__ c10::Float8_e4m3fn +vec_conversion( + const float& a, const __nv_fp8_interpretation_t fp8_type) { + #if defined(__CUDA_ARCH__) && __CUDA_ARCH__ < 800 + return static_cast(a); + #else + return c10::Float8_e4m3fn(__nv_cvt_float_to_fp8(a, __NV_SATFINITE, fp8_type), + c10::Float8_e4m3fn::from_bits()); + #endif +} + + #if 0 // Disable the following code to reduce the binary size. +// fp8 -> half +template <> +__inline__ __device__ uint16_t vec_conversion( + const uint8_t &a, const __nv_fp8_interpretation_t fp8_type) { + __half_raw res = __nv_cvt_fp8_to_halfraw(a, fp8_type); + return res.x; +} + +// fp8x2 -> half2 +template <> +__inline__ __device__ uint32_t vec_conversion( + const uint16_t &a, const __nv_fp8_interpretation_t fp8_type) { + union { + uint16_t u16[2]; + uint32_t u32; + } tmp; + __half2_raw res = __nv_cvt_fp8x2_to_halfraw2(a, fp8_type); + tmp.u16[0] = res.x; + tmp.u16[1] = res.y; + return tmp.u32; +} + +// fp8x4 -> half2x2 +template <> +__inline__ __device__ uint2 vec_conversion( + const uint32_t &a, const __nv_fp8_interpretation_t fp8_type) { + union { + uint2 u32x2; + uint32_t u32[2]; + } tmp; + tmp.u32[0] = vec_conversion((uint16_t)a, fp8_type); + tmp.u32[1] = + vec_conversion((uint16_t)(a >> 16U), fp8_type); + return tmp.u32x2; +} + +// fp8x8 -> half2x4 +template <> +__inline__ __device__ uint4 vec_conversion( + const uint2 &a, const __nv_fp8_interpretation_t fp8_type) { + union { + uint4 u64x2; + uint2 u64[2]; + } tmp; + tmp.u64[0] = vec_conversion(a.x, fp8_type); + tmp.u64[1] = vec_conversion(a.y, fp8_type); + return tmp.u64x2; +} + +// fp8 -> __nv_bfloat16 +template <> +__inline__ __device__ __nv_bfloat16 vec_conversion<__nv_bfloat16, uint8_t>( + const uint8_t &a, const __nv_fp8_interpretation_t fp8_type) { + // Note there is no direct convert function from fp8 to bf16. + // fp8 -> half + __half_raw res = __nv_cvt_fp8_to_halfraw(a, fp8_type); + // half -> float -> bf16 + float tmp = half_to_float(res.x); + return __float2bfloat16(tmp); +} + +// fp8x2 -> __nv_bfloat162 +template <> +__inline__ __device__ __nv_bfloat162 vec_conversion<__nv_bfloat162, uint16_t>( + const uint16_t &a, const __nv_fp8_interpretation_t fp8_type) { + __nv_bfloat162 res; + res.x = vec_conversion<__nv_bfloat16, uint8_t>((uint8_t)a, fp8_type); + res.y = vec_conversion<__nv_bfloat16, uint8_t>((uint8_t)(a >> 8U), fp8_type); + return res; +} + +// fp8x4 -> bf16_4_t +template <> +__inline__ __device__ bf16_4_t vec_conversion( + const uint32_t &a, const __nv_fp8_interpretation_t fp8_type) { + bf16_4_t res; + res.x = vec_conversion<__nv_bfloat162, uint16_t>((uint16_t)a, fp8_type); + res.y = + vec_conversion<__nv_bfloat162, uint16_t>((uint16_t)(a >> 16U), fp8_type); + return res; +} + +// fp8x8 -> bf16_8_t +template <> +__inline__ __device__ bf16_8_t vec_conversion( + const uint2 &a, const __nv_fp8_interpretation_t fp8_type) { + bf16_4_t tmp1, tmp2; + tmp1 = vec_conversion(a.x, fp8_type); + tmp2 = vec_conversion(a.y, fp8_type); + bf16_8_t res; + res.x = tmp1.x; + res.y = tmp1.y; + res.z = tmp2.x; + res.w = tmp2.y; + return res; +} + +// fp8 -> float +template <> +__inline__ __device__ float +vec_conversion(const uint8_t &a, + const __nv_fp8_interpretation_t fp8_type) { + // fp8 -> half + uint16_t tmp = vec_conversion(a, fp8_type); + // half -> float + return half_to_float(tmp); +} + +// fp8x2 -> float2 +template <> +__inline__ __device__ float2 vec_conversion( + const uint16_t &a, const __nv_fp8_interpretation_t fp8_type) { + // fp8x2 -> half2 + uint32_t tmp = vec_conversion(a, fp8_type); + // half2 -> float2 + return half2_to_float2(tmp); +} + +// fp8x4 -> float4 +template <> +__inline__ __device__ Float4_ vec_conversion( + const uint32_t &a, const __nv_fp8_interpretation_t fp8_type) { + Float4_ res; + res.x = vec_conversion((uint16_t)a, fp8_type); + res.y = vec_conversion((uint16_t)(a >> 16U), fp8_type); + return res; +} + +// fp8x8 -> float8 +template <> +__inline__ __device__ Float8_ vec_conversion( + const uint2 &a, const __nv_fp8_interpretation_t fp8_type) { + Float4_ tmp1, tmp2; + tmp1 = vec_conversion(a.x, fp8_type); + tmp2 = vec_conversion(a.y, fp8_type); + Float8_ res; + res.x = tmp1.x; + res.y = tmp1.y; + res.z = tmp2.x; + res.w = tmp2.y; + return res; +} + +// half -> fp8 +template <> +__inline__ __device__ uint8_t vec_conversion( + const uint16_t &a, const __nv_fp8_interpretation_t fp8_type) { + __half_raw tmp; + tmp.x = a; + __nv_fp8_storage_t res = + __nv_cvt_halfraw_to_fp8(tmp, __NV_SATFINITE, fp8_type); + return (uint8_t)res; +} + +// bf16 -> fp8 +template <> +__inline__ __device__ uint8_t vec_conversion( + const __nv_bfloat16 &a, const __nv_fp8_interpretation_t fp8_type) { + #if defined(__CUDA_ARCH__) && __CUDA_ARCH__ < 800 + assert(false); + #else + __nv_fp8_storage_t res = __nv_cvt_bfloat16raw_to_fp8( + __nv_bfloat16_raw(a), __NV_SATFINITE, fp8_type); + return (uint8_t)res; + #endif +} + +// float -> fp8 +template <> +__inline__ __device__ uint8_t vec_conversion( + const float &a, const __nv_fp8_interpretation_t fp8_type) { + __nv_fp8_storage_t res = __nv_cvt_float_to_fp8(a, __NV_SATFINITE, fp8_type); + return (uint8_t)res; +} + +// fp8x4 -> float4 +template <> +__inline__ __device__ float4 vec_conversion( + const uint32_t &a, const __nv_fp8_interpretation_t fp8_type) { + Float4_ tmp = vec_conversion(a, fp8_type); + float4 res = make_float4(tmp.x.x, tmp.x.y, tmp.y.x, tmp.y.y); + return res; +} + +template <> +__inline__ __device__ uint32_t vec_conversion( + const float2 &a, const __nv_fp8_interpretation_t fp8_type) { + union { + half2 float16; + uint32_t uint32; + }; + + float16 = __float22half2_rn(a); + return uint32; +} + +template <> +__inline__ __device__ uint2 vec_conversion( + const Float4_ &a, const __nv_fp8_interpretation_t fp8_type) { + uint2 b; + float2 val; + val.x = a.x.x; + val.y = a.x.y; + b.x = vec_conversion(val, fp8_type); + + val.x = a.y.x; + val.y = a.y.y; + b.y = vec_conversion(val, fp8_type); + + return b; +} + +template <> +__inline__ __device__ float4 vec_conversion( + const Float4_ &a, const __nv_fp8_interpretation_t fp8_type) { + float4 b; + b.x = a.x.x; + b.y = a.x.y; + b.z = a.y.x; + b.w = a.y.y; + return b; +} + +template <> +__inline__ __device__ uint4 vec_conversion( + const Float8_ &a, const __nv_fp8_interpretation_t fp8_type) { + uint4 b; + b.x = vec_conversion(a.x, fp8_type); + b.y = vec_conversion(a.y, fp8_type); + b.z = vec_conversion(a.z, fp8_type); + b.w = vec_conversion(a.w, fp8_type); + return b; +} + +template <> +__inline__ __device__ __nv_bfloat162 vec_conversion<__nv_bfloat162, float2>( + const float2 &a, const __nv_fp8_interpretation_t fp8_type) { + __nv_bfloat162 b; + from_float(b, a); + return b; +} + +template <> +__inline__ __device__ bf16_4_t vec_conversion( + const Float4_ &a, const __nv_fp8_interpretation_t fp8_type) { + bf16_4_t b; + from_float(b, a); + return b; +} + +template <> +__inline__ __device__ bf16_8_t vec_conversion( + const Float8_ &a, const __nv_fp8_interpretation_t fp8_type) { + bf16_8_t b; + from_float(b, a); + return b; +} + #endif + +/* Scaled and vectorized conversions, for data exchange between high and low + precision domains Convention of the scale in API, e.g: FP8_data = + Quantization( High_Precision_data / scale ) s.t. Quantize(HP / scale) => FP8 + Dequant(FP8) * scale => HP + */ + +template +__inline__ __device__ Tout scaled_vec_conversion( + const Tin& x, const float scale, const __nv_fp8_interpretation_t fp8_type) { + return x; +} + +// fp8 -> half +template <> +__inline__ __device__ uint16_t scaled_vec_conversion( + const uint8_t& a, const float scale, + const __nv_fp8_interpretation_t fp8_type) { + __half_raw tmp = __nv_cvt_fp8_to_halfraw(a, fp8_type); + return float_to_half(half_to_float(tmp.x) * scale); +} + +// fp8x2 -> half2 +template <> +__inline__ __device__ uint32_t scaled_vec_conversion( + const uint16_t& a, const float scale, + const __nv_fp8_interpretation_t fp8_type) { + union { + uint16_t u16[2]; + uint32_t u32; + } tmp; + __half2_raw res = __nv_cvt_fp8x2_to_halfraw2(a, fp8_type); + tmp.u16[0] = float_to_half(half_to_float(res.x) * scale); + tmp.u16[1] = float_to_half(half_to_float(res.y) * scale); + return tmp.u32; +} + +// fp8x4 -> half2x2 +template <> +__inline__ __device__ uint2 scaled_vec_conversion( + const uint32_t& a, const float scale, + const __nv_fp8_interpretation_t fp8_type) { + union { + uint2 u32x2; + uint32_t u32[2]; + } tmp; + tmp.u32[0] = + scaled_vec_conversion((uint16_t)a, scale, fp8_type); + tmp.u32[1] = scaled_vec_conversion((uint16_t)(a >> 16U), + scale, fp8_type); + return tmp.u32x2; +} + +// fp8x8 -> half2x4 +template <> +__inline__ __device__ uint4 +scaled_vec_conversion(const uint2& a, const float scale, + const __nv_fp8_interpretation_t fp8_type) { + union { + uint4 u64x2; + uint2 u64[2]; + } tmp; + tmp.u64[0] = scaled_vec_conversion(a.x, scale, fp8_type); + tmp.u64[1] = scaled_vec_conversion(a.y, scale, fp8_type); + return tmp.u64x2; +} + +// fp8 -> __nv_bfloat16 +template <> +__inline__ __device__ __nv_bfloat16 +scaled_vec_conversion<__nv_bfloat16, uint8_t>( + const uint8_t& a, const float scale, + const __nv_fp8_interpretation_t fp8_type) { + // Note there is no direct convert function from fp8 to bf16. + // fp8 -> half + __half_raw res = __nv_cvt_fp8_to_halfraw(a, fp8_type); + // half -> float -> bf16 + float tmp = half_to_float(res.x); + return __float2bfloat16(tmp * scale); +} + +// fp8x2 -> __nv_bfloat162 +template <> +__inline__ __device__ __nv_bfloat162 +scaled_vec_conversion<__nv_bfloat162, uint16_t>( + const uint16_t& a, const float scale, + const __nv_fp8_interpretation_t fp8_type) { + __nv_bfloat162 res; + res.x = scaled_vec_conversion<__nv_bfloat16, uint8_t>((uint8_t)a, scale, + fp8_type); + res.y = scaled_vec_conversion<__nv_bfloat16, uint8_t>((uint8_t)(a >> 8U), + scale, fp8_type); + return res; +} + +// fp8x4 -> bf16_4_t +template <> +__inline__ __device__ bf16_4_t scaled_vec_conversion( + const uint32_t& a, const float scale, + const __nv_fp8_interpretation_t fp8_type) { + bf16_4_t res; + res.x = scaled_vec_conversion<__nv_bfloat162, uint16_t>((uint16_t)a, scale, + fp8_type); + res.y = scaled_vec_conversion<__nv_bfloat162, uint16_t>((uint16_t)(a >> 16U), + scale, fp8_type); + return res; +} + +// fp8x8 -> bf16_8_t +template <> +__inline__ __device__ bf16_8_t scaled_vec_conversion( + const uint2& a, const float scale, + const __nv_fp8_interpretation_t fp8_type) { + bf16_4_t tmp1, tmp2; + tmp1 = scaled_vec_conversion(a.x, scale, fp8_type); + tmp2 = scaled_vec_conversion(a.y, scale, fp8_type); + bf16_8_t res; + res.x = tmp1.x; + res.y = tmp1.y; + res.z = tmp2.x; + res.w = tmp2.y; + return res; +} + +// fp8 -> float +template <> +__inline__ __device__ float scaled_vec_conversion( + const uint8_t& a, const float scale, + const __nv_fp8_interpretation_t fp8_type) { + // fp8 -> half + __half_raw res = __nv_cvt_fp8_to_halfraw(a, fp8_type); + uint16_t tmp = res.x; + + // half -> float + return half_to_float(tmp) * scale; +} + +// fp8x2 -> float2 +template <> +__inline__ __device__ float2 scaled_vec_conversion( + const uint16_t& a, const float scale, + const __nv_fp8_interpretation_t fp8_type) { + // fp8x2 -> half2 + uint32_t tmp = scaled_vec_conversion(a, scale, fp8_type); + // half2 -> float2 + return half2_to_float2(tmp); +} + +// fp8x4 -> float4 +template <> +__inline__ __device__ Float4_ scaled_vec_conversion( + const uint32_t& a, const float scale, + const __nv_fp8_interpretation_t fp8_type) { + Float4_ res; + res.x = scaled_vec_conversion((uint16_t)a, scale, fp8_type); + res.y = scaled_vec_conversion((uint16_t)(a >> 16U), scale, + fp8_type); + return res; +} + +// fp8x8 -> float8 +template <> +__inline__ __device__ Float8_ scaled_vec_conversion( + const uint2& a, const float scale, + const __nv_fp8_interpretation_t fp8_type) { + Float4_ tmp1, tmp2; + tmp1 = scaled_vec_conversion(a.x, scale, fp8_type); + tmp2 = scaled_vec_conversion(a.y, scale, fp8_type); + Float8_ res; + res.x = tmp1.x; + res.y = tmp1.y; + res.z = tmp2.x; + res.w = tmp2.y; + return res; +} + +// half -> fp8 +template <> +__inline__ __device__ uint8_t scaled_vec_conversion( + const uint16_t& a, const float scale, + const __nv_fp8_interpretation_t fp8_type) { + __nv_fp8_storage_t res = + __nv_cvt_float_to_fp8(half_to_float(a) / scale, __NV_SATFINITE, fp8_type); + return (uint8_t)res; +} + +// bf16 -> fp8 +template <> +__inline__ __device__ uint8_t scaled_vec_conversion( + const __nv_bfloat16& a, const float scale, + const __nv_fp8_interpretation_t fp8_type) { + #if defined(__CUDA_ARCH__) && __CUDA_ARCH__ < 800 + assert(false); + #else + __nv_fp8_storage_t res = __nv_cvt_float_to_fp8(__bfloat162float(a) / scale, + __NV_SATFINITE, fp8_type); + return (uint8_t)res; + #endif + __builtin_unreachable(); // Suppress missing return statement warning +} + +// float -> fp8 +template <> +__inline__ __device__ uint8_t scaled_vec_conversion( + const float& a, const float scale, + const __nv_fp8_interpretation_t fp8_type) { + __nv_fp8_storage_t res = + __nv_cvt_float_to_fp8(a / scale, __NV_SATFINITE, fp8_type); + return (uint8_t)res; +} + +// fp8x4 -> float4 +template <> +__inline__ __device__ float4 scaled_vec_conversion( + const uint32_t& a, const float scale, + const __nv_fp8_interpretation_t fp8_type) { + Float4_ tmp = scaled_vec_conversion(a, scale, fp8_type); + float4 res = make_float4(tmp.x.x, tmp.x.y, tmp.y.x, tmp.y.y); + return res; +} + #endif // ENABLE_FP8 + +template +__inline__ __device__ Tout convert(const Tin& x) { + #if 0 // Disable the following code to reduce the binary size. + if constexpr (kv_dt == Fp8KVCacheDataType::kFp8E4M3) { + return vec_conversion(x, __NV_E4M3); + } else if constexpr (kv_dt == Fp8KVCacheDataType::kFp8E5M2) { + return vec_conversion(x, __NV_E5M2); + } + #endif + assert(false); + __builtin_unreachable(); // Suppress missing return statement warning +} + +template +__inline__ __device__ Tout scaled_convert(const Tin& x, const float scale) { + #ifdef ENABLE_FP8 + if constexpr (kv_dt == Fp8KVCacheDataType::kFp8E4M3) { + return scaled_vec_conversion(x, scale, __NV_E4M3); + } else if constexpr (kv_dt == Fp8KVCacheDataType::kFp8E5M2) { + return scaled_vec_conversion(x, scale, __NV_E5M2); + } + #endif + assert(false); + __builtin_unreachable(); // Suppress missing return statement warning +} + + // The following macro is used to dispatch the conversion function based on + // the data type of the key and value cache. The FN is a macro that calls a + // function with template. + #define DISPATCH_BY_KV_CACHE_DTYPE(SRC_DTYPE, KV_DTYPE, FN) \ + if (KV_DTYPE == "auto") { \ + if (SRC_DTYPE == at::ScalarType::Float) { \ + FN(float, float, vllm::Fp8KVCacheDataType::kAuto); \ + } else if (SRC_DTYPE == at::ScalarType::Half) { \ + FN(uint16_t, uint16_t, vllm::Fp8KVCacheDataType::kAuto); \ + } else if (SRC_DTYPE == at::ScalarType::BFloat16) { \ + FN(__nv_bfloat16, __nv_bfloat16, vllm::Fp8KVCacheDataType::kAuto); \ + } else { \ + TORCH_CHECK(false, "Unsupported input type of kv cache: ", SRC_DTYPE); \ + } \ + } else { \ + if (KV_DTYPE == "fp8" || KV_DTYPE == "fp8_e4m3") { \ + if (SRC_DTYPE == at::ScalarType::Float) { \ + FN(float, uint8_t, vllm::Fp8KVCacheDataType::kFp8E4M3); \ + } else if (SRC_DTYPE == at::ScalarType::Half) { \ + FN(uint16_t, uint8_t, vllm::Fp8KVCacheDataType::kFp8E4M3); \ + } else if (SRC_DTYPE == at::ScalarType::BFloat16) { \ + FN(__nv_bfloat16, uint8_t, vllm::Fp8KVCacheDataType::kFp8E4M3); \ + } else { \ + TORCH_CHECK(false, \ + "Unsupported input type of kv cache: ", SRC_DTYPE); \ + } \ + } else if (KV_DTYPE == "fp8_e5m2") { \ + if (SRC_DTYPE == at::ScalarType::Float) { \ + FN(float, uint8_t, vllm::Fp8KVCacheDataType::kFp8E5M2); \ + } else if (SRC_DTYPE == at::ScalarType::Half) { \ + FN(uint16_t, uint8_t, vllm::Fp8KVCacheDataType::kFp8E5M2); \ + } else if (SRC_DTYPE == at::ScalarType::BFloat16) { \ + FN(__nv_bfloat16, uint8_t, vllm::Fp8KVCacheDataType::kFp8E5M2); \ + } else { \ + TORCH_CHECK(false, \ + "Unsupported input type of kv cache: ", SRC_DTYPE); \ + } \ + } else if (KV_DTYPE == "fp8_ds_mla") { \ + if (SRC_DTYPE == at::ScalarType::Float) { \ + FN(float, uint8_t, vllm::Fp8KVCacheDataType::kFp8E4M3); \ + } else if (SRC_DTYPE == at::ScalarType::Half) { \ + FN(uint16_t, uint8_t, vllm::Fp8KVCacheDataType::kFp8E4M3); \ + } else if (SRC_DTYPE == at::ScalarType::BFloat16) { \ + FN(__nv_bfloat16, uint8_t, vllm::Fp8KVCacheDataType::kFp8E4M3); \ + } else { \ + TORCH_CHECK(false, \ + "Unsupported input type of kv cache: ", SRC_DTYPE); \ + } \ + } else { \ + TORCH_CHECK(false, "Unsupported data type of kv cache: ", KV_DTYPE); \ + } \ + } + +} // namespace fp8 +#endif // not USE_ROCM +} // namespace vllm diff --git a/tasks/hip2hip/campaign20/wvsplitk/src/rocm/skinny_gemms.cu b/tasks/hip2hip/campaign20/wvsplitk/src/rocm/skinny_gemms.cu new file mode 100644 index 00000000..60e10e53 --- /dev/null +++ b/tasks/hip2hip/campaign20/wvsplitk/src/rocm/skinny_gemms.cu @@ -0,0 +1,2327 @@ +#include +#include +#include + +#include +#include +#include + +#include +#include + +#include "../cuda_compat.h" +#include "dispatch_utils.h" +#include "quantization/w8a8/fp8/common.cuh" +#include "core/batch_invariant.hpp" + +// TODO(rasmith): The kernels in this file are susceptible to integer overflow +// issues, do not take strides, and are unable to handle PyTorch tensors that +// return is_contiguous() as False (the tensors may actually be contiguous +// in memory). +// +// However, it may be possible to fix these kernels to handle both issues. + +#if defined(__HIPCC__) && \ + (defined(__gfx90a__) || defined(__gfx942__) || defined(__gfx950__)) + #define __HIP__GFX9__ +#endif + +#if defined(__HIPCC__) && \ + (defined(__gfx1100__) || defined(__gfx1101__) || defined(__gfx1150__) || \ + defined(__gfx1151__) || defined(__gfx1200__) || defined(__gfx1201__)) + #define __HIP__GFX1X__ +#endif + +#if defined(__HIPCC__) && (defined(__gfx1200__) || defined(__gfx1201__)) + #define __HIP__GFX12__ +#endif + +#if defined(__HIPCC__) && (defined(__gfx942__) || defined(__gfx950__)) + #define __HIP__MI3XX__ +#endif + +#if defined(__gfx950__) + #define LDS_SIZE 160 * 1024 +#else + #define LDS_SIZE 64 * 1024 +#endif + +int get_lds_size() { + static const int result = [] { + const auto* dprops = at::cuda::getCurrentDeviceProperties(); + const std::string device_arch = dprops->gcnArchName; + return device_arch.find("gfx95") == std::string::npos ? 64 * 1024 + : 160 * 1024; + }(); + return result; +} + +bool on_gfx1x() { + static const bool result = [] { + const auto* dprops = at::cuda::getCurrentDeviceProperties(); + const std::string device_arch = dprops->gcnArchName; + return device_arch.find("gfx11") != std::string::npos || + device_arch.find("gfx12") != std::string::npos; + }(); + return result; +} + +bool on_gfx12() { + static const bool result = [] { + const auto* dprops = at::cuda::getCurrentDeviceProperties(); + const std::string device_arch = dprops->gcnArchName; + return device_arch.find("gfx12") != std::string::npos; + }(); + return result; +} + +#if defined(NDEBUG) + #undef NDEBUG + #include + #define UNREACHABLE_CODE assert(false); + #define NDEBUG +#else + #define UNREACHABLE_CODE assert(false); +#endif + +template +struct scalar {}; + +template +struct scalar2 {}; + +template +__device__ __forceinline__ float2 __s22float2(T v); + +template +__device__ __forceinline__ T __float2s(float v); + +template +__device__ __forceinline__ T __float22s2_rn(float2 v); + +// Definitions and cvt functions for fp16 +template <> +struct scalar { + using type = half; +}; + +template <> +struct scalar2 { + using type = __half2; +}; + +template <> +__device__ __forceinline__ half __float2s(float v) { + return __float2half(v); +} + +template <> +__device__ __forceinline__ float2 __s22float2(__half2 v) { + return __half22float2(v); +} + +template <> +__device__ __forceinline__ __half2 __float22s2_rn(float2 v) { + return __float22half2_rn(v); +} + +// Definitions and cvt functions for bf16 +template <> +struct scalar { + using type = __hip_bfloat16; +}; + +template <> +struct scalar2 { + using type = __hip_bfloat162; +}; + +template <> +__device__ __forceinline__ __hip_bfloat16 __float2s(float v) { + return __float2bfloat16(v); +} + +template <> +__device__ __forceinline__ float2 __s22float2(__hip_bfloat162 v) { + return __bfloat1622float2(v); +} + +template <> +__device__ __forceinline__ __hip_bfloat162 __float22s2_rn(float2 v) { + return __float22bfloat162_rn(v); +} + +template +__device__ __forceinline__ T loadnt(T* addr) { + return __builtin_nontemporal_load(addr); +} + +__device__ __forceinline__ float4 load_ntmprl(const float4* addr) { + auto addr_alias = reinterpret_cast(addr); + auto dat0 = loadnt(addr_alias); + auto dat1 = loadnt(addr_alias + 1); + auto dat2 = loadnt(addr_alias + 2); + auto dat3 = loadnt(addr_alias + 3); + return make_float4(dat0, dat1, dat2, dat3); +} + +// TBlock fetches entire rows of A, and entire col of B (K dimension); assume +// N=1 for time being grid is M/A_NUM_ROWS blocks +template +__global__ void LLGemm1_kernel(const scalar_t* in_a, const scalar_t* in_b, + scalar_t* out_c, const int K) { + using scalar2_t = typename scalar2::type; + auto af4 = reinterpret_cast(in_a); + auto bf4 = reinterpret_cast(in_b); + auto c = reinterpret_cast(out_c); + __shared__ float red_smem[NUM_A_ROWS_PER_BLOCK][WARP_SIZE]; + const int row_addr = blockIdx.x * NUM_A_ROWS_PER_BLOCK * K / 8; + const int threadid = threadIdx.x; + const int warp = threadIdx.x / WARP_SIZE; + const int lane = threadIdx.x % WARP_SIZE; + const int num_warps = blockDim.x / WARP_SIZE; + const int qwarpid = threadid / 16; + const int qthreadid = threadid % 16; + float4 rowA_elem4[NUM_A_ROWS_PER_BLOCK]; + scalar2_t colB_elem4x, colB_elem4y, colB_elem4z, colB_elem4w; + float acc[NUM_A_ROWS_PER_BLOCK]; + scalar2_t acch2; + scalar2_t oval; + + // As we later use warp shuffle operations, we may have more threads in the + // block than the actual available data, hence the if guard here. + if (threadid * 8 < K) { +#pragma unroll + for (int i = 0; i < NUM_A_ROWS_PER_BLOCK; i++) { + // rowA_elem4[i] holds 8 * half numbers seen as a single float4. + rowA_elem4[i] = load_ntmprl(&af4[row_addr + threadid + K / 8 * i]); + } + colB_elem4x = bf4[threadid * 4 + 0]; + colB_elem4y = bf4[threadid * 4 + 1]; + colB_elem4z = bf4[threadid * 4 + 2]; + colB_elem4w = bf4[threadid * 4 + 3]; + } + + scalar2_t Af2; + float2 S; + + auto Ah2ptr = reinterpret_cast(&rowA_elem4); + scalar2_t* ah2lptr; + +#pragma unroll + for (int i = 0; i < NUM_A_ROWS_PER_BLOCK; i++) { + // Multiply-add on 8 scalar_t. + ah2lptr = Ah2ptr + i * 4; + Af2 = *(ah2lptr); + acch2 = __hmul2(Af2, colB_elem4x); + Af2 = *(ah2lptr + 1); + acch2 = __hfma2(Af2, colB_elem4y, acch2); + Af2 = *(ah2lptr + 2); + acch2 = __hfma2(Af2, colB_elem4z, acch2); + Af2 = *(ah2lptr + 3); + acch2 = __hfma2(Af2, colB_elem4w, acch2); + S = __s22float2(acch2); + + // See comment above concerning the if guard. + acc[i] = (threadid * 8 < K ? S.x + S.y : 0.f); + } + +// all reduce across warp. +#pragma unroll + for (int mask = WARP_SIZE / 2; mask >= 1; mask /= 2) { +#pragma unroll + for (int i = 0; i < NUM_A_ROWS_PER_BLOCK; i++) { + acc[i] += __shfl_xor(acc[i], mask); + } + } + + // Warp leaders store the data to shared memory. + if (lane < NUM_A_ROWS_PER_BLOCK) { + red_smem[lane][warp] = acc[lane]; + } + + // Make sure the data is in shared memory. + __syncthreads(); + + if (qwarpid < NUM_A_ROWS_PER_BLOCK) { + acc[qwarpid] = qthreadid < num_warps ? red_smem[qwarpid][qthreadid] : 0.f; +#pragma unroll + for (int mask = 16 / 2; mask >= 1; mask /= 2) { + acc[qwarpid] += __shfl_xor(acc[qwarpid], mask); + } + float oval2 = __shfl_xor(acc[qwarpid], 16); + + if (lane % 32 == 0) { + oval = __float22s2_rn(make_float2(acc[qwarpid], oval2)); + c[blockIdx.x * NUM_A_ROWS_PER_BLOCK / 2 + qwarpid / 2] = oval; + } + } +} + +torch::Tensor LLMM1(at::Tensor& in_a, at::Tensor& in_b, + const int64_t rows_per_block) { + auto M = in_a.size(0); + auto K = in_a.size(1); + auto N = in_b.size(0); + + TORCH_CHECK(N == 1, "Row number of activation tensor must be 1."); + TORCH_CHECK(in_a.dtype() == in_b.dtype()); + TORCH_CHECK(in_b.dtype() == torch::kFloat16 || + in_b.dtype() == torch::kBFloat16); + + auto out_c = torch::empty( + {N, M}, torch::TensorOptions().dtype(in_b.dtype()).device(in_b.device())); + + // NUM_TREADS need to be a multiple of WARP_SIZE, as we are using warp shuffle + // operations. + const int NUM_THREADS = + max(rows_per_block * 16, + K * 2 / 16 % WARP_SIZE == 0 + ? K * 2 / 16 + : K * 2 / 16 + (WARP_SIZE - K * 2 / 16 % WARP_SIZE)); + + int NUM_BLOCKS = M / rows_per_block; + + const at::cuda::OptionalCUDAGuard device_guard(device_of(in_b)); + const cudaStream_t stream = at::cuda::getCurrentCUDAStream(); + + // call the kernel function... + AT_DISPATCH_REDUCED_FLOATING_TYPES(in_b.scalar_type(), "LLGemm1", [&] { + auto a_ptr = in_a.data_ptr(); + auto b_ptr = in_b.data_ptr(); + auto c_ptr = out_c.data_ptr(); + if (rows_per_block == 2) { + LLGemm1_kernel + <<>>(a_ptr, b_ptr, c_ptr, K); + } else if (rows_per_block == 4) { + LLGemm1_kernel + <<>>(a_ptr, b_ptr, c_ptr, K); + } else if (rows_per_block == 8) { + LLGemm1_kernel + <<>>(a_ptr, b_ptr, c_ptr, K); + } else if (rows_per_block == 16) { + LLGemm1_kernel + <<>>(a_ptr, b_ptr, c_ptr, K); + } else { + NUM_BLOCKS = M / 4; + LLGemm1_kernel + <<>>(a_ptr, b_ptr, c_ptr, K); + } + }); + + return out_c; +} + +#if defined(__HIP__GFX9__) && !defined(__HIP__GFX1X__) + #define DOT2C(V0, V2, V3) \ + if constexpr (std::is_same_v) { \ + asm("v_dot2c_f32_f16 %0, %2, %3" \ + : "=v"(V0) \ + : "0"(V0), "v"(V2), "v"(V3)); \ + } else if constexpr (std::is_same_v) { \ + float2 s = __bfloat1622float2(*((__hip_bfloat162*)(&(V2)))) * \ + __bfloat1622float2(*((__hip_bfloat162*)(&(V3)))); \ + V0 += (s.x + s.y); \ + } +#elif defined(__HIP__GFX1X__) + // gfx1x: v_dot2_f32_f16 (VOP3-P, dot10-insts, available on gfx11+gfx12) + #define DOT2C(V0, V2, V3) \ + if constexpr (std::is_same_v) { \ + asm("v_dot2_f32_f16 %0, %1, %2, %0" : "+v"(V0) : "v"(V2), "v"(V3)); \ + } else if constexpr (std::is_same_v) { \ + float2 s = __bfloat1622float2(*((__hip_bfloat162*)(&(V2)))) * \ + __bfloat1622float2(*((__hip_bfloat162*)(&(V3)))); \ + V0 += (s.x + s.y); \ + } +#endif + +// To avoid LLVM silently upcasting to double +__device__ inline unsigned int min__(uint32_t a, uint32_t b) { + return min(a, b); +} + +#if defined(__HIP__GFX9__) || defined(__HIP__GFX1X__) +// This version targets cases where A[] fits LDS capacity +template +__global__ void __launch_bounds__(WvPrGrp* THRDS) + wvSplitK_hf_sml_(const int K, const int Kbp, const int Kap, const int M, + const int Bx, const int By, const scalar_t* B, + const scalar_t* __restrict__ A, + const scalar_t* __restrict__ BIAS, scalar_t* C, + const int _WvPrGrp, const int CuCount) { + constexpr int max_lds_len = LDS_SIZE / 2; + #if defined(__HIP__MI3XX__) + constexpr bool use_mfma = (std::is_same_v); + #else + constexpr bool use_mfma = false; + #endif + using scalar8 = + __attribute__((__vector_size__((A_CHUNK / 2) * sizeof(float)))) float; + using half4 = + __attribute__((__vector_size__((A_CHUNK / 2) * sizeof(__bf16)))) __bf16; + union bigType { + scalar_t h[A_CHUNK]; + float f[A_CHUNK / 2]; + float2 f2[A_CHUNK / 4]; + double d[A_CHUNK / 4]; + half4 h4[A_CHUNK / 4]; + scalar8 h8; + }; + + //---------------------------------------------------- + // Reserving 64/160 KB of LDS to have 1 WG / CU + // Goal is to bring the activation matrix A to the LDS + // and use it across the lifetime of the work group + // TODO: When activation matrix is larger than 64 KB + // then this is not going to work! + //---------------------------------------------------- + __shared__ scalar_t s[max_lds_len]; + + //---------------------------------------------------- + // Fetch the activation matrix to LDS + // Loop iteration: + // - Each thread (lane) is fetching 8 elements (A_Chunk) + // - Each wave will fetch 64*8=> 512 elements + // - Each WG will fetch 512 * 16 => 8K elements + // - Then the WG will move to another 8 K elements + // TODO: Logic below will only work when K is multiple of 8 + //---------------------------------------------------- + for (uint32_t k = (threadIdx.y * THRDS + threadIdx.x) * A_CHUNK; + k < min__(Kap * N, max_lds_len); k += THRDS * WvPrGrp * A_CHUNK) { + #if defined(__gfx950__) + __builtin_amdgcn_global_load_lds((int*)(&A[k]), (int*)(&s[k]), 16, 0, 0); + #else + *((bigType*)(&s[k])) = *((bigType*)(&A[k])); + #endif + } + __syncthreads(); + + if (threadIdx.y >= _WvPrGrp) return; + + uint32_t m = (blockIdx.x * _WvPrGrp + (threadIdx.y % _WvPrGrp)) * YTILE; + + //---------------------------------------------------- + // Each wave works on a single column of weight matrix. + // There are 16 waves per WG, and hence, each WG is + // working on 16 columns of weight matrix. Moreover, + // we tile in column direction by YTILE, so when YTILE=1 + // the above math is right, however, when YTILE=2 then + // each wave will be working on 2 columns and WG will + // be working on 32 columns. + // + // Top level loop that makes WGs persistent! + // - WGs iterates across columns of weight matrix + // - Each wave within WG works on a given column(s) + // - After completing first set of columns, WGs start + // working on the next set of available columns + //---------------------------------------------------- + while (m < M) { + //---------------------------------------------------- + // 'sum' accumulates the matrix A x B computation + // split across 64 lanes. + // + // YTILE represents how many column of weight matrix + // are being worked on by each wave. + //---------------------------------------------------- + float sum[N][YTILE] = {}; + scalar8 sum4[N][YTILE] = {}; + + for (uint32_t k1 = 0; k1 < K; k1 += THRDS * A_CHUNK * UNRL) { + bigType bigA[N][UNRL] = {}; + bigType bigB[YTILE][UNRL]; + // Fetch the weight matrix from memory! + #pragma unroll + for (uint32_t k2 = 0; k2 < UNRL; k2++) { + uint32_t k = k1 + k2 * THRDS * A_CHUNK; + uint32_t k_ = k + threadIdx.x * A_CHUNK; + const scalar_t* B_ = &B[min__(k_, K - A_CHUNK)]; + for (int y = 0; y < YTILE; y++) + bigB[y][k2].h8 = (loadnt((scalar8*)(&B_[min__(y + m, M - 1) * Kbp]))); + } + + // Fetch activation matrix from either just LDS or from both LDS / memory + #pragma unroll + for (uint32_t k2 = 0; k2 < UNRL; k2++) { + uint32_t k = k1 + k2 * THRDS * A_CHUNK; + uint32_t k_ = k + threadIdx.x * A_CHUNK; + if (k_ >= K) break; + for (int n = 0; n < N; n++) { + bigA[n][k2] = *((const bigType*)(&(s[k_ + Kap * n]))); + } + } + + // Do the matrix multiplication in interleaved manner + for (uint32_t k2 = 0; k2 < UNRL; k2++) { + for (uint32_t n = 0; n < N; n++) { + for (int y = 0; y < YTILE; y++) { + if constexpr (!use_mfma) + for (uint32_t b = 0; b < A_CHUNK / 2; b++) { + DOT2C(sum[n][y], bigA[n][k2].f[b], bigB[y][k2].f[b]) + } + else + for (uint32_t b = 0; b < A_CHUNK / 4; b++) + sum4[n][y] = __builtin_amdgcn_mfma_f32_4x4x4bf16_1k( + bigA[n][k2].h4[b], bigB[y][k2].h4[b], sum4[n][y], 0, 0, 0); + } + } + } + } + __builtin_amdgcn_sched_barrier(0); + //---------------------------------------------------- + // Final reduction step using shuffle + //---------------------------------------------------- + if constexpr (!use_mfma) { + for (int n = 0; n < N; n++) { + for (int y = 0; y < YTILE; y++) { + sum[n][y] += __builtin_amdgcn_mov_dpp(sum[n][y], 0x118, 0xf, 0xf, + 1); // row_shr8 + sum[n][y] += __builtin_amdgcn_mov_dpp(sum[n][y], 0x114, 0xf, 0xf, + 1); // row_shr4 + sum[n][y] += __builtin_amdgcn_mov_dpp(sum[n][y], 0x112, 0xf, 0xf, + 1); // row_shr2 + sum[n][y] += __builtin_amdgcn_mov_dpp(sum[n][y], 0x111, 0xf, 0xf, + 1); // row_shr1 + #if defined(__HIP__GFX9__) + sum[n][y] += __builtin_amdgcn_mov_dpp(sum[n][y], 0x142, 0xf, 0xf, + 1); // ROW_BCAST15 + sum[n][y] += __builtin_amdgcn_mov_dpp(sum[n][y], 0x143, 0xf, 0xf, + 1); // ROW_BCAST31 + #else + sum[n][y] += __shfl_xor(sum[n][y], 16); + #endif + } + } + + if (threadIdx.x == (THRDS - 1)) { + scalar_t biases[N][YTILE] = {}; + if (BIAS) + for (int n = 0; n < N; n++) { + for (int y = 0; y < YTILE; y++) { + biases[n][y] = BIAS[(m + y) % Bx + (n % By) * Bx]; + } + } + for (int n = 0; n < N; n++) { + for (int y = 0; y < YTILE; y++) { + if constexpr (std::is_same_v) { + sum[n][y] += __half2float(biases[n][y]); + } else if constexpr (std::is_same_v) { + sum[n][y] += __bfloat162float(biases[n][y]); + } + C[m + y + n * M] = __float2s(sum[n][y]); + } + } + } + } else { + #ifdef __HIP__GFX9__ + #pragma unroll + for (int n = 0; n < N; n++) { + #pragma unroll + for (int y = 0; y < YTILE; y++) { + /*float accm1 = 0; + for (int i=0; i<64; i++) + accm1 += __shfl(sum4[n][y][i%4], i); + sum4[n][y][0] = accm1;*/ + float accm = sum4[n][y][0]; + accm += __builtin_amdgcn_mov_dpp(sum4[n][y][1], 0x101, 0xf, 0xf, + 1); // row_shl1 + accm += __builtin_amdgcn_mov_dpp(sum4[n][y][2], 0x102, 0xf, 0xf, + 1); // row_shl2 + accm += __builtin_amdgcn_mov_dpp(sum4[n][y][3], 0x103, 0xf, 0xf, + 1); // row_shl3 + accm += __builtin_amdgcn_mov_dpp(accm, 0x104, 0xf, 0xf, + 1); // row_shl4 + accm += __builtin_amdgcn_mov_dpp(accm, 0x108, 0xf, 0xf, + 1); // row_shl8 + accm = __builtin_amdgcn_mov_dpp(accm, 0x11f, 0xf, 0xf, + 1); // row_shr15 + accm += __builtin_amdgcn_mov_dpp(accm, 0x142, 0xf, 0xf, + 1); // ROW_BCAST15 + accm += __builtin_amdgcn_mov_dpp(accm, 0x143, 0xf, 0xf, + 1); // ROW_BCAST31 + + sum4[n][y][0] = accm; + } + } + if (threadIdx.x == (THRDS - 1)) { + scalar_t biases[N][YTILE] = {}; + if (BIAS) + for (int n = 0; n < N; n++) { + for (int y = 0; y < YTILE; y++) { + biases[n][y] = BIAS[(m + y) % Bx + (n % By) * Bx]; + } + } + for (int n = 0; n < N; n++) { + for (int y = 0; y < YTILE; y++) { + sum4[n][y][0] += __bfloat162float(biases[n][y]); + C[m + y + n * M] = __float2bfloat16(sum4[n][y][0]); + } + } + } + #endif // __HIP__GFX9__ (MFMA path) + } + m += CuCount * _WvPrGrp * YTILE; + } +} +#else +template +__global__ void wvSplitK_hf_sml_(const int K, const int Kbp, const int Kap, + const int M, const int Bx, const int By, + const scalar_t* B, + const scalar_t* __restrict__ A, + const scalar_t* __restrict__ BIAS, scalar_t* C, + const int _WvPrGrp, const int CuCount) { + UNREACHABLE_CODE +} +#endif + +#if defined(__HIP__GFX9__) || defined(__HIP__GFX1X__) +// This version targets cases where A[] marginally exceeds LDS capacity +template +__global__ void __launch_bounds__(WvPrGrp* THRDS) + wvSplitK_hf_(const int K, const int Kbp, const int Kap, const int M, + const int Bx, const int By, const scalar_t* B, + const scalar_t* __restrict__ A, + const scalar_t* __restrict__ BIAS, scalar_t* C, + const int _WvPrGrp, const int CuCount) { + constexpr int max_lds_len = LDS_SIZE / 2; + #if defined(__HIP__MI3XX__) + constexpr bool use_mfma = (std::is_same_v); + #else + constexpr bool use_mfma = false; + #endif + + using scalar8 = + __attribute__((__vector_size__((A_CHUNK / 2) * sizeof(float)))) float; + using half4 = + __attribute__((__vector_size__((A_CHUNK / 2) * sizeof(__bf16)))) __bf16; + union bigType { + scalar_t h[A_CHUNK]; + float f[A_CHUNK / 2]; + float2 f2[A_CHUNK / 4]; + double d[A_CHUNK / 4]; + half4 h4[A_CHUNK / 4]; + scalar8 h8; + }; + + __shared__ scalar_t s[max_lds_len]; + + //---------------------------------------------------- + // Computation of columns that need to be committed to memory! + //---------------------------------------------------- + uint32_t commitColumn[YTILE]; + for (uint32_t i = 0; i < YTILE; i++) { + commitColumn[i] = 1; + } + + uint32_t m = (blockIdx.x * _WvPrGrp + threadIdx.y) * YTILE; + + // Check whether there will be fragmentation! + // This will happen only for the last wave! + if (m < M && (m + YTILE) >= M) { + uint32_t startColumn = M - YTILE; + for (uint32_t i = 0; i < (m - startColumn); i++) { + commitColumn[i] = 0; + } + m = startColumn; + } + + for (uint32_t k = (threadIdx.y * THRDS + threadIdx.x) * A_CHUNK; + k < min__(Kap * N, max_lds_len); k += THRDS * WvPrGrp * A_CHUNK) { + #if defined(__gfx950__) + __builtin_amdgcn_global_load_lds((int*)(&A[k]), (int*)(&s[k]), 16, 0, 0); + #else + *((bigType*)(&s[k])) = *((bigType*)(&A[k])); + #endif + } + + __syncthreads(); + + if (threadIdx.y >= _WvPrGrp) return; + + while (m < M) { + float sum[N][YTILE] = {}; + scalar8 sum4[N][YTILE] = {}; + + for (uint32_t k1 = 0; k1 < K; k1 += THRDS * A_CHUNK * UNRL) { + bigType bigA[N][UNRL] = {}; + bigType bigB[YTILE][UNRL]; + // Fetch the weight matrix from memory! + #pragma unroll + for (uint32_t k2 = 0; k2 < UNRL; k2++) { + uint32_t k = k1 + k2 * THRDS * A_CHUNK; + uint32_t k_ = k + threadIdx.x * A_CHUNK; + const scalar_t* B_ = &B[min__(k_, K - A_CHUNK)]; + for (int y = 0; y < YTILE; y++) + bigB[y][k2].h8 = (loadnt((scalar8*)(&B_[min__(y + m, M - 1) * Kbp]))); + } + + // Fetch activation matrix from either just LDS or from both LDS / memory + #pragma unroll + for (uint32_t k2 = 0; k2 < UNRL; k2++) { + uint32_t k = k1 + k2 * THRDS * A_CHUNK; + uint32_t k_ = k + threadIdx.x * A_CHUNK; + if (k_ >= K) break; + for (int n = 0; n < N; n++) { + if (k_ + Kap * n < max_lds_len) + bigA[n][k2] = *((const bigType*)(&(s[k_ + Kap * n]))); + else + bigA[n][k2] = *((const bigType*)(&(A[k_ + Kap * n]))); + } + } + + // Do the matrix multiplication in interleaved manner + for (uint32_t n = 0; n < N; n++) { + for (uint32_t k2 = 0; k2 < UNRL; k2++) { + for (int y = 0; y < YTILE; y++) { + if constexpr (!use_mfma) + for (uint32_t b = 0; b < A_CHUNK / 2; b++) { + DOT2C(sum[n][y], bigA[n][k2].f[b], bigB[y][k2].f[b]) + } + else + for (uint32_t b = 0; b < A_CHUNK / 4; b++) + sum4[n][y] = __builtin_amdgcn_mfma_f32_4x4x4bf16_1k( + bigA[n][k2].h4[b], bigB[y][k2].h4[b], sum4[n][y], 0, 0, 0); + } + } + } + } + + //---------------------------------------------------- + // Final reduction step using shuffle + //---------------------------------------------------- + if constexpr (!use_mfma) { + for (int n = 0; n < N; n++) { + for (int y = 0; y < YTILE; y++) { + sum[n][y] += __builtin_amdgcn_mov_dpp(sum[n][y], 0x118, 0xf, 0xf, + 1); // row_shr8 + sum[n][y] += __builtin_amdgcn_mov_dpp(sum[n][y], 0x114, 0xf, 0xf, + 1); // row_shr4 + sum[n][y] += __builtin_amdgcn_mov_dpp(sum[n][y], 0x112, 0xf, 0xf, + 1); // row_shr2 + sum[n][y] += __builtin_amdgcn_mov_dpp(sum[n][y], 0x111, 0xf, 0xf, + 1); // row_shr1 + #if defined(__HIP__GFX9__) + sum[n][y] += __builtin_amdgcn_mov_dpp(sum[n][y], 0x142, 0xf, 0xf, + 1); // ROW_BCAST15 + sum[n][y] += __builtin_amdgcn_mov_dpp(sum[n][y], 0x143, 0xf, 0xf, + 1); // ROW_BCAST31 + #else + sum[n][y] += __shfl_xor(sum[n][y], 16); + #endif + } + } + + if (threadIdx.x == (THRDS - 1)) { + scalar_t biases[N][YTILE] = {}; + if (BIAS) + for (int n = 0; n < N; n++) { + for (int y = 0; y < YTILE; y++) { + biases[n][y] = BIAS[(m + y) % Bx + (n % By) * Bx]; + } + } + for (int n = 0; n < N; n++) { + for (int y = 0; y < YTILE; y++) { + if (commitColumn[y]) { + if constexpr (std::is_same_v) { + sum[n][y] += __half2float(biases[n][y]); + } else if constexpr (std::is_same_v) { + sum[n][y] += __bfloat162float(biases[n][y]); + } + C[m + y + n * M] = __float2s(sum[n][y]); + } + } + } + } + } else { + #ifdef __HIP__GFX9__ + #pragma unroll + for (int n = 0; n < N; n++) { + #pragma unroll + for (int y = 0; y < YTILE; y++) { + // float accm1 = 0; + // for (int i=0; i<64; i++) + // accm1 += __shfl(sum4[n][y][i%4], i); + float accm = sum4[n][y][0]; + accm += __builtin_amdgcn_mov_dpp(sum4[n][y][1], 0x101, 0xf, 0xf, + 1); // row_shl1 + accm += __builtin_amdgcn_mov_dpp(sum4[n][y][2], 0x102, 0xf, 0xf, + 1); // row_shl2 + accm += __builtin_amdgcn_mov_dpp(sum4[n][y][3], 0x103, 0xf, 0xf, + 1); // row_shl3 + accm += __builtin_amdgcn_mov_dpp(accm, 0x104, 0xf, 0xf, + 1); // row_shl4 + accm += __builtin_amdgcn_mov_dpp(accm, 0x108, 0xf, 0xf, + 1); // row_shl8 + accm = __builtin_amdgcn_mov_dpp(accm, 0x11f, 0xf, 0xf, + 1); // row_shr15 + accm += __builtin_amdgcn_mov_dpp(accm, 0x142, 0xf, 0xf, + 1); // ROW_BCAST15 + accm += __builtin_amdgcn_mov_dpp(accm, 0x143, 0xf, 0xf, + 1); // ROW_BCAST31 + sum4[n][y][0] = accm; + } + } + if (threadIdx.x == (THRDS - 1)) { + scalar_t biases[N][YTILE] = {}; + if (BIAS) + for (int n = 0; n < N; n++) { + for (int y = 0; y < YTILE; y++) { + biases[n][y] = BIAS[(m + y) % Bx + (n % By) * Bx]; + } + } + for (int n = 0; n < N; n++) { + for (int y = 0; y < YTILE; y++) { + if (commitColumn[y]) { + sum4[n][y][0] += __bfloat162float(biases[n][y]); + C[m + y + n * M] = __float2bfloat16(sum4[n][y][0]); + } + } + } + } + #endif // __HIP__GFX9__ (MFMA path) + } + + m += CuCount * _WvPrGrp * YTILE; + + // Check whether there will be fragmentation! + // This will happen only for the last wave! + if (m < M && (m + YTILE) >= M) { + uint32_t startColumn = M - YTILE; + for (uint32_t i = 0; i < (m - startColumn); i++) { + commitColumn[i] = 0; + } + m = startColumn; + } + } +} + +#else +template +__global__ void wvSplitK_hf_(const int K, const int Kbp, const int Kap, + const int M, const int Bx, const int By, + const scalar_t* B, const scalar_t* __restrict__ A, + const scalar_t* __restrict__ BIAS, scalar_t* C, + const int _WvPrGrp, const int CuCount) { + UNREACHABLE_CODE +} +#endif + +#if defined(__HIP__GFX9__) || defined(__HIP__GFX1X__) +// This version targets big A[] cases, where it is much larger than LDS capacity +template +__global__ void __launch_bounds__(WvPrGrp* THRDS) + wvSplitK_hf_big_(const int K, const int Kbp, const int Kap, const int M, + const int Bx, const int By, const scalar_t* B, + const scalar_t* __restrict__ A, + const scalar_t* __restrict__ BIAS, scalar_t* C, + const int _WvPrGrp, const int CuCount) { + constexpr int max_lds_len = LDS_SIZE / 2; + #if defined(__HIP__MI3XX__) + constexpr bool use_mfma = (std::is_same_v); + #else + constexpr bool use_mfma = false; + #endif + + using scalar8 = + __attribute__((__vector_size__((A_CHUNK / 2) * sizeof(float)))) float; + using half4 = + __attribute__((__vector_size__((A_CHUNK / 2) * sizeof(__bf16)))) __bf16; + union bigType { + scalar_t h[A_CHUNK]; + float f[A_CHUNK / 2]; + float2 f2[A_CHUNK / 4]; + double d[A_CHUNK / 4]; + half4 h4[A_CHUNK / 4]; + scalar8 h8; + }; + + //---------------------------------------------------- + // Reserving 64/160 KB of LDS to have 1 WG / CU + // Goal is to bring the activation matrix A to the LDS + // and use it across the lifetime of the work group + // TODO: When activation matrix is larger than 64 KB + // then this is not going to work! + //---------------------------------------------------- + __shared__ scalar_t s[max_lds_len]; + + //---------------------------------------------------- + // Computation of columns that need to be committed to memory! + //---------------------------------------------------- + uint32_t commitColumn[YTILE]; + for (uint32_t i = 0; i < YTILE; i++) { + commitColumn[i] = 1; + } + + // int _WvPrGrp = mindiv(N, CuCount * YTILE, WvPrGrp); + if (threadIdx.y >= _WvPrGrp) return; + + //---------------------------------------------------- + // Indexing function into the column of weight matrix B + // Algorithm does 64 lane k-splitting / wave and uses + // WG ID and Thread ID to find the index. + //---------------------------------------------------- + uint32_t m = (blockIdx.x * _WvPrGrp + threadIdx.y) * YTILE; + + // Check whether there will be fragmentation! + // This will happen only for the last wave! + if (m < M && (m + YTILE) >= M) { + uint32_t startColumn = M - YTILE; + for (uint32_t i = 0; i < (m - startColumn); i++) { + commitColumn[i] = 0; + } + m = startColumn; + } + + //---------------------------------------------------- + // Fetch the activation matrix to LDS + // Loop iteration: + // - Each thread (lane) is fetching 8 elements (A_Chunk) + // - Each wave will fetch 64*8=> 512 elements + // - Each WG will fetch 512 * 16 => 8K elements + // - Then the WG will move to another 8 K elements + // TODO: Logic below will only work when K is multiple of 8 + //---------------------------------------------------- + #define PCML + #ifndef PCML + for (uint32_t k = (threadIdx.y * THRDS + threadIdx.x) * A_CHUNK; + k < min__(Kap * N, max_lds_len); k += THRDS * WvPrGrp * A_CHUNK) { + #if defined(__gfx950__) + __builtin_amdgcn_global_load_lds((int*)(&A[k]), (int*)(&s[k]), 16, 0, 0); + #else + *((bigType*)(&s[k])) = *((bigType*)(&A[k])); + #endif + } + __syncthreads(); + #endif + + #define TUC (THRDS * UNRL * A_CHUNK) + uint32_t kBase = 0; + // find biggest k size that fits in LDS + uint32_t kFit = (max_lds_len) / N; + // kFit = (kFit%TWC==0) ? kFit : (kFit-kFit%TWC+TWC); //round up to multiple + // of TUC + kFit = (kFit % TUC == 0) + ? kFit + : (kFit - kFit % TUC); // round up to multiple of TUC + // if (kFit == 0) kFit = TUC; + kFit = min__(kFit, Kap); + + //---------------------------------------------------- + // Each wave works on a single column of weight matrix. + // There are 16 waves per WG, and hence, each WG is + // working on 16 columns of weight matrix. Moreover, + // we tile in column direction by YTILE, so when YTILE=1 + // the above math is right, however, when YTILE=2 then + // each wave will be working on 2 columns and WG will + // be working on 32 columns. + // + // Top level loop that makes WGs persistent! + // - WGs iterates across columns of weight matrix + // - Each wave within WG works on a given column(s) + // - After completing first set of columns, WGs start + // working on the next set of available columns + //---------------------------------------------------- + #ifdef PCML + int YW = (YTILE * _WvPrGrp); + uint32_t Mrndp = (M % YW == 0) ? M : (M - M % YW + YW); + while (m < Mrndp) { + #else + while (m < M) { + #endif + //---------------------------------------------------- + // 'sum' accumulates the matrix A x B computation + // split across 64 lanes. + // + // YTILE represents how many column of weight matrix + // are being worked on by each wave. + //---------------------------------------------------- + float sum[N][YTILE] = {}; + scalar8 sum4[N][YTILE] = {}; + + //---------------------------------------------------- + // Fetch weight matrix B in interleaved K-split! + // - Each thread (lane) is fetching 8 elements (A_Chunk) + // - Each wave will fetch 64*8=> 512 elements (1024B) + // - YTILE represents the number of column being serviced + // by wave + // - Loop for fetching weight matrix (B) are unrolled + // + // Fetch activation matrix A from LDS + // - Loop for fetching activation matrix (A) are unrolled + // + // Finally, do the matrix multiplication in an unrolled + // fashion. This provides lot of food for compiler + // scheduling. + // + // TODO: Logic below will only work when K is multiple of 8 + //---------------------------------------------------- + for (uint32_t k1 = 0; k1 < K; k1 += THRDS * A_CHUNK * UNRL) { + bigType bigA[N][UNRL] = {}; + bigType bigB[YTILE][UNRL]; + + #ifdef PCML + if ((k1 == 0) || (k1 == kBase + kFit)) { // load next chunk of A[] to LDS + if (k1 != 0) kBase += kFit; + __syncthreads(); + for (uint32_t k = 0; k < kFit; k += THRDS * _WvPrGrp * A_CHUNK) { + uint32_t kOff = k + ((threadIdx.y * THRDS + threadIdx.x) * A_CHUNK); + if (kBase + kOff >= Kap) break; + if (kOff >= kFit) break; + for (uint32_t n = 0; n < N; n++) { + uint32_t k_in = kBase + n * Kap + kOff; + uint32_t k_ot = n * kFit + kOff; + #if defined(__gfx950__) + __builtin_amdgcn_global_load_lds((int*)(&A[k_in]), (int*)(&s[k_ot]), + 16, 0, 0); + #else + *((bigType*)(&s[k_ot])) = *((bigType*)(&A[k_in])); + #endif + } + } + __syncthreads(); + } + if (m >= M) continue; + #endif + + // Fetch the weight matrix from memory! + #pragma unroll + for (uint32_t k2 = 0; k2 < UNRL; k2++) { + uint32_t k = k1 + k2 * THRDS * A_CHUNK; + uint32_t k_ = k + threadIdx.x * A_CHUNK; + const scalar_t* B_ = &B[min__(k_, K - A_CHUNK)]; + for (int y = 0; y < YTILE; y++) + bigB[y][k2].h8 = (loadnt((scalar8*)(&B_[min__(y + m, M - 1) * Kbp]))); + } + + // Fetch activation matrix from either just LDS or from both LDS / memory + #pragma unroll + for (uint32_t k2 = 0; k2 < UNRL; k2++) { + uint32_t k = k1 + k2 * THRDS * A_CHUNK; + uint32_t k_ = k + threadIdx.x * A_CHUNK; + if (k_ >= K) break; + for (int n = 0; n < N; n++) { + #ifdef PCML + bigA[n][k2] = *((const bigType*)(&(s[k_ - kBase + kFit * n]))); + #else + if (k_ + Kap * n < max_lds_len) + bigA[n][k2] = *((const bigType*)(&(s[k_ + Kap * n]))); + else + bigA[n][k2] = *((const bigType*)(&(A[k_ + Kap * n]))); + #endif + } + } + + // Do the matrix multiplication in interleaved manner + #pragma unroll + for (uint32_t k2 = 0; k2 < UNRL; k2++) { + for (uint32_t n = 0; n < N; n++) { + for (int y = 0; y < YTILE; y++) { + if constexpr (!use_mfma) + for (uint32_t b = 0; b < A_CHUNK / 2; b++) { + DOT2C(sum[n][y], bigA[n][k2].f[b], bigB[y][k2].f[b]) + } + else + for (uint32_t b = 0; b < A_CHUNK / 4; b++) + sum4[n][y] = __builtin_amdgcn_mfma_f32_4x4x4bf16_1k( + bigA[n][k2].h4[b], bigB[y][k2].h4[b], sum4[n][y], 0, 0, 0); + } + } + } + } + + #ifdef PCML + if (m >= M) { + m += CuCount * _WvPrGrp * YTILE; + kBase = 0; + continue; + } + #endif + + //---------------------------------------------------- + // Final reduction step using shuffle + //---------------------------------------------------- + if constexpr (!use_mfma) { + for (int n = 0; n < N; n++) { + for (int y = 0; y < YTILE; y++) { + sum[n][y] += __builtin_amdgcn_mov_dpp(sum[n][y], 0x118, 0xf, 0xf, + 1); // row_shr8 + sum[n][y] += __builtin_amdgcn_mov_dpp(sum[n][y], 0x114, 0xf, 0xf, + 1); // row_shr4 + sum[n][y] += __builtin_amdgcn_mov_dpp(sum[n][y], 0x112, 0xf, 0xf, + 1); // row_shr2 + sum[n][y] += __builtin_amdgcn_mov_dpp(sum[n][y], 0x111, 0xf, 0xf, + 1); // row_shr1 + #if defined(__HIP__GFX9__) + sum[n][y] += __builtin_amdgcn_mov_dpp(sum[n][y], 0x142, 0xf, 0xf, + 1); // ROW_BCAST15 + sum[n][y] += __builtin_amdgcn_mov_dpp(sum[n][y], 0x143, 0xf, 0xf, + 1); // ROW_BCAST31 + #else + sum[n][y] += __shfl_xor(sum[n][y], 16); + #endif + } + } + + if (threadIdx.x == (THRDS - 1)) { + scalar_t biases[N][YTILE] = {}; + if (BIAS) + for (int n = 0; n < N; n++) { + for (int y = 0; y < YTILE; y++) { + biases[n][y] = BIAS[(m + y) % Bx + (n % By) * Bx]; + } + } + for (int n = 0; n < N; n++) { + for (int y = 0; y < YTILE; y++) { + if (commitColumn[y]) { + if constexpr (std::is_same_v) { + sum[n][y] += __half2float(biases[n][y]); + } else if constexpr (std::is_same_v) { + sum[n][y] += __bfloat162float(biases[n][y]); + } + C[m + y + n * M] = __float2s(sum[n][y]); + } + } + } + } + } else { + #ifdef __HIP__GFX9__ + #pragma unroll + for (int n = 0; n < N; n++) { + #pragma unroll + for (int y = 0; y < YTILE; y++) { + float accm = sum4[n][y][0]; + accm += __builtin_amdgcn_mov_dpp(sum4[n][y][1], 0x101, 0xf, 0xf, + 1); // row_shl1 + accm += __builtin_amdgcn_mov_dpp(sum4[n][y][2], 0x102, 0xf, 0xf, + 1); // row_shl2 + accm += __builtin_amdgcn_mov_dpp(sum4[n][y][3], 0x103, 0xf, 0xf, + 1); // row_shl3 + accm += __builtin_amdgcn_mov_dpp(accm, 0x104, 0xf, 0xf, + 1); // row_shl4 + accm += __builtin_amdgcn_mov_dpp(accm, 0x108, 0xf, 0xf, + 1); // row_shl8 + accm = __builtin_amdgcn_mov_dpp(accm, 0x11f, 0xf, 0xf, + 1); // row_shr15 + accm += __builtin_amdgcn_mov_dpp(accm, 0x142, 0xf, 0xf, + 1); // ROW_BCAST15 + accm += __builtin_amdgcn_mov_dpp(accm, 0x143, 0xf, 0xf, + 1); // ROW_BCAST31 + sum4[n][y][0] = accm; + } + } + if (threadIdx.x == (THRDS - 1)) { + scalar_t biases[N][YTILE] = {}; + if (BIAS) + for (int n = 0; n < N; n++) { + for (int y = 0; y < YTILE; y++) { + biases[n][y] = BIAS[(m + y) % Bx + (n % By) * Bx]; + } + } + for (int n = 0; n < N; n++) { + for (int y = 0; y < YTILE; y++) { + if (commitColumn[y]) { + sum4[n][y][0] += __bfloat162float(biases[n][y]); + C[m + y + n * M] = __float2bfloat16(sum4[n][y][0]); + } + } + } + } + #endif // __HIP__GFX9__ (MFMA path) + } + + m += CuCount * _WvPrGrp * YTILE; + kBase = 0; + + // Check whether there will be fragmentation! + // This will happen only for the last wave! + if (m < M && (m + YTILE) >= M) { + uint32_t startColumn = M - YTILE; + for (uint32_t i = 0; i < (m - startColumn); i++) { + commitColumn[i] = 0; + } + m = startColumn; + } + } +} +#else +template +__global__ void wvSplitK_hf_big_(const int K, const int Kbp, const int Kap, + const int M, const int Bx, const int By, + const scalar_t* B, + const scalar_t* __restrict__ A, + const scalar_t* __restrict__ BIAS, scalar_t* C, + const int _WvPrGrp, const int CuCount) { + UNREACHABLE_CODE +} +#endif + +// Find the min val of div2 that doesn't increase N/(div1*div2) +int mindiv(int N, int div1, int div2) { + int nPrRnd = div1 * div2; + int rnds[13]; + for (int i = 0; i < 13; i++) { + rnds[i] = (N + nPrRnd - 1) / nPrRnd; + nPrRnd -= div1; + } + for (int i = 12; i >= 0; i--) + if (rnds[0] == rnds[i]) return (div2 - i); + return 0; +} + +torch::Tensor wvSplitK(const at::Tensor& in_a, const at::Tensor& in_b, + const std::optional& in_bias, + const int64_t CuCount) { + auto M_in = in_a.size(0); + auto K_in = in_a.size(1); + auto N_in = in_b.size(0); + auto Kap_in = in_a.stride(0); + auto Kbp_in = in_b.stride(0); + auto Bx_in = + (in_bias.has_value() && in_bias->numel() > 0) + ? (in_bias->sizes().size() == 2) ? in_bias->size(1) : in_bias->size(0) + : 1; + auto By_in = (in_bias.has_value() && in_bias->numel() > 0 && + in_bias->sizes().size() == 2) + ? in_bias->size(0) + : 1; + + TORCH_CHECK(in_a.dtype() == in_b.dtype()); + TORCH_CHECK(K_in % 8 == 0, "k % 8 == 0"); + TORCH_CHECK(in_a.dtype() == torch::kFloat16 || + in_a.dtype() == torch::kBFloat16); + + auto out_c = torch::empty( + {N_in, M_in}, + torch::TensorOptions().dtype(in_b.dtype()).device(in_b.device())); + + dim3 grid(CuCount); + + const at::cuda::OptionalCUDAGuard device_guard(device_of(in_a)); + const cudaStream_t stream = at::cuda::getCurrentCUDAStream(); + const int max_lds_len = get_lds_size() / 2; + +#define WVSPLITK_CFG(_THRDS, _WVPRGRP, _YTILE, _UNRL, _N) \ + { \ + dim3 block(_THRDS, _WVPRGRP); \ + int __wvPrGrp = mindiv(M_in, CuCount * _YTILE, _WVPRGRP); \ + if ((Kbp_in * N_in <= max_lds_len) && (M_in % _YTILE == 0)) \ + wvSplitK_hf_sml_ \ + <<>>(K_in, Kap_in, Kbp_in, M_in, Bx_in, \ + By_in, af4, bf4, biasf4, c, __wvPrGrp, \ + CuCount); \ + else if (Kbp_in * N_in <= max_lds_len * 1.2) \ + wvSplitK_hf_ \ + <<>>(K_in, Kap_in, Kbp_in, M_in, Bx_in, \ + By_in, af4, bf4, biasf4, c, __wvPrGrp, \ + CuCount); \ + else \ + wvSplitK_hf_big_ \ + <<>>(K_in, Kap_in, Kbp_in, M_in, Bx_in, \ + By_in, af4, bf4, biasf4, c, __wvPrGrp, \ + CuCount); \ + } + +#define WVSPLIT_TILE_CFG(_THRDS, _WVPRGRP, _sYT, __N) \ + { \ + bool fit_lds = (Kbp_in * N_in <= max_lds_len); \ + if (_sYT <= 1) \ + WVSPLITK_CFG(_THRDS, _WVPRGRP, 1, 4, __N) \ + else if ((__N == 1) || (!fit_lds) || (_sYT <= 4 * 2)) \ + WVSPLITK_CFG(_THRDS, _WVPRGRP, 2, 2, __N) \ + else if (_sYT <= 4 * 3) \ + WVSPLITK_CFG(_THRDS, _WVPRGRP, 3, 2, __N) \ + else if (__N == 4) \ + WVSPLITK_CFG(_THRDS, _WVPRGRP, 4, 1, __N) \ + else \ + WVSPLITK_CFG(_THRDS, _WVPRGRP, 4, 2, __N) \ + } + + AT_DISPATCH_REDUCED_FLOATING_TYPES(in_b.scalar_type(), "wvSplitK", [&] { + using fptype = typename scalar::type; + fptype* af4 = reinterpret_cast(in_a.data_ptr()); + const fptype* bf4 = reinterpret_cast(in_b.data_ptr()); + const fptype* biasf4 = + (in_bias.has_value() && in_bias->numel() > 0) + ? reinterpret_cast(in_bias->data_ptr()) + : nullptr; + fptype* c = reinterpret_cast(out_c.data_ptr()); + + // first shoot for biggest tile-size that keeps all simd busy, + // then cut the active waves to balance their distribution... + int sYT = (M_in + CuCount * 4 - 1) / (CuCount * 4); + + const bool use_wave32 = on_gfx1x(); + switch (N_in) { + case 1: + if (use_wave32) + WVSPLIT_TILE_CFG(32, 16, sYT, 1) + else + WVSPLIT_TILE_CFG(64, 16, sYT, 1) + break; + case 2: + if (use_wave32) + WVSPLIT_TILE_CFG(32, 16, sYT, 2) + else + WVSPLIT_TILE_CFG(64, 16, sYT, 2) + break; + case 3: + if (use_wave32) + WVSPLIT_TILE_CFG(32, 16, sYT, 3) + else + WVSPLIT_TILE_CFG(64, 16, sYT, 3) + break; + case 4: + if (use_wave32) + WVSPLIT_TILE_CFG(32, 16, sYT, 4) + else + WVSPLIT_TILE_CFG(64, 16, sYT, 4) + break; + default: + throw std::runtime_error( + "Unsupported N value: " + std::to_string(M_in) + "," + + std::to_string(K_in) + "," + std::to_string(N_in)); + } + }); + return out_c; +} + +// This version targets cases skinny where CUs are not filled +// Wave-SplitK is used with reduction done via atomics. +#if defined(__gfx950__) + #define WVSPLITKRC_1KPASS +template +__global__ void __launch_bounds__(WvPrGrp* THRDS) + __attribute__((amdgpu_waves_per_eu(1, 1))) + wvSplitKrc_(const int actlN, const int K, const int Kap, const int M, + const int Bx, const int By, const scalar_t* __restrict__ A, + const scalar_t* __restrict__ B, + const scalar_t* __restrict__ BIAS, float* glbl, int* cntr, + scalar_t* C, const int CuCount) { + constexpr int NTILE = 16; + constexpr int APAD = 1; + constexpr int ASTRD = 64; + constexpr int BPAD = 1; + constexpr int WVLDS_ = THRDS * A_CHUNK / CHUNKK; + constexpr int WVLDS = ((WVLDS_ + A_CHUNK * BPAD)) * YTILE; + + constexpr int max_lds_len = LDS_SIZE / 2; + + using scalar16 = + __attribute__((__vector_size__((A_CHUNK * 2) * sizeof(float)))) float; + using scalar8 = + __attribute__((__vector_size__((A_CHUNK / 2) * sizeof(float)))) float; + using half4 = + __attribute__((__vector_size__((A_CHUNK / 2) * sizeof(__bf16)))) __bf16; + union bigType { + scalar_t h[A_CHUNK]; + float f[A_CHUNK / 2]; + unsigned int i[A_CHUNK / 2]; + float2 f2[A_CHUNK / 4]; + unsigned long l[A_CHUNK / 4]; + double d[A_CHUNK / 4]; + half4 h4[A_CHUNK / 4]; + scalar8 h8; + }; + using big4 = __attribute__((__vector_size__(4 * sizeof(bigType)))) __bf16; + + __shared__ scalar_t stg[WvPrGrp * WVLDS / GrpsShrB]; + unsigned int* myStg = (unsigned int*)(&stg[WVLDS * (threadIdx.y / GrpsShrB)]); + __shared__ scalar_t s[max_lds_len - WvPrGrp * WVLDS / GrpsShrB]; + + #ifndef WVSPLITKRC_1KPASS + constexpr int TUC_ = (THRDS * UNRL * A_CHUNK); + // find biggest k size that fits padded into LDS + constexpr uint32_t kFit__ = (max_lds_len - WvPrGrp * WVLDS / GrpsShrB) / N; + constexpr uint32_t kFit_ = (kFit__ * ASTRD) / (APAD + ASTRD); + uint32_t kFit = kFit_ - (kFit_ % TUC_); + uint32_t kfitsPerRdc = (K + kFit - 1) / kFit; + + // find best k split to fill the CUs + if (((K + kfitsPerRdc * kFit - 1) / (kfitsPerRdc * kFit)) * numCuWithFullK <= + CuCount) + while (true) { + while (kFit > TUC_) { + uint32_t kFit_ = kFit - TUC_; + if (((K + (kfitsPerRdc * kFit_ - 1)) / (kfitsPerRdc * kFit_)) * + numCuWithFullK > + CuCount) + break; + kFit = kFit_; + } + if (((K + ((kfitsPerRdc - 1) * kFit - 1)) / ((kfitsPerRdc - 1) * kFit)) * + numCuWithFullK <= + CuCount) + kfitsPerRdc--; + else + break; + } + #else + int constexpr kFit = 512 / CHUNKK; + int constexpr kfitsPerRdc = 1; + #endif + + bool doRdc = true; // Assuming (kfitsPerRdc * kFit < K) is always true + uint32_t numCuWithFullK = + ((M + (WvPrGrp * YTILE / GrpsShrB) - 1) / (WvPrGrp * YTILE / GrpsShrB)); + uint32_t Mmod = numCuWithFullK * (WvPrGrp * YTILE / GrpsShrB); + + // given above k-split, find this wave's position + uint32_t kFitPdd = kFit * CHUNKK + ((kFit * CHUNKK) / ASTRD) * APAD; + uint32_t m0 = (blockIdx.x * WvPrGrp / GrpsShrB) * YTILE; + uint32_t m1 = ((threadIdx.y % WvPrGrp) / GrpsShrB) * YTILE; + uint32_t m = (m0 + m1) % Mmod; + const uint32_t k_str = (m0 / Mmod) * kFit * kfitsPerRdc; + uint32_t k_end = (m0 / Mmod + 1) * kFit * kfitsPerRdc; + const uint32_t k_rnd = (K + kFit * kfitsPerRdc - 1) / (kFit * kfitsPerRdc); + + scalar8 sum4[N / NTILE / GrpsShrB][1] = {0}; + bigType bigB_[YTILE / GrpsShrB / CHUNKK][UNRL]; + const uint32_t bLoader = (threadIdx.y % GrpsShrB); + uint32_t kBase = 0; + if (k_str >= K) return; + if (m >= Mmod) return; + + bool noreloada = false; + constexpr bool FAST_UNSAFE_RDC_INIT = false; + + #ifdef WVSPLITKRC_1KPASS + // Early glbl init, B[] loading, if 1KPASS + if constexpr (FAST_UNSAFE_RDC_INIT) { + if (m + (threadIdx.x % 16) < M) + if (doRdc) + if (k_str == 0) { + int mindx = m + (threadIdx.x % 16); + int nindx_ = (0 + (threadIdx.x / 16) * 4) + 0 * NTILE + + (N / GrpsShrB) * (threadIdx.y % GrpsShrB); + int adr_ = mindx + M * nindx_ / 4; + __hip_atomic_store(&cntr[adr_], 0, __ATOMIC_RELAXED, + __HIP_MEMORY_SCOPE_AGENT); + for (uint32_t nt = 0; nt < N / NTILE / GrpsShrB; nt++) { + for (uint32_t j = 0; j < 4; j++) { + int nindx = (j + (threadIdx.x / 16) * 4) + nt * NTILE + + (N / GrpsShrB) * (threadIdx.y % GrpsShrB); + int adr = mindx + M * nindx; + __hip_atomic_store(&glbl[adr], 0, __ATOMIC_RELAXED, + __HIP_MEMORY_SCOPE_AGENT); + } + } + } + } + + // Load first B[] chunk + #pragma unroll + for (uint32_t k2 = 0; k2 < UNRL; k2++) { + uint32_t k = k_str + k2 * THRDS * A_CHUNK; + uint32_t k_ = k + (threadIdx.x % (THRDS / CHUNKK)) * A_CHUNK; + const scalar_t* B_ = &B[min__(k_, K - A_CHUNK)]; + #pragma unroll + for (uint32_t y = 0; y < YTILE / GrpsShrB; y += CHUNKK) + bigB_[y / CHUNKK][k2].h8 = (loadnt( + (scalar8*)(&B_[min__((y + threadIdx.x / (THRDS / CHUNKK)) * GrpsShrB + + bLoader + m, + M - 1) * + K]))); + } + { + #else + while (m < Mmod) { + #endif + + #ifndef WVSPLITKRC_1KPASS + if constexpr (FAST_UNSAFE_RDC_INIT) { + if (m + (threadIdx.x % 16) < M) + if (doRdc) + if (k_str == 0) { + int mindx = m + (threadIdx.x % 16); + int nindx_ = (0 + (threadIdx.x / 16) * 4) + 0 * NTILE + + (N / GrpsShrB) * (threadIdx.y % GrpsShrB); + int adr_ = mindx + M * nindx_ / 4; + __hip_atomic_store(&cntr[adr_], 0, __ATOMIC_RELAXED, + __HIP_MEMORY_SCOPE_AGENT); + for (uint32_t nt = 0; nt < N / NTILE / GrpsShrB; nt++) { + for (uint32_t j = 0; j < 4; j++) { + int nindx = (j + (threadIdx.x / 16) * 4) + nt * NTILE + + (N / GrpsShrB) * (threadIdx.y % GrpsShrB); + int adr = mindx + M * nindx; + __hip_atomic_store(&glbl[adr], 0, __ATOMIC_RELAXED, + __HIP_MEMORY_SCOPE_AGENT); + } + } + } + } + + #endif + + #ifndef WVSPLITKRC_1KPASS + for (uint32_t k1 = k_str; k1 < k_end; k1 += THRDS * A_CHUNK * UNRL) { + #else + const uint32_t k1 = k_str; + { + #endif + #ifndef WVSPLITKRC_1KPASS + const bool reloada = (!noreloada) && + ((k1 == k_str) || (k1 == k_str + kBase + kFit)) && + (k1 < k_end); + // load next chunk of A[] to LDS + if (reloada) { + if (k1 != k_str) kBase += kFit; + __syncthreads(); + #else + const bool reloada = (!noreloada) && + ((k1 == k_str) || (k1 == k_str + kBase + kFit)) && + (k1 < k_end); + if (reloada) { + #endif + constexpr int sprdN = 4; + const uint32_t thrd = threadIdx.x % (THRDS / CHUNKK); + + #ifndef WVSPLITKRC_1KPASS + #pragma unroll + for (int k = 0; k < kFit; + k += (THRDS * (WvPrGrp / sprdN) * A_CHUNK) / CHUNKK) { + #else + const unsigned int k = 0; + { + #endif + unsigned int kOff = k + (thrd * A_CHUNK); + unsigned int kOffcp = min__(K - A_CHUNK, k_str + kOff); + for (unsigned int n = 0; n < N; n += CHUNKK * sprdN) { + __builtin_amdgcn_global_load_lds( + (int*)(&A[min__(Kap * actlN - A_CHUNK, + kOffcp + Kap * (n / CHUNKK + + (N / CHUNKK) * (threadIdx.x / + (64 / CHUNKK)) + + (threadIdx.y % sprdN)))]), + (int*)(&s[(k + + kFitPdd * ((n / CHUNKK) + (threadIdx.y % sprdN)))]), + 16, 0, 0); + } + + // Stage loaded B[] to LDS for MFMA swizzling... + for (uint32_t k2 = 0; k2 < UNRL; k2++) { + uint32_t k = k1 + k2 * THRDS * A_CHUNK; + uint32_t k_ = k + (threadIdx.x % (THRDS / CHUNKK)) * A_CHUNK; + const bool oob_k = (k_ >= K); + for (uint32_t y = 0; y < YTILE / GrpsShrB; y += CHUNKK) { + uint32_t idx = + (threadIdx.x % (THRDS / CHUNKK)) * 4 + + ((y + threadIdx.x / (THRDS / CHUNKK)) * GrpsShrB + bLoader) * + ((THRDS / CHUNKK + BPAD) * 4); + // zero out if oob + *((scalar8*)&myStg[idx]) = + (oob_k) // TODO: ever necessary (y*GrpsShrB+bLoader+m>=M) ? + ? 0 + : bigB_[y / CHUNKK][k2].h8; + } + } + } + } + } + #ifndef WVSPLITKRC_1KPASS + // Fire load of next B[] chunk... + if ((k1 + THRDS * A_CHUNK * UNRL < k_end) && + (k1 + THRDS * A_CHUNK * UNRL < K)) + #pragma unroll + for (uint32_t k2 = 0; k2 < UNRL; k2++) { + uint32_t k = k1 + THRDS * A_CHUNK * UNRL + k2 * THRDS * A_CHUNK; + uint32_t k_ = k + threadIdx.x * A_CHUNK; + const scalar_t* B_ = &B[min__(k_, K - A_CHUNK)]; + #pragma unroll + for (uint32_t y = 0; y < YTILE / GrpsShrB; y += CHUNKK) + bigB_[y / CHUNKK][k2].h8 = (loadnt( + (scalar8*)(&B_[min__((y + threadIdx.x / (THRDS / CHUNKK)) * + GrpsShrB + + bLoader + m, + M - 1) * + K]))); + } + #endif + + // B[] staging is cooperative across GrpsShrB, so sync here before reading + // back. This wait is currently inserted by compiler, but not guaranteed. + asm volatile("s_waitcnt 0"); + __syncthreads(); + + // read back B[] swizzled for MFMA... + bigType bigB[YTILE / CHUNKK][UNRL]; + for (uint32_t k2 = 0; k2 < UNRL; k2++) { + for (uint32_t y = 0; y < YTILE / CHUNKK; y++) { + unsigned int idx = + (threadIdx.x % YTILE) * ((THRDS / CHUNKK + BPAD) * 4) + + (threadIdx.x / YTILE) * 4 + y * 16; + bigB[y][k2].h8 = *((scalar8*)&myStg[idx]); + } + } + + // rReadback A[] swizzled for MFMA... + bigType bigA[N / GrpsShrB / CHUNKK][UNRL]; + #pragma unroll + for (uint32_t k2 = 0; k2 < UNRL; k2++) { + uint32_t k = k1 + k2 * THRDS * A_CHUNK - kBase - k_str; + #pragma unroll + for (uint32_t nt = 0; nt < N / GrpsShrB; nt += NTILE) + #pragma unroll + for (uint32_t n = 0; n < NTILE / CHUNKK; n++) { + uint32_t idxa = + ((nt + (N / GrpsShrB) * (threadIdx.y % GrpsShrB)) % (N / CHUNKK) + + (threadIdx.x % NTILE)) * + kFitPdd + + ((nt + (N / GrpsShrB) * (threadIdx.y % GrpsShrB)) / + (N / CHUNKK)) * + A_CHUNK * (64 / CHUNKK) + + A_CHUNK * ((threadIdx.x / NTILE) + n * 4) + k; + bigA[nt / CHUNKK + n][k2] = *((const bigType*)(&(s[idxa]))); + } + } + + // Do the MFMAs + #pragma unroll + for (uint32_t k2 = 0; k2 < UNRL; k2++) { + #pragma unroll + for (uint32_t nt = 0; nt < N / NTILE / GrpsShrB; nt++) { + #pragma unroll + for (uint32_t j = 0; j < YTILE / CHUNKK; j++) { + if constexpr (std::is_same_v) { + sum4[nt][0] = __builtin_amdgcn_mfma_f32_16x16x32_f16( + bigA[nt * (YTILE / CHUNKK) + j][k2].h8, bigB[j][k2].h8, + sum4[nt][0], 0, 0, 0); + } else { // bf16 + sum4[nt][0] = __builtin_amdgcn_mfma_f32_16x16x32_bf16( + bigA[nt * (YTILE / CHUNKK) + j][k2].h8, bigB[j][k2].h8, + sum4[nt][0], 0, 0, 0); + } + } + } + } + } + + union flt4 { + scalar8 s8; + float2 f2[2]; + float4 f4; + }; + if (m + (threadIdx.x % 16) < M) { + int my_cntr; + int mindx = m + (threadIdx.x % 16); + int g_mindx = m * 4 + (threadIdx.x % 64); // coalesced atomic reduction + scalar_t biases[N / NTILE / GrpsShrB][4] = {}; + // Atomic add the output, read biases + for (uint32_t nt = 0; nt < N / NTILE / GrpsShrB; nt++) { + int g_nindx = + (nt * NTILE + (N / GrpsShrB) * (threadIdx.y % GrpsShrB)) / 4; + int g_adr = g_mindx * 4 + 0 + M * g_nindx * 4; + if (DTRMNSTC) { + flt4 flt4_ = {.s8 = sum4[nt][0]}; + __hip_atomic_store((float2*)&glbl[g_adr + M * N * (m0 / Mmod)], + flt4_.f2[0], __ATOMIC_RELAXED, + __HIP_MEMORY_SCOPE_AGENT); + __hip_atomic_store((float2*)&glbl[g_adr + 2 + M * N * (m0 / Mmod)], + flt4_.f2[1], __ATOMIC_RELAXED, + __HIP_MEMORY_SCOPE_AGENT); + } else { + for (uint32_t j = 0; j < 4; j++) + atomicAdd((&glbl[g_adr + j]), sum4[nt][0][j]); + } + } + + __atomic_signal_fence(__ATOMIC_SEQ_CST); + asm volatile("s_waitcnt vmcnt(0)" ::: "memory"); + __atomic_signal_fence(__ATOMIC_SEQ_CST); + + int nindx_ = (0 + (threadIdx.x / 16) * 4) + 0 * NTILE + + (N / GrpsShrB) * (threadIdx.y % GrpsShrB); + int adr_ = mindx + M * nindx_ / 4; + my_cntr = atomicAdd(&cntr[adr_], 1); + + // make sure LDS is free for write out staging + if (DTRMNSTC) __syncthreads(); + + // Update the complete counter + flt4 vals[N / NTILE / GrpsShrB] = {}; + // If we're the last k-shard, read back the value and convert... + if (my_cntr + 1 == k_rnd) { + cntr[adr_] = 0; // clear for next round + if constexpr (DTRMNSTC) { + #pragma unroll + for (int ks = 0; ks < k_rnd; ks++) { + for (uint32_t nt = 0; nt < N / NTILE / GrpsShrB; nt++) { + int g_nindx = + (nt * NTILE + (N / GrpsShrB) * (threadIdx.y % GrpsShrB)) / 4; + int g_adr = g_mindx * 4 + 0 + M * g_nindx * 4; + __builtin_amdgcn_global_load_lds( + (float4*)(&glbl[g_adr + M * N * ks]), + &(((float4*)s)[(threadIdx.y * THRDS) + ks * THRDS * 4 + + nt * THRDS * 4 * k_rnd]), + 16, 0, 0); + } + } + if (BIAS) + for (uint32_t nt = 0; nt < N / NTILE / GrpsShrB; nt++) { + for (uint32_t j = 0; j < 4; j++) { + int nindx = (j + (threadIdx.x / 16) * 4) + nt * NTILE + + (N / GrpsShrB) * (threadIdx.y % GrpsShrB); + biases[nt][j] = BIAS[(mindx % Bx) + (nindx % By) * Bx]; + } + } + asm volatile("s_waitcnt 0"); + for (int ks = 0; ks < k_rnd; ks++) { + for (uint32_t nt = 0; nt < N / NTILE / GrpsShrB; nt++) { + float4 eval = ((float4*)s)[(threadIdx.x + threadIdx.y * THRDS) + + ks * THRDS * 4 + nt * THRDS * 4 * k_rnd]; + vals[nt].f4 += eval; + } + } + } else { + for (uint32_t nt = 0; nt < N / NTILE / GrpsShrB; nt++) { + int g_nindx = + (nt * NTILE + (N / GrpsShrB) * (threadIdx.y % GrpsShrB)) / 4; + int g_adr = g_mindx * 4 + 0 + M * g_nindx * 4; + vals[nt].f4 = *(float4*)(&glbl[g_adr]); + *(float4*)(&glbl[g_adr]) = {}; // clear out for next round + } + if (BIAS) + for (uint32_t nt = 0; nt < N / NTILE / GrpsShrB; nt++) { + for (uint32_t j = 0; j < 4; j++) { + int nindx = (j + (threadIdx.x / 16) * 4) + nt * NTILE + + (N / GrpsShrB) * (threadIdx.y % GrpsShrB); + biases[nt][j] = BIAS[(mindx % Bx) + (nindx % By) * Bx]; + } + } + } + __builtin_amdgcn_sched_barrier(0); + for (uint32_t nt = 0; nt < N / NTILE / GrpsShrB; nt++) { + for (uint32_t j = 0; j < 4; j++) { + int nindx = (j + (threadIdx.x / 16) * 4) + nt * NTILE + + (N / GrpsShrB) * (threadIdx.y % GrpsShrB); + if (nindx < actlN) { + int adr = mindx + M * nindx; + if constexpr (std::is_same_v) { + vals[nt].s8[j] += __bfloat162float(biases[nt][j]); + C[adr] = __float2bfloat16(vals[nt].s8[j]); + } else { + vals[nt].s8[j] += __half2float(biases[nt][j]); + C[adr] = __float2half(vals[nt].s8[j]); + } + } + } + } + } + + #ifndef WVSPLITKRC_1KPASS + m0 += CuCount * WvPrGrp * YTILE / GrpsShrB; + m = (m0 + m1) % Mmod; + k_str = (m0 / Mmod) * kFit * kfitsPerRdc; + k_end = (m0 / Mmod + 1) * kFit * kfitsPerRdc; + if (k_str >= K) break; + kBase = 0; + #endif + } +} +#else +template +__global__ void wvSplitKrc_(const int actlN, const int K, const int Kap, + const int M, const int Bx, const int By, + const scalar_t* B, const scalar_t* __restrict__ A, + const scalar_t* __restrict__ BIAS, float* glbl, + int* cntr, scalar_t* C, + const int CuCount){UNREACHABLE_CODE} +#endif // defined(__HIP__GFX9__) TODO: Add NAVI support + +torch::Tensor wvSplitKrc(const at::Tensor& in_a, const at::Tensor& in_b, + const std::optional& in_bias, + const int64_t CuCount) { + int _DTRMNSTC = 1; // vllm::vllm_is_batch_invariant(); + + auto M_in = in_b.size(0); + auto N_in = in_a.size(0); + auto K_in = in_b.size(1); + auto Kap_in = in_a.stride(0); + + auto Bx_in = + (in_bias.has_value() && in_bias->numel() > 0) + ? (in_bias->sizes().size() == 2) ? in_bias->size(1) : in_bias->size(0) + : 1; + auto By_in = (in_bias.has_value() && in_bias->numel() > 0 && + in_bias->sizes().size() == 2) + ? in_bias->size(0) + : 1; + + TORCH_CHECK(in_a.dtype() == in_b.dtype()); + TORCH_CHECK(K_in % 8 == 0, "k % 8 == 0"); + TORCH_CHECK(in_a.dtype() == torch::kFloat16 || + in_a.dtype() == torch::kBFloat16); + + const at::cuda::OptionalCUDAGuard device_guard(device_of(in_a)); + + auto out_c = torch::empty( + {N_in, M_in}, + torch::TensorOptions().dtype(in_a.dtype()).device(in_a.device())); + + auto N_p2 = 1U << (32 - __builtin_clz(N_in - 1)); + + dim3 grid(CuCount); + + const cudaStream_t stream = at::cuda::getCurrentCUDAStream(); + // const int max_lds_len = get_lds_size() / 2; + + // With 64 Ms per CU (each of 4 SIMDs working on a 16x16 tile), + // and each working on a 512-shard of K, how many CUs would we need? + int rndup_cus = ((M_in + 64 - 1) / 64) * ((K_in + 512 - 1) / 512); + + // How many of 4 waves in a group can work on same 16 Ms at same time? First + // try to maximize this. This reduces the Ms each group works on, i.e. + // increasing the number of CUs needed. + int GrpsShrB = min(N_p2 / 16, 4); + + // Given the above, how many CUs would we need? + int CuNeeded = rndup_cus * GrpsShrB; + + if (CuNeeded > CuCount) throw std::runtime_error("Invalid wvSplitKrc size"); + + // Can we increase SplitK by shrinking the K-shared to 256? + int chunkk = (CuNeeded * 2 <= CuCount) ? 2 : 1; + + static torch::Tensor axl_glbl = + torch::zeros( + 128 * 1024 * (_DTRMNSTC ? 12 : 1), + torch::TensorOptions().dtype(torch::kFloat32).device(in_a.device())) + .detach(); + static torch::Tensor axl_cntr = + torch::zeros( + 128 * 1024 * (_DTRMNSTC ? 12 : 1) / 4, + torch::TensorOptions().dtype(torch::kInt).device(in_a.device())) + .detach(); + auto glbl = axl_glbl.data_ptr(); + auto cntr = axl_cntr.data_ptr(); + +#define WVSPLITKrc(_N, _GrpsShrB, _CHUNKK) \ + { \ + dim3 block(64, 4); \ + if (_DTRMNSTC) \ + wvSplitKrc_ \ + <<>>(N_in, K_in, Kap_in, M_in, Bx_in, By_in, \ + af4, bf4, biasf4, glbl, cntr, c, \ + CuCount); \ + else \ + wvSplitKrc_ \ + <<>>(N_in, K_in, Kap_in, M_in, Bx_in, By_in, \ + af4, bf4, biasf4, glbl, cntr, c, \ + CuCount); \ + } + + AT_DISPATCH_REDUCED_FLOATING_TYPES(in_a.scalar_type(), "wvSplitKrc", [&] { + using fptype = typename scalar::type; + const fptype* af4 = reinterpret_cast(in_a.data_ptr()); + const fptype* bf4 = reinterpret_cast(in_b.data_ptr()); + const fptype* biasf4 = + (in_bias.has_value() && in_bias->numel() > 0) + ? reinterpret_cast(in_bias->data_ptr()) + : nullptr; + fptype* c = reinterpret_cast(out_c.data_ptr()); + + switch (N_p2) { + case 16: + WVSPLITKrc(16, 1, 1) break; + case 32: + if (chunkk == 2) WVSPLITKrc(32, 2, 2) else WVSPLITKrc(32, 2, 1) break; + case 64: + if (chunkk == 2) WVSPLITKrc(64, 4, 2) else WVSPLITKrc(64, 4, 1) break; + case 128: + if (chunkk == 2) WVSPLITKrc(128, 4, 2) else WVSPLITKrc(128, 4, 1) break; + default: + throw std::runtime_error( + "Unsupported N value: " + std::to_string(M_in) + "," + + std::to_string(K_in) + "," + std::to_string(N_in)); + } + }); + return out_c; +} + +#if defined(__HIP__MI3XX__) || defined(__HIP__GFX12__) +template +__global__ void __launch_bounds__(WvPrGrp* THRDS) + wvSplitKQ_hf_sml_(const int K, const int Kap, const int Kbp, const int M, + const int Bx, const int By, const fp8_t* B, + const fp8_t* __restrict__ A, + const scalar_t* __restrict__ BIAS, scalar_t* C, + const float* __restrict__ s_A, + const float* __restrict__ s_B, const int _WvPrGrp, + const int CuCount) { + constexpr int max_lds_len = LDS_SIZE; + using scalar8 = + __attribute__((__vector_size__((A_CHUNK / 4) * sizeof(float)))) float; + using intx2 = __attribute__((__vector_size__(2 * sizeof(int)))) int; + using intx4 = __attribute__((__vector_size__(4 * sizeof(int)))) int; + union bigType { + char f8[A_CHUNK]; + char2 c2[A_CHUNK / 2]; + scalar_t h[A_CHUNK / 2]; + float f[A_CHUNK / 4]; + int i[A_CHUNK / 4]; + long l[A_CHUNK / 8]; + intx4 l2[A_CHUNK / 16]; + scalar8 h8; + }; + + __shared__ fp8_t s[max_lds_len]; + + for (uint32_t k = (threadIdx.y * THRDS + threadIdx.x) * A_CHUNK; + k < min__(Kap * N, max_lds_len); k += THRDS * WvPrGrp * A_CHUNK) { + #if defined(__gfx950__) + __builtin_amdgcn_global_load_lds((int*)(&A[k]), (int*)(&s[k]), 16, 0, 0); + #else + *((bigType*)(&s[k])) = *((bigType*)(&A[k])); + #endif + } + asm volatile("s_waitcnt vmcnt(0)"); + __syncthreads(); + + if (threadIdx.y >= _WvPrGrp) return; + + uint32_t m = (blockIdx.x * _WvPrGrp + (threadIdx.y % _WvPrGrp)) * YTILE; + + float sA = *s_A; + float sB = *s_B; + + while (m < M) { + #ifdef __HIP__GFX12__ + // gfx12: per-lane scalar accumulation via v_dot4_f32_fp8_fp8 + float sum[N][YTILE] = {}; + #else + // gfx9: MFMA accumulation + scalar8 sum[N][YTILE] = {}; + #endif + for (uint32_t k1 = 0; k1 < K; k1 += THRDS * A_CHUNK * UNRL) { + bigType bigA[N][UNRL] = {}; + bigType bigB[YTILE][UNRL]; + + // Fetch the weight matrix from memory! + #pragma unroll + for (uint32_t k2 = 0; k2 < UNRL; k2++) { + uint32_t k = k1 + k2 * THRDS * A_CHUNK; + uint32_t k_ = k + threadIdx.x * A_CHUNK; + const fp8_t* B_ = &B[min__(k_, K - A_CHUNK)]; + #pragma unroll + for (uint32_t y = 0; y < YTILE; ++y) { + bigB[y][k2].h8 = (loadnt((scalar8*)(&B_[min__(y + m, M - 1) * Kbp]))); + } + } + + // Fetch activation matrix from either just LDS or from both LDS / memory + #pragma unroll + for (uint32_t k2 = 0; k2 < UNRL; k2++) { + uint32_t k = k1 + k2 * THRDS * A_CHUNK; + uint32_t k_ = k + threadIdx.x * A_CHUNK; + if (k_ >= K) break; + for (int n = 0; n < N; n++) { + bigA[n][k2] = *((const bigType*)(&(s[k_ + Kap * n]))); + } + } + + // Do the matrix multiplication in interleaved manner + #pragma unroll + for (uint32_t k2 = 0; k2 < UNRL; k2++) { + for (uint32_t n = 0; n < N; n++) { + #ifdef __HIP__GFX12__ + // gfx12: 4 x dot4 per A_CHUNK=16 bytes (4 FP8 per dot4) + for (int y = 0; y < YTILE; ++y) { + #pragma unroll + for (int i = 0; i < A_CHUNK / 4; i++) { + sum[n][y] = __builtin_amdgcn_dot4_f32_fp8_fp8( + bigA[n][k2].i[i], bigB[y][k2].i[i], sum[n][y]); + } + } + #else + // gfx9: MFMA path + for (int i = 0; i < A_CHUNK; i += 8) { + for (int y = 0; y < YTILE; ++y) { + sum[n][y] = __builtin_amdgcn_mfma_f32_16x16x32_fp8_fp8( + bigA[n][k2].l[i / 8], bigB[y][k2].l[i / 8], sum[n][y], 0, 0, + 0); + } + } + #endif + } + } + } + + // Final reduction + #ifdef __HIP__GFX12__ + // gfx12 wave32: DPP row_shr within 16-lane rows + cross-row shuffle + for (int n = 0; n < N; n++) { + for (int y = 0; y < YTILE; y++) { + asm("s_nop 0\n\tv_add_f32 %0, %2, %3 row_shr:8 bound_ctrl:0 " + : "=v"(sum[n][y]) + : "0"(sum[n][y]), "v"(sum[n][y]), "v"(sum[n][y])); + asm("s_nop 0\n\tv_add_f32 %0, %2, %3 row_shr:4 bound_ctrl:0 " + : "=v"(sum[n][y]) + : "0"(sum[n][y]), "v"(sum[n][y]), "v"(sum[n][y])); + asm("s_nop 0\n\tv_add_f32 %0, %2, %3 row_shr:2 bound_ctrl:0 " + : "=v"(sum[n][y]) + : "0"(sum[n][y]), "v"(sum[n][y]), "v"(sum[n][y])); + asm("s_nop 0\n\tv_add_f32 %0, %2, %3 row_shr:1 bound_ctrl:0 " + : "=v"(sum[n][y]) + : "0"(sum[n][y]), "v"(sum[n][y]), "v"(sum[n][y])); + sum[n][y] += __shfl_xor(sum[n][y], 16); + } + } + #else + // gfx9 MFMA reduction + for (int n = 0; n < N; n++) { + for (int y = 0; y < YTILE; y++) { + float accm0 = sum[n][y][0]; + accm0 += __builtin_amdgcn_mov_dpp(sum[n][y][1], 0x101, 0xf, 0xf, + 1); // row_shl1 + accm0 += __builtin_amdgcn_mov_dpp(sum[n][y][2], 0x102, 0xf, 0xf, + 1); // row_shl2 + accm0 += __builtin_amdgcn_mov_dpp(sum[n][y][3], 0x103, 0xf, 0xf, + 1); // row_shl3 + accm0 += __shfl_down(accm0, 20); + accm0 += __shfl_down(accm0, 40); + sum[n][y][0] = accm0; + } + } + #endif + + const bool writeback_lane = + #ifdef __HIP__GFX12__ + threadIdx.x == (THRDS - 1); + #else + threadIdx.x == 0; + #endif + if (writeback_lane) { + scalar_t biases[N][YTILE] = {}; + if (BIAS) + for (int n = 0; n < N; n++) { + for (int y = 0; y < YTILE; y++) { + biases[n][y] = BIAS[(m + y) % Bx + (n % By) * Bx]; + } + } + for (int n = 0; n < N; n++) { + for (int y = 0; y < YTILE; y++) { + if (y + m >= M) break; // To avoid mem access fault. + #ifdef __HIP__GFX12__ + float result = sum[n][y] * sA * sB; + #else + float result = sum[n][y][0] * sA * sB; + #endif + if constexpr (std::is_same_v) { + result += __half2float(biases[n][y]); + } else if constexpr (std::is_same_v) { + result += __bfloat162float(biases[n][y]); + } + C[m + y + n * M] = __float2s(result); + } + } + } + + m += CuCount * _WvPrGrp * YTILE; + } +} +#else // !defined(__HIP__MI3XX__) && !defined(__HIP__GFX12__) +template +__global__ void wvSplitKQ_hf_sml_(const int K, const int Kap, const int Kbp, + const int M, const int Bx, const int By, + const fp8_t* B, const fp8_t* __restrict__ A, + const scalar_t* __restrict__ BIAS, + scalar_t* C, const float* __restrict__ s_A, + const float* __restrict__ s_B, + const int _WvPrGrp, const int CuCount) { + UNREACHABLE_CODE +} +#endif // defined(__HIP__MI3XX__) || defined(__HIP__GFX12__) + +#if defined(__HIP__MI3XX__) || defined(__HIP__GFX12__) +template +__global__ void __launch_bounds__(WvPrGrp* THRDS) + wvSplitKQ_hf_(const int K, const int Kap, const int Kbp, const int M, + const int Bx, const int By, const fp8_t* B, + const fp8_t* __restrict__ A, + const scalar_t* __restrict__ BIAS, scalar_t* C, + const float* __restrict__ s_A, const float* __restrict__ s_B, + const int _WvPrGrp, const int CuCount) { + constexpr int max_lds_len = LDS_SIZE; + using scalar8 = + __attribute__((__vector_size__((A_CHUNK / 4) * sizeof(float)))) float; + using intx2 = __attribute__((__vector_size__(2 * sizeof(int)))) int; + using intx4 = __attribute__((__vector_size__(4 * sizeof(int)))) int; + union bigType { + char f8[A_CHUNK]; + char2 c2[A_CHUNK / 2]; + scalar_t h[A_CHUNK / 2]; + float f[A_CHUNK / 4]; + int i[A_CHUNK / 4]; + long l[A_CHUNK / 8]; + intx4 l2[A_CHUNK / 16]; + scalar8 h8; + }; + + __shared__ fp8_t s[max_lds_len]; + + for (uint32_t k = (threadIdx.y * THRDS + threadIdx.x) * A_CHUNK; + k < min__(Kap * N, max_lds_len); k += THRDS * WvPrGrp * A_CHUNK) { + #if defined(__gfx950__) + __builtin_amdgcn_global_load_lds((int*)(&A[k]), (int*)(&s[k]), 16, 0, 0); + #else + *((bigType*)(&s[k])) = *((bigType*)(&A[k])); + #endif + } + asm volatile("s_waitcnt vmcnt(0)"); + __syncthreads(); + + if (threadIdx.y >= _WvPrGrp) return; + + uint32_t m = (blockIdx.x * _WvPrGrp + (threadIdx.y % _WvPrGrp)) * YTILE; + + float sA = *s_A; + float sB = *s_B; + + while (m < M) { + #ifdef __HIP__GFX12__ + // gfx12: per-lane scalar accumulation via v_dot4_f32_fp8_fp8 + float sum[N][YTILE] = {}; + #else + // gfx9: MFMA accumulation + scalar8 sum[N][YTILE] = {}; + #endif + for (uint32_t k1 = 0; k1 < K; k1 += THRDS * A_CHUNK * UNRL) { + bigType bigA[N][UNRL] = {}; + bigType bigB[YTILE][UNRL]; + + // Fetch the weight matrix from memory! + #pragma unroll + for (uint32_t k2 = 0; k2 < UNRL; k2++) { + uint32_t k = k1 + k2 * THRDS * A_CHUNK; + uint32_t k_ = k + threadIdx.x * A_CHUNK; + const fp8_t* B_ = &B[min__(k_, K - A_CHUNK)]; + for (int y = 0; y < YTILE; ++y) { + bigB[y][k2].h8 = (loadnt((scalar8*)(&B_[min__(y + m, M - 1) * Kbp]))); + } + } + + // Fetch activation matrix from either just LDS or from both LDS / memory + #pragma unroll + for (uint32_t k2 = 0; k2 < UNRL; k2++) { + uint32_t k = k1 + k2 * THRDS * A_CHUNK; + uint32_t k_ = k + threadIdx.x * A_CHUNK; + if (k_ >= K) break; + for (int n = 0; n < N; n++) { + if (k_ + Kap * n < max_lds_len) + bigA[n][k2] = *((const bigType*)(&(s[k_ + Kap * n]))); + else + bigA[n][k2] = *((const bigType*)(&(A[k_ + Kap * n]))); + } + } + + // Do the matrix multiplication in interleaved manner + #pragma unroll + for (uint32_t k2 = 0; k2 < UNRL; k2++) { + for (uint32_t n = 0; n < N; n++) { + #ifdef __HIP__GFX12__ + // gfx12: 4 x dot4 per A_CHUNK=16 bytes (4 FP8 per dot4) + for (int y = 0; y < YTILE; ++y) { + #pragma unroll + for (int i = 0; i < A_CHUNK / 4; i++) { + sum[n][y] = __builtin_amdgcn_dot4_f32_fp8_fp8( + bigA[n][k2].i[i], bigB[y][k2].i[i], sum[n][y]); + } + } + #else + // gfx9: MFMA path + for (int i = 0; i < A_CHUNK; i += 8) { + for (int y = 0; y < YTILE; ++y) { + sum[n][y] = __builtin_amdgcn_mfma_f32_16x16x32_fp8_fp8( + bigA[n][k2].l[i / 8], bigB[y][k2].l[i / 8], sum[n][y], 0, 0, + 0); + } + } + #endif + } + } + } + + // Final reduction + #ifdef __HIP__GFX12__ + // gfx12 wave32: DPP row_shr within 16-lane rows + cross-row shuffle + for (int n = 0; n < N; n++) { + for (int y = 0; y < YTILE; y++) { + asm("s_nop 0\n\tv_add_f32 %0, %2, %3 row_shr:8 bound_ctrl:0 " + : "=v"(sum[n][y]) + : "0"(sum[n][y]), "v"(sum[n][y]), "v"(sum[n][y])); + asm("s_nop 0\n\tv_add_f32 %0, %2, %3 row_shr:4 bound_ctrl:0 " + : "=v"(sum[n][y]) + : "0"(sum[n][y]), "v"(sum[n][y]), "v"(sum[n][y])); + asm("s_nop 0\n\tv_add_f32 %0, %2, %3 row_shr:2 bound_ctrl:0 " + : "=v"(sum[n][y]) + : "0"(sum[n][y]), "v"(sum[n][y]), "v"(sum[n][y])); + asm("s_nop 0\n\tv_add_f32 %0, %2, %3 row_shr:1 bound_ctrl:0 " + : "=v"(sum[n][y]) + : "0"(sum[n][y]), "v"(sum[n][y]), "v"(sum[n][y])); + sum[n][y] += __shfl_xor(sum[n][y], 16); + } + } + #else + // gfx9 MFMA reduction + for (int n = 0; n < N; n++) { + for (int y = 0; y < YTILE; y++) { + float accm0 = sum[n][y][0]; + accm0 += __builtin_amdgcn_mov_dpp(sum[n][y][1], 0x101, 0xf, 0xf, + 1); // row_shl1 + accm0 += __builtin_amdgcn_mov_dpp(sum[n][y][2], 0x102, 0xf, 0xf, + 1); // row_shl2 + accm0 += __builtin_amdgcn_mov_dpp(sum[n][y][3], 0x103, 0xf, 0xf, + 1); // row_shl3 + accm0 += __shfl_down(accm0, 20); + accm0 += __shfl_down(accm0, 40); + sum[n][y][0] = accm0; + } + } + #endif + + const bool writeback_lane = + #ifdef __HIP__GFX12__ + threadIdx.x == (THRDS - 1); + #else + threadIdx.x == 0; + #endif + if (writeback_lane) { + scalar_t biases[N][YTILE] = {}; + if (BIAS) + for (int n = 0; n < N; n++) { + for (int y = 0; y < YTILE; y++) { + biases[n][y] = BIAS[(m + y) % Bx + (n % By) * Bx]; + } + } + for (int n = 0; n < N; n++) { + for (int y = 0; y < YTILE; y++) { + if (y + m >= M) break; // To avoid mem access fault. + #ifdef __HIP__GFX12__ + float result = sum[n][y] * sA * sB; + #else + float result = sum[n][y][0] * sA * sB; + #endif + if constexpr (std::is_same_v) { + result += __half2float(biases[n][y]); + } else if constexpr (std::is_same_v) { + result += __bfloat162float(biases[n][y]); + } + C[m + y + n * M] = __float2s(result); + } + } + } + + m += CuCount * _WvPrGrp * YTILE; + } +} +#else // !defined(__HIP__MI3XX__) && !defined(__HIP__GFX12__) +template +__global__ void wvSplitKQ_hf_(const int K, const int Kap, const int Kbp, + const int M, const int Bx, const int By, + const fp8_t* B, const fp8_t* __restrict__ A, + const scalar_t* __restrict__ BIAS, scalar_t* C, + const float* __restrict__ s_A, + const float* __restrict__ s_B, const int _WvPrGrp, + const int CuCount) { + UNREACHABLE_CODE +} +#endif // defined(__HIP__MI3XX__) || defined(__HIP__GFX12__) + +void wvSplitKQ(const at::Tensor& in_b, const at::Tensor& in_a, + const std::optional& in_bias, at::Tensor& out_c, + const at::Tensor& scale_a, const at::Tensor& scale_b, + const int64_t CuCount) { + static c10::ScalarType kFp8Type = is_fp8_ocp() + ? c10::ScalarType::Float8_e4m3fn + : c10::ScalarType::Float8_e4m3fnuz; + auto M_in = in_b.size(0); + auto K_in = in_b.size(1); + auto N_in = in_a.size(0); + auto Kap_in = in_a.stride(0); + auto Kbp_in = in_b.stride(0); + auto Bx_in = + (in_bias.has_value() && in_bias->numel() > 0) + ? (in_bias->sizes().size() == 2) ? in_bias->size(1) : in_bias->size(0) + : 1; + auto By_in = (in_bias.has_value() && in_bias->numel() > 0 && + in_bias->sizes().size() == 2) + ? in_bias->size(0) + : 1; + + TORCH_CHECK(K_in % 16 == 0, "k % 16 == 0"); + TORCH_CHECK(in_a.dtype() == in_b.dtype() && in_a.dtype() == kFp8Type); + TORCH_CHECK(out_c.dtype() == torch::kFloat16 || + out_c.dtype() == torch::kBFloat16); + + dim3 grid(CuCount); + const at::cuda::OptionalCUDAGuard device_guard(device_of(in_a)); + const cudaStream_t stream = at::cuda::getCurrentCUDAStream(); + const int max_lds_len = get_lds_size(); + +#define WVSPLITKQ_IMPL(_THRDS, _WvPrGrp, _YTILEs, _YTILEm, _UNRLs, _UNRLm, _N) \ + { \ + dim3 block(_THRDS, _WvPrGrp); \ + if ((Kap_in * N_in <= max_lds_len) && (M_in % _YTILEs == 0)) { \ + int __wvPrGrp = min(_WvPrGrp, mindiv(M_in, CuCount * _YTILEs, 16)); \ + wvSplitKQ_hf_sml_<<>>( \ + K_in, Kap_in, Kbp_in, M_in, Bx_in, By_in, b_ptr, a_ptr, bias_ptr, \ + c_ptr, s_a, s_b, __wvPrGrp, CuCount); \ + } else { \ + int __wvPrGrp = min(_WvPrGrp, mindiv(M_in, CuCount * _YTILEm, 16)); \ + wvSplitKQ_hf_ \ + <<>>(K_in, Kap_in, Kbp_in, M_in, Bx_in, \ + By_in, b_ptr, a_ptr, bias_ptr, c_ptr, \ + s_a, s_b, __wvPrGrp, CuCount); \ + } \ + } + +#define WVSPLITKQ(_WvPrGrp, _YTILEs, _YTILEm, _UNRLs, _UNRLm, _N) \ + if (on_gfx12()) \ + WVSPLITKQ_IMPL(32, _WvPrGrp, _YTILEs, _YTILEm, _UNRLs, _UNRLm, _N) \ + else \ + WVSPLITKQ_IMPL(64, _WvPrGrp, _YTILEs, _YTILEm, _UNRLs, _UNRLm, _N) + + AT_DISPATCH_REDUCED_FLOATING_TYPES(out_c.scalar_type(), "wvSplitKQ", [&] { + using fptype = typename scalar::type; + auto c_ptr = reinterpret_cast(out_c.data_ptr()); + auto s_a = scale_a.data_ptr(); + auto s_b = scale_b.data_ptr(); + VLLM_DISPATCH_FP8_TYPES(in_a.scalar_type(), "wvSplitKQ", [&] { + auto a_ptr = in_a.data_ptr(); + auto b_ptr = in_b.data_ptr(); + auto bias_ptr = (in_bias.has_value() && in_bias->numel() > 0) + ? reinterpret_cast(in_bias->data_ptr()) + : nullptr; + switch (N_in) { + case 1: + WVSPLITKQ(16, 2, 2, 2, 2, 1) + break; + case 2: + WVSPLITKQ(16, 2, 2, 2, 2, 2) + break; + case 3: + WVSPLITKQ(16, 2, 2, 1, 1, 3) + break; + case 4: + WVSPLITKQ(16, 2, 2, 1, 1, 4) + break; + default: + throw std::runtime_error( + "Unsupported N value: " + std::to_string(M_in) + "," + + std::to_string(K_in) + "," + std::to_string(N_in)); + } + }); + }); +} diff --git a/tasks/hip2hip/campaign20/wvsplitk/test_cases.json b/tasks/hip2hip/campaign20/wvsplitk/test_cases.json new file mode 100644 index 00000000..fa84d26a --- /dev/null +++ b/tasks/hip2hip/campaign20/wvsplitk/test_cases.json @@ -0,0 +1,993 @@ +[ + { + "test_case_id": "sig_91eb270fe708", + "count": 95880, + "args_sig": [ + { + "kind": "tensor", + "dtype": "bfloat16", + "shape": [ + 128, + 4096 + ], + "stride": [ + 4096, + 1 + ], + "device": "cuda:0" + }, + { + "kind": "tensor", + "dtype": "bfloat16", + "shape": [ + 1, + 4096 + ], + "stride": [ + 4096, + 1 + ], + "device": "cuda:0" + }, + { + "kind": "scalar", + "value": null, + "py_type": "NoneType" + }, + { + "kind": "scalar", + "value": 304, + "py_type": "int" + } + ], + "kwargs_sig": {}, + "params_repr": {}, + "args_names": [ + "in_a", + "in_b", + "in_bias", + "CuCount" + ], + "sources": [ + { + "model": "Qwen_Qwen3-235B-A22B-FP8", + "engine": "vllm", + "mode": "default", + "combo": "Qwen_Qwen3-235B-A22B-FP8__vllm_default", + "count": 95880 + } + ] + }, + { + "test_case_id": "sig_a6f816730e0a", + "count": 1024, + "args_sig": [ + { + "kind": "tensor", + "dtype": "bfloat16", + "shape": [ + 37984, + 4096 + ], + "stride": [ + 4096, + 1 + ], + "device": "cuda:0" + }, + { + "kind": "tensor", + "dtype": "bfloat16", + "shape": [ + 1, + 4096 + ], + "stride": [ + 4096, + 1 + ], + "device": "cuda:0" + }, + { + "kind": "scalar", + "value": null, + "py_type": "NoneType" + }, + { + "kind": "scalar", + "value": 304, + "py_type": "int" + } + ], + "kwargs_sig": {}, + "params_repr": {}, + "args_names": [ + "in_a", + "in_b", + "in_bias", + "CuCount" + ], + "sources": [ + { + "model": "Qwen_Qwen3-235B-A22B-FP8", + "engine": "vllm", + "mode": "default", + "combo": "Qwen_Qwen3-235B-A22B-FP8__vllm_default", + "count": 1024 + } + ] + }, + { + "test_case_id": "sig_34c6787c2d55", + "count": 98304, + "args_sig": [ + { + "kind": "tensor", + "dtype": "bfloat16", + "shape": [ + 512, + 2048 + ], + "stride": [ + 2048, + 1 + ], + "device": "cuda:0" + }, + { + "kind": "tensor", + "dtype": "bfloat16", + "shape": [ + 1, + 2048 + ], + "stride": [ + 2048, + 1 + ], + "device": "cuda:0" + }, + { + "kind": "scalar", + "value": null, + "py_type": "NoneType" + }, + { + "kind": "scalar", + "value": 304, + "py_type": "int" + } + ], + "kwargs_sig": {}, + "params_repr": {}, + "args_names": [ + "in_a", + "in_b", + "in_bias", + "CuCount" + ], + "sources": [ + { + "model": "Qwen_Qwen3-Next-80B-A3B-Instruct-FP8", + "engine": "vllm", + "mode": "default", + "combo": "Qwen_Qwen3-Next-80B-A3B-Instruct-FP8__vllm_default", + "count": 49152 + }, + { + "model": "Qwen_Qwen3-Next-80B-A3B-Instruct-FP8", + "engine": "vllm", + "mode": "hip", + "combo": "Qwen_Qwen3-Next-80B-A3B-Instruct-FP8__vllm_hip", + "count": 49152 + } + ] + }, + { + "test_case_id": "sig_f376ca31172d", + "count": 73728, + "args_sig": [ + { + "kind": "tensor", + "dtype": "bfloat16", + "shape": [ + 16, + 2048 + ], + "stride": [ + 2048, + 1 + ], + "device": "cuda:0" + }, + { + "kind": "tensor", + "dtype": "bfloat16", + "shape": [ + 1, + 2048 + ], + "stride": [ + 2048, + 1 + ], + "device": "cuda:0" + }, + { + "kind": "scalar", + "value": null, + "py_type": "NoneType" + }, + { + "kind": "scalar", + "value": 304, + "py_type": "int" + } + ], + "kwargs_sig": {}, + "params_repr": {}, + "args_names": [ + "in_a", + "in_b", + "in_bias", + "CuCount" + ], + "sources": [ + { + "model": "Qwen_Qwen3-Next-80B-A3B-Instruct-FP8", + "engine": "vllm", + "mode": "default", + "combo": "Qwen_Qwen3-Next-80B-A3B-Instruct-FP8__vllm_default", + "count": 36864 + }, + { + "model": "Qwen_Qwen3-Next-80B-A3B-Instruct-FP8", + "engine": "vllm", + "mode": "hip", + "combo": "Qwen_Qwen3-Next-80B-A3B-Instruct-FP8__vllm_hip", + "count": 36864 + } + ] + }, + { + "test_case_id": "sig_ae3e0c7506c3", + "count": 2056, + "args_sig": [ + { + "kind": "tensor", + "dtype": "bfloat16", + "shape": [ + 37984, + 2048 + ], + "stride": [ + 2048, + 1 + ], + "device": "cuda:0" + }, + { + "kind": "tensor", + "dtype": "bfloat16", + "shape": [ + 1, + 2048 + ], + "stride": [ + 2048, + 1 + ], + "device": "cuda:0" + }, + { + "kind": "scalar", + "value": null, + "py_type": "NoneType" + }, + { + "kind": "scalar", + "value": 304, + "py_type": "int" + } + ], + "kwargs_sig": {}, + "params_repr": {}, + "args_names": [ + "in_a", + "in_b", + "in_bias", + "CuCount" + ], + "sources": [ + { + "model": "Qwen_Qwen3-Next-80B-A3B-Instruct-FP8", + "engine": "vllm", + "mode": "default", + "combo": "Qwen_Qwen3-Next-80B-A3B-Instruct-FP8__vllm_default", + "count": 1028 + }, + { + "model": "Qwen_Qwen3-Next-80B-A3B-Instruct-FP8", + "engine": "vllm", + "mode": "hip", + "combo": "Qwen_Qwen3-Next-80B-A3B-Instruct-FP8__vllm_hip", + "count": 1028 + } + ] + }, + { + "test_case_id": "sig_06de0e61e142", + "count": 24480, + "args_sig": [ + { + "kind": "tensor", + "dtype": "bfloat16", + "shape": [ + 128, + 2048 + ], + "stride": [ + 2048, + 1 + ], + "device": "cuda:0" + }, + { + "kind": "tensor", + "dtype": "bfloat16", + "shape": [ + 1, + 2048 + ], + "stride": [ + 2048, + 1 + ], + "device": "cuda:0" + }, + { + "kind": "scalar", + "value": null, + "py_type": "NoneType" + }, + { + "kind": "scalar", + "value": 304, + "py_type": "int" + } + ], + "kwargs_sig": {}, + "params_repr": {}, + "args_names": [ + "in_a", + "in_b", + "in_bias", + "CuCount" + ], + "sources": [ + { + "model": "Qwen_Qwen3.5-35B-A3B", + "engine": "vllm", + "mode": "default", + "combo": "Qwen_Qwen3.5-35B-A3B__vllm_default", + "count": 12240 + }, + { + "model": "Qwen_Qwen3.5-35B-A3B", + "engine": "vllm", + "mode": "hip", + "combo": "Qwen_Qwen3.5-35B-A3B__vllm_hip", + "count": 12240 + } + ] + }, + { + "test_case_id": "sig_f97e2b0e3611", + "count": 512, + "args_sig": [ + { + "kind": "tensor", + "dtype": "bfloat16", + "shape": [ + 151936, + 2048 + ], + "stride": [ + 2048, + 1 + ], + "device": "cuda:0" + }, + { + "kind": "tensor", + "dtype": "bfloat16", + "shape": [ + 1, + 2048 + ], + "stride": [ + 2048, + 1 + ], + "device": "cuda:0" + }, + { + "kind": "scalar", + "value": null, + "py_type": "NoneType" + }, + { + "kind": "scalar", + "value": 304, + "py_type": "int" + } + ], + "kwargs_sig": {}, + "params_repr": {}, + "args_names": [ + "in_a", + "in_b", + "in_bias", + "CuCount" + ], + "sources": [ + { + "model": "Qwen_Qwen3.5-35B-A3B", + "engine": "vllm", + "mode": "default", + "combo": "Qwen_Qwen3.5-35B-A3B__vllm_default", + "count": 256 + }, + { + "model": "Qwen_Qwen3.5-35B-A3B", + "engine": "vllm", + "mode": "hip", + "combo": "Qwen_Qwen3.5-35B-A3B__vllm_hip", + "count": 256 + } + ] + }, + { + "test_case_id": "sig_a6a0fd1cdc57", + "count": 709920, + "args_sig": [ + { + "kind": "tensor", + "dtype": "bfloat16", + "shape": [ + 256, + 7168 + ], + "stride": [ + 7168, + 1 + ], + "device": "cuda:0" + }, + { + "kind": "tensor", + "dtype": "bfloat16", + "shape": [ + 1, + 7168 + ], + "stride": [ + 7168, + 1 + ], + "device": "cuda:0" + }, + { + "kind": "scalar", + "value": null, + "py_type": "NoneType" + }, + { + "kind": "scalar", + "value": 304, + "py_type": "int" + } + ], + "kwargs_sig": {}, + "params_repr": {}, + "args_names": [ + "in_a", + "in_b", + "in_bias", + "CuCount" + ], + "sources": [ + { + "model": "deepseek-ai_DeepSeek-R1", + "engine": "vllm", + "mode": "default", + "combo": "deepseek-ai_DeepSeek-R1__vllm_default", + "count": 118320 + }, + { + "model": "deepseek-ai_DeepSeek-R1", + "engine": "vllm", + "mode": "hip", + "combo": "deepseek-ai_DeepSeek-R1__vllm_hip", + "count": 118320 + }, + { + "model": "deepseek-ai_DeepSeek-R1", + "engine": "vllm", + "mode": "triton_mla", + "combo": "deepseek-ai_DeepSeek-R1__vllm_triton_mla", + "count": 118320 + }, + { + "model": "deepseek-ai_DeepSeek-V3.2", + "engine": "vllm", + "mode": "default", + "combo": "deepseek-ai_DeepSeek-V3.2__vllm_default", + "count": 118320 + }, + { + "model": "deepseek-ai_DeepSeek-V3.2", + "engine": "vllm", + "mode": "hip", + "combo": "deepseek-ai_DeepSeek-V3.2__vllm_hip", + "count": 118320 + }, + { + "model": "deepseek-ai_DeepSeek-V3.2", + "engine": "vllm", + "mode": "triton_mla", + "combo": "deepseek-ai_DeepSeek-V3.2__vllm_triton_mla", + "count": 118320 + } + ] + }, + { + "test_case_id": "sig_c423f8c6d387", + "count": 12288, + "args_sig": [ + { + "kind": "tensor", + "dtype": "bfloat16", + "shape": [ + 16160, + 7168 + ], + "stride": [ + 7168, + 1 + ], + "device": "cuda:0" + }, + { + "kind": "tensor", + "dtype": "bfloat16", + "shape": [ + 1, + 7168 + ], + "stride": [ + 7168, + 1 + ], + "device": "cuda:0" + }, + { + "kind": "scalar", + "value": null, + "py_type": "NoneType" + }, + { + "kind": "scalar", + "value": 304, + "py_type": "int" + } + ], + "kwargs_sig": {}, + "params_repr": {}, + "args_names": [ + "in_a", + "in_b", + "in_bias", + "CuCount" + ], + "sources": [ + { + "model": "deepseek-ai_DeepSeek-R1", + "engine": "vllm", + "mode": "default", + "combo": "deepseek-ai_DeepSeek-R1__vllm_default", + "count": 2048 + }, + { + "model": "deepseek-ai_DeepSeek-R1", + "engine": "vllm", + "mode": "hip", + "combo": "deepseek-ai_DeepSeek-R1__vllm_hip", + "count": 2048 + }, + { + "model": "deepseek-ai_DeepSeek-R1", + "engine": "vllm", + "mode": "triton_mla", + "combo": "deepseek-ai_DeepSeek-R1__vllm_triton_mla", + "count": 2048 + }, + { + "model": "deepseek-ai_DeepSeek-V3.2", + "engine": "vllm", + "mode": "default", + "combo": "deepseek-ai_DeepSeek-V3.2__vllm_default", + "count": 2048 + }, + { + "model": "deepseek-ai_DeepSeek-V3.2", + "engine": "vllm", + "mode": "hip", + "combo": "deepseek-ai_DeepSeek-V3.2__vllm_hip", + "count": 2048 + }, + { + "model": "deepseek-ai_DeepSeek-V3.2", + "engine": "vllm", + "mode": "triton_mla", + "combo": "deepseek-ai_DeepSeek-V3.2__vllm_triton_mla", + "count": 2048 + } + ] + }, + { + "test_case_id": "sig_228a0a69990f", + "count": 373320, + "args_sig": [ + { + "kind": "tensor", + "dtype": "bfloat16", + "shape": [ + 64, + 7168 + ], + "stride": [ + 7168, + 1 + ], + "device": "cuda:0" + }, + { + "kind": "tensor", + "dtype": "bfloat16", + "shape": [ + 1, + 7168 + ], + "stride": [ + 7168, + 1 + ], + "device": "cuda:0" + }, + { + "kind": "scalar", + "value": null, + "py_type": "NoneType" + }, + { + "kind": "scalar", + "value": 304, + "py_type": "int" + } + ], + "kwargs_sig": {}, + "params_repr": {}, + "args_names": [ + "in_a", + "in_b", + "in_bias", + "CuCount" + ], + "sources": [ + { + "model": "deepseek-ai_DeepSeek-V3.2", + "engine": "vllm", + "mode": "default", + "combo": "deepseek-ai_DeepSeek-V3.2__vllm_default", + "count": 124440 + }, + { + "model": "deepseek-ai_DeepSeek-V3.2", + "engine": "vllm", + "mode": "hip", + "combo": "deepseek-ai_DeepSeek-V3.2__vllm_hip", + "count": 124440 + }, + { + "model": "deepseek-ai_DeepSeek-V3.2", + "engine": "vllm", + "mode": "triton_mla", + "combo": "deepseek-ai_DeepSeek-V3.2__vllm_triton_mla", + "count": 124440 + } + ] + }, + { + "test_case_id": "sig_291b8dd0b92c", + "count": 18360, + "args_sig": [ + { + "kind": "tensor", + "dtype": "bfloat16", + "shape": [ + 5120, + 2880 + ], + "stride": [ + 2880, + 1 + ], + "device": "cuda:0" + }, + { + "kind": "tensor", + "dtype": "bfloat16", + "shape": [ + 1, + 2880 + ], + "stride": [ + 2880, + 1 + ], + "device": "cuda:0" + }, + { + "kind": "tensor", + "dtype": "bfloat16", + "shape": [ + 5120 + ], + "stride": [ + 1 + ], + "device": "cuda:0" + }, + { + "kind": "scalar", + "value": 304, + "py_type": "int" + } + ], + "kwargs_sig": {}, + "params_repr": {}, + "args_names": [ + "in_a", + "in_b", + "in_bias", + "CuCount" + ], + "sources": [ + { + "model": "openai_gpt-oss-120b", + "engine": "vllm", + "mode": "default", + "combo": "openai_gpt-oss-120b__vllm_default", + "count": 9180 + }, + { + "model": "openai_gpt-oss-120b", + "engine": "vllm", + "mode": "hip", + "combo": "openai_gpt-oss-120b__vllm_hip", + "count": 9180 + } + ] + }, + { + "test_case_id": "sig_d257625fefea", + "count": 18360, + "args_sig": [ + { + "kind": "tensor", + "dtype": "bfloat16", + "shape": [ + 2880, + 4096 + ], + "stride": [ + 4096, + 1 + ], + "device": "cuda:0" + }, + { + "kind": "tensor", + "dtype": "bfloat16", + "shape": [ + 1, + 4096 + ], + "stride": [ + 4096, + 1 + ], + "device": "cuda:0" + }, + { + "kind": "tensor", + "dtype": "bfloat16", + "shape": [ + 2880 + ], + "stride": [ + 1 + ], + "device": "cuda:0" + }, + { + "kind": "scalar", + "value": 304, + "py_type": "int" + } + ], + "kwargs_sig": {}, + "params_repr": {}, + "args_names": [ + "in_a", + "in_b", + "in_bias", + "CuCount" + ], + "sources": [ + { + "model": "openai_gpt-oss-120b", + "engine": "vllm", + "mode": "default", + "combo": "openai_gpt-oss-120b__vllm_default", + "count": 9180 + }, + { + "model": "openai_gpt-oss-120b", + "engine": "vllm", + "mode": "hip", + "combo": "openai_gpt-oss-120b__vllm_hip", + "count": 9180 + } + ] + }, + { + "test_case_id": "sig_feea1eea6fbf", + "count": 18360, + "args_sig": [ + { + "kind": "tensor", + "dtype": "bfloat16", + "shape": [ + 128, + 2880 + ], + "stride": [ + 2880, + 1 + ], + "device": "cuda:0" + }, + { + "kind": "tensor", + "dtype": "bfloat16", + "shape": [ + 1, + 2880 + ], + "stride": [ + 2880, + 1 + ], + "device": "cuda:0" + }, + { + "kind": "tensor", + "dtype": "bfloat16", + "shape": [ + 128 + ], + "stride": [ + 1 + ], + "device": "cuda:0" + }, + { + "kind": "scalar", + "value": 304, + "py_type": "int" + } + ], + "kwargs_sig": {}, + "params_repr": {}, + "args_names": [ + "in_a", + "in_b", + "in_bias", + "CuCount" + ], + "sources": [ + { + "model": "openai_gpt-oss-120b", + "engine": "vllm", + "mode": "default", + "combo": "openai_gpt-oss-120b__vllm_default", + "count": 9180 + }, + { + "model": "openai_gpt-oss-120b", + "engine": "vllm", + "mode": "hip", + "combo": "openai_gpt-oss-120b__vllm_hip", + "count": 9180 + } + ] + }, + { + "test_case_id": "sig_6688e58ae019", + "count": 512, + "args_sig": [ + { + "kind": "tensor", + "dtype": "bfloat16", + "shape": [ + 201088, + 2880 + ], + "stride": [ + 2880, + 1 + ], + "device": "cuda:0" + }, + { + "kind": "tensor", + "dtype": "bfloat16", + "shape": [ + 1, + 2880 + ], + "stride": [ + 2880, + 1 + ], + "device": "cuda:0" + }, + { + "kind": "scalar", + "value": null, + "py_type": "NoneType" + }, + { + "kind": "scalar", + "value": 304, + "py_type": "int" + } + ], + "kwargs_sig": {}, + "params_repr": {}, + "args_names": [ + "in_a", + "in_b", + "in_bias", + "CuCount" + ], + "sources": [ + { + "model": "openai_gpt-oss-120b", + "engine": "vllm", + "mode": "default", + "combo": "openai_gpt-oss-120b__vllm_default", + "count": 256 + }, + { + "model": "openai_gpt-oss-120b", + "engine": "vllm", + "mode": "hip", + "combo": "openai_gpt-oss-120b__vllm_hip", + "count": 256 + } + ] + } +] \ No newline at end of file diff --git a/tasks/triton2triton/campaign20/gemm_a16_w16_kernel/Makefile b/tasks/triton2triton/campaign20/gemm_a16_w16_kernel/Makefile new file mode 100644 index 00000000..0f8b4f36 --- /dev/null +++ b/tasks/triton2triton/campaign20/gemm_a16_w16_kernel/Makefile @@ -0,0 +1,6 @@ +# Auto-generated. +all: + python3 scripts/task_runner.py compile + +clean: + rm -rf build diff --git a/tasks/triton2triton/campaign20/gemm_a16_w16_kernel/config.yaml b/tasks/triton2triton/campaign20/gemm_a16_w16_kernel/config.yaml new file mode 100644 index 00000000..245b3c12 --- /dev/null +++ b/tasks/triton2triton/campaign20/gemm_a16_w16_kernel/config.yaml @@ -0,0 +1,29 @@ +source_file_path: +- source/triton__gemm_a16_w16_kernel.py +target_kernel_functions: +- _gemm_a16_w16_kernel +compile_command: +- python3 scripts/task_runner.py compile +correctness_command: +- python3 scripts/task_runner.py correctness +performance_command: +- python3 scripts/task_runner.py performance +task_type: triton2triton +platform_support: + required_arch: gfx950 + status: active + skip_reason: null +task_result_template: null +prompt: + source_code: null + instructions: Optimize the Triton kernel `_gemm_a16_w16_kernel` (originally aiter.ops.triton._triton_kernels.gemm.basic.gemm_a16w16). + Captured via inspect.getsource. Tests generate inputs from test_cases.json on + the fly. + cheatsheet: null +extracted_from: + engine: vllm + kernel_id: triton:aiter.ops.triton._triton_kernels.gemm.basic.gemm_a16w16._gemm_a16_w16_kernel + triton_fn: _gemm_a16_w16_kernel + origin_module: aiter.ops.triton._triton_kernels.gemm.basic.gemm_a16w16 + source: triton_inspect + test_cases: 3 diff --git a/tasks/triton2triton/campaign20/gemm_a16_w16_kernel/scripts/harness_run.py b/tasks/triton2triton/campaign20/gemm_a16_w16_kernel/scripts/harness_run.py new file mode 100644 index 00000000..caba8f43 --- /dev/null +++ b/tasks/triton2triton/campaign20/gemm_a16_w16_kernel/scripts/harness_run.py @@ -0,0 +1,227 @@ +#!/usr/bin/env python3 +"""Real launcher + benchmark harness for _gemm_a16_w16_kernel (triton2triton). + +Compile-only kernel: raw @triton.jit with no launcher. This module builds a +launcher from the captured arg schema, injects the helper symbols the kernel +references from its module globals (remap_xcd, pid_grid), regenerates the 3 +workload-regime test cases, runs golden-vs-editable correctness, and times +each case with CUDA events. + +Workload regime (token-parallel GEMM): + - num_tokens M = B*1024 for B in {2,32,64} + - model dims kept from captured base case: K=2880, N=5120 + - prefill token count semantics (seqlen=1024) +""" +import os +import json +import importlib.util + +import torch +import triton +import triton.language as tl + +TASK_DIR = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) +SOURCE_FILE = os.path.join(TASK_DIR, "source", "triton__gemm_a16_w16_kernel.py") +GOLDEN_FILE = os.path.join(TASK_DIR, "source_golden", "triton__gemm_a16_w16_kernel.py") +KERNEL_NAME = "_gemm_a16_w16_kernel" +TEST_CASES = os.path.join(TASK_DIR, "test_cases.json") +BUILD_DIR = os.path.join(TASK_DIR, "build") + +# ---- model dims from captured base case ---- +K_DIM = 2880 +N_DIM = 5120 +SEQLEN = 1024 +CONCURRENCY = [2, 32, 64] + +# ---- meta-params from the captured base case kwargs_sig ---- +META = dict( + BLOCK_SIZE_M=64, + BLOCK_SIZE_N=128, + BLOCK_SIZE_K=128, + GROUP_SIZE_M=1, + NUM_KSPLIT=1, + SPLITK_BLOCK_SIZE=K_DIM, # == K so EVEN_K stays consistent with single split + num_warps=4, + num_stages=2, + waves_per_eu=2, + matrix_instr_nonkdim=16, + cache_modifier=".cg", +) + + +def regime_cases(): + cases = [] + for b in CONCURRENCY: + m = b * SEQLEN + cases.append({ + "test_case_id": f"c{b}", + "B": b, + "M": m, + "N": N_DIM, + "K": K_DIM, + "params": {"B": b, "seqlen": SEQLEN, "M": m, "N": N_DIM, "K": K_DIM}, + }) + return cases + + +def write_test_cases(): + json.dump(regime_cases(), open(TEST_CASES, "w"), indent=2) + + +# ---- helper symbols the kernel references from module globals ---- +@triton.jit +def remap_xcd(pid, GRID_MN, NUM_XCDS: tl.constexpr = 8): # noqa: F821 + pids_per_xcd = (GRID_MN + NUM_XCDS - 1) // NUM_XCDS + tall_xcds = GRID_MN % NUM_XCDS + tall_xcds = NUM_XCDS if tall_xcds == 0 else tall_xcds + xcd = pid % NUM_XCDS + local_pid = pid // NUM_XCDS + if xcd < tall_xcds: + pid = xcd * pids_per_xcd + local_pid + else: + pid = ( + tall_xcds * pids_per_xcd + + (xcd - tall_xcds) * (pids_per_xcd - 1) + + local_pid + ) + return pid + + +@triton.jit +def pid_grid(pid, num_pid_m, num_pid_n, GROUP_SIZE_M: tl.constexpr = 1): # noqa: F821 + if GROUP_SIZE_M == 1: + pid_m = pid // num_pid_n + pid_n = pid % num_pid_n + else: + num_pid_in_group = GROUP_SIZE_M * num_pid_n + group_id = pid // num_pid_in_group + first_pid_m = group_id * GROUP_SIZE_M + group_size_m = min(num_pid_m - first_pid_m, GROUP_SIZE_M) + tl.assume(group_size_m >= 0) # noqa: F821 + pid_m = first_pid_m + (pid % group_size_m) + pid_n = (pid % num_pid_in_group) // group_size_m + return pid_m, pid_n + + +def _load_kernel(path): + spec = importlib.util.spec_from_file_location("k_" + str(abs(hash(path))), path) + mod = importlib.util.module_from_spec(spec) + spec.loader.exec_module(mod) + # inject helper symbols + tl alias into the kernel's module globals so the + # @triton.jit body can resolve them at JIT time. + import triton.language as _tl + inj = mod.__dict__ + inj.setdefault("tl", _tl) + inj.setdefault("triton", triton) + inj["remap_xcd"] = remap_xcd + inj["pid_grid"] = pid_grid + kern = getattr(mod, KERNEL_NAME) + # also patch the JITFunction's own global namespace + try: + kern.fn.__globals__.update({ + "remap_xcd": remap_xcd, + "pid_grid": pid_grid, + "tl": _tl, + "triton": triton, + }) + except Exception: + pass + return kern + + +def build_inputs(M, N, K, seed=42, dtype=torch.bfloat16, device="cuda"): + g = torch.Generator(device=device).manual_seed(seed) + # A: (M, K) row-major ; W stored (N, K) then transposed to (K, N) like the wrapper + a = torch.randn((M, K), generator=g, dtype=dtype, device=device) + w = torch.randn((N, K), generator=g, dtype=dtype, device=device) + b = w.T # (K, N), strides (1, K) + bias = torch.randn((N,), generator=g, dtype=dtype, device=device) + c = torch.empty((M, N), dtype=dtype, device=device) + return a, b, bias, c + + +def _grid(M, N, meta): + return (meta["NUM_KSPLIT"] + * triton.cdiv(M, meta["BLOCK_SIZE_M"]) + * triton.cdiv(N, meta["BLOCK_SIZE_N"]),) + + +def launch(kern, a, b, bias, c, M, N, K, meta): + grid = _grid(M, N, meta) + kern[grid]( + a, b, bias, c, + M, N, K, + a.stride(0), a.stride(1), + b.stride(0), b.stride(1), + 0, # stride_ck (NUM_KSPLIT==1) + c.stride(0), c.stride(1), + BLOCK_SIZE_M=meta["BLOCK_SIZE_M"], + BLOCK_SIZE_N=meta["BLOCK_SIZE_N"], + BLOCK_SIZE_K=meta["BLOCK_SIZE_K"], + GROUP_SIZE_M=meta["GROUP_SIZE_M"], + NUM_KSPLIT=meta["NUM_KSPLIT"], + SPLITK_BLOCK_SIZE=meta["SPLITK_BLOCK_SIZE"], + cache_modifier=meta["cache_modifier"], + activation="", + use_activation=False, + ADD_BIAS=True, + SKIP_REDUCE=False, + num_warps=meta["num_warps"], + num_stages=meta["num_stages"], + waves_per_eu=meta["waves_per_eu"], + matrix_instr_nonkdim=meta["matrix_instr_nonkdim"], + ) + return c + + +def _cos(x, y): + x = x.float().flatten() + y = y.float().flatten() + return torch.nn.functional.cosine_similarity(x, y, dim=0).item() + + +def run_correctness(): + edit = _load_kernel(SOURCE_FILE) + gold = _load_kernel(GOLDEN_FILE) + for tc in regime_cases(): + M, N, K = tc["M"], tc["N"], tc["K"] + ae = build_inputs(M, N, K, seed=42) + ag = build_inputs(M, N, K, seed=42) + out_e = launch(edit, *ae, M, N, K, META) + out_g = launch(gold, *ag, M, N, K, META) + torch.cuda.synchronize() + cos = _cos(out_e, out_g) + denom = out_g.float().abs().max().item() + max_rel = (out_e.float() - out_g.float()).abs().max().item() / (denom + 1e-6) + if not (cos >= 0.99 and max_rel < 1e-2): + return False, f"{tc['test_case_id']}: cos={cos:.5f} max_rel={max_rel:.4e}" + return True, None + + +def compile_smoke(): + """JIT-compile and launch the smallest workload case.""" + kern = _load_kernel(SOURCE_FILE) + tc = regime_cases()[0] + args = build_inputs(tc["M"], tc["N"], tc["K"], seed=42) + launch(kern, *args, tc["M"], tc["N"], tc["K"], META) + torch.cuda.synchronize() + + +def run_performance(benchmark): + edit = _load_kernel(SOURCE_FILE) + out = [] + for tc in regime_cases(): + M, N, K = tc["M"], tc["N"], tc["K"] + args = build_inputs(M, N, K, seed=42) + avg, metadata = benchmark( + lambda: launch(edit, *args, M, N, K, META), + warmup=10, + repetition=100, + ) + out.append({ + "test_case_id": tc["test_case_id"], + "execution_time_ms": avg, + **metadata, + "params": tc["params"], + }) + return out diff --git a/tasks/triton2triton/campaign20/gemm_a16_w16_kernel/scripts/task_runner.py b/tasks/triton2triton/campaign20/gemm_a16_w16_kernel/scripts/task_runner.py new file mode 100644 index 00000000..49d2378a --- /dev/null +++ b/tasks/triton2triton/campaign20/gemm_a16_w16_kernel/scripts/task_runner.py @@ -0,0 +1,116 @@ +#!/usr/bin/env python3 +"""Auto-generated task runner for triton__gemm_a16_w16_kernel (Triton). + +Inputs are generated each run from the shape/dtype signatures in +test_cases.json. The launcher symbol expected in source/triton__gemm_a16_w16_kernel.py is +``_gemm_a16_w16_kernel`` (for raw @triton.jit kernels) or ``None`` (a +wrapper that handles grid + meta resolution if present). +""" +import sys, os, json, argparse, glob, importlib.util +sys.path.insert(0, os.path.dirname(os.path.abspath(__file__))) + +# >>> AKA-GENERATED: shared CUDA-graph benchmark helpers - edit src/tools/perf/vllm_cuda_graph_block.py then run `make sync-perf-helpers` >>> +def _measure_cuda_event_fallback(*args, **kwargs): + raise RuntimeError( + "CUDA-graph benchmark helpers were not materialized. " + "Run this task through AgentKernelArena so setup_workspace() can inject " + "src/tools/perf/vllm_cuda_graph_block.py into the workspace." + ) + + +def _benchmark_cuda_graph_or_events(*args, **kwargs): + raise RuntimeError( + "CUDA-graph benchmark helpers were not materialized. " + "Run this task through AgentKernelArena so setup_workspace() can inject " + "src/tools/perf/vllm_cuda_graph_block.py into the workspace." + ) +# <<< AKA-GENERATED <<< + + +TASK_DIR = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) +os.chdir(TASK_DIR) +TASK_NAME = "triton2triton/campaign20/gemm_a16_w16_kernel" +SOURCE_FILE = os.path.join(TASK_DIR, "source", "triton__gemm_a16_w16_kernel.py") +KERNEL_NAME = "_gemm_a16_w16_kernel" +LAUNCHER_NAME = "None" +REF_SOURCE = "triton" +TEST_CASES = os.path.join(TASK_DIR, "test_cases.json") + + +def _load_module(): + spec = importlib.util.spec_from_file_location("triton_kernel", SOURCE_FILE) + mod = importlib.util.module_from_spec(spec) + spec.loader.exec_module(mod) + return mod + + +def _test_cases(): + if not os.path.isfile(TEST_CASES): + return [] + with open(TEST_CASES) as f: + return json.load(f) + + +def _resolve_callable(mod): + # Prefer a wrapping launcher if it was found alongside the @triton.jit kernel. + for n in (LAUNCHER_NAME, KERNEL_NAME): + if n and hasattr(mod, n) and n != "None": + return getattr(mod, n) + raise AttributeError(f"neither {LAUNCHER_NAME} nor {KERNEL_NAME} found in source") + + +def run_compile(): + try: + import ast + ast.parse(open(SOURCE_FILE).read()) + mod = _load_module() + _resolve_callable(mod) + import harness_run as hr + hr.compile_smoke() + return True, None + except Exception as e: + return False, str(e) + + +def run_correctness(): + """Golden-vs-editable correctness via the real launcher harness.""" + import harness_run as hr + hr.write_test_cases() + return hr.run_correctness() + + +def run_performance(): + """Canonical CUDA-graph timing with event fallback.""" + import harness_run as hr + hr.write_test_cases() + return hr.run_performance(_benchmark_cuda_graph_or_events) + + +def main(): + ap = argparse.ArgumentParser() + ap.add_argument("mode", choices=["compile", "correctness", "performance"]) + args = ap.parse_args() + build_dir = os.path.join(TASK_DIR, "build") + os.makedirs(build_dir, exist_ok=True) + if args.mode == "compile": + ok, err = run_compile() + json.dump({"status": "ok" if ok else "fail", "error": err}, open(os.path.join(build_dir, "compile_report.json"), "w")) + print(f"Compilation: {'PASS' if ok else 'FAIL'}") + if err: print("Error:", err) + sys.exit(0 if ok else 1) + if args.mode == "correctness": + ok, err = run_correctness() + json.dump({"status": "ok" if ok else "fail", "error": err}, open(os.path.join(build_dir, "correctness_report.json"), "w")) + print(f"Correctness: {'PASS' if ok else 'FAIL'}") + if err: print("Error:", err) + sys.exit(0 if ok else 1) + cases = run_performance() + json.dump({"test_cases": cases}, open(os.path.join(build_dir, "performance_report.json"), "w"), indent=2) + for c in cases: + print(f"Performance: {c['execution_time_ms']:.4f} ms ({c['test_case_id']})") + valid = bool(cases) and all(c.get("execution_time_ms", -1) > 0 for c in cases) + sys.exit(0 if valid else 1) + + +if __name__ == "__main__": + main() diff --git a/tasks/triton2triton/campaign20/gemm_a16_w16_kernel/source/triton__gemm_a16_w16_kernel.py b/tasks/triton2triton/campaign20/gemm_a16_w16_kernel/source/triton__gemm_a16_w16_kernel.py new file mode 100644 index 00000000..3ff222c0 --- /dev/null +++ b/tasks/triton2triton/campaign20/gemm_a16_w16_kernel/source/triton__gemm_a16_w16_kernel.py @@ -0,0 +1,141 @@ +import triton +import triton.language as tl + +@triton.heuristics( + { + "EVEN_K": lambda args: (args["K"] % (args["SPLITK_BLOCK_SIZE"]) == 0) + and (args["SPLITK_BLOCK_SIZE"] % args["BLOCK_SIZE_K"] == 0), + "GRID_MN": lambda args: triton.cdiv(args["M"], args["BLOCK_SIZE_M"]) + * triton.cdiv(args["N"], args["BLOCK_SIZE_N"]), + } +) +@triton.jit +def _gemm_a16_w16_kernel( + a_ptr, + b_ptr, + bias_ptr, + c_ptr, + M, + N, + K, + stride_am, + stride_ak, + stride_bk, + stride_bn, + stride_ck, + stride_cm, + stride_cn, + # Meta-parameters + BLOCK_SIZE_M: tl.constexpr, + BLOCK_SIZE_N: tl.constexpr, + BLOCK_SIZE_K: tl.constexpr, + GROUP_SIZE_M: tl.constexpr, + NUM_KSPLIT: tl.constexpr, + SPLITK_BLOCK_SIZE: tl.constexpr, + EVEN_K: tl.constexpr, + GRID_MN: tl.constexpr, + cache_modifier: tl.constexpr, + activation: tl.constexpr, + use_activation: tl.constexpr, + ADD_BIAS: tl.constexpr, + SKIP_REDUCE: tl.constexpr, +): + """Kernel for computing the matmul C = A x B. + A has shape (M, K), B has shape (K, N) and C has shape (M, N) + """ + + tl.assume(stride_am > 0) + tl.assume(stride_ak > 0) + tl.assume(stride_bk > 0) + tl.assume(stride_bn > 0) + tl.assume(stride_ck > 0) + tl.assume(stride_cm > 0) + tl.assume(stride_cn > 0) + + # ----------------------------------------------------------- + # Map program ids `pid` to the block of C it should compute. + # This is done in a grouped ordering to promote L2 data reuse. + pid_unified = tl.program_id(axis=0) + pid_unified = remap_xcd(pid_unified, GRID_MN * NUM_KSPLIT, NUM_XCDS=8) + pid_k = pid_unified % NUM_KSPLIT + pid = pid_unified // NUM_KSPLIT + num_pid_m = tl.cdiv(M, BLOCK_SIZE_M) + num_pid_n = tl.cdiv(N, BLOCK_SIZE_N) + + if NUM_KSPLIT == 1: + pid_m, pid_n = pid_grid(pid, num_pid_m, num_pid_n, GROUP_SIZE_M=GROUP_SIZE_M) + else: + pid_m = pid // num_pid_n + pid_n = pid % num_pid_n + + tl.assume(pid_m >= 0) + tl.assume(pid_n >= 0) + tl.assume(pid_k >= 0) + + split_k_start = pid_k * SPLITK_BLOCK_SIZE + if split_k_start < K: + # Create pointers for first block of A and B input matrices + offs_k = tl.arange(0, BLOCK_SIZE_K) + offs_k_split = split_k_start + offs_k + offs_am = (pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M)) % M + offs_bn = (pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N)) % N + + a_ptrs = a_ptr + ( + offs_am[:, None] * stride_am + offs_k_split[None, :] * stride_ak + ) + b_ptrs = b_ptr + ( + offs_k_split[:, None] * stride_bk + offs_bn[None, :] * stride_bn + ) + + acc_dtype = tl.float32 if c_ptr.type.element_ty != tl.int8 else tl.int32 + if ADD_BIAS: + if NUM_KSPLIT == 1 or (SKIP_REDUCE and pid_k == 0): + accumulator = tl.load(bias_ptr + offs_bn).to(dtype=acc_dtype) + accumulator = tl.broadcast_to( + accumulator[None, :], (BLOCK_SIZE_M, BLOCK_SIZE_N) + ) + else: + accumulator = tl.zeros((BLOCK_SIZE_M, BLOCK_SIZE_N), dtype=acc_dtype) + else: + accumulator = tl.zeros((BLOCK_SIZE_M, BLOCK_SIZE_N), dtype=acc_dtype) + + split_k_end = tl.minimum(split_k_start + SPLITK_BLOCK_SIZE, K) + k_span = split_k_end - split_k_start + num_k_iter = tl.cdiv(k_span, BLOCK_SIZE_K) + + for k in range(num_k_iter): + # Load the next block of A and B, generate a mask by checking the K dimension. + # If it is out of bounds, set it to 0. + if EVEN_K: + a = tl.load(a_ptrs) + b = tl.load(b_ptrs, cache_modifier=cache_modifier) + else: + a = tl.load( + a_ptrs, mask=offs_k[None, :] < k_span - k * BLOCK_SIZE_K, other=0.0 + ) + b = tl.load( + b_ptrs, + mask=offs_k[:, None] < k_span - k * BLOCK_SIZE_K, + other=0.0, + cache_modifier=cache_modifier, + ) + accumulator += tl.dot(a, b, input_precision="ieee") + # Advance the ptrs to the next K block. + a_ptrs += BLOCK_SIZE_K * stride_ak + b_ptrs += BLOCK_SIZE_K * stride_bk + + if use_activation and NUM_KSPLIT == 1: + accumulator = activation(accumulator) + + # Write back the block of the output matrix C with masks. + c = accumulator.to(c_ptr.type.element_ty) + offs_cm = pid_m.to(tl.int64) * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M) + offs_cn = pid_n.to(tl.int64) * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N) + c_ptrs = ( + c_ptr + + stride_cm * offs_cm[:, None] + + stride_cn * offs_cn[None, :] + + pid_k * stride_ck + ) + c_mask = (offs_cm[:, None] < M) & (offs_cn[None, :] < N) + tl.store(c_ptrs, c, mask=c_mask) diff --git a/tasks/triton2triton/campaign20/gemm_a16_w16_kernel/source_golden/triton__gemm_a16_w16_kernel.py b/tasks/triton2triton/campaign20/gemm_a16_w16_kernel/source_golden/triton__gemm_a16_w16_kernel.py new file mode 100644 index 00000000..3ff222c0 --- /dev/null +++ b/tasks/triton2triton/campaign20/gemm_a16_w16_kernel/source_golden/triton__gemm_a16_w16_kernel.py @@ -0,0 +1,141 @@ +import triton +import triton.language as tl + +@triton.heuristics( + { + "EVEN_K": lambda args: (args["K"] % (args["SPLITK_BLOCK_SIZE"]) == 0) + and (args["SPLITK_BLOCK_SIZE"] % args["BLOCK_SIZE_K"] == 0), + "GRID_MN": lambda args: triton.cdiv(args["M"], args["BLOCK_SIZE_M"]) + * triton.cdiv(args["N"], args["BLOCK_SIZE_N"]), + } +) +@triton.jit +def _gemm_a16_w16_kernel( + a_ptr, + b_ptr, + bias_ptr, + c_ptr, + M, + N, + K, + stride_am, + stride_ak, + stride_bk, + stride_bn, + stride_ck, + stride_cm, + stride_cn, + # Meta-parameters + BLOCK_SIZE_M: tl.constexpr, + BLOCK_SIZE_N: tl.constexpr, + BLOCK_SIZE_K: tl.constexpr, + GROUP_SIZE_M: tl.constexpr, + NUM_KSPLIT: tl.constexpr, + SPLITK_BLOCK_SIZE: tl.constexpr, + EVEN_K: tl.constexpr, + GRID_MN: tl.constexpr, + cache_modifier: tl.constexpr, + activation: tl.constexpr, + use_activation: tl.constexpr, + ADD_BIAS: tl.constexpr, + SKIP_REDUCE: tl.constexpr, +): + """Kernel for computing the matmul C = A x B. + A has shape (M, K), B has shape (K, N) and C has shape (M, N) + """ + + tl.assume(stride_am > 0) + tl.assume(stride_ak > 0) + tl.assume(stride_bk > 0) + tl.assume(stride_bn > 0) + tl.assume(stride_ck > 0) + tl.assume(stride_cm > 0) + tl.assume(stride_cn > 0) + + # ----------------------------------------------------------- + # Map program ids `pid` to the block of C it should compute. + # This is done in a grouped ordering to promote L2 data reuse. + pid_unified = tl.program_id(axis=0) + pid_unified = remap_xcd(pid_unified, GRID_MN * NUM_KSPLIT, NUM_XCDS=8) + pid_k = pid_unified % NUM_KSPLIT + pid = pid_unified // NUM_KSPLIT + num_pid_m = tl.cdiv(M, BLOCK_SIZE_M) + num_pid_n = tl.cdiv(N, BLOCK_SIZE_N) + + if NUM_KSPLIT == 1: + pid_m, pid_n = pid_grid(pid, num_pid_m, num_pid_n, GROUP_SIZE_M=GROUP_SIZE_M) + else: + pid_m = pid // num_pid_n + pid_n = pid % num_pid_n + + tl.assume(pid_m >= 0) + tl.assume(pid_n >= 0) + tl.assume(pid_k >= 0) + + split_k_start = pid_k * SPLITK_BLOCK_SIZE + if split_k_start < K: + # Create pointers for first block of A and B input matrices + offs_k = tl.arange(0, BLOCK_SIZE_K) + offs_k_split = split_k_start + offs_k + offs_am = (pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M)) % M + offs_bn = (pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N)) % N + + a_ptrs = a_ptr + ( + offs_am[:, None] * stride_am + offs_k_split[None, :] * stride_ak + ) + b_ptrs = b_ptr + ( + offs_k_split[:, None] * stride_bk + offs_bn[None, :] * stride_bn + ) + + acc_dtype = tl.float32 if c_ptr.type.element_ty != tl.int8 else tl.int32 + if ADD_BIAS: + if NUM_KSPLIT == 1 or (SKIP_REDUCE and pid_k == 0): + accumulator = tl.load(bias_ptr + offs_bn).to(dtype=acc_dtype) + accumulator = tl.broadcast_to( + accumulator[None, :], (BLOCK_SIZE_M, BLOCK_SIZE_N) + ) + else: + accumulator = tl.zeros((BLOCK_SIZE_M, BLOCK_SIZE_N), dtype=acc_dtype) + else: + accumulator = tl.zeros((BLOCK_SIZE_M, BLOCK_SIZE_N), dtype=acc_dtype) + + split_k_end = tl.minimum(split_k_start + SPLITK_BLOCK_SIZE, K) + k_span = split_k_end - split_k_start + num_k_iter = tl.cdiv(k_span, BLOCK_SIZE_K) + + for k in range(num_k_iter): + # Load the next block of A and B, generate a mask by checking the K dimension. + # If it is out of bounds, set it to 0. + if EVEN_K: + a = tl.load(a_ptrs) + b = tl.load(b_ptrs, cache_modifier=cache_modifier) + else: + a = tl.load( + a_ptrs, mask=offs_k[None, :] < k_span - k * BLOCK_SIZE_K, other=0.0 + ) + b = tl.load( + b_ptrs, + mask=offs_k[:, None] < k_span - k * BLOCK_SIZE_K, + other=0.0, + cache_modifier=cache_modifier, + ) + accumulator += tl.dot(a, b, input_precision="ieee") + # Advance the ptrs to the next K block. + a_ptrs += BLOCK_SIZE_K * stride_ak + b_ptrs += BLOCK_SIZE_K * stride_bk + + if use_activation and NUM_KSPLIT == 1: + accumulator = activation(accumulator) + + # Write back the block of the output matrix C with masks. + c = accumulator.to(c_ptr.type.element_ty) + offs_cm = pid_m.to(tl.int64) * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M) + offs_cn = pid_n.to(tl.int64) * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N) + c_ptrs = ( + c_ptr + + stride_cm * offs_cm[:, None] + + stride_cn * offs_cn[None, :] + + pid_k * stride_ck + ) + c_mask = (offs_cm[:, None] < M) & (offs_cn[None, :] < N) + tl.store(c_ptrs, c, mask=c_mask) diff --git a/tasks/triton2triton/campaign20/gemm_a16_w16_kernel/test_cases.json b/tasks/triton2triton/campaign20/gemm_a16_w16_kernel/test_cases.json new file mode 100644 index 00000000..76c945e7 --- /dev/null +++ b/tasks/triton2triton/campaign20/gemm_a16_w16_kernel/test_cases.json @@ -0,0 +1,44 @@ +[ + { + "test_case_id": "c2", + "B": 2, + "M": 2048, + "N": 5120, + "K": 2880, + "params": { + "B": 2, + "seqlen": 1024, + "M": 2048, + "N": 5120, + "K": 2880 + } + }, + { + "test_case_id": "c32", + "B": 32, + "M": 32768, + "N": 5120, + "K": 2880, + "params": { + "B": 32, + "seqlen": 1024, + "M": 32768, + "N": 5120, + "K": 2880 + } + }, + { + "test_case_id": "c64", + "B": 64, + "M": 65536, + "N": 5120, + "K": 2880, + "params": { + "B": 64, + "seqlen": 1024, + "M": 65536, + "N": 5120, + "K": 2880 + } + } +] \ No newline at end of file diff --git a/tasks/triton2triton/campaign20/gemm_a8w8_blockscale_kernel/Makefile b/tasks/triton2triton/campaign20/gemm_a8w8_blockscale_kernel/Makefile new file mode 100644 index 00000000..0f8b4f36 --- /dev/null +++ b/tasks/triton2triton/campaign20/gemm_a8w8_blockscale_kernel/Makefile @@ -0,0 +1,6 @@ +# Auto-generated. +all: + python3 scripts/task_runner.py compile + +clean: + rm -rf build diff --git a/tasks/triton2triton/campaign20/gemm_a8w8_blockscale_kernel/config.yaml b/tasks/triton2triton/campaign20/gemm_a8w8_blockscale_kernel/config.yaml new file mode 100644 index 00000000..112829f8 --- /dev/null +++ b/tasks/triton2triton/campaign20/gemm_a8w8_blockscale_kernel/config.yaml @@ -0,0 +1,29 @@ +source_file_path: +- source/triton__gemm_a8w8_blockscale_kernel.py +target_kernel_functions: +- _gemm_a8w8_blockscale_kernel +compile_command: +- python3 scripts/task_runner.py compile +correctness_command: +- python3 scripts/task_runner.py correctness +performance_command: +- python3 scripts/task_runner.py performance +task_type: triton2triton +platform_support: + required_arch: gfx950 + status: active + skip_reason: null +task_result_template: null +prompt: + source_code: null + instructions: Optimize the Triton kernel `_gemm_a8w8_blockscale_kernel` (originally + aiter.ops.triton._triton_kernels.gemm.basic.gemm_a8w8_blockscale). Captured via + inspect.getsource. Tests generate inputs from test_cases.json on the fly. + cheatsheet: null +extracted_from: + engine: sglang + kernel_id: triton:aiter.ops.triton._triton_kernels.gemm.basic.gemm_a8w8_blockscale._gemm_a8w8_blockscale_kernel + triton_fn: _gemm_a8w8_blockscale_kernel + origin_module: aiter.ops.triton._triton_kernels.gemm.basic.gemm_a8w8_blockscale + source: triton_inspect + test_cases: 3 diff --git a/tasks/triton2triton/campaign20/gemm_a8w8_blockscale_kernel/scripts/harness_run.py b/tasks/triton2triton/campaign20/gemm_a8w8_blockscale_kernel/scripts/harness_run.py new file mode 100644 index 00000000..47e686b1 --- /dev/null +++ b/tasks/triton2triton/campaign20/gemm_a8w8_blockscale_kernel/scripts/harness_run.py @@ -0,0 +1,219 @@ +#!/usr/bin/env python3 +"""Real launcher + benchmark for the raw @triton.jit kernel +``_gemm_a8w8_blockscale_kernel`` (a8w8 block-scale GEMM, aiter). + +Workload regime: seqlen = 1024, concurrency B in {2, 32, 64}. +This is a token-parallel GEMM (per-token activation rows). We map: + M (num_tokens / rows of A) = B * 1024 (prefill token count) +Model dims held from the captured base case: + K = 6144 (hidden), N = 2624 (intermediate-ish out feature) + GROUP_K = GROUP_N = BLOCK_SIZE_K = 128, group_size = 128 +Scales: + a_scale : (M, ceil(K/128) = 48) fp32 + b_scale : (ceil(K/128) = 48, ceil(N/128) = 21) fp32 (col-major) +A : (M, K) fp8_e4m3fnuz row-major +B : (K, N) fp8_e4m3fnuz col-major (stride [1, K]) +C : (M, N) bf16 row-major (output, written in-place) + +The kernel references module globals ``remap_xcd`` and ``pid_grid``. Small +self-contained implementations are provided below so the task does not depend +on an external AITER checkout. +""" +import importlib.util +import os +import sys + +import torch +import triton +import triton.language as tl + +HERE = os.path.dirname(os.path.abspath(__file__)) +TASK_DIR = os.path.dirname(HERE) +KERNEL_NAME = "_gemm_a8w8_blockscale_kernel" + +@triton.jit +def remap_xcd(pid, GRID_MN, NUM_XCDS: tl.constexpr = 8): + """Distribute consecutive program IDs across XCDs.""" + pids_per_xcd = (GRID_MN + NUM_XCDS - 1) // NUM_XCDS + tall_xcds = GRID_MN % NUM_XCDS + tall_xcds = NUM_XCDS if tall_xcds == 0 else tall_xcds + xcd = pid % NUM_XCDS + local_pid = pid // NUM_XCDS + if xcd < tall_xcds: + pid = xcd * pids_per_xcd + local_pid + else: + pid = ( + tall_xcds * pids_per_xcd + + (xcd - tall_xcds) * (pids_per_xcd - 1) + + local_pid + ) + return pid + + +@triton.jit +def pid_grid(pid, num_pid_m, num_pid_n, GROUP_SIZE_M: tl.constexpr = 1): + """Map a linear program ID to grouped GEMM M/N coordinates.""" + if GROUP_SIZE_M == 1: + pid_m = pid // num_pid_n + pid_n = pid % num_pid_n + else: + num_pid_in_group = GROUP_SIZE_M * num_pid_n + group_id = pid // num_pid_in_group + first_pid_m = group_id * GROUP_SIZE_M + group_size_m = min(num_pid_m - first_pid_m, GROUP_SIZE_M) + tl.assume(group_size_m >= 0) + pid_m = first_pid_m + (pid % group_size_m) + pid_n = (pid % num_pid_in_group) // group_size_m + return pid_m, pid_n + +# Fixed model dims from captured base case +K_DIM = 6144 +N_DIM = 2624 +GROUP = 128 # GROUP_K == GROUP_N == BLOCK_SIZE_K +BLOCK_M = 128 +BLOCK_N = 128 +BLOCK_K = 128 +GROUP_SIZE_M = 1 +NUM_KSPLIT = 1 +NUM_WARPS = 4 +NUM_STAGES = 2 +CACHE_MOD = ".cg" + +SEQLEN = 1024 +CONCURRENCY = [2, 32, 64] + + +def _load_kernel(source_file): + spec = importlib.util.spec_from_file_location("kmod", source_file) + mod = importlib.util.module_from_spec(spec) + # Inject the helper jit functions the kernel references as globals. + mod.pid_grid = pid_grid + mod.remap_xcd = remap_xcd + spec.loader.exec_module(mod) + # exec_module rebinds the module dict; re-inject after exec to be safe. + mod.pid_grid = pid_grid + mod.remap_xcd = remap_xcd + kern = getattr(mod, KERNEL_NAME) + # the JIT fn captures globals from its own __globals__; patch those too. + kern.fn.__globals__.setdefault("pid_grid", pid_grid) + kern.fn.__globals__.setdefault("remap_xcd", remap_xcd) + kern.fn.__globals__["pid_grid"] = pid_grid + kern.fn.__globals__["remap_xcd"] = remap_xcd + return kern + + +def _build_inputs(B, seed=42): + """Deterministic seeded inputs for concurrency B. Returns (args, meta, c).""" + g = torch.Generator(device="cuda").manual_seed(seed) + M = B * SEQLEN + K = K_DIM + N = N_DIM + scale_k = (K + GROUP - 1) // GROUP # 48 + scale_n = (N + GROUP - 1) // GROUP # 21 + fp8 = getattr(torch, "float8_e4m3fnuz") + + a_f = (torch.randn(M, K, generator=g, device="cuda", dtype=torch.float32) * 0.2) + b_f = (torch.randn(K, N, generator=g, device="cuda", dtype=torch.float32) * 0.2) + a = a_f.to(fp8) + # B is col-major (stride [1, K]) as in capture + b = b_f.to(fp8).t().contiguous().t() + c = torch.empty(M, N, device="cuda", dtype=torch.bfloat16) + + a_scale = (torch.rand(M, scale_k, generator=g, device="cuda", dtype=torch.float32) * 0.01 + 0.005) + # b_scale col-major (stride [1, scale_k]) as in capture + b_scale = (torch.rand(scale_k, scale_n, generator=g, device="cuda", dtype=torch.float32) * 0.01 + 0.005) + b_scale = b_scale.t().contiguous().t() + + args = ( + a, b, c, a_scale, b_scale, + M, N, K, + a.stride(0), a.stride(1), + b.stride(0), b.stride(1), + 0, # stride_ck (NUM_KSPLIT==1) + c.stride(0), c.stride(1), + a_scale.stride(0), a_scale.stride(1), + b_scale.stride(0), b_scale.stride(1), + ) + meta = dict( + GROUP_K=GROUP, GROUP_N=GROUP, + BLOCK_SIZE_M=BLOCK_M, BLOCK_SIZE_N=BLOCK_N, BLOCK_SIZE_K=BLOCK_K, + GROUP_SIZE_M=GROUP_SIZE_M, + NUM_KSPLIT=NUM_KSPLIT, + SPLITK_BLOCK_SIZE=(K + NUM_KSPLIT - 1) // NUM_KSPLIT, + cache_modifier=CACHE_MOD, + num_warps=NUM_WARPS, num_stages=NUM_STAGES, + ) + return args, meta, c, M, N + + +def _grid(M, N): + return (NUM_KSPLIT * ((M + BLOCK_M - 1) // BLOCK_M) * ((N + BLOCK_N - 1) // BLOCK_N),) + + +def _launch(kern, B, seed=42): + args, meta, c, M, N = _build_inputs(B, seed=seed) + grid = _grid(M, N) + c.zero_() + kern[grid](*args, **meta) + torch.cuda.synchronize() + return c + + +def cases(): + return [(f"c{b}", b) for b in CONCURRENCY] + + +def run_correctness(): + src_edit = os.path.join(TASK_DIR, "source", "triton__gemm_a8w8_blockscale_kernel.py") + src_gold = os.path.join(TASK_DIR, "source_golden", "triton__gemm_a8w8_blockscale_kernel.py") + kern_e = _load_kernel(src_edit) + kern_g = _load_kernel(src_gold) + for cid, B in cases(): + out_e = _launch(kern_e, B, seed=42).float() + out_g = _launch(kern_g, B, seed=42).float() + cos = torch.nn.functional.cosine_similarity( + out_e.flatten().unsqueeze(0), out_g.flatten().unsqueeze(0) + ).item() + denom = out_g.abs().max().clamp_min(1e-6) + max_rel = (out_e - out_g).abs().max().item() / denom.item() + if not (cos >= 0.99 and max_rel <= 1e-2): + return False, f"{cid}: cos={cos:.5f} max_rel={max_rel:.5f}" + return True, None + + +def compile_smoke(): + """JIT-compile and launch the smallest workload case.""" + src = os.path.join( + TASK_DIR, "source", "triton__gemm_a8w8_blockscale_kernel.py" + ) + _launch(_load_kernel(src), CONCURRENCY[0], seed=42) + + +def run_performance(benchmark): + src_edit = os.path.join(TASK_DIR, "source", "triton__gemm_a8w8_blockscale_kernel.py") + kern = _load_kernel(src_edit) + results = [] + for cid, B in cases(): + args, meta, c, M, N = _build_inputs(B, seed=42) + grid = _grid(M, N) + avg, metadata = benchmark( + lambda: kern[grid](*args, **meta), + warmup=10, + repetition=100, + ) + results.append({ + "test_case_id": cid, + "execution_time_ms": avg, + **metadata, + "params": {"B": B, "M": M, "N": N, "K": K_DIM, "seqlen": SEQLEN}, + }) + return results + + +if __name__ == "__main__": + mode = sys.argv[1] if len(sys.argv) > 1 else "correctness" + if mode == "correctness": + ok, err = run_correctness() + print("Correctness:", "PASS" if ok else "FAIL", err or "") + else: + raise SystemExit("run performance through scripts/task_runner.py") diff --git a/tasks/triton2triton/campaign20/gemm_a8w8_blockscale_kernel/scripts/task_runner.py b/tasks/triton2triton/campaign20/gemm_a8w8_blockscale_kernel/scripts/task_runner.py new file mode 100644 index 00000000..566a6e64 --- /dev/null +++ b/tasks/triton2triton/campaign20/gemm_a8w8_blockscale_kernel/scripts/task_runner.py @@ -0,0 +1,126 @@ +#!/usr/bin/env python3 +"""Auto-generated task runner for triton__gemm_a8w8_blockscale_kernel (Triton). + +Inputs are generated each run from the shape/dtype signatures in +test_cases.json. The launcher symbol expected in source/triton__gemm_a8w8_blockscale_kernel.py is +``_gemm_a8w8_blockscale_kernel`` (for raw @triton.jit kernels) or ``None`` (a +wrapper that handles grid + meta resolution if present). +""" +import sys, os, json, argparse, glob, importlib.util +sys.path.insert(0, os.path.dirname(os.path.abspath(__file__))) + +# >>> AKA-GENERATED: shared CUDA-graph benchmark helpers - edit src/tools/perf/vllm_cuda_graph_block.py then run `make sync-perf-helpers` >>> +def _measure_cuda_event_fallback(*args, **kwargs): + raise RuntimeError( + "CUDA-graph benchmark helpers were not materialized. " + "Run this task through AgentKernelArena so setup_workspace() can inject " + "src/tools/perf/vllm_cuda_graph_block.py into the workspace." + ) + + +def _benchmark_cuda_graph_or_events(*args, **kwargs): + raise RuntimeError( + "CUDA-graph benchmark helpers were not materialized. " + "Run this task through AgentKernelArena so setup_workspace() can inject " + "src/tools/perf/vllm_cuda_graph_block.py into the workspace." + ) +# <<< AKA-GENERATED <<< + + +TASK_DIR = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) +os.chdir(TASK_DIR) +TASK_NAME = "triton2triton/campaign20/gemm_a8w8_blockscale_kernel" +SOURCE_FILE = os.path.join(TASK_DIR, "source", "triton__gemm_a8w8_blockscale_kernel.py") +KERNEL_NAME = "_gemm_a8w8_blockscale_kernel" +LAUNCHER_NAME = "None" +REF_SOURCE = "triton" +TEST_CASES = os.path.join(TASK_DIR, "test_cases.json") + + +def _load_module(): + spec = importlib.util.spec_from_file_location("triton_kernel", SOURCE_FILE) + mod = importlib.util.module_from_spec(spec) + spec.loader.exec_module(mod) + return mod + + +def _test_cases(): + if not os.path.isfile(TEST_CASES): + return [] + with open(TEST_CASES) as f: + return json.load(f) + + +def _resolve_callable(mod): + # Prefer a wrapping launcher if it was found alongside the @triton.jit kernel. + for n in (LAUNCHER_NAME, KERNEL_NAME): + if n and hasattr(mod, n) and n != "None": + return getattr(mod, n) + raise AttributeError(f"neither {LAUNCHER_NAME} nor {KERNEL_NAME} found in source") + + +def run_compile(): + try: + import ast + ast.parse(open(SOURCE_FILE).read()) + mod = _load_module() + _resolve_callable(mod) + import harness_run as hr + hr.compile_smoke() + return True, None + except Exception as e: + return False, str(e) + + +def run_correctness(): + """Real golden-vs-editable correctness via the custom launcher harness. + The raw @triton.jit kernel has no wrapper, so we build a real launcher in + scripts/harness_run.py (grid formula + meta-args + seeded inputs) and run + both the editable source and the frozen golden copy, comparing outputs.""" + import harness_run as hr + try: + return hr.run_correctness() + except Exception as e: + import traceback + return False, traceback.format_exc() + + +def run_performance(): + """Canonical CUDA-graph timing with event fallback.""" + import harness_run as hr + try: + return hr.run_performance(_benchmark_cuda_graph_or_events) + except Exception as e: + import traceback + return [{"test_case_id": "error", "execution_time_ms": -1.0, + "params": {"error": traceback.format_exc()[-200:]}}] + + +def main(): + ap = argparse.ArgumentParser() + ap.add_argument("mode", choices=["compile", "correctness", "performance"]) + args = ap.parse_args() + build_dir = os.path.join(TASK_DIR, "build") + os.makedirs(build_dir, exist_ok=True) + if args.mode == "compile": + ok, err = run_compile() + json.dump({"status": "ok" if ok else "fail", "error": err}, open(os.path.join(build_dir, "compile_report.json"), "w")) + print(f"Compilation: {'PASS' if ok else 'FAIL'}") + if err: print("Error:", err) + sys.exit(0 if ok else 1) + if args.mode == "correctness": + ok, err = run_correctness() + json.dump({"status": "ok" if ok else "fail", "error": err}, open(os.path.join(build_dir, "correctness_report.json"), "w")) + print(f"Correctness: {'PASS' if ok else 'FAIL'}") + if err: print("Error:", err) + sys.exit(0 if ok else 1) + cases = run_performance() + json.dump({"test_cases": cases}, open(os.path.join(build_dir, "performance_report.json"), "w"), indent=2) + for c in cases: + print(f"Performance: {c['execution_time_ms']:.4f} ms ({c['test_case_id']})") + valid = bool(cases) and all(c.get("execution_time_ms", -1) > 0 for c in cases) + sys.exit(0 if valid else 1) + + +if __name__ == "__main__": + main() diff --git a/tasks/triton2triton/campaign20/gemm_a8w8_blockscale_kernel/source/triton__gemm_a8w8_blockscale_kernel.py b/tasks/triton2triton/campaign20/gemm_a8w8_blockscale_kernel/source/triton__gemm_a8w8_blockscale_kernel.py new file mode 100644 index 00000000..b9ec4243 --- /dev/null +++ b/tasks/triton2triton/campaign20/gemm_a8w8_blockscale_kernel/source/triton__gemm_a8w8_blockscale_kernel.py @@ -0,0 +1,180 @@ +import triton +import triton.language as tl + +@triton.heuristics( + { + "EVEN_K": lambda args: args["K"] % args["BLOCK_SIZE_K"] == 0, + "GRID_MN": lambda args: triton.cdiv(args["M"], args["BLOCK_SIZE_M"]) + * triton.cdiv(args["N"], args["BLOCK_SIZE_N"]), + } +) +@triton.jit +def _gemm_a8w8_blockscale_kernel( + # Pointers to matrices + a_ptr, + b_ptr, + c_ptr, + a_scale_ptr, + b_scale_ptr, + # Matrix dimensions + M, + N, + K, + # The stride variables represent how much to increase the ptr by when + # moving by 1 element in a particular dimension. E.g. `stride_am` is + # how much to increase `a_ptr` by to get the element one row down + # (A has M rows). + stride_am, + stride_ak, + stride_bk, + stride_bn, + stride_ck, + stride_cm, + stride_cn, + stride_ascale_m, + stride_ascale_k, + stride_bscale_k, + stride_bscale_n, + # Meta-parameters + GROUP_K: tl.constexpr, + GROUP_N: tl.constexpr, + BLOCK_SIZE_M: tl.constexpr, + BLOCK_SIZE_N: tl.constexpr, + BLOCK_SIZE_K: tl.constexpr, + GROUP_SIZE_M: tl.constexpr, + NUM_KSPLIT: tl.constexpr, + SPLITK_BLOCK_SIZE: tl.constexpr, + EVEN_K: tl.constexpr, + GRID_MN: tl.constexpr, + cache_modifier: tl.constexpr, +): + """ + Note: this is Triton jited function and not meant to be called directly. Call gemm_a8w8_blockscale function + below + + Computes the 8 bit matmul C = A x B using the block-scale quantization approach. + + Key parameters: + - A: Matrix A with shape (M, K). + - B: Matrix B with shape (K, N). + - C: Matrix C with shape (M, N). + - A_scale: Scale tensor for A with shape (M, *scale_k). + - B_scale: Scale tensor for B with shape (*scale_k, **scale_n). + + *scale_k = (K + GROUP_K - 1) // GROUP_K + **scale_n = (N + GROUP_N - 1) // GROUP_N + + For this kernel implementation, GROUP_K must equal BLOCK_K. + """ + + tl.assume(stride_am > 0) + tl.assume(stride_ak > 0) + tl.assume(stride_bk > 0) + tl.assume(stride_bn > 0) + tl.assume(stride_ck > 0) + tl.assume(stride_cm > 0) + tl.assume(stride_cn > 0) + tl.assume(stride_ascale_m > 0) + tl.assume(stride_ascale_k > 0) + tl.assume(stride_bscale_k > 0) + tl.assume(stride_bscale_n > 0) + + # ----------------------------------------------------------- + # Map program ids `pid` to the block of C it should compute. + # This is done in a grouped ordering to promote L2 data reuse. + pid_unified = tl.program_id(axis=0) + pid_k = pid_unified % NUM_KSPLIT + pid = pid_unified // NUM_KSPLIT + num_pid_m = tl.cdiv(M, BLOCK_SIZE_M) + num_pid_n = tl.cdiv(N, BLOCK_SIZE_N) + + if NUM_KSPLIT == 1: + remap_xcd(pid, GRID_MN) + + pid_m, pid_n = pid_grid(pid, num_pid_m, num_pid_n, GROUP_SIZE_M=GROUP_SIZE_M) + else: + pid_m = pid // num_pid_n + pid_n = pid % num_pid_n + + tl.assume(pid_m >= 0) + tl.assume(pid_n >= 0) + tl.assume(pid_k >= 0) + + if (pid_k * SPLITK_BLOCK_SIZE) < K: + + # SPLITK_BLOCK_SIZE = tl.cdiv(K, NUM_KSPLIT) + num_k_iter = tl.cdiv(SPLITK_BLOCK_SIZE, BLOCK_SIZE_K) + # ^ Number of K blocks within our split-K partition + + # Create pointers for first block of A and B input matrices + offs_k = tl.arange(0, BLOCK_SIZE_K) + offs_k_split = pid_k * SPLITK_BLOCK_SIZE + offs_k + offs_am = (pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M)) % M + offs_bn = (pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N)) % N + a_ptrs = a_ptr + ( + offs_am[:, None] * stride_am + offs_k_split[None, :] * stride_ak + ) + b_ptrs = b_ptr + ( + offs_k_split[:, None] * stride_bk + offs_bn[None, :] * stride_bn + ) + + # Create pointers for the scales + offs_k_scale = (pid_k * SPLITK_BLOCK_SIZE) // GROUP_K + a_scale_ptrs = ( + a_scale_ptr + offs_am * stride_ascale_m + offs_k_scale * stride_ascale_k + ) + offs_b_scale_n = offs_bn // GROUP_N + b_scale_ptrs = ( + b_scale_ptr + + offs_k_scale * stride_bscale_k + + offs_b_scale_n * stride_bscale_n + ) + offs_ks_step = BLOCK_SIZE_K // GROUP_K + + acc_dtype = tl.float32 if c_ptr.type.element_ty != tl.int8 else tl.int32 + accumulator = tl.zeros((BLOCK_SIZE_M, BLOCK_SIZE_N), dtype=acc_dtype) + + for k in range(pid_k * num_k_iter, (pid_k + 1) * num_k_iter): + # Load the next block of A and B, generate a mask by checking the K dimension. + # If it is out of bounds, set it to 0. + if EVEN_K: + a = tl.load(a_ptrs) + b = tl.load(b_ptrs, cache_modifier=cache_modifier) + else: + a = tl.load( + a_ptrs, mask=offs_k[None, :] < K - k * BLOCK_SIZE_K, other=0.0 + ) + b = tl.load( + b_ptrs, mask=offs_k[:, None] < K - k * BLOCK_SIZE_K, other=0.0 + ) + + a_scale = tl.load(a_scale_ptrs) + b_scale = tl.load(b_scale_ptrs) + + # Perform dot operation and apply scale + accumulator += ( + tl.dot(a, b, input_precision="ieee") + * a_scale[:, None] + * b_scale[None, :] + ) + + # Advance the ptrs to the next K block. + a_ptrs += BLOCK_SIZE_K * stride_ak + b_ptrs += BLOCK_SIZE_K * stride_bk + + a_scale_ptrs += offs_ks_step * stride_ascale_k + b_scale_ptrs += offs_ks_step * stride_bscale_k + + c = accumulator.to(c_ptr.type.element_ty) + + # Write back the block of the output matrix C with masks. + offs_cm = pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M).to(tl.int64) + offs_cn = pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N).to(tl.int64) + c_ptrs = ( + c_ptr + + stride_cm * offs_cm[:, None] + + stride_cn * offs_cn[None, :] + + pid_k * stride_ck + ) + c_mask = (offs_cm[:, None] < M) & (offs_cn[None, :] < N) + tl.store(c_ptrs, c, mask=c_mask) diff --git a/tasks/triton2triton/campaign20/gemm_a8w8_blockscale_kernel/source_golden/triton__gemm_a8w8_blockscale_kernel.py b/tasks/triton2triton/campaign20/gemm_a8w8_blockscale_kernel/source_golden/triton__gemm_a8w8_blockscale_kernel.py new file mode 100644 index 00000000..b9ec4243 --- /dev/null +++ b/tasks/triton2triton/campaign20/gemm_a8w8_blockscale_kernel/source_golden/triton__gemm_a8w8_blockscale_kernel.py @@ -0,0 +1,180 @@ +import triton +import triton.language as tl + +@triton.heuristics( + { + "EVEN_K": lambda args: args["K"] % args["BLOCK_SIZE_K"] == 0, + "GRID_MN": lambda args: triton.cdiv(args["M"], args["BLOCK_SIZE_M"]) + * triton.cdiv(args["N"], args["BLOCK_SIZE_N"]), + } +) +@triton.jit +def _gemm_a8w8_blockscale_kernel( + # Pointers to matrices + a_ptr, + b_ptr, + c_ptr, + a_scale_ptr, + b_scale_ptr, + # Matrix dimensions + M, + N, + K, + # The stride variables represent how much to increase the ptr by when + # moving by 1 element in a particular dimension. E.g. `stride_am` is + # how much to increase `a_ptr` by to get the element one row down + # (A has M rows). + stride_am, + stride_ak, + stride_bk, + stride_bn, + stride_ck, + stride_cm, + stride_cn, + stride_ascale_m, + stride_ascale_k, + stride_bscale_k, + stride_bscale_n, + # Meta-parameters + GROUP_K: tl.constexpr, + GROUP_N: tl.constexpr, + BLOCK_SIZE_M: tl.constexpr, + BLOCK_SIZE_N: tl.constexpr, + BLOCK_SIZE_K: tl.constexpr, + GROUP_SIZE_M: tl.constexpr, + NUM_KSPLIT: tl.constexpr, + SPLITK_BLOCK_SIZE: tl.constexpr, + EVEN_K: tl.constexpr, + GRID_MN: tl.constexpr, + cache_modifier: tl.constexpr, +): + """ + Note: this is Triton jited function and not meant to be called directly. Call gemm_a8w8_blockscale function + below + + Computes the 8 bit matmul C = A x B using the block-scale quantization approach. + + Key parameters: + - A: Matrix A with shape (M, K). + - B: Matrix B with shape (K, N). + - C: Matrix C with shape (M, N). + - A_scale: Scale tensor for A with shape (M, *scale_k). + - B_scale: Scale tensor for B with shape (*scale_k, **scale_n). + + *scale_k = (K + GROUP_K - 1) // GROUP_K + **scale_n = (N + GROUP_N - 1) // GROUP_N + + For this kernel implementation, GROUP_K must equal BLOCK_K. + """ + + tl.assume(stride_am > 0) + tl.assume(stride_ak > 0) + tl.assume(stride_bk > 0) + tl.assume(stride_bn > 0) + tl.assume(stride_ck > 0) + tl.assume(stride_cm > 0) + tl.assume(stride_cn > 0) + tl.assume(stride_ascale_m > 0) + tl.assume(stride_ascale_k > 0) + tl.assume(stride_bscale_k > 0) + tl.assume(stride_bscale_n > 0) + + # ----------------------------------------------------------- + # Map program ids `pid` to the block of C it should compute. + # This is done in a grouped ordering to promote L2 data reuse. + pid_unified = tl.program_id(axis=0) + pid_k = pid_unified % NUM_KSPLIT + pid = pid_unified // NUM_KSPLIT + num_pid_m = tl.cdiv(M, BLOCK_SIZE_M) + num_pid_n = tl.cdiv(N, BLOCK_SIZE_N) + + if NUM_KSPLIT == 1: + remap_xcd(pid, GRID_MN) + + pid_m, pid_n = pid_grid(pid, num_pid_m, num_pid_n, GROUP_SIZE_M=GROUP_SIZE_M) + else: + pid_m = pid // num_pid_n + pid_n = pid % num_pid_n + + tl.assume(pid_m >= 0) + tl.assume(pid_n >= 0) + tl.assume(pid_k >= 0) + + if (pid_k * SPLITK_BLOCK_SIZE) < K: + + # SPLITK_BLOCK_SIZE = tl.cdiv(K, NUM_KSPLIT) + num_k_iter = tl.cdiv(SPLITK_BLOCK_SIZE, BLOCK_SIZE_K) + # ^ Number of K blocks within our split-K partition + + # Create pointers for first block of A and B input matrices + offs_k = tl.arange(0, BLOCK_SIZE_K) + offs_k_split = pid_k * SPLITK_BLOCK_SIZE + offs_k + offs_am = (pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M)) % M + offs_bn = (pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N)) % N + a_ptrs = a_ptr + ( + offs_am[:, None] * stride_am + offs_k_split[None, :] * stride_ak + ) + b_ptrs = b_ptr + ( + offs_k_split[:, None] * stride_bk + offs_bn[None, :] * stride_bn + ) + + # Create pointers for the scales + offs_k_scale = (pid_k * SPLITK_BLOCK_SIZE) // GROUP_K + a_scale_ptrs = ( + a_scale_ptr + offs_am * stride_ascale_m + offs_k_scale * stride_ascale_k + ) + offs_b_scale_n = offs_bn // GROUP_N + b_scale_ptrs = ( + b_scale_ptr + + offs_k_scale * stride_bscale_k + + offs_b_scale_n * stride_bscale_n + ) + offs_ks_step = BLOCK_SIZE_K // GROUP_K + + acc_dtype = tl.float32 if c_ptr.type.element_ty != tl.int8 else tl.int32 + accumulator = tl.zeros((BLOCK_SIZE_M, BLOCK_SIZE_N), dtype=acc_dtype) + + for k in range(pid_k * num_k_iter, (pid_k + 1) * num_k_iter): + # Load the next block of A and B, generate a mask by checking the K dimension. + # If it is out of bounds, set it to 0. + if EVEN_K: + a = tl.load(a_ptrs) + b = tl.load(b_ptrs, cache_modifier=cache_modifier) + else: + a = tl.load( + a_ptrs, mask=offs_k[None, :] < K - k * BLOCK_SIZE_K, other=0.0 + ) + b = tl.load( + b_ptrs, mask=offs_k[:, None] < K - k * BLOCK_SIZE_K, other=0.0 + ) + + a_scale = tl.load(a_scale_ptrs) + b_scale = tl.load(b_scale_ptrs) + + # Perform dot operation and apply scale + accumulator += ( + tl.dot(a, b, input_precision="ieee") + * a_scale[:, None] + * b_scale[None, :] + ) + + # Advance the ptrs to the next K block. + a_ptrs += BLOCK_SIZE_K * stride_ak + b_ptrs += BLOCK_SIZE_K * stride_bk + + a_scale_ptrs += offs_ks_step * stride_ascale_k + b_scale_ptrs += offs_ks_step * stride_bscale_k + + c = accumulator.to(c_ptr.type.element_ty) + + # Write back the block of the output matrix C with masks. + offs_cm = pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M).to(tl.int64) + offs_cn = pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N).to(tl.int64) + c_ptrs = ( + c_ptr + + stride_cm * offs_cm[:, None] + + stride_cn * offs_cn[None, :] + + pid_k * stride_ck + ) + c_mask = (offs_cm[:, None] < M) & (offs_cn[None, :] < N) + tl.store(c_ptrs, c, mask=c_mask) diff --git a/tasks/triton2triton/campaign20/gemm_a8w8_blockscale_kernel/test_cases.json b/tasks/triton2triton/campaign20/gemm_a8w8_blockscale_kernel/test_cases.json new file mode 100644 index 00000000..c07ae4d2 --- /dev/null +++ b/tasks/triton2triton/campaign20/gemm_a8w8_blockscale_kernel/test_cases.json @@ -0,0 +1,107 @@ +[ + { + "test_case_id": "c2", + "count": 1, + "regime": { + "B": 2, + "seqlen": 1024, + "M": 2048, + "N": 2624, + "K": 6144, + "group_size": 128, + "scale_k": 48, + "scale_n": 21, + "role": "prefill token-parallel GEMM" + }, + "meta": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "NUM_KSPLIT": 1, + "SPLITK_BLOCK_SIZE": 6144, + "GROUP_K": 128, + "GROUP_N": 128, + "num_warps": 4, + "num_stages": 2, + "cache_modifier": ".cg" + }, + "grid_formula": "NUM_KSPLIT * cdiv(M,BLOCK_SIZE_M) * cdiv(N,BLOCK_SIZE_N)", + "params_repr": { + "B": 2, + "M": 2048, + "N": 2624, + "K": 6144 + } + }, + { + "test_case_id": "c32", + "count": 1, + "regime": { + "B": 32, + "seqlen": 1024, + "M": 32768, + "N": 2624, + "K": 6144, + "group_size": 128, + "scale_k": 48, + "scale_n": 21, + "role": "prefill token-parallel GEMM" + }, + "meta": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "NUM_KSPLIT": 1, + "SPLITK_BLOCK_SIZE": 6144, + "GROUP_K": 128, + "GROUP_N": 128, + "num_warps": 4, + "num_stages": 2, + "cache_modifier": ".cg" + }, + "grid_formula": "NUM_KSPLIT * cdiv(M,BLOCK_SIZE_M) * cdiv(N,BLOCK_SIZE_N)", + "params_repr": { + "B": 32, + "M": 32768, + "N": 2624, + "K": 6144 + } + }, + { + "test_case_id": "c64", + "count": 1, + "regime": { + "B": 64, + "seqlen": 1024, + "M": 65536, + "N": 2624, + "K": 6144, + "group_size": 128, + "scale_k": 48, + "scale_n": 21, + "role": "prefill token-parallel GEMM" + }, + "meta": { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "NUM_KSPLIT": 1, + "SPLITK_BLOCK_SIZE": 6144, + "GROUP_K": 128, + "GROUP_N": 128, + "num_warps": 4, + "num_stages": 2, + "cache_modifier": ".cg" + }, + "grid_formula": "NUM_KSPLIT * cdiv(M,BLOCK_SIZE_M) * cdiv(N,BLOCK_SIZE_N)", + "params_repr": { + "B": 64, + "M": 65536, + "N": 2624, + "K": 6144 + } + } +] \ No newline at end of file diff --git a/tasks/triton2triton/campaign20/topk_forward/Makefile b/tasks/triton2triton/campaign20/topk_forward/Makefile new file mode 100644 index 00000000..0f8b4f36 --- /dev/null +++ b/tasks/triton2triton/campaign20/topk_forward/Makefile @@ -0,0 +1,6 @@ +# Auto-generated. +all: + python3 scripts/task_runner.py compile + +clean: + rm -rf build diff --git a/tasks/triton2triton/campaign20/topk_forward/THIRD_PARTY_NOTICES.md b/tasks/triton2triton/campaign20/topk_forward/THIRD_PARTY_NOTICES.md new file mode 100644 index 00000000..3ec887d0 --- /dev/null +++ b/tasks/triton2triton/campaign20/topk_forward/THIRD_PARTY_NOTICES.md @@ -0,0 +1,27 @@ +# Triton streaming TopK + +The helper chain in `source/triton__topk_forward.py` and its frozen golden copy +comes from `triton-lang/triton` commit +`2046eb542a9c30e5bc770b7c6671f03f9adbdf55`, file +`python/triton_kernels/triton_kernels/topk_details/_topk_forward.py`. + +Copyright 2018-2020 Philippe Tillet +Copyright 2020-2022 OpenAI + +Permission is hereby granted, free of charge, to any person obtaining a copy of +this software and associated documentation files (the "Software"), to deal in +the Software without restriction, including without limitation the rights to +use, copy, modify, merge, publish, distribute, sublicense, and/or sell copies of +the Software, and to permit persons to whom the Software is furnished to do so, +subject to the following conditions: + +The above copyright notice and this permission notice shall be included in all +copies or substantial portions of the Software. + +THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +SOFTWARE. diff --git a/tasks/triton2triton/campaign20/topk_forward/config.yaml b/tasks/triton2triton/campaign20/topk_forward/config.yaml new file mode 100644 index 00000000..44749513 --- /dev/null +++ b/tasks/triton2triton/campaign20/topk_forward/config.yaml @@ -0,0 +1,29 @@ +source_file_path: +- source/triton__topk_forward.py +target_kernel_functions: +- _topk_forward +compile_command: +- python3 scripts/task_runner.py compile +correctness_command: +- python3 scripts/task_runner.py correctness +performance_command: +- python3 scripts/task_runner.py performance +task_type: triton2triton +platform_support: + required_arch: gfx950 + status: active + skip_reason: null +task_result_template: null +prompt: + source_code: null + instructions: Optimize the Triton kernel `_topk_forward` (originally triton_kernels.topk_details._topk_forward). + Captured via inspect.getsource. Tests generate inputs from test_cases.json on + the fly. + cheatsheet: null +extracted_from: + engine: vllm + kernel_id: triton:triton_kernels.topk_details._topk_forward._topk_forward + triton_fn: _topk_forward + origin_module: triton_kernels.topk_details._topk_forward + source: triton_inspect + test_cases: 3 diff --git a/tasks/triton2triton/campaign20/topk_forward/scripts/harness_run.py b/tasks/triton2triton/campaign20/topk_forward/scripts/harness_run.py new file mode 100644 index 00000000..0fd4aa14 --- /dev/null +++ b/tasks/triton2triton/campaign20/topk_forward/scripts/harness_run.py @@ -0,0 +1,165 @@ +#!/usr/bin/env python3 +"""Real launcher + benchmark for the compile-only triton kernel `_topk_forward` +(multi-peer variant captured from triton_kernels.topk_details). + +The editable source and frozen golden both include the complete +``streaming_topk`` helper chain, so this harness has no dependency on an +installed ``triton_kernels`` package. + +Workload regime (token-parallel router op): + n_rows (M) = B * 1024 for B in {2, 32, 64} -> ids c2, c32, c64 +Model dims kept from the captured base case: + n_expts_tot = 128, N_EXPTS_ACT (topk k) = 4, N_EXPTS_PAD = 128, + BLOCK_M = 32, BLOCK_N = 32, APPLY_SOFTMAX = True, USE_PROVIDED_INDX = False. +grid = (cdiv(n_rows, BLOCK_M),) -- recomputed per regime (not the captured const). +""" +import os, sys, json, importlib.util +import torch +import triton + +TASK_DIR = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) +SOURCE_FILE = os.path.join(TASK_DIR, "source", "triton__topk_forward.py") +GOLDEN_FILE = os.path.join(TASK_DIR, "source_golden", "triton__topk_forward.py") +BUILD_DIR = os.path.join(TASK_DIR, "build") +KERNEL_NAME = "_topk_forward" + +# Fixed model dims (from captured base case) +N_EXPTS_TOT = 128 +N_EXPTS_ACT = 4 +N_EXPTS_PAD = 128 +BLOCK_M = 32 +BLOCK_N = 32 +APPLY_SOFTMAX = True +USE_PROVIDED_INDX = False +DST_OFFS_M = 0 + +# Regime: B in {2,32,64}; M = B*1024 +REGIME = [("c2", 2), ("c32", 32), ("c64", 64)] + + +def _cdiv(a, b): + return (a + b - 1) // b + + +def _load_kernel(path, modname): + """Load the self-contained kernel module.""" + spec = importlib.util.spec_from_file_location(modname, path) + mod = importlib.util.module_from_spec(spec) + spec.loader.exec_module(mod) + return getattr(mod, KERNEL_NAME) + + +def _build_inputs(n_rows, seed=42): + """Deterministic seeded inputs for one regime case.""" + dev = "cuda:0" + g = torch.Generator(device=dev).manual_seed(seed) + X = torch.randn((n_rows, N_EXPTS_TOT), dtype=torch.bfloat16, device=dev, generator=g) + # outputs (mutated in place) -- one peer + Yv = torch.zeros((n_rows, N_EXPTS_ACT), dtype=torch.bfloat16, device=dev) + Yi = torch.zeros((n_rows, N_EXPTS_ACT), dtype=torch.int16, device=dev) + n_cols_words = _cdiv(N_EXPTS_TOT, 32) # = 4 + # bitmatrix stored [words, n_rows] then conceptually transposed; kernel uses + # stride_rm=1 (row stride over rows) and stride_rn=n_rows (word stride). + Bits = torch.zeros((n_cols_words, n_rows), dtype=torch.uint32, device=dev) + PeerYvs = (Yv,) + PeerYis = (Yi,) + PeerBits = (Bits,) + stride_xm = N_EXPTS_TOT + stride_ym = N_EXPTS_ACT + stride_rm = 1 + stride_rn = n_rows + args = [ + X, stride_xm, + PeerYvs, PeerYis, stride_ym, + USE_PROVIDED_INDX, PeerBits, stride_rm, stride_rn, + n_rows, N_EXPTS_TOT, + DST_OFFS_M, + ] + kwargs = dict( + APPLY_SOFTMAX=APPLY_SOFTMAX, + BLOCK_M=BLOCK_M, N_EXPTS_PAD=N_EXPTS_PAD, + N_EXPTS_ACT=N_EXPTS_ACT, BLOCK_N=BLOCK_N, + ) + return args, kwargs, (Yv, Yi, Bits) + + +def _launch(kern, n_rows, seed=42): + args, kwargs, outs = _build_inputs(n_rows, seed=seed) + grid = (max(_cdiv(n_rows, BLOCK_M), 1),) + kern[grid](*args, **kwargs) + torch.cuda.synchronize() + return outs + + +def _compare(a_outs, b_outs): + """Compare (Yv float, Yi index, Bits uint) edited-vs-golden. + Values: cosine>=0.99 & small max-rel; indices/bits: exact.""" + Yv_a, Yi_a, Bits_a = a_outs + Yv_b, Yi_b, Bits_b = b_outs + # exact for integer index + bitmatrix + if not torch.equal(Yi_a, Yi_b): + nmis = (Yi_a != Yi_b).sum().item() + return f"index mismatch: {nmis} differing entries" + if not torch.equal(Bits_a, Bits_b): + nmis = (Bits_a != Bits_b).sum().item() + return f"bitmatrix mismatch: {nmis} differing words" + fa = Yv_a.to(torch.float32).flatten() + fb = Yv_b.to(torch.float32).flatten() + cos = torch.nn.functional.cosine_similarity(fa, fb, dim=0).item() + denom = fb.abs().clamp_min(1e-6) + max_rel = ((fa - fb).abs() / denom).max().item() + if cos < 0.99: + return f"values cosine {cos:.5f} < 0.99" + if max_rel > 1e-2: + return f"values max_rel {max_rel:.4e} > 1e-2" + return None + + +def run_correctness(): + edit = _load_kernel(SOURCE_FILE, "topk_edit") + gold = _load_kernel(GOLDEN_FILE, "topk_gold") + for cid, B in REGIME: + n_rows = B * 1024 + out_e = _launch(edit, n_rows, seed=42) + out_g = _launch(gold, n_rows, seed=42) + err = _compare(out_e, out_g) + if err: + return False, f"{cid} (n_rows={n_rows}): {err}" + return True, None + + +def compile_smoke(): + """JIT-compile and launch the smallest workload case.""" + _launch(_load_kernel(SOURCE_FILE, "topk_compile"), REGIME[0][1] * 1024) + + +def run_performance(benchmark): + edit = _load_kernel(SOURCE_FILE, "topk_perf") + results = [] + for cid, B in REGIME: + n_rows = B * 1024 + args, kwargs, _ = _build_inputs(n_rows, seed=42) + grid = (max(_cdiv(n_rows, BLOCK_M), 1),) + avg, metadata = benchmark( + lambda: edit[grid](*args, **kwargs), + warmup=10, + repetition=100, + ) + results.append({ + "test_case_id": cid, + "execution_time_ms": avg, + **metadata, + "params": {"B": B, "n_rows": n_rows, "n_expts_tot": N_EXPTS_TOT, + "topk": N_EXPTS_ACT, "BLOCK_M": BLOCK_M, "BLOCK_N": BLOCK_N}, + }) + return results + + +if __name__ == "__main__": + mode = sys.argv[1] if len(sys.argv) > 1 else "correctness" + os.makedirs(BUILD_DIR, exist_ok=True) + if mode == "correctness": + ok, err = run_correctness() + print("Correctness:", "PASS" if ok else "FAIL", err or "") + else: + raise SystemExit("run performance through scripts/task_runner.py") diff --git a/tasks/triton2triton/campaign20/topk_forward/scripts/task_runner.py b/tasks/triton2triton/campaign20/topk_forward/scripts/task_runner.py new file mode 100644 index 00000000..88eaa11f --- /dev/null +++ b/tasks/triton2triton/campaign20/topk_forward/scripts/task_runner.py @@ -0,0 +1,118 @@ +#!/usr/bin/env python3 +"""Auto-generated task runner for triton__topk_forward (Triton). + +Inputs are generated each run from the shape/dtype signatures in +test_cases.json. The launcher symbol expected in source/triton__topk_forward.py is +``_topk_forward`` (for raw @triton.jit kernels) or ``None`` (a +wrapper that handles grid + meta resolution if present). +""" +import sys, os, json, argparse, glob, importlib.util +sys.path.insert(0, os.path.dirname(os.path.abspath(__file__))) + +# >>> AKA-GENERATED: shared CUDA-graph benchmark helpers - edit src/tools/perf/vllm_cuda_graph_block.py then run `make sync-perf-helpers` >>> +def _measure_cuda_event_fallback(*args, **kwargs): + raise RuntimeError( + "CUDA-graph benchmark helpers were not materialized. " + "Run this task through AgentKernelArena so setup_workspace() can inject " + "src/tools/perf/vllm_cuda_graph_block.py into the workspace." + ) + + +def _benchmark_cuda_graph_or_events(*args, **kwargs): + raise RuntimeError( + "CUDA-graph benchmark helpers were not materialized. " + "Run this task through AgentKernelArena so setup_workspace() can inject " + "src/tools/perf/vllm_cuda_graph_block.py into the workspace." + ) +# <<< AKA-GENERATED <<< + + +TASK_DIR = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) +os.chdir(TASK_DIR) +TASK_NAME = "triton2triton/campaign20/topk_forward" +SOURCE_FILE = os.path.join(TASK_DIR, "source", "triton__topk_forward.py") +KERNEL_NAME = "_topk_forward" +LAUNCHER_NAME = "None" +REF_SOURCE = "triton" +TEST_CASES = os.path.join(TASK_DIR, "test_cases.json") + + +def _load_module(): + spec = importlib.util.spec_from_file_location("triton_kernel", SOURCE_FILE) + mod = importlib.util.module_from_spec(spec) + spec.loader.exec_module(mod) + return mod + + +def _test_cases(): + if not os.path.isfile(TEST_CASES): + return [] + with open(TEST_CASES) as f: + return json.load(f) + + +def _resolve_callable(mod): + # Prefer a wrapping launcher if it was found alongside the @triton.jit kernel. + for n in (LAUNCHER_NAME, KERNEL_NAME): + if n and hasattr(mod, n) and n != "None": + return getattr(mod, n) + raise AttributeError(f"neither {LAUNCHER_NAME} nor {KERNEL_NAME} found in source") + + +def run_compile(): + try: + import ast + ast.parse(open(SOURCE_FILE).read()) + mod = _load_module() + _resolve_callable(mod) + import harness_run + harness_run.compile_smoke() + return True, None + except Exception as e: + return False, str(e) + + +def run_correctness(): + """Real launcher-based correctness for this compile-only @triton.jit kernel. + Delegates to scripts/harness_run.py which builds a real launcher, runs the + editable source vs the frozen golden copy on identical seeded inputs, and + compares outputs (values cosine>=0.99/small max-rel; indices+bitmatrix + exact).""" + import harness_run + return harness_run.run_correctness() + + +def run_performance(): + """Canonical CUDA-graph timing with event fallback.""" + import harness_run + return harness_run.run_performance(_benchmark_cuda_graph_or_events) + + +def main(): + ap = argparse.ArgumentParser() + ap.add_argument("mode", choices=["compile", "correctness", "performance"]) + args = ap.parse_args() + build_dir = os.path.join(TASK_DIR, "build") + os.makedirs(build_dir, exist_ok=True) + if args.mode == "compile": + ok, err = run_compile() + json.dump({"status": "ok" if ok else "fail", "error": err}, open(os.path.join(build_dir, "compile_report.json"), "w")) + print(f"Compilation: {'PASS' if ok else 'FAIL'}") + if err: print("Error:", err) + sys.exit(0 if ok else 1) + if args.mode == "correctness": + ok, err = run_correctness() + json.dump({"status": "ok" if ok else "fail", "error": err}, open(os.path.join(build_dir, "correctness_report.json"), "w")) + print(f"Correctness: {'PASS' if ok else 'FAIL'}") + if err: print("Error:", err) + sys.exit(0 if ok else 1) + cases = run_performance() + json.dump({"test_cases": cases}, open(os.path.join(build_dir, "performance_report.json"), "w"), indent=2) + for c in cases: + print(f"Performance: {c['execution_time_ms']:.4f} ms ({c['test_case_id']})") + valid = bool(cases) and all(c.get("execution_time_ms", -1) > 0 for c in cases) + sys.exit(0 if valid else 1) + + +if __name__ == "__main__": + main() diff --git a/tasks/triton2triton/campaign20/topk_forward/source/triton__topk_forward.py b/tasks/triton2triton/campaign20/topk_forward/source/triton__topk_forward.py new file mode 100644 index 00000000..2f95457c --- /dev/null +++ b/tasks/triton2triton/campaign20/topk_forward/source/triton__topk_forward.py @@ -0,0 +1,158 @@ +"""Self-contained Triton streaming TopK kernel. + +The helper chain was restored from triton-lang/triton commit +2046eb542a9c30e5bc770b7c6671f03f9adbdf55 so the task does not depend on an +installed ``triton_kernels`` package. That upstream code is MIT licensed. +""" + +import triton +import triton.language as tl + + +@triton.jit +def get_topmask_and_fullmask(x): + tl.static_assert( + x.dtype.is_int_unsigned(), + "floating-point value must be passed as bits", + ) + tm: tl.constexpr = 1 << (-1 + x.dtype.primitive_bitwidth) + fm: tl.constexpr = (1 << x.dtype.primitive_bitwidth) - 1 + tm_arr = tl.full(x.shape, tm, dtype=x.dtype) + fm_arr = tl.full(x.shape, fm, dtype=x.dtype) + return tm_arr, fm_arr + + +@triton.jit +def fpval_to_key(x): + tm, fm = get_topmask_and_fullmask(x) + return x ^ tl.where((x & tm) != 0, fm, tm) + + +@triton.jit +def key_to_fpval(x): + tm, fm = get_topmask_and_fullmask(x) + return x ^ tl.where((x & tm) == 0, fm, tm) + + +@triton.jit +def indx_to_key(indx, N_EXPTS_PAD: tl.constexpr): + """Prefer the smaller expert index when values tie.""" + return N_EXPTS_PAD - indx + + +@triton.jit +def key_to_indx(indx, N_EXPTS_PAD: tl.constexpr): + return N_EXPTS_PAD - indx + + +@triton.jit +def streaming_topk( + X, + stride_xm, + n_expts_tot, + offs_m, + mask_m, + N_EXPTS_PAD: tl.constexpr, + N_EXPTS_ACT: tl.constexpr, + BLOCK_N: tl.constexpr, +): + x_nbits: tl.constexpr = X.dtype.element_ty.primitive_bitwidth + x_utype: tl.constexpr = tl.dtype(f"uint{x_nbits}") + y_nbits: tl.constexpr = 32 if x_nbits < 16 else x_nbits * 2 + x_ultype: tl.constexpr = tl.dtype(f"uint{y_nbits}") + x_dtype: tl.constexpr = X.dtype.element_ty + + loop_iterations: tl.constexpr = N_EXPTS_PAD // BLOCK_N - 1 + offs_x_n = loop_iterations * BLOCK_N + tl.arange(0, BLOCK_N) + mask_n = offs_x_n[None, :] < n_expts_tot + + X_ptrs = X + offs_m[:, None] * stride_xm + offs_x_n[None, :] + x = tl.load(X_ptrs, mask=(mask_m & mask_n), other=float("-inf")) + x = fpval_to_key(x.to(x_utype, bitcast=True)) + x = (x.to(x_ultype) << 16) | indx_to_key( + offs_x_n, N_EXPTS_PAD + )[None, :] + x = tl.where(mask_n, x, 0) + acc = tl.topk(x, N_EXPTS_ACT, dim=1) + + for _i in (tl.static_range if loop_iterations <= 4 else range)( + loop_iterations + ): + acc = tl.bitonic_merge(acc) + X_ptrs -= BLOCK_N + offs_x_n -= BLOCK_N + x = tl.load(X_ptrs, mask=mask_m, other=float("-inf")) + x = fpval_to_key(x.to(x_utype, bitcast=True)) + x = (x.to(x_ultype) << 16) | indx_to_key( + offs_x_n, N_EXPTS_PAD + )[None, :] + acc = tl.maximum(acc, tl.topk(x, N_EXPTS_ACT, dim=1)) + + acc = tl.sort(acc, dim=1, descending=True) + y_indices_raw = (acc & 0xFFFF).to(tl.uint32) + y_indices = key_to_indx(y_indices_raw, N_EXPTS_PAD) + y_values_raw = (acc >> 16).to(x_utype) + y_values = key_to_fpval(y_values_raw).to(x_dtype, bitcast=True) + return y_values, y_indices + +@triton.jit +def _topk_forward(X, stride_xm, # inputs + PeerYvs, PeerYis, stride_ym, # topk values/indices + USE_PROVIDED_INDX: tl.constexpr, PeerBits, stride_rm: tl.constexpr, + stride_rn: tl.constexpr, # bitmatrix + n_rows, n_expts_tot, # shape + dst_offs_m, APPLY_SOFTMAX: tl.constexpr, # constant + BLOCK_M: tl.constexpr, N_EXPTS_PAD: tl.constexpr, N_EXPTS_ACT: tl.constexpr, BLOCK_N: tl.constexpr): + + N_PEERS: tl.constexpr = len(PeerYvs) + + pid = tl.program_id(0) + if isinstance(n_rows, tl.tensor) and n_rows.dtype.is_ptr(): + n_rows = tl.load(n_rows) + + if pid * BLOCK_M >= n_rows: + # early exit: + return + + tl.static_assert(BLOCK_N % 32 == 0) + tl.static_assert(N_EXPTS_PAD % BLOCK_N == 0) + x_dtype: tl.constexpr = X.dtype.element_ty + + # load logits + offs_m = pid * BLOCK_M + tl.arange(0, BLOCK_M) + offs_y_n = tl.arange(0, N_EXPTS_ACT) + mask_m = offs_m[:, None] < n_rows + if USE_PROVIDED_INDX: + tl.static_assert(len(PeerYis) == 1) + Yi_ptrs = PeerYis[0] + (dst_offs_m + offs_m[:, None]) * stride_ym + offs_y_n[None, :] + y_indices = tl.load(Yi_ptrs, mask=mask_m) + Xv_ptrs = X + offs_m[:, None] * stride_xm + y_indices + y_values = tl.load(Xv_ptrs, mask=mask_m) + else: + y_values, y_indices = streaming_topk(X, stride_xm, n_expts_tot, offs_m, mask_m, # + N_EXPTS_PAD, N_EXPTS_ACT, BLOCK_N) + + # normalize selected values + if APPLY_SOFTMAX: + y_values = tl.softmax(y_values.to(tl.float32), dim=1, keep_dims=True).to(x_dtype) + + # write back + for rank in tl.static_range(N_PEERS): + Yv_ptrs = PeerYvs[rank] + (dst_offs_m + offs_m[:, None]) * stride_ym + offs_y_n[None, :] + tl.store(Yv_ptrs, y_values, mask=mask_m) + if not USE_PROVIDED_INDX: + for rank in tl.static_range(N_PEERS): + Yi_ptrs = PeerYis[rank] + (dst_offs_m + offs_m[:, None]) * stride_ym + offs_y_n[None, :] + tl.store(Yi_ptrs, y_indices, mask=mask_m) + + # pack into bitmatrix + y_div = y_indices // 32 + y_rem = y_indices % 32 + loop_iterations = N_EXPTS_PAD // BLOCK_N + for i in range(loop_iterations): + offs_r_n = tl.arange(0, BLOCK_N // 32) + i * (BLOCK_N // 32) + y2 = tl.where(y_div[:, :, None] == offs_r_n[None, None, :], (1 << y_rem)[:, :, None], 0) + r = tl.reduce_or(y2, axis=1) + for rank in tl.static_range(N_PEERS): + BitsPtrs = PeerBits[rank] + (dst_offs_m + offs_m[:, None]) * stride_rm + offs_r_n[None, :] * stride_rn + tl.store(BitsPtrs, r, mask=mask_m) diff --git a/tasks/triton2triton/campaign20/topk_forward/source_golden/triton__topk_forward.py b/tasks/triton2triton/campaign20/topk_forward/source_golden/triton__topk_forward.py new file mode 100644 index 00000000..2f95457c --- /dev/null +++ b/tasks/triton2triton/campaign20/topk_forward/source_golden/triton__topk_forward.py @@ -0,0 +1,158 @@ +"""Self-contained Triton streaming TopK kernel. + +The helper chain was restored from triton-lang/triton commit +2046eb542a9c30e5bc770b7c6671f03f9adbdf55 so the task does not depend on an +installed ``triton_kernels`` package. That upstream code is MIT licensed. +""" + +import triton +import triton.language as tl + + +@triton.jit +def get_topmask_and_fullmask(x): + tl.static_assert( + x.dtype.is_int_unsigned(), + "floating-point value must be passed as bits", + ) + tm: tl.constexpr = 1 << (-1 + x.dtype.primitive_bitwidth) + fm: tl.constexpr = (1 << x.dtype.primitive_bitwidth) - 1 + tm_arr = tl.full(x.shape, tm, dtype=x.dtype) + fm_arr = tl.full(x.shape, fm, dtype=x.dtype) + return tm_arr, fm_arr + + +@triton.jit +def fpval_to_key(x): + tm, fm = get_topmask_and_fullmask(x) + return x ^ tl.where((x & tm) != 0, fm, tm) + + +@triton.jit +def key_to_fpval(x): + tm, fm = get_topmask_and_fullmask(x) + return x ^ tl.where((x & tm) == 0, fm, tm) + + +@triton.jit +def indx_to_key(indx, N_EXPTS_PAD: tl.constexpr): + """Prefer the smaller expert index when values tie.""" + return N_EXPTS_PAD - indx + + +@triton.jit +def key_to_indx(indx, N_EXPTS_PAD: tl.constexpr): + return N_EXPTS_PAD - indx + + +@triton.jit +def streaming_topk( + X, + stride_xm, + n_expts_tot, + offs_m, + mask_m, + N_EXPTS_PAD: tl.constexpr, + N_EXPTS_ACT: tl.constexpr, + BLOCK_N: tl.constexpr, +): + x_nbits: tl.constexpr = X.dtype.element_ty.primitive_bitwidth + x_utype: tl.constexpr = tl.dtype(f"uint{x_nbits}") + y_nbits: tl.constexpr = 32 if x_nbits < 16 else x_nbits * 2 + x_ultype: tl.constexpr = tl.dtype(f"uint{y_nbits}") + x_dtype: tl.constexpr = X.dtype.element_ty + + loop_iterations: tl.constexpr = N_EXPTS_PAD // BLOCK_N - 1 + offs_x_n = loop_iterations * BLOCK_N + tl.arange(0, BLOCK_N) + mask_n = offs_x_n[None, :] < n_expts_tot + + X_ptrs = X + offs_m[:, None] * stride_xm + offs_x_n[None, :] + x = tl.load(X_ptrs, mask=(mask_m & mask_n), other=float("-inf")) + x = fpval_to_key(x.to(x_utype, bitcast=True)) + x = (x.to(x_ultype) << 16) | indx_to_key( + offs_x_n, N_EXPTS_PAD + )[None, :] + x = tl.where(mask_n, x, 0) + acc = tl.topk(x, N_EXPTS_ACT, dim=1) + + for _i in (tl.static_range if loop_iterations <= 4 else range)( + loop_iterations + ): + acc = tl.bitonic_merge(acc) + X_ptrs -= BLOCK_N + offs_x_n -= BLOCK_N + x = tl.load(X_ptrs, mask=mask_m, other=float("-inf")) + x = fpval_to_key(x.to(x_utype, bitcast=True)) + x = (x.to(x_ultype) << 16) | indx_to_key( + offs_x_n, N_EXPTS_PAD + )[None, :] + acc = tl.maximum(acc, tl.topk(x, N_EXPTS_ACT, dim=1)) + + acc = tl.sort(acc, dim=1, descending=True) + y_indices_raw = (acc & 0xFFFF).to(tl.uint32) + y_indices = key_to_indx(y_indices_raw, N_EXPTS_PAD) + y_values_raw = (acc >> 16).to(x_utype) + y_values = key_to_fpval(y_values_raw).to(x_dtype, bitcast=True) + return y_values, y_indices + +@triton.jit +def _topk_forward(X, stride_xm, # inputs + PeerYvs, PeerYis, stride_ym, # topk values/indices + USE_PROVIDED_INDX: tl.constexpr, PeerBits, stride_rm: tl.constexpr, + stride_rn: tl.constexpr, # bitmatrix + n_rows, n_expts_tot, # shape + dst_offs_m, APPLY_SOFTMAX: tl.constexpr, # constant + BLOCK_M: tl.constexpr, N_EXPTS_PAD: tl.constexpr, N_EXPTS_ACT: tl.constexpr, BLOCK_N: tl.constexpr): + + N_PEERS: tl.constexpr = len(PeerYvs) + + pid = tl.program_id(0) + if isinstance(n_rows, tl.tensor) and n_rows.dtype.is_ptr(): + n_rows = tl.load(n_rows) + + if pid * BLOCK_M >= n_rows: + # early exit: + return + + tl.static_assert(BLOCK_N % 32 == 0) + tl.static_assert(N_EXPTS_PAD % BLOCK_N == 0) + x_dtype: tl.constexpr = X.dtype.element_ty + + # load logits + offs_m = pid * BLOCK_M + tl.arange(0, BLOCK_M) + offs_y_n = tl.arange(0, N_EXPTS_ACT) + mask_m = offs_m[:, None] < n_rows + if USE_PROVIDED_INDX: + tl.static_assert(len(PeerYis) == 1) + Yi_ptrs = PeerYis[0] + (dst_offs_m + offs_m[:, None]) * stride_ym + offs_y_n[None, :] + y_indices = tl.load(Yi_ptrs, mask=mask_m) + Xv_ptrs = X + offs_m[:, None] * stride_xm + y_indices + y_values = tl.load(Xv_ptrs, mask=mask_m) + else: + y_values, y_indices = streaming_topk(X, stride_xm, n_expts_tot, offs_m, mask_m, # + N_EXPTS_PAD, N_EXPTS_ACT, BLOCK_N) + + # normalize selected values + if APPLY_SOFTMAX: + y_values = tl.softmax(y_values.to(tl.float32), dim=1, keep_dims=True).to(x_dtype) + + # write back + for rank in tl.static_range(N_PEERS): + Yv_ptrs = PeerYvs[rank] + (dst_offs_m + offs_m[:, None]) * stride_ym + offs_y_n[None, :] + tl.store(Yv_ptrs, y_values, mask=mask_m) + if not USE_PROVIDED_INDX: + for rank in tl.static_range(N_PEERS): + Yi_ptrs = PeerYis[rank] + (dst_offs_m + offs_m[:, None]) * stride_ym + offs_y_n[None, :] + tl.store(Yi_ptrs, y_indices, mask=mask_m) + + # pack into bitmatrix + y_div = y_indices // 32 + y_rem = y_indices % 32 + loop_iterations = N_EXPTS_PAD // BLOCK_N + for i in range(loop_iterations): + offs_r_n = tl.arange(0, BLOCK_N // 32) + i * (BLOCK_N // 32) + y2 = tl.where(y_div[:, :, None] == offs_r_n[None, None, :], (1 << y_rem)[:, :, None], 0) + r = tl.reduce_or(y2, axis=1) + for rank in tl.static_range(N_PEERS): + BitsPtrs = PeerBits[rank] + (dst_offs_m + offs_m[:, None]) * stride_rm + offs_r_n[None, :] * stride_rn + tl.store(BitsPtrs, r, mask=mask_m) diff --git a/tasks/triton2triton/campaign20/topk_forward/test_cases.json b/tasks/triton2triton/campaign20/topk_forward/test_cases.json new file mode 100644 index 00000000..3f86807e --- /dev/null +++ b/tasks/triton2triton/campaign20/topk_forward/test_cases.json @@ -0,0 +1,50 @@ +[ + { + "test_case_id": "c2", + "params_repr": { + "B": 2, + "n_rows": 2048, + "n_expts_tot": 128, + "topk": 4, + "BLOCK_M": 32, + "BLOCK_N": 32, + "N_EXPTS_PAD": 128, + "N_EXPTS_ACT": 4, + "APPLY_SOFTMAX": true, + "USE_PROVIDED_INDX": false, + "grid": "(cdiv(n_rows,BLOCK_M),) = (64,)" + } + }, + { + "test_case_id": "c32", + "params_repr": { + "B": 32, + "n_rows": 32768, + "n_expts_tot": 128, + "topk": 4, + "BLOCK_M": 32, + "BLOCK_N": 32, + "N_EXPTS_PAD": 128, + "N_EXPTS_ACT": 4, + "APPLY_SOFTMAX": true, + "USE_PROVIDED_INDX": false, + "grid": "(cdiv(n_rows,BLOCK_M),) = (1024,)" + } + }, + { + "test_case_id": "c64", + "params_repr": { + "B": 64, + "n_rows": 65536, + "n_expts_tot": 128, + "topk": 4, + "BLOCK_M": 32, + "BLOCK_N": 32, + "N_EXPTS_PAD": 128, + "N_EXPTS_ACT": 4, + "APPLY_SOFTMAX": true, + "USE_PROVIDED_INDX": false, + "grid": "(cdiv(n_rows,BLOCK_M),) = (2048,)" + } + } +] diff --git a/tasks/triton2triton/campaign20/write_req_to_token_pool_triton/Makefile b/tasks/triton2triton/campaign20/write_req_to_token_pool_triton/Makefile new file mode 100644 index 00000000..0f8b4f36 --- /dev/null +++ b/tasks/triton2triton/campaign20/write_req_to_token_pool_triton/Makefile @@ -0,0 +1,6 @@ +# Auto-generated. +all: + python3 scripts/task_runner.py compile + +clean: + rm -rf build diff --git a/tasks/triton2triton/campaign20/write_req_to_token_pool_triton/config.yaml b/tasks/triton2triton/campaign20/write_req_to_token_pool_triton/config.yaml new file mode 100644 index 00000000..8c69a71d --- /dev/null +++ b/tasks/triton2triton/campaign20/write_req_to_token_pool_triton/config.yaml @@ -0,0 +1,29 @@ +source_file_path: +- source/triton_write_req_to_token_pool_triton.py +target_kernel_functions: +- write_req_to_token_pool_triton +compile_command: +- python3 scripts/task_runner.py compile +correctness_command: +- python3 scripts/task_runner.py correctness +performance_command: +- python3 scripts/task_runner.py performance +task_type: triton2triton +platform_support: + required_arch: gfx950 + status: active + skip_reason: null +task_result_template: null +prompt: + source_code: null + instructions: Optimize the Triton kernel `write_req_to_token_pool_triton` (originally + sglang.srt.mem_cache.common). Captured via inspect.getsource. Tests generate inputs + from test_cases.json on the fly. + cheatsheet: null +extracted_from: + engine: sglang + kernel_id: triton:sglang.srt.mem_cache.common.write_req_to_token_pool_triton + triton_fn: write_req_to_token_pool_triton + origin_module: sglang.srt.mem_cache.common + source: triton_inspect + test_cases: 3 diff --git a/tasks/triton2triton/campaign20/write_req_to_token_pool_triton/scripts/harness_run.py b/tasks/triton2triton/campaign20/write_req_to_token_pool_triton/scripts/harness_run.py new file mode 100644 index 00000000..608de2cc --- /dev/null +++ b/tasks/triton2triton/campaign20/write_req_to_token_pool_triton/scripts/harness_run.py @@ -0,0 +1,176 @@ +#!/usr/bin/env python3 +"""Real launcher + benchmark for the compile-only @triton.jit kernel +``write_req_to_token_pool_triton`` (sglang.srt.mem_cache.common). + +WORKLOAD REGIME: input seqlen = output seqlen = 1024; concurrency B in {2,32,64}. +This is a token-parallel KV-write op. grid = (B,) (one program per request). +We run it prefill-style: each request writes seq_len=1024 tokens into its +req_to_token row. We use a small prefix (pre_len=PREFIX_LEN) so BOTH the +prefix-copy loop and the extend-copy loop are exercised, and provide REAL +device pointers for ``prefix_tensors`` so the int64* dereference is valid. + +shape_mapping: + B in {2,32,64} -> grid=(B,), one program per request (token-parallel). + seq_len = 1024 per request (prefill: input==output==1024). + pre_len = PREFIX_LEN (small captured-style prefix), extend_len = seq_len - pre_len. + out_cache_loc total = sum(extend_lens) = B * (1024 - PREFIX_LEN). + Model dims kept from captured base case: + req_to_token_ptr columns (max_context_len) = 202756, stride = 202756. + req_to_token_ptr rows (max_batch) = 2048. +""" +import importlib.util +import json +import os +import sys + +import torch + +TASK_DIR = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) +SOURCE_FILE = os.path.join(TASK_DIR, "source", "triton_write_req_to_token_pool_triton.py") +GOLDEN_FILE = os.path.join(TASK_DIR, "source_golden", "triton_write_req_to_token_pool_triton.py") +KERNEL_NAME = "write_req_to_token_pool_triton" + +# captured model dims +MAX_BATCH = 2048 +MAX_CONTEXT_LEN = 202756 +STRIDE = 202756 +SEQ_LEN = 1024 +PREFIX_LEN = 16 # small prefix so both loops run; pointer deref is exercised + +CASES = [("c2", 2), ("c32", 32), ("c64", 64)] + + +def _load(path, name): + spec = importlib.util.spec_from_file_location(name, path) + mod = importlib.util.module_from_spec(spec) + spec.loader.exec_module(mod) + return getattr(mod, KERNEL_NAME) + + +def build_inputs(B, seed=42): + """Deterministic, seeded inputs for one concurrency case. + + Returns (args_tuple, grid, keep_alive) where keep_alive holds python refs + to the per-request prefix tensors so their data_ptr() stay valid during + the launch. + """ + g = torch.Generator(device="cuda").manual_seed(seed + B) + + # output buffer (the thing the kernel writes into) -- mutated in place + req_to_token = torch.zeros((B, MAX_CONTEXT_LEN), dtype=torch.int32, device="cuda") + + # which row each request writes to: a unique permutation of rows + req_pool_indices = torch.randperm(B, generator=g, device="cuda").to(torch.int64) + + pre_lens = torch.full((B,), PREFIX_LEN, dtype=torch.int64, device="cuda") + seq_lens = torch.full((B,), SEQ_LEN, dtype=torch.int64, device="cuda") + extend_lens = (seq_lens - pre_lens).to(torch.int64) # all == SEQ_LEN-PREFIX_LEN + + total_extend = int(extend_lens.sum().item()) + out_cache_loc = torch.randint( + 0, MAX_CONTEXT_LEN, (total_extend,), dtype=torch.int64, device="cuda", generator=g + ) + + # prefix_tensors: a uint64 array holding device pointers to per-request + # int64 prefix buffers of length pre_len. Keep python refs alive. + keep_alive = [] + ptrs = torch.empty((B,), dtype=torch.uint64, device="cuda") + ptr_vals = [] + for i in range(B): + buf = torch.randint( + 0, MAX_CONTEXT_LEN, (PREFIX_LEN,), dtype=torch.int64, device="cuda", generator=g + ) + keep_alive.append(buf) + ptr_vals.append(buf.data_ptr()) + ptrs.copy_(torch.tensor(ptr_vals, dtype=torch.uint64, device="cuda")) + + args = ( + req_to_token, + req_pool_indices, + ptrs, # prefix_tensors + pre_lens, + seq_lens, + extend_lens, + out_cache_loc, + STRIDE, # req_to_token_ptr_stride (constexpr int) + ) + grid = (B,) + return args, grid, keep_alive + + +def reference_output(B, seed=42): + """Pure-torch golden of what the kernel should write into req_to_token, + computed from the SAME seeded inputs. Returns the mutated req_to_token.""" + args, grid, keep = build_inputs(B, seed=seed) + req_to_token, req_pool_indices, ptrs, pre_lens, seq_lens, extend_lens, out_cache_loc, stride = args + out = torch.zeros_like(req_to_token) + cumsum = 0 + for pid in range(B): + row = int(req_pool_indices[pid].item()) + pl = int(pre_lens[pid].item()) + sl = int(seq_lens[pid].item()) + # prefix region + prefix = keep[pid][:pl].to(torch.int32) + out[row, 0:pl] = prefix + # extend region + n = sl - pl + seg = out_cache_loc[cumsum:cumsum + n].to(torch.int32) + out[row, pl:sl] = seg + cumsum += int(extend_lens[pid].item()) + return out + + +def run_kernel(kern, B, seed=42): + args, grid, keep = build_inputs(B, seed=seed) + kern[grid](*args) + torch.cuda.synchronize() + return args[0] # req_to_token (mutated in place) + + +def correctness(): + kern_edit = _load(SOURCE_FILE, "k_edit") + kern_gold = _load(GOLDEN_FILE, "k_gold") + for cid, B in CASES: + got = run_kernel(kern_edit, B, seed=42) + # golden run with identical inputs (same seed -> same tensors) + gold = run_kernel(kern_gold, B, seed=42) + ref = reference_output(B, seed=42) + # integer/index kernel: require EXACT match + if not torch.equal(got, gold): + return False, f"{cid}: edited vs golden mismatch ({(got != gold).sum().item()} elems)" + if not torch.equal(got, ref): + return False, f"{cid}: edited vs torch-reference mismatch ({(got != ref).sum().item()} elems)" + return True, None + + +def compile_smoke(): + """JIT-compile and launch the smallest workload case.""" + run_kernel(_load(SOURCE_FILE, "k_compile"), CASES[0][1], seed=42) + + +def performance(benchmark): + kern = _load(SOURCE_FILE, "k_perf") + results = [] + for cid, B in CASES: + args, grid, keep = build_inputs(B, seed=42) + avg, metadata = benchmark( + lambda: kern[grid](*args), + warmup=10, + repetition=100, + ) + results.append({ + "test_case_id": cid, + "execution_time_ms": avg, + **metadata, + "params": {"B": B, "seq_len": SEQ_LEN, "pre_len": PREFIX_LEN, "grid": list(grid)}, + }) + return results + + +if __name__ == "__main__": + mode = sys.argv[1] if len(sys.argv) > 1 else "correctness" + if mode == "correctness": + ok, err = correctness() + print("CORRECTNESS", "PASS" if ok else "FAIL", err or "") + elif mode == "performance": + raise SystemExit("run performance through scripts/task_runner.py") diff --git a/tasks/triton2triton/campaign20/write_req_to_token_pool_triton/scripts/task_runner.py b/tasks/triton2triton/campaign20/write_req_to_token_pool_triton/scripts/task_runner.py new file mode 100644 index 00000000..a93a31e3 --- /dev/null +++ b/tasks/triton2triton/campaign20/write_req_to_token_pool_triton/scripts/task_runner.py @@ -0,0 +1,127 @@ +#!/usr/bin/env python3 +"""Auto-generated task runner for triton_write_req_to_token_pool_triton (Triton). + +Inputs are generated each run from the shape/dtype signatures in +test_cases.json. The launcher symbol expected in source/triton_write_req_to_token_pool_triton.py is +``write_req_to_token_pool_triton`` (for raw @triton.jit kernels) or ``None`` (a +wrapper that handles grid + meta resolution if present). +""" +import sys, os, json, argparse, glob, importlib.util +sys.path.insert(0, os.path.dirname(os.path.abspath(__file__))) + +# >>> AKA-GENERATED: shared CUDA-graph benchmark helpers - edit src/tools/perf/vllm_cuda_graph_block.py then run `make sync-perf-helpers` >>> +def _measure_cuda_event_fallback(*args, **kwargs): + raise RuntimeError( + "CUDA-graph benchmark helpers were not materialized. " + "Run this task through AgentKernelArena so setup_workspace() can inject " + "src/tools/perf/vllm_cuda_graph_block.py into the workspace." + ) + + +def _benchmark_cuda_graph_or_events(*args, **kwargs): + raise RuntimeError( + "CUDA-graph benchmark helpers were not materialized. " + "Run this task through AgentKernelArena so setup_workspace() can inject " + "src/tools/perf/vllm_cuda_graph_block.py into the workspace." + ) +# <<< AKA-GENERATED <<< + + +TASK_DIR = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) +os.chdir(TASK_DIR) +TASK_NAME = "triton2triton/campaign20/write_req_to_token_pool_triton" +SOURCE_FILE = os.path.join(TASK_DIR, "source", "triton_write_req_to_token_pool_triton.py") +KERNEL_NAME = "write_req_to_token_pool_triton" +LAUNCHER_NAME = "None" +REF_SOURCE = "triton" +TEST_CASES = os.path.join(TASK_DIR, "test_cases.json") + + +def _load_module(): + spec = importlib.util.spec_from_file_location("triton_kernel", SOURCE_FILE) + mod = importlib.util.module_from_spec(spec) + spec.loader.exec_module(mod) + return mod + + +def _test_cases(): + if not os.path.isfile(TEST_CASES): + return [] + with open(TEST_CASES) as f: + return json.load(f) + + +def _resolve_callable(mod): + # Prefer a wrapping launcher if it was found alongside the @triton.jit kernel. + for n in (LAUNCHER_NAME, KERNEL_NAME): + if n and hasattr(mod, n) and n != "None": + return getattr(mod, n) + raise AttributeError(f"neither {LAUNCHER_NAME} nor {KERNEL_NAME} found in source") + + +def run_compile(): + try: + import ast + ast.parse(open(SOURCE_FILE).read()) + mod = _load_module() + _resolve_callable(mod) + _harness().compile_smoke() + return True, None + except Exception as e: + return False, str(e) + + +def _harness(): + import importlib.util + hpath = os.path.join(os.path.dirname(os.path.abspath(__file__)), "harness_run.py") + spec = importlib.util.spec_from_file_location("harness_run", hpath) + mod = importlib.util.module_from_spec(spec) + spec.loader.exec_module(mod) + return mod + + +def run_correctness(): + """Compile-only @triton.jit kernel: delegate to the real launcher harness, + which runs edited-vs-golden (frozen original) on identical seeded inputs and + requires EXACT match (integer/index kernel).""" + try: + return _harness().correctness() + except Exception as e: + return False, str(e) + + +def run_performance(): + try: + return _harness().performance(_benchmark_cuda_graph_or_events) + except Exception as e: + return [{"test_case_id": "error", "execution_time_ms": -1.0, "params": {"error": str(e)[:200]}}] + + +def main(): + ap = argparse.ArgumentParser() + ap.add_argument("mode", choices=["compile", "correctness", "performance"]) + args = ap.parse_args() + build_dir = os.path.join(TASK_DIR, "build") + os.makedirs(build_dir, exist_ok=True) + if args.mode == "compile": + ok, err = run_compile() + json.dump({"status": "ok" if ok else "fail", "error": err}, open(os.path.join(build_dir, "compile_report.json"), "w")) + print(f"Compilation: {'PASS' if ok else 'FAIL'}") + if err: print("Error:", err) + sys.exit(0 if ok else 1) + if args.mode == "correctness": + ok, err = run_correctness() + json.dump({"status": "ok" if ok else "fail", "error": err}, open(os.path.join(build_dir, "correctness_report.json"), "w")) + print(f"Correctness: {'PASS' if ok else 'FAIL'}") + if err: print("Error:", err) + sys.exit(0 if ok else 1) + cases = run_performance() + json.dump({"test_cases": cases}, open(os.path.join(build_dir, "performance_report.json"), "w"), indent=2) + for c in cases: + print(f"Performance: {c['execution_time_ms']:.4f} ms ({c['test_case_id']})") + valid = bool(cases) and all(c.get("execution_time_ms", -1) > 0 for c in cases) + sys.exit(0 if valid else 1) + + +if __name__ == "__main__": + main() diff --git a/tasks/triton2triton/campaign20/write_req_to_token_pool_triton/source/triton_write_req_to_token_pool_triton.py b/tasks/triton2triton/campaign20/write_req_to_token_pool_triton/source/triton_write_req_to_token_pool_triton.py new file mode 100644 index 00000000..ea0c3ed5 --- /dev/null +++ b/tasks/triton2triton/campaign20/write_req_to_token_pool_triton/source/triton_write_req_to_token_pool_triton.py @@ -0,0 +1,51 @@ +import triton +import triton.language as tl + +@triton.jit +def write_req_to_token_pool_triton( + req_to_token_ptr, # [max_batch, max_context_len] + req_pool_indices, + prefix_tensors, + pre_lens, + seq_lens, + extend_lens, + out_cache_loc, + req_to_token_ptr_stride: tl.constexpr, +): + BLOCK_SIZE: tl.constexpr = 512 + pid = tl.program_id(0) + + req_pool_index = tl.load(req_pool_indices + pid) + pre_len = tl.load(pre_lens + pid) + seq_len = tl.load(seq_lens + pid) + prefix_tensor = tl.load(prefix_tensors + pid).to(tl.pointer_type(tl.int64)) + + # write prefix + num_loop = tl.cdiv(pre_len, BLOCK_SIZE) + for i in range(num_loop): + offset = tl.arange(0, BLOCK_SIZE) + i * BLOCK_SIZE + mask = offset < pre_len + value = tl.load(prefix_tensor + offset, mask=mask) + tl.store( + req_to_token_ptr + req_pool_index * req_to_token_ptr_stride + offset, + value, + mask=mask, + ) + # NOTE: This can be slow for large bs + cumsum_start = tl.cast(0, tl.int64) + for i in range(pid): + cumsum_start += tl.load(extend_lens + i) + + num_loop = tl.cdiv(seq_len - pre_len, BLOCK_SIZE) + for i in range(num_loop): + offset = tl.arange(0, BLOCK_SIZE) + i * BLOCK_SIZE + mask = offset < (seq_len - pre_len) + value = tl.load(out_cache_loc + cumsum_start + offset, mask=mask) + tl.store( + req_to_token_ptr + + req_pool_index * req_to_token_ptr_stride + + offset + + pre_len, + value, + mask=mask, + ) diff --git a/tasks/triton2triton/campaign20/write_req_to_token_pool_triton/source_golden/triton_write_req_to_token_pool_triton.py b/tasks/triton2triton/campaign20/write_req_to_token_pool_triton/source_golden/triton_write_req_to_token_pool_triton.py new file mode 100644 index 00000000..ea0c3ed5 --- /dev/null +++ b/tasks/triton2triton/campaign20/write_req_to_token_pool_triton/source_golden/triton_write_req_to_token_pool_triton.py @@ -0,0 +1,51 @@ +import triton +import triton.language as tl + +@triton.jit +def write_req_to_token_pool_triton( + req_to_token_ptr, # [max_batch, max_context_len] + req_pool_indices, + prefix_tensors, + pre_lens, + seq_lens, + extend_lens, + out_cache_loc, + req_to_token_ptr_stride: tl.constexpr, +): + BLOCK_SIZE: tl.constexpr = 512 + pid = tl.program_id(0) + + req_pool_index = tl.load(req_pool_indices + pid) + pre_len = tl.load(pre_lens + pid) + seq_len = tl.load(seq_lens + pid) + prefix_tensor = tl.load(prefix_tensors + pid).to(tl.pointer_type(tl.int64)) + + # write prefix + num_loop = tl.cdiv(pre_len, BLOCK_SIZE) + for i in range(num_loop): + offset = tl.arange(0, BLOCK_SIZE) + i * BLOCK_SIZE + mask = offset < pre_len + value = tl.load(prefix_tensor + offset, mask=mask) + tl.store( + req_to_token_ptr + req_pool_index * req_to_token_ptr_stride + offset, + value, + mask=mask, + ) + # NOTE: This can be slow for large bs + cumsum_start = tl.cast(0, tl.int64) + for i in range(pid): + cumsum_start += tl.load(extend_lens + i) + + num_loop = tl.cdiv(seq_len - pre_len, BLOCK_SIZE) + for i in range(num_loop): + offset = tl.arange(0, BLOCK_SIZE) + i * BLOCK_SIZE + mask = offset < (seq_len - pre_len) + value = tl.load(out_cache_loc + cumsum_start + offset, mask=mask) + tl.store( + req_to_token_ptr + + req_pool_index * req_to_token_ptr_stride + + offset + + pre_len, + value, + mask=mask, + ) diff --git a/tasks/triton2triton/campaign20/write_req_to_token_pool_triton/test_cases.json b/tasks/triton2triton/campaign20/write_req_to_token_pool_triton/test_cases.json new file mode 100644 index 00000000..b13a4bc3 --- /dev/null +++ b/tasks/triton2triton/campaign20/write_req_to_token_pool_triton/test_cases.json @@ -0,0 +1,17 @@ +[ + { + "test_case_id": "c2", + "regime": {"B": 2, "seq_len": 1024, "pre_len": 16, "extend_len": 1008, "grid": [2]}, + "note": "token-parallel KV-write, prefill-style; grid=(B,); shapes driven by scripts/harness_run.py" + }, + { + "test_case_id": "c32", + "regime": {"B": 32, "seq_len": 1024, "pre_len": 16, "extend_len": 1008, "grid": [32]}, + "note": "token-parallel KV-write, prefill-style; grid=(B,); shapes driven by scripts/harness_run.py" + }, + { + "test_case_id": "c64", + "regime": {"B": 64, "seq_len": 1024, "pre_len": 16, "extend_len": 1008, "grid": [64]}, + "note": "token-parallel KV-write, prefill-style; grid=(B,); shapes driven by scripts/harness_run.py" + } +] diff --git a/tests/test_visualization.py b/tests/test_visualization.py index 6c7cccc8..963fbb81 100644 --- a/tests/test_visualization.py +++ b/tests/test_visualization.py @@ -49,6 +49,7 @@ def test_local_reports_are_default_and_workspace_reports_are_opt_in(self) -> Non local_report = reports_root / "manual_baseline" workspace_report = ( project_root + / "experiments" / "workspace_MI300_claude_code" / "run_20260715_120000" / "reports" @@ -59,6 +60,7 @@ def test_local_reports_are_default_and_workspace_reports_are_opt_in(self) -> Non with mock.patch.multiple( build_data, PROJECT_ROOT=project_root, + EXPERIMENTS_ROOT=project_root / "experiments", REPORTS_ROOT=reports_root, DATA_ROOT=data_root, OUTPUT_JSON=data_root / "data.json", @@ -92,7 +94,7 @@ def test_local_reports_are_default_and_workspace_reports_are_opt_in(self) -> Non for report in full_dataset["reports"] } self.assertIn( - "artifacts/workspace_MI300_claude_code/" + "artifacts/experiments/workspace_MI300_claude_code/" "run_20260715_120000/reports/overall_summary.csv", source_paths, ) @@ -131,8 +133,8 @@ def test_routes_static_generated_and_report_files(self) -> None: self.reports / "manual" / "overall_summary.csv", ) self.assertEqual( - self.resolve("/artifacts/workspace_x/run_x/reports/overall_report.txt"), - self.root / "workspace_x" / "run_x" / "reports" / "overall_report.txt", + self.resolve("/artifacts/experiments/workspace_x/run_x/reports/overall_report.txt"), + self.root / "experiments" / "workspace_x" / "run_x" / "reports" / "overall_report.txt", ) def test_rejects_traversal_hidden_paths_and_non_report_artifacts(self) -> None: