diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index d93673410f..02aa92f993 100644 --- a/.pre-commit-config.yaml +++ b/.pre-commit-config.yaml @@ -17,6 +17,8 @@ repos: rev: v1.14.1 hooks: - id: mypy + args: [--ignore-missing-imports] + additional_dependencies: [types-PyYAML, types-tabulate] files: \.py$ - repo: https://github.com/timothycrosley/isort diff --git a/benchmark/kvbench/README.md b/benchmark/kvbench/README.md index a2eeb5e431..19fe20627b 100644 --- a/benchmark/kvbench/README.md +++ b/benchmark/kvbench/README.md @@ -1,7 +1,8 @@ # NIXL KVBench -A utility for generating NIXL Bench commands that test KVCache transfer across various LLM architectures and access patterns (including block and layer approaches). +A comprehensive utility for generating NIXL Bench commands that test KVCache transfer across various LLM architectures and access patterns, plus custom traffic performance testing for asymmetric traffic flows. ## Table of Contents +- [Overview](#overview) - [Supported LLM Architectures](#supported-llm-architectures) - [Building](#building) - [Docker](#docker) @@ -13,17 +14,22 @@ A utility for generating NIXL Bench commands that test KVCache transfer across v - [CLI Override Arguments](#cli-override-arguments) - [Plan Command Arguments](#plan-command-arguments) - [Shared Benchmark Arguments](#shared-benchmark-arguments) + - [CTP Command Arguments](#ctp-command-arguments) - [Command Descriptions](#command-descriptions) - - [Plan Command](#plan-command) - - [Profile Command](#profile-command) - - [KVCache Command](#kvcache-command) - - [IOSize Command](#iosize-command) -- [Examples](#example-deepseek-r1-with-block-access-tp1-pp16) - - [DeepSeek R1 with Block Access](#example-deepseek-r1-with-block-access-tp1-pp16) - - [DeepSeek R1 with Layer Access](#example-deepseek-r1-with-layer-access-tp1-pp16) - - [Overriding Model Configuration](#example-overriding-model-configuration-with-cli-args) + - [KVBench Commands](#kvbench-commands) + - [CTP Commands](#ctp-commands) +- [Examples](#examples) + - [KVBench Examples](#kvbench-examples) + - [CTP Examples](#ctp-examples) - [Developer Guides](#developer-guides) +## Overview + +NIXL KVBench provides two main categories of functionality: + +1. **KVBench Commands**: Test KV cache transfers across various LLM architectures with different access patterns (block and layer approaches) +2. **CTP Commands**: Custom Traffic Performance Testing for measuring asymmetric traffic patterns using transfer matrices + ## Supported LLM Architectures - DeepSeek R1 - LLama 3.1 @@ -48,34 +54,34 @@ source venv/bin/activate pip install uv uv sync --active ``` + ## Usage ### Basic Usage ```bash python main.py --help -usage: main.py [-h] {plan,profile,kvcache,iosize} ... +Usage: main.py [OPTIONS] COMMAND [ARGS]... -KVBench + KVBench - NIXL Performance Testing CLI -positional arguments: - {plan,profile,kvcache,iosize} - Available commands - plan Display the recommended configuration for nixlbench - profile Run nixlbench - kvcache Display kvcache information - iosize Display io size information +Options: + --debug / --no-debug Enable debug logging + --help Show this message and exit. -options: - -h, --help show this help message and exit +Commands: + ct-perftest Run custom traffic performance test using patterns... + io-size Display IO size information + kvcache Display kvcache information + plan Display the recommended configuration for nixlbench + profile Run nixlbench + sequential-ct-perftest Run sequential custom traffic performance test... ``` ## Command Line Arguments -KVBench supports various argument groups that apply to different commands: - ### Common Arguments -These arguments are shared across all KVBench commands (plan, kvcache, profile): +These arguments are shared across KVBench commands (plan, kvcache, profile, io-size): | Argument | Description | | -------- | ----------- | @@ -139,11 +145,22 @@ These arguments are used by both `plan` and `profile` commands: | `--gds_filepath` | File path for GDS operations | | `--enable_vmm` | Enable VMM memory allocation when DRAM is requested | +### CTP Command Arguments + +Specific to CTP (Custom Traffic Performance) commands: + +| Argument | Description | +| -------- | ----------- | +| `config_file` | Path to YAML configuration file (required) | +| `--verify-buffers / --no-verify-buffers` | Verify buffer contents after transfer (default: False) | +| `--print-recv-buffers / --no-print-recv-buffers` | Print received buffer contents (default: False) | +| `--json-output-path` | Path to save JSON output (sequential-ct-perftest only) | + ## Command Descriptions -KVBench provides four main commands: +### KVBench Commands -### Plan Command +#### Plan Command The `plan` command generates and displays recommended `nixlbench` command configurations based on your model architecture and parameters. It helps you prepare optimal benchmark settings without running the benchmark itself. @@ -151,7 +168,7 @@ The `plan` command generates and displays recommended `nixlbench` command config python main.py plan --model ./examples/model_deepseek_r1.yaml --model_config "./examples/block-tp1-pp8.yaml" ``` -### Profile Command +#### Profile Command The `profile` command actually runs the benchmark with the specified configuration using `nixlbench`, collecting performance data across various KV cache operations and access patterns. @@ -159,7 +176,7 @@ The `profile` command actually runs the benchmark with the specified configurati python main.py profile --model ./examples/model_deepseek_r1.yaml --model_config "./examples/block-tp1-pp8.yaml" ``` -### KVCache Command +#### KVCache Command The `kvcache` command analyzes and displays detailed information about the KV cache for a specified model configuration, including model type, sequence lengths, batch sizes, and I/O sizes. @@ -171,23 +188,55 @@ Batch Size : 298 IO Size : 1.12 MB ``` -### IOSize Command +#### IO-Size Command -The `iosize` command displays information about the I/O size requirements for a specified model configuration, helping you understand memory usage and data transfer needs. +The `io-size` command displays information about the I/O size requirements for a specified model configuration, helping you understand memory usage and data transfer needs. ```bash -python main.py iosize --model ./examples/model_deepseek_r1.yaml --model_config "./examples/block-tp1-pp8.yaml" -IO Size per GPU: 4608 +python main.py io-size --model ./examples/model_deepseek_r1.yaml --model_config "./examples/block-tp1-pp8.yaml" +Model: DEEPSEEK_R1 +Page Size: 256.0 B +Input Sequence Length: 10000 +Batch Size: 298 +IO Size: 1.12 MB +``` + +### CTP Commands + +#### Sequential CT Perftest + +Benchmark the performance of a continuum of traffic patterns one after the other. Before running each pattern, all ranks do a barrier, optionally sleep for a given amount of time, then run the pattern and measure execution time. + +**Reports**: Sequential CT Perftest reports the total latency per matrix execution, along with their latency when run isolated, which can be used to evaluate how well the network reacts to congestion. + +``` + Transfer size (GB) Latency (ms) Isolated Latency (ms) Num Senders +-------------------- -------------- ----------------------- ------------- + 4.945 35.047 35.421 4 + 3.230 21.152 21.800 4 + 1.104 8.222 8.280 4 + ... ... ... ... + 0.129 2.147 2.386 4 ``` -## Example: DeepSeek R1 with Block Access (TP=1, PP=16) +#### CT Perftest + +Benchmark the performance of one traffic pattern. The pattern is run in multiple iterations and then metrics are reported. Useful for optimizing specific patterns. + +**Reports**: CT Perftest reports total latency (time elapsed between the first rank started until the last rank finished), average time per iteration, total size sent over the network, and average bandwidth by rank. + +## Examples + +### KVBench Examples + +#### DeepSeek R1 with Block Access (TP=1, PP=16) ```bash python main.py plan \ --model ./examples/model_deepseek_r1.yaml \ --model_config ./examples/block-tp1-pp16.yaml \ --backend GDS \ --source gpu \ - --etcd-endpoint "http://localhost:2379" + --etcd-endpoints "http://localhost:2379" ================================================================================ Model Config: ./examples/block-tp1-pp16.yaml ISL: 10000 tokens @@ -205,14 +254,14 @@ nixlbench \ --target_seg_type VRAM ``` -## Example: DeepSeek R1 with Layer Access (TP=1, PP=16) +#### DeepSeek R1 with Layer Access (TP=1, PP=16) ```bash python main.py plan \ --model ./examples/model_deepseek_r1.yaml \ --model_config ./examples/layer-tp1-pp16.yaml \ --backend GDS \ --source gpu \ - --etcd-endpoint "http://localhost:2379" + --etcd-endpoints "http://localhost:2379" ================================================================================ Model Config: ./examples/layer-tp1-pp16.yaml ISL: 10000 tokens @@ -230,14 +279,14 @@ nixlbench \ --target_seg_type VRAM ``` -## Example: Overriding model configuration with cli args +#### Overriding Model Configuration with CLI Args ```bash python main.py plan \ --model ./examples/model_deepseek_r1.yaml \ --model_config ./examples/block-tp1-pp8.yaml \ --backend GDS \ --source gpu \ - --etcd-endpoint "http://localhost:2379" \ + --etcd-endpoints "http://localhost:2379" \ --pp 32 \ --num_requests 100 ================================================================================ @@ -257,6 +306,126 @@ nixlbench \ --target_seg_type VRAM ``` +### CTP Examples + +#### Configuration Files + +CTP tests are defined using YAML configuration files. + +**CT Perftest Configuration** (single traffic pattern): +```yaml +iters: 50 +warmup_iters: 10 +traffic_pattern: + matrix_file: "/path/to/matrix.txt" + shards: 1 + mem_type: "cuda" + xfer_op: "WRITE" +``` + +**Sequential CT Perftest Configuration** (multiple traffic patterns): +```yaml +traffic_patterns: +- matrix_file: /path/to/matrices/matrix_0.txt + metadata: + isl: 38328 + sleep_after_launch_sec: 16.480753057792 +- matrix_file: /path/to/matrices/matrix_1.txt + metadata: + isl: 25034 + sleep_after_launch_sec: 71.875102179328 +``` + +**Traffic Pattern Parameters**: +- `matrix_file`: File containing the transfer matrix (required) +- `shards`: Number of chunks to shard the buffer into (default: 1) +- `mem_type`: Memory type, currently supports "cuda" (default: "cuda") +- `xfer_op`: Transfer operation, "READ" or "WRITE" (default: "WRITE") +- `sleep_after_launch_sec`: Seconds to sleep before running pattern (default: 0) + +**Matrix File Format**: +Matrix cells are separated by whitespaces and contain either bytes as integers or standard units (K, M, G). + +Example matrix file: +``` +0 0 1M 1M +0 0 1M 1M +0 0 0 5K +0 0 0 5K +``` + +#### Generate Traffic Pattern Matrices + +Optionally, generate matrices using the inference workload matrix generation tool: +```bash +python test/inference_workload_matgen.py generate \ + --num-user-requests 10 \ + --num-prefill-nodes 16 \ + --num-decode-nodes 16 \ + --prefill-tp 8 \ + --prefill-pp 1 \ + --prefill-cp 1 \ + --decode-tp 8 \ + --decode-pp 1 \ + --decode-cp 1 \ + --results-dir $PWD/matrices \ + --model llama-405b +``` + +#### Running CTP Tests + +**Sequential CT Perftest**: +```bash +# Basic usage +python main.py sequential-ct-perftest ./config.yaml + +# With options +python main.py sequential-ct-perftest ./config.yaml \ + --verify-buffers \ + --json-output-path ./results.json + +# With debug logging +python main.py --debug sequential-ct-perftest ./config.yaml \ + --verify-buffers \ + --json-output-path ./results.json + +# With Slurm +srun python main.py sequential-ct-perftest ./config.yaml \ + --verify-buffers \ + --json-output-path ./results.json +``` + +**CT Perftest**: +```bash +# Basic usage +python main.py ct-perftest ./config.yaml + +# With options +python main.py ct-perftest ./config.yaml \ + --verify-buffers \ + --print-recv-buffers + +# With debug logging +python main.py --debug ct-perftest ./config.yaml \ + --verify-buffers +``` + +## Implementation Details + +### CTP Implementation +- Custom traffic patterns defined using transfer matrices where cell [i,j] defines message size from rank i to rank j +- `benchmark.kvbench.test.custom_traffic_perftest.py` implements CT Perftest and TrafficPattern dataclass +- `benchmark.kvbench.test.sequential_custom_traffic_perftest.py` implements Sequential CT Perftest +- Utilizes common utilities for distributed testing support +- `benchmark.kvbench.test.traffic_pattern.py` defines abstraction for traffic patterns + +### Known Issues +- The nixl xfer preparation currently takes significant time (the `_prepare_tp()` method) + +### Next Steps +- [ ] Support more memory types beyond CUDA +- [ ] Optimize transfer preparation performance + ## Developer Guides For more detailed information, please refer to the following documentation: - [Tutorial with GDS](docs/tutorial-gds.md) - Quick tutorial for running NIXLBench with GDS diff --git a/benchmark/kvbench/__init__.py b/benchmark/kvbench/__init__.py new file mode 100644 index 0000000000..3159bfe656 --- /dev/null +++ b/benchmark/kvbench/__init__.py @@ -0,0 +1,14 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. diff --git a/benchmark/kvbench/commands/__init__.py b/benchmark/kvbench/commands/__init__.py index 96d87d46fb..3159bfe656 100644 --- a/benchmark/kvbench/commands/__init__.py +++ b/benchmark/kvbench/commands/__init__.py @@ -12,11 +12,3 @@ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. -from commands import io_size, kvcache, plan, profile - -available_commands = [ - plan.Command(), - profile.Command(), - kvcache.Command(), - io_size.Command(), -] diff --git a/benchmark/kvbench/commands/args.py b/benchmark/kvbench/commands/args.py index ba75ff3768..3ddebf0f0b 100644 --- a/benchmark/kvbench/commands/args.py +++ b/benchmark/kvbench/commands/args.py @@ -13,148 +13,185 @@ # See the License for the specific language governing permissions and # limitations under the License. -import argparse +import click -def add_common_args(subparser: argparse.ArgumentParser): - subparser.add_argument("--model", type=str, help="Model name (e.g., 'llama3.1-8b')") - subparser.add_argument( - "--model_config", type=str, help="Path to a single model config YAML file" +def common_args(func): + """Decorator for common model arguments""" + func = click.option("--model", type=str, help="Model name (e.g., 'llama3.1-8b')")( + func ) - subparser.add_argument( + func = click.option( + "--model_config", type=str, help="Path to a single model config YAML file" + )(func) + func = click.option( "--model_configs", type=str, help="Path to multiple model config YAML files (supports glob patterns like 'configs/*.yaml')", - ) + )(func) + return func -def add_cli_args(subparser: argparse.ArgumentParser): - subparser.add_argument("--pp", type=int, help="Pipeline parallelism size") - subparser.add_argument("--tp", type=int, help="Tensor parallelism size") - subparser.add_argument("--isl", type=int, help="Input sequence length") - subparser.add_argument("--osl", type=int, help="Output sequence length") - subparser.add_argument("--num_requests", type=int, help="Number of requests") - subparser.add_argument("--page_size", type=int, help="Page size") - subparser.add_argument("--access_pattern", type=str, help="Access pattern") +def cli_args(func): + """Decorator for CLI-specific arguments""" + func = click.option("--pp", type=int, help="Pipeline parallelism size")(func) + func = click.option("--tp", type=int, help="Tensor parallelism size")(func) + func = click.option("--isl", type=int, help="Input sequence length")(func) + func = click.option("--osl", type=int, help="Output sequence length")(func) + func = click.option("--num_requests", type=int, help="Number of requests")(func) + func = click.option("--page_size", type=int, help="Page size")(func) + func = click.option("--access_pattern", type=str, help="Access pattern")(func) + return func -def add_plan_args(subparser: argparse.ArgumentParser): - subparser.add_argument( +def plan_args(func): + """Decorator for plan-specific arguments""" + func = click.option( "--format", default="text", type=str, help="Output of the nixl command [text, json, csv] (default: text)", - ) + )(func) + return func -def add_nixl_bench_args(subparser: argparse.ArgumentParser): - subparser.add_argument( +def nixl_bench_args(func): + """Decorator for NIXL benchmark arguments""" + func = click.option( "--source", default="file", type=str, help="Source of the nixl descriptors [file, memory, gpu] (default: file)", - ) - subparser.add_argument( + )(func) + func = click.option( "--destination", default="memory", type=str, help="Destination of the nixl descriptors [file, memory, gpu] (default: memory)", - ) - subparser.add_argument( + )(func) + func = click.option( "--backend", type=str, help="Communication backend [UCX, UCX_MO] (default: UCX)" - ) - subparser.add_argument( + )(func) + func = click.option( "--worker_type", type=str, help="Worker to use to transfer data [nixl, nvshmem] (default: nixl)", - ) - subparser.add_argument( + )(func) + func = click.option( "--initiator_seg_type", type=str, help="Memory segment type for initiator [DRAM, VRAM] (default: DRAM)", - ) - subparser.add_argument( + )(func) + func = click.option( "--target_seg_type", type=str, help="Memory segment type for target [DRAM, VRAM] (default: DRAM)", - ) - subparser.add_argument( + )(func) + func = click.option( "--scheme", type=str, help="Communication scheme [pairwise, manytoone, onetomany, tp] (default: pairwise)", - ) - subparser.add_argument( + )(func) + func = click.option( "--mode", type=str, help="Process mode [SG (Single GPU per proc), MG (Multi GPU per proc)] (default: SG)", - ) - subparser.add_argument( + )(func) + func = click.option( "--op_type", type=str, help="Operation type [READ, WRITE] (default: WRITE)" - ) - subparser.add_argument( - "--check_consistency", action="store_true", help="Enable consistency checking" - ) - subparser.add_argument( + )(func) + func = click.option( + "--check_consistency", is_flag=True, help="Enable consistency checking" + )(func) + func = click.option( "--total_buffer_size", type=int, help="Total buffer size (default: 8GiB)" - ) - subparser.add_argument( + )(func) + func = click.option( "--start_block_size", type=int, help="Starting block size (default: 4KiB)" - ) - subparser.add_argument( + )(func) + func = click.option( "--max_block_size", type=int, help="Maximum block size (default: 64MiB)" - ) - subparser.add_argument( + )(func) + func = click.option( "--start_batch_size", type=int, help="Starting batch size (default: 1)" - ) - subparser.add_argument( + )(func) + func = click.option( "--max_batch_size", type=int, help="Maximum batch size (default: 1)" - ) - subparser.add_argument( + )(func) + func = click.option( "--num_iter", type=int, help="Number of iterations (default: 1000)" - ) - subparser.add_argument( + )(func) + func = click.option( "--warmup_iter", type=int, help="Number of warmup iterations (default: 100)" - ) - subparser.add_argument( + )(func) + func = click.option( "--num_threads", type=int, help="Number of threads used by benchmark (default: 1)", - ) - subparser.add_argument( + )(func) + func = click.option( "--num_initiator_dev", type=int, help="Number of devices in initiator processes (default: 1)", - ) - subparser.add_argument( + )(func) + func = click.option( "--num_target_dev", type=int, help="Number of devices in target processes (default: 1)", + )(func) + func = click.option("--enable_pt", is_flag=True, help="Enable progress thread")( + func ) - subparser.add_argument( - "--enable_pt", action="store_true", help="Enable progress thread" - ) - subparser.add_argument( + func = click.option( "--device_list", type=str, help="Comma-separated device names (default: all)" - ) - subparser.add_argument( + )(func) + func = click.option( "--runtime_type", type=str, help="Type of runtime to use [ETCD] (default: ETCD)" - ) - subparser.add_argument( + )(func) + func = click.option( "--etcd-endpoints", type=str, help="ETCD server URL for coordination (default: http://localhost:2379)", - ) - subparser.add_argument( + )(func) + func = click.option( "--storage_enable_direct", type=bool, help="Enable direct I/O for storage operations (only used with POSIX backend)", default=False, - ) - subparser.add_argument( - "--gds_filepath", type=str, help="(File path for GDS operations" - ) - subparser.add_argument( - "--enable_vmm", - action="store_true", + )(func) + func = click.option( + "--gds_filepath", type=str, help="File path for GDS operations" + )(func) + func = click.option( + "--enable-vmm", + type=bool, help="Enable VMM memory allocation when DRAM is requested", - ) + )(func) + return func + + +def ctp_args(func): + """Decorator for CTP (Custom Traffic Pattern) arguments""" + func = click.option( + "--verify-buffers/--no-verify-buffers", + default=False, + help="Verify buffer contents after transfer", + )(func) + func = click.option( + "--print-recv-buffers/--no-print-recv-buffers", + default=False, + help="Print received buffer contents", + )(func) + func = click.option( + "--json-output-path", + type=click.Path(), + help="Path to save JSON output", + default=None, + )(func) + func = click.option( + "--config-file", + type=click.Path(exists=True), + required=True, + help="Path to YAML config file", + )(func) + return func diff --git a/benchmark/kvbench/commands/io_size.py b/benchmark/kvbench/commands/io_size.py deleted file mode 100644 index a7cd5bbbb2..0000000000 --- a/benchmark/kvbench/commands/io_size.py +++ /dev/null @@ -1,75 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. - -import argparse - -from commands.args import add_common_args -from models.model_config import ModelConfig -from models.models import BaseModelArch - - -class Command: - """ - Command handler for the 'iops' subcommand. - - This command calculates and displays the number of IOPs for a given model and configuration. - """ - - def __init__(self): - """ - Initialize the iops command. - - Sets the command name and help text for the command-line interface. - """ - self.name = "iosize" - self.help = "Display io size information" - - def add_arguments( - self, subparser: argparse.ArgumentParser - ) -> argparse.ArgumentParser: - """ - Add command-specific arguments to the argument parser. - - Args: - subparser (argparse.ArgumentParser): The parser for this command. - - Returns: - argparse.ArgumentParser: The updated argument parser with added arguments. - """ - add_common_args(subparser) - return subparser - - def execute(self, args: argparse.Namespace): - """ - Execute the iops command with the provided arguments. - - Loads the model architecture and configuration from the specified files, - calculates the number of IOPs for a given model and configuration, and - displays the information in a formatted output. - """ - if not args.model or not args.model_config: - print("Error: --model and --model_config are required") - return -1 - - # Load model architecture - model = BaseModelArch.from_yaml(args.model, None) - - # Load model configuration - model_config = ModelConfig.from_yaml(args.model_config) - - # Set model_config on the model instance using the new method - model.set_model_config(model_config) - io_size = model.get_io_size() - print(f"IO Size per GPU: {io_size}") diff --git a/benchmark/kvbench/commands/kvcache.py b/benchmark/kvbench/commands/kvcache.py deleted file mode 100644 index 7c9667b75c..0000000000 --- a/benchmark/kvbench/commands/kvcache.py +++ /dev/null @@ -1,104 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. - -import argparse - -from commands.args import add_cli_args, add_common_args -from models.model_config import ModelConfig -from models.models import BaseModelArch -from models.utils import get_batch_size, override_yaml_args - - -class Command: - """ - Command handler for the 'kvcache' subcommand. - - This command analyzes and displays key-value cache size information for the - specified model architecture and configuration, including per-token size, - total size, and related metrics. - """ - - def __init__(self): - """ - Initialize the kvcache command. - - Sets the command name and help text for the command-line interface. - """ - self.name = "kvcache" - self.help = "Display kvcache information" - - def add_arguments( - self, subparser: argparse.ArgumentParser - ) -> argparse.ArgumentParser: - """ - Add command-specific arguments to the argument parser. - - Args: - subparser (argparse.ArgumentParser): The parser for this command. - - Returns: - argparse.ArgumentParser: The updated argument parser with added arguments. - """ - add_common_args(subparser) - add_cli_args(subparser) - return subparser - - def execute(self, args: argparse.Namespace): - """ - Execute the kvcache command with the provided arguments. - - Loads the model architecture and configuration from the specified files, - calculates KV cache metrics (per-token size, total size, page size), and - displays the information in a formatted output. - - Args: - args (argparse.Namespace): Command-line arguments. - - Returns: - int: -1 if required arguments are missing, otherwise None. - """ - if not args.model or not args.model_config: - print("Error: --model and --model_config are required") - return -1 - - # Load model architecture - model = BaseModelArch.from_yaml(args.model, None) - - # Load model configuration - model_config = ModelConfig.from_yaml(args.model_config) - override_yaml_args(model_config, args) - # Set model_config on the model instance using the new method - model.set_model_config(model_config) - - from math import floor, log - - def format_bytes(size): - power = 0 if size <= 0 else floor(log(size, 1024)) - return f"{round(size / 1024 ** power, 2)} {['B', 'KB', 'MB', 'GB', 'TB'][int(power)]}" - - labels = [ - "KV Cache Size Per Token", - "IO Size", - "Batch Size", - "Model", - "Input Sequence Length", - ] - max_width = max(len(label) for label in labels) - io_size = model.get_io_size(model_config.system.page_size) - batch_size = get_batch_size(model, model_config, io_size) - print(f"{'Model':{max_width}}: {model.model}") - print(f"{'Input Sequence Length':{max_width}}: {model_config.runtime.isl}") - print(f"{'Batch Size':{max_width}}: {batch_size}") - print(f"{'IO Size':{max_width}}: {format_bytes(io_size)}") diff --git a/benchmark/kvbench/commands/plan.py b/benchmark/kvbench/commands/plan.py deleted file mode 100644 index bb741277de..0000000000 --- a/benchmark/kvbench/commands/plan.py +++ /dev/null @@ -1,214 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. - -import argparse -import csv -import glob -import io -import json -import os - -from commands.args import ( - add_cli_args, - add_common_args, - add_nixl_bench_args, - add_plan_args, -) -from commands.nixlbench import NIXLBench -from models.model_config import ModelConfig -from models.models import BaseModelArch -from models.utils import get_batch_size, override_yaml_args - - -class Command: - """ - Command handler for the 'plan' subcommand. - - This command displays the recommended configuration for nixlbench based on - the provided model architecture and model configuration files, showing both - ISL (Input Sequence Length) and OSL (Output Sequence Length) versions. - """ - - def __init__(self): - """ - Initialize the plan command. - - Sets the command name and help text for the command-line interface. - """ - self.name = "plan" - self.help = "Display the recommended configuration for nixlbench" - - def add_arguments( - self, subparser: argparse.ArgumentParser - ) -> argparse.ArgumentParser: - """ - Add command-specific arguments to the argument parser. - - Args: - subparser (argparse.ArgumentParser): The parser for this command. - - Returns: - argparse.ArgumentParser: The updated argument parser with added arguments. - """ - add_common_args(subparser) - add_plan_args(subparser) - add_nixl_bench_args(subparser) - add_cli_args(subparser) - return subparser - - def execute(self, args: argparse.Namespace): - """ - Execute the plan command with the provided arguments. - - Loads the model architecture and configuration from the specified files, - creates NIXLBench instances for both ISL and OSL configurations, and - generates nixlbench command plans for both sequence types. - - Args: - args (argparse.Namespace): Command-line arguments. - - Returns: - int: -1 if required arguments are missing, otherwise None. - """ - if not args.model: - print("Error: --model is required") - return -1 - - if not args.model_config and not args.model_configs: - print("Error: either --model_config or --model_configs is required") - return -1 - - # Load model architecture - model = BaseModelArch.from_yaml(args.model, None) - - # Get list of model config files - config_files = [] - - if args.model_config: - config_files.append(args.model_config) - - if args.model_configs: - # Expand glob patterns into list of files - expanded_files = glob.glob(args.model_configs) - if not expanded_files: - print(f"Warning: No files matched the pattern: {args.model_configs}") - config_files.extend(expanded_files) - - if not config_files: - print("Error: No valid model config files specified") - return -1 - - # Filter out duplicate paths - config_files = list(dict.fromkeys(config_files)) - - filtered_args = { - k: v for k, v in args.__dict__.items() if k in NIXLBench.defaults() - } - - # Process each model config - all_plans = [] - for config_file in config_files: - # Skip if file doesn't exist - if not os.path.exists(config_file): - print(f"Warning: Config file not found: {config_file}") - continue - - try: - # Load model configuration - model_config = ModelConfig.from_yaml(config_file) - # override yaml args with cli args if supplied - # i.e. if --pp is supplied, use it to override the pp size in the yaml file - override_yaml_args(model_config, args) - model.set_model_config(model_config) - - separator = "=" * 80 - isl_nixl_bench = NIXLBench(model, model_config, **filtered_args) - - io_size = model.get_io_size(model_config.system.page_size) - batch_size = get_batch_size(model, model_config, io_size) - isl_nixl_bench.set_io_size(io_size) - isl_nixl_bench.set_batch_size(batch_size) - - isl_nixl_bench.configure_scheme(direction="isl") - isl_nixl_bench.configure_segment_type( - args.backend, args.source, args.destination - ) - - # Generate plan - plan = isl_nixl_bench.plan(format=args.format) - - # For JSON format, add config filename to the output - if args.format == "json": - plan_with_config = plan.copy() if isinstance(plan, dict) else {} - plan_with_config["config_file"] = config_file - all_plans.append(plan_with_config) - elif args.format == "csv": - plan_data = plan - # Add metadata - plan_data["config_file"] = config_file - plan_data["model"] = model.to_dict().get("model") - - # Add all model_config parameters with proper prefixes - model_config_dict = model_config.to_dict() - - # Add strategy parameters - for key, value in model_config_dict.get("strategy", {}).items(): - plan_data[f"model_strategy_{key}"] = value - - # Add runtime parameters - for key, value in model_config_dict.get("runtime", {}).items(): - plan_data[f"model_runtime_{key}"] = value - - # Add system parameters - for key, value in model_config_dict.get("system", {}).items(): - plan_data[f"model_system_{key}"] = value - - all_plans.append(plan_data) - else: - print(separator) - print(f"Model Config: {config_file}") - print(f"ISL: {model_config.runtime.isl} tokens") - print(f"Page Size: {model_config.system.page_size}") - print(f"Requests: {model_config.runtime.num_requests}") - print(f"TP: {model_config.model.tp_size}") - print(f"PP: {model_config.model.pp_size}") - print(separator) - print(plan) - print() - except Exception as e: - print(f"Error processing config file {config_file}: {str(e)}") - - # For JSON format, output all plans as an array - if args.format == "json" and all_plans: - print(json.dumps(all_plans, indent=2)) - # For CSV format, output all plans as CSV - elif args.format == "csv" and all_plans: - # Get all unique keys from all plans - fieldnames = set() - for plan in all_plans: - fieldnames.update(plan.keys()) - fieldnames = set(sorted(list(fieldnames))) - - # Write CSV to string buffer - output = io.StringIO() - writer = csv.DictWriter(output, fieldnames=fieldnames) - writer.writeheader() - for plan in all_plans: - writer.writerow(plan) - - # Print CSV output - print(output.getvalue()) - - return 0 diff --git a/benchmark/kvbench/commands/profile.py b/benchmark/kvbench/commands/profile.py deleted file mode 100644 index 13b3264da9..0000000000 --- a/benchmark/kvbench/commands/profile.py +++ /dev/null @@ -1,105 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. - -import argparse - -from commands.args import add_cli_args, add_common_args, add_nixl_bench_args -from commands.nixlbench import NIXLBench -from models.model_config import ModelConfig -from models.models import BaseModelArch -from models.utils import get_batch_size, override_yaml_args - - -class Command: - """ - Command handler for the 'plan' subcommand. - - This command displays the recommended configuration for nixlbench based on - the provided model architecture and model configuration files. - """ - - def __init__(self): - """ - Initialize the plan command. - - Sets the command name and help text for the command-line interface. - """ - self.name = "profile" - self.help = "Run nixlbench" - - def add_arguments( - self, subparser: argparse.ArgumentParser - ) -> argparse.ArgumentParser: - """ - Add command-specific arguments to the argument parser. - - Args: - subparser (argparse.ArgumentParser): The parser for this command. - - Returns: - argparse.ArgumentParser: The updated argument parser with added arguments. - """ - add_common_args(subparser) - add_nixl_bench_args(subparser) - add_cli_args(subparser) - return subparser - - def execute(self, args: argparse.Namespace): - """ - Execute the plan command with the provided arguments. - - Loads the model architecture and configuration from the specified files, - creates a NIXLBench instance with the provided arguments, and generates - a nixlbench command plan. - - Args: - args (argparse.Namespace): Command-line arguments. - - Returns: - int: -1 if required arguments are missing, otherwise None. - """ - if not args.model or not args.model_config: - print("Error: --model and --model_config are required") - return -1 - - if args.model: - model = BaseModelArch.from_yaml(args.model, None) - if args.model_config: - model_config = ModelConfig.from_yaml(args.model_config) - override_yaml_args(model_config, args) - model.set_model_config(model_config) - - filtered_args = { - k: v for k, v in args.__dict__.items() if k in NIXLBench.defaults() - } - nixl_bench = NIXLBench(model, model_config, **filtered_args) - io_size = model.get_io_size(model_config.system.page_size) - batch_size = get_batch_size(model, model_config, io_size) - nixl_bench.set_io_size(io_size) - nixl_bench.set_batch_size(batch_size) - nixl_bench.configure_buffer_size() - - nixl_bench.configure_scheme(direction="isl") - nixl_bench.configure_segment_type(args.backend, args.source, args.destination) - separator = "=" * 80 - - print(f"Model Config: {args.model_config}") - print(f"ISL: {model_config.runtime.isl} tokens") - print(f"Page Size: {model_config.system.page_size}") - print(f"Requests: {model_config.runtime.num_requests}") - print(f"TP: {model_config.model.tp_size}") - print(f"PP: {model_config.model.pp_size}") - print(separator) - nixl_bench.profile() diff --git a/benchmark/kvbench/docs/ct-perftest.md b/benchmark/kvbench/docs/ct-perftest.md new file mode 100644 index 0000000000..1925c2e3a1 --- /dev/null +++ b/benchmark/kvbench/docs/ct-perftest.md @@ -0,0 +1,145 @@ +# CTPerftest - Custom Traffic Performance Test Implementation + +This refers to the commands `ct-perftest` and `sequential-ct-perftest` + + +## Overview +Benchmark tool to measure the performance of customizable traffic patterns, where each rank sends a custom message size to every other rank (can be 0). + +The purpose of such a benchmark is to simulate asymmetric traffic flow, like the KV caches transfers flow in a disaggregated inference system. + +Such a pattern is defined using a transfer matrix, i.e a matrix where cell [i.j] defines the size of the message sent by rank i to rank j. + +The following benchmarks are available: +- Sequential CT Perftest: Benchmark the performance of a continuum of traffic patterns one after the other, before running each pattern, all the ranks involved in the pattern do a barrier, then optionally sleep a given amount of time, and then run the pattern and measure the execution time. +- CT Perftest: Benchmark the performance of one traffic pattern, the pattern is ran in multiple iterations and then metrics are reported. It is useful to optimize specific patterns + + +## Implementation Details +- `custom_traffic_perftest.py` implements CT Perftest and the TrafficPattern dataclass +- `sequential_custom_traffic_perftest.py` implements Sequential CT Perftest +- Utilizes common utilities from `runtime`, mostly for collective communication +- `common.py` defines an abstraction for NixlBuffers + + +# Sequential CT Perftest + +### Reports +Sequential CT Perftest reports the total latency per matrix execution, along with their latency when ran isolated, which can be used to evaluate how good the network react to congestion. The report is both saved as a JSON file (controlled by the `--json-output-path` parameter) and displayed in stdout as: + +``` + Transfer size (GB) Latency (ms) Isolated Latency (ms) Num Senders +-------------------- -------------- ----------------------- ------------- + 4.945 35.047 35.421 4 + 3.230 21.152 21.800 4 + 1.104 8.222 8.280 4 + ... ... ... ... + 0.129 2.147 2.386 4 +``` + + +### Usage +Tests can be defined using YAML configuration files. + +CT Perftest define a single traffic pattern, as well as number of iterations and warmup iterations: +```yaml +iters: 50 +warmup_iters: 10 +traffic_pattern: + matrix_file: "/path/to/matrix.txt" + shards: 1 + mem_type: "cuda" + xfer_op: "WRITE" +``` + +Sequential CT Perftest configuration defines a sequence of traffic patterns: + +```yaml +traffic_patterns: +- matrix_file: /swgwork/eshukrun/nixl/tools/perf/run/llama-405b/prefill_tp_4_decode_tp_8/matrices/matrix_0.txt + metadata: + isl: 38328 + sleep_before_launch_sec: 16.480753057792 +- matrix_file: /swgwork/eshukrun/nixl/tools/perf/run/llama-405b/prefill_tp_4_decode_tp_8/matrices/matrix_1.txt + metadata: + isl: 25034 + sleep_before_launch_sec: 71.875102179328 +``` +`traffic_patterns` can contain multiple elements that run sequentially. See `TrafficPattern` in `common.py` for default values. + +- **matrix_file**: The file containing the matrix, the matrix cells should be separated by whitespaces and contain either a number of bytes as integer or use a standard unit like K, M and G. +- **shards**: Number of chunks the buffer has to be sharded into. +- **mem_type**: For now support only cuda, but it should follow nixl memory types +- **xfer_op**: xfer operation, can be READ or WRITE +- **sleep_before_launch_sec**: number of seconds to sleep before running this traffic pattern, can be used for example to simulate computation. + +Example of a matrix file: +``` +0 0 1M 1M +0 0 1M 1M +0 0 0 5K +0 0 0 5K +``` + +### Usage +Install requirements +```bash +pip install -r requirements.txt +``` + +Optionally, generate matrices using the `benchmark/tools/inference_workload_matgen.py` script +```bash +python benchmark/tools/inference_workload_matgen.py generate \ + --num-user-requests 10 \ + --num-prefill-nodes 16 \ + --num-decode-nodes 16 \ + --prefill-tp 8 \ + --prefill-pp 1 \ + --prefill-cp 1 \ + --decode-tp 8 \ + --decode-pp 1 \ + --decode-cp 1 \ + --results-dir $PWD/matrices \ + --model llama-405b +``` + +Run the test using the CLI: +```bash +python main.py sequential-ct-perftest ./matrices/metadata.yaml --verify-buffers --json-output-path ./results.json +# Or with slurm +srun python main.py sequential-ct-perftest ./matrices/metadata.yaml --verify-buffers --json-output-path ./results.json +``` + +### Next Steps +- [ ] Support more memory types + +### Known Issues +- The nixl xfer preparation currently takes a lot of time (the `_prepare_tp()` method). + + +# Sequential CT Perftest + +### Reports +CT Perftest reports total latency (the time elapsed between the first rank started until the last rank finished), the average time per iteration, the total size sent over the network and the average bandwidth by rank. + + +### Usage +Tests can be defined using YAML configuration files. + + +The configuration define a single traffic pattern, as well as number of iterations and warmup iterations: +```yaml +iters: 50 +warmup_iters: 10 +traffic_pattern: + matrix_file: "/path/to/matrix.txt" + shards: 1 + mem_type: "cuda" + xfer_op: "WRITE" +``` + +Then run it with + +```bash +python tools/perf/main.py ct-perftest path/to/config.yaml +``` \ No newline at end of file diff --git a/benchmark/kvbench/main.py b/benchmark/kvbench/main.py index 0c02476dbd..42e9bd62c0 100644 --- a/benchmark/kvbench/main.py +++ b/benchmark/kvbench/main.py @@ -13,29 +13,415 @@ # See the License for the specific language governing permissions and # limitations under the License. -import argparse +import csv +import glob +import io +import json +import logging +import os +from pathlib import Path +from test.custom_traffic_perftest import CTPerftest +from test.sequential_custom_traffic_perftest import SequentialCTPerftest +from test.traffic_pattern import TrafficPattern -import commands +import click +import numpy as np +import yaml +from commands.args import cli_args, common_args, nixl_bench_args, plan_args +from commands.nixlbench import NIXLBench +from models.model_config import ModelConfig +from models.models import BaseModelArch +from models.utils import get_batch_size, override_yaml_args -def main(): - parser = argparse.ArgumentParser(description="KVBench") - subparsers = parser.add_subparsers(dest="command", help="Available commands") +def parse_size(nbytes: str) -> int: + """Convert formatted string with unit to bytes""" - for command in commands.available_commands: - subparser = subparsers.add_parser(command.name, help=command.help) - command.add_arguments(subparser) + options = {"g": 1024 * 1024 * 1024, "m": 1024 * 1024, "k": 1024, "b": 1} + unit = 1 + key = nbytes[-1].lower() + if key in options: + unit = options[key] + value = float(nbytes[:-1]) + else: + value = float(nbytes) + count = int(unit * value) + return count - args = parser.parse_args() - if args.command: - for command in commands.available_commands: - if command.name == args.command: - command.execute(args) - break - else: - parser.print_help() +def load_matrix(matrix_file) -> np.ndarray: + """Load traffic pattern matrix from file""" + matrix = [] + with open(matrix_file, "r") as f: + for line in f: + row = line.strip().split() + matrix.append([parse_size(x) for x in row]) + mat = np.array(matrix) + return mat + + +@click.group() +@click.option("--debug/--no-debug", default=False, help="Enable debug logging") +def cli(debug): + """KVBench - NIXL Performance Testing CLI""" + log_level = logging.DEBUG if debug else logging.INFO + + # Configure root logger + logging.basicConfig( + level=log_level, format="%(asctime)s - %(name)s - %(levelname)s - %(message)s" + ) + + # Set level for all existing loggers + for logger_name in logging.root.manager.loggerDict: + logger = logging.getLogger(logger_name) + logger.setLevel(log_level) + + +@cli.command("plan") +@common_args +@plan_args +@nixl_bench_args +@cli_args +def plan_command(model, model_config, model_configs, format, **kwargs): + """Display the recommended configuration for nixlbench""" + if not model: + click.echo("Error: --model is required") + return + + if not model_config and not model_configs: + click.echo("Error: either --model_config or --model_configs is required") + return + + # Load model architecture + model_arch = BaseModelArch.from_yaml(model, None) + + # Get list of model config files + config_files = [] + + if model_config: + config_files.append(model_config) + + if model_configs: + # Expand glob patterns into list of files + expanded_files = glob.glob(model_configs) + if not expanded_files: + click.echo(f"Warning: No files matched the pattern: {model_configs}") + config_files.extend(expanded_files) + + if not config_files: + click.echo("Error: No valid model config files specified") + return + + # Filter out duplicate paths + config_files = list(dict.fromkeys(config_files)) + + # Filter arguments for NIXLBench + filtered_args = { + k: v for k, v in kwargs.items() if k in NIXLBench.defaults() and v is not None + } + + # Process each model config + all_plans = [] + for config_file in config_files: + # Skip if file doesn't exist + if not os.path.exists(config_file): + click.echo(f"Warning: Config file not found: {config_file}") + continue + + try: + # Load model configuration + model_configuration = ModelConfig.from_yaml(config_file) + # Override yaml args with cli args if supplied + override_yaml_args(model_configuration, type("Args", (), kwargs)()) + model_arch.set_model_config(model_configuration) + + separator = "=" * 80 + isl_nixl_bench = NIXLBench(model_arch, model_configuration, **filtered_args) + + io_size = model_arch.get_io_size(model_configuration.system.page_size) + batch_size = get_batch_size(model_arch, model_configuration, io_size) + isl_nixl_bench.set_io_size(io_size) + isl_nixl_bench.set_batch_size(batch_size) + + isl_nixl_bench.configure_scheme(direction="isl") + isl_nixl_bench.configure_segment_type( + kwargs.get("backend"), kwargs.get("source"), kwargs.get("destination") + ) + + # Generate plan + plan = isl_nixl_bench.plan(format=format) + + # For JSON format, add config filename to the output + if format == "json": + plan_with_config = plan.copy() if isinstance(plan, dict) else {} + plan_with_config["config_file"] = config_file + all_plans.append(plan_with_config) + elif format == "csv": + plan_data = plan + # Add metadata + plan_data["config_file"] = config_file + plan_data["model"] = model_arch.to_dict().get("model") + + # Add all model_config parameters with proper prefixes + model_config_dict = model_configuration.to_dict() + + # Add strategy parameters + for key, value in model_config_dict.get("strategy", {}).items(): + plan_data[f"model_strategy_{key}"] = value + + # Add runtime parameters + for key, value in model_config_dict.get("runtime", {}).items(): + plan_data[f"model_runtime_{key}"] = value + + # Add system parameters + for key, value in model_config_dict.get("system", {}).items(): + plan_data[f"model_system_{key}"] = value + + all_plans.append(plan_data) + else: + click.echo(separator) + click.echo(f"Model Config: {config_file}") + click.echo(f"ISL: {model_configuration.runtime.isl} tokens") + click.echo(f"Page Size: {model_configuration.system.page_size}") + click.echo(f"Requests: {model_configuration.runtime.num_requests}") + click.echo(f"TP: {model_configuration.model.tp_size}") + click.echo(f"PP: {model_configuration.model.pp_size}") + click.echo(separator) + click.echo(plan) + click.echo() + except Exception as e: + click.echo(f"Error processing config file {config_file}: {str(e)}") + + # For JSON format, output all plans as an array + if format == "json" and all_plans: + click.echo(json.dumps(all_plans, indent=2)) + # For CSV format, output all plans as CSV + elif format == "csv" and all_plans: + # Get all unique keys from all plans + fieldnames = set() + for plan in all_plans: + fieldnames.update(plan.keys()) + fieldnames = set(sorted(list(fieldnames))) + + # Write CSV to string buffer + output = io.StringIO() + writer = csv.DictWriter(output, fieldnames=fieldnames) + writer.writeheader() + for plan in all_plans: + writer.writerow(plan) + + # Print CSV output + click.echo(output.getvalue()) + + +@cli.command("profile") +@common_args +@nixl_bench_args +@cli_args +def profile_command(model, model_config, **kwargs): + """Run nixlbench""" + if not model or not model_config: + click.echo("Error: --model and --model_config are required") + return + + model_arch = BaseModelArch.from_yaml(model, None) + model_configuration = ModelConfig.from_yaml(model_config) + override_yaml_args(model_configuration, type("Args", (), kwargs)()) + model_arch.set_model_config(model_configuration) + + filtered_args = { + k: v for k, v in kwargs.items() if k in NIXLBench.defaults() and v is not None + } + nixl_bench = NIXLBench(model_arch, model_configuration, **filtered_args) + io_size = model_arch.get_io_size(model_configuration.system.page_size) + batch_size = get_batch_size(model_arch, model_configuration, io_size) + nixl_bench.set_io_size(io_size) + nixl_bench.set_batch_size(batch_size) + nixl_bench.configure_buffer_size() + + nixl_bench.configure_scheme(direction="isl") + nixl_bench.configure_segment_type( + kwargs.get("backend"), kwargs.get("source"), kwargs.get("destination") + ) + separator = "=" * 80 + + click.echo(f"Model Config: {model_config}") + click.echo(f"ISL: {model_configuration.runtime.isl} tokens") + click.echo(f"Page Size: {model_configuration.system.page_size}") + click.echo(f"Requests: {model_configuration.runtime.num_requests}") + click.echo(f"TP: {model_configuration.model.tp_size}") + click.echo(f"PP: {model_configuration.model.pp_size}") + click.echo(separator) + nixl_bench.profile() + + +@cli.command("kvcache") +@common_args +@cli_args +def kvcache_command(model, model_config, **kwargs): + """Display kvcache information""" + if not model or not model_config: + click.echo("Error: --model and --model_config are required") + return + + # Load model architecture + model_arch = BaseModelArch.from_yaml(model, None) + + # Load model configuration + model_configuration = ModelConfig.from_yaml(model_config) + override_yaml_args(model_configuration, type("Args", (), kwargs)()) + # Set model_config on the model instance using the new method + model_arch.set_model_config(model_configuration) + + from math import floor, log + + def format_bytes(size): + power = 0 if size <= 0 else floor(log(size, 1024)) + return f"{round(size / 1024 ** power, 2)} {['B', 'KB', 'MB', 'GB', 'TB'][int(power)]}" + + labels = [ + "KV Cache Size Per Token", + "IO Size", + "Batch Size", + "Model", + "Input Sequence Length", + ] + max_width = max(len(label) for label in labels) + io_size = model_arch.get_io_size(model_configuration.system.page_size) + batch_size = get_batch_size(model_arch, model_configuration, io_size) + click.echo(f"{'Model':{max_width}}: {model_arch.model}") + click.echo( + f"{'Input Sequence Length':{max_width}}: {model_configuration.runtime.isl}" + ) + click.echo(f"{'Batch Size':{max_width}}: {batch_size}") + click.echo(f"{'IO Size':{max_width}}: {format_bytes(io_size)}") + + +@cli.command("io-size") +@common_args +@cli_args +def io_size_command(model, model_config, **kwargs): + """Display IO size information""" + if not model or not model_config: + click.echo("Error: --model and --model_config are required") + return + + # Load model architecture + model_arch = BaseModelArch.from_yaml(model, None) + + # Load model configuration + model_configuration = ModelConfig.from_yaml(model_config) + override_yaml_args(model_configuration, type("Args", (), kwargs)()) + model_arch.set_model_config(model_configuration) + + from math import floor, log + + def format_bytes(size): + power = 0 if size <= 0 else floor(log(size, 1024)) + return f"{round(size / 1024 ** power, 2)} {['B', 'KB', 'MB', 'GB', 'TB'][int(power)]}" + + io_size = model_arch.get_io_size(model_configuration.system.page_size) + batch_size = get_batch_size(model_arch, model_configuration, io_size) + + click.echo(f"Model: {model_arch.model}") + click.echo(f"Page Size: {format_bytes(model_configuration.system.page_size)}") + click.echo(f"Input Sequence Length: {model_configuration.runtime.isl}") + click.echo(f"Batch Size: {batch_size}") + click.echo(f"IO Size: {format_bytes(io_size)}") + + +@cli.command("sequential-ct-perftest") +@click.argument("config_file", type=click.Path(exists=True)) +@click.option( + "--verify-buffers/--no-verify-buffers", + default=False, + help="Verify buffer contents after transfer", +) +@click.option( + "--print-recv-buffers/--no-print-recv-buffers", + default=False, + help="Print received buffer contents", +) +@click.option( + "--json-output-path", + type=click.Path(), + help="Path to save JSON output", + default=None, +) +def sequential_ct_perftest( + config_file, verify_buffers, print_recv_buffers, json_output_path +): + """Run sequential custom traffic performance test using patterns defined in YAML config""" + with open(config_file, "r") as f: + config = yaml.safe_load(f) + + if "traffic_patterns" not in config: + raise ValueError("Config file must contain 'traffic_patterns' key") + + patterns = [] + for instruction_config in config["traffic_patterns"]: + tp_config = instruction_config + required_fields = ["matrix_file"] + missing_fields = [field for field in required_fields if field not in tp_config] + + if missing_fields: + raise ValueError( + f"Traffic pattern missing required fields: {missing_fields}" + ) + + pattern = TrafficPattern( + matrix=load_matrix(Path(tp_config["matrix_file"])), + shards=tp_config.get("shards", 1), + mem_type=tp_config.get("mem_type", "cuda").lower(), + xfer_op=tp_config.get("xfer_op", "WRITE").upper(), + sleep_after_launch_sec=tp_config.get("sleep_after_launch_sec", 0), + ) + patterns.append(pattern) + + output_path = json_output_path + + perftest = SequentialCTPerftest(patterns) + perftest.run( + verify_buffers=verify_buffers, + print_recv_buffers=print_recv_buffers, + json_output_path=output_path, + ) + + +@cli.command("ct-perftest") +@click.argument("config_file", type=click.Path(exists=True)) +@click.option( + "--verify-buffers/--no-verify-buffers", + default=False, + help="Verify buffer contents after transfer", +) +@click.option( + "--print-recv-buffers/--no-print-recv-buffers", + default=False, + help="Print received buffer contents", +) +def ct_perftest(config_file, verify_buffers, print_recv_buffers): + """Run custom traffic performance test using patterns defined in YAML config""" + with open(config_file, "r") as f: + config = yaml.safe_load(f) + + tp_config = config.get("traffic_pattern") + if tp_config is None: + raise ValueError("Config file must contain 'traffic_pattern' key") + + iters = config.get("iters", 1) + warmup_iters = config.get("warmup_iters", 0) + + pattern = TrafficPattern( + matrix=load_matrix(Path(tp_config["matrix_file"])), + shards=tp_config.get("shards", 1), + mem_type=tp_config.get("mem_type", "cuda").lower(), + xfer_op=tp_config.get("xfer_op", "WRITE").upper(), + ) + + perftest = CTPerftest(pattern, iters=iters, warmup_iters=warmup_iters) + perftest.run(verify_buffers=verify_buffers, print_recv_buffers=print_recv_buffers) if __name__ == "__main__": - main() + cli() diff --git a/benchmark/kvbench/pyproject.toml b/benchmark/kvbench/pyproject.toml index ac43c90894..a8fa5afa88 100644 --- a/benchmark/kvbench/pyproject.toml +++ b/benchmark/kvbench/pyproject.toml @@ -14,12 +14,27 @@ # limitations under the License. [project] -name = "kvbench" -version = "0.1.0" -description = "Add your description here" +name = "NIXLKVBench" +version = "0.2.0" +description = "Benchmarking utility for testing KVCache transfers in large language models (LLMs) using the NIXL runtime." readme = "README.md" requires-python = ">=3.9" dependencies = [ - "pytest>=8.3.5", + "asposestorage==1.0.2", + "click==8.1.7", + "numpy==2.3.1", + "pytest==7.4.4", "pyyaml>=6.0.2", + "etcd3>=0.12.0", + "tabulate==0.9.0", + "torch==2.7.0", + "tqdm==4.66.5", + "nixl", ] + +[tool.setuptools.packages.find] +include = ["models*", "runtime*", "commands*"] + +[build-system] +requires = ["setuptools>=61.0"] +build-backend = "setuptools.build_meta" diff --git a/benchmark/kvbench/runtime/etcd_rt.py b/benchmark/kvbench/runtime/etcd_rt.py new file mode 100644 index 0000000000..4f3adecdbc --- /dev/null +++ b/benchmark/kvbench/runtime/etcd_rt.py @@ -0,0 +1,249 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import logging +import os +import pickle +import re +import time +from collections import defaultdict +from typing import Any, List, Optional + +import etcd3 + +from .rt_base import ReduceOp, _RTUtils + +log = logging.getLogger(__name__) + + +def int_to_bytes(val: int) -> bytes: + return int.to_bytes(val, length=4) + + +class _EtcdDistUtils(_RTUtils): + """ETCD-based MPI utilities - NOT PERFORMANCE OPTIMIZED (for control path only)""" + + def __init__( + self, + etcd_endpoints: str = "http://localhost:2379", + prefix: str = "/nixl/kvbench", + ): + super().__init__() + self.prefix = prefix + self.ops_counter: dict[str, dict[Any, int]] = defaultdict( + lambda: defaultdict(int) + ) + + # Initialize rank & world size + if os.environ.get("SLURM_PROCID"): + self.rank = int(os.environ["SLURM_PROCID"]) + self.world_size = int(os.environ["SLURM_NTASKS"]) + elif os.environ.get("OMPI_COMM_WORLD_RANK"): + self.rank = int(os.environ["OMPI_COMM_WORLD_RANK"]) + self.world_size = int(os.environ["OMPI_COMM_WORLD_SIZE"]) + elif os.environ.get("RANK"): + self.rank = int(os.environ["RANK"]) + self.world_size = int(os.environ["WORLD_SIZE"]) + else: + raise ValueError( + "Rank and world size not found in environment variables SLURM_PROCID/SLURM_NTASKS or RANK/WORLD_SIZE" + ) + + # Parse endpoint host & port + url_pattern = r"^(https?://)?([^:]+)(?::(\d+))?$" + match = re.match(url_pattern, etcd_endpoints) + + if match: + protocol = match.group(1) or "http://" + protocol = protocol.rstrip("://") + host = match.group(2) + port = int(match.group(3)) if match.group(3) else 2379 + else: + raise ValueError( + f"Invalid etcd endpoint format: {etcd_endpoints}, expected format is [http://]host[:port]" + ) + + log.info(f"ETCD client initialized with host {host} & port {port}") + + try: + self.client = etcd3.client(host=host, port=port) + except Exception as e: + raise ValueError(f"Failed to initialize ETCD client: {e}") + + if self.rank == 0: + log.info(f"Wiping ETCD prefix {self.prefix}") + self.client.delete_prefix(self.prefix) + + def destroy_dist(self): + self.client.delete_prefix(self.prefix) + + def _get_int_val(self, key: str) -> int | None: + val = self.client.get(key)[0] + if val is None: + return None + return int.from_bytes(val) + + def barrier(self, ranks: Optional[List[int]] = None, timeout_sec=600): + """Barrier for a group of ranks using etcd barrier""" + if ranks is None: + ranks = list(range(self.world_size)) + + if self.rank not in ranks: + return + + ranks = sorted(ranks) + root = ranks[0] + # Create barrier for specific group of ranks + group_id = self._get_group_id(ranks) + group_size = len(ranks) + barrier_ix = self.ops_counter["barrier"][group_id] + self.ops_counter["barrier"][group_id] += 1 + + key = f"{self.prefix}/barrier/{group_id}/{barrier_ix}" + start_time = time.time() + + if self.rank == root: + # Fan in - count from 1 to len(ranks) + self.client.put(key, int_to_bytes(1)) + + # Fan out - put it back to 0 to signal that all ranks have entered the barrier + while not self.client.replace( + key, int_to_bytes(len(ranks)), int_to_bytes(0) + ): + if timeout_sec and time.time() - start_time > timeout_sec: + raise TimeoutError( + f"[Rank {self.rank}] ROOT - Barrier {key} timed out after {timeout_sec} seconds, current value: {self.client.get(key)}, waiting for val={len(ranks)} (i.e all the ranks have entered the barrier), (ranks: {ranks})" + ) + else: + my_index = ranks.index(self.rank) + # Fan in - count from 1 to len(ranks) + while not self.client.replace( + key, int_to_bytes(my_index), int_to_bytes(my_index + 1) + ): + if timeout_sec and time.time() - start_time > timeout_sec: + raise TimeoutError( + f"[Rank {self.rank}] Barrier {key} timed out after {timeout_sec} seconds, current value: {self.client.get(key)}, waiting for val={my_index} (i.e rank {ranks[my_index - 1]} entered barrier), (ranks: {ranks})" + ) + # Fan out - wait for root to set 0 again + while not self._get_int_val(key) == 0: + if timeout_sec and time.time() - start_time > timeout_sec: + raise TimeoutError( + f"[Rank {self.rank}] Barrier {key} timed out after {timeout_sec} seconds, current value: {self.client.get(key)}, waiting for val={group_size} (i.e all the ranks have entered the barrier), (ranks: {ranks})" + ) + + def get_rank(self) -> int: + return self.rank + + def get_world_size(self) -> int: + return self.world_size + + def allgather_obj(self, obj: Any) -> List[Any]: + allgather_ix = self.ops_counter["allgather"]["world"] + self.ops_counter["allgather"]["world"] += 1 + + result = [None for _ in range(self.world_size)] + # Serialize the object + serialized_obj = pickle.dumps(obj) + + self.client.put( + f"{self.prefix}/allgather/{allgather_ix}/{self.rank}", serialized_obj + ) + + for dest_rank in range(self.world_size): + val = None + while val is None: + val = self.client.get( + f"{self.prefix}/allgather/{allgather_ix}/{dest_rank}" + )[0] + + result[dest_rank] = pickle.loads(val) + + return result + + def alltoall_obj(self, send_objs: List[Any]) -> List[Any]: + result = [None for _ in range(self.world_size)] + serialized_objs = [pickle.dumps(obj) for obj in send_objs] + + self.barrier() + for dest_rank in range(self.world_size): + self.client.put( + f"{self.prefix}/alltoall/{self.rank}_to_{dest_rank}", + serialized_objs[dest_rank], + ) + + self.barrier() + for src_rank in range(self.world_size): + val = self.client.get(f"{self.prefix}/alltoall/{src_rank}_to_{self.rank}")[ + 0 + ] + result[src_rank] = pickle.loads(val) + + return result + + def all_reduce( + self, vals: List[float | int], op: ReduceOp, root: int = 0 + ) -> List[float | int]: + self.barrier() + self.client.put(f"{self.prefix}/all_reduce/{self.rank}", pickle.dumps(vals)) + if self.rank == root: + self.client.delete(f"{self.prefix}/all_reduce/result") + self.barrier() + + if self.rank == root: + vals = [] + for dest_rank in range(self.world_size): + val = self.client.get(f"{self.prefix}/all_reduce/{dest_rank}")[0] + vals.append(pickle.loads(val)) + + print(vals) + if op == ReduceOp.SUM: + final_val = [sum(col) for col in zip(*vals)] + elif op == ReduceOp.AVG: + final_val = [sum(col) / self.world_size for col in zip(*vals)] + elif op == ReduceOp.MIN: + final_val = [min(col) for col in zip(*vals)] + elif op == ReduceOp.MAX: + final_val = [max(col) for col in zip(*vals)] + else: + raise ValueError(f"Unsupported reduce operation: {op}") + + self.client.put(f"{self.prefix}/all_reduce/result", pickle.dumps(final_val)) + else: + while self.client.get(f"{self.prefix}/all_reduce/result")[0] is None: + pass + val = self.client.get(f"{self.prefix}/all_reduce/result")[0] + final_val = pickle.loads(val) + + return final_val + + def _get_group_id(self, ranks: List[int]) -> int: + """Get the id for a group of ranks""" + key = tuple(sorted(ranks)) + return hash(key) + + +if not os.environ.get("NIXL_ETCD_NAMESPACE"): + log.warning( + "Environment variable NIXL_ETCD_NAMESPACE is not set, using default prefix /nixl/kvbench. " + "Note that it can lead to conflicts if multiple instances of KVBench are running. " + "To avoid this, set NIXL_ETCD_NAMESPACE to a unique value for each instance of KVBench. " + 'For example, export NIXL_ETCD_NAMESPACE="/nixl/kvbench/$(uuidgen)"' + ) + + +etcd_dist_utils = _EtcdDistUtils( + etcd_endpoints=os.environ.get("NIXL_ETCD_ENDPOINTS", "http://localhost:2379"), + prefix=os.environ.get("NIXL_ETCD_NAMESPACE", "/nixl/kvbench"), +) diff --git a/benchmark/kvbench/runtime/rt_base.py b/benchmark/kvbench/runtime/rt_base.py new file mode 100644 index 0000000000..f44dea8180 --- /dev/null +++ b/benchmark/kvbench/runtime/rt_base.py @@ -0,0 +1,68 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from enum import Enum +from typing import Any, List, Optional + + +class ReduceOp(Enum): + """Reduction operations for distributed computing""" + + SUM = "sum" + AVG = "avg" + MIN = "min" + MAX = "max" + + +class _RTUtils: + """Base class for distributed runtime utilities""" + + def __init__(self): + pass + + def get_rank(self) -> int: + """Get the rank of the current process""" + raise NotImplementedError("Subclasses must implement get_rank") + + def get_world_size(self) -> int: + """Get the total number of processes""" + raise NotImplementedError("Subclasses must implement get_world_size") + + def barrier(self, ranks: Optional[List[int]] = None, timeout_sec: int = 600): + """Synchronization barrier for processes""" + raise NotImplementedError("Subclasses must implement barrier") + + def allgather_obj(self, obj: Any) -> List[Any]: + """All-gather operation for objects""" + raise NotImplementedError("Subclasses must implement allgather_obj") + + def alltoall_obj(self, send_objs: List[Any]) -> List[Any]: + """All-to-all operation for objects""" + raise NotImplementedError("Subclasses must implement alltoall_obj") + + def all_reduce( + self, vals: List[float | int], op: ReduceOp, root: int = 0 + ) -> List[float | int]: + """All-reduce operation with specified reduction operation""" + raise NotImplementedError("Subclasses must implement all_reduce") + + def destroy_dist(self): + """Clean up distributed resources""" + raise NotImplementedError("Subclasses must implement destroy_dist") + + def _get_group_id(self, ranks: List[int]) -> int: + """Get the id for a group of ranks""" + key = tuple(sorted(ranks)) + return hash(key) diff --git a/benchmark/kvbench/test/__init__.py b/benchmark/kvbench/test/__init__.py new file mode 100644 index 0000000000..3159bfe656 --- /dev/null +++ b/benchmark/kvbench/test/__init__.py @@ -0,0 +1,14 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. diff --git a/benchmark/kvbench/test/custom_traffic_perftest.py b/benchmark/kvbench/test/custom_traffic_perftest.py new file mode 100644 index 0000000000..6a3bcf2cf1 --- /dev/null +++ b/benchmark/kvbench/test/custom_traffic_perftest.py @@ -0,0 +1,388 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import logging +import time +from test.traffic_pattern import TrafficPattern +from typing import Literal, Optional, Tuple + +import numpy as np +import torch +from runtime.etcd_rt import etcd_dist_utils as dist_rt +from tabulate import tabulate + +from nixl._api import nixl_agent + +log = logging.getLogger(__name__) + + +class NixlHandle: + def __init__(self, remote_rank, handle, traffic_pattern): + self.remote_rank = remote_rank + self.handle = handle + self.tp = traffic_pattern + + def __str__(self): + return f"to:{self.remote_rank}" + + +class NixlBuffer: + """Can be sharded""" + + def __init__( + self, + size: int, + mem_type: str, + nixl_agent: nixl_agent, + shards=1, + fill_value=0, + dtype: torch.dtype = torch.int8, + ): + self.size = size + self.nixl_agent = nixl_agent + if mem_type in ("cuda", "vram"): + device = torch.device("cuda") + elif mem_type in ("cpu", "dram"): + device = torch.device("cpu") + else: + raise ValueError(f"Unsupported memory type: {mem_type}") + + if shards > 1: + raise ValueError("Sharding is not supported yet") + + log.debug( + f"[Rank {dist_rt.get_rank()}] Initializing NixlBuffer with size {size}, device {device}, shards {shards}, fill_value {fill_value}" + ) + self.buf = torch.full((size,), fill_value, dtype=dtype, device=device) + + log.debug( + f"[Rank {dist_rt.get_rank()}] Registering memory for buffer {self.buf}" + ) + self.reg_descs = nixl_agent.get_reg_descs(self.buf) + assert ( + nixl_agent.register_memory(self.reg_descs) is not None + ), "Failed to register memory" + + def get_chunk(self, size, offset): + if offset + size > self.size: + raise ValueError( + f"Offset {offset} + size {size} is greater than buffer size {self.size}" + ) + return self.buf[offset : offset + size] + + def destroy(self): + self.nixl_agent.deregister_memory(self.reg_descs) + if hasattr(self.buf, "is_cuda") and self.buf.is_cuda: + del self.buf + torch.cuda.empty_cache() + + +class CTPerftest: + def __init__( + self, traffic_pattern: TrafficPattern, iters: int = 1, warmup_iters: int = 0 + ): + """ + Args: + traffic_pattern: The communication pattern to test + iters: Number of test iterations + warmup_iters: Number of warmup iterations before timing + """ + self.my_rank = dist_rt.get_rank() + self.world_size = dist_rt.get_world_size() + self.traffic_pattern = traffic_pattern + self.iters = iters + self.warmup_iters = warmup_iters + + self.nixl_agent = nixl_agent(f"{self.my_rank}") + + """Initialize the buffers, one big send and recv buffer is used for all the transfers + it has to be chunked inside each transfer to get buffers per ranks + the buffer is big enough to handle any of the transfers + For now, support only CUDA/CPU buffers + """ + self.send_buf: NixlBuffer = NixlBuffer( + self.traffic_pattern.total_src_size(self.my_rank), + mem_type=self.traffic_pattern.mem_type, + nixl_agent=self.nixl_agent, + dtype=self.traffic_pattern.dtype, + ) + self.recv_buf: NixlBuffer = NixlBuffer( + self.traffic_pattern.total_dst_size(self.my_rank), + mem_type=self.traffic_pattern.mem_type, + nixl_agent=self.nixl_agent, + dtype=self.traffic_pattern.dtype, + ) + + self._check_tp_config(traffic_pattern) + assert "UCX" in self.nixl_agent.get_plugin_list(), "UCX plugin is not loaded" + + def _barrier_tp(self, tp: TrafficPattern, senders_only=True): + """Barrier for a traffic pattern""" + if senders_only: + dist_rt.barrier(tp.senders_ranks()) + else: + dist_rt.barrier(tp.ranks()) + + def _share_md(self) -> None: + """Share agent metadata between all ranks. (Need to be run after registering buffers)""" + log.debug(f"[Rank {self.my_rank}] Sharing MD") + md = self.nixl_agent.get_agent_metadata() + mds = dist_rt.allgather_obj(md) + for other_rank, metadata in enumerate(mds): + if other_rank == self.my_rank: + continue + self.nixl_agent.add_remote_agent(metadata) + + dist_rt.barrier() + + def _share_recv_buf_descs(self, my_recv_bufs: list[Optional[NixlBuffer]]) -> list: + """Share receive buffer descriptors between all ranks, in alltoall style. + Args: + my_recv_bufs: List of receive buffers for current rank + Returns: + List of buffer descriptors from all ranks + """ + xfer_descs = [ + self.nixl_agent.get_xfer_descs(buf) if buf is not None else None + for buf in my_recv_bufs + ] + + my_recv_bufs_serdes = [ + self.nixl_agent.get_serialized_descs(xfer_descs) + for xfer_descs in xfer_descs + ] + + dst_bufs_serdes = dist_rt.alltoall_obj(my_recv_bufs_serdes) + + dst_bufs_descs = [ + self.nixl_agent.deserialize_descs(serdes) for serdes in dst_bufs_serdes + ] + return dst_bufs_descs + + def _get_bufs(self, tp: TrafficPattern): + """Returns lists of buffers where bufs[i] is the send/recv buffer for rank i""" + send_bufs = [None for _ in range(self.world_size)] + recv_bufs = [None for _ in range(self.world_size)] + send_offset = recv_offset = 0 + + for other_rank in range(self.world_size): + send_size = tp.matrix[self.my_rank][other_rank] + recv_size = tp.matrix[other_rank][self.my_rank] + send_buf = recv_buf = None + if send_size > 0: + send_buf = self.send_buf.get_chunk(send_size, send_offset) + send_offset += send_size + if recv_size > 0: + recv_buf = self.recv_buf.get_chunk(recv_size, recv_offset) + recv_offset += recv_size + + send_bufs[other_rank] = send_buf + recv_bufs[other_rank] = recv_buf + + return send_bufs, recv_bufs + + def _prepare_tp( + self, tp: TrafficPattern + ) -> Tuple[ + list[NixlHandle], list[Optional[NixlBuffer]], list[Optional[NixlBuffer]] + ]: + """Timing everything in this function because it takes a lot of time""" + + send_bufs, recv_bufs = self._get_bufs(tp) + + log.debug(f"[Rank {self.my_rank}] Sharing recv buf descs") + dst_bufs_descs = self._share_recv_buf_descs(recv_bufs) + + handles: list[NixlHandle] = [] + for other, buf in enumerate(send_bufs): + if buf is None: + continue + + xfer_desc = self.nixl_agent.get_xfer_descs(buf) + + log.debug( + f"[Rank {self.my_rank}] Initializing xfer for {other} - xfer desc: {xfer_desc}, dst buf desc: {dst_bufs_descs[other]}" + ) + handle = self.nixl_agent.initialize_xfer( + "WRITE", + xfer_desc, + dst_bufs_descs[other], + f"{other}", + f"{tp.id}_{self.my_rank}_{other}", + ) + handles.append(NixlHandle(other, handle, tp)) + + return handles, send_bufs, recv_bufs + + def _warmup( + self, + iters=15, + fill_value: int = 100000, + mem_type: Literal["cuda", "vram", "cpu", "dram"] = "cuda", + ): + full_matrix = np.full((self.world_size, self.world_size), fill_value=fill_value) + tp = TrafficPattern(matrix=full_matrix, mem_type=mem_type) + handles, send_bufs, recv_bufs = self._prepare_tp(tp) + for _ in range(iters): + self._run_tp(handles) + self._wait(handles) + + def _run_tp(self, handles: list[NixlHandle], blocking=False) -> list[NixlHandle]: + pending = [] + for handle in handles: + status = self.nixl_agent.transfer(handle.handle) + assert status != "ERR", "Transfer failed" + if status != "DONE": + pending.append(handle) + + if not blocking: + return pending + else: + self._wait(pending) + return [] + + def _wait(self, handles: list[NixlHandle]): + # Wait for transfers to complete + while True: + pending = [] + for handle in handles: + state = self.nixl_agent.check_xfer_state(handle.handle) + assert state != "ERR", "Transfer got to Error state." + if state != "DONE": + pending.append(handle) + + if not pending: + break + handles = pending + + def _destroy(self, handles: list[NixlHandle]): + log.debug(f"[Rank {self.my_rank}] Releasing XFER handles") + for handle in handles: + self.nixl_agent.release_xfer_handle(handle.handle) + + log.debug(f"[Rank {self.my_rank}] Removing remote agents") + for other_rank in range(self.world_size): + if other_rank == self.my_rank: + continue + self.nixl_agent.remove_remote_agent(f"{other_rank}") + + self._destroy_buffers() + + def _destroy_buffers(self): + log.debug(f"[Rank {self.my_rank}] Destroying buffers") + self.send_buf.destroy() + self.recv_buf.destroy() + + def _check_tp_config(self, tp: TrafficPattern): + # Matrix size should be world * world + assert tp.matrix.shape == ( + self.world_size, + self.world_size, + ), f"Matrix size is not the same as world size, got {tp.matrix.shape}, world_size={self.world_size}" + + def _verify_tp( + self, + tp: TrafficPattern, + recv_bufs: list[Optional[NixlBuffer]], + print_recv_buffers: bool = False, + ): + for r, recv_buf in enumerate(recv_bufs): + if recv_buf is None: + if tp.matrix[r][self.my_rank] > 0: + log.error( + f"Rank {self.my_rank} expected {tp.matrix[r][self.my_rank]} bytes from rank {r}, but got 0" + ) + raise RuntimeError("Buffer verification failed") + continue + + if print_recv_buffers: + log.info(f"Recv buffer {r}:\n{recv_buf.buf}") + + # recv_buf has to be filled with the rank of the sender + # and its size has to be the same as matrix[r][my_rank] + full_recv_buf = recv_buf.buf + expected = torch.full_like(full_recv_buf, r) + assert torch.all( + full_recv_buf == expected + ), f"Vector not equal to {r}, got {full_recv_buf}" + assert ( + full_recv_buf.size(0) == tp.matrix[r][self.my_rank] + ), f"Size of vector {r} is not the same as matrix[r][{self.my_rank}], got {full_recv_buf.size(0)}" + + def _get_tp_total_size(self, tp: TrafficPattern) -> int: + """Return total size of matrix in bytes""" + return np.sum(tp.matrix, axis=(0, 1)) * tp.dtype.itemsize + + def run( + self, verify_buffers: bool = False, print_recv_buffers: bool = False + ) -> float: + """Execute the performance test. + + Args: + verify_buffers: Whether to verify buffer contents after transfer + print_recv_buffers: Whether to print receive buffer contents + + Returns: + Total execution time in seconds + """ + log.debug(f"[Rank {self.my_rank}] Running CT perftest") + self._share_md() + + handles, send_bufs, recv_bufs = self._prepare_tp(self.traffic_pattern) + + for _ in range(self.warmup_iters): + pending_handles = self._run_tp(handles) + self._wait(pending_handles) + + start = time.time() + for i in range(self.iters): + pending_handles = self._run_tp(handles) + self._wait(pending_handles) + end = time.time() + + # Metrics report + start_times = dist_rt.allgather_obj(start) + end_times = dist_rt.allgather_obj(end) + + global_time = max(end_times) - min(start_times) + avg_time_per_iter_sec = global_time / self.iters + + total_size_gb = self._get_tp_total_size(self.traffic_pattern) / 1e9 + + # Print metrics as a table + if self.my_rank == 0: + headers = [ + "Iters", + "Total time (s)", + "Avg latency/iter (s)", + "Total size (GB)", + ] + data = [ + [ + self.iters, + global_time, + avg_time_per_iter_sec, + total_size_gb, + ] + ] + print(tabulate(data, headers=headers, floatfmt=".6f")) + + if verify_buffers: + self._verify_tp(self.traffic_pattern, recv_bufs, print_recv_buffers) + + # Destroy + self._destroy(handles) + + return end - start diff --git a/benchmark/kvbench/test/inference_workload_matgen.py b/benchmark/kvbench/test/inference_workload_matgen.py new file mode 100644 index 0000000000..42d42a0ebe --- /dev/null +++ b/benchmark/kvbench/test/inference_workload_matgen.py @@ -0,0 +1,624 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Utils to generate matrices that represent the communication patterns of an inference workload + +Disclaimers: +- For now there is only support for TP and CP +- The compute time estimation is very naive and assumes 36 TFlops (h100) +- The batching is very naive and just aggregates requests into batches until it exceeds max_batch_mem or batch_size + +Example usage: +python inference_workload_matgen.py generate \ + --num-user-requests 10 \ + --batch-size 1 \ + --num-prefill-nodes 54 \ + --num-decode-nodes 54 \ + --prefill-tp 8 \ + --prefill-pp 1 \ + --prefill-cp 1 \ + --decode-tp 8 \ + --decode-pp 1 \ + --decode-cp 1 \ + --results-dir /tmp/matrices \ + --isl-mean 16000 \ + --isl-scale 10000 \ + --min-isl 1000 \ + --max-isl 128000 \ + --max-batch-mem 100000000000 \ + --hidden-size 16384 \ + --num-layers 126 \ + --num-heads 128 \ + --num-kv-heads 8 \ + --dtype-size 2 \ +OR instead of hidden-size/num-layers/etc.. Use a preconfigured model: + --model llama-405b +""" + +from dataclasses import dataclass +from itertools import cycle +from os import PathLike +from pathlib import Path +from typing import Any, List, Optional + +import numpy as np +import yaml +from tqdm import tqdm + + +@dataclass +class ModelConfig: + """Configuration for a language model.""" + + hidden_size: int # Model's hidden dimension (H) + num_layers: int # Number of layers (L) + num_heads: int = 1 # Number of attention heads (N_heads) + num_kv_heads: Optional[int] = None # Number of key/value heads (for MQA/GQA) + dtype_size: float = 2 # Size in bytes (2 for FP16, 4 for FP32) + + @property + def head_dim(self): + return self.hidden_size // self.num_heads + + def bytes_per_token(self): + # return 2 * self.hidden_size * self.num_layers * self.dtype_size * self.num_heads + if self.num_kv_heads is not None: + return ( + 2 + * self.head_dim + * self.num_kv_heads + * self.num_layers + * self.dtype_size + ) + else: + return 2 * self.hidden_size * self.num_layers * self.dtype_size + + def kv_cache_size(self, isl): + """KV cache size in bytes""" + return isl * self.bytes_per_token() + + +@dataclass +class TaskConfig: + """Configuration for an inference task.""" + + isl_mean: int = 0 # Context/sequence length (S) + isl_scale: int = 10000 # Context/sequence length scale + min_isl: int = 1000 # Minimum context/sequence length + max_isl: int = 128000 # Maximum context/sequence length + batch_size: int = 1 # Batch size (B) + max_batch_mem: float = 100e9 # Maximum batch memory (100GB) + + +@dataclass +class WorkerConfig: + """Configuration for the GPU cluster.""" + + tp: int + pp: int = 1 + cp: int = 1 + ep: int = 1 + + def __post_init__(self): + assert self.pp == 1, "PP is not supported yet" + assert self.ep == 1, "EP is not supported yet" + + +@dataclass +class UserRequest: + isl: int + + @classmethod + def rand( + cls, + mean: int = 16000, + scale: int = 10000, + min_isl: int = 1000, + max_isl: int = 128000, + ): + # Sample and clip to ensure we stay within bounds + isl = int(np.random.normal(mean, scale)) + # Ensure ISL is positive + isl = min(max(min_isl, isl), max_isl) + return cls(isl=isl) + + +@dataclass +class Batch: + user_requests: List[UserRequest] + + def kv_cache_size(self, model_config: ModelConfig): + """KV cache size in bytes""" + + return sum(model_config.kv_cache_size(req.isl) for req in self.user_requests) + + @property + def size(self): + return len(self.user_requests) + + @property + def total_isl(self): + return sum(req.isl for req in self.user_requests) + + +@dataclass +class TransferMatrix: + matrix: np.ndarray + compute_time: float + isl: int + + +def gen_batches( + num_user_requests: int, + task_config: TaskConfig, + model_config: ModelConfig, + max_batch_mem: float = 100e9, # 100GB - capacity of a gpu +): + """ + For now very naive, aggregate requests into batches until it exceeds max_batch_mem or batch_size + Args: + - num_user_requests: Number of user requests + - task_config: Task configuration + """ + batches = [] + curr = [] + curr_mem = 0 + for _ in range(num_user_requests): + req = UserRequest.rand( + task_config.isl_mean, + task_config.isl_scale, + task_config.min_isl, + task_config.max_isl, + ) + curr_mem += model_config.kv_cache_size(req.isl) + curr.append(req) + if curr_mem > task_config.max_batch_mem or len(curr) >= task_config.batch_size: + batches.append(Batch(user_requests=curr)) + curr = [] + curr_mem = 0 + if curr: + batches.append(Batch(user_requests=curr)) + print(f"Last batch is incomplete, his size is {len(curr)}") + + return batches + + +def gen_matrices_and_compute_time( + batches: List[Batch], + prefill_workers: List[List[int]], + decode_workers: List[List[int]], + model_config: ModelConfig, + prefill_worker_config: WorkerConfig, + decode_worker_config: WorkerConfig, +) -> List[TransferMatrix]: + """ + Args: + - batches: List of batches + """ + # For now, every prefill worker is bound to a single decode worker + assert len(prefill_workers) == len( + decode_workers + ), f"Prefill and decode workers must have the same number of workers, got {len(prefill_workers)} and {len(decode_workers)}" + + # Assertions + all_ranks = list(r for worker in prefill_workers + decode_workers for r in worker) + world_size = max(all_ranks) + 1 + assert set(all_ranks) == set(range(world_size)), "Ranks are missing" + + workers_coupling = list(zip(prefill_workers, decode_workers)) + + workers_pool = cycle(workers_coupling) + matrices = [] + + for batch in tqdm(batches, desc="Generating matrices"): + prefill_worker, decode_worker = next(workers_pool) + mat = gen_matrix( + batch, + world_size, + prefill_worker, + decode_worker, + model_config, + prefill_worker_config, + decode_worker_config, + ) + + compute_time = estimate_compute_time(batch, model_config) + matrix_obj = TransferMatrix( + matrix=mat, compute_time=compute_time, isl=batch.total_isl + ) + matrices.append(matrix_obj) + + return matrices + + +def gen_matrix( + batch: Batch, + world_size: int, + prefill_worker: List[int], + decode_worker: List[int], + model_config: ModelConfig, + prefill_worker_config: WorkerConfig, + decode_worker_config: WorkerConfig, +): + kv_size = batch.kv_cache_size(model_config) + kv_slice_size = ( + kv_size + / prefill_worker_config.tp + / prefill_worker_config.pp + / prefill_worker_config.cp + ) + + num_peers = ( + decode_worker_config.tp / prefill_worker_config.tp / prefill_worker_config.cp + ) + if num_peers % 1 != 0: + raise ValueError("Prefill TP*Prefill CP must be a divisor of decode TP") + num_peers = int(num_peers) + buf_size = kv_slice_size / num_peers + + # print(f"kv_size: {format_size(kv_size)}, kv_slice_size: {format_size(kv_slice_size)}, buf_size: {format_size(buf_size)}, num_peers: {num_peers}") + + mat = np.zeros((world_size, world_size)) + + dst_iter = iter(decode_worker) + for rank in prefill_worker: + for _ in range(num_peers): + dst = next(dst_iter) + mat[rank, dst] = buf_size + return mat + + +def estimate_compute_time( + batch: Batch, + model_config: ModelConfig, + flops: float = 36 * 1e12, # 36 TFlops (h100) +): + """Estimate the compute time of a batch, in seconds + Very approximative and assumes 36 TFlops (h100) + The formula comes from: + https://medium.com/@plienhar/llm-inference-series-3-kv-caching-unveiled-048152e461c8 + """ + flop = ( + 2 + * batch.size + * model_config.num_layers + * model_config.hidden_size + * batch.total_isl**2 + ) + return flop / flops + + +def format_size(nbytes: float, precision=2) -> str: + if nbytes == 0: + return "0" + + units = ["B", "K", "M", "G"] + units_ix = 0 + while nbytes / 1024 >= 1 and units_ix < len(units) - 1: + nbytes /= 1024 + units_ix += 1 + + nbytes = round(nbytes, precision) + return f"{nbytes:g}{units[units_ix]}" + + +def main( + num_user_requests: int, + task_config: TaskConfig, + num_prefill_gpus: int, + num_decode_gpus: int, + prefill_worker_config: WorkerConfig, + decode_worker_config: WorkerConfig, + model_config: ModelConfig, + results_dir: Optional[PathLike] = None, + rail_optimized: bool = False, +): + """ + Args: + - prefill_gpus: List of GPUs ranks that are used for prefill + - rail_optimized: Whether to reorder the decode workers to match rail-optimized communication (assumption: 8 nic per nodes, nic 0 is connected to nic 0 and 4 of other nodes, 1 to 1/5 etc) + Only supported for 4 GPUs per prefill worker and 8 GPUs per decode worker + + Returns: + matrices + """ + # Rules of thumb + assert ( + prefill_worker_config.tp <= decode_worker_config.tp + ), "Prefill TP must be less than or equal to decode TP" + assert ( + prefill_worker_config.pp >= decode_worker_config.pp + ), "Prefill PP must be more or equal to decode PP" + assert ( + prefill_worker_config.cp >= decode_worker_config.cp + ), "Prefill CP must be more or equal to decode CP" + assert decode_worker_config.cp == 1, "Decode CP must be 1" + assert ( + prefill_worker_config.ep <= decode_worker_config.ep + ), "Prefill EP must be less or equal to decode EP" + if rail_optimized: + assert ( + decode_worker_config.tp == 8 + ), "Rail optimized communication is only supported when decode worker is a full node (8 GPUs)" + assert ( + prefill_worker_config.tp == 4 + ), "Rail optimized communication is only supported when prefill worker is half a node (4 GPUs)" + + # Create workers - group of gpus that do prefill/decode + prefill_worker_size = ( + prefill_worker_config.tp * prefill_worker_config.pp * prefill_worker_config.cp + ) + decode_worker_size = ( + decode_worker_config.tp * decode_worker_config.pp * decode_worker_config.cp + ) + world_size = num_prefill_gpus + num_decode_gpus + + # Create list of all GPU ranks + prefill_ranks = list(range(num_prefill_gpus)) + decode_ranks = list(range(num_prefill_gpus, num_prefill_gpus + num_decode_gpus)) + + # Chunk the ranks into worker groups + prefill_workers = [ + prefill_ranks[i : i + prefill_worker_size] + for i in range(0, len(prefill_ranks), prefill_worker_size) + ] + + decode_workers = [ + decode_ranks[i : i + decode_worker_size] + for i in range(0, len(decode_ranks), decode_worker_size) + ] + if rail_optimized: + # Reorder the decode workers to match rail-optimized communication + reordered = [] + order = [0, 4, 1, 5, 2, 6, 3, 7] + for worker in decode_workers: + new_worker = [worker[ix] for ix in order] + reordered.append(new_worker) + + decode_workers = reordered + + print(f"Prefill workers: {prefill_workers}") + print(f"Decode workers: {decode_workers}") + + batches = gen_batches(num_user_requests, task_config, model_config) + print(f"Generated {len(batches)} batches") + matrices = gen_matrices_and_compute_time( + batches, + prefill_workers, + decode_workers, + model_config, + prefill_worker_config, + decode_worker_config, + ) + + # Save matrices and metadata to files + results_dir = results_dir or Path(f"matrices_{world_size}ranks") + results_dir = Path(results_dir) + print(f"Saving {len(matrices)} matrices to {results_dir}") + results_dir.mkdir(parents=True, exist_ok=True) + + metadata: dict[str, Any] = { + "traffic_patterns": [], + } + + for idx, matrix in enumerate(tqdm(matrices, desc="Saving matrices")): + # Save matrix to npy file + matrix_path = results_dir / f"matrix_{idx}.txt" + with open(matrix_path, "w") as f: + for row in matrix.matrix: + f.write(" ".join(f"{format_size(val)}" for val in row) + "\n") + + # Add metadata + metadata["traffic_patterns"].append( + { + "matrix_file": matrix_path.absolute().as_posix(), + "sleep_before_launch_sec": matrix.compute_time, + "metadata": { + "isl": matrix.isl, + }, + } + ) + + # Save metadata to YAML + metadata_path = results_dir / "metadata.yaml" + with open(metadata_path, "w") as f: + yaml.dump(metadata, f) + print(f"Saved metadata to {metadata_path}") + + +if __name__ == "__main__": + import click + + PREDEFINED_MODELS = { + "llama-405b": ModelConfig( + hidden_size=16384, + num_layers=126, + num_heads=128, + num_kv_heads=8, + dtype_size=2, + ), + "qwen3-30B": ModelConfig( + hidden_size=32768, num_layers=48, num_heads=32, num_kv_heads=4, dtype_size=2 + ), + "deepseek-r1-distill-llama-70b": ModelConfig( + hidden_size=8192, num_layers=80, num_heads=64, num_kv_heads=8, dtype_size=2 + ), + "deepseek-r1": ModelConfig( + hidden_size=12288, + num_layers=100, + num_heads=96, + num_kv_heads=12, + dtype_size=2, + ), + } + + @click.group() + def cli(): + """Generate communication matrices for inference workloads""" + pass + + @cli.command() + @click.option( + "--num-user-requests", + type=int, + default=1000, + help="Number of user requests to simulate", + ) + @click.option("--batch-size", type=int, default=1, help="Batch size for requests") + @click.option( + "--num-prefill-nodes", + type=int, + required=True, + help="Number of nodes for prefill", + ) + @click.option( + "--num-decode-nodes", type=int, required=True, help="Number of nodes for decode" + ) + @click.option( + "--prefill-tp", type=int, default=1, help="Tensor parallelism for prefill" + ) + @click.option( + "--prefill-pp", type=int, default=1, help="Pipeline parallelism for prefill" + ) + @click.option( + "--prefill-cp", + type=int, + default=1, + help="Communication parallelism for prefill", + ) + @click.option( + "--decode-tp", type=int, default=1, help="Tensor parallelism for decode" + ) + @click.option( + "--decode-pp", type=int, default=1, help="Pipeline parallelism for decode" + ) + @click.option( + "--decode-cp", type=int, default=1, help="Communication parallelism for decode" + ) + @click.option("--model", type=str, help="Name of predefined model") + @click.option( + "--hidden-size", type=int, help="Model hidden size (for custom model)" + ) + @click.option( + "--num-layers", type=int, help="Number of model layers (for custom model)" + ) + @click.option( + "--num-heads", type=int, help="Number of attention heads (for custom model)" + ) + @click.option( + "--num-kv-heads", + type=int, + help="Number of KV attention heads (for custom model)", + ) + @click.option( + "--dtype-size", type=int, help="Size of model dtype in bytes (for custom model)" + ) + @click.option("--results-dir", type=str, help="Directory to save results") + @click.option( + "--isl-mean", default=16000, type=int, help="Mean context/sequence length" + ) + @click.option( + "--isl-scale", default=10000, type=int, help="Scale context/sequence length" + ) + @click.option( + "--min-isl", default=1000, type=int, help="Minimum context/sequence length" + ) + @click.option( + "--max-isl", default=128000, type=int, help="Maximum context/sequence length" + ) + @click.option( + "--max-batch-mem", default=100e9, type=float, help="Maximum batch memory" + ) + @click.option( + "--rail-optimized/--no-rail-optimized", + default=False, + help="Whether to use rail optimization", + ) + @click.option("--ppn", default=8, type=int, help="Number of GPUs per node") + def generate( + num_user_requests, + batch_size, + num_prefill_nodes, + num_decode_nodes, + prefill_tp, + prefill_pp, + prefill_cp, + decode_tp, + decode_pp, + decode_cp, + model, + hidden_size, + num_layers, + num_heads, + num_kv_heads, + dtype_size, + results_dir, + isl_mean, + isl_scale, + min_isl, + max_isl, + max_batch_mem, + rail_optimized, + ppn, + ): + """Generate communication matrices for given configuration""" + + if model: + if model not in PREDEFINED_MODELS: + raise click.BadParameter( + f"Unknown model {model}. Available models: {list(PREDEFINED_MODELS.keys())}" + ) + model_config = PREDEFINED_MODELS[model] + else: + if not all([hidden_size, num_layers, num_heads, num_kv_heads, dtype_size]): + raise click.BadParameter( + "Must specify either --model or all custom model parameters" + ) + model_config = ModelConfig( + hidden_size=hidden_size, + num_layers=num_layers, + num_heads=num_heads, + num_kv_heads=num_kv_heads, + dtype_size=dtype_size, + ) + + task_config = TaskConfig( + isl_mean=isl_mean, + isl_scale=isl_scale, + min_isl=min_isl, + max_isl=max_isl, + batch_size=batch_size, + max_batch_mem=max_batch_mem, + ) + + # world_size = num_prefill_nodes * prefill_tp + num_decode_nodes * decode_tp + # print(f"World size: {world_size}") + # print(f"Model config: {model_config}") + + main( + num_user_requests=num_user_requests, + task_config=task_config, + num_prefill_gpus=num_prefill_nodes * ppn, + num_decode_gpus=num_decode_nodes * ppn, + prefill_worker_config=WorkerConfig( + tp=prefill_tp, pp=prefill_pp, cp=prefill_cp + ), + decode_worker_config=WorkerConfig(tp=decode_tp, pp=decode_pp, cp=decode_cp), + model_config=model_config, + results_dir=results_dir, + rail_optimized=rail_optimized, + ) + + cli() diff --git a/benchmark/kvbench/test/sequential_custom_traffic_perftest.py b/benchmark/kvbench/test/sequential_custom_traffic_perftest.py new file mode 100644 index 0000000000..278ac7476c --- /dev/null +++ b/benchmark/kvbench/test/sequential_custom_traffic_perftest.py @@ -0,0 +1,416 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Sequential is different from multi in that every rank processes only one TP at a time, but they can process different ones""" + +import json +import logging +import time +from collections import defaultdict +from itertools import chain +from test.custom_traffic_perftest import CTPerftest, NixlBuffer +from test.traffic_pattern import TrafficPattern +from typing import Any, Dict, List, Optional + +import yaml +from runtime.etcd_rt import etcd_dist_utils as dist_rt +from tabulate import tabulate + +from nixl._api import nixl_agent + +log = logging.getLogger(__name__) + + +class SequentialCTPerftest(CTPerftest): + """Extends CTPerftest to handle multiple traffic patterns sequentially. + The patterns are executed in sequence, and the results are aggregated. + + Allows testing multiple communication patterns sequentially between distributed processes. + """ + + def __init__( + self, + traffic_patterns: list[TrafficPattern], + n_iters: int = 3, + n_isolation_iters=30, + warmup_iters=30, + ) -> None: + """Initialize multi-pattern performance test. + + Args: + traffic_patterns: List of traffic patterns to test simultaneously + """ + self.my_rank = dist_rt.get_rank() + self.world_size = dist_rt.get_world_size() + self.traffic_patterns = traffic_patterns + self.n_iters = n_iters + self.n_isolation_iters = n_isolation_iters + self.warmup_iters = warmup_iters + + log.debug(f"[Rank {self.my_rank}] Initializing Nixl agent") + self.nixl_agent = nixl_agent(f"{self.my_rank}") + + for tp in self.traffic_patterns: + self._check_tp_config(tp) + assert "UCX" in self.nixl_agent.get_plugin_list(), "UCX plugin is not loaded" + + # NixlBuffer caches buffers and reuse them if they are big enough, let's initialize them once, with the largest needed size + self.send_buf_by_mem_type: dict[str, NixlBuffer] = {} + self.recv_buf_by_mem_type: dict[str, NixlBuffer] = {} + + def _init_buffers(self): + log.debug(f"[Rank {self.my_rank}] Initializing buffers") + max_src_by_mem_type = defaultdict(int) + max_dst_by_mem_type = defaultdict(int) + + for tp in self.traffic_patterns: + max_src_by_mem_type[tp.mem_type] = max( + max_src_by_mem_type[tp.mem_type], tp.total_src_size(self.my_rank) + ) + max_dst_by_mem_type[tp.mem_type] = max( + max_dst_by_mem_type[tp.mem_type], tp.total_dst_size(self.my_rank) + ) + + for mem_type, size in max_src_by_mem_type.items(): + if not size: + continue + self.send_buf_by_mem_type[mem_type] = NixlBuffer( + size, mem_type=mem_type, nixl_agent=self.nixl_agent + ) + + for mem_type, size in max_dst_by_mem_type.items(): + if not size: + continue + self.recv_buf_by_mem_type[mem_type] = NixlBuffer( + size, mem_type=mem_type, nixl_agent=self.nixl_agent + ) + + def _destroy_buffers(self): + log.debug(f"[Rank {self.my_rank}] Destroying buffers") + for buf in chain( + self.send_buf_by_mem_type.values(), self.recv_buf_by_mem_type.values() + ): + buf.destroy() + + def _get_bufs(self, tp: TrafficPattern): + log.debug(f"[Rank {self.my_rank}] Getting buffers for TP {tp.id}") + + send_bufs = [None for _ in range(self.world_size)] + recv_bufs = [None for _ in range(self.world_size)] + + send_offset_by_memtype: dict[str, int] = defaultdict(int) + recv_offset_by_memtype: dict[str, int] = defaultdict(int) + + for other_rank in range(self.world_size): + send_size = tp.matrix[self.my_rank][other_rank] + recv_size = tp.matrix[other_rank][self.my_rank] + send_buf = recv_buf = None + + if send_size > 0: + send_buf = self.send_buf_by_mem_type[tp.mem_type].get_chunk( + send_size, send_offset_by_memtype[tp.mem_type] + ) + send_offset_by_memtype[tp.mem_type] += send_size + if recv_size > 0: + recv_buf = self.recv_buf_by_mem_type[tp.mem_type].get_chunk( + recv_size, recv_offset_by_memtype[tp.mem_type] + ) + recv_offset_by_memtype[tp.mem_type] += recv_size + + send_bufs[other_rank] = send_buf + recv_bufs[other_rank] = recv_buf + + return send_bufs, recv_bufs + + def run( + self, + verify_buffers: bool = False, + print_recv_buffers: bool = False, + json_output_path: Optional[str] = None, + ): + """ + Args: + verify_buffers: Whether to verify buffer contents after transfer + print_recv_buffers: Whether to print receive buffer contents + yaml_output_path: Path to save results in YAML format + + Returns: + Total execution time in seconds + + This method initializes and executes multiple traffic patterns simultaneously, + measures their performance, and optionally verifies the results. + """ + log.debug(f"[Rank {self.my_rank}] Running sequential CT perftest") + self._init_buffers() + self._share_md() + + results: Dict[str, Any] = { + "iterations_results": [], + "metadata": { + "ts": time.time(), + "iters": [{} for _ in range(self.n_iters)], + }, + } + + tp_handles: list[list] = [] + tp_bufs = [] + + s = time.time() + log.info(f"[Rank {self.my_rank}] Preparing TPs") + for i, tp in enumerate(self.traffic_patterns): + handles, send_bufs, recv_bufs = self._prepare_tp(tp) + tp_bufs.append((send_bufs, recv_bufs)) + tp_handles.append(handles) + + results["metadata"]["prepare_tp_time"] = time.time() - s + + # Warmup + warm_dsts: set[int] = set() + for tp_ix, handles in enumerate(tp_handles): + dsts = set(tp.receivers_ranks(from_ranks=[self.my_rank])) + if dsts.issubset(warm_dsts): + # All the dsts have been warmed up + continue + for _ in range(self.warmup_iters): + self._run_tp(handles, blocking=True) + warm_dsts.update(dsts) + + dist_rt.barrier() + + # Isolated mode - Measure SOL for every matrix + log.info( + f"[Rank {self.my_rank}] Running isolated benchmark (to measure perf without noise)" + ) + my_isolated_tp_latencies: list[float] = [0 for _ in tp_handles] + + results["metadata"]["sol_calculation_ts"] = time.time() + for tp_ix, handles in enumerate(tp_handles): + tp = self.traffic_patterns[tp_ix] + dist_rt.barrier() + if self.my_rank not in tp.senders_ranks(): + continue + + self._barrier_tp(tp) + + for _ in range(self.n_isolation_iters): + t = time.time() + self._run_tp(handles, blocking=True) + e = time.time() + my_isolated_tp_latencies[tp_ix] += e - t + self._barrier_tp(tp) + + log.debug( + f"[Rank {self.my_rank}] Ran {self.n_isolation_iters} isolated iters for tp {tp_ix}/{len(tp_handles)}, took {e - t} secs" + ) + + my_isolated_tp_latencies[tp_ix] /= self.n_isolation_iters + + # Store isolated results + isolated_tp_latencies_by_ranks = dist_rt.allgather_obj(my_isolated_tp_latencies) + isolated_tp_latencies_ms = [] + for i in range(len(self.traffic_patterns)): + tp_lats = [ + rank_lats[i] + for rank_lats in isolated_tp_latencies_by_ranks + if rank_lats[i] > 0 + ] + + if tp_lats: + isolated_tp_latencies_ms.append(max(tp_lats) * 1e3) + + log.info(f"[Rank {self.my_rank}] Running workload benchmark") + + # Workload mode - Measure perf of the matrices while running the full workload + for iter_ix in range(self.n_iters): + log.debug( + f"[Rank {self.my_rank}] Running iteration {iter_ix + 1}/{self.n_iters}" + ) + iter_metadata = results["metadata"]["iters"][iter_ix] + + tp_starts: list[float | None] = [None] * len(tp_handles) + tp_ends: list[float | None] = [None] * len(tp_handles) + log.debug(f"[Rank {self.my_rank}] Warmup done.") + dist_rt.barrier(timeout_sec=None) + + iter_metadata["start_ts"] = time.time() + for tp_ix, handles in enumerate(tp_handles): + tp = self.traffic_patterns[tp_ix] + + if self.my_rank not in tp.senders_ranks(): + continue + + self._barrier_tp(tp) + if tp.sleep_before_launch_sec is not None: + time.sleep(tp.sleep_before_launch_sec) + + # Run TP + log.debug(f"[Rank {self.my_rank}] Running TP {tp_ix}/{len(tp_handles)}") + + tp_start_ts = time.time() + self._run_tp(handles, blocking=True) + tp_end_ts = time.time() + + log.debug( + f"[Rank {self.my_rank}] TP {tp_ix} took {tp_end_ts - tp_start_ts} seconds" + ) + + tp_starts[tp_ix] = tp_start_ts + tp_ends[tp_ix] = tp_end_ts + + if tp.sleep_after_launch_sec is not None: + time.sleep(tp.sleep_after_launch_sec) + + iter_metadata["tps_start_ts"] = tp_starts.copy() + iter_metadata["tps_end_ts"] = tp_ends.copy() + + tp_starts_by_ranks = dist_rt.allgather_obj(tp_starts) + tp_ends_by_ranks = dist_rt.allgather_obj(tp_ends) + + tp_latencies_ms: list[float | None] = [] + + tp_sizes_gb = [ + self._get_tp_total_size(tp) / 1e9 for tp in self.traffic_patterns + ] + + for i, tp in enumerate(self.traffic_patterns): + starts = [ + tp_starts_by_ranks[rank][i] + for rank in range(len(tp_starts_by_ranks)) + ] + ends = [ + tp_ends_by_ranks[rank][i] for rank in range(len(tp_ends_by_ranks)) + ] + starts = [x for x in starts if x is not None] + ends = [x for x in ends if x is not None] + + if not ends or not starts: + tp_latencies_ms.append(None) + else: + tp_latencies_ms.append((max(ends) - min(starts)) * 1e3) + + if self.my_rank == 0: + headers = [ + "Transfer size (GB)", + "Latency (ms)", + "Isolated Latency (ms)", + "Num Senders", + ] + data = [ + [ + tp_sizes_gb[i], + tp_latencies_ms[i], + isolated_tp_latencies_ms[i], + len(tp.senders_ranks()), + ] + for i, tp in enumerate(self.traffic_patterns) + ] + print( + f"Iteration {iter_ix + 1}/{self.n_iters}\n", + tabulate(data, headers=headers, floatfmt=".3f"), + ) + + if verify_buffers: + for i, tp in enumerate(self.traffic_patterns): + send_bufs, recv_bufs = tp_bufs[i] + self._verify_tp(tp, recv_bufs, print_recv_buffers) + + iter_results = [ + { + "size": tp_sizes_gb[i], + "latency": tp_latencies_ms[i], + "isolated_latency": isolated_tp_latencies_ms[i], + "num_senders": len(tp.senders_ranks()), + "min_start_ts": min( + filter( + None, + ( + tp_starts_by_ranks[rank][i] + for rank in range(len(tp_starts_by_ranks)) + ), + ) + ), + "max_end_ts": max( + filter( + None, + ( + tp_ends_by_ranks[rank][i] + for rank in range(len(tp_ends_by_ranks)) + ), + ) + ), + } + for i, tp in enumerate(self.traffic_patterns) + ] + results["iterations_results"].append(iter_results) + + results["metadata"]["finished_ts"] = time.time() + if json_output_path and self.my_rank == 0: + log.info(f"Saving results to {json_output_path}") + with open(json_output_path, "w") as f: + json.dump(results, f) + + # Destroy + log.info(f"[Rank {self.my_rank}] Finished run, destroying objects") + self._destroy(handles) + + def _write_yaml_results( + self, + output_path: str, + headers: List[str], + data: List[List], + traffic_patterns: List[TrafficPattern], + ) -> None: + """Write performance test results to a YAML file. + + Args: + output_path: Path to save the YAML file + headers: Column headers for the results + data: Performance data rows + traffic_patterns: List of traffic patterns tested + """ + results: Dict[str, Any] = { + "performance_results": { + "timestamp": time.time(), + "world_size": self.world_size, + "traffic_patterns": [], + } + } + + for i in range(len(traffic_patterns)): + tp_data = {} + for j, header in enumerate(headers): + # Convert header to a valid YAML key + key = header.lower().replace(" ", "_").replace("(", "").replace(")", "") + # Format floating point values to 2 decimal places for readability + if isinstance(data[i][j], float): + tp_data[key] = round(data[i][j], 2) + else: + tp_data[key] = data[i][j] + + # Add traffic pattern name or index for reference + tp_data["pattern_index"] = i + + # You can add more pattern-specific information here if needed + # For example: + # tp_data["sender_ranks"] = list(tp.senders_ranks()) + + results["performance_results"]["traffic_patterns"].append(tp_data) + + try: + with open(output_path, "w") as f: + yaml.dump(results, f, default_flow_style=False, sort_keys=False) + log.info(f"Results saved to YAML file: {output_path}") + except Exception as e: + log.error(f"Failed to write YAML results to {output_path}: {e}") diff --git a/benchmark/kvbench/test/traffic_pattern.py b/benchmark/kvbench/test/traffic_pattern.py new file mode 100644 index 0000000000..e771e58551 --- /dev/null +++ b/benchmark/kvbench/test/traffic_pattern.py @@ -0,0 +1,99 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +import logging +from dataclasses import dataclass, field +from typing import ClassVar, Literal, Optional + +import numpy as np +import torch + +log = logging.getLogger(__name__) + + +@dataclass +class TrafficPattern: + """Represents a communication pattern between distributed processes. + + Attributes: + matrix: Communication matrix as numpy array + mem_type: Type of memory to use + xfer_op: Transfer operation type + shards: Number of shards for distributed processing + dtype: PyTorch data type for the buffers + sleep_before_launch_sec: Number of seconds to sleep before launch + sleep_after_launch_sec: Number of seconds to sleep after launch + id: Unique identifier for this traffic pattern + """ + + matrix: np.ndarray + mem_type: Literal["cuda", "vram", "cpu", "dram"] + xfer_op: Literal["WRITE", "READ"] = "WRITE" + shards: int = 1 + dtype: torch.dtype = torch.int8 + sleep_before_launch_sec: Optional[int] = None + sleep_after_launch_sec: Optional[int] = None + + id: int = field(default_factory=lambda: TrafficPattern._get_next_id()) + _id_counter: ClassVar[int] = 0 + + @classmethod + def _get_next_id(cls) -> int: + """Get the next available ID and increment the counter""" + current_id = cls._id_counter + cls._id_counter += 1 + return current_id + + def senders_ranks(self): + """Return the ranks that send messages""" + senders_ranks = [] + for i in range(self.matrix.shape[0]): + for j in range(self.matrix.shape[1]): + if self.matrix[i, j] > 0: + senders_ranks.append(i) + break + return list(set(senders_ranks)) + + def receivers_ranks(self, from_ranks: Optional[list[int]] = None): + """Return the ranks that receive messages""" + if from_ranks is None: + from_ranks = list(range(self.matrix.shape[0])) + receivers_ranks = [] + for i in from_ranks: + for j in range(self.matrix.shape[1]): + if self.matrix[i, j] > 0: + receivers_ranks.append(j) + break + return list(set(receivers_ranks)) + + def ranks(self): + """Return all ranks that are involved in the traffic pattern""" + return list(set(self.senders_ranks() + self.receivers_ranks())) + + def buf_size(self, src, dst): + return self.matrix[src, dst] + + def total_src_size(self, rank): + """Return the sum of the sizes received by """ + total_src_size = 0 + for other_rank in range(self.matrix.shape[0]): + total_src_size += self.matrix[rank][other_rank] + return total_src_size + + def total_dst_size(self, rank): + """Return the sum of the sizes received by """ + total_dst_size = 0 + for other_rank in range(self.matrix.shape[0]): + total_dst_size += self.matrix[other_rank][rank] + return total_dst_size diff --git a/benchmark/nixlbench/contrib/Dockerfile b/benchmark/nixlbench/contrib/Dockerfile index 658bf0fbfb..4c5ff9396c 100644 --- a/benchmark/nixlbench/contrib/Dockerfile +++ b/benchmark/nixlbench/contrib/Dockerfile @@ -120,7 +120,7 @@ ENV LD_LIBRARY_PATH=/usr/local/lib:$LD_LIBRARY_PATH ENV VIRTUAL_ENV=/workspace/nixl/.venv RUN uv venv $VIRTUAL_ENV --python $DEFAULT_PYTHON_VERSION && \ # pybind11 pip install needed for ubuntu 22.04 - uv pip install --upgrade meson pybind11 patchelf pyYAML + uv pip install --upgrade meson pybind11 patchelf pyYAML click tabulate RUN rm -rf build && \ mkdir build && \ @@ -154,9 +154,23 @@ RUN echo $LD_LIBRARY_PATH RUN ldconfig RUN ls -ll /workspace/nixlbench + RUN rm -rf build && \ mkdir build && \ - uv run meson setup build -Dnixl_path=/usr/local/nixl/ && \ - cd build && ninja -ENV PATH=/workspace/nixlbench/build:$PATH + uv run meson setup build -Dnixl_path=/usr/local/nixl/ -Dprefix=/usr/local/nixlbench && \ + cd build && ninja && ninja install + +WORKDIR /workspace/nixl + +RUN ls -ll benchmark/kvbench + +# Install dependencies for benchmarks +RUN uv pip install -e benchmark/kvbench + +ENV PATH=/usr/local/nixlbench/bin:$PATH +ENV LD_LIBRARY_PATH=/usr/local/nixlbench/lib:$LD_LIBRARY_PATH +ENV PYTHON_PATH=/usr/local/nixlbench/lib/python3/dist-packages/nixlbench/ + +# Fix for etcd proto dependencies +ENV PROTOCOL_BUFFERS_PYTHON_IMPLEMENTATION=python WORKDIR /workspace/nixl/benchmark/kvbench diff --git a/benchmark/nixlbench/src/runtime/etcd/README.md b/benchmark/nixlbench/src/runtime/etcd/README.md new file mode 100644 index 0000000000..183f0caead --- /dev/null +++ b/benchmark/nixlbench/src/runtime/etcd/README.md @@ -0,0 +1,116 @@ +# ETCD Python Runtime + +This directory contains the ETCD-based runtime implementation that replaces the MPI Python runtime with a new Python API that uses C++ etcd functions. + +## Overview + +The ETCD runtime provides distributed coordination and communication using etcd as the backend instead of MPI. This includes: + +- **C++ Runtime** (`etcd_rt.cpp/h`): Core ETCD-based communication primitives +- **Python Bindings** (`python_bindings.cpp`): pybind11 bindings to expose C++ functions to Python +- **Python Runtime** (`../../../kvbench/runtime/etcd_rt.py`): High-level Python API that implements the `_RTUtils` interface + +## Features + +- Point-to-point communication (sendInt, recvInt, sendChar, recvChar) +- Collective operations (broadcast, reduce) +- Barrier synchronization +- Distributed rank management +- Object serialization for Python data structures + +## Building + +### Prerequisites + +- etcd server running (default: `http://localhost:2379`) +- etcd-cpp-api library +- pybind11 +- Python 3.6+ + +### Option 1: Using Meson (Recommended) + +The Python bindings are automatically built as part of the main meson build if pybind11 and Python dependencies are found. + +```bash +# From the main build directory +meson setup builddir +ninja -C builddir +``` + +### Option 2: Using setuptools + +```bash +cd /path/to/etcd +pip install pybind11 +python setup.py build_ext --inplace +``` + +## Usage + +The Python runtime provides ETCD-based distributed coordination: + +### Option 1: Direct instantiation +```python +from nixl.benchmark.kvbench.runtime.etcd_rt import _EtcdDistUtils + +# Create runtime instance using ETCD +runtime = _EtcdDistUtils(etcd_endpoints="http://localhost:2379", size=4) + +# Use the distributed runtime API +rank = runtime.get_rank() +world_size = runtime.get_world_size() + +# Collective operations +data = runtime.allgather_obj({"rank": rank, "data": "hello"}) +runtime.barrier() +``` + +### Option 2: Using the global instance (configured via environment variables) +```python +from nixl.benchmark.kvbench.runtime import dist_rt + +# Runtime is automatically configured from ETCD_ENDPOINTS and WORLD_SIZE env vars +rank = dist_rt.get_rank() +world_size = dist_rt.get_world_size() + +# Collective operations +data = dist_rt.allgather_obj({"rank": rank, "data": "hello"}) +dist_rt.barrier() +``` + +## Configuration + +Environment variables: +- `ETCD_ENDPOINTS`: Comma-separated list of etcd endpoints (default: "http://localhost:2379") +- `WORLD_SIZE`: Number of distributed processes (default: 1) + +## Architecture + +``` +Python Application + ↓ +_EtcdDistUtils (Python) + ↓ +etcd_runtime (pybind11) + ↓ +xferBenchEtcdRT (C++) + ↓ +etcd-cpp-api + ↓ +ETCD Server +``` + +## Limitations + +1. **Data Size**: Character data transfer has practical size limits +2. **Reduce Operations**: Only SUM and AVG are fully implemented (MIN/MAX use placeholders) +3. **Error Handling**: Basic error handling with timeouts and retries +4. **Performance**: May be slower than native MPI for high-frequency operations + +## Future Improvements + +- Implement proper MIN/MAX reduce operations +- Add support for custom serialization protocols +- Optimize data transfer for large objects +- Add connection pooling and reconnection logic +- Support for hierarchical barriers and reductions diff --git a/benchmark/nixlbench/src/runtime/etcd/etcd_rt.cpp b/benchmark/nixlbench/src/runtime/etcd/etcd_rt.cpp index 6d0ad89be1..6a8b3dd32d 100644 --- a/benchmark/nixlbench/src/runtime/etcd/etcd_rt.cpp +++ b/benchmark/nixlbench/src/runtime/etcd/etcd_rt.cpp @@ -30,25 +30,28 @@ // ETCD Runtime implementation xferBenchEtcdRT::xferBenchEtcdRT(const std::string& etcd_endpoints, const int size, int *terminate_input) { + // Store parameters for later use in setup + stored_etcd_endpoints = etcd_endpoints.empty() ? ETCD_EP_DEFAULT : etcd_endpoints; + global_size = size; + terminate = terminate_input; - std::string use_etcd_ep = ETCD_EP_DEFAULT; - - // Parse command line arguments to get ETCD endpoints - if (!etcd_endpoints.empty()) { - use_etcd_ep = etcd_endpoints; - } - + // Initialize other members namespace_prefix = "xferbench/"; // Namespace for XFER benchmark barrier_gen = 0; - terminate = terminate_input; + my_rank = -1; // Will be set in setup() + + // Client will be initialized in setup() +} +int +xferBenchEtcdRT::setup() { // Connect to ETCD try { - std::cout << "Connecting to ETCD at " << use_etcd_ep << std::endl; - client = std::make_unique(use_etcd_ep); + std::cout << "Connecting to ETCD at " << stored_etcd_endpoints << std::endl; + client = std::make_unique (stored_etcd_endpoints); } catch (const std::exception& e) { std::cerr << "Failed to connect to ETCD: " << e.what() << std::endl; - exit(EXIT_FAILURE); + return -1; } // Registration process - get a unique rank @@ -58,7 +61,7 @@ xferBenchEtcdRT::xferBenchEtcdRT(const std::string& etcd_endpoints, const int si auto lock_response = client->lock(lock_key).get(); if (lock_response.error_code() != 0) { std::cerr << "Failed to acquire lock: " << lock_response.error_message() << std::endl; - exit(EXIT_FAILURE); + return -1; } // Get the current size - number of processes that have registered @@ -69,9 +72,6 @@ xferBenchEtcdRT::xferBenchEtcdRT(const std::string& etcd_endpoints, const int si my_rank = 0; } - // Set the global size to the sum of initiator and target devices - global_size = size; - // Update registration information client->put(makeKey("size"), std::to_string(my_rank + 1)).get(); client->put(makeKey("rank", my_rank), "active").get(); @@ -88,11 +88,13 @@ xferBenchEtcdRT::xferBenchEtcdRT(const std::string& etcd_endpoints, const int si << " is greater than or equal to global size " << global_size << std::endl; client->rmdir(makeKey(""), true).get(); - exit(EXIT_FAILURE); + return -1; } std::cout << "ETCD Runtime: Registered as rank " << my_rank << " item " << my_rank + 1 << " of " << global_size << std::endl; + + return 0; } xferBenchEtcdRT::~xferBenchEtcdRT() { diff --git a/benchmark/nixlbench/src/runtime/etcd/etcd_rt.h b/benchmark/nixlbench/src/runtime/etcd/etcd_rt.h index c068ff7703..93c7d00f32 100644 --- a/benchmark/nixlbench/src/runtime/etcd/etcd_rt.h +++ b/benchmark/nixlbench/src/runtime/etcd/etcd_rt.h @@ -44,7 +44,7 @@ enum xferBenchEtcdMsgType { class xferBenchEtcdRT: public xferBenchRT { private: // ETCD connection settings - std::string etcd_endpoints; + std::string stored_etcd_endpoints; std::string namespace_prefix; std::unique_ptr client; @@ -76,6 +76,10 @@ class xferBenchEtcdRT: public xferBenchRT { int *terminate = nullptr); ~xferBenchEtcdRT(); + // Setup function to initialize connection and perform registration + int + setup(); + int getRank() const; int getSize() const; // Communication methods diff --git a/benchmark/nixlbench/src/runtime/etcd/meson.build b/benchmark/nixlbench/src/runtime/etcd/meson.build index 005499d4af..ca4f8b9c0e 100644 --- a/benchmark/nixlbench/src/runtime/etcd/meson.build +++ b/benchmark/nixlbench/src/runtime/etcd/meson.build @@ -24,3 +24,18 @@ etcd_rt_lib = static_library('etcd_rt', dependencies: [etcd_dep], install: true, ) + +# Python bindings +python = import('python').find_installation('python3', required: false) +pybind11_dep = dependency('pybind11', required: false) + +if python.found() and pybind11_dep.found() + py_mod = python.extension_module('etcd_runtime', + 'python_bindings.cpp', + include_directories: inc_dir, + dependencies: [etcd_dep, pybind11_dep], + link_with: etcd_rt_lib, + install: true, + subdir: 'nixlbench' + ) +endif diff --git a/benchmark/nixlbench/src/runtime/etcd/python_bindings.cpp b/benchmark/nixlbench/src/runtime/etcd/python_bindings.cpp new file mode 100644 index 0000000000..477762dd18 --- /dev/null +++ b/benchmark/nixlbench/src/runtime/etcd/python_bindings.cpp @@ -0,0 +1,71 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#include +#include +#include +#include "etcd_rt.h" + +namespace py = pybind11; + +PYBIND11_MODULE (etcd_runtime, m) { + m.doc() = "Python bindings for ETCD runtime"; + + py::class_ (m, "EtcdRuntime") + .def (py::init(), + py::arg ("etcd_endpoints"), + py::arg ("size"), + py::arg ("terminate") = nullptr) + .def ("setup", &xferBenchEtcdRT::setup) + .def ("get_rank", &xferBenchEtcdRT::getRank) + .def ("get_size", &xferBenchEtcdRT::getSize) + .def ("send_int", + [] (xferBenchEtcdRT &self, int value, int dest_rank) { + return self.sendInt (&value, dest_rank); + }) + .def ("recv_int", + [] (xferBenchEtcdRT &self, int src_rank) { + int value; + int result = self.recvInt (&value, src_rank); + return py::make_tuple (result, value); + }) + .def ("send_char", + [] (xferBenchEtcdRT &self, const std::string &data, int dest_rank) { + return self.sendChar (const_cast (data.c_str()), data.size(), dest_rank); + }) + .def ("recv_char", + [] (xferBenchEtcdRT &self, size_t count, int src_rank) { + std::vector buffer (count); + int result = self.recvChar (buffer.data(), count, src_rank); + std::string data (buffer.begin(), buffer.end()); + return py::make_tuple (result, data); + }) + .def ("broadcast_int", + [] (xferBenchEtcdRT &self, py::array_t buffer, int root_rank) { + py::buffer_info buf_info = buffer.request(); + int *ptr = static_cast (buf_info.ptr); + size_t count = buf_info.size; + return self.broadcastInt (ptr, count, root_rank); + }) + .def ("reduce_sum_double", + [] (xferBenchEtcdRT &self, double local_value, int dest_rank) { + double global_value; + int result = self.reduceSumDouble (&local_value, &global_value, dest_rank); + return py::make_tuple (result, global_value); + }) + .def ("barrier", &xferBenchEtcdRT::barrier); +} diff --git a/benchmark/nixlbench/src/runtime/etcd/setup.py b/benchmark/nixlbench/src/runtime/etcd/setup.py new file mode 100755 index 0000000000..56ef74490f --- /dev/null +++ b/benchmark/nixlbench/src/runtime/etcd/setup.py @@ -0,0 +1,48 @@ +#!/usr/bin/env python3 +# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import pybind11 +from pybind11.setup_helpers import Pybind11Extension, build_ext +from setuptools import setup + +# Define the extension module +ext_modules = [ + Pybind11Extension( + "etcd_runtime", + [ + "python_bindings.cpp", + "etcd_rt.cpp", + ], + include_dirs=[ + # Path to pybind11 headers + pybind11.get_include(), + # Path to the parent directory (for runtime.h) + "..", + # Path to current directory + ".", + ], + libraries=["etcd-cpp-api"], + language="c++", + ), +] + +setup( + name="etcd_runtime", + ext_modules=ext_modules, + cmdclass={"build_ext": build_ext}, + zip_safe=False, + python_requires=">=3.6", +) diff --git a/benchmark/nixlbench/src/runtime/etcd/test_etcd_runtime.py b/benchmark/nixlbench/src/runtime/etcd/test_etcd_runtime.py new file mode 100755 index 0000000000..d869dd99b4 --- /dev/null +++ b/benchmark/nixlbench/src/runtime/etcd/test_etcd_runtime.py @@ -0,0 +1,74 @@ +#!/usr/bin/env python3 +# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +""" +Test script for ETCD Python runtime +Run multiple instances of this script to test distributed functionality +""" + +import os +import sys + +# Add the kvbench runtime path +sys.path.insert(0, os.path.join(os.path.dirname(__file__), "../../../kvbench/runtime")) + +try: + from etcd_rt import _EtcdDistUtils + + def test_basic_functionality(): + """Test basic rank and size functionality""" + print("Testing basic functionality...") + + # Initialize runtime - modify size based on how many processes you're running + runtime = _EtcdDistUtils(etcd_endpoints="http://localhost:2379", size=2) + + rank = runtime.get_rank() + world_size = runtime.get_world_size() + + print(f"Rank: {rank}, World Size: {world_size}") + + # Test barrier + print(f"Rank {rank}: Before barrier") + runtime.barrier() + print(f"Rank {rank}: After barrier") + + # Test allgather + my_data = {"rank": rank, "message": f"Hello from rank {rank}"} + print(f"Rank {rank}: Gathering data...") + + try: + all_data = runtime.allgather_obj(my_data) + print(f"Rank {rank}: Gathered data from all ranks:") + for i, data in enumerate(all_data): + print(f" Rank {i}: {data}") + except Exception as e: + print(f"Rank {rank}: Allgather failed: {e}") + + # Test barrier again + runtime.barrier() + print(f"Rank {rank}: Test completed successfully!") + + if __name__ == "__main__": + test_basic_functionality() + +except ImportError as e: + print(f"Import error: {e}") + print("Make sure the etcd_runtime module is built and accessible") + print("Also ensure the etcd server is running at http://localhost:2379") + sys.exit(1) +except Exception as e: + print(f"Runtime error: {e}") + sys.exit(1) diff --git a/benchmark/nixlbench/src/worker/worker.cpp b/benchmark/nixlbench/src/worker/worker.cpp index 703b99109a..e54819d504 100644 --- a/benchmark/nixlbench/src/worker/worker.cpp +++ b/benchmark/nixlbench/src/worker/worker.cpp @@ -31,7 +31,14 @@ static xferBenchRT *createRT(int *terminate) { if (xferBenchConfig::isStorageBackend()) { total = 1; } - return new xferBenchEtcdRT(xferBenchConfig::etcd_endpoints, total, terminate); + xferBenchEtcdRT *etcd_rt = + new xferBenchEtcdRT (xferBenchConfig::etcd_endpoints, total, terminate); + if (etcd_rt->setup() != 0) { + std::cerr << "Failed to setup ETCD runtime" << std::endl; + delete etcd_rt; + exit (EXIT_FAILURE); + } + return etcd_rt; } std::cerr << "Invalid runtime: " << xferBenchConfig::runtime_type << std::endl;