From e7d1b582e8bba2ab53b8cba1698ebe55ae698fce Mon Sep 17 00:00:00 2001 From: yihonglie Date: Mon, 8 Jun 2026 20:46:04 -0500 Subject: [PATCH 1/2] Optimize JitFunction direct CallState dispatch --- python/flydsl/compiler/jit_function.py | 116 +++++++++++++++++------ tests/unit/test_jit_direct_call_state.py | 35 +++++++ 2 files changed, 124 insertions(+), 27 deletions(-) create mode 100644 tests/unit/test_jit_direct_call_state.py diff --git a/python/flydsl/compiler/jit_function.py b/python/flydsl/compiler/jit_function.py index a2332d711..5d3c37987 100644 --- a/python/flydsl/compiler/jit_function.py +++ b/python/flydsl/compiler/jit_function.py @@ -1344,21 +1344,48 @@ def _ensure_cache_manager(self, owner_cls=None): self.cache_manager = JitCacheManager(cache_dir) self.cache_manager.load_all() - def _resolve_and_make_cache_key(self, bound_args): - """Resolve raw call values into JitArgument instances *in place* and - build the tuple cache key from them. - - Side effect: entries in ``bound_args`` whose annotation is neither - ``Constexpr[T]`` nor ``Type[T]`` are replaced with their resolved - ``JitArgument`` instance (e.g. ``int`` → ``Int32``, - - * Annotation-driven (``Constexpr[T]`` / ``Type[T]``): value or type - baked directly into the key, ``bound_args`` left untouched. - * JitArgument-driven: the call value is (or is wrapped into) a - ``JitArgument`` and its ``cache_signature()`` is appended. + @staticmethod + def _raw_arg_cache_signature(arg, ann): + """Cache signature for a raw runtime argument without wrapping it. + + The hot launch path may receive plain Python values for annotated + runtime parameters, e.g. ``int`` for ``Int32`` or ``torch.cuda.Stream`` + for ``Stream``. Those values do not affect generated code, so use the + JitArgument type's value-independent signature directly instead of + constructing a fresh JitArgument on every launch. """ from .jit_argument import JitArgumentRegistry + if isinstance(arg, JitArgument): + return cache_signature(arg) + if hasattr(arg, "__cache_signature__"): + return arg.__cache_signature__() + + if ann is not inspect.Parameter.empty and isinstance(ann, type): + if issubclass(ann, JitArgument): + return (ann,) + if hasattr(ann, "__get_c_pointers__"): + return (ann,) + + if isinstance(arg, tuple): + return tuple(JitFunction._raw_arg_cache_signature(a, inspect.Parameter.empty) for a in arg) + if isinstance(arg, list): + return tuple(JitFunction._raw_arg_cache_signature(a, inspect.Parameter.empty) for a in arg) + + ctor, _ = JitArgumentRegistry.get(type(arg)) + if ctor is None: + raise TypeError( + f"{type(arg).__name__} is neither a JitArgument nor has a registered " + f"constructor; cannot derive cache signature." + ) + if hasattr(ctor, "raw_cache_signature"): + return ctor.raw_cache_signature(arg) + if isinstance(arg, (int, float, bool)) and isinstance(ctor, type): + return (ctor,) + return cache_signature(ctor(arg)) + + def _resolve_and_make_cache_key(self, bound_args): + """Build the tuple cache key without resolving runtime args in place.""" sig = self._sig # Re-read env vars on every call. key_parts = [("_env_", _cache_invalidating_env_values()), ("_target_", self._backend_target)] @@ -1377,21 +1404,7 @@ def _resolve_and_make_cache_key(self, bound_args): key_parts.append((name, arg)) continue - if isinstance(arg, JitArgument): - jit_arg = arg - elif isinstance(ann, type) and issubclass(ann, JitArgument): - jit_arg = ann(arg) - else: - ctor, _ = JitArgumentRegistry.get(type(arg)) - if ctor is None: - raise TypeError( - f"{name}: {type(arg).__name__} is neither a JitArgument nor has a registered " - f"constructor; cannot derive cache signature." - ) - jit_arg = ctor(arg) - - bound_args[name] = jit_arg - key_parts.append((name, cache_signature(jit_arg))) + key_parts.append((name, self._raw_arg_cache_signature(arg, ann))) return tuple(key_parts) @@ -1420,6 +1433,46 @@ def _build_full_cache_key(self, bound_arguments, *, owner_cls=None, bound_self=N cache_key = (("_self_type_", type(bound_self)),) + cache_key return cache_key + def _can_direct_call_state(self, args_tuple, *, bound_self=None) -> bool: + if bound_self is not None: + return False + for (_name, param), _arg in zip(self._sig.parameters.items(), args_tuple): + ann = param.annotation + if ann is not inspect.Parameter.empty: + if Constexpr.is_constexpr_annotation(ann) or is_type_param_annotation(ann): + return False + return True + + def _set_direct_call_state(self, args_tuple, state, *, bound_self=None): + if self._can_direct_call_state(args_tuple, bound_self=bound_self): + params = tuple(self._sig.parameters.values()) + names = tuple(p.name for p in params) + name_to_index = {name: idx for idx, name in enumerate(names)} + missing = inspect.Parameter.empty + defaults = tuple(p.default if p.default is not inspect.Parameter.empty else missing for p in params) + self._direct_call_state = (len(args_tuple), state, names, name_to_index, defaults, missing) + + def _direct_call_args_tuple(self, args, kwargs, direct): + expected_len, _state, _names, name_to_index, defaults, missing = direct + if not kwargs: + return args if len(args) == expected_len else None + if len(args) > expected_len: + return None + + values = list(defaults) + for idx, value in enumerate(args): + values[idx] = value + + for name, value in kwargs.items(): + idx = name_to_index.get(name) + if idx is None or idx < len(args): + return None + values[idx] = value + + if any(value is missing for value in values): + return None + return tuple(values) + @staticmethod def _cache_key_to_str(cache_key) -> str: """Convert tuple cache key to string for disk cache.""" @@ -1429,6 +1482,12 @@ def __call__(self, *args, **kwargs): if ir.Context.current is not None: return self.func(*args, **kwargs) + direct = getattr(self, "_direct_call_state", None) + if direct is not None: + direct_args = self._direct_call_args_tuple(args, kwargs, direct) + if direct_args is not None: + return direct[1](direct_args) + self._ensure_sig() bound_self = None @@ -1464,6 +1523,7 @@ def __call__(self, *args, **kwargs): if call_state is not None: if env.compile.compile_only: return None + self._set_direct_call_state(args_tuple, call_state, bound_self=bound_self) return call_state(args_tuple) # Normal path: check in-process cache first, then optional disk cache. @@ -1497,6 +1557,7 @@ def __call__(self, *args, **kwargs): state = None if state is not None: self._call_state_cache[cache_key] = state + self._set_direct_call_state(args_tuple, state, bound_self=bound_self) return state(args_tuple) # Fallback: run through DLPack (should not happen for static layout) @@ -1660,6 +1721,7 @@ def __call__(self, *args, **kwargs): state = None if state is not None: self._call_state_cache[cache_key] = state + self._set_direct_call_state(args_tuple, state, bound_self=bound_self) return state(args_tuple) # Fallback: run through DLPack diff --git a/tests/unit/test_jit_direct_call_state.py b/tests/unit/test_jit_direct_call_state.py new file mode 100644 index 000000000..a9a525f0e --- /dev/null +++ b/tests/unit/test_jit_direct_call_state.py @@ -0,0 +1,35 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2025 FlyDSL Project Contributors + +import flydsl.compiler as flyc +import flydsl.expr as fx +from flydsl.compiler.jit_function import JitFunction + + +def test_raw_arg_cache_signature_uses_annotation_type_for_scalar_runtime_args(): + assert JitFunction._raw_arg_cache_signature(7, fx.Int32) == (fx.Int32,) + + +def test_direct_call_state_reconstructs_kwargs_after_warmup(): + @flyc.jit + def launch(a: fx.Int32, b: fx.Int32 = 1): + pass + + launch._ensure_sig() + launch._set_direct_call_state((1, 2), object()) + direct = launch._direct_call_state + + assert launch._direct_call_args_tuple((3, 4), {}, direct) == (3, 4) + assert launch._direct_call_args_tuple((3,), {"b": 4}, direct) == (3, 4) + assert launch._direct_call_args_tuple((), {"a": 3, "b": 4}, direct) == (3, 4) + assert launch._direct_call_args_tuple((3,), {"a": 4}, direct) is None + + +def test_direct_call_state_skips_constexpr_launchers(): + @flyc.jit + def launch(a: fx.Int32, n: fx.Constexpr[int]): + pass + + launch._ensure_sig() + + assert not launch._can_direct_call_state((1, 2)) From 610d9f94ddb05c582268c3f1c87f74511389109d Mon Sep 17 00:00:00 2001 From: yihonglie Date: Tue, 9 Jun 2026 04:15:56 -0500 Subject: [PATCH 2/2] bench: add FlyDSL hotpath replay repro --- benchmarks/dsv4_hotpath_repro/.gitignore | 2 + benchmarks/dsv4_hotpath_repro/README.md | 114 +++ benchmarks/dsv4_hotpath_repro/baselines.md | 63 ++ .../bench_flydsl_hotpath.py | 738 ++++++++++++++++++ benchmarks/dsv4_hotpath_repro/run_matrix.sh | 123 +++ 5 files changed, 1040 insertions(+) create mode 100644 benchmarks/dsv4_hotpath_repro/.gitignore create mode 100644 benchmarks/dsv4_hotpath_repro/README.md create mode 100644 benchmarks/dsv4_hotpath_repro/baselines.md create mode 100755 benchmarks/dsv4_hotpath_repro/bench_flydsl_hotpath.py create mode 100755 benchmarks/dsv4_hotpath_repro/run_matrix.sh diff --git a/benchmarks/dsv4_hotpath_repro/.gitignore b/benchmarks/dsv4_hotpath_repro/.gitignore new file mode 100644 index 000000000..b45f55965 --- /dev/null +++ b/benchmarks/dsv4_hotpath_repro/.gitignore @@ -0,0 +1,2 @@ +results/ +__pycache__/ diff --git a/benchmarks/dsv4_hotpath_repro/README.md b/benchmarks/dsv4_hotpath_repro/README.md new file mode 100644 index 000000000..9c8ccd45e --- /dev/null +++ b/benchmarks/dsv4_hotpath_repro/README.md @@ -0,0 +1,114 @@ +# FlyDSL-only hotpath repro + +This directory contains a minimal repro for the DeepSeek-V4-Pro non-DPA +`1k/1k c=256` regression. It intentionally does not import AITER, ATOM, model +weights, or a serving stack. + +The benchmark builds small FlyDSL kernels with AITER-like launcher signatures +and replays the rank-local call mix observed in the original trace: + +| launcher kind | calls per prefill window | +| --- | ---: | +| `qk_norm_rope_quant_like` | 81 | +| `fused_compress_attn_like` | 80 | +| `hca_compress_forward_like` | 41 | +| `hca_norm_rope_scatter_like` | 41 | + +The GPU work inside each synthetic kernel is intentionally small. The signal +is host-side launch overhead after JIT compilation: signature binding, +cache-key construction, TensorAdaptor / PointerAdaptor handling, and CallState +reuse. This is the part that is amplified by the high call count in the real +serving trace. + +## Quick run + +Run one FlyDSL tree directly: + +```bash +PYTHONPATH=/path/to/FlyDSL/build-fly/python_packages:/path/to/FlyDSL/python:/path/to/FlyDSL \ +python benchmarks/dsv4_hotpath_repro/bench_flydsl_hotpath.py \ + --label original \ + --case dsv4-c256 \ + --tokens 991 \ + --windows 16 \ + --output results/original.json +``` + +Compare original FlyDSL and the PR/fixed FlyDSL: + +```bash +benchmarks/dsv4_hotpath_repro/run_matrix.sh \ + --original /path/to/flydsl-original \ + --fixed /path/to/flydsl-fixed +``` + +`--fixed` defaults to the current checkout, so from the fixed PR branch this is +usually enough: + +```bash +benchmarks/dsv4_hotpath_repro/run_matrix.sh \ + --original /path/to/flydsl-original +``` + +For a shorter smoke run: + +```bash +benchmarks/dsv4_hotpath_repro/run_matrix.sh \ + --original /path/to/flydsl-original \ + -- --windows 1 --gpu-event-calls 20 +``` + +## What to look at + +Each JSON result records the imported FlyDSL path, version, git head, GPU, ROCm +arch, call counts, per-kernel metrics, and mixed replay metrics. + +The most important fields are: + +- `per_kernel.qk.paths.jit_keyword_stream.host_wall_us_per_call` +- `per_kernel.qk.paths.compiled_positional.host_wall_us_per_call` +- `mixed_replay.jit.host_wall_us_per_call` +- `mixed_replay.compiled.host_wall_us_per_call` +- `gpu_event_us_per_call` + +Expected interpretation: + +- If `gpu_event_us_per_call` is similar but `jit_keyword_stream` is much slower + than `compiled_positional`, the problem is host launch overhead, not the GPU + kernel body. +- If the fixed FlyDSL lowers `jit_keyword_stream` and mixed replay time versus + original FlyDSL, the PR is addressing the right hotpath. +- `qk_norm_rope_quant_like` is the primary case. MoE can be enabled with + `--include-moe`, but the regression is not MoE-only. + +## Local smoke baseline + +A 1-window smoke run in the `yhl_dev` ROCm container on gfx950 produced: + +| stack | mixed jit host us/call | mixed compiled host us/call | qk jit host us/call | qk compiled host us/call | +| --- | ---: | ---: | ---: | ---: | +| original FlyDSL `0.2.0-pristine` | 104.23 | 10.09 | 68.68 | 8.63 | +| fixed `directraw-kw-patch` | 91.36 | 10.01 | 64.89 | 8.66 | + +Use these only as a direction baseline. Absolute values vary by host CPU, +Python build, ROCm stack, and GPU queue state. + +## Relation to the full E2E baseline + +Known local full-stack c=256 results: + +| stack | total tok/s | output tok/s | TPOT | TTFT | +| --- | ---: | ---: | ---: | ---: | +| old good `68a2d29` old stack | 8496.85 | 4250.48 | 58.04 ms | 719.77 ms | +| bad original `2984891` new stack | 7810.44 | 3907.11 | 63.32 ms | 798.97 ms | +| full hotpath fix | 8450.04 | 4227.06 | 58.30 ms | 768.43 ms | +| reviewer-safe nocachesig fix | 8225.67 | 4114.83 | 59.95 ms | 734.25 ms | + +This FlyDSL-only repro does not try to reproduce model math or full serving +throughput. It isolates the necessary FlyDSL-side condition: warm repeated +`@flyc.jit` launches with many tensor/pointer arguments must not rebuild the +same expensive host-side state on every call. + +The cache-signature semantics are not part of this repro requirement. The +script compares original behavior with the fixed hotpath behavior while keeping +the benchmark independent of AITER-side changes. diff --git a/benchmarks/dsv4_hotpath_repro/baselines.md b/benchmarks/dsv4_hotpath_repro/baselines.md new file mode 100644 index 000000000..7dd466894 --- /dev/null +++ b/benchmarks/dsv4_hotpath_repro/baselines.md @@ -0,0 +1,63 @@ +# Baselines + +## Full E2E c=256 background + +These are the local non-DPA DeepSeek-V4-Pro `1k/1k c=256` results that motivated +the FlyDSL-only repro. + +| stack | ATOM / deps | total tok/s | output tok/s | TPOT | TTFT | +| --- | --- | ---: | ---: | ---: | ---: | +| old good | `68a2d29`, old aiter, FlyDSL `0.1.9.dev599` | 8496.85 | 4250.48 | 58.04 ms | 719.77 ms | +| bad original | `2984891`, aiter `d988eaa`, FlyDSL `0.2.0` | 7810.44 | 3907.11 | 63.32 ms | 798.97 ms | +| full hotpath fix | FlyDSL + AITER qk/rawptr hotpath fixes | 8450.04 | 4227.06 | 58.30 ms | 768.43 ms | +| reviewer-safe nocachesig fix | FlyDSL direct-state only, no cache-sig change | 8225.67 | 4114.83 | 59.95 ms | 734.25 ms | + +## Short hotpath background + +Prior short c=256 hotpath checks: + +| stack | output tok/s | TPOT | +| --- | ---: | ---: | +| bad pristine | 4042.2 | 56.95 ms | +| direct raw only | 4115.7 | 55.96 ms | +| direct raw + kwargs | 4478.6 | 51.50 ms | +| direct raw + kwargs + qkcache | 4550.2 | 50.58 ms | +| direct raw + kwargs + qkcache + rawptr | 4616.4 | 49.75 ms | + +## FlyDSL-only repro baseline policy + +Record fresh results from this script in `results/` for the machine under test. +Absolute numbers depend on CPU, Python, ROCm, GPU, and FlyDSL build options. +The pass/fail signal is the relative gap: + +- original FlyDSL `jit_keyword_stream` vs fixed FlyDSL `jit_keyword_stream` +- `jit_keyword_stream` vs `compiled_positional` +- mixed replay original vs fixed + +For the regression to be considered reproduced, `gpu_event_us_per_call` should +stay in the same range while host wall time changes materially. + +## Local smoke baseline from this repro + +Command shape: + +```bash +python bench_flydsl_hotpath.py --windows 1 --warmup-windows 0 --gpu-event-calls 2 +``` + +Environment: + +| field | value | +| --- | --- | +| container | `yhl_dev` | +| GPU | `AMD Radeon Graphics` | +| arch | `gfx950` | +| original path | `local-c256-verify/flydsl-0.2.0-pristine` | +| fixed path | `local-c256-verify/flydsl-0.2.0-directraw-kw-patch` | + +| stack | qk jit | qk compiled | fused jit | fused compiled | hca compress jit | hca compress compiled | hca scatter jit | hca scatter compiled | mixed jit | mixed compiled | +| --- | ---: | ---: | ---: | ---: | ---: | ---: | ---: | ---: | ---: | ---: | +| original | 68.68 | 8.63 | 152.45 | 13.27 | 106.40 | 9.25 | 91.67 | 10.14 | 104.23 | 10.09 | +| fixed directraw-kw | 64.89 | 8.66 | 129.27 | 12.93 | 90.80 | 8.70 | 78.31 | 9.10 | 91.36 | 10.01 | + +All values are host wall `us/call`. diff --git a/benchmarks/dsv4_hotpath_repro/bench_flydsl_hotpath.py b/benchmarks/dsv4_hotpath_repro/bench_flydsl_hotpath.py new file mode 100755 index 000000000..0948cb0f5 --- /dev/null +++ b/benchmarks/dsv4_hotpath_repro/bench_flydsl_hotpath.py @@ -0,0 +1,738 @@ +#!/usr/bin/env python3 +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 FlyDSL Project Contributors + +"""FlyDSL-only hotpath replay for the DeepSeek-V4-Pro c=256 regression. + +This benchmark intentionally does not import AITER or ATOM. It builds small +FlyDSL kernels with AITER-like launcher signatures, then replays the same +high-frequency host call pattern seen in the non-DPA c=256 trace. + +The measured signal is host-side launch overhead after compilation: +JitFunction/cache-key/TensorAdaptor/CallState behavior, not model math. +""" + +import argparse +import json +import math +import os +import platform +import statistics +import subprocess +import sys +import time +from dataclasses import dataclass +from pathlib import Path +from typing import Callable + +import torch + +import flydsl +import flydsl.compiler as flyc +import flydsl.expr as fx + +try: + from flydsl.runtime.device import get_rocm_arch +except Exception: # pragma: no cover - older FlyDSL variants may not expose it + get_rocm_arch = None + + +BLOCK_DIM = 256 +VEC_WIDTH = 4 +TILE_ELEMS = BLOCK_DIM * VEC_WIDTH + +PROFILE_CALLS = { + # Rank-local mean calls per prefill window from the c=256 trace. + # These counts are used to reproduce the same host-launch pressure without + # requiring AITER, ATOM, model weights, or a serving process. + "dsv4-c256": { + "qk": 81, + "fused_compress": 80, + "hca_compress": 41, + "hca_scatter": 41, + "moe": 0, + } +} + + +@flyc.kernel +def _pointer_add_kernel( + a: fx.Pointer, + b: fx.Pointer, + out: fx.Pointer, + n: fx.Int32, +): + idx = fx.block_idx.x * fx.block_dim.x + fx.thread_idx.x + if idx < n: + out[idx] = a[idx] + b[idx] + + +@flyc.jit +def launch_qk_norm_rope_quant_like( + q_in: fx.Pointer, + kv_in: fx.Pointer, + q_weight: fx.Tensor, + kv_weight: fx.Tensor, + cos_cache: fx.Tensor, + sin_cache: fx.Tensor, + positions: fx.Pointer, + q_out: fx.Pointer, + kv_out: fx.Pointer, + q_scale: fx.Pointer, + kv_scale: fx.Pointer, + kv_in_row_stride: fx.Int32, + num_tokens: fx.Int32, + block_dim: fx.Constexpr[int], + stream: fx.Stream = fx.Stream(None), +): + grid_x = (num_tokens + block_dim - 1) // block_dim + _pointer_add_kernel(q_in, kv_in, q_out, num_tokens).launch( + grid=(grid_x, 1, 1), + block=(block_dim, 1, 1), + stream=stream, + ) + + +@flyc.kernel +def _tensor_add_kernel( + a: fx.Tensor, + b: fx.Tensor, + out: fx.Tensor, + block_dim: fx.Constexpr[int], + vec_width: fx.Constexpr[int], +): + bid = fx.block_idx.x + tid = fx.thread_idx.x + tile_elems = block_dim * vec_width + + ta = fx.logical_divide(a, fx.make_layout(tile_elems, 1)) + tb = fx.logical_divide(b, fx.make_layout(tile_elems, 1)) + tout = fx.logical_divide(out, fx.make_layout(tile_elems, 1)) + ta = fx.slice(ta, (None, bid)) + tb = fx.slice(tb, (None, bid)) + tout = fx.slice(tout, (None, bid)) + + ta = fx.logical_divide(ta, fx.make_layout(vec_width, 1)) + tb = fx.logical_divide(tb, fx.make_layout(vec_width, 1)) + tout = fx.logical_divide(tout, fx.make_layout(vec_width, 1)) + + copy_bits = vec_width * 32 + copy_atom = fx.make_copy_atom(fx.UniversalCopy(copy_bits), fx.Float32) + + ra = fx.make_rmem_tensor(vec_width, fx.Float32) + rb = fx.make_rmem_tensor(vec_width, fx.Float32) + rout = fx.make_rmem_tensor(vec_width, fx.Float32) + + fx.copy_atom_call(copy_atom, fx.slice(ta, (None, tid)), ra) + fx.copy_atom_call(copy_atom, fx.slice(tb, (None, tid)), rb) + fx.memref_store_vec(fx.arith.addf(fx.memref_load_vec(ra), fx.memref_load_vec(rb)), rout) + fx.copy_atom_call(copy_atom, rout, fx.slice(tout, (None, tid))) + + +@flyc.jit +def launch_fused_compress_attn_like( + kv_in: fx.Tensor, + kv_in_row_stride: fx.Int32, + score_in: fx.Tensor, + score_in_row_stride: fx.Int32, + plan: fx.Tensor, + kv_state: fx.Tensor, + kv_state_slot_stride: fx.Int32, + kv_state_pos_stride: fx.Int32, + score_state: fx.Tensor, + score_state_slot_stride: fx.Int32, + score_state_pos_stride: fx.Int32, + state_slot_mapping: fx.Tensor, + ape: fx.Tensor, + rms_weight: fx.Tensor, + cos_cache: fx.Tensor, + sin_cache: fx.Tensor, + kv_cache: fx.Tensor, + kv_cache_block_stride: fx.Int32, + kv_cache_token_stride: fx.Int32, + cache_scale: fx.Tensor, + cache_scale_block_stride: fx.Int32, + block_table: fx.Tensor, + block_table_seq_stride: fx.Int32, + plan_capacity: fx.Int32, + block_dim: fx.Constexpr[int], + vec_width: fx.Constexpr[int], + stream: fx.Stream = fx.Stream(None), +): + grid_x = (plan_capacity + block_dim * vec_width - 1) // (block_dim * vec_width) + _tensor_add_kernel(kv_in, score_in, kv_state, block_dim, vec_width).launch( + grid=(grid_x, 1, 1), + block=(block_dim, 1, 1), + stream=stream, + ) + + +@flyc.jit +def launch_hca_compress_forward_like( + kv_in: fx.Tensor, + kv_in_row_stride: fx.Int32, + score_in: fx.Tensor, + score_in_row_stride: fx.Int32, + plan: fx.Tensor, + kv_state: fx.Tensor, + kv_state_slot_stride: fx.Int32, + kv_state_pos_stride: fx.Int32, + score_state: fx.Tensor, + score_state_slot_stride: fx.Int32, + score_state_pos_stride: fx.Int32, + state_slot_mapping: fx.Tensor, + ape: fx.Tensor, + kv_compressed: fx.Tensor, + kv_compressed_row_stride: fx.Int32, + plan_capacity: fx.Int32, + block_dim: fx.Constexpr[int], + vec_width: fx.Constexpr[int], + stream: fx.Stream = fx.Stream(None), +): + grid_x = (plan_capacity + block_dim * vec_width - 1) // (block_dim * vec_width) + _tensor_add_kernel(kv_in, score_in, kv_compressed, block_dim, vec_width).launch( + grid=(grid_x, 1, 1), + block=(block_dim, 1, 1), + stream=stream, + ) + + +@flyc.jit +def launch_hca_norm_rope_scatter_like( + kv_compressed: fx.Tensor, + kv_compressed_row_stride: fx.Int32, + plan: fx.Tensor, + rms_weight: fx.Tensor, + cos_cache: fx.Tensor, + sin_cache: fx.Tensor, + kv_cache: fx.Tensor, + kv_cache_block_stride: fx.Int32, + kv_cache_token_stride: fx.Int32, + block_table: fx.Tensor, + block_table_seq_stride: fx.Int32, + plan_capacity: fx.Int32, + block_dim: fx.Constexpr[int], + vec_width: fx.Constexpr[int], + stream: fx.Stream = fx.Stream(None), +): + grid_x = (plan_capacity + block_dim * vec_width - 1) // (block_dim * vec_width) + _tensor_add_kernel(kv_compressed, kv_cache, kv_cache, block_dim, vec_width).launch( + grid=(grid_x, 1, 1), + block=(block_dim, 1, 1), + stream=stream, + ) + + +@flyc.jit +def launch_moe_like( + hidden: fx.Tensor, + weight0: fx.Tensor, + weight1: fx.Tensor, + expert_ids: fx.Tensor, + topk_weight: fx.Tensor, + out: fx.Tensor, + num_tokens: fx.Int32, + block_dim: fx.Constexpr[int], + vec_width: fx.Constexpr[int], + stream: fx.Stream = fx.Stream(None), +): + grid_x = (num_tokens + block_dim * vec_width - 1) // (block_dim * vec_width) + _tensor_add_kernel(hidden, weight0, out, block_dim, vec_width).launch( + grid=(grid_x, 1, 1), + block=(block_dim, 1, 1), + stream=stream, + ) + + +@dataclass +class LauncherSpec: + name: str + fn: Callable + args: tuple + args_without_stream: tuple + + +def _round_up(x: int, multiple: int) -> int: + return ((x + multiple - 1) // multiple) * multiple + + +def _percentile(values: list[float], pct: float) -> float: + if not values: + return float("nan") + ordered = sorted(values) + idx = int(math.ceil((pct / 100.0) * len(ordered))) - 1 + return ordered[min(max(idx, 0), len(ordered) - 1)] + + +def _bench_host(fn: Callable, calls: int, warmup_calls: int) -> dict: + for _ in range(warmup_calls): + fn() + torch.cuda.synchronize() + + samples = [] + t0 = time.perf_counter_ns() + for _ in range(calls): + s = time.perf_counter_ns() + fn() + e = time.perf_counter_ns() + samples.append((e - s) / 1000.0) + t1 = time.perf_counter_ns() + torch.cuda.synchronize() + t2 = time.perf_counter_ns() + + return { + "calls": calls, + "host_wall_total_ms": (t1 - t0) / 1e6, + "host_wall_us_per_call": (t1 - t0) / calls / 1000.0 if calls else float("nan"), + "host_return_mean_us": statistics.mean(samples) if samples else float("nan"), + "host_return_p50_us": statistics.median(samples) if samples else float("nan"), + "host_return_p95_us": _percentile(samples, 95.0), + "host_return_p99_us": _percentile(samples, 99.0), + "host_return_max_us": max(samples) if samples else float("nan"), + "sync_tail_ms": (t2 - t1) / 1e6, + } + + +def _bench_gpu_event(fn: Callable, calls: int, warmup_calls: int) -> dict: + start = torch.cuda.Event(enable_timing=True) + end = torch.cuda.Event(enable_timing=True) + for _ in range(warmup_calls): + fn() + torch.cuda.synchronize() + start.record() + for _ in range(calls): + fn() + end.record() + torch.cuda.synchronize() + total_us = start.elapsed_time(end) * 1000.0 + return { + "calls": calls, + "gpu_event_total_us": total_us, + "gpu_event_us_per_call": total_us / calls if calls else float("nan"), + } + + +def _git_head(path: str) -> str | None: + try: + out = subprocess.check_output( + ["git", "-C", path, "rev-parse", "--short", "HEAD"], + stderr=subprocess.DEVNULL, + text=True, + ) + return out.strip() + except Exception: + return None + + +def _make_ptr(dtype, tensor): + return flyc.from_c_void_p(dtype, tensor.data_ptr()) + + +def build_specs(tokens: int, stream_arg) -> dict[str, LauncherSpec]: + n = max(1, tokens) + tensor_elems = _round_up(max(tokens, TILE_ELEMS), TILE_ELEMS) + max_pos = 8192 + + q = torch.randn(n, device="cuda", dtype=torch.float32) + kv = torch.randn(n, device="cuda", dtype=torch.float32) + q_out = torch.empty_like(q) + kv_out = torch.empty_like(kv) + q_scale = torch.empty(n, device="cuda", dtype=torch.float32) + kv_scale = torch.empty(n, device="cuda", dtype=torch.float32) + positions = torch.arange(n, device="cuda", dtype=torch.int32) + + q_weight = torch.randn(16, 512, device="cuda", dtype=torch.bfloat16) + kv_weight = torch.randn(512, device="cuda", dtype=torch.bfloat16) + cos_cache = torch.randn(max_pos, 32, device="cuda", dtype=torch.bfloat16) + sin_cache = torch.randn(max_pos, 32, device="cuda", dtype=torch.bfloat16) + + qk_args = ( + _make_ptr(fx.Float32, q), + _make_ptr(fx.Float32, kv), + q_weight, + kv_weight, + cos_cache, + sin_cache, + _make_ptr(fx.Int32, positions), + _make_ptr(fx.Float32, q_out), + _make_ptr(fx.Float32, kv_out), + _make_ptr(fx.Float32, q_scale), + _make_ptr(fx.Float32, kv_scale), + 512, + n, + BLOCK_DIM, + stream_arg, + ) + + kv_in = torch.randn(tensor_elems, device="cuda", dtype=torch.float32) + score_in = torch.randn(tensor_elems, device="cuda", dtype=torch.float32) + plan = torch.arange(tensor_elems, device="cuda", dtype=torch.int32) + kv_state = torch.empty(tensor_elems, device="cuda", dtype=torch.float32) + score_state = torch.empty(tensor_elems, device="cuda", dtype=torch.float32) + state_slot_mapping = torch.arange(tensor_elems, device="cuda", dtype=torch.int32) + ape = torch.randn(tensor_elems, device="cuda", dtype=torch.float32) + rms_weight = torch.randn(tensor_elems, device="cuda", dtype=torch.float32) + kv_cache = torch.empty(tensor_elems, device="cuda", dtype=torch.float32) + cache_scale = torch.empty(tensor_elems, device="cuda", dtype=torch.float32) + block_table = torch.arange(tensor_elems, device="cuda", dtype=torch.int32) + kv_compressed = torch.empty(tensor_elems, device="cuda", dtype=torch.float32) + + fused_args = ( + kv_in, + 512, + score_in, + 128, + plan, + kv_state, + 2048, + 512, + score_state, + 1024, + 256, + state_slot_mapping, + ape, + rms_weight, + cos_cache, + sin_cache, + kv_cache, + 4096, + 256, + cache_scale, + 64, + block_table, + 512, + tensor_elems, + BLOCK_DIM, + VEC_WIDTH, + stream_arg, + ) + + hca_compress_args = ( + kv_in, + 512, + score_in, + 128, + plan, + kv_state, + 2048, + 512, + score_state, + 1024, + 256, + state_slot_mapping, + ape, + kv_compressed, + 512, + tensor_elems, + BLOCK_DIM, + VEC_WIDTH, + stream_arg, + ) + + hca_scatter_args = ( + kv_compressed, + 512, + plan, + rms_weight, + cos_cache, + sin_cache, + kv_cache, + 4096, + 256, + block_table, + 512, + tensor_elems, + BLOCK_DIM, + VEC_WIDTH, + stream_arg, + ) + + moe_args = ( + kv_in, + score_in, + ape, + state_slot_mapping, + rms_weight, + kv_cache, + tensor_elems, + BLOCK_DIM, + VEC_WIDTH, + stream_arg, + ) + + return { + "qk": LauncherSpec("qk", launch_qk_norm_rope_quant_like, qk_args, qk_args[:-1]), + "fused_compress": LauncherSpec( + "fused_compress", + launch_fused_compress_attn_like, + fused_args, + fused_args[:-1], + ), + "hca_compress": LauncherSpec( + "hca_compress", + launch_hca_compress_forward_like, + hca_compress_args, + hca_compress_args[:-1], + ), + "hca_scatter": LauncherSpec( + "hca_scatter", + launch_hca_norm_rope_scatter_like, + hca_scatter_args, + hca_scatter_args[:-1], + ), + "moe": LauncherSpec("moe", launch_moe_like, moe_args, moe_args[:-1]), + } + + +def _make_invokers(spec: LauncherSpec, stream_arg, call_style: str) -> dict[str, Callable]: + invokers = {} + if call_style in ("positional", "both"): + invokers["jit_positional"] = lambda spec=spec: spec.fn(*spec.args) + if call_style in ("keyword-stream", "both"): + invokers["jit_keyword_stream"] = lambda spec=spec: spec.fn(*spec.args_without_stream, stream=stream_arg) + return invokers + + +def _compile_spec(spec: LauncherSpec) -> tuple[Callable | None, float | None, str | None]: + t0 = time.perf_counter_ns() + try: + compiled = flyc.compile(spec.fn, *spec.args) + torch.cuda.synchronize() + t1 = time.perf_counter_ns() + return compiled, (t1 - t0) / 1000.0, None + except Exception as exc: + return None, None, repr(exc) + + +def bench_spec( + spec: LauncherSpec, + calls: int, + warmup_calls: int, + stream_arg, + call_style: str, + gpu_event_calls: int, +) -> dict: + compiled, compile_once_us, compile_error = _compile_spec(spec) + result = { + "name": spec.name, + "calls": calls, + "compile_once_us": compile_once_us, + "compile_error": compile_error, + "paths": {}, + } + if compile_error: + return result + + invokers = _make_invokers(spec, stream_arg, call_style) + if compiled is not None: + invokers["compiled_positional"] = lambda compiled=compiled, spec=spec: compiled(*spec.args) + + for path_name, fn in invokers.items(): + host = _bench_host(fn, calls, warmup_calls) + gpu = _bench_gpu_event(fn, min(calls, gpu_event_calls), min(warmup_calls, gpu_event_calls)) + merged = dict(host) + merged.update(gpu) + result["paths"][path_name] = merged + + return result + + +def bench_mixed( + specs: dict[str, LauncherSpec], + calls_by_kernel: dict[str, int], + warmup_windows: int, + stream_arg, + call_style: str, + use_compiled: bool, +) -> dict: + compiled = {} + if use_compiled: + for name, spec in specs.items(): + if calls_by_kernel.get(name, 0) <= 0: + continue + compiled_fn, _, err = _compile_spec(spec) + if err: + return {"error": f"compile failed for {name}: {err}"} + compiled[name] = compiled_fn + + def call_one(name: str): + spec = specs[name] + if use_compiled: + compiled[name](*spec.args) + elif call_style == "keyword-stream": + spec.fn(*spec.args_without_stream, stream=stream_arg) + else: + spec.fn(*spec.args) + + warmup_counts = {name: count * warmup_windows for name, count in calls_by_kernel.items()} + for name, count in warmup_counts.items(): + for _ in range(count): + call_one(name) + torch.cuda.synchronize() + + total_calls = sum(calls_by_kernel.values()) + t0 = time.perf_counter_ns() + for name, count in calls_by_kernel.items(): + for _ in range(count): + call_one(name) + t1 = time.perf_counter_ns() + torch.cuda.synchronize() + t2 = time.perf_counter_ns() + + path = "compiled_positional" if use_compiled else f"jit_{call_style}" + return { + "path": path, + "calls_by_kernel": calls_by_kernel, + "total_calls": total_calls, + "host_wall_total_ms": (t1 - t0) / 1e6, + "host_wall_us_per_call": (t1 - t0) / total_calls / 1000.0 if total_calls else float("nan"), + "sync_tail_ms": (t2 - t1) / 1e6, + } + + +def collect_env(label: str | None) -> dict: + flydsl_root = str(Path(flydsl.__file__).resolve().parents[1]) + return { + "label": label, + "python": sys.version.split()[0], + "platform": platform.platform(), + "torch_version": torch.__version__, + "torch_cuda": getattr(torch.version, "cuda", None), + "torch_hip": getattr(torch.version, "hip", None), + "gpu": torch.cuda.get_device_name(0) if torch.cuda.is_available() else None, + "rocm_arch": get_rocm_arch() if get_rocm_arch is not None else None, + "flydsl_file": str(Path(flydsl.__file__).resolve()), + "flydsl_version": getattr(flydsl, "__version__", None), + "flydsl_git_head": _git_head(flydsl_root), + "flydsl_runtime_cache_dir": os.environ.get("FLYDSL_RUNTIME_CACHE_DIR"), + "flydsl_runtime_enable_cache": os.environ.get("FLYDSL_RUNTIME_ENABLE_CACHE"), + } + + +def parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--label", default=None, help="Label stored in the JSON result.") + parser.add_argument("--case", default="dsv4-c256", choices=sorted(PROFILE_CALLS)) + parser.add_argument("--tokens", type=int, default=991) + parser.add_argument("--windows", type=int, default=16) + parser.add_argument("--warmup-windows", type=int, default=1) + parser.add_argument("--kernels", default="qk,fused_compress,hca_compress,hca_scatter") + parser.add_argument("--include-moe", action="store_true") + parser.add_argument("--call-style", choices=["positional", "keyword-stream", "both"], default="both") + parser.add_argument("--gpu-event-calls", type=int, default=200) + parser.add_argument("--output", default=None) + return parser.parse_args() + + +def main() -> int: + args = parse_args() + if not torch.cuda.is_available(): + print("CUDA/ROCm GPU is not available.", file=sys.stderr) + return 2 + + stream_arg = torch.cuda.current_stream().cuda_stream + specs = build_specs(args.tokens, stream_arg) + profile = dict(PROFILE_CALLS[args.case]) + if args.include_moe: + profile["moe"] = 16 + + wanted = [k.strip() for k in args.kernels.split(",") if k.strip()] + if args.include_moe and "moe" not in wanted: + wanted.append("moe") + unknown = sorted(set(wanted) - set(specs)) + if unknown: + raise ValueError(f"unknown kernels: {unknown}") + + calls_by_kernel = { + name: profile.get(name, 0) * args.windows + for name in wanted + if profile.get(name, 0) > 0 + } + warmup_by_kernel = { + name: profile.get(name, 0) * args.warmup_windows + for name in wanted + if profile.get(name, 0) > 0 + } + + result = { + "env": collect_env(args.label), + "profile": { + "case": args.case, + "tokens": args.tokens, + "windows": args.windows, + "warmup_windows": args.warmup_windows, + "source": "rank-local c=256 trace call counts", + "calls_per_window": profile, + "calls_by_kernel": calls_by_kernel, + }, + "per_kernel": {}, + "mixed_replay": {}, + } + + print(f"FlyDSL: {result['env']['flydsl_file']}") + print(f"GPU: {result['env']['gpu']} arch={result['env']['rocm_arch']}") + print(f"calls_by_kernel: {calls_by_kernel}") + + for name in wanted: + calls = calls_by_kernel.get(name, 0) + if calls <= 0: + continue + print(f"\n[{name}] calls={calls}") + spec_result = bench_spec( + specs[name], + calls, + warmup_by_kernel.get(name, 0), + stream_arg, + args.call_style, + args.gpu_event_calls, + ) + result["per_kernel"][name] = spec_result + for path_name, metrics in spec_result.get("paths", {}).items(): + print( + f" {path_name:<22s} " + f"host={metrics['host_wall_us_per_call']:.2f} us/call " + f"gpu={metrics['gpu_event_us_per_call']:.2f} us/call" + ) + + if calls_by_kernel: + mixed_style = "keyword-stream" if args.call_style in ("keyword-stream", "both") else "positional" + result["mixed_replay"]["jit"] = bench_mixed( + specs, + calls_by_kernel, + args.warmup_windows, + stream_arg, + mixed_style, + use_compiled=False, + ) + result["mixed_replay"]["compiled"] = bench_mixed( + specs, + calls_by_kernel, + args.warmup_windows, + stream_arg, + mixed_style, + use_compiled=True, + ) + print("\n[mixed]") + for name, metrics in result["mixed_replay"].items(): + if "error" in metrics: + print(f" {name}: {metrics['error']}") + else: + print( + f" {name:<8s} {metrics['path']:<24s} " + f"host={metrics['host_wall_us_per_call']:.2f} us/call " + f"total={metrics['host_wall_total_ms']:.2f} ms" + ) + + if args.output: + out = Path(args.output) + out.parent.mkdir(parents=True, exist_ok=True) + out.write_text(json.dumps(result, indent=2, sort_keys=True) + "\n") + print(f"\nwrote {out}") + else: + print(json.dumps(result, indent=2, sort_keys=True)) + + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/benchmarks/dsv4_hotpath_repro/run_matrix.sh b/benchmarks/dsv4_hotpath_repro/run_matrix.sh new file mode 100755 index 000000000..2443d6b3d --- /dev/null +++ b/benchmarks/dsv4_hotpath_repro/run_matrix.sh @@ -0,0 +1,123 @@ +#!/usr/bin/env bash +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 FlyDSL Project Contributors + +set -euo pipefail + +SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" +ROOT_DIR="$(cd "${SCRIPT_DIR}/../.." && pwd)" + +ORIGINAL="${FLYDSL_ORIGINAL:-}" +FIXED="${FLYDSL_FIXED:-${ROOT_DIR}}" +OUT_DIR="${SCRIPT_DIR}/results" +EXTRA_ARGS=() + +pythonpath_for_repo() { + local repo="$1" + local path="" + + if [[ -d "${repo}/build-fly/python_packages" ]]; then + path="${repo}/build-fly/python_packages" + elif [[ -d "${repo}/build/python_packages" ]]; then + path="${repo}/build/python_packages" + fi + + if [[ -d "${repo}/python" ]]; then + path="${path:+${path}:}${repo}/python" + fi + + path="${path:+${path}:}${repo}" + printf '%s' "${path}" +} + +while [[ $# -gt 0 ]]; do + case "$1" in + --original) + ORIGINAL="$2" + shift 2 + ;; + --fixed) + FIXED="$2" + shift 2 + ;; + --out-dir) + OUT_DIR="$2" + shift 2 + ;; + --) + shift + EXTRA_ARGS+=("$@") + break + ;; + *) + EXTRA_ARGS+=("$1") + shift + ;; + esac +done + +run_one() { + local label="$1" + local repo="$2" + local out="${OUT_DIR}/${label}.json" + + if [[ ! -d "${repo}" ]]; then + echo "missing FlyDSL repo for ${label}: ${repo}" >&2 + exit 2 + fi + + mkdir -p "${OUT_DIR}/cache-${label}" + echo + echo "=== ${label} ===" + echo "repo: ${repo}" + echo "out: ${out}" + + PYTHONPATH="$(pythonpath_for_repo "${repo}")${PYTHONPATH:+:${PYTHONPATH}}" \ + FLYDSL_RUNTIME_CACHE_DIR="${OUT_DIR}/cache-${label}" \ + python "${SCRIPT_DIR}/bench_flydsl_hotpath.py" \ + --label "${label}" \ + --output "${out}" \ + "${EXTRA_ARGS[@]}" +} + +if [[ -z "${ORIGINAL}" ]]; then + echo "missing original FlyDSL path; pass --original /path/to/original or set FLYDSL_ORIGINAL" >&2 + exit 2 +fi + +mkdir -p "${OUT_DIR}" +run_one original "${ORIGINAL}" +run_one fixed "${FIXED}" + +python - "${OUT_DIR}/original.json" "${OUT_DIR}/fixed.json" <<'PY' +import json +import sys +from pathlib import Path + +orig = json.loads(Path(sys.argv[1]).read_text()) +fixed = json.loads(Path(sys.argv[2]).read_text()) + +def get(obj, path, default=None): + cur = obj + for key in path: + if not isinstance(cur, dict) or key not in cur: + return default + cur = cur[key] + return cur + +print("\n=== summary ===") +for kernel in sorted(orig.get("per_kernel", {})): + o = get(orig, ["per_kernel", kernel, "paths", "jit_keyword_stream", "host_wall_us_per_call"]) + f = get(fixed, ["per_kernel", kernel, "paths", "jit_keyword_stream", "host_wall_us_per_call"]) + oc = get(orig, ["per_kernel", kernel, "paths", "compiled_positional", "host_wall_us_per_call"]) + fc = get(fixed, ["per_kernel", kernel, "paths", "compiled_positional", "host_wall_us_per_call"]) + if o is not None and f is not None: + print(f"{kernel:<16s} jit original={o:8.2f} us fixed={f:8.2f} us speedup={o / f:5.2f}x") + if oc is not None and fc is not None: + print(f"{'':<16s} cmp original={oc:8.2f} us fixed={fc:8.2f} us speedup={oc / fc:5.2f}x") + +om = get(orig, ["mixed_replay", "jit", "host_wall_us_per_call"]) +fm = get(fixed, ["mixed_replay", "jit", "host_wall_us_per_call"]) +if om is not None and fm is not None: + print(f"mixed jit original={om:8.2f} us fixed={fm:8.2f} us speedup={om / fm:5.2f}x") +PY