[K3 Perf] Flash kda out kernel for prefill, 1.1~1.4x kernel performance improvement

[yewentao256]

Purpose

Use workspace manager to preallocate the memory for _flashkda_prefill, avoid re-allocate each time when we call the kernel.

Test

Acc covered in unit tests

Perf can be seen in this AI generated script

import statistics
import time

import torch
import vllm._flashkda_C  # noqa: F401

from vllm.models.kimi_k3.nvidia.kda import _flashkda_prefill

TOKENS = (1, 32, 64, 128, 256, 512, 1024, 2048)
H, D, N = 12, 128, 8
WARMUP, ITERATIONS, REPEATS = 10, 100, 5
DEVICE = torch.device("cuda")


def benchmark(fn) -> float:
    for _ in range(WARMUP):
        fn()
    torch.accelerator.synchronize()

    samples = []
    for _ in range(REPEATS):
        start = time.perf_counter()
        for _ in range(ITERATIONS):
            fn()
        torch.accelerator.synchronize()
        samples.append((time.perf_counter() - start) * 1000 / ITERATIONS)
    return statistics.median(samples)


def run_case(T: int) -> tuple[float, float]:
    num_seqs = min(N, T)
    q, k, v, g = [
        torch.randn(1, T, H, D, dtype=torch.bfloat16, device=DEVICE) for _ in range(4)
    ]
    beta = torch.randn(1, T, H, dtype=torch.bfloat16, device=DEVICE)
    A_log = torch.randn(H, dtype=torch.float32, device=DEVICE)
    dt_bias = torch.randn(H, D, dtype=torch.float32, device=DEVICE)
    initial_state = torch.randn(num_seqs, H, D, D, dtype=torch.float32, device=DEVICE)
    cu_seqlens = torch.arange(num_seqs + 1, dtype=torch.int32, device=DEVICE) * (
        T // num_seqs
    )
    workspace_size = torch.ops._flashkda_C.get_workspace_size(T, H, num_seqs)

    core_out = torch.empty_like(v)
    state_cache = torch.empty_like(initial_state)
    final_state = torch.empty_like(initial_state)
    workspace = torch.empty(workspace_size, dtype=torch.uint8, device=DEVICE)

    def before() -> None:
        out = torch.empty_like(v)
        final = torch.empty_like(initial_state)
        ws = torch.empty(workspace_size, dtype=torch.uint8, device=DEVICE)
        torch.ops._flashkda_C.fwd(
            q,
            k,
            v,
            g,
            beta,
            D**-0.5,
            out,
            ws,
            A_log,
            dt_bias,
            -3.0,
            initial_state,
            final,
            cu_seqlens,
        )
        core_out.copy_(out)
        state_cache.copy_(final)

    def after() -> None:
        _, final = _flashkda_prefill(
            q,
            k,
            v,
            g,
            beta,
            A_log,
            dt_bias,
            -3.0,
            initial_state,
            cu_seqlens,
            core_out,
            final_state,
            workspace,
        )
        state_cache.copy_(final)

    return benchmark(before), benchmark(after)


@torch.inference_mode()
def main() -> None:
    print(f"{'tokens':>8} {'before (ms)':>12} {'after (ms)':>11} {'speedup':>9}")
    for tokens in TOKENS:
        before_ms, after_ms = run_case(tokens)
        print(
            f"{tokens:>8} {before_ms:>12.3f} {after_ms:>11.3f} "
            f"{before_ms / after_ms:>8.3f}x"
        )


if __name__ == "__main__":
    main()

And we can get

  tokens  before (ms)  after (ms)   speedup
       1        0.037       0.026    1.406x
      32        0.042       0.034    1.263x
      64        0.044       0.033    1.307x
     128        0.042       0.033    1.278x
     256        0.043       0.033    1.285x
     512        0.043       0.033    1.324x
    1024        0.042       0.033    1.283x
    2048        0.048       0.044    1.113x

[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

[yewentao256]

/ci run

[github-actions]

✅ Triggered Buildkite CI #82874 for commit 7c7d49f0389b.

[yewentao256]

/ci run

[github-actions]

✅ Triggered Buildkite CI #83165 for commit 7a116dccb56c.

[yewentao256]

/ci run

[github-actions]

✅ Triggered Buildkite CI #83350 for commit d0d3d24b7b22.

[yewentao256]

/ci run

[github-actions]

✅ Triggered Buildkite CI #83424 for commit 3880076b3f1d.

[yewentao256]

/ci run

[github-actions]

✅ Triggered Buildkite CI #83561 for commit 16288d9c6f60.

[gau-nernst]

FlashKDA requires contiguous output buffer

https://github.com/vllm-project/FlashKDA/blob/053de1b716ef3255873e02d2d28f4adf09951978/csrc/flash_kda.cpp#L49

The slicing here might cause runtime errors.