Add Custom Traffic Pattern benchmark by x41lakazam · Pull Request #263 · ai-dynamo/nixl

x41lakazam · 2025-05-07T13:35:57Z

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 asymetric 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.

The PR also introduce a tool script to generate the matrices that would simulate a flow of KV cache transfer, in tools/inference_workload_matgen.py.

Implementation Details

ct_perftest/sequential_custom_traffic_perftest.py implements Sequential CT Perftest
Utilizes common utilities from ct_perftest/utils.py and ct_perftest/dist_utils.py for distributed testing support
Command-line interface extensions in ct_perftest/cli.py for test configuration
ct_perftest/common.py defines an abstraction for NixlBuffers
tools/inference_workload_matgen.py - Small script to generate transfer matrices of KV transfers, with very basic computation simulation, the output of this script will be a lot of matrices and a metadata.yaml file that can be directly fed to Sequential CT Benchmark.

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:

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:

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:

Usage

Install requirements

pip install -r requirements.txt

Optionally, generate matrices using the inference_workload_matgen.py script

python 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:

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).

copy-pr-bot · 2025-05-07T13:36:00Z

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

github-actions · 2025-05-07T13:36:09Z

👋 Hi x41lakazam! Thank you for contributing to ai-dynamo/nixl.

Your PR reviewers will review your contribution then trigger the CI to test your changes.

🚀

yosefe · 2025-05-07T13:57:30Z

@x41lakazam thanks for the contribution!
Can you pls add running the new benchmark to CI scripts?

sb17v

Thanks for the PR!
Please check the pre-commit hook failure and fix the formatting errors. Initial comments are included.

sb17v · 2025-05-09T19:53:28Z

+                )
+            )
+
+        self.reg_descs = nixl_agent.get_reg_descs(self.bufs, mem_type=mem_type)


Instead of registering all individual sharded chunks, can we check if chunks are contiguous and register the entire region? We could time this part and compare the difference.

The whole point is to simulate sharded bufs

sb17v · 2025-05-09T19:54:24Z

+            )
+
+        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)


It will also be better to time this part.

Isn't this supposed to happen in parallel of the KV computation ?

x41lakazam · 2025-05-12T10:00:22Z

+    def _share_md(self) -> None:
+        """Share agent metadata between all ranks. (Need to be run after registering buffers)"""
+        md = self.nixl_agent.get_agent_metadata()
+        mds = dist_utils.allgather_obj(md)
+        for other_rank, metadata in enumerate(mds):
+            if other_rank == self.my_rank:
+                continue
+            self.nixl_agent.add_remote_agent(metadata)


@sb17v @aranadive This part takes a really really long time (sometimes 20 mins), is it the good way to do it ?

@x41lakazam, How many ranks are you using?
The steps look good to me.
Do you have the split time measurement for:

get_agent_metadata

allgather

individual and loop measurement for add_remote_agent

aranadive · 2025-05-15T07:35:48Z

@x41lakazam please fix pre-commit issues

aranadive · 2025-05-27T19:37:48Z

Next steps for this PR:
CTP will be divided into 3 parts and combined with current KVBench and made NIXLKVBench.py. We cant merge this as-is and will require some fix ups.

New Runner - Run either CTP/KVBench
Model Parameters Input - merged with what KVBench has
Workload Matrix Generator - merged as is
CTP Runner - merged as is

@aranadive will look at creating a new skeleton merge PR.

aranadive · 2025-07-13T18:43:27Z

This is merged with PR #433.

pull-request-size Bot added the size/XXL label May 7, 2025

github-actions Bot added the external-contribution label May 7, 2025

x41lakazam force-pushed the main branch 2 times, most recently from b96ed65 to 83eae3c Compare May 7, 2025 13:41

Add Custom Traffic Pattern benchmark

2db5546

x41lakazam force-pushed the main branch from 83eae3c to 2db5546 Compare May 7, 2025 13:43

Fix CI

91d35fe

sb17v reviewed May 9, 2025

View reviewed changes

sb17v requested review from aranadive, mkhazraee and tstamler May 9, 2025 19:59

x41lakazam mentioned this pull request May 11, 2025

Add python perf test engine with custom traffic perftest #84

Closed

1 task

ci/cd

1e4395b

x41lakazam force-pushed the main branch from d269ce5 to 1e4395b Compare May 12, 2025 09:59

x41lakazam commented May 12, 2025

View reviewed changes

x41lakazam added 2 commits May 14, 2025 12:36

Increase warmup and improve latency calculation of isolated

a60b33f

Add ISL params to matgen script

6566407

x41lakazam and others added 3 commits May 26, 2025 12:46

add debug log

a440c0d

Merge branch 'ai-dynamo:main' into main

ac42063

set nccl as torch dist backend

e9e2170

x41lakazam added 5 commits May 29, 2025 12:44

fix typo

89f2c47

Dockerfile

de59ef7

Use a single buffer for every transfer

81f81ae

better logging

dcdd3ec

add mpi4py to reqs

670475c

remove debug logs

2809683

aranadive mentioned this pull request Jun 27, 2025

Merge CTPerf into NIXLKVBench #433

Merged

aranadive closed this Jul 13, 2025

Conversation

x41lakazam commented May 7, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

CTPerftest - Custom Traffic Performance Test Implementation

Overview

Implementation Details

Sequential CT Perftest

Reports

Usage

Usage

Next Steps

Known Issues

Uh oh!

copy-pr-bot Bot commented May 7, 2025

Uh oh!

github-actions Bot commented May 7, 2025

Uh oh!

yosefe commented May 7, 2025

Uh oh!

sb17v left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

sb17v May 9, 2025

Choose a reason for hiding this comment

Uh oh!

x41lakazam May 14, 2025

Choose a reason for hiding this comment

Uh oh!

sb17v May 9, 2025

Choose a reason for hiding this comment

Uh oh!

x41lakazam May 14, 2025

Choose a reason for hiding this comment

Uh oh!

x41lakazam May 12, 2025

Choose a reason for hiding this comment

Uh oh!

sb17v May 12, 2025

Choose a reason for hiding this comment

Uh oh!

aranadive commented May 15, 2025

Uh oh!

aranadive commented May 27, 2025

Uh oh!

aranadive commented Jul 13, 2025

Uh oh!

Reviewers

Uh oh!

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants

x41lakazam commented May 7, 2025 •

edited

Loading