From bbaca61f6e230aaf167b51a1f0e218a2e9417f1e Mon Sep 17 00:00:00 2001 From: AMD-yanfeiwang Date: Sun, 6 Sep 2026 08:09:21 +0000 Subject: [PATCH 1/2] [AMD][DSv4] Add experimental AITER FP4 streaming TopK --- .../docs/references/environment_variables.mdx | 5 + .../ops/attention/dsv4/fp4_indexer_hip.py | 169 +++++++++++++++++- .../dsv4/fp4_indexer_schedule_hip.py | 36 ++-- python/sglang/srt/environ.py | 4 + .../deepseek_v4_backend_hip_radix.py | 40 ++++- .../srt/layers/attention/dsv4/indexer.py | 90 +++++++--- 6 files changed, 305 insertions(+), 39 deletions(-) diff --git a/docs/docs/references/environment_variables.mdx b/docs/docs/references/environment_variables.mdx index 99bf8c935bb8..c88406ffc5e4 100644 --- a/docs/docs/references/environment_variables.mdx +++ b/docs/docs/references/environment_variables.mdx @@ -710,6 +710,11 @@ SGLang supports various environment variables that can be used to configure its Whether DeepSeek V4 experts use FP4. Set to false when using an FP4-to-FP8 converted DeepSeek V4 checkpoint. true + + SGLANG_DSV4_FP4_FUSED_TOPK + Experimental: on gfx950 eager prefill with the SGL top-k backend, use AITER's streaming FP4 indexer score and top-k operator when supported, avoiding the full FP32 logits tensor. Decode and graph capture retain the legacy path. The current prototype trades substantial prefill performance for lower scratch memory. + false + SGLANG_DSV4_REASONING_EFFORT Default reasoning_effort for the DeepSeek V4 chat encoder when a request does not set it. The preview profile accepts high and max; the official profile accepts low, high, and max. The profile is detected from the bundled encoder. Override it with --json-model-override-args '{"dsv4_reasoning_effort_profile":"official"}'. diff --git a/python/sglang/kernels/ops/attention/dsv4/fp4_indexer_hip.py b/python/sglang/kernels/ops/attention/dsv4/fp4_indexer_hip.py index 3982898a309a..aa45640f1eff 100644 --- a/python/sglang/kernels/ops/attention/dsv4/fp4_indexer_hip.py +++ b/python/sglang/kernels/ops/attention/dsv4/fp4_indexer_hip.py @@ -3,7 +3,8 @@ from __future__ import annotations import os -from typing import TYPE_CHECKING, NamedTuple, Optional, Tuple, Union +from functools import lru_cache +from typing import TYPE_CHECKING, Any, Callable, NamedTuple, Optional, Tuple, Union import torch @@ -41,6 +42,29 @@ _LOGITS_POOL: dict = {} +@lru_cache(maxsize=1) +def get_aiter_fp4_streaming_topk() -> Optional[Callable]: + """Resolve the complete optional AITER API once, before graph capture.""" + try: + from aiter.ops.flydsl import ( + FP4PrefillTopKResult, + allocate_fp4_prefill_topk_workspace, + flydsl_pa_mqa_topk_fp4_prefill, + ) + except (AttributeError, ImportError): + return None + required = ( + flydsl_pa_mqa_topk_fp4_prefill, + allocate_fp4_prefill_topk_workspace, + FP4PrefillTopKResult, + ) + return flydsl_pa_mqa_topk_fp4_prefill if all(map(callable, required)) else None + + +def aiter_fp4_streaming_topk_available() -> bool: + return get_aiter_fp4_streaming_topk() is not None + + class FP4DecodeWorkspace(NamedTuple): guarded_page_table: torch.Tensor c4_seq_lens: torch.Tensor @@ -65,6 +89,17 @@ class FP4PrefillWorkspace(NamedTuple): schedule_buffers: Optional[PrefillScheduleBuffers] = None +class FP4StreamingTopKScratch(NamedTuple): + """Bounded AITER candidate scratch reused across C4 layers on one stream.""" + + workspace: Any + values: torch.Tensor + raw_indices: torch.Tensor + counts: torch.Tensor + parallel_unit_num: int + stream_id: int + + class FP4KWriteMetadata(NamedTuple): positions: torch.Tensor slots: torch.Tensor @@ -284,6 +319,52 @@ def prepare_fp4_prefill_workspace( return workspace +def prepare_fp4_streaming_topk_scratch( + *, + rows: int, + topk: int, + device: torch.device, + scratch: Optional[FP4StreamingTopKScratch] = None, +) -> FP4StreamingTopKScratch: + """Create or reuse context-independent AITER TopK candidate scratch.""" + from aiter.ops.flydsl import allocate_fp4_prefill_topk_workspace + + stream = torch.cuda.current_stream(device) + stream_id = stream.cuda_stream + parallel_unit_num = max(512, rows) + if ( + scratch is not None + and scratch.values.shape == (rows, topk) + and scratch.values.device == device + and scratch.parallel_unit_num == parallel_unit_num + and scratch.stream_id == stream_id + ): + return scratch + + with torch.cuda.stream(stream): + return FP4StreamingTopKScratch( + workspace=allocate_fp4_prefill_topk_workspace( + rows, + parallel_unit_num, + topk, + device, + ), + values=torch.empty( + (rows, topk), + dtype=torch.float32, + device=device, + ), + raw_indices=torch.empty( + (rows, topk), + dtype=torch.int32, + device=device, + ), + counts=torch.empty(rows, dtype=torch.int32, device=device), + parallel_unit_num=parallel_unit_num, + stream_id=stream_id, + ) + + def aiter_fp4_paged_mqa_logits( *, q_fp4: torch.Tensor, @@ -422,6 +503,92 @@ def aiter_fp4_paged_mqa_logits( return logits +def aiter_fp4_paged_mqa_topk( + *, + 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, + out_page_indices: torch.Tensor, + out_raw_indices: Optional[torch.Tensor] = None, + prefill_workspace: Optional[FP4PrefillWorkspace] = None, + streaming_scratch: Optional[FP4StreamingTopKScratch] = None, +) -> FP4StreamingTopKScratch: + """Run AITER's bounded FP4 score + exact TopK prefill operator.""" + fused_topk = get_aiter_fp4_streaming_topk() + if fused_topk is None: + raise RuntimeError("The installed AITER does not expose FP4 streaming top-k") + + from aiter.ops.flydsl import FP4PrefillTopKResult + + num_tokens = q_fp4.shape[0] + c4_seq_lens = _as_int32_1d(c4_seq_lens) + expected_width = _guarded_pages(page_table.shape[1]) * _KV_BLOCK_SIZE + workspace = prefill_workspace + if workspace is not None and ( + workspace.guarded_page_table.shape[0] != num_tokens + or workspace.max_seq_len != expected_width + ): + workspace = None + + if workspace is None: + guarded_page_table, max_seq_len = _guard_page_table(page_table) + row_to_batch = torch.arange( + num_tokens, + dtype=torch.int32, + device=page_table.device, + ) + local_starts = torch.zeros_like(c4_seq_lens) + else: + guarded_page_table = workspace.guarded_page_table + max_seq_len = workspace.max_seq_len + row_to_batch = workspace.row_to_batch + local_starts = workspace.local_starts + + streaming_scratch = prepare_fp4_streaming_topk_scratch( + rows=num_tokens, + topk=out_page_indices.shape[1], + device=q_fp4.device, + scratch=streaming_scratch, + ) + out = FP4PrefillTopKResult( + values=streaming_scratch.values, + raw_indices=( + out_raw_indices + if out_raw_indices is not None + else streaming_scratch.raw_indices + ), + physical_indices=out_page_indices, + counts=streaming_scratch.counts, + ) + fused_topk( + q_fp4.view(torch.uint8), + q_scale, + k_payload.view(torch.uint8), + k_scale, + guarded_page_table, + weights, + row_to_batch, + local_starts, + c4_seq_lens, + max_seq_len, + topk=out_page_indices.shape[1], + weight_scale=weight_scale, + block_k=256, + kv_block_size=_KV_BLOCK_SIZE, + num_warps=4, + parallel_unit_num=streaming_scratch.parallel_unit_num, + workspace=streaming_scratch.workspace, + out=out, + stream=torch.cuda.current_stream(q_fp4.device), + ) + return streaming_scratch + + def prepare_fp4_k_write_metadata( plan: Union[CompressorDecodePlan, CompressorPrefillPlan], out_loc: torch.Tensor, diff --git a/python/sglang/kernels/ops/attention/dsv4/fp4_indexer_schedule_hip.py b/python/sglang/kernels/ops/attention/dsv4/fp4_indexer_schedule_hip.py index cd930dd373f1..44e4725bbf84 100644 --- a/python/sglang/kernels/ops/attention/dsv4/fp4_indexer_schedule_hip.py +++ b/python/sglang/kernels/ops/attention/dsv4/fp4_indexer_schedule_hip.py @@ -93,9 +93,10 @@ def _prefill_schedule_prep_kernel( off = tl.arange(0, BLOCK_T) live = off < T le = tl.load(le_ptr + off, mask=live, other=0) + le = tl.minimum(tl.maximum(le, 0), s_max * BLOCK_K) # ceil(le / block_k); a non-positive length contributes no chunks, which # matches the reference's clamp(floor_div(le + block_k - 1), min=0). - chunks = tl.where(live, (tl.maximum(le, 0) + (BLOCK_K - 1)) // BLOCK_K, 0) + chunks = tl.where(live, (le + (BLOCK_K - 1)) // BLOCK_K, 0) # A split factor s fits when the persistent grid can host every # (row, split) pair: sum_i ceil(chunks_i / s) <= P. ceil(c/s) is @@ -348,18 +349,33 @@ def build_prefill_schedule( ) BLOCK_P = 256 - _prefill_cta_info_kernel[(triton.cdiv(parallel_unit_num, BLOCK_P),)]( + cta_args = [ buffers.incl, buffers.excl, buffers.chunks, - buffers.row_to_batch, - buffers.local_starts, - local_ends, - buffers.safe, - buffers.total_splits, - cta_info_out, - total_rows, - parallel_unit_num, + ] + # AITER's window-aware scheduler added first_chunks_ptr before rb_ptr. + # SGLang's fused prep always emits local_starts == 0, so that same zero + # buffer is the exact first-chunk vector. Keep the old signature working + # for released AITER builds. + if "first_chunks_ptr" in getattr(_prefill_cta_info_kernel, "arg_names", ()): + cta_args.append(buffers.local_starts) + cta_args.extend( + [ + buffers.row_to_batch, + buffers.local_starts, + local_ends, + buffers.safe, + buffers.total_splits, + cta_info_out, + total_rows, + parallel_unit_num, + ] + ) + if "max_seq_len" in getattr(_prefill_cta_info_kernel, "arg_names", ()): + cta_args.append(max_seq_len) + _prefill_cta_info_kernel[(triton.cdiv(parallel_unit_num, BLOCK_P),)]( + *cta_args, BLOCK_P=BLOCK_P, ) return guarded_out, buffers diff --git a/python/sglang/srt/environ.py b/python/sglang/srt/environ.py index e06ed21a2b3f..54f9bbe72278 100644 --- a/python/sglang/srt/environ.py +++ b/python/sglang/srt/environ.py @@ -889,6 +889,10 @@ class Envs: # go back to the unfused chain on the verify path. SGLANG_OPT_FUSED_QK_NORM_ROPE_VERIFY = EnvBool(True) SGLANG_OPT_USE_AITER_INDEXER = EnvBool(False) + # Experimental opt-in for AITER's FP4 paged-MQA streaming top-k. The legacy + # materialized-logits path remains the default until the fused prefill + # performance gap is closed. + SGLANG_DSV4_FP4_FUSED_TOPK = EnvBool(False) # =================================================================== # Apple Silicon and MLX 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 3f3fde172960..2d36861e54fb 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 @@ -23,6 +23,7 @@ ) from sglang.srt.environ import envs from sglang.srt.layers.attention.base_attn_backend import AttentionBackend +from sglang.srt.layers.attention.dsa.dsa_topk_backend import DSATopKBackend from sglang.srt.layers.attention.dsv4.compressor_v2 import ( CompressorBackendMixin, FusedCompressMetadata, @@ -51,6 +52,7 @@ FP4DecodeWorkspace, FP4KWriteMetadata, FP4PrefillWorkspace, + FP4StreamingTopKScratch, ) from sglang.srt.layers.radix_attention import RadixAttention from sglang.srt.model_executor.model_runner import ModelRunner @@ -365,6 +367,9 @@ class DSV4Metadata: fp4_prefill_workspace: Optional[FP4PrefillWorkspace] = field( default=None, repr=False ) + fp4_streaming_topk_scratch: Optional[FP4StreamingTopKScratch] = 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 @@ -449,6 +454,7 @@ def __init__( speculative_num_steps=0, ): super().__init__() + self.dsa_topk_backend = DSATopKBackend.resolve(model_runner) self.device = torch.device(model_runner.device) head_dim = model_runner.model_config.head_dim assert head_dim == 512, ( @@ -476,6 +482,22 @@ def __init__( self.enable_deepseek_v4_fp4_indexer: bool = ( model_runner.server_args.enable_deepseek_v4_fp4_indexer ) + self.use_aiter_fp4_streaming_topk = False + if ( + self.enable_deepseek_v4_fp4_indexer + and envs.SGLANG_DSV4_FP4_FUSED_TOPK.get() + and self.dsa_topk_backend.is_sgl_kernel() + ): + from sglang.kernels.ops.attention.dsv4.fp4_indexer_hip import ( + aiter_fp4_streaming_topk_available, + ) + + self.use_aiter_fp4_streaming_topk = aiter_fp4_streaming_topk_available() + if self.enable_deepseek_v4_fp4_indexer: + logger.info( + "AITER FP4 streaming top-k: %s", + "enabled" if self.use_aiter_fp4_streaming_topk else "legacy fallback", + ) self.topk = get_spec().speculative_eagle_topk or 0 assert self.topk in [0, 1], "MTP Topk > 1 not supported for DeepSeek V4" self.mtp_enabled = self.topk > 0 @@ -504,12 +526,18 @@ def _move_to_device(self, x: List[int]) -> torch.Tensor: pin_tensor = torch.tensor(x, dtype=torch.int32, pin_memory=True) return pin_tensor.to(self.device, non_blocking=True) - def init_forward_metadata_indexer(self, core_attn_metadata: DSV4AttnMetadata): + def init_forward_metadata_indexer( + self, + core_attn_metadata: DSV4AttnMetadata, + *, + use_prefill_cuda_graph: bool = False, + ): return PagedIndexerMetadata( page_size=self.page_size, page_table=core_attn_metadata.page_table, c4_seq_lens=core_attn_metadata.c4_topk_lengths_raw, use_topk_v2=self.dsa_topk_backend.should_use_topk_v2(), + use_prefill_cuda_graph=use_prefill_cuda_graph, ) def init_forward_metadata_decode( @@ -596,7 +624,10 @@ def init_forward_metadata_prefill( core_attn_metadata, req_pool_indices_repeated ) indexer_metadata = ( - self.init_forward_metadata_indexer(core_attn_metadata) + self.init_forward_metadata_indexer( + core_attn_metadata, + use_prefill_cuda_graph=use_prefill_cuda_graph, + ) if need_compress else None ) @@ -752,7 +783,10 @@ def make_forward_metadata_from_raw_verify( out_loc=out_cache_loc, need_compress=True, ) - indexer_metadata = self.init_forward_metadata_indexer(core_attn_metadata) + indexer_metadata = self.init_forward_metadata_indexer( + core_attn_metadata, + use_prefill_cuda_graph=True, + ) create = functools.partial( create_paged_compressor_data, is_prefill=True, diff --git a/python/sglang/srt/layers/attention/dsv4/indexer.py b/python/sglang/srt/layers/attention/dsv4/indexer.py index 1be1910403a2..061bbde6f262 100644 --- a/python/sglang/srt/layers/attention/dsv4/indexer.py +++ b/python/sglang/srt/layers/attention/dsv4/indexer.py @@ -25,6 +25,7 @@ ) from sglang.kernels.ops.attention.dsv4.fp4_indexer_hip import ( aiter_fp4_paged_mqa_logits, + aiter_fp4_paged_mqa_topk, aiter_q_indexer_fp4, logits_rows_per_chunk, ) @@ -41,6 +42,7 @@ ) from sglang.srt.layers.linear import ReplicatedLinear from sglang.srt.model_executor.forward_batch_info import ForwardMode +from sglang.srt.model_executor.runner import get_is_capture_mode from sglang.srt.model_executor.runner_backend_utils.breakable_cuda_graph.context import ( is_in_breakable_cuda_graph, ) @@ -74,6 +76,22 @@ _arange_cache = {} +def _should_use_aiter_fp4_streaming_topk( + forward_mode: ForwardMode, + enabled: bool, + use_prefill_cuda_graph: bool, +) -> bool: + return ( + not forward_mode.is_decode() + and enabled + and not use_prefill_cuda_graph + and not get_is_capture_mode() + and not is_in_tc_piecewise_cuda_graph() + and not is_in_breakable_cuda_graph() + and not torch.cuda.is_current_stream_capturing() + ) + + def fp8_paged_mqa_logits_torch( q_fp8: torch.Tensor, kvcache_fp8: torch.Tensor, @@ -906,6 +924,11 @@ def run_topk_transform(rows: slice, logits: torch.Tensor) -> None: elif use_aiter_fp4: q_fp4, q_scale = q is_decode = forward_batch.forward_mode.is_decode() + use_streaming_topk = _should_use_aiter_fp4_streaming_topk( + forward_batch.forward_mode, + getattr(self, "use_aiter_fp4_streaming_topk", False), + indexer_metadata.use_prefill_cuda_graph, + ) # Hoisted: these await this layer's KV transfer, which every chunk # would otherwise re-await. k_payload = token_to_kv_pool.get_index_k_fp4_payload_buffer( @@ -913,38 +936,55 @@ def run_topk_transform(rows: slice, logits: torch.Tensor) -> None: ) k_scale = token_to_kv_pool.get_index_k_fp4_scale_buffer(c4_indexer.layer_id) - def run_fp4_indexer(rows: slice) -> None: - logits = aiter_fp4_paged_mqa_logits( - q_fp4=q_fp4[rows], - q_scale=q_scale[rows], + if use_streaming_topk: + metadata.fp4_streaming_topk_scratch = aiter_fp4_paged_mqa_topk( + q_fp4=q_fp4, + q_scale=q_scale, k_payload=k_payload, k_scale=k_scale, - weights=weights[rows], - page_table=page_table[rows], - c4_seq_lens=c4_seq_lens[rows], + weights=weights, + page_table=page_table, + c4_seq_lens=c4_seq_lens, weight_scale=c4_indexer.weight_scale, - is_decode=is_decode, - decode_workspace=metadata.fp4_decode_workspace, + out_page_indices=c4_sparse_page_indices, + out_raw_indices=raw_indices, prefill_workspace=metadata.fp4_prefill_workspace, + streaming_scratch=metadata.fp4_streaming_topk_scratch, ) - run_topk_transform(rows, logits) - - # The scores are the layer's largest transient and their width tracks - # context length, so prefill splits the rows into whatever fits the - # pooled logits block and reduces each chunk before the next one - # reuses it. Rows are scored and reduced independently, so this - # matches a single pass. Decode's rectangle is bounded by its capture - # shapes, so it always stays whole. - rows_per_chunk = ( - query_rows if is_decode else logits_rows_per_chunk(page_table) - ) - if rows_per_chunk >= query_rows: - run_fp4_indexer(all_rows) else: - for start in range(0, query_rows, max(1, rows_per_chunk)): - run_fp4_indexer( - slice(start, min(start + rows_per_chunk, query_rows)) + + def run_fp4_indexer(rows: slice) -> None: + logits = aiter_fp4_paged_mqa_logits( + q_fp4=q_fp4[rows], + q_scale=q_scale[rows], + k_payload=k_payload, + k_scale=k_scale, + weights=weights[rows], + page_table=page_table[rows], + c4_seq_lens=c4_seq_lens[rows], + weight_scale=c4_indexer.weight_scale, + is_decode=is_decode, + decode_workspace=metadata.fp4_decode_workspace, + prefill_workspace=metadata.fp4_prefill_workspace, ) + run_topk_transform(rows, logits) + + # The scores are the layer's largest transient and their width tracks + # context length, so prefill splits the rows into whatever fits the + # pooled logits block and reduces each chunk before the next one + # reuses it. Rows are scored and reduced independently, so this + # matches a single pass. Decode's rectangle is bounded by its capture + # shapes, so it always stays whole. + rows_per_chunk = ( + query_rows if is_decode else logits_rows_per_chunk(page_table) + ) + if rows_per_chunk >= query_rows: + run_fp4_indexer(all_rows) + else: + for start in range(0, query_rows, max(1, rows_per_chunk)): + run_fp4_indexer( + slice(start, min(start + rows_per_chunk, query_rows)) + ) else: c4_indexer_kv_cache = token_to_kv_pool.get_index_k_with_scale_buffer( layer_id=c4_indexer.layer_id, From 096074438671360c3c79eedde9d43a4e59cc2963 Mon Sep 17 00:00:00 2001 From: AMD-yanfeiwang Date: Sun, 6 Sep 2026 08:09:21 +0000 Subject: [PATCH 2/2] [AMD][DSv4] Cover AITER FP4 streaming TopK --- .../ops/attention/test_fp4_indexer_hip.py | 206 ++++++++++++++++++ .../layers/test_dsv4_fp4_streaming_topk.py | 167 ++++++++++++++ 2 files changed, 373 insertions(+) create mode 100644 test/registered/unit/layers/test_dsv4_fp4_streaming_topk.py diff --git a/test/registered/kernels/ops/attention/test_fp4_indexer_hip.py b/test/registered/kernels/ops/attention/test_fp4_indexer_hip.py index cf6aacb6ab39..71613d306910 100644 --- a/test/registered/kernels/ops/attention/test_fp4_indexer_hip.py +++ b/test/registered/kernels/ops/attention/test_fp4_indexer_hip.py @@ -16,10 +16,13 @@ from __future__ import annotations import sys +from types import SimpleNamespace +from unittest import mock import pytest import torch +from sglang.kernels.ops.attention.dsv4 import fp4_indexer_schedule_hip as schedule_hip from sglang.kernels.ops.attention.deepseek_v4_rope import precompute_freqs_cis from sglang.kernels.ops.attention.dsv4 import ( CompressorDecodePlan, @@ -31,6 +34,8 @@ _decode_cta_count, _guard_page_table, aiter_fp4_paged_mqa_logits, + aiter_fp4_paged_mqa_topk, + aiter_fp4_streaming_topk_available, aiter_k_indexer_fp4_cache_write, aiter_q_indexer_fp4, prepare_fp4_decode_workspace, @@ -400,6 +405,106 @@ def test_fp4_fused_q_indexer_rope_hadamard_quant(batch_size: int) -> None: # --------------------------------------------------------------------------- +_LEGACY_CTA_ARG_NAMES = ( + "incl_ptr", + "excl_ptr", + "chunks_ptr", + "rb_ptr", + "ls_ptr", + "le_ptr", + "safe_ptr", + "total_splits_ptr", + "cta_info_ptr", + "T", + "P", + "BLOCK_P", +) +_WINDOWED_CTA_ARG_NAMES = ( + *_LEGACY_CTA_ARG_NAMES[:3], + "first_chunks_ptr", + *_LEGACY_CTA_ARG_NAMES[3:-1], + "max_seq_len", + "BLOCK_P", +) + + +@pytest.mark.parametrize( + ("arg_names", "window_aware"), + [ + pytest.param(_LEGACY_CTA_ARG_NAMES, False, id="legacy-aiter"), + pytest.param(_WINDOWED_CTA_ARG_NAMES, True, id="window-aware-aiter"), + ], +) +def test_prefill_schedule_forwards_versioned_aiter_cta_args( + monkeypatch: pytest.MonkeyPatch, + arg_names: tuple[str, ...], + window_aware: bool, +) -> None: + from aiter.ops.flydsl.kernels.mqa_logits import ( + pa_mqa_logits_fp4_prefill as aiter_prefill, + ) + + prep_kernel = mock.MagicMock() + cta_kernel = mock.MagicMock() + cta_kernel.arg_names = arg_names + monkeypatch.setattr(schedule_hip, "_prefill_schedule_prep_kernel", prep_kernel) + monkeypatch.setattr(aiter_prefill, "_prefill_cta_info_kernel", cta_kernel) + + buffers = SimpleNamespace( + fused=True, + incl=object(), + excl=object(), + chunks=object(), + row_to_batch=object(), + local_starts=object(), + safe=object(), + total_splits=object(), + ) + page_table = torch.zeros((1, 1), dtype=torch.int32) + local_ends = torch.tensor([64], dtype=torch.int32) + guarded_out = torch.empty((1, 8), dtype=torch.int32) + cta_info_out = torch.empty((257, 6), dtype=torch.int32) + max_seq_len = 8192 + + schedule_hip.build_prefill_schedule( + page_table=page_table, + local_ends=local_ends, + cta_info_out=cta_info_out, + parallel_unit_num=257, + max_seq_len=max_seq_len, + guarded_out=guarded_out, + buffers=buffers, + ) + + expected = [buffers.incl, buffers.excl, buffers.chunks] + if window_aware: + expected.append(buffers.local_starts) + expected.extend( + [ + buffers.row_to_batch, + buffers.local_starts, + local_ends, + buffers.safe, + buffers.total_splits, + cta_info_out, + 1, + 257, + ] + ) + if window_aware: + expected.append(max_seq_len) + + launch = cta_kernel.__getitem__.return_value + launch.assert_called_once() + assert launch.call_args.kwargs == {"BLOCK_P": 256} + assert len(launch.call_args.args) == len(expected) + for actual, wanted in zip(launch.call_args.args, expected, strict=True): + if isinstance(wanted, int): + assert actual == wanted + else: + assert actual is wanted + + @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 @@ -825,5 +930,106 @@ def test_row_chunks_reproduce_the_unsplit_batch() -> None: ) +@pytest.mark.parametrize("topk", [512, 1024]) +def test_streaming_topk_writes_raw_and_physical_indices(topk: int) -> None: + if not aiter_fp4_streaming_topk_available(): + pytest.skip("installed AITER does not expose FP4 streaming top-k") + + torch.manual_seed(41 + topk) + contexts = [0, topk, topk + 64] + case = _build_logits_case( + len(contexts), + topk + 64, + ctx_lens=contexts, + shuffle_pages=True, + ) + out_page_indices = torch.empty( + (len(contexts), topk), + dtype=torch.int32, + device=get_device(), + ) + out_raw_indices = torch.empty_like(out_page_indices) + workspace = prepare_fp4_prefill_workspace( + case["page_table"], + case["c4_seq_lens"], + ) + legacy_logits = _run_logits( + case, + is_decode=False, + prefill_ws=workspace, + ).clone() + + scratch = aiter_fp4_paged_mqa_topk( + 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"], + out_page_indices=out_page_indices, + out_raw_indices=out_raw_indices, + prefill_workspace=workspace, + ) + torch.cuda.synchronize() + scratch_ptrs = ( + scratch.values.data_ptr(), + scratch.raw_indices.data_ptr(), + scratch.counts.data_ptr(), + scratch.workspace.candidate_values.data_ptr(), + ) + reused = aiter_fp4_paged_mqa_topk( + 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"], + out_page_indices=out_page_indices, + out_raw_indices=out_raw_indices, + prefill_workspace=workspace, + streaming_scratch=scratch, + ) + torch.cuda.synchronize() + assert reused is scratch + assert scratch_ptrs == ( + reused.values.data_ptr(), + reused.raw_indices.data_ptr(), + reused.counts.data_ptr(), + reused.workspace.candidate_values.data_ptr(), + ) + + for row, context in enumerate(contexts): + valid = min(context, topk) + raw = out_raw_indices[row, :valid].long() + assert torch.all((raw >= 0) & (raw < context)) + assert torch.unique(raw).numel() == valid + if context <= topk: + expected_raw = torch.arange(context, device=raw.device) + else: + expected_raw = torch.argsort( + legacy_logits[row, :context], + descending=True, + stable=True, + )[:topk] + torch.testing.assert_close(raw, expected_raw, rtol=0, atol=0) + expected_slots = ( + case["page_table"][row, raw // PAGE_SIZE].long() * PAGE_SIZE + + raw % PAGE_SIZE + ) + torch.testing.assert_close( + out_page_indices[row, :valid].long(), + expected_slots, + rtol=0, + atol=0, + ) + if valid < topk: + assert torch.all(out_raw_indices[row, valid:] == -1) + assert torch.all(out_page_indices[row, valid:] == -1) + + if __name__ == "__main__": sys.exit(pytest.main([__file__, "-v"])) diff --git a/test/registered/unit/layers/test_dsv4_fp4_streaming_topk.py b/test/registered/unit/layers/test_dsv4_fp4_streaming_topk.py new file mode 100644 index 000000000000..27cbe950b930 --- /dev/null +++ b/test/registered/unit/layers/test_dsv4_fp4_streaming_topk.py @@ -0,0 +1,167 @@ +import sys +import types +import unittest +from unittest import mock + +from sglang.kernels.ops.attention.dsv4 import fp4_indexer_hip +from sglang.srt.layers.attention.deepseek_v4_backend_hip_radix import ( + DeepseekV4HipRadixBackend, +) +from sglang.srt.layers.attention.dsv4 import indexer +from sglang.srt.model_executor.forward_batch_info import ForwardMode +from sglang.test.ci.ci_register import register_cpu_ci + +register_cpu_ci(est_time=2, suite="base-a-test-cpu") + + +def _fake_aiter_modules(entrypoint=None, allocator=None, result_type=None): + aiter = types.ModuleType("aiter") + ops = types.ModuleType("aiter.ops") + flydsl = types.ModuleType("aiter.ops.flydsl") + if entrypoint is not None: + flydsl.flydsl_pa_mqa_topk_fp4_prefill = entrypoint + if allocator is not None: + flydsl.allocate_fp4_prefill_topk_workspace = allocator + if result_type is not None: + flydsl.FP4PrefillTopKResult = result_type + aiter.ops = ops + ops.flydsl = flydsl + return { + "aiter": aiter, + "aiter.ops": ops, + "aiter.ops.flydsl": flydsl, + } + + +class TestAiterFP4StreamingTopKCapability(unittest.TestCase): + def tearDown(self): + fp4_indexer_hip.get_aiter_fp4_streaming_topk.cache_clear() + + def test_detects_exact_entrypoint(self): + entrypoint = lambda: None + with mock.patch.dict( + sys.modules, + _fake_aiter_modules(entrypoint, lambda: None, lambda: None), + ): + self.assertIs( + fp4_indexer_hip.get_aiter_fp4_streaming_topk(), + entrypoint, + ) + self.assertTrue(fp4_indexer_hip.aiter_fp4_streaming_topk_available()) + + def test_missing_entrypoint_uses_legacy_path(self): + with mock.patch.dict(sys.modules, _fake_aiter_modules()): + self.assertIsNone(fp4_indexer_hip.get_aiter_fp4_streaming_topk()) + self.assertFalse(fp4_indexer_hip.aiter_fp4_streaming_topk_available()) + + def test_capability_is_frozen_after_first_probe(self): + entrypoint = lambda: None + with mock.patch.dict(sys.modules, _fake_aiter_modules()): + self.assertIsNone(fp4_indexer_hip.get_aiter_fp4_streaming_topk()) + with mock.patch.dict( + sys.modules, + _fake_aiter_modules(entrypoint, lambda: None, lambda: None), + ): + self.assertIsNone(fp4_indexer_hip.get_aiter_fp4_streaming_topk()) + + def test_partial_api_uses_legacy_path(self): + entrypoint = lambda: None + cases = ( + _fake_aiter_modules(entrypoint), + _fake_aiter_modules(entrypoint, lambda: None), + _fake_aiter_modules(entrypoint, result_type=lambda: None), + ) + for modules in cases: + with self.subTest(exports=tuple(vars(modules["aiter.ops.flydsl"]))): + fp4_indexer_hip.get_aiter_fp4_streaming_topk.cache_clear() + with mock.patch.dict(sys.modules, modules): + self.assertIsNone(fp4_indexer_hip.get_aiter_fp4_streaming_topk()) + + +class TestAiterFP4StreamingTopKSelection(unittest.TestCase): + def _is_selected( + self, + *, + mode=ForwardMode.EXTEND, + prefill_graph=False, + capture_mode=False, + piecewise_graph=False, + breakable_graph=False, + capturing=False, + ): + with ( + mock.patch.object( + indexer, "get_is_capture_mode", return_value=capture_mode + ), + mock.patch.object( + indexer, + "is_in_tc_piecewise_cuda_graph", + return_value=piecewise_graph, + ), + mock.patch.object( + indexer, + "is_in_breakable_cuda_graph", + return_value=breakable_graph, + ), + mock.patch( + "torch.cuda.is_current_stream_capturing", + return_value=capturing, + ), + ): + return indexer._should_use_aiter_fp4_streaming_topk( + mode, + True, + prefill_graph, + ) + + def test_eager_prefill_uses_streaming_topk(self): + self.assertTrue(self._is_selected()) + + def test_capture_warmup_uses_legacy_topk(self): + self.assertFalse(self._is_selected(capture_mode=True)) + + def test_prefill_graph_uses_legacy_topk(self): + self.assertFalse(self._is_selected(prefill_graph=True)) + + def test_piecewise_graph_uses_legacy_topk(self): + self.assertFalse(self._is_selected(piecewise_graph=True)) + + def test_breakable_graph_uses_legacy_topk(self): + self.assertFalse(self._is_selected(breakable_graph=True)) + + def test_active_capture_uses_legacy_topk(self): + self.assertFalse(self._is_selected(capturing=True)) + + def test_decode_uses_legacy_topk(self): + self.assertFalse(self._is_selected(mode=ForwardMode.DECODE)) + + def test_hip_indexer_metadata_preserves_prefill_graph_mode(self): + backend = types.SimpleNamespace( + page_size=256, + dsa_topk_backend=types.SimpleNamespace(should_use_topk_v2=lambda: False), + ) + core_metadata = types.SimpleNamespace( + page_table=mock.sentinel.page_table, + c4_topk_lengths_raw=mock.sentinel.c4_seq_lens, + ) + with mock.patch.object( + sys.modules["sglang.srt.layers.attention.deepseek_v4_backend_hip_radix"], + "PagedIndexerMetadata", + ) as metadata_type: + DeepseekV4HipRadixBackend.init_forward_metadata_indexer( + backend, + core_metadata, + use_prefill_cuda_graph=True, + ) + + metadata_type.assert_called_once_with( + page_size=256, + page_table=mock.sentinel.page_table, + c4_seq_lens=mock.sentinel.c4_seq_lens, + use_topk_v2=False, + use_prefill_cuda_graph=True, + ) + + +if __name__ == "__main__": + unittest.main()