add hash key and early stop mechanism for libtuner

[meinie0826]

PR Category

Other

Type of Change

New Feature

Description

Issue

The current libtuner lacks a field representing the code, which may result in enabling outdated cache after modifying the code. To address this issue, this PR has added the md5 value of the source code as the hash key.
There is a phenomenon of excessive tuning in the way libtuner searches for the optimal config, which requires excessive computation to improve performance by 1%. Therefore, this PR has added an early stop mechanism.

Progress

• Change is properly reviewed (1 reviewer required, 2 recommended).
• Change is responded to an issue.
• Change is fully covered by a UT.

Performance

Test by the following code on H800, Through the following code test, on H800, the first end-to-end tuning time decreased from 11.51 seconds to 8.53 seconds, while the second end-to-end inference time increased from 1.16 seconds to 1.67 seconds.

# SPDX-License-Identifier: Apache-2.0
import random
import time

import numpy as np
import torch
from vllm import LLM, SamplingParams

import flag_gems


flag_gems.enable()  # Enable gems for PyTorch (aten) operators
flag_gems.apply_gems_patches_to_vllm(verbose=True)  # Patch vLLM custom ops


def set_seed(seed: int = 42):
    random.seed(seed)
    np.random.seed(seed)
    torch.manual_seed(seed)
    torch.cuda.manual_seed(seed)
    torch.cuda.manual_seed_all(seed)
    torch.backends.cudnn.deterministic = True
    torch.backends.cudnn.benchmark = False


set_seed(42)


# Sample prompts.
prompts = [
    "Hello, my name is",
    "The president of the United States is",
    "The capital of France is",
    "The future of AI is",
]

# Create a sampling params object.
sampling_params = SamplingParams(temperature=0.8, top_p=0.95, max_tokens=120)


def main():
    # Create an LLM.
    llm = LLM(
        model="Qwen/Qwen2.5-7B-Instruct",
        tensor_parallel_size=1,
        max_model_len=1024,
        gpu_memory_utilization=0.5,
        # enforce_eager=True,
    )

    # test time 2 times
    for _ in range(2):
        start_time = time.time()
        # Generate texts from the prompts. The output is a list of RequestOutput
        # objects that contain the prompt, generated text, and other information.
        outputs = llm.generate(prompts, sampling_params)
        end_time = time.time()
        print(f"Time taken for generation: {end_time - start_time:.2f} seconds")



if __name__ == "__main__":
    main()

[StrongSpoon]

It would be beneficial if you could demonstrate a performance comparison between the matmul function before and after enabling early_stop.

[StrongSpoon]

I suggest also storing the hash key of the configuration space. If the configuration space expands or shrinks, the optimal configuration may change.

[StrongSpoon]

There's no need to compare again.

[StrongSpoon]

I prefer to apply the hash key when retrieving self.cache from libcache.

[StrongSpoon]

length of timings is equal to the value of configs_tested. I suggest maintaining two scalars optimal_perf and suboptimal_perf, instead of sorting the timings list every time.

[meinie0826]

It would be beneficial if you could demonstrate a performance comparison between the matmul function before and after enabling early_stop.

As I mentioned earlier, by running the script above, the test results before and after opening are as follows:
Before:

