[kernel][perf] support uncontiguous input for rms_norm kernel

[izhuhaoran]

Purpose

currently, main branch has a todo:

vllm/vllm/_custom_ops.py

Lines 331 to 332 in 14a125a

	# TODO: Remove this contiguous call when the kernel is updated to support non-contiguous input
	input_contiguous = input.contiguous()

As titiled, this pr support uncontiguous input for norm kernel and solve this todo, which is introduced by #17735 .

Previously, the RMS norm kernel required .contiguous() because q/k tensors used in qk-norm are sliced from qkv via .split() and then reshaped to [num_tokens, num_heads, head_dim]. This results in non-contiguous tensors where the first dimension has qkv's original stride. The original kernel used input.view({-1, hidden_size}), which fails or produces incorrect results for such tensors.

This PR extends the kernel to accept explicit stride information and supports both 2D and 3D non-contiguous inputs (with the last dimension required to be contiguous).

BTW, this PR should be merged after #27165

Test Result

timeline profile trace

• Main

• This PR

bench serve

setting: qwen3-0.6b, tp1, num_requests=32, max_concurrency=8, in_len=out_len=1024
result: TTFT from 86.76ms to 85.73ms, TPOT from 3.46ms to 3.31ms

• Main

============ Serving Benchmark Result ============
Successful requests:                     32        
Failed requests:                         0         
Maximum request concurrency:             8         
Benchmark duration (s):                  14.49     
Total input tokens:                      32768     
Total generated tokens:                  32768     
Request throughput (req/s):              2.21      
Output token throughput (tok/s):         2261.23   
Peak output token throughput (tok/s):    2472.00   
Peak concurrent requests:                16.00     
Total Token throughput (tok/s):          4522.47   
---------------Time to First Token----------------
Mean TTFT (ms):                          86.76     
Median TTFT (ms):                        95.76     
P99 TTFT (ms):                           100.81    
-----Time per Output Token (excl. 1st token)------
Mean TPOT (ms):                          3.46      
Median TPOT (ms):                        3.45      
P99 TPOT (ms):                           3.52      
---------------Inter-token Latency----------------
Mean ITL (ms):                           3.46      
Median ITL (ms):                         3.48      
P99 ITL (ms):                            4.00      
----------------End-to-end Latency----------------
Mean E2EL (ms):                          3621.38   
Median E2EL (ms):                        3618.82   
P99 E2EL (ms):                           3634.57   
==================================================

• This PR

============ Serving Benchmark Result ============
Successful requests:                     32        
Failed requests:                         0         
Maximum request concurrency:             8         
Benchmark duration (s):                  13.91     
Total input tokens:                      32768     
Total generated tokens:                  32768     
Request throughput (req/s):              2.30      
Output token throughput (tok/s):         2356.32   
Peak output token throughput (tok/s):    2528.00   
Peak concurrent requests:                16.00     
Total Token throughput (tok/s):          4712.63   
---------------Time to First Token----------------
Mean TTFT (ms):                          85.73     
Median TTFT (ms):                        93.82     
P99 TTFT (ms):                           97.15     
-----Time per Output Token (excl. 1st token)------
Mean TPOT (ms):                          3.31      
Median TPOT (ms):                        3.31      
P99 TPOT (ms):                           3.38      
---------------Inter-token Latency----------------
Mean ITL (ms):                           3.31      
Median ITL (ms):                         3.33      
P99 ITL (ms):                            3.87      
----------------End-to-end Latency----------------
Mean E2EL (ms):                          3475.61   
Median E2EL (ms):                        3476.51   
P99 E2EL (ms):                           3485.74   
==================================================

lm_eval

lm_eval --model local-completions --tasks gsm8k --batch_size 128 --model_args model=/mnt/data/nas/zhr/models/Qwen3-0.6B,base_url=http://localhost:8000/v1/completions,max_retries=3

• Main

