chore: import upstream snapshot with attribution

python/sglang/srt/layers/utils/__init__.py

@@ -0,0 +1,3 @@
+# Temp workaround, make layer utils more fine-grained later
+from sglang.srt.layers.utils.common import *
+from sglang.srt.layers.utils.multi_platform import MultiPlatformOp

python/sglang/srt/layers/utils/common.py

@@ -0,0 +1,127 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+import logging
+import re
+
+import torch
+from torch.nn.parameter import Parameter
+
+logger = logging.getLogger(__name__)
+
+
+def get_layer_id(weight_name):
+    # example weight name: model.layers.10.self_attn.qkv_proj.weight
+    match = re.search(r"layers\.(\d+)\.", weight_name)
+    if match:
+        return int(match.group(1))
+    return None
+
+
+def pad_or_narrow_weight(
+    loaded_weight: torch.Tensor, input_dim: int, start_idx: int, shard_size: int
+) -> torch.Tensor:
+    # Padding with zeros for special case such as qwen2_5_VL's mlp which is not 8-aligned
+    valid_size = max(loaded_weight.shape[input_dim] - start_idx, 0)
+
+    if valid_size > 0:
+        loaded_slice = loaded_weight.narrow(input_dim, start_idx, valid_size)
+        pad_shape = list(loaded_weight.shape)
+        pad_shape[input_dim] = shard_size - valid_size
+        pad = torch.zeros(
+            pad_shape, dtype=loaded_weight.dtype, device=loaded_weight.device
+        )
+        return torch.cat([loaded_slice, pad], dim=input_dim)
+
+    # All padding
+    pad_shape = list(loaded_weight.shape)
+    pad_shape[input_dim] = shard_size
+    return torch.zeros(
+        pad_shape, dtype=loaded_weight.dtype, device=loaded_weight.device
+    )
+
+
+def is_strict_contiguous(x: torch.Tensor) -> bool:
+    expected_stride = 1
+    for size, stride in zip(reversed(x.shape), reversed(x.stride())):
+        if stride != expected_stride:
+            return False
+        expected_stride *= size
+    return True
+
+
+def strict_contiguous(x: torch.Tensor) -> torch.Tensor:
+    if is_strict_contiguous(x):
+        return x
+    return x.clone(memory_format=torch.contiguous_format)
+
+
+def copy_or_rebind_param(
+    module: torch.nn.Module, name: str, new_value: torch.Tensor
+) -> None:
+    """Keep parameter identities stable for CUDA graph reuse and hot reload."""
+    new_value = new_value.detach()
+    param = getattr(module, name, None)
+    if isinstance(param, Parameter):
+        if param.data.shape == new_value.shape and param.data.dtype == new_value.dtype:
+            param.data.copy_(new_value)
+        else:
+            param.data = new_value
+        param.requires_grad_(False)
+    else:
+        setattr(module, name, Parameter(new_value, requires_grad=False))
+
+
+def alias_or_bind_derived_param(
+    module: torch.nn.Module,
+    source_name: str,
+    derived_name: str,
+    derived_value: torch.Tensor,
+) -> None:
+    """Bind a post-processed (derived) tensor to a derived attribute name.
+
+    When `derived_value` is broadcastable to the source Parameter's shape (and
+    dtype matches), write it broadcast-filled into the source's storage in
+    place and register `derived_name` as an alias of the source Parameter. The
+    two attribute names then share one underlying buffer, so:
+      - apply() can read via `derived_name`
+      - update_weights_from_disk can keep refilling `source_name` (the loader
+        re-runs process_weights_after_loading which re-derives in place)
+      - peak GPU memory is the source size, not source + derived.
+
+    When the shapes are not broadcast-compatible, fall back to allocating a
+    separate Parameter under `derived_name` via copy_or_rebind_param.
+    """
+    derived_value = derived_value.detach()
+    source = getattr(module, source_name, None)
+    if isinstance(source, Parameter) and source.data.dtype == derived_value.dtype:
+        try:
+            broadcast = torch.broadcast_to(derived_value, source.data.shape)
+        except RuntimeError:
+            broadcast = None
+        if broadcast is not None:
+            source.data.copy_(broadcast)
+            source.requires_grad_(False)
+            setattr(module, derived_name, source)
+            return
+    copy_or_rebind_param(module, derived_name, derived_value)
+
+
+class PPMissingLayer(torch.nn.Identity):
+    # Adapted from
+    # https://github.com/vllm-project/vllm/blob/18ed3132d2bfe1df9a74729457b69243955221e8/vllm/model_executor/models/utils.py#L468C1-L486C1
+    """
+    A placeholder layer for missing layers in a pipeline parallel model.
+    """
+
+    def __init__(self, *args, **kwargs):
+        super().__init__()
+        self.return_tuple = kwargs.get("return_tuple", False)
+
+    def forward(self, *args, **kwargs):
+        """
+        Return the first arg from args or the first value from kwargs.
+
+        Wraps the input in a tuple if `self.return_tuple` is True.
+        """
+        input = args[0] if args else next(iter(kwargs.values()))
+        return (input,) if self.return_tuple else input

python/sglang/srt/layers/utils/cp_utils.py

