dsv4.1: standalone kernels and Python wrappers (#39646)
This commit is contained in:
Liangsheng Yin
@@ -0,0 +1,94 @@
|
||||
import unittest
|
||||
|
||||
import torch
|
||||
|
||||
from sglang.kernels.ops.attention.dsv4.elementwise import fused_rope_inplace
|
||||
from sglang.kernels.ops.attention.dsv4.q_rope_store import q_rope_store
|
||||
from sglang.test.ci.ci_register import register_cuda_ci
|
||||
from sglang.test.test_utils import CustomTestCase
|
||||
|
||||
register_cuda_ci(est_time=40, stage="base-b-kernel-unit", runner_config="1-gpu-large")
|
||||
|
||||
|
||||
class TestQRopeStore(CustomTestCase):
|
||||
def test_exact_output_and_padding(self):
|
||||
torch.manual_seed(911)
|
||||
freqs = torch.polar(
|
||||
torch.ones(8192, 32, device="cuda"), torch.randn(8192, 32, device="cuda")
|
||||
)
|
||||
for rows in (1, 2, 5, 6, 8):
|
||||
for heads in (8, 16, 32):
|
||||
for dtype in (torch.int32, torch.int64):
|
||||
q = torch.randn(
|
||||
rows, heads + 1, 512, device="cuda", dtype=torch.bfloat16
|
||||
)[:, :heads]
|
||||
padding = torch.full(
|
||||
(rows, 64, 512), 7.0, device="cuda", dtype=q.dtype
|
||||
)
|
||||
output = padding[:, :heads]
|
||||
positions = torch.randint(
|
||||
0, 8192, (rows,), device="cuda", dtype=dtype
|
||||
)
|
||||
original = q.clone()
|
||||
expected = q.clone()
|
||||
fused_rope_inplace(expected[..., 448:], None, freqs, positions)
|
||||
q_rope_store(q, output, freqs, positions)
|
||||
torch.testing.assert_close(output, expected, rtol=0, atol=0)
|
||||
torch.testing.assert_close(q, original, rtol=0, atol=0)
|
||||
self.assertTrue((padding[:, heads:] == 7).all().item())
|
||||
|
||||
def test_large_prefill_exact_output_and_padding(self):
|
||||
torch.manual_seed(911)
|
||||
freqs = torch.polar(
|
||||
torch.ones(8192, 32, device="cuda"), torch.randn(8192, 32, device="cuda")
|
||||
)
|
||||
for rows in (4096, 4097, 65536):
|
||||
for dtype in (torch.int32, torch.int64):
|
||||
with self.subTest(rows=rows, dtype=dtype):
|
||||
q = torch.randn(rows, 17, 512, device="cuda", dtype=torch.bfloat16)[
|
||||
:, :16
|
||||
]
|
||||
original = q.clone()
|
||||
expected = q.clone()
|
||||
padding = torch.full(
|
||||
(rows, 64, 512), 7.0, device="cuda", dtype=q.dtype
|
||||
)
|
||||
positions = torch.randint(
|
||||
0, 8192, (rows,), device="cuda", dtype=dtype
|
||||
)
|
||||
fused_rope_inplace(expected[..., 448:], None, freqs, positions)
|
||||
q_rope_store(q, padding[:, :16], freqs, positions)
|
||||
torch.testing.assert_close(
|
||||
padding[:, :16], expected, rtol=0, atol=0
|
||||
)
|
||||
torch.testing.assert_close(q, original, rtol=0, atol=0)
|
||||
self.assertTrue((padding[:, 16:] == 7).all().item())
|
||||
|
||||
def _check_graph_replay(self, rows):
|
||||
q = torch.randn(rows, 16, 512, device="cuda", dtype=torch.bfloat16)
|
||||
output = torch.zeros(rows, 64, 512, device="cuda", dtype=q.dtype)[:, :16]
|
||||
freqs = torch.polar(
|
||||
torch.ones(8192, 32, device="cuda"), torch.randn(8192, 32, device="cuda")
|
||||
)
|
||||
positions = torch.arange(rows, device="cuda") % 8192
|
||||
q_rope_store(q, output, freqs, positions)
|
||||
graph = torch.cuda.CUDAGraph()
|
||||
with torch.cuda.graph(graph):
|
||||
q_rope_store(q, output, freqs, positions)
|
||||
for _ in range(3):
|
||||
q.normal_()
|
||||
positions.random_(0, 8192)
|
||||
graph.replay()
|
||||
expected = q.clone()
|
||||
fused_rope_inplace(expected[..., 448:], None, freqs, positions)
|
||||
torch.testing.assert_close(output, expected, rtol=0, atol=0)
|
||||
|
||||
def test_graph_replay(self):
|
||||
self._check_graph_replay(6)
|
||||
|
||||
def test_large_prefill_graph_replay(self):
|
||||
self._check_graph_replay(4097)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -0,0 +1,167 @@
|
||||
import sys
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
from sglang.srt.runtime_context import get_platform
|
||||
from sglang.test.ci.ci_register import register_cuda_ci
|
||||
|
||||
register_cuda_ci(est_time=25, stage="base-b-kernel-unit", runner_config="4-gpu-b200")
|
||||
|
||||
pytestmark = pytest.mark.skipif(
|
||||
not torch.cuda.is_available()
|
||||
or torch.version.cuda is None
|
||||
or not get_platform().is_blackwell,
|
||||
reason="the FlashMLA split-KV schedule is Blackwell-only here",
|
||||
)
|
||||
|
||||
H_Q, D_QK, D_V = 64, 512, 512
|
||||
# The only cache format this build's sparse decode takes for both caches, and a
|
||||
# page size that keeps page * bytes_per_token a multiple of 576.
|
||||
BYTES_PER_TOKEN, PAGE = 584, 288
|
||||
BLOCK_SIZE_N, FIXED_OVERHEAD = 64, 5
|
||||
|
||||
|
||||
def _cache():
|
||||
n = 8192 // PAGE + 16
|
||||
return torch.randint(
|
||||
0, 200, (n, PAGE, 1, BYTES_PER_TOKEN), device="cuda", dtype=torch.uint8
|
||||
)
|
||||
|
||||
|
||||
def _call(kv, b, s_q, topk, topk_length, *, extra=None, meta=None):
|
||||
import sgl_kernel.flash_mla as flash_mla
|
||||
|
||||
g = torch.Generator(device="cuda").manual_seed(b * 7 + s_q * 13 + topk)
|
||||
q = torch.randn(
|
||||
(b, s_q, H_Q, D_QK), device="cuda", dtype=torch.bfloat16, generator=g
|
||||
)
|
||||
indices = torch.randint(
|
||||
0, 5120, (b, s_q, topk), device="cuda", dtype=torch.int32, generator=g
|
||||
)
|
||||
sink = torch.randn((H_Q,), device="cuda", dtype=torch.float32, generator=g)
|
||||
kwargs = {}
|
||||
if extra is not None:
|
||||
extra_kv, extra_indices, extra_topk_length = extra
|
||||
kwargs = dict(
|
||||
extra_k_cache=extra_kv,
|
||||
extra_indices_in_kvcache=extra_indices,
|
||||
extra_topk_length=extra_topk_length,
|
||||
)
|
||||
sched = flash_mla.FlashMLASchedMeta()
|
||||
if meta is not None:
|
||||
sched.tile_scheduler_metadata, sched.num_splits = meta
|
||||
out, lse = flash_mla.flash_mla_with_kvcache(
|
||||
q,
|
||||
kv,
|
||||
None,
|
||||
None,
|
||||
D_V,
|
||||
sched,
|
||||
indices=indices,
|
||||
is_fp8_kvcache=True,
|
||||
softmax_scale=0.1,
|
||||
causal=False,
|
||||
topk_length=topk_length,
|
||||
attn_sink=sink,
|
||||
**kwargs,
|
||||
)
|
||||
torch.cuda.synchronize()
|
||||
return out, lse, sched
|
||||
|
||||
|
||||
def _ours(
|
||||
like_meta, like_splits, topk_length, topk, *, extra_topk_length=None, extra_topk=0
|
||||
):
|
||||
from sglang.kernels.ops.attention.dsv4.flashmla_sched_meta import (
|
||||
flashmla_sched_meta,
|
||||
)
|
||||
|
||||
meta = torch.empty_like(like_meta)
|
||||
splits = torch.empty_like(like_splits)
|
||||
flashmla_sched_meta(
|
||||
meta,
|
||||
splits,
|
||||
topk_length=topk_length,
|
||||
extra_topk_length=extra_topk_length,
|
||||
block_size_n=BLOCK_SIZE_N,
|
||||
fixed_overhead_num_blocks=FIXED_OVERHEAD,
|
||||
topk=topk,
|
||||
extra_topk=extra_topk,
|
||||
)
|
||||
return meta, splits
|
||||
|
||||
|
||||
def _lengths(b, topk, mode, seed):
|
||||
g = torch.Generator(device="cuda").manual_seed(seed)
|
||||
if mode == "full":
|
||||
return torch.full((b,), topk, device="cuda", dtype=torch.int32)
|
||||
if mode == "ones":
|
||||
return torch.ones((b,), device="cuda", dtype=torch.int32)
|
||||
if mode == "zeros":
|
||||
return torch.zeros((b,), device="cuda", dtype=torch.int32)
|
||||
lengths = torch.randint(
|
||||
0, topk + 1, (b,), device="cuda", dtype=torch.int32, generator=g
|
||||
)
|
||||
if mode == "mixed":
|
||||
lengths[0] = 0
|
||||
lengths[-1] = topk
|
||||
return lengths
|
||||
|
||||
|
||||
# FlashMLA's DecodingSchedMeta ends in a `_pad` word it never writes, so the
|
||||
# reference carries whatever torch::empty left there.
|
||||
DEFINED = slice(0, 7)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("b", [1, 6, 64])
|
||||
@pytest.mark.parametrize("s_q", [1, 6])
|
||||
@pytest.mark.parametrize("topk", [2048])
|
||||
@pytest.mark.parametrize("mode", ["full", "random", "zeros", "ones", "mixed"])
|
||||
def test_matches_flashmla_schedule(b: int, s_q: int, topk: int, mode: str):
|
||||
kv = _cache()
|
||||
topk_length = _lengths(b, topk, mode, b * 1000 + s_q * 37 + topk + len(mode))
|
||||
_, _, sched = _call(kv, b, s_q, topk, topk_length)
|
||||
if sched.tile_scheduler_metadata is None:
|
||||
pytest.skip("FlashMLA did not split the KV for this shape")
|
||||
meta, splits = _ours(
|
||||
sched.tile_scheduler_metadata, sched.num_splits, topk_length, topk
|
||||
)
|
||||
assert torch.equal(meta[:, DEFINED], sched.tile_scheduler_metadata[:, DEFINED])
|
||||
assert torch.equal(splits, sched.num_splits)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("b", [1, 8])
|
||||
@pytest.mark.parametrize("topk,extra_topk", [(512, 512), (2048, 512), (512, 2048)])
|
||||
def test_extra_cache_schedule(b: int, topk: int, extra_topk: int):
|
||||
kv, extra_kv = _cache(), _cache()
|
||||
s_q = 1
|
||||
g = torch.Generator(device="cuda").manual_seed(b + topk + extra_topk)
|
||||
topk_length = _lengths(b, topk, "random", b + topk)
|
||||
extra_topk_length = torch.randint(
|
||||
1, extra_topk + 1, (b,), device="cuda", dtype=torch.int32, generator=g
|
||||
)
|
||||
extra_indices = torch.randint(
|
||||
0, 4096, (b, s_q, extra_topk), device="cuda", dtype=torch.int32, generator=g
|
||||
)
|
||||
extra = (extra_kv, extra_indices, extra_topk_length)
|
||||
ref_out, ref_lse, sched = _call(kv, b, s_q, topk, topk_length, extra=extra)
|
||||
if sched.tile_scheduler_metadata is None:
|
||||
pytest.skip("FlashMLA did not split the KV for this shape")
|
||||
meta, splits = _ours(
|
||||
sched.tile_scheduler_metadata,
|
||||
sched.num_splits,
|
||||
topk_length,
|
||||
topk,
|
||||
extra_topk_length=extra_topk_length,
|
||||
extra_topk=extra_topk,
|
||||
)
|
||||
assert torch.equal(meta[:, DEFINED], sched.tile_scheduler_metadata[:, DEFINED])
|
||||
assert torch.equal(splits, sched.num_splits)
|
||||
out, lse, _ = _call(kv, b, s_q, topk, topk_length, extra=extra, meta=(meta, splits))
|
||||
assert torch.equal(out.view(torch.int16), ref_out.view(torch.int16))
|
||||
assert torch.equal(lse.view(torch.int32), ref_lse.view(torch.int32))
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
sys.exit(pytest.main([__file__]))
|
||||
@@ -0,0 +1,153 @@
|
||||
"""MXFP8 epilogues stay bitwise identical to FlashInfer's standalone quantizer."""
|
||||
|
||||
import sys
|
||||
|
||||
import flashinfer
|
||||
import pytest
|
||||
import torch
|
||||
from flashinfer import mxfp8_quantize
|
||||
|
||||
from sglang.kernels.ops.attention.dsv4.wo_a_bf16 import (
|
||||
_quantize_partial,
|
||||
_wo_a_reduce,
|
||||
wo_a_bf16_small_batch,
|
||||
wo_a_bf16_small_batch_mxfp8,
|
||||
)
|
||||
from sglang.kernels.ops.layernorm.mxfp8_epilogue import rmsnorm_mxfp8
|
||||
from sglang.srt.runtime_context import get_platform
|
||||
from sglang.test.ci.ci_register import register_cuda_ci
|
||||
|
||||
register_cuda_ci(est_time=120, stage="base-b-kernel-unit", runner_config="4-gpu-b200")
|
||||
|
||||
pytestmark = pytest.mark.skipif(
|
||||
not torch.cuda.is_available()
|
||||
or torch.version.cuda is None
|
||||
or not get_platform().is_blackwell,
|
||||
reason="the MXFP8 reference quantizer is Blackwell-only",
|
||||
)
|
||||
|
||||
DEVICE = "cuda"
|
||||
|
||||
HIDDEN = 5120
|
||||
STREAMS = 4
|
||||
|
||||
|
||||
@pytest.mark.parametrize("m", [1, 2, 5, 6, 8])
|
||||
@pytest.mark.parametrize("scale", [1e-3, 1.0, 1e3])
|
||||
@pytest.mark.parametrize("backend", ["cuda", "cute-dsl"])
|
||||
def test_bitwise_identical_to_norm_then_quantize(m: int, scale: float, backend: str):
|
||||
from sglang.kernels.ops.layernorm.hc_combine_norm import (
|
||||
hc_combine_norm,
|
||||
hc_combine_norm_mxfp8,
|
||||
)
|
||||
from sglang.srt.layers.quantization.fp8_utils import flashinfer_mxfp8_quantize
|
||||
|
||||
g = torch.Generator(device="cuda").manual_seed(m * 31 + int(scale * 1000))
|
||||
x = (
|
||||
torch.randn(
|
||||
(m, STREAMS * HIDDEN), device="cuda", dtype=torch.bfloat16, generator=g
|
||||
)
|
||||
* scale
|
||||
)
|
||||
pre = torch.randn(
|
||||
(m, STREAMS), device="cuda", dtype=torch.bfloat16, generator=g
|
||||
).contiguous()
|
||||
w = torch.randn((HIDDEN,), device="cuda", dtype=torch.bfloat16, generator=g)
|
||||
eps = 1e-6
|
||||
|
||||
y_ref = hc_combine_norm(x, pre, w, eps)
|
||||
q_ref, sf_ref = flashinfer_mxfp8_quantize(y_ref, True, 32, backend)
|
||||
y, q, sf = hc_combine_norm_mxfp8(x, pre, w, eps)
|
||||
|
||||
assert torch.equal(y, y_ref)
|
||||
assert torch.equal(
|
||||
q.reshape(-1).view(torch.uint8), q_ref.reshape(-1).view(torch.uint8)
|
||||
)
|
||||
assert sf.shape == sf_ref.reshape(-1).shape
|
||||
assert torch.equal(sf, sf_ref.reshape(-1))
|
||||
|
||||
|
||||
def check(x, w, got):
|
||||
y, q, sf = got
|
||||
expected = flashinfer.norm.rmsnorm(x, w, 1e-6)
|
||||
eq, esf = flashinfer.mxfp8_quantize(expected, is_sf_swizzled_layout=True)
|
||||
assert torch.equal(y.view(torch.int16), expected.view(torch.int16))
|
||||
assert torch.equal(q.view(torch.uint8), eq.view(torch.uint8))
|
||||
m = x.shape[0]
|
||||
g = torch.arange(40, device=x.device)
|
||||
row = torch.arange(m, device=x.device)[:, None]
|
||||
offsets = (g // 4) * 512 + ((row % 32) * 4 + row // 32) * 4 + g % 4
|
||||
assert torch.equal(sf[offsets], esf.flatten()[offsets])
|
||||
pad = torch.ones_like(sf, dtype=torch.bool)
|
||||
pad[offsets] = False
|
||||
assert torch.count_nonzero(sf[pad]) == 0
|
||||
|
||||
|
||||
@pytest.mark.parametrize("m", [1, 5, 6, 8])
|
||||
@pytest.mark.parametrize("stride", [1280, 1792])
|
||||
def test_dynamic_graph(m, stride):
|
||||
torch.manual_seed(941)
|
||||
x = torch.randn((m, stride), device="cuda", dtype=torch.bfloat16)[:, :1280]
|
||||
w = torch.randn(1280, device="cuda", dtype=torch.bfloat16)
|
||||
for _ in range(3):
|
||||
check(x, w, rmsnorm_mxfp8(x, w, 1e-6))
|
||||
graph = torch.cuda.CUDAGraph()
|
||||
with torch.cuda.graph(graph):
|
||||
got = rmsnorm_mxfp8(x, w, 1e-6)
|
||||
for scale in [0, 1e-4, 1, 100]:
|
||||
x.copy_(torch.randn_like(x) * scale)
|
||||
w.copy_(torch.randn_like(w))
|
||||
graph.replay()
|
||||
check(x, w, got)
|
||||
|
||||
|
||||
def test_wo_a_partial_quant_matches_quantize():
|
||||
torch.manual_seed(0)
|
||||
for rows in range(2, 9):
|
||||
for magnitude in (0.0, 1e-37, 1e-7, 1.0, 448.0, 1e10):
|
||||
partial = torch.randn(8, rows, 2, 1024, device=DEVICE) * magnitude
|
||||
bf16 = torch.empty(rows, 2048, dtype=torch.bfloat16, device=DEVICE)
|
||||
_wo_a_reduce[(rows * 8,)](partial, bf16, rows * 2048, num_warps=4)
|
||||
expected_q, expected_s = mxfp8_quantize(bf16, True, alignment=32)
|
||||
actual_q, actual_s = _quantize_partial(partial)
|
||||
torch.testing.assert_close(
|
||||
actual_q.view(torch.uint8),
|
||||
expected_q.view(torch.uint8),
|
||||
rtol=0,
|
||||
atol=0,
|
||||
)
|
||||
torch.testing.assert_close(actual_s, expected_s, rtol=0, atol=0)
|
||||
|
||||
x = torch.randn(rows, 64, 512, device=DEVICE, dtype=torch.bfloat16)[
|
||||
:, :16
|
||||
].view(rows, 2, 4096)
|
||||
wo_a = torch.randn(2, 1024, 4096, device=DEVICE, dtype=torch.bfloat16) * 0.02
|
||||
bf16 = wo_a_bf16_small_batch(x, wo_a).flatten(1)
|
||||
q, s = wo_a_bf16_small_batch_mxfp8(x, wo_a)
|
||||
expected_q, expected_s = mxfp8_quantize(bf16, True, alignment=32)
|
||||
torch.testing.assert_close(
|
||||
q.view(torch.uint8), expected_q.view(torch.uint8), rtol=0, atol=0
|
||||
)
|
||||
torch.testing.assert_close(s, expected_s, rtol=0, atol=0)
|
||||
|
||||
partial = torch.randn(8, 6, 2, 1024, device=DEVICE)
|
||||
_quantize_partial(partial)
|
||||
graph = torch.cuda.CUDAGraph()
|
||||
with torch.cuda.graph(graph):
|
||||
q, s = _quantize_partial(partial)
|
||||
for _ in range(3):
|
||||
partial.normal_()
|
||||
# The replay must regenerate scale padding as well as live rows.
|
||||
s.fill_(255)
|
||||
graph.replay()
|
||||
bf16 = torch.empty(6, 2048, dtype=torch.bfloat16, device=DEVICE)
|
||||
_wo_a_reduce[(48,)](partial, bf16, 6 * 2048, num_warps=4)
|
||||
eq, es = mxfp8_quantize(bf16, True, alignment=32)
|
||||
torch.testing.assert_close(
|
||||
q.view(torch.uint8), eq.view(torch.uint8), rtol=0, atol=0
|
||||
)
|
||||
torch.testing.assert_close(s, es, rtol=0, atol=0)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
sys.exit(pytest.main([__file__]))
|
||||
Reference in New Issue
Block a user