From f8cbf000f4a5bfd86d3fb7c1e2d6c8fb12339d0e Mon Sep 17 00:00:00 2001 From: Xinyi Song Date: Wed, 2 Sep 2026 09:45:08 -0700 Subject: [PATCH] [AMD] Enable FP4 indexer for Deepseek V4 (#37353) Co-authored-by: 1am9trash <1am9trash@gmail.com> Co-authored-by: AMD-yanfeiwang <256076023+AMD-yanfeiwang@users.noreply.github.com> Co-authored-by: Thomas Wang --- docker/rocm.Dockerfile | 3 + .../kernels/ops/attention/dsv4/compress.py | 26 +- .../ops/attention/dsv4/fp4_indexer_hip.py | 400 +++++++++ python/sglang/srt/arg_groups/serving_hook.py | 7 +- .../deepseek_v4_backend_hip_radix.py | 144 +++- .../srt/layers/attention/dsv4/compressor.py | 35 +- .../layers/attention/dsv4/compressor_v2.py | 39 +- .../srt/layers/attention/dsv4/indexer.py | 83 +- .../srt/layers/attention/dsv4/metadata.py | 4 +- .../srt/mem_cache/deepseek_v4_memory_pool.py | 72 +- .../sglang/srt/mem_cache/hicache_storage.py | 3 + .../hybrid_cache/hybrid_pool_assembler.py | 115 ++- .../hybrid_cache/linker_pool_assembler.py | 16 +- .../sglang/srt/mem_cache/memory_pool_host.py | 18 + .../storage/mooncake_store/mooncake_store.py | 1 + .../srt/model_executor/pool_configurator.py | 17 +- python/sglang/srt/models/deepseek_v4.py | 9 + rust/sglang-radix-tree/src/python_bindings.rs | 2 + .../src/unified_tree_core.rs | 1 + .../ops/attention/test_fp4_indexer_hip.py | 776 ++++++++++++++++++ .../unit/mem_cache/test_mem_pool_host.py | 31 + 21 files changed, 1727 insertions(+), 75 deletions(-) create mode 100644 python/sglang/kernels/ops/attention/dsv4/fp4_indexer_hip.py create mode 100644 test/registered/kernels/ops/attention/test_fp4_indexer_hip.py diff --git a/docker/rocm.Dockerfile b/docker/rocm.Dockerfile index 872eb564c..da4d68da0 100644 --- a/docker/rocm.Dockerfile +++ b/docker/rocm.Dockerfile @@ -569,10 +569,13 @@ RUN pip uninstall -y aiter # produced by a fresh `git clone` above, so there are no real user changes to # preserve. # cherry pick 8578af1 commit for v4 fp4 indexer kv-cache fix, may be removed in next aiter upgrade +# apply fix for v4 fp4 indexer, may be removed in next aiter upgrade RUN git clone ${AITER_REPO} \ && cd aiter \ && git checkout -f ${AITER_COMMIT} \ && git cherry-pick --no-commit 8578af153f4fa1e007fede7e3c1e1b373f07af4c \ + && sed -i 's/from functools import lru_cache/from functools import cache/' aiter/ops/flydsl/kernels/mqa_logits/pa_mqa_logits_fp4_prefill.py \ + && sed -i 's/@lru_cache(maxsize=32)/@cache/' aiter/ops/flydsl/kernels/mqa_logits/pa_mqa_logits_fp4_prefill.py \ && git submodule update --init --recursive \ && pip install -r requirements.txt \ && if [ "${GPU_ARCH_LIST}" = "gfx1250" ]; then \ diff --git a/python/sglang/kernels/ops/attention/dsv4/compress.py b/python/sglang/kernels/ops/attention/dsv4/compress.py index 99b9597c9..a9e915d67 100644 --- a/python/sglang/kernels/ops/attention/dsv4/compress.py +++ b/python/sglang/kernels/ops/attention/dsv4/compress.py @@ -1,6 +1,6 @@ from __future__ import annotations -from typing import TYPE_CHECKING, Literal, NamedTuple, Optional, Union +from typing import TYPE_CHECKING, Literal, NamedTuple, Optional, Union, cast import torch @@ -430,9 +430,33 @@ def compress_norm_rope_store( page_size: int, use_fp4: bool = False, bf16_store: bool = False, + # HIP FP4 uses split scale storage and precomputed BF16 RoPE tables. + kvcache_scale: Optional[torch.Tensor] = None, + rope_cache: Optional[tuple[torch.Tensor, torch.Tensor]] = None, + fp4_k_write_metadata=None, ) -> None: if use_fp4: assert kv.shape[-1] == 128 + if is_hip() and use_fp4: + from sglang.kernels.ops.attention.dsv4.fp4_indexer_hip import ( + aiter_k_indexer_fp4_cache_write, + ) + + cos, sin = cast(tuple[torch.Tensor, torch.Tensor], rope_cache) + aiter_k_indexer_fp4_cache_write( + k=kv, + norm_weight=norm_weight, + norm_epsilon=norm_eps, + cos=cos, + sin=sin, + plan=plan, + out_loc=out_loc, + k_payload=kvcache, + k_scale=cast(torch.Tensor, kvcache_scale), + write_metadata=fp4_k_write_metadata, + ) + return + freq_cis = torch.view_as_real(freq_cis).flatten(-2) if _is_xpu: compress_norm_rope_store_xpu( diff --git a/python/sglang/kernels/ops/attention/dsv4/fp4_indexer_hip.py b/python/sglang/kernels/ops/attention/dsv4/fp4_indexer_hip.py new file mode 100644 index 000000000..5380f104e --- /dev/null +++ b/python/sglang/kernels/ops/attention/dsv4/fp4_indexer_hip.py @@ -0,0 +1,400 @@ +"""AITER adapters for the DeepSeek-V4 FP4 indexer on HIP.""" + +from __future__ import annotations + +from typing import TYPE_CHECKING, NamedTuple, Optional, Tuple, Union + +import torch + +if TYPE_CHECKING: + from sglang.kernels.ops.attention.dsv4.compress import ( + CompressorDecodePlan, + CompressorPrefillPlan, + ) + + +_HEADS = 64 +_HEAD_DIM = 128 +_ROPE_DIM = 64 +_GROUP_SIZE = 32 +_KV_BLOCK_SIZE = 64 +_Q_SCALE_SHAPE = (1, 4, 16, 4) +# gfx950 has 256 CUs; target four persistent CTAs per CU. +_DECODE_BASE_CTA_TARGET = 1024 +# Preserve per-query parallelism when the batch itself exceeds one CTA per CU. +_DECODE_CTAS_PER_QUERY = 4 +_PREFILL_BASE_CTA_TARGET = 1024 +# AITER varctx cta_info row: [batch_packed, chunk_start, chunk_count, ctx_len]. +_DECODE_CTA_INFO_WIDTH = 4 + + +class FP4DecodeWorkspace(NamedTuple): + guarded_page_table: torch.Tensor + c4_seq_lens: torch.Tensor + cta_info: torch.Tensor + cta_count: int + max_seq_len: int + # Held only so AITER's schedule scratch never returns to the graph memory + # pool: the captured builder writes it again on every replay. + schedule_scratch: torch.Tensor + + +class FP4PrefillWorkspace(NamedTuple): + guarded_page_table: torch.Tensor + row_to_batch: torch.Tensor + local_starts: torch.Tensor + cta_info: torch.Tensor + cta_count: int + max_seq_len: int + + +class FP4KWriteMetadata(NamedTuple): + positions: torch.Tensor + slots: torch.Tensor + + +def aiter_q_indexer_fp4( + q: torch.Tensor, + cos: torch.Tensor, + sin: torch.Tensor, + positions: torch.Tensor, +) -> Tuple[torch.Tensor, torch.Tensor]: + """Apply RoPE and Hadamard rotation, then quantize indexer Q to FP4.""" + import aiter + + num_tokens = q.shape[0] + # AITER asserts int64 positions; the caller normally widens once per forward. + if ( + positions.dtype is not torch.int64 + or positions.device != q.device + or not positions.is_contiguous() + ): + positions = positions.to(device=q.device, dtype=torch.int64).contiguous() + q_fp4 = torch.empty( + (num_tokens, _HEADS, _HEAD_DIM // 2), + dtype=aiter.dtypes.fp4x2, + device=q.device, + ) + q_scale = torch.empty( + (num_tokens, *_Q_SCALE_SHAPE), dtype=torch.uint8, device=q.device + ) + aiter.rope_rotate_activation( + q_fp4, + q, + cos, + sin, + positions, + rope_dim=_ROPE_DIM, + out_scale=q_scale, + group_size=_GROUP_SIZE, + shuffle_scale=True, + do_rotate_act=True, + ) + return q_fp4, q_scale + + +def _as_int32_1d(t: torch.Tensor) -> torch.Tensor: + """Normalize a length vector without dispatching when it already matches. + + Called once per C4 layer, so the no-op fast path matters: the metadata + builder already hands us a 1-D contiguous int32 tensor. + """ + if t.dim() == 1 and t.dtype is torch.int32 and t.is_contiguous(): + return t + return t.reshape(-1).to(torch.int32).contiguous() + + +def _decode_cta_count(num_queries: int, max_seq_len: int) -> int: + """Choose a bounded persistent grid without exceeding available KV chunks.""" + chunks_per_seq = max(1, (max_seq_len + 255) // 256) + available_ctas = num_queries * chunks_per_seq + target_ctas = max(_DECODE_BASE_CTA_TARGET, num_queries * _DECODE_CTAS_PER_QUERY) + return min(available_ctas, target_ctas) + + +def _guard_page_table(page_table: torch.Tensor, out: Optional[torch.Tensor] = None): + """Pad page tables for 256-token scheduling and one-chunk lookahead.""" + page_table = page_table.to(dtype=torch.int32).contiguous() + rows, logical_width = page_table.shape + padded_width = max(4, (logical_width + 3) // 4 * 4) + if out is None: + out = page_table.new_zeros((rows, padded_width + 4)) + else: + assert out.shape == (rows, padded_width + 4), f"{out.shape=} {rows=}" + out[:, :logical_width].copy_(page_table) + return out, padded_width * _KV_BLOCK_SIZE + + +def prepare_fp4_decode_workspace( + page_table: torch.Tensor, + c4_seq_lens: torch.Tensor, +) -> FP4DecodeWorkspace: + """Build the decode page-table, schedule, and logits buffers. + + Safe to run under CUDA-graph capture: every tensor the captured schedule + kernel touches is reachable from the returned workspace, so none of it can + be handed out again by a later capture sharing the graph memory pool. + """ + from aiter.ops.flydsl.kernels.mqa_logits.pa_mqa_logits_fp4 import ( + compute_varctx_schedule, + ) + + guarded, max_seq_len = _guard_page_table(page_table) + c4_seq_lens = _as_int32_1d(c4_seq_lens) + num_queries = guarded.shape[0] + cta_count = _decode_cta_count(num_queries, max_seq_len) + cta_info = torch.empty( + (cta_count, _DECODE_CTA_INFO_WIDTH), + dtype=torch.int32, + device=guarded.device, + ) + schedule_scratch, _, _ = compute_varctx_schedule( + c4_seq_lens, + block_k=256, + parallel_unit_num=cta_count, + max_seq_len=max_seq_len, + next_n=1, + cta_info_out=cta_info, + ) + return FP4DecodeWorkspace( + guarded, c4_seq_lens, cta_info, cta_count, max_seq_len, schedule_scratch + ) + + +def prepare_fp4_prefill_workspace( + page_table: torch.Tensor, + c4_seq_lens: torch.Tensor, + workspace: Optional[FP4PrefillWorkspace] = None, +) -> FP4PrefillWorkspace: + """Build or refresh the prefill page-table, schedule, and logits buffers. + + Must run OUTSIDE CUDA-graph capture. AITER's prefill scheduler frees its own + scratch when it returns, and its schedule kernel reads that scratch, so a + captured build would replay against recycled graph-pool memory. Callers + instead refresh this workspace per step and let the graph read only the + pinned ``cta_info`` / ``logits`` / page-table buffers. + """ + from aiter.ops.flydsl.kernels.mqa_logits.pa_mqa_logits_fp4_prefill import ( + CTA_INFO_WIDTH, + compute_prefill_schedule, + ) + + c4_seq_lens = _as_int32_1d(c4_seq_lens) + if workspace is None: + guarded, max_seq_len = _guard_page_table(page_table) + num_queries = guarded.shape[0] + cta_count = max(_PREFILL_BASE_CTA_TARGET, num_queries) + workspace = FP4PrefillWorkspace( + guarded_page_table=guarded, + row_to_batch=torch.arange( + num_queries, device=guarded.device, dtype=torch.int32 + ), + local_starts=torch.zeros( + num_queries, device=guarded.device, dtype=torch.int32 + ), + cta_info=torch.empty( + (cta_count, CTA_INFO_WIDTH), dtype=torch.int32, device=guarded.device + ), + cta_count=cta_count, + max_seq_len=max_seq_len, + ) + else: + _guard_page_table(page_table, out=workspace.guarded_page_table) + + assert c4_seq_lens.shape[0] == workspace.row_to_batch.shape[0], ( + f"c4_seq_lens rows {c4_seq_lens.shape[0]} do not match the workspace's " + f"{workspace.row_to_batch.shape[0]}; the schedule kernel indexes both by row" + ) + compute_prefill_schedule( + workspace.row_to_batch, + workspace.local_starts, + c4_seq_lens, + block_k=256, + parallel_unit_num=workspace.cta_count, + max_seq_len=workspace.max_seq_len, + cta_info_out=workspace.cta_info, + ) + return workspace + + +def aiter_fp4_paged_mqa_logits( + *, + q_fp4: torch.Tensor, + q_scale: torch.Tensor, + k_payload: torch.Tensor, + k_scale: torch.Tensor, + weights: torch.Tensor, + page_table: torch.Tensor, + c4_seq_lens: torch.Tensor, + weight_scale: float, + is_decode: bool, + decode_workspace: Optional[FP4DecodeWorkspace] = None, + prefill_workspace: Optional[FP4PrefillWorkspace] = None, +) -> torch.Tensor: + """Compute FP4 Q/K indexer logits with the decode or prefill FlyDSL kernel.""" + from aiter.ops.flydsl import ( + flydsl_pa_mqa_logits_fp4, + flydsl_pa_mqa_logits_fp4_prefill, + ) + + num_tokens = q_fp4.shape[0] + c4_seq_lens = _as_int32_1d(c4_seq_lens) + workspace = decode_workspace if is_decode else prefill_workspace + # A workspace is bound to one row count. DP padding or truncated activations + # can leave it stale, in which case fall back to building the schedule here. + if workspace is not None and workspace.guarded_page_table.shape[0] != num_tokens: + workspace = None + if workspace is not None: + page_table = workspace.guarded_page_table + max_seq_len = workspace.max_seq_len + else: + page_table, max_seq_len = _guard_page_table(page_table) + q_payload = q_fp4.view(torch.uint8) + k_payload = k_payload.view(torch.uint8) + # Scored write-once and freed with this call. Recycling it through the + # allocator costs nothing because a pinned cta_info makes the kernel skip + # its -inf pre-fill and the length-aware top-k reads only [0, c4_seq_len). + logits = torch.empty( + (num_tokens, max_seq_len), dtype=torch.float32, device=q_fp4.device + ) + common = { + "weight_scale": weight_scale, + "block_k": 256, + "kv_block_size": _KV_BLOCK_SIZE, + "num_warps": 4, + "out": logits, + } + + if is_decode: + pinned = ( + {} + if workspace is None + else { + "cta_info": workspace.cta_info, + "total_ctas": workspace.cta_count, + } + ) + logits = flydsl_pa_mqa_logits_fp4( + q_payload.reshape(num_tokens, 1, _HEADS, _HEAD_DIM // 2), + q_scale.reshape(num_tokens, 1, *_Q_SCALE_SHAPE), + k_payload, + k_scale, + page_table, + weights, + c4_seq_lens, + max_seq_len, + next_n=1, + parallel_unit_num=None, + **pinned, + **common, + ) + else: + if workspace is None: + pinned = {} + row_to_batch = torch.arange( + num_tokens, device=q_fp4.device, dtype=torch.int32 + ) + local_starts = torch.zeros( + num_tokens, device=q_fp4.device, dtype=torch.int32 + ) + else: + pinned = { + "cta_info": workspace.cta_info, + "n_ctas": workspace.cta_count, + } + row_to_batch = workspace.row_to_batch + local_starts = workspace.local_starts + logits = flydsl_pa_mqa_logits_fp4_prefill( + q_payload, + q_scale, + k_payload, + k_scale, + page_table, + weights, + row_to_batch, + local_starts, + c4_seq_lens, + max_seq_len, + parallel_unit_num=max(_PREFILL_BASE_CTA_TARGET, num_tokens), + **pinned, + **common, + ) + + return logits + + +def prepare_fp4_k_write_metadata( + plan: Union[CompressorDecodePlan, CompressorPrefillPlan], + out_loc: torch.Tensor, + rope_table_len: int, +) -> FP4KWriteMetadata: + """ + Build RoPE positions and cache slots from a compressor plan. + """ + plan_words = plan[1].view(torch.int32) + seq_lens = plan_words[:, 0].to(torch.int64) + positions = seq_lens - plan.compress_ratio + valid = (positions >= 0) & (positions < rope_table_len) + positions = torch.where(valid, positions, torch.zeros_like(positions)) + valid &= seq_lens % plan.compress_ratio == 0 + + out_loc = out_loc.to(dtype=torch.int64) + if plan.is_decode: + slots = out_loc + elif out_loc.shape[0] == 0: + slots = torch.full_like(seq_lens, -1) + valid.zero_() + else: + ragged_ids = plan_words[:, 1].bitwise_and(0xFFFF).to(torch.int64) + valid &= ragged_ids < out_loc.shape[0] + slots = out_loc[ragged_ids.clamp(max=out_loc.shape[0] - 1)] + slots = torch.where(valid, slots, torch.full_like(slots, -1)) + return FP4KWriteMetadata(positions.contiguous(), slots.contiguous()) + + +def aiter_k_indexer_fp4_cache_write( + *, + k: torch.Tensor, + norm_weight: torch.Tensor, + norm_epsilon: float, + cos: torch.Tensor, + sin: torch.Tensor, + plan: Union[CompressorDecodePlan, CompressorPrefillPlan], + out_loc: torch.Tensor, + k_payload: torch.Tensor, + k_scale: torch.Tensor, + write_metadata: Optional[FP4KWriteMetadata] = None, +) -> None: + """ + Map compressed K rows to cache slots and run the fused AITER FP4 writer. + """ + num_rows = k.shape[0] + if num_rows == 0: + return + + assert write_metadata is not None, "FP4 K-write metadata is missing." + + positions, slots = write_metadata + # The compressor normally hands over its BF16 mirror; convert only when some + # caller still passes the FP32 parameter. + if norm_weight.dtype is not torch.bfloat16 or norm_weight.device != k.device: + norm_weight = norm_weight.to(device=k.device, dtype=torch.bfloat16).contiguous() + + import aiter + + aiter.rmsnorm_rope_rotate_activation_fp4quant_kvcache( + k_payload, + k_scale, + k.view(num_rows, 1, _HEAD_DIM), + norm_weight, + cos, + sin, + positions, + slots, + norm_epsilon, + rope_dim=_ROPE_DIM, + kv_block_size=_KV_BLOCK_SIZE, + group_size=_GROUP_SIZE, + shuffle_scale=True, + do_rotate_act=True, + ) diff --git a/python/sglang/srt/arg_groups/serving_hook.py b/python/sglang/srt/arg_groups/serving_hook.py index fe7713124..39c3f63ca 100644 --- a/python/sglang/srt/arg_groups/serving_hook.py +++ b/python/sglang/srt/arg_groups/serving_hook.py @@ -22,6 +22,7 @@ from sglang.srt.runtime_context import get_platform from sglang.srt.utils.common import ( configure_media_url_security, get_device, + is_gfx95_supported, is_mnnvl_fabric_device, ) from sglang.utils import is_in_ci @@ -419,11 +420,11 @@ def handle_environment_variables(server_args: Any): "All operations will run eagerly through the graph capture/replay path." ) if cfg.enable_deepseek_v4_fp4_indexer and not ( - get_platform().is_sm100 or get_platform().is_sm120 + get_platform().is_sm100 or get_platform().is_sm120 or is_gfx95_supported() ): raise ValueError( - "--enable-deepseek-v4-fp4-indexer requires SM100 or SM120 GPUs with " - "DeepGEMM FP4 indexer support." + "--enable-deepseek-v4-fp4-indexer requires SM100, SM120, or gfx95 GPUs " + "with FP4 indexer support." ) # FP8 W_o GEMM needs DeepGEMM JIT. Enable exactly where the runtime can run # it, mirroring the forward scale split: the ue8m0 path diff --git a/python/sglang/srt/layers/attention/deepseek_v4_backend_hip_radix.py b/python/sglang/srt/layers/attention/deepseek_v4_backend_hip_radix.py index cabc08389..92c39e57e 100644 --- a/python/sglang/srt/layers/attention/deepseek_v4_backend_hip_radix.py +++ b/python/sglang/srt/layers/attention/deepseek_v4_backend_hip_radix.py @@ -47,6 +47,11 @@ from sglang.srt.utils import ceil_align if TYPE_CHECKING: from sgl_kernel.flash_mla import FlashMLASchedMeta + from sglang.kernels.ops.attention.dsv4.fp4_indexer_hip import ( + FP4DecodeWorkspace, + FP4KWriteMetadata, + FP4PrefillWorkspace, + ) from sglang.srt.layers.radix_attention import RadixAttention from sglang.srt.model_executor.model_runner import ModelRunner @@ -353,6 +358,18 @@ class DSV4Metadata: c4_compress_metadata: Optional[FusedCompressMetadata] = None c128_compress_metadata: Optional[FusedCompressMetadata] = None + # FP4 indexer buffers that captured kernels bind by address. Deliberately + # absent from copy_: the addresses must stay pinned across replays, and the + # workspace builders refresh their contents instead. + fp4_decode_workspace: Optional[FP4DecodeWorkspace] = field(default=None, repr=False) + fp4_prefill_workspace: Optional[FP4PrefillWorkspace] = field( + default=None, repr=False + ) + # Derived by the first C4 layer of a forward and reused by the rest. + fp4_k_write_metadata: Optional[FP4KWriteMetadata] = field(default=None, repr=False) + # AITER's rope kernels require int64 positions while the core metadata keeps + # them int32, so widen once per forward instead of once per C4 layer. + fp4_q_positions: Optional[torch.Tensor] = field(default=None, repr=False) @property def core_metadata(self) -> DSV4AttnMetadata: @@ -655,6 +672,8 @@ class DeepseekV4HipRadixBackend( ragged_layout=None, ) -> DSV4Metadata: batch_size = len(seq_lens) + # Verify tokens may cross into the next page beyond the accepted prefix. + max_seq_len += self.target_verify_num_draft_tokens extend_start_loc = None if ragged_layout is not None: verify_lens_dev = ragged_layout.verify_lens.to( @@ -760,6 +779,10 @@ class DeepseekV4HipRadixBackend( req_pool_indices = raw_metadata.req_pool_indices seq_lens = raw_metadata.seq_lens out_cache_loc = raw_metadata.out_cache_loc + if self.topk > 0 and self.speculative_num_steps > 1: + # Each EAGLE draft step appends one token while ForwardBatch keeps + # the accepted-prefix lengths unchanged across the captured loop. + seq_lens = seq_lens + self.speculative_step_id + 1 core_attn_metadata = self.make_core_attn_metadata( req_to_token=self.req_to_token, @@ -860,6 +883,83 @@ class DeepseekV4HipRadixBackend( ) ) + if self._fp4_workspaces_enabled(metadata): + from sglang.kernels.ops.attention.dsv4.fp4_indexer_hip import ( + prepare_fp4_k_write_metadata, + ) + + metadata.fp4_k_write_metadata = prepare_fp4_k_write_metadata( + metadata.c4_compress_metadata, + metadata.core_attn_metadata.c4_out_loc, + self.MAX_SEQ_LEN_FOR_CAPTURE, + ) + metadata.fp4_q_positions = metadata.core_attn_metadata.positions.to( + torch.int64 + ) + + # Decode's schedule builder is capture-safe because the workspace pins + # the scratch it reads, so it can stay next to the metadata it consumes. + # Prefill/target-verify cannot; see _refresh_fp4_prefill_workspace. + if self._fp4_workspaces_enabled(metadata) and ( + forward_batch.forward_mode.is_decode() + ): + from sglang.kernels.ops.attention.dsv4.fp4_indexer_hip import ( + prepare_fp4_decode_workspace, + ) + + indexer_metadata = metadata.indexer_metadata + metadata.fp4_decode_workspace = prepare_fp4_decode_workspace( + indexer_metadata.page_table, + indexer_metadata.c4_seq_lens, + ) + + def _fp4_workspaces_enabled(self, metadata) -> bool: + return ( + self.enable_deepseek_v4_fp4_indexer + # Draft-step backends drive the NextN layer, which is built with + # compress_ratio_override=0 and so owns no C4 indexer. Their + # workspaces would be built, scheduled, and never read. + and self.speculative_num_steps == 0 + and isinstance(metadata, DSV4Metadata) + and metadata.indexer_metadata is not None + and metadata.c4_compress_metadata is not None + ) + + def _refresh_fp4_prefill_workspace(self, forward_batch: ForwardBatch) -> None: + """Rebuild the FP4 prefill schedule outside CUDA-graph capture. + + AITER's prefill scheduler frees the scratch that its schedule kernel + reads, so recording the build into a graph would leave every replay + reading recycled graph-pool memory. Only the pinned buffers it fills + (cta_info / logits / guarded page table) may be read from the graph. + """ + metadata = self.forward_metadata + if not self._fp4_workspaces_enabled(metadata): + return + if forward_batch.forward_mode not in ( + ForwardMode.EXTEND, + ForwardMode.MIXED, + ForwardMode.TARGET_VERIFY, + ): + return + if ( + get_parallel().attn_cp_size != 1 + or getattr(forward_batch, "tbo_children", None) + or getattr(forward_batch, "tbo_parent_token_range", None) is not None + ): + return + + from sglang.kernels.ops.attention.dsv4.fp4_indexer_hip import ( + prepare_fp4_prefill_workspace, + ) + + indexer_metadata = metadata.indexer_metadata + metadata.fp4_prefill_workspace = prepare_fp4_prefill_workspace( + indexer_metadata.page_table, + indexer_metadata.c4_seq_lens, + workspace=metadata.fp4_prefill_workspace, + ) + def init_forward_metadata_out_graph( self, forward_batch: ForwardBatch, @@ -984,6 +1084,7 @@ class DeepseekV4HipRadixBackend( self.replay_cuda_graph_metadata_from( bs=graph_key, temp_metadata=temp_metadata, bucket=bucket ) + self._refresh_fp4_prefill_workspace(forward_batch) if in_capture: metadata = self.forward_metadata @@ -1065,6 +1166,7 @@ class DeepseekV4HipRadixBackend( self.forward_metadata = metadata self.init_forward_metadata_in_graph(forward_batch) + self._refresh_fp4_prefill_workspace(forward_batch) def init_cuda_graph_state(self, max_bs: int, max_num_tokens: int) -> None: self.cuda_graph_metadata_of_bucket_and_bs: Dict[ @@ -1767,14 +1869,32 @@ class DeepseekV4MultiStepBackend(DeepseekV4HipRadixBackend): ): from types import SimpleNamespace + actual_forward_mode = getattr( + forward_batch, "actual_forward_mode", forward_batch.forward_mode + ) + out_cache_loc = getattr(forward_batch, "out_cache_loc", None) + step_out_cache_locs = None + # C4/C128 write locations are baked into each child backend's metadata, + # so every speculative step must consume its own cache-location row. + if ( + actual_forward_mode != ForwardMode.IDLE + and out_cache_loc is not None + and self.topk > 0 + and self.speculative_num_steps > 1 + ): + step_out_cache_locs = per_step_draft_out_cache_loc( + out_cache_loc, + forward_batch.batch_size, + self.topk, + self.speculative_num_steps, + ) + inner_fb = SimpleNamespace( batch_size=forward_batch.batch_size, forward_mode=ForwardMode.DECODE, # Propagate the real runtime mode so inner backends can detect IDLE # and apply their idle substitution. - actual_forward_mode=getattr( - forward_batch, "actual_forward_mode", forward_batch.forward_mode - ), + actual_forward_mode=actual_forward_mode, input_ids=getattr(forward_batch, "input_ids", None), positions=getattr(forward_batch, "positions", None), req_pool_indices=forward_batch.req_pool_indices, @@ -1782,23 +1902,37 @@ class DeepseekV4MultiStepBackend(DeepseekV4HipRadixBackend): seq_lens_sum=forward_batch.seq_lens_sum, seq_lens_cpu=forward_batch.seq_lens_cpu, encoder_lens=None, - out_cache_loc=getattr(forward_batch, "out_cache_loc", None), + out_cache_loc=out_cache_loc, spec_info=forward_batch.spec_info, ) if in_capture: for i in range(self.speculative_num_steps): + if step_out_cache_locs is not None: + inner_fb.out_cache_loc = step_out_cache_locs[i] self.attn_backends[i].init_forward_metadata_out_graph( inner_fb, in_capture=True ) else: if self.speculative_num_steps == 1: return + if step_out_cache_locs is not None: + inner_fb.out_cache_loc = step_out_cache_locs[0] self.attn_backends[0].init_forward_metadata_out_graph(inner_fb) temp_metadata = self.attn_backends[0].forward_metadata + if step_out_cache_locs is not None: + assert isinstance(temp_metadata, DSV4RawDecodeMetadata) for i in range(1, self.speculative_num_steps - 1): + if step_out_cache_locs is None: + step_metadata = temp_metadata + else: + step_metadata = DSV4RawDecodeMetadata( + req_pool_indices=temp_metadata.req_pool_indices, + seq_lens=temp_metadata.seq_lens, + out_cache_loc=step_out_cache_locs[i], + ) self.attn_backends[i].replay_cuda_graph_metadata_from( bs=forward_batch.batch_size, - temp_metadata=temp_metadata, + temp_metadata=step_metadata, bucket=_GraphBucket.DECODE_OR_IDLE, ) diff --git a/python/sglang/srt/layers/attention/dsv4/compressor.py b/python/sglang/srt/layers/attention/dsv4/compressor.py index 31fec7b84..9b213c1ba 100644 --- a/python/sglang/srt/layers/attention/dsv4/compressor.py +++ b/python/sglang/srt/layers/attention/dsv4/compressor.py @@ -32,7 +32,7 @@ from sglang.srt.mem_cache.deepseek_v4_compress_state import ( from sglang.srt.mem_cache.deepseek_v4_memory_pool import DeepSeekV4TokenToKVPool from sglang.srt.model_executor.forward_context import get_attn_backend from sglang.srt.models.deepseek_v2 import _is_hip -from sglang.srt.runtime_context import get_parallel +from sglang.srt.runtime_context import get_exec, get_parallel from sglang.srt.utils import add_prefix, is_npu, set_weight_attrs _is_npu = is_npu() @@ -374,11 +374,44 @@ class Compressor(BaseFusedOp): self.norm = RMSNorm( self.head_dim, eps=config.rms_norm_eps, weight_dtype=torch.float32 ) + if ( + is_in_indexer + and _is_hip + and get_exec().kernel.enable_deepseek_v4_fp4_indexer + ): + self._init_fp4_norm_weight() self.rotary_emb = rotary_emb self.freqs_cis = freqs_cis self.ape_converted = False + def _init_fp4_norm_weight(self) -> None: + """Mirror the FP32 norm weight in BF16 for the AITER FP4 K writer. + + The FP8 path feeds the FP32 weight straight to its kernel; AITER wants + BF16, and converting at the call site costs one copy per C4 layer per + forward. A buffer keeps the conversion out of the forward and survives + module ``_apply``, and the loader below re-derives it so online weight + updates propagate -- they land as ``param.data.copy_``, which leaves the + parameter's identity and ``_version`` untouched and would silently + defeat any cache keyed on those. Same reach as ``load_ape_weight``: + ``update_weights_from_tensor(load_format="direct")`` calls + ``default_weight_loader`` itself and so skips both hooks. + """ + self.norm.register_buffer( + "fp4_weight_bf16", + self.norm.weight.detach().to(torch.bfloat16).contiguous(), + persistent=False, + ) + set_weight_attrs(self.norm.weight, {"weight_loader": self.load_norm_weight}) + + def load_norm_weight( + self, param: torch.Tensor, loaded_weight: torch.Tensor + ) -> None: + assert param is self.norm.weight + param.data.copy_(loaded_weight) + self.norm.fp4_weight_bf16.copy_(param.data) + def _apply_ape_hotfix(self): self.ape_converted = True diff --git a/python/sglang/srt/layers/attention/dsv4/compressor_v2.py b/python/sglang/srt/layers/attention/dsv4/compressor_v2.py index b9268e8c6..707efb74d 100644 --- a/python/sglang/srt/layers/attention/dsv4/compressor_v2.py +++ b/python/sglang/srt/layers/attention/dsv4/compressor_v2.py @@ -156,6 +156,8 @@ class CompressorBackendMixin: out_loc: torch.Tensor, use_fp4_indexer: bool = False, bf16_store: bool = False, + kv_scale_cache: Optional[torch.Tensor] = None, + rope_cache: Optional[tuple[torch.Tensor, torch.Tensor]] = None, ) -> None: assert compress_ratio == 4 or compress_ratio == 128 assert rotate == is_indexer == (head_dim == 128) @@ -175,6 +177,11 @@ class CompressorBackendMixin: kv_score_buffer = kv_score_buffer.view(-1, compress_ratio, last_dim) # Step 1: compress_forward + compress_out = None + if _is_hip and use_fp4_indexer: + compress_out = kv_score_input.new_empty( + (plan[1].shape[0], head_dim), dtype=torch.bfloat16 + ) kv_compressed = compress_forward( kv_score_buffer=kv_score_buffer, kv_score_input=kv_score_input, @@ -182,14 +189,21 @@ class CompressorBackendMixin: plan=plan, compress_ratio=compress_ratio, head_dim=head_dim, + out=compress_out, is_online=is_online, ) + # The AITER FP4 writer takes BF16; the compressor mirrors its FP32 norm + # weight so the conversion stays out of the per-layer forward. + norm_weight = norm.weight + if _is_hip and use_fp4_indexer: + norm_weight = getattr(norm, "fp4_weight_bf16", norm_weight) + # Step 2: norm + rope + store compress_norm_rope_store( kv_compressed, plan, - norm_weight=norm.weight, + norm_weight=norm_weight, norm_eps=norm.variance_epsilon, freq_cis=freqs_cis_cache, out_loc=out_loc, @@ -197,6 +211,15 @@ class CompressorBackendMixin: page_size=page_size, use_fp4=use_fp4_indexer, bf16_store=bf16_store, + kvcache_scale=kv_scale_cache, + rope_cache=rope_cache, + # Derived once per forward by the backend; every C4 layer writes the + # same rows to the same slots. + fp4_k_write_metadata=( + getattr(self.forward_metadata, "fp4_k_write_metadata", None) + if _is_hip and use_fp4_indexer + else None + ), ) def forward_unified( @@ -222,10 +245,16 @@ class CompressorBackendMixin: use_fp4_indexer = ( compressor.is_in_indexer and self.enable_deepseek_v4_fp4_indexer ) + use_hip_fp4 = _is_hip and use_fp4_indexer bf16_store = False + kv_scale_cache = None if compressor.is_in_indexer: - kv_cache = token_to_kv_pool.get_index_k_with_scale_buffer(layer_id) page_size = token_to_kv_pool.get_index_k_page_size() + if use_hip_fp4: + kv_cache = token_to_kv_pool.get_index_k_fp4_payload_buffer(layer_id) + kv_scale_cache = token_to_kv_pool.get_index_k_fp4_scale_buffer(layer_id) + else: + kv_cache = token_to_kv_pool.get_index_k_with_scale_buffer(layer_id) elif is_unified_kv_triton(): kv_cache = token_to_kv_pool.get_unified_kv(layer_id) page_size = 1 @@ -248,7 +277,7 @@ class CompressorBackendMixin: head_dim=compressor.head_dim, norm=compressor.norm, freqs_cis_cache=compressor.freqs_cis, - kv_cache=kv_cache.view(dtype=torch.uint8), + kv_cache=kv_cache if use_hip_fp4 else kv_cache.view(dtype=torch.uint8), is_indexer=compressor.is_in_indexer, rotate=compressor.rotate, compress_ratio=compressor.ratio, @@ -256,6 +285,10 @@ class CompressorBackendMixin: out_loc=out_loc, use_fp4_indexer=use_fp4_indexer, bf16_store=bf16_store, + kv_scale_cache=kv_scale_cache, + rope_cache=( + (compressor.fp4_cos, compressor.fp4_sin) if use_hip_fp4 else None + ), ) online_c128_mtp = getattr(self, "online_c128_mtp", None) if online_c128_mtp is not None: diff --git a/python/sglang/srt/layers/attention/dsv4/indexer.py b/python/sglang/srt/layers/attention/dsv4/indexer.py index 3568d9f47..f9fccebe8 100644 --- a/python/sglang/srt/layers/attention/dsv4/indexer.py +++ b/python/sglang/srt/layers/attention/dsv4/indexer.py @@ -22,6 +22,10 @@ from sglang.kernels.ops.attention.dsv4 import ( topk_transform_paged, topk_transform_paged_v2, ) +from sglang.kernels.ops.attention.dsv4.fp4_indexer_hip import ( + aiter_fp4_paged_mqa_logits, + aiter_q_indexer_fp4, +) from sglang.kernels.ops.quantization.fp8_kernel import is_fp8_fnuz from sglang.srt.configs.deepseek_v4 import DeepSeekV4Config from sglang.srt.environ import envs @@ -694,8 +698,13 @@ class C4IndexerBackendMixin: core_metadata = metadata.core_metadata assert isinstance(indexer_metadata, PagedIndexerMetadata) + use_aiter_fp4 = c4_indexer.use_fp4_indexer and is_hip() positions = core_metadata.positions + if use_aiter_fp4: + widened = getattr(metadata, "fp4_q_positions", None) + if widened is not None and widened.shape == positions.shape: + positions = widened num_queries = min(x.shape[0], q_lora.shape[0], positions.shape[0]) if x.shape[0] != num_queries: x = x[:num_queries] @@ -727,7 +736,9 @@ class C4IndexerBackendMixin: use_fp4_indexer = c4_indexer.use_fp4_indexer - if use_fp4_indexer: + if use_aiter_fp4: + q = q_indexer + elif use_fp4_indexer: q_fp4, q_sf = q_indexer assert len(q_fp4.shape) == 3 assert len(q_sf.shape) == 2 @@ -736,13 +747,18 @@ class C4IndexerBackendMixin: assert len(q_indexer.shape) == 3 q = q_indexer.unsqueeze(1) - assert len(weights.shape) == 3 - weights = weights.squeeze(2) - if use_fp4_indexer: + if use_aiter_fp4: + weights = weights.contiguous() + else: + assert len(weights.shape) == 3 + weights = weights.squeeze(2) + if use_fp4_indexer and not use_aiter_fp4: weights = weights.float() if envs.SGLANG_OPT_USE_TILELANG_INDEXER.get(): raise RuntimeError("DeepSeek V4 FP4 indexer requires DeepGEMM indexer.") from deep_gemm import fp8_fp4_paged_mqa_logits as fn + elif use_aiter_fp4: + fn = None elif envs.SGLANG_OPT_USE_TILELANG_INDEXER.get(): from sglang.kernels.ops.attention.dsa.tilelang_kernel import ( tilelang_fp8_paged_mqa_logits as fn, @@ -773,7 +789,9 @@ class C4IndexerBackendMixin: pad = (0, 0) * (tensor.dim() - 1) + (0, query_rows - tensor.shape[0]) return F.pad(tensor, pad, value=value) - c4_seq_lens = match_num_queries(indexer_metadata.c4_seq_lens, value=1) + c4_seq_lens = match_num_queries( + indexer_metadata.c4_seq_lens, value=0 if use_aiter_fp4 else 1 + ) _c4sl = c4_seq_lens page_table = match_num_queries(indexer_metadata.page_table, value=0) c4_sparse_page_indices = match_num_queries( @@ -791,17 +809,39 @@ class C4IndexerBackendMixin: and not _use_tilelang and not _use_aiter and not _use_torch_fn + and not use_aiter_fp4 ): _c4sl = _c4sl.unsqueeze(-1) - nonpaged_plan = self._get_nonpaged_indexer_plan( - c4_indexer=c4_indexer, - forward_batch=forward_batch, - indexer_metadata=indexer_metadata, - page_table=page_table, - c4_seq_lens=c4_seq_lens, - query_rows=query_rows, - ) - if nonpaged_plan is not None: + nonpaged_plan = None + if not use_aiter_fp4: + nonpaged_plan = self._get_nonpaged_indexer_plan( + c4_indexer=c4_indexer, + forward_batch=forward_batch, + indexer_metadata=indexer_metadata, + page_table=page_table, + c4_seq_lens=c4_seq_lens, + query_rows=query_rows, + ) + if use_aiter_fp4: + q_fp4, q_scale = q + logits = aiter_fp4_paged_mqa_logits( + q_fp4=q_fp4, + q_scale=q_scale, + k_payload=token_to_kv_pool.get_index_k_fp4_payload_buffer( + c4_indexer.layer_id + ), + k_scale=token_to_kv_pool.get_index_k_fp4_scale_buffer( + c4_indexer.layer_id + ), + weights=weights, + page_table=page_table, + c4_seq_lens=c4_seq_lens, + weight_scale=c4_indexer.weight_scale, + is_decode=forward_batch.forward_mode.is_decode(), + decode_workspace=metadata.fp4_decode_workspace, + prefill_workspace=metadata.fp4_prefill_workspace, + ) + elif nonpaged_plan is not None: assert isinstance(q_indexer, torch.Tensor) logits = self._forward_nonpaged_indexer( q_indexer=q_indexer, @@ -926,6 +966,8 @@ class C4Indexer(nn.Module): prefix: str = "", alt_streams: Optional[List[torch.cuda.Stream]] = None, rotary_emb=None, + fp4_cos: Optional[torch.Tensor] = None, + fp4_sin: Optional[torch.Tensor] = None, ): super().__init__() self.layer_id = layer_id @@ -937,6 +979,7 @@ class C4Indexer(nn.Module): self.q_lora_rank = config.q_lora_rank self.softmax_scale = self.head_dim**-0.5 self.n_local_heads = self.n_heads + self.use_fp4_indexer = get_exec().kernel.enable_deepseek_v4_fp4_indexer self.wq_b = ReplicatedLinear( self.q_lora_rank, self.n_heads * self.head_dim, @@ -970,11 +1013,13 @@ class C4Indexer(nn.Module): quant_config=expert_pack_quant_config, rotary_emb=rotary_emb, ) + if self.use_fp4_indexer and is_hip(): + self.compressor.fp4_cos = fp4_cos + self.compressor.fp4_sin = fp4_sin self.rotary_emb = rotary_emb self.freqs_cis = freqs_cis self.weight_scale: float = self.softmax_scale * self.n_heads**-0.5 - self.use_fp4_indexer = get_exec().kernel.enable_deepseek_v4_fp4_indexer self.alt_streams = alt_streams def compute_q( @@ -985,6 +1030,14 @@ class C4Indexer(nn.Module): ) -> Tuple[IndexerQuery, torch.Tensor]: q, _ = self.wq_b(q_lora) q = q.view(-1, self.n_local_heads, self.head_dim) + if self.use_fp4_indexer and is_hip(): + q_fp4, q_scale = aiter_q_indexer_fp4( + q.contiguous(), + self.compressor.fp4_cos, + self.compressor.fp4_sin, + positions, + ) + return (q_fp4, q_scale), weight if self.use_fp4_indexer: return fused_q_indexer_rope_hadamard_fp4_quant( q.contiguous(), weight, self.weight_scale, self.freqs_cis, positions diff --git a/python/sglang/srt/layers/attention/dsv4/metadata.py b/python/sglang/srt/layers/attention/dsv4/metadata.py index 0d01451fc..56c762cf1 100644 --- a/python/sglang/srt/layers/attention/dsv4/metadata.py +++ b/python/sglang/srt/layers/attention/dsv4/metadata.py @@ -123,9 +123,7 @@ class PagedIndexerMetadata: def __post_init__(self): if ( - envs.SGLANG_FP8_PAGED_MQA_LOGITS_TORCH.get() - or is_xpu() - or envs.SGLANG_OPT_USE_AITER_INDEXER.get() + is_hip() or is_xpu() or envs.SGLANG_FP8_PAGED_MQA_LOGITS_TORCH.get() ) and not self.force_deep_gemm_metadata: self.deep_gemm_metadata = None else: diff --git a/python/sglang/srt/mem_cache/deepseek_v4_memory_pool.py b/python/sglang/srt/mem_cache/deepseek_v4_memory_pool.py index 54eabcb7e..ec2adad9f 100644 --- a/python/sglang/srt/mem_cache/deepseek_v4_memory_pool.py +++ b/python/sglang/srt/mem_cache/deepseek_v4_memory_pool.py @@ -31,6 +31,13 @@ _is_hip = is_hip() ONLINE_C128 = not _is_hip and envs.SGLANG_OPT_USE_ONLINE_COMPRESS.get() +def get_dsv4_indexer_bytes_per_token(index_head_dim: int, use_fp4_indexer: bool) -> int: + """Return payload and quant-scale bytes for one compressed indexer token.""" + if use_fp4_indexer: + return index_head_dim // 2 + index_head_dim // 32 + return index_head_dim + index_head_dim // 128 * 4 + + def get_compress_state_ring_size( compress_ratio: int, is_speculative: bool = False ) -> int: @@ -285,25 +292,47 @@ class DeepSeekV4IndexerPool(KVCache): ) self.index_head_dim = index_head_dim self.use_fp4_indexer = get_exec().kernel.enable_deepseek_v4_fp4_indexer + self.uses_aiter_fp4_layout = _is_hip and self.use_fp4_indexer self._create_buffer() def get_bytes_per_token(self) -> int: - if self.use_fp4_indexer: - return self.index_head_dim // 2 + 4 - return self.index_head_dim + 4 + return get_dsv4_indexer_bytes_per_token( + self.index_head_dim, self.use_fp4_indexer + ) def _create_buffer(self): page_bytes = self.page_size * self.get_bytes_per_token() + num_pages = (self.size + self.page_size + 1) // self.page_size with self.memory_saver_adapter.region(GPU_MEMORY_TYPE_KV_CACHE): with ( torch.cuda.use_mem_pool(self.custom_mem_pool) if self.custom_mem_pool else nullcontext() ): + if self.uses_aiter_fp4_layout: + self.index_k_payload_buffer = [ + torch.zeros( + (num_pages, 1, 4, self.page_size, 16), + dtype=torch.uint8, + device=self.device, + ).view(torch.float4_e2m1fn_x2) + for _ in range(self.layer_num) + ] + self.index_k_scale_buffer = [ + torch.zeros( + (num_pages, 1, 4, self.page_size), + dtype=torch.uint8, + device=self.device, + ) + for _ in range(self.layer_num) + ] + self.index_k_with_scale_buffer = None + return + self.index_k_with_scale_buffer = [ torch.zeros( - (self.size + self.page_size + 1) // self.page_size, + num_pages, page_bytes, dtype=self.index_k_with_scale_buffer_dtype, device=self.device, @@ -326,6 +355,25 @@ class DeepSeekV4IndexerPool(KVCache): def get_index_k_with_scale_buffer(self, layer_id: int) -> torch.Tensor: return self.index_k_with_scale_buffer[layer_id] + def contiguous_page_row_buffers(self) -> List[torch.Tensor]: + """Every indexer buffer as 2D page rows, for PD and HiCache transfer. + + FP8 keeps key and scale fused in one buffer per layer; the FP4 layout + stores payload and scale separately, so it yields two buffers per layer. + """ + if self.index_k_with_scale_buffer is not None: + return self.index_k_with_scale_buffer + return [ + buf.view(torch.uint8).flatten(1) + for buf in (*self.index_k_payload_buffer, *self.index_k_scale_buffer) + ] + + def get_index_k_fp4_payload_buffer(self, layer_id: int) -> torch.Tensor: + return self.index_k_payload_buffer[layer_id] + + def get_index_k_fp4_scale_buffer(self, layer_id: int) -> torch.Tensor: + return self.index_k_scale_buffer[layer_id] + def get_index_k_scale_buffer( self, layer_id: int, @@ -702,7 +750,7 @@ class DeepSeekV4TokenToKVPool(BaseSWAKVPool): for i in c4_locals: _append_compressed_entry(i, 4) - for buf in self.c4_indexer_kv_pool.index_k_with_scale_buffer: + for buf in self.c4_indexer_kv_pool.contiguous_page_row_buffers(): assert buf.ndim == 2, f"expected 2D buffer, got {buf.ndim}D" data_ptrs.append(buf.data_ptr()) data_lens.append(buf.nbytes) @@ -714,7 +762,7 @@ class DeepSeekV4TokenToKVPool(BaseSWAKVPool): buf_groups = [ self.c4_kv_pool.kv_buffer, - self.c4_indexer_kv_pool.index_k_with_scale_buffer, + self.c4_indexer_kv_pool.contiguous_page_row_buffers(), self.c128_kv_pool.kv_buffer, ] @@ -1109,6 +1157,18 @@ class DeepSeekV4TokenToKVPool(BaseSWAKVPool): assert compress_ratio == 4, f"only c4 has indexer, got {compress_ratio = }" return self.c4_indexer_kv_pool.get_index_k_with_scale_buffer(compress_layer_id) + def get_index_k_fp4_payload_buffer(self, layer_id: int) -> torch.Tensor: + self.wait_layer_transfer(layer_id) + compress_ratio, compress_layer_id, _ = self.layer_mapping[layer_id] + assert compress_ratio == 4, f"only c4 has indexer, got {compress_ratio = }" + return self.c4_indexer_kv_pool.get_index_k_fp4_payload_buffer(compress_layer_id) + + def get_index_k_fp4_scale_buffer(self, layer_id: int) -> torch.Tensor: + self.wait_layer_transfer(layer_id) + compress_ratio, compress_layer_id, _ = self.layer_mapping[layer_id] + assert compress_ratio == 4, f"only c4 has indexer, got {compress_ratio = }" + return self.c4_indexer_kv_pool.get_index_k_fp4_scale_buffer(compress_layer_id) + def get_index_k_scale_buffer( self, layer_id: int, diff --git a/python/sglang/srt/mem_cache/hicache_storage.py b/python/sglang/srt/mem_cache/hicache_storage.py index 951cc58f6..e8d403f07 100644 --- a/python/sglang/srt/mem_cache/hicache_storage.py +++ b/python/sglang/srt/mem_cache/hicache_storage.py @@ -66,6 +66,9 @@ class PoolName(str, Enum): # 'COMPRESSED_KV / COMPRESSED_INDEXER / COMPRESSED_STATE' in the next PR. DEEPSEEK_V4_C4 = "deepseek_v4_c4" DEEPSEEK_V4_C4_INDEXER = "deepseek_v4_c4_indexer" + # FP4 indexer splits the indexer cache into separate payload/scale buffers, + # so it needs a second pool alongside DEEPSEEK_V4_C4_INDEXER. + DEEPSEEK_V4_C4_INDEXER_SCALE = "deepseek_v4_c4_indexer_scale" DEEPSEEK_V4_C128 = "deepseek_v4_c128" DEEPSEEK_V4_C4_STATE = "deepseek_v4_c4_state" DEEPSEEK_V4_C4_INDEXER_STATE = "deepseek_v4_c4_indexer_state" diff --git a/python/sglang/srt/mem_cache/hybrid_cache/hybrid_pool_assembler.py b/python/sglang/srt/mem_cache/hybrid_cache/hybrid_pool_assembler.py index cef0ef92b..b3ed6f84d 100644 --- a/python/sglang/srt/mem_cache/hybrid_cache/hybrid_pool_assembler.py +++ b/python/sglang/srt/mem_cache/hybrid_cache/hybrid_pool_assembler.py @@ -470,6 +470,72 @@ def _dsv4_compressed_region_buffers(kvcache: Any, ratio: int) -> tuple[list, int return pool.kv_buffer, pool.bytes_per_page_padded +@dataclass(frozen=True) +class _IndexerRegion: + """One page-contiguous indexer buffer group to mirror on the host.""" + + name: PoolName + device_buffers: list + item_bytes: int + # FP4 page rows group their slots instead of laying tokens out flat, so the + # fused-row token-granular copy does not apply and transfers must be whole + # pages. The fused FP8 row has no such restriction. + page_aligned_only: bool + + +def _dsv4_indexer_regions(kvcache: Any, page_size: int) -> list[_IndexerRegion]: + """ + Resolve the indexer HiCache regions, hiding the FP8/FP4 split from the + stack builder. FP8 keeps key and scale fused in one buffer, while FP4 + stores payload and scale separately, so it maps to two host pools. + """ + import torch + + pool = kvcache.c4_indexer_kv_pool + fused = pool.index_k_with_scale_buffer + if fused is not None: + return [ + _IndexerRegion( + name=PoolName.DEEPSEEK_V4_C4_INDEXER, + device_buffers=fused, + item_bytes=fused[0].shape[1] * fused[0].element_size(), + page_aligned_only=False, + ) + ] + + payload_ref = pool.index_k_payload_buffer[0] + scale_ref = pool.index_k_scale_buffer[0] + # A page row covers ``page_slots`` C4 slots, i.e. one tree page of tokens + # after 4:1 compression. + page_slots = payload_ref.shape[3] + if scale_ref.shape[3] != page_slots: + raise ValueError( + "FP4 indexer payload and scale must agree on slots per page: " + f"payload={page_slots}, scale={scale_ref.shape[3]}" + ) + if page_size % page_slots != 0: + raise ValueError( + f"Tree page size {page_size} must be a multiple of the FP4 indexer " + f"slots per page {page_slots}" + ) + payload = [b.view(torch.uint8).flatten(1) for b in pool.index_k_payload_buffer] + scale = [b.view(torch.uint8).flatten(1) for b in pool.index_k_scale_buffer] + return [ + _IndexerRegion( + name=PoolName.DEEPSEEK_V4_C4_INDEXER, + device_buffers=payload, + item_bytes=payload[0].shape[1], + page_aligned_only=True, + ), + _IndexerRegion( + name=PoolName.DEEPSEEK_V4_C4_INDEXER_SCALE, + device_buffers=scale, + item_bytes=scale[0].shape[1], + page_aligned_only=True, + ), + ] + + def build_deepseek_v4_hicache_stack( *, params: CacheInitParams, @@ -582,36 +648,34 @@ def build_deepseek_v4_hicache_stack( layout=get_memory().hicache_mem_layout, allocator_type=_get_allocator_type(), ) - c4_indexer_host_pool = DeepSeekV4PagedHostPool( - pool_name=str(PoolName.DEEPSEEK_V4_C4_INDEXER), - device_buffers=kvcache.c4_indexer_kv_pool.index_k_with_scale_buffer, - item_bytes=( - kvcache.c4_indexer_kv_pool.index_k_with_scale_buffer[0].shape[1] - * kvcache.c4_indexer_kv_pool.index_k_with_scale_buffer[0].element_size() - ), - num_host_pages=num_host_pages, - slot_page_size=page_size, - layout=get_memory().hicache_mem_layout, - allocator_type=_get_allocator_type(), + entries.append( + build_pool_entry( + name=PoolName.DEEPSEEK_V4_C4, + host_pool=c4_host_pool, + device_pool=kvcache.c4_kv_pool, + layer_mapping=c4_layer_mapping, + transfer_layer_num=transfer_layer_num, + ) ) - entries.extend( - [ + for region in _dsv4_indexer_regions(kvcache, page_size): + entries.append( build_pool_entry( - name=PoolName.DEEPSEEK_V4_C4, - host_pool=c4_host_pool, - device_pool=kvcache.c4_kv_pool, - layer_mapping=c4_layer_mapping, - transfer_layer_num=transfer_layer_num, - ), - build_pool_entry( - name=PoolName.DEEPSEEK_V4_C4_INDEXER, - host_pool=c4_indexer_host_pool, + name=region.name, + host_pool=DeepSeekV4PagedHostPool( + pool_name=str(region.name), + device_buffers=region.device_buffers, + item_bytes=region.item_bytes, + num_host_pages=num_host_pages, + slot_page_size=page_size, + layout=get_memory().hicache_mem_layout, + allocator_type=_get_allocator_type(), + page_aligned_only=region.page_aligned_only, + ), device_pool=kvcache.c4_indexer_kv_pool, layer_mapping=c4_layer_mapping, transfer_layer_num=transfer_layer_num, - ), - ] - ) + ) + ) if not is_unified_kv: c4_state_host_pool = DeepSeekV4StateHostPool( @@ -1276,6 +1340,7 @@ class _DeepSeekV4Strategy(StackStrategy): for name, src in ( (PoolName.DEEPSEEK_V4_C4, PoolName.KV), (PoolName.DEEPSEEK_V4_C4_INDEXER, PoolName.KV), + (PoolName.DEEPSEEK_V4_C4_INDEXER_SCALE, PoolName.KV), (PoolName.DEEPSEEK_V4_C128, PoolName.KV), (PoolName.DEEPSEEK_V4_C4_STATE, PoolName.SWA), (PoolName.DEEPSEEK_V4_C4_INDEXER_STATE, PoolName.SWA), diff --git a/python/sglang/srt/mem_cache/hybrid_cache/linker_pool_assembler.py b/python/sglang/srt/mem_cache/hybrid_cache/linker_pool_assembler.py index 9fbf01554..bf8d90fcb 100644 --- a/python/sglang/srt/mem_cache/hybrid_cache/linker_pool_assembler.py +++ b/python/sglang/srt/mem_cache/hybrid_cache/linker_pool_assembler.py @@ -239,6 +239,7 @@ def _build_deepseek_v4_device_pool_group( ) -> DevicePoolGroup: from sglang.srt.mem_cache.deepseek_v4_memory_pool import HiSparseC4DevicePool from sglang.srt.mem_cache.hybrid_cache.hybrid_pool_assembler import ( + _dsv4_indexer_regions, _resolve_deepseek_v4_layer_mappings, ) @@ -288,13 +289,14 @@ def _build_deepseek_v4_device_pool_group( kvcache.c4_kv_pool.kv_buffer, mappings.c4, ) - add( - PoolName.DEEPSEEK_V4_C4_INDEXER, - PoolName.KV, - kvcache.c4_indexer_kv_pool, - kvcache.c4_indexer_kv_pool.index_k_with_scale_buffer, - mappings.c4, - ) + for region in _dsv4_indexer_regions(kvcache, page_size): + add( + region.name, + PoolName.KV, + kvcache.c4_indexer_kv_pool, + region.device_buffers, + mappings.c4, + ) add( PoolName.DEEPSEEK_V4_C128, PoolName.KV, diff --git a/python/sglang/srt/mem_cache/memory_pool_host.py b/python/sglang/srt/mem_cache/memory_pool_host.py index f6a83dab4..006d0dbde 100644 --- a/python/sglang/srt/mem_cache/memory_pool_host.py +++ b/python/sglang/srt/mem_cache/memory_pool_host.py @@ -185,10 +185,15 @@ class DeepSeekV4PagedHostPool(HiSparseHostPoolMixin, HostKVCache): device: str = "cpu", pin_memory: bool = True, allocator_type: str = "default", + page_aligned_only: bool = False, ): self.pool_name = pool_name self.layer_num = len(device_buffers) self.item_bytes = item_bytes + # A page row of the FP4 indexer buffers is a grouped slot layout rather + # than a flat token array, so the token-granular copy used for fused + # DSv4 C4 rows does not apply and only whole pages may move. + self.page_aligned_only = page_aligned_only self.num_host_pages = num_host_pages self.slot_page_size = slot_page_size self.dtype = torch.uint8 @@ -305,6 +310,15 @@ class DeepSeekV4PagedHostPool(HiSparseHostPoolMixin, HostKVCache): def _to_page_indices(self, indices: torch.Tensor) -> torch.Tensor: return indices.reshape(-1, self.slot_page_size)[:, 0] // self.slot_page_size + def _unaligned_transfer_error( + self, host_indices: torch.Tensor, device_indices: torch.Tensor + ) -> ValueError: + return ValueError( + f"{self.pool_name} expects page-aligned indices: got " + f"{host_indices.numel()} host and {device_indices.numel()} device " + f"indices for page size {self.slot_page_size}." + ) + def _has_transfer_indices( self, host_indices: torch.Tensor | None, device_indices: torch.Tensor | None ) -> bool: @@ -375,6 +389,8 @@ class DeepSeekV4PagedHostPool(HiSparseHostPoolMixin, HostKVCache): # Token-granular DSV4 C4 copy needs this helper because a token is # not one contiguous byte range in the paged row: # [value0..value63][scale0..scale63]. + if self.page_aligned_only: + raise self._unaligned_transfer_error(host_indices, device_indices) transfer_cache_dsv4_mla( src_ptrs=self.device_ptrs, dst_ptrs=self.data_ptrs, @@ -453,6 +469,8 @@ class DeepSeekV4PagedHostPool(HiSparseHostPoolMixin, HostKVCache): ): # Same DSV4 C4 layout issue as backup: this is token-granular # preload, so it cannot use the normal HiCache page-row copy. + if self.page_aligned_only: + raise self._unaligned_transfer_error(host_indices, device_indices) transfer_cache_dsv4_mla( src_ptrs=self.data_ptrs[layer_id : layer_id + 1], dst_ptrs=self.device_ptrs[layer_id : layer_id + 1], diff --git a/python/sglang/srt/mem_cache/storage/mooncake_store/mooncake_store.py b/python/sglang/srt/mem_cache/storage/mooncake_store/mooncake_store.py index a9362a1c0..6266059cf 100644 --- a/python/sglang/srt/mem_cache/storage/mooncake_store/mooncake_store.py +++ b/python/sglang/srt/mem_cache/storage/mooncake_store/mooncake_store.py @@ -799,6 +799,7 @@ class MooncakeStore(HiCacheStorage, MooncakeBaseStore): PoolName.DRAFT_INDEXER, PoolName.DEEPSEEK_V4_C4, PoolName.DEEPSEEK_V4_C4_INDEXER, + PoolName.DEEPSEEK_V4_C4_INDEXER_SCALE, PoolName.DEEPSEEK_V4_C128, PoolName.DEEPSEEK_V4_C4_STATE, PoolName.DEEPSEEK_V4_C4_INDEXER_STATE, diff --git a/python/sglang/srt/model_executor/pool_configurator.py b/python/sglang/srt/model_executor/pool_configurator.py index 38b7fccd4..8bfd033b6 100644 --- a/python/sglang/srt/model_executor/pool_configurator.py +++ b/python/sglang/srt/model_executor/pool_configurator.py @@ -36,6 +36,7 @@ from sglang.srt.mem_cache.allocation_sizing import get_alloc_len_per_decode from sglang.srt.mem_cache.deepseek_v4_memory_pool import ( get_compress_state_ring_size, get_compress_state_write_pad, + get_dsv4_indexer_bytes_per_token, ) from sglang.srt.mem_cache.memory_pool import DSATokenToKVPool from sglang.srt.runtime_context import ( @@ -50,9 +51,12 @@ from sglang.srt.utils.common import ( ceil_align, ceil_div, is_float4_e2m1fn_x2, + is_hip, spec_decode_alloc_len_per_request, ) +_is_hip = is_hip() + @dataclass class MemoryPoolConfig: @@ -774,6 +778,12 @@ class DSV4PoolConfigurator(MemoryPoolConfigurator): self.qk_nope_head_dim = cfg.qk_nope_head_dim self.qk_rope_head_dim = cfg.qk_rope_head_dim self.indexer_head_dim = cfg.index_head_dim + # HIP takes the FP4-accurate byte count here. The NVIDIA FP4 path + # keeps the FP8 estimate. + self.indexer_bytes_per_token = get_dsv4_indexer_bytes_per_token( + self.indexer_head_dim, + _is_hip and kvc.server_args.enable_deepseek_v4_fp4_indexer, + ) self.context_len = kvc.model_config.context_len # PP-local slice; matches DeepSeekV4TokenToKVPool's stage_ratios. self.compression_ratios = cfg.compress_ratios[ @@ -884,11 +894,6 @@ class DSV4PoolConfigurator(MemoryPoolConfigurator): def _get_bytes_per_full_token(self) -> float: kv_bytes = self.qk_nope_head_dim + self.qk_rope_head_dim * 2 + 8 - quant_block_size = 128 - indexer_bytes = ( - self.indexer_head_dim + self.indexer_head_dim // quant_block_size * 4 - ) - attn_head_dim = self.qk_nope_head_dim + self.qk_rope_head_dim c4_state_dtype_size, c128_state_dtype_size = ( _get_dsv4_compress_state_dtype_sizes() @@ -914,7 +919,7 @@ class DSV4PoolConfigurator(MemoryPoolConfigurator): self.swa_ratio * kv_bytes * self.num_layers_total + c4_frac * kv_bytes * self.num_layers_ca4 + 1 / 128 * kv_bytes * self.num_layers_ca128 - + 1 / 4 * indexer_bytes * self.num_layers_ca4 + + 1 / 4 * self.indexer_bytes_per_token * self.num_layers_ca4 + self.swa_ratio * c4_state_ratio * c4_state_bytes * self.num_layers_ca4 + c128_state_ratio * c128_state_bytes * self.num_layers_ca128 + self.swa_ratio diff --git a/python/sglang/srt/models/deepseek_v4.py b/python/sglang/srt/models/deepseek_v4.py index 9492090de..fba2702e7 100644 --- a/python/sglang/srt/models/deepseek_v4.py +++ b/python/sglang/srt/models/deepseek_v4.py @@ -891,6 +891,8 @@ class MQALayer(MqaAttentionBase): prefix=add_prefix("indexer", prefix), alt_streams=self.alt_streams_indexer, rotary_emb=self.rotary_emb, + fp4_cos=(self.cos_cache[:, 0, 0, :] if _is_hip else None), + fp4_sin=(self.sin_cache[:, 0, 0, :] if _is_hip else None), ) self.attn_mqa = RadixAttention( @@ -911,6 +913,13 @@ class MQALayer(MqaAttentionBase): # (`_compute_kv_to_cache`), so the legacy "overlap store cache" flag # has no effect here -- the fused path is on by default. + def _apply(self, fn, recurse=True): + result = super()._apply(fn, recurse=recurse) + if self.indexer is not None and hasattr(self.indexer.compressor, "fp4_cos"): + self.indexer.compressor.fp4_cos = self.cos_cache[:, 0, 0, :] + self.indexer.compressor.fp4_sin = self.sin_cache[:, 0, 0, :] + return result + def _get_npu_rope_position_cache( self, positions: torch.Tensor, dtype: torch.dtype, inverse: bool = False ) -> Tuple[torch.Tensor, torch.Tensor]: diff --git a/rust/sglang-radix-tree/src/python_bindings.rs b/rust/sglang-radix-tree/src/python_bindings.rs index ce6695c6c..6b6bc3823 100644 --- a/rust/sglang-radix-tree/src/python_bindings.rs +++ b/rust/sglang-radix-tree/src/python_bindings.rs @@ -235,6 +235,7 @@ fn pool_name_str(name: PoolName) -> &'static str { PoolName::Indexer => "indexer", PoolName::DeepseekV4C4 => "deepseek_v4_c4", PoolName::DeepseekV4C4Indexer => "deepseek_v4_c4_indexer", + PoolName::DeepseekV4C4IndexerScale => "deepseek_v4_c4_indexer_scale", PoolName::DeepseekV4C128 => "deepseek_v4_c128", PoolName::DeepseekV4C4State => "deepseek_v4_c4_state", PoolName::DeepseekV4C4IndexerState => "deepseek_v4_c4_indexer_state", @@ -254,6 +255,7 @@ fn parse_pool_name(name: &str) -> PyResult { "indexer" => Ok(PoolName::Indexer), "deepseek_v4_c4" => Ok(PoolName::DeepseekV4C4), "deepseek_v4_c4_indexer" => Ok(PoolName::DeepseekV4C4Indexer), + "deepseek_v4_c4_indexer_scale" => Ok(PoolName::DeepseekV4C4IndexerScale), "deepseek_v4_c128" => Ok(PoolName::DeepseekV4C128), "deepseek_v4_c4_state" => Ok(PoolName::DeepseekV4C4State), "deepseek_v4_c4_indexer_state" => Ok(PoolName::DeepseekV4C4IndexerState), diff --git a/rust/sglang-radix-tree/src/unified_tree_core.rs b/rust/sglang-radix-tree/src/unified_tree_core.rs index a54bad262..55989fd07 100644 --- a/rust/sglang-radix-tree/src/unified_tree_core.rs +++ b/rust/sglang-radix-tree/src/unified_tree_core.rs @@ -312,6 +312,7 @@ pub enum PoolName { Indexer, DeepseekV4C4, DeepseekV4C4Indexer, + DeepseekV4C4IndexerScale, DeepseekV4C128, DeepseekV4C4State, DeepseekV4C4IndexerState, diff --git a/test/registered/kernels/ops/attention/test_fp4_indexer_hip.py b/test/registered/kernels/ops/attention/test_fp4_indexer_hip.py new file mode 100644 index 000000000..11b877b97 --- /dev/null +++ b/test/registered/kernels/ops/attention/test_fp4_indexer_hip.py @@ -0,0 +1,776 @@ +"""HIP counterpart of ``test_fp4_indexer.py`` for the AITER FP4 DeepSeek-V4 indexer. + +The CUDA path exposes the quantizer, the cache store and the fused +norm/RoPE/store as separate Triton entry points. On HIP all three collapse into +``aiter_k_indexer_fp4_cache_write``, and the query side into +``aiter_q_indexer_fp4``, so the tests below mirror the CUDA file's four cases +through those two ops. The FP4 grid and the UE8M0 scale rule are identical on +both targets, so the reference helpers are shared verbatim. + +Two layout details differ from CUDA and are pinned here because nothing else +checks them: the K cache keeps payload and scale in separate buffers with the +scale token axis shuffled, and the Q scale is emitted preshuffled into the +logits kernel's ABI layout. +""" + +from __future__ import annotations + +import sys + +import pytest +import torch + +from sglang.kernels.ops.attention.deepseek_v4_rope import precompute_freqs_cis +from sglang.kernels.ops.attention.dsv4 import ( + CompressorDecodePlan, + compress_norm_rope_store, +) +from sglang.kernels.ops.attention.dsv4.compress import CompressorPrefillPlan +from sglang.kernels.ops.attention.dsv4.fp4_indexer_hip import ( + FP4KWriteMetadata, + _decode_cta_count, + _guard_page_table, + aiter_fp4_paged_mqa_logits, + aiter_k_indexer_fp4_cache_write, + aiter_q_indexer_fp4, + prepare_fp4_decode_workspace, + prepare_fp4_k_write_metadata, + prepare_fp4_prefill_workspace, +) +from sglang.srt.utils import get_device, is_gfx95_supported, is_hip +from sglang.test.ci.ci_register import register_amd_ci + +register_amd_ci(est_time=120, suite="stage-b-test-1-gpu-small-amd-mi35x") + +pytestmark = pytest.mark.skipif( + not (is_hip() and is_gfx95_supported()), + reason="The FP4 indexer adapters wrap AITER CDNA4 (gfx95x) kernels.", +) + +HEAD_DIM = 128 +FP4_DIM = HEAD_DIM // 2 +GROUP_SIZE = 32 +SCALE_GROUPS = HEAD_DIM // GROUP_SIZE +PAGE_SIZE = 64 +E2M1_MAX = 6.0 +NUM_HEADS = 64 +ROPE_DIM = 64 +NORM_EPS = 1.0e-6 +# Tokens per group along the shuffled scale axis; see _ref_store_fp4_index_cache. +SCALE_SHUFFLE_TILE = 16 + +_E2M1_GRID = [0.0, 0.5, 1.0, 1.5, 2.0, 3.0, 4.0, 6.0] + + +def _ceil_ue8m0_exp_ref(x: torch.Tensor) -> torch.Tensor: + bits = x.to(torch.float32).contiguous().view(torch.int32) + exp = (bits >> 23) & 0xFF + mantissa = bits & 0x7FFFFF + exp = exp + (mantissa != 0).to(torch.int32) + return exp.clamp(1, 254) + + +def _fp4_e2m1_code_ref(x: torch.Tensor) -> torch.Tensor: + ax = torch.minimum(x.abs(), torch.tensor(E2M1_MAX, device=x.device)) + idx = torch.zeros_like(ax, dtype=torch.uint8) + for threshold in (0.25, 0.75, 1.25, 1.75, 2.5, 3.5, 5.0): + idx += (ax > threshold).to(torch.uint8) + sign = ((x < 0) & (idx != 0)).to(torch.uint8) * 8 + return idx | sign + + +def _ref_quantize_fp4_indexer(x: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: + """Nibble-packed E2M1 payload and per-group UE8M0 exponents. + + Same rule as the CUDA reference, except the exponents stay one byte per + group instead of being packed into an int32: the HIP cache stores them in a + separate buffer, one byte per (group, token). + """ + x = x.contiguous().view(-1, HEAD_DIM).float() + groups = x.view(-1, SCALE_GROUPS, GROUP_SIZE) + scale_raw = (groups.abs().amax(dim=-1) / E2M1_MAX).clamp_min(1.0e-4) + scale_exp = _ceil_ue8m0_exp_ref(scale_raw) + scale = (scale_exp << 23).contiguous().view(torch.float32) + + scaled = (groups / scale.unsqueeze(-1)).view(-1, HEAD_DIM) + code = _fp4_e2m1_code_ref(scaled) + packed = (code[:, 0::2].to(torch.int16) | (code[:, 1::2].to(torch.int16) << 4)).to( + torch.uint8 + ) + return packed, scale_exp.to(torch.uint8) + + +def _canonical_zero(packed: torch.Tensor) -> torch.Tensor: + """Fold negative zero onto positive zero in both nibbles. + + AITER keeps the sign bit when a lane quantizes to zero magnitude, while the + reference clears it. Both decode to 0.0, so normalize before comparing the + packed bytes. + """ + lo, hi = packed & 0x0F, packed >> 4 + zero = torch.zeros_like(lo) + lo = torch.where(lo == 0x8, zero, lo) + hi = torch.where(hi == 0x8, zero, hi) + return lo | (hi << 4) + + +def _ref_dequantize_fp4_indexer( + packed: torch.Tensor, scale_exp: torch.Tensor +) -> torch.Tensor: + """Inverse of :func:`_ref_quantize_fp4_indexer`, for value comparisons.""" + packed = packed.reshape(-1, FP4_DIM) + codes = torch.stack([packed & 0x0F, packed >> 4], dim=-1).long() + values = torch.tensor(_E2M1_GRID + [-v for v in _E2M1_GRID], device=packed.device)[ + codes.reshape(-1, SCALE_GROUPS, GROUP_SIZE) + ] + factor = (scale_exp.reshape(-1, SCALE_GROUPS).to(torch.int32) << 23).view( + torch.float32 + ) + return (values * factor.unsqueeze(-1)).reshape(-1, HEAD_DIM) + + +def _empty_index_k_cache(num_pages: int) -> tuple[torch.Tensor, torch.Tensor]: + """Allocate the split buffers ``uses_aiter_fp4_layout`` creates per layer.""" + payload = torch.zeros( + (num_pages, 1, SCALE_GROUPS, PAGE_SIZE, GROUP_SIZE // 2), + dtype=torch.uint8, + device=get_device(), + ).view(torch.float4_e2m1fn_x2) + scale = torch.zeros( + (num_pages, 1, SCALE_GROUPS, PAGE_SIZE), dtype=torch.uint8, device=get_device() + ) + return payload, scale + + +def _ref_store_fp4_index_cache( + x_fp4: torch.Tensor, + x_sf: torch.Tensor, + loc: torch.Tensor, + num_pages: int, +) -> tuple[torch.Tensor, torch.Tensor]: + """Build the payload and scale buffers a correct writer would produce. + + Payload rows stay in slot order. The scale buffer stores each page's tokens + as the transpose of a ``SCALE_SHUFFLE_TILE x 4`` tile, so token ``t`` lands + at ``(t % 16) * 4 + t // 16``. Rows whose ``loc`` is negative are skipped and + must be left at zero. + """ + payload = torch.zeros( + (num_pages, 1, SCALE_GROUPS, PAGE_SIZE, GROUP_SIZE // 2), + dtype=torch.uint8, + device=x_fp4.device, + ) + scale = torch.zeros( + (num_pages, 1, SCALE_GROUPS, PAGE_SIZE), dtype=torch.uint8, device=x_fp4.device + ) + for token_id in range(x_fp4.shape[0]): + cache_loc = int(loc[token_id].item()) + if cache_loc < 0: + continue + page, offset = divmod(cache_loc, PAGE_SIZE) + shuffled = (offset % SCALE_SHUFFLE_TILE) * 4 + offset // SCALE_SHUFFLE_TILE + for group in range(SCALE_GROUPS): + lo = group * (GROUP_SIZE // 2) + payload[page, 0, group, offset] = x_fp4[token_id, lo : lo + GROUP_SIZE // 2] + scale[page, 0, group, shuffled] = x_sf[token_id, group] + return payload, scale + + +def _read_index_k_cache(payload, scale, loc: torch.Tensor): + """Gather ``loc`` back out of the cache as (packed nibbles, exponents).""" + page, offset = loc // PAGE_SIZE, loc % PAGE_SIZE + shuffled = (offset % SCALE_SHUFFLE_TILE) * 4 + offset // SCALE_SHUFFLE_TILE + packed = payload.view(torch.uint8)[page, 0, :, offset].reshape(-1, FP4_DIM) + return packed, scale[page, 0, :, shuffled] + + +def _rope_tables( + max_pos: int, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """The bf16 cos/sin pair ``DeepseekV4AttentionMLA`` hands the FP4 adapters.""" + freqs_cis = precompute_freqs_cis(ROPE_DIM, max_pos, 0, 10000, 1, 32, 1).to( + get_device() + ) + return ( + freqs_cis.real.to(torch.bfloat16), + freqs_cis.imag.to(torch.bfloat16), + freqs_cis, + ) + + +def _ref_apply_rope(x, cos, sin, positions: torch.Tensor) -> torch.Tensor: + """Interleaved (non-neox) RoPE over the trailing ``ROPE_DIM`` lanes.""" + shape = (positions.shape[0], *(1,) * (x.dim() - 2), ROPE_DIM // 2) + c = cos.float()[positions].reshape(shape) + s = sin.float()[positions].reshape(shape) + pairs = x[..., ROPE_DIM:].reshape(*x.shape[:-1], ROPE_DIM // 2, 2) + even, odd = pairs[..., 0], pairs[..., 1] + rotated = torch.stack([even * c - odd * s, even * s + odd * c], dim=-1) + return torch.cat([x[..., :ROPE_DIM], rotated.flatten(-2)], dim=-1) + + +def _ref_hadamard(x: torch.Tensor) -> torch.Tensor: + """Orthonormal Sylvester Hadamard, the rotation ``do_rotate_act`` applies. + + Written out in torch instead of reusing ``ops.quantization.hadamard`` + because that kernel is CUDA-only and does not build under ROCm. + """ + h = torch.ones(1, 1, device=get_device(), dtype=torch.float32) + while h.shape[0] < HEAD_DIM: + h = torch.cat([torch.cat([h, h], 1), torch.cat([h, -h], 1)], 0) + return x.float() @ (h * HEAD_DIM**-0.5) + + +def _ref_k_transform(k, norm_weight, cos, sin, positions) -> torch.Tensor: + x = k.float() + x = x * torch.rsqrt((x * x).mean(dim=-1, keepdim=True) + NORM_EPS) + x = x * norm_weight.float() + return _ref_hadamard(_ref_apply_rope(x, cos, sin, positions)) + + +def _ref_q_transform(q, cos, sin, positions) -> torch.Tensor: + return _ref_hadamard(_ref_apply_rope(q.float(), cos, sin, positions)) + + +def _write_index_k_cache( + k, norm_weight, cos, sin, positions, loc, payload, scale +) -> None: + aiter_k_indexer_fp4_cache_write( + k=k, + norm_weight=norm_weight, + norm_epsilon=NORM_EPS, + cos=cos, + sin=sin, + plan=None, + out_loc=None, + k_payload=payload, + k_scale=scale, + write_metadata=FP4KWriteMetadata(positions, loc), + ) + + +# --------------------------------------------------------------------------- +# The four cases mirrored from the CUDA file +# --------------------------------------------------------------------------- + + +@pytest.mark.parametrize("num_tokens", [1, 7, 96]) +def test_quantize_fp4_indexer_tensor(num_tokens: int) -> None: + """The fused writer's quantization matches the shared FP4 reference exactly.""" + torch.manual_seed(num_tokens) + cos, sin, _ = _rope_tables(512) + payload, scale = _empty_index_k_cache(-(-num_tokens // PAGE_SIZE) + 1) + x = torch.randn(num_tokens, HEAD_DIM, device=get_device(), dtype=torch.bfloat16) + x[0, :8] = torch.tensor( + [-8.0, -6.0, -3.0, -1.5, 0.0, 0.5, 2.0, 8.0], + device=get_device(), + dtype=torch.bfloat16, + ) + norm_weight = torch.randn(HEAD_DIM, device=get_device(), dtype=torch.bfloat16) + positions = torch.arange(num_tokens, device=get_device(), dtype=torch.int64) * 3 + loc = torch.arange(num_tokens, device=get_device(), dtype=torch.int64) + + _write_index_k_cache(x, norm_weight, cos, sin, positions, loc, payload, scale) + + ref_fp4, ref_sf = _ref_quantize_fp4_indexer( + _ref_k_transform(x, norm_weight, cos, sin, positions) + ) + stored_fp4, stored_sf = _read_index_k_cache(payload, scale, loc) + torch.testing.assert_close(stored_sf, ref_sf) + torch.testing.assert_close(_canonical_zero(stored_fp4), _canonical_zero(ref_fp4)) + + +@pytest.mark.parametrize("num_tokens", [1, 16, 96]) +def test_fp4_index_cache_store_layout(num_tokens: int) -> None: + """Scattered slots land in the paged layout and touch nothing else.""" + torch.manual_seed(num_tokens + 50) + cos, sin, _ = _rope_tables(512) + num_pages = max(2, -(-num_tokens // PAGE_SIZE) + 1) + payload, scale = _empty_index_k_cache(num_pages) + x = torch.randn(num_tokens, HEAD_DIM, device=get_device(), dtype=torch.bfloat16) + norm_weight = torch.randn(HEAD_DIM, device=get_device(), dtype=torch.bfloat16) + positions = torch.arange(num_tokens, device=get_device(), dtype=torch.int64) * 3 + loc = torch.randperm(num_pages * PAGE_SIZE, device=get_device())[:num_tokens].to( + torch.int64 + ) + # A masked row must be dropped rather than written to some default slot. + loc[num_tokens // 2] = -1 + + _write_index_k_cache(x, norm_weight, cos, sin, positions, loc, payload, scale) + + ref_fp4, ref_sf = _ref_quantize_fp4_indexer( + _ref_k_transform(x, norm_weight, cos, sin, positions) + ) + expected_payload, expected_scale = _ref_store_fp4_index_cache( + ref_fp4, ref_sf, loc, num_pages + ) + torch.testing.assert_close( + _canonical_zero(payload.view(torch.uint8)), _canonical_zero(expected_payload) + ) + torch.testing.assert_close(scale, expected_scale) + + +# 17 and 33 straddle the 16-token scale shuffle tile, 65 the 64-token page. +@pytest.mark.parametrize("num_tokens", [1, 16, 17, 33, 65, 96]) +def test_fp4_fused_norm_rope_store_layout(num_tokens: int) -> None: + """The real ``compress_norm_rope_store`` entry point, plan and metadata included.""" + torch.manual_seed(num_tokens + 100) + num_pages = -(-num_tokens // PAGE_SIZE) + 1 + compress_ratio = 4 + kv = torch.randn(num_tokens, HEAD_DIM, device=get_device(), dtype=torch.bfloat16) + norm_weight = torch.randn(HEAD_DIM, device=get_device(), dtype=torch.bfloat16) + seq_lens = ( + torch.arange(1, num_tokens + 1, device=get_device(), dtype=torch.int64) + * compress_ratio + ) + req_pool_indices = torch.arange(num_tokens, device=get_device(), dtype=torch.int64) + plan = CompressorDecodePlan.generate_legacy( + compress_ratio, req_pool_indices, seq_lens + ) + loc = torch.arange(num_tokens, device=get_device(), dtype=torch.int64) + rope_len = int(seq_lens.max().item()) + 1 + cos, sin, freqs_cis = _rope_tables(rope_len) + payload, scale = _empty_index_k_cache(num_pages) + + metadata = prepare_fp4_k_write_metadata(plan, loc, rope_len) + compress_norm_rope_store( + kv.clone(), + plan, + norm_weight=norm_weight, + norm_eps=NORM_EPS, + freq_cis=freqs_cis, + out_loc=loc, + kvcache=payload, + page_size=PAGE_SIZE, + use_fp4=True, + kvcache_scale=scale, + rope_cache=(cos, sin), + fp4_k_write_metadata=metadata, + ) + + # The plan drives RoPE off the compression boundary, not the token index. + torch.testing.assert_close(metadata.positions, seq_lens - compress_ratio) + ref_fp4, ref_sf = _ref_quantize_fp4_indexer( + _ref_k_transform(kv, norm_weight, cos, sin, metadata.positions) + ) + expected_payload, expected_scale = _ref_store_fp4_index_cache( + ref_fp4, ref_sf, metadata.slots, num_pages + ) + torch.testing.assert_close( + _canonical_zero(payload.view(torch.uint8)), _canonical_zero(expected_payload) + ) + torch.testing.assert_close(scale, expected_scale) + + +@pytest.mark.parametrize("batch_size", [1, 5, 17]) +def test_fp4_fused_q_indexer_rope_hadamard_quant(batch_size: int) -> None: + torch.manual_seed(batch_size + 200) + cos, sin, _ = _rope_tables(256) + q = torch.randn( + batch_size, NUM_HEADS, HEAD_DIM, device=get_device(), dtype=torch.bfloat16 + ) + positions = ( + torch.arange(batch_size, device=get_device(), dtype=torch.int64) * 7 + ) % 63 + + q_fp4, q_sf = aiter_q_indexer_fp4(q.contiguous(), cos, sin, positions) + + ref_fp4, ref_sf = _ref_quantize_fp4_indexer( + _ref_q_transform(q, cos, sin, positions) + ) + torch.testing.assert_close( + _canonical_zero(q_fp4.view(torch.uint8).reshape(-1, FP4_DIM)), + _canonical_zero(ref_fp4), + ) + # Unlike the CUDA path, the scales are emitted already preshuffled into the + # logits kernel's ABI: heads split as (m_tiles, 16) and moved behind the + # group axis. Nothing downstream reorders them, and every candidate layout + # shares the tensor's shape, so a regression here is silent. + m_tiles, k_tiles = NUM_HEADS // 16, HEAD_DIM // 128 + expected_sf = ( + ref_sf.reshape(batch_size, m_tiles, 16, k_tiles, SCALE_GROUPS) + .permute(0, 3, 4, 2, 1) + .contiguous() + ) + torch.testing.assert_close(q_sf.reshape(expected_sf.shape), expected_sf) + + +# --------------------------------------------------------------------------- +# HIP-only surface: schedule bookkeeping and the paged MQA logits kernel +# --------------------------------------------------------------------------- + + +@pytest.mark.parametrize("logical_width", [1, 3, 4, 5, 8, 33]) +def test_guard_page_table_pads_to_schedule_granularity(logical_width: int) -> None: + rows = 3 + page_table = torch.arange( + 1, rows * logical_width + 1, device=get_device(), dtype=torch.int32 + ).reshape(rows, logical_width) + + guarded, max_seq_len = _guard_page_table(page_table) + + padded_width = max(4, -(-logical_width // 4) * 4) + assert guarded.shape == (rows, padded_width + 4) + assert guarded.dtype is torch.int32 + assert max_seq_len == padded_width * PAGE_SIZE + torch.testing.assert_close(guarded[:, :logical_width], page_table) + assert guarded[:, logical_width:].eq(0).all() + + +def test_guard_page_table_refreshes_reused_buffer() -> None: + rows, width = 2, 6 + first = torch.arange(rows * width, device=get_device(), dtype=torch.int32).reshape( + rows, width + ) + guarded, _ = _guard_page_table(first) + + refreshed, _ = _guard_page_table(first + 100, out=guarded) + + assert refreshed.data_ptr() == guarded.data_ptr() + torch.testing.assert_close(refreshed[:, :width], first + 100) + assert refreshed[:, width:].eq(0).all() + + +@pytest.mark.parametrize( + "num_queries,max_seq_len", [(1, 256), (8, 4096), (512, 256), (4096, 65536)] +) +def test_decode_cta_count_stays_within_available_chunks( + num_queries: int, max_seq_len: int +) -> None: + chunks_per_seq = max(1, -(-max_seq_len // 256)) + + cta_count = _decode_cta_count(num_queries, max_seq_len) + + assert 1 <= cta_count <= num_queries * chunks_per_seq + assert cta_count <= max(1024, num_queries * 4) + + +def _decode_plan(seq_lens: torch.Tensor, compress_ratio: int) -> CompressorDecodePlan: + """Hand-build the 16-byte decode plan rows the metadata builder reads.""" + words = torch.zeros((seq_lens.shape[0], 4), dtype=torch.int32, device=get_device()) + words[:, 0] = seq_lens.to(torch.int32) + return CompressorDecodePlan(compress_ratio, words.view(torch.uint8)) + + +def _prefill_plan( + seq_lens: torch.Tensor, ragged_ids: torch.Tensor, compress_ratio: int +) -> CompressorPrefillPlan: + words = torch.zeros((seq_lens.shape[0], 4), dtype=torch.int32, device=get_device()) + words[:, 0] = seq_lens.to(torch.int32) + words[:, 1] = ragged_ids.to(torch.int32) + return CompressorPrefillPlan( + compress_ratio, + words.view(torch.uint8), + torch.zeros((seq_lens.shape[0], 8), dtype=torch.uint8, device=get_device()), + ) + + +def test_k_write_metadata_decode_masks_unaligned_and_out_of_range() -> None: + compress_ratio, rope_len = 4, 4096 + seq_lens = torch.tensor( + [8, 9, 0, rope_len + compress_ratio], device=get_device(), dtype=torch.int64 + ) + out_loc = torch.tensor([11, 22, 33, 44], device=get_device(), dtype=torch.int64) + + meta = prepare_fp4_k_write_metadata( + _decode_plan(seq_lens, compress_ratio), out_loc, rope_len + ) + + # Row 0 is the only aligned, in-range row; 1 is unaligned, 2 has a negative + # RoPE position and 3 runs past the table. Only the slot mask has to cover + # all three: an out-of-range position is additionally clamped to 0 to keep + # the RoPE gather in bounds, but an unaligned row keeps its position and is + # dropped by its -1 slot alone. + torch.testing.assert_close( + meta.slots, + torch.tensor([11, -1, -1, -1], device=get_device(), dtype=torch.int64), + ) + torch.testing.assert_close( + meta.positions, + torch.tensor([4, 5, 0, 0], device=get_device(), dtype=torch.int64), + ) + + +def test_k_write_metadata_prefill_gathers_ragged_slots() -> None: + compress_ratio = 4 + seq_lens = torch.tensor([4, 8, 12, 16], device=get_device(), dtype=torch.int64) + ragged_ids = torch.tensor([2, 0, 5, 9], device=get_device(), dtype=torch.int64) + out_loc = torch.arange(6, device=get_device(), dtype=torch.int64) * 7 + + meta = prepare_fp4_k_write_metadata( + _prefill_plan(seq_lens, ragged_ids, compress_ratio), out_loc, 4096 + ) + + # ragged_id 9 is past the end of out_loc and must be dropped, not clamped. + torch.testing.assert_close( + meta.slots, + torch.tensor([14, 0, 35, -1], device=get_device(), dtype=torch.int64), + ) + torch.testing.assert_close(meta.positions, seq_lens - compress_ratio) + + +def test_k_write_metadata_prefill_with_empty_out_loc_writes_nothing() -> None: + seq_lens = torch.tensor([4, 8, 12], device=get_device(), dtype=torch.int64) + plan = _prefill_plan(seq_lens, torch.zeros_like(seq_lens), 4) + + meta = prepare_fp4_k_write_metadata( + plan, torch.empty(0, device=get_device(), dtype=torch.int64), 4096 + ) + + assert meta.slots.eq(-1).all() + + +def _build_logits_case( + batch: int, + seq_len: int, + *, + ctx_lens: list[int] | None = None, + shuffle_pages: bool = False, +): + """Populate an FP4 K cache and quantized Q for one synthetic indexer step. + + Every slot of every page is written, including the tail past ``ctx_lens``, + so a last block that is only partly in context still has live neighbours + that must not leak into the scored range. + """ + pages_per_seq = -(-seq_len // PAGE_SIZE) + padded_len = pages_per_seq * PAGE_SIZE + num_pages = batch * pages_per_seq + cos, sin, _ = _rope_tables(max(padded_len, 256)) + payload, scale = _empty_index_k_cache(num_pages) + + physical = ( + torch.randperm(num_pages, device=get_device()) + if shuffle_pages + else torch.arange(num_pages, device=get_device()) + ) + page_table = physical.to(torch.int32).reshape(batch, pages_per_seq) + context = torch.tensor( + ctx_lens if ctx_lens is not None else [seq_len] * batch, + device=get_device(), + dtype=torch.int32, + ) + + kv_positions = ( + torch.arange(padded_len, device=get_device(), dtype=torch.int64) + .repeat(batch) + .reshape(batch, padded_len) + ) + loc = ( + page_table.long()[:, :, None] * PAGE_SIZE + + torch.arange(PAGE_SIZE, device=get_device(), dtype=torch.int64)[None, None, :] + ).reshape(batch, padded_len) + + k = torch.randn( + batch, padded_len, HEAD_DIM, device=get_device(), dtype=torch.bfloat16 + ) + norm_weight = torch.randn(HEAD_DIM, device=get_device(), dtype=torch.bfloat16) + _write_index_k_cache( + k.reshape(-1, HEAD_DIM), + norm_weight, + cos, + sin, + kv_positions.reshape(-1), + loc.reshape(-1), + payload, + scale, + ) + k_ref = _ref_k_transform( + k.reshape(-1, HEAD_DIM), norm_weight, cos, sin, kv_positions.reshape(-1) + ).reshape(batch, padded_len, HEAD_DIM) + + q = torch.randn( + batch, NUM_HEADS, HEAD_DIM, device=get_device(), dtype=torch.bfloat16 + ) + q_positions = (context.long() - 1).clamp_min(0) + q_fp4, q_scale = aiter_q_indexer_fp4(q.contiguous(), cos, sin, q_positions) + q_ref = _ref_q_transform(q, cos, sin, q_positions) + + # ``C4Indexer.compute_weights`` runs a bf16 projection and the adapter + # forwards the result unconverted, so the kernel is fed bf16 weights. + weights = torch.randn(batch, NUM_HEADS, device=get_device(), dtype=torch.bfloat16) + weight_scale = HEAD_DIM**-0.5 * NUM_HEADS**-0.5 + + q_dq = _ref_dequantize_fp4_indexer( + q_fp4.view(torch.uint8), _ref_quantize_fp4_indexer(q_ref)[1] + ).reshape(batch, NUM_HEADS, HEAD_DIM) + k_dq = _ref_dequantize_fp4_indexer( + *_read_index_k_cache(payload, scale, loc.reshape(-1)) + ).reshape(batch, padded_len, HEAD_DIM) + + def _logits_from(q_src, k_src) -> torch.Tensor: + # The indexer scores each head separately, clamps it at zero and only + # then takes the weighted sum; dropping the ReLU changes the result. + per_head = torch.einsum("qhd,qsd->qhs", q_src.float(), k_src.float()) + return weight_scale * torch.einsum( + "qhs,qh->qs", per_head.relu(), weights.float() + ) + + return { + "q_fp4": q_fp4, + "q_scale": q_scale, + "payload": payload, + "scale": scale, + "weights": weights, + "weight_scale": weight_scale, + "page_table": page_table, + "c4_seq_lens": context, + # Against exactly the FP4 operands the kernel read, and against the + # unquantized bf16 model. + "ref_logits_fp4": _logits_from(q_dq, k_dq), + "ref_logits_bf16": _logits_from(q_ref, k_ref), + "context": context, + "seq_len": seq_len, + } + + +def _run_logits(case, *, is_decode: bool, decode_ws=None, prefill_ws=None): + return aiter_fp4_paged_mqa_logits( + q_fp4=case["q_fp4"], + q_scale=case["q_scale"], + k_payload=case["payload"], + k_scale=case["scale"], + weights=case["weights"], + page_table=case["page_table"], + c4_seq_lens=case["c4_seq_lens"], + weight_scale=case["weight_scale"], + is_decode=is_decode, + decode_workspace=decode_ws, + prefill_workspace=prefill_ws, + ) + + +def _assert_logits_agree(logits: torch.Tensor, case) -> None: + """Check each row over its own context, against FP4 operands and bf16. + + The FP4 comparison is the tight one: the reference is fed exactly what the + kernel read, so only reduction order differs. The bf16 comparison is a + coarse guard that FP4 has not disturbed the ranking the indexer is about to + top-k; a dropped ReLU or a mispaired scale lands near 0.7 there, well clear + of the quantization noise floor. Positions past a row's context are left + undefined by design, so they are never compared. + """ + for row, ctx in enumerate(case["context"].tolist()): + if ctx == 0: + # A padded row owns no valid position; all it must do is leave the + # rest of the batch alone, which the other iterations cover. + continue + got = logits[row, :ctx] + exact = case["ref_logits_fp4"][row, :ctx] + bf16 = case["ref_logits_bf16"][row, :ctx] + torch.testing.assert_close( + got, exact, rtol=2.0e-3, atol=2.0e-3, msg=f"row {row} (ctx={ctx})" + ) + cosine = torch.nn.functional.cosine_similarity(got, bf16, dim=-1).item() + assert cosine > 0.95, f"row {row} (ctx={ctx}) cosine vs bf16 is {cosine:.4f}" + topk = min(64, ctx) + overlap = ( + len( + set(got.topk(topk).indices.tolist()) + & set(bf16.topk(topk).indices.tolist()) + ) + / topk + ) + assert overlap > 0.75, f"row {row} top-{topk} overlap is {overlap:.3f}" + + +@pytest.mark.parametrize("batch,seq_len", [(1, 256), (2, 384), (4, 512)]) +def test_decode_paged_mqa_logits(batch: int, seq_len: int) -> None: + torch.manual_seed(batch * 100 + seq_len) + case = _build_logits_case(batch, seq_len) + workspace = prepare_fp4_decode_workspace(case["page_table"], case["c4_seq_lens"]) + + logits = _run_logits(case, is_decode=True, decode_ws=workspace) + + _assert_logits_agree(logits, case) + + +@pytest.mark.parametrize("batch,seq_len", [(1, 256), (3, 512)]) +def test_prefill_paged_mqa_logits(batch: int, seq_len: int) -> None: + torch.manual_seed(batch * 200 + seq_len) + case = _build_logits_case(batch, seq_len) + workspace = prepare_fp4_prefill_workspace(case["page_table"], case["c4_seq_lens"]) + + logits = _run_logits(case, is_decode=False, prefill_ws=workspace) + + _assert_logits_agree(logits, case) + + +@pytest.mark.parametrize("is_decode", [True, False]) +def test_logits_with_ragged_context_lengths(is_decode: bool) -> None: + """Sizing the persistent grid for uneven contexts is the scheduler's job. + + A uniform batch hides an unbalanced chunk assignment: every row gets the + same number of KV chunks, so an off-by-one in the split still covers each + row exactly once. + """ + torch.manual_seed(21 if is_decode else 22) + # Deliberately not multiples of the 64-token page or the 256-token chunk, + # so the last block of each row is only partly in context. + # A 0 stands for a padded row: ``match_num_queries`` pads c4_seq_lens with + # 0 on the FP4 path, and such a row must not disturb its neighbours. + ctx_lens = [17, 512, 1, 300, 0, 64, 129, 511] + case = _build_logits_case(len(ctx_lens), 512, ctx_lens=ctx_lens) + if is_decode: + ws = prepare_fp4_decode_workspace(case["page_table"], case["c4_seq_lens"]) + logits = _run_logits(case, is_decode=True, decode_ws=ws) + else: + ws = prepare_fp4_prefill_workspace(case["page_table"], case["c4_seq_lens"]) + logits = _run_logits(case, is_decode=False, prefill_ws=ws) + + _assert_logits_agree(logits, case) + + +@pytest.mark.parametrize("is_decode", [True, False]) +def test_logits_follow_shuffled_page_table(is_decode: bool) -> None: + """Pages of a sequence are neither contiguous nor ordered under a radix cache.""" + torch.manual_seed(31 if is_decode else 32) + case = _build_logits_case(4, 384, ctx_lens=[384, 300, 129, 384], shuffle_pages=True) + if is_decode: + ws = prepare_fp4_decode_workspace(case["page_table"], case["c4_seq_lens"]) + logits = _run_logits(case, is_decode=True, decode_ws=ws) + else: + ws = prepare_fp4_prefill_workspace(case["page_table"], case["c4_seq_lens"]) + logits = _run_logits(case, is_decode=False, prefill_ws=ws) + + _assert_logits_agree(logits, case) + + +@pytest.mark.parametrize("is_decode", [True, False]) +def test_pinned_schedule_matches_unpinned_logits(is_decode: bool) -> None: + """A pinned workspace only preplans the grid; the logits must not move.""" + torch.manual_seed(11 if is_decode else 12) + case = _build_logits_case(2, 384) + if is_decode: + workspace = prepare_fp4_decode_workspace( + case["page_table"], case["c4_seq_lens"] + ) + pinned = _run_logits(case, is_decode=True, decode_ws=workspace) + else: + workspace = prepare_fp4_prefill_workspace( + case["page_table"], case["c4_seq_lens"] + ) + pinned = _run_logits(case, is_decode=False, prefill_ws=workspace) + unpinned = _run_logits(case, is_decode=is_decode) + + seq_len = case["seq_len"] + torch.testing.assert_close(pinned[:, :seq_len], unpinned[:, :seq_len]) + + +def test_stale_workspace_row_count_falls_back_to_inline_schedule() -> None: + """DP padding can leave a workspace sized for a different row count.""" + torch.manual_seed(13) + case = _build_logits_case(2, 256) + stale = prepare_fp4_decode_workspace( + case["page_table"][:1], case["c4_seq_lens"][:1] + ) + + with_stale = _run_logits(case, is_decode=True, decode_ws=stale) + without = _run_logits(case, is_decode=True) + + seq_len = case["seq_len"] + torch.testing.assert_close(with_stale[:, :seq_len], without[:, :seq_len]) + + +if __name__ == "__main__": + sys.exit(pytest.main([__file__, "-v"])) diff --git a/test/registered/unit/mem_cache/test_mem_pool_host.py b/test/registered/unit/mem_cache/test_mem_pool_host.py index 2352e0127..61c4a0931 100644 --- a/test/registered/unit/mem_cache/test_mem_pool_host.py +++ b/test/registered/unit/mem_cache/test_mem_pool_host.py @@ -136,6 +136,17 @@ class TestLazyHostPoolRelease(CustomTestCase): def _make_logical_pool(): return LogicalHostPool(size=8, page_size=2) + @staticmethod + def _make_transfer_pool(*, page_aligned_only): + pool = DeepSeekV4PagedHostPool.__new__(DeepSeekV4PagedHostPool) + pool.pool_name = str(PoolName.DEEPSEEK_V4_C4_INDEXER) + pool.slot_page_size = 4 + pool.layer_num = 1 + pool.page_aligned_only = page_aligned_only + pool.device_ptrs = [0] + pool.data_ptrs = [0] + return pool + def _assert_lazy_release(self, pool): self.assertEqual(pool.free(torch.empty(0, dtype=torch.int64)), 0) self.assertEqual(pool.num_release_slots, 0) @@ -190,6 +201,26 @@ class TestLazyHostPoolRelease(CustomTestCase): pool.clear() self.assertEqual(len(pool.alloc(1)), 2) + def test_grouped_page_rows_reject_unaligned_transfers(self): + # FP4 indexer rows group their slots, so a partial page has no + # well-defined token-granular copy and must not silently fall back. + pool = self._make_transfer_pool(page_aligned_only=True) + unaligned = torch.arange(3, dtype=torch.int64) + with self.assertRaisesRegex(ValueError, "page-aligned"): + pool.backup_from_device_all_layer(None, unaligned, unaligned, "direct") + with self.assertRaisesRegex(ValueError, "page-aligned"): + pool.load_to_device_per_layer(None, unaligned, unaligned, 0, "direct") + + def test_fused_page_rows_keep_token_granular_transfers(self): + pool = self._make_transfer_pool(page_aligned_only=False) + unaligned = torch.arange(3, dtype=torch.int64) + with unittest.mock.patch( + "sglang.srt.mem_cache.memory_pool_host.transfer_cache_dsv4_mla" + ) as transfer: + pool.backup_from_device_all_layer(None, unaligned, unaligned, "direct") + pool.load_to_device_per_layer(None, unaligned, unaligned, 0, "direct") + self.assertEqual(transfer.call_count, 2) + def test_logical_pool_lazy_release(self): pool = self._make_logical_pool() self._assert_lazy_release(pool)