@@ -0,0 +1,694 @@
+from dataclasses import dataclass
+from itertools import accumulate
+from typing import Callable, List
+
+import torch
+import torch.nn.functional as F
+
+from sglang.srt.distributed.device_communicators.pynccl_allocator import (
+    use_symmetric_memory,
+)
+from sglang.srt.layers.dp_attention import (
+    attn_cp_all_gather_into_tensor,
+    is_allocation_symmetric,
+)
+from sglang.srt.layers.moe import get_moe_a2a_backend
+from sglang.srt.mem_cache.memory_pool import KVWriteLoc
+from sglang.srt.model_executor.forward_context import get_token_to_kv_pool
+from sglang.srt.runtime_context import get_parallel, get_server_args
+
+
+@dataclass
+class ContextParallelMetadata:
+    # Layout lists have length bs * cp_segment_num (= bs * 2 * cp_size).
+    split_list: List[int] = None
+    zigzag_index: List[int] = None
+    cp_reverse_index: List[int] = None
+    reverse_split_len: List[int] = None
+
+    # Per-rank-aggregate lists have length cp_size.
+    # max_rank_len is a list of cp_size copies of max(per_rank_actual_token),
+    # kept as a list for torch.split() bucket sizes.
+    per_rank_actual_token: List[int] = None
+    max_rank_len: List[int] = None
+
+    # Per-sequence FlashAttention tensors (shape [bs] or [bs+1]).
+    kv_len_prev_tensor: torch.Tensor = None  # [bs] int32 CUDA
+    kv_len_next_tensor: torch.Tensor = None  # [bs] int32 CUDA
+    actual_seq_q_prev_tensor: torch.Tensor = None  # [bs] int32 CUDA
+    actual_seq_q_next_tensor: torch.Tensor = None  # [bs] int32 CUDA
+    cu_seqlens_q_prev_tensor: torch.Tensor = None  # [bs+1] int32 CUDA
+    cu_seqlens_q_next_tensor: torch.Tensor = None  # [bs+1] int32 CUDA
+
+    # Scalars derived from the per-sequence lists above.
+    total_q_prev_tokens: int = 0
+    total_q_next_tokens: int = 0
+    max_seqlen_q_prev: int = 0
+    max_seqlen_q_next: int = 0
+
+    # Per-seq CPU lists (useful for NSA indexer and diagnostics).
+    kv_len_prev_list: List[int] = None
+    kv_len_next_list: List[int] = None
+    actual_seq_q_prev_list: List[int] = None
+    actual_seq_q_next_list: List[int] = None
+
+    # Aggregate sum of extend_seq_lens across the batch.
+    total_seq_lens: int = 0
+    bs: int = 1
+
+
+def is_prefill_context_parallel_enabled():
+    return get_server_args().enable_prefill_context_parallel
+
+
+def is_prefill_cp_in_seq_split():
+    return (
+        is_prefill_context_parallel_enabled()
+        and get_server_args().prefill_cp_mode == "in-seq-split"
+    )
+
+
+def get_cp_padding_align_size() -> int:
+    """Token-count alignment for CP padding of global_num_tokens: 2 * cp_size
+    for zigzag (in-seq-split) CP, otherwise cp_size (1 when CP is off, so the
+    padding is a no-op; extra padding breaks EAGLE/MTP draft prefill, see
+    #23269). Keep prepare_mlp_sync_batch and cal_padded_tokens consistent
+    through this helper.
+    """
+    from sglang.srt.layers.attention.dsa.utils import is_dsa_prefill_cp_in_seq_split
+
+    attn_cp_size = get_parallel().attn_cp_size
+    if is_prefill_cp_in_seq_split() or is_dsa_prefill_cp_in_seq_split():
+        return attn_cp_size * 2
+    return attn_cp_size
+
+
+def is_mla_prefill_cp_enabled() -> bool:
+    sa = get_server_args()
+    return sa.enable_prefill_context_parallel and sa.use_mla_backend()
+
+
+def mla_use_prefill_cp(forward_batch, mla_enable_prefill_cp=None):
+    if mla_enable_prefill_cp is None:
+        mla_enable_prefill_cp = is_mla_prefill_cp_enabled()
+    return (
+        forward_batch.attn_cp_metadata is not None
+        and mla_enable_prefill_cp
+        and forward_batch.forward_mode.is_context_parallel_extend()
+    )
+
+
+def can_cp_split(seq_len: int, cp_size: int, forward_batch):
+    # Base conditions: CP must be enabled, size > 1, and this must be a
+    # CP-extend (prefill) step. The seq_len // (cp_size * 2) check ensures
+    # the load-balancing split into 2 * cp_size blocks is non-degenerate.
+    from sglang.srt.model_executor.forward_batch_info import ForwardMode
+
+    cur_cp_seq_len = seq_len // (cp_size * 2)
+    if not (
+        cur_cp_seq_len != 0
+        and cp_size > 1
+        # prepare_context_parallel_metadata hard-codes bs_per_cp_group = 1;
+        # guard explicitly to avoid silent mis-partitioning under continuous batching.
+        and forward_batch.forward_mode.is_context_parallel_extend()
+        # is_context_parallel_extend() returns True for MIXED (prefill+decode
+        # in one step), but the zigzag split only makes sense on pure extend.
+        and forward_batch.forward_mode != ForwardMode.MIXED
+        and is_prefill_context_parallel_enabled()
+    ):
+        return False
+
+    # Per-sequence guards for bs > 1. Every sequence must be long enough for
+    # the 2*cp_size-way split. A sub-threshold request reaching this point
+    # means the scheduler failed to filter it out and a silent non-CP
+    # fallback would have masked the bug -- raise instead. Per-sequence
+    # radix-cache prefix is supported: prefix is baked into kv_len_prev/next
+    # via prefix_offsets[s] inside prepare_context_parallel_metadata.
+    extend_lens = getattr(forward_batch, "extend_seq_lens_cpu", None)
+    if extend_lens is None:
+        return True
+
+    cp_min = cp_size * 2
+    for L in extend_lens:
+        if L < cp_min:
+            # A sub-threshold request cannot be zigzag-split into 2*cp_size
+            # blocks; fall back to a normal (non-CP) prefill for this batch
+            # instead of failing. Happens e.g. when a radix-cache prefix hit
+            # leaves only a few unique extend tokens.
+            return False
+
+    return True
+
+
+def cp_split_and_rebuild_data(forward_batch, input_: torch.Tensor):
+    from sglang.srt.layers.attention.dsa.utils import (
+        dsa_cp_round_robin_split_data,
+        is_dsa_prefill_cp_round_robin_split,
+    )
+
+    if is_dsa_prefill_cp_round_robin_split():
+        cp_size = get_parallel().attn_cp_size
+        assert (
+            input_.shape[0] % cp_size == 0
+        ), f"Expect input shape 0 can divided by cp size, but got input shape {input_.shape}, cp size {cp_size}"
+        return dsa_cp_round_robin_split_data(input_)
+
+    input_list = list(
+        torch.split(input_, forward_batch.attn_cp_metadata.split_list, dim=0)
+    )
+    result = torch.cat(
+        [input_list[i] for i in forward_batch.attn_cp_metadata.zigzag_index], dim=0
+    ).view(-1, input_.shape[-1])
+    return result
+
+
+def cp_split_and_rebuild_position(forward_batch, positions: torch.Tensor):
+    from sglang.srt.layers.attention.dsa.utils import (
+        dsa_cp_round_robin_split_data,
+        is_dsa_prefill_cp_round_robin_split,
+    )
+
+    if is_dsa_prefill_cp_round_robin_split():
+        cp_size = get_parallel().attn_cp_size
+        assert positions.shape[0] % cp_size == 0, (
+            f"Expect positions shape 0 can divided by cp size, but got positions shape {positions.shape}, "
+            f"cp size {cp_size}"
+        )
+        return dsa_cp_round_robin_split_data(positions)
+
+    position_id_list = list(
+        torch.split(positions, forward_batch.attn_cp_metadata.split_list, dim=-1)
+    )
+    positions = torch.cat(
+        [position_id_list[i] for i in forward_batch.attn_cp_metadata.zigzag_index],
+        dim=-1,
+    )
+    return positions
+
+
+def cp_round_robin_input_ids(input_ids):
+    """
+    input input_ids:
+    rank0~7: 0,1,2,3,4,5,...
+
+    output input_ids:
+    a2a none:
+    rank0~7: 0,8,16,...,1,9,17,...,2,10,18,...
+
+    not a2a none:
+    rank0: 0,8,16,...
+    rank1: 1,9,17,...
+    rank2: 2,10,18,...
+    ...
+    """
+    cp_size = get_parallel().attn_cp_size
+    cp_rank = get_parallel().attn_cp_rank
+    if get_moe_a2a_backend().is_none():
+        input_ids = input_ids.reshape(-1, cp_size).T.flatten()
+    else:
+        input_ids = input_ids[cp_rank::cp_size].contiguous()
+    return input_ids
+
+
+def cp_all_gather_reorganized_into_tensor(input_tensor, cp_size, forward_batch, stream):
+    """
+    Allgather communication for context_parallel(kv_cache, index_k, hidden_states).
+    This implementation mainly consists of three parts:
+    Step 1, padding the input shape to unify the shape for allgather communication (the shape must be the same).
+    Step 2, allgather communication(async).
+    Step 3, removing the padding and reassembling the data according to the actual tokens.
+    """
+    max_len = forward_batch.attn_cp_metadata.max_rank_len[0]
+    pad_size = max_len - input_tensor.shape[0]
+    if pad_size > 0:
+        input_tensor = F.pad(
+            input_tensor, (0, 0, 0, pad_size), mode="constant", value=0
+        )
+    with use_symmetric_memory(
+        get_parallel().attn_cp_group, disabled=not is_allocation_symmetric()
+    ):
+        input_tensor_full = torch.empty(
+            max_len * cp_size,
+            input_tensor.shape[1],
+            device=input_tensor.device,
+            dtype=input_tensor.dtype,
+        )
+
+    get_parallel().attn_cp_group.cp_all_gather_into_tensor_async(
+        input_tensor_full, input_tensor, stream
+    )
+
+    outputs_list_max = list(
+        torch.split(
+            input_tensor_full, forward_batch.attn_cp_metadata.max_rank_len, dim=0
+        )
+    )
+    outputs = torch.cat(
+        [
+            outputs_list_max[index][:per_rank_len]
+            for index, per_rank_len in enumerate(
+                forward_batch.attn_cp_metadata.per_rank_actual_token
+            )
+        ],
+        dim=0,
+    )
+
+    return outputs
+
+
+def cp_all_gather_reorganized_into_tensor_kv_cache(
+    input_tensor, cp_size, forward_batch, stream
+):
+    """
+    Allgather communication for context_parallel KV cache.
+    Handles multi-dimensional tensors (e.g., [seq_len, num_heads, head_dim]).
+    """
+    max_len = forward_batch.attn_cp_metadata.max_rank_len[0]
+    pad_size = max_len - input_tensor.shape[0]
+    if pad_size > 0:
+        # Pad the first dimension (seq_len). F.pad expects padding in reverse dimension order.
+        # For n dimensional tensor, we need 2*n values: (last_dim_left, last_dim_right, ..., first_dim_left, first_dim_right)
+        # To pad only the first dimension: [0, 0] * (ndim - 1) + [0, pad_size]
+        padding = [0, 0] * (input_tensor.ndim - 1) + [0, pad_size]
+        input_tensor = F.pad(input_tensor, padding, mode="constant", value=0)
+
+    # Create output tensor with proper shape for all dimensions
+    with use_symmetric_memory(
+        get_parallel().attn_cp_group, disabled=not is_allocation_symmetric()
+    ):
+        input_tensor_full = torch.empty(
+            max_len * cp_size,
+            *input_tensor.shape[1:],
+            device=input_tensor.device,
+            dtype=input_tensor.dtype,
+        )
+
+    get_parallel().attn_cp_group.cp_all_gather_into_tensor_async(
+        input_tensor_full, input_tensor, stream
+    )
+
+    outputs_list_max = list(
+        torch.split(
+            input_tensor_full, forward_batch.attn_cp_metadata.max_rank_len, dim=0
+        )
+    )
+    outputs = torch.cat(
+        [
+            outputs_list_max[index][:per_rank_len]
+            for index, per_rank_len in enumerate(
+                forward_batch.attn_cp_metadata.per_rank_actual_token
+            )
+        ],
+        dim=0,
+    )
+
+    return outputs
+
+
+def cp_all_gather_rerange_output(input_tensor, cp_size, forward_batch, stream):
+    """
+    # for in-seq-split
+    |   +-----------before allgather------------+|
+    |   | dp_atten_tp0: block0, block7 |
+    |   | dp_atten_tp1: block1, block6 |
+    |   | dp_atten_tp2: block2, block5 |
+    |   | dp_atten_tp3: block3, block4 |
+    |
+    |   +----------before rerange---------------+|
+    | block0 | block7 | block1 | block6 | block2 | block5 | block3 | block4 |
+    |
+    |   +--------------result-------------------+
+    | block0 | block1 | block2 | block3 | block4 | block5 | block6 | block7 |
+    |   +-------------------------+
+
+    # for round-robin-split
+    |   +-----------before allgather------------+|
+    | dp_atten_tp0: token0, token4, token8, token12, token16, ... |
+    | dp_atten_tp1: token1, token5, token9, token13, token17, ... |
+    | dp_atten_tp2: token2, token6, token10, token14, token18, ... |
+    | dp_atten_tp3: token3, token7, token11, token15, token19, ... |
+    |
+    |   +--------------result-------------------+
+    | token0, token1, token2, token3, token4, token5, token6, token7, ...
+    |   +-------------------------+
+    """
+    from sglang.srt.layers.attention.dsa.utils import (
+        is_dsa_prefill_cp_round_robin_split,
+    )
+
+    if is_dsa_prefill_cp_round_robin_split():
+        with use_symmetric_memory(
+            get_parallel().attn_cp_group, disabled=not is_allocation_symmetric()
+        ):
+            output_tensor = input_tensor.new_empty(
+                (input_tensor.shape[0] * cp_size, *input_tensor.shape[1:]),
+            )
+        attn_cp_all_gather_into_tensor(
+            output_tensor,
+            input_tensor,
+        )
+        out_shape = output_tensor.shape
+        output_tensor = (
+            output_tensor.view(cp_size, -1, *out_shape[1:])
+            .transpose(0, 1)
+            .reshape(out_shape)
+        )
+        return output_tensor
+
+    # TODO: Do we need to remove the padding here?
+    bs_seq_len, hidden_size = input_tensor.shape
+    output_tensor = cp_all_gather_reorganized_into_tensor(
+        input_tensor,
+        cp_size,
+        forward_batch,
+        stream,
+    )
+    outputs_list = list(
+        torch.split(
+            output_tensor, forward_batch.attn_cp_metadata.reverse_split_len, dim=0
+        )
+    )
+    output_tensor = torch.cat(
+        [outputs_list[i] for i in forward_batch.attn_cp_metadata.cp_reverse_index],
+        dim=0,
+    )
+    output_tensor = output_tensor.view(-1, hidden_size)
+    return output_tensor
+
+
+def cp_all_gather_rerange_kv_cache(input_tensor, cp_size, forward_batch, stream):
+    """
+    Allgather and reorganize KV cache from all ranks in context parallel group.
+
+    # for in-seq-split
+    |   +-----------before allgather------------+|
+    |   | dp_atten_tp0: block0, block7 |
+    |   | dp_atten_tp1: block1, block6 |
+    |   | dp_atten_tp2: block2, block5 |
+    |   | dp_atten_tp3: block3, block4 |
+    |
+    |   +----------before rerange---------------+|
+    | block0 | block7 | block1 | block6 | block2 | block5 | block3 | block4 |
+    |
+    |   +--------------result-------------------+
+    | block0 | block1 | block2 | block3 | block4 | block5 | block6 | block7 |
+    |   +-------------------------+
+    """
+    output_tensor = cp_all_gather_reorganized_into_tensor_kv_cache(
+        input_tensor,
+        cp_size,
+        forward_batch,
+        stream,
+    )
+    outputs_list = list(
+        torch.split(
+            output_tensor, forward_batch.attn_cp_metadata.reverse_split_len, dim=0
+        )
+    )
+    output_tensor = torch.cat(
+        [outputs_list[i] for i in forward_batch.attn_cp_metadata.cp_reverse_index],
+        dim=0,
+    )
+    # No need to reshape - output_tensor already has the correct shape [seq_len, ...]
+    return output_tensor
+
+
+def cp_allgather_and_save_kv_cache(forward_batch, layer, k, v, cp_size, swa_loc=None):
+    """
+    Allgather KV cache from all CP ranks and write the full result
+    into each rank's local memory pool.
+
+    swa_loc is the pre-translated full->SWA write target for hybrid SWA pools.
+    """
+    cache_loc = (
+        forward_batch.out_cache_loc
+        if not layer.is_cross_attention
+        else forward_batch.encoder_out_cache_loc
+    )
+
+    k = k.contiguous()
+    v = v.contiguous()
+
+    key_cache_full = cp_all_gather_rerange_kv_cache(
+        k, cp_size, forward_batch, torch.cuda.current_stream()
+    )
+    value_cache_full = cp_all_gather_rerange_kv_cache(
+        v, cp_size, forward_batch, torch.cuda.current_stream()
+    )
+
+    get_token_to_kv_pool().set_kv_buffer(
+        layer,
+        KVWriteLoc(cache_loc, swa_loc),
+        key_cache_full,
+        value_cache_full,
+        layer.k_scale,
+        layer.v_scale,
+    )
+
+
+def cp_attn_forward_extend(
+    forward_batch,
+    q: torch.Tensor,
+    device: torch.device,
+    attn_fn: Callable[[torch.Tensor, torch.Tensor, torch.Tensor, int], torch.Tensor],
+) -> torch.Tensor:
+    """
+    Split q into prev/next zigzag halves based on CP metadata, call the
+    backend-specific attention function twice with appropriate per-half
+    metadata, and concatenate the results.
+
+    For bs > 1, q is laid out as [all_prev_tokens_across_seqs,
+    all_next_tokens_across_seqs]; the split point is total_q_prev_tokens.
+    cu_seqlens_q_prev/next tensors have shape [bs+1] and carry the
+    per-sequence boundaries through FlashAttention's variable-length API.
+
+    attn_fn signature:
+        attn_fn(q, cu_seqlens_q, cache_seqlens, max_seqlen_q) -> result
+    where only these four CP-varying parameters differ between halves.
+    All other backend-specific args should be captured in the closure.
+    """
+    cp_meta = forward_batch.attn_cp_metadata
+
+    q_prev = q[: cp_meta.total_q_prev_tokens]
+    q_next = q[cp_meta.total_q_prev_tokens :]
+
+    result_prev = attn_fn(
+        q_prev,
+        cp_meta.cu_seqlens_q_prev_tensor,
+        cp_meta.kv_len_prev_tensor,
+        cp_meta.max_seqlen_q_prev,
+    )
+    result_next = attn_fn(
+        q_next,
+        cp_meta.cu_seqlens_q_next_tensor,
+        cp_meta.kv_len_next_tensor,
+        cp_meta.max_seqlen_q_next,
+    )
+
+    return torch.concat([result_prev, result_next], dim=0)
+
+
+def prepare_context_parallel_metadata(
+    kv_len,
+    cp_rank,
+    cp_size,
+    seqs_len,
+    extend_seqs_len=None,
+    device="cuda",
+):
+    from sglang.srt.layers.attention.dsa.utils import (
+        is_dsa_prefill_cp_round_robin_split,
+    )
+
+    if is_dsa_prefill_cp_round_robin_split():
+        return ContextParallelMetadata()
+
+    """prepare_input_dp_with_cp_dsa-zigzag index
+    Example (DP_ATTENT_TP == CP_SIZE == 4, single sequence):
+        block0 | block1 | block2 | block3 | block4 | block5 | block6 | block7
+        rank 0: block0, block7
+        rank 1: block1, block6
+        rank 2: block2, block5
+        rank 3: block3, block4
+    For bs > 1, each sequence is split into cp_segment_num = 2 * cp_size
+    blocks independently; per-rank layout becomes:
+        [s0.block_r, s1.block_r, ..., s_{bs-1}.block_r,
+         s0.block_{2*cp_size-1-r}, ..., s_{bs-1}.block_{2*cp_size-1-r}]
+    i.e. all prev blocks first, then all next blocks -- so torch.split at
+    total_q_prev_tokens cleanly separates them.
+    """
+    assert extend_seqs_len is not None
+    extend_seqs_len = [int(x) for x in extend_seqs_len]
+
+    # Update the extend_seqs_len to the padded length.
+    pad_len = int(kv_len) - sum(extend_seqs_len)
+    if pad_len > 0:
+        extend_seqs_len[-1] += pad_len
+        if seqs_len is not None and len(seqs_len) == len(extend_seqs_len):
+            seqs_len = list(seqs_len)
+            seqs_len[-1] += pad_len
+
+    bs = len(extend_seqs_len)
+    cp_segment_num = cp_size * 2
+
+    # Prefix offset (radix cache hit length) per sequence. For non-NSA
+    # (FlashAttention) the prefix is baked into kv_len_prev/next via
+    # prefix_offsets[s] below, so cache_seqlens correctly covers the cached
+    # prefix. NSA leaves bare cumulatives so its indexer can re-add the
+    # offset itself.
+    if seqs_len is not None and len(seqs_len) == bs:
+        prefix_offsets = [
+            max(int(seqs_len[s]) - extend_seqs_len[s], 0) for s in range(bs)
+        ]
+    else:
+        prefix_offsets = [0] * bs
+
+    # Per-sequence block sizes: first (L % cp_segment_num) blocks get +1.
+    per_seq_block_sizes: List[List[int]] = []
+    split_list: List[int] = []
+    for s in range(bs):
+        L = extend_seqs_len[s]
+        base = L // cp_segment_num
+        rem = L % cp_segment_num
+        blk = [base + 1 if i < rem else base for i in range(cp_segment_num)]
+        per_seq_block_sizes.append(blk)
+        split_list.extend(blk)
+
+    # Per-rank aggregate: this rank owns block r and block (2*cp_size-1-r)
+    # of every sequence.
+    per_rank_actual_token = [0] * cp_size
+    for r in range(cp_size):
+        total = 0
+        for s in range(bs):
+            total += (
+                per_seq_block_sizes[s][r]
+                + per_seq_block_sizes[s][cp_segment_num - 1 - r]
+            )
+        per_rank_actual_token[r] = total
+    max_single_rank = max(per_rank_actual_token) if per_rank_actual_token else 0
+    # Kept as cp_size copies so downstream torch.split(x, max_rank_len) still
+    # works directly. All entries intentionally identical.
+    max_rank_len = [max_single_rank] * cp_size
+
+    # Zigzag index selecting which of split_list's bs * cp_segment_num pieces
+    # this rank owns, in the order [all_prevs, all_nexts].
+    zigzag_index = list(
+        range(cp_rank, cp_rank + bs * cp_segment_num, cp_segment_num)
+    ) + list(
+        range(
+            cp_segment_num - cp_rank - 1,
+            bs * cp_segment_num,
+            cp_segment_num,
+        )
+    )
+
+    # Reverse index: given the post-allgather concatenation
+    #   [rank0_prevs_all_seqs, rank0_nexts_all_seqs,
+    #    rank1_prevs_all_seqs, rank1_nexts_all_seqs, ...]
+    # produce a permutation that restores [s0_b0..s0_bN, s1_b0..s1_bN, ...].
+    cp_reverse_index: List[int] = []
+    for batch_id in range(bs):
+        cp_reverse_index.extend(
+            list(range(batch_id, cp_segment_num * bs, 2 * bs))
+            + list(
+                range(
+                    (cp_segment_num - 1) * bs + batch_id,
+                    0,
+                    -2 * bs,
+                )
+            )
+        )
+
+    # Split sizes matching the post-allgather concatenation order above.
+    reverse_split_len: List[int] = []
+    for r in range(cp_size):
+        for s in range(bs):
+            reverse_split_len.append(per_seq_block_sizes[s][r])
+        for s in range(bs):
+            reverse_split_len.append(per_seq_block_sizes[s][cp_segment_num - 1 - r])
+
+    # Per-sequence cumulatives used for FA cache_seqlens.
+    #   kv_len_prev[s] = sum of seq s's blocks [0..cp_rank] (inclusive).
+    #   kv_len_next[s] = sum of seq s's blocks [0..cp_segment_num-cp_rank-1] (inclusive).
+    from sglang.srt.layers.attention.dsa.utils import is_dsa_enable_prefill_cp
+
+    nsa_mode = is_dsa_enable_prefill_cp()
+    kv_len_prev_list: List[int] = []
+    kv_len_next_list: List[int] = []
+    actual_seq_q_prev_list: List[int] = []
+    actual_seq_q_next_list: List[int] = []
+    for s in range(bs):
+        blk = per_seq_block_sizes[s]
+        cum_prev = sum(blk[: cp_rank + 1])
+        cum_next = sum(blk[: cp_segment_num - cp_rank])
+        # NSA indexer re-adds prefix offset itself; leave bare cumulative.
+        # For non-NSA (FlashAttention), bake prefix into cache_seqlens.
+        if nsa_mode:
+            kv_len_prev_list.append(cum_prev)
+            kv_len_next_list.append(cum_next)
+        else:
+            kv_len_prev_list.append(prefix_offsets[s] + cum_prev)
+            kv_len_next_list.append(prefix_offsets[s] + cum_next)
+        actual_seq_q_prev_list.append(blk[cp_rank])
+        actual_seq_q_next_list.append(blk[cp_segment_num - cp_rank - 1])
+
+    # FlashAttention CUDA tensors (device parameterized for unit tests).
+    kv_len_prev_tensor = torch.tensor(
+        kv_len_prev_list, device=device, dtype=torch.int32
+    )
+    kv_len_next_tensor = torch.tensor(
+        kv_len_next_list, device=device, dtype=torch.int32
+    )
+    actual_seq_q_prev_tensor = torch.tensor(
+        actual_seq_q_prev_list, device=device, dtype=torch.int32
+    )
+    actual_seq_q_next_tensor = torch.tensor(
+        actual_seq_q_next_list, device=device, dtype=torch.int32
+    )
+    cu_prev = [0] + list(accumulate(actual_seq_q_prev_list))
+    cu_next = [0] + list(accumulate(actual_seq_q_next_list))
+    cu_seqlens_q_prev_tensor = torch.tensor(cu_prev, device=device, dtype=torch.int32)
+    cu_seqlens_q_next_tensor = torch.tensor(cu_next, device=device, dtype=torch.int32)
+
+    total_q_prev_tokens = cu_prev[-1]
+    total_q_next_tokens = cu_next[-1]
+    max_seqlen_q_prev = max(actual_seq_q_prev_list) if actual_seq_q_prev_list else 0
+    max_seqlen_q_next = max(actual_seq_q_next_list) if actual_seq_q_next_list else 0
+    total_seq_lens = sum(extend_seqs_len)
+
+    # Cheap invariants: metadata must be a valid permutation spec.
+    # - split_list has bs * cp_segment_num pieces (all blocks, all seqs).
+    # - zigzag_index has 2 * bs entries (this rank's prev + next per seq).
+    # - cp_reverse_index has bs * cp_segment_num entries (reorders the
+    #   full allgathered stream back to per-seq-original order).
+    assert len(split_list) == bs * cp_segment_num
+    assert sum(split_list) == total_seq_lens
+    assert len(zigzag_index) == 2 * bs
+    assert len(cp_reverse_index) == bs * cp_segment_num
+    assert sorted(cp_reverse_index) == list(range(bs * cp_segment_num))
+    assert sum(per_rank_actual_token) == total_seq_lens
+
+    return ContextParallelMetadata(
+        split_list=split_list,
+        zigzag_index=zigzag_index,
+        cp_reverse_index=cp_reverse_index,
+        reverse_split_len=reverse_split_len,
+        per_rank_actual_token=per_rank_actual_token,
+        max_rank_len=max_rank_len,
+        kv_len_prev_tensor=kv_len_prev_tensor,
+        kv_len_next_tensor=kv_len_next_tensor,
+        actual_seq_q_prev_tensor=actual_seq_q_prev_tensor,
+        actual_seq_q_next_tensor=actual_seq_q_next_tensor,
+        cu_seqlens_q_prev_tensor=cu_seqlens_q_prev_tensor,
+        cu_seqlens_q_next_tensor=cu_seqlens_q_next_tensor,
+        total_q_prev_tokens=total_q_prev_tokens,
+        total_q_next_tokens=total_q_next_tokens,
+        max_seqlen_q_prev=max_seqlen_q_prev,
+        max_seqlen_q_next=max_seqlen_q_next,
+        kv_len_prev_list=kv_len_prev_list,
+        kv_len_next_list=kv_len_next_list,
+        actual_seq_q_prev_list=actual_seq_q_prev_list,
+        actual_seq_q_next_list=actual_seq_q_next_list,
+        total_seq_lens=total_seq_lens,
+        bs=bs,
+    )

