[AMD][DSv4] Add experimental AITER FP4 streaming TopK

docs/docs/references/environment_variables.mdx

@@ -710,6 +710,11 @@ SGLang supports various environment variables that can be used to configure its
       <td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}>Whether DeepSeek V4 experts use FP4. Set to <code>false</code> when using an FP4-to-FP8 converted DeepSeek V4 checkpoint.</td>
       <td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.02)"}}><code>true</code></td>
     </tr>
+    <tr>
+      <td style={{padding: "9px 12px", fontWeight: 500, backgroundColor: "rgba(255,255,255,0.02)"}}><code>SGLANG_DSV4_FP4_FUSED_TOPK</code></td>
+      <td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}>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.</td>
+      <td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.02)"}}><code>false</code></td>
+    </tr>
     <tr>
       <td style={{padding: "9px 12px", fontWeight: 500, backgroundColor: "rgba(255,255,255,0.02)"}}><code>SGLANG_DSV4_REASONING_EFFORT</code></td>
       <td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}>Default <code>reasoning_effort</code> for the DeepSeek V4 chat encoder when a request does not set it. The preview profile accepts <code>high</code> and <code>max</code>; the official profile accepts <code>low</code>, <code>high</code>, and <code>max</code>. The profile is detected from the bundled encoder. Override it with <code>--json-model-override-args '&#123;"dsv4_reasoning_effort_profile":"official"&#125;'</code>.</td>

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,

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

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

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,