[Kernel] Support DCP for Triton backend

[frank-wei]

Purpose

As a follow up for #23734, this PR made some changes to support triton backend for DCP.
Specifically, 1) return the LSE from triton kernel 2) fix a bug in deepseekV2 which could potentially modify the residual variable.

Test Plan

export CUDA_VISIBLE_DEVICES=4,5,6,7
export VLLM_USE_V1=1
export VLLM_ATTENTION_BACKEND=TRITON_MLA
export VLLM_LOG_LEVEL=DEBUG
pytest tests/distributed/test_context_parallel.py -s

Test Result

=============================== warnings summary ===============================
:488
:488: DeprecationWarning: builtin type SwigPyPacked has no module attribute

:488
:488: DeprecationWarning: builtin type SwigPyObject has no module attribute

-- Docs: https://docs.pytest.org/en/stable/how-to/capture-warnings.html
================== 2 passed, 2 warnings in 212.42s (0:03:32) ===================

Essential Elements of an Effective PR Description Checklist

• The purpose of the PR, such as "Fix some issue (link existing issues this PR will resolve)".
• The test plan, such as providing test command.
• The test results, such as pasting the results comparison before and after, or e2e results
• (Optional) The necessary documentation update, such as updating supported_models.md and examples for a new model.
• (Optional) Release notes update. If your change is user facing, please update the release notes draft in the Google Doc.

[gemini-code-assist]

Code Review

This pull request extends Distributed Context Parallelism (DCP) support to the Triton backend by enabling the return of Log-Sum-Exp (LSE) values from attention kernels. The changes are generally well-implemented, but I have identified a critical issue where the Multi-Head Attention (MHA) path appears to be broken due to an incomplete function signature update. Additionally, a temporary test script with user-specific configurations seems to have been included by mistake and should be removed.

[houseroad]

Do we have a perf comparison?

[frank-wei]

Do we have a perf comparison?

I do not have a perf comparison. This is more on functional support especially adding LSE return in triton kernel. It won't impact the existing triton kernel performance. Now, I think vllm has completed flashMLA, FA, triton backend support for CP.

[youkaichao]

thanks! this will help running MLA on ampere or lower end GPUs.

please fix the pre-commit errors.

[mergify]

This pull request has merge conflicts that must be resolved before it can be
merged. Please rebase the PR, @frank-wei.

https://docs.github.com/en/pull-requests/collaborating-with-pull-requests/working-with-forks/syncing-a-fork