diff --git a/python/sglang/kernels/ops/attention/minimax_sparse/decode/verify_topk_gfx950.py b/python/sglang/kernels/ops/attention/minimax_sparse/decode/verify_topk_gfx950.py index 9927e1128cdd..5f424acddc4a 100644 --- a/python/sglang/kernels/ops/attention/minimax_sparse/decode/verify_topk_gfx950.py +++ b/python/sglang/kernels/ops/attention/minimax_sparse/decode/verify_topk_gfx950.py @@ -1,15 +1,15 @@ """Fused target-verify indexer for gfx950: per-block max index score with a streaming top-k. -A request's draft rows (num_draft x index heads, at most 16) form one MFMA tile. Each -warp of a program streams its own index-K blocks of one context split and keeps every -row's best ``KP`` blocks in registers; the program reduces its warps' lists to the -split's best ``KP``, and a merge kernel reduces the splits' lists, adds the forced -init/local blocks and writes ascending ids. The scores are the 16x16x32 MFMA products -and scaling of ``_verify_score_kernel`` (bf16), or for an NVFP4 cache of +A request's draft rows (num_draft x index heads, at most 16; 32 under indexer CP) form +one MFMA tile. Each warp of a program streams its own index-K blocks of one context +split and keeps every row's best ``KP`` blocks in registers; the program reduces its +warps' lists to the split's best ``KP``, and a merge kernel reduces the splits' lists, +adds the forced init/local blocks and writes ascending ids. The scores are the 16x16x32 +MFMA products and scaling of ``_verify_score_kernel`` (bf16), or for an NVFP4 cache of ``_verify_score_nvfp4_kernel`` (fp16, even dims then odd dims). Blocks rank by score, -then lower block id, in the radix top-k's order; every list keeps its blocks' top -``KP`` in that order, so the ids equal the score + top-k path's whatever the split of -the blocks. +then lower block id, in the radix top-k's order; every list keeps its blocks' top ``KP`` +in that order, so the ids equal the score + top-k path's whatever the split of the +blocks. """ import torch @@ -54,6 +54,16 @@ def _add(a, b): return a + b +@gluon.jit +def _local_range(init_blocks, req_hi, WORLD: gl.constexpr, RANK: gl.constexpr): + """Local ids [lo, hi) of blocks l * WORLD + RANK in [init_blocks, req_hi).""" + if WORLD == 1: + return init_blocks, req_hi + lo = (gl.maximum(init_blocks - RANK, 0) + WORLD - 1) // WORLD + hi = (gl.maximum(req_hi - RANK, 0) + WORLD - 1) // WORLD + return lo, gl.maximum(hi, lo) + + @gluon.jit def _split_plan( seq_lens_ptr, @@ -70,6 +80,8 @@ def _split_plan( MAX_SPLITS: gl.constexpr, MIN_SPLIT: gl.constexpr, W: gl.constexpr, + WORLD: gl.constexpr = 1, + RANK: gl.constexpr = 0, ): """This program's (request, split, request's splits, candidate end, req_to_token row). @@ -77,7 +89,8 @@ def _split_plan( splits of at most S blocks (at least one split, at most MAX_SPLITS), with S = max(cdiv(total, PROGRAMS), MIN_SPLIT) when those splits fit in PROGRAMS, else the same over PROGRAMS - num_reqs, which always fits; the programs take the splits - in request order, so a long request gets more programs than a short one. + in request order, so a long request gets more programs than a short one. With + WORLD > 1 only this rank's blocks (RANK, RANK + WORLD, ...) count. """ plan_l: gl.constexpr = gl.BlockedLayout([1, ND_P2], [64, 1], [W, 1], [1, 0]) r = gl.arange(0, R_P2, layout=gl.SliceLayout(1, plan_l)) @@ -91,7 +104,8 @@ def _split_plan( gl.where(ok, gl.maximum(n_blk - local_blocks, init_blocks), init_blocks), axis=1 ) slots = gl.load(slot_ids_ptr + r * NUM_DRAFT, mask=valid, other=0) - blocks = gl.where(valid, req_hi - init_blocks, 0) + lo, hi = _local_range(init_blocks, req_hi, WORLD, RANK) + blocks = gl.where(valid, hi - lo, 0) total = gl.sum(blocks, axis=0) size = gl.maximum(gl.cdiv(total, PROGRAMS), MIN_SPLIT) n = gl.where(valid, gl.minimum(gl.maximum(gl.cdiv(blocks, size), 1), MAX_SPLITS), 0) @@ -200,8 +214,10 @@ def _score_round( M: gl.constexpr, BLOCK_N: gl.constexpr, K_FP4: gl.constexpr, + WORLD: gl.constexpr = 1, + RANK: gl.constexpr = 0, ): - """Score round ``it`` (block b0 + it * W + w for warp w) into each row's list. + """Score round ``it`` (local block b0 + it * W + w for warp w) into each row's list. Without ``K_FP4``, ``q_even`` is the whole Q operand and ``q_odd`` is unused. """ @@ -219,9 +235,10 @@ def _score_round( # NaN ranks below every real score yet above an empty entry, as the # radix top-k's NaN -> -inf does against its unscored tail. s = gl.where(s != s, -3.4028234663852886e38, s) - blk = b0 + it * W + warp + local = b0 + it * W + warp + blk = local * WORLD + RANK s = gl.where( - (blk[:, None] < b1) & (blk[:, None] < cand_hi[None, :]), + (local[:, None] < b1) & (blk[:, None] < cand_hi[None, :]), s, float("-inf"), ) @@ -274,11 +291,17 @@ def _verify_index_topk_partial_kernel( R_P2: gl.constexpr, ND_P2: gl.constexpr, num_warps: gl.constexpr, + # Rows of a request's tile (NUM_DRAFT * gqa_group_size, padded); 32 fits the + # indexer CP's four gathered heads, as two 16-row MFMA tiles of the same K. + M: gl.constexpr = 16, + # Indexer CP: this rank scores local blocks l as physical blocks l * WORLD + RANK. + WORLD: gl.constexpr = 1, + RANK: gl.constexpr = 0, ): - M: gl.constexpr = 16 req, split, n_split, req_hi, slot = _split_plan( seq_lens_ptr, slot_ids_ptr, num_reqs, max_blocks, init_blocks, local_blocks, NUM_DRAFT, BLOCK_N, R_P2, ND_P2, PROGRAMS, MAX_SPLITS, W * MIN_BLOCKS, W, + WORLD, RANK, ) # fmt: skip if req >= num_reqs: return @@ -332,9 +355,11 @@ def _verify_index_topk_partial_kernel( cand_hi = gl.where(m_ok, gl.maximum(n_blk - local_blocks, init_blocks), 0) if (split == 0) & (pid_kh == 0): gl.store(part_n_ptr + req, n_split) - per_split = gl.cdiv(req_hi - init_blocks, n_split) - b0 = init_blocks + split * per_split - b1 = gl.minimum(req_hi, b0 + per_split) + # b0, b1 and the round indices are local ids (the blocks themselves for one rank) + lo, hi = _local_range(init_blocks, req_hi, WORLD, RANK) + per_split = gl.cdiv(hi - lo, n_split) + b0 = lo + split * per_split + b1 = gl.minimum(hi, b0 + per_split) qw = gl.arange(0, W, layout=gl.SliceLayout(1, gl.SliceLayout(2, q_load))) qm = gl.arange(0, M, layout=gl.SliceLayout(0, gl.SliceLayout(2, q_load))) @@ -389,7 +414,7 @@ def _verify_index_topk_partial_kernel( # the current round is scored. n_rounds = gl.cdiv(b1 - b0, W) if n_rounds > 0: - base = gl.load(r2t_row + gl.minimum(b0 + kw, b1 - 1) * BLOCK_N) + base = gl.load(r2t_row + (gl.minimum(b0 + kw, b1 - 1) * WORLD + RANK) * BLOCK_N) k_next = _load_k_tile( k_cache_ptr, k_scale_ptr, @@ -404,8 +429,12 @@ def _verify_index_topk_partial_kernel( # Unrolled by two with two register sets, so no prefetched tile is # register-copied (the copy would wait on its load). An odd last round # rereads block b1 - 1 and scores nothing. - base_b = gl.load(r2t_row + gl.minimum(b0 + W + kw, b1 - 1) * BLOCK_N) - base = gl.load(r2t_row + gl.minimum(b0 + 2 * W + kw, b1 - 1) * BLOCK_N) + base_b = gl.load( + r2t_row + (gl.minimum(b0 + W + kw, b1 - 1) * WORLD + RANK) * BLOCK_N + ) + base = gl.load( + r2t_row + (gl.minimum(b0 + 2 * W + kw, b1 - 1) * WORLD + RANK) * BLOCK_N + ) for it in range(0, n_rounds, 2): k_b = _load_k_tile( k_cache_ptr, @@ -418,12 +447,14 @@ def _verify_index_topk_partial_kernel( K_FP4, ) base_b = gl.load( - r2t_row + gl.minimum(b0 + (it + 3) * W + kw, b1 - 1) * BLOCK_N + r2t_row + + (gl.minimum(b0 + (it + 3) * W + kw, b1 - 1) * WORLD + RANK) + * BLOCK_N ) top_k = _score_round( k_next, it, q_even, q_odd, grp, b0, b1, warp, cand_hi, qk_scale, top_k, scale_layout, group_layout, mfma_layout, dot_k, - row_layout, M, BLOCK_N, K_FP4, + row_layout, M, BLOCK_N, K_FP4, WORLD, RANK, ) # fmt: skip k_next = _load_k_tile( k_cache_ptr, @@ -436,15 +467,19 @@ def _verify_index_topk_partial_kernel( K_FP4, ) base = gl.load( - r2t_row + gl.minimum(b0 + (it + 4) * W + kw, b1 - 1) * BLOCK_N + r2t_row + + (gl.minimum(b0 + (it + 4) * W + kw, b1 - 1) * WORLD + RANK) + * BLOCK_N ) top_k = _score_round( k_b, it + 1, q_even, q_odd, grp, b0, b1, warp, cand_hi, qk_scale, top_k, scale_layout, group_layout, mfma_layout, dot_k, - row_layout, M, BLOCK_N, K_FP4, + row_layout, M, BLOCK_N, K_FP4, WORLD, RANK, ) # fmt: skip else: - base = gl.load(r2t_row + gl.minimum(b0 + W + kw, b1 - 1) * BLOCK_N) + base = gl.load( + r2t_row + (gl.minimum(b0 + W + kw, b1 - 1) * WORLD + RANK) * BLOCK_N + ) for it in range(n_rounds): k = k_next k_next = _load_k_tile( @@ -458,12 +493,14 @@ def _verify_index_topk_partial_kernel( K_FP4, ) base = gl.load( - r2t_row + gl.minimum(b0 + (it + 2) * W + kw, b1 - 1) * BLOCK_N + r2t_row + + (gl.minimum(b0 + (it + 2) * W + kw, b1 - 1) * WORLD + RANK) + * BLOCK_N ) top_k = _score_round( k, it, q_even, q_odd, grp, b0, b1, warp, cand_hi, qk_scale, top_k, scale_layout, group_layout, mfma_layout, dot_k, - row_layout, M, BLOCK_N, K_FP4, + row_layout, M, BLOCK_N, K_FP4, WORLD, RANK, ) # fmt: skip # The split's best KP per row from its W warp lists, best first. @@ -485,6 +522,30 @@ def _verify_index_topk_partial_kernel( gl.store(part_k_ptr + out, best, mask=(m2 < NUM_DRAFT * gqa_group_size)[:, None]) +@triton.jit +def _select_block_ids(key, n_blocks, topk, init_blocks, local_blocks, KT: tl.constexpr): + """A row's topk block ids: the forced init and local blocks, then the best candidate + keys (NO_KEY = empty); ascending, with the -1 tail the sparse decode kernels expect.""" + NO_KEY: tl.constexpr = -9223372036854775808 + BIG: tl.constexpr = 0x7FFFFFFF + n_init = tl.minimum(init_blocks, n_blocks) + local_start = tl.maximum(tl.maximum(n_blocks - local_blocks, 0), n_init) + n_forced = n_init + n_blocks - local_start + t = tl.arange(0, KT) + kept = tl.full([KT], BIG, tl.int32) + for j in tl.static_range(KT): + if j < topk - n_forced: + best = tl.max(key, axis=0) + pick = tl.where(best != NO_KEY, (~best).to(tl.int32), BIG) + kept = tl.where(t == j, pick, kept) + key = tl.where(key == best, NO_KEY, key) + forced = tl.where( + t < n_forced, tl.where(t < n_init, t, local_start + t - n_init), BIG + ) + ids = -tl.topk(-tl.reshape(tl.join(kept, forced), [2 * KT]), KT) + return tl.where(ids == BIG, -1, ids) + + @triton.jit def _verify_index_topk_merge_kernel( part_k_ptr, # [num_heads * rows, NC] int64 keys @@ -519,27 +580,8 @@ def _verify_index_topk_merge_kernel( ) # Empty entries have negative ids. key = tl.where((~key).to(tl.int32) >= 0, key, NO_KEY) - BIG: tl.constexpr = 0x7FFFFFFF - - # Forced blocks: init [0, init_blocks) and local [n_blocks - local_blocks, n_blocks). - n_init = tl.minimum(init_blocks, n_blocks) - local_start = tl.maximum(tl.maximum(n_blocks - local_blocks, 0), n_init) - n_forced = n_init + n_blocks - local_start + ids = _select_block_ids(key, n_blocks, topk, init_blocks, local_blocks, KT) t = tl.arange(0, KT) - kept = tl.full([KT], BIG, tl.int32) - # Highest score first, lower block id on ties. - for j in tl.static_range(KT): - if j < topk - n_forced: - best = tl.max(key, axis=0) - pick = tl.where(best != NO_KEY, (~best).to(tl.int32), BIG) - kept = tl.where(t == j, pick, kept) - key = tl.where(key == best, NO_KEY, key) - forced = tl.where( - t < n_forced, tl.where(t < n_init, t, local_start + t - n_init), BIG - ) - # Ascending ids with the -1 tail the sparse decode kernels expect. - ids = -tl.topk(-tl.reshape(tl.join(kept, forced), [2 * KT]), KT) - ids = tl.where(ids == BIG, -1, ids) tl.store(topk_ptr + r.to(tl.int64) * stride_t + t, ids, mask=t < topk) @@ -713,3 +755,381 @@ def fused_verify_topk_idx( num_warps=1, ) return out + + +def _cp_split_config(num_reqs: int) -> tuple[int, int, int, bool]: + """``_split_config`` for a rank's quarter of the blocks, every head's rows in one tile: + one wave of short splits up to 96 requests, then two waves of longer splits.""" + if num_reqs <= 96: + return 256, 16, 4, num_reqs <= 64 + return max(512, 2 * num_reqs), 8, 16, False + + +def _cp_partial( + q_heads: torch.Tensor, # [heads, rows, 128]: every rank's index heads, gathered + k_cache: Nvfp4Cache, + req_to_token: torch.Tensor, + seq_lens: torch.Tensor, # [rows] + max_seqlen: int, + slot_ids: torch.Tensor, # [rows] + num_draft: int, + block_size: int, + init_blocks: int, + local_blocks: int, + sm_scale: float, # q and k scales folded in, as verify_topk_idx does for NVFP4 + world: int, + rank: int, +): + """The fused verify partial pass over this rank's blocks, all heads in one tile: + (split lists [heads, rows, nc], splits per request, nc).""" + heads, rows, head_dim = q_heads.shape + num_reqs = rows // num_draft + max_blocks = triton.cdiv(min(max_seqlen, req_to_token.shape[1]), block_size) + programs, max_splits, min_blocks, paired = _cp_split_config(num_reqs) + nc = triton.next_power_of_2(max_splits * _KP) + part_k = torch.empty((heads, rows, nc), dtype=torch.int64, device=q_heads.device) + part_n = torch.empty(num_reqs, dtype=torch.int32, device=q_heads.device) + _verify_index_topk_partial_kernel[(programs, 1)]( + q_heads, + k_cache.payload.view(torch.int16), + k_cache.scales.view(torch.int32), + req_to_token, + seq_lens, + slot_ids, + part_k, + part_n, + num_reqs, + sm_scale, + 1.0, + max_blocks, + init_blocks, + local_blocks, + heads, + q_heads.stride(1), + q_heads.stride(0), + *(s // 2 for s in k_cache.payload_strides()), + *(s // 4 for s in k_cache.scale_strides()[:2]), + req_to_token.stride(0), + part_k.stride(0), + part_k.stride(1), + NUM_DRAFT=num_draft, + MAX_SPLITS=max_splits, + HEAD_DIM=head_dim, + BLOCK_N=block_size, + KP=_KP, + W=_WARPS, + MIN_BLOCKS=min_blocks, + PAGE_SIZE=k_cache.page_size, + K_FP4=True, + PAIRED=paired, + PROGRAMS=programs, + R_P2=triton.next_power_of_2(num_reqs), + ND_P2=triton.next_power_of_2(num_draft), + num_warps=_WARPS, + M=max(16, triton.next_power_of_2(num_draft * heads)), + WORLD=world, + RANK=rank, + ) + return part_k, part_n, nc + + +# ---- Indexer CP candidate exchange over IPC (no collective) ---- +# A rank writes its best KP candidates of head h into rank h's receive buffer +# [CP_PARITIES, world, CP_MAX_ROWS, KP] int64. Each word carries its call's tag, its own +# arrival flag: n % 255 + 1 for a row's n-th call, so a zeroed word reads "not arrived". +# The kernels count calls per row, so graph replay keeps ranks in step; counts wrap at 510, +# where tag and parity both repeat. Two parity slots suffice: a rank writes call n + 2 only +# after merging call n + 1, which needed every peer's call n + 1 writes. +CP_PARITIES = 2 +CP_MAX_ROWS = 4096 +# ~2.5 min of polls, under the scheduler watchdog's 300 s: a dead peer traps, not hangs +CP_MAX_SPINS = 1 << 30 + + +@triton.jit +def _wire(key, tag): + """A native key (ordered score << 32 | ~block) as an exchange word. + + The word is ordered score, tag, then 0xFFFFFF - block, so that within one call the + words order as the keys do. An empty key becomes the lowest score with block field 0. + """ + NO_KEY: tl.constexpr = -9223372036854775808 + blk = (~key).to(tl.int32) + score = (key >> 32).to(tl.int32) + empty = key == NO_KEY + score = tl.where(empty, -2147483648, score) + blk_field = tl.where(empty, 0, 0xFFFFFF - blk) + return ( + (score.to(tl.int64) << 32) | (tag.to(tl.int64) << 24) | blk_field.to(tl.int64) + ) + + +@triton.jit +def _cp_local_keys_ipc_kernel( + part_k_ptr, # [num_heads * rows, NC] int64 keys of a request's splits + part_n_ptr, # [num_reqs] splits of each request + peer_ptrs, # [world] int64 device addresses of every rank's receive buffer + counters_ptr, # [num_heads, CP_MAX_ROWS] int32, this rank's per-(head, row) calls + rows, + stride_p, + RANK: tl.constexpr, + WORLD: tl.constexpr, + NUM_DRAFT: tl.constexpr, + KP: tl.constexpr, + NC: tl.constexpr, + MAX_ROWS: tl.constexpr, +): + """This rank's best KP candidates of (head, row), forced blocks excluded, written into + the head owner's receive buffer: what the native merge would pick from these blocks.""" + pid = tl.program_id(0) + head, r = pid // rows, pid % rows + req = r // NUM_DRAFT + c = tl.arange(0, NC) + NO_KEY: tl.constexpr = -9223372036854775808 + key = tl.load( + part_k_ptr + pid.to(tl.int64) * stride_p + c, + mask=c < tl.load(part_n_ptr + req) * KP, + other=NO_KEY, + ) + key = tl.where((~key).to(tl.int32) >= 0, key, NO_KEY) + best = tl.topk(key, KP) + n = tl.load(counters_ptr + head * MAX_ROWS + r) + tl.store(counters_ptr + head * MAX_ROWS + r, (n + 1) % 510) + tag = n % 255 + 1 + parity = n % 2 + dst = tl.load(peer_ptrs + head).to(tl.pointer_type(tl.int64)) + off = ((parity * WORLD + RANK) * MAX_ROWS + r).to(tl.int64) * KP + tl.store(dst + off + tl.arange(0, KP), _wire(best, tag), cache_modifier=".wt") + + +@triton.jit +def _cp_merge_ipc_kernel( + recv_ptr, # this rank's receive buffer [CP_PARITIES, WORLD, MAX_ROWS, KP] int64 + counters_ptr, # [MAX_ROWS] int32, this rank's per-row merges + topk_ptr, # [rows, topk] int32 + seq_lens_ptr, + max_blocks, + topk, + init_blocks, + local_blocks, + WORLD: tl.constexpr, + KP: tl.constexpr, + BLOCK_N: tl.constexpr, + KT: tl.constexpr, + MAX_ROWS: tl.constexpr, + MAX_SPINS: tl.constexpr, +): + """Waits for every rank's KP words of this row to carry this call's tag (trapping + after MAX_SPINS polls), then selects as _verify_index_topk_merge_kernel: rank r scored + blocks r, r + WORLD, ..., so together they hold the row's best.""" + r = tl.program_id(0) + n = tl.load(counters_ptr + r) + tl.store(counters_ptr + r, (n + 1) % 510) + tag = (n % 255 + 1).to(tl.int64) + parity = n % 2 + c = tl.arange(0, WORLD * KP) + src, j = c // KP, c % KP + ptrs = recv_ptr + ((parity * WORLD + src) * MAX_ROWS + r).to(tl.int64) * KP + j + word = tl.load(ptrs, cache_modifier=".cv", volatile=True) + spins = 0 + while tl.min((((word >> 24) & 0xFF) == tag).to(tl.int32), axis=0) == 0: + word = tl.load(ptrs, cache_modifier=".cv", volatile=True) + spins += 1 + if spins >= MAX_SPINS: + tl.device_print( + "indexer CP exchange: a peer's candidates never came; row, tag", + r, + tag, + ) + tl.inline_asm_elementwise( + "s_trap 2", "=r", [], dtype=tl.int32, is_pure=False, pack=1 + ) + # back to the native key order: ordered score, then the lower block id + blk_field = word & 0xFFFFFF + blk = 0xFFFFFF - blk_field + NO_KEY: tl.constexpr = -9223372036854775808 + key = tl.where( + blk_field == 0, + NO_KEY, + ((word >> 32) << 32) | ((~blk.to(tl.int32)).to(tl.int64) & 0xFFFFFFFF), + ) + n_blocks = tl.minimum( + (tl.load(seq_lens_ptr + r).to(tl.int32) + BLOCK_N - 1) // BLOCK_N, max_blocks + ) + ids = _select_block_ids(key, n_blocks, topk, init_blocks, local_blocks, KT) + t = tl.arange(0, KT) + tl.store(topk_ptr + r.to(tl.int64) * topk + t, ids, mask=t < topk) + + +class CpExchange: + """One rank's receive buffer, peer addresses and call counters of the IPC exchange. + + Allocate once, before graph capture, on every rank of ``group`` (a collective); every + rank must then make the same calls. + """ + + def __init__(self, group, device): + import torch.distributed as dist + + self.world = group.world_size + self.rank = group.rank_in_group + nbytes = CP_PARITIES * self.world * CP_MAX_ROWS * _KP * 8 + addrs = _uncached_ipc_buffers(nbytes, group.cpu_group, self.rank, self.world) + self.recv = _RawPointer(addrs[self.rank], torch.int64) + self.peer_ptrs = torch.tensor(addrs, dtype=torch.int64, device=device) + self.write_calls = torch.zeros( + (self.world, CP_MAX_ROWS), dtype=torch.int32, device=device + ) + self.merge_calls = torch.zeros(CP_MAX_ROWS, dtype=torch.int32, device=device) + torch.cuda.synchronize(device) + dist.barrier(group=group.cpu_group) + + +def _hip_runtime(): + """The HIP runtime this process already loaded, through ctypes. + + Found in /proc/self/maps because ROCm wheels ship only the versioned soname. + """ + import ctypes + + with open("/proc/self/maps") as maps: + for line in maps: + if "libamdhip64.so" in line: + return ctypes.CDLL(line.split()[-1]) + raise RuntimeError("the HIP runtime is not loaded in this process") + + +def _uncached_ipc_buffers(nbytes, cpu_group, rank, world): + """This rank's addresses of every rank's zeroed ``nbytes`` of uncached device memory. + + Uncached, because a peer's writes land in this GPU's memory behind its L2: a cached + buffer would keep serving the waiting merge a stale line. + """ + import ctypes + + import torch.distributed as dist + + class _IpcHandle(ctypes.Structure): + _fields_ = [("reserved", ctypes.c_byte * 64)] + + hip = _hip_runtime() + + def check(err, what): + if err != 0: + raise RuntimeError(f"{what} failed with hipError {err}") + + ptr = ctypes.c_void_p() + hip_uncached = 3 # hipDeviceMallocUncached + check( + hip.hipExtMallocWithFlags( + ctypes.byref(ptr), ctypes.c_size_t(nbytes), ctypes.c_uint(hip_uncached) + ), + "hipExtMallocWithFlags", + ) + check(hip.hipMemset(ptr, 0, ctypes.c_size_t(nbytes)), "hipMemset") + check(hip.hipDeviceSynchronize(), "hipDeviceSynchronize") + handle = _IpcHandle() + check(hip.hipIpcGetMemHandle(ctypes.byref(handle), ptr), "hipIpcGetMemHandle") + handles = [None] * world + dist.all_gather_object(handles, bytes(handle.reserved), group=cpu_group) + hip.hipIpcOpenMemHandle.argtypes = [ + ctypes.POINTER(ctypes.c_void_p), + _IpcHandle, + ctypes.c_uint, + ] + addrs = [] + for r, raw in enumerate(handles): + if r == rank: + addrs.append(ptr.value) + continue + peer = _IpcHandle() + ctypes.memmove(ctypes.byref(peer), raw, 64) + opened = ctypes.c_void_p() + check( + hip.hipIpcOpenMemHandle(ctypes.byref(opened), peer, ctypes.c_uint(1)), + "hipIpcOpenMemHandle", + ) + addrs.append(opened.value) + return addrs + + +class _RawPointer: + """A device address Triton takes as a pointer argument.""" + + def __init__(self, addr: int, dtype: torch.dtype): + self.addr, self.dtype = addr, dtype + + def data_ptr(self) -> int: + return self.addr + + +def cp_verify_exchange( + q_heads, + k_cache, + req_to_token, + seq_lens, + max_seqlen, + slot_ids, + num_draft, + block_size, + topk, + init_blocks, + local_blocks, + sm_scale, + exchange: CpExchange, +): + """Indexer CP: score this rank's blocks, send each head's best _KP to its owner over + IPC and merge this rank's head: [1, rows, topk] int32, as verify_topk_idx.""" + heads, rows, head_dim = q_heads.shape + assert rows <= CP_MAX_ROWS and heads == exchange.world + num_reqs = rows // num_draft + max_blocks = triton.cdiv(min(max_seqlen, req_to_token.shape[1]), block_size) + part_k, part_n, nc = _cp_partial( + q_heads, + k_cache, + req_to_token, + seq_lens, + max_seqlen, + slot_ids, + num_draft, + block_size, + init_blocks, + local_blocks, + sm_scale, + exchange.world, + exchange.rank, + ) + _cp_local_keys_ipc_kernel[(heads * rows,)]( + part_k, + part_n, + exchange.peer_ptrs, + exchange.write_calls, + rows, + part_k.stride(1), + RANK=exchange.rank, + WORLD=exchange.world, + NUM_DRAFT=num_draft, + KP=_KP, + NC=nc, + MAX_ROWS=CP_MAX_ROWS, + num_warps=1, + ) + out = torch.empty((1, rows, topk), dtype=torch.int32, device=q_heads.device) + _cp_merge_ipc_kernel[(rows,)]( + exchange.recv, + exchange.merge_calls, + out, + seq_lens, + max_blocks, + topk, + init_blocks, + local_blocks, + WORLD=exchange.world, + KP=_KP, + BLOCK_N=block_size, + KT=triton.next_power_of_2(topk), + MAX_ROWS=CP_MAX_ROWS, + MAX_SPINS=CP_MAX_SPINS, + num_warps=1, + ) + return out diff --git a/python/sglang/srt/layers/attention/minimax_sparse_backend.py b/python/sglang/srt/layers/attention/minimax_sparse_backend.py index 892af9aa8d19..a086b2de2f43 100644 --- a/python/sglang/srt/layers/attention/minimax_sparse_backend.py +++ b/python/sglang/srt/layers/attention/minimax_sparse_backend.py @@ -1487,9 +1487,10 @@ def forward_extend( out_quant=out_quant, ) disable_value = layer.layer_id in self.disable_value_layer_ids - kv_cached_by_fusion = not save_kv_cache or self._is_sparse_kv_cached_by_fusion( - forward_batch, layer.layer_id - ) + # ROCm: save_kv_cache=False means the fused QK-norm/RoPE store already wrote K/V + kv_cached_by_fusion = ( + self.is_hip and not save_kv_cache + ) or self._is_sparse_kv_cached_by_fusion(forward_batch, layer.layer_id) if not kv_cached_by_fusion: self.kv_pool.set_fused_kv_index_buffer( layer, diff --git a/python/sglang/srt/layers/attention/minimax_sparse_ops/indexer_cp.py b/python/sglang/srt/layers/attention/minimax_sparse_ops/indexer_cp.py index 6c5029d43045..92b851f2ed4a 100644 --- a/python/sglang/srt/layers/attention/minimax_sparse_ops/indexer_cp.py +++ b/python/sglang/srt/layers/attention/minimax_sparse_ops/indexer_cp.py @@ -1,9 +1,10 @@ """Opt-in, decode-only context partitioning of MiniMax's replicated index K. This first implementation gathers the existing post-RoPE index queries rather -than changing checkpoint loading or fused projections. Both gathers go through +than changing checkpoint loading or fused projections. The query gather goes through SGLang's communicator so AITER registration and graph capture remain owned by -the runtime. Communication cost must be included in performance comparisons. +the runtime; the candidates return through a second gather, or on the NVFP4 index cache +through an IPC exchange. Communication cost must be included in performance comparisons. """ import logging @@ -12,6 +13,10 @@ logger = logging.getLogger(__name__) +# NVFP4 verify batches below this many requests keep the native indexer: there the Q +# gather and candidate exchange cost more than the index-K reads CP saves (TP4, graphs) +CP_MIN_VERIFY_REQS = 12 + def unsupported_reason( *, @@ -65,12 +70,20 @@ def unsupported_reason( def draft_is_chain_layout(algorithm, eagle_topk): """Whether verify rows reach the indexer as independent chain rows. - Allowlisted, not tree-denylisted: an unrecognized algorithm (DSPARK's ragged - verify, NGRAM's tree-in-mask) must disable CP rather than silently mis-score. + Allowlisted, not tree-denylisted: an unrecognized algorithm (NGRAM's + tree-in-mask) must disable CP rather than silently mis-score. DSPARK + qualifies only under the static (verify-all) schedule, where every request + verifies its whole draft block; ragged verify trims rows per request. """ + from sglang.srt.speculative.ragged_verify import ( + RaggedVerifyMode, + read_ragged_verify_mode, + ) from sglang.srt.speculative.spec_info import SpeculativeAlgorithm algo = SpeculativeAlgorithm.from_string(algorithm) + if algo.is_dspark(): + return read_ragged_verify_mode() is RaggedVerifyMode.STATIC return algo.is_none() or (algo.is_eagle() and (eagle_topk or 1) == 1) @@ -117,11 +130,19 @@ def make_indexer_cp(backend, runner, sparse_cfg): logger.warning("MiniMax indexer CP disabled: %s", reason) return None cp = MiniMaxIndexerCP(parallel.attn_tp_group) + # only the NVFP4 index cache reads the exchange; it is a collective on every rank + if backend.kv_is_nvfp4: + from sglang.kernels.ops.attention.minimax_sparse.decode.verify_topk_gfx950 import ( + CpExchange, + ) + + cp.exchange = CpExchange(cp.group, cp.group.device) cp.warmup() logger.info( "MiniMax indexer CP enabled: TP%d, decode-only, replicated index cache, " - "query gather + candidate gather, top-k 16; index-value layers use TP", + "query gather + %s, top-k 16; index-value layers use TP", cp.world, + "IPC candidate exchange" if cp.exchange is not None else "candidate gather", ) return cp @@ -133,6 +154,7 @@ def __init__(self, group): self.group = group self.world = group.world_size self.rank = group.rank_in_group + self.exchange = None def warmup(self): """Initialize both runtime and PyNCCL gather paths before graph capture.""" @@ -147,19 +169,47 @@ def warmup(self): comm.all_gather(output, probe) torch.cuda.synchronize(probe.device) - def supports(self, q, k_cache, max_seqlen, req_to_token): + def supports( + self, q, k_cache, max_seqlen, req_to_token, packed_queries=1, local_blocks=1 + ): # Shape/dtype/config predicates only. A rank-local exception must never # choose a different collective sequence from its peers. - return ( + from sglang.kernels.ops.attention.minimax_sparse.common.nvfp4 import ( + Nvfp4Cache, + ) + + common = ( q.ndim == 3 and q.shape[1:] == (1, 128) + and req_to_token.stride(1) == 1 + and 0 < max_seqlen <= min(req_to_token.shape[1], 16384 * 128) + ) + if isinstance(k_cache, Nvfp4Cache): + from sglang.kernels.ops.attention.minimax_sparse.decode.verify_topk_gfx950 import ( + CP_MAX_ROWS, + can_fuse_verify_topk, + ) + + return ( + common + and packed_queries > 1 + and q.shape[0] % packed_queries == 0 + and q.shape[0] // packed_queries >= CP_MIN_VERIFY_REQS + # the fused verify indexer scores all four heads' rows in one 32-row tile + and packed_queries * 4 <= 32 + # the exchange's receive buffers hold CP_MAX_ROWS rows per head + and q.shape[0] <= CP_MAX_ROWS + and can_fuse_verify_topk( + q, k_cache, packed_queries, 16, local_blocks, k_cache.page_size, 128 + ) + ) + return ( + common and q.dtype in (torch.bfloat16, torch.float16) and k_cache.ndim == 3 and k_cache.shape[1:] == (1, 128) and k_cache.dtype in (q.dtype, torch.float8_e4m3fn, torch.float8_e4m3fnuz, torch.float8_e5m2) - and req_to_token.stride(1) == 1 - and 0 < max_seqlen <= min(req_to_token.shape[1], 16384 * 128) ) def __call__( @@ -177,11 +227,18 @@ def __call__( k_scale=None, packed_queries=1, ): + from sglang.kernels.ops.attention.minimax_sparse.common.nvfp4 import ( + Nvfp4Cache, + ) + from sglang.kernels.ops.attention.minimax_sparse.common.utils import unit_scale from sglang.kernels.ops.attention.minimax_sparse.decode.indexer_cp import ( merge_candidates, score_local_blocks, select_local_candidates, ) + from sglang.kernels.ops.attention.minimax_sparse.decode.verify_topk_gfx950 import ( + cp_verify_exchange, + ) batch = q.shape[0] if batch == 0: @@ -190,10 +247,28 @@ def __call__( world = self.world gathered_q = torch.empty((world * batch, 128), dtype=q.dtype, device=q.device) self.group.all_gather_into_tensor(gathered_q, local_q) + q_heads = gathered_q.view(world, batch, 128) sm_scale = 128**-0.5 if sm_scale is None else sm_scale + if isinstance(k_cache, Nvfp4Cache): + # sm_scale folds q and k scales as verify_topk_idx does, so the ids match it + return cp_verify_exchange( + q_heads=q_heads, + k_cache=k_cache, + req_to_token=req_to_token, + seq_lens=seq_lens, + max_seqlen=max_seqlen, + slot_ids=slot_ids, + num_draft=packed_queries, + block_size=128, + topk=16, + init_blocks=init_blocks, + local_blocks=local_blocks, + sm_scale=sm_scale * unit_scale(q_scale) * unit_scale(k_scale), + exchange=self.exchange, + ) sm_scale *= 1.0 if q_scale is None else q_scale scores = score_local_blocks( - gathered_q.view(world, batch, 128), + q_heads, k_cache, req_to_token, slot_ids, diff --git a/python/sglang/srt/layers/attention/minimax_sparse_ops/minimax_sparse.py b/python/sglang/srt/layers/attention/minimax_sparse_ops/minimax_sparse.py index 81ad63903ffd..7673cde6bee9 100644 --- a/python/sglang/srt/layers/attention/minimax_sparse_ops/minimax_sparse.py +++ b/python/sglang/srt/layers/attention/minimax_sparse_ops/minimax_sparse.py @@ -384,7 +384,14 @@ def minimax_sparse_decode( and block_size_k == 128 and topk == 16 and dense_main_attn_fn is None - and indexer_cp.supports(idx_q, idx_k_cache, max_seqlen, req_to_token) + and indexer_cp.supports( + idx_q, + idx_k_cache, + max_seqlen, + req_to_token, + packed_queries=packed_queries, + local_blocks=local_blocks, + ) ): idx_o = real_seq_lens = None topk_idx = indexer_cp( diff --git a/python/sglang/srt/layers/quantization/quark/quark.py b/python/sglang/srt/layers/quantization/quark/quark.py index b9d2b8b85d68..ab25ed4ed20c 100644 --- a/python/sglang/srt/layers/quantization/quark/quark.py +++ b/python/sglang/srt/layers/quantization/quark/quark.py @@ -284,6 +284,8 @@ def _build_excluded_fp8_config(config: Dict[str, Any]) -> Optional["Fp8Config"]: logger = logging.getLogger(__name__) +_is_hip = is_hip() + _MOE_SHARED_EXPERT_QUANT_LAYER0_BASES: tuple[str, ...] = ( "model.layers.0", "model.language_model.layers.0", @@ -454,9 +456,12 @@ def _get_block_fp8_config( ): return None - # Quark's [1, 32] float8_e8m0fnu scales are uint8 e8m0, not fp32, so they load as MXFP8 - use_mxfp8 = block_size == [1, 32] and "e8m0" in str( - weight_config.get("scale_type") or "" + # Quark's [1, 32] float8_e8m0fnu scales are uint8 e8m0, not fp32, so they load as + # MXFP8; ROCm only, so CUDA keeps its earlier handling of these checkpoints + use_mxfp8 = ( + _is_hip + and block_size == [1, 32] + and "e8m0" in str(weight_config.get("scale_type") or "") ) return Fp8Config( is_checkpoint_fp8_serialized=True, @@ -589,7 +594,8 @@ def from_config(cls, config: dict[str, Any]) -> "QuarkConfig": from sglang.srt.layers.quantization.quark import mxfp8_base - if mxfp8_base.is_mxfp8_base_nvfp4_moe_source(config): + # ROCm only: CUDA keeps raising for this source below + if _is_hip and mxfp8_base.is_mxfp8_base_nvfp4_moe_source(config): return mxfp8_base.QuarkMxfp8BaseConfig.from_source(config, hf_config) nvfp4_src = _detect_nvfp4_source(config) diff --git a/python/sglang/srt/models/minimax_m3.py b/python/sglang/srt/models/minimax_m3.py index 7e232f1e0e49..1ba5552e3748 100644 --- a/python/sglang/srt/models/minimax_m3.py +++ b/python/sglang/srt/models/minimax_m3.py @@ -2185,9 +2185,12 @@ def load_weights(self, weights: Iterable[Tuple[str, torch.Tensor]]): params_dict = dict(self.named_parameters()) loaded_params: Set[str] = set() - # minimax_m3_nvfp4 reads weight_global_scale itself; remapping it inverts the scale twice + # minimax_m3_nvfp4 reads weight_global_scale itself; remapping it inverts the scale twice. + # The Quark remap and scale alias below are ROCm-only, like Quark's MXFP8 loading. remap_llm_compressor_experts = ( - self.quant_config is not None and self.quant_config.get_name() == "quark" + _is_hip + and self.quant_config is not None + and self.quant_config.get_name() == "quark" ) for name, loaded_weight in weights: layer_id = get_layer_id(name) @@ -2229,7 +2232,8 @@ def load_weights(self, weights: Iterable[Tuple[str, torch.Tensor]]): name = name.replace(weight_name, param_name) if name.endswith(".bias") and name not in params_dict: continue - name = mxfp8_linear_scale_alias(name, params_dict) + if _is_hip: + name = mxfp8_linear_scale_alias(name, params_dict) if name not in params_dict: continue @@ -2272,7 +2276,8 @@ def load_weights(self, weights: Iterable[Tuple[str, torch.Tensor]]): if name is None: continue - name = mxfp8_linear_scale_alias(name, params_dict) + if _is_hip: + name = mxfp8_linear_scale_alias(name, params_dict) if name in params_dict: param = params_dict[name] weight_loader = getattr( diff --git a/python/sglang/srt/models/minimax_m3_vl.py b/python/sglang/srt/models/minimax_m3_vl.py index 0537df70c57a..48eb76324643 100644 --- a/python/sglang/srt/models/minimax_m3_vl.py +++ b/python/sglang/srt/models/minimax_m3_vl.py @@ -432,9 +432,11 @@ def _load_llm_weight( ): return - # minimax_m3_nvfp4 reads weight_global_scale itself; remapping it inverts the scale twice + # minimax_m3_nvfp4 reads weight_global_scale itself; remapping it inverts the scale twice. + # The Quark remap and scale alias are ROCm-only, like Quark's MXFP8 loading. if ( - self.quant_config is not None + _is_hip + and self.quant_config is not None and self.quant_config.get_name() == "quark" and "mlp.experts." in name ): @@ -448,7 +450,8 @@ def _load_llm_weight( new_name = name.replace(weight_name, param_name) if new_name.endswith(".bias") and new_name not in params_dict: continue - new_name = mxfp8_linear_scale_alias(new_name, params_dict) + if _is_hip: + new_name = mxfp8_linear_scale_alias(new_name, params_dict) if new_name not in params_dict: continue param = params_dict[new_name] @@ -481,7 +484,8 @@ def _load_llm_weight( remapped = maybe_remap_kv_scale_name(name, params_dict) if remapped is None: return - remapped = mxfp8_linear_scale_alias(remapped, params_dict) + if _is_hip: + remapped = mxfp8_linear_scale_alias(remapped, params_dict) if remapped not in params_dict: logger.warning(f"Parameter {remapped} not found in params_dict") return diff --git a/test/registered/amd/test_minimax_indexer_cp_ipc.py b/test/registered/amd/test_minimax_indexer_cp_ipc.py new file mode 100644 index 000000000000..9f8e86909d5d --- /dev/null +++ b/test/registered/amd/test_minimax_indexer_cp_ipc.py @@ -0,0 +1,201 @@ +"""The IPC candidate exchange must deliver each head's native NVFP4 verify selection +on four ranks, for DSpark's 8-row and EAGLE's 4-row chains, tied scores and rows shorter +than top-k, across call counters and their wrap, row-count changes and graph replay.""" + +import multiprocessing +import types +import unittest + +import torch +import torch.distributed as dist + +from sglang.srt.utils import is_gfx95_supported, is_hip +from sglang.srt.utils.network import get_open_port +from sglang.test.ci.ci_register import register_amd_ci +from sglang.test.test_utils import CustomTestCase + +# backend-specific: the exchange is gfx950 HIP IPC; the suite's MI35x box supplies 4 GPUs +register_amd_ci(est_time=60, suite="stage-c-test-large-8-gpu-amd-mi35x") + +WORLD = 4 +BLOCK, TOPK, HEADS, DIM = 128, 16, 4, 128 +SM_SCALE, K_SCALE = 0.07, 0.5 +CALLS = 5 +# (requests, draft rows, tied pages, init blocks, local blocks): init != local so a +# swapped argument shows; the second case also has rows shorter than top-k +CASES = ((64, 8, False, 2, 1), (17, 4, True, 0, 1)) + + +def _nvfp4_verify_inputs(prefixes, num_draft, seed, tied_pages=False): + """DSpark verify rows over a shuffled NVFP4 index-K pool of whole 128-token pages. + + With ``tied_pages`` most pages copy page 1, so most blocks tie. + """ + from sglang.kernels.ops.attention.minimax_sparse.common.nvfp4 import Nvfp4Cache + from sglang.test.quant_ref_utils import nvfp4_kv_quantize + + device = torch.device("cuda") + gen = torch.Generator(device=device).manual_seed(seed) + max_len = max(prefixes) + num_draft + pages_per_req = -(-max_len // BLOCK) + num_pages = len(prefixes) * pages_per_req + pages = torch.randperm(num_pages, device=device, generator=gen) + 1 + offs = torch.arange(BLOCK, device=device) + table = ( + (pages.view(len(prefixes), pages_per_req, 1) * BLOCK + offs) + .view(len(prefixes), -1)[:, :max_len] + .to(torch.int32) + ) + k = torch.randn((num_pages + 1) * BLOCK, DIM, device=device, generator=gen) + if tied_pages: + k_pages = k.view(num_pages + 1, BLOCK * DIM) + rare = torch.arange(num_pages + 1, device=device) % 13 == 0 + k_pages.copy_(k_pages[1 + rare.long()]) + payload, scales = nvfp4_kv_quantize( + k.to(torch.bfloat16), torch.ones(1, device=device) + ) + cache = Nvfp4Cache( + payload=payload.view(num_pages + 1, BLOCK, DIM // 2), + scales=scales.view(num_pages + 1, BLOCK, DIM // 16), + ) + rows = len(prefixes) * num_draft + q = torch.randn(rows, HEADS, DIM, device=device, generator=gen).to(torch.bfloat16) + lengths = torch.tensor( + [p + j + 1 for p in prefixes for j in range(num_draft)], + dtype=torch.int64, + device=device, + ) + slots = torch.arange(len(prefixes), device=device).repeat_interleave(num_draft) + # graph capture sizes the score columns past the real lengths + return q, cache, table, slots, lengths, max_len + 4 * BLOCK + + +def _run_exchange(rank: int, port: int) -> None: + from sglang.kernels.ops.attention.minimax_sparse.decode import verify_topk_idx + from sglang.kernels.ops.attention.minimax_sparse.decode.verify_topk_gfx950 import ( + CpExchange, + cp_verify_exchange, + ) + + device = torch.device(f"cuda:{rank}") + torch.cuda.set_device(device) + dist.init_process_group( + backend="gloo", + init_method=f"tcp://127.0.0.1:{port}", + rank=rank, + world_size=WORLD, + ) + try: + group = types.SimpleNamespace( + world_size=WORLD, rank_in_group=rank, cpu_group=dist.group.WORLD + ) + exchange = CpExchange(group, device) + # every rank builds the same pool; a call's row count changes between cases + for seed, (num_reqs, num_draft, tied, init, local) in enumerate(CASES): + prefixes = [20000 - 97 * r - 1000 * seed for r in range(num_reqs)] + if num_reqs > 5: + prefixes[3], prefixes[5] = 3 * BLOCK - 2, 777 # <= top-k blocks + q, cache, table, slots, lengths, _ = _nvfp4_verify_inputs( + prefixes, num_draft, seed=seed, tied_pages=tied + ) + max_len = table.shape[1] + # distinct queries per call, so a word left from an earlier call never + # passes for this one's + gen = torch.Generator(device=device).manual_seed(100 + seed) + queries = [q] + [ + torch.randn(q.shape, device=device, generator=gen).to(q.dtype) + for _ in range(CALLS - 1) + ] + natives = [ + verify_topk_idx.verify_topk_idx( + q=query[:, rank : rank + 1].contiguous(), + k_cache=cache, + req_to_token=table, + seq_lens=lengths, + max_seqlen=max_len, + slot_ids=slots, + num_draft=num_draft, + block_size=BLOCK, + topk=TOPK, + init_blocks=init, + local_blocks=local, + sm_scale=SM_SCALE, + q_scale=None, + k_scale=K_SCALE, + page_size=BLOCK, + ) + for query in queries + ] + q_heads = torch.empty_like(q.permute(1, 0, 2).contiguous()) + + def exchange_call(): + return cp_verify_exchange( + q_heads, cache, table, lengths, max_len, slots, num_draft, + BLOCK, TOPK, init, local, SM_SCALE * K_SCALE, exchange, + ) # fmt: skip + + # seed 1 starts two calls short of the counters' wrap at 510 + if seed == 1: + dist.barrier() + exchange.write_calls.fill_(508) + exchange.merge_calls.fill_(508) + torch.cuda.synchronize() + dist.barrier() + # more calls than parities, so a stale parity or tag would be read back + for call, (query, native) in enumerate(zip(queries, natives)): + q_heads.copy_(query.permute(1, 0, 2)) + out = exchange_call() + torch.cuda.synchronize() + assert torch.equal(out, native), ( + f"rank {rank} reqs {num_reqs} call {call}: " + f"{(out != native).sum()} IDs differ" + ) + if seed == 0: + graph = torch.cuda.CUDAGraph() + stream = torch.cuda.Stream() + with torch.cuda.stream(stream): + exchange_call() + torch.cuda.synchronize() + with torch.cuda.graph(graph, stream=stream): + replayed = exchange_call() + for replay, (query, native) in enumerate(zip(queries, natives)): + q_heads.copy_(query.permute(1, 0, 2)) + graph.replay() + torch.cuda.synchronize() + assert torch.equal(replayed, native), ( + f"rank {rank} replay {replay}: " + f"{(replayed != native).sum()} IDs differ" + ) + dist.barrier() + finally: + dist.destroy_process_group() + + +def _spawn(target, port): + context = multiprocessing.get_context("spawn") + procs = [context.Process(target=target, args=(rank, port)) for rank in range(WORLD)] + for proc in procs: + proc.start() + return procs + + +@unittest.skipUnless( + is_hip() and is_gfx95_supported(), "MiniMax indexer CP is gfx950-only" +) +class TestMiniMaxIndexerCPExchange(CustomTestCase): + def test_exchange_matches_the_native_selector_on_four_ranks(self): + if torch.cuda.device_count() < WORLD: + self.skipTest(f"needs {WORLD} GPUs") + procs = _spawn(_run_exchange, get_open_port()) + # a buffer that serves a stale line spins the merge forever instead of failing + for proc in procs: + proc.join(timeout=600) + hung = [proc for proc in procs if proc.is_alive()] + for proc in hung: + proc.kill() + self.assertFalse(hung, "an exchange rank hung") + self.assertEqual([proc.exitcode for proc in procs], [0] * WORLD) + + +if __name__ == "__main__": + unittest.main() diff --git a/test/registered/unit/layers/attention/test_minimax_indexer_cp_gate.py b/test/registered/unit/layers/attention/test_minimax_indexer_cp_gate.py index d797c6868c28..cd032755a3a1 100644 --- a/test/registered/unit/layers/attention/test_minimax_indexer_cp_gate.py +++ b/test/registered/unit/layers/attention/test_minimax_indexer_cp_gate.py @@ -1,9 +1,11 @@ import unittest +from sglang.srt.environ import envs from sglang.srt.layers.attention.minimax_sparse_ops.indexer_cp import ( draft_is_chain_layout, unsupported_reason, ) +from sglang.srt.speculative.ragged_verify import RaggedVerifyMode from sglang.test.ci.ci_register import register_cpu_ci from sglang.test.test_utils import CustomTestCase @@ -42,7 +44,6 @@ def test_only_validated_draft_layouts_enable_cp(self): ("EAGLE", 1): True, ("EAGLE3", 1): True, ("EAGLE3", 4): False, # tree draft - ("DSPARK", None): False, # ragged verify lengths ("NGRAM", None): False, # tree lives in the verify mask ("DFLASH", None): False, ("STANDALONE", 1): False, @@ -56,6 +57,14 @@ def test_only_validated_draft_layouts_enable_cp(self): } ) self.assertEqual(reason is None, expected, reason) + # ragged verify trims rows per request, which CP would score as whole chains + for mode in RaggedVerifyMode: + with ( + self.subTest(algorithm="DSPARK", mode=mode), + envs.SGLANG_RAGGED_VERIFY_MODE.override(mode.value), + ): + chain = draft_is_chain_layout("DSPARK", None) + self.assertEqual(chain, mode is RaggedVerifyMode.STATIC) def test_sm90_enables_only_validated_topologies(self): """Each rank must serve exactly one index head of the four; any other split, diff --git a/test/registered/unit/models/test_minimax_m31_model.py b/test/registered/unit/models/test_minimax_m31_model.py index 6e3f85a8493f..085c06816b88 100644 --- a/test/registered/unit/models/test_minimax_m31_model.py +++ b/test/registered/unit/models/test_minimax_m31_model.py @@ -23,7 +23,7 @@ from sglang.srt.environ import envs from sglang.srt.layers.layer_boundary.fusions.moe_finalize import MoeDeferredFinalize from sglang.srt.layers.layer_boundary.residual import batch as residual_batch -from sglang.srt.models import minimax_m3 +from sglang.srt.models import minimax_m3, minimax_m3_vl from sglang.srt.models.dspark_minimax import DSparkMiniMaxDraftModel from sglang.srt.models.minimax_m3 import ( MiniMaxM3Model, @@ -296,6 +296,17 @@ def test_rocm_caps_nvfp4_tp_and_quark_attention_tp(self): _minimax_m3_overrides(rocm_server_args(**quark), MXFP8) _minimax_m3_overrides(rocm_server_args(attn_dp_size=2, **quark), MXFP8) + def test_gfx950_keeps_custom_all_reduce_only_for_the_m31_quark_export(self): + """Custom all-reduce is measured faster for the M3.1 Quark export; the M3 export, + whose linears are bf16, stays on NCCL.""" + with gfx950(): + m31 = _minimax_m3_overrides( + rocm_server_args(quantization="quark"), QUARK_M31 + ) + m3 = _minimax_m3_overrides(rocm_server_args(quantization="quark"), QUARK_M3) + self.assertNotIn("disable_custom_all_reduce", m31) + self.assertTrue(m3["disable_custom_all_reduce"]) + def test_gfx950_quick_all_reduce_takes_int4_only_where_unset(self): """M3.1 gets INT4 when no level is set; a level already set stays, INT8 included, since the ROCm images' INT8 cannot be told from one the user chose.""" @@ -422,8 +433,14 @@ def scale_param(key): k = torch.full((2,), 127, dtype=torch.uint8) o = torch.full((2,), 126, dtype=torch.uint8) prefix = "model.layers.0.self_attn." - model._load_llm_weight(prefix + "k_proj.weight_scale", k, params, mapping, []) - model._load_llm_weight(prefix + "o_proj.weight_scale", o, params, mapping, []) + # the alias is ROCm-only; this CPU test runs the HIP loader path + with patch.object(minimax_m3_vl, "_is_hip", True): + model._load_llm_weight( + prefix + "k_proj.weight_scale", k, params, mapping, [] + ) + model._load_llm_weight( + prefix + "o_proj.weight_scale", o, params, mapping, [] + ) self.assertEqual(loaded["qkv"][0], "k") self.assertTrue(torch.equal(loaded["qkv"][1], k)) self.assertIsNone(loaded["o"][0])