Improving Low-Concurrency Inference Performance in vLLM

Introduction

In my previous article, I evaluated the performance and accuracy of gpt-oss-20b using both llama.cpp and vLLM on the AMD Ryzen AI Max+ 395.

https://zenn.dev/anjn/articles/5afd0b2e525c27

The results showed that vLLM's performance at low request concurrency was significantly lower than that of llama.cpp. I was not satisfied with this result, so I investigated the cause of vLLM's poor inference performance.

Investigation via PyTorch Profiler

Although I will skip the details, I used the PyTorch Profiler to capture the timeline of generating a single token in vLLM. The repeating pattern of light blue and purple indicates that the 24 Transformer blocks are being executed.

Zooming in, we can see that the execution of the matmul (matrix multiplication) kernel is taking significant time.

The following chart summarizes the execution time per kernel.

The name of the kernel taking the most time is _matmul_ogs_NNT_bf16xbf16xmxfp4_32x256x128x1_swiglu, which likely implies the following:

• Input and output types: bf16
• Weight type: mxfp4
• block_m = 32
• block_n = 256
• block_k = 128
• split_k = 1
• Fused SwiGLU activation function

The Transformer's FFN calculates products of large matrices. Instead of calculating them directly, it splits them into tiles. block_{m,n,k} represents the size of these tiles.

Optimization of Matrix Multiplication Kernels

The matrix multiplication mentioned above is performed by Triton, and it seems that how the tiles are split depending on the matrix size is determined by a heuristic method. I found that changing the block size can significantly alter the kernel execution time, potentially leading to speed improvements.

To find the optimal block size, I measured the kernel execution time while varying the concurrency (conc) and block sizes.

$ python sweep_tiles.py 
conc=1 bm=16 bn=32 bk=128 sk=1  median=10.101 ms  samples=[10.100704000000007]
conc=1 bm=16 bn=64 bk=128 sk=1  median=11.037 ms  samples=[11.036509000000002]
conc=1 bm=16 bn=128 bk=128 sk=1  median=16.656 ms  samples=[16.65572499999999]
conc=1 bm=32 bn=32 bk=128 sk=1  median=11.430 ms  samples=[11.430138999999992]
conc=1 bm=32 bn=64 bk=128 sk=1  median=24.284 ms  samples=[24.283716000000002]
conc=1 bm=32 bn=128 bk=128 sk=1  median=28.853 ms  samples=[28.853489000000042]
BEST 1 {'bm': 16, 'bn': 32, 'bk': 128, 'sk': 1, 'median_ms': 10.100704000000007}
conc=2 bm=16 bn=32 bk=128 sk=1  median=15.147 ms  samples=[15.147296000000022]
conc=2 bm=16 bn=64 bk=128 sk=1  median=17.775 ms  samples=[17.775126000000007]
conc=2 bm=16 bn=128 bk=128 sk=1  median=17.025 ms  samples=[17.024792000000016]
conc=2 bm=32 bn=32 bk=128 sk=1  median=11.497 ms  samples=[11.496612999999986]
conc=2 bm=32 bn=64 bk=128 sk=1  median=37.420 ms  samples=[37.42023899999997]
conc=2 bm=32 bn=128 bk=128 sk=1  median=29.012 ms  samples=[29.011639999999993]
BEST 2 {'bm': 32, 'bn': 32, 'bk': 128, 'sk': 1, 'median_ms': 11.496612999999986}
conc=3 bm=16 bn=32 bk=128 sk=1  median=15.170 ms  samples=[15.170093999999994]
conc=3 bm=16 bn=64 bk=128 sk=1  median=17.639 ms  samples=[17.638684999999995]
conc=3 bm=16 bn=128 bk=128 sk=1  median=24.430 ms  samples=[24.430322000000032]
conc=3 bm=32 bn=32 bk=128 sk=1  median=11.533 ms  samples=[11.532536999999977]
conc=3 bm=32 bn=64 bk=128 sk=1  median=37.682 ms  samples=[37.68203199999997]
conc=3 bm=32 bn=128 bk=128 sk=1  median=46.550 ms  samples=[46.54976800000002]
BEST 3 {'bm': 32, 'bn': 32, 'bk': 128, 'sk': 1, 'median_ms': 11.532536999999977}
conc=4 bm=16 bn=32 bk=128 sk=1  median=15.256 ms  samples=[15.256412]
conc=4 bm=16 bn=64 bk=128 sk=1  median=17.777 ms  samples=[17.77655900000001]
conc=4 bm=16 bn=128 bk=128 sk=1  median=24.520 ms  samples=[24.520478999999973]
conc=4 bm=32 bn=32 bk=128 sk=1  median=11.505 ms  samples=[11.505040000000015]
conc=4 bm=32 bn=64 bk=128 sk=1  median=37.579 ms  samples=[37.578684999999986]
conc=4 bm=32 bn=128 bk=128 sk=1  median=46.486 ms  samples=[46.486446]
BEST 4 {'bm': 32, 'bn': 32, 'bk': 128, 'sk': 1, 'median_ms': 11.505040000000015}
conc=5 bm=16 bn=32 bk=128 sk=1  median=80.182 ms  samples=[80.18241600000006]
conc=5 bm=16 bn=64 bk=128 sk=1  median=94.947 ms  samples=[94.94713600000009]
conc=5 bm=16 bn=128 bk=128 sk=1  median=15.922 ms  samples=[15.922151999999999]
conc=5 bm=32 bn=32 bk=128 sk=1  median=43.256 ms  samples=[43.25561900000002]
conc=5 bm=32 bn=64 bk=128 sk=1  median=23.605 ms  samples=[23.605318999999987]
conc=5 bm=32 bn=128 bk=128 sk=1  median=19.858 ms  samples=[19.85752999999998]
BEST 5 {'bm': 16, 'bn': 128, 'bk': 128, 'sk': 1, 'median_ms': 15.922151999999999}
conc=6 bm=16 bn=32 bk=128 sk=1  median=80.110 ms  samples=[80.10993799999994]
conc=6 bm=16 bn=64 bk=128 sk=1  median=95.799 ms  samples=[95.79853300000006]
conc=6 bm=16 bn=128 bk=128 sk=1  median=15.987 ms  samples=[15.987370999999992]
conc=6 bm=32 bn=32 bk=128 sk=1  median=43.207 ms  samples=[43.20698200000002]
conc=6 bm=32 bn=64 bk=128 sk=1  median=23.584 ms  samples=[23.583591999999996]
conc=6 bm=32 bn=128 bk=128 sk=1  median=19.790 ms  samples=[19.789950999999984]
BEST 6 {'bm': 16, 'bn': 128, 'bk': 128, 'sk': 1, 'median_ms': 15.987370999999992}
conc=7 bm=16 bn=32 bk=128 sk=1  median=80.658 ms  samples=[80.657629]
conc=7 bm=16 bn=64 bk=128 sk=1  median=96.565 ms  samples=[96.56531200000012]
conc=7 bm=16 bn=128 bk=128 sk=1  median=15.989 ms  samples=[15.989277000000014]
conc=7 bm=32 bn=32 bk=128 sk=1  median=43.456 ms  samples=[43.45628800000001]
conc=7 bm=32 bn=64 bk=128 sk=1  median=23.845 ms  samples=[23.84533799999998]
conc=7 bm=32 bn=128 bk=128 sk=1  median=20.238 ms  samples=[20.23824]
BEST 7 {'bm': 16, 'bn': 128, 'bk': 128, 'sk': 1, 'median_ms': 15.989277000000014}
conc=8 bm=16 bn=32 bk=128 sk=1  median=81.026 ms  samples=[81.02575999999995]
conc=8 bm=16 bn=64 bk=128 sk=1  median=95.400 ms  samples=[95.39969900000006]
conc=8 bm=16 bn=128 bk=128 sk=1  median=16.175 ms  samples=[16.175329999999978]
conc=8 bm=32 bn=32 bk=128 sk=1  median=43.418 ms  samples=[43.41767900000001]
conc=8 bm=32 bn=64 bk=128 sk=1  median=23.697 ms  samples=[23.697203000000005]
conc=8 bm=32 bn=128 bk=128 sk=1  median=19.975 ms  samples=[19.975342000000005]
BEST 8 {'bm': 16, 'bn': 128, 'bk': 128, 'sk': 1, 'median_ms': 16.175329999999978}

The block size for matrix multiplication is determined by the following function:

https://github.com/triton-lang/triton/blob/v3.5.0/python/triton_kernels/triton_kernels/matmul_ogs_details/opt_flags.py#L33

I determined the parameters based on the measurement results and forced the block size according to the concurrency by adding the following code at the beginning of the function.

    constraints = dict(constraints)
    if m <= 16:
        constraints["block_m"] = 32
        constraints["block_n"] = 32
        constraints["block_k"] = 128
        constraints["split_k"] = 1
    else:
        constraints["block_m"] = 16
        constraints["block_n"] = 128
        constraints["block_k"] = 128
        constraints["split_k"] = 1

I then measured the inference performance in the same way as before.

Conc.	RPS	Avg Lat.(s)	P99 Lat.(s)	Gen. toks/s	Avg TTFT(s)	P99 TTFT(s)	Avg TPOT(s)	P99 TPOT(s)	Success Rate
1	0.06	15.851	15.851	32.30	0.116	0.116	0.031	0.031	100.0%
2	0.10	19.188	19.188	53.37	0.177	0.177	0.037	0.037	100.0%
3	0.12	25.627	25.640	59.91	0.275	0.336	0.050	0.050	100.0%
4	0.16	24.882	24.895	82.26	0.326	0.385	0.048	0.048	100.0%
5	0.20	25.403	25.418	100.71	0.389	0.448	0.049	0.049	100.0%
6	0.22	26.923	26.934	113.95	0.396	0.458	0.052	0.052	100.0%
7	0.25	27.975	28.005	127.98	0.487	0.564	0.054	0.054	100.0%
8	0.29	27.591	27.617	148.31	0.551	0.626	0.053	0.054	100.0%
16	0.46	34.785	34.824	235.23	1.517	1.643	0.065	0.068	100.0%
32	0.82	39.085	39.120	418.71	1.965	2.064	0.073	0.076	100.0%

The token generation speed is plotted as follows (green indicates after optimization).

By changing the parameters of the matrix multiplication, token generation performance improved across all concurrency levels. TTFT also improved, and RPS increased significantly.

Conclusion

It turned out that the low inference performance of vLLM at low concurrency was due to suboptimal matrix multiplication parameters determined by Triton. Adjusting these parameters improved performance by about 2 to 4 times.

Compared to llama.cpp, it only falls significantly behind when concurrency is 1. Since there are other kernels taking time besides the matmul I optimized this time, there might still be room for improvement there.

Backup (Results when code modification had issues and failed to optimize properly)

I determined parameters based on the measurement results and simply added the following code.

    if tokens_per_expt <= 16:
        block_m, block_n, block_k, split_k = 32, 32, 128, 1
    elif tokens_per_expt <= 32:
        block_m, block_n, block_k, split_k = 16, 128, 128, 1    

Optimization Results

I measured inference performance in the same way as before.

Conc.	RPS	Avg Lat.(s)	P99 Lat.(s)	Gen. toks/s	Avg TTFT(s)	P99 TTFT(s)	Avg TPOT(s)	P99 TPOT(s)	Success Rate
1	0.06	15.963	15.963	32.07	0.218	0.218	0.031	0.031	100.0%
2	0.10	19.529	19.529	52.43	0.177	0.177	0.038	0.038	100.0%
3	0.12	26.022	26.035	59.00	0.545	0.686	0.050	0.050	100.0%
4	0.16	24.700	24.701	82.91	0.584	0.584	0.047	0.047	100.0%
5	0.08	58.950	58.966	43.41	0.730	0.847	0.114	0.115	100.0%
6	0.10	62.247	62.283	49.32	0.815	0.941	0.120	0.121	100.0%
7	0.10	71.727	71.762	49.94	0.911	1.032	0.139	0.140	100.0%
8	0.13	63.062	63.125	64.88	0.882	1.149	0.122	0.122	100.0%
16	0.21	74.561	74.605	109.81	3.317	3.554	0.139	0.145	100.0%
32	0.33	97.719	97.785	167.48	3.381	4.189	0.185	0.189	100.0%

It is plotted as follows (green indicates after optimization).

Inference performance improved significantly when concurrency was 4 or less. Performance dropped for other concurrency levels for some reason, so I plan to restart from a clean state.

Since there are other kernels taking time besides the Triton matmul, there might be room for optimization there as well.

Thoughts on Triton

vLLM v0.13.0 uses Triton v3.5.0's triton_kernels. Since there are changes to the logic for determining matmul block sizes in Triton v3.5.1, performance might change just by changing the Triton version. I plan to test this as well. -> Addendum: This change appears to be only for Nvidia GPUs and likely not introduced in v3.5.1.

https://github.com/triton-lang/triton/commit/0a2e3a391cbb9e13d29bf12a2a0005e358102d74