[2/3] [EAGLE] perf: Fuse TP vocab-parallel embedding (#30948)
This commit is contained in:
Kaixi
@@ -0,0 +1,185 @@
|
||||
import torch
|
||||
import torch.nn.functional as F
|
||||
import triton
|
||||
import triton.testing
|
||||
|
||||
from sglang.jit_kernel.benchmark.utils import get_benchmark_range, run_benchmark
|
||||
from sglang.kernels.ops.embeddings.vocab_parallel_embedding import (
|
||||
vocab_parallel_embedding,
|
||||
)
|
||||
from sglang.srt.layers.vocab_parallel_embedding import get_masked_input_and_mask
|
||||
from sglang.test.ci.ci_register import register_cuda_ci
|
||||
|
||||
register_cuda_ci(
|
||||
est_time=10, stage="base-b-kernel-benchmark", runner_config="1-gpu-large"
|
||||
)
|
||||
|
||||
# Key order must match the perf_report x_names.
|
||||
DEFAULTS = dict(
|
||||
batch_size=120,
|
||||
hidden_size=6144,
|
||||
vocab_size=128256,
|
||||
tp_size=8,
|
||||
token_pattern="uniform",
|
||||
dtype="bf16",
|
||||
)
|
||||
|
||||
# One-at-a-time star sweep around DEFAULTS (the full product would be 1170
|
||||
# configs). Each entry overrides one field (or one coupled field group).
|
||||
SWEEPS = [
|
||||
(
|
||||
"batch_size",
|
||||
get_benchmark_range(
|
||||
[1, 2, 4, 8, 16, 32, 64, 120, 256, 512, 1024, 2048, 4096],
|
||||
ci_range=[1, 120],
|
||||
),
|
||||
),
|
||||
("hidden_size", get_benchmark_range([4096, 6144, 7168], ci_range=[])),
|
||||
(
|
||||
("vocab_size", "tp_size"),
|
||||
get_benchmark_range(
|
||||
[(32000, 4), (32000, 8), (128256, 4), (128256, 8), (154880, 8)],
|
||||
ci_range=[],
|
||||
),
|
||||
),
|
||||
(
|
||||
"token_pattern",
|
||||
get_benchmark_range(["uniform", "all_local", "all_remote"], ci_range=[]),
|
||||
),
|
||||
("dtype", get_benchmark_range(["bf16", "fp16"], ci_range=[])),
|
||||
]
|
||||
|
||||
|
||||
def _make_benchmark_configs():
|
||||
# Dict keying dedupes the all-defaults config each sweep re-produces.
|
||||
configs = {}
|
||||
for keys, values in SWEEPS:
|
||||
for value in values:
|
||||
override = (
|
||||
dict(zip(keys, value)) if isinstance(keys, tuple) else {keys: value}
|
||||
)
|
||||
config = {**DEFAULTS, **override}
|
||||
configs[tuple(config.values())] = None
|
||||
return list(configs)
|
||||
|
||||
|
||||
BENCHMARK_CONFIGS = _make_benchmark_configs()
|
||||
|
||||
|
||||
def _dtype_from_name(dtype: str) -> torch.dtype:
|
||||
if dtype == "bf16":
|
||||
return torch.bfloat16
|
||||
if dtype == "fp16":
|
||||
return torch.float16
|
||||
raise ValueError(f"Unknown dtype: {dtype}")
|
||||
|
||||
|
||||
def _make_input_ids(
|
||||
batch_size: int, vocab_size: int, tp_size: int, token_pattern: str
|
||||
) -> torch.Tensor:
|
||||
assert vocab_size % tp_size == 0
|
||||
per_partition = vocab_size // tp_size
|
||||
if token_pattern == "uniform":
|
||||
return torch.randint(
|
||||
0, vocab_size, (batch_size,), dtype=torch.int64, device="cuda"
|
||||
)
|
||||
if token_pattern == "all_local":
|
||||
return torch.randint(
|
||||
0,
|
||||
per_partition,
|
||||
(batch_size,),
|
||||
dtype=torch.int64,
|
||||
device="cuda",
|
||||
)
|
||||
if token_pattern == "all_remote":
|
||||
return torch.randint(
|
||||
per_partition,
|
||||
vocab_size,
|
||||
(batch_size,),
|
||||
dtype=torch.int64,
|
||||
device="cuda",
|
||||
)
|
||||
raise ValueError(f"Unknown token_pattern: {token_pattern}")
|
||||
|
||||
|
||||
def _torch_vocab_parallel_embedding(
|
||||
input_ids: torch.Tensor, weight: torch.Tensor, vocab_size: int
|
||||
):
|
||||
masked_input, input_mask = get_masked_input_and_mask(
|
||||
input_ids,
|
||||
0,
|
||||
weight.shape[0],
|
||||
0,
|
||||
vocab_size,
|
||||
vocab_size,
|
||||
)
|
||||
output = F.embedding(masked_input.long(), weight)
|
||||
output.masked_fill_(input_mask.unsqueeze(-1), 0)
|
||||
return output
|
||||
|
||||
|
||||
def _triton_vocab_parallel_embedding(
|
||||
input_ids: torch.Tensor, weight: torch.Tensor, vocab_size: int
|
||||
):
|
||||
return vocab_parallel_embedding(
|
||||
input_ids,
|
||||
weight,
|
||||
0,
|
||||
weight.shape[0],
|
||||
0,
|
||||
vocab_size,
|
||||
vocab_size,
|
||||
)
|
||||
|
||||
|
||||
@triton.testing.perf_report(
|
||||
triton.testing.Benchmark(
|
||||
x_names=[
|
||||
"batch_size",
|
||||
"hidden_size",
|
||||
"vocab_size",
|
||||
"tp_size",
|
||||
"token_pattern",
|
||||
"dtype",
|
||||
],
|
||||
x_vals=BENCHMARK_CONFIGS,
|
||||
line_arg="provider",
|
||||
line_vals=["triton", "torch"],
|
||||
line_names=["Fused Triton", "Compiled mask + embedding + masked_fill"],
|
||||
styles=[("blue", "-"), ("red", "--")],
|
||||
ylabel="us",
|
||||
plot_name="vocab-parallel-embedding-performance",
|
||||
args={},
|
||||
)
|
||||
)
|
||||
def benchmark(
|
||||
batch_size: int,
|
||||
hidden_size: int,
|
||||
vocab_size: int,
|
||||
tp_size: int,
|
||||
token_pattern: str,
|
||||
dtype: str,
|
||||
provider: str,
|
||||
):
|
||||
assert vocab_size % tp_size == 0
|
||||
torch_dtype = _dtype_from_name(dtype)
|
||||
per_partition = vocab_size // tp_size
|
||||
input_ids = _make_input_ids(batch_size, vocab_size, tp_size, token_pattern)
|
||||
weight = torch.randn((per_partition, hidden_size), dtype=torch_dtype, device="cuda")
|
||||
|
||||
expected = _torch_vocab_parallel_embedding(input_ids, weight, vocab_size)
|
||||
actual = _triton_vocab_parallel_embedding(input_ids, weight, vocab_size)
|
||||
torch.testing.assert_close(actual, expected, rtol=0, atol=0)
|
||||
|
||||
if provider == "triton":
|
||||
fn = lambda: _triton_vocab_parallel_embedding(input_ids, weight, vocab_size)
|
||||
elif provider == "torch":
|
||||
fn = lambda: _torch_vocab_parallel_embedding(input_ids, weight, vocab_size)
|
||||
else:
|
||||
raise ValueError(f"Unknown provider: {provider}")
|
||||
|
||||
return run_benchmark(fn)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
benchmark.run(print_data=True)
|
||||
@@ -68,6 +68,7 @@ EXPECTED_OPS = {
|
||||
"grammar.apply_token_bitmask_inplace_triton": {"triton"},
|
||||
"memory.alloc_extend_kernel": {"triton"},
|
||||
"attention.decode_attention_fwd": {"triton"},
|
||||
"embeddings.vocab_parallel_embedding": {"triton"},
|
||||
"kvcache.create_flashinfer_kv_indices_triton": {"triton"},
|
||||
"speculative.draft_topk1_postprocess": {"triton"},
|
||||
"speculative.gather_spec_extras": {"triton"},
|
||||
@@ -118,6 +119,7 @@ ALL_GROUPS = [
|
||||
"attention",
|
||||
"communication",
|
||||
"diffusion",
|
||||
"embeddings",
|
||||
"gemm",
|
||||
"grammar",
|
||||
"kvcache",
|
||||
|
||||
@@ -0,0 +1,142 @@
|
||||
import types
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
import torch.nn.functional as F
|
||||
|
||||
from sglang.kernels.ops.embeddings.vocab_parallel_embedding import (
|
||||
vocab_parallel_embedding,
|
||||
)
|
||||
from sglang.srt.layers.quantization.unquant import UnquantizedEmbeddingMethod
|
||||
from sglang.srt.layers.vocab_parallel_embedding import (
|
||||
VocabParallelEmbedding,
|
||||
get_masked_input_and_mask,
|
||||
)
|
||||
from sglang.test.ci.ci_register import register_cuda_ci
|
||||
|
||||
register_cuda_ci(est_time=10, stage="base-b", runner_config="1-gpu-small")
|
||||
|
||||
pytestmark = pytest.mark.skipif(
|
||||
not torch.cuda.is_available(), reason="CUDA is required for this test."
|
||||
)
|
||||
|
||||
|
||||
def _reference_embedding(input_ids, weight, cfg):
|
||||
# The kernel's contract is bit-parity with the eager production path it
|
||||
# replaces, so use that exact path as the oracle.
|
||||
masked_input, input_mask = get_masked_input_and_mask(input_ids, **cfg)
|
||||
output = F.embedding(masked_input.long(), weight)
|
||||
output.masked_fill_(input_mask.unsqueeze(-1), 0)
|
||||
return output
|
||||
|
||||
|
||||
def _run_case(input_ids, weight, cfg):
|
||||
expected = _reference_embedding(input_ids, weight, cfg)
|
||||
actual = vocab_parallel_embedding(input_ids, weight, **cfg)
|
||||
assert actual.dtype == expected.dtype
|
||||
assert actual.shape == expected.shape
|
||||
assert actual.is_contiguous()
|
||||
torch.testing.assert_close(actual, expected, rtol=0, atol=0)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float16, torch.float32])
|
||||
@pytest.mark.parametrize("input_dtype", [torch.int32, torch.int64])
|
||||
@pytest.mark.parametrize("hidden_dim", [7, 128, 6144])
|
||||
def test_vocab_parallel_embedding_no_added_vocab(dtype, input_dtype, hidden_dim):
|
||||
cfg = dict(
|
||||
org_vocab_start_index=16,
|
||||
org_vocab_end_index=32,
|
||||
num_org_vocab_padding=0,
|
||||
added_vocab_start_index=64,
|
||||
added_vocab_end_index=64,
|
||||
)
|
||||
weight = torch.randn((16, hidden_dim), dtype=dtype, device="cuda")
|
||||
# Include negative and garbage ids: the mask is the only defense against
|
||||
# them when the OOB probe is disabled, so they must produce zero rows.
|
||||
token_ids = [-100, -1, 0, 16, 17, 31, 32, 63, 2**30]
|
||||
if input_dtype == torch.int64:
|
||||
token_ids += [torch.iinfo(torch.int64).min, torch.iinfo(torch.int64).max]
|
||||
input_ids = torch.tensor(token_ids, dtype=input_dtype, device="cuda")
|
||||
_run_case(input_ids, weight, cfg)
|
||||
|
||||
|
||||
def test_vocab_parallel_embedding_added_vocab_with_padding():
|
||||
cfg = dict(
|
||||
org_vocab_start_index=10,
|
||||
org_vocab_end_index=18,
|
||||
num_org_vocab_padding=4,
|
||||
added_vocab_start_index=100,
|
||||
added_vocab_end_index=103,
|
||||
)
|
||||
weight = torch.randn((16, 257), dtype=torch.bfloat16, device="cuda")
|
||||
input_ids = torch.tensor(
|
||||
[[9, 10, 17], [18, 100, 102]], dtype=torch.int64, device="cuda"
|
||||
)
|
||||
_run_case(input_ids, weight, cfg)
|
||||
|
||||
|
||||
def test_vocab_parallel_embedding_strided_weight():
|
||||
cfg = dict(
|
||||
org_vocab_start_index=10,
|
||||
org_vocab_end_index=18,
|
||||
num_org_vocab_padding=4,
|
||||
added_vocab_start_index=100,
|
||||
added_vocab_end_index=103,
|
||||
)
|
||||
# Column slice of a wider buffer: stride(0) != hidden_dim, stride(1) == 1.
|
||||
weight = torch.randn((16, 300), dtype=torch.bfloat16, device="cuda")[:, :257]
|
||||
assert weight.stride(0) == 300 and weight.stride(1) == 1
|
||||
input_ids = torch.tensor(
|
||||
[[9, 10, 17], [18, 100, 102]], dtype=torch.int64, device="cuda"
|
||||
)
|
||||
_run_case(input_ids, weight, cfg)
|
||||
|
||||
|
||||
def test_vocab_parallel_embedding_empty_input():
|
||||
cfg = dict(
|
||||
org_vocab_start_index=0,
|
||||
org_vocab_end_index=8,
|
||||
num_org_vocab_padding=0,
|
||||
added_vocab_start_index=8,
|
||||
added_vocab_end_index=8,
|
||||
)
|
||||
weight = torch.randn((8, 64), dtype=torch.bfloat16, device="cuda")
|
||||
input_ids = torch.empty((0,), dtype=torch.int64, device="cuda")
|
||||
_run_case(input_ids, weight, cfg)
|
||||
|
||||
|
||||
def _stub_layer(**overrides):
|
||||
# The gate reads only tp_size, quant_method, and weight, so a stub avoids
|
||||
# needing distributed init for tp_size > 1.
|
||||
layer = types.SimpleNamespace(
|
||||
tp_size=2,
|
||||
quant_method=UnquantizedEmbeddingMethod(),
|
||||
weight=torch.empty((16, 32), dtype=torch.bfloat16, device="cuda"),
|
||||
)
|
||||
for name, value in overrides.items():
|
||||
setattr(layer, name, value)
|
||||
return layer
|
||||
|
||||
|
||||
def test_use_triton_embedding_gate():
|
||||
# The gate's failure direction is silent (the eager fallback is
|
||||
# numerically correct), so pin eligibility and the exclusions that define
|
||||
# the fused kernel's scope.
|
||||
gate = VocabParallelEmbedding._use_triton_embedding
|
||||
input_ids = torch.zeros((4,), dtype=torch.int64, device="cuda")
|
||||
assert gate(_stub_layer(), input_ids)
|
||||
assert not gate(_stub_layer(tp_size=1), input_ids)
|
||||
assert not gate(_stub_layer(quant_method=object()), input_ids)
|
||||
unsupported_weights = (
|
||||
torch.empty((16, 32), dtype=torch.bfloat16, device="cpu"),
|
||||
torch.empty((16,), dtype=torch.bfloat16, device="cuda"),
|
||||
torch.empty((16, 64), dtype=torch.bfloat16, device="cuda")[:, ::2],
|
||||
)
|
||||
for weight in unsupported_weights:
|
||||
assert not gate(_stub_layer(weight=weight), input_ids)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
import sys
|
||||
|
||||
sys.exit(pytest.main([__file__, "-v"]))
|
||||
Reference in New Issue
Block a user