[K3 Perf] Optimize k3 mamba metadata preparation, 6.6~7.6x kernel performance improvement

[yewentao256]

Purpose

Optimize Kimi K3 Mamba align metadata preparation by computing aligned state indices for all KV-cache groups in one Triton launch.

Before:

prepare_attn
  -> builder group 0 -> allocate buffer -> launch kernel
  -> builder group 1 -> allocate buffer -> launch kernel
  -> builder group N -> allocate buffer -> launch kernel

After:

prepare_attn
  -> persistent output buffer
  -> one multi-group Triton kernel
  -> directly use the buffer

Test

Acc covered in current unit tests

Perf can be seen in this AI generated script

from types import SimpleNamespace

import torch

from vllm.models.kimi_k3.nvidia.kda_metadata import (
    _mamba_get_block_table_tensor,
)
from vllm.triton_utils import triton
from vllm.v1.worker.mamba_utils import (
    get_aligned_state_indices_multi_group_kernel,
)

# Decode-only benchmark: one token per request, so tokens == batch size.
TOKENS = [1, 2, 4, 8, 16, 32, 64, 128, 256]
NUM_GROUPS = 4
BLOCK_SIZE = 32
MAX_BLOCKS = 1024
BLOCK_ROWS = 32
DEVICE = "cuda"


def main() -> None:
    max_tokens = max(TOKENS)
    seq_lens = torch.randint(
        1,
        BLOCK_SIZE * MAX_BLOCKS + 1,
        (max_tokens,),
        dtype=torch.int32,
        device=DEVICE,
    )
    block_tables = [
        torch.randint(
            0,
            1_000_000,
            (max_tokens, MAX_BLOCKS),
            dtype=torch.int32,
            device=DEVICE,
        )
        for _ in range(NUM_GROUPS)
    ]
    block_table_ptrs = torch.tensor(
        [table.data_ptr() for table in block_tables],
        dtype=torch.int64,
        device=DEVICE,
    )
    output = torch.empty((NUM_GROUPS, max_tokens, 1), dtype=torch.int32, device=DEVICE)
    spec = SimpleNamespace(block_size=BLOCK_SIZE, num_speculative_blocks=0)

    print("tokens  old (us)  new (us)  saved (us)  speedup")
    for tokens in TOKENS:
        current_seq_lens = seq_lens[:tokens]

        def old(
            tokens: int = tokens,
            current_seq_lens: torch.Tensor = current_seq_lens,
        ) -> None:
            for table in block_tables:
                _mamba_get_block_table_tensor(
                    table[:tokens], current_seq_lens, spec, "align"
                )

        def new(
            tokens: int = tokens,
            current_seq_lens: torch.Tensor = current_seq_lens,
        ) -> None:
            get_aligned_state_indices_multi_group_kernel[
                (triton.cdiv(tokens, BLOCK_ROWS),)
            ](
                block_table_ptrs,
                current_seq_lens,
                output,
                block_tables[0].stride(0),
                current_seq_lens.stride(0),
                output.stride(0),
                output.stride(1),
                output.stride(2),
                tokens,
                CACHE_BLOCK_SIZE=BLOCK_SIZE,
                NUM_GROUPS=NUM_GROUPS,
                BLOCK_GROUPS=triton.next_power_of_2(NUM_GROUPS),
                NUM_STATE_SLOTS=1,
                BLOCK_STATE_SLOTS=1,
                BLOCK_ROWS=BLOCK_ROWS,
                num_warps=1,
            )

        expected = torch.stack(
            [
                _mamba_get_block_table_tensor(
                    table[:tokens], current_seq_lens, spec, "align"
                )
                for table in block_tables
            ]
        )
        new()
        torch.testing.assert_close(output[:, :tokens], expected)

        old_us = triton.testing.do_bench(old, warmup=100, rep=500) * 1000
        new_us = triton.testing.do_bench(new, warmup=100, rep=500) * 1000
        print(
            f"{tokens:>6}  {old_us:>8.2f}  {new_us:>8.2f}  "
            f"{old_us - new_us:>10.2f}  {old_us / new_us:>7.2f}x"
        )


if __name__ == "__main__":
    main()

And we get

tokens  old (us)  new (us)  saved (us)  speedup
     1     49.76      7.12       42.65     6.99x
     2     48.91      7.14       41.77     6.85x
     4     48.40      7.14       41.26     6.78x
     8     50.25      7.12       43.12     7.05x
    16     48.42      7.12       41.30     6.80x
    32     47.47      7.12       40.35     6.67x
    64     50.72      7.12       43.60     7.12x
   128     49.79      7.13       42.66     6.98x
   256     53.86      7.13       46.72     7.55x

[claude]

Claude Code Review

This repository is configured for manual code reviews. Comment @claude review for a one-time review, or @claude review always to subscribe this PR to a review on every future push.

_{Tip: disable this comment in your organization's Code Review settings.}

[yewentao256]

/ci run

[github-actions]

✅ Triggered Buildkite CI #83973 for commit beba6f92b2b9.

[mergify]

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

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

[yewentao256]

/ci run

[github-actions]

✅ Triggered Buildkite CI #84243 for commit 5413a661e54d.

[yewentao256]

/ci run

[github-actions]

✅ Triggered Buildkite CI #84383 for commit fa59e4d6f9aa.

[mergify]

Hi @yewentao256, 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.

[yewentao256]

/ci run

[github-actions]

✅ Triggered Buildkite CI #84443 for commit 8e3f0aed52dc.

[tlrmchlsmth]

Apparently when VLLM_USE_V2_MODEL_RUNNER=0 and mamba_cache_mode == "align", the V1 model runner will hit this. Could you take a look?

[yewentao256]

Nice catch, fixed!

[tlrmchlsmth]

Wondering if there's any cleanup we can do now, or in the future once MRV1 goes away -- Can we remove _mamba_get_block_table_tensor?

[yewentao256]

Let's keep it for now as fallback will need it, I also add a TODO to remove them when MRv2 is used by default.

[yewentao256]

/ci run

[github-actions]

✅ Triggered Buildkite CI #85384 for commit 32e358483856.