diff --git a/CONTRIBUTING.md b/CONTRIBUTING.md index 9170a77..ca062df 100644 --- a/CONTRIBUTING.md +++ b/CONTRIBUTING.md @@ -37,6 +37,73 @@ If any step fails, the commit will be aborted. find csrc -name "*.cu" -o -name "*.cuh" | xargs clang-format -i ``` +## Benchmark Annotations + +Kernel benchmarks are auto-generated from annotations in `.cuh` headers. Add `// @bench` annotations above dispatch functions (not `__global__` kernels) to generate benchmarks automatically. + +### Quick example + +```cpp +// @bench rmsnorm +// @title: RMSNorm +// @configs: (N,C) = [(128,4096),(512,4096),(1024,4096)] +// @in: inp float[N * C] random(-1,1), weight floatX[C] bf16(0.75,1.25) +// @out: out float[N * C] +// @scalar: eps=1e-5f +// @bandwidth: (size_t)N * C * sizeof(float) + (size_t)C * sizeof(floatX) + (size_t)N * C * sizeof(float) +// @profile: rmsnorm_forward_profiled N=1024 C=4096 grid=N +inline void rmsnorm_forward(float* out, const float* inp, const floatX* weight, + int N, int C, float eps, cudaStream_t stream = nullptr) { +``` + +### Annotation fields + +| Field | Required | Description | +|-------|----------|-------------| +| `@bench` | Yes | Marks function for codegen. Optional name: `@bench custom_name` | +| `@configs` | Yes | Dimension tuples: `(M,N) = [(32,32),(64,64)]` | +| `@in` | Yes | Input buffers: `name type[size] fill`. Comma-separated for multiple | +| `@out` | Yes | Output buffers: `name type[size]` | +| `@title` | No | Section title in benchmark output | +| `@scalar` | No | Scalar args: `eps=1e-5f` | +| `@bandwidth` | No | Custom bytes expression (default: sum of buffer sizes) | +| `@flops` | No | FLOPs expression for GFLOPS metric | +| `@profile` | No | Profiled variant call with config overrides | +| `@group` | No | Group multiple `@bench` blocks into one output file | +| `@extra-buf` | No | Buffers not in the function signature | +| `@setup` | No | Custom C++ setup code (runs after device allocs) | +| `@pre-setup` | No | Setup code before device allocs (e.g. workspace size calc) | +| `@pre-iter` | No | Per-iteration reset: `buf.zero()` | +| `@runs` | No | Iteration count (default: 100) | + +### Fill types + +- `random(lo,hi)` — fp32 random fill, optional seed: `random(-1,1,43)` +- `bf16(lo,hi)` — bf16 random fill +- `zero` — device-side zero (no host allocation) +- `ones(val)` — constant fill: `ones(1.0/n)` +- `computed` — filled by `@setup` code + +### Multi-line annotations + +Long lines can wrap using `//` continuation: + +```cpp +// @configs: (N,C) = +// [(128,4096),(256,4096),(512,4096),(1024,4096)] +// @bandwidth: (size_t)N * C * sizeof(float) + +// (size_t)C * sizeof(floatX) +``` + +### Generating and checking + +```bash +make benchgen # regenerate all benchmark .cu files +make benchgen-check # verify generated files are up to date (for CI) +``` + +Generated files go to `src/csrc/benchmarks/bench_.cu` and start with `// AUTO-GENERATED by benchgen.py`. Do not edit them manually. + ## Code Style - 4 space indents diff --git a/Makefile b/Makefile index d56cb2c..fda470b 100644 --- a/Makefile +++ b/Makefile @@ -1,4 +1,4 @@ -.PHONY: all build test test-python bench bench-python clean format lint install install-dev force-install +.PHONY: all build test test-python bench bench-python clean format lint install install-dev force-install benchgen benchgen-check CUDA_ARCH ?= "80;86;89;90;100" BUILD_DIR = build @@ -55,12 +55,18 @@ test: build test-python: install LD_PRELOAD=/usr/lib/x86_64-linux-gnu/libstdc++.so.6 pytest src/python/tests -v -bench: build +bench: benchgen build @for b in $(BUILD_DIR)/bench_*; do echo "Running $$b..."; $$b; done bench-python: install LD_PRELOAD=/usr/lib/x86_64-linux-gnu/libstdc++.so.6 python src/python/benchmarks/bench_layer.py --all-configs --backward +benchgen: + python3 scripts/benchgen.py --force + +benchgen-check: + python3 scripts/benchgen.py --check + clean: rm -rf $(BUILD_DIR) $(STAMP_DIR) diff --git a/README.md b/README.md index afee297..61f8aa4 100644 --- a/README.md +++ b/README.md @@ -15,6 +15,8 @@ Once the image builds the first time, it will be cached and will not require a r ### Benchmark +Run benchmark suite + ```bash # python bench modal run runmodal.py --gpu h100 --mode bench --scope python @@ -26,6 +28,16 @@ modal run runmodal.py --gpu h100 --mode bench --scope native modal run runmodal.py --gpu h100 --mode bench --scope all ``` +Generate benchmark files (this is automatically run in the above) + +```bash +# generate the benchmark files +make benchgen + +# check status of benchmark files +make benchgen-check +``` + ### Test ```bash @@ -108,29 +120,19 @@ in section A.1): | Batch | Hidden | n | Forward | Backward | |-------|--------|----|---------|----------| -| 320 | 1280 | 4 | 13.2x | 11.1x | -| 512 | 1920 | 4 | 9.0x | 7.7x | -| 1280 | 2560 | 4 | 5.1x | 3.6x | -| 2560 | 1280 | 4 | 5.0x | 3.5x | -| 128 | 1280 | 8 | 13.6x | 11.5x | -| 256 | 1280 | 8 | 10.3x | 9.8x | -| 32 | 1280 | 32 | 5.8x | 2.9x | -| 64 | 1280 | 32 | 4.7x | 2.2x | -| 128 | 1280 | 32 | 3.5x | 1.5x | +| 320 | 1280 | 4 | 15.20x | 10.07x | +| 512 | 1920 | 4 | 10.52x | 9.20x | +| 1280 | 2560 | 4 | 5.66x | 4.34x | +| 2560 | 1280 | 4 | 5.66x | 4.21x | **Dynamic H Path** (per-batch H values computed via Equations 7-9 from paper): | Batch | Hidden | n | Forward | Backward | |-------|--------|----|---------|----------| -| 320 | 1280 | 4 | 6.7x | 11.0x | -| 512 | 1920 | 4 | 6.8x | 9.0x | -| 1280 | 2560 | 4 | 4.6x | 5.1x | -| 2560 | 1280 | 4 | 4.6x | 5.0x | -| 128 | 1280 | 8 | 6.4x | 11.2x | -| 256 | 1280 | 8 | 6.1x | 10.6x | -| 32 | 1280 | 32 | 1.9x | 3.1x | -| 64 | 1280 | 32 | 1.8x | 2.5x | -| 128 | 1280 | 32 | 1.7x | 1.9x | +| 320 | 1280 | 4 | 7.39x | 3.35x | +| 512 | 1920 | 4 | 7.38x | 3.47x | +| 1280 | 2560 | 4 | 5.33x | 3.07x | +| 2560 | 1280 | 4 | 5.21x | 3.02x | ## Format diff --git a/scripts/benchgen.py b/scripts/benchgen.py new file mode 100644 index 0000000..14650c0 --- /dev/null +++ b/scripts/benchgen.py @@ -0,0 +1,938 @@ +#!/usr/bin/env python3 +"""benchgen.py — Generate benchmark .cu files from @bench annotations in kernel headers. + +Scans src/csrc/kernels/*.cuh for // @bench annotation blocks above dispatch functions, +then generates complete benchmark files into src/csrc/benchmarks/. + +Usage: + python3 scripts/benchgen.py # Generate/overwrite benchmark files + python3 scripts/benchgen.py --check # Check if generated files are up-to-date + python3 scripts/benchgen.py --dry-run # Print generated code without writing +""" + +import argparse +import os +import re +import sys +import tempfile +from dataclasses import dataclass, field +from pathlib import Path + +REPO_ROOT = Path(__file__).resolve().parent.parent +KERNELS_DIR = REPO_ROOT / "src" / "csrc" / "kernels" +BENCH_DIR = REPO_ROOT / "src" / "csrc" / "benchmarks" +GENERATED_HEADER = "// AUTO-GENERATED by benchgen.py — do not edit manually" + + +@dataclass +class Buffer: + name: str + type: str # "float" or "floatX" + size_expr: str # C++ expression e.g. "N*C" + fill: str = ( + "" # "random(lo,hi)", "bf16(lo,hi)", "zero", "ones(val)", "computed", "" + ) + is_input: bool = True + + +@dataclass +class ProfileSpec: + func_name: str + config_vars: dict # e.g. {"N": "1024", "C": "4096"} + max_entries: str = "8" + grid: str = "" # expression for num_blocks in profile_kernel + timeline: bool = True + timeline_blocks: str = "1" + + +@dataclass +class BenchSpec: + name: str # benchmark name (becomes bench_.cu if no group) + group: str # group name — grouped specs share a file + func_name: str # dispatch function name + source_file: str # relative path to source .cuh + configs_dims: list # e.g. ["N", "C"] or ["M", "N", "num_iters"] + configs_values: list # e.g. [(128,4096), (256,4096), ...] + buffers: list # list of Buffer + scalars: dict # e.g. {"eps": "1e-5f"} + flops_expr: str = "" # C++ expression for FLOPs + bandwidth_expr: str = "" # C++ expression for bytes (default: auto-computed) + profile: ProfileSpec | None = None + runs: int = 100 + pre_iter: str = "" # per-iteration setup e.g. "d_d_weight.zero()" + pre_setup_lines: list = field( + default_factory=list + ) # code before device allocs (e.g. workspace calc) + setup_lines: list = field( + default_factory=list + ) # code after device allocs (e.g. host compute + upload) + call_args: list = field(default_factory=list) # ordered arg names for function call + call_expr: str = ( + "" # override for the bench_kernel call (e.g. "fused.forward(...)") + ) + cleanup_lines: list = field( + default_factory=list + ) # code after bench loop (e.g. "fused.destroy()") + tflops_expr: str = "" # FLOPs expression for TFLOPS metric (separate from GFLOPS) + metrics: list = field(default_factory=list) # custom metrics: [(name, expr), ...] + title: str = "" # section title override + + +def parse_annotation_block(lines: list[str]) -> dict[str, list[str]]: + """Parse a block of // @key: value lines into a dict. + + Continuation lines (starting with // but not // @) are appended to the + last value of the current key, allowing long annotations to wrap. + """ + result = {} + current_key = None + for line in lines: + line = line.strip() + if line.startswith("// @"): + line = line[4:] # strip "// @" + # Check for key: value + m = re.match(r"(\S+?):\s?(.*)", line) + if m: + current_key = m.group(1) + val = m.group(2) + # Strip for non-setup keys (setup/pre-setup/cleanup preserve indentation) + if current_key in ("setup", "pre-setup", "cleanup"): + val = val.rstrip() + else: + val = val.strip() + if current_key not in result: + result[current_key] = [] + if val: + result[current_key].append(val) + else: + # Key without colon — e.g. "@bench" or "@bench custom_name" + parts = line.split(None, 1) + current_key = parts[0] + val = parts[1].strip() if len(parts) > 1 else "" + if current_key not in result: + result[current_key] = [] + if val: + result[current_key].append(val) + elif line.startswith("//") and current_key is not None: + # Continuation line — append to the last value of current key + cont = line[2:].strip() + if cont: + if current_key in result and result[current_key]: + result[current_key][-1] += " " + cont + else: + if current_key not in result: + result[current_key] = [] + result[current_key].append(cont) + return result + + +def parse_configs(config_str: str) -> tuple[list[str], list[tuple]]: + """Parse config string like '(N,C) = [(128,4096),(256,4096)]' into dims and values.""" + m = re.match(r"\(([^)]+)\)\s*=\s*\[(.+)\]", config_str) + if not m: + raise ValueError(f"Invalid @configs format: {config_str}") + dims = [d.strip() for d in m.group(1).split(",")] + values_str = m.group(2) + values = [] + for tup in re.findall(r"\(([^)]+)\)", values_str): + vals = tuple(v.strip() for v in tup.split(",")) + values.append(vals) + return dims, values + + +def parse_buffer(buf_str: str, is_input: bool) -> Buffer: + """Parse buffer spec like 'inp float[N*C] random(-1,1)' or 'out float[N*C]'.""" + # Pattern: name type[size] [fill] + m = re.match(r"(\w+)\s+(\w+)\[([^\]]+)\](?:\s+(.+))?", buf_str.strip()) + if not m: + raise ValueError(f"Invalid buffer spec: {buf_str}") + name = m.group(1) + typ = m.group(2) + size = m.group(3) + fill_str = m.group(4).strip() if m.group(4) else "" + return Buffer(name=name, type=typ, size_expr=size, fill=fill_str, is_input=is_input) + + +def parse_profile(profile_str: str) -> ProfileSpec: + """Parse profile spec like 'rmsnorm_forward_profiled N=1024 C=4096 max_entries=8 grid=N'.""" + parts = profile_str.split() + func_name = parts[0] + config_vars = {} + max_entries = "8" + grid = "" + timeline = False + timeline_blocks = "1" + for part in parts[1:]: + if "=" not in part: + continue + k, v = part.split("=", 1) + if k == "max_entries": + max_entries = v + elif k == "grid": + grid = v + elif k == "timeline": + timeline = v.lower() == "true" + elif k == "timeline_blocks": + timeline_blocks = v + else: + config_vars[k] = v + return ProfileSpec( + func_name=func_name, + config_vars=config_vars, + max_entries=max_entries, + grid=grid, + timeline=timeline, + timeline_blocks=timeline_blocks, + ) + + +def parse_func_signature(sig_lines: str) -> tuple[str, list[tuple[str, str]]]: + """Parse function signature to extract name and (type, name) params.""" + sig = " ".join(sig_lines.split()) + # Extract function name + m = re.search(r"(\w+)\s*\(", sig) + if not m: + raise ValueError(f"Cannot find function name in: {sig}") + func_name = m.group(1) + # Extract params + paren_start = sig.index("(") + paren_end = sig.rindex(")") + params_str = sig[paren_start + 1 : paren_end] + params = [] + for p in split_params(params_str): + p = p.strip() + if not p: + continue + # Remove default value + if "=" in p: + p = p[: p.index("=")].strip() + # Split type and name + parts = p.rsplit(None, 1) + if len(parts) == 2: + params.append((parts[0].strip(), parts[1].strip())) + return func_name, params + + +def split_params(s: str) -> list[str]: + """Split comma-separated params, respecting nested angle brackets and parens.""" + depth = 0 + parts = [] + current = [] + for c in s: + if c in "(<": + depth += 1 + elif c in ")>": + depth -= 1 + elif c == "," and depth == 0: + parts.append("".join(current)) + current = [] + continue + current.append(c) + if current: + parts.append("".join(current)) + return parts + + +def scan_kernel_file(filepath: Path) -> list[BenchSpec]: + """Scan a .cuh file for @bench annotation blocks and return BenchSpecs.""" + lines = filepath.read_text().splitlines() + specs = [] + + i = 0 + while i < len(lines): + line = lines[i].strip() + if not line.startswith("// @bench"): + i += 1 + continue + + # Found a @bench block — collect all contiguous // @ lines and // continuation lines + block_start = i + block_lines = [] + while i < len(lines): + stripped = lines[i].strip() + if stripped.startswith("// @") or ( + stripped.startswith("//") + and not stripped.startswith("///") + and block_lines + ): + block_lines.append(lines[i]) + i += 1 + else: + break + + # Parse annotation block first to check for @call + ann = parse_annotation_block(block_lines) + + # If @call is present, function signature is optional + call_expr = " ".join(ann.get("call", [])) + if call_expr: + # No function signature to parse — don't advance i + func_name = "" + params = [] + else: + # Skip blank lines between annotation and function + while i < len(lines) and not lines[i].strip(): + i += 1 + + # Collect function signature (until opening brace) + sig_lines = [] + while i < len(lines): + sig_lines.append(lines[i]) + if "{" in lines[i]: + break + i += 1 + + sig_text = " ".join(l.strip() for l in sig_lines) + # Remove the { and everything after + sig_text = sig_text[: sig_text.index("{")] + + func_name, params = parse_func_signature(sig_text) + + # Build BenchSpec + bench_name = ann.get("bench", [""])[0] or func_name + group = ann.get("group", [""])[0] if "group" in ann else "" + + # Parse configs + config_str = " ".join(ann.get("configs", [])) + dims, values = parse_configs(config_str) + + # Parse buffers + buffers = [] + for buf_str in ann.get("in", []): + # Multiple buffers can be comma-separated + for b in split_buf_list(buf_str): + buffers.append(parse_buffer(b, is_input=True)) + for buf_str in ann.get("out", []): + for b in split_buf_list(buf_str): + buffers.append(parse_buffer(b, is_input=False)) + for buf_str in ann.get("extra-buf", []): + for b in split_buf_list(buf_str): + buffers.append(parse_buffer(b, is_input=True)) + + # Parse scalars + scalars = {} + for s in ann.get("scalar", []): + for pair in s.split(","): + pair = pair.strip() + if "=" in pair: + k, v = pair.split("=", 1) + scalars[k.strip()] = v.strip() + + # Parse profile + profile = None + if "profile" in ann: + profile = parse_profile(" ".join(ann["profile"])) + + # Parse other fields + flops_expr = " ".join(ann.get("flops", [])) + tflops_expr = " ".join(ann.get("tflops", [])) + bandwidth_expr = " ".join(ann.get("bandwidth", [])) + runs = int(ann.get("runs", ["100"])[0]) + pre_iter = " ".join(ann.get("pre-iter", [])) + pre_setup_lines = ann.get("pre-setup", []) + setup_lines = ann.get("setup", []) + cleanup_lines = ann.get("cleanup", []) + title = " ".join(ann.get("title", [])) + + # Parse custom metrics: "name=expression" + metrics = [] + for m in ann.get("metric", []): + if "=" in m: + mname, mexpr = m.split("=", 1) + metrics.append((mname.strip(), mexpr.strip())) + + # Determine call args from function params + call_args = [] + for _ptype, pname in params: + if pname == "stream": + continue # skip stream param (uses default nullptr) + call_args.append(pname) + + rel_path = filepath.relative_to(REPO_ROOT / "src" / "csrc") + + spec = BenchSpec( + name=bench_name, + group=group, + func_name=func_name, + source_file=str(rel_path), + configs_dims=dims, + configs_values=values, + buffers=buffers, + scalars=scalars, + flops_expr=flops_expr, + tflops_expr=tflops_expr, + bandwidth_expr=bandwidth_expr, + profile=profile, + runs=runs, + pre_iter=pre_iter, + pre_setup_lines=pre_setup_lines, + setup_lines=setup_lines, + call_args=call_args, + call_expr=call_expr, + cleanup_lines=cleanup_lines, + metrics=metrics, + title=title, + ) + specs.append(spec) + i += 1 + + return specs + + +def split_buf_list(buf_str: str) -> list[str]: + """Split comma-separated buffer specs, respecting brackets.""" + depth = 0 + parts = [] + current = [] + for c in buf_str: + if c in "([": + depth += 1 + elif c in ")]": + depth -= 1 + elif c == "," and depth == 0: + parts.append("".join(current).strip()) + current = [] + continue + current.append(c) + if current: + s = "".join(current).strip() + if s: + parts.append(s) + return parts + + +def float_literal(val: str) -> str: + """Ensure a value is a valid C++ float literal (e.g. -1 -> -1.0f, 0.75 -> 0.75f).""" + val = val.strip() + if val.endswith("f"): + return val + # Check if it has a decimal point + core = val.lstrip("-") + if "." not in core: + return val + ".0f" + return val + "f" + + +def gen_fill_code(buf: Buffer, host_name: str, size_expr: str) -> list[str]: + """Generate fill code for a host buffer. Zero fills are handled device-side, not here.""" + code = [] + fill = buf.fill + if not fill or fill == "computed" or fill == "zero": + return code + if fill.startswith("ones("): + val = fill[5:-1] + code.append(f" for (int _i = 0; _i < {size_expr}; _i++)") + code.append(f" {host_name}.ptr[_i] = ({buf.type})({val});") + elif fill.startswith("random("): + args = fill[7:-1] + parts = [a.strip() for a in args.split(",")] + lo, hi = float_literal(parts[0]), float_literal(parts[1]) + seed = parts[2] if len(parts) > 2 else "42" + if seed == "42": + code.append(f" fill_random({host_name}, {size_expr}, {lo}, {hi});") + else: + code.append( + f" fill_random({host_name}, {size_expr}, {lo}, {hi}, {seed});" + ) + elif fill.startswith("bf16("): + args = fill[5:-1] + parts = [a.strip() for a in args.split(",")] + lo, hi = float_literal(parts[0]), float_literal(parts[1]) + seed = parts[2] if len(parts) > 2 else "42" + if seed == "42": + code.append( + f" fill_random_bf16({host_name}, {size_expr}, {lo}, {hi});" + ) + else: + code.append( + f" fill_random_bf16({host_name}, {size_expr}, {lo}, {hi}, {seed});" + ) + return code + + +def compute_bandwidth_expr(buffers: list[Buffer]) -> str: + """Auto-compute bandwidth expression from buffer sizes.""" + read_parts = [] + write_parts = [] + for buf in buffers: + sizeof = f"sizeof({buf.type})" + term = f"(size_t){buf.size_expr} * {sizeof}" + if buf.is_input: + read_parts.append(term) + else: + write_parts.append(term) + all_parts = read_parts + write_parts + if not all_parts: + return "" + return " + ".join(all_parts) + + +def build_call_args(spec: BenchSpec) -> str: + """Build the function call argument string.""" + args = [] + for arg in spec.call_args: + # Check if it's a buffer name + buf = next((b for b in spec.buffers if b.name == arg), None) + if buf: + args.append(f"d_{arg}") + elif arg in spec.scalars: + args.append(spec.scalars[arg]) + elif arg in spec.configs_dims: + args.append(arg) + else: + # Check if it's a scalar from annotation or a config dim + args.append(arg) + return ", ".join(args) + + +def generate_section(spec: BenchSpec, is_grouped: bool) -> list[str]: + """Generate a benchmark section for one BenchSpec.""" + code = [] + fn = spec.func_name + + # Section header for grouped benchmarks + if is_grouped: + title = spec.title or fn + code.append(f' printf("\\n{title}\\n");') + + # Determine which columns to print + dims = spec.configs_dims + has_flops = bool(spec.flops_expr) + has_tflops = bool(spec.tflops_expr) + has_per_iter = "num_iters" in dims or "iters" in dims + has_bw = bool(spec.bandwidth_expr) or bool(spec.buffers) + + # Compute column widths and format strings + col_widths = [] + fmt_parts_v = [] + header_args_h = [] + + for d in dims: + w = 8 if d not in ("iters", "num_iters") else 6 + col_widths.append(w) + fmt_parts_v.append(f"%{w}d") + header_args_h.append(f'"{d}"') + + col_widths.append(12) + fmt_parts_v.append("%12.2f") + header_args_h.append('"Time (us)"') + + if has_per_iter: + col_widths.append(10) + fmt_parts_v.append("%10.2f") + header_args_h.append('"us/iter"') + + if has_tflops: + col_widths.append(12) + fmt_parts_v.append("%12.2f") + header_args_h.append('"TFLOPS"') + elif has_flops: + col_widths.append(10) + fmt_parts_v.append("%10.2f") + header_args_h.append('"GFLOPS"') + + if has_bw: + col_widths.append(12) + fmt_parts_v.append("%12.2f" if (has_flops or has_tflops) else "%12.0f") + header_args_h.append('"Bandwidth (GB/s)"') + + for mname, _mexpr in spec.metrics: + col_widths.append(14) + fmt_parts_v.append("%14.0f") + header_args_h.append(f'"{mname}"') + + # Build format strings (space-separated) + fmt_str_h = " ".join(f"%{w}s" for w in col_widths) + fmt_str_v = " ".join(fmt_parts_v) + + code.append(f' printf("{fmt_str_h}\\n", {", ".join(header_args_h)});') + sep_len = sum(col_widths) + len(col_widths) - 1 + code.append(f' printf("{"-" * sep_len}\\n");') + code.append("") + + # Config loop + code.append(" for (int c = 0; c < num_configs; c++) {") + for i, d in enumerate(dims): + code.append(f" int {d} = configs[c][{i}];") + code.append("") + + # Categorize buffers + input_bufs = [b for b in spec.buffers if b.is_input] + output_bufs = [b for b in spec.buffers if not b.is_input] + + # Host allocations (for any input buffer that needs host data) + for buf in input_bufs: + if buf.fill and buf.fill != "zero": + code.append(f" HostMem<{buf.type}> h_{buf.name}({buf.size_expr});") + + # Host fills (auto-generated — computed buffers are filled by @setup) + for buf in input_bufs: + fill_lines = gen_fill_code(buf, f"h_{buf.name}", buf.size_expr) + code.extend(fill_lines) + + if any(b.fill and b.fill not in ("zero", "computed", "") for b in input_bufs): + code.append("") + + # Pre-setup code (before device allocs — e.g. workspace size calculation) + if spec.pre_setup_lines: + for line in spec.pre_setup_lines: + code.append(f" {line}") + code.append("") + + # Device allocations + for buf in input_bufs: + code.append(f" DeviceMem<{buf.type}> d_{buf.name}({buf.size_expr});") + for buf in output_bufs: + code.append(f" DeviceMem<{buf.type}> d_{buf.name}({buf.size_expr});") + + # Uploads (non-computed, non-zero) + for buf in input_bufs: + if buf.fill == "zero": + code.append(f" d_{buf.name}.zero();") + elif buf.fill and buf.fill != "computed": + code.append(f" d_{buf.name}.upload(h_{buf.name});") + + # Post-alloc setup code (after device allocs — e.g. host computation + upload, forward pass) + if spec.setup_lines: + code.append("") + for setup_line in spec.setup_lines: + code.append(f" {setup_line}") + + code.append("") + + # Bandwidth computation + bw_expr = spec.bandwidth_expr or ( + compute_bandwidth_expr(spec.buffers) if spec.buffers else "" + ) + if bw_expr: + code.append(f" size_t total_bytes = {bw_expr};") + code.append("") + + # Flops computation (declare before bench_kernel for TFLOPS) + if has_tflops: + code.append(f" double flops = {spec.tflops_expr};") + code.append("") + + # Bench kernel call + if spec.call_expr: + call = spec.call_expr + else: + call = f"{fn}({build_call_args(spec)})" + if spec.pre_iter: + code.append(" float avg_time_ms = bench_kernel(") + code.append(f" [&]() {{ {call}; }},") + code.append(f" bench_runs, flusher, [&]() {{ {spec.pre_iter}; }});") + else: + code.append(" float avg_time_ms = bench_kernel(") + code.append(f" [&]() {{ {call}; }},") + code.append(" bench_runs, flusher);") + code.append("") + + # Compute metrics + code.append(" float time_us = avg_time_ms * 1000.0f;") + + iter_dim = None + for d in ("num_iters", "iters"): + if d in dims: + iter_dim = d + break + + if has_per_iter: + code.append(f" float time_per_iter_us = time_us / {iter_dim};") + + if has_tflops: + code.append(" float tflops = (flops / 1e12f) / (avg_time_ms / 1e3f);") + elif has_flops: + code.append( + f" float gflops = (({spec.flops_expr}) / 1e9f) / (avg_time_ms / 1e3f);" + ) + + if bw_expr: + code.append(" float bw = (total_bytes / 1e9f) / (avg_time_ms / 1e3f);") + + for mname, mexpr in spec.metrics: + varname = mname.lower().replace("/", "_per_").replace(" ", "_") + code.append(f" float {varname} = {mexpr};") + + code.append("") + + # Printf values + val_args = [] + for d in dims: + val_args.append(d) + val_args.append("time_us") + if has_per_iter: + val_args.append("time_per_iter_us") + if has_tflops: + val_args.append("tflops") + elif has_flops: + val_args.append("gflops") + if bw_expr: + val_args.append("bw") + for mname, _mexpr in spec.metrics: + varname = mname.lower().replace("/", "_per_").replace(" ", "_") + val_args.append(varname) + + code.append(f' printf("{fmt_str_v}\\n", {", ".join(val_args)});') + + # Cleanup code (per-config, e.g. fused.destroy()) + if spec.cleanup_lines: + code.append("") + for cl in spec.cleanup_lines: + code.append(f" {cl}") + + code.append(" }") + code.append("") + + # Profile section + if spec.profile: + p = spec.profile + config_desc = ", ".join(f"{k}={v}" for k, v in p.config_vars.items()) + code.append(f' printf("\\n--- Phase Breakdown ({config_desc}) ---\\n");') + code.append(" {") + + # Declare config vars + for k, v in p.config_vars.items(): + code.append(f" int {k} = {v};") + code.append("") + + # Host allocations for profile section + for buf in input_bufs: + if buf.fill and buf.fill != "zero": + code.append( + f" HostMem<{buf.type}> h_{buf.name}({buf.size_expr});" + ) + + for buf in input_bufs: + fill_lines = gen_fill_code(buf, f"h_{buf.name}", buf.size_expr) + code.extend(fill_lines) + + if any(b.fill and b.fill not in ("zero", "computed", "") for b in input_bufs): + code.append("") + + # Pre-setup before device allocs + if spec.pre_setup_lines: + for line in spec.pre_setup_lines: + code.append(f" {line}") + code.append("") + + for buf in input_bufs: + code.append(f" DeviceMem<{buf.type}> d_{buf.name}({buf.size_expr});") + for buf in output_bufs: + code.append(f" DeviceMem<{buf.type}> d_{buf.name}({buf.size_expr});") + + for buf in input_bufs: + if buf.fill == "zero": + code.append(f" d_{buf.name}.zero();") + elif buf.fill and buf.fill != "computed": + code.append(f" d_{buf.name}.upload(h_{buf.name});") + + # Post-alloc setup for profile section + if spec.setup_lines: + code.append("") + for setup_line in spec.setup_lines: + code.append(f" {setup_line}") + + code.append("") + + # Build profile call args + profile_call_args = build_call_args(spec) + # Replace function-specific args with profiled variant's args + # The profiled function has extra profiler_buf, max_entries params + profile_call = f"{p.func_name}({profile_call_args}, profiler_buf, max_ent)" + + grid = p.grid or "1" + code.append(" profile_kernel(") + code.append(" [&](int64_t* profiler_buf, int max_ent) {") + code.append(f" {profile_call};") + code.append(" },") + if p.timeline: + code.append( + f" {grid}, max_entries, flusher, true, {p.timeline_blocks});" + ) + else: + code.append(f" {grid}, max_entries, flusher);") + code.append(" }") + code.append("") + + return code + + +def generate_benchmark_file( + specs: list[BenchSpec], +) -> str: + """Generate a complete benchmark .cu file from one or more BenchSpecs.""" + # Determine source file (use first spec's source for include path) + source_file = specs[0].source_file + include_path = f"../{source_file}" + + # File name + if specs[0].group: + file_name = specs[0].group + else: + file_name = specs[0].name + + lines = [] + lines.append(GENERATED_HEADER) + lines.append(f"// Source: src/csrc/{source_file}") + lines.append("// Regenerate: make benchgen") + lines.append("") + lines.append("#include ") + lines.append("#include ") + lines.append("") + lines.append(f'#include "{include_path}"') + lines.append('#include "../include/bench_harness.cuh"') + lines.append("") + lines.append("using namespace mhc;") + lines.append("") + lines.append("int main() {") + lines.append(f" const int bench_runs = {specs[0].runs};") + + # Check if any spec uses scalars + all_scalars = {} + for spec in specs: + all_scalars.update(spec.scalars) + for k, v in all_scalars.items(): + lines.append(f" const float {k} = {v};") + + # Check if any spec has profiling + has_profile = any(s.profile for s in specs) + if has_profile: + lines.append(" const int max_entries = 8;") + + lines.append("") + lines.append(" L2Flusher flusher;") + lines.append("") + + # Configs — use the first spec's configs (all grouped specs should share) + dims = specs[0].configs_dims + values = specs[0].configs_values + ndims = len(dims) + + lines.append(f" int configs[][{ndims}] = {{") + for v in values: + vals = ", ".join(str(x) for x in v) + lines.append(f" {{{vals}}},") + lines.append(" };") + lines.append(" int num_configs = sizeof(configs) / sizeof(configs[0]);") + lines.append("") + + is_grouped = len(specs) > 1 + + # Title — for grouped benchmarks, use group name (sections have their own titles) + if is_grouped: + title = specs[0].group.replace("_", " ").title() + else: + title = specs[0].title or specs[0].name.replace("_", " ").title() + lines.append(f' printf("{title} Benchmark\\n");') + lines.append(f' printf("{"=" * 70}\\n");') + + for spec in specs: + section = generate_section(spec, is_grouped) + lines.extend(section) + + lines.append(" return 0;") + lines.append("}") + lines.append("") + + return "\n".join(lines) + + +def main(): + parser = argparse.ArgumentParser( + description="Generate benchmark files from annotations" + ) + parser.add_argument( + "--check", action="store_true", help="Check if generated files are up-to-date" + ) + parser.add_argument( + "--dry-run", action="store_true", help="Print generated code without writing" + ) + parser.add_argument( + "--force", + action="store_true", + help="Overwrite even if file lacks auto-generated header", + ) + args = parser.parse_args() + + # Scan all kernel files + all_specs = [] + for cuh_file in sorted(KERNELS_DIR.glob("*.cuh")): + specs = scan_kernel_file(cuh_file) + if specs: + all_specs.extend(specs) + print( + f" Found {len(specs)} @bench annotation(s) in {cuh_file.name}", + file=sys.stderr, + ) + + if not all_specs: + print("No @bench annotations found in kernel files.", file=sys.stderr) + return + + # Group specs by group name or individual name + groups = {} + for spec in all_specs: + key = spec.group if spec.group else spec.name + if key not in groups: + groups[key] = [] + groups[key].append(spec) + + # Generate files + generated = {} + for group_name, specs in groups.items(): + content = generate_benchmark_file(specs) + out_name = f"bench_{group_name}.cu" + generated[out_name] = content + + if args.dry_run: + for name, content in sorted(generated.items()): + print(f"=== {name} ===") + print(content) + print() + return + + if args.check: + ok = True + for name, content in sorted(generated.items()): + out_path = BENCH_DIR / name + if not out_path.exists(): + print(f"MISSING: {out_path}", file=sys.stderr) + ok = False + continue + existing = out_path.read_text() + if existing != content: + print(f"OUT OF DATE: {out_path}", file=sys.stderr) + ok = False + else: + print(f" OK: {name}", file=sys.stderr) + if not ok: + print( + "\nGenerated benchmarks are out of date. Run 'make benchgen' to update.", + file=sys.stderr, + ) + sys.exit(1) + else: + print("\nAll generated benchmarks are up to date.", file=sys.stderr) + return + + # Write files + for name, content in sorted(generated.items()): + out_path = BENCH_DIR / name + # Only overwrite if it's auto-generated, doesn't exist, or --force + if out_path.exists() and not args.force: + first_line = out_path.read_text().split("\n", 1)[0] + if not first_line.startswith(GENERATED_HEADER[:30]): + print( + f" SKIP: {name} (not auto-generated, use --force to overwrite)", + file=sys.stderr, + ) + continue + out_path.write_text(content) + print(f" Generated: {name}", file=sys.stderr) + + print(f"\nGenerated {len(generated)} benchmark file(s).", file=sys.stderr) + + +if __name__ == "__main__": + main() diff --git a/src/csrc/benchmarks/bench_fused_rmsnorm_matmul.cu b/src/csrc/benchmarks/bench_fused_rmsnorm_matmul.cu index b906ac1..a3db6d3 100644 --- a/src/csrc/benchmarks/bench_fused_rmsnorm_matmul.cu +++ b/src/csrc/benchmarks/bench_fused_rmsnorm_matmul.cu @@ -1,11 +1,12 @@ +// AUTO-GENERATED by benchgen.py — do not edit manually +// Source: src/csrc/kernels/fused_rmsnorm_matmul.cuh +// Regenerate: make benchgen + #include -#include -#include #include -#include "../include/mhc_types.h" -#include "../include/utils.cuh" #include "../kernels/fused_rmsnorm_matmul.cuh" +#include "../include/bench_harness.cuh" using namespace mhc; @@ -14,83 +15,54 @@ int main() { L2Flusher flusher; - struct Config { - int M; - int N; - int K; - }; - - Config configs[] = { + int configs[][3] = { {128, 4096, 4096}, {256, 4096, 4096}, {512, 4096, 4096}, {1024, 4096, 4096}, {2048, 4096, 4096}, {1024, 8192, 4096}, {2048, 8192, 4096}, }; int num_configs = sizeof(configs) / sizeof(configs[0]); printf("Fused RMSNorm + MatMul Benchmark\n"); - printf("==========================================================================\n"); + printf("======================================================================\n"); printf("%8s %8s %8s %12s %12s %12s\n", "M", "N", "K", "Time (us)", "TFLOPS", "Bandwidth (GB/s)"); - printf("--------------------------------------------------------------------------\n"); + printf("-----------------------------------------------------------------\n"); for (int c = 0; c < num_configs; c++) { - int M = configs[c].M; - int N = configs[c].N; - int K = configs[c].K; - - floatX* h_inp = (floatX*)malloc(M * K * sizeof(floatX)); - floatX* h_weight = (floatX*)malloc(N * K * sizeof(floatX)); - - srand(42); - for (int i = 0; i < M * K; i++) { - h_inp[i] = (floatX)((float)rand() / RAND_MAX * 2.0f - 1.0f); - } - for (int i = 0; i < N * K; i++) { - h_weight[i] = (floatX)((float)rand() / RAND_MAX * 0.5f + 0.75f); - } - - floatX *d_inp, *d_weight; - float* d_out; - CHECK_CUDA(cudaMalloc(&d_inp, M * K * sizeof(floatX))); - CHECK_CUDA(cudaMalloc(&d_weight, N * K * sizeof(floatX))); - CHECK_CUDA(cudaMalloc(&d_out, M * N * sizeof(float))); - - CHECK_CUDA(cudaMemcpy(d_inp, h_inp, M * K * sizeof(floatX), cudaMemcpyHostToDevice)); - CHECK_CUDA(cudaMemcpy(d_weight, h_weight, N * K * sizeof(floatX), cudaMemcpyHostToDevice)); + int M = configs[c][0]; + int N = configs[c][1]; + int K = configs[c][2]; - FusedRMSNormMatmul fused; - fused.init(M, N, K); + HostMem h_inp(M * K); + HostMem h_weight(N * K); + fill_random_bf16(h_inp, M * K, -1.0f, 1.0f); + fill_random_bf16(h_weight, N * K, 0.75f, 1.25f); - double flops = 2.0 * (double)M * (double)N * (double)K; + DeviceMem d_inp(M * K); + DeviceMem d_weight(N * K); + DeviceMem d_out(M * N); + d_inp.upload(h_inp); + d_weight.upload(h_weight); - size_t bytes_read = M * K * sizeof(floatX) + N * K * sizeof(floatX); - size_t bytes_write = M * N * sizeof(float); - size_t total_bytes = bytes_read + bytes_write; + FusedRMSNormMatmul fused; + fused.init(M, N, K); - BenchTimer timer; - float total_time = 0.0f; + size_t total_bytes = + M * K * sizeof(floatX) + N * K * sizeof(floatX) + M * N * sizeof(float); - for (int i = 0; i < bench_runs; i++) { - flusher.flush(); + double flops = 2.0 * (double)M * (double)N * (double)K; - timer.record_start(); - fused.forward(d_out, d_inp, d_weight); - timer.record_stop(); - total_time += timer.elapsed_ms(); - } + float avg_time_ms = + bench_kernel([&]() { fused.forward(d_out, d_inp, d_weight); }, bench_runs, flusher); - float avg_time_ms = total_time / bench_runs; + float time_us = avg_time_ms * 1000.0f; float tflops = (flops / 1e12f) / (avg_time_ms / 1e3f); float bw = (total_bytes / 1e9f) / (avg_time_ms / 1e3f); - printf("%8d %8d %8d %12.2f %12.2f %12.2f\n", M, N, K, avg_time_ms * 1000.0f, tflops, bw); + printf("%8d %8d %8d %12.2f %12.2f %12.2f\n", M, N, K, time_us, tflops, bw); fused.destroy(); - cudaFree(d_inp); - cudaFree(d_weight); - cudaFree(d_out); - free(h_inp); - free(h_weight); } + return 0; } diff --git a/src/csrc/benchmarks/bench_fused_rmsnorm_matmul_backward.cu b/src/csrc/benchmarks/bench_fused_rmsnorm_matmul_backward.cu index 654c2e4..bc47823 100644 --- a/src/csrc/benchmarks/bench_fused_rmsnorm_matmul_backward.cu +++ b/src/csrc/benchmarks/bench_fused_rmsnorm_matmul_backward.cu @@ -1,11 +1,12 @@ +// AUTO-GENERATED by benchgen.py — do not edit manually +// Source: src/csrc/kernels/fused_rmsnorm_matmul.cuh +// Regenerate: make benchgen + #include -#include -#include #include -#include "../include/mhc_types.h" -#include "../include/utils.cuh" #include "../kernels/fused_rmsnorm_matmul.cuh" +#include "../include/bench_harness.cuh" using namespace mhc; @@ -14,108 +15,72 @@ int main() { L2Flusher flusher; - struct Config { - int M; - int N; - int K; - }; - - Config configs[] = { + int configs[][3] = { {128, 4096, 4096}, {256, 4096, 4096}, {512, 4096, 4096}, {1024, 4096, 4096}, {2048, 4096, 4096}, {1024, 8192, 4096}, {2048, 8192, 4096}, }; int num_configs = sizeof(configs) / sizeof(configs[0]); printf("Fused RMSNorm + MatMul Backward Benchmark\n"); - printf("==========================================================================\n"); + printf("======================================================================\n"); printf("%8s %8s %8s %12s %12s %12s\n", "M", "N", "K", "Time (us)", "TFLOPS", "Bandwidth (GB/s)"); - printf("--------------------------------------------------------------------------\n"); + printf("-----------------------------------------------------------------\n"); for (int c = 0; c < num_configs; c++) { - int M = configs[c].M; - int N = configs[c].N; - int K = configs[c].K; - - floatX* h_inp = (floatX*)malloc(M * K * sizeof(floatX)); - floatX* h_weight = (floatX*)malloc(N * K * sizeof(floatX)); - float* h_grad = (float*)malloc(M * N * sizeof(float)); - float* h_rms = (float*)malloc(M * sizeof(float)); - - srand(42); - for (int i = 0; i < M * K; i++) { - h_inp[i] = (floatX)((float)rand() / RAND_MAX * 2.0f - 1.0f); - } - for (int i = 0; i < N * K; i++) { - h_weight[i] = (floatX)((float)rand() / RAND_MAX * 0.5f + 0.75f); - } - for (int i = 0; i < M * N; i++) { - h_grad[i] = (float)rand() / RAND_MAX * 2.0f - 1.0f; - } + int M = configs[c][0]; + int N = configs[c][1]; + int K = configs[c][2]; + + HostMem h_inp(M * K); + HostMem h_weight(N * K); + HostMem h_grad(M * N); + HostMem h_rms(M); + fill_random_bf16(h_inp, M * K, -1.0f, 1.0f); + fill_random_bf16(h_weight, N * K, 0.75f, 1.25f); + fill_random(h_grad, M * N, -1.0f, 1.0f, 43); + + DeviceMem d_inp(M * K); + DeviceMem d_weight(N * K); + DeviceMem d_grad(M * N); + DeviceMem d_rms(M); + DeviceMem d_dW(N * K); + DeviceMem d_dx(M * K); + d_inp.upload(h_inp); + d_weight.upload(h_weight); + d_grad.upload(h_grad); + for (int i = 0; i < M; i++) { float sum_sq = 0.0f; for (int j = 0; j < K; j++) { - float v = (float)h_inp[i * K + j]; + float v = (float)h_inp.ptr[i * K + j]; sum_sq += v * v; } - h_rms[i] = sqrtf(sum_sq / (float)K + 1e-5f); + h_rms.ptr[i] = sqrtf(sum_sq / (float)K + 1e-5f); } - - floatX *d_inp, *d_weight; - float *d_grad, *d_rms, *d_dW, *d_dx; - CHECK_CUDA(cudaMalloc(&d_inp, M * K * sizeof(floatX))); - CHECK_CUDA(cudaMalloc(&d_weight, N * K * sizeof(floatX))); - CHECK_CUDA(cudaMalloc(&d_grad, M * N * sizeof(float))); - CHECK_CUDA(cudaMalloc(&d_rms, M * sizeof(float))); - CHECK_CUDA(cudaMalloc(&d_dW, N * K * sizeof(float))); - CHECK_CUDA(cudaMalloc(&d_dx, M * K * sizeof(float))); - - CHECK_CUDA(cudaMemcpy(d_inp, h_inp, M * K * sizeof(floatX), cudaMemcpyHostToDevice)); - CHECK_CUDA(cudaMemcpy(d_weight, h_weight, N * K * sizeof(floatX), cudaMemcpyHostToDevice)); - CHECK_CUDA(cudaMemcpy(d_grad, h_grad, M * N * sizeof(float), cudaMemcpyHostToDevice)); - CHECK_CUDA(cudaMemcpy(d_rms, h_rms, M * sizeof(float), cudaMemcpyHostToDevice)); - + d_rms.upload(h_rms); FusedRMSNormMatmulBackward backward; backward.init(M, N, K); - double flops = 4.0 * (double)M * (double)N * (double)K; - - size_t bytes_read = M * K * sizeof(floatX) + N * K * sizeof(floatX) + - M * N * sizeof(float) + M * sizeof(float); - size_t bytes_write = N * K * sizeof(float) + M * K * sizeof(float); - size_t total_bytes = bytes_read + bytes_write; + size_t total_bytes = M * K * sizeof(floatX) + N * K * sizeof(floatX) + + M * N * sizeof(float) + M * sizeof(float) + N * K * sizeof(float) + + M * K * sizeof(float); - BenchTimer timer; - float total_time = 0.0f; - - for (int i = 0; i < bench_runs; i++) { - flusher.flush(); - CHECK_CUDA(cudaMemset(d_dW, 0, N * K * sizeof(float))); + double flops = 4.0 * (double)M * (double)N * (double)K; - timer.record_start(); - backward.backward(d_dW, d_dx, d_grad, d_inp, d_weight, d_rms); - timer.record_stop(); - total_time += timer.elapsed_ms(); - } + float avg_time_ms = + bench_kernel([&]() { backward.backward(d_dW, d_dx, d_grad, d_inp, d_weight, d_rms); }, + bench_runs, flusher, [&]() { d_dW.zero(); }); - float avg_time_ms = total_time / bench_runs; + float time_us = avg_time_ms * 1000.0f; float tflops = (flops / 1e12f) / (avg_time_ms / 1e3f); float bw = (total_bytes / 1e9f) / (avg_time_ms / 1e3f); - printf("%8d %8d %8d %12.2f %12.2f %12.2f\n", M, N, K, avg_time_ms * 1000.0f, tflops, bw); + printf("%8d %8d %8d %12.2f %12.2f %12.2f\n", M, N, K, time_us, tflops, bw); backward.destroy(); - cudaFree(d_inp); - cudaFree(d_weight); - cudaFree(d_grad); - cudaFree(d_rms); - cudaFree(d_dW); - cudaFree(d_dx); - free(h_inp); - free(h_weight); - free(h_grad); - free(h_rms); } + return 0; } diff --git a/src/csrc/benchmarks/bench_mhc_layer.cu b/src/csrc/benchmarks/bench_mhc_layer.cu index 5855c1e..fc5ffa8 100644 --- a/src/csrc/benchmarks/bench_mhc_layer.cu +++ b/src/csrc/benchmarks/bench_mhc_layer.cu @@ -1,228 +1,150 @@ +// AUTO-GENERATED by benchgen.py — do not edit manually +// Source: src/csrc/kernels/mhc_layer.cuh +// Regenerate: make benchgen + #include -#include #include -#include -#include "mhc_layer.cuh" -#include "mhc_types.h" -#include "utils.cuh" - -using namespace mhc; - -void run_static_benchmark(int B, int C, int n, int bench_runs, L2Flusher& flusher, bool use_tc) { - float* h_x_expanded = (float*)malloc(B * n * C * sizeof(float)); - floatX* h_rmsnorm_weight = (floatX*)malloc(C * sizeof(floatX)); - float* h_H_pre = (float*)malloc(n * sizeof(float)); - float* h_H_post = (float*)malloc(n * sizeof(float)); - float* h_H_res = (float*)malloc(n * n * sizeof(float)); - - srand(42); - for (int i = 0; i < B * n * C; i++) { - h_x_expanded[i] = (float)rand() / RAND_MAX * 2.0f - 1.0f; - } - for (int i = 0; i < C; i++) { - h_rmsnorm_weight[i] = (floatX)((float)rand() / RAND_MAX * 0.5f + 0.75f); - } - for (int i = 0; i < n; i++) { - h_H_pre[i] = 0.0f; - h_H_post[i] = 0.0f; - } - for (int i = 0; i < n * n; i++) { - h_H_res[i] = 0.01f * ((float)rand() / RAND_MAX * 2.0f - 1.0f); - } - float* d_x_expanded; - CHECK_CUDA(cudaMalloc(&d_x_expanded, B * n * C * sizeof(float))); - CHECK_CUDA( - cudaMemcpy(d_x_expanded, h_x_expanded, B * n * C * sizeof(float), cudaMemcpyHostToDevice)); +#include "../kernels/mhc_layer.cuh" +#include "../include/bench_harness.cuh" - MHCLayerConfig cfg; - cfg.batch_size = B; - cfg.hidden_dim = C; - cfg.expansion_rate = n; - cfg.sinkhorn_iters = 20; - cfg.eps = 1e-5f; - cfg.use_pdl = true; - cfg.use_dynamic_h = false; - - MHCLayer layer; - layer.init(cfg); - layer.use_tc_mix = use_tc; - layer.set_weights(h_rmsnorm_weight, h_H_pre, h_H_post, h_H_res); - layer.sync(); +using namespace mhc; - layer.forward_device(d_x_expanded); - layer.sync(); +int main() { + const int bench_runs = 50; - size_t bytes_io = (size_t)B * n * C * sizeof(float) * 3; + L2Flusher flusher; - BenchTimer timer; - float total_time = 0.0f; + int configs[][3] = { + {64, 1280, 4}, {128, 1280, 4}, {256, 1280, 4}, {320, 1280, 4}, + {64, 1920, 4}, {128, 1920, 4}, {64, 2560, 4}, {128, 2560, 4}, + }; + int num_configs = sizeof(configs) / sizeof(configs[0]); - for (int i = 0; i < bench_runs; i++) { - flusher.flush(); + printf("Mhc Layer Benchmark\n"); + printf("======================================================================\n"); + printf("\nStatic H Path\n"); + printf("%8s %8s %8s %12s %12s %14s\n", "B", "C", "n", "Time (us)", "Bandwidth (GB/s)", + "Samples/sec"); + printf("-------------------------------------------------------------------\n"); - timer.record_start(); + for (int c = 0; c < num_configs; c++) { + int B = configs[c][0]; + int C = configs[c][1]; + int n = configs[c][2]; + + + HostMem h_x_expanded(B * n * C); + HostMem h_rmsnorm_weight(C); + HostMem h_H_pre(n); + HostMem h_H_post(n); + HostMem h_H_res(n * n); + fill_random(h_x_expanded, B * n * C); + fill_random_bf16(h_rmsnorm_weight, C, 0.75f, 1.25f); + for (int i = 0; i < n; i++) { + h_H_pre.ptr[i] = 0.0f; + h_H_post.ptr[i] = 0.0f; + } + srand(42); + for (int i = 0; i < n * n; i++) { + h_H_res.ptr[i] = 0.01f * ((float)rand() / RAND_MAX * 2.0f - 1.0f); + } + DeviceMem d_x_expanded(B * n * C); + d_x_expanded.upload(h_x_expanded); + MHCLayerConfig cfg; + cfg.batch_size = B; + cfg.hidden_dim = C; + cfg.expansion_rate = n; + cfg.sinkhorn_iters = 20; + cfg.eps = 1e-5f; + cfg.use_pdl = true; + cfg.use_dynamic_h = false; + MHCLayer layer; + layer.init(cfg); + layer.set_weights(h_rmsnorm_weight, h_H_pre, h_H_post, h_H_res); + layer.sync(); layer.forward_device(d_x_expanded); - timer.record_stop(); - total_time += timer.elapsed_ms(); - } + layer.sync(); - float avg_time_ms = total_time / bench_runs; - float throughput = B / (avg_time_ms / 1000.0f); - float bw = (bytes_io / 1e9f) / (avg_time_ms / 1e3f); + size_t total_bytes = (size_t)B * n * C * sizeof(float) * 3; - printf("%6d %6d %4d %8s %6s %12.2f %14.0f %14.0f\n", B, C, n, "static", - use_tc ? "TC" : "CUDA CORE", avg_time_ms * 1000.0f, throughput, bw); + float avg_time_ms = + bench_kernel([&]() { layer.forward_device(d_x_expanded); }, bench_runs, flusher); - layer.destroy(); - cudaFree(d_x_expanded); - free(h_x_expanded); - free(h_rmsnorm_weight); - free(h_H_pre); - free(h_H_post); - free(h_H_res); -} + float time_us = avg_time_ms * 1000.0f; + float bw = (total_bytes / 1e9f) / (avg_time_ms / 1e3f); + float samples_per_sec = B / (avg_time_ms / 1000.0f); -void run_dynamic_benchmark(int B, int C, int n, int bench_runs, L2Flusher& flusher) { - float* h_x_expanded = (float*)malloc(B * n * C * sizeof(float)); - floatX* h_rmsnorm_weight = (floatX*)malloc(C * sizeof(floatX)); + printf("%8d %8d %8d %12.2f %12.0f %14.0f\n", B, C, n, time_us, bw, samples_per_sec); - int nC = n * C; - int total_H_dim = n + n + n * n; - floatX* h_phi = (floatX*)malloc(total_H_dim * nC * sizeof(floatX)); - float* h_b_pre = (float*)malloc(n * sizeof(float)); - float* h_b_post = (float*)malloc(n * sizeof(float)); - float* h_b_res = (float*)malloc(n * n * sizeof(float)); - - srand(42); - for (int i = 0; i < B * n * C; i++) { - h_x_expanded[i] = (float)rand() / RAND_MAX * 2.0f - 1.0f; - } - for (int i = 0; i < C; i++) { - h_rmsnorm_weight[i] = (floatX)((float)rand() / RAND_MAX * 0.5f + 0.75f); - } - for (int i = 0; i < total_H_dim * nC; i++) { - h_phi[i] = (floatX)((float)rand() / RAND_MAX * 0.1f - 0.05f); + layer.destroy(); } - for (int i = 0; i < n; i++) { - h_b_pre[i] = 0.0f; - h_b_post[i] = 0.0f; - } - for (int i = 0; i < n * n; i++) { - h_b_res[i] = 0.01f * ((float)rand() / RAND_MAX * 2.0f - 1.0f); - } - - float* d_x_expanded; - CHECK_CUDA(cudaMalloc(&d_x_expanded, B * n * C * sizeof(float))); - CHECK_CUDA( - cudaMemcpy(d_x_expanded, h_x_expanded, B * n * C * sizeof(float), cudaMemcpyHostToDevice)); - MHCLayerConfig cfg; - cfg.batch_size = B; - cfg.hidden_dim = C; - cfg.expansion_rate = n; - cfg.sinkhorn_iters = 20; - cfg.eps = 1e-5f; - cfg.use_pdl = true; - cfg.use_dynamic_h = true; - MHCLayer layer; - layer.init(cfg); + printf("\nDynamic H Path\n"); + printf("%8s %8s %8s %12s %12s %14s\n", "B", "C", "n", "Time (us)", "Bandwidth (GB/s)", + "Samples/sec"); + printf("-------------------------------------------------------------------\n"); - floatX* h_phi_pre = h_phi; - floatX* h_phi_post = h_phi + n * nC; - floatX* h_phi_res = h_phi + 2 * n * nC; - - layer.set_weights_dynamic(h_rmsnorm_weight, h_phi_pre, h_phi_post, h_phi_res, h_b_pre, h_b_post, - h_b_res, 0.01f, 0.01f, 0.01f); - layer.sync(); - - layer.forward_device(d_x_expanded); - layer.sync(); - - size_t bytes_io = (size_t)B * n * C * sizeof(float) * 3; - - BenchTimer timer; - float total_time = 0.0f; - - for (int i = 0; i < bench_runs; i++) { - flusher.flush(); - - timer.record_start(); + for (int c = 0; c < num_configs; c++) { + int B = configs[c][0]; + int C = configs[c][1]; + int n = configs[c][2]; + + + int nC = n * C; + int total_H_dim = n + n + n * n; + HostMem h_x_expanded(B * n * C); + HostMem h_rmsnorm_weight(C); + HostMem h_phi(total_H_dim * nC); + HostMem h_b_pre(n); + HostMem h_b_post(n); + HostMem h_b_res(n * n); + fill_random(h_x_expanded, B * n * C); + fill_random_bf16(h_rmsnorm_weight, C, 0.75f, 1.25f); + fill_random_bf16(h_phi, total_H_dim * nC, -0.05f, 0.05f, 43); + for (int i = 0; i < n; i++) { + h_b_pre.ptr[i] = 0.0f; + h_b_post.ptr[i] = 0.0f; + } + srand(42); + for (int i = 0; i < n * n; i++) { + h_b_res.ptr[i] = 0.01f * ((float)rand() / RAND_MAX * 2.0f - 1.0f); + } + DeviceMem d_x_expanded(B * n * C); + d_x_expanded.upload(h_x_expanded); + MHCLayerConfig cfg; + cfg.batch_size = B; + cfg.hidden_dim = C; + cfg.expansion_rate = n; + cfg.sinkhorn_iters = 20; + cfg.eps = 1e-5f; + cfg.use_pdl = true; + cfg.use_dynamic_h = true; + MHCLayer layer; + layer.init(cfg); + floatX* phi_base = h_phi; + layer.set_weights_dynamic(h_rmsnorm_weight, phi_base, phi_base + n * nC, + phi_base + 2 * n * nC, h_b_pre, h_b_post, h_b_res, 0.01f, 0.01f, + 0.01f); + layer.sync(); layer.forward_device(d_x_expanded); - timer.record_stop(); - total_time += timer.elapsed_ms(); - } - - float avg_time_ms = total_time / bench_runs; - float throughput = B / (avg_time_ms / 1000.0f); - float bw = (bytes_io / 1e9f) / (avg_time_ms / 1e3f); - - printf("%6d %6d %4d %8s %6s %12.2f %14.0f %14.0f\n", B, C, n, "dynamic", "CUDA", - avg_time_ms * 1000.0f, throughput, bw); - - layer.destroy(); - cudaFree(d_x_expanded); - free(h_x_expanded); - free(h_rmsnorm_weight); - free(h_phi); - free(h_b_pre); - free(h_b_post); - free(h_b_res); -} - -int main() { - const int bench_runs = 50; + layer.sync(); - L2Flusher flusher; + size_t total_bytes = (size_t)B * n * C * sizeof(float) * 3; - struct Config { - int B; - int C; - int n; - }; + float avg_time_ms = + bench_kernel([&]() { layer.forward_device(d_x_expanded); }, bench_runs, flusher); - Config configs[] = { - {64, 1280, 4}, {128, 1280, 4}, {256, 1280, 4}, {320, 1280, 4}, - {64, 1920, 4}, {128, 1920, 4}, {64, 2560, 4}, {128, 2560, 4}, - {32, 1280, 32}, {64, 1280, 32}, {128, 1280, 32}, - }; - int num_configs = sizeof(configs) / sizeof(configs[0]); + float time_us = avg_time_ms * 1000.0f; + float bw = (total_bytes / 1e9f) / (avg_time_ms / 1e3f); + float samples_per_sec = B / (avg_time_ms / 1000.0f); - printf("MHC Layer End-to-End Benchmark\n"); - printf("==========================================================\n"); - printf("Pipeline: Aggregate(H_pre) -> RMSNorm -> Distribute(H_post) + Mix(M)\n"); - printf("Static H: H coefficients shared across batch (Sinkhorn on H_res)\n"); - printf("Dynamic H: H coefficients computed per-sample via RMSNorm + MatMul + activations\n"); - printf("Input shape: [B, n, C]\n"); - printf("PDL path: %s\n\n", -#ifdef MHC_ENABLE_PDL - "Enabled" -#else - "Disabled" -#endif - ); - - printf("%6s %6s %4s %8s %6s %12s %14s %14s\n", "Batch", "Hidden", "n", "Mode", "Implementation", - "Time (us)", "Samples/sec", "Bandwidth (GB/s)"); - printf( - "--------------------------------------------------------------------------------------\n"); + printf("%8d %8d %8d %12.2f %12.0f %14.0f\n", B, C, n, time_us, bw, samples_per_sec); - for (int c = 0; c < num_configs; c++) { - int B = configs[c].B; - int C = configs[c].C; - int n = configs[c].n; - - run_static_benchmark(B, C, n, bench_runs, flusher, false); - - if (n >= 32) { - run_static_benchmark(B, C, n, bench_runs, flusher, true); - } - - run_dynamic_benchmark(B, C, n, bench_runs, flusher); - - printf("\n"); + layer.destroy(); } + return 0; } diff --git a/src/csrc/benchmarks/bench_rmsnorm.cu b/src/csrc/benchmarks/bench_rmsnorm.cu index 785cbc2..b44c140 100644 --- a/src/csrc/benchmarks/bench_rmsnorm.cu +++ b/src/csrc/benchmarks/bench_rmsnorm.cu @@ -1,99 +1,15 @@ +// AUTO-GENERATED by benchgen.py — do not edit manually +// Source: src/csrc/kernels/rmsnorm.cuh +// Regenerate: make benchgen + #include -#include -#include -#include -#include #include -#include "../include/mhc_types.h" -#include "../include/utils.cuh" +#include "../kernels/rmsnorm.cuh" +#include "../include/bench_harness.cuh" -namespace cg = cooperative_groups; using namespace mhc; -template -__global__ void rmsnorm_profiled_kernel(floatX* __restrict__ out, const floatX* __restrict__ inp, - const floatX* __restrict__ weight, int N, int C, float eps, - int64_t* profiler_buf, int max_entries) { - cg::thread_block block = cg::this_thread_block(); - cg::thread_block_tile<32> warp = cg::tiled_partition<32>(block); - - int idx = blockIdx.x; - if (idx >= N) - return; - - DeviceProfiler profiler; - if constexpr (DO_PROFILE) { - if (threadIdx.x == 0) { - profiler.init(max_entries, profiler_buf, idx); - profiler.start(TagLoad); - } - } - - const floatX* x = inp + idx * C; - floatX* o = out + idx * C; - - extern __shared__ float shared[]; - float* s_sum_sq = shared; - - float thread_sum_sq = 0.0f; - for (int i = threadIdx.x; i < C; i += BLOCK_SIZE) { - float val = (float)x[i]; - thread_sum_sq += val * val; - } - - if constexpr (DO_PROFILE) { - if (threadIdx.x == 0) { - profiler.stop(); - profiler.start(TagReduce); - } - } - - float warp_sum = cg::reduce(warp, thread_sum_sq, cg::plus()); - - int warp_id = threadIdx.x / 32; - int lane_id = threadIdx.x % 32; - int num_warps = BLOCK_SIZE / 32; - - if (lane_id == 0) { - s_sum_sq[warp_id] = warp_sum; - } - __syncthreads(); - - if (warp_id == 0) { - float val = (lane_id < num_warps) ? s_sum_sq[lane_id] : 0.0f; - float block_sum = cg::reduce(warp, val, cg::plus()); - - if (lane_id == 0) { - float rms = rsqrtf(block_sum / (float)C + eps); - s_sum_sq[0] = rms; - } - } - __syncthreads(); - - if constexpr (DO_PROFILE) { - if (threadIdx.x == 0) { - profiler.stop(); - profiler.start(TagStore); - } - } - - float rms_inv = s_sum_sq[0]; - - for (int i = threadIdx.x; i < C; i += BLOCK_SIZE) { - float val = (float)x[i]; - float w = (float)weight[i]; - o[i] = (floatX)(val * rms_inv * w); - } - - if constexpr (DO_PROFILE) { - if (threadIdx.x == 0) { - profiler.stop(); - profiler.flush(); - } - } -} - int main() { const int bench_runs = 100; const float eps = 1e-5f; @@ -108,107 +24,59 @@ int main() { int num_configs = sizeof(configs) / sizeof(configs[0]); printf("RMSNorm Benchmark\n"); - printf("====================================\n"); + printf("======================================================================\n"); printf("%8s %8s %12s %12s\n", "N", "C", "Time (us)", "Bandwidth (GB/s)"); - printf("---------------------------------------------------\n"); + printf("-------------------------------------------\n"); for (int c = 0; c < num_configs; c++) { int N = configs[c][0]; int C = configs[c][1]; - floatX* h_inp = (floatX*)malloc(N * C * sizeof(floatX)); - floatX* h_weight = (floatX*)malloc(C * sizeof(floatX)); - - srand(42); - for (int i = 0; i < N * C; i++) { - h_inp[i] = (floatX)((float)rand() / RAND_MAX * 2.0f - 1.0f); - } - for (int i = 0; i < C; i++) { - h_weight[i] = (floatX)((float)rand() / RAND_MAX * 0.5f + 0.75f); - } + HostMem h_inp(N * C); + HostMem h_weight(C); + fill_random(h_inp, N * C, -1.0f, 1.0f); + fill_random_bf16(h_weight, C, 0.75f, 1.25f); - floatX *d_inp, *d_weight, *d_out; - CHECK_CUDA(cudaMalloc(&d_inp, N * C * sizeof(floatX))); - CHECK_CUDA(cudaMalloc(&d_weight, C * sizeof(floatX))); - CHECK_CUDA(cudaMalloc(&d_out, N * C * sizeof(floatX))); + DeviceMem d_inp(N * C); + DeviceMem d_weight(C); + DeviceMem d_out(N * C); + d_inp.upload(h_inp); + d_weight.upload(h_weight); - CHECK_CUDA(cudaMemcpy(d_inp, h_inp, N * C * sizeof(floatX), cudaMemcpyHostToDevice)); - CHECK_CUDA(cudaMemcpy(d_weight, h_weight, C * sizeof(floatX), cudaMemcpyHostToDevice)); + size_t total_bytes = (size_t)N * C * sizeof(float) + (size_t)C * sizeof(floatX) + + (size_t)N * C * sizeof(float); - size_t bytes_read = (size_t)N * C * sizeof(floatX) + (size_t)C * sizeof(floatX); - size_t bytes_written = (size_t)N * C * sizeof(floatX); - size_t total_bytes = bytes_read + bytes_written; + float avg_time_ms = bench_kernel( + [&]() { rmsnorm_forward(d_out, d_inp, d_weight, N, C, 1e-5f); }, bench_runs, flusher); - constexpr int BLOCK_SIZE = 512; - int num_warps = BLOCK_SIZE / 32; - size_t shared_mem = num_warps * sizeof(float); - - BenchTimer timer; - float total_time = 0.0f; - - for (int i = 0; i < bench_runs; i++) { - flusher.flush(); - - timer.record_start(); - rmsnorm_profiled_kernel - <<>>(d_out, d_inp, d_weight, N, C, eps, nullptr, 0); - timer.record_stop(); - total_time += timer.elapsed_ms(); - } - - float avg_time_ms = total_time / bench_runs; + float time_us = avg_time_ms * 1000.0f; float bw = (total_bytes / 1e9f) / (avg_time_ms / 1e3f); - printf("%8d %8d %12.2f %12.0f\n", N, C, avg_time_ms * 1000.0f, bw); - - cudaFree(d_inp); - cudaFree(d_weight); - cudaFree(d_out); - free(h_inp); - free(h_weight); + printf("%8d %8d %12.2f %12.0f\n", N, C, time_us, bw); } - printf("\n--- Step Breakdown (N=1024, C=4096) ---\n"); + printf("\n--- Phase Breakdown (N=1024, C=4096) ---\n"); { - int N = 1024, C = 4096; - - floatX* h_inp = (floatX*)malloc(N * C * sizeof(floatX)); - floatX* h_weight = (floatX*)malloc(C * sizeof(floatX)); - - srand(42); - for (int i = 0; i < N * C; i++) { - h_inp[i] = (floatX)((float)rand() / RAND_MAX * 2.0f - 1.0f); - } - for (int i = 0; i < C; i++) { - h_weight[i] = (floatX)((float)rand() / RAND_MAX * 0.5f + 0.75f); - } - - floatX *d_inp, *d_weight, *d_out; - CHECK_CUDA(cudaMalloc(&d_inp, N * C * sizeof(floatX))); - CHECK_CUDA(cudaMalloc(&d_weight, C * sizeof(floatX))); - CHECK_CUDA(cudaMalloc(&d_out, N * C * sizeof(floatX))); - - CHECK_CUDA(cudaMemcpy(d_inp, h_inp, N * C * sizeof(floatX), cudaMemcpyHostToDevice)); - CHECK_CUDA(cudaMemcpy(d_weight, h_weight, C * sizeof(floatX), cudaMemcpyHostToDevice)); - - constexpr int BLOCK_SIZE = 512; - int num_warps = BLOCK_SIZE / 32; - size_t shared_mem = num_warps * sizeof(float); - - HostProfiler profiler(N, max_entries); - - flusher.flush(); - rmsnorm_profiled_kernel<<>>( - d_out, d_inp, d_weight, N, C, eps, profiler.device_ptr(), max_entries); - CHECK_CUDA(cudaDeviceSynchronize()); - - profiler.print_summary(); - - cudaFree(d_inp); - cudaFree(d_weight); - cudaFree(d_out); - free(h_inp); - free(h_weight); + int N = 1024; + int C = 4096; + + HostMem h_inp(N * C); + HostMem h_weight(C); + fill_random(h_inp, N * C, -1.0f, 1.0f); + fill_random_bf16(h_weight, C, 0.75f, 1.25f); + + DeviceMem d_inp(N * C); + DeviceMem d_weight(C); + DeviceMem d_out(N * C); + d_inp.upload(h_inp); + d_weight.upload(h_weight); + + profile_kernel( + [&](int64_t* profiler_buf, int max_ent) { + rmsnorm_forward_profiled(d_out, d_inp, d_weight, N, C, 1e-5f, profiler_buf, + max_ent); + }, + N, max_entries, flusher); } return 0; diff --git a/src/csrc/benchmarks/bench_rmsnorm_backward.cu b/src/csrc/benchmarks/bench_rmsnorm_backward.cu index 915a22c..2f37e75 100644 --- a/src/csrc/benchmarks/bench_rmsnorm_backward.cu +++ b/src/csrc/benchmarks/bench_rmsnorm_backward.cu @@ -1,16 +1,18 @@ +// AUTO-GENERATED by benchgen.py — do not edit manually +// Source: src/csrc/kernels/rmsnorm.cuh +// Regenerate: make benchgen + #include -#include -#include #include -#include "../include/mhc_types.h" -#include "../include/utils.cuh" #include "../kernels/rmsnorm.cuh" +#include "../include/bench_harness.cuh" using namespace mhc; int main() { const int bench_runs = 100; + const int max_entries = 8; L2Flusher flusher; @@ -21,85 +23,95 @@ int main() { int num_configs = sizeof(configs) / sizeof(configs[0]); printf("RMSNorm Backward Benchmark\n"); - printf("====================================\n"); + printf("======================================================================\n"); printf("%8s %8s %12s %12s\n", "N", "C", "Time (us)", "Bandwidth (GB/s)"); - printf("---------------------------------------------------\n"); + printf("-------------------------------------------\n"); for (int c = 0; c < num_configs; c++) { int N = configs[c][0]; int C = configs[c][1]; - float* h_inp = (float*)malloc(N * C * sizeof(float)); - floatX* h_weight = (floatX*)malloc(C * sizeof(floatX)); - float* h_grad = (float*)malloc(N * C * sizeof(float)); - float* h_rms = (float*)malloc(N * sizeof(float)); + HostMem h_grad(N * C); + HostMem h_inp(N * C); + HostMem h_weight(C); + HostMem h_rms(N); + fill_random(h_grad, N * C, -1.0f, 1.0f, 43); + fill_random(h_inp, N * C, -1.0f, 1.0f); + fill_random_bf16(h_weight, C, 0.75f, 1.25f); + + DeviceMem d_grad(N * C); + DeviceMem d_inp(N * C); + DeviceMem d_weight(C); + DeviceMem d_rms(N); + DeviceMem d_d_inp(N * C); + DeviceMem d_d_weight(C); + d_grad.upload(h_grad); + d_inp.upload(h_inp); + d_weight.upload(h_weight); - srand(42); - for (int i = 0; i < N * C; i++) { - h_inp[i] = (float)rand() / RAND_MAX * 2.0f - 1.0f; - h_grad[i] = (float)rand() / RAND_MAX * 2.0f - 1.0f; - } - for (int i = 0; i < C; i++) { - h_weight[i] = (floatX)((float)rand() / RAND_MAX * 0.5f + 0.75f); - } for (int i = 0; i < N; i++) { float sum_sq = 0.0f; for (int j = 0; j < C; j++) { - float v = h_inp[i * C + j]; + float v = h_inp.ptr[i * C + j]; sum_sq += v * v; } - h_rms[i] = sqrtf(sum_sq / (float)C + 1e-5f); + h_rms.ptr[i] = sqrtf(sum_sq / (float)C + 1e-5f); } + d_rms.upload(h_rms); - float* d_inp; - floatX* d_weight; - float *d_grad, *d_rms, *d_d_inp, *d_d_weight; - - CHECK_CUDA(cudaMalloc(&d_inp, N * C * sizeof(float))); - CHECK_CUDA(cudaMalloc(&d_weight, C * sizeof(floatX))); - CHECK_CUDA(cudaMalloc(&d_grad, N * C * sizeof(float))); - CHECK_CUDA(cudaMalloc(&d_rms, N * sizeof(float))); - CHECK_CUDA(cudaMalloc(&d_d_inp, N * C * sizeof(float))); - CHECK_CUDA(cudaMalloc(&d_d_weight, C * sizeof(float))); - - CHECK_CUDA(cudaMemcpy(d_inp, h_inp, N * C * sizeof(float), cudaMemcpyHostToDevice)); - CHECK_CUDA(cudaMemcpy(d_weight, h_weight, C * sizeof(floatX), cudaMemcpyHostToDevice)); - CHECK_CUDA(cudaMemcpy(d_grad, h_grad, N * C * sizeof(float), cudaMemcpyHostToDevice)); - CHECK_CUDA(cudaMemcpy(d_rms, h_rms, N * sizeof(float), cudaMemcpyHostToDevice)); - - size_t bytes_read = (size_t)N * C * sizeof(float) + (size_t)C * sizeof(floatX) + - (size_t)N * C * sizeof(float) + (size_t)N * sizeof(float); - size_t bytes_write = (size_t)N * C * sizeof(float) + (size_t)C * sizeof(float); - size_t total_bytes = bytes_read + bytes_write; - - BenchTimer timer; - float total_time = 0.0f; - - for (int i = 0; i < bench_runs; i++) { - flusher.flush(); - CHECK_CUDA(cudaMemset(d_d_weight, 0, C * sizeof(float))); - - timer.record_start(); - rmsnorm_backward(d_d_inp, d_d_weight, d_grad, d_inp, d_weight, d_rms, N, C); - timer.record_stop(); - total_time += timer.elapsed_ms(); - } + size_t total_bytes = (size_t)N * C * sizeof(float) + (size_t)C * sizeof(floatX) + + (size_t)N * C * sizeof(float) + (size_t)N * sizeof(float) + + (size_t)N * C * sizeof(float) + (size_t)C * sizeof(float); - float avg_time_ms = total_time / bench_runs; + float avg_time_ms = bench_kernel( + [&]() { rmsnorm_backward(d_d_inp, d_d_weight, d_grad, d_inp, d_weight, d_rms, N, C); }, + bench_runs, flusher, [&]() { d_d_weight.zero(); }); + + float time_us = avg_time_ms * 1000.0f; float bw = (total_bytes / 1e9f) / (avg_time_ms / 1e3f); - printf("%8d %8d %12.2f %12.0f\n", N, C, avg_time_ms * 1000.0f, bw); - - cudaFree(d_inp); - cudaFree(d_weight); - cudaFree(d_grad); - cudaFree(d_rms); - cudaFree(d_d_inp); - cudaFree(d_d_weight); - free(h_inp); - free(h_weight); - free(h_grad); - free(h_rms); + printf("%8d %8d %12.2f %12.0f\n", N, C, time_us, bw); + } + + printf("\n--- Phase Breakdown (N=1024, C=4096) ---\n"); + { + int N = 1024; + int C = 4096; + + HostMem h_grad(N * C); + HostMem h_inp(N * C); + HostMem h_weight(C); + HostMem h_rms(N); + fill_random(h_grad, N * C, -1.0f, 1.0f, 43); + fill_random(h_inp, N * C, -1.0f, 1.0f); + fill_random_bf16(h_weight, C, 0.75f, 1.25f); + + DeviceMem d_grad(N * C); + DeviceMem d_inp(N * C); + DeviceMem d_weight(C); + DeviceMem d_rms(N); + DeviceMem d_d_inp(N * C); + DeviceMem d_d_weight(C); + d_grad.upload(h_grad); + d_inp.upload(h_inp); + d_weight.upload(h_weight); + + for (int i = 0; i < N; i++) { + float sum_sq = 0.0f; + for (int j = 0; j < C; j++) { + float v = h_inp.ptr[i * C + j]; + sum_sq += v * v; + } + h_rms.ptr[i] = sqrtf(sum_sq / (float)C + 1e-5f); + } + d_rms.upload(h_rms); + + profile_kernel( + [&](int64_t* profiler_buf, int max_ent) { + rmsnorm_backward_profiled(d_d_inp, d_d_weight, d_grad, d_inp, d_weight, d_rms, N, C, + profiler_buf, max_ent); + }, + N, max_entries, flusher); } return 0; diff --git a/src/csrc/benchmarks/bench_sinkhorn_knopp.cu b/src/csrc/benchmarks/bench_sinkhorn_knopp.cu index 2364b98..67f186e 100644 --- a/src/csrc/benchmarks/bench_sinkhorn_knopp.cu +++ b/src/csrc/benchmarks/bench_sinkhorn_knopp.cu @@ -1,106 +1,15 @@ +// AUTO-GENERATED by benchgen.py — do not edit manually +// Source: src/csrc/kernels/sinkhorn_knopp.cuh +// Regenerate: make benchgen + #include -#include #include -#include "../include/mhc_types.h" -#include "../include/utils.cuh" #include "../kernels/sinkhorn_knopp.cuh" +#include "../include/bench_harness.cuh" using namespace mhc; -template -__global__ void -sinkhorn_knopp_profiled_kernel(float* __restrict__ out, const float* __restrict__ inp, int M, int N, - int num_iters, float eps, int64_t* profiler_buf, int max_entries) { - extern __shared__ float smem[]; - float* tile = smem; - float* row_sums = smem + MAX_DIM * MAX_DIM; - float* col_sums = row_sums + MAX_DIM; - - int total_elems = M * N; - - DeviceProfiler profiler; - if constexpr (DO_PROFILE) { - if (threadIdx.x == 0) { - profiler.init(max_entries, profiler_buf, blockIdx.x); - profiler.start(TagLoad); - } - } - - for (int i = threadIdx.x; i < total_elems; i += BLOCK_SIZE) { - tile[i] = inp[i]; - } - __syncthreads(); - - if constexpr (DO_PROFILE) { - if (threadIdx.x == 0) { - profiler.stop(); - } - } - - for (int iter = 0; iter < num_iters; iter++) { - if constexpr (DO_PROFILE) { - if (threadIdx.x == 0 && iter == 0) { - profiler.start(TagCompute); - } - } - - for (int r = threadIdx.x; r < M; r += BLOCK_SIZE) { - float sum = 0.0f; - for (int c = 0; c < N; c++) { - sum += tile[r * N + c]; - } - row_sums[r] = sum; - } - __syncthreads(); - - for (int i = threadIdx.x; i < total_elems; i += BLOCK_SIZE) { - int r = i / N; - float row_sum = row_sums[r]; - if (row_sum > eps) { - tile[i] /= row_sum; - } - } - __syncthreads(); - - for (int c = threadIdx.x; c < N; c += BLOCK_SIZE) { - float sum = 0.0f; - for (int r = 0; r < M; r++) { - sum += tile[r * N + c]; - } - col_sums[c] = sum; - } - __syncthreads(); - - for (int i = threadIdx.x; i < total_elems; i += BLOCK_SIZE) { - int c = i % N; - float col_sum = col_sums[c]; - if (col_sum > eps) { - tile[i] /= col_sum; - } - } - __syncthreads(); - } - - if constexpr (DO_PROFILE) { - if (threadIdx.x == 0) { - profiler.stop(); - profiler.start(TagStore); - } - } - - for (int i = threadIdx.x; i < total_elems; i += BLOCK_SIZE) { - out[i] = tile[i]; - } - - if constexpr (DO_PROFILE) { - if (threadIdx.x == 0) { - profiler.stop(); - profiler.flush(); - } - } -} - int main() { const int bench_runs = 100; const float eps = 1e-8f; @@ -108,163 +17,122 @@ int main() { L2Flusher flusher; - struct Config { - int M; - int N; - int iters; - }; - - Config configs[] = { + int configs[][3] = { {32, 32, 5}, {32, 32, 10}, {32, 32, 20}, {64, 64, 5}, {64, 64, 10}, {64, 64, 20}, {128, 128, 5}, {128, 128, 10}, {128, 128, 20}, }; int num_configs = sizeof(configs) / sizeof(configs[0]); - printf("Sinkhorn-Knopp Forward Benchmark\n"); - printf("=======================================================================\n"); - printf("%6s %6s %6s %12s %10s %10s %12s\n", "M", "N", "Iters", "Time (us)", "us/iter", "GFLOPS", - "Bandwidth (GB/s)"); - printf("-----------------------------------------------------------------------\n"); + printf("Sinkhorn Knopp Benchmark\n"); + printf("======================================================================\n"); + printf("\nSinkhorn-Knopp Forward\n"); + printf("%8s %8s %6s %12s %10s %10s %12s\n", "M", "N", "num_iters", "Time (us)", "us/iter", + "GFLOPS", "Bandwidth (GB/s)"); + printf("------------------------------------------------------------------------\n"); for (int c = 0; c < num_configs; c++) { - int M = configs[c].M; - int N = configs[c].N; - int num_iters = configs[c].iters; - - float* h_inp = (float*)malloc(M * N * sizeof(float)); - - srand(42); - for (int i = 0; i < M * N; i++) { - h_inp[i] = (float)rand() / RAND_MAX + 0.1f; - } - - float *d_inp, *d_out; - CHECK_CUDA(cudaMalloc(&d_inp, M * N * sizeof(float))); - CHECK_CUDA(cudaMalloc(&d_out, M * N * sizeof(float))); - CHECK_CUDA(cudaMemcpy(d_inp, h_inp, M * N * sizeof(float), cudaMemcpyHostToDevice)); - - constexpr int BLOCK_SIZE = 256; - constexpr int MAX_DIM = 128; - size_t smem_size = - MAX_DIM * MAX_DIM * sizeof(float) + MAX_DIM * sizeof(float) + MAX_DIM * sizeof(float); + int M = configs[c][0]; + int N = configs[c][1]; + int num_iters = configs[c][2]; - auto kernel = sinkhorn_knopp_profiled_kernel; - cudaFuncSetAttribute(kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, smem_size); + HostMem h_inp(M * N); + fill_random(h_inp, M * N, 0.1f, 1.1f); - BenchTimer timer; - float total_time = 0.0f; + DeviceMem d_inp(M * N); + DeviceMem d_out(M * N); + d_inp.upload(h_inp); - for (int i = 0; i < bench_runs; i++) { - flusher.flush(); + size_t total_bytes = (size_t)M * N * sizeof(float) + (size_t)M * N * sizeof(float); - timer.record_start(); - kernel<<<1, BLOCK_SIZE, smem_size>>>(d_out, d_inp, M, N, num_iters, eps, nullptr, 0); - timer.record_stop(); - total_time += timer.elapsed_ms(); - } + float avg_time_ms = + bench_kernel([&]() { sinkhorn_knopp_forward(d_out, d_inp, M, N, num_iters, 1e-8f); }, + bench_runs, flusher); - float avg_time_ms = total_time / bench_runs; float time_us = avg_time_ms * 1000.0f; float time_per_iter_us = time_us / num_iters; + float gflops = ((4.0 * M * N * num_iters) / 1e9f) / (avg_time_ms / 1e3f); + float bw = (total_bytes / 1e9f) / (avg_time_ms / 1e3f); - double flops_per_iter = 4.0 * M * N; - double total_flops = flops_per_iter * num_iters; - float gflops = (total_flops / 1e9f) / (avg_time_ms / 1e3f); - - size_t bytes_io = 2 * M * N * sizeof(float); - float bw = (bytes_io / 1e9f) / (avg_time_ms / 1e3f); - - printf("%6d %6d %6d %12.2f %10.2f %10.2f %12.2f\n", M, N, num_iters, time_us, + printf("%8d %8d %6d %12.2f %10.2f %10.2f %12.2f\n", M, N, num_iters, time_us, time_per_iter_us, gflops, bw); - - cudaFree(d_inp); - cudaFree(d_out); - free(h_inp); } - printf("\n--- Phase Breakdown (64 x 64, 10 iterations) ---\n"); + printf("\n--- Phase Breakdown (M=64, N=64, num_iters=10) ---\n"); { - int M = 64, N = 64, num_iters = 10; - - float* h_inp = (float*)malloc(M * N * sizeof(float)); - srand(42); - for (int i = 0; i < M * N; i++) { - h_inp[i] = (float)rand() / RAND_MAX + 0.1f; - } - - float *d_inp, *d_out; - CHECK_CUDA(cudaMalloc(&d_inp, M * N * sizeof(float))); - CHECK_CUDA(cudaMalloc(&d_out, M * N * sizeof(float))); - CHECK_CUDA(cudaMemcpy(d_inp, h_inp, M * N * sizeof(float), cudaMemcpyHostToDevice)); - - constexpr int BLOCK_SIZE = 256; - constexpr int MAX_DIM = 128; - size_t smem_size = - MAX_DIM * MAX_DIM * sizeof(float) + MAX_DIM * sizeof(float) + MAX_DIM * sizeof(float); - - auto kernel = sinkhorn_knopp_profiled_kernel; - cudaFuncSetAttribute(kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, smem_size); - - HostProfiler profiler(1, max_entries); - - flusher.flush(); - kernel<<<1, BLOCK_SIZE, smem_size>>>(d_out, d_inp, M, N, num_iters, eps, - profiler.device_ptr(), max_entries); - CHECK_CUDA(cudaDeviceSynchronize()); - - profiler.print_summary(); - profiler.print_timeline(1); - - cudaFree(d_inp); - cudaFree(d_out); - free(h_inp); + int M = 64; + int N = 64; + int num_iters = 10; + + HostMem h_inp(M * N); + fill_random(h_inp, M * N, 0.1f, 1.1f); + + DeviceMem d_inp(M * N); + DeviceMem d_out(M * N); + d_inp.upload(h_inp); + + profile_kernel( + [&](int64_t* profiler_buf, int max_ent) { + sinkhorn_knopp_forward_profiled(d_out, d_inp, M, N, num_iters, 1e-8f, profiler_buf, + max_ent); + }, + 1, max_entries, flusher, true, 1); } - printf("\n--- Fused Exp Version (as used in MHCLayer) ---\n"); - printf("sinkhorn_knopp_forward_fused_exp applies exp() to input before iterations.\n"); - printf("%6s %6s %6s %12s %10s\n", "M", "N", "Iters", "Time (us)", "us/iter"); - printf("-----------------------------------------------------------------------\n"); + printf("\nFused Exp Version (as used in MHCLayer)\n"); + printf("%8s %8s %6s %12s %10s %12s\n", "M", "N", "num_iters", "Time (us)", "us/iter", + "Bandwidth (GB/s)"); + printf("-------------------------------------------------------------\n"); for (int c = 0; c < num_configs; c++) { - int M = configs[c].M; - int N = configs[c].N; - int num_iters = configs[c].iters; - - float* h_inp = (float*)malloc(M * N * sizeof(float)); - - srand(42); - for (int i = 0; i < M * N; i++) { - h_inp[i] = 0.01f * ((float)rand() / RAND_MAX * 2.0f - 1.0f); - } + int M = configs[c][0]; + int N = configs[c][1]; + int num_iters = configs[c][2]; - float *d_inp, *d_out, *d_exp; - CHECK_CUDA(cudaMalloc(&d_inp, M * N * sizeof(float))); - CHECK_CUDA(cudaMalloc(&d_out, M * N * sizeof(float))); - CHECK_CUDA(cudaMalloc(&d_exp, M * N * sizeof(float))); - CHECK_CUDA(cudaMemcpy(d_inp, h_inp, M * N * sizeof(float), cudaMemcpyHostToDevice)); + HostMem h_inp(M * N); + fill_random(h_inp, M * N, -0.01f, 0.01f); - BenchTimer timer; - float total_time = 0.0f; + DeviceMem d_inp(M * N); + DeviceMem d_H_res_exp(M * N); + DeviceMem d_out(M * N); + d_inp.upload(h_inp); - for (int i = 0; i < bench_runs; i++) { - flusher.flush(); + size_t total_bytes = (size_t)M * N * sizeof(float) + (size_t)M * N * sizeof(float) + + (size_t)M * N * sizeof(float); - timer.record_start(); - sinkhorn_knopp_forward_fused_exp(d_out, d_exp, d_inp, M, N, num_iters, eps); - timer.record_stop(); - total_time += timer.elapsed_ms(); - } + float avg_time_ms = bench_kernel( + [&]() { + sinkhorn_knopp_forward_fused_exp(d_out, d_H_res_exp, d_inp, M, N, num_iters, 1e-8f); + }, + bench_runs, flusher); - float avg_time_ms = total_time / bench_runs; float time_us = avg_time_ms * 1000.0f; float time_per_iter_us = time_us / num_iters; + float bw = (total_bytes / 1e9f) / (avg_time_ms / 1e3f); - printf("%6d %6d %6d %12.2f %10.2f\n", M, N, num_iters, time_us, time_per_iter_us); + printf("%8d %8d %6d %12.2f %10.2f %12.0f\n", M, N, num_iters, time_us, time_per_iter_us, + bw); + } - cudaFree(d_inp); - cudaFree(d_out); - cudaFree(d_exp); - free(h_inp); + printf("\n--- Phase Breakdown (M=64, N=64, num_iters=10) ---\n"); + { + int M = 64; + int N = 64; + int num_iters = 10; + + HostMem h_inp(M * N); + fill_random(h_inp, M * N, -0.01f, 0.01f); + + DeviceMem d_inp(M * N); + DeviceMem d_H_res_exp(M * N); + DeviceMem d_out(M * N); + d_inp.upload(h_inp); + + profile_kernel( + [&](int64_t* profiler_buf, int max_ent) { + sinkhorn_knopp_forward_fused_exp_profiled(d_out, d_H_res_exp, d_inp, M, N, + num_iters, 1e-8f, profiler_buf, max_ent); + }, + 1, max_entries, flusher, true, 1); } return 0; diff --git a/src/csrc/benchmarks/bench_sinkhorn_knopp_backward.cu b/src/csrc/benchmarks/bench_sinkhorn_knopp_backward.cu index b424606..b726f81 100644 --- a/src/csrc/benchmarks/bench_sinkhorn_knopp_backward.cu +++ b/src/csrc/benchmarks/bench_sinkhorn_knopp_backward.cu @@ -1,94 +1,97 @@ +// AUTO-GENERATED by benchgen.py — do not edit manually +// Source: src/csrc/kernels/sinkhorn_knopp.cuh +// Regenerate: make benchgen + #include -#include #include -#include "../include/mhc_types.h" -#include "../include/utils.cuh" #include "../kernels/sinkhorn_knopp.cuh" +#include "../include/bench_harness.cuh" using namespace mhc; int main() { const int bench_runs = 100; const float eps = 1e-8f; + const int max_entries = 8; L2Flusher flusher; - struct Config { - int M; - int N; - int iters; - }; - - Config configs[] = { + int configs[][3] = { {32, 32, 5}, {32, 32, 10}, {32, 32, 20}, {64, 64, 5}, {64, 64, 10}, {64, 64, 20}, }; int num_configs = sizeof(configs) / sizeof(configs[0]); printf("Sinkhorn-Knopp Backward Benchmark\n"); - printf("=======================================================================\n"); - printf("%6s %6s %6s %12s %10s %10s %12s\n", "M", "N", "Iters", "Time (us)", "us/iter", "GFLOPS", - "Bandwidth (GB/s)"); - printf("-----------------------------------------------------------------------\n"); + printf("======================================================================\n"); + printf("%8s %8s %6s %12s %10s %10s %12s\n", "M", "N", "num_iters", "Time (us)", "us/iter", + "GFLOPS", "Bandwidth (GB/s)"); + printf("------------------------------------------------------------------------\n"); for (int c = 0; c < num_configs; c++) { - int M = configs[c].M; - int N = configs[c].N; - int num_iters = configs[c].iters; - - float* h_inp = (float*)malloc(M * N * sizeof(float)); - float* h_grad = (float*)malloc(M * N * sizeof(float)); - - srand(42); - for (int i = 0; i < M * N; i++) { - h_inp[i] = (float)rand() / RAND_MAX + 0.1f; - h_grad[i] = (float)rand() / RAND_MAX * 2.0f - 1.0f; - } - - float *d_inp, *d_out, *d_grad, *d_dinp; - CHECK_CUDA(cudaMalloc(&d_inp, M * N * sizeof(float))); - CHECK_CUDA(cudaMalloc(&d_out, M * N * sizeof(float))); - CHECK_CUDA(cudaMalloc(&d_grad, M * N * sizeof(float))); - CHECK_CUDA(cudaMalloc(&d_dinp, M * N * sizeof(float))); - - CHECK_CUDA(cudaMemcpy(d_inp, h_inp, M * N * sizeof(float), cudaMemcpyHostToDevice)); - CHECK_CUDA(cudaMemcpy(d_grad, h_grad, M * N * sizeof(float), cudaMemcpyHostToDevice)); - - sinkhorn_knopp_forward(d_out, d_inp, M, N, num_iters, eps); + int M = configs[c][0]; + int N = configs[c][1]; + int num_iters = configs[c][2]; + + HostMem h_M_inp(M * N); + HostMem h_grad(M * N); + fill_random(h_M_inp, M * N, 0.1f, 1.1f); + fill_random(h_grad, M * N, -1.0f, 1.0f, 43); + + DeviceMem d_M_inp(M * N); + DeviceMem d_grad(M * N); + DeviceMem d_M_out(M * N); + DeviceMem d_d_inp(M * N); + d_M_inp.upload(h_M_inp); + d_grad.upload(h_grad); + + sinkhorn_knopp_forward(d_M_out, d_M_inp, M, N, num_iters, eps); CHECK_CUDA(cudaDeviceSynchronize()); - BenchTimer timer; - float total_time = 0.0f; - - for (int i = 0; i < bench_runs; i++) { - flusher.flush(); + size_t total_bytes = 4 * M * N * sizeof(float); - timer.record_start(); - sinkhorn_knopp_backward(d_dinp, d_grad, d_out, d_inp, N, num_iters, eps); - timer.record_stop(); - total_time += timer.elapsed_ms(); - } + float avg_time_ms = bench_kernel( + [&]() { + sinkhorn_knopp_backward(d_d_inp, d_grad, d_M_out, d_M_inp, N, num_iters, 1e-8f); + }, + bench_runs, flusher); - float avg_time_ms = total_time / bench_runs; float time_us = avg_time_ms * 1000.0f; float time_per_iter_us = time_us / num_iters; + float gflops = ((8.0 * M * N * num_iters) / 1e9f) / (avg_time_ms / 1e3f); + float bw = (total_bytes / 1e9f) / (avg_time_ms / 1e3f); - double flops_per_iter = 8.0 * M * N; - double total_flops = flops_per_iter * num_iters; - float gflops = (total_flops / 1e9f) / (avg_time_ms / 1e3f); - - size_t bytes_io = 4 * M * N * sizeof(float); - float bw = (bytes_io / 1e9f) / (avg_time_ms / 1e3f); - - printf("%6d %6d %6d %12.2f %10.2f %10.2f %12.2f\n", M, N, num_iters, time_us, + printf("%8d %8d %6d %12.2f %10.2f %10.2f %12.2f\n", M, N, num_iters, time_us, time_per_iter_us, gflops, bw); + } + + printf("\n--- Phase Breakdown (M=64, N=64, num_iters=10) ---\n"); + { + int M = 64; + int N = 64; + int num_iters = 10; + + HostMem h_M_inp(M * N); + HostMem h_grad(M * N); + fill_random(h_M_inp, M * N, 0.1f, 1.1f); + fill_random(h_grad, M * N, -1.0f, 1.0f, 43); + + DeviceMem d_M_inp(M * N); + DeviceMem d_grad(M * N); + DeviceMem d_M_out(M * N); + DeviceMem d_d_inp(M * N); + d_M_inp.upload(h_M_inp); + d_grad.upload(h_grad); + + sinkhorn_knopp_forward(d_M_out, d_M_inp, M, N, num_iters, eps); + CHECK_CUDA(cudaDeviceSynchronize()); - cudaFree(d_inp); - cudaFree(d_out); - cudaFree(d_grad); - cudaFree(d_dinp); - free(h_inp); - free(h_grad); + profile_kernel( + [&](int64_t* profiler_buf, int max_ent) { + sinkhorn_knopp_backward_profiled(d_d_inp, d_grad, d_M_out, d_M_inp, N, num_iters, + 1e-8f, profiler_buf, max_ent); + }, + 1, max_entries, flusher, true, 1); } return 0; diff --git a/src/csrc/benchmarks/bench_stream_ops.cu b/src/csrc/benchmarks/bench_stream_ops.cu index acb4e44..8833db3 100644 --- a/src/csrc/benchmarks/bench_stream_ops.cu +++ b/src/csrc/benchmarks/bench_stream_ops.cu @@ -1,11 +1,12 @@ +// AUTO-GENERATED by benchgen.py — do not edit manually +// Source: src/csrc/kernels/stream_ops.cuh +// Regenerate: make benchgen + #include -#include #include -#include -#include "../include/mhc_types.h" -#include "../include/utils.cuh" #include "../kernels/stream_ops.cuh" +#include "../include/bench_harness.cuh" using namespace mhc; @@ -23,130 +24,86 @@ int main() { int num_configs = sizeof(configs) / sizeof(configs[0]); printf("Stream Ops Benchmark\n"); - printf("==========================================\n"); - - printf("stream_aggregate_bf16_fused_sigmoid\n"); + printf("======================================================================\n"); + printf("\nstream_aggregate_bf16_fused_sigmoid\n"); printf("%8s %8s %8s %12s %12s\n", "B", "n", "C", "Time (us)", "Bandwidth (GB/s)"); - printf("-----------------------------------------------------------\n"); + printf("----------------------------------------------------\n"); for (int c = 0; c < num_configs; c++) { int B = configs[c][0]; int n = configs[c][1]; int C = configs[c][2]; - float* h_x = (float*)malloc(B * n * C * sizeof(float)); - float* h_H = (float*)malloc(n * sizeof(float)); - - srand(42); - for (int i = 0; i < B * n * C; i++) - h_x[i] = (float)rand() / RAND_MAX * 2.0f - 1.0f; - for (int i = 0; i < n; i++) - h_H[i] = 0.0f; - - float *d_x, *d_H, *d_H_activated; - float* d_out; - CHECK_CUDA(cudaMalloc(&d_x, B * n * C * sizeof(float))); - CHECK_CUDA(cudaMalloc(&d_H, n * sizeof(float))); - CHECK_CUDA(cudaMalloc(&d_H_activated, n * sizeof(float))); - CHECK_CUDA(cudaMalloc(&d_out, B * C * sizeof(float))); - - CHECK_CUDA(cudaMemcpy(d_x, h_x, B * n * C * sizeof(float), cudaMemcpyHostToDevice)); - CHECK_CUDA(cudaMemcpy(d_H, h_H, n * sizeof(float), cudaMemcpyHostToDevice)); - - size_t bytes = (B * n * C + n) * sizeof(float) + B * C * sizeof(float); - - BenchTimer timer; - float total_time = 0.0f; - - for (int i = 0; i < bench_runs; i++) { - flusher.flush(); - timer.record_start(); - stream_aggregate_bf16_fused_sigmoid(d_out, d_H_activated, d_x, d_H, B, n, C); - timer.record_stop(); - total_time += timer.elapsed_ms(); - } - - float avg_time_ms = total_time / bench_runs; - float bw = (bytes / 1e9f) / (avg_time_ms / 1e3f); - - printf("%8d %8d %8d %12.2f %12.0f\n", B, n, C, avg_time_ms * 1000.0f, bw); - - cudaFree(d_x); - cudaFree(d_H); - cudaFree(d_H_activated); - cudaFree(d_out); - free(h_x); - free(h_H); + HostMem h_inp(B * n * C); + fill_random(h_inp, B * n * C, -1.0f, 1.0f); + + DeviceMem d_inp(B * n * C); + DeviceMem d_H_pre_raw(n); + DeviceMem d_H_pre_activated(n); + DeviceMem d_out(B * C); + d_inp.upload(h_inp); + d_H_pre_raw.zero(); + + size_t total_bytes = (B * n * C + n) * sizeof(float) + B * C * sizeof(float); + + float avg_time_ms = bench_kernel( + [&]() { + stream_aggregate_bf16_fused_sigmoid(d_out, d_H_pre_activated, d_inp, d_H_pre_raw, B, + n, C); + }, + bench_runs, flusher); + + float time_us = avg_time_ms * 1000.0f; + float bw = (total_bytes / 1e9f) / (avg_time_ms / 1e3f); + + printf("%8d %8d %8d %12.2f %12.0f\n", B, n, C, time_us, bw); } + printf("\nstream_distribute_mix_add_fused\n"); printf("%8s %8s %8s %12s %12s\n", "B", "n", "C", "Time (us)", "Bandwidth (GB/s)"); - printf("-----------------------------------------------------------\n"); + printf("----------------------------------------------------\n"); for (int c = 0; c < num_configs; c++) { int B = configs[c][0]; int n = configs[c][1]; int C = configs[c][2]; - float* h_x = (float*)malloc(B * n * C * sizeof(float)); - float* h_y_norm = (float*)malloc(B * C * sizeof(float)); - float* h_H = (float*)malloc(n * sizeof(float)); - float* h_M = (float*)malloc(n * n * sizeof(float)); - - srand(42); - for (int i = 0; i < B * n * C; i++) - h_x[i] = (float)rand() / RAND_MAX * 2.0f - 1.0f; - for (int i = 0; i < B * C; i++) - h_y_norm[i] = (float)rand() / RAND_MAX * 2.0f - 1.0f; - for (int i = 0; i < n; i++) - h_H[i] = 0.0f; - for (int i = 0; i < n * n; i++) - h_M[i] = 1.0f / n; - - float *d_x, *d_H, *d_H_activated, *d_M, *d_out; - float* d_y_norm; - CHECK_CUDA(cudaMalloc(&d_x, B * n * C * sizeof(float))); - CHECK_CUDA(cudaMalloc(&d_y_norm, B * C * sizeof(float))); - CHECK_CUDA(cudaMalloc(&d_H, n * sizeof(float))); - CHECK_CUDA(cudaMalloc(&d_H_activated, n * sizeof(float))); - CHECK_CUDA(cudaMalloc(&d_M, n * n * sizeof(float))); - CHECK_CUDA(cudaMalloc(&d_out, B * n * C * sizeof(float))); - - CHECK_CUDA(cudaMemcpy(d_x, h_x, B * n * C * sizeof(float), cudaMemcpyHostToDevice)); - CHECK_CUDA(cudaMemcpy(d_y_norm, h_y_norm, B * C * sizeof(float), cudaMemcpyHostToDevice)); - CHECK_CUDA(cudaMemcpy(d_H, h_H, n * sizeof(float), cudaMemcpyHostToDevice)); - CHECK_CUDA(cudaMemcpy(d_M, h_M, n * n * sizeof(float), cudaMemcpyHostToDevice)); - - size_t bytes = (B * n * C + n * n + n) * sizeof(float) + B * C * sizeof(float) + - B * n * C * sizeof(float); - - BenchTimer timer; - float total_time = 0.0f; - - for (int i = 0; i < bench_runs; i++) { - flusher.flush(); - timer.record_start(); - stream_distribute_mix_add_fused(d_out, d_H_activated, d_x, d_y_norm, d_H, d_M, B, n, C); - timer.record_stop(); - total_time += timer.elapsed_ms(); - } - - float avg_time_ms = total_time / bench_runs; - float bw = (bytes / 1e9f) / (avg_time_ms / 1e3f); - - printf("%8d %8d %8d %12.2f %12.0f\n", B, n, C, avg_time_ms * 1000.0f, bw); - - cudaFree(d_x); - cudaFree(d_y_norm); - cudaFree(d_H); - cudaFree(d_H_activated); - cudaFree(d_M); - cudaFree(d_out); - free(h_x); - free(h_y_norm); - free(h_H); - free(h_M); + HostMem h_x_inp(B * n * C); + HostMem h_y_norm(B * C); + HostMem h_M(n * n); + fill_random(h_x_inp, B * n * C, -1.0f, 1.0f); + fill_random(h_y_norm, B * C, -1.0f, 1.0f, 43); + for (int _i = 0; _i < n * n; _i++) + h_M.ptr[_i] = (float)(1.0 / n); + + DeviceMem d_x_inp(B * n * C); + DeviceMem d_y_norm(B * C); + DeviceMem d_H_post_raw(n); + DeviceMem d_M(n * n); + DeviceMem d_H_post_activated(n); + DeviceMem d_out(B * n * C); + d_x_inp.upload(h_x_inp); + d_y_norm.upload(h_y_norm); + d_H_post_raw.zero(); + d_M.upload(h_M); + + size_t total_bytes = (B * n * C + n * n + n) * sizeof(float) + B * C * sizeof(float) + + B * n * C * sizeof(float); + + float avg_time_ms = bench_kernel( + [&]() { + stream_distribute_mix_add_fused(d_out, d_H_post_activated, d_x_inp, d_y_norm, + d_H_post_raw, d_M, B, n, C); + }, + bench_runs, flusher); + + float time_us = avg_time_ms * 1000.0f; + float bw = (total_bytes / 1e9f) / (avg_time_ms / 1e3f); + + printf("%8d %8d %8d %12.2f %12.0f\n", B, n, C, time_us, bw); } + return 0; } diff --git a/src/csrc/benchmarks/bench_stream_ops_backward.cu b/src/csrc/benchmarks/bench_stream_ops_backward.cu index aa1d884..cb5db21 100644 --- a/src/csrc/benchmarks/bench_stream_ops_backward.cu +++ b/src/csrc/benchmarks/bench_stream_ops_backward.cu @@ -1,15 +1,18 @@ +// AUTO-GENERATED by benchgen.py — do not edit manually +// Source: src/csrc/kernels/stream_ops.cuh +// Regenerate: make benchgen + #include -#include #include -#include "../include/mhc_types.h" -#include "../include/utils.cuh" #include "../kernels/stream_ops.cuh" +#include "../include/bench_harness.cuh" using namespace mhc; int main() { const int bench_runs = 100; + L2Flusher flusher; int configs[][3] = { @@ -21,104 +24,110 @@ int main() { int num_configs = sizeof(configs) / sizeof(configs[0]); printf("Stream Ops Backward Benchmark\n"); - printf("==============================================\n\n"); - - printf("stream_aggregate_backward\n"); + printf("======================================================================\n"); + printf("\nstream_aggregate_backward\n"); printf("%8s %8s %8s %12s %12s\n", "B", "n", "C", "Time (us)", "Bandwidth (GB/s)"); - printf("-----------------------------------------------------------\n"); + printf("----------------------------------------------------\n"); for (int c = 0; c < num_configs; c++) { - int B = configs[c][0], n = configs[c][1], C = configs[c][2]; + int B = configs[c][0]; + int n = configs[c][1]; + int C = configs[c][2]; + + HostMem h_inp(B * n * C); + HostMem h_grad(B * C); + fill_random(h_inp, B * n * C, -1.0f, 1.0f); + fill_random(h_grad, B * C, -1.0f, 1.0f, 43); + constexpr int BLOCK = 256; int workspace_num_blocks = std::min(128, (B * C + BLOCK - 1) / BLOCK); - float *d_x, *d_H, *d_grad, *d_dx, *d_dH, *d_workspace; - CHECK_CUDA(cudaMalloc(&d_x, B * n * C * sizeof(float))); - CHECK_CUDA(cudaMalloc(&d_H, n * sizeof(float))); - CHECK_CUDA(cudaMalloc(&d_grad, B * C * sizeof(float))); - CHECK_CUDA(cudaMalloc(&d_dx, B * n * C * sizeof(float))); - CHECK_CUDA(cudaMalloc(&d_dH, n * sizeof(float))); - CHECK_CUDA(cudaMalloc(&d_workspace, workspace_num_blocks * n * sizeof(float))); - - size_t bytes = (B * n * C + B * C + n) * sizeof(float) + (B * n * C + n) * sizeof(float); - BenchTimer timer; - float total_time = 0.0f; - - for (int i = 0; i < bench_runs; i++) { - flusher.flush(); - timer.record_start(); - stream_aggregate_backward(d_dx, d_dH, d_grad, d_x, d_H, B, n, C, d_workspace, - workspace_num_blocks); - timer.record_stop(); - total_time += timer.elapsed_ms(); - } - - float avg_ms = total_time / bench_runs; - printf("%8d %8d %8d %12.2f %12.0f\n", B, n, C, avg_ms * 1000.0f, - (bytes / 1e9f) / (avg_ms / 1e3f)); - - cudaFree(d_x); - cudaFree(d_H); - cudaFree(d_grad); - cudaFree(d_dx); - cudaFree(d_dH); - cudaFree(d_workspace); + DeviceMem d_inp(B * n * C); + DeviceMem d_H_pre(n); + DeviceMem d_grad(B * C); + DeviceMem d_workspace(workspace_num_blocks * n); + DeviceMem d_d_inp(B * n * C); + DeviceMem d_d_H_pre(n); + d_inp.upload(h_inp); + d_H_pre.zero(); + d_grad.upload(h_grad); + d_workspace.zero(); + + size_t total_bytes = + (B * n * C + B * C + n) * sizeof(float) + (B * n * C + n) * sizeof(float); + + float avg_time_ms = bench_kernel( + [&]() { + stream_aggregate_backward(d_d_inp, d_d_H_pre, d_grad, d_inp, d_H_pre, B, n, C, + d_workspace, workspace_num_blocks); + }, + bench_runs, flusher); + + float time_us = avg_time_ms * 1000.0f; + float bw = (total_bytes / 1e9f) / (avg_time_ms / 1e3f); + + printf("%8d %8d %8d %12.2f %12.0f\n", B, n, C, time_us, bw); } + printf("\nstream_distribute_mix_backward_fused\n"); - printf("%8s %8s %8s %12s %12s\n", "B", "n", "C", "Time (us)", "BW (GB/s)"); - printf("-----------------------------------------------------------\n"); + printf("%8s %8s %8s %12s %12s\n", "B", "n", "C", "Time (us)", "Bandwidth (GB/s)"); + printf("----------------------------------------------------\n"); for (int c = 0; c < num_configs; c++) { - int B = configs[c][0], n = configs[c][1], C = configs[c][2]; + int B = configs[c][0]; + int n = configs[c][1]; + int C = configs[c][2]; + + HostMem h_x(B * n * C); + HostMem h_y_norm(B * C); + HostMem h_M(n * n); + HostMem h_grad(B * n * C); + fill_random(h_x, B * n * C, -1.0f, 1.0f); + fill_random(h_y_norm, B * C, -1.0f, 1.0f, 43); + for (int _i = 0; _i < n * n; _i++) + h_M.ptr[_i] = (float)(1.0 / n); + fill_random(h_grad, B * n * C, -1.0f, 1.0f, 44); + constexpr int BLOCK = 256; int workspace_num_blocks = std::min(128, (B * C + BLOCK - 1) / BLOCK); - float *d_x, *d_y, *d_H, *d_M, *d_grad; - float *d_dx, *d_dy, *d_dM, *d_dH, *d_workspace_M, *d_workspace_H; - CHECK_CUDA(cudaMalloc(&d_x, B * n * C * sizeof(float))); - CHECK_CUDA(cudaMalloc(&d_y, B * C * sizeof(float))); - CHECK_CUDA(cudaMalloc(&d_H, n * sizeof(float))); - CHECK_CUDA(cudaMalloc(&d_M, n * n * sizeof(float))); - CHECK_CUDA(cudaMalloc(&d_grad, B * n * C * sizeof(float))); - CHECK_CUDA(cudaMalloc(&d_dx, B * n * C * sizeof(float))); - CHECK_CUDA(cudaMalloc(&d_dy, B * C * sizeof(float))); - CHECK_CUDA(cudaMalloc(&d_dM, n * n * sizeof(float))); - CHECK_CUDA(cudaMalloc(&d_dH, n * sizeof(float))); - CHECK_CUDA(cudaMalloc(&d_workspace_M, workspace_num_blocks * n * n * sizeof(float))); - CHECK_CUDA(cudaMalloc(&d_workspace_H, workspace_num_blocks * n * sizeof(float))); - - size_t bytes = (2 * B * n * C + B * C + n * n + n) * sizeof(float) + - (B * n * C + B * C + n * n + n) * sizeof(float); - BenchTimer timer; - float total_time = 0.0f; - - for (int i = 0; i < bench_runs; i++) { - flusher.flush(); - timer.record_start(); - stream_distribute_mix_backward_fused(d_dx, d_dy, d_dM, d_dH, d_grad, d_x, d_y, d_M, d_H, - B, n, C, d_workspace_M, d_workspace_H, - workspace_num_blocks); - timer.record_stop(); - total_time += timer.elapsed_ms(); - } - - float avg_ms = total_time / bench_runs; - printf("%8d %8d %8d %12.2f %12.0f\n", B, n, C, avg_ms * 1000.0f, - (bytes / 1e9f) / (avg_ms / 1e3f)); - - cudaFree(d_x); - cudaFree(d_y); - cudaFree(d_H); - cudaFree(d_M); - cudaFree(d_grad); - cudaFree(d_dx); - cudaFree(d_dy); - cudaFree(d_dM); - cudaFree(d_dH); - cudaFree(d_workspace_M); - cudaFree(d_workspace_H); + DeviceMem d_x(B * n * C); + DeviceMem d_y_norm(B * C); + DeviceMem d_M(n * n); + DeviceMem d_H_post(n); + DeviceMem d_grad(B * n * C); + DeviceMem d_workspace_M(workspace_num_blocks * n * n); + DeviceMem d_workspace_H(workspace_num_blocks * n); + DeviceMem d_d_x(B * n * C); + DeviceMem d_d_y_norm(B * C); + DeviceMem d_d_M(n * n); + DeviceMem d_d_H_post(n); + d_x.upload(h_x); + d_y_norm.upload(h_y_norm); + d_M.upload(h_M); + d_H_post.zero(); + d_grad.upload(h_grad); + d_workspace_M.zero(); + d_workspace_H.zero(); + + size_t total_bytes = (2 * B * n * C + B * C + n * n + n) * sizeof(float) + + (B * n * C + B * C + n * n + n) * sizeof(float); + + float avg_time_ms = bench_kernel( + [&]() { + stream_distribute_mix_backward_fused( + d_d_x, d_d_y_norm, d_d_M, d_d_H_post, d_grad, d_x, d_y_norm, d_M, d_H_post, B, + n, C, d_workspace_M, d_workspace_H, workspace_num_blocks); + }, + bench_runs, flusher); + + float time_us = avg_time_ms * 1000.0f; + float bw = (total_bytes / 1e9f) / (avg_time_ms / 1e3f); + + printf("%8d %8d %8d %12.2f %12.0f\n", B, n, C, time_us, bw); } + return 0; } diff --git a/src/csrc/include/bench_harness.cuh b/src/csrc/include/bench_harness.cuh new file mode 100644 index 0000000..0b5855b --- /dev/null +++ b/src/csrc/include/bench_harness.cuh @@ -0,0 +1,167 @@ +#pragma once + +#include +#include +#include +#include + +#include "mhc_types.h" +#include "utils.cuh" + +namespace mhc { + +// RAII wrapper for device memory (cudaMalloc / cudaFree). +template struct DeviceMem { + T* ptr; + size_t count; + + DeviceMem() : ptr(nullptr), count(0) {} + + explicit DeviceMem(size_t n) : ptr(nullptr), count(n) { + CHECK_CUDA(cudaMalloc(&ptr, n * sizeof(T))); + } + + ~DeviceMem() { + if (ptr) + cudaFree(ptr); + } + + DeviceMem(const DeviceMem&) = delete; + DeviceMem& operator=(const DeviceMem&) = delete; + + DeviceMem(DeviceMem&& o) noexcept : ptr(o.ptr), count(o.count) { + o.ptr = nullptr; + o.count = 0; + } + + DeviceMem& operator=(DeviceMem&& o) noexcept { + if (this != &o) { + if (ptr) + cudaFree(ptr); + ptr = o.ptr; + count = o.count; + o.ptr = nullptr; + o.count = 0; + } + return *this; + } + + void upload(const T* host_data) { + CHECK_CUDA(cudaMemcpy(ptr, host_data, count * sizeof(T), cudaMemcpyHostToDevice)); + } + + void zero() { CHECK_CUDA(cudaMemset(ptr, 0, count * sizeof(T))); } + + // cppcheck-suppress unusedFunction + size_t bytes() const { return count * sizeof(T); } + + operator T*() { return ptr; } // cppcheck-suppress noExplicitConstructor + operator const T*() const { return ptr; } // cppcheck-suppress noExplicitConstructor +}; + +// RAII wrapper for host memory (malloc / free). +template struct HostMem { + T* ptr; + size_t count; + + HostMem() : ptr(nullptr), count(0) {} + + explicit HostMem(size_t n) : ptr(nullptr), count(n) { ptr = (T*)malloc(n * sizeof(T)); } + + ~HostMem() { + if (ptr) + free(ptr); + } + + HostMem(const HostMem&) = delete; + HostMem& operator=(const HostMem&) = delete; + + HostMem(HostMem&& o) noexcept : ptr(o.ptr), count(o.count) { + o.ptr = nullptr; + o.count = 0; + } + + HostMem& operator=(HostMem&& o) noexcept { + if (this != &o) { + if (ptr) + free(ptr); + ptr = o.ptr; + count = o.count; + o.ptr = nullptr; + o.count = 0; + } + return *this; + } + + size_t bytes() const { return count * sizeof(T); } + + operator T*() { return ptr; } // cppcheck-suppress noExplicitConstructor + operator const T*() const { return ptr; } // cppcheck-suppress noExplicitConstructor +}; + +// Run a kernel lambda `runs` times with L2 flush. Returns average time in ms. +template float bench_kernel(KernelFn fn, int runs, L2Flusher& flusher) { + BenchTimer timer; + float total_time = 0.0f; + + for (int i = 0; i < runs; i++) { + flusher.flush(); + timer.record_start(); + fn(); + timer.record_stop(); + total_time += timer.elapsed_ms(); + } + return total_time / runs; +} + +// Run with per-iteration setup (e.g. zeroing gradient buffers). +template +float bench_kernel(KernelFn fn, int runs, L2Flusher& flusher, PreFn pre) { + BenchTimer timer; + float total_time = 0.0f; + + for (int i = 0; i < runs; i++) { + flusher.flush(); + pre(); + timer.record_start(); + fn(); + timer.record_stop(); + total_time += timer.elapsed_ms(); + } + return total_time / runs; +} + +// Run a profiled kernel once and print phase breakdown. +// KernelFn signature: void(int64_t* profiler_buf, int max_entries) +template +void profile_kernel(KernelFn fn, int num_blocks, int max_entries, L2Flusher& flusher, + bool print_timeline = false, int timeline_blocks = 4) { + HostProfiler profiler(num_blocks, max_entries); + flusher.flush(); + fn(profiler.device_ptr(), max_entries); + CHECK_CUDA(cudaDeviceSynchronize()); + profiler.print_summary(); + if (print_timeline) { + profiler.print_timeline(timeline_blocks); + } +} + +// Fill buffer with random floats in [lo, hi]. +inline void fill_random(float* buf, int n, float lo = -1.0f, float hi = 1.0f, int seed = 42) { + srand(seed); + float range = hi - lo; + for (int i = 0; i < n; i++) { + buf[i] = lo + (float)rand() / RAND_MAX * range; + } +} + +// Fill buffer with random bf16 values in [lo, hi]. +inline void fill_random_bf16(floatX* buf, int n, float lo = -1.0f, float hi = 1.0f, int seed = 42) { + srand(seed); + float range = hi - lo; + for (int i = 0; i < n; i++) { + buf[i] = (floatX)(lo + (float)rand() / RAND_MAX * range); + } +} + +} // namespace mhc diff --git a/src/csrc/include/profiling.cuh b/src/csrc/include/profiling.cuh new file mode 100644 index 0000000..2c10daf --- /dev/null +++ b/src/csrc/include/profiling.cuh @@ -0,0 +1,64 @@ +#pragma once + +#include "utils.cuh" + +namespace mhc { + +// Phase profiling macros for production CUDA kernels. +// +// Usage: Add to __global__ kernels templated with `bool DO_PROFILE = false`. +// Kernel must accept `int64_t* profiler_buf, int max_entries` as final params. +// +// When DO_PROFILE=false, `if constexpr` eliminates all profiling code at compile +// time — zero register pressure, zero instructions, zero overhead. +// +// Example: +// template +// __global__ void my_kernel(float* out, const float* inp, int N, +// int64_t* profiler_buf, int max_entries) { +// MHC_PROFILE_INIT(blockIdx.x); +// MHC_PROFILE_START(TagLoad); +// // ... load phase ... +// MHC_PROFILE_PHASE(TagCompute); +// // ... compute phase ... +// MHC_PROFILE_PHASE(TagStore); +// // ... store phase ... +// MHC_PROFILE_END(); +// } + +// Declare DeviceProfiler and initialize it for the given block. +#define MHC_PROFILE_INIT(block_id) \ + DeviceProfiler _mhc_profiler; \ + if constexpr (DO_PROFILE) { \ + if (threadIdx.x == 0) { \ + _mhc_profiler.init(max_entries, profiler_buf, block_id); \ + } \ + } + +// Begin a new profiling phase. Use for the first phase after INIT. +#define MHC_PROFILE_START(tag) \ + if constexpr (DO_PROFILE) { \ + if (threadIdx.x == 0) { \ + _mhc_profiler.start(tag); \ + } \ + } + +// Transition from one phase to the next (stop current, start new). +#define MHC_PROFILE_PHASE(tag) \ + if constexpr (DO_PROFILE) { \ + if (threadIdx.x == 0) { \ + _mhc_profiler.stop(); \ + _mhc_profiler.start(tag); \ + } \ + } + +// Stop the final phase and flush profiling data. +#define MHC_PROFILE_END() \ + if constexpr (DO_PROFILE) { \ + if (threadIdx.x == 0) { \ + _mhc_profiler.stop(); \ + _mhc_profiler.flush(); \ + } \ + } + +} // namespace mhc diff --git a/src/csrc/kernels/fused_rmsnorm_matmul.cuh b/src/csrc/kernels/fused_rmsnorm_matmul.cuh index de8828a..f139af0 100644 --- a/src/csrc/kernels/fused_rmsnorm_matmul.cuh +++ b/src/csrc/kernels/fused_rmsnorm_matmul.cuh @@ -305,6 +305,18 @@ inline void matmul_forward(MatmulDescriptors& desc, float* out, const floatX* A, &desc.heuristic.algo, desc.workspace, desc.workspace_size, stream)); } +// @bench fused_rmsnorm_matmul +// @title: Fused RMSNorm + MatMul +// @configs: (M,N,K) = +// [(128,4096,4096),(256,4096,4096),(512,4096,4096),(1024,4096,4096),(2048,4096,4096),(1024,8192,4096),(2048,8192,4096)] +// @in: inp floatX[M * K] bf16(-1,1), weight floatX[N * K] bf16(0.75,1.25) +// @out: out float[M * N] +// @tflops: 2.0 * (double)M * (double)N * (double)K +// @bandwidth: M * K * sizeof(floatX) + N * K * sizeof(floatX) + M * N * sizeof(float) +// @setup: FusedRMSNormMatmul fused; +// @setup: fused.init(M, N, K); +// @call: fused.forward(d_out, d_inp, d_weight) +// @cleanup: fused.destroy(); struct FusedRMSNormMatmul { MatmulDescriptors matmul_desc; float* rms_buffer; @@ -533,6 +545,30 @@ __global__ void rms_correction_kernel(float* __restrict__ dx, const float* __res } } +// @bench fused_rmsnorm_matmul_backward +// @title: Fused RMSNorm + MatMul Backward +// @configs: (M,N,K) = +// [(128,4096,4096),(256,4096,4096),(512,4096,4096),(1024,4096,4096),(2048,4096,4096),(1024,8192,4096),(2048,8192,4096)] +// @in: inp floatX[M * K] bf16(-1,1), weight floatX[N * K] bf16(0.75,1.25), +// grad float[M * N] random(-1,1,43), rms float[M] computed +// @out: dW float[N * K], dx float[M * K] +// @setup: for (int i = 0; i < M; i++) { +// @setup: float sum_sq = 0.0f; +// @setup: for (int j = 0; j < K; j++) { +// @setup: float v = (float)h_inp.ptr[i * K + j]; +// @setup: sum_sq += v * v; +// @setup: } +// @setup: h_rms.ptr[i] = sqrtf(sum_sq / (float)K + 1e-5f); +// @setup: } +// @setup: d_rms.upload(h_rms); +// @setup: FusedRMSNormMatmulBackward backward; +// @setup: backward.init(M, N, K); +// @call: backward.backward(d_dW, d_dx, d_grad, d_inp, d_weight, d_rms) +// @cleanup: backward.destroy(); +// @pre-iter: d_dW.zero() +// @tflops: 4.0 * (double)M * (double)N * (double)K +// @bandwidth: M * K * sizeof(floatX) + N * K * sizeof(floatX) + +// M * N * sizeof(float) + M * sizeof(float) + N * K * sizeof(float) + M * K * sizeof(float) struct FusedRMSNormMatmulBackward { cublasLtHandle_t handle; cublasLtMatmulDesc_t dW_matmul_desc; diff --git a/src/csrc/kernels/mhc_layer.cuh b/src/csrc/kernels/mhc_layer.cuh index c640042..cc14226 100644 --- a/src/csrc/kernels/mhc_layer.cuh +++ b/src/csrc/kernels/mhc_layer.cuh @@ -418,6 +418,77 @@ inline void apply_dynamic_h_activations(float* H_pre_out, float* H_post_out, flo alpha_res, B, n); } +// @bench mhc_layer_static +// @group: mhc_layer +// @title: Static H Path +// @configs: (B,C,n) = +// [(64,1280,4),(128,1280,4),(256,1280,4),(320,1280,4),(64,1920,4),(128,1920,4),(64,2560,4),(128,2560,4)] +// @runs: 50 +// @setup: HostMem h_x_expanded(B * n * C); +// @setup: HostMem h_rmsnorm_weight(C); +// @setup: HostMem h_H_pre(n); +// @setup: HostMem h_H_post(n); +// @setup: HostMem h_H_res(n * n); +// @setup: fill_random(h_x_expanded, B * n * C); +// @setup: fill_random_bf16(h_rmsnorm_weight, C, 0.75f, 1.25f); +// @setup: for (int i = 0; i < n; i++) { h_H_pre.ptr[i] = 0.0f; h_H_post.ptr[i] = 0.0f; } +// @setup: srand(42); +// @setup: for (int i = 0; i < n * n; i++) { h_H_res.ptr[i] = 0.01f * ((float)rand() / RAND_MAX +// * 2.0f - 1.0f); } +// @setup: DeviceMem d_x_expanded(B * n * C); +// @setup: d_x_expanded.upload(h_x_expanded); +// @setup: MHCLayerConfig cfg; +// @setup: cfg.batch_size = B; cfg.hidden_dim = C; cfg.expansion_rate = n; +// @setup: cfg.sinkhorn_iters = 20; cfg.eps = 1e-5f; cfg.use_pdl = true; cfg.use_dynamic_h = false; +// @setup: MHCLayer layer; +// @setup: layer.init(cfg); +// @setup: layer.set_weights(h_rmsnorm_weight, h_H_pre, h_H_post, h_H_res); +// @setup: layer.sync(); +// @setup: layer.forward_device(d_x_expanded); +// @setup: layer.sync(); +// @bandwidth: (size_t)B * n * C * sizeof(float) * 3 +// @metric: Samples/sec = B / (avg_time_ms / 1000.0f) +// @call: layer.forward_device(d_x_expanded) +// @cleanup: layer.destroy(); + +// @bench mhc_layer_dynamic +// @group: mhc_layer +// @title: Dynamic H Path +// @configs: (B,C,n) = +// [(64,1280,4),(128,1280,4),(256,1280,4),(320,1280,4),(64,1920,4),(128,1920,4),(64,2560,4),(128,2560,4)] +// @runs: 50 +// @setup: int nC = n * C; +// @setup: int total_H_dim = n + n + n * n; +// @setup: HostMem h_x_expanded(B * n * C); +// @setup: HostMem h_rmsnorm_weight(C); +// @setup: HostMem h_phi(total_H_dim * nC); +// @setup: HostMem h_b_pre(n); +// @setup: HostMem h_b_post(n); +// @setup: HostMem h_b_res(n * n); +// @setup: fill_random(h_x_expanded, B * n * C); +// @setup: fill_random_bf16(h_rmsnorm_weight, C, 0.75f, 1.25f); +// @setup: fill_random_bf16(h_phi, total_H_dim * nC, -0.05f, 0.05f, 43); +// @setup: for (int i = 0; i < n; i++) { h_b_pre.ptr[i] = 0.0f; h_b_post.ptr[i] = 0.0f; } +// @setup: srand(42); +// @setup: for (int i = 0; i < n * n; i++) { h_b_res.ptr[i] = 0.01f * ((float)rand() / RAND_MAX +// * 2.0f - 1.0f); } +// @setup: DeviceMem d_x_expanded(B * n * C); +// @setup: d_x_expanded.upload(h_x_expanded); +// @setup: MHCLayerConfig cfg; +// @setup: cfg.batch_size = B; cfg.hidden_dim = C; cfg.expansion_rate = n; +// @setup: cfg.sinkhorn_iters = 20; cfg.eps = 1e-5f; cfg.use_pdl = true; cfg.use_dynamic_h = true; +// @setup: MHCLayer layer; +// @setup: layer.init(cfg); +// @setup: floatX* phi_base = h_phi; +// @setup: layer.set_weights_dynamic(h_rmsnorm_weight, phi_base, phi_base + n * nC, phi_base + 2 * n +// * nC, h_b_pre, h_b_post, h_b_res, 0.01f, 0.01f, 0.01f); +// @setup: layer.sync(); +// @setup: layer.forward_device(d_x_expanded); +// @setup: layer.sync(); +// @bandwidth: (size_t)B * n * C * sizeof(float) * 3 +// @metric: Samples/sec = B / (avg_time_ms / 1000.0f) +// @call: layer.forward_device(d_x_expanded) +// @cleanup: layer.destroy(); struct MHCLayer { MHCLayerConfig config; MHCLayerWeights weights; diff --git a/src/csrc/kernels/rmsnorm.cuh b/src/csrc/kernels/rmsnorm.cuh index d3e0c12..5f7a4dd 100644 --- a/src/csrc/kernels/rmsnorm.cuh +++ b/src/csrc/kernels/rmsnorm.cuh @@ -5,17 +5,17 @@ #include #include #include "../include/mhc_types.h" +#include "../include/profiling.cuh" namespace cg = cooperative_groups; namespace mhc { -template -__global__ __launch_bounds__(BLOCK_SIZE, 2) void rmsnorm_kernel(float* __restrict__ out, - float* __restrict__ rms_out, - const float* __restrict__ inp, - const floatX* __restrict__ weight, - int N, int C, float eps) { +template +__global__ __launch_bounds__(BLOCK_SIZE, 2) void rmsnorm_kernel( + float* __restrict__ out, float* __restrict__ rms_out, const float* __restrict__ inp, + const floatX* __restrict__ weight, int N, int C, float eps, int64_t* profiler_buf = nullptr, + int max_entries = 0) { cg::thread_block block = cg::this_thread_block(); cg::thread_block_tile<32> warp = cg::tiled_partition<32>(block); @@ -23,6 +23,9 @@ __global__ __launch_bounds__(BLOCK_SIZE, 2) void rmsnorm_kernel(float* __restric if (idx >= N) return; + MHC_PROFILE_INIT(idx); + MHC_PROFILE_START(TagLoad); + const float* x = inp + idx * C; float* o = out + idx * C; @@ -35,6 +38,8 @@ __global__ __launch_bounds__(BLOCK_SIZE, 2) void rmsnorm_kernel(float* __restric thread_sum_sq += val * val; } + MHC_PROFILE_PHASE(TagReduce); + float warp_sum = cg::reduce(warp, thread_sum_sq, cg::plus()); int warp_id = threadIdx.x / 32; @@ -61,6 +66,8 @@ __global__ __launch_bounds__(BLOCK_SIZE, 2) void rmsnorm_kernel(float* __restric } __syncthreads(); + MHC_PROFILE_PHASE(TagStore); + float rms_inv = s_sum_sq[0]; for (int i = threadIdx.x; i < C; i += BLOCK_SIZE) { @@ -68,12 +75,15 @@ __global__ __launch_bounds__(BLOCK_SIZE, 2) void rmsnorm_kernel(float* __restric float w = (float)weight[i]; o[i] = val * rms_inv * w; } + + MHC_PROFILE_END(); } -template +template __global__ __launch_bounds__(BLOCK_SIZE, 2) void rmsnorm_kernel_vectorized( float* __restrict__ out, float* __restrict__ rms_out, const float* __restrict__ inp, - const floatX* __restrict__ weight, int N, int C, float eps) { + const floatX* __restrict__ weight, int N, int C, float eps, int64_t* profiler_buf = nullptr, + int max_entries = 0) { cg::thread_block block = cg::this_thread_block(); cg::thread_block_tile<32> warp = cg::tiled_partition<32>(block); @@ -81,6 +91,9 @@ __global__ __launch_bounds__(BLOCK_SIZE, 2) void rmsnorm_kernel_vectorized( if (idx >= N) return; + MHC_PROFILE_INIT(idx); + MHC_PROFILE_START(TagLoad); + const float* x = inp + idx * C; float* o = out + idx * C; @@ -105,6 +118,8 @@ __global__ __launch_bounds__(BLOCK_SIZE, 2) void rmsnorm_kernel_vectorized( thread_sum_sq += val * val; } + MHC_PROFILE_PHASE(TagReduce); + float warp_sum = cg::reduce(warp, thread_sum_sq, cg::plus()); int warp_id = threadIdx.x / 32; @@ -131,6 +146,8 @@ __global__ __launch_bounds__(BLOCK_SIZE, 2) void rmsnorm_kernel_vectorized( } __syncthreads(); + MHC_PROFILE_PHASE(TagStore); + float rms_inv = s_sum_sq[0]; // Output as fp32: read fp32 input, bf16 weight, compute in fp32, write fp32 @@ -168,8 +185,20 @@ __global__ __launch_bounds__(BLOCK_SIZE, 2) void rmsnorm_kernel_vectorized( float w = (float)weight[i]; o[i] = val * rms_inv * w; } + + MHC_PROFILE_END(); } +// @bench rmsnorm +// @title: RMSNorm +// @configs: (N,C) = +// [(128,4096),(256,4096),(512,4096),(1024,4096),(2048,4096),(1024,8192),(2048,8192)] +// @in: inp float[N * C] random(-1,1), weight floatX[C] bf16(0.75,1.25) +// @out: out float[N * C] +// @scalar: eps=1e-5f +// @bandwidth: (size_t)N * C * sizeof(float) + (size_t)C * sizeof(floatX) + (size_t)N * C * +// sizeof(float) +// @profile: rmsnorm_forward_profiled N=1024 C=4096 grid=N inline void rmsnorm_forward(float* out, const float* inp, const floatX* weight, int N, int C, float eps, cudaStream_t stream = nullptr) { constexpr int BLOCK_SIZE = 512; @@ -213,10 +242,10 @@ inline void rmsnorm_forward_with_rms(float* out, float* rms_out, const float* in if (C % 8 == 0 && C >= 64) { cudaLaunchKernelEx(&config, rmsnorm_kernel_vectorized, out, rms_out, inp, - weight, N, C, eps); + weight, N, C, eps, (int64_t*)nullptr, 0); } else { cudaLaunchKernelEx(&config, rmsnorm_kernel, out, rms_out, inp, weight, N, - C, eps); + C, eps, (int64_t*)nullptr, 0); } #else if (C % 8 == 0 && C >= 64) { @@ -228,11 +257,11 @@ inline void rmsnorm_forward_with_rms(float* out, float* rms_out, const float* in } #endif } -template +template __global__ __launch_bounds__(BLOCK_SIZE, 2) void rmsnorm_backward_kernel( float* __restrict__ d_inp, float* __restrict__ d_weight, const float* __restrict__ grad, const float* __restrict__ inp, const floatX* __restrict__ weight, const float* __restrict__ rms, - int N, int C) { + int N, int C, int64_t* profiler_buf = nullptr, int max_entries = 0) { cg::thread_block block = cg::this_thread_block(); cg::thread_block_tile<32> warp = cg::tiled_partition<32>(block); @@ -240,6 +269,9 @@ __global__ __launch_bounds__(BLOCK_SIZE, 2) void rmsnorm_backward_kernel( if (idx >= N) return; + MHC_PROFILE_INIT(idx); + MHC_PROFILE_START(TagLoad); + const float* x = inp + idx * C; const float* g = grad + idx * C; float* dx = d_inp + idx * C; @@ -257,6 +289,8 @@ __global__ __launch_bounds__(BLOCK_SIZE, 2) void rmsnorm_backward_kernel( thread_dot += g_val * w_val * x_val; } + MHC_PROFILE_PHASE(TagReduce); + float warp_dot = cg::reduce(warp, thread_dot, cg::plus()); int warp_id = threadIdx.x / 32; @@ -277,6 +311,8 @@ __global__ __launch_bounds__(BLOCK_SIZE, 2) void rmsnorm_backward_kernel( } __syncthreads(); + MHC_PROFILE_PHASE(TagStore); + float dot_sum = s_reduce[0]; float correction = dot_sum / ((float)C * r * r); @@ -289,8 +325,31 @@ __global__ __launch_bounds__(BLOCK_SIZE, 2) void rmsnorm_backward_kernel( atomicAdd(&d_weight[i], g_val * x_val * r_inv); } + + MHC_PROFILE_END(); } +// @bench rmsnorm_backward +// @title: RMSNorm Backward +// @configs: (N,C) = +// [(128,4096),(256,4096),(512,4096),(1024,4096),(2048,4096),(1024,8192),(2048,8192)] +// @in: grad float[N * C] random(-1,1,43), inp float[N * C] random(-1,1), weight floatX[C] +// bf16(0.75,1.25), rms float[N] computed +// @out: d_inp float[N * C], d_weight float[C] +// @setup: for (int i = 0; i < N; i++) { +// @setup: float sum_sq = 0.0f; +// @setup: for (int j = 0; j < C; j++) { +// @setup: float v = h_inp.ptr[i * C + j]; +// @setup: sum_sq += v * v; +// @setup: } +// @setup: h_rms.ptr[i] = sqrtf(sum_sq / (float)C + 1e-5f); +// @setup: } +// @setup: d_rms.upload(h_rms); +// @pre-iter: d_d_weight.zero() +// @bandwidth: (size_t)N * C * sizeof(float) + (size_t)C * sizeof(floatX) + (size_t)N * C * +// sizeof(float) + (size_t)N * sizeof(float) + (size_t)N * C * sizeof(float) + (size_t)C * +// sizeof(float) +// @profile: rmsnorm_backward_profiled N=1024 C=4096 grid=N inline void rmsnorm_backward(float* d_inp, float* d_weight, const float* grad, const float* inp, const floatX* weight, const float* rms, int N, int C, cudaStream_t stream = nullptr) { @@ -302,4 +361,37 @@ inline void rmsnorm_backward(float* d_inp, float* d_weight, const float* grad, c <<>>(d_inp, d_weight, grad, inp, weight, rms, N, C); } +// Profiled dispatch variants — call production kernels with DO_PROFILE=true. + +inline void rmsnorm_forward_profiled(float* out, const float* inp, const floatX* weight, int N, + int C, float eps, int64_t* profiler_buf, int max_entries, + cudaStream_t stream = nullptr) { + constexpr int BLOCK_SIZE = 512; + int num_warps = BLOCK_SIZE / 32; + size_t shared_mem = num_warps * sizeof(float); + + dim3 grid(N); + dim3 block(BLOCK_SIZE); + + if (C % 8 == 0 && C >= 64) { + rmsnorm_kernel_vectorized<<>>( + out, nullptr, inp, weight, N, C, eps, profiler_buf, max_entries); + } else { + rmsnorm_kernel<<>>( + out, nullptr, inp, weight, N, C, eps, profiler_buf, max_entries); + } +} + +inline void rmsnorm_backward_profiled(float* d_inp, float* d_weight, const float* grad, + const float* inp, const floatX* weight, const float* rms, + int N, int C, int64_t* profiler_buf, int max_entries, + cudaStream_t stream = nullptr) { + constexpr int BLOCK_SIZE = 512; + int num_warps = BLOCK_SIZE / 32; + size_t shared_mem = num_warps * sizeof(float); + + rmsnorm_backward_kernel<<>>( + d_inp, d_weight, grad, inp, weight, rms, N, C, profiler_buf, max_entries); +} + } // namespace mhc diff --git a/src/csrc/kernels/sinkhorn_knopp.cuh b/src/csrc/kernels/sinkhorn_knopp.cuh index 8c1354d..4570ad7 100644 --- a/src/csrc/kernels/sinkhorn_knopp.cuh +++ b/src/csrc/kernels/sinkhorn_knopp.cuh @@ -6,6 +6,7 @@ #include #include "../include/mhc_types.h" #include "../include/utils.cuh" +#include "../include/profiling.cuh" namespace cg = cooperative_groups; @@ -301,10 +302,10 @@ __global__ __launch_bounds__(BLOCK_SIZE, } } -template +template __global__ __launch_bounds__(BLOCK_SIZE, 2) void sinkhorn_knopp_single_block_kernel( - float* __restrict__ out, const float* __restrict__ inp, int M, int N, int num_iters, - float eps) { + float* __restrict__ out, const float* __restrict__ inp, int M, int N, int num_iters, float eps, + int64_t* profiler_buf = nullptr, int max_entries = 0) { extern __shared__ float smem[]; float* tile = smem; float* row_sums = smem + MAX_DIM * MAX_DIM; @@ -312,11 +313,16 @@ __global__ __launch_bounds__(BLOCK_SIZE, 2) void sinkhorn_knopp_single_block_ker int total_elems = M * N; + MHC_PROFILE_INIT(blockIdx.x); + MHC_PROFILE_START(TagLoad); + for (int i = threadIdx.x; i < total_elems; i += BLOCK_SIZE) { tile[i] = inp[i]; } __syncthreads(); + MHC_PROFILE_PHASE(TagCompute); + for (int iter = 0; iter < num_iters; iter++) { for (int r = threadIdx.x; r < M; r += BLOCK_SIZE) { float sum = 0.0f; @@ -355,15 +361,19 @@ __global__ __launch_bounds__(BLOCK_SIZE, 2) void sinkhorn_knopp_single_block_ker __syncthreads(); } + MHC_PROFILE_PHASE(TagStore); + for (int i = threadIdx.x; i < total_elems; i += BLOCK_SIZE) { out[i] = tile[i]; } + + MHC_PROFILE_END(); } -template +template __global__ __launch_bounds__(BLOCK_SIZE, 2) void sinkhorn_knopp_single_block_fused_exp_kernel( float* __restrict__ out, float* __restrict__ H_res_exp, const float* __restrict__ inp, int M, - int N, int num_iters, float eps) { + int N, int num_iters, float eps, int64_t* profiler_buf = nullptr, int max_entries = 0) { extern __shared__ float smem[]; float* tile = smem; float* row_sums = smem + MAX_DIM * MAX_DIM; @@ -371,6 +381,9 @@ __global__ __launch_bounds__(BLOCK_SIZE, 2) void sinkhorn_knopp_single_block_fus int total_elems = M * N; + MHC_PROFILE_INIT(blockIdx.x); + MHC_PROFILE_START(TagLoad); + for (int i = threadIdx.x; i < total_elems; i += BLOCK_SIZE) { float val = fast_exp(inp[i]); tile[i] = val; @@ -379,6 +392,8 @@ __global__ __launch_bounds__(BLOCK_SIZE, 2) void sinkhorn_knopp_single_block_fus } __syncthreads(); + MHC_PROFILE_PHASE(TagCompute); + for (int iter = 0; iter < num_iters; iter++) { for (int r = threadIdx.x; r < M; r += BLOCK_SIZE) { float sum = 0.0f; @@ -417,11 +432,26 @@ __global__ __launch_bounds__(BLOCK_SIZE, 2) void sinkhorn_knopp_single_block_fus __syncthreads(); } + MHC_PROFILE_PHASE(TagStore); + for (int i = threadIdx.x; i < total_elems; i += BLOCK_SIZE) { out[i] = tile[i]; } + + MHC_PROFILE_END(); } +// @bench sinkhorn_knopp_forward +// @group: sinkhorn_knopp +// @title: Sinkhorn-Knopp Forward +// @configs: (M,N,num_iters) = +// [(32,32,5),(32,32,10),(32,32,20),(64,64,5),(64,64,10),(64,64,20),(128,128,5),(128,128,10),(128,128,20)] +// @in: inp float[M * N] random(0.1,1.1) +// @out: out float[M * N] +// @scalar: eps=1e-8f +// @flops: 4.0 * M * N * num_iters +// @profile: sinkhorn_knopp_forward_profiled M=64 N=64 num_iters=10 grid=1 timeline=true +// timeline_blocks=1 inline void sinkhorn_knopp_forward(float* out, const float* inp, int M, int N, int num_iters, float eps, cudaStream_t stream = nullptr) { constexpr int BLOCK_SIZE = 256; @@ -448,7 +478,7 @@ inline void sinkhorn_knopp_forward(float* out, const float* inp, int M, int N, i auto kernel = sinkhorn_knopp_single_block_kernel; cudaFuncSetAttribute(kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, smem_size); - kernel<<<1, BLOCK_SIZE, smem_size, stream>>>(out, inp, M, N, num_iters, eps); + kernel<<<1, BLOCK_SIZE, smem_size, stream>>>(out, inp, M, N, num_iters, eps, nullptr, 0); } else { constexpr int TILE_SIZE = 32; dim3 grid((N + TILE_SIZE - 1) / TILE_SIZE, (M + TILE_SIZE - 1) / TILE_SIZE); @@ -460,6 +490,17 @@ inline void sinkhorn_knopp_forward(float* out, const float* inp, int M, int N, i } } +// @bench sinkhorn_knopp_forward_fused_exp +// @group: sinkhorn_knopp +// @title: Fused Exp Version (as used in MHCLayer) +// @configs: (M,N,num_iters) = +// [(32,32,5),(32,32,10),(32,32,20),(64,64,5),(64,64,10),(64,64,20),(128,128,5),(128,128,10),(128,128,20)] +// @in: inp float[M * N] random(-0.01,0.01) +// @out: out float[M * N] +// @extra-buf: H_res_exp float[M * N] +// @scalar: eps=1e-8f +// @profile: sinkhorn_knopp_forward_fused_exp_profiled M=64 N=64 num_iters=10 grid=1 timeline=true +// timeline_blocks=1 inline void sinkhorn_knopp_forward_fused_exp(float* out, float* H_res_exp, const float* inp, int M, int N, int num_iters, float eps, cudaStream_t stream = nullptr) { @@ -485,7 +526,7 @@ inline void sinkhorn_knopp_forward_fused_exp(float* out, float* H_res_exp, const cudaLaunchKernelEx(&config, sinkhorn_knopp_single_block_fused_exp_kernel, out, - H_res_exp, inp, M, N, num_iters, eps); + H_res_exp, inp, M, N, num_iters, eps, (int64_t*)nullptr, 0); #else sinkhorn_knopp_single_block_fused_exp_kernel <<<1, BLOCK_SIZE, smem_size, stream>>>(out, H_res_exp, inp, M, N, num_iters, eps); @@ -511,9 +552,11 @@ inline void sinkhorn_knopp_forward_fused_exp(float* out, float* H_res_exp, const config.dynamicSmemBytes = smem_size; config.stream = stream; - cudaLaunchKernelEx(&config, kernel, out, H_res_exp, inp, M, N, num_iters, eps); + cudaLaunchKernelEx(&config, kernel, out, H_res_exp, inp, M, N, num_iters, eps, + (int64_t*)nullptr, 0); #else - kernel<<<1, BLOCK_SIZE, smem_size, stream>>>(out, H_res_exp, inp, M, N, num_iters, eps); + kernel<<<1, BLOCK_SIZE, smem_size, stream>>>(out, H_res_exp, inp, M, N, num_iters, eps, + nullptr, 0); #endif } else { fprintf(stderr, "sinkhorn_knopp_forward_fused_exp: M > 128 or N > 128 not supported\n"); @@ -690,10 +733,11 @@ __global__ __launch_bounds__(BLOCK_SIZE, 2) void sinkhorn_knopp_backward_checkpo } } -template +template __global__ __launch_bounds__(BLOCK_SIZE, 2) void sinkhorn_knopp_backward_kernel( float* __restrict__ d_inp, const float* __restrict__ grad, const float* __restrict__ M_out, - const float* __restrict__ M_inp, int N, int num_iters, float eps) { + const float* __restrict__ M_inp, int N, int num_iters, float eps, + int64_t* profiler_buf = nullptr, int max_entries = 0) { extern __shared__ float smem[]; float* d_tile = smem; float* row_buffer = smem + MAX_DIM * MAX_DIM; @@ -704,11 +748,16 @@ __global__ __launch_bounds__(BLOCK_SIZE, 2) void sinkhorn_knopp_backward_kernel( int total = N * N; + MHC_PROFILE_INIT(blockIdx.x); + MHC_PROFILE_START(TagLoad); + for (int i = threadIdx.x; i < total; i += BLOCK_SIZE) { d_tile[i] = grad[i]; } __syncthreads(); + MHC_PROFILE_PHASE(TagCompute); + for (int iter = num_iters - 1; iter >= 0; iter--) { for (int i = threadIdx.x; i < total; i += BLOCK_SIZE) { tile_fwd[i] = M_inp[i]; @@ -793,11 +842,28 @@ __global__ __launch_bounds__(BLOCK_SIZE, 2) void sinkhorn_knopp_backward_kernel( __syncthreads(); } + MHC_PROFILE_PHASE(TagStore); + for (int i = threadIdx.x; i < total; i += BLOCK_SIZE) { d_inp[i] = d_tile[i]; } + + MHC_PROFILE_END(); } +// @bench sinkhorn_knopp_backward +// @title: Sinkhorn-Knopp Backward +// @configs: (M,N,num_iters) = [(32,32,5),(32,32,10),(32,32,20),(64,64,5),(64,64,10),(64,64,20)] +// @in: M_inp float[M * N] random(0.1,1.1), grad float[M * N] random(-1,1,43) +// @out: d_inp float[M * N] +// @extra-buf: M_out float[M * N] +// @scalar: eps=1e-8f +// @flops: 8.0 * M * N * num_iters +// @bandwidth: 4 * M * N * sizeof(float) +// @setup: sinkhorn_knopp_forward(d_M_out, d_M_inp, M, N, num_iters, eps); +// @setup: CHECK_CUDA(cudaDeviceSynchronize()); +// @profile: sinkhorn_knopp_backward_profiled M=64 N=64 num_iters=10 grid=1 timeline=true +// timeline_blocks=1 inline void sinkhorn_knopp_backward(float* d_inp, const float* grad, const float* M_out, const float* M_inp, int N, int num_iters, float eps, cudaStream_t stream = nullptr) { @@ -1236,4 +1302,76 @@ inline void sinkhorn_knopp_backward_batched(float* d_inp, const float* grad, con } } +// Profiled dispatch variants — call production kernels with DO_PROFILE=true. + +inline void sinkhorn_knopp_forward_profiled(float* out, const float* inp, int M, int N, + int num_iters, float eps, int64_t* profiler_buf, + int max_entries, cudaStream_t stream = nullptr) { + constexpr int BLOCK_SIZE = 256; + + if (M <= 64 && N <= 64) { + constexpr int MAX_DIM = 64; + size_t smem_size = + MAX_DIM * MAX_DIM * sizeof(float) + MAX_DIM * sizeof(float) + MAX_DIM * sizeof(float); + + sinkhorn_knopp_single_block_kernel + <<<1, BLOCK_SIZE, smem_size, stream>>>(out, inp, M, N, num_iters, eps, profiler_buf, + max_entries); + } else if (M <= 128 && N <= 128) { + constexpr int MAX_DIM = 128; + size_t smem_size = + MAX_DIM * MAX_DIM * sizeof(float) + MAX_DIM * sizeof(float) + MAX_DIM * sizeof(float); + + auto kernel = sinkhorn_knopp_single_block_kernel; + cudaFuncSetAttribute(kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, smem_size); + + kernel<<<1, BLOCK_SIZE, smem_size, stream>>>(out, inp, M, N, num_iters, eps, profiler_buf, + max_entries); + } +} + +inline void sinkhorn_knopp_forward_fused_exp_profiled(float* out, float* H_res_exp, + const float* inp, int M, int N, int num_iters, + float eps, int64_t* profiler_buf, + int max_entries, + cudaStream_t stream = nullptr) { + constexpr int BLOCK_SIZE = 256; + + if (M <= 64 && N <= 64) { + constexpr int MAX_DIM = 64; + size_t smem_size = + MAX_DIM * MAX_DIM * sizeof(float) + MAX_DIM * sizeof(float) + MAX_DIM * sizeof(float); + + sinkhorn_knopp_single_block_fused_exp_kernel + <<<1, BLOCK_SIZE, smem_size, stream>>>(out, H_res_exp, inp, M, N, num_iters, eps, + profiler_buf, max_entries); + } else if (M <= 128 && N <= 128) { + constexpr int MAX_DIM = 128; + size_t smem_size = + MAX_DIM * MAX_DIM * sizeof(float) + MAX_DIM * sizeof(float) + MAX_DIM * sizeof(float); + + auto kernel = sinkhorn_knopp_single_block_fused_exp_kernel; + cudaFuncSetAttribute(kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, smem_size); + + kernel<<<1, BLOCK_SIZE, smem_size, stream>>>(out, H_res_exp, inp, M, N, num_iters, eps, + profiler_buf, max_entries); + } +} + +inline void sinkhorn_knopp_backward_profiled(float* d_inp, const float* grad, const float* M_out, + const float* M_inp, int N, int num_iters, float eps, + int64_t* profiler_buf, int max_entries, + cudaStream_t stream = nullptr) { + constexpr int BLOCK_SIZE = 256; + + if (N <= 64) { + constexpr int MAX_DIM = 64; + size_t smem_size = 2 * MAX_DIM * MAX_DIM * sizeof(float) + 4 * MAX_DIM * sizeof(float); + + sinkhorn_knopp_backward_kernel + <<<1, BLOCK_SIZE, smem_size, stream>>>(d_inp, grad, M_out, M_inp, N, num_iters, eps, + profiler_buf, max_entries); + } +} + } // namespace mhc diff --git a/src/csrc/kernels/stream_ops.cuh b/src/csrc/kernels/stream_ops.cuh index 2107157..a1764f1 100644 --- a/src/csrc/kernels/stream_ops.cuh +++ b/src/csrc/kernels/stream_ops.cuh @@ -322,6 +322,14 @@ class StreamMixTC { } }; +// @bench stream_aggregate_bf16_fused_sigmoid +// @group: stream_ops +// @title: stream_aggregate_bf16_fused_sigmoid +// @configs: (B,n,C) = [(320,4,1280),(512,4,1920),(1280,4,2560),(2560,4,1280)] +// @in: inp float[B * n * C] random(-1,1), H_pre_raw float[n] zero +// @out: out float[B * C] +// @extra-buf: H_pre_activated float[n] +// @bandwidth: (B * n * C + n) * sizeof(float) + B * C * sizeof(float) inline void stream_aggregate_bf16_fused_sigmoid(float* out, float* H_pre_activated, const float* inp, const float* H_pre_raw, int B, int n, int C, cudaStream_t stream = nullptr) { @@ -399,6 +407,16 @@ inline void stream_aggregate_bf16_fused_sigmoid(float* out, float* H_pre_activat } } +// @bench stream_distribute_mix_add_fused +// @group: stream_ops +// @title: stream_distribute_mix_add_fused +// @configs: (B,n,C) = [(320,4,1280),(512,4,1920),(1280,4,2560),(2560,4,1280)] +// @in: x_inp float[B * n * C] random(-1,1), y_norm float[B * C] random(-1,1,43), H_post_raw +// float[n] zero, M float[n * n] ones(1.0/n) +// @out: out float[B * n * C] +// @extra-buf: H_post_activated float[n] +// @bandwidth: (B * n * C + n * n + n) * sizeof(float) + B * C * sizeof(float) + B * n * C * +// sizeof(float) inline void stream_distribute_mix_add_fused(float* out, float* H_post_activated, const float* x_inp, const float* y_norm, const float* H_post_raw, const float* M, int B, int n, int C, @@ -910,6 +928,16 @@ __global__ __launch_bounds__(256, 4) void reduce_partials_kernel(float* __restri } } +// @bench stream_aggregate_backward +// @group: stream_ops_backward +// @title: stream_aggregate_backward +// @configs: (B,n,C) = [(320,4,1280),(512,4,1920),(1280,4,2560),(2560,4,1280)] +// @in: inp float[B * n * C] random(-1,1), H_pre float[n] zero, grad float[B * C] random(-1,1,43) +// @out: d_inp float[B * n * C], d_H_pre float[n] +// @pre-setup: constexpr int BLOCK = 256; +// @pre-setup: int workspace_num_blocks = std::min(128, (B * C + BLOCK - 1) / BLOCK); +// @extra-buf: workspace float[workspace_num_blocks * n] zero +// @bandwidth: (B * n * C + B * C + n) * sizeof(float) + (B * n * C + n) * sizeof(float) inline void stream_aggregate_backward(float* d_inp, float* d_H_pre, const float* grad, const float* inp, const float* H_pre, int B, int n, int C, float* workspace, int workspace_num_blocks, @@ -1134,6 +1162,19 @@ __global__ __launch_bounds__(256, 4) void reduce_partials_matrix_kernel( } } +// @bench stream_distribute_mix_backward_fused +// @group: stream_ops_backward +// @title: stream_distribute_mix_backward_fused +// @configs: (B,n,C) = [(320,4,1280),(512,4,1920),(1280,4,2560),(2560,4,1280)] +// @in: x float[B * n * C] random(-1,1), y_norm float[B * C] random(-1,1,43), M float[n * n] +// ones(1.0/n), H_post float[n] zero, grad float[B * n * C] random(-1,1,44) +// @out: d_x float[B * n * C], d_y_norm float[B * C], d_M float[n * n], d_H_post float[n] +// @pre-setup: constexpr int BLOCK = 256; +// @pre-setup: int workspace_num_blocks = std::min(128, (B * C + BLOCK - 1) / BLOCK); +// @extra-buf: workspace_M float[workspace_num_blocks * n * n] zero, workspace_H +// float[workspace_num_blocks * n] zero +// @bandwidth: (2 * B * n * C + B * C + n * n + n) * sizeof(float) + (B * n * C + B * C + n * n + n) +// * sizeof(float) inline void stream_distribute_mix_backward_fused(float* d_x, float* d_y_norm, float* d_M, float* d_H_post, const float* grad, const float* x, const float* y_norm, const float* M, diff --git a/src/python/benchmarks/bench_layer.py b/src/python/benchmarks/bench_layer.py index 445c415..c49f309 100644 --- a/src/python/benchmarks/bench_layer.py +++ b/src/python/benchmarks/bench_layer.py @@ -266,11 +266,6 @@ def main(): (512, 1920, 4), (1280, 2560, 4), (2560, 1280, 4), - (128, 1280, 8), - (256, 1280, 8), - (32, 1280, 32), - (64, 1280, 32), - (128, 1280, 32), ] else: configs = [(args.batch, args.hidden, args.expansion)] diff --git a/src/python/mhc/ops.py b/src/python/mhc/ops.py index 1432730..0f436e2 100644 --- a/src/python/mhc/ops.py +++ b/src/python/mhc/ops.py @@ -459,7 +459,7 @@ def mhc_layer_fused_dynamic_inference( float(alpha_res.item()) if torch.is_tensor(alpha_res) else float(alpha_res) ) x_f32 = x_expanded.float().contiguous() - H_pre_activated, H_post_activated, M = mhc_cuda.mhc_dynamic_h_fwd( + H_pre_activated, H_post_activated, M, _rms_h = mhc_cuda.mhc_dynamic_h_fwd( x_f32, phi_concat, alpha_pre_val,