minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

feat(unified-memory): read unified pool from attention backends fa3/flashinfer/trtllm_mha/flashmla (#34613)

Browse Source 
Co-authored-by: Caihua Li <caihua.li@bytedance.com>
Co-authored-by: Claude Fable 5 <noreply@anthropic.com>
Co-authored-by: Cheng Wan <cheng.wan@radixark.ai>
This commit is contained in:
caihuali95
2026-08-30 23:58:24 -07:00
committed by GitHub
co-authored by Caihua Li Claude Fable 5 Cheng Wan
parent 29578d5578
commit 8bb776dc48
31 changed files with 1182 additions and 757 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/kernels/ops/attention/metadata.py
+52 -52
View File
@@ -193,13 +193,9 @@ def _fused_metadata_kernel_general(
use_swa: tl.constexpr,
SHIFT: tl.constexpr,
BLOCK_COLS: tl.constexpr,
# Unified-memory per-layer-view path (page-major envelope shared with the mamba
# sub-pool). Both default to the identity for the statically-partitioned
# pool, where req_to_token already holds physical ids.
v2p_ptr=None,
PAGE_MULT: tl.constexpr = 1,
# Unified SWA puts its independent v2p table in the legacy mapping slot.
SWA_MAPPING_IS_V2P: tl.constexpr = False,
# 1: the two table pointers carry PAGE-granular, already kernel-facing
# read tables; emit verbatim -- no >>SHIFT, no v2p, no mapping gather.
SRC_IS_KERNEL_PAGE_TABLE: tl.constexpr = 0,
):
pid_b = tl.program_id(0) # batch index
pid_c = tl.program_id(1) # column chunk index
@@ -240,8 +236,11 @@ def _fused_metadata_kernel_general(
col_offsets = col_start + tl.arange(0, BLOCK_COLS)
mask = col_offsets < num_live_pages
# Compute column indices in the source tensor (token offset)
if page_size == 1:
# Compute column indices in the source tensor (token offset; page offset
# when the source is already the page-granular canonical)
if SRC_IS_KERNEL_PAGE_TABLE:
col_idx = col_offsets
elif page_size == 1:
col_idx = col_offsets
else:
col_idx = col_offsets << SHIFT # faster than multiplication for power-of-two
@@ -253,37 +252,35 @@ def _fused_metadata_kernel_general(
)
# Compute page_table
if page_size == 1:
if SRC_IS_KERNEL_PAGE_TABLE:
page_table_val = page_index # read-table entries are the page ids
elif page_size == 1:
page_table_val = page_index
else:
page_table_val = page_index >> SHIFT
# Unified memory: virtual page -> physical page -> that layer's dense block.
# Derived from page_table_val, NOT page_index, which the SWA branch below
# still needs in virtual space. Masked so padded lanes never index the table.
if v2p_ptr is not None:
page_table_val = tl.load(v2p_ptr + page_table_val, mask=mask, other=0)
page_table_val = page_table_val * PAGE_MULT
# Store to page_table
pt_offsets = i * page_table_stride_0 + col_offsets * page_table_stride_1
tl.store(page_table + pt_offsets, page_table_val, mask=mask, cache_modifier=".cg")
if use_swa:
if SWA_MAPPING_IS_V2P:
if page_size == 1:
swa_mapping_index = page_index
else:
swa_mapping_index = page_index >> SHIFT
else:
swa_mapping_index = page_index * full_to_swa_mapping_stride_0
swa_slot = tl.load(
full_to_swa_mapping + swa_mapping_index,
if SRC_IS_KERNEL_PAGE_TABLE:
# The swa canonical shares the full canonical's shape and strides,
# so the SAME rt_offsets address the matching swa entry.
swa_val = tl.load(
full_to_swa_mapping + rt_offsets,
mask=mask,
other=0,
cache_modifier=".cg",
)
if page_size == 1 or SWA_MAPPING_IS_V2P:
else:
swa_slot = tl.load(
full_to_swa_mapping + page_index * full_to_swa_mapping_stride_0,
mask=mask,
other=0,
cache_modifier=".cg",
)
if page_size == 1:
swa_val = swa_slot
else:
swa_val = swa_slot >> SHIFT
@@ -316,9 +313,6 @@ def _fused_metadata_kernel_ps1_no_swa(
max_seq_pages,
seq_len_delta: tl.constexpr,
BLOCK_COLS: tl.constexpr,
# Unified-memory per-layer-view path; identity defaults for the static pool.
v2p_ptr=None,
PAGE_MULT: tl.constexpr = 1,
):
pid_b = tl.program_id(0) # batch index
pid_c = tl.program_id(1) # column chunk index
@@ -361,11 +355,6 @@ def _fused_metadata_kernel_ps1_no_swa(
req_to_token + rt_offsets, mask=mask, other=0, cache_modifier=".cg"
)
# page_table = page_index // 1 = page_index
# Unified memory: at page_size 1 the virtual token id IS the virtual page id.
if v2p_ptr is not None:
page_index = tl.load(v2p_ptr + page_index, mask=mask, other=0)
page_index = page_index * PAGE_MULT
pt_offsets = i * page_table_stride_0 + col_offsets * page_table_stride_1
tl.store(page_table + pt_offsets, page_index, mask=mask, cache_modifier=".cg")
@@ -586,15 +575,14 @@ def normal_decode_set_metadata(
page_table: torch.Tensor,
req_to_token: torch.Tensor,
req_pool_indices: torch.Tensor,
strided_indices: torch.Tensor,
max_seq_pages: torch.Tensor,
seq_lens: torch.Tensor,
seq_len_delta: int,
page_size: int,
swa_page_table: Optional[torch.Tensor] = None,
token_to_kv_pool: Optional["SWAKVPool"] = None,
v2p_page_table: Optional[torch.Tensor] = None,
kernel_page_multiplier: int = 1,
src_is_read_table: bool = False,
swa_src_table: Optional[torch.Tensor] = None,
):
"""
Fused Triton implementation that replaces 4-5 sequential CUDA kernels with 1-2 kernels:
@@ -602,14 +590,15 @@ def normal_decode_set_metadata(
2. cu_seqlens_k = cumsum(cache_seqlens) (prefix-sum)
3. page_indices = req_to_token[pool_idx, stride_idx] (2-D gather)
4. page_table = page_indices // page_size (floor-divide)
4b. (unified memory) page_table = v2p_page_table[page] * kernel_page_multiplier
5. (optional) swa_page_table via the legacy full->SWA map or the unified
SWA pool's independent page map
5. (optional) swa_page_table for sliding window attention
Step 4b is folded in rather than applied afterwards so the capture-stable
page_table is written already translated: no separate pass a caller could
forget, and no temporary to keep pointer-stable across cuda-graph replays.
Identity (None / 1) for the statically-partitioned pool.
Unified pool (``src_is_read_table=True``): ``req_to_token`` /
``req_pool_indices`` carry the translator's PAGE-granular read table and its
row indices instead (entries already kernel-facing; ``swa_src_table`` is
the swa canonical, same shape and strides); steps 3-5 become verbatim
copies of the read table's rows' live prefixes, folded into the same launch
so the capture-stable page_table is written translated with no separate
pass a caller could forget.
Achieves ~5.2x speedup on H200 hardware for typical decode workloads.
@@ -639,7 +628,9 @@ def normal_decode_set_metadata(
page_table_stride_0 = page_table.stride(0)
page_table_stride_1 = page_table.stride(1)
use_swa = swa_page_table is not None and token_to_kv_pool is not None
use_swa = swa_page_table is not None and (
token_to_kv_pool is not None or swa_src_table is not None
)
# Unified SWA uses an independent SWA v2p table.
swa_v2p_page_table = None
@@ -681,15 +672,26 @@ def normal_decode_set_metadata(
max_seq_pages,
seq_len_delta,
BLOCK_COLS=BLOCK_COLS,
v2p_ptr=v2p_page_table,
PAGE_MULT=kernel_page_multiplier,
num_warps=8,
num_stages=3,
)
else:
# General kernel for page_size > 1 or SWA cases
# SWA parameters
if use_swa:
if use_swa and src_is_read_table:
# Unified pool: the swa canonical rides in the mapping slot; the
# kernel addresses it with the SAME row/col offsets as the full
# canonical, so their layouts must match exactly.
assert swa_src_table is not None
assert (
swa_src_table.stride() == req_to_token.stride()
), "swa canonical must share the full canonical's strides"
swa_page_table = swa_page_table.contiguous()
swa_page_table_stride_0 = swa_page_table.stride(0)
swa_page_table_stride_1 = swa_page_table.stride(1)
full_to_swa_mapping = swa_src_table
full_to_swa_mapping_stride_0 = 0 # unused under the canonical source
elif use_swa:
from sglang.srt.mem_cache.swa_memory_pool import SWAKVPool
swa_page_table = swa_page_table.contiguous()
@@ -749,9 +751,7 @@ def normal_decode_set_metadata(
use_swa,
shift,
BLOCK_COLS=BLOCK_COLS,
v2p_ptr=v2p_page_table,
PAGE_MULT=kernel_page_multiplier,
SWA_MAPPING_IS_V2P=swa_uses_v2p,
SRC_IS_KERNEL_PAGE_TABLE=1 if src_is_read_table else 0,
num_warps=4,
num_stages=3,
)
python/sglang/kernels/ops/kvcache/kv_indices.py
+3 -15
View File
@@ -129,16 +129,9 @@ def create_flashmla_kv_indices_triton(
req_to_token_ptr_stride: tl.constexpr,
kv_indices_ptr_stride: tl.constexpr,
PAGED_SIZE: tl.constexpr = 64,
# Unified-memory per-layer-view path (page-major envelope shared with the mamba
# sub-pool). req_to_token holds VIRTUAL token ids; the block table the MLA
# kernel consumes must hold kernel-facing page ids. When v2p_ptr is given, map each
# virtual page through it to the physical page, then scale by PAGE_MULT
# (= num MLA layers) so the entry addresses the layer's per-page block
# in the (num_pages*L, page_size, kv_cache_dim) reshaped view. Both default
# to the identity (v2p_ptr None, PAGE_MULT 1) for the static pool.
v2p_ptr=None,
PAGE_MULT: tl.constexpr = 1,
):
# Static-pool builder only: token ids here are physical, entry = token//ps.
# The unified pool's block table is filled by KVIndexTranslator instead.
NUM_PAGE_PER_BLOCK: tl.constexpr = (
FLASHMLA_CREATE_KV_BLOCK_SIZE_TRITON // PAGED_SIZE
)
@@ -178,13 +171,8 @@ def create_flashmla_kv_indices_triton(
+ paged_offset,
mask=mask,
)
page = data // PAGED_SIZE
if v2p_ptr is not None:
# virtual page -> physical page (page-level v2p); masked so padded
# lanes never index the table out of bounds.
page = tl.load(v2p_ptr + page, mask=mask_out, other=0)
tl.store(
kv_indices_ptr + pid * kv_indices_ptr_stride + paged_offset_out,
page * PAGE_MULT,
data // PAGED_SIZE,
mask=mask_out,
)
python/sglang/kernels/ops/kvcache/trtllm_mha_graph_metadata.py
+19 -2
View File
@@ -62,19 +62,25 @@ def update_trtllm_mha_graph_metadata_kernel(
Q_MODE: tl.constexpr,
PAGE_BLOCK: tl.constexpr,
BS_BLOCK: tl.constexpr,
# 1: the page tables are refreshed out-of-graph, so rebuild only the
# seqlen metadata. 0 = static pool, full rebuild.
SKIP_PAGE_TABLE: tl.constexpr = 0,
):
pid = tl.program_id(axis=0)
if pid < bs:
# One program per batch row: cache_seqlens + page table row(s).
req_pool_index = tl.load(req_pool_indices_ptr + pid).to(tl.int64)
seqlen = (tl.load(seq_lens_ptr + pid) + seqlen_offset).to(tl.int32)
tl.store(cache_seqlens_ptr + pid, seqlen)
if not SKIP_PAGE_TABLE:
req_pool_index = tl.load(req_pool_indices_ptr + pid).to(tl.int64)
row_in = req_to_token_ptr + req_pool_index * req_to_token_stride
row_out = page_table_ptr + pid.to(tl.int64) * page_table_stride
if HAS_SWA:
swa_row_out = swa_page_table_ptr + pid.to(tl.int64) * swa_page_table_stride
swa_row_out = (
swa_page_table_ptr + pid.to(tl.int64) * swa_page_table_stride
)
# Self-guard on the device-side seqlen: pages past cdiv(cache_seqlen,
# PAGE_SIZE) keep stale values the attention kernels never read.
num_live_pages = tl.minimum(tl.cdiv(seqlen, PAGE_SIZE), max_seq_pages)
@@ -145,12 +151,18 @@ def update_trtllm_mha_graph_metadata(
qlens=None,
q_stride: int = 0,
q_mode: int = Q_MODE_NONE,
skip_page_table: bool = False,
):
"""Launch the fused metadata update (one kernel for the whole replay init).
Contract: only the live prefix (cdiv(cache_seqlens, page_size) pages) of each
page_table / swa_page_table row is (re)written; the tail keeps stale values
across replays, so consumers must bound reads by cache_seqlens.
``skip_page_table=True`` (the unified memory pool): the seqlen metadata is
still rebuilt in one launch, but the page-table writes are compiled out --
the bound tables are capture-stable read tables the translator
refreshes out-of-graph. ``page_table`` may then be None.
"""
if bs == 0:
return
@@ -160,6 +172,10 @@ def update_trtllm_mha_graph_metadata(
# set small enough to stay off the register-pressure / occupancy cliff while
# being wide enough to cover the static page-table width in few iterations.
PAGE_BLOCK = 512
if skip_page_table:
# Dead pointers under SKIP_PAGE_TABLE=1; pass a valid dummy for codegen.
page_table = cache_seqlens
swa_page_table = None
has_swa = swa_page_table is not None
has_swa_out = swa_out_cache_loc is not None
@@ -203,4 +219,5 @@ def update_trtllm_mha_graph_metadata(
Q_MODE=q_mode,
PAGE_BLOCK=PAGE_BLOCK,
BS_BLOCK=triton.next_power_of_2(bs),
SKIP_PAGE_TABLE=1 if skip_page_table else 0,
)
python/sglang/srt/arg_groups/kv_cache_hook.py
+23 -16
View File
@@ -318,18 +318,14 @@ def handle_page_major_kv_layout(server_args: Any):
"pool's per-layer views require a uniform row width; run "
"this model without --enable-unified-memory."
)
# Only the Triton attention kernels read the strided 4-D envelope K/V
# views; FA3 / FlashInfer do not. EXCEPTION: the unified-memory MLA pool
# exposes each layer as a contiguous per-layer view
# (build_mla_views), which the paged MLA kernels consume directly,
# with their kv_indices / block tables remapped to kernel-facing ids. Names below
# are the RESOLVED ids from attention_backends_of: "flashinfer" is
# FlashInferMLAAttnBackend for an MLA model, "trtllm_mla" the trtllm
# decode kernel; "cutedsl_mla" and "tokenspeed_mla" subclass
# TRTLLMMLABackend and inherit its read/write path; "fa3" remaps its
# page_table (in-kernel for captured decode, one funnel for eager).
# flashmla / cutlass_mla share the create_flashmla block-table path and
# can be added the same way once exercised.
# Allow-list. Every backend below reads through the translator, so what
# gates one is only whether its kernels can address the per-layer views:
# * MLA models: the full paged MLA family, incl. flashmla (ps=64
# snap). cutlass_mla stays rejected (never exercised).
# * MHA/SWA models: fa3 / fa4 / flashinfer / trtllm_mha alongside
# Triton. fa4 is the fa3 class.
# * Without the unified pool, plain page-major stays Triton-only.
# Names are the RESOLVED ids from attention_backends_of.
if cfg.enable_unified_memory and use_mla_backend(server_args):
allowed_full = {
"triton",
@@ -338,16 +334,27 @@ def handle_page_major_kv_layout(server_args: Any):
"flashinfer",
"cutedsl_mla",
"tokenspeed_mla",
"flashmla",
}
elif cfg.enable_unified_memory:
allowed_full = {
"triton",
"fa3",
"fa4",
"flashinfer",
"trtllm_mha",
}
else:
allowed_full = {"triton"}
backends = set(attention_backends_of(resolved_view(server_args)))
backends.discard(None)
assert backends <= allowed_full, (
"--enable-page-major-kv-layout requires the Triton attention backend "
"for the full-attention layers (unified-memory MLA also allows the "
f"paged MLA backends); got {sorted(backends)}, allowed "
f"{sorted(allowed_full)}. Pass a compatible --attention-backend."
"--enable-page-major-kv-layout: the resolved attention backends "
f"{sorted(backends)} are not in the allowed set "
f"{sorted(allowed_full)} for this configuration (unified memory "
"allows the per-layer-view families; plain page-major keeps the "
"envelope-strided views only Triton reads). Pass a compatible "
"--attention-backend."
)
# The Mamba/KDA state is stored in envelope-strided views; only
# stride-audited kernels may read it (Stage 4 audit, per slot):
python/sglang/srt/layers/attention/base_attn_backend.py
+6
View File
@@ -79,6 +79,12 @@ class AttentionBackend(ABC):
# (metadata glue graph) can read it off any backend without hasattr.
forward_metadata: Optional[object] = None
# The runner's KVIndexTranslator; backends that read through it set the
# instance attribute in __init__. None means "no translate" -- a backend
# that never set it cannot serve the unified pool, which the server-args
# allow-list enforces.
kv_index_translator = None
def init_forward_metadata(self, forward_batch: ForwardBatch):
"""Eager entry point. Default = ``_out_graph(fb) + _in_graph(fb)``.
python/sglang/srt/layers/attention/cutedsl_mla_backend.py
+1 -1
View File
@@ -364,7 +364,7 @@ class CuteDslMLABackend(TRTLLMMLABackend):
if (
save_kv_cache
and self._fused_set_kv_concat_q_fp8
and not self._unified_mla
and not self.kv_index_translator.is_translating
):
# Static pool: out_cache_loc is already the physical loc.
# Fused: bf16->fp8 quantize + KV scatter + q concat in one
python/sglang/srt/layers/attention/flashattention_backend.py
+71 -46
View File
@@ -18,7 +18,6 @@ from sglang.kernels.ops.kvcache.trtllm_mha_page_table import (
)
from sglang.srt.configs.model_config import AttentionArch
from sglang.srt.layers.attention.base_attn_backend import AttentionBackend
from sglang.srt.layers.attention.unified_mem_hooks import unified_mla_hooks
from sglang.srt.layers.attention.verify_mask import VerifyMask, maybe_create_verify_mask
from sglang.srt.layers.cp.base import CPAttentionBackendKind, get_cp_strategy
from sglang.srt.layers.cp.utils import is_cp_v2_active
@@ -192,11 +191,8 @@ class FlashAttentionBackend(AttentionBackend):
# seq_lens_cpu / seq_lens_sum D2H sync is ever needed.
self.needs_cpu_seq_lens = False
self.use_mla = model_runner.model_config.attention_arch == AttentionArch.MLA
# Unified pool: req_to_token holds VIRTUAL ids but the MLA per-layer views
# are kernel-facing, so every page_table needs remapping. MLA-only -- the MHA/SWA
# sub-pools keep the strided envelope layout FA3 cannot read at all.
self._unified_hooks = unified_mla_hooks(model_runner.token_to_kv_pool_allocator)
self._unified_dense = self._unified_hooks.enabled and self.use_mla
self.kv_index_translator = model_runner.kv_index_translator
self.kv_read_tables = None
self.skip_prefill = skip_prefill
self.attn_cp_size = model_runner.ps.attn_cp_size
self._verify_mask = None
@@ -229,6 +225,16 @@ class FlashAttentionBackend(AttentionBackend):
# Local attention settings
self.has_local_attention = model_runner.model_config.is_local_attention_model
# Local (chunked) attention derives its page table by re-translating
# metadata.page_table through the static full->swa map -- meaningless
# on the unified pool's kernel-facing tables, and no unified-eligible
# model uses it. Fail loud rather than silently double-translate.
assert not (
self.kv_index_translator.is_translating and self.has_local_attention
), (
"--enable-unified-memory does not support local-attention models "
"on the fa3/fa4 backend."
)
if self.has_local_attention:
assert (
model_runner.attention_chunk_size is not None
@@ -248,6 +254,12 @@ class FlashAttentionBackend(AttentionBackend):
"Prefill-aware SWA requires page_size=1, "
f"got page_size={self.page_size}"
)
# Its page-table builder indexes prefill_lens by POOL SLOT,
# incompatible with the batch-row canonical source.
assert not self.kv_index_translator.is_translating, (
"--enable-unified-memory does not support the prefill-aware "
"SWA decode mode; disable it for this model."
)
# Indexed by raw req_pool_idx values (see the write below and
# _build_pa_page_table), which range over [0, size] (row 0 is the
# reserved padding slot) -- so this needs size+1, not size.
@@ -527,6 +539,7 @@ class FlashAttentionBackend(AttentionBackend):
spec_info=spec_info,
seq_lens_cpu=seq_lens_cpu,
out_cache_loc=out_cache_loc,
in_capture=True,
)
if forward_mode.is_decode_or_idle() and spec_info is None:
@@ -1081,7 +1094,25 @@ class FlashAttentionBackend(AttentionBackend):
text_row, text_col
]
# Safe to rebind: every eager branch above produced a fresh tensor.
_unified_read = (
self.kv_index_translator.is_translating and metadata.page_table is not None
)
if _unified_read:
kv_view = self.kv_index_translator.index_table_for_batch(forward_batch)
metadata.page_table = kv_view.ids
if self.use_sliding_window_kv_pool:
metadata.swa_page_table = kv_view.sliding_window_ids
if forward_batch.out_cache_loc is not None:
# The swa write loc was computed from the still-VIRTUAL
# loc at ForwardBatch construction; re-running the
# full->swa map on the kernel-facing loc would be garbage.
metadata.swa_out_cache_loc = (
self.kv_index_translator.sliding_window_write_loc_for(
forward_batch.out_cache_loc
)
)
elif self.use_sliding_window_kv_pool:
# FA3 requires an int32 page_table.
metadata.swa_page_table = (
self.token_to_kv_pool.translate_loc_from_full_to_swa(
@@ -1095,28 +1126,8 @@ class FlashAttentionBackend(AttentionBackend):
)
)
# Unified pool: one remap for every eager branch above, which all filled
# page_table with VIRTUAL token ids. Rebinding is safe here because those
# branches each produced a fresh tensor; the captured path instead folds
# the remap into normal_decode_set_metadata, which must write in place.
#
# Placed BEFORE the `// page_size` reduction, in token space: since
# kernel_id(t) = phys_page * (ps * L) + t % ps, dense(page_start) // ps is
# phys_page * L, the dense page id the kernel wants. One site then serves
# both page sizes, and it inherits translate_kv_loc_for_kernel's tombstone
# clamp so an unwritten req_to_token slot lands in the page-0 sink.
if self._unified_dense and metadata.page_table is not None:
# Flattened: the page_size == 1 translate path uses index_select,
# which rejects a 2-D index.
pt = metadata.page_table
metadata.page_table = (
self._unified_hooks.translate_kv_loc_for_kernel(pt.reshape(-1))
.to(torch.int32)
.view(pt.shape)
)
# Convert the page table to a strided format which is needed by FA3 API
if self.page_size > 1:
if self.page_size > 1 and not _unified_read:
self.strided_indices = torch.arange(
0, metadata.page_table.shape[1], self.page_size, device=self.device
)
@@ -2157,6 +2168,12 @@ class FlashAttentionBackend(AttentionBackend):
"""
max_num_pages = (self.max_context_len + self.page_size - 1) // self.page_size
if self.kv_index_translator.is_translating:
# Zero-filled: slot 0 is the reserved sink in every id space.
self.kv_read_tables = self.kv_index_translator.make_capture_tables(
max_bs=max_bs, max_context_len=self.max_context_len
)
# This is being used by normal decode and draft decode when topk == 1
self.decode_cuda_graph_metadata = {
"cache_seqlens": torch.zeros(max_bs, dtype=torch.int32, device=self.device),
@@ -2709,6 +2726,7 @@ class FlashAttentionBackend(AttentionBackend):
spec_info: Optional[SpecInput],
seq_lens_cpu: Optional[torch.Tensor],
out_cache_loc: Optional[torch.Tensor] = None,
in_capture: bool = False,
):
"""Shared capture+replay body for the cuda-graph init path.
@@ -2731,8 +2749,13 @@ class FlashAttentionBackend(AttentionBackend):
if self.use_sliding_window_kv_pool and out_cache_loc is not None:
n = out_cache_loc.shape[0]
self.cuda_graph_swa_out_cache_loc[n:].zero_()
if in_capture and self.kv_index_translator.is_translating:
# A capture batch never went through `init_new`, so there is no
# rebound write loc; zeros are the page-0 sink.
self.cuda_graph_swa_out_cache_loc[:n].zero_()
else:
self.cuda_graph_swa_out_cache_loc[:n].copy_(
self.token_to_kv_pool.translate_loc_from_full_to_swa(out_cache_loc)
self.kv_index_translator.sliding_window_write_loc_for(out_cache_loc)
)
if forward_mode.is_decode_or_idle():
@@ -2744,13 +2767,19 @@ class FlashAttentionBackend(AttentionBackend):
# Page table built on-device (self-guards on cache_seqlens);
# max_seq_len_k left unset -- unread here (scheduler_metadata
# is normal-decode-only).
# Spec is asserted off under the unified pool, so this
# captured view is always the passthrough (req_to_token).
kv_view = self.kv_index_translator.build_index_table(
req_pool_indices=req_pool_indices,
seq_lens=seq_lens,
into=self.kv_read_tables,
)
normal_decode_set_metadata(
metadata.cache_seqlens_int32,
metadata.cu_seqlens_k,
metadata.page_table,
self.req_to_token,
req_pool_indices,
self.decode_cuda_graph_metadata["strided_indices"],
kv_view.ids,
kv_view.row_ids,
self.max_num_pages,
seq_lens,
self.speculative_step_id + 1,
@@ -2761,12 +2790,8 @@ class FlashAttentionBackend(AttentionBackend):
if self.use_sliding_window_kv_pool
else None
),
v2p_page_table=(
self._unified_hooks.v2p_page_table
if self._unified_dense
else None
),
kernel_page_multiplier=self._unified_hooks.kernel_page_multiplier,
src_is_read_table=kv_view.is_translated,
swa_src_table=kv_view.sliding_window_ids,
)
else:
@@ -2866,13 +2891,17 @@ class FlashAttentionBackend(AttentionBackend):
if seq_lens_cpu is not None
else self.max_context_len
)
kv_view = self.kv_index_translator.build_index_table(
req_pool_indices=req_pool_indices,
seq_lens=seq_lens,
into=self.kv_read_tables,
)
normal_decode_set_metadata(
metadata.cache_seqlens_int32,
metadata.cu_seqlens_k,
metadata.page_table,
self.req_to_token,
req_pool_indices,
self.decode_cuda_graph_metadata["strided_indices"],
kv_view.ids,
kv_view.row_ids,
self.max_num_pages,
seq_lens,
0,
@@ -2883,12 +2912,8 @@ class FlashAttentionBackend(AttentionBackend):
if self.use_sliding_window_kv_pool
else None
),
v2p_page_table=(
self._unified_hooks.v2p_page_table
if self._unified_dense
else None
),
kernel_page_multiplier=self._unified_hooks.kernel_page_multiplier,
src_is_read_table=kv_view.is_translated,
swa_src_table=kv_view.sliding_window_ids,
)
self._maybe_update_local_attn_metadata_for_replay(
python/sglang/srt/layers/attention/flashinfer_backend.py
+89 -24
View File
@@ -36,6 +36,7 @@ from sglang.srt.layers.quantization.fp4_kv_cache_quant_method import (
)
from sglang.srt.layers.radix_attention import AttentionType
from sglang.srt.mem_cache.base_swa_memory_pool import BaseSWAKVPool
from sglang.srt.mem_cache.kv_index_translator import KVIndexTable
from sglang.srt.mem_cache.memory_pool import KVWriteLoc
from sglang.srt.model_executor.cuda_graph_config import (
Backend,
@@ -310,6 +311,8 @@ class FlashInferAttnBackend(AttentionBackend):
self.req_to_token_pool = model_runner.req_to_token_pool
self.token_to_kv_pool = model_runner.token_to_kv_pool
self.kv_index_translator = model_runner.kv_index_translator
self.kv_read_tables = None
self._swa_kv_pool: Optional[BaseSWAKVPool] = self._resolve_swa_kv_pool(
model_runner
)
@@ -721,9 +724,16 @@ class FlashInferAttnBackend(AttentionBackend):
num_tokens = forward_batch.positions.numel()
self._prepare_cuda_graph_metadata(bs, num_tokens, forward_mode, spec_info)
# All flashinfer gathers run OUT-of-graph (plan time), so the
# capture-stable read table is buffer reuse, not pointer stability.
kv_view = self.kv_index_translator.build_index_table(
req_pool_indices=req_pool_indices[:bs],
seq_lens=seq_lens[:bs],
into=self.kv_read_tables,
)
if forward_mode.is_decode_or_idle():
self.indices_updater_decode.update(
req_pool_indices[:bs],
seq_lens[:bs],
seq_lens_cpu[:bs] if seq_lens_cpu is not None else None,
seq_lens_sum,
@@ -732,6 +742,7 @@ class FlashInferAttnBackend(AttentionBackend):
spec_info=spec_info,
fixed_split_size=None,
disable_split_kv=self.disable_cuda_graph_kv_split,
kv_view=kv_view,
)
elif forward_mode.is_target_verify():
self.indices_updater_prefill.update(
@@ -744,6 +755,7 @@ class FlashInferAttnBackend(AttentionBackend):
use_ragged=False,
encoder_lens=encoder_lens[:bs] if encoder_lens is not None else None,
spec_info=spec_info,
kv_view=kv_view,
)
elif forward_mode.is_dllm_extend():
self.indices_updater_prefill.update(
@@ -756,6 +768,7 @@ class FlashInferAttnBackend(AttentionBackend):
use_ragged=not self.use_paged,
encoder_lens=encoder_lens[:bs] if encoder_lens is not None else None,
spec_info=None,
kv_view=kv_view,
)
elif forward_mode.is_draft_extend_v2():
self.indices_updater_prefill.update(
@@ -768,6 +781,7 @@ class FlashInferAttnBackend(AttentionBackend):
use_ragged=False,
encoder_lens=encoder_lens[:bs] if encoder_lens is not None else None,
spec_info=spec_info,
kv_view=kv_view,
)
elif forward_mode.is_extend():
# Plain EXTEND under full prefill CUDA graph. plan() runs
@@ -786,6 +800,7 @@ class FlashInferAttnBackend(AttentionBackend):
use_ragged=False,
encoder_lens=encoder_lens[:bs] if encoder_lens is not None else None,
spec_info=None,
kv_view=kv_view,
)
else:
raise ValueError("Invalid forward mode")
@@ -840,11 +855,16 @@ class FlashInferAttnBackend(AttentionBackend):
# Refill the SWA write-target buffer from the live out_cache_loc before
# replay (bound onto the metadata at capture below).
if self.use_sliding_window_kv_pool and forward_batch.out_cache_loc is not None:
assert self._swa_kv_pool is not None
n = forward_batch.out_cache_loc.shape[0]
self.cuda_graph_swa_out_cache_loc[n:].zero_()
if in_capture and self.kv_index_translator.is_translating:
# A runner-built capture batch never went through `init_new`,
# so there is no prepared write loc to resolve -- and zeros are the
# page-0 sink in every id space. Replay refills below.
self.cuda_graph_swa_out_cache_loc[:n].zero_()
else:
self.cuda_graph_swa_out_cache_loc[:n].copy_(
self._swa_kv_pool.translate_loc_from_full_to_swa(
self.kv_index_translator.sliding_window_write_loc_for(
forward_batch.out_cache_loc
)
)
@@ -934,16 +954,15 @@ class FlashInferAttnBackend(AttentionBackend):
return layer.k_scale, layer.v_scale
def init_forward_metadata(self, forward_batch: ForwardBatch):
kv_view = self.kv_index_translator.index_table_for_batch(forward_batch)
swa_out_cache_loc = None
if self.use_sliding_window_kv_pool and forward_batch.out_cache_loc is not None:
assert self._swa_kv_pool is not None
swa_out_cache_loc = self._swa_kv_pool.translate_loc_from_full_to_swa(
swa_out_cache_loc = self.kv_index_translator.sliding_window_write_loc_for(
forward_batch.out_cache_loc
)
if forward_batch.forward_mode.is_decode_or_idle():
self.indices_updater_decode.update(
forward_batch.req_pool_indices,
forward_batch.seq_lens,
forward_batch.seq_lens_cpu,
forward_batch.seq_lens_sum,
@@ -952,6 +971,7 @@ class FlashInferAttnBackend(AttentionBackend):
spec_info=forward_batch.spec_info,
fixed_split_size=self.decode_split_tile_size,
disable_split_kv=False,
kv_view=kv_view,
)
self.forward_metadata = DecodeMetadata(
self.decode_wrappers, swa_out_cache_loc=swa_out_cache_loc
@@ -967,6 +987,7 @@ class FlashInferAttnBackend(AttentionBackend):
use_ragged=False,
encoder_lens=forward_batch.encoder_lens,
spec_info=forward_batch.spec_info,
kv_view=kv_view,
)
self.forward_metadata = PrefillMetadata(
self.prefill_wrappers_verify,
@@ -1020,6 +1041,7 @@ class FlashInferAttnBackend(AttentionBackend):
cross_attention_custom_mask=forward_batch.cross_attention_custom_mask,
extend_prefix_lens_cpu=forward_batch.extend_prefix_lens_cpu,
custom_kv_indices=self.dq_page_table,
kv_view=kv_view,
)
self.forward_metadata = PrefillMetadata(
self.prefill_wrappers_paged,
@@ -1035,6 +1057,9 @@ class FlashInferAttnBackend(AttentionBackend):
max_num_tokens: int,
kv_indices_buf: Optional[torch.Tensor] = None,
):
self.kv_read_tables = self.kv_index_translator.make_capture_tables(
max_bs=max_bs, max_context_len=self.max_context_len
)
if kv_indices_buf is None:
cuda_graph_kv_indices = torch.zeros(
(max_num_tokens * self.max_context_len,),
@@ -1539,7 +1564,6 @@ class FlashInferIndicesUpdaterDecode:
# Buffers and wrappers
self.kv_indptr = attn_backend.kv_indptr
self.kv_last_page_len = attn_backend.kv_last_page_len
self.req_to_token = model_runner.req_to_token_pool.req_to_token
self._swa_kv_pool = attn_backend._swa_kv_pool
# Dispatch the update function
@@ -1553,7 +1577,6 @@ class FlashInferIndicesUpdaterDecode:
def update(
self,
req_pool_indices: torch.Tensor,
seq_lens: torch.Tensor,
seq_lens_cpu: Optional[torch.Tensor],
seq_lens_sum: int,
@@ -1562,13 +1585,14 @@ class FlashInferIndicesUpdaterDecode:
spec_info: Optional[SpecInput],
fixed_split_size: Optional[int] = None,
disable_split_kv: Optional[bool] = None,
*,
kv_view: KVIndexTable,
):
# Keep the signature for type checking. It will be assigned during runtime.
raise NotImplementedError()
def update_single_wrapper(
self,
req_pool_indices: torch.Tensor,
seq_lens: torch.Tensor,
seq_lens_cpu: Optional[torch.Tensor],
seq_lens_sum: int,
@@ -1577,11 +1601,12 @@ class FlashInferIndicesUpdaterDecode:
spec_info: Optional[SpecInput],
fixed_split_size: Optional[int] = None,
disable_split_kv: Optional[bool] = None,
*,
kv_view: KVIndexTable,
):
decode_wrappers = decode_wrappers or self.decode_wrappers
self.call_begin_forward(
decode_wrappers[0],
req_pool_indices,
seq_lens,
seq_lens_sum,
self.kv_indptr[0],
@@ -1590,11 +1615,11 @@ class FlashInferIndicesUpdaterDecode:
seq_lens_cpu,
fixed_split_size=fixed_split_size,
disable_split_kv=disable_split_kv,
kv_view=kv_view,
)
def update_sliding_window(
self,
req_pool_indices: torch.Tensor,
seq_lens: torch.Tensor,
seq_lens_cpu: Optional[torch.Tensor],
seq_lens_sum: int,
@@ -1603,6 +1628,8 @@ class FlashInferIndicesUpdaterDecode:
spec_info: Optional[SpecInput],
fixed_split_size: Optional[int] = None,
disable_split_kv: Optional[bool] = None,
*,
kv_view: KVIndexTable,
):
assert self.sliding_window_size is not None
for wrapper_id in range(2):
@@ -1632,7 +1659,6 @@ class FlashInferIndicesUpdaterDecode:
self.call_begin_forward(
decode_wrappers[wrapper_id],
req_pool_indices,
paged_kernel_lens_tmp,
paged_kernel_lens_sum_tmp,
self.kv_indptr[wrapper_id],
@@ -1642,11 +1668,11 @@ class FlashInferIndicesUpdaterDecode:
use_sliding_window_kv_pool=use_sliding_window_kv_pool,
fixed_split_size=fixed_split_size,
disable_split_kv=disable_split_kv,
kv_view=kv_view,
)
def update_cross_attention(
self,
req_pool_indices: torch.Tensor,
seq_lens: torch.Tensor,
seq_lens_cpu: Optional[torch.Tensor],
seq_lens_sum: int,
@@ -1655,6 +1681,8 @@ class FlashInferIndicesUpdaterDecode:
spec_info: Optional[SpecInput],
fixed_split_size: Optional[int] = None,
disable_split_kv: Optional[bool] = None,
*,
kv_view: KVIndexTable,
):
# Cache encoder_lens on CPU to avoid GPU→CPU transfer per call
encoder_lens_cpu = encoder_lens.cpu() if encoder_lens is not None else None
@@ -1672,7 +1700,6 @@ class FlashInferIndicesUpdaterDecode:
self.call_begin_forward(
decode_wrappers[wrapper_id],
req_pool_indices,
paged_kernel_lens,
seq_lens_sum,
self.kv_indptr[wrapper_id],
@@ -1681,12 +1708,12 @@ class FlashInferIndicesUpdaterDecode:
seq_lens_cpu=kv_lens_cpu,
fixed_split_size=fixed_split_size,
disable_split_kv=disable_split_kv,
kv_view=kv_view,
)
def call_begin_forward(
self,
wrapper: BatchDecodeWithPagedKVCacheWrapper,
req_pool_indices: torch.Tensor,
paged_kernel_lens: torch.Tensor,
paged_kernel_lens_sum: int,
kv_indptr: torch.Tensor,
@@ -1696,9 +1723,15 @@ class FlashInferIndicesUpdaterDecode:
use_sliding_window_kv_pool: bool = False,
fixed_split_size: Optional[int] = None,
disable_split_kv: Optional[bool] = None,
*,
kv_view: KVIndexTable,
):
# Unified SWA wrapper-0: gather from the swa canonical directly -- its
# entries are already swa-side kernel-facing ids, so the in-place
# full->swa translate below must not run on top of them.
use_swa_source = use_sliding_window_kv_pool and kv_view.is_translated
if spec_info is None or getattr(spec_info, "kv_indptr", None) is None:
bs = len(req_pool_indices)
bs = len(paged_kernel_lens)
kv_indptr[1 : bs + 1] = torch.cumsum(paged_kernel_lens, dim=0)
kv_indptr = kv_indptr[: bs + 1]
@@ -1710,20 +1743,26 @@ class FlashInferIndicesUpdaterDecode:
paged_kernel_lens_sum, dtype=torch.int32, device="cuda"
)
if use_swa_source:
assert kv_view.sliding_window_ids is not None
src_table = kv_view.sliding_window_ids
else:
src_table = kv_view.ids
create_flashinfer_kv_indices_triton[(bs,)](
self.req_to_token,
req_pool_indices,
src_table,
kv_view.row_ids,
paged_kernel_lens,
kv_indptr,
kv_start_idx,
kv_indices,
self.req_to_token.shape[1],
kv_view.row_stride,
ENTRY_PAGE_SIZE=kv_view.entry_page_size,
)
else:
kv_indptr, kv_indices = spec_info.kv_indptr, spec_info.kv_indices
bs = kv_indptr.shape[0] - 1
if use_sliding_window_kv_pool:
if use_sliding_window_kv_pool and not use_swa_source:
assert self._swa_kv_pool is not None
kv_last_index = kv_indptr[-1]
kv_indices[:kv_last_index] = (
@@ -1811,6 +1850,9 @@ class FlashInferIndicesUpdaterPrefill:
self.kv_indptr = attn_backend.kv_indptr
self.kv_last_page_len = attn_backend.kv_last_page_len
self.qo_indptr = attn_backend.qo_indptr
# Kept ONLY for the spec-info branches (generate_attn_arg_prefill),
# which are static-pool-only: unified memory asserts spec off. The
# normal builders source from the per-batch KVIndexTable.
self.req_to_token = model_runner.req_to_token_pool.req_to_token
self._swa_kv_pool = attn_backend._swa_kv_pool
self.prefill_wrapper_ragged = attn_backend.prefill_wrapper_ragged
@@ -1840,6 +1882,8 @@ class FlashInferIndicesUpdaterPrefill:
cross_attention_custom_mask: Optional[torch.Tensor] = None,
extend_prefix_lens_cpu: Optional[List[int]] = None,
custom_kv_indices: Optional[torch.Tensor] = None,
*,
kv_view: KVIndexTable,
):
# Keep the signature for type checking. It will be assigned during runtime.
raise NotImplementedError()
@@ -1860,6 +1904,8 @@ class FlashInferIndicesUpdaterPrefill:
cross_attention_custom_mask: Optional[torch.Tensor] = None,
extend_prefix_lens_cpu: Optional[List[int]] = None,
custom_kv_indices: Optional[torch.Tensor] = None,
*,
kv_view: KVIndexTable,
):
if use_ragged:
assert prefix_lens is not None
@@ -1890,6 +1936,7 @@ class FlashInferIndicesUpdaterPrefill:
multi_item_params=multi_item_params,
seq_lens_cpu=seq_lens_cpu,
custom_kv_indices=custom_kv_indices,
kv_view=kv_view,
)
def update_sliding_window(
@@ -1908,6 +1955,8 @@ class FlashInferIndicesUpdaterPrefill:
cross_attention_custom_mask: Optional[torch.Tensor] = None,
extend_prefix_lens_cpu: Optional[List[int]] = None,
custom_kv_indices: Optional[torch.Tensor] = None,
*,
kv_view: KVIndexTable,
):
if custom_kv_indices is not None:
raise RuntimeError(
@@ -1983,6 +2032,7 @@ class FlashInferIndicesUpdaterPrefill:
if (wrapper_id == 0 and not use_ragged and spec_info is None)
else -1
),
kv_view=kv_view,
)
def _build_swa_prefix_custom_mask(
@@ -2042,6 +2092,8 @@ class FlashInferIndicesUpdaterPrefill:
cross_attention_custom_mask: Optional[torch.Tensor] = None,
extend_prefix_lens_cpu: Optional[List[int]] = None,
custom_kv_indices: Optional[torch.Tensor] = None,
*,
kv_view: KVIndexTable,
):
if custom_kv_indices is not None:
raise RuntimeError(
@@ -2077,6 +2129,7 @@ class FlashInferIndicesUpdaterPrefill:
cross_attention_custom_mask=(
cross_attention_custom_mask if wrapper_id == 1 else None
),
kv_view=kv_view,
)
def call_begin_forward(
@@ -2100,8 +2153,14 @@ class FlashInferIndicesUpdaterPrefill:
seq_lens_cpu: Optional[torch.Tensor] = None,
custom_kv_indices: Optional[torch.Tensor] = None,
window_left: int = -1,
*,
kv_view: KVIndexTable,
):
bs = len(seq_lens)
# Unified SWA wrapper-0: gather from the swa canonical directly -- its
# entries are already swa-side kernel-facing ids, so the in-place
# full->swa translate below must not run on top of them.
use_swa_source = use_sliding_window_kv_pool and kv_view.is_translated
if spec_info is None:
assert prefix_lens is not None
assert len(seq_lens) == len(req_pool_indices)
@@ -2127,14 +2186,20 @@ class FlashInferIndicesUpdaterPrefill:
dtype=torch.int32,
device=req_pool_indices.device,
)
if use_swa_source:
assert kv_view.sliding_window_ids is not None
src_table = kv_view.sliding_window_ids
else:
src_table = kv_view.ids
create_flashinfer_kv_indices_triton[(bs,)](
self.req_to_token,
req_pool_indices,
src_table,
kv_view.row_ids,
paged_kernel_lens,
kv_indptr,
kv_start_idx,
kv_indices,
self.req_to_token.shape[1],
kv_view.row_stride,
ENTRY_PAGE_SIZE=kv_view.entry_page_size,
)
qo_indptr[1 : bs + 1] = torch.cumsum(seq_lens - prefix_lens, dim=0)
qo_indptr = qo_indptr[: bs + 1]
@@ -2173,7 +2238,7 @@ class FlashInferIndicesUpdaterPrefill:
q_data_type=self.q_data_type,
)
if use_sliding_window_kv_pool:
if use_sliding_window_kv_pool and not use_swa_source:
assert self._swa_kv_pool is not None
kv_last_index = kv_indptr[-1]
kv_indices[:kv_last_index] = (
python/sglang/srt/layers/attention/flashinfer_mla_backend.py
+49 -46
View File
@@ -30,12 +30,12 @@ from sglang.srt.layers.attention.base_attn_backend import AttentionBackend
from sglang.srt.layers.attention.flashinfer_backend import (
create_flashinfer_kv_indices_triton,
)
from sglang.srt.layers.attention.unified_mem_hooks import unified_mla_hooks
from sglang.srt.layers.dcp import (
DecodeContextParallelMetadata,
update_local_kv_lens_for_dcp,
)
from sglang.srt.layers.dcp.planner import plan_dcp_decode_metadata
from sglang.srt.mem_cache.kv_index_translator import KVIndexTable
from sglang.srt.model_executor.forward_batch_info import ForwardBatch, ForwardMode
from sglang.srt.model_executor.runner_backend_utils.breakable_cuda_graph import (
is_in_breakable_cuda_graph,
@@ -238,6 +238,8 @@ class FlashInferMLAAttnBackend(AttentionBackend):
# corresponding ForwardBatch fields.
self.req_to_token_pool = model_runner.req_to_token_pool
self.token_to_kv_pool = model_runner.token_to_kv_pool
self.kv_index_translator = model_runner.kv_index_translator
self.kv_read_tables = None
self.enable_chunk_kv = (
not skip_prefill
and get_disagg().disaggregation_mode != "decode"
@@ -342,6 +344,14 @@ class FlashInferMLAAttnBackend(AttentionBackend):
forward_mode = forward_batch.forward_mode
spec_info = forward_batch.spec_info
# All flashinfer gathers run OUT-of-graph (plan time), so the
# capture-stable table is buffer reuse, not pointer stability.
kv_view = self.kv_index_translator.build_index_table(
req_pool_indices=req_pool_indices[:bs],
seq_lens=seq_lens[:bs],
into=self.kv_read_tables,
)
if in_capture:
num_tokens = forward_batch.positions.numel()
seq_lens_sum = seq_lens.sum().item()
@@ -358,12 +368,12 @@ class FlashInferMLAAttnBackend(AttentionBackend):
backend="auto",
)
self.indices_updater_decode.update(
req_pool_indices,
seq_lens,
seq_lens_sum,
decode_wrapper=decode_wrapper,
init_metadata_replay=False,
spec_info=spec_info,
kv_view=kv_view,
)
self.decode_cuda_graph_metadata[bs] = decode_wrapper
self.forward_metadata = DecodeMetadata(decode_wrapper)
@@ -396,6 +406,7 @@ class FlashInferMLAAttnBackend(AttentionBackend):
spec_info=spec_info,
seq_lens_cpu=seq_lens_cpu,
in_capture=True,
kv_view=kv_view,
)
if forward_mode.is_target_verify() and (
spec_info is None
@@ -412,16 +423,18 @@ class FlashInferMLAAttnBackend(AttentionBackend):
forward_mode=forward_mode,
spec_info=spec_info,
seq_lens_cpu=forward_batch.seq_lens_cpu,
kv_view=kv_view,
)
def init_forward_metadata(self, forward_batch: ForwardBatch):
kv_view = self.kv_index_translator.index_table_for_batch(forward_batch)
if forward_batch.forward_mode.is_decode_or_idle():
self.indices_updater_decode.update(
forward_batch.req_pool_indices,
forward_batch.seq_lens,
forward_batch.seq_lens_sum,
decode_wrapper=self.decode_wrapper,
init_metadata_replay=False,
kv_view=kv_view,
)
self.forward_metadata = DecodeMetadata(self.decode_wrapper)
elif forward_batch.forward_mode.is_target_verify():
@@ -433,6 +446,7 @@ class FlashInferMLAAttnBackend(AttentionBackend):
prefill_wrapper_paged=self.prefill_wrapper_verify,
use_ragged=False,
spec_info=forward_batch.spec_info,
kv_view=kv_view,
)
self.forward_metadata = PrefillMetadata(self.prefill_wrapper_verify, False)
else:
@@ -481,6 +495,7 @@ class FlashInferMLAAttnBackend(AttentionBackend):
qo_indptr_cpu=qo_indptr_cpu,
kv_indptr_cpu=kv_indptr_cpu,
kv_len_arr_cpu=kv_len_arr_cpu,
kv_view=kv_view,
)
self.forward_metadata = PrefillMetadata(
self.prefill_wrapper_paged, use_ragged
@@ -492,6 +507,9 @@ class FlashInferMLAAttnBackend(AttentionBackend):
max_num_tokens: int,
kv_indices_buf: Optional[torch.Tensor] = None,
):
self.kv_read_tables = self.kv_index_translator.make_capture_tables(
max_bs=max_bs, max_context_len=self.max_context_len
)
if kv_indices_buf is None:
cuda_graph_kv_indices = torch.zeros(
(max_bs * self.max_context_len,),
@@ -535,6 +553,7 @@ class FlashInferMLAAttnBackend(AttentionBackend):
forward_mode: ForwardMode,
spec_info: Optional[SpecInput],
seq_lens_cpu: Optional[torch.Tensor],
kv_view: KVIndexTable,
in_capture: bool = False,
):
"""Shared capture+replay body for the cuda-graph init path.
@@ -557,12 +576,12 @@ class FlashInferMLAAttnBackend(AttentionBackend):
)
self.indices_updater_decode.update(
req_pool_indices[:bs],
seq_lens[:bs],
seq_lens_sum,
decode_wrapper=self.decode_cuda_graph_metadata[bs],
init_metadata_replay=True,
spec_info=spec_info,
kv_view=kv_view,
**self.fast_decode_kwargs,
)
elif forward_mode.is_target_verify():
@@ -613,6 +632,7 @@ class FlashInferMLAAttnBackend(AttentionBackend):
if use_generic_fast_plan
else None
),
kv_view=kv_view,
)
else:
raise ValueError(f"Invalid forward mode: {forward_mode=}")
@@ -845,84 +865,70 @@ class FlashInferMLAIndicesUpdaterDecode:
# Buffers and wrappers
self.kv_indptr = attn_backend.kv_indptr
self.req_to_token = model_runner.req_to_token_pool.req_to_token
self.q_indptr = attn_backend.q_indptr_decode
# Unified dense MLA pool: VIRTUAL -> DENSE kv_indices (see prefill updater).
self._translate_kv_loc_dense = unified_mla_hooks(
model_runner.token_to_kv_pool_allocator
).translate_kv_loc_for_kernel
def update(
self,
req_pool_indices: torch.Tensor,
seq_lens: torch.Tensor,
seq_lens_sum: int,
decode_wrapper: BatchMLAPagedAttentionWrapper,
init_metadata_replay: bool = False,
spec_info: Optional[SpecInput] = None,
*,
kv_view: KVIndexTable,
**fast_decode_kwargs,
):
decode_wrapper = decode_wrapper or self.decode_wrapper
self.call_begin_forward(
decode_wrapper,
req_pool_indices,
seq_lens,
seq_lens_sum,
self.q_indptr,
self.kv_indptr,
init_metadata_replay,
spec_info,
kv_view=kv_view,
**fast_decode_kwargs,
)
def call_begin_forward(
self,
wrapper: BatchMLAPagedAttentionWrapper,
req_pool_indices: torch.Tensor,
paged_kernel_lens: torch.Tensor,
paged_kernel_lens_sum: int,
q_indptr: torch.Tensor,
kv_indptr: torch.Tensor,
init_metadata_replay: bool = False,
spec_info: Optional[SpecInput] = None,
*,
kv_view: KVIndexTable,
**fast_decode_kwargs,
):
bs = len(req_pool_indices)
bs = len(paged_kernel_lens)
q_indptr = q_indptr[: bs + 1]
kv_lens = paged_kernel_lens.to(torch.int32)
sm_scale = self.scaling
if spec_info is None:
kv_indptr[1 : bs + 1] = torch.cumsum(paged_kernel_lens, dim=0)
kv_indptr = kv_indptr[: bs + 1]
# On replay `kv_indices` IS the capture-stable buffer the captured
# wrapper reads -- never rebind it. The builder fills only the
# [:paged_kernel_lens_sum] prefix; the stale tail is unread.
kv_indices = (
torch.empty(paged_kernel_lens_sum, dtype=torch.int32, device="cuda")
if not init_metadata_replay
else fast_decode_kwargs["kv_indices"]
)
create_flashinfer_kv_indices_triton[(bs,)](
self.req_to_token,
req_pool_indices,
kv_view.ids,
kv_view.row_ids,
paged_kernel_lens,
kv_indptr,
None,
kv_indices,
self.req_to_token.shape[1],
kv_view.row_stride,
ENTRY_PAGE_SIZE=kv_view.entry_page_size,
)
# Unified pool: VIRTUAL -> DENSE, written back IN PLACE.
#
# On the cuda-graph replay path `kv_indices` IS the capture-stable
# buffer (fast_decode_kwargs["kv_indices"] == cuda_graph_kv_indices)
# that the captured wrapper reads, and `fast_mla_decode_plan` ignores
# the kv_indices argument entirely -- rebinding the local name to a
# fresh tensor would leave the graph reading VIRTUAL ids. Only the
# [:paged_kernel_lens_sum] prefix the index kernel just filled is
# translated; the stale tail is left alone so it can never index the
# v2p table out of bounds. The int64 translate result narrows back to
# the buffer's int32 on copy_ (flashinfer requires int32; kernel-facing ids
# fit comfortably).
if self._translate_kv_loc_dense is not None:
valid = kv_indices[:paged_kernel_lens_sum]
valid.copy_(self._translate_kv_loc_dense(valid))
if get_parallel().dcp_enabled:
plan_dcp_decode_metadata(
@@ -986,14 +992,11 @@ class FlashInferMLAIndicesUpdaterPrefill:
# Buffers and wrappers
self.kv_indptr = attn_backend.kv_indptr
self.qo_indptr = attn_backend.qo_indptr
# Kept ONLY for the spec-info branch (generate_attn_arg_prefill), which
# is static-pool-only: unified memory asserts spec off. The normal
# builder sources from the per-batch KVIndexTable.
self.req_to_token = model_runner.req_to_token_pool.req_to_token
self.prefill_wrapper_ragged = attn_backend.prefill_wrapper_ragged
# Unified dense MLA pool: kv_indices built from req_to_token are VIRTUAL;
# the paged wrapper reads the per-layer view, so remap them to kernel-facing
# token ids. None (identity) unless the unified MLA pool is active.
self._translate_kv_loc_dense = unified_mla_hooks(
model_runner.token_to_kv_pool_allocator
).translate_kv_loc_for_kernel
def update(
self,
@@ -1006,6 +1009,8 @@ class FlashInferMLAIndicesUpdaterPrefill:
spec_info: Optional[SpecInput] = None,
attn_dcp_metadata: Optional[DecodeContextParallelMetadata] = None,
fast_verify_plan_kwargs: Optional[dict] = None,
*,
kv_view: KVIndexTable,
qo_indptr_cpu: Optional[torch.Tensor] = None,
kv_indptr_cpu: Optional[torch.Tensor] = None,
kv_len_arr_cpu: Optional[torch.Tensor] = None,
@@ -1034,6 +1039,7 @@ class FlashInferMLAIndicesUpdaterPrefill:
qo_indptr_cpu=qo_indptr_cpu,
kv_indptr_cpu=kv_indptr_cpu,
kv_len_arr_cpu=kv_len_arr_cpu,
kv_view=kv_view,
)
def call_begin_forward(
@@ -1051,6 +1057,8 @@ class FlashInferMLAIndicesUpdaterPrefill:
spec_info: Optional[SpecInput] = None,
attn_dcp_metadata: Optional[DecodeContextParallelMetadata] = None,
fast_verify_plan_kwargs: Optional[dict] = None,
*,
kv_view: KVIndexTable,
qo_indptr_cpu: Optional[torch.Tensor] = None,
kv_indptr_cpu: Optional[torch.Tensor] = None,
kv_len_arr_cpu: Optional[torch.Tensor] = None,
@@ -1068,20 +1076,15 @@ class FlashInferMLAIndicesUpdaterPrefill:
device=req_pool_indices.device,
)
create_flashinfer_kv_indices_triton[(bs,)](
self.req_to_token,
req_pool_indices,
kv_view.ids,
kv_view.row_ids,
paged_kernel_lens,
kv_indptr,
None,
kv_indices,
self.req_to_token.shape[1],
kv_view.row_stride,
ENTRY_PAGE_SIZE=kv_view.entry_page_size,
)
# Unified pool: VIRTUAL -> DENSE token ids for the paged wrapper.
# Prefill is not cuda-graph captured under unified memory, so an eager
# gather is safe. Dense ids fit int32 (max = full_slots*num_layers ~
# 1e7 << 2^31); the flashinfer wrapper requires int32.
if self._translate_kv_loc_dense is not None:
kv_indices = self._translate_kv_loc_dense(kv_indices).to(torch.int32)
qo_indptr[1 : bs + 1] = torch.cumsum(seq_lens - prefix_lens, dim=0)
qo_indptr = qo_indptr[: bs + 1]
custom_mask = None
python/sglang/srt/layers/attention/flashmla_backend.py
+16
View File
@@ -162,6 +162,14 @@ class FlashMLABackend(FlashInferMLAAttnBackend):
if forward_batch.forward_mode.is_decode_or_idle():
max_seqlen_pad = triton.cdiv(eager_max_k, PAGE_SIZE)
block_kv_indices = self._eager_block_kv_indices(bs, max_seqlen_pad)
if self.kv_index_translator.is_translating:
assert self.page_size == PAGE_SIZE
self.kv_index_translator.fill_read_table(
out=block_kv_indices,
req_pool_indices=forward_batch.req_pool_indices,
seq_lens=forward_batch.seq_lens,
)
else:
create_flashmla_kv_indices_triton[
(bs, get_num_kv_index_blocks_flashmla(max_seqlen_pad, PAGE_SIZE))
](
@@ -328,6 +336,14 @@ class FlashMLABackend(FlashInferMLAAttnBackend):
else:
max_seqlen_pad = self.cuda_graph_kv_indices.shape[1]
if self.kv_index_translator.is_translating:
assert self.page_size == PAGE_SIZE
self.kv_index_translator.fill_read_table(
out=self.cuda_graph_kv_indices,
req_pool_indices=req_pool_indices[:bs],
seq_lens=seq_lens,
)
else:
create_flashmla_kv_indices_triton[
(
bs,
python/sglang/srt/layers/attention/tbo_backend.py
+1
View File
@@ -23,6 +23,7 @@ class TboAttnBackend(AttentionBackend):
# reads through TboAttnBackend resolve to the underlying pool.
self.token_to_kv_pool = primary.token_to_kv_pool
self.req_to_token_pool = primary.req_to_token_pool
self.kv_index_translator = primary.kv_index_translator
self.extend_dummy_seqs_capped_by_req_pool = getattr(
primary, "extend_dummy_seqs_capped_by_req_pool", False
)
python/sglang/srt/layers/attention/trtllm_mha_backend.py
+55 -6
View File
@@ -198,8 +198,12 @@ class TRTLLMHAAttnBackend(FlashInferAttnBackend):
# separate index spaces; SWA layers need a translated page_table.
self._swa_kv_pool: Optional[SWAKVPool] = self._resolve_swa_kv_pool(model_runner)
# Raw full->swa index mapping tensor for the fused cuda-graph
# metadata kernel (gather + // page_size happen on device).
if self._swa_kv_pool is not None:
# metadata kernel (gather + // page_size happen on device). The unified
# pool has no token-level mapping, so this is a static-pool mechanism.
if (
self._swa_kv_pool is not None
and not self.kv_index_translator.is_translating
):
self._swa_full_to_swa_mapping = self._swa_kv_pool.full_to_swa_index_mapping
assert self._swa_full_to_swa_mapping is not None, (
"SWA pool must register full_to_swa_index_mapping before "
@@ -465,6 +469,9 @@ class TRTLLMHAAttnBackend(FlashInferAttnBackend):
kv_indices_buf: Optional[torch.Tensor] = None,
):
"""Initialize CUDA graph state for TRTLLM MHA."""
self.kv_read_tables = self.kv_index_translator.make_capture_tables(
max_bs=max_bs, max_context_len=self.max_context_len
)
max_num_pages = self.max_num_pages
self.decode_cuda_graph_metadata = {
"cache_seqlens": torch.zeros(max_bs, dtype=torch.int32, device=self.device),
@@ -766,25 +773,27 @@ class TRTLLMHAAttnBackend(FlashInferAttnBackend):
# bounds real KV reads by cache_seqlens, so this is a fixed loop
# bound only — never a host max / seq_lens_cpu D2H sync.
max_seq_pages = self.max_num_pages
unified = self.kv_index_translator.is_translating
update_trtllm_mha_graph_metadata(
req_pool_indices=req_pool_indices,
seq_lens=seq_lens,
req_to_token=self.req_to_token,
cache_seqlens=metadata.cache_seqlens_int32,
cu_seqlens_k=metadata.cu_seqlens_k,
page_table=metadata.page_table,
page_table=None if unified else metadata.page_table,
bs=bs,
seqlen_offset=seqlen_offset,
max_seq_pages=max_seq_pages,
page_size=self.page_size,
swa_mapping=self._swa_full_to_swa_mapping,
swa_page_table=metadata.swa_page_table,
swa_page_table=None if unified else metadata.swa_page_table,
out_cache_loc=out_cache_loc,
swa_out_cache_loc=metadata.swa_out_cache_loc,
swa_out_cache_loc=None if unified else metadata.swa_out_cache_loc,
cu_seqlens_q=cu_seqlens_q,
qlens=qlens,
q_stride=q_stride,
q_mode=q_mode,
skip_page_table=unified,
)
if self._needs_encoder_only_expand(forward_mode, metadata):
@@ -888,6 +897,39 @@ class TRTLLMHAAttnBackend(FlashInferAttnBackend):
f"Invalid forward mode: {forward_mode=} for CUDA Graph replay."
)
if self.kv_index_translator.is_translating:
# Unified pool: refresh the capture-stable read table (this runs
# out-of-graph on BOTH capture and every replay-prep; the recorded
# fused kernel skips its page-table writes so the graph reads the
# refreshed content through pointers baked at capture).
kv_view = self.kv_index_translator.build_index_table(
req_pool_indices=forward_batch.req_pool_indices[:bs],
seq_lens=forward_batch.seq_lens[:bs],
into=self.kv_read_tables,
)
metadata = self.forward_metadata
if in_capture:
# Bind ONCE: the attention kernels bake these pointers at capture.
metadata.page_table = kv_view.ids[:bs]
if kv_view.sliding_window_ids is not None:
metadata.swa_page_table = kv_view.sliding_window_ids[:bs]
# A capture batch carries no prepared write loc; zeros are the
# page-0 sink.
if (
self.use_sliding_window_kv_pool
and forward_batch.out_cache_loc is not None
):
n = forward_batch.out_cache_loc.shape[0]
self.cuda_graph_swa_out_cache_loc[n:].zero_()
if in_capture and self.kv_index_translator.is_translating:
self.cuda_graph_swa_out_cache_loc[:n].zero_()
else:
self.cuda_graph_swa_out_cache_loc[:n].copy_(
self.kv_index_translator.sliding_window_write_loc_for(
forward_batch.out_cache_loc
)
)
def _assert_ragged_verify_supported(self) -> None:
if self.is_xqa_impl:
raise NotImplementedError(
@@ -1035,6 +1077,13 @@ class TRTLLMHAAttnBackend(FlashInferAttnBackend):
else:
metadata.cu_seqlens_q = metadata.cu_seqlens_k
kv_view = self.kv_index_translator.index_table_for_batch(forward_batch)
if kv_view.is_translated:
# No fill kernel: the kernels take the tensor's own width/stride
# and bound their reads by cache_seqlens.
metadata.page_table = kv_view.ids
metadata.swa_page_table = kv_view.sliding_window_ids
else:
has_swa = self._swa_kv_pool is not None
metadata.page_table = torch.empty(
(batch_size, self.max_num_pages), dtype=torch.int32, device=device
@@ -1063,7 +1112,7 @@ class TRTLLMHAAttnBackend(FlashInferAttnBackend):
# int64 scatter index (unlike the int32 read page table above).
if self.use_sliding_window_kv_pool and forward_batch.out_cache_loc is not None:
metadata.swa_out_cache_loc = (
self.token_to_kv_pool.translate_loc_from_full_to_swa(
self.kv_index_translator.sliding_window_write_loc_for(
forward_batch.out_cache_loc
)
)
python/sglang/srt/layers/attention/trtllm_mla_backend.py
+40 -30
View File
@@ -49,7 +49,6 @@ from sglang.srt.layers.attention.flashinfer_mla_backend import (
FlashInferMLAAttnBackend,
FlashInferMLAMultiStepDraftBackend,
)
from sglang.srt.layers.attention.unified_mem_hooks import unified_mla_hooks
from sglang.srt.layers.attention.verify_mask import VerifyMask, maybe_create_verify_mask
from sglang.srt.model_executor.forward_batch_info import ForwardBatch, ForwardMode
from sglang.srt.model_executor.runner_backend_utils.breakable_cuda_graph import (
@@ -285,17 +284,8 @@ class TRTLLMMLABackend(FlashInferMLAAttnBackend):
# Tree-mask scratch is fetched from the target backend only.
self.is_draft_runner = model_runner.is_draft_worker
# Unified-memory per-layer-view hooks (None on the static pool). req_to_token
# holds VIRTUAL token ids; the block table needs kernel-facing page ids, so the
# kv-index kernels gather virtual->physical page through `_v2p_page_table`
# then scale by `_kernel_page_multiplier` (= num MLA layers). See
# build_mla_views / create_flashmla_kv_indices_triton.
_hooks = unified_mla_hooks(model_runner.token_to_kv_pool_allocator)
self._v2p_page_table = _hooks.v2p_page_table
self._kernel_page_multiplier = _hooks.kernel_page_multiplier
self._unified_mla = _hooks.enabled
# Per-forward kernel-facing write loc ([:n] view of a capture-stable buffer);
# None on the eager path, which passes out_cache_loc straight through.
# [:n] view of a capture-stable buffer on the cuda-graph path; None on
# the eager path, which passes forward_batch.out_cache_loc through.
self._decode_kernel_loc: Optional[torch.Tensor] = None
self.cuda_graph_out_cache_loc_kernel: Optional[torch.Tensor] = None
# Fused KV-scatter + q-concat on the decode dense-loc path (one launch
@@ -368,6 +358,13 @@ class TRTLLMMLABackend(FlashInferMLAAttnBackend):
(batch_size, max_blocks), -1, dtype=torch.int32, device=device
)
if self.kv_index_translator.is_translating:
self.kv_index_translator.fill_read_table(
out=block_kv_indices,
req_pool_indices=req_pool_indices,
seq_lens=seq_lens,
)
else:
create_flashmla_kv_indices_triton[
(
batch_size,
@@ -382,8 +379,6 @@ class TRTLLMMLABackend(FlashInferMLAAttnBackend):
self.req_to_token.stride(0),
max_blocks,
PAGED_SIZE=self.page_size,
v2p_ptr=self._v2p_page_table,
PAGE_MULT=self._kernel_page_multiplier,
)
return block_kv_indices
@@ -404,7 +399,7 @@ class TRTLLMMLABackend(FlashInferMLAAttnBackend):
# Unified pool: capture-stable buffer for the DENSE KV write loc, filled
# out-of-graph in init_forward_metadata_out_graph so the in-graph
# set_mla_kv_buffer captures no translate.
if self._unified_mla:
if self.kv_index_translator.is_translating:
self.cuda_graph_out_cache_loc_kernel = torch.zeros(
max_num_tokens, dtype=torch.int64, device=self.device
)
@@ -545,6 +540,13 @@ class TRTLLMMLABackend(FlashInferMLAAttnBackend):
metadata.seq_lens_k.copy_(seq_lens[:bs])
# Update block indices for new sequences.
if self.kv_index_translator.is_translating:
self.kv_index_translator.fill_read_table(
out=metadata.block_kv_indices,
req_pool_indices=req_pool_indices[:bs],
seq_lens=seq_lens,
)
else:
create_flashmla_kv_indices_triton[
(
bs,
@@ -561,8 +563,6 @@ class TRTLLMMLABackend(FlashInferMLAAttnBackend):
self.req_to_token.stride(0),
metadata.block_kv_indices.shape[1],
PAGED_SIZE=self.page_size,
v2p_ptr=self._v2p_page_table,
PAGE_MULT=self._kernel_page_multiplier,
)
def get_cuda_graph_seq_len_fill_value(self) -> int:
@@ -623,11 +623,9 @@ class TRTLLMMLABackend(FlashInferMLAAttnBackend):
forward_mode=forward_mode,
)
# Unified pool: precompute the DENSE KV write loc into the capture-stable
# buffer (both capture and each replay-prep run this out of the graph),
# so the in-graph set_mla_kv_buffer writes a dense loc without capturing
# a translate.
if self._unified_mla and (
# Out-of-graph on capture AND every replay-prep, so the in-graph
# set_mla_kv_buffer captures no translate.
if self.kv_index_translator.is_translating and (
forward_mode.is_decode_or_idle() or forward_mode.is_target_verify()
):
out_cache_loc = forward_batch.out_cache_loc
@@ -646,6 +644,24 @@ class TRTLLMMLABackend(FlashInferMLAAttnBackend):
else:
self._decode_kernel_loc = None
def _resolve_fused_write_loc(
self, forward_batch: ForwardBatch
) -> Optional[torch.Tensor]:
"""Write loc for the fused fp8 KV scatter, or None when this batch is
not covered by it.
Captured decode refills `_decode_kernel_loc` out of the graph, and the
captured kernel must read that buffer. Eager decode on a unified pool
has no such buffer, and the caller falls back to the unfused path.
"""
if self._decode_kernel_loc is not None:
return self._decode_kernel_loc
return (
None
if self.kv_index_translator.is_translating
else forward_batch.out_cache_loc
)
def init_forward_metadata(self, forward_batch: ForwardBatch):
"""Initialize the metadata for a forward pass."""
self._decode_kernel_loc = None
@@ -1050,13 +1066,7 @@ class TRTLLMMLABackend(FlashInferMLAAttnBackend):
assert q_rope is not None and k_rope is not None
if cos_sin_cache is None:
if save_kv_cache and self._fused_set_kv_concat_q_fp8:
loc = (
self._decode_kernel_loc
if self._decode_kernel_loc is not None
else (
None if self._unified_mla else forward_batch.out_cache_loc
)
)
loc = self._resolve_fused_write_loc(forward_batch)
if loc is not None:
# Fused: bf16->fp8 quantize + KV scatter + q concat
# in one launch; None when not covered.
@@ -1113,7 +1123,7 @@ class TRTLLMMLABackend(FlashInferMLAAttnBackend):
if (
merge_query
and self._fused_set_kv_concat_q
and not self._unified_mla
and not self.kv_index_translator.is_translating
):
# Static pool only, conservatively.
query = self._set_kv_and_concat_q_fused(
python/sglang/srt/layers/attention/unified_mem_hooks.py
-72
View File
@@ -1,72 +0,0 @@
# Copyright 2023-2026 SGLang Team
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
"""Allocator hooks the paged MLA attention backends need under the unified
memory pool.
Lives in its own module because three unrelated backend families consume it
(fa3, flashinfer_mla, and trtllm_mla with its cutedsl_mla / tokenspeed_mla
subclasses) and none of them should have to import another's module to get it.
"""
from __future__ import annotations
from typing import Callable, Optional
import msgspec
import torch
class UnifiedMLAHooks(msgspec.Struct, frozen=True):
"""Dense-view hooks for one KV allocator.
All-``None``/1/``False`` for the statically-partitioned pool, where
``req_to_token`` already holds physical ids and no translation is needed.
"""
# Page-level virtual->physical table, gathered through by block-table kernels.
v2p_page_table: Optional[torch.Tensor]
# Virtual token id -> DENSE kernel-facing id (tombstones clamped to the sink).
translate_kv_loc_for_kernel: Optional[Callable[..., torch.Tensor]]
# Dense page stride scale (= number of full-attention MLA layers).
kernel_page_multiplier: int
enabled: bool
_STATIC_POOL = UnifiedMLAHooks(
v2p_page_table=None,
translate_kv_loc_for_kernel=None,
kernel_page_multiplier=1,
enabled=False,
)
def unified_mla_hooks(allocator) -> UnifiedMLAHooks:
"""Probe ``allocator`` for the unified-pool per-layer-view hooks.
Detection keys on the v2p table, NOT on ``kernel_page_multiplier > 1``: a
rank owning exactly ONE full-attention layer has multiplier 1 while its
``req_to_token`` is still virtual. There the kernel-facing id collapses onto the
physical id, so the v2p gather alone is the whole translation.
"""
v2p = getattr(allocator, "full_v2p_page_table", None)
if v2p is None:
return _STATIC_POOL
return UnifiedMLAHooks(
v2p_page_table=v2p,
translate_kv_loc_for_kernel=getattr(
allocator, "translate_kv_loc_for_kernel", None
),
kernel_page_multiplier=getattr(allocator, "kernel_page_multiplier", 1),
enabled=True,
)
python/sglang/srt/mem_cache/memory_pool.py
+13 -13
View File
@@ -1602,13 +1602,15 @@ class KVWriteLoc:
KERNEL-FACING on every pool: physical by allocation on non-unified
pools, rebound at ForwardBatch construction (``rebind_write_loc``) on
the unified pool.
- ``swa_loc``: the pre-resolved SWA-sub-pool location for hybrid SWA pools
(``None`` otherwise).
- ``full_loc``: the full-attention-sub-pool location for the unified
memory pool (``None`` otherwise), carried in attention metadata
(``ForwardMetadata.out_cache_loc_full_physical``). Since the
construction-time rebind it is the SAME id space as ``loc``; the shared
full pool writes it directly and never translates.
- ``swa_loc``: the SWA-sub-pool location for hybrid SWA pools (``None``
otherwise); under the unified pool the translator derives it from the
same rebound loc (``sliding_window_write_loc_for``).
- ``full_loc``: OPTIONAL full-attention-sub-pool location. Since the
construction-time rebind it is the SAME id space as ``loc``, so pools
fall back to ``loc`` when it is ``None`` -- only triton's captured path
still passes its capture-stable
``ForwardMetadata.out_cache_loc_full_physical`` buffer here (a
same-space alias slated for collapse).
``swa_loc`` and ``full_loc`` are the parallel pair (each a pre-resolved
loc into its sub-pool, mirroring ``swa_kv_pool`` / ``full_kv_pool``);
@@ -3665,11 +3667,6 @@ class HybridLinearKVPool(KVCache):
# virtual->physical mamba-slot translate for the HiCache offload path;
# identity for a static pool, the allocator's `translate` for the unified pool.
self._mamba_translate = lambda ids: ids
# The MLA doors take DIFFERENT id spaces: `get_mla_kv_buffer` gets
# ForwardBatch-built read indices (prefix_chunk_kv_indices /
# fetch_mha_one_shot_kv_indices), still VIRTUAL, so it translates;
# `set_mla_kv_buffer` gets out_cache_loc, already kernel-facing.
self._full_translate = lambda ids: ids
self.use_mla = use_mla
if full_kv_pool is not None:
# Shared-KV-pool path: the caller built a UnifiedMHATokenToKVPool
@@ -3967,7 +3964,10 @@ class HybridLinearKVPool(KVCache):
dst_dtype: Optional[torch.dtype] = None,
):
assert self.use_mla, "get_mla_kv_buffer called when use_mla is False"
loc = self._full_translate(loc)
# Read door -- same kernel-facing contract as the write door: `loc` is
# a read-index tensor already translated at its production site
# (fetch_mha_one_shot_kv_indices / prepare_chunked_kv_indices); the
# pool never translates.
with self._transfer_id_context(layer):
return self.full_kv_pool.get_mla_kv_buffer(layer, loc, dst_dtype)
python/sglang/srt/mem_cache/unified_memory_pool.py
+9 -9
View File
@@ -1244,8 +1244,9 @@ def init_unified_mamba_pools(
# `_mamba_translate` feeds the HiCache offload path, GATED OFF here — wired but inert.
req_to_token_pool.mamba_allocator = mamba_slot_allocator
token_to_kv_pool._mamba_translate = mamba_slot_allocator.translate
if use_mla_backend:
token_to_kv_pool._full_translate = allocator.translate_kv_loc_for_kernel
# No full-KV translate hook is wired: both MLA doors now receive
# KERNEL-FACING ids -- writes from the ForwardBatch rebind, reads
# translated at their production sites.
logger.info(
"[unified-memory-pool] ============================================================"
@@ -1466,7 +1467,7 @@ class UnifiedSWAKVPool(SWAKVPool):
"""Route to the right sub-pool. Both `swa_loc` and `full_loc` are PHYSICAL
(pre-translated once per forward by the attention backend); never translates here.
"""
_, swa_loc, full_loc = unwrap_write_loc(loc_info)
loc, swa_loc, full_loc = unwrap_write_loc(loc_info)
layer_id = layer.layer_id
pool_layer_id, is_swa = self.layers_mapping[layer_id]
if is_swa:
@@ -1486,12 +1487,11 @@ class UnifiedSWAKVPool(SWAKVPool):
layer_id_override=pool_layer_id,
)
return
# Full layer: full_loc is full-physical, always precomputed (eager + cuda-graph).
assert full_loc is not None, (
"UnifiedSWAKVPool.set_kv_buffer: full layer received no full_loc; "
"ForwardMetadata.out_cache_loc_full_physical must be precomputed for "
"the unified memory pool."
)
# Full layer: `loc` is already the full-side kernel-facing id, so an
# explicit full_loc is a same-space alias -- only triton's captured path
# passes one (its capture-stable buffer).
if full_loc is None:
full_loc = loc
self.full_kv_pool.set_kv_buffer(
None,
full_loc,
python/sglang/srt/model_executor/forward_batch_deepseek_mha_mixin.py
+9
View File
@@ -12,6 +12,7 @@ from sglang.kernels.ops.kvcache.kv_indices import (
from sglang.srt.environ import envs
from sglang.srt.layers.dcp.layout import filter_dcp_local_chunk_kv_indices
from sglang.srt.model_executor.forward_context import (
get_attn_backend,
get_req_to_token_pool,
get_token_to_kv_pool,
)
@@ -90,6 +91,10 @@ class ForwardBatchDeepSeekMHAMixin:
self.prefix_chunk_starts_cpu[idx],
self.prefix_chunk_seq_lens_cpu[idx],
)
# None on a backend that never bound a translator.
src = get_attn_backend().kv_index_translator
if src is not None:
chunk_kv_indices = src.translate_full_attn_ids(chunk_kv_indices)
self.prefix_chunk_kv_indices.append(chunk_kv_indices)
# Here we suppose the length of each chunk is equal
@@ -235,5 +240,9 @@ class ForwardBatchDeepSeekMHAMixin:
kv_indices,
req_to_token.shape[1],
)
# None on a backend that never bound a translator.
src = get_attn_backend().kv_index_translator
if src is not None:
kv_indices = src.translate_full_attn_ids(kv_indices)
self.mha_one_shot_kv_indices = kv_indices
return kv_indices
python/sglang/srt/server_args.py
+1 -1
View File
@@ -3762,7 +3762,7 @@ class ServerArgs:
LANGUAGE_MODEL_ONLY_ARCHITECTURES = ("MuseGlimmerForConditionalGeneration",)
# The strided-layout Triton requirement is enforced via
# The attention-backend allow-list is enforced via
# --enable-page-major-kv-layout (implied by the unified pool in
# _handle_page_major_kv_layout); the model-family gate is enforced at pool
# construction in model_runner_kv_cache_mixin._init_pools.
test/registered/attention/test_normal_decode_set_metadata.py
+5 -6
View File
@@ -31,7 +31,6 @@ def reference_normal_decode_set_metadata(
page_table: torch.Tensor,
req_to_token: torch.Tensor,
req_pool_indices: torch.Tensor,
strided_indices: torch.Tensor,
max_seq_pages: int,
seq_lens: torch.Tensor,
seq_len_delta: int,
@@ -45,6 +44,11 @@ def reference_normal_decode_set_metadata(
"""
cache_seqlens_int32.copy_(seq_lens + seq_len_delta)
cu_seqlens_k[1:].copy_(torch.cumsum(cache_seqlens_int32, dim=0, dtype=torch.int32))
# Page-start columns, derived internally (the wrapper's dead
# strided_indices parameter was removed alongside its v2p args).
strided_indices = torch.arange(
0, req_to_token.shape[1], page_size, device=req_to_token.device
)
page_indices = req_to_token[
req_pool_indices[:, None],
strided_indices[:max_seq_pages][None, :],
@@ -213,7 +217,6 @@ class TestNormalDecodeSetMetadata(CustomTestCase):
ref_data["page_table"],
test_data["req_to_token"],
test_data["req_pool_indices"],
test_data["strided_indices"],
test_data["max_seq_pages"],
test_data["seq_lens"],
test_data["seq_len_delta"],
@@ -229,7 +232,6 @@ class TestNormalDecodeSetMetadata(CustomTestCase):
test_data["page_table"],
test_data["req_to_token"],
test_data["req_pool_indices"],
test_data["strided_indices"],
test_data["max_seq_pages"],
test_data["seq_lens"],
test_data["seq_len_delta"],
@@ -351,7 +353,6 @@ class TestNormalDecodeSetMetadata(CustomTestCase):
test_data["page_table"],
test_data["req_to_token"],
test_data["req_pool_indices"],
test_data["strided_indices"],
test_data["max_seq_pages"],
test_data["seq_lens"],
test_data["seq_len_delta"],
@@ -408,7 +409,6 @@ class TestNormalDecodeSetMetadata(CustomTestCase):
ref_data["page_table"],
test_data["req_to_token"],
test_data["req_pool_indices"],
test_data["strided_indices"],
test_data["max_seq_pages"],
test_data["seq_lens"],
0,
@@ -423,7 +423,6 @@ class TestNormalDecodeSetMetadata(CustomTestCase):
test_data["page_table"],
test_data["req_to_token"],
test_data["req_pool_indices"],
test_data["strided_indices"],
test_data["max_seq_pages"],
test_data["seq_lens"],
0,
test/registered/attention/test_trtllm_mha_graph_metadata.py
+72
View File
@@ -30,6 +30,8 @@ PAGE_SIZE = 128
def _make_backend_for_hook_test(speculative_num_draft_tokens=None):
from sglang.srt.mem_cache.kv_index_translator import KVIndexTranslator
backend = TRTLLMHAAttnBackend.__new__(TRTLLMHAAttnBackend)
backend.device = torch.device("cpu")
backend.max_context_len = 1024
@@ -45,6 +47,15 @@ def _make_backend_for_hook_test(speculative_num_draft_tokens=None):
backend.decode_cuda_graph_metadata = {}
backend.target_verify_metadata = {}
backend.draft_extend_metadata = {}
# Passthrough source (static pool): every unified-arm branch stays off,
# matching the real __init__'s parent binding.
backend.kv_index_translator = KVIndexTranslator(
req_to_token=backend.req_to_token,
token_to_kv_pool_allocator=SimpleNamespace(),
token_to_kv_pool=SimpleNamespace(),
page_size=PAGE_SIZE,
device="cpu",
)
backend.init_cuda_graph_state(max_bs=4, max_num_tokens=16)
return backend
@@ -543,6 +554,67 @@ def test_metadata_correctness(bs, seqlen_offset, q_mode, with_swa, static_width)
torch.testing.assert_close(swa_out_cache_loc, out_ref, rtol=0, atol=0)
@pytest.mark.parametrize("pass_tables", [False, True])
def test_skip_page_table_updates_seqlens_only(pass_tables):
"""The unified-memory arm: skip_page_table=True must still rebuild the
seqlen metadata in-graph but leave every page-table byte alone -- the bound
tables are capture-stable read tables the translator refreshes
out-of-graph, and an in-graph write would clobber them with virtual-derived
pages. Covers both call shapes: page_table=None (what the backend passes)
and a real sentinel-filled table (pins that the writes are compiled out,
not just unpassed)."""
if not torch.cuda.is_available():
pytest.skip("CUDA required")
bs, seqlen_offset, seed = 5, 1, 4242
pool_size, max_num_pages = 64, 16
seq_max = (max_num_pages - 2) * PAGE_SIZE
(
req_to_token,
req_pool_indices,
seq_lens,
_stride,
_cap,
) = _build_inputs(bs, pool_size, max_num_pages, None, seq_max, seed)
cache_seqlens = torch.zeros(bs, dtype=torch.int32, device=DEVICE)
cu_seqlens_k = torch.zeros(bs + 1, dtype=torch.int32, device=DEVICE)
sentinel_pt = None
sentinel_swa = None
if pass_tables:
sentinel_pt = torch.full(
(bs, max_num_pages), 777, dtype=torch.int32, device=DEVICE
)
sentinel_swa = torch.full(
(bs, max_num_pages), 888, dtype=torch.int32, device=DEVICE
)
update_trtllm_mha_graph_metadata(
req_pool_indices=req_pool_indices,
seq_lens=seq_lens,
req_to_token=req_to_token,
cache_seqlens=cache_seqlens,
cu_seqlens_k=cu_seqlens_k,
page_table=sentinel_pt,
bs=bs,
seqlen_offset=seqlen_offset,
max_seq_pages=max_num_pages,
page_size=PAGE_SIZE,
swa_page_table=sentinel_swa,
skip_page_table=True,
)
torch.cuda.synchronize()
cache_seqlens_ref = _ref_cache_seqlens(seq_lens, seqlen_offset)
torch.testing.assert_close(cache_seqlens, cache_seqlens_ref, rtol=0, atol=0)
cu_k_ref = torch.zeros(bs + 1, dtype=torch.int32, device=DEVICE)
cu_k_ref[1:] = torch.cumsum(cache_seqlens_ref, dim=0, dtype=torch.int32)
torch.testing.assert_close(cu_seqlens_k, cu_k_ref, rtol=0, atol=0)
if pass_tables:
assert bool((sentinel_pt == 777).all()), "page_table written despite skip"
assert bool((sentinel_swa == 888).all()), "swa_page_table written despite skip"
def test_bs_zero_noop():
if not torch.cuda.is_available():
pytest.skip("CUDA required")
test/registered/models_e2e/test_kimi_linear_unified_memory.py
+15 -1
View File
@@ -32,7 +32,7 @@ from sglang.test.kits.eval_accuracy_kit import GSM8KMixin
from sglang.test.kits.prefix_cache_branching_kit import PrefixCacheBranchingMixin
from sglang.test.server_fixtures.default_fixture import DefaultServerBase
register_cuda_ci(est_time=570, stage="nightly", runner_config="4-gpu-h100")
register_cuda_ci(est_time=1200, stage="nightly", runner_config="4-gpu-h100")
KIMI_LINEAR_MODEL = "moonshotai/Kimi-Linear-48B-A3B-Instruct"
@@ -55,5 +55,19 @@ class TestKimiLinearUnifiedMemory(
]
class TestKimiLinearUnifiedMemoryFlashMLA(TestKimiLinearUnifiedMemory):
"""flashmla at its ps=64 snap: the canonical block-table route
(KVIndexTranslator.build_into into flashmla's padded tables) plus the ps=64
sub-pool sizing (64-token sink floor, dense-view tail pad) end to end.
Hopper-only, like the rest of this nightly suite."""
other_args = TestKimiLinearUnifiedMemory.other_args + [
"--attention-backend",
"flashmla",
"--page-size",
"64",
]
if __name__ == "__main__":
unittest.main()
test/registered/page_major/test_page_major_gpt_oss.py
+21 -10
View File
@@ -1,15 +1,12 @@
"""
End-to-end accuracy test for the unified memory pool on a hybrid-SWA MoE model.
"""Unified memory pool on a hybrid-SWA MoE model, across the backend matrix.
Launches gpt-oss-20b with ``--enable-unified-memory`` on the Triton attention
backend and checks that GSM8K accuracy holds. This exercises the SWA +
full-attention KV sub-pools stored as per-layer views in the unified
page-major envelope.
gpt-oss-20b is uniform-row hybrid-SWA, so its MHA and SWA sub-pools are
per-layer views and the fa3 cell reads them through the translator's read
tables. The resolved-default cell pins the no-pin path, since a pinned
backend hides default-resolution breakage by construction. flashinfer is
absent on purpose: gpt-oss uses attention sinks, which it does not support.
Registered to the label-gated ``run-ci-extra`` suite (opt-in, not per-commit).
Usage:
python3 -m unittest test_page_major_gpt_oss
"""
import unittest
@@ -20,7 +17,7 @@ from sglang.test.ci.ci_register import register_cuda_ci
from sglang.test.server_fixtures.default_fixture import DefaultServerBase
from sglang.test.test_utils import DEFAULT_MODEL_NAME_FOR_TEST_MXFP4_WITH_MOE
register_cuda_ci(est_time=420, stage="extra-a", runner_config="1-gpu-large")
register_cuda_ci(est_time=1500, stage="extra-a", runner_config="1-gpu-large")
_UNIFIED_COMMON_ARGS = [
"--enable-unified-memory",
@@ -64,5 +61,19 @@ class TestUnifiedGptOssTriton(DefaultServerBase):
self.assertGreaterEqual(metrics["accuracy"], self.gsm8k_threshold)
class TestUnifiedGptOssFa3(TestUnifiedGptOssTriton):
"""fa3 pinned: the per-layer views read through the translator's read
tables (eager direct-bind + captured fused copy)."""
other_args = _UNIFIED_COMMON_ARGS + ["--attention-backend", "fa3"]
class TestUnifiedGptOssResolvedDefault(TestUnifiedGptOssTriton):
"""No backend pin: whatever the host resolves must be in the allow-list,
or the server fails to boot under its own defaults."""
other_args = _UNIFIED_COMMON_ARGS
if __name__ == "__main__":
unittest.main()
test/registered/page_major/test_page_major_qwen_hybrid.py
+29 -12
View File
@@ -1,17 +1,14 @@
"""
End-to-end accuracy test for the unified memory pool on a GDN-hybrid model.
"""Unified memory pool on a GDN-hybrid model, across the backend matrix.
Launches Qwen3.5-4B (a gated-delta-net / linear-attention hybrid) with
``--enable-unified-memory`` on the Triton attention + linear-attn + Mamba
backends and checks that GSM8K accuracy holds. This exercises the unified
envelope's most bug-prone path: the Mamba conv/SSM state stored as a strided
envelope view, plus the full-attention KV stored as per-layer views,
both read/written by the GDN prefill and decode kernels.
Qwen3.5-4B is a gated-delta-net / linear-attention hybrid, which exercises the
path most prone to subtle bugs: the Mamba conv/SSM state stays a strided
envelope view (its kernels are stride-aware by design) while the
full-attention KV is per-layer views, which the fa3 / flashinfer cells read
through the translator's read tables. The resolved-default cell pins the
no-pin path, since a pinned backend hides default-resolution breakage by
construction.
Registered to the label-gated ``run-ci-extra`` suite (opt-in, not per-commit).
Usage:
python3 -m unittest test_page_major_qwen_hybrid
"""
import unittest
@@ -22,7 +19,7 @@ from sglang.test.ci.ci_register import register_cuda_ci
from sglang.test.server_fixtures.default_fixture import DefaultServerBase
from sglang.test.test_utils import DEFAULT_HYBRID_GDN_SMALL_MODEL_NAME_FOR_TEST
register_cuda_ci(est_time=300, stage="extra-a", runner_config="1-gpu-large")
register_cuda_ci(est_time=1600, stage="extra-a", runner_config="1-gpu-large")
_UNIFIED_COMMON_ARGS = [
"--trust-remote-code",
@@ -74,5 +71,25 @@ class TestUnifiedQwenHybridTriton(DefaultServerBase):
self.assertGreaterEqual(metrics["accuracy"], self.gsm8k_threshold)
class TestUnifiedQwenHybridFa3(TestUnifiedQwenHybridTriton):
"""fa3 pinned: read tables, eager direct-bind + captured fused copy."""
other_args = _UNIFIED_COMMON_ARGS + ["--attention-backend", "fa3"]
class TestUnifiedQwenHybridFlashinfer(TestUnifiedQwenHybridTriton):
"""flashinfer pinned: token ids reconstructed from the read table by the
ENTRY_PAGE_SIZE CSR builder."""
other_args = _UNIFIED_COMMON_ARGS + ["--attention-backend", "flashinfer"]
class TestUnifiedQwenHybridResolvedDefault(TestUnifiedQwenHybridTriton):
"""No backend pin: whatever the host resolves must be in the allow-list,
or the server fails to boot under its own defaults."""
other_args = _UNIFIED_COMMON_ARGS
if __name__ == "__main__":
unittest.main()
test/registered/unit/layers/attention/test_flashattention_graph_metadata.py
+10
View File
@@ -4,6 +4,7 @@ from types import SimpleNamespace
import torch
from sglang.srt.layers.attention.flashattention_backend import FlashAttentionBackend
from sglang.srt.mem_cache.kv_index_translator import KVIndexTranslator
from sglang.srt.model_executor.forward_batch_info import ForwardMode
from sglang.test.ci.ci_register import register_cuda_ci
from sglang.test.test_utils import CustomTestCase
@@ -46,6 +47,15 @@ class TestFlashAttentionGraphMetadata(CustomTestCase):
backend.req_to_token_pool = SimpleNamespace(
req_to_token=torch.zeros((1, 16), dtype=torch.int32)
)
# A real source over the stub pool: the probe disables it, giving the
# backend the strict passthrough view it now reads its tables from.
backend.kv_index_translator = KVIndexTranslator(
req_to_token=backend.req_to_token_pool.req_to_token,
token_to_kv_pool_allocator=SimpleNamespace(),
token_to_kv_pool=SimpleNamespace(),
page_size=1,
device="cpu",
)
backend.is_prefill_aware_swa = False
backend.has_swa = False
backend.use_sliding_window_kv_pool = False
test/registered/unit/layers/attention/test_flashattention_pa_swa_prefill_lens_size.py
+14 -4
View File
@@ -24,6 +24,7 @@ import torch
from sglang.srt.configs.model_config import AttentionArch
from sglang.srt.layers.attention.flashattention_backend import FlashAttentionBackend
from sglang.srt.mem_cache.kv_index_translator import KVIndexTranslator
from sglang.srt.runtime_context import get_context
from sglang.test.ci.ci_register import register_cuda_ci
from sglang.test.test_utils import CustomTestCase
@@ -62,15 +63,24 @@ def _make_prefill_aware_swa_runner(*, pool_size: int, max_context_len: int = 64)
enable_prefill_cp=False,
enable_dp_attention=False,
)
token_to_kv_pool = object()
token_to_kv_pool_allocator = object()
return SimpleNamespace(
sliding_window_size=None,
model_config=model_config,
device=device,
req_to_token_pool=req_to_token_pool,
token_to_kv_pool=object(), # not a SWAKVPool instance -> use_sliding_window_kv_pool=False
# getattr(..., "full_v2p_page_table", None) is None -> unified_mla_hooks
# falls back to the static (disabled) hook set.
token_to_kv_pool_allocator=object(),
token_to_kv_pool=token_to_kv_pool, # not a SWAKVPool -> use_sliding_window_kv_pool=False
token_to_kv_pool_allocator=token_to_kv_pool_allocator,
# A real KVIndexTranslator over this non-unified pool: the probe finds no
# unified composite, so it is the strict passthrough the backend reads.
kv_index_translator=KVIndexTranslator(
req_to_token=req_to_token_pool.req_to_token,
token_to_kv_pool_allocator=token_to_kv_pool_allocator,
token_to_kv_pool=token_to_kv_pool,
page_size=1,
device=device,
),
kv_cache_dtype=torch.float16,
kv_cache_dtype_str="auto",
page_size=1,
test/registered/unit/layers/attention/test_kv_translate_ownership.py
+83
View File
@@ -0,0 +1,83 @@
"""Nothing under layers/attention may translate KV ids for itself.
Ownership is exactly two places: `KVIndexTranslator` for READS (indices are
born kernel-facing, backends consume its tables) and the ForwardBatch rebind
(`rebind_write_loc`) for WRITES. Virtual and physical ids share a value range,
so a backend that forgets a translate -- or does one twice -- reads the wrong
rows and nothing crashes. This scan makes both unrepresentable.
Out of scope, deliberately: the allocator-internal implementations
(`multi_ended_allocator` / `unified_memory_pool`), which ARE the mechanism the
translator calls; the PD transfer plane's `translate_kv_indices_for_transfer`,
which stages for RDMA outside the forward path; and the STATIC SWA pool's
legacy full->swa slot map, a different mapping kind with no virtual/physical
ambiguity -- its call sites are count-pinned below so new ones are added
consciously.
python3 -m pytest test/registered/unit/layers/attention/test_kv_translate_ownership.py -v
"""
import os
import re
import unittest
from sglang.srt.layers.attention import triton_backend as _anchor_module
from sglang.test.ci.ci_register import register_cpu_ci
from sglang.test.test_utils import CustomTestCase
register_cpu_ci(est_time=5, suite="base-a-test-cpu")
# The attention package is a namespace package (no __init__), so anchor the
# scan on a concrete module inside it.
_ATTN_DIR = os.path.dirname(os.path.abspath(_anchor_module.__file__))
def _iter_sources():
for root, _dirs, files in os.walk(_ATTN_DIR):
for name in sorted(files):
if not name.endswith(".py"):
continue
path = os.path.join(root, name)
with open(path, encoding="utf-8") as fh:
yield os.path.relpath(path, _ATTN_DIR), fh.read()
class TestUnifiedTranslateBanned(CustomTestCase):
def test_no_unified_translate_calls(self):
"""No backend calls the unified translate surfaces. A hit here means
a backend re-grew its own id-space transition -- the design whose two
failure modes (forgotten translate, duplicated translate) this scan
exists to prevent. Route reads through KVIndexTranslator views and
writes through the ForwardBatch rebind instead."""
banned = re.compile(r"\.translate_kv_loc(_kernel_id)?\(")
hits = [
f"{rel}: {m.group(0)}"
for rel, src in _iter_sources()
for m in banned.finditer(src)
]
self.assertEqual(hits, [])
def test_no_translate_capability_probing(self):
"""No backend probes an allocator for translate capability -- the
getattr-hook pattern is how per-backend translation grew the first
time."""
probing = re.compile(r"""getattr\([^)]*['"]translate_kv_loc""")
hits = [rel for rel, src in _iter_sources() if probing.search(src)]
self.assertEqual(hits, [])
def test_hooks_module_deleted_and_unimported(self):
"""The per-backend hooks module (the previous owner of backend-side
v2p knowledge) stays deleted, and nothing imports it."""
self.assertFalse(
os.path.exists(os.path.join(_ATTN_DIR, "unified_mem_hooks.py"))
)
hits = [
rel
for rel, src in _iter_sources()
if "unified_mem_hooks" in src or "unified_mla_hooks" in src
]
self.assertEqual(hits, [])
if __name__ == "__main__":
unittest.main()
test/registered/unit/mem_cache/test_full_loc_fast_path.py
+33 -33
View File
@@ -59,9 +59,11 @@ class _RecordingPool:
class TestUnifiedSWARouting(unittest.TestCase):
"""`UnifiedSWAKVPool.set_kv_buffer` routing: full layers write the full-physical
`full_loc`; SWA layers write the swa-physical `swa_loc`. Both come from the
write metadata; the pool never translates."""
"""`UnifiedSWAKVPool.set_kv_buffer` routing: full layers write `full_loc`
when present (triton's capture-stable buffer), else the rebound generic
`loc` -- the same id space once the loc is rebound; SWA layers write the swa-physical
`swa_loc`, which has no fallback (a different id space). The pool never
translates."""
def _make_bare_pool(self):
from sglang.srt.mem_cache.unified_memory_pool import UnifiedSWAKVPool
@@ -96,22 +98,30 @@ class TestUnifiedSWARouting(unittest.TestCase):
self.assertIsNot(forwarded, virtual_loc)
self.assertNotIn("already_physical", kwargs)
def test_full_layer_requires_full_loc(self):
def test_full_layer_falls_back_to_generic_loc(self):
"""Bug regression: fa3 x unified-SWA crashed at gpt-oss
cuda-graph capture because every backend except triton bundles the
2-arg KVWriteLoc(loc, swa) and the full-layer door demanded an explicit
full_loc. Once the loc is rebound the generic `loc` IS the full-side kernel-facing id
(rebind_write_loc runs at ForwardBatch construction),
so the door must fall back to it -- the pool still never translates."""
pool = self._make_bare_pool()
virtual_loc = torch.tensor([10, 11, 12], dtype=torch.int64)
rebound_loc = torch.tensor([10, 11, 12], dtype=torch.int64)
swa_phys = torch.tensor([1, 2, 0], dtype=torch.int64)
layer = types.SimpleNamespace(layer_id=0)
# No full_loc precomputed -> fail loud (the unified memory pool must precompute
# out_cache_loc_full_physical) rather than write a virtual loc as physical.
with self.assertRaises(AssertionError):
pool.set_kv_buffer(
layer,
_loc_info(virtual_loc, swa_phys),
_loc_info(rebound_loc, swa_phys),
torch.zeros(3, 4, 8),
torch.zeros(3, 4, 8),
)
self.assertEqual(len(pool.full_kv_pool.calls), 1)
forwarded, kwargs = pool.full_kv_pool.calls[0]
self.assertIs(forwarded, rebound_loc)
self.assertNotIn("already_physical", kwargs)
def test_swa_layer_writes_swa_loc(self):
pool = self._make_bare_pool()
virtual_loc = torch.tensor([10, 11, 12], dtype=torch.int64)
@@ -264,19 +274,17 @@ class TestHybridLinearMLARouting(unittest.TestCase):
- `set_kv_buffer` (MLA branch) mirrors the MHA branch — write the
pre-translated `KVWriteLoc.full_loc` when present (unified pool, where it
carries the DENSE loc), else the raw `loc` (static pool, already physical).
- `set_mla_kv_buffer` forwards `loc` untouched (kernel-facing since the
ForwardBatch rebind); `get_mla_kv_buffer` applies `_full_translate`
exactly once (its indices are req_to_token-produced, virtual under the
unified pool)."""
- `set_mla_kv_buffer` / `get_mla_kv_buffer` forward `loc` untouched:
writes are kernel-facing since the ForwardBatch rebind, and read
indices are translated at their production sites."""
def _make_bare_pool(self, translate=None):
def _make_bare_pool(self):
from sglang.srt.mem_cache.memory_pool import HybridLinearKVPool
pool = object.__new__(HybridLinearKVPool)
pool.full_kv_pool = _RecordingMLAPool()
pool.use_mla = True
pool.full_attention_layer_id_mapping = {0: 0}
pool._full_translate = translate if translate is not None else (lambda ids: ids)
return pool
def test_mla_writes_full_loc_from_write_loc(self):
@@ -328,28 +336,20 @@ class TestHybridLinearMLARouting(unittest.TestCase):
self.assertEqual(len(pool.full_kv_pool.mla_set_calls), 1)
self.assertIs(pool.full_kv_pool.mla_set_calls[0], loc)
def test_get_mla_kv_buffer_translates_exactly_once(self):
"""READ door: `loc` is produced from req_to_token (VIRTUAL under the
unified pool), so the get side still translates here — exactly once.
The WRITE door (case above) never translates: the split is the write
flip's contract."""
calls = []
def translate(ids):
calls.append(ids)
return ids + 100
pool = self._make_bare_pool(translate=translate)
virtual_loc = torch.tensor([4, 5], dtype=torch.int64)
def test_get_mla_kv_buffer_door_never_translates(self):
"""Kernel-facing contract, read side: `loc` is a read-index tensor
already translated at its production site
(fetch_mha_one_shot_kv_indices / prepare_chunked_kv_indices); the
door forwards it UNTOUCHED -- a re-added door translate would
double-translate every unified MLA prefix read."""
pool = self._make_bare_pool()
loc = torch.tensor([104, 105], dtype=torch.int64)
layer = types.SimpleNamespace(layer_id=0)
pool.get_mla_kv_buffer(layer, virtual_loc)
pool.get_mla_kv_buffer(layer, loc)
self.assertEqual(len(calls), 1)
self.assertEqual(len(pool.full_kv_pool.mla_get_calls), 1)
self.assertTrue(
torch.all(pool.full_kv_pool.mla_get_calls[0] == virtual_loc + 100)
)
self.assertIs(pool.full_kv_pool.mla_get_calls[0], loc)
if __name__ == "__main__":
test/registered/unit/mem_cache/test_multi_ended_allocator.py
+76
View File
@@ -31,11 +31,13 @@ import torch
from sglang.srt.mem_cache.multi_ended_allocator import (
MultiEndedAllocator,
UnifiedMambaTokenToKVPoolAllocator,
UnifiedSWATokenToKVPoolAllocator,
)
from sglang.srt.mem_cache.unified_memory_pool import (
MambaSubPoolSpec,
MHASubPoolSpec,
MLASubPoolSpec,
UnifiedKVPool,
)
@@ -2746,5 +2748,79 @@ class TestSWACompositeDenseSurface(unittest.TestCase):
self.assertTrue(bool((got[in_tomb] == 0).all().item()))
class TestPs64MLACompositeFeasibility(unittest.TestCase):
"""The Kimi/flashmla shape: MLA + mamba composite at page_size=64 (the
flashmla arg snap). Large pages stress every sizing derivation at once —
the 64-token sink-page floor, the ps*entry_bytes per-layer-view tail pad, and
the page-granular alloc — so this pins that the factory-shaped
construction stays FEASIBLE and the dense surface stays on-formula when
the page size jumps from the usual 1..4 to 64."""
PS = 64
LAYERS = 3
def _build(self):
full = MLASubPoolSpec(
name="full",
layer_num=self.LAYERS,
kv_lora_rank=64,
qk_rope_head_dim=16,
store_dtype=torch.float16,
grow_direction="down",
)
mamba = MambaSubPoolSpec(
name="mamba",
layer_num=2,
conv_state_shapes=((8, 16),),
conv_dtype=torch.bfloat16,
temporal_state_shape=(4, 8, 8),
temporal_dtype=torch.float32,
grow_direction="up",
)
n_full = 8 * self.PS # 8 pages incl. the sink page
total = n_full * full.entry_bytes() + 16 * mamba.entry_bytes()
pool = UnifiedKVPool(
total_bytes=total,
sub_pool_specs=[full, mamba],
device=_DEV,
enable_memory_saver=False,
page_size=self.PS,
)
full_kv = _FakeKVCache(pool.max_slots("full"))
full_kv.attach_allocator = lambda allocator: None
mamba_kv = _FakeKVCache(pool.max_slots("mamba"))
mamba_kv.attach_allocator = lambda allocator: None
mamba_kv._copy_from_physical = lambda src, dst: None
class _FakeHybridLinearKVPool:
full_kv_pool = full_kv
mamba_pool = mamba_kv
return UnifiedMambaTokenToKVPoolAllocator(
unified_buffer=pool,
kvcache=_FakeHybridLinearKVPool(),
device=_DEV,
page_size=self.PS,
need_sort=False,
forward_stream=None,
)
def test_construction_alloc_and_dense_formula(self):
a = self._build()
# MLA: one latent row per layer, so the spec reports LAYERS blocks.
self.assertEqual(a.kernel_page_multiplier, self.LAYERS)
v = a.alloc(2 * self.PS)
self.assertIsNotNone(v, "2-page alloc infeasible at ps=64")
# Page-aligned virtual run (page-granular allocator invariant).
self.assertEqual(int(v[0].item()) % self.PS, 0)
# Dense translate follows the affine formula at ps=64, and every id
# fits int32 (the canonical narrows on store).
v2p = a.full_v2p_page_table
want = v2p[v // self.PS] * (self.PS * self.LAYERS) + v % self.PS
got = a.translate_kv_loc_for_kernel(v)
self.assertTrue(torch.equal(got, want), "kernel-facing formula broke at ps=64")
self.assertTrue(bool((got < 2**31).all().item()))
if __name__ == "__main__":
unittest.main()
test/registered/unit/mem_cache/test_unified_mla_block_table.py
+95 -216
View File
@@ -20,20 +20,26 @@ block table filled with kernel-facing page ids:
dense_page(virtual_page) = v2p[virtual_page] * layer_num
Three backend families reach that same formula by different routes:
- `create_flashmla_kv_indices_triton` in-kernel via `v2p_ptr` / `PAGE_MULT`
(trtllm_mla / cutedsl_mla / tokenspeed_mla);
- the flashinfer_mla updaters, post-gathering `translate_kv_loc_for_kernel` over the
token-level kv_indices;
- `normal_decode_set_metadata` in-kernel, for fa3's captured-decode page table.
Since the read-path translator, ONE builder computes that formula for every
family — `build_kv_read_table` (the canonical) — and the backends only
differ in how they consume it:
- trtllm_mla / cutedsl_mla / tokenspeed_mla / flashmla: rows filled straight
into their padded block tables (`KVIndexTranslator.build_into`, prefix-only so
the backends' own -1 / stale tail sentinels survive);
- the flashinfer updaters: token ids reconstructed from the canonical by
`create_flashinfer_kv_indices_triton[ENTRY_PAGE_SIZE=ps]`;
- fa3's captured decode: `normal_decode_set_metadata` copies the canonical
rows' live prefixes (src_is_read_table=True).
Covered here:
- kernel identity: `v2p_ptr=None, PAGE_MULT=1` is byte-identical to main;
- kernel dense mapping against the python reference, for several page sizes,
ragged sequence lengths and a non-identity v2p permutation;
- padded block-table lanes never index the v2p table out of bounds;
- the token-level kernel-facing translate the flashinfer updaters apply agrees with the
page-level block table the trtllm path builds;
- the static `create_flashmla_kv_indices_triton` (no id-space knowledge left)
still matches the plain token//ps reference;
- the canonical route against the python dense reference, for several page
sizes, ragged sequence lengths and a non-identity v2p permutation;
- lanes past a row's live prefix keep the backend's -1 sentinel (prefix-only
discipline — the trtllm/flashmla tail contract);
- the token-level kernel-facing translate the flashinfer updaters used to apply
agrees with the canonical page table (page-affinity of the id space);
- fa3's fused metadata kernels agree with the same reference, on both the
page_size == 1 fast path (which is what Kimi-Linear takes: fa3 imposes no
page-size constraint) and the general path.
@@ -57,6 +63,28 @@ _LAYERS = 24 # K3 MLA full-attention layer count
def _fill_block_table(
req_to_token, req_pool_indices, seq_lens, page_size, *, v2p, mult
):
"""The unified route: canonical builder into a -1-filled block table
(exactly what KVIndexTranslator.build_into does for trtllm_mla/flashmla)."""
from sglang.kernels.ops.kvcache.kv_read_table import build_kv_read_table
bs = req_pool_indices.shape[0]
max_blocks = (int(seq_lens.max().item()) + page_size - 1) // page_size
out = torch.full((bs, max_blocks), -1, dtype=torch.int32, device=_DEV)
build_kv_read_table(
req_to_token=req_to_token,
req_pool_indices=req_pool_indices,
seq_lens=seq_lens.to(torch.int64),
v2p=v2p,
multiplier=mult,
page_size=page_size,
max_pages=max_blocks,
out=out,
)
return out
def _fill_block_table_static(req_to_token, req_pool_indices, seq_lens, page_size):
"""The static-pool route: the stripped flashmla kernel, token//ps verbatim."""
from sglang.kernels.ops.kvcache.kv_indices import (
create_flashmla_kv_indices_triton,
get_num_kv_index_blocks_flashmla,
@@ -75,8 +103,6 @@ def _fill_block_table(
req_to_token.stride(0),
max_blocks,
PAGED_SIZE=page_size,
v2p_ptr=v2p,
PAGE_MULT=mult,
)
return out
@@ -122,11 +148,12 @@ class TestDenseBlockTable(unittest.TestCase):
v2p[0] = 0 # page 0 is the reserved sink
return req_to_token, req_pool_indices, seq_lens, v2p
def test_identity_when_hooks_absent(self):
"""v2p_ptr=None / PAGE_MULT=1 must reproduce the pre-change behaviour."""
def test_static_kernel_matches_reference(self):
"""The stripped (id-space-free) flashmla kernel is byte-identical to the
plain token//ps reference -- guards the v2p-arg removal itself."""
for page_size in (1, 32, 64):
rt, rpi, sl, _ = self._make_batch(page_size)
got = _fill_block_table(rt, rpi, sl, page_size, v2p=None, mult=1)
got = _fill_block_table_static(rt, rpi, sl, page_size)
want = _reference(rt, rpi, sl, page_size, v2p=None, mult=1)
self.assertTrue(
torch.equal(got.long(), want), f"page_size={page_size}: {got} != {want}"
@@ -166,8 +193,9 @@ class TestDenseBlockTable(unittest.TestCase):
)
def test_padded_lanes_stay_untouched(self):
"""Lanes past a request's page count keep the -1 fill: the masked v2p load
must not write a translated value (nor read out of bounds)."""
"""Lanes past a request's page count keep the -1 fill: the prefix-only
canonical build must never write a backend's tail sentinel (the
trtllm/flashmla block-table contract)."""
page_size = 64
rt, rpi, sl, v2p = self._make_batch(page_size)
got = _fill_block_table(rt, rpi, sl, page_size, v2p=v2p, mult=_LAYERS)
@@ -207,43 +235,77 @@ class TestDenseBlockTable(unittest.TestCase):
@unittest.skipUnless(_HAS_CUDA, "requires CUDA")
class TestFa3MetadataDenseBlockTable(unittest.TestCase):
"""fa3 folds the unified remap into `normal_decode_set_metadata`, the fused
gather that writes its captured-decode page table, so the kernel itself has
to get the mapping right. Two kernels back it: a page_size == 1 / no-SWA fast
path (what Kimi-Linear takes, since fa3 imposes no page-size constraint) and
a general one.
"""fa3's captured-decode page table is written by `normal_decode_set_metadata`
fed with the translator's read table kernel page table
(src_is_read_table=True): the fused kernel copies the canonical
rows' live prefixes into the capture-stable buffer. Pinned END-TO-END:
build_kv_read_table -> wrapper -> page_table must equal the python
reference of the kernel-facing formula, on both the page_size == 1 / no-SWA fast
path (what Kimi-Linear takes) and the general kernel. The static call
(no source flag) stays byte-identical to the pre-translator kernel.
"""
def _run(self, page_size, *, v2p, mult, bs=5, max_ctx=2048):
from sglang.kernels.ops.attention.metadata import normal_decode_set_metadata
from sglang.kernels.ops.kvcache.kv_read_table import (
build_kv_read_table,
)
maker = TestDenseBlockTable._make_batch
rt, rpi, sl, v2p_full = maker(self, page_size, bs=bs, max_ctx=max_ctx)
v2p_arg = v2p_full if v2p else None
max_pages = (max_ctx + page_size - 1) // page_size
page_table = torch.zeros((bs, max_pages), dtype=torch.int32, device=_DEV)
cache_seqlens = torch.zeros((bs,), dtype=torch.int32, device=_DEV)
cu_seqlens_k = torch.zeros((bs + 1,), dtype=torch.int32, device=_DEV)
strided = torch.arange(0, max_ctx, page_size, device=_DEV)
max_seq_pages = (int(sl.max().item()) + page_size - 1) // page_size
if v2p:
# The translator's canonical, then the wrapper copies its rows.
canonical = torch.zeros((bs, max_pages), dtype=torch.int32, device=_DEV)
build_kv_read_table(
req_to_token=rt,
req_pool_indices=rpi,
seq_lens=sl.to(torch.int64),
v2p=v2p_full,
multiplier=mult,
page_size=page_size,
max_pages=max_pages,
out=canonical,
)
rows = torch.arange(bs, dtype=torch.int64, device=_DEV)
normal_decode_set_metadata(
cache_seqlens,
cu_seqlens_k,
page_table,
canonical,
rows,
max_seq_pages,
sl.to(torch.int64),
0,
page_size,
None,
None,
src_is_read_table=True,
)
else:
normal_decode_set_metadata(
cache_seqlens,
cu_seqlens_k,
page_table,
rt,
rpi,
strided,
max_seq_pages,
sl.to(torch.int64),
0,
page_size,
v2p_page_table=v2p_arg,
kernel_page_multiplier=mult,
None,
None,
)
torch.cuda.synchronize()
want = _reference(rt, rpi, sl, page_size, v2p=v2p_arg, mult=mult)
want = _reference(
rt, rpi, sl, page_size, v2p=(v2p_full if v2p else None), mult=mult
)
return page_table, want, sl
def _assert_live_prefix(self, got, want, sl, page_size):
@@ -287,192 +349,9 @@ class TestFa3MetadataDenseBlockTable(unittest.TestCase):
"test batch degenerated: v2p is the identity on the pages used",
)
def test_agrees_with_flashmla_block_table(self):
"""fa3 and trtllm_mla build the same table two different ways; a
disagreement means one family is addressing the wrong pages."""
for page_size in (1, 64):
got, _, sl = self._run(page_size, v2p=True, mult=_LAYERS)
rt, rpi, sl2, v2p = TestDenseBlockTable._make_batch(self, page_size)
other = _fill_block_table(
rt, rpi, sl2, page_size, v2p=v2p, mult=_LAYERS
).long()
for r in range(got.shape[0]):
n_pages = (int(sl[r].item()) + page_size - 1) // page_size
self.assertTrue(
torch.equal(got[r, :n_pages].long(), other[r, :n_pages]),
f"fa3 and flashmla block tables disagree (row {r}, ps={page_size})",
)
class TestUnifiedMLAHookDetection(unittest.TestCase):
"""`unified_mla_hooks` decides whether the paged MLA backends translate at
all. Getting the predicate wrong is silent: the block table and KV write loc
stay in virtual id space and address the wrong pages once virtual and
physical diverge (e.g. after compaction)."""
@staticmethod
def _probe(**attrs):
from sglang.srt.layers.attention.unified_mem_hooks import (
unified_mla_hooks,
)
class _Alloc:
pass
alloc = _Alloc()
for k, v in attrs.items():
setattr(alloc, k, v)
return unified_mla_hooks(alloc)
def test_static_pool_disables_every_hook(self):
"""No v2p table -> statically-partitioned pool; req_to_token is already
physical, so all hooks must stay off (byte-identical to pre-change)."""
hooks = self._probe()
self.assertFalse(hooks.enabled)
self.assertIsNone(hooks.v2p_page_table)
self.assertIsNone(hooks.translate_kv_loc_for_kernel)
self.assertEqual(hooks.kernel_page_multiplier, 1)
def test_multi_layer_unified_pool(self):
table = torch.arange(8)
hooks = self._probe(
full_v2p_page_table=table,
translate_kv_loc_for_kernel=lambda x, **kw: x,
kernel_page_multiplier=_LAYERS,
)
self.assertTrue(hooks.enabled)
self.assertIs(hooks.v2p_page_table, table)
self.assertIsNotNone(hooks.translate_kv_loc_for_kernel)
self.assertEqual(hooks.kernel_page_multiplier, _LAYERS)
def test_single_full_attention_layer_pool_is_still_unified(self):
"""REGRESSION: `kernel_page_multiplier == 1` does NOT mean static.
A hybrid MLA config with exactly one full-attention layer (e.g. a
pipeline-parallel rank owning a single MLA layer) has multiplier 1, yet
its locs are still virtual. Detecting on `multiplier > 1` would disable
the v2p gather here and corrupt reads/writes after compaction.
"""
table = torch.arange(8)
hooks = self._probe(
full_v2p_page_table=table,
translate_kv_loc_for_kernel=lambda x, **kw: x,
kernel_page_multiplier=1,
)
self.assertTrue(hooks.enabled, "single-layer unified pool read as static")
self.assertIs(hooks.v2p_page_table, table)
self.assertIsNotNone(hooks.translate_kv_loc_for_kernel)
# Multiplier stays 1: kernel-facing id == physical id, so the v2p gather alone is
# the whole translation and PAGE_MULT must not scale it.
self.assertEqual(hooks.kernel_page_multiplier, 1)
@unittest.skipUnless(_HAS_CUDA, "requires CUDA")
class TestInPlaceKvIndicesTranslate(unittest.TestCase):
"""The flashinfer decode updater must translate kv_indices IN PLACE.
Under cuda-graph replay the `kv_indices` it is handed IS the capture-stable
buffer the captured wrapper reads (`fast_decode_kwargs["kv_indices"]`), and
`fast_mla_decode_plan` ignores its `kv_indices` argument -- so rebinding the
local name to a fresh translated tensor leaves the graph reading VIRTUAL ids.
These pin the write-back contract that fix relies on.
"""
def _allocator(self, page_size=1, n_full_tokens=4096):
from sglang.srt.mem_cache.multi_ended_allocator import MultiEndedAllocator
from sglang.srt.mem_cache.unified_memory_pool import (
MambaSubPoolSpec,
MLASubPoolSpec,
UnifiedKVPool,
)
full = MLASubPoolSpec(
name="full",
layer_num=_LAYERS,
kv_lora_rank=512,
qk_rope_head_dim=64,
store_dtype=torch.bfloat16,
grow_direction="down",
)
mamba = MambaSubPoolSpec(
name="mamba",
layer_num=2,
conv_state_shapes=((8, 16),),
conv_dtype=torch.bfloat16,
temporal_state_shape=(4, 8, 8),
temporal_dtype=torch.float32,
grow_direction="up",
)
pool = UnifiedKVPool(
total_bytes=full.entry_bytes() * n_full_tokens + mamba.entry_bytes() * 16,
sub_pool_specs=[full, mamba],
device=_DEV,
enable_memory_saver=False,
page_size=page_size,
)
class _Stub:
def move_kv_cache(self, dst, src):
pass
full_alloc = MultiEndedAllocator(
kvcache=_Stub(),
unified_buffer=pool,
sub_pool_name="full",
device=_DEV,
is_id_owner=True,
page_size=page_size,
kernel_page_multiplier=_LAYERS,
)
mamba_alloc = MultiEndedAllocator(
kvcache=_Stub(),
unified_buffer=pool,
sub_pool_name="mamba",
device=_DEV,
is_id_owner=True,
)
full_alloc.bind_peer(mamba_alloc)
mamba_alloc.bind_peer(full_alloc)
return full_alloc
def test_int32_buffer_prefix_translated_tail_untouched(self):
"""Mirrors the updater: an int32 capture-stable buffer holding VIRTUAL
ids in [:n] gets the kernel-facing ids written back in place, narrowed to int32,
with the stale tail left alone (it must never index the v2p table)."""
alloc = self._allocator()
virt = alloc.alloc(64)
self.assertIsNotNone(virt)
n = virt.numel()
# Capture-stable int32 buffer: [:n] freshly filled with virtual ids by
# create_flashinfer_kv_indices_triton, tail = stale junk from a bigger replay.
buf = torch.full((n * 3,), 2**30, dtype=torch.int32, device=_DEV)
buf[:n] = virt.to(torch.int32)
tail_before = buf[n:].clone()
valid = buf[:n]
valid.copy_(alloc.translate_kv_loc_for_kernel(valid))
expected = alloc.translate_kv_loc_for_kernel(virt)
self.assertEqual(buf.dtype, torch.int32)
self.assertTrue(
torch.equal(buf[:n].long(), expected),
"in-place translate did not land kernel-facing ids in the stable buffer",
)
self.assertTrue(
torch.equal(buf[n:], tail_before),
"stale tail was modified -- it can hold ids outside the v2p table",
)
def test_dense_ids_differ_from_virtual(self):
"""Guard the guard: if dense == virtual the in-place test proves nothing."""
alloc = self._allocator()
virt = alloc.alloc(64)
self.assertIsNotNone(virt)
self.assertFalse(
torch.equal(alloc.translate_kv_loc_for_kernel(virt), virt),
"kernel-facing ids coincide with virtual ids; pick a different allocation",
)
# (The old fa3<->flashmla agreement case is gone: both families now
# consume the SAME canonical builder, so cross-family agreement holds by
# construction and the per-family cases above cover the two consumers.)
if __name__ == "__main__":
test/registered/unit/model_executor/test_unified_out_cache_loc_rebind.py
+100
View File
@@ -27,6 +27,7 @@ import inspect
import textwrap
import unittest
from types import SimpleNamespace
from unittest.mock import create_autospec
import torch
@@ -177,5 +178,104 @@ class TestPadComposesWithDerivation(CustomTestCase):
self.assertEqual(src._swa_write_loc_unified(fb.out_cache_loc).numel(), 0)
class TestReadRailTranslatesAtProduction(CustomTestCase):
"""The model-door READ indices (req_to_token-derived, VIRTUAL under the
unified pool) are translated at their PRODUCTION site -- the cache then
holds the kernel-facing result and the pool door never translates."""
def _fb_for_one_shot(self):
fb = _make_fb(torch.tensor([1, 2], dtype=torch.int64))
fb.batch_size = 2
fb.seq_lens = torch.tensor([2, 3], dtype=torch.int64)
fb.seq_lens_cpu = torch.tensor([2, 3], dtype=torch.int32)
fb.req_pool_indices = torch.tensor([0, 1], dtype=torch.int64)
return fb
def test_one_shot_indices_translated_once_and_cached(self):
from unittest.mock import patch
from sglang.srt.model_executor import forward_batch_deepseek_mha_mixin as mix
calls = []
sentinel = torch.arange(5, dtype=torch.int64) + 5000
def translate(t):
calls.append(t)
return sentinel
fb = self._fb_for_one_shot()
fake_pool = SimpleNamespace(
req_to_token=torch.zeros((4, 16), dtype=torch.int32)
)
# autospec, not a bare namespace: setting a name the translator does
# not have raises, so renaming the method breaks this test loudly.
fake_translator = create_autospec(KVIndexTranslator, instance=True)
fake_translator.translate_full_attn_ids = translate
fake_backend = SimpleNamespace(kv_index_translator=fake_translator)
with (
patch.object(mix, "get_req_to_token_pool", return_value=fake_pool),
patch.object(mix, "get_attn_backend", return_value=fake_backend),
patch.object(mix, "create_flashinfer_kv_indices_triton"),
):
r1 = fb.fetch_mha_one_shot_kv_indices()
r2 = fb.fetch_mha_one_shot_kv_indices()
self.assertIs(r1, sentinel) # production site translated
self.assertIs(r2, sentinel) # cache holds the TRANSLATED result
self.assertEqual(len(calls), 1) # translated exactly once
self.assertEqual(calls[0].dtype, torch.int32) # raw producer output
def test_one_shot_indices_noop_on_unmigrated_backend(self):
from unittest.mock import patch
from sglang.srt.model_executor import forward_batch_deepseek_mha_mixin as mix
fb = self._fb_for_one_shot()
fake_pool = SimpleNamespace(
req_to_token=torch.zeros((4, 16), dtype=torch.int32)
)
# A backend that never set the attribute inherits the base-class None.
fake_backend = SimpleNamespace(kv_index_translator=None)
with (
patch.object(mix, "get_req_to_token_pool", return_value=fake_pool),
patch.object(mix, "get_attn_backend", return_value=fake_backend),
patch.object(mix, "create_flashinfer_kv_indices_triton"),
):
r = fb.fetch_mha_one_shot_kv_indices()
# The raw int32 producer output passes through untouched.
self.assertEqual(r.dtype, torch.int32)
def test_get_mla_kv_buffer_door_passes_loc_untranslated(self):
from sglang.srt.mem_cache.memory_pool import HybridLinearKVPool
recorded = {}
class _RecordingLeafPool:
def get_mla_kv_buffer(self, layer, loc, dst_dtype):
recorded["loc"] = loc
return None, None
def get_kv_size_bytes(self):
return 0
pool = HybridLinearKVPool(
size=16,
dtype=torch.float16,
page_size=1,
head_num=1,
head_dim=8,
full_attention_layer_ids=[0],
device=_DEV,
mamba_pool=SimpleNamespace(get_size_per_token=lambda: 0),
enable_memory_saver=False,
use_mla=True,
start_layer=0,
full_kv_pool=_RecordingLeafPool(),
)
loc = torch.tensor([9, 10], dtype=torch.int64)
pool.get_mla_kv_buffer(SimpleNamespace(layer_id=0), loc, torch.float16)
self.assertIs(recorded["loc"], loc)
if __name__ == "__main__":
unittest.main()
test/registered/unit/server_args/test_page_major_backend_allowlist.py
+57 -27
View File
@@ -13,18 +13,28 @@
# ==============================================================================
"""`--enable-page-major-kv-layout` full-attention backend allowlist.
The page-major envelope K/V views are strided, which only the Triton attention
kernels read. The one exception is the unified-memory MLA pool: it exposes each
layer as a contiguous view (`build_mla_views`), so the paged MLA
backends can read it directly once their kv_indices / block tables are remapped
to kernel-facing ids -- `fa3`, `flashinfer`'s MLA backend, and `trtllm_mla` with its
`cutedsl_mla` / `tokenspeed_mla` subclasses.
Two-way gate (see `_handle_page_major_kv_layout`), because the unified pool
exposes per-layer views and nothing else:
* unified-memory MLA models (`build_mla_views`) allow the whole wired
paged MLA family -- `fa3`, `flashinfer`'s MLA backend, `trtllm_mla` with
its `cutedsl_mla` / `tokenspeed_mla` subclasses, and `flashmla` (ps=64
snap);
* unified-memory MHA/SWA models (`build_mha_views`) allow `fa3` /
`fa4` / `flashinfer` / `trtllm_mha` alongside Triton;
* plain `--enable-page-major-kv-layout` without the unified pool keeps the
envelope-strided 4-D views only the stride-aware Triton kernels read.
Pinned here so the exception cannot silently widen to a backend that has no
dense-id remapping (`flashmla`, `cutlass_mla`, ...) or leak into the MHA path.
`fa3` matters most: it is the resolved default on pre-Blackwell hosts, so it is
the one entry whose absence used to make `--enable-unified-memory` fail to boot
under its own default configuration.
The same handler also screens the pool itself: the dense MHA/SWA views need
uniform K/V rows, so an asymmetric-K/V model (MiMoV2: head_dim 192 !=
v_head_dim 128) cannot run `--enable-unified-memory` at all and is rejected on
EVERY backend, Triton included. MLA models are exempt -- their sub-pool keeps
one latent row per layer, and several MLA configs (Kimi-Linear: head_dim 72,
v_head_dim 128) report asymmetric dims while running the unified pool today.
Pinned here so no arm silently widens to an unwired backend (`cutlass_mla`,
`aiter`) and no arm silently narrows: `fa3` is the resolved default on
pre-Blackwell hosts, so its absence from an arm makes `--enable-unified-memory`
fail to boot under its own default configuration.
python -m pytest test/registered/unit/server_args/test_page_major_backend_allowlist.py -v
"""
@@ -96,13 +106,23 @@ class TestPageMajorBackendAllowlist(unittest.TestCase):
"flashinfer",
"cutedsl_mla",
"tokenspeed_mla",
"flashmla",
)
# No dense-id remapping: must stay rejected until they get one.
UNWIRED_BACKENDS = ("flashmla", "cutlass_mla", "trtllm_mha", "aiter")
# Wired for the dense per-layer MHA/SWA views (uniform-row models).
DENSE_MHA_BACKENDS = ("fa3", "fa4", "flashinfer", "trtllm_mha")
# MLA-family kernels that must never leak into the MHA arm.
MLA_ONLY_BACKENDS = ("trtllm_mla", "cutedsl_mla", "tokenspeed_mla", "flashmla")
# No dense-id wiring anywhere: must stay rejected until they get one.
UNWIRED_BACKENDS = ("cutlass_mla", "aiter")
def test_triton_always_allowed(self):
def test_triton_allowed_on_every_arm(self):
"""Triton reads both view families, so it is the one backend neither
the MLA nor the MHA arm can narrow away."""
for use_mla in (True, False):
self.assertTrue(_accepts("triton", use_mla=use_mla))
# The uniform-row screen is a property of the model, not of the
# backend, so it rejects even Triton.
self.assertFalse(_accepts("triton", use_mla=False, has_asymmetric_kv=True))
def test_dense_mla_backends_allowed_under_unified_mla(self):
for backend in self.DENSE_MLA_BACKENDS:
@@ -111,21 +131,18 @@ class TestPageMajorBackendAllowlist(unittest.TestCase):
f"{backend} should be allowed with the unified-memory MLA pool",
)
def test_dense_mla_backends_rejected_for_mha(self):
"""The per-layer-view exception is MLA-only -- MHA sub-pools stay strided."""
for backend in self.DENSE_MLA_BACKENDS:
self.assertFalse(
def test_dense_mha_backends_allowed_for_uniform_row_models(self):
for backend in self.DENSE_MHA_BACKENDS:
self.assertTrue(
_accepts(backend, use_mla=False),
f"{backend} must stay rejected for a non-MLA model",
f"{backend} should be allowed for a uniform-row MHA model",
)
def test_dense_mla_backends_rejected_without_unified_memory(self):
"""Plain --enable-page-major-kv-layout (no unified pool) keeps the
strided views, so only Triton can read them."""
for backend in self.DENSE_MLA_BACKENDS:
def test_mla_only_backends_rejected_for_mha(self):
for backend in self.MLA_ONLY_BACKENDS:
self.assertFalse(
_accepts(backend, use_mla=True, unified=False),
f"{backend} must stay rejected without --enable-unified-memory",
_accepts(backend, use_mla=False),
f"{backend} is an MLA kernel and must stay out of the MHA arm",
)
def test_plain_page_major_arm_is_gated_at_boot(self):
@@ -143,7 +160,7 @@ class TestPageMajorBackendAllowlist(unittest.TestCase):
"""head_dim != v_head_dim (MiMoV2): no uniform rows, so no per-layer views
and no unified pool. The rejection is the POOL's, not a backend's, so
it must fire on every backend -- Triton included."""
for backend in ("triton",) + self.DENSE_MLA_BACKENDS:
for backend in ("triton",) + self.DENSE_MHA_BACKENDS:
self.assertFalse(
_accepts(backend, use_mla=False, has_asymmetric_kv=True),
f"--enable-unified-memory + {backend} must be rejected for an "
@@ -162,12 +179,25 @@ class TestPageMajorBackendAllowlist(unittest.TestCase):
"K/V head dims",
)
def test_page_major_rejected_without_unified_memory(self):
"""--enable-page-major-kv-layout without the unified pool is rejected
outright, Triton included: the static page-major arm went away with
the strided views and awaits its per-layer-view reimplementation."""
for backend in ("triton",) + tuple(
set(self.DENSE_MLA_BACKENDS + self.DENSE_MHA_BACKENDS)
):
for use_mla in (True, False):
self.assertFalse(
_accepts(backend, use_mla=use_mla, unified=False),
f"{backend} must stay rejected without --enable-unified-memory",
)
def test_unwired_backends_always_rejected(self):
for backend in self.UNWIRED_BACKENDS:
for use_mla in (True, False):
self.assertFalse(
_accepts(backend, use_mla=use_mla),
f"{backend} has no dense-id remapping and must be rejected",
f"{backend} has no dense-id wiring and must be rejected",
)
def test_helion_linear_attention_is_kda_only(self):