Skip to content

[Ascend] feat: update do_bench_npu #669

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
winter77-x wants to merge 1 commit into
base: triton_v3.2.x
Choose a base branch
from
Open

[Ascend] feat: update do_bench_npu #669

Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
76 changes: 23 additions & 53 deletions third_party/ascend/backend/testing.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -28,16 +28,18 @@
from triton._C.clear_l2 import do_bench_clear


def do_bench_npu(funcs, warmup=25, active=100, prof_dir=None, clear_l2_cache=False, keep_res=False, collect_prof=True,
return_mode="mean", quantiles=None):
def do_bench_npu(funcs, warmup=5, active=30, prof_dir=None, clear_l2_cache=False, keep_res=False):
import torch
import torch_npu

assert return_mode in ["min", "max", "mean", "median", "all"]

if not isinstance(funcs, list):
funcs = [funcs]

# warmup kernel
for fn in funcs:
fn()
torch.npu.synchronize()

experimental_config = torch_npu.profiler._ExperimentalConfig(
aic_metrics=torch_npu.profiler.AiCMetrics.PipeUtilization,
profiler_level=torch_npu.profiler.ProfilerLevel.Level1,
Expand All @@ -52,74 +54,42 @@ def do_bench_npu(funcs, warmup=25, active=100, prof_dir=None, clear_l2_cache=Fal
pid = process.pid
process_name = process.name
timestamp = datetime.now(tz=timezone.utc).strftime("%Y%m%d_%H%M%S")
base_path = os.path.join(runtime.cache.get_home_dir(), ".triton", "profile_results")
base_path = os.path.join(
runtime.cache.get_home_dir(), ".triton", "profile_results"
)
torch_path = os.path.join(base_path, f"prof_{timestamp}_{process_name}-{pid}")

if clear_l2_cache:
device = triton.runtime.driver.active.get_current_device()
stream = triton.runtime.driver.active.get_current_stream(device)
buffer = runtime.driver.active.get_empty_cache_for_benchmark()
do_bench_clear(buffer.data_ptr(), buffer.numel(), stream)
buffer = buffer.float() # to avoid type cast
buffer.sum()

@neoskypig neoskypig Jun 8, 2026

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

usually，we use a CPU "wrapper" to call a triton-kernel，it is probably that using the "torch.sum" in the wrapper function， and it will cause the time is wrong because this "torch.sum" time is captured by msprof。

@winter77-x winter77-x Jun 10, 2026

Copy link
Copy Markdown
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

In the do_bench_npu function, buffer.sum() is used to forcefully clear the NPU’s L2 cache before each iteration of the benchmark (when clear_l2_cache=True).

torch.npu.synchronize() # shake out of any npu error

# cal warmup num
di = runtime.driver.active.get_device_interface()
for fn in funcs:
fn()
di.synchronize()

cache = runtime.driver.active.get_empty_cache_for_benchmark()
start_event = di.Event(enable_timing=True)
end_event = di.Event(enable_timing=True)
start_event.record()
for fn in funcs:
for _ in range(5):
cache.zero_()
fn()
di.synchronize()
end_event.record()
estimate_ms = start_event.elapsed_time(end_event) / 5

n_warmup = min(5, max(1, int(warmup / estimate_ms)))
n_repeat = min(30, max(1, int(active / estimate_ms)))

# cal warmup num
total = n_warmup + n_repeat
print(f"total={total}, n_warmup={n_warmup}, n_repeat={n_repeat}")

# Run for 300 μs to raise the frequency to 800.
mat_a = torch.randn(4096, 4096).to(dtype=torch.bfloat16).npu()
mat_b = torch.randn(4096, 4096).to(dtype=torch.bfloat16).npu()
mat_c = torch.matmul(mat_a, mat_b)

@neoskypig neoskypig Jun 8, 2026

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

according to experience, it needs a ">= 300us" kernel which make NPU works on a high frequency。
this method is used by some open-source libraries, e.g. bytedance's "mojoOpset".

@winter77-x winter77-x Jun 10, 2026

Copy link
Copy Markdown
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The warmup in do_bench_npu is implemented in two stages:

  1. Pre‑warmup (1 call per function) – ensures the runtime environment is “ready” without being measured
  2. Profiled warmup (warmup calls per function) – allows the device to reach a steady state, but these calls are excluded from the final results.

mat_c.cpu()
total = warmup + active
with torch_npu.profiler.profile(
activities=[torch_npu.profiler.ProfilerActivity.NPU],
on_trace_ready=torch_npu.profiler.tensorboard_trace_handler(torch_path),
record_shapes=False,
profile_memory=False,
with_stack=False,
with_flops=False,
with_modules=False,
experimental_config=experimental_config,
activities=[torch_npu.profiler.ProfilerActivity.NPU],
on_trace_ready=torch_npu.profiler.tensorboard_trace_handler(torch_path),
record_shapes=False,
profile_memory=False,
with_stack=False,
with_flops=False,
with_modules=False,
experimental_config=experimental_config,
) as prof:
for fn in funcs:
for _ in builtins.range(total):
if clear_l2_cache:
do_bench_clear(buffer.data_ptr(), buffer.numel(), stream)
buffer.sum() # use buffer read to clear l2 cache
torch.npu.synchronize()
fn()
torch.npu.synchronize()
if clear_l2_cache:
del buffer
if collect_prof:
time_cost = _collect_prof_result(torch_path, funcs, n_warmup, n_repeat, return_mode=return_mode,
quantiles=quantiles) # read kernel_details.csv
else:
time_cost = _collect_single(torch_path, return_mode=return_mode) # read op_static.csv

time_cost = _collect_prof_result(torch_path, funcs, warmup, active)
_rm_dic(keep_res, torch_path)
return time_cost


# keep the original behavior to get the statistics for the specified kernel func
def _collect_single(base_dir: str, key: str = None, print_flag=True, return_mode="mean") -> float:
if not os.path.exists(base_dir):
Expand Down
Loading