Adding requests: 100%|███████████████████████████████████████████████████| 4/4 [00:00<00:00, 110.27it/s]
Processed prompts: 100%|█| 4/4 [00:11<00:00,  2.87s/it, est. speed input: 1.92 toks/s, output: 38.20 tok
Time taken for generation: 11.51 seconds
Adding requests: 100%|██████████████████████████████████████████████████| 4/4 [00:00<00:00, 1592.67it/s]
Processed prompts: 100%|█| 4/4 [00:01<00:00,  3.46it/s, est. speed input: 19.05 toks/s, output: 415.69 t
Time taken for generation: 1.16 seconds

After:

Adding requests: 100%|███████████████████████████████████████████████████| 4/4 [00:00<00:00, 288.53it/s]
Processed prompts: 100%|█| 4/4 [00:08<00:00,  2.13s/it, est. speed input: 2.58 toks/s, output: 51.45 tok
Time taken for generation: 8.53 seconds
Adding requests: 100%|██████████████████████████████████████████████████| 4/4 [00:00<00:00, 1959.27it/s]
Processed prompts: 100%|█| 4/4 [00:01<00:00,  2.40it/s, est. speed input: 13.18 toks/s, output: 287.53 t
Time taken for generation: 1.67 seconds

There are two running times before and after, the first time represents the running time without cache, and the second time represents the running time after tuning. In fact, this is related to the two hyperparameters early_stop_threshold and early_stop_min_comfigs. We can also consider adding a flag to choose whether to enable the early stop function. I think this function is meaningful for non long and small tasks.

[meinie0826]

Based on your suggestion, I retested the end-to-end latency

Adding requests: 100%|████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 4/4 [00:00<00:00, 294.44it/s]
Processed prompts: 100%|██████████████████████████████████████████████████████████████| 4/4 [00:08<00:00,  2.12s/it, est. speed input: 2.59 toks/s, output: 51.56 toks/s]
Time taken for generation: 8.51 seconds
Adding requests: 100%|███████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 4/4 [00:00<00:00, 1281.88it/s]
Processed prompts: 100%|████████████████████████████████████████████████████████████| 4/4 [00:01<00:00,  2.41it/s, est. speed input: 13.28 toks/s, output: 289.70 toks/s]
Time taken for generation: 1.66 seconds

[meinie0826]

I have retest on the follow script:

import torch

import flag_gems

flag_gems.enable()


def test_mm_configs():
    device = torch.device("cuda")
    nk_combinations = [
        [3584, 3584],
        [18944, 3584],
        [3584, 18944],
        [152064, 3584],
        [37888, 3584],
    ]

    m_values = list(range(1, 1025, 16))  

    total_times = [0, 0]  # [tuning_time, cached_time]

    for round_idx in range(6):
        print("=" * 50)
        for N, K in nk_combinations:
            print(f"test N={N}, K={K}")
            print("-" * 30)

            for M in m_values:
                A = torch.randn(M, K, device=device, dtype=torch.float16)
                B = torch.randn(K, N, device=device, dtype=torch.float16)
                torch.cuda.synchronize()
                start_time = torch.cuda.Event(enable_timing=True)
                end_time = torch.cuda.Event(enable_timing=True)

                start_time.record()
                C = torch.mm(A, B)
                end_time.record()
                torch.cuda.synchronize()

                exec_time = start_time.elapsed_time(end_time) 
                total_times[round_idx] += exec_time

                del A, B, C
                torch.cuda.empty_cache()

        print(f"for {round_idx+1} round time: {total_times[round_idx]:.2f}ms")

    print(f"- first round time (auto-tuning): {total_times[0]:.2f}ms")

    sum = 0
    for _ in range(5):
        sum += total_times[_+1] 
    sum /= 5
    print(f"- second round time (cached, with repeat five times): {sum:.2f}ms")

if __name__ == "__main__":
    test_mm_configs()

the result as follow, disable early stop:

- first round time (auto-tuning): 503138.81ms
- second round time (cached, with repeat five times): 167.11ms

enable early stop:

- first round time (auto-tuning): 252487.97ms
- second round time (cached, with repeat five times): 175.43ms

Therefore, we can see that early stop can greatly reduce the cost of tuning, while having a relatively small impact on the performance of the operator.

[sgjzfzzf]

I have retest on the follow script:

import torch

import flag_gems

flag_gems.enable()


def test_mm_configs():
    device = torch.device("cuda")
    nk_combinations = [
        [3584, 3584],
        [18944, 3584],
        [3584, 18944],
        [152064, 3584],
        [37888, 3584],
    ]

    m_values = list(range(1, 1025, 16))  

    total_times = [0, 0]  # [tuning_time, cached_time]

    for round_idx in range(6):
        print("=" * 50)
        for N, K in nk_combinations:
            print(f"test N={N}, K={K}")
            print("-" * 30)

            for M in m_values:
                A = torch.randn(M, K, device=device, dtype=torch.float16)
                B = torch.randn(K, N, device=device, dtype=torch.float16)
                torch.cuda.synchronize()
                start_time = torch.cuda.Event(enable_timing=True)
                end_time = torch.cuda.Event(enable_timing=True)

                start_time.record()
                C = torch.mm(A, B)
                end_time.record()
                torch.cuda.synchronize()

                exec_time = start_time.elapsed_time(end_time) 
                total_times[round_idx] += exec_time

                del A, B, C
                torch.cuda.empty_cache()

        print(f"for {round_idx+1} round time: {total_times[round_idx]:.2f}ms")

    print(f"- first round time (auto-tuning): {total_times[0]:.2f}ms")

    sum = 0
    for _ in range(5):
        sum += total_times[_+1] 
    sum /= 5
    print(f"- second round time (cached, with repeat five times): {sum:.2f}ms")

if __name__ == "__main__":
    test_mm_configs()

the result as follow, disable early stop:

- first round time (auto-tuning): 503138.81ms
- second round time (cached, with repeat five times): 167.11ms

enable early stop:

- first round time (auto-tuning): 252487.97ms
- second round time (cached, with repeat five times): 175.43ms

Therefore, we can see that early stop can greatly reduce the cost of tuning, while having a relatively small impact on the performance of the operator.

What would happen with the increase of round number? If the selected config isn't the optimal, we may lose more than the benefits we gain from the auto-tuning. Will there be any mechanism to help decide when to enable early stopping?