python/sglang/srt/layers/utils/hash.py

@@ -0,0 +1,121 @@
+import torch
+import triton
+import triton.language as tl
+
+
+@triton.jit
+def rotl32(x, r: tl.constexpr) -> tl.uint32:
+    """
+    rotate left 32-bit integer x by r bits
+    e.g. x = 01110001, r = 2 -> 11000101
+    """
+    x = x.to(tl.uint64)
+    return ((x << r) | (x >> (32 - r))) & 0xFFFFFFFF
+
+
+@triton.jit
+def fmix32(h: tl.uint32) -> tl.uint32:
+    """
+    final mix of 32-bit hash value for MurmurHash
+    """
+    h ^= h >> 16
+    h = (h * 0x85EBCA6B) & 0xFFFFFFFF
+    h ^= h >> 13
+    h = (h * 0xC2B2AE35) & 0xFFFFFFFF
+    h ^= h >> 16
+    return h
+
+
+@triton.jit
+def murmur3_mix(h: tl.uint32, k: tl.uint32) -> tl.uint32:
+    """
+    Mixes a 32-bit key into the hash state.
+    """
+    c1: tl.uint32 = 0xCC9E2D51
+    c2: tl.uint32 = 0x1B873593
+    r1: tl.constexpr = 15
+    r2: tl.constexpr = 13
+    mm: tl.uint32 = 5
+    nn: tl.uint32 = 0xE6546B64
+
+    k = (k * c1) & 0xFFFFFFFF
+    k = rotl32(k, r1)
+    k = (k * c2) & 0xFFFFFFFF
+    h ^= k
+    h = rotl32(h, r2)
+    h = (h * mm + nn) & 0xFFFFFFFF
+    return h
+
+
+@triton.jit
+def murmur_hash32_kernel(
+    seed_ptr,
+    positions_ptr,
+    col_indices_ptr,
+    output_ptr,
+    num_rows,
+    num_cols,
+    BLOCK_SIZE: tl.constexpr,
+):
+    """
+    MurmurHash 32-bit implementation for Triton.
+    Reference:
+    - https://medium.com/@thealonemusk/murmurhash-the-scrappy-algorithm-that-secretly-powers-half-the-internet-2d3f79b4509b
+    - https://en.wikipedia.org/wiki/MurmurHash
+
+    We treat 64-bit seed, 32-bit position, and 32-bit col_index as 4 4-byte blocks, and bit-blend them together.
+    """
+    pid_row = tl.program_id(0)
+    pid_col = tl.program_id(1)
+
+    row_idx = pid_row
+    col_offsets = pid_col * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)
+    mask = col_offsets < num_cols
+
+    # Load inputs
+    seed = tl.load(seed_ptr + row_idx).to(tl.uint64)
+    pos = tl.load(positions_ptr + row_idx).to(tl.uint32)
+    col = tl.load(col_indices_ptr + col_offsets, mask=mask, other=0).to(tl.uint32)
+
+    h: tl.uint32 = 0  # hash accumulator
+
+    # Process seed_low
+    k: tl.uint32 = (seed & 0xFFFFFFFF).to(tl.uint32)
+    h = murmur3_mix(h, k)
+
+    # Process seed_high
+    k = ((seed >> 32) & 0xFFFFFFFF).to(tl.uint32)
+    h = murmur3_mix(h, k)
+
+    # Process position block starting from seed32
+    h = murmur3_mix(h, pos)
+
+    # Process col block
+    h = murmur3_mix(h, col)
+
+    # Finalize (len=16 for seed + pos + col)
+    h ^= 16
+    h = fmix32(h)
+
+    # Store result as uint32
+    tl.store(output_ptr + row_idx * num_cols + col_offsets, h, mask=mask)
+
+
+def murmur_hash32(seed, positions, col_indices):
+    assert (
+        seed.shape == positions.shape
+    ), "Seed and positions must have the same shape (n,)"
+    assert (
+        len(seed.shape) == 1 and len(col_indices.shape) == 1
+    ), f"Inputs must be 1D tensors {seed.shape=} {col_indices.shape=}"
+    n = seed.shape[0]
+    m = col_indices.shape[0]
+    device = seed.device
+    hashed = torch.empty((n, m), dtype=torch.uint32, device=device)
+
+    BLOCK_SIZE = 1024
+    grid = (n, triton.cdiv(m, BLOCK_SIZE))
+    murmur_hash32_kernel[grid](
+        seed, positions, col_indices, hashed, n, m, BLOCK_SIZE=BLOCK_SIZE
+    )
+    return hashed

