[Cleanup] Deduplicate kernel tests, diffusion fixtures and benchmark helpers (#40265)
This commit is contained in:
Xiaoyu Zhang
@@ -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
|
||||
|
||||
@@ -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()
|
||||
@@ -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,
|
||||
|
||||
@@ -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)
|
||||
|
||||
@@ -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)
|
||||
|
||||
@@ -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)
|
||||
|
||||
@@ -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
|
||||
):
|
||||
|
||||
@@ -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"]
|
||||
|
||||
@@ -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
|
||||
@@ -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)
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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():
|
||||
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().to(device=inputs.device),
|
||||
norm.bias.float().to(device=inputs.device),
|
||||
norm.weight.float(),
|
||||
norm.bias.float(),
|
||||
norm.eps,
|
||||
).to(inputs.dtype)
|
||||
|
||||
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"]
|
||||
|
||||
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]
|
||||
|
||||
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)
|
||||
|
||||
@@ -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
|
||||
|
||||
|
||||
@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
|
||||
lu_mod, "apply_group_norm_silu", side_effect=lambda x, norm, act: act(norm(x))
|
||||
):
|
||||
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)
|
||||
return upsampler(latent)
|
||||
|
||||
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)
|
||||
|
||||
|
||||
@@ -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):
|
||||
|
||||
@@ -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 (
|
||||
|
||||
@@ -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)
|
||||
@@ -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, (
|
||||
|
||||
@@ -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
|
||||
)
|
||||
|
||||
|
||||
@@ -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
|
||||
)
|
||||
|
||||
|
||||
@@ -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
|
||||
)
|
||||
|
||||
# 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"
|
||||
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
|
||||
)
|
||||
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"
|
||||
|
||||
|
||||
# =============================================================================
|
||||
@@ -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}",
|
||||
)
|
||||
|
||||
|
||||
|
||||
@@ -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
|
||||
|
||||
|
||||
|
||||
@@ -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(
|
||||
"shape,gate_shape",
|
||||
[((1, 9, 64), (1, 1, 64)), ((1, 9, 64), (1, 9, 64)), PER_TOKEN_GATE_CASES[0]],
|
||||
)
|
||||
|
||||
|
||||
@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)
|
||||
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):
|
||||
|
||||
@@ -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)])
|
||||
|
||||
@@ -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])
|
||||
|
||||
Reference in New Issue
Block a user