Add Custom Traffic Pattern benchmark

benchmark/ct_perftest/README.md

@@ -0,0 +1,113 @@
+# CTPerftest - Custom Traffic Performance Test Implementation
+
+## 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.
+
+
+## Implementation Details
+- `sequential_custom_traffic_perftest.py` implements Sequential CT Perftest
+- Utilizes common utilities from `utils.py` and `dist_utils.py` for distributed testing support
+- Command-line interface extensions in `cli.py` for test configuration
+- `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 cli.py sequential-ct-perftest ./matrices/metadata.yaml --verify-buffers --json-output-path ./results.json
+# Or with slurm
+srun <params> python cli.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).

benchmark/ct_perftest/cli.py

@@ -0,0 +1,98 @@
+# 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 pathlib import Path
+from utils import load_matrix
+import click
+import yaml
+from common import TrafficPattern
+from sequential_custom_traffic_perftest import SequentialCTPerftest
+from dist_utils import dist_utils
+
+log = logging.getLogger(__name__)
+
+@click.group()
+@click.option("--debug/--no-debug", default=False, help="Enable debug logging")
+def cli(debug):
+    """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()
+@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 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", "dram").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)
+    dist_utils.destroy_dist()
+
+
+if __name__ == "__main__":
+    cli()

benchmark/ct_perftest/common.py

@@ -0,0 +1,123 @@
+# 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 uuid
+from dataclasses import dataclass
+from typing import Literal
+
+import numpy as np
+import torch
+from dist_utils import dist_utils
+
+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=2,
+        fill_value=0,
+        dtype: torch.dtype = torch.int8,
+    ):
+        self.size = size
+        self.nixl_agent = nixl_agent
+        if mem_type in ("cuda", "vram"):
+            device = "cuda"
+        elif mem_type in ("cpu", "dram"):
+            device = "cpu"
+        else:
+            raise ValueError(f"Unsupported memory type: {mem_type}")
+
+        log.debug(
+            f"[Rank {dist_utils.get_rank()}] Initializing NixlBuffer with size {size}, device {device}, shards {shards}, fill_value {fill_value}"
+        )
+        self.bufs = []
+        chunk_size = size // shards
+        self.bufs = [
+            torch.full((chunk_size,), fill_value, dtype=dtype, device=device)
+            for _ in range(shards)
+        ]
+        if size % chunk_size != 0:
+            self.bufs.append(
+                torch.full(
+                    (size % chunk_size,), fill_value, dtype=dtype, device=device
+                )
+            )
+
+        self.reg_descs = nixl_agent.get_reg_descs(self.bufs, mem_type=mem_type)
+        self.xfer_descs = nixl_agent.get_xfer_descs(self.bufs, mem_type=mem_type)
+
+        log.debug(
+            f"[Rank {dist_utils.get_rank()}] Registering memory for bufs {self.bufs}"
+        )
+        assert (
+            nixl_agent.register_memory(self.reg_descs) is not None
+        ), "Failed to register memory"
+
+    def deregister(self):
+        self.nixl_agent.deregister_memory(self.reg_descs)
+
+
+@dataclass
+class TrafficPattern:
+    """Represents a communication pattern between distributed processes.
+
+    Attributes:
+        matrix_file: Path to the file containing the communication matrix
+        shards: Number of shards for distributed processing
+        mem_type: Type of memory to use
+        xfer_op: Transfer operation type
+        dtype: PyTorch data type for the buffers
+        sleep_sec: Number of seconds to sleep after finish
+        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: int = 0
+    sleep_after_launch_sec: int = 0
+
+    id: str = str(uuid.uuid4())
+
+    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 senders_ranks