python/sglang/srt/layers/utils/logprob.py

@@ -0,0 +1,357 @@
+from __future__ import annotations
+
+import dataclasses
+from enum import Enum, auto
+from typing import TYPE_CHECKING, List, Optional
+
+import torch
+
+from sglang.srt.environ import envs
+
+if TYPE_CHECKING:
+    from sglang.srt.layers.logits_processor import LogitsMetadata
+
+
+class LogprobStage(Enum):
+    PREFILL = auto()
+    DECODE = auto()
+
+
+@dataclasses.dataclass
+class InputLogprobsResult:
+    input_token_logprobs: torch.Tensor
+    input_top_logprobs_val: Optional[List] = None
+    input_top_logprobs_idx: Optional[List] = None
+    input_token_ids_logprobs_val: Optional[List] = None
+    input_token_ids_logprobs_idx: Optional[List] = None
+
+
+def get_top_logprobs_raw(
+    logprobs: torch.Tensor,
+    top_logprobs_nums: List[int],
+    stage: LogprobStage,
+    extend_logprob_pruned_lens_cpu: Optional[List[int]] = None,
+    no_copy_to_cpu: bool = False,
+):
+    max_k = max(top_logprobs_nums)
+    values, indices = logprobs.topk(max_k, dim=-1)
+    if not no_copy_to_cpu:
+        values = values.tolist()
+        indices = indices.tolist()
+
+    top_logprobs_val = []
+    top_logprobs_idx = []
+
+    if stage == LogprobStage.DECODE:
+        for i, k in enumerate(top_logprobs_nums):
+            top_logprobs_val.append(values[i][:k])
+            top_logprobs_idx.append(indices[i][:k])
+    else:
+        pt = 0
+        for k, pruned_len in zip(top_logprobs_nums, extend_logprob_pruned_lens_cpu):
+            if pruned_len <= 0:
+                top_logprobs_val.append([])
+                top_logprobs_idx.append([])
+                continue
+
+            top_logprobs_val.append([values[pt + j][:k] for j in range(pruned_len)])
+            top_logprobs_idx.append([indices[pt + j][:k] for j in range(pruned_len)])
+            pt += pruned_len
+
+    return top_logprobs_val, top_logprobs_idx
+
+
+def get_top_logprobs_prefill(
+    all_logprobs: torch.Tensor, logits_metadata: LogitsMetadata
+):
+    return get_top_logprobs_raw(
+        all_logprobs,
+        logits_metadata.top_logprobs_nums,
+        stage=LogprobStage.PREFILL,
+        extend_logprob_pruned_lens_cpu=logits_metadata.extend_logprob_pruned_lens_cpu,
+    )
+
+
+def get_top_logprobs(
+    logprobs: torch.Tensor,
+    top_logprobs_nums: List[int],
+    no_copy_to_cpu: bool = False,
+):
+    return get_top_logprobs_raw(
+        logprobs,
+        top_logprobs_nums,
+        stage=LogprobStage.DECODE,
+        no_copy_to_cpu=no_copy_to_cpu,
+    )
+
+
+def get_token_ids_logprobs_raw(
+    logprobs: torch.Tensor,
+    token_ids_logprobs_list: List[Optional[List[int]]],
+    stage: LogprobStage,
+    extend_logprob_pruned_lens_cpu: Optional[List[int]] = None,
+    no_copy_to_cpu: bool = False,
+):
+    vals, idxs = [], []
+    if stage == LogprobStage.DECODE:
+        for i, token_ids in enumerate(token_ids_logprobs_list):
+            if token_ids is None:
+                vals.append([])
+                idxs.append([])
+            else:
+                token_ids_tensor = torch.tensor(token_ids, dtype=torch.long).to(
+                    logprobs.device, non_blocking=True
+                )
+                row = logprobs[i, token_ids_tensor]
+                vals.append(row if no_copy_to_cpu else row.tolist())
+                idxs.append(token_ids)
+    else:  # prefill
+        pt = 0
+        for i, (token_ids, pruned_len) in enumerate(
+            zip(token_ids_logprobs_list, extend_logprob_pruned_lens_cpu)
+        ):
+            if pruned_len <= 0:
+                vals.append([])
+                idxs.append([])
+                continue
+            token_ids_tensor = torch.tensor(token_ids, dtype=torch.long).to(
+                logprobs.device, non_blocking=True
+            )
+            pos_logprobs = logprobs[pt : pt + pruned_len, token_ids_tensor]
+            vals.append(pos_logprobs if no_copy_to_cpu else pos_logprobs.tolist())
+            idxs.append([token_ids for _ in range(pruned_len)])
+            pt += pruned_len
+    return vals, idxs
+
+
+def get_token_ids_logprobs_prefill(
+    all_logprobs, logits_metadata: LogitsMetadata, no_copy_to_cpu=False
+):
+    return get_token_ids_logprobs_raw(
+        all_logprobs,
+        logits_metadata.token_ids_logprobs,
+        stage=LogprobStage.PREFILL,
+        extend_logprob_pruned_lens_cpu=logits_metadata.extend_logprob_pruned_lens_cpu,
+        no_copy_to_cpu=no_copy_to_cpu,
+    )
+
+
+def get_token_ids_logprobs(logprobs, token_ids_logprobs, no_copy_to_cpu=False):
+    return get_token_ids_logprobs_raw(
+        logprobs,
+        token_ids_logprobs,
+        stage=LogprobStage.DECODE,
+        no_copy_to_cpu=no_copy_to_cpu,
+    )
+
+
+def get_top_logprobs_chunk(
+    logprobs: torch.Tensor,
+    logits_metadata: LogitsMetadata,
+    top_k_nums: List[int],
+    pruned_lens: List[int],
+    input_top_logprobs_val: List,
+    input_top_logprobs_idx: List,
+    split_pruned_len: int,
+) -> int:
+    """Get top-k logprobs for each sequence in the chunk.
+
+    Args:
+        logprobs: Log probabilities tensor of shape [seq_len, vocab_size]
+        logits_metadata: Metadata containing top-k and pruned length info
+        top_k_nums: List of top-k numbers for each sequence
+        pruned_lens: List of pruned lengths for each sequence
+        input_top_logprobs_val: List to store top-k logprob values
+        input_top_logprobs_idx: List to store top-k token indices
+        split_pruned_len: Length of pruned tokens from previous chunk
+
+    Returns:
+        int: Number of remaining tokens to process in next chunk
+    """
+    # No sequences in the chunk
+    if logprobs.shape[0] == 0:
+        return 0
+
+    max_k = max(logits_metadata.top_logprobs_nums)
+    ret = logprobs.topk(max_k, dim=1)
+    values = ret.values.tolist()
+    indices = ret.indices.tolist()
+
+    pt = 0
+    next_split_pruned_len = 0
+    for n, (k, pruned_len) in enumerate(zip(top_k_nums, pruned_lens)):
+        if n == 0:
+            # For the first sequence, adjust the pruned length
+            pruned_len -= split_pruned_len
+        else:
+            # After the first sequence, no split in the middle
+            split_pruned_len = 0
+
+        if pruned_len <= 0:
+            # if pruned length is less than or equal to 0,
+            # there is no top-k logprobs to process
+            input_top_logprobs_val.append([])
+            input_top_logprobs_idx.append([])
+            continue
+
+        # Get the top-k logprobs
+        val = []
+        idx = []
+        for j in range(pruned_len):
+            # Handle remaining tokens in next chunk if any
+            if pt + j >= len(values):
+                next_split_pruned_len = split_pruned_len + j
+                break
+            # Append the top-k logprobs
+            val.append(values[pt + j][:k])
+            idx.append(indices[pt + j][:k])
+
+        # Append or extend based on whether the sequence was split across chunks
+        if len(val) > 0:
+            if split_pruned_len > 0:
+                input_top_logprobs_val[-1].extend(val)
+                input_top_logprobs_idx[-1].extend(idx)
+            else:
+                input_top_logprobs_val.append(val)
+                input_top_logprobs_idx.append(idx)
+
+        pt += pruned_len
+    return next_split_pruned_len
+
+
+def get_token_ids_logprobs_chunk(
+    logprobs: torch.Tensor,
+    token_ids_logprobs: List[int],
+    pruned_lens: List[int],
+    input_token_ids_logprobs_val: List,
+    input_token_ids_logprobs_idx: List,
+    split_pruned_len: int = 0,
+):
+    """Get token_ids logprobs for each sequence in the chunk.
+
+    Args:
+        logprobs: Log probabilities tensor of shape [seq_len, vocab_size]
+        logits_metadata: Metadata containing token IDs and pruned length info
+        token_ids_logprobs: List of token IDs for each sequence
+        pruned_lens: List of pruned lengths for each sequence
+        input_token_ids_logprobs_val: List to store token logprob values
+        input_token_ids_logprobs_idx: List to store token indices
+        split_pruned_len: Length of pruned tokens from previous chunk
+
+    Returns:
+        int: Number of remaining tokens to process in next chunk
+    """
+
+    # No sequences in the chunk
+    if logprobs.shape[0] == 0:
+        return 0
+
+    pt = 0
+    next_split_pruned_len = 0
+    for n, (token_ids, pruned_len) in enumerate(
+        zip(
+            token_ids_logprobs,
+            pruned_lens,
+        )
+    ):
+        # Adjust pruned length for first sequence
+        if n == 0:
+            pruned_len -= split_pruned_len
+        else:
+            split_pruned_len = 0
+
+        if pruned_len <= 0:
+            # if pruned length is less than or equal to 0,
+            # there is no token ids logprobs to process
+            input_token_ids_logprobs_val.append([])
+            input_token_ids_logprobs_idx.append([])
+            continue
+
+        # Get the token ids logprobs
+        val = []
+        idx = []
+        for j in range(pruned_len):
+            # Handle remaining tokens in next chunk if any
+            if pt + j >= logprobs.shape[0]:
+                next_split_pruned_len = split_pruned_len + j
+                break
+            if token_ids is not None:
+                val.append(logprobs[pt + j, token_ids].tolist())
+                idx.append(token_ids)
+
+        # Append or extend based on whether the sequence was split across chunks
+        if len(val) > 0:
+            if split_pruned_len > 0:
+                input_token_ids_logprobs_val[-1].extend(val)
+                input_token_ids_logprobs_idx[-1].extend(idx)
+            else:
+                input_token_ids_logprobs_val.append(val)
+                input_token_ids_logprobs_idx.append(idx)
+
+        pt += pruned_len
+    return next_split_pruned_len
+
+
+def compute_spec_v2_logprobs(
+    batch,
+    logits_output,
+    predict: torch.Tensor,
+    accept_index: torch.Tensor,
+    speculative_num_steps: int,
+):
+    """Compute logprobs for accepted tokens after spec v2 verify sampling.
+
+    Gathers logits at accepted positions, applies log_softmax (temperature-scaled
+    if not greedy), and populates logits_output.next_token_logprobs (plus optional
+    top-k / token-ids logprobs) so they flow through copy_to_cpu().
+    """
+    bs = len(batch.seq_lens)
+    max_accept = speculative_num_steps + 1
+    device = predict.device
+
+    flat_accept_idx = accept_index.long().reshape(-1)
+    gathered_logits = logits_output.next_token_logits[flat_accept_idx]
+
+    if batch.sampling_info.is_all_greedy or envs.SGLANG_RETURN_ORIGINAL_LOGPROB.get():
+        gathered_logprobs = torch.nn.functional.log_softmax(gathered_logits, dim=-1)
+    else:
+        temperatures = torch.repeat_interleave(
+            batch.sampling_info.temperatures,
+            max_accept,
+            dim=0,
+        )
+        gathered_logprobs = torch.nn.functional.log_softmax(
+            gathered_logits / temperatures, dim=-1
+        )
+    gathered_logprobs.clamp_(min=torch.finfo(gathered_logprobs.dtype).min)
+
+    accepted_token_ids = predict[flat_accept_idx]
+    token_logprobs = gathered_logprobs[
+        torch.arange(bs * max_accept, device=device),
+        accepted_token_ids.long(),
+    ]
+    logits_output.next_token_logprobs = token_logprobs.reshape(bs, max_accept)
+
+    if batch.top_logprobs_nums and any(x > 0 for x in batch.top_logprobs_nums):
+        top_logprobs_nums_expanded = [
+            num for num in batch.top_logprobs_nums for _ in range(max_accept)
+        ]
+        (
+            logits_output.next_token_top_logprobs_val,
+            logits_output.next_token_top_logprobs_idx,
+        ) = get_top_logprobs(
+            gathered_logprobs, top_logprobs_nums_expanded, no_copy_to_cpu=True
+        )
+
+    if batch.token_ids_logprobs and any(
+        x is not None for x in batch.token_ids_logprobs
+    ):
+        token_ids_logprobs_expanded = [
+            ids for ids in batch.token_ids_logprobs for _ in range(max_accept)
+        ]
+        (
+            logits_output.next_token_token_ids_logprobs_val,
+            logits_output.next_token_token_ids_logprobs_idx,
+        ) = get_token_ids_logprobs(
+            gathered_logprobs, token_ids_logprobs_expanded, no_copy_to_cpu=True
+        )

