diff --git a/python/sglang/kernels/ops/attention/metadata.py b/python/sglang/kernels/ops/attention/metadata.py index 54762daff..7e305da9e 100644 --- a/python/sglang/kernels/ops/attention/metadata.py +++ b/python/sglang/kernels/ops/attention/metadata.py @@ -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,40 +252,38 @@ 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 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", + ) + 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_mapping_index = page_index + swa_val = swa_slot 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, - mask=mask, - other=0, - cache_modifier=".cg", - ) - if page_size == 1 or SWA_MAPPING_IS_V2P: - swa_val = swa_slot - else: - swa_val = swa_slot >> SHIFT + swa_val = swa_slot >> SHIFT swa_offsets = ( i * swa_page_table_stride_0 + col_offsets * swa_page_table_stride_1 ) @@ -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, ) diff --git a/python/sglang/kernels/ops/kvcache/kv_indices.py b/python/sglang/kernels/ops/kvcache/kv_indices.py index 9f2de2e45..fc6d6879c 100644 --- a/python/sglang/kernels/ops/kvcache/kv_indices.py +++ b/python/sglang/kernels/ops/kvcache/kv_indices.py @@ -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, ) diff --git a/python/sglang/kernels/ops/kvcache/trtllm_mha_graph_metadata.py b/python/sglang/kernels/ops/kvcache/trtllm_mha_graph_metadata.py index b1a062fc2..b869c2f16 100644 --- a/python/sglang/kernels/ops/kvcache/trtllm_mha_graph_metadata.py +++ b/python/sglang/kernels/ops/kvcache/trtllm_mha_graph_metadata.py @@ -62,38 +62,44 @@ 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) - 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 - # 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) - for i in range(tl.cdiv(num_live_pages, PAGE_BLOCK)): - page_idx = i * PAGE_BLOCK + tl.arange(0, PAGE_BLOCK) - mask = page_idx < num_live_pages - token = tl.load( - row_in + page_idx.to(tl.int64) * PAGE_SIZE, mask=mask, other=0 - ) - tl.store(row_out + page_idx, token // PAGE_SIZE, mask=mask) + 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: - token64 = token.to(tl.int64) - # Real req_to_token slots are >=0; the token>=0 guard + other=-1 mirror - # the swa_out_cache_loc -1 sentinel (uniform handling, no wrap). - swa_token = tl.load( - swa_mapping_ptr + token64, mask=mask & (token64 >= 0), other=-1 + swa_row_out = ( + swa_page_table_ptr + pid.to(tl.int64) * swa_page_table_stride ) - swa_page = tl.where(swa_token < 0, -1, swa_token // PAGE_SIZE) - tl.store(swa_row_out + page_idx, swa_page.to(tl.int32), mask=mask) + # 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) + for i in range(tl.cdiv(num_live_pages, PAGE_BLOCK)): + page_idx = i * PAGE_BLOCK + tl.arange(0, PAGE_BLOCK) + mask = page_idx < num_live_pages + token = tl.load( + row_in + page_idx.to(tl.int64) * PAGE_SIZE, mask=mask, other=0 + ) + tl.store(row_out + page_idx, token // PAGE_SIZE, mask=mask) + if HAS_SWA: + token64 = token.to(tl.int64) + # Real req_to_token slots are >=0; the token>=0 guard + other=-1 mirror + # the swa_out_cache_loc -1 sentinel (uniform handling, no wrap). + swa_token = tl.load( + swa_mapping_ptr + token64, mask=mask & (token64 >= 0), other=-1 + ) + swa_page = tl.where(swa_token < 0, -1, swa_token // PAGE_SIZE) + tl.store(swa_row_out + page_idx, swa_page.to(tl.int32), mask=mask) elif pid == bs: # Single program: cu_seqlens_k (+ optional cu_seqlens_q) cumsum. offs = tl.arange(0, BS_BLOCK) @@ -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, ) diff --git a/python/sglang/srt/arg_groups/kv_cache_hook.py b/python/sglang/srt/arg_groups/kv_cache_hook.py index 1200ec1ac..c05c4b26d 100644 --- a/python/sglang/srt/arg_groups/kv_cache_hook.py +++ b/python/sglang/srt/arg_groups/kv_cache_hook.py @@ -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): diff --git a/python/sglang/srt/layers/attention/base_attn_backend.py b/python/sglang/srt/layers/attention/base_attn_backend.py index cf76f34b2..b988c9a55 100644 --- a/python/sglang/srt/layers/attention/base_attn_backend.py +++ b/python/sglang/srt/layers/attention/base_attn_backend.py @@ -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)``. diff --git a/python/sglang/srt/layers/attention/cutedsl_mla_backend.py b/python/sglang/srt/layers/attention/cutedsl_mla_backend.py index 137718c99..bd93ed57e 100644 --- a/python/sglang/srt/layers/attention/cutedsl_mla_backend.py +++ b/python/sglang/srt/layers/attention/cutedsl_mla_backend.py @@ -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 diff --git a/python/sglang/srt/layers/attention/flashattention_backend.py b/python/sglang/srt/layers/attention/flashattention_backend.py index f3cd3b29f..2c81c5937 100644 --- a/python/sglang/srt/layers/attention/flashattention_backend.py +++ b/python/sglang/srt/layers/attention/flashattention_backend.py @@ -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 ] - if self.use_sliding_window_kv_pool: + # 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,9 +2749,14 @@ 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_() - self.cuda_graph_swa_out_cache_loc[:n].copy_( - self.token_to_kv_pool.translate_loc_from_full_to_swa(out_cache_loc) - ) + 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.kv_index_translator.sliding_window_write_loc_for(out_cache_loc) + ) if forward_mode.is_decode_or_idle(): if spec_info is not None: @@ -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( diff --git a/python/sglang/srt/layers/attention/flashinfer_backend.py b/python/sglang/srt/layers/attention/flashinfer_backend.py index 80ec2bb22..4d89de822 100644 --- a/python/sglang/srt/layers/attention/flashinfer_backend.py +++ b/python/sglang/srt/layers/attention/flashinfer_backend.py @@ -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,14 +855,19 @@ 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_() - self.cuda_graph_swa_out_cache_loc[:n].copy_( - self._swa_kv_pool.translate_loc_from_full_to_swa( - forward_batch.out_cache_loc + 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.kv_index_translator.sliding_window_write_loc_for( + forward_batch.out_cache_loc + ) ) - ) if in_capture: self.forward_metadata.swa_out_cache_loc = ( self.cuda_graph_swa_out_cache_loc[:n] @@ -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] = ( diff --git a/python/sglang/srt/layers/attention/flashinfer_mla_backend.py b/python/sglang/srt/layers/attention/flashinfer_mla_backend.py index 295cd60a0..39f8510f5 100644 --- a/python/sglang/srt/layers/attention/flashinfer_mla_backend.py +++ b/python/sglang/srt/layers/attention/flashinfer_mla_backend.py @@ -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 diff --git a/python/sglang/srt/layers/attention/flashmla_backend.py b/python/sglang/srt/layers/attention/flashmla_backend.py index 92fa2bf29..301a06d79 100644 --- a/python/sglang/srt/layers/attention/flashmla_backend.py +++ b/python/sglang/srt/layers/attention/flashmla_backend.py @@ -162,17 +162,25 @@ 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) - create_flashmla_kv_indices_triton[ - (bs, get_num_kv_index_blocks_flashmla(max_seqlen_pad, PAGE_SIZE)) - ]( - self.req_to_token, - forward_batch.req_pool_indices, - forward_batch.seq_lens, - None, - block_kv_indices, - self.req_to_token.stride(0), - block_kv_indices.stride(0), - ) + 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)) + ]( + self.req_to_token, + forward_batch.req_pool_indices, + forward_batch.seq_lens, + None, + block_kv_indices, + self.req_to_token.stride(0), + block_kv_indices.stride(0), + ) mla_metadata, num_splits = get_mla_metadata( forward_batch.seq_lens.to(torch.int32), self.num_q_heads, @@ -328,22 +336,30 @@ class FlashMLABackend(FlashInferMLAAttnBackend): else: max_seqlen_pad = self.cuda_graph_kv_indices.shape[1] - create_flashmla_kv_indices_triton[ - ( - bs, - get_num_kv_index_blocks_flashmla( - self.cuda_graph_kv_indices.stride(0), PAGE_SIZE - ), + 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, + get_num_kv_index_blocks_flashmla( + self.cuda_graph_kv_indices.stride(0), PAGE_SIZE + ), + ) + ]( + self.req_to_token, + req_pool_indices[:bs], + seq_lens, + None, + self.cuda_graph_kv_indices, + self.req_to_token.stride(0), + self.cuda_graph_kv_indices.stride(0), ) - ]( - self.req_to_token, - req_pool_indices[:bs], - seq_lens, - None, - self.cuda_graph_kv_indices, - self.req_to_token.stride(0), - self.cuda_graph_kv_indices.stride(0), - ) q_head_mult = ( self.num_draft_tokens diff --git a/python/sglang/srt/layers/attention/tbo_backend.py b/python/sglang/srt/layers/attention/tbo_backend.py index d8713e9f5..77615cf50 100644 --- a/python/sglang/srt/layers/attention/tbo_backend.py +++ b/python/sglang/srt/layers/attention/tbo_backend.py @@ -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 ) diff --git a/python/sglang/srt/layers/attention/trtllm_mha_backend.py b/python/sglang/srt/layers/attention/trtllm_mha_backend.py index e60083627..681b29b80 100644 --- a/python/sglang/srt/layers/attention/trtllm_mha_backend.py +++ b/python/sglang/srt/layers/attention/trtllm_mha_backend.py @@ -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,20 +1077,27 @@ class TRTLLMHAAttnBackend(FlashInferAttnBackend): else: metadata.cu_seqlens_q = metadata.cu_seqlens_k - 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 - ) - metadata.swa_page_table = ( - torch.empty( + 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 ) - if has_swa - else None - ) - self._fill_page_table_device( - metadata, forward_batch.req_pool_indices, metadata.cache_seqlens_int32 - ) + metadata.swa_page_table = ( + torch.empty( + (batch_size, self.max_num_pages), dtype=torch.int32, device=device + ) + if has_swa + else None + ) + self._fill_page_table_device( + metadata, forward_batch.req_pool_indices, metadata.cache_seqlens_int32 + ) self._maybe_build_cp_zigzag_page_tables(metadata, forward_batch) if self._needs_encoder_only_expand(forward_batch.forward_mode, metadata): @@ -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 ) ) diff --git a/python/sglang/srt/layers/attention/trtllm_mla_backend.py b/python/sglang/srt/layers/attention/trtllm_mla_backend.py index 24820d17c..4cac93462 100755 --- a/python/sglang/srt/layers/attention/trtllm_mla_backend.py +++ b/python/sglang/srt/layers/attention/trtllm_mla_backend.py @@ -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,23 +358,28 @@ class TRTLLMMLABackend(FlashInferMLAAttnBackend): (batch_size, max_blocks), -1, dtype=torch.int32, device=device ) - create_flashmla_kv_indices_triton[ - ( - batch_size, - get_num_kv_index_blocks_flashmla(max_blocks, self.page_size), + 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, + get_num_kv_index_blocks_flashmla(max_blocks, self.page_size), + ) + ]( + self.req_to_token, + req_pool_indices, + seq_lens, + None, + block_kv_indices, + self.req_to_token.stride(0), + max_blocks, + PAGED_SIZE=self.page_size, ) - ]( - self.req_to_token, - req_pool_indices, - seq_lens, - None, - block_kv_indices, - 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,25 +540,30 @@ class TRTLLMMLABackend(FlashInferMLAAttnBackend): metadata.seq_lens_k.copy_(seq_lens[:bs]) # Update block indices for new sequences. - create_flashmla_kv_indices_triton[ - ( - bs, - get_num_kv_index_blocks_flashmla( - metadata.block_kv_indices.shape[1], self.page_size - ), + 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, + get_num_kv_index_blocks_flashmla( + metadata.block_kv_indices.shape[1], self.page_size + ), + ) + ]( + self.req_to_token, + req_pool_indices[:bs], + seq_lens, + None, + metadata.block_kv_indices, + self.req_to_token.stride(0), + metadata.block_kv_indices.shape[1], + PAGED_SIZE=self.page_size, ) - ]( - self.req_to_token, - req_pool_indices[:bs], - seq_lens, - None, - metadata.block_kv_indices, - 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: """Get the fill value for sequence lengths in CUDA graph.""" @@ -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( diff --git a/python/sglang/srt/layers/attention/unified_mem_hooks.py b/python/sglang/srt/layers/attention/unified_mem_hooks.py deleted file mode 100644 index 6839926f6..000000000 --- a/python/sglang/srt/layers/attention/unified_mem_hooks.py +++ /dev/null @@ -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, - ) diff --git a/python/sglang/srt/mem_cache/memory_pool.py b/python/sglang/srt/mem_cache/memory_pool.py index 3f809ca84..1cbce3c17 100644 --- a/python/sglang/srt/mem_cache/memory_pool.py +++ b/python/sglang/srt/mem_cache/memory_pool.py @@ -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) diff --git a/python/sglang/srt/mem_cache/unified_memory_pool.py b/python/sglang/srt/mem_cache/unified_memory_pool.py index 72daecffc..fb84d9263 100644 --- a/python/sglang/srt/mem_cache/unified_memory_pool.py +++ b/python/sglang/srt/mem_cache/unified_memory_pool.py @@ -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, diff --git a/python/sglang/srt/model_executor/forward_batch_deepseek_mha_mixin.py b/python/sglang/srt/model_executor/forward_batch_deepseek_mha_mixin.py index 4ee470352..0a623bb26 100644 --- a/python/sglang/srt/model_executor/forward_batch_deepseek_mha_mixin.py +++ b/python/sglang/srt/model_executor/forward_batch_deepseek_mha_mixin.py @@ -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 diff --git a/python/sglang/srt/server_args.py b/python/sglang/srt/server_args.py index b83057c04..263bf3d35 100644 --- a/python/sglang/srt/server_args.py +++ b/python/sglang/srt/server_args.py @@ -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. diff --git a/test/registered/attention/test_normal_decode_set_metadata.py b/test/registered/attention/test_normal_decode_set_metadata.py index b5b28fa1b..d9d4f927e 100644 --- a/test/registered/attention/test_normal_decode_set_metadata.py +++ b/test/registered/attention/test_normal_decode_set_metadata.py @@ -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, diff --git a/test/registered/attention/test_trtllm_mha_graph_metadata.py b/test/registered/attention/test_trtllm_mha_graph_metadata.py index bad6dd76a..6d46cbe17 100644 --- a/test/registered/attention/test_trtllm_mha_graph_metadata.py +++ b/test/registered/attention/test_trtllm_mha_graph_metadata.py @@ -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") diff --git a/test/registered/models_e2e/test_kimi_linear_unified_memory.py b/test/registered/models_e2e/test_kimi_linear_unified_memory.py index 45320227a..d77854d28 100644 --- a/test/registered/models_e2e/test_kimi_linear_unified_memory.py +++ b/test/registered/models_e2e/test_kimi_linear_unified_memory.py @@ -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() diff --git a/test/registered/page_major/test_page_major_gpt_oss.py b/test/registered/page_major/test_page_major_gpt_oss.py index 5b7ae0ebb..1c60f48f1 100644 --- a/test/registered/page_major/test_page_major_gpt_oss.py +++ b/test/registered/page_major/test_page_major_gpt_oss.py @@ -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() diff --git a/test/registered/page_major/test_page_major_qwen_hybrid.py b/test/registered/page_major/test_page_major_qwen_hybrid.py index 1bd98f82e..d0096b0e4 100644 --- a/test/registered/page_major/test_page_major_qwen_hybrid.py +++ b/test/registered/page_major/test_page_major_qwen_hybrid.py @@ -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() diff --git a/test/registered/unit/layers/attention/test_flashattention_graph_metadata.py b/test/registered/unit/layers/attention/test_flashattention_graph_metadata.py index e44db9903..93e2a679d 100644 --- a/test/registered/unit/layers/attention/test_flashattention_graph_metadata.py +++ b/test/registered/unit/layers/attention/test_flashattention_graph_metadata.py @@ -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 diff --git a/test/registered/unit/layers/attention/test_flashattention_pa_swa_prefill_lens_size.py b/test/registered/unit/layers/attention/test_flashattention_pa_swa_prefill_lens_size.py index f320700ec..f00bda40e 100644 --- a/test/registered/unit/layers/attention/test_flashattention_pa_swa_prefill_lens_size.py +++ b/test/registered/unit/layers/attention/test_flashattention_pa_swa_prefill_lens_size.py @@ -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, diff --git a/test/registered/unit/layers/attention/test_kv_translate_ownership.py b/test/registered/unit/layers/attention/test_kv_translate_ownership.py new file mode 100644 index 000000000..d81b55af0 --- /dev/null +++ b/test/registered/unit/layers/attention/test_kv_translate_ownership.py @@ -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() diff --git a/test/registered/unit/mem_cache/test_full_loc_fast_path.py b/test/registered/unit/mem_cache/test_full_loc_fast_path.py index 2a3a47d0e..d04ab4875 100644 --- a/test/registered/unit/mem_cache/test_full_loc_fast_path.py +++ b/test/registered/unit/mem_cache/test_full_loc_fast_path.py @@ -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,21 +98,29 @@ 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), - torch.zeros(3, 4, 8), - torch.zeros(3, 4, 8), - ) + pool.set_kv_buffer( + layer, + _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() @@ -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__": diff --git a/test/registered/unit/mem_cache/test_multi_ended_allocator.py b/test/registered/unit/mem_cache/test_multi_ended_allocator.py index e1a02127b..f927892b8 100644 --- a/test/registered/unit/mem_cache/test_multi_ended_allocator.py +++ b/test/registered/unit/mem_cache/test_multi_ended_allocator.py @@ -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() diff --git a/test/registered/unit/mem_cache/test_unified_mla_block_table.py b/test/registered/unit/mem_cache/test_unified_mla_block_table.py index e4c582fe7..04a2a4bda 100644 --- a/test/registered/unit/mem_cache/test_unified_mla_block_table.py +++ b/test/registered/unit/mem_cache/test_unified_mla_block_table.py @@ -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 - 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, - ) + 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, + max_seq_pages, + sl.to(torch.int64), + 0, + page_size, + 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__": diff --git a/test/registered/unit/model_executor/test_unified_out_cache_loc_rebind.py b/test/registered/unit/model_executor/test_unified_out_cache_loc_rebind.py index e11706057..3d0c19f37 100644 --- a/test/registered/unit/model_executor/test_unified_out_cache_loc_rebind.py +++ b/test/registered/unit/model_executor/test_unified_out_cache_loc_rebind.py @@ -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() diff --git a/test/registered/unit/server_args/test_page_major_backend_allowlist.py b/test/registered/unit/server_args/test_page_major_backend_allowlist.py index 3c1752017..d1f3e291a 100644 --- a/test/registered/unit/server_args/test_page_major_backend_allowlist.py +++ b/test/registered/unit/server_args/test_page_major_backend_allowlist.py @@ -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):