[ROCm] Support decode context parallel (DCP) for GLM-5 / DeepSeek-V3.2

python/sglang/kernels/ops/attention/dcp_kernels.py

@@ -231,6 +231,208 @@ def update_kv_lens_and_indices(
     tl.store(local_kv_indices + local_kv_indices_offsets, kv_values, mask=mask)
 
 
+# ---------------------------------------------------------------------------
+# Sparse (DSA) decode: this rank's share of a widened top-k slot table.
+# ---------------------------------------------------------------------------
+@triton.jit
+def _dcp_compact_owned_slots_kernel(
+    table_ptr,  # [bs, topk] widened slots, -1 = invalid
+    out_ptr,  # [bs, topk] local rows, owned first, -1 tail
+    len_ptr,  # [bs] number of owned slots
+    table_stride,
+    topk: tl.constexpr,
+    TOPK_POW2: tl.constexpr,
+    DCP_SIZE: tl.constexpr,
+    DCP_RANK: tl.constexpr,
+):
+    b = tl.program_id(0)
+    offs = tl.arange(0, TOPK_POW2)
+    in_range = offs < topk
+    slot = tl.load(table_ptr + b * table_stride + offs, mask=in_range, other=-1)
+    owned = (slot >= 0) & (slot % DCP_SIZE == DCP_RANK)
+    dst = tl.cumsum(owned.to(tl.int32), axis=0) - 1
+    n = tl.sum(owned.to(tl.int32), axis=0)
+    tl.store(out_ptr + b * topk + dst, (slot // DCP_SIZE).to(tl.int32), mask=owned)
+    tl.store(out_ptr + b * topk + offs, -1, mask=in_range & (offs >= n))
+    tl.store(len_ptr + b, n)
+
+
+def dcp_compact_owned_slots(
+    table: torch.Tensor, dcp_size: int, dcp_rank: int
+) -> Tuple[torch.Tensor, torch.Tensor]:
+    """Widened [bs, topk] slot table -> (local rows packed to the front, counts).
+
+    Owner rule: slot % dcp_size == dcp_rank, local row = slot // dcp_size. Score
+    order is kept; the tail past each row's count is -1.
+    """
+    if table.stride(1) != 1:
+        table = table.contiguous()
+    bs, topk = table.shape
+    out = torch.empty((bs, topk), dtype=torch.int32, device=table.device)
+    lens = torch.empty((bs,), dtype=torch.int32, device=table.device)
+    if bs > 0:
+        _dcp_compact_owned_slots_kernel[(bs,)](
+            table,
+            out,
+            lens,
+            table.stride(0),
+            topk=topk,
+            TOPK_POW2=triton.next_power_of_2(topk),
+            DCP_SIZE=dcp_size,
+            DCP_RANK=dcp_rank,
+        )
+    return out, lens
+
+
+@triton.jit
+def _dcp_topk_pack_kernel(
+    logits_ptr,
+    idx_ptr,
+    len_ptr,
+    row_start_ptr,
+    send_ptr,
+    send_int_ptr,
+    logits_stride,
+    idx_stride,
+    rows,
+    topk: tl.constexpr,
+    BLOCK: tl.constexpr,
+    DCP_SIZE: tl.constexpr,
+    DCP_RANK: tl.constexpr,
+    HAS_ROW_START: tl.constexpr,
+):
+    b = tl.program_id(0)
+    offs = tl.program_id(1) * BLOCK + tl.arange(0, BLOCK)
+    in_range = offs < topk
+    idx = tl.load(idx_ptr + b * idx_stride + offs, mask=in_range, other=-1)
+    valid = (idx >= 0) & (idx < tl.load(len_ptr + b))
+    col = idx
+    if HAS_ROW_START:
+        col = idx + tl.load(row_start_ptr + b)
+    score = tl.load(
+        logits_ptr + b.to(tl.int64) * logits_stride + col,
+        mask=in_range & valid,
+        other=float("-inf"),
+    )
+    gid = tl.where(valid, idx * DCP_SIZE + DCP_RANK, -1)
+    tl.store(send_ptr + b.to(tl.int64) * topk + offs, score, mask=in_range)
+    tl.store(send_int_ptr + (rows + b).to(tl.int64) * topk + offs, gid, mask=in_range)
+
+
+def dcp_topk_pack(
+    logits: torch.Tensor,
+    local_idx: torch.Tensor,
+    local_lens: torch.Tensor,
+    dcp_size: int,
+    dcp_rank: int,
+    row_starts: Optional[torch.Tensor] = None,
+) -> torch.Tensor:
+    """Local top-k -> [2, rows, topk] fp32 send plane: (score, global pos as int32 bits).
+
+    Invalid picks get score -inf and position -1; global pos = local * W + rank.
+    ``local_idx`` is relative to ``row_starts`` (packed prefill logits) when given.
+    """
+    rows, topk = local_idx.shape
+    send = torch.empty((2, rows, topk), dtype=torch.float32, device=logits.device)
+    if rows > 0:
+        block = min(1024, triton.next_power_of_2(topk))
+        _dcp_topk_pack_kernel[(rows, triton.cdiv(topk, block))](
+            logits,
+            local_idx,
+            local_lens,
+            row_starts if row_starts is not None else local_lens,
+            send,
+            send.view(torch.int32),
+            logits.stride(0),
+            local_idx.stride(0),
+            rows,
+            topk=topk,
+            BLOCK=block,
+            DCP_SIZE=dcp_size,
+            DCP_RANK=dcp_rank,
+            HAS_ROW_START=row_starts is not None,
+        )
+    return send
+
+
+@triton.jit
+def _dcp_topk_gather_scores_kernel(
+    recv_ptr,
+    scores_ptr,
+    rows,
+    topk: tl.constexpr,
+    BLOCK: tl.constexpr,
+    DCP_SIZE: tl.constexpr,
+):
+    b = tl.program_id(0)
+    col = tl.program_id(1) * BLOCK + tl.arange(0, BLOCK)
+    in_range = col < DCP_SIZE * topk
+    src = col // topk
+    score = tl.load(
+        recv_ptr + (src * 2 * rows + b).to(tl.int64) * topk + col % topk,
+        mask=in_range,
+    )
+    tl.store(scores_ptr + b.to(tl.int64) * DCP_SIZE * topk + col, score, mask=in_range)
+
+
+@triton.jit
+def _dcp_topk_finalize_kernel(
+    recv_ptr,
+    recv_int_ptr,
+    pick_ptr,
+    out_ptr,
+    pick_stride,
+    rows,
+    topk: tl.constexpr,
+    BLOCK: tl.constexpr,
+):
+    b = tl.program_id(0)
+    offs = tl.program_id(1) * BLOCK + tl.arange(0, BLOCK)
+    in_range = offs < topk
+    pick = tl.load(
+        pick_ptr + b.to(tl.int64) * pick_stride + offs, mask=in_range, other=-1
+    )
+    ok = in_range & (pick >= 0)
+    pick = tl.where(ok, pick, 0)
+    base = ((pick // topk) * 2 * rows + b).to(tl.int64) * topk + pick % topk
+    score = tl.load(recv_ptr + base, mask=ok, other=float("-inf"))
+    gid = tl.load(recv_int_ptr + base + rows * topk, mask=ok, other=-1)
+    tl.store(
+        out_ptr + b.to(tl.int64) * topk + offs,
+        tl.where(ok & (score > float("-inf")), gid, -1),
+        mask=in_range,
+    )
+
+
+def dcp_topk_merge(recv: torch.Tensor, dcp_size: int, topk_func) -> torch.Tensor:
+    """All-gathered [W * 2, rows, topk] send planes -> global top-k positions (-1 padded)."""
+    _, rows, topk = recv.shape
+    recv = recv.contiguous()
+    width = dcp_size * topk
+    scores = torch.empty((rows, width), dtype=torch.float32, device=recv.device)
+    out = torch.empty((rows, topk), dtype=torch.int32, device=recv.device)
+    if rows == 0:
+        return out
+    block = 1024
+    _dcp_topk_gather_scores_kernel[(rows, triton.cdiv(width, block))](
+        recv, scores, rows, topk=topk, BLOCK=block, DCP_SIZE=dcp_size
+    )
+    lens = torch.full((rows,), width, dtype=torch.int32, device=recv.device)
+    pick = topk_func(scores, lens, topk)
+    block = min(1024, triton.next_power_of_2(topk))
+    _dcp_topk_finalize_kernel[(rows, triton.cdiv(topk, block))](
+        recv,
+        recv.view(torch.int32),
+        pick,
+        out,
+        pick.stride(0),
+        rows,
+        topk=topk,
+        BLOCK=block,
+    )
+    return out
+
+
 # ---------------------------------------------------------------------------
 # Partial-attention LSE correction (PR #14194, MLA path).
 # ---------------------------------------------------------------------------

python/sglang/kernels/ops/attention/dsa/triton_sparse_mla.py

@@ -983,6 +983,8 @@ def _sparse_mla_reduce_kernel(
     ACTIVE_SPLITS_POW2: tl.constexpr,
     D_CHUNK: tl.constexpr,
     BLOCK_K: tl.constexpr,
+    lse_out_ptr=None,  # [seq, H] fp32 natural-log LSE; written when STORE_LSE
+    STORE_LSE: tl.constexpr = False,
 ):
     """Reduce split-K partials via log-space combine. grid=(seq, H, d_v_chunks)."""
     t = tl.program_id(0)
@@ -1028,6 +1030,11 @@ def _sparse_mla_reduce_kernel(
         out.to(tl.bfloat16),
         mask=d_mask,
     )
+    if STORE_LSE:
+        if dc == 0:
+            # Partials are base-2 over log2(e)-scaled scores; * ln(2) is base-e.
+            lse = (lse_max + tl.log2(tl.maximum(w_sum, 1.0e-30))) * 0.6931471805599453
+            tl.store(lse_out_ptr + t * H + h, lse)
 
 
 def _triton_sparse_mla_fwd_splitk(

python/sglang/kernels/ops/attention/dsa/triton_sparse_mla_decode.py

@@ -310,6 +310,9 @@ def _sparse_mla_decode_split_kernel(
     KV_SPLITS: tl.constexpr,
     BLOCK_H: tl.constexpr,
     BLOCK_K: tl.constexpr,
+    len_ptr=None,  # [N] valid prefix of each idx row; read when USE_LENGTH
+    USE_LENGTH: tl.constexpr = False,
+    LSE_SCALE: tl.constexpr = 1.0,  # ln(2) when a lone split stores the final LSE
 ):
     t = tl.program_id(0)
     pid_h = tl.program_id(1)
@@ -357,12 +360,20 @@ def _sparse_mla_decode_split_kernel(
         other=0.0,
     ).to(input_type)
 
-    tiles_per_segment = tl.cdiv(topk, KV_SPLITS * BLOCK_K)
-    if pid_k * tiles_per_segment * BLOCK_K >= topk:
-        return
-    num_tiles = tl.cdiv(topk, BLOCK_K)
-    tile_start = pid_k * tiles_per_segment
-    tile_end = tl.minimum((pid_k + 1) * tiles_per_segment, num_tiles)
+    if USE_LENGTH:
+        # Spread only the row's valid tiles over every split; a split left with
+        # none still stores an empty partial, since the reduce reads all splits.
+        row_tiles = tl.cdiv(tl.load(len_ptr + t), BLOCK_K)
+        tiles_per_segment = tl.cdiv(row_tiles, KV_SPLITS)
+        tile_start = pid_k * tiles_per_segment
+        tile_end = tl.minimum(tile_start + tiles_per_segment, row_tiles)
+    else:
+        tiles_per_segment = tl.cdiv(topk, KV_SPLITS * BLOCK_K)
+        if pid_k * tiles_per_segment * BLOCK_K >= topk:
+            return
+        num_tiles = tl.cdiv(topk, BLOCK_K)
+        tile_start = pid_k * tiles_per_segment
+        tile_end = tl.minimum((pid_k + 1) * tiles_per_segment, num_tiles)
 
     neg_large = -3.4028234663852886e38
     m_i = tl.full((BLOCK_H,), neg_large, dtype=tl.float32)
@@ -464,7 +475,7 @@ def _sparse_mla_decode_split_kernel(
     if NUM_GROUPS >= 4:
         acc3 = tl.where(has_data[:, None], acc3 * inv_denom[:, None], 0.0)
 
-    lse = tl.where(has_data, tl.log2(l_i) + m_i, neg_large)
+    lse = tl.where(has_data, tl.log2(l_i) + m_i, neg_large) * LSE_SCALE
 
     H_padded = tl.cdiv(H, BLOCK_H) * BLOCK_H
     lse_base = t * KV_SPLITS * H_padded + pid_k * H_padded
@@ -505,15 +516,21 @@ def triton_sparse_mla_decode_splitk(
     d_v: int = 512,
     kv_splits: int | None = None,
     workspace: _SplitKWorkspace | None = None,
-) -> torch.Tensor:
+    return_lse: bool = False,
+    lengths: torch.Tensor | None = None,
+) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
     """Split-K Triton sparse MLA decode (DSv4 pattern).
 
     q_nope:  [bs, H, d_v] fp8/bf16
     q_rope:  [bs, H, d_tail] fp8/bf16
     kv:      [num_pages, 1, DIM] fp8/bf16
     indices: [bs, 1, topk] int32
     workspace: backend-owned grow-only split-K buffers for graph-safe reuse
-    returns: [1, bs, H, d_v] bf16
+    lengths: optional [bs] int32; row b attends only indices[b, 0, :lengths[b]],
+        and the splits share that prefix instead of the full topk width
+    returns: [1, bs, H, d_v] bf16, plus a natural-log [bs, H] fp32 LSE when
+        ``return_lse``. A row with no valid index gets a zero output and a large
+        finite negative LSE, so it carries no weight in a cross-rank merge.
     """
     is_fp8 = _validate_input_dtypes(q_nope, q_rope, kv)
     use_fp8_dot = is_fp8
@@ -583,7 +600,8 @@ def triton_sparse_mla_decode_splitk(
 
     qk_scale = float(sm_scale) * LOG2E
 
-    if kv_splits == 1:
+    # The fused one-split kernel takes neither an LSE output nor row lengths.
+    if kv_splits == 1 and not return_lse and lengths is None:
         fused_num_warps = 4
         if optimize_gfx950_fp8:
             fused_num_warps = _gfx950_sparse_mla_num_warps(base_ctas, 1, num_cu)
@@ -626,14 +644,23 @@ def triton_sparse_mla_decode_splitk(
         tiles_per_split * BLOCK_K
     )
     active_splits = min(active_splits, kv_splits)
+    if lengths is not None:
+        active_splits = kv_splits
     split_num_warps = 4
     if optimize_gfx950_fp8:
         split_num_warps = _gfx950_sparse_mla_num_warps(base_ctas, active_splits, num_cu)
 
-    lse_partial, acc_partial = _get_splitk_bufs(
-        bs, kv_splits, h_padded, d_v, q_nope.device, workspace
-    )
     out = torch.empty(bs, H, d_v, device=q_nope.device, dtype=torch.bfloat16)
+    # A lone split already holds the final output, laid out as [N, 1, H, D_V]
+    # when no head padding is needed; store it and the LSE in place.
+    direct = kv_splits == 1 and return_lse and h_padded == H
+    if direct:
+        lse = torch.empty(bs, H, device=q_nope.device, dtype=torch.float32)
+        lse_partial, acc_partial = lse, out
+    else:
+        lse_partial, acc_partial = _get_splitk_bufs(
+            bs, kv_splits, h_padded, d_v, q_nope.device, workspace
+        )
 
     grid_split = (bs, n_head_blocks, kv_splits)
     with _no_async_copy():
@@ -660,12 +687,22 @@ def triton_sparse_mla_decode_splitk(
             KV_SPLITS=kv_splits,
             BLOCK_H=BLOCK_H,
             BLOCK_K=BLOCK_K,
+            len_ptr=lengths,
+            USE_LENGTH=lengths is not None,
+            LSE_SCALE=0.6931471805599453 if direct else 1.0,
             num_warps=split_num_warps,
             num_stages=2,
         )
+    if direct:
+        return out.unsqueeze(0), lse
 
     D_CHUNK = _reduce_d_chunk(active_splits, bs * H) if optimize_gfx950_fp8 else 64
     grid_reduce = (bs, H, (d_v + D_CHUNK - 1) // D_CHUNK)
+    lse = (
+        torch.empty(bs, H, device=q_nope.device, dtype=torch.float32)
+        if return_lse
+        else None
+    )
     _sparse_mla_reduce_kernel[grid_reduce](
         lse_partial,
         acc_partial,
@@ -677,9 +714,13 @@ def triton_sparse_mla_decode_splitk(
         ACTIVE_SPLITS_POW2=_next_pow2(active_splits),
         D_CHUNK=D_CHUNK,
         BLOCK_K=BLOCK_K,
+        lse_out_ptr=lse,
+        STORE_LSE=return_lse,
         num_warps=4,
     )
-    return out.unsqueeze(0)
+    if not return_lse:
+        return out.unsqueeze(0)
+    return out.unsqueeze(0), lse
 
 
 # ---------------------------------------------------------------------------