python/sglang/srt/layers/utils/multi_platform.py

@@ -0,0 +1,134 @@
+from typing import Callable, ClassVar
+
+from torch import nn
+
+from sglang.kernel_api_logging import debug_kernel_api
+from sglang.srt.platforms import current_platform
+from sglang.srt.utils import (
+    cpu_has_amx_support,
+    is_cpu,
+    is_cuda,
+    is_hip,
+    is_musa,
+    is_npu,
+    is_xpu,
+)
+
+_is_cuda = is_cuda()
+_is_hip = is_hip()
+_is_cpu = is_cpu()
+_is_cpu_amx_available = cpu_has_amx_support()
+_is_npu = is_npu()
+_is_xpu = is_xpu()
+_is_musa = is_musa()
+
+
+class MultiPlatformOp(nn.Module):
+
+    # OOT forward registry: maps dispatch_key -> {op_cls -> forward_fn}
+    _oot_forward_registry: ClassVar[dict[str, dict[type, Callable]]] = {}
+
+    @classmethod
+    def register_oot_forward(cls, op_cls: type, fn: Callable, platform_key: str):
+        """Register an OOT forward implementation for a specific op class and platform."""
+        cls._oot_forward_registry.setdefault(platform_key, {})[op_cls] = fn
+
+    def __init__(self):
+        super().__init__()
+        self._forward_method: Callable = self.dispatch_forward()
+
+        # States for torch.compile
+        self._original_forward_method = None
+        self.is_torch_compile = False
+
+    def enter_torch_compile(self, num_tokens: int):
+        # Skip if Op is already entered compile mode.
+        # NOTE(alcanderian): Some Ops(for example RotaryEmbedding) will be reused
+        # among layers and `enter_torch_compile` will be called many times.
+        # We should prevent `self._original_forward_method` from being overridden when
+        # it is not the first time `enter_torch_compile` called.
+        if self.is_torch_compile:
+            return
+
+        self._original_forward_method = self._forward_method
+        # NOTE: Temporarily workaround MoE
+        # The performance of torch.compile on this layer is not always good when bs > 1,
+        # so we decide to only use torch.compile when bs=1
+        if "FusedMoE" in self.__class__.__name__:
+            if num_tokens == 1:
+                from sglang.srt.layers.moe.fused_moe_native import (
+                    fused_moe_forward_native,
+                )
+
+                self._forward_method = fused_moe_forward_native
+        elif "TopK" in self.__class__.__name__:
+            if num_tokens == 1:
+                self._forward_method = self.forward_native
+        else:
+            self._forward_method = self.forward_native
+        self.is_torch_compile = True
+
+    def leave_torch_compile(self):
+        # Skip if Op is already exited compile mode.
+        if not self.is_torch_compile:
+            return
+
+        self._forward_method = self._original_forward_method
+        self._original_forward_method = None
+        self.is_torch_compile = False
+
+    # Please do not override this method, because `self._forward_method` can change when in torch compile mode
+    @debug_kernel_api
+    def forward(self, *args, **kwargs):
+        return self._forward_method(*args, **kwargs)
+
+    def forward_native(self, *args, **kwargs):
+        raise NotImplementedError
+
+    def forward_cuda(self, *args, **kwargs):
+        raise NotImplementedError
+
+    def forward_npu(self, *args, **kwargs):
+        return self.forward_native(*args, **kwargs)
+
+    def forward_hip(self, *args, **kwargs):
+        return self.forward_cuda(*args, **kwargs)
+
+    def forward_xpu(self, *args, **kwargs):
+        return self.forward_native(*args, **kwargs)
+
+    def forward_musa(self, *args, **kwargs):
+        return self.forward_cuda(*args, **kwargs)
+
+    def forward_hpu(self, *args, **kwargs):
+        return self.forward_native(*args, **kwargs)
+
+    def forward_cpu(self, *args, **kwargs):
+        return self.forward_native(*args, **kwargs)
+
+    def dispatch_forward(self):
+        # OOT platform dispatch: check registry then method lookup
+        if current_platform.is_out_of_tree():
+            key = current_platform.get_dispatch_key_name()
+            oot = self._oot_forward_registry.get(key, {})
+            if type(self) in oot:
+                return oot[type(self)].__get__(self)
+            method = getattr(self, f"forward_{key}", None)
+            if method is not None:
+                return method
+            return self.forward_native
+
+        if _is_cuda:
+            return self.forward_cuda
+        elif _is_hip:
+            return self.forward_hip
+        elif _is_cpu and _is_cpu_amx_available:
+            return self.forward_cpu
+        elif _is_npu:
+            return self.forward_npu
+        elif _is_xpu:
+            return self.forward_xpu
+        elif _is_musa:
+            return self.forward_musa
+        else:
+            return self.forward_native