From c825e1d235896694ca50160286440859ed46d7d3 Mon Sep 17 00:00:00 2001 From: fululi12 Date: Thu, 15 Oct 2026 09:02:53 -0700 Subject: [PATCH] [Perf] Fuse the mamba align block-table gather into one Triton kernel, 7 kernels -> 1, 6x speedup on all shapes for this kernel (#58737) Signed-off-by: lifulu Signed-off-by: lifulu Co-authored-by: Claude Opus 5 (1M context) Co-authored-by: lifulu --- .../attention/test_mamba_align_block_table.py | 149 ++++++++++++++++++ vllm/v1/attention/backends/utils.py | 136 ++++++++++++++-- 2 files changed, 273 insertions(+), 12 deletions(-) create mode 100644 tests/v1/attention/test_mamba_align_block_table.py diff --git a/tests/v1/attention/test_mamba_align_block_table.py b/tests/v1/attention/test_mamba_align_block_table.py new file mode 100644 index 0000000000..0380771388 --- /dev/null +++ b/tests/v1/attention/test_mamba_align_block_table.py @@ -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 diff --git a/vllm/v1/attention/backends/utils.py b/vllm/v1/attention/backends/utils.py index f93156d53f..7e0a7fb644 100644 --- a/vllm/v1/attention/backends/utils.py +++ b/vllm/v1/attention/backends/utils.py @@ -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