diff --git a/python/sglang/kernels/jit/csrc/deepseek_v4/decoding_sched_meta.cuh b/python/sglang/kernels/jit/csrc/deepseek_v4/decoding_sched_meta.cuh new file mode 100644 index 000000000000..bcf0eda3d179 --- /dev/null +++ b/python/sglang/kernels/jit/csrc/deepseek_v4/decoding_sched_meta.cuh @@ -0,0 +1,296 @@ +#include +#include + +#include + +#include + +#include + +// Tile-scheduler metadata for FlashMLA's split-KV decode. +// +// FlashMLA computes this itself when it is handed no metadata, in a <<<1, 32>>> +// kernel whose whole partition loop runs on thread 0 and stores each 32-byte +// entry straight to global memory. The loop is over `num_sm_parts`, which is +// `num_sms / s_q`, so on GB300 a BS=1 draft step walks 152 iterations at global +// store latency: 25.1 us measured, against 5.7 us for the 25 parts of a 6-row +// verify. Inside a cuda graph it is fully exposed on the critical path, and it +// cannot be hoisted out because the schedule depends on the per-request +// `topk_length` of the step being replayed. +// +// The walk itself is sequential -- each part starts where the last one stopped -- +// and moving its stores off global memory barely helps, because what it actually +// costs is a chain of dependent shared loads run by one thread with no ILP to +// hide them. Two things fix that without changing a single output: +// +// * the request being consumed changes only `batch_size` times over the whole +// walk, so its three values live in registers and are reloaded on that event +// rather than re-read from shared memory once per partition; +// * the walk stops as soon as the last request is consumed. Every partition +// after that one describes an empty range, all of them identical, and the +// whole block fills them in parallel. At BS=1 that is most of them: 152 +// partitions over 5 or 6 query rows leaves fewer than 50 doing any work. +// +// Feed the result to FlashMLA as the cached `tile_scheduler_metadata` / +// `num_splits` and it skips its own kernel. +// +// The shared-memory layout matches FlashMLA's exactly. It has to: when the +// walk finishes early it reads `first_block_idx_shared[batch_size]`, one past +// the end, which lands on `last_block_idx_shared[0]` -- those entries belong to +// parts with no work, but the bytes still have to agree. + +namespace sglang { + +namespace dsv4_sched_meta { + +// sizeof(DecodingSchedMeta)/4: begin/end req, begin/end block, begin split, +// the two per-end split flags, one pad word. +constexpr int kMetaInts = 8; +constexpr int kBlockSize = 256; + +struct Params { + int b; + int block_size_n; + int fixed_overhead_num_blocks; + int topk; // -1 for a dense model + int extra_topk; // 0 when there is no extra cache + const int* __restrict__ topk_length; + const int* __restrict__ extra_topk_length; + const int* __restrict__ seqlens_k; // dense only + int* __restrict__ tile_scheduler_metadata; + int* __restrict__ num_splits; + int num_sm_parts; +}; + +__device__ __forceinline__ int ceil_div_i(int a, int b) { + return (a + b - 1) / b; +} + +// ku::ceil: round up to a multiple of b. +__device__ __forceinline__ int ceil_to_i(int a, int b) { + return (a + b - 1) / b * b; +} + +__global__ void __launch_bounds__(kBlockSize) decoding_sched_meta_kernel(__grid_constant__ const Params p) { + extern __shared__ int smem[]; + const int b = p.b; + int* num_blocks_shared = smem; // [b] + int* num_splits_shared = smem + b; // [b + 1] + int* seqlens_k_shared = smem + b * 2 + 1; // [b] + int* first_block_idx_shared = smem + b * 3 + 1; // [b] + int* last_block_idx_shared = smem + b * 4 + 1; // [b] + int* out_shared = smem + b * 5 + 1; // [num_sm_parts * kMetaInts] + + __shared__ int total_num_blocks_shared; + + int partial = 0; + for (int i = threadIdx.x; i < b; i += kBlockSize) { + int cur_s_k; + if (p.topk == -1) { + cur_s_k = __ldg(p.seqlens_k + i); + } else { + cur_s_k = p.topk_length ? __ldg(p.topk_length + i) : p.topk; + if (cur_s_k == 0) cur_s_k = 1; // the main loop must never be empty + if (p.extra_topk) { + cur_s_k = ceil_to_i(cur_s_k, p.block_size_n); + cur_s_k += p.extra_topk_length ? __ldg(p.extra_topk_length + i) : p.extra_topk; + } + } + seqlens_k_shared[i] = cur_s_k; + const int last_token_idx = max(cur_s_k - 1, 0); + const int cur_first_block_idx = 0; // first_token_idx is always 0 + const int cur_last_block_idx = last_token_idx / p.block_size_n; + const int num_blocks = cur_last_block_idx - cur_first_block_idx + 1; + partial += num_blocks + p.fixed_overhead_num_blocks; + num_blocks_shared[i] = num_blocks; + first_block_idx_shared[i] = cur_first_block_idx; + last_block_idx_shared[i] = cur_last_block_idx; + } + + // Integer sum, so the tree order does not change the result. + for (int offset = 16; offset >= 1; offset /= 2) { + partial += __shfl_xor_sync(uint32_t(-1), partial, offset); + } + __shared__ int warp_sums[kBlockSize / 32]; + if ((threadIdx.x & 31) == 0) warp_sums[threadIdx.x >> 5] = partial; + __syncthreads(); + if (threadIdx.x == 0) { + int total = 0; +#pragma unroll + for (int w = 0; w < kBlockSize / 32; ++w) + total += warp_sums[w]; + total_num_blocks_shared = total; + } + __syncthreads(); + + const int fixed_overhead_num_blocks = p.fixed_overhead_num_blocks; + __shared__ int first_idle_part_shared; + + if (threadIdx.x == 0) { + const int payload = ceil_div_i(total_num_blocks_shared, p.num_sm_parts) + fixed_overhead_num_blocks; + + int now_req_idx = 0, now_block = 0, now_n_split_idx = 0, cum_num_splits = 0; + // The request being consumed, and the one before it, held across partitions. + int cur_c = num_blocks_shared[0], cur_lb = last_block_idx_shared[0], cur_sk = seqlens_k_shared[0]; + int prev_lb = 0, prev_sk = 0; + num_splits_shared[0] = 0; + int i = 0; + for (; i < p.num_sm_parts; ++i) { + int* meta = out_shared + i * kMetaInts; + const int begin_req_idx = now_req_idx; + // first_block_idx is 0 for every request: the first token index is 0. + const int begin_block_idx = now_block; + const int begin_split_idx = now_n_split_idx; + int is_first_req_splitted = (now_block != 0); + int remain_payload = payload; + while (now_req_idx < b) { + const int now_remain_blocks = cur_c - now_block; + if (remain_payload >= now_remain_blocks + fixed_overhead_num_blocks) { + cum_num_splits += now_n_split_idx + 1; + num_splits_shared[now_req_idx + 1] = cum_num_splits; + remain_payload -= now_remain_blocks + fixed_overhead_num_blocks; + ++now_req_idx; + now_block = 0; + now_n_split_idx = 0; + prev_lb = cur_lb; + prev_sk = cur_sk; + if (now_req_idx < b) { + cur_c = num_blocks_shared[now_req_idx]; + cur_lb = last_block_idx_shared[now_req_idx]; + cur_sk = seqlens_k_shared[now_req_idx]; + } + } else { + if (remain_payload - fixed_overhead_num_blocks > 0) { + now_block += remain_payload - fixed_overhead_num_blocks; + ++now_n_split_idx; + remain_payload = 0; + } + break; + } + } + const int split_open = now_block > 0; + const int end_req_idx = split_open ? now_req_idx : now_req_idx - 1; + const int end_lb = split_open ? cur_lb : prev_lb; + const int end_sk = split_open ? cur_sk : prev_sk; + const int end_block_idx = split_open ? now_block : (end_sk == 0 ? 0 : end_lb + 1); + int is_last_req_splitted = (end_block_idx != end_lb + 1) && (end_sk != 0); + if (begin_req_idx == end_req_idx) { + is_first_req_splitted = is_last_req_splitted = is_first_req_splitted || is_last_req_splitted; + } + meta[0] = begin_req_idx; + meta[1] = end_req_idx; + meta[2] = begin_block_idx; + meta[3] = end_block_idx; + meta[4] = begin_split_idx; + meta[5] = is_first_req_splitted; + meta[6] = is_last_req_splitted; + meta[7] = 0; + if (now_req_idx == b) { + ++i; + break; + } + } + first_idle_part_shared = i; + } + __syncthreads(); + + // Every partition past the walk describes the same empty range. + { + const int lb_last = last_block_idx_shared[b - 1]; + const int sk_last = seqlens_k_shared[b - 1]; + const int end_block_idx = (sk_last == 0) ? 0 : lb_last + 1; + // FlashMLA reads first_block_idx_shared[batch_size] for these, one past the + // end of that array, which aliases last_block_idx_shared[0]. + const int begin_block_idx = last_block_idx_shared[0]; + const int is_last_req_splitted = (end_block_idx != lb_last + 1) && (sk_last != 0); + for (int i = first_idle_part_shared + threadIdx.x; i < p.num_sm_parts; i += kBlockSize) { + int* meta = out_shared + i * kMetaInts; + meta[0] = b; + meta[1] = b - 1; + meta[2] = begin_block_idx; + meta[3] = end_block_idx; + meta[4] = 0; + meta[5] = 0; + meta[6] = is_last_req_splitted; + meta[7] = 0; + } + } + __syncthreads(); + + const int meta_words = p.num_sm_parts * kMetaInts; + for (int i = threadIdx.x; i < meta_words; i += kBlockSize) { + p.tile_scheduler_metadata[i] = out_shared[i]; + } + for (int i = threadIdx.x; i <= b; i += kBlockSize) { + p.num_splits[i] = num_splits_shared[i]; + } +} + +} // namespace dsv4_sched_meta + +void decoding_sched_meta( + tvm::ffi::TensorView tile_scheduler_metadata, + tvm::ffi::TensorView num_splits, + tvm::ffi::Optional topk_length, + tvm::ffi::Optional extra_topk_length, + tvm::ffi::Optional seqlens_k, + int64_t block_size_n, + int64_t fixed_overhead_num_blocks, + int64_t topk, + int64_t extra_topk) { + using namespace host; + using namespace dsv4_sched_meta; + + auto parts = SymbolicSize{"num_sm_parts"}; + auto meta_ints = SymbolicSize{"meta_ints"}; + auto b_plus_one = SymbolicSize{"batch_size_plus_one"}; + auto device_ = SymbolicDevice{}; + device_.set_options(); + + TensorMatcher({parts, meta_ints}).with_dtype().with_device(device_).verify(tile_scheduler_metadata); + TensorMatcher({b_plus_one}).with_strides({1}).with_dtype().with_device(device_).verify(num_splits); + + const int num_sm_parts = static_cast(parts.unwrap()); + const int b = static_cast(b_plus_one.unwrap()) - 1; + RuntimeCheck( + static_cast(meta_ints.unwrap()) == kMetaInts, + "tile_scheduler_metadata must be [num_sm_parts, ", + kMetaInts, + "], got last dim ", + meta_ints.unwrap()); + RuntimeCheck(b >= 1, "batch size must be positive, got ", b); + RuntimeCheck(num_sm_parts >= 1, "num_sm_parts must be positive, got ", num_sm_parts); + RuntimeCheck(block_size_n >= 1, "block_size_n must be positive, got ", block_size_n); + + auto opt_ptr = [&](const tvm::ffi::Optional& t, const char* name) -> const int* { + if (!t.has_value()) return nullptr; + auto n = SymbolicSize{"batch_size"}; + auto dev = SymbolicDevice{}; + dev.set_options(); + TensorMatcher({n}).with_strides({1}).with_dtype().with_device(dev).verify(t.value()); + RuntimeCheck(static_cast(n.unwrap()) == b, name, " must have ", b, " entries, got ", n.unwrap()); + return static_cast(t.value().data_ptr()); + }; + + RuntimeCheck(topk != -1 || seqlens_k.has_value(), "a dense schedule (topk == -1) needs seqlens_k"); + + const Params p{ + b, + static_cast(block_size_n), + static_cast(fixed_overhead_num_blocks), + static_cast(topk), + static_cast(extra_topk), + opt_ptr(topk_length, "topk_length"), + opt_ptr(extra_topk_length, "extra_topk_length"), + opt_ptr(seqlens_k, "seqlens_k"), + static_cast(tile_scheduler_metadata.data_ptr()), + static_cast(num_splits.data_ptr()), + num_sm_parts, + }; + + const std::size_t smem = sizeof(int) * static_cast(b * 5 + 1 + num_sm_parts * kMetaInts); + RuntimeCheck(smem <= 48 * 1024, "schedule does not fit in shared memory: ", smem, " bytes"); + LaunchKernel(1, kBlockSize, device_.unwrap(), smem)(decoding_sched_meta_kernel, p); +} + +} // namespace sglang diff --git a/python/sglang/kernels/ops/attention/dsv4/decoding_sched_meta.py b/python/sglang/kernels/ops/attention/dsv4/decoding_sched_meta.py new file mode 100644 index 000000000000..72da22f7ed1c --- /dev/null +++ b/python/sglang/kernels/ops/attention/dsv4/decoding_sched_meta.py @@ -0,0 +1,73 @@ +from __future__ import annotations + +from typing import TYPE_CHECKING, Optional + +import torch + +from sglang.kernels.jit.utils import cache_once, load_jit + +from .utils import make_name + +if TYPE_CHECKING: + from tvm_ffi.module import Module + +# sizeof(DecodingSchedMeta) / 4, fixed by FlashMLA's params.h. +META_INTS = 8 + + +@cache_once +def _jit_decoding_sched_meta_module() -> Module: + return load_jit( + make_name("decoding_sched_meta"), + cuda_files=["deepseek_v4/decoding_sched_meta.cuh"], + cuda_wrappers=[("decoding_sched_meta", "decoding_sched_meta")], + ) + + +def decoding_sched_meta( + tile_scheduler_metadata: torch.Tensor, + num_splits: torch.Tensor, + *, + topk_length: Optional[torch.Tensor] = None, + extra_topk_length: Optional[torch.Tensor] = None, + seqlens_k: Optional[torch.Tensor] = None, + block_size_n: int, + fixed_overhead_num_blocks: int, + topk: int, + extra_topk: int = 0, +) -> None: + """Fill FlashMLA's split-KV tile-scheduler metadata in place. + + Same schedule FlashMLA computes for itself when handed no metadata, but the + per-request pass and the write-out run over a whole block instead of one + warp, and the sequential partition walk writes to shared memory rather than + storing each 32-byte entry to global. At ``num_sm_parts = 152`` (a BS=1 + draft step on GB300) that is the difference between 25.1 us and ~2 us, all + of it on the critical path inside the decode graph. + + Pass the filled tensors to FlashMLA as the cached + ``tile_scheduler_metadata`` / ``num_splits`` and it skips its own kernel. + + Args: + tile_scheduler_metadata: ``[num_sm_parts, 8]`` int32, written. + num_splits: ``[batch_size + 1]`` int32, written. + topk_length: ``[batch_size]`` int32 per-request candidate count, or None + to use ``topk`` for every request. + extra_topk_length: the same for the extra cache, when ``extra_topk``. + seqlens_k: ``[batch_size]`` int32, required only for a dense schedule. + block_size_n: the kernel's KV block size. + fixed_overhead_num_blocks: the implementation's per-request overhead. + topk: the sparse top-k, or -1 for a dense model. + extra_topk: the extra cache's top-k, 0 when there is none. + """ + _jit_decoding_sched_meta_module().decoding_sched_meta( + tile_scheduler_metadata, + num_splits, + topk_length, + extra_topk_length, + seqlens_k, + block_size_n, + fixed_overhead_num_blocks, + topk, + extra_topk, + ) diff --git a/python/sglang/kernels/ops/layernorm/mxfp8_epilogue.py b/python/sglang/kernels/ops/layernorm/mxfp8_epilogue.py new file mode 100644 index 000000000000..aeca24f6ba60 --- /dev/null +++ b/python/sglang/kernels/ops/layernorm/mxfp8_epilogue.py @@ -0,0 +1,130 @@ +"""MXFP8 quantization written as an epilogue of the kernel that produces the row. + +At the speculative BS=1 shapes (<=8 rows) the standalone FlashInfer +``mxfp8_quantize`` launch costs about as much as the norm it follows, even +though it only re-reads what that norm just wrote. These kernels do the norm +and the quantization in one pass; the quantized values come from the same BF16 +rounding the standalone pair produces, so both outputs are bitwise identical to +``rmsnorm`` followed by ``mxfp8_quantize(..., is_sf_swizzled_layout=True)``. + +The scale factors use FlashInfer's 128x4 swizzle: +``off = (g // 4) * 512 + ((r % 32) * 4 + ((r // 32) % 4)) * 4 + (g % 4)`` +over a row count padded to a multiple of 128, with the UE8M0 conversion +(positive rounding, subnormals included) that ``mxfp8_quantize`` uses. +""" + +from __future__ import annotations + +import torch +import triton +import triton.language as tl + + +@triton.jit +def _mxfp8_epilogue( + y, row, Q, S, K: tl.constexpr, BLOCK: tl.constexpr, GROUPS: tl.constexpr, g_lo, g_hi +): + """Quantize one BF16 row tile (already masked to zero past K) into Q/S. + + Only groups in ``[g_lo, g_hi)`` are stored, so the row can be split across + CTAs that each recompute the shared statistic and write their own slice. + """ + GP: tl.constexpr = BLOCK // 32 + g = tl.arange(0, GP) + gmask = (g < GROUPS) & (g >= g_lo) & (g < g_hi) + e = tl.arange(0, 32) + idx = g[:, None] * 32 + e[None, :] + v = tl.reshape(y.to(tl.float32), (GP, 32)) + amax = tl.max(tl.abs(v), 1) + normalized = amax * (1.0 / 448.0) + bits = normalized.to(tl.int32, bitcast=True) + exponent = (bits >> 23) & 255 + mantissa = bits & 0x7FFFFF + bump = (mantissa != 0) & ~((exponent == 0) & (mantissa <= 0x400000)) + sf = tl.where(normalized <= 0, 0, tl.minimum(exponent + bump.to(tl.int32), 254)) + scale = tl.where(sf == 0, 0, ((254 - sf) << 23)).to(tl.float32, bitcast=True) + q = tl.minimum(tl.maximum(v * scale[:, None], -448.0), 448.0).to(tl.float8e4nv) + tl.store(Q + row * K + idx, q, gmask[:, None]) + off = (g // 4) * 512 + ((row % 32) * 4 + ((row // 32) % 4)) * 4 + (g % 4) + tl.store(S + off, sf.to(tl.uint8), gmask) + + +@triton.jit +def _hc_combine_norm_mxfp8_kernel( + X, + P, + W, + Y, + Q, + S, + SX: tl.constexpr, + SP: tl.constexpr, + EPS: tl.constexpr, + K: tl.constexpr, + BLOCK: tl.constexpr, + GROUPS: tl.constexpr, + PARTS: tl.constexpr, + SLICE: tl.constexpr, +): + row, part = tl.program_id(0), tl.program_id(1) + h = tl.arange(0, BLOCK) + m = h < K + value = tl.full((BLOCK,), 0, tl.float32) + for c in tl.static_range(4): + pre = tl.load(P + row * SP + c).to(tl.float32) + x = tl.load(X + row * SX + c * K + h, m, 0).to(tl.float32) + value += x * pre + # The unfused combine stores BF16 before RMSNorm reads it. + value = value.to(tl.bfloat16).to(tl.float32) + inv_rms = tl.rsqrt(tl.sum(value * value, 0) / K + EPS) + weight = tl.load(W + h, m, 0).to(tl.float32) + y = (value * inv_rms * weight).to(tl.bfloat16) + tl.store(Y + row * K + h, y, m & (h >= part * SLICE) & (h < (part + 1) * SLICE)) + _mxfp8_epilogue( + y, row, Q, S, K, BLOCK, GROUPS, part * (SLICE // 32), (part + 1) * (SLICE // 32) + ) + + +def _parts_for(m: int, k: int) -> int: + """Row splits: recomputing the statistic beats running 6 CTAs on 148 SMs.""" + parts = 4 + while parts > 1 and (k % (parts * 32)): + parts //= 2 + return parts + + +def _alloc(m, k, device): + q = torch.empty((m, k), dtype=torch.float8_e4m3fn, device=device) + s = torch.zeros( + (k // 32) * (triton.cdiv(m, 128) * 128), dtype=torch.uint8, device=device + ) + return q, s + + +def hc_combine_norm_mxfp8( + x: torch.Tensor, pre: torch.Tensor, weight: torch.Tensor, eps: float +): + """Four-stream combine + RMSNorm returning ``(y_bf16, y_q, y_sf)``.""" + m = x.shape[0] + k = x.shape[1] // 4 + y = torch.empty((m, k), dtype=x.dtype, device=x.device) + q, s = _alloc(m, k, x.device) + parts = _parts_for(m, k) + _hc_combine_norm_mxfp8_kernel[(m, parts)]( + x, + pre, + weight, + y, + q, + s, + SX=x.stride(0), + SP=pre.stride(0), + EPS=eps, + K=k, + BLOCK=triton.next_power_of_2(k), + GROUPS=k // 32, + PARTS=parts, + SLICE=k // parts, + num_warps=8, + ) + return y, q, s diff --git a/python/sglang/kernels/ops/speculative/dspark/dspark_accept.py b/python/sglang/kernels/ops/speculative/dspark/dspark_accept.py index a63748c0db36..cbc52989e86b 100644 --- a/python/sglang/kernels/ops/speculative/dspark/dspark_accept.py +++ b/python/sglang/kernels/ops/speculative/dspark/dspark_accept.py @@ -605,9 +605,7 @@ def accept_greedy( cutoff_verify_lens: Optional[torch.Tensor] = None, ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: bs = candidates.shape[0] - target_predict = torch.argmax(target_logits, dim=-1).view( - bs, verify_num_draft_tokens - ) + target_predict = _row_argmax(target_logits).view(bs, verify_num_draft_tokens) correct_len, bonus = compute_dflash_correct_drafts_and_bonus( candidates=candidates, target_predict=target_predict, @@ -649,6 +647,24 @@ def gather_row_bonus_triton(*, table: torch.Tensor, idx: torch.Tensor) -> torch. return out +def _row_argmax(logits: torch.Tensor) -> torch.Tensor: + """``logits.argmax(-1)``; the speculative shape is few rows over a wide vocab, + where ``torch.argmax``'s single-block-per-row reduction is ~7x off the memory + the reduction touches. Falls back for anything the split kernel does not cover.""" + if ( + logits.is_cuda + and logits.dim() == 2 + and logits.dtype == torch.float32 + and logits.stride(1) == 1 + and logits.shape[0] <= 64 + and logits.shape[1] >= 4096 + ): + from sglang.kernels.ops.speculative.dspark.fast_argmax import fast_row_argmax + + return fast_row_argmax(logits) + return torch.argmax(logits, dim=-1) + + def accept_greedy_triton( *, candidates: torch.Tensor, @@ -657,9 +673,7 @@ def accept_greedy_triton( cutoff_verify_lens: Optional[torch.Tensor] = None, ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: bs = candidates.shape[0] - target_predict = torch.argmax(target_logits, dim=-1).view( - bs, verify_num_draft_tokens - ) + target_predict = _row_argmax(target_logits).view(bs, verify_num_draft_tokens) correct_len, bonus = compute_dflash_correct_drafts_and_bonus( candidates=candidates, target_predict=target_predict, diff --git a/python/sglang/kernels/ops/speculative/dspark/fast_argmax.py b/python/sglang/kernels/ops/speculative/dspark/fast_argmax.py new file mode 100644 index 000000000000..ab948a32bf8a --- /dev/null +++ b/python/sglang/kernels/ops/speculative/dspark/fast_argmax.py @@ -0,0 +1,68 @@ +"""Row-wise argmax for the tall-and-thin speculative verify logits. + +``torch.argmax`` dispatches ``at::native::reduce_kernel`` here, which needs +~24 us for the [6, 129280] FP32 verify logits on GB300 -- two orders of +magnitude off the 3 MB the reduction actually reads. The rows are few and very +wide, so a flat two-stage split over the vocabulary saturates the machine +instead of leaving it to one block per row. + +Ties resolve to the lowest index, matching ``ArgMaxOps``' strict ``>``. +""" + +import torch +import triton +import triton.language as tl + + +@triton.jit +def _argmax_partial_kernel( + X, OUTV, OUTI, N, SX, SPLITS: tl.constexpr, BLOCK: tl.constexpr +): + row = tl.program_id(0) + part = tl.program_id(1) + per = tl.cdiv(N, SPLITS) + start = part * per + best_v = float("-inf") + best_i = N + for off in tl.range(start, tl.minimum(start + per, N), BLOCK): + idx = off + tl.arange(0, BLOCK) + m = idx < tl.minimum(start + per, N) + v = tl.load(X + row * SX + idx, m, float("-inf")) + cur_v = tl.max(v, 0) + # lowest index among the maxima of this tile + cur_i = tl.min(tl.where(v == cur_v, idx, N), 0) + take = (cur_v > best_v) | ((cur_v == best_v) & (cur_i < best_i)) + best_i = tl.where(take, cur_i, best_i) + best_v = tl.where(take, cur_v, best_v) + tl.store(OUTV + row * SPLITS + part, best_v) + tl.store(OUTI + row * SPLITS + part, best_i) + + +@triton.jit +def _argmax_final_kernel(INV, INI, OUT, SPLITS: tl.constexpr, BLOCK: tl.constexpr): + row = tl.program_id(0) + o = tl.arange(0, BLOCK) + m = o < SPLITS + v = tl.load(INV + row * SPLITS + o, m, float("-inf")) + i = tl.load(INI + row * SPLITS + o, m, 0x7FFFFFFF) + best_v = tl.max(v, 0) + best_i = tl.min(tl.where(v == best_v, i, 0x7FFFFFFF), 0) + tl.store(OUT + row, best_i.to(tl.int64)) + + +_SPLITS = 64 + + +def fast_row_argmax(x: torch.Tensor, out: torch.Tensor | None = None) -> torch.Tensor: + """``x.argmax(dim=-1)`` for a 2D FP32 tensor with few rows and a wide vocab.""" + assert x.dim() == 2 and x.dtype == torch.float32 and x.stride(1) == 1 + rows, n = x.shape + if out is None: + out = torch.empty((rows,), dtype=torch.int64, device=x.device) + pv = torch.empty((rows, _SPLITS), dtype=torch.float32, device=x.device) + pi = torch.empty((rows, _SPLITS), dtype=torch.int32, device=x.device) + _argmax_partial_kernel[(rows, _SPLITS)]( + x, pv, pi, n, x.stride(0), SPLITS=_SPLITS, BLOCK=2048, num_warps=8 + ) + _argmax_final_kernel[(rows,)](pv, pi, out, SPLITS=_SPLITS, BLOCK=64, num_warps=2) + return out diff --git a/python/sglang/srt/environ.py b/python/sglang/srt/environ.py index 77f6eb8e0540..3cea1c1d0ec2 100644 --- a/python/sglang/srt/environ.py +++ b/python/sglang/srt/environ.py @@ -1443,6 +1443,10 @@ class Envs: SGLANG_DSV41_DEEP_GEMM_CANDIDATE_INDEXER = EnvBool(False) # use multistream to overlap the publish-side with other computation SGLANG_DSV41_DEEP_GEMM_CANDIDATE_OVERLAP = EnvBool(True) + # Compute FlashMLA's split-KV tile-scheduler metadata with the wide DSV4 + # kernel instead of FlashMLA's one-warp one, which walks num_sm_parts + # serially on a single thread inside the decode graph. Off = FlashMLA's. + SGLANG_DSV41_FAST_FLASHMLA_SCHED = EnvBool(True) # Keep the DeepSeek-V4.1 engram tables in host memory (layout below) and gather # rows from the GPU instead of sharding them over HBM. SGLANG_ENABLE_DSV41_ENGRAM_HOST_TABLE = EnvBool(False) diff --git a/python/sglang/srt/layers/attention/deepseek_v4_backend.py b/python/sglang/srt/layers/attention/deepseek_v4_backend.py index 41c7dcabc5b3..e85b61df27f8 100644 --- a/python/sglang/srt/layers/attention/deepseek_v4_backend.py +++ b/python/sglang/srt/layers/attention/deepseek_v4_backend.py @@ -199,6 +199,74 @@ def _create_flashmla_metadata(): return flash_mla.get_mla_metadata()[0] +# The head64 sm100 decode scheduling constants, and the partition count +# `num_sm_parts` that goes with them. Not exported, so the fast schedule only +# runs for the shape they are known for and FlashMLA's own shape check is what +# catches it if they ever stop matching. +_FLASHMLA_SCHED_BLOCK_SIZE_N = 64 +_FLASHMLA_SCHED_FIXED_OVERHEAD = 5 + + +@functools.lru_cache(maxsize=None) +def _num_sms(device_index: int) -> int: + return torch.cuda.get_device_properties(device_index).multi_processor_count + + +def _fast_flashmla_sched_shape(q: torch.Tensor) -> bool: + return q.is_cuda and get_platform().is_blackwell and q.shape[-2] == 64 + + +def _maybe_precompute_flashmla_sched_meta( + flashmla_metadata, + *, + q: torch.Tensor, + indices: torch.Tensor, + topk_length: Optional[torch.Tensor], + extra_indices: Optional[torch.Tensor], + extra_topk_length: Optional[torch.Tensor], +) -> None: + """Compute FlashMLA's split-KV schedule before it has to. + + `sparse_decode_fwd` builds the schedule itself whenever it is handed none, + in a `<<<1, 32>>>` kernel whose partition loop runs on thread 0 and stores + each 32-byte entry to global memory. At the 152 partitions of a BS=1 step + that is 28 us, and a decode graph replays it on the critical path. Filling + the buffers here instead means FlashMLA finds them already populated and + skips its kernel; `decoding_sched_meta` produces the same schedule, bit for + bit, in about 9 us. + + Only fires where FlashMLA would have computed -- when the scheduler holds no + buffers yet -- so this does not add work to the calls that already reuse one. + """ + if flashmla_metadata is None or not envs.SGLANG_DSV41_FAST_FLASHMLA_SCHED.get(): + return + if getattr(flashmla_metadata, "tile_scheduler_metadata", None) is not None: + return + if not _fast_flashmla_sched_shape(q): + return + from sglang.kernels.ops.attention.dsv4.decoding_sched_meta import ( + META_INTS, + decoding_sched_meta, + ) + + b, s_q = q.shape[0], q.shape[1] + num_sm_parts = max(_num_sms(q.device.index) // s_q, 1) + meta = torch.empty((num_sm_parts, META_INTS), dtype=torch.int32, device=q.device) + num_splits = torch.empty((b + 1,), dtype=torch.int32, device=q.device) + decoding_sched_meta( + meta, + num_splits, + topk_length=topk_length, + extra_topk_length=extra_topk_length, + block_size_n=_FLASHMLA_SCHED_BLOCK_SIZE_N, + fixed_overhead_num_blocks=_FLASHMLA_SCHED_FIXED_OVERHEAD, + topk=indices.shape[-1], + extra_topk=0 if extra_indices is None else extra_indices.shape[-1], + ) + flashmla_metadata.tile_scheduler_metadata = meta + flashmla_metadata.num_splits = num_splits + + def _expand_index_page_table( page_table: torch.Tensor, *, @@ -3731,6 +3799,14 @@ def match_num_queries(x, value): else: from sgl_kernel.flash_mla import flash_mla_with_kvcache + _maybe_precompute_flashmla_sched_meta( + flashmla_metadata, + q=q, + indices=swa_page_indices, + topk_length=swa_topk_lengths, + extra_indices=extra_indices, + extra_topk_length=extra_topk_lengths, + ) o = flash_mla_with_kvcache( q=q, k_cache=swa_k_cache, diff --git a/python/sglang/srt/models/deepseek_v4.py b/python/sglang/srt/models/deepseek_v4.py index 9f5bc2d6e391..12385b9f6788 100644 --- a/python/sglang/srt/models/deepseek_v4.py +++ b/python/sglang/srt/models/deepseek_v4.py @@ -1158,6 +1158,36 @@ def _get_npu_rope_position_cache( inverse=inverse, ) + def accepts_mxfp8_swizzled_input(self) -> bool: + """Whether the first projection consumes a 128x4 MXFP8 activation tuple.""" + cached = getattr(self, "_accepts_mxfp8_swizzled_input", None) + if cached is not None: + return cached + if self.fuse_wqa_wkv: + linears = [getattr(self, "wqkv_a", None)] + else: + # Both projections read the same activation on this path. + linears = [getattr(self, "wq_a", None), getattr(self, "wkv", None)] + + def _takes_swizzled(linear) -> bool: + method = getattr(linear, "quant_method", None) + return bool( + linear is not None + and getattr(method, "mxfp8_dense_backend", None) + in ( + Mxfp8DenseGemmBackend.FLASHINFER_CUTEDSL, + Mxfp8DenseGemmBackend.FLASHINFER_CUTLASS, + ) + and ( + getattr(method, "use_mxfp8", False) + or getattr(linear, "block_fp8_mxfp8_ready", False) + ) + ) + + ok = all(_takes_swizzled(linear) for linear in linears) + self._accepts_mxfp8_swizzled_input = ok + return ok + def _compute_q_a( self, x: torch.Tensor, @@ -2737,10 +2767,16 @@ def _hc_mix_and_combine( apply_pre: Optional[torch.Tensor], norm: RMSNorm, stats_stream: Optional[torch.cuda.Stream] = None, + quantized: Optional[list] = None, ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: """Mixing coefficients come from x; the sublayer input is x collapsed with apply_pre (None selects copy 0), then RMS-normalized. - Returns (y, pre, post, comb).""" + Returns (y, pre, post, comb). + + ``quantized`` (a list) opts into the fused MXFP8 epilogue: the collapsed, + normalized row is also emitted as a ``Mxfp8SwizzledInput`` appended to it, + so the consuming projection skips its own quantization launch.""" + quantize = quantized is not None from sglang.kernels.ops.layernorm.mhc import ( hc_combine, hc_mix_stats, @@ -2767,6 +2803,16 @@ def combine_and_norm(): and norm.variance_size_override is None and not is_batch_invariant_mode_enabled() ): + if quantize: + from sglang.kernels.ops.layernorm.mxfp8_epilogue import ( + hc_combine_norm_mxfp8, + ) + + y, y_q, y_sf = hc_combine_norm_mxfp8( + x_flat, apply_pre, norm.weight, norm.variance_epsilon + ) + quantized.append(Mxfp8SwizzledInput(y_q, y_sf)) + return y from sglang.kernels.ops.layernorm.hc_combine_norm import hc_combine_norm return hc_combine_norm( @@ -2869,6 +2915,9 @@ def forward_hc_pre_from_prev( the FFN consumes this attention's. Returns (hidden_states, ffn_pre).""" stats_stream = self._get_hc_stats_stream(hidden_states, forward_batch) residual = hidden_states + attn_quantized: Optional[list] = ( + [] if self.self_attn.accepts_mxfp8_swizzled_input() else None + ) x, attn_pre, attn_post, attn_comb = self._hc_mix_and_combine( hidden_states, self.hc_attn_fn, @@ -2877,10 +2926,14 @@ def forward_hc_pre_from_prev( apply_pre=prev_pre, norm=self.input_layernorm, stats_stream=stats_stream, + quantized=attn_quantized, ) with self.self_attn.maybe_use_decode_attn_tp(forward_batch): x = self.self_attn( - x=x, positions=positions, forward_batch=forward_batch, x_quant=None + x=x, + positions=positions, + forward_batch=forward_batch, + x_quant=attn_quantized[0] if attn_quantized else None, ) if stats_stream is not None: torch.cuda.current_stream().wait_stream(stats_stream) diff --git a/python/sglang/srt/multimodal/dsv41/vl_routing.py b/python/sglang/srt/multimodal/dsv41/vl_routing.py index abbe98fd2827..176e45f10c8f 100644 --- a/python/sglang/srt/multimodal/dsv41/vl_routing.py +++ b/python/sglang/srt/multimodal/dsv41/vl_routing.py @@ -4,6 +4,7 @@ from sglang.srt.layers.moe.topk import ( _RENORMALIZE_SUM_EPSILON, StandardTopKOutput, + StandardTopKOutputPacked, _mask_topk_ids_padded_region, _zero_topk_weights_padded_region, ) @@ -31,6 +32,20 @@ def vision_topk(moe, logits, input_ids, num_token_non_padded=None): ) if is_cuda(): from sglang.kernels.ops.moe.moe_fused_gate import moe_fused_gate + from sglang.srt.layers.moe.utils import get_moe_runner_backend + + # Same admission as _fused_gate_emits_packed_ids on the text path: only + # flashinfer_mxfp4 consumes the packed form, and nothing may rewrite ids + # or weights after the router -- which rules out the shared-expert slots + # that _scale_fused_shared_weights rescales below. + packed_topk = None + if ( + num_fused_shared_experts == 0 + and get_moe_runner_backend().is_flashinfer_mxfp4() + ): + packed_topk = torch.empty( + (logits.shape[0], config.top_k), dtype=torch.int32, device=logits.device + ) weights, indices = moe_fused_gate( logits, @@ -46,12 +61,15 @@ def vision_topk(moe, logits, input_ids, num_token_non_padded=None): routed_scaling_factor=config.routed_scaling_factor, apply_routed_scaling_factor_on_output=config.apply_routed_scaling_factor_on_output, num_token_non_padded=num_token_non_padded, + packed_out=packed_topk, ) weights = _scale_fused_shared_weights( weights, num_fused_shared_experts, config.fused_shared_experts_scaling_factor, ) + if packed_topk is not None: + return StandardTopKOutputPacked(weights, indices, logits, packed_topk) return StandardTopKOutput(weights, indices, logits) scores = F.softplus(logits.float()).sqrt() if input_ids is None: diff --git a/test/registered/kernel/attention/test_decoding_sched_meta.py b/test/registered/kernel/attention/test_decoding_sched_meta.py new file mode 100644 index 000000000000..a620401f8e7c --- /dev/null +++ b/test/registered/kernel/attention/test_decoding_sched_meta.py @@ -0,0 +1,183 @@ +import sys + +import pytest +import torch + +from sglang.srt.runtime_context import get_platform +from sglang.test.ci.ci_register import register_cuda_ci + +register_cuda_ci(est_time=25, stage="base-b-kernel-unit", runner_config="4-gpu-b200") + +pytestmark = pytest.mark.skipif( + not torch.cuda.is_available() + or torch.version.cuda is None + or not get_platform().is_blackwell, + reason="the FlashMLA split-KV schedule is Blackwell-only here", +) + +H_Q, D_QK, D_V = 64, 512, 512 +# The only cache format this build's sparse decode takes for both caches, and a +# page size that keeps page * bytes_per_token a multiple of 576. +BYTES_PER_TOKEN, PAGE = 584, 288 +BLOCK_SIZE_N, FIXED_OVERHEAD = 64, 5 + + +def _cache(): + n = 8192 // PAGE + 16 + return torch.randint( + 0, 200, (n, PAGE, 1, BYTES_PER_TOKEN), device="cuda", dtype=torch.uint8 + ) + + +def _call(kv, b, s_q, topk, topk_length, *, extra=None, meta=None): + import sgl_kernel.flash_mla as flash_mla + + g = torch.Generator(device="cuda").manual_seed(b * 7 + s_q * 13 + topk) + q = torch.randn( + (b, s_q, H_Q, D_QK), device="cuda", dtype=torch.bfloat16, generator=g + ) + indices = torch.randint( + 0, 5120, (b, s_q, topk), device="cuda", dtype=torch.int32, generator=g + ) + sink = torch.randn((H_Q,), device="cuda", dtype=torch.float32, generator=g) + kwargs = {} + if extra is not None: + extra_kv, extra_indices, extra_topk_length = extra + kwargs = dict( + extra_k_cache=extra_kv, + extra_indices_in_kvcache=extra_indices, + extra_topk_length=extra_topk_length, + ) + sched = flash_mla.FlashMLASchedMeta() + if meta is not None: + sched.tile_scheduler_metadata, sched.num_splits = meta + out, lse = flash_mla.flash_mla_with_kvcache( + q, + kv, + None, + None, + D_V, + sched, + indices=indices, + is_fp8_kvcache=True, + softmax_scale=0.1, + causal=False, + topk_length=topk_length, + attn_sink=sink, + **kwargs, + ) + torch.cuda.synchronize() + return out, lse, sched + + +def _ours( + like_meta, like_splits, topk_length, topk, *, extra_topk_length=None, extra_topk=0 +): + from sglang.kernels.ops.attention.dsv4.decoding_sched_meta import ( + decoding_sched_meta, + ) + + meta = torch.empty_like(like_meta) + splits = torch.empty_like(like_splits) + decoding_sched_meta( + meta, + splits, + topk_length=topk_length, + extra_topk_length=extra_topk_length, + block_size_n=BLOCK_SIZE_N, + fixed_overhead_num_blocks=FIXED_OVERHEAD, + topk=topk, + extra_topk=extra_topk, + ) + return meta, splits + + +def _lengths(b, topk, mode, seed): + g = torch.Generator(device="cuda").manual_seed(seed) + if mode == "full": + return torch.full((b,), topk, device="cuda", dtype=torch.int32) + if mode == "ones": + return torch.ones((b,), device="cuda", dtype=torch.int32) + if mode == "zeros": + return torch.zeros((b,), device="cuda", dtype=torch.int32) + lengths = torch.randint( + 0, topk + 1, (b,), device="cuda", dtype=torch.int32, generator=g + ) + if mode == "mixed": + lengths[0] = 0 + lengths[-1] = topk + return lengths + + +# FlashMLA's DecodingSchedMeta ends in a `_pad` word it never writes, so the +# reference carries whatever torch::empty left there. +DEFINED = slice(0, 7) + + +@pytest.mark.parametrize("b", [1, 2, 5, 6, 17, 64]) +@pytest.mark.parametrize("s_q", [1, 6]) +@pytest.mark.parametrize("topk", [512, 2048]) +@pytest.mark.parametrize("mode", ["full", "random", "zeros", "ones", "mixed"]) +def test_matches_flashmla_schedule(b: int, s_q: int, topk: int, mode: str): + kv = _cache() + topk_length = _lengths(b, topk, mode, b * 1000 + s_q * 37 + topk + len(mode)) + _, _, sched = _call(kv, b, s_q, topk, topk_length) + if sched.tile_scheduler_metadata is None: + pytest.skip("FlashMLA did not split the KV for this shape") + meta, splits = _ours( + sched.tile_scheduler_metadata, sched.num_splits, topk_length, topk + ) + assert torch.equal(meta[:, DEFINED], sched.tile_scheduler_metadata[:, DEFINED]) + assert torch.equal(splits, sched.num_splits) + + +@pytest.mark.parametrize("b", [1, 2, 8]) +@pytest.mark.parametrize("s_q", [1, 6]) +@pytest.mark.parametrize("topk", [512, 2048]) +def test_attention_output_is_unchanged(b: int, s_q: int, topk: int): + kv = _cache() + topk_length = _lengths(b, topk, "random", b + s_q + topk) + ref_out, ref_lse, sched = _call(kv, b, s_q, topk, topk_length) + if sched.tile_scheduler_metadata is None: + pytest.skip("FlashMLA did not split the KV for this shape") + meta = _ours(sched.tile_scheduler_metadata, sched.num_splits, topk_length, topk) + out, lse, _ = _call(kv, b, s_q, topk, topk_length, meta=meta) + # Bitwise: a random fp8 cache legitimately decodes to NaN in places. + assert torch.equal(out.view(torch.int16), ref_out.view(torch.int16)) + assert torch.equal(lse.view(torch.int32), ref_lse.view(torch.int32)) + + +@pytest.mark.parametrize("b", [1, 2, 8]) +@pytest.mark.parametrize("topk,extra_topk", [(512, 512), (2048, 512), (512, 2048)]) +def test_extra_cache_schedule(b: int, topk: int, extra_topk: int): + kv, extra_kv = _cache(), _cache() + s_q = 1 + g = torch.Generator(device="cuda").manual_seed(b + topk + extra_topk) + topk_length = _lengths(b, topk, "random", b + topk) + extra_topk_length = torch.randint( + 1, extra_topk + 1, (b,), device="cuda", dtype=torch.int32, generator=g + ) + extra_indices = torch.randint( + 0, 4096, (b, s_q, extra_topk), device="cuda", dtype=torch.int32, generator=g + ) + extra = (extra_kv, extra_indices, extra_topk_length) + ref_out, ref_lse, sched = _call(kv, b, s_q, topk, topk_length, extra=extra) + if sched.tile_scheduler_metadata is None: + pytest.skip("FlashMLA did not split the KV for this shape") + meta, splits = _ours( + sched.tile_scheduler_metadata, + sched.num_splits, + topk_length, + topk, + extra_topk_length=extra_topk_length, + extra_topk=extra_topk, + ) + assert torch.equal(meta[:, DEFINED], sched.tile_scheduler_metadata[:, DEFINED]) + assert torch.equal(splits, sched.num_splits) + out, lse, _ = _call(kv, b, s_q, topk, topk_length, extra=extra, meta=(meta, splits)) + assert torch.equal(out.view(torch.int16), ref_out.view(torch.int16)) + assert torch.equal(lse.view(torch.int32), ref_lse.view(torch.int32)) + + +if __name__ == "__main__": + sys.exit(pytest.main([__file__])) diff --git a/test/registered/kernel/hyperconnection/test_hc_combine_norm_mxfp8.py b/test/registered/kernel/hyperconnection/test_hc_combine_norm_mxfp8.py new file mode 100644 index 000000000000..e6cef5423681 --- /dev/null +++ b/test/registered/kernel/hyperconnection/test_hc_combine_norm_mxfp8.py @@ -0,0 +1,56 @@ +import sys + +import pytest +import torch + +from sglang.srt.runtime_context import get_platform +from sglang.test.ci.ci_register import register_cuda_ci + +register_cuda_ci(est_time=20, stage="base-b-kernel-unit", runner_config="4-gpu-b200") + +pytestmark = pytest.mark.skipif( + not torch.cuda.is_available() + or torch.version.cuda is None + or not get_platform().is_blackwell, + reason="the MXFP8 reference quantizer is Blackwell-only", +) + +HIDDEN = 5120 +STREAMS = 4 + + +@pytest.mark.parametrize("m", [1, 2, 5, 6, 8]) +@pytest.mark.parametrize("scale", [1e-3, 1.0, 1e3]) +@pytest.mark.parametrize("backend", ["cuda", "cute-dsl"]) +def test_bitwise_identical_to_norm_then_quantize(m: int, scale: float, backend: str): + from sglang.kernels.ops.layernorm.hc_combine_norm import hc_combine_norm + from sglang.kernels.ops.layernorm.mxfp8_epilogue import hc_combine_norm_mxfp8 + from sglang.srt.layers.quantization.fp8_utils import flashinfer_mxfp8_quantize + + g = torch.Generator(device="cuda").manual_seed(m * 31 + int(scale * 1000)) + x = ( + torch.randn( + (m, STREAMS * HIDDEN), device="cuda", dtype=torch.bfloat16, generator=g + ) + * scale + ) + pre = torch.randn( + (m, STREAMS), device="cuda", dtype=torch.bfloat16, generator=g + ).contiguous() + w = torch.randn((HIDDEN,), device="cuda", dtype=torch.bfloat16, generator=g) + eps = 1e-6 + + y_ref = hc_combine_norm(x, pre, w, eps) + q_ref, sf_ref = flashinfer_mxfp8_quantize(y_ref, True, 32, backend) + y, q, sf = hc_combine_norm_mxfp8(x, pre, w, eps) + + assert torch.equal(y, y_ref) + assert torch.equal( + q.reshape(-1).view(torch.uint8), q_ref.reshape(-1).view(torch.uint8) + ) + assert sf.shape == sf_ref.reshape(-1).shape + assert torch.equal(sf, sf_ref.reshape(-1)) + + +if __name__ == "__main__": + sys.exit(pytest.main([__file__])) diff --git a/test/registered/kernel/speculative/test_dspark_fast_argmax.py b/test/registered/kernel/speculative/test_dspark_fast_argmax.py new file mode 100644 index 000000000000..6b206c58879a --- /dev/null +++ b/test/registered/kernel/speculative/test_dspark_fast_argmax.py @@ -0,0 +1,75 @@ +import sys + +import pytest +import torch + +from sglang.kernels.ops.speculative.dspark.fast_argmax import fast_row_argmax +from sglang.test.ci.ci_register import register_cuda_ci + +register_cuda_ci(est_time=20, stage="base-b-kernel-unit", runner_config="1-gpu-small") + +pytestmark = pytest.mark.skipif( + not torch.cuda.is_available(), reason="fast_row_argmax requires CUDA" +) + +VOCAB = 129280 + + +def _check(x: torch.Tensor) -> None: + got = fast_row_argmax(x) + want = torch.argmax(x, dim=-1) + assert got.dtype == want.dtype == torch.int64 + assert torch.equal(got, want), (got - want).nonzero() + + +@pytest.mark.parametrize("rows", [1, 6, 8, 64]) +@pytest.mark.parametrize("vocab", [4096, 32000, VOCAB]) +def test_matches_torch_argmax(rows: int, vocab: int): + g = torch.Generator(device="cuda").manual_seed(rows * 1000 + vocab) + x = torch.randn((rows, vocab), device="cuda", dtype=torch.float32, generator=g) + _check(x) + + +def test_ties_resolve_to_the_lowest_index(): + # Every column equal: torch.argmax returns 0, and so must the split kernel. + x = torch.zeros((6, VOCAB), device="cuda", dtype=torch.float32) + _check(x) + # A tie spanning two different partials, plus one strictly larger value in + # a third, so the final stage has to break a tie and pick a winner. + x = torch.full((6, VOCAB), -1.0, device="cuda", dtype=torch.float32) + x[:, 100] = 5.0 + x[:, VOCAB // 2] = 5.0 + x[3, VOCAB - 7] = 9.0 + _check(x) + + +def test_infinities(): + x = torch.randn((6, VOCAB), device="cuda", dtype=torch.float32) + x[0, :] = float("-inf") # an all -inf row still has to return an index + x[1, 77] = float("inf") + x[2, VOCAB - 1] = float("inf") + x[3, 5] = float("inf") + x[3, 6] = float("inf") # first +inf wins + x[4, :] = float("-inf") + x[4, VOCAB - 2] = 0.0 + _check(x) + + +def test_strided_rows(): + # A column slice of a padded buffer: only the innermost stride is unit. + buf = torch.randn((6, VOCAB + 64), device="cuda", dtype=torch.float32) + x = buf[:, :VOCAB] + assert x.stride(1) == 1 and not x.is_contiguous() + _check(x) + + +def test_out_parameter_is_filled(): + x = torch.randn((6, VOCAB), device="cuda", dtype=torch.float32) + out = torch.empty((6,), device="cuda", dtype=torch.int64) + got = fast_row_argmax(x, out=out) + assert got.data_ptr() == out.data_ptr() + assert torch.equal(out, torch.argmax(x, dim=-1)) + + +if __name__ == "__main__": + sys.exit(pytest.main([__file__]))