diff --git a/scripts/profile_dsv3_fake.sh b/scripts/profile_dsv3_fake.sh new file mode 100755 index 0000000000..82371206bf --- /dev/null +++ b/scripts/profile_dsv3_fake.sh @@ -0,0 +1,369 @@ +#!/bin/bash +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +set -euo pipefail + +REPO_ROOT="$(cd "$(dirname "${BASH_SOURCE[0]}")/.." && pwd)" +cd "$REPO_ROOT" + +# Select TRAINER=graph (default) or TRAINER=eager. +TRAINER="${TRAINER:-graph}" +NGPU="${NGPU:-256}" +TRAINING_STEPS="${TRAINING_STEPS:-10}" +LOCAL_BATCH_SIZE="${LOCAL_BATCH_SIZE:-24}" +SEQ_LEN="${SEQ_LEN:-4096}" +DATASET="${DATASET:-c4_test}" +DP_SHARD_DEGREE="${DP_SHARD_DEGREE:-256}" +EP_DEGREE="${EP_DEGREE:-64}" +PROFILE_FREQ="${PROFILE_FREQ:-10}" +PROFILER_WARMUP="${PROFILER_WARMUP:-3}" +PROFILER_ACTIVE="${PROFILER_ACTIVE:-1}" +RUN_ID="${RUN_ID:-$(date +%Y%m%d-%H%M%S)}" + +case "$TRAINER" in + graph) + MODULE="${MODULE:-graph_trainer.deepseek_v3}" + CONFIG="${CONFIG:-graph_trainer_deepseek_v3_671b}" + ENABLE_TLPARSE="${ENABLE_TLPARSE:-1}" + ;; + eager) + MODULE="${MODULE:-deepseek_v3}" + CONFIG="${CONFIG:-deepseek_v3_671b}" + ENABLE_TLPARSE="${ENABLE_TLPARSE:-0}" + ;; + *) + echo "TRAINER must be 'graph' or 'eager', got: $TRAINER" >&2 + exit 2 + ;; +esac + +TRACE_SUBDIR="${TRACE_SUBDIR:-profiling/dsv3_fake/$TRAINER/$RUN_ID}" +PROFILE_DIR="$REPO_ROOT/outputs/$TRACE_SUBDIR" +MEMORY_SUBDIR="${MEMORY_SUBDIR:-$TRACE_SUBDIR/memory_snapshot}" +MEMORY_DIR="$REPO_ROOT/outputs/$MEMORY_SUBDIR" +TORCH_TRACE_DIR="${TORCH_TRACE_DIR:-$PROFILE_DIR/torch_trace}" +TLPARSE_OUTPUT_DIR="${TLPARSE_OUTPUT_DIR:-$PROFILE_DIR/tlparse}" +LOG_FILE="${LOG_FILE:-$PROFILE_DIR/profile.log}" +PROFILER_UPLOAD_LOG="$PROFILE_DIR/profiler_upload.log" +MEMORY_UPLOAD_LOG="$PROFILE_DIR/memory_upload.log" +TLPARSE_MANIFOLD_LOG="$PROFILE_DIR/tlparse_manifold_upload.log" +TLPARSE_UPLOAD_LOG="$PROFILE_DIR/tlparse_artifact_uploads.log" +PERFETTO_UPLOADER="${PERFETTO_UPLOADER:-$REPO_ROOT/scripts/share_trace.py}" +UPLOAD_TRACE="${UPLOAD_TRACE:-1}" +UPLOAD_MEMORY_SNAPSHOT="${UPLOAD_MEMORY_SNAPSHOT:-$UPLOAD_TRACE}" +UPLOAD_TLPARSE="${UPLOAD_TLPARSE:-1}" +UPLOAD_LOG="${UPLOAD_LOG:-1}" + +if (( PROFILE_FREQ < PROFILER_WARMUP + PROFILER_ACTIVE )); then + echo "PROFILE_FREQ must be at least PROFILER_WARMUP + PROFILER_ACTIVE" >&2 + exit 2 +fi + +if (( TRAINING_STEPS < PROFILE_FREQ )); then + echo "TRAINING_STEPS must be at least PROFILE_FREQ to produce a trace" >&2 + exit 2 +fi + +if { [[ "$UPLOAD_TRACE" == "1" ]] || [[ "$UPLOAD_MEMORY_SNAPSHOT" == "1" ]]; } \ + && [[ ! -f "$PERFETTO_UPLOADER" ]]; then + echo "Trace uploader not found: $PERFETTO_UPLOADER" >&2 + exit 2 +fi + +if [[ "$ENABLE_TLPARSE" == "1" ]] && ! command -v tlparse >/dev/null 2>&1; then + echo "ENABLE_TLPARSE=1 requires tlparse on PATH" >&2 + exit 2 +fi + +mkdir -p "$PROFILE_DIR" "$MEMORY_DIR" "$(dirname "$LOG_FILE")" +if [[ "$ENABLE_TLPARSE" == "1" ]]; then + mkdir -p "$TORCH_TRACE_DIR" "$TLPARSE_OUTPUT_DIR" +fi + +upload_tlparse_passes() +( + set -u + local dir="${1:?Usage: upload_tlparse_passes }" + local artifact_dir="$dir/-_-_-_-" + local upload_status=0 + + if [[ ! -d "$artifact_dir" ]]; then + echo "No tlparse artifact directory found at $artifact_dir" >&2 + exit 2 + fi + + echo "=== Tlparse pass artifacts ===" + + local traced traced_name traced_output + local -a traced_files=("$artifact_dir"/make_fx_graph_traced_*.txt) + traced="${traced_files[0]}" + if [[ -f "$traced" ]]; then + traced_name="${traced##*/}" + if traced_output="$(pastry -t "${traced_name%.txt}" -l python -q \ + < "$traced" 2>/dev/null)"; then + echo "${traced_name%.txt}: $traced_output" + else + echo "Failed to upload $traced" >&2 + upload_status=1 + fi + fi + + local before_file basename pass_name after_file after_basename + local before_output after_output before_paste after_paste diff_url + local -a after_files + for before_file in "$artifact_dir"/before_*_pass_*.txt; do + [[ -f "$before_file" ]] || continue + basename="${before_file##*/}" + pass_name="${basename#before_}" + pass_name="${pass_name%_pass_*.txt}" + after_files=("$artifact_dir"/after_"${pass_name}"_pass_*.txt) + after_file="${after_files[0]}" + + if [[ ! -f "$after_file" ]]; then + echo "WARN: no after file for $pass_name, skipping" >&2 + continue + fi + + if diff -q "$before_file" "$after_file" >/dev/null 2>&1; then + echo "$pass_name: no changes, skipping" + continue + fi + + after_basename="${after_file##*/}" + if ! before_output="$(pastry -t "$basename" -l python -q \ + < "$before_file" 2>/dev/null)"; then + echo "Failed to upload $before_file" >&2 + upload_status=1 + continue + fi + if ! after_output="$(pastry -t "$after_basename" -l python -q \ + < "$after_file" 2>/dev/null)"; then + echo "Failed to upload $after_file" >&2 + upload_status=1 + continue + fi + + if [[ "$before_output" =~ P([0-9]+) ]]; then + before_paste="${BASH_REMATCH[1]}" + else + echo "Could not extract a paste number from: $before_output" >&2 + upload_status=1 + continue + fi + if [[ "$after_output" =~ P([0-9]+) ]]; then + after_paste="${BASH_REMATCH[1]}" + else + echo "Could not extract a paste number from: $after_output" >&2 + upload_status=1 + continue + fi + + diff_url="https://www.internalfb.com/intern/diffing/?before_paste_number=${before_paste}&after_paste_number=${after_paste}®ex_remove_pattern=&enable_regex_remove=0&strip_empty_lines=0&line_wrap=0&selected_tab=plain_diff" + echo "$pass_name: $diff_url" + done + + echo "=== Standalone tlparse artifacts ===" + local artifact artifact_name artifact_output + for artifact in \ + "$artifact_dir"/activation_memory_policy_*.txt \ + "$artifact_dir"/fx_codegen_*.txt \ + "$artifact_dir"/fx_collectives_analytical_estimation_*.txt \ + "$artifact_dir"/fx_compute_nodes_runtime_estimation_*.txt; do + [[ -f "$artifact" ]] || continue + artifact_name="${artifact##*/}" + if artifact_output="$(pastry -t "$artifact_name" -l python -q \ + < "$artifact" 2>/dev/null)"; then + echo "$artifact_name: $artifact_output" + else + echo "Failed to upload $artifact" >&2 + upload_status=1 + fi + done + + exit "$upload_status" +) + +train_args=( + --module "$MODULE" + --config "$CONFIG" + --comm.mode=fake_backend + --parallelism.data_parallel_shard_degree "$DP_SHARD_DEGREE" + --parallelism.expert_parallel_degree "$EP_DEGREE" + --compile.enable + --training.local_batch_size "$LOCAL_BATCH_SIZE" + --training.seq_len "$SEQ_LEN" + --training.steps "$TRAINING_STEPS" + --dataloader.dataset "$DATASET" + --debug.seed 42 + --debug.deterministic + --debug.moe_force_load_balance + --profiler.enable_profiling + --profiler.enable_memory_snapshot + --profiler.save_traces_folder "$TRACE_SUBDIR" + --profiler.save_memory_snapshot_folder "$MEMORY_SUBDIR" + --profiler.profile_freq "$PROFILE_FREQ" + --profiler.profiler_warmup "$PROFILER_WARMUP" + --profiler.profiler_active "$PROFILER_ACTIVE" + --profiler.profiler_repeat 1 +) + +if [[ "$TRAINER" == "graph" ]]; then + train_args+=( + --compile.mode aot_fx_trace + --compile.memory_policy full + --compile.debug_graph_passes + ) +fi +train_args+=("$@") +if [[ "$TRAINER" == "eager" ]]; then + train_args+=(activation-checkpoint:full) +fi + +run_env=( + NGPU="$NGPU" + LOCAL_RANK=0 + LOG_RANK=0 + PYTORCH_ALLOC_CONF="${PYTORCH_ALLOC_CONF:-expandable_segments:True}" +) +if [[ "$ENABLE_TLPARSE" == "1" ]]; then + run_env+=(TORCH_TRACE="$TORCH_TRACE_DIR") +fi + +set +e +{ + echo "Trainer: $TRAINER" + echo "Model: $MODULE / $CONFIG" + echo "Parallelism: FSDP=$DP_SHARD_DEGREE EP=$EP_DEGREE virtual GPUs=$NGPU" + echo "Batch/sequence: local batch=$LOCAL_BATCH_SIZE sequence=$SEQ_LEN" + echo "Dataset: $DATASET" + echo "Profile directory: $PROFILE_DIR" + echo "Training log: $LOG_FILE" + printf "Command:" + printf " %q" env "${run_env[@]}" python -m torchtitan.train "${train_args[@]}" + printf "\n" + + run_status=0 + env "${run_env[@]}" python -m torchtitan.train "${train_args[@]}" \ + || run_status=$? + training_status=$run_status + + if [[ "$ENABLE_TLPARSE" == "1" ]]; then + mapfile -t torch_trace_files < <( + rg --files "$TORCH_TRACE_DIR" -g '*rank_0*' | sort + ) + if (( ${#torch_trace_files[@]} == 0 )); then + echo "No rank-0 Torch trace found in $TORCH_TRACE_DIR" >&2 + (( training_status != 0 )) || run_status=2 + else + torch_trace_file="${torch_trace_files[0]}" + echo "Parsing Torch trace: $torch_trace_file" + + if [[ "$UPLOAD_TLPARSE" == "1" ]]; then + tlparse "$torch_trace_file" --overwrite-manifold \ + 2>&1 | tee "$TLPARSE_MANIFOLD_LOG" + tlparse_status=${PIPESTATUS[0]} + (( tlparse_status == 0 )) || run_status=$tlparse_status + fi + + tlparse parse "$torch_trace_file" -o "$TLPARSE_OUTPUT_DIR" --overwrite + tlparse_status=$? + (( tlparse_status == 0 )) || run_status=$tlparse_status + + if (( tlparse_status == 0 )) && [[ "$UPLOAD_TLPARSE" == "1" ]]; then + if command -v pastry >/dev/null 2>&1; then + upload_tlparse_passes "$TLPARSE_OUTPUT_DIR" \ + 2>&1 | tee "$TLPARSE_UPLOAD_LOG" + upload_status=${PIPESTATUS[0]} + (( upload_status == 0 )) || run_status=$upload_status + else + echo "pastry not found; skipping tlparse artifact uploads" >&2 + run_status=2 + fi + fi + fi + fi + + if (( training_status == 0 )) && [[ "$UPLOAD_TRACE" == "1" ]]; then + mapfile -t trace_files < <(rg --files "$PROFILE_DIR" -g 'rank0_*' | sort) + if (( ${#trace_files[@]} > 0 )); then + echo "Uploading ${trace_files[0]}" + python3 "$PERFETTO_UPLOADER" "${trace_files[0]}" \ + >"$PROFILER_UPLOAD_LOG" 2>&1 + upload_status=$? + cat "$PROFILER_UPLOAD_LOG" + (( upload_status == 0 )) || run_status=$upload_status + else + echo "No rank-0 profiler trace found in $PROFILE_DIR" >&2 + run_status=2 + fi + fi + + if (( training_status == 0 )) && [[ "$UPLOAD_MEMORY_SNAPSHOT" == "1" ]]; then + mapfile -t memory_files < <(rg --files "$MEMORY_DIR" -g '*.pickle' | sort) + if (( ${#memory_files[@]} > 0 )); then + echo "Uploading ${memory_files[0]}" + python3 "$PERFETTO_UPLOADER" --is-memory-snapshot \ + "${memory_files[0]}" >"$MEMORY_UPLOAD_LOG" 2>&1 + upload_status=$? + cat "$MEMORY_UPLOAD_LOG" + (( upload_status == 0 )) || run_status=$upload_status + else + echo "No memory snapshot found in $MEMORY_DIR" >&2 + run_status=2 + fi + fi + + exit "$run_status" +} 2>&1 | tee "$LOG_FILE" +run_status=${PIPESTATUS[0]} +set -e + +if [[ "$UPLOAD_LOG" == "1" ]]; then + if command -v pastry >/dev/null 2>&1; then + PASTRY_LINK="$(pastry < "$LOG_FILE")" + echo "Pastry link: $PASTRY_LINK" + else + echo "pastry not found; log remains at $LOG_FILE" >&2 + fi +fi + +profiler_link="" +memory_link="" +tlparse_link="" +if [[ -f "$PROFILER_UPLOAD_LOG" ]]; then + profiler_link="$(awk '/^Perfetto UI:$/ { getline; print; exit }' \ + "$PROFILER_UPLOAD_LOG")" +fi +if [[ -f "$MEMORY_UPLOAD_LOG" ]]; then + memory_link="$(awk '/^Memory snapshot:$/ { getline; print; exit }' \ + "$MEMORY_UPLOAD_LOG")" +fi +if [[ -f "$TLPARSE_MANIFOLD_LOG" ]]; then + tlparse_link="$(awk ' + { + for (i = 1; i <= NF; i++) { + if ($i ~ /^https:\/\//) { + gsub(/[),]$/, "", $i) + print $i + exit + } + } + } + ' "$TLPARSE_MANIFOLD_LOG")" +fi + +if [[ -n "$profiler_link" || -n "$memory_link" || -n "$tlparse_link" ]]; then + echo "Artifact links:" + [[ -z "$profiler_link" ]] || echo "Profiler trace: $profiler_link" + [[ -z "$memory_link" ]] || echo "Memory snapshot: $memory_link" + [[ -z "$tlparse_link" ]] || echo "Tlparse report: $tlparse_link" +fi +if [[ -f "$TLPARSE_UPLOAD_LOG" ]]; then + echo "Tlparse artifact links: $TLPARSE_UPLOAD_LOG" +fi + +exit "$run_status" diff --git a/scripts/share_trace.py b/scripts/share_trace.py new file mode 100755 index 0000000000..52d1dad903 --- /dev/null +++ b/scripts/share_trace.py @@ -0,0 +1,140 @@ +#!/usr/bin/env python3 +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +import argparse +import getpass +import logging +import os +import subprocess +import sys +import urllib.parse +import uuid + + +PERFETTO_OPEN_TRACE_URL = "https://www.internalfb.com/intern/perfetto/open_trace/" +PERFETTO_UI_ROOT_URL_META_INSIGHTS = "https://www.internalfb.com/intern/metainsights" +MEMORY_SNAPSHOT_ROOT_URL = "https://www.internalfb.com/pytorch_memory_visualizer" +MANIFOLD_BUCKET = "perfetto_internal_traces" +MANIFOLD_TRACE_DIR = "tree/shared_trace" +DEFAULT_TTL_SEC = 28 * 24 * 60 * 60 + + +def upload_trace_file(local_path: str, ttl_sec: int) -> str | None: + file_name = os.path.basename(local_path) + trace_path = "/".join( + [MANIFOLD_TRACE_DIR, f"{getpass.getuser()}_{uuid.uuid4()}_{file_name}"] + ) + manifold_path = f"{MANIFOLD_BUCKET}/{trace_path}" + result = subprocess.run( + [ + "manifold", + "put", + local_path, + manifold_path, + "--ttl", + str(ttl_sec), + "--userData", + "false", + ], + capture_output=True, + text=True, + ) + if result.returncode != 0: + logging.error("Upload failed:\n%s", result.stderr) + return None + + logging.info("Upload trace successfully.") + return trace_path + + +def get_perfetto_ui_url(trace_path: str, use_meta_insights: bool) -> str: + manifold_path = f"{MANIFOLD_BUCKET}/{trace_path}" + if use_meta_insights: + return ( + f"{PERFETTO_UI_ROOT_URL_META_INSIGHTS}#!/?url=" + "https://interncache-all.fbcdn.net/manifold/" + f"{urllib.parse.quote_plus(manifold_path)}" + ) + query = urllib.parse.urlencode({"manifold_path": manifold_path}) + return f"{PERFETTO_OPEN_TRACE_URL}?{query}" + + +def get_memory_snapshot_url(trace_path: str) -> str: + return f"{MEMORY_SNAPSHOT_ROOT_URL}/{MANIFOLD_BUCKET}/{trace_path}" + + +def parse_args(argv: list[str]) -> argparse.Namespace: + parser = argparse.ArgumentParser() + parser.add_argument("local_path", help="Trace file or directory to upload.") + parser.add_argument( + "-mi", + "--meta-insights", + action="store_true", + help="Open execution traces with Meta Insights.", + ) + parser.add_argument( + "--is-memory-snapshot", + action="store_true", + help="Open the uploaded file with the PyTorch memory visualizer.", + ) + parser.add_argument( + "-t", + "--ttl", + type=int, + default=DEFAULT_TTL_SEC, + help="Manifold object TTL in seconds.", + ) + return parser.parse_args(argv[1:]) + + +def get_upload_paths(local_path: str, is_memory_snapshot: bool) -> list[str]: + if not os.path.isdir(local_path): + return [local_path] + + suffix = ".pickle" if is_memory_snapshot else "trace.json" + return [ + os.path.join(local_path, filename) + for filename in sorted(os.listdir(local_path)) + if suffix in filename + ] + + +def main(argv: list[str]) -> int: + args = parse_args(argv) + logging.basicConfig(format="%(levelname)s: %(message)s", level=logging.INFO) + + if not os.path.exists(args.local_path): + logging.error("The trace path does not exist: %s", args.local_path) + return 1 + + paths = get_upload_paths(args.local_path, args.is_memory_snapshot) + if not paths: + logging.error("No uploadable files found in %s", args.local_path) + return 1 + + upload_failed = False + for path in paths: + logging.info("Uploading %s", path) + trace_path = upload_trace_file(path, args.ttl) + if trace_path is None: + upload_failed = True + continue + + print(f"Manifold path:\n{MANIFOLD_BUCKET}/{trace_path}") + if args.is_memory_snapshot: + print(f"Memory snapshot:\n{get_memory_snapshot_url(trace_path)}") + else: + print( + "Perfetto UI:\n" + f"{get_perfetto_ui_url(trace_path, args.meta_insights)}" + ) + + return int(upload_failed) + + +if __name__ == "__main__": + sys.exit(main(sys.argv)) diff --git a/tests/unit_tests/test_coda_fusion_microbench.py b/tests/unit_tests/test_coda_fusion_microbench.py new file mode 100644 index 0000000000..05fbe579ac --- /dev/null +++ b/tests/unit_tests/test_coda_fusion_microbench.py @@ -0,0 +1,142 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +import argparse +import unittest + +from torchtitan.experiments.graph_trainer.benchmarks.coda_fusion_autotune import ( + _full_sm100_configs, + _initial_configs, + _microbenchmark_command, + _search_configs, + INITIAL_CONFIGS, + INITIAL_CONFIGS_16B, + PRIORITY_CONFIGS, +) +from torchtitan.experiments.graph_trainer.benchmarks.coda_fusion_microbench import ( + _kernel_options, + CASES, +) +from torchtitan.experiments.graph_trainer.benchmarks.coda_fusion_microbench_16b import ( + CASES as CASES_16B, +) + + +class TestCodaFusionMicrobench(unittest.TestCase): + def test_case_inventory(self) -> None: + self.assertEqual(len(CASES), 12) + self.assertEqual( + {case.pattern for case in CASES.values()}, + { + "B1", + "B2", + "B4", + "B5", + "B6", + "B7", + "F2-Q", + "F2-KV", + "F3-A", + "F3-B", + "F4", + "F6", + }, + ) + + def test_tuning_and_fast_math_defaults(self) -> None: + case = CASES["f4_shared_expert_swiglu"] + options = _kernel_options( + case, + configs=(' {"tile_m": 128, "cluster_n": 1}',), + ) + self.assertEqual( + options, + ( + { + "backend": "QUACK", + "tuned": True, + "fast_math": True, + "config": {"tile_m": 128, "cluster_n": 1}, + }, + { + "backend": "QUACK", + "tuned": True, + "config": {"tile_m": 128, "cluster_n": 1}, + }, + ), + ) + + def test_16b_case_inventory(self) -> None: + self.assertEqual(len(CASES_16B), 13) + self.assertEqual( + {case.pattern for case in CASES_16B.values()}, + {"B1", "B2", "B4", "B5", "B6", "B7", "F2-KV", "F3-A", "F3-B", "F4"}, + ) + self.assertNotIn("f2_q_rmsnorm", CASES_16B) + self.assertNotIn("f6_router_sigmoid_bias", CASES_16B) + self.assertEqual( + CASES_16B["b1_lm_head_input_grad_cast"].shape, + "(2048, 102400) @ (102400, 2048); 8 chunks per step", + ) + + def test_every_case_always_uses_tuned_mode(self) -> None: + for case in CASES.values(): + with self.subTest(case=case.name): + options = _kernel_options(case, configs=()) + self.assertTrue(all(option["tuned"] for option in options)) + self.assertEqual( + tuple( + index + for index, option in enumerate(options) + if option.get("fast_math") + ), + case.fast_math_flex_gemms, + ) + + def test_config_count_must_match_flex_gemm_count(self) -> None: + case = CASES["f2_q_rmsnorm"] + with self.assertRaisesRegex(ValueError, "zero, one, or 2"): + _kernel_options( + case, + configs=("{}", "{}", "{}"), + ) + + def test_autotune_search_spaces(self) -> None: + self.assertEqual(len(PRIORITY_CONFIGS), 12) + self.assertEqual(len(_full_sm100_configs()), 74) + self.assertEqual(len(_search_configs("full", None)), 74) + for config in _full_sm100_configs(): + self.assertFalse(config["swap_ab"]) + self.assertFalse(config["use_tma_gather"]) + + def test_autotune_initial_configs_cover_every_case(self) -> None: + self.assertEqual(INITIAL_CONFIGS.keys(), CASES.keys()) + for name, case in CASES.items(): + args = argparse.Namespace(case=name, base_config=[]) + self.assertEqual(len(_initial_configs(args)), case.num_flex_gemms) + + self.assertEqual(INITIAL_CONFIGS_16B.keys(), CASES_16B.keys()) + for name, case in CASES_16B.items(): + args = argparse.Namespace( + suite="16b", + case=name, + base_config=[], + ) + self.assertEqual(len(_initial_configs(args)), case.num_flex_gemms) + + def test_autotune_uses_packaged_microbenchmark(self) -> None: + self.assertEqual( + _microbenchmark_command()[-1], + "torchtitan.experiments.graph_trainer.benchmarks.coda_fusion_microbench", + ) + self.assertEqual( + _microbenchmark_command("16b")[-1], + "torchtitan.experiments.graph_trainer.benchmarks.coda_fusion_microbench_16b", + ) + + +if __name__ == "__main__": + unittest.main() diff --git a/torchtitan/experiments/graph_trainer/CODA_FUSION_RESULTS.md b/torchtitan/experiments/graph_trainer/CODA_FUSION_RESULTS.md new file mode 100644 index 0000000000..bb6973c4e7 --- /dev/null +++ b/torchtitan/experiments/graph_trainer/CODA_FUSION_RESULTS.md @@ -0,0 +1,1553 @@ +# GraphTrainer CODA fusion results + +This file records the graph evidence and isolated performance result for each +CODA fusion pass. Full tlparse dumps stay under `outputs/` because they are +generated artifacts; the grounded samples and exact artifact paths are listed +here. + +## Standalone 12-pattern tuning suite + +The standalone suite is implemented in +`benchmarks/coda_fusion_microbench.py`. Every case has an explicit plain +PyTorch eager function and an explicit FlexGEMM function at a real shape from +the DSV3-671B joint graph. The primary baseline is the same eager function +compiled with Inductor. Source eager is also reported to keep the epilogue +implementation visible and independently measurable. + +The primary timing contract uses fixed-pointer CUDA graph replay, 25 warmups, +10 alternating candidate-order rounds, and 200 replays per round. F2-Q and +F2-KV used 100 replays per round because of their larger live input sets. +Compilation and tuning are outside the timed region. Every generated source is +checked for the expected number of `flex_gemm_epilogue` calls, and the selected +QuACK configurations are extracted from the generated code. All outputs pass +elementwise tolerance checks plus max-absolute, mean-absolute, and relative-L2 +reporting. + +Environment: one NVIDIA GB300 (SM 10.3), PyTorch +`2.14.0.dev20260811+cu130` at `50e2fa0ee83b`, CUTLASS DSL `4.6.2`. Results are +milliseconds; speedup is compiled eager divided by FlexGEMM. + +| Pattern | Source eager | Compiled eager | FlexGEMM | Speedup | Selected config(s) | +| --- | ---: | ---: | ---: | ---: | --- | +| B1 | 13.443 | 13.636 | 14.226 | 0.959x | `256x256 c2x1 dynamic` | +| B2 | 6.140 | 5.743 | 5.797 | 0.991x | `128x256 c2x1 dynamic`; `256x256 c2x2 dynamic` | +| B4 | 7.205 | 6.799 | 0.587 | 11.582x | `256x256 c2x2 dynamic` | +| B5 | 4.408 | 4.522 | 4.841 | 0.934x | `256x256 c2x1 dynamic` | +| B6 | 1.713 | 1.748 | 1.772 | 0.987x | `256x256 c2x1 dynamic` | +| B7 | 2.302 | 2.086 | 2.058 | 1.013x | `256x256 c2x1 dynamic` | +| F2-KV | 3.052 | 3.048 | 3.135 | 0.972x | `256x192 c2x1 dynamic`; `256x256 c2x1 dynamic` | +| F2-Q | 5.946 | 6.078 | 6.062 | 1.003x | `256x512 c2x1 static`; `256x256 c2x1 dynamic` | +| F3-A | 14.190 | 14.385 | 14.932 | 0.963x | `256x512 c2x1 static` | +| F3-B | 3.042 | 2.558 | 2.720 | 0.940x | `256x512 c2x1 static` | +| F4 | 3.725 | 3.730 | 3.598 | 1.037x | two `256x256 c2x1 dynamic` | +| F6 | 5.454 | 5.404 | 0.573 | 9.428x | `256x256 c2x1 dynamic` | + +Every FlexGEMM uses `tuned=True`. The F4 SiLU FlexGEMM and F6 sigmoid +FlexGEMM also use `fast_math=True`; the non-SiLU F4 FlexGEMM does not. B4's +BF16 result has `max_abs=4.883e-4` and `relative_l2=1.707e-3` versus source +eager. F6's two FP32 results have `max_abs=1.392e-5` and relative L2 below +`5e-6`. The large B4 and F6 speedups therefore satisfy the pattern-specific +tolerances but are not bitwise-equivalence claims. + +### Automatic tuning behavior + +FlexGEMM can tune without an explicit configuration: + +```python +kernel_options={"backend": "QUACK", "tuned": True} +``` + +With this form, the lowering obtains all device-compatible QuACK candidates, +filters them for the epilogue and local-reduction geometry, passes the resulting +template choices to Inductor's `autotune_select_algorithm()`, and caches the +winner. An explicit partial `config` constrains that candidate set while still +using the tuned path. + +The results above use `benchmarks/coda_fusion_autotune.py`, which supplies one +fully constrained configuration to each fresh child process and performs the +search outside Inductor. This is a reliability workaround, not a FlexGEMM API +requirement. On this build, unconstrained in-process tuning can execute an SM100 +candidate that raises `cudaErrorNoKernelImageForDevice` on SM103 or hangs in +the kernel. Inductor subprocess tuning cannot currently transport the resulting +TVM-FFI exception. Process isolation lets the suite time out or reject one bad +candidate without losing the entire search. + +The tuner searches 12 measured-priority configurations on every pattern, +remeasures per-GPU finalists sequentially on GPU 0, and performs a final full +timing run in another fresh process. Multi-FlexGEMM patterns use coordinate +descent. A full 74-configuration non-TMA, non-transposed SM100 search was also +run for B4 and F6. It retained the priority winners shown above, so the broader +space did not improve either decisive win. + +Artifacts: + +```text +outputs/coda_fusion_microbench/20260811_tournament_autotune +outputs/coda_fusion_microbench/20260811_full_autotune +``` + +Representative commands: + +```bash +python -m torchtitan.experiments.graph_trainer.benchmarks.coda_fusion_autotune \ + --case f4_shared_expert_swiglu --search priority --devices 0,1,2,3 + +python -m torchtitan.experiments.graph_trainer.benchmarks.coda_fusion_autotune \ + --case f6_router_sigmoid_bias --search full --devices 0,1,2,3 +``` + +## B1 LM-head input-gradient cast + +Pattern name: `b1_lm_head_input_grad_cast` + +The full B1 region contains eight chunked LM-head input-gradient GEMMs, FP32 +chunk writes and accumulation, a final BF16 conversion, and RMSNorm backward. +The supported FlexGEMM boundary is intentionally smaller: each +`BF16 mm -> reshape -> alias -> FP32 cast` chain is fused only when it is an +LM-head backward GEMM and the cast is the source of the corresponding chunk +`copy_`. The chunk writes, their cross-chunk accumulation, final conversion, +and RMSNorm backward remain unchanged. + +Real before sample: + +```python +mm_548 = torch.ops.aten.mm.default(view_4384, _unsafe_view_2737) +view_4385 = torch.ops.aten.reshape.default(mm_548, [24, 512, 7168]) +alias_436 = torch.ops.aten.alias.default(view_4385) +_to_copy_1737 = torch.ops.aten._to_copy.default( + alias_436, dtype=torch.float32 +) +copy_ = torch.ops.aten.copy_.default(slice_798, _to_copy_1737) +``` + +Real after sample: + +```python +_coda_b1_lm_head_input_grad_body_0 = ( + self._coda_b1_lm_head_input_grad_body_0 +) +flex_gemm = torch.ops.higher_order.flex_gemm( + torch.ops.aten.mm.default, + _coda_b1_lm_head_input_grad_body_0, + (view_4384, _unsafe_view_2737), + {}, + {"backend": "QUACK"}, +) +getitem_25338 = flex_gemm[0] +reshape_default = torch.ops.aten.reshape.default( + getitem_25338, [24, 512, 7168] +) +copy_ = torch.ops.aten.copy_.default(slice_798, reshape_default) +``` + +The FlexGEMM body explicitly models accumulator FP32 -> BF16 -> FP32 +conversion, preserving the original BF16 GEMM store rounding while producing +the FP32 chunk value. + +### Graph proof + +Configuration: DSV3-671B, fake communication backend, FSDP 256, EP 64, local +batch 24, sequence length 4096, full activation checkpointing, `c4_test`. + +Artifact root: + +```text +outputs/profiling/dsv3_fake/graph/coda-b1-proof-20260810/tlparse/-_-_-_- +``` + +Before and after dumps: + +```text +before_fuse_b1_lm_head_input_grad_cast_pass_274.txt +after_fuse_b1_lm_head_input_grad_cast_pass_275.txt +``` + +Pass diff: + +| Operation | Before | After | Delta | +| --- | ---: | ---: | ---: | +| root `mm.default` | 2,148 | 2,140 | -8 | +| root `_to_copy.default` | 2,486 | 2,478 | -8 | +| root `alias.default` | 2,825 | 2,817 | -8 | +| root `flex_gemm` | 0 | 8 | +8 | +| root `get_attr` | 427 | 435 | +8 | +| root `getitem` | 24,620 | 24,628 | +8 | + +All eight grounded chunk chains fused. Root reshape and `copy_` counts are +unchanged because each removed pre-cast reshape is replaced by a post-FlexGEMM +reshape. The proof run disabled regional Inductor and the CUDA graph pass so +tlparse could record the rewrite in isolation; execution subsequently OOMed in +unfused FlexAttention after all before/after artifacts had been written. + +### GB300 microbenchmark + +The exact fused boundary was benchmarked at the real chunk shape: +`(12,288, 129,280) @ (129,280, 7,168)`, BF16 inputs and FP32 output. Five +warmups and 20 CUDA-event iterations were run on one GB300. + +| Implementation | Median | Min | Max | +| --- | ---: | ---: | ---: | +| eager `mm` + cast | 13.265 ms | 12.952 ms | 13.502 ms | +| QUACK FlexGEMM | 15.874 ms | 15.226 ms | 17.246 ms | + +Speedup: `0.84x` (FlexGEMM is about 20% slower). This is retained because the +project accepts FlexGEMM fusions without a speedup. The result was bitwise +identical to eager (`max_abs_error=0`). The benchmark excludes the unchanged +chunk `copy_`. + +## B6 BF16 weight-gradient cast + +Pattern name: `b6_bf16_weight_grad_cast` + +The matcher was derived from the DSV3-671B post-bucketing graph. It requires an +`aten.mm.default` with BF16 output, exactly one user, and an +`aten._to_copy.default(dtype=torch.float32)` user. Multi-use GEMMs and the FP32 +router round trip are not matched. + +Real before sample: + +```python +t_702 = torch.ops.aten.t.default(view_4444) +mm_570 = torch.ops.aten.mm.default(t_702, view_4377_recomputed) +_to_copy_1775 = torch.ops.aten._to_copy.default( + mm_570, dtype=torch.float32 +) +``` + +Real after sample: + +```python +_coda_b6_body_8 = self._coda_b6_body_8 +flex_gemm_8 = torch.ops.higher_order.flex_gemm( + torch.ops.aten.mm.default, + _coda_b6_body_8, + (t_702, view_4377_recomputed), + {}, + {"backend": "QUACK"}, +) +getitem_25346 = flex_gemm_8[0] +``` + +The body explicitly models accumulator FP32 -> BF16 -> FP32 conversion. This +preserves the original BF16 GEMM store rounding while returning the FP32 value +expected by FSDP reduce-scatter bucketing. + +### Graph proof + +Configuration: DSV3-671B, fake communication backend, FSDP 256, EP 64, local +batch 24, sequence length 4096, full activation checkpointing, `c4_test`. + +Artifact root: + +```text +outputs/profiling/dsv3_fake/graph/coda-b6-bf16-proof-20260810-142424/tlparse/-_-_-_- +``` + +Before and after dumps: + +```text +before_fuse_b6_bf16_weight_grad_cast_pass_279.txt +after_fuse_b6_bf16_weight_grad_cast_pass_280.txt +``` + +Pass diff: + +| Operation | Before | After | Delta | +| --- | ---: | ---: | ---: | +| root `mm.default` | 2,148 | 1,652 | -496 | +| root `_to_copy.default` | 2,486 | 1,990 | -496 | +| root `flex_gemm` | 0 | 496 | +496 | +| root `get_attr` | 427 | 923 | +496 | +| root `getitem` | 24,620 | 25,116 | +496 | + +All 496 grounded BF16 chains fused. The 58 FP32 router chains remain in the +after graph. The proof run disabled regional Inductor so tlparse could record +the rewrite in isolation; execution subsequently OOMed in unfused +FlexAttention after all before/after artifacts had been written. + +### GB300 microbenchmark + +Representative shape: `(2048, 98304) @ (98304, 7168)`, BF16 inputs, FP32 +output. This is the most frequent B6 shape, with 116 occurrences in the graph. +Five warmups and 20 CUDA-event iterations were run on one GB300. + +| Implementation | Median | Min | Max | +| --- | ---: | ---: | ---: | +| eager `mm` + cast | 1.534 ms | 1.511 ms | 1.553 ms | +| QUACK FlexGEMM | 2.094 ms | 2.062 ms | 2.390 ms | + +Speedup: `0.73x` (FlexGEMM is 37% slower). This is retained because the project +accepts FlexGEMM fusions without a speedup. The output was bitwise identical to +eager (`max_abs_error=0`) and every FP32 output value was BF16-representable. + +### Dependencies and exclusions + +The pass requires PyTorch commit `bd2911838e0`, which preserves marked +FlexGEMM accumulator conversions through joint-graph cleanup and makes explicit +CuTeDSL casts use physical dtypes. The active environment also has the PyTorch +CI-pinned `nvidia-cutlass-dsl==4.5.2` and `apache-tvm-ffi==0.1.11` packages. + +The 58 FP32 router `mm -> BF16 -> FP32` chains are intentionally excluded. +QUACK compilation for the real `(256, 98304) @ (98304, 7168)` shape failed on +GB300 with `cudaErrorNoKernelImageForDevice`. No custom router kernel is added +without evidence that it is faster than the original implementation. + +## F6 router sigmoid and expert bias + +Pattern name: `f6_router_sigmoid_bias` + +The matcher follows the real FP32 router chain through its canonical reshape: +`mm -> reshape -> sigmoid`. The GEMM and reshape must each have one user, the +last reshape dimension must equal the GEMM output width, and an expert bias is +captured only when it is the exact one-dimensional output-width tensor. Other +sigmoid users are preserved. + +That last point is required by the training graph. The original forward score +is consumed both by the biased top-k path and by forced-load-balance routing; +the recomputed score is also consumed by sigmoid backward. FlexGEMM therefore +returns the raw sigmoid as its main output and, for the 58 forward cases, the +biased score as a same-shape auxiliary output. + +Real before sample: + +```python +mm_29 = torch.ops.aten.mm.default(view_109, t_29) +_unsafe_view_29 = torch.ops.aten.reshape.default(mm_29, [24, 4096, 256]) +sigmoid = torch.ops.aten.sigmoid.default(_unsafe_view_29) +add_7 = torch.ops.aten.add.Tensor(sigmoid, arg1938_1) +``` + +Real after sample: + +```python +reshape_default = torch.ops.aten.reshape.default(arg1938_1, [1, 256]) +_coda_f6_body_0 = self._coda_f6_body_0 +flex_gemm = torch.ops.higher_order.flex_gemm( + torch.ops.aten.mm.default, + _coda_f6_body_0, + (view_109, t_29, reshape_default), + {}, + {"backend": "QUACK"}, +) +raw_scores = torch.ops.aten.reshape.default(flex_gemm[0], [24, 4096, 256]) +biased_scores = torch.ops.aten.reshape.default(flex_gemm[1], [24, 4096, 256]) +``` + +### Graph proof + +Configuration: DSV3-671B, fake communication backend, FSDP 256, EP 64, local +batch 24, sequence length 4096, full activation checkpointing, `c4_test`. + +Artifact root: + +```text +outputs/profiling/dsv3_fake/graph/coda-f6-proof-20260810-144205/tlparse/-_-_-_- +``` + +Before and after dumps: + +```text +before_fuse_f6_router_sigmoid_bias_pass_274.txt +after_fuse_f6_router_sigmoid_bias_pass_275.txt +``` + +Pass diff for the root graph: + +| Operation | Before | After | Delta | +| --- | ---: | ---: | ---: | +| root `mm.default` | 2,148 | 2,032 | -116 | +| root `sigmoid.default` | 116 | 0 | -116 | +| root `add.Tensor` | 1,017 | 959 | -58 | +| root `flex_gemm` | 0 | 116 | +116 | +| root `get_attr` | 427 | 543 | +116 | +| root `getitem` | 24,620 | 24,794 | +174 | +| root `reshape.default` | 6,822 | 6,938 | +116 | + +All 58 original and 58 recomputed router sigmoid chains fused. The 58 original +expert-bias adds became auxiliary epilogue outputs. The proof run disabled +regional Inductor to retain the rewrite as FX text, then OOMed in unfused +FlexAttention after all pass artifacts had been written. + +### GB300 microbenchmark + +Representative shape: `(98304, 7168) @ (7168, 256)`, FP32 inputs and outputs, +with both raw and biased router scores returned. Five warmups and 20 CUDA-event +iterations were run on one GB300. + +| Implementation | Median | Min | Max | +| --- | ---: | ---: | ---: | +| eager `mm` + sigmoid + bias | 5.534 ms | 5.513 ms | 5.611 ms | +| QUACK FlexGEMM | 0.983 ms | 0.966 ms | 1.032 ms | + +Speedup: `5.63x`. For random FP32 inputs scaled by `0.02`, both raw and biased +outputs had `max_abs_error=1.395e-5` and `mean_abs_error=1.984e-6` versus eager. + +## F4 dense and shared-expert SwiGLU + +Pattern name: `f4_dense_swiglu` + +The matcher requires the grounded BF16 two-GEMM layout: + +```text +mm(W1) -> reshape -> silu --+ + mul +mm(W3) -> reshape ----------+ +``` + +Both GEMMs and their reshapes must be shape-compatible, and each GEMM must feed +only its reshape. Routed `_grouped_mm` SwiGLU is excluded and remains assigned +to distMoE. + +The first FlexGEMM emits SiLU. The second captures that same-shape tile and +emits both the BF16-rounded W3 result and the product. Recomputed W1 +preactivations have an additional `silu_backward` consumer, so those first +FlexGEMMs also return the rounded preactivation as an auxiliary output. The 61 +original forward cases do not create that unused auxiliary output. Explicit +FP32 -> BF16 conversions inside both bodies preserve the original GEMM store +rounding. + +Real before sample: + +```python +mm_5 = torch.ops.aten.mm.default(view_25, t_5) +_unsafe_view_5 = torch.ops.aten.reshape.default(mm_5, [24, 4096, 18432]) +silu = torch.ops.aten.silu.default(_unsafe_view_5) +mm_6 = torch.ops.aten.mm.default(view_27, t_6) +_unsafe_view_6 = torch.ops.aten.reshape.default(mm_6, [24, 4096, 18432]) +mul_2 = torch.ops.aten.mul.Tensor(silu, _unsafe_view_6) +``` + +Real after sample: + +```python +flex_gemm = torch.ops.higher_order.flex_gemm( + torch.ops.aten.mm.default, + _coda_f4_silu_body_0, + (view_25, t_5), + {}, + {"backend": "QUACK"}, +) +silu_2d = flex_gemm[0] +flex_gemm_1 = torch.ops.higher_order.flex_gemm( + torch.ops.aten.mm.default, + _coda_f4_mul_body_0, + (view_27, t_6, silu_2d), + {}, + {"backend": "QUACK"}, +) +gate_2d = flex_gemm_1[0] +product_2d = flex_gemm_1[1] +``` + +The first recomputed case has two outputs instead: `flex_gemm_122[0]` is SiLU +and `flex_gemm_122[1]` is the saved W1 preactivation. + +### Graph proof + +Configuration: DSV3-671B, fake communication backend, FSDP 256, EP 64, local +batch 24, sequence length 4096, full activation checkpointing, `c4_test`. + +Artifact root: + +```text +outputs/profiling/dsv3_fake/graph/coda-f4-full-proof-20260810-150132/tlparse/-_-_-_- +``` + +Before and after dumps: + +```text +before_fuse_f4_dense_swiglu_pass_274.txt +after_fuse_f4_dense_swiglu_pass_275.txt +``` + +Pass diff for the root graph: + +| Operation | Before | After | Delta | +| --- | ---: | ---: | ---: | +| root `mm.default` | 2,148 | 1,904 | -244 | +| root `silu.default` | 238 | 116 | -122 | +| root `mul.Tensor` | 1,249 | 1,127 | -122 | +| root `flex_gemm` | 0 | 244 | +244 | +| root `get_attr` | 427 | 671 | +244 | +| root `getitem` | 24,620 | 25,047 | +427 | +| root `reshape.default` | 6,822 | 7,005 | +183 | + +All 6 dense and 116 shared-expert original/recomputed chains fused. The 116 +routed grouped-GEMM chains remain unchanged. The proof run disabled regional +Inductor to retain the rewrite as FX text, then OOMed in unfused FlexAttention +after all pass artifacts had been written. + +### GB300 microbenchmark + +Representative shared-expert shape: two `(98304, 7168) @ (7168, 2048)` BF16 +GEMMs, returning SiLU, the raw W3 gate, and their product. Five warmups and 20 +CUDA-event iterations were run on one GB300. + +| Implementation | Median | Min | Max | +| --- | ---: | ---: | ---: | +| eager two GEMMs + SiLU + multiply | 3.613 ms | 3.332 ms | 3.739 ms | +| two QUACK FlexGEMMs | 3.711 ms | 3.619 ms | 4.069 ms | + +Speedup: `0.97x` (FlexGEMM is 2.7% slower). This is retained under the project +policy that accepts FlexGEMM fusions without a speedup. SiLU, raw gate, and +product were bitwise identical to eager on the full representative shape. + +## B2 dense and shared-expert SwiGLU backward + +Pattern name: `b2_dense_swiglu_backward` + +The matcher covers the two grounded BF16 GEMM epilogues in a dense or shared +SwiGLU backward block. The W2 input-gradient GEMM feeds two derivatives: + +```text +mm(W2 input gradient) -> reshape --+-> mul(saved SiLU) + +-> mul(saved W3) -> silu_backward(saved W1) +``` + +The first FlexGEMM captures the three saved activations and returns both +derivatives. A second FlexGEMM captures the already-computed W3 input-gradient +GEMM and folds its add into the W1 input-gradient GEMM. The second match also +requires sibling `w3` and `w1` module FQNs from the same feed-forward block, +with the captured W3 GEMM preceding W1 in the graph. Routed grouped GEMMs are +excluded and remain assigned to distMoE. + +Real before sample from layer 60: + +```python +mm_571 = torch.ops.aten.mm.default(view_4444, _unsafe_view_2734) +view_4445 = torch.ops.aten.reshape.default(mm_571, [24, 4096, 2048]) +mul_357 = torch.ops.aten.mul.Tensor(view_4445, silu_118_recomputed) +mul_358 = torch.ops.aten.mul.Tensor(view_4445, _unsafe_view_660_recomputed) +silu_backward = torch.ops.aten.silu_backward.default( + mul_358, _unsafe_view_659_recomputed +) + +mm_573 = torch.ops.aten.mm.default(view_4447, _unsafe_view_2735) +view_4448 = torch.ops.aten.reshape.default(mm_573, [24, 4096, 7168]) +mm_575 = torch.ops.aten.mm.default(view_4450, _unsafe_view_2733) +view_4451 = torch.ops.aten.reshape.default(mm_575, [24, 4096, 7168]) +add_362 = torch.ops.aten.add.Tensor(view_4448, view_4451) +``` + +Real after sample: + +```python +flex_gemm = torch.ops.higher_order.flex_gemm( + torch.ops.aten.mm.default, + _coda_b2_branch_body_0, + ( + view_4444, + _unsafe_view_2734, + reshape_default, + reshape_default_1, + reshape_default_2, + ), + {}, + {"backend": "QUACK"}, +) +gate_grad = flex_gemm[0] +silu_grad = flex_gemm[1] + +flex_gemm_61 = torch.ops.higher_order.flex_gemm( + torch.ops.aten.mm.default, + _coda_b2_input_add_body_0, + (view_4450, _unsafe_view_2733, mm_573), + {}, + {"backend": "QUACK"}, +) +input_grad = flex_gemm_61[0] +``` + +Both bodies explicitly model the original BF16 GEMM store rounding with an +FP32 -> BF16 conversion before evaluating their epilogues. + +### Graph proof + +Configuration: DSV3-671B, fake communication backend, FSDP 256, EP 64, local +batch 24, sequence length 4096, full activation checkpointing, `c4_test`. + +Artifact root: + +```text +outputs/profiling/dsv3_fake/graph/coda-b2-proof-20260810-continued/tlparse/-_-_-_- +``` + +Before and after dumps: + +```text +before_fuse_b2_dense_swiglu_backward_pass_274.txt +after_fuse_b2_dense_swiglu_backward_pass_275.txt +``` + +Pass diff for the root graph: + +| Operation | Before | After | Delta | +| --- | ---: | ---: | ---: | +| root `mm.default` | 2,148 | 2,026 | -122 | +| root `mul.Tensor` | 1,249 | 1,127 | -122 | +| root `silu_backward.default` | 119 | 58 | -61 | +| root `add.Tensor` | 1,017 | 956 | -61 | +| root `flex_gemm` | 0 | 122 | +122 | +| root `get_attr` | 427 | 549 | +122 | +| root `getitem` | 24,620 | 24,803 | +183 | +| root `reshape.default` | 6,822 | 7,005 | +183 | + +All 61 dense/shared branch-derivative chains and all 61 matching +input-gradient adds fused. The proof run disabled regional Inductor to retain +the rewrite as FX text, then OOMed on the known 192 GiB unfused FlexAttention +allocation after all pass artifacts had been written. + +### GB300 microbenchmark + +The full representative shared-expert backward block used BF16 tensors with +`M=98,304`, model width `7,168`, and shared-expert width `2,048`. It includes +the W2 input-gradient GEMM, both derivative branches, both W3/W1 input-gradient +GEMMs, and their final add. Five warmups and 20 CUDA-event iterations were run +on one GB300. + +| Implementation | Median | Min | Max | +| --- | ---: | ---: | ---: | +| eager three GEMMs + derivative epilogues | 6.046 ms | 5.663 ms | 6.277 ms | +| two QUACK FlexGEMMs + one eager GEMM | 6.955 ms | 6.769 ms | 7.899 ms | + +Speedup: `0.87x` (FlexGEMM is 15% slower). This is retained under the project +policy that accepts FlexGEMM fusions without a speedup. The gate derivative was +bitwise identical to eager. The SiLU derivative and final input gradient had +maximum absolute errors of `3.052e-5` and `1.526e-5`, respectively, with random +BF16 inputs scaled by `0.02`. + +## F2 MLA Q projection and RMSNorm + +Pattern name: `f2_q_rmsnorm` + +This is an opt-in, numerics-changing CODA reparameterization of the grounded +MLA Q chain: + +```text +mm(wq_a) -> RMSNorm(q_norm) -> mm(wq_b) +``` + +The first FlexGEMM applies the 1,536-element norm weight without the row scale +and emits three 512-column partial mean-square values per token. Root-graph +pointwise ops finalize one inverse-RMS value per token. The second FlexGEMM +captures that row value and applies it after the `wq_b` accumulation. + +The pass rewrites both the 61 original and 61 rematerialized Q chains. The +rematerialized first FlexGEMM additionally returns raw BF16 Q. The pass +reconstructs normalized Q and reshapes `rstd` for the existing `wq_b` weight +gradient and RMSNorm backward consumers, so activation checkpointing does not +mix CODA forward values with an unrelated recomputation path. + +Real before sample: + +```python +mm = torch.ops.aten.mm.default(view_4, t) +_unsafe_view = torch.ops.aten.reshape.default(mm, [24, 4096, 1536]) +_fused_rms_norm_1 = torch.ops.aten._fused_rms_norm.default( + _unsafe_view, [1536], _unsafe_view_1273, 1e-05 +) +getitem_2 = _fused_rms_norm_1[0] +view_7 = torch.ops.aten.reshape.default(getitem_2, [98304, 1536]) +mm_1 = torch.ops.aten.mm.default(view_7, t_1) +``` + +Real original-forward after sample: + +```python +flex_gemm = torch.ops.higher_order.flex_gemm( + torch.ops.aten.mm.default, + _coda_f2_q_first_body_0, + (view_4, t, reshape_default), + {}, + {"backend": "QUACK"}, +) +weighted_q = flex_gemm[0] +partial_mean_square = flex_gemm[1] +mean_square = torch.ops.aten.mean.dim(partial_mean_square, [-1], True) +rstd = torch.ops.aten.rsqrt.default( + torch.ops.aten.add.Scalar(mean_square, 1e-05) +) +flex_gemm_1 = torch.ops.higher_order.flex_gemm( + torch.ops.aten.mm.default, + _coda_f2_q_second_body_0, + (weighted_q, t_1, rstd), + {}, + {"backend": "QUACK"}, +) +q = flex_gemm_1[0] +``` + +The rematerialized first HOP has a third output: raw BF16 Q is output 1 and +the partial statistics move to output 2. Generated-kernel validation confirmed +that QUACK supports this combination of a same-shape auxiliary output and a +compressed physical reduction. + +### Graph proof + +Configuration: DSV3-671B, fake communication backend, FSDP 256, EP 64, local +batch 24, sequence length 4096, full activation checkpointing, `c4_test`. + +Artifact root: + +```text +outputs/profiling/dsv3_fake/graph/coda-f2-q-full-proof-20260810-2/tlparse/-_-_-_- +``` + +Before and after dumps: + +```text +before_fuse_f2_q_rmsnorm_pass_274.txt +after_fuse_f2_q_rmsnorm_pass_275.txt +``` + +Pass diff for the root graph: + +| Operation | Before | After | Delta | +| --- | ---: | ---: | ---: | +| root `mm.default` | 2,148 | 1,904 | -244 | +| root `_fused_rms_norm.default` | 490 | 368 | -122 | +| root `flex_gemm` | 0 | 244 | +244 | +| root `get_attr` | 427 | 671 | +244 | +| root `getitem` | 24,620 | 24,864 | +244 | +| root `mean.dim` | 0 | 122 | +122 | +| root `add.Scalar` | 0 | 122 | +122 | +| root `rsqrt.default` | 0 | 122 | +122 | +| root `_to_copy.default` | 2,486 | 2,608 | +122 | +| root `mul.Tensor` | 1,249 | 1,371 | +122 | + +All 61 original and 61 rematerialized Q chains fused. The remaining 368 norms +include KV, residual/FFN, final, and their rematerialized copies. The proof run +disabled regional Inductor to retain the rewrite as FX text, then OOMed on the +known 192 GiB unfused FlexAttention allocation after all pass artifacts had +been written. + +### GB300 microbenchmarks + +The representative Q chain uses BF16 tensors with `M=98,304`, input width +`7,168`, Q low-rank width `1,536`, and projected width `24,576`. Five warmups +and 20 CUDA-event iterations were run on one GB300. + +| Variant | Eager | CODA FlexGEMM | Speedup | +| --- | ---: | ---: | ---: | +| original forward | 5.816 ms | 6.020 ms | 0.97x | +| rematerialized with saved values | 5.919 ms | 6.294 ms | 0.94x | + +Both variants are retained under the project policy that accepts FlexGEMM +fusions without a speedup. On the full original-forward shape, projected Q had +`max_abs_error=0.03125` and `mean_abs_error=0.001034`. In the rematerialized +variant, raw Q was bitwise exact, saved `rstd` had `max_abs_error=5.722e-6`, +and reconstructed normalized Q had `max_abs_error=0.015625` and +`mean_abs_error=3.656e-8`. + +This F2-Q pass requires convergence validation because moving the row scale +across BF16 GEMMs changes rounding. The separate F2-KV pass below handles the +512-of-576 KV chain without moving RoPE into the GEMM epilogue. + +## F2 MLA KV projection and segmented RMSNorm + +Pattern name: `f2_kv_rmsnorm` + +This pass rewrites the grounded segmented MLA KV chain: + +```text +mm(wkv_a) -> split [512, 64] -> RMSNorm(kv_norm) -> mm(wkv_b) + -> 64-column RoPE tail +``` + +The first FlexGEMM retains its physical 576-column output. The 512-element norm +weight is padded with 64 ones, so the epilogue emits a gamma-weighted full-width +value, a raw full-width auxiliary, and nine 64-column mean-square partials. The +root graph uses the first eight partials to form `rstd`; the ninth corresponds +to the unnormalized RoPE tail. The second FlexGEMM consumes the first 512 +weighted columns and applies `rstd` after `wkv_b` accumulation. The raw +auxiliary continues through the original split, preserving the RoPE tail and +the rematerialized RMSNorm-backward input. + +Real before sample: + +```python +mm_2 = torch.ops.aten.mm.default(view_10, t_2) +kv = torch.ops.aten.reshape.default(mm_2, [24, 4096, 576]) +kv, k_pe = torch.ops.aten.split_with_sizes.default(kv, [512, 64], -1) +norm = torch.ops.aten._fused_rms_norm.default( + kv, [512], kv_norm_weight, 1e-05 +) +mm_3 = torch.ops.aten.mm.default(norm[0].reshape(98304, 512), t_3) +``` + +Real after sample: + +```python +gamma = torch.ops.aten.reshape.default(kv_norm_weight, [1, 512]) +gamma_full = torch.ops.aten.constant_pad_nd.default(gamma, [0, 64], 1.0) +first = torch.ops.higher_order.flex_gemm( + torch.ops.aten.mm.default, + _coda_f2_kv_first_body_0, + (view_10, t_2, gamma_full), + {}, + {"backend": "QUACK"}, +) +weighted_full, raw_full, partials = first +weighted = torch.ops.aten.slice.Tensor(weighted_full, 1, 0, 512) +active_partials = torch.ops.aten.slice.Tensor(partials, 1, 0, 8) +rstd = torch.ops.aten.rsqrt.default( + torch.ops.aten.add.Scalar( + torch.ops.aten.mean.dim(active_partials, [-1], True), 1e-05 + ) +) +second = torch.ops.higher_order.flex_gemm( + torch.ops.aten.mm.default, + _coda_f2_kv_second_body_0, + (weighted, t_3, rstd), + {}, + {"backend": "QUACK"}, +) +``` + +### Graph proof + +Configuration: DSV3-671B, fake communication backend, FSDP 256, EP 64, local +batch 24, sequence length 4096, full activation checkpointing, `c4_test`. + +Artifact root: + +```text +outputs/profiling/dsv3_fake/graph/coda-f2-kv-proof-20260810/tlparse/-_-_-_- +``` + +Before and after dumps: + +```text +before_fuse_f2_kv_rmsnorm_pass_274.txt +after_fuse_f2_kv_rmsnorm_pass_275.txt +``` + +Pass diff for the root graph: + +| Operation | Before | After | Delta | +| --- | ---: | ---: | ---: | +| root `mm.default` | 2,148 | 1,904 | -244 | +| root `_fused_rms_norm.default` | 490 | 368 | -122 | +| root `flex_gemm` | 0 | 244 | +244 | +| root `get_attr` | 427 | 671 | +244 | +| root `getitem` | 24,620 | 24,864 | +244 | +| root `mean.dim` | 0 | 122 | +122 | +| root `add.Scalar` | 0 | 122 | +122 | +| root `rsqrt.default` | 0 | 122 | +122 | +| root `_to_copy.default` | 2,486 | 2,608 | +122 | +| root `constant_pad_nd.default` | 3,904 | 4,026 | +122 | +| root `mul.Tensor` | 1,249 | 1,371 | +122 | +| root `reshape.default` | 6,822 | 6,883 | +61 | +| root `slice.Tensor` | 615 | 920 | +305 | + +All 61 original and 61 rematerialized KV chains fused. The raw 576-column +auxiliary preserves all 122 existing split/RoPE paths. The proof run disabled +regional Inductor to retain the rewrite as FX text, then OOMed on the known +192 GiB unfused FlexAttention allocation after all artifacts were written. + +### GB300 microbenchmark + +The representative chain uses BF16 tensors with `M=98,304`, input width +`7,168`, WKV-A width `576`, active RMSNorm width `512`, and WKV-B output width +`32,768`. Five warmups and 20 CUDA-event iterations were run on one GB300. + +| Implementation | Median | Min | Max | +| --- | ---: | ---: | ---: | +| eager WKV-A + segmented RMSNorm + WKV-B | 2.816 ms | 2.654 ms | 3.602 ms | +| two QUACK FlexGEMMs + partial reduction | 3.516 ms | 3.435 ms | 4.559 ms | + +Speedup: `0.80x`. This is retained under the project policy that accepts +FlexGEMM fusions without a speedup. The raw RoPE tail was bitwise exact. The +WKV-B output had `max_abs_error=0.015625` and `mean_abs_error=0.000596412`. +Moving `rstd` across WKV-B requires convergence validation. This result does +not justify a custom kernel. + +## B4 router input-gradient cast and add + +Pattern name: `b4_router_input_grad_add` + +This pattern fuses the router input-gradient GEMM's FP32 -> BF16 store rounding +and the addition of the expert-path input gradient. The matcher requires the +GEMM's module FQN to end in `.moe.router.gate`, an FP32 two-dimensional GEMM, +a sole reshape and BF16 cast chain, and a BF16 residual with the same final +shape. These constraints exclude unrelated FP32 linears and multi-use values. + +Real before sample from layer 60: + +```python +view_4467 = torch.ops.aten.reshape.default(sigmoid_backward, [98304, 256]) +mm_577 = torch.ops.aten.mm.default(view_4467, _to_copy_1718_recomputed) +view_4468 = torch.ops.aten.reshape.default(mm_577, [24, 4096, 7168]) +_to_copy_1784 = torch.ops.aten._to_copy.default( + view_4468, + dtype=torch.bfloat16, + layout=torch.strided, + device=torch.device("cuda:0"), +) +add_366 = torch.ops.aten.add.Tensor(add_364, _to_copy_1784) +``` + +Real after sample: + +```python +reshape_default = torch.ops.aten.reshape.default(add_364, [98304, 7168]) +flex_gemm = torch.ops.higher_order.flex_gemm( + torch.ops.aten.mm.default, + _coda_b4_router_body_0, + (view_4467, _to_copy_1718_recomputed, reshape_default), + {}, + {"backend": "QUACK"}, +) +router_input_grad = flex_gemm[0] +reshape_default_1 = torch.ops.aten.reshape.default( + router_input_grad, [24, 4096, 7168] +) +``` + +The FlexGEMM body performs the original BF16 conversion before the add, so the +expert-path gradient is not added to an unrounded FP32 accumulator. +`sigmoid_backward` feeds the GEMM and is therefore a prologue, not an epilogue; +the following `_fused_rms_norm_backward` consumes the fused result but is not +part of this fusion. + +### Graph proof + +Configuration: DSV3-671B, fake communication backend, FSDP 256, EP 64, local +batch 24, sequence length 4096, full activation checkpointing, `c4_test`. + +Artifact root: + +```text +outputs/profiling/dsv3_fake/graph/coda-b4-proof-20260810/tlparse/-_-_-_- +``` + +Before and after dumps: + +```text +before_fuse_b4_router_input_grad_add_pass_274.txt +after_fuse_b4_router_input_grad_add_pass_275.txt +``` + +Pass diff for the root graph: + +| Operation | Before | After | Delta | +| --- | ---: | ---: | ---: | +| root `mm.default` | 2,148 | 2,090 | -58 | +| root `_to_copy.default` | 2,486 | 2,428 | -58 | +| root `add.Tensor` | 1,017 | 959 | -58 | +| root `flex_gemm` | 0 | 58 | +58 | +| root `get_attr` | 427 | 485 | +58 | +| root `getitem` | 24,620 | 24,678 | +58 | +| root `reshape.default` | 6,822 | 6,880 | +58 | + +All 58 router input-gradient chains fused. The proof run disabled regional +Inductor to retain the rewrite as FX text, then OOMed on the known 192 GiB +unfused FlexAttention allocation after all pass artifacts had been written. + +### GB300 microbenchmark + +The representative full shape is FP32 `(98,304, 256) @ (256, 7,168)` followed +by a BF16 conversion and addition to a BF16 `(98,304, 7,168)` residual. Inputs +were scaled by `0.02`; five warmups and 20 CUDA-event iterations were run on one +GB300. + +| Implementation | Median | Min | Max | +| --- | ---: | ---: | ---: | +| eager GEMM + BF16 cast + add | 7.276 ms | 7.266 ms | 7.300 ms | +| QUACK FlexGEMM | 0.784 ms | 0.761 ms | 0.865 ms | + +Speedup: `9.29x`. The fused output had `max_abs_error=0.00048828125` and +`mean_abs_error=1.494e-6` relative to eager. + +## B7 Q/KV attention input-gradient merge + +Pattern name: `b7_attention_grad_merge` + +The attention backward graph computes separate BF16 input gradients for +`wkv_a` and `wq_a`, reshapes both to the residual-stream shape, and adds them +before attention RMSNorm backward. The matcher requires both branches to be +backward `mm` nodes from the same transformer layer, with exact +`attention.wkv_a` and `attention.wq_a` module annotations, one reshape user per +GEMM, and one common BF16 add. This prevents unrelated same-shape gradient +adds from matching. + +The KV GEMM remains independent. The Q FlexGEMM captures its 2D BF16 output, +rounds the Q accumulator to BF16 exactly where the original GEMM stored it, +and performs the add in the original operand order. Its root-graph output is +2D and the original 3D reshape remains after the HOP. The following +`_fused_rms_norm_backward` is not part of this fusion. + +Real before sample: + +```python +mm_583 = torch.ops.aten.mm.default(view_4484, slice_787_recomputed) +view_4485 = torch.ops.aten.reshape.default(mm_583, [24, 4096, 7168]) +mm_587 = torch.ops.aten.mm.default(view_4492, _unsafe_view_2724) +view_4493 = torch.ops.aten.reshape.default(mm_587, [24, 4096, 7168]) +add_368 = torch.ops.aten.add.Tensor(view_4485, view_4493) +``` + +Real after sample: + +```python +flex_gemm = torch.ops.higher_order.flex_gemm( + torch.ops.aten.mm.default, + _coda_b7_attention_grad_merge_body_0, + (view_4492, _unsafe_view_2724, mm_583), + {}, + {"backend": "QUACK"}, +) +getitem_25338 = flex_gemm[0] # bf16[98304, 7168] +reshape_default = torch.ops.aten.reshape.default( + getitem_25338, [24, 4096, 7168] +) +``` + +### Graph proof + +Configuration: DSV3-671B, fake communication backend, FSDP 256, EP 64, local +batch 24, sequence length 4096, full activation checkpointing, `c4_test`. + +Artifact root: + +```text +outputs/profiling/dsv3_fake/graph/coda-b7-proof-final-20260810/tlparse/-_-_-_- +``` + +Before and after dumps: + +```text +before_fuse_b7_attention_grad_merge_pass_274.txt +after_fuse_b7_attention_grad_merge_pass_275.txt +``` + +Pass diff for the root graph: + +| Operation | Before | After | Delta | +| --- | ---: | ---: | ---: | +| root `mm.default` | 2,148 | 2,087 | -61 | +| root `add.Tensor` | 1,017 | 956 | -61 | +| root `reshape.default` | 6,822 | 6,761 | -61 | +| root `flex_gemm` | 0 | 61 | +61 | +| root `get_attr` | 427 | 488 | +61 | +| root `getitem` | 24,620 | 24,681 | +61 | + +All 61 transformer-layer Q/KV input-gradient merges fused. The proof run +disabled regional Inductor to retain the rewrite as FX text, then OOMed on the +known 192 GiB unfused FlexAttention allocation after all pass artifacts had +been written. + +### GB300 microbenchmark + +The representative full shape uses a KV GEMM `(98,304, 576) @ (576, 7,168)` +and a Q GEMM `(98,304, 1,536) @ (1,536, 7,168)`, followed by their BF16 add. +Inputs were scaled by `0.02`; five warmups and 20 CUDA-event iterations were +run on one GB300. + +| Implementation | Median | Min | Max | +| --- | ---: | ---: | ---: | +| eager two GEMMs + add | 2.195 ms | 2.149 ms | 2.285 ms | +| KV GEMM + Q FlexGEMM add epilogue | 2.115 ms | 1.999 ms | 2.156 ms | + +Speedup: `1.04x`. The fused output was bitwise identical to eager +(`max_abs_error=0`). + +## F3 residual RMSNorm between GEMMs + +Pattern name: `f3_residual_rmsnorm` + +### Current terminal implementation + +The current pass implements two GEMM-residual-RMSNorm epilogues and stops at +the normalized activation: + +```text +F3-A: layers.L.attention.wo + residual -> layers.L.ffn_norm + +F3-B: layers.L.feed_forward.w2 + residual + -> layers.L+1.attention_norm + +F3-B: layers.L.moe.shared_experts.w2 + routed output + residual + -> layers.L+1.attention_norm +``` + +Each match creates one FlexGEMM. Its body preserves the source graph's BF16 +GEMM store and addition order, then returns the raw residual sum and +512-column FP32 mean-square partials. The root graph reduces the partials, +forms `rstd`, applies `rstd` and gamma to the raw sum in FP32, and stores the +normalized BF16 activation. The raw sum and optional saved `rstd` remain +available for checkpointed backward. + +This terminal form does not inspect or rewrite consumers of the norm. The +following Q/KV, dense SwiGLU, router, shared-expert, and routed grouped GEMMs +remain available to F2, F4, F6, and distMoE. F3 and F4 consequently have no +pass-order dependency. + +The focused CPU tests cover both boundary families, the shared-expert two-add +form, saved forward values, arbitrary downstream consumers, F4 composition, +and invalid module-role pairs. The full CODA pass test file has 37 passing +tests. A refreshed DSV3-671B fake-backend graph proof and GB300 benchmark are +still required; the counts and timings below predate the terminal rewrite. + +### Superseded cross-GEMM implementation + +The following evidence records the earlier implementation of the paper's +central GEMM-residual-RMSNorm-GEMM +reparameterization. It follows PyTorch's end-to-end FlexGEMM coverage in +`test_mm_coda_rmsnorm_rewrite_e2e`: the first FlexGEMM emits a gamma-weighted +activation and 512-column mean-square partials, a small root-graph reduction +forms the row-wise `rstd`, and each downstream FlexGEMM applies `rstd` in its +epilogue. The TorchTitan body additionally emits the raw residual sum because +the model reuses that value as the residual stream and checkpointed backward +requires it. + +The matcher covers three grounded boundaries: + +```text +layers.L.attention.wo -> residual -> layers.L.ffn_norm + -> layers.L.feed_forward.{w1,w3} + +layers.L.feed_forward.w2 -> residual -> layers.L+1.attention_norm + -> layers.L+1.attention.{wq_a,wkv_a} + +layers.L.moe.shared_experts.w2 + routed output -> residual + -> layers.L+1.attention_norm + -> layers.L+1.attention.{wq_a,wkv_a} +``` + +The first form is restricted to the three dense FFN layers. MoE normalized +activations also feed routing and grouped GEMMs, so that boundary remains +assigned to distMoE. The third form starts at the shared expert's ordinary W2 +GEMM after the routed expert collective. It captures both the routed result and +the residual in their original BF16 addition order; it does not rewrite either +grouped expert GEMM. + +For the first form, F3 also recognizes the exact dense SwiGLU epilogues handled +by F4. Its W1 FlexGEMM applies `rstd` and SiLU, then the W3 FlexGEMM applies +`rstd` and multiplies by the captured W1 activation. This keeps both fusions on +the six overlapping original/recomputed boundaries. When both patterns are +enabled, `f3_residual_rmsnorm` must therefore precede `f4_dense_swiglu`; config +validation rejects the reverse order instead of silently losing F3 coverage. + +Real before sample from layer 0: + +```python +mm_4 = torch.ops.aten.mm.default(view_22, t_4) +attention_out = torch.ops.aten.reshape.default(mm_4, [24, 4096, 7168]) +residual = torch.ops.aten.add.Tensor(embedding, attention_out) +norm = torch.ops.aten._fused_rms_norm.default( + residual, [7168], ffn_norm_weight, 1e-05 +) +normalized = norm[0] +w1 = torch.ops.aten.mm.default(normalized.reshape(98304, 7168), t_5) +w3 = torch.ops.aten.mm.default(normalized.reshape(98304, 7168), t_6) +``` + +Real after sample: + +```python +residual_2d = torch.ops.aten.reshape.default(embedding, [98304, 7168]) +gamma_2d = torch.ops.aten.reshape.default(ffn_norm_weight, [1, 7168]) +first = torch.ops.higher_order.flex_gemm( + torch.ops.aten.mm.default, + _coda_f3_residual_first_body_0, + (view_22, t_4, residual_2d, gamma_2d), + {}, + {"backend": "QUACK"}, +) +weighted = first[0] +residual = first[1] +partial_mean_square = first[2] +rstd = torch.ops.aten.rsqrt.default( + torch.ops.aten.add.Scalar( + torch.ops.aten.mean.dim(partial_mean_square, [-1], True), 1e-05 + ) +) +activated = torch.ops.higher_order.flex_gemm( + torch.ops.aten.mm.default, + _coda_f3_f4_silu_body_0, + (weighted, t_5, rstd), + {}, + {"backend": "QUACK"}, +)[0] +gate, product = torch.ops.higher_order.flex_gemm( + torch.ops.aten.mm.default, + _coda_f3_f4_mul_body_0, + (weighted, t_6, rstd, activated), + {}, + {"backend": "QUACK"}, +) +``` + +Real MoE-output before and after samples from layer 3: + +```python +# Before +shared = torch.ops.aten.mm.default(view_159, t_34) +shared = torch.ops.aten.reshape.default(shared, [24, 4096, 7168]) +moe_output = torch.ops.aten.add.Tensor(routed_output, shared) +residual = torch.ops.aten.add.Tensor(previous_residual, moe_output) +norm = torch.ops.aten._fused_rms_norm.default( + residual, [7168], attention_norm_weight, 1e-05 +) + +# After +routed_2d = torch.ops.aten.reshape.default(routed_output, [98304, 7168]) +residual_2d = torch.ops.aten.reshape.default(previous_residual, [98304, 7168]) +gamma_2d = torch.ops.aten.reshape.default(attention_norm_weight, [1, 7168]) +first = torch.ops.higher_order.flex_gemm( + torch.ops.aten.mm.default, + _coda_f3_residual_first_body_6, + (view_159, t_34, routed_2d, residual_2d, gamma_2d), + {}, + {"backend": "QUACK"}, +) +``` + +FSDP bucketing can define the next norm's unsharded gamma after the producing +GEMM. The new first HOP is therefore placed immediately before the norm, where +the GEMM inputs, residual, and gamma all dominate it. A late-gamma unit test +covers this real graph ordering. + +### Graph proof + +Configuration: DSV3-671B, fake communication backend, FSDP 256, EP 64, local +batch 24, sequence length 4096, full activation checkpointing, `c4_test`. + +Artifact root: + +```text +outputs/profiling/dsv3_fake/graph/coda-all-composed-f3-f4-20260810/tlparse/-_-_-_- +``` + +Before and after dumps: + +```text +before_fuse_f3_residual_rmsnorm_pass_278.txt +after_fuse_f3_residual_rmsnorm_pass_279.txt +``` + +Pass diff for the root graph: + +| Operation | Before | After | Delta | +| --- | ---: | ---: | ---: | +| root `mm.default` | 2,024 | 1,826 | -198 | +| root `_fused_rms_norm.default` | 490 | 424 | -66 | +| root `add.Tensor` | 959 | 836 | -123 | +| root `flex_gemm` | 124 | 322 | +198 | +| root `get_attr` | 551 | 749 | +198 | +| root `getitem` | 24,802 | 25,072 | +270 | +| root `mean.dim` | 0 | 66 | +66 | +| root `add.Scalar` | 0 | 66 | +66 | +| root `rsqrt.default` | 0 | 66 | +66 | +| root `_to_copy.default` | 2,478 | 2,481 | +3 | +| root `reshape.default` | 6,938 | 7,007 | +69 | +| root `silu.default` | 238 | 232 | -6 | + +All six original dense boundaries, 57 original MoE-output boundaries, and three +within-layer recomputed dense boundaries fused, covering 132 downstream +projections and six dense SwiGLU epilogues. The following F4 pass fused all 116 +remaining, non-overlapping SwiGLU chains. Cross-layer residual outputs are +`MUST_SAVE` under full activation checkpointing, so their backward-side norms +have no recomputed producer GEMM to fuse. The proof run then OOMed on the known +192 GiB unfused FlexAttention allocation after all pass artifacts were written. + +### GB300 microbenchmarks + +The representative attention-to-dense-FFN boundary uses BF16 tensors with +`M=98,304`, attention width `16,384`, model width `7,168`, and two FFN +projections of width `18,432`. It includes all three GEMMs, the residual add, +RMSNorm, partial-statistics reduction, row-scale epilogues, SiLU, and the final +gate multiply. Inputs were scaled by `0.02`; five warmups and 20 CUDA-event +iterations were run on one GB300. + +| Implementation | Median | Min | Max | +| --- | ---: | ---: | ---: | +| eager full boundary | 48.657 ms | 46.389 ms | 51.643 ms | +| three composed QUACK FlexGEMMs | 48.837 ms | 46.274 ms | 51.238 ms | + +Speedup: `1.00x` (`0.996x`; FlexGEMM is 0.4% slower). This remains under the +project policy that accepts FlexGEMM fusions without a speedup. The residual +output was bitwise exact. The final SwiGLU output had `max_abs_error=0.5` and +`mean_abs_error=0.00283090`. Moving `rstd` across the downstream GEMMs changes +rounding, so this pattern requires convergence validation. + +The representative MoE-output boundary uses a shared-expert W2 GEMM with +`M=98,304`, input width `2,048`, and model width `7,168`. Its first epilogue +captures both model-width routed and residual tensors, and the two downstream +attention projections have widths `1,536` and `576`. The benchmark otherwise +uses the same input scaling, warmups, and iteration count. + +| Implementation | Median | Min | Max | +| --- | ---: | ---: | ---: | +| eager shared W2 + two adds + RMSNorm + projections | 4.464 ms | 4.284 ms | 5.673 ms | +| three QUACK FlexGEMMs + two captures + partial reduction | 5.120 ms | 4.999 ms | 6.504 ms | + +Speedup: `0.87x`. The residual output was bitwise exact. Both projected outputs +had `max_abs_error=0.0625`; their mean absolute errors were `0.00223414` and +`0.00223534`. This does not justify a custom kernel under the project policy. + +## B5 MLA projection and RMSNorm backward + +Pattern name: `b5_mla_rmsnorm_backward`. + +The pass matches only backward BF16 input-gradient GEMMs from +`layers.*.attention.wkv_b` and `layers.*.attention.wq_b` whose sole reshape +feeds `_fused_rms_norm_backward` with both `dx` and `dweight` requested. The +FlexGEMM preserves the original BF16 GEMM store and emits 128-column partials +of `x_hat * grad_x_hat`. Regional Inductor then completes the row dot product, +`dx`, and the independent token-axis `dweight` reduction. The latter cannot be +put in the same FlexGEMM body because the two reductions use different grouped +layouts. + +Real before sample from the layer-60 KV path: + +```python +mm_581 = torch.ops.aten.mm.default(view_4476, _unsafe_view_2729) +view_4477 = torch.ops.aten.reshape.default(mm_581, [24, 4096, 512]) +_fused_rms_norm_backward_2 = ( + torch.ops.aten._fused_rms_norm_backward.default( + view_4477, + getitem_8778_recomputed, + [512], + alias_473, + _unsafe_view_2728, + [True, True], + ) +) +``` + +Real after sample: + +```python +flex_gemm = torch.ops.higher_order.flex_gemm( + torch.ops.aten.mm.default, + _coda_b5_mla_rmsnorm_body_0, + ( + view_4476, + _unsafe_view_2729, + reshape_default, + reshape_default_1, + reshape_default_2, + ), + {}, + {"backend": "QUACK"}, +) +rounded = flex_gemm[0] +partial_row_dot = flex_gemm[1] # f32[98304, 4] +row_dot = torch.ops.aten.sum.dim_IntList(partial_row_dot, [-1], True) +grad_input = torch.ops.aten.sub.Tensor(grad_x_hat, correction) +grad_weight = torch.ops.aten.sum.dim_IntList(grad_weight_terms, [0]) +``` + +Proof artifacts: + +- Before: + `outputs/profiling/dsv3_fake/graph/coda-b5-proof-20260810/tlparse/-_-_-_-/before_fuse_b5_mla_rmsnorm_backward_pass_274.txt` +- After: + `outputs/profiling/dsv3_fake/graph/coda-b5-proof-20260810/tlparse/-_-_-_-/after_fuse_b5_mla_rmsnorm_backward_pass_275.txt` + +The DSV3-671B fake-backend proof used local batch 24, sequence 4,096, EP64, +FSDP256, full activation checkpointing, and `c4_test`. It fused all 122 grounded +chains, two per transformer layer. Root `mm` and +`_fused_rms_norm_backward` counts each fell by 122; 122 FlexGEMMs were added. +The remaining 123 RMSNorm backward nodes are outside the MLA pattern. As in the +other proof runs, execution later OOMed on the known 192 GiB unfused +FlexAttention allocation after all graph artifacts were written. + +### GB300 microbenchmarks + +Both real shapes used `M=98,304`, BF16 inputs scaled by `0.02`, `rstd` sampled +from `[0.5, 1.5]`, five warmups, and 20 CUDA-event samples on one GB300. The KV +case is `(98,304, 32,768) @ (32,768, 512)`; the Q case is +`(98,304, 24,576) @ (24,576, 1,536)`. + +| Shape | Implementation | Median | Min | Max | +| --- | --- | ---: | ---: | ---: | +| KV, `N=512` | eager GEMM + fused RMSNorm backward | 1.975 ms | 1.966 ms | 2.100 ms | +| KV, `N=512` | FlexGEMM row partial + compiled completion | 2.558 ms | 2.523 ms | 2.915 ms | +| Q, `N=1,536` | eager GEMM + fused RMSNorm backward | 4.218 ms | 3.960 ms | 4.410 ms | +| Q, `N=1,536` | FlexGEMM row partial + compiled completion | 5.454 ms | 5.194 ms | 5.733 ms | + +The FlexGEMM path is `0.772x` eager for KV and `0.773x` for Q, about 29% slower +in both cases. It remains accepted under the project policy that permits a +FlexGEMM pattern without a speedup. KV `dweight` was bitwise exact; KV `dx` had +`max_abs_error=7.451e-9` and `mean_abs_error=1.573e-16`. Q `dx` had +`max_abs_error=5.821e-11` and `mean_abs_error=5.951e-19`; Q `dweight` had +`max_abs_error=4.768e-7` and `mean_abs_error=3.104e-10`. These differences come +from the 128-column partial reduction order and require convergence validation. + +## Historical composed graph proof + +This proof predates the terminal F3 implementation above. Its F3 row, total +FlexGEMM count, and composed F3/F4 conclusions must not be used to validate the +current pass. All 11 FlexGEMM patterns were enabled together after +`joint_transformer_block_bucketing_reordering_pass`. The DSV3-671B fake-backend +run used FSDP256, EP64, local batch 24, sequence length 4,096, full activation +checkpointing, force-balanced MoE routing, deterministic seed 42, and +`c4_test`. + +Artifact root: + +```text +outputs/profiling/dsv3_fake/graph/coda-all-proof-20260810/tlparse/-_-_-_- +``` + +| Pass | Grounded matches | FlexGEMMs added | +| --- | ---: | ---: | +| B1 LM-head input-gradient cast | 8 | 8 | +| F6 router sigmoid and bias | 116 | 116 | +| F3 residual RMSNorm | 66 boundaries across 132 projections | 198 | +| F4 dense/shared-expert SwiGLU | 116 chains | 232 | +| B2 dense/shared-expert SwiGLU backward | 61 derivative + 61 add chains | 122 | +| F2 MLA Q RMSNorm | 62 | 124 | +| F2 MLA KV RMSNorm | 62 | 124 | +| B4 router input-gradient add | 58 | 58 | +| B5 MLA RMSNorm backward | 122 | 122 | +| B6 BF16 weight-gradient cast | 496 | 496 | +| B7 attention input-gradient merge | 61 | 61 | + +The final root graph has 1,661 FlexGEMMs. Relative to the post-bucketing graph, +root `mm.default` fell from 2,148 to 487, `sigmoid.default` from 116 to zero, +`_fused_rms_norm.default` from 490 to 300, and +`_fused_rms_norm_backward.default` from 245 to 123. F3 ran before F4, so its +six dense SwiGLU epilogues were retained inside the composed F3 boundary while +F4 handled the remaining 116 shared-expert chains. Every isolated pattern +retained its expected match count. + +The run disabled regional Inductor and CUDA graphs to preserve every pass dump. +After all 19 configured graph passes and artifacts completed, execution hit the +known 192 GiB allocation in unfused FlexAttention. This failure is downstream +of the rewrite proof and is not a CODA pass failure. + +## CODA-kernels comparison + +CODA-kernels was measured from `~/local/coda-kernels` on branch +`gb300-perf-v061`, commit `b5afe0d`. The branch starts from CODA commit +`c9c4447`, the last commit with a coherent Quack 0.6.1 API, and adds SM100/110 +GEMM dispatch. CODA `main` at `8c7c4d5` starts a Quack 0.6.2 migration but mixes +the old batch-last ABI with new epilogue imports and does not run as checked +out. + +The two implementations require dependency-isolated runs: + +- PyTorch FlexGEMM: CUTLASS DSL 4.5.2 and PyTorch's vendored Quack. +- CODA-kernels: Quack 0.6.1 and CUTLASS DSL 4.6.1 with the CUDA 13 library + package, needed to import `GemmSm100` on this CUDA 13.0 host. + +Both runs used BF16 tensors scaled by `0.02`, five warmups, 20 CUDA-event +samples, and one GB300. The tables therefore compare steady-state GPU time but +not a single shared Python process. This matters for small differences: the F4 +eager median was `3.482 ms` in the CODA process and `3.613 ms` in the earlier +FlexGEMM process. + +### F1 LM-head and cross entropy + +The real per-chunk shape is `(12,288, 7,168) @ (7,168, 129,280)`. The eager +boundary is the exact graph sequence: BF16 GEMM, FP32 log-softmax, summed NLL, +FP32 log-softmax backward, BF16 `dlogits`, and the two activation/weight +gradient GEMMs. The loss is divided by the full 98,304-token local batch. + +FlexGEMM has no valid implementation for this row. Its local-reduction API +cannot combine online max/sum state across the full 129,280-column vocabulary +or select the target column. CODA instead uses `gemm_lse`, overwrites the BF16 +logits buffer with `dlogits` using `cross_entropy_fwd_bwd_`, and applies the loss +scale in the two gradient GEMM epilogues. This avoids the 6.35 GB FP32 +log-softmax tensor while retaining the backward data. + +| Implementation | Median | Min | Max | +| --- | ---: | ---: | ---: | +| eager full forward/backward boundary | 46.735 ms | 45.926 ms | 51.256 ms | +| PyTorch QUACK FlexGEMM | unsupported | unsupported | unsupported | +| CODA fused-LSE boundary | 42.794 ms | 41.780 ms | 43.763 ms | + +CODA is `1.092x` faster than eager. This is a fixed-config result, not an +autotuned result: `tile_m=256`, `tile_n=256`, `cluster_m=2`, `cluster_n=1`, +dynamic persistent scheduling, and no ping-pong. CODA's generic LSE autotuner +currently includes SM100 layouts whose epilogue has more than one N warp and +fails its `warps_in_N == 1` invariant, so the known-valid configuration was +selected directly. + +Against eager, CODA had `max_abs_error=1.192e-7` in loss and +`max_abs_error=7.451e-9` in both gradients. Mean absolute errors were +`3.164e-10` for the activation gradient and `2.963e-11` for the weight +gradient. CODA computes LSE from the FP32 GEMM accumulator and stores unscaled +BF16 `dlogits`, then moves the loss scale into the gradient GEMM epilogues; the +source graph rounds BF16 logits before FP32 log-softmax and rounds scaled +`dlogits` before the GEMMs. The faster kernel is therefore a performance +candidate, not a numerically identical replacement. + +This kernel was not wired into the GraphTrainer pass pipeline. CODA requires +Quack 0.6.1 and CUTLASS DSL 4.6.1, while the PyTorch build that provides the +landed FlexGEMM passes uses its vendored Quack and CUTLASS DSL 4.5.2. Loading +CODA directly would make the two implementations share an incompatible +`cutlass` Python package. The benchmark used the isolated environment +`~/local/coda-kernels/.venv-coda061`. + +### F4 shared-expert SwiGLU + +Real shape: `M=98,304`, `K=7,168`, `P=2,048`. CODA interleaves W1/W3 columns +and executes one `gemm_swiglu` with a `K x 2P` weight. This performs the same +GEMM FLOPs as the two eager/FlexGEMM projections and returns the combined raw +preactivation plus the final `P`-wide product. + +| Implementation | Median | Min | Max | +| --- | ---: | ---: | ---: | +| eager two GEMMs + SiLU + multiply | 3.482 ms | 3.359 ms | 3.743 ms | +| two PyTorch QUACK FlexGEMMs | 3.711 ms | 3.619 ms | 4.069 ms | +| CODA `gemm_swiglu` (autotuned) | 3.583 ms | 3.529 ms | 3.853 ms | + +CODA is `1.036x` faster than FlexGEMM but `0.972x` versus its paired eager +baseline. That difference is within the observed cross-process eager drift, so +there is no defensible F4 winner. CODA's preactivation was bitwise exact versus +the paired eager GEMMs. Its final product had `max_abs_error=6.104e-5` and +`mean_abs_error=8.497e-7`; CODA applies SwiGLU to the FP32 accumulator before +the BF16 output store, unlike the graph's explicit BF16 intermediate. + +### Historical composed F3 and F4 dense boundary + +Real shape: `(98,304, 16,384) @ (16,384, 7,168)`, followed by two +`(98,304, 7,168) @ (7,168, 18,432)` projections. CODA uses +`gemm_residual_partial_rmsnorm` followed by one interleaved +`gemm_rmsnorm_swiglu`, reducing three GEMMs to two without changing the total +FLOPs. + +| Implementation | Median | Min | Max | +| --- | ---: | ---: | ---: | +| eager full graph boundary | 48.657 ms | 46.389 ms | 51.643 ms | +| three composed PyTorch QUACK FlexGEMMs | 48.837 ms | 46.274 ms | 51.238 ms | +| two CODA GEMMs, SM100 default config | 46.715 ms | 45.555 ms | 47.510 ms | + +The CODA row is `1.045x` faster than FlexGEMM and `1.042x` faster than the +recorded eager median. It is not an autotuned result: CODA's full search was +stopped after more than 20 minutes while tuning the second GEMM. The measured +config was its SM100 default, `tile_m=256`, `tile_n=256`, `cluster_m=2`, and +`cluster_n=1`. + +This CODA path is not currently a drop-in numerical replacement. It adds the +residual in the GEMM accumulator, while the source graph stores BF16 before the +residual add. Against an `addmm`-ordered reference, the residual had +`max_abs_error=0.001953` and `mean_abs_error=6.117e-5`; the final product had +`max_abs_error=0.0001221` and `mean_abs_error=2.161e-6`. The result is a useful +performance bound, but landing it requires an explicit numerical policy and +convergence validation. + +### Pattern coverage + +| GraphTrainer pattern | Existing CODA-kernels analogue | Comparison status | +| --- | --- | --- | +| F1 LM-head plus cross entropy | `gemm_lse` + `cross_entropy_fwd_bwd_` + scaled GEMMs | measured above; faster external candidate, FlexGEMM unsupported | +| `b1_lm_head_input_grad_cast` | none with the required FP32 post-store cast | FlexGEMM/eager only | +| `b6_bf16_weight_grad_cast` | none with the required FP32 post-store cast | FlexGEMM/eager only | +| `f6_router_sigmoid_bias` | no public sigmoid-plus-bias GEMM | FlexGEMM/eager only | +| `f4_dense_swiglu` | `gemm_swiglu` | measured above | +| `b2_dense_swiglu_backward` | `gemm_swiglu_bwd_zdz` | ABI mismatch: CODA emits packed interleaved `dZ` and a full-row `ZdZ`; the graph consumes separate W1/W3 branches | +| `f2_q_rmsnorm` | QKV square-sum and RMS-scaled GEMM primitives | no exact 1,536-wide segmented boundary wrapper | +| `f2_kv_rmsnorm` | RMS-scaled GEMM primitives | no exact 512-plus-64 tail/RoPE wrapper | +| `b4_router_input_grad_add` | no public residual-add-only GEMM | FlexGEMM/eager only | +| `b5_mla_rmsnorm_backward` | `gemm_residual_partial_rmsnorm_bwd` has a different residual/`ZdZ` contract | FlexGEMM/eager measured above; no exact CODA row | +| `b7_attention_grad_merge` | no public captured-add-only GEMM | FlexGEMM/eager only | +| `f3_residual_rmsnorm` | residual partial RMSNorm plus RMS-scaled GEMM | measured with F4 above | + +The routed grouped-GEMM opportunities remain assigned to distMoE. They were +not included in this comparison. diff --git a/torchtitan/experiments/graph_trainer/CODA_INVESTIGATION_HANDOFF.md b/torchtitan/experiments/graph_trainer/CODA_INVESTIGATION_HANDOFF.md new file mode 100644 index 0000000000..f28967fa36 --- /dev/null +++ b/torchtitan/experiments/graph_trainer/CODA_INVESTIGATION_HANDOFF.md @@ -0,0 +1,292 @@ +# CODA investigation handoff + +Last updated: 2026-08-21 + +## Resume point + +The work is preserved on the local branch `coda-flex-gemm-passes`. Its base is +TorchTitan commit `9228564523aa63f78c5e3e038068a886572e90de`. Resume with: + +```bash +cd /home/bahuang/local/torchtitan +git switch coda-flex-gemm-passes +git log --oneline --decorate -20 +``` + +The branch is local and has no configured upstream. Do not delete it when +starting unrelated work from `main`. + +The main implementation is in: + +```text +/home/bahuang/local/torchtitan/torchtitan/experiments/graph_trainer/coda_passes.py +/home/bahuang/local/torchtitan/torchtitan/experiments/graph_trainer/configs.py +/home/bahuang/local/torchtitan/torchtitan/experiments/graph_trainer/passes.py +/home/bahuang/local/torchtitan/torchtitan/experiments/graph_trainer/tests/test_coda_passes.py +``` + +The graph evidence and pattern-by-pattern status are in: + +```text +/home/bahuang/local/torchtitan/torchtitan/experiments/graph_trainer/CODA_FUSION_RESULTS.md +``` + +## Branch history + +The CODA branch contains these commits after its TorchTitan base: + +```text +73408dc18 Fuse B6 BF16 weight-gradient casts with FlexGEMM +d2cf35a47 Fuse F6 router sigmoid and bias with FlexGEMM +32ede641b Fuse F4 dense SwiGLU with FlexGEMM +538cf666c Fuse B2 dense SwiGLU backward with FlexGEMM +ccd9d2e6f Fuse F2 MLA Q RMSNorm with FlexGEMM +dff1bbe10 Fuse B4 router input gradient with FlexGEMM +f2efa2187 Fuse F3 residual RMSNorm with FlexGEMM +421675a7b Fuse F2 MLA KV RMSNorm with FlexGEMM +f11d8c709 Fuse B7 attention gradient merge with FlexGEMM +6eb058388 Fuse B1 LM-head input gradient casts with FlexGEMM +841ad78d6 Compose F3 RMSNorm and F4 SwiGLU FlexGEMMs +8aeed44d6 Document CODA kernel GB300 comparisons +08b7cc582 Fuse B5 MLA RMSNorm backward with FlexGEMM +dab40eeb2 Document CODA linear cross-entropy benchmark +38f0ad020 Document composed CODA graph proof +09c7c538f Implement terminal F3 residual RMSNorm fusions +388cdae19 Add CODA FlexGEMM autotuning benchmarks +100b595ef Add reusable DSV3 profiling workflow +``` + +The handoff document itself is in the next branch commit. Use `git log` rather +than relying on this list if more commits are added after resuming. + +## Implemented pass scope + +The branch implements GraphTrainer rewrites for: + +- B1: LM-head input-gradient BF16 store rounding plus FP32 cast. +- B2: dense/shared-expert SwiGLU backward branch derivatives and input add. +- B4: router input-gradient GEMM, BF16 store rounding, and expert-gradient add. +- B5: MLA projection plus RMSNorm backward. +- B6: BF16 weight-gradient store rounding plus FP32 cast. +- B7: Q and KV attention input-gradient merge. +- F2-Q: Q low-rank projection, RMSNorm, and Q-B projection. +- F2-KV: segmented KV low-rank projection, RMSNorm, and KV-B projection. +- F3-A: attention WO projection, residual add, and following RMSNorm. +- F3-B: FFN/shared-expert W2 projection, residual path, and following RMSNorm. +- F4: dense/shared-expert W1/W3 SwiGLU forward. +- F6: router GEMM, sigmoid, and optional expert bias. + +F3 and F4 composition is implemented so the terminal F3 producer can stop at +RMSNorm without consuming the next layer's projection. The pass runs after +`joint_transformer_block_bucketing_reordering_pass` because the patterns and +FSDP bucketing boundaries must be final before matching. + +Routed `aten._grouped_mm` chains are intentionally excluded. They belong to +DistMoE, especially under MinimalAsyncEP. Do not add dense FlexGEMM rewrites +around those routed grouped GEMMs without coordinating with the DistMoE +implementation. + +F1 linear cross-entropy is documented but is not implemented as a FlexGEMM +pass. CODA kernels can reduce the full vocabulary materialization and remain a +better external candidate for that pattern. + +## FX graph evidence + +The original DSV3-671B traced graph is: + +```text +/home/bahuang/local/torchtitan/outputs/profiling/dsv3_fake/graph/20260805-110153/tlparse/-_-_-_-/make_fx_graph_traced_257.txt +``` + +The post-bucketing graph used for the CODA analysis is: + +```text +/home/bahuang/local/torchtitan/outputs/profiling/dsv3_fake/graph/20260805-110153/tlparse/-_-_-_-/after_joint_transformer_block_bucketing_reordering_pass_273.txt +``` + +The most useful composed F3/F4 proof graph is: + +```text +/home/bahuang/local/torchtitan/outputs/profiling/dsv3_fake/graph/coda-all-composed-f3-f4-20260810/tlparse/-_-_-_-/after_joint_transformer_block_bucketing_reordering_pass_273.txt +``` + +The unfused DSV3-16B post-bucketing graph used to extract the smaller-model +benchmark shapes is: + +```text +/home/bahuang/local/torchtitan/outputs/profiling/dsv3_fake/graph/dsv3-16b-unfused-shapes-20260812/tlparse/-_-_-_-/after_joint_transformer_block_bucketing_reordering_pass_142.txt +``` + +Pattern-specific proof directories are listed under each `Graph proof` +section of `CODA_FUSION_RESULTS.md`. + +## Benchmark code + +The standalone source-eager, compiled-eager, and FlexGEMM suite is here: + +```text +/home/bahuang/local/torchtitan/torchtitan/experiments/graph_trainer/benchmarks/coda_fusion_microbench.py +/home/bahuang/local/torchtitan/torchtitan/experiments/graph_trainer/benchmarks/coda_fusion_microbench_16b.py +/home/bahuang/local/torchtitan/torchtitan/experiments/graph_trainer/benchmarks/coda_fusion_autotune.py +/home/bahuang/local/torchtitan/tests/unit_tests/test_coda_fusion_microbench.py +``` + +Every FlexGEMM uses `tuned: true`. Epilogues containing SiLU or sigmoid use +fast math. The external tuner isolates every explicit QUACK configuration in +a fresh process because some SM100 configurations fail or hang on SM103. +Multi-GEMM cases use coordinate descent; a full search evaluates all 74 +configurations per FlexGEMM, not the Cartesian product of configurations. + +Run a 16B case with the known CUDA paths pinned: + +```bash +cd /home/bahuang/local/torchtitan +CUDA_ROOT=/home/bahuang/local/venvs/torch-cu132-nightly/lib/python3.12/site-packages/nvidia/cu13 +export CUDA_HOME="$CUDA_ROOT" +export PATH="$CUDA_ROOT/bin:/usr/local/bin:/usr/bin:/bin" +export LD_LIBRARY_PATH="$CUDA_ROOT/lib" +export TORCH_NATIVE_SKIP_VERSION_CHECK=1 + +python -m torchtitan.experiments.graph_trainer.benchmarks.coda_fusion_autotune \ + --suite 16b \ + --case f4_shared_expert_swiglu \ + --devices 0,1,2,3 \ + --search full \ + --passes 1 +``` + +Before rerunning, verify that this environment really reports CUDA 13.2. As of +2026-08-21, the venv at that path has drifted and reports PyTorch +`2.15.0a0+git1b5baff`, CUDA 13.0, and git revision +`1b5baff2649da3d5c57ad14d72cc2dbc24dbdf72`. It is not the environment that +produced the recorded CUDA 13.2 results. + +The recorded 16B exhaustive run used: + +```text +PyTorch: 2.15.0.dev20260812+cu132 +PyTorch git: 3eb0e5d0968d26971fdd6684ab5c9b605bfec4a6 +CUDA: 13.2 +CUTLASS DSL: 4.6.2 +Triton: 3.8.0+git675c5987 +GPU: NVIDIA GB300, compute capability 10.3 +``` + +## Performance conclusions + +The DSV3-671B 12-pattern results are recorded in +`CODA_FUSION_RESULTS.md`. The strongest isolated results were B4 at `11.582x` +and F6 at `9.428x` versus compiled eager. F4 was `1.037x`; most large-GEMM and +RMSNorm patterns were neutral or slower. + +The DSV3-16B exhaustive run covered 13 cases, 18 FlexGEMMs, and 1,332 +candidate evaluations. The primary long-run results versus compiled eager +were: + +| Pattern | Speedup | Conclusion | +| --- | ---: | --- | +| B4 router input-gradient add | `4.531x` | Strong, stable win | +| F4 shared-expert SwiGLU | `1.191x` | Useful win | +| B7 attention input-gradient merge | `1.039x` | Small win | +| F4 dense SwiGLU | `1.015x` | Marginal; verify end to end | +| B1 LM-head cast | `1.014x` | Effectively neutral | +| B6 shared weight-gradient cast | `1.063x` final, `0.979x` verification | Noisy/inconclusive | +| F2-KV RMSNorm | `0.981x` | Slight regression | +| B5 KV RMSNorm backward | `0.891x` | Regression | +| B2 shared/dense backward | `0.944x` / `0.953x` | Regression | +| F3 attention/MoE/dense | `0.849x` / `0.779x` / `0.959x` | Regression | + +The complete 16B report and raw results are: + +```text +/home/bahuang/local/torchtitan/outputs/coda_fusion_microbench/dsv3_16b_exhaustive_autotune/FULL_REPORT.md +/home/bahuang/local/torchtitan/outputs/coda_fusion_microbench/dsv3_16b_exhaustive_autotune/16b +``` + +The portable 671B autotuning handoff is: + +```text +/home/bahuang/local/torchtitan/outputs/coda_fusion_autotune_handoff +/home/bahuang/local/torchtitan/outputs/coda_fusion_autotune_handoff.zip +``` + +## Known issues + +1. Reduction configuration coverage is the largest FlexGEMM limitation. F3 + accepts only 2 of the 74 generic configurations because the local RMSNorm + reduction requires group 512. B5 and F2-KV have similar restrictions. +2. Some configuration-specific QUACK kernels raise + `cudaErrorNoKernelImageForDevice` on SM103. This does not prevent other + FlexGEMM configurations from running on GB300. +3. Native unconstrained `tuned: true` can lose the whole search when a bad + SM103 candidate fails. The process-isolated tuner is a workaround; the + backend should reject unsupported candidates before execution. +4. Small gains below roughly 3% are sensitive to GPU load and clocks. Require + repeated fixed-GPU and end-to-end evidence before enabling those patterns. +5. Moving RMSNorm state such as `rstd` across a projection changes scheduling + and requires convergence validation, even when isolated outputs pass the + microbenchmark tolerance. +6. The branch is based on an older TorchTitan `main` and must be rebased or + replayed deliberately before upstreaming. Audit the pass pipeline and all + config callsites after rebasing. + +## Recommended next work + +1. Rebase the branch onto a current TorchTitan `main`, keeping each logical + pattern commit separate so regressions can be bisected. +2. Run GraphTrainer end-to-end A/B tests with only B4 enabled, then only F6, + then both. Use `c4_test`, at least 10 steps, seed 42, deterministic mode, + identical parallelism, and compare full-precision loss and grad norm. +3. Add a minimized PyTorch FlexGEMM test for SM103 candidates that currently + produce `cudaErrorNoKernelImageForDevice`, then fix candidate filtering. +4. Extend local-reduction layout support for group 512 and retune F3. More + generic search is not useful until additional layouts are legal. +5. Compare specialized CODA kernels against FlexGEMM for F1 and the + reduction-heavy F2/F3/B5 cases using the exact FX shapes. +6. Keep routed grouped GEMMs under DistMoE ownership. Test any dense epilogue + changes with standard EP and MinimalAsyncEP separately. +7. Performance-gate broad graph rewrites. The evidence supports prioritizing + B4 and F6; the remaining patterns need stronger end-to-end justification. + +The local CODA kernel checkout used for earlier comparisons is: + +```text +/home/bahuang/local/coda-kernels +branch: gb300-perf-v061 +commit: b5afe0d9572b66845445efbc0399a4151e51235a +``` + +## Figures + +```text +/home/bahuang/local/torchtitan/outputs/coda_fusion_patterns/coda_fusion_patterns.png +/home/bahuang/local/torchtitan/outputs/dsv3_one_layer_graph/dsv3_layer60_gemms.png +/home/bahuang/local/torchtitan/outputs/dsv3_one_layer_graph/dsv3_layer60_coda_forward_patterns.png +/home/bahuang/local/torchtitan/outputs/dsv3_one_layer_graph/dsv3_cross_layer_f3.png +/home/bahuang/local/torchtitan/outputs/dsv3_one_layer_graph/dsv3_layer60_coda_backward_patterns.png +``` + +## Validation state + +Before the final handoff commit: + +- `tests/unit_tests/test_coda_fusion_microbench.py`: 8 passed. +- Benchmark Python files: `flake8`, `ufmt`, `pydoclint`, `codespell`, and + targeted Pyrefly passed. +- Profiling shell script: `bash -n` and ShellCheck passed. +- Trace uploader: Python compilation, `flake8`, `ufmt`, `pydoclint`, + `codespell`, and targeted Pyrefly passed. +- Repository `pre-commit` could not bootstrap because fetching + `pre-commit-hooks` from GitHub returned proxy HTTP 403. +- Repository-wide Pyrefly reaches unrelated optional-dependency errors for + `transformers` and `rich`; targeted checks for the added Python files pass. + +Run the core pass tests again after rebasing because the pass manager and FX +graph details may have changed on current `main`: + +```bash +cd /home/bahuang/local/torchtitan +python -m pytest \ + torchtitan/experiments/graph_trainer/tests/test_coda_passes.py \ + torchtitan/experiments/graph_trainer/tests/test_passes.py -q +``` diff --git a/torchtitan/experiments/graph_trainer/benchmarks/DSV3_16B_CODA_BENCHMARK_SHAPES.md b/torchtitan/experiments/graph_trainer/benchmarks/DSV3_16B_CODA_BENCHMARK_SHAPES.md new file mode 100644 index 0000000000..f169a7b791 --- /dev/null +++ b/torchtitan/experiments/graph_trainer/benchmarks/DSV3_16B_CODA_BENCHMARK_SHAPES.md @@ -0,0 +1,236 @@ +# DeepSeek-V3 16B CODA benchmark shapes + +## Capture + +The source is an unfused GraphTrainer joint graph captured with: + +- model: `graph_trainer_deepseek_v3_16b` (15.706B parameters, 2.661B active) +- fake world: FSDP4, EP4, TP1, standard EP +- batch: local batch 4, sequence length 4096 (`M=16384`) +- activation checkpointing: full +- dataset: `c4_test` +- graph passes: `coda_patterns: []` + +The analyzed graph is: + +```text +outputs/profiling/dsv3_fake/graph/dsv3-16b-unfused-shapes-20260812/ + tlparse/-_-_-_-/after_joint_transformer_block_bucketing_reordering_pass_142.txt +``` + +It contains 858 `aten.mm` nodes and 312 `aten._grouped_mm` nodes. It contains +no FlexGEMM or CODA nodes. Routed expert token dimensions remain symbolic after +EP all-to-all, so the grouped GEMMs are excluded from this dense FlexGEMM suite +and remain owned by DistMoE. + +## Benchmark inventory + +| Case | Pattern | Real GEMM shape | Occurrence | +| --- | --- | --- | --- | +| `b1_lm_head_input_grad_cast` | B1 | `(2048, 102400) @ (102400, 2048)` | 8 loss chunks | +| `b2_shared_expert_swiglu_backward` | B2 | `M=16384, D=2048, P=2816`, 3 GEMMs | 26 MoE layers | +| `b2_dense_ffn_swiglu_backward` | B2 | `M=16384, D=2048, P=10944`, 3 GEMMs | dense layer 0 | +| `b4_router_input_grad_add` | B4 | `(16384, 64) @ (64, 2048)` | 26 MoE layers | +| `b5_mla_kv_rmsnorm_backward` | B5 | `(16384, 4096) @ (4096, 512)` | 27 layers | +| `b6_shared_expert_weight_grad_cast` | B6 | `(2816, 16384) @ (16384, 2048)` | shared W1/W3 | +| `b7_attention_input_grad_merge` | B7 | `(16384, 576) @ (576, 2048)` + `(16384, 3072) @ (3072, 2048)` | 27 layers | +| `f2_kv_rmsnorm` | F2-KV | `2048 -> 576`, RMSNorm(512), `512 -> 4096`, `M=16384` | 27 layers | +| `f3_attention_output` | F3-A | `(16384, 2048) @ (2048, 2048)` | 27 layers | +| `f3_moe_output` | F3-B | `(16384, 2816) @ (2816, 2048)` | 26 MoE layers | +| `f3_dense_ffn_output` | F3-B | `(16384, 10944) @ (10944, 2048)` | dense layer 0 | +| `f4_shared_expert_swiglu` | F4 | two `(16384, 2048) @ (2048, 2816)` | 26 MoE layers | +| `f4_dense_ffn_swiglu` | F4 | two `(16384, 2048) @ (2048, 10944)` | dense layer 0 | + +F2-Q and F6 are intentionally absent. This 16B configuration has a direct Q +projection (`2048 -> 3072`) rather than Q LoRA plus RMSNorm, and its router is +softmax without the sigmoid-plus-expert-bias epilogue. + +## Extracted forward samples + +These are real nodes copied from the post-bucketing FX graph. + +### F2-KV and direct Q + +```python +mm_7: "bf16[16384, 3072]" = torch.ops.aten.mm.default(view_29, t_7) + +mm_8: "bf16[16384, 576]" = torch.ops.aten.mm.default(view_32, t_8) +split_with_sizes_4 = torch.ops.aten.split_with_sizes.default( + _unsafe_view_8, [512, 64], -1 +) +_fused_rms_norm_4 = torch.ops.aten._fused_rms_norm.default( + getitem_19, [512], _unsafe_view_538, 1e-05 +) +mm_9: "bf16[16384, 4096]" = torch.ops.aten.mm.default(view_40, t_9) +``` + +The `3072` direct-Q GEMM feeds B7 backward but does not form F2-Q. + +### F3-A attention output + +```python +mm_3: "bf16[16384, 2048]" = torch.ops.aten.mm.default(view_19, t_3) +_unsafe_view_3 = torch.ops.aten.reshape.default(mm_3, [4, 4096, 2048]) +add = torch.ops.aten.add.Tensor(embedding, _unsafe_view_3) +_fused_rms_norm_2 = torch.ops.aten._fused_rms_norm.default( + add, [2048], _unsafe_view_522, 1e-05 +) +``` + +### Dense-layer F4 and F3-B + +```python +mm_4: "bf16[16384, 10944]" = torch.ops.aten.mm.default(view_22, t_4) +silu = torch.ops.aten.silu.default(_unsafe_view_4) +mm_5: "bf16[16384, 10944]" = torch.ops.aten.mm.default(view_24, t_5) +mul_2 = torch.ops.aten.mul.Tensor(silu, _unsafe_view_5) + +mm_6: "bf16[16384, 2048]" = torch.ops.aten.mm.default(view_26, t_6) +add_1 = torch.ops.aten.add.Tensor(add, _unsafe_view_6) +_fused_rms_norm_3 = torch.ops.aten._fused_rms_norm.default( + add_1, [2048], _unsafe_view_541, 1e-05 +) +``` + +### Shared-expert F4 and F3-B + +```python +mm_12: "bf16[16384, 2816]" = torch.ops.aten.mm.default(view_91, t_14) +silu_2 = torch.ops.aten.silu.default(_unsafe_view_14) +mm_13: "bf16[16384, 2816]" = torch.ops.aten.mm.default(view_93, t_15) +mul_8 = torch.ops.aten.mul.Tensor(silu_2, _unsafe_view_15) + +mm_14: "bf16[16384, 2048]" = torch.ops.aten.mm.default(view_95, t_16) +add_5 = torch.ops.aten.add.Tensor(view_83, _unsafe_view_14) +add_6 = torch.ops.aten.add.Tensor(add_2, add_5) +``` + +## Extracted backward samples + +### B1 chunked LM head + +```python +view_1827: "bf16[2048, 102400]" = torch.ops.aten.reshape.default( + view_1826, [2048, 102400] +) +mm_217: "bf16[2048, 2048]" = torch.ops.aten.mm.default( + view_1827, _unsafe_view_1109 +) +_to_copy_723: "f32[4, 512, 2048]" = torch.ops.aten._to_copy.default( + alias_173, dtype=torch.float32 +) +``` + +The graph repeats this sequence eight times because the loss is chunked into +512 tokens per local batch element. + +### B2 and B6 shared expert + +```python +mm_240: "bf16[16384, 2816]" = torch.ops.aten.mm.default( + view_1887, _unsafe_view_1106 +) +mul_159 = torch.ops.aten.mul.Tensor(view_1888, silu_52_recomputed) +mul_160 = torch.ops.aten.mul.Tensor(view_1888, _unsafe_view_265_recomputed) + +mm_241: "bf16[2816, 2048]" = torch.ops.aten.mm.default( + t_311, view_1818_recomputed +) +mm_242: "bf16[16384, 2048]" = torch.ops.aten.mm.default( + view_1890, _unsafe_view_1107 +) +silu_backward = torch.ops.aten.silu_backward.default( + mul_160, _unsafe_view_264_recomputed +) +mm_243: "bf16[2816, 2048]" = torch.ops.aten.mm.default( + t_315, view_1816_recomputed +) +mm_244: "bf16[16384, 2048]" = torch.ops.aten.mm.default( + view_1893, _unsafe_view_1105 +) +add_140 = torch.ops.aten.add.Tensor(view_1891, view_1894) +_to_copy_763: "f32[2816, 2048]" = torch.ops.aten._to_copy.default( + mm_243, dtype=torch.float32 +) +``` + +The layer-0 dense backward has the same topology with `P=10944`; it is kept as +a separate B2 case because its GEMM aspect ratios and tuning choices differ. + +### B4 router input gradient + +```python +mm_246: "f32[16384, 2048]" = torch.ops.aten.mm.default( + view_1910, _to_copy_704_recomputed +) +_to_copy_770: "bf16[4, 4096, 2048]" = torch.ops.aten._to_copy.default( + view_1911, dtype=torch.bfloat16 +) +add_143 = torch.ops.aten.add.Tensor(add_142, _to_copy_770) +``` + +### B5 KV projection plus RMSNorm backward + +```python +mm_250: "bf16[16384, 512]" = torch.ops.aten.mm.default( + view_1919, _unsafe_view_1100 +) +view_1920 = torch.ops.aten.reshape.default(mm_250, [4, 4096, 512]) +_fused_rms_norm_backward_2 = torch.ops.aten._fused_rms_norm_backward.default( + view_1920, + getitem_694_recomputed, + [512], + alias_210, + _unsafe_view_1099, + [True, True], +) +``` + +### B7 direct-Q and KV input-gradient merge + +```python +mm_252: "bf16[16384, 2048]" = torch.ops.aten.mm.default( + view_1927, _unsafe_view_1098 +) +mm_254: "bf16[16384, 2048]" = torch.ops.aten.mm.default( + view_1931, _unsafe_view_1097 +) +add_145 = torch.ops.aten.add.Tensor(view_1928, view_1932) +``` + +## Running the suite + +For the CUDA 13.2 nightly used during validation, pin the toolkit libraries to +the venv. Loading this wheel against the host CUDA 13.0 libraries can produce +`cudaErrorNoKernelImageForDevice` on SM103. + +```bash +CUDA_ROOT=/home/bahuang/local/venvs/torch-cu132-nightly/lib/python3.12/site-packages/nvidia/cu13 +export CUDA_HOME="$CUDA_ROOT" +export PATH="$CUDA_ROOT/bin:$PATH" +export LD_LIBRARY_PATH="$CUDA_ROOT/lib" + +python -m torchtitan.experiments.graph_trainer.benchmarks.coda_fusion_microbench_16b --list + +python -m torchtitan.experiments.graph_trainer.benchmarks.coda_fusion_microbench_16b \ + --case f4_shared_expert_swiglu \ + --warmup 5 \ + --rounds 5 \ + --iterations 20 +``` + +Run different cases in fresh processes on SM103. QuACK CUDA dialect +initialization can make a later case fail if several FlexGEMM programs run in +one process. The tuner already provides this isolation: + +```bash +python -m torchtitan.experiments.graph_trainer.benchmarks.coda_fusion_autotune \ + --suite 16b \ + --case f4_shared_expert_swiglu \ + --devices 0,1,2,3 \ + --search full +``` + +Every FlexGEMM uses `tuned: true`; F4's SiLU-producing GEMM also uses +`fast_math: true`. Pass `--config` once per FlexGEMM to force explicit QuACK +configurations during tuning. diff --git a/torchtitan/experiments/graph_trainer/benchmarks/coda_fusion_autotune.py b/torchtitan/experiments/graph_trainer/benchmarks/coda_fusion_autotune.py new file mode 100644 index 0000000000..58bf3b884a --- /dev/null +++ b/torchtitan/experiments/graph_trainer/benchmarks/coda_fusion_autotune.py @@ -0,0 +1,727 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +"""Process-isolated SM100 tuning for the DSV3 CODA FlexGEMM benchmarks. + +QuACK configurations are compiled in separate processes because a broken +candidate can fail or hang inside a CUDA kernel. Multi-GEMM patterns use +coordinate descent: one FlexGEMM configuration is swept while the other +winning configurations remain fixed. + +Select the graph-grounded model shape inventory with ``--suite 671b`` or +``--suite 16b``. +""" + +import argparse +import dataclasses +import hashlib +import json +import os +import queue +import shutil +import signal +import subprocess +import sys +from collections.abc import Sequence +from concurrent.futures import as_completed, ThreadPoolExecutor +from pathlib import Path +from typing import Any + +if __package__: + from torchtitan.experiments.graph_trainer.benchmarks.coda_fusion_microbench import ( + CASES as CASES_671B, + ) + from torchtitan.experiments.graph_trainer.benchmarks.coda_fusion_microbench_16b import ( + CASES as CASES_16B, + ) +else: + from coda_fusion_microbench import ( # pyrefly: ignore [missing-import] + CASES as CASES_671B, + ) + from coda_fusion_microbench_16b import ( # pyrefly: ignore [missing-import] + CASES as CASES_16B, + ) + + +CASES = CASES_671B +SUITES = {"671b": CASES_671B, "16b": CASES_16B} + + +Config = dict[str, Any] + + +def _microbenchmark_command(suite: str = "671b") -> list[str]: + module_suffix = ( + "coda_fusion_microbench_16b" if suite == "16b" else "coda_fusion_microbench" + ) + if __package__: + return [ + sys.executable, + "-m", + f"torchtitan.experiments.graph_trainer.benchmarks.{module_suffix}", + ] + return [sys.executable, str(Path(__file__).with_name(f"{module_suffix}.py"))] + + +def _config( + tile_m: int, + tile_n: int, + cluster_m: int, + cluster_n: int, + is_dynamic_persistent: bool, +) -> Config: + return { + "tile_m": tile_m, + "tile_n": tile_n, + "is_dynamic_persistent": is_dynamic_persistent, + "cluster_m": cluster_m, + "cluster_n": cluster_n, + "swap_ab": False, + "use_tma_gather": False, + } + + +PRIORITY_CONFIGS = ( + _config(128, 256, 2, 1, True), + _config(128, 192, 2, 1, True), + _config(256, 256, 2, 1, True), + _config(256, 256, 2, 2, True), + _config(256, 192, 2, 1, True), + _config(128, 128, 1, 1, False), + _config(128, 256, 1, 1, True), + _config(128, 256, 1, 1, False), + _config(128, 128, 2, 1, True), + _config(256, 128, 2, 1, True), + _config(128, 224, 1, 1, True), + _config(128, 160, 1, 1, True), +) + +WIDE_REDUCTION_CONFIG = _config(256, 512, 2, 1, False) +LARGE_GEMM_CONFIG = _config(256, 256, 2, 2, True) + +INITIAL_CONFIGS: dict[str, tuple[Config, ...]] = { + "b1_lm_head_input_grad_cast": (LARGE_GEMM_CONFIG,), + "b2_shared_expert_swiglu_backward": ( + LARGE_GEMM_CONFIG, + LARGE_GEMM_CONFIG, + ), + "b4_router_input_grad_add": (PRIORITY_CONFIGS[1],), + "b5_mla_q_rmsnorm_backward": (PRIORITY_CONFIGS[0],), + "b6_weight_grad_cast": (LARGE_GEMM_CONFIG,), + "b7_attention_input_grad_merge": (LARGE_GEMM_CONFIG,), + "f2_q_rmsnorm": (WIDE_REDUCTION_CONFIG, LARGE_GEMM_CONFIG), + "f2_kv_rmsnorm": (PRIORITY_CONFIGS[0], LARGE_GEMM_CONFIG), + "f3_attention_output": (WIDE_REDUCTION_CONFIG,), + "f3_moe_output": (WIDE_REDUCTION_CONFIG,), + "f4_shared_expert_swiglu": (LARGE_GEMM_CONFIG, LARGE_GEMM_CONFIG), + "f6_router_sigmoid_bias": (PRIORITY_CONFIGS[1],), +} + +INITIAL_CONFIGS_16B: dict[str, tuple[Config, ...]] = { + "b1_lm_head_input_grad_cast": (LARGE_GEMM_CONFIG,), + "b2_shared_expert_swiglu_backward": ( + LARGE_GEMM_CONFIG, + LARGE_GEMM_CONFIG, + ), + "b2_dense_ffn_swiglu_backward": ( + LARGE_GEMM_CONFIG, + LARGE_GEMM_CONFIG, + ), + "b4_router_input_grad_add": (PRIORITY_CONFIGS[1],), + "b5_mla_kv_rmsnorm_backward": (WIDE_REDUCTION_CONFIG,), + "b6_shared_expert_weight_grad_cast": (LARGE_GEMM_CONFIG,), + "b7_attention_input_grad_merge": (LARGE_GEMM_CONFIG,), + "f2_kv_rmsnorm": (PRIORITY_CONFIGS[0], LARGE_GEMM_CONFIG), + "f3_attention_output": (WIDE_REDUCTION_CONFIG,), + "f3_moe_output": (WIDE_REDUCTION_CONFIG,), + "f3_dense_ffn_output": (WIDE_REDUCTION_CONFIG,), + "f4_shared_expert_swiglu": (LARGE_GEMM_CONFIG, LARGE_GEMM_CONFIG), + "f4_dense_ffn_swiglu": (LARGE_GEMM_CONFIG, LARGE_GEMM_CONFIG), +} + + +@dataclasses.dataclass(frozen=True) +class CandidateResult: + index: int + configs: tuple[Config, ...] + status: str + device: int + median_ms: float | None + compiled_eager_ms: float | None + flex_to_compiled_eager: float | None + result_path: str + log_path: str + error: str | None = None + + +def _full_sm100_configs() -> tuple[Config, ...]: + tile_cluster_shapes = ( + *( + (128, tile_n, cluster_m, cluster_n) + for tile_n in (64, 128, 160, 192, 224, 256) + for cluster_m, cluster_n in ((1, 1), (1, 2), (2, 1), (2, 2)) + ), + *( + (256, tile_n, cluster_m, cluster_n) + for tile_n in (64, 128, 160, 192, 224, 256) + for cluster_m, cluster_n in ((2, 1), (2, 2)) + ), + (256, 512, 2, 1), + ) + configs = [ + _config(tile_m, tile_n, cluster_m, cluster_n, is_dynamic_persistent) + for tile_m, tile_n, cluster_m, cluster_n in tile_cluster_shapes + for is_dynamic_persistent in (True, False) + ] + priority_keys = {_config_key(config) for config in PRIORITY_CONFIGS} + configs.sort( + key=lambda config: ( + 0 if _config_key(config) in priority_keys else 1, + config["tile_m"], + config["tile_n"], + config["cluster_m"], + config["cluster_n"], + not config["is_dynamic_persistent"], + ) + ) + return tuple(configs) + + +def _config_key(config: Config) -> str: + return json.dumps(config, sort_keys=True, separators=(",", ":")) + + +def _deduplicate_configs(configs: Sequence[Config]) -> tuple[Config, ...]: + unique = {} + for config in configs: + unique.setdefault(_config_key(config), config) + return tuple(unique.values()) + + +def _deduplicate_config_sets( + config_sets: Sequence[tuple[Config, ...]], +) -> tuple[tuple[Config, ...], ...]: + unique = {} + for configs in config_sets: + key = json.dumps(configs, sort_keys=True, separators=(",", ":")) + unique.setdefault(key, configs) + return tuple(unique.values()) + + +def _search_configs(search: str, max_candidates: int | None) -> tuple[Config, ...]: + configs = ( + PRIORITY_CONFIGS + if search == "priority" + else _deduplicate_configs((*PRIORITY_CONFIGS, *_full_sm100_configs())) + ) + return configs if max_candidates is None else configs[:max_candidates] + + +def _parse_config(value: str) -> Config: + config = json.loads(value) + if not isinstance(config, dict): + raise argparse.ArgumentTypeError("configuration must be a JSON object") + return config + + +def _initial_configs(args: argparse.Namespace) -> tuple[Config, ...]: + suite = getattr(args, "suite", "671b") + cases = SUITES[suite] + initial_configs = INITIAL_CONFIGS_16B if suite == "16b" else INITIAL_CONFIGS + case = cases[args.case] + if args.base_config: + configs = tuple(args.base_config) + if len(configs) == 1: + configs *= case.num_flex_gemms + else: + configs = initial_configs[args.case] + if len(configs) != case.num_flex_gemms: + raise ValueError( + f"{args.case} requires {case.num_flex_gemms} base configurations, " + f"got {len(configs)}" + ) + return configs + + +def _result_metrics(path: Path) -> tuple[float, float, float]: + with path.open(encoding="utf-8") as result_file: + results = json.load(result_file) + if not isinstance(results, list) or len(results) != 1: + raise ValueError(f"expected one benchmark result in {path}") + result = results[0] + if not all(report["passed"] for report in result["correctness"]["flex_gemm"]): + raise ValueError(f"correctness failed in {path}") + flex_ms = float(result["flex_gemm"]["median_ms"]) + compiled_eager_ms = float(result["compiled_eager"]["median_ms"]) + return flex_ms, compiled_eager_ms, flex_ms / compiled_eager_ms + + +def _candidate_name(index: int, configs: Sequence[Config]) -> str: + digest = hashlib.sha1( + json.dumps(configs, sort_keys=True).encode("utf-8"), usedforsecurity=False + ).hexdigest()[:10] + return f"candidate_{index:03d}_{digest}" + + +def _run_candidate( + *, + index: int, + configs: tuple[Config, ...], + device_queue: queue.Queue[int], + output_dir: Path, + args: argparse.Namespace, +) -> CandidateResult: + device = device_queue.get() + candidate_name = _candidate_name(index, configs) + result_path = output_dir / f"{candidate_name}.json" + log_path = output_dir / f"{candidate_name}.log" + metadata_path = output_dir / f"{candidate_name}.metadata.json" + cache_path = output_dir / f".{candidate_name}_cache" + try: + if args.resume and result_path.exists() and metadata_path.exists(): + try: + with metadata_path.open(encoding="utf-8") as metadata_file: + metadata = json.load(metadata_file) + median_ms, compiled_eager_ms, normalized_ratio = _result_metrics( + result_path + ) + return CandidateResult( + index=index, + configs=configs, + status="passed", + device=int(metadata["physical_device"]), + median_ms=median_ms, + compiled_eager_ms=compiled_eager_ms, + flex_to_compiled_eager=normalized_ratio, + result_path=str(result_path), + log_path=str(log_path), + ) + except (KeyError, TypeError, ValueError, json.JSONDecodeError): + result_path.unlink() + + command = [ + *_microbenchmark_command(args.suite), + "--case", + args.case, + "--device", + "0", + "--warmup", + str(args.warmup), + "--rounds", + str(args.rounds), + "--iterations", + str(args.iterations), + "--output", + str(result_path), + ] + for config in configs: + command.extend(("--config", json.dumps(config, separators=(",", ":")))) + + environment = os.environ.copy() + environment.update( + { + "CUDA_VISIBLE_DEVICES": str(device), + "TORCH_NATIVE_SKIP_VERSION_CHECK": "1", + "TORCHINDUCTOR_CACHE_DIR": str(cache_path), + } + ) + with log_path.open("w", encoding="utf-8") as log_file: + process = subprocess.Popen( + command, + stdout=log_file, + stderr=subprocess.STDOUT, + env=environment, + start_new_session=True, + text=True, + ) + try: + return_code = process.wait(timeout=args.timeout) + except subprocess.TimeoutExpired: + os.killpg(process.pid, signal.SIGTERM) + try: + process.wait(timeout=10) + except subprocess.TimeoutExpired: + os.killpg(process.pid, signal.SIGKILL) + process.wait() + return CandidateResult( + index=index, + configs=configs, + status="timeout", + device=device, + median_ms=None, + compiled_eager_ms=None, + flex_to_compiled_eager=None, + result_path=str(result_path), + log_path=str(log_path), + error=f"exceeded {args.timeout}s", + ) + if return_code != 0: + return CandidateResult( + index=index, + configs=configs, + status="failed", + device=device, + median_ms=None, + compiled_eager_ms=None, + flex_to_compiled_eager=None, + result_path=str(result_path), + log_path=str(log_path), + error=f"exit code {return_code}", + ) + median_ms, compiled_eager_ms, normalized_ratio = _result_metrics(result_path) + _write_json( + metadata_path, + { + "physical_device": device, + "configs": configs, + "median_ms": median_ms, + "compiled_eager_ms": compiled_eager_ms, + "flex_to_compiled_eager": normalized_ratio, + }, + ) + return CandidateResult( + index=index, + configs=configs, + status="passed", + device=device, + median_ms=median_ms, + compiled_eager_ms=compiled_eager_ms, + flex_to_compiled_eager=normalized_ratio, + result_path=str(result_path), + log_path=str(log_path), + ) + except Exception as error: + return CandidateResult( + index=index, + configs=configs, + status="failed", + device=device, + median_ms=None, + compiled_eager_ms=None, + flex_to_compiled_eager=None, + result_path=str(result_path), + log_path=str(log_path), + error=f"{type(error).__name__}: {error}", + ) + finally: + if not args.keep_cache: + shutil.rmtree(cache_path, ignore_errors=True) + device_queue.put(device) + + +def _write_json(path: Path, value: Any) -> None: + with path.open("w", encoding="utf-8") as output_file: + json.dump(value, output_file, indent=2) + output_file.write("\n") + + +def _tune_flex_gemm( + *, + flex_index: int, + pass_index: int, + current_configs: tuple[Config, ...], + candidates: tuple[Config, ...], + device_queue: queue.Queue[int], + root_output_dir: Path, + args: argparse.Namespace, +) -> tuple[Config, ...]: + output_dir = root_output_dir / f"pass_{pass_index}" / f"flex_{flex_index}" + output_dir.mkdir(parents=True, exist_ok=True) + candidates = _deduplicate_configs((current_configs[flex_index], *candidates)) + jobs = [] + for candidate in candidates: + configs = list(current_configs) + configs[flex_index] = candidate + jobs.append(tuple(configs)) + + results = [] + with ThreadPoolExecutor(max_workers=len(args.devices)) as executor: + futures = { + executor.submit( + _run_candidate, + index=index, + configs=configs, + device_queue=device_queue, + output_dir=output_dir, + args=args, + ): index + for index, configs in enumerate(jobs) + } + for future in as_completed(futures): + result = future.result() + results.append(result) + timing = "" if result.median_ms is None else f" {result.median_ms:.6f} ms" + ratio = ( + "" + if result.flex_to_compiled_eager is None + else f" ratio={result.flex_to_compiled_eager:.6f}" + ) + print( + f"pass={pass_index} flex={flex_index} candidate={result.index} " + f"gpu={result.device} {result.status}{timing}{ratio}", + flush=True, + ) + + results.sort(key=lambda result: result.index) + _write_json( + output_dir / "summary.json", + [dataclasses.asdict(result) for result in results], + ) + passed = [result for result in results if result.flex_to_compiled_eager is not None] + if not passed: + raise RuntimeError( + f"no valid configuration for {args.case} FlexGEMM {flex_index}; " + f"see {output_dir}" + ) + verification_inputs = [] + for device in args.devices: + device_results = sorted( + (result for result in passed if result.device == device), + key=lambda result: result.flex_to_compiled_eager, + ) + verification_inputs.extend(device_results[: args.verify_per_device]) + verification_configs = _deduplicate_config_sets( + tuple(result.configs for result in verification_inputs) + ) + if not verification_configs: + raise RuntimeError("no candidates selected for reference-GPU verification") + + verification_dir = output_dir / "reference_verification" + verification_dir.mkdir(parents=True, exist_ok=True) + reference_device_queue: queue.Queue[int] = queue.Queue() + reference_device_queue.put(args.devices[0]) + verified = [] + for index, configs in enumerate(verification_configs): + attempts = [] + for attempt in range(args.verify_retries): + result = _run_candidate( + index=index + attempt * len(verification_configs), + configs=configs, + device_queue=reference_device_queue, + output_dir=verification_dir, + args=args, + ) + attempts.append(result) + if result.median_ms is not None: + break + result = attempts[-1] + verified.append(result) + timing = "" if result.median_ms is None else f" {result.median_ms:.6f} ms" + ratio = ( + "" + if result.flex_to_compiled_eager is None + else f" ratio={result.flex_to_compiled_eager:.6f}" + ) + print( + f"verify pass={pass_index} flex={flex_index} candidate={index} " + f"gpu={result.device} {result.status}{timing}{ratio}", + flush=True, + ) + _write_json( + verification_dir / "summary.json", + [dataclasses.asdict(result) for result in verified], + ) + verified_passed = [ + result for result in verified if result.flex_to_compiled_eager is not None + ] + if not verified_passed: + raise RuntimeError( + f"all reference-GPU verification runs failed for {args.case} " + f"FlexGEMM {flex_index}; see {verification_dir}" + ) + winner = min( + verified_passed, + key=lambda result: result.median_ms, + ) + assert winner.median_ms is not None + assert winner.flex_to_compiled_eager is not None + print( + f"winner pass={pass_index} flex={flex_index}: " + f"{winner.median_ms:.6f} ms ratio={winner.flex_to_compiled_eager:.6f} " + f"{winner.configs[flex_index]}", + flush=True, + ) + return winner.configs + + +def _run_final( + configs: tuple[Config, ...], + output_dir: Path, + args: argparse.Namespace, +) -> tuple[CandidateResult, ...]: + final_args = argparse.Namespace(**vars(args)) + final_args.warmup = args.final_warmup + final_args.rounds = args.final_rounds + final_args.iterations = args.final_iterations + final_args.timeout = args.final_timeout + final_args.resume = False + device_queue: queue.Queue[int] = queue.Queue() + device_queue.put(args.devices[0]) + attempts = [] + for index in range(args.final_retries): + result = _run_candidate( + index=index, + configs=configs, + device_queue=device_queue, + output_dir=output_dir / "final", + args=final_args, + ) + attempts.append(result) + if result.median_ms is not None: + break + print( + f"final attempt {index + 1}/{args.final_retries} failed: " + f"{result.error}", + flush=True, + ) + return tuple(attempts) + + +def _parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--suite", choices=SUITES, default="671b") + parser.add_argument("--case", required=True) + parser.add_argument( + "--devices", + default="0,1,2,3", + help="comma-separated physical CUDA device indices", + ) + parser.add_argument("--search", choices=("priority", "full"), default="priority") + parser.add_argument( + "--max-candidates", + type=int, + help="limit the selected search space for debugging", + ) + parser.add_argument("--passes", type=int, default=1) + parser.add_argument( + "--flex-index", + type=int, + action="append", + help="tune only this FlexGEMM index; repeat to select multiple indices", + ) + parser.add_argument( + "--base-config", + type=_parse_config, + action="append", + default=[], + help="starting JSON config; repeat once per FlexGEMM", + ) + parser.add_argument("--warmup", type=int, default=3) + parser.add_argument("--rounds", type=int, default=3) + parser.add_argument("--iterations", type=int, default=20) + parser.add_argument("--timeout", type=int, default=300) + parser.add_argument( + "--verify-per-device", + type=int, + default=3, + help="remeasure this many candidates per search GPU on the first device", + ) + parser.add_argument("--verify-retries", type=int, default=3) + parser.add_argument("--final-warmup", type=int, default=25) + parser.add_argument("--final-rounds", type=int, default=10) + parser.add_argument("--final-iterations", type=int, default=200) + parser.add_argument("--final-timeout", type=int, default=900) + parser.add_argument("--final-retries", type=int, default=3) + parser.add_argument( + "--output-dir", + type=Path, + default=Path("outputs/coda_fusion_microbench/autotune"), + ) + parser.add_argument( + "--resume", + action=argparse.BooleanOptionalAction, + default=True, + ) + parser.add_argument("--keep-cache", action="store_true") + parser.add_argument("--skip-final", action="store_true") + args = parser.parse_args() + if args.case not in SUITES[args.suite]: + parser.error( + f"--case must be one of {', '.join(SUITES[args.suite])} " + f"for --suite {args.suite}" + ) + try: + args.devices = tuple(int(device) for device in args.devices.split(",")) + except ValueError as error: + parser.error(f"--devices must contain integers: {error}") + if not args.devices: + parser.error("--devices cannot be empty") + if args.max_candidates is not None and args.max_candidates <= 0: + parser.error("--max-candidates must be positive") + if args.passes <= 0: + parser.error("--passes must be positive") + if args.verify_per_device <= 0: + parser.error("--verify-per-device must be positive") + if args.verify_retries <= 0: + parser.error("--verify-retries must be positive") + if args.final_retries <= 0: + parser.error("--final-retries must be positive") + for name in ("warmup", "rounds", "iterations", "timeout"): + if getattr(args, name) <= 0: + parser.error(f"--{name.replace('_', '-')} must be positive") + return args + + +def main() -> None: + args = _parse_args() + case = SUITES[args.suite][args.case] + current_configs = _initial_configs(args) + candidates = _search_configs(args.search, args.max_candidates) + flex_indices = ( + tuple(args.flex_index) + if args.flex_index is not None + else tuple(range(case.num_flex_gemms)) + ) + if any(index < 0 or index >= case.num_flex_gemms for index in flex_indices): + raise ValueError( + f"FlexGEMM indices for {args.case} must be in " + f"[0, {case.num_flex_gemms})" + ) + + output_dir = args.output_dir / args.case + if args.suite != "671b": + output_dir = args.output_dir / args.suite / args.case + output_dir.mkdir(parents=True, exist_ok=True) + device_queue: queue.Queue[int] = queue.Queue() + for device in args.devices: + device_queue.put(device) + + print( + f"tuning {args.suite}/{args.case}: {len(candidates)} candidates, " + f"FlexGEMMs={flex_indices}, GPUs={args.devices}", + flush=True, + ) + for pass_index in range(args.passes): + for flex_index in flex_indices: + current_configs = _tune_flex_gemm( + flex_index=flex_index, + pass_index=pass_index, + current_configs=current_configs, + candidates=candidates, + device_queue=device_queue, + root_output_dir=output_dir, + args=args, + ) + + _write_json(output_dir / "best_configs.json", current_configs) + if args.skip_final: + return + (output_dir / "final").mkdir(parents=True, exist_ok=True) + final_attempts = _run_final(current_configs, output_dir, args) + _write_json( + output_dir / "final_summary.json", + [dataclasses.asdict(result) for result in final_attempts], + ) + final = final_attempts[-1] + if final.median_ms is None: + raise RuntimeError( + f"final benchmark failed: {final.error}; see {final.log_path}" + ) + print(f"final: {final.median_ms:.6f} ms; result={final.result_path}") + + +if __name__ == "__main__": + main() diff --git a/torchtitan/experiments/graph_trainer/benchmarks/coda_fusion_microbench.py b/torchtitan/experiments/graph_trainer/benchmarks/coda_fusion_microbench.py new file mode 100644 index 0000000000..c4c6a1567f --- /dev/null +++ b/torchtitan/experiments/graph_trainer/benchmarks/coda_fusion_microbench.py @@ -0,0 +1,1322 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +"""Standalone eager-versus-FlexGEMM microbenchmarks for DSV3 CODA patterns. + +Every case uses a shape grounded in the DSV3-671B joint FX graph captured with +local batch 24 and sequence length 4096. The eager functions intentionally use +plain PyTorch operations so the source epilogue remains visible. The FlexGEMM +functions spell out the corresponding epilogue callback and are compiled as a +full graph with Inductor. + +Run one case before attempting the full suite because several cases allocate +multiple gigabytes of inputs:: + + python -m torchtitan.experiments.graph_trainer.benchmarks.coda_fusion_microbench \ + --case f3_attention_output + +Pass ``--config`` once to force one QuACK configuration for every FlexGEMM in +a case, or once per FlexGEMM to tune a multi-GEMM case independently:: + + --config '{"tile_m": 256, "tile_n": 256, "cluster_m": 2, "cluster_n": 1}' +""" + +import argparse +import ast +import dataclasses +import importlib.metadata +import inspect +import json +import re +import statistics +import time +from collections.abc import Callable, Mapping, Sequence +from typing import Any + +import torch +from torch._higher_order_ops import flex_gemm +from torch._inductor.utils import run_and_get_code + + +M = 24 * 4096 +MODEL_WIDTH = 7168 +SHARED_EXPERT_WIDTH = 2048 +RMSNORM_GROUP = 512 +RMSNORM_BACKWARD_GROUP = 128 +EPS = 1e-5 + +TensorTree = torch.Tensor | tuple["TensorTree", ...] | list["TensorTree"] +BenchmarkFn = Callable[..., TensorTree] +InputFactory = Callable[[torch.device], tuple[torch.Tensor, ...]] +FlexFactory = Callable[[tuple[dict[str, Any], ...]], BenchmarkFn] + + +@dataclasses.dataclass(frozen=True) +class BenchmarkCase: + name: str + pattern: str + description: str + shape: str + num_flex_gemms: int + make_inputs: InputFactory + eager: BenchmarkFn + make_flex: FlexFactory + atol: float = 0.05 + rtol: float = 0.02 + fast_math_flex_gemms: tuple[int, ...] = () + + +@dataclasses.dataclass(frozen=True) +class Timing: + median_ms: float + minimum_ms: float + maximum_ms: float + round_means_ms: tuple[float, ...] + samples_ms: tuple[float, ...] + + +def _randn( + shape: Sequence[int], + *, + device: torch.device, + dtype: torch.dtype, + scale: float = 0.02, +) -> torch.Tensor: + return torch.empty(tuple(shape), device=device, dtype=dtype).normal_(std=scale) + + +def _bf16(shape: Sequence[int], device: torch.device) -> torch.Tensor: + return _randn(shape, device=device, dtype=torch.bfloat16) + + +def _fp32(shape: Sequence[int], device: torch.device) -> torch.Tensor: + return _randn(shape, device=device, dtype=torch.float32) + + +def eager_b1_lm_head_input_grad_cast( + grad: torch.Tensor, + weight: torch.Tensor, +) -> torch.Tensor: + input_grad = torch.mm(grad, weight) + return input_grad.float() + + +def make_flex_b1_lm_head_input_grad_cast( + options: tuple[dict[str, Any], ...], +) -> BenchmarkFn: + def flex_b1_lm_head_input_grad_cast( + grad: torch.Tensor, + weight: torch.Tensor, + ) -> torch.Tensor: + def epilogue(accumulator: torch.Tensor) -> torch.Tensor: + return accumulator.float().bfloat16().float() + + return flex_gemm( + torch.mm, + (grad, weight), + epilogue, + kernel_options=options[0], + ) + + return flex_b1_lm_head_input_grad_cast + + +def eager_b6_weight_grad_cast( + lhs: torch.Tensor, + rhs: torch.Tensor, +) -> torch.Tensor: + weight_grad = torch.mm(lhs, rhs) + return weight_grad.float() + + +def make_flex_b6_weight_grad_cast( + options: tuple[dict[str, Any], ...], +) -> BenchmarkFn: + def flex_b6_weight_grad_cast( + lhs: torch.Tensor, + rhs: torch.Tensor, + ) -> torch.Tensor: + def epilogue(accumulator: torch.Tensor) -> torch.Tensor: + return accumulator.float().bfloat16().float() + + return flex_gemm( + torch.mm, + (lhs, rhs), + epilogue, + kernel_options=options[0], + ) + + return flex_b6_weight_grad_cast + + +def eager_f6_router_sigmoid_bias( + tokens: torch.Tensor, + weight: torch.Tensor, + expert_bias: torch.Tensor, +) -> tuple[torch.Tensor, torch.Tensor]: + raw_scores = torch.sigmoid(torch.mm(tokens, weight)) + return raw_scores, raw_scores + expert_bias + + +def make_flex_f6_router_sigmoid_bias( + options: tuple[dict[str, Any], ...], +) -> BenchmarkFn: + def flex_f6_router_sigmoid_bias( + tokens: torch.Tensor, + weight: torch.Tensor, + expert_bias: torch.Tensor, + ) -> tuple[torch.Tensor, torch.Tensor]: + bias_2d = expert_bias.view(1, -1) + + def epilogue( + accumulator: torch.Tensor, + ) -> tuple[torch.Tensor, torch.Tensor]: + raw_scores = torch.sigmoid(accumulator) + return raw_scores, raw_scores + bias_2d + + return flex_gemm( + torch.mm, + (tokens, weight), + epilogue, + kernel_options=options[0], + ) + + return flex_f6_router_sigmoid_bias + + +def eager_f4_shared_expert_swiglu( + tokens: torch.Tensor, + w1: torch.Tensor, + w3: torch.Tensor, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + w1_output = torch.mm(tokens, w1) + activated = torch.nn.functional.silu(w1_output) + gate = torch.mm(tokens, w3) + return activated, gate, activated * gate + + +def make_flex_f4_shared_expert_swiglu( + options: tuple[dict[str, Any], ...], +) -> BenchmarkFn: + def flex_f4_shared_expert_swiglu( + tokens: torch.Tensor, + w1: torch.Tensor, + w3: torch.Tensor, + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + def silu_epilogue(accumulator: torch.Tensor) -> torch.Tensor: + rounded = accumulator.float().bfloat16() + return torch.nn.functional.silu(rounded) + + activated = flex_gemm( + torch.mm, + (tokens, w1), + silu_epilogue, + kernel_options=options[0], + ) + + def gate_epilogue( + accumulator: torch.Tensor, + ) -> tuple[torch.Tensor, torch.Tensor]: + gate = accumulator.float().bfloat16() + return gate, activated * gate + + gate, product = flex_gemm( + torch.mm, + (tokens, w3), + gate_epilogue, + kernel_options=options[1], + ) + return activated, gate, product + + return flex_f4_shared_expert_swiglu + + +def eager_b2_shared_expert_swiglu_backward( + output_grad: torch.Tensor, + w2: torch.Tensor, + saved_silu: torch.Tensor, + saved_gate: torch.Tensor, + saved_w1: torch.Tensor, + w3: torch.Tensor, + w1: torch.Tensor, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + branch_grad = torch.mm(output_grad, w2) + gate_grad = branch_grad * saved_silu + silu_grad = torch.ops.aten.silu_backward.default( + branch_grad * saved_gate, + saved_w1, + ) + w3_input_grad = torch.mm(gate_grad, w3) + w1_input_grad = torch.mm(silu_grad, w1) + return gate_grad, silu_grad, w3_input_grad + w1_input_grad + + +def make_flex_b2_shared_expert_swiglu_backward( + options: tuple[dict[str, Any], ...], +) -> BenchmarkFn: + def flex_b2_shared_expert_swiglu_backward( + output_grad: torch.Tensor, + w2: torch.Tensor, + saved_silu: torch.Tensor, + saved_gate: torch.Tensor, + saved_w1: torch.Tensor, + w3: torch.Tensor, + w1: torch.Tensor, + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + def branch_epilogue( + accumulator: torch.Tensor, + ) -> tuple[torch.Tensor, torch.Tensor]: + branch_grad = accumulator.float().bfloat16() + gate_grad = branch_grad * saved_silu + silu_grad = torch.ops.aten.silu_backward.default( + branch_grad * saved_gate, + saved_w1, + ) + return gate_grad, silu_grad + + gate_grad, silu_grad = flex_gemm( + torch.mm, + (output_grad, w2), + branch_epilogue, + kernel_options=options[0], + ) + w3_input_grad = torch.mm(gate_grad, w3) + + def input_grad_epilogue(accumulator: torch.Tensor) -> torch.Tensor: + w1_input_grad = accumulator.float().bfloat16() + return w3_input_grad + w1_input_grad + + input_grad = flex_gemm( + torch.mm, + (silu_grad, w1), + input_grad_epilogue, + kernel_options=options[1], + ) + return gate_grad, silu_grad, input_grad + + return flex_b2_shared_expert_swiglu_backward + + +def eager_b4_router_input_grad_add( + score_grad: torch.Tensor, + router_weight: torch.Tensor, + expert_input_grad: torch.Tensor, +) -> torch.Tensor: + router_input_grad = torch.mm(score_grad, router_weight).bfloat16() + return expert_input_grad + router_input_grad + + +def make_flex_b4_router_input_grad_add( + options: tuple[dict[str, Any], ...], +) -> BenchmarkFn: + def flex_b4_router_input_grad_add( + score_grad: torch.Tensor, + router_weight: torch.Tensor, + expert_input_grad: torch.Tensor, + ) -> torch.Tensor: + def epilogue(accumulator: torch.Tensor) -> torch.Tensor: + router_input_grad = accumulator.bfloat16() + return expert_input_grad + router_input_grad + + return flex_gemm( + torch.mm, + (score_grad, router_weight), + epilogue, + kernel_options=options[0], + ) + + return flex_b4_router_input_grad_add + + +def eager_b5_mla_rmsnorm_backward( + output_grad: torch.Tensor, + projection_weight: torch.Tensor, + norm_input: torch.Tensor, + rstd: torch.Tensor, + gamma: torch.Tensor, +) -> tuple[torch.Tensor, torch.Tensor]: + grad = torch.mm(output_grad, projection_weight) + return torch.ops.aten._fused_rms_norm_backward.default( + grad, + norm_input, + [grad.shape[-1]], + rstd, + gamma, + [True, True], + ) + + +def make_flex_b5_mla_rmsnorm_backward( + options: tuple[dict[str, Any], ...], +) -> BenchmarkFn: + def flex_b5_mla_rmsnorm_backward( + output_grad: torch.Tensor, + projection_weight: torch.Tensor, + norm_input: torch.Tensor, + rstd: torch.Tensor, + gamma: torch.Tensor, + ) -> tuple[torch.Tensor, torch.Tensor]: + gamma_2d = gamma.view(1, -1) + + def epilogue( + accumulator: torch.Tensor, + ) -> tuple[torch.Tensor, torch.Tensor]: + rounded = accumulator.float().bfloat16() + rounded_fp32 = rounded.float() + x_hat = norm_input.float() * rstd + grad_x_hat = rounded_fp32 * gamma_2d.float() + row_products = x_hat * grad_x_hat + partial_row_dot = row_products.view( + row_products.shape[0], + -1, + RMSNORM_BACKWARD_GROUP, + ).sum(-1) + return rounded, partial_row_dot + + rounded, partial_row_dot = flex_gemm( + torch.mm, + (output_grad, projection_weight), + epilogue, + kernel_options=options[0], + ) + rounded_fp32 = rounded.float() + x_hat = norm_input.float() * rstd + grad_x_hat = rounded_fp32 * gamma_2d.float() + row_dot = partial_row_dot.sum(-1, keepdim=True) + correction = (x_hat / rounded.shape[-1]) * row_dot + grad_input = ((grad_x_hat - correction) * rstd).bfloat16() + grad_weight = (rounded_fp32 * x_hat).sum(0).bfloat16() + return grad_input, grad_weight + + return flex_b5_mla_rmsnorm_backward + + +def eager_b7_attention_input_grad_merge( + kv_grad: torch.Tensor, + kv_weight: torch.Tensor, + q_grad: torch.Tensor, + q_weight: torch.Tensor, +) -> torch.Tensor: + kv_input_grad = torch.mm(kv_grad, kv_weight) + q_input_grad = torch.mm(q_grad, q_weight) + return kv_input_grad + q_input_grad + + +def make_flex_b7_attention_input_grad_merge( + options: tuple[dict[str, Any], ...], +) -> BenchmarkFn: + def flex_b7_attention_input_grad_merge( + kv_grad: torch.Tensor, + kv_weight: torch.Tensor, + q_grad: torch.Tensor, + q_weight: torch.Tensor, + ) -> torch.Tensor: + kv_input_grad = torch.mm(kv_grad, kv_weight) + + def epilogue(accumulator: torch.Tensor) -> torch.Tensor: + q_input_grad = accumulator.float().bfloat16() + return kv_input_grad + q_input_grad + + return flex_gemm( + torch.mm, + (q_grad, q_weight), + epilogue, + kernel_options=options[0], + ) + + return flex_b7_attention_input_grad_merge + + +def eager_f2_q_rmsnorm( + tokens: torch.Tensor, + wq_a: torch.Tensor, + gamma: torch.Tensor, + wq_b: torch.Tensor, +) -> torch.Tensor: + q_low_rank = torch.mm(tokens, wq_a) + normalized = torch.nn.functional.rms_norm( + q_low_rank, + (q_low_rank.shape[-1],), + gamma, + EPS, + ) + return torch.mm(normalized, wq_b) + + +def make_flex_f2_q_rmsnorm( + options: tuple[dict[str, Any], ...], +) -> BenchmarkFn: + def flex_f2_q_rmsnorm( + tokens: torch.Tensor, + wq_a: torch.Tensor, + gamma: torch.Tensor, + wq_b: torch.Tensor, + ) -> torch.Tensor: + gamma_2d = gamma.view(1, -1) + + def first_epilogue( + accumulator: torch.Tensor, + ) -> tuple[torch.Tensor, torch.Tensor]: + rounded = accumulator.float().bfloat16() + rounded_fp32 = rounded.float() + weighted = (rounded_fp32 * gamma_2d).bfloat16() + partial_mean_square = ( + rounded_fp32.view(rounded.shape[0], -1, RMSNORM_GROUP).square().mean(-1) + ) + return weighted, partial_mean_square + + weighted, partial_mean_square = flex_gemm( + torch.mm, + (tokens, wq_a), + first_epilogue, + kernel_options=options[0], + ) + rstd = (partial_mean_square.mean(-1, keepdim=True) + EPS).rsqrt() + + def second_epilogue(accumulator: torch.Tensor) -> torch.Tensor: + return (accumulator.float() * rstd).bfloat16() + + return flex_gemm( + torch.mm, + (weighted, wq_b), + second_epilogue, + kernel_options=options[1], + ) + + return flex_f2_q_rmsnorm + + +def eager_f2_kv_rmsnorm( + tokens: torch.Tensor, + wkv_a: torch.Tensor, + gamma: torch.Tensor, + wkv_b: torch.Tensor, +) -> tuple[torch.Tensor, torch.Tensor]: + kv_low_rank = torch.mm(tokens, wkv_a) + active = kv_low_rank[:, :512] + rope_tail = kv_low_rank[:, 512:] + normalized = torch.nn.functional.rms_norm(active, (512,), gamma, EPS) + return torch.mm(normalized, wkv_b), rope_tail + + +def make_flex_f2_kv_rmsnorm( + options: tuple[dict[str, Any], ...], +) -> BenchmarkFn: + def flex_f2_kv_rmsnorm( + tokens: torch.Tensor, + wkv_a: torch.Tensor, + gamma: torch.Tensor, + wkv_b: torch.Tensor, + ) -> tuple[torch.Tensor, torch.Tensor]: + gamma_full = torch.nn.functional.pad( + gamma.view(1, 512), + (0, 64), + value=1.0, + ) + + def first_epilogue( + accumulator: torch.Tensor, + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + raw = accumulator.float().bfloat16() + raw_fp32 = raw.float() + weighted_full = (raw_fp32 * gamma_full).bfloat16() + partial_mean_square = raw_fp32.view(raw.shape[0], -1, 64).square().mean(-1) + return weighted_full, raw, partial_mean_square + + weighted_full, raw, partial_mean_square = flex_gemm( + torch.mm, + (tokens, wkv_a), + first_epilogue, + kernel_options=options[0], + ) + weighted = weighted_full[:, :512] + active_partials = partial_mean_square[:, :8] + rstd = (active_partials.mean(-1, keepdim=True) + EPS).rsqrt() + + def second_epilogue(accumulator: torch.Tensor) -> torch.Tensor: + return (accumulator.float() * rstd).bfloat16() + + output = flex_gemm( + torch.mm, + (weighted, wkv_b), + second_epilogue, + kernel_options=options[1], + ) + return output, raw[:, 512:] + + return flex_f2_kv_rmsnorm + + +def eager_f3_attention_output( + attention: torch.Tensor, + wo: torch.Tensor, + residual: torch.Tensor, + gamma: torch.Tensor, +) -> tuple[torch.Tensor, torch.Tensor]: + attention_output = torch.mm(attention, wo) + total = residual + attention_output + normalized = torch.nn.functional.rms_norm( + total, + (total.shape[-1],), + gamma, + EPS, + ) + return normalized, total + + +def make_flex_f3_attention_output( + options: tuple[dict[str, Any], ...], +) -> BenchmarkFn: + def flex_f3_attention_output( + attention: torch.Tensor, + wo: torch.Tensor, + residual: torch.Tensor, + gamma: torch.Tensor, + ) -> tuple[torch.Tensor, torch.Tensor]: + def epilogue( + accumulator: torch.Tensor, + ) -> tuple[torch.Tensor, torch.Tensor]: + attention_output = accumulator.float().bfloat16() + total = residual + attention_output + partial_mean_square = ( + total.float().view(total.shape[0], -1, RMSNORM_GROUP).square().mean(-1) + ) + return total, partial_mean_square + + total, partial_mean_square = flex_gemm( + torch.mm, + (attention, wo), + epilogue, + kernel_options=options[0], + ) + rstd = (partial_mean_square.mean(-1, keepdim=True) + EPS).rsqrt() + normalized = (total.float() * rstd * gamma.float()).bfloat16() + return normalized, total + + return flex_f3_attention_output + + +def eager_f3_moe_output( + shared_activation: torch.Tensor, + shared_w2: torch.Tensor, + routed_output: torch.Tensor, + residual: torch.Tensor, + gamma: torch.Tensor, +) -> tuple[torch.Tensor, torch.Tensor]: + shared_output = torch.mm(shared_activation, shared_w2) + moe_output = routed_output + shared_output + total = residual + moe_output + normalized = torch.nn.functional.rms_norm( + total, + (total.shape[-1],), + gamma, + EPS, + ) + return normalized, total + + +def make_flex_f3_moe_output( + options: tuple[dict[str, Any], ...], +) -> BenchmarkFn: + def flex_f3_moe_output( + shared_activation: torch.Tensor, + shared_w2: torch.Tensor, + routed_output: torch.Tensor, + residual: torch.Tensor, + gamma: torch.Tensor, + ) -> tuple[torch.Tensor, torch.Tensor]: + def epilogue( + accumulator: torch.Tensor, + ) -> tuple[torch.Tensor, torch.Tensor]: + shared_output = accumulator.float().bfloat16() + moe_output = routed_output + shared_output + total = residual + moe_output + partial_mean_square = ( + total.float().view(total.shape[0], -1, RMSNORM_GROUP).square().mean(-1) + ) + return total, partial_mean_square + + total, partial_mean_square = flex_gemm( + torch.mm, + (shared_activation, shared_w2), + epilogue, + kernel_options=options[0], + ) + rstd = (partial_mean_square.mean(-1, keepdim=True) + EPS).rsqrt() + normalized = (total.float() * rstd * gamma.float()).bfloat16() + return normalized, total + + return flex_f3_moe_output + + +def _make_b1_inputs(device: torch.device) -> tuple[torch.Tensor, ...]: + return _bf16((12288, 129280), device), _bf16((129280, MODEL_WIDTH), device) + + +def _make_b6_inputs(device: torch.device) -> tuple[torch.Tensor, ...]: + return _bf16((SHARED_EXPERT_WIDTH, M), device), _bf16((M, MODEL_WIDTH), device) + + +def _make_f6_inputs(device: torch.device) -> tuple[torch.Tensor, ...]: + return ( + _fp32((M, MODEL_WIDTH), device), + _fp32((MODEL_WIDTH, 256), device), + _fp32((256,), device), + ) + + +def _make_f4_inputs(device: torch.device) -> tuple[torch.Tensor, ...]: + return ( + _bf16((M, MODEL_WIDTH), device), + _bf16((MODEL_WIDTH, SHARED_EXPERT_WIDTH), device), + _bf16((MODEL_WIDTH, SHARED_EXPERT_WIDTH), device), + ) + + +def _make_b2_inputs(device: torch.device) -> tuple[torch.Tensor, ...]: + return ( + _bf16((M, MODEL_WIDTH), device), + _bf16((MODEL_WIDTH, SHARED_EXPERT_WIDTH), device), + _bf16((M, SHARED_EXPERT_WIDTH), device), + _bf16((M, SHARED_EXPERT_WIDTH), device), + _bf16((M, SHARED_EXPERT_WIDTH), device), + _bf16((SHARED_EXPERT_WIDTH, MODEL_WIDTH), device), + _bf16((SHARED_EXPERT_WIDTH, MODEL_WIDTH), device), + ) + + +def _make_b4_inputs(device: torch.device) -> tuple[torch.Tensor, ...]: + return ( + _fp32((M, 256), device), + _fp32((256, MODEL_WIDTH), device), + _bf16((M, MODEL_WIDTH), device), + ) + + +def _make_b5_inputs(device: torch.device) -> tuple[torch.Tensor, ...]: + return ( + _bf16((M, 24576), device), + _bf16((24576, 1536), device), + _bf16((M, 1536), device), + torch.empty((M, 1), device=device, dtype=torch.float32).uniform_(0.5, 1.5), + _bf16((1536,), device), + ) + + +def _make_b7_inputs(device: torch.device) -> tuple[torch.Tensor, ...]: + return ( + _bf16((M, 576), device), + _bf16((576, MODEL_WIDTH), device), + _bf16((M, 1536), device), + _bf16((1536, MODEL_WIDTH), device), + ) + + +def _make_f2_q_inputs(device: torch.device) -> tuple[torch.Tensor, ...]: + return ( + _bf16((M, MODEL_WIDTH), device), + _bf16((MODEL_WIDTH, 1536), device), + torch.ones((1536,), device=device, dtype=torch.bfloat16), + _bf16((1536, 24576), device), + ) + + +def _make_f2_kv_inputs(device: torch.device) -> tuple[torch.Tensor, ...]: + return ( + _bf16((M, MODEL_WIDTH), device), + _bf16((MODEL_WIDTH, 576), device), + torch.ones((512,), device=device, dtype=torch.bfloat16), + _bf16((512, 32768), device), + ) + + +def _make_f3_attention_inputs(device: torch.device) -> tuple[torch.Tensor, ...]: + return ( + _bf16((M, 16384), device), + _bf16((16384, MODEL_WIDTH), device), + _bf16((M, MODEL_WIDTH), device), + torch.ones((MODEL_WIDTH,), device=device, dtype=torch.bfloat16), + ) + + +def _make_f3_moe_inputs(device: torch.device) -> tuple[torch.Tensor, ...]: + return ( + _bf16((M, SHARED_EXPERT_WIDTH), device), + _bf16((SHARED_EXPERT_WIDTH, MODEL_WIDTH), device), + _bf16((M, MODEL_WIDTH), device), + _bf16((M, MODEL_WIDTH), device), + torch.ones((MODEL_WIDTH,), device=device, dtype=torch.bfloat16), + ) + + +CASES: dict[str, BenchmarkCase] = { + case.name: case + for case in ( + BenchmarkCase( + "b1_lm_head_input_grad_cast", + "B1", + "LM-head input-gradient BF16 store followed by FP32 cast", + "(12288, 129280) @ (129280, 7168)", + 1, + _make_b1_inputs, + eager_b1_lm_head_input_grad_cast, + make_flex_b1_lm_head_input_grad_cast, + ), + BenchmarkCase( + "b2_shared_expert_swiglu_backward", + "B2", + "Shared-expert SwiGLU branch derivatives and input-gradient merge", + "M=98304, D=7168, P=2048; three BF16 GEMMs", + 2, + _make_b2_inputs, + eager_b2_shared_expert_swiglu_backward, + make_flex_b2_shared_expert_swiglu_backward, + 3.1e-5, + 1e-2, + ), + BenchmarkCase( + "b4_router_input_grad_add", + "B4", + "FP32 router input-gradient GEMM, BF16 store, and expert-gradient add", + "(98304, 256) @ (256, 7168)", + 1, + _make_b4_inputs, + eager_b4_router_input_grad_add, + make_flex_b4_router_input_grad_add, + 4.9e-4, + 1e-2, + ), + BenchmarkCase( + "b5_mla_q_rmsnorm_backward", + "B5", + "MLA Q input-gradient projection plus RMSNorm backward", + "(98304, 24576) @ (24576, 1536)", + 1, + _make_b5_inputs, + eager_b5_mla_rmsnorm_backward, + make_flex_b5_mla_rmsnorm_backward, + 0.05, + 0.02, + ), + BenchmarkCase( + "b6_weight_grad_cast", + "B6", + "Frequent shared-expert BF16 weight-gradient GEMM and FP32 cast", + "(2048, 98304) @ (98304, 7168)", + 1, + _make_b6_inputs, + eager_b6_weight_grad_cast, + make_flex_b6_weight_grad_cast, + ), + BenchmarkCase( + "b7_attention_input_grad_merge", + "B7", + "KV and Q input-gradient GEMMs followed by BF16 add", + "(98304, 576) @ (576, 7168) + (98304, 1536) @ (1536, 7168)", + 1, + _make_b7_inputs, + eager_b7_attention_input_grad_merge, + make_flex_b7_attention_input_grad_merge, + ), + BenchmarkCase( + "f2_q_rmsnorm", + "F2-Q", + "MLA Q low-rank projection, RMSNorm, and expanded projection", + "7168 -> 1536 -> 24576 at M=98304", + 2, + _make_f2_q_inputs, + eager_f2_q_rmsnorm, + make_flex_f2_q_rmsnorm, + 3.2e-2, + 2e-2, + ), + BenchmarkCase( + "f2_kv_rmsnorm", + "F2-KV", + "Segmented MLA KV projection, RMSNorm, and expanded projection", + "7168 -> 576, RMSNorm(512), 512 -> 32768 at M=98304", + 2, + _make_f2_kv_inputs, + eager_f2_kv_rmsnorm, + make_flex_f2_kv_rmsnorm, + 1.6e-2, + 2e-2, + ), + BenchmarkCase( + "f3_attention_output", + "F3-A", + "Attention WO projection, residual add, and FFN RMSNorm", + "(98304, 16384) @ (16384, 7168)", + 1, + _make_f3_attention_inputs, + eager_f3_attention_output, + make_flex_f3_attention_output, + 0.07, + 2e-2, + ), + BenchmarkCase( + "f3_moe_output", + "F3-B", + "Shared W2 projection, routed add, residual add, and next RMSNorm", + "(98304, 2048) @ (2048, 7168)", + 1, + _make_f3_moe_inputs, + eager_f3_moe_output, + make_flex_f3_moe_output, + 0.07, + 2e-2, + ), + BenchmarkCase( + "f4_shared_expert_swiglu", + "F4", + "Shared-expert W1/W3 GEMMs with SiLU and multiply", + "two (98304, 7168) @ (7168, 2048) GEMMs", + 2, + _make_f4_inputs, + eager_f4_shared_expert_swiglu, + make_flex_f4_shared_expert_swiglu, + fast_math_flex_gemms=(0,), + ), + BenchmarkCase( + "f6_router_sigmoid_bias", + "F6", + "FP32 router GEMM with sigmoid and expert bias", + "(98304, 7168) @ (7168, 256)", + 1, + _make_f6_inputs, + eager_f6_router_sigmoid_bias, + make_flex_f6_router_sigmoid_bias, + 1.5e-5, + 1e-4, + fast_math_flex_gemms=(0,), + ), + ) +} + + +def _kernel_options( + case: BenchmarkCase, + *, + configs: Sequence[str], +) -> tuple[dict[str, Any], ...]: + parsed_configs = tuple(json.loads(config) for config in configs) + if any(not isinstance(config, dict) for config in parsed_configs): + raise ValueError("every --config value must decode to a JSON object") + if len(parsed_configs) not in (0, 1, case.num_flex_gemms): + raise ValueError( + f"{case.name} has {case.num_flex_gemms} FlexGEMMs; pass zero, one, " + f"or {case.num_flex_gemms} --config values" + ) + if len(parsed_configs) == 1: + parsed_configs = parsed_configs * case.num_flex_gemms + if not parsed_configs: + parsed_configs = ({},) * case.num_flex_gemms + + options = [] + for index, config in enumerate(parsed_configs): + value: dict[str, Any] = {"backend": "QUACK", "tuned": True} + if index in case.fast_math_flex_gemms: + value["fast_math"] = True + if config: + value["config"] = config + options.append(value) + return tuple(options) + + +def _flatten(value: TensorTree) -> list[torch.Tensor]: + if isinstance(value, torch.Tensor): + return [value] + tensors = [] + for child in value: + tensors.extend(_flatten(child)) + return tensors + + +def _correctness( + actual: TensorTree, + expected: TensorTree, + *, + atol: float, + rtol: float, +) -> list[dict[str, Any]]: + actual_tensors = _flatten(actual) + expected_tensors = _flatten(expected) + if len(actual_tensors) != len(expected_tensors): + raise AssertionError( + f"output count differs: {len(actual_tensors)} != {len(expected_tensors)}" + ) + + reports = [] + for index, (actual_tensor, expected_tensor) in enumerate( + zip(actual_tensors, expected_tensors, strict=True) + ): + if actual_tensor.shape != expected_tensor.shape: + raise AssertionError( + f"output {index} shape differs: " + f"{actual_tensor.shape} != {expected_tensor.shape}" + ) + difference = (actual_tensor.float() - expected_tensor.float()).abs() + max_abs = difference.max().item() + mean_abs = difference.mean().item() + allowed = atol + rtol * expected_tensor.float().abs() + passed = (difference <= allowed).all().item() + relative_l2 = ( + (difference.square().sum() / expected_tensor.float().square().sum()) + .sqrt() + .item() + ) + reports.append( + { + "output": index, + "shape": list(actual_tensor.shape), + "dtype": str(actual_tensor.dtype), + "max_abs": max_abs, + "mean_abs": mean_abs, + "relative_l2": relative_l2, + "atol": atol, + "rtol": rtol, + "passed": passed, + } + ) + if not passed: + raise AssertionError( + f"output {index} violates atol={atol}, rtol={rtol}; " + f"max_abs={max_abs}, relative_l2={relative_l2}" + ) + return reports + + +def _benchmark( + fn: BenchmarkFn, + inputs: tuple[torch.Tensor, ...], + *, + warmup: int, + iterations: int, + flush_cache_mb: int, +) -> Timing: + cache = None + if flush_cache_mb: + cache = torch.empty( + flush_cache_mb * 1024 * 1024 // 4, + device=inputs[0].device, + dtype=torch.float32, + ) + + for _ in range(warmup): + if cache is not None: + cache.zero_() + fn(*inputs) + torch.cuda.synchronize(inputs[0].device) + + starts = [torch.cuda.Event(enable_timing=True) for _ in range(iterations)] + ends = [torch.cuda.Event(enable_timing=True) for _ in range(iterations)] + for start, end in zip(starts, ends, strict=True): + if cache is not None: + cache.zero_() + start.record() + fn(*inputs) + end.record() + torch.cuda.synchronize(inputs[0].device) + samples = tuple(start.elapsed_time(end) for start, end in zip(starts, ends)) + return Timing( + median_ms=statistics.median(samples), + minimum_ms=min(samples), + maximum_ms=max(samples), + round_means_ms=(statistics.mean(samples),), + samples_ms=samples, + ) + + +def _make_cuda_graph_replay( + fn: BenchmarkFn, + inputs: tuple[torch.Tensor, ...], + *, + warmup: int, +) -> Callable[[], None]: + device = inputs[0].device + warmup_stream = torch.cuda.Stream(device=device) + warmup_stream.wait_stream(torch.cuda.current_stream(device)) + with torch.cuda.stream(warmup_stream): + for _ in range(max(1, warmup)): + fn(*inputs) + torch.cuda.current_stream(device).wait_stream(warmup_stream) + torch.cuda.synchronize(device) + + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + static_output = fn(*inputs) + + def replay() -> None: + graph.replay() + # Keep graph-owned output storage alive for the lifetime of the replay. + _ = static_output + + return replay + + +def _benchmark_cuda_graphs( + candidates: tuple[tuple[str, Callable[[], None]], ...], + *, + device: torch.device, + rounds: int, + iterations: int, +) -> dict[str, Timing]: + round_samples: dict[str, list[tuple[float, ...]]] = { + name: [] for name, _ in candidates + } + for round_index in range(rounds): + ordered = candidates if round_index % 2 == 0 else candidates[::-1] + for name, replay in ordered: + start = torch.cuda.Event(enable_timing=True) + end = torch.cuda.Event(enable_timing=True) + start.record() + for _ in range(iterations): + replay() + end.record() + torch.cuda.synchronize(device) + elapsed_ms = start.elapsed_time(end) + round_samples[name].append((elapsed_ms / iterations,)) + + timings = {} + for name, samples_by_round in round_samples.items(): + samples = tuple(value for group in samples_by_round for value in group) + round_means = tuple(statistics.mean(group) for group in samples_by_round) + timings[name] = Timing( + median_ms=statistics.median(round_means), + minimum_ms=min(samples), + maximum_ms=max(samples), + round_means_ms=round_means, + samples_ms=samples, + ) + return timings + + +def _run_case( + case: BenchmarkCase, + args: argparse.Namespace, +) -> dict[str, Any]: + device = torch.device(f"cuda:{args.device}") + torch.cuda.set_device(device) + torch.manual_seed(args.seed) + inputs = case.make_inputs(device) + options = _kernel_options( + case, + configs=args.config, + ) + flex = case.make_flex(options) + + eager_compile_start = time.perf_counter() + compiled_eager = torch.compile(case.eager, backend="inductor", fullgraph=True) + compiled_expected = compiled_eager(*inputs) + torch.cuda.synchronize(device) + eager_compile_seconds = time.perf_counter() - eager_compile_start + + flex_compile_start = time.perf_counter() + compiled_flex = torch.compile(flex, backend="inductor", fullgraph=True) + actual, generated_sources = run_and_get_code(compiled_flex, *inputs) + torch.cuda.synchronize(device) + flex_compile_seconds = time.perf_counter() - flex_compile_start + + generated_source = "\n".join(generated_sources) + num_flex_gemm_calls = generated_source.count("flex_gemm_epilogue(") + if num_flex_gemm_calls != case.num_flex_gemms: + raise RuntimeError( + f"expected {case.num_flex_gemms} generated FlexGEMM calls, " + f"found {num_flex_gemm_calls}" + ) + selected_configs = tuple( + dict(ast.literal_eval(config_key)) + for config_key in re.findall( + r"config_key=(\(.*?\)), config_is_lowering_validated", + generated_source, + ) + ) + if len(selected_configs) != case.num_flex_gemms: + raise RuntimeError( + f"expected {case.num_flex_gemms} generated QuACK configs, " + f"found {len(selected_configs)}" + ) + + expected = case.eager(*inputs) + torch.cuda.synchronize(device) + flex_correctness = _correctness( + actual, + expected, + atol=case.atol, + rtol=case.rtol, + ) + compiled_eager_correctness = _correctness( + compiled_expected, + expected, + atol=case.atol, + rtol=case.rtol, + ) + + source_eager_timing = _benchmark( + case.eager, + inputs, + warmup=args.warmup, + iterations=args.iterations, + flush_cache_mb=args.flush_cache_mb, + ) + compiled_eager_replay = _make_cuda_graph_replay( + compiled_eager, + inputs, + warmup=args.warmup, + ) + flex_replay = _make_cuda_graph_replay( + compiled_flex, + inputs, + warmup=args.warmup, + ) + replay_timings = _benchmark_cuda_graphs( + ( + ("compiled_eager", compiled_eager_replay), + ("flex_gemm", flex_replay), + ), + device=device, + rounds=args.rounds, + iterations=args.iterations, + ) + compiled_eager_timing = replay_timings["compiled_eager"] + flex_timing = replay_timings["flex_gemm"] + result = { + "case": case.name, + "pattern": case.pattern, + "shape": case.shape, + "device": torch.cuda.get_device_name(device), + "device_capability": list(torch.cuda.get_device_capability(device)), + "torch_version": torch.__version__, + "torch_git_version": torch.version.git_version, + "cutlass_dsl_version": importlib.metadata.version("nvidia-cutlass-dsl"), + "triton_version": importlib.metadata.version("triton"), + "kernel_options": options, + "selected_quack_configs": selected_configs, + "benchmark_contract": { + "primary": "fixed-pointer CUDA-graph replay", + "rounds": args.rounds, + "iterations_per_round": args.iterations, + "alternating_candidate_order": True, + }, + "compile_seconds": { + "compiled_eager": eager_compile_seconds, + "flex_gemm": flex_compile_seconds, + }, + "generated_flex_gemm_calls": num_flex_gemm_calls, + "correctness": { + "compiled_eager": compiled_eager_correctness, + "flex_gemm": flex_correctness, + }, + "source_eager": dataclasses.asdict(source_eager_timing), + "compiled_eager": dataclasses.asdict(compiled_eager_timing), + "flex_gemm": dataclasses.asdict(flex_timing), + "speedup_vs_compiled_eager": ( + compiled_eager_timing.median_ms / flex_timing.median_ms + ), + "speedup_vs_source_eager": ( + source_eager_timing.median_ms / flex_timing.median_ms + ), + } + print(json.dumps(result, indent=2)) + return result + + +def _show_cases(cases: Mapping[str, BenchmarkCase] = CASES) -> None: + print("case | pattern | FlexGEMMs | real shape") + print("--- | --- | ---: | ---") + for case in cases.values(): + print(f"{case.name} | {case.pattern} | {case.num_flex_gemms} | {case.shape}") + + +def _show_source(case: BenchmarkCase) -> None: + print(f"# {case.name}: eager") + print(inspect.getsource(case.eager)) + print(f"# {case.name}: FlexGEMM factory") + print(inspect.getsource(case.make_flex)) + + +def _parse_args( + cases: Mapping[str, BenchmarkCase] = CASES, + *, + description: str | None = None, +) -> argparse.Namespace: + parser = argparse.ArgumentParser(description=description or __doc__) + parser.add_argument("--case", choices=(*cases, "all")) + parser.add_argument("--list", action="store_true", help="list all cases") + parser.add_argument( + "--show-source", + action="store_true", + help="print the selected eager and FlexGEMM implementations", + ) + parser.add_argument("--device", type=int, default=0) + parser.add_argument("--warmup", type=int, default=5) + parser.add_argument("--rounds", type=int, default=5) + parser.add_argument("--iterations", type=int, default=20) + parser.add_argument("--seed", type=int, default=42) + parser.add_argument( + "--config", + action="append", + default=[], + help="JSON QuACK config; repeat once per FlexGEMM for independent tuning", + ) + parser.add_argument( + "--flush-cache-mb", + type=int, + default=0, + help="zero this many MiB before every timed call", + ) + parser.add_argument("--output", help="write all JSON results to this path") + args = parser.parse_args() + if not args.list and args.case is None: + parser.error("--case is required unless --list is used") + if args.warmup < 0 or args.rounds <= 0 or args.iterations <= 0: + parser.error( + "--warmup must be nonnegative; --rounds and --iterations must be positive" + ) + return args + + +def run_benchmark_suite( + cases: Mapping[str, BenchmarkCase], + *, + description: str | None = None, +) -> None: + args = _parse_args(cases, description=description) + if args.list: + _show_cases(cases) + return + + names = tuple(cases) if args.case == "all" else (args.case,) + assert all(name is not None for name in names) + if args.show_source: + for name in names: + assert name is not None + _show_source(cases[name]) + return + + if not torch.cuda.is_available(): + raise RuntimeError("the CODA microbenchmarks require CUDA") + if torch.cuda.get_device_capability(args.device) < (10, 0): + raise RuntimeError("QUACK FlexGEMM requires an SM100-or-later GPU") + + results = [] + for name in names: + assert name is not None + results.append(_run_case(cases[name], args)) + torch._dynamo.reset() + torch.cuda.empty_cache() + if args.output: + with open(args.output, "w", encoding="utf-8") as output_file: + json.dump(results, output_file, indent=2) + output_file.write("\n") + + +def main() -> None: + run_benchmark_suite(CASES) + + +if __name__ == "__main__": + main() diff --git a/torchtitan/experiments/graph_trainer/benchmarks/coda_fusion_microbench_16b.py b/torchtitan/experiments/graph_trainer/benchmarks/coda_fusion_microbench_16b.py new file mode 100644 index 0000000000..77f9577eb2 --- /dev/null +++ b/torchtitan/experiments/graph_trainer/benchmarks/coda_fusion_microbench_16b.py @@ -0,0 +1,348 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +"""DSV3-16B CODA eager-versus-FlexGEMM microbenchmarks. + +The shapes come from an unfused GraphTrainer joint FX graph for local batch 4, +sequence length 4096, FSDP4, EP4, and TP1. The run used full activation +checkpointing and the standard EP implementation. CODA graph passes were +disabled while collecting the graph. + +Run one case with:: + + python -m torchtitan.experiments.graph_trainer.benchmarks.coda_fusion_microbench_16b \ + --case f2_kv_rmsnorm + +Use ``--show-source`` to print the plain eager epilogue and FlexGEMM callback. +""" + +import torch + +from torchtitan.experiments.graph_trainer.benchmarks.coda_fusion_microbench import ( + _bf16, + _fp32, + BenchmarkCase, + eager_b1_lm_head_input_grad_cast, + eager_b2_shared_expert_swiglu_backward, + eager_b4_router_input_grad_add, + eager_b5_mla_rmsnorm_backward, + eager_b6_weight_grad_cast, + eager_b7_attention_input_grad_merge, + eager_f2_kv_rmsnorm, + eager_f3_attention_output, + eager_f3_moe_output, + eager_f4_shared_expert_swiglu, + make_flex_b1_lm_head_input_grad_cast, + make_flex_b2_shared_expert_swiglu_backward, + make_flex_b4_router_input_grad_add, + make_flex_b5_mla_rmsnorm_backward, + make_flex_b6_weight_grad_cast, + make_flex_b7_attention_input_grad_merge, + make_flex_f2_kv_rmsnorm, + make_flex_f3_attention_output, + make_flex_f3_moe_output, + make_flex_f4_shared_expert_swiglu, + run_benchmark_suite, +) + + +TOKENS = 4 * 4096 +MODEL_WIDTH = 2048 +NUM_EXPERTS = 64 +KV_PROJECTION_WIDTH = 4096 +KV_LOW_RANK_WIDTH = 512 +KV_ROPE_WIDTH = 64 +Q_PROJECTION_WIDTH = 3072 +SHARED_EXPERT_WIDTH = 2816 +DENSE_FFN_WIDTH = 10944 +VOCAB_SIZE = 102400 +LOSS_CHUNKS = 8 +LOSS_CHUNK_TOKENS = TOKENS // LOSS_CHUNKS + + +def _make_b1_inputs(device: torch.device) -> tuple[torch.Tensor, ...]: + return ( + _bf16((LOSS_CHUNK_TOKENS, VOCAB_SIZE), device), + _bf16((VOCAB_SIZE, MODEL_WIDTH), device), + ) + + +def _make_b2_inputs( + device: torch.device, + ffn_width: int, +) -> tuple[torch.Tensor, ...]: + return ( + _bf16((TOKENS, MODEL_WIDTH), device), + _bf16((MODEL_WIDTH, ffn_width), device), + _bf16((TOKENS, ffn_width), device), + _bf16((TOKENS, ffn_width), device), + _bf16((TOKENS, ffn_width), device), + _bf16((ffn_width, MODEL_WIDTH), device), + _bf16((ffn_width, MODEL_WIDTH), device), + ) + + +def _make_b2_shared_inputs(device: torch.device) -> tuple[torch.Tensor, ...]: + return _make_b2_inputs(device, SHARED_EXPERT_WIDTH) + + +def _make_b2_dense_inputs(device: torch.device) -> tuple[torch.Tensor, ...]: + return _make_b2_inputs(device, DENSE_FFN_WIDTH) + + +def _make_b4_inputs(device: torch.device) -> tuple[torch.Tensor, ...]: + return ( + _fp32((TOKENS, NUM_EXPERTS), device), + _fp32((NUM_EXPERTS, MODEL_WIDTH), device), + _bf16((TOKENS, MODEL_WIDTH), device), + ) + + +def _make_b5_inputs(device: torch.device) -> tuple[torch.Tensor, ...]: + return ( + _bf16((TOKENS, KV_PROJECTION_WIDTH), device), + _bf16((KV_PROJECTION_WIDTH, KV_LOW_RANK_WIDTH), device), + _bf16((TOKENS, KV_LOW_RANK_WIDTH), device), + torch.empty((TOKENS, 1), device=device, dtype=torch.float32).uniform_(0.5, 1.5), + _bf16((KV_LOW_RANK_WIDTH,), device), + ) + + +def _make_b6_inputs(device: torch.device) -> tuple[torch.Tensor, ...]: + return ( + _bf16((SHARED_EXPERT_WIDTH, TOKENS), device), + _bf16((TOKENS, MODEL_WIDTH), device), + ) + + +def _make_b7_inputs(device: torch.device) -> tuple[torch.Tensor, ...]: + return ( + _bf16((TOKENS, KV_LOW_RANK_WIDTH + KV_ROPE_WIDTH), device), + _bf16((KV_LOW_RANK_WIDTH + KV_ROPE_WIDTH, MODEL_WIDTH), device), + _bf16((TOKENS, Q_PROJECTION_WIDTH), device), + _bf16((Q_PROJECTION_WIDTH, MODEL_WIDTH), device), + ) + + +def _make_f2_kv_inputs(device: torch.device) -> tuple[torch.Tensor, ...]: + return ( + _bf16((TOKENS, MODEL_WIDTH), device), + _bf16((MODEL_WIDTH, KV_LOW_RANK_WIDTH + KV_ROPE_WIDTH), device), + torch.ones((KV_LOW_RANK_WIDTH,), device=device, dtype=torch.bfloat16), + _bf16((KV_LOW_RANK_WIDTH, KV_PROJECTION_WIDTH), device), + ) + + +def _make_f3_projection_inputs( + device: torch.device, + projection_width: int, +) -> tuple[torch.Tensor, ...]: + return ( + _bf16((TOKENS, projection_width), device), + _bf16((projection_width, MODEL_WIDTH), device), + _bf16((TOKENS, MODEL_WIDTH), device), + torch.ones((MODEL_WIDTH,), device=device, dtype=torch.bfloat16), + ) + + +def _make_f3_attention_inputs(device: torch.device) -> tuple[torch.Tensor, ...]: + return _make_f3_projection_inputs(device, MODEL_WIDTH) + + +def _make_f3_dense_inputs(device: torch.device) -> tuple[torch.Tensor, ...]: + return _make_f3_projection_inputs(device, DENSE_FFN_WIDTH) + + +def _make_f3_moe_inputs(device: torch.device) -> tuple[torch.Tensor, ...]: + return ( + _bf16((TOKENS, SHARED_EXPERT_WIDTH), device), + _bf16((SHARED_EXPERT_WIDTH, MODEL_WIDTH), device), + _bf16((TOKENS, MODEL_WIDTH), device), + _bf16((TOKENS, MODEL_WIDTH), device), + torch.ones((MODEL_WIDTH,), device=device, dtype=torch.bfloat16), + ) + + +def _make_f4_inputs( + device: torch.device, + ffn_width: int, +) -> tuple[torch.Tensor, ...]: + return ( + _bf16((TOKENS, MODEL_WIDTH), device), + _bf16((MODEL_WIDTH, ffn_width), device), + _bf16((MODEL_WIDTH, ffn_width), device), + ) + + +def _make_f4_shared_inputs(device: torch.device) -> tuple[torch.Tensor, ...]: + return _make_f4_inputs(device, SHARED_EXPERT_WIDTH) + + +def _make_f4_dense_inputs(device: torch.device) -> tuple[torch.Tensor, ...]: + return _make_f4_inputs(device, DENSE_FFN_WIDTH) + + +CASES: dict[str, BenchmarkCase] = { + case.name: case + for case in ( + BenchmarkCase( + "b1_lm_head_input_grad_cast", + "B1", + "One chunked LM-head input-gradient BF16 store and FP32 cast", + "(2048, 102400) @ (102400, 2048); 8 chunks per step", + 1, + _make_b1_inputs, + eager_b1_lm_head_input_grad_cast, + make_flex_b1_lm_head_input_grad_cast, + ), + BenchmarkCase( + "b2_shared_expert_swiglu_backward", + "B2", + "Shared-expert SwiGLU branch derivatives and input-gradient merge", + "M=16384, D=2048, P=2816; three BF16 GEMMs", + 2, + _make_b2_shared_inputs, + eager_b2_shared_expert_swiglu_backward, + make_flex_b2_shared_expert_swiglu_backward, + 3.1e-5, + 1e-2, + ), + BenchmarkCase( + "b2_dense_ffn_swiglu_backward", + "B2", + "Dense first-layer SwiGLU branch derivatives and input-gradient merge", + "M=16384, D=2048, P=10944; three BF16 GEMMs", + 2, + _make_b2_dense_inputs, + eager_b2_shared_expert_swiglu_backward, + make_flex_b2_shared_expert_swiglu_backward, + 3.1e-5, + 1e-2, + ), + BenchmarkCase( + "b4_router_input_grad_add", + "B4", + "FP32 router input-gradient GEMM, BF16 store, and expert-gradient add", + "(16384, 64) @ (64, 2048)", + 1, + _make_b4_inputs, + eager_b4_router_input_grad_add, + make_flex_b4_router_input_grad_add, + 4.9e-4, + 1e-2, + ), + BenchmarkCase( + "b5_mla_kv_rmsnorm_backward", + "B5", + "MLA KV input-gradient projection plus RMSNorm backward", + "(16384, 4096) @ (4096, 512)", + 1, + _make_b5_inputs, + eager_b5_mla_rmsnorm_backward, + make_flex_b5_mla_rmsnorm_backward, + 0.05, + 0.02, + ), + BenchmarkCase( + "b6_shared_expert_weight_grad_cast", + "B6", + "Frequent shared-expert BF16 weight-gradient GEMM and FP32 cast", + "(2816, 16384) @ (16384, 2048)", + 1, + _make_b6_inputs, + eager_b6_weight_grad_cast, + make_flex_b6_weight_grad_cast, + ), + BenchmarkCase( + "b7_attention_input_grad_merge", + "B7", + "KV and direct-Q input-gradient GEMMs followed by BF16 add", + "(16384, 576) @ (576, 2048) + (16384, 3072) @ (3072, 2048)", + 1, + _make_b7_inputs, + eager_b7_attention_input_grad_merge, + make_flex_b7_attention_input_grad_merge, + ), + BenchmarkCase( + "f2_kv_rmsnorm", + "F2-KV", + "Segmented MLA KV projection, RMSNorm, and expanded projection", + "2048 -> 576, RMSNorm(512), 512 -> 4096 at M=16384", + 2, + _make_f2_kv_inputs, + eager_f2_kv_rmsnorm, + make_flex_f2_kv_rmsnorm, + 1.6e-2, + 2e-2, + ), + BenchmarkCase( + "f3_attention_output", + "F3-A", + "Attention WO projection, residual add, and FFN RMSNorm", + "(16384, 2048) @ (2048, 2048)", + 1, + _make_f3_attention_inputs, + eager_f3_attention_output, + make_flex_f3_attention_output, + 0.07, + 2e-2, + ), + BenchmarkCase( + "f3_moe_output", + "F3-B", + "Shared W2 projection, routed add, residual add, and next RMSNorm", + "(16384, 2816) @ (2816, 2048)", + 1, + _make_f3_moe_inputs, + eager_f3_moe_output, + make_flex_f3_moe_output, + 0.07, + 2e-2, + ), + BenchmarkCase( + "f3_dense_ffn_output", + "F3-B", + "Dense W2 projection, residual add, and next attention RMSNorm", + "(16384, 10944) @ (10944, 2048)", + 1, + _make_f3_dense_inputs, + eager_f3_attention_output, + make_flex_f3_attention_output, + 0.07, + 2e-2, + ), + BenchmarkCase( + "f4_shared_expert_swiglu", + "F4", + "Shared-expert W1/W3 GEMMs with SiLU and multiply", + "two (16384, 2048) @ (2048, 2816) GEMMs", + 2, + _make_f4_shared_inputs, + eager_f4_shared_expert_swiglu, + make_flex_f4_shared_expert_swiglu, + fast_math_flex_gemms=(0,), + ), + BenchmarkCase( + "f4_dense_ffn_swiglu", + "F4", + "Dense first-layer W1/W3 GEMMs with SiLU and multiply", + "two (16384, 2048) @ (2048, 10944) GEMMs", + 2, + _make_f4_dense_inputs, + eager_f4_shared_expert_swiglu, + make_flex_f4_shared_expert_swiglu, + fast_math_flex_gemms=(0,), + ), + ) +} + + +def main() -> None: + run_benchmark_suite(CASES, description=__doc__) + + +if __name__ == "__main__": + main() diff --git a/torchtitan/experiments/graph_trainer/coda_passes.py b/torchtitan/experiments/graph_trainer/coda_passes.py new file mode 100644 index 0000000000..fbcd59d33a --- /dev/null +++ b/torchtitan/experiments/graph_trainer/coda_passes.py @@ -0,0 +1,3357 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +"""CODA-style GEMM fusion passes for GraphTrainer.""" + +from __future__ import annotations + +import math +import operator +from collections.abc import Callable, Iterable + +import torch +from torch._higher_order_ops.flex_gemm import ( + apply_flex_gemm_body_graph_passes, + flex_gemm_hop, +) + +from torchtitan.tools.logging import logger + + +_MM = torch.ops.aten.mm.default +_ADD = torch.ops.aten.add.Tensor +_ADD_SCALAR = torch.ops.aten.add.Scalar +_ALIAS = torch.ops.aten.alias.default +_CONSTANT_PAD_ND = torch.ops.aten.constant_pad_nd.default +_COPY_ = torch.ops.aten.copy_.default +_FUSED_RMS_NORM = torch.ops.aten._fused_rms_norm.default +_FUSED_RMS_NORM_BACKWARD = torch.ops.aten._fused_rms_norm_backward.default +_DIV_SCALAR = torch.ops.aten.div.Scalar +_MEAN_DIM = torch.ops.aten.mean.dim +_MUL = torch.ops.aten.mul.Tensor +_POW_SCALAR = torch.ops.aten.pow.Tensor_Scalar +_RESHAPE = torch.ops.aten.reshape.default +_RSQRT = torch.ops.aten.rsqrt.default +_SLICE = torch.ops.aten.slice.Tensor +_SPLIT_WITH_SIZES = torch.ops.aten.split_with_sizes.default +_SIGMOID = torch.ops.aten.sigmoid.default +_SILU = torch.ops.aten.silu.default +_SILU_BACKWARD = torch.ops.aten.silu_backward.default +_SUB = torch.ops.aten.sub.Tensor +_SUM_DIM = torch.ops.aten.sum.dim_IntList +_TO_COPY = torch.ops.aten._to_copy.default +_VIEW = torch.ops.aten.view.default +_COMPILE_WITH_INDUCTOR = "compile_with_inductor" +_CODA_RMSNORM_BACKWARD_GROUP = 128 +_CODA_RMSNORM_GROUP = 512 +_CODA_KV_RMSNORM_GROUP = 64 +_CODA_KV_ACTIVE_WIDTH = 512 +_CODA_KV_TAIL_WIDTH = 64 + + +def _node_dtype(node: torch.fx.Node) -> torch.dtype | None: + value = node.meta.get("val") + if isinstance(value, torch.Tensor): + return value.dtype + tensor_meta = node.meta.get("tensor_meta") + return getattr(tensor_meta, "dtype", None) + + +def _node_shape(node: torch.fx.Node) -> tuple | None: + value = node.meta.get("val") + if isinstance(value, torch.Tensor): + return tuple(value.shape) + tensor_meta = node.meta.get("tensor_meta") + shape = getattr(tensor_meta, "shape", None) + return None if shape is None else tuple(shape) + + +def _module_fqn(node: torch.fx.Node) -> str | None: + custom = node.meta.get("custom", {}) + module_fqn = custom.get("module_fqn") + return module_fqn if isinstance(module_fqn, str) else None + + +def _sole_user(node: torch.fx.Node) -> torch.fx.Node | None: + if len(node.users) != 1: + return None + return next(iter(node.users)) + + +def _is_cast(node: torch.fx.Node, dtype: torch.dtype) -> bool: + return node.target == _TO_COPY and node.kwargs.get("dtype") == dtype + + +def _copy_meta(*nodes: torch.fx.Node) -> dict: + """Merge metadata in dataflow order without copying FakeTensor values.""" + merged: dict = {} + custom: dict = {} + for node in nodes: + merged.update(node.meta) + custom.update(node.meta.get("custom", {})) + if custom: + merged["custom"] = custom + return merged + + +def _copy_meta_with_value_from( + value_source: torch.fx.Node, + *nodes: torch.fx.Node, +) -> dict: + merged = _copy_meta(*nodes) + for key in ("val", "tensor_meta"): + if key in value_source.meta: + merged[key] = value_source.meta[key] + else: + merged.pop(key, None) + return merged + + +def _copy_meta_with_dtype( + value_source: torch.fx.Node, + dtype: torch.dtype, + *nodes: torch.fx.Node, +) -> dict: + merged = _copy_meta(*nodes) + value = value_source.meta.get("val") + if isinstance(value, torch.Tensor): + merged["val"] = value.to(dtype=dtype) + else: + merged.pop("val", None) + merged.pop("tensor_meta", None) + return merged + + +def _copy_meta_with_value(value: torch.Tensor, *nodes: torch.fx.Node) -> dict: + merged = _copy_meta(*nodes) + merged["val"] = value + merged.pop("tensor_meta", None) + return merged + + +def _tag_for_regional_inductor(node: torch.fx.Node) -> None: + node.meta.setdefault("custom", {})[_COMPILE_WITH_INDUCTOR] = {} + + +def _next_submodule_name(gm: torch.fx.GraphModule, prefix: str) -> str: + index = 0 + while hasattr(gm, f"{prefix}_{index}"): + index += 1 + return f"{prefix}_{index}" + + +def _build_bf16_mm_fp32_body( + mm: torch.fx.Node, + *metadata_nodes: torch.fx.Node, +) -> torch.fx.GraphModule: + graph = torch.fx.Graph() + inputs = [] + for index, arg in enumerate(mm.args): + placeholder = graph.placeholder(f"arg{index}") + if isinstance(arg, torch.fx.Node): + placeholder.meta = dict(arg.meta) + inputs.append(placeholder) + + body_mm = graph.call_function(_MM, tuple(inputs), dict(mm.kwargs)) + body_mm.meta = dict(mm.meta) + # FlexGEMM exposes its accumulator as FP32. Preserve the original BF16 + # GEMM store rounding before returning the FP32 gradient. + to_fp32 = graph.call_function(_TO_COPY, (body_mm,), {"dtype": torch.float32}) + to_fp32.meta = _copy_meta_with_dtype(mm, torch.float32, mm, *metadata_nodes) + to_bf16 = graph.call_function(_TO_COPY, (to_fp32,), {"dtype": torch.bfloat16}) + to_bf16.meta = _copy_meta_with_value_from(mm, mm, *metadata_nodes) + value = graph.call_function(_TO_COPY, (to_bf16,), {"dtype": torch.float32}) + value.meta = _copy_meta_with_dtype(mm, torch.float32, mm, *metadata_nodes) + + # FlexGEMM lowering returns an ordered output tuple. Keeping the singleton + # tuple explicit gives regional Inductor a getitem node to consume. + graph.output((value,)) + body = torch.fx.GraphModule(torch.nn.Module(), graph) + apply_flex_gemm_body_graph_passes(body, _MM) + for node in body.graph.nodes: + _tag_for_regional_inductor(node) + return body + + +def _build_f6_router_body( + mm: torch.fx.Node, + sigmoid: torch.fx.Node, + bias_2d: torch.fx.Node | None, + bias_add: torch.fx.Node | None, +) -> torch.fx.GraphModule: + graph = torch.fx.Graph() + inputs = [] + for index, arg in enumerate(mm.args): + placeholder = graph.placeholder(f"arg{index}") + if isinstance(arg, torch.fx.Node): + placeholder.meta = dict(arg.meta) + inputs.append(placeholder) + + bias_input = None + if bias_2d is not None: + bias_input = graph.placeholder(f"arg{len(inputs)}") + bias_input.meta = dict(bias_2d.meta) + inputs.append(bias_input) + + body_mm = graph.call_function(_MM, tuple(inputs[:2]), dict(mm.kwargs)) + body_mm.meta = dict(mm.meta) + raw_scores = graph.call_function(_SIGMOID, (body_mm,)) + raw_scores.meta = _copy_meta_with_value_from(mm, mm, sigmoid) + outputs = [raw_scores] + if bias_input is not None and bias_add is not None: + biased_scores = graph.call_function(_ADD, (raw_scores, bias_input)) + biased_scores.meta = _copy_meta_with_value_from(mm, mm, sigmoid, bias_add) + outputs.append(biased_scores) + + graph.output(tuple(outputs)) + body = torch.fx.GraphModule(torch.nn.Module(), graph) + apply_flex_gemm_body_graph_passes(body, _MM) + for node in body.graph.nodes: + _tag_for_regional_inductor(node) + return body + + +def _build_f4_silu_body( + mm: torch.fx.Node, + reshape: torch.fx.Node, + silu: torch.fx.Node, + preserve_preactivation: bool, +) -> torch.fx.GraphModule: + graph = torch.fx.Graph() + inputs = [] + for index, arg in enumerate(mm.args): + placeholder = graph.placeholder(f"arg{index}") + if isinstance(arg, torch.fx.Node): + placeholder.meta = dict(arg.meta) + inputs.append(placeholder) + + body_mm = graph.call_function(_MM, tuple(inputs), dict(mm.kwargs)) + body_mm.meta = dict(mm.meta) + to_fp32 = graph.call_function(_TO_COPY, (body_mm,), {"dtype": torch.float32}) + to_fp32.meta = _copy_meta_with_dtype(mm, torch.float32, mm, reshape, silu) + rounded = graph.call_function(_TO_COPY, (to_fp32,), {"dtype": torch.bfloat16}) + rounded.meta = _copy_meta_with_value_from(mm, mm, reshape, silu) + silu_2d = graph.call_function(_SILU, (rounded,)) + silu_2d.meta = _copy_meta_with_value_from(mm, mm, reshape, silu) + + outputs = (silu_2d, rounded) if preserve_preactivation else (silu_2d,) + graph.output(outputs) + body = torch.fx.GraphModule(torch.nn.Module(), graph) + apply_flex_gemm_body_graph_passes(body, _MM) + for node in body.graph.nodes: + _tag_for_regional_inductor(node) + return body + + +def _build_f4_mul_body( + mm: torch.fx.Node, + reshape: torch.fx.Node, + silu_2d: torch.fx.Node, + mul: torch.fx.Node, +) -> torch.fx.GraphModule: + graph = torch.fx.Graph() + inputs = [] + for index, arg in enumerate(mm.args): + placeholder = graph.placeholder(f"arg{index}") + if isinstance(arg, torch.fx.Node): + placeholder.meta = dict(arg.meta) + inputs.append(placeholder) + silu_input = graph.placeholder(f"arg{len(inputs)}") + silu_input.meta = dict(silu_2d.meta) + inputs.append(silu_input) + + body_mm = graph.call_function(_MM, tuple(inputs[:2]), dict(mm.kwargs)) + body_mm.meta = dict(mm.meta) + to_fp32 = graph.call_function(_TO_COPY, (body_mm,), {"dtype": torch.float32}) + to_fp32.meta = _copy_meta_with_dtype(mm, torch.float32, mm, reshape, mul) + rounded = graph.call_function(_TO_COPY, (to_fp32,), {"dtype": torch.bfloat16}) + rounded.meta = _copy_meta_with_value_from(mm, mm, reshape, mul) + product = graph.call_function(_MUL, (rounded, silu_input)) + product.meta = _copy_meta_with_value_from(mm, mm, reshape, mul) + + graph.output((rounded, product)) + body = torch.fx.GraphModule(torch.nn.Module(), graph) + apply_flex_gemm_body_graph_passes(body, _MM) + for node in body.graph.nodes: + _tag_for_regional_inductor(node) + return body + + +def _build_b2_branch_body( + mm: torch.fx.Node, + grad_view: torch.fx.Node, + saved_silu_2d: torch.fx.Node, + saved_gate_2d: torch.fx.Node, + saved_preactivation_2d: torch.fx.Node, + gate_grad: torch.fx.Node, + silu_grad: torch.fx.Node, +) -> torch.fx.GraphModule: + graph = torch.fx.Graph() + inputs = [] + sources = ( + *mm.args, + saved_silu_2d, + saved_gate_2d, + saved_preactivation_2d, + ) + for index, source in enumerate(sources): + placeholder = graph.placeholder(f"arg{index}") + if isinstance(source, torch.fx.Node): + placeholder.meta = dict(source.meta) + inputs.append(placeholder) + + body_mm = graph.call_function(_MM, tuple(inputs[:2]), dict(mm.kwargs)) + body_mm.meta = dict(mm.meta) + to_fp32 = graph.call_function(_TO_COPY, (body_mm,), {"dtype": torch.float32}) + to_fp32.meta = _copy_meta_with_dtype(mm, torch.float32, mm, grad_view) + rounded = graph.call_function(_TO_COPY, (to_fp32,), {"dtype": torch.bfloat16}) + rounded.meta = _copy_meta_with_value_from(mm, mm, grad_view) + + gate_grad_2d = graph.call_function(_MUL, (rounded, inputs[2])) + gate_grad_2d.meta = _copy_meta_with_value_from(mm, mm, grad_view, gate_grad) + gated_grad_2d = graph.call_function(_MUL, (rounded, inputs[3])) + gated_grad_2d.meta = _copy_meta_with_value_from(mm, mm, grad_view) + silu_grad_2d = graph.call_function( + _SILU_BACKWARD, + (gated_grad_2d, inputs[4]), + ) + silu_grad_2d.meta = _copy_meta_with_value_from(mm, mm, grad_view, silu_grad) + + graph.output((gate_grad_2d, silu_grad_2d)) + body = torch.fx.GraphModule(torch.nn.Module(), graph) + apply_flex_gemm_body_graph_passes(body, _MM) + for node in body.graph.nodes: + _tag_for_regional_inductor(node) + return body + + +def _build_b2_input_add_body( + mm: torch.fx.Node, + reshape: torch.fx.Node, + captured_branch_2d: torch.fx.Node, + add: torch.fx.Node, +) -> torch.fx.GraphModule: + graph = torch.fx.Graph() + inputs = [] + sources = (*mm.args, captured_branch_2d) + for index, source in enumerate(sources): + placeholder = graph.placeholder(f"arg{index}") + if isinstance(source, torch.fx.Node): + placeholder.meta = dict(source.meta) + inputs.append(placeholder) + + body_mm = graph.call_function(_MM, tuple(inputs[:2]), dict(mm.kwargs)) + body_mm.meta = dict(mm.meta) + to_fp32 = graph.call_function(_TO_COPY, (body_mm,), {"dtype": torch.float32}) + to_fp32.meta = _copy_meta_with_dtype(mm, torch.float32, mm, reshape, add) + rounded = graph.call_function(_TO_COPY, (to_fp32,), {"dtype": torch.bfloat16}) + rounded.meta = _copy_meta_with_value_from(mm, mm, reshape, add) + total = graph.call_function(_ADD, (inputs[2], rounded)) + total.meta = _copy_meta_with_value_from(mm, mm, reshape, add) + + graph.output((total,)) + body = torch.fx.GraphModule(torch.nn.Module(), graph) + apply_flex_gemm_body_graph_passes(body, _MM) + for node in body.graph.nodes: + _tag_for_regional_inductor(node) + return body + + +def _build_b4_router_input_grad_body( + mm: torch.fx.Node, + cast: torch.fx.Node, + residual_2d: torch.fx.Node, + add: torch.fx.Node, +) -> torch.fx.GraphModule: + graph = torch.fx.Graph() + inputs = [] + for index, source in enumerate((*mm.args, residual_2d)): + placeholder = graph.placeholder(f"arg{index}") + if isinstance(source, torch.fx.Node): + placeholder.meta = dict(source.meta) + inputs.append(placeholder) + + body_mm = graph.call_function(_MM, tuple(inputs[:2]), dict(mm.kwargs)) + body_mm.meta = dict(mm.meta) + rounded = graph.call_function(_TO_COPY, (body_mm,), dict(cast.kwargs)) + rounded.meta = _copy_meta_with_dtype(mm, torch.bfloat16, mm, cast) + total = graph.call_function(_ADD, (inputs[2], rounded)) + total.meta = _copy_meta_with_dtype(mm, torch.bfloat16, mm, cast, add) + + graph.output((total,)) + body = torch.fx.GraphModule(torch.nn.Module(), graph) + apply_flex_gemm_body_graph_passes(body, _MM) + for node in body.graph.nodes: + _tag_for_regional_inductor(node) + return body + + +def _build_b7_attention_grad_merge_body( + mm: torch.fx.Node, + other_branch_2d: torch.fx.Node, + mm_is_lhs: bool, +) -> torch.fx.GraphModule: + graph = torch.fx.Graph() + inputs = [] + for index, source in enumerate((*mm.args, other_branch_2d)): + placeholder = graph.placeholder(f"arg{index}") + if isinstance(source, torch.fx.Node): + placeholder.meta = dict(source.meta) + inputs.append(placeholder) + + body_mm = graph.call_function(_MM, tuple(inputs[:2]), dict(mm.kwargs)) + body_mm.meta = dict(mm.meta) + accumulator = graph.call_function(_TO_COPY, (body_mm,), {"dtype": torch.float32}) + accumulator.meta = _copy_meta_with_dtype(mm, torch.float32, mm) + rounded = graph.call_function(_TO_COPY, (accumulator,), {"dtype": torch.bfloat16}) + rounded.meta = _copy_meta_with_value_from(mm, mm) + add_args = (rounded, inputs[2]) if mm_is_lhs else (inputs[2], rounded) + total = graph.call_function(_ADD, add_args) + total.meta = _copy_meta_with_value_from(mm, mm, other_branch_2d) + + graph.output((total,)) + body = torch.fx.GraphModule(torch.nn.Module(), graph) + apply_flex_gemm_body_graph_passes(body, _MM) + for node in body.graph.nodes: + _tag_for_regional_inductor(node) + return body + + +def _build_b5_mla_rmsnorm_backward_body( + mm: torch.fx.Node, + norm_input_2d: torch.fx.Node, + rstd_2d: torch.fx.Node, + gamma_2d: torch.fx.Node, +) -> torch.fx.GraphModule: + graph = torch.fx.Graph() + inputs = [] + for index, source in enumerate((*mm.args, norm_input_2d, rstd_2d, gamma_2d)): + placeholder = graph.placeholder(f"arg{index}") + if isinstance(source, torch.fx.Node): + placeholder.meta = dict(source.meta) + inputs.append(placeholder) + + body_mm = graph.call_function(_MM, tuple(inputs[:2]), dict(mm.kwargs)) + body_mm.meta = dict(mm.meta) + accumulator = graph.call_function(_TO_COPY, (body_mm,), {"dtype": torch.float32}) + accumulator.meta = _copy_meta_with_dtype(mm, torch.float32, mm) + rounded = graph.call_function(_TO_COPY, (accumulator,), {"dtype": torch.bfloat16}) + rounded.meta = _copy_meta_with_value_from(mm, mm) + rounded_fp32 = graph.call_function(_TO_COPY, (rounded,), {"dtype": torch.float32}) + rounded_fp32.meta = _copy_meta_with_dtype(mm, torch.float32, mm) + norm_input_fp32 = graph.call_function( + _TO_COPY, (inputs[2],), {"dtype": torch.float32} + ) + norm_input_fp32.meta = _copy_meta_with_dtype( + norm_input_2d, torch.float32, norm_input_2d + ) + x_hat = graph.call_function(_MUL, (norm_input_fp32, inputs[3])) + x_hat.meta = _copy_meta_with_dtype( + norm_input_2d, torch.float32, norm_input_2d, rstd_2d + ) + gamma_fp32 = graph.call_function(_TO_COPY, (inputs[4],), {"dtype": torch.float32}) + gamma_fp32.meta = _copy_meta_with_dtype(gamma_2d, torch.float32, gamma_2d) + grad_x_hat = graph.call_function(_MUL, (rounded_fp32, gamma_fp32)) + grad_x_hat.meta = _copy_meta_with_dtype(mm, torch.float32, mm, gamma_2d) + row_products = graph.call_function(_MUL, (x_hat, grad_x_hat)) + row_products.meta = _copy_meta_with_dtype( + mm, torch.float32, mm, norm_input_2d, rstd_2d, gamma_2d + ) + + mm_shape = _node_shape(mm) + assert mm_shape is not None + grouped = graph.call_function( + _VIEW, + ( + row_products, + [mm_shape[0], -1, _CODA_RMSNORM_BACKWARD_GROUP], + ), + ) + row_products_value = row_products.meta.get("val") + if isinstance(row_products_value, torch.Tensor): + grouped.meta = _copy_meta_with_value( + row_products_value.reshape(mm_shape[0], -1, _CODA_RMSNORM_BACKWARD_GROUP), + mm, + norm_input_2d, + ) + else: + grouped.meta = _copy_meta(mm, norm_input_2d) + partial_row_dot = graph.call_function(_SUM_DIM, (grouped, [-1])) + grouped_value = grouped.meta.get("val") + if isinstance(grouped_value, torch.Tensor): + partial_row_dot.meta = _copy_meta_with_value( + grouped_value.sum(-1), mm, norm_input_2d, rstd_2d, gamma_2d + ) + else: + partial_row_dot.meta = _copy_meta(mm, norm_input_2d, rstd_2d, gamma_2d) + + graph.output((rounded, partial_row_dot)) + body = torch.fx.GraphModule(torch.nn.Module(), graph) + apply_flex_gemm_body_graph_passes(body, _MM) + for node in body.graph.nodes: + _tag_for_regional_inductor(node) + return body + + +def _build_f3_residual_rmsnorm_body( + mm: torch.fx.Node, + addends_2d: tuple[torch.fx.Node, ...], + group: int, + accumulated_value_is_lhs: tuple[bool, ...], +) -> torch.fx.GraphModule: + graph = torch.fx.Graph() + inputs = [] + for index, source in enumerate((*mm.args, *addends_2d)): + placeholder = graph.placeholder(f"arg{index}") + if isinstance(source, torch.fx.Node): + placeholder.meta = dict(source.meta) + inputs.append(placeholder) + + body_mm = graph.call_function(_MM, tuple(inputs[:2]), dict(mm.kwargs)) + body_mm.meta = dict(mm.meta) + accumulator = graph.call_function(_TO_COPY, (body_mm,), {"dtype": torch.float32}) + accumulator.meta = _copy_meta_with_dtype(mm, torch.float32, mm) + rounded = graph.call_function(_TO_COPY, (accumulator,), {"dtype": torch.bfloat16}) + rounded.meta = _copy_meta_with_value_from(mm, mm) + total = rounded + for input_index, (addend, value_is_lhs) in enumerate( + zip(addends_2d, accumulated_value_is_lhs, strict=True), start=2 + ): + add_args = ( + (total, inputs[input_index]) + if value_is_lhs + else ( + inputs[input_index], + total, + ) + ) + total = graph.call_function(_ADD, add_args) + total.meta = _copy_meta_with_value_from(mm, mm, *addends_2d[: input_index - 1]) + total_fp32 = graph.call_function(_TO_COPY, (total,), {"dtype": torch.float32}) + total_fp32.meta = _copy_meta_with_dtype(mm, torch.float32, mm, *addends_2d) + + mm_shape = _node_shape(mm) + assert mm_shape is not None + grouped = graph.call_function( + _VIEW, + (total_fp32, [mm_shape[0], -1, group]), + ) + total_value = total_fp32.meta.get("val") + if isinstance(total_value, torch.Tensor): + grouped.meta = _copy_meta_with_value( + total_value.reshape(mm_shape[0], -1, group), mm, *addends_2d + ) + else: + grouped.meta = _copy_meta(mm, *addends_2d) + squared = graph.call_function(_POW_SCALAR, (grouped, 2)) + grouped_value = grouped.meta.get("val") + if isinstance(grouped_value, torch.Tensor): + squared.meta = _copy_meta_with_value(grouped_value.square(), mm, *addends_2d) + else: + squared.meta = _copy_meta(mm, *addends_2d) + partial_mean_square = graph.call_function(_MEAN_DIM, (squared, [-1])) + squared_value = squared.meta.get("val") + if isinstance(squared_value, torch.Tensor): + partial_mean_square.meta = _copy_meta_with_value( + squared_value.mean(-1), mm, *addends_2d + ) + else: + partial_mean_square.meta = _copy_meta(mm, *addends_2d) + + graph.output((total, partial_mean_square)) + body = torch.fx.GraphModule(torch.nn.Module(), graph) + apply_flex_gemm_body_graph_passes(body, _MM) + for node in body.graph.nodes: + _tag_for_regional_inductor(node) + return body + + +def _build_f2_q_rmsnorm_first_body( + mm: torch.fx.Node, + gamma_2d: torch.fx.Node, + group: int, + preserve_raw: bool, +) -> torch.fx.GraphModule: + graph = torch.fx.Graph() + inputs = [] + for index, source in enumerate((*mm.args, gamma_2d)): + placeholder = graph.placeholder(f"arg{index}") + if isinstance(source, torch.fx.Node): + placeholder.meta = dict(source.meta) + inputs.append(placeholder) + + body_mm = graph.call_function(_MM, tuple(inputs[:2]), dict(mm.kwargs)) + body_mm.meta = dict(mm.meta) + accumulator = graph.call_function(_TO_COPY, (body_mm,), {"dtype": torch.float32}) + accumulator.meta = _copy_meta_with_dtype(mm, torch.float32, mm) + rounded = graph.call_function(_TO_COPY, (accumulator,), {"dtype": torch.bfloat16}) + rounded.meta = _copy_meta_with_value_from(mm, mm) + rounded_fp32 = graph.call_function(_TO_COPY, (rounded,), {"dtype": torch.float32}) + rounded_fp32.meta = _copy_meta_with_dtype(mm, torch.float32, mm) + weighted_fp32 = graph.call_function(_MUL, (rounded_fp32, inputs[2])) + weighted_fp32.meta = _copy_meta_with_dtype(mm, torch.float32, mm, gamma_2d) + weighted = graph.call_function( + _TO_COPY, (weighted_fp32,), {"dtype": torch.bfloat16} + ) + weighted.meta = _copy_meta_with_value_from(mm, mm, gamma_2d) + + mm_shape = _node_shape(mm) + assert mm_shape is not None + grouped = graph.call_function( + _VIEW, + (rounded_fp32, [mm_shape[0], -1, group]), + ) + rounded_value = rounded_fp32.meta.get("val") + if isinstance(rounded_value, torch.Tensor): + grouped.meta = _copy_meta_with_value( + rounded_value.reshape(mm_shape[0], -1, group), mm + ) + else: + grouped.meta = _copy_meta(mm) + squared = graph.call_function(_POW_SCALAR, (grouped, 2)) + grouped_value = grouped.meta.get("val") + if isinstance(grouped_value, torch.Tensor): + squared.meta = _copy_meta_with_value(grouped_value.square(), mm) + else: + squared.meta = _copy_meta(mm) + partial_mean_square = graph.call_function(_MEAN_DIM, (squared, [-1])) + squared_value = squared.meta.get("val") + if isinstance(squared_value, torch.Tensor): + partial_mean_square.meta = _copy_meta_with_value(squared_value.mean(-1), mm) + else: + partial_mean_square.meta = _copy_meta(mm) + + outputs = ( + (weighted, rounded, partial_mean_square) + if preserve_raw + else (weighted, partial_mean_square) + ) + graph.output(outputs) + body = torch.fx.GraphModule(torch.nn.Module(), graph) + apply_flex_gemm_body_graph_passes(body, _MM) + for node in body.graph.nodes: + _tag_for_regional_inductor(node) + return body + + +def _build_f2_q_rmsnorm_second_body( + mm: torch.fx.Node, + rstd_2d: torch.fx.Node, +) -> torch.fx.GraphModule: + graph = torch.fx.Graph() + inputs = [] + for index, source in enumerate((*mm.args, rstd_2d)): + placeholder = graph.placeholder(f"arg{index}") + if isinstance(source, torch.fx.Node): + placeholder.meta = dict(source.meta) + inputs.append(placeholder) + + body_mm = graph.call_function(_MM, tuple(inputs[:2]), dict(mm.kwargs)) + body_mm.meta = dict(mm.meta) + accumulator = graph.call_function(_TO_COPY, (body_mm,), {"dtype": torch.float32}) + accumulator.meta = _copy_meta_with_dtype(mm, torch.float32, mm) + scaled = graph.call_function(_MUL, (accumulator, inputs[2])) + scaled.meta = _copy_meta_with_dtype(mm, torch.float32, mm, rstd_2d) + rounded = graph.call_function(_TO_COPY, (scaled,), {"dtype": torch.bfloat16}) + rounded.meta = _copy_meta_with_value_from(mm, mm) + + graph.output((rounded,)) + body = torch.fx.GraphModule(torch.nn.Module(), graph) + apply_flex_gemm_body_graph_passes(body, _MM) + for node in body.graph.nodes: + _tag_for_regional_inductor(node) + return body + + +def _match_b1_lm_head_input_grad_cast( + cast: torch.fx.Node, +) -> tuple[torch.fx.Node, torch.fx.Node, torch.fx.Node] | None: + if not _is_cast(cast, torch.float32) or not cast.args: + return None + alias = cast.args[0] + if ( + not isinstance(alias, torch.fx.Node) + or alias.target != _ALIAS + or not alias.args + or _sole_user(alias) is not cast + ): + return None + reshape = alias.args[0] + if ( + not isinstance(reshape, torch.fx.Node) + or reshape.target not in (_RESHAPE, _VIEW) + or not reshape.args + or not isinstance(reshape.args[0], torch.fx.Node) + or _sole_user(reshape) is not alias + ): + return None + mm = reshape.args[0] + if ( + mm.target != _MM + or _sole_user(mm) is not reshape + or mm.meta.get("autograd_backward") is not True + or _module_fqn(mm) != "lm_head" + ): + return None + copy = _sole_user(cast) + if ( + copy is None + or copy.target != _COPY_ + or len(copy.args) < 2 + or copy.args[1] is not cast + or not isinstance(copy.args[0], torch.fx.Node) + ): + return None + + mm_shape = _node_shape(mm) + cast_shape = _node_shape(cast) + mm_value = mm.meta.get("val") + cast_value = cast.meta.get("val") + if ( + mm_shape is None + or len(mm_shape) != 2 + or cast_shape is None + or not cast_shape + or mm_shape[-1] != cast_shape[-1] + or _node_shape(reshape) != cast_shape + or _node_shape(alias) != cast_shape + or _node_shape(copy.args[0]) != cast_shape + or _node_dtype(mm) != torch.bfloat16 + or _node_dtype(cast) != torch.float32 + or not isinstance(mm_value, torch.Tensor) + or not isinstance(cast_value, torch.Tensor) + or mm_value.numel() != cast_value.numel() + ): + return None + return mm, reshape, alias + + +def fuse_b1_lm_head_input_grad_cast_pass( + gm: torch.fx.GraphModule, + example_inputs: tuple | None = None, +) -> torch.fx.GraphModule: + """Fuse chunked LM-head input-gradient BF16-to-FP32 conversion.""" + del example_inputs + num_fused = 0 + for cast in list(gm.graph.nodes): + match = _match_b1_lm_head_input_grad_cast(cast) + if match is None: + continue + mm, reshape, alias = match + mm_shape = _node_shape(mm) + cast_shape = _node_shape(cast) + mm_value = mm.meta["val"] + cast_value = cast.meta["val"] + assert mm_shape is not None + assert cast_shape is not None + assert isinstance(mm_value, torch.Tensor) + assert isinstance(cast_value, torch.Tensor) + + body = _build_bf16_mm_fp32_body(mm, reshape, alias, cast) + body_name = _next_submodule_name(gm, "_coda_b1_lm_head_input_grad_body") + gm.add_module(body_name, body) + with gm.graph.inserting_before(mm): + body_ref = gm.graph.get_attr(body_name) + _tag_for_regional_inductor(body_ref) + fused = gm.graph.call_function( + flex_gemm_hop, + ( + _MM, + body_ref, + tuple(mm.args), + dict(mm.kwargs), + {"backend": "QUACK"}, + ), + ) + output_2d_meta = _copy_meta_with_value( + mm_value.float(), mm, reshape, alias, cast + ) + fused.meta = _copy_meta(mm, reshape, alias, cast) + fused.meta["val"] = (output_2d_meta["val"],) + fused.meta.pop("tensor_meta", None) + _tag_for_regional_inductor(fused) + output_2d = gm.graph.call_function(operator.getitem, (fused, 0)) + output_2d.meta = output_2d_meta + _tag_for_regional_inductor(output_2d) + output = gm.graph.call_function( + _RESHAPE, + (output_2d, list(cast_shape)), + ) + output.meta = _copy_meta_with_value(cast_value.reshape(cast_shape), cast) + _tag_for_regional_inductor(output) + + cast.replace_all_uses_with(output) + gm.graph.erase_node(cast) + gm.graph.erase_node(alias) + gm.graph.erase_node(reshape) + gm.graph.erase_node(mm) + num_fused += 1 + + gm.graph.lint() + gm.recompile() + logger.info(f"B1 fused {num_fused} chunked LM-head input-gradient casts") + return gm + + +def _match_b6_bf16_cast(mm: torch.fx.Node) -> torch.fx.Node | None: + if mm.target != _MM: + return None + first = _sole_user(mm) + if first is None: + return None + + if _node_dtype(mm) == torch.bfloat16 and _is_cast(first, torch.float32): + return first + return None + + +def fuse_b6_bf16_weight_grad_cast_pass( + gm: torch.fx.GraphModule, + example_inputs: tuple | None = None, +) -> torch.fx.GraphModule: + """Fuse BF16 weight-gradient ``mm`` plus FP32 cast into FlexGEMM.""" + del example_inputs + num_fused = 0 + for mm in list(gm.graph.nodes): + cast = _match_b6_bf16_cast(mm) + if cast is None: + continue + + body = _build_bf16_mm_fp32_body(mm, cast) + body_name = _next_submodule_name(gm, "_coda_b6_body") + gm.add_module(body_name, body) + with gm.graph.inserting_before(mm): + body_ref = gm.graph.get_attr(body_name) + _tag_for_regional_inductor(body_ref) + fused = gm.graph.call_function( + flex_gemm_hop, + ( + _MM, + body_ref, + tuple(mm.args), + dict(mm.kwargs), + {"backend": "QUACK"}, + ), + ) + fused.meta = _copy_meta(mm, cast) + for key in ("val", "tensor_meta"): + if key in fused.meta: + fused.meta[key] = (fused.meta[key],) + _tag_for_regional_inductor(fused) + output = gm.graph.call_function(operator.getitem, (fused, 0)) + output.meta = _copy_meta(mm, cast) + _tag_for_regional_inductor(output) + + cast.replace_all_uses_with(output) + gm.graph.erase_node(cast) + gm.graph.erase_node(mm) + num_fused += 1 + + gm.graph.lint() + gm.recompile() + logger.info(f"B6 fused {num_fused} BF16 weight-gradient GEMM cast chains") + return gm + + +def _f6_bias_add( + sigmoid: torch.fx.Node, + output_width, +) -> tuple[torch.fx.Node, torch.fx.Node] | None: + matches = [] + for user in sigmoid.users: + if user.target != _ADD or user.args[0] is not sigmoid: + continue + if user.kwargs.get("alpha", 1) != 1 or len(user.args) < 2: + continue + bias = user.args[1] + if not isinstance(bias, torch.fx.Node): + continue + if _node_dtype(bias) != torch.float32 or _node_shape(bias) != (output_width,): + continue + matches.append((user, bias)) + if len(matches) != 1: + return None + return matches[0] + + +def _match_f6_router_sigmoid( + sigmoid: torch.fx.Node, +) -> tuple[ + torch.fx.Node, torch.fx.Node, torch.fx.Node | None, torch.fx.Node | None +] | None: + if sigmoid.target != _SIGMOID or _node_dtype(sigmoid) != torch.float32: + return None + if len(sigmoid.args) != 1 or not isinstance(sigmoid.args[0], torch.fx.Node): + return None + reshape = sigmoid.args[0] + if reshape.target not in (_RESHAPE, _VIEW) or _sole_user(reshape) is not sigmoid: + return None + if not reshape.args or not isinstance(reshape.args[0], torch.fx.Node): + return None + mm = reshape.args[0] + if mm.target != _MM or _sole_user(mm) is not reshape: + return None + if _node_dtype(mm) != torch.float32: + return None + + mm_shape = _node_shape(mm) + reshape_shape = _node_shape(reshape) + if ( + mm_shape is None + or len(mm_shape) != 2 + or reshape_shape is None + or len(reshape_shape) < 2 + or reshape_shape[-1] != mm_shape[-1] + ): + return None + + bias_match = _f6_bias_add(sigmoid, mm_shape[-1]) + if bias_match is None: + return mm, reshape, None, None + bias_add, bias = bias_match + return mm, reshape, bias_add, bias + + +def fuse_f6_router_sigmoid_bias_pass( + gm: torch.fx.GraphModule, + example_inputs: tuple | None = None, +) -> torch.fx.GraphModule: + """Fuse router GEMM sigmoid and optional expert bias into FlexGEMM.""" + del example_inputs + num_fused = 0 + num_bias_fused = 0 + for sigmoid in list(gm.graph.nodes): + match = _match_f6_router_sigmoid(sigmoid) + if match is None: + continue + mm, reshape, bias_add, bias = match + mm_shape = _node_shape(mm) + output_shape = _node_shape(sigmoid) + assert mm_shape is not None + assert output_shape is not None + + bias_2d = None + if bias is not None: + with gm.graph.inserting_before(mm): + bias_2d = gm.graph.call_function( + _RESHAPE, + (bias, [1, mm_shape[-1]]), + ) + bias_2d.meta = _copy_meta(bias) + bias_value = bias.meta.get("val") + if isinstance(bias_value, torch.Tensor): + bias_2d.meta["val"] = bias_value.reshape(1, -1) + bias_2d.meta.pop("tensor_meta", None) + _tag_for_regional_inductor(bias_2d) + + body = _build_f6_router_body(mm, sigmoid, bias_2d, bias_add) + body_name = _next_submodule_name(gm, "_coda_f6_body") + gm.add_module(body_name, body) + with gm.graph.inserting_before(mm): + body_ref = gm.graph.get_attr(body_name) + _tag_for_regional_inductor(body_ref) + hop_args = tuple(mm.args) + (() if bias_2d is None else (bias_2d,)) + fused = gm.graph.call_function( + flex_gemm_hop, + ( + _MM, + body_ref, + hop_args, + dict(mm.kwargs), + {"backend": "QUACK"}, + ), + ) + fused.meta = _copy_meta(mm, reshape, sigmoid) + value_meta = _copy_meta_with_value_from(mm, mm, reshape, sigmoid) + output_metas = [value_meta] + if bias_add is not None: + output_metas.append( + _copy_meta_with_value_from(mm, mm, reshape, sigmoid, bias_add) + ) + for key in ("val", "tensor_meta"): + values = [meta[key] for meta in output_metas if key in meta] + if len(values) == len(output_metas): + fused.meta[key] = tuple(values) + else: + fused.meta.pop(key, None) + _tag_for_regional_inductor(fused) + + raw_2d = gm.graph.call_function(operator.getitem, (fused, 0)) + raw_2d.meta = value_meta + _tag_for_regional_inductor(raw_2d) + raw_scores = gm.graph.call_function( + _RESHAPE, + (raw_2d, list(output_shape)), + ) + raw_scores.meta = _copy_meta(mm, reshape, sigmoid) + _tag_for_regional_inductor(raw_scores) + + biased_scores = None + if bias_add is not None: + biased_2d = gm.graph.call_function(operator.getitem, (fused, 1)) + biased_2d.meta = output_metas[1] + _tag_for_regional_inductor(biased_2d) + biased_scores = gm.graph.call_function( + _RESHAPE, + (biased_2d, list(output_shape)), + ) + biased_scores.meta = _copy_meta(mm, reshape, sigmoid, bias_add) + _tag_for_regional_inductor(biased_scores) + + sigmoid.replace_all_uses_with(raw_scores) + if bias_add is not None: + assert biased_scores is not None + bias_add.replace_all_uses_with(biased_scores) + gm.graph.erase_node(bias_add) + num_bias_fused += 1 + gm.graph.erase_node(sigmoid) + gm.graph.erase_node(reshape) + gm.graph.erase_node(mm) + num_fused += 1 + + gm.graph.lint() + gm.recompile() + logger.info( + f"F6 fused {num_fused} router GEMM sigmoid chains, " + f"including {num_bias_fused} expert-bias epilogues" + ) + return gm + + +def _match_f4_swiglu( + mul: torch.fx.Node, +) -> tuple[ + torch.fx.Node, + torch.fx.Node, + torch.fx.Node, + torch.fx.Node, + torch.fx.Node, +] | None: + if mul.target != _MUL or len(mul.args) < 2: + return None + silu, gate_reshape = mul.args[:2] + if not isinstance(silu, torch.fx.Node) or silu.target != _SILU: + return None + if not isinstance(gate_reshape, torch.fx.Node) or gate_reshape.target not in ( + _RESHAPE, + _VIEW, + ): + return None + if len(silu.args) != 1 or not isinstance(silu.args[0], torch.fx.Node): + return None + silu_reshape = silu.args[0] + if silu_reshape.target not in (_RESHAPE, _VIEW): + return None + if not silu_reshape.args or not isinstance(silu_reshape.args[0], torch.fx.Node): + return None + if not gate_reshape.args or not isinstance(gate_reshape.args[0], torch.fx.Node): + return None + silu_mm = silu_reshape.args[0] + gate_mm = gate_reshape.args[0] + if silu_mm.target != _MM or gate_mm.target != _MM: + return None + if _sole_user(silu_mm) is not silu_reshape: + return None + if _sole_user(gate_mm) is not gate_reshape: + return None + if any( + _node_dtype(node) != torch.bfloat16 + for node in (silu_mm, silu_reshape, silu, gate_mm, gate_reshape, mul) + ): + return None + + silu_mm_shape = _node_shape(silu_mm) + gate_mm_shape = _node_shape(gate_mm) + silu_shape = _node_shape(silu) + gate_shape = _node_shape(gate_reshape) + if ( + silu_mm_shape is None + or len(silu_mm_shape) != 2 + or silu_mm_shape != gate_mm_shape + or silu_shape is None + or silu_shape != gate_shape + or silu_shape != _node_shape(mul) + or silu_shape[-1] != silu_mm_shape[-1] + ): + return None + return silu_mm, silu_reshape, silu, gate_mm, gate_reshape + + +def fuse_f4_dense_swiglu_pass( + gm: torch.fx.GraphModule, + example_inputs: tuple | None = None, +) -> torch.fx.GraphModule: + """Fuse dense/shared-expert two-GEMM SwiGLU pointwise work.""" + del example_inputs + num_fused = 0 + for mul in list(gm.graph.nodes): + match = _match_f4_swiglu(mul) + if match is None: + continue + silu_mm, silu_reshape, silu, gate_mm, gate_reshape = match + output_shape = _node_shape(mul) + assert output_shape is not None + preserve_preactivation = any(user is not silu for user in silu_reshape.users) + + silu_body = _build_f4_silu_body( + silu_mm, + silu_reshape, + silu, + preserve_preactivation, + ) + silu_body_name = _next_submodule_name(gm, "_coda_f4_silu_body") + gm.add_module(silu_body_name, silu_body) + with gm.graph.inserting_before(silu_mm): + silu_body_ref = gm.graph.get_attr(silu_body_name) + _tag_for_regional_inductor(silu_body_ref) + silu_fused = gm.graph.call_function( + flex_gemm_hop, + ( + _MM, + silu_body_ref, + tuple(silu_mm.args), + dict(silu_mm.kwargs), + {"backend": "QUACK"}, + ), + ) + silu_2d_meta = _copy_meta_with_value_from( + silu_mm, silu_mm, silu_reshape, silu + ) + silu_fused.meta = _copy_meta(silu_mm, silu_reshape, silu) + for key in ("val", "tensor_meta"): + values = [silu_2d_meta[key]] if key in silu_2d_meta else [] + if preserve_preactivation and key in silu_2d_meta: + values.append(silu_2d_meta[key]) + if len(values) == 1 + int(preserve_preactivation): + silu_fused.meta[key] = tuple(values) + else: + silu_fused.meta.pop(key, None) + _tag_for_regional_inductor(silu_fused) + silu_2d = gm.graph.call_function(operator.getitem, (silu_fused, 0)) + silu_2d.meta = silu_2d_meta + _tag_for_regional_inductor(silu_2d) + silu_output = gm.graph.call_function( + _RESHAPE, + (silu_2d, list(output_shape)), + ) + silu_output.meta = _copy_meta(silu_mm, silu_reshape, silu) + _tag_for_regional_inductor(silu_output) + + preactivation_output = None + if preserve_preactivation: + preactivation_2d = gm.graph.call_function( + operator.getitem, + (silu_fused, 1), + ) + preactivation_2d.meta = silu_2d_meta + _tag_for_regional_inductor(preactivation_2d) + preactivation_output = gm.graph.call_function( + _RESHAPE, + (preactivation_2d, list(output_shape)), + ) + preactivation_output.meta = _copy_meta(silu_mm, silu_reshape) + _tag_for_regional_inductor(preactivation_output) + + gate_body = _build_f4_mul_body( + gate_mm, + gate_reshape, + silu_2d, + mul, + ) + gate_body_name = _next_submodule_name(gm, "_coda_f4_mul_body") + gm.add_module(gate_body_name, gate_body) + with gm.graph.inserting_before(gate_mm): + gate_body_ref = gm.graph.get_attr(gate_body_name) + _tag_for_regional_inductor(gate_body_ref) + gate_fused = gm.graph.call_function( + flex_gemm_hop, + ( + _MM, + gate_body_ref, + tuple(gate_mm.args) + (silu_2d,), + dict(gate_mm.kwargs), + {"backend": "QUACK"}, + ), + ) + gate_2d_meta = _copy_meta_with_value_from(gate_mm, gate_mm, gate_reshape) + product_2d_meta = _copy_meta_with_value_from( + gate_mm, silu, gate_mm, gate_reshape, mul + ) + gate_fused.meta = _copy_meta(silu, gate_mm, gate_reshape, mul) + for key in ("val", "tensor_meta"): + values = [ + meta[key] for meta in (gate_2d_meta, product_2d_meta) if key in meta + ] + if len(values) == 2: + gate_fused.meta[key] = tuple(values) + else: + gate_fused.meta.pop(key, None) + _tag_for_regional_inductor(gate_fused) + + gate_2d = gm.graph.call_function(operator.getitem, (gate_fused, 0)) + gate_2d.meta = gate_2d_meta + _tag_for_regional_inductor(gate_2d) + gate_output = gm.graph.call_function( + _RESHAPE, + (gate_2d, list(output_shape)), + ) + gate_output.meta = _copy_meta(gate_mm, gate_reshape) + _tag_for_regional_inductor(gate_output) + + product_2d = gm.graph.call_function(operator.getitem, (gate_fused, 1)) + product_2d.meta = product_2d_meta + _tag_for_regional_inductor(product_2d) + product = gm.graph.call_function( + _RESHAPE, + (product_2d, list(output_shape)), + ) + product.meta = _copy_meta(silu, gate_mm, gate_reshape, mul) + _tag_for_regional_inductor(product) + + silu.replace_all_uses_with(silu_output) + if preactivation_output is not None: + silu_reshape.replace_all_uses_with(preactivation_output) + gate_reshape.replace_all_uses_with(gate_output) + mul.replace_all_uses_with(product) + gm.graph.erase_node(mul) + gm.graph.erase_node(silu) + gm.graph.erase_node(silu_reshape) + gm.graph.erase_node(silu_mm) + gm.graph.erase_node(gate_reshape) + gm.graph.erase_node(gate_mm) + num_fused += 1 + + gm.graph.lint() + gm.recompile() + logger.info(f"F4 fused {num_fused} dense/shared-expert SwiGLU chains") + return gm + + +def _match_b2_branch_derivatives( + silu_grad: torch.fx.Node, +) -> tuple[ + torch.fx.Node, + torch.fx.Node, + torch.fx.Node, + torch.fx.Node, + torch.fx.Node, + torch.fx.Node, + torch.fx.Node, +] | None: + if silu_grad.target != _SILU_BACKWARD or len(silu_grad.args) < 2: + return None + gated_grad, saved_preactivation = silu_grad.args[:2] + if not isinstance(gated_grad, torch.fx.Node) or gated_grad.target != _MUL: + return None + if not isinstance(saved_preactivation, torch.fx.Node): + return None + if _sole_user(gated_grad) is not silu_grad or len(gated_grad.args) < 2: + return None + grad_view, saved_gate = gated_grad.args[:2] + if not isinstance(grad_view, torch.fx.Node) or grad_view.target not in ( + _RESHAPE, + _VIEW, + ): + return None + if not isinstance(saved_gate, torch.fx.Node): + return None + if not grad_view.args or not isinstance(grad_view.args[0], torch.fx.Node): + return None + mm = grad_view.args[0] + if mm.target != _MM or _sole_user(mm) is not grad_view: + return None + + sibling_grads = [ + user + for user in grad_view.users + if user is not gated_grad + and user.target == _MUL + and len(user.args) >= 2 + and user.args[0] is grad_view + and isinstance(user.args[1], torch.fx.Node) + ] + if len(grad_view.users) != 2 or len(sibling_grads) != 1: + return None + gate_grad = sibling_grads[0] + saved_silu = gate_grad.args[1] + assert isinstance(saved_silu, torch.fx.Node) + + nodes = ( + mm, + grad_view, + gate_grad, + gated_grad, + silu_grad, + saved_silu, + saved_gate, + saved_preactivation, + ) + if any(_node_dtype(node) != torch.bfloat16 for node in nodes): + return None + mm_shape = _node_shape(mm) + output_shape = _node_shape(grad_view) + if ( + mm_shape is None + or len(mm_shape) != 2 + or output_shape is None + or output_shape[-1] != mm_shape[-1] + or any(_node_shape(node) != output_shape for node in nodes[2:]) + ): + return None + return ( + mm, + grad_view, + gate_grad, + gated_grad, + silu_grad, + saved_silu, + saved_gate, + ) + + +def _match_b2_input_grad_add( + add: torch.fx.Node, +) -> tuple[torch.fx.Node, torch.fx.Node, torch.fx.Node, torch.fx.Node,] | None: + if add.target != _ADD or len(add.args) < 2: + return None + lhs, rhs = add.args[:2] + if not all( + isinstance(node, torch.fx.Node) and node.target in (_RESHAPE, _VIEW) + for node in (lhs, rhs) + ): + return None + assert isinstance(lhs, torch.fx.Node) + assert isinstance(rhs, torch.fx.Node) + if not lhs.args or not rhs.args: + return None + lhs_mm, rhs_mm = lhs.args[0], rhs.args[0] + if not all( + isinstance(node, torch.fx.Node) and node.target == _MM + for node in (lhs_mm, rhs_mm) + ): + return None + assert isinstance(lhs_mm, torch.fx.Node) + assert isinstance(rhs_mm, torch.fx.Node) + if _sole_user(lhs_mm) is not lhs or _sole_user(rhs_mm) is not rhs: + return None + if _sole_user(lhs) is not add or _sole_user(rhs) is not add: + return None + + lhs_fqn = _module_fqn(lhs_mm) + rhs_fqn = _module_fqn(rhs_mm) + if lhs_fqn is None or rhs_fqn is None or "." not in lhs_fqn or "." not in rhs_fqn: + return None + lhs_parent, lhs_role = lhs_fqn.rsplit(".", 1) + rhs_parent, rhs_role = rhs_fqn.rsplit(".", 1) + if lhs_parent != rhs_parent or {lhs_role, rhs_role} != {"w1", "w3"}: + return None + if lhs_role == "w3": + captured_mm, captured_reshape = lhs_mm, lhs + fused_mm, fused_reshape = rhs_mm, rhs + else: + captured_mm, captured_reshape = rhs_mm, rhs + fused_mm, fused_reshape = lhs_mm, lhs + seen_captured_mm = False + for node in add.graph.nodes: + if node is captured_mm: + seen_captured_mm = True + if node is fused_mm: + break + if not seen_captured_mm: + return None + + if any( + _node_dtype(node) != torch.bfloat16 + for node in (captured_mm, captured_reshape, fused_mm, fused_reshape, add) + ): + return None + if ( + _node_shape(captured_mm) != _node_shape(fused_mm) + or _node_shape(captured_reshape) != _node_shape(fused_reshape) + or _node_shape(captured_reshape) != _node_shape(add) + ): + return None + return captured_mm, captured_reshape, fused_mm, fused_reshape + + +def fuse_b2_dense_swiglu_backward_pass( + gm: torch.fx.GraphModule, + example_inputs: tuple | None = None, +) -> torch.fx.GraphModule: + """Fuse dense/shared-expert SwiGLU backward GEMM epilogues.""" + del example_inputs + num_branch_fused = 0 + for silu_grad in list(gm.graph.nodes): + match = _match_b2_branch_derivatives(silu_grad) + if match is None: + continue + ( + mm, + grad_view, + gate_grad, + gated_grad, + silu_grad, + saved_silu, + saved_gate, + ) = match + saved_preactivation = silu_grad.args[1] + assert isinstance(saved_preactivation, torch.fx.Node) + mm_shape = _node_shape(mm) + output_shape = _node_shape(grad_view) + assert mm_shape is not None + assert output_shape is not None + + captures = [] + with gm.graph.inserting_before(mm): + for saved in (saved_silu, saved_gate, saved_preactivation): + captured = gm.graph.call_function( + _RESHAPE, + (saved, list(mm_shape)), + ) + captured.meta = _copy_meta_with_value_from(mm, saved) + _tag_for_regional_inductor(captured) + captures.append(captured) + + body = _build_b2_branch_body( + mm, + grad_view, + captures[0], + captures[1], + captures[2], + gate_grad, + silu_grad, + ) + body_name = _next_submodule_name(gm, "_coda_b2_branch_body") + gm.add_module(body_name, body) + with gm.graph.inserting_before(mm): + body_ref = gm.graph.get_attr(body_name) + _tag_for_regional_inductor(body_ref) + fused = gm.graph.call_function( + flex_gemm_hop, + ( + _MM, + body_ref, + tuple(mm.args) + tuple(captures), + dict(mm.kwargs), + {"backend": "QUACK"}, + ), + ) + gate_grad_2d_meta = _copy_meta_with_value_from(mm, mm, grad_view, gate_grad) + silu_grad_2d_meta = _copy_meta_with_value_from(mm, mm, grad_view, silu_grad) + fused.meta = _copy_meta(mm, grad_view, gate_grad, silu_grad) + for key in ("val", "tensor_meta"): + values = [ + meta[key] + for meta in (gate_grad_2d_meta, silu_grad_2d_meta) + if key in meta + ] + if len(values) == 2: + fused.meta[key] = tuple(values) + else: + fused.meta.pop(key, None) + _tag_for_regional_inductor(fused) + + gate_grad_2d = gm.graph.call_function(operator.getitem, (fused, 0)) + gate_grad_2d.meta = gate_grad_2d_meta + _tag_for_regional_inductor(gate_grad_2d) + gate_grad_output = gm.graph.call_function( + _RESHAPE, + (gate_grad_2d, list(output_shape)), + ) + gate_grad_output.meta = _copy_meta(mm, grad_view, gate_grad) + _tag_for_regional_inductor(gate_grad_output) + + silu_grad_2d = gm.graph.call_function(operator.getitem, (fused, 1)) + silu_grad_2d.meta = silu_grad_2d_meta + _tag_for_regional_inductor(silu_grad_2d) + silu_grad_output = gm.graph.call_function( + _RESHAPE, + (silu_grad_2d, list(output_shape)), + ) + silu_grad_output.meta = _copy_meta(mm, grad_view, silu_grad) + _tag_for_regional_inductor(silu_grad_output) + + gate_grad.replace_all_uses_with(gate_grad_output) + silu_grad.replace_all_uses_with(silu_grad_output) + gm.graph.erase_node(silu_grad) + gm.graph.erase_node(gated_grad) + gm.graph.erase_node(gate_grad) + gm.graph.erase_node(grad_view) + gm.graph.erase_node(mm) + num_branch_fused += 1 + + num_input_add_fused = 0 + for add in list(gm.graph.nodes): + match = _match_b2_input_grad_add(add) + if match is None: + continue + captured_mm, captured_reshape, fused_mm, fused_reshape = match + output_shape = _node_shape(add) + assert output_shape is not None + + body = _build_b2_input_add_body( + fused_mm, + fused_reshape, + captured_mm, + add, + ) + body_name = _next_submodule_name(gm, "_coda_b2_input_add_body") + gm.add_module(body_name, body) + with gm.graph.inserting_before(fused_mm): + body_ref = gm.graph.get_attr(body_name) + _tag_for_regional_inductor(body_ref) + fused = gm.graph.call_function( + flex_gemm_hop, + ( + _MM, + body_ref, + tuple(fused_mm.args) + (captured_mm,), + dict(fused_mm.kwargs), + {"backend": "QUACK"}, + ), + ) + total_2d_meta = _copy_meta_with_value_from( + fused_mm, captured_reshape, fused_mm, fused_reshape, add + ) + fused.meta = _copy_meta(captured_reshape, fused_mm, fused_reshape, add) + for key in ("val", "tensor_meta"): + if key in total_2d_meta: + fused.meta[key] = (total_2d_meta[key],) + else: + fused.meta.pop(key, None) + _tag_for_regional_inductor(fused) + total_2d = gm.graph.call_function(operator.getitem, (fused, 0)) + total_2d.meta = total_2d_meta + _tag_for_regional_inductor(total_2d) + total = gm.graph.call_function( + _RESHAPE, + (total_2d, list(output_shape)), + ) + total.meta = _copy_meta(captured_reshape, fused_reshape, add) + _tag_for_regional_inductor(total) + + add.replace_all_uses_with(total) + gm.graph.erase_node(add) + gm.graph.erase_node(captured_reshape) + gm.graph.erase_node(fused_reshape) + gm.graph.erase_node(fused_mm) + num_input_add_fused += 1 + + gm.graph.lint() + gm.recompile() + logger.info( + f"B2 fused {num_branch_fused} branch-derivative GEMMs and " + f"{num_input_add_fused} input-gradient GEMM adds" + ) + return gm + + +def _match_b4_router_input_grad_add( + add: torch.fx.Node, +) -> tuple[torch.fx.Node, torch.fx.Node, torch.fx.Node, torch.fx.Node,] | None: + if add.target != _ADD or len(add.args) < 2: + return None + lhs, rhs = add.args[:2] + cast = next( + ( + node + for node in (lhs, rhs) + if isinstance(node, torch.fx.Node) and _is_cast(node, torch.bfloat16) + ), + None, + ) + if cast is None: + return None + residual = rhs if cast is lhs else lhs + if not isinstance(residual, torch.fx.Node): + return None + if _sole_user(cast) is not add or not cast.args: + return None + reshape = cast.args[0] + if not isinstance(reshape, torch.fx.Node) or reshape.target not in ( + _RESHAPE, + _VIEW, + ): + return None + if _sole_user(reshape) is not cast or not reshape.args: + return None + mm = reshape.args[0] + if not isinstance(mm, torch.fx.Node) or mm.target != _MM: + return None + if _sole_user(mm) is not reshape: + return None + + module_fqn = _module_fqn(mm) + if module_fqn is None or not module_fqn.endswith(".moe.router.gate"): + return None + mm_shape = _node_shape(mm) + output_shape = _node_shape(add) + if ( + mm_shape is None + or len(mm_shape) != 2 + or output_shape is None + or output_shape[-1] != mm_shape[-1] + or _node_shape(reshape) != output_shape + or _node_shape(cast) != output_shape + or _node_shape(residual) != output_shape + ): + return None + if _node_dtype(mm) != torch.float32 or _node_dtype(reshape) != torch.float32: + return None + if any(_node_dtype(node) != torch.bfloat16 for node in (cast, residual, add)): + return None + if not all( + isinstance(node.meta.get("val"), torch.Tensor) + for node in (mm, cast, residual, add) + ): + return None + return mm, reshape, cast, residual + + +def fuse_b4_router_input_grad_add_pass( + gm: torch.fx.GraphModule, + example_inputs: tuple | None = None, +) -> torch.fx.GraphModule: + """Fuse the router input-gradient cast and expert-gradient add.""" + del example_inputs + num_fused = 0 + for add in list(gm.graph.nodes): + match = _match_b4_router_input_grad_add(add) + if match is None: + continue + mm, reshape, cast, residual = match + mm_shape = _node_shape(mm) + output_shape = _node_shape(add) + assert mm_shape is not None + assert output_shape is not None + residual_value = residual.meta["val"] + cast_value = cast.meta["val"] + assert isinstance(residual_value, torch.Tensor) + assert isinstance(cast_value, torch.Tensor) + + with gm.graph.inserting_before(mm): + residual_2d = gm.graph.call_function( + _RESHAPE, + (residual, list(mm_shape)), + ) + residual_2d.meta = _copy_meta_with_value( + residual_value.reshape(mm_shape), residual + ) + _tag_for_regional_inductor(residual_2d) + body = _build_b4_router_input_grad_body(mm, cast, residual_2d, add) + body_name = _next_submodule_name(gm, "_coda_b4_router_body") + gm.add_module(body_name, body) + body_ref = gm.graph.get_attr(body_name) + _tag_for_regional_inductor(body_ref) + fused = gm.graph.call_function( + flex_gemm_hop, + ( + _MM, + body_ref, + tuple(mm.args) + (residual_2d,), + dict(mm.kwargs), + {"backend": "QUACK"}, + ), + ) + total_2d_meta = _copy_meta_with_value( + cast_value.reshape(mm_shape), mm, reshape, cast, residual, add + ) + fused.meta = _copy_meta(mm, reshape, cast, residual, add) + fused.meta["val"] = (total_2d_meta["val"],) + fused.meta.pop("tensor_meta", None) + _tag_for_regional_inductor(fused) + total_2d = gm.graph.call_function(operator.getitem, (fused, 0)) + total_2d.meta = total_2d_meta + _tag_for_regional_inductor(total_2d) + total = gm.graph.call_function( + _RESHAPE, + (total_2d, list(output_shape)), + ) + total.meta = _copy_meta(mm, reshape, cast, residual, add) + _tag_for_regional_inductor(total) + + add.replace_all_uses_with(total) + gm.graph.erase_node(add) + gm.graph.erase_node(cast) + gm.graph.erase_node(reshape) + gm.graph.erase_node(mm) + num_fused += 1 + + gm.graph.lint() + gm.recompile() + logger.info(f"B4 fused {num_fused} router input-gradient cast/add chains") + return gm + + +def _match_b7_attention_grad_merge( + add: torch.fx.Node, +) -> tuple[torch.fx.Node, torch.fx.Node, torch.fx.Node, torch.fx.Node, bool,] | None: + if add.target != _ADD or len(add.args) < 2 or add.kwargs.get("alpha", 1) != 1: + return None + + branches = [] + for index, operand in enumerate(add.args[:2]): + if ( + not isinstance(operand, torch.fx.Node) + or operand.target not in (_RESHAPE, _VIEW) + or not operand.args + or not isinstance(operand.args[0], torch.fx.Node) + ): + return None + mm = operand.args[0] + if ( + mm.target != _MM + or _sole_user(mm) is not operand + or _sole_user(operand) is not add + or mm.meta.get("autograd_backward") is not True + ): + return None + role = _layer_module_role(_module_fqn(mm)) + if role is None: + return None + branches.append((index, mm, operand, role)) + + by_role = {role: (index, mm, reshape) for index, mm, reshape, (_, role) in branches} + if set(by_role) != {"attention.wkv_a", "attention.wq_a"}: + return None + kv_index, kv_mm, kv_reshape = by_role["attention.wkv_a"] + q_index, q_mm, q_reshape = by_role["attention.wq_a"] + kv_layer = next(layer for _, mm, _, (layer, _) in branches if mm is kv_mm) + q_layer = next(layer for _, mm, _, (layer, _) in branches if mm is q_mm) + if kv_layer != q_layer: + return None + + q_shape = _node_shape(q_mm) + add_shape = _node_shape(add) + if ( + q_shape is None + or len(q_shape) != 2 + or _node_shape(kv_mm) != q_shape + or add_shape is None + or _node_shape(kv_reshape) != add_shape + or _node_shape(q_reshape) != add_shape + or _node_dtype(kv_mm) != torch.bfloat16 + or _node_dtype(q_mm) != torch.bfloat16 + or _node_dtype(add) != torch.bfloat16 + or not all( + isinstance(node.meta.get("val"), torch.Tensor) for node in (kv_mm, q_mm) + ) + ): + return None + assert kv_index != q_index + return q_mm, q_reshape, kv_mm, kv_reshape, q_index == 0 + + +def fuse_b7_attention_grad_merge_pass( + gm: torch.fx.GraphModule, + example_inputs: tuple | None = None, +) -> torch.fx.GraphModule: + """Fuse the Q/KV attention input-gradient merge into the Q GEMM.""" + del example_inputs + num_fused = 0 + for add in list(gm.graph.nodes): + match = _match_b7_attention_grad_merge(add) + if match is None: + continue + q_mm, q_reshape, kv_mm, kv_reshape, q_is_lhs = match + q_shape = _node_shape(q_mm) + add_shape = _node_shape(add) + assert q_shape is not None + assert add_shape is not None + total_value = add.meta["val"] + assert isinstance(total_value, torch.Tensor) + + body = _build_b7_attention_grad_merge_body(q_mm, kv_mm, q_is_lhs) + body_name = _next_submodule_name(gm, "_coda_b7_attention_grad_merge_body") + gm.add_module(body_name, body) + with gm.graph.inserting_before(q_mm): + body_ref = gm.graph.get_attr(body_name) + _tag_for_regional_inductor(body_ref) + fused = gm.graph.call_function( + flex_gemm_hop, + ( + _MM, + body_ref, + tuple(q_mm.args) + (kv_mm,), + dict(q_mm.kwargs), + {"backend": "QUACK"}, + ), + ) + total_2d_meta = _copy_meta_with_value( + total_value.reshape(q_shape), + kv_mm, + kv_reshape, + q_mm, + q_reshape, + add, + ) + fused.meta = _copy_meta(kv_mm, kv_reshape, q_mm, q_reshape, add) + fused.meta["val"] = (total_2d_meta["val"],) + fused.meta.pop("tensor_meta", None) + _tag_for_regional_inductor(fused) + total_2d = gm.graph.call_function(operator.getitem, (fused, 0)) + total_2d.meta = total_2d_meta + _tag_for_regional_inductor(total_2d) + total = gm.graph.call_function(_RESHAPE, (total_2d, list(add_shape))) + total.meta = _copy_meta_with_value(total_value.reshape(add_shape), add) + _tag_for_regional_inductor(total) + + add.replace_all_uses_with(total) + gm.graph.erase_node(add) + gm.graph.erase_node(q_reshape) + gm.graph.erase_node(q_mm) + gm.graph.erase_node(kv_reshape) + num_fused += 1 + + gm.graph.lint() + gm.recompile() + logger.info(f"B7 fused {num_fused} Q/KV attention input-gradient merges") + return gm + + +def _match_b5_mla_rmsnorm_backward( + norm: torch.fx.Node, +) -> tuple[ + torch.fx.Node, + torch.fx.Node, + torch.fx.Node, + torch.fx.Node, + torch.fx.Node, + torch.fx.Node, + torch.fx.Node, +] | None: + if norm.target != _FUSED_RMS_NORM_BACKWARD or len(norm.args) < 6: + return None + grad_view, norm_input, normalized_shape, rstd, gamma, output_mask = norm.args[:6] + if ( + not isinstance(grad_view, torch.fx.Node) + or grad_view.target not in (_RESHAPE, _VIEW) + or not grad_view.args + or not isinstance(grad_view.args[0], torch.fx.Node) + or _sole_user(grad_view) is not norm + or not isinstance(norm_input, torch.fx.Node) + or not isinstance(rstd, torch.fx.Node) + or not isinstance(gamma, torch.fx.Node) + or not isinstance(normalized_shape, (list, tuple)) + or len(normalized_shape) != 1 + or list(output_mask) != [True, True] + or norm.meta.get("autograd_backward") is not True + ): + return None + mm = grad_view.args[0] + module_fqn = _module_fqn(mm) + if ( + mm.target != _MM + or _sole_user(mm) is not grad_view + or mm.meta.get("autograd_backward") is not True + or module_fqn is None + or not module_fqn.endswith((".attention.wkv_b", ".attention.wq_b")) + ): + return None + + mm_shape = _node_shape(mm) + grad_shape = _node_shape(grad_view) + input_shape = _node_shape(norm_input) + rstd_shape = _node_shape(rstd) + gamma_shape = _node_shape(gamma) + if ( + mm_shape is None + or len(mm_shape) != 2 + or grad_shape is None + or input_shape != grad_shape + or rstd_shape != (*grad_shape[:-1], 1) + or gamma_shape != (grad_shape[-1],) + or normalized_shape[0] != grad_shape[-1] + or mm_shape != (math.prod(grad_shape[:-1]), grad_shape[-1]) + or grad_shape[-1] % _CODA_RMSNORM_BACKWARD_GROUP != 0 + or _node_dtype(mm) != torch.bfloat16 + or _node_dtype(grad_view) != torch.bfloat16 + or _node_dtype(norm_input) != torch.bfloat16 + or _node_dtype(rstd) != torch.float32 + or _node_dtype(gamma) != torch.bfloat16 + ): + return None + + outputs: dict[int, torch.fx.Node] = {} + for user in norm.users: + if ( + user.target is not operator.getitem + or len(user.args) < 2 + or user.args[0] is not norm + or user.args[1] not in (0, 1) + ): + return None + outputs[user.args[1]] = user + if set(outputs) != {0, 1}: + return None + if ( + _node_shape(outputs[0]) != input_shape + or _node_dtype(outputs[0]) != torch.bfloat16 + or _node_shape(outputs[1]) != gamma_shape + or _node_dtype(outputs[1]) != torch.bfloat16 + ): + return None + return mm, grad_view, norm_input, rstd, gamma, outputs[0], outputs[1] + + +def fuse_b5_mla_rmsnorm_backward_pass( + gm: torch.fx.GraphModule, + example_inputs: tuple | None = None, +) -> torch.fx.GraphModule: + """Fuse MLA input-gradient GEMMs with RMSNorm backward row partials.""" + del example_inputs + num_fused = 0 + for norm in list(gm.graph.nodes): + match = _match_b5_mla_rmsnorm_backward(norm) + if match is None: + continue + mm, grad_view, norm_input, rstd, gamma, grad_input, grad_weight = match + mm_shape = _node_shape(mm) + input_shape = _node_shape(norm_input) + assert mm_shape is not None + assert input_shape is not None + mm_value = mm.meta["val"] + norm_input_value = norm_input.meta["val"] + rstd_value = rstd.meta["val"] + gamma_value = gamma.meta["val"] + assert isinstance(mm_value, torch.Tensor) + assert isinstance(norm_input_value, torch.Tensor) + assert isinstance(rstd_value, torch.Tensor) + assert isinstance(gamma_value, torch.Tensor) + + with gm.graph.inserting_before(norm): + norm_input_2d = gm.graph.call_function( + _RESHAPE, (norm_input, list(mm_shape)) + ) + norm_input_2d_value = norm_input_value.reshape(mm_shape) + norm_input_2d.meta = _copy_meta_with_value( + norm_input_2d_value, norm_input, norm + ) + _tag_for_regional_inductor(norm_input_2d) + rstd_2d = gm.graph.call_function(_RESHAPE, (rstd, [mm_shape[0], 1])) + rstd_2d_value = rstd_value.reshape(mm_shape[0], 1) + rstd_2d.meta = _copy_meta_with_value(rstd_2d_value, rstd, norm) + _tag_for_regional_inductor(rstd_2d) + gamma_2d = gm.graph.call_function(_RESHAPE, (gamma, [1, mm_shape[1]])) + gamma_2d_value = gamma_value.reshape(1, mm_shape[1]) + gamma_2d.meta = _copy_meta_with_value(gamma_2d_value, gamma, norm) + _tag_for_regional_inductor(gamma_2d) + + body = _build_b5_mla_rmsnorm_backward_body( + mm, norm_input_2d, rstd_2d, gamma_2d + ) + body_name = _next_submodule_name(gm, "_coda_b5_mla_rmsnorm_body") + gm.add_module(body_name, body) + body_ref = gm.graph.get_attr(body_name) + _tag_for_regional_inductor(body_ref) + fused = gm.graph.call_function( + flex_gemm_hop, + ( + _MM, + body_ref, + tuple(mm.args) + (norm_input_2d, rstd_2d, gamma_2d), + dict(mm.kwargs), + {"backend": "QUACK"}, + ), + ) + + x_hat_value = norm_input_2d_value.float() * rstd_2d_value + grad_x_hat_value = mm_value.float() * gamma_2d_value.float() + row_products_value = x_hat_value * grad_x_hat_value + partial_value = row_products_value.reshape( + mm_shape[0], -1, _CODA_RMSNORM_BACKWARD_GROUP + ).sum(-1) + fused.meta = _copy_meta(mm, norm_input, rstd, gamma, norm) + fused.meta["val"] = (mm_value, partial_value) + fused.meta.pop("tensor_meta", None) + _tag_for_regional_inductor(fused) + rounded = gm.graph.call_function(operator.getitem, (fused, 0)) + rounded.meta = _copy_meta_with_value(mm_value, mm, grad_view, norm) + _tag_for_regional_inductor(rounded) + partial = gm.graph.call_function(operator.getitem, (fused, 1)) + partial.meta = _copy_meta_with_value( + partial_value, mm, norm_input, rstd, gamma, norm + ) + _tag_for_regional_inductor(partial) + + row_dot = gm.graph.call_function(_SUM_DIM, (partial, [-1], True)) + row_dot_value = partial_value.sum(-1, keepdim=True) + row_dot.meta = _copy_meta_with_value( + row_dot_value, mm, norm_input, rstd, gamma, norm + ) + _tag_for_regional_inductor(row_dot) + rounded_fp32 = gm.graph.call_function( + _TO_COPY, (rounded,), {"dtype": torch.float32} + ) + rounded_fp32.meta = _copy_meta_with_value(mm_value.float(), mm, norm) + _tag_for_regional_inductor(rounded_fp32) + norm_input_fp32 = gm.graph.call_function( + _TO_COPY, (norm_input_2d,), {"dtype": torch.float32} + ) + norm_input_fp32.meta = _copy_meta_with_value( + norm_input_2d_value.float(), norm_input, norm + ) + _tag_for_regional_inductor(norm_input_fp32) + x_hat = gm.graph.call_function(_MUL, (norm_input_fp32, rstd_2d)) + x_hat.meta = _copy_meta_with_value(x_hat_value, norm_input, rstd, norm) + _tag_for_regional_inductor(x_hat) + gamma_fp32 = gm.graph.call_function( + _TO_COPY, (gamma_2d,), {"dtype": torch.float32} + ) + gamma_fp32.meta = _copy_meta_with_value(gamma_2d_value.float(), gamma, norm) + _tag_for_regional_inductor(gamma_fp32) + grad_x_hat = gm.graph.call_function(_MUL, (rounded_fp32, gamma_fp32)) + grad_x_hat.meta = _copy_meta_with_value(grad_x_hat_value, mm, gamma, norm) + _tag_for_regional_inductor(grad_x_hat) + scaled_x_hat = gm.graph.call_function(_DIV_SCALAR, (x_hat, mm_shape[1])) + scaled_x_hat_value = x_hat_value / mm_shape[1] + scaled_x_hat.meta = _copy_meta_with_value( + scaled_x_hat_value, norm_input, norm + ) + _tag_for_regional_inductor(scaled_x_hat) + correction = gm.graph.call_function(_MUL, (scaled_x_hat, row_dot)) + correction_value = scaled_x_hat_value * row_dot_value + correction.meta = _copy_meta_with_value( + correction_value, mm, norm_input, rstd, gamma, norm + ) + _tag_for_regional_inductor(correction) + centered = gm.graph.call_function(_SUB, (grad_x_hat, correction)) + centered_value = grad_x_hat_value - correction_value + centered.meta = _copy_meta_with_value( + centered_value, mm, norm_input, rstd, gamma, norm + ) + _tag_for_regional_inductor(centered) + grad_input_fp32 = gm.graph.call_function(_MUL, (centered, rstd_2d)) + grad_input_fp32_value = centered_value * rstd_2d_value + grad_input_fp32.meta = _copy_meta_with_value( + grad_input_fp32_value, mm, norm_input, rstd, gamma, norm + ) + _tag_for_regional_inductor(grad_input_fp32) + grad_input_2d = gm.graph.call_function( + _TO_COPY, (grad_input_fp32,), {"dtype": torch.bfloat16} + ) + grad_input_2d_value = grad_input_fp32_value.to(torch.bfloat16) + grad_input_2d.meta = _copy_meta_with_value(grad_input_2d_value, grad_input) + _tag_for_regional_inductor(grad_input_2d) + new_grad_input = gm.graph.call_function( + _RESHAPE, (grad_input_2d, list(input_shape)) + ) + new_grad_input.meta = _copy_meta_with_value( + grad_input_2d_value.reshape(input_shape), grad_input + ) + _tag_for_regional_inductor(new_grad_input) + + grad_weight_terms = gm.graph.call_function(_MUL, (rounded_fp32, x_hat)) + grad_weight_terms_value = mm_value.float() * x_hat_value + grad_weight_terms.meta = _copy_meta_with_value( + grad_weight_terms_value, mm, norm_input, rstd, norm + ) + _tag_for_regional_inductor(grad_weight_terms) + grad_weight_fp32 = gm.graph.call_function( + _SUM_DIM, (grad_weight_terms, [0]) + ) + grad_weight_fp32_value = grad_weight_terms_value.sum(0) + grad_weight_fp32.meta = _copy_meta_with_value( + grad_weight_fp32_value, grad_weight + ) + _tag_for_regional_inductor(grad_weight_fp32) + new_grad_weight = gm.graph.call_function( + _TO_COPY, (grad_weight_fp32,), {"dtype": torch.bfloat16} + ) + new_grad_weight.meta = _copy_meta_with_value( + grad_weight_fp32_value.to(torch.bfloat16), grad_weight + ) + _tag_for_regional_inductor(new_grad_weight) + + grad_input.replace_all_uses_with(new_grad_input) + grad_weight.replace_all_uses_with(new_grad_weight) + gm.graph.erase_node(grad_input) + gm.graph.erase_node(grad_weight) + gm.graph.erase_node(norm) + gm.graph.erase_node(grad_view) + gm.graph.erase_node(mm) + num_fused += 1 + + gm.graph.lint() + gm.recompile() + logger.info(f"B5 fused {num_fused} MLA GEMM plus RMSNorm backward chains") + return gm + + +def _layer_module_role(module_fqn: str | None) -> tuple[int, str] | None: + if module_fqn is None: + return None + parts = module_fqn.split(".") + if len(parts) < 3 or parts[0] != "layers" or not parts[1].isdigit(): + return None + return int(parts[1]), ".".join(parts[2:]) + + +def _is_f3_producer_norm_pair( + first_mm: torch.fx.Node, + norm: torch.fx.Node, +) -> bool: + first = _layer_module_role(_module_fqn(first_mm)) + normalized = _layer_module_role(_module_fqn(norm)) + if first is None or normalized is None: + return False + first_layer, first_role = first + norm_layer, norm_role = normalized + if ( + first_role == "attention.wo" + and norm_layer == first_layer + and norm_role == "ffn_norm" + ): + return True + if ( + first_role == "feed_forward.w2" + and norm_layer == first_layer + 1 + and norm_role == "attention_norm" + ): + return True + if ( + first_role == "moe.shared_experts.w2" + and norm_layer == first_layer + 1 + and norm_role == "attention_norm" + ): + return True + return False + + +def _match_f3_add_tree( + add: torch.fx.Node, + norm: torch.fx.Node, + depth: int = 1, +) -> list[ + tuple[ + torch.fx.Node, + torch.fx.Node, + tuple[torch.fx.Node, ...], + tuple[bool, ...], + tuple[torch.fx.Node, ...], + ] +]: + if add.target != _ADD or len(add.args) < 2 or add.kwargs.get("alpha", 1) != 1: + return [] + + matches = [] + for index, operand in enumerate(add.args[:2]): + sibling = add.args[1 - index] + if not isinstance(operand, torch.fx.Node) or not isinstance( + sibling, torch.fx.Node + ): + continue + if ( + operand.target in (_RESHAPE, _VIEW) + and operand.args + and isinstance(operand.args[0], torch.fx.Node) + ): + candidate_mm = operand.args[0] + if ( + candidate_mm.target == _MM + and _sole_user(candidate_mm) is operand + and _sole_user(operand) is add + and _is_f3_producer_norm_pair(candidate_mm, norm) + ): + matches.append( + ( + candidate_mm, + operand, + (sibling,), + (index == 0,), + (add,), + ) + ) + if depth == 1 or operand.target != _ADD or _sole_user(operand) is not add: + continue + for candidate_mm, reshape, addends, orders, add_nodes in _match_f3_add_tree( + operand, norm, depth - 1 + ): + matches.append( + ( + candidate_mm, + reshape, + (*addends, sibling), + (*orders, index == 0), + (*add_nodes, add), + ) + ) + return matches + + +def _match_f3_residual_rmsnorm( + norm: torch.fx.Node, +) -> ( + tuple[ + torch.fx.Node, + torch.fx.Node, + tuple[torch.fx.Node, ...], + tuple[bool, ...], + tuple[torch.fx.Node, ...], + torch.fx.Node, + torch.fx.Node, + torch.fx.Node, + torch.fx.Node | None, + float, + ] + | None +): + if norm.target != _FUSED_RMS_NORM or len(norm.args) < 4: + return None + add, normalized_shape, gamma, eps = norm.args[:4] + if ( + not isinstance(add, torch.fx.Node) + or add.target != _ADD + or len(add.args) < 2 + or add.kwargs.get("alpha", 1) != 1 + or not isinstance(gamma, torch.fx.Node) + or not isinstance(normalized_shape, (list, tuple)) + or len(normalized_shape) != 1 + or not isinstance(eps, float) + ): + return None + + producer_matches = _match_f3_add_tree(add, norm, depth=2) + if len(producer_matches) != 1: + return None + ( + first_mm, + first_reshape, + addends, + accumulated_value_is_lhs, + add_nodes, + ) = producer_matches[0] + + norm_outputs: dict[int, torch.fx.Node] = {} + for user in norm.users: + if ( + user.target is not operator.getitem + or len(user.args) < 2 + or not isinstance(user.args[1], int) + or user.args[1] in norm_outputs + ): + return None + norm_outputs[user.args[1]] = user + if set(norm_outputs) not in ({0}, {0, 1}): + return None + norm_output = norm_outputs[0] + rstd_output = norm_outputs.get(1) + + first_shape = _node_shape(first_mm) + output_shape = _node_shape(add) + if ( + first_shape is None + or len(first_shape) != 2 + or output_shape is None + or output_shape[-1] != first_shape[-1] + or first_shape[-1] != normalized_shape[0] + or first_shape[-1] % _CODA_RMSNORM_GROUP != 0 + or _node_shape(first_reshape) != output_shape + or any(_node_shape(addend) != output_shape for addend in addends) + or _node_shape(norm_output) != output_shape + or _node_shape(gamma) != (first_shape[-1],) + ): + return None + if rstd_output is not None and ( + _node_dtype(rstd_output) != torch.float32 + or _node_shape(rstd_output) != (*output_shape[:-1], 1) + ): + return None + tensor_nodes = ( + first_mm, + first_reshape, + *addends, + *add_nodes, + gamma, + norm_output, + ) + if any(_node_dtype(node) != torch.bfloat16 for node in tensor_nodes): + return None + if not all( + isinstance(node.meta.get("val"), torch.Tensor) + for node in (first_mm, *addends, gamma) + ): + return None + return ( + first_mm, + first_reshape, + addends, + accumulated_value_is_lhs, + add_nodes, + norm, + norm_output, + gamma, + rstd_output, + eps, + ) + + +def fuse_f3_residual_rmsnorm_pass( + gm: torch.fx.GraphModule, + example_inputs: tuple | None = None, +) -> torch.fx.GraphModule: + """Fuse a producer GEMM, residual additions, and RMSNorm preparation.""" + del example_inputs + num_attention_output_fused = 0 + num_layer_output_fused = 0 + for norm in list(gm.graph.nodes): + match = _match_f3_residual_rmsnorm(norm) + if match is None: + continue + ( + first_mm, + first_reshape, + addends, + accumulated_value_is_lhs, + add_nodes, + norm, + norm_output, + gamma, + rstd_output, + eps, + ) = match + add = add_nodes[-1] + first_shape = _node_shape(first_mm) + output_shape = _node_shape(add) + assert first_shape is not None + assert output_shape is not None + first_value = first_mm.meta["val"] + addend_values = tuple(addend.meta["val"] for addend in addends) + gamma_value = gamma.meta["val"] + assert isinstance(first_value, torch.Tensor) + assert all(isinstance(value, torch.Tensor) for value in addend_values) + assert isinstance(gamma_value, torch.Tensor) + # Bucketing may materialize the next layer's RMSNorm weight after the + # producing GEMM. Insert at the norm boundary where every capture is + # guaranteed to dominate the new HOP. + with gm.graph.inserting_before(norm): + addends_2d = [] + addend_values_2d = [] + for addend, addend_value in zip(addends, addend_values, strict=True): + addend_2d = gm.graph.call_function( + _RESHAPE, + (addend, list(first_shape)), + ) + addend_value_2d = addend_value.reshape(first_shape) + addend_2d.meta = _copy_meta_with_value( + addend_value_2d, addend, *add_nodes + ) + _tag_for_regional_inductor(addend_2d) + addends_2d.append(addend_2d) + addend_values_2d.append(addend_value_2d) + gamma_2d = gm.graph.call_function( + _RESHAPE, + (gamma, [1, first_shape[-1]]), + ) + gamma_2d_value = gamma_value.reshape(1, first_shape[-1]) + gamma_2d.meta = _copy_meta_with_value(gamma_2d_value, gamma) + _tag_for_regional_inductor(gamma_2d) + + body = _build_f3_residual_rmsnorm_body( + first_mm, + tuple(addends_2d), + _CODA_RMSNORM_GROUP, + accumulated_value_is_lhs, + ) + body_name = _next_submodule_name(gm, "_coda_f3_residual_rmsnorm_body") + gm.add_module(body_name, body) + body_ref = gm.graph.get_attr(body_name) + _tag_for_regional_inductor(body_ref) + fused = gm.graph.call_function( + flex_gemm_hop, + ( + _MM, + body_ref, + tuple(first_mm.args) + tuple(addends_2d), + dict(first_mm.kwargs), + {"backend": "QUACK"}, + ), + ) + total_value = first_value + for addend_value, value_is_lhs in zip( + addend_values_2d, accumulated_value_is_lhs, strict=True + ): + total_value = ( + total_value + addend_value + if value_is_lhs + else addend_value + total_value + ) + partial_value = ( + total_value.float() + .reshape(first_shape[0], -1, _CODA_RMSNORM_GROUP) + .square() + .mean(-1) + ) + total_meta = _copy_meta_with_value( + total_value, first_mm, first_reshape, *addends, *add_nodes + ) + partial_meta = _copy_meta_with_value( + partial_value, + first_mm, + first_reshape, + *addends, + *add_nodes, + norm, + ) + fused.meta = _copy_meta( + first_mm, + first_reshape, + *addends, + *add_nodes, + norm, + norm_output, + ) + fused.meta["val"] = ( + total_meta["val"], + partial_meta["val"], + ) + fused.meta.pop("tensor_meta", None) + _tag_for_regional_inductor(fused) + + total_2d = gm.graph.call_function(operator.getitem, (fused, 0)) + total_2d.meta = total_meta + _tag_for_regional_inductor(total_2d) + partial = gm.graph.call_function(operator.getitem, (fused, 1)) + partial.meta = partial_meta + _tag_for_regional_inductor(partial) + total = gm.graph.call_function( + _RESHAPE, + (total_2d, list(output_shape)), + ) + total.meta = _copy_meta_with_value(total_value.reshape(output_shape), add) + _tag_for_regional_inductor(total) + mean_square = gm.graph.call_function(_MEAN_DIM, (partial, [-1], True)) + mean_square_value = partial_value.mean(-1, keepdim=True) + mean_square.meta = _copy_meta_with_value( + mean_square_value, first_mm, add, norm + ) + _tag_for_regional_inductor(mean_square) + stabilized = gm.graph.call_function(_ADD_SCALAR, (mean_square, eps)) + stabilized_value = mean_square_value + eps + stabilized.meta = _copy_meta_with_value( + stabilized_value, first_mm, add, norm + ) + _tag_for_regional_inductor(stabilized) + rstd = gm.graph.call_function(_RSQRT, (stabilized,)) + rstd_value = stabilized_value.rsqrt() + rstd.meta = _copy_meta_with_value(rstd_value, first_mm, add, norm) + _tag_for_regional_inductor(rstd) + + total_fp32 = gm.graph.call_function( + _TO_COPY, + (total_2d,), + {"dtype": torch.float32}, + ) + total_fp32.meta = _copy_meta_with_value( + total_value.float(), add, norm, norm_output + ) + _tag_for_regional_inductor(total_fp32) + normalized_fp32 = gm.graph.call_function(_MUL, (total_fp32, rstd)) + normalized_value = total_value.float() * rstd_value + normalized_fp32.meta = _copy_meta_with_value( + normalized_value, add, norm, norm_output + ) + _tag_for_regional_inductor(normalized_fp32) + normalized_weighted_fp32 = gm.graph.call_function( + _MUL, + (normalized_fp32, gamma_2d), + ) + normalized_weighted_value = normalized_value * gamma_2d_value + normalized_weighted_fp32.meta = _copy_meta_with_value( + normalized_weighted_value, add, norm, norm_output, gamma + ) + _tag_for_regional_inductor(normalized_weighted_fp32) + normalized_2d = gm.graph.call_function( + _TO_COPY, + (normalized_weighted_fp32,), + {"dtype": torch.bfloat16}, + ) + normalized_2d.meta = _copy_meta_with_value( + normalized_weighted_value.to(torch.bfloat16), norm_output + ) + _tag_for_regional_inductor(normalized_2d) + normalized = gm.graph.call_function( + _RESHAPE, + (normalized_2d, list(output_shape)), + ) + normalized.meta = _copy_meta_with_value( + normalized_weighted_value.to(torch.bfloat16).reshape(output_shape), + norm_output, + ) + _tag_for_regional_inductor(normalized) + + saved_rstd_output = None + if rstd_output is not None: + rstd_output_shape = _node_shape(rstd_output) + assert rstd_output_shape is not None + saved_rstd_output = gm.graph.call_function( + _RESHAPE, + (rstd, list(rstd_output_shape)), + ) + saved_rstd_output.meta = _copy_meta_with_value( + rstd_value.reshape(rstd_output_shape), rstd_output + ) + _tag_for_regional_inductor(saved_rstd_output) + + add.replace_all_uses_with(total) + norm_output.replace_all_uses_with(normalized) + if rstd_output is not None: + assert saved_rstd_output is not None + rstd_output.replace_all_uses_with(saved_rstd_output) + gm.graph.erase_node(rstd_output) + first_role = _layer_module_role(_module_fqn(first_mm)) + assert first_role is not None + gm.graph.erase_node(norm_output) + gm.graph.erase_node(norm) + for add_node in reversed(add_nodes): + gm.graph.erase_node(add_node) + gm.graph.erase_node(first_reshape) + gm.graph.erase_node(first_mm) + + if first_role[1] == "attention.wo": + num_attention_output_fused += 1 + else: + num_layer_output_fused += 1 + + gm.graph.lint() + gm.recompile() + logger.info( + f"F3 fused {num_attention_output_fused} attention-output and " + f"{num_layer_output_fused} layer-output residual RMSNorm boundaries" + ) + return gm + + +def _match_f2_q_rmsnorm( + norm: torch.fx.Node, +) -> tuple[ + torch.fx.Node, + torch.fx.Node, + torch.fx.Node, + torch.fx.Node, + torch.fx.Node, + torch.fx.Node, + torch.fx.Node, + torch.fx.Node | None, + float, +] | None: + if norm.target != _FUSED_RMS_NORM or len(norm.args) < 4: + return None + norm_input, normalized_shape, gamma, eps = norm.args[:4] + if ( + not isinstance(norm_input, torch.fx.Node) + or norm_input.target not in (_RESHAPE, _VIEW) + or not isinstance(gamma, torch.fx.Node) + or not isinstance(normalized_shape, (list, tuple)) + or len(normalized_shape) != 1 + or not isinstance(eps, float) + ): + return None + if not norm_input.args or not isinstance(norm_input.args[0], torch.fx.Node): + return None + first_mm = norm_input.args[0] + if first_mm.target != _MM or _sole_user(first_mm) is not norm_input: + return None + + norm_outputs: dict[int, torch.fx.Node] = {} + for user in norm.users: + if ( + user.target is not operator.getitem + or len(user.args) < 2 + or not isinstance(user.args[1], int) + or user.args[1] in norm_outputs + ): + return None + norm_outputs[user.args[1]] = user + if set(norm_outputs) not in ({0}, {0, 1}): + return None + norm_output = norm_outputs[0] + rstd_output = norm_outputs.get(1) + second_input = _sole_user(norm_output) + if second_input is None or second_input.target not in (_RESHAPE, _VIEW): + return None + second_mms = [ + user + for user in second_input.users + if user.target == _MM and user.args and user.args[0] is second_input + ] + if len(second_mms) != 1: + return None + second_mm = second_mms[0] + + is_recomputed = rstd_output is not None + if not is_recomputed and ( + _sole_user(norm_input) is not norm or _sole_user(second_input) is not second_mm + ): + return None + if rstd_output is not None: + rstd_shape = _node_shape(rstd_output) + norm_input_shape = _node_shape(norm_input) + if ( + _node_dtype(rstd_output) != torch.float32 + or norm_input_shape is None + or rstd_shape != (*norm_input_shape[:-1], 1) + ): + return None + + fqns = tuple(_module_fqn(node) for node in (first_mm, norm, second_mm)) + if any(fqn is None or "." not in fqn for fqn in fqns): + return None + first_fqn, norm_fqn, second_fqn = fqns + assert first_fqn is not None + assert norm_fqn is not None + assert second_fqn is not None + first_parent, first_role = first_fqn.rsplit(".", 1) + norm_parent, norm_role = norm_fqn.rsplit(".", 1) + second_parent, second_role = second_fqn.rsplit(".", 1) + if ( + first_parent != norm_parent + or first_parent != second_parent + or (first_role, norm_role, second_role) != ("wq_a", "q_norm", "wq_b") + ): + return None + + first_shape = _node_shape(first_mm) + norm_input_shape = _node_shape(norm_input) + second_input_shape = _node_shape(second_input) + gamma_shape = _node_shape(gamma) + if ( + first_shape is None + or len(first_shape) != 2 + or norm_input_shape is None + or first_shape[-1] != normalized_shape[0] + or first_shape[-1] % _CODA_RMSNORM_GROUP != 0 + or norm_input_shape[-1:] != first_shape[-1:] + or _node_shape(norm_output) != _node_shape(norm_input) + or second_input_shape != first_shape + or gamma_shape != (first_shape[-1],) + ): + return None + tensor_nodes = ( + first_mm, + norm_input, + gamma, + norm_output, + second_input, + second_mm, + ) + if any(_node_dtype(node) != torch.bfloat16 for node in tensor_nodes): + return None + if not all( + isinstance(node.meta.get("val"), torch.Tensor) + for node in (first_mm, gamma, second_mm) + ): + return None + return ( + first_mm, + norm_input, + norm, + norm_output, + second_input, + second_mm, + gamma, + rstd_output, + eps, + ) + + +def fuse_f2_q_rmsnorm_pass( + gm: torch.fx.GraphModule, + example_inputs: tuple | None = None, +) -> torch.fx.GraphModule: + """Reparameterize MLA Q RMSNorm across two GEMMs.""" + del example_inputs + num_original_fused = 0 + num_recomputed_fused = 0 + for norm in list(gm.graph.nodes): + match = _match_f2_q_rmsnorm(norm) + if match is None: + continue + ( + first_mm, + norm_input, + norm, + norm_output, + second_input, + second_mm, + gamma, + rstd_output, + eps, + ) = match + first_shape = _node_shape(first_mm) + assert first_shape is not None + first_value = first_mm.meta["val"] + gamma_value = gamma.meta["val"] + second_value = second_mm.meta["val"] + assert isinstance(first_value, torch.Tensor) + assert isinstance(gamma_value, torch.Tensor) + assert isinstance(second_value, torch.Tensor) + preserve_saved_values = rstd_output is not None + + with gm.graph.inserting_before(first_mm): + gamma_2d = gm.graph.call_function( + _RESHAPE, + (gamma, [1, first_shape[-1]]), + ) + gamma_2d.meta = _copy_meta_with_value( + gamma_value.reshape(1, first_shape[-1]), gamma + ) + _tag_for_regional_inductor(gamma_2d) + + first_body = _build_f2_q_rmsnorm_first_body( + first_mm, + gamma_2d, + _CODA_RMSNORM_GROUP, + preserve_saved_values, + ) + first_body_name = _next_submodule_name(gm, "_coda_f2_q_first_body") + gm.add_module(first_body_name, first_body) + first_body_ref = gm.graph.get_attr(first_body_name) + _tag_for_regional_inductor(first_body_ref) + first_fused = gm.graph.call_function( + flex_gemm_hop, + ( + _MM, + first_body_ref, + tuple(first_mm.args) + (gamma_2d,), + dict(first_mm.kwargs), + {"backend": "QUACK"}, + ), + ) + weighted_meta = _copy_meta_with_value_from( + first_mm, first_mm, norm_input, norm, norm_output, gamma + ) + partial_value = ( + first_value.float() + .reshape(first_shape[0], -1, _CODA_RMSNORM_GROUP) + .square() + .mean(-1) + ) + partial_meta = _copy_meta_with_value( + partial_value, first_mm, norm_input, norm + ) + first_fused.meta = _copy_meta(first_mm, norm_input, norm, norm_output) + raw_meta = _copy_meta_with_value_from(first_mm, first_mm, norm_input, norm) + first_fused_values = [weighted_meta["val"]] + if preserve_saved_values: + first_fused_values.append(raw_meta["val"]) + first_fused_values.append(partial_meta["val"]) + first_fused.meta["val"] = tuple(first_fused_values) + first_fused.meta.pop("tensor_meta", None) + _tag_for_regional_inductor(first_fused) + + weighted = gm.graph.call_function(operator.getitem, (first_fused, 0)) + weighted.meta = weighted_meta + _tag_for_regional_inductor(weighted) + raw = None + partial_index = 1 + if preserve_saved_values: + raw = gm.graph.call_function(operator.getitem, (first_fused, 1)) + raw.meta = raw_meta + _tag_for_regional_inductor(raw) + partial_index = 2 + partial = gm.graph.call_function( + operator.getitem, (first_fused, partial_index) + ) + partial.meta = partial_meta + _tag_for_regional_inductor(partial) + mean_square = gm.graph.call_function(_MEAN_DIM, (partial, [-1], True)) + mean_square_value = partial_value.mean(-1, keepdim=True) + mean_square.meta = _copy_meta_with_value(mean_square_value, first_mm, norm) + _tag_for_regional_inductor(mean_square) + stabilized = gm.graph.call_function(_ADD_SCALAR, (mean_square, eps)) + stabilized_value = mean_square_value + eps + stabilized.meta = _copy_meta_with_value(stabilized_value, first_mm, norm) + _tag_for_regional_inductor(stabilized) + rstd = gm.graph.call_function(_RSQRT, (stabilized,)) + rstd_value = stabilized_value.rsqrt() + rstd.meta = _copy_meta_with_value(rstd_value, first_mm, norm) + _tag_for_regional_inductor(rstd) + + saved_norm_input = None + saved_norm_output = None + saved_rstd_output = None + if preserve_saved_values: + assert raw is not None + assert rstd_output is not None + norm_input_shape = _node_shape(norm_input) + rstd_output_shape = _node_shape(rstd_output) + assert norm_input_shape is not None + assert rstd_output_shape is not None + + raw_fp32 = gm.graph.call_function( + _TO_COPY, (raw,), {"dtype": torch.float32} + ) + raw_fp32.meta = _copy_meta_with_dtype( + first_mm, torch.float32, first_mm, norm_input, norm + ) + _tag_for_regional_inductor(raw_fp32) + normalized_fp32 = gm.graph.call_function(_MUL, (raw_fp32, rstd)) + normalized_fp32_value = first_value.float() * rstd_value + normalized_fp32.meta = _copy_meta_with_value( + normalized_fp32_value, first_mm, norm_input, norm + ) + _tag_for_regional_inductor(normalized_fp32) + normalized_weighted_fp32 = gm.graph.call_function( + _MUL, (normalized_fp32, gamma_2d) + ) + normalized_weighted_value = normalized_fp32_value * gamma_value.reshape( + 1, first_shape[-1] + ) + normalized_weighted_fp32.meta = _copy_meta_with_value( + normalized_weighted_value, first_mm, norm_input, norm, gamma + ) + _tag_for_regional_inductor(normalized_weighted_fp32) + saved_norm_output = gm.graph.call_function( + _TO_COPY, + (normalized_weighted_fp32,), + {"dtype": torch.bfloat16}, + ) + saved_norm_output.meta = _copy_meta_with_value( + normalized_weighted_value.to(torch.bfloat16), + norm_output, + second_input, + ) + _tag_for_regional_inductor(saved_norm_output) + saved_norm_input = gm.graph.call_function( + _RESHAPE, (raw, list(norm_input_shape)) + ) + saved_norm_input.meta = _copy_meta_with_value( + first_value.reshape(norm_input_shape), norm_input + ) + _tag_for_regional_inductor(saved_norm_input) + saved_rstd_output = gm.graph.call_function( + _RESHAPE, (rstd, list(rstd_output_shape)) + ) + saved_rstd_output.meta = _copy_meta_with_value( + rstd_value.reshape(rstd_output_shape), rstd_output + ) + _tag_for_regional_inductor(saved_rstd_output) + + second_body = _build_f2_q_rmsnorm_second_body(second_mm, rstd) + second_body_name = _next_submodule_name(gm, "_coda_f2_q_second_body") + gm.add_module(second_body_name, second_body) + with gm.graph.inserting_before(second_mm): + second_body_ref = gm.graph.get_attr(second_body_name) + _tag_for_regional_inductor(second_body_ref) + second_fused = gm.graph.call_function( + flex_gemm_hop, + ( + _MM, + second_body_ref, + (weighted, second_mm.args[1], rstd), + dict(second_mm.kwargs), + {"backend": "QUACK"}, + ), + ) + second_meta = _copy_meta_with_value( + second_value, norm, norm_output, second_input, second_mm + ) + second_fused.meta = _copy_meta(norm, norm_output, second_input, second_mm) + second_fused.meta["val"] = (second_meta["val"],) + second_fused.meta.pop("tensor_meta", None) + _tag_for_regional_inductor(second_fused) + output = gm.graph.call_function(operator.getitem, (second_fused, 0)) + output.meta = second_meta + _tag_for_regional_inductor(output) + + second_mm.replace_all_uses_with(output) + gm.graph.erase_node(second_mm) + if preserve_saved_values: + assert saved_norm_input is not None + assert saved_norm_output is not None + assert saved_rstd_output is not None + assert rstd_output is not None + second_input.replace_all_uses_with(saved_norm_output) + norm_input.replace_all_uses_with(saved_norm_input) + rstd_output.replace_all_uses_with(saved_rstd_output) + gm.graph.erase_node(rstd_output) + num_recomputed_fused += 1 + else: + num_original_fused += 1 + gm.graph.erase_node(second_input) + gm.graph.erase_node(norm_output) + gm.graph.erase_node(norm) + gm.graph.erase_node(norm_input) + gm.graph.erase_node(first_mm) + + gm.graph.lint() + gm.recompile() + logger.info( + f"F2 fused {num_original_fused} original and " + f"{num_recomputed_fused} recomputed MLA Q RMSNorm chains" + ) + return gm + + +def _match_f2_kv_rmsnorm( + norm: torch.fx.Node, +) -> tuple[ + torch.fx.Node, + torch.fx.Node, + torch.fx.Node, + torch.fx.Node, + torch.fx.Node, + torch.fx.Node, + torch.fx.Node, + torch.fx.Node, + torch.fx.Node, + torch.fx.Node, + torch.fx.Node | None, + float, +] | None: + if norm.target != _FUSED_RMS_NORM or len(norm.args) < 4: + return None + norm_input, normalized_shape, gamma, eps = norm.args[:4] + if ( + not isinstance(norm_input, torch.fx.Node) + or norm_input.target is not operator.getitem + or len(norm_input.args) < 2 + or norm_input.args[1] != 0 + or not isinstance(gamma, torch.fx.Node) + or normalized_shape != [_CODA_KV_ACTIVE_WIDTH] + or not isinstance(eps, float) + ): + return None + split = norm_input.args[0] + if ( + not isinstance(split, torch.fx.Node) + or split.target != _SPLIT_WITH_SIZES + or len(split.args) < 3 + or not isinstance(split.args[1], (list, tuple)) + or list(split.args[1]) != [_CODA_KV_ACTIVE_WIDTH, _CODA_KV_TAIL_WIDTH] + or split.args[2] != -1 + ): + return None + first_reshape = split.args[0] + if ( + not isinstance(first_reshape, torch.fx.Node) + or first_reshape.target not in (_RESHAPE, _VIEW) + or not first_reshape.args + or not isinstance(first_reshape.args[0], torch.fx.Node) + or _sole_user(first_reshape) is not split + ): + return None + first_mm = first_reshape.args[0] + if first_mm.target != _MM or _sole_user(first_mm) is not first_reshape: + return None + + split_outputs = {} + for user in split.users: + if ( + user.target is not operator.getitem + or len(user.args) < 2 + or not isinstance(user.args[1], int) + or user.args[1] in split_outputs + ): + return None + split_outputs[user.args[1]] = user + if set(split_outputs) != {0, 1} or split_outputs[0] is not norm_input: + return None + tail = split_outputs[1] + + norm_outputs: dict[int, torch.fx.Node] = {} + for user in norm.users: + if ( + user.target is not operator.getitem + or len(user.args) < 2 + or not isinstance(user.args[1], int) + or user.args[1] in norm_outputs + ): + return None + norm_outputs[user.args[1]] = user + if set(norm_outputs) not in ({0}, {0, 1}): + return None + norm_output = norm_outputs[0] + rstd_output = norm_outputs.get(1) + is_recomputed = rstd_output is not None + second_input = _sole_user(norm_output) + if second_input is None or second_input.target not in (_RESHAPE, _VIEW): + return None + second_mms = [ + user + for user in second_input.users + if user.target == _MM and user.args and user.args[0] is second_input + ] + if len(second_mms) != 1: + return None + second_mm = second_mms[0] + + if not is_recomputed and ( + _sole_user(norm_input) is not norm or _sole_user(second_input) is not second_mm + ): + return None + if rstd_output is not None: + rstd_shape = _node_shape(rstd_output) + norm_input_shape = _node_shape(norm_input) + if ( + _node_dtype(rstd_output) != torch.float32 + or norm_input_shape is None + or rstd_shape != (*norm_input_shape[:-1], 1) + ): + return None + + fqns = tuple(_module_fqn(node) for node in (first_mm, norm, second_mm)) + if any(fqn is None or "." not in fqn for fqn in fqns): + return None + first_fqn, norm_fqn, second_fqn = fqns + assert first_fqn is not None + assert norm_fqn is not None + assert second_fqn is not None + first_parent, first_role = first_fqn.rsplit(".", 1) + norm_parent, norm_role = norm_fqn.rsplit(".", 1) + second_parent, second_role = second_fqn.rsplit(".", 1) + if ( + first_parent != norm_parent + or first_parent != second_parent + or (first_role, norm_role, second_role) != ("wkv_a", "kv_norm", "wkv_b") + ): + return None + + first_shape = _node_shape(first_mm) + first_reshape_shape = _node_shape(first_reshape) + norm_input_shape = _node_shape(norm_input) + second_input_shape = _node_shape(second_input) + if ( + first_shape is None + or len(first_shape) != 2 + or first_shape[-1] != _CODA_KV_ACTIVE_WIDTH + _CODA_KV_TAIL_WIDTH + or first_reshape_shape is None + or first_reshape_shape[-1] != first_shape[-1] + or norm_input_shape is None + or norm_input_shape[-1] != _CODA_KV_ACTIVE_WIDTH + or _node_shape(tail) != (*norm_input_shape[:-1], _CODA_KV_TAIL_WIDTH) + or _node_shape(norm_output) != norm_input_shape + or second_input_shape != (first_shape[0], _CODA_KV_ACTIVE_WIDTH) + or _node_shape(gamma) != (_CODA_KV_ACTIVE_WIDTH,) + ): + return None + tensor_nodes = ( + first_mm, + first_reshape, + norm_input, + tail, + gamma, + norm_output, + second_input, + second_mm, + ) + if any(_node_dtype(node) != torch.bfloat16 for node in tensor_nodes): + return None + if not all( + isinstance(node.meta.get("val"), torch.Tensor) + for node in (first_mm, gamma, second_mm) + ): + return None + return ( + first_mm, + first_reshape, + split, + norm_input, + tail, + norm, + norm_output, + second_input, + second_mm, + gamma, + rstd_output, + eps, + ) + + +def fuse_f2_kv_rmsnorm_pass( + gm: torch.fx.GraphModule, + example_inputs: tuple | None = None, +) -> torch.fx.GraphModule: + """Reparameterize segmented MLA KV RMSNorm across two GEMMs.""" + del example_inputs + num_original_fused = 0 + num_recomputed_fused = 0 + for norm in list(gm.graph.nodes): + match = _match_f2_kv_rmsnorm(norm) + if match is None: + continue + ( + first_mm, + first_reshape, + split, + norm_input, + tail, + norm, + norm_output, + second_input, + second_mm, + gamma, + rstd_output, + eps, + ) = match + del first_reshape, split, tail + first_shape = _node_shape(first_mm) + norm_input_shape = _node_shape(norm_input) + assert first_shape is not None + assert norm_input_shape is not None + first_value = first_mm.meta["val"] + gamma_value = gamma.meta["val"] + second_value = second_mm.meta["val"] + assert isinstance(first_value, torch.Tensor) + assert isinstance(gamma_value, torch.Tensor) + assert isinstance(second_value, torch.Tensor) + preserve_saved_values = rstd_output is not None + + with gm.graph.inserting_before(first_mm): + gamma_2d = gm.graph.call_function( + _RESHAPE, + (gamma, [1, _CODA_KV_ACTIVE_WIDTH]), + ) + gamma_2d_value = gamma_value.reshape(1, _CODA_KV_ACTIVE_WIDTH) + gamma_2d.meta = _copy_meta_with_value(gamma_2d_value, gamma) + _tag_for_regional_inductor(gamma_2d) + gamma_full = gm.graph.call_function( + _CONSTANT_PAD_ND, + (gamma_2d, [0, _CODA_KV_TAIL_WIDTH], 1.0), + ) + gamma_full_value = torch.nn.functional.pad( + gamma_2d_value, (0, _CODA_KV_TAIL_WIDTH), value=1.0 + ) + gamma_full.meta = _copy_meta_with_value(gamma_full_value, gamma, gamma_2d) + _tag_for_regional_inductor(gamma_full) + + first_body = _build_f2_q_rmsnorm_first_body( + first_mm, + gamma_full, + _CODA_KV_RMSNORM_GROUP, + preserve_raw=True, + ) + first_body_name = _next_submodule_name(gm, "_coda_f2_kv_first_body") + gm.add_module(first_body_name, first_body) + first_body_ref = gm.graph.get_attr(first_body_name) + _tag_for_regional_inductor(first_body_ref) + first_fused = gm.graph.call_function( + flex_gemm_hop, + ( + _MM, + first_body_ref, + tuple(first_mm.args) + (gamma_full,), + dict(first_mm.kwargs), + {"backend": "QUACK"}, + ), + ) + weighted_full_value = (first_value.float() * gamma_full_value).to( + torch.bfloat16 + ) + partial_value = ( + first_value.float() + .reshape(first_shape[0], -1, _CODA_KV_RMSNORM_GROUP) + .square() + .mean(-1) + ) + weighted_full_meta = _copy_meta_with_value( + weighted_full_value, first_mm, norm_input, norm, norm_output, gamma + ) + raw_meta = _copy_meta_with_value_from(first_mm, first_mm, norm_input, norm) + partial_meta = _copy_meta_with_value( + partial_value, first_mm, norm_input, norm + ) + first_fused.meta = _copy_meta(first_mm, norm_input, norm, norm_output) + first_fused.meta["val"] = ( + weighted_full_meta["val"], + raw_meta["val"], + partial_meta["val"], + ) + first_fused.meta.pop("tensor_meta", None) + _tag_for_regional_inductor(first_fused) + + weighted_full = gm.graph.call_function(operator.getitem, (first_fused, 0)) + weighted_full.meta = weighted_full_meta + _tag_for_regional_inductor(weighted_full) + raw = gm.graph.call_function(operator.getitem, (first_fused, 1)) + raw.meta = raw_meta + _tag_for_regional_inductor(raw) + partial = gm.graph.call_function(operator.getitem, (first_fused, 2)) + partial.meta = partial_meta + _tag_for_regional_inductor(partial) + + weighted = gm.graph.call_function( + _SLICE, + (weighted_full, 1, 0, _CODA_KV_ACTIVE_WIDTH), + ) + weighted_value = weighted_full_value[:, :_CODA_KV_ACTIVE_WIDTH] + weighted.meta = _copy_meta_with_value( + weighted_value, first_mm, norm_input, norm, norm_output, gamma + ) + _tag_for_regional_inductor(weighted) + partial_active = gm.graph.call_function( + _SLICE, + ( + partial, + 1, + 0, + _CODA_KV_ACTIVE_WIDTH // _CODA_KV_RMSNORM_GROUP, + ), + ) + partial_active_value = partial_value[ + :, : _CODA_KV_ACTIVE_WIDTH // _CODA_KV_RMSNORM_GROUP + ] + partial_active.meta = _copy_meta_with_value( + partial_active_value, first_mm, norm_input, norm + ) + _tag_for_regional_inductor(partial_active) + mean_square = gm.graph.call_function( + _MEAN_DIM, (partial_active, [-1], True) + ) + mean_square_value = partial_active_value.mean(-1, keepdim=True) + mean_square.meta = _copy_meta_with_value(mean_square_value, first_mm, norm) + _tag_for_regional_inductor(mean_square) + stabilized = gm.graph.call_function(_ADD_SCALAR, (mean_square, eps)) + stabilized_value = mean_square_value + eps + stabilized.meta = _copy_meta_with_value(stabilized_value, first_mm, norm) + _tag_for_regional_inductor(stabilized) + rstd = gm.graph.call_function(_RSQRT, (stabilized,)) + rstd_value = stabilized_value.rsqrt() + rstd.meta = _copy_meta_with_value(rstd_value, first_mm, norm) + _tag_for_regional_inductor(rstd) + + saved_norm_output = None + saved_rstd_output = None + if preserve_saved_values: + assert rstd_output is not None + raw_active = gm.graph.call_function( + _SLICE, + (raw, 1, 0, _CODA_KV_ACTIVE_WIDTH), + ) + raw_active_value = first_value[:, :_CODA_KV_ACTIVE_WIDTH] + raw_active.meta = _copy_meta_with_value( + raw_active_value, first_mm, norm_input, norm + ) + _tag_for_regional_inductor(raw_active) + raw_active_fp32 = gm.graph.call_function( + _TO_COPY, (raw_active,), {"dtype": torch.float32} + ) + raw_active_fp32.meta = _copy_meta_with_value( + raw_active_value.float(), first_mm, norm_input, norm + ) + _tag_for_regional_inductor(raw_active_fp32) + normalized_fp32 = gm.graph.call_function(_MUL, (raw_active_fp32, rstd)) + normalized_fp32_value = raw_active_value.float() * rstd_value + normalized_fp32.meta = _copy_meta_with_value( + normalized_fp32_value, first_mm, norm_input, norm + ) + _tag_for_regional_inductor(normalized_fp32) + normalized_weighted_fp32 = gm.graph.call_function( + _MUL, (normalized_fp32, gamma_2d) + ) + normalized_weighted_value = normalized_fp32_value * gamma_2d_value + normalized_weighted_fp32.meta = _copy_meta_with_value( + normalized_weighted_value, norm_input, norm, norm_output, gamma + ) + _tag_for_regional_inductor(normalized_weighted_fp32) + saved_norm_output = gm.graph.call_function( + _TO_COPY, + (normalized_weighted_fp32,), + {"dtype": torch.bfloat16}, + ) + saved_norm_output.meta = _copy_meta_with_value( + normalized_weighted_value.to(torch.bfloat16), second_input + ) + _tag_for_regional_inductor(saved_norm_output) + rstd_output_shape = _node_shape(rstd_output) + assert rstd_output_shape is not None + saved_rstd_output = gm.graph.call_function( + _RESHAPE, (rstd, list(rstd_output_shape)) + ) + saved_rstd_output.meta = _copy_meta_with_value( + rstd_value.reshape(rstd_output_shape), rstd_output + ) + _tag_for_regional_inductor(saved_rstd_output) + + second_body = _build_f2_q_rmsnorm_second_body(second_mm, rstd) + second_body_name = _next_submodule_name(gm, "_coda_f2_kv_second_body") + gm.add_module(second_body_name, second_body) + with gm.graph.inserting_before(second_mm): + second_body_ref = gm.graph.get_attr(second_body_name) + _tag_for_regional_inductor(second_body_ref) + second_fused = gm.graph.call_function( + flex_gemm_hop, + ( + _MM, + second_body_ref, + (weighted, second_mm.args[1], rstd), + dict(second_mm.kwargs), + {"backend": "QUACK"}, + ), + ) + second_meta = _copy_meta_with_value( + second_value, norm, norm_output, second_input, second_mm + ) + second_fused.meta = _copy_meta(norm, norm_output, second_input, second_mm) + second_fused.meta["val"] = (second_meta["val"],) + second_fused.meta.pop("tensor_meta", None) + _tag_for_regional_inductor(second_fused) + output = gm.graph.call_function(operator.getitem, (second_fused, 0)) + output.meta = second_meta + _tag_for_regional_inductor(output) + + second_mm.replace_all_uses_with(output) + gm.graph.erase_node(second_mm) + if preserve_saved_values: + assert saved_norm_output is not None + assert saved_rstd_output is not None + assert rstd_output is not None + second_input.replace_all_uses_with(saved_norm_output) + rstd_output.replace_all_uses_with(saved_rstd_output) + gm.graph.erase_node(rstd_output) + num_recomputed_fused += 1 + else: + num_original_fused += 1 + gm.graph.erase_node(second_input) + gm.graph.erase_node(norm_output) + gm.graph.erase_node(norm) + if not preserve_saved_values: + gm.graph.erase_node(norm_input) + first_mm.replace_all_uses_with(raw) + gm.graph.erase_node(first_mm) + + gm.graph.lint() + gm.recompile() + logger.info( + f"F2-KV fused {num_original_fused} original and " + f"{num_recomputed_fused} recomputed segmented RMSNorm chains" + ) + return gm + + +CODA_PATTERN_PASSES: dict[str, Callable] = { + "b1_lm_head_input_grad_cast": fuse_b1_lm_head_input_grad_cast_pass, + "b2_dense_swiglu_backward": fuse_b2_dense_swiglu_backward_pass, + "b4_router_input_grad_add": fuse_b4_router_input_grad_add_pass, + "b5_mla_rmsnorm_backward": fuse_b5_mla_rmsnorm_backward_pass, + "b6_bf16_weight_grad_cast": fuse_b6_bf16_weight_grad_cast_pass, + "b7_attention_grad_merge": fuse_b7_attention_grad_merge_pass, + "f3_residual_rmsnorm": fuse_f3_residual_rmsnorm_pass, + "f2_q_rmsnorm": fuse_f2_q_rmsnorm_pass, + "f2_kv_rmsnorm": fuse_f2_kv_rmsnorm_pass, + "f4_dense_swiglu": fuse_f4_dense_swiglu_pass, + "f6_router_sigmoid_bias": fuse_f6_router_sigmoid_bias_pass, +} + + +def get_coda_pattern_passes(patterns: Iterable[str]) -> list[Callable]: + """Resolve configured CODA pattern names to independently logged passes.""" + pattern_list = list(patterns) + duplicates = sorted( + pattern for pattern in set(pattern_list) if pattern_list.count(pattern) > 1 + ) + if duplicates: + raise ValueError(f"Duplicate --compile.coda_patterns entries: {duplicates}") + + unknown = sorted(set(pattern_list) - CODA_PATTERN_PASSES.keys()) + if unknown: + supported = sorted(CODA_PATTERN_PASSES) + raise ValueError( + f"Unknown --compile.coda_patterns entries: {unknown}; " + f"supported patterns: {supported}" + ) + return [CODA_PATTERN_PASSES[pattern] for pattern in pattern_list] diff --git a/torchtitan/experiments/graph_trainer/configs.py b/torchtitan/experiments/graph_trainer/configs.py index 4d89e8b808..4d04d15d88 100644 --- a/torchtitan/experiments/graph_trainer/configs.py +++ b/torchtitan/experiments/graph_trainer/configs.py @@ -125,6 +125,13 @@ class GraphTrainerCompileConfig(CompileConfig): """Enable passes that improve performance but may change numerics compared to the uncompiled path (e.g. RMSNorm Inductor fusion).""" + coda_patterns: list[str] = field(default_factory=list) + """CODA fusion patterns to apply after distributed scheduling. + + Each entry enables one independently logged pass for graph-dump and + performance ablation. Supported entries are defined in ``coda_passes``. + """ + cpu_offload_prefetch_n_layers: int = 1 """Prefetch reloads this many layers ahead in the backward graph to overlap H2D transfers with compute.""" diff --git a/torchtitan/experiments/graph_trainer/passes.py b/torchtitan/experiments/graph_trainer/passes.py index 7b0e9d5567..bc8748bb49 100644 --- a/torchtitan/experiments/graph_trainer/passes.py +++ b/torchtitan/experiments/graph_trainer/passes.py @@ -31,6 +31,8 @@ import torch +from torchtitan.experiments.graph_trainer.coda_passes import get_coda_pattern_passes + from torchtitan.experiments.graph_trainer.configs import ( GraphTrainerCompileConfig, MOE_BLOCK_FQN, @@ -337,6 +339,8 @@ def compile_time_passes( if config.parallelism.enable_async_tensor_parallel: passes.append(async_tensor_parallel_pass) + passes.extend(get_coda_pattern_passes(config.compile.coda_patterns)) + if not include_inductor: return passes diff --git a/torchtitan/experiments/graph_trainer/tests/test_coda_passes.py b/torchtitan/experiments/graph_trainer/tests/test_coda_passes.py new file mode 100644 index 0000000000..91cb62fe1c --- /dev/null +++ b/torchtitan/experiments/graph_trainer/tests/test_coda_passes.py @@ -0,0 +1,1151 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +import operator + +import torch +from torch.fx.experimental.proxy_tensor import make_fx +from torch.fx.passes.fake_tensor_prop import FakeTensorProp +from torch.testing._internal.common_utils import TestCase + +from torchtitan.experiments.graph_trainer.coda_passes import ( + fuse_b1_lm_head_input_grad_cast_pass, + fuse_b2_dense_swiglu_backward_pass, + fuse_b4_router_input_grad_add_pass, + fuse_b5_mla_rmsnorm_backward_pass, + fuse_b6_bf16_weight_grad_cast_pass, + fuse_b7_attention_grad_merge_pass, + fuse_f2_kv_rmsnorm_pass, + fuse_f2_q_rmsnorm_pass, + fuse_f3_residual_rmsnorm_pass, + fuse_f4_dense_swiglu_pass, + fuse_f6_router_sigmoid_bias_pass, + get_coda_pattern_passes, +) + + +class TestB1LMHeadInputGradCastPass(TestCase): + def _trace(self, *, module_fqn="lm_head", backward=True): + m, k, n = 8, 16, 32 + + def fn(grad, weight, destination): + input_grad = torch.ops.aten.mm.default(grad, weight) + input_grad = torch.ops.aten.reshape.default(input_grad, [2, 4, n]) + input_grad = torch.ops.aten.alias.default(input_grad) + input_grad = torch.ops.aten._to_copy.default( + input_grad, dtype=torch.float32 + ) + target = torch.ops.aten.slice.Tensor(destination, 1, 0, 4) + torch.ops.aten.copy_.default(target, input_grad) + return destination + + inputs = ( + torch.randn(m, k, dtype=torch.bfloat16) * 0.02, + torch.randn(k, n, dtype=torch.bfloat16) * 0.02, + torch.zeros(2, 8, n, dtype=torch.float32), + ) + gm = make_fx(fn)(*inputs) + mm = next( + node for node in gm.graph.nodes if node.target == torch.ops.aten.mm.default + ) + mm.meta["custom"] = {"module_fqn": module_fqn} + if backward: + mm.meta["autograd_backward"] = True + return gm, inputs + + def _flex_gemm_nodes(self, gm): + return [ + node + for node in gm.graph.nodes + if node.target == torch.ops.higher_order.flex_gemm + ] + + def test_fuses_lm_head_input_gradient_cast(self): + gm, inputs = self._trace() + expected = gm(*(value.clone() for value in inputs)) + + fuse_b1_lm_head_input_grad_cast_pass(gm) + + actual = gm(*(value.clone() for value in inputs)) + self.assertEqual(actual, expected, exact_dtype=True) + fused = self._flex_gemm_nodes(gm) + self.assertEqual(len(fused), 1) + fused_output = next(iter(fused[0].users)) + self.assertEqual(fused_output.meta["val"].shape, (8, 32)) + self.assertEqual(fused_output.meta["val"].dtype, torch.float32) + body = getattr(gm, fused[0].args[1].target) + body_targets = [node.target for node in body.graph.nodes] + self.assertEqual(body_targets.count(torch.ops.aten._to_copy.default), 3) + root_targets = [node.target for node in gm.graph.nodes] + self.assertEqual(root_targets.count(torch.ops.aten.copy_.default), 1) + + def test_does_not_fuse_unrelated_module(self): + gm, _ = self._trace(module_fqn="layers.3.attention.wq_a") + + fuse_b1_lm_head_input_grad_cast_pass(gm) + + self.assertEqual(self._flex_gemm_nodes(gm), []) + + def test_does_not_fuse_forward_graph(self): + gm, _ = self._trace(backward=False) + + fuse_b1_lm_head_input_grad_cast_pass(gm) + + self.assertEqual(self._flex_gemm_nodes(gm), []) + + +class TestB6WeightGradCastPass(TestCase): + def _trace(self, fn, *inputs): + return make_fx(fn)(*inputs) + + def _flex_gemm_nodes(self, gm): + return [ + node + for node in gm.graph.nodes + if node.target == torch.ops.higher_order.flex_gemm + ] + + def _body(self, gm, fused): + body_ref = fused.args[1] + self.assertEqual(body_ref.op, "get_attr") + return getattr(gm, body_ref.target) + + def test_fuses_bf16_mm_to_fp32(self): + x = torch.randn(8, 16, dtype=torch.bfloat16) + weight = torch.randn(16, 12, dtype=torch.bfloat16) + gm = self._trace(lambda a, b: torch.mm(a, b).float(), x, weight) + expected = gm(x, weight) + + fuse_b6_bf16_weight_grad_cast_pass(gm) + + self.assertEqual(gm(x, weight), expected, exact_dtype=True) + fused = self._flex_gemm_nodes(gm) + self.assertEqual(len(fused), 1) + body_targets = [node.target for node in self._body(gm, fused[0]).graph.nodes] + self.assertEqual(body_targets.count(torch.ops.aten.mm.default), 1) + self.assertEqual(body_targets.count(torch.ops.aten._to_copy.default), 3) + + def test_does_not_fuse_unsupported_fp32_router_round_trip(self): + x = torch.randn(8, 16) + weight = torch.randn(16, 12) + gm = self._trace(lambda a, b: torch.mm(a, b).bfloat16().float(), x, weight) + fuse_b6_bf16_weight_grad_cast_pass(gm) + self.assertEqual(self._flex_gemm_nodes(gm), []) + + def test_does_not_fuse_multi_use_mm(self): + x = torch.randn(8, 16, dtype=torch.bfloat16) + weight = torch.randn(16, 12, dtype=torch.bfloat16) + + def fn(a, b): + mm = torch.mm(a, b) + return mm.float(), mm + 1 + + gm = self._trace(fn, x, weight) + fuse_b6_bf16_weight_grad_cast_pass(gm) + self.assertEqual(self._flex_gemm_nodes(gm), []) + + def test_preserves_metadata_and_tags_regional_inductor(self): + x = torch.randn(8, 16, dtype=torch.bfloat16) + weight = torch.randn(16, 12, dtype=torch.bfloat16) + gm = self._trace(lambda a, b: torch.mm(a, b).float(), x, weight) + mm = next( + node for node in gm.graph.nodes if node.target == torch.ops.aten.mm.default + ) + cast = next( + node + for node in gm.graph.nodes + if node.target == torch.ops.aten._to_copy.default + ) + mm.meta["custom"] = {"module_fqn": "layers.0.moe", "EP": "compute"} + cast.meta["custom"] = {"autograd_backward": True} + + fuse_b6_bf16_weight_grad_cast_pass(gm) + + fused = self._flex_gemm_nodes(gm)[0] + self.assertEqual(fused.meta["custom"]["module_fqn"], "layers.0.moe") + self.assertEqual(fused.meta["custom"]["EP"], "compute") + self.assertTrue(fused.meta["custom"]["autograd_backward"]) + self.assertIn("compile_with_inductor", fused.meta["custom"]) + output = next(iter(fused.users)) + self.assertEqual(output.target, operator.getitem) + self.assertIn("compile_with_inductor", output.meta["custom"]) + body_ref = fused.args[1] + self.assertIn("compile_with_inductor", body_ref.meta["custom"]) + for node in self._body(gm, fused).graph.nodes: + self.assertIn("compile_with_inductor", node.meta["custom"]) + + +class TestF6RouterSigmoidBiasPass(TestCase): + def _flex_gemm_nodes(self, gm): + return [ + node + for node in gm.graph.nodes + if node.target == torch.ops.higher_order.flex_gemm + ] + + def test_fuses_sigmoid_and_bias_while_preserving_raw_scores(self): + x = torch.randn(8, 16) + weight = torch.randn(16, 5) + bias = torch.randn(5) + + def fn(a, b, expert_bias): + scores = torch.sigmoid(torch.mm(a, b).reshape(2, 4, 5)) + return scores, scores + expert_bias, scores * 2 + + gm = make_fx(fn)(x, weight, bias) + expected = gm(x, weight, bias) + fuse_f6_router_sigmoid_bias_pass(gm) + + self.assertEqual(gm(x, weight, bias), expected, exact_dtype=True) + fused = self._flex_gemm_nodes(gm) + self.assertEqual(len(fused), 1) + self.assertEqual(len(fused[0].args[2]), 3) + body_ref = fused[0].args[1] + body = getattr(gm, body_ref.target) + targets = [node.target for node in body.graph.nodes] + self.assertEqual(targets.count(torch.ops.aten.mm.default), 1) + self.assertEqual(targets.count(torch.ops.aten.sigmoid.default), 1) + self.assertEqual(targets.count(torch.ops.aten.add.Tensor), 1) + + def test_fuses_recomputed_sigmoid_without_bias(self): + x = torch.randn(8, 16) + weight = torch.randn(16, 5) + + def fn(a, b): + scores = torch.sigmoid(torch.mm(a, b).reshape(2, 4, 5)) + return scores, scores * 2 + + gm = make_fx(fn)(x, weight) + expected = gm(x, weight) + fuse_f6_router_sigmoid_bias_pass(gm) + + self.assertEqual(gm(x, weight), expected, exact_dtype=True) + fused = self._flex_gemm_nodes(gm) + self.assertEqual(len(fused), 1) + self.assertEqual(len(fused[0].args[2]), 2) + + def test_does_not_fuse_multi_use_mm(self): + x = torch.randn(8, 16) + weight = torch.randn(16, 5) + + def fn(a, b): + mm = torch.mm(a, b) + return torch.sigmoid(mm.reshape(2, 4, 5)), mm + 1 + + gm = make_fx(fn)(x, weight) + fuse_f6_router_sigmoid_bias_pass(gm) + self.assertEqual(self._flex_gemm_nodes(gm), []) + + +class TestF4DenseSwiGLUPass(TestCase): + def _flex_gemm_nodes(self, gm): + return [ + node + for node in gm.graph.nodes + if node.target == torch.ops.higher_order.flex_gemm + ] + + def test_fuses_two_gemms_and_preserves_saved_activations(self): + x = torch.randn(8, 16, dtype=torch.bfloat16) + w1 = torch.randn(16, 12, dtype=torch.bfloat16) + w3 = torch.randn(16, 12, dtype=torch.bfloat16) + + def fn(a, first_weight, gate_weight): + first = torch.mm(a, first_weight).reshape(2, 4, 12) + activated = torch.nn.functional.silu(first) + gate = torch.mm(a, gate_weight).reshape(2, 4, 12) + product = activated * gate + return first, activated, gate, product, product + 1 + + gm = make_fx(fn)(x, w1, w3) + expected = gm(x, w1, w3) + fuse_f4_dense_swiglu_pass(gm) + + self.assertEqual(gm(x, w1, w3), expected, exact_dtype=True) + fused = self._flex_gemm_nodes(gm) + self.assertEqual(len(fused), 2) + self.assertEqual(len(fused[0].args[2]), 2) + self.assertEqual(len(fused[1].args[2]), 3) + self.assertEqual(len(fused[0].meta["val"]), 2) + first_body = getattr(gm, fused[0].args[1].target) + second_body = getattr(gm, fused[1].args[1].target) + first_targets = [node.target for node in first_body.graph.nodes] + second_targets = [node.target for node in second_body.graph.nodes] + self.assertEqual(first_targets.count(torch.ops.aten.silu.default), 1) + self.assertEqual(first_targets.count(torch.ops.aten._to_copy.default), 2) + self.assertEqual(second_targets.count(torch.ops.aten.mul.Tensor), 1) + self.assertEqual(second_targets.count(torch.ops.aten._to_copy.default), 2) + + def test_does_not_fuse_fp32_swiglu(self): + x = torch.randn(8, 16) + w1 = torch.randn(16, 12) + w3 = torch.randn(16, 12) + + def fn(a, first_weight, gate_weight): + activated = torch.nn.functional.silu( + torch.mm(a, first_weight).reshape(2, 4, 12) + ) + gate = torch.mm(a, gate_weight).reshape(2, 4, 12) + return activated * gate + + gm = make_fx(fn)(x, w1, w3) + fuse_f4_dense_swiglu_pass(gm) + self.assertEqual(self._flex_gemm_nodes(gm), []) + + def test_does_not_fuse_multi_use_first_gemm(self): + x = torch.randn(8, 16, dtype=torch.bfloat16) + w1 = torch.randn(16, 12, dtype=torch.bfloat16) + w3 = torch.randn(16, 12, dtype=torch.bfloat16) + + def fn(a, first_weight, gate_weight): + first_mm = torch.mm(a, first_weight) + activated = torch.nn.functional.silu(first_mm.reshape(2, 4, 12)) + gate = torch.mm(a, gate_weight).reshape(2, 4, 12) + return activated * gate, first_mm + 1 + + gm = make_fx(fn)(x, w1, w3) + fuse_f4_dense_swiglu_pass(gm) + self.assertEqual(self._flex_gemm_nodes(gm), []) + + +class TestB2DenseSwiGLUBackwardPass(TestCase): + def _flex_gemm_nodes(self, gm): + return [ + node + for node in gm.graph.nodes + if node.target == torch.ops.higher_order.flex_gemm + ] + + def test_fuses_branch_derivatives_and_input_gradient_add(self): + grad_out = torch.randn(8, 16, dtype=torch.bfloat16) + w2 = torch.randn(16, 12, dtype=torch.bfloat16) + saved_silu = torch.randn(2, 4, 12, dtype=torch.bfloat16) + saved_gate = torch.randn(2, 4, 12, dtype=torch.bfloat16) + saved_preactivation = torch.randn(2, 4, 12, dtype=torch.bfloat16) + w3 = torch.randn(12, 7, dtype=torch.bfloat16) + w1 = torch.randn(12, 7, dtype=torch.bfloat16) + + def fn( + grad, + output_weight, + silu, + gate, + preactivation, + gate_weight, + first_weight, + ): + branch_grad = torch.mm(grad, output_weight).reshape(2, 4, 12) + gate_grad = branch_grad * silu + silu_grad = torch.ops.aten.silu_backward.default( + branch_grad * gate, preactivation + ) + w3_input_grad = torch.mm(gate_grad.reshape(8, 12), gate_weight) + w1_input_grad = torch.mm(silu_grad.reshape(8, 12), first_weight) + input_grad = w3_input_grad.reshape(2, 4, 7) + w1_input_grad.reshape(2, 4, 7) + return gate_grad, silu_grad, input_grad + + inputs = ( + grad_out, + w2, + saved_silu, + saved_gate, + saved_preactivation, + w3, + w1, + ) + gm = make_fx(fn)(*inputs) + expected = gm(*inputs) + mm_nodes = [ + node for node in gm.graph.nodes if node.target == torch.ops.aten.mm.default + ] + self.assertEqual(len(mm_nodes), 3) + mm_nodes[1].meta["custom"] = { + "module_fqn": "layers.0.feed_forward.w3", + "autograd_backward": True, + } + mm_nodes[2].meta["custom"] = { + "module_fqn": "layers.0.feed_forward.w1", + "autograd_backward": True, + } + + fuse_b2_dense_swiglu_backward_pass(gm) + + self.assertEqual(gm(*inputs), expected, exact_dtype=True) + fused = self._flex_gemm_nodes(gm) + self.assertEqual(len(fused), 2) + self.assertEqual(len(fused[0].args[2]), 5) + self.assertEqual(len(fused[1].args[2]), 3) + + branch_body = getattr(gm, fused[0].args[1].target) + branch_targets = [node.target for node in branch_body.graph.nodes] + self.assertEqual(branch_targets.count(torch.ops.aten.mul.Tensor), 2) + self.assertEqual(branch_targets.count(torch.ops.aten.silu_backward.default), 1) + self.assertEqual(branch_targets.count(torch.ops.aten._to_copy.default), 2) + + input_add_body = getattr(gm, fused[1].args[1].target) + input_add_targets = [node.target for node in input_add_body.graph.nodes] + self.assertEqual(input_add_targets.count(torch.ops.aten.add.Tensor), 1) + self.assertEqual(input_add_targets.count(torch.ops.aten._to_copy.default), 2) + + def test_does_not_fuse_unrelated_input_gradient_add(self): + x = torch.randn(8, 12, dtype=torch.bfloat16) + first_weight = torch.randn(12, 7, dtype=torch.bfloat16) + second_weight = torch.randn(12, 7, dtype=torch.bfloat16) + + def fn(a, first, second): + lhs = torch.mm(a, first).reshape(2, 4, 7) + rhs = torch.mm(a, second).reshape(2, 4, 7) + return lhs + rhs + + gm = make_fx(fn)(x, first_weight, second_weight) + for index, node in enumerate( + node for node in gm.graph.nodes if node.target == torch.ops.aten.mm.default + ): + node.meta["custom"] = {"module_fqn": f"layers.0.linear{index}"} + + fuse_b2_dense_swiglu_backward_pass(gm) + + self.assertEqual(self._flex_gemm_nodes(gm), []) + + +class TestF3ResidualRMSNormPass(TestCase): + def _trace_chain(self, *, cross_layer=False, return_saved=False): + m, k, n, p, q = 8, 16, 512, 64, 96 + + def fn(a, first_weight, residual, gamma, second_weight, third_weight): + first = torch.ops.aten.mm.default(a, first_weight) + first = torch.ops.aten.reshape.default(first, [2, 4, n]) + total = torch.ops.aten.add.Tensor(residual, first) + gamma_for_norm = torch.ops.aten.clone.default(gamma) + normalized, rstd = torch.ops.aten._fused_rms_norm.default( + total, [n], gamma_for_norm, 1e-5 + ) + second_input = torch.ops.aten.reshape.default(normalized, [m, n]) + second = torch.ops.aten.mm.default(second_input, second_weight) + third_input = torch.ops.aten.reshape.default(normalized, [m, n]) + third = torch.ops.aten.mm.default(third_input, third_weight) + if return_saved: + return second, third, total, rstd, second_input, third_input + return second, third, total + + inputs = ( + torch.randn(m, k, dtype=torch.bfloat16) * 0.02, + torch.randn(k, n, dtype=torch.bfloat16) * 0.02, + torch.randn(2, 4, n, dtype=torch.bfloat16) * 0.02, + torch.ones(n, dtype=torch.bfloat16), + torch.randn(n, p, dtype=torch.bfloat16) * 0.02, + torch.randn(n, q, dtype=torch.bfloat16) * 0.02, + ) + gm = torch.fx.symbolic_trace(fn) + gm.graph.eliminate_dead_code() + gm.recompile() + FakeTensorProp(gm).propagate(*inputs) + mm_nodes = [ + node for node in gm.graph.nodes if node.target == torch.ops.aten.mm.default + ] + norm = next( + node + for node in gm.graph.nodes + if node.target == torch.ops.aten._fused_rms_norm.default + ) + if cross_layer: + mm_nodes[0].meta["custom"] = {"module_fqn": "layers.2.feed_forward.w2"} + norm.meta["custom"] = {"module_fqn": "layers.3.attention_norm"} + mm_nodes[1].meta["custom"] = {"module_fqn": "layers.3.attention.wq_a"} + mm_nodes[2].meta["custom"] = {"module_fqn": "layers.3.attention.wkv_a"} + else: + mm_nodes[0].meta["custom"] = {"module_fqn": "layers.2.attention.wo"} + norm.meta["custom"] = {"module_fqn": "layers.2.ffn_norm"} + mm_nodes[1].meta["custom"] = {"module_fqn": "layers.2.feed_forward.w1"} + mm_nodes[2].meta["custom"] = {"module_fqn": "layers.2.feed_forward.w3"} + return gm, inputs + + def _flex_gemm_nodes(self, gm): + return [ + node + for node in gm.graph.nodes + if node.target == torch.ops.higher_order.flex_gemm + ] + + def _plain_mm_nodes(self, gm): + return [ + node for node in gm.graph.nodes if node.target == torch.ops.aten.mm.default + ] + + def _assert_outputs_close(self, actual, expected): + self.assertEqual(len(actual), len(expected)) + for actual_tensor, expected_tensor in zip(actual, expected, strict=True): + torch.testing.assert_close( + actual_tensor, + expected_tensor, + atol=1e-3, + rtol=2e-2, + ) + + def _trace_swiglu_chain(self, *, return_saved=False): + m, k, n, p = 8, 16, 512, 64 + + def fn(a, first_weight, residual, gamma, w1_weight, w3_weight): + first = torch.ops.aten.mm.default(a, first_weight) + first = torch.ops.aten.reshape.default(first, [2, 4, n]) + total = torch.ops.aten.add.Tensor(residual, first) + normalized, rstd = torch.ops.aten._fused_rms_norm.default( + total, [n], gamma, 1e-5 + ) + w1_input = torch.ops.aten.reshape.default(normalized, [m, n]) + w1 = torch.ops.aten.mm.default(w1_input, w1_weight) + w1 = torch.ops.aten.reshape.default(w1, [2, 4, p]) + activated = torch.ops.aten.silu.default(w1) + w3_input = torch.ops.aten.reshape.default(normalized, [m, n]) + w3 = torch.ops.aten.mm.default(w3_input, w3_weight) + w3 = torch.ops.aten.reshape.default(w3, [2, 4, p]) + product = torch.ops.aten.mul.Tensor(activated, w3) + return (product, total, rstd, w1, w3) if return_saved else (product, total) + + inputs = ( + torch.randn(m, k, dtype=torch.bfloat16) * 0.02, + torch.randn(k, n, dtype=torch.bfloat16) * 0.02, + torch.randn(2, 4, n, dtype=torch.bfloat16) * 0.02, + torch.ones(n, dtype=torch.bfloat16), + torch.randn(n, p, dtype=torch.bfloat16) * 0.02, + torch.randn(n, p, dtype=torch.bfloat16) * 0.02, + ) + gm = torch.fx.symbolic_trace(fn) + gm.graph.eliminate_dead_code() + gm.recompile() + FakeTensorProp(gm).propagate(*inputs) + mm_nodes = [ + node for node in gm.graph.nodes if node.target == torch.ops.aten.mm.default + ] + norm = next( + node + for node in gm.graph.nodes + if node.target == torch.ops.aten._fused_rms_norm.default + ) + mm_nodes[0].meta["custom"] = {"module_fqn": "layers.2.attention.wo"} + norm.meta["custom"] = {"module_fqn": "layers.2.ffn_norm"} + mm_nodes[1].meta["custom"] = {"module_fqn": "layers.2.feed_forward.w1"} + mm_nodes[2].meta["custom"] = {"module_fqn": "layers.2.feed_forward.w3"} + return gm, inputs + + def test_fuses_attention_output_residual_into_ffn_norm(self): + gm, inputs = self._trace_chain() + expected = gm(*inputs) + + fuse_f3_residual_rmsnorm_pass(gm) + + actual = gm(*inputs) + self._assert_outputs_close(actual, expected) + self.assertEqual(actual[2], expected[2]) + self.assertEqual(len(self._flex_gemm_nodes(gm)), 1) + self.assertEqual(len(self._plain_mm_nodes(gm)), 2) + self.assertFalse( + any( + node.target == torch.ops.aten._fused_rms_norm.default + for node in gm.graph.nodes + ) + ) + + def test_composes_dense_swiglu_epilogues(self): + gm, inputs = self._trace_swiglu_chain() + expected = gm(*inputs) + + fuse_f3_residual_rmsnorm_pass(gm) + fuse_f4_dense_swiglu_pass(gm) + + actual = gm(*inputs) + self._assert_outputs_close(actual, expected) + self.assertEqual(actual[1], expected[1]) + self.assertEqual(len(self._flex_gemm_nodes(gm)), 3) + body_targets = [ + [node.target for node in getattr(gm, fused.args[1].target).graph.nodes] + for fused in self._flex_gemm_nodes(gm) + ] + self.assertTrue( + any( + targets.count(torch.ops.aten.silu.default) == 1 + for targets in body_targets + ) + ) + self.assertTrue( + any( + targets.count(torch.ops.aten.mul.Tensor) == 1 + for targets in body_targets + ) + ) + + def test_composes_recomputed_dense_swiglu_and_preserves_saved_values(self): + gm, inputs = self._trace_swiglu_chain(return_saved=True) + expected = gm(*inputs) + + fuse_f3_residual_rmsnorm_pass(gm) + + actual = gm(*inputs) + self._assert_outputs_close(actual, expected) + self.assertEqual(actual[1], expected[1]) + self.assertEqual(len(self._flex_gemm_nodes(gm)), 1) + self.assertEqual(len(self._plain_mm_nodes(gm)), 2) + + def test_fuses_feed_forward_output_into_next_attention_norm(self): + gm, inputs = self._trace_chain(cross_layer=True, return_saved=True) + expected = gm(*inputs) + + fuse_f3_residual_rmsnorm_pass(gm) + + actual = gm(*inputs) + self._assert_outputs_close(actual, expected) + self.assertEqual(actual[2], expected[2]) + self.assertEqual(len(self._flex_gemm_nodes(gm)), 1) + self.assertEqual(len(self._plain_mm_nodes(gm)), 2) + + def test_fuses_shared_expert_output_into_next_attention_norm(self): + m, k, n, p, q = 8, 16, 512, 64, 96 + + def fn( + a, + first_weight, + routed_output, + residual, + gamma, + second_weight, + third_weight, + ): + shared_output = torch.ops.aten.mm.default(a, first_weight) + shared_output = torch.ops.aten.reshape.default(shared_output, [2, 4, n]) + moe_output = torch.ops.aten.add.Tensor(routed_output, shared_output) + total = torch.ops.aten.add.Tensor(residual, moe_output) + gamma_for_norm = torch.ops.aten.clone.default(gamma) + normalized, _ = torch.ops.aten._fused_rms_norm.default( + total, [n], gamma_for_norm, 1e-5 + ) + second_input = torch.ops.aten.reshape.default(normalized, [m, n]) + second = torch.ops.aten.mm.default(second_input, second_weight) + third_input = torch.ops.aten.reshape.default(normalized, [m, n]) + third = torch.ops.aten.mm.default(third_input, third_weight) + return second, third, total + + inputs = ( + torch.randn(m, k, dtype=torch.bfloat16) * 0.02, + torch.randn(k, n, dtype=torch.bfloat16) * 0.02, + torch.randn(2, 4, n, dtype=torch.bfloat16) * 0.02, + torch.randn(2, 4, n, dtype=torch.bfloat16) * 0.02, + torch.ones(n, dtype=torch.bfloat16), + torch.randn(n, p, dtype=torch.bfloat16) * 0.02, + torch.randn(n, q, dtype=torch.bfloat16) * 0.02, + ) + gm = torch.fx.symbolic_trace(fn) + gm.graph.eliminate_dead_code() + gm.recompile() + FakeTensorProp(gm).propagate(*inputs) + mm_nodes = [ + node for node in gm.graph.nodes if node.target == torch.ops.aten.mm.default + ] + norm = next( + node + for node in gm.graph.nodes + if node.target == torch.ops.aten._fused_rms_norm.default + ) + mm_nodes[0].meta["custom"] = {"module_fqn": "layers.3.moe.shared_experts.w2"} + norm.meta["custom"] = {"module_fqn": "layers.4.attention_norm"} + mm_nodes[1].meta["custom"] = {"module_fqn": "layers.4.attention.wq_a"} + mm_nodes[2].meta["custom"] = {"module_fqn": "layers.4.attention.wkv_a"} + expected = gm(*inputs) + + fuse_f3_residual_rmsnorm_pass(gm) + + actual = gm(*inputs) + self._assert_outputs_close(actual, expected) + self.assertEqual(actual[2], expected[2]) + self.assertEqual(len(self._flex_gemm_nodes(gm)), 1) + self.assertEqual(len(self._plain_mm_nodes(gm)), 2) + + def test_does_not_require_downstream_projection_roles(self): + gm, inputs = self._trace_chain() + mm_nodes = [ + node for node in gm.graph.nodes if node.target == torch.ops.aten.mm.default + ] + mm_nodes[1].meta["custom"] = {"module_fqn": "unrelated.first"} + mm_nodes[2].meta["custom"] = {"module_fqn": "layers.2.feed_forward.w2"} + expected = gm(*inputs) + + fuse_f3_residual_rmsnorm_pass(gm) + + self._assert_outputs_close(gm(*inputs), expected) + self.assertEqual(len(self._flex_gemm_nodes(gm)), 1) + self.assertEqual(len(self._plain_mm_nodes(gm)), 2) + + def test_does_not_fuse_mismatched_boundary_roles(self): + gm, _ = self._trace_chain() + norm = next( + node + for node in gm.graph.nodes + if node.target == torch.ops.aten._fused_rms_norm.default + ) + norm.meta["custom"] = {"module_fqn": "layers.3.attention_norm"} + + fuse_f3_residual_rmsnorm_pass(gm) + + self.assertEqual(self._flex_gemm_nodes(gm), []) + + +class TestF2QRMSNormPass(TestCase): + def _trace_q_chain(self, *, return_rstd=False): + m, k, n, p = 8, 16, 512, 64 + + def fn(a, first_weight, gamma, second_weight): + first = torch.ops.aten.mm.default(a, first_weight) + first = torch.ops.aten.reshape.default(first, [2, 4, n]) + normalized, rstd = torch.ops.aten._fused_rms_norm.default( + first, [n], gamma, 1e-5 + ) + normalized = torch.ops.aten.reshape.default(normalized, [m, n]) + output = torch.ops.aten.mm.default(normalized, second_weight) + return (output, rstd, first, normalized) if return_rstd else output + + inputs = ( + torch.randn(m, k, dtype=torch.bfloat16) * 0.02, + torch.randn(k, n, dtype=torch.bfloat16) * 0.02, + torch.ones(n, dtype=torch.bfloat16), + torch.randn(n, p, dtype=torch.bfloat16) * 0.02, + ) + gm = torch.fx.symbolic_trace(fn) + gm.graph.eliminate_dead_code() + gm.recompile() + FakeTensorProp(gm).propagate(*inputs) + mm_nodes = [ + node for node in gm.graph.nodes if node.target == torch.ops.aten.mm.default + ] + norm = next( + node + for node in gm.graph.nodes + if node.target == torch.ops.aten._fused_rms_norm.default + ) + mm_nodes[0].meta["custom"] = {"module_fqn": "layers.0.attention.wq_a"} + norm.meta["custom"] = {"module_fqn": "layers.0.attention.q_norm"} + mm_nodes[1].meta["custom"] = {"module_fqn": "layers.0.attention.wq_b"} + return gm, inputs + + def _flex_gemm_nodes(self, gm): + return [ + node + for node in gm.graph.nodes + if node.target == torch.ops.higher_order.flex_gemm + ] + + def test_reparameterizes_original_forward_q_norm(self): + gm, inputs = self._trace_q_chain() + a, first_weight, gamma, second_weight = inputs + first = torch.mm(a, first_weight) + weighted = (first.float() * gamma).bfloat16() + partial_mean_square = first.float().reshape(8, -1, 512).square().mean(-1) + rstd = (partial_mean_square.mean(-1, keepdim=True) + 1e-5).rsqrt() + expected = (torch.mm(weighted, second_weight).float() * rstd).bfloat16() + + fuse_f2_q_rmsnorm_pass(gm) + + actual = gm(*inputs) + self.assertEqual(actual.dtype, torch.bfloat16) + torch.testing.assert_close(actual, expected, atol=1e-3, rtol=2e-2) + fused = self._flex_gemm_nodes(gm) + self.assertEqual(len(fused), 2) + self.assertEqual(len(fused[0].args[2]), 3) + self.assertEqual(len(fused[1].args[2]), 3) + self.assertFalse( + any( + node.target == torch.ops.aten._fused_rms_norm.default + for node in gm.graph.nodes + ) + ) + first_body = getattr(gm, fused[0].args[1].target) + first_targets = [node.target for node in first_body.graph.nodes] + self.assertEqual(first_targets.count(torch.ops.aten.mean.dim), 1) + self.assertEqual(first_targets.count(torch.ops.aten.pow.Tensor_Scalar), 1) + + def test_rewrites_recomputed_norm_and_preserves_saved_values(self): + gm, inputs = self._trace_q_chain(return_rstd=True) + a, first_weight, gamma, second_weight = inputs + first = torch.mm(a, first_weight) + partial_mean_square = first.float().reshape(8, -1, 512).square().mean(-1) + rstd_2d = (partial_mean_square.mean(-1, keepdim=True) + 1e-5).rsqrt() + weighted = (first.float() * gamma).bfloat16() + output = (torch.mm(weighted, second_weight).float() * rstd_2d).bfloat16() + normalized = (first.float() * rstd_2d * gamma).bfloat16() + + fuse_f2_q_rmsnorm_pass(gm) + + actual_output, actual_rstd, actual_first, actual_normalized = gm(*inputs) + torch.testing.assert_close(actual_output, output, atol=1e-3, rtol=2e-2) + self.assertEqual(actual_rstd, rstd_2d.reshape(2, 4, 1)) + self.assertEqual(actual_first, first.reshape(2, 4, 512)) + self.assertEqual(actual_normalized, normalized) + self.assertEqual(len(self._flex_gemm_nodes(gm)), 2) + self.assertFalse( + any( + node.target == torch.ops.aten._fused_rms_norm.default + for node in gm.graph.nodes + ) + ) + + +class TestF2KVRMSNormPass(TestCase): + def _trace_kv_chain(self, *, return_saved=False): + m, k, full_width, active_width, output_width = 8, 16, 576, 512, 96 + + def fn(a, first_weight, gamma, second_weight): + first = torch.ops.aten.mm.default(a, first_weight) + first = torch.ops.aten.reshape.default(first, [2, 4, full_width]) + active, tail = torch.ops.aten.split_with_sizes.default( + first, [active_width, full_width - active_width], -1 + ) + normalized, rstd = torch.ops.aten._fused_rms_norm.default( + active, [active_width], gamma, 1e-5 + ) + second_input = torch.ops.aten.reshape.default(normalized, [m, active_width]) + output = torch.ops.aten.mm.default(second_input, second_weight) + if return_saved: + return output, tail, rstd, active, second_input + return output, tail + + inputs = ( + torch.randn(m, k, dtype=torch.bfloat16) * 0.02, + torch.randn(k, full_width, dtype=torch.bfloat16) * 0.02, + torch.ones(active_width, dtype=torch.bfloat16), + torch.randn(active_width, output_width, dtype=torch.bfloat16) * 0.02, + ) + gm = torch.fx.symbolic_trace(fn) + gm.graph.eliminate_dead_code() + gm.recompile() + FakeTensorProp(gm).propagate(*inputs) + mm_nodes = [ + node for node in gm.graph.nodes if node.target == torch.ops.aten.mm.default + ] + norm = next( + node + for node in gm.graph.nodes + if node.target == torch.ops.aten._fused_rms_norm.default + ) + mm_nodes[0].meta["custom"] = {"module_fqn": "layers.0.attention.wkv_a"} + norm.meta["custom"] = {"module_fqn": "layers.0.attention.kv_norm"} + mm_nodes[1].meta["custom"] = {"module_fqn": "layers.0.attention.wkv_b"} + return gm, inputs + + def _expected(self, inputs): + a, first_weight, gamma, second_weight = inputs + first = torch.mm(a, first_weight) + active = first[:, :512] + tail = first[:, 512:].reshape(2, 4, 64) + gamma_full = torch.nn.functional.pad(gamma.reshape(1, 512), (0, 64), value=1) + weighted = (first.float() * gamma_full).bfloat16()[:, :512] + partial = first.float().reshape(8, -1, 64).square().mean(-1)[:, :8] + rstd = (partial.mean(-1, keepdim=True) + 1e-5).rsqrt() + output = (torch.mm(weighted, second_weight).float() * rstd).bfloat16() + normalized = (active.float() * rstd * gamma).bfloat16() + return output, tail, rstd, active, normalized + + def _flex_gemm_nodes(self, gm): + return [ + node + for node in gm.graph.nodes + if node.target == torch.ops.higher_order.flex_gemm + ] + + def test_reparameterizes_segmented_kv_norm(self): + gm, inputs = self._trace_kv_chain() + expected_output, expected_tail, _, _, _ = self._expected(inputs) + + fuse_f2_kv_rmsnorm_pass(gm) + + actual_output, actual_tail = gm(*inputs) + torch.testing.assert_close(actual_output, expected_output, atol=1e-3, rtol=2e-2) + self.assertEqual(actual_tail, expected_tail) + self.assertEqual(len(self._flex_gemm_nodes(gm)), 2) + self.assertFalse( + any( + node.target == torch.ops.aten._fused_rms_norm.default + for node in gm.graph.nodes + ) + ) + + def test_preserves_recomputed_values_and_raw_tail(self): + gm, inputs = self._trace_kv_chain(return_saved=True) + expected_output, expected_tail, rstd, active, normalized = self._expected( + inputs + ) + + fuse_f2_kv_rmsnorm_pass(gm) + + actual_output, actual_tail, actual_rstd, actual_active, actual_normalized = gm( + *inputs + ) + torch.testing.assert_close(actual_output, expected_output, atol=1e-3, rtol=2e-2) + self.assertEqual(actual_tail, expected_tail) + self.assertEqual(actual_rstd, rstd.reshape(2, 4, 1)) + self.assertEqual(actual_active, active.reshape(2, 4, 512)) + self.assertEqual(actual_normalized, normalized) + self.assertEqual(len(self._flex_gemm_nodes(gm)), 2) + + +class TestB4RouterInputGradAddPass(TestCase): + def _trace(self, module_fqn): + x = torch.randn(8, 16) + weight = torch.randn(16, 12) + residual = torch.randn(2, 4, 12, dtype=torch.bfloat16) + + def fn(a, b, other_grad): + router_grad = torch.mm(a, b).reshape(2, 4, 12).bfloat16() + return other_grad + router_grad + + gm = make_fx(fn)(x, weight, residual) + mm = next( + node for node in gm.graph.nodes if node.target == torch.ops.aten.mm.default + ) + mm.meta["custom"] = { + "module_fqn": module_fqn, + "autograd_backward": True, + } + return gm, (x, weight, residual) + + def _flex_gemm_nodes(self, gm): + return [ + node + for node in gm.graph.nodes + if node.target == torch.ops.higher_order.flex_gemm + ] + + def test_fuses_router_cast_and_expert_gradient_add(self): + gm, inputs = self._trace("layers.3.moe.router.gate") + expected = gm(*inputs) + + fuse_b4_router_input_grad_add_pass(gm) + + self.assertEqual(gm(*inputs), expected, exact_dtype=True) + fused = self._flex_gemm_nodes(gm) + self.assertEqual(len(fused), 1) + self.assertEqual(len(fused[0].args[2]), 3) + body = getattr(gm, fused[0].args[1].target) + targets = [node.target for node in body.graph.nodes] + self.assertEqual(targets.count(torch.ops.aten._to_copy.default), 1) + self.assertEqual(targets.count(torch.ops.aten.add.Tensor), 1) + + def test_does_not_fuse_unrelated_fp32_linear(self): + gm, _ = self._trace("layers.3.attention.wo") + + fuse_b4_router_input_grad_add_pass(gm) + + self.assertEqual(self._flex_gemm_nodes(gm), []) + + +class TestB5MLARMSNormBackwardPass(TestCase): + def _trace(self, *, module_fqn="layers.3.attention.wkv_b"): + m, k, n = 8, 16, 512 + + def fn(a, b, x, rstd, gamma): + grad = torch.ops.aten.mm.default(a, b) + grad = torch.ops.aten.reshape.default(grad, [2, 4, n]) + return torch.ops.aten._fused_rms_norm_backward.default( + grad, x, [n], rstd, gamma, [True, True] + ) + + inputs = ( + torch.randn(m, k, dtype=torch.bfloat16) * 0.02, + torch.randn(k, n, dtype=torch.bfloat16) * 0.02, + torch.randn(2, 4, n, dtype=torch.bfloat16) * 0.02, + torch.rand(2, 4, 1, dtype=torch.float32) + 0.5, + torch.randn(n, dtype=torch.bfloat16) * 0.02, + ) + gm = make_fx(fn, tracing_mode="fake")(*inputs) + mm = next( + node for node in gm.graph.nodes if node.target == torch.ops.aten.mm.default + ) + mm.meta["custom"] = {"module_fqn": module_fqn} + mm.meta["autograd_backward"] = True + norm = next( + node + for node in gm.graph.nodes + if node.target == torch.ops.aten._fused_rms_norm_backward.default + ) + norm.meta["autograd_backward"] = True + return gm, inputs + + def _expected(self, inputs): + a, b, x, rstd, gamma = inputs + n = x.shape[-1] + grad = torch.mm(a, b).reshape_as(x).float() + x_hat = x.float() * rstd + grad_x_hat = grad * gamma.float() + row_dot = (x_hat * grad_x_hat).sum(-1, keepdim=True) + grad_input = ((grad_x_hat - (x_hat / n) * row_dot) * rstd).bfloat16() + grad_weight = (grad * x_hat).reshape(-1, n).sum(0).bfloat16() + return grad_input, grad_weight + + def _flex_gemm_nodes(self, gm): + return [ + node + for node in gm.graph.nodes + if node.target == torch.ops.higher_order.flex_gemm + ] + + def test_fuses_mla_projection_and_rmsnorm_backward_partials(self): + gm, inputs = self._trace() + expected = self._expected(inputs) + + fuse_b5_mla_rmsnorm_backward_pass(gm) + + actual = gm(*inputs) + self.assertEqual(actual, expected, exact_dtype=True) + fused = self._flex_gemm_nodes(gm) + self.assertEqual(len(fused), 1) + self.assertEqual(len(fused[0].args[2]), 5) + self.assertFalse( + any( + node.target == torch.ops.aten._fused_rms_norm_backward.default + for node in gm.graph.nodes + ) + ) + body = getattr(gm, fused[0].args[1].target) + targets = [node.target for node in body.graph.nodes] + self.assertEqual(targets.count(torch.ops.aten.sum.dim_IntList), 1) + + def test_does_not_fuse_non_mla_projection(self): + gm, _ = self._trace(module_fqn="layers.3.feed_forward.w2") + + fuse_b5_mla_rmsnorm_backward_pass(gm) + + self.assertEqual(self._flex_gemm_nodes(gm), []) + + +class TestB7AttentionGradMergePass(TestCase): + def _trace(self, *, q_layer=3): + m, kv_width, q_width, model_width = 8, 16, 24, 512 + + def fn(kv_grad, kv_weight, q_grad, q_weight): + kv_input_grad = torch.ops.aten.mm.default(kv_grad, kv_weight) + kv_input_grad = torch.ops.aten.reshape.default( + kv_input_grad, [2, 4, model_width] + ) + q_input_grad = torch.ops.aten.mm.default(q_grad, q_weight) + q_input_grad = torch.ops.aten.reshape.default( + q_input_grad, [2, 4, model_width] + ) + return torch.ops.aten.add.Tensor(kv_input_grad, q_input_grad) + + inputs = ( + torch.randn(m, kv_width, dtype=torch.bfloat16) * 0.02, + torch.randn(kv_width, model_width, dtype=torch.bfloat16) * 0.02, + torch.randn(m, q_width, dtype=torch.bfloat16) * 0.02, + torch.randn(q_width, model_width, dtype=torch.bfloat16) * 0.02, + ) + gm = make_fx(fn)(*inputs) + mm_nodes = [ + node for node in gm.graph.nodes if node.target == torch.ops.aten.mm.default + ] + mm_nodes[0].meta["custom"] = { + "module_fqn": "layers.3.attention.wkv_a", + } + mm_nodes[0].meta["autograd_backward"] = True + mm_nodes[1].meta["custom"] = { + "module_fqn": f"layers.{q_layer}.attention.wq_a", + } + mm_nodes[1].meta["autograd_backward"] = True + return gm, inputs + + def _flex_gemm_nodes(self, gm): + return [ + node + for node in gm.graph.nodes + if node.target == torch.ops.higher_order.flex_gemm + ] + + def test_fuses_q_kv_input_gradient_add(self): + gm, inputs = self._trace() + expected = gm(*inputs) + + fuse_b7_attention_grad_merge_pass(gm) + + self.assertEqual(gm(*inputs), expected) + fused = self._flex_gemm_nodes(gm) + self.assertEqual(len(fused), 1) + self.assertEqual(len(fused[0].args[2]), 3) + fused_output = next(iter(fused[0].users)) + self.assertEqual(fused_output.meta["val"].shape, (8, 512)) + body = getattr(gm, fused[0].args[1].target) + targets = [node.target for node in body.graph.nodes] + self.assertEqual(targets.count(torch.ops.aten.add.Tensor), 1) + + def test_does_not_fuse_different_layers(self): + gm, _ = self._trace(q_layer=4) + + fuse_b7_attention_grad_merge_pass(gm) + + self.assertEqual(self._flex_gemm_nodes(gm), []) + + def test_does_not_fuse_forward_graph(self): + gm, _ = self._trace() + for node in gm.graph.nodes: + node.meta.pop("autograd_backward", None) + + fuse_b7_attention_grad_merge_pass(gm) + + self.assertEqual(self._flex_gemm_nodes(gm), []) + + +class TestCodaPatternRegistry(TestCase): + def test_resolves_configured_order(self): + passes = get_coda_pattern_passes( + [ + "b1_lm_head_input_grad_cast", + "f6_router_sigmoid_bias", + "f3_residual_rmsnorm", + "f4_dense_swiglu", + "b2_dense_swiglu_backward", + "f2_q_rmsnorm", + "f2_kv_rmsnorm", + "b4_router_input_grad_add", + "b5_mla_rmsnorm_backward", + "b6_bf16_weight_grad_cast", + "b7_attention_grad_merge", + ] + ) + self.assertEqual( + passes, + [ + fuse_b1_lm_head_input_grad_cast_pass, + fuse_f6_router_sigmoid_bias_pass, + fuse_f3_residual_rmsnorm_pass, + fuse_f4_dense_swiglu_pass, + fuse_b2_dense_swiglu_backward_pass, + fuse_f2_q_rmsnorm_pass, + fuse_f2_kv_rmsnorm_pass, + fuse_b4_router_input_grad_add_pass, + fuse_b5_mla_rmsnorm_backward_pass, + fuse_b6_bf16_weight_grad_cast_pass, + fuse_b7_attention_grad_merge_pass, + ], + ) + + def test_rejects_unknown_pattern(self): + with self.assertRaisesRegex(ValueError, "Unknown.*not_a_pattern"): + get_coda_pattern_passes(["not_a_pattern"]) + + def test_rejects_duplicate_pattern(self): + with self.assertRaisesRegex(ValueError, "Duplicate.*b6_bf16_weight_grad_cast"): + get_coda_pattern_passes( + ["b6_bf16_weight_grad_cast", "b6_bf16_weight_grad_cast"] + ) + + def test_allows_f4_before_f3(self): + passes = get_coda_pattern_passes(["f4_dense_swiglu", "f3_residual_rmsnorm"]) + + self.assertEqual( + passes, + [fuse_f4_dense_swiglu_pass, fuse_f3_residual_rmsnorm_pass], + ) + + +if __name__ == "__main__": + from torch.testing._internal.common_utils import run_tests + + run_tests() diff --git a/torchtitan/experiments/graph_trainer/tests/test_passes.py b/torchtitan/experiments/graph_trainer/tests/test_passes.py index fb69c3c127..0b28ac191b 100644 --- a/torchtitan/experiments/graph_trainer/tests/test_passes.py +++ b/torchtitan/experiments/graph_trainer/tests/test_passes.py @@ -3420,6 +3420,7 @@ def pass_name(pass_fn): def test_ep_overlap_pass_pipeline_order(self): traced_result, config = self._compile_config_for_ep_overlap_test() + config.compile.coda_patterns = ["b6_bf16_weight_grad_cast"] names = self._compile_pass_names(traced_result, config) dead_code_indices = [ i for i, name in enumerate(names) if name == "eliminate_dead_code_pass" @@ -3463,6 +3464,10 @@ def test_ep_overlap_pass_pipeline_order(self): ) self.assertLess( names.index("concretize_ep_chunk_symbolic_shapes_pass"), + names.index("fuse_b6_bf16_weight_grad_cast_pass"), + ) + self.assertLess( + names.index("fuse_b6_bf16_weight_grad_cast_pass"), names.index("full_inductor_compilation_pass"), )