[TRITON][TEST]Add nightly test for triton split_qkv_rmsnorm_rope

[Angazenn]

What this PR does / why we need it?

Add nightly test for triton split_rmsnorm_rope

Does this PR introduce any user-facing change?

How was this patch tested?

• vLLM version: release/v0.13.0
• vLLM main: vllm-project/vllm@ad32e3e

[gemini-code-assist]

Code Review

This pull request adds a nightly test for the Triton split_rmsnorm_rope operator. While the intention is good, the new test file contains several critical bugs that will prevent it from running correctly. These include missing imports, incorrect logic in the reference implementation, incorrect variable calculations, and numerous undefined variables and typos. I have provided detailed comments and code suggestions to fix these issues. Addressing them will make the tests functional and provide valuable coverage for the operator.

[gemini-code-assist]

If norm_bias is None, this line will raise a TypeError because you cannot add None to a numpy array. You should handle the case where norm_bias is None.

Suggested change

	out = input * reciprocal_std * norm_weight + norm_bias
	out = input * reciprocal_std * norm_weight + (norm_bias if norm_bias is not None else 0.0)

[gemini-code-assist]

The kv_hidden_size is incorrectly calculated. It should be the product of num_kv_heads and head_size. The multiplication by 2 is incorrect here, as the total size of k and v tensors is accounted for later when creating the qkv tensor (kv_hidden_size * 2).

Suggested change

	kv_hidden_size = num_kv_heads * head_size * 2
	kv_hidden_size = num_kv_heads * head_size

[gemini-code-assist]

This block of code contains several undefined variables that will cause the test to fail:

• bsz is not defined. It should be num_tokens.
• cus_q and cus_k are not defined. These are likely typos for q_gold and k_gold.
• v_gold is not defined. The reference tensor for v should be _v from the split operation.

The call to custom_rope also needs to be updated to pass head_size as suggested in another comment.

    _q = _q.reshape(num_tokens, 1, -1, head_size)
    _k = _k.reshape(num_tokens, 1, -1, head_size)

    # rope
    q_gold, k_gold = custom_rope(_q, _k, sin, cos, head_size)
    q_gold = q_gold.reshape(num_tokens, -1)
    k_gold = k_gold.reshape(num_tokens, -1)

    # Compare the results.
    torch.testing.assert_close(q.to(torch.float32),
                               q_gold,
                               atol=DEFAULT_ATOL,
                               rtol=DEFAULT_RTOL)

    torch.testing.assert_close(k.to(torch.float32),
                               k_gold,
                               atol=DEFAULT_ATOL,
                               rtol=DEFAULT_RTOL)

    torch.testing.assert_close(v.to(torch.float32),
                               _v.to(torch.float32),
                               atol=DEFAULT_ATOL,
                               rtol=DEFAULT_RTOL)

[gemini-code-assist]

The kv_hidden_size is incorrectly calculated. It should be the product of num_kv_heads and head_size. The multiplication by 2 is incorrect here, as the total size of k and v tensors is accounted for later when creating the qkv tensor (kv_hidden_size * 2).

Suggested change

	kv_hidden_size = num_kv_heads * head_size * 2
	kv_hidden_size = num_kv_heads * head_size

[gemini-code-assist]

This block of code contains several undefined variables that will cause the test to fail:

• bsz is not defined. It should be num_tokens.
• cus_q and cus_k are not defined. These are likely typos for q_gold and k_gold.
• v_gold is not defined. The reference tensor for v should be _v from the split operation.

The call to custom_rope also needs to be updated to pass head_size as suggested in another comment.

    _q = _q.reshape(num_tokens, 1, -1, head_size)
    _k = _k.reshape(num_tokens, 1, -1, head_size)

    # rope
    q_gold, k_gold = custom_rope(_q, _k, sin, cos, head_size)
    q_gold = q_gold.reshape(num_tokens, -1)
    k_gold = k_gold.reshape(num_tokens, -1)

    # Compare the results.
    torch.testing.assert_close(q.to(torch.float32),
                               q_gold,
                               atol=DEFAULT_ATOL,
                               rtol=DEFAULT_RTOL)

    torch.testing.assert_close(k.to(torch.float32),
                               k_gold,
                               atol=DEFAULT_ATOL,
                               rtol=DEFAULT_RTOL)

    torch.testing.assert_close(v.to(torch.float32),
                               _v.to(torch.float32),
                               atol=DEFAULT_ATOL,
                               rtol=DEFAULT_RTOL)

[github-actions]

👋 Hi! Thank you for contributing to the vLLM Ascend project. The following points will speed up your PR merge:‌‌

• A PR should do only one thing, smaller PRs enable faster reviews.
• Every PR should include unit tests and end-to-end tests ‌to ensure it works and is not broken by other future PRs.
• Write the commit message by fulfilling the PR description to help reviewer and future developers understand.

If CI fails, you can run linting and testing checks locally according Contributing and Testing.

[Angazenn]

parallel test in #5320 .