local-completions (model=/mnt/data/nas/zhr/models/Qwen3-0.6B,base_url=http://localhost:8000/v1/completions,max_retries=3), gen_kwargs: (None), limit: None, num_fewshot: None, batch_size: 128
|Tasks|Version|     Filter     |n-shot|  Metric   |   |Value |   |Stderr|
|-----|------:|----------------|-----:|-----------|---|-----:|---|-----:|
|gsm8k|      3|flexible-extract|     5|exact_match|↑  |0.4071|±  |0.0135|
|     |       |strict-match    |     5|exact_match|↑  |0.4094|±  |0.0135|

• This PR

local-completions (model=/mnt/data/nas/zhr/models/Qwen3-0.6B,base_url=http://localhost:8000/v1/completions,max_retries=3), gen_kwargs: (None), limit: None, num_fewshot: None, batch_size: 128
|Tasks|Version|     Filter     |n-shot|  Metric   |   |Value |   |Stderr|
|-----|------:|----------------|-----:|-----------|---|-----:|---|-----:|
|gsm8k|      3|flexible-extract|     5|exact_match|↑  |0.4086|±  |0.0135|
|     |       |strict-match    |     5|exact_match|↑  |0.4124|±  |0.0136|

[gemini-code-assist]

Code Review

This pull request successfully adds support for non-contiguous inputs to the rms_norm kernel, which resolves a TODO and improves performance by avoiding an explicit .contiguous() call. The changes in the CUDA kernel to handle 2D and 3D tensors with explicit strides are well-implemented. My review includes one suggestion to refactor duplicated code in the C++ dispatcher function to improve maintainability.

[gemini-code-assist]

Code Review

This pull request successfully adds support for non-contiguous inputs to the RMS norm kernel, which removes a .contiguous() call and provides a performance improvement. The changes in the CUDA kernel to handle 2D and 3D non-contiguous tensors using explicit strides are well-implemented.

However, I've identified a critical issue where the output tensor out is not guaranteed to be contiguous, which will cause a runtime failure in the C++ kernel. I've left a specific comment with details on how to address this. Once that is fixed, this PR should be in a great shape.

[gemini-code-assist]

Code Review

This pull request successfully adds support for non-contiguous inputs to the RMS norm kernel, which removes a .contiguous() call and provides a performance improvement as shown in the benchmarks. The changes in the CUDA kernel correctly handle both 2D and 3D non-contiguous tensors by using explicit stride information. The Python and C++ wrapper code is updated accordingly. My feedback includes one suggestion to refactor the C++ code to improve maintainability by reducing code duplication.

[izhuhaoran]

@ProExpertProg , would you please take a look when you have time ?

[izhuhaoran]

@ProExpertProg I think this PR is ready for review. Would you please have a look?

[ProExpertProg]

Just one nit vis-a-vis dispatching

[ProExpertProg]

cc @yewentao256

[yewentao256]

Could you also test for a larger model, eg. R1 using lm_eval and vllm bench as well?

[izhuhaoran]

Could you also test for a larger model, eg. R1 using lm_eval and vllm bench as well?

Actually, could we test Qwen3-235b-fp8 instead? R1-fp8 is too large for my current hardware and would result in an OOM.
Here are the test results of Qwen3-235b-fp8:

bench serve

• Main

============ Serving Benchmark Result ============
Successful requests:                     32        
Failed requests:                         0         
Maximum request concurrency:             8         
Benchmark duration (s):                  327.34    
Total input tokens:                      32768     
Total generated tokens:                  32768     
Request throughput (req/s):              0.10      
Output token throughput (tok/s):         100.10    
Peak output token throughput (tok/s):    104.00    
Peak concurrent requests:                16.00     
Total Token throughput (tok/s):          200.21    
---------------Time to First Token----------------
Mean TTFT (ms):                          670.19    
Median TTFT (ms):                        730.73    
P99 TTFT (ms):                           746.87    
-----Time per Output Token (excl. 1st token)------
Mean TPOT (ms):                          79.34     
Median TPOT (ms):                        79.32     
P99 TPOT (ms):                           79.81     
---------------Inter-token Latency----------------
Mean ITL (ms):                           79.34     
Median ITL (ms):                         79.25     
P99 ITL (ms):                            80.00     
----------------End-to-end Latency----------------
Mean E2EL (ms):                          81831.64  
Median E2EL (ms):                        81859.35  
P99 E2EL (ms):                           81888.10  
==================================================

• this PR

============ Serving Benchmark Result ============
Successful requests:                     32        
Failed requests:                         0         
Maximum request concurrency:             8         
Benchmark duration (s):                  325.45    
Total input tokens:                      32768     
Total generated tokens:                  32768     
Request throughput (req/s):              0.10      
Output token throughput (tok/s):         100.68    
Peak output token throughput (tok/s):    104.00    
Peak concurrent requests:                16.00     
Total Token throughput (tok/s):          201.37    
---------------Time to First Token----------------
Mean TTFT (ms):                          663.54    
Median TTFT (ms):                        723.32    
P99 TTFT (ms):                           739.82    
-----Time per Output Token (excl. 1st token)------
Mean TPOT (ms):                          78.88     
Median TPOT (ms):                        78.86     
P99 TPOT (ms):                           79.35     
---------------Inter-token Latency----------------
Mean ITL (ms):                           78.88     
Median ITL (ms):                         78.80     
P99 ITL (ms):                            79.56     
----------------End-to-end Latency----------------
Mean E2EL (ms):                          81360.23  
Median E2EL (ms):                        81384.90  
P99 E2EL (ms):                           81407.92  
==================================================

lm_eval

• Main

local-completions (model=/mnt/data/nas/models/Qwen3-235B-A22B-Thinking-2507-FP8,base_url=http://localhost:8000/v1/completions,max_retries=3), gen_kwargs: (None), limit: None, num_fewshot: None, batch_size: 128
|Tasks|Version|     Filter     |n-shot|  Metric   |   |Value |   |Stderr|
|-----|------:|----------------|-----:|-----------|---|-----:|---|-----:|
|gsm8k|      3|flexible-extract|     5|exact_match|↑  |0.6732|±  |0.0129|
|     |       |strict-match    |     5|exact_match|↑  |0.6209|±  |0.0134|

• This PR

local-completions (model=/mnt/data/nas/models/Qwen3-235B-A22B-Thinking-2507-FP8,base_url=http://localhost:8000/v1/completions,max_retries=3), gen_kwargs: (None), limit: None, num_fewshot: None, batch_size: 128
|Tasks|Version|     Filter     |n-shot|  Metric   |   |Value |   |Stderr|
|-----|------:|----------------|-----:|-----------|---|-----:|---|-----:|
|gsm8k|      3|flexible-extract|     5|exact_match|↑  |0.6823|±  |0.0128|
|     |       |strict-match    |     5|exact_match|↑  |0.6217|±  |0.0134|

[izhuhaoran]

@yewentao256 I've updated the test results for the larger models and fixed the CI issues—please take a look when you have time. Also cc @ProExpertProg

BTW, there's currently a CI error: "ValueError: No available memory for the cache blocks. Try increasing gpu_memory_utilization when initializing the engine." This appears unrelated to this PR.

[yewentao256]

LGTM, thanks for the work!

[izhuhaoran]

@yewentao256 The CI failure is unrelated to this PR. The failing Plamo3 test is also failing on main and should be fixed by #29092