diff --git a/benchmark/kernels/attention/fa4_benchmark_utils.py b/benchmark/kernels/attention/fa4_benchmark_utils.py index 8f5884777..749ca50a7 100644 --- a/benchmark/kernels/attention/fa4_benchmark_utils.py +++ b/benchmark/kernels/attention/fa4_benchmark_utils.py @@ -32,250 +32,3 @@ def benchmark_forward( if verbose: print(m) return t, m - - -def benchmark_backward( - fn, - *inputs, - grad=None, - repeats=10, - desc="", - verbose=True, - amp=False, - amp_dtype=torch.float16, - **kwinputs, -): - """Use Pytorch Benchmark on the backward pass of an arbitrary function.""" - if verbose: - print(desc, "- Backward pass") - with torch.autocast(device_type="cuda", dtype=amp_dtype, enabled=amp): - y = fn(*inputs, **kwinputs) - if type(y) is tuple: - y = y[0] - if grad is None: - grad = torch.randn_like(y) - else: - if grad.shape != y.shape: - raise RuntimeError("Grad shape does not match output shape") - - def f(*inputs, y, grad): - # Set .grad to None to avoid extra operation of gradient accumulation - for x in inputs: - if isinstance(x, torch.Tensor): - x.grad = None - y.backward(grad, retain_graph=True) - - t = benchmark.Timer( - stmt="f(*inputs, y=y, grad=grad)", - globals={"f": f, "inputs": inputs, "y": y, "grad": grad}, - num_threads=torch.get_num_threads(), - ) - m = t.timeit(repeats) - if verbose: - print(m) - return t, m - - -def benchmark_combined( - fn, - *inputs, - grad=None, - repeats=10, - desc="", - verbose=True, - amp=False, - amp_dtype=torch.float16, - **kwinputs, -): - """Use Pytorch Benchmark on the forward+backward pass of an arbitrary function.""" - if verbose: - print(desc, "- Forward + Backward pass") - with torch.autocast(device_type="cuda", dtype=amp_dtype, enabled=amp): - y = fn(*inputs, **kwinputs) - if type(y) is tuple: - y = y[0] - if grad is None: - grad = torch.randn_like(y) - else: - if grad.shape != y.shape: - raise RuntimeError("Grad shape does not match output shape") - - def f(grad, *inputs, **kwinputs): - for x in inputs: - if isinstance(x, torch.Tensor): - x.grad = None - with torch.autocast(device_type="cuda", dtype=amp_dtype, enabled=amp): - y = fn(*inputs, **kwinputs) - if type(y) is tuple: - y = y[0] - y.backward(grad, retain_graph=True) - - t = benchmark.Timer( - stmt="f(grad, *inputs, **kwinputs)", - globals={ - "f": f, - "fn": fn, - "inputs": inputs, - "grad": grad, - "kwinputs": kwinputs, - }, - num_threads=torch.get_num_threads(), - ) - m = t.timeit(repeats) - if verbose: - print(m) - return t, m - - -def benchmark_fwd_bwd( - fn, - *inputs, - grad=None, - repeats=10, - desc="", - verbose=True, - amp=False, - amp_dtype=torch.float16, - **kwinputs, -): - """Use Pytorch Benchmark on the forward+backward pass of an arbitrary function.""" - return ( - benchmark_forward( - fn, - *inputs, - repeats=repeats, - desc=desc, - verbose=verbose, - amp=amp, - amp_dtype=amp_dtype, - **kwinputs, - ), - benchmark_backward( - fn, - *inputs, - grad=grad, - repeats=repeats, - desc=desc, - verbose=verbose, - amp=amp, - amp_dtype=amp_dtype, - **kwinputs, - ), - ) - - -def benchmark_all( - fn, - *inputs, - grad=None, - repeats=10, - desc="", - verbose=True, - amp=False, - amp_dtype=torch.float16, - **kwinputs, -): - """Use Pytorch Benchmark on the forward+backward pass of an arbitrary function.""" - return ( - benchmark_forward( - fn, - *inputs, - repeats=repeats, - desc=desc, - verbose=verbose, - amp=amp, - amp_dtype=amp_dtype, - **kwinputs, - ), - benchmark_backward( - fn, - *inputs, - grad=grad, - repeats=repeats, - desc=desc, - verbose=verbose, - amp=amp, - amp_dtype=amp_dtype, - **kwinputs, - ), - benchmark_combined( - fn, - *inputs, - grad=grad, - repeats=repeats, - desc=desc, - verbose=verbose, - amp=amp, - amp_dtype=amp_dtype, - **kwinputs, - ), - ) - - -def pytorch_profiler( - fn, - *inputs, - trace_filename=None, - backward=False, - amp=False, - amp_dtype=torch.float16, - cpu=False, - verbose=True, - **kwinputs, -): - """Wrap benchmark functions in Pytorch profiler to see CUDA information.""" - if backward: - with torch.autocast(device_type="cuda", dtype=amp_dtype, enabled=amp): - out = fn(*inputs, **kwinputs) - if type(out) is tuple: - out = out[0] - g = torch.randn_like(out) - for _ in range(30): # Warm up - if backward: - for x in inputs: - if isinstance(x, torch.Tensor): - x.grad = None - with torch.autocast(device_type="cuda", dtype=amp_dtype, enabled=amp): - out = fn(*inputs, **kwinputs) - if type(out) is tuple: - out = out[0] - # Backward should be done outside autocast - if backward: - out.backward(g, retain_graph=True) - activities = ([torch.profiler.ProfilerActivity.CPU] if cpu else []) + [ - torch.profiler.ProfilerActivity.CUDA - ] - with torch.profiler.profile( - activities=activities, - record_shapes=True, - # profile_memory=True, - with_stack=True, - ) as prof: - if backward: - for x in inputs: - if isinstance(x, torch.Tensor): - x.grad = None - with torch.autocast(device_type="cuda", dtype=amp_dtype, enabled=amp): - out = fn(*inputs, **kwinputs) - if type(out) is tuple: - out = out[0] - if backward: - out.backward(g, retain_graph=True) - if verbose: - # print(prof.key_averages().table(sort_by="self_cuda_time_total", row_limit=50)) - print(prof.key_averages().table(row_limit=50)) - if trace_filename is not None: - prof.export_chrome_trace(trace_filename) - - -def benchmark_memory(fn, *inputs, desc="", verbose=True, **kwinputs): - torch.cuda.empty_cache() - torch.cuda.reset_peak_memory_stats() - torch.cuda.synchronize() - fn(*inputs, **kwinputs) - torch.cuda.synchronize() - mem = torch.cuda.max_memory_allocated() / ((2**20) * 1000) - if verbose: - print(f"{desc} max memory: {mem}GB") - torch.cuda.empty_cache() - return mem diff --git a/benchmark/kernels/deepep/deepep_utils.py b/benchmark/kernels/deepep/deepep_utils.py deleted file mode 100644 index 573d68deb..000000000 --- a/benchmark/kernels/deepep/deepep_utils.py +++ /dev/null @@ -1,218 +0,0 @@ -# ADAPTED FROM https://github.com/deepseek-ai/DeepEP/blob/main/tests/utils.py - -import os -import sys -from typing import Optional - -import numpy as np -import torch -import torch.distributed as dist - - -def init_dist(local_rank: int, num_local_ranks: int, args): - ip = args.master_addr - port = args.master_port - num_nodes = args.nnodes - node_rank = args.node_rank - assert (num_local_ranks < 8 and num_nodes == 1) or num_local_ranks == 8 - - dist.init_process_group( - backend="nccl", - init_method=f"tcp://{ip}:{port}", - world_size=num_nodes * num_local_ranks, - rank=node_rank * num_local_ranks + local_rank, - ) - torch.set_default_dtype(torch.bfloat16) - torch.set_default_device("cuda") - torch.cuda.set_device(local_rank) - - return ( - dist.get_rank(), - dist.get_world_size(), - dist.new_group(list(range(num_local_ranks * num_nodes))), - ) - - -def calc_diff(x: torch.Tensor, y: torch.Tensor): - x, y = x.double() + 1, y.double() + 1 - denominator = (x * x + y * y).sum() - sim = 2 * (x * y).sum() / denominator - return (1 - sim).item() - - -def per_token_cast_to_fp8(x: torch.Tensor): - assert x.dim() == 2 and x.size(1) % 128 == 0 - m, n = x.shape - x_view = x.view(m, -1, 128) - x_amax = x_view.abs().float().amax(dim=2).view(m, -1).clamp(1e-4) - return (x_view * (448.0 / x_amax.unsqueeze(2))).to(torch.float8_e4m3fn).view( - m, n - ), (x_amax / 448.0).view(m, -1) - - -def per_token_cast_back(x_fp8: torch.Tensor, x_scales: torch.Tensor): - x_fp32 = x_fp8.to(torch.float32).view(x_fp8.size(0), -1, 128) - x_scales = x_scales.view(x_fp8.size(0), -1, 1) - return (x_fp32 * x_scales).view(x_fp8.shape).to(torch.bfloat16) - - -def inplace_unique(x: torch.Tensor, num_slots: int): - assert x.dim() == 2 - mask = x < 0 - x_padded = x.masked_fill(mask, num_slots) - bin_count = torch.zeros((x.size(0), num_slots + 1), dtype=x.dtype, device=x.device) - bin_count.scatter_add_(1, x_padded, torch.ones_like(x_padded)) - bin_count = bin_count[:, :num_slots] - sorted_bin_count, sorted_bin_idx = torch.sort(bin_count, dim=-1, descending=True) - sorted_bin_idx.masked_fill_(sorted_bin_count == 0, -1) - sorted_bin_idx = torch.sort(sorted_bin_idx, descending=True, dim=-1).values - x[:, :].fill_(-1) - valid_len = min(num_slots, x.size(1)) - x[:, :valid_len] = sorted_bin_idx[:, :valid_len] - - -def create_grouped_scores( - scores: torch.Tensor, group_idx: torch.Tensor, num_groups: int -): - num_tokens, num_experts = scores.shape - scores = scores.view(num_tokens, num_groups, -1) - mask = torch.zeros((num_tokens, num_groups), dtype=torch.bool, device=scores.device) - mask = mask.scatter_(1, group_idx, True).unsqueeze(-1).expand_as(scores) - return (scores * mask).view(num_tokens, num_experts) - - -def bench(fn, num_warmups: int = 20, num_tests: int = 30, post_fn=None): - # Flush L2 cache with 256 MB data - torch.cuda.synchronize() - cache = torch.empty(int(256e6 // 4), dtype=torch.int, device="cuda") - - # Warmup - for _ in range(num_warmups): - fn() - - # Flush L2 - cache.zero_() - - # Testing - start_events = [torch.cuda.Event(enable_timing=True) for _ in range(num_tests)] - end_events = [torch.cuda.Event(enable_timing=True) for _ in range(num_tests)] - for i in range(num_tests): - # Record - start_events[i].record() - fn() - end_events[i].record() - if post_fn is not None: - post_fn() - torch.cuda.synchronize() - - times = np.array( - [s.elapsed_time(e) / 1e3 for s, e in zip(start_events, end_events)] - )[1:] - return np.average(times), np.min(times), np.max(times) - - -class empty_suppress: - def __enter__(self): - return self - - def __exit__(self, *_): - pass - - -class suppress_stdout_stderr: - def __enter__(self): - self.outnull_file = open(os.devnull, "w") - self.errnull_file = open(os.devnull, "w") - - self.old_stdout_fileno_undup = sys.stdout.fileno() - self.old_stderr_fileno_undup = sys.stderr.fileno() - - self.old_stdout_fileno = os.dup(sys.stdout.fileno()) - self.old_stderr_fileno = os.dup(sys.stderr.fileno()) - - self.old_stdout = sys.stdout - self.old_stderr = sys.stderr - - os.dup2(self.outnull_file.fileno(), self.old_stdout_fileno_undup) - os.dup2(self.errnull_file.fileno(), self.old_stderr_fileno_undup) - - sys.stdout = self.outnull_file - sys.stderr = self.errnull_file - return self - - def __exit__(self, *_): - sys.stdout = self.old_stdout - sys.stderr = self.old_stderr - - os.dup2(self.old_stdout_fileno, self.old_stdout_fileno_undup) - os.dup2(self.old_stderr_fileno, self.old_stderr_fileno_undup) - - os.close(self.old_stdout_fileno) - os.close(self.old_stderr_fileno) - - self.outnull_file.close() - self.errnull_file.close() - - -def bench_kineto( - fn, - kernel_names, - num_tests: int = 30, - suppress_kineto_output: bool = False, - trace_path: Optional[str] = None, - barrier_comm_profiling: bool = False, -): - # Profile - suppress = suppress_stdout_stderr if suppress_kineto_output else empty_suppress - with suppress(): - schedule = torch.profiler.schedule(wait=0, warmup=1, active=1, repeat=1) - with torch.profiler.profile( - activities=[torch.profiler.ProfilerActivity.CUDA], schedule=schedule - ) as prof: - for i in range(2): - # NOTES: use a large kernel and a barrier to eliminate the unbalanced CPU launch overhead - if barrier_comm_profiling: - lhs = torch.randn((8192, 8192), dtype=torch.float, device="cuda") - rhs = torch.randn((8192, 8192), dtype=torch.float, device="cuda") - lhs @ rhs - dist.all_reduce(torch.ones(1, dtype=torch.float, device="cuda")) - for _ in range(num_tests): - fn() - prof.step() - - # Parse the profiling table - assert isinstance(kernel_names, str) or isinstance(kernel_names, tuple) - is_tupled = isinstance(kernel_names, tuple) - prof_lines = ( - prof.key_averages() - .table(sort_by="cuda_time_total", max_name_column_width=100) - .split("\n") - ) - kernel_names = (kernel_names,) if isinstance(kernel_names, str) else kernel_names - assert all([isinstance(name, str) for name in kernel_names]) - for name in kernel_names: - assert sum([name in line for line in prof_lines]) == 1, ( - f"Errors of the kernel {name} in the profiling table" - ) - - # Save chrome traces - if trace_path is not None: - prof.export_chrome_trace(trace_path) - - # Return average kernel times - units = {"ms": 1e3, "us": 1e6} - kernel_times = [] - for name in kernel_names: - for line in prof_lines: - if name in line: - time_str = line.split()[-2] - for unit, scale in units.items(): - if unit in time_str: - kernel_times.append(float(time_str.replace(unit, "")) / scale) - break - break - return tuple(kernel_times) if is_tupled else kernel_times[0] - - -def hash_tensor(t: torch.Tensor): - return t.view(torch.int64).sum().item() diff --git a/benchmark/kernels/deepep/tuning_deepep.py b/benchmark/kernels/deepep/tuning_deepep.py index 6a1401408..97dd14504 100644 --- a/benchmark/kernels/deepep/tuning_deepep.py +++ b/benchmark/kernels/deepep/tuning_deepep.py @@ -16,17 +16,41 @@ from pathlib import Path import deep_ep import torch import torch.distributed as dist -from deepep_utils import ( + +from sglang.test.test_deepep_utils import ( bench, calc_diff, create_grouped_scores, - init_dist, inplace_unique, per_token_cast_back, per_token_cast_to_fp8, ) +def init_dist(local_rank: int, num_local_ranks: int, args): + ip = args.master_addr + port = args.master_port + num_nodes = args.nnodes + node_rank = args.node_rank + assert (num_local_ranks < 8 and num_nodes == 1) or num_local_ranks == 8 + + dist.init_process_group( + backend="nccl", + init_method=f"tcp://{ip}:{port}", + world_size=num_nodes * num_local_ranks, + rank=node_rank * num_local_ranks + local_rank, + ) + torch.set_default_dtype(torch.bfloat16) + torch.set_default_device("cuda") + torch.cuda.set_device(local_rank) + + return ( + dist.get_rank(), + dist.get_world_size(), + dist.new_group(list(range(num_local_ranks * num_nodes))), + ) + + def test_main( num_sms: int, local_rank: int, diff --git a/benchmark/kernels/deepseek/benchmark_deepgemm_fp8_gemm.py b/benchmark/kernels/deepseek/benchmark_deepgemm_fp8_gemm.py index bae6904e2..e994ea5df 100644 --- a/benchmark/kernels/deepseek/benchmark_deepgemm_fp8_gemm.py +++ b/benchmark/kernels/deepseek/benchmark_deepgemm_fp8_gemm.py @@ -1,17 +1,16 @@ -from typing import Tuple - import deep_gemm import tilelang import tilelang.language as T import torch import triton -from deep_gemm import ceil_div from deep_gemm.utils.layout import get_mn_major_tma_aligned_tensor -from vllm.model_executor.layers.quantization.utils.fp8_utils import ( - w8a8_block_fp8_matmul as vllm_w8a8_block_fp8_matmul, -) from sglang.benchmark.bench_utils import run_bench +from sglang.benchmark.deepseek_utils import ( + get_weight_shapes, + per_block_cast_to_fp8, + per_token_cast_to_fp8, +) from sglang.kernels.ops.quantization.fp8_kernel import ( w8a8_block_fp8_matmul_deepgemm as w8a8_block_fp8_matmul, ) @@ -94,31 +93,6 @@ def tl_gemm( return main -def per_token_cast_to_fp8(x: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]: - assert x.dim() == 2 and x.size(1) % 128 == 0 - m, n = x.shape - x_view = x.view(m, -1, 128) - x_amax = x_view.abs().float().amax(dim=2).view(m, -1).clamp(1e-4) - return (x_view * (448.0 / x_amax.unsqueeze(2))).to(torch.float8_e4m3fn).view( - m, n - ), (x_amax / 448.0).view(m, -1) - - -def per_block_cast_to_fp8(x: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]: - assert x.dim() == 2 - m, n = x.shape - x_padded = torch.zeros( - (ceil_div(m, 128) * 128, ceil_div(n, 128) * 128), dtype=x.dtype, device=x.device - ) - x_padded[:m, :n] = x - x_view = x_padded.view(-1, 128, x_padded.size(1) // 128, 128) - x_amax = x_view.abs().float().amax(dim=(1, 3), keepdim=True).clamp(1e-4) - x_scaled = (x_view * (448.0 / x_amax)).to(torch.float8_e4m3fn) - return x_scaled.view_as(x_padded)[:m, :n].contiguous(), (x_amax / 448.0).view( - x_view.size(0), x_view.size(2) - ) - - def fp8_gemm_deepgemm( x_fp8: torch.Tensor, x_scale: torch.Tensor, @@ -155,25 +129,6 @@ def fp8_gemm_sglang( return out -def fp8_gemm_vllm( - x_fp8: torch.Tensor, - x_scale: torch.Tensor, - y_fp8: torch.Tensor, - y_scale: torch.Tensor, - m: int, - n: int, - k: int, -): - """vLLM implementation of FP8 GEMM""" - block_size = [128, 128] # Matches the block size in per_block_cast_to_fp8 - - # Run vLLM kernel - out = vllm_w8a8_block_fp8_matmul( - x_fp8, y_fp8, x_scale, y_scale, block_size, torch.bfloat16 - ) - return out - - def calculate_diff(m: int, n: int, k: int): x = torch.randn((m, k), device="cuda", dtype=torch.bfloat16) y = torch.randn((n, k), device="cuda", dtype=torch.bfloat16) @@ -232,39 +187,6 @@ def calculate_diff(m: int, n: int, k: int): print(f" - TileLang vs SGLang: {'✅' if tilelang_sglang_match else '❌'}\n") -def get_weight_shapes(tp_size): - # cannot TP - total = [ - (512 + 64, 7168), - ((128 + 64) * 128, 7168), - (128 * (128 + 128), 512), - (7168, 16384), - (7168, 18432), - ] - # N can TP - n_tp = [ - (18432 * 2, 7168), - ((128 + 64) * 128, 7168), - (128 * (128 + 128), 512), - (24576, 1536), - (4096, 7168), - ] - # K can TP - k_tp = [(7168, 18432), (7168, 16384), (7168, 2048)] - - weight_shapes = [] - for t in total: - weight_shapes.append(t) - for n_t in n_tp: - new_t = (n_t[0] // tp_size, n_t[1]) - weight_shapes.append(new_t) - for k_t in k_tp: - new_t = (k_t[0], k_t[1] // tp_size) - weight_shapes.append(new_t) - - return weight_shapes - - def create_benchmark_configs(tp_size): configs = [] weight_shapes = get_weight_shapes(tp_size) diff --git a/benchmark/kernels/deepseek/benchmark_deepgemm_fp8_gemm_blackwell.py b/benchmark/kernels/deepseek/benchmark_deepgemm_fp8_gemm_blackwell.py index 280f6b97d..8fae14b19 100644 --- a/benchmark/kernels/deepseek/benchmark_deepgemm_fp8_gemm_blackwell.py +++ b/benchmark/kernels/deepseek/benchmark_deepgemm_fp8_gemm_blackwell.py @@ -1,12 +1,14 @@ import argparse -from typing import Tuple import torch import triton -from deep_gemm import ceil_div from flashinfer.gemm import gemm_fp8_nt_groupwise from sglang.benchmark.bench_utils import run_bench +from sglang.benchmark.deepseek_utils import ( + get_weight_shapes, + per_block_cast_to_fp8, +) from sglang.kernels.ops.quantization.fp8_kernel import ( sglang_per_token_group_quant_fp8, w8a8_block_fp8_matmul_deepgemm, @@ -16,55 +18,6 @@ from sglang.srt.layers.quantization.fp8_utils import requant_weight_ue8m0 BLOCK_SIZE = 128 -def per_block_cast_to_fp8(x: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]: - assert x.dim() == 2 - assert BLOCK_SIZE == 128 - m, n = x.shape - x_padded = torch.zeros( - (ceil_div(m, 128) * 128, ceil_div(n, 128) * 128), dtype=x.dtype, device=x.device - ) - x_padded[:m, :n] = x - x_view = x_padded.view(-1, 128, x_padded.size(1) // 128, 128) - x_amax = x_view.abs().float().amax(dim=(1, 3), keepdim=True).clamp(1e-4) - x_scaled = (x_view * (448.0 / x_amax)).to(torch.float8_e4m3fn) - return x_scaled.view_as(x_padded)[:m, :n].contiguous(), (x_amax / 448.0).view( - x_view.size(0), x_view.size(2) - ) - - -def get_weight_shapes(tp_size): - # cannot TP - total = [ - (512 + 64, 7168), - ((128 + 64) * 128, 7168), - (128 * (128 + 128), 512), - (7168, 16384), - (7168, 18432), - ] - # N can TP - n_tp = [ - (18432 * 2, 7168), - ((128 + 64) * 128, 7168), - (128 * (128 + 128), 512), - (24576, 1536), - (4096, 7168), - ] - # K can TP - k_tp = [(7168, 18432), (7168, 16384), (7168, 2048)] - - weight_shapes = [] - for t in total: - weight_shapes.append(t) - for n_t in n_tp: - new_t = (n_t[0] // tp_size, n_t[1]) - weight_shapes.append(new_t) - for k_t in k_tp: - new_t = (k_t[0], k_t[1] // tp_size) - weight_shapes.append(new_t) - - return weight_shapes - - def create_benchmark_configs(tp_size): configs = [] weight_shapes = get_weight_shapes(tp_size) diff --git a/benchmark/kernels/deepseek/benchmark_deepgemm_fp8_group_gemm.py b/benchmark/kernels/deepseek/benchmark_deepgemm_fp8_group_gemm.py index 1aace8087..4e2eaa864 100644 --- a/benchmark/kernels/deepseek/benchmark_deepgemm_fp8_group_gemm.py +++ b/benchmark/kernels/deepseek/benchmark_deepgemm_fp8_group_gemm.py @@ -4,12 +4,12 @@ import deep_gemm import torch import triton import triton.language as tl -from deep_gemm import calc_diff +from deep_gemm.testing import calc_diff from deep_gemm.utils.layout import get_mn_major_tma_aligned_tensor -# Import shared functionality from the regular GEMM benchmark from sglang.benchmark.bench_utils import run_bench -from sglang.benchmark.kernels.deepseek.benchmark_deepgemm_fp8_gemm import ( +from sglang.benchmark.deepseek_utils import ( + get_weight_shapes, per_block_cast_to_fp8, per_token_cast_to_fp8, ) @@ -318,39 +318,6 @@ def calculate_diff(m: int, n: int, k: int, num_groups: int): ) -def get_weight_shapes(tp_size): - # cannot TP - total = [ - (512 + 64, 7168), - ((128 + 64) * 128, 7168), - (128 * (128 + 128), 512), - (7168, 16384), - (7168, 18432), - ] - # N can TP - n_tp = [ - (18432 * 2, 7168), - ((128 + 64) * 128, 7168), - (128 * (128 + 128), 512), - (24576, 1536), - (4096, 7168), - ] - # K can TP - k_tp = [(7168, 18432), (7168, 16384), (7168, 2048)] - - weight_shapes = [] - for t in total: - weight_shapes.append(t) - for n_t in n_tp: - new_t = (n_t[0] // tp_size, n_t[1]) - weight_shapes.append(new_t) - for k_t in k_tp: - new_t = (k_t[0], k_t[1] // tp_size) - weight_shapes.append(new_t) - - return weight_shapes - - def create_benchmark_configs(tp_size): configs = [] weight_shapes = get_weight_shapes(tp_size) diff --git a/benchmark/kernels/quantization/tuning_block_wise_kernel.py b/benchmark/kernels/quantization/tuning_block_wise_kernel.py index 1c232fa38..f2c4cca0a 100644 --- a/benchmark/kernels/quantization/tuning_block_wise_kernel.py +++ b/benchmark/kernels/quantization/tuning_block_wise_kernel.py @@ -27,6 +27,7 @@ from tqdm import tqdm mp.set_start_method("spawn", force=True) +from sglang.benchmark.deepseek_utils import get_weight_shapes from sglang.kernels.ops.quantization.fp8_kernel import ( _w8a8_block_fp8_matmul, _w8a8_block_fp8_matmul_unrolledx4, @@ -192,39 +193,6 @@ def get_configs_compute_bound(): return configs -def get_weight_shapes(tp_size): - # NOTE(HandH1998): The weight shapes only works for DeepSeek-V3. Modify them, if you tune for another different model. - # cannot TP - total = [ - (512 + 64, 7168), - ((128 + 64) * 128, 7168), - (128 * (128 + 128), 512), - (7168, 16384), - (7168, 18432), - ] - # N can TP - n_tp = [ - (18432 * 2, 7168), - ((128 + 64) * 128, 7168), - (128 * (128 + 128), 512), - (24576, 1536), - (4096, 7168), - ] - # K can TP - k_tp = [(7168, 18432), (7168, 16384), (7168, 2048)] - - weight_shapes = [] - for t in total: - weight_shapes.append(t) - for n_t in n_tp: - new_t = (n_t[0] // tp_size, n_t[1]) - weight_shapes.append(new_t) - for k_t in k_tp: - new_t = (k_t[0], k_t[1] // tp_size) - weight_shapes.append(new_t) - return weight_shapes - - def benchmark_config( A, B, As, Bs, block_size, config, out_dtype=torch.float16, num_iters=10 ): diff --git a/benchmark/mmmu/data_utils.py b/benchmark/mmmu/data_utils.py index 8c36768d0..aca202969 100644 --- a/benchmark/mmmu/data_utils.py +++ b/benchmark/mmmu/data_utils.py @@ -75,22 +75,6 @@ CAT_SHORT2LONG = { } -def get_multi_choice_info(options): - """ - Given the list of options for multiple choice question - Return the index2ans and all_choices - """ - - start_chr = "A" - all_choices = [] - index2ans = {} - for i, option in enumerate(options): - index2ans[chr(ord(start_chr) + i)] = option - all_choices.append(chr(ord(start_chr) + i)) - - return index2ans, all_choices - - def load_yaml(file_path): with open(file_path, "r") as stream: try: @@ -142,33 +126,6 @@ def save_json(filename, ds): json.dump(ds, f, indent=4) -def save_jsonl(filename, data): - """ - Save a dictionary of data to a JSON Lines file with the filename as key and caption as value. - - Args: - filename (str): The path to the file where the data should be saved. - data (dict): The dictionary containing the data to save where key is the image path and value is the caption. - """ - with open(filename, "w", encoding="utf-8") as f: - for img_path, caption in data.items(): - # Extract the base filename without the extension - base_filename = os.path.basename(img_path) - # Create a JSON object with the filename as the key and caption as the value - json_record = json.dumps({base_filename: caption}, ensure_ascii=False) - # Write the JSON object to the file, one per line - f.write(json_record + "\n") - - -def save_args(args, path_dir): - argsDict = args.__dict__ - with open(path_dir + "setting.txt", "w") as f: - f.writelines("------------------ start ------------------" + "\n") - for eachArg, value in argsDict.items(): - f.writelines(eachArg + " : " + str(value) + "\n") - f.writelines("------------------- end -------------------") - - # DATA PROCESSING def construct_prompt(sample, config): question = sample["question"] diff --git a/python/sglang/benchmark/deepseek_utils.py b/python/sglang/benchmark/deepseek_utils.py new file mode 100644 index 000000000..6bbcab5d0 --- /dev/null +++ b/python/sglang/benchmark/deepseek_utils.py @@ -0,0 +1,65 @@ +"""Shared DeepSeek-V3 benchmark shapes and FP8 input preparation.""" + +from typing import Tuple + +import torch +from triton import cdiv + + +def per_token_cast_to_fp8(x: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]: + assert x.dim() == 2 and x.size(1) % 128 == 0 + m, n = x.shape + x_view = x.view(m, -1, 128) + x_amax = x_view.abs().float().amax(dim=2).view(m, -1).clamp(1e-4) + return (x_view * (448.0 / x_amax.unsqueeze(2))).to(torch.float8_e4m3fn).view( + m, n + ), (x_amax / 448.0).view(m, -1) + + +def per_block_cast_to_fp8(x: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]: + assert x.dim() == 2 + m, n = x.shape + x_padded = torch.zeros( + (cdiv(m, 128) * 128, cdiv(n, 128) * 128), dtype=x.dtype, device=x.device + ) + x_padded[:m, :n] = x + x_view = x_padded.view(-1, 128, x_padded.size(1) // 128, 128) + x_amax = x_view.abs().float().amax(dim=(1, 3), keepdim=True).clamp(1e-4) + x_scaled = (x_view * (448.0 / x_amax)).to(torch.float8_e4m3fn) + return x_scaled.view_as(x_padded)[:m, :n].contiguous(), (x_amax / 448.0).view( + x_view.size(0), x_view.size(2) + ) + + +def get_weight_shapes(tp_size): + """Return the DeepSeek-V3 (N, K) shapes, including TP-sharded projections.""" + # cannot TP + total = [ + (512 + 64, 7168), + ((128 + 64) * 128, 7168), + (128 * (128 + 128), 512), + (7168, 16384), + (7168, 18432), + ] + # N can TP + n_tp = [ + (18432 * 2, 7168), + ((128 + 64) * 128, 7168), + (128 * (128 + 128), 512), + (24576, 1536), + (4096, 7168), + ] + # K can TP + k_tp = [(7168, 18432), (7168, 16384), (7168, 2048)] + + weight_shapes = [] + for t in total: + weight_shapes.append(t) + for n_t in n_tp: + new_t = (n_t[0] // tp_size, n_t[1]) + weight_shapes.append(new_t) + for k_t in k_tp: + new_t = (k_t[0], k_t[1] // tp_size) + weight_shapes.append(new_t) + + return weight_shapes diff --git a/python/sglang/multimodal_gen/test/unit/sana_wm/test_realtime_chain.py b/python/sglang/multimodal_gen/test/unit/sana_wm/test_realtime_chain.py index 66a3ec82a..4a1a5c08b 100644 --- a/python/sglang/multimodal_gen/test/unit/sana_wm/test_realtime_chain.py +++ b/python/sglang/multimodal_gen/test/unit/sana_wm/test_realtime_chain.py @@ -15,7 +15,6 @@ import pytest import torch from PIL import Image -from sglang.multimodal_gen.runtime import server_args as _sa_mod from sglang.multimodal_gen.runtime.pipelines_core.stages.model_specific_stages.sana_wm.realtime_chain import ( SanaWMCameraCondStage, SanaWMNoiseState, @@ -29,7 +28,6 @@ from sglang.multimodal_gen.runtime.realtime.session import RealtimeSession from sglang.multimodal_gen.runtime.realtime.states import ( get_realtime_causal_dit_state, ) -from sglang.multimodal_gen.runtime.server_args import set_global_server_args MC = 8 @@ -39,23 +37,6 @@ class _TestRealtimeStage(SanaWMRealtimeStage): raise NotImplementedError -@pytest.fixture -def _global_args(): - prev = _sa_mod._global_server_args - set_global_server_args( - SimpleNamespace( - comfyui_mode=False, - enable_cfg_parallel=False, - enable_torch_compile=False, - attention_backend=None, - ) - ) - try: - yield - finally: - set_global_server_args(prev) - - def _prep_stage(): return SanaWMRealtimeLatentPrepStage( use_refiner=True, transformer=None, vae=None, model_path="" @@ -153,7 +134,7 @@ def test_realtime_camera_conditioning_uses_requested_size(): assert plucker.shape == (1, 48, 3, 15, 26) -def test_latent_prep_plan_and_noise_discipline(_global_args): +def test_latent_prep_plan_and_noise_discipline(): stage = _prep_stage() session = RealtimeSession() fl = torch.ones(1, MC, 1, 2, 2, dtype=torch.float32) @@ -193,7 +174,7 @@ def test_latent_prep_plan_and_noise_discipline(_global_args): assert torch.isfinite(batch.latents).all() -def test_latent_prep_open_ended_uniform_chunk0(_global_args): +def test_latent_prep_open_ended_uniform_chunk0(): stage = _prep_stage() session = RealtimeSession() fl = torch.ones(1, MC, 1, 2, 2, dtype=torch.float32) diff --git a/python/sglang/multimodal_gen/test/unit/sana_wm/test_streaming_forward_long.py b/python/sglang/multimodal_gen/test/unit/sana_wm/test_streaming_forward_long.py index 78c04201a..280d080ef 100644 --- a/python/sglang/multimodal_gen/test/unit/sana_wm/test_streaming_forward_long.py +++ b/python/sglang/multimodal_gen/test/unit/sana_wm/test_streaming_forward_long.py @@ -12,14 +12,12 @@ from __future__ import annotations from types import SimpleNamespace -import pytest import torch from sglang.multimodal_gen.configs.models.dits.sana_wm import ( SanaWMArchConfig, SanaWMConfig, ) -from sglang.multimodal_gen.runtime import server_args as _sa_mod from sglang.multimodal_gen.runtime.models.dits.sana_wm import ( _CACHE_TYPE_STATE, _SLOT_CAM_K, @@ -35,7 +33,6 @@ from sglang.multimodal_gen.runtime.models.dits.sana_wm import ( _slice_rope_to_current_chunk, process_camera_conditions_ucpe, ) -from sglang.multimodal_gen.runtime.server_args import set_global_server_args HEAD_DIM = 112 H, W = 2, 3 @@ -236,23 +233,6 @@ def test_forward_long_gdn_reduces_to_dense_with_camera(): # --------------------------------------------------------------------------- # -@pytest.fixture -def _global_args(): - prev = _sa_mod._global_server_args - set_global_server_args( - SimpleNamespace( - comfyui_mode=False, - enable_cfg_parallel=False, - enable_torch_compile=False, - attention_backend=None, - ) - ) - try: - yield - finally: - set_global_server_args(prev) - - class _ZeroCross(torch.nn.Module): def forward(self, x, y, mask=None): return torch.zeros_like(x) @@ -284,7 +264,7 @@ def _block(): return b -def test_block_forward_long_reduces_to_dense(_global_args): +def test_block_forward_long_reduces_to_dense(): block = _block() x = _x() y = torch.randn(AB, 4, AC, dtype=torch.float64) @@ -349,7 +329,7 @@ def _model_inputs(): ) -def test_model_forward_long_single_chunk_reduces_to_dense(_global_args): +def test_model_forward_long_single_chunk_reduces_to_dense(): m = _tiny_model() inp = _model_inputs() with torch.no_grad(): @@ -364,7 +344,7 @@ def test_model_forward_long_single_chunk_reduces_to_dense(_global_args): assert cache[0][_SLOT_FFN_TCONV] is not None -def test_model_forward_long_two_chunks_runs_and_windows(_global_args): +def test_model_forward_long_two_chunks_runs_and_windows(): m = _tiny_model() inp = _model_inputs() split = 2 diff --git a/python/sglang/multimodal_gen/test/unit/sana_wm/test_streaming_stage.py b/python/sglang/multimodal_gen/test/unit/sana_wm/test_streaming_stage.py index c2471df0a..d318eb3ca 100644 --- a/python/sglang/multimodal_gen/test/unit/sana_wm/test_streaming_stage.py +++ b/python/sglang/multimodal_gen/test/unit/sana_wm/test_streaming_stage.py @@ -9,16 +9,12 @@ dim=2 would concat the head axis and silently corrupt every softmax block). from __future__ import annotations -from types import SimpleNamespace - -import pytest import torch from sglang.multimodal_gen.configs.models.dits.sana_wm import ( SanaWMArchConfig, SanaWMConfig, ) -from sglang.multimodal_gen.runtime import server_args as _sa_mod from sglang.multimodal_gen.runtime.models.dits.sana_wm import ( _CACHE_TYPE_CONCAT, _CACHE_TYPE_STATE, @@ -35,7 +31,6 @@ from sglang.multimodal_gen.runtime.models.dits.sana_wm import ( from sglang.multimodal_gen.runtime.pipelines_core.stages.model_specific_stages.sana_wm.streaming import ( SanaWMStreamingDenoisingStage as Stage, ) -from sglang.multimodal_gen.runtime.server_args import set_global_server_args B, Hh, D = 1, 2, 4 @@ -160,23 +155,6 @@ class _ZeroCross(torch.nn.Module): return torch.zeros_like(x) -@pytest.fixture -def _global_args(): - prev = _sa_mod._global_server_args - set_global_server_args( - SimpleNamespace( - comfyui_mode=False, - enable_cfg_parallel=False, - enable_torch_compile=False, - attention_backend=None, - ) - ) - try: - yield - finally: - set_global_server_args(prev) - - def _depth4_model(): arch = SanaWMArchConfig( in_channels=MC, @@ -201,7 +179,7 @@ def _depth4_model(): return m -def test_streaming_loop_runs_and_accumulates_concat_block(_global_args): +def test_streaming_loop_runs_and_accumulates_concat_block(): """Drive forward_long chunk-by-chunk (the stage's core loop) on a depth-4 model: accumulate -> denoise(save=False) -> clean(save=True). Verify finite output, threaded GDN state, and that the softmax block (idx 3) accumulates a diff --git a/python/sglang/multimodal_gen/test/unit/test_fp32_layernorm.py b/python/sglang/multimodal_gen/test/unit/test_fp32_layernorm.py index ed4616e7d..b4ff004ac 100644 --- a/python/sglang/multimodal_gen/test/unit/test_fp32_layernorm.py +++ b/python/sglang/multimodal_gen/test/unit/test_fp32_layernorm.py @@ -4,61 +4,40 @@ import torch.nn.functional as F from sglang.multimodal_gen.runtime.layers.layernorm import FP32LayerNorm +pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") -@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") -def test_fp32_layernorm_cache_matches_reference(): + +def test_fp32_layernorm_cache_reuse_and_invalidation(): norm = FP32LayerNorm(16, eps=1e-5).cuda().to(torch.bfloat16) inputs = torch.randn(4, 16, device="cuda", dtype=torch.bfloat16) with torch.no_grad(): - actual = norm(inputs) - expected = F.layer_norm( - inputs.float(), - norm.normalized_shape, - norm.weight.float().to(device=inputs.device), - norm.bias.float().to(device=inputs.device), - norm.eps, - ).to(inputs.dtype) + for updated in (False, True): + if updated: + previous = norm.__dict__["_weight_fp32_cache"] + norm.weight.add_(1.0) + actual = norm(inputs) + expected = F.layer_norm( + inputs.float(), + norm.normalized_shape, + norm.weight.float(), + norm.bias.float(), + norm.eps, + ).to(inputs.dtype) + torch.testing.assert_close(actual, expected) + weight_cache = norm.__dict__["_weight_fp32_cache"] + bias_cache = norm.__dict__["_bias_fp32_cache"] + if updated: + assert weight_cache[0] != previous[0] + assert weight_cache[1] is not previous[1] + norm(inputs) + assert norm.__dict__["_weight_fp32_cache"][1] is weight_cache[1] + assert norm.__dict__["_bias_fp32_cache"][1] is bias_cache[1] - torch.testing.assert_close(actual, expected) - - -@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") -def test_fp32_layernorm_cache_reuses_converted_params(): - norm = FP32LayerNorm(16, eps=1e-5).cuda().to(torch.bfloat16) - inputs = torch.randn(4, 16, device="cuda", dtype=torch.bfloat16) - - with torch.no_grad(): - norm(inputs) - weight_cache = norm.__dict__["_weight_fp32_cache"] - bias_cache = norm.__dict__["_bias_fp32_cache"] - - norm(inputs) - - assert norm.__dict__["_weight_fp32_cache"][1] is weight_cache[1] - assert norm.__dict__["_bias_fp32_cache"][1] is bias_cache[1] assert "_weight_fp32_cache" not in norm.state_dict() assert "_bias_fp32_cache" not in norm.state_dict() -@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") -def test_fp32_layernorm_cache_invalidates_on_param_update(): - norm = FP32LayerNorm(16, eps=1e-5).cuda().to(torch.bfloat16) - inputs = torch.randn(4, 16, device="cuda", dtype=torch.bfloat16) - - with torch.no_grad(): - norm(inputs) - first_key, first_weight = norm.__dict__["_weight_fp32_cache"] - - norm.weight.add_(1.0) - norm(inputs) - second_key, second_weight = norm.__dict__["_weight_fp32_cache"] - - assert second_key != first_key - assert second_weight is not first_weight - - -@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") def test_fp32_layernorm_grad_mode_preserves_autograd_path(): norm = FP32LayerNorm(16, eps=1e-5).cuda().to(torch.bfloat16) inputs = torch.randn(4, 16, device="cuda", dtype=torch.bfloat16, requires_grad=True) diff --git a/python/sglang/multimodal_gen/test/unit/test_latent_upsampler_group_norm_silu.py b/python/sglang/multimodal_gen/test/unit/test_latent_upsampler_group_norm_silu.py index 88615fbc4..2a9700930 100644 --- a/python/sglang/multimodal_gen/test/unit/test_latent_upsampler_group_norm_silu.py +++ b/python/sglang/multimodal_gen/test/unit/test_latent_upsampler_group_norm_silu.py @@ -10,7 +10,6 @@ from sglang.kernels.ops.diffusion import apply_group_norm_silu from sglang.multimodal_gen.runtime.models.upsampler.latent_upsampler import ( LatentUpsampler, ResBlock, - SpatialRationalResampler, ) @@ -21,199 +20,87 @@ def _resblock_eager_reference(block: ResBlock, x: torch.Tensor) -> torch.Tensor: return block.activation(x + residual) -def _latent_upsampler_eager_reference( - upsampler: LatentUpsampler, latent: torch.Tensor -) -> torch.Tensor: - from einops import rearrange - - b, _, f, _, _ = latent.shape - if upsampler.dims == 2: - x = rearrange(latent, "b c f h w -> (b f) c h w") - x = upsampler.initial_activation( - upsampler.initial_norm(upsampler.initial_conv(x)) - ) - for block in upsampler.res_blocks: - x = _resblock_eager_reference(block, x) - x = upsampler.upsampler(x) - for block in upsampler.post_upsample_res_blocks: - x = _resblock_eager_reference(block, x) - x = upsampler.final_conv(x) - return rearrange(x, "(b f) c h w -> b c f h w", b=b, f=f) - - x = upsampler.initial_activation( - upsampler.initial_norm(upsampler.initial_conv(latent)) - ) - for block in upsampler.res_blocks: - x = _resblock_eager_reference(block, x) - if upsampler.temporal_upsample: - x = upsampler.upsampler(x)[:, :, 1:, :, :] - elif isinstance(upsampler.upsampler, SpatialRationalResampler): - x = upsampler.upsampler(x) - else: - x = rearrange(x, "b c f h w -> (b f) c h w") - x = upsampler.upsampler(x) - x = rearrange(x, "(b f) c h w -> b c f h w", b=b, f=f) - for block in upsampler.post_upsample_res_blocks: - x = _resblock_eager_reference(block, x) - return upsampler.final_conv(x) - - -@pytest.mark.parametrize( - "batch,channels,dims,spatial", - [ - (1, 64, 2, (16, 16)), - (2, 64, 2, (8, 24)), - (1, 128, 3, (2, 8, 8)), - ], -) -def test_resblock_forward_parity(batch, channels, dims, spatial): - torch.manual_seed(0) - block = ResBlock(channels=channels, dims=dims).eval() - torch.manual_seed(1) - x = torch.randn(batch, channels, *spatial, dtype=torch.float32) - - with torch.no_grad(): - out = block(x) - ref = _resblock_eager_reference(block, x) - - torch.testing.assert_close(out, ref, atol=0.0, rtol=0.0) - - -@pytest.mark.parametrize( - "dims,latent_shape,mid_channels,num_blocks_per_stage,rational_resampler", - [ - (2, (1, 32, 2, 16, 16), 64, 2, False), - (3, (1, 32, 2, 16, 16), 64, 2, False), - (3, (1, 32, 2, 16, 16), 64, 2, True), - ], -) -def test_latent_upsampler_forward_parity( - dims, latent_shape, mid_channels, num_blocks_per_stage, rational_resampler -): - torch.manual_seed(2) - upsampler = LatentUpsampler( - in_channels=latent_shape[1], - mid_channels=mid_channels, - num_blocks_per_stage=num_blocks_per_stage, - dims=dims, - spatial_upsample=True, - temporal_upsample=False, - spatial_scale=2.0, - rational_resampler=rational_resampler, - ).eval() - torch.manual_seed(3) - latent = torch.randn(*latent_shape, dtype=torch.float32) - - with torch.no_grad(): - out = upsampler(latent) - ref = _latent_upsampler_eager_reference(upsampler, latent) - - torch.testing.assert_close(out, ref, atol=0.0, rtol=0.0) - - -def test_resblock_fuses_exactly_one_site(): - torch.manual_seed(4) - block = ResBlock(channels=64, dims=2).eval() - x = torch.randn(1, 64, 16, 16, dtype=torch.float32) - +def _latent_upsampler_eager_reference(upsampler, latent): with patch.object( - lu_mod, "apply_group_norm_silu", wraps=lu_mod.apply_group_norm_silu - ) as spy: - with torch.no_grad(): - block(x) - assert spy.call_count == 1 + lu_mod, "apply_group_norm_silu", side_effect=lambda x, norm, act: act(norm(x)) + ): + return upsampler(latent) -@pytest.mark.parametrize( - "dims,num_blocks_per_stage,rational_resampler,expected_calls", - [ - (2, 2, False, 1 + 2 * 2), - (2, 4, False, 1 + 4 * 2), - (3, 2, False, 1 + 2 * 2), - (3, 2, True, 1 + 2 * 2), - ], -) -def test_latent_upsampler_fuses_expected_sites( - dims, num_blocks_per_stage, rational_resampler, expected_calls -): - torch.manual_seed(5) - upsampler = LatentUpsampler( - in_channels=32, - mid_channels=64, - num_blocks_per_stage=num_blocks_per_stage, - dims=dims, - spatial_upsample=True, - temporal_upsample=False, - spatial_scale=2.0, - rational_resampler=rational_resampler, - ).eval() - latent = torch.randn(1, 32, 2, 16, 16, dtype=torch.float32) - - with patch.object( - lu_mod, "apply_group_norm_silu", wraps=lu_mod.apply_group_norm_silu - ) as spy: - with torch.no_grad(): - upsampler(latent) - assert spy.call_count == expected_calls - - -# CUDA Triton fast path ------------------------------------------------------- - requires_cuda = pytest.mark.skipif( - not torch.cuda.is_available(), - reason="Triton fused group_norm_silu requires CUDA", + not torch.cuda.is_available(), reason="CUDA required" ) -# bf16 keeps kernel-level tolerance because its fp32-equivalent exponent range -# absorbs multi-layer conv drift; fp16 needs a looser tolerance on the e2e -# upsampler test where 8+ downstream convs amplify fused-vs-eager rounding. -_RESBLOCK_TOL = {torch.bfloat16: (7e-2, 2e-2), torch.float16: (3e-3, 3e-3)} -_UPSAMPLER_TOL = {torch.bfloat16: (7e-2, 2e-2), torch.float16: (2e-2, 1e-1)} + +def _parity_cases(common, cpu_only): + return [ + pytest.param( + device, dtype, *case, marks=requires_cuda if device == "cuda" else () + ) + for device, dtype in [ + ("cpu", torch.float32), + ("cuda", torch.bfloat16), + ("cuda", torch.float16), + ] + for case in common + (cpu_only if device == "cpu" else []) + ] + + +# Downstream convolutions amplify fp16 drift in the complete upsampler. +_RESBLOCK_TOL = { + torch.float32: (0, 0), + torch.bfloat16: (7e-2, 2e-2), + torch.float16: (3e-3, 3e-3), +} +_UPSAMPLER_TOL = { + torch.float32: (0, 0), + torch.bfloat16: (7e-2, 2e-2), + torch.float16: (2e-2, 1e-1), +} -@requires_cuda -@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float16]) @pytest.mark.parametrize( - "batch,channels,dims,spatial", - [ - (1, 64, 2, (16, 16)), - (1, 128, 3, (2, 8, 8)), - ], + "device,dtype,batch,channels,dims,spatial", + _parity_cases( + [(1, 64, 2, (16, 16)), (1, 128, 3, (2, 8, 8))], + cpu_only=[(2, 64, 2, (8, 24))], + ), ) -def test_resblock_forward_parity_cuda(dtype, batch, channels, dims, spatial): +def test_resblock_forward_parity(device, dtype, batch, channels, dims, spatial): torch.manual_seed(0) - device = torch.device("cuda") block = ResBlock(channels=channels, dims=dims).to(device=device, dtype=dtype).eval() torch.manual_seed(1) x = torch.randn(batch, channels, *spatial, device=device, dtype=dtype) - with torch.no_grad(): + with ( + torch.no_grad(), + patch.object( + lu_mod, "apply_group_norm_silu", wraps=lu_mod.apply_group_norm_silu + ) as fused, + ): out = block(x) + assert fused.call_count == 1 + with torch.no_grad(): ref = _resblock_eager_reference(block, x) - atol, rtol = _RESBLOCK_TOL[dtype] torch.testing.assert_close(out, ref, atol=atol, rtol=rtol) -@requires_cuda -@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float16]) @pytest.mark.parametrize( - "dims,latent_shape,mid_channels,num_blocks_per_stage,rational_resampler", - [ - (2, (1, 32, 2, 16, 16), 64, 2, False), - (3, (1, 32, 2, 16, 16), 64, 2, False), - (3, (1, 32, 2, 16, 16), 64, 2, True), - ], + "device,dtype,dims,num_blocks_per_stage,rational_resampler", + _parity_cases( + [(2, 2, False), (3, 2, False), (3, 2, True)], + cpu_only=[(2, 4, False)], + ), ) -def test_latent_upsampler_forward_parity_cuda( - dtype, dims, latent_shape, mid_channels, num_blocks_per_stage, rational_resampler +def test_latent_upsampler_forward_parity( + device, dtype, dims, num_blocks_per_stage, rational_resampler ): torch.manual_seed(2) - device = torch.device("cuda") upsampler = ( LatentUpsampler( - in_channels=latent_shape[1], - mid_channels=mid_channels, + in_channels=32, + mid_channels=64, num_blocks_per_stage=num_blocks_per_stage, dims=dims, spatial_upsample=True, @@ -225,12 +112,18 @@ def test_latent_upsampler_forward_parity_cuda( .eval() ) torch.manual_seed(3) - latent = torch.randn(*latent_shape, device=device, dtype=dtype) + latent = torch.randn(1, 32, 2, 16, 16, device=device, dtype=dtype) - with torch.no_grad(): + with ( + torch.no_grad(), + patch.object( + lu_mod, "apply_group_norm_silu", wraps=lu_mod.apply_group_norm_silu + ) as fused, + ): out = upsampler(latent) + assert fused.call_count == 1 + 2 * num_blocks_per_stage + with torch.no_grad(): ref = _latent_upsampler_eager_reference(upsampler, latent) - atol, rtol = _UPSAMPLER_TOL[dtype] torch.testing.assert_close(out, ref, atol=atol, rtol=rtol) diff --git a/python/sglang/multimodal_gen/test/unit/test_layernorm_cutedsl_dispatch.py b/python/sglang/multimodal_gen/test/unit/test_layernorm_cutedsl_dispatch.py index 948f630d5..479e9d2d9 100644 --- a/python/sglang/multimodal_gen/test/unit/test_layernorm_cutedsl_dispatch.py +++ b/python/sglang/multimodal_gen/test/unit/test_layernorm_cutedsl_dispatch.py @@ -14,41 +14,25 @@ _CUTEDSL_MODULE = "sglang.kernels.ops.diffusion.norm.scale_residual_norm_cutedsl @pytest.mark.parametrize("hidden_size", [257, 8448]) -def test_norm_scale_shift_cuda_falls_back_for_unsupported_hidden_size(hidden_size): - layer = RMSNormScaleShift(hidden_size) - x = torch.empty(1, 1, hidden_size) - shift = torch.empty(1, 1, hidden_size) - scale = torch.empty(1, 1, hidden_size) +@pytest.mark.parametrize( + "layer_cls,num_inputs", + [(RMSNormScaleShift, 3), (ScaleResidualRMSNormScaleShift, 5)], +) +def test_cuda_falls_back_for_unsupported_hidden_size( + hidden_size, layer_cls, num_inputs +): + layer = layer_cls(hidden_size) + inputs = [torch.empty(1, 1, hidden_size) for _ in range(num_inputs)] expected = object() with ( patch.object(layer, "forward_native", return_value=expected) as native, pytest.warns(UserWarning, match="native fallback"), ): - actual = layer.forward_cuda(x, shift, scale) + actual = layer.forward_cuda(*inputs) assert actual is expected - native.assert_called_once_with(x, shift, scale) - - -@pytest.mark.parametrize("hidden_size", [257, 8448]) -def test_scale_residual_cuda_falls_back_for_unsupported_hidden_size(hidden_size): - layer = ScaleResidualRMSNormScaleShift(hidden_size) - residual = torch.empty(1, 1, hidden_size) - x = torch.empty(1, 1, hidden_size) - gate = torch.empty(1, 1, hidden_size) - shift = torch.empty(1, 1, hidden_size) - scale = torch.empty(1, 1, hidden_size) - expected = object() - - with ( - patch.object(layer, "forward_native", return_value=expected) as native, - pytest.warns(UserWarning, match="native fallback"), - ): - actual = layer.forward_cuda(residual, x, gate, shift, scale) - - assert actual is expected - native.assert_called_once_with(residual, x, gate, shift, scale) + native.assert_called_once_with(*inputs) def test_norm_scale_shift_cuda_uses_cutedsl_for_supported_hidden_size(monkeypatch): diff --git a/python/sglang/multimodal_gen/test/unit/test_vae_loader.py b/python/sglang/multimodal_gen/test/unit/test_vae_loader.py index 889c6e83c..897717535 100644 --- a/python/sglang/multimodal_gen/test/unit/test_vae_loader.py +++ b/python/sglang/multimodal_gen/test/unit/test_vae_loader.py @@ -647,59 +647,26 @@ class TestVAELoader(unittest.TestCase): self.assertNotIn("latents_mean", loaded) self.assertNotIn("latents_std", loaded) - def test_channels_last_3d_defaults_true_for_qwen_image_on_cuda(self): + def test_channels_last_3d_cuda_model_defaults(self): + cases = [ + (QwenImagePipelineConfig, 1, "vae", True), + (WanT2V480PConfig, 1, "video_vae", True), + (FastWan2_2_TI2V_5B_Config, 1, "video_vae", True), + (Wan2_2_I2V_A14B_Config, 2, "video_vae", False), + (LTX2PipelineConfig, 1, "video_vae", True), + (LTX2PipelineConfig, 2, "video_vae", False), + ] with ( patch.dict("os.environ", {}, clear=True), patch.object(vae_loader.current_platform, "is_cuda", return_value=True), patch.object(vae_loader.current_platform, "is_rocm", return_value=False), ): - server_args = _FakeServerArgs(QwenImagePipelineConfig()) - self.assertTrue(_should_use_channels_last_3d(server_args, "vae")) - - def test_channels_last_3d_defaults_true_for_single_gpu_wan_on_cuda(self): - with ( - patch.dict("os.environ", {}, clear=True), - patch.object(vae_loader.current_platform, "is_cuda", return_value=True), - patch.object(vae_loader.current_platform, "is_rocm", return_value=False), - ): - server_args = _FakeServerArgs(WanT2V480PConfig(), num_gpus=1) - self.assertTrue(_should_use_channels_last_3d(server_args, "video_vae")) - - def test_channels_last_3d_defaults_true_for_single_gpu_fast_wan_on_cuda(self): - with ( - patch.dict("os.environ", {}, clear=True), - patch.object(vae_loader.current_platform, "is_cuda", return_value=True), - patch.object(vae_loader.current_platform, "is_rocm", return_value=False), - ): - server_args = _FakeServerArgs(FastWan2_2_TI2V_5B_Config(), num_gpus=1) - self.assertTrue(_should_use_channels_last_3d(server_args, "video_vae")) - - def test_channels_last_3d_defaults_false_for_multi_gpu_wan_on_cuda(self): - with ( - patch.dict("os.environ", {}, clear=True), - patch.object(vae_loader.current_platform, "is_cuda", return_value=True), - patch.object(vae_loader.current_platform, "is_rocm", return_value=False), - ): - server_args = _FakeServerArgs(Wan2_2_I2V_A14B_Config(), num_gpus=2) - self.assertFalse(_should_use_channels_last_3d(server_args, "video_vae")) - - def test_channels_last_3d_defaults_true_for_single_gpu_ltx_on_cuda(self): - with ( - patch.dict("os.environ", {}, clear=True), - patch.object(vae_loader.current_platform, "is_cuda", return_value=True), - patch.object(vae_loader.current_platform, "is_rocm", return_value=False), - ): - server_args = _FakeServerArgs(LTX2PipelineConfig(), num_gpus=1) - self.assertTrue(_should_use_channels_last_3d(server_args, "video_vae")) - - def test_channels_last_3d_defaults_false_for_multi_gpu_ltx_on_cuda(self): - with ( - patch.dict("os.environ", {}, clear=True), - patch.object(vae_loader.current_platform, "is_cuda", return_value=True), - patch.object(vae_loader.current_platform, "is_rocm", return_value=False), - ): - server_args = _FakeServerArgs(LTX2PipelineConfig(), num_gpus=2) - self.assertFalse(_should_use_channels_last_3d(server_args, "video_vae")) + for config_cls, num_gpus, component, expected in cases: + with self.subTest(config=config_cls.__name__, num_gpus=num_gpus): + server_args = _FakeServerArgs(config_cls(), num_gpus=num_gpus) + self.assertEqual( + _should_use_channels_last_3d(server_args, component), expected + ) def test_channels_last_3d_can_be_disabled_by_env(self): with ( diff --git a/python/sglang/test/kernels/paged_mqa.py b/python/sglang/test/kernels/paged_mqa.py new file mode 100644 index 000000000..a521280b6 --- /dev/null +++ b/python/sglang/test/kernels/paged_mqa.py @@ -0,0 +1,140 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + + +"""Shared FP8 paged MQA reference and inputs for the CuTe DSL and DeepGEMM tests.""" + +import torch + +from sglang.srt.layers.attention.dsa.utils import ( + fp8_mqa_logits_ceil_to_ue8m0, + fp8_mqa_logits_make_fused_kv, +) + +BLOCK_KV = 64 +HEAD_DIM = 128 + + +def ref_fp8_paged_mqa_logits( + q_fp8, + kv_fp8, + kv_scales, + weights, + context_lens, + block_table, + max_model_len, + block_kv, +): + B, next_n, H, D = q_fp8.shape + device = q_fp8.device + + logits = torch.full( + (B * next_n, max_model_len), float("-inf"), device=device, dtype=torch.float32 + ) + q_f32 = q_fp8.float() + + for b in range(B): + ctx_len = context_lens[b].item() + q_positions = torch.arange(ctx_len - next_n, ctx_len, device=device) + w = weights[b * next_n : (b + 1) * next_n, :] + + for blk_idx in range((ctx_len + block_kv - 1) // block_kv): + phys_blk = block_table[b, blk_idx].item() + k_f32 = kv_fp8[phys_blk].float() + scales = kv_scales[phys_blk] + + k_positions = torch.arange( + blk_idx * block_kv, (blk_idx + 1) * block_kv, device=device + ) + mask = (k_positions[None, :] < ctx_len) & ( + k_positions[None, :] <= q_positions[:, None] + ) + + qk = torch.matmul(q_f32[b].permute(1, 0, 2), k_f32.T) + qk = torch.where(mask[None, :, :], qk, torch.zeros(1, device=device)) + qk = torch.relu(qk) + + weighted = (w.T[:, :, None] * qk).sum(dim=0) + weighted = weighted * scales[None, :] + + start_pos = blk_idx * block_kv + end_pos = start_pos + block_kv + logits[b * next_n : (b + 1) * next_n, start_pos:end_pos] = torch.where( + mask, + weighted, + torch.tensor(float("-inf"), device=device, dtype=torch.float32), + ) + + return logits + + +def generate_paged_mqa_test_data( + batch_size, + next_n, + num_heads, + avg_context_len, + max_model_len, + device="cuda", +): + torch.manual_seed(42) + torch.cuda.manual_seed(42) + context_lens = torch.randint( + max(BLOCK_KV, int(0.7 * avg_context_len)), + int(1.3 * avg_context_len) + 1, + (batch_size,), + dtype=torch.int32, + device="cpu", + ).clamp(max=max_model_len) + + max_blocks_per_seq = (max_model_len + BLOCK_KV - 1) // BLOCK_KV + total_blocks = ((context_lens + BLOCK_KV - 1) // BLOCK_KV).sum().item() + num_phys_blocks = total_blocks + batch_size * 2 + + block_table = torch.full( + (batch_size, max_blocks_per_seq), 0, dtype=torch.int32, device=device + ) + blk_offset = 0 + for i in range(batch_size): + n_blks = (context_lens[i].item() + BLOCK_KV - 1) // BLOCK_KV + block_table[i, :n_blks] = torch.arange( + blk_offset, blk_offset + n_blks, dtype=torch.int32, device=device + ) + blk_offset += n_blks + + q_bf16 = torch.randn(batch_size, next_n, num_heads, HEAD_DIM, device=device) + q_fp8 = q_bf16.to(torch.float8_e4m3fn) + + kv_bf16 = torch.randn(num_phys_blocks, BLOCK_KV, HEAD_DIM, device=device) + kv_amax = kv_bf16.abs().float().amax(dim=-1, keepdim=True).clamp(1e-4) + kv_scale = fp8_mqa_logits_ceil_to_ue8m0(kv_amax / 448.0).squeeze(-1) + kv_fp8 = (kv_bf16 / kv_scale.unsqueeze(-1)).to(torch.float8_e4m3fn) + + weights = torch.randn( + batch_size * next_n, num_heads, device=device, dtype=torch.float32 + ) + kv_fused = fp8_mqa_logits_make_fused_kv(kv_fp8, kv_scale, BLOCK_KV, HEAD_DIM) + + return { + "q_fp8": q_fp8, + "kv_fp8": kv_fp8, + "kv_scales": kv_scale, + "kv_fused": kv_fused, + "weights": weights, + "context_lens": context_lens.to(device), + "block_table": block_table, + } + + +def assert_paged_mqa_matches_ref( + logits, ref_logits, context_lens, B, next_n, max_model_len +): + device = logits.device + positions = torch.arange(max_model_len, device=device).unsqueeze(0) + row_indices = torch.arange(B * next_n, device=device) // next_n + next_n_offset = torch.arange(B * next_n, device=device) % next_n + end_pos = context_lens[row_indices] - next_n + next_n_offset + mask = positions <= end_pos.unsqueeze(1) + + logits_masked = logits.float().masked_fill(~mask, 0) + ref_masked = ref_logits.float().masked_fill(~mask, 0) + torch.testing.assert_close(logits_masked, ref_masked, atol=5e-5, rtol=1e-5) diff --git a/python/sglang/test/test_marlin_utils.py b/python/sglang/test/test_marlin_utils.py index 6bf4693a0..b87958775 100644 --- a/python/sglang/test/test_marlin_utils.py +++ b/python/sglang/test/test_marlin_utils.py @@ -21,11 +21,33 @@ from sglang.srt.layers.quantization.marlin_utils import ( from sglang.srt.layers.quantization.utils import ( get_pack_factor, gptq_quantize_weights, + pack_cols, quantize_weights, sort_weights, ) +def awq_pack( + q_w: torch.Tensor, + num_bits: int, + size_k: int, + size_n: int, +): + assert q_w.shape == (size_k, size_n) + + if num_bits == 4: + interleave = np.array([0, 2, 4, 6, 1, 3, 5, 7]) + elif num_bits == 8: + interleave = np.array([0, 2, 1, 3]) + else: + raise Exception("num_bits must be 4 or 8, got {}".format(num_bits)) + + q_w = q_w.reshape((-1, len(interleave)))[:, interleave].ravel() + q_w = q_w.reshape((-1, size_n)).contiguous() + + return pack_cols(q_w, num_bits, size_k, size_n) + + class MarlinWorkspace: def __init__(self, out_features, min_thread_n, max_parallel): assert out_features % min_thread_n == 0, ( diff --git a/test/registered/kernels/ops/attention/test_cute_dsl_fp8_paged_mqa_logits.py b/test/registered/kernels/ops/attention/test_cute_dsl_fp8_paged_mqa_logits.py index 112581f91..ee6f3f741 100644 --- a/test/registered/kernels/ops/attention/test_cute_dsl_fp8_paged_mqa_logits.py +++ b/test/registered/kernels/ops/attention/test_cute_dsl_fp8_paged_mqa_logits.py @@ -7,141 +7,18 @@ import pytest import torch from sglang.kernels.ops.attention.dsa import cutedsl_paged_mqa_logits, pick_dsl_expand -from sglang.srt.layers.attention.dsa.utils import ( - fp8_mqa_logits_ceil_to_ue8m0, - fp8_mqa_logits_make_fused_kv, -) from sglang.srt.utils import is_sm100_supported from sglang.test.ci.ci_register import register_cuda_ci +from sglang.test.kernels.paged_mqa import ( + BLOCK_KV, + HEAD_DIM, + assert_paged_mqa_matches_ref, + generate_paged_mqa_test_data, + ref_fp8_paged_mqa_logits, +) register_cuda_ci(est_time=180, stage="nightly", runner_config="4-gpu-b200") -BLOCK_KV = 64 -HEAD_DIM = 128 - - -def _ref_fp8_paged_mqa_logits( - q_fp8, - kv_fp8, - kv_scales, - weights, - context_lens, - block_table, - max_model_len, - block_kv, -): - B, next_n, H, D = q_fp8.shape - device = q_fp8.device - - logits = torch.full( - (B * next_n, max_model_len), float("-inf"), device=device, dtype=torch.float32 - ) - q_f32 = q_fp8.float() - - for b in range(B): - ctx_len = context_lens[b].item() - q_positions = torch.arange(ctx_len - next_n, ctx_len, device=device) - w = weights[b * next_n : (b + 1) * next_n, :] - - for blk_idx in range((ctx_len + block_kv - 1) // block_kv): - phys_blk = block_table[b, blk_idx].item() - k_f32 = kv_fp8[phys_blk].float() - scales = kv_scales[phys_blk] - - k_positions = torch.arange( - blk_idx * block_kv, (blk_idx + 1) * block_kv, device=device - ) - mask = (k_positions[None, :] < ctx_len) & ( - k_positions[None, :] <= q_positions[:, None] - ) - - qk = torch.matmul(q_f32[b].permute(1, 0, 2), k_f32.T) - qk = torch.where(mask[None, :, :], qk, torch.zeros(1, device=device)) - qk = torch.relu(qk) - - weighted = (w.T[:, :, None] * qk).sum(dim=0) - weighted = weighted * scales[None, :] - - start_pos = blk_idx * block_kv - end_pos = start_pos + block_kv - logits[b * next_n : (b + 1) * next_n, start_pos:end_pos] = torch.where( - mask, - weighted, - torch.tensor(float("-inf"), device=device, dtype=torch.float32), - ) - - return logits - - -def _generate_test_data( - batch_size, - next_n, - num_heads, - avg_context_len, - max_model_len, - device="cuda", -): - torch.manual_seed(42) - torch.cuda.manual_seed(42) - context_lens = torch.randint( - max(BLOCK_KV, int(0.7 * avg_context_len)), - int(1.3 * avg_context_len) + 1, - (batch_size,), - dtype=torch.int32, - device="cpu", - ).clamp(max=max_model_len) - - max_blocks_per_seq = (max_model_len + BLOCK_KV - 1) // BLOCK_KV - total_blocks = ((context_lens + BLOCK_KV - 1) // BLOCK_KV).sum().item() - num_phys_blocks = total_blocks + batch_size * 2 - - block_table = torch.full( - (batch_size, max_blocks_per_seq), 0, dtype=torch.int32, device=device - ) - blk_offset = 0 - for i in range(batch_size): - n_blks = (context_lens[i].item() + BLOCK_KV - 1) // BLOCK_KV - block_table[i, :n_blks] = torch.arange( - blk_offset, blk_offset + n_blks, dtype=torch.int32, device=device - ) - blk_offset += n_blks - - q_bf16 = torch.randn(batch_size, next_n, num_heads, HEAD_DIM, device=device) - q_fp8 = q_bf16.to(torch.float8_e4m3fn) - - kv_bf16 = torch.randn(num_phys_blocks, BLOCK_KV, HEAD_DIM, device=device) - kv_amax = kv_bf16.abs().float().amax(dim=-1, keepdim=True).clamp(1e-4) - kv_scale = fp8_mqa_logits_ceil_to_ue8m0(kv_amax / 448.0).squeeze(-1) - kv_fp8 = (kv_bf16 / kv_scale.unsqueeze(-1)).to(torch.float8_e4m3fn) - - weights = torch.randn( - batch_size * next_n, num_heads, device=device, dtype=torch.float32 - ) - kv_fused = fp8_mqa_logits_make_fused_kv(kv_fp8, kv_scale, BLOCK_KV, HEAD_DIM) - - return { - "q_fp8": q_fp8, - "kv_fp8": kv_fp8, - "kv_scales": kv_scale, - "kv_fused": kv_fused, - "weights": weights, - "context_lens": context_lens.to(device), - "block_table": block_table, - } - - -def _assert_matches_ref(logits, ref_logits, context_lens, B, next_n, max_model_len): - device = logits.device - positions = torch.arange(max_model_len, device=device).unsqueeze(0) - row_indices = torch.arange(B * next_n, device=device) // next_n - next_n_offset = torch.arange(B * next_n, device=device) % next_n - end_pos = context_lens[row_indices] - next_n + next_n_offset - mask = positions <= end_pos.unsqueeze(1) - - logits_masked = logits.float().masked_fill(~mask, 0) - ref_masked = ref_logits.float().masked_fill(~mask, 0) - torch.testing.assert_close(logits_masked, ref_masked, atol=5e-5, rtol=1e-5) - def _run_cutedsl_paged_mqa_logits( data, batch_size, next_n, num_heads, max_model_len, is_target_verify @@ -204,7 +81,9 @@ def _run_cutedsl_paged_mqa_logits( @pytest.mark.parametrize("avg_ctx", [128, 1024, 4096, 16384]) def test_cutedsl_paged_mqa_logits(batch_size, next_n, num_heads, avg_ctx): max_model_len = max(avg_ctx * 2, 2048) - data = _generate_test_data(batch_size, next_n, num_heads, avg_ctx, max_model_len) + data = generate_paged_mqa_test_data( + batch_size, next_n, num_heads, avg_ctx, max_model_len + ) logits = _run_cutedsl_paged_mqa_logits( data, @@ -215,7 +94,7 @@ def test_cutedsl_paged_mqa_logits(batch_size, next_n, num_heads, avg_ctx): is_target_verify=next_n >= 2, ) - ref_logits = _ref_fp8_paged_mqa_logits( + ref_logits = ref_fp8_paged_mqa_logits( data["q_fp8"], data["kv_fp8"], data["kv_scales"], @@ -225,7 +104,7 @@ def test_cutedsl_paged_mqa_logits(batch_size, next_n, num_heads, avg_ctx): max_model_len, BLOCK_KV, ) - _assert_matches_ref( + assert_paged_mqa_matches_ref( logits, ref_logits, data["context_lens"], batch_size, next_n, max_model_len ) diff --git a/test/registered/kernels/ops/attention/test_deepgemm_paged_mqa_logits.py b/test/registered/kernels/ops/attention/test_deepgemm_paged_mqa_logits.py index a28ebbd07..6bb953a2e 100644 --- a/test/registered/kernels/ops/attention/test_deepgemm_paged_mqa_logits.py +++ b/test/registered/kernels/ops/attention/test_deepgemm_paged_mqa_logits.py @@ -10,141 +10,18 @@ from sglang.kernels.ops.attention.dsa import ( deepgemm_paged_mqa_logits_native, deepgemm_paged_mqa_logits_split, ) -from sglang.srt.layers.attention.dsa.utils import ( - fp8_mqa_logits_ceil_to_ue8m0, - fp8_mqa_logits_make_fused_kv, -) from sglang.srt.utils import is_sm90_supported, is_sm100_supported from sglang.test.ci.ci_register import register_cuda_ci +from sglang.test.kernels.paged_mqa import ( + BLOCK_KV, + HEAD_DIM, + assert_paged_mqa_matches_ref, + generate_paged_mqa_test_data, + ref_fp8_paged_mqa_logits, +) register_cuda_ci(est_time=40, stage="nightly", runner_config="4-gpu-b200") -BLOCK_KV = 64 -HEAD_DIM = 128 - - -def _ref_fp8_paged_mqa_logits( - q_fp8, - kv_fp8, - kv_scales, - weights, - context_lens, - block_table, - max_model_len, - block_kv, -): - B, next_n, H, D = q_fp8.shape - device = q_fp8.device - - logits = torch.full( - (B * next_n, max_model_len), float("-inf"), device=device, dtype=torch.float32 - ) - q_f32 = q_fp8.float() - - for b in range(B): - ctx_len = context_lens[b].item() - q_positions = torch.arange(ctx_len - next_n, ctx_len, device=device) - w = weights[b * next_n : (b + 1) * next_n, :] - - for blk_idx in range((ctx_len + block_kv - 1) // block_kv): - phys_blk = block_table[b, blk_idx].item() - k_f32 = kv_fp8[phys_blk].float() - scales = kv_scales[phys_blk] - - k_positions = torch.arange( - blk_idx * block_kv, (blk_idx + 1) * block_kv, device=device - ) - mask = (k_positions[None, :] < ctx_len) & ( - k_positions[None, :] <= q_positions[:, None] - ) - - qk = torch.matmul(q_f32[b].permute(1, 0, 2), k_f32.T) - qk = torch.where(mask[None, :, :], qk, torch.zeros(1, device=device)) - qk = torch.relu(qk) - - weighted = (w.T[:, :, None] * qk).sum(dim=0) - weighted = weighted * scales[None, :] - - start_pos = blk_idx * block_kv - end_pos = start_pos + block_kv - logits[b * next_n : (b + 1) * next_n, start_pos:end_pos] = torch.where( - mask, - weighted, - torch.tensor(float("-inf"), device=device, dtype=torch.float32), - ) - - return logits - - -def _generate_test_data( - batch_size, - next_n, - num_heads, - avg_context_len, - max_model_len, - device="cuda", -): - torch.manual_seed(42) - torch.cuda.manual_seed(42) - context_lens = torch.randint( - max(BLOCK_KV, int(0.7 * avg_context_len)), - int(1.3 * avg_context_len) + 1, - (batch_size,), - dtype=torch.int32, - device="cpu", - ).clamp(max=max_model_len) - - max_blocks_per_seq = (max_model_len + BLOCK_KV - 1) // BLOCK_KV - total_blocks = ((context_lens + BLOCK_KV - 1) // BLOCK_KV).sum().item() - num_phys_blocks = total_blocks + batch_size * 2 - - block_table = torch.full( - (batch_size, max_blocks_per_seq), 0, dtype=torch.int32, device=device - ) - blk_offset = 0 - for i in range(batch_size): - n_blks = (context_lens[i].item() + BLOCK_KV - 1) // BLOCK_KV - block_table[i, :n_blks] = torch.arange( - blk_offset, blk_offset + n_blks, dtype=torch.int32, device=device - ) - blk_offset += n_blks - - q_bf16 = torch.randn(batch_size, next_n, num_heads, HEAD_DIM, device=device) - q_fp8 = q_bf16.to(torch.float8_e4m3fn) - - kv_bf16 = torch.randn(num_phys_blocks, BLOCK_KV, HEAD_DIM, device=device) - kv_amax = kv_bf16.abs().float().amax(dim=-1, keepdim=True).clamp(1e-4) - kv_scale = fp8_mqa_logits_ceil_to_ue8m0(kv_amax / 448.0).squeeze(-1) - kv_fp8 = (kv_bf16 / kv_scale.unsqueeze(-1)).to(torch.float8_e4m3fn) - - weights = torch.randn( - batch_size * next_n, num_heads, device=device, dtype=torch.float32 - ) - kv_fused = fp8_mqa_logits_make_fused_kv(kv_fp8, kv_scale, BLOCK_KV, HEAD_DIM) - - return { - "q_fp8": q_fp8, - "kv_fp8": kv_fp8, - "kv_scales": kv_scale, - "kv_fused": kv_fused, - "weights": weights, - "context_lens": context_lens.to(device), - "block_table": block_table, - } - - -def _assert_matches_ref(logits, ref_logits, context_lens, B, next_n, max_model_len): - device = logits.device - positions = torch.arange(max_model_len, device=device).unsqueeze(0) - row_indices = torch.arange(B * next_n, device=device) // next_n - next_n_offset = torch.arange(B * next_n, device=device) % next_n - end_pos = context_lens[row_indices] - next_n + next_n_offset - mask = positions <= end_pos.unsqueeze(1) - - logits_masked = logits.float().masked_fill(~mask, 0) - ref_masked = ref_logits.float().masked_fill(~mask, 0) - torch.testing.assert_close(logits_masked, ref_masked, atol=5e-5, rtol=1e-5) - def _run_deepgemm_paged_mqa_logits(data, batch_size, next_n, num_heads, max_model_len): """Mirrors the DEEPGEMM dispatch in @@ -208,13 +85,15 @@ def _run_deepgemm_paged_mqa_logits(data, batch_size, next_n, num_heads, max_mode @pytest.mark.parametrize("avg_ctx", [128, 1024, 4096, 16384]) def test_deepgemm_paged_mqa_logits(batch_size, next_n, num_heads, avg_ctx): max_model_len = max(avg_ctx * 2, 2048) - data = _generate_test_data(batch_size, next_n, num_heads, avg_ctx, max_model_len) + data = generate_paged_mqa_test_data( + batch_size, next_n, num_heads, avg_ctx, max_model_len + ) logits = _run_deepgemm_paged_mqa_logits( data, batch_size, next_n, num_heads, max_model_len ) - ref_logits = _ref_fp8_paged_mqa_logits( + ref_logits = ref_fp8_paged_mqa_logits( data["q_fp8"], data["kv_fp8"], data["kv_scales"], @@ -224,7 +103,7 @@ def test_deepgemm_paged_mqa_logits(batch_size, next_n, num_heads, avg_ctx): max_model_len, BLOCK_KV, ) - _assert_matches_ref( + assert_paged_mqa_matches_ref( logits, ref_logits, data["context_lens"], batch_size, next_n, max_model_len ) diff --git a/test/registered/kernels/ops/attention/test_fused_metadata_copy.py b/test/registered/kernels/ops/attention/test_fused_metadata_copy.py index 805352dda..a4f093c24 100644 --- a/test/registered/kernels/ops/attention/test_fused_metadata_copy.py +++ b/test/registered/kernels/ops/attention/test_fused_metadata_copy.py @@ -1,15 +1,6 @@ -""" -Comprehensive tests for JIT-compiled fused metadata copy kernels. - -This test suite verifies: -1. Single-backend fused kernel (fused_metadata_copy_cuda) - all forward modes -2. Multi-backend fused kernel (fused_metadata_copy_multi_cuda) - 3 backends at once -3. Correctness against reference implementations -4. Performance benchmarks and speedup measurements -""" +"""Compare single- and multi-backend metadata copies with PyTorch references.""" import sys -import time import pytest import torch @@ -151,62 +142,6 @@ def reference_copy_decode(src, dst, max_len): dst["flashmla_metadata"].copy_(src["flashmla_metadata"]) -def reference_copy_target_verify(src, dst, max_seqlen_k, seqlens_expanded_size): - """Reference implementation: individual .copy_() for TARGET_VERIFY mode.""" - bs = src["cache_seqlens"].shape[0] - dst["cache_seqlens"].copy_(src["cache_seqlens"]) - dst["cu_seqlens_k"][1:].copy_(src["cu_seqlens_k"][1:]) - - rows, cols = src["page_indices"].shape - dst["page_table_1"][:rows, :cols].copy_(src["page_indices"]) - dst["dsa_seqlens_expanded"][:seqlens_expanded_size].copy_(src["seqlens_expanded"]) - dst["dsa_cache_seqlens"][:seqlens_expanded_size].copy_(src["dsa_cache_seqlens"]) - dst["dsa_cu_seqlens_k"][1 : seqlens_expanded_size + 1].copy_( - src["dsa_cu_seqlens_k"][1 : seqlens_expanded_size + 1] - ) - - if src["real_page_table"] is not None: - rows, cols = src["real_page_table"].shape - dst["real_page_table"][:rows, :cols].copy_(src["real_page_table"]) - - if src["flashmla_num_splits"] is not None: - flashmla_size = seqlens_expanded_size + 1 - dst["flashmla_num_splits"][:flashmla_size].copy_( - src["flashmla_num_splits"][:flashmla_size] - ) - - if src["flashmla_metadata"] is not None: - dst["flashmla_metadata"].copy_(src["flashmla_metadata"]) - - -def reference_copy_draft_extend(src, dst, max_seqlen_k, seqlens_expanded_size): - """Reference implementation: individual .copy_() for DRAFT_EXTEND mode.""" - bs = src["cache_seqlens"].shape[0] - dst["cache_seqlens"].copy_(src["cache_seqlens"]) - dst["cu_seqlens_k"][1:].copy_(src["cu_seqlens_k"][1:]) - - rows, cols = src["page_indices"].shape - dst["page_table_1"][:rows, :cols].copy_(src["page_indices"]) - dst["dsa_seqlens_expanded"][:seqlens_expanded_size].copy_(src["seqlens_expanded"]) - dst["dsa_cache_seqlens"][:seqlens_expanded_size].copy_(src["dsa_cache_seqlens"]) - dst["dsa_cu_seqlens_k"][1 : seqlens_expanded_size + 1].copy_( - src["dsa_cu_seqlens_k"][1 : seqlens_expanded_size + 1] - ) - - if src["real_page_table"] is not None: - rows, cols = src["real_page_table"].shape - dst["real_page_table"][:rows, :cols].copy_(src["real_page_table"]) - - if src["flashmla_num_splits"] is not None: - flashmla_size = seqlens_expanded_size + 1 - dst["flashmla_num_splits"][:flashmla_size].copy_( - src["flashmla_num_splits"][:flashmla_size] - ) - - if src["flashmla_metadata"] is not None: - dst["flashmla_metadata"].copy_(src["flashmla_metadata"]) - - # ============================================================================= # Single-Backend Kernel Tests # ============================================================================= @@ -321,13 +256,17 @@ def test_fused_metadata_copy_dtype_validation(): ) -@pytest.mark.parametrize("bs", [1, 2, 4, 8]) @pytest.mark.parametrize( - "forward_mode", [0] -) # DECODE mode only (other modes not fully tested yet) -@pytest.mark.parametrize("has_real_page_table", [False, True]) -@pytest.mark.parametrize("has_flashmla", [False, True]) -def test_fused_metadata_copy(bs, forward_mode, has_real_page_table, has_flashmla): + "bs,has_real_page_table,has_flashmla", + [ + (bs, page, mla) + for bs in (1, 2, 4, 8) + for page in (False, True) + for mla in (False, True) + ] + + [(16, True, True), (32, True, True)], +) +def test_fused_metadata_copy(bs, has_real_page_table, has_flashmla): """Test fused metadata copy kernel against reference implementation.""" if not torch.cuda.is_available(): pytest.skip("CUDA not available") @@ -336,9 +275,10 @@ def test_fused_metadata_copy(bs, forward_mode, has_real_page_table, has_flashmla fused_metadata_copy_cuda, ) + forward_mode = 0 # DECODE max_len = 128 max_seqlen_k = 256 - seqlens_expanded_size = bs if forward_mode == 0 else bs * 2 + seqlens_expanded_size = bs # Create test data data = create_test_metadata( @@ -356,17 +296,7 @@ def test_fused_metadata_copy(bs, forward_mode, has_real_page_table, has_flashmla k: v.clone() if v is not None else None for k, v in data["dst"].items() } - # Run reference implementation - if forward_mode == 0: # DECODE - reference_copy_decode(data["src"], dst_ref, max_len) - elif forward_mode == 1: # TARGET_VERIFY - reference_copy_target_verify( - data["src"], dst_ref, max_seqlen_k, seqlens_expanded_size - ) - else: # DRAFT_EXTEND - reference_copy_draft_extend( - data["src"], dst_ref, max_seqlen_k, seqlens_expanded_size - ) + reference_copy_decode(data["src"], dst_ref, max_len) # Run fused kernel fused_metadata_copy_cuda( @@ -395,101 +325,11 @@ def test_fused_metadata_copy(bs, forward_mode, has_real_page_table, has_flashmla seqlens_expanded_size, ) - # Compare results - assert torch.equal(dst_ref["cache_seqlens"], dst_fused["cache_seqlens"]), ( - "cache_seqlens mismatch" - ) - assert torch.equal(dst_ref["cu_seqlens_k"], dst_fused["cu_seqlens_k"]), ( - "cu_seqlens_k mismatch" - ) - assert torch.equal(dst_ref["page_table_1"], dst_fused["page_table_1"]), ( - "page_table_1 mismatch" - ) - assert torch.equal(dst_ref["dsa_cache_seqlens"], dst_fused["dsa_cache_seqlens"]), ( - "dsa_cache_seqlens mismatch" - ) - assert torch.equal( - dst_ref["dsa_seqlens_expanded"], dst_fused["dsa_seqlens_expanded"] - ), "dsa_seqlens_expanded mismatch" - assert torch.equal(dst_ref["dsa_cu_seqlens_k"], dst_fused["dsa_cu_seqlens_k"]), ( - "dsa_cu_seqlens_k mismatch" - ) - - if has_real_page_table: - assert torch.equal(dst_ref["real_page_table"], dst_fused["real_page_table"]), ( - "real_page_table mismatch" - ) - - if has_flashmla: - assert torch.equal( - dst_ref["flashmla_num_splits"], dst_fused["flashmla_num_splits"] - ), "flashmla_num_splits mismatch" - assert torch.equal( - dst_ref["flashmla_metadata"], dst_fused["flashmla_metadata"] - ), "flashmla_metadata mismatch" - - -@pytest.mark.parametrize("bs", [16, 32]) -def test_fused_metadata_copy_large_batch(bs): - """Test with larger batch sizes.""" - if not torch.cuda.is_available(): - pytest.skip("CUDA not available") - - from sglang.kernels.ops.attention.fused_metadata_copy import ( - fused_metadata_copy_cuda, - ) - - forward_mode = 0 # DECODE - max_len = 128 - max_seqlen_k = 256 - seqlens_expanded_size = bs - - data = create_test_metadata( - bs=bs, - max_len=max_len, - max_seqlen_k=max_seqlen_k, - seqlens_expanded_size=seqlens_expanded_size, - has_real_page_table=True, - has_flashmla=True, - ) - - dst_ref = {k: v.clone() if v is not None else None for k, v in data["dst"].items()} - dst_fused = { - k: v.clone() if v is not None else None for k, v in data["dst"].items() - } - - reference_copy_decode(data["src"], dst_ref, max_len) - - fused_metadata_copy_cuda( - data["src"]["cache_seqlens"], - data["src"]["cu_seqlens_k"], - data["src"]["page_indices"], - data["src"]["dsa_cache_seqlens"], - data["src"]["seqlens_expanded"], - data["src"]["dsa_cu_seqlens_k"], - data["src"]["real_page_table"], - data["src"]["flashmla_num_splits"], - data["src"]["flashmla_metadata"], - dst_fused["cache_seqlens"], - dst_fused["cu_seqlens_k"], - dst_fused["page_table_1"], - dst_fused["dsa_cache_seqlens"], - dst_fused["dsa_seqlens_expanded"], - dst_fused["dsa_cu_seqlens_k"], - dst_fused["real_page_table"], - dst_fused["flashmla_num_splits"], - dst_fused["flashmla_metadata"], - forward_mode, - bs, - max_len, - max_seqlen_k, - seqlens_expanded_size, - ) - - # Verify all tensors match - for key in dst_ref: - if dst_ref[key] is not None: - assert torch.equal(dst_ref[key], dst_fused[key]), f"{key} mismatch" + for key, expected in dst_ref.items(): + if expected is not None: + torch.testing.assert_close( + dst_fused[key], expected, rtol=0, atol=0, msg=key + ) # ============================================================================= @@ -725,9 +565,16 @@ def test_fused_metadata_copy_multi_dtype_validation(): ) -@pytest.mark.parametrize("bs", [1, 2, 4, 8, 16]) -@pytest.mark.parametrize("has_real_page_table", [False, True]) -@pytest.mark.parametrize("has_flashmla", [False, True]) +@pytest.mark.parametrize( + "bs,has_real_page_table,has_flashmla", + [ + (bs, page, mla) + for bs in (1, 2, 4, 8, 16) + for page in (False, True) + for mla in (False, True) + ] + + [(32, True, True), (64, True, True)], +) def test_fused_metadata_copy_multi(bs, has_real_page_table, has_flashmla): """Test fused multi-backend metadata copy kernel against for-loop version.""" if not torch.cuda.is_available(): @@ -749,38 +596,16 @@ def test_fused_metadata_copy_multi(bs, has_real_page_table, has_flashmla): has_flashmla=has_flashmla, ) - # Create separate destination tensors for reference (for-loop) and fused kernel - dst_ref_0 = { - k: v.clone() if v is not None else None for k, v in data["dst0"].items() - } - dst_ref_1 = { - k: v.clone() if v is not None else None for k, v in data["dst1"].items() - } - dst_ref_2 = { - k: v.clone() if v is not None else None for k, v in data["dst2"].items() - } + dst_ref = [ + {k: v.clone() if v is not None else None for k, v in data[f"dst{i}"].items()} + for i in range(3) + ] + dst_fused = [ + {k: v.clone() if v is not None else None for k, v in data[f"dst{i}"].items()} + for i in range(3) + ] + reference_copy_for_loop(data["src"], dst_ref, bs, max_len) - dst_fused_0 = { - k: v.clone() if v is not None else None for k, v in data["dst0"].items() - } - dst_fused_1 = { - k: v.clone() if v is not None else None for k, v in data["dst1"].items() - } - dst_fused_2 = { - k: v.clone() if v is not None else None for k, v in data["dst2"].items() - } - - # Run reference implementation (for-loop) - torch.cuda.synchronize() - loop_start = time.perf_counter() - reference_copy_for_loop(data["src"], [dst_ref_0, dst_ref_1, dst_ref_2], bs, max_len) - torch.cuda.synchronize() - loop_end = time.perf_counter() - loop_time = loop_end - loop_start - - # Run fused kernel - torch.cuda.synchronize() - fused_start = time.perf_counter() fused_metadata_copy_multi_cuda( # Source tensors data["src"]["cache_seqlens"], @@ -792,296 +617,46 @@ def test_fused_metadata_copy_multi(bs, has_real_page_table, has_flashmla): data["src"]["flashmla_num_splits"], data["src"]["flashmla_metadata"], # Destination tensors for backend 0 - dst_fused_0["cache_seqlens_int32"], - dst_fused_0["cu_seqlens_k"], - dst_fused_0["page_table_1"], - dst_fused_0["dsa_cache_seqlens_int32"], - dst_fused_0["dsa_cu_seqlens_k"], - dst_fused_0["real_page_table"], - dst_fused_0["flashmla_num_splits"], - dst_fused_0["flashmla_metadata"], + dst_fused[0]["cache_seqlens_int32"], + dst_fused[0]["cu_seqlens_k"], + dst_fused[0]["page_table_1"], + dst_fused[0]["dsa_cache_seqlens_int32"], + dst_fused[0]["dsa_cu_seqlens_k"], + dst_fused[0]["real_page_table"], + dst_fused[0]["flashmla_num_splits"], + dst_fused[0]["flashmla_metadata"], # Destination tensors for backend 1 - dst_fused_1["cache_seqlens_int32"], - dst_fused_1["cu_seqlens_k"], - dst_fused_1["page_table_1"], - dst_fused_1["dsa_cache_seqlens_int32"], - dst_fused_1["dsa_cu_seqlens_k"], - dst_fused_1["real_page_table"], - dst_fused_1["flashmla_num_splits"], - dst_fused_1["flashmla_metadata"], + dst_fused[1]["cache_seqlens_int32"], + dst_fused[1]["cu_seqlens_k"], + dst_fused[1]["page_table_1"], + dst_fused[1]["dsa_cache_seqlens_int32"], + dst_fused[1]["dsa_cu_seqlens_k"], + dst_fused[1]["real_page_table"], + dst_fused[1]["flashmla_num_splits"], + dst_fused[1]["flashmla_metadata"], # Destination tensors for backend 2 - dst_fused_2["cache_seqlens_int32"], - dst_fused_2["cu_seqlens_k"], - dst_fused_2["page_table_1"], - dst_fused_2["dsa_cache_seqlens_int32"], - dst_fused_2["dsa_cu_seqlens_k"], - dst_fused_2["real_page_table"], - dst_fused_2["flashmla_num_splits"], - dst_fused_2["flashmla_metadata"], + dst_fused[2]["cache_seqlens_int32"], + dst_fused[2]["cu_seqlens_k"], + dst_fused[2]["page_table_1"], + dst_fused[2]["dsa_cache_seqlens_int32"], + dst_fused[2]["dsa_cu_seqlens_k"], + dst_fused[2]["real_page_table"], + dst_fused[2]["flashmla_num_splits"], + dst_fused[2]["flashmla_metadata"], # Parameters bs, max_len, seqlens_expanded_size, ) - torch.cuda.synchronize() - fused_end = time.perf_counter() - fused_time = fused_end - fused_start - - # Compare results for all 3 backends - speedup = loop_time / fused_time if fused_time > 0 else 0 - print( - f"\n[VERIFY] bs={bs}, real_page_table={has_real_page_table}, flashmla={has_flashmla}" - ) - print( - f"[VERIFY] Fused time: {fused_time * 1000:.3f}ms, Loop time: {loop_time * 1000:.3f}ms, Speedup: {speedup:.2f}x" - ) - - max_diff = 0.0 - all_match = True - - for backend_idx, (dst_ref, dst_fused) in enumerate( - [ - (dst_ref_0, dst_fused_0), - (dst_ref_1, dst_fused_1), - (dst_ref_2, dst_fused_2), - ] - ): - for key in [ - "cache_seqlens_int32", - "cu_seqlens_k", - "page_table_1", - "dsa_cache_seqlens_int32", - "dsa_cu_seqlens_k", - ]: - if not torch.equal(dst_ref[key], dst_fused[key]): - diff = ( - (dst_ref[key].float() - dst_fused[key].float()).abs().max().item() - ) - max_diff = max(max_diff, diff) - all_match = False - print( - f"[ERROR] Backend {backend_idx} {key}: MISMATCH! Max diff: {diff}" - ) - - if has_real_page_table and dst_ref["real_page_table"] is not None: - if not torch.equal( - dst_ref["real_page_table"], dst_fused["real_page_table"] - ): - diff = ( - ( - dst_ref["real_page_table"].float() - - dst_fused["real_page_table"].float() - ) - .abs() - .max() - .item() - ) - max_diff = max(max_diff, diff) - all_match = False - print( - f"[ERROR] Backend {backend_idx} real_page_table: MISMATCH! Max diff: {diff}" - ) - - if has_flashmla: - if dst_ref["flashmla_num_splits"] is not None and not torch.equal( - dst_ref["flashmla_num_splits"], dst_fused["flashmla_num_splits"] - ): - diff = ( - ( - dst_ref["flashmla_num_splits"].float() - - dst_fused["flashmla_num_splits"].float() - ) - .abs() - .max() - .item() - ) - max_diff = max(max_diff, diff) - all_match = False - print( - f"[ERROR] Backend {backend_idx} flashmla_num_splits: MISMATCH! Max diff: {diff}" - ) - - if dst_ref["flashmla_metadata"] is not None and not torch.equal( - dst_ref["flashmla_metadata"], dst_fused["flashmla_metadata"] - ): - diff = ( - ( - dst_ref["flashmla_metadata"].float() - - dst_fused["flashmla_metadata"].float() - ) - .abs() - .max() - .item() - ) - max_diff = max(max_diff, diff) - all_match = False - print( - f"[ERROR] Backend {backend_idx} flashmla_metadata: MISMATCH! Max diff: {diff}" - ) - - if not all_match: - error_msg = ( - f"Fused metadata copy verification FAILED! " - f"Maximum difference: {max_diff}. " - f"The fused kernel produces different results than the for-loop version." - ) - print(f"[ERROR] {error_msg}") - raise AssertionError(error_msg) - - print(f"[VERIFY] Verification PASSED - all tensors match!") - - -@pytest.mark.parametrize("bs", [32, 64]) -def test_fused_metadata_copy_multi_large_batch(bs): - """Test with larger batch sizes and timing comparison.""" - if not torch.cuda.is_available(): - pytest.skip("CUDA not available") - - from sglang.kernels.ops.attention.fused_metadata_copy import ( - fused_metadata_copy_multi_cuda, - ) - - max_len = 128 - seqlens_expanded_size = bs - - data = create_test_metadata_multi( - bs=bs, - max_len=max_len, - seqlens_expanded_size=seqlens_expanded_size, - has_real_page_table=True, - has_flashmla=True, - ) - - dst_ref_0 = { - k: v.clone() if v is not None else None for k, v in data["dst0"].items() - } - dst_ref_1 = { - k: v.clone() if v is not None else None for k, v in data["dst1"].items() - } - dst_ref_2 = { - k: v.clone() if v is not None else None for k, v in data["dst2"].items() - } - - dst_fused_0 = { - k: v.clone() if v is not None else None for k, v in data["dst0"].items() - } - dst_fused_1 = { - k: v.clone() if v is not None else None for k, v in data["dst1"].items() - } - dst_fused_2 = { - k: v.clone() if v is not None else None for k, v in data["dst2"].items() - } - - # Warmup - for _ in range(5): - reference_copy_for_loop( - data["src"], [dst_ref_0, dst_ref_1, dst_ref_2], bs, max_len - ) - fused_metadata_copy_multi_cuda( - data["src"]["cache_seqlens"], - data["src"]["cu_seqlens_k"], - data["src"]["page_indices"], - data["src"]["dsa_cache_seqlens"], - data["src"]["dsa_cu_seqlens_k"], - data["src"]["real_page_table"], - data["src"]["flashmla_num_splits"], - data["src"]["flashmla_metadata"], - dst_fused_0["cache_seqlens_int32"], - dst_fused_0["cu_seqlens_k"], - dst_fused_0["page_table_1"], - dst_fused_0["dsa_cache_seqlens_int32"], - dst_fused_0["dsa_cu_seqlens_k"], - dst_fused_0["real_page_table"], - dst_fused_0["flashmla_num_splits"], - dst_fused_0["flashmla_metadata"], - dst_fused_1["cache_seqlens_int32"], - dst_fused_1["cu_seqlens_k"], - dst_fused_1["page_table_1"], - dst_fused_1["dsa_cache_seqlens_int32"], - dst_fused_1["dsa_cu_seqlens_k"], - dst_fused_1["real_page_table"], - dst_fused_1["flashmla_num_splits"], - dst_fused_1["flashmla_metadata"], - dst_fused_2["cache_seqlens_int32"], - dst_fused_2["cu_seqlens_k"], - dst_fused_2["page_table_1"], - dst_fused_2["dsa_cache_seqlens_int32"], - dst_fused_2["dsa_cu_seqlens_k"], - dst_fused_2["real_page_table"], - dst_fused_2["flashmla_num_splits"], - dst_fused_2["flashmla_metadata"], - bs, - max_len, - seqlens_expanded_size, - ) - torch.cuda.synchronize() - - # Actual timing - torch.cuda.synchronize() - loop_start = time.perf_counter() - reference_copy_for_loop(data["src"], [dst_ref_0, dst_ref_1, dst_ref_2], bs, max_len) - torch.cuda.synchronize() - loop_time = time.perf_counter() - loop_start - - torch.cuda.synchronize() - fused_start = time.perf_counter() - fused_metadata_copy_multi_cuda( - data["src"]["cache_seqlens"], - data["src"]["cu_seqlens_k"], - data["src"]["page_indices"], - data["src"]["dsa_cache_seqlens"], - data["src"]["dsa_cu_seqlens_k"], - data["src"]["real_page_table"], - data["src"]["flashmla_num_splits"], - data["src"]["flashmla_metadata"], - dst_fused_0["cache_seqlens_int32"], - dst_fused_0["cu_seqlens_k"], - dst_fused_0["page_table_1"], - dst_fused_0["dsa_cache_seqlens_int32"], - dst_fused_0["dsa_cu_seqlens_k"], - dst_fused_0["real_page_table"], - dst_fused_0["flashmla_num_splits"], - dst_fused_0["flashmla_metadata"], - dst_fused_1["cache_seqlens_int32"], - dst_fused_1["cu_seqlens_k"], - dst_fused_1["page_table_1"], - dst_fused_1["dsa_cache_seqlens_int32"], - dst_fused_1["dsa_cu_seqlens_k"], - dst_fused_1["real_page_table"], - dst_fused_1["flashmla_num_splits"], - dst_fused_1["flashmla_metadata"], - dst_fused_2["cache_seqlens_int32"], - dst_fused_2["cu_seqlens_k"], - dst_fused_2["page_table_1"], - dst_fused_2["dsa_cache_seqlens_int32"], - dst_fused_2["dsa_cu_seqlens_k"], - dst_fused_2["real_page_table"], - dst_fused_2["flashmla_num_splits"], - dst_fused_2["flashmla_metadata"], - bs, - max_len, - seqlens_expanded_size, - ) - torch.cuda.synchronize() - fused_time = time.perf_counter() - fused_start - - speedup = loop_time / fused_time if fused_time > 0 else 0 - print( - f"\n[PERF] Large batch (bs={bs}): Fused={fused_time * 1000:.3f}ms, Loop={loop_time * 1000:.3f}ms, Speedup={speedup:.2f}x" - ) - - # Verify correctness - for backend_idx, (dst_ref, dst_fused) in enumerate( - [ - (dst_ref_0, dst_fused_0), - (dst_ref_1, dst_fused_1), - (dst_ref_2, dst_fused_2), - ] - ): - for key in dst_ref: - if dst_ref[key] is not None and dst_fused[key] is not None: - assert torch.equal(dst_ref[key], dst_fused[key]), ( - f"Backend {backend_idx} {key} mismatch" + for backend_idx, (expected, actual) in enumerate(zip(dst_ref, dst_fused)): + for key, tensor in expected.items(): + if tensor is not None: + torch.testing.assert_close( + actual[key], + tensor, + rtol=0, + atol=0, + msg=f"Backend {backend_idx} {key}", ) diff --git a/test/registered/kernels/ops/diffusion/test_flux2_token_cat_fp8.py b/test/registered/kernels/ops/diffusion/test_flux2_token_cat_fp8.py index 230bb7d8b..3fb5e2b4c 100644 --- a/test/registered/kernels/ops/diffusion/test_flux2_token_cat_fp8.py +++ b/test/registered/kernels/ops/diffusion/test_flux2_token_cat_fp8.py @@ -36,19 +36,14 @@ def test_flux2_token_cat_fp8_is_bit_exact(tokens: int) -> None: assert torch.equal(actual, expected) -def test_flux2_token_cat_fp8_rejects_compile() -> None: +@pytest.mark.parametrize( + "guard", ["torch.compiler.is_compiling", "torch.cuda.is_current_stream_capturing"] +) +def test_flux2_token_cat_fp8_rejects_capture(guard) -> None: attention = torch.empty((1, 1, 16), device="cuda", dtype=torch.bfloat16) mlp = torch.empty((1, 1, 48), device="cuda", dtype=torch.bfloat16) scale = torch.ones((1,), device="cuda", dtype=torch.float32) - with patch("torch.compiler.is_compiling", return_value=True): - assert try_flux2_token_cat_fp8(attention, mlp, scale) is None - - -def test_flux2_token_cat_fp8_rejects_cuda_graph_capture() -> None: - attention = torch.empty((1, 1, 16), device="cuda", dtype=torch.bfloat16) - mlp = torch.empty((1, 1, 48), device="cuda", dtype=torch.bfloat16) - scale = torch.ones((1,), device="cuda", dtype=torch.float32) - with patch("torch.cuda.is_current_stream_capturing", return_value=True): + with patch(guard, return_value=True): assert try_flux2_token_cat_fp8(attention, mlp, scale) is None diff --git a/test/registered/kernels/ops/diffusion/test_modulate.py b/test/registered/kernels/ops/diffusion/test_modulate.py index 436e060e4..c4bee5102 100644 --- a/test/registered/kernels/ops/diffusion/test_modulate.py +++ b/test/registered/kernels/ops/diffusion/test_modulate.py @@ -146,26 +146,6 @@ GATE_CASES = [ ] -def _assert_gate_add(out, ref): - if ref.dtype == torch.float32: - # fp32 has no rounding boundary to reproduce; the kernel keeps the - # accumulation in fp32 and only order may differ. - torch.testing.assert_close(out, ref, atol=1e-5, rtol=1e-5) - else: - assert torch.equal(out, ref) - - -@pytest.mark.parametrize("residual_shape,gate_shape", GATE_CASES) -def test_residual_gate_add_matches_torch(residual_shape, gate_shape): - residual = torch.randn(residual_shape, device=DEVICE, dtype=torch.bfloat16) - update = torch.randn_like(residual) - gate = torch.randn(gate_shape, device=DEVICE, dtype=torch.bfloat16) - - ref = residual + update * gate - _assert_gate_add(residual_gate_add_cuda(residual, update, gate), ref) - assert torch.equal(residual_gate_add(residual, update, gate), ref) - - # LingBot per-token gates are [B, S, 1]: one scalar per token, broadcast # along the hidden dimension. PER_TOKEN_GATE_CASES = [ @@ -175,8 +155,17 @@ PER_TOKEN_GATE_CASES = [ ] -@pytest.mark.parametrize("residual_shape,gate_shape", PER_TOKEN_GATE_CASES) -def test_residual_gate_add_per_token_matches_torch(residual_shape, gate_shape): +def _assert_gate_add(out, ref): + if ref.dtype == torch.float32: + # fp32 has no rounding boundary to reproduce; the kernel keeps the + # accumulation in fp32 and only order may differ. + torch.testing.assert_close(out, ref, atol=1e-5, rtol=1e-5) + else: + assert torch.equal(out, ref) + + +@pytest.mark.parametrize("residual_shape,gate_shape", GATE_CASES + PER_TOKEN_GATE_CASES) +def test_residual_gate_add_matches_torch(residual_shape, gate_shape): residual = torch.randn(residual_shape, device=DEVICE, dtype=torch.bfloat16) update = torch.randn_like(residual) gate = torch.randn(gate_shape, device=DEVICE, dtype=torch.bfloat16) @@ -188,19 +177,12 @@ def test_residual_gate_add_per_token_matches_torch(residual_shape, gate_shape): @pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16, torch.float32]) -def test_residual_gate_add_per_token_dtypes(dtype): - residual = torch.randn((1, 2560, 512), device=DEVICE, dtype=dtype) - update = torch.randn_like(residual) - gate = torch.randn((1, 2560, 1), device=DEVICE, dtype=dtype) - _assert_gate_add( - residual_gate_add_cuda(residual, update, gate), residual + update * gate - ) - - -@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16, torch.float32]) -@pytest.mark.parametrize("gate_shape", [(1, 1, 64), (1, 9, 64)]) -def test_residual_gate_add_dtypes(dtype, gate_shape): - residual = torch.randn((1, 9, 64), device=DEVICE, dtype=dtype) +@pytest.mark.parametrize( + "shape,gate_shape", + [((1, 9, 64), (1, 1, 64)), ((1, 9, 64), (1, 9, 64)), PER_TOKEN_GATE_CASES[0]], +) +def test_residual_gate_add_dtypes(dtype, shape, gate_shape): + residual = torch.randn(shape, device=DEVICE, dtype=dtype) update = torch.randn_like(residual) gate = torch.randn(gate_shape, device=DEVICE, dtype=dtype) _assert_gate_add( @@ -248,10 +230,12 @@ def test_residual_gate_add_transposed_storage_offsets(): assert torch.equal(out, residual + update * gate) -def test_residual_gate_add_transposed_torch_compile_fullgraph(): - residual = torch.randn((1, 128, 32), device=DEVICE, dtype=torch.bfloat16).transpose( - 1, 2 - ) +@pytest.mark.parametrize("transposed", [False, True]) +def test_residual_gate_add_torch_compile_fullgraph(transposed): + shape = (1, 128, 32) if transposed else (1, 32, 128) + residual = torch.randn(shape, device=DEVICE, dtype=torch.bfloat16) + if transposed: + residual = residual.transpose(1, 2) update = torch.randn_like(residual, memory_format=torch.contiguous_format) gate = torch.randn((1, 1, 128), device=DEVICE, dtype=torch.bfloat16) compiled = torch.compile(residual_gate_add, fullgraph=True) @@ -309,14 +293,6 @@ def test_residual_gate_add_guards_and_eager_fallback(): ) -def test_residual_gate_add_torch_compile_fullgraph(): - residual = torch.randn((1, 32, 128), device=DEVICE, dtype=torch.bfloat16) - update = torch.randn_like(residual) - gate = torch.randn((1, 1, 128), device=DEVICE, dtype=torch.bfloat16) - compiled = torch.compile(residual_gate_add, fullgraph=True) - assert torch.equal(compiled(residual, update, gate), residual + update * gate) - - @torch.no_grad() @pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16]) def test_scaled_residual_add_is_bit_exact(dtype): diff --git a/test/registered/kernels/ops/quantization/test_awq_marlin_moe_repack.py b/test/registered/kernels/ops/quantization/test_awq_marlin_moe_repack.py index 6c71e43ad..d407afee9 100644 --- a/test/registered/kernels/ops/quantization/test_awq_marlin_moe_repack.py +++ b/test/registered/kernels/ops/quantization/test_awq_marlin_moe_repack.py @@ -1,6 +1,5 @@ import sys -import numpy as np import pytest import torch from sgl_kernel.scalar_type import scalar_types @@ -8,8 +7,9 @@ from sgl_kernel.scalar_type import scalar_types from sglang.kernels.ops.quantization.awq_marlin_repack import ( awq_marlin_moe_repack as jit_awq_marlin_moe_repack, ) -from sglang.srt.layers.quantization.utils import pack_cols, quantize_weights +from sglang.srt.layers.quantization.utils import quantize_weights from sglang.test.ci.ci_register import register_cuda_ci +from sglang.test.test_marlin_utils import awq_pack register_cuda_ci(est_time=10, stage="base-b-kernel-unit", runner_config="1-gpu-large") @@ -23,27 +23,6 @@ def _has_aot_awq_marlin_moe_repack() -> bool: AOT_AVAILABLE = _has_aot_awq_marlin_moe_repack() -def awq_pack( - q_w: torch.Tensor, - num_bits: int, - size_k: int, - size_n: int, -): - assert q_w.shape == (size_k, size_n) - - if num_bits == 4: - interleave = np.array([0, 2, 4, 6, 1, 3, 5, 7]) - elif num_bits == 8: - interleave = np.array([0, 2, 1, 3]) - else: - raise Exception("num_bits must be 4 or 8, got {}".format(num_bits)) - - q_w = q_w.reshape((-1, len(interleave)))[:, interleave].ravel() - q_w = q_w.reshape((-1, size_n)).contiguous() - - return pack_cols(q_w, num_bits, size_k, size_n) - - @pytest.mark.parametrize("num_bits", [4]) @pytest.mark.parametrize("num_experts", [2, 4, 8]) @pytest.mark.parametrize("k_tiles,n_tiles", [(1, 1), (2, 2), (4, 4)]) diff --git a/test/registered/kernels/ops/quantization/test_awq_marlin_repack.py b/test/registered/kernels/ops/quantization/test_awq_marlin_repack.py index bc3a8526e..10b8efa83 100644 --- a/test/registered/kernels/ops/quantization/test_awq_marlin_repack.py +++ b/test/registered/kernels/ops/quantization/test_awq_marlin_repack.py @@ -1,6 +1,5 @@ import sys -import numpy as np import pytest import torch from sgl_kernel.scalar_type import scalar_types @@ -8,9 +7,9 @@ from sgl_kernel.scalar_type import scalar_types from sglang.kernels.ops.quantization.awq_marlin_repack import ( awq_marlin_repack as jit_awq_marlin_repack, ) -from sglang.srt.layers.quantization.utils import pack_cols, quantize_weights +from sglang.srt.layers.quantization.utils import quantize_weights from sglang.test.ci.ci_register import register_cuda_ci -from sglang.test.test_marlin_utils import get_weight_perm, marlin_weights +from sglang.test.test_marlin_utils import awq_pack, get_weight_perm, marlin_weights register_cuda_ci(est_time=10, stage="base-b-kernel-unit", runner_config="1-gpu-large") @@ -24,27 +23,6 @@ def _has_aot_awq_marlin_repack() -> bool: AOT_AVAILABLE = _has_aot_awq_marlin_repack() -def awq_pack( - q_w: torch.Tensor, - num_bits: int, - size_k: int, - size_n: int, -): - assert q_w.shape == (size_k, size_n) - - if num_bits == 4: - interleave = np.array([0, 2, 4, 6, 1, 3, 5, 7]) - elif num_bits == 8: - interleave = np.array([0, 2, 1, 3]) - else: - raise Exception("num_bits must be 4 or 8, got {}".format(num_bits)) - - q_w = q_w.reshape((-1, len(interleave)))[:, interleave].ravel() - q_w = q_w.reshape((-1, size_n)).contiguous() - - return pack_cols(q_w, num_bits, size_k, size_n) - - @pytest.mark.parametrize("num_bits", [4, 8]) @pytest.mark.parametrize("k_tiles,n_tiles", [(1, 1), (2, 2), (4, 4)]) @pytest.mark.parametrize("group_size", [16, 32])