flashinfer.testing

此模块为 FlashInfer 中的基准测试和性能分析提供全面的测试工具。

测试环境设置

set_seed(random_seed)	在测试期间设置随机种子以实现可重复性。
sleep_after_kernel_run(execution_time)	内核运行后休眠。

性能分析

FLOPS 计算

attention_flops(batch_size, qo_seqlen, ...)	计算给定注意力层的 FLOPs。
attention_flops_with_actual_seq_lens(...)	计算给定注意力层的 FLOPs，其中提供实际序列长度作为 1D 张量。
attention_tflops_per_sec(batch_size, ...)	计算给定注意力层的每秒 TFLOPS。
attention_tflops_per_sec_with_actual_seq_lens(...)	计算给定注意力层的每秒 TFLOPS，其中包含实际序列长度。

吞吐量分析

attention_tb_per_sec(batch_size, qo_seqlen, ...)	计算给定注意力层实现的每秒 TB。
attention_tb_per_sec_with_actual_seq_lens(...)	计算给定注意力层实现的每秒 TB，其中包含实际序列长度。

GPU 基准测试

bench_gpu_time(fn[, dry_run_iters, ...])	具有可配置计时后端的统一 GPU 基准测试接口。
bench_gpu_time_with_cuda_event(fn[, ...])	使用 CUDA 事件（无 CUDA 图）基准测试内核执行时间。
bench_gpu_time_with_cudagraph(fn[, ...])	使用摊销内核启动开销的 CUDA 图基准测试 GPU 时间。
bench_gpu_time_with_cupti(fn[, ...])	使用 CUPTI 活动跟踪进行精确的内核计时基准测试 GPU 时间。