wehub-resource-sync
59a0a3844c
PR Test AMD / cancel-on-close (push) Has been skipped
PR Test NVIDIA ARM / scan (push) Has been skipped
PR Test NVIDIA / cancel-on-close (push) Has been skipped
PR Test AMD / scan (push) Has been skipped
PR Test NVIDIA ARM / cancel-on-close (push) Has been skipped
PR Test NVIDIA / scan (push) Has been skipped
Release Docker Images / build (cu129-torch-2.11.0) (push) Has been skipped
Release Docker Images / build (cu130-torch-2.11.0) (push) Has been skipped
Release PyPI / publish (push) Has been skipped
Scheduler Python Test / test (push) Successful in 27m19s
Docs / build (push) Successful in 28m8s
Scheduler C++ Test / test (push) Successful in 28m19s
Scheduler C++ Test / test-flat (push) Successful in 28m18s
Docs / deploy (push) Has been cancelled
PR Test AMD / finish (push) Has been cancelled
PR Test NVIDIA / finish (push) Has been cancelled
PR Test NVIDIA ARM / finish (push) Has been cancelled
PR Test NVIDIA ARM / ${{ matrix.name }} (${{ matrix.runner }}) (push) Has been cancelled
PR Test AMD / ${{ matrix.name }} (${{ matrix.runner }}) (push) Has been cancelled
PR Test NVIDIA / ${{ matrix.name }} (${{ matrix.runner }}) (push) Has been cancelled
966 lines
38 KiB
Python
966 lines
38 KiB
Python
# Copyright (c) 2026 LightSeek Foundation
|
|
#
|
|
# Permission is hereby granted, free of charge, to any person obtaining a copy
|
|
# of this software and associated documentation files (the "Software"), to deal
|
|
# in the Software without restriction, including without limitation the rights
|
|
# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
|
|
# copies of the Software, and to permit persons to whom the Software is
|
|
# furnished to do so, subject to the following conditions:
|
|
#
|
|
# The above copyright notice and this permission notice shall be included in
|
|
# all copies or substantial portions of the Software.
|
|
#
|
|
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
|
|
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
|
|
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
|
|
# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
|
|
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
|
|
# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
|
|
# SOFTWARE.
|
|
|
|
"""
|
|
MHA attention backend for TokenSpeed scheduling.
|
|
Uses fused kernels optimized for SM100 (Blackwell).
|
|
Supports sliding window, attention sinks, and FP8 KV cache.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
from collections.abc import Sequence
|
|
from dataclasses import dataclass
|
|
from typing import TYPE_CHECKING
|
|
|
|
import torch
|
|
from tokenspeed_kernel.ops.attention.flashinfer import (
|
|
trtllm_batch_context_with_kv_cache,
|
|
trtllm_batch_decode_with_kv_cache,
|
|
)
|
|
from tokenspeed_kernel.ops.kvcache.triton import (
|
|
fused_fp8_set_kv_buffer,
|
|
gather_page_table_with_padding,
|
|
)
|
|
|
|
from tokenspeed.runtime.configs.model_config import AttentionArch
|
|
from tokenspeed.runtime.execution.forward_batch_info import ForwardMode
|
|
from tokenspeed.runtime.layers.attention.backends.base import AttentionBackend
|
|
from tokenspeed.runtime.layers.attention.backends.flat_groups import (
|
|
FlatCacheGroupsMixin,
|
|
)
|
|
from tokenspeed.runtime.layers.attention.configs.mha import MHAConfig
|
|
from tokenspeed.runtime.layers.attention.registry import register_backend
|
|
from tokenspeed.runtime.layers.common import fp8_cast_contiguous
|
|
from tokenspeed.runtime.utils import get_colorful_logger
|
|
|
|
if TYPE_CHECKING:
|
|
from tokenspeed.runtime.layers.paged_attention import PagedAttention
|
|
|
|
logger = get_colorful_logger(__name__)
|
|
|
|
# Workspace buffer shared across all trtllm_mha wrappers.
|
|
_global_workspace_buffer: torch.Tensor | None = None
|
|
TRTLLM_MHA_WORKSPACE = 512 * 1024 * 1024
|
|
|
|
|
|
def canonicalize_stride(tensor: torch.Tensor) -> torch.Tensor:
|
|
"""Adjust degenerate strides for a tensor, make it canonical.
|
|
|
|
When a dimension has size=1, PyTorch may use the same stride as the next dim.
|
|
This causes TMA desc validation failures in the trtllm_mha backend.
|
|
See: https://github.com/flashinfer-ai/flashinfer/issues/2232
|
|
"""
|
|
sizes = tensor.size()
|
|
strides = tensor.stride()
|
|
ndim = tensor.dim()
|
|
|
|
need_fix = any(
|
|
sizes[i] == 1 and strides[i] == strides[i + 1] for i in range(ndim - 1)
|
|
)
|
|
|
|
if not need_fix:
|
|
return tensor
|
|
|
|
new_strides = [0] * ndim
|
|
new_strides[-1] = 1
|
|
for i in range(ndim - 2, -1, -1):
|
|
new_strides[i] = new_strides[i + 1] * sizes[i + 1]
|
|
|
|
return tensor.as_strided(sizes, new_strides)
|
|
|
|
|
|
@dataclass
|
|
class TRTLLMMHAMetadata:
|
|
cache_seqlens_int32: torch.Tensor = None
|
|
max_seq_len_q: int = 1
|
|
max_seq_len_k: int = 0
|
|
cu_seqlens_q: torch.Tensor = None
|
|
cu_seqlens_k: torch.Tensor = None
|
|
# page_table is None on the flat path (per-group page_tables route reads).
|
|
page_table: torch.Tensor = None
|
|
# Flat per-group tables/write-locs, keyed by group id (see flat_groups).
|
|
page_tables: dict[str, torch.Tensor] | None = None
|
|
out_cache_locs: dict[str, torch.Tensor] | None = None
|
|
|
|
|
|
class TRTLLMMHAAttnBackend(FlatCacheGroupsMixin, AttentionBackend):
|
|
"""trtllm_mha attention backend optimized for SM100 (Blackwell)."""
|
|
|
|
# Per-group flat tables: reads and writes route by layer.group_id, so
|
|
# slab-aliased pools (e.g. gpt-oss sliding+full pairing) are safe.
|
|
uses_flat_cache_groups: bool = True
|
|
# Graph-buffer column tails pad with the zero-init dummy page, matching
|
|
# the radix replay contract (gather_page_table_with_padding dummy_slot=0).
|
|
flat_table_tail_pad: int = 0
|
|
|
|
def support_kv_cache_prewrite(
|
|
self, forward_mode: ForwardMode | None = None
|
|
) -> bool:
|
|
return True
|
|
|
|
def _prewrite_metadata(self, forward_mode):
|
|
# Prewrite fires on extend too (unlike MHA): route it to the slot
|
|
# init_forward_metadata filled for this forward. Target verify is
|
|
# DECODE mode but its multi-token metadata lives in the prefill slot.
|
|
if forward_mode is not None and forward_mode.is_extend_or_mixed():
|
|
return self.forward_prefill_metadata
|
|
if self.spec_num_tokens > 1 and not self.is_draft:
|
|
return self.forward_prefill_metadata
|
|
return self.forward_decode_metadata
|
|
|
|
@property
|
|
def sinks_dtype(self) -> torch.dtype:
|
|
return torch.float32
|
|
|
|
def __init__(self, config: MHAConfig):
|
|
super().__init__(config)
|
|
|
|
self.page_size = config.page_size
|
|
self.max_context_len = config.context_len
|
|
self.kv_cache_dtype = config.kv_cache_dtype
|
|
max_bs = config.max_bs
|
|
|
|
# Shared workspace buffer (allocated once per process).
|
|
global _global_workspace_buffer
|
|
if _global_workspace_buffer is None:
|
|
_global_workspace_buffer = torch.zeros(
|
|
TRTLLM_MHA_WORKSPACE,
|
|
dtype=torch.uint8,
|
|
device=config.device,
|
|
)
|
|
self.workspace_buffer = _global_workspace_buffer
|
|
|
|
# Max pages per request.
|
|
self.max_num_pages = (config.context_len + self.page_size - 1) // self.page_size
|
|
|
|
# Persistent buffers for page table construction.
|
|
self.page_table_buf = torch.zeros(
|
|
(max_bs, self.max_num_pages),
|
|
dtype=torch.int32,
|
|
device=config.device,
|
|
)
|
|
self.cache_seqlens_buf = torch.zeros(
|
|
(max_bs,), dtype=torch.int32, device=config.device
|
|
)
|
|
# KV seqlens clamped to >= spec_num_tokens for the MTP verify path.
|
|
# Padded decode rows have seq_len=1 (InputBuffer); with q_len=spec_num_tokens
|
|
# they'd hit an empty causal span and the kernel returns NaN. Mirrors mha.py.
|
|
self.spec_cache_seqlens_buf = torch.zeros(
|
|
(max_bs,), dtype=torch.int32, device=config.device
|
|
)
|
|
self.cu_seqlens_q_buf = torch.zeros(
|
|
(max_bs + 1,), dtype=torch.int32, device=config.device
|
|
)
|
|
self.cu_seqlens_k_buf = torch.zeros(
|
|
(max_bs + 1,), dtype=torch.int32, device=config.device
|
|
)
|
|
|
|
# Separate slots for prefill-kernel vs decode-kernel forward paths.
|
|
# forward_extend reads prefill; forward_decode reads decode.
|
|
self.forward_prefill_metadata: TRTLLMMHAMetadata | None = None
|
|
self.forward_decode_metadata: TRTLLMMHAMetadata | None = None
|
|
|
|
# CUDA graph state — per-slot dicts.
|
|
self.cuda_graph_prefill_metadata: dict[int, TRTLLMMHAMetadata] = {}
|
|
self.cuda_graph_decode_metadata: dict[int, TRTLLMMHAMetadata] = {}
|
|
|
|
# DFLASH draft: the drafter predicts a whole block of spec_num_tokens
|
|
# per decode forward and needs non-causal (block-diffusion) attention.
|
|
# Instead of a non-causal mask, expand each request into spec_num_tokens
|
|
# single-query rows sharing the SAME block-end seq_len, so each row
|
|
# attends over the whole block. Mirrors the MHA draft_block_decode path;
|
|
# target verify and ordinary trtllm decode are untouched.
|
|
self.draft_block_decode = bool(getattr(config, "draft_block_decode", False))
|
|
|
|
# ------------------------------------------------------------------
|
|
# Page table helpers
|
|
# ------------------------------------------------------------------
|
|
|
|
def _build_page_table(
|
|
self,
|
|
req_pool_indices: torch.Tensor,
|
|
seq_lens: torch.Tensor,
|
|
bs: int,
|
|
req_to_page: torch.Tensor,
|
|
page_table_buf: torch.Tensor,
|
|
) -> torch.Tensor:
|
|
"""Build page table in [bs, max_pages] format from req_to_page.
|
|
|
|
req_to_page is [req_pool_size+1, max_pages] containing page IDs.
|
|
"""
|
|
page_table_buf[:bs].copy_(
|
|
req_to_page[req_pool_indices[:bs], : self.max_num_pages]
|
|
)
|
|
return page_table_buf[:bs]
|
|
|
|
# ------------------------------------------------------------------
|
|
# KV cache helpers
|
|
# ------------------------------------------------------------------
|
|
|
|
def _get_kv_cache_permuted(self, layer: PagedAttention, token_to_kv_pool):
|
|
"""Get KV cache in [num_pages, num_kv_heads, page_size, head_dim] layout."""
|
|
k_cache, v_cache = token_to_kv_pool.get_kv_buffer(layer.layer_id)
|
|
k_cache = k_cache.view(
|
|
-1, self.page_size, layer.tp_k_head_num, layer.head_dim
|
|
).permute(0, 2, 1, 3)
|
|
v_cache = v_cache.view(
|
|
-1, self.page_size, layer.tp_v_head_num, layer.head_dim
|
|
).permute(0, 2, 1, 3)
|
|
|
|
if layer.tp_k_head_num == 1:
|
|
k_cache = canonicalize_stride(k_cache)
|
|
if layer.tp_v_head_num == 1:
|
|
v_cache = canonicalize_stride(v_cache)
|
|
|
|
return k_cache, v_cache
|
|
|
|
def _compute_scales(self, layer: PagedAttention):
|
|
"""Compute bmm1/bmm2 scales for the fused kernel."""
|
|
q_scale = 1.0
|
|
k_scale = (
|
|
layer.k_scale_float
|
|
if getattr(layer, "k_scale_float", None) is not None
|
|
else 1.0
|
|
)
|
|
bmm1_scale = q_scale * k_scale * layer.scaling
|
|
bmm2_scale = 1.0
|
|
return bmm1_scale, bmm2_scale
|
|
|
|
def _should_use_fused_fp8_path(self, save_kv_cache: bool, k) -> bool:
|
|
return (
|
|
save_kv_cache
|
|
and k is not None
|
|
and self.kv_cache_dtype == torch.float8_e4m3fn
|
|
)
|
|
|
|
# ------------------------------------------------------------------
|
|
# Forward
|
|
# ------------------------------------------------------------------
|
|
|
|
def _save_kv_and_prepare_q(
|
|
self, q, k, v, layer, out_cache_loc, token_to_kv_pool, save_kv_cache
|
|
):
|
|
if self._should_use_fused_fp8_path(save_kv_cache, k):
|
|
k_cache, v_cache = token_to_kv_pool.get_kv_buffer(layer.layer_id)
|
|
fused_fp8_set_kv_buffer(
|
|
k=k.view(-1, layer.tp_k_head_num, layer.head_dim),
|
|
v=v.view(-1, layer.tp_k_head_num, layer.head_dim),
|
|
k_cache=k_cache,
|
|
v_cache=v_cache,
|
|
cache_loc=out_cache_loc,
|
|
k_scale=layer.k_scale,
|
|
v_scale=layer.v_scale,
|
|
page_size=self.page_size,
|
|
)
|
|
elif save_kv_cache and k is not None:
|
|
token_to_kv_pool.set_kv_buffer(
|
|
layer, out_cache_loc, k, v, layer.k_scale, layer.v_scale
|
|
)
|
|
|
|
if self.kv_cache_dtype == torch.float8_e4m3fn:
|
|
q = fp8_cast_contiguous(q)
|
|
else:
|
|
q = q.contiguous()
|
|
|
|
return q.view(-1, layer.tp_q_head_num, layer.head_dim)
|
|
|
|
def forward_decode(
|
|
self,
|
|
q: torch.Tensor,
|
|
k: torch.Tensor,
|
|
v: torch.Tensor,
|
|
layer: PagedAttention,
|
|
out_cache_loc: torch.Tensor,
|
|
token_to_kv_pool,
|
|
bs: int,
|
|
save_kv_cache: bool = True,
|
|
**kwargs,
|
|
) -> torch.Tensor:
|
|
if self.draft_block_decode and self.spec_num_tokens > 1:
|
|
# DFLASH draft block: metadata is expanded to bs*spec_num_tokens
|
|
# single-query rows, so use the decode slot directly. Inferring
|
|
# q_len_per_req from q.shape[0]//bs would be spec_num_tokens and
|
|
# wrongly pick the prefill slot.
|
|
metadata = self.forward_decode_metadata
|
|
else:
|
|
# Multi-token decode (q_len > 1) reads the prefill slot's
|
|
# uniform-stride metadata; plain decode reads the single-token slot.
|
|
q_len_per_req = q.shape[0] // bs if bs > 0 else 1
|
|
metadata = (
|
|
self.forward_prefill_metadata
|
|
if q_len_per_req > 1
|
|
else self.forward_decode_metadata
|
|
)
|
|
|
|
out_cache_loc = self._select_out_cache_loc(
|
|
layer, metadata, out_cache_loc, prefer_caller=self.is_draft
|
|
)
|
|
q = self._save_kv_and_prepare_q(
|
|
q, k, v, layer, out_cache_loc, token_to_kv_pool, save_kv_cache
|
|
)
|
|
k_cache, v_cache = self._get_kv_cache_permuted(layer, token_to_kv_pool)
|
|
bmm1_scale, bmm2_scale = self._compute_scales(layer)
|
|
|
|
attention_sink = kwargs.get("sinks", None)
|
|
if attention_sink is not None:
|
|
attention_sink = attention_sink.float()
|
|
|
|
o = trtllm_batch_decode_with_kv_cache(
|
|
query=q,
|
|
kv_cache=(k_cache, v_cache),
|
|
workspace_buffer=self.workspace_buffer,
|
|
block_tables=self._select_page_table(layer, metadata),
|
|
seq_lens=metadata.cache_seqlens_int32,
|
|
max_seq_len=self.max_context_len,
|
|
bmm1_scale=bmm1_scale,
|
|
bmm2_scale=bmm2_scale,
|
|
window_left=layer.sliding_window_size,
|
|
sinks=attention_sink,
|
|
out_dtype=self.dtype,
|
|
q_len_per_req=metadata.max_seq_len_q,
|
|
)
|
|
return o.view(-1, layer.tp_q_head_num * layer.head_dim)
|
|
|
|
def forward_extend(
|
|
self,
|
|
q: torch.Tensor,
|
|
k: torch.Tensor,
|
|
v: torch.Tensor,
|
|
layer: PagedAttention,
|
|
out_cache_loc: torch.Tensor,
|
|
token_to_kv_pool,
|
|
bs: int,
|
|
save_kv_cache: bool = True,
|
|
**kwargs,
|
|
) -> torch.Tensor:
|
|
metadata = self.forward_prefill_metadata
|
|
out_cache_loc = self._select_out_cache_loc(layer, metadata, out_cache_loc)
|
|
q = self._save_kv_and_prepare_q(
|
|
q, k, v, layer, out_cache_loc, token_to_kv_pool, save_kv_cache
|
|
)
|
|
k_cache, v_cache = self._get_kv_cache_permuted(layer, token_to_kv_pool)
|
|
bmm1_scale, bmm2_scale = self._compute_scales(layer)
|
|
|
|
attention_sink = kwargs.get("sinks", None)
|
|
if attention_sink is not None:
|
|
attention_sink = attention_sink.float()
|
|
|
|
o = trtllm_batch_context_with_kv_cache(
|
|
query=q,
|
|
kv_cache=(k_cache, v_cache),
|
|
workspace_buffer=self.workspace_buffer,
|
|
block_tables=self._select_page_table(layer, metadata),
|
|
seq_lens=metadata.cache_seqlens_int32,
|
|
max_q_len=metadata.max_seq_len_q,
|
|
max_kv_len=self.max_context_len,
|
|
bmm1_scale=bmm1_scale,
|
|
bmm2_scale=bmm2_scale,
|
|
batch_size=metadata.cu_seqlens_q.shape[0] - 1,
|
|
cum_seq_lens_q=metadata.cu_seqlens_q,
|
|
cum_seq_lens_kv=metadata.cu_seqlens_k,
|
|
window_left=layer.sliding_window_size,
|
|
sinks=attention_sink,
|
|
out_dtype=self.dtype,
|
|
)
|
|
return o.view(-1, layer.tp_q_head_num * layer.head_dim)
|
|
|
|
# ------------------------------------------------------------------
|
|
# Metadata initialisation
|
|
# ------------------------------------------------------------------
|
|
|
|
def init_forward_metadata(
|
|
self,
|
|
bs: int,
|
|
req_pool_indices: torch.Tensor,
|
|
seq_lens: torch.Tensor,
|
|
forward_mode: ForwardMode,
|
|
req_to_page: torch.Tensor,
|
|
extend_with_prefix: bool = False,
|
|
extend_prefix_lens: torch.Tensor | None = None,
|
|
extend_prefix_lens_cpu: torch.Tensor | None = None,
|
|
extend_seq_lens_cpu: torch.Tensor | None = None,
|
|
spec_info=None,
|
|
use_cuda_graph: bool = False,
|
|
flat_block_tables: dict[str, torch.Tensor] | None = None,
|
|
**kwargs,
|
|
):
|
|
flat_page_tables = self._shed_state_groups(flat_block_tables)
|
|
flat_out_cache_locs = None
|
|
if flat_page_tables:
|
|
# Verify keeps [bs]-row tables; only DFLASH expands rows. TODO(flat+dflash).
|
|
assert not (
|
|
self.draft_block_decode and self.spec_num_tokens > 1
|
|
), "flat cache groups are unsupported with DFLASH block decode"
|
|
if forward_mode.is_extend_or_mixed():
|
|
assert extend_prefix_lens_cpu is not None
|
|
assert extend_seq_lens_cpu is not None
|
|
flat_out_cache_locs = self._compute_flat_extend_out_cache_locs(
|
|
flat_page_tables,
|
|
extend_prefix_lens_cpu[:bs],
|
|
extend_seq_lens_cpu[:bs],
|
|
self.page_size,
|
|
)
|
|
else:
|
|
flat_out_cache_locs = self._compute_flat_decode_out_cache_locs(
|
|
flat_page_tables,
|
|
seq_lens[:bs],
|
|
self.page_size,
|
|
self._flat_verify_tokens(),
|
|
)
|
|
self._maybe_check_flat_write_locs(
|
|
flat_page_tables, flat_out_cache_locs, self.page_size
|
|
)
|
|
|
|
if forward_mode.is_extend_or_mixed():
|
|
self._init_extend_metadata(
|
|
bs,
|
|
req_pool_indices,
|
|
seq_lens,
|
|
req_to_page,
|
|
extend_with_prefix=extend_with_prefix,
|
|
extend_prefix_lens=extend_prefix_lens,
|
|
extend_prefix_lens_cpu=extend_prefix_lens_cpu,
|
|
extend_seq_lens_cpu=extend_seq_lens_cpu,
|
|
flat_page_tables=flat_page_tables,
|
|
flat_out_cache_locs=flat_out_cache_locs,
|
|
)
|
|
# Drafter: also fill decode_metadata so step 1+ multi-step has
|
|
# metadata under EXTEND/MIXED target. seq_lens is the drafter's
|
|
# live alias buffer (wrapper pre-writes before this call).
|
|
if self.is_draft:
|
|
self._init_decode_metadata(bs, req_pool_indices, seq_lens, req_to_page)
|
|
return
|
|
|
|
if self.draft_block_decode and self.spec_num_tokens > 1:
|
|
# DFLASH draft block (eager): expand to spec_num_tokens single-query
|
|
# rows per request; seq_lens is the block-end length the drafter
|
|
# already wrote.
|
|
self._init_block_decode_metadata(
|
|
bs, req_pool_indices, seq_lens, req_to_page
|
|
)
|
|
return
|
|
|
|
if self.spec_num_tokens > 1:
|
|
self._init_multi_token_metadata(
|
|
bs,
|
|
self.spec_num_tokens,
|
|
req_pool_indices,
|
|
seq_lens,
|
|
req_to_page,
|
|
flat_page_tables=flat_page_tables,
|
|
flat_out_cache_locs=flat_out_cache_locs,
|
|
)
|
|
if self.is_draft:
|
|
# Drafter's N-1 single-token steps after the first.
|
|
self._init_decode_metadata(
|
|
bs,
|
|
req_pool_indices,
|
|
seq_lens,
|
|
req_to_page,
|
|
flat_page_tables=flat_page_tables,
|
|
flat_out_cache_locs=flat_out_cache_locs,
|
|
)
|
|
else:
|
|
self._init_decode_metadata(
|
|
bs,
|
|
req_pool_indices,
|
|
seq_lens,
|
|
req_to_page,
|
|
flat_page_tables=flat_page_tables,
|
|
flat_out_cache_locs=flat_out_cache_locs,
|
|
)
|
|
|
|
def _init_decode_metadata(
|
|
self,
|
|
bs: int,
|
|
req_pool_indices: torch.Tensor,
|
|
seq_lens: torch.Tensor,
|
|
req_to_page: torch.Tensor,
|
|
flat_page_tables: dict[str, torch.Tensor] | None = None,
|
|
flat_out_cache_locs: dict[str, torch.Tensor] | None = None,
|
|
):
|
|
assert (
|
|
seq_lens.dtype == torch.int32
|
|
), f"seq_lens must be int32, got {seq_lens.dtype}"
|
|
device = seq_lens.device
|
|
# Alias seq_lens (no copy, no mutation). cu_seqlens_k omitted:
|
|
# the decode kernel doesn't read it. On the flat path the per-group
|
|
# tables route every read; the radix single table would be dead work.
|
|
self.forward_decode_metadata = TRTLLMMHAMetadata(
|
|
cache_seqlens_int32=seq_lens[:bs],
|
|
max_seq_len_q=1,
|
|
max_seq_len_k=self.max_context_len,
|
|
cu_seqlens_q=torch.arange(0, bs + 1, dtype=torch.int32, device=device),
|
|
page_table=(
|
|
None
|
|
if flat_page_tables
|
|
else self._build_page_table(
|
|
req_pool_indices, seq_lens, bs, req_to_page, self.page_table_buf
|
|
)
|
|
),
|
|
page_tables=flat_page_tables,
|
|
out_cache_locs=flat_out_cache_locs,
|
|
)
|
|
|
|
def _replicate_block_page_table(
|
|
self,
|
|
out: torch.Tensor,
|
|
req_pool_indices: torch.Tensor,
|
|
bs: int,
|
|
req_to_page: torch.Tensor,
|
|
) -> None:
|
|
"""Replicate each request's page table to its spec_num_tokens block rows.
|
|
|
|
``out`` is the [bs*spec_num_tokens, max_num_pages] destination. All block
|
|
rows of a request share its pages (decode only reads KV), so a single
|
|
broadcast copy suffices.
|
|
"""
|
|
spec = self.spec_num_tokens
|
|
base_page_table = req_to_page[req_pool_indices[:bs], : self.max_num_pages]
|
|
out[: bs * spec, :].view(bs, spec, self.max_num_pages).copy_(
|
|
base_page_table[:, None, :]
|
|
)
|
|
|
|
def _init_block_decode_metadata(
|
|
self,
|
|
bs: int,
|
|
req_pool_indices: torch.Tensor,
|
|
seq_lens: torch.Tensor,
|
|
req_to_page: torch.Tensor,
|
|
):
|
|
"""Eager DFLASH draft-block metadata: spec_num_tokens single-query rows
|
|
per request, all carrying the block-end seq_len (prefix + spec_num_tokens)
|
|
so each query attends over the whole block. Allocates fresh buffers (the
|
|
cuda-graph path uses persistent ones), mirroring the MHA backend.
|
|
"""
|
|
assert (
|
|
seq_lens.dtype == torch.int32
|
|
), f"seq_lens must be int32, got {seq_lens.dtype}"
|
|
spec = self.spec_num_tokens
|
|
device = seq_lens.device
|
|
expanded_bs = bs * spec
|
|
|
|
page_table = torch.empty(
|
|
(expanded_bs, self.max_num_pages), dtype=torch.int32, device=device
|
|
)
|
|
self._replicate_block_page_table(page_table, req_pool_indices, bs, req_to_page)
|
|
|
|
# Clamp the block-end length so the decode never asks for more page-table
|
|
# columns than exist (prefix + spec_num_tokens can exceed max_context_len).
|
|
cache_seqlens = (
|
|
seq_lens[:bs]
|
|
.clamp(spec, self.max_context_len)
|
|
.unsqueeze(1)
|
|
.expand(bs, spec)
|
|
.reshape(expanded_bs)
|
|
.contiguous()
|
|
)
|
|
|
|
self.forward_decode_metadata = TRTLLMMHAMetadata(
|
|
cache_seqlens_int32=cache_seqlens,
|
|
max_seq_len_q=1,
|
|
max_seq_len_k=self.max_context_len,
|
|
cu_seqlens_q=torch.arange(
|
|
0, expanded_bs + 1, dtype=torch.int32, device=device
|
|
),
|
|
page_table=page_table,
|
|
)
|
|
|
|
def _clamped_spec_seqlens(
|
|
self, seq_lens: torch.Tensor, bs: int, spec_num_tokens: int
|
|
) -> torch.Tensor:
|
|
"""Return KV seqlens clamped to >= spec_num_tokens for the MTP verify path.
|
|
|
|
Writes into the persistent spec_cache_seqlens_buf (CUDA-graph safe)
|
|
to avoid NaN from empty causal spans on padded rows (seq_len=1).
|
|
"""
|
|
dst = self.spec_cache_seqlens_buf[:bs]
|
|
torch.clamp_min(seq_lens[:bs], spec_num_tokens, out=dst)
|
|
return dst
|
|
|
|
def _flat_verify_tokens(self) -> int:
|
|
# Write locs per request on the flat path: N for target verify
|
|
# ([bs*N], token-major), 1 elsewhere (draft chains use caller locs).
|
|
return (
|
|
self.spec_num_tokens
|
|
if self.spec_num_tokens > 1 and not self.is_draft
|
|
else 1
|
|
)
|
|
|
|
def _init_multi_token_metadata(
|
|
self,
|
|
bs: int,
|
|
spec_num_tokens: int,
|
|
req_pool_indices: torch.Tensor,
|
|
seq_lens: torch.Tensor,
|
|
req_to_page: torch.Tensor,
|
|
flat_page_tables: dict[str, torch.Tensor] | None = None,
|
|
flat_out_cache_locs: dict[str, torch.Tensor] | None = None,
|
|
):
|
|
"""Prefill-slot metadata for multi-token decode (uniform q_len per
|
|
request). Routes through the decode kernel via q_len_per_req; the
|
|
kernel doesn't read cu_seqlens_k."""
|
|
assert (
|
|
seq_lens.dtype == torch.int32
|
|
), f"seq_lens must be int32, got {seq_lens.dtype}"
|
|
device = seq_lens.device
|
|
self.forward_prefill_metadata = TRTLLMMHAMetadata(
|
|
cache_seqlens_int32=self._clamped_spec_seqlens(
|
|
seq_lens, bs, spec_num_tokens
|
|
),
|
|
max_seq_len_q=spec_num_tokens,
|
|
max_seq_len_k=self.max_context_len,
|
|
cu_seqlens_q=torch.arange(
|
|
0,
|
|
bs * spec_num_tokens + 1,
|
|
spec_num_tokens,
|
|
dtype=torch.int32,
|
|
device=device,
|
|
),
|
|
page_table=(
|
|
None
|
|
if flat_page_tables
|
|
else self._build_page_table(
|
|
req_pool_indices, seq_lens, bs, req_to_page, self.page_table_buf
|
|
)
|
|
),
|
|
page_tables=flat_page_tables,
|
|
out_cache_locs=flat_out_cache_locs,
|
|
)
|
|
|
|
def _init_extend_metadata(
|
|
self,
|
|
bs: int,
|
|
req_pool_indices: torch.Tensor,
|
|
seq_lens: torch.Tensor,
|
|
req_to_page: torch.Tensor,
|
|
extend_with_prefix: bool = False,
|
|
extend_prefix_lens: torch.Tensor | None = None,
|
|
extend_prefix_lens_cpu=None,
|
|
extend_seq_lens_cpu=None,
|
|
flat_page_tables: dict[str, torch.Tensor] | None = None,
|
|
flat_out_cache_locs: dict[str, torch.Tensor] | None = None,
|
|
):
|
|
"""Populate prefill slot for regular EXTEND (ragged query)."""
|
|
assert (
|
|
seq_lens.dtype == torch.int32
|
|
), f"seq_lens must be int32, got {seq_lens.dtype}"
|
|
assert (
|
|
extend_seq_lens_cpu is not None
|
|
), "trtllm extend requires extend_seq_lens_cpu (pinned-CPU mirror) to avoid GPU sync"
|
|
cache_seqlens_int32 = seq_lens[:bs]
|
|
cu_seqlens_k = torch.nn.functional.pad(
|
|
torch.cumsum(seq_lens, dim=0, dtype=torch.int32), (1, 0)
|
|
)
|
|
# Flat path: per-group tables route every read (see _init_decode_metadata).
|
|
page_table = (
|
|
None
|
|
if flat_page_tables
|
|
else self._build_page_table(
|
|
req_pool_indices, seq_lens, bs, req_to_page, self.page_table_buf
|
|
)
|
|
)
|
|
|
|
# Read the max from the pinned-CPU mirror — avoids a per-iter
|
|
# GPU->CPU sync that would block the host on the previous step's
|
|
# forward and erase prefill/decode overlap. Both branches want
|
|
# max(new tokens per request); for a no-prefix extend that's
|
|
# seq_lens, for a prefix-cached extend it's seq_lens-prefix_lens —
|
|
# extend_seq_lens_cpu holds those new-token counts in either case.
|
|
max_seq_len_q = int(extend_seq_lens_cpu[:bs].max().item())
|
|
|
|
if extend_with_prefix and (
|
|
(extend_prefix_lens_cpu is not None and any(extend_prefix_lens_cpu))
|
|
or (extend_prefix_lens is not None and any(extend_prefix_lens.tolist()))
|
|
):
|
|
if extend_prefix_lens is None:
|
|
raise RuntimeError(
|
|
"TRTLLMMHAAttnBackend requires extend_prefix_lens tensor "
|
|
"when extend_with_prefix is true."
|
|
)
|
|
extend_seq_lens = seq_lens - extend_prefix_lens
|
|
cu_seqlens_q = torch.nn.functional.pad(
|
|
torch.cumsum(extend_seq_lens, dim=0, dtype=torch.int32), (1, 0)
|
|
)
|
|
else:
|
|
cu_seqlens_q = cu_seqlens_k
|
|
|
|
self.forward_prefill_metadata = TRTLLMMHAMetadata(
|
|
cache_seqlens_int32=cache_seqlens_int32,
|
|
max_seq_len_q=max_seq_len_q,
|
|
max_seq_len_k=self.max_context_len,
|
|
cu_seqlens_q=cu_seqlens_q,
|
|
cu_seqlens_k=cu_seqlens_k,
|
|
page_table=page_table,
|
|
page_tables=flat_page_tables,
|
|
out_cache_locs=flat_out_cache_locs,
|
|
)
|
|
|
|
# ------------------------------------------------------------------
|
|
# CUDA graph support
|
|
# ------------------------------------------------------------------
|
|
|
|
def init_cuda_graph_state(
|
|
self,
|
|
max_bs: int,
|
|
seq_lens_buf: torch.Tensor,
|
|
paged_cache_group_specs: Sequence = (),
|
|
**kwargs,
|
|
):
|
|
assert (
|
|
seq_lens_buf.dtype == torch.int32
|
|
and seq_lens_buf.dim() == 1
|
|
and seq_lens_buf.shape[0] >= max_bs
|
|
), (
|
|
f"seq_lens_buf must be int32 with shape[0] >= {max_bs}, "
|
|
f"got {seq_lens_buf.dtype} {tuple(seq_lens_buf.shape)}"
|
|
)
|
|
self.cuda_graph_prefill_metadata = {}
|
|
self.cuda_graph_decode_metadata = {}
|
|
# Flat per-group persistent buffers + state-group shed; before the
|
|
# DFLASH early return (replay reads the dict for the stale guard).
|
|
self._learn_flat_state_groups(paged_cache_group_specs)
|
|
self._init_flat_graph_buffers(max_bs)
|
|
if self.draft_block_decode and self.spec_num_tokens > 1:
|
|
# DFLASH draft block: spec_num_tokens decode rows per request. Unlike
|
|
# the plain path, cache_seqlens is a dedicated buffer (NOT aliasing
|
|
# seq_lens_buf): it is filled in-graph by fill_block_decode_seq_lens.
|
|
self.cuda_graph_page_table = torch.zeros(
|
|
(max_bs * self.spec_num_tokens, self.max_num_pages),
|
|
dtype=torch.int32,
|
|
device=self.device,
|
|
)
|
|
self.cuda_graph_cache_seqlens = torch.full(
|
|
(max_bs * self.spec_num_tokens,),
|
|
self.spec_num_tokens,
|
|
dtype=torch.int32,
|
|
device=self.device,
|
|
)
|
|
return
|
|
# Alias controller's seq_lens_buf — backend never mutates it.
|
|
self.cuda_graph_page_table = torch.zeros(
|
|
(max_bs, self.max_num_pages), dtype=torch.int32, device=self.device
|
|
)
|
|
self.cuda_graph_cache_seqlens = seq_lens_buf
|
|
|
|
def init_forward_metadata_capture_cuda_graph(
|
|
self,
|
|
bs: int,
|
|
req_pool_indices: torch.Tensor,
|
|
seq_lens: torch.Tensor,
|
|
forward_mode: ForwardMode,
|
|
flat_cache_group_ids: tuple[str, ...] = (),
|
|
**kwargs,
|
|
):
|
|
if forward_mode.is_extend_or_mixed():
|
|
raise NotImplementedError(
|
|
f"trtllm CUDA graph capture not supported for {forward_mode}"
|
|
)
|
|
|
|
# Real tables only arrive at replay: capture lazily allocates
|
|
# persistent per-group buffers and records metadata views into them.
|
|
if flat_cache_group_ids:
|
|
# Verify keeps [bs]-row tables + [bs*N] loc views. TODO(flat+dflash).
|
|
assert not (
|
|
self.draft_block_decode and self.spec_num_tokens > 1
|
|
), "flat_cache_group_ids is unsupported with DFLASH block decode"
|
|
page_tables, out_cache_locs = self._flat_capture_group_views(
|
|
bs, flat_cache_group_ids, tokens_per_req=self._flat_verify_tokens()
|
|
)
|
|
|
|
if self.draft_block_decode and self.spec_num_tokens > 1:
|
|
self._init_block_decode_metadata_capture(bs)
|
|
return
|
|
|
|
if self.spec_num_tokens > 1:
|
|
self._init_multi_token_metadata_capture(
|
|
bs, self.spec_num_tokens, page_tables, out_cache_locs
|
|
)
|
|
if self.is_draft:
|
|
self._init_decode_metadata_capture(
|
|
bs, seq_lens, page_tables, out_cache_locs
|
|
)
|
|
else:
|
|
self._init_decode_metadata_capture(
|
|
bs, seq_lens, page_tables, out_cache_locs
|
|
)
|
|
|
|
def _init_block_decode_metadata_capture(self, bs: int):
|
|
"""DFLASH draft block (cuda-graph capture): spec_num_tokens single-query
|
|
rows per request over the persistent expanded buffers. seq_lens are
|
|
filled in-graph by fill_block_decode_seq_lens; seed a safe baseline here
|
|
so the capture run stays in range before that op records."""
|
|
expanded_bs = bs * self.spec_num_tokens
|
|
self.cuda_graph_cache_seqlens[:expanded_bs].fill_(self.spec_num_tokens)
|
|
metadata = TRTLLMMHAMetadata(
|
|
cache_seqlens_int32=self.cuda_graph_cache_seqlens[:expanded_bs],
|
|
max_seq_len_q=1,
|
|
max_seq_len_k=self.max_context_len,
|
|
cu_seqlens_q=torch.arange(
|
|
0, expanded_bs + 1, dtype=torch.int32, device=self.device
|
|
),
|
|
page_table=self.cuda_graph_page_table[:expanded_bs, :],
|
|
)
|
|
self.cuda_graph_decode_metadata[bs] = metadata
|
|
self.forward_decode_metadata = metadata
|
|
|
|
def _init_decode_metadata_capture(
|
|
self,
|
|
bs: int,
|
|
seq_lens: torch.Tensor,
|
|
page_tables: dict[str, torch.Tensor] | None = None,
|
|
out_cache_locs: dict[str, torch.Tensor] | None = None,
|
|
):
|
|
# cache_seqlens aliases seq_lens_buf (set in init_cuda_graph_state).
|
|
# Flat captures route reads through the per-group buffer views and
|
|
# replay never fills the radix single table, so record page_table=None
|
|
# instead of a slice of the never-filled zero buffer.
|
|
metadata = TRTLLMMHAMetadata(
|
|
cache_seqlens_int32=self.cuda_graph_cache_seqlens[:bs],
|
|
max_seq_len_q=1,
|
|
max_seq_len_k=self.max_context_len,
|
|
cu_seqlens_q=torch.arange(0, bs + 1, dtype=torch.int32, device=self.device),
|
|
page_table=(
|
|
None if page_tables is not None else self.cuda_graph_page_table[:bs, :]
|
|
),
|
|
page_tables=page_tables,
|
|
out_cache_locs=out_cache_locs,
|
|
)
|
|
self.cuda_graph_decode_metadata[bs] = metadata
|
|
self.forward_decode_metadata = metadata
|
|
|
|
def _init_multi_token_metadata_capture(
|
|
self,
|
|
bs: int,
|
|
spec_num_tokens: int,
|
|
page_tables: dict[str, torch.Tensor] | None = None,
|
|
out_cache_locs: dict[str, torch.Tensor] | None = None,
|
|
):
|
|
# Multi-token decode: seed spec_cache_seqlens_buf (clamped to >=
|
|
# spec_num_tokens) at capture so padded rows (seq_len=1) avoid NaN.
|
|
# The replay path refreshes it each step.
|
|
cache_seqlens = self._clamped_spec_seqlens(
|
|
self.cuda_graph_cache_seqlens, bs, spec_num_tokens
|
|
)
|
|
metadata = TRTLLMMHAMetadata(
|
|
cache_seqlens_int32=cache_seqlens,
|
|
max_seq_len_q=spec_num_tokens,
|
|
max_seq_len_k=self.max_context_len,
|
|
cu_seqlens_q=torch.arange(
|
|
0,
|
|
bs * spec_num_tokens + 1,
|
|
spec_num_tokens,
|
|
dtype=torch.int32,
|
|
device=self.device,
|
|
),
|
|
page_table=(
|
|
None if page_tables is not None else self.cuda_graph_page_table[:bs, :]
|
|
),
|
|
page_tables=page_tables,
|
|
out_cache_locs=out_cache_locs,
|
|
)
|
|
self.cuda_graph_prefill_metadata[bs] = metadata
|
|
self.forward_prefill_metadata = metadata
|
|
|
|
def init_forward_metadata_replay_cuda_graph(
|
|
self,
|
|
bs: int,
|
|
req_pool_indices: torch.Tensor,
|
|
seq_lens: torch.Tensor,
|
|
forward_mode: ForwardMode,
|
|
req_to_page: torch.Tensor = None,
|
|
flat_block_tables: dict[str, torch.Tensor] | None = None,
|
|
**kwargs,
|
|
):
|
|
if forward_mode.is_extend_or_mixed():
|
|
raise NotImplementedError(
|
|
f"trtllm CUDA graph replay not supported for {forward_mode}"
|
|
)
|
|
|
|
# Fail loudly instead of replaying over stale/zero page tables.
|
|
self._flat_replay_stale_guard(bs, flat_block_tables)
|
|
|
|
if self.draft_block_decode and self.spec_num_tokens > 1:
|
|
# DFLASH draft block: replicate the page table to each request's
|
|
# block rows. seq_lens are re-derived in-graph, so not touched here.
|
|
if req_to_page is not None:
|
|
self._replicate_block_page_table(
|
|
self.cuda_graph_page_table, req_pool_indices, bs, req_to_page
|
|
)
|
|
if bs in self.cuda_graph_decode_metadata:
|
|
self.forward_decode_metadata = self.cuda_graph_decode_metadata[bs]
|
|
return
|
|
|
|
# cache_seqlens aliases seq_lens_buf; only page tables need refresh.
|
|
# Flat captures read only the per-group buffers; the radix single
|
|
# table would be dead work there (and req_to_page is unpopulated on
|
|
# a flat scheduler build).
|
|
if not self.cuda_graph_flat_page_tables and req_to_page is not None:
|
|
gather_page_table_with_padding(
|
|
req_to_page=req_to_page,
|
|
req_pool_indices=req_pool_indices,
|
|
seq_lens=seq_lens,
|
|
out=self.cuda_graph_page_table,
|
|
bs=bs,
|
|
max_num_pages=self.max_num_pages,
|
|
page_size=self.page_size,
|
|
dummy_slot=0,
|
|
)
|
|
if flat_block_tables:
|
|
# cuda_graph_cache_seqlens aliases the controller's seq_lens_buf,
|
|
# filled by input prep BEFORE this call.
|
|
self._flat_replay_fill(
|
|
bs,
|
|
flat_block_tables,
|
|
self.cuda_graph_cache_seqlens,
|
|
tokens_per_req=self._flat_verify_tokens(),
|
|
)
|
|
|
|
# Refresh for both verify and draft: draft step 1 is multi-token
|
|
# and reads spec_cache_seqlens_buf; later single-token steps don't.
|
|
if self.spec_num_tokens > 1:
|
|
self._clamped_spec_seqlens(seq_lens, bs, self.spec_num_tokens)
|
|
|
|
if bs in self.cuda_graph_prefill_metadata:
|
|
self.forward_prefill_metadata = self.cuda_graph_prefill_metadata[bs]
|
|
if bs in self.cuda_graph_decode_metadata:
|
|
self.forward_decode_metadata = self.cuda_graph_decode_metadata[bs]
|
|
|
|
def fill_block_decode_seq_lens(self, bs: int, block_seq_lens: torch.Tensor) -> None:
|
|
"""DFLASH: broadcast each request's block-end length to its
|
|
spec_num_tokens cuda-graph decode rows.
|
|
|
|
Called by the drafter inside the captured graph so every replay
|
|
re-derives cache_seqlens from the live draft length. Mirrors the MHA
|
|
backend method of the same name.
|
|
|
|
Args:
|
|
bs: Number of draft requests.
|
|
block_seq_lens: ``[bs]`` per-request block-end lengths.
|
|
"""
|
|
spec = self.spec_num_tokens
|
|
self.cuda_graph_cache_seqlens[: bs * spec].view(bs, spec).copy_(
|
|
block_seq_lens[:bs].clamp(spec, self.max_context_len).unsqueeze(1)
|
|
)
|
|
|
|
|
|
register_backend("trtllm", {AttentionArch.MHA}, TRTLLMMHAAttnBackend)
|