Hi @morrison-turnansky, the pre-commit checks have failed. Please run:

uv pip install pre-commit>=4.5.1
pre-commit install
pre-commit run --all-files

Then, commit the changes and push to your branch.

For future commits, pre-commit will run automatically on changed files before each commit.

# For mypy (substitute "3.10" with the failing version if needed)
pre-commit run --hook-stage manual mypy-3.10

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

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

current performance:

vllm serve deepseek-ai/DeepSeek-V2-Lite
--trust-remote-code
-cc.mode=VLLM_COMPILE
-cc.cudagraph_mode=FULL_AND_PIECEWISE
-cc.use_inductor_graph_partition=true
--profiler-config '{"profiler": "torch", "torch_profiler_dir": "./vllm_profile"}'

vllm bench serve --backend vllm --model deepseek-ai/DeepSeek-V2-Lite --dataset-name sharegpt --num-prompts 128 --profile --dataset-path /workspaces/vllm-dev/vllm/sharegpt.json

profiler_out_0.txt

============ Serving Benchmark Result ============

Successful requests: 128
Failed requests: 0
Benchmark duration (s): 9.83
Total input tokens: 31879
Total generated tokens: 27379
Request throughput (req/s): 13.02
Output token throughput (tok/s): 2784.11
Peak output token throughput (tok/s): 6243.00
Peak concurrent requests: 128.00
Total token throughput (tok/s): 6025.82
---------------Time to First Token----------------
Mean TTFT (ms): 543.37
Median TTFT (ms): 625.09
P99 TTFT (ms): 645.95
-----Time per Output Token (excl. 1st token)------
Mean TPOT (ms): 17.00
Median TPOT (ms): 12.66
P99 TPOT (ms): 86.88
---------------Inter-token Latency----------------
Mean ITL (ms): 11.92
Median ITL (ms): 11.58
P99 ITL (ms): 23.41

==================================================

pr performance:

for repro:
VLLM_MLA_EXPOSED_SPLIT=1
vllm serve deepseek-ai/DeepSeek-V2-Lite
--trust-remote-code
-cc.mode=VLLM_COMPILE
-cc.cudagraph_mode=FULL_AND_PIECEWISE
-cc.use_inductor_graph_partition=true
--profiler-config '{"profiler": "torch", "torch_profiler_dir": "./vllm_profile"}'

vllm bench serve --backend vllm --model deepseek-ai/DeepSeek-V2-Lite --dataset-name sharegpt --num-prompts 128 --profile --dataset-path /workspaces/vllm-dev/vllm/sharegpt.json

profiler_out_0.txt

============ Serving Benchmark Result ============

Successful requests: 128
Failed requests: 0
Benchmark duration (s): 15.27
Total input tokens: 31879
Total generated tokens: 26575
Request throughput (req/s): 8.38
Output token throughput (tok/s): 1740.69
Peak output token throughput (tok/s): 6028.00
Peak concurrent requests: 128.00
Total token throughput (tok/s): 3828.79
---------------Time to First Token----------------
Mean TTFT (ms): 3135.06
Median TTFT (ms): 3175.92
P99 TTFT (ms): 3304.33
-----Time per Output Token (excl. 1st token)------
Mean TPOT (ms): 37.26
Median TPOT (ms): 24.28
P99 TPOT (ms): 129.87
---------------Inter-token Latency----------------
Mean ITL (ms): 21.72
Median ITL (ms): 11.85
P99 ITL (ms): 28.04

==================================================

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

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

Hi @morrison-turnansky, the pre-commit checks have failed. Please run:

uv pip install pre-commit>=4.5.1
pre-commit install
pre-commit run --all-files

Then, commit the changes and push to your branch.

For future commits, pre-commit will run automatically on changed files before each commit.

# For mypy (substitute "3.10" with the failing version if needed)
pre-commit run --hook-stage manual mypy-3.10

Hi @morrison-turnansky, the pre-commit checks have failed. Please run:

uv pip install pre-commit>=4.5.1
pre-commit install
pre-commit run --all-files

Then, commit the changes and push to your branch.

For future commits, pre-commit will run automatically on changed files before each commit.

# For mypy (substitute "3.10" with the failing version if needed)
pre-commit run --hook-stage manual mypy-3.10

Hi @morrison-turnansky, the pre-commit checks have failed. Please run:

uv pip install pre-commit>=4.5.1
pre-commit install
pre-commit run --all-files

Then, commit the changes and push to your branch.

For future commits, pre-commit will run automatically on changed files before each commit.

# For mypy (substitute "3.10" with the failing version if needed)
pre-commit run --hook-stage manual mypy-3.10

Hi @morrison-turnansky, the pre-commit checks have failed. Please run:

uv pip install pre-commit>=4.5.1
pre-commit install
pre-commit run --all-files

Then, commit the changes and push to your branch.

For future commits, pre-commit will run automatically on changed files before each commit.

# For mypy (substitute "3.10" with the failing version if needed)
pre-commit run --hook-stage manual mypy-3.10

Btw I had this idea on how to unify the exposed/non-exposed paths, using the wrap_if_exposed decorator:

def wrap_if_exposed(op_name: str):
    def decorator(func):
        # could optionally register the custom op automatically here as well, 
        # but that might be more effort than is worth

        @wraps(func)
        def wrapper(self: "MLAAttention", *args, **kwargs):
            if not self.exposed:
                return func(self, *args, **kwargs)
            return getattr(torch.ops.vllm, op_name)(self.layer_name, *args, **kwargs)
        return wrapper
    return decorator


class MLAAttention(...):
    def __init__(self):
        self.exposed = ...
        ...

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        if self.exposed:
            return self.inner_forward(x)
        return torch.ops.vllm.mla_attention(self.layer_name, x)

    def inner_forward(self, x: torch.Tensor) -> torch.Tensor:
        """Inner forward with decomposed method calls."""
        num_prefills = self.split_batch()

        output = torch.empty_like(x)
        output[num_prefills:] = self.forward_mha(x[num_prefills:])
        output[:num_prefills] = self.forward_mqa(x[:num_prefills])

        return output

    def forward_mqa(self, x: torch.Tensor) -> torch.Tensor:
        ...

    @wrap_if_exposed("mla_forward_mha")
    def forward_mha(self, x: torch.Tensor) -> torch.Tensor:
        ...

Hi @morrison-turnansky, the pre-commit checks have failed. Please run:

uv pip install pre-commit>=4.5.1
pre-commit install
pre-commit run --all-files

Then, commit the changes and push to your branch.

For future commits, pre-commit will run automatically on changed files before each commit.

# For mypy (substitute "3.10" with the failing version if needed)
pre-commit run --hook-stage manual mypy-3.10

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

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