diff --git a/python/sglang/kernels/jit/csrc/minimax/minimax_decode_topk.cuh b/python/sglang/kernels/jit/csrc/minimax/minimax_decode_topk.cuh index 38b4368d52d1..b2e95de76117 100644 --- a/python/sglang/kernels/jit/csrc/minimax/minimax_decode_topk.cuh +++ b/python/sglang/kernels/jit/csrc/minimax/minimax_decode_topk.cuh @@ -62,7 +62,6 @@ struct TopKTraitImpl { struct Smem { uint32_t warp_sum[kNumWarps]; alignas(128) uint32_t counter; - alignas(128) uint32_t counter_final; alignas(128) uint32_t threshold_bin; uint32_t equal_count; uint32_t above_count; @@ -372,6 +371,20 @@ struct TopKTraitImpl { } __syncthreads(); }; + // Number of threads with a lower tx whose `flag` is set; called by the whole + // CTA. Final-round ties are exactly equal scores; ranking them by block id + // makes the selection deterministic (lower id wins, as in the small path). + const auto block_rank = [&](bool flag) -> uint32_t { + using namespace device; + const uint32_t ballot = __ballot_sync(kWarpSyncMask, flag); + if (lane_id == 0) smem->warp_sum[warp_id] = __popc(ballot); + __syncthreads(); + uint32_t base = 0; + for (uint32_t w = 0; w < warp_id; ++w) + base += smem->warp_sum[w]; + __syncthreads(); // warp_sum is reused by the next call + return base + __popc(ballot & ((1u << lane_id) - 1u)); + }; if (num_blocks <= kSmallThreshold) { // O(n^2) compare: each block's rank = #blocks that outrank it; the ones @@ -417,7 +430,7 @@ struct TopKTraitImpl { uint32_t topk_remain = topk; uint32_t write_pos = topk; // sentinel: not selected if (tx < kRadixSize) smem->histogram[0][tx] = 0; - if (tx == kRadixSize) smem->counter = smem->counter_final = 0; + if (tx == kRadixSize) smem->counter = 0; __syncthreads(); uint32_t total_active = num_blocks; @@ -441,17 +454,24 @@ struct TopKTraitImpl { if (round < 3) total_active = equal_count; topk_remain -= above_count; - // scatter: above -> selected now; equal at the last pass -> keep the rest + // scatter: above -> selected now; equal at the last pass -> lowest ids first if (active) { if (bin > threshold_bin) { write_pos = atomicAdd(&smem->counter, 1); active = false; } else if (bin < threshold_bin) { active = false; - } else if (round == 3) { - write_pos = topk - topk_remain + atomicAdd(&smem->counter_final, 1); } - // bin == threshold && round < 3: stay active for the next pass + // bin == threshold: stays active (for the next pass, or as a tie below) + } + if (round == 3) { + if (equal_count <= topk_remain) { + // Every tie is selected, so slot order is free (the output is sorted). + if (active) write_pos = atomicAdd(&smem->counter, 1); + } else { + const auto rank = block_rank(active); + if (active && rank < topk_remain) write_pos = topk - topk_remain + rank; + } } if (round == 3 || topk_remain == 0) break; @@ -480,7 +500,7 @@ struct TopKTraitImpl { } } if (tx < kRadixSize) smem->histogram[0][tx] = 0; - if (tx == kRadixSize) smem->counter = smem->counter_final = 0; + if (tx == kRadixSize) smem->counter = 0; __syncthreads(); uint32_t topk_remain = topk; @@ -514,11 +534,26 @@ struct TopKTraitImpl { active &= ~(1u << i); } else if (bin < threshold_bin) { active &= ~(1u << i); - } else if (round == 3) { - const auto pos = topk - topk_remain + atomicAdd(&smem->counter_final, 1); - if (pos < topk) topk_out[pos] = i * kCTASize + tx; } - // bin == threshold && round < 3: slot stays live for the next pass + // bin == threshold: slot stays live (for the next pass, or as a tie below) + } + } + if (round == 3 && equal_count <= topk_remain) { + // Every tie is selected, so slot order is free (the output is sorted). +#pragma unroll + for (uint32_t i = 0; i < kIters; ++i) + if (active & (1u << i)) topk_out[atomicAdd(&smem->counter, 1)] = i * kCTASize + tx; + } else if (round == 3) { + // Ties in block-id order: slot i holds ids i * kCTASize + tx, so slots + // are ranked in turn and each slot's ties by tx. + uint32_t taken = 0; + for (uint32_t i = 0; i < kIters && taken < topk_remain; ++i) { + const bool tie = (active >> i) & 1u; + const uint32_t num_ties = __syncthreads_count(tie); + if (num_ties == 0) continue; + const auto rank = taken + block_rank(tie); + if (tie && rank < topk_remain) topk_out[topk - topk_remain + rank] = i * kCTASize + tx; + taken += num_ties; } } diff --git a/test/registered/kernels/ops/attention/test_minimax_decode_topk.py b/test/registered/kernels/ops/attention/test_minimax_decode_topk.py index 5f464180bd3a..402224b2f6a4 100644 --- a/test/registered/kernels/ops/attention/test_minimax_decode_topk.py +++ b/test/registered/kernels/ops/attention/test_minimax_decode_topk.py @@ -118,6 +118,29 @@ def test_decode_topk_adversarial(kind): torch.testing.assert_close(a, b, rtol=0, atol=0) +@pytest.mark.parametrize("nb", [300, 512, 2000, 3000, 9000]) +@pytest.mark.parametrize("levels", [2, 4, 1]) +def test_decode_topk_ties_pick_lowest_ids(nb, levels): + """Exactly equal scores are broken by block id (lower id wins), the same rule as + the small-row path, so repeated calls select the same blocks: one-element-per- + thread rows (nb <= 512) and register-bucket rows (nb > 512).""" + torch.manual_seed(nb + levels) + block_size = 128 + H, B, topk = 2, 3, 16 + score = torch.randint(0, levels, (H, B, nb), device="cuda").float() + seq_lens = torch.full((B,), nb * block_size, device="cuda", dtype=torch.int32) + # reference: order by (score desc, id asc), keep topk, emit ascending + want = torch.sort( + torch.argsort(-score, dim=-1, stable=True)[..., :topk], dim=-1 + ).values.to(torch.int32) + first = minimax_decode_topk(score, seq_lens, block_size, topk) + assert torch.equal(first, want), (first, want) + for _ in range(20): + assert torch.equal( + minimax_decode_topk(score, seq_lens, block_size, topk), first + ) + + @pytest.mark.parametrize("seq_len", [1, 128, 129, 2048, 2049]) def test_decode_topk_small_num_blocks(seq_len): # num_blocks around / below topk -> naive identity path and boundary.