iHore/vLLM
8
0
Fork 0
Code Issues Pull Requests Actions 1 Packages Projects Releases Wiki Activity

[Perf] Fuse the mamba align block-table gather into one Triton kernel, 7 kernels -> 1, 6x speedup on all shapes for this kernel (#58737)

Browse Source 
Signed-off-by: lifulu <fululi12@amd.com>
Signed-off-by: lifulu <lifulu@users.noreply.github.com>
Co-authored-by: Claude Opus 5 (1M context) <noreply@anthropic.com>
Co-authored-by: lifulu <lifulu@users.noreply.github.com>
This commit is contained in:
fululi12
2026-10-15 11:02:53 -05:00
committed by GitHub
co-authored by Claude Opus 5 lifulu
parent a97d19d792
commit c825e1d235
2 changed files with 273 additions and 12 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tests/v1/attention/test_mamba_align_block_table.py
+149
View File
@@ -0,0 +1,149 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""Equivalence tests for the fused mamba "align" block-table gather.
`mamba_get_block_table_tensor`'s align branch used to be five PyTorch
statements that expanded to seven kernel launches over a `(num_reqs,)` tensor.
It is now one Triton kernel. These tests pin the fused kernel to the exact
semantics of the expression it replaced, including cases a single-request
decode never exercises but other callers can hit.
`block_table` is deliberately tested non-contiguous: the original
`torch.gather(block_table, 1, ...)` honours any layout, and this is a shared
utility (`gdn_attn`, `mamba_attn`, `short_conv_attn` and `linear_attn` all call
it), so a column stride of 1 must not be assumed.
"""
import pytest
import torch
from vllm.v1.attention.backends.utils import mamba_get_block_table_tensor
from vllm.v1.kv_cache_interface import MambaSpec
WIDTH = 64
def _spec(block_size: int, num_speculative_blocks: int) -> MambaSpec:
# Only block_size and num_speculative_blocks are read by the align branch;
# shapes/dtypes just have to be well formed for the frozen dataclass.
return MambaSpec(
block_size=block_size,
shapes=((1,),),
dtypes=(torch.float32,),
num_speculative_blocks=num_speculative_blocks,
)
def _reference(block_table, seq_lens, block_size, num_spec_blocks):
"""The align branch exactly as it was before the fusion."""
start_indices = (seq_lens - 1) // block_size
start_indices = start_indices.clamp(min=0)
offsets = torch.arange(
1 + num_spec_blocks, device=block_table.device, dtype=torch.int32
)
indices_to_gather = (start_indices.unsqueeze(1) + offsets).to(torch.int64)
return torch.gather(block_table, 1, indices_to_gather)
def _make(nreq, block_size, nspec, device, noncontig=False):
torch.manual_seed(0)
width = WIDTH * 2 if noncontig else WIDTH
bt = torch.randint(0, 5000, (nreq, width), dtype=torch.int32, device=device)
if noncontig:
# Every other column: stride(1) == 2, which a kernel assuming a unit
# column stride would read straight through, silently returning the
# wrong blocks.
bt = bt[:, ::2]
assert not bt.is_contiguous()
# Keep start + nspec inside the row, which the gather already requires.
hi = max(1, (WIDTH - 1 - nspec) * block_size)
sl = torch.randint(0, hi, (nreq,), dtype=torch.int32, device=device)
return bt, sl
@pytest.mark.skipif(not torch.cuda.is_available(), reason="needs a GPU")
@pytest.mark.parametrize("nreq", [1, 2, 8, 64])
@pytest.mark.parametrize("block_size", [16, 64, 256])
@pytest.mark.parametrize("nspec", [0, 1, 2, 3, 4, 7])
def test_matches_reference(nreq, block_size, nspec):
bt, sl = _make(nreq, block_size, nspec, "cuda")
got = mamba_get_block_table_tensor(bt, sl, _spec(block_size, nspec), "align")
want = _reference(bt, sl, block_size, nspec)
assert got.shape == want.shape
assert got.dtype == want.dtype
torch.testing.assert_close(got, want, atol=0, rtol=0)
@pytest.mark.skipif(not torch.cuda.is_available(), reason="needs a GPU")
@pytest.mark.parametrize("nspec", [0, 3, 7])
def test_non_contiguous_block_table(nspec):
"""A column stride != 1 must be honoured, not assumed away."""
bt, sl = _make(8, 64, nspec, "cuda", noncontig=True)
got = mamba_get_block_table_tensor(bt, sl, _spec(64, nspec), "align")
want = _reference(bt, sl, 64, nspec)
torch.testing.assert_close(got, want, atol=0, rtol=0)
@pytest.mark.skipif(not torch.cuda.is_available(), reason="needs a GPU")
def test_zero_seq_len_clamps_to_first_row():
"""CUDA-graph padding rows carry seq_len 0; the old code clamped to row 0."""
bt = torch.randint(0, 5000, (4, WIDTH), dtype=torch.int32, device="cuda")
sl = torch.zeros(4, dtype=torch.int32, device="cuda")
got = mamba_get_block_table_tensor(bt, sl, _spec(64, 3), "align")
torch.testing.assert_close(got, _reference(bt, sl, 64, 3), atol=0, rtol=0)
# start clamps to 0, so this is just the first 1 + nspec columns.
torch.testing.assert_close(got, bt[:, :4], atol=0, rtol=0)
@pytest.mark.skipif(not torch.cuda.is_available(), reason="needs a GPU")
def test_zero_requests_returns_empty():
"""An empty batch must not launch a (0,) grid."""
bt = torch.zeros((0, WIDTH), dtype=torch.int32, device="cuda")
sl = torch.zeros(0, dtype=torch.int32, device="cuda")
got = mamba_get_block_table_tensor(bt, sl, _spec(64, 3), "align")
assert got.shape == (0, 4)
assert got.dtype == bt.dtype
@pytest.mark.skipif(not torch.cuda.is_available(), reason="needs a GPU")
def test_none_mode_returns_the_input_unchanged():
"""`MambaCacheMode` is Literal["align", "none"]; only "none" passes through.
An earlier version of this test also parametrised a "mode" of "all", which
is not a member of that Literal. Anything that is not "none" takes the
align branch and returns a NEW tensor, so asserting identity for it failed
-- correctly. The passthrough contract applies to "none" alone.
"""
bt = torch.randint(0, 5000, (4, WIDTH), dtype=torch.int32, device="cuda")
sl = torch.randint(0, 1000, (4,), dtype=torch.int32, device="cuda")
assert mamba_get_block_table_tensor(bt, sl, _spec(64, 3), "none") is bt
# ---------------------------------------------------------------------------
# CPU tensors. Every test above runs on "cuda", which is exactly how the
# regression reached review: `mamba_get_block_table_tensor` is a SHARED
# utility and its callers legitimately pass CPU tensors -- the GDN and
# linear-attention metadata builders are unit-tested with
# `DEVICE = torch.device("cpu")`. A Triton kernel cannot run on those at all,
# while the torch expression it replaced could. These run without a GPU.
# ---------------------------------------------------------------------------
@pytest.mark.parametrize("nspec", [0, 1, 3])
@pytest.mark.parametrize("block_size", [1, 16, 64])
def test_cpu_tensors_match_the_reference(nspec, block_size):
"""The align branch must work on CPU, not just on the accelerator."""
width = 2 + nspec + 8
bt = torch.arange(5 * width, dtype=torch.int32).reshape(5, width)
sl = torch.tensor([0, 1, block_size, block_size * 3 + 1, 7], dtype=torch.int32)
got = mamba_get_block_table_tensor(bt, sl, _spec(block_size, nspec), "align")
want = _reference(bt, sl, block_size, nspec)
assert torch.equal(got, want), f"cpu mismatch: {got.tolist()} != {want.tolist()}"
assert got.device == bt.device
assert got.dtype == bt.dtype
def test_cpu_none_mode_passes_through():
bt = torch.arange(12, dtype=torch.int32).reshape(3, 4)
sl = torch.tensor([1, 2, 3], dtype=torch.int32)
assert mamba_get_block_table_tensor(bt, sl, _spec(16, 3), "none") is bt
vllm/v1/attention/backends/utils.py
+124 -12
View File
@@ -18,6 +18,7 @@ from typing_extensions import runtime_checkable
from vllm.config import CacheConfig, VllmConfig, get_layers_from_vllm_config
from vllm.config.cache import _layout_from_name
from vllm.triton_utils import tl, triton
from vllm.utils.gpu_sync_debug import gpu_sync_allowed
from vllm.utils.math_utils import cdiv
from vllm.utils.torch_utils import PIN_MEMORY, async_tensor_h2d, np_to_pinned_tensor
@@ -1160,6 +1161,75 @@ def get_dcp_local_seq_lens(
return dcp_local_seq_lens
@triton.jit(
do_not_specialize_on_alignment=["out_ptr", "block_table_ptr", "seq_lens_ptr"]
)
def _mamba_align_block_table_kernel(
out_ptr,
block_table_ptr,
block_table_stride_0: tl.int64,
block_table_stride_1,
seq_lens_ptr,
seq_lens_stride,
block_size,
n_out,
BLOCK: tl.constexpr,
):
"""Gather the last ``n_out`` block-table entries of each request.
One program per request. This replaces a five-statement PyTorch expression
that expanded to seven kernels over a ``(num_reqs,)`` tensor -- one element
in a single-request decode -- so almost all of its cost was launch
overhead rather than work.
Both block-table strides are passed explicitly. The expression this replaces
was ``torch.gather(block_table, 1, ...)``, which honours any layout, and
this is a shared utility -- ``gdn_attn``, ``mamba_attn``, ``short_conv_attn``
and ``linear_attn`` all call it -- so assuming a column stride of 1 would
silently read the wrong slots for a caller holding a non-contiguous view.
``kimi_k3/nvidia/kda_metadata.py`` passes both strides for the same reason.
Two annotations above are load-bearing and easy to drop by accident:
``block_table_stride_0: tl.int64``
Without it ``req_idx * block_table_stride_0`` is computed in i32 --
``program_id`` is i32 and Triton specializes a stride below 2^31 to
i32 -- so the product wraps once a row offset reaches 2^31 elements
(a 4096 x 625000 table faults from row 3436 on MI355X).
``torch.gather``, which this replaces, indexes in 64-bit and has no
such limit, so narrowing it would be a regression.
``mamba_utils.py:43`` annotates its row stride the same way.
``do_not_specialize_on_alignment`` on the three pointers
Triton puts pointer alignment in the specialization key, and the MRV1
DBO path hands this kernel ``seq_lens[request_slice]`` /
``block_table_tensor[request_slice]`` VIEWS (ubatch_utils.py:278/316).
When a ubatch starts at a request index that is not a multiple of 4
the view is differently aligned and a fresh variant JIT-compiles
mid-step -- 25-28 ms cold, ~1.5 ms warm, and an error under
``jit_monitor_mode=error``. The torch expression this replaced never
compiled anything, so that recompile would be a regression introduced
by this change.
"""
req_idx = tl.program_id(0)
seq_len = tl.load(seq_lens_ptr + req_idx * seq_lens_stride)
# A zero-length request in a CUDA graph has an invalid block table; clamp
# to 0 so it reads a valid row, exactly as the previous clamp_ did.
start = tl.maximum((seq_len - 1) // block_size, 0)
offs = tl.arange(0, BLOCK)
mask = offs < n_out
vals = tl.load(
block_table_ptr
+ req_idx * block_table_stride_0
+ (start + offs) * block_table_stride_1,
mask=mask,
other=0,
)
# `out` is freshly allocated by the caller and therefore contiguous, so its
# row stride is exactly n_out.
tl.store(out_ptr + req_idx * n_out + offs, vals, mask=mask)
def mamba_get_block_table_tensor(
block_table: torch.Tensor,
seq_lens: torch.Tensor,
@@ -1180,16 +1250,58 @@ def mamba_get_block_table_tensor(
return block_table
else:
assert isinstance(kv_cache_spec, MambaSpec)
# NOTE: For 0-length requests in CUDA graph, use a start_index of 0
# to handle the invalid block table.
start_indices = (seq_lens - 1) // kv_cache_spec.block_size
start_indices.clamp_(min=0)
# Use int32 for arithmetic to avoid dtype promotion overhead,
# then convert to int64 for gather (which requires Long indices)
offsets = torch.arange(
1 + kv_cache_spec.num_speculative_blocks,
device=block_table.device,
dtype=torch.int32,
# The arithmetic here is trivial -- a divide, a clamp, an arange, an
# add, a cast and a gather -- but as PyTorch ops it was seven kernel
# launches over a (num_reqs,) tensor, which holds a single element in a
# one-request decode. One Triton kernel does the same work in one
# launch. Zero-length requests in a CUDA graph still clamp to row 0.
num_reqs = block_table.shape[0]
n_out = 1 + kv_cache_spec.num_speculative_blocks
# The expression this replaced took its row count from `seq_lens` (a
# `torch.gather` whose index tensor is built from it), so a mismatch
# raised and a cross-device pair raised. Taking `num_reqs` from
# `block_table` instead silently changed all of that: with 4 block-table
# rows and 2 `seq_lens`, the old code returned 2 correct rows and the new
# kernel returns 4, two of them read from memory past the end of
# `seq_lens`; a pinned-CPU `seq_lens` with a GPU `block_table` used to
# raise and would now be read zero-copy from host memory.
#
# No in-tree caller does any of that today, but this helper is shared by
# gdn_attn, mamba_attn, short_conv_attn and linear_attn. Both are host
# metadata, so this costs no synchronisation, and it keeps the CPU
# branch below and the GPU kernel agreeing on the same inputs.
assert seq_lens.shape[0] == num_reqs, (
f"seq_lens has {seq_lens.shape[0]} rows, block_table has {num_reqs}"
)
indices_to_gather = (start_indices.unsqueeze(1) + offsets).to(torch.int64)
return torch.gather(block_table, 1, indices_to_gather)
assert seq_lens.device == block_table.device, (
f"seq_lens is on {seq_lens.device}, block_table on {block_table.device}"
)
if not block_table.is_cuda:
# Triton runs on the accelerator only, and this helper IS called
# with CPU tensors: the GDN and linear-attention metadata builders
# are unit-tested on CPU (`DEVICE = torch.device("cpu")`). The
# expression this replaced was plain torch and handled that, so
# keep it for the CPU path rather than narrowing the contract.
# kimi_k3/nvidia/kda_metadata.py asserts `is_cuda` instead; a
# shared utility cannot, because its callers legitimately pass CPU.
start_indices = (seq_lens - 1) // kv_cache_spec.block_size
start_indices = start_indices.clamp(min=0)
offsets = torch.arange(n_out, device=block_table.device, dtype=torch.int32)
return torch.gather(
block_table, 1, (start_indices.unsqueeze(1) + offsets).to(torch.int64)
)
out = block_table.new_empty((num_reqs, n_out))
if num_reqs == 0:
return out
_mamba_align_block_table_kernel[(num_reqs,)](
out,
block_table,
block_table.stride(0),
block_table.stride(1),
seq_lens,
seq_lens.stride(0),
kv_cache_spec.block_size,
n_out,
BLOCK=triton.next_power_of_2(n_out